Repository: sgl-project/sglang Branch: main Commit: c82d20d48ecc Files: 4115 Total size: 40.1 MB Directory structure: gitextract_8m4rfrdr/ ├── .claude/ │ └── skills/ │ ├── add-jit-kernel/ │ │ └── SKILL.md │ ├── add-sgl-kernel/ │ │ └── SKILL.md │ ├── sglang-bisect-ci-regression/ │ │ └── SKILL.md │ └── write-sglang-test/ │ └── SKILL.md ├── .codespellrc ├── .coveragerc ├── .devcontainer/ │ ├── Dockerfile │ └── devcontainer.json ├── .github/ │ ├── CI_PERMISSIONS.json │ ├── CODEOWNERS │ ├── FOLDER_README.md │ ├── ISSUE_TEMPLATE/ │ │ ├── 1-bug-report.yml │ │ └── 2-feature-request.yml │ ├── MAINTAINER.md │ ├── actions/ │ │ ├── upload-cuda-coredumps/ │ │ │ └── action.yml │ │ └── wait-for-jobs/ │ │ └── action.yml │ ├── labeler.yml │ ├── pull_request_template.md │ ├── update_ci_permission.py │ └── workflows/ │ ├── amd-aiter-scout.yml │ ├── amd-ci-job-monitor.yml │ ├── auto-tune.yml │ ├── bot-bump-flashinfer-version.yml │ ├── bot-bump-kernel-version-to-sglang.yml │ ├── bot-bump-kernel-version.yml │ ├── bot-bump-sglang-version.yml │ ├── bot-cherry-pick.yml │ ├── cancel-pr-workflow-on-merge.yml │ ├── cancel-unfinished-pr-tests.yml │ ├── ci-coverage-overview.yml │ ├── ci-failure-monitor.yml │ ├── close-inactive-issues.yml │ ├── diffusion-ci-gt-gen.yml │ ├── execute-notebook.yml │ ├── labeler.yml │ ├── lint.yml │ ├── list-active-pr-runs.yml.yml │ ├── nightly-release-gateway.yml │ ├── nightly-test-amd-rocm720.yml │ ├── nightly-test-amd.yml │ ├── nightly-test-intel.yml │ ├── nightly-test-npu.yml │ ├── nightly-test-nvidia.yml │ ├── open-pr-copy-from-oss.yml │ ├── open-pr-copy-to-oss.yml │ ├── patch-docker-dev.yml │ ├── pr-benchmark-rust.yml │ ├── pr-gate.yml │ ├── pr-test-amd-rocm720.yml │ ├── pr-test-amd.yml │ ├── pr-test-npu.yml │ ├── pr-test-rust.yml │ ├── pr-test-xeon.yml │ ├── pr-test-xpu.yml │ ├── pr-test.yml │ ├── release-branch-cut.yml │ ├── release-docker-amd-nightly.yml │ ├── release-docker-amd-rocm720-nightly.yml │ ├── release-docker-amd.yml │ ├── release-docker-cu13-framework.yml │ ├── release-docker-dev.yml │ ├── release-docker-gateway.yml │ ├── release-docker-npu-nightly.yml │ ├── release-docker-npu.yml │ ├── release-docker-xeon.yml │ ├── release-docker.yml │ ├── release-docs.yml │ ├── release-pypi-gateway.yml │ ├── release-pypi-nightly.yml │ ├── release-pypi-pr.yml │ ├── release-pypi.yml │ ├── release-tag.yml │ ├── release-whl-kernel.yml │ ├── rerun-ut.yml │ ├── retag-docker.yml │ ├── runner-utilization.yml │ ├── slash-command-handler.yml │ ├── stress-test.yml │ └── weekly-test-nvidia.yml ├── .gitignore ├── .isort.cfg ├── .pre-commit-config.yaml ├── 3rdparty/ │ └── amd/ │ ├── profiling/ │ │ ├── PROFILING.md │ │ ├── client.sh │ │ ├── install_rpd.sh │ │ ├── loadTracer.sh │ │ ├── rpd.patch │ │ ├── rpd_profile_server_enable.patch │ │ ├── rpd_profile_server_enable_wCPU_activities.patch │ │ ├── server.sh │ │ └── torch_profiler.patch │ ├── tuning/ │ │ ├── TUNING.md │ │ └── benchmark_moe_rocm.py │ └── wheel/ │ ├── README.md │ ├── sgl-kernel/ │ │ ├── CMakeLists_rocm.txt │ │ ├── build_rocm.sh │ │ ├── rename_wheels_rocm.sh │ │ └── rocm_hipify.py │ └── sglang/ │ └── pyproject.toml ├── LICENSE ├── README.md ├── benchmark/ │ ├── asr/ │ │ ├── README.md │ │ └── bench_sglang.py │ ├── bench_attention_sink/ │ │ └── bench_attention_sink_triton.py │ ├── bench_in_batch_prefix/ │ │ └── bench_in_batch_prefix.py │ ├── bench_linear_attention/ │ │ ├── bench_gdn_decode.py │ │ └── bench_gdn_prefill.py │ ├── bench_rope/ │ │ └── benchmark_rope_index.py │ ├── benchmark_batch/ │ │ ├── benchmark_batch.py │ │ └── benchmark_tokenizer.py │ ├── benchmark_vllm_060/ │ │ └── README.md │ ├── blog_v0_2/ │ │ ├── 405b_sglang.sh │ │ ├── 405b_trt.sh │ │ ├── 405b_vllm.sh │ │ ├── README.md │ │ └── config.md │ ├── boolq/ │ │ ├── README.md │ │ ├── bench_sglang.py │ │ ├── convert_parquet_to_json.py │ │ └── parquet_to_json.sh │ ├── ceval/ │ │ ├── README.md │ │ └── bench_sglang.py │ ├── deepseek_v3/ │ │ └── README.md │ ├── dspy/ │ │ ├── README.md │ │ └── bench_dspy_intro.py │ ├── fla/ │ │ └── benchmark_layernorm_gated.py │ ├── generative_agents/ │ │ ├── README.md │ │ ├── agent_functions.py │ │ ├── bench_other.py │ │ └── bench_sglang.py │ ├── gpt_oss/ │ │ └── README.md │ ├── gsm8k/ │ │ ├── README.md │ │ ├── bench_other.py │ │ └── bench_sglang.py │ ├── hellaswag/ │ │ ├── README.md │ │ ├── bench_other.py │ │ └── bench_sglang.py │ ├── hf3fs/ │ │ ├── bench.sh │ │ ├── bench_client.py │ │ ├── bench_storage.py │ │ └── bench_zerocopy.py │ ├── hicache/ │ │ ├── README.md │ │ ├── bench_long_context.py │ │ ├── bench_mix.py │ │ ├── bench_mix.sh │ │ ├── bench_multiturn.py │ │ ├── bench_serving.py │ │ ├── data_processing.py │ │ ├── download.sh │ │ ├── nextqa.py │ │ └── perf.py │ ├── json_decode_regex/ │ │ ├── README.md │ │ ├── bench_other.py │ │ ├── bench_sglang.py │ │ └── build_dataset.py │ ├── json_jump_forward/ │ │ ├── README.md │ │ ├── bench_other.py │ │ ├── bench_sglang.py │ │ ├── build_dataset.py │ │ └── dataset.txt │ ├── json_schema/ │ │ ├── README.md │ │ └── bench_sglang.py │ ├── kernels/ │ │ ├── all_reduce/ │ │ │ ├── benchmark_aiter.py │ │ │ ├── benchmark_all_reduce.py │ │ │ ├── benchmark_fused_ar_rms_amd.py │ │ │ ├── benchmark_mscclpp.py │ │ │ └── benchmark_torch_symm_mem.py │ │ ├── decoding_attention_triton/ │ │ │ └── triton_flashinfer_cudnn.py │ │ ├── deepep/ │ │ │ ├── deepep_utils.py │ │ │ └── tuning_deepep.py │ │ ├── deepseek/ │ │ │ ├── README.md │ │ │ ├── benchmark_deepgemm_fp8_gemm.py │ │ │ ├── benchmark_deepgemm_fp8_gemm_blackwell.py │ │ │ └── benchmark_deepgemm_fp8_group_gemm.py │ │ ├── elementwise/ │ │ │ └── benchmark_concat_mla.py │ │ ├── flashinfer_allreduce_fusion/ │ │ │ ├── README.md │ │ │ └── benchmark_fused_collective.py │ │ ├── fused_moe_triton/ │ │ │ ├── README.md │ │ │ ├── benchmark_sglang_fused_moe_triton.py │ │ │ ├── benchmark_torch_compile_fused_moe.py │ │ │ ├── benchmark_vllm_vs_sglang_fused_moe_triton.py │ │ │ ├── common_utils.py │ │ │ ├── tuning_client.py │ │ │ ├── tuning_fused_moe_triton.py │ │ │ ├── tuning_fused_moe_triton_sep.py │ │ │ └── tuning_text.json │ │ ├── quantization/ │ │ │ ├── README.md │ │ │ ├── bench_fp4_quant.py │ │ │ ├── bench_int8_quant.py │ │ │ └── tuning_block_wise_kernel.py │ │ ├── scheduler_batch/ │ │ │ ├── benchmark_get_last_loc_triton.py │ │ │ └── benchmark_write_req_to_token_pool_triton.py │ │ └── sliding_window_attention_triton/ │ │ └── bench_triton_swa_kernel.py │ ├── line_retrieval/ │ │ ├── README.md │ │ ├── bench_sglang.py │ │ └── gen_data.py │ ├── llava_bench/ │ │ ├── README.md │ │ ├── bench_hf_llava_bench.sh │ │ ├── bench_hf_mme.sh │ │ ├── bench_sglang.py │ │ ├── bench_sglang_mme.sh │ │ └── download_images.py │ ├── llm_judge/ │ │ ├── README.md │ │ ├── bench_other.py │ │ └── bench_sglang.py │ ├── long_json_decode/ │ │ ├── README.md │ │ ├── bench_other.py │ │ ├── bench_sglang.py │ │ └── build_dataset.py │ ├── lora/ │ │ ├── launch_server.py │ │ └── lora_bench.py │ ├── mmlu/ │ │ ├── README.md │ │ ├── bench_other.py │ │ ├── bench_sglang.py │ │ └── download_data.sh │ ├── mmmu/ │ │ ├── README.md │ │ ├── bench_hf.py │ │ ├── bench_sglang.py │ │ ├── data_utils.py │ │ ├── eval_utils.py │ │ └── prompt_format.yaml │ ├── mtbench/ │ │ ├── README.md │ │ ├── bench_other.py │ │ ├── bench_sglang.py │ │ └── bench_sglang_eagle.py │ ├── multi_chain_reasoning/ │ │ ├── README.md │ │ ├── bench_other.py │ │ └── bench_sglang.py │ ├── multi_document_qa/ │ │ ├── README.md │ │ ├── bench_other.py │ │ ├── bench_sglang.py │ │ └── build_dataset.py │ ├── multi_turn_chat/ │ │ ├── README.md │ │ ├── bench_other.py │ │ ├── bench_sglang.py │ │ ├── data_gen.py │ │ └── long_prompt_multi_turn.py │ ├── prefill_only/ │ │ ├── bench_embeddings.py │ │ ├── bench_score.py │ │ └── util.py │ ├── react/ │ │ ├── README.md │ │ ├── bench_other.py │ │ └── bench_sglang.py │ ├── reasoning_benchmark/ │ │ ├── README.md │ │ ├── answer_extraction.py │ │ ├── bench_sglang.py │ │ └── eval_utils.py │ ├── tip_suggestion/ │ │ ├── .gitignore │ │ ├── README.md │ │ ├── bench_other.py │ │ ├── bench_sglang.py │ │ ├── lmql_funcs.py │ │ └── topic.jsonl │ ├── tree_of_thought_deep/ │ │ ├── README.md │ │ ├── bench_other.py │ │ ├── bench_sglang.py │ │ └── lmql_funcs.py │ └── tree_of_thought_v0/ │ ├── README.md │ ├── bench_other.py │ └── bench_sglang.py ├── docker/ │ ├── Dockerfile │ ├── compose.yaml │ ├── configs/ │ │ ├── .zshrc │ │ ├── opt/ │ │ │ ├── .gitconfig │ │ │ ├── .tmux.conf │ │ │ └── .vimrc │ │ └── yank │ ├── diffusion.Dockerfile │ ├── gateway.Dockerfile │ ├── k8s-sglang-distributed-sts.yaml │ ├── k8s-sglang-service.yaml │ ├── npu.Dockerfile │ ├── rocm.Dockerfile │ ├── sagemaker.Dockerfile │ ├── serve │ ├── xeon.Dockerfile │ └── xpu.Dockerfile ├── docs/ │ ├── Makefile │ ├── README.md │ ├── _static/ │ │ └── css/ │ │ ├── custom_log.css │ │ └── readthedocs.css │ ├── advanced_features/ │ │ ├── attention_backend.md │ │ ├── checkpoint_engine.md │ │ ├── cuda_graph_for_multi_modal_encoder.md │ │ ├── deterministic_inference.md │ │ ├── dp_dpa_smg_guide.md │ │ ├── dp_for_multi_modal_encoder.md │ │ ├── epd_disaggregation.md │ │ ├── expert_parallelism.md │ │ ├── forward_hooks.md │ │ ├── hicache.rst │ │ ├── hicache_best_practices.md │ │ ├── hicache_design.md │ │ ├── hicache_storage_runtime_attach_detach.md │ │ ├── hyperparameter_tuning.md │ │ ├── lora.ipynb │ │ ├── observability.md │ │ ├── pd_disaggregation.md │ │ ├── piecewise_cuda_graph.md │ │ ├── pipeline_parallelism.md │ │ ├── quantization.md │ │ ├── quantized_kv_cache.md │ │ ├── rfork.md │ │ ├── separate_reasoning.ipynb │ │ ├── server_arguments.md │ │ ├── sgl_model_gateway.md │ │ ├── sglang_for_rl.md │ │ ├── speculative_decoding.md │ │ ├── structured_outputs.ipynb │ │ ├── structured_outputs_for_reasoning_models.ipynb │ │ ├── tool_parser.ipynb │ │ └── vlm_query.ipynb │ ├── basic_usage/ │ │ ├── deepseek_ocr.md │ │ ├── deepseek_v3.md │ │ ├── deepseek_v32.md │ │ ├── glm45.md │ │ ├── glmv.md │ │ ├── gpt_oss.md │ │ ├── llama4.md │ │ ├── minimax_m2.md │ │ ├── native_api.ipynb │ │ ├── offline_engine_api.ipynb │ │ ├── ollama_api.md │ │ ├── openai_api.rst │ │ ├── openai_api_completions.ipynb │ │ ├── openai_api_embeddings.ipynb │ │ ├── openai_api_vision.ipynb │ │ ├── popular_model_usage.rst │ │ ├── qwen3.md │ │ ├── qwen3_5.md │ │ ├── qwen3_vl.md │ │ ├── sampling_params.md │ │ └── send_request.ipynb │ ├── conf.py │ ├── deploy.py │ ├── developer_guide/ │ │ ├── bench_serving.md │ │ ├── benchmark_and_profiling.md │ │ ├── contribution_guide.md │ │ ├── development_guide_using_docker.md │ │ ├── development_jit_kernel_guide.md │ │ ├── evaluating_new_models.md │ │ ├── release_process.md │ │ └── setup_github_runner.md │ ├── diffusion/ │ │ ├── api/ │ │ │ ├── cli.md │ │ │ ├── openai_api.md │ │ │ └── post_processing.md │ │ ├── ci_perf.md │ │ ├── compatibility_matrix.md │ │ ├── contributing.md │ │ ├── environment_variables.md │ │ ├── index.md │ │ ├── installation.md │ │ ├── performance/ │ │ │ ├── attention_backends.md │ │ │ ├── cache/ │ │ │ │ ├── cache_dit.md │ │ │ │ ├── index.md │ │ │ │ └── teacache.md │ │ │ ├── index.md │ │ │ └── profiling.md │ │ └── support_new_models.md │ ├── get_started/ │ │ └── install.md │ ├── index.rst │ ├── performance_dashboard/ │ │ ├── README.md │ │ ├── app.js │ │ ├── fetch_metrics.py │ │ ├── index.html │ │ └── server.py │ ├── platforms/ │ │ ├── amd_gpu.md │ │ ├── apple_metal.md │ │ ├── ascend_contribution_guide.md │ │ ├── ascend_npu.md │ │ ├── ascend_npu_best_practice.md │ │ ├── ascend_npu_deepseek_example.md │ │ ├── ascend_npu_environment_variables.md │ │ ├── ascend_npu_glm5_examples.md │ │ ├── ascend_npu_quantization.md │ │ ├── ascend_npu_qwen3_5_examples.md │ │ ├── ascend_npu_qwen3_examples.md │ │ ├── ascend_npu_support.rst │ │ ├── ascend_npu_support_features.md │ │ ├── ascend_npu_support_models.md │ │ ├── cpu_server.md │ │ ├── mindspore_backend.md │ │ ├── mthreads_gpu.md │ │ ├── nvidia_jetson.md │ │ ├── tpu.md │ │ └── xpu.md │ ├── references/ │ │ ├── custom_chat_template.md │ │ ├── environment_variables.md │ │ ├── faq.md │ │ ├── frontend/ │ │ │ ├── choices_methods.md │ │ │ ├── frontend_index.rst │ │ │ └── frontend_tutorial.ipynb │ │ ├── learn_more.md │ │ ├── multi_node_deployment/ │ │ │ ├── deploy_on_k8s.md │ │ │ ├── lws_pd/ │ │ │ │ ├── lws-examples/ │ │ │ │ │ ├── d-svc.yaml │ │ │ │ │ ├── d.yaml │ │ │ │ │ ├── lb.yaml │ │ │ │ │ ├── p-svc.yaml │ │ │ │ │ └── p.yaml │ │ │ │ └── lws_pd_deploy.md │ │ │ ├── multi_node.md │ │ │ ├── multi_node_index.rst │ │ │ └── rbg_pd/ │ │ │ └── deepseekv32_pd.md │ │ ├── post_training_integration.md │ │ ├── production_metrics.md │ │ ├── production_request_trace.md │ │ ├── release_lookup.rst │ │ └── torch_compile_cache.md │ ├── release_lookup/ │ │ ├── README.md │ │ ├── generate_index.py │ │ ├── index.html │ │ └── release_index.json │ ├── requirements.txt │ ├── serve.sh │ ├── supported_models/ │ │ ├── extending/ │ │ │ ├── index.rst │ │ │ ├── mindspore_models.md │ │ │ ├── modelscope.md │ │ │ ├── support_new_models.md │ │ │ └── transformers_fallback.md │ │ ├── index.rst │ │ ├── retrieval_ranking/ │ │ │ ├── classify_models.md │ │ │ ├── embedding_models.md │ │ │ ├── index.rst │ │ │ └── rerank_models.md │ │ ├── specialized/ │ │ │ ├── index.rst │ │ │ └── reward_models.md │ │ └── text_generation/ │ │ ├── diffusion_language_models.md │ │ ├── generative_models.md │ │ ├── index.rst │ │ └── multimodal_language_models.md │ └── wrap_run_llm.py ├── examples/ │ ├── assets/ │ │ └── .gitignore │ ├── chat_template/ │ │ ├── qwen3_reranker.jinja │ │ ├── qwen3_vl_reranker.jinja │ │ ├── tool_chat_template_deepseekr1.jinja │ │ ├── tool_chat_template_deepseekv3.jinja │ │ ├── tool_chat_template_deepseekv31.jinja │ │ ├── tool_chat_template_deepseekv32.jinja │ │ ├── tool_chat_template_llama3.1_json.jinja │ │ ├── tool_chat_template_llama4_pythonic.jinja │ │ └── vision_template_sarashina_vl.jinja │ ├── checkpoint_engine/ │ │ └── update.py │ ├── frontend_language/ │ │ ├── quick_start/ │ │ │ ├── anthropic_example_chat.py │ │ │ ├── anthropic_example_complete.py │ │ │ ├── azure_openai_example_chat.py │ │ │ ├── gemini_example_chat.py │ │ │ ├── gemini_example_complete.py │ │ │ ├── gemini_example_multimodal_chat.py │ │ │ ├── local_example_chat.py │ │ │ ├── local_example_complete.py │ │ │ ├── local_example_llava_next.py │ │ │ ├── openai_example_chat.py │ │ │ ├── openai_example_complete.py │ │ │ ├── openai_example_n.py │ │ │ ├── openai_example_o1.py │ │ │ ├── openrouter_example_chat.py │ │ │ ├── together_example_chat.py │ │ │ └── together_example_complete.py │ │ └── usage/ │ │ ├── chinese_regex.py │ │ ├── choices_logprob.py │ │ ├── cot_decoding.py │ │ ├── json_decode.py │ │ ├── json_logprobs.py │ │ ├── llava_video/ │ │ │ ├── srt_example_llava_v.py │ │ │ └── srt_example_llava_v.sh │ │ ├── openai_chat_speculative.py │ │ ├── openai_speculative.py │ │ ├── parallel_sample.py │ │ ├── rag_using_parea/ │ │ │ └── trace_and_evaluate_rag_using_parea.ipynb │ │ ├── readme_examples.py │ │ ├── sgl_gen_min_tokens.py │ │ ├── streaming.py │ │ └── triton/ │ │ ├── Dockerfile │ │ ├── README.md │ │ └── models/ │ │ └── character_generation/ │ │ ├── 1/ │ │ │ └── model.py │ │ └── config.pbtxt │ ├── monitoring/ │ │ ├── README.md │ │ ├── docker-compose.yaml │ │ ├── grafana/ │ │ │ ├── dashboards/ │ │ │ │ ├── config/ │ │ │ │ │ └── dashboard.yaml │ │ │ │ └── json/ │ │ │ │ └── sglang-dashboard.json │ │ │ └── datasources/ │ │ │ └── datasource.yaml │ │ ├── opentelemetry.yaml │ │ ├── prometheus.yaml │ │ └── tracing_compose.yaml │ ├── profiler/ │ │ └── nsys_profile_tools/ │ │ ├── README.md │ │ ├── gputrc2graph.py │ │ └── sglang_engine_model.json │ ├── runtime/ │ │ ├── README.md │ │ ├── engine/ │ │ │ ├── custom_server.py │ │ │ ├── embedding.py │ │ │ ├── fastapi_engine_inference.py │ │ │ ├── launch_engine.py │ │ │ ├── offline_batch_inference.py │ │ │ ├── offline_batch_inference_async.py │ │ │ ├── offline_batch_inference_eagle.py │ │ │ ├── offline_batch_inference_qwen_1m.py │ │ │ ├── offline_batch_inference_vlm.py │ │ │ ├── readme.md │ │ │ ├── save_remote_state.py │ │ │ └── save_sharded_state.py │ │ ├── hidden_states/ │ │ │ ├── hidden_states_engine.py │ │ │ └── hidden_states_server.py │ │ ├── lora.py │ │ ├── multimodal/ │ │ │ ├── llama3_llava_server.py │ │ │ ├── llava_onevision_server.py │ │ │ ├── pixtral_server.py │ │ │ └── qwen_llava_server.py │ │ ├── multimodal_embedding.py │ │ ├── openai_chat_with_response_prefill.py │ │ ├── qwen3_vl_reranker.py │ │ ├── reward_model.py │ │ ├── token_in_token_out/ │ │ │ ├── token_in_token_out_llm_engine.py │ │ │ ├── token_in_token_out_llm_server.py │ │ │ ├── token_in_token_out_vlm_engine.py │ │ │ └── token_in_token_out_vlm_server.py │ │ └── vertex_predict.py │ ├── sagemaker/ │ │ └── deploy_and_serve_endpoint.py │ └── usage/ │ └── modelopt_quantize_and_export.py ├── python/ │ ├── pyproject.toml │ ├── pyproject_cpu.toml │ ├── pyproject_npu.toml │ ├── pyproject_other.toml │ ├── pyproject_xpu.toml │ └── sglang/ │ ├── README.md │ ├── __init__.py │ ├── _mps_stub.py │ ├── _triton_stub.py │ ├── bench_offline_throughput.py │ ├── bench_one_batch.py │ ├── bench_one_batch_server.py │ ├── bench_serving.py │ ├── benchmark/ │ │ ├── __init__.py │ │ ├── bench_utils.py │ │ ├── datasets/ │ │ │ ├── __init__.py │ │ │ ├── common.py │ │ │ ├── custom.py │ │ │ ├── generated_shared_prefix.py │ │ │ ├── image.py │ │ │ ├── mmmu.py │ │ │ ├── mooncake.py │ │ │ ├── openai_dataset.py │ │ │ ├── random.py │ │ │ └── sharegpt.py │ │ └── utils.py │ ├── check_env.py │ ├── cli/ │ │ ├── __init__.py │ │ ├── generate.py │ │ ├── main.py │ │ ├── serve.py │ │ └── utils.py │ ├── compile_deep_gemm.py │ ├── eval/ │ │ ├── llama3_eval.py │ │ └── loogle_eval.py │ ├── global_config.py │ ├── jit_kernel/ │ │ ├── .clang-format │ │ ├── __init__.py │ │ ├── __main__.py │ │ ├── add_constant.py │ │ ├── awq_dequantize.py │ │ ├── awq_marlin_repack.py │ │ ├── benchmark/ │ │ │ ├── bench_awq_dequantize.py │ │ │ ├── bench_awq_marlin_moe_repack.py │ │ │ ├── bench_awq_marlin_repack.py │ │ │ ├── bench_concat_mla.py │ │ │ ├── bench_fused_add_rmsnorm.py │ │ │ ├── bench_fused_norm_scale_shift.py │ │ │ ├── bench_gptq_marlin.py │ │ │ ├── bench_gptq_marlin_repack.py │ │ │ ├── bench_hadamard.py │ │ │ ├── bench_hicache.py │ │ │ ├── bench_moe_wna16_marlin.py │ │ │ ├── bench_norm.py │ │ │ ├── bench_norm_impls.py │ │ │ ├── bench_nvfp4_blockwise_moe.py │ │ │ ├── bench_nvfp4_quant.py │ │ │ ├── bench_nvfp4_scaled_mm.py │ │ │ ├── bench_per_tensor_quant_fp8.py │ │ │ ├── bench_per_token_group_quant_8bit.py │ │ │ ├── bench_qknorm.py │ │ │ ├── bench_qknorm_across_heads.py │ │ │ ├── bench_qwen_image_modulation.py │ │ │ ├── bench_renorm.py │ │ │ ├── bench_rmsnorm.py │ │ │ ├── bench_rope.py │ │ │ ├── bench_store_cache.py │ │ │ └── utils.py │ │ ├── concat_mla.py │ │ ├── csrc/ │ │ │ ├── add_constant.cuh │ │ │ ├── diffusion/ │ │ │ │ └── timestep_embedding.cuh │ │ │ ├── elementwise/ │ │ │ │ ├── concat_mla.cuh │ │ │ │ ├── fused_add_rmsnorm.cuh │ │ │ │ ├── fused_metadata_copy.cuh │ │ │ │ ├── kvcache.cuh │ │ │ │ ├── pos_enc.cuh │ │ │ │ ├── qknorm.cuh │ │ │ │ ├── qknorm_across_heads.cuh │ │ │ │ ├── rmsnorm.cuh │ │ │ │ └── rope.cuh │ │ │ ├── fast-hadamard-transform/ │ │ │ │ ├── code_gen.py │ │ │ │ ├── fast_hadamard_transform.h │ │ │ │ ├── fast_hadamard_transform_common.h │ │ │ │ ├── fast_hadamard_transform_special.h │ │ │ │ ├── hadamard_jit.cuh │ │ │ │ └── static_switch.h │ │ │ ├── gemm/ │ │ │ │ ├── awq_dequantize.cuh │ │ │ │ ├── marlin/ │ │ │ │ │ ├── awq_marlin_repack.cuh │ │ │ │ │ ├── dequant.h │ │ │ │ │ ├── gptq_marlin.cuh │ │ │ │ │ ├── gptq_marlin_repack.cuh │ │ │ │ │ ├── kernel.h │ │ │ │ │ ├── marlin.cuh │ │ │ │ │ ├── marlin_dtypes.cuh │ │ │ │ │ └── marlin_template.h │ │ │ │ ├── marlin_moe/ │ │ │ │ │ ├── kernel.h │ │ │ │ │ ├── marlin_template.h │ │ │ │ │ └── moe_wna16_marlin.cuh │ │ │ │ ├── nvfp4/ │ │ │ │ │ ├── nvfp4_expert_quant.cuh │ │ │ │ │ ├── nvfp4_quant.cuh │ │ │ │ │ ├── nvfp4_quant_entry.cuh │ │ │ │ │ ├── nvfp4_quant_kernels.cuh │ │ │ │ │ ├── nvfp4_scaled_mm_entry.cuh │ │ │ │ │ └── nvfp4_scaled_mm_kernels.cuh │ │ │ │ ├── per_tensor_quant_fp8.cuh │ │ │ │ └── per_token_group_quant_8bit.cuh │ │ │ ├── hicache.cuh │ │ │ ├── lora/ │ │ │ │ └── moe_lora_align_kernel.cu │ │ │ ├── moe/ │ │ │ │ └── nvfp4_blockwise_moe.cuh │ │ │ ├── ngram_embedding.cuh │ │ │ └── nsa/ │ │ │ └── fused_store_index_cache.cuh │ │ ├── cutedsl_gdn.py │ │ ├── diffusion/ │ │ │ ├── cutedsl/ │ │ │ │ ├── common/ │ │ │ │ │ ├── norm_fusion.py │ │ │ │ │ └── reduce.py │ │ │ │ ├── scale_residual_norm_scale_shift.py │ │ │ │ └── utils.py │ │ │ └── triton/ │ │ │ ├── mps_fallback.py │ │ │ ├── norm.py │ │ │ ├── npu_fallback.py │ │ │ ├── rmsnorm_onepass.py │ │ │ ├── rotary.py │ │ │ └── scale_shift.py │ │ ├── flash_attention_v4.py │ │ ├── fused_metadata_copy.py │ │ ├── fused_store_index_cache.py │ │ ├── gptq_marlin.py │ │ ├── gptq_marlin_repack.py │ │ ├── hadamard.py │ │ ├── hicache.py │ │ ├── include/ │ │ │ └── sgl_kernel/ │ │ │ ├── atomic.cuh │ │ │ ├── cta.cuh │ │ │ ├── impl/ │ │ │ │ └── norm.cuh │ │ │ ├── math.cuh │ │ │ ├── runtime.cuh │ │ │ ├── scalar_type.hpp │ │ │ ├── source_location.h │ │ │ ├── tensor.h │ │ │ ├── tile.cuh │ │ │ ├── type.cuh │ │ │ ├── utils.cuh │ │ │ ├── utils.h │ │ │ ├── vec.cuh │ │ │ └── warp.cuh │ │ ├── kvcache.py │ │ ├── moe_lora_align.py │ │ ├── moe_wna16_marlin.py │ │ ├── ngram_embedding.py │ │ ├── norm.py │ │ ├── nvfp4.py │ │ ├── per_tensor_quant_fp8.py │ │ ├── per_token_group_quant_8bit.py │ │ ├── rope.py │ │ ├── tests/ │ │ │ ├── test_add_constant.py │ │ │ ├── test_awq_dequantize.py │ │ │ ├── test_awq_marlin_moe_repack.py │ │ │ ├── test_awq_marlin_repack.py │ │ │ ├── test_concat_mla.py │ │ │ ├── test_cutedsl_gdn.py │ │ │ ├── test_flash_attention_4.py │ │ │ ├── test_fused_add_rmsnorm.py │ │ │ ├── test_fused_metadata_copy.py │ │ │ ├── test_fused_norm_scale_shift.py │ │ │ ├── test_fused_store_index_cache.py │ │ │ ├── test_fused_verify_triton_gdn.py │ │ │ ├── test_gptq_marlin.py │ │ │ ├── test_gptq_marlin_repack.py │ │ │ ├── test_hadamard_jit.py │ │ │ ├── test_moe_lora_align_block_size.py │ │ │ ├── test_moe_wna16_marlin.py │ │ │ ├── test_norm_jit.py │ │ │ ├── test_nvfp4_blockwise_moe.py │ │ │ ├── test_nvfp4_gemm.py │ │ │ ├── test_nvfp4_quant.py │ │ │ ├── test_per_tensor_quant_fp8.py │ │ │ ├── test_per_token_group_quant_8bit.py │ │ │ ├── test_pos_enc.py │ │ │ ├── test_qknorm.py │ │ │ ├── test_qknorm_across_heads.py │ │ │ ├── test_qwen_image_modulation.py │ │ │ ├── test_renorm.py │ │ │ ├── test_rmsnorm.py │ │ │ ├── test_rope.py │ │ │ ├── test_store_cache.py │ │ │ └── test_timestep_embedding.py │ │ ├── timestep_embedding.py │ │ └── utils.py │ ├── lang/ │ │ ├── api.py │ │ ├── backend/ │ │ │ ├── anthropic.py │ │ │ ├── base_backend.py │ │ │ ├── litellm.py │ │ │ ├── openai.py │ │ │ ├── runtime_endpoint.py │ │ │ └── vertexai.py │ │ ├── chat_template.py │ │ ├── choices.py │ │ ├── interpreter.py │ │ ├── ir.py │ │ └── tracer.py │ ├── launch_server.py │ ├── multimodal_gen/ │ │ ├── .claude/ │ │ │ ├── CLAUDE.md │ │ │ └── skills/ │ │ │ ├── diffusion-kernel/ │ │ │ │ ├── SKILL.md │ │ │ │ ├── add-cuda-kernel.md │ │ │ │ ├── add-triton-kernel.md │ │ │ │ ├── diffusion-benchmark-and-profile.md │ │ │ │ ├── nsight-profiler.md │ │ │ │ ├── references/ │ │ │ │ │ ├── a100-optimization-guide.md │ │ │ │ │ ├── h100-optimization-guide.md │ │ │ │ │ ├── kernel-templates.md │ │ │ │ │ ├── t4-optimization-guide.md │ │ │ │ │ └── troubleshooting.md │ │ │ │ ├── scripts/ │ │ │ │ │ ├── bench_diffusion_denoise.py │ │ │ │ │ ├── bench_diffusion_rmsnorm.py │ │ │ │ │ └── diffusion_skill_env.py │ │ │ │ └── use-efficient-diffusion-kernels.md │ │ │ ├── diffusion-optimal-perf/ │ │ │ │ └── SKILL.md │ │ │ └── support-new-model/ │ │ │ └── SKILL.md │ │ ├── README.md │ │ ├── __init__.py │ │ ├── apps/ │ │ │ ├── ComfyUI_SGLDiffusion/ │ │ │ │ ├── README.md │ │ │ │ ├── __init__.py │ │ │ │ ├── core/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── generator.py │ │ │ │ │ ├── model_patcher.py │ │ │ │ │ └── server_api.py │ │ │ │ ├── executors/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── base.py │ │ │ │ │ ├── flux.py │ │ │ │ │ ├── qwen_image.py │ │ │ │ │ └── zimage.py │ │ │ │ ├── nodes.py │ │ │ │ ├── test/ │ │ │ │ │ ├── README.md │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── test_flux_pipeline.py │ │ │ │ │ ├── test_qwen_image_edit_pipeline.py │ │ │ │ │ ├── test_qwen_image_pipeline.py │ │ │ │ │ └── test_zimage_pipeline.py │ │ │ │ ├── utils.py │ │ │ │ └── workflows/ │ │ │ │ ├── flux_sgld_sp.json │ │ │ │ ├── qwen_image_sgld.json │ │ │ │ ├── sgld_image2video.json │ │ │ │ ├── sgld_text2img.json │ │ │ │ └── z-image_sgld.json │ │ │ └── webui/ │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ └── main.py │ │ ├── benchmarks/ │ │ │ ├── bench_offline_throughput.py │ │ │ ├── bench_serving.py │ │ │ ├── compare_perf.py │ │ │ └── datasets.py │ │ ├── configs/ │ │ │ ├── __init__.py │ │ │ ├── backend/ │ │ │ │ └── vmoba/ │ │ │ │ ├── wan_1.3B_77_448_832.json │ │ │ │ └── wan_1.3B_77_480_832.json │ │ │ ├── models/ │ │ │ │ ├── __init__.py │ │ │ │ ├── adapter/ │ │ │ │ │ ├── base.py │ │ │ │ │ └── ltx_2_connector.py │ │ │ │ ├── base.py │ │ │ │ ├── bridges/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── mova_dual_tower.py │ │ │ │ ├── dits/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── base.py │ │ │ │ │ ├── flux.py │ │ │ │ │ ├── glmimage.py │ │ │ │ │ ├── helios.py │ │ │ │ │ ├── hunyuan3d.py │ │ │ │ │ ├── hunyuanvideo.py │ │ │ │ │ ├── ltx_2.py │ │ │ │ │ ├── mova_audio.py │ │ │ │ │ ├── mova_video.py │ │ │ │ │ ├── qwenimage.py │ │ │ │ │ ├── sana.py │ │ │ │ │ ├── wanvideo.py │ │ │ │ │ └── zimage.py │ │ │ │ ├── encoders/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── base.py │ │ │ │ │ ├── clip.py │ │ │ │ │ ├── gemma2.py │ │ │ │ │ ├── gemma_3.py │ │ │ │ │ ├── llama.py │ │ │ │ │ ├── qwen3.py │ │ │ │ │ ├── qwen_image.py │ │ │ │ │ └── t5.py │ │ │ │ ├── vaes/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── base.py │ │ │ │ │ ├── dac.py │ │ │ │ │ ├── flux.py │ │ │ │ │ ├── glmimage.py │ │ │ │ │ ├── hunyuan3d.py │ │ │ │ │ ├── hunyuanvae.py │ │ │ │ │ ├── ltx_audio.py │ │ │ │ │ ├── ltx_video.py │ │ │ │ │ ├── qwenimage.py │ │ │ │ │ ├── sana.py │ │ │ │ │ └── wanvae.py │ │ │ │ └── vocoder/ │ │ │ │ ├── __init__.py │ │ │ │ ├── base.py │ │ │ │ └── ltx_vocoder.py │ │ │ ├── pipeline_configs/ │ │ │ │ ├── __init__.py │ │ │ │ ├── base.py │ │ │ │ ├── diffusers_generic.py │ │ │ │ ├── flux.py │ │ │ │ ├── flux_finetuned.py │ │ │ │ ├── glm_image.py │ │ │ │ ├── helios.py │ │ │ │ ├── hunyuan.py │ │ │ │ ├── hunyuan3d.py │ │ │ │ ├── ltx_2.py │ │ │ │ ├── mova.py │ │ │ │ ├── qwen_image.py │ │ │ │ ├── sana.py │ │ │ │ ├── wan.py │ │ │ │ └── zimage.py │ │ │ ├── quantization.py │ │ │ ├── sample/ │ │ │ │ ├── __init__.py │ │ │ │ ├── diffusers_generic.py │ │ │ │ ├── flux.py │ │ │ │ ├── glmimage.py │ │ │ │ ├── helios.py │ │ │ │ ├── hunyuan.py │ │ │ │ ├── hunyuan3d.py │ │ │ │ ├── ltx_2.py │ │ │ │ ├── mova.py │ │ │ │ ├── qwenimage.py │ │ │ │ ├── sampling_params.py │ │ │ │ ├── sana.py │ │ │ │ ├── teacache.py │ │ │ │ ├── wan.py │ │ │ │ └── zimage.py │ │ │ └── utils.py │ │ ├── csrc/ │ │ │ ├── attn/ │ │ │ │ └── vmoba_attn/ │ │ │ │ ├── README.md │ │ │ │ ├── setup.py │ │ │ │ ├── tests/ │ │ │ │ │ └── test_vmoba_attn.py │ │ │ │ └── vmoba/ │ │ │ │ ├── __init__.py │ │ │ │ └── vmoba.py │ │ │ └── render/ │ │ │ ├── hunyuan3d_rasterizer/ │ │ │ │ ├── __init__.py │ │ │ │ ├── rasterizer.cpp │ │ │ │ ├── rasterizer.h │ │ │ │ └── rasterizer_gpu.cu │ │ │ └── mesh_processor/ │ │ │ ├── __init__.py │ │ │ └── mesh_processor.cpp │ │ ├── docs/ │ │ │ └── quantization.md │ │ ├── envs.py │ │ ├── registry.py │ │ ├── runtime/ │ │ │ ├── cache/ │ │ │ │ ├── __init__.py │ │ │ │ ├── cache_dit_integration.py │ │ │ │ └── teacache.py │ │ │ ├── distributed/ │ │ │ │ ├── __init__.py │ │ │ │ ├── communication_op.py │ │ │ │ ├── device_communicators/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── base_device_communicator.py │ │ │ │ │ ├── cpu_communicator.py │ │ │ │ │ ├── cuda_communicator.py │ │ │ │ │ ├── pynccl.py │ │ │ │ │ └── pynccl_wrapper.py │ │ │ │ ├── group_coordinator.py │ │ │ │ ├── parallel_groups.py │ │ │ │ ├── parallel_state.py │ │ │ │ └── utils.py │ │ │ ├── entrypoints/ │ │ │ │ ├── __init__.py │ │ │ │ ├── cli/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── cli_types.py │ │ │ │ │ ├── generate.py │ │ │ │ │ ├── main.py │ │ │ │ │ ├── serve.py │ │ │ │ │ └── utils.py │ │ │ │ ├── diffusion_generator.py │ │ │ │ ├── http_server.py │ │ │ │ ├── openai/ │ │ │ │ │ ├── common_api.py │ │ │ │ │ ├── image_api.py │ │ │ │ │ ├── mesh_api.py │ │ │ │ │ ├── protocol.py │ │ │ │ │ ├── storage.py │ │ │ │ │ ├── stores.py │ │ │ │ │ ├── utils.py │ │ │ │ │ └── video_api.py │ │ │ │ ├── post_training/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── io_struct.py │ │ │ │ │ └── weights_api.py │ │ │ │ └── utils.py │ │ │ ├── launch_server.py │ │ │ ├── layers/ │ │ │ │ ├── __init__.py │ │ │ │ ├── activation.py │ │ │ │ ├── attention/ │ │ │ │ │ ├── STA_configuration.py │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── backends/ │ │ │ │ │ │ ├── __init__.py │ │ │ │ │ │ ├── aiter.py │ │ │ │ │ │ ├── aiter_sage.py │ │ │ │ │ │ ├── attention_backend.py │ │ │ │ │ │ ├── flash_attn.py │ │ │ │ │ │ ├── flash_attn_2.py │ │ │ │ │ │ ├── sage_attn.py │ │ │ │ │ │ ├── sage_attn3.py │ │ │ │ │ │ ├── sdpa.py │ │ │ │ │ │ ├── sliding_tile_attn.py │ │ │ │ │ │ ├── sparse_linear_attn.py │ │ │ │ │ │ ├── sparse_video_gen_2_attn.py │ │ │ │ │ │ ├── video_sparse_attn.py │ │ │ │ │ │ └── vmoba.py │ │ │ │ │ ├── layer.py │ │ │ │ │ ├── selector.py │ │ │ │ │ └── turbo_layer.py │ │ │ │ ├── custom_op.py │ │ │ │ ├── elementwise.py │ │ │ │ ├── layernorm.py │ │ │ │ ├── linear.py │ │ │ │ ├── lora/ │ │ │ │ │ └── linear.py │ │ │ │ ├── mlp.py │ │ │ │ ├── quantization/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── configs/ │ │ │ │ │ │ ├── base_config.py │ │ │ │ │ │ └── nunchaku_config.py │ │ │ │ │ ├── fp8.py │ │ │ │ │ ├── modelslim.py │ │ │ │ │ └── nunchaku_linear.py │ │ │ │ ├── rotary_embedding/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── base.py │ │ │ │ │ ├── factory.py │ │ │ │ │ ├── mrope.py │ │ │ │ │ └── utils.py │ │ │ │ ├── usp.py │ │ │ │ ├── utils.py │ │ │ │ ├── visual_embedding.py │ │ │ │ └── vocab_parallel_embedding.py │ │ │ ├── loader/ │ │ │ │ ├── component_loaders/ │ │ │ │ │ ├── adapter_loader.py │ │ │ │ │ ├── bridge_loader.py │ │ │ │ │ ├── component_loader.py │ │ │ │ │ ├── image_encoder_loader.py │ │ │ │ │ ├── scheduler_loader.py │ │ │ │ │ ├── text_encoder_loader.py │ │ │ │ │ ├── transformer_loader.py │ │ │ │ │ ├── vae_loader.py │ │ │ │ │ ├── vl_encoder_loader.py │ │ │ │ │ └── vocoder_loader.py │ │ │ │ ├── fsdp_load.py │ │ │ │ ├── utils.py │ │ │ │ ├── weight_utils.py │ │ │ │ └── weights_updater.py │ │ │ ├── managers/ │ │ │ │ ├── forward_context.py │ │ │ │ ├── gpu_worker.py │ │ │ │ └── scheduler.py │ │ │ ├── models/ │ │ │ │ ├── __init__.py │ │ │ │ ├── adapter/ │ │ │ │ │ └── ltx_2_connector.py │ │ │ │ ├── bridges/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── mova_dual_tower.py │ │ │ │ ├── dits/ │ │ │ │ │ ├── base.py │ │ │ │ │ ├── causal_wanvideo.py │ │ │ │ │ ├── flux.py │ │ │ │ │ ├── flux_2.py │ │ │ │ │ ├── glm_image.py │ │ │ │ │ ├── helios.py │ │ │ │ │ ├── hunyuan3d.py │ │ │ │ │ ├── hunyuanvideo.py │ │ │ │ │ ├── ltx_2.py │ │ │ │ │ ├── mova_audio_dit.py │ │ │ │ │ ├── mova_video_dit.py │ │ │ │ │ ├── qwen_image.py │ │ │ │ │ ├── sana.py │ │ │ │ │ ├── wanvideo.py │ │ │ │ │ └── zimage.py │ │ │ │ ├── encoders/ │ │ │ │ │ ├── base.py │ │ │ │ │ ├── bert.py │ │ │ │ │ ├── clip.py │ │ │ │ │ ├── gemma2.py │ │ │ │ │ ├── gemma_3.py │ │ │ │ │ ├── hunyuan3d.py │ │ │ │ │ ├── llama.py │ │ │ │ │ ├── mistral_3.py │ │ │ │ │ ├── qwen2_5vl.py │ │ │ │ │ ├── qwen3.py │ │ │ │ │ ├── t5.py │ │ │ │ │ └── vision.py │ │ │ │ ├── parameter.py │ │ │ │ ├── registry.py │ │ │ │ ├── schedulers/ │ │ │ │ │ ├── base.py │ │ │ │ │ ├── flow_match_pair.py │ │ │ │ │ ├── hunyuan3d_scheduler.py │ │ │ │ │ ├── scheduling_comfyui_passthrough.py │ │ │ │ │ ├── scheduling_dpm_solver_multistep.py │ │ │ │ │ ├── scheduling_flow_match_euler_discrete.py │ │ │ │ │ ├── scheduling_flow_unipc_multistep.py │ │ │ │ │ ├── scheduling_helios.py │ │ │ │ │ ├── scheduling_self_forcing_flow_match.py │ │ │ │ │ └── scheduling_unipc_multistep.py │ │ │ │ ├── utils.py │ │ │ │ ├── vaes/ │ │ │ │ │ ├── autoencoder.py │ │ │ │ │ ├── autoencoder_dc.py │ │ │ │ │ ├── autoencoder_kl_flux2.py │ │ │ │ │ ├── autoencoder_kl_qwenimage.py │ │ │ │ │ ├── common.py │ │ │ │ │ ├── dac.py │ │ │ │ │ ├── hunyuan3d_vae.py │ │ │ │ │ ├── hunyuanvae.py │ │ │ │ │ ├── ltx_2_audio.py │ │ │ │ │ ├── ltx_2_vae.py │ │ │ │ │ ├── parallel/ │ │ │ │ │ │ ├── wan_common_utils.py │ │ │ │ │ │ └── wan_dist_utils.py │ │ │ │ │ └── wanvae.py │ │ │ │ ├── vision_utils.py │ │ │ │ └── vocoder/ │ │ │ │ └── ltx_2_vocoder.py │ │ │ ├── pipelines/ │ │ │ │ ├── __init__.py │ │ │ │ ├── comfyui_flux_pipeline.py │ │ │ │ ├── comfyui_qwen_image_pipeline.py │ │ │ │ ├── comfyui_zimage_pipeline.py │ │ │ │ ├── diffusers_pipeline.py │ │ │ │ ├── flux.py │ │ │ │ ├── flux_2.py │ │ │ │ ├── flux_2_klein.py │ │ │ │ ├── glm_image.py │ │ │ │ ├── helios_pipeline.py │ │ │ │ ├── hunyuan3d_pipeline.py │ │ │ │ ├── hunyuan_pipeline.py │ │ │ │ ├── ltx_2_pipeline.py │ │ │ │ ├── mova_pipeline.py │ │ │ │ ├── qwen_image.py │ │ │ │ ├── sana.py │ │ │ │ ├── wan_causal_dmd_pipeline.py │ │ │ │ ├── wan_dmd_pipeline.py │ │ │ │ ├── wan_i2v_dmd_pipeline.py │ │ │ │ ├── wan_i2v_pipeline.py │ │ │ │ ├── wan_pipeline.py │ │ │ │ └── zimage_pipeline.py │ │ │ ├── pipelines_core/ │ │ │ │ ├── __init__.py │ │ │ │ ├── composed_pipeline_base.py │ │ │ │ ├── executors/ │ │ │ │ │ ├── parallel_executor.py │ │ │ │ │ ├── pipeline_executor.py │ │ │ │ │ └── sync_executor.py │ │ │ │ ├── lora_format_adapter.py │ │ │ │ ├── lora_pipeline.py │ │ │ │ ├── schedule_batch.py │ │ │ │ └── stages/ │ │ │ │ ├── __init__.py │ │ │ │ ├── base.py │ │ │ │ ├── causal_denoising.py │ │ │ │ ├── comfyui_latent_preparation.py │ │ │ │ ├── decoding.py │ │ │ │ ├── decoding_av.py │ │ │ │ ├── denoising.py │ │ │ │ ├── denoising_av.py │ │ │ │ ├── denoising_dmd.py │ │ │ │ ├── encoding.py │ │ │ │ ├── hunyuan3d_paint.py │ │ │ │ ├── hunyuan3d_shape.py │ │ │ │ ├── image_encoding.py │ │ │ │ ├── input_validation.py │ │ │ │ ├── latent_preparation.py │ │ │ │ ├── latent_preparation_av.py │ │ │ │ ├── model_specific_stages/ │ │ │ │ │ ├── glm_image.py │ │ │ │ │ ├── helios_decoding.py │ │ │ │ │ ├── helios_denoising.py │ │ │ │ │ ├── mova.py │ │ │ │ │ └── qwen_image_layered.py │ │ │ │ ├── text_connector.py │ │ │ │ ├── text_encoding.py │ │ │ │ ├── timestep_preparation.py │ │ │ │ └── validators.py │ │ │ ├── platforms/ │ │ │ │ ├── __init__.py │ │ │ │ ├── cpu.py │ │ │ │ ├── cuda.py │ │ │ │ ├── interface.py │ │ │ │ ├── mps.py │ │ │ │ ├── musa.py │ │ │ │ ├── npu.py │ │ │ │ └── rocm.py │ │ │ ├── postprocess/ │ │ │ │ ├── __init__.py │ │ │ │ ├── realesrgan_upscaler.py │ │ │ │ └── rife_interpolator.py │ │ │ ├── scheduler_client.py │ │ │ ├── server_args.py │ │ │ └── utils/ │ │ │ ├── common.py │ │ │ ├── distributed.py │ │ │ ├── hf_diffusers_utils.py │ │ │ ├── layerwise_offload.py │ │ │ ├── logging_utils.py │ │ │ ├── mesh3d_utils.py │ │ │ ├── perf_logger.py │ │ │ ├── profiler.py │ │ │ └── quantization_utils.py │ │ ├── test/ │ │ │ ├── __init__.py │ │ │ ├── cli/ │ │ │ │ ├── test_generate_common.py │ │ │ │ ├── test_generate_i2i.py │ │ │ │ └── test_generate_t2i_perf.py │ │ │ ├── run_suite.py │ │ │ ├── scripts/ │ │ │ │ ├── gen_diffusion_ci_outputs.py │ │ │ │ └── gen_perf_baselines.py │ │ │ ├── server/ │ │ │ │ ├── ascend/ │ │ │ │ │ ├── perf_baselines_npu.json │ │ │ │ │ ├── test_server_1_npu.py │ │ │ │ │ ├── test_server_2_npu.py │ │ │ │ │ ├── test_server_8_npu.py │ │ │ │ │ └── testcase_configs_npu.py │ │ │ │ ├── conftest.py │ │ │ │ ├── perf_baselines.json │ │ │ │ ├── test_server_2_gpu_a.py │ │ │ │ ├── test_server_2_gpu_b.py │ │ │ │ ├── test_server_a.py │ │ │ │ ├── test_server_b.py │ │ │ │ ├── test_server_common.py │ │ │ │ ├── test_server_utils.py │ │ │ │ ├── test_update_weights_from_disk.py │ │ │ │ └── testcase_configs.py │ │ │ ├── slack_utils.py │ │ │ ├── test_files/ │ │ │ │ ├── launch_flux.json │ │ │ │ └── launch_wan.json │ │ │ ├── test_utils.py │ │ │ └── unit/ │ │ │ ├── test_lora_format_adapter.py │ │ │ ├── test_sampling_params.py │ │ │ ├── test_server_args.py │ │ │ └── test_storage.py │ │ ├── third_party/ │ │ │ ├── __init__.py │ │ │ └── pynvml.py │ │ ├── tools/ │ │ │ ├── convert_hf_to_fp8.py │ │ │ └── wan_repack.py │ │ └── utils.py │ ├── profiler.py │ ├── srt/ │ │ ├── batch_invariant_ops/ │ │ │ ├── __init__.py │ │ │ └── batch_invariant_ops.py │ │ ├── batch_overlap/ │ │ │ ├── operations.py │ │ │ ├── operations_strategy.py │ │ │ ├── single_batch_overlap.py │ │ │ └── two_batch_overlap.py │ │ ├── checkpoint_engine/ │ │ │ ├── __init__.py │ │ │ ├── checkpoint_engine_worker.py │ │ │ └── update.py │ │ ├── compilation/ │ │ │ ├── backend.py │ │ │ ├── compilation_config.py │ │ │ ├── compilation_counter.py │ │ │ ├── compile.py │ │ │ ├── compiler_interface.py │ │ │ ├── cuda_piecewise_backend.py │ │ │ ├── fix_functionalization.py │ │ │ ├── fx_utils.py │ │ │ ├── inductor_pass.py │ │ │ ├── npu_piecewise_backend.py │ │ │ ├── pass_manager.py │ │ │ ├── piecewise_context_manager.py │ │ │ └── weak_ref_tensor.py │ │ ├── configs/ │ │ │ ├── __init__.py │ │ │ ├── afmoe.py │ │ │ ├── bailing_hybrid.py │ │ │ ├── chatglm.py │ │ │ ├── dbrx.py │ │ │ ├── deepseek_ocr.py │ │ │ ├── deepseekvl2.py │ │ │ ├── device_config.py │ │ │ ├── dots_ocr.py │ │ │ ├── dots_vlm.py │ │ │ ├── exaone.py │ │ │ ├── falcon_h1.py │ │ │ ├── granitemoehybrid.py │ │ │ ├── internvl.py │ │ │ ├── janus_pro.py │ │ │ ├── jet_nemotron.py │ │ │ ├── jet_vlm.py │ │ │ ├── kimi_k25.py │ │ │ ├── kimi_linear.py │ │ │ ├── kimi_vl.py │ │ │ ├── kimi_vl_moonvit.py │ │ │ ├── lfm2.py │ │ │ ├── lfm2_moe.py │ │ │ ├── load_config.py │ │ │ ├── longcat_flash.py │ │ │ ├── mamba_utils.py │ │ │ ├── model_config.py │ │ │ ├── modelopt_config.py │ │ │ ├── nano_nemotron_vl.py │ │ │ ├── nemotron_h.py │ │ │ ├── olmo3.py │ │ │ ├── points_v15_chat.py │ │ │ ├── qwen3_5.py │ │ │ ├── qwen3_next.py │ │ │ ├── qwen3_omni.py │ │ │ ├── qwen3_vl.py │ │ │ ├── radio.py │ │ │ ├── step3_vl.py │ │ │ ├── step3p5.py │ │ │ ├── update_config.py │ │ │ └── utils.py │ │ ├── connector/ │ │ │ ├── __init__.py │ │ │ ├── base_connector.py │ │ │ ├── redis.py │ │ │ ├── remote_instance.py │ │ │ ├── s3.py │ │ │ ├── serde/ │ │ │ │ ├── __init__.py │ │ │ │ ├── safe_serde.py │ │ │ │ └── serde.py │ │ │ └── utils.py │ │ ├── constants.py │ │ ├── constrained/ │ │ │ ├── base_grammar_backend.py │ │ │ ├── grammar_manager.py │ │ │ ├── llguidance_backend.py │ │ │ ├── outlines_backend.py │ │ │ ├── outlines_jump_forward.py │ │ │ ├── reasoner_grammar_backend.py │ │ │ ├── triton_ops/ │ │ │ │ └── bitmask_ops.py │ │ │ ├── utils.py │ │ │ └── xgrammar_backend.py │ │ ├── debug_utils/ │ │ │ ├── __init__.py │ │ │ ├── comparator/ │ │ │ │ ├── __init__.py │ │ │ │ ├── __main__.py │ │ │ │ ├── aligner/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── axis_aligner.py │ │ │ │ │ ├── entrypoint/ │ │ │ │ │ │ ├── __init__.py │ │ │ │ │ │ ├── executor.py │ │ │ │ │ │ ├── planner.py │ │ │ │ │ │ ├── traced_types.py │ │ │ │ │ │ └── types.py │ │ │ │ │ ├── reorderer/ │ │ │ │ │ │ ├── __init__.py │ │ │ │ │ │ ├── executor.py │ │ │ │ │ │ ├── planner.py │ │ │ │ │ │ └── types.py │ │ │ │ │ ├── token_aligner/ │ │ │ │ │ │ ├── __init__.py │ │ │ │ │ │ ├── concat_steps/ │ │ │ │ │ │ │ ├── __init__.py │ │ │ │ │ │ │ ├── executor.py │ │ │ │ │ │ │ └── thd_seq_lens_loader.py │ │ │ │ │ │ ├── entrypoint.py │ │ │ │ │ │ └── smart/ │ │ │ │ │ │ ├── __init__.py │ │ │ │ │ │ ├── aux_loader.py │ │ │ │ │ │ ├── aux_plugins.py │ │ │ │ │ │ ├── executor.py │ │ │ │ │ │ ├── planner.py │ │ │ │ │ │ ├── seq_info_builder.py │ │ │ │ │ │ └── types.py │ │ │ │ │ └── unsharder/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── executor.py │ │ │ │ │ ├── parallel_info.py │ │ │ │ │ ├── planner.py │ │ │ │ │ └── types.py │ │ │ │ ├── bundle_comparator.py │ │ │ │ ├── bundle_matcher.py │ │ │ │ ├── dims_spec/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── comment_parser.py │ │ │ │ │ ├── dim_parser.py │ │ │ │ │ ├── dims_parser.py │ │ │ │ │ ├── modifier_parser.py │ │ │ │ │ ├── tensor_naming.py │ │ │ │ │ └── types.py │ │ │ │ ├── display.py │ │ │ │ ├── dp_utils.py │ │ │ │ ├── entrypoint.py │ │ │ │ ├── log_sink.py │ │ │ │ ├── meta_overrider.py │ │ │ │ ├── output_formatter.py │ │ │ │ ├── output_types.py │ │ │ │ ├── per_token_visualizer.py │ │ │ │ ├── preset.py │ │ │ │ ├── report_sink.py │ │ │ │ ├── tensor_comparator/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── comparator.py │ │ │ │ │ ├── formatter.py │ │ │ │ │ └── types.py │ │ │ │ ├── utils.py │ │ │ │ └── visualizer/ │ │ │ │ ├── __init__.py │ │ │ │ ├── figure.py │ │ │ │ ├── panels.py │ │ │ │ └── preprocessing.py │ │ │ ├── cuda_coredump.py │ │ │ ├── dump_comparator.py │ │ │ ├── dump_loader.py │ │ │ ├── dumper.py │ │ │ ├── log_parser.py │ │ │ ├── model_truncator.py │ │ │ ├── schedule_simulator/ │ │ │ │ ├── __init__.py │ │ │ │ ├── __main__.py │ │ │ │ ├── data_source/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── data_loader.py │ │ │ │ │ └── data_synthesis.py │ │ │ │ ├── entrypoint.py │ │ │ │ ├── gpu_state.py │ │ │ │ ├── metrics.py │ │ │ │ ├── request.py │ │ │ │ ├── routers/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── base.py │ │ │ │ │ ├── random_router.py │ │ │ │ │ ├── round_robin_router.py │ │ │ │ │ └── sticky_router.py │ │ │ │ ├── schedulers/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── base.py │ │ │ │ │ └── fifo_scheduler.py │ │ │ │ └── simulator.py │ │ │ ├── source_patcher/ │ │ │ │ ├── __init__.py │ │ │ │ ├── code_patcher.py │ │ │ │ ├── source_editor.py │ │ │ │ └── types.py │ │ │ ├── tensor_dump_forward_hook.py │ │ │ └── text_comparator.py │ │ ├── disaggregation/ │ │ │ ├── ascend/ │ │ │ │ ├── __init__.py │ │ │ │ ├── conn.py │ │ │ │ └── transfer_engine.py │ │ │ ├── base/ │ │ │ │ ├── __init__.py │ │ │ │ └── conn.py │ │ │ ├── common/ │ │ │ │ ├── __init__.py │ │ │ │ ├── conn.py │ │ │ │ └── utils.py │ │ │ ├── decode.py │ │ │ ├── decode_kvcache_offload_manager.py │ │ │ ├── decode_schedule_batch_mixin.py │ │ │ ├── encode_grpc_server.py │ │ │ ├── encode_receiver.py │ │ │ ├── encode_server.py │ │ │ ├── fake/ │ │ │ │ ├── __init__.py │ │ │ │ └── conn.py │ │ │ ├── kv_events.py │ │ │ ├── mooncake/ │ │ │ │ ├── __init__.py │ │ │ │ ├── conn.py │ │ │ │ └── utils.py │ │ │ ├── mori/ │ │ │ │ ├── __init__.py │ │ │ │ └── conn.py │ │ │ ├── nixl/ │ │ │ │ ├── __init__.py │ │ │ │ └── conn.py │ │ │ ├── prefill.py │ │ │ └── utils.py │ │ ├── distributed/ │ │ │ ├── __init__.py │ │ │ ├── communication_op.py │ │ │ ├── device_communicators/ │ │ │ │ ├── all_reduce_utils.py │ │ │ │ ├── cuda_wrapper.py │ │ │ │ ├── custom_all_reduce.py │ │ │ │ ├── custom_all_reduce_ops.py │ │ │ │ ├── custom_all_reduce_utils.py │ │ │ │ ├── hpu_communicator.py │ │ │ │ ├── mooncake_transfer_engine.py │ │ │ │ ├── npu_communicator.py │ │ │ │ ├── pymscclpp.py │ │ │ │ ├── pynccl.py │ │ │ │ ├── pynccl_allocator.py │ │ │ │ ├── pynccl_wrapper.py │ │ │ │ ├── quick_all_reduce.py │ │ │ │ ├── shm_broadcast.py │ │ │ │ ├── torch_symm_mem.py │ │ │ │ └── xpu_communicator.py │ │ │ ├── naive_distributed.py │ │ │ ├── parallel_state.py │ │ │ └── utils.py │ │ ├── dllm/ │ │ │ ├── algorithm/ │ │ │ │ ├── __init__.py │ │ │ │ ├── base.py │ │ │ │ ├── joint_threshold.py │ │ │ │ └── low_confidence.py │ │ │ ├── config.py │ │ │ └── mixin/ │ │ │ ├── req.py │ │ │ └── scheduler.py │ │ ├── elastic_ep/ │ │ │ ├── elastic_ep.py │ │ │ ├── expert_backup_client.py │ │ │ └── expert_backup_manager.py │ │ ├── entrypoints/ │ │ │ ├── EngineBase.py │ │ │ ├── anthropic/ │ │ │ │ ├── __init__.py │ │ │ │ ├── protocol.py │ │ │ │ └── serving.py │ │ │ ├── context.py │ │ │ ├── engine.py │ │ │ ├── grpc_server.py │ │ │ ├── harmony_utils.py │ │ │ ├── http_server.py │ │ │ ├── http_server_engine.py │ │ │ ├── ollama/ │ │ │ │ ├── README.md │ │ │ │ ├── __init__.py │ │ │ │ ├── protocol.py │ │ │ │ ├── serving.py │ │ │ │ └── smart_router.py │ │ │ ├── openai/ │ │ │ │ ├── __init__.py │ │ │ │ ├── encoding_dsv32.py │ │ │ │ ├── protocol.py │ │ │ │ ├── serving_base.py │ │ │ │ ├── serving_chat.py │ │ │ │ ├── serving_classify.py │ │ │ │ ├── serving_completions.py │ │ │ │ ├── serving_embedding.py │ │ │ │ ├── serving_rerank.py │ │ │ │ ├── serving_responses.py │ │ │ │ ├── serving_score.py │ │ │ │ ├── serving_tokenize.py │ │ │ │ ├── serving_transcription.py │ │ │ │ ├── tool_server.py │ │ │ │ ├── usage_processor.py │ │ │ │ └── utils.py │ │ │ ├── ssl_utils.py │ │ │ ├── tool.py │ │ │ ├── v1_loads.py │ │ │ └── warmup.py │ │ ├── environ.py │ │ ├── eplb/ │ │ │ ├── __init__.py │ │ │ ├── eplb_algorithms/ │ │ │ │ ├── __init__.py │ │ │ │ ├── deepseek.py │ │ │ │ ├── deepseek_vec.py │ │ │ │ └── elasticity_aware.py │ │ │ ├── eplb_manager.py │ │ │ ├── eplb_simulator/ │ │ │ │ ├── __init__.py │ │ │ │ └── reader.py │ │ │ ├── expert_distribution.py │ │ │ ├── expert_location.py │ │ │ ├── expert_location_dispatch.py │ │ │ └── expert_location_updater.py │ │ ├── function_call/ │ │ │ ├── base_format_detector.py │ │ │ ├── core_types.py │ │ │ ├── deepseekv31_detector.py │ │ │ ├── deepseekv32_detector.py │ │ │ ├── deepseekv3_detector.py │ │ │ ├── function_call_parser.py │ │ │ ├── gigachat3_detector.py │ │ │ ├── glm47_moe_detector.py │ │ │ ├── glm4_moe_detector.py │ │ │ ├── gpt_oss_detector.py │ │ │ ├── hermes_detector.py │ │ │ ├── internlm_detector.py │ │ │ ├── json_array_parser.py │ │ │ ├── kimik2_detector.py │ │ │ ├── lfm2_detector.py │ │ │ ├── llama32_detector.py │ │ │ ├── mimo_detector.py │ │ │ ├── minimax_m2.py │ │ │ ├── mistral_detector.py │ │ │ ├── pythonic_detector.py │ │ │ ├── qwen25_detector.py │ │ │ ├── qwen3_coder_detector.py │ │ │ ├── step3_detector.py │ │ │ ├── trinity_detector.py │ │ │ └── utils.py │ │ ├── grpc/ │ │ │ └── __init__.py │ │ ├── hardware_backend/ │ │ │ └── npu/ │ │ │ ├── allocator_npu.py │ │ │ ├── attention/ │ │ │ │ ├── ascend_backend.py │ │ │ │ ├── ascend_torch_native_backend.py │ │ │ │ └── mla_preprocess.py │ │ │ ├── cmo.py │ │ │ ├── graph_runner/ │ │ │ │ ├── eagle_draft_extend_npu_graph_runner.py │ │ │ │ ├── eagle_draft_npu_graph_runner.py │ │ │ │ ├── npu_graph_runner.py │ │ │ │ └── vit_npu_graph_runner.py │ │ │ ├── memory_pool_npu.py │ │ │ ├── modules/ │ │ │ │ ├── deepseek_v2_attention_mla_npu.py │ │ │ │ └── qwen_vl_processor.py │ │ │ ├── moe/ │ │ │ │ └── topk.py │ │ │ ├── quantization/ │ │ │ │ ├── fused_moe_method_npu.py │ │ │ │ └── linear_method_npu.py │ │ │ └── utils.py │ │ ├── layers/ │ │ │ ├── activation.py │ │ │ ├── amx_utils.py │ │ │ ├── attention/ │ │ │ │ ├── aiter_backend.py │ │ │ │ ├── attention_registry.py │ │ │ │ ├── base_attn_backend.py │ │ │ │ ├── cutlass_mla_backend.py │ │ │ │ ├── double_sparsity_backend.py │ │ │ │ ├── dual_chunk_flashattention_backend.py │ │ │ │ ├── fla/ │ │ │ │ │ ├── chunk.py │ │ │ │ │ ├── chunk_delta_h.py │ │ │ │ │ ├── chunk_o.py │ │ │ │ │ ├── chunk_scaled_dot_kkt.py │ │ │ │ │ ├── cumsum.py │ │ │ │ │ ├── fused_gdn_gating.py │ │ │ │ │ ├── fused_norm_gate.py │ │ │ │ │ ├── fused_recurrent.py │ │ │ │ │ ├── fused_sigmoid_gating_recurrent.py │ │ │ │ │ ├── index.py │ │ │ │ │ ├── kda.py │ │ │ │ │ ├── l2norm.py │ │ │ │ │ ├── layernorm_gated.py │ │ │ │ │ ├── op.py │ │ │ │ │ ├── solve_tril.py │ │ │ │ │ ├── utils.py │ │ │ │ │ └── wy_fast.py │ │ │ │ ├── flashattention_backend.py │ │ │ │ ├── flashinfer_backend.py │ │ │ │ ├── flashinfer_mla_backend.py │ │ │ │ ├── flashmla_backend.py │ │ │ │ ├── hybrid_attn_backend.py │ │ │ │ ├── hybrid_linear_attn_backend.py │ │ │ │ ├── intel_amx_backend.py │ │ │ │ ├── linear/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── gdn_backend.py │ │ │ │ │ ├── kda_backend.py │ │ │ │ │ ├── kernels/ │ │ │ │ │ │ ├── __init__.py │ │ │ │ │ │ ├── gdn_cutedsl.py │ │ │ │ │ │ ├── gdn_flashinfer.py │ │ │ │ │ │ ├── gdn_triton.py │ │ │ │ │ │ ├── kda_triton.py │ │ │ │ │ │ └── kernel_backend.py │ │ │ │ │ ├── lightning_attn.py │ │ │ │ │ ├── lightning_backend.py │ │ │ │ │ ├── linear_metadata.py │ │ │ │ │ ├── seg_la.py │ │ │ │ │ └── utils.py │ │ │ │ ├── mamba/ │ │ │ │ │ ├── causal_conv1d.py │ │ │ │ │ ├── causal_conv1d_triton.py │ │ │ │ │ ├── mamba.py │ │ │ │ │ ├── mamba2_metadata.py │ │ │ │ │ ├── mamba_state_scatter_triton.py │ │ │ │ │ ├── mixer2_rms_norm_gated.py │ │ │ │ │ └── ops/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── layernorm_gated.py │ │ │ │ │ ├── mamba_ssm.py │ │ │ │ │ ├── ssd_bmm.py │ │ │ │ │ ├── ssd_chunk_scan.py │ │ │ │ │ ├── ssd_chunk_state.py │ │ │ │ │ ├── ssd_combined.py │ │ │ │ │ ├── ssd_state_passing.py │ │ │ │ │ └── ssu_dispatch.py │ │ │ │ ├── merge_state.py │ │ │ │ ├── nsa/ │ │ │ │ │ ├── dequant_k_cache.py │ │ │ │ │ ├── index_buf_accessor.py │ │ │ │ │ ├── nsa_backend_mtp_precompute.py │ │ │ │ │ ├── nsa_indexer.py │ │ │ │ │ ├── nsa_mtp_verification.py │ │ │ │ │ ├── quant_k_cache.py │ │ │ │ │ ├── tilelang_kernel.py │ │ │ │ │ ├── transform_index.py │ │ │ │ │ ├── triton_kernel.py │ │ │ │ │ └── utils.py │ │ │ │ ├── nsa_backend.py │ │ │ │ ├── tbo_backend.py │ │ │ │ ├── torch_flex_backend.py │ │ │ │ ├── torch_native_backend.py │ │ │ │ ├── triton_backend.py │ │ │ │ ├── triton_ops/ │ │ │ │ │ ├── decode_attention.py │ │ │ │ │ ├── double_sparsity_attention.py │ │ │ │ │ ├── extend_attention.py │ │ │ │ │ ├── merge_state.py │ │ │ │ │ ├── prefill_attention.py │ │ │ │ │ ├── rocm_mla_decode_rope.py │ │ │ │ │ └── trtllm_fp8_kv_kernel.py │ │ │ │ ├── trtllm_mha_backend.py │ │ │ │ ├── trtllm_mla_backend.py │ │ │ │ ├── utils.py │ │ │ │ ├── vision.py │ │ │ │ ├── vision_utils.py │ │ │ │ ├── wave_backend.py │ │ │ │ ├── wave_ops/ │ │ │ │ │ ├── decode_attention.py │ │ │ │ │ ├── extend_attention.py │ │ │ │ │ └── prefill_attention.py │ │ │ │ └── xpu_backend.py │ │ │ ├── communicator.py │ │ │ ├── communicator_nsa_cp.py │ │ │ ├── conv.py │ │ │ ├── deep_gemm_wrapper/ │ │ │ │ ├── __init__.py │ │ │ │ ├── compile_utils.py │ │ │ │ ├── configurer.py │ │ │ │ └── entrypoint.py │ │ │ ├── dp_attention.py │ │ │ ├── elementwise.py │ │ │ ├── flashinfer_comm_fusion.py │ │ │ ├── int4fp8_utils.py │ │ │ ├── layernorm.py │ │ │ ├── linear.py │ │ │ ├── logits_processor.py │ │ │ ├── model_parallel.py │ │ │ ├── modelopt_utils.py │ │ │ ├── moe/ │ │ │ │ ├── __init__.py │ │ │ │ ├── cutlass_moe.py │ │ │ │ ├── cutlass_moe_params.py │ │ │ │ ├── cutlass_w4a8_moe.py │ │ │ │ ├── ep_moe/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── kernels.py │ │ │ │ │ └── layer.py │ │ │ │ ├── flashinfer_cutedsl_moe.py │ │ │ │ ├── flashinfer_trtllm_moe.py │ │ │ │ ├── fused_moe_native.py │ │ │ │ ├── fused_moe_triton/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── configs/ │ │ │ │ │ │ ├── README.md │ │ │ │ │ │ ├── triton_3_1_0/ │ │ │ │ │ │ │ ├── E=1,N=14336,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json │ │ │ │ │ │ │ ├── E=1,N=14336,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=1,N=1792,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json │ │ │ │ │ │ │ ├── E=1,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=1,N=3072,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json │ │ │ │ │ │ │ ├── E=1,N=3072,device_name=NVIDIA_H100_80GB_HBM3,dtype=int8_w8a16.json │ │ │ │ │ │ │ ├── E=1,N=3072,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=1,N=3584,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json │ │ │ │ │ │ │ ├── E=1,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=1,N=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json │ │ │ │ │ │ │ ├── E=1,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=144,N=512,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=16,N=1024,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=16,N=1024,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=16,N=1344,device_name=NVIDIA_A100-SXM4-40GB.json │ │ │ │ │ │ │ ├── E=16,N=1344,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=16,N=1344,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=16,N=14336,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json │ │ │ │ │ │ │ ├── E=16,N=14336,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=16,N=1792,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json │ │ │ │ │ │ │ ├── E=16,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=16,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=16,N=2688,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=16,N=2688,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=16,N=3072,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json │ │ │ │ │ │ │ ├── E=16,N=3072,device_name=NVIDIA_H100_80GB_HBM3,dtype=int8_w8a16.json │ │ │ │ │ │ │ ├── E=16,N=3200,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=16,N=3584,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json │ │ │ │ │ │ │ ├── E=16,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=16,N=6400,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=16,N=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a16.json │ │ │ │ │ │ │ ├── E=16,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=16,N=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=int8_w8a16.json │ │ │ │ │ │ │ ├── E=16,N=800,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=160,N=192,device_name=NVIDIA_A800-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=20,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=24,N=1024,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=256,N=128,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=128,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8.json │ │ │ │ │ │ │ ├── E=256,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json │ │ │ │ │ │ │ ├── E=256,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=128,device_name=NVIDIA_H20,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=128,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=256,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=256,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=256,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=256,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=64,device_name=NVIDIA_A800-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=256,N=64,device_name=NVIDIA_L20,dtype=int8_w8a8.json │ │ │ │ │ │ │ ├── E=256,N=64,device_name=NVIDIA_L40S,dtype=int8_w8a8.json │ │ │ │ │ │ │ ├── E=64,N=1024,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=64,N=1280,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=64,N=1280,device_name=NVIDIA_A800-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=64,N=1280,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=64,N=1280,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=64,N=1280,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=64,N=1280,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=64,N=2560,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=64,N=2560,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=64,N=2560,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=64,N=320,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=64,N=320,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=64,N=320,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=64,N=320,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=64,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=64,N=640,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=64,N=640,device_name=NVIDIA_A800-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=64,N=640,device_name=NVIDIA_GeForce_RTX_4090,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=64,N=640,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=64,N=640,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=64,N=640,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=64,N=640,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=8,N=14336,device_name=AMD_Instinct_MI300X.json │ │ │ │ │ │ │ ├── E=8,N=14336,device_name=AMD_Instinct_MI325X.json │ │ │ │ │ │ │ ├── E=8,N=14336,device_name=AMD_Radeon_Graphics.json │ │ │ │ │ │ │ ├── E=8,N=14336,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=14336,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=14336,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=8,N=1792,device_name=AMD_Instinct_MI300X.json │ │ │ │ │ │ │ ├── E=8,N=1792,device_name=AMD_Instinct_MI325X.json │ │ │ │ │ │ │ ├── E=8,N=1792,device_name=AMD_Radeon_Graphics.json │ │ │ │ │ │ │ ├── E=8,N=1792,device_name=NVIDIA_A100-SXM4-40GB.json │ │ │ │ │ │ │ ├── E=8,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=8,N=1792,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=8,N=1792,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=1792,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=8,N=2048,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=8,N=2048,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=8,N=2048,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=2048,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=8,N=3584,device_name=AMD_Instinct_MI300X.json │ │ │ │ │ │ │ ├── E=8,N=3584,device_name=AMD_Instinct_MI325X.json │ │ │ │ │ │ │ ├── E=8,N=3584,device_name=AMD_Radeon_Graphics.json │ │ │ │ │ │ │ ├── E=8,N=3584,device_name=NVIDIA_A100-SXM4-40GB.json │ │ │ │ │ │ │ ├── E=8,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=8,N=3584,device_name=NVIDIA_GeForce_RTX_4090,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=3584,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=3584,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=8,N=3584,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=3584,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=8,N=3584,device_name=NVIDIA_L40S.json │ │ │ │ │ │ │ ├── E=8,N=4096,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=4096,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=4096,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=4096,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=8,N=4096,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=4096,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=8,N=4096,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=4096,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=8,N=7168,device_name=AMD_Instinct_MI300X.json │ │ │ │ │ │ │ ├── E=8,N=7168,device_name=AMD_Instinct_MI325X.json │ │ │ │ │ │ │ ├── E=8,N=7168,device_name=AMD_Radeon_Graphics.json │ │ │ │ │ │ │ ├── E=8,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=8,N=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=7168,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=8,N=8192,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=8192,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=8192,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=8,N=8192,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8.json │ │ │ │ │ │ │ └── E=8,N=8192,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ ├── triton_3_2_0/ │ │ │ │ │ │ │ ├── E=128,N=192,device_name=NVIDIA_A800-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=128,N=192,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=128,N=192,device_name=NVIDIA_H20.json │ │ │ │ │ │ │ ├── E=128,N=192,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=128,N=384,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=128,N=384,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=128,N=384,device_name=NVIDIA_H20.json │ │ │ │ │ │ │ ├── E=128,N=384,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=128,N=384,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=128,N=512,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=128,N=768,device_name=NVIDIA_A800-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=128,N=768,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=128,N=768,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=128,N=768,device_name=NVIDIA_H20.json │ │ │ │ │ │ │ ├── E=128,N=768,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=128,N=768,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=128,N=96,device_name=NVIDIA_H20.json │ │ │ │ │ │ │ ├── E=129,N=352,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json │ │ │ │ │ │ │ ├── E=160,N=320,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json │ │ │ │ │ │ │ ├── E=161,N=192,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json │ │ │ │ │ │ │ ├── E=257,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=257,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=264,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json │ │ │ │ │ │ │ ├── E=264,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=264,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=264,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=272,N=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8.json │ │ │ │ │ │ │ ├── E=272,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=272,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=272,N=64,device_name=NVIDIA_A800-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=288,N=64,device_name=NVIDIA_A800-SXM4-80GB.json │ │ │ │ │ │ │ └── E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── triton_3_3_0/ │ │ │ │ │ │ │ └── E=16,N=1024,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── triton_3_3_1/ │ │ │ │ │ │ │ ├── E=128,N=352,device_name=NVIDIA_RTX_6000_Ada_Generation,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=128,N=384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=128,N=384,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=128,N=768,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=128,N=768,device_name=NVIDIA_H20.json │ │ │ │ │ │ │ ├── E=160,N=192,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=160,N=320,device_name=NVIDIA_H20-3e.json │ │ │ │ │ │ │ ├── E=160,N=384,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=160,N=640,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=257,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=257,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_H20-3e,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=384,N=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=384,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=384,N=256,device_name=NVIDIA_H20-3e,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=385,N=128,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=385,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ └── E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── triton_3_4_0/ │ │ │ │ │ │ │ ├── E=128,N=1856,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=128,N=1856,device_name=NVIDIA_L40S.json │ │ │ │ │ │ │ ├── E=128,N=192,device_name=NVIDIA_H200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=128,N=352,device_name=NVIDIA_RTX_5880_Ada_Generation,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=128,N=384,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=128,N=928,device_name=NVIDIA_L40S.json │ │ │ │ │ │ │ ├── E=129,N=352,device_name=NVIDIA_B200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=129,N=352,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Max-Q_Workstation_Edition,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=129,N=704,device_name=NVIDIA_B200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=160,N=384,device_name=NVIDIA_B200,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=161,N=192,device_name=NVIDIA_H200,dtype=fp8_w8a8,per_channel_quant=True.json │ │ │ │ │ │ │ ├── E=161,N=384,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Max-Q_Workstation_Edition,dtype=fp8_w8a8.json │ │ │ │ │ │ │ ├── E=256,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=256,device_name=NVIDIA_B200.json │ │ │ │ │ │ │ ├── E=256,N=256,device_name=NVIDIA_H800,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=256,N=512,device_name=NVIDIA_B200.json │ │ │ │ │ │ │ ├── E=256,N=512,device_name=NVIDIA_H20.json │ │ │ │ │ │ │ ├── E=257,N=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128]_down.json │ │ │ │ │ │ │ ├── E=257,N=64,device_name=NVIDIA_A100-SXM4-80GB.json │ │ │ │ │ │ │ ├── E=384,N=128,device_name=,dtype=int4_w4a16.json │ │ │ │ │ │ │ ├── E=384,N=128,device_name=,dtype=int4_w4a16_down.json │ │ │ │ │ │ │ ├── E=384,N=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=512,N=128,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ ├── E=512,N=128,device_name=NVIDIA_H20-3e.json │ │ │ │ │ │ │ ├── E=512,N=128,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=512,N=128,device_name=NVIDIA_H800,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ │ ├── E=512,N=256,device_name=NVIDIA_B200.json │ │ │ │ │ │ │ ├── E=512,N=256,device_name=NVIDIA_H20-3e.json │ │ │ │ │ │ │ ├── E=512,N=256,device_name=NVIDIA_H200.json │ │ │ │ │ │ │ ├── E=512,N=64,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ │ └── E=512,N=64,device_name=NVIDIA_H200.json │ │ │ │ │ │ └── triton_3_5_1/ │ │ │ │ │ │ ├── E=128,N=1344,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=128,N=1344,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=128,N=1856,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=128,N=1856,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=128,N=232,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=128,N=232,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=128,N=2688,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=128,N=2688,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=128,N=464,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=128,N=464,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=128,N=928,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=128,N=928,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=16,N=1856,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=16,N=1856,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=16,N=2048,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=161,N=192,device_name=NVIDIA_B200,dtype=fp8_w8a8,per_channel_quant=True.json │ │ │ │ │ │ ├── E=161,N=192,device_name=NVIDIA_H20,dtype=fp8_w8a8,per_channel_quant=True.json │ │ │ │ │ │ ├── E=161,N=192,device_name=NVIDIA_H20-3e,dtype=fp8_w8a8,per_channel_quant=True.json │ │ │ │ │ │ ├── E=161,N=192,device_name=NVIDIA_H200,dtype=fp8_w8a8,per_channel_quant=True.json │ │ │ │ │ │ ├── E=161,N=192,device_name=NVIDIA_H200.json │ │ │ │ │ │ ├── E=161,N=192,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Max-Q_Workstation_Edition,dtype=fp8_w8a8,per_channel_quant=True.json │ │ │ │ │ │ ├── E=161,N=192,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Max-Q_Workstation_Edition.json │ │ │ │ │ │ ├── E=161,N=384,device_name=NVIDIA_B200,dtype=fp8_w8a8,per_channel_quant=True.json │ │ │ │ │ │ ├── E=161,N=384,device_name=NVIDIA_H200,dtype=fp8_w8a8,per_channel_quant=True.json │ │ │ │ │ │ ├── E=20,N=1536,device_name=NVIDIA_H200,dtype=fp8_w8a8,per_channel_quant=True.json │ │ │ │ │ │ ├── E=20,N=1536,device_name=NVIDIA_H200.json │ │ │ │ │ │ ├── E=256,N=1344,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=256,N=1344,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=256,N=2688,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=256,N=2688,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=256,N=384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ ├── E=256,N=384,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ ├── E=256,N=672,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=256,N=672,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128]_down.json │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_H20-3e,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_H20-3e,dtype=fp8_w8a8,block_shape=[128, 128]_down.json │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ ├── E=257,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128]_down.json │ │ │ │ │ │ ├── E=32,N=1856,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=32,N=1856,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=32,N=928,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=32,N=928,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=40,N=1536,device_name=NVIDIA_H200,dtype=fp8_w8a8,per_channel_quant=True.json │ │ │ │ │ │ ├── E=512,N=128,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=512,N=128,device_name=NVIDIA_H200.json │ │ │ │ │ │ ├── E=512,N=1344,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=512,N=1344,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=512,N=256,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=512,N=256,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ ├── E=512,N=256,device_name=NVIDIA_H200.json │ │ │ │ │ │ ├── E=512,N=2688,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=512,N=2688,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=512,N=336,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=512,N=336,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=512,N=672,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=512,N=672,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=64,N=1856,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=64,N=1856,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=64,N=2688,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=64,N=2688,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=64,N=464,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=64,N=464,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=64,N=928,device_name=NVIDIA_B200.json │ │ │ │ │ │ ├── E=64,N=928,device_name=NVIDIA_H100_80GB_HBM3.json │ │ │ │ │ │ ├── E=80,N=640,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ │ └── E=80,N=640,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128]_down.json │ │ │ │ │ ├── fused_marlin_moe.py │ │ │ │ │ ├── fused_moe.py │ │ │ │ │ ├── fused_moe_triton_config.py │ │ │ │ │ ├── fused_moe_triton_kernels.py │ │ │ │ │ ├── layer.py │ │ │ │ │ ├── moe_align_block_size.py │ │ │ │ │ └── triton_kernels_moe.py │ │ │ │ ├── kt_ep_wrapper.py │ │ │ │ ├── moe_runner/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── base.py │ │ │ │ │ ├── deep_gemm.py │ │ │ │ │ ├── flashinfer_trtllm.py │ │ │ │ │ ├── marlin.py │ │ │ │ │ ├── runner.py │ │ │ │ │ ├── triton.py │ │ │ │ │ └── triton_kernels.py │ │ │ │ ├── rocm_moe_utils.py │ │ │ │ ├── routed_experts_capturer.py │ │ │ │ ├── router.py │ │ │ │ ├── token_dispatcher/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── base.py │ │ │ │ │ ├── deepep.py │ │ │ │ │ ├── flashinfer.py │ │ │ │ │ ├── flashinfer_utils.py │ │ │ │ │ ├── fuseep.py │ │ │ │ │ ├── mooncake.py │ │ │ │ │ ├── moriep.py │ │ │ │ │ ├── nixl.py │ │ │ │ │ └── standard.py │ │ │ │ ├── topk.py │ │ │ │ └── utils.py │ │ │ ├── multimodal.py │ │ │ ├── n_gram_embedding.py │ │ │ ├── parameter.py │ │ │ ├── pooler.py │ │ │ ├── quantization/ │ │ │ │ ├── __init__.py │ │ │ │ ├── auto_round.py │ │ │ │ ├── awq.py │ │ │ │ ├── awq_triton.py │ │ │ │ ├── base_config.py │ │ │ │ ├── base_scheme.py │ │ │ │ ├── bitsandbytes.py │ │ │ │ ├── blockwise_int8.py │ │ │ │ ├── compressed_tensors/ │ │ │ │ │ ├── README.md │ │ │ │ │ ├── compressed_tensors.py │ │ │ │ │ ├── schemes/ │ │ │ │ │ │ ├── __init__.py │ │ │ │ │ │ ├── compressed_tensors_scheme.py │ │ │ │ │ │ ├── compressed_tensors_w4a4_mxint4_moe.py │ │ │ │ │ │ ├── compressed_tensors_w4a4_nvfp4.py │ │ │ │ │ │ ├── compressed_tensors_w4a4_nvfp4_moe.py │ │ │ │ │ │ ├── compressed_tensors_w4a8_int8_moe.py │ │ │ │ │ │ ├── compressed_tensors_w8a16_fp8.py │ │ │ │ │ │ ├── compressed_tensors_w8a8_fp8.py │ │ │ │ │ │ ├── compressed_tensors_w8a8_fp8_moe.py │ │ │ │ │ │ ├── compressed_tensors_w8a8_int8.py │ │ │ │ │ │ ├── compressed_tensors_w8a8_int8_moe.py │ │ │ │ │ │ ├── compressed_tensors_wNa16.py │ │ │ │ │ │ └── compressed_tensors_wNa16_moe.py │ │ │ │ │ └── utils.py │ │ │ │ ├── configs/ │ │ │ │ │ ├── N=1280,K=5120,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=1536,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=1536,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=1536,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=1536,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=1536,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=1536,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=2048,K=4096,device_name=NVIDIA_L40,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=2048,K=512,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=2048,K=512,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=2048,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=2048,K=512,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=2048,K=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=2048,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=2304,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=2304,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=2304,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=2304,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=2304,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=2304,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=24576,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=24576,K=1536,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=24576,K=1536,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=24576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=24576,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=24576,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=24576,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=24576,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=24576,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=24576,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=24576,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=256,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=256,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=256,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=256,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=256,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=3072,K=1536,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=3072,K=1536,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=3072,K=1536,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=3072,K=1536,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=3072,K=1536,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=3072,K=1536,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=3072,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=3072,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=3072,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=3072,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=32768,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=32768,K=512,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=32768,K=512,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=32768,K=512,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=32768,K=512,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=32768,K=512,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=32768,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=32768,K=512,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=32768,K=512,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=32768,K=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=32768,K=512,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=36864,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=36864,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4096,K=512,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4096,K=512,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4096,K=512,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4096,K=512,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4096,K=512,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4096,K=512,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4096,K=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4608,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4608,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4608,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4608,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4608,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4608,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=4608,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=512,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=512,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=512,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=512,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=512,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=512,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=5120,K=1024,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=5120,K=2048,device_name=NVIDIA_L40,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=5120,K=3200,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=576,K=7168,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=576,K=7168,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=576,K=7168,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=576,K=7168,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=576,K=7168,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=576,K=7168,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=576,K=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=576,K=7168,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=576,K=7168,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=576,K=7168,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=576,K=7168,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=6400,K=5120,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=1024,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=1024,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=1024,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=1024,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=1024,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=1024,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=1152,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=1152,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=1152,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=1152,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=1152,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=1152,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=128,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=128,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=128,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=128,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=128,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=16384,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=16384,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=16384,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=16384,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=16384,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=16384,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=16384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=16384,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=16384,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=16384,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=16384,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=18432,device_name=NVIDIA_A100-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=18432,device_name=NVIDIA_A800-SXM4-80GB,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=18432,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=18432,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=18432,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=18432,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=18432,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=18432,device_name=NVIDIA_L20Y,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2048,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2048,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2048,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2048,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2048,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2048,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2048,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2304,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2304,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2304,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2304,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2304,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2304,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=2304,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=256,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=256,device_name=AMD_Instinct_MI325X,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=256,device_name=AMD_Radeon_Graphics,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=256,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=256,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json │ │ │ │ │ ├── N=7168,K=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json │ │ │ │ │ └── README.md │ │ │ │ ├── fp4_utils.py │ │ │ │ ├── fp8.py │ │ │ │ ├── fp8_kernel.py │ │ │ │ ├── fp8_utils.py │ │ │ │ ├── fpgemm_fp8.py │ │ │ │ ├── gguf.py │ │ │ │ ├── gptq.py │ │ │ │ ├── int8_kernel.py │ │ │ │ ├── int8_utils.py │ │ │ │ ├── kv_cache.py │ │ │ │ ├── kvfp4_tensor.py │ │ │ │ ├── marlin_utils.py │ │ │ │ ├── marlin_utils_fp8.py │ │ │ │ ├── modelopt_quant.py │ │ │ │ ├── modelslim/ │ │ │ │ │ ├── README.md │ │ │ │ │ ├── modelslim.py │ │ │ │ │ └── schemes/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── modelslim_scheme.py │ │ │ │ │ ├── modelslim_w4a4_int4.py │ │ │ │ │ ├── modelslim_w4a4_int4_moe.py │ │ │ │ │ ├── modelslim_w4a8_int8_moe.py │ │ │ │ │ ├── modelslim_w8a8_int8.py │ │ │ │ │ └── modelslim_w8a8_int8_moe.py │ │ │ │ ├── moe_wna16.py │ │ │ │ ├── mxfp4.py │ │ │ │ ├── mxfp4_tensor.py │ │ │ │ ├── petit.py │ │ │ │ ├── petit_utils.py │ │ │ │ ├── qoq.py │ │ │ │ ├── quark/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── quark.py │ │ │ │ │ ├── schemes/ │ │ │ │ │ │ ├── __init__.py │ │ │ │ │ │ ├── quark_scheme.py │ │ │ │ │ │ ├── quark_w4a4_mxfp4.py │ │ │ │ │ │ ├── quark_w4a4_mxfp4_moe.py │ │ │ │ │ │ ├── quark_w8a8_fp8.py │ │ │ │ │ │ └── quark_w8a8_fp8_moe.py │ │ │ │ │ └── utils.py │ │ │ │ ├── quark_int4fp8_moe.py │ │ │ │ ├── rocm_mxfp4_utils.py │ │ │ │ ├── unquant.py │ │ │ │ ├── utils.py │ │ │ │ ├── w4afp8.py │ │ │ │ ├── w8a8_fp8.py │ │ │ │ └── w8a8_int8.py │ │ │ ├── radix_attention.py │ │ │ ├── radix_linear_attention.py │ │ │ ├── rocm_linear_utils.py │ │ │ ├── rotary_embedding/ │ │ │ │ ├── __init__.py │ │ │ │ ├── base.py │ │ │ │ ├── factory.py │ │ │ │ ├── mrope.py │ │ │ │ ├── mrope_rope_index.py │ │ │ │ ├── rope_variant.py │ │ │ │ ├── triton_kernels.py │ │ │ │ ├── utils.py │ │ │ │ └── yarn.py │ │ │ ├── sampler.py │ │ │ ├── sparse_pooler.py │ │ │ ├── torchao_utils.py │ │ │ ├── utils/ │ │ │ │ ├── __init__.py │ │ │ │ ├── common.py │ │ │ │ ├── hash.py │ │ │ │ ├── logprob.py │ │ │ │ └── multi_platform.py │ │ │ └── vocab_parallel_embedding.py │ │ ├── lora/ │ │ │ ├── backend/ │ │ │ │ ├── ascend_backend.py │ │ │ │ ├── base_backend.py │ │ │ │ ├── chunked_backend.py │ │ │ │ ├── lmhead_mixing.py │ │ │ │ ├── lora_registry.py │ │ │ │ ├── torch_backend.py │ │ │ │ └── triton_backend.py │ │ │ ├── eviction_policy.py │ │ │ ├── layers.py │ │ │ ├── lora.py │ │ │ ├── lora_config.py │ │ │ ├── lora_manager.py │ │ │ ├── lora_overlap_loader.py │ │ │ ├── lora_registry.py │ │ │ ├── mem_pool.py │ │ │ ├── torch_ops/ │ │ │ │ ├── __init__.py │ │ │ │ └── lora_ops.py │ │ │ ├── triton_ops/ │ │ │ │ ├── __init__.py │ │ │ │ ├── chunked_embedding_lora_a.py │ │ │ │ ├── chunked_sgmv_expand.py │ │ │ │ ├── chunked_sgmv_shrink.py │ │ │ │ ├── embedding_lora_a.py │ │ │ │ ├── fused_moe_lora_kernel.py │ │ │ │ ├── gate_up_lora_b.py │ │ │ │ ├── qkv_lora_b.py │ │ │ │ ├── sgemm_lora_a.py │ │ │ │ └── sgemm_lora_b.py │ │ │ └── utils.py │ │ ├── managers/ │ │ │ ├── async_dynamic_batch_tokenizer.py │ │ │ ├── async_mm_data_processor.py │ │ │ ├── cache_controller.py │ │ │ ├── configure_logging.py │ │ │ ├── data_parallel_controller.py │ │ │ ├── detokenizer_manager.py │ │ │ ├── disagg_service.py │ │ │ ├── io_struct.py │ │ │ ├── mm_utils.py │ │ │ ├── multi_tokenizer_mixin.py │ │ │ ├── multimodal_processor.py │ │ │ ├── overlap_utils.py │ │ │ ├── prefill_delayer.py │ │ │ ├── schedule_batch.py │ │ │ ├── schedule_policy.py │ │ │ ├── scheduler.py │ │ │ ├── scheduler_dp_attn_mixin.py │ │ │ ├── scheduler_input_blocker.py │ │ │ ├── scheduler_output_processor_mixin.py │ │ │ ├── scheduler_pp_mixin.py │ │ │ ├── scheduler_profiler_mixin.py │ │ │ ├── scheduler_recv_skipper.py │ │ │ ├── scheduler_runtime_checker_mixin.py │ │ │ ├── scheduler_update_weights_mixin.py │ │ │ ├── session_controller.py │ │ │ ├── template_manager.py │ │ │ ├── tokenizer_communicator_mixin.py │ │ │ ├── tokenizer_manager.py │ │ │ ├── tokenizer_manager_multiitem_mixin.py │ │ │ ├── tp_worker.py │ │ │ └── utils.py │ │ ├── mem_cache/ │ │ │ ├── allocator.py │ │ │ ├── base_prefix_cache.py │ │ │ ├── cache_init_params.py │ │ │ ├── chunk_cache.py │ │ │ ├── common.py │ │ │ ├── cpp_radix_tree/ │ │ │ │ ├── common.h │ │ │ │ ├── radix_tree.py │ │ │ │ ├── tree_v2.cpp │ │ │ │ ├── tree_v2.h │ │ │ │ ├── tree_v2_binding.cpp │ │ │ │ ├── tree_v2_debug.cpp │ │ │ │ ├── tree_v2_impl.h │ │ │ │ └── tree_v2_node.h │ │ │ ├── evict_policy.py │ │ │ ├── flush_cache.py │ │ │ ├── hi_mamba_radix_cache.py │ │ │ ├── hicache_storage.py │ │ │ ├── hiradix_cache.py │ │ │ ├── mamba_radix_cache.py │ │ │ ├── memory_pool.py │ │ │ ├── memory_pool_host.py │ │ │ ├── multimodal_cache.py │ │ │ ├── radix_cache.py │ │ │ ├── radix_cache_cpp.py │ │ │ ├── session_aware_cache.py │ │ │ ├── sparsity/ │ │ │ │ ├── __init__.py │ │ │ │ ├── algorithms/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── base_algorithm.py │ │ │ │ │ ├── deepseek_nsa.py │ │ │ │ │ └── quest_algorithm.py │ │ │ │ ├── backend/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── backend_adaptor.py │ │ │ │ ├── core/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── sparse_coordinator.py │ │ │ │ └── factory.py │ │ │ ├── storage/ │ │ │ │ ├── __init__.py │ │ │ │ ├── aibrix_kvcache/ │ │ │ │ │ ├── README.md │ │ │ │ │ ├── aibrix_kvcache_storage.py │ │ │ │ │ └── unit_test.py │ │ │ │ ├── backend_factory.py │ │ │ │ ├── eic/ │ │ │ │ │ ├── README.md │ │ │ │ │ ├── eic_storage.py │ │ │ │ │ └── test_unit.py │ │ │ │ ├── hf3fs/ │ │ │ │ │ ├── docs/ │ │ │ │ │ │ ├── README.md │ │ │ │ │ │ ├── deploy_sglang_3fs_multinode.md │ │ │ │ │ │ └── setup_usrbio_client.md │ │ │ │ │ ├── hf3fs_client.py │ │ │ │ │ ├── hf3fs_usrbio_client.py │ │ │ │ │ ├── hf3fs_utils.cpp │ │ │ │ │ ├── mini_3fs_metadata_server.py │ │ │ │ │ ├── storage_hf3fs.py │ │ │ │ │ └── test_hf3fs_utils.py │ │ │ │ ├── lmcache/ │ │ │ │ │ ├── README.md │ │ │ │ │ ├── example_config.yaml │ │ │ │ │ ├── lmc_radix_cache.py │ │ │ │ │ └── unit_test.py │ │ │ │ ├── mooncake_store/ │ │ │ │ │ ├── README.md │ │ │ │ │ ├── embedding_cache_controller.py │ │ │ │ │ ├── mooncake_embedding_store.py │ │ │ │ │ ├── mooncake_store.py │ │ │ │ │ └── test_mooncake_store.py │ │ │ │ └── nixl/ │ │ │ │ ├── README.md │ │ │ │ ├── hicache_nixl.py │ │ │ │ ├── nixl.config.toml.sample │ │ │ │ ├── nixl_utils.py │ │ │ │ └── test_hicache_nixl_storage.py │ │ │ ├── swa_memory_pool.py │ │ │ ├── swa_radix_cache.py │ │ │ └── utils.py │ │ ├── model_executor/ │ │ │ ├── cpu_graph_runner.py │ │ │ ├── cuda_graph_runner.py │ │ │ ├── forward_batch_deepseek_mha_mixin.py │ │ │ ├── forward_batch_info.py │ │ │ ├── hook_manager.py │ │ │ ├── input_buffers.py │ │ │ ├── mindspore_runner.py │ │ │ ├── model_runner.py │ │ │ ├── model_runner_kv_cache_mixin.py │ │ │ └── piecewise_cuda_graph_runner.py │ │ ├── model_loader/ │ │ │ ├── __init__.py │ │ │ ├── ci_weight_validation.py │ │ │ ├── loader.py │ │ │ ├── remote_instance_weight_loader_utils.py │ │ │ ├── utils.py │ │ │ └── weight_utils.py │ │ ├── models/ │ │ │ ├── afmoe.py │ │ │ ├── apertus.py │ │ │ ├── arcee.py │ │ │ ├── baichuan.py │ │ │ ├── bailing_moe.py │ │ │ ├── bailing_moe_linear.py │ │ │ ├── bailing_moe_nextn.py │ │ │ ├── bert.py │ │ │ ├── chatglm.py │ │ │ ├── clip.py │ │ │ ├── commandr.py │ │ │ ├── dbrx.py │ │ │ ├── deepseek.py │ │ │ ├── deepseek_common/ │ │ │ │ ├── __init__.py │ │ │ │ ├── attention_backend_handler.py │ │ │ │ ├── attention_forward_methods/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── forward_methods.py │ │ │ │ │ ├── forward_mha.py │ │ │ │ │ ├── forward_mla.py │ │ │ │ │ ├── forward_mla_fused_rope_cpu.py │ │ │ │ │ └── forward_mla_fused_rope_rocm.py │ │ │ │ ├── deepseek_weight_loader.py │ │ │ │ └── utils.py │ │ │ ├── deepseek_janus_pro.py │ │ │ ├── deepseek_nextn.py │ │ │ ├── deepseek_ocr.py │ │ │ ├── deepseek_v2.py │ │ │ ├── deepseek_vl2.py │ │ │ ├── dots_ocr.py │ │ │ ├── dots_vlm.py │ │ │ ├── dots_vlm_vit.py │ │ │ ├── ernie4.py │ │ │ ├── ernie45_moe_vl.py │ │ │ ├── ernie45_vl.py │ │ │ ├── ernie4_eagle.py │ │ │ ├── exaone.py │ │ │ ├── exaone4.py │ │ │ ├── exaone_moe.py │ │ │ ├── exaone_moe_mtp.py │ │ │ ├── falcon_h1.py │ │ │ ├── gemma.py │ │ │ ├── gemma2.py │ │ │ ├── gemma2_reward.py │ │ │ ├── gemma3_causal.py │ │ │ ├── gemma3_mm.py │ │ │ ├── gemma3n_audio.py │ │ │ ├── gemma3n_causal.py │ │ │ ├── gemma3n_mm.py │ │ │ ├── glm4.py │ │ │ ├── glm4_moe.py │ │ │ ├── glm4_moe_lite.py │ │ │ ├── glm4_moe_nextn.py │ │ │ ├── glm4v.py │ │ │ ├── glm4v_moe.py │ │ │ ├── glm_ocr.py │ │ │ ├── glm_ocr_nextn.py │ │ │ ├── glmasr.py │ │ │ ├── gpt2.py │ │ │ ├── gpt_bigcode.py │ │ │ ├── gpt_j.py │ │ │ ├── gpt_oss.py │ │ │ ├── granite.py │ │ │ ├── granitemoe.py │ │ │ ├── granitemoehybrid.py │ │ │ ├── grok.py │ │ │ ├── hunyuan.py │ │ │ ├── idefics2.py │ │ │ ├── internlm2.py │ │ │ ├── internlm2_reward.py │ │ │ ├── interns1.py │ │ │ ├── interns1pro.py │ │ │ ├── internvl.py │ │ │ ├── iquest_loopcoder.py │ │ │ ├── jet_nemotron.py │ │ │ ├── jet_vlm.py │ │ │ ├── kimi_k25.py │ │ │ ├── kimi_linear.py │ │ │ ├── kimi_vl.py │ │ │ ├── kimi_vl_moonvit.py │ │ │ ├── lfm2.py │ │ │ ├── lfm2_moe.py │ │ │ ├── lightonocr.py │ │ │ ├── llada2.py │ │ │ ├── llama.py │ │ │ ├── llama4.py │ │ │ ├── llama_classification.py │ │ │ ├── llama_eagle.py │ │ │ ├── llama_eagle3.py │ │ │ ├── llama_embedding.py │ │ │ ├── llama_reward.py │ │ │ ├── llava.py │ │ │ ├── llavavid.py │ │ │ ├── longcat_flash.py │ │ │ ├── longcat_flash_nextn.py │ │ │ ├── midashenglm.py │ │ │ ├── mimo.py │ │ │ ├── mimo_mtp.py │ │ │ ├── mimo_v2_flash.py │ │ │ ├── mimo_v2_flash_nextn.py │ │ │ ├── mindspore.py │ │ │ ├── minicpm.py │ │ │ ├── minicpm3.py │ │ │ ├── minicpmo.py │ │ │ ├── minicpmv.py │ │ │ ├── minimax_m2.py │ │ │ ├── ministral3.py │ │ │ ├── mistral.py │ │ │ ├── mistral_large_3.py │ │ │ ├── mistral_large_3_eagle.py │ │ │ ├── mixtral.py │ │ │ ├── mixtral_quant.py │ │ │ ├── mllama.py │ │ │ ├── mllama4.py │ │ │ ├── nano_nemotron_vl.py │ │ │ ├── nemotron_h.py │ │ │ ├── nemotron_h_mtp.py │ │ │ ├── nemotron_nas.py │ │ │ ├── nvila.py │ │ │ ├── nvila_lite.py │ │ │ ├── olmo.py │ │ │ ├── olmo2.py │ │ │ ├── olmoe.py │ │ │ ├── opt.py │ │ │ ├── orion.py │ │ │ ├── paddleocr_vl.py │ │ │ ├── persimmon.py │ │ │ ├── phi.py │ │ │ ├── phi3_small.py │ │ │ ├── phi4mm.py │ │ │ ├── phi4mm_audio.py │ │ │ ├── phi4mm_utils.py │ │ │ ├── phimoe.py │ │ │ ├── pixtral.py │ │ │ ├── points_v15_chat.py │ │ │ ├── qwen.py │ │ │ ├── qwen2.py │ │ │ ├── qwen2_5_vl.py │ │ │ ├── qwen2_audio.py │ │ │ ├── qwen2_classification.py │ │ │ ├── qwen2_eagle.py │ │ │ ├── qwen2_moe.py │ │ │ ├── qwen2_rm.py │ │ │ ├── qwen2_vl.py │ │ │ ├── qwen3.py │ │ │ ├── qwen3_5.py │ │ │ ├── qwen3_5_mtp.py │ │ │ ├── qwen3_classification.py │ │ │ ├── qwen3_moe.py │ │ │ ├── qwen3_next.py │ │ │ ├── qwen3_next_mtp.py │ │ │ ├── qwen3_omni_moe.py │ │ │ ├── qwen3_rm.py │ │ │ ├── qwen3_vl.py │ │ │ ├── qwen3_vl_moe.py │ │ │ ├── radio.py │ │ │ ├── registry.py │ │ │ ├── roberta.py │ │ │ ├── sarashina2_vision.py │ │ │ ├── sarvam_moe.py │ │ │ ├── sdar.py │ │ │ ├── sdar_moe.py │ │ │ ├── siglip.py │ │ │ ├── solar.py │ │ │ ├── stablelm.py │ │ │ ├── starcoder2.py │ │ │ ├── step3_vl.py │ │ │ ├── step3_vl_10b.py │ │ │ ├── step3p5.py │ │ │ ├── step3p5_mtp.py │ │ │ ├── teleflm.py │ │ │ ├── torch_native_llama.py │ │ │ ├── transformers.py │ │ │ ├── utils.py │ │ │ ├── whisper.py │ │ │ ├── xverse.py │ │ │ ├── xverse_moe.py │ │ │ └── yivl.py │ │ ├── multimodal/ │ │ │ ├── customized_mm_processor_utils.py │ │ │ ├── evs/ │ │ │ │ ├── README.md │ │ │ │ ├── __init__.py │ │ │ │ ├── evs_core.py │ │ │ │ ├── evs_module.py │ │ │ │ └── evs_processor.py │ │ │ ├── internvl_utils.py │ │ │ ├── internvl_vit_cuda_graph_runner.py │ │ │ ├── mm_utils.py │ │ │ ├── processors/ │ │ │ │ ├── base_processor.py │ │ │ │ ├── clip.py │ │ │ │ ├── deepseek_ocr.py │ │ │ │ ├── deepseek_vl_v2.py │ │ │ │ ├── dots_vlm.py │ │ │ │ ├── ernie45_vl.py │ │ │ │ ├── gemma3.py │ │ │ │ ├── gemma3n.py │ │ │ │ ├── glm4v.py │ │ │ │ ├── glmasr.py │ │ │ │ ├── interns1pro.py │ │ │ │ ├── internvl.py │ │ │ │ ├── janus_pro.py │ │ │ │ ├── kimi_k25.py │ │ │ │ ├── kimi_vl.py │ │ │ │ ├── lightonocr.py │ │ │ │ ├── llava.py │ │ │ │ ├── midashenglm.py │ │ │ │ ├── minicpm.py │ │ │ │ ├── mlama.py │ │ │ │ ├── mllama4.py │ │ │ │ ├── nano_nemotron_vl.py │ │ │ │ ├── nvila.py │ │ │ │ ├── paddleocr_vlm.py │ │ │ │ ├── phi4mm.py │ │ │ │ ├── pixtral.py │ │ │ │ ├── points_v15_chat.py │ │ │ │ ├── qwen_audio.py │ │ │ │ ├── qwen_vl.py │ │ │ │ ├── sarashina2_vision.py │ │ │ │ ├── step3_vl.py │ │ │ │ └── whisper.py │ │ │ └── vit_cuda_graph_runner.py │ │ ├── multiplex/ │ │ │ ├── multiplexing_mixin.py │ │ │ └── pdmux_context.py │ │ ├── observability/ │ │ │ ├── cpu_monitor.py │ │ │ ├── func_timer.py │ │ │ ├── label_transform.py │ │ │ ├── metrics_collector.py │ │ │ ├── req_time_stats.py │ │ │ ├── request_metrics_exporter.py │ │ │ ├── scheduler_metrics_mixin.py │ │ │ ├── startup_func_log_and_timer.py │ │ │ ├── trace.py │ │ │ └── utils.py │ │ ├── parser/ │ │ │ ├── code_completion_parser.py │ │ │ ├── conversation.py │ │ │ ├── harmony_parser.py │ │ │ ├── jinja_template_utils.py │ │ │ └── reasoning_parser.py │ │ ├── ray/ │ │ │ ├── __init__.py │ │ │ ├── engine.py │ │ │ ├── http_server.py │ │ │ └── scheduler_actor.py │ │ ├── sampling/ │ │ │ ├── custom_logit_processor.py │ │ │ ├── penaltylib/ │ │ │ │ ├── __init__.py │ │ │ │ ├── frequency_penalty.py │ │ │ │ ├── min_new_tokens.py │ │ │ │ ├── orchestrator.py │ │ │ │ └── presence_penalty.py │ │ │ ├── sampling_batch_info.py │ │ │ └── sampling_params.py │ │ ├── server_args.py │ │ ├── server_args_config_parser.py │ │ ├── speculative/ │ │ │ ├── base_spec_worker.py │ │ │ ├── cpp_ngram/ │ │ │ │ ├── .clang-format │ │ │ │ ├── ngram.cpp │ │ │ │ ├── ngram.h │ │ │ │ ├── ngram_cache.py │ │ │ │ ├── ngram_cache_binding.cpp │ │ │ │ ├── param.h │ │ │ │ └── queue.h │ │ │ ├── draft_utils.py │ │ │ ├── eagle_draft_cuda_graph_runner.py │ │ │ ├── eagle_draft_extend_cuda_graph_runner.py │ │ │ ├── eagle_info.py │ │ │ ├── eagle_info_v2.py │ │ │ ├── eagle_utils.py │ │ │ ├── eagle_worker.py │ │ │ ├── eagle_worker_v2.py │ │ │ ├── multi_layer_eagle_draft_extend_cuda_graph_runner.py │ │ │ ├── multi_layer_eagle_utils.py │ │ │ ├── multi_layer_eagle_worker.py │ │ │ ├── multi_layer_eagle_worker_v2.py │ │ │ ├── ngram_info.py │ │ │ ├── ngram_worker.py │ │ │ ├── spec_info.py │ │ │ ├── spec_utils.py │ │ │ ├── standalone_worker.py │ │ │ └── standalone_worker_v2.py │ │ ├── tokenizer/ │ │ │ └── tiktoken_tokenizer.py │ │ ├── utils/ │ │ │ ├── __init__.py │ │ │ ├── aio_rwlock.py │ │ │ ├── auth.py │ │ │ ├── bench_utils.py │ │ │ ├── common.py │ │ │ ├── cuda_ipc_transport_utils.py │ │ │ ├── custom_op.py │ │ │ ├── device_timer.py │ │ │ ├── gauge_histogram.py │ │ │ ├── hf_transformers_utils.py │ │ │ ├── host_shared_memory.py │ │ │ ├── json_response.py │ │ │ ├── log_utils.py │ │ │ ├── mistral_utils.py │ │ │ ├── model_file_verifier.py │ │ │ ├── multi_stream_utils.py │ │ │ ├── network.py │ │ │ ├── numa_utils.py │ │ │ ├── nvtx_pytorch_hooks.py │ │ │ ├── offloader.py │ │ │ ├── patch_tokenizer.py │ │ │ ├── patch_torch.py │ │ │ ├── poll_based_barrier.py │ │ │ ├── profile_merger.py │ │ │ ├── profile_utils.py │ │ │ ├── request_logger.py │ │ │ ├── rpd_utils.py │ │ │ ├── scheduler_status_logger.py │ │ │ ├── slow_rank_detector.py │ │ │ ├── torch_memory_saver_adapter.py │ │ │ ├── video_decoder.py │ │ │ ├── watchdog.py │ │ │ └── weight_checker.py │ │ └── weight_sync/ │ │ ├── tensor_bucket.py │ │ └── utils.py │ ├── test/ │ │ ├── __init__.py │ │ ├── accuracy_test_runner.py │ │ ├── ascend/ │ │ │ ├── __init__.py │ │ │ ├── disaggregation_utils.py │ │ │ ├── gsm8k_ascend_mixin.py │ │ │ ├── test_ascend_utils.py │ │ │ └── vlm_utils.py │ │ ├── attention/ │ │ │ ├── __init__.py │ │ │ ├── test_flashattn_backend.py │ │ │ ├── test_flashattn_mla_backend.py │ │ │ ├── test_prefix_chunk_info.py │ │ │ └── test_trtllm_mla_backend.py │ │ ├── bench_one_batch_server_internal.py │ │ ├── ci/ │ │ │ ├── __init__.py │ │ │ ├── ci_register.py │ │ │ ├── ci_stress_utils.py │ │ │ ├── ci_utils.py │ │ │ └── run_with_retry.py │ │ ├── doc_patch.py │ │ ├── external_models/ │ │ │ └── custom_qwen2_vl.py │ │ ├── few_shot_gsm8k.py │ │ ├── few_shot_gsm8k_engine.py │ │ ├── get_logits_ut.py │ │ ├── gpt_oss_common.py │ │ ├── kits/ │ │ │ ├── abort_timeout_kit.py │ │ │ ├── cache_hit_kit.py │ │ │ ├── ebnf_constrained_kit.py │ │ │ ├── gsm8k_accuracy_kit.py │ │ │ ├── json_constrained_kit.py │ │ │ ├── kl_divergence_kit.py │ │ │ ├── lm_eval_kit.py │ │ │ ├── matched_stop_kit.py │ │ │ ├── mmmu_vlm_kit.py │ │ │ ├── prefix_cache_branching_kit.py │ │ │ ├── radix_cache_server_kit.py │ │ │ ├── regex_constrained_kit.py │ │ │ └── spec_decoding_kit.py │ │ ├── kl_test_utils.py │ │ ├── long_prompt.txt │ │ ├── longbench_v2/ │ │ │ ├── __init__.py │ │ │ ├── longbench_v2_evaluation.md │ │ │ ├── test_longbench_v2_eval.py │ │ │ ├── validate_longbench_v2.py │ │ │ └── validate_longbench_v2_standalone.py │ │ ├── lora_utils.py │ │ ├── nightly_bench_utils.py │ │ ├── nightly_utils.py │ │ ├── performance_test_runner.py │ │ ├── run_combined_tests.py │ │ ├── run_eval.py │ │ ├── runners.py │ │ ├── send_one.py │ │ ├── server_fixtures/ │ │ │ ├── default_fixture.py │ │ │ ├── disaggregation_fixture.py │ │ │ ├── eagle_fixture.py │ │ │ └── mmmu_fixture.py │ │ ├── simple_eval_aime25.py │ │ ├── simple_eval_common.py │ │ ├── simple_eval_gpqa.py │ │ ├── simple_eval_gsm8k.py │ │ ├── simple_eval_humaneval.py │ │ ├── simple_eval_longbench_v2.py │ │ ├── simple_eval_math.py │ │ ├── simple_eval_mgsm.py │ │ ├── simple_eval_mmlu.py │ │ ├── simple_eval_mmmu_vlm.py │ │ ├── speculative/ │ │ │ └── test_spec_utils.py │ │ ├── test_activation.py │ │ ├── test_block_fp8.py │ │ ├── test_block_fp8_deep_gemm_blackwell.py │ │ ├── test_custom_ops.py │ │ ├── test_cutlass_moe.py │ │ ├── test_cutlass_w16a16_moe.py │ │ ├── test_cutlass_w4a8_moe.py │ │ ├── test_deepep_utils.py │ │ ├── test_deterministic.py │ │ ├── test_deterministic_utils.py │ │ ├── test_dump_metric.py │ │ ├── test_dynamic_grad_mode.py │ │ ├── test_flashinfer_dispatcher.py │ │ ├── test_http_server_auth.py │ │ ├── test_kvfp4_quant_dequant.py │ │ ├── test_layernorm.py │ │ ├── test_marlin_utils.py │ │ ├── test_programs.py │ │ ├── test_utils.py │ │ ├── tool_call_test_runner.py │ │ └── vlm_utils.py │ ├── utils.py │ └── version.py ├── scripts/ │ ├── check_vram_clear.sh │ ├── ci/ │ │ ├── amd/ │ │ │ ├── amd_ci_exec.sh │ │ │ ├── amd_ci_install_dependency.sh │ │ │ ├── amd_ci_start_container.sh │ │ │ ├── amd_ci_start_container_disagg.sh │ │ │ ├── amd_ci_warmup_aiter.py │ │ │ └── test_rccl_multi_gpu.py │ │ ├── cuda/ │ │ │ ├── ci_download_flashinfer_cubin.sh │ │ │ ├── ci_install_deepep.sh │ │ │ ├── ci_install_dependency.sh │ │ │ ├── ci_install_gateway_dependencies.sh │ │ │ ├── ci_start_disaggregation_servers.sh │ │ │ ├── prepare_runner.sh │ │ │ ├── warmup_deep_gemm.py │ │ │ └── warmup_server.py │ │ ├── musa/ │ │ │ ├── musa_install_dependency.sh │ │ │ └── rename_wheels_musa.sh │ │ ├── npu/ │ │ │ ├── npu_ci_install_dependency.sh │ │ │ └── npu_log_print.sh │ │ └── utils/ │ │ ├── ci_coverage_report.py │ │ ├── cleanup_hf_cache.py │ │ ├── merge_metrics.py │ │ ├── prevalidate_cached_models.py │ │ ├── publish_diffusion_gt.py │ │ ├── publish_traces.py │ │ ├── query_job_status.py │ │ ├── runner_utilization_report.py │ │ ├── save_diffusion_metrics.py │ │ ├── save_metrics.py │ │ └── slash_command_handler.py │ ├── ci_monitor/ │ │ ├── README.md │ │ ├── ci_failures_analysis.py │ │ └── post_ci_failures_to_slack.py │ ├── code_sync/ │ │ ├── check_commits.py │ │ ├── copy_from_oss.py │ │ ├── copy_to_oss.py │ │ ├── guideline.md │ │ ├── install_github_cli.sh │ │ └── utils.py │ ├── convert_otel_2_perfetto.py │ ├── ensure_vram_clear.sh │ ├── export_deepseek_nextn.py │ ├── killall_sglang.sh │ ├── playground/ │ │ ├── bench_speculative.py │ │ ├── disaggregation/ │ │ │ ├── cli-logprob.py │ │ │ ├── cli-so.py │ │ │ └── cli.py │ │ ├── frontend_reasoning.ipynb │ │ ├── load_tokenizer.py │ │ ├── long_context_example.py │ │ ├── lora/ │ │ │ ├── analyzer.py │ │ │ ├── lora_hf_play.py │ │ │ └── lora_vllm_play.py │ │ ├── reference_hf.py │ │ ├── replay_request_dump.py │ │ └── router/ │ │ ├── test_tree.py │ │ └── tree.py │ ├── release/ │ │ ├── README.md │ │ ├── bump_flashinfer_version.py │ │ ├── bump_kernel_version.py │ │ ├── bump_kernel_version_to_sglang.py │ │ ├── bump_sglang_version.py │ │ ├── check_kernel_version_to_sglang.py │ │ ├── commit_and_pr.sh │ │ ├── commit_and_pr_kernel_to_sglang.sh │ │ ├── test_utils.py │ │ └── utils.py │ ├── sort_testcases_alphabetically.py │ ├── update_kernel_whl_index.py │ ├── update_nightly_whl_index.py │ ├── update_pr_whl_index.py │ └── version_branch_to_tag.sh ├── sgl-kernel/ │ ├── .clang-format │ ├── CMakeLists.txt │ ├── Dockerfile │ ├── LICENSE │ ├── Makefile │ ├── README.md │ ├── THIRDPARTYNOTICES.txt │ ├── analyze_whl_kernel_sizes.py │ ├── benchmark/ │ │ ├── bench_activation.py │ │ ├── bench_amd_deterministic_allreduce.py │ │ ├── bench_awq_dequant.py │ │ ├── bench_cutlass_mla.py │ │ ├── bench_dsv3_fused_a_gemm.py │ │ ├── bench_dsv3_router_gemm.py │ │ ├── bench_es_fp8_blockwise_grouped_gemm.py │ │ ├── bench_fp4_gemm.py │ │ ├── bench_fp8_blockwise_gemm.py │ │ ├── bench_fp8_blockwise_group_gemm.py │ │ ├── bench_fp8_gemm.py │ │ ├── bench_int8_gemm.py │ │ ├── bench_kimi_k2_moe_fused_gate.py │ │ ├── bench_moe_align_block_size.py │ │ ├── bench_moe_ep_post_reorder.py │ │ ├── bench_moe_fused_gate.py │ │ ├── bench_moe_topk_sigmoid.py │ │ ├── bench_moe_topk_softmax.py │ │ ├── bench_mrope.py │ │ ├── bench_per_tensor_quant_fp8.py │ │ ├── bench_per_token_group_quant_8bit.py │ │ ├── bench_per_token_quant_fp8.py │ │ ├── bench_qserve_w4a8_gemm.py │ │ ├── bench_rmsnorm.py │ │ ├── bench_rotary_embedding.py │ │ ├── bench_sum_scale.py │ │ └── bench_top_k_top_p_sampling.py │ ├── build.sh │ ├── cmake/ │ │ ├── flashmla.cmake │ │ └── utils.cmake │ ├── csrc/ │ │ ├── allreduce/ │ │ │ ├── custom_all_reduce.cu │ │ │ ├── custom_all_reduce.cuh │ │ │ ├── custom_all_reduce.hip │ │ │ ├── custom_all_reduce_hip.cuh │ │ │ ├── deterministic_all_reduce.hip │ │ │ ├── mscclpp_allreduce.cu │ │ │ ├── mscclpp_allreduce.cuh │ │ │ ├── quick_all_reduce.cu │ │ │ ├── quick_all_reduce.cuh │ │ │ ├── quick_all_reduce.h │ │ │ ├── quick_all_reduce_base.h │ │ │ └── test_mscclpp_allreduce.cu │ │ ├── attention/ │ │ │ ├── cascade.cu │ │ │ ├── cutlass_mla_kernel.cu │ │ │ ├── cutlass_sm100_mla/ │ │ │ │ ├── device/ │ │ │ │ │ └── sm100_mla.hpp │ │ │ │ └── kernel/ │ │ │ │ ├── sm100_fmha_mla_reduction.hpp │ │ │ │ ├── sm100_fmha_mla_tma_warpspecialized.hpp │ │ │ │ └── sm100_mla_tile_scheduler.hpp │ │ │ ├── merge_attn_states.cu │ │ │ └── vertical_slash_index.cu │ │ ├── common_extension.cc │ │ ├── common_extension_musa.cc │ │ ├── common_extension_rocm.cc │ │ ├── cpu/ │ │ │ ├── CMakeLists.txt │ │ │ ├── aarch64/ │ │ │ │ └── shm.h │ │ │ ├── activation.cpp │ │ │ ├── bmm.cpp │ │ │ ├── common.h │ │ │ ├── conv3d.cpp │ │ │ ├── decode.cpp │ │ │ ├── extend.cpp │ │ │ ├── flash_attn.cpp │ │ │ ├── flash_attn.h │ │ │ ├── gemm.cpp │ │ │ ├── gemm.h │ │ │ ├── gemm_fp8.cpp │ │ │ ├── gemm_int4.cpp │ │ │ ├── gemm_int8.cpp │ │ │ ├── interface.cpp │ │ │ ├── mamba/ │ │ │ │ ├── conv.cpp │ │ │ │ └── fla.cpp │ │ │ ├── model/ │ │ │ │ └── qwen3.cpp │ │ │ ├── moe.cpp │ │ │ ├── moe_fp8.cpp │ │ │ ├── moe_int4.cpp │ │ │ ├── moe_int8.cpp │ │ │ ├── norm.cpp │ │ │ ├── numa_utils.cpp │ │ │ ├── preprocessor.cpp │ │ │ ├── qkv_proj.cpp │ │ │ ├── rope.cpp │ │ │ ├── shm.cpp │ │ │ ├── shm.h │ │ │ ├── topk.cpp │ │ │ ├── torch_extension_cpu.cpp │ │ │ ├── vec.h │ │ │ ├── vec_pack.h │ │ │ └── x86_64/ │ │ │ └── shm.h │ │ ├── cutlass_extensions/ │ │ │ ├── common.hpp │ │ │ ├── detail/ │ │ │ │ └── collective/ │ │ │ │ └── mixed_input_utils.hpp │ │ │ ├── epilogue/ │ │ │ │ └── epilogue_per_row_per_col_scale.h │ │ │ └── gemm/ │ │ │ ├── collective/ │ │ │ │ ├── builders/ │ │ │ │ │ └── sm90_gmma_builder_mixed_input.inl │ │ │ │ ├── collective_builder_mixed_input.hpp │ │ │ │ ├── collective_mma_array_mixed_input.hpp │ │ │ │ └── sm90_mma_array_tma_gmma_rs_warpspecialized_mixed_input_.hpp │ │ │ ├── cutlass_gemm_caller.cuh │ │ │ ├── dispatch_policy.hpp │ │ │ ├── fp8_blockwise_gemm_sm90_dispatch.cuh │ │ │ ├── gemm_universal_base_compat.h │ │ │ └── gemm_with_epilogue_visitor.h │ │ ├── elementwise/ │ │ │ ├── activation.cu │ │ │ ├── cast.cu │ │ │ ├── concat_mla.cu │ │ │ ├── copy.cu │ │ │ ├── fused_add_rms_norm_kernel.cu │ │ │ ├── pos_enc.cu │ │ │ ├── pos_enc.cuh │ │ │ ├── topk.cu │ │ │ └── utils.cuh │ │ ├── expert_specialization/ │ │ │ ├── es_fp8_blockwise.cu │ │ │ ├── es_fp8_blockwise_functor.cuh │ │ │ ├── es_fp8_blockwise_launcher.cuh │ │ │ ├── es_fp8_blockwise_traits.cuh │ │ │ ├── es_sm100_mxfp8_blockscaled.cu │ │ │ ├── es_sm100_mxfp8_blockscaled_functor.cuh │ │ │ ├── es_sm100_mxfp8_blockscaled_group_quant.cu │ │ │ ├── es_sm100_mxfp8_blockscaled_group_quant.cuh │ │ │ ├── es_sm100_mxfp8_blockscaled_launcher.cuh │ │ │ └── es_sm100_mxfp8_blockscaled_traits.cuh │ │ ├── flash_extension.cc │ │ ├── flashmla_extension.cc │ │ ├── gemm/ │ │ │ ├── awq_kernel.cu │ │ │ ├── bmm_fp8.cu │ │ │ ├── dsv3_fused_a_gemm.cu │ │ │ ├── dsv3_router_gemm_bf16_out.cu │ │ │ ├── dsv3_router_gemm_entry.cu │ │ │ ├── dsv3_router_gemm_float_out.cu │ │ │ ├── fp8_blockwise_gemm_kernel.cu │ │ │ ├── fp8_gemm_kernel.cu │ │ │ ├── gptq/ │ │ │ │ ├── compat.cuh │ │ │ │ ├── gptq_kernel.cu │ │ │ │ ├── matrix_view.cuh │ │ │ │ ├── qdq_2.cuh │ │ │ │ ├── qdq_3.cuh │ │ │ │ ├── qdq_4.cuh │ │ │ │ ├── qdq_8.cuh │ │ │ │ └── qdq_util.cuh │ │ │ ├── int8_gemm_kernel.cu │ │ │ ├── marlin/ │ │ │ │ ├── dequant.h │ │ │ │ ├── kernel.h │ │ │ │ ├── marlin.cuh │ │ │ │ ├── marlin_dtypes.cuh │ │ │ │ └── marlin_template.h │ │ │ ├── math.hpp │ │ │ ├── per_token_group_quant_8bit.cu │ │ │ ├── per_token_group_quant_8bit_v2.cu │ │ │ ├── per_token_quant_fp8.cu │ │ │ ├── qserve_w4a8_per_chn_gemm.cu │ │ │ └── qserve_w4a8_per_group_gemm.cu │ │ ├── grammar/ │ │ │ └── apply_token_bitmask_inplace_cuda.cu │ │ ├── kvcacheio/ │ │ │ └── transfer.cu │ │ ├── mamba/ │ │ │ ├── causal_conv1d.cu │ │ │ └── causal_conv1d.h │ │ ├── memory/ │ │ │ └── weak_ref_tensor.cpp │ │ ├── moe/ │ │ │ ├── cutlass_moe/ │ │ │ │ └── w4a8/ │ │ │ │ ├── scaled_mm_entry.cu │ │ │ │ ├── w4a8_get_group_starts.cuh │ │ │ │ ├── w4a8_grouped_mm_c3x.cu │ │ │ │ ├── w4a8_grouped_mm_c3x.cuh │ │ │ │ └── w4a8_moe_data.cu │ │ │ ├── cutlass_moe_helper.cu │ │ │ ├── fp8_blockwise_moe_kernel.cu │ │ │ ├── fused_qknorm_rope_kernel.cu │ │ │ ├── kimi_k2_moe_fused_gate.cu │ │ │ ├── moe_align_kernel.cu │ │ │ ├── moe_fused_gate.cu │ │ │ ├── moe_sum.cu │ │ │ ├── moe_sum_reduce.cu │ │ │ ├── moe_topk_sigmoid_kernels.cu │ │ │ ├── moe_topk_softmax_kernels.cu │ │ │ └── prepare_moe_input.cu │ │ ├── quantization/ │ │ │ └── gguf/ │ │ │ ├── dequantize.cuh │ │ │ ├── ggml-common.h │ │ │ ├── gguf_kernel.cu │ │ │ ├── mmq.cuh │ │ │ ├── mmvq.cuh │ │ │ ├── moe.cuh │ │ │ ├── moe_vec.cuh │ │ │ └── vecdotq.cuh │ │ ├── spatial/ │ │ │ ├── cuda_utils.h │ │ │ ├── greenctx_stream.cu │ │ │ └── greenctx_stream.h │ │ ├── spatial_extension.cc │ │ └── speculative/ │ │ ├── eagle_utils.cu │ │ ├── ngram_utils.cu │ │ ├── packbit.cu │ │ ├── speculative_sampling.cu │ │ └── speculative_sampling.cuh │ ├── include/ │ │ ├── hip/ │ │ │ ├── hip_act_and_mul.cuh │ │ │ ├── hip_math_def.h │ │ │ ├── hip_vec_dtypes.h │ │ │ └── impl/ │ │ │ ├── hip_vec_bf16_impl.h │ │ │ ├── hip_vec_fp32_impl.h │ │ │ └── hip_vec_half_impl.h │ │ ├── pytorch_extension_utils_rocm.h │ │ ├── scalar_type.hpp │ │ ├── sgl_flash_kernel_ops.h │ │ ├── sgl_kernel_ops.h │ │ ├── sgl_kernel_torch_shim.h │ │ └── utils.h │ ├── kernel-runner-setup.sh │ ├── pyproject.toml │ ├── pyproject_cpu.toml │ ├── pyproject_musa.toml │ ├── pyproject_rocm.toml │ ├── python/ │ │ └── sgl_kernel/ │ │ ├── __init__.py │ │ ├── _fa4_interface.py │ │ ├── allreduce.py │ │ ├── attention.py │ │ ├── cutlass_moe.py │ │ ├── elementwise.py │ │ ├── expert_specialization.py │ │ ├── flash_attn.py │ │ ├── flash_mla.py │ │ ├── gemm.py │ │ ├── grammar.py │ │ ├── kvcacheio.py │ │ ├── load_utils.py │ │ ├── mamba.py │ │ ├── memory.py │ │ ├── moe.py │ │ ├── quantization/ │ │ │ ├── __init__.py │ │ │ └── gguf.py │ │ ├── sampling.py │ │ ├── scalar_type.py │ │ ├── sparse_flash_attn.py │ │ ├── spatial.py │ │ ├── speculative.py │ │ ├── test_utils.py │ │ ├── testing/ │ │ │ ├── __init__.py │ │ │ └── rotary_embedding.py │ │ ├── top_k.py │ │ ├── utils.py │ │ └── version.py │ ├── rename_wheels.sh │ ├── setup_musa.py │ ├── setup_rocm.py │ └── tests/ │ ├── conftest.py │ ├── spatial/ │ │ └── test_greenctx_stream.py │ ├── speculative/ │ │ ├── test_eagle_utils.py │ │ ├── test_ngram_utils.py │ │ └── test_speculative_sampling.py │ ├── test_activation.py │ ├── test_amd_deterministic_custom_allreduce.py │ ├── test_amd_nccl_allreduce_determinism.py │ ├── test_apply_token_bitmask_inplace.py │ ├── test_awq_dequant.py │ ├── test_bmm_fp8.py │ ├── test_causal_conv1d.py │ ├── test_copy.py │ ├── test_custom_allreduce.py │ ├── test_cutlass_mla.py │ ├── test_cutlass_w4a8_moe_mm.py │ ├── test_dsv3_fused_a_gemm.py │ ├── test_dsv3_router_gemm.py │ ├── test_es_fp8_blockwise_moe.py │ ├── test_es_mxfp8_blockscaled_moe.py │ ├── test_flash_attention.py │ ├── test_flash_attn_sparse.py │ ├── test_flashmla.py │ ├── test_fp8_blockwise_gemm.py │ ├── test_fp8_blockwise_moe.py │ ├── test_fp8_gemm.py │ ├── test_fused_qk_norm_rope.py │ ├── test_gguf.py │ ├── test_gptq_kernel.py │ ├── test_hadamard.py │ ├── test_int8_gemm.py │ ├── test_kimi_k2_moe_fused_gate.py │ ├── test_kvcacheio.py │ ├── test_merge_state.py │ ├── test_merge_state_v2.py │ ├── test_moe_align.py │ ├── test_moe_fused_gate.py │ ├── test_moe_topk_sigmoid.py │ ├── test_moe_topk_softmax.py │ ├── test_mscclpp.py │ ├── test_norm.py │ ├── test_per_token_group_quant_8bit.py │ ├── test_per_token_quant_fp8.py │ ├── test_qserve_w4a8_per_chn_gemm.py │ ├── test_qserve_w4a8_per_group_gemm.py │ ├── test_sampling.py │ ├── test_topk.py │ ├── test_torch_defaults_reset.py │ └── utils.py ├── sgl-model-gateway/ │ ├── .cargo/ │ │ └── config.toml │ ├── Cargo.toml │ ├── Makefile │ ├── README.md │ ├── benches/ │ │ ├── consistent_hash_bench.rs │ │ ├── manual_policy_benchmark.rs │ │ ├── request_processing.rs │ │ ├── router_registry_bench.rs │ │ ├── tree_benchmark.rs │ │ └── wasm_middleware_latency.rs │ ├── bindings/ │ │ ├── golang/ │ │ │ ├── .gitignore │ │ │ ├── Cargo.toml │ │ │ ├── Makefile │ │ │ ├── README.md │ │ │ ├── client.go │ │ │ ├── client_test.go │ │ │ ├── examples/ │ │ │ │ ├── oai_server/ │ │ │ │ │ ├── Makefile │ │ │ │ │ ├── README.md │ │ │ │ │ ├── config/ │ │ │ │ │ │ └── config.go │ │ │ │ │ ├── docs/ │ │ │ │ │ │ └── benchmark_result.md │ │ │ │ │ ├── go.sum │ │ │ │ │ ├── handlers/ │ │ │ │ │ │ ├── chat.go │ │ │ │ │ │ ├── health.go │ │ │ │ │ │ └── models.go │ │ │ │ │ ├── logger/ │ │ │ │ │ │ └── logger.go │ │ │ │ │ ├── main.go │ │ │ │ │ ├── models/ │ │ │ │ │ │ └── chat.go │ │ │ │ │ ├── run.sh │ │ │ │ │ ├── scripts/ │ │ │ │ │ │ ├── analyze_tpot.sh │ │ │ │ │ │ ├── pprof_analysis.sh │ │ │ │ │ │ ├── pprof_quick.sh │ │ │ │ │ │ ├── pprof_test.sh │ │ │ │ │ │ └── profile_tpot.sh │ │ │ │ │ ├── service/ │ │ │ │ │ │ └── sglang.go │ │ │ │ │ └── utils/ │ │ │ │ │ └── utils.go │ │ │ │ ├── simple/ │ │ │ │ │ ├── main.go │ │ │ │ │ └── run.sh │ │ │ │ └── streaming/ │ │ │ │ ├── main.go │ │ │ │ └── run.sh │ │ │ ├── go.sum │ │ │ ├── integration_test.go │ │ │ ├── internal/ │ │ │ │ ├── ffi/ │ │ │ │ │ ├── batch_postprocessor.go │ │ │ │ │ ├── client.go │ │ │ │ │ ├── grpc_converter.go │ │ │ │ │ ├── postprocessor.go │ │ │ │ │ └── preprocessor.go │ │ │ │ ├── grpc/ │ │ │ │ │ └── client_grpc.go │ │ │ │ └── proto/ │ │ │ │ ├── sglang_scheduler.pb.go │ │ │ │ └── sglang_scheduler_grpc.pb.go │ │ │ └── src/ │ │ │ ├── client.rs │ │ │ ├── error.rs │ │ │ ├── grpc_converter.rs │ │ │ ├── lib.rs │ │ │ ├── memory.rs │ │ │ ├── postprocessor.rs │ │ │ ├── preprocessor.rs │ │ │ ├── stream.rs │ │ │ ├── tokenizer.rs │ │ │ ├── tool_parser.rs │ │ │ └── utils.rs │ │ └── python/ │ │ ├── .coveragerc │ │ ├── Cargo.toml │ │ ├── MANIFEST.in │ │ ├── README.md │ │ ├── pyproject.toml │ │ ├── setup.py │ │ ├── src/ │ │ │ ├── lib.rs │ │ │ └── sglang_router/ │ │ │ ├── __init__.py │ │ │ ├── __main__.py │ │ │ ├── cli.py │ │ │ ├── launch_router.py │ │ │ ├── launch_server.py │ │ │ ├── mini_lb.py │ │ │ ├── router.py │ │ │ ├── router_args.py │ │ │ └── version.py │ │ └── tests/ │ │ ├── conftest.py │ │ ├── test_arg_parser.py │ │ ├── test_router_config.py │ │ ├── test_startup_sequence.py │ │ └── test_validation.py │ ├── build.rs │ ├── e2e_test/ │ │ ├── __init__.py │ │ ├── benchmarks/ │ │ │ ├── __init__.py │ │ │ ├── conftest.py │ │ │ ├── results.py │ │ │ ├── summarize.py │ │ │ ├── test_pd_perf.py │ │ │ └── test_regular_perf.py │ │ ├── chat_completions/ │ │ │ ├── __init__.py │ │ │ ├── test_enable_thinking.py │ │ │ ├── test_function_calling.py │ │ │ ├── test_openai_server.py │ │ │ ├── test_reasoning_content.py │ │ │ └── test_validation.py │ │ ├── conftest.py │ │ ├── embeddings/ │ │ │ ├── __init__.py │ │ │ ├── test_basic.py │ │ │ └── test_correctness.py │ │ ├── fixtures/ │ │ │ ├── __init__.py │ │ │ ├── hooks.py │ │ │ ├── markers.py │ │ │ ├── pool.py │ │ │ ├── ports.py │ │ │ └── setup_backend.py │ │ ├── infra/ │ │ │ ├── __init__.py │ │ │ ├── constants.py │ │ │ ├── gateway.py │ │ │ ├── gpu_allocator.py │ │ │ ├── gpu_monitor.py │ │ │ ├── model_pool.py │ │ │ ├── model_specs.py │ │ │ ├── process_utils.py │ │ │ ├── run_eval.py │ │ │ ├── simple_eval_common.py │ │ │ └── simple_eval_mmlu.py │ │ ├── pyproject.toml │ │ ├── responses/ │ │ │ ├── __init__.py │ │ │ ├── test_basic_crud.py │ │ │ ├── test_state_management.py │ │ │ ├── test_streaming_events.py │ │ │ ├── test_structured_output.py │ │ │ └── test_tools_call.py │ │ └── router/ │ │ ├── __init__.py │ │ ├── test_mmlu.py │ │ ├── test_pd_mmlu.py │ │ └── test_worker_api.py │ ├── examples/ │ │ └── wasm/ │ │ ├── .gitignore │ │ ├── README.md │ │ ├── wasm-guest-auth/ │ │ │ ├── Cargo.toml │ │ │ ├── README.md │ │ │ ├── build.sh │ │ │ └── src/ │ │ │ └── lib.rs │ │ ├── wasm-guest-logging/ │ │ │ ├── Cargo.toml │ │ │ ├── README.md │ │ │ ├── build.sh │ │ │ └── src/ │ │ │ └── lib.rs │ │ └── wasm-guest-ratelimit/ │ │ ├── Cargo.toml │ │ ├── README.md │ │ ├── build.sh │ │ └── src/ │ │ └── lib.rs │ ├── pytest.ini │ ├── rustfmt.toml │ ├── scripts/ │ │ ├── generate_gateway_release_notes.sh │ │ ├── generate_vision_golden.py │ │ ├── run_benchmarks.py │ │ └── setup-sccache.sh │ ├── src/ │ │ ├── app_context.rs │ │ ├── config/ │ │ │ ├── builder.rs │ │ │ ├── mod.rs │ │ │ ├── types.rs │ │ │ └── validation.rs │ │ ├── core/ │ │ │ ├── circuit_breaker.rs │ │ │ ├── error.rs │ │ │ ├── job_queue.rs │ │ │ ├── metrics_aggregator.rs │ │ │ ├── mod.rs │ │ │ ├── model_card.rs │ │ │ ├── model_type.rs │ │ │ ├── retry.rs │ │ │ ├── steps/ │ │ │ │ ├── mcp_registration.rs │ │ │ │ ├── mod.rs │ │ │ │ ├── tokenizer_registration.rs │ │ │ │ ├── wasm_module_registration.rs │ │ │ │ ├── wasm_module_removal.rs │ │ │ │ ├── worker/ │ │ │ │ │ ├── external/ │ │ │ │ │ │ ├── create_workers.rs │ │ │ │ │ │ ├── discover_models.rs │ │ │ │ │ │ └── mod.rs │ │ │ │ │ ├── local/ │ │ │ │ │ │ ├── create_worker.rs │ │ │ │ │ │ ├── detect_connection.rs │ │ │ │ │ │ ├── discover_dp.rs │ │ │ │ │ │ ├── discover_metadata.rs │ │ │ │ │ │ ├── find_worker_to_update.rs │ │ │ │ │ │ ├── find_workers_to_remove.rs │ │ │ │ │ │ ├── mod.rs │ │ │ │ │ │ ├── remove_from_policy_registry.rs │ │ │ │ │ │ ├── remove_from_worker_registry.rs │ │ │ │ │ │ ├── submit_tokenizer_job.rs │ │ │ │ │ │ ├── update_policies_for_worker.rs │ │ │ │ │ │ ├── update_remaining_policies.rs │ │ │ │ │ │ └── update_worker_properties.rs │ │ │ │ │ ├── mod.rs │ │ │ │ │ └── shared/ │ │ │ │ │ ├── activate.rs │ │ │ │ │ ├── mod.rs │ │ │ │ │ ├── register.rs │ │ │ │ │ └── update_policies.rs │ │ │ │ ├── workflow_data.rs │ │ │ │ └── workflow_engines.rs │ │ │ ├── token_bucket.rs │ │ │ ├── worker.rs │ │ │ ├── worker_builder.rs │ │ │ ├── worker_manager.rs │ │ │ ├── worker_registry.rs │ │ │ └── worker_service.rs │ │ ├── lib.rs │ │ ├── main.rs │ │ ├── middleware.rs │ │ ├── observability/ │ │ │ ├── events.rs │ │ │ ├── gauge_histogram.rs │ │ │ ├── inflight_tracker.rs │ │ │ ├── logging.rs │ │ │ ├── metrics.rs │ │ │ ├── mod.rs │ │ │ └── otel_trace.rs │ │ ├── policies/ │ │ │ ├── bucket.rs │ │ │ ├── cache_aware.rs │ │ │ ├── consistent_hashing.rs │ │ │ ├── factory.rs │ │ │ ├── manual.rs │ │ │ ├── mod.rs │ │ │ ├── power_of_two.rs │ │ │ ├── prefix_hash.rs │ │ │ ├── random.rs │ │ │ ├── registry.rs │ │ │ ├── round_robin.rs │ │ │ ├── tree.rs │ │ │ └── utils.rs │ │ ├── routers/ │ │ │ ├── conversations/ │ │ │ │ ├── handlers.rs │ │ │ │ └── mod.rs │ │ │ ├── error.rs │ │ │ ├── factory.rs │ │ │ ├── grpc/ │ │ │ │ ├── client.rs │ │ │ │ ├── common/ │ │ │ │ │ ├── mod.rs │ │ │ │ │ ├── response_collection.rs │ │ │ │ │ ├── response_formatting.rs │ │ │ │ │ ├── responses/ │ │ │ │ │ │ ├── context.rs │ │ │ │ │ │ ├── handlers.rs │ │ │ │ │ │ ├── mod.rs │ │ │ │ │ │ ├── streaming.rs │ │ │ │ │ │ └── utils.rs │ │ │ │ │ └── stages/ │ │ │ │ │ ├── client_acquisition.rs │ │ │ │ │ ├── dispatch_metadata.rs │ │ │ │ │ ├── helpers.rs │ │ │ │ │ ├── mod.rs │ │ │ │ │ ├── request_execution.rs │ │ │ │ │ └── worker_selection.rs │ │ │ │ ├── context.rs │ │ │ │ ├── harmony/ │ │ │ │ │ ├── builder.rs │ │ │ │ │ ├── detector.rs │ │ │ │ │ ├── mod.rs │ │ │ │ │ ├── parser.rs │ │ │ │ │ ├── processor.rs │ │ │ │ │ ├── responses/ │ │ │ │ │ │ ├── common.rs │ │ │ │ │ │ ├── execution.rs │ │ │ │ │ │ ├── mod.rs │ │ │ │ │ │ ├── non_streaming.rs │ │ │ │ │ │ └── streaming.rs │ │ │ │ │ ├── stages/ │ │ │ │ │ │ ├── mod.rs │ │ │ │ │ │ ├── preparation.rs │ │ │ │ │ │ ├── request_building.rs │ │ │ │ │ │ └── response_processing.rs │ │ │ │ │ ├── streaming.rs │ │ │ │ │ └── types.rs │ │ │ │ ├── mod.rs │ │ │ │ ├── pd_router.rs │ │ │ │ ├── pipeline.rs │ │ │ │ ├── proto_wrapper.rs │ │ │ │ ├── regular/ │ │ │ │ │ ├── mod.rs │ │ │ │ │ ├── processor.rs │ │ │ │ │ ├── responses/ │ │ │ │ │ │ ├── common.rs │ │ │ │ │ │ ├── conversions.rs │ │ │ │ │ │ ├── handlers.rs │ │ │ │ │ │ ├── mod.rs │ │ │ │ │ │ ├── non_streaming.rs │ │ │ │ │ │ └── streaming.rs │ │ │ │ │ ├── stages/ │ │ │ │ │ │ ├── chat/ │ │ │ │ │ │ │ ├── mod.rs │ │ │ │ │ │ │ ├── preparation.rs │ │ │ │ │ │ │ ├── request_building.rs │ │ │ │ │ │ │ └── response_processing.rs │ │ │ │ │ │ ├── classify/ │ │ │ │ │ │ │ ├── mod.rs │ │ │ │ │ │ │ └── response_processing.rs │ │ │ │ │ │ ├── embedding/ │ │ │ │ │ │ │ ├── mod.rs │ │ │ │ │ │ │ ├── preparation.rs │ │ │ │ │ │ │ ├── request_building.rs │ │ │ │ │ │ │ └── response_processing.rs │ │ │ │ │ │ ├── generate/ │ │ │ │ │ │ │ ├── mod.rs │ │ │ │ │ │ │ ├── preparation.rs │ │ │ │ │ │ │ ├── request_building.rs │ │ │ │ │ │ │ └── response_processing.rs │ │ │ │ │ │ ├── mod.rs │ │ │ │ │ │ ├── preparation.rs │ │ │ │ │ │ ├── request_building.rs │ │ │ │ │ │ └── response_processing.rs │ │ │ │ │ └── streaming.rs │ │ │ │ ├── router.rs │ │ │ │ └── utils.rs │ │ │ ├── header_utils.rs │ │ │ ├── http/ │ │ │ │ ├── mod.rs │ │ │ │ ├── pd_router.rs │ │ │ │ ├── pd_types.rs │ │ │ │ └── router.rs │ │ │ ├── mcp_utils.rs │ │ │ ├── mesh/ │ │ │ │ ├── handlers.rs │ │ │ │ └── mod.rs │ │ │ ├── mod.rs │ │ │ ├── openai/ │ │ │ │ ├── context.rs │ │ │ │ ├── mod.rs │ │ │ │ ├── provider.rs │ │ │ │ ├── responses/ │ │ │ │ │ ├── accumulator.rs │ │ │ │ │ ├── common.rs │ │ │ │ │ ├── mcp.rs │ │ │ │ │ ├── mod.rs │ │ │ │ │ ├── non_streaming.rs │ │ │ │ │ ├── streaming.rs │ │ │ │ │ ├── tool_handler.rs │ │ │ │ │ └── utils.rs │ │ │ │ └── router.rs │ │ │ ├── parse/ │ │ │ │ ├── handlers.rs │ │ │ │ └── mod.rs │ │ │ ├── persistence_utils.rs │ │ │ ├── router_manager.rs │ │ │ └── tokenize/ │ │ │ ├── handlers.rs │ │ │ └── mod.rs │ │ ├── server.rs │ │ ├── service_discovery.rs │ │ ├── version.rs │ │ └── wasm/ │ │ ├── mod.rs │ │ └── route.rs │ └── tests/ │ ├── api/ │ │ ├── api_endpoints_test.rs │ │ ├── mod.rs │ │ ├── parser_endpoints_test.rs │ │ ├── request_formats_test.rs │ │ ├── responses_api_test.rs │ │ └── streaming_tests.rs │ ├── api_tests.rs │ ├── common/ │ │ ├── mock_mcp_server.rs │ │ ├── mock_openai_server.rs │ │ ├── mock_worker.rs │ │ ├── mod.rs │ │ ├── redis_test_server.rs │ │ ├── streaming_helpers.rs │ │ ├── test_app.rs │ │ ├── test_certs.rs │ │ ├── test_config.rs │ │ └── tls_mock_worker.rs │ ├── inflight_tracker_test.rs │ ├── load_guard_raii_test.rs │ ├── mcp_test.rs │ ├── metrics_aggregator_test.rs │ ├── otel_tracing_test.rs │ ├── reliability/ │ │ ├── circuit_breaker_test.rs │ │ ├── fault_tolerance_test.rs │ │ ├── mod.rs │ │ ├── rate_limiting_test.rs │ │ └── retries_test.rs │ ├── reliability_tests.rs │ ├── routing/ │ │ ├── cache_aware_backward_compat_test.rs │ │ ├── header_forwarding_test.rs │ │ ├── load_balancing_test.rs │ │ ├── manual_routing_test.rs │ │ ├── mod.rs │ │ ├── payload_size_test.rs │ │ ├── pd_routing_test.rs │ │ ├── policy_registry_integration.rs │ │ ├── power_of_two_test.rs │ │ ├── service_discovery_test.rs │ │ ├── test_openai_routing.rs │ │ ├── test_pd_routing.rs │ │ └── worker_management_test.rs │ ├── routing_tests.rs │ ├── security/ │ │ ├── auth_integration_test.rs │ │ ├── auth_test.rs │ │ ├── mod.rs │ │ └── mtls_test.rs │ ├── security_tests.rs │ ├── spec/ │ │ ├── chat_completion.rs │ │ ├── chat_message.rs │ │ ├── embedding.rs │ │ ├── mod.rs │ │ ├── rerank.rs │ │ └── responses.rs │ ├── spec_test.rs │ └── wasm_test.rs └── test/ ├── README.md ├── lm_eval_configs/ │ ├── NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml │ ├── NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml │ └── Qwen3.5-397B-A17B.yaml ├── manual/ │ ├── ascend/ │ │ ├── test_ascend_deepseek_mtp.py │ │ ├── test_ascend_w8a8_quantization.py │ │ └── test_mindspore_models.py │ ├── cpu/ │ │ └── test_comm.py │ ├── debug_utils/ │ │ └── test_log_parser.py │ ├── entrypoints/ │ │ └── http_server/ │ │ └── test_abort_request.py │ ├── ep/ │ │ ├── test_deepep_internode.py │ │ ├── test_deepep_intranode.py │ │ ├── test_deepep_low_latency.py │ │ ├── test_eplb.py │ │ ├── test_moe_deepep.py │ │ ├── test_moe_deepep_eval_accuracy_large.py │ │ ├── test_mooncake_expert_backup.py │ │ └── test_nixl_ep.py │ ├── hicache/ │ │ ├── test_disaggregation_hicache.py │ │ └── test_pp_with_hicache.py │ ├── kv_transfer/ │ │ └── test_mooncake_transfer_engine.py │ ├── lang_frontend/ │ │ ├── test_bind_cache.py │ │ ├── test_choices.py │ │ ├── test_jump_forward.py │ │ ├── test_openai_backend.py │ │ ├── test_separate_reasoning.py │ │ └── test_separate_reasoning_execution.py │ ├── layers/ │ │ ├── attention/ │ │ │ └── nsa/ │ │ │ ├── test_act_quant_triton.py │ │ │ ├── test_get_k_scale_triton_kernel.py │ │ │ └── test_index_buf_accessor.py │ │ └── moe/ │ │ ├── test_moe_runners_1gpu.py │ │ └── test_moe_runners_4gpu.py │ ├── lora/ │ │ ├── test_lora_cuda_graph.py │ │ ├── test_lora_llama4.py │ │ ├── test_lora_ops.py │ │ ├── test_lora_qwen3_vl.py │ │ ├── test_lora_spec_decoding.py │ │ └── test_torch_backend.py │ ├── models/ │ │ ├── test_clip_models.py │ │ ├── test_falcon_h1_models.py │ │ ├── test_gme_qwen_models.py │ │ ├── test_grok_models.py │ │ ├── test_kimi_k2_models.py │ │ ├── test_llama4_models.py │ │ ├── test_mistral_large3_basic.py │ │ ├── test_mtp_models.py │ │ └── test_unsloth_models.py │ ├── nightly/ │ │ ├── test_deepseek_v31_perf.py │ │ ├── test_deepseek_v32_perf.py │ │ ├── test_text_models_gsm8k_eval.py │ │ ├── test_text_models_perf.py │ │ ├── test_vlms_mmmu_eval.py │ │ ├── test_vlms_perf.py │ │ ├── test_vlms_piecewise_cuda_graph.py │ │ ├── test_vlms_vit_cuda_graph.py │ │ └── test_vlms_vit_flashinfer_cudnn.py │ ├── openai_server/ │ │ └── features/ │ │ ├── test_cache_report.py │ │ ├── test_continuous_usage_stats.py │ │ └── test_structural_tag.py │ ├── piecewise_cudagraph/ │ │ └── test_disaggregation_piecewise_cuda_graph.py │ ├── quant/ │ │ └── test_fp8_kvcache.py │ ├── test_async_dynamic_batch_tokenizer.py │ ├── test_async_mm_data_processor.py │ ├── test_config_integration.py │ ├── test_custom_allreduce.py │ ├── test_deepseek_chat_templates.py │ ├── test_double_sparsity.py │ ├── test_expert_distribution.py │ ├── test_expert_location_updater.py │ ├── test_fim_completion.py │ ├── test_forward_split_prefill.py │ ├── test_get_weights_by_name.py │ ├── test_health_check.py │ ├── test_kv_events.py │ ├── test_logprobs.py │ ├── test_mla_tp.py │ ├── test_modelopt.py │ ├── test_modelopt_fp8kvcache.py │ ├── test_models_from_modelscope.py │ ├── test_mori_transfer_engine_e2e.py │ ├── test_mscclpp.py │ ├── test_quick_allreduce.py │ ├── test_ray_engine.py │ ├── test_sagemaker_server.py │ ├── test_schedule_policy.py │ ├── test_srt_engine_with_quant_args.py │ ├── test_tokenizer_batch_encode.py │ ├── test_tokenizer_manager.py │ ├── test_torch_flex_attention_backend.py │ ├── test_torch_tp.py │ ├── test_tracing.py │ ├── test_triton_attention_rocm_mla.py │ ├── test_triton_moe_wna16.py │ ├── test_trtllm_fp8_kv_kernel.py │ ├── test_two_batch_overlap.py │ ├── test_vertex_endpoint.py │ ├── test_vlm_accuracy.py │ ├── test_wave_attention_backend.py │ ├── test_weight_validation.py │ ├── test_weight_version.py │ └── vlm/ │ └── test_anthropic_vision.py ├── pytest.ini ├── registered/ │ ├── 4-gpu-models/ │ │ ├── test_deepseek_v3_cutedsl_4gpu.py │ │ ├── test_gpt_oss_4gpu.py │ │ ├── test_nvidia_nemotron_3_super_nvfp4.py │ │ ├── test_qwen35_models.py │ │ ├── test_qwen3_next_models.py │ │ └── test_qwen3_next_models_mtp.py │ ├── 8-gpu-models/ │ │ ├── test_deepseek_v31.py │ │ ├── test_deepseek_v32.py │ │ ├── test_deepseek_v32_basic.py │ │ ├── test_deepseek_v32_cp_single_node.py │ │ ├── test_deepseek_v32_mtp.py │ │ ├── test_deepseek_v3_basic.py │ │ ├── test_deepseek_v3_mtp.py │ │ ├── test_glm_46.py │ │ ├── test_glm_46_fp8.py │ │ ├── test_gpt_oss_120b.py │ │ ├── test_kimi_k25.py │ │ ├── test_llama4.py │ │ ├── test_mimo_models.py │ │ ├── test_minimax_m25.py │ │ ├── test_mistral_large3.py │ │ ├── test_nvidia_nemotron_3_super_bf16.py │ │ ├── test_nvidia_nemotron_3_super_nightly.py │ │ ├── test_qwen35.py │ │ ├── test_qwen3_235b.py │ │ └── test_ring_2_5_1t.py │ ├── README.md │ ├── amd/ │ │ ├── accuracy/ │ │ │ ├── mi30x/ │ │ │ │ ├── test_deepseek_r1_eval_amd.py │ │ │ │ ├── test_deepseek_v31_eval_amd.py │ │ │ │ ├── test_deepseek_v32_dp_eval_amd.py │ │ │ │ ├── test_deepseek_v32_eval_amd.py │ │ │ │ ├── test_deepseek_v32_mtp_eval_amd.py │ │ │ │ ├── test_deepseek_v32_tc_eval_amd.py │ │ │ │ ├── test_glm5_eval_amd.py │ │ │ │ ├── test_gpt_oss_eval_amd.py │ │ │ │ ├── test_grok1_fp8_eval_amd.py │ │ │ │ ├── test_grok1_int4_eval_amd.py │ │ │ │ ├── test_grok2_eval_amd.py │ │ │ │ ├── test_grok_eval_amd.py │ │ │ │ ├── test_gsm8k_eval_amd.py │ │ │ │ ├── test_kimi_k25_eval_amd.py │ │ │ │ ├── test_kimi_k2_eval_amd.py │ │ │ │ ├── test_minimax_m25_eval_amd.py │ │ │ │ ├── test_qwen35_eval_amd.py │ │ │ │ └── test_vlms_mmmu_eval_amd.py │ │ │ └── mi35x/ │ │ │ ├── test_deepseek_r1_eval_mi35x.py │ │ │ ├── test_deepseek_r1_mxfp4_ar_fusion_eval_mi35x.py │ │ │ ├── test_deepseek_r1_mxfp4_eval_mi35x.py │ │ │ ├── test_deepseek_r1_mxfp4_kv_fp8_eval_mi35x.py │ │ │ ├── test_deepseek_v32_dp_eval_mi35x.py │ │ │ ├── test_deepseek_v32_eval_mi35x.py │ │ │ ├── test_deepseek_v32_mtp_eval_mi35x.py │ │ │ ├── test_glm5_eval_mi35x.py │ │ │ ├── test_gpt_oss_eval_mi35x.py │ │ │ ├── test_grok1_int4_eval_mi35x.py │ │ │ ├── test_grok2_eval_mi35x.py │ │ │ ├── test_kimi_k25_aiter_mla_eval_mi35x.py │ │ │ ├── test_kimi_k25_eval_mi35x.py │ │ │ ├── test_kimi_k25_mxfp4_eval_mi35x.py │ │ │ ├── test_kimi_k2_eval_mi35x.py │ │ │ ├── test_minimax_m25_eval_mi35x.py │ │ │ ├── test_qwen35_eval_mi35x.py │ │ │ └── test_qwen3_coder_next_eval_mi35x.py │ │ ├── disaggregation/ │ │ │ ├── test_disaggregation_basic.py │ │ │ └── test_disaggregation_pp.py │ │ ├── perf/ │ │ │ ├── mi30x/ │ │ │ │ ├── test_deepseek_v31_perf.py │ │ │ │ ├── test_deepseek_v32_basic_perf_amd.py │ │ │ │ ├── test_deepseek_v32_mtp_perf_amd.py │ │ │ │ ├── test_deepseek_v3_perf.py │ │ │ │ ├── test_grok1_fp8_perf.py │ │ │ │ ├── test_grok1_int4_perf.py │ │ │ │ ├── test_grok2_perf.py │ │ │ │ ├── test_text_models_perf_amd.py │ │ │ │ └── test_vlms_perf_amd.py │ │ │ └── mi35x/ │ │ │ ├── test_deepseek_r1_mxfp4_ar_fusion_perf_mi35x.py │ │ │ ├── test_deepseek_r1_mxfp4_kv_fp8_perf_mi35x.py │ │ │ ├── test_deepseek_r1_mxfp4_perf_mi35x.py │ │ │ ├── test_deepseek_v32_basic_perf_mi35x.py │ │ │ ├── test_deepseek_v32_mtp_perf_mi35x.py │ │ │ ├── test_grok1_int4_perf_mi35x.py │ │ │ └── test_grok2_perf_mi35x.py │ │ ├── test_deepseek_r1_mxfp4_8gpu.py │ │ ├── test_deepseek_v32_basic.py │ │ ├── test_deepseek_v32_mtp.py │ │ ├── test_deepseek_v3_basic.py │ │ ├── test_deepseek_v3_basic_kv_fp8.py │ │ ├── test_deepseek_v3_mtp.py │ │ ├── test_deepseek_v3_mtp_kv_fp8.py │ │ ├── test_kimi_k25_mxfp4.py │ │ ├── test_kimi_k2_instruct.py │ │ ├── test_moriep_small.py │ │ ├── test_qwen3_coder_next_8gpu.py │ │ ├── test_qwen3_instruct.py │ │ ├── test_qwen3_instruct_fp8.py │ │ ├── test_qwen3_instruct_mxfp4.py │ │ └── test_zimage_turbo.py │ ├── ascend/ │ │ ├── basic_function/ │ │ │ ├── HiCache/ │ │ │ │ ├── test_npu_hierarchical_cache.py │ │ │ │ ├── test_npu_hierarchical_cache_mla.py │ │ │ │ ├── test_npu_hierarchical_cache_mutually_exclusive.py │ │ │ │ ├── test_npu_hierarchical_cache_ttft_mha.py │ │ │ │ └── test_npu_radix_cache.py │ │ │ ├── parallel_strategy/ │ │ │ │ └── expert_parallelism/ │ │ │ │ ├── test_npu_deepep_auto_deepseek_v3_2_w8a8.py │ │ │ │ ├── test_npu_deepep_auto_qwen3_480b.py │ │ │ │ ├── test_npu_deepep_auto_qwen3_next.py │ │ │ │ ├── test_npu_deepep_low_latency_deepseek_v3_2_w8a8.py │ │ │ │ ├── test_npu_deepep_low_latency_qwen3_480b.py │ │ │ │ └── test_npu_deepep_low_latency_qwen3_next.py │ │ │ ├── parameter/ │ │ │ │ ├── deepseek_coder.json │ │ │ │ ├── test_npu_fim_completion.py │ │ │ │ ├── test_npu_log_level.py │ │ │ │ ├── test_npu_no_chunked_prefill.py │ │ │ │ ├── test_npu_no_overlap_scheduler.py │ │ │ │ ├── test_npu_original_logprobs.py │ │ │ │ └── test_npu_warmups.py │ │ │ └── speculative_inference/ │ │ │ └── test_npu_eagle3.py │ │ ├── embedding_models/ │ │ │ └── test_npu_bge_large_en_v1_5.py │ │ ├── interface/ │ │ │ ├── test_npu_api.py │ │ │ ├── test_npu_api_abort_request.py │ │ │ ├── test_npu_api_encode.py │ │ │ ├── test_npu_enable_thinking.py │ │ │ ├── test_npu_matched_stop.py │ │ │ ├── test_npu_openai_function_calling.py │ │ │ ├── test_npu_openai_server_ignore_eos.py │ │ │ └── test_npu_penalty.py │ │ ├── llm_models/ │ │ │ ├── test_npu_afm_4_5b.py │ │ │ ├── test_npu_baichuan2_13b_chat.py │ │ │ ├── test_npu_c4ai_command_r_v01.py │ │ │ ├── test_npu_chatglm2_6b.py │ │ │ ├── test_npu_deepseek_v3_2_exp_w8a8.py │ │ │ ├── test_npu_exaone_3.py │ │ │ ├── test_npu_gemma_3_4b_it_llm.py │ │ │ ├── test_npu_glm4_9b_chat.py │ │ │ ├── test_npu_granite_3_0_3b_a800m.py │ │ │ ├── test_npu_granite_3_1_8b.py │ │ │ ├── test_npu_grok_2.py │ │ │ ├── test_npu_internlm2_7b.py │ │ │ ├── test_npu_ling_lite.py │ │ │ ├── test_npu_llama4_scount_17b_16e.py │ │ │ ├── test_npu_llama_2_7b.py │ │ │ ├── test_npu_mimo_7b_rl.py │ │ │ ├── test_npu_minicpm3_4b.py │ │ │ ├── test_npu_mistral_7b.py │ │ │ ├── test_npu_persimmon_8b_chat.py │ │ │ ├── test_npu_phi_4_multimodal_llm.py │ │ │ ├── test_npu_qwen3_0_6b.py │ │ │ ├── test_npu_qwen3_1_7b_gptq_int8.py │ │ │ ├── test_npu_qwen3_235b_a22b_w8a8.py │ │ │ ├── test_npu_qwen3_30b.py │ │ │ ├── test_npu_qwen3_30b_w4a4.py │ │ │ ├── test_npu_qwen3_32b.py │ │ │ ├── test_npu_qwen3_coder_480b_a35b.py │ │ │ ├── test_npu_qwq_32b_w8a8.py │ │ │ ├── test_npu_smollm_1_7b.py │ │ │ ├── test_npu_stablelm_2_1_6b.py │ │ │ └── tool_chat_template_c4ai_command_r_v01.jinja │ │ ├── rerank_models/ │ │ │ └── test_npu_bge_reranker_v2_m3.py │ │ ├── reward_models/ │ │ │ ├── test_npu_gemma_2_27b_v0_2.py │ │ │ ├── test_npu_internlm2_7b_reward.py │ │ │ └── test_npu_llama_3_1_8b_v0_2.py │ │ ├── test_npu_memory_consumption.py │ │ └── vlm_models/ │ │ ├── mmmu-val.yaml │ │ ├── test_npu_deepseek_vl2.py │ │ ├── test_npu_gemma_3_4b_it.py │ │ ├── test_npu_janus_pro_1b.py │ │ ├── test_npu_janus_pro_7b.py │ │ ├── test_npu_kimi_vl_a3b_instruct.py │ │ ├── test_npu_llama_3_2_11b_vision_instruct.py │ │ ├── test_npu_mimo_vl_7b_rl.py │ │ ├── test_npu_minicpm_o_2_6.py │ │ ├── test_npu_minicpm_v_2_6.py │ │ ├── test_npu_mistral_small_3_1_24b_instruct_2503.py │ │ ├── test_npu_phi4_multimodal_instruct.py │ │ ├── test_npu_qwen2_5_vl_3b_instruct.py │ │ ├── test_npu_qwen2_5_vl_72b_instruct.py │ │ ├── test_npu_qwen3_vl_235b_a22b_instruct.py │ │ ├── test_npu_qwen3_vl_30b_a3b_instruct.py │ │ ├── test_npu_qwen3_vl_4b_instruct.py │ │ └── test_npu_qwen3_vl_8b_instruct.py │ ├── attention/ │ │ ├── test_chunk_gated_delta_rule.py │ │ ├── test_create_kvindices.py │ │ ├── test_fa3.py │ │ ├── test_flash_attention_4.py │ │ ├── test_hybrid_attn_backend.py │ │ ├── test_kda_kernels.py │ │ ├── test_local_attn.py │ │ ├── test_torch_native_attention_backend.py │ │ ├── test_triton_attention_backend.py │ │ ├── test_triton_attention_kernels.py │ │ ├── test_triton_sliding_window.py │ │ └── test_wave_attention_kernels.py │ ├── backends/ │ │ ├── test_deepseek_r1_fp8_trtllm_backend.py │ │ ├── test_deepseek_v3_fp4_cutlass_moe.py │ │ ├── test_flashinfer_trtllm_gen_attn_backend.py │ │ ├── test_flashinfer_trtllm_gen_moe_backend.py │ │ ├── test_qwen3_fp4_trtllm_gen_moe.py │ │ └── test_torch_compile.py │ ├── bench_fn/ │ │ ├── test_bench_serving_functionality.py │ │ └── test_benchmark_datasets_api.py │ ├── constrained_decoding/ │ │ └── test_constrained_decoding.py │ ├── core/ │ │ ├── test_cpp_radix_cache.py │ │ ├── test_deepseek_v3_deterministic.py │ │ ├── test_deterministic.py │ │ ├── test_gpt_oss_1gpu.py │ │ ├── test_gpt_oss_sm120.py │ │ ├── test_hidden_states.py │ │ ├── test_page_size.py │ │ ├── test_qwen3_next_deterministic.py │ │ ├── test_request_queue_validation.py │ │ ├── test_score_api.py │ │ ├── test_srt_endpoint.py │ │ └── test_srt_engine.py │ ├── debug_utils/ │ │ ├── comparator/ │ │ │ ├── __init__.py │ │ │ ├── aligner/ │ │ │ │ ├── __init__.py │ │ │ │ ├── conftest.py │ │ │ │ ├── entrypoint/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── conftest.py │ │ │ │ │ ├── test_executor.py │ │ │ │ │ └── test_planner.py │ │ │ │ ├── reorderer/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── conftest.py │ │ │ │ │ ├── test_executor.py │ │ │ │ │ └── test_planner.py │ │ │ │ ├── test_axis_aligner.py │ │ │ │ ├── token_aligner/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── conftest.py │ │ │ │ │ ├── test_aux_loader.py │ │ │ │ │ ├── test_aux_plugins.py │ │ │ │ │ ├── test_concat_steps.py │ │ │ │ │ ├── test_executor.py │ │ │ │ │ ├── test_planner.py │ │ │ │ │ └── test_thd_seq_lens_loader.py │ │ │ │ └── unsharder/ │ │ │ │ ├── __init__.py │ │ │ │ ├── conftest.py │ │ │ │ ├── test_executor.py │ │ │ │ ├── test_parallel_info.py │ │ │ │ └── test_planner.py │ │ │ ├── conftest.py │ │ │ ├── dims_spec/ │ │ │ │ ├── __init__.py │ │ │ │ ├── test_dim_parser.py │ │ │ │ ├── test_dims_parser.py │ │ │ │ ├── test_tensor_naming.py │ │ │ │ └── test_types.py │ │ │ ├── tensor_comparator/ │ │ │ │ ├── __init__.py │ │ │ │ ├── conftest.py │ │ │ │ ├── test_comparator.py │ │ │ │ ├── test_formatter.py │ │ │ │ └── test_types.py │ │ │ ├── test_bundle_comparator.py │ │ │ ├── test_bundle_matcher.py │ │ │ ├── test_display.py │ │ │ ├── test_dp_utils.py │ │ │ ├── test_dump_loader.py │ │ │ ├── test_entrypoint.py │ │ │ ├── test_log_sink.py │ │ │ ├── test_manually_verify.py │ │ │ ├── test_meta_overrider.py │ │ │ ├── test_model_validation.py │ │ │ ├── test_output_types.py │ │ │ ├── test_per_token_visualizer.py │ │ │ ├── test_preset.py │ │ │ ├── test_utils.py │ │ │ ├── test_visualizer.py │ │ │ └── testing_helpers.py │ │ ├── source_patcher/ │ │ │ ├── conftest.py │ │ │ ├── test_code_patcher.py │ │ │ ├── test_dumper_integration.py │ │ │ └── test_source_editor.py │ │ ├── test_crash_dump.py │ │ ├── test_cuda_coredump_smoke.py │ │ ├── test_dump_comparator.py │ │ ├── test_dump_loader.py │ │ ├── test_dumper.py │ │ ├── test_engine_dumper_comparator_e2e.py │ │ ├── test_schedule_simulator.py │ │ ├── test_soft_watchdog.py │ │ └── test_tensor_dump_forward_hook.py │ ├── disaggregation/ │ │ ├── test_disaggregation_basic.py │ │ ├── test_disaggregation_decode_offload.py │ │ └── test_specv2_kvcache_offloading.py │ ├── distributed/ │ │ ├── test_data_parallelism.py │ │ ├── test_disaggregation_aarch64.py │ │ ├── test_disaggregation_different_tp.py │ │ ├── test_disaggregation_dp_attention.py │ │ ├── test_disaggregation_hybrid_attention.py │ │ ├── test_disaggregation_pp.py │ │ ├── test_dp_attention.py │ │ ├── test_dp_attention_large.py │ │ ├── test_epd_disaggregation.py │ │ ├── test_load_weights_from_remote_instance.py │ │ ├── test_load_weights_from_remote_instance_npu.py │ │ ├── test_parallel_state.py │ │ └── test_pp_single_node.py │ ├── dllm/ │ │ ├── test_llada2_mini.py │ │ └── test_llada2_mini_amd.py │ ├── embedding/ │ │ ├── test_embedding_models.py │ │ ├── test_encoder_embedding_models.py │ │ ├── test_input_embeddings.py │ │ ├── test_input_embeds_chunked.py │ │ └── test_openai_embedding.py │ ├── ep/ │ │ ├── test_deepep_large.py │ │ ├── test_deepep_small.py │ │ └── test_mooncake_ep_small.py │ ├── eval/ │ │ ├── test_eval_accuracy_large.py │ │ ├── test_moe_eval_accuracy_large.py │ │ ├── test_text_models_gsm8k_eval.py │ │ └── test_vlms_mmmu_eval.py │ ├── function_call/ │ │ └── test_kimik2_detector.py │ ├── hicache/ │ │ ├── test_hicache_storage.py │ │ ├── test_hicache_storage_3fs_backend.py │ │ ├── test_hicache_storage_file_backend.py │ │ ├── test_hicache_storage_mooncake_backend.py │ │ ├── test_hicache_storage_runtime_attach_detach.py │ │ └── test_hicache_variants.py │ ├── kernels/ │ │ ├── test_fp4_moe.py │ │ ├── test_fused_topk_deepseek.py │ │ └── test_nsa_indexer.py │ ├── layers/ │ │ ├── mamba/ │ │ │ ├── conftest.py │ │ │ ├── test_causal_conv1d.py │ │ │ ├── test_mamba2_mixer.py │ │ │ ├── test_mamba_ssm.py │ │ │ └── test_mamba_ssm_ssd.py │ │ └── test_fla_layernorm_guard.py │ ├── lora/ │ │ ├── test_chunked_sgmv_backend.py │ │ ├── test_embedding_lora_support.py │ │ ├── test_fused_moe_lora_kernel.py │ │ ├── test_lora_backend.py │ │ ├── test_lora_eviction.py │ │ ├── test_lora_eviction_policy.py │ │ ├── test_lora_hf_sgl_logprob_diff.py │ │ ├── test_lora_openai_api.py │ │ ├── test_lora_openai_compatible.py │ │ ├── test_lora_overlap_loading.py │ │ ├── test_lora_qwen3.py │ │ ├── test_lora_radix_cache.py │ │ ├── test_lora_tied_lm_head.py │ │ ├── test_lora_tp.py │ │ ├── test_lora_update.py │ │ └── test_multi_lora_backend.py │ ├── metrics/ │ │ ├── test_metrics.py │ │ └── test_priority_metrics.py │ ├── mla/ │ │ ├── test_flashmla.py │ │ ├── test_mla.py │ │ ├── test_mla_deepseek_v3.py │ │ ├── test_mla_flashinfer.py │ │ ├── test_mla_fp8.py │ │ └── test_mla_int8_deepseek_v3.py │ ├── model_loading/ │ │ ├── test_external_models.py │ │ └── test_utils_update_weights.py │ ├── models/ │ │ ├── test_compressed_tensors_models.py │ │ ├── test_cross_encoder_models.py │ │ ├── test_dummy_grok_models.py │ │ ├── test_generation_models.py │ │ ├── test_gpt_oss_models_pcg.py │ │ ├── test_kimi_linear_models.py │ │ ├── test_kimi_linear_models_pcg.py │ │ ├── test_ministral3_models.py │ │ ├── test_nvidia_nemotron_3_nano.py │ │ ├── test_nvidia_nemotron_nano_v2.py │ │ ├── test_nvidia_nemotron_nano_v2_vl.py │ │ ├── test_qwen3_next_models_fp4.py │ │ ├── test_qwen3_next_models_pcg.py │ │ ├── test_qwen_models.py │ │ ├── test_reward_models.py │ │ ├── test_transformers_models.py │ │ └── test_vlm_models.py │ ├── moe/ │ │ ├── test_cutedsl_moe.py │ │ ├── test_fused_moe.py │ │ ├── test_glm4_moe_models.py │ │ ├── test_moe_ep.py │ │ ├── test_torch_compile_moe.py │ │ ├── test_triton_fused_moe.py │ │ └── test_triton_moe_channel_fp8_kernel.py │ ├── openai_server/ │ │ ├── basic/ │ │ │ ├── test_anthropic_server.py │ │ │ ├── test_openai_server.py │ │ │ ├── test_protocol.py │ │ │ ├── test_serving_chat.py │ │ │ ├── test_serving_completions.py │ │ │ └── test_serving_rerank.py │ │ ├── features/ │ │ │ ├── test_enable_thinking.py │ │ │ ├── test_json_mode.py │ │ │ ├── test_openai_server_ebnf.py │ │ │ ├── test_openai_server_hidden_states.py │ │ │ └── test_reasoning_content.py │ │ ├── function_call/ │ │ │ ├── test_anthropic_tool_use.py │ │ │ ├── test_openai_function_calling.py │ │ │ └── test_tool_choice.py │ │ └── validation/ │ │ ├── test_large_max_new_tokens.py │ │ ├── test_matched_stop.py │ │ ├── test_openai_server_ignore_eos.py │ │ └── test_request_length_validation.py │ ├── ops/ │ │ ├── test_aiter_allreduce_fusion_amd.py │ │ └── test_repeat_interleave.py │ ├── perf/ │ │ ├── test_bench_one_batch_1gpu.py │ │ ├── test_bench_one_batch_2gpu.py │ │ ├── test_bench_serving_1gpu_large.py │ │ ├── test_bench_serving_1gpu_part1.py │ │ ├── test_bench_serving_1gpu_part2.py │ │ ├── test_bench_serving_2gpu.py │ │ ├── test_dpsk_r1_fp4_4gpu_perf.py │ │ ├── test_gpt_oss_4gpu_perf.py │ │ ├── test_text_models_perf.py │ │ ├── test_vlm_perf_5090.py │ │ └── test_vlms_perf.py │ ├── piecewise_cuda_graph/ │ │ └── test_piecewise_cuda_graph_support_1_gpu.py │ ├── profiling/ │ │ ├── test_profile_v2.py │ │ └── test_start_profile.py │ ├── quant/ │ │ ├── test_autoround.py │ │ ├── test_awq.py │ │ ├── test_awq_dequant.py │ │ ├── test_block_int8.py │ │ ├── test_bnb.py │ │ ├── test_deepseek_v32_fp4_4gpu.py │ │ ├── test_deepseek_v32_fp4_mtp_4gpu.py │ │ ├── test_deepseek_v3_fp4_4gpu.py │ │ ├── test_eval_fp8_accuracy.py │ │ ├── test_fp8_blockwise_gemm.py │ │ ├── test_fp8_kernel.py │ │ ├── test_fp8_utils.py │ │ ├── test_fp8kv_triton.py │ │ ├── test_fused_rms_fp8_group_quant.py │ │ ├── test_gguf.py │ │ ├── test_gptqmodel_dynamic.py │ │ ├── test_int4fp8_moe.py │ │ ├── test_int8_kernel.py │ │ ├── test_marlin_moe.py │ │ ├── test_modelopt_fp8.py │ │ ├── test_nvfp4_gemm.py │ │ ├── test_quant_config_parsing.py │ │ ├── test_quantization.py │ │ ├── test_torchao.py │ │ ├── test_triton_scaled_mm.py │ │ ├── test_w4a8_deepseek_v3.py │ │ └── test_w8a8_quantization.py │ ├── radix_cache/ │ │ ├── test_radix_attention.py │ │ ├── test_radix_cache_hit.py │ │ └── test_swa_radix_cache_kl.py │ ├── rl/ │ │ ├── test_fp32_lm_head.py │ │ ├── test_lora_load_from_tensor.py │ │ ├── test_multi_instance_release_memory_occupation.py │ │ ├── test_patch_torch.py │ │ ├── test_release_memory_occupation.py │ │ ├── test_return_routed_experts.py │ │ ├── test_update_weights_from_disk.py │ │ ├── test_update_weights_from_distributed.py │ │ └── test_update_weights_from_tensor.py │ ├── rotary/ │ │ ├── test_mrope.py │ │ └── test_rope_rocm.py │ ├── sampling/ │ │ ├── test_original_logprobs.py │ │ ├── test_penalty.py │ │ └── test_pytorch_sampling_backend.py │ ├── scheduler/ │ │ ├── test_abort.py │ │ ├── test_chunked_prefill.py │ │ ├── test_no_chunked_prefill.py │ │ ├── test_no_overlap_scheduler.py │ │ ├── test_prefill_delayer.py │ │ ├── test_priority_scheduling.py │ │ ├── test_retract_decode.py │ │ └── test_routing_key_scheduling.py │ ├── sessions/ │ │ ├── test_session_control.py │ │ ├── test_session_latency.py │ │ └── test_streaming_session.py │ ├── spec/ │ │ ├── eagle/ │ │ │ ├── test_deepseek_v3_fp4_mtp_small.py │ │ │ ├── test_eagle3_basic.py │ │ │ ├── test_eagle_constrained_decoding.py │ │ │ ├── test_eagle_dp_attention.py │ │ │ ├── test_eagle_infer_a.py │ │ │ ├── test_eagle_infer_b.py │ │ │ ├── test_eagle_infer_beta.py │ │ │ ├── test_eagle_infer_beta_dp_attention.py │ │ │ └── test_eagle_infer_beta_dp_attention_large.py │ │ ├── test_constrained_decoding_spec_reasoning.py │ │ ├── test_ngram_speculative_decoding.py │ │ ├── test_standalone_speculative_decoding.py │ │ └── utils/ │ │ └── test_build_eagle_tree.py │ ├── stress/ │ │ ├── test_stress_deepseek_v3.py │ │ ├── test_stress_glm_4_6.py │ │ ├── test_stress_kimi_k2.py │ │ └── test_stress_qwen3_235b.py │ ├── test_hybrid_dp_ep_tp_mtp.py │ ├── test_srt_backend.py │ ├── tokenizer/ │ │ ├── test_multi_tokenizer.py │ │ └── test_skip_tokenizer_init.py │ ├── unit/ │ │ ├── README.md │ │ ├── batch_invariant_ops/ │ │ │ └── test_batch_invariant_ops.py │ │ ├── entrypoints/ │ │ │ ├── openai/ │ │ │ │ └── test_serving_embedding.py │ │ │ └── test_ssl_cert_refresher.py │ │ ├── function_call/ │ │ │ ├── test_function_call_parser.py │ │ │ ├── test_glm47_moe_detector.py │ │ │ ├── test_json_schema_constraint.py │ │ │ ├── test_parallel_tool_calls.py │ │ │ └── test_unknown_tool_name.py │ │ ├── layers/ │ │ │ ├── test_conv_layer.py │ │ │ └── test_mamba_state_scatter_triton.py │ │ ├── managers/ │ │ │ ├── test_io_struct.py │ │ │ ├── test_prefill_adder.py │ │ │ └── test_profile_merger_http_api.py │ │ ├── mem_cache/ │ │ │ ├── test_evict_policy.py │ │ │ ├── test_mamba_unittest.py │ │ │ ├── test_nsa_pool_host_unit.py │ │ │ ├── test_radix_cache_slru_accuracy.py │ │ │ ├── test_radix_cache_unit.py │ │ │ └── test_swa_unittest.py │ │ ├── model_executor/ │ │ │ └── test_model_hooks.py │ │ ├── model_loader/ │ │ │ ├── test_modelopt_export.py │ │ │ └── test_modelopt_loader.py │ │ ├── observability/ │ │ │ ├── test_cpu_monitor.py │ │ │ └── test_metrics_utils.py │ │ ├── parser/ │ │ │ ├── test_harmony_parser.py │ │ │ ├── test_jinja_template_utils.py │ │ │ └── test_reasoning_parser.py │ │ ├── server_args/ │ │ │ └── test_server_args.py │ │ └── utils/ │ │ ├── test_gauge_histogram.py │ │ ├── test_json_response.py │ │ ├── test_patch_tokenizer.py │ │ └── test_profile_merger.py │ ├── utils/ │ │ ├── test_bench_typebaseddispatcher.py │ │ ├── test_log_utils.py │ │ ├── test_model_file_verifier.py │ │ ├── test_network_address.py │ │ ├── test_request_logger.py │ │ ├── test_scheduler_status_logger.py │ │ ├── test_socket_utils.py │ │ └── test_type_based_dispatcher.py │ └── vlm/ │ ├── test_encoder_dp.py │ ├── test_evs.py │ ├── test_patch_embed_perf.py │ ├── test_video_utils.py │ ├── test_vision_chunked_prefill.py │ ├── test_vision_openai_server_a.py │ └── test_vlm_input_format.py ├── run_suite.py ├── run_suite_nightly.py ├── show_partitions.py └── srt/ ├── ascend/ │ ├── test_ascend_autoround_dense.py │ ├── test_ascend_autoround_moe.py │ ├── test_ascend_compile_graph_tp1_bf16.py │ ├── test_ascend_deepep.py │ ├── test_ascend_gptq_moe.py │ ├── test_ascend_graph_tp1_bf16.py │ ├── test_ascend_graph_tp2_bf16.py │ ├── test_ascend_hicache_mha.py │ ├── test_ascend_hicache_mla.py │ ├── test_ascend_mla_fia_w8a8int8.py │ ├── test_ascend_mla_w8a8int8.py │ ├── test_ascend_piecewise_graph_prefill.py │ ├── test_ascend_sampling_backend.py │ ├── test_ascend_tp1_bf16.py │ ├── test_ascend_tp2_bf16.py │ ├── test_ascend_tp2_fia_bf16.py │ ├── test_ascend_tp4_bf16.py │ ├── test_ascend_w4a4_quantization.py │ ├── test_ascend_w8a8_quantization.py │ └── test_llada2_mini_ascend.py ├── configs/ │ ├── deepseek_v3.yaml │ ├── deepseek_v3_long_context.yaml │ ├── llama_405b.yaml │ ├── random_config.yaml │ ├── random_flashinfer_vs_triton_config.yaml │ └── sharegpt_config.yaml ├── cpu/ │ ├── test_activation.py │ ├── test_binding.py │ ├── test_bmm.py │ ├── test_causal_conv1d.py │ ├── test_cpu_graph.py │ ├── test_decode.py │ ├── test_extend.py │ ├── test_flash_attn.py │ ├── test_gemm.py │ ├── test_intel_amx_attention_backend_a.py │ ├── test_intel_amx_attention_backend_b.py │ ├── test_intel_amx_attention_backend_c.py │ ├── test_mamba.py │ ├── test_mla.py │ ├── test_moe.py │ ├── test_norm.py │ ├── test_qkv_proj_with_rope.py │ ├── test_qwen3.py │ ├── test_rope.py │ ├── test_shared_expert.py │ ├── test_topk.py │ └── utils.py ├── double-sparsity-config-Llama-3.1-8B-Instruct.json ├── experiment_runner.py ├── kv_cache_scales_llama3_1_8b.json ├── kv_cache_scales_llama3_8b.json ├── kv_cache_scales_qwen2_1_5b.json ├── models/ │ └── compare.py ├── parse_results.py ├── run_suite.py ├── test_embed_interpolate_unittest.py └── xpu/ ├── test_deepseek_ocr.py └── test_intel_xpu_backend.py ================================================ FILE CONTENTS ================================================ ================================================ FILE: .claude/skills/add-jit-kernel/SKILL.md ================================================ --- name: add-jit-kernel description: Step-by-step tutorial for adding a new lightweight JIT CUDA kernel to sglang's jit_kernel module --- # Tutorial: Adding a New JIT Kernel to SGLang This tutorial walks through adding a simple element-wise scale operation as a JIT kernel. We'll implement `scale(x, factor) = x * factor` to demonstrate the complete workflow. ## Goal Add a new operation that scales each element of a tensor by a scalar factor: - Input: tensor `x` (CUDA) and scalar `factor` (float, passed at runtime) - Output: `x * factor` (element-wise), allocated internally - Supported dtypes: **FP16 (`torch.float16`), BF16 (`torch.bfloat16`), FP32 (`torch.float32`)** ## When to use JIT vs AOT (`sgl-kernel`) - **JIT (`jit_kernel`)**: prefer this first for kernels that do **not** depend on CUTLASS or another large C++ project. It is the default choice for lightweight kernels that benefit from rapid iteration and first-use compilation. - **AOT (`sgl-kernel`)**: prefer this when the kernel **does** depend on CUTLASS or another large C++ project, or when it should live in `sgl-kernel/` and participate in the wheel build / torch op registration flow. - **Exception**: kernels that depend on `flashinfer`, or on CUTLASS that is already provided through `flashinfer`, can still be implemented as `jit_kernel`. --- ## Common Abstractions in `python/sglang/jit_kernel/include/sgl_kernel/` **Always prefer these abstractions over raw CUDA primitives.** They provide safety, readability, and consistency with the rest of the codebase. **Important include rule:** for every `#include ` line, add a short trailing comment explaining why that header is included (for example `// For TensorMatcher, SymbolicSize, SymbolicDevice`). This matches the current JIT kernel style and keeps include usage self-documenting. ### `utils.h` — Host-side utilities ```cpp #include ``` - **`host::RuntimeCheck(cond, args...)`** — Assert a condition at runtime; throws `PanicError` with file/line info on failure. Prefer this over bare `assert`. - **`host::Panic(args...)`** — Unconditionally throw a `PanicError` with a descriptive message. - **`host::div_ceil(a, b)`** — Integer ceiling division `(a + b - 1) / b`. - **`host::irange(n)`** / **`host::irange(start, end)`** — Range views for cleaner loops. - **`host::pointer::offset(ptr, offsets...)`** — Byte-safe pointer arithmetic on `void*`. Use this instead of raw casts. ### `utils.cuh` — Device-side utilities + `LaunchKernel` ```cpp #include ``` - **Type aliases**: `fp16_t`, `bf16_t`, `fp32_t`, `fp8_e4m3_t`, `fp8_e5m2_t` and their packed variants `fp16x2_t`, `bf16x2_t`, `fp32x2_t`, etc. - **`SGL_DEVICE`** — Expands to `__forceinline__ __device__`. Use on all device functions. - **`device::kWarpThreads`** — Constant `32`. - **`device::load_as(ptr, offset)`** / **`device::store_as(ptr, val, offset)`** — Type-safe loads/stores from `void*`. - **`device::pointer::offset(ptr, offsets...)`** — Pointer arithmetic on device. - **`host::LaunchKernel(grid, block, device_or_stream [, smem])`** — RAII kernel launcher that: - Resolves the CUDA stream from a `DLDevice` via TVM-FFI automatically. - Checks the CUDA error with file/line info after launch via `operator()(kernel, args...)`. - Supports `.enable_pdl(bool)` for PDL (Programmatic Dependent Launch, SM90+). - **`host::RuntimeDeviceCheck(cudaError_t)`** — Check a CUDA error; throw on failure. ### `tensor.h` — Tensor validation (`TensorMatcher`, Symbolic types) ```cpp #include ``` This is the **primary validation API** for all kernel launchers. Use it to validate every `tvm::ffi::TensorView` argument. - **`host::SymbolicSize{"name"}`** — A named symbolic dimension. Call `.set_value(n)` to pin it, `.unwrap()` to extract after verification. - **`host::SymbolicDType`** — Symbolic dtype. Use `.set_options()` to restrict allowed types. - **`host::SymbolicDevice`** — Symbolic device. Use `.set_options()` to restrict to CUDA. - **`host::TensorMatcher({dims...})`** — Fluent builder for tensor validation: - `.with_dtype()` — require a specific C++ type (e.g. `fp16_t`) - `.with_dtype()` — allow a set of types - `.with_device(device_sym)` — require CUDA and bind the checked device to a `SymbolicDevice` - `.with_strides({strides...})` — validate strides (omit to require contiguous) - `.verify(tensor_view)` — execute the check; throws `PanicError` with full context on failure; **chainable** (`verify(a).verify(b)` to check multiple tensors with the same shape) **Typical pattern:** ```cpp auto N = SymbolicSize{"num_elements"}; auto device = SymbolicDevice{}; device.set_options(); TensorMatcher({N}) // .with_dtype() .with_device(device) .verify(dst) .verify(src); // same shape, dtype, device as dst const size_t n = N.unwrap(); const DLDevice dev = device.unwrap(); ``` ### `type.cuh` — `dtype_trait` and `packed_t` ```cpp #include ``` - **`dtype_trait`** — Static trait struct for each scalar type. Provides: - `dtype_trait::from(value)` — convert from another type (e.g. `fp32_t` → `fp16_t`) - `dtype_trait::abs/sqrt/rsqrt/exp/sin/cos(x)` — type-dispatched unary math (primarily for `fp32_t`) - `dtype_trait::max/min(x, y)` — type-dispatched binary math (primarily for `fp32_t`) - **`packed_t`** — Two-element packed alias: `packed_t` = `fp16x2_t`, `packed_t` = `bf16x2_t`, `packed_t` = `fp32x2_t`. Use for vectorized loads/stores. - **`device::cast(value)`** — Type-safe cast using `dtype_trait`, e.g. `cast(v)`. ### `vec.cuh` — Vectorized memory access (`AlignedVector`) ```cpp #include ``` - **`device::AlignedVector`** — Aligned storage for N elements of type T. N must be a power of two, `sizeof(T)*N <= 32`. Enables vectorized loads/stores for bandwidth efficiency. In terms of API/codegen constraints, the upper bound is 256-bit; in practice, 128-bit is the portable default, while 256-bit vectorization is typically only viable on `SM100+` and should be gated by an architecture check when needed. - `.load(ptr, offset)` — vectorized load from `ptr[offset]` - `.store(ptr, offset)` — vectorized store to `ptr[offset]` - `.fill(value)` — fill all lanes - `operator[](i)` — element access ### `tile.cuh` — `tile::Memory` (strided memory access pattern) ```cpp #include ``` - `tile::Memory` is fundamentally a **1D cooperative accessor** over a contiguous region. - **`device::tile::Memory::cta(blockDim.x)`** — Creates a tile accessor where each thread handles `tid = threadIdx.x` with stride `tsize` (for `cta(blockDim.x)`, this is `blockDim.x`). Common for loops over a 1D array. - **`.load(ptr, offset)`** — loads `ptr[tid + offset * tsize]` - **`.store(ptr, val, offset)`** — stores to `ptr[tid + offset * tsize]` - **`.in_bound(n, offset)`** — boundary check For a **2D tile**, either flatten `(row, col)` into a linear tile index first, or compute the address manually with `ptr[row * stride + col]` using your thread/block coordinates. ### `math.cuh` — Device math (`device::math::`) ```cpp #include ``` - `device::math::max/min(a, b)` — type-dispatched binary math via `dtype_trait` - `device::math::abs/sqrt/rsqrt/exp/sin/cos(x)` — type-dispatched unary math via `dtype_trait` ### `warp.cuh` — Warp-level primitives ```cpp #include ``` - `device::warp::reduce_sum(value)` — warp-level sum reduction via `__shfl_xor_sync` - `device::warp::reduce_max(value)` — warp-level max reduction ### `cta.cuh` — CTA-level primitives ```cpp #include ``` - `device::cta::reduce_max(value, smem, min_value)` — CTA-wide max using shared memory + warp reduction. Caller is responsible for a `__syncthreads()` after if the result in `smem[0]` is needed. ### `atomic.cuh` — Atomic operations ```cpp #include ``` - `device::atomic::max(float* addr, float value)` — float atomic max (handles negative values correctly via bit tricks). ### `runtime.cuh` — Occupancy and device info ```cpp #include ``` - `host::runtime::get_blocks_per_sm(kernel, block_dim)` — max active blocks per SM (occupancy) - `host::runtime::get_sm_count(device_id)` — number of SMs on the device - `host::runtime::get_cc_major(device_id)` — compute capability major version **Persistent kernel pattern** (cap blocks to SM count × occupancy): ```cpp static const uint32_t max_occ = runtime::get_blocks_per_sm(kernel, kBlockSize); static const uint32_t num_sm = runtime::get_sm_count(device.unwrap().device_id); const auto num_blocks = std::min(num_sm * max_occ, div_ceil(n, kBlockSize)); LaunchKernel(num_blocks, kBlockSize, device.unwrap())(kernel, params); ``` --- ## Step 0 (optional): Generate a `.clangd` config for better IDE support ```bash python -m sglang.jit_kernel ``` --- ## Step 1: Implement the CUDA kernel in `jit_kernel/csrc/` Create `python/sglang/jit_kernel/csrc/elementwise/scale.cuh`. The implementation fully uses the project abstractions described above: ```cpp #include // For TensorMatcher, SymbolicSize, SymbolicDevice #include // For dtype_trait, fp16_t, bf16_t, fp32_t #include // For RuntimeCheck, div_ceil #include // For LaunchKernel, SGL_DEVICE #include // For AlignedVector #include #include namespace { // ---------------------------------------------------------------- // Kernel: element-wise scale using vectorized 128-bit loads/stores // T = fp16_t | bf16_t | fp32_t // kVecN = number of elements per vector load (e.g. 8 for fp16) // factor = runtime scale factor // ---------------------------------------------------------------- template __global__ void scale_kernel(T* __restrict__ dst, const T* __restrict__ src, float factor, uint32_t n_total) { using vec_t = device::AlignedVector; const uint32_t n_vecs = n_total / kVecN; // --- vectorised body --- const uint32_t vec_stride = blockDim.x * gridDim.x; for (uint32_t vi = blockIdx.x * blockDim.x + threadIdx.x; vi < n_vecs; vi += vec_stride) { vec_t v; v.load(src, vi); #pragma unroll for (int i = 0; i < kVecN; ++i) { v[i] = static_cast(static_cast(v[i]) * factor); } v.store(dst, vi); } // --- scalar tail --- const uint32_t base = n_vecs * kVecN; const uint32_t scalar_stride = blockDim.x * gridDim.x; for (uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; base + i < n_total; i += scalar_stride) { dst[base + i] = static_cast(static_cast(src[base + i]) * factor); } } // ---------------------------------------------------------------- // Launcher: validates tensors, selects vector width, launches kernel // ---------------------------------------------------------------- template void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) { using namespace host; // 1. Validate input tensors with TensorMatcher SymbolicSize N = {"num_elements"}; SymbolicDevice device_; device_.set_options(); TensorMatcher({N}) // .with_dtype() .with_device(device_) .verify(dst) .verify(src); // same shape / dtype / device as dst const uint32_t n = static_cast(N.unwrap()); const DLDevice device = device_.unwrap(); RuntimeCheck(n > 0, "scale: num_elements must be > 0, got ", n); // 2. Choose vector width for 128-bit loads (16 bytes) // fp16/bf16: 8 elements × 2 bytes = 16 bytes // fp32: 4 elements × 4 bytes = 16 bytes constexpr int kVecN = 16 / sizeof(T); const uint32_t n_work_items = div_ceil(n, static_cast(kVecN)); // 3. Launch constexpr uint32_t kBlockSize = 256; const uint32_t grid = div_ceil(n_work_items, kBlockSize); LaunchKernel(grid, kBlockSize, device)( scale_kernel, static_cast(dst.data_ptr()), static_cast(src.data_ptr()), factor, n); } } // namespace ``` **Key points:** - Include headers from `sgl_kernel/` — **not** raw CUDA headers for anything already covered - Add a short trailing `// For ...` explanation to every `#include ` line - Use `TensorMatcher` for all tensor validation; never manually check shape/dtype/device - Use `AlignedVector` for vectorised 128-bit loads/stores — significant bandwidth win - Use `LaunchKernel` — it resolves the stream and checks errors automatically - Use `RuntimeCheck` for runtime assertions with useful error messages - Prefer passing runtime scalars like `factor` directly unless compile-time specialisation is genuinely required - `fp16_t` / `bf16_t` / `fp32_t` are the project's type aliases (from `utils.cuh`) - `device::cast` or `dtype_trait::from(val)` for cross-type conversions - `device::math::` functions for device math instead of bare `__` intrinsics --- ## Step 2: Add the Python wrapper in `jit_kernel/` Create `python/sglang/jit_kernel/scale.py`: ```python from __future__ import annotations from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_scale_module(dtype: torch.dtype) -> Module: """Compile and cache the JIT scale module for a given dtype.""" args = make_cpp_args(dtype) return load_jit( "scale", *args, cuda_files=["elementwise/scale.cuh"], cuda_wrappers=[("scale", f"scale<{args}>")], ) def scale(src: torch.Tensor, factor: float, out: torch.Tensor | None = None) -> torch.Tensor: """ Element-wise scale: dst = src * factor. Supported dtypes: torch.float16, torch.bfloat16, torch.float32. Parameters ---------- src : CUDA tensor (FP16 / BF16 / FP32) factor : scale factor out : optional pre-allocated output tensor (same shape/dtype as src) Returns ------- Scaled tensor (dst = src * factor). """ if not src.is_cuda: raise RuntimeError("src must be a CUDA tensor") if src.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise RuntimeError( f"Unsupported dtype {src.dtype}. Supported: float16, bfloat16, float32" ) if out is None: out = torch.empty_like(src) else: if out.shape != src.shape: raise RuntimeError("out shape must match src") if out.dtype != src.dtype: raise RuntimeError("out dtype must match src") if out.device != src.device: raise RuntimeError("out device must match src") # Keep the Python wrapper thin, but still enforce the basic preconditions # that the current JIT/FFI path does not reject safely on its own. module = _jit_scale_module(src.dtype) module.scale(out, src, factor) return out ``` **Key points:** - Use `cache_once` — **not** `functools.lru_cache` (incompatible with `torch.compile`) - `load_jit` first arg(s) form the unique build marker; same marker = same cached binary - Only include compile-time specialisation knobs in the build marker; runtime values like `factor` should stay runtime unless the kernel truly needs templating - `cuda_wrappers`: `(export_name, kernel_symbol)` — `export_name` is called from Python - `make_cpp_args(dtype, ...)` converts `torch.dtype` to C++ type alias: - Keep Python launchers thin, but still validate the basic invariants (`is_cuda`, supported dtype, `out` metadata). In the current JIT/FFI path, invalid tensors are not always rejected safely before launch | `torch.dtype` | C++ type | |--------------------|------------| | `torch.float16` | `fp16_t` | | `torch.bfloat16` | `bf16_t` | | `torch.float32` | `fp32_t` | --- ## Step 3 (optional): Tune JIT build flags ```python return load_jit( "scale", *args, cuda_files=["elementwise/scale.cuh"], cuda_wrappers=[("scale", f"scale<{args}>")], extra_cuda_cflags=["-O3", "--use_fast_math"], ) ``` If your kernel requires SM90+, raise a clear Python error before calling `load_jit`: ```python if torch.cuda.get_device_capability()[0] < 9: raise RuntimeError("This kernel requires SM90 (Hopper) or later") ``` --- ## Step 4: Write tests (required) Create `python/sglang/jit_kernel/tests/test_scale.py`: ```python import pytest import torch from sglang.jit_kernel.scale import scale @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @pytest.mark.parametrize("size", [1, 127, 128, 1024, 4097]) # cover tail remainder @pytest.mark.parametrize("factor", [0.5, 1.0, 2.0, 3.0]) def test_scale_correctness(dtype, size, factor): src = torch.randn(size, dtype=dtype, device="cuda") out = scale(src, factor) expected = src * factor rtol, atol = (1e-5, 1e-6) if dtype == torch.float32 else (1e-2, 1e-2) torch.testing.assert_close(out, expected, rtol=rtol, atol=atol) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) def test_scale_out_param(dtype): src = torch.randn(1024, dtype=dtype, device="cuda") out = torch.empty_like(src) result = scale(src, 2.0, out=out) assert result is out torch.testing.assert_close(out, src * 2.0, rtol=1e-2, atol=1e-2) def test_scale_cpu_error(): src = torch.randn(128, dtype=torch.float16) # CPU tensor with pytest.raises(RuntimeError, match="CUDA"): scale(src, 2.0) def test_scale_unsupported_dtype(): src = torch.randint(0, 10, (128,), dtype=torch.int32, device="cuda") with pytest.raises(RuntimeError, match="dtype"): scale(src, 2.0) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ``` --- ## Step 5: Add a benchmark (required) Create `python/sglang/jit_kernel/benchmark/bench_scale.py`: ```python import itertools import torch import triton import triton.testing from sglang.jit_kernel.benchmark.utils import ( DEFAULT_DEVICE, DEFAULT_DTYPE, get_benchmark_range, run_benchmark, ) from sglang.jit_kernel.scale import scale as jit_scale SIZE_LIST = get_benchmark_range( full_range=[2**n for n in range(10, 20)], # 1K … 512K elements ci_range=[4096, 65536], ) configs = list(itertools.product(SIZE_LIST)) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["size"], x_vals=configs, line_arg="provider", line_vals=["jit", "torch"], line_names=["SGL JIT Kernel", "PyTorch"], styles=[("blue", "-"), ("red", "--")], ylabel="us", plot_name="scale-performance", args={}, ) ) def benchmark(size: int, provider: str): src = torch.randn(size, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) factor = 2.0 if provider == "jit": fn = lambda: jit_scale(src, factor) else: fn = lambda: src * factor return run_benchmark(fn) if __name__ == "__main__": benchmark.run(print_data=True) ``` Run: ```bash python python/sglang/jit_kernel/benchmark/bench_scale.py ``` --- ## Troubleshooting - **JIT compilation fails**: ensure the `.cuh` file is under `python/sglang/jit_kernel/csrc/`; reduce template argument combinations - **CUDA crash / illegal memory access**: `CUDA_LAUNCH_BLOCKING=1`; `compute-sanitizer --tool memcheck python ...` - **Unstable benchmark results**: `run_benchmark` uses CUDA-graph-based timing by default --- ## References - `docs/developer_guide/development_jit_kernel_guide.md` - `python/sglang/jit_kernel/utils.py` — `cache_once`, `load_jit`, `make_cpp_args` - `python/sglang/jit_kernel/include/sgl_kernel/tensor.h` — `TensorMatcher`, `SymbolicSize/DType/Device` - `python/sglang/jit_kernel/include/sgl_kernel/utils.cuh` — type aliases, `LaunchKernel`, `SGL_DEVICE` - `python/sglang/jit_kernel/include/sgl_kernel/vec.cuh` — `AlignedVector` - `python/sglang/jit_kernel/include/sgl_kernel/tile.cuh` — `tile::Memory` - `python/sglang/jit_kernel/include/sgl_kernel/type.cuh` — `dtype_trait`, `packed_t`, `device::cast` - `python/sglang/jit_kernel/include/sgl_kernel/math.cuh` — `device::math::` - `python/sglang/jit_kernel/include/sgl_kernel/warp.cuh` — `warp::reduce_sum/max` - `python/sglang/jit_kernel/include/sgl_kernel/cta.cuh` — `cta::reduce_max` - `python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh` — `atomic::max` - `python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh` — occupancy / SM count helpers - `python/sglang/jit_kernel/csrc/add_constant.cuh` — minimal runnable reference - `python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh` — real example using `TensorMatcher` + `LaunchKernel` + `tile::Memory` - `python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh` — real example using `runtime::get_blocks_per_sm` + persistent kernel pattern - `python/sglang/jit_kernel/benchmark/utils.py` — benchmark helpers ## Summary of Files Created ``` python/sglang/jit_kernel/csrc/elementwise/scale.cuh # NEW: CUDA kernel python/sglang/jit_kernel/scale.py # NEW: Python wrapper python/sglang/jit_kernel/tests/test_scale.py # NEW: Tests python/sglang/jit_kernel/benchmark/bench_scale.py # NEW: Benchmark ``` ================================================ FILE: .claude/skills/add-sgl-kernel/SKILL.md ================================================ --- name: add-sgl-kernel description: Step-by-step tutorial for adding a heavyweight AOT CUDA/C++ kernel to sgl-kernel (including tests & benchmarks) --- # Tutorial: Adding a New Kernel to `sgl-kernel` (AOT / Heavyweight) This tutorial walks through adding a simple element-wise scale operation as an AOT kernel. We'll implement `scale(x, factor) = x * factor` to demonstrate the complete workflow. ## Goal Add a new operation that scales each element of a tensor by a scalar factor: - Input: tensor `x` (CUDA) and scalar `factor` (float) - Output: `x * factor` (element-wise, in-place or into pre-allocated `out`) - Supported dtypes: **FP16 (`torch.float16`), BF16 (`torch.bfloat16`), FP32 (`torch.float32`)** - Dispatched via `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16` macro (defined in `sgl-kernel/include/utils.h`) ## Two rules of thumb (must follow) 1. **Prefer `python/sglang/jit_kernel` first** when the kernel does **not** depend on CUTLASS or another large C++ project. This is the default path for lightweight kernels that benefit from rapid iteration. 2. **Prefer `sgl-kernel`** when the kernel **does** depend on CUTLASS or another large C++ project, or when it should be part of the AOT wheel / torch op registration flow. 3. **Exception**: if the dependency is `flashinfer`, or CUTLASS that is already provided through `flashinfer`, the kernel can still be implemented as `jit_kernel`. In addition, every new kernel must ship with: - **Tests** (pytest) - **A benchmark script** (triton.testing) --- ## Repository integration map You will typically touch these files/areas: - Implementation: `sgl-kernel/csrc/elementwise/scale.cu` (pick the right subdirectory) - Public declarations: `sgl-kernel/include/sgl_kernel_ops.h` - Torch extension registration: `sgl-kernel/csrc/common_extension.cc` - Build: `sgl-kernel/CMakeLists.txt` (`set(SOURCES ...)`) - Python API: `sgl-kernel/python/sgl_kernel/` and `sgl-kernel/python/sgl_kernel/__init__.py` - Tests: `sgl-kernel/tests/test_scale.py` - Benchmarks: `sgl-kernel/benchmark/bench_scale.py` --- ## Step 1: Implement the kernel in `csrc/` Pick the right subdirectory: - `csrc/elementwise/` — for element-wise ops (our example) - `csrc/gemm/`, `csrc/attention/`, `csrc/moe/` — for other categories Create `sgl-kernel/csrc/elementwise/scale.cu`: ```cpp #include #include #include #include "utils.h" // DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16 // scale_kernel: out[i] = input[i] * factor // Supports float, half (__half), __nv_bfloat16 via template T template __global__ void scale_kernel(T* __restrict__ out, const T* __restrict__ input, float factor, int64_t n) { int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (idx < n) { out[idx] = static_cast(static_cast(input[idx]) * factor); } } void scale(at::Tensor& out, const at::Tensor& input, double factor) { TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor"); TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); TORCH_CHECK(out.is_cuda(), "out must be a CUDA tensor"); TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); TORCH_CHECK(out.sizes() == input.sizes(), "out and input must have the same shape"); TORCH_CHECK(out.scalar_type() == input.scalar_type(), "out and input must have the same dtype"); const int64_t n = input.numel(); const int threads = 256; const int blocks = (n + threads - 1) / threads; const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); // Dispatches over float, float16, bfloat16 DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(input.scalar_type(), c_type, [&] { scale_kernel<<>>( static_cast(out.data_ptr()), static_cast(input.data_ptr()), static_cast(factor), n); cudaError_t status = cudaGetLastError(); TORCH_CHECK(status == cudaSuccess, "scale_kernel launch failed: ", cudaGetErrorString(status)); return true; }); } ``` **Key points:** - Use `at::Tensor` (PyTorch tensors), `TORCH_CHECK` for validation, `at::cuda::getCurrentCUDAStream()` for stream - Keep Python wrappers thin; do shape/dtype/device validation in C++ right around the launch path - `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16` covers `float`, `half` (FP16), `__nv_bfloat16` (BF16) - Add device error checking after every kernel launch - If a kernel only works on certain architectures, enforce that with `TORCH_CHECK` and skip logic in tests --- ## Step 2: Add a C++ declaration in `include/sgl_kernel_ops.h` Edit `sgl-kernel/include/sgl_kernel_ops.h`, add to the elementwise section: ```cpp void scale(at::Tensor& out, const at::Tensor& input, double factor); ``` --- ## Step 3: Register the op in `csrc/common_extension.cc` Edit `sgl-kernel/csrc/common_extension.cc`, inside `TORCH_LIBRARY_FRAGMENT(sgl_kernel, m)`: ```cpp // From csrc/elementwise m.def("scale(Tensor! out, Tensor input, float factor) -> ()"); m.impl("scale", torch::kCUDA, &scale); ``` **Key points:** - `Tensor!` means in-place / mutable output argument - The schema is important for `torch.compile` and for consistent call signatures - Keep the torch schema in PyTorch scalar types (`float` here), but note that the C++ launcher signature still needs `double` for scalar arguments accepted by `torch::Library` --- ## Step 4: Add the new source file to `CMakeLists.txt` Edit `sgl-kernel/CMakeLists.txt`, add to `set(SOURCES ...)`: ```cmake csrc/elementwise/scale.cu ``` **Key points:** - Keep the list **alphabetically sorted** (the file explicitly requires this) - If the kernel has arch constraints, reflect that in tests/benchmarks via skip logic --- ## Step 5: Expose a Python API under `sgl-kernel/python/sgl_kernel/` Prefer following the existing module organization first. For elementwise kernels, the usual pattern is: - implement the Python wrapper in `sgl-kernel/python/sgl_kernel/elementwise.py` - then re-export it from `sgl-kernel/python/sgl_kernel/__init__.py` For example, in `sgl-kernel/python/sgl_kernel/elementwise.py`, add: ```python import torch def scale( input: torch.Tensor, factor: float, out: torch.Tensor | None = None, ) -> torch.Tensor: """ Element-wise scale: out = input * factor. Supported dtypes: torch.float16, torch.bfloat16, torch.float32. Parameters ---------- input : CUDA input tensor factor : scale factor (float) out : optional pre-allocated CUDA output tensor (same shape/dtype as input) """ if out is None: out = torch.empty_like(input) torch.ops.sgl_kernel.scale.default(out, input, factor) return out ``` Then re-export it from `sgl-kernel/python/sgl_kernel/__init__.py` following the existing import style used by other kernels. --- ## Step 6: Write tests (required) Create `sgl-kernel/tests/test_scale.py`: ```python import pytest import torch import sgl_kernel @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @pytest.mark.parametrize("size", [128, 1024, 4096, 65536]) @pytest.mark.parametrize("factor", [0.5, 1.0, 2.0]) def test_scale_correctness(dtype, size, factor): input = torch.randn(size, dtype=dtype, device="cuda") out = torch.empty_like(input) result = sgl_kernel.scale(input, factor, out=out) assert result is out expected = input * factor rtol, atol = (1e-5, 1e-6) if dtype == torch.float32 else (1e-2, 1e-2) torch.testing.assert_close(out, expected, rtol=rtol, atol=atol) def test_scale_shape_mismatch(): input = torch.randn(128, dtype=torch.float16, device="cuda") out = torch.empty(256, dtype=torch.float16, device="cuda") with pytest.raises(RuntimeError, match="same shape"): sgl_kernel.scale(input, 2.0, out=out) def test_scale_cpu_input(): input = torch.randn(128, dtype=torch.float16) # CPU out = torch.empty_like(input) with pytest.raises(RuntimeError, match="CUDA"): sgl_kernel.scale(input, 2.0, out=out) if __name__ == "__main__": pytest.main([__file__, "-q"]) ``` --- ## Step 7: Add a benchmark (required) Create `sgl-kernel/benchmark/bench_scale.py`: ```python import itertools import os import torch import triton import triton.testing import sgl_kernel IS_CI = ( os.getenv("CI", "false").lower() == "true" or os.getenv("GITHUB_ACTIONS", "false").lower() == "true" ) dtypes = [torch.float16] if IS_CI else [torch.float16, torch.bfloat16, torch.float32] sizes = [4096] if IS_CI else [2**n for n in range(10, 20)] # 1K … 512K factors = [2.0] configs = list(itertools.product(dtypes, sizes)) def torch_scale(input: torch.Tensor, factor: float) -> torch.Tensor: return input * factor @triton.testing.perf_report( triton.testing.Benchmark( x_names=["dtype", "size"], x_vals=configs, line_arg="provider", line_vals=["sglang", "torch"], line_names=["SGL Kernel", "PyTorch"], styles=[("green", "-"), ("red", "--")], ylabel="µs (median)", plot_name="scale-performance", args={}, ) ) def benchmark(dtype, size, provider): input = torch.randn(size, dtype=dtype, device="cuda") out = torch.empty_like(input) factor = 2.0 if provider == "sglang": fn = lambda: sgl_kernel.scale(input, factor, out=out) else: fn = lambda: torch_scale(input, factor) ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( fn, quantiles=[0.5, 0.2, 0.8] ) return 1000 * ms, 1000 * max_ms, 1000 * min_ms if __name__ == "__main__": benchmark.run(print_data=True) ``` --- ## Step 8: Build Build: ```bash cd sgl-kernel make build -j16 ``` If you need to limit host resource usage: ```bash cd sgl-kernel make build -j1 MAX_JOBS=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1" ``` --- ## Step 9: Validate After building successfully, run the test and benchmark: ```bash pytest sgl-kernel/tests/test_scale.py -q python sgl-kernel/benchmark/bench_scale.py ``` --- ## Troubleshooting - **Async CUDA errors**: `CUDA_LAUNCH_BLOCKING=1` - **Memory errors**: `compute-sanitizer --tool memcheck python ...` - **Build is too slow / OOM**: reduce `MAX_JOBS` and `SGL_KERNEL_COMPILE_THREADS` - **Binary bloat**: use `sgl-kernel/analyze_whl_kernel_sizes.py` - **CMake sources list**: if your `.cu` file is missing from `SOURCES`, the symbol will be undefined at link time --- ## References - `sgl-kernel/README.md` - `sgl-kernel/include/sgl_kernel_ops.h` - `sgl-kernel/csrc/common_extension.cc` - `sgl-kernel/CMakeLists.txt` - `sgl-kernel/include/utils.h` — `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16` macro and friends - `sgl-kernel/csrc/elementwise/activation.cu` — reference for the FP16/BF16/FP32 dispatch pattern ## Summary of Files Created/Modified ``` sgl-kernel/csrc/elementwise/scale.cu # NEW: CUDA kernel + launcher sgl-kernel/include/sgl_kernel_ops.h # MODIFIED: C++ declaration sgl-kernel/csrc/common_extension.cc # MODIFIED: schema + dispatch registration sgl-kernel/CMakeLists.txt # MODIFIED: add source file (alphabetical) sgl-kernel/python/sgl_kernel/elementwise.py # MODIFIED: Python wrapper sgl-kernel/python/sgl_kernel/__init__.py # MODIFIED: re-export Python API sgl-kernel/tests/test_scale.py # NEW: tests sgl-kernel/benchmark/bench_scale.py # NEW: benchmark ``` ================================================ FILE: .claude/skills/sglang-bisect-ci-regression/SKILL.md ================================================ # SGLang Bisect CI Regression Investigate a consistently failing CI test to find the root cause - whether it's a code regression from a specific PR, a hardware/runner-specific issue, or an environment change. Optionally reproduce the failure on a remote GPU server. ## Slash Command `/sglang-bisect-ci-regression [ssh_target] [docker_container]` ## When to Use This Skill - A CI test is failing consistently on main (scheduled runs) - You need to find which PR introduced a regression - You suspect a runner-specific or GPU-specific issue - You want to reproduce a CI failure on a remote server ## Arguments - **First argument (required)**: Test file name (e.g. `test_lora_tp.py`) or a GitHub Actions job URL - **Second argument (optional)**: SSH target for remote reproduction (e.g. `user@host`) - **Third argument (optional)**: Docker container name on the SSH target (e.g. `sglang_dev`) If SSH target and docker container are not provided, the skill will only perform the CI log analysis and bisection, without remote reproduction. **Ask the user** for these if reproduction is needed and they weren't provided. ## Background: Scheduled CI Runs SGLang uses the `pr-test.yml` workflow with **scheduled runs** (cron-triggered) to periodically test the `main` branch. These runs are the primary data source for detecting regressions: - **Workflow**: `pr-test.yml` with `event: schedule` - **Branch**: `main` - **Dashboard**: https://github.com/sgl-project/sglang/actions/workflows/pr-test.yml?query=event%3Aschedule - **Frequency**: Runs multiple times daily, each pinned to the HEAD of `main` at trigger time - **Purpose**: Catches regressions that slip through PR-level CI (e.g., interaction bugs between merged PRs, hardware-specific issues) Always use these scheduled runs (not PR-triggered runs) when bisecting regressions on `main`. The `--event schedule` filter in `gh run list` ensures you only see these periodic main-branch runs. ## Workflow ### Phase 1: Extract the Failure Signature 1. **Get the failing test details from CI logs.** If given a URL, fetch logs directly. If given a test name, find recent scheduled runs of `pr-test.yml` on `main` that failed: ```bash # List recent scheduled runs targeting main (the primary source of truth for regressions) # These are cron-triggered runs visible at: # https://github.com/sgl-project/sglang/actions/workflows/pr-test.yml?query=event%3Aschedule gh run list --repo sgl-project/sglang --workflow="pr-test.yml" --event schedule --branch main --limit 20 --json databaseId,conclusion,createdAt,headSha # Find the job containing the test gh run view {RUN_ID} --repo sgl-project/sglang --json jobs --jq '.jobs[] | select(.conclusion == "failure") | {name, conclusion, databaseId}' # Get the failure details gh run view {RUN_ID} --repo sgl-project/sglang --job {JOB_ID} --log 2>&1 | grep -E -B 5 -A 30 "AssertionError|FAIL|Error|{TEST_NAME}" ``` 2. **Record the failure signature:** - Exact error message and assertion - Affected test method name - Model/config involved - Numeric values (e.g., tolerance diffs, scores) - Whether the failure is deterministic (same values across runs) ### Phase 2: Temporal Bisection 3. **Find the boundary between passing and failing runs.** Walk through the scheduled run history (from the `pr-test.yml` schedule runs on `main`) to identify: - Last known PASSING run (sha + date) - First known FAILING run (sha + date) ```bash # For each scheduled run, check the specific partition/job status gh run view {RUN_ID} --repo sgl-project/sglang --json jobs --jq '.jobs[] | select(.name == "{JOB_NAME}") | {conclusion, databaseId}' # Verify a specific test passed or failed in a run gh run view {RUN_ID} --repo sgl-project/sglang --job {JOB_ID} --log 2>&1 | grep -E "{TEST_NAME}|PASSED|FAILED|logprobs mismatch" | head -10 ``` 4. **List commits between the boundary:** ```bash git log --oneline {LAST_PASS_SHA}..{FIRST_FAIL_SHA} ``` 5. **Filter for relevant commits** that touch files related to the failing test (model layers, kernels, test utilities, etc.): ```bash git log --oneline {LAST_PASS_SHA}..{FIRST_FAIL_SHA} -- {relevant_paths} ``` ### Phase 3: Runner/Hardware Analysis 6. **Check if the failure is runner-specific.** Extract the runner identity from each failing and passing run: ```bash # Get runner name and machine gh run view {RUN_ID} --repo sgl-project/sglang --job {JOB_ID} --log 2>&1 | grep -E "Runner name|Machine name" | head -5 # Get GPU/driver info gh run view {RUN_ID} --repo sgl-project/sglang --job {JOB_ID} --log 2>&1 | grep -i -E "NVIDIA-SMI|Driver Version|CUDA Version" | head -5 # Get package versions gh run view {RUN_ID} --repo sgl-project/sglang --job {JOB_ID} --log 2>&1 | grep -E "sgl.kernel.*==|flashinfer.*==" | head -5 ``` 7. **Correlate runners with pass/fail outcomes.** Build a table: | Run ID | Date | Runner | GPU Type | Driver | Result | |--------|------|--------|----------|--------|--------| If all failures map to a specific runner type/GPU and all passes map to another, the issue is **hardware-specific**, not a code regression. ### Phase 4: Code Analysis 8. **If a code regression is suspected** (failures not runner-specific), examine the candidate commits: - Read the changed files - Understand how the changes could affect the failing test - Look for prefill-vs-decode differences, TP-specific paths, kernel changes 9. **If a hardware issue is suspected**, analyze: - Kernel compatibility (CUDA compute capability) - Driver version differences - All-reduce / NCCL behavior differences - CUDA graph capture differences across GPU architectures ### Phase 5: Remote Reproduction (Optional) Only if SSH target and docker container were provided. 10. **Verify the remote environment:** ```bash ssh {SSH_TARGET} "docker exec {CONTAINER} nvidia-smi --query-gpu=name,driver_version --format=csv" ssh {SSH_TARGET} "docker exec {CONTAINER} pip show sgl-kernel sglang flashinfer-python 2>&1 | grep -E 'Name:|Version:'" ``` 11. **Ensure latest code is installed.** If the container is stale, update: ```bash # Try fetching latest main ssh {SSH_TARGET} "docker exec {CONTAINER} bash -c 'cd /path/to/sglang && git fetch origin main && git checkout origin/main'" # Or download and install from tarball if git auth fails ssh {SSH_TARGET} "docker exec {CONTAINER} bash -c 'cd /tmp && curl -L https://github.com/sgl-project/sglang/archive/refs/heads/main.tar.gz | tar xz && cd sglang-main && pip install -e \"python[all]\"'" # Reinstall (after git fetch) ssh {SSH_TARGET} "docker exec {CONTAINER} bash -c 'cd /path/to/sglang && pip install -e \"python[all]\"'" # Install test dependencies if needed ssh {SSH_TARGET} "docker exec {CONTAINER} pip install peft rouge-score" ``` 12. **Create a minimal reproduction script** that: - Uses `if __name__ == '__main__'` with `mp.set_start_method("spawn")` - Runs the specific failing test configuration - Prints key metrics (diffs, scores, outputs) - Exits with code 1 on failure 13. **Copy and run the reproduction script:** ```bash scp /tmp/repro_script.py {SSH_TARGET}:/tmp/ ssh {SSH_TARGET} "docker cp /tmp/repro_script.py {CONTAINER}:/tmp/" ssh {SSH_TARGET} "docker exec -e CUDA_VISIBLE_DEVICES=0,1 {CONTAINER} python3 /tmp/repro_script.py" ``` 14. **Run control experiments** to isolate the variable: - If suspecting TP issue: run with TP=1 as control - If suspecting GPU issue: compare same code on different GPU - If suspecting a specific commit: test before/after that commit ### Phase 6: Report 15. **Produce a structured report:** ```markdown ## CI Regression Bisection Report ### Failure Signature - **Test**: {test_file}::{test_method} - **Error**: {exact error message} - **Key metrics**: {numeric values} - **Deterministic**: Yes/No ### Root Cause Classification One of: - **Code Regression**: PR #{number} introduced the bug - **Hardware-Specific**: Fails on {GPU_TYPE}, passes on others - **Environment Change**: New runner/driver/package version - **Pre-existing Flakiness**: Intermittent, not a new regression ### Evidence | Condition | Result | |-----------|--------| | {condition1} | PASS/FAIL | | {condition2} | PASS/FAIL | ### Timeline - {date}: Last known pass ({sha}, {runner}) - {date}: First known fail ({sha}, {runner}) - {date}: Confirmed reproduction on {server} ### Recommended Fix - **Short-term**: {workaround} - **Long-term**: {proper fix} ``` ## Key Patterns to Recognize | Pattern | Diagnosis | |---------|-----------| | Same SHA passes on runner A, fails on runner B | Hardware/runner-specific | | All runners fail after commit X | Code regression from commit X | | Intermittent - same runner sometimes passes/fails | Flaky test or race condition | | Prefill OK but decode fails | TP/all-reduce issue in decode path | | Works with TP=1, fails with TP>1 | Tensor parallelism bug | | Exact same numeric diff every time | Deterministic bug, not flakiness | ## Important Notes - **Always check runner identity** before concluding it's a code regression. Many "consistent" failures are actually runner-specific. - **Test partition assignments change over time** as tests are added/removed. A test may move between partitions, landing on different runner types. - **H200 runners** use `/root/actions-runner/` path and machine names like `gpu-h200-worker-*`. Non-H200 runners use `/public_sglang_ci/runner-*` paths. - When running remote reproduction, use `run_in_background` for long-running tests and check output with `TaskOutput`. - Container environments may be stale - always verify package versions match CI before drawing conclusions. ================================================ FILE: .claude/skills/write-sglang-test/SKILL.md ================================================ --- name: write-sglang-test description: Guide for writing SGLang CI/UT tests following project conventions. Covers CustomTestCase, CI registration, server fixtures, model selection, and test placement. Use when creating new tests, adding CI test cases, writing unit tests, or when the user asks to add tests for SGLang features. --- # Writing SGLang CI / UT Tests ## Core Rules 1. **Always use `CustomTestCase`** — never raw `unittest.TestCase` 2. **Place tests in `test/registered//`** — only use `test/manual/` for debugging / non-CI tests 3. **Reuse server fixtures** — inherit from `DefaultServerBase` or write `setUpClass`/`tearDownClass` with `popen_launch_server` 4. **Smallest model for model-agnostic functionality** — use `DEFAULT_SMALL_MODEL_NAME_FOR_TEST` (Llama-3.2-1B-Instruct) for basic features that don't depend on model size 5. **8B for general performance** — use `DEFAULT_MODEL_NAME_FOR_TEST` (Llama-3.1-8B-Instruct, single-node) for performance tests that don't involve spec / DP / parallelism 6. **Bigger features → discuss case by case** — spec, DP attention, tensor/pipeline parallelism etc. may need multi-GPU suites and specific models --- ## Test File Template ### Functional correctness test (small model) ```python import unittest import requests from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import ( DEFAULT_SMALL_MODEL_NAME_FOR_TEST, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, popen_launch_server, ) register_cuda_ci(est_time=60, suite="stage-b-test-small-1-gpu") class TestMyFeature(CustomTestCase): @classmethod def setUpClass(cls): cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST cls.base_url = DEFAULT_URL_FOR_TEST cls.process = popen_launch_server( cls.model, cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, other_args=["--arg1", "value1"], # feature-specific args ) @classmethod def tearDownClass(cls): kill_process_tree(cls.process.pid) def test_basic_functionality(self): response = requests.post( self.base_url + "/generate", json={"text": "Hello", "sampling_params": {"max_new_tokens": 32}}, ) self.assertEqual(response.status_code, 200) if __name__ == "__main__": unittest.main(verbosity=3) ``` ### General performance test (8B model, single node, no spec/DP/parallelism) ```python import time import unittest import requests from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import ( DEFAULT_MODEL_NAME_FOR_TEST, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, popen_launch_server, ) register_cuda_ci(est_time=300, suite="stage-b-test-large-1-gpu") class TestMyFeaturePerf(CustomTestCase): @classmethod def setUpClass(cls): cls.model = DEFAULT_MODEL_NAME_FOR_TEST cls.base_url = DEFAULT_URL_FOR_TEST cls.process = popen_launch_server( cls.model, cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, ) @classmethod def tearDownClass(cls): kill_process_tree(cls.process.pid) def test_latency(self): start = time.perf_counter() response = requests.post( self.base_url + "/generate", json={"text": "Hello", "sampling_params": {"max_new_tokens": 128}}, ) elapsed = time.perf_counter() - start self.assertEqual(response.status_code, 200) self.assertLess(elapsed, 5.0, "Latency exceeded threshold") if __name__ == "__main__": unittest.main(verbosity=3) ``` --- ## Server Fixture Reuse For tests that only need a standard server, inherit from `DefaultServerBase` and override class attributes: ```python from sglang.test.server_fixtures.default_fixture import DefaultServerBase class TestMyFeature(DefaultServerBase): model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST other_args = ["--enable-my-feature"] def test_something(self): ... ``` Available fixtures in `python/sglang/test/server_fixtures/`: | Fixture | Use case | |---------|----------| | `DefaultServerBase` | Standard single-server tests | | `EagleServerBase` | EAGLE speculative decoding | | `PDDisaggregationServerBase` | Disaggregated prefill/decode | | `MMMUServerBase` | Multimodal VLM tests | --- ## CI Registration Every test file in `test/registered/` **must** call a registration function at module level: ```python from sglang.test.ci.ci_register import register_cuda_ci, register_amd_ci register_cuda_ci(est_time=60, suite="stage-b-test-small-1-gpu") register_amd_ci(est_time=60, suite="stage-b-test-small-1-gpu-amd") # optional ``` Parameters: - `est_time`: estimated runtime in seconds (used for CI partitioning) - `suite`: which CI suite to run in (see below) - `nightly=True`: for nightly-only tests (default `False` = per-commit) - `disabled="reason"`: temporarily disable with explanation ### Suite selection guide **Default cases (1 GPU):** | Scenario | Model | Suite | |----------|-------|-------| | Model-agnostic basic functionality | 1B (smallest) | `stage-b-test-small-1-gpu` | | General performance (no spec/DP/parallelism) | 8B | `stage-b-test-large-1-gpu` | **Bigger features (case by case):** | Scenario | Suite | |----------|-------| | 2 GPU (e.g. TP=2) | `stage-b-test-large-2-gpu` | | 4 GPU (H100) | `stage-c-test-4-gpu-h100` | | 8 GPU (H200) | `stage-c-test-8-gpu-h200` | | Nightly, 1 GPU | `nightly-1-gpu` | | Nightly, 8 GPU | `nightly-8-gpu` | For spec, DP attention, parallelism, disaggregation, etc., discuss with the team to determine the appropriate suite and GPU configuration. --- ## Model Constants All defined in `python/sglang/test/test_utils.py`: | Constant | Model | When to use | |----------|-------|-------------| | `DEFAULT_SMALL_MODEL_NAME_FOR_TEST` | Llama-3.2-1B-Instruct | Model-agnostic basic functionality | | `DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE` | Llama-3.2-1B | Base (non-instruct) model tests | | `DEFAULT_MODEL_NAME_FOR_TEST` | Llama-3.1-8B-Instruct | General performance (single node) | | `DEFAULT_MOE_MODEL_NAME_FOR_TEST` | Mixtral-8x7B-Instruct | MoE-specific tests | | `DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST` | — | Embedding tests | | `DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST` | — | Vision-language tests | --- ## Test Placement ``` test/ ├── registered/ # CI tests (auto-discovered by run_suite.py) │ ├── sampling/ # test_penalty.py, test_sampling_params.py ... │ ├── sessions/ # test_session_control.py ... │ ├── openai_server/ # basic/, features/, validation/ ... │ ├── spec/ # eagle/, utils/ ... │ ├── models/ # model-specific accuracy tests │ ├── perf/ # performance benchmarks │ └── / # create new category if needed ├── manual/ # Non-CI: debugging, one-off, manual verification └── run_suite.py # CI runner (scans registered/ only) ``` **Decision rule**: if the test should run in CI → `registered/`. If it's for local debugging or requires special hardware not in CI → `manual/`. --- ## Key Utilities ```python from sglang.test.test_utils import ( CustomTestCase, # base class with retry logic popen_launch_server, # launch server subprocess DEFAULT_URL_FOR_TEST, # auto-configured base URL DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, # 600s default run_bench_serving, # benchmark helper (launch + bench) ) from sglang.srt.utils import kill_process_tree # cleanup server ``` --- ## Checklist Before submitting a test: - [ ] Inherits from `CustomTestCase` (not `unittest.TestCase`) - [ ] Has `register_*_ci(...)` call at module level - [ ] Placed in `test/registered//` - [ ] Model selection: smallest for model-agnostic features, 8B for general perf, case-by-case for other complex features - [ ] `setUpClass` launches server, `tearDownClass` kills it - [ ] Has `if __name__ == "__main__": unittest.main(verbosity=3)` - [ ] `est_time` is reasonable (measure locally) ================================================ FILE: .codespellrc ================================================ [codespell] ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS skip = *.json,*.jsonl,*.patch,*.txt ================================================ FILE: .coveragerc ================================================ [run] source = python/sglang/srt omit = */test/* */__pycache__/* [report] show_missing = true exclude_lines = pragma: no cover if __name__ == .__main__.: raise NotImplementedError if TYPE_CHECKING [html] directory = htmlcov ================================================ FILE: .devcontainer/Dockerfile ================================================ FROM lmsysorg/sglang:dev # Create non-root user with specified UID and GID # NOTE: Replace with your own UID and GID. This is a workaround from https://github.com/microsoft/vscode-remote-release/issues/49#issuecomment-489060908. ARG HOST_UID=1003 ARG HOST_GID=1003 RUN groupadd -g $HOST_GID devuser && \ useradd -m -u $HOST_UID -g $HOST_GID -s /bin/zsh devuser # Give devuser sudo access RUN apt-get update && apt-get install -y sudo && \ echo "devuser ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/devuser && \ rm -rf /var/lib/apt/lists/* && \ apt-get clean # Set up oh-my-zsh for devuser RUN cp -r /root/.oh-my-zsh /home/devuser/.oh-my-zsh && \ cp /root/.zshrc /home/devuser/.zshrc && \ cp /root/.vimrc /home/devuser/.vimrc && \ cp /root/.tmux.conf /home/devuser/.tmux.conf && \ sed -i 's|/root/.oh-my-zsh|/home/devuser/.oh-my-zsh|g' /home/devuser/.zshrc && \ chown -R devuser:devuser /home/devuser/ # Set workspace directory and ownership WORKDIR /sgl-workspace/sglang RUN chown -R devuser:devuser /sgl-workspace # Switch to devuser USER devuser # Install uv RUN curl -LsSf https://astral.sh/uv/install.sh | sh # Install rust RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y ================================================ FILE: .devcontainer/devcontainer.json ================================================ { "name": "sglang", "build": { "dockerfile": "Dockerfile" }, "remoteUser": "devuser", "customizations": { "vscode": { "extensions": [ // Python development "ms-python.python", "charliermarsh.ruff", // Rust development "rust-lang.rust-analyzer", "tamasfe.even-better-toml" ] } }, "forwardPorts": [], "runArgs": [ "--gpus", "all" ], // The two lines below ensures that your local changes in the sglang // repo is automatically synced to the sglang pip package installed // in the dev docker container. You can remove / comment out these // two lines if you prefer to sync code changes manually. "workspaceMount": "source=${localWorkspaceFolder},target=/sgl-workspace/sglang,type=bind", "workspaceFolder": "/sgl-workspace/sglang" } ================================================ FILE: .github/CI_PERMISSIONS.json ================================================ { "1pikachu": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "Alcanderian": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "AniZpZ": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "BBuf": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "BHZ-BER": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "ByronHsu": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "CaoE": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "CatherineSue": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "Chen-0210": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "ClawSeven": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "ConnorLi96": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "DarkSharpness": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "Edwardf0t1": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "FlamingoPg": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "FrankLeeeee": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "Fridge003": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "HaiShaw": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "HanHan009527": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "HandH1998": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "Hanrui-Wang": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "HydraQYH": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "JeremieMelo": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "Johnsonms": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "JustinTong0323": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "Kangyan-Zhou": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "LorrinWWW": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "MingxuZh": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "Oasis-Git": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "Prozac614": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "Qiaolin-Yu": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "Qihang-Zhang": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "Ratish1": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "RubiaCx": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "ShangmingCai": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "Shunkangz": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "SimonCqk": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "TianQiLin666666": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "Ubospica": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "Valentine233": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "Xia-Weiwen": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "XiaotongJiang": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "XucSh": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "YAMY1234": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "Ying1123": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "ZailiWang": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "ZhengWG": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "ZhengdQin": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "acelyc111": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "adarshxs": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "airMeng": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "alisonshao": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "alphabetc1": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "amysaq2023": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "attack204": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "ayrnb": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "azhurkevich": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "b8zhong": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "blzheng": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "byjiang1996": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "cctry": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "ch-wan": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "chunyuan-w": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "cicirori": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "cyb70289": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "dongjiyingdjy": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "dougyster": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "elfiegg": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "fy1214": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "fzyzcjy": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "gaopengff": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "gongwei-130": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "gongy": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "guapisolo": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "guoyuhong": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "hanming-lu": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "harrisonlimh": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "harvenstar": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "hebiao064": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "hlu1": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "hnyls2002": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "huaiyuzh": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "huangtingwei9988": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "hubertlu-tw": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "hyhieu": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "hzh0425": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "iforgetmyname": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "ishandhanani": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "ispobock": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "jason-fxz": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "jasperjiaguo": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "jhinpan": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "jianan-gu": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "jinleic": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "jinmingyi1998": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "kaixih": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "kevin85421": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "key4ng": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "kkHuang-amd": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "kpham-sgl": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "kssteven418": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "kushanam": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "lanking520": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "lifuhuang": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "liusy58": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "liz-badada": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "merrymercy": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "mickqian": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "mingfeima": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "minleminzui": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "mmangkad": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "narutolhy": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "netanel-haber": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "nvcastet": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "ocss884": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "pansicheng": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "pavanimajety": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "pdasgup": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "ping1jing2": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "pranavm-nvidia": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "pyc96": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "qingquansong": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "qywu": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "rainj-me": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "ravi03071991": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "rkooo567": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "saienduri": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "samuellees": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "scottjlee": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "sglang-bot": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "sglang-npu-bot": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "shaharmor98": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "shanyu-sys": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "shuaills": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "sleepcoo": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "slin1237": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "stmatengss": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "strgrb": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "sufeng-buaa": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "sundar24295s": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "sunjiweiswift": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "sunxxuns": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "thecodingwizard": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "timmy-feng": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "trevor-m": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "vincentzed": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "wenscarl": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "whybeyoung": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "wisclmy0611": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "xiezhq-hermann": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "xutizhou": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "xyjixyjixyji": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "yanbing-j": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "yangsijia-serena": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "yeahdongcn": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "yhyang201": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "yilian49": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "yinghai": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "yingluosanqian": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "yizhang2077": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "ykcombat": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "ynwang007": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "yuan-luo": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "yundai424": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "yushengsu-thu": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "yyihuang": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "yzh119": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "zhaochenyang20": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true }, "zhijian-liu": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "zhuzilin": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "zhyncs": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "zminglei": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor", "can_rerun_stage": true }, "zyksir": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override", "can_rerun_stage": true }, "zyzshishui": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "custom override", "can_rerun_stage": true } } ================================================ FILE: .github/CODEOWNERS ================================================ .github @merrymercy @Fridge003 @ispobock @Kangyan-Zhou @bingxche /docker @Fridge003 @ispobock @HaiShaw @ishandhanani @yctseng0211 /docker/npu.Dockerfile @ping1jing2 @iforgetmyname /python/pyproject.toml @merrymercy @Fridge003 @ispobock /python/sglang/jit_kernel @DarkSharpness @BBuf @celve @HydraQYH @yuan-luo /python/sglang/jit_kernel/diffusion @yingluosanqian @BBuf @mickqian /python/sglang/multimodal_gen @mickqian @yhyang201 @ping1jing2 /python/sglang/multimodal_gen/runtime/cache @DefTruth /python/sglang/multimodal_gen/runtime/layers @mickqian @yhyang201 @BBuf @yingluosanqian @ping1jing2 /python/sglang/multimodal_gen/runtime/models/dits @mickqian @yhyang201 @BBuf @yingluosanqian @ping1jing2 /python/sglang/srt/batch_invariant_ops @Fridge003 @hebiao064 /python/sglang/srt/constrained @hnyls2002 @DarkSharpness /python/sglang/srt/compilation @hebiao064 /python/sglang/srt/disaggregation @ByronHsu @hnyls2002 @ShangmingCai /python/sglang/srt/disaggregation/ascend @ping1jing2 @iforgetmyname /python/sglang/srt/distributed @yizhang2077 @merrymercy @ch-wan /python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py @ShangmingCai @stmatengss /python/sglang/srt/dllm @ClawSeven @btw616 /python/sglang/srt/entrypoints @ispobock @CatherineSue @slin1237 @merrymercy @JustinTong0323 /python/sglang/srt/entrypoints/grpc_server.py @CatherineSue @slin1237 /python/sglang/srt/eplb @fzyzcjy @ch-wan /python/sglang/srt/function_call @CatherineSue @JustinTong0323 /python/sglang/srt/grpc @CatherineSue @slin1237 /python/sglang/srt/hardware_backend/npu @ping1jing2 @iforgetmyname /python/sglang/srt/hardware_backend/npu/quantization @OrangeRedeng @TamirBaydasov @iforgetmyname /python/sglang/srt/layers @merrymercy @Ying1123 @Fridge003 @ispobock @HaiShaw @ch-wan @BBuf @Edwardf0t1 /python/sglang/srt/layers/attention @merrymercy @Fridge003 @ispobock @Qiaolin-Yu @hebiao064 @HaiShaw /python/sglang/srt/layers/attention/fla @yizhang2077 @hebiao064 /python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @yizhang2077 @hebiao064 @hanming-lu /python/sglang/srt/layers/attention/mamba @yizhang2077 @hebiao064 /python/sglang/srt/layers/attention/nsa @1am9trash @hubertlu-tw @kkHuang-amd @HaiShaw @Fridge003 @hlu1 @rainj-me /python/sglang/srt/layers/attention/vision.py @mickqian @yuan-luo @yhyang201 /python/sglang/srt/layers/quantization @ch-wan @BBuf @Edwardf0t1 @FlamingoPg @AniZpZ @HaiShaw @b8zhong /python/sglang/srt/layers/quantization/quark @kkHuang-amd @yichiche @hubertlu-tw @1am9trash @BowenBao /python/sglang/srt/lora @Ying1123 @Fridge003 @lifuhuang @yushengsu-thu /python/sglang/srt/managers @merrymercy @Ying1123 @hnyls2002 @xiezhq-hermann /python/sglang/srt/managers/scheduler_pp_mixin.py @ShangmingCai @XucSh /python/sglang/srt/mem_cache @merrymercy @Ying1123 @hnyls2002 @xiezhq-hermann @hanming-lu @yizhang2077 @hzh0425 @ispobock /python/sglang/srt/model_executor @merrymercy @Ying1123 @hnyls2002 @Fridge003 @ispobock /python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @hebiao064 /python/sglang/srt/models/deepseek_common @Fridge003 @ispobock @fzyzcjy @ch-wan /python/sglang/srt/models/deepseek_v2.py @fzyzcjy @zhyncs @ispobock @ch-wan @merrymercy @Fridge003 /python/sglang/srt/multimodal @mickqian @JustinTong0323 @yhyang201 @yuan-luo /python/sglang/srt/observability @merrymercy @fzyzcjy @sufeng-buaa /python/sglang/srt/ray @Qiaolin-Yu @xyuzh /python/sglang/srt/speculative @Ying1123 @merrymercy @hnyls2002 /sgl-kernel @ispobock @BBuf @yizhang2077 @merrymercy @FlamingoPg @HaiShaw /sgl-model-gateway @slin1237 @CatherineSue /sgl-model-gateway/benches @slin1237 /sgl-model-gateway/bindings/python @CatherineSue @key4ng @slin1237 /sgl-model-gateway/e2e_test @CatherineSue @key4ng /sgl-model-gateway/src/config @slin1237 /sgl-model-gateway/src/core @slin1237 /sgl-model-gateway/src/data_connector @key4ng /sgl-model-gateway/src/grpc_client @CatherineSue @slin1237 /sgl-model-gateway/src/mcp @key4ng @slin1237 /sgl-model-gateway/src/policies @slin1237 @ByronHsu /sgl-model-gateway/src/proto @CatherineSue @slin1237 /sgl-model-gateway/src/protocols @CatherineSue @key4ng /sgl-model-gateway/src/reasoning_parser @CatherineSue /sgl-model-gateway/src/routers @CatherineSue @key4ng @slin1237 /sgl-model-gateway/src/tokenizer @slin1237 @CatherineSue /sgl-model-gateway/src/tool_parser @slin1237 @CatherineSue /sgl-model-gateway/src/wasm @slin1237 /sgl-model-gateway/examples/wasm @slin1237 /test/srt/ascend @ping1jing2 @iforgetmyname /test/srt/test_modelopt* @Edwardf0t1 ================================================ FILE: .github/FOLDER_README.md ================================================ # Maintenance Tools This folder contains tools and workflows for automating maintenance tasks. ## CI Permissions `CI_PERMISSIONS.json` defines the CI permissions granted to each user. Maintainers can directly edit the file to add entries with `"reason": "custom override"`. Maintainers can also run `update_ci_permission.py` to update it with some auto rules (e.g., top contributors in the last 90 days get full permissions). ## Others - `MAINTAINER.md` defines the code maintenance model. ================================================ FILE: .github/ISSUE_TEMPLATE/1-bug-report.yml ================================================ name: 🐞 Bug report description: Report a bug to help us reproduce and fix it. title: "[Bug] " labels: ['Bug'] body: - type: checkboxes attributes: label: Checklist options: - label: I searched related issues but found no solution. - label: The bug persists in the latest version. - label: Issues without environment info and a minimal reproducible demo are hard to resolve and may receive no feedback. - label: If this is not a bug report but a general question, please start a discussion at https://github.com/sgl-project/sglang/discussions. Otherwise, it will be closed. - label: Please use English. Otherwise, it will be closed. - type: textarea attributes: label: Describe the bug description: A clear, concise description of the bug. validations: required: true - type: textarea attributes: label: Reproduction description: Command/script run and model used. placeholder: Paste the command here. validations: required: true - type: textarea attributes: label: Environment description: Run `python3 -m sglang.check_env` and paste output here. Issues without this will be closed. placeholder: Paste environment output here. validations: required: true ================================================ FILE: .github/ISSUE_TEMPLATE/2-feature-request.yml ================================================ name: 🚀 Feature request description: Suggest an idea for this project title: "[Feature] " body: - type: checkboxes attributes: label: Checklist options: - label: If this is not a feature request but a general question, please start a discussion at https://github.com/sgl-project/sglang/discussions. Otherwise, it will be closed. - label: Please use English. Otherwise, it will be closed. - type: textarea attributes: label: Motivation description: | Clearly and concisely describe the feature's motivation. validations: required: true - type: textarea attributes: label: Related resources description: | Provide official releases or third-party implementations if available. ================================================ FILE: .github/MAINTAINER.md ================================================ # SGLang Code Maintenance Model This document describes the code maintenance model for the SGLang project. Since SGLang is a large project involving multiple organizations and hardware platforms, we designed this model with the following goals: - Ensure a responsive and smooth review process. - Allow for fast iteration, so maintainers can sometimes bypass flaky CI tests for important PRs. ## Role Descriptions There are four roles in this maintenance model. Some are custom roles, while others are predefined by GitHub. - **Merge Oncall**: The person who drives the PR merge process. They have strong area-specific expertise and uphold a high bar for code quality. - Permission: Merge PRs. Bypass branch protection rules if needed. - Responsibility: Shepherd the merge of PRs assigned to their area. Revert or hotfix any issues related to their merge (especially if they bypass). - **Codeowner**: The person who protects critical code. Without a bypass, each PR needs at least one Codeowner approval for each modified file protected by [CODEOWNERS](./CODEOWNERS). Please note that this role is not an honor but a significant responsibility because PRs cannot be merged without your approval (except when bypassed by a Merge Oncall). - Permission: Approve PRs, allowing them to be merged without a bypass. - Responsibility: Review PRs in a timely manner. - **Write**: A person with write permission to the SGLang repo. - Permission: Merge PRs if they have passed required tests and been approved by Codeowners. This role cannot bypass branch protection rules. - Responsibility: Review and merge PRs in a timely manner. - **CI Oncall**: A person who manages CI runners for specific hardware platforms. - Permission: Add CI runners. - Responsibility: Keep the CI runners up and running. __Note__: Difference between Merge Oncall and Codeowner - The Merge Oncall is an active role held by someone who actively tries to help merge PRs and can bypass CI if needed. - The Codeowner is a passive protection role provided by GitHub; it prevents accidental changes to critical code. - The list of Merge Oncalls is attached below. The list of Codeowners is in the [CODEOWNERS](./CODEOWNERS) file. __Note__: The permissions to trigger CI tests are defined separately according to these [rules](https://docs.sglang.io/developer_guide/contribution_guide.html#how-to-trigger-ci-tests). ## Pull Request Merge Process 1. The author submits a pull request (PR) and fills out the PR checklist. 2. A bot assigns this PR to a Merge Oncall and @-mentions them. At the same time, GitHub will automatically request reviews from Codeowners. 3. Someone tags the PR with a `run-ci` label ([help](https://docs.sglang.io/developer_guide/contribution_guide.html#how-to-trigger-ci-tests)). Then the author can trigger CI by pushing new commits. 4. The Merge Oncall coordinates the review (e.g., asking people to review) and approves the PR; the Codeowners also approve the PR. If the assigned Merge Oncall is not responsive, the author can ping other related Merge Oncalls and Reviewers in the list below. 5. The code can now be merged: - **Ideal case:** For each modified file, one Codeowner has approved the PR. The PR has also passed the required CI tests. Then, anyone with write permission can merge the PR. - **Exception:** In cases where it is difficult to meet all requirements (due to flaky CI or slow responses), a Merge Oncall can bypass branch protection to merge the PR. If you meet any issues during the merge, you can discuss in [slack channels](https://slack.sglang.io/): #dev, #pull-request, and #ci-cd-build-release. ## The List of Merge Oncalls and Reviewers The format is @github-username (Slack username). TODO: fill in the list. Now we have many Merge Oncalls mainly because the CI is flaky and the CODEOWNERS is too coarse-grained. In the future, we hope the CI can be improved and we only need bypass rarely. After that, most Merge Oncalls can be converted back to Write and CODEOWNERS. This list is based on the current situation. If you or someone you know would like to take on more responsibility and are qualified, please ping @Lianmin Zheng and @Ying Sheng in the Slack channel. They will start a nomination and internal review process. ## The List of CI Oncalls The format is @github-username (Slack username). ### NVIDIA GPUs @merrymercy (Lianmin Zheng), @Kangyan-Zhou (Kangyan Zhou), @ch-wan (Cheng Wan), @HanHan009527 (hanhan), @ishandhanani (Ishan Dhanani), @key4ng (Keyang Ru), @slin1237 (Simo Lin), @ShangmingCai (Shangming Cai) ### AMD GPUs @saienduri (Sai Enduri), @HaiShaw (Henry HAI) ### Intel CPU and XPU @mingfeima (Mingfei Ma), @DiweiSun (Diwei Sun) ### Ascend NPUs @iforgetmyname (Even Zhou) This list is based on the current situation. If you or someone you know would like to donate machines for CI, they can serve as the CI oncalls for their machines. Please ping @Lianmin Zheng and @Ying Sheng in the Slack channel. They will start a nomination and internal review process. ================================================ FILE: .github/actions/upload-cuda-coredumps/action.yml ================================================ name: Upload CUDA Coredumps description: Upload CUDA coredump files as artifacts and clean up the directory. inputs: artifact-suffix: description: Suffix appended to the artifact name (e.g. matrix partition id) required: false default: "" retention-days: description: Number of days to retain the artifact required: false default: "7" runs: using: composite steps: - name: Upload CUDA coredumps uses: actions/upload-artifact@v4 with: name: cuda-coredumps-${{ github.job }}${{ inputs.artifact-suffix && format('-{0}', inputs.artifact-suffix) }} path: ${{ env.SGLANG_CUDA_COREDUMP_DIR || '/tmp/sglang_cuda_coredumps' }}/ retention-days: ${{ inputs.retention-days }} if-no-files-found: ignore - name: Cleanup CUDA coredumps shell: bash run: rm -rf "${{ env.SGLANG_CUDA_COREDUMP_DIR || '/tmp/sglang_cuda_coredumps' }}" ================================================ FILE: .github/actions/wait-for-jobs/action.yml ================================================ name: Wait for Jobs description: Poll and wait for specified jobs in the current workflow run to complete inputs: stage-name: description: 'Human-readable stage name for log messages (e.g. "stage-a")' required: true jobs: description: | JSON array of job specs to wait for. Each element is either: - a string: exact job name (e.g. "stage-a-test-1") - an object { "prefix": "...", "expected_count": N }: for matrix jobs required: true max-wait-minutes: description: 'Maximum time to wait before timing out' required: false default: '240' poll-interval-seconds: description: 'Seconds between polling attempts' required: false default: '120' github-token: description: 'GitHub token for API calls' required: false default: ${{ github.token }} outputs: result: description: 'Overall result: success, failure, or timeout' value: ${{ steps.wait.outputs.result }} runs: using: composite steps: - name: Wait for jobs to complete id: wait uses: actions/github-script@v7 env: INPUT_STAGE_NAME: ${{ inputs.stage-name }} INPUT_JOBS: ${{ inputs.jobs }} INPUT_MAX_WAIT_MINUTES: ${{ inputs.max-wait-minutes }} INPUT_POLL_INTERVAL_SECONDS: ${{ inputs.poll-interval-seconds }} with: github-token: ${{ inputs.github-token }} script: | const stageName = process.env.INPUT_STAGE_NAME; const jobSpecs = JSON.parse(process.env.INPUT_JOBS); const maxWaitMinutes = parseInt(process.env.INPUT_MAX_WAIT_MINUTES); const pollIntervalSeconds = parseInt(process.env.INPUT_POLL_INTERVAL_SECONDS); const maxAttempts = (maxWaitMinutes * 60) / pollIntervalSeconds; // Normalize job specs into a uniform format const normalizedSpecs = jobSpecs.map(spec => { if (typeof spec === 'string') { return { prefix: spec, expected_count: 1, exact: true }; } return { ...spec, exact: false }; }); const totalExpectedJobs = normalizedSpecs.reduce((sum, s) => sum + s.expected_count, 0); // Match job name: exact match or prefix + " (" for matrix jobs const matchesSpec = (jobName, spec) => { if (spec.exact) { return jobName === spec.prefix; } return jobName === spec.prefix || jobName.startsWith(spec.prefix + ' ('); }; for (let attempt = 0; attempt < maxAttempts; attempt++) { const jobs = await github.paginate(github.rest.actions.listJobsForWorkflowRun, { owner: context.repo.owner, repo: context.repo.repo, run_id: context.runId, per_page: 100, }); let allCompleted = true; let failedJobs = []; let completedCount = 0; let totalCount = 0; for (const spec of normalizedSpecs) { const matchingJobs = jobs.filter(job => matchesSpec(job.name, spec)); for (const job of matchingJobs) { totalCount++; console.log(`${job.name}: status=${job.status}, conclusion=${job.conclusion}`); if (job.status === 'completed') { completedCount++; if (job.conclusion !== 'success' && job.conclusion !== 'skipped') { failedJobs.push(job.name); } } else { allCompleted = false; } } if (matchingJobs.length < spec.expected_count) { console.log(`${spec.prefix}: found ${matchingJobs.length}/${spec.expected_count} jobs (waiting for more)`); allCompleted = false; } } console.log(`[${stageName}] Progress: ${completedCount}/${totalCount} jobs completed (expected ${totalExpectedJobs})`); // Fail fast if any jobs failed if (failedJobs.length > 0) { core.setOutput('result', 'failure'); core.setFailed(`${stageName} jobs failed: ${failedJobs.join(', ')}`); return; } if (allCompleted && totalCount >= totalExpectedJobs) { core.setOutput('result', 'success'); return; } console.log(`Waiting ${pollIntervalSeconds}s... (attempt ${attempt + 1}/${maxAttempts})`); await new Promise(resolve => setTimeout(resolve, pollIntervalSeconds * 1000)); } core.setFailed(`Timeout waiting for ${stageName} jobs`); core.setOutput('result', 'timeout'); ================================================ FILE: .github/labeler.yml ================================================ # Configuration for the GitHub Labeler action # Automatically adds labels to PRs based on the files changed # Router specific (Rust code in sgl-model-gateway) model-gateway: - changed-files: - any-glob-to-any-file: 'sgl-model-gateway/**/*' # Kernel specific sgl-kernel: - changed-files: - any-glob-to-any-file: 'sgl-kernel/**/*' # JIT kernel specific jit-kernel: - changed-files: - any-glob-to-any-file: 'python/sglang/jit_kernel/**/*' # Documentation documentation: - changed-files: - any-glob-to-any-file: - '**/*.md' - 'docs/**/*' - 'README*' # Dependencies dependencies: - changed-files: - any-glob-to-any-file: - '**/requirements*.txt' - '**/Cargo.toml' - '**/Cargo.lock' - '**/pyproject*.toml' - '**/setup.py' - '**/poetry.lock' - '**/package.json' - '**/package-lock.json' # Multi-modal Multi-modal: - changed-files: - any-glob-to-any-file: - '**/*multimodal*' - '**/*vision*' - '**/*vlm*' # Diffusion diffusion: - changed-files: - any-glob-to-any-file: 'python/sglang/multimodal_gen/**/*' # LoRA lora: - changed-files: - any-glob-to-any-file: - '**/*lora*' # Quantization quant: - changed-files: - any-glob-to-any-file: - '**/*quant*' - '**/*quantization*' # Speculative decoding speculative-decoding: - changed-files: - any-glob-to-any-file: - '**/*speculative*' # AMD specific amd: - changed-files: - any-glob-to-any-file: - '**/*amd*' - '**/*rocm*' # NPU specific npu: - changed-files: - any-glob-to-any-file: - '**/*npu*' - '**/*ascend*' # Blackwell blackwell: - changed-files: - any-glob-to-any-file: - '**/*nvfp4*' - 'sgl-kernel/csrc/attention/cutlass_sm100_mla/**/*' - 'python/sglang/srt/layers/attention/trtllm_mla_backend.py' - 'python/sglang/srt/layers/attention/trtllm_mha_backend.py' # DeepSeek specific deepseek: - changed-files: - any-glob-to-any-file: - '**/*deepseek*' # HiCache hicache: - changed-files: - any-glob-to-any-file: - '**/*hicache*' # Deterministic deterministic: - changed-files: - any-glob-to-any-file: 'python/sglang/srt/batch_invariant_ops/**/*' # Piecewise CUDA Graph piecewise-cuda-graph: - changed-files: - any-glob-to-any-file: 'python/sglang/srt/compilation/**/*' # Moore Threads specific mthreads: - changed-files: - any-glob-to-any-file: - '**/*mthreads*' - '**/*musa*' ================================================ FILE: .github/pull_request_template.md ================================================ ## Motivation ## Modifications ## Accuracy Tests ## Benchmarking and Profiling ## Checklist - [ ] Format your code according to the [Format code with pre-commit](https://docs.sglang.io/developer_guide/contribution_guide.html#format-code-with-pre-commit). - [ ] Add unit tests according to the [Run and add unit tests](https://docs.sglang.io/developer_guide/contribution_guide.html#run-and-add-unit-tests). - [ ] Update documentation according to [Write documentations](https://docs.sglang.io/developer_guide/contribution_guide.html#write-documentations). - [ ] Provide accuracy and speed benchmark results according to [Test the accuracy](https://docs.sglang.io/developer_guide/contribution_guide.html#test-the-accuracy) and [Benchmark the speed](https://docs.sglang.io/developer_guide/contribution_guide.html#benchmark-the-speed). - [ ] Follow the SGLang code style [guidance](https://docs.sglang.io/developer_guide/contribution_guide.html#code-style-guidance). ## Review Process 1. Ping Merge Oncalls to start the PR flow. See the [PR Merge Process](https://github.com/sgl-project/sglang/blob/main/.github/MAINTAINER.md#pull-request-merge-process). 2. Get approvals from [CODEOWNERS](https://github.com/sgl-project/sglang/blob/main/.github/CODEOWNERS) and other reviewers. 3. Trigger CI tests with [comments](https://docs.sglang.io/developer_guide/contribution_guide.html#how-to-trigger-ci-tests) or contact authorized users to do so. - `/tag-run-ci-label`, `/rerun-failed-ci`, `/tag-and-rerun-ci` 4. After green CI and required approvals, ask Merge Oncalls to merge. ================================================ FILE: .github/update_ci_permission.py ================================================ """ Update the CI permissions configuration file. This script updates the `CI_PERMISSIONS.json` file, which defines the CI permissions granted to each user. The format of `CI_PERMISSIONS.json` is as follows: { "username1": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, "reason": "top contributor" }, "username2": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override" } } Permissions are assigned according to the following rules: 1. Add the top 50 contributors from the last 90 days with full permissions, no cooldown, and the reason "top contributor". 2. Load all users from the existing `CI_PERMISSIONS.json` file and update their entries as follows: - If a user is already covered by rule 1, skip that user. - If the old reason of a user is "top contributor" but they are not in the current top contributors list, change their configuration to: { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 60, "reason": "custom override" } - For all other cases, preserve the original configuration unchanged. 3. All other users receive no permissions and a 120-minute cooldown (they are omitted from the file). Usage: export GH_TOKEN="your_github_token" python3 update_ci_permission.py # Sort-only mode (no network calls, no GH_TOKEN required) python3 update_ci_permission.py --sort-only """ import argparse import json import os from collections import Counter from datetime import datetime, timedelta, timezone try: import requests except ImportError: requests = None # Only needed for non-sort-only runs # Configuration REPO_OWNER = "sgl-project" REPO_NAME = "sglang" FILE_NAME = os.path.join(os.path.dirname(__file__), "CI_PERMISSIONS.json") HEADERS = {} def github_api_get(endpoint, params=None): """Helper to make paginated GitHub API requests.""" if requests is None: raise RuntimeError( "The requests package is required. Install it or use --sort-only." ) if not HEADERS: raise RuntimeError( "GitHub headers not initialized. Set GH_TOKEN or use --sort-only." ) results = [] url = f"https://api.github.com/repos/{REPO_OWNER}/{REPO_NAME}/{endpoint}" while url: response = requests.get(url, headers=HEADERS, params=params) if response.status_code != 200: print(f"Error fetching {url}: {response.status_code} {response.text}") # If we fail to fetch, strictly return what we have or empty to avoid crashing logic break data = response.json() if isinstance(data, list): results.extend(data) else: return data # Non-list response (not paginated usually) # Handle pagination url = None if "link" in response.headers: links = response.headers["link"].split(", ") for link in links: if 'rel="next"' in link: url = link[link.find("<") + 1 : link.find(">")] params = None # Params are included in the next link break return results def get_write_access_users(): """Fetches users with push (write) or admin access.""" print("Fetching collaborators with write access...") # Note: This endpoint usually requires admin rights on the token. collaborators = github_api_get("collaborators", params={"per_page": 100}) writers = set() for col in collaborators: perms = col.get("permissions", {}) # Check for admin, maintain, or push rights if perms.get("admin") or perms.get("maintain") or perms.get("push"): writers.add(col["login"]) print(f"Found {len(writers)} users with write access.") return writers def get_top_contributors(days=90, limit=50): """Fetches top contributors based on commit count in the last N days.""" print(f"Fetching commits from the last {days} days...") since_date = (datetime.now(timezone.utc) - timedelta(days=days)).isoformat() # Fetch commits commits = github_api_get("commits", params={"since": since_date, "per_page": 100}) author_counts = Counter() for commit in commits: # commit['author'] contains the GitHub user object (can be None if not linked) if commit.get("author") and "login" in commit["author"]: author_counts[commit["author"]["login"]] += 1 top_users = [user for user, _ in author_counts.most_common(limit)] print(f"Found {len(top_users)} active contributors in the last {days} days.") return set(top_users) def load_existing_permissions(): if os.path.exists(FILE_NAME): try: with open(FILE_NAME, "r") as f: return json.load(f) except json.JSONDecodeError: print(f"Warning: {FILE_NAME} is invalid JSON. Starting fresh.") return {} def sort_permissions_file(): """Sort the existing CI permissions file alphabetically and exit.""" if not os.path.exists(FILE_NAME): print(f"{FILE_NAME} not found. Nothing to sort.") return old_permissions = load_existing_permissions() sorted_permissions = dict(sorted(old_permissions.items())) with open(FILE_NAME, "w") as f: json.dump(sorted_permissions, f, indent=4) f.write("\n") print(f"Sorted {FILE_NAME}. Total users: {len(sorted_permissions)}") def main(): parser = argparse.ArgumentParser(description="Update or sort CI permissions.") parser.add_argument( "--sort-only", action="store_true", help="Only sort CI_PERMISSIONS.json alphabetically without fetching data.", ) args = parser.parse_args() if args.sort_only: sort_permissions_file() return gh_token = os.getenv("GH_TOKEN") if not gh_token: raise ValueError("Error: GH_TOKEN environment variable is not set.") global HEADERS HEADERS = { "Authorization": f"Bearer {gh_token}", "Accept": "application/vnd.github+json", "X-GitHub-Api-Version": "2022-11-28", } # Gather Data try: write_access_users = get_write_access_users() except Exception as e: print(f"Warning: Could not fetch collaborators (check token scope). Error: {e}") write_access_users = set() top_contributors = get_top_contributors(days=90, limit=50) old_permissions = load_existing_permissions() new_permissions = {} # Rule 1: Add Top 50 Contributors for user in top_contributors: new_permissions[user] = { "can_tag_run_ci_label": True, "can_rerun_failed_ci": True, "cooldown_interval_minutes": 0, "reason": "top contributor", } # Rule 2: Process Existing Users (Merge Logic) for user, config in old_permissions.items(): if user in new_permissions: # Already handled by Rule 1 or 2 continue old_reason = config.get("reason", "") # If they fell off the top contributor list if old_reason in ["top contributor"]: new_permissions[user] = { "can_tag_run_ci_label": True, "can_rerun_failed_ci": True, "cooldown_interval_minutes": 60, "reason": "custom override", } else: # Preserve custom overrides new_permissions[user] = config # Save and Sort # Sorting keys for cleaner diffs sorted_permissions = dict(sorted(new_permissions.items())) with open(FILE_NAME, "w") as f: json.dump(sorted_permissions, f, indent=4) f.write("\n") # Add trailing newline print(f"Successfully updated {FILE_NAME}. Total users: {len(sorted_permissions)}") if __name__ == "__main__": main() ================================================ FILE: .github/workflows/amd-aiter-scout.yml ================================================ name: AMD AITER Scout on: schedule: - cron: '0 20 * * 1' # Monday 20:00 UTC - cron: '0 20 * * 4' # Thursday 20:00 UTC workflow_dispatch: inputs: aiter_ref: description: 'AITER git ref (branch, tag, or SHA). Default: main (latest commit)' required: false type: string default: 'main' job_filter: description: 'Comma-separated workflows to run: nightly-amd, nightly-amd-rocm720, pr-test-amd, pr-test-amd-rocm720. Default: all' required: false type: string default: 'all' continue_on_error: description: 'Continue running other workflows even if one fails' required: false type: boolean default: true concurrency: group: amd-aiter-scout-${{ github.run_id }} cancel-in-progress: true jobs: resolve-aiter: runs-on: ubuntu-latest outputs: aiter_sha: ${{ steps.resolve.outputs.sha }} run_nightly_amd: ${{ steps.parse.outputs.run_nightly_amd }} run_nightly_amd_rocm720: ${{ steps.parse.outputs.run_nightly_amd_rocm720 }} run_pr_test_amd: ${{ steps.parse.outputs.run_pr_test_amd }} run_pr_test_amd_rocm720: ${{ steps.parse.outputs.run_pr_test_amd_rocm720 }} steps: - name: Resolve AITER commit id: resolve run: | REF="${{ inputs.aiter_ref || 'main' }}" echo "Resolving AITER ref: ${REF}" SHA=$(git ls-remote https://github.com/ROCm/aiter.git "refs/heads/${REF}" | head -1 | cut -f1) if [ -z "$SHA" ]; then SHA=$(git ls-remote https://github.com/ROCm/aiter.git "refs/tags/${REF}" | head -1 | cut -f1) fi if [ -z "$SHA" ]; then SHA=$(git ls-remote https://github.com/ROCm/aiter.git "${REF}" | head -1 | cut -f1) fi if [ -z "$SHA" ]; then SHA="${REF}" fi echo "sha=${SHA}" >> $GITHUB_OUTPUT echo "### AITER Ref Resolution" >> $GITHUB_STEP_SUMMARY echo "- **Requested ref:** \`${REF}\`" >> $GITHUB_STEP_SUMMARY echo "- **Resolved SHA:** \`${SHA}\`" >> $GITHUB_STEP_SUMMARY echo "- **AITER commit:** https://github.com/ROCm/aiter/commit/${SHA}" >> $GITHUB_STEP_SUMMARY - name: Parse job filter id: parse run: | FILTER="${{ inputs.job_filter || 'all' }}" echo "Job filter: ${FILTER}" if [[ "$FILTER" == "all" ]]; then echo "run_nightly_amd=true" >> $GITHUB_OUTPUT echo "run_nightly_amd_rocm720=true" >> $GITHUB_OUTPUT echo "run_pr_test_amd=true" >> $GITHUB_OUTPUT echo "run_pr_test_amd_rocm720=true" >> $GITHUB_OUTPUT else # Wrap with commas for exact substring matching (avoids "nightly-amd" matching "nightly-amd-rocm720") PADDED=",${FILTER// /}," echo "run_nightly_amd=$(echo "$PADDED" | grep -q ',nightly-amd,' && echo true || echo false)" >> $GITHUB_OUTPUT echo "run_nightly_amd_rocm720=$(echo "$PADDED" | grep -q ',nightly-amd-rocm720,' && echo true || echo false)" >> $GITHUB_OUTPUT echo "run_pr_test_amd=$(echo "$PADDED" | grep -q ',pr-test-amd,' && echo true || echo false)" >> $GITHUB_OUTPUT echo "run_pr_test_amd_rocm720=$(echo "$PADDED" | grep -q ',pr-test-amd-rocm720,' && echo true || echo false)" >> $GITHUB_OUTPUT fi echo "### Job Filter" >> $GITHUB_STEP_SUMMARY echo "- **Filter:** \`${FILTER}\`" >> $GITHUB_STEP_SUMMARY call-nightly-amd: if: needs.resolve-aiter.outputs.run_nightly_amd == 'true' needs: resolve-aiter uses: ./.github/workflows/nightly-test-amd.yml secrets: inherit with: ref: ${{ github.sha }} aiter_ref: ${{ needs.resolve-aiter.outputs.aiter_sha }} job_filter: 'all' continue_on_error: ${{ inputs.continue_on_error == '' && true || inputs.continue_on_error }} call-nightly-amd-rocm720: if: needs.resolve-aiter.outputs.run_nightly_amd_rocm720 == 'true' needs: resolve-aiter uses: ./.github/workflows/nightly-test-amd-rocm720.yml secrets: inherit with: ref: ${{ github.sha }} aiter_ref: ${{ needs.resolve-aiter.outputs.aiter_sha }} job_filter: 'all' continue_on_error: ${{ inputs.continue_on_error == '' && true || inputs.continue_on_error }} call-pr-test-amd: if: needs.resolve-aiter.outputs.run_pr_test_amd == 'true' needs: resolve-aiter uses: ./.github/workflows/pr-test-amd.yml secrets: inherit with: run_all_tests: true aiter_ref: ${{ needs.resolve-aiter.outputs.aiter_sha }} continue_on_error: ${{ inputs.continue_on_error == '' && true || inputs.continue_on_error }} call-pr-test-amd-rocm720: if: needs.resolve-aiter.outputs.run_pr_test_amd_rocm720 == 'true' needs: resolve-aiter uses: ./.github/workflows/pr-test-amd-rocm720.yml secrets: inherit with: run_all_tests: true aiter_ref: ${{ needs.resolve-aiter.outputs.aiter_sha }} continue_on_error: ${{ inputs.continue_on_error == '' && true || inputs.continue_on_error }} check-all-jobs: if: always() needs: - resolve-aiter - call-nightly-amd - call-nightly-amd-rocm720 - call-pr-test-amd - call-pr-test-amd-rocm720 runs-on: ubuntu-latest steps: - name: Summary run: | echo "## AMD AITER Scout Results" >> $GITHUB_STEP_SUMMARY echo "" >> $GITHUB_STEP_SUMMARY echo "- **AITER SHA:** \`${{ needs.resolve-aiter.outputs.aiter_sha }}\`" >> $GITHUB_STEP_SUMMARY echo "- **AITER commit:** https://github.com/ROCm/aiter/commit/${{ needs.resolve-aiter.outputs.aiter_sha }}" >> $GITHUB_STEP_SUMMARY echo "" >> $GITHUB_STEP_SUMMARY echo "| Workflow | Result |" >> $GITHUB_STEP_SUMMARY echo "|----------|--------|" >> $GITHUB_STEP_SUMMARY echo "| Nightly AMD (AITER Latest) | \`${{ needs.call-nightly-amd.result }}\` |" >> $GITHUB_STEP_SUMMARY echo "| Nightly AMD ROCm 7.2 | \`${{ needs.call-nightly-amd-rocm720.result }}\` |" >> $GITHUB_STEP_SUMMARY echo "| PR Test AMD (AITER Latest) | \`${{ needs.call-pr-test-amd.result }}\` |" >> $GITHUB_STEP_SUMMARY echo "| PR Test AMD ROCm 7.2 | \`${{ needs.call-pr-test-amd-rocm720.result }}\` |" >> $GITHUB_STEP_SUMMARY - name: Check if any job failed run: | if [[ "${{ contains(needs.*.result, 'failure') }}" == "true" ]]; then echo "One or more workflows failed" exit 1 fi if [[ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]]; then echo "One or more workflows were cancelled" exit 1 fi echo "All workflows passed" ================================================ FILE: .github/workflows/amd-ci-job-monitor.yml ================================================ name: AMD CI Job Monitor on: schedule: - cron: '0 0 * * *' # Daily at midnight UTC pull_request: paths: - '.github/workflows/amd-ci-job-monitor.yml' - 'scripts/ci/utils/query_job_status.py' workflow_dispatch: inputs: hours: description: 'Time window in hours' required: false default: '24' type: string job_filter: description: 'Job name filter (leave empty for all AMD jobs)' required: false type: string jobs: # Single job filter mode custom-report: name: Custom Job Report if: ${{ inputs.job_filter }} runs-on: ubuntu-latest env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} steps: - name: Checkout code uses: actions/checkout@v4 - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Install dependencies run: pip install tabulate - name: Generate Custom Job Report timeout-minutes: 30 run: | python scripts/ci/utils/query_job_status.py \ --repo ${{ github.repository }} \ --job "${{ inputs.job_filter }}" \ --workflow "pr-test-amd.yml" \ --hours ${{ inputs.hours || '24' }} \ --summary # Parse workflow files to get job names dynamically parse-workflows: name: Parse Workflow Jobs if: ${{ !inputs.job_filter }} runs-on: ubuntu-latest outputs: pr_jobs: ${{ steps.parse.outputs.pr_jobs }} nightly_jobs: ${{ steps.parse.outputs.nightly_jobs }} steps: - name: Checkout code uses: actions/checkout@v4 - name: Parse workflow files id: parse run: | # Parse pr-test-amd.yml and extract job names (exclude utility jobs) # Excluded: call-gate, check-changes, pr-test-amd-finish, cancel, check-all-jobs pr_jobs=$(yq -r '.jobs | keys | .[]' .github/workflows/pr-test-amd.yml | \ grep -v -E '^(call-gate|check-changes|pr-test-amd-finish|cancel|check-all-jobs)$' | \ jq -R -s -c 'split("\n") | map(select(length > 0))') echo "pr_jobs=$pr_jobs" >> $GITHUB_OUTPUT echo "PR jobs: $pr_jobs" # Parse nightly-test-amd.yml and extract job names (exclude utility jobs) # Excluded: check-all-jobs nightly_jobs=$(yq -r '.jobs | keys | .[]' .github/workflows/nightly-test-amd.yml | \ grep -v -E '^(check-all-jobs)$' | \ jq -R -s -c 'split("\n") | map(select(length > 0))') echo "nightly_jobs=$nightly_jobs" >> $GITHUB_OUTPUT echo "Nightly jobs: $nightly_jobs" # PR CI reports using dynamic matrix pr-ci-reports: name: PR - ${{ matrix.job_name }} needs: parse-workflows if: ${{ !inputs.job_filter }} runs-on: ubuntu-latest strategy: fail-fast: false matrix: job_name: ${{ fromJson(needs.parse-workflows.outputs.pr_jobs) }} env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} steps: - name: Checkout code uses: actions/checkout@v4 - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Install dependencies run: pip install tabulate - name: Generate Report timeout-minutes: 15 run: | python scripts/ci/utils/query_job_status.py \ --repo ${{ github.repository }} \ --job "${{ matrix.job_name }}" \ --workflow "pr-test-amd.yml" \ --hours ${{ inputs.hours || '24' }} \ --summary # Nightly AMD test reports using dynamic matrix nightly-reports: name: Nightly - ${{ matrix.job_name }} needs: parse-workflows if: ${{ !inputs.job_filter }} runs-on: ubuntu-latest strategy: fail-fast: false matrix: job_name: ${{ fromJson(needs.parse-workflows.outputs.nightly_jobs) }} env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} steps: - name: Checkout code uses: actions/checkout@v4 - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Install dependencies run: pip install tabulate - name: Generate Nightly Report timeout-minutes: 15 run: | python scripts/ci/utils/query_job_status.py \ --repo ${{ github.repository }} \ --job "${{ matrix.job_name }}" \ --workflow "nightly-test-amd.yml" \ --hours ${{ inputs.hours || '24' }} \ --summary ================================================ FILE: .github/workflows/auto-tune.yml ================================================ name: Auto tune on: workflow_dispatch: jobs: lint: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 ================================================ FILE: .github/workflows/bot-bump-flashinfer-version.yml ================================================ name: Bot Bump Flashinfer Version on: workflow_dispatch: inputs: new_version: description: 'New flashinfer version (e.g., 0.6.4)' required: true type: string permissions: contents: write pull-requests: write jobs: bump-flashinfer-version: runs-on: ubuntu-latest steps: - name: Checkout code uses: actions/checkout@v4 with: token: ${{ secrets.GITHUB_TOKEN }} - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Install Python dependencies run: | pip install tomli - name: Configure Git and branch run: | git config user.name "sglang-bot" git config user.email "sglang-bot@users.noreply.github.com" RANDOM_SUFFIX=$(echo $RANDOM | md5sum | head -c 4) BRANCH_NAME="bot/bump-flashinfer-version-${{ github.event.inputs.new_version }}-${RANDOM_SUFFIX}" git checkout -b "$BRANCH_NAME" echo "BRANCH_NAME=$BRANCH_NAME" >> $GITHUB_ENV - name: Run flashinfer version bump script run: | python scripts/release/bump_flashinfer_version.py "${{ github.event.inputs.new_version }}" - name: Commit and create PR env: GH_TOKEN: ${{ secrets.GH_PAT_FOR_PULL_REQUEST }} run: | bash scripts/release/commit_and_pr.sh "flashinfer" "${{ github.event.inputs.new_version }}" "$BRANCH_NAME" ================================================ FILE: .github/workflows/bot-bump-kernel-version-to-sglang.yml ================================================ name: Bot Bump Kernel Version to SGLang on: workflow_dispatch: permissions: contents: write pull-requests: write jobs: bump-kernel-version-to-sglang: runs-on: ubuntu-latest outputs: branch_name: ${{ steps.set_output.outputs.branch_name }} needs_sync: ${{ steps.check_sync.outputs.needs_sync }} steps: - name: Checkout code uses: actions/checkout@v4 with: token: ${{ secrets.GITHUB_TOKEN }} - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Install Python dependencies run: | pip install tomli - name: Check if sync is needed id: check_sync run: | python scripts/release/check_kernel_version_to_sglang.py - name: Configure Git and branch if: steps.check_sync.outputs.needs_sync == 'true' id: set_output run: | git config user.name "sglang-bot" git config user.email "sglang-bot@users.noreply.github.com" RANDOM_SUFFIX=$(echo $RANDOM | md5sum | head -c 4) KERNEL_VERSION="${{ steps.check_sync.outputs.kernel_version }}" BRANCH_NAME="bot/bump-kernel-version-to-sglang-${KERNEL_VERSION}-${RANDOM_SUFFIX}" git checkout -b "$BRANCH_NAME" echo "BRANCH_NAME=$BRANCH_NAME" >> $GITHUB_ENV echo "KERNEL_VERSION=$KERNEL_VERSION" >> $GITHUB_ENV echo "branch_name=$BRANCH_NAME" >> $GITHUB_OUTPUT - name: Run kernel version bump script if: steps.check_sync.outputs.needs_sync == 'true' run: | python scripts/release/bump_kernel_version_to_sglang.py - name: Commit and create PR if: steps.check_sync.outputs.needs_sync == 'true' env: GH_TOKEN: ${{ secrets.GH_PAT_FOR_PULL_REQUEST }} run: | bash scripts/release/commit_and_pr_kernel_to_sglang.sh "$KERNEL_VERSION" "$BRANCH_NAME" run-nightly-tests-nvidia: needs: bump-kernel-version-to-sglang if: needs.bump-kernel-version-to-sglang.outputs.needs_sync == 'true' uses: ./.github/workflows/nightly-test-nvidia.yml with: ref: ${{ needs.bump-kernel-version-to-sglang.outputs.branch_name }} secrets: inherit run-nightly-tests-amd: needs: bump-kernel-version-to-sglang if: needs.bump-kernel-version-to-sglang.outputs.needs_sync == 'true' uses: ./.github/workflows/nightly-test-amd.yml with: ref: ${{ needs.bump-kernel-version-to-sglang.outputs.branch_name }} secrets: inherit run-nightly-tests-npu: needs: bump-kernel-version-to-sglang if: needs.bump-kernel-version-to-sglang.outputs.needs_sync == 'true' uses: ./.github/workflows/nightly-test-npu.yml with: ref: ${{ needs.bump-kernel-version-to-sglang.outputs.branch_name }} secrets: inherit run-pr-tests-xeon: needs: bump-kernel-version-to-sglang if: needs.bump-kernel-version-to-sglang.outputs.needs_sync == 'true' uses: ./.github/workflows/pr-test-xeon.yml with: ref: ${{ needs.bump-kernel-version-to-sglang.outputs.branch_name }} secrets: inherit run-pr-tests-xpu: needs: bump-kernel-version-to-sglang if: needs.bump-kernel-version-to-sglang.outputs.needs_sync == 'true' uses: ./.github/workflows/pr-test-xpu.yml with: ref: ${{ needs.bump-kernel-version-to-sglang.outputs.branch_name }} secrets: inherit ================================================ FILE: .github/workflows/bot-bump-kernel-version.yml ================================================ name: Bot Bump Kernel Version on: workflow_dispatch: inputs: new_version: description: 'New sgl-kernel version (e.g., 0.3.12)' required: true type: string permissions: contents: write pull-requests: write jobs: bump-kernel-version: runs-on: ubuntu-latest steps: - name: Checkout code uses: actions/checkout@v4 with: token: ${{ secrets.GITHUB_TOKEN }} - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Install Python dependencies run: | pip install tomli - name: Configure Git and branch run: | git config user.name "sglang-bot" git config user.email "sglang-bot@users.noreply.github.com" RANDOM_SUFFIX=$(echo $RANDOM | md5sum | head -c 4) BRANCH_NAME="bot/bump-kernel-version-${{ github.event.inputs.new_version }}-${RANDOM_SUFFIX}" git checkout -b "$BRANCH_NAME" echo "BRANCH_NAME=$BRANCH_NAME" >> $GITHUB_ENV - name: Run kernel version bump script run: | python scripts/release/bump_kernel_version.py "${{ github.event.inputs.new_version }}" - name: Commit and create PR env: GH_TOKEN: ${{ secrets.GH_PAT_FOR_PULL_REQUEST }} run: | bash scripts/release/commit_and_pr.sh "sgl-kernel" "${{ github.event.inputs.new_version }}" "$BRANCH_NAME" ================================================ FILE: .github/workflows/bot-bump-sglang-version.yml ================================================ name: Bot Bump SGLang Version on: workflow_dispatch: inputs: new_version: description: 'New SGLang version (e.g., 0.5.3 or 0.5.3rc0)' required: true type: string permissions: contents: write pull-requests: write jobs: bump-sglang-version: runs-on: ubuntu-latest outputs: branch_name: ${{ steps.set_output.outputs.branch_name }} steps: - name: Checkout code uses: actions/checkout@v4 with: token: ${{ secrets.GITHUB_TOKEN }} - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Install Python dependencies run: | pip install tomli - name: Configure Git and branch id: set_output run: | git config user.name "sglang-bot" git config user.email "sglang-bot@users.noreply.github.com" RANDOM_SUFFIX=$(echo $RANDOM | md5sum | head -c 4) BRANCH_NAME="bot/bump-sglang-version-${{ github.event.inputs.new_version }}-${RANDOM_SUFFIX}" git checkout -b "$BRANCH_NAME" echo "BRANCH_NAME=$BRANCH_NAME" >> $GITHUB_ENV echo "branch_name=$BRANCH_NAME" >> $GITHUB_OUTPUT - name: Run SGLang version bump script run: | python scripts/release/bump_sglang_version.py "${{ github.event.inputs.new_version }}" - name: Commit and create PR env: GH_TOKEN: ${{ secrets.GH_PAT_FOR_PULL_REQUEST }} run: | bash scripts/release/commit_and_pr.sh "SGLang" "${{ github.event.inputs.new_version }}" "$BRANCH_NAME" run-nightly-tests-nvidia: needs: bump-sglang-version uses: ./.github/workflows/nightly-test-nvidia.yml with: ref: ${{ needs.bump-sglang-version.outputs.branch_name }} secrets: inherit run-nightly-tests-amd: needs: bump-sglang-version uses: ./.github/workflows/nightly-test-amd.yml with: ref: ${{ needs.bump-sglang-version.outputs.branch_name }} secrets: inherit run-nightly-tests-npu: needs: bump-sglang-version uses: ./.github/workflows/nightly-test-npu.yml with: ref: ${{ needs.bump-sglang-version.outputs.branch_name }} secrets: inherit run-pr-tests-xeon: needs: bump-sglang-version uses: ./.github/workflows/pr-test-xeon.yml with: ref: ${{ needs.bump-sglang-version.outputs.branch_name }} secrets: inherit run-pr-tests-xpu: needs: bump-sglang-version uses: ./.github/workflows/pr-test-xpu.yml with: ref: ${{ needs.bump-sglang-version.outputs.branch_name }} secrets: inherit ================================================ FILE: .github/workflows/bot-cherry-pick.yml ================================================ name: Bot Cherry Pick to Release Branch on: workflow_dispatch: inputs: commit_sha: description: 'Commit SHA to cherry-pick (full or short hash)' required: true type: string target_branch: description: 'Target release branch (e.g., release/v0.5.7)' required: true type: string create_pr: description: 'Create a PR instead of pushing directly' required: false type: boolean default: true permissions: contents: write pull-requests: write concurrency: group: cherry-pick-${{ github.event.inputs.target_branch }} cancel-in-progress: false jobs: cherry-pick: if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-latest environment: 'prod' steps: - name: Validate inputs env: TARGET_BRANCH: ${{ github.event.inputs.target_branch }} run: | if [[ ! "$TARGET_BRANCH" =~ ^release/v[0-9]+\.[0-9]+(\.[0-9]+)?$ ]]; then echo "::error::Target branch must match pattern 'release/vX.Y' or 'release/vX.Y.Z' (e.g., release/v0.5.7)" exit 1 fi - name: Checkout code uses: actions/checkout@v4 with: fetch-depth: 0 token: ${{ secrets.GH_PAT_FOR_PULL_REQUEST }} - name: Configure Git run: | git config user.name "sglang-bot" git config user.email "sglang-bot@users.noreply.github.com" - name: Validate target branch exists env: TARGET_BRANCH: ${{ github.event.inputs.target_branch }} run: | git fetch origin if ! git ls-remote --exit-code --heads origin "$TARGET_BRANCH" > /dev/null 2>&1; then echo "::error::Target branch '$TARGET_BRANCH' does not exist on remote" exit 1 fi - name: Get commit info id: commit_info env: COMMIT_SHA_INPUT: ${{ github.event.inputs.commit_sha }} run: | # Verify commit exists if ! git cat-file -t "$COMMIT_SHA_INPUT" > /dev/null 2>&1; then echo "::error::Commit SHA '$COMMIT_SHA_INPUT' does not exist" exit 1 fi # Get full SHA if short hash provided FULL_SHA=$(git rev-parse "$COMMIT_SHA_INPUT") COMMIT_TITLE=$(git log -1 --format="%s" "$FULL_SHA") SHORT_SHA=$(git rev-parse --short "$FULL_SHA") echo "full_sha=$FULL_SHA" >> $GITHUB_OUTPUT echo "short_sha=$SHORT_SHA" >> $GITHUB_OUTPUT # Use delimiter for multiline-safe output { echo "commit_title<> $GITHUB_OUTPUT echo "Cherry-picking commit: $SHORT_SHA - $COMMIT_TITLE" - name: Cherry-pick commit id: cherry_pick env: TARGET_BRANCH: ${{ github.event.inputs.target_branch }} FULL_SHA: ${{ steps.commit_info.outputs.full_sha }} SHORT_SHA: ${{ steps.commit_info.outputs.short_sha }} CREATE_PR: ${{ github.event.inputs.create_pr }} run: | if [[ "$CREATE_PR" == "true" ]]; then # Create a new branch for the PR RANDOM_SUFFIX=$(head -c 4 /dev/urandom | xxd -p) NEW_BRANCH="cherry-pick/${SHORT_SHA}-to-${TARGET_BRANCH#release/}-${RANDOM_SUFFIX}" git checkout -b "$NEW_BRANCH" "origin/$TARGET_BRANCH" echo "new_branch=$NEW_BRANCH" >> $GITHUB_OUTPUT else # Checkout target branch directly git checkout "$TARGET_BRANCH" fi # Attempt cherry-pick if git cherry-pick "$FULL_SHA"; then echo "cherry_pick_success=true" >> $GITHUB_OUTPUT else echo "::error::Cherry-pick failed due to conflicts. Please resolve manually." git cherry-pick --abort || true echo "cherry_pick_success=false" >> $GITHUB_OUTPUT exit 1 fi - name: Push changes if: steps.cherry_pick.outputs.cherry_pick_success == 'true' env: CREATE_PR: ${{ github.event.inputs.create_pr }} TARGET_BRANCH: ${{ github.event.inputs.target_branch }} NEW_BRANCH: ${{ steps.cherry_pick.outputs.new_branch }} run: | if [[ "$CREATE_PR" == "true" ]]; then git push origin "$NEW_BRANCH" else git push origin "$TARGET_BRANCH" fi - name: Create Pull Request if: steps.cherry_pick.outputs.cherry_pick_success == 'true' && github.event.inputs.create_pr == 'true' env: GH_TOKEN: ${{ secrets.GH_PAT_FOR_PULL_REQUEST }} TARGET_BRANCH: ${{ github.event.inputs.target_branch }} SHORT_SHA: ${{ steps.commit_info.outputs.short_sha }} COMMIT_TITLE: ${{ steps.commit_info.outputs.commit_title }} FULL_SHA: ${{ steps.commit_info.outputs.full_sha }} NEW_BRANCH: ${{ steps.cherry_pick.outputs.new_branch }} run: | PR_TITLE="[Cherry-pick] ${COMMIT_TITLE} to ${TARGET_BRANCH}" gh pr create \ --title "$PR_TITLE" \ --base "$TARGET_BRANCH" \ --head "$NEW_BRANCH" \ --label "cherry-pick" \ --body-file - <> $GITHUB_STEP_SUMMARY echo "" >> $GITHUB_STEP_SUMMARY echo "- **Triggered by:** @${ACTOR}" >> $GITHUB_STEP_SUMMARY echo "- **Commit:** ${FULL_SHA}" >> $GITHUB_STEP_SUMMARY echo "- **Title:** ${COMMIT_TITLE}" >> $GITHUB_STEP_SUMMARY echo "- **Target Branch:** ${TARGET_BRANCH}" >> $GITHUB_STEP_SUMMARY if [[ "$CHERRY_PICK_SUCCESS" == "true" ]]; then echo "- **Status:** ✅ Success" >> $GITHUB_STEP_SUMMARY else echo "- **Status:** ❌ Failed" >> $GITHUB_STEP_SUMMARY fi if [[ "$CREATE_PR" == "true" && "$CHERRY_PICK_SUCCESS" == "true" ]]; then echo "- **PR Branch:** ${NEW_BRANCH}" >> $GITHUB_STEP_SUMMARY fi ================================================ FILE: .github/workflows/cancel-pr-workflow-on-merge.yml ================================================ name: Cancel PR Workflows on Merge on: pull_request_target: types: - closed permissions: actions: write jobs: cancel: if: github.event.pull_request.merged == true runs-on: ubuntu-latest steps: - name: Cancel Previous Runs uses: styfle/cancel-workflow-action@0.12.1 with: workflow_id: all access_token: ${{ secrets.GITHUB_TOKEN }} ignore_sha: true pr_number: ${{ github.event.pull_request.number }} ================================================ FILE: .github/workflows/cancel-unfinished-pr-tests.yml ================================================ name: Cancel Unfinished PR Runs on: workflow_dispatch: inputs: workflows: description: 'Space-separated list of workflow filenames to cancel' required: true type: string default: 'pr-test.yml' permissions: actions: write # Needed to cancel runs contents: read # Needed to read repo info pull-requests: read # needed for gh pr view (labels) jobs: cancel-unfinished-pr-runs: runs-on: ubuntu-latest steps: - name: Install GitHub CLI run: sudo apt-get install -y gh jq - name: Cancel unfinished PR-associated runs (skip high-priority PRs) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} REPO: ${{ github.repository }} WORKFLOWS: ${{ github.event.inputs.workflows || 'pr-test.yml' }} shell: bash run: | set -euo pipefail # Read the space-separated string from the input into a bash array read -r -a WORKFLOW_FILES <<< "${WORKFLOWS}" echo "Targeting ${#WORKFLOW_FILES[@]} workflow(s): ${WORKFLOWS}" echo "" for workflow_file in "${WORKFLOW_FILES[@]}"; do echo "=========================================" echo "Workflow: $workflow_file" echo "=========================================" # Get all unfinished runs all_runs=$(gh run list \ --repo "$REPO" \ --workflow "$workflow_file" \ --json databaseId,status,event,url,createdAt \ --limit 1000 \ | jq -c '.[] | select(.status=="queued" or .status=="waiting" or .status=="in_progress")') if [ -z "$all_runs" ]; then echo "✅ No unfinished runs found" echo "" continue fi # Count runs by event type total_runs=$(echo "$all_runs" | wc -l) pr_runs=$(echo "$all_runs" | jq -s '[.[] | select(.event=="pull_request")] | length') other_runs=$(echo "$all_runs" | jq -s '[.[] | select(.event!="pull_request")] | length') echo "📊 Summary: $total_runs unfinished runs ($pr_runs PR-related, $other_runs other)" echo "" # Process non-PR runs first if [ "$other_runs" -gt 0 ]; then echo "--- Non-PR Runs ---" echo "$all_runs" | jq -c 'select(.event!="pull_request")' | while read -r run; do run_url=$(echo "$run" | jq -r '.url') run_event=$(echo "$run" | jq -r '.event') run_status=$(echo "$run" | jq -r '.status') echo " • $run_event ($run_status): $run_url" done echo "" fi # Process PR runs if [ "$pr_runs" -gt 0 ]; then echo "--- PR Runs (checking for cancellation) ---" echo "$all_runs" | jq -c 'select(.event=="pull_request")' | while read -r run; do run_id=$(echo "$run" | jq -r '.databaseId') run_url=$(echo "$run" | jq -r '.url') run_status=$(echo "$run" | jq -r '.status') echo "" echo "Run ($run_status): $run_url" # Fetch full run details to get head repository and branch info run_details=$(gh api -H "Accept: application/vnd.github+json" \ "repos/$REPO/actions/runs/$run_id" 2>/dev/null || true) if [ -z "$run_details" ]; then echo " ⚠️ Could not fetch run details, skipping" continue fi # Get head owner and branch (works for both fork and non-fork PRs) head_owner=$(echo "$run_details" | jq -r '.head_repository.owner.login // empty') head_branch=$(echo "$run_details" | jq -r '.head_branch // empty') if [ -z "$head_owner" ] || [ -z "$head_branch" ]; then echo " ⚠️ Missing head info, skipping" continue fi echo " Branch: ${head_owner}:${head_branch}" # Find PR by searching with head=owner:branch pr_number=$(gh api -H "Accept: application/vnd.github+json" \ "repos/$REPO/pulls?state=open&head=${head_owner}:${head_branch}" \ --jq '.[0].number // empty' 2>/dev/null || true) if [ -z "$pr_number" ]; then echo " ⚠️ No open PR found, skipping" continue fi pr_url="https://github.com/$REPO/pull/$pr_number" echo " PR: $pr_url" # Check for high priority label labels=$(gh pr view "$pr_number" --repo "$REPO" --json labels \ | jq -r '.labels[].name' 2>/dev/null || true) if echo "$labels" | grep -Fxq "high priority"; then echo " 🛑 Skipping (high priority label)" continue fi echo " 🚫 Cancelling..." gh run cancel "$run_id" --repo "$REPO" || echo " ⚠️ Cancellation failed" done fi echo "" done echo "=========================================" echo "✅ Processing complete" echo "=========================================" ================================================ FILE: .github/workflows/ci-coverage-overview.yml ================================================ name: CI Coverage Overview on: schedule: - cron: '0 6 * * *' # Daily at 6 AM UTC pull_request: paths: - '.github/workflows/ci-coverage-overview.yml' - 'scripts/ci/utils/ci_coverage_report.py' - 'test/registered/**' workflow_dispatch: inputs: output_format: description: 'Output format' required: false default: 'markdown' type: choice options: - markdown - json jobs: summary: name: Summary runs-on: ubuntu-latest steps: - name: Checkout code uses: actions/checkout@v4 - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Generate Summary Report run: | python scripts/ci/utils/ci_coverage_report.py --section summary by-folder: name: Tests by Folder runs-on: ubuntu-latest steps: - name: Checkout code uses: actions/checkout@v4 - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Generate Tests by Folder Report run: | python scripts/ci/utils/ci_coverage_report.py --section by-folder by-suite: name: Tests by Suite runs-on: ubuntu-latest steps: - name: Checkout code uses: actions/checkout@v4 - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Generate Tests by Suite Report run: | python scripts/ci/utils/ci_coverage_report.py --section by-suite unit-test-coverage: name: Unit Test Code Coverage if: github.event_name != 'pull_request' runs-on: 1-gpu-runner timeout-minutes: 30 steps: - name: Checkout code uses: actions/checkout@v4 - name: Install dependencies timeout-minutes: 10 run: | pip install -e "python/[test]" - name: Run unit tests with coverage timeout-minutes: 10 run: | pytest test/registered/unit/ \ --cov --cov-config=.coveragerc \ --cov-report=term-missing:skip-covered \ --continue-on-collection-errors \ -v | tee coverage_output.txt - name: Write coverage to summary if: always() run: | echo "## Unit Test Code Coverage" >> $GITHUB_STEP_SUMMARY echo "" >> $GITHUB_STEP_SUMMARY echo "**Commit:** \`${GITHUB_SHA::8}\` | **Branch:** \`${GITHUB_REF_NAME}\`" >> $GITHUB_STEP_SUMMARY echo "" >> $GITHUB_STEP_SUMMARY # Test result line (e.g., "== 42 passed, 1 failed in 23.5s ==") echo '```' >> $GITHUB_STEP_SUMMARY grep -E '^=+.*passed' coverage_output.txt >> $GITHUB_STEP_SUMMARY || true echo "" >> $GITHUB_STEP_SUMMARY # Coverage total grep -E '^TOTAL ' coverage_output.txt >> $GITHUB_STEP_SUMMARY || true echo '```' >> $GITHUB_STEP_SUMMARY # Partially covered core modules (1-49%) — most actionable for contributors # Only show modules with testable logic; skip configs, models, layers, etc. LOW_COV=$(awk '/^python\/.*%/ { for (i=1; i<=NF; i++) { if ($i ~ /^[0-9]+%$/) { pct = $i + 0 if (pct >= 1 && pct < 50) printf "%-80s %5s %s\n", $1, $(i-2), $i break } } }' coverage_output.txt \ | grep -E '/(mem_cache|managers|sampling|parser|observability|function_call|entrypoints|speculative|multimodal|utils)/' \ | head -40 || true) if [ -n "$LOW_COV" ]; then echo "" >> $GITHUB_STEP_SUMMARY echo "
Core modules with coverage below 50% — good candidates for more unit tests" >> $GITHUB_STEP_SUMMARY echo "" >> $GITHUB_STEP_SUMMARY echo '```' >> $GITHUB_STEP_SUMMARY echo "$LOW_COV" >> $GITHUB_STEP_SUMMARY echo '```' >> $GITHUB_STEP_SUMMARY echo "
" >> $GITHUB_STEP_SUMMARY fi json-export: name: JSON Export runs-on: ubuntu-latest if: inputs.output_format == 'json' steps: - name: Checkout code uses: actions/checkout@v4 - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Generate JSON Report run: | python scripts/ci/utils/ci_coverage_report.py --output-format json > ci_coverage.json - name: Upload JSON artifact uses: actions/upload-artifact@v4 with: name: ci-coverage-report path: ci_coverage.json ================================================ FILE: .github/workflows/ci-failure-monitor.yml ================================================ name: CI Failure Monitor on: schedule: - cron: '0 */12 * * *' # Every 12 hour workflow_dispatch: concurrency: group: ci-failure-monitor-${{ github.ref }} cancel-in-progress: true permissions: contents: read actions: read jobs: failure-analysis: if: github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' runs-on: ubuntu-latest steps: - name: Checkout code uses: actions/checkout@v4 - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.14' - name: Install dependencies run: | python -m pip install --upgrade pip pip install requests slack_sdk - name: Run Failure Analysis env: GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} GH_PAT_FOR_RUNNER_ADMIN: ${{ secrets.GH_PAT_FOR_RUNNER_ADMIN }} PYTHONUNBUFFERED: 1 PYTHONIOENCODING: utf-8 run: | cd scripts/ci_monitor python ci_failures_analysis.py \ --token $GITHUB_TOKEN \ --limit 100 \ --output ci_failure_analysis_$(date +%Y%m%d_%H%M%S).json - name: Upload Analysis Results uses: actions/upload-artifact@v4 with: name: ci-failure-analysis-${{ github.run_number }} path: | scripts/ci_monitor/ci_failure_analysis_*.json retention-days: 7 - name: Send Slack Notification if: always() env: SGLANG_DIFFUSION_SLACK_TOKEN: ${{ secrets.SGLANG_DIFFUSION_SLACK_TOKEN }} run: | cd scripts/ci_monitor LATEST_REPORT=$(ls -t ci_failure_analysis_*.json | head -1) if [ ! -f "$LATEST_REPORT" ]; then echo "No report found, so skipping Slack notification" exit 0 fi if [ -n "$SGLANG_DIFFUSION_SLACK_TOKEN" ]; then python3 post_ci_failures_to_slack.py --report-file "$LATEST_REPORT" else echo "SGLANG_DIFFUSION_SLACK_TOKEN not configured, skipping notification" fi ================================================ FILE: .github/workflows/close-inactive-issues.yml ================================================ name: Close Inactive Issues on: schedule: - cron: '0 0 * * *' workflow_dispatch: permissions: issues: write contents: read jobs: close-inactive-issues: if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-latest steps: - name: Check and close inactive issues uses: actions/github-script@v6 with: github-token: ${{secrets.GITHUB_TOKEN}} script: | const sixtyDaysAgo = new Date(Date.now() - 60 * 24 * 60 * 60 * 1000); const [owner, repo] = process.env.GITHUB_REPOSITORY.split('/'); console.log(`Owner: ${owner}, Repo: ${repo}`); async function fetchIssues(page = 1) { console.log(`Fetching issues for ${owner}/${repo}, page ${page}`); return await github.rest.issues.listForRepo({ owner, repo, state: 'open', sort: 'updated', direction: 'asc', per_page: 100, page: page }); } async function processIssues() { console.log('Starting to process issues'); console.log(`Repository: ${owner}/${repo}`); let page = 1; let hasMoreIssues = true; while (hasMoreIssues) { try { const issues = await fetchIssues(page); console.log(`Fetched ${issues.data.length} issues on page ${page}`); if (issues.data.length === 0) { hasMoreIssues = false; break; } for (const issue of issues.data) { // Skip if the issue has 'good first issue' label if (issue.labels.some(label => label.name === 'good first issue')) { console.log(`Skipping issue #${issue.number} as it's marked as 'good first issue'`); continue; } if (new Date(issue.updated_at) < sixtyDaysAgo) { try { await github.rest.issues.update({ owner, repo, issue_number: issue.number, state: 'closed', labels: [...issue.labels.map(l => l.name), 'inactive'] }); await github.rest.issues.createComment({ owner, repo, issue_number: issue.number, body: 'This issue has been automatically closed due to inactivity. Please feel free to reopen it if needed.' }); console.log(`Closed issue #${issue.number} due to inactivity.`); } catch (error) { console.error(`Failed to close issue #${issue.number}: ${error.message}`); } } else { console.log(`Issue #${issue.number} is still active. Stopping processing.`); hasMoreIssues = false; break; } } page += 1; } catch (error) { console.error(`Error fetching issues on page ${page}: ${error.message}`); hasMoreIssues = false; } } console.log('Finished processing issues'); } await processIssues(); ================================================ FILE: .github/workflows/diffusion-ci-gt-gen.yml ================================================ name: Diffusion CI Ground Truth Generation on: workflow_dispatch: inputs: ref: description: 'Git ref to checkout' required: false default: '' type: string case_ids: description: 'Specific case IDs to run (space-separated, optional)' required: false default: '' type: string concurrency: group: diffusion-ci-gt-gen-${{ github.ref }} cancel-in-progress: true permissions: contents: write actions: read jobs: multimodal-diffusion-gen-1gpu: if: github.repository == 'sgl-project/sglang' runs-on: 1-gpu-runner strategy: matrix: part: [0, 1] timeout-minutes: 150 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: bash scripts/ci/cuda/ci_install_dependency.sh diffusion - name: Generate outputs run: | cd python python -m sglang.multimodal_gen.test.scripts.gen_diffusion_ci_outputs \ --suite 1-gpu \ --partition-id ${{ matrix.part }} \ --total-partitions 2 \ --out-dir ./diffusion-ci-outputs \ --continue-on-error \ ${{ inputs.case_ids != '' && format('--case-ids {0}', inputs.case_ids) || '' }} - name: Upload artifact uses: actions/upload-artifact@v4 with: name: diffusion-gen-1gpu-part${{ matrix.part }} path: python/diffusion-ci-outputs retention-days: 7 multimodal-diffusion-gen-2gpu: if: github.repository == 'sgl-project/sglang' runs-on: 2-gpu-runner strategy: matrix: part: [0, 1] timeout-minutes: 150 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: bash scripts/ci/cuda/ci_install_dependency.sh diffusion - name: Generate outputs run: | cd python python -m sglang.multimodal_gen.test.scripts.gen_diffusion_ci_outputs \ --suite 2-gpu \ --partition-id ${{ matrix.part }} \ --total-partitions 2 \ --out-dir ./diffusion-ci-outputs \ --continue-on-error \ ${{ inputs.case_ids != '' && format('--case-ids {0}', inputs.case_ids) || '' }} - name: Upload artifact uses: actions/upload-artifact@v4 with: name: diffusion-gen-2gpu-part${{ matrix.part }} path: python/diffusion-ci-outputs retention-days: 7 diffusion-ci-push: needs: [multimodal-diffusion-gen-1gpu, multimodal-diffusion-gen-2gpu] if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-latest steps: - name: Checkout code uses: actions/checkout@v4 - name: Download artifacts uses: actions/download-artifact@v4 with: pattern: diffusion-gen-* path: combined merge-multiple: true - name: Collect image files run: | mkdir -p gt_images find combined \( -name "*.png" -o -name "*.jpg" -o -name "*.jpeg" -o -name "*.webp" \) -type f -exec cp -f {} gt_images/ \; - name: Publish GT images to sglang-bot/sglang-ci-data env: GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} run: python scripts/ci/utils/publish_diffusion_gt.py --source-dir gt_images ================================================ FILE: .github/workflows/execute-notebook.yml ================================================ name: Execute Notebooks on: pull_request: branches: [ main ] types: [opened, synchronize, reopened, labeled] paths: - "python/sglang/**" - "docs/**" - "!python/sglang/**/*.md" - "!docs/**/*.md" workflow_dispatch: concurrency: group: execute-notebook-${{ github.ref }} cancel-in-progress: true env: SGLANG_IS_IN_CI: true jobs: call-gate: # Align with PR Test: fail fast if PR doesn't have run-ci label. # This makes /tag-and-rerun-ci work by rerunning this failed workflow. uses: ./.github/workflows/pr-gate.yml secrets: inherit run-all-notebooks: needs: [call-gate] runs-on: 1-gpu-runner if: github.event_name != 'pull_request' || needs.call-gate.result == 'success' steps: - name: Checkout code uses: actions/checkout@v4 - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh pip install -r docs/requirements.txt apt-get update && apt-get install -y pandoc parallel retry ln -sf "$(which python3)" /usr/bin/python - name: Setup Jupyter Kernel run: | python -m ipykernel install --user --name python3 --display-name "Python 3" - name: Execute notebooks timeout-minutes: 40 run: | cd docs make clean make compile notebook-finish: needs: [ call-gate, run-all-notebooks ] runs-on: ubuntu-latest if: always() && needs.run-all-notebooks.result != 'skipped' steps: - name: Check all dependent job statuses run: | results=(${{ join(needs.*.result, ' ') }}) for result in "${results[@]}"; do if [ "$result" = "failure" ] || [ "$result" = "cancelled" ]; then echo "Job failed with result: $result" exit 1 fi done echo "All jobs completed successfully" exit 0 ================================================ FILE: .github/workflows/labeler.yml ================================================ name: Auto Label PRs on: pull_request_target: types: [opened, synchronize, reopened] permissions: contents: read pull-requests: write jobs: label: runs-on: ubuntu-latest steps: - name: Auto-label by file changes uses: actions/labeler@v5 with: repo-token: "${{ secrets.GITHUB_TOKEN }}" configuration-path: .github/labeler.yml sync-labels: false ================================================ FILE: .github/workflows/lint.yml ================================================ name: Lint on: push: branches: [main] pull_request: branches: [main] jobs: lint: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - name: Set up Python uses: actions/setup-python@v4 with: python-version: "3.12" - name: Install pre-commit hook run: | python -m pip install pre-commit pre-commit install - name: Run pre-commit checks run: SKIP=no-commit-to-branch pre-commit run --all-files --show-diff-on-failure - name: Run sgl-kernel clang-format checks uses: DoozyX/clang-format-lint-action@v0.20 with: source: sgl-kernel extensions: h,c,cpp,hpp,cu,cuh,cc clangFormatVersion: 20 style: file ================================================ FILE: .github/workflows/list-active-pr-runs.yml.yml ================================================ name: List Active Runs on: workflow_dispatch: inputs: workflows: description: 'Space-separated list of workflow filenames to check' required: false type: string default: 'pr-test.yml' permissions: actions: read contents: read pull-requests: read jobs: list-active-runs: runs-on: ubuntu-latest steps: - name: Install GitHub CLI run: sudo apt-get install -y gh jq - name: List active runs grouped by PR env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} REPO: ${{ github.repository }} WORKFLOWS: ${{ github.event.inputs.workflows || 'pr-test.yml' }} shell: bash run: | set -euo pipefail echo "=========================================" echo "🔍 Active Workflow Runs Report" echo "=========================================" echo "" # Get all workflows or specific ones read -r -a workflow_files <<< "${WORKFLOWS}" echo "📋 Checking specified workflows: ${WORKFLOWS}" echo "" # Create a temporary file to store PR data pr_data_file=$(mktemp) # Process each workflow for workflow_file in ${workflow_files[@]}; do echo "Scanning workflow: $workflow_file" # Get all active runs (queued, waiting, in_progress) active_runs=$(gh run list \ --repo "$REPO" \ --workflow "$workflow_file" \ --json databaseId,status,event,headBranch,createdAt,updatedAt,headSha,number,attempt \ --limit 500 \ | jq -c '.[] | select(.status=="queued" or .status=="waiting" or .status=="in_progress")') if [ -z "$active_runs" ]; then continue fi # Process each run echo "$active_runs" | while read -r run; do run_id=$(echo "$run" | jq -r '.databaseId') run_status=$(echo "$run" | jq -r '.status') run_event=$(echo "$run" | jq -r '.event') created_at=$(echo "$run" | jq -r '.createdAt') head_sha=$(echo "$run" | jq -r '.headSha') run_number=$(echo "$run" | jq -r '.number') run_attempt=$(echo "$run" | jq -r '.attempt // 1') # Get detailed run information including jobs run_details=$(gh api "repos/$REPO/actions/runs/$run_id" 2>/dev/null || true) if [ -z "$run_details" ]; then continue fi head_owner=$(echo "$run_details" | jq -r '.head_repository.owner.login // empty') head_branch=$(echo "$run_details" | jq -r '.head_branch // empty') if [ -z "$head_owner" ] || [ -z "$head_branch" ]; then continue fi # Find PR number (may be empty for non-PR runs) pr_number=$(gh api "repos/$REPO/pulls?state=open&head=${head_owner}:${head_branch}" \ --jq '.[0].number // empty' 2>/dev/null || true) if [ -z "$pr_number" ]; then pr_number="NO_PR" fi # Get jobs for this run (with pagination to avoid missing jobs) jobs=$(gh api "repos/$REPO/actions/runs/$run_id/jobs" --paginate --jq '.jobs[]' | jq -s '.') running_jobs=$(echo "$jobs" | jq '[.[] | select(.status=="in_progress")] | length') queued_jobs=$(echo "$jobs" | jq '[.[] | select(.status=="queued" or .status=="waiting")] | length') # Get runner info for running jobs runners=$(echo "$jobs" | jq -r '.[] | select(.status=="in_progress") | .runner_name // "N/A"' | paste -sd "," -) # Calculate queue time current_time=$(date -u +%s) created_time=$(date -u -d "$created_at" +%s 2>/dev/null || echo "$current_time") queue_time=$((current_time - created_time)) queue_minutes=$((queue_time / 60)) # Store data in temporary file (unified format with event and branch) echo "$pr_number|$workflow_file|$run_id|$run_status|$running_jobs|$queued_jobs|$runners|$queue_minutes|$created_at|$head_sha|$run_attempt|$run_event|$head_branch" >> "$pr_data_file" done done echo "" echo "=========================================" echo "📊 Active Runs Summary" echo "=========================================" echo "" if [ ! -s "$pr_data_file" ]; then echo "✅ No active runs found" rm -f "$pr_data_file" exit 0 fi # Get unique PR numbers (exclude NO_PR entries) pr_numbers=$(cut -d'|' -f1 < "$pr_data_file" | grep -v '^NO_PR$' | sort -u || true) # Separate high priority and normal PRs high_priority_prs=() normal_prs=() for pr_num in $pr_numbers; do labels=$(gh pr view "$pr_num" --repo "$REPO" --json labels \ | jq -r '.labels[].name' 2>/dev/null || true) if echo "$labels" | grep -Fxq "high priority"; then high_priority_prs+=($pr_num) else normal_prs+=($pr_num) fi done # Combine: high priority first, then normal sorted_pr_numbers=("${high_priority_prs[@]}" "${normal_prs[@]}") pr_count=0 total_running=0 total_queued=0 for pr_num in "${sorted_pr_numbers[@]}"; do pr_count=$((pr_count + 1)) # Get PR details pr_info=$(gh pr view "$pr_num" --repo "$REPO" --json title,author,labels,url 2>/dev/null || true) if [ -z "$pr_info" ]; then continue fi pr_title=$(echo "$pr_info" | jq -r '.title') pr_author=$(echo "$pr_info" | jq -r '.author.login') pr_url=$(echo "$pr_info" | jq -r '.url') pr_labels=$(echo "$pr_info" | jq -r '.labels[].name' | paste -sd ", " -) if [ -z "$pr_labels" ]; then pr_labels="(no labels)" fi # Add priority indicator priority_indicator="" if echo "$pr_labels" | grep -q "high priority"; then priority_indicator="🔴 [HIGH PRIORITY] " fi echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" echo "🔗 ${priority_indicator}PR #$pr_num: $pr_title" echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" echo "👤 Author: $pr_author" echo "🏷️ Labels: $pr_labels" echo "🔗 URL: $pr_url" echo "" # Get all runs for this PR pr_runs=$(grep "^$pr_num|" "$pr_data_file") pr_running_total=0 pr_queued_total=0 echo "$pr_runs" | while read -r line; do workflow=$(echo "$line" | cut -d'|' -f2) run_id=$(echo "$line" | cut -d'|' -f3) status=$(echo "$line" | cut -d'|' -f4) running=$(echo "$line" | cut -d'|' -f5) queued=$(echo "$line" | cut -d'|' -f6) runners=$(echo "$line" | cut -d'|' -f7) queue_min=$(echo "$line" | cut -d'|' -f8) created=$(echo "$line" | cut -d'|' -f9) attempt=$(echo "$line" | cut -d'|' -f11) pr_running_total=$((pr_running_total + running)) pr_queued_total=$((pr_queued_total + queued)) run_url="https://github.com/$REPO/actions/runs/$run_id" # Calculate retry count for this specific run retry_count=$((attempt - 1)) # Show retry indicator retry_indicator="" if [ "$retry_count" -gt 0 ]; then retry_indicator=" 🔄 Retry #$retry_count" fi echo " 📦 Workflow: $workflow (Run #$run_id)$retry_indicator" echo " Status: $status" echo " 🟢 Running jobs: $running" echo " 🟡 Queued jobs: $queued" if [ "$running" -gt 0 ] && [ "$runners" != "" ]; then echo " 🖥️ Runners: $runners" fi if [ "$queue_min" -gt 0 ]; then echo " ⏱️ Queue time: ${queue_min} minutes" fi echo " 🔗 Run URL: $run_url" echo "" done # Summary for this PR pr_running_total=$(grep "^$pr_num|" "$pr_data_file" | cut -d'|' -f5 | awk '{sum+=$1} END {print sum+0}') pr_queued_total=$(grep "^$pr_num|" "$pr_data_file" | cut -d'|' -f6 | awk '{sum+=$1} END {print sum+0}') total_running=$((total_running + pr_running_total)) total_queued=$((total_queued + pr_queued_total)) echo " 📊 PR Total: $pr_running_total running, $pr_queued_total queued" echo "" done # --- Non-PR Runs Section --- non_pr_runs=$(grep '^NO_PR|' "$pr_data_file" 2>/dev/null || true) non_pr_running=0 non_pr_queued=0 if [ -n "$non_pr_runs" ]; then echo "=========================================" echo "📦 Non-PR Runs (manual / scheduled / other)" echo "=========================================" echo "" echo "$non_pr_runs" | while read -r line; do workflow=$(echo "$line" | cut -d'|' -f2) run_id=$(echo "$line" | cut -d'|' -f3) status=$(echo "$line" | cut -d'|' -f4) running=$(echo "$line" | cut -d'|' -f5) queued=$(echo "$line" | cut -d'|' -f6) runners=$(echo "$line" | cut -d'|' -f7) queue_min=$(echo "$line" | cut -d'|' -f8) created=$(echo "$line" | cut -d'|' -f9) attempt=$(echo "$line" | cut -d'|' -f11) event=$(echo "$line" | cut -d'|' -f12) branch=$(echo "$line" | cut -d'|' -f13) run_url="https://github.com/$REPO/actions/runs/$run_id" retry_count=$((attempt - 1)) retry_indicator="" if [ "$retry_count" -gt 0 ]; then retry_indicator=" 🔄 Retry #$retry_count" fi echo " 📦 Workflow: $workflow (Run #$run_id)$retry_indicator" echo " Event: $event" echo " Branch: $branch" echo " Status: $status" echo " 🟢 Running jobs: $running" echo " 🟡 Queued jobs: $queued" if [ "$running" -gt 0 ] && [ "$runners" != "" ]; then echo " 🖥️ Runners: $runners" fi if [ "$queue_min" -gt 0 ]; then echo " ⏱️ Queue time: ${queue_min} minutes" fi echo " 🔗 Run URL: $run_url" echo "" done non_pr_running=$(echo "$non_pr_runs" | cut -d'|' -f5 | awk '{sum+=$1} END {print sum+0}') non_pr_queued=$(echo "$non_pr_runs" | cut -d'|' -f6 | awk '{sum+=$1} END {print sum+0}') non_pr_count=$(echo "$non_pr_runs" | wc -l | tr -d ' ') total_running=$((total_running + non_pr_running)) total_queued=$((total_queued + non_pr_queued)) echo " 📊 Non-PR Total: $non_pr_running running, $non_pr_queued queued" echo "" fi # Overall summary echo "=========================================" echo "📈 Overall Summary" echo "=========================================" echo "Total PRs with active runs: $pr_count" echo "Total non-PR active runs: ${non_pr_count:-0}" echo "Total running jobs: $total_running" echo "Total queued jobs: $total_queued" echo "=========================================" # Cleanup rm -f "$pr_data_file" ================================================ FILE: .github/workflows/nightly-release-gateway.yml ================================================ # Nightly release workflow for SGLang Model Gateway name: Nightly Release SGLang Model Gateway to PyPI on: schedule: # Run at 2 AM UTC every day - cron: '0 2 * * *' workflow_dispatch: # Allow manual trigger jobs: build: name: build on ${{ matrix.platform || matrix.os }} (${{ matrix.target }} - ${{ matrix.manylinux || 'auto' }}) runs-on: ${{ matrix.os }}-latest strategy: fail-fast: false matrix: os: [ubuntu, macos, windows] target: [x86_64, aarch64] manylinux: [auto] include: - os: ubuntu platform: linux - os: windows ls: dir target: x86_64 python-architecture: x64 interpreter: 3.9 3.10 3.11 3.12 3.13 - os: macos target: aarch64 interpreter: 3.9 3.10 3.11 3.12 3.13 - os: ubuntu platform: linux target: aarch64 # musllinux - os: ubuntu platform: linux target: x86_64 manylinux: musllinux_1_1 - os: ubuntu platform: linux target: aarch64 manylinux: musllinux_1_1 exclude: - os: windows target: aarch64 steps: - uses: actions/checkout@v4 with: path: sglang-repo - name: Move sgl-model-gateway folder to root and delete sglang-repo run: | mv sglang-repo/sgl-model-gateway/* . rm -rf sglang-repo ls -alt shell: bash - name: Set up Python uses: actions/setup-python@v5 with: python-version: "3.13" architecture: ${{ matrix.python-architecture || 'x64' }} - name: Modify version for nightly release run: | # Get current version from pyproject.toml CURRENT_VERSION=$(python -c "import tomllib; print(tomllib.load(open('bindings/python/pyproject.toml', 'rb'))['project']['version'])" 2>/dev/null || python -c "import tomli; print(tomli.load(open('bindings/python/pyproject.toml', 'rb'))['project']['version'])") # Create nightly version with date: e.g., 0.2.1.dev20250128 NIGHTLY_VERSION="${CURRENT_VERSION}.dev$(date +%Y%m%d)" echo "Nightly version: $NIGHTLY_VERSION" # Update pyproject.toml with nightly version (temporary, not committed) sed -i.bak "s/version = \"${CURRENT_VERSION}\"/version = \"${NIGHTLY_VERSION}\"/" bindings/python/pyproject.toml # Verify the change cat bindings/python/pyproject.toml | grep "^version" shell: bash - name: Install twine and tomli run: pip install -U twine tomli - name: Install protoc (macOS) if: matrix.os == 'macos' run: brew install protobuf - name: Install protoc (Windows) if: matrix.os == 'windows' run: choco install protoc -y - name: Build wheels uses: PyO3/maturin-action@v1 with: working-directory: bindings/python target: ${{ matrix.target }} manylinux: ${{ matrix.manylinux || 'auto' }} args: --release --out dist --features vendored-openssl --interpreter ${{ matrix.interpreter || '3.9 3.10 3.11 3.12 3.13 3.14' }} rust-toolchain: stable docker-options: -e CI -e CC_aarch64_unknown_linux_gnu=aarch64-linux-gnu-gcc -e CXX_aarch64_unknown_linux_gnu=aarch64-linux-gnu-g++ before-script-linux: | # Install build dependencies (perl/make for vendored OpenSSL, protoc for gRPC) if command -v yum &> /dev/null; then yum update -y && yum install -y wget unzip gcc gcc-c++ perl-core make # Install cross-compilation toolchain for aarch64 if needed if [ "${{ matrix.target }}" = "aarch64" ]; then yum install -y gcc-aarch64-linux-gnu gcc-c++-aarch64-linux-gnu || true fi elif command -v apt-get &> /dev/null; then apt-get update && apt-get install -y wget unzip gcc g++ perl make # Install cross-compilation toolchain for aarch64 if needed if [ "${{ matrix.target }}" = "aarch64" ]; then apt-get install -y gcc-aarch64-linux-gnu g++-aarch64-linux-gnu || true fi fi (cd /tmp && \ wget https://github.com/protocolbuffers/protobuf/releases/download/v32.0/protoc-32.0-linux-x86_64.zip && \ unzip protoc-32.0-linux-x86_64.zip -d /usr/local && \ rm protoc-32.0-linux-x86_64.zip) protoc --version - name: List built packages run: ${{ matrix.ls || 'ls -lh' }} bindings/python/dist/ - name: Check packages run: twine check --strict bindings/python/dist/* - uses: actions/upload-artifact@v4 with: name: packages-${{ matrix.os }}-${{ matrix.target }}-${{ matrix.manylinux || 'auto' }} path: bindings/python/dist/ build-sdist: name: Build SDist runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 with: path: sglang-repo - name: Move sgl-model-gateway folder to root and delete sglang-repo run: | mv sglang-repo/sgl-model-gateway/* . rm -rf sglang-repo ls -alt - name: Set up Python uses: actions/setup-python@v5 with: python-version: "3.13" - name: Modify version for nightly release run: | # Get current version from pyproject.toml CURRENT_VERSION=$(python -c "import tomllib; print(tomllib.load(open('bindings/python/pyproject.toml', 'rb'))['project']['version'])" 2>/dev/null || python -c "import tomli; print(tomli.load(open('bindings/python/pyproject.toml', 'rb'))['project']['version'])") # Create nightly version with date: e.g., 0.2.1.dev20250128 NIGHTLY_VERSION="${CURRENT_VERSION}.dev$(date +%Y%m%d)" echo "Nightly version: $NIGHTLY_VERSION" # Update pyproject.toml with nightly version (temporary, not committed) sed -i "s/version = \"${CURRENT_VERSION}\"/version = \"${NIGHTLY_VERSION}\"/" bindings/python/pyproject.toml # Verify the change cat bindings/python/pyproject.toml | grep "^version" - name: Build SDist uses: PyO3/maturin-action@v1 with: working-directory: bindings/python command: sdist args: --out dist rust-toolchain: stable - uses: actions/upload-artifact@v4 with: name: sdist path: bindings/python/dist/*.tar.gz upload: name: Upload to TestPyPI if: github.repository == 'sgl-project/sglang' # Ensure this job only runs for the sgl-project/sglang repository needs: [build, build-sdist] runs-on: ubuntu-latest steps: - uses: actions/download-artifact@v4 with: path: dist merge-multiple: true - name: Upload to TestPyPI env: TWINE_USERNAME: __token__ TWINE_PASSWORD: ${{ secrets.TEST_PYPI_TOKEN_ROUTER }} run: | pip install twine twine upload --repository testpypi dist/* --verbose ================================================ FILE: .github/workflows/nightly-test-amd-rocm720.yml ================================================ name: Nightly Test (AMD ROCm 7.2) on: schedule: - cron: '30 17 * * *' push: branches: - main paths: - "python/sglang/version.py" workflow_dispatch: inputs: aiter_ref: description: 'Override AITER commit (optional, leave empty to use Dockerfile default)' required: false type: string default: '' continue_on_error: description: 'Continue on error (do not fail the workflow on test failures)' required: false type: boolean default: true job_select: description: 'Select a job to run from dropdown (choose "all" to run all jobs)' required: false type: choice default: 'all' options: - 'all' - nightly-test-1-gpu-unit-rocm720 - nightly-accuracy-2-gpu-rocm720 - nightly-accuracy-2-gpu-vlm-rocm720 - nightly-perf-2-gpu-text-rocm720 - nightly-perf-2-gpu-vlm-rocm720 - nightly-accuracy-8-gpu-rocm720 - nightly-8-gpu-grok1-int4-rocm720 - nightly-8-gpu-grok2-rocm720 - nightly-8-gpu-deepseek-v31-rocm720 - nightly-8-gpu-deepseek-v32-rocm720 - nightly-8-gpu-deepseek-v32-mtp-rocm720 - nightly-8-gpu-deepseek-v3-kv-fp8-rocm720 - nightly-8-gpu-kimi-k25-rocm720 - nightly-8-gpu-qwen3-235b-rocm720 - nightly-8-gpu-qwen35-rocm720 - nightly-8-gpu-glm5-rocm720 - nightly-8-gpu-minimax-m25-rocm720 - nightly-1-gpu-zimage-turbo-rocm720 - nightly-test-1-gpu-mi35x-rocm720 - nightly-accuracy-8-gpu-mi35x-rocm720 - nightly-8-gpu-mi35x-grok1-int4-rocm720 - nightly-8-gpu-mi35x-grok2-rocm720 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-rocm720 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8-rocm720 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion-rocm720 - nightly-accuracy-8-gpu-mi35x-deepseek-v32-rocm720 - nightly-accuracy-8-gpu-mi35x-deepseek-v32-mtp-rocm720 - nightly-perf-8-gpu-mi35x-deepseek-v32-basic-rocm720 - nightly-perf-8-gpu-mi35x-deepseek-v32-mtp-rocm720 - nightly-8-gpu-mi35x-kimi-k25-rocm720 - nightly-8-gpu-mi35x-qwen3-235b-mxfp4-rocm720 - nightly-8-gpu-mi35x-qwen35-rocm720 - nightly-8-gpu-mi35x-glm5-rocm720 - nightly-8-gpu-mi35x-minimax-m25-rocm720 job_filter: description: 'Or type comma-separated job names (overrides dropdown if non-empty)' required: false type: string default: '' workflow_call: inputs: ref: description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.' required: false type: string default: '' aiter_ref: description: 'Override AITER commit (optional, leave empty to use Dockerfile default)' required: false type: string default: '' job_filter: description: 'Select which job to run (leave empty or "all" to run all jobs)' required: false type: string default: 'all' continue_on_error: description: 'Continue on error (do not fail the workflow on test failures)' required: false type: boolean default: true env: AITER_COMMIT_OVERRIDE: ${{ inputs.aiter_ref }} concurrency: # When called via workflow_call with ref set, use a unique group per caller run to avoid # collisions with direct schedule/push triggers. We use inputs.ref (not github.event_name) # to detect this, because github.event_name inherits from the caller in workflow_call. group: nightly-test-amd-rocm720-${{ inputs.ref && format('caller-{0}', github.run_id) || github.ref }} cancel-in-progress: ${{ !inputs.ref && github.event_name != 'workflow_call' }} jobs: # ============================================== MI30x ROCm 7.2 Unit Tests ============================================== # 1-GPU Unit Tests - LoRA, debug utils, scheduler, etc. (MI30x ROCm 7.2) nightly-test-1-gpu-unit-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-test-1-gpu-unit-rocm720,')) runs-on: linux-mi325-1gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Nightly Unit Test ROCm 7.2 (1-GPU) timeout-minutes: 90 run: | bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-1-gpu --nightly --timeout-per-file 900 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # ============================================== MI30x ROCm 7.2 Accuracy Tests ============================================== # 2-GPU Accuracy Tests - GSM8K eval (MI30x ROCm 7.2) nightly-accuracy-2-gpu-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-2-gpu-rocm720,')) runs-on: linux-mi325-2gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Nightly Test ROCm 7.2 (2-GPU) run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 2-GPU VLM Accuracy Tests - Vision-Language Models MMMU evaluation (ROCm 7.2) nightly-accuracy-2-gpu-vlm-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-2-gpu-vlm-rocm720,')) runs-on: linux-mi325-2gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Nightly Accuracy Test ROCm 7.2 (2-GPU VLM MMMU) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-2-gpu-vlm --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 2-GPU Text Models Performance Tests (ROCm 7.2) nightly-perf-2-gpu-text-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-perf-2-gpu-text-rocm720,')) runs-on: linux-mi325-2gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Performance Test ROCm 7.2 (2-GPU Text Models) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_USE_AITER=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-perf-text-2-gpu --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 2-GPU VLM Performance Tests (ROCm 7.2) nightly-perf-2-gpu-vlm-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-perf-2-gpu-vlm-rocm720,')) runs-on: linux-mi325-2gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Performance Test ROCm 7.2 (2-GPU VLM Models) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_USE_AITER=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-perf-vlm-2-gpu --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU Accuracy Tests - GPT-OSS, Grok1-FP8 (ROCm 7.2) nightly-accuracy-8-gpu-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-8-gpu-rocm720,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps - name: Accuracy Test ROCm 7.2 (8-GPU GPT-OSS) timeout-minutes: 180 run: | bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-gpt-oss --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Accuracy Test ROCm 7.2 (8-GPU Grok1-FP8) timeout-minutes: 60 run: | bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-grok1-fp8 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # ============================================== MI30x ROCm 7.2 Combined Accuracy + Performance Tests ============================================== # 8-GPU Grok1-INT4 (Accuracy + Performance) ROCm 7.2 nightly-8-gpu-grok1-int4-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-grok1-int4-rocm720,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps - name: Accuracy Test ROCm 7.2 (8-GPU Grok1-INT4) timeout-minutes: 60 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-grok1-int4 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test ROCm 7.2 (8-GPU Grok1-INT4) timeout-minutes: 60 continue-on-error: true run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-grok1-int4 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU Grok2 (Accuracy + Performance) ROCm 7.2 nightly-8-gpu-grok2-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-grok2-rocm720,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps - name: Accuracy Test ROCm 7.2 (8-GPU Grok2) timeout-minutes: 60 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-grok2 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test ROCm 7.2 (8-GPU Grok2) timeout-minutes: 60 continue-on-error: true run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-grok2 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU DeepSeek-V3.1 (Accuracy + Performance) ROCm 7.2 nightly-8-gpu-deepseek-v31-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-deepseek-v31-rocm720,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps - name: Accuracy Test ROCm 7.2 (8-GPU DeepSeek-V3.1) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_USE_AITER=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-deepseek-v31 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test ROCm 7.2 (8-GPU DeepSeek-V3.1) timeout-minutes: 300 continue-on-error: true run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_USE_ROCM700A=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-deepseek-v31 --nightly --timeout-per-file 18000 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU DeepSeek-V3.2 (Basic Accuracy + Perf) ROCm 7.2 nightly-8-gpu-deepseek-v32-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-deepseek-v32-rocm720,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps - name: Accuracy Test ROCm 7.2 (8-GPU DeepSeek-V3.2 Basic) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-deepseek-v32 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test ROCm 7.2 (8-GPU DeepSeek-V3.2 Basic) timeout-minutes: 150 continue-on-error: true run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-deepseek-v32-basic --nightly --timeout-per-file 5400 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU DeepSeek-V3.2 MTP (MTP Accuracy + Perf) ROCm 7.2 nightly-8-gpu-deepseek-v32-mtp-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-deepseek-v32-mtp-rocm720,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps - name: Accuracy Test ROCm 7.2 (8-GPU DeepSeek-V3.2 MTP) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-deepseek-v32-mtp --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test ROCm 7.2 (8-GPU DeepSeek-V3.2 MTP) timeout-minutes: 180 continue-on-error: true run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-deepseek-v32-mtp --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU DeepSeek-V3 KV FP8 (Basic + MTP with --kv-cache-dtype fp8_e4m3) ROCm 7.2 nightly-8-gpu-deepseek-v3-kv-fp8-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-deepseek-v3-kv-fp8-rocm720,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps - name: DeepSeek-V3 KV FP8 Test ROCm 7.2 (8-GPU Basic + MTP) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-deepseek-v3-kv-fp8 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU Kimi-K2.5 (Accuracy) ROCm 7.2 nightly-8-gpu-kimi-k25-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-kimi-k25-rocm720,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps - name: Accuracy Test ROCm 7.2 (8-GPU Kimi-K2.5) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-kimi-k25 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU Qwen3-235B (Accuracy + Performance) ROCm 7.2 nightly-8-gpu-qwen3-235b-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-qwen3-235b-rocm720,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps - name: Accuracy Test + Performance Test ROCm 7.2 (8-GPU Qwen3) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-8-gpu-qwen3-235b --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU Qwen 3.5 (Accuracy) ROCm 7.2 nightly-8-gpu-qwen35-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-qwen35-rocm720,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-aiter-build --skip-test-time-deps bash scripts/ci/amd/amd_ci_exec.sh pip install mistral-common "lm-eval[api]" - name: Accuracy Test ROCm 7.2 (8-GPU Qwen 3.5) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-qwen35 --nightly --timeout-per-file 3600 --continue-on-error || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU GLM-5 (Accuracy) ROCm 7.2 nightly-8-gpu-glm5-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-glm5-rocm720,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps bash scripts/ci/amd/amd_ci_exec.sh pip install git+https://github.com/huggingface/transformers.git@96f807a33b75 - name: Accuracy Test ROCm 7.2 (8-GPU GLM-5 NSA) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-glm5 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU MiniMax-M2.5 (Accuracy) ROCm 7.2 nightly-8-gpu-minimax-m25-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-minimax-m25-rocm720,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps - name: Accuracy Test ROCm 7.2 (8-GPU MiniMax-M2.5) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_USE_AITER=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-minimax-m25 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # ============================================== MI30x ROCm 7.2 Diffusion Tests ============================================== # 1-GPU Z-Image-Turbo (Diffusion T2I) ROCm 7.2 nightly-1-gpu-zimage-turbo-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-1-gpu-zimage-turbo-rocm720,')) runs-on: linux-mi325-1gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Z-Image-Turbo Diffusion Test ROCm 7.2 (1-GPU) timeout-minutes: 45 run: | bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ -e SGLANG_DIFFUSION_ARTIFACT_DIR="/sglang-checkout/diffusion-artifacts" \ pytest test/registered/amd/test_zimage_turbo.py -v -s ${{ inputs.continue_on_error && '|| true' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Upload generated images if: always() uses: actions/upload-artifact@v4 with: name: zimage-turbo-outputs-rocm720 path: diffusion-artifacts/ if-no-files-found: ignore retention-days: 30 # ============================================== MI35x ROCm 7.2 Tests ============================================== # MI35x 1-GPU ROCm 7.2 tests nightly-test-1-gpu-mi35x-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-test-1-gpu-mi35x-rocm720,')) runs-on: linux-mi35x-gpu-1 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Nightly Test MI35x ROCm 7.2 (1-GPU) timeout-minutes: 90 run: | bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-1-gpu-mi35x --nightly --timeout-per-file 900 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU Accuracy Tests - GPT-OSS (ROCm 7.2) nightly-accuracy-8-gpu-mi35x-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-8-gpu-mi35x-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x ROCm 7.2 (8-GPU GPT-OSS) timeout-minutes: 180 run: | bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU Grok1-INT4 (Accuracy + Performance) ROCm 7.2 nightly-8-gpu-mi35x-grok1-int4-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-grok1-int4-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x ROCm 7.2 (8-GPU Grok1-INT4) timeout-minutes: 60 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-mi35x-grok1-int4 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test MI35x ROCm 7.2 (8-GPU Grok1-INT4) timeout-minutes: 60 continue-on-error: true run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-mi35x-grok1-int4 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU Grok2 (Accuracy + Performance) ROCm 7.2 nightly-8-gpu-mi35x-grok2-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-grok2-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x ROCm 7.2 (8-GPU Grok2) timeout-minutes: 60 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-mi35x-grok2 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test MI35x ROCm 7.2 (8-GPU Grok2) timeout-minutes: 60 continue-on-error: true run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-mi35x-grok2 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-R1-MXFP4 (Accuracy + Performance) ROCm 7.2 nightly-8-gpu-mi35x-deepseek-r1-mxfp4-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-deepseek-r1-mxfp4-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x ROCm 7.2 (8-GPU DeepSeek-R1-MXFP4) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-r1-mxfp4 --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test MI35x ROCm 7.2 (8-GPU DeepSeek-R1-MXFP4) timeout-minutes: 300 continue-on-error: true run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 registered/amd/perf/mi35x/test_deepseek_r1_mxfp4_perf_mi35x.py || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-R1-MXFP4 KV FP8 (Accuracy + Performance) ROCm 7.2 nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x ROCm 7.2 (8-GPU DeepSeek-R1-MXFP4 KV FP8) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8 --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test MI35x ROCm 7.2 (8-GPU DeepSeek-R1-MXFP4 KV FP8) timeout-minutes: 300 continue-on-error: true run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 registered/amd/perf/mi35x/test_deepseek_r1_mxfp4_kv_fp8_perf_mi35x.py || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-R1-MXFP4 AllReduce Fusion (Accuracy + Performance) ROCm 7.2 nightly-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x ROCm 7.2 (8-GPU DeepSeek-R1-MXFP4 AllReduce Fusion) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test MI35x ROCm 7.2 (8-GPU DeepSeek-R1-MXFP4 AllReduce Fusion) timeout-minutes: 300 continue-on-error: true run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 registered/amd/perf/mi35x/test_deepseek_r1_mxfp4_ar_fusion_perf_mi35x.py || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-V3.2 Accuracy Test (ROCm 7.2) nightly-accuracy-8-gpu-mi35x-deepseek-v32-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-8-gpu-mi35x-deepseek-v32-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x ROCm 7.2 (8-GPU DeepSeek-V3.2) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-v32 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-V3.2 TP+MTP Accuracy Test (ROCm 7.2) nightly-accuracy-8-gpu-mi35x-deepseek-v32-mtp-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-8-gpu-mi35x-deepseek-v32-mtp-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x ROCm 7.2 (8-GPU DeepSeek-V3.2 TP+MTP) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-mi35x-deepseek-v32-mtp --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-V3.2 Performance Test (Basic) ROCm 7.2 nightly-perf-8-gpu-mi35x-deepseek-v32-basic-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-perf-8-gpu-mi35x-deepseek-v32-basic-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Performance Test MI35x ROCm 7.2 (8-GPU DeepSeek-V3.2 Basic) timeout-minutes: 150 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-mi35x-deepseek-v32-basic --nightly --timeout-per-file 5400 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU Kimi-K2.5 (Accuracy) ROCm 7.2 nightly-8-gpu-mi35x-kimi-k25-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-kimi-k25-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x ROCm 7.2 (8-GPU Kimi-K2.5) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-mi35x-kimi-k25 --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU Qwen3-235B-MXFP4 (Accuracy + Performance) ROCm 7.2 nightly-8-gpu-mi35x-qwen3-235b-mxfp4-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-qwen3-235b-mxfp4-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test + Performance Test MI35x ROCm 7.2 (8-GPU Qwen3-235B-MXFP4) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-8-gpu-mi35x-qwen3-235b-mxfp4 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU Qwen 3.5 (Accuracy) ROCm 7.2 nightly-8-gpu-mi35x-qwen35-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-qwen35-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-aiter-build --skip-test-time-deps bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate bash scripts/ci/amd/amd_ci_exec.sh pip install mistral-common "lm-eval[api]" - name: Accuracy Test MI35x ROCm 7.2 (8-GPU Qwen 3.5) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-mi35x-qwen35 --nightly --timeout-per-file 3600 --continue-on-error || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} nightly-8-gpu-mi35x-glm5-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-glm5-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate bash scripts/ci/amd/amd_ci_exec.sh pip install git+https://github.com/huggingface/transformers.git@96f807a33b75 - name: Accuracy Test MI35x ROCm 7.2 (8-GPU GLM-5 NSA) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-glm5 --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU MiniMax-M2.5 (Accuracy) ROCm 7.2 nightly-8-gpu-mi35x-minimax-m25-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-minimax-m25-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x ROCm 7.2 (8-GPU MiniMax-M2.5) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_USE_AITER=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-minimax-m25 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-V3.2 Performance Test (MTP) ROCm 7.2 nightly-perf-8-gpu-mi35x-deepseek-v32-mtp-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-perf-8-gpu-mi35x-deepseek-v32-mtp-rocm720,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker (ROCm 7.2) run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Performance Test MI35x ROCm 7.2 (8-GPU DeepSeek-V3.2 MTP) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-mi35x-deepseek-v32-mtp --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} check-all-jobs: if: always() && (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch') needs: # MI30x ROCm 7.2 Unit Tests - nightly-test-1-gpu-unit-rocm720 # MI30x ROCm 7.2 Accuracy Tests - nightly-accuracy-2-gpu-rocm720 - nightly-accuracy-2-gpu-vlm-rocm720 # MI30x ROCm 7.2 Performance Tests - nightly-perf-2-gpu-text-rocm720 - nightly-perf-2-gpu-vlm-rocm720 - nightly-accuracy-8-gpu-rocm720 # MI30x ROCm 7.2 Combined Accuracy + Performance Tests - nightly-8-gpu-grok1-int4-rocm720 - nightly-8-gpu-grok2-rocm720 - nightly-8-gpu-deepseek-v31-rocm720 - nightly-8-gpu-deepseek-v32-rocm720 - nightly-8-gpu-deepseek-v32-mtp-rocm720 - nightly-8-gpu-deepseek-v3-kv-fp8-rocm720 - nightly-8-gpu-kimi-k25-rocm720 - nightly-8-gpu-qwen3-235b-rocm720 - nightly-8-gpu-qwen35-rocm720 - nightly-8-gpu-glm5-rocm720 - nightly-8-gpu-minimax-m25-rocm720 # MI30x ROCm 7.2 Diffusion Tests - nightly-1-gpu-zimage-turbo-rocm720 # MI35x ROCm 7.2 jobs - nightly-test-1-gpu-mi35x-rocm720 - nightly-accuracy-8-gpu-mi35x-rocm720 - nightly-8-gpu-mi35x-grok1-int4-rocm720 - nightly-8-gpu-mi35x-grok2-rocm720 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-rocm720 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8-rocm720 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion-rocm720 - nightly-accuracy-8-gpu-mi35x-deepseek-v32-rocm720 - nightly-accuracy-8-gpu-mi35x-deepseek-v32-mtp-rocm720 - nightly-perf-8-gpu-mi35x-deepseek-v32-basic-rocm720 - nightly-perf-8-gpu-mi35x-deepseek-v32-mtp-rocm720 - nightly-8-gpu-mi35x-kimi-k25-rocm720 - nightly-8-gpu-mi35x-qwen3-235b-mxfp4-rocm720 - nightly-8-gpu-mi35x-qwen35-rocm720 - nightly-8-gpu-mi35x-glm5-rocm720 - nightly-8-gpu-mi35x-minimax-m25-rocm720 runs-on: ubuntu-latest steps: - name: Check if any job failed run: | if [[ "${{ contains(needs.*.result, 'failure') }}" == "true" ]]; then echo "One or more ROCm 7.2 nightly test jobs failed" exit 1 fi if [[ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]]; then echo "One or more ROCm 7.2 nightly test jobs were cancelled" exit 1 fi echo "All ROCm 7.2 nightly test jobs passed" ================================================ FILE: .github/workflows/nightly-test-amd.yml ================================================ name: Nightly Test (AMD) on: schedule: - cron: '30 17 * * *' push: branches: - main paths: - "python/sglang/version.py" workflow_dispatch: inputs: aiter_ref: description: 'Override AITER commit (optional, leave empty to use Dockerfile default)' required: false type: string default: '' continue_on_error: description: 'Continue on error (do not fail the workflow on test failures)' required: false type: boolean default: true job_select: description: 'Select a job to run from dropdown (choose "all" to run all jobs)' required: false type: choice default: 'all' options: - 'all' - nightly-test-1-gpu-unit - nightly-accuracy-2-gpu - nightly-accuracy-2-gpu-vlm - nightly-perf-2-gpu-text - nightly-perf-2-gpu-vlm - nightly-accuracy-8-gpu - nightly-8-gpu-grok1-int4 - nightly-8-gpu-grok2 - nightly-8-gpu-deepseek-v31 - nightly-8-gpu-deepseek-v32 - nightly-8-gpu-deepseek-v32-mtp - nightly-8-gpu-deepseek-v3-kv-fp8 - nightly-8-gpu-kimi-k25 - nightly-8-gpu-qwen3-235b - nightly-8-gpu-qwen35 - nightly-8-gpu-glm5 - nightly-8-gpu-minimax-m25 - nightly-1-gpu-zimage-turbo - nightly-test-1-gpu-mi35x - nightly-accuracy-8-gpu-mi35x - nightly-8-gpu-mi35x-grok1-int4 - nightly-8-gpu-mi35x-grok2 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion - nightly-accuracy-8-gpu-mi35x-deepseek-v32 - nightly-accuracy-8-gpu-mi35x-deepseek-v32-mtp - nightly-perf-8-gpu-mi35x-deepseek-v32-basic - nightly-perf-8-gpu-mi35x-deepseek-v32-mtp - nightly-8-gpu-mi35x-kimi-k25 - nightly-8-gpu-mi35x-qwen3-235b-mxfp4 - nightly-8-gpu-mi35x-qwen35 - nightly-8-gpu-mi35x-glm5 - nightly-8-gpu-mi35x-minimax-m25 job_filter: description: 'Or type comma-separated job names (overrides dropdown if non-empty)' required: false type: string default: '' workflow_call: inputs: ref: description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.' required: false type: string default: '' aiter_ref: description: 'Override AITER commit (optional, leave empty to use Dockerfile default)' required: false type: string default: '' job_filter: description: 'Select which job to run (leave empty or "all" to run all jobs)' required: false type: string default: 'all' continue_on_error: description: 'Continue on error (do not fail the workflow on test failures)' required: false type: boolean default: true env: AITER_COMMIT_OVERRIDE: ${{ inputs.aiter_ref }} concurrency: # When called via workflow_call with ref set, use a unique group per caller run to avoid # collisions with direct schedule/push triggers. We use inputs.ref (not github.event_name) # to detect this, because github.event_name inherits from the caller in workflow_call. group: nightly-test-amd-${{ inputs.ref && format('caller-{0}', github.run_id) || github.ref }} cancel-in-progress: ${{ !inputs.ref && github.event_name != 'workflow_call' }} jobs: # ============================================== MI30x Unit Tests ============================================== # 1-GPU Unit Tests - LoRA, debug utils, scheduler, etc. (MI30x only) nightly-test-1-gpu-unit: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-test-1-gpu-unit,')) runs-on: linux-mi325-1gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Nightly Unit Test (1-GPU) timeout-minutes: 90 run: | bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-1-gpu --nightly --timeout-per-file 900 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # ============================================== MI30x Accuracy Tests ============================================== # 2-GPU Accuracy Tests - GSM8K eval (MI30x only) nightly-accuracy-2-gpu: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-2-gpu,')) runs-on: linux-mi325-2gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Nightly Test (2-GPU) run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 2-GPU VLM Accuracy Tests - Vision-Language Models MMMU evaluation nightly-accuracy-2-gpu-vlm: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-2-gpu-vlm,')) runs-on: linux-mi325-2gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Nightly Accuracy Test (2-GPU VLM MMMU) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-2-gpu-vlm --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 2-GPU Text Models Performance Tests nightly-perf-2-gpu-text: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-perf-2-gpu-text,')) runs-on: linux-mi325-2gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Performance Test (2-GPU Text Models) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_USE_AITER=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-perf-text-2-gpu --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 2-GPU VLM Performance Tests nightly-perf-2-gpu-vlm: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-perf-2-gpu-vlm,')) runs-on: linux-mi325-2gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Performance Test (2-GPU VLM Models) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_USE_AITER=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-perf-vlm-2-gpu --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU Accuracy Tests - GPT-OSS, Grok1-FP8 (accuracy only) nightly-accuracy-8-gpu: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-8-gpu,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Accuracy Test (8-GPU GPT-OSS) timeout-minutes: 180 run: | bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-gpt-oss --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Accuracy Test (8-GPU Grok1-FP8) timeout-minutes: 60 run: | bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-grok1-fp8 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # ============================================== MI30x Combined Accuracy + Performance Tests ============================================== # 8-GPU Grok1-INT4 (Accuracy + Performance combined) nightly-8-gpu-grok1-int4: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-grok1-int4,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Accuracy Test (8-GPU Grok1-INT4) timeout-minutes: 60 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-grok1-int4 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test (8-GPU Grok1-INT4) timeout-minutes: 60 continue-on-error: true # Perf test failure doesn't fail the job if accuracy passed run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-grok1-int4 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU Grok2 (Accuracy + Performance combined) nightly-8-gpu-grok2: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-grok2,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Accuracy Test (8-GPU Grok2) timeout-minutes: 60 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-grok2 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test (8-GPU Grok2) timeout-minutes: 60 continue-on-error: true # Perf test failure doesn't fail the job if accuracy passed run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-grok2 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU DeepSeek-V3.1 (Accuracy + Performance combined) nightly-8-gpu-deepseek-v31: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-deepseek-v31,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Accuracy Test (8-GPU DeepSeek-V3.1) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_USE_AITER=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-deepseek-v31 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test (8-GPU DeepSeek-V3.1) timeout-minutes: 300 continue-on-error: true # Perf test failure doesn't fail the job if accuracy passed run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_USE_ROCM700A=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-deepseek-v31 --nightly --timeout-per-file 18000 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU DeepSeek-V3.2 (Basic Accuracy + Perf) nightly-8-gpu-deepseek-v32: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-deepseek-v32,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Accuracy Test (8-GPU DeepSeek-V3.2 Basic) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-deepseek-v32 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test (8-GPU DeepSeek-V3.2 Basic) timeout-minutes: 150 continue-on-error: true # Perf test failure doesn't fail the job if accuracy passed run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-deepseek-v32-basic --nightly --timeout-per-file 5400 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU DeepSeek-V3.2 MTP (MTP Accuracy + Perf) nightly-8-gpu-deepseek-v32-mtp: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-deepseek-v32-mtp,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Accuracy Test (8-GPU DeepSeek-V3.2 MTP) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-deepseek-v32-mtp --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test (8-GPU DeepSeek-V3.2 MTP) timeout-minutes: 180 continue-on-error: true # Perf test failure doesn't fail the job if accuracy passed run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-deepseek-v32-mtp --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU DeepSeek-V3 KV FP8 (Basic + MTP with --kv-cache-dtype fp8_e4m3) nightly-8-gpu-deepseek-v3-kv-fp8: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-deepseek-v3-kv-fp8,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: DeepSeek-V3 KV FP8 Test (8-GPU Basic + MTP) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-deepseek-v3-kv-fp8 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU Kimi-K2.5 (Accuracy) nightly-8-gpu-kimi-k25: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-kimi-k25,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Accuracy Test (8-GPU Kimi-K2.5) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-kimi-k25 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} nightly-8-gpu-qwen3-235b: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-qwen3-235b,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Accuracy Test + Performance Test (8-GPU Qwen3) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-8-gpu-qwen3-235b --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU Qwen 3.5 (Accuracy) nightly-8-gpu-qwen35: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-qwen35,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh bash scripts/ci/amd/amd_ci_exec.sh pip install mistral-common "lm-eval[api]" - name: Accuracy Test (8-GPU Qwen 3.5) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-qwen35 --nightly --timeout-per-file 3600 || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} nightly-8-gpu-glm5: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-glm5,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh bash scripts/ci/amd/amd_ci_exec.sh pip install git+https://github.com/huggingface/transformers.git@96f807a33b75 - name: Accuracy Test (8-GPU GLM-5 NSA) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-glm5 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # 8-GPU MiniMax-M2.5 (Accuracy) nightly-8-gpu-minimax-m25: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-minimax-m25,')) runs-on: linux-mi325-8gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Accuracy Test (8-GPU MiniMax-M2.5) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_USE_AITER=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-minimax-m25 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # ============================================== MI30x Diffusion Tests ============================================== # 1-GPU Z-Image-Turbo (Diffusion T2I) nightly-1-gpu-zimage-turbo: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-1-gpu-zimage-turbo,')) runs-on: linux-mi325-1gpu-sglang steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Z-Image-Turbo Diffusion Test (1-GPU) timeout-minutes: 45 run: | bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ -e SGLANG_DIFFUSION_ARTIFACT_DIR="/sglang-checkout/diffusion-artifacts" \ pytest test/registered/amd/test_zimage_turbo.py -v -s ${{ inputs.continue_on_error && '|| true' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Upload generated images if: always() uses: actions/upload-artifact@v4 with: name: zimage-turbo-outputs path: diffusion-artifacts/ if-no-files-found: ignore retention-days: 30 # ============================================== MI35x Tests ============================================== # MI35x 1-GPU tests - platform-agnostic tests that may work on CDNA4 (gfx950) nightly-test-1-gpu-mi35x: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-test-1-gpu-mi35x,')) runs-on: linux-mi35x-gpu-1 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Nightly Test MI35x (1-GPU) timeout-minutes: 90 run: | bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-1-gpu-mi35x --nightly --timeout-per-file 900 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU Accuracy Tests - GPT-OSS (accuracy only) nightly-accuracy-8-gpu-mi35x: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-8-gpu-mi35x,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x (8-GPU GPT-OSS) timeout-minutes: 180 run: | bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU Grok1-INT4 (Accuracy + Performance combined) nightly-8-gpu-mi35x-grok1-int4: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-grok1-int4,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x (8-GPU Grok1-INT4) timeout-minutes: 90 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-mi35x-grok1-int4 --nightly --timeout-per-file 5400 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test MI35x (8-GPU Grok1-INT4) timeout-minutes: 60 continue-on-error: true # Perf test failure doesn't fail the job if accuracy passed run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-mi35x-grok1-int4 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU Grok2 (Accuracy + Performance combined) nightly-8-gpu-mi35x-grok2: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-grok2,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x (8-GPU Grok2) timeout-minutes: 60 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-mi35x-grok2 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test MI35x (8-GPU Grok2) timeout-minutes: 60 continue-on-error: true # Perf test failure doesn't fail the job if accuracy passed run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e RCCL_MSCCL_ENABLE=0 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-mi35x-grok2 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-R1-MXFP4 (Accuracy + Performance combined) nightly-8-gpu-mi35x-deepseek-r1-mxfp4: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-deepseek-r1-mxfp4,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x (8-GPU DeepSeek-R1-MXFP4) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-r1-mxfp4 --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test MI35x (8-GPU DeepSeek-R1-MXFP4) timeout-minutes: 300 continue-on-error: true # Perf test failure doesn't fail the job if accuracy passed run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 registered/amd/perf/mi35x/test_deepseek_r1_mxfp4_perf_mi35x.py || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-R1-MXFP4 KV FP8 (Accuracy + Performance combined) nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x (8-GPU DeepSeek-R1-MXFP4 KV FP8) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8 --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test MI35x (8-GPU DeepSeek-R1-MXFP4 KV FP8) timeout-minutes: 300 continue-on-error: true # Perf test failure doesn't fail the job if accuracy passed run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 registered/amd/perf/mi35x/test_deepseek_r1_mxfp4_kv_fp8_perf_mi35x.py || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-R1-MXFP4 AllReduce Fusion (Accuracy + Performance combined) nightly-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x (8-GPU DeepSeek-R1-MXFP4 AllReduce Fusion) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} - name: Performance Test MI35x (8-GPU DeepSeek-R1-MXFP4 AllReduce Fusion) timeout-minutes: 300 continue-on-error: true # Perf test failure doesn't fail the job if accuracy passed run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 registered/amd/perf/mi35x/test_deepseek_r1_mxfp4_ar_fusion_perf_mi35x.py || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-V3.2 Accuracy Test nightly-accuracy-8-gpu-mi35x-deepseek-v32: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-8-gpu-mi35x-deepseek-v32,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x (8-GPU DeepSeek-V3.2) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-v32 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-V3.2 TP+MTP Accuracy Test nightly-accuracy-8-gpu-mi35x-deepseek-v32-mtp: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-8-gpu-mi35x-deepseek-v32-mtp,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x (8-GPU DeepSeek-V3.2 TP+MTP) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-mi35x-deepseek-v32-mtp --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-V3.2 Performance Test (Basic) nightly-perf-8-gpu-mi35x-deepseek-v32-basic: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-perf-8-gpu-mi35x-deepseek-v32-basic,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Performance Test MI35x (8-GPU DeepSeek-V3.2 Basic) timeout-minutes: 150 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-mi35x-deepseek-v32-basic --nightly --timeout-per-file 5400 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU Kimi-K2.5 (Accuracy) nightly-8-gpu-mi35x-kimi-k25: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-kimi-k25,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x (8-GPU Kimi-K2.5) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-mi35x-kimi-k25 --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU Qwen3-235B-MXFP4 (Accuracy + Performance) nightly-8-gpu-mi35x-qwen3-235b-mxfp4: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-qwen3-235b-mxfp4,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test + Performance Test MI35x (8-GPU Qwen3-235B-MXFP4) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-8-gpu-mi35x-qwen3-235b-mxfp4 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU Qwen 3.5 (Accuracy) nightly-8-gpu-mi35x-qwen35: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-qwen35,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate bash scripts/ci/amd/amd_ci_exec.sh pip install mistral-common "lm-eval[api]" - name: Accuracy Test MI35x (8-GPU Qwen 3.5) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-mi35x-qwen35 --nightly --timeout-per-file 3600 || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} nightly-8-gpu-mi35x-glm5: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-glm5,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate bash scripts/ci/amd/amd_ci_exec.sh pip install git+https://github.com/huggingface/transformers.git@96f807a33b75 - name: Accuracy Test MI35x (8-GPU GLM-5 NSA) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-glm5 --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU MiniMax-M2.5 (Accuracy) nightly-8-gpu-mi35x-minimax-m25: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-minimax-m25,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Accuracy Test MI35x (8-GPU MiniMax-M2.5) timeout-minutes: 120 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_USE_AITER=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-minimax-m25 --nightly --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} # MI35x 8-GPU DeepSeek-V3.2 Performance Test (MTP) nightly-perf-8-gpu-mi35x-deepseek-v32-mtp: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-perf-8-gpu-mi35x-deepseek-v32-mtp,')) runs-on: linux-mi35x-gpu-8 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Setup docker run: | touch github_summary.md bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh # Install tabulate for run_suite.py (missing in MI35x container) bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate - name: Performance Test MI35x (8-GPU DeepSeek-V3.2 MTP) timeout-minutes: 180 run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-perf-8-gpu-mi35x-deepseek-v32-mtp --nightly --timeout-per-file 7200 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} check-all-jobs: if: always() && (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch') needs: # MI30x Unit Tests - nightly-test-1-gpu-unit # MI30x Accuracy Tests - nightly-accuracy-2-gpu - nightly-accuracy-2-gpu-vlm - nightly-accuracy-8-gpu # MI30x Performance Tests - excluded from check (perf failures don't block CI) # - nightly-perf-2-gpu-text # - nightly-perf-2-gpu-vlm # MI30x Combined Accuracy + Performance Tests - nightly-8-gpu-grok1-int4 - nightly-8-gpu-grok2 - nightly-8-gpu-deepseek-v31 - nightly-8-gpu-deepseek-v32 - nightly-8-gpu-deepseek-v32-mtp - nightly-8-gpu-deepseek-v3-kv-fp8 - nightly-8-gpu-kimi-k25 - nightly-8-gpu-qwen3-235b - nightly-8-gpu-qwen35 - nightly-8-gpu-glm5 - nightly-8-gpu-minimax-m25 # MI30x Diffusion Tests - nightly-1-gpu-zimage-turbo # MI35x jobs - nightly-test-1-gpu-mi35x - nightly-accuracy-8-gpu-mi35x - nightly-8-gpu-mi35x-grok1-int4 - nightly-8-gpu-mi35x-grok2 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion - nightly-accuracy-8-gpu-mi35x-deepseek-v32 - nightly-accuracy-8-gpu-mi35x-deepseek-v32-mtp - nightly-8-gpu-mi35x-kimi-k25 - nightly-8-gpu-mi35x-qwen3-235b-mxfp4 - nightly-8-gpu-mi35x-qwen35 - nightly-8-gpu-mi35x-glm5 - nightly-8-gpu-mi35x-minimax-m25 # MI35x perf jobs excluded from check - perf failures don't block CI # - nightly-perf-8-gpu-mi35x-deepseek-v32-basic # - nightly-perf-8-gpu-mi35x-deepseek-v32-mtp runs-on: ubuntu-latest steps: - name: Check if any job failed run: | if [[ "${{ contains(needs.*.result, 'failure') }}" == "true" ]]; then echo "One or more nightly test jobs failed" exit 1 fi if [[ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]]; then echo "One or more nightly test jobs were cancelled" exit 1 fi echo "All nightly test jobs passed" ================================================ FILE: .github/workflows/nightly-test-intel.yml ================================================ name: Nightly Test (Intel) on: schedule: - cron: '0 0 * * *' push: branches: - main paths: - "python/sglang/version.py" workflow_dispatch: workflow_call: inputs: ref: description: "Branch, tag or SHA to checkout" required: false type: string default: "" concurrency: group: nightly-test-intel-${{ inputs.ref || github.ref }} cancel-in-progress: ${{ github.event_name != 'workflow_call' }} jobs: # Placeholder for Intel GPU tests # Add Intel-specific nightly test workflows here when available placeholder: if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-latest steps: - name: Placeholder run: echo "Intel nightly tests will be added here" ================================================ FILE: .github/workflows/nightly-test-npu.yml ================================================ name: Nightly Test (NPU) on: schedule: - cron: '0 17 * * *' # Execute at 1:00 a.m. Beijing Time every day pull_request: branches: - main paths: - ".github/workflows/nightly-test-npu.yml" workflow_dispatch: workflow_call: inputs: ref: description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.' required: false type: string default: '' job_filter: description: 'Select which job to run (leave empty or "all" to run all jobs)' required: false type: string default: 'all' concurrency: group: nightly-test-npu-${{ inputs.ref || github.ref }} cancel-in-progress: ${{ github.event_name != 'workflow_call' }} jobs: nightly-1-npu-a3: if: ${{ (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') }} runs-on: linux-aarch64-a3-2 strategy: fail-fast: false matrix: part: [0, 1] container: image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.5.0-a3-ubuntu22.04-py3.11 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | # speed up by using infra cache services CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple" pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn" bash scripts/ci/npu/npu_ci_install_dependency.sh a3 # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Print Log Information run: | bash scripts/ci/npu/npu_log_print.sh - name: Run test timeout-minutes: 240 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" STREAMS_PER_DEVICE: 32 run: | pip install sglang_router hf download lmms-lab/MMMU --repo-type dataset pip install sentence_transformers torchaudio==2.8.0 pip install protobuf==6.31.1 zss pre-commit wandb>=0.16.0 tenacity==8.3.0 loguru openpyxl latex2sympy2 zstandard transformers-stream-generator tqdm-multiprocess pycocoevalcap pip install yt-dlp sentencepiece==0.1.99 nltk av ftfy sqlitedict==2.1.0 sacrebleu>=1.5.0 pytablewriter black==24.1.0 isort==5.13.2 peft>=0.2.0 accelerate>=0.29.1 pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.1.25 numpy==1.26.4 dotenv git clone --branch v0.3.3 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git cd ./lmms-eval nohup pip install . > lmmslog.txt 2>&1 & sleep 120 export PYTHONPATH=$PYTHONPATH:$(pwd) cd ../ cd test python3 run_suite.py --hw npu --suite nightly-1-npu-a3 --nightly --continue-on-error --timeout-per-file 3600 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 nightly-2-npu-a3: if: ${{ (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') }} runs-on: linux-aarch64-a3-2 strategy: fail-fast: false matrix: part: [0] container: image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.5.0-a3-ubuntu22.04-py3.11 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | # speed up by using infra cache services CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple" pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn" bash scripts/ci/npu/npu_ci_install_dependency.sh a3 # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Print Log Information run: | bash scripts/ci/npu/npu_log_print.sh - name: Run test timeout-minutes: 240 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" STREAMS_PER_DEVICE: 32 run: | pip install sglang_router hf download lmms-lab/MMMU --repo-type dataset pip install sentence_transformers torchaudio==2.8.0 pip install protobuf==6.31.1 zss pre-commit wandb>=0.16.0 tenacity==8.3.0 loguru openpyxl latex2sympy2 zstandard transformers-stream-generator tqdm-multiprocess pycocoevalcap pip install yt-dlp sentencepiece==0.1.99 nltk av ftfy sqlitedict==2.1.0 sacrebleu>=1.5.0 pytablewriter black==24.1.0 isort==5.13.2 peft>=0.2.0 accelerate>=0.29.1 pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.1.25 numpy==1.26.4 dotenv git clone --branch v0.3.3 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git cd ./lmms-eval nohup pip install . > lmmslog.txt 2>&1 & sleep 120 export PYTHONPATH=$PYTHONPATH:$(pwd) cd ../ cd test python3 run_suite.py --hw npu --suite nightly-2-npu-a3 --nightly --continue-on-error --timeout-per-file 3600 --auto-partition-id ${{ matrix.part }} --auto-partition-size 1 nightly-4-npu-a3: if: ${{ (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') }} runs-on: linux-aarch64-a3-4 strategy: fail-fast: false matrix: part: [0] container: image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.5.0-a3-ubuntu22.04-py3.11 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | # speed up by using infra cache services CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple" pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn" bash scripts/ci/npu/npu_ci_install_dependency.sh a3 # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Print Log Information run: | bash scripts/ci/npu/npu_log_print.sh - name: Run test timeout-minutes: 240 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" STREAMS_PER_DEVICE: 32 run: | pip install sglang_router hf download lmms-lab/MMMU --repo-type dataset pip install sentence_transformers torchaudio==2.8.0 pip install protobuf==6.31.1 zss pre-commit wandb>=0.16.0 tenacity==8.3.0 loguru openpyxl latex2sympy2 zstandard transformers-stream-generator tqdm-multiprocess pycocoevalcap pip install yt-dlp sentencepiece==0.1.99 nltk av ftfy sqlitedict==2.1.0 sacrebleu>=1.5.0 pytablewriter black==24.1.0 isort==5.13.2 peft>=0.2.0 accelerate>=0.29.1 pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.1.25 numpy==1.26.4 dotenv git clone --branch v0.3.3 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git cd ./lmms-eval nohup pip install . > lmmslog.txt 2>&1 & sleep 120 export PYTHONPATH=$PYTHONPATH:$(pwd) cd ../ cd test python3 run_suite.py --hw npu --suite nightly-4-npu-a3 --nightly --continue-on-error --timeout-per-file 3600 --auto-partition-id ${{ matrix.part }} --auto-partition-size 1 nightly-8-npu-a3: if: ${{ (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') }} runs-on: linux-aarch64-a3-8 strategy: fail-fast: false matrix: part: [0] container: image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.5.0-a3-ubuntu22.04-py3.11 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | # speed up by using infra cache services CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple" pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn" bash scripts/ci/npu/npu_ci_install_dependency.sh a3 # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Print Log Information run: | bash scripts/ci/npu/npu_log_print.sh - name: Run test timeout-minutes: 240 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" STREAMS_PER_DEVICE: 32 run: | pip install sglang_router hf download lmms-lab/MMMU --repo-type dataset pip install sentence_transformers torchaudio==2.8.0 pip install protobuf==6.31.1 zss pre-commit wandb>=0.16.0 tenacity==8.3.0 loguru openpyxl latex2sympy2 zstandard transformers-stream-generator tqdm-multiprocess pycocoevalcap pip install yt-dlp sentencepiece==0.1.99 nltk av ftfy sqlitedict==2.1.0 sacrebleu>=1.5.0 pytablewriter black==24.1.0 isort==5.13.2 peft>=0.2.0 accelerate>=0.29.1 pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.1.25 numpy==1.26.4 dotenv git clone --branch v0.3.3 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git cd ./lmms-eval nohup pip install . > lmmslog.txt 2>&1 & sleep 120 export PYTHONPATH=$PYTHONPATH:$(pwd) cd ../ cd test python3 run_suite.py --hw npu --suite nightly-8-npu-a3 --nightly --continue-on-error --timeout-per-file 3600 --auto-partition-id ${{ matrix.part }} --auto-partition-size 1 nightly-16-npu-a3: if: ${{ (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') }} runs-on: linux-aarch64-a3-16 strategy: fail-fast: false matrix: part: [0, 1] container: image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.5.0-a3-ubuntu22.04-py3.11 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | # speed up by using infra cache services CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple" pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn" bash scripts/ci/npu/npu_ci_install_dependency.sh a3 # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Print Log Information run: | bash scripts/ci/npu/npu_log_print.sh - name: Run test timeout-minutes: 240 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" STREAMS_PER_DEVICE: 32 run: | pip install sglang_router hf download lmms-lab/MMMU --repo-type dataset pip install sentence_transformers torchaudio==2.8.0 pip install protobuf==6.31.1 zss pre-commit wandb>=0.16.0 tenacity==8.3.0 loguru openpyxl latex2sympy2 zstandard transformers-stream-generator tqdm-multiprocess pycocoevalcap pip install yt-dlp sentencepiece==0.1.99 nltk av ftfy sqlitedict==2.1.0 sacrebleu>=1.5.0 pytablewriter black==24.1.0 isort==5.13.2 peft>=0.2.0 accelerate>=0.29.1 pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.1.25 numpy==1.26.4 dotenv git clone --branch v0.3.3 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git cd ./lmms-eval nohup pip install . > lmmslog.txt 2>&1 & sleep 120 export PYTHONPATH=$PYTHONPATH:$(pwd) cd ../ cd test python3 run_suite.py --hw npu --suite nightly-16-npu-a3 --nightly --continue-on-error --timeout-per-file 3600 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 check-all-jobs: if: github.repository == 'sgl-project/sglang' && always() needs: - nightly-1-npu-a3 - nightly-2-npu-a3 - nightly-4-npu-a3 - nightly-8-npu-a3 - nightly-16-npu-a3 runs-on: ubuntu-latest container: image: docker.m.daocloud.io/ubuntu:22.04 steps: - name: Check if any job failed run: | if [[ "${{ contains(needs.*.result, 'failure') }}" == "true" ]]; then echo "One or more nightly test jobs failed" exit 1 fi if [[ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]]; then echo "One or more nightly test jobs were cancelled" exit 1 fi echo "All nightly test jobs passed" ================================================ FILE: .github/workflows/nightly-test-nvidia.yml ================================================ name: Nightly Test (Nvidia) on: schedule: - cron: '0 0 * * *' workflow_dispatch: inputs: job_filter: description: 'Select which job to run (leave empty or "all" to run all jobs)' required: false type: choice default: 'all' options: - 'all' - 'nightly-test-general-1-gpu-runner' - 'nightly-test-general-4-gpu-h100' - 'nightly-test-general-8-gpu-h200' - 'nightly-test-general-8-gpu-h20' - 'nightly-test-general-8-gpu-b200' - 'nightly-test-text-accuracy-2-gpu-runner' - 'nightly-test-text-perf-2-gpu-runner' - 'nightly-test-vlm-accuracy-2-gpu-runner' - 'nightly-test-vlm-perf-2-gpu-runner' - 'nightly-test-multimodal-server-1-gpu' - 'nightly-test-multimodal-server-2-gpu' - 'nightly-test-perf-4-gpu-b200' - 'nightly-test-perf-8-gpu-b200' workflow_call: inputs: ref: description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.' required: false type: string default: '' job_filter: description: 'Select which job to run (leave empty or "all" to run all jobs)' required: false type: string default: 'all' concurrency: group: nightly-test-nvidia-${{ inputs.ref || github.ref }} cancel-in-progress: ${{ github.event_name != 'workflow_call' }} env: SGLANG_IS_IN_CI: true SGLANG_CUDA_COREDUMP: "1" HF_HUB_DOWNLOAD_TIMEOUT: 300 HF_HUB_ETAG_TIMEOUT: 300 jobs: # General tests - 1 GPU nightly-test-general-1-gpu-runner: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-general-1-gpu-runner') runs-on: 1-gpu-runner steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh - name: Run test timeout-minutes: 60 run: | cd test python3 run_suite.py --hw cuda --suite nightly-1-gpu --nightly --continue-on-error - uses: ./.github/actions/upload-cuda-coredumps if: always() # General tests - 4 GPU H100 nightly-test-general-4-gpu-h100: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-general-4-gpu-h100') runs-on: 4-gpu-h100 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | cd test python3 run_suite.py --hw cuda --suite nightly-4-gpu --nightly --continue-on-error - uses: ./.github/actions/upload-cuda-coredumps if: always() # General tests - 8 GPU H200 nightly-test-general-8-gpu-h200: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-general-8-gpu-h200') runs-on: 8-gpu-h200 strategy: fail-fast: false matrix: partition: [0, 1, 2, 3] env: RUNNER_LABELS: 8-gpu-h200 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh - name: Run common 8-GPU model tests if: always() timeout-minutes: 300 env: TRACE_BASE_URL: https://raw.githubusercontent.com/sglang-bot/sglang-ci-data/main/traces/${{ github.run_id }} PERFETTO_RELAY_URL: ${{ vars.PERFETTO_RELAY_URL }} GPU_CONFIG: "8-gpu-h200" IS_H200: "1" run: | cd test python3 run_suite.py --hw cuda --suite nightly-8-gpu-common --nightly --timeout-per-file=18000 --continue-on-error --auto-partition-id=${{ matrix.partition }} --auto-partition-size=4 - name: Publish traces to storage repo if: always() continue-on-error: true env: GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} GITHUB_RUN_ID: ${{ github.run_id }} GITHUB_RUN_NUMBER: ${{ github.run_number }} run: | TRACE_ARGS="" for dir in test/performance_profiles_*/; do [ -d "$dir" ] && TRACE_ARGS="$TRACE_ARGS --traces-dir $dir" done if [ -n "$TRACE_ARGS" ]; then python3 scripts/ci/utils/publish_traces.py $TRACE_ARGS find test/performance_profiles_*/ -name '*.json.gz' -delete else echo "No trace directories found, skipping publish" fi - name: Run test timeout-minutes: 30 env: GPU_CONFIG: "8-gpu-h200" run: | cd test python3 run_suite.py --hw cuda --suite nightly-8-gpu-h200 --nightly --continue-on-error - name: Collect performance metrics if: always() run: | python3 scripts/ci/utils/save_metrics.py \ --gpu-config 8-gpu-h200 \ --partition ${{ matrix.partition }} \ --run-id ${{ github.run_id }} \ --output test/metrics-8gpu-h200-partition-${{ matrix.partition }}.json \ --search-dir test/performance_profiles_8_gpu \ --search-dir test - name: Upload partition metrics if: always() uses: actions/upload-artifact@v4 with: name: metrics-8gpu-h200-partition-${{ matrix.partition }} path: test/metrics-8gpu-h200-partition-${{ matrix.partition }}.json retention-days: 5 if-no-files-found: ignore - uses: ./.github/actions/upload-cuda-coredumps if: always() with: artifact-suffix: ${{ matrix.partition }} # General tests - 8 GPU H20 nightly-test-general-8-gpu-h20: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-general-8-gpu-h20') runs-on: 8-gpu-h20 env: SGLANG_CI_RDMA_ALL_DEVICES: "mlx5_1,mlx5_2,mlx5_3,mlx5_4" steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh - name: Run test timeout-minutes: 30 env: GPU_CONFIG: "8-gpu-h20" run: | cd test python3 run_suite.py --hw cuda --suite nightly-8-gpu-h20 --nightly --continue-on-error - uses: ./.github/actions/upload-cuda-coredumps if: always() # General tests - 8 GPU B200 nightly-test-general-8-gpu-b200: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-general-8-gpu-b200') runs-on: 8-gpu-b200 strategy: fail-fast: false matrix: partition: [0, 1, 2, 3] steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | IS_BLACKWELL=1 bash scripts/ci/cuda/ci_install_dependency.sh - name: Run common 8-GPU model tests if: always() timeout-minutes: 300 env: TRACE_BASE_URL: https://raw.githubusercontent.com/sglang-bot/sglang-ci-data/main/traces/${{ github.run_id }} PERFETTO_RELAY_URL: ${{ vars.PERFETTO_RELAY_URL }} GPU_CONFIG: "8-gpu-b200" run: | cd test IS_BLACKWELL=1 python3 run_suite.py --hw cuda --suite nightly-8-gpu-common --nightly --timeout-per-file=12000 --continue-on-error --auto-partition-id=${{ matrix.partition }} --auto-partition-size=4 - name: Publish traces to storage repo if: always() continue-on-error: true env: GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} GITHUB_RUN_ID: ${{ github.run_id }} GITHUB_RUN_NUMBER: ${{ github.run_number }} run: | TRACE_ARGS="" for dir in test/performance_profiles_*/; do [ -d "$dir" ] && TRACE_ARGS="$TRACE_ARGS --traces-dir $dir" done if [ -n "$TRACE_ARGS" ]; then python3 scripts/ci/utils/publish_traces.py $TRACE_ARGS find test/performance_profiles_*/ -name '*.json.gz' -delete else echo "No trace directories found, skipping publish" fi - name: Collect performance metrics if: always() run: | python3 scripts/ci/utils/save_metrics.py \ --gpu-config 8-gpu-b200 \ --partition ${{ matrix.partition }} \ --run-id ${{ github.run_id }} \ --output test/metrics-8gpu-b200-partition-${{ matrix.partition }}.json \ --search-dir test/performance_profiles_8_gpu \ --search-dir test - name: Upload partition metrics if: always() uses: actions/upload-artifact@v4 with: name: metrics-8gpu-b200-partition-${{ matrix.partition }} path: test/metrics-8gpu-b200-partition-${{ matrix.partition }}.json retention-days: 5 if-no-files-found: ignore - uses: ./.github/actions/upload-cuda-coredumps if: always() with: artifact-suffix: ${{ matrix.partition }} # Text model accuracy tests nightly-test-text-accuracy-2-gpu-runner: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-text-accuracy-2-gpu-runner') runs-on: 2-gpu-runner steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh - name: Run eval test for text models timeout-minutes: 120 run: | cd test python3 run_suite.py --hw cuda --suite nightly-eval-text-2-gpu --nightly --continue-on-error --timeout-per-file 4500 - uses: ./.github/actions/upload-cuda-coredumps if: always() # Text model performance tests nightly-test-text-perf-2-gpu-runner: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-text-perf-2-gpu-runner') runs-on: 2-gpu-runner steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh - name: Run performance test for text models timeout-minutes: 180 env: TRACE_BASE_URL: https://raw.githubusercontent.com/sglang-bot/sglang-ci-data/main/traces/${{ github.run_id }} PERFETTO_RELAY_URL: ${{ vars.PERFETTO_RELAY_URL }} GPU_CONFIG: "2-gpu-runner" run: | cd test rm -rf performance_profiles_text_models/ python3 run_suite.py --hw cuda --suite nightly-perf-text-2-gpu --nightly --continue-on-error --timeout-per-file 3600 - name: Publish traces to storage repo env: GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} GITHUB_RUN_ID: ${{ github.run_id }} GITHUB_RUN_NUMBER: ${{ github.run_number }} run: | python3 scripts/ci/utils/publish_traces.py --traces-dir test/performance_profiles_text_models - uses: ./.github/actions/upload-cuda-coredumps if: always() # VLM accuracy tests nightly-test-vlm-accuracy-2-gpu-runner: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-vlm-accuracy-2-gpu-runner') runs-on: 2-gpu-runner steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh - name: Run eval test for VLM models (fixed MMMU-100) timeout-minutes: 240 run: | cd test python3 run_suite.py --hw cuda --suite nightly-eval-vlm-2-gpu --nightly --continue-on-error --timeout-per-file 9000 - uses: ./.github/actions/upload-cuda-coredumps if: always() # VLM performance tests nightly-test-vlm-perf-2-gpu-runner: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-vlm-perf-2-gpu-runner') runs-on: 2-gpu-runner steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh - name: Run perf test for VLM models (MMMU) timeout-minutes: 240 env: TRACE_BASE_URL: https://raw.githubusercontent.com/sglang-bot/sglang-ci-data/main/traces/${{ github.run_id }} PERFETTO_RELAY_URL: ${{ vars.PERFETTO_RELAY_URL }} GPU_CONFIG: "2-gpu-runner" run: | cd test rm -rf performance_profiles_vlms/ python3 run_suite.py --hw cuda --suite nightly-perf-vlm-2-gpu --nightly --continue-on-error --timeout-per-file 3600 - name: Publish traces to storage repo env: GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }} GITHUB_RUN_ID: ${{ github.run_id }} GITHUB_RUN_NUMBER: ${{ github.run_number }} run: | python3 scripts/ci/utils/publish_traces.py --traces-dir test/performance_profiles_vlms - uses: ./.github/actions/upload-cuda-coredumps if: always() # diffusion performance tests nightly-test-multimodal-server-1-gpu: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-multimodal-server-1-gpu') runs-on: 1-gpu-runner strategy: fail-fast: false max-parallel: 5 matrix: part: [0, 1] steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh diffusion pip install slack_sdk - name: Run diffusion server tests env: SGLANG_DIFFUSION_SLACK_TOKEN: ${{ secrets.SGLANG_DIFFUSION_SLACK_TOKEN }} GITHUB_RUN_ID: ${{ github.run_id }} GPU_CONFIG: "1-gpu-runner" timeout-minutes: 60 run: | cd python python3 sglang/multimodal_gen/test/run_suite.py \ --suite 1-gpu \ --partition-id ${{ matrix.part }} \ --total-partitions 2 - name: Collect diffusion performance metrics if: always() run: | python3 scripts/ci/utils/save_diffusion_metrics.py \ --gpu-config 1-gpu-runner \ --run-id ${{ github.run_id }} \ --output python/diffusion-metrics-1gpu-partition-${{ matrix.part }}.json \ --results-json python/diffusion-results.json - name: Upload diffusion metrics if: always() uses: actions/upload-artifact@v4 with: name: diffusion-metrics-1gpu-partition-${{ matrix.part }} path: python/diffusion-metrics-1gpu-partition-${{ matrix.part }}.json retention-days: 90 if-no-files-found: ignore - uses: ./.github/actions/upload-cuda-coredumps if: always() with: artifact-suffix: ${{ matrix.part }} nightly-test-multimodal-server-2-gpu: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-multimodal-server-2-gpu') runs-on: 2-gpu-runner strategy: fail-fast: false max-parallel: 5 matrix: part: [0, 1] steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh diffusion pip install slack_sdk - name: Run diffusion server tests env: SGLANG_DIFFUSION_SLACK_TOKEN: ${{ secrets.SGLANG_DIFFUSION_SLACK_TOKEN }} GITHUB_RUN_ID: ${{ github.run_id }} GPU_CONFIG: "2-gpu-runner" timeout-minutes: 60 run: | cd python python3 sglang/multimodal_gen/test/run_suite.py \ --suite 2-gpu \ --partition-id ${{ matrix.part }} \ --total-partitions 2 - name: Collect diffusion performance metrics if: always() run: | python3 scripts/ci/utils/save_diffusion_metrics.py \ --gpu-config 2-gpu-runner \ --run-id ${{ github.run_id }} \ --output python/diffusion-metrics-2gpu-partition-${{ matrix.part }}.json \ --results-json python/diffusion-results.json - name: Upload diffusion metrics if: always() uses: actions/upload-artifact@v4 with: name: diffusion-metrics-2gpu-partition-${{ matrix.part }} path: python/diffusion-metrics-2gpu-partition-${{ matrix.part }}.json retention-days: 90 if-no-files-found: ignore - uses: ./.github/actions/upload-cuda-coredumps if: always() with: artifact-suffix: ${{ matrix.part }} # B200 Performance tests - 4 GPU nightly-test-perf-4-gpu-b200: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-perf-4-gpu-b200') runs-on: 4-gpu-b200 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | IS_BLACKWELL=1 bash scripts/ci/cuda/ci_install_dependency.sh - name: Run test timeout-minutes: 300 run: | cd test python3 run_suite.py --hw cuda --suite nightly-4-gpu-b200 --nightly --continue-on-error --timeout-per-file 12000 - uses: ./.github/actions/upload-cuda-coredumps if: always() # Specialized B200 tests - 8 GPU, for specific backends and configs nightly-test-specialized-8-gpu-b200: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-perf-8-gpu-b200') runs-on: 8-gpu-b200 env: RUNNER_LABELS: 8-gpu-b200 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies run: | IS_BLACKWELL=1 bash scripts/ci/cuda/ci_install_dependency.sh - name: Run test timeout-minutes: 120 env: GPU_CONFIG: "8-gpu-b200" run: | cd test python3 run_suite.py --hw cuda --suite nightly-8-gpu-b200 --nightly --continue-on-error --timeout-per-file 2400 - uses: ./.github/actions/upload-cuda-coredumps if: always() # Consolidate performance metrics from all jobs consolidate-metrics: if: github.repository == 'sgl-project/sglang' && always() needs: - nightly-test-general-8-gpu-h200 - nightly-test-general-8-gpu-b200 - nightly-test-multimodal-server-1-gpu - nightly-test-multimodal-server-2-gpu runs-on: ubuntu-latest steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Download all partition metrics uses: actions/download-artifact@v4 with: pattern: "*metrics-*" path: metrics/ merge-multiple: true - name: List downloaded metrics run: | echo "Downloaded metrics files:" find metrics/ -name "*.json" -type f 2>/dev/null || echo "No metrics files found" - name: Merge metrics run: | python3 scripts/ci/utils/merge_metrics.py \ --input-dir metrics/ \ --output consolidated-metrics-${{ github.run_id }}.json \ --run-id ${{ github.run_id }} \ --commit-sha ${{ github.sha }} \ --branch ${{ github.ref_name }} - name: Upload consolidated metrics uses: actions/upload-artifact@v4 with: name: consolidated-metrics-${{ github.run_id }} path: consolidated-metrics-${{ github.run_id }}.json retention-days: 90 if-no-files-found: warn # Final check job check-all-jobs: if: github.repository == 'sgl-project/sglang' && always() needs: - nightly-test-general-1-gpu-runner - nightly-test-general-4-gpu-h100 - nightly-test-general-8-gpu-h200 - nightly-test-general-8-gpu-h20 - nightly-test-general-8-gpu-b200 - nightly-test-text-accuracy-2-gpu-runner - nightly-test-text-perf-2-gpu-runner - nightly-test-vlm-accuracy-2-gpu-runner - nightly-test-vlm-perf-2-gpu-runner - nightly-test-multimodal-server-1-gpu - nightly-test-multimodal-server-2-gpu - nightly-test-perf-4-gpu-b200 - nightly-test-specialized-8-gpu-b200 - consolidate-metrics runs-on: ubuntu-latest steps: - name: Check if any job failed run: | if [[ "${{ contains(needs.*.result, 'failure') }}" == "true" ]]; then echo "One or more nightly test jobs failed" exit 1 fi if [[ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]]; then echo "One or more nightly test jobs were cancelled" exit 1 fi echo "All nightly test jobs passed" ================================================ FILE: .github/workflows/open-pr-copy-from-oss.yml ================================================ name: Open A PR to Copy Code From OSS on: workflow_dispatch: # schedule: # - cron: '0 10 * * *' permissions: contents: write jobs: copy: runs-on: ubuntu-latest steps: - name: Checkout repository uses: actions/checkout@v4 with: ref: 'main' - name: Install GitHub CLI (if not present) run: | bash scripts/code_sync/install_github_cli.sh - name: Copy from OSS code env: GH_TOKEN: ${{ secrets.GH_PAT_FOR_OPEN_PR_TO_PRIVATE }} run: | python3 scripts/code_sync/copy_from_oss.py ================================================ FILE: .github/workflows/open-pr-copy-to-oss.yml ================================================ name: Open A PR to Copy Diff To OSS on: workflow_dispatch: inputs: commit_sha: description: 'The commit SHA to copy. Defaults to LAST to copy the latest commit.' required: false default: 'LAST' permissions: contents: write jobs: copy: runs-on: ubuntu-latest steps: - name: Checkout repository uses: actions/checkout@v4 with: fetch-depth: 0 - name: Install GitHub CLI (if not present) run: | bash scripts/code_sync/install_github_cli.sh - name: Copy to OSS code env: GH_TOKEN: ${{ secrets.GH_PAT_FOR_OPEN_PR_TO_OSS }} run: | python3 scripts/code_sync/copy_to_oss.py --commit ${{ github.event.inputs.commit_sha }} ================================================ FILE: .github/workflows/patch-docker-dev.yml ================================================ name: Patch Docker Image on: workflow_dispatch: inputs: pr_numbers: description: "Comma-separated PR numbers to apply (e.g. 18962,19010)" required: false default: "" image_tag: description: "Base image tag to patch (e.g. dev-x86, dev-x86-cu13)" required: true concurrency: group: patch-docker-${{ inputs.image_tag }} cancel-in-progress: true jobs: patch: if: github.repository == 'sgl-project/sglang' runs-on: x64-docker-build-node steps: - name: Checkout repository uses: actions/checkout@v4 with: fetch-depth: 0 - name: Login to Docker Hub uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Pull base image and extract commit run: | IMAGE="lmsysorg/sglang:${{ inputs.image_tag }}" docker pull "${IMAGE}" if BASE_SHA=$(docker run --rm "${IMAGE}" git -C /sgl-workspace/sglang rev-parse HEAD 2>/dev/null); then echo "Image built from commit: ${BASE_SHA}" else BASE_SHA="" echo "::warning::Image has no .git directory — cannot extract base commit" fi echo "BASE_SHA=${BASE_SHA}" >> "$GITHUB_ENV" - name: Generate patches run: | git config --global --add safe.directory "$GITHUB_WORKSPACE" git fetch origin main mkdir -p /tmp/patch-ctx if [ -n "${{ inputs.pr_numbers }}" ]; then IFS=',' read -ra PRS <<< "${{ inputs.pr_numbers }}" for pr in "${PRS[@]}"; do pr=$(echo "${pr}" | xargs) echo "Fetching PR #${pr}" git fetch origin "pull/${pr}/head:pr-${pr}" MERGE_BASE=$(git merge-base origin/main "pr-${pr}") echo " PR #${pr}: merge-base=${MERGE_BASE}" git diff "${MERGE_BASE}..pr-${pr}" > "/tmp/patch-ctx/${pr}.patch" echo " PR #${pr}: $(wc -l < /tmp/patch-ctx/${pr}.patch) lines" done elif [ -n "${BASE_SHA}" ]; then echo "Generating diff: image ${BASE_SHA} → latest main" git fetch origin "${BASE_SHA}" git diff "${BASE_SHA}..origin/main" > /tmp/patch-ctx/main.patch echo " main: $(wc -l < /tmp/patch-ctx/main.patch) lines" else echo "::error::No PR numbers specified and image has no .git — cannot generate diff against main" exit 1 fi TOTAL=$(cat /tmp/patch-ctx/*.patch | wc -l) if [ "${TOTAL}" -eq 0 ]; then echo "::warning::All patches are empty — image is already up to date" echo "SKIP_BUILD=true" >> "$GITHUB_ENV" fi - name: Build patched image if: env.SKIP_BUILD != 'true' run: | IMAGE="lmsysorg/sglang:${{ inputs.image_tag }}" cat <<'DOCKERFILE' > /tmp/patch-ctx/Dockerfile ARG BASE_IMAGE FROM ${BASE_IMAGE} COPY *.patch /tmp/patches/ RUN cd /sgl-workspace/sglang \ && for p in /tmp/patches/*.patch; do \ if [ ! -s "${p}" ]; then \ echo "Skipping ${p} (empty)"; \ else \ echo "Applying ${p}..." \ && patch -p1 --fuzz=2 --no-backup-if-mismatch -f < "${p}" \ || { echo "ERROR: Failed to apply ${p}"; exit 1; }; \ fi; \ done \ && rm -rf /tmp/patches DOCKERFILE docker build \ --no-cache \ --build-arg BASE_IMAGE="${IMAGE}" \ -t "${IMAGE}" \ /tmp/patch-ctx/ - name: Push patched image if: env.SKIP_BUILD != 'true' run: | IMAGE="lmsysorg/sglang:${{ inputs.image_tag }}" docker push "${IMAGE}" echo "### Patched \`${IMAGE}\`" >> "$GITHUB_STEP_SUMMARY" echo "- **Base commit:** \`${BASE_SHA:-unknown (no .git)}\`" >> "$GITHUB_STEP_SUMMARY" echo "- **Source:** ${{ inputs.pr_numbers && format('PRs: {0}', inputs.pr_numbers) || 'latest main' }}" >> "$GITHUB_STEP_SUMMARY" ================================================ FILE: .github/workflows/pr-benchmark-rust.yml ================================================ name: PR Benchmark (SMG Components) on: push: branches: [ main ] paths: - "sgl-model-gateway/**" pull_request: branches: [ main ] paths: - "sgl-model-gateway/**" workflow_dispatch: concurrency: group: pr-benchmark-rust-${{ github.ref }} cancel-in-progress: true env: RUSTC_WRAPPER: sccache SCCACHE_GHA_ENABLED: "true" permissions: contents: read pull-requests: write issues: write jobs: benchmark-compile-check: name: Benchmark Compilation Check runs-on: ubuntu-latest steps: - name: Checkout code uses: actions/checkout@v4 - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_gateway_dependencies.sh - name: Configure sccache uses: mozilla-actions/sccache-action@v0.0.9 with: version: "v0.12.0" disable_annotations: true - name: Rust cache uses: Swatinem/rust-cache@v2 with: workspaces: sgl-model-gateway shared-key: "rust-cache" save-if: true cache-all-crates: true cache-on-failure: true - name: Check benchmarks compile run: | source "$HOME/.cargo/env" cd sgl-model-gateway/ cargo check --benches - name: Show sccache stats if: always() run: sccache --show-stats benchmark: name: Benchmark - ${{ matrix.name }} if: | github.repository == 'sgl-project/sglang' && (github.event_name == 'push' || github.event_name == 'workflow_dispatch' || (contains(github.event.pull_request.labels.*.name, 'router-benchmark') && contains(github.event.pull_request.labels.*.name, 'run-ci'))) strategy: fail-fast: false matrix: include: - name: Request Processing bench_name: request_processing bench_args: "benchmark_summary --exact" runner: ubuntu-latest sccache_version: "v0.12.0" artifact_name: request-processing-results artifact_path: criterion/benchmark_summary/ - name: Manual Policy bench_name: manual_policy_benchmark bench_args: "" runner: ubuntu-latest sccache_version: "v0.12.0" artifact_name: manual-policy-results artifact_path: criterion/manual_policy*/ runs-on: ${{ matrix.runner }} steps: - name: Checkout code uses: actions/checkout@v4 with: fetch-depth: 100 - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_gateway_dependencies.sh - name: Configure sccache uses: mozilla-actions/sccache-action@v0.0.9 with: version: ${{ matrix.sccache_version }} disable_annotations: true - name: Rust cache uses: Swatinem/rust-cache@v2 with: workspaces: sgl-model-gateway shared-key: "rust-cache" cache-all-crates: true cache-on-failure: true save-if: true - name: Run benchmark timeout-minutes: 30 run: | source "$HOME/.cargo/env" cd sgl-model-gateway/ if command -v sccache &> /dev/null; then echo "Testing sccache availability..." export RUSTC_WRAPPER=sccache export SCCACHE_GHA_ENABLED="true" if sccache --start-server 2>/dev/null && sccache --show-stats 2>/dev/null; then echo "sccache is working, using it for compilation" else echo "sccache failed to start, falling back to regular cargo" unset RUSTC_WRAPPER unset SCCACHE_GHA_ENABLED fi else echo "sccache not available, using regular cargo" fi cargo bench --bench ${{ matrix.bench_name }} -- ${{ matrix.bench_args }} 2>&1 | tee benchmark_output.txt - name: Upload benchmark results if: always() uses: actions/upload-artifact@v4 with: name: ${{ matrix.artifact_name }}-${{ github.sha }} path: | sgl-model-gateway/target/${{ matrix.artifact_path }} sgl-model-gateway/benchmark_output.txt retention-days: 30 - name: Show sccache stats if: always() run: sccache --show-stats benchmark-summary: name: Benchmark Summary needs: [benchmark] if: always() && (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') runs-on: ubuntu-latest steps: - name: Download all benchmark results uses: actions/download-artifact@v4 with: pattern: '*-results-${{ github.sha }}' path: benchmark-results - name: Generate summary run: | generate_section() { local title="$1" dir_name="$2" lines="${3:-100}" local dir="benchmark-results/${dir_name}-${{ github.sha }}" echo "### $title" >> summary.md if [ -d "$dir" ]; then echo "✅ **Completed**" >> summary.md if [ -f "$dir/benchmark_output.txt" ]; then echo -e "\n
\nView Results\n\n\`\`\`" >> summary.md tail -"$lines" "$dir/benchmark_output.txt" >> summary.md echo -e "\`\`\`\n
" >> summary.md fi else echo "❌ Failed or skipped" >> summary.md fi echo "" >> summary.md } echo "## 🚀 Benchmark Results Summary" > summary.md echo "" >> summary.md generate_section "Request Processing" "request-processing-results" 60 generate_section "Manual Policy (Sticky Sessions)" "manual-policy-results" 100 echo -e "---\n_Generated at $(date -u '+%Y-%m-%d %H:%M:%S UTC')_" >> summary.md cat summary.md cat summary.md >> $GITHUB_STEP_SUMMARY - name: Upload summary uses: actions/upload-artifact@v4 with: name: benchmark-summary-${{ github.sha }} path: summary.md retention-days: 30 ================================================ FILE: .github/workflows/pr-gate.yml ================================================ on: workflow_call: inputs: require-run-ci: description: "Whether the PR must have the run-ci label" type: boolean default: true cool-down-minutes: description: "Cooldown period in minutes for low-permission users; 0 disables rate limiting" type: number default: 120 jobs: pr-gate: # 1. for commits on main: no gating needed # 2. for workflow_dispatch: this can only be triggered by users with write access runs-on: ubuntu-latest steps: - name: Fetch latest PR info if: github.event_name == 'pull_request' id: pr uses: actions/github-script@v7 with: github-token: ${{ secrets.GITHUB_TOKEN }} script: | const pr = await github.rest.pulls.get({ owner: context.repo.owner, repo: context.repo.repo, pull_number: context.issue.number }); core.setOutput("labels", JSON.stringify(pr.data.labels.map(l => l.name))); core.setOutput("draft", pr.data.draft); core.setOutput("user", pr.data.user.login); - name: Log PR info if: github.event_name == 'pull_request' run: | echo "===== PR Info =====" echo "PR Event: ${{ github.event_name }}" echo "PR Labels: ${{ steps.pr.outputs.labels }}" echo "PR Draft: ${{ steps.pr.outputs.draft }}" echo "PR User: ${{ steps.pr.outputs.user }}" echo "Require run-ci: ${{ inputs.require-run-ci }}" echo "Cool down minutes: ${{ inputs.cool-down-minutes }}" echo "===================" - name: Block draft PR if: github.event_name == 'pull_request' && fromJson(steps.pr.outputs.draft) run: | echo "PR is draft. Blocking CI." exit 1 - name: Require run-ci label (optional) if: github.event_name == 'pull_request' && inputs.require-run-ci == true run: | labels='${{ steps.pr.outputs.labels }}' if [[ "${{ contains(fromJson(steps.pr.outputs.labels), 'run-ci') }}" == "false" ]]; then echo "Missing required label 'run-ci'. See https://docs.sglang.io/developer_guide/contribution_guide.html#how-to-trigger-ci-tests for more details." exit 1 fi - name: Enforce rate limit for low-permission actors (optional) if: github.event_name == 'pull_request' && inputs.cool-down-minutes > 0 uses: actions/github-script@v7 with: github-token: ${{ secrets.GITHUB_TOKEN }} script: | const DEFAULT_MINUTES = Number("${{ inputs.cool-down-minutes }}"); const owner = context.repo.owner; const repo = context.repo.repo; const eventName = context.eventName; const curRun = await github.rest.actions.getWorkflowRun({ owner, repo, run_id: context.runId }); let triggeringActor = curRun.data.triggering_actor?.login || context.actor; if (triggeringActor === "github-actions[bot]") { triggeringActor = `${{ steps.pr.outputs.user }}`; core.info( `triggering_actor is github-actions[bot]; substituting PR author '${triggeringActor}'.` ); } async function hasHighPermission(username) { try { const { data } = await github.rest.repos.getCollaboratorPermissionLevel({ owner, repo, username }); const perm = data.permission || 'none'; return perm === 'write' || perm === 'maintain' || perm === 'admin'; } catch (e) { if (e.status === 404 || e.status === 403) return false; throw e; } } if (await hasHighPermission(triggeringActor)) { core.info(`Triggering user '${triggeringActor}' has high permission. No rate limit applied.`); return; } let effectiveCooldownMinutes = DEFAULT_MINUTES; let perUserCooldownMinutes = null; try { const contentResp = await github.rest.repos.getContent({ owner, repo, path: ".github/CI_PERMISSIONS.json", ref: "main", }); if (!Array.isArray(contentResp.data) && contentResp.data && "content" in contentResp.data) { const raw = Buffer.from( contentResp.data.content, contentResp.data.encoding || "base64" ).toString(); const ciPermissions = JSON.parse(raw); const userPerm = ciPermissions[triggeringActor]; if (userPerm && typeof userPerm.cooldown_interval_minutes === "number") { perUserCooldownMinutes = userPerm.cooldown_interval_minutes; core.info( `Per-user cooldown for '${triggeringActor}' from CI_PERMISSIONS.json: ${perUserCooldownMinutes} minutes.` ); } else { core.info(`No per-user cooldown found for '${triggeringActor}' in CI_PERMISSIONS.json.`); } } else { core.info("CI_PERMISSIONS.json content response is not a file; skipping per-user cooldown."); } } catch (e) { core.info(`CI_PERMISSIONS.json not found or unreadable: ${e.message}. Using default rate limit only.`); } if (perUserCooldownMinutes !== null) { effectiveCooldownMinutes = Math.min(effectiveCooldownMinutes, perUserCooldownMinutes); } if (effectiveCooldownMinutes <= 0) { core.info( `Effective cooldown for '${triggeringActor}' is 0 minutes; no rate limit enforced for this user.` ); return; } const cutoff = new Date(Date.now() - effectiveCooldownMinutes * 60 * 1000); core.info( `Checking for workflow runs since ${cutoff.toISOString()} (last ${effectiveCooldownMinutes} minutes) for event '${eventName}'.` ); const { data } = await github.rest.actions.listWorkflowRuns({ owner, repo, workflow_id: 'pr-test.yml', event: eventName, per_page: 100, }); const runs = data.workflow_runs || []; // Rate Limiting Logic: // We only count workflow runs that actually consumed CI resources (i.e., passed the gate). // A run "passes the gate" if any jobs beyond the gate jobs (check-changes, pr-gate, call-gate) // actually executed (not skipped/cancelled). This prevents scenarios where: // - User has PR A with missing 'run-ci' label (fails at gate) // - User opens PR B with 'run-ci' label // - PR B should be able to run even though PR A triggered a run recently // Helper function to check if a run passed the gate (i.e., actually consumed CI resources) async function didRunPassGate(run) { try { // Note: Fetching up to 100 jobs (API maximum). If a workflow has >100 jobs, // we may miss some, but this is unlikely in practice. const { data: jobsData } = await github.rest.actions.listJobsForWorkflowRun({ owner, repo, run_id: run.id, per_page: 100 }); const jobs = jobsData.jobs || []; // If no jobs exist yet, the run hasn't started consuming resources if (jobs.length === 0) { core.info(`Run ${run.id} has no jobs yet; not counting against rate limit.`); return false; } // Gate jobs that don't consume significant CI resources const gateJobs = ['check-changes', 'pr-gate', 'call-gate', 'pr-test-finish']; const jobsBeyondGate = jobs.filter(j => !gateJobs.some(g => j.name === g || j.name.startsWith(g + ' '))); // A job "ran" if it reached a terminal conclusion state that indicates actual execution const ranStates = ['success', 'failure', 'timed_out', 'action_required']; const hasJobsThatRan = jobsBeyondGate.some(j => j.conclusion && ranStates.includes(j.conclusion)); return hasJobsThatRan; } catch (e) { core.warning(`Could not check jobs for run ${run.id}: ${e.message}`); // If it's a rate limit error, count it conservatively to prevent abuse if (e.status === 429) { core.warning(`Hit rate limit checking run ${run.id}; counting it to be safe.`); return true; } // For cancelled/skipped runs, they likely didn't consume resources if (run.conclusion === 'cancelled' || run.conclusion === 'skipped') { return false; } // Default to counting it to prevent abuse return true; } } // Limit the number of runs we'll check in detail to avoid API rate limits const MAX_RUNS_TO_CHECK = 5; let runsChecked = 0; let runsSkippedAtGate = 0; let recentFound = null; for (const run of runs) { if (String(run.id) === String(context.runId)) continue; if (new Date(run.created_at) < cutoff) continue; const isUserRun = (run.actor?.login === triggeringActor) || (run.triggering_actor?.login === triggeringActor); if (!isUserRun) continue; runsChecked++; core.info(`Checking run ${run.id} (created: ${run.created_at}, conclusion: ${run.conclusion})`); // Safety limit: if we've checked too many runs, assume the next one passed to be conservative if (runsChecked > MAX_RUNS_TO_CHECK) { core.warning(`Checked ${MAX_RUNS_TO_CHECK} runs; assuming this one passed gate to avoid API limits.`); recentFound = run; break; } // Only count runs that actually passed the gate and consumed CI resources if (await didRunPassGate(run)) { recentFound = run; core.info(`Found recent run ${run.id} that passed gate.`); break; } else { runsSkippedAtGate++; core.info(`Run ${run.id} failed at gate; not counting against rate limit.`); } } core.info(`Rate limit check summary: checked ${runsChecked} runs, ${runsSkippedAtGate} failed at gate.`); if (recentFound) { core.setFailed( `User '${triggeringActor}' already triggered '${context.workflow}' via '${eventName}' at ${recentFound.created_at}. ` + `Please wait ${effectiveCooldownMinutes} minutes before triggering again.` ); } else { core.info( `No recent runs detected for '${triggeringActor}' within the last ${effectiveCooldownMinutes} minutes; proceeding.` ); } ================================================ FILE: .github/workflows/pr-test-amd-rocm720.yml ================================================ name: PR Test ROCm 7.2 (AMD) # Dynamic run-name for /rerun-stage commands to enable URL lookup # Format: "[stage-name] sha" for fork PRs, "[stage-name]" for non-fork, default for normal runs run-name: ${{ (inputs.target_stage || inputs.target_stage_select) && (inputs.pr_head_sha && format('[{0}] {1}', inputs.target_stage || inputs.target_stage_select, inputs.pr_head_sha) || format('[{0}]', inputs.target_stage || inputs.target_stage_select)) || '' }} on: schedule: - cron: '30 17 * * *' # push: # branches: [ main ] # paths: # - "python/**" # - "scripts/ci/**" # - "test/**" # - "sgl-kernel/**" # - ".github/workflows/pr-test-amd-rocm720.yml" # - "docker/rocm.Dockerfile" # pull_request: # branches: [ main ] # paths: # - "python/**" # - "scripts/ci/**" # - "test/**" # - "sgl-kernel/**" # - ".github/workflows/pr-test-amd-rocm720.yml" # - "docker/rocm.Dockerfile" workflow_dispatch: inputs: target_stage_select: description: "Select a stage to run from dropdown (leave empty for auto-detect)" required: false type: choice default: '' options: - '' - sgl-kernel-unit-test-amd - sgl-kernel-unit-test-2-gpu-amd - stage-a-test-1-amd - jit-kernel-unit-test-amd - stage-b-test-small-1-gpu-amd - stage-b-test-small-1-gpu-amd-nondeterministic - stage-b-test-small-1-gpu-amd-mi35x - stage-b-test-large-1-gpu-amd - stage-b-test-large-2-gpu-amd - multimodal-gen-test-1-gpu-amd - multimodal-gen-test-2-gpu-amd - stage-c-test-large-8-gpu-amd - stage-c-test-large-8-gpu-amd-mi35x - stage-b-test-large-8-gpu-disaggregation-amd target_stage: description: "Or type comma-separated stage names (overrides dropdown if non-empty)" required: false type: string default: "" pr_head_sha: description: "PR head SHA to checkout (for /rerun-stage on fork PRs)" required: false type: string default: "" aiter_ref: description: 'Override AITER commit (optional, leave empty to use Dockerfile default)' required: false type: string default: '' continue_on_error: description: 'Continue on error (do not fail the workflow on test failures)' required: false type: boolean default: true workflow_call: inputs: ref: description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.' required: false type: string default: '' run_all_tests: description: "Run all tests (for releasing or testing purpose)" required: false type: boolean default: false aiter_ref: description: 'Override AITER commit (optional, leave empty to use Dockerfile default)' required: false type: string default: '' continue_on_error: description: 'Continue on error (do not fail the workflow on test failures)' required: false type: boolean default: true env: AITER_COMMIT_OVERRIDE: ${{ inputs.aiter_ref }} concurrency: # When called via workflow_call with run_all_tests=true, use a unique group per run to # avoid collisions with direct schedule/workflow_dispatch triggers. We use run_all_tests # (not github.event_name) to detect this, because github.event_name inherits from the caller. group: pr-test-amd-rocm720-${{ inputs.run_all_tests && format('full-{0}', github.run_id) || inputs.pr_head_sha || inputs.ref || github.ref }} cancel-in-progress: ${{ !inputs.run_all_tests && github.event_name != 'workflow_call' }} jobs: call-gate: uses: ./.github/workflows/pr-gate.yml secrets: inherit check-changes: needs: [call-gate] runs-on: ubuntu-latest outputs: main_package: ${{ steps.filter.outputs.main_package || steps.run-mode.outputs.run_all_tests }} sgl_kernel: ${{ steps.filter.outputs.sgl_kernel || steps.run-mode.outputs.run_all_tests }} jit_kernel: ${{ steps.filter.outputs.jit_kernel || steps.run-mode.outputs.run_all_tests }} multimodal_gen: ${{ steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Determine run mode id: run-mode run: | # Run all tests for workflow_call (when ref input is provided) # Note: github.event_name is inherited from caller, so we detect workflow_call by checking inputs.ref if [[ "${{ inputs.run_all_tests }}" == "true" ]]; then echo "run_all_tests=true" >> $GITHUB_OUTPUT echo "Run mode: ALL TESTS (run_all_tests=${{ inputs.run_all_tests }})" else echo "run_all_tests=false" >> $GITHUB_OUTPUT echo "Run mode: FILTERED (triggered by ${{ github.event_name }})" fi - name: Detect file changes id: filter uses: dorny/paths-filter@v3 if: steps.run-mode.outputs.run_all_tests != 'true' with: filters: | main_package: - "python/sglang/!(multimodal_gen)/**" - "python/pyproject_rocm.toml" - "python/pyproject_other.toml" - "scripts/ci/amd/*" - "scripts/ci/utils/*" - "test/**" - ".github/workflows/pr-test-amd-rocm720.yml" sgl_kernel: - "sgl-kernel/**" - ".github/workflows/pr-test-amd-rocm720.yml" jit_kernel: - "python/sglang/jit_kernel/**" - ".github/workflows/pr-test-amd-rocm720.yml" multimodal_gen: - "python/sglang/multimodal_gen/**" - "python/sglang/cli/**" - "python/sglang/jit_kernel/diffusion/**" - "python/pyproject_rocm.toml" - "python/pyproject_other.toml" # =============================================== sgl-kernel ==================================================== sgl-kernel-unit-test-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',sgl-kernel-unit-test-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && needs.check-changes.outputs.sgl_kernel == 'true' ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-1gpu-sglang] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 14 run: | docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_moe_align.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_moe_topk_softmax.py docker exec -w /sglang-checkout/sgl-kernel/tests/speculative ci_sglang python3 -m pytest test_eagle_utils.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_apply_token_bitmask_inplace.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_activation.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_topk.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_kvcacheio.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_moe_topk_sigmoid.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_torch_defaults_reset.py sgl-kernel-unit-test-2-gpu-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',sgl-kernel-unit-test-2-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && needs.check-changes.outputs.sgl_kernel == 'true' ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-2gpu-sglang] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 20 run: | docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_amd_deterministic_custom_allreduce.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_amd_nccl_allreduce_determinism.py # =============================================== primary ==================================================== stage-a-test-1-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-a-test-1-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-1gpu-sglang] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 10 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-a-test-1-amd ${{ inputs.continue_on_error && '--continue-on-error' || '' }} jit-kernel-unit-test-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',jit-kernel-unit-test-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && needs.check-changes.outputs.jit_kernel == 'true' ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-1gpu-sglang] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run JIT kernel unit tests timeout-minutes: 10 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout" python3 -m pytest -q python/sglang/jit_kernel/tests/test_store_cache.py stage-b-test-small-1-gpu-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-b-test-small-1-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-1gpu-sglang] part: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-small-1-gpu-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 14 --timeout-per-file 1800 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} stage-b-test-small-1-gpu-amd-nondeterministic: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-b-test-small-1-gpu-amd-nondeterministic,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-1gpu-sglang] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-small-1-gpu-amd-nondeterministic --timeout-per-file 1800 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} stage-b-test-small-1-gpu-amd-mi35x: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-b-test-small-1-gpu-amd-mi35x,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi35x-gpu-1] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-small-1-gpu-amd-mi35x ${{ inputs.continue_on_error && '--continue-on-error' || '' }} stage-b-test-large-1-gpu-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-b-test-large-1-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-1gpu-sglang] part: [0, 1] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-large-1-gpu-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 --timeout-per-file 1800 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} stage-b-test-large-2-gpu-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-b-test-large-2-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-2gpu-sglang] part: [0, 1] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-large-2-gpu-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 --timeout-per-file 1800 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} multimodal-gen-test-1-gpu-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',multimodal-gen-test-1-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false max-parallel: 1 # Run one at a time to avoid eviction from resource exhaustion during AITER kernel JIT matrix: runner: [linux-mi325-1gpu-sglang] part: [0, 1, 2, 3] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh diffusion docker exec ci_sglang pip install amdsmi - name: Setup kernel caches run: | # Use the persistent /sgl-data directory (mounted from /home/runner/sgl-data) # This directory persists across container restarts on the self-hosted runner docker exec ci_sglang mkdir -p /sgl-data/aiter-kernels /sgl-data/miopen-cache /sgl-data/hf-cache/hub # Clear pre-built AITER kernels from Docker image to avoid segfaults # The image may have stale/incompatible kernels at /sgl-workspace/aiter/aiter/jit/ echo "Clearing pre-built AITER kernels from Docker image..." docker exec ci_sglang rm -rf /sgl-workspace/aiter/aiter/jit/*.so 2>/dev/null || true docker exec ci_sglang rm -rf /sgl-data/aiter-kernels/*.so 2>/dev/null || true echo "AITER kernels cleared - will be rebuilt on first use" # Create persistent cache marker if /sgl-data is a real mount (not ephemeral) # This tells the test cleanup code to NOT delete downloaded models if docker exec ci_sglang test -d /sgl-data && docker exec ci_sglang mountpoint -q /sgl-data 2>/dev/null; then docker exec ci_sglang touch /sgl-data/hf-cache/.persistent_cache echo "Created .persistent_cache marker - HF cache will persist" else echo "WARNING: /sgl-data is not a mount point - models will be cleaned up after each test" fi # Check MIOpen cache (VAE convolution kernels) miopen_files=$(docker exec ci_sglang find /sgl-data/miopen-cache -name "*.udb" 2>/dev/null | wc -l || echo "0") echo "Found ${miopen_files} MIOpen cache files" - name: Diagnose HF cache and system resources run: | echo "=== System Memory Status ===" free -h echo "" echo "=== Disk Space ===" df -h /home/runner/sgl-data 2>/dev/null || df -h echo "" echo "=== HF Cache Directory Structure ===" docker exec ci_sglang ls -la /sgl-data/hf-cache/ 2>/dev/null || echo "HF cache dir not found" docker exec ci_sglang ls -la /sgl-data/hf-cache/hub/ 2>/dev/null || echo "HF hub cache not found" echo "" echo "=== Checking for cached diffusion models (1-GPU tests) ===" # Models used in 1-GPU tests: Wan2.1-T2V-1.3B, HunyuanVideo, Qwen-Image, FLUX.1, FLUX.2 for model in "Wan-AI--Wan2.1-T2V-1.3B-Diffusers" "tencent--HunyuanVideo" "Qwen--Qwen-Image" "black-forest-labs--FLUX.1-dev" "black-forest-labs--FLUX.2-dev"; do cache_path="/sgl-data/hf-cache/hub/models--${model}" if docker exec ci_sglang test -d "$cache_path"; then size=$(docker exec ci_sglang du -sh "$cache_path" 2>/dev/null | cut -f1) echo "✓ CACHED: $model ($size)" else echo "✗ NOT CACHED: $model" fi done echo "" echo "=== GPU Memory Status ===" docker exec ci_sglang rocm-smi --showmeminfo vram 2>/dev/null || echo "rocm-smi not available" - name: Run diffusion server tests (1-GPU) timeout-minutes: 60 run: | # AMD CI: All 1-GPU tests except FLUX.2 (FLUX.1 covers same code path) # Tests: T2V, T2I, I2V, LoRA # # HF download env vars: # - HF_HUB_ENABLE_HF_TRANSFER=1: Use faster hf_transfer for downloads (if available) # - HF_HUB_DISABLE_SYMLINKS_WARNING=1: Suppress symlink warnings docker exec \ -e SGLANG_E2E_TOLERANCE=0.3 \ -e SGLANG_STAGE_TIME_TOLERANCE=0.2 \ -e SGLANG_NON_DENOISE_STAGE_TIME_TOLERANCE=0.6 \ -e SGLANG_DENOISE_STEP_TOLERANCE=0.6 \ -e SGLANG_DENOISE_AGG_TOLERANCE=0.3 \ -e SGLANG_TEST_NUM_INFERENCE_STEPS=5 \ -e AITER_JIT_DIR=/sgl-data/aiter-kernels \ -e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \ -e HF_HUB_ENABLE_HF_TRANSFER=1 \ -e HF_HUB_DISABLE_SYMLINKS_WARNING=1 \ -w /sglang-checkout/python \ ci_sglang python3 sglang/multimodal_gen/test/run_suite.py \ --suite 1-gpu \ --partition-id ${{ matrix.part }} \ --total-partitions 4 \ -k "not flux_2" # Post-test diagnostics echo "=== Post-test System Memory Status ===" free -h multimodal-gen-test-2-gpu-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',multimodal-gen-test-2-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false max-parallel: 1 # Run one at a time to avoid eviction from resource exhaustion during AITER kernel JIT matrix: runner: [linux-mi325-2gpu-sglang] part: [0, 1] # 2 partitions: 9 tests ÷ 2 = ~4-5 tests each runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh diffusion docker exec ci_sglang pip install amdsmi - name: Setup kernel caches run: | # Use the persistent /sgl-data directory (mounted from /home/runner/sgl-data) docker exec ci_sglang mkdir -p /sgl-data/aiter-kernels /sgl-data/miopen-cache /sgl-data/hf-cache/hub # Clear pre-built AITER kernels from Docker image to avoid segfaults # The image may have stale/incompatible kernels at /sgl-workspace/aiter/aiter/jit/ echo "Clearing pre-built AITER kernels from Docker image..." docker exec ci_sglang rm -rf /sgl-workspace/aiter/aiter/jit/*.so 2>/dev/null || true docker exec ci_sglang rm -rf /sgl-data/aiter-kernels/*.so 2>/dev/null || true echo "AITER kernels cleared - will be rebuilt on first use" # Create persistent cache marker if /sgl-data is a real mount (not ephemeral) # This tells the test cleanup code to NOT delete downloaded models if docker exec ci_sglang test -d /sgl-data && docker exec ci_sglang mountpoint -q /sgl-data 2>/dev/null; then docker exec ci_sglang touch /sgl-data/hf-cache/.persistent_cache echo "Created .persistent_cache marker - HF cache will persist" else echo "WARNING: /sgl-data is not a mount point - models will be cleaned up after each test" fi # Check MIOpen cache (VAE convolution kernels) miopen_files=$(docker exec ci_sglang find /sgl-data/miopen-cache -name "*.udb" 2>/dev/null | wc -l || echo "0") echo "Found ${miopen_files} MIOpen cache files" - name: Diagnose HF cache and system resources run: | echo "=== System Memory Status ===" free -h echo "" echo "=== Disk Space ===" df -h /home/runner/sgl-data 2>/dev/null || df -h echo "" echo "=== HF Cache Directory Structure ===" docker exec ci_sglang ls -la /sgl-data/hf-cache/ 2>/dev/null || echo "HF cache dir not found" docker exec ci_sglang ls -la /sgl-data/hf-cache/hub/ 2>/dev/null || echo "HF hub cache not found" echo "" echo "=== Checking for cached diffusion models (2-GPU tests) ===" # Models used in 2-GPU tests: Wan2.2-T2V-A14B, Wan2.1-T2V-14B, Qwen-Image, FLUX.1 for model in "Wan-AI--Wan2.2-T2V-A14B-Diffusers" "Wan-AI--Wan2.1-T2V-14B-Diffusers" "Qwen--Qwen-Image" "black-forest-labs--FLUX.1-dev"; do cache_path="/sgl-data/hf-cache/hub/models--${model}" if docker exec ci_sglang test -d "$cache_path"; then size=$(docker exec ci_sglang du -sh "$cache_path" 2>/dev/null | cut -f1) echo "✓ CACHED: $model ($size)" else echo "✗ NOT CACHED: $model" fi done echo "" echo "=== GPU Memory Status ===" docker exec ci_sglang rocm-smi --showmeminfo vram 2>/dev/null || echo "rocm-smi not available" - name: Run diffusion server tests (2-GPU) timeout-minutes: 80 run: | # AMD CI: All 2-GPU tests including LoRA # Tests: T2V, T2I, I2V, LoRA # # HF download env vars: # - HF_HUB_ENABLE_HF_TRANSFER=1: Use faster hf_transfer for downloads (if available) # - HF_HUB_DISABLE_SYMLINKS_WARNING=1: Suppress symlink warnings docker exec \ -e SGLANG_E2E_TOLERANCE=0.3 \ -e SGLANG_STAGE_TIME_TOLERANCE=0.2 \ -e SGLANG_NON_DENOISE_STAGE_TIME_TOLERANCE=0.6 \ -e SGLANG_DENOISE_STEP_TOLERANCE=0.6 \ -e SGLANG_DENOISE_AGG_TOLERANCE=0.3 \ -e SGLANG_TEST_NUM_INFERENCE_STEPS=5 \ -e AITER_JIT_DIR=/sgl-data/aiter-kernels \ -e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \ -e HF_HUB_ENABLE_HF_TRANSFER=1 \ -e HF_HUB_DISABLE_SYMLINKS_WARNING=1 \ -w /sglang-checkout/python \ ci_sglang python3 sglang/multimodal_gen/test/run_suite.py \ --suite 2-gpu \ --partition-id ${{ matrix.part }} \ --total-partitions 2 # Post-test diagnostics echo "=== Post-test System Memory Status ===" free -h stage-c-test-large-8-gpu-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-c-test-large-8-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) env: RUNNER_LABELS: linux-mi325-8gpu-sglang strategy: fail-fast: false matrix: runner: [linux-mi325-8gpu-sglang] part: [0, 1, 2] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Test RCCL multi-GPU communication timeout-minutes: 5 run: | echo "Testing RCCL multi-GPU communication with debug info..." docker exec ci_sglang bash -c "cd /sglang-checkout && NCCL_DEBUG=INFO RCCL_DEBUG=INFO torchrun --nproc_per_node=8 scripts/ci/amd/test_rccl_multi_gpu.py" - name: Run test timeout-minutes: 60 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-c-test-large-8-gpu-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} stage-c-test-large-8-gpu-amd-mi35x: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-c-test-large-8-gpu-amd-mi35x,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi35x-gpu-8] part: [0, 1] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 60 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-c-test-large-8-gpu-amd-mi35x --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} # =============================================== Disaggregation ==================================================== stage-b-test-large-8-gpu-35x-disaggregation-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-b-test-large-8-gpu-disaggregation-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi35x-gpu-8.fabric] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Check Host RDMA Environment id: rdma_detect run: | set +e echo "=== Checking Host RDMA Environment ===" echo "" echo "=== 1. Ionic driver library check ===" ls -l /usr/lib/x86_64-linux-gnu/libibverbs/libionic* 2>/dev/null || echo "libionic not found in standard path" echo "" echo "=== 2. Infiniband devices ===" ls -la /dev/infiniband/ 2>/dev/null || echo "/dev/infiniband not found" ls -la /sys/class/infiniband/ 2>/dev/null || echo "/sys/class/infiniband not found" echo "" echo "=== 3. ibv_devinfo ===" which ibv_devinfo 2>/dev/null && ibv_devinfo 2>&1 || echo "ibv_devinfo not available" echo "" echo "=== 4. Kernel modules ===" lsmod 2>/dev/null | grep -E "ib_|rdma|ionic" || echo "No RDMA kernel modules loaded" echo "" echo "=== 5. Detect RDMA Devices for test environment ===" if [ -d "/sys/class/infiniband" ]; then RDMA_DEVS=$(ls /sys/class/infiniband | paste -sd "," -) echo "Detected RDMA Devices: $RDMA_DEVS" echo "SGLANG_TEST_RDMA_DEVICE=$RDMA_DEVS" >> $GITHUB_ENV else echo "No RDMA devices found in /sys/class/infiniband" echo "SGLANG_TEST_RDMA_DEVICE=" >> $GITHUB_ENV fi echo "" echo "=== Host RDMA Check Complete ===" - name: Start Special Container run: bash scripts/ci/amd/amd_ci_start_container_disagg.sh --rocm-version rocm720 env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Verify RDMA in Container run: | docker exec -u root ci_sglang bash -c ' echo "=== Container RDMA Verification ===" echo "Device nodes:" ls -la /dev/infiniband/ echo "" echo "Provider libraries:" ls /usr/lib/x86_64-linux-gnu/libibverbs/ | grep -E "ionic|mlx" || echo "No Ionic/Mellanox providers" echo "" echo "HCA devices:" HCA_COUNT=$(ibv_devinfo -list 2>&1 | grep -oE "^[0-9]+ HCAs? found" | grep -oE "^[0-9]+" || echo "0") ibv_devinfo -list if [ "$HCA_COUNT" -gt 0 ]; then echo "" echo "=== SUCCESS: RDMA setup complete. Found $HCA_COUNT HCA(s) ===" else echo "" echo "=== WARNING: No HCAs detected. RDMA tests may fail ===" fi ' - name: Run Aiter Op Test (RMSNorm) timeout-minutes: 10 run: | echo "Running pre-check: test_rmsnorm2d.py" docker exec \ -e MAX_JOBS=192 \ ci_sglang \ python /sgl-workspace/aiter/op_tests/test_rmsnorm2d.py - name: Run test_disaggregation timeout-minutes: 60 run: | bash scripts/ci/amd/amd_ci_exec.sh \ -e SGLANG_TEST_RDMA_DEVICE="${{ env.SGLANG_TEST_RDMA_DEVICE }}" \ -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-large-8-gpu-35x-disaggregation-amd --timeout-per-file 1800 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} pr-test-amd-finish: needs: [ call-gate, check-changes, sgl-kernel-unit-test-amd, sgl-kernel-unit-test-2-gpu-amd, multimodal-gen-test-1-gpu-amd, multimodal-gen-test-2-gpu-amd, stage-a-test-1-amd, jit-kernel-unit-test-amd, stage-b-test-small-1-gpu-amd, stage-b-test-small-1-gpu-amd-nondeterministic, stage-b-test-small-1-gpu-amd-mi35x, stage-b-test-large-1-gpu-amd, stage-b-test-large-2-gpu-amd, stage-b-test-large-8-gpu-35x-disaggregation-amd, stage-c-test-large-8-gpu-amd, stage-c-test-large-8-gpu-amd-mi35x, ] if: always() runs-on: ubuntu-latest steps: - name: Check all dependent job statuses run: | # Convert the 'needs' context to a JSON string json_needs='${{ toJson(needs) }}' # Get a list of all job names from the JSON keys job_names=$(echo "$json_needs" | jq -r 'keys_unsorted[]') for job in $job_names; do # For each job, extract its result result=$(echo "$json_needs" | jq -r --arg j "$job" '.[$j].result') # Print the job name and its result echo "$job: $result" # Check for failure or cancellation and exit if found if [[ "$result" == "failure" || "$result" == "cancelled" ]]; then echo "The above jobs failed." exit 1 fi done # If the loop completes, all jobs were successful echo "All jobs completed successfully" exit 0 ================================================ FILE: .github/workflows/pr-test-amd.yml ================================================ name: PR Test (AMD) # Dynamic run-name for /rerun-stage commands to enable URL lookup # Format: "[stage-name] sha" for fork PRs, "[stage-name]" for non-fork, default for normal runs run-name: ${{ (inputs.target_stage || inputs.target_stage_select) && (inputs.pr_head_sha && format('[{0}] {1}', inputs.target_stage || inputs.target_stage_select, inputs.pr_head_sha) || format('[{0}]', inputs.target_stage || inputs.target_stage_select)) || '' }} on: push: branches: [ main ] paths: - "python/**" - "scripts/ci/**" - "test/**" - "sgl-kernel/**" - ".github/workflows/pr-test-amd.yml" - "docker/rocm.Dockerfile" pull_request: branches: [ main ] paths: - "python/**" - "scripts/ci/**" - "test/**" - "sgl-kernel/**" - ".github/workflows/pr-test-amd.yml" - "docker/rocm.Dockerfile" workflow_dispatch: inputs: target_stage_select: description: "Select a stage to run from dropdown (leave empty for auto-detect)" required: false type: choice default: '' options: - '' - sgl-kernel-unit-test-amd - sgl-kernel-unit-test-2-gpu-amd - stage-a-test-1-amd - jit-kernel-unit-test-amd - stage-b-test-small-1-gpu-amd - stage-b-test-small-1-gpu-amd-nondeterministic - stage-b-test-small-1-gpu-amd-mi35x - stage-b-test-large-1-gpu-amd - stage-b-test-large-2-gpu-amd - multimodal-gen-test-1-gpu-amd - multimodal-gen-test-2-gpu-amd - stage-c-test-large-8-gpu-amd - stage-c-test-large-8-gpu-amd-mi35x - stage-b-test-large-8-gpu-disaggregation-amd target_stage: description: "Or type comma-separated stage names (overrides dropdown if non-empty)" required: false type: string default: "" pr_head_sha: description: "PR head SHA to checkout (for /rerun-stage on fork PRs)" required: false type: string default: "" aiter_ref: description: 'Override AITER commit (optional, leave empty to use Dockerfile default)' required: false type: string default: '' continue_on_error: description: 'Continue on error (do not fail the workflow on test failures)' required: false type: boolean default: false workflow_call: inputs: ref: description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.' required: false type: string default: '' run_all_tests: description: "Run all tests (for releasing or testing purpose)" required: false type: boolean default: false aiter_ref: description: 'Override AITER commit (optional, leave empty to use Dockerfile default)' required: false type: string default: '' continue_on_error: description: 'Continue on error (do not fail the workflow on test failures)' required: false type: boolean default: false env: AITER_COMMIT_OVERRIDE: ${{ inputs.aiter_ref }} concurrency: # When called via workflow_call with run_all_tests=true, use a unique group per run to # avoid collisions with direct push/PR triggers. We use run_all_tests (not github.event_name) # to detect this, because github.event_name inherits from the caller in workflow_call. group: pr-test-amd-${{ inputs.run_all_tests && format('full-{0}', github.run_id) || inputs.pr_head_sha || inputs.ref || github.ref }} cancel-in-progress: ${{ !inputs.run_all_tests && github.event_name != 'workflow_call' }} jobs: call-gate: uses: ./.github/workflows/pr-gate.yml secrets: inherit check-changes: needs: [call-gate] runs-on: ubuntu-latest outputs: main_package: ${{ steps.filter.outputs.main_package || steps.run-mode.outputs.run_all_tests }} sgl_kernel: ${{ steps.filter.outputs.sgl_kernel || steps.run-mode.outputs.run_all_tests }} jit_kernel: ${{ steps.filter.outputs.jit_kernel || steps.run-mode.outputs.run_all_tests }} multimodal_gen: ${{ steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Determine run mode id: run-mode run: | # Run all tests for workflow_call (when ref input is provided) # Note: github.event_name is inherited from caller, so we detect workflow_call by checking inputs.ref if [[ "${{ inputs.run_all_tests }}" == "true" ]]; then echo "run_all_tests=true" >> $GITHUB_OUTPUT echo "Run mode: ALL TESTS (run_all_tests=${{ inputs.run_all_tests }})" else echo "run_all_tests=false" >> $GITHUB_OUTPUT echo "Run mode: FILTERED (triggered by ${{ github.event_name }})" fi - name: Detect file changes id: filter uses: dorny/paths-filter@v3 if: steps.run-mode.outputs.run_all_tests != 'true' with: filters: | main_package: - "python/sglang/!(multimodal_gen)/**" - "python/pyproject_rocm.toml" - "python/pyproject_other.toml" - "scripts/ci/amd/*" - "scripts/ci/utils/*" - "test/**" - ".github/workflows/pr-test-amd.yml" sgl_kernel: - "sgl-kernel/**" - ".github/workflows/pr-test-amd.yml" jit_kernel: - "python/sglang/jit_kernel/**" - ".github/workflows/pr-test-amd.yml" multimodal_gen: - "python/sglang/multimodal_gen/**" - "python/sglang/cli/**" - "python/sglang/jit_kernel/diffusion/**" - "python/pyproject_rocm.toml" - "python/pyproject_other.toml" # =============================================== sgl-kernel ==================================================== sgl-kernel-unit-test-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',sgl-kernel-unit-test-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && needs.check-changes.outputs.sgl_kernel == 'true' ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-1gpu-sglang] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 14 run: | docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_moe_align.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_moe_topk_softmax.py docker exec -w /sglang-checkout/sgl-kernel/tests/speculative ci_sglang python3 -m pytest test_eagle_utils.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_apply_token_bitmask_inplace.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_activation.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_topk.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_kvcacheio.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_moe_topk_sigmoid.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_torch_defaults_reset.py sgl-kernel-unit-test-2-gpu-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',sgl-kernel-unit-test-2-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && needs.check-changes.outputs.sgl_kernel == 'true' ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-2gpu-sglang] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 20 run: | docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_amd_deterministic_custom_allreduce.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_amd_nccl_allreduce_determinism.py # =============================================== primary ==================================================== stage-a-test-1-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-a-test-1-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-1gpu-sglang] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 10 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-a-test-1-amd ${{ inputs.continue_on_error && '--continue-on-error' || '' }} jit-kernel-unit-test-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',jit-kernel-unit-test-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && needs.check-changes.outputs.jit_kernel == 'true' ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-1gpu-sglang] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run JIT kernel unit tests timeout-minutes: 10 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout" python3 -m pytest -q python/sglang/jit_kernel/tests/test_store_cache.py stage-b-test-small-1-gpu-amd: needs: [check-changes, stage-a-test-1-amd] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-b-test-small-1-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-1gpu-sglang] part: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-small-1-gpu-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 14 --timeout-per-file 1800 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} stage-b-test-small-1-gpu-amd-nondeterministic: needs: [check-changes, stage-a-test-1-amd] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-b-test-small-1-gpu-amd-nondeterministic,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-1gpu-sglang] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-small-1-gpu-amd-nondeterministic --timeout-per-file 1800 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} stage-b-test-small-1-gpu-amd-mi35x: needs: [check-changes, stage-a-test-1-amd] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-b-test-small-1-gpu-amd-mi35x,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi35x-gpu-1] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-small-1-gpu-amd-mi35x ${{ inputs.continue_on_error && '--continue-on-error' || '' }} stage-b-test-large-1-gpu-amd: needs: [check-changes, stage-a-test-1-amd] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-b-test-large-1-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-1gpu-sglang] part: [0, 1] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-large-1-gpu-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 --timeout-per-file 1800 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} stage-b-test-large-2-gpu-amd: needs: [check-changes, stage-a-test-1-amd] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-b-test-large-2-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi325-2gpu-sglang] part: [0, 1] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-large-2-gpu-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 --timeout-per-file 1800 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} multimodal-gen-test-1-gpu-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',multimodal-gen-test-1-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && needs.check-changes.outputs.multimodal_gen == 'true' ) ) strategy: fail-fast: false max-parallel: 1 # Run one at a time to avoid eviction from resource exhaustion during AITER kernel JIT matrix: runner: [linux-mi325-1gpu-sglang] part: [0, 1, 2, 3] # 2 partitions: 11 tests ÷ 2 = ~5-6 tests each runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh diffusion - name: Setup kernel caches run: | # Use the persistent /sgl-data directory (mounted from /home/runner/sgl-data) # This directory persists across container restarts on the self-hosted runner docker exec ci_sglang mkdir -p /sgl-data/aiter-kernels /sgl-data/miopen-cache /sgl-data/hf-cache/hub # Clear pre-built AITER kernels from Docker image to avoid segfaults # The image may have stale/incompatible kernels at /sgl-workspace/aiter/aiter/jit/ echo "Clearing pre-built AITER kernels from Docker image..." docker exec ci_sglang rm -rf /sgl-workspace/aiter/aiter/jit/*.so 2>/dev/null || true docker exec ci_sglang rm -rf /sgl-data/aiter-kernels/*.so 2>/dev/null || true echo "AITER kernels cleared - will be rebuilt on first use" # Create persistent cache marker if /sgl-data is a real mount (not ephemeral) # This tells the test cleanup code to NOT delete downloaded models if docker exec ci_sglang test -d /sgl-data && docker exec ci_sglang mountpoint -q /sgl-data 2>/dev/null; then docker exec ci_sglang touch /sgl-data/hf-cache/.persistent_cache echo "Created .persistent_cache marker - HF cache will persist" else echo "WARNING: /sgl-data is not a mount point - models will be cleaned up after each test" fi # Check MIOpen cache (VAE convolution kernels) miopen_files=$(docker exec ci_sglang find /sgl-data/miopen-cache -name "*.udb" 2>/dev/null | wc -l || echo "0") echo "Found ${miopen_files} MIOpen cache files" - name: Diagnose HF cache and system resources run: | echo "=== System Memory Status ===" free -h echo "" echo "=== Disk Space ===" df -h /home/runner/sgl-data 2>/dev/null || df -h echo "" echo "=== HF Cache Directory Structure ===" docker exec ci_sglang ls -la /sgl-data/hf-cache/ 2>/dev/null || echo "HF cache dir not found" docker exec ci_sglang ls -la /sgl-data/hf-cache/hub/ 2>/dev/null || echo "HF hub cache not found" echo "" echo "=== Checking for cached diffusion models (1-GPU tests) ===" # Models used in 1-GPU tests: Wan2.1-T2V-1.3B, HunyuanVideo, Qwen-Image, FLUX.1, FLUX.2 for model in "Wan-AI--Wan2.1-T2V-1.3B-Diffusers" "tencent--HunyuanVideo" "Qwen--Qwen-Image" "black-forest-labs--FLUX.1-dev" "black-forest-labs--FLUX.2-dev"; do cache_path="/sgl-data/hf-cache/hub/models--${model}" if docker exec ci_sglang test -d "$cache_path"; then size=$(docker exec ci_sglang du -sh "$cache_path" 2>/dev/null | cut -f1) echo "✓ CACHED: $model ($size)" else echo "✗ NOT CACHED: $model" fi done echo "" echo "=== GPU Memory Status ===" docker exec ci_sglang rocm-smi --showmeminfo vram 2>/dev/null || echo "rocm-smi not available" - name: Run diffusion server tests (1-GPU) timeout-minutes: 90 run: | # AMD CI: All 1-GPU tests except FLUX.2 (FLUX.1 covers same code path) # Tests: T2V, T2I, I2V, LoRA # # HF download env vars: # - HF_HUB_ENABLE_HF_TRANSFER=1: Use faster hf_transfer for downloads (if available) # - HF_HUB_DISABLE_SYMLINKS_WARNING=1: Suppress symlink warnings docker exec \ -e SGLANG_E2E_TOLERANCE=0.3 \ -e SGLANG_STAGE_TIME_TOLERANCE=0.2 \ -e SGLANG_NON_DENOISE_STAGE_TIME_TOLERANCE=0.6 \ -e SGLANG_DENOISE_STEP_TOLERANCE=0.6 \ -e SGLANG_DENOISE_AGG_TOLERANCE=0.3 \ -e SGLANG_TEST_NUM_INFERENCE_STEPS=5 \ -e AITER_JIT_DIR=/sgl-data/aiter-kernels \ -e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \ -e HF_HUB_ENABLE_HF_TRANSFER=1 \ -e HF_HUB_DISABLE_SYMLINKS_WARNING=1 \ -w /sglang-checkout/python \ ci_sglang python3 sglang/multimodal_gen/test/run_suite.py \ --suite 1-gpu \ --partition-id ${{ matrix.part }} \ --total-partitions 4 \ -k "not flux_2" # Post-test diagnostics echo "=== Post-test System Memory Status ===" free -h multimodal-gen-test-2-gpu-amd: needs: [check-changes] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',multimodal-gen-test-2-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && needs.check-changes.outputs.multimodal_gen == 'true' ) ) strategy: fail-fast: false max-parallel: 1 # Run one at a time to avoid eviction from resource exhaustion during AITER kernel JIT matrix: runner: [linux-mi325-2gpu-sglang] part: [0, 1] # 2 partitions: 9 tests ÷ 2 = ~4-5 tests each runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: | bash scripts/ci/amd/amd_ci_install_dependency.sh diffusion - name: Setup kernel caches run: | # Use the persistent /sgl-data directory (mounted from /home/runner/sgl-data) docker exec ci_sglang mkdir -p /sgl-data/aiter-kernels /sgl-data/miopen-cache /sgl-data/hf-cache/hub # Clear pre-built AITER kernels from Docker image to avoid segfaults # The image may have stale/incompatible kernels at /sgl-workspace/aiter/aiter/jit/ echo "Clearing pre-built AITER kernels from Docker image..." docker exec ci_sglang rm -rf /sgl-workspace/aiter/aiter/jit/*.so 2>/dev/null || true docker exec ci_sglang rm -rf /sgl-data/aiter-kernels/*.so 2>/dev/null || true echo "AITER kernels cleared - will be rebuilt on first use" # Create persistent cache marker if /sgl-data is a real mount (not ephemeral) # This tells the test cleanup code to NOT delete downloaded models if docker exec ci_sglang test -d /sgl-data && docker exec ci_sglang mountpoint -q /sgl-data 2>/dev/null; then docker exec ci_sglang touch /sgl-data/hf-cache/.persistent_cache echo "Created .persistent_cache marker - HF cache will persist" else echo "WARNING: /sgl-data is not a mount point - models will be cleaned up after each test" fi # Check MIOpen cache (VAE convolution kernels) miopen_files=$(docker exec ci_sglang find /sgl-data/miopen-cache -name "*.udb" 2>/dev/null | wc -l || echo "0") echo "Found ${miopen_files} MIOpen cache files" - name: Diagnose HF cache and system resources run: | echo "=== System Memory Status ===" free -h echo "" echo "=== Disk Space ===" df -h /home/runner/sgl-data 2>/dev/null || df -h echo "" echo "=== HF Cache Directory Structure ===" docker exec ci_sglang ls -la /sgl-data/hf-cache/ 2>/dev/null || echo "HF cache dir not found" docker exec ci_sglang ls -la /sgl-data/hf-cache/hub/ 2>/dev/null || echo "HF hub cache not found" echo "" echo "=== Checking for cached diffusion models (2-GPU tests) ===" # Models used in 2-GPU tests: Wan2.2-T2V-A14B, Wan2.1-T2V-14B, Qwen-Image, FLUX.1 for model in "Wan-AI--Wan2.2-T2V-A14B-Diffusers" "Wan-AI--Wan2.1-T2V-14B-Diffusers" "Qwen--Qwen-Image" "black-forest-labs--FLUX.1-dev"; do cache_path="/sgl-data/hf-cache/hub/models--${model}" if docker exec ci_sglang test -d "$cache_path"; then size=$(docker exec ci_sglang du -sh "$cache_path" 2>/dev/null | cut -f1) echo "✓ CACHED: $model ($size)" else echo "✗ NOT CACHED: $model" fi done echo "" echo "=== GPU Memory Status ===" docker exec ci_sglang rocm-smi --showmeminfo vram 2>/dev/null || echo "rocm-smi not available" - name: Run diffusion server tests (2-GPU) timeout-minutes: 80 run: | # AMD CI: All 2-GPU tests including LoRA # Tests: T2V, T2I, I2V, LoRA # # HF download env vars: # - HF_HUB_ENABLE_HF_TRANSFER=1: Use faster hf_transfer for downloads (if available) # - HF_HUB_DISABLE_SYMLINKS_WARNING=1: Suppress symlink warnings docker exec \ -e SGLANG_E2E_TOLERANCE=0.3 \ -e SGLANG_STAGE_TIME_TOLERANCE=0.2 \ -e SGLANG_NON_DENOISE_STAGE_TIME_TOLERANCE=0.6 \ -e SGLANG_DENOISE_STEP_TOLERANCE=0.6 \ -e SGLANG_DENOISE_AGG_TOLERANCE=0.3 \ -e SGLANG_TEST_NUM_INFERENCE_STEPS=5 \ -e AITER_JIT_DIR=/sgl-data/aiter-kernels \ -e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \ -e HF_HUB_ENABLE_HF_TRANSFER=1 \ -e HF_HUB_DISABLE_SYMLINKS_WARNING=1 \ -w /sglang-checkout/python \ ci_sglang python3 sglang/multimodal_gen/test/run_suite.py \ --suite 2-gpu \ --partition-id ${{ matrix.part }} \ --total-partitions 2 # Post-test diagnostics echo "=== Post-test System Memory Status ===" free -h stage-c-test-large-8-gpu-amd: needs: [check-changes, call-gate, stage-b-test-small-1-gpu-amd, stage-b-test-large-2-gpu-amd] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-c-test-large-8-gpu-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) env: RUNNER_LABELS: linux-mi325-8gpu-sglang strategy: fail-fast: false matrix: runner: [linux-mi325-8gpu-sglang] part: [0, 1, 2] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Test RCCL multi-GPU communication timeout-minutes: 5 run: | echo "Testing RCCL multi-GPU communication with debug info..." docker exec ci_sglang bash -c "cd /sglang-checkout && NCCL_DEBUG=INFO RCCL_DEBUG=INFO torchrun --nproc_per_node=8 scripts/ci/amd/test_rccl_multi_gpu.py" - name: Run test timeout-minutes: 60 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-c-test-large-8-gpu-amd --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} stage-c-test-large-8-gpu-amd-mi35x: needs: [check-changes, call-gate, stage-b-test-small-1-gpu-amd, stage-b-test-large-2-gpu-amd] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-c-test-large-8-gpu-amd-mi35x,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi35x-gpu-8] part: [0, 1] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Start CI container run: bash scripts/ci/amd/amd_ci_start_container.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test timeout-minutes: 60 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-c-test-large-8-gpu-amd-mi35x --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 --timeout-per-file 3600 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} # =============================================== Disaggregation ==================================================== stage-b-test-large-8-gpu-35x-disaggregation-amd: needs: [check-changes, stage-a-test-1-amd] if: | always() && ( (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',stage-b-test-large-8-gpu-disaggregation-amd,')) || ( !(inputs.target_stage || inputs.target_stage_select) && (!failure() && !cancelled()) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) strategy: fail-fast: false matrix: runner: [linux-mi35x-gpu-8.fabric] runs-on: ${{matrix.runner}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} - name: Ensure VRAM is clear run: bash scripts/ensure_vram_clear.sh rocm - name: Check Host RDMA Environment id: rdma_detect run: | set +e echo "=== Checking Host RDMA Environment ===" echo "" echo "=== 1. Ionic driver library check ===" ls -l /usr/lib/x86_64-linux-gnu/libibverbs/libionic* 2>/dev/null || echo "libionic not found in standard path" echo "" echo "=== 2. Infiniband devices ===" ls -la /dev/infiniband/ 2>/dev/null || echo "/dev/infiniband not found" ls -la /sys/class/infiniband/ 2>/dev/null || echo "/sys/class/infiniband not found" echo "" echo "=== 3. ibv_devinfo ===" which ibv_devinfo 2>/dev/null && ibv_devinfo 2>&1 || echo "ibv_devinfo not available" echo "" echo "=== 4. Kernel modules ===" lsmod 2>/dev/null | grep -E "ib_|rdma|ionic" || echo "No RDMA kernel modules loaded" echo "" echo "=== 5. Detect RDMA Devices for test environment ===" if [ -d "/sys/class/infiniband" ]; then RDMA_DEVS=$(ls /sys/class/infiniband | paste -sd "," -) echo "Detected RDMA Devices: $RDMA_DEVS" echo "SGLANG_TEST_RDMA_DEVICE=$RDMA_DEVS" >> $GITHUB_ENV else echo "No RDMA devices found in /sys/class/infiniband" echo "SGLANG_TEST_RDMA_DEVICE=" >> $GITHUB_ENV fi echo "" echo "=== Host RDMA Check Complete ===" - name: Start Special Container run: bash scripts/ci/amd/amd_ci_start_container_disagg.sh env: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Install dependencies run: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Verify RDMA in Container run: | docker exec -u root ci_sglang bash -c ' echo "=== Container RDMA Verification ===" echo "Device nodes:" ls -la /dev/infiniband/ echo "" echo "Provider libraries:" ls /usr/lib/x86_64-linux-gnu/libibverbs/ | grep -E "ionic|mlx" || echo "No Ionic/Mellanox providers" echo "" echo "HCA devices:" HCA_COUNT=$(ibv_devinfo -list 2>&1 | grep -oE "^[0-9]+ HCAs? found" | grep -oE "^[0-9]+" || echo "0") ibv_devinfo -list if [ "$HCA_COUNT" -gt 0 ]; then echo "" echo "=== SUCCESS: RDMA setup complete. Found $HCA_COUNT HCA(s) ===" else echo "" echo "=== WARNING: No HCAs detected. RDMA tests may fail ===" fi ' - name: Run Aiter Op Test (RMSNorm) timeout-minutes: 10 run: | echo "Running pre-check: test_rmsnorm2d.py" docker exec \ -e MAX_JOBS=192 \ ci_sglang \ python /sgl-workspace/aiter/op_tests/test_rmsnorm2d.py - name: Run test_disaggregation timeout-minutes: 60 run: | bash scripts/ci/amd/amd_ci_exec.sh \ -e SGLANG_TEST_RDMA_DEVICE="${{ env.SGLANG_TEST_RDMA_DEVICE }}" \ -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-b-test-large-8-gpu-35x-disaggregation-amd --timeout-per-file 1800 ${{ inputs.continue_on_error && '--continue-on-error' || '' }} pr-test-amd-finish: needs: [ call-gate, check-changes, sgl-kernel-unit-test-amd, sgl-kernel-unit-test-2-gpu-amd, multimodal-gen-test-1-gpu-amd, multimodal-gen-test-2-gpu-amd, stage-a-test-1-amd, jit-kernel-unit-test-amd, stage-b-test-small-1-gpu-amd, stage-b-test-small-1-gpu-amd-nondeterministic, stage-b-test-small-1-gpu-amd-mi35x, stage-b-test-large-1-gpu-amd, stage-b-test-large-2-gpu-amd, stage-b-test-large-8-gpu-35x-disaggregation-amd, stage-c-test-large-8-gpu-amd, stage-c-test-large-8-gpu-amd-mi35x, ] if: always() runs-on: ubuntu-latest steps: - name: Check all dependent job statuses run: | # Convert the 'needs' context to a JSON string json_needs='${{ toJson(needs) }}' # Get a list of all job names from the JSON keys job_names=$(echo "$json_needs" | jq -r 'keys_unsorted[]') for job in $job_names; do # For each job, extract its result result=$(echo "$json_needs" | jq -r --arg j "$job" '.[$j].result') # Print the job name and its result echo "$job: $result" # Check for failure or cancellation and exit if found if [[ "$result" == "failure" || "$result" == "cancelled" ]]; then echo "The above jobs failed." exit 1 fi done # If the loop completes, all jobs were successful echo "All jobs completed successfully" exit 0 ================================================ FILE: .github/workflows/pr-test-npu.yml ================================================ name: PR Test (NPU) on: push: branches: [ main ] pull_request: branches: [ main ] workflow_dispatch: workflow_call: inputs: ref: description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.' required: false type: string default: '' run_all_tests: description: "Run all tests (for releasing or testing purpose)" required: false type: boolean default: false concurrency: group: pr-test-npu-${{ inputs.ref || github.ref }} cancel-in-progress: ${{ github.event_name != 'workflow_call' }} jobs: # ==================== Check Changes ==================== # check-changes: runs-on: ubuntu-latest outputs: changes_exist: ${{ steps.filter.outputs.main_package == 'true' || steps.filter.outputs.multimodal_gen == 'true' || steps.run-mode.outputs.run_all_tests == 'true'}} main_package: ${{ steps.filter.outputs.main_package == 'true' || steps.run-mode.outputs.run_all_tests == 'true' }} multimodal_gen: ${{ steps.filter.outputs.multimodal_gen == 'true' || steps.run-mode.outputs.run_all_tests == 'true' }} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Determine run mode id: run-mode run: | # Run all tests for workflow_call (when ref input is provided) # Note: github.event_name is inherited from caller, so we detect workflow_call by checking inputs.ref if [[ "${{ inputs.run_all_tests }}" == "true" ]]; then echo "run_all_tests=true" >> $GITHUB_OUTPUT echo "Run mode: ALL TESTS (run_all_tests=${{ inputs.run_all_tests }})" else echo "run_all_tests=false" >> $GITHUB_OUTPUT echo "Run mode: FILTERED (triggered by ${{ github.event_name }})" fi - name: Detect file changes id: filter uses: dorny/paths-filter@v3 if: steps.run-mode.outputs.run_all_tests != 'true' with: filters: | main_package: - "python/sglang/!(multimodal_gen)/**" - "python/pyproject_npu.toml" - "scripts/ci/npu/npu_ci_install_dependency.sh" - "test/srt/ascend/**" - ".github/workflows/pr-test-npu.yml" multimodal_gen: - "python/sglang/multimodal_gen/**" - "python/sglang/srt/**" - "python/pyproject_npu.toml" - "scripts/ci/npu/npu_ci_install_dependency.sh" - ".github/workflows/pr-test-npu.yml" # ==================== PR Gate ==================== # pr-gate: needs: check-changes if: needs.check-changes.outputs.changes_exist == 'true' uses: ./.github/workflows/pr-gate.yml secrets: inherit per-commit-1-npu-a2: needs: [check-changes, pr-gate] if: needs.check-changes.outputs.main_package == 'true' runs-on: linux-aarch64-a2-1 strategy: fail-fast: false matrix: part: [ 0, 1 ] container: image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.5.0-910b-ubuntu22.04-py3.11 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies env: TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" run: | # speed up by using infra cache services CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple pip config set global.trusted-host "${CACHING_URL}" bash scripts/ci/npu/npu_ci_install_dependency.sh 910b # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Run registered test timeout-minutes: 240 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" STREAMS_PER_DEVICE: 32 run: | cd test python3 run_suite.py --hw npu --suite per-commit-1-npu-a2 --continue-on-error --timeout-per-file 3600 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 - name: Run test timeout-minutes: 60 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" STREAMS_PER_DEVICE: 32 run: | cd test/srt python3 run_suite.py --suite per-commit-1-npu-a2 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 per-commit-2-npu-a2: needs: [check-changes, pr-gate] if: needs.check-changes.outputs.main_package == 'true' runs-on: linux-aarch64-a2-2 strategy: fail-fast: true matrix: part: [0, 1] container: image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.5.0-910b-ubuntu22.04-py3.11 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies env: TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" run: | # speed up by using infra cache services CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple pip config set global.trusted-host "${CACHING_URL}" bash scripts/ci/npu/npu_ci_install_dependency.sh 910b # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Run test timeout-minutes: 60 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" STREAMS_PER_DEVICE: 32 run: | cd test/srt python3 run_suite.py --suite per-commit-2-npu-a2 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 per-commit-4-npu-a3: needs: [check-changes, pr-gate] if: needs.check-changes.outputs.main_package == 'true' runs-on: linux-aarch64-a3-4 container: image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.5.0-a3-ubuntu22.04-py3.11 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies env: TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" run: | # speed up by using infra cache services CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple pip config set global.trusted-host "${CACHING_URL}" bash scripts/ci/npu/npu_ci_install_dependency.sh a3 # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Run test timeout-minutes: 60 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" STREAMS_PER_DEVICE: 32 run: | cd test/srt python3 run_suite.py --suite per-commit-4-npu-a3 --timeout-per-file 3600 per-commit-16-npu-a3: needs: [check-changes, pr-gate] if: needs.check-changes.outputs.main_package == 'true' runs-on: linux-aarch64-a3-16 container: image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.5.0-a3-ubuntu22.04-py3.11 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies env: TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" run: | # speed up by using infra cache services CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple pip config set global.trusted-host "${CACHING_URL}" bash scripts/ci/npu/npu_ci_install_dependency.sh a3 # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Run test timeout-minutes: 60 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" STREAMS_PER_DEVICE: 32 run: | cd test/srt python3 run_suite.py --suite per-commit-16-npu-a3 --timeout-per-file 3600 multimodal-gen-test-1-npu-a3: needs: [check-changes, pr-gate] if: needs.check-changes.outputs.multimodal_gen == 'true' runs-on: linux-aarch64-a3-2 container: image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.3.rc2-a3-ubuntu22.04-py3.11 steps: - name: Checkout code uses: actions/checkout@v4 - name: Install dependencies env: TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" run: | # speed up by using infra cache services CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple pip config set global.trusted-host "${CACHING_URL}" bash scripts/ci/npu/npu_ci_install_dependency.sh a3 diffusion # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Run test timeout-minutes: 60 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" STREAMS_PER_DEVICE: 32 run: | export PATH="/usr/local/Ascend/8.3.RC1/compiler/bishengir/bin:${PATH}" cd python python3 sglang/multimodal_gen/test/run_suite.py --suite 1-npu multimodal-gen-test-2-npu-a3: needs: [check-changes, pr-gate] if: needs.check-changes.outputs.multimodal_gen == 'true' runs-on: linux-aarch64-a3-16 container: image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.3.rc2-a3-ubuntu22.04-py3.11 steps: - name: Checkout code uses: actions/checkout@v4 - name: Install dependencies env: TORCH_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/whl/cpu" PYPI_CACHE_URL: "http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple" GITHUB_PROXY_URL: "https://gh-proxy.test.osinfra.cn/" run: | # speed up by using infra cache services CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple pip config set global.trusted-host "${CACHING_URL}" bash scripts/ci/npu/npu_ci_install_dependency.sh a3 diffusion # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Run test timeout-minutes: 60 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" STREAMS_PER_DEVICE: 32 run: | export PATH="/usr/local/Ascend/8.3.RC1/compiler/bishengir/bin:${PATH}" cd python python3 sglang/multimodal_gen/test/run_suite.py --suite 2-npu multimodal-gen-test-8-npu-a3: needs: [check-changes, pr-gate] if: needs.check-changes.outputs.multimodal_gen == 'true' runs-on: linux-aarch64-a3-16 container: image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.5.0-a3-ubuntu22.04-py3.11 steps: - name: Checkout code uses: actions/checkout@v4 - name: Install dependencies run: | # speed up by using infra cache services CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple" pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn" bash scripts/ci/npu/npu_ci_install_dependency.sh a3 diffusion # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Run test timeout-minutes: 60 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" STREAMS_PER_DEVICE: 32 run: | cd python python3 sglang/multimodal_gen/test/run_suite.py --suite 8-npu ================================================ FILE: .github/workflows/pr-test-rust.yml ================================================ name: PR Test (SMG) on: push: branches: [ main ] paths: - "sgl-model-gateway/**" pull_request: branches: [ main ] types: [opened, synchronize, reopened, labeled] paths: - "sgl-model-gateway/**" workflow_dispatch: concurrency: group: gateway-tests-${{ github.ref }} cancel-in-progress: true env: RUSTC_WRAPPER: sccache SCCACHE_GHA_ENABLED: "true" jobs: build-wheel: if: | github.event_name != 'pull_request' || (github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) || (github.event.action == 'labeled' && github.event.label.name == 'run-ci') runs-on: 4-gpu-a10 steps: - name: Checkout code uses: actions/checkout@v4 - name: Install rust dependencies run: | bash scripts/ci/cuda/ci_install_gateway_dependencies.sh - name: Configure sccache uses: mozilla-actions/sccache-action@v0.0.9 with: version: "v0.12.0" disable_annotations: true - name: Rust cache uses: Swatinem/rust-cache@v2 with: workspaces: sgl-model-gateway shared-key: "rust-cache" cache-all-crates: true cache-on-failure: true save-if: true - name: Build python binding run: | source "$HOME/.cargo/env" export RUSTC_WRAPPER=sccache cd sgl-model-gateway/bindings/python python3 -m pip install --upgrade pip maturin maturin build --profile ci --features vendored-openssl --out dist - name: List built wheel run: ls -lh sgl-model-gateway/bindings/python/dist/ - name: Upload wheel artifact uses: actions/upload-artifact@v4 with: name: smg-wheel path: sgl-model-gateway/bindings/python/dist/*.whl retention-days: 1 - name: Test wheel install run: | pip install sgl-model-gateway/bindings/python/dist/*.whl python3 -c "import sglang_router; print('Python package: OK')" python3 -c "from sglang_router.sglang_router_rs import Router; print('Rust extension: OK')" python3 -m sglang_router.launch_router --help > /dev/null && echo "Entry point: OK" python-unit-tests: needs: build-wheel runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 with: path: sglang-repo - name: Move sgl-model-gateway folder to root run: | mv sglang-repo/sgl-model-gateway/* . rm -rf sglang-repo - name: Set up Python uses: actions/setup-python@v5 with: python-version: "3.13" - name: Download wheel artifact uses: actions/download-artifact@v4 with: name: smg-wheel path: dist/ - name: Install wheel run: pip install dist/*.whl - name: Run Python unit tests run: | cd bindings/python python3 -m pip install pytest pytest-cov pytest-xdist pytest -q tests --cov=sglang_router --cov-config=.coveragerc --cov-report=term-missing --cov-fail-under=80 unit-tests: if: | github.event_name != 'pull_request' || (github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) || (github.event.action == 'labeled' && github.event.label.name == 'run-ci') runs-on: ubuntu-latest steps: - name: Checkout code uses: actions/checkout@v4 - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_gateway_dependencies.sh - name: Configure sccache uses: mozilla-actions/sccache-action@v0.0.9 with: version: "v0.12.0" disable_annotations: true - name: Rust cache uses: Swatinem/rust-cache@v2 with: workspaces: sgl-model-gateway shared-key: "rust-cache" cache-all-crates: true cache-on-failure: true save-if: true - name: Run lint run: | source "$HOME/.cargo/env" cd sgl-model-gateway/ rustup component add clippy cargo clippy --all-targets --all-features -- -D warnings - name: Run fmt run: | source "$HOME/.cargo/env" cd sgl-model-gateway/ rustup component add --toolchain nightly-x86_64-unknown-linux-gnu rustfmt rustup toolchain install nightly --profile minimal cargo +nightly fmt -- --check - name: Generate vision golden fixtures run: | pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu pip install transformers pillow numpy scipy pip install transformers pillow numpy cd sgl-model-gateway/ python scripts/generate_vision_golden.py - name: Run Rust tests timeout-minutes: 20 run: | source "$HOME/.cargo/env" cd sgl-model-gateway/ cargo test - name: Show sccache stats if: always() run: sccache --show-stats gateway-e2e: name: ${{ matrix.name }} needs: build-wheel if: | github.event_name != 'pull_request' || (github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) || (github.event.action == 'labeled' && github.event.label.name == 'run-ci') strategy: fail-fast: false matrix: include: - name: benchmarks timeout: 32 test_dirs: "e2e_test/benchmarks" extra_deps: "genai-bench==0.0.3" env_vars: "" reruns: "" upload_benchmarks: true parallel_opts: "" # No parallel for benchmarks (performance measurement) - name: responses timeout: 45 test_dirs: "e2e_test/responses" extra_deps: "" env_vars: "SHOW_WORKER_LOGS=0 SHOW_ROUTER_LOGS=1" reruns: "--reruns 2 --reruns-delay 5" setup_oracle: true setup_brave: true parallel_opts: "" # Cloud backend tests not compatible with parallel execution - name: e2e timeout: 45 test_dirs: "e2e_test/router e2e_test/embeddings" extra_deps: "pytest-parallel py" # py is required for pytest-parallel with newer pytest env_vars: "SHOW_WORKER_LOGS=0 SHOW_ROUTER_LOGS=1" reruns: "--reruns 2 --reruns-delay 5" parallel_opts: "--workers 1 --tests-per-worker 4" # Thread-based parallelism - name: chat-completions timeout: 45 test_dirs: "e2e_test/chat_completions" extra_deps: "" env_vars: "SHOW_WORKER_LOGS=0 SHOW_ROUTER_LOGS=1" reruns: "--reruns 2 --reruns-delay 5" parallel_opts: "" runs-on: 4-gpu-a10 timeout-minutes: ${{ matrix.timeout }} steps: - name: Checkout code uses: actions/checkout@v4 - name: Install SGLang dependencies run: | sudo --preserve-env=PATH bash scripts/ci/cuda/ci_install_dependency.sh - name: Setup Oracle Instant Client if: matrix.setup_oracle run: | sudo apt-get install -y unzip INSTANT_CLIENT_DIR="/home/ubuntu/instant-client" INSTANT_CLIENT_ZIP="instantclient-basic-linux.x64-23.9.0.25.07.zip" if [ ! -d "$INSTANT_CLIENT_DIR/instantclient_23_9" ]; then echo "Downloading Oracle Instant Client..." mkdir -p "$INSTANT_CLIENT_DIR" cd "$INSTANT_CLIENT_DIR" wget https://download.oracle.com/otn_software/linux/instantclient/2390000/$INSTANT_CLIENT_ZIP unzip $INSTANT_CLIENT_ZIP rm $INSTANT_CLIENT_ZIP else echo "Oracle Instant Client already exists, skipping download" fi echo "LD_LIBRARY_PATH=/home/ubuntu/instant-client/instantclient_23_9:\$LD_LIBRARY_PATH" >> $GITHUB_ENV - name: Start Oracle Database if: matrix.setup_oracle run: | docker run -d -p 1521:1521 -e ORACLE_PASSWORD=oracle --name oracle-db gvenzl/oracle-xe:21-slim echo "Starting Oracle DB..." # Export Oracle connection environment variables echo "ATP_USER=system" >> $GITHUB_ENV echo "ATP_PASSWORD=oracle" >> $GITHUB_ENV echo "ATP_DSN=localhost:1521/XEPDB1" >> $GITHUB_ENV - name: Start Brave MCP Server if: matrix.setup_brave run: | docker run -d --rm \ -p 8001:8080 \ -e BRAVE_API_KEY \ --name brave-search-server \ shoofio/brave-search-mcp-sse:1.0.10 echo "Starting Brave MCP Server..." sleep 2 curl -f --max-time 1 http://localhost:8001/sse > /dev/null 2>&1 && echo "Brave MCP Server is healthy!" || echo "Brave MCP Server responded" - name: Download wheel artifact uses: actions/download-artifact@v4 with: name: smg-wheel path: wheel/ - name: Install wheel run: | pip uninstall -y sglang-router || true pip install wheel/*.whl - name: Install e2e test dependencies run: | python3 -m pip install pytest pytest-rerunfailures httpx openai grpcio grpcio-health-checking numpy if [ -n "${{ matrix.extra_deps }}" ]; then python3 -m pip --no-cache-dir install --upgrade ${{ matrix.extra_deps }} fi - name: Run E2E tests run: | bash scripts/killall_sglang.sh all cd sgl-model-gateway ${{ matrix.env_vars }} ROUTER_LOCAL_MODEL_PATH="/home/ubuntu/models" pytest ${{ matrix.reruns }} ${{ matrix.parallel_opts }} ${{ matrix.test_dirs }} -s -vv -o log_cli=true --log-cli-level=INFO - name: Upload benchmark results if: matrix.upload_benchmarks && success() uses: actions/upload-artifact@v4 with: name: genai-bench-results-all-policies path: sgl-model-gateway/benchmark_**/ - name: Cleanup Brave MCP Server if: always() && matrix.setup_brave run: | docker stop brave-search-server || true docker rm brave-search-server || true - name: Cleanup Oracle Database if: always() && matrix.setup_oracle run: | docker stop oracle-db || true docker rm oracle-db || true docker-build-test: if: | github.event_name != 'pull_request' || (github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) || (github.event.action == 'labeled' && github.event.label.name == 'run-ci') runs-on: ubuntu-24.04 steps: - name: Checkout repository uses: actions/checkout@v4 - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - name: Build Docker image (no push) uses: docker/build-push-action@v5 with: context: . file: docker/gateway.Dockerfile push: false tags: sgl-model-gateway:test cache-from: type=gha cache-to: type=gha,mode=max finish: needs: [build-wheel, python-unit-tests, unit-tests, gateway-e2e, docker-build-test] runs-on: ubuntu-latest steps: - name: Finish run: echo "This is an empty step to ensure that all jobs are completed." summarize-benchmarks: needs: gateway-e2e runs-on: ubuntu-latest if: success() steps: - name: Checkout code uses: actions/checkout@v4 - name: Download benchmark results uses: actions/download-artifact@v4 with: name: genai-bench-results-all-policies - name: Create benchmark summary run: python3 sgl-model-gateway/e2e_test/benchmarks/summarize.py . ================================================ FILE: .github/workflows/pr-test-xeon.yml ================================================ name: PR Test (Xeon) on: push: branches: [ main ] pull_request: branches: [ main ] workflow_dispatch: workflow_call: inputs: ref: description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.' required: false type: string default: '' run_all_tests: description: "Run all tests (for releasing or testing purpose)" required: false type: boolean default: false concurrency: group: pr-test-xeon-${{ inputs.ref || github.ref }} cancel-in-progress: false jobs: # ==================== Check Changes ==================== # check-changes: runs-on: ubuntu-latest outputs: main_package: ${{ steps.filter.outputs.main_package || steps.run-mode.outputs.run_all_tests}} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Determine run mode id: run-mode run: | # Run all tests for workflow_call (when ref input is provided) # Note: github.event_name is inherited from caller, so we detect workflow_call by checking inputs.ref if [[ "${{ inputs.run_all_tests }}" == "true" ]]; then echo "run_all_tests=true" >> $GITHUB_OUTPUT echo "Run mode: ALL TESTS (run_all_tests=${{ inputs.run_all_tests }})" else echo "run_all_tests=false" >> $GITHUB_OUTPUT echo "Run mode: FILTERED (triggered by ${{ github.event_name }})" fi - name: Detect file changes id: filter uses: dorny/paths-filter@v3 if: steps.run-mode.outputs.run_all_tests != 'true' with: filters: | main_package: - "python/sglang/!(multimodal_gen)/**" - "python/pyproject_cpu.toml" - "test/**" - "sgl-kernel/**" - ".github/workflows/pr-test-xeon.yml" - "docker/xeon.Dockerfile" # ==================== PR Gate ==================== # pr-gate: needs: check-changes if: needs.check-changes.outputs.main_package == 'true' uses: ./.github/workflows/pr-gate.yml secrets: inherit build-test: needs: [check-changes, pr-gate] if: needs.check-changes.outputs.main_package == 'true' runs-on: xeon-gnr env: HF_HOME: /home/sdp/.cache/huggingface strategy: matrix: build_type: ['all'] steps: - name: Checkout repository uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Build and Push run: | version=$(cat python/sglang/version.py | cut -d'"' -f2) tag=v${version}-xeon PR_REPO=${{ github.event.pull_request.head.repo.clone_url }} PR_HEAD_REF=${{ github.head_ref }} docker build \ ${PR_REPO:+--build-arg SGLANG_REPO=$PR_REPO} \ ${PR_HEAD_REF:+--build-arg VER_SGLANG=$PR_HEAD_REF} \ . -f docker/xeon.Dockerfile -t sglang_xeon --no-cache - name: Run container run: | docker run -dt \ -v ${{ github.workspace }}:/sglang-checkout/ --ipc=host \ -v ${HF_HOME}:/root/.cache/huggingface \ --name ci_sglang_xeon \ sglang_xeon - name: Check AMX support id: check_amx timeout-minutes: 5 run: | docker exec -w /sglang-checkout/ ci_sglang_xeon \ bash -c "source /opt/.venv/bin/activate && python3 -c 'import torch; import sgl_kernel; assert torch._C._cpu._is_amx_tile_supported(); assert hasattr(torch.ops.sgl_kernel, \"convert_weight_packed\"); '" - name: Run unit tests timeout-minutes: 36 run: | docker exec -w /sglang-checkout/ ci_sglang_xeon \ bash -c "source /opt/.venv/bin/activate && cd ./test/srt && python3 run_suite.py --suite per-commit-cpu --timeout-per-file 1500" - name: Change permission timeout-minutes: 2 run: | docker exec -u root ci_sglang_xeon bash -c " rm -rf /tmp/ci-home && chown -R $(id -u):$(id -g) /sglang-checkout/ 2>/dev/null || true " - name: Cleanup container if: always() run: | docker rm -f ci_sglang_xeon || true ================================================ FILE: .github/workflows/pr-test-xpu.yml ================================================ name: PR Test (XPU) on: push: branches: [ main ] pull_request: branches: [ main ] workflow_dispatch: workflow_call: inputs: ref: description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.' required: false type: string default: '' run_all_tests: description: "Run all tests (for releasing or testing purpose)" required: false type: boolean default: false concurrency: group: pr-test-xpu-${{ inputs.ref || github.ref }} cancel-in-progress: ${{ github.event_name != 'workflow_call' }} jobs: # ==================== Check Changes ==================== # check-changes: runs-on: ubuntu-latest outputs: main_package: ${{ steps.filter.outputs.main_package || steps.run-mode.outputs.run_all_tests }} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.ref || github.ref }} - name: Determine run mode id: run-mode run: | # Run all tests for workflow_call (when ref input is provided) # Note: github.event_name is inherited from caller, so we detect workflow_call by checking inputs.ref if [[ "${{ inputs.run_all_tests }}" == "true" ]]; then echo "run_all_tests=true" >> $GITHUB_OUTPUT echo "Run mode: ALL TESTS (run_all_tests=${{ inputs.run_all_tests }})" else echo "run_all_tests=false" >> $GITHUB_OUTPUT echo "Run mode: FILTERED (triggered by ${{ github.event_name }})" fi - name: Detect file changes id: filter uses: dorny/paths-filter@v3 if: steps.run-mode.outputs.run_all_tests != 'true' with: filters: | main_package: - "python/sglang/!(multimodal_gen)/**" - "python/pyproject_xpu.toml" - "test/**" - "sgl-kernel/**" - ".github/workflows/pr-test-xpu.yml" - "docker/xpu.Dockerfile" # ==================== PR Gate ==================== # pr-gate: needs: check-changes if: needs.check-changes.outputs.main_package == 'true' uses: ./.github/workflows/pr-gate.yml secrets: inherit build-and-test: needs: [check-changes, pr-gate] if: needs.check-changes.outputs.main_package == 'true' runs-on: intel-bmg env: HF_HOME: /home/sdp/.cache/huggingface steps: - name: Checkout code uses: actions/checkout@v4 with: fetch-depth: 0 ref: ${{ inputs.ref || github.ref }} - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - name: Build Docker image run: | PR_REPO=${{ github.event.pull_request.head.repo.clone_url }} PR_HEAD_REF=${{ github.head_ref }} docker build \ ${PR_REPO:+--build-arg SG_LANG_REPO=$PR_REPO} \ ${PR_HEAD_REF:+--build-arg SG_LANG_BRANCH=$PR_HEAD_REF} \ --no-cache --progress=plain -f docker/xpu.Dockerfile -t xpu_sglang_main:bmg . - name: Run container id: start_container run: | container_id=$(docker run -dt \ --group-add 992 \ --group-add $(getent group video | cut -d: -f3) \ -v ${HF_HOME}:/root/.cache/huggingface \ --device /dev/dri \ -e HF_TOKEN="$(cat ~/huggingface_token.txt)" \ xpu_sglang_main:bmg) echo "Started container: $container_id" echo "container_id=$container_id" >> "$GITHUB_OUTPUT" - name: Install Dependency timeout-minutes: 20 run: | cid="${{ steps.start_container.outputs.container_id }}" docker exec "$cid" /home/sdp/miniforge3/envs/py3.10/bin/python3 -m pip install --upgrade pip docker exec "$cid" /home/sdp/miniforge3/envs/py3.10/bin/python3 -m pip install pytest expecttest ray huggingface_hub docker exec "$cid" /home/sdp/miniforge3/envs/py3.10/bin/python3 -m pip uninstall -y flashinfer-python docker exec "$cid" /bin/bash -c '/home/sdp/miniforge3/envs/py3.10/bin/hf auth login --token ${HF_TOKEN} ' - name: Run E2E Bfloat16 tests timeout-minutes: 20 run: | cid="${{ steps.start_container.outputs.container_id }}" docker exec "$cid" bash -c "source /home/sdp/miniforge3/bin/activate && conda activate py3.10 && cd /home/sdp/sglang/test/srt && python3 run_suite.py --suite per-commit-xpu" - name: Cleanup container if: always() run: | cid="${{ steps.start_container.outputs.container_id }}" docker rm -f "$cid" || true finish: if: always() needs: [build-and-test, pr-gate] runs-on: ubuntu-latest steps: - name: Check job status run: | result="${{ needs.build-and-test.result }}" if [ "$result" != "success" ] && [ "$result" != "skipped" ]; then echo "Job failed with result: $result" exit 1 fi echo "All jobs completed successfully (result: $result)" exit 0 ================================================ FILE: .github/workflows/pr-test.yml ================================================ name: PR Test # Dynamic run-name for /rerun-stage commands to enable URL lookup # Format: "[stage-name] sha" for fork PRs, "[stage-name]" for non-fork, default for normal runs run-name: ${{ inputs.target_stage && (inputs.pr_head_sha && format('[{0}] {1}', inputs.target_stage, inputs.pr_head_sha) || format('[{0}]', inputs.target_stage)) || '' }} on: schedule: - cron: '0 */12 * * *' # Run every 12 hours pull_request: branches: [main] workflow_dispatch: inputs: target_stage: description: "Specific stage to run (optional, for quick testing)" required: false type: string default: "" force_continue_on_error: description: "Force continue-on-error (test scheduled CI behavior)" required: false type: boolean default: false pr_head_sha: description: "PR head SHA to checkout (for /rerun-stage on fork PRs)" required: false type: string default: "" test_parallel_dispatch: description: "Test parallel dispatch behavior (simulates scheduled run)" required: false type: boolean default: false workflow_call: inputs: git_ref: description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.' required: false type: string default: '' run_all_tests: description: "Run all tests (for releasing or testing purpose)" required: false type: boolean default: false concurrency: # Concurrency group structure: pr-test-{event}-{branch}-{pr_sha}-{stage} # - event_name prevents scheduled runs from colliding with fork PRs whose branch is named 'main' # (without it, both resolve the branch segment to 'main' and block each other) # - github.head_ref (pull_request) or github.ref_name (workflow_dispatch) normalizes to branch name # - pr_head_sha isolates /rerun-stage from main branch runs # - target_stage allows parallel stage dispatches to run independently group: pr-test-${{ github.event_name }}-${{ github.head_ref || github.ref_name || 'default' }}-${{ inputs.pr_head_sha || 'current' }}-${{ inputs.target_stage || inputs.git_ref || 'all' }} cancel-in-progress: ${{ github.event_name != 'workflow_call' }} env: SGLANG_IS_IN_CI: true SGLANG_CUDA_COREDUMP: "1" SGLANG_JIT_DEEPGEMM_FAST_WARMUP: true permissions: actions: write contents: read pull-requests: read jobs: # =============================================== check changes ==================================================== check-changes: runs-on: ubuntu-latest outputs: # Use API-based detection for target_stage mode (filter-api), otherwise use dorny/paths-filter (filter) main_package: ${{ steps.filter-api.outputs.main_package || steps.filter.outputs.main_package || steps.run-mode.outputs.run_all_tests }} # sgl_kernel is forced to false when target_stage is set, since sgl-kernel-build-wheels won't run # This prevents CUSTOM_BUILD_SGL_KERNEL=true when the wheel artifacts aren't available # Note: If PR has kernel changes AND target_stage is set, the validate-target-stage step will fail sgl_kernel: ${{ !inputs.target_stage && (steps.filter-api.outputs.sgl_kernel || steps.filter.outputs.sgl_kernel) }} # Raw sgl_kernel value before target_stage override (used for validation) sgl_kernel_raw: ${{ steps.filter-api.outputs.sgl_kernel || steps.filter.outputs.sgl_kernel }} jit_kernel: ${{ steps.filter-api.outputs.jit_kernel || steps.filter.outputs.jit_kernel || steps.run-mode.outputs.run_all_tests }} multimodal_gen: ${{ steps.filter-api.outputs.multimodal_gen || steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }} max_parallel: ${{ steps.set-parallel.outputs.max_parallel }} b200_runner: ${{ steps.set-runner.outputs.b200_runner }} enable_retry: ${{ steps.set-retry.outputs.enable_retry }} continue_on_error: ${{ steps.set-continue-on-error.outputs.continue_on_error }} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Show test partition assignments continue-on-error: true run: python3 test/show_partitions.py - name: Determine run mode id: run-mode run: | # Run all tests for scheduled runs and workflow_call (when ref input is provided) # Note: github.event_name is inherited from caller, so we detect workflow_call by checking inputs.git_ref if [[ "${{ github.event_name }}" == "schedule" || "${{ inputs.run_all_tests }}" == "true" ]]; then echo "run_all_tests=true" >> $GITHUB_OUTPUT echo "Run mode: ALL TESTS (schedule=${{ github.event_name == 'schedule' }}, run_all_tests=${{ inputs.run_all_tests }})" else echo "run_all_tests=false" >> $GITHUB_OUTPUT echo "Run mode: FILTERED (triggered by ${{ github.event_name }})" fi - name: Detect file changes id: filter uses: dorny/paths-filter@v3 # Only use paths-filter for pull_request events (where it works correctly) # For workflow_dispatch with target_stage, we use GitHub API in the next step if: steps.run-mode.outputs.run_all_tests != 'true' && !inputs.target_stage with: filters: | main_package: - ".github/workflows/pr-test.yml" - "python/pyproject.toml" - "python/sglang/!(multimodal_gen)/**" - "scripts/ci/cuda/*" - "scripts/ci/utils/*" - "test/**" multimodal_gen: - ".github/workflows/pr-test.yml" - "python/pyproject.toml" - "python/sglang/multimodal_gen/**" - "python/sglang/jit_kernel/**" - "python/sglang/cli/**" jit_kernel: - ".github/workflows/pr-test.yml" - "python/pyproject.toml" - "python/sglang/jit_kernel/**" sgl_kernel: - "sgl-kernel/**" # For /rerun-stage (workflow_dispatch with target_stage), dorny/paths-filter doesn't work # correctly because it falls back to "last commit" detection which breaks for merge commits. # Instead, we use the GitHub API to compare the PR commit against main. - name: Detect file changes via API (for target_stage) id: filter-api if: inputs.target_stage && inputs.pr_head_sha env: GH_TOKEN: ${{ github.token }} run: | echo "Detecting file changes via GitHub API for target_stage mode..." echo "PR head SHA: ${{ inputs.pr_head_sha }}" # Get the list of changed files by comparing PR commit against main # This correctly handles merge commits by looking at the actual PR diff CHANGED_FILES=$(gh api "repos/${{ github.repository }}/compare/main...${{ inputs.pr_head_sha }}" \ --jq '[.files[].filename] | .[]' 2>/dev/null || echo "") if [ -z "$CHANGED_FILES" ]; then echo "Warning: Could not fetch changed files from API, assuming no changes" echo "sgl_kernel=false" >> $GITHUB_OUTPUT echo "main_package=false" >> $GITHUB_OUTPUT echo "jit_kernel=false" >> $GITHUB_OUTPUT echo "multimodal_gen=false" >> $GITHUB_OUTPUT exit 0 fi echo "Changed files:" echo "$CHANGED_FILES" | head -20 echo "..." # Check for sgl-kernel changes if echo "$CHANGED_FILES" | grep -q "^sgl-kernel/"; then echo "sgl_kernel=true" >> $GITHUB_OUTPUT echo "Detected sgl-kernel changes" else echo "sgl_kernel=false" >> $GITHUB_OUTPUT fi # Check for main_package changes (excluding multimodal_gen) # Note: Need to filter out multimodal_gen before checking, not pipe grep -q output MAIN_PKG_FILES=$(echo "$CHANGED_FILES" | grep -E "^(python/sglang/|python/pyproject\.toml|scripts/ci/cuda/|scripts/ci/utils/|test/|\.github/workflows/pr-test\.yml)" | grep -v "^python/sglang/multimodal_gen/" || true) if [ -n "$MAIN_PKG_FILES" ]; then echo "main_package=true" >> $GITHUB_OUTPUT echo "Detected main_package changes" else echo "main_package=false" >> $GITHUB_OUTPUT fi # Check for jit_kernel changes if echo "$CHANGED_FILES" | grep -qE "^(python/sglang/jit_kernel/|python/pyproject\.toml|\.github/workflows/pr-test\.yml)"; then echo "jit_kernel=true" >> $GITHUB_OUTPUT echo "Detected jit_kernel changes" else echo "jit_kernel=false" >> $GITHUB_OUTPUT fi # Check for multimodal_gen changes if echo "$CHANGED_FILES" | grep -qE "^(python/sglang/multimodal_gen/|python/sglang/cli/|python/pyproject\.toml|\.github/workflows/pr-test\.yml)"; then echo "multimodal_gen=true" >> $GITHUB_OUTPUT echo "Detected multimodal_gen changes" else echo "multimodal_gen=false" >> $GITHUB_OUTPUT fi - name: Set max-parallel based on run type id: set-parallel env: GH_TOKEN: ${{ github.token }} run: | # Scheduled runs and high-priority PRs get full parallelism if [[ "${{ github.event_name }}" == "schedule" ]]; then echo "max_parallel=14" >> $GITHUB_OUTPUT echo "Scheduled run detected, setting max_parallel to 14" elif [[ "${{ github.event_name }}" == "pull_request" && "${{ contains(github.event.pull_request.labels.*.name, 'high priority') }}" == "true" ]]; then echo "max_parallel=14" >> $GITHUB_OUTPUT echo "High priority PR detected, setting max_parallel to 14" elif [[ -n "${{ inputs.target_stage }}" ]]; then # /rerun-stage (workflow_dispatch): query PR labels via GitHub API # Try SHA lookup first (fork PRs), fallback to branch name (non-fork PRs) LABELS="" PR_HEAD_SHA="${{ inputs.pr_head_sha }}" if [[ -n "$PR_HEAD_SHA" ]]; then LABELS=$(gh api "repos/${{ github.repository }}/commits/${PR_HEAD_SHA}/pulls" \ --jq '.[0].labels[].name' 2>/dev/null || true) fi if [[ -z "$LABELS" ]]; then LABELS=$(gh pr list --head "${{ github.ref_name }}" --repo "${{ github.repository }}" \ --json labels --jq '.[0].labels[].name' 2>/dev/null || true) fi echo "PR labels: ${LABELS:-"(none)"}" if echo "$LABELS" | grep -Fxq "high priority"; then echo "max_parallel=14" >> $GITHUB_OUTPUT echo "High priority PR detected via API (/rerun-stage), setting max_parallel to 14" else echo "max_parallel=3" >> $GITHUB_OUTPUT echo "Using default max_parallel of 3 (/rerun-stage, no high priority label)" fi else echo "max_parallel=3" >> $GITHUB_OUTPUT echo "Using default max_parallel of 3" fi - name: Set B200 runner tag id: set-runner run: | # Use kernel-build runner only when sgl_kernel changes are detected AND we're not in target_stage mode # (target_stage skips wheel builds, so we can't use custom kernels) # Use API-based detection (filter-api) for target_stage mode, otherwise use dorny/paths-filter (filter) sgl_kernel="${{ steps.filter-api.outputs.sgl_kernel || steps.filter.outputs.sgl_kernel }}" target_stage="${{ inputs.target_stage }}" if [[ "$sgl_kernel" == "true" && -z "$target_stage" ]]; then echo "b200_runner=4-gpu-b200-kernel" >> $GITHUB_OUTPUT else echo "b200_runner=4-gpu-b200" >> $GITHUB_OUTPUT fi - name: Enable retry for CI id: set-retry run: | echo "enable_retry=true" >> $GITHUB_OUTPUT echo "Retry logic enabled for CI" - name: Set continue-on-error for full test runs id: set-continue-on-error run: | if [[ "${{ steps.run-mode.outputs.run_all_tests }}" == "true" || "${{ inputs.force_continue_on_error }}" == "true" ]]; then echo "continue_on_error=true" >> $GITHUB_OUTPUT echo "Full test run or force flag detected, enabling continue-on-error to run all tests" else echo "continue_on_error=false" >> $GITHUB_OUTPUT echo "Filtered run, continue-on-error disabled" fi - name: Validate target_stage with kernel changes # Use API-based detection (filter-api) for target_stage mode, otherwise use dorny/paths-filter (filter) if: inputs.target_stage && (steps.filter-api.outputs.sgl_kernel == 'true' || steps.filter.outputs.sgl_kernel == 'true') run: | echo "::error::Cannot use /rerun-stage when PR has sgl-kernel changes." echo "::error::The sgl-kernel-build-wheels job is skipped in target_stage mode, but this PR modifies sgl-kernel/ files." echo "::error::Please use /tag-and-rerun-ci to run the full workflow including kernel builds." echo "" echo "ERROR: Cannot use /rerun-stage when PR has sgl-kernel changes." echo "" echo "This PR modifies files in sgl-kernel/, which requires building custom kernel wheels." echo "The /rerun-stage command skips the wheel build job, so the test would run against" echo "the wrong (PyPI) version of sgl-kernel instead of your changes." echo "" echo "To properly test your kernel changes, use one of these commands instead:" echo " /tag-and-rerun-ci - Re-run the full workflow including kernel builds" echo " /rerun-ci - Re-run the full workflow" echo "" exit 1 - name: Show filter results in summary (table) run: | { echo "## Change Detection" echo "" echo "| Component | Changed |" echo "|-------------------|---------|" echo "| main_package | ${{ steps.filter-api.outputs.main_package || steps.filter.outputs.main_package || steps.run-mode.outputs.run_all_tests }} |" echo "| sgl_kernel (raw) | ${{ steps.filter-api.outputs.sgl_kernel || steps.filter.outputs.sgl_kernel }} |" echo "| sgl_kernel (used) | ${{ !inputs.target_stage && (steps.filter-api.outputs.sgl_kernel || steps.filter.outputs.sgl_kernel) }} |" echo "| jit_kernel | ${{ steps.filter-api.outputs.jit_kernel || steps.filter.outputs.jit_kernel || steps.run-mode.outputs.run_all_tests }} |" echo "| multimodal_gen | ${{ steps.filter-api.outputs.multimodal_gen || steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }} |" echo "| target_stage | ${{ inputs.target_stage || '(none)' }} |" echo "| detection_method | ${{ inputs.target_stage && 'GitHub API' || 'dorny/paths-filter' }} |" echo "| max_parallel | ${{ steps.set-parallel.outputs.max_parallel }} |" echo "| b200_runner | ${{ steps.set-runner.outputs.b200_runner }} |" echo "| enable_retry | ${{ steps.set-retry.outputs.enable_retry }} |" echo "| continue_on_error | ${{ steps.set-continue-on-error.outputs.continue_on_error }} |" } >> $GITHUB_STEP_SUMMARY # =============================================== Wait Jobs for Sequential PR Execution ==================================================== # These jobs poll GitHub API to wait for previous stages to complete. # For PR runs: wait jobs run and enforce sequential execution via polling. # For scheduled runs: wait jobs are skipped, enabling parallel execution for easier retry. wait-for-stage-a: needs: [check-changes, call-gate] # Only run for PRs (not scheduled) and when not targeting a specific stage # Skip if call-gate failed (stage-a jobs will be skipped, nothing to wait for) # !cancelled() ensures this job respects workflow cancellation from concurrency group if: | always() && !cancelled() && github.event_name == 'pull_request' && !inputs.target_stage && inputs.test_parallel_dispatch != true && (needs.check-changes.outputs.main_package == 'true' || needs.check-changes.outputs.sgl_kernel == 'true') && (needs.call-gate.result == 'success' || needs.call-gate.result == 'skipped') runs-on: ubuntu-latest outputs: stage_a_result: ${{ steps.wait.outputs.result }} steps: - uses: actions/checkout@v4 - uses: ./.github/actions/wait-for-jobs id: wait with: stage-name: stage-a jobs: '["stage-a-test-1", "stage-a-cpu-only"]' max-wait-minutes: '240' wait-for-stage-b: needs: [check-changes, call-gate, wait-for-stage-a] # Only run for PRs (not scheduled) and when not targeting a specific stage # Skip if call-gate failed (stage-b jobs will be skipped, nothing to wait for) if: | always() && !cancelled() && github.event_name == 'pull_request' && !inputs.target_stage && inputs.test_parallel_dispatch != true && (needs.check-changes.outputs.main_package == 'true' || needs.check-changes.outputs.sgl_kernel == 'true') && (needs.wait-for-stage-a.result == 'success' || needs.wait-for-stage-a.result == 'skipped') && (needs.call-gate.result == 'success' || needs.call-gate.result == 'skipped') runs-on: ubuntu-latest outputs: stage_b_result: ${{ steps.wait.outputs.result }} steps: - uses: actions/checkout@v4 - uses: ./.github/actions/wait-for-jobs id: wait with: stage-name: stage-b jobs: | [ {"prefix": "stage-b-test-small-1-gpu", "expected_count": 8}, {"prefix": "stage-b-test-large-1-gpu", "expected_count": 14}, {"prefix": "stage-b-test-large-2-gpu", "expected_count": 4}, {"prefix": "stage-b-test-4-gpu-b200", "expected_count": 1} ] max-wait-minutes: '480' # =============================================== PR Gate ==================================================== call-gate: needs: check-changes # Skip for scheduled runs (they run all tests) and when target_stage is specified if: | github.event_name != 'schedule' && inputs.test_parallel_dispatch != true && !inputs.target_stage && ( needs.check-changes.outputs.main_package == 'true' || needs.check-changes.outputs.sgl_kernel == 'true' || needs.check-changes.outputs.jit_kernel == 'true' || needs.check-changes.outputs.multimodal_gen == 'true' ) uses: ./.github/workflows/pr-gate.yml secrets: inherit # =============================================== sgl-kernel ==================================================== sgl-kernel-build-wheels: needs: [check-changes, call-gate] # Skip for scheduled runs (they run stages independently) and when target_stage is set if: github.event_name != 'schedule' && inputs.test_parallel_dispatch != true && !inputs.target_stage && needs.check-changes.outputs.sgl_kernel == 'true' runs-on: x64-kernel-build-node timeout-minutes: 240 strategy: matrix: include: - python-version: "3.10" cuda-version: "12.9" # Add back when CUDA 13.0 is supported on CI # - python-version: "3.10" # cuda-version: "13.0" name: Build Wheel steps: - name: Cleanup run: | sudo rm -rf $GITHUB_WORKSPACE/* || true - uses: actions/checkout@v4 with: submodules: "recursive" ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Build wheel for Python ${{ matrix.python-version }} and CUDA ${{ matrix.cuda-version }} run: | cd sgl-kernel ./build.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" env: USE_CCACHE: 1 - name: Verify wheel artifacts run: | ls -alh sgl-kernel/dist ls -alh sgl-kernel/dist/*.whl - name: Upload artifacts uses: actions/upload-artifact@v4 with: name: wheel-python${{ matrix.python-version }}-cuda${{ matrix.cuda-version }} path: sgl-kernel/dist/* if-no-files-found: error sgl-kernel-build-wheels-arm: needs: [check-changes, call-gate] # Skip for scheduled runs (they run stages independently) and when target_stage is set if: github.event_name != 'schedule' && inputs.test_parallel_dispatch != true && !inputs.target_stage && needs.check-changes.outputs.sgl_kernel == 'true' runs-on: arm-kernel-build-node timeout-minutes: 240 strategy: matrix: include: - python-version: "3.10" cuda-version: "12.9" name: Build Wheel Arm steps: - name: Cleanup run: | if [ -d "$GITHUB_WORKSPACE" ]; then sudo rm -rf "$GITHUB_WORKSPACE"/* || true else echo "$GITHUB_WORKSPACE does not exist, nothing to clean" fi - uses: actions/checkout@v4 with: submodules: "recursive" ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Build wheel for Python ${{ matrix.python-version }} and CUDA ${{ matrix.cuda-version }} run: | cd sgl-kernel ./build.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" env: USE_CCACHE: 1 - name: Verify wheel artifacts run: | ls -alh sgl-kernel/dist ls -alh sgl-kernel/dist/*.whl - name: Upload artifacts uses: actions/upload-artifact@v4 with: name: wheel-python${{ matrix.python-version }}-cuda${{ matrix.cuda-version }}-aarch64 path: sgl-kernel/dist/* if-no-files-found: error sgl-kernel-unit-test: needs: [check-changes, call-gate, sgl-kernel-build-wheels] # Skip for scheduled runs and when target_stage is set if: | github.event_name != 'schedule' && inputs.test_parallel_dispatch != true && !inputs.target_stage && needs.check-changes.outputs.sgl_kernel == 'true' runs-on: 1-gpu-runner timeout-minutes: 240 steps: - uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Cleanup run: | ls -alh sgl-kernel/dist || true rm -rf sgl-kernel/dist/* || true - name: Download artifacts uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion - name: Run test timeout-minutes: 30 run: | cd sgl-kernel pytest tests/ sgl-kernel-mla-test: needs: [check-changes, call-gate, sgl-kernel-build-wheels] # Skip for scheduled runs and when target_stage is set if: | github.event_name != 'schedule' && inputs.test_parallel_dispatch != true && !inputs.target_stage && needs.check-changes.outputs.sgl_kernel == 'true' runs-on: 1-gpu-runner timeout-minutes: 240 steps: - uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Cleanup run: | ls -alh sgl-kernel/dist || true rm -rf sgl-kernel/dist/* || true - name: Download artifacts uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | cd test/registered/mla python3 test_mla_deepseek_v3.py sgl-kernel-benchmark-test: needs: [check-changes, call-gate, sgl-kernel-build-wheels] # Skip for scheduled runs and when target_stage is set if: | github.event_name != 'schedule' && inputs.test_parallel_dispatch != true && !inputs.target_stage && needs.check-changes.outputs.sgl_kernel == 'true' runs-on: 1-gpu-runner timeout-minutes: 240 env: CI: true steps: - uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Cleanup run: | ls -alh sgl-kernel/dist || true rm -rf sgl-kernel/dist/* || true - name: Download artifacts uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh - name: Run benchmark tests timeout-minutes: 45 run: | cd sgl-kernel/benchmark echo "Running sgl-kernel benchmark tests in CI mode..." echo "CI environment variable: $CI" echo "GITHUB_ACTIONS environment variable: $GITHUB_ACTIONS" for bench_file in bench_*.py; do echo "Testing $bench_file..." timeout 60 python3 "$bench_file" || echo "Warning: $bench_file timed out or failed, continuing..." echo "Completed $bench_file" echo "---" done echo "All benchmark tests completed!" sgl-kernel-b200-test: needs: [check-changes, sgl-kernel-build-wheels] # Skip for scheduled runs and when target_stage is set if: | github.event_name != 'schedule' && inputs.test_parallel_dispatch != true && !inputs.target_stage && needs.check-changes.outputs.sgl_kernel == 'true' runs-on: ${{ needs.check-changes.outputs.b200_runner }} timeout-minutes: 240 steps: - uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Cleanup run: | ls -alh sgl-kernel/dist || true rm -rf sgl-kernel/dist/* || true - name: Download artifacts uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} IS_BLACKWELL=1 bash scripts/ci/cuda/ci_install_dependency.sh diffusion - name: Run sgl-kernel unit tests on B200 timeout-minutes: 30 run: | cd sgl-kernel pytest tests/ # Adding a single CUDA13 smoke test to verify that the kernel builds and runs # TODO: Add back this test when it can pass on CI # cuda13-kernel-smoke-test: # needs: [check-changes, sgl-kernel-build-wheels] # if: needs.check-changes.outputs.sgl_kernel == 'true' # runs-on: x64-cu13-kernel-tests # steps: # - uses: actions/checkout@v4 # - name: Cleanup # run: | # ls -alh sgl-kernel/dist || true # rm -rf sgl-kernel/dist/* || true # - name: Download CUDA 13.0 artifacts # uses: actions/download-artifact@v4 # with: # path: sgl-kernel/dist/ # merge-multiple: true # pattern: wheel-python3.10-cuda13.0 # - name: Install dependencies # run: | # CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh # - name: Run kernel unit tests # timeout-minutes: 30 # run: | # cd sgl-kernel # pytest tests/ # =============================================== jit-kernel ==================================================== jit-kernel-unit-test: needs: [check-changes, call-gate] # Skip for scheduled runs and when target_stage is set if: | github.event_name != 'schedule' && inputs.test_parallel_dispatch != true && !inputs.target_stage && needs.check-changes.outputs.jit_kernel == 'true' runs-on: 1-gpu-runner timeout-minutes: 240 steps: - uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Install dependencies timeout-minutes: 20 run: | bash scripts/ci/cuda/ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | cd python/sglang/jit_kernel pytest tests/ jit-kernel-unit-test-nightly: needs: [check-changes] if: | github.event_name == 'schedule' && needs.check-changes.outputs.jit_kernel == 'true' runs-on: 1-gpu-runner timeout-minutes: 240 env: SGLANG_JIT_KERNEL_RUN_FULL_TESTS: "1" steps: - uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Install dependencies timeout-minutes: 20 run: | bash scripts/ci/cuda/ci_install_dependency.sh - name: Run full nightly test timeout-minutes: 60 run: | cd python/sglang/jit_kernel pytest tests/ jit-kernel-benchmark-test: needs: [check-changes, call-gate] # Skip for scheduled runs and when target_stage is set if: | github.event_name != 'schedule' && inputs.test_parallel_dispatch != true && !inputs.target_stage && needs.check-changes.outputs.jit_kernel == 'true' runs-on: 1-gpu-runner timeout-minutes: 240 env: CI: true steps: - uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Install dependencies timeout-minutes: 20 run: | bash scripts/ci/cuda/ci_install_dependency.sh - name: Run benchmark tests timeout-minutes: 45 run: | cd python/sglang/jit_kernel/benchmark echo "Running jit-kernel benchmark tests in CI mode..." failures=() for bench_file in bench_*.py; do echo "Testing $bench_file..." if ! timeout 120 python3 "$bench_file"; then failures+=("$bench_file") fi echo "Completed $bench_file" echo "---" done if [ ${#failures[@]} -ne 0 ]; then echo "The following benchmark tests failed: ${failures[*]}" exit 1 fi echo "All jit-kernel benchmark tests completed successfully!" # =============================================== primary ==================================================== stage-a-test-1: needs: [check-changes, call-gate, sgl-kernel-build-wheels] if: | always() && ( (inputs.target_stage == 'stage-a-test-1') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) runs-on: 1-gpu-runner timeout-minutes: 240 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh - name: Run test timeout-minutes: 10 run: | cd test/ CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi python3 run_suite.py --hw cuda --suite stage-a-test-1 $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: always() stage-a-cpu-only: needs: [check-changes, call-gate] if: | always() && ( (inputs.target_stage == 'stage-a-cpu-only') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && (needs.check-changes.outputs.main_package == 'true') ) ) runs-on: ubuntu-latest timeout-minutes: 240 steps: - name: Free disk space run: | sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc df -h - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Install dependencies timeout-minutes: 20 run: | pip install -e "python/[dev]" - name: Run test timeout-minutes: 10 run: | cd test/ CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi python3 run_suite.py --hw cpu --suite stage-a-cpu-only $CONTINUE_ON_ERROR_FLAG # Runs on 5090 (32GB, SM120) stage-b-test-small-1-gpu: needs: [check-changes, call-gate, wait-for-stage-a, sgl-kernel-build-wheels] if: | always() && ( (inputs.target_stage == 'stage-b-test-small-1-gpu') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) runs-on: 1-gpu-5090 timeout-minutes: 240 env: IS_BLACKWELL: "1" strategy: fail-fast: false max-parallel: 8 matrix: partition: [0, 1, 2, 3, 4, 5, 6, 7] steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | source /etc/profile.d/sglang-ci.sh CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh git clone https://github.com/merrymercy/human-eval.git cd human-eval pip install -e . --no-build-isolation - name: Run test timeout-minutes: 30 run: | source /etc/profile.d/sglang-ci.sh cd test/ CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi python3 run_suite.py --hw cuda --suite stage-b-test-small-1-gpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size 8 $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: always() with: artifact-suffix: ${{ matrix.partition }} # Runs on H100 (80GB, SM90) - tests that don't pass on 5090 (FA3, FP8, high VRAM, etc.) stage-b-test-large-1-gpu: needs: [check-changes, call-gate, wait-for-stage-a, sgl-kernel-build-wheels] if: | always() && ( (inputs.target_stage == 'stage-b-test-large-1-gpu') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) runs-on: 1-gpu-runner timeout-minutes: 240 strategy: fail-fast: false max-parallel: ${{ fromJson(needs.check-changes.outputs.max_parallel) }} matrix: partition: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | cd test/ CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi python3 run_suite.py --hw cuda --suite stage-b-test-large-1-gpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size 14 --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: always() with: artifact-suffix: ${{ matrix.partition }} stage-b-test-large-2-gpu: needs: [check-changes, call-gate, wait-for-stage-a, sgl-kernel-build-wheels] if: | always() && ( (inputs.target_stage == 'stage-b-test-large-2-gpu') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) runs-on: 2-gpu-runner timeout-minutes: 240 strategy: fail-fast: false matrix: partition: [0, 1, 2, 3] steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh git clone https://github.com/merrymercy/human-eval.git cd human-eval pip install -e . --no-build-isolation - name: Run test timeout-minutes: 30 run: | cd test/ CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi python3 run_suite.py --hw cuda --suite stage-b-test-large-2-gpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size 4 $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: always() with: artifact-suffix: ${{ matrix.partition }} stage-b-test-4-gpu-b200: needs: [check-changes, call-gate, wait-for-stage-a, sgl-kernel-build-wheels] if: | always() && ( (inputs.target_stage == 'stage-b-test-4-gpu-b200') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) runs-on: ${{ needs.check-changes.outputs.b200_runner }} timeout-minutes: 240 strategy: fail-fast: false steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v6 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} IS_BLACKWELL=1 bash scripts/ci/cuda/ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | cd test CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi IS_BLACKWELL=1 python3 run_suite.py --hw cuda --suite stage-b-test-4-gpu-b200 $CONTINUE_ON_ERROR_FLAG - name: Run FA4 jit_kernel tests (SM100+) timeout-minutes: 10 run: | IS_BLACKWELL=1 python3 -m pytest -q python/sglang/jit_kernel/tests/test_flash_attention_4.py - uses: ./.github/actions/upload-cuda-coredumps if: always() multimodal-gen-test-1-gpu: needs: [check-changes, call-gate, sgl-kernel-build-wheels] if: | always() && ( (inputs.target_stage == 'multimodal-gen-test-1-gpu') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && needs.check-changes.outputs.multimodal_gen == 'true' ) ) runs-on: 1-gpu-runner timeout-minutes: 240 strategy: fail-fast: false matrix: part: [0, 1] steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion - name: Run diffusion server tests timeout-minutes: 240 env: RUNAI_STREAMER_MEMORY_LIMIT: 0 run: | cd python CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi python3 sglang/multimodal_gen/test/run_suite.py \ --suite 1-gpu \ --partition-id ${{ matrix.part }} \ --total-partitions 2 \ $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: always() with: artifact-suffix: ${{ matrix.part }} multimodal-gen-test-2-gpu: needs: [check-changes, call-gate, sgl-kernel-build-wheels] if: | always() && ( (inputs.target_stage == 'multimodal-gen-test-2-gpu') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && needs.check-changes.outputs.multimodal_gen == 'true' ) ) runs-on: 2-gpu-runner timeout-minutes: 240 strategy: fail-fast: false matrix: part: [0, 1] steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion - name: Run diffusion server tests timeout-minutes: 240 env: RUNAI_STREAMER_MEMORY_LIMIT: 0 run: | cd python CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi python3 sglang/multimodal_gen/test/run_suite.py \ --suite 2-gpu \ --partition-id ${{ matrix.part }} \ --total-partitions 2 \ $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: always() with: artifact-suffix: ${{ matrix.part }} multimodal-gen-unit-test: needs: [check-changes, call-gate, sgl-kernel-build-wheels] if: | always() && ( (inputs.target_stage == 'multimodal-gen-unit-test') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && needs.check-changes.outputs.multimodal_gen == 'true' ) ) runs-on: 1-gpu-runner timeout-minutes: 120 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion - name: Run diffusion unit tests timeout-minutes: 60 run: | cd python python3 sglang/multimodal_gen/test/run_suite.py --suite unit stage-c-test-4-gpu-h100: needs: [check-changes, call-gate, wait-for-stage-b] if: | always() && ( (inputs.target_stage == 'stage-c-test-4-gpu-h100') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) runs-on: 4-gpu-h100 timeout-minutes: 240 strategy: fail-fast: false matrix: part: [0, 1, 2] steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | cd test CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi python3 run_suite.py --hw cuda --suite stage-c-test-4-gpu-h100 --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: always() with: artifact-suffix: ${{ matrix.part }} stage-c-test-8-gpu-h200: needs: [check-changes, call-gate, wait-for-stage-b] if: | always() && ( (inputs.target_stage == 'stage-c-test-8-gpu-h200') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) runs-on: 8-gpu-h200 timeout-minutes: 240 strategy: fail-fast: false matrix: part: [0, 1, 2, 3] steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh - name: Warmup DeepGEMM JIT Compilation timeout-minutes: 25 run: | python3 scripts/ci/cuda/warmup_deep_gemm.py \ deepseek-ai/DeepSeek-V3-0324:8 \ deepseek-ai/DeepSeek-V3.2-Exp:8 - name: Warmup Server CUDA Graphs timeout-minutes: 25 run: | python3 scripts/ci/cuda/warmup_server.py \ deepseek-ai/DeepSeek-V3-0324:8 \ inclusionAI/Ring-2.5-1T:8 - name: Run test timeout-minutes: 30 run: | cd test CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi python3 run_suite.py --hw cuda --suite stage-c-test-8-gpu-h200 --auto-partition-id ${{ matrix.part }} --auto-partition-size 4 $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: always() with: artifact-suffix: ${{ matrix.part }} stage-c-test-8-gpu-h20: needs: [check-changes, call-gate, wait-for-stage-b] if: | always() && ( (inputs.target_stage == 'stage-c-test-8-gpu-h20') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) runs-on: 8-gpu-h20 timeout-minutes: 240 env: SGLANG_CI_RDMA_ALL_DEVICES: "mlx5_1,mlx5_2,mlx5_3,mlx5_4" strategy: fail-fast: false matrix: part: [0, 1] steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_deepep.sh - name: Run test timeout-minutes: 30 run: | cd test CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi python3 run_suite.py --hw cuda --suite stage-c-test-8-gpu-h20 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: always() with: artifact-suffix: ${{ matrix.part }} stage-c-test-deepep-4-gpu: needs: [check-changes, call-gate, wait-for-stage-b] if: | always() && ( (inputs.target_stage == 'stage-c-test-deepep-4-gpu') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) runs-on: 4-gpu-h100 timeout-minutes: 240 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_deepep.sh - name: Warmup DeepGEMM JIT Compilation timeout-minutes: 25 run: | python3 scripts/ci/cuda/warmup_deep_gemm.py \ lmsys/sglang-ci-dsv3-test:4 - name: Warmup Server CUDA Graphs timeout-minutes: 25 run: | python3 scripts/ci/cuda/warmup_server.py \ lmsys/sglang-ci-dsv3-test:4 - name: Run test timeout-minutes: 30 run: | cd test CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi python3 run_suite.py --hw cuda --suite stage-c-test-deepep-4-gpu $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: always() stage-c-test-deepep-8-gpu-h200: needs: [check-changes, call-gate, wait-for-stage-b] if: | always() && ( (inputs.target_stage == 'stage-c-test-deepep-8-gpu-h200') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) runs-on: 8-gpu-h200 timeout-minutes: 240 steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_deepep.sh - name: Warmup DeepGEMM JIT Compilation timeout-minutes: 25 run: | python3 scripts/ci/cuda/warmup_deep_gemm.py \ deepseek-ai/DeepSeek-V3-0324:8 \ deepseek-ai/DeepSeek-V3.2-Exp:8 - name: Warmup Server CUDA Graphs timeout-minutes: 25 run: | python3 scripts/ci/cuda/warmup_server.py \ deepseek-ai/DeepSeek-V3-0324:8 - name: Run test timeout-minutes: 45 run: | cd test CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi python3 run_suite.py --hw cuda --suite stage-c-test-deepep-8-gpu-h200 $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: always() stage-c-test-4-gpu-b200: needs: [check-changes, call-gate, wait-for-stage-b] if: | always() && ( (inputs.target_stage == 'stage-c-test-4-gpu-b200') || ( !inputs.target_stage && ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) ) ) runs-on: ${{ needs.check-changes.outputs.b200_runner }} timeout-minutes: 240 strategy: fail-fast: false matrix: part: [0, 1, 2, 3] steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} - name: Download artifacts if: needs.check-changes.outputs.sgl_kernel == 'true' uses: actions/download-artifact@v6 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 run: | CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} IS_BLACKWELL=1 bash scripts/ci/cuda/ci_install_dependency.sh - name: Run test timeout-minutes: 30 run: | cd test CONTINUE_ON_ERROR_FLAG="" if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then CONTINUE_ON_ERROR_FLAG="--continue-on-error" fi IS_BLACKWELL=1 python3 run_suite.py --hw cuda --suite stage-c-test-4-gpu-b200 --auto-partition-id ${{ matrix.part }} --auto-partition-size 4 --timeout-per-file 1800 $CONTINUE_ON_ERROR_FLAG - uses: ./.github/actions/upload-cuda-coredumps if: always() with: artifact-suffix: ${{ matrix.part }} # NOTE: GB200 stage temporarily disabled — no company-owned GB200 runner available yet. # Re-enable when a 4-gpu-gb200 runner is provisioned. # stage-c-test-4-gpu-gb200: # needs: [check-changes, call-gate, wait-for-stage-b, sgl-kernel-build-wheels-arm] # if: | # always() && # ( # (inputs.target_stage == 'stage-c-test-4-gpu-gb200') || # ( # !inputs.target_stage && # ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) && # ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true')) # ) # ) # runs-on: 4-gpu-gb200 # timeout-minutes: 240 # strategy: # fail-fast: false # steps: # - name: Checkout code # uses: actions/checkout@v4 # with: # ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} # # - name: Download artifacts # if: needs.check-changes.outputs.sgl_kernel == 'true' # uses: actions/download-artifact@v4 # with: # path: sgl-kernel/dist/ # merge-multiple: true # pattern: wheel-python3.10-cuda12.9-aarch64 # # - name: Install dependencies # timeout-minutes: 20 # run: | # CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} IS_BLACKWELL=1 GRACE_BLACKWELL=1 bash scripts/ci/cuda/ci_install_deepep.sh # # - name: Run test # timeout-minutes: 45 # run: | # cd test # CONTINUE_ON_ERROR_FLAG="" # if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then # CONTINUE_ON_ERROR_FLAG="--continue-on-error" # fi # python3 run_suite.py --hw cuda --suite stage-c-test-4-gpu-gb200 --timeout-per-file 3600 $CONTINUE_ON_ERROR_FLAG # # - uses: ./.github/actions/upload-cuda-coredumps # if: always() pr-test-finish: needs: [ call-gate, check-changes, sgl-kernel-build-wheels, sgl-kernel-unit-test, sgl-kernel-mla-test, sgl-kernel-benchmark-test, sgl-kernel-b200-test, wait-for-stage-a, wait-for-stage-b, jit-kernel-unit-test, jit-kernel-unit-test-nightly, jit-kernel-benchmark-test, multimodal-gen-unit-test, multimodal-gen-test-1-gpu, multimodal-gen-test-2-gpu, stage-a-test-1, stage-a-cpu-only, stage-b-test-small-1-gpu, stage-b-test-large-1-gpu, stage-b-test-large-2-gpu, stage-b-test-4-gpu-b200, stage-c-test-4-gpu-h100, stage-c-test-8-gpu-h20, stage-c-test-8-gpu-h200, stage-c-test-deepep-4-gpu, stage-c-test-deepep-8-gpu-h200, stage-c-test-4-gpu-b200, # stage-c-test-4-gpu-gb200, # Temporarily disabled — no GB200 runner ] if: always() runs-on: ubuntu-latest steps: - name: Check all dependent job statuses run: | # Convert the 'needs' context to a JSON string json_needs='${{ toJson(needs) }}' # Get a list of all job names from the JSON keys job_names=$(echo "$json_needs" | jq -r 'keys_unsorted[]') for job in $job_names; do # For each job, extract its result result=$(echo "$json_needs" | jq -r --arg j "$job" '.[$j].result') # Print the job name and its result echo "$job: $result" # Check for failure or cancellation and exit if found if [[ "$result" == "failure" || "$result" == "cancelled" ]]; then echo "The above jobs failed." exit 1 fi done # If the loop completes, all jobs were successful echo "All jobs completed successfully" exit 0 ================================================ FILE: .github/workflows/release-branch-cut.yml ================================================ name: Release Branch Cut on: workflow_dispatch: inputs: branch_name: description: 'Branch name to create (e.g., release/v0.5.7)' required: true type: string commit_sha: description: 'Commit SHA from main to cut the release branch from (defaults to latest main)' required: false type: string default: '' permissions: actions: write contents: write pull-requests: read jobs: cut-release-branch: if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-latest environment: 'prod' outputs: branch_name: ${{ steps.set_output.outputs.branch_name }} steps: - name: Checkout repository uses: actions/checkout@v4 with: ref: main fetch-depth: 0 token: ${{ secrets.GITHUB_TOKEN }} - name: Validate branch name run: | BRANCH_NAME="${{ github.event.inputs.branch_name }}" if [ -z "$BRANCH_NAME" ]; then echo "::error::Branch name is required" exit 1 fi # Validate branch name format (should start with release/) if [[ ! "$BRANCH_NAME" =~ ^release/ ]]; then echo "::warning::Branch name '$BRANCH_NAME' does not follow convention 'release/vX.Y.Z'" fi echo "Branch name: $BRANCH_NAME" - name: Validate commit SHA id: validate run: | COMMIT_SHA="${{ github.event.inputs.commit_sha }}" # If no commit SHA provided, use latest main if [ -z "$COMMIT_SHA" ]; then COMMIT_SHA=$(git rev-parse HEAD) echo "No commit SHA provided, using latest main: $COMMIT_SHA" fi # Verify the commit exists and is on main if ! git cat-file -t "$COMMIT_SHA" > /dev/null 2>&1; then echo "::error::Commit SHA '$COMMIT_SHA' does not exist" exit 1 fi # Check if commit is an ancestor of main (i.e., is on main branch) if ! git merge-base --is-ancestor "$COMMIT_SHA" main; then echo "::error::Commit SHA '$COMMIT_SHA' is not on the main branch" exit 1 fi echo "COMMIT_SHA=$COMMIT_SHA" >> $GITHUB_OUTPUT echo "Validated commit SHA: $COMMIT_SHA" - name: Check if branch already exists run: | BRANCH_NAME="${{ github.event.inputs.branch_name }}" if git ls-remote --heads origin "$BRANCH_NAME" | grep -q "$BRANCH_NAME"; then echo "::error::Branch '$BRANCH_NAME' already exists" exit 1 fi echo "Branch '$BRANCH_NAME' does not exist, proceeding with creation" - name: Create release branch id: set_output run: | COMMIT_SHA="${{ steps.validate.outputs.COMMIT_SHA }}" BRANCH_NAME="${{ github.event.inputs.branch_name }}" git config user.name "sglang-bot" git config user.email "sglang-bot@users.noreply.github.com" # Create branch from the specified commit git checkout -b "$BRANCH_NAME" "$COMMIT_SHA" echo "branch_name=$BRANCH_NAME" >> $GITHUB_OUTPUT echo "Successfully created branch '$BRANCH_NAME' from commit '$COMMIT_SHA'" - name: Update version references in documentation run: | BRANCH_NAME="${{ github.event.inputs.branch_name }}" # Extract version from branch name (e.g., release/v0.5.8 -> v0.5.8) VERSION=$(echo "$BRANCH_NAME" | sed 's/release\///') # Update git clone version references in docs sed -i "s/git clone -b v[0-9]\+\.[0-9]\+\.[0-9]\+\.\?post\?[0-9]*/git clone -b $VERSION/" docs/get_started/install.md sed -i "s/git clone -b v[0-9]\+\.[0-9]\+\.[0-9]\+\.\?post\?[0-9]*/git clone -b $VERSION/" docs/platforms/amd_gpu.md # Check if any changes were made if git diff --quiet; then echo "No version references needed updating" else git add docs/get_started/install.md docs/platforms/amd_gpu.md git commit -m "docs: update version references to $VERSION" echo "Updated version references to $VERSION" fi - name: Push release branch run: | BRANCH_NAME="${{ steps.set_output.outputs.branch_name }}" git push origin "$BRANCH_NAME" echo "Successfully pushed branch '$BRANCH_NAME'" - name: Summary run: | COMMIT_SHA="${{ steps.validate.outputs.COMMIT_SHA }}" BRANCH_NAME="${{ github.event.inputs.branch_name }}" echo "## Release Branch Cut Summary" >> $GITHUB_STEP_SUMMARY echo "" >> $GITHUB_STEP_SUMMARY echo "| Property | Value |" >> $GITHUB_STEP_SUMMARY echo "|----------|-------|" >> $GITHUB_STEP_SUMMARY echo "| Branch | \`$BRANCH_NAME\` |" >> $GITHUB_STEP_SUMMARY echo "| Commit | \`$COMMIT_SHA\` |" >> $GITHUB_STEP_SUMMARY echo "| Triggered by | @${{ github.actor }} |" >> $GITHUB_STEP_SUMMARY echo "" >> $GITHUB_STEP_SUMMARY echo "### Next Steps" >> $GITHUB_STEP_SUMMARY echo "1. Tests are automatically triggered on the release branch" >> $GITHUB_STEP_SUMMARY echo "2. Apply any hotfixes if needed" >> $GITHUB_STEP_SUMMARY echo "3. Create a tag to trigger release: \`gh workflow run release-tag.yml -f version=X.Y.Z -f ref=$BRANCH_NAME\`" >> $GITHUB_STEP_SUMMARY run-pr-tests-nvidia: needs: cut-release-branch uses: ./.github/workflows/pr-test.yml with: git_ref: ${{ needs.cut-release-branch.outputs.branch_name }} run_all_tests: true secrets: inherit run-pr-tests-amd: needs: cut-release-branch uses: ./.github/workflows/pr-test-amd.yml with: ref: ${{ needs.cut-release-branch.outputs.branch_name }} run_all_tests: true secrets: inherit run-pr-test-npu: needs: cut-release-branch uses: ./.github/workflows/pr-test-npu.yml with: ref: ${{ needs.cut-release-branch.outputs.branch_name }} run_all_tests: true secrets: inherit run-pr-tests-xeon: needs: cut-release-branch uses: ./.github/workflows/pr-test-xeon.yml with: ref: ${{ needs.cut-release-branch.outputs.branch_name }} run_all_tests: true secrets: inherit run-pr-tests-xpu: needs: cut-release-branch uses: ./.github/workflows/pr-test-xpu.yml with: ref: ${{ needs.cut-release-branch.outputs.branch_name }} run_all_tests: true secrets: inherit run-nightly-tests-nvidia: needs: cut-release-branch uses: ./.github/workflows/nightly-test-nvidia.yml with: ref: ${{ needs.cut-release-branch.outputs.branch_name }} secrets: inherit run-nightly-tests-amd: needs: cut-release-branch uses: ./.github/workflows/nightly-test-amd.yml with: ref: ${{ needs.cut-release-branch.outputs.branch_name }} secrets: inherit run-nightly-tests-npu: needs: cut-release-branch uses: ./.github/workflows/nightly-test-npu.yml with: ref: ${{ needs.cut-release-branch.outputs.branch_name }} secrets: inherit run-nightly-tests-intel: needs: cut-release-branch uses: ./.github/workflows/nightly-test-intel.yml with: ref: ${{ needs.cut-release-branch.outputs.branch_name }} secrets: inherit ================================================ FILE: .github/workflows/release-docker-amd-nightly.yml ================================================ name: Release Docker Images Nightly (AMD) on: workflow_dispatch: schedule: - cron: '0 12 * * *' concurrency: # A PR number if a pull request and otherwise the commit hash. This cancels # queued and in-progress runs for the same PR (presubmit) or commit # (postsubmit). The workflow name is prepended to avoid conflicts between # different workflows. group: ${{ github.workflow }}-${{ github.event.number || github.sha }} cancel-in-progress: true jobs: publish: if: github.repository == 'sgl-project/sglang' runs-on: amd-docker-scale environment: 'prod' strategy: fail-fast: false matrix: gpu_arch: ['gfx942', 'gfx950'] build_type: ['all'] steps: - name: Checkout repository uses: actions/checkout@v4 with: fetch-depth: 0 # Required for git describe to find tags - name: "Set Date" run: | echo "DATE=$(date +%Y%m%d)" >> $GITHUB_ENV - name: Get version from latest tag id: version run: | # Get the latest version tag sorted by version number (e.g., v0.5.7 -> 0.5.7) VERSION=$(git tag -l 'v[0-9]*' --sort=-v:refname | head -1 | sed 's/^v//') if [ -z "$VERSION" ]; then echo "::error::Could not determine version from git tags" exit 1 fi # Get short commit hash of current HEAD COMMIT_HASH=$(git rev-parse --short HEAD) # Compose pretend version for setuptools_scm: e.g., 0.5.8.dev20260129+g1a2b3c4 PRETEND_VERSION="${VERSION}.dev${{ env.DATE }}+g${COMMIT_HASH}" echo "version=${VERSION}" >> $GITHUB_OUTPUT echo "pretend_version=${PRETEND_VERSION}" >> $GITHUB_OUTPUT echo "Detected version: ${VERSION}" echo "Pretend version for pip: ${PRETEND_VERSION}" - name: Login to Docker Hub (AMD) uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_AMD_USERNAME }} password: ${{ secrets.DOCKERHUB_AMD_TOKEN }} - name: Build and Push to rocm/sgl-dev run: | version=${{ steps.version.outputs.version }} pretend_version=${{ steps.version.outputs.pretend_version }} echo "Version: ${version}" echo "Pretend version: ${pretend_version}" if [ "${{ matrix.gpu_arch }}" = "gfx942" ]; then rocm_tag="rocm700-mi30x" elif [ "${{ matrix.gpu_arch }}" = "gfx950" ]; then rocm_tag="rocm700-mi35x" else echo "Unsupported gfx arch" exit 1 fi tag=v${version}-${rocm_tag} echo "IMAGE_TAG=${tag}-${{ env.DATE }}" >> $GITHUB_ENV docker build . -f docker/rocm.Dockerfile --build-arg SGL_BRANCH=${{ github.ref_name }} --build-arg BUILD_TYPE=${{ matrix.build_type }} --build-arg GPU_ARCH=${{ matrix.gpu_arch }} --build-arg ENABLE_MORI=1 --build-arg NIC_BACKEND=ainic --build-arg SETUPTOOLS_SCM_PRETEND_VERSION=${pretend_version} -t rocm/sgl-dev:${tag}-${{ env.DATE }} --no-cache docker push rocm/sgl-dev:${tag}-${{ env.DATE }} - name: Login to Docker Hub (lmsys) uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Push to lmsysorg/sglang-rocm run: | docker tag rocm/sgl-dev:${{ env.IMAGE_TAG }} lmsysorg/sglang-rocm:${{ env.IMAGE_TAG }} docker push lmsysorg/sglang-rocm:${{ env.IMAGE_TAG }} # Temporarily disable docker cache seeding until performant storage is in place cache: if: false # if: always() && github.repository == 'sgl-project/sglang' runs-on: linux-mi300-gpu-1 environment: 'prod' needs: publish strategy: fail-fast: false matrix: gpu_arch: ['gfx942'] build_type: ['all'] steps: - name: Checkout repository uses: actions/checkout@v4 with: fetch-depth: 0 # Required for git describe to find tags - name: "Set Date" run: | echo "DATE=$(date +%Y%m%d)" >> $GITHUB_ENV - name: Get version from latest tag id: version run: | # Get the latest version tag sorted by version number (e.g., v0.5.7 -> 0.5.7) VERSION=$(git tag -l 'v[0-9]*' --sort=-v:refname | head -1 | sed 's/^v//') if [ -z "$VERSION" ]; then echo "::error::Could not determine version from git tags" exit 1 fi echo "version=${VERSION}" >> $GITHUB_OUTPUT echo "Detected version: ${VERSION}" - name: Login to Docker Hub uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_AMD_USERNAME }} password: ${{ secrets.DOCKERHUB_AMD_TOKEN }} - name: Pull and Save Docker Image to Cache run: | set -euxo pipefail version=${{ steps.version.outputs.version }} echo "Version: ${version}" if [ "${{ matrix.gpu_arch }}" = "gfx942" ]; then rocm_tag="rocm700-mi30x" else echo "Unsupported gfx arch" exit 1 fi tag=v${version}-${rocm_tag} if [ "${{ matrix.build_type }}" = "all" ]; then tag_suffix="" else echo "Unsupported build type" exit 1 fi image="rocm/sgl-dev:${tag}-${{ env.DATE }}${tag_suffix}" # Determine target cache file name based on ROCm variant if [[ "${rocm_tag}" == rocm700* ]]; then final_path="/home/runner/sgl-data/docker/image-700.tar" else echo "Unexpected ROCm tag: ${rocm_tag}" exit 1 fi tmp_path="${final_path}.tmp" echo "Pulling image: ${image}" docker pull "${image}" echo "Saving to temp file: ${tmp_path}" docker save "${image}" -o "${tmp_path}" echo "Moving to final path: ${final_path}" mv -f "${tmp_path}" "${final_path}" echo "Cache populated successfully at ${final_path}" ================================================ FILE: .github/workflows/release-docker-amd-rocm720-nightly.yml ================================================ name: Release Docker Images ROCm 7.2.0 Nightly Preview (AMD) on: workflow_dispatch: schedule: - cron: '0 12 * * *' concurrency: # A PR number if a pull request and otherwise the commit hash. This cancels # queued and in-progress runs for the same PR (presubmit) or commit # (postsubmit). The workflow name is prepended to avoid conflicts between # different workflows. group: ${{ github.workflow }}-${{ github.event.number || github.sha }} cancel-in-progress: True jobs: publish: if: github.repository == 'sgl-project/sglang' runs-on: amd-docker-scale environment: 'prod' strategy: fail-fast: false matrix: gpu_arch: ['gfx942-rocm720', 'gfx950-rocm720'] build_type: ['all'] steps: - name: Checkout repository uses: actions/checkout@v4 with: fetch-depth: 0 # Required for git describe to find tags - name: "Set Date" run: | echo "DATE=$(date +%Y%m%d)" >> $GITHUB_ENV - name: Get version from latest tag id: version run: | # Get the latest version tag sorted by version number (e.g., v0.5.7 -> 0.5.7) VERSION=$(git tag -l 'v[0-9]*' --sort=-v:refname | head -1 | sed 's/^v//') if [ -z "$VERSION" ]; then echo "::error::Could not determine version from git tags" exit 1 fi # Get short commit hash of current HEAD COMMIT_HASH=$(git rev-parse --short HEAD) # Compose pretend version for setuptools_scm: e.g., 0.5.8.post1.dev20260211+g1a2b3c4 PRETEND_VERSION="${VERSION}.dev${{ env.DATE }}+g${COMMIT_HASH}" echo "version=${VERSION}" >> $GITHUB_OUTPUT echo "pretend_version=${PRETEND_VERSION}" >> $GITHUB_OUTPUT echo "Detected version: ${VERSION}" echo "Pretend version for pip: ${PRETEND_VERSION}" - name: Login to Docker Hub uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_AMD_USERNAME }} password: ${{ secrets.DOCKERHUB_AMD_TOKEN }} - name: Build and Push to rocm/sgl-dev run: | version=${{ steps.version.outputs.version }} pretend_version=${{ steps.version.outputs.pretend_version }} echo "Version: ${version}" echo "Pretend version: ${pretend_version}" if [ "${{ matrix.gpu_arch }}" = "gfx942-rocm720" ]; then rocm_tag="rocm720-mi30x" elif [ "${{ matrix.gpu_arch }}" = "gfx950-rocm720" ]; then rocm_tag="rocm720-mi35x" else echo "Unsupported gfx arch" exit 1 fi tag=v${version}-${rocm_tag} echo "IMAGE_TAG=${tag}-${{ env.DATE }}" >> $GITHUB_ENV docker build . -f docker/rocm.Dockerfile --build-arg SGL_BRANCH=${{ github.ref_name }} --build-arg BUILD_TYPE=${{ matrix.build_type }} --build-arg GPU_ARCH=${{ matrix.gpu_arch }} --build-arg ENABLE_MORI=1 --build-arg NIC_BACKEND=ainic --build-arg SETUPTOOLS_SCM_PRETEND_VERSION=${pretend_version} -t rocm/sgl-dev:${tag}-${{ env.DATE }} --no-cache docker push rocm/sgl-dev:${tag}-${{ env.DATE }} - name: Login to Docker Hub (lmsys) uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Push to lmsysorg/sglang-rocm run: | docker tag rocm/sgl-dev:${{ env.IMAGE_TAG }} lmsysorg/sglang-rocm:${{ env.IMAGE_TAG }} docker push lmsysorg/sglang-rocm:${{ env.IMAGE_TAG }} ================================================ FILE: .github/workflows/release-docker-amd.yml ================================================ name: Release Docker Images (AMD) on: push: tags: - 'v[0-9]+.*' workflow_dispatch: inputs: version: description: 'Version to build (without v prefix, e.g., 0.5.7)' required: true jobs: publish: if: github.repository == 'sgl-project/sglang' runs-on: amd-docker-scale environment: 'prod' strategy: matrix: rocm_version: ['rocm700', 'rocm720'] gpu_arch: ['gfx942', 'gfx950'] build_type: ['all'] steps: - name: Checkout repository uses: actions/checkout@v4 - name: Login to Docker Hub uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Get version from tag id: version run: | if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then VERSION="${{ github.event.inputs.version }}" else # Extract version from tag (e.g., v0.5.7 -> 0.5.7) VERSION="${GITHUB_REF_NAME#v}" fi # Validate version format if [ -z "$VERSION" ]; then echo "::error::Version is empty" exit 1 fi if ! echo "$VERSION" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]+'; then echo "::error::Invalid version format: $VERSION (expected: X.Y.Z)" exit 1 fi echo "version=${VERSION}" >> $GITHUB_OUTPUT - name: Build and Push run: | version=${{ steps.version.outputs.version }} echo "Version: ${version}" gpu_arch_suffix="" if [ "${{ matrix.rocm_version }}" = "rocm700" ]; then if [ "${{ matrix.gpu_arch }}" = "gfx942" ]; then rocm_tag="rocm700-mi30x" elif [ "${{ matrix.gpu_arch }}" = "gfx950" ]; then rocm_tag="rocm700-mi35x" else echo "Unsupported gfx arch" exit 1 fi elif [ "${{ matrix.rocm_version }}" = "rocm720" ]; then gpu_arch_suffix="-${{ matrix.rocm_version }}" if [ "${{ matrix.gpu_arch }}" = "gfx942" ]; then rocm_tag="rocm720-mi30x" elif [ "${{ matrix.gpu_arch }}" = "gfx950" ]; then rocm_tag="rocm720-mi35x" else echo "Unsupported gfx arch" exit 1 fi else echo "Unsupported rocm version" exit 1 fi tag=v${version}-${rocm_tag} # rocm.Dockerfile expects SGL_BRANCH with 'v' prefix for git tag checkout docker build . -f docker/rocm.Dockerfile --build-arg BUILD_TYPE=${{ matrix.build_type }} --build-arg GPU_ARCH=${{ matrix.gpu_arch }}${gpu_arch_suffix} --build-arg SGL_BRANCH=v${version} --build-arg ENABLE_MORI=1 --build-arg NIC_BACKEND=ainic -t lmsysorg/sglang:${tag} --no-cache docker push lmsysorg/sglang:${tag} ================================================ FILE: .github/workflows/release-docker-cu13-framework.yml ================================================ name: Release CUDA 13 Framework Docker Images (Temporary) # Temporary workflow to build only versioned cu13 framework images # Can be deleted after use on: workflow_dispatch: inputs: version: description: "Version to build (without v prefix, e.g., 0.5.8)" required: true jobs: publish-x86: if: github.repository == 'sgl-project/sglang' runs-on: x64-docker-build-node steps: - name: Delete huge unnecessary tools folder run: rm -rf /opt/hostedtoolcache - name: Checkout repository uses: actions/checkout@v4 - name: Free disk space uses: jlumbroso/free-disk-space@main with: tool-cache: false docker-images: false android: true dotnet: true haskell: true large-packages: true swap-storage: false - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - name: Login to Docker Hub uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Validate version id: version run: | VERSION="${{ github.event.inputs.version }}" if [ -z "$VERSION" ]; then echo "::error::Version is empty" exit 1 fi if ! echo "$VERSION" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]+'; then echo "::error::Invalid version format: $VERSION (expected: X.Y.Z)" exit 1 fi echo "version=${VERSION}" >> $GITHUB_OUTPUT - name: Build and Push AMD64 Framework (CUDA 13) run: | version=${{ steps.version.outputs.version }} docker buildx build \ --target framework \ --platform linux/amd64 \ --output type=image,name=lmsysorg/sglang,push-by-digest=true,name-canonical=true,push=true \ -f docker/Dockerfile \ --build-arg CUDA_VERSION=13.0.1 \ --build-arg BUILD_TYPE=all \ --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \ --build-arg GRACE_BLACKWELL=0 \ --build-arg SGL_VERSION=${version} \ --metadata-file /tmp/metadata.json \ --no-cache \ . DIGEST=$(python3 -c "import json; print(json.load(open('/tmp/metadata.json'))['containerimage.digest'])") echo "Pushed digest: ${DIGEST}" echo "${DIGEST}" > /tmp/digest-cu130-amd64-framework.txt - name: Upload digest uses: actions/upload-artifact@v4 with: name: digest-cu130-amd64 path: /tmp/digest-cu130-amd64-framework.txt retention-days: 1 publish-arm64: if: github.repository == 'sgl-project/sglang' runs-on: arm-docker-build-node steps: - name: Delete huge unnecessary tools folder run: rm -rf /opt/hostedtoolcache - name: Checkout repository uses: actions/checkout@v4 - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - name: Login to Docker Hub uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Validate version id: version run: | VERSION="${{ github.event.inputs.version }}" if [ -z "$VERSION" ]; then echo "::error::Version is empty" exit 1 fi if ! echo "$VERSION" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]+'; then echo "::error::Invalid version format: $VERSION (expected: X.Y.Z)" exit 1 fi echo "version=${VERSION}" >> $GITHUB_OUTPUT - name: Build and Push ARM64 Framework (CUDA 13) run: | version=${{ steps.version.outputs.version }} docker buildx build \ --target framework \ --platform linux/arm64 \ --output type=image,name=lmsysorg/sglang,push-by-digest=true,name-canonical=true,push=true \ -f docker/Dockerfile \ --build-arg CUDA_VERSION=13.0.1 \ --build-arg BUILD_TYPE=all \ --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \ --build-arg GRACE_BLACKWELL=1 \ --build-arg SGL_VERSION=${version} \ --metadata-file /tmp/metadata.json \ --no-cache \ . DIGEST=$(python3 -c "import json; print(json.load(open('/tmp/metadata.json'))['containerimage.digest'])") echo "Pushed digest: ${DIGEST}" echo "${DIGEST}" > /tmp/digest-cu130-arm64-framework.txt - name: Upload digest uses: actions/upload-artifact@v4 with: name: digest-cu130-arm64 path: /tmp/digest-cu130-arm64-framework.txt retention-days: 1 create-manifest: runs-on: ubuntu-22.04 needs: [publish-x86, publish-arm64] if: github.repository == 'sgl-project/sglang' steps: - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - name: Login to Docker Hub uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Download amd64 digest uses: actions/download-artifact@v4 with: name: digest-cu130-amd64 path: /tmp/digests/amd64 - name: Download arm64 digest uses: actions/download-artifact@v4 with: name: digest-cu130-arm64 path: /tmp/digests/arm64 - name: Create multi-arch manifest run: | version=${{ github.event.inputs.version }} AMD64_DIGEST=$(cat /tmp/digests/amd64/digest-cu130-amd64-framework.txt) ARM64_DIGEST=$(cat /tmp/digests/arm64/digest-cu130-arm64-framework.txt) # Create versioned CUDA 13 framework manifest docker buildx imagetools create \ -t lmsysorg/sglang:v${version}-cu130 \ lmsysorg/sglang@${AMD64_DIGEST} \ lmsysorg/sglang@${ARM64_DIGEST} # Create latest CUDA 13 framework manifest docker buildx imagetools create \ -t lmsysorg/sglang:latest-cu130 \ lmsysorg/sglang@${AMD64_DIGEST} \ lmsysorg/sglang@${ARM64_DIGEST} ================================================ FILE: .github/workflows/release-docker-dev.yml ================================================ name: Build and Push Development Docker Images on: workflow_dispatch: inputs: pr_number: description: "PR number to build from (leave empty to use current branch)" required: false default: "" tag: description: "Custom tag suffix (overrides pr_number in tag). E.g. 'my-test' → dev-my-test, dev-cu13-my-test, etc." required: false default: "" schedule: - cron: "0 0 * * *" concurrency: group: release-docker-dev-${{ inputs.tag || inputs.pr_number || 'nightly' }} cancel-in-progress: true jobs: build-dev: if: ${{ github.repository == 'sgl-project/sglang' }} runs-on: ${{ matrix.runner }} strategy: matrix: include: - runner: x64-docker-build-node platform: linux/amd64 build_type: all grace_blackwell: 0 arch_tag: x86 version: 12.9.1 - runner: arm-docker-build-node platform: linux/arm64 build_type: all grace_blackwell: 1 arch_tag: arm64 version: 12.9.1 - runner: x64-docker-build-node platform: linux/amd64 build_type: all grace_blackwell: 0 arch_tag: x86-cu13 version: 13.0.1 - runner: arm-docker-build-node platform: linux/arm64 build_type: all grace_blackwell: 1 arch_tag: arm64-cu13 version: 13.0.1 steps: - name: Delete huge unnecessary tools folder run: rm -rf /opt/hostedtoolcache - name: Checkout repository uses: actions/checkout@v4 with: ref: ${{ inputs.pr_number && format('refs/pull/{0}/head', inputs.pr_number) || github.ref }} - name: Free disk space uses: jlumbroso/free-disk-space@main with: tool-cache: true docker-images: true android: true dotnet: true haskell: true large-packages: true swap-storage: true - name: Prune Docker to reclaim disk space run: | docker buildx prune --filter "until=72h" -f docker system prune -af --filter "until=72h" docker volume prune -af - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - name: Login to Docker Hub uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Build and Push Dev Image run: | # Nightly (schedule) installs latest release; manual dispatch builds from checked-out source if [ "${{ github.event_name }}" = "schedule" ]; then SOURCE_ARG="--build-arg USE_LATEST_SGLANG=1" else SOURCE_ARG="--build-arg BRANCH_TYPE=local" fi docker buildx build \ --platform ${{ matrix.platform }} \ --output type=image,name=lmsysorg/sglang,push-by-digest=true,name-canonical=true,push=true \ --target framework \ -f docker/Dockerfile \ --build-arg CUDA_VERSION=${{ matrix.version }} \ --build-arg BUILD_TYPE=${{ matrix.build_type }} \ --build-arg CMAKE_BUILD_PARALLEL_LEVEL=$(nproc) \ --build-arg GRACE_BLACKWELL=${{ matrix.grace_blackwell }} \ ${SOURCE_ARG} \ --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \ --metadata-file /tmp/metadata.json \ --no-cache \ . DIGEST=$(python3 -c "import json; print(json.load(open('/tmp/metadata.json'))['containerimage.digest'])") echo "Pushed digest: ${DIGEST}" echo "${DIGEST}" > /tmp/digest.txt - name: Upload digest uses: actions/upload-artifact@v4 with: name: digest-${{ matrix.arch_tag }} path: /tmp/digest.txt retention-days: 1 create-manifests: runs-on: ubuntu-22.04 needs: [build-dev] if: ${{ github.repository == 'sgl-project/sglang' }} strategy: matrix: variant: - base: dev x86: x86 arm64: arm64 - base: dev-cu13 x86: x86-cu13 arm64: arm64-cu13 steps: - uses: docker/setup-buildx-action@v3 - uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Download x86 digest uses: actions/download-artifact@v4 with: name: digest-${{ matrix.variant.x86 }} path: /tmp/digests/x86 - name: Download arm64 digest uses: actions/download-artifact@v4 with: name: digest-${{ matrix.variant.arm64 }} path: /tmp/digests/arm64 - name: Create multi-arch manifest run: | X86_DIGEST=$(cat /tmp/digests/x86/digest.txt) ARM64_DIGEST=$(cat /tmp/digests/arm64/digest.txt) SUFFIX="" if [ -n "${{ inputs.tag }}" ]; then SUFFIX="-${{ inputs.tag }}" elif [ -n "${{ inputs.pr_number }}" ]; then SUFFIX="-pr-${{ inputs.pr_number }}" fi TAG="${{ matrix.variant.base }}${SUFFIX}" # For nightly (no suffix), also stamp a dated tag EXTRA_TAG="" if [ -z "${SUFFIX}" ]; then SHORT_SHA="${{ github.sha }}" EXTRA_TAG="-t lmsysorg/sglang:nightly-${TAG}-$(date +%Y%m%d)-${SHORT_SHA:0:8}" fi docker buildx imagetools create \ -t lmsysorg/sglang:${TAG} \ ${EXTRA_TAG} \ lmsysorg/sglang@${X86_DIGEST} \ lmsysorg/sglang@${ARM64_DIGEST} echo "✓ Published lmsysorg/sglang:${TAG}" - name: Cleanup Old Nightly Builds if: ${{ !inputs.tag && !inputs.pr_number }} run: | TOKEN=$(curl -s -H "Content-Type: application/json" \ -X POST -d '{"username": "${{ secrets.DOCKERHUB_USERNAME }}", "password": "${{ secrets.DOCKERHUB_TOKEN }}"}' \ https://hub.docker.com/v2/users/login/ | jq -r .token) TAGS_RESPONSE=$(curl -s -H "Authorization: JWT $TOKEN" \ "https://hub.docker.com/v2/repositories/lmsysorg/sglang/tags/?page_size=100") TAGS=$(echo "$TAGS_RESPONSE" | jq -r \ '.results[] | select(.name | test("^nightly-${{ matrix.variant.base }}-[0-9]")) | "\(.last_updated)|\(.name)"' \ | sort -r | cut -d'|' -f2) TAG_COUNT=$(echo "$TAGS" | wc -l) if [ "$TAG_COUNT" -gt 14 ]; then echo "Found $TAG_COUNT nightly builds, keeping only the 14 most recent" TAGS_TO_DELETE=$(echo "$TAGS" | tail -n +15) for tag in $TAGS_TO_DELETE; do echo "Deleting tag: $tag" curl -X DELETE -H "Authorization: JWT $TOKEN" \ "https://hub.docker.com/v2/repositories/lmsysorg/sglang/tags/$tag/" done else echo "Only $TAG_COUNT nightly builds found, no cleanup needed" fi ================================================ FILE: .github/workflows/release-docker-gateway.yml ================================================ name: Release SGLang Model Gateway Docker Image on: push: branches: - main paths: - sgl-model-gateway/bindings/python/pyproject.toml workflow_dispatch: jobs: publish: if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-24.04 steps: - name: Checkout repository uses: actions/checkout@v4 - name: Set up QEMU uses: docker/setup-qemu-action@v3 - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - name: Login to Docker Hub uses: docker/login-action@v3 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Build and Push run: | version=$(cat sgl-model-gateway/bindings/python/src/sglang_router/version.py | cut -d'"' -f2) tag=v${version} docker buildx build . -f docker/gateway.Dockerfile \ --platform linux/amd64,linux/arm64 \ -t lmsysorg/sgl-model-gateway:${tag} \ -t lmsysorg/sgl-model-gateway:latest \ --push ================================================ FILE: .github/workflows/release-docker-npu-nightly.yml ================================================ name: Release Docker Images Nightly (NPU) on: pull_request: branches: - 'main' paths: - '.github/workflows/release-docker-npu-nightly.yml' - 'docker/npu.Dockerfile' workflow_dispatch: schedule: - cron: "0 0 * * *" concurrency: group: ${{ github.workflow }}-${{ github.sha }} cancel-in-progress: true jobs: build: runs-on: ubuntu-22.04-arm strategy: matrix: cann_version: ["8.5.0"] device_type: ["910b", "a3"] steps: - name: Checkout repository uses: actions/checkout@v4 - name: Free up disk space uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be # v1.3.1 with: tool-cache: true docker-images: false - name: Setup Docker buildx uses: docker/setup-buildx-action@v3 - name: Docker meta id: meta uses: docker/metadata-action@v5 with: images: | lmsysorg/sglang # push with schedule event # push with workflow_dispatch event tags: | type=ref,event=pr type=ref,event=branch type=schedule,pattern=main flavor: | latest=false suffix=-cann${{ matrix.cann_version }}-${{ matrix.device_type }},onlatest=true # Login against a Docker registry except on PR # https://github.com/docker/login-action - name: Log into docker hub uses: docker/login-action@v3 if: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }} with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} # Enable Docker multi-architecture build environment # Emulate non-native architectures - name: Set up QEMU uses: docker/setup-qemu-action@v3 # Required for building and pushing multi-arch Docker images - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 # Build and push Docker image with Buildx (don't push on PR) # https://github.com/docker/build-push-action - name: Build and push Docker image id: build-and-push uses: docker/build-push-action@v6 with: context: docker file: docker/npu.Dockerfile platforms: linux/arm64,linux/amd64 labels: ${{ steps.meta.outputs.labels }} tags: ${{ steps.meta.outputs.tags }} push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }} provenance: false build-args: | SGLANG_KERNEL_NPU_TAG=2026.03.10.rc1 CANN_VERSION=${{ matrix.cann_version }} DEVICE_TYPE=${{ matrix.device_type }} ================================================ FILE: .github/workflows/release-docker-npu.yml ================================================ name: Release Docker Images (NPU) on: push: tags: - 'v[0-9]+.*' workflow_dispatch: inputs: version: description: 'Version to build (without v prefix, e.g., 0.5.7)' required: true jobs: build: runs-on: ubuntu-22.04-arm strategy: matrix: cann_version: ["8.5.0"] device_type: ["910b", "a3"] steps: - name: Checkout repository uses: actions/checkout@v4 - name: Free up disk space uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be # v1.3.1 with: tool-cache: true docker-images: false # push with tag - name: Docker meta id: meta uses: docker/metadata-action@v5 with: images: | lmsysorg/sglang tags: | type=ref,event=pr flavor: | latest=false # Login against a Docker registry except on PR # https://github.com/docker/login-action - name: Login to Docker Hub uses: docker/login-action@v2 if: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }} with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Get version from tag id: version run: | if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then VERSION="${{ github.event.inputs.version }}" else # Extract version from tag (e.g., v0.5.7 -> 0.5.7) VERSION="${GITHUB_REF_NAME#v}" fi # Validate version format if [ -z "$VERSION" ]; then echo "::error::Version is empty" exit 1 fi if ! echo "$VERSION" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]+'; then echo "::error::Invalid version format: $VERSION (expected: X.Y.Z)" exit 1 fi echo "version=v${VERSION}" >> $GITHUB_OUTPUT echo "TAG=lmsysorg/sglang:v${VERSION}-cann${{ matrix.cann_version }}-${{ matrix.device_type }}" >> $GITHUB_OUTPUT # Enable Docker multi-architecture build environment # Emulate non-native architectures - name: Set up QEMU uses: docker/setup-qemu-action@v3 # Required for building and pushing multi-arch Docker images - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - name: Build and push Docker image id: build-and-push uses: docker/build-push-action@v6 with: context: docker file: docker/npu.Dockerfile platforms: linux/arm64,linux/amd64 labels: ${{ steps.meta.outputs.labels }} tags: ${{ steps.meta.outputs.tags || steps.version.outputs.TAG }} push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }} provenance: false build-args: | SGLANG_KERNEL_NPU_TAG=2026.03.10.rc1 CANN_VERSION=${{ matrix.cann_version }} DEVICE_TYPE=${{ matrix.device_type }} SGLANG_TAG=${{ steps.version.outputs.version }} ================================================ FILE: .github/workflows/release-docker-xeon.yml ================================================ name: Release Docker Xeon Images on: push: tags: - 'v[0-9]+.*' workflow_dispatch: inputs: version: description: 'Version to build (without v prefix, e.g., 0.5.7)' required: true jobs: publish: if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-24.04 environment: 'prod' strategy: matrix: build_type: ['all'] steps: - name: Checkout repository uses: actions/checkout@v4 - name: Login to Docker Hub uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Get version from tag id: version run: | if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then VERSION="${{ github.event.inputs.version }}" else # Extract version from tag (e.g., v0.5.7 -> 0.5.7) VERSION="${GITHUB_REF_NAME#v}" fi # Validate version format if [ -z "$VERSION" ]; then echo "::error::Version is empty" exit 1 fi if ! echo "$VERSION" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]+'; then echo "::error::Invalid version format: $VERSION (expected: X.Y.Z)" exit 1 fi echo "version=${VERSION}" >> $GITHUB_OUTPUT - name: Build and Push run: | version=${{ steps.version.outputs.version }} tag=v${version}-xeon docker build . -f docker/xeon.Dockerfile \ --build-arg VER_SGLANG=v${version} \ -t lmsysorg/sglang:${tag} \ --no-cache docker push lmsysorg/sglang:${tag} ================================================ FILE: .github/workflows/release-docker.yml ================================================ name: Release Docker Images # # This workflow builds and publishes both framework and runtime Docker images: # # Framework images (full development environment): # - lmsysorg/sglang:v{version}, lmsysorg/sglang:latest # - lmsysorg/sglang:v{version}-cu130, lmsysorg/sglang:latest-cu130 # # Runtime images (production-optimized, ~50% smaller): # - lmsysorg/sglang:v{version}-runtime, lmsysorg/sglang:latest-runtime # - lmsysorg/sglang:v{version}-cu130-runtime, lmsysorg/sglang:latest-cu130-runtime # on: push: tags: - "v[0-9]+.*" workflow_dispatch: inputs: version: description: "Version to build (without v prefix, e.g., 0.5.7)" required: true jobs: publish-x86: if: github.repository == 'sgl-project/sglang' environment: "prod" strategy: matrix: variant: - cuda_version: "12.9.1" build_type: "all" grace_blackwell: 0 runs-on: x64-docker-build-node steps: - name: Delete huge unnecessary tools folder run: rm -rf /opt/hostedtoolcache - name: Checkout repository uses: actions/checkout@v4 - name: Free disk space uses: jlumbroso/free-disk-space@main with: tool-cache: false docker-images: false android: true dotnet: true haskell: true large-packages: true swap-storage: false - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - name: Login to Docker Hub uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Get version from tag id: version run: | if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then VERSION="${{ github.event.inputs.version }}" else # Extract version from tag (e.g., v0.5.7 -> 0.5.7) VERSION="${GITHUB_REF_NAME#v}" fi # Validate version format if [ -z "$VERSION" ]; then echo "::error::Version is empty" exit 1 fi if ! echo "$VERSION" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]+'; then echo "::error::Invalid version format: $VERSION (expected: X.Y.Z)" exit 1 fi echo "version=${VERSION}" >> $GITHUB_OUTPUT - name: Build AMD64 Framework run: | version=${{ steps.version.outputs.version }} docker buildx build \ --target framework \ --platform linux/amd64 \ --output type=image,name=lmsysorg/sglang,push-by-digest=true,name-canonical=true,push=true \ -f docker/Dockerfile \ --build-arg CUDA_VERSION=${{ matrix.variant.cuda_version }} \ --build-arg BUILD_TYPE=${{ matrix.variant.build_type }} \ --build-arg GRACE_BLACKWELL=${{ matrix.variant.grace_blackwell }} \ --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \ --build-arg SGL_VERSION=${version} \ --metadata-file /tmp/metadata-cu129-framework.json \ --no-cache \ . DIGEST=$(python3 -c "import json; print(json.load(open('/tmp/metadata-cu129-framework.json'))['containerimage.digest'])") echo "Pushed digest: ${DIGEST}" echo "${DIGEST}" > /tmp/digest-cu129-amd64-framework.txt - name: Build and Push AMD64 Runtime run: | version=${{ steps.version.outputs.version }} docker buildx build \ --target runtime \ --platform linux/amd64 \ --output type=image,name=lmsysorg/sglang,push-by-digest=true,name-canonical=true,push=true \ -f docker/Dockerfile \ --build-arg CUDA_VERSION=${{ matrix.variant.cuda_version }} \ --build-arg BUILD_TYPE=${{ matrix.variant.build_type }} \ --build-arg GRACE_BLACKWELL=${{ matrix.variant.grace_blackwell }} \ --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \ --build-arg SGL_VERSION=${version} \ --metadata-file /tmp/metadata-cu129-runtime.json \ --no-cache \ . DIGEST=$(python3 -c "import json; print(json.load(open('/tmp/metadata-cu129-runtime.json'))['containerimage.digest'])") echo "Pushed digest: ${DIGEST}" echo "${DIGEST}" > /tmp/digest-cu129-amd64-runtime.txt - name: Build and Push AMD64 Framework (CUDA 13) run: | version=${{ steps.version.outputs.version }} docker buildx build \ --target framework \ --platform linux/amd64 \ --output type=image,name=lmsysorg/sglang,push-by-digest=true,name-canonical=true,push=true \ -f docker/Dockerfile \ --build-arg CUDA_VERSION=13.0.1 \ --build-arg BUILD_TYPE=${{ matrix.variant.build_type }} \ --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \ --build-arg GRACE_BLACKWELL=0 \ --build-arg SGL_VERSION=${version} \ --metadata-file /tmp/metadata-cu130-framework.json \ --no-cache \ . DIGEST=$(python3 -c "import json; print(json.load(open('/tmp/metadata-cu130-framework.json'))['containerimage.digest'])") echo "Pushed digest: ${DIGEST}" echo "${DIGEST}" > /tmp/digest-cu130-amd64-framework.txt - name: Build and Push AMD64 Runtime (CUDA 13) run: | version=${{ steps.version.outputs.version }} docker buildx build \ --target runtime \ --platform linux/amd64 \ --output type=image,name=lmsysorg/sglang,push-by-digest=true,name-canonical=true,push=true \ -f docker/Dockerfile \ --build-arg CUDA_VERSION=13.0.1 \ --build-arg BUILD_TYPE=${{ matrix.variant.build_type }} \ --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \ --build-arg GRACE_BLACKWELL=0 \ --build-arg SGL_VERSION=${version} \ --metadata-file /tmp/metadata-cu130-runtime.json \ --no-cache \ . DIGEST=$(python3 -c "import json; print(json.load(open('/tmp/metadata-cu130-runtime.json'))['containerimage.digest'])") echo "Pushed digest: ${DIGEST}" echo "${DIGEST}" > /tmp/digest-cu130-amd64-runtime.txt - name: Upload digests uses: actions/upload-artifact@v4 with: name: digests-amd64 path: /tmp/digest-*.txt retention-days: 1 publish-arm64: if: github.repository == 'sgl-project/sglang' environment: "prod" strategy: matrix: variant: - cuda_version: "12.9.1" build_type: "all" grace_blackwell: 1 runs-on: arm-docker-build-node steps: - name: Delete huge unnecessary tools folder run: rm -rf /opt/hostedtoolcache - name: Checkout repository uses: actions/checkout@v4 - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - name: Login to Docker Hub uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Get version from tag id: version run: | if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then VERSION="${{ github.event.inputs.version }}" else # Extract version from tag (e.g., v0.5.7 -> 0.5.7) VERSION="${GITHUB_REF_NAME#v}" fi # Validate version format if [ -z "$VERSION" ]; then echo "::error::Version is empty" exit 1 fi if ! echo "$VERSION" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]+'; then echo "::error::Invalid version format: $VERSION (expected: X.Y.Z)" exit 1 fi echo "version=${VERSION}" >> $GITHUB_OUTPUT - name: Build ARM64 Framework run: | version=${{ steps.version.outputs.version }} docker buildx build \ --target framework \ --platform linux/arm64 \ --output type=image,name=lmsysorg/sglang,push-by-digest=true,name-canonical=true,push=true \ -f docker/Dockerfile \ --build-arg CUDA_VERSION=${{ matrix.variant.cuda_version }} \ --build-arg BUILD_TYPE=${{ matrix.variant.build_type }} \ --build-arg GRACE_BLACKWELL=${{ matrix.variant.grace_blackwell }} \ --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \ --build-arg SGL_VERSION=${version} \ --metadata-file /tmp/metadata-cu129-framework.json \ --no-cache \ . DIGEST=$(python3 -c "import json; print(json.load(open('/tmp/metadata-cu129-framework.json'))['containerimage.digest'])") echo "Pushed digest: ${DIGEST}" echo "${DIGEST}" > /tmp/digest-cu129-arm64-framework.txt - name: Build and Push ARM64 Runtime run: | version=${{ steps.version.outputs.version }} docker buildx build \ --target runtime \ --platform linux/arm64 \ --output type=image,name=lmsysorg/sglang,push-by-digest=true,name-canonical=true,push=true \ -f docker/Dockerfile \ --build-arg CUDA_VERSION=${{ matrix.variant.cuda_version }} \ --build-arg BUILD_TYPE=${{ matrix.variant.build_type }} \ --build-arg GRACE_BLACKWELL=${{ matrix.variant.grace_blackwell }} \ --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \ --build-arg SGL_VERSION=${version} \ --metadata-file /tmp/metadata-cu129-runtime.json \ --no-cache \ . DIGEST=$(python3 -c "import json; print(json.load(open('/tmp/metadata-cu129-runtime.json'))['containerimage.digest'])") echo "Pushed digest: ${DIGEST}" echo "${DIGEST}" > /tmp/digest-cu129-arm64-runtime.txt - name: Build and Push ARM64 Framework (CUDA 13) run: | version=${{ steps.version.outputs.version }} docker buildx build \ --target framework \ --platform linux/arm64 \ --output type=image,name=lmsysorg/sglang,push-by-digest=true,name-canonical=true,push=true \ -f docker/Dockerfile \ --build-arg CUDA_VERSION=13.0.1 \ --build-arg BUILD_TYPE=${{ matrix.variant.build_type }} \ --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \ --build-arg GRACE_BLACKWELL=1 \ --build-arg SGL_VERSION=${version} \ --metadata-file /tmp/metadata-cu130-framework.json \ --no-cache \ . DIGEST=$(python3 -c "import json; print(json.load(open('/tmp/metadata-cu130-framework.json'))['containerimage.digest'])") echo "Pushed digest: ${DIGEST}" echo "${DIGEST}" > /tmp/digest-cu130-arm64-framework.txt - name: Build and Push ARM64 Runtime (CUDA 13) run: | version=${{ steps.version.outputs.version }} docker buildx build \ --target runtime \ --platform linux/arm64 \ --output type=image,name=lmsysorg/sglang,push-by-digest=true,name-canonical=true,push=true \ -f docker/Dockerfile \ --build-arg CUDA_VERSION=13.0.1 \ --build-arg BUILD_TYPE=${{ matrix.variant.build_type }} \ --build-arg GRACE_BLACKWELL=1 \ --build-arg SGL_VERSION=${version} \ --metadata-file /tmp/metadata-cu130-runtime.json \ --no-cache \ . DIGEST=$(python3 -c "import json; print(json.load(open('/tmp/metadata-cu130-runtime.json'))['containerimage.digest'])") echo "Pushed digest: ${DIGEST}" echo "${DIGEST}" > /tmp/digest-cu130-arm64-runtime.txt - name: Upload digests uses: actions/upload-artifact@v4 with: name: digests-arm64 path: /tmp/digest-*.txt retention-days: 1 create-manifests: runs-on: ubuntu-22.04 needs: [publish-x86, publish-arm64] if: github.repository == 'sgl-project/sglang' environment: "prod" steps: - name: Checkout repository uses: actions/checkout@v4 - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - name: Login to Docker Hub uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Get version from tag id: version run: | if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then VERSION="${{ github.event.inputs.version }}" else # Extract version from tag (e.g., v0.5.7 -> 0.5.7) VERSION="${GITHUB_REF_NAME#v}" fi # Validate version format if [ -z "$VERSION" ]; then echo "::error::Version is empty" exit 1 fi if ! echo "$VERSION" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]+'; then echo "::error::Invalid version format: $VERSION (expected: X.Y.Z)" exit 1 fi echo "version=${VERSION}" >> $GITHUB_OUTPUT - name: Download amd64 digests uses: actions/download-artifact@v4 with: name: digests-amd64 path: /tmp/digests/amd64 - name: Download arm64 digests uses: actions/download-artifact@v4 with: name: digests-arm64 path: /tmp/digests/arm64 - name: Create multi-arch manifests run: | version=${{ steps.version.outputs.version }} # Load all digests CU129_AMD64_FW=$(cat /tmp/digests/amd64/digest-cu129-amd64-framework.txt) CU129_AMD64_RT=$(cat /tmp/digests/amd64/digest-cu129-amd64-runtime.txt) CU130_AMD64_FW=$(cat /tmp/digests/amd64/digest-cu130-amd64-framework.txt) CU130_AMD64_RT=$(cat /tmp/digests/amd64/digest-cu130-amd64-runtime.txt) CU129_ARM64_FW=$(cat /tmp/digests/arm64/digest-cu129-arm64-framework.txt) CU129_ARM64_RT=$(cat /tmp/digests/arm64/digest-cu129-arm64-runtime.txt) CU130_ARM64_FW=$(cat /tmp/digests/arm64/digest-cu130-arm64-framework.txt) CU130_ARM64_RT=$(cat /tmp/digests/arm64/digest-cu130-arm64-runtime.txt) # Create versioned framework manifest (default) docker buildx imagetools create \ -t lmsysorg/sglang:v${version} \ lmsysorg/sglang@${CU129_AMD64_FW} \ lmsysorg/sglang@${CU129_ARM64_FW} # Create latest framework manifest (default) docker buildx imagetools create \ -t lmsysorg/sglang:latest \ lmsysorg/sglang@${CU129_AMD64_FW} \ lmsysorg/sglang@${CU129_ARM64_FW} # Create versioned runtime manifest docker buildx imagetools create \ -t lmsysorg/sglang:v${version}-runtime \ lmsysorg/sglang@${CU129_AMD64_RT} \ lmsysorg/sglang@${CU129_ARM64_RT} # Create latest runtime manifest docker buildx imagetools create \ -t lmsysorg/sglang:latest-runtime \ lmsysorg/sglang@${CU129_AMD64_RT} \ lmsysorg/sglang@${CU129_ARM64_RT} # Create versioned CUDA 13 framework manifest docker buildx imagetools create \ -t lmsysorg/sglang:v${version}-cu130 \ lmsysorg/sglang@${CU130_AMD64_FW} \ lmsysorg/sglang@${CU130_ARM64_FW} # Create latest CUDA 13 framework manifest docker buildx imagetools create \ -t lmsysorg/sglang:latest-cu130 \ lmsysorg/sglang@${CU130_AMD64_FW} \ lmsysorg/sglang@${CU130_ARM64_FW} # Create versioned CUDA 13 runtime manifest docker buildx imagetools create \ -t lmsysorg/sglang:v${version}-cu130-runtime \ lmsysorg/sglang@${CU130_AMD64_RT} \ lmsysorg/sglang@${CU130_ARM64_RT} # Create latest CUDA 13 runtime manifest docker buildx imagetools create \ -t lmsysorg/sglang:latest-cu130-runtime \ lmsysorg/sglang@${CU130_AMD64_RT} \ lmsysorg/sglang@${CU130_ARM64_RT} ================================================ FILE: .github/workflows/release-docs.yml ================================================ name: Release Documentation on: release: types: [published] push: branches: - main paths: - "docs/**" - "python/sglang/version.py" - "python/sglang/**" workflow_dispatch: concurrency: group: release-docs-${{ github.ref }} cancel-in-progress: true env: SGLANG_IS_IN_CI: true jobs: execute-and-deploy: runs-on: 1-gpu-runner if: github.repository == 'sgl-project/sglang' steps: - name: Checkout code uses: actions/checkout@v4 - name: Fetch full git history for release index if: github.event_name == 'release' run: | git fetch --prune --unshallow || git fetch --prune --depth=0 - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh pip install -r docs/requirements.txt apt-get update && apt-get install -y pandoc parallel retry ln -sf "$(which python3)" /usr/bin/python - name: Setup Jupyter Kernel run: | python -m ipykernel install --user --name python3 --display-name "Python 3" - name: Execute notebooks timeout-minutes: 40 run: | cd docs make clean make compile - name: Push HTML to sgl-project.github.io timeout-minutes: 30 env: GITHUB_TOKEN: ${{ secrets.GH_PAT_FOR_DOCUMENTATION }} run: | cd docs make html make markdown python3 wrap_run_llm.py if [[ "${{ github.event_name }}" == "release" ]]; then python3 release_lookup/generate_index.py --output release_lookup/release_index.json # Copy release lookup tool for official docs on published releases. mkdir -p _build/html/release_lookup cp release_lookup/index.html _build/html/release_lookup/ cp release_lookup/release_index.json _build/html/release_lookup/ fi cd _build/html git clone https://$GITHUB_TOKEN@github.com/sgl-project/sgl-project.github.io.git ../sgl-project.github.io --depth 1 if [[ "${{ github.event_name }}" == "release" ]]; then find ../sgl-project.github.io/ -mindepth 1 -not -path "../sgl-project.github.io/.git*" -not -name CNAME -not -name ".jekyll" -not -name ".nojekyll" -delete else find ../sgl-project.github.io/ -mindepth 1 -not -path "../sgl-project.github.io/.git*" -not -path "../sgl-project.github.io/release_lookup*" -not -name CNAME -not -name ".jekyll" -not -name ".nojekyll" -delete fi cp -r * ../sgl-project.github.io cp ../../README.md ../sgl-project.github.io/README.md cd ../sgl-project.github.io git config user.name "sglang-bot" git config user.email "sglangbot@gmail.com" git add . git commit -m "Update $(date +'%Y-%m-%d %H:%M:%S')" git push https://$GITHUB_TOKEN@github.com/sgl-project/sgl-project.github.io.git main cd .. rm -rf sgl-project.github.io ================================================ FILE: .github/workflows/release-pypi-gateway.yml ================================================ name: Release SGLang Model Gateway to PyPI on: push: branches: - main paths: - sgl-model-gateway/bindings/python/pyproject.toml workflow_dispatch: jobs: build: name: build on ${{ matrix.platform || matrix.os }} (${{ matrix.target }} - ${{ matrix.manylinux || 'auto' }}) runs-on: ${{ matrix.os }}-latest strategy: fail-fast: false matrix: os: [ubuntu, macos, windows] target: [x86_64, aarch64] manylinux: [auto] include: - os: ubuntu platform: linux - os: windows ls: dir target: x86_64 python-architecture: x64 interpreter: 3.9 3.10 3.11 3.12 3.13 - os: macos target: aarch64 interpreter: 3.9 3.10 3.11 3.12 3.13 - os: ubuntu platform: linux target: aarch64 # musllinux - os: ubuntu platform: linux target: x86_64 manylinux: musllinux_1_1 - os: ubuntu platform: linux target: aarch64 manylinux: musllinux_1_1 exclude: - os: windows target: aarch64 steps: - uses: actions/checkout@v4 with: path: sglang-repo - name: Move sgl-model-gateway folder to root and delete sglang-repo run: | mv sglang-repo/sgl-model-gateway/* . rm -rf sglang-repo ls -alt shell: bash - name: Set up Python uses: actions/setup-python@v5 with: python-version: "3.13" architecture: ${{ matrix.python-architecture || 'x64' }} - name: Install twine run: pip install -U twine - name: Install protoc (macOS) if: matrix.os == 'macos' run: brew install protobuf - name: Install protoc (Windows) if: matrix.os == 'windows' run: choco install protoc -y - name: Build wheels uses: PyO3/maturin-action@v1 with: working-directory: bindings/python target: ${{ matrix.target }} manylinux: ${{ matrix.manylinux || 'auto' }} args: --release --out dist --features vendored-openssl --interpreter ${{ matrix.interpreter || '3.9 3.10 3.11 3.12 3.13 3.14' }} rust-toolchain: stable docker-options: -e CI -e CC_aarch64_unknown_linux_gnu=aarch64-linux-gnu-gcc -e CXX_aarch64_unknown_linux_gnu=aarch64-linux-gnu-g++ before-script-linux: | # Install build dependencies (perl/make for vendored OpenSSL, protoc for gRPC) if command -v yum &> /dev/null; then yum update -y && yum install -y wget unzip gcc gcc-c++ perl-core make # Install cross-compilation toolchain for aarch64 if needed if [ "${{ matrix.target }}" = "aarch64" ]; then yum install -y gcc-aarch64-linux-gnu gcc-c++-aarch64-linux-gnu || true fi elif command -v apt-get &> /dev/null; then apt-get update && apt-get install -y wget unzip gcc g++ perl make # Install cross-compilation toolchain for aarch64 if needed if [ "${{ matrix.target }}" = "aarch64" ]; then apt-get install -y gcc-aarch64-linux-gnu g++-aarch64-linux-gnu || true fi fi (cd /tmp && \ wget https://github.com/protocolbuffers/protobuf/releases/download/v32.0/protoc-32.0-linux-x86_64.zip && \ unzip protoc-32.0-linux-x86_64.zip -d /usr/local && \ rm protoc-32.0-linux-x86_64.zip) protoc --version - name: List built packages run: ${{ matrix.ls || 'ls -lh' }} bindings/python/dist/ - name: Check packages run: twine check --strict bindings/python/dist/* - uses: actions/upload-artifact@v4 with: name: packages-${{ matrix.os }}-${{ matrix.target }}-${{ matrix.manylinux || 'auto' }} path: bindings/python/dist/ build-sdist: name: Build SDist runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 with: path: sglang-repo - name: Move sgl-model-gateway folder to root and delete sglang-repo run: | mv sglang-repo/sgl-model-gateway/* . rm -rf sglang-repo ls -alt - name: Set up Python uses: actions/setup-python@v5 with: python-version: "3.13" - name: Build SDist uses: PyO3/maturin-action@v1 with: working-directory: bindings/python command: sdist args: --out dist rust-toolchain: stable - uses: actions/upload-artifact@v4 with: name: sdist path: bindings/python/dist/*.tar.gz upload: name: Upload to PyPI if: github.repository == 'sgl-project/sglang' # Ensure this job only runs for the sgl-project/sglang repository needs: [build, build-sdist] runs-on: ubuntu-latest steps: - uses: actions/download-artifact@v4 with: path: dist merge-multiple: true - name: Upload to PyPI env: TWINE_USERNAME: __token__ TWINE_PASSWORD: ${{ secrets.PYPI_TOKEN_ROUTER }} run: | pip install twine twine upload dist/* --verbose ================================================ FILE: .github/workflows/release-pypi-nightly.yml ================================================ name: Release PyPI Nightly Wheels on: # Run daily at 2 AM UTC schedule: - cron: '0 2 * * *' # Triggered by nightly Docker workflow to use same commit repository_dispatch: types: [nightly-release] # Manual trigger for testing workflow_dispatch: inputs: commit_sha: description: 'Specific commit SHA to build (leave empty for latest)' required: false type: string cuda_version: description: 'CUDA version (e.g., 129 or 130)' required: false default: '129' type: string concurrency: group: release-pypi-nightly-${{ github.ref }} cancel-in-progress: true jobs: build-nightly-wheel: if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-latest outputs: nightly_version: ${{ steps.build.outputs.nightly_version }} commit_hash: ${{ steps.build.outputs.commit_hash }} build_date: ${{ steps.build.outputs.build_date }} steps: - uses: actions/checkout@v4 with: # Use commit from: 1) Docker workflow, 2) manual input, 3) latest main ref: ${{ github.event.client_payload.commit_sha || inputs.commit_sha || github.sha }} fetch-depth: 0 # Need full history for setuptools-scm - name: Set up Python uses: actions/setup-python@v5 with: python-version: "3.10" - name: Install build dependencies run: | pip install build wheel setuptools setuptools-scm - name: Build wheel id: build run: | cd python cp ../README.md ../LICENSE . # Parse git describe output to get latest tag # Use same command as pyproject.toml to ensure version consistency DESC=$(git tag --list --sort=-version:refname 'v*.*.*' | head -1 | xargs git describe --tags --long 2>/dev/null || echo 'v0.0.0-0-g0000000') TAG=$(echo "$DESC" | cut -d- -f1) HASH="g$(git rev-parse --short HEAD)" BUILD_DATE=$(date -u +%Y%m%d) # Increment patch version for nightlies (e.g., v0.5.8 -> 0.5.9) VERSION=${TAG#v} # Remove 'v' prefix MAJOR=$(echo "$VERSION" | cut -d. -f1) MINOR=$(echo "$VERSION" | cut -d. -f2) PATCH=$(echo "$VERSION" | cut -d. -f3) NEXT_PATCH=$((PATCH + 1)) NEXT_VERSION="${MAJOR}.${MINOR}.${NEXT_PATCH}" # Use date-based dev number for correct chronological sorting # e.g., 0.5.9.dev20260215+g4cf4f0859 > 0.5.9.dev20260214+g45a4697d4 FORCE_VERSION="${NEXT_VERSION}.dev${BUILD_DATE}+${HASH}" echo "Forcing nightly version to: $FORCE_VERSION" export SETUPTOOLS_SCM_PRETEND_VERSION="$FORCE_VERSION" # Build wheel python3 -m build --wheel # Extract version from built wheel filename WHEEL_FILE=$(ls dist/*.whl) NIGHTLY_VERSION=$(echo "$WHEEL_FILE" | sed 's/.*sglang-\(.*\)-py3.*/\1/') # Get commit info COMMIT_HASH=$(git rev-parse --short HEAD) BUILD_DATE=$(date -u +%Y-%m-%d) echo "Built wheel: $WHEEL_FILE" echo "Nightly version: ${NIGHTLY_VERSION}" echo "Commit: ${COMMIT_HASH}" echo "Build date: ${BUILD_DATE}" echo "nightly_version=${NIGHTLY_VERSION}" >> $GITHUB_OUTPUT echo "commit_hash=${COMMIT_HASH}" >> $GITHUB_OUTPUT echo "build_date=${BUILD_DATE}" >> $GITHUB_OUTPUT - name: Upload wheel artifact uses: actions/upload-artifact@v4 with: name: nightly-wheel path: python/dist/*.whl retention-days: 7 release-nightly: needs: build-nightly-wheel runs-on: ubuntu-latest environment: 'prod' steps: - uses: actions/checkout@v4 - name: Download wheel artifact uses: actions/download-artifact@v4 with: name: nightly-wheel path: dist/ - name: List downloaded wheels run: | echo "Downloaded wheel:" ls -lh dist/ - name: Create GitHub Release for nightly wheel uses: softprops/action-gh-release@v2 with: tag_name: nightly-${{ needs.build-nightly-wheel.outputs.build_date }}-${{ needs.build-nightly-wheel.outputs.commit_hash }} name: Nightly Build ${{ needs.build-nightly-wheel.outputs.build_date }} (${{ needs.build-nightly-wheel.outputs.commit_hash }}) repository: sgl-project/whl token: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} prerelease: true body: | Nightly build from commit ${{ github.sha }} Build date: ${{ needs.build-nightly-wheel.outputs.build_date }} Version: ${{ needs.build-nightly-wheel.outputs.nightly_version }} files: | dist/*.whl - name: Clone wheel index repository run: | git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl cd sgl-whl git config --local user.name "sglang-bot" git config --local user.email "sglangbot@gmail.com" env: WHL_TOKEN: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} - name: Set up Python uses: actions/setup-python@v5 with: python-version: "3.10" - name: Update wheel index run: | python3 scripts/update_nightly_whl_index.py \ --commit-hash ${{ needs.build-nightly-wheel.outputs.commit_hash }} \ --nightly-version ${{ needs.build-nightly-wheel.outputs.nightly_version }} \ --cuda-version ${{ inputs.cuda_version || '129' }} \ --build-date ${{ needs.build-nightly-wheel.outputs.build_date }} - name: Push wheel index run: | cd sgl-whl git add -A git diff --staged --quiet || git commit -m "Update nightly wheel index for commit ${{ needs.build-nightly-wheel.outputs.commit_hash }}" git push ================================================ FILE: .github/workflows/release-pypi-pr.yml ================================================ name: Release PyPI PR Wheels on: workflow_dispatch: inputs: pr_number: description: 'PR number to build wheel for (works with both internal and fork PRs)' required: true type: string concurrency: group: build-pr-wheel-${{ github.event.inputs.pr_number }} cancel-in-progress: true jobs: build-pr-wheel: if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-latest outputs: wheel_version: ${{ steps.gen_version.outputs.wheel_version }} commit_hash: ${{ steps.gen_version.outputs.commit_hash }} build_date: ${{ steps.gen_version.outputs.build_date }} steps: - uses: actions/checkout@v4 with: ref: refs/pull/${{ inputs.pr_number }}/head fetch-depth: 0 # Need full history for version generation - name: Set up Python uses: actions/setup-python@v5 with: python-version: "3.10" - name: Generate PR wheel version id: gen_version run: | # Get base version from the latest v*.*.* git tag directly # Note: We cannot use setuptools_scm here because the [tool.setuptools_scm] # config (with custom git_describe_command) lives in python/pyproject.toml, # not at the repo root. Without that config, setuptools_scm falls back to # default git describe which finds gateway-* tags instead of v*.*.* release tags. LATEST_TAG=$(git tag --list --sort=-version:refname 'v*.*.*' | head -1) BASE_VERSION=${LATEST_TAG#v} echo "Latest release tag: ${LATEST_TAG}" # Get commit info COMMIT_HASH=$(git rev-parse --short HEAD) COMMIT_COUNT=$(git rev-list --count HEAD) # Get current date in YYYY-MM-DD format BUILD_DATE=$(date -u +%Y-%m-%d) # Always use pr-{number} format for suffix SUFFIX="pr-${{ inputs.pr_number }}" # Generate PR wheel version following PEP 440 # Format: {base_version}.dev{commit_count}+pr-{number}.g{commit_hash} WHEEL_VERSION="${BASE_VERSION}.dev${COMMIT_COUNT}+${SUFFIX}.g${COMMIT_HASH}" echo "Base version: ${BASE_VERSION}" echo "PR wheel version: ${WHEEL_VERSION}" echo "Commit: ${COMMIT_HASH}" echo "Build date: ${BUILD_DATE}" echo "wheel_version=${WHEEL_VERSION}" >> $GITHUB_OUTPUT echo "commit_hash=${COMMIT_HASH}" >> $GITHUB_OUTPUT echo "base_version=${BASE_VERSION}" >> $GITHUB_OUTPUT echo "build_date=${BUILD_DATE}" >> $GITHUB_OUTPUT - name: Update pyproject.toml with PR wheel version run: | cd python WHEEL_VERSION="${{ steps.gen_version.outputs.wheel_version }}" # Update pyproject.toml to use static version instead of dynamic # Remove 'version' from dynamic list and add static version sed -i 's/dynamic = \["version"\]/dynamic = []/' pyproject.toml sed -i "/^name = \"sglang\"/a version = \"${WHEEL_VERSION}\"" pyproject.toml # Verify update echo "Updated version in pyproject.toml:" grep "^version" pyproject.toml grep "^dynamic" pyproject.toml - name: Install build dependencies run: | cd python pip install build wheel setuptools - name: Build wheel run: | cd python cp ../README.md ../LICENSE . python3 -m build --wheel # List built wheels echo "Built wheel:" ls -lh dist/ - name: Upload wheel artifact uses: actions/upload-artifact@v4 with: name: pr-wheel-${{ inputs.pr_number }} path: python/dist/*.whl retention-days: 30 release-pr-wheel: needs: build-pr-wheel runs-on: ubuntu-latest environment: 'prod' steps: - uses: actions/checkout@v4 - name: Download wheel artifact uses: actions/download-artifact@v4 with: name: pr-wheel-${{ inputs.pr_number }} path: dist/ - name: List downloaded wheels run: | echo "Downloaded wheel:" ls -lh dist/ - name: Create GitHub Release for PR wheel uses: softprops/action-gh-release@v2 with: tag_name: pr-${{ inputs.pr_number }}-${{ needs.build-pr-wheel.outputs.build_date }}-${{ needs.build-pr-wheel.outputs.commit_hash }} name: "PR #${{ inputs.pr_number }} Build (${{ needs.build-pr-wheel.outputs.commit_hash }})" repository: sgl-project/whl token: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} prerelease: true body: | PR wheel build from PR #${{ inputs.pr_number }} Commit: ${{ github.sha }} Build date: ${{ needs.build-pr-wheel.outputs.build_date }} Version: ${{ needs.build-pr-wheel.outputs.wheel_version }} **Installation via index (pip):** ```bash pip install sglang==${{ needs.build-pr-wheel.outputs.wheel_version }} --index-url https://sgl-project.github.io/whl/pr/ ``` **Installation via index (uv):** ```bash uv pip install sglang==${{ needs.build-pr-wheel.outputs.wheel_version }} --index-url https://sgl-project.github.io/whl/pr/ --extra-index-url https://pypi.org/simple --index-strategy unsafe-best-match ``` **Direct installation:** ```bash pip install https://github.com/sgl-project/whl/releases/download/pr-${{ inputs.pr_number }}-${{ needs.build-pr-wheel.outputs.build_date }}-${{ needs.build-pr-wheel.outputs.commit_hash }}/sglang-${{ needs.build-pr-wheel.outputs.wheel_version }}-py3-none-any.whl ``` files: | dist/*.whl - name: Clone wheel index repository run: | git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl cd sgl-whl git config --local user.name "sglang-bot" git config --local user.email "sglangbot@gmail.com" env: WHL_TOKEN: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} - name: Set up Python uses: actions/setup-python@v5 with: python-version: "3.10" - name: Update wheel index run: | python3 scripts/update_pr_whl_index.py \ --pr-number ${{ inputs.pr_number }} \ --commit-hash ${{ needs.build-pr-wheel.outputs.commit_hash }} \ --wheel-version ${{ needs.build-pr-wheel.outputs.wheel_version }} \ --build-date ${{ needs.build-pr-wheel.outputs.build_date }} - name: Push wheel index run: | cd sgl-whl git add -A git diff --staged --quiet || git commit -m "Update PR wheel index for PR #${{ inputs.pr_number }} (commit ${{ needs.build-pr-wheel.outputs.commit_hash }})" git push ================================================ FILE: .github/workflows/release-pypi.yml ================================================ name: Release PyPI on: push: tags: - 'v[0-9]+.*' workflow_dispatch: jobs: publish: if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-latest environment: "prod" steps: - name: Set up Python uses: actions/setup-python@v4 with: python-version: "3.10" - name: Checkout repository uses: actions/checkout@v4 with: fetch-depth: 0 # Required for setuptools-scm to determine version from tags - name: Upload to pypi run: | cd python cp ../README.md ../LICENSE . pip install build wheel setuptools setuptools-scm python3 -m build pip install twine python3 -m twine upload dist/* -u __token__ -p ${{ secrets.PYPI_TOKEN }} ================================================ FILE: .github/workflows/release-tag.yml ================================================ name: Release Tag # Creates a git tag to trigger release workflows (PyPI, Docker) # Use this after testing on a release branch is complete on: workflow_dispatch: inputs: version: description: 'Version to tag (without v prefix, e.g., 0.5.7)' required: true type: string ref: description: 'Branch or commit to tag (e.g., release/v0.5.7, main, or commit SHA)' required: false default: 'main' type: string permissions: contents: write jobs: create-tag: if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-latest environment: 'prod' steps: - name: Validate version format run: | VERSION="${{ github.event.inputs.version }}" if [ -z "$VERSION" ]; then echo "::error::Version is required" exit 1 fi if ! echo "$VERSION" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]+'; then echo "::error::Invalid version format: $VERSION (expected: X.Y.Z or X.Y.Z.postN)" exit 1 fi echo "Version validated: v$VERSION" - name: Checkout repository uses: actions/checkout@v4 with: ref: ${{ github.event.inputs.ref }} fetch-depth: 0 token: ${{ secrets.GITHUB_TOKEN }} - name: Check if tag already exists run: | TAG="v${{ github.event.inputs.version }}" if git rev-parse "$TAG" >/dev/null 2>&1; then echo "::error::Tag $TAG already exists" exit 1 fi echo "Tag $TAG does not exist, proceeding..." - name: Create and push tag run: | TAG="v${{ github.event.inputs.version }}" REF="${{ github.event.inputs.ref }}" git config user.name "sglang-bot" git config user.email "sglang-bot@users.noreply.github.com" echo "Creating tag $TAG on ref $REF (commit: $(git rev-parse HEAD))" git tag -a "$TAG" -m "Release $TAG" git push origin "$TAG" echo "::notice::Successfully created and pushed tag $TAG" echo "This will trigger the release workflows (PyPI, Docker)" ================================================ FILE: .github/workflows/release-whl-kernel.yml ================================================ name: Release SGLang Kernels on: push: branches: - main paths: - sgl-kernel/python/sgl_kernel/version.py workflow_dispatch: inputs: target: type: choice description: 'Build target' required: false default: 'all' options: - 'all' - 'cu129' - 'cu130' - 'rocm700' - 'rocm720' - 'musa43' tag_name: type: string required: false pr_number: description: "PR number to build from (e.g. 12345)" type: string required: false concurrency: group: release-sglang-kernels-${{ github.ref }} cancel-in-progress: true jobs: build-cu129-matrix: if: | github.repository == 'sgl-project/sglang' && (github.event.inputs.target == 'all' || github.event.inputs.target == 'cu129') strategy: matrix: python-version: ["3.10"] cuda-version: ["12.9"] arch: [x86_64, aarch64] include: - arch: x86_64 runner: x64-kernel-build-node - arch: aarch64 runner: arm-kernel-build-node runs-on: ${{ matrix.runner }} steps: - uses: actions/checkout@v4 with: submodules: "recursive" ref: ${{ inputs.pr_number && format('refs/pull/{0}/head', inputs.pr_number) || '' }} - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Build wheels run: | cd sgl-kernel chmod +x ./build.sh ./build.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" ${{ matrix.arch == 'aarch64' && 'aarch64' || '' }} env: BUILD_JOBS: 64 NVCC_THREADS: 8 - name: Upload to PyPI working-directory: sgl-kernel run: | pip install twine python3 -m twine upload --skip-existing dist/* -u __token__ -p ${{ secrets.PYPI_TOKEN_SGLANG_KERNEL }} - name: Upload artifacts uses: actions/upload-artifact@v4 with: name: wheel-python${{ matrix.python-version }}-cuda${{ matrix.cuda-version }}${{ matrix.arch == 'aarch64' && '-aarch64' || '' }} path: sgl-kernel/dist/* release-cu129: needs: build-cu129-matrix runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 with: ref: ${{ inputs.pr_number && format('refs/pull/{0}/head', inputs.pr_number) || '' }} - name: Download artifacts uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-* - name: Set tag name id: set_tag_name run: | if [ -z "${{ inputs.tag_name }}" ]; then TAG_NAME="v$(cat sgl-kernel/python/sgl_kernel/version.py | cut -d'"' -f2)" echo "tag_name=$TAG_NAME" >> $GITHUB_OUTPUT else echo "tag_name=${{ inputs.tag_name }}" >> $GITHUB_OUTPUT fi - name: Release uses: softprops/action-gh-release@v2 with: tag_name: ${{ steps.set_tag_name.outputs.tag_name }} repository: sgl-project/whl token: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} files: | sgl-kernel/dist/* - name: Clone wheel index run: git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl env: WHL_TOKEN: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} - name: Update wheel index run: python3 scripts/update_kernel_whl_index.py --cuda 129 - name: Push wheel index run: | cd sgl-whl git config --local user.name "sglang-bot" git config --local user.email "sglangbot@gmail.com" git add -A git commit -m "update whl index" git push # for now we do not release CUDA 13.0 wheels to pypi build-cu130-matrix: if: | github.repository == 'sgl-project/sglang' && (github.event.inputs.target == 'all' || github.event.inputs.target == 'cu130') strategy: matrix: python-version: ["3.10"] cuda-version: ["13.0"] arch: [x86_64, aarch64] include: - arch: x86_64 runner: x64-kernel-build-node - arch: aarch64 runner: arm-kernel-build-node runs-on: ${{ matrix.runner }} steps: - uses: actions/checkout@v4 with: submodules: "recursive" ref: ${{ inputs.pr_number && format('refs/pull/{0}/head', inputs.pr_number) || '' }} - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Build wheels run: | cd sgl-kernel chmod +x ./build.sh ./build.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" ${{ matrix.arch == 'aarch64' && 'aarch64' || '' }} env: BUILD_JOBS: 64 NVCC_THREADS: 8 - name: Upload artifacts uses: actions/upload-artifact@v4 with: name: wheel-python${{ matrix.python-version }}-cuda${{ matrix.cuda-version }}${{ matrix.arch == 'aarch64' && '-aarch64' || '' }} path: sgl-kernel/dist/* release-cu130: needs: build-cu130-matrix runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 with: ref: ${{ inputs.pr_number && format('refs/pull/{0}/head', inputs.pr_number) || '' }} - name: Download artifacts uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-* - name: Set tag name id: set_tag_name run: | if [ -z "${{ inputs.tag_name }}" ]; then TAG_NAME="v$(cat sgl-kernel/python/sgl_kernel/version.py | cut -d'"' -f2)" echo "tag_name=$TAG_NAME" >> $GITHUB_OUTPUT else echo "tag_name=${{ inputs.tag_name }}" >> $GITHUB_OUTPUT fi - name: Release uses: softprops/action-gh-release@v2 with: tag_name: ${{ steps.set_tag_name.outputs.tag_name }} repository: sgl-project/whl token: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} files: | sgl-kernel/dist/* - name: Clone wheel index run: git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl env: WHL_TOKEN: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} - name: Update wheel index run: python3 scripts/update_kernel_whl_index.py --cuda 130 - name: Push wheel index run: | cd sgl-whl git config --local user.name "sglang-bot" git config --local user.email "sglangbot@gmail.com" git add -A git commit -m "update whl index" git push build-rocm-matrix: if: | github.repository == 'sgl-project/sglang' && (github.event.inputs.target == 'all' || github.event.inputs.target == 'rocm700' || github.event.inputs.target == 'rocm720') runs-on: amd-docker-scale strategy: matrix: python-version: ["3.10"] rocm-version: ["700", "720"] steps: - uses: actions/checkout@v4 with: submodules: "recursive" ref: ${{ inputs.pr_number && format('refs/pull/{0}/head', inputs.pr_number) || '' }} - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Build wheels run: | cp 3rdparty/amd/wheel/sgl-kernel/* sgl-kernel/ cd sgl-kernel chmod +x ./build_rocm.sh ./build_rocm.sh "${{ matrix.rocm-version }}" - name: Upload artifacts uses: actions/upload-artifact@v4 with: name: wheel-python${{ matrix.python-version }}-rocm${{ matrix.rocm-version }} path: sgl-kernel/dist/* release-rocm700: needs: build-rocm-matrix runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 with: ref: ${{ inputs.pr_number && format('refs/pull/{0}/head', inputs.pr_number) || '' }} - name: Download artifacts uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-*-rocm700 - name: Set tag name id: set_tag_name run: | if [ -z "${{ inputs.tag_name }}" ]; then TAG_NAME="v$(cat sgl-kernel/python/sgl_kernel/version.py | cut -d'"' -f2)" echo "tag_name=$TAG_NAME" >> $GITHUB_OUTPUT else echo "tag_name=${{ inputs.tag_name }}" >> $GITHUB_OUTPUT fi - name: Release uses: softprops/action-gh-release@v2 with: tag_name: ${{ steps.set_tag_name.outputs.tag_name }} repository: sgl-project/whl token: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} files: | sgl-kernel/dist/* - name: Clone wheel index run: git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl env: WHL_TOKEN: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} - name: Update wheel index run: python3 scripts/update_kernel_whl_index.py --rocm 700 - name: Push wheel index run: | cd sgl-whl git config --local user.name "sglang-bot" git config --local user.email "sglangbot@gmail.com" git add -A git commit -m "update whl index" git push release-rocm720: needs: build-rocm-matrix runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 with: ref: ${{ inputs.pr_number && format('refs/pull/{0}/head', inputs.pr_number) || '' }} - name: Download artifacts uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-*-rocm720 - name: Set tag name id: set_tag_name run: | if [ -z "${{ inputs.tag_name }}" ]; then TAG_NAME="v$(cat sgl-kernel/python/sgl_kernel/version.py | cut -d'"' -f2)" echo "tag_name=$TAG_NAME" >> $GITHUB_OUTPUT else echo "tag_name=${{ inputs.tag_name }}" >> $GITHUB_OUTPUT fi - name: Release uses: softprops/action-gh-release@v2 with: tag_name: ${{ steps.set_tag_name.outputs.tag_name }} repository: sgl-project/whl token: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} files: | sgl-kernel/dist/* - name: Clone wheel index run: git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl env: WHL_TOKEN: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} - name: Update wheel index run: python3 scripts/update_kernel_whl_index.py --rocm 720 - name: Push wheel index run: | cd sgl-whl git config --local user.name "sglang-bot" git config --local user.email "sglangbot@gmail.com" git add -A git commit -m "update whl index" git push build-musa43: if: | github.repository == 'sgl-project/sglang' && (github.event.inputs.target == 'all' || github.event.inputs.target == 'musa43') runs-on: kernel-build-node-musa strategy: matrix: python-version: ["3.10"] musa-version: ["43"] steps: - uses: actions/checkout@v4 with: submodules: "recursive" - name: Build wheels run: | cd sgl-kernel mv pyproject_musa.toml pyproject.toml python setup_musa.py sdist bdist_wheel - name: Rename MUSA wheels run: | bash scripts/ci/musa/rename_wheels_musa.sh ${{ matrix.musa-version }} sgl-kernel/dist - name: Upload artifacts uses: actions/upload-artifact@v4 with: name: wheel-python${{ matrix.python-version }}-musa${{ matrix.musa-version }} path: sgl-kernel/dist/* release-musa43: needs: build-musa43 runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - name: Download artifacts uses: actions/download-artifact@v4 with: path: sgl-kernel/dist/ merge-multiple: true pattern: wheel-* - name: Set tag name id: set_tag_name run: | if [ -z "${{ inputs.tag_name }}" ]; then TAG_NAME="v$(cat sgl-kernel/python/sgl_kernel/version.py | cut -d'"' -f2)" echo "tag_name=$TAG_NAME" >> $GITHUB_OUTPUT else echo "tag_name=${{ inputs.tag_name }}" >> $GITHUB_OUTPUT fi - name: Release uses: softprops/action-gh-release@v2 with: tag_name: ${{ steps.set_tag_name.outputs.tag_name }} repository: sgl-project/whl token: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} files: | sgl-kernel/dist/* - name: Clone wheel index run: git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl env: WHL_TOKEN: ${{ secrets.GH_PAT_FOR_WHL_RELEASE }} - name: Update wheel index run: python3 scripts/update_kernel_whl_index.py --musa 43 - name: Push wheel index run: | cd sgl-whl git config --local user.name "sglang-bot" git config --local user.email "sglangbot@gmail.com" git add -A git commit -m "update whl index" git push ================================================ FILE: .github/workflows/rerun-ut.yml ================================================ name: Rerun UT run-name: ${{ inputs.pr_head_sha && format('[rerun-ut] {0}', inputs.pr_head_sha) || '[rerun-ut]' }} on: workflow_dispatch: inputs: test_command: description: "Test command to run (e.g. 'registered/core/test_srt_endpoint.py TestSRTEndpoint.test_simple_decode')" required: true type: string runner_label: description: "Runner label" required: true type: choice options: - 1-gpu-runner - 1-gpu-5090 - 2-gpu-runner - 4-gpu-h100 - 4-gpu-a10 - 4-gpu-b200 - 8-gpu-h200 - 8-gpu-h20 - 8-gpu-b200 pr_head_sha: description: "PR head SHA to checkout (for /rerun-ut on fork PRs)" required: false type: string default: "" use_deepep: description: "Use ci_install_deepep.sh instead of ci_install_dependency.sh" required: false type: string default: "false" env: SGLANG_IS_IN_CI: true SGLANG_CUDA_COREDUMP: "1" SGLANG_JIT_DEEPGEMM_FAST_WARMUP: true permissions: actions: write contents: read jobs: rerun-ut-cuda: runs-on: ${{ inputs.runner_label }} timeout-minutes: 120 env: RUNNER_LABELS: ${{ inputs.runner_label }} IS_BLACKWELL: ${{ (inputs.runner_label == '1-gpu-5090' || contains(inputs.runner_label, 'b200')) && '1' || '0' }} SGLANG_CI_RDMA_ALL_DEVICES: ${{ inputs.runner_label == '8-gpu-h20' && 'mlx5_1,mlx5_2,mlx5_3,mlx5_4' || '' }} steps: - name: Checkout code uses: actions/checkout@v4 with: ref: ${{ inputs.pr_head_sha || github.sha }} - name: Install dependencies timeout-minutes: 20 run: | if [[ "${{ inputs.runner_label }}" == "1-gpu-5090" ]]; then source /etc/profile.d/sglang-ci.sh fi if [[ "${{ inputs.use_deepep }}" == "true" ]]; then bash scripts/ci/cuda/ci_install_deepep.sh else bash scripts/ci/cuda/ci_install_dependency.sh fi - name: Run test timeout-minutes: 60 run: | if [[ "${{ inputs.runner_label }}" == "1-gpu-5090" ]]; then source /etc/profile.d/sglang-ci.sh fi cd test/ python3 ${{ inputs.test_command }} - uses: ./.github/actions/upload-cuda-coredumps if: always() ================================================ FILE: .github/workflows/retag-docker.yml ================================================ name: Retag Docker Image on: workflow_dispatch: inputs: source_tag: description: "Existing image tag (e.g., v0.4.7-cu129-amd64)" required: true target_tag: description: "New tag to apply (e.g., latest)" required: true jobs: retag: if: github.repository == 'sgl-project/sglang' runs-on: ubuntu-22.04 environment: "prod" steps: - name: Login to Docker Hub uses: docker/login-action@v2 with: username: ${{ secrets.DOCKERHUB_USERNAME }} password: ${{ secrets.DOCKERHUB_TOKEN }} - name: Retag image run: | echo "Retagging lmsysorg/sglang:${{ inputs.source_tag }} -> lmsysorg/sglang:${{ inputs.target_tag }}" docker buildx imagetools create \ -t lmsysorg/sglang:${{ inputs.target_tag }} \ lmsysorg/sglang:${{ inputs.source_tag }} ================================================ FILE: .github/workflows/runner-utilization.yml ================================================ name: Runner Utilization Report on: schedule: - cron: '0 8 * * *' # Daily at 8 AM UTC pull_request: paths: - '.github/workflows/runner-utilization.yml' - 'scripts/ci/utils/runner_utilization_report.py' workflow_dispatch: inputs: hours: description: 'Time window in hours' required: false default: '24' type: string filter: description: 'Filter runner labels (e.g., 5090, h200)' required: false type: string jobs: report: name: Generate Report runs-on: ubuntu-latest steps: - name: Checkout code uses: actions/checkout@v4 - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Generate Utilization Report timeout-minutes: 30 env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} run: | python scripts/ci/utils/runner_utilization_report.py \ --repo ${{ github.repository }} \ --hours ${{ inputs.hours || '24' }} \ ${{ inputs.filter && format('--filter {0}', inputs.filter) || '' }} ================================================ FILE: .github/workflows/slash-command-handler.yml ================================================ name: Slash Command Handler on: issue_comment: types: [created, edited] permissions: contents: read pull-requests: write # Required to add labels and reactions actions: write # Required to rerun workflows issues: write # Required for comment reactions in some contexts jobs: slash_command: # Only run if it is a PR and the comment contains a recognized command # Use contains() since startsWith() can't handle leading whitespace/newlines if: > github.event.issue.pull_request && (contains(github.event.comment.body, '/tag-run-ci-label') || contains(github.event.comment.body, '/rerun-failed-ci') || contains(github.event.comment.body, '/tag-and-rerun-ci') || contains(github.event.comment.body, '/rerun-stage') || contains(github.event.comment.body, '/rerun-ut')) runs-on: ubuntu-latest steps: # SECURITY: This workflow runs on issue_comment trigger with elevated permissions # (pull-requests: write, actions: write). For non-fork PRs, we can safely checkout # the PR branch to allow testing changes to this handler. For fork PRs, we MUST # stay on main to prevent untrusted code execution with these elevated permissions. - name: Get PR details id: pr shell: bash env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} run: | PR_DATA=$(gh pr view ${{ github.event.issue.number }} --repo ${{ github.repository }} --json headRefName,headRepositoryOwner) || { echo "::error::Failed to fetch PR data" exit 1 } # Use 'empty' filter to handle null/missing values (e.g., deleted forks) HEAD_OWNER=$(echo "$PR_DATA" | jq -r '.headRepositoryOwner.login // empty') REPO_OWNER="${{ github.repository_owner }}" # Treat missing/null owner as fork for security (fail-safe) if [[ -z "$HEAD_OWNER" || "$HEAD_OWNER" != "$REPO_OWNER" ]]; then IS_FORK="true" else IS_FORK="false" fi echo "is_fork=$IS_FORK" >> $GITHUB_OUTPUT echo "ref=$(echo "$PR_DATA" | jq -r '.headRefName')" >> $GITHUB_OUTPUT echo "PR owner: $HEAD_OWNER, Repo owner: $REPO_OWNER, Is fork: $IS_FORK" - name: Checkout code uses: actions/checkout@v4 with: # For non-fork PRs, checkout PR branch to allow testing handler changes # For fork PRs, stay on main for security (don't run untrusted code with elevated permissions) ref: ${{ steps.pr.outputs.is_fork == 'false' && steps.pr.outputs.ref || '' }} - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Install dependencies run: | pip install PyGithub - name: Handle Slash Command env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} REPO_FULL_NAME: ${{ github.repository }} PR_NUMBER: ${{ github.event.issue.number }} COMMENT_ID: ${{ github.event.comment.id }} COMMENT_BODY: ${{ github.event.comment.body }} USER_LOGIN: ${{ github.event.comment.user.login }} run: | python scripts/ci/utils/slash_command_handler.py ================================================ FILE: .github/workflows/stress-test.yml ================================================ name: Stress Test on: workflow_dispatch: inputs: num_prompts: description: 'Number of prompts per model' required: true default: '50000' type: string duration_minutes: description: 'Timeout per model in minutes' required: true default: '45' type: string jobs: stress-test: if: github.repository == 'sgl-project/sglang' runs-on: 8-gpu-h200 steps: - name: Checkout code uses: actions/checkout@v4 - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh - name: Run stress tests timeout-minutes: 210 env: NUM_PROMPTS: ${{ inputs.num_prompts }} DURATION_MINUTES: ${{ inputs.duration_minutes }} run: | cd test python3 run_suite.py --hw cuda --suite stress - name: Upload results if: always() uses: actions/upload-artifact@v4 with: name: stress-test-results path: | stress_test_*.jsonl ================================================ FILE: .github/workflows/weekly-test-nvidia.yml ================================================ name: Weekly Test (Nvidia) on: schedule: - cron: '0 0 * * 0' # Run every Sunday at midnight UTC workflow_dispatch: inputs: job_filter: description: 'Select which job to run (leave empty or "all" to run all jobs)' required: false type: choice default: 'all' options: - 'all' - 'weekly-test-8-gpu-h200' concurrency: group: weekly-test-nvidia-${{ github.ref }} cancel-in-progress: true env: SGLANG_IS_IN_CI: true HF_HUB_DOWNLOAD_TIMEOUT: 300 HF_HUB_ETAG_TIMEOUT: 300 jobs: # Weekly tests - 8 GPU H200 weekly-test-8-gpu-h200: if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'weekly-test-8-gpu-h200') runs-on: 8-gpu-h200 timeout-minutes: 120 env: RUNNER_LABELS: 8-gpu-h200 steps: - name: Checkout code uses: actions/checkout@v4 - name: Install dependencies run: | bash scripts/ci/cuda/ci_install_dependency.sh - name: Run weekly 8-GPU H200 tests timeout-minutes: 120 env: GPU_CONFIG: "8-gpu-h200" IS_H200: "1" run: | cd test python3 run_suite.py --hw cuda --suite weekly-8-gpu-h200 --nightly --continue-on-error --timeout-per-file 7200 ================================================ FILE: .gitignore ================================================ # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class # C extensions *.so # Distribution / packaging .Python **/build/ **/develop-eggs/ **/dist/ **/downloads/ **/eggs/ .eggs/ **/lib/ **/lib64/ **/parts/ **/sdist/ **/var/ **/wheels/ **/share/python-wheels/ *.egg-info/ .installed.cfg *.egg MANIFEST # PyInstaller # Usually these files are written by a python script from a template # before PyInstaller builds the exe, so as to inject date/other infos into it. *.manifest *.spec # Installer logs pip-log.txt pip-delete-this-directory.txt # Unit test / coverage reports htmlcov/ .tox/ .nox/ .coverage .coverage.* .cache nosetests.xml coverage.xml *.cover *.py,cover .hypothesis/ # Tokenizer cache for tests .tokenizer_cache/ .pytest_cache/ cover/ # Translations *.mo *.pot # Django stuff: *.log local_settings.py db.sqlite3 db.sqlite3-journal # Flask stuff: instance/ .webassets-cache # Scrapy stuff: .scrapy # Sphinx documentation docs/_build/ # PyBuilder .pybuilder/ target/ # Jupyter Notebook .ipynb_checkpoints # IPython profile_default/ ipython_config.py # pyenv # For a library or package, you might want to ignore these files since the code is # intended to run in multiple environments; otherwise, check them in: # .python-version # pipenv # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. # However, in case of collaboration, if having platform-specific dependencies or dependencies # having no cross-platform support, pipenv may install dependencies that don't work, or not # install all needed dependencies. #Pipfile.lock # poetry # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. # This is especially recommended for binary packages to ensure reproducibility, and is more # commonly ignored for libraries. # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control #poetry.lock # pdm # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. #pdm.lock # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it # in version control. # https://pdm.fming.dev/#use-with-ide .pdm.toml # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm __pypackages__/ # Celery stuff celerybeat-schedule celerybeat.pid # SageMath parsed files *.sage.py # Environments .env .venv env/ venv/ ENV/ env.bak/ venv.bak/ # Spyder project settings .spyderproject .spyproject # Rope project settings .ropeproject # mkdocs documentation /site # mypy .mypy_cache/ .dmypy.json dmypy.json # Pyre type checker .pyre/ # pytype static type analyzer .pytype/ # Cython debug symbols cython_debug/ # PyCharm # JetBrains specific template is maintained in a separate JetBrains.gitignore that can # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. .idea/ # MacOS .DS_Store # Vim *.swp # Documentation docs/_build # SGL benchmark/mmlu/data benchmark/mmlu/data.tar benchmark/llava_bench/images benchmark/llava_bench/mme_pack *.jsonl tmp*.txt # Torch Compile logs tl_out/ # Plots *.png *.pdf # personnal work_dirs/ *.csv !logo.png # Prerequisites *.d # Compiled Object files *.slo *.lo *.o *.obj # Precompiled Headers *.gch *.pch # Compiled Dynamic libraries *.so *.dylib *.dll # Fortran module files *.mod *.smod # Compiled Static libraries *.lai *.la *.a *.lib # Executables *.exe *.out *.app *.iml # VSCode .vscode # Autoenv .env.leave # Rust lib Cargo.lock # Generated vision test fixtures (regenerate with: python scripts/generate_vision_golden.py) sgl-model-gateway/tests/fixtures/golden/ # Other repos lmms-eval **/.serena/ ctags/ outputs/ inputs/ # Eval Cache .longbench_cache/ # CUDA kernel develop, profile and debug .clangd *.nsys-rep *.ncu-rep *.nvcudmp # setuptools-scm generated version file python/sglang/_version.py # MUSA section # Generated source files by torchada sgl-kernel/csrc_musa/ sgl-kernel/include_musa/ sgl-kernel/csrc/**/*_musa/ # MUSA core dump files *.mudmp # Others # diffusion 3D outputs *.glb *.ply *.npz ================================================ FILE: .isort.cfg ================================================ [settings] profile=black known_first_party=sglang ================================================ FILE: .pre-commit-config.yaml ================================================ default_stages: [pre-commit, pre-push, manual] exclude: ^(python/sglang/multimodal_gen/csrc|python/sglang/jit_kernel/flash_attention/cute) repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v6.0.0 hooks: - id: check-symlinks - id: destroyed-symlinks - id: trailing-whitespace - id: end-of-file-fixer - id: check-yaml args: [--allow-multiple-documents] - id: check-toml - id: check-ast - id: check-added-large-files - id: check-merge-conflict - id: check-shebang-scripts-are-executable - id: detect-private-key exclude: ^sgl-model-gateway/tests/.*_test\.rs$ - id: debug-statements - id: no-commit-to-branch - repo: https://github.com/PyCQA/isort rev: 7.0.0 hooks: - id: isort exclude: '^python/sglang/srt/grpc/.*_pb2\.py$|^python/sglang/srt/grpc/.*_pb2_grpc\.py$|^python/sglang/srt/grpc/.*_pb2\.pyi$|^python/sglang/srt/grpc/.*_pb2_grpc\.pyi$' - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.15.1 hooks: - id: ruff args: - --select=F401,F821 - --fix files: ^(benchmark/|docs/|examples/|python/sglang/|sgl-model-gateway/py_*|test/) exclude: | (?x)^( .*/__init__\.py$| .*\.ipynb$| python/sglang/srt/grpc/.*_pb2\.py$| python/sglang/srt/grpc/.*_pb2_grpc\.py$| python/sglang/srt/grpc/.*_pb2\.pyi$| python/sglang/srt/grpc/.*_pb2_grpc\.pyi$| )$ - repo: https://github.com/psf/black rev: 26.1.0 hooks: - id: black-jupyter exclude: '^python/sglang/srt/grpc/.*_pb2\.py$|^python/sglang/srt/grpc/.*_pb2_grpc\.py$|^python/sglang/srt/grpc/.*_pb2\.pyi$|^python/sglang/srt/grpc/.*_pb2_grpc\.pyi$' - repo: https://github.com/codespell-project/codespell rev: v2.4.1 hooks: - id: codespell args: ['--config', '.codespellrc'] - repo: https://github.com/pre-commit/mirrors-clang-format rev: v20.1.7 hooks: - id: clang-format types_or: [c++, cuda] args: [--style=file, --verbose] - repo: https://github.com/kynan/nbstripout rev: 0.9.0 hooks: - id: nbstripout args: - '--keep-output' - '--extra-keys=metadata.kernelspec metadata.language_info.version' - repo: local hooks: - id: check-chinese-characters name: check chinese characters in multimodal_gen entry: >- python3 -c 'import sys, re; p=re.compile(r"[\u4e00-\u9fff]"); ec=0; [ ([(print(f"{f}:{i+1}: {l.strip()}") or (ec:=1)) for i,l in enumerate(open(f, "r", encoding="utf-8", errors="ignore")) if p.search(l)]) for f in sys.argv[1:] ]; sys.exit(ec)' language: system files: ^python/sglang/multimodal_gen/.* exclude: ^(python/sglang/multimodal_gen/configs/sample|python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows|python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages)(/|$) types_or: [python, markdown, json, text] - id: sort-ci-permissions name: sort CI_PERMISSIONS.json entry: python3 .github/update_ci_permission.py --sort-only language: system files: ^\.github/CI_PERMISSIONS\.json$ pass_filenames: false ================================================ FILE: 3rdparty/amd/profiling/PROFILING.md ================================================ ## Profiling SGLang Infer System with AMD GPUs This AppNote describes the SGLang profiling technical, code augment and running steps for systems with AMD Instinct GPUs, nevertheless the same procedure may work with Nvidia GPUs too. Examples and steps are provided in detail, to facilitate easy reproduce and use to localize performance problem towards optimizations. Two primary methods are covered: - [RPD](https://github.com/ROCm/rocmProfileData.git) - [PyTorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) ### Profiling SGLang Infer System with RPD Profiler RPD profiler is a low-overhead cross-platform profiler. Therefore, the same RPD code augment not only works for profiling on ROCm/AMD GPUs, but also works for profiling on CUDA/Nvidia GPUs as well. To do RPD profiling on SGLang repository, please use scripts and patch files included in this directory and follow the steps below: 1. Install RPD with rpd.patch applied during installation using install_rpd.sh, both files are in this directory. install_rpd.sh ```bash # download and install RPD apt update && apt install -y sqlite3 libsqlite3-dev libfmt-dev # install rpd module git clone https://github.com/ROCmSoftwarePlatform/rocmProfileData cd rocmProfileData git checkout 976899e9c6dbc6dd2bccf770818e4e44125590ac git apply rpd.patch make && make install cd rocpd_python && python setup.py install && cd .. cd rpd_tracer && make clean;make install && python setup.py install && cd .. ``` rpd.patch ```bash diff --git a/rpd_tracer/Makefile b/rpd_tracer/Makefile index e9d9feb..b2e9e1a 100644 --- a/rpd_tracer/Makefile +++ b/rpd_tracer/Makefile @@ -16,7 +16,7 @@ ifneq (,$(HIP_PATH)) $(info Building with roctracer) RPD_LIBS += -L/opt/rocm/lib -lroctracer64 -lroctx64 -lamdhip64 -lrocm_smi64 RPD_INCLUDES += -I/opt/rocm/include -I/opt/rocm/include/roctracer -I/opt/rocm/include/hsa - RPD_SRCS += RoctracerDataSource.cpp RocmSmiDataSource.cpp + RPD_SRCS += RoctracerDataSource.cpp RPD_INCLUDES += -D__HIP_PLATFORM_AMD__ endif ``` 2. Add loadTracer.sh file included in this directory to /sglang/python/sglang. loadTracer.sh ```bash #!/bin/bash ################################################################################ # Copyright (c) 2021 - 2023 Advanced Micro Devices, Inc. All rights reserved. # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN # THE SOFTWARE. ################################################################################ OUTPUT_FILE="trace.rpd" if [ "$1" = "-o" ] ; then OUTPUT_FILE=$2 shift shift fi if [ -e ${OUTPUT_FILE} ] ; then rm ${OUTPUT_FILE} fi python3 -m rocpd.schema --create ${OUTPUT_FILE} if [ $? != 0 ] ; then echo "Error: Could not create rpd file. Please run 'python setup.py install' from the rocpd_python dir" exit fi export RPDT_FILENAME=${OUTPUT_FILE} export RPDT_AUTOSTART=0 LD_PRELOAD=librocm-smi_64:librpd_tracer.so "$@" ``` 3. Apply patch (provided in this directory) with "git apply rpd_profile_server_enable.patch" if the main profiling purpose is to get info on gpu kernels as well as limited cpu activity info. #### Common Notes 1 Please note that although we are doing TP=8 in the example, we purposely only log RPD profiling on 2 ranks in the patch file (i.e.tp_rank=0/1) for profiling/visualization convenience, as even Perfetto streaming mode can only load maximal 8GB json file for visualization. With 2 ranks logged in RPD profiling, we could still check whether there are issues among ranks (e.g. load imbalance issue, nccl issue), and at the same time, we could log relatively longer time duration before the json file generated from RPD file hits 8GB size. rpd_profile_server_enable.patch ```bash diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 62d1ff9..9021c01 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -71,6 +71,8 @@ from sglang.srt.utils import ( suppress_other_loggers, ) from sglang.utils import get_exception_traceback +from rpdTracerControl import rpdTracerControl +rpdTracerControl.skipCreate() logger = logging.getLogger(__name__) @@ -245,6 +247,7 @@ class Scheduler: ], with_stack=True, ) + self.rpd = rpdTracerControl() @torch.inference_mode() def event_loop(self): @@ -1027,15 +1030,24 @@ class Scheduler: def start_profile(self) -> None: if self.profiler is None: raise RuntimeError("Profiler is not enabled.") - self.profiler.start() + #self.profiler.start() #block pytorch profiler for rpd profiler enabling + if self.tp_rank == 0 or self.tp_rank == 1: + self.rpd.start() + self.rpd.rangePush("", "rpd profile range", "") + logger.info("rpd is enabled") def stop_profile(self) -> None: if self.profiler is None: raise RuntimeError("Profiler is not enabled.") - self.profiler.stop() - self.profiler.export_chrome_trace( - self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" - ) + #self.profiler.stop() + #self.profiler.export_chrome_trace( + # self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" + #) + if self.tp_rank ==0 or self.tp_rank ==1: + self.rpd.rangePop() + self.rpd.stop() + self.rpd.flush() + logger.info("rpd is done") logger.info("Profiler is done") ``` #### Advanced Debugging with RPD Profiler Sometimes, we want to use rpd profiler to capture more CPU and python activities in order to debug some challenging issues (e.g. root cause of load imbalance across gpu processes, root cause of bubbles, etc). Only in such cases, we need to apply patch "git apply rpd_profile_server_enable_wCPU_activities.patch", where 3 files are modified. rpd_profile_server_enable_wCPU_activities.patch ```bash diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 62d1ff9..2edb427 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -71,6 +71,8 @@ from sglang.srt.utils import ( suppress_other_loggers, ) from sglang.utils import get_exception_traceback +from rpdTracerControl import rpdTracerControl +rpdTracerControl.skipCreate() logger = logging.getLogger(__name__) @@ -245,6 +247,7 @@ class Scheduler: ], with_stack=True, ) + self.rpd = rpdTracerControl() @torch.inference_mode() def event_loop(self): @@ -1027,15 +1030,26 @@ class Scheduler: def start_profile(self) -> None: if self.profiler is None: raise RuntimeError("Profiler is not enabled.") - self.profiler.start() + #self.profiler.start() + logger.info("torch profiler is disabled") + if self.tp_rank == 0 or self.tp_rank == 1: + self.rpd.setPythonTrace(True) + self.rpd.start() + self.rpd.rangePush("", "scheduler", "") + logger.info("rpd is enabled inside scheduler profiling") def stop_profile(self) -> None: if self.profiler is None: raise RuntimeError("Profiler is not enabled.") - self.profiler.stop() - self.profiler.export_chrome_trace( - self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" - ) + #self.profiler.stop() + #self.profiler.export_chrome_trace( + # self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" + #) + if self.tp_rank ==0 or self.tp_rank ==1: + self.rpd.rangePop() + self.rpd.stop() + self.rpd.flush() + logger.info("rpd is done inside scheduler") logger.info("Profiler is done") diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 2621ccd..181df85 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -58,6 +58,10 @@ from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.utils import is_generation_model, is_multimodal_model +from rpdTracerControl import rpdTracerControl +rpdTracerControl.skipCreate() + + asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) logger = logging.getLogger(__name__) @@ -514,10 +518,20 @@ class TokenizerManager: self.send_to_scheduler.send_pyobj(req) def start_profile(self): + rpd = rpdTracerControl() + rpd.setPythonTrace(True) + rpd.start() + rpd.rangePush("", "tokenizer_manager", "") + logger.info("tokenizer_manager rpd profiling started!") req = ProfileReq.START_PROFILE self.send_to_scheduler.send_pyobj(req) def stop_profile(self): + rpd = rpdTracerControl() + rpd.rangePop() + rpd.stop() + rpd.flush() + logger.info("rpd profiling is done inside tokenizer_manager!") req = ProfileReq.STOP_PROFILE self.send_to_scheduler.send_pyobj(req) diff --git a/python/sglang/srt/server.py b/python/sglang/srt/server.py index 7111c93..2bd722c 100644 --- a/python/sglang/srt/server.py +++ b/python/sglang/srt/server.py @@ -30,6 +30,8 @@ import threading import time from http import HTTPStatus from typing import Dict, List, Optional, Union +from rpdTracerControl import rpdTracerControl +rpdTracerControl.skipCreate() # Fix a bug of Python threading setattr(threading, "_register_atexit", lambda *args, **kwargs: None) @@ -152,6 +154,11 @@ async def flush_cache(): @app.post("/start_profile") async def start_profile(): """Start profiling.""" + rpd = rpdTracerControl() + rpd.setPythonTrace(True) + rpd.start() + rpd.rangePush("", "server rpd profile range", "") + logger.info("rpd profiling started in server.py!") tokenizer_manager.start_profile() return Response( content="Start profiling.\n", @@ -164,6 +171,11 @@ async def start_profile(): async def stop_profile(): """Stop profiling.""" tokenizer_manager.stop_profile() + rpd = rpdTracerControl() + rpd.rangePop() + rpd.stop() + rpd.flush() + logger.info("rpd profiling is done in server.py!") return Response( content="Stop profiling. This will take some time.\n", status_code=200, ``` 4. As an example for grok1 profiling, we create a dummy_grok1 directory with config.json (see content below) inside this directory and copy this directory to the right path for "--model-path" if you want to use the example server.sh file provided. ```bash cat ../dummy_grok1/config.json { "architectures": [ "Grok1ModelForCausalLM" ], "embedding_multiplier_scale": 78.38367176906169, "output_multiplier_scale": 0.5773502691896257, "vocab_size": 131072, "hidden_size": 6144, "intermediate_size": 32768, "max_position_embeddings": 8192, "num_experts_per_tok": 2, "num_local_experts": 8, "num_attention_heads": 48, "num_hidden_layers": 64, "num_key_value_heads": 8, "head_dim": 128, "rms_norm_eps": 1e-05, "rope_theta": 10000.0, "model_type": "mixtral", "torch_dtype": "bfloat16" } ``` 5. Launch server with rpd enabled script ./server.sh in one terminal inside the docker container. #### Common Notes 2 - Remember to change model-path to the correct path - loadTracer.sh is needed to conduct profiling - SGLANG_TORCH_PROFILER_DIR is used for default torch profiler - Do not use loadTracer.sh if you are using the torch profiler, simply use python3 -m sglang.launch_server. server.sh ```bash #!/bin/bash # export SGLANG_TORCH_PROFILER_DIR=/data/sglang/ export SGLANG_TORCH_PROFILER_DIR=/sgl-workspace/sglang/profile/ # Get the current timestamp TIMESTAMP=$(date +"%Y%m%d_%H%M%S") # Define the log file with a timestamp LOGFILE="sglang_server_log_$TIMESTAMP.json" # Run the Python command and save the output to the log file loadTracer.sh python3 -m sglang.launch_server \ --model-path /sgl-workspace/sglang/dummy_grok1 \ --tokenizer-path Xenova/grok-1-tokenizer \ --load-format dummy \ --quantization fp8 \ --tp 8 \ --port 30000 \ --disable-radix-cache 2>&1 | tee "$LOGFILE" ``` 6. Open another terminal for the same docker container, and run the rpd enabled ./client.sh after you see "The server is fired up and is ready to roll!" message from server side terminal. #### Common Notes 3 - Use curl http://localhost:30000/start_profile & curl http://localhost:30000/stop_profile to control the start and end of profiling. Check sglang/python/sglang/srt/managers/scheduler.py for more details. - Please don't use RPD profiler together with PyTorch profiler to avoid interference. - The rocmProfileData/tools/rpd2tracing.py file is used to generate json file from RPD file. client.sh ```bash #!/bin/bash # Start profiling via API curl http://localhost:30000/start_profile -H "Content-Type: application/json" # Benchmark serving using sglang with random dataset and tokenizer # Define the log file with a timestamp TIMESTAMP=$(date +%Y%m%d_%H%M%S) LOGFILE="sglang_client_log_$TIMESTAMP.json" # Run the benchmark with specified parameters and save logs python3 -m sglang.bench_serving \ --backend sglang \ --tokenizer Xenova/grok-1-tokenizer \ --dataset-name random \ --random-input 1024\ --random-output 1024 \ --num-prompts 120 \ --request-rate 8 \ --output-file online.jsonl 2>&1 | tee "$LOGFILE" # Stop profiling via API curl http://localhost:30000/stop_profile -H "Content-Type: application/json" # Convert tracing file to csv & json sqlite3 trace.rpd ".mode csv" ".header on" ".output trace.csv" "select * from top;" ".output stdout" python3 ./rocmProfileData/tools/rpd2tracing.py trace.rpd trace.json ``` 7. Follow [Perfetto docs](https://perfetto.dev/docs/visualization/large-traces) to visualize large json files. Try to adjust parameters so that the trace.json file size is less than 9GB. ### Profiling SGLang Infer System with PyTorch Profiler Please use the steps as follows: 1. Apply the patch torch_profiler.patch. Note that you can modify "if self.tp_rank == 0" in the patch to allow more ranks be recorded in profiling. torch_profiler.patch ```bash diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 62d1ff9..6ecd78c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -240,7 +240,6 @@ class Scheduler: ) self.profiler = torch.profiler.profile( activities=[ - torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], with_stack=True, @@ -1033,9 +1032,11 @@ class Scheduler: if self.profiler is None: raise RuntimeError("Profiler is not enabled.") self.profiler.stop() - self.profiler.export_chrome_trace( - self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" - ) + if self.tp_rank == 0: + with open(f"stats_repro_{int(time.time())}.txt", "w") as f: + print(self.profiler.key_averages(group_by_input_shape=True).table(sort_by="cuda_time_total", row_limit=-1), file=f) + print("Profiling stats done.") + logger.info("Profiler is done") ``` 2. Create the model path directory and copy it to the right path for "--model-path" if you want to use the server.sh file provided. 3. Modify the included server.sh by removing "loadTracer.sh" before python command and launch script ./server.sh in one terminal inside the docker container. 4. Similar to step 6 in RPD profiling section, but remove the last 2 lines in client.sh, which converted rpd file into csv and json files. Run modified client.sh for PyTorch profiling. ------- - [Torch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) ================================================ FILE: 3rdparty/amd/profiling/client.sh ================================================ #!/bin/bash # Start profiling via API curl http://localhost:30000/start_profile -H "Content-Type: application/json" # Benchmark serving using sglang with random dataset and tokenizer # Define the log file with a timestamp TIMESTAMP=$(date +%Y%m%d_%H%M%S) LOGFILE="sglang_client_log_$TIMESTAMP.json" # Run the benchmark with specified parameters and save logs python3 -m sglang.bench_serving \ --backend sglang \ --tokenizer Xenova/grok-1-tokenizer \ --dataset-name random \ --random-input 1024\ --random-output 1024 \ --num-prompts 240 \ --request-rate 8 \ --output-file online.jsonl 2>&1 | tee "$LOGFILE" # Stop profiling via API curl http://localhost:30000/stop_profile -H "Content-Type: application/json" # Convert tracing file to csv & json sqlite3 trace.rpd ".mode csv" ".header on" ".output trace.csv" "select * from top;" ".output stdout" python3 /sgl-workspace/rocmProfileData/tools/rpd2tracing.py trace.rpd trace.json ================================================ FILE: 3rdparty/amd/profiling/install_rpd.sh ================================================ # download and install RPD apt update && apt install -y sqlite3 libsqlite3-dev libfmt-dev # install rpd module git clone https://github.com/ROCmSoftwarePlatform/rocmProfileData cd rocmProfileData git apply rpd.patch make && make install cd rocpd_python && python setup.py install && cd .. cd rpd_tracer && make clean;make install && python setup.py install && cd .. ================================================ FILE: 3rdparty/amd/profiling/loadTracer.sh ================================================ #!/bin/bash ################################################################################ # Copyright (c) 2021 - 2023 Advanced Micro Devices, Inc. All rights reserved. # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN # THE SOFTWARE. ################################################################################ OUTPUT_FILE="trace.rpd" if [ "$1" = "-o" ] ; then OUTPUT_FILE=$2 shift shift fi if [ -e ${OUTPUT_FILE} ] ; then rm ${OUTPUT_FILE} fi python3 -m rocpd.schema --create ${OUTPUT_FILE} if [ $? != 0 ] ; then echo "Error: Could not create rpd file. Please run 'python setup.py install' from the rocpd_python dir" exit fi export RPDT_FILENAME=${OUTPUT_FILE} export RPDT_AUTOSTART=0 LD_PRELOAD=librocm-smi_64:librpd_tracer.so "$@" ================================================ FILE: 3rdparty/amd/profiling/rpd.patch ================================================ diff --git a/rpd_tracer/Makefile b/rpd_tracer/Makefile index e9d9feb..b2e9e1a 100644 --- a/rpd_tracer/Makefile +++ b/rpd_tracer/Makefile @@ -16,7 +16,7 @@ ifneq (,$(HIP_PATH)) $(info Building with roctracer) RPD_LIBS += -L/opt/rocm/lib -lroctracer64 -lroctx64 -lamdhip64 -lrocm_smi64 RPD_INCLUDES += -I/opt/rocm/include -I/opt/rocm/include/roctracer -I/opt/rocm/include/hsa - RPD_SRCS += RoctracerDataSource.cpp RocmSmiDataSource.cpp + RPD_SRCS += RoctracerDataSource.cpp RPD_INCLUDES += -D__HIP_PLATFORM_AMD__ endif ================================================ FILE: 3rdparty/amd/profiling/rpd_profile_server_enable.patch ================================================ diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 62d1ff9..9021c01 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -71,6 +71,8 @@ from sglang.srt.utils import ( suppress_other_loggers, ) from sglang.utils import get_exception_traceback +from rpdTracerControl import rpdTracerControl +rpdTracerControl.skipCreate() logger = logging.getLogger(__name__) @@ -245,6 +247,7 @@ class Scheduler: ], with_stack=True, ) + self.rpd = rpdTracerControl() @torch.inference_mode() def event_loop(self): @@ -1027,15 +1030,24 @@ class Scheduler: def start_profile(self) -> None: if self.profiler is None: raise RuntimeError("Profiler is not enabled.") - self.profiler.start() + #self.profiler.start() #block pytorch profiler for rpd profiler enabling + if self.tp_rank == 0 or self.tp_rank == 1: + self.rpd.start() + self.rpd.rangePush("", "rpd profile range", "") + logger.info("rpd is enabled") def stop_profile(self) -> None: if self.profiler is None: raise RuntimeError("Profiler is not enabled.") - self.profiler.stop() - self.profiler.export_chrome_trace( - self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" - ) + #self.profiler.stop() + #self.profiler.export_chrome_trace( + # self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" + #) + if self.tp_rank ==0 or self.tp_rank ==1: + self.rpd.rangePop() + self.rpd.stop() + self.rpd.flush() + logger.info("rpd is done") logger.info("Profiler is done") ================================================ FILE: 3rdparty/amd/profiling/rpd_profile_server_enable_wCPU_activities.patch ================================================ diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 62d1ff9..2edb427 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -71,6 +71,8 @@ from sglang.srt.utils import ( suppress_other_loggers, ) from sglang.utils import get_exception_traceback +from rpdTracerControl import rpdTracerControl +rpdTracerControl.skipCreate() logger = logging.getLogger(__name__) @@ -245,6 +247,7 @@ class Scheduler: ], with_stack=True, ) + self.rpd = rpdTracerControl() @torch.inference_mode() def event_loop(self): @@ -1027,15 +1030,26 @@ class Scheduler: def start_profile(self) -> None: if self.profiler is None: raise RuntimeError("Profiler is not enabled.") - self.profiler.start() + #self.profiler.start() + logger.info("torch profiler is disabled") + if self.tp_rank == 0 or self.tp_rank == 1: + self.rpd.setPythonTrace(True) + self.rpd.start() + self.rpd.rangePush("", "scheduler", "") + logger.info("rpd is enabled inside scheduler profiling") def stop_profile(self) -> None: if self.profiler is None: raise RuntimeError("Profiler is not enabled.") - self.profiler.stop() - self.profiler.export_chrome_trace( - self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" - ) + #self.profiler.stop() + #self.profiler.export_chrome_trace( + # self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" + #) + if self.tp_rank ==0 or self.tp_rank ==1: + self.rpd.rangePop() + self.rpd.stop() + self.rpd.flush() + logger.info("rpd is done inside scheduler") logger.info("Profiler is done") diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 2621ccd..181df85 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -58,6 +58,10 @@ from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.utils import is_generation_model, is_multimodal_model +from rpdTracerControl import rpdTracerControl +rpdTracerControl.skipCreate() + + asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) logger = logging.getLogger(__name__) @@ -514,10 +518,20 @@ class TokenizerManager: self.send_to_scheduler.send_pyobj(req) def start_profile(self): + rpd = rpdTracerControl() + rpd.setPythonTrace(True) + rpd.start() + rpd.rangePush("", "tokenizer_manager", "") + logger.info("tokenizer_manager rpd profiling started!") req = ProfileReq.START_PROFILE self.send_to_scheduler.send_pyobj(req) def stop_profile(self): + rpd = rpdTracerControl() + rpd.rangePop() + rpd.stop() + rpd.flush() + logger.info("rpd profiling is done inside tokenizer_manager!") req = ProfileReq.STOP_PROFILE self.send_to_scheduler.send_pyobj(req) diff --git a/python/sglang/srt/server.py b/python/sglang/srt/server.py index 7111c93..2bd722c 100644 --- a/python/sglang/srt/server.py +++ b/python/sglang/srt/server.py @@ -30,6 +30,8 @@ import threading import time from http import HTTPStatus from typing import Dict, List, Optional, Union +from rpdTracerControl import rpdTracerControl +rpdTracerControl.skipCreate() # Fix a bug of Python threading setattr(threading, "_register_atexit", lambda *args, **kwargs: None) @@ -152,6 +154,11 @@ async def flush_cache(): @app.post("/start_profile") async def start_profile(): """Start profiling.""" + rpd = rpdTracerControl() + rpd.setPythonTrace(True) + rpd.start() + rpd.rangePush("", "server rpd profile range", "") + logger.info("rpd profiling started in server.py!") tokenizer_manager.start_profile() return Response( content="Start profiling.\n", @@ -164,6 +171,11 @@ async def start_profile(): async def stop_profile(): """Stop profiling.""" tokenizer_manager.stop_profile() + rpd = rpdTracerControl() + rpd.rangePop() + rpd.stop() + rpd.flush() + logger.info("rpd profiling is done in server.py!") return Response( content="Stop profiling. This will take some time.\n", status_code=200, ================================================ FILE: 3rdparty/amd/profiling/server.sh ================================================ #!/bin/bash # export SGLANG_TORCH_PROFILER_DIR=/data/sglang/ export SGLANG_TORCH_PROFILER_DIR=/sgl-workspace/sglang/profile/ # Get the current timestamp TIMESTAMP=$(date +"%Y%m%d_%H%M%S") # Define the log file with a timestamp LOGFILE="sglang_server_log_$TIMESTAMP.json" # Run the Python command and save the output to the log file loadTracer.sh python3 -m sglang.launch_server \ --model-path /sgl-workspace/sglang/dummy_grok1 \ --tokenizer-path Xenova/grok-1-tokenizer \ --load-format dummy \ --quantization fp8 \ --tp 8 \ --port 30000 \ --disable-radix-cache 2>&1 | tee "$LOGFILE" ================================================ FILE: 3rdparty/amd/profiling/torch_profiler.patch ================================================ diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 62d1ff9..6ecd78c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -240,7 +240,6 @@ class Scheduler: ) self.profiler = torch.profiler.profile( activities=[ - torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], with_stack=True, @@ -1033,9 +1032,11 @@ class Scheduler: if self.profiler is None: raise RuntimeError("Profiler is not enabled.") self.profiler.stop() - self.profiler.export_chrome_trace( - self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" - ) + if self.tp_rank == 0: + with open(f"stats_repro_{int(time.time())}.txt", "w") as f: + print(self.profiler.key_averages(group_by_input_shape=True).table(sort_by="cuda_time_total", row_limit=-1), file=f) + print("Profiling stats done.") + logger.info("Profiler is done") ================================================ FILE: 3rdparty/amd/tuning/TUNING.md ================================================ ## Tuning SGLang Infer System with AMD GPUs This AppNote describes the SGLang performance tuning technical, code harness and running steps for systems with AMD Instinct GPUs. Harness code, examples and steps are provided in detail, to facilitate easy reproduce & use to tune performance towards workloads. Three primary runtime areas are covered: ## 1. Triton Kernels To maximize Triton kernel efficiency, several strategies can be employed: ### Key Environment Variables: - **num_stages**: Adjusts the number of pipeline stages to optimize kernel efficiency based on the specific type of operations (e.g., General Matrix Multiplication - GEMM). - **waves_per_eu**: Controls the usage of Vector General Purpose Registers (VGPR) to enhance occupancy, thereby improving latency or throughput. - **BLOCK_M, BLOCK_N, BLOCK_K**: Tunable tile sizes that assist in balancing memory transfer and computational efficiency. - **matrix_instr_nonkdim**: Optimizes the usage of Matrix-Fused Multiply-Add (MFMA) instructions for specific kernel types, such as Flash Attention. - **OPTIMIZE_EPILOGUE**: An environment variable that can be set to `1` to enhance performance by eliminating the `convert_layout` operation in the kernel's epilogue. ```python @triton.autotune(configs=[ triton.Config({'waves_per_eu': 1}, num_warps=4, num_stages=1), triton.Config({'waves_per_eu': 1}, num_warps=8, num_stages=1), triton.Config({'waves_per_eu': 1}, num_warps=16, num_stages=1), triton.Config({'waves_per_eu': 2}, num_warps=4, num_stages=1), triton.Config({'waves_per_eu': 2}, num_warps=8, num_stages=1), triton.Config({'waves_per_eu': 2}, num_warps=16, num_stages=1), triton.Config({'waves_per_eu': 4}, num_warps=4, num_stages=1), triton.Config({'waves_per_eu': 4}, num_warps=8, num_stages=1), triton.Config({'waves_per_eu': 4}, num_warps=16, num_stages=1), ], key=['BLOCK_N', 'NUM_TOKEN_BLKS'], use_cuda_graph=True) @triton.jit def _triton_kernel_function(): ... ``` ## 2. Torch Tunable Operations **TunableOp** is a feature in PyTorch that allows for the definition and optimization of custom kernels with tunable parameters. This feature is particularly useful for enhancing the performance of kernels by experimenting with different configurations. ### Key Environment Variables: 1. **PYTORCH_TUNABLEOP_ENABLED**: - Default: `0` - Set to `1` to enable TunableOp. 2. **PYTORCH_TUNABLEOP_TUNING**: - Default: `1` - Set to `0` to disable tuning. If a tuned entry is not found, it will run the tuning step and record the entry when PYTORCH_TUNABLEOP_ENABLED is enabled. 3. **PYTORCH_TUNABLEOP_VERBOSE**: - Default: `0` - Set to `1` to enable verbose output for TunableOp. ### Usage Example: To enable TunableOp and tuning, and optionally enable verbose mode, you can run the following command in your terminal: ```bash #Tuning PYTORCH_TUNABLEOP_ENABLED=1 PYTORCH_TUNABLEOP_TUNING=1 your_script.sh #Inference with tuning op PYTORCH_TUNABLEOP_ENABLED=1 PYTORCH_TUNABLEOP_TUNING=0 your_script.sh #Print out the log PYTORCH_TUNABLEOP_ENABLED=1 PYTORCH_TUNABLEOP_TUNING=0 PYTORCH_TUNABLEOP_VERBOSE=1 your_script.sh ``` ## 3. Torch Compilation The following are suggestions for optimizing matrix multiplication (GEMM) and convolution (conv) operations in PyTorch using Inductor, a part of the PyTorch compilation framework. The goal is to leverage Triton to achieve better performance. To tune Triton kernels with GEMM and convolution ops (conv), use the `torch.compile` function with the max-autotune mode. This benchmarks a predefined list of Triton configurations and selects the fastest one for each shape. ### Key Configurations: 1. **Max Autotune**: - Set `torch._inductor.config.max_autotune = True` or `TORCHINDUCTOR_MAX_AUTOTUNE=1`. 2. **Fine-Grained Control**: - Enable GEMM tuning: `torch._inductor.config.max_autotune_gemm = True`. - Enable tuning for pointwise/reduction ops: `torch._inductor.config.max_autotune.pointwise = True`. 3. **Backend Selection**: - Use `torch._inductor.max_autotune_gemm_backends` to limit backends to TRITON for better performance. 4. **Freezing for Inference**: - Use `torch._inductor.config.freezing=True` to enable constant folding optimizations. 5. **Debugging**: - Set `TORCH_COMPILE_DEBUG=1` to extract Triton kernels generated by Inductor. ### Example Code Block: ```bash #Gemm Tuning TORCHINDUCTOR_MAX_AUTOTUNE=1 TORCHINDUCTOR_COORDINATE_DESCENT_TUNING=1 your_script.sh #Specify your backend to TRITON for Gemm Tuning TORCHINDUCTOR_MAX_AUTOTUNE=1 TORCHINDUCTOR_COORDINATE_DESCENT_TUNING=1 TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS=TRITON your_script.sh #Inference with large improvement on AMD GPU TORCHINDUCTOR_FREEZING=1 your_script.sh ``` ## 4. Fused MOE kernel To maximize moe kernel efficiency, need to use below scripts to find out the best launch configuration ### Key parameters: - **--model**: what moe model type to do tuning, it will automatically decide the size of d_model, model_intermediate_size, num_layers - **--tp-size**: simulate the whole model run configuration to set the dimension size using tp correctly - **--batch**: M dimension size of moe kernel, for prefill moe kernel the value is batch*input_len, for decode moe kernel the value is batch - **--dtype**: computation type ```bash #Tuning #for example, we have one case like this "python3 -m sglang.bench_latency --model dummy_grok1/ --load-format dummy --tokenizer-path Xenova/grok-1-tokenizer --tp 8 --batch-size 32 --input 1024 --output 8 --attention-backend triton --sampling-backend pytorch --quantization fp8" to run, it defined batch-size 32 input length 1024 and output length 8, from "--batch" in moe view point, the prefill batch is 32*1024 = 32768, the decode batch is 32*1(only one output token generated in each run). #so we can tune decode moe use below command python benchmark_moe_rocm.py --model grok1 --tp-size 8 --dtype float8 --batch "32" # and use this command to tune prefill moe python benchmark_moe_rocm.py --model grok1 --tp-size 8 --dtype float8 --batch "32768" ``` ## Reference For more detailed information on tuning SGLang performance with AMD GPUs, please refer to the following link: [ROCm Documentation: Triton Kernel Performance Optimization](https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html#triton-kernel-performance-optimization) ================================================ FILE: 3rdparty/amd/tuning/benchmark_moe_rocm.py ================================================ import argparse import json import os import sys import torch import torch.nn.functional as F import triton import triton.language as tl from tqdm import tqdm from transformers import AutoConfig from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( fused_moe, get_config_file_name, ) padding_size = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0 def main(model, tp_size, dtype: str, batches): method = fused_moe for bs in batches: run_grid(int(bs), model=model, method=method, tp_size=tp_size, dtype=dtype) def prune_configs(M, N, K, configs): pruned_configs = [] elemBytes_a = 1 # [DV Note] Hard-coded for float16 (2 bytes) elemBytes_b = 1 # [DV Note] Hard-coded for float16 (2 bytes) mfma = 16 if M < 32 or N < 32 else 32 # TODO (zhanglx): figure out the boundary between large and small gemms large_gemm = False if M >= 2048 and N >= 2048: large_gemm = True for config in configs: BLOCK_SIZE_M = config.get("BLOCK_SIZE_M") BLOCK_SIZE_N = config.get("BLOCK_SIZE_N") BLOCK_SIZE_K = config.get("BLOCK_SIZE_K") num_warps = config.get("num_warps") matrix_instr_nonkdim = config.get("matrix_instr_nonkdim") # kpack = config.get("kpack") if matrix_instr_nonkdim > mfma: continue if mfma == 4 and BLOCK_SIZE_K < 64: continue # some layouts could not work properly in case # number elements per thread is less 1 if BLOCK_SIZE_M * BLOCK_SIZE_N < 64: continue SPLIT_K = 1 # config.get("SPLIT_K") GROUP_M = config.get("GROUP_SIZE_M") if matrix_instr_nonkdim > BLOCK_SIZE_M or matrix_instr_nonkdim > BLOCK_SIZE_N: continue if matrix_instr_nonkdim >= M and matrix_instr_nonkdim != BLOCK_SIZE_M: continue if matrix_instr_nonkdim >= N and matrix_instr_nonkdim != BLOCK_SIZE_N: continue # Skip BLOCK_SIZE that is too large compare to M/N # unless BLOCK_SIZE is already small enough if M * 2 < BLOCK_SIZE_M and BLOCK_SIZE_M != 16: continue if N * 2 < BLOCK_SIZE_N and BLOCK_SIZE_N != 16: continue # skip large split_k when not necessary if SPLIT_K != 1 and not need_split_k(M, N, K): continue # skip split_k that leads to EVEN_K = false leap = SPLIT_K * BLOCK_SIZE_K modv = K % leap if modv != 0: continue # skip large GROUP_M if GROUP_M * BLOCK_SIZE_M > M and GROUP_M != 1: continue # out of shared memory resource # TODO (zhanglx): This does not consider the LDS usage in the epilogue LDS = ( BLOCK_SIZE_K * BLOCK_SIZE_M * elemBytes_a + BLOCK_SIZE_K * BLOCK_SIZE_N * elemBytes_b ) if LDS > 65536: continue # Skip small block sizes and num_warps for large gemm # For fp16 and f8, we want to only use BLOCK_SIZE >= 64 if large_gemm: if BLOCK_SIZE_M < 64 or BLOCK_SIZE_N < 64: continue if BLOCK_SIZE_K < 64: continue if num_warps < 4: continue pruned_configs.append(config) return pruned_configs def union_of_list_of_dicts(l1, l2): result = [] temp_list = l1.copy() temp_list.extend(l2) for myDict in temp_list: if myDict not in result: result.append(myDict) return result def run_grid(bs, model, method, tp_size, dtype: str): config = AutoConfig.from_pretrained(model) top_k = config.num_experts_per_tok d_model = config.hidden_size model_intermediate_size = config.intermediate_size num_layers = config.num_hidden_layers hidden_states_dtype = config.torch_dtype if config.num_experts_per_tok: if config.architectures[0] == "Grok1ModelForCausalLM": num_total_experts = config.num_experts else: num_total_experts = config.num_local_experts else: raise ValueError(f"Unsupported Mixtral model {model}") # tp_size = 2 num_warmup_calls = 10 num_calls = 30 num_warmup_trials = 1 num_trials = 1 full_configs = [] block_m_range = [16, 32, 64, 128, 256] block_n_range = [16, 32, 64, 128, 256] block_k_range = [32, 64, 128, 256] # MUST >= 32 num_warps_range = [1, 2, 4, 8] group_m_range = [1, 4, 8, 16, 32] # For now we see better perf with num_stages=0 for all gemm configs we care # But keep this explicit so that we do not forget we may need to set it to # other values in the future num_stage_range = [2] waves_per_eu_range = [0, 1, 2, 4, 8] # Remove 32 because of triton compiling error matrix_instr_nonkdim_range = [16] kpack_range = [1, 2] for block_size_m in block_m_range: for block_size_n in block_n_range: for block_size_k in block_k_range: for group_size_m in group_m_range: for num_warps in num_warps_range: for num_stages in num_stage_range: for waves_per_eu in waves_per_eu_range: for matrix_instr_nonkdim in matrix_instr_nonkdim_range: for kpack in kpack_range: full_configs.append( { "BLOCK_SIZE_M": block_size_m, "BLOCK_SIZE_N": block_size_n, "BLOCK_SIZE_K": block_size_k, "GROUP_SIZE_M": group_size_m, "num_warps": num_warps, "num_stages": num_stages, "waves_per_eu": waves_per_eu, "matrix_instr_nonkdim": matrix_instr_nonkdim, "kpack": kpack, } ) M1 = bs * 2 N1 = model_intermediate_size * 2 // tp_size K1 = d_model prune_configs_1 = prune_configs(M1, N1, K1, full_configs) M2 = bs * 2 N2 = d_model K2 = model_intermediate_size // tp_size prune_configs_2 = prune_configs(M2, N2, K2, full_configs) configs = union_of_list_of_dicts(prune_configs_1, prune_configs_2) print(f"{bs=} || {len(full_configs)=} | {len(prune_configs_1)=} | \ {len(prune_configs_2)=} | {len(configs)=}") best_config = None best_time_us = 1e20 print(f"{tp_size=} {bs=}") for config in tqdm(configs): # warmup try: print(config) for _ in range(num_warmup_trials): run_timing( num_calls=num_warmup_calls, bs=bs, d_model=d_model, num_total_experts=num_total_experts, top_k=top_k, tp_size=tp_size, model_intermediate_size=model_intermediate_size, method=method, config=config, dtype=dtype, hidden_states_dtype=hidden_states_dtype, ) except triton.runtime.autotuner.OutOfResources: continue # trial for _ in range(num_trials): kernel_dur_ms = run_timing( num_calls=num_calls, bs=bs, d_model=d_model, num_total_experts=num_total_experts, top_k=top_k, tp_size=tp_size, model_intermediate_size=model_intermediate_size, method=method, config=config, dtype=dtype, hidden_states_dtype=hidden_states_dtype, ) kernel_dur_us = 1000 * kernel_dur_ms model_dur_ms = kernel_dur_ms * num_layers if kernel_dur_us < best_time_us: best_config = config best_time_us = kernel_dur_us tqdm.write( f"{kernel_dur_us=:.1f} {model_dur_ms=:.1f}" f" {bs=} {tp_size=} {top_k=} {num_total_experts=} " f"{d_model=} {model_intermediate_size=} {num_layers=}" ) print("best_time_us", best_time_us) print("best_config", best_config) # holds Dict[str, Dict[str, int]] filename = get_config_file_name( num_total_experts, model_intermediate_size // tp_size, "float8" if dtype == "float8" else None, ) print(f"writing config to file {filename}") existing_content = {} if os.path.exists(filename): with open(filename, "r") as f: existing_content = json.load(f) existing_content[str(bs)] = best_config with open(filename, "w") as f: json.dump(existing_content, f, indent=4) f.write("\n") def run_timing( num_calls: int, bs: int, d_model: int, num_total_experts: int, top_k: int, tp_size: int, model_intermediate_size: int, method, config, dtype: str, hidden_states_dtype, ) -> float: shard_intermediate_size = model_intermediate_size // tp_size hidden_states = torch.rand( (bs, d_model), device="cuda:0", dtype=hidden_states_dtype, ) w1 = torch.rand( (num_total_experts, 2 * shard_intermediate_size, d_model + padding_size), device=hidden_states.device, dtype=hidden_states.dtype, ) w2 = torch.rand( (num_total_experts, d_model, shard_intermediate_size + padding_size), device=hidden_states.device, dtype=hidden_states.dtype, ) w1_scale = None w2_scale = None a1_scale = None a2_scale = None if dtype == "float8": w1 = w1.to(torch.float8_e4m3fnuz) w2 = w2.to(torch.float8_e4m3fnuz) w1_scale = torch.ones( num_total_experts, device=hidden_states.device, dtype=torch.float32 ) w2_scale = torch.ones( num_total_experts, device=hidden_states.device, dtype=torch.float32 ) a1_scale = torch.ones(1, device=hidden_states.device, dtype=torch.float32) a2_scale = torch.ones(1, device=hidden_states.device, dtype=torch.float32) gating_output = F.softmax( torch.rand( (num_calls, bs, num_total_experts), device=hidden_states.device, dtype=torch.float32, ), dim=-1, ) ################################## start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) start_event.record() for i in range(num_calls): hidden_states = method( hidden_states=hidden_states, w1=w1, w2=w2, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, gating_output=gating_output[0], topk=top_k, renormalize=True, inplace=True, override_config=config, use_fp8=dtype == "float8", ) end_event.record() end_event.synchronize() dur_ms = start_event.elapsed_time(end_event) / num_calls return dur_ms if __name__ == "__main__": parser = argparse.ArgumentParser( prog="benchmark_mixtral_moe", description="Benchmark and tune the fused_moe kernel", ) parser.add_argument( "--dtype", type=str, default="auto", choices=["float8", "float16", "bfloat16"], help="Data type used for fused_moe kernel computations", ) parser.add_argument("--model", type=str, default="hpcai-tech/grok-1") parser.add_argument("--tp-size", type=int, default=2, help="Tensor paralleli size") parser.add_argument("-b", "--batches", type=str) args = parser.parse_args() batches = args.batches.split(",") sys.exit(main(args.model, args.tp_size, args.dtype, batches)) ================================================ FILE: 3rdparty/amd/wheel/README.md ================================================ # sglang-kernel (prior sgl-kernel) Building and releasing `sglang-kernel` as a wheel is a part of the release workflow. Check [release-whl-kernel.yml](https://github.com/sgl-project/sglang/blob/main/.github/workflows/release-whl-kernel.yml) for details. # sglang `3rdparty/amd/wheel/sglang/pyproject.toml` is the AMD-specific pyproject for building the `amd-sglang` wheel. It extends `python/pyproject_other.toml` with two ROCm-version extras (`rocm700`, `rocm720`) that pin the matching torch/triton/torchaudio/torchvision/`sglang-kernel` wheels, and renames the package to `amd-sglang`. ## Operation to build sglang wheel ``` $ git clone https://github.com/sgl-project/sglang.git && cd sglang $ cp 3rdparty/amd/wheel/sglang/pyproject.toml python/pyproject.toml $ cd python && python -m build ``` ## Installation ### v0.5.9 ROCm 7.0.0: ``` pip uninstall sglang-kernel sglang amd-sglang pip install "amd-sglang[all-hip,rocm700]" -i https://pypi.amd.com/rocm-7.0.0/simple --extra-index-url https://pypi.org/simple ``` ROCm 7.2.0: ``` pip uninstall sglang-kernel sglang amd-sglang pip install "amd-sglang[all-hip,rocm720]" -i https://pypi.amd.com/rocm-7.2.0/simple --extra-index-url https://pypi.org/simple ``` Note: You must resolve the two dependencies, AITER and triton, below. Others are optional depending on your applications. ## Manual Dependency Resolution ### Resolving AITER [AITER](https://github.com/ROCm/aiter) is a fundamental dependency. Wheel-izing it is ongoing. Until we can pin it reliably, install it manually (typically following the [ROCm docker recipe](https://github.com/sgl-project/sglang/blob/main/docker/rocm.Dockerfile#L106). ### Revolving triton To avoid known issues in triton 3.5.1 installed by default, we recommend upgrading triton after installation. In ROCm 7.0.0 environment, ``` pip install triton==3.6.0 ``` or ROCm 7.2.0, ``` pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/triton-3.6.0%2Brocm7.2.0.gitba5c1517-cp310-cp310-linux_x86_64.whl ``` #### `torch._inductor.exc.InductorError: AttributeError: 'KernelMetadata' object has no attribute 'cluster_dims'` After upgrading, you may hit this error during inference when PyTorch Inductor interacts with Triton metadata. A pragmatic workaround is to guard the metadata access in Inductor's Triton heuristics so it only reads `cluster_dims` when the attribute exists: ```diff --- a/opt/venv/lib/python3.10/site-packages/torch/_inductor/runtime/triton_heuristics.py +++ b/opt/venv/lib/python3.10/site-packages/torch/_inductor/runtime/triton_heuristics.py @@ -1759,6 +1759,8 @@ else ( (binary.metadata.num_ctas, *binary.metadata.cluster_dims) if hasattr(binary, "metadata") + and hasattr(binary.metadata, "num_ctas") + and hasattr(binary.metadata, "cluster_dims") else () ) ), ``` ### Resolving Dependencies for Distributed Inference #### sgl-model-gateway Install sgl-model-gateway as follows: ``` $ apt install openssl libssl-dev protobuf $ export PATH="/$HOME/.cargo/bin:${PATH}" \ && curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ && rustc --version && cargo --version # Prepare for a rust toolchain $ python3 -m pip install --no-cache-dir setuptools-rust \ && cd /sgl-workspace/sglang/sgl-model-gateway/bindings/python \ && cargo build --release \ && python3 -m pip install --no-cache-dir . \ && rm -rf /root/.cache # Build and install sgl-model-gateway ``` #### [Mori](https://github.com/sgl-project/sglang/blob/main/docker/rocm.Dockerfile#L381) ### Resolving Dependencies for DeepSeek-V3.2 #### [TileLang](https://github.com/sgl-project/sglang/blob/main/docker/rocm.Dockerfile#L216) #### [FHT (fast-hadamard-transform)](https://github.com/sgl-project/sglang/blob/main/docker/rocm.Dockerfile#L300) ================================================ FILE: 3rdparty/amd/wheel/sgl-kernel/CMakeLists_rocm.txt ================================================ cmake_minimum_required(VERSION 3.24 FATAL_ERROR) project(sgl_kernel LANGUAGES CXX) # Cmake set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_POSITION_INDEPENDENT_CODE ON) set(CMAKE_SHARED_LIBRARY_PREFIX "") set(CMAKE_COLOR_DIAGNOSTICS ON) set(CMAKE_VERBOSE_MAKEFILE ON CACHE BOOL "ON") # Python / Torch find_package(Python COMPONENTS Interpreter Development.Module ${SKBUILD_SABI_COMPONENT} REQUIRED) execute_process( COMMAND ${Python_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)" OUTPUT_VARIABLE TORCH_PY_PREFIX OUTPUT_STRIP_TRAILING_WHITESPACE ) set(Torch_DIR "${TORCH_PY_PREFIX}/Torch") list(APPEND CMAKE_PREFIX_PATH "${TORCH_PY_PREFIX}/Torch") find_package(Torch REQUIRED) execute_process( COMMAND ${Python_EXECUTABLE} -c "import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))" OUTPUT_VARIABLE TORCH_CXX11_ABI OUTPUT_STRIP_TRAILING_WHITESPACE ) if(TORCH_CXX11_ABI STREQUAL "0") add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=0) else() add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=1) endif() # ROCm/HIP enable_language(HIP) list(APPEND CMAKE_PREFIX_PATH "/opt/rocm/lib/cmake/hip-lang") find_package(hip REQUIRED CONFIG) # Determine AMDGPU target from environment variable or default to gfx942 set(AMDGPU_TARGET_ENV "$ENV{AMDGPU_TARGET}") if(AMDGPU_TARGET_ENV) # Use environment variable if specified set(AMDGPU_TARGETS "${AMDGPU_TARGET_ENV}") message(STATUS "Using AMDGPU_TARGET from environment: ${AMDGPU_TARGETS}") else() # Default to gfx942 only set(AMDGPU_TARGETS "gfx942") message(STATUS "AMDGPU_TARGET not set, defaulting to gfx942") endif() # Set HIP architectures set(CMAKE_HIP_ARCHITECTURES ${AMDGPU_TARGETS}) # FP8 macro selection # Always define HIP_FP8_TYPE_FNUZ=1 (for gfx942 and host compilation) # Additionally define HIP_FP8_TYPE_E4M3=1 when building for gfx950 # The existing utils.h logic will pick the right one based on architecture set(SGL_FP8_MACROS "-DHIP_FP8_TYPE_FNUZ=1") if(AMDGPU_TARGETS MATCHES "gfx950") list(APPEND SGL_FP8_MACROS "-DHIP_FP8_TYPE_E4M3=1") message(STATUS "Multi-arch build: Enabling both HIP_FP8_TYPE_FNUZ (gfx942) and HIP_FP8_TYPE_E4M3 (gfx950)") elseif(AMDGPU_TARGETS MATCHES "gfx942") message(STATUS "Single-arch build: Enabling HIP_FP8_TYPE_FNUZ for gfx942") else() message(FATAL_ERROR "Unsupported AMDGPU_TARGET '${AMDGPU_TARGETS}'. Expected 'gfx942' or 'gfx950' or both.") endif() # TopK dynamic smem bytes # Dynamic shared-memory budget for the TopK kernels. # - gfx942 (MI300/MI325): LDS is typically 64KB per workgroup -> keep dynamic smem <= ~48KB # (leaves room for static shared allocations in the kernel). # - gfx95x (MI350): LDS is larger (e.g. 160KB per CU) -> allow the original 128KB dynamic smem. if(AMDGPU_TARGET_ONE STREQUAL "gfx942") math(EXPR SGL_TOPK_DYNAMIC_SMEM_BYTES "48 * 1024") else() math(EXPR SGL_TOPK_DYNAMIC_SMEM_BYTES "32 * 1024 * 4") endif() set(SGL_TOPK_MACROS "-DSGL_TOPK_DYNAMIC_SMEM_BYTES=${SGL_TOPK_DYNAMIC_SMEM_BYTES}") # Paths / includes set(PROJ_ROOT ${CMAKE_CURRENT_LIST_DIR}) set(SGL_INCLUDE_DIRS ${PROJ_ROOT}/include ${PROJ_ROOT}/include/impl ${PROJ_ROOT}/csrc ${TORCH_INCLUDE_DIRS} ) # Platform-specific library directory set(PLAT_LIB_DIR "/usr/lib/x86_64-linux-gnu") link_directories(${PLAT_LIB_DIR}) # Sources set(SOURCES ${PROJ_ROOT}/csrc/allreduce/custom_all_reduce.hip ${PROJ_ROOT}/csrc/allreduce/deterministic_all_reduce.hip ${PROJ_ROOT}/csrc/allreduce/quick_all_reduce.hip ${PROJ_ROOT}/csrc/common_extension_rocm.cc ${PROJ_ROOT}/csrc/elementwise/activation.hip ${PROJ_ROOT}/csrc/elementwise/pos_enc.hip ${PROJ_ROOT}/csrc/elementwise/topk.hip ${PROJ_ROOT}/csrc/grammar/apply_token_bitmask_inplace_hip.hip ${PROJ_ROOT}/csrc/kvcacheio/transfer.hip ${PROJ_ROOT}/csrc/memory/weak_ref_tensor.cpp ${PROJ_ROOT}/csrc/moe/moe_align_kernel.hip ${PROJ_ROOT}/csrc/moe/moe_topk_softmax_kernels.hip ${PROJ_ROOT}/csrc/moe/moe_topk_sigmoid_kernels.hip ${PROJ_ROOT}/csrc/speculative/eagle_utils.hip ) set_source_files_properties( ${SOURCES} PROPERTIES LANGUAGE HIP ) # Compile / Link flags add_compile_options($<$:-O3>) set(SGL_HIP_FLAGS -DNDEBUG -DOPERATOR_NAMESPACE=sgl_kernel -O3 -std=c++17 -DENABLE_BF16 -DENABLE_FP8 ${SGL_FP8_MACROS} -Wno-pass-failed -Wundefined-internal ${SGL_TOPK_MACROS} ) # Python extension Python_add_library(common_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${SOURCES}) target_include_directories(common_ops PRIVATE ${SGL_INCLUDE_DIRS}) # Apply per-language flags target_compile_options(common_ops PRIVATE $<$:${SGL_HIP_FLAGS}> ) target_link_libraries(common_ops PRIVATE ${TORCH_LIBRARIES} hip::device hip::host hiprtc amdhip64 ) target_link_options(common_ops PRIVATE "SHELL:-Wl,-rpath,'\$ORIGIN/../../torch/lib'" ) install(TARGETS common_ops LIBRARY DESTINATION sgl_kernel ) ================================================ FILE: 3rdparty/amd/wheel/sgl-kernel/build_rocm.sh ================================================ #!/bin/bash set -euo pipefail ROCM_VERSION=${1:-} if [[ "${ROCM_VERSION}" == "700" ]]; then IMAGE="lmsysorg/sglang:v0.5.8.post1-rocm700-mi35x" elif [[ "${ROCM_VERSION}" == "720" ]]; then IMAGE="rocm/pytorch:rocm7.2_ubuntu22.04_py3.10_pytorch_release_2.9.1" else echo "ERROR: Unsupported ROCM_VERSION='${ROCM_VERSION}'. Only '700' and '720' are supported." >&2 exit 1 fi PYTHON_ROOT_PATH="/opt/venv/bin" AMDGPU_TARGET="gfx942;gfx950" # Pull and run the latest image echo "Pulling Docker image: ${IMAGE}" docker pull "${IMAGE}" docker run --rm \ -v $(pwd):/sgl-kernel \ -e AMDGPU_TARGET="${AMDGPU_TARGET}" \ -e PYTORCH_ROCM_ARCH="${AMDGPU_TARGET}" \ ${IMAGE} \ bash -c " # Install torch, triton, and friends, depending on the ROCm version if [[ "${ROCM_VERSION}" == "700" ]]; then ${PYTHON_ROOT_PATH}/pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torch-2.9.1.dev20251204%2Brocm7.0.2.lw.git351ff442-cp310-cp310-linux_x86_64.whl https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/triton-3.5.1%2Brocm7.0.2.gita272dfa8-cp310-cp310-linux_x86_64.whl https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torchaudio-2.9.0%2Brocm7.0.2.gite3c6ee2b-cp310-cp310-linux_x86_64.whl https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torchvision-0.24.0%2Brocm7.0.2.gitb919bd0c-cp310-cp310-linux_x86_64.whl elif [[ "${ROCM_VERSION}" == "720" ]]; then ${PYTHON_ROOT_PATH}/pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torch-2.9.1%2Brocm7.2.0.lw.git7e1940d4-cp310-cp310-linux_x86_64.whl https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/triton-3.5.1%2Brocm7.2.0.gita272dfa8-cp310-cp310-linux_x86_64.whl https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchaudio-2.9.0%2Brocm7.2.0.gite3c6ee2b-cp310-cp310-linux_x86_64.whl https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.24.0%2Brocm7.2.0.gitb919bd0c-cp310-cp310-linux_x86_64.whl fi # Install CMake (version >= 3.26) - Robust Installation export CMAKE_VERSION_MAJOR=3.31 export CMAKE_VERSION_MINOR=1 echo \"Downloading CMake from: https://cmake.org/files/v\${CMAKE_VERSION_MAJOR}/cmake-\${CMAKE_VERSION_MAJOR}.\${CMAKE_VERSION_MINOR}-linux-x86_64.tar.gz\" wget https://cmake.org/files/v\${CMAKE_VERSION_MAJOR}/cmake-\${CMAKE_VERSION_MAJOR}.\${CMAKE_VERSION_MINOR}-linux-x86_64.tar.gz tar -xzf cmake-\${CMAKE_VERSION_MAJOR}.\${CMAKE_VERSION_MINOR}-linux-x86_64.tar.gz mv cmake-\${CMAKE_VERSION_MAJOR}.\${CMAKE_VERSION_MINOR}-linux-x86_64 /opt/cmake export PATH=/opt/cmake/bin:\$PATH ${PYTHON_ROOT_PATH}/pip install --no-cache-dir ninja setuptools wheel numpy uv scikit-build-core && \ cd /sgl-kernel && \ rm -rf CMakeLists.txt && mv CMakeLists_rocm.txt CMakeLists.txt && \ ${PYTHON_ROOT_PATH}/python rocm_hipify.py && \ ${PYTHON_ROOT_PATH}/python -m uv build --wheel -Cbuild-dir=build . --color=always --no-build-isolation && \ ./rename_wheels_rocm.sh " ================================================ FILE: 3rdparty/amd/wheel/sgl-kernel/rename_wheels_rocm.sh ================================================ #!/usr/bin/env bash set -ex WHEEL_DIR="dist" wheel_files=($WHEEL_DIR/*.whl) for wheel in "${wheel_files[@]}"; do intermediate_wheel="${wheel/linux/manylinux2014}" [[ "$intermediate_wheel" == *"+rocm"* ]] && continue # Extract the current python version from the wheel name if [[ $intermediate_wheel =~ -cp([0-9]+)- ]]; then cp_version="${BASH_REMATCH[1]}" else echo "Could not extract Python version from wheel name: $intermediate_wheel" continue fi # Detect ROCm version and add appropriate suffix ver_abrv=$(realpath /opt/rocm-* | sed -e 's/.*-//' -e 's/\.//g') new_wheel=${intermediate_wheel/-cp${cp_version}/+rocm${ver_abrv}-cp${cp_version}} if [[ "$wheel" != "$new_wheel" ]]; then echo "Renaming $wheel to $new_wheel" mv -- "$wheel" "$new_wheel" fi done echo "Wheel renaming completed." ================================================ FILE: 3rdparty/amd/wheel/sgl-kernel/rocm_hipify.py ================================================ from pathlib import Path import torch from torch.utils.cpp_extension import CUDAExtension root = Path(__file__).parent.resolve() include_dirs = [ root / "include", root / "include" / "impl", root / "csrc", ] sources = [ "csrc/allreduce/custom_all_reduce.hip", "csrc/allreduce/deterministic_all_reduce.hip", "csrc/allreduce/quick_all_reduce.cu", "csrc/common_extension_rocm.cc", "csrc/elementwise/activation.cu", "csrc/elementwise/pos_enc.cu", "csrc/elementwise/topk.cu", "csrc/grammar/apply_token_bitmask_inplace_cuda.cu", "csrc/kvcacheio/transfer.cu", "csrc/memory/weak_ref_tensor.cpp", "csrc/moe/moe_align_kernel.cu", "csrc/moe/moe_topk_softmax_kernels.cu", "csrc/moe/moe_topk_sigmoid_kernels.cu", "csrc/speculative/eagle_utils.cu", ] libraries = ["hiprtc", "amdhip64", "c10", "torch", "torch_python"] ext_modules = [ CUDAExtension( name="sgl_kernel.common_ops", sources=sources, include_dirs=include_dirs, libraries=libraries, py_limited_api=False, ), ] ================================================ FILE: 3rdparty/amd/wheel/sglang/pyproject.toml ================================================ [build-system] requires = ["setuptools>=61.0", "setuptools-scm>=8.0", "wheel"] build-backend = "setuptools.build_meta" [project] name = "amd-sglang" dynamic = ["version"] description = "SGLang is a fast serving framework for large language models and vision language models." readme = "README.md" requires-python = ">=3.10" license = { file = "LICENSE" } classifiers = [ "Programming Language :: Python :: 3", "License :: OSI Approved :: Apache Software License", ] dependencies = ["aiohttp", "requests", "tqdm", "numpy", "IPython", "setproctitle"] [project.optional-dependencies] runtime_common = [ "IPython", "aiohttp", "anthropic>=0.20.0", "blobfile==3.0.0", "build", "compressed-tensors", "decord2", "datasets", "einops", "fastapi", "gguf", "hf_transfer", "huggingface_hub", "interegular", "llguidance>=0.7.11,<0.8.0", "modelscope", "msgspec", "ninja", "numpy", "openai-harmony==0.0.4", "openai==2.6.1", "orjson", "outlines==0.1.11", "packaging", "partial_json_parser", "pillow", "prometheus-client>=0.20.0", "psutil", "py-spy", "pybase64", "pydantic", "python-multipart", "pyzmq>=25.1.2", "requests", "scipy", "sentencepiece", "setproctitle", "soundfile==0.13.1", "tiktoken", "timm==1.0.16", "torchao==0.9.0", "tqdm", "transformers==4.57.1", "uvicorn", "uvloop", "xgrammar==0.1.27", "smg-grpc-servicer>=0.5.0", ] # ROCm specific packages (https://repo.radeon.com/rocm/manylinux/) # Existing practice for daily rocm700 docker images relies on 700-rc # versions of software that are not public available. Here we pin some # from rocm702 as the closest set as daily rocm700 images. rocm700 = [ "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torch-2.9.1.dev20251204%2Brocm7.0.2.lw.git351ff442-cp310-cp310-linux_x86_64.whl", "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/triton-3.5.1%2Brocm7.0.2.gita272dfa8-cp310-cp310-linux_x86_64.whl", "torchaudio @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torchaudio-2.9.0%2Brocm7.0.2.gite3c6ee2b-cp310-cp310-linux_x86_64.whl", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torchvision-0.24.0%2Brocm7.0.2.gitb919bd0c-cp310-cp310-linux_x86_64.whl", "mooncake-transfer-engine-non-cuda==0.3.8.post1", "sglang-kernel @ https://github.com/sgl-project/whl/releases/download/v0.4.0/sglang_kernel-0.4.0+rocm700-cp310-abi3-manylinux2014_x86_64.whl", ] rocm720 = [ "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torch-2.9.1%2Brocm7.2.0.lw.git7e1940d4-cp310-cp310-linux_x86_64.whl", "triton @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/triton-3.5.1%2Brocm7.2.0.gita272dfa8-cp310-cp310-linux_x86_64.whl", "torchaudio @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchaudio-2.9.0%2Brocm7.2.0.gite3c6ee2b-cp310-cp310-linux_x86_64.whl", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.24.0%2Brocm7.2.0.gitb919bd0c-cp310-cp310-linux_x86_64.whl", "mooncake-transfer-engine-non-cuda==0.3.8.post1", "sglang-kernel @ https://github.com/sgl-project/whl/releases/download/v0.4.0/sglang_kernel-0.4.0+rocm720-cp310-abi3-manylinux2014_x86_64.whl", ] # HIP (Heterogeneous-computing Interface for Portability) for AMD # Install with one of: # pip install "amd-sglang[srt_hip,rocm700]" # pip install "amd-sglang[srt_hip,rocm720]" srt_hip = [ "amd-sglang[runtime_common]", "petit_kernel==0.0.2", "wave-lang==3.8.2", ] diffusion_hip = [ "PyYAML==6.0.1", "cloudpickle", "diffusers==0.37.0", "imageio==2.36.0", "imageio-ffmpeg==0.5.1", "moviepy>=2.0.0", "opencv-python-headless==4.10.0.84", "remote-pdb", "st_attn==0.0.7", "vsa==0.0.4", "runai_model_streamer>=0.15.5", "cache-dit==1.1.8", "addict", ] # For Intel Gaudi(device : hpu) follow the installation guide # https://docs.vllm.ai/en/latest/getting_started/gaudi-installation.html srt_hpu = ["sglang[runtime_common]"] # https://docs.sglang.io/platforms/mthreads_gpu.md srt_musa = [ "sglang[runtime_common]", "torch", "torch_musa", "torchada>=0.1.25", "mthreads-ml-py", "numpy<2.0", ] diffusion_musa = [ "PyYAML==6.0.1", "cloudpickle", "diffusers==0.37.0", "imageio==2.36.0", "imageio-ffmpeg==0.5.1", "moviepy>=2.0.0", "opencv-python-headless==4.10.0.84", "remote-pdb", "st_attn==0.0.7", "vsa==0.0.4", "runai_model_streamer>=0.15.5", "cache-dit==1.1.8", "addict", ] tracing = [ "opentelemetry-sdk", "opentelemetry-api", "opentelemetry-exporter-otlp", "opentelemetry-exporter-otlp-proto-grpc", ] test = [ "accelerate", "expecttest", "gguf", "jsonlines", "matplotlib", "pandas", "peft", "pytest", "sentence_transformers", "tabulate", ] all_hip = ["amd-sglang[srt_hip]", "amd-sglang[diffusion_hip]"] all_hpu = ["sglang[srt_hpu]"] all_musa = ["sglang[srt_musa]", "sglang[diffusion_musa]"] dev_hip = ["amd-sglang[all_hip]", "amd-sglang[test]"] dev_hpu = ["sglang[all_hpu]", "sglang[test]"] dev_musa = ["sglang[all_musa]", "sglang[test]"] [project.urls] "Homepage" = "https://github.com/sgl-project/sglang" "Bug Tracker" = "https://github.com/sgl-project/sglang/issues" [project.scripts] sglang = "sglang.cli.main:main" [tool.setuptools.package-data] "sglang" = [ "srt/**/*", "jit_kernel/**/*", ] [tool.setuptools.packages.find] exclude = [ "assets*", "benchmark*", "docs*", "dist*", "playground*", "scripts*", "tests*", ] [tool.wheel] exclude = [ "assets*", "benchmark*", "docs*", "dist*", "playground*", "scripts*", "tests*", ] [tool.setuptools_scm] root = ".." version_file = "sglang/_version.py" git_describe_command = ["git", "describe", "--tags", "--long", "--match", "v*"] ================================================ FILE: LICENSE ================================================ Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright 2023-2024 SGLang Team Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ================================================ FILE: README.md ================================================
logo [![PyPI](https://img.shields.io/pypi/v/sglang)](https://pypi.org/project/sglang) ![PyPI - Downloads](https://static.pepy.tech/badge/sglang?period=month) [![license](https://img.shields.io/github/license/sgl-project/sglang.svg)](https://github.com/sgl-project/sglang/tree/main/LICENSE) [![issue resolution](https://img.shields.io/github/issues-closed-raw/sgl-project/sglang)](https://github.com/sgl-project/sglang/issues) [![open issues](https://img.shields.io/github/issues-raw/sgl-project/sglang)](https://github.com/sgl-project/sglang/issues) [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/sgl-project/sglang)
--------------------------------------------------------------------------------

Blog | Documentation | Roadmap | Join Slack | Weekly Dev Meeting | Slides

## News - [2026/02] 🔥 Unlocking 25x Inference Performance with SGLang on NVIDIA GB300 NVL72 ([blog](https://lmsys.org/blog/2026-02-20-gb300-inferencex/)). - [2026/01] 🔥 SGLang Diffusion accelerates video and image generation ([blog](https://lmsys.org/blog/2026-01-16-sglang-diffusion/)). - [2025/12] SGLang provides day-0 support for latest open models ([MiMo-V2-Flash](https://lmsys.org/blog/2025-12-16-mimo-v2-flash/), [Nemotron 3 Nano](https://lmsys.org/blog/2025-12-15-run-nvidia-nemotron-3-nano/), [Mistral Large 3](https://github.com/sgl-project/sglang/pull/14213), [LLaDA 2.0 Diffusion LLM](https://lmsys.org/blog/2025-12-19-diffusion-llm/), [MiniMax M2](https://lmsys.org/blog/2025-11-04-miminmax-m2/)). - [2025/10] 🔥 SGLang now runs natively on TPU with the SGLang-Jax backend ([blog](https://lmsys.org/blog/2025-10-29-sglang-jax/)). - [2025/09] Deploying DeepSeek on GB200 NVL72 with PD and Large Scale EP (Part II): 3.8x Prefill, 4.8x Decode Throughput ([blog](https://lmsys.org/blog/2025-09-25-gb200-part-2/)). - [2025/09] SGLang Day 0 Support for DeepSeek-V3.2 with Sparse Attention ([blog](https://lmsys.org/blog/2025-09-29-deepseek-V32/)). - [2025/08] SGLang x AMD SF Meetup on 8/22: Hands-on GPU workshop, tech talks by AMD/xAI/SGLang, and networking ([Roadmap](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_sglang_roadmap.pdf), [Large-scale EP](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_sglang_ep.pdf), [Highlights](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_highlights.pdf), [AITER/MoRI](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_aiter_mori.pdf), [Wave](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_wave.pdf)).
More - [2025/11] SGLang Diffusion accelerates video and image generation ([blog](https://lmsys.org/blog/2025-11-07-sglang-diffusion/)). - [2025/10] PyTorch Conference 2025 SGLang Talk ([slide](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/sglang_pytorch_2025.pdf)). - [2025/10] SGLang x Nvidia SF Meetup on 10/2 ([recap](https://x.com/lmsysorg/status/1975339501934510231)). - [2025/08] SGLang provides day-0 support for OpenAI gpt-oss model ([instructions](https://github.com/sgl-project/sglang/issues/8833)) - [2025/06] SGLang, the high-performance serving infrastructure powering trillions of tokens daily, has been awarded the third batch of the Open Source AI Grant by a16z ([a16z blog](https://a16z.com/advancing-open-source-ai-through-benchmarks-and-bold-experimentation/)). - [2025/05] Deploying DeepSeek with PD Disaggregation and Large-scale Expert Parallelism on 96 H100 GPUs ([blog](https://lmsys.org/blog/2025-05-05-large-scale-ep/)). - [2025/06] Deploying DeepSeek on GB200 NVL72 with PD and Large Scale EP (Part I): 2.7x Higher Decoding Throughput ([blog](https://lmsys.org/blog/2025-06-16-gb200-part-1/)). - [2025/03] Supercharge DeepSeek-R1 Inference on AMD Instinct MI300X ([AMD blog](https://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1-Part2/README.html)) - [2025/03] SGLang Joins PyTorch Ecosystem: Efficient LLM Serving Engine ([PyTorch blog](https://pytorch.org/blog/sglang-joins-pytorch/)) - [2025/02] Unlock DeepSeek-R1 Inference Performance on AMD Instinct™ MI300X GPU ([AMD blog](https://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1_Perf/README.html)) - [2025/01] SGLang provides day one support for DeepSeek V3/R1 models on NVIDIA and AMD GPUs with DeepSeek-specific optimizations. ([instructions](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3), [AMD blog](https://www.amd.com/en/developer/resources/technical-articles/amd-instinct-gpus-power-deepseek-v3-revolutionizing-ai-development-with-sglang.html), [10+ other companies](https://x.com/lmsysorg/status/1887262321636221412)) - [2024/12] v0.4 Release: Zero-Overhead Batch Scheduler, Cache-Aware Load Balancer, Faster Structured Outputs ([blog](https://lmsys.org/blog/2024-12-04-sglang-v0-4/)). - [2024/10] The First SGLang Online Meetup ([slides](https://github.com/sgl-project/sgl-learning-materials?tab=readme-ov-file#the-first-sglang-online-meetup)). - [2024/09] v0.3 Release: 7x Faster DeepSeek MLA, 1.5x Faster torch.compile, Multi-Image/Video LLaVA-OneVision ([blog](https://lmsys.org/blog/2024-09-04-sglang-v0-3/)). - [2024/07] v0.2 Release: Faster Llama3 Serving with SGLang Runtime (vs. TensorRT-LLM, vLLM) ([blog](https://lmsys.org/blog/2024-07-25-sglang-llama3/)). - [2024/02] SGLang enables **3x faster JSON decoding** with compressed finite state machine ([blog](https://lmsys.org/blog/2024-02-05-compressed-fsm/)). - [2024/01] SGLang provides up to **5x faster inference** with RadixAttention ([blog](https://lmsys.org/blog/2024-01-17-sglang/)). - [2024/01] SGLang powers the serving of the official **LLaVA v1.6** release demo ([usage](https://github.com/haotian-liu/LLaVA?tab=readme-ov-file#demo)).
## About SGLang is a high-performance serving framework for large language models and multimodal models. It is designed to deliver low-latency and high-throughput inference across a wide range of setups, from a single GPU to large distributed clusters. Its core features include: - **Fast Runtime**: Provides efficient serving with RadixAttention for prefix caching, a zero-overhead CPU scheduler, prefill-decode disaggregation, speculative decoding, continuous batching, paged attention, tensor/pipeline/expert/data parallelism, structured outputs, chunked prefill, quantization (FP4/FP8/INT4/AWQ/GPTQ), and multi-LoRA batching. - **Broad Model Support**: Supports a wide range of language models (Llama, Qwen, DeepSeek, Kimi, GLM, GPT, Gemma, Mistral, etc.), embedding models (e5-mistral, gte, mcdse), reward models (Skywork), and diffusion models (WAN, Qwen-Image), with easy extensibility for adding new models. Compatible with most Hugging Face models and OpenAI APIs. - **Extensive Hardware Support**: Runs on NVIDIA GPUs (GB200/B300/H100/A100/Spark), AMD GPUs (MI355/MI300), Intel Xeon CPUs, Google TPUs, Ascend NPUs, and more. - **Active Community**: SGLang is open-source and supported by a vibrant community with widespread industry adoption, powering over 400,000 GPUs worldwide. - **RL & Post-Training Backbone**: SGLang is a proven rollout backend used for training many frontier models, with native RL integrations and adoption by well-known post-training frameworks such as [**AReaL**](https://github.com/inclusionAI/AReaL), [**Miles**](https://github.com/radixark/miles), [**slime**](https://github.com/THUDM/slime), [**Tunix**](https://github.com/google/tunix), [**verl**](https://github.com/volcengine/verl) and more. ## Getting Started - [Install SGLang](https://docs.sglang.io/get_started/install.html) - [Quick Start](https://docs.sglang.io/basic_usage/send_request.html) - [Backend Tutorial](https://docs.sglang.io/basic_usage/openai_api_completions.html) - [Frontend Tutorial](https://docs.sglang.io/references/frontend/frontend_tutorial.html) - [Contribution Guide](https://docs.sglang.io/developer_guide/contribution_guide.html) ## Benchmark and Performance Learn more in the release blogs: [v0.2 blog](https://lmsys.org/blog/2024-07-25-sglang-llama3/), [v0.3 blog](https://lmsys.org/blog/2024-09-04-sglang-v0-3/), [v0.4 blog](https://lmsys.org/blog/2024-12-04-sglang-v0-4/), [Large-scale expert parallelism](https://lmsys.org/blog/2025-05-05-large-scale-ep/), [GB200 rack-scale parallelism](https://lmsys.org/blog/2025-09-25-gb200-part-2/). ## Adoption and Sponsorship SGLang has been deployed at large scale, generating trillions of tokens in production each day. It is trusted and adopted by a wide range of leading enterprises and institutions, including xAI, AMD, NVIDIA, Intel, LinkedIn, Cursor, Oracle Cloud, Google Cloud, Microsoft Azure, AWS, Atlas Cloud, Voltage Park, Nebius, DataCrunch, Novita, InnoMatrix, MIT, UCLA, the University of Washington, Stanford, UC Berkeley, Tsinghua University, Jam & Tea Studios, Baseten, and other major technology organizations across North America and Asia. As an open-source LLM inference engine, SGLang has become the de facto industry standard, with deployments running on over 400,000 GPUs worldwide. SGLang is currently hosted under the non-profit open-source organization [LMSYS](https://lmsys.org/about/). logo ## Contact Us For enterprises interested in adopting or deploying SGLang at scale, including technical consulting, sponsorship opportunities, or partnership inquiries, please contact us at sglang@lmsys.org ## Acknowledgment We learned the design and reused code from the following projects: [Guidance](https://github.com/guidance-ai/guidance), [vLLM](https://github.com/vllm-project/vllm), [LightLLM](https://github.com/ModelTC/lightllm), [FlashInfer](https://github.com/flashinfer-ai/flashinfer), [Outlines](https://github.com/outlines-dev/outlines), and [LMQL](https://github.com/eth-sri/lmql). ================================================ FILE: benchmark/asr/README.md ================================================ # ASR Benchmark This benchmark evaluates the performance and accuracy (Word Error Rate - WER) of Automatic Speech Recognition (ASR) models served via SGLang. ## Supported Models - `openai/whisper-large-v3` - `openai/whisper-large-v3-turbo` ## Setup Install the required dependencies: ```bash apt install ffmpeg pip install librosa soundfile datasets evaluate jiwer transformers openai torchcodec torch ``` ## Running the Benchmark ### 1. Start SGLang Server Launch the SGLang server with a Whisper model: ```bash python -m sglang.launch_server --model-path openai/whisper-large-v3 --port 30000 ``` ### 2. Run the Benchmark Script Basic usage (using chat completions API): ```bash python bench_sglang.py --base-url http://localhost:30000 --model openai/whisper-large-v3 --n-examples 10 ``` Using the OpenAI-compatible transcription API: ```bash python bench_sglang.py \ --base-url http://localhost:30000 \ --model openai/whisper-large-v3 \ --api-type transcription \ --language English \ --n-examples 10 ``` Run with streaming and show real-time output: ```bash python bench_sglang.py \ --base-url http://localhost:30000 \ --model openai/whisper-large-v3 \ --api-type transcription \ --stream \ --show-predictions \ --concurrency 1 ``` Run with higher concurrency and save results: ```bash python bench_sglang.py \ --base-url http://localhost:30000 \ --model openai/whisper-large-v3 \ --concurrency 8 \ --n-examples 100 \ --output results.json \ --show-predictions ``` ## Arguments | Argument | Description | Default | |----------|-------------|---------| | `--base-url` | SGLang server URL | `http://localhost:30000` | | `--model` | Model name on the server | `openai/whisper-large-v3` | | `--dataset` | HuggingFace dataset for evaluation | `D4nt3/esb-datasets-earnings22-validation-tiny-filtered` | | `--split` | Dataset split to use | `validation` | | `--concurrency` | Number of concurrent requests | `4` | | `--n-examples` | Number of examples to process (`-1` for all) | `-1` | | `--output` | Path to save results as JSON | `None` | | `--show-predictions` | Display sample predictions | `False` | | `--print-n` | Number of samples to display | `5` | | `--api-type` | API to use: `chat` (chat completions) or `transcription` (audio transcriptions) | `chat` | | `--language` | Language for transcription API (e.g., `English`, `en`) | `None` | | `--stream` | Enable streaming mode for transcription API | `False` | ## Metrics The benchmark outputs: | Metric | Description | |--------|-------------| | **Total Requests** | Number of successful ASR requests processed | | **WER** | Word Error Rate (lower is better), computed using the `evaluate` library | | **Average Latency** | Mean time per request (seconds) | | **Median Latency** | 50th percentile latency (seconds) | | **95th Latency** | 95th percentile latency (seconds) | | **Throughput** | Requests processed per second | | **Token Throughput** | Output tokens per second | ## Example Output ```bash python bench_sglang.py --api-type transcription --concurrency 128 --model openai/whisper-large-v3 --show-predictions Loading dataset: D4nt3/esb-datasets-earnings22-validation-tiny-filtered... Using API type: transcription Repo card metadata block was not found. Setting CardData to empty. WARNING:huggingface_hub.repocard:Repo card metadata block was not found. Setting CardData to empty. Performing warmup... Processing 511 samples... ------------------------------ Results for openai/whisper-large-v3: Total Requests: 511 WER: 12.7690 Average Latency: 1.3602s Median Latency: 1.2090s 95th Latency: 2.9986s Throughput: 19.02 req/s Token Throughput: 354.19 tok/s Total Test Time: 26.8726s ------------------------------ ==================== Sample Predictions ==================== Sample 1: REF: on the use of taxonomy i you know i think it is it is early days for us to to make any clear indications to the market about the proportion that would fall under that requirement PRED: on the eu taxonomy i think it is early days for us to make any clear indications to the market about the proportion that would fall under that requirement ---------------------------------------- Sample 2: REF: so within fiscal year 2021 say 120 a 100 depending on what the micro will do and next year it is not necessarily payable in q one is we will look at what the cash flows for 2022 look like PRED: so within fiscal year 2021 say $120000 $100000 depending on what the macro will do and next year it is not necessarily payable in q one is we will look at what the cash flows for 2022 look like ---------------------------------------- Sample 3: REF: we talked about 4.7 gigawatts PRED: we talked about 4.7 gigawatts ---------------------------------------- Sample 4: REF: and you know depending on that working capital build we will we will see what that yields PRED: and depending on that working capital build we will see what that yields what ---------------------------------------- Sample 5: REF: so on on sinopec what we have agreed with sinopec way back then is that free cash flows after paying all capexs are distributed out 30 70% PRED: so on sinopec what we have agreed with sinopec way back then is that free cash flows after paying all capexes are distributed out 30% 70% ---------------------------------------- ============================================================ ``` ## Notes - Audio samples longer than 30 seconds are automatically filtered out (Whisper limitation) - The benchmark performs a warmup request before measuring performance - Results are normalized using the model's tokenizer when available - When using `--stream` with `--show-predictions`, use `--concurrency 1` for clean sequential output - The `--language` option accepts both full names (e.g., `English`) and ISO 639-1 codes (e.g., `en`) ## Troubleshooting **Server connection refused** - Ensure the SGLang server is running and accessible at the specified `--base-url` - Check that the port is not blocked by a firewall **Out of memory errors** - Reduce `--concurrency` to lower GPU memory usage - Use a smaller Whisper model variant ================================================ FILE: benchmark/asr/bench_sglang.py ================================================ import argparse import asyncio import base64 import io import json import time from statistics import mean, median import httpx import librosa import numpy as np import soundfile from datasets import load_dataset from evaluate import load from openai import AsyncOpenAI, OpenAI from transformers import AutoTokenizer def to_bytes(y, sr): buffer = io.BytesIO() soundfile.write(buffer, y, sr, format="WAV") buffer.seek(0) return buffer async def run_asr_chat(client, model_name, y, sr): """Use chat completions API with audio_url for ASR.""" with to_bytes(y, sr) as f: audio_bytes = f.read() audio_base64 = base64.b64encode(audio_bytes).decode("utf-8") start_time = time.perf_counter() response = await client.chat.completions.create( model=model_name, messages=[ { "role": "user", "content": [ { "type": "audio_url", "audio_url": {"url": f"data:audio/wav;base64,{audio_base64}"}, } ], } ], temperature=0.0, ) end_time = time.perf_counter() asr_text = response.choices[0].message.content latency = end_time - start_time return latency, asr_text def run_asr_transcription_sync(client, model_name, y, sr, language=None): """Use audio transcriptions API for ASR (sync version).""" audio_buffer = to_bytes(y, sr) audio_buffer.name = "audio.wav" # OpenAI client needs a name attribute start_time = time.perf_counter() kwargs = { "model": model_name, "file": audio_buffer, } if language: kwargs["language"] = language transcription = client.audio.transcriptions.create(**kwargs) end_time = time.perf_counter() latency = end_time - start_time return latency, transcription.text def run_asr_transcription_stream_sync( base_url, model_name, y, sr, language=None, show_stream=False ): """Use audio transcriptions API with streaming for ASR.""" audio_buffer = to_bytes(y, sr) audio_bytes = audio_buffer.read() data = { "model": model_name, "response_format": "json", "stream": "true", } if language: data["language"] = language start_time = time.perf_counter() text_chunks = [] if show_stream: print("[STREAM] ", end="", flush=True) with httpx.stream( "POST", f"{base_url}/v1/audio/transcriptions", data=data, files={"file": ("audio.wav", audio_bytes, "audio/wav")}, timeout=60.0, ) as response: for line in response.iter_lines(): if line.startswith("data: ") and not line.startswith("data: [DONE]"): try: chunk = json.loads(line[6:]) if "choices" in chunk and chunk["choices"]: delta = chunk["choices"][0].get("delta", {}) content = delta.get("content", "") if content: text_chunks.append(content) if show_stream: print(content, end="", flush=True) except json.JSONDecodeError: pass if show_stream: print() # newline after stream end_time = time.perf_counter() latency = end_time - start_time return latency, "".join(text_chunks) async def run_asr_transcription( client, model_name, y, sr, language=None, stream=False, base_url=None, show_stream=False, ): """Async wrapper for transcription API (runs sync call in executor).""" loop = asyncio.get_event_loop() if stream: return await loop.run_in_executor( None, run_asr_transcription_stream_sync, base_url, model_name, y, sr, language, show_stream, ) return await loop.run_in_executor( None, run_asr_transcription_sync, client, model_name, y, sr, language ) async def bound_asr( sem, client, model_name, tokenizer, audio, reference, api_type="chat", language=None, stream=False, base_url=None, show_stream=False, ): async with sem: try: if api_type == "transcription": latency, text = await run_asr_transcription( client, model_name, *audio, language=language, stream=stream, base_url=base_url, show_stream=show_stream, ) else: latency, text = await run_asr_chat(client, model_name, *audio) # Calculate tokens for throughput metrics num_output_tokens = len(tokenizer(text, add_special_tokens=False).input_ids) # Normalize for WER evaluation # Whisper tokenizer has a normalize method if hasattr(tokenizer, "normalize"): out = tokenizer.normalize(text) ref = tokenizer.normalize(reference) else: out = text.lower().strip() ref = reference.lower().strip() return latency, num_output_tokens, out, ref except Exception as e: print(f"Error during ASR: {e}") return None async def process_dataset( model_name, client, data, concurrent_request, api_type="chat", language=None, stream=False, base_url=None, show_predictions=False, ): sem = asyncio.Semaphore(concurrent_request) tokenizer = AutoTokenizer.from_pretrained(model_name) # Warmup print("Performing warmup...") audio_warmup, sr_warmup = ( data[0]["audio"]["array"], data[0]["audio"]["sampling_rate"], ) await bound_asr( sem, client, model_name, tokenizer, (audio_warmup, sr_warmup), "", api_type=api_type, language=language, stream=stream, base_url=base_url, show_stream=False, # Don't show stream during warmup ) tasks = [] print(f"Processing {len(data)} samples...") for sample in data: audio, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"] tasks.append( asyncio.create_task( bound_asr( sem, client, model_name, tokenizer, (audio, sr), sample["text"], api_type=api_type, language=language, stream=stream, base_url=base_url, show_stream=show_predictions and stream, ) ) ) results = await asyncio.gather(*tasks) return [r for r in results if r is not None] def run_evaluation(args): # Use sync client for transcription API, async for chat API if args.api_type == "transcription": client = OpenAI(base_url=f"{args.base_url}/v1", api_key="None") else: client = AsyncOpenAI(base_url=f"{args.base_url}/v1", api_key="None") print(f"Loading dataset: {args.dataset}...") print(f"Using API type: {args.api_type}" + (f" (streaming)" if args.stream else "")) dataset = load_dataset(args.dataset, split=args.split) # Filter by duration if needed (Whisper max is 30s) def add_duration(sample): y, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"] sample["duration_ms"] = librosa.get_duration(y=y, sr=sr) * 1000 return sample if "duration_ms" not in dataset.column_names: dataset = dataset.map(add_duration) dataset = dataset.filter(lambda x: x["duration_ms"] < 30000) if args.n_examples > 0: dataset = dataset.select(range(min(args.n_examples, len(dataset)))) start = time.perf_counter() results = asyncio.run( process_dataset( args.model, client, dataset, args.concurrency, api_type=args.api_type, language=args.language, stream=args.stream, base_url=args.base_url, show_predictions=args.show_predictions, ) ) total_test_time = time.perf_counter() - start if not results: print("No successful results to evaluate.") return # Metrics latencies = [res[0] for res in results] total_tokens = sum([res[1] for res in results]) predictions = [res[2] for res in results] references = [res[3] for res in results] wer_metric = load("wer") wer_score = 100 * wer_metric.compute(references=references, predictions=predictions) print("-" * 30) print(f"Results for {args.model}:") print(f"Total Requests: {len(results)}") print(f"WER: {wer_score:.4f}") print(f"Average Latency: {mean(latencies):.4f}s") print(f"Median Latency: {median(latencies):.4f}s") print(f"95th Latency: {np.percentile(latencies, 95):.4f}s") print(f"Throughput: {len(results) / total_test_time:.2f} req/s") print(f"Token Throughput: {total_tokens / total_test_time:.2f} tok/s") print(f"Total Test Time: {total_test_time:.4f}s") print("-" * 30) if args.output: with open(args.output, "w") as f: import json json.dump( { "model": args.model, "dataset": args.dataset, "wer": wer_score, "avg_latency": mean(latencies), "throughput": len(results) / total_test_time, "token_throughput": total_tokens / total_test_time, }, f, indent=2, ) if args.show_predictions: print("\n" + "=" * 20 + " Sample Predictions " + "=" * 20) num_to_show = min(args.print_n, len(results)) for i in range(num_to_show): print(f"Sample {i+1}:") print(f" REF: {references[i]}") print(f" PRED: {predictions[i]}") print("-" * 40) print("=" * 60) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Benchmark sGLang ASR performance.") parser.add_argument( "--base-url", default="http://localhost:30000", help="sGLang server base URL" ) parser.add_argument( "--model", default="openai/whisper-large-v3", help="Model name on the server" ) parser.add_argument( "--dataset", default="D4nt3/esb-datasets-earnings22-validation-tiny-filtered", help="HF dataset repo", ) parser.add_argument("--split", default="validation", help="Dataset split") parser.add_argument( "--concurrency", type=int, default=4, help="Number of concurrent requests" ) parser.add_argument( "--n-examples", "-n", type=int, default=-1, help="Number of examples to test (-1 for all)", ) parser.add_argument("--output", help="Path to save results in JSON") parser.add_argument( "--show-predictions", action="store_true", help="Print sample predictions and references", ) parser.add_argument( "--print-n", type=int, default=5, help="Number of sample predictions to print" ) parser.add_argument( "--api-type", choices=["chat", "transcription"], default="chat", help="API type to use: 'chat' for chat completions with audio_url, 'transcription' for audio.transcriptions API", ) parser.add_argument( "--language", default=None, help="Language code for transcription API (e.g., 'en')", ) parser.add_argument( "--stream", action="store_true", help="Use streaming mode for transcription API", ) args = parser.parse_args() run_evaluation(args) ================================================ FILE: benchmark/bench_attention_sink/bench_attention_sink_triton.py ================================================ import argparse import torch import triton from sglang.srt.layers.attention.triton_ops.decode_attention import ( decode_attention_fwd_grouped, ) from sglang.srt.layers.attention.triton_ops.extend_attention import extend_attention_fwd # gpt oss head_num = 64 head_dim = 64 head_kv_num = 8 @triton.testing.perf_report( triton.testing.Benchmark( x_names=["S"], # sequence length on x-axis x_vals=[128, 256, 512, 1024, 2048, 4096], x_log=True, line_arg="B", # batch size as different lines line_vals=[1, 8, 32, 128], line_names=["B=1", "B=8", "B=32", "B=128"], styles=[ ("blue", "-"), ("green", "-"), ("red", "-"), ("cyan", "-"), ], ylabel="TFLOPS", plot_name="attention-sink-triton-decode", args={}, ) ) def benchmark_decode(B, S, H_Q, H_KV, D): D_V = D dtype = torch.bfloat16 seq_len = S total_tokens = B * seq_len device = torch.device("cuda") sm_scale = 1.0 / (D**0.5) max_kv_splits = 8 num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device="cuda") # q represents the new token being generated, one per batch q = torch.randn(B, H_Q, D, dtype=dtype, device="cuda") # k_buffer and v_buffer represent all previous tokens k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda") v_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda") o = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda") b_seq_len = torch.full((B,), seq_len, device="cuda") kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda") kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len, dim=0) kv_indices = torch.arange(total_tokens, device="cuda") attn_logits1 = torch.empty( (B, H_Q, max_kv_splits, D_V), dtype=torch.float32, device="cuda", ) attn_lse1 = torch.empty( (B, H_Q, max_kv_splits, D_V), dtype=torch.float32, device="cuda", ) sink = torch.randn(H_Q, device=device, dtype=torch.float32) # warmup for _ in range(5): decode_attention_fwd_grouped( q, k_buffer, v_buffer, o, kv_indptr, kv_indices, attn_logits1, attn_lse1, num_kv_splits, max_kv_splits, sm_scale, logit_cap=0.0, sinks=sink, ) # benchmark run_step = 500 start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) start_event.record() for _ in range(run_step): decode_attention_fwd_grouped( q, k_buffer, v_buffer, o, kv_indptr, kv_indices, attn_logits1, attn_lse1, num_kv_splits, max_kv_splits, sm_scale, logit_cap=0.0, sinks=sink, ) end_event.record() end_event.synchronize() torch.cuda.synchronize() ms = start_event.elapsed_time(end_event) / run_step tflops = lambda ms: (2 * B * S * H_Q * D) * 1e-9 / ms # must be causal return tflops(ms) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["S"], # sequence length on x-axis x_vals=[128, 256, 512, 1024, 2048, 4096], x_log=True, line_arg="B", # batch size as different lines line_vals=[1, 8, 32, 128], line_names=["B=1", "B=8", "B=32", "B=128"], styles=[ ("blue", "-"), ("green", "-"), ("red", "-"), ("cyan", "-"), ], ylabel="TFLOPS", plot_name="attention-sink-triton-extend", args={}, ) ) def benchmark_extend(B, S, H_Q, H_KV, D): # S here represents N_CTX from the test dtype = torch.bfloat16 device = "cuda" # Split S into prefix and extend lengths prefill_len = S // 2 # Similar to test's N_CTX // 2 extend_len = S // 4 # Make extend length smaller than prefix # Calculate total tokens and extend tokens total_extend_tokens = B * extend_len total_prefix_tokens = B * prefill_len # Create query, key, value tensors for extension q_extend = torch.randn(total_extend_tokens, H_Q, D, dtype=dtype, device=device) k_extend = torch.randn(total_extend_tokens, H_KV, D, dtype=dtype, device=device) v_extend = torch.randn(total_extend_tokens, H_KV, D, dtype=dtype, device=device) o_extend = torch.empty_like(q_extend) # Create key-value buffers for prefix k_buffer = torch.randn(total_prefix_tokens, H_KV, D, dtype=dtype, device=device) v_buffer = torch.randn(total_prefix_tokens, H_KV, D, dtype=dtype, device=device) # Create index pointers qo_indptr = torch.arange(0, (B + 1) * extend_len, extend_len, device=device).to( torch.int32 ) kv_indptr = torch.arange(0, (B + 1) * prefill_len, prefill_len, device=device).to( torch.int32 ) kv_indices = torch.arange(0, total_prefix_tokens, device=device).to(torch.int32) sm_scale = 1.0 / (D**0.5) # sliding_window = 128 # From GPT-OSS config, skip for now sliding_window = -1 sink = torch.randn(H_Q, device=device, dtype=torch.float32) # warmup for _ in range(5): extend_attention_fwd( q_extend, k_extend, v_extend, o_extend, k_buffer, v_buffer, qo_indptr, kv_indptr, kv_indices, custom_mask=None, is_causal=True, mask_indptr=None, max_len_extend=extend_len, sm_scale=sm_scale, sliding_window_size=sliding_window, sinks=sink, ) # benchmark run_step = 500 start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) start_event.record() for _ in range(run_step): extend_attention_fwd( q_extend, k_extend, v_extend, o_extend, k_buffer, v_buffer, qo_indptr, kv_indptr, kv_indices, custom_mask=None, is_causal=True, mask_indptr=None, max_len_extend=extend_len, sm_scale=sm_scale, sliding_window_size=sliding_window, sinks=sink, ) end_event.record() end_event.synchronize() torch.cuda.synchronize() ms = start_event.elapsed_time(end_event) / run_step # FLOPS calculation: each attention operation requires 2 multiplications per element total_flops = 2 * total_extend_tokens * H_Q * (prefill_len + extend_len / 2) * D tflops = lambda ms: total_flops * 1e-12 / (ms * 1e-3) # convert to TFLOPS return tflops(ms) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--bench", type=str, default="all", help="all, extend, decode") args = parser.parse_args() kwargs = { "H_Q": head_num, "H_KV": head_kv_num, "D": head_dim, } if args.bench in ["all", "decode"]: benchmark_decode.run(print_data=True, show_plots=False, **kwargs) if args.bench in ["all", "extend"]: benchmark_extend.run(print_data=True, show_plots=False, **kwargs) print("Benchmark finished!") ================================================ FILE: benchmark/bench_in_batch_prefix/bench_in_batch_prefix.py ================================================ # Benchmark with lots of common prefixes. Used to benchmark prefix caching performance. # # Launch a server: # python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --log-level-http warning import random import string import time from tqdm import tqdm from transformers import AutoTokenizer import sglang as sgl from sglang import set_default_backend from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint def generate_random_string(token_length: int) -> str: random_string = "".join( random.choices(string.ascii_letters + string.digits, k=token_length * 100) ) tokenized_output = tokenizer.encode(random_string, add_special_tokens=False)[ :token_length ] if len(tokenized_output) < token_length: tokenized_output = tokenized_output + [tokenizer.pad_token_id] * ( token_length - len(tokenized_output) ) decoded_string = tokenizer.decode(tokenized_output, skip_special_tokens=False) return decoded_string def generate_unique_prefix(base_text, index): return str(index) + base_text[len(str(index)) :] @sgl.function def text_qa(s, question, gen_len): s += "Q: " + question + "\n" s += "A:" + sgl.gen("answer", stop="\n", temperature=0, max_tokens=gen_len) def prepare_prompts(num_prefix, num_samples_per_prefix, prefix_length, suffix_length): base_prefix = generate_random_string(prefix_length) tot_input_len = 0 all_prompts = [] for i in tqdm(range(num_prefix), desc="prepare prompts"): unique_prefix = generate_unique_prefix(base_prefix, i) prompt_list = [] for j in range(num_samples_per_prefix): suffix = generate_random_string(suffix_length) prompt = unique_prefix + suffix prompt_list.append(prompt) tot_input_len += len(tokenizer.encode(prompt)) all_prompts.append(prompt_list) return all_prompts, tot_input_len def test_batch_by_batch(all_prompts, gen_len): backend.flush_cache() tot_time = 0 for i in range(len(all_prompts)): tic = time.perf_counter() text_qa.run_batch( list(zip(all_prompts[i], [gen_len] * len(all_prompts[i]))), ) tot_time += time.perf_counter() - tic return tot_time def test_batch_by_batch_with_hint(all_prompts, gen_len): backend.flush_cache() tot_time = 0 for i in range(len(all_prompts)): tic = time.perf_counter() # Send a hint to cache the prefix text_qa.run_batch(list(zip(all_prompts[i][:1], [gen_len]))) # Send the batch text_qa.run_batch(list(zip(all_prompts[i], [gen_len] * len(all_prompts[i])))) tot_time += time.perf_counter() - tic return tot_time def test_send_all(all_prompts, gen_len): backend.flush_cache() all_prompts = [x for prompt_list in all_prompts for x in prompt_list] tic = time.perf_counter() text_qa.run_batch( list(zip(all_prompts, [gen_len] * len(all_prompts))), ) tot_time = time.perf_counter() - tic return tot_time if __name__ == "__main__": tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer") backend = RuntimeEndpoint("http://127.0.0.1:30000") set_default_backend(backend) random.seed(0) num_prefix = 10 num_samples_per_prefix = 32 prefix_length = 1024 suffix_length = 128 gen_len = 1 all_prompts, tot_input_len = prepare_prompts( num_prefix, num_samples_per_prefix, prefix_length, suffix_length ) print(f"Total input token length: {tot_input_len}\n") cost = test_batch_by_batch(all_prompts, gen_len) print(f"Latency of test_batch_by_batch : {cost:.4f} s\n") cost = test_batch_by_batch_with_hint(all_prompts, gen_len) print(f"Latency of test_batch_by_batch_with_hint: {cost:.4f} s\n") cost = test_send_all(all_prompts, gen_len) print(f"Latency of test_send_all : {cost:.4f} s\n") ================================================ FILE: benchmark/bench_linear_attention/bench_gdn_decode.py ================================================ """ Benchmark & Correctness: GDN Packed Decode vs Baseline Decode. Compares: - Baseline: split(mixed_qkv) → view → fused_sigmoid_gating_delta_rule_update - Packed: fused_recurrent_gated_delta_rule_packed_decode (single kernel) The packed path eliminates: - torch.split() + .view() tensor materialization - Separate gating kernel launches - Intermediate tensor allocations Reports correctness (output & state matching) and performance (ms, speedup). Usage: python bench_gdn_decode.py # default sweep python bench_gdn_decode.py --mode bench # benchmark only python bench_gdn_decode.py --mode correctness # correctness only python bench_gdn_decode.py --preset qwen3.5-35b # Qwen3.5-35B-A3B config """ import argparse import os import sys import time sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python")) import torch import triton from sglang.srt.layers.attention.fla.fused_recurrent import ( fused_recurrent_gated_delta_rule_packed_decode, ) from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) # --------------------------------------------------------------------------- # Input factory # --------------------------------------------------------------------------- def make_inputs( B: int, H: int, HV: int, K: int, V: int, pool_size: int, device: str, dtype: torch.dtype, seed: int = 42, ): """Create all input tensors for a single benchmark / correctness run.""" torch.manual_seed(seed) qkv_dim = 2 * H * K + HV * V mixed_qkv = torch.randn(B, qkv_dim, device=device, dtype=dtype) a = torch.randn(B, HV, device=device, dtype=dtype) b = torch.randn(B, HV, device=device, dtype=dtype) A_log = torch.randn(HV, device=device, dtype=dtype) dt_bias = torch.randn(HV, device=device, dtype=dtype) ssm_states = torch.randn(pool_size, HV, V, K, device=device, dtype=dtype) * 0.1 cache_indices = torch.arange(B, device=device, dtype=torch.int32) cu_seqlens = torch.arange(B + 1, device=device, dtype=torch.long) return dict( B=B, H=H, HV=HV, K=K, V=V, qkv_dim=qkv_dim, pool_size=pool_size, mixed_qkv=mixed_qkv, a=a, b=b, A_log=A_log, dt_bias=dt_bias, ssm_states=ssm_states, cache_indices=cache_indices, cu_seqlens=cu_seqlens, ) # --------------------------------------------------------------------------- # Runner wrappers # --------------------------------------------------------------------------- def run_baseline(inp): """Baseline path: split → view → fused_sigmoid_gating_delta_rule_update. This mirrors the FULL original decode path in GDNAttnBackend.forward_decode, including the split, view, and kernel call. """ B, H, HV, K, V = inp["B"], inp["H"], inp["HV"], inp["K"], inp["V"] mixed_qkv = inp["mixed_qkv"] ssm_states = inp["ssm_states"].clone() # Step 1: split (same as forward_decode) q_flat, k_flat, v_flat = torch.split(mixed_qkv, [H * K, H * K, HV * V], dim=-1) # Step 2: view + reshape (same as forward_decode) q = q_flat.view(1, B, H, K) k = k_flat.view(1, B, H, K) v = v_flat.view(1, B, HV, V) # Step 3: fused gating + recurrent update o = fused_sigmoid_gating_delta_rule_update( A_log=inp["A_log"], dt_bias=inp["dt_bias"], q=q, k=k, v=v, a=inp["a"], b=inp["b"], initial_state_source=ssm_states, initial_state_indices=inp["cache_indices"], cu_seqlens=inp["cu_seqlens"], use_qk_l2norm_in_kernel=True, softplus_beta=1.0, softplus_threshold=20.0, ) return o, ssm_states def run_packed(inp): """Packed path: single fused kernel directly on mixed_qkv.""" B, HV, K, V = inp["B"], inp["HV"], inp["K"], inp["V"] ssm_states = inp["ssm_states"].clone() out = inp["mixed_qkv"].new_empty(B, 1, HV, V) fused_recurrent_gated_delta_rule_packed_decode( mixed_qkv=inp["mixed_qkv"], a=inp["a"], b=inp["b"], A_log=inp["A_log"], dt_bias=inp["dt_bias"], scale=inp["K"] ** -0.5, initial_state=ssm_states, out=out, ssm_state_indices=inp["cache_indices"], use_qk_l2norm_in_kernel=True, ) # Convert [B, 1, HV, V] → [1, B, HV, V] to match baseline layout return out.transpose(0, 1), ssm_states # --------------------------------------------------------------------------- # Correctness check # --------------------------------------------------------------------------- def check_correctness(B, H, HV, K, V, pool_size, device, dtype, seed=42): """Run correctness check for a single config. Returns True if PASS.""" tag = f"B={B:>4} H={H:>2} HV={HV:>2} K={K:>3} V={V:>3} pool={pool_size:>4}" inp = make_inputs(B, H, HV, K, V, pool_size, device, dtype, seed=seed) o_baseline, state_baseline = run_baseline(inp) o_packed, state_packed = run_packed(inp) # Output comparison atol = 2e-2 if dtype != torch.float32 else 1e-4 rtol = 1e-2 if dtype != torch.float32 else 1e-4 try: torch.testing.assert_close(o_packed, o_baseline, atol=atol, rtol=rtol) output_ok = True except AssertionError as e: output_ok = False out_diff = (o_packed - o_baseline).abs().max().item() # State comparison (only for slots that were updated) indices = inp["cache_indices"] try: torch.testing.assert_close( state_packed[indices], state_baseline[indices], atol=atol, rtol=rtol ) state_ok = True except AssertionError: state_ok = False st_diff = (state_packed[indices] - state_baseline[indices]).abs().max().item() passed = output_ok and state_ok if passed: print(f" [PASS] {tag}") else: details = [] if not output_ok: details.append(f"output max_diff={out_diff:.6f}") if not state_ok: details.append(f"state max_diff={st_diff:.6f}") print(f" [FAIL] {tag} ({', '.join(details)})") return passed # --------------------------------------------------------------------------- # Benchmark # --------------------------------------------------------------------------- def bench_shape(B, H, HV, K, V, pool_size, device, dtype): """Benchmark baseline vs packed for a single config.""" inp = make_inputs(B, H, HV, K, V, pool_size, device, dtype) # ── Baseline: full path including split + view ── def fn_baseline(): q_flat, k_flat, v_flat = torch.split( inp["mixed_qkv"], [H * K, H * K, HV * V], dim=-1 ) q = q_flat.view(1, B, H, K) k = k_flat.view(1, B, H, K) v = v_flat.view(1, B, HV, V) fused_sigmoid_gating_delta_rule_update( A_log=inp["A_log"], dt_bias=inp["dt_bias"], q=q, k=k, v=v, a=inp["a"], b=inp["b"], initial_state_source=inp["ssm_states"], initial_state_indices=inp["cache_indices"], cu_seqlens=inp["cu_seqlens"], use_qk_l2norm_in_kernel=True, softplus_beta=1.0, softplus_threshold=20.0, ) # ── Packed: single kernel ── out_buf = inp["mixed_qkv"].new_empty(B, 1, HV, V) def fn_packed(): fused_recurrent_gated_delta_rule_packed_decode( mixed_qkv=inp["mixed_qkv"], a=inp["a"], b=inp["b"], A_log=inp["A_log"], dt_bias=inp["dt_bias"], scale=K**-0.5, initial_state=inp["ssm_states"], out=out_buf, ssm_state_indices=inp["cache_indices"], use_qk_l2norm_in_kernel=True, ) # Warmup for _ in range(10): fn_baseline() fn_packed() torch.cuda.synchronize() quantiles = [0.5, 0.2, 0.8] try: ms_baseline, ms_base_lo, ms_base_hi = triton.testing.do_bench( fn_baseline, quantiles=quantiles, warmup=50, rep=200 ) ms_packed, ms_pack_lo, ms_pack_hi = triton.testing.do_bench( fn_packed, quantiles=quantiles, warmup=50, rep=200 ) except Exception: # Fallback to manual timing torch.cuda.synchronize() N = 200 start = time.perf_counter() for _ in range(N): fn_baseline() torch.cuda.synchronize() ms_baseline = (time.perf_counter() - start) / N * 1000 start = time.perf_counter() for _ in range(N): fn_packed() torch.cuda.synchronize() ms_packed = (time.perf_counter() - start) / N * 1000 speedup = ms_baseline / ms_packed if ms_packed > 0 else float("inf") saved_us = (ms_baseline - ms_packed) * 1000 print( f" {B:>5} {H:>3} {HV:>3} {K:>3} {V:>3} | " f"{ms_baseline * 1000:>10.1f} | " f"{ms_packed * 1000:>10.1f} | " f"{speedup:>7.2f}x | " f"{saved_us:>+9.1f}" ) # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def run_correctness(device, dtype): print("=" * 70) print("Correctness: Baseline GDN Decode vs Packed GDN Decode") print("=" * 70) shapes = [ # (B, H, HV, K, V, pool_size) # --- Qwen3.5-35B-A3B style (TP=2: H=8, HV=16) --- (1, 8, 16, 128, 128, 32), (4, 8, 16, 128, 128, 32), (16, 8, 16, 128, 128, 64), (32, 8, 16, 128, 128, 128), (64, 8, 16, 128, 128, 128), (128, 8, 16, 128, 128, 256), (256, 8, 16, 128, 128, 512), # --- Qwen3.5-35B-A3B style (TP=1: H=16, HV=32) --- (1, 16, 32, 128, 128, 32), (32, 16, 32, 128, 128, 128), (64, 16, 32, 128, 128, 128), # --- Qwen3-Next-80B-A3B style --- (32, 16, 16, 128, 128, 128), (64, 16, 16, 128, 128, 128), # --- With PAD_SLOT_ID --- (32, 8, 16, 128, 128, 128), # some indices may be padded # --- Edge cases --- (1, 8, 16, 128, 128, 32), (2, 8, 16, 128, 128, 32), ] all_pass = True for B, H, HV, K, V, pool_size in shapes: if not check_correctness(B, H, HV, K, V, pool_size, device, dtype): all_pass = False # PAD_SLOT_ID test print("\n PAD_SLOT_ID test (indices with -1):") inp = make_inputs(32, 8, 16, 128, 128, 128, device, dtype) o_baseline, st_baseline = run_baseline(inp) o_packed, st_packed = run_packed(inp) try: torch.testing.assert_close(o_packed, o_baseline, atol=2e-2, rtol=1e-2) print(" [PASS] PAD_SLOT_ID=-1 handling") except AssertionError: print(" [FAIL] PAD_SLOT_ID=-1 handling") all_pass = False print() if all_pass: print("ALL PASSED.") else: print("SOME FAILED.") return all_pass def run_benchmark(device, dtype, args): print() print("=" * 85) print("Benchmark: Baseline GDN Decode vs Packed GDN Decode") print("=" * 85) K = args.head_size_k V = args.head_size_v pool_size = args.pool_size if args.preset == "qwen3.5-35b": # Qwen3.5-35B-A3B: H_qk=16, H_v=32, K=128, V=128 # After TP=2: H=8, HV=16 bench_configs = [ # (B, H, HV) — TP=2 config (1, 8, 16), (2, 8, 16), (4, 8, 16), (8, 8, 16), (16, 8, 16), (32, 8, 16), (64, 8, 16), (128, 8, 16), (256, 8, 16), (512, 8, 16), # TP=1 config (full heads) (1, 16, 32), (8, 16, 32), (32, 16, 32), (64, 16, 32), (128, 16, 32), (256, 16, 32), ] elif args.preset == "qwen3-next-80b": bench_configs = [ # Qwen3-Next-80B-A3B: all same H=HV=16 after TP (1, 16, 16), (8, 16, 16), (32, 16, 16), (64, 16, 16), (128, 16, 16), (256, 16, 16), ] else: bench_configs = [] for B in args.batch_sizes: for H in args.num_q_heads: for HV in args.num_v_heads: bench_configs.append((B, H, HV)) print(f" Config: K={K}, V={V}, pool_size={pool_size}, dtype={dtype}") print( f" {'B':>5} {'H':>3} {'HV':>3} {'K':>3} {'V':>3} | " f"{'base (μs)':>10} | " f"{'packed (μs)':>10} | " f"{'speedup':>8} | " f"{'saved (μs)':>10}" ) print(" " + "-" * 75) for B, H, HV in bench_configs: actual_pool = max(pool_size, B + 16) bench_shape(B, H, HV, K, V, actual_pool, device, dtype) def main(): parser = argparse.ArgumentParser( description="Benchmark & Correctness: GDN Packed Decode vs Baseline" ) parser.add_argument( "--mode", choices=["all", "correctness", "bench"], default="all", help="Run mode (default: all)", ) parser.add_argument( "--preset", choices=["qwen3.5-35b", "qwen3-next-80b", "custom"], default="qwen3.5-35b", help="Preset config (default: qwen3.5-35b)", ) parser.add_argument( "--dtype", choices=["float16", "bfloat16", "float32"], default="bfloat16", ) parser.add_argument("--head-size-k", type=int, default=128) parser.add_argument("--head-size-v", type=int, default=128) parser.add_argument("--pool-size", type=int, default=512) parser.add_argument( "--batch-sizes", type=int, nargs="+", default=[1, 4, 8, 16, 32, 64, 128, 256, 512], ) parser.add_argument( "--num-q-heads", type=int, nargs="+", default=[8, 16], ) parser.add_argument( "--num-v-heads", type=int, nargs="+", default=[16, 32], ) args = parser.parse_args() device = "cuda" dtype = getattr(torch, args.dtype) cap = torch.cuda.get_device_capability() dev_name = torch.cuda.get_device_name() print(f"Device: {dev_name} (SM {cap[0]}{cap[1]})") if args.mode in ("all", "correctness"): all_pass = run_correctness(device, dtype) if not all_pass and args.mode == "all": print("\nSkipping benchmark due to correctness failures.") return 1 if args.mode in ("all", "bench"): run_benchmark(device, dtype, args) return 0 if __name__ == "__main__": sys.exit(main()) ================================================ FILE: benchmark/bench_linear_attention/bench_gdn_prefill.py ================================================ """ Benchmark & Correctness: Triton GDN vs FlashInfer GDN (prefill). Compares: - Triton: sglang's chunk_gated_delta_rule (K-contiguous pool, pool-indexed) - FlashInfer: flashinfer's chunk_gated_delta_rule (gather/scatter, 3D tensors) The two kernels have different APIs: - Triton: q/k/v=[1,T,H,D], g=logsigmoid, beta=sigmoid, has initial_state_indices - FlashInfer: q/k/v=[T,H,D], g=alpha(float32), beta=float32, no indices (gathered state) Reports correctness (output & state matching) and performance (ms, TFLOPS, TB/s). Usage: python benchmark_gdn_prefill.py # default sweep python benchmark_gdn_prefill.py --mode bench # benchmark only python benchmark_gdn_prefill.py --mode correctness # correctness only python benchmark_gdn_prefill.py --preset qwen3-next # Qwen3-Next config """ import argparse import os import sys sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python")) import torch from flashinfer.gdn_prefill import ( chunk_gated_delta_rule as flashinfer_chunk_gated_delta_rule, ) from sglang.srt.layers.attention.fla.chunk import ( chunk_gated_delta_rule as triton_chunk_gated_delta_rule, ) from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def make_k_contiguous(t: torch.Tensor) -> torch.Tensor: """ Given a V-contiguous tensor [..., K, V], return a K-contiguous view of the same logical shape [..., K, V] (physically [..., V, K], K-last). """ return t.transpose(-2, -1).contiguous().transpose(-2, -1) def gdn_flops( total_seq_len: int, num_heads: int, head_size_k: int, head_size_v: int, ) -> int: """ FLOPs for GDN prefill (delta rule). Per token per head: 1. k @ v^T (outer product): 2 * K * V 2. q @ state (output): 2 * K * V """ outer_product_flops = 2 * total_seq_len * num_heads * head_size_k * head_size_v output_flops = 2 * total_seq_len * num_heads * head_size_k * head_size_v return outer_product_flops + output_flops def gdn_bytes( total_seq_len: int, num_q_heads: int, num_v_heads: int, head_size_k: int, head_size_v: int, num_seqs: int, dtype: torch.dtype, ) -> int: """Memory bytes accessed (inputs + outputs + state).""" num_o_heads = max(num_q_heads, num_v_heads) elem = dtype.itemsize q_bytes = total_seq_len * num_q_heads * head_size_k * elem k_bytes = total_seq_len * num_v_heads * head_size_k * elem v_bytes = total_seq_len * num_v_heads * head_size_v * elem o_bytes = total_seq_len * num_o_heads * head_size_v * elem # state (float32): read + write state_bytes = 2 * num_seqs * num_o_heads * head_size_k * head_size_v * 4 # g, beta (float32) g_bytes = total_seq_len * num_o_heads * 4 beta_bytes = total_seq_len * num_o_heads * 4 return q_bytes + k_bytes + v_bytes + o_bytes + state_bytes + g_bytes + beta_bytes # --------------------------------------------------------------------------- # Input factory # --------------------------------------------------------------------------- def make_inputs( B: int, T_per_seq: int, H: int, K: int, V: int, pool_size: int, device: str, dtype: torch.dtype, sequential_indices: bool = False, seed: int = 42, ): """Create all input tensors for a single benchmark / correctness run. Returns a dict with both Triton-format and FlashInfer-format tensors. """ T = B * T_per_seq torch.manual_seed(seed) if sequential_indices: cache_indices = torch.arange(B, dtype=torch.int32, device=device) else: perm = torch.randperm(pool_size, device=device)[:B] cache_indices = perm.to(torch.int32) pool_init = torch.randn(pool_size, H, K, V, dtype=dtype, device=device) * 0.1 cu_seqlens = torch.arange( 0, (B + 1) * T_per_seq, T_per_seq, dtype=torch.long, device=device ) # Triton format: [1, T, H, D] q = torch.randn(1, T, H, K, dtype=dtype, device=device) k = torch.randn(1, T, H, K, dtype=dtype, device=device) v = torch.randn(1, T, H, V, dtype=dtype, device=device) # g (logsigmoid) and beta (sigmoid) in Triton format: [1, T, H] g_raw = torch.randn(1, T, H, dtype=dtype, device=device) g_triton = torch.nn.functional.logsigmoid(g_raw) # logsigmoid for Triton beta_triton = torch.sigmoid(torch.randn(1, T, H, dtype=dtype, device=device)) return dict( B=B, T=T, T_per_seq=T_per_seq, H=H, K=K, V=V, pool_size=pool_size, cache_indices=cache_indices, pool_init=pool_init, cu_seqlens=cu_seqlens, q=q, k=k, v=v, g_triton=g_triton, beta_triton=beta_triton, ) # --------------------------------------------------------------------------- # Runner wrappers # --------------------------------------------------------------------------- def run_triton(inp): """Triton path: K-contiguous pool, pool-indexed, [1,T,H,D] tensors.""" pool = make_k_contiguous(inp["pool_init"].clone()) o, _, h = triton_chunk_gated_delta_rule( q=inp["q"], k=inp["k"], v=inp["v"], g=inp["g_triton"], beta=inp["beta_triton"], initial_state=pool, initial_state_indices=inp["cache_indices"], cu_seqlens=inp["cu_seqlens"], head_first=False, use_qk_l2norm_in_kernel=True, ) return o, pool, h def run_flashinfer(inp): """FlashInfer path: matches sglang FlashInferGDNKernel.extend() exactly. Key differences from Triton path: - q, k are L2-normalized BEFORE calling the kernel - use_qk_l2norm_in_kernel=False (kernel skips internal normalization) - Tensors are [T, H, D] (no batch dim) - g is alpha = exp(logsigmoid(...)) = sigmoid(...), float32 - beta is float32 - initial_state is gathered from pool (no pool-index support) - Uses keyword arguments (matching sglang production code) NOTE: FlashInfer GDN requires K == V (square head_size). """ K = inp["K"] V = inp["V"] assert K == V, f"FlashInfer GDN requires K == V, got K={K}, V={V}" pool = make_k_contiguous(inp["pool_init"].clone()) cache_indices = inp["cache_indices"] # Gather states from K-contiguous pool -> K-contiguous float32 # In production, ssm_states is already float32 so .float() is no-op. # Here pool_init is bf16, so .float() loses K-contiguous layout. gathered = pool[cache_indices] initial_state = make_k_contiguous(gathered.float().contiguous()) q_fi = l2norm_fwd(inp["q"][0].contiguous()) k_fi = l2norm_fwd(inp["k"][0].contiguous()) v_fi = inp["v"][0].contiguous() # g -> alpha (exp of logsigmoid = sigmoid), float32 alpha_fi = torch.exp(inp["g_triton"][0].to(torch.float32)) # beta -> float32 beta_fi = inp["beta_triton"][0].to(torch.float32) cu_seqlens_fi = inp["cu_seqlens"].to(torch.int64) # Call FlashInfer with keyword args (matching sglang production code) # use_qk_l2norm_in_kernel=False because we pre-normalized above o_fi, state_fi = flashinfer_chunk_gated_delta_rule( q=q_fi, k=k_fi, v=v_fi, g=alpha_fi, beta=beta_fi, scale=None, initial_state=initial_state, output_final_state=True, cu_seqlens=cu_seqlens_fi, use_qk_l2norm_in_kernel=False, ) # Scatter updated states back to K-contiguous pool pool[cache_indices] = state_fi.to(pool.dtype) # Reshape output: [T, H, D] -> [1, T, H, D] to match Triton o_out = o_fi.unsqueeze(0) return o_out, pool, state_fi # --------------------------------------------------------------------------- # Correctness check # --------------------------------------------------------------------------- def check_shape( B, T_per_seq, H, K, V, pool_size, device, dtype, sequential_indices=False, seed=42, ): """Run correctness check for a single shape config. Returns True if PASS. Pass/fail is based on OUTPUT comparison only (atol=5e-2). Pool state diff is reported as informational — state divergence over many tokens is expected due to different chunk sizes and accumulation order. """ tag = ( f"B={B:>3} T/seq={T_per_seq:>4} H={H:>2} K={K:>3} V={V:>3} pool={pool_size:>4}" ) idx_tag = " (seq)" if sequential_indices else "" # FlashInfer GDN requires K == V (square head_size) if K != V: print(f" [SKIP] {tag}{idx_tag} (FlashInfer requires K==V)") return True # FlashInfer GDN CUTLASS kernels are only compiled for head_size=128. # Running with other sizes causes illegal memory access that poisons # the CUDA context (unrecoverable), so we must skip upfront. FLASHINFER_SUPPORTED_HEAD_SIZES = {128} if K not in FLASHINFER_SUPPORTED_HEAD_SIZES: print( f" [SKIP] {tag}{idx_tag} (FlashInfer only supports head_size={FLASHINFER_SUPPORTED_HEAD_SIZES})" ) return True inp = make_inputs( B, T_per_seq, H, K, V, pool_size, device, dtype, sequential_indices=sequential_indices, seed=seed, ) o_triton, pool_triton, h_triton = run_triton(inp) # FlashInfer may not support all head_size values (e.g., only 128). # CUDA errors from unsupported configs are often asynchronous, so we # must synchronize inside the try block to catch them here. try: o_fi, pool_fi, _ = run_flashinfer(inp) torch.cuda.synchronize() except Exception as e: # Catch RuntimeError, torch.AcceleratorError, etc. # Reset CUDA error state so subsequent tests can proceed try: torch.cuda.synchronize() except Exception: pass print(f" [SKIP] {tag}{idx_tag} (FlashInfer error: {e})") return True cache_indices = inp["cache_indices"] # --- Output comparison --- # bf16 prefill with L2norm + chunked accumulation torch.testing.assert_close(o_triton, o_fi, atol=5e-2, rtol=1e-2) # --- Stride check --- def strides_ok(pool): s = pool.stride() return s[-2] == 1 and s[-1] == K strides_triton = strides_ok(pool_triton) strides_fi = strides_ok(pool_fi) passed = strides_triton and strides_fi # Build detail string details = [] if not strides_triton: details.append("triton strides bad") if not strides_fi: details.append("flashinfer strides bad") status = "PASS" if passed else "FAIL" detail_str = f" [{', '.join(details)}]" print(f" [{status}] {tag}{idx_tag}") return passed # --------------------------------------------------------------------------- # Benchmark # --------------------------------------------------------------------------- def bench_shape(B, H, T_per_seq, K, V, pool_size, device, dtype): """Benchmark Triton vs FlashInfer for a single config. Requires K == V.""" import triton.testing assert K == V, f"FlashInfer GDN requires K == V, got K={K}, V={V}" T = B * T_per_seq inp = make_inputs(B, T_per_seq, H, K, V, pool_size, device, dtype) # -- Shared read-only tensors -- q, k_t, v = inp["q"], inp["k"], inp["v"] g_triton, beta_triton = inp["g_triton"], inp["beta_triton"] cu_seqlens = inp["cu_seqlens"] cache_indices = inp["cache_indices"] seq_indices = torch.arange(B, dtype=torch.int32, device=device) pool_v = inp["pool_init"] def fn_triton(): pool = make_k_contiguous(pool_v.clone()) triton_chunk_gated_delta_rule( q=q, k=k_t, v=v, g=g_triton, beta=beta_triton, initial_state=pool, initial_state_indices=cache_indices, cu_seqlens=cu_seqlens, head_first=False, use_qk_l2norm_in_kernel=True, ) def fn_flashinfer(): # -- Pre-compute FlashInfer format tensors (outside timing) -- # Pre-normalize q and k (matching sglang production: l2norm_fwd) # q_fi = torch.nn.functional.normalize(q[0].contiguous().float(), p=2.0, dim=-1).to( # dtype # ) # k_fi = torch.nn.functional.normalize(k_t[0].contiguous().float(), p=2.0, dim=-1).to( # dtype # ) q_fi = l2norm_fwd(q[0].contiguous()) k_fi = l2norm_fwd(k_t[0].contiguous()) v_fi = v[0].contiguous() alpha_fi = torch.exp(g_triton[0].to(torch.float32)) beta_fi = beta_triton[0].to(torch.float32) cu_seqlens_fi = cu_seqlens.to(torch.int64) pool = make_k_contiguous(pool_v.clone()) gathered = pool[cache_indices] initial_state = make_k_contiguous(gathered.float().contiguous()) flashinfer_chunk_gated_delta_rule( q=q_fi, k=k_fi, v=v_fi, g=alpha_fi, beta=beta_fi, scale=None, initial_state=initial_state, output_final_state=True, cu_seqlens=cu_seqlens_fi, use_qk_l2norm_in_kernel=False, ) quantiles = [0.5, 0.2, 0.8] # Warmup fn_triton() fn_flashinfer() torch.cuda.synchronize() ms_triton, _, _ = triton.testing.do_bench_cudagraph(fn_triton, quantiles=quantiles) ms_fi, _, _ = triton.testing.do_bench_cudagraph(fn_flashinfer, quantiles=quantiles) # Metrics num_o_heads = H flops = gdn_flops(T, num_o_heads, K, V) mem_bytes = gdn_bytes(T, H, H, K, V, B, dtype) tflops_triton = flops / ms_triton / 1e9 tflops_fi = flops / ms_fi / 1e9 tb_s_triton = mem_bytes / ms_triton / 1e9 tb_s_fi = mem_bytes / ms_fi / 1e9 speedup = ms_triton / ms_fi if ms_fi > 0 else float("inf") print( f" {B:>5} {H:>3} {T_per_seq:>6} {T:>7} | " f"{ms_triton:>8.3f} {tflops_triton:>7.2f} {tb_s_triton:>7.2f} | " f"{ms_fi:>8.3f} {tflops_fi:>7.2f} {tb_s_fi:>7.2f} | " f"{speedup:>7.2f}x" ) # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def run_correctness(device, dtype): print("=" * 78) print("Correctness sweep: Triton vs FlashInfer") print("=" * 78) shapes = [ # (B, T_per_seq, H, K, V, pool_size) # --- baseline (Qwen3-Next style) --- (4, 64, 16, 128, 128, 32), (4, 256, 16, 128, 128, 32), # --- different batch sizes --- (1, 128, 16, 128, 128, 32), (8, 128, 16, 128, 128, 64), (16, 64, 16, 128, 128, 128), (32, 32, 16, 128, 128, 256), # --- different head counts --- (4, 128, 4, 128, 128, 32), (4, 128, 8, 128, 128, 32), (4, 128, 16, 64, 64, 32), (4, 128, 32, 128, 128, 32), (4, 128, 64, 128, 128, 32), # --- short sequences --- (4, 1, 16, 128, 128, 32), (4, 7, 16, 128, 128, 32), (4, 16, 16, 128, 128, 32), # --- large pool (sparse access) --- (4, 128, 16, 128, 128, 512), # --- combined stress --- (32, 128, 32, 128, 128, 256), ] shapes_seq = [ (8, 128, 16, 128, 128, 8), (4, 128, 32, 128, 128, 4), (4, 128, 64, 128, 128, 4), (32, 128, 32, 128, 128, 32), ] all_pass = True for B, T_per_seq, H, K, V, pool_size in shapes: if not check_shape(B, T_per_seq, H, K, V, pool_size, device, dtype): all_pass = False print() print("Sequential-index variants:") for B, T_per_seq, H, K, V, pool_size in shapes_seq: if not check_shape( B, T_per_seq, H, K, V, pool_size, device, dtype, sequential_indices=True, ): all_pass = False print() if all_pass: print("ALL PASSED.") else: print("SOME FAILED.") return all_pass def run_benchmark(device, dtype, args): print() print("=" * 105) print("Benchmark: Triton GDN vs FlashInfer GDN (do_bench_cudagraph)") print("=" * 105) K = args.head_size_k V = args.head_size_v pool_size = args.pool_size if args.preset == "qwen3-next": bench_configs = [ # (B, H, T_per_seq) (4, 16, 256), (4, 32, 256), (16, 16, 256), (16, 32, 256), (32, 16, 256), (32, 32, 256), (64, 16, 256), (64, 32, 256), (128, 16, 256), (128, 32, 256), # longer sequences (4, 16, 1024), (4, 32, 1024), (32, 16, 1024), (32, 32, 1024), ] else: bench_configs = [] for B in args.batch_sizes: for H in args.num_heads: for T_per_seq in args.seq_lens: bench_configs.append((B, H, T_per_seq)) print(f" Config: K={K}, V={V}, pool_size={pool_size}, dtype={dtype}") print( f" {'B':>5} {'H':>3} {'T/seq':>6} {'T_tot':>7} | " f"{'tri(ms)':>8} {'TFLOPS':>7} {'TB/s':>7} | " f"{'fi(ms)':>8} {'TFLOPS':>7} {'TB/s':>7} | " f"{'speedup':>8}" ) print(" " + "-" * 98) for B, H, T_per_seq in bench_configs: actual_pool = max(pool_size, B) bench_shape(B, H, T_per_seq, K, V, actual_pool, device, dtype) def main(): parser = argparse.ArgumentParser( description="Benchmark & Correctness: Triton GDN vs FlashInfer GDN" ) parser.add_argument( "--mode", choices=["all", "correctness", "bench"], default="all", help="Run mode (default: all)", ) parser.add_argument( "--preset", choices=["qwen3-next", "custom"], default="qwen3-next", help="Preset config (default: qwen3-next)", ) parser.add_argument( "--dtype", choices=["float16", "bfloat16"], default="bfloat16", ) parser.add_argument("--head-size-k", type=int, default=128) parser.add_argument("--head-size-v", type=int, default=128) parser.add_argument("--pool-size", type=int, default=256) parser.add_argument( "--batch-sizes", type=int, nargs="+", default=[4, 16, 32, 64, 128], ) parser.add_argument( "--num-heads", type=int, nargs="+", default=[16, 32], ) parser.add_argument( "--seq-lens", type=int, nargs="+", default=[128, 256, 512, 1024], ) args = parser.parse_args() if args.preset == "qwen3-next": args.head_size_k = 128 args.head_size_v = 128 device = "cuda" dtype = getattr(torch, args.dtype) # Check SM version cap = torch.cuda.get_device_capability() dev_name = torch.cuda.get_device_name() print(f"Device: {dev_name} (SM {cap[0]}{cap[1]})") if args.mode in ("all", "correctness"): all_pass = run_correctness(device, dtype) if not all_pass and args.mode == "all": print("\nSkipping benchmark due to correctness failures.") return 1 if args.mode in ("all", "bench"): run_benchmark(device, dtype, args) return 0 if __name__ == "__main__": sys.exit(main()) ================================================ FILE: benchmark/bench_rope/benchmark_rope_index.py ================================================ # This script benchmarks MRotaryEmbedding.get_rope_index_glm4v (GLM4V mrope index builder). # It generates synthetic multimodal input_ids + attention_mask (+ optional image/video grids), # runs benchmarks. # # == Usage Examples == # # python3 benchmark_rope_index.py --device cuda --num-tokens 1024 2048 --benchmark-iter 200 import argparse import math import time from dataclasses import dataclass, field from typing import Any import numpy as np import torch from sglang.srt.layers.rotary_embedding import MRotaryEmbedding # ----------------------------- # Minimal config objects # ----------------------------- @dataclass class DummyVisionConfig: spatial_merge_size: int = 2 @dataclass class DummyHFConfig: image_token_id: int = 32000 video_start_token_id: int = 32001 video_end_token_id: int = 32002 vision_config: DummyVisionConfig = field( default_factory=lambda: DummyVisionConfig(spatial_merge_size=2) ) # ----------------------------- # Helpers # ----------------------------- def calculate_stats(times: list[float]) -> dict[str, float]: """Calculate statistics from a list of times.""" times_array = np.array(times, dtype=np.float64) return { "mean": float(np.mean(times_array)), "median": float(np.median(times_array)), "p99": float(np.percentile(times_array, 99)), "min": float(np.min(times_array)), "max": float(np.max(times_array)), } def _sync(device: torch.device): if device.type == "cuda": torch.cuda.synchronize() def _approx_hw(patches: int, merge: int) -> tuple[int, int]: # want (h/merge)*(w/merge) ~= patches gh = int(math.sqrt(max(1, patches))) gw = max(1, patches // max(1, gh)) return gh * merge, gw * merge def generate_test_data( num_tokens: int, batch_size: int, hf_config: DummyHFConfig, dtype: torch.dtype, device: torch.device, pad_ratio: float, num_images_per_sample: int, image_patch_tokens: int, num_videos_per_sample: int, video_patch_tokens: int, seed: int, ): """ Generate synthetic (input_ids, attention_mask, image_grid_thw, video_grid_thw). NOTE: - image_grid_thw / video_grid_thw are global lists across the entire batch in encounter order, matching the function's image_index/video_index behavior. - image patches are represented by repeated image_token_id. - video patches are represented by image_token_id wrapped with start/end tokens. """ torch.manual_seed(seed) forbidden = { 0, hf_config.image_token_id, hf_config.video_start_token_id, hf_config.video_end_token_id, } vocab_size = 50000 def rand_text(n: int) -> torch.Tensor: # generate random ids not in forbidden out = torch.randint(1, vocab_size, (n,), device=device, dtype=torch.long) # fix forbidden by +1 until ok (cheap, deterministic enough for benchmark data) for bad in forbidden: out = torch.where(out == bad, out + 1, out) return out image_grids: list[list[int]] = [] video_grids: list[list[int]] = [] input_ids = torch.zeros((batch_size, num_tokens), device=device, dtype=torch.long) attention_mask = torch.zeros( (batch_size, num_tokens), device=device, dtype=torch.long ) eff_len = int(round(num_tokens * (1.0 - pad_ratio))) eff_len = max(1, min(num_tokens, eff_len)) min_needed = 1 min_needed += num_images_per_sample * image_patch_tokens min_needed += num_videos_per_sample * (2 + video_patch_tokens) if eff_len < min_needed: num_images_per_sample = 0 num_videos_per_sample = 0 for b in range(batch_size): blocks: list[torch.Tensor] = [] reserved = ( num_images_per_sample * image_patch_tokens + num_videos_per_sample * (2 + video_patch_tokens) ) reserved = min(reserved, max(0, eff_len - 1)) text_budget = max(1, eff_len - reserved) n_text_chunks = num_images_per_sample + num_videos_per_sample + 1 base = text_budget // n_text_chunks rem = text_budget % n_text_chunks text_chunks = [base + (1 if i < rem else 0) for i in range(n_text_chunks)] tci = 0 for _ in range(num_images_per_sample): blocks.append(rand_text(text_chunks[tci])) tci += 1 blocks.append( torch.full( (image_patch_tokens,), hf_config.image_token_id, device=device, dtype=torch.long, ) ) h, w = _approx_hw( image_patch_tokens, hf_config.vision_config.spatial_merge_size ) image_grids.append([1, h, w]) for _ in range(num_videos_per_sample): blocks.append(rand_text(text_chunks[tci])) tci += 1 blocks.append( torch.tensor( [hf_config.video_start_token_id], device=device, dtype=torch.long ) ) blocks.append( torch.full( (video_patch_tokens,), hf_config.image_token_id, device=device, dtype=torch.long, ) ) blocks.append( torch.tensor( [hf_config.video_end_token_id], device=device, dtype=torch.long ) ) h, w = _approx_hw( video_patch_tokens, hf_config.vision_config.spatial_merge_size ) # first field = group count used by code; set to 1 video_grids.append([1, h, w]) blocks.append(rand_text(text_chunks[tci])) tokens = torch.cat(blocks, dim=0)[:eff_len] pad = torch.zeros( (num_tokens - tokens.numel(),), device=device, dtype=torch.long ) ids = torch.cat([tokens, pad], dim=0) mask = torch.cat( [ torch.ones((tokens.numel(),), device=device, dtype=torch.long), torch.zeros( (num_tokens - tokens.numel(),), device=device, dtype=torch.long ), ], dim=0, ) input_ids[b] = ids attention_mask[b] = mask image_grid_thw = ( torch.tensor(image_grids, device=device, dtype=torch.long) if len(image_grids) else None ) video_grid_thw = ( torch.tensor(video_grids, device=device, dtype=torch.long) if len(video_grids) else None ) return ( input_ids.to(dtype=torch.long), attention_mask.to(dtype=torch.long), image_grid_thw, video_grid_thw, ) def benchmark_rope_index( model_name: str, tp_size: int, num_tokens: int, batch_size: int, pad_ratio: float, spatial_merge_size: int, num_images: int, image_patch_tokens: int, num_videos: int, video_patch_tokens: int, dtype: torch.dtype, seed: int, warmup_iter: int, benchmark_iter: int, device: torch.device, ): torch.manual_seed(seed) hf_config = DummyHFConfig( image_token_id=32000, video_start_token_id=32001, video_end_token_id=32002, vision_config=DummyVisionConfig(spatial_merge_size=spatial_merge_size), ) print(80 * "=") print( f"Evaluating: {model_name} tp_size={tp_size} " f"num_tokens={num_tokens} batch={batch_size} pad_ratio={pad_ratio} " f"images/sample={num_images} image_patch_tokens={image_patch_tokens} " f"videos/sample={num_videos} video_patch_tokens={video_patch_tokens} " f"dtype={dtype} device={device}" ) input_ids, attention_mask, image_grid_thw, video_grid_thw = generate_test_data( num_tokens=num_tokens, batch_size=batch_size, hf_config=hf_config, dtype=dtype, device=device, pad_ratio=pad_ratio, num_images_per_sample=num_images, image_patch_tokens=image_patch_tokens, num_videos_per_sample=num_videos, video_patch_tokens=video_patch_tokens, seed=seed, ) # Smoke test has_mm = (image_grid_thw is not None) or (video_grid_thw is not None) if has_mm: pos, delta = MRotaryEmbedding.get_rope_index_glm4v( input_ids=input_ids, hf_config=hf_config, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, ) assert pos.shape == (3, batch_size, num_tokens) assert delta.shape == (batch_size, 1) # Warm up for _ in range(warmup_iter): if has_mm: MRotaryEmbedding.get_rope_index_glm4v( input_ids=input_ids, hf_config=hf_config, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, ) MRotaryEmbedding.get_rope_index_glm4v( input_ids=input_ids, hf_config=hf_config, image_grid_thw=None, video_grid_thw=None, attention_mask=attention_mask, ) _sync(device) # Time multimodal branch multimodal_times = [] for _ in range(benchmark_iter): _sync(device) start = time.time() MRotaryEmbedding.get_rope_index_glm4v( input_ids=input_ids, hf_config=hf_config, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, ) _sync(device) multimodal_times.append(time.time() - start) # Time fallback branch fallback_times = [] for _ in range(benchmark_iter): _sync(device) start = time.time() MRotaryEmbedding.get_rope_index_glm4v( input_ids=input_ids, hf_config=hf_config, image_grid_thw=None, video_grid_thw=None, attention_mask=attention_mask, ) _sync(device) fallback_times.append(time.time() - start) multimodal_stats = calculate_stats(multimodal_times) fallback_stats = calculate_stats(fallback_times) print(f"\nPerformance for config (B={batch_size}, T={num_tokens}):") print( f"Multimodal: mean={multimodal_stats['mean']:.8f}s, " f"median={multimodal_stats['median']:.8f}s, " f"p99={multimodal_stats['p99']:.8f}s" ) print( f"Fallback: mean={fallback_stats['mean']:.8f}s, " f"median={fallback_stats['median']:.8f}s, " f"p99={fallback_stats['p99']:.8f}s" ) if has_mm: speedup = ( multimodal_stats["mean"] / fallback_stats["mean"] if fallback_stats["mean"] > 0 else float("inf") ) print(f"Fallback Speedup over Multimodal: {speedup:.8f}x") else: speedup = float("nan") print( "[INFO] num_tokens too small for multimodal segments; skip multimodal benchmark." ) print(f"Fallback Speedup over Multimodal: {speedup:.8f}x") return multimodal_stats, fallback_stats, speedup if __name__ == "__main__": parser = argparse.ArgumentParser( description="Benchmark GLM4V get_rope_index_glm4v." ) parser.add_argument("--model-name", type=str, default="GLM4V") parser.add_argument("--tp-size", type=int, default=1) parser.add_argument( "--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu" ) parser.add_argument("--warmup-iter", type=int, default=10) parser.add_argument("--benchmark-iter", type=int, default=100) parser.add_argument("--dtype", type=str, choices=["int64"], default="int64") parser.add_argument("--seed", type=int, default=0) # token length sweep parser.add_argument("--num-tokens", type=int, nargs="+", required=False) # data shape knobs parser.add_argument("--batch-size", type=int, default=1) parser.add_argument("--pad-ratio", type=float, default=0.0) parser.add_argument("--spatial-merge-size", type=int, default=2) parser.add_argument("--num-images", type=int, default=1) parser.add_argument("--image-patch-tokens", type=int, default=256) parser.add_argument("--num-videos", type=int, default=1) parser.add_argument("--video-patch-tokens", type=int, default=256) # output parser.add_argument("--out-dir", type=str, default=".") args = parser.parse_args() print(args) device = torch.device(args.device) if args.num_tokens is None: num_tokens_list = [2**i for i in range(0, 18)] else: num_tokens_list = args.num_tokens rows: list[dict[str, Any]] = [] for num_tokens in num_tokens_list: multimodal_stats, fallback_stats, speedup = benchmark_rope_index( model_name=args.model_name, tp_size=args.tp_size, num_tokens=num_tokens, batch_size=args.batch_size, pad_ratio=args.pad_ratio, spatial_merge_size=args.spatial_merge_size, num_images=args.num_images, image_patch_tokens=args.image_patch_tokens, num_videos=args.num_videos, video_patch_tokens=args.video_patch_tokens, dtype=getattr(torch, args.dtype), seed=args.seed, warmup_iter=args.warmup_iter, benchmark_iter=args.benchmark_iter, device=device, ) ================================================ FILE: benchmark/benchmark_batch/benchmark_batch.py ================================================ import concurrent.futures import os import random import time from concurrent.futures import ProcessPoolExecutor from statistics import mean import requests from tqdm import tqdm from transformers import AutoTokenizer from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint ############################################################################### # CONFIG ############################################################################### ENDPOINT_URL = "http://127.0.0.1:30000" TOKENIZER_DIR = "/models/meta-llama/Llama-3.2-3B" # Benchmark configurations NUM_REQUESTS = 10 # Total number of requests (each with BATCH_SIZE prompts) NUM_TOKENS = 32000 # Tokens per prompt BATCH_SIZE = 8 # Number of prompts per request GEN_TOKENS = 0 # Tokens to generate per prompt ############################################################################### # REQUEST GENERATION (in parallel) ############################################################################### def generate_random_prompt(index, tokenizer_dir, num_tokens): """Generate a single random prompt with specified token count.""" tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir) vocab_size = tokenizer.vocab_size def generate_random_text(num_toks): random_token_ids = [random.randint(0, vocab_size - 1) for _ in range(num_toks)] return tokenizer.decode(random_token_ids, clean_up_tokenization_spaces=True) random_text = generate_random_text(num_tokens) return f"Prompt {index}: {random_text}" def prepare_all_prompts(num_requests, batch_size, num_tokens, tokenizer_dir): """Generate prompts for all requests in parallel.""" total_prompts = num_requests * batch_size all_prompts = [None] * total_prompts max_workers = min(os.cpu_count() or 1, total_prompts) with ProcessPoolExecutor(max_workers=max_workers) as executor: futures = [ executor.submit(generate_random_prompt, i, tokenizer_dir, num_tokens) for i in range(total_prompts) ] for future in tqdm( concurrent.futures.as_completed(futures), total=total_prompts, desc="Generating prompts", ): index = futures.index(future) all_prompts[index] = future.result() batched_prompts = [ all_prompts[i * batch_size : (i + 1) * batch_size] for i in range(num_requests) ] print( f"Generated {total_prompts} prompts with {num_tokens} tokens each, grouped into {num_requests} requests of {batch_size} prompts.\n" ) return batched_prompts ############################################################################### # HTTP CALLS ############################################################################### def send_batch_request(endpoint, prompts, gen_tokens, request_id): """Send a batch of prompts to the /generate endpoint synchronously.""" sampling_params = { "max_new_tokens": gen_tokens, "temperature": 0.7, "stop": "\n", } data = {"text": prompts, "sampling_params": sampling_params} start_time = time.perf_counter() try: response = requests.post( endpoint.base_url + "/generate", json=data, timeout=3600 ) if response.status_code != 200: error = response.json() raise RuntimeError(f"Request {request_id} failed: {error}") result = response.json() elapsed_time = (time.perf_counter() - start_time) * 1000 # Convert to ms avg_per_prompt = elapsed_time / len(prompts) if prompts else 0 return request_id, elapsed_time, avg_per_prompt, True, len(prompts) except Exception as e: print(f"[Request] Error for request {request_id}: {e}") return request_id, 0, 0, False, len(prompts) def run_benchmark(endpoint, batched_prompts, batch_size, gen_tokens): """Run the benchmark sequentially.""" results = [] num_requests = len(batched_prompts) # Record start time for total latency benchmark_start_time = time.perf_counter() for i, batch_prompts in enumerate(batched_prompts): request_id = i + 1 assert ( len(batch_prompts) == batch_size ), f"Request {request_id} should have {batch_size} prompts, got {len(batch_prompts)}" print( f"[Request] Sending request {request_id}/{num_requests} with {len(batch_prompts)} prompts at {int(time.time()*1000)}" ) result = send_batch_request(endpoint, batch_prompts, gen_tokens, request_id) results.append(result) # Calculate total latency total_latency = (time.perf_counter() - benchmark_start_time) * 1000 # Convert to ms return results, total_latency ############################################################################### # RESULTS ############################################################################### def process_results(results, total_latency, num_requests): """Process and display benchmark results.""" total_time = 0 successful_requests = 0 failed_requests = 0 request_latencies = [] per_prompt_latencies = [] total_prompts = 0 for request_id, elapsed_time, avg_per_prompt, success, batch_size in results: if success: successful_requests += 1 total_prompts += batch_size request_latencies.append(elapsed_time) per_prompt_latencies.append(avg_per_prompt) total_time += elapsed_time / 1000 # Convert to seconds else: failed_requests += 1 avg_request_latency = mean(request_latencies) if request_latencies else 0 avg_per_prompt_latency = mean(per_prompt_latencies) if per_prompt_latencies else 0 throughput = total_prompts / total_time if total_time > 0 else 0 print("\nBenchmark Summary:") print(f" Total requests sent: {len(results)}") print(f" Total prompts sent: {total_prompts}") print(f" Successful requests: {successful_requests}") print(f" Failed requests: {failed_requests}") print(f" Total latency (all requests): {total_latency:.2f} ms") print(f" Avg per request latency: {avg_request_latency:.2f} ms") print(f" Avg per prompt latency: {avg_per_prompt_latency:.2f} ms") print(f" Throughput: {throughput:.2f} prompts/second\n") ############################################################################### # MAIN ############################################################################### def main(): # Initialize endpoint endpoint = RuntimeEndpoint(ENDPOINT_URL) # Generate prompts batched_prompts = prepare_all_prompts( NUM_REQUESTS, BATCH_SIZE, NUM_TOKENS, TOKENIZER_DIR ) # Flush cache before benchmark # endpoint.flush_cache() # Run benchmark print( f"Starting benchmark: NUM_TOKENS={NUM_TOKENS}, BATCH_SIZE={BATCH_SIZE}, NUM_REQUESTS={NUM_REQUESTS}\n" ) results, total_latency = run_benchmark( endpoint, batched_prompts, BATCH_SIZE, GEN_TOKENS ) # Process and display results process_results(results, total_latency, NUM_REQUESTS) if __name__ == "__main__": random.seed(0) main() ================================================ FILE: benchmark/benchmark_batch/benchmark_tokenizer.py ================================================ import argparse import random import time from statistics import mean from transformers import AutoTokenizer from sglang.srt.utils.patch_tokenizer import patch_tokenizer def main(): args = parse_args() print("Tokenizer Benchmark: Sequential vs Batch Processing") print("-" * 60) print(f"Tokenizer: {args.tokenizer}") print(f"Functions: {', '.join(args.function)}") print(f"Tokens per prompt: {args.num_tokens}") print(f"Number of runs per batch size: {args.num_runs}") print(f"Batch mode: {', '.join(args.batch_mode)}") print("-" * 60) tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) tokenizer = patch_tokenizer(tokenizer) max_batch_size = max(args.batch_sizes) token_ids = generate_random_token_ids( num_prompts=max_batch_size, num_tokens=args.num_tokens, tokenizer=tokenizer ) if "encode" in args.function: prompts = [ tokenizer.decode(ids, clean_up_tokenization_spaces=True) for ids in token_ids ] run_benchmark( name="encode", data=prompts, sequential_fn=lambda batch: [tokenizer.encode(p) for p in batch], batch_fn=lambda batch: tokenizer(batch), batch_sizes=args.batch_sizes, num_runs=args.num_runs, batch_mode=args.batch_mode, ) if "decode" in args.function: # mimic DetokenizerManager's usual case decode_kwargs = dict( skip_special_tokens=True, spaces_between_special_tokens=True, ) run_benchmark( name="decode", data=token_ids, sequential_fn=lambda batch: [ tokenizer.decode(ids, **decode_kwargs) for ids in batch ], batch_fn=lambda batch: tokenizer.batch_decode(batch, **decode_kwargs), batch_sizes=args.batch_sizes, num_runs=args.num_runs, batch_mode=args.batch_mode, ) def run_benchmark( *, name, data, sequential_fn, batch_fn, batch_sizes, num_runs, batch_mode ): print("\n" + "=" * 60) print(f"{name.upper()} BENCHMARK") print("=" * 60) results = [ benchmark( data=data, batch_size=bs, sequential_fn=sequential_fn, batch_fn=batch_fn, num_runs=num_runs, batch_mode=batch_mode, ) for bs in batch_sizes ] print_results(results=results, func_name=name, batch_mode=batch_mode) def benchmark(*, data, batch_size, sequential_fn, batch_fn, num_runs, batch_mode): batch_data = data[:batch_size] run_single = "single" in batch_mode run_batch = "batch" in batch_mode out = {"batch_size": batch_size} if run_single: sequential_times = measure_times( fn=lambda: sequential_fn(batch_data), num_runs=num_runs ) out |= { "avg_sequential_ms": mean(sequential_times), "sequential_runs": sequential_times, } if run_batch: batch_times = measure_times(fn=lambda: batch_fn(batch_data), num_runs=num_runs) out |= { "avg_batch_ms": mean(batch_times), "batch_runs": batch_times, } if run_single and run_batch: out["speedup_factor"] = ( out["avg_sequential_ms"] / out["avg_batch_ms"] if out["avg_batch_ms"] > 0 else 0 ) return out def print_results(*, results, func_name, batch_mode): run_single = "single" in batch_mode run_batch = "batch" in batch_mode for r in results: print(f"\nBatch size: {r['batch_size']}") if run_single: print_runs( label=f"Sequential {func_name}", runs=r["sequential_runs"], avg=r["avg_sequential_ms"], ) if run_batch: print_runs( label=f"Batch {func_name}", runs=r["batch_runs"], avg=r["avg_batch_ms"] ) if run_single and run_batch: print(f" Speedup factor: {r['speedup_factor']:.2f}x") print("\n" + "=" * 60) print(f"SUMMARY: {func_name.upper()}") print("=" * 60) headers = ["Batch Size"] if run_single: headers.append("Sequential (ms)") if run_batch: headers.append("Batch (ms)") if run_single and run_batch: headers.append("Speedup") print("".join(f"{h:<18}" for h in headers)) print("-" * (18 * len(headers))) for r in results: row = [f"{r['batch_size']}"] if run_single: row.append(f"{r['avg_sequential_ms']:.2f} ms") if run_batch: row.append(f"{r['avg_batch_ms']:.2f} ms") if run_single and run_batch: row.append(f"{r['speedup_factor']:.2f}x") print("".join(f"{v:<18}" for v in row)) def print_runs(*, label, runs, avg): print(f" {label}:") for i, t in enumerate(runs): print(f" Run {i+1}: {t:.2f} ms") print(f" Average: {avg:.2f} ms") def measure_times(*, fn, num_runs): times = [] for _ in range(num_runs): start = time.perf_counter() fn() times.append((time.perf_counter() - start) * 1000) return times def generate_random_token_ids(*, num_prompts, num_tokens, tokenizer): vocab_size = tokenizer.vocab_size print(f"Generating {num_prompts} random sequences with {num_tokens} tokens each...") return [ [random.randint(0, vocab_size - 1) for _ in range(num_tokens)] for _ in range(num_prompts) ] def parse_args(): parser = argparse.ArgumentParser( description="Tokenizer Benchmark: Sequential vs Batch Processing" ) parser.add_argument( "--tokenizer", type=str, required=True, help="Tokenizer name or path (e.g. nvidia/Kimi-K2-Thinking-NVFP4)", ) parser.add_argument( "--function", type=str, nargs="+", choices=["encode", "decode"], default=["encode", "decode"], help="Functions to benchmark (default: encode decode)", ) parser.add_argument( "--num-tokens", type=int, default=20000, help="Number of tokens per prompt (default: 20000)", ) parser.add_argument( "--batch-sizes", type=int, nargs="+", default=[1, 2, 4, 8], help="Batch sizes to test (default: 1 2 4 8)", ) parser.add_argument( "--batch-mode", nargs="+", choices=["single", "batch"], default=["single", "batch"], help="Benchmark modes to run (default: single batch)", ) parser.add_argument( "--num-runs", type=int, default=5, help="Number of runs per batch size (default: 5)", ) return parser.parse_args() if __name__ == "__main__": random.seed(0) main() ================================================ FILE: benchmark/benchmark_vllm_060/README.md ================================================ ## How to reproduce the benchmark results for SGLang v0.3.0 compared to vLLM v0.6.0 In short, with multi step enabled, in online scenarios that we benchmarked, the Median TTFT of vLLM is **3 times** that of SGLang, and the Median ITL is **10 times** that of SGLang. Lower Median TTFT and ITL are better. vLLM's multi-step optimization did not improve throughput while ensuring lower Median TTFT and ITL. Also, under maximum throughput benchmark, if vLLM does not set gpu util to 0.95 separately and uses the default configuration instead, its maximum throughput is **lower** than that of SGLang. ## Online benchmark results ### Llama 3.1 8B Instruct 1 x A100 80G | RPS | Num prompts | Engine | Median E2E Latency | Median TTFT | Median TPOT | Median ITL | |------|-------------|--------|--------------------|-------------|-------------|------------| | 4 | 1200 | SGLang | 1564.17 | **31.98** | 13.17 | **11.93** | | 4 | 1200 | vLLM | 1691.97 | **100.48** | 14.14 | **129.32** | | 8 | 2400 | SGLang | 2175.02 | **35.68** | 17.85 | **14.41** | | 8 | 2400 | vLLM | 2137.16 | **120.39** | 17.09 | **158.63** | ### Llama 3.1 70B Insruct 4 x H100 80G | RPS | Num Prompts | Engine | Median E2E Latency | Median TTFT | Median TPOT | Median ITL | |------|-------------|--------|--------------------|-------------|-------------|------------| | 4 | 1200 | SGLang | 3005.24 | **53.94** | 25.03 | **21.67** | | 4 | 1200 | vLLM | 2915.60 | **179.15** | 23.58 | **231.23** | | 8 | 2400 | SGLang | 4064.98 | **58.11** | 33.07 | **24.45** | | 8 | 2400 | vLLM | 3752.38 | **207.12** | 29.15 | **275.32** | ## Offline benchmark results ### Llama 3.1 8B Instruct 1 x A100 80G | RPS | Num Prompts | Engine | Request throughput | Output token throughput | |------|-------------|--------|--------------------|-------------------------| | inf | 5000 | SGLang | 22.03 | **4281.51** | | inf | 5000 | vLLM | 21.27 | **4132.37** | ### Llama 3.1 70B Insruct 4 x H100 80G | RPS | Num Prompts | Engine | Request throughput | Output token throughput | |------|-------------|--------|--------------------|-------------------------| | inf | 5000 | SGLang | 19.84 | **3856.01** | | inf | 5000 | vLLM | 19.04 | **3700.64** | ## Installation ```bash # install sglang v0.3.0 pip install --upgrade pip pip install "sglang[all]"==0.3.0 pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.4/ # install vllm v0.6.0 pip install vllm==0.6.0 ``` ## Notes We referred to the reproduction method in https://github.com/vllm-project/vllm/issues/8176, and added the `--num-scheduler-steps 10` parameter when starting the vLLM server. The `gpu_memory_utilization` of vLLM is by default 0.9 at both TP 1 and TP 4, while SGLang's `mem_frac` is 0.88 at TP 1 and 0.85 at TP 4, so we manually set it to 0.88 at TP 4. ## Online benchmarks ```bash # Llama 3.1 8B Instruct on 1 x A100 python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --enable-torch-compile --disable-radix-cache python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.1-8B-Instruct --disable-log-requests --num-scheduler-steps 10 --max_model_len 4096 # Llama 3.1 70B Instruct on 4 x H100 python -m sglang.launch_server --model-path meta-llama/Llama-3.1-70B-Instruct --disable-radix-cache --tp 4 python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.1-70B-Instruct --disable-log-requests --num-scheduler-steps 10 --tensor 4 --max_model_len 4096 # bench serving python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompts 1200 --request-rate 4 python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompts 2400 --request-rate 8 python3 -m sglang.bench_serving --backend vllm --dataset-name sharegpt --num-prompts 1200 --request-rate 4 python3 -m sglang.bench_serving --backend vllm --dataset-name sharegpt --num-prompts 2400 --request-rate 8 ``` ## Offline benchmarks ```bash # Llama 3.1 8B Instruct on 1 x A100 python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --enable-torch-compile --disable-radix-cache python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.1-8B-Instruct --disable-log-requests --num-scheduler-steps 10 --max_model_len 4096 # Llama 3.1 70B Instruct on 4 x H100 python -m sglang.launch_server --model-path meta-llama/Llama-3.1-70B-Instruct --disable-radix-cache --tp 4 --mem-frac 0.88 python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.1-70B-Instruct --disable-log-requests --num-scheduler-steps 10 --tensor 4 --max_model_len 4096 # bench serving python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompts 5000 python3 -m sglang.bench_serving --backend vllm --dataset-name sharegpt --num-prompts 5000 ``` ================================================ FILE: benchmark/blog_v0_2/405b_sglang.sh ================================================ # Create dummy weights: # 1. Create a folder `~/llama-3.1-405b-fp8-dummy` and create `config.json` and tokenizer under this folder. # 2. Get `config.json`` from ./config.md # 3. Download the tokenizer # wget https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8/resolve/main/tokenizer.json # wget https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8/resolve/main/tokenizer_config.json # Launch sglang # python -m sglang.launch_server --model-path ~/llama-3.1-405b-fp8-dummy/ --load-format dummy --tp 8 --quantization fp8 --disable-radix --mem-frac 0.87 # offline python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 3000 --random-input 1024 --random-output 1024 > sglang_log11 python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 4000 --random-input 1024 --random-output 512 > sglang_log12 python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 800 --random-input 4096 --random-output 2048 > sglang_log13 python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 1500 --random-input 4096 --random-output 1024 > sglang_log14 python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 6000 --random-input 256 --random-output 512 > sglang_log15 python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompt 2000 > sglang_log21 # online python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 300 --request-rate 1 --random-input 1024 --random-output 1024 > sglang_log31 python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 600 --request-rate 2 --random-input 1024 --random-output 1024 > sglang_log32 python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 1200 --request-rate 4 --random-input 1024 --random-output 1024 > sglang_log33 python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 2400 --request-rate 8 --random-input 1024 --random-output 1024 > sglang_log34 python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 3200 --request-rate 16 --random-input 1024 --random-output 1024 > sglang_log35 ================================================ FILE: benchmark/blog_v0_2/405b_trt.sh ================================================ # Launch trtllm # https://github.com/sgl-project/tensorrt-demo # offline python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 3000 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log11 python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 4000 --random-input 1024 --random-output 512 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log12 python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 800 --random-input 4096 --random-output 2048 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log13 python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 1500 --random-input 4096 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log14 python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 6000 --random-input 256 --random-output 512 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log15 python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name sharegpt --num-prompt 2000 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log21 # online python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 300 --request-rate 1 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log31 python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 600 --request-rate 2 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log32 python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 1200 --request-rate 4 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log33 python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 2400 --request-rate 8 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log34 python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 3200 --request-rate 16 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log35 ================================================ FILE: benchmark/blog_v0_2/405b_vllm.sh ================================================ # Create dummy weights: # 1. Create a folder `~/llama-3.1-405b-fp8-dummy` and create `config.json` and tokenizer under this folder. # 2. Get `config.json`` from ./config.md # 3. Download the tokenizer # wget https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8/resolve/main/tokenizer.json # wget https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8/resolve/main/tokenizer_config.json # Launch vllm # python3 -m vllm.entrypoints.openai.api_server --model ~/llama-3.1-405b-fp8-dummy/ --load-format dummy --disable-log-requests --tensor-parallel-size 8 --max-model-len 10000 # offline python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 3000 --random-input 1024 --random-output 1024 > vllm_log11 python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 4000 --random-input 1024 --random-output 512 > vllm_log12 python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 800 --random-input 4096 --random-output 2048 > vllm_log13 python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 1500 --random-input 4096 --random-output 1024 > vllm_log14 python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 6000 --random-input 256 --random-output 512 > vllm_log15 python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name sharegpt --num-prompt 2000 > vllm_log21 # online python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 300 --request-rate 1 --random-input 1024 --random-output 1024 > vllm_log31 python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 600 --request-rate 2 --random-input 1024 --random-output 1024 > vllm_log32 python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 1200 --request-rate 4 --random-input 1024 --random-output 1024 > vllm_log33 python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 2400 --request-rate 8 --random-input 1024 --random-output 1024 > vllm_log34 python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 3200 --request-rate 16 --random-input 1024 --random-output 1024 > vllm_log35 ================================================ FILE: benchmark/blog_v0_2/README.md ================================================ # How to reproduce the benchmark results of SGLang ## Prerequisite ### Install the latest SGLang ```bash git clone https://github.com/sgl-project/sglang.git cd sglang git checkout v0.2.7 pip install --upgrade pip pip install -e "python[all]" pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.3/ ``` ### Set up ulimit and HF_TOKEN ```bash ulimit -n 65535 # Change the token to a real and usable one, with access permissions for the Llama 3 models. export HF_TOKEN=hf_token ``` ### Launch the server ```bash # Meta-Llama-3.1-8B-Instruct python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --enable-torch-compile --disable-radix-cache # Meta-Llama-3.1-70B-Instruct python -m sglang.launch_server --model-path meta-llama/Llama-3.1-70B-Instruct --disable-radix-cache --tp 8 # Meta-Llama-3-70B-Instruct-FP8 python -m sglang.launch_server --model-path neuralmagic/Meta-Llama-3-70B-Instruct-FP8 --disable-radix-cache --tp 8 ``` ## Benchmark ### Hardware Requirements - 8B models: Single NVIDIA A100 80GB GPU - 70B models: 8 x NVIDIA A100 80GB GPUs with Tensor Parallelism (TP) 8 - 70B FP8 models: 8 x NVIDIA H100 GPUs with Tensor Parallelism (TP) 8 Please ensure you have the appropriate hardware before running the benchmarks. #### Offline benchmark ```bash python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 4000 --random-input 1024 --random-output 1024 --output-file offline.jsonl python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 5000 --random-input 1024 --random-output 512 --output-file offline.jsonl python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 1000 --random-input 4096 --random-output 2048 --output-file offline.jsonl python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 2000 --random-input 4096 --random-output 1024 --output-file offline.jsonl python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 6000 --random-input 256 --random-output 512 --output-file offline.jsonl python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompts 3000 --output-file offline.jsonl cat offline.jsonl | cut -d':' -f12 | cut -d',' -f1 ``` #### Online benchmark ```bash python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 300 --request-rate 1 --output-file online.jsonl python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 600 --request-rate 2 --output-file online.jsonl python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 1200 --request-rate 4 --output-file online.jsonl python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 2400 --request-rate 8 --output-file online.jsonl python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 3200 --request-rate 16 --output-file online.jsonl cat online.jsonl | cut -d':' -f9 | cut -d',' -f1 ``` ## Other We tried using vLLM 0.5.3.post1, but it often crashes under high loads, and it seems to have similar or worse performance compared to vLLM 0.5.2 from our partial benchmarking, so we are using the older version, vLLM 0.5.2. Preparation for TensorRT LLM can refer to https://github.com/sgl-project/tensorrt-demo. Specifically, we used a batch size of 512, a max input length of 8192, and a max number of tokens of 8192. The instance count for preprocessing and postprocessing in Triton Server is 16. ```bash # vLLM pip install vllm==0.5.2 pip install jsonschema==4.21.1 # Meta-Llama-3-8B-Instruct python -m vllm.entrypoints.openai.api_server --model meta-llama/Meta-Llama-3-8B-Instruct --disable-log-requests # meta-llama/Meta-Llama-3-70B-Instruct python -m vllm.entrypoints.openai.api_server --model meta-llama/Meta-Llama-3-70B-Instruct --disable-log-requests --tensor 8 # neuralmagic/Meta-Llama-3-70B-Instruct-FP8 python -m vllm.entrypoints.openai.api_server --model neuralmagic/Meta-Llama-3-70B-Instruct-FP8 --disable-log-requests --tensor 8 ``` ```bash wget https://raw.githubusercontent.com/sgl-project/sglang/main/python/sglang/bench_serving.py ``` ```bash # vLLM Offline python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 4000 --random-input 1024 --random-output 1024 --output-file offline_vllm.jsonl python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 5000 --random-input 1024 --random-output 512 --output-file offline_vllm.jsonl python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 1000 --random-input 4096 --random-output 2048 --output-file offline_vllm.jsonl python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 2000 --random-input 4096 --random-output 1024 --output-file offline_vllm.jsonl python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 6000 --random-input 256 --random-output 512 --output-file offline_vllm.jsonl python3 bench_serving.py --backend vllm --dataset-name sharegpt --num-prompts 3000 --output-file offline_vllm.jsonl cat offline_vllm.jsonl | cut -d':' -f12 | cut -d',' -f1 ``` ```bash # vLLM Online python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 300 --request-rate 1 --output-file online_vllm.jsonl python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 600 --request-rate 2 --output-file online_vllm.jsonl python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 1200 --request-rate 4 --output-file online_vllm.jsonl python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 2400 --request-rate 8 --output-file online_vllm.jsonl python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 3200 --request-rate 16 --output-file online_vllm.jsonl cat online_vllm.jsonl | cut -d':' -f9 | cut -d',' -f1 ``` ```bash # TensorRT LLM Offline 8B python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --num-prompts 4000 --random-input 1024 --random-output 1024 --output-file offline_trt_8b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --num-prompts 5000 --random-input 1024 --random-output 512 --output-file offline_trt_8b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --num-prompts 1000 --random-input 4096 --random-output 2048 --output-file offline_trt_8b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --num-prompts 2000 --random-input 4096 --random-output 1024 --output-file offline_trt_8b.jsonl python3 bench_serving.py --backend trt --dataset-name random --num-prompts 6000 --random-input 256 --random-output 512 --output-file offline_trt_8b.jsonl --model meta-llama/Meta-Llama-3-8B-Instruct python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name sharegpt --num-prompts 3000 --output-file offline_trt_8b.jsonl cat offline_trt_8b.jsonl | cut -d':' -f12 | cut -d',' -f1 ``` ```bash # TensorRT LLM Online 8B python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 300 --request-rate 1 --output-file online_trt_8b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 600 --request-rate 2 --output-file online_trt_8b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 1200 --request-rate 4 --output-file online_trt_8b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 2400 --request-rate 8 --output-file online_trt_8b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 3200 --request-rate 16 --output-file online_trt_8b.jsonl cat online_trt_8b.jsonl | cut -d':' -f9 | cut -d',' -f1 ``` ```bash # TensorRT LLM Offline 70B python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --num-prompts 4000 --random-input 1024 --random-output 1024 --output-file offline_trt_70b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --num-prompts 5000 --random-input 1024 --random-output 512 --output-file offline_trt_70b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --num-prompts 1000 --random-input 4096 --random-output 2048 --output-file offline_trt_70b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --num-prompts 2000 --random-input 4096 --random-output 1024 --output-file offline_trt_70b.jsonl python3 bench_serving.py --backend trt --dataset-name random --num-prompts 6000 --random-input 256 --random-output 512 --output-file offline_trt_70b.jsonl --model meta-llama/Meta-Llama-3-70B-Instruct python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name sharegpt --num-prompts 3000 --output-file offline_trt_70b.jsonl cat offline_trt_70b.jsonl | cut -d':' -f12 | cut -d',' -f1 ``` ```bash # TensorRT LLM Online 70B python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 300 --request-rate 1 --output-file online_trt_70b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 600 --request-rate 2 --output-file online_trt_70b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 1200 --request-rate 4 --output-file online_trt_70b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 2400 --request-rate 8 --output-file online_trt_70b.jsonl python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 3200 --request-rate 16 --output-file online_trt_70b.jsonl cat online_trt_70b.jsonl | cut -d':' -f9 | cut -d',' -f1 ``` ================================================ FILE: benchmark/blog_v0_2/config.md ================================================ ### used for TensorRT LLM ``` { "architecture": "LlamaForCausalLM", "dtype": "float16", "logits_dtype": "float32", "vocab_size": 128256, "max_position_embeddings": 8192, "hidden_size": 16384, "num_hidden_layers": 126, "num_attention_heads": 128, "num_key_value_heads": 16, "head_size": 128, "qk_layernorm": false, "hidden_act": "silu", "intermediate_size": 53248, "norm_epsilon": 1e-05, "position_embedding_type": "rope_gpt_neox", "use_parallel_embedding": false, "embedding_sharding_dim": 0, "share_embedding_table": false, "mapping": { "world_size": 8, "tp_size": 8, "pp_size": 1, "gpus_per_node": 8 }, "quantization": { "quant_algo": "FP8", "kv_cache_quant_algo": null, "group_size": 128, "smoothquant_val": null, "has_zero_point": false, "pre_quant_scale": false, "exclude_modules": [ "lm_head" ] }, "kv_dtype": "float16", "rotary_scaling": null, "residual_mlp": false, "moe_normalization_mode": null, "rotary_base": 500000.0, "moe_num_experts": 0, "moe_top_k": 0, "moe_tp_mode": 2, "attn_bias": false, "disable_weight_only_quant_plugin": false, "mlp_bias": false } ``` ### used for vLLM and SGLang ``` { "_name_or_path": "dummy_fp8", "architectures": [ "LlamaForCausalLM" ], "attention_bias": false, "attention_dropout": 0.0, "bos_token_id": 128000, "eos_token_id": 128009, "hidden_act": "silu", "hidden_size": 16384, "initializer_range": 0.02, "intermediate_size": 53248, "mlp_bias": false, "model_type": "llama", "num_attention_heads": 128, "num_hidden_layers": 126, "num_key_value_heads": 8, "pretraining_tp": 1, "quantization_config": { "activation_scheme": "static", "ignored_layers": [ "lm_head" ], "quant_method": "fp8" }, "rope_scaling": { "factor": 8.0, "low_freq_factor": 1.0, "high_freq_factor": 4.0, "original_max_position_embeddings": 8192, "rope_type": "llama3" }, "max_position_embeddings": 131072, "rms_norm_eps": 1e-05, "rope_scaling": null, "rope_theta": 500000.0, "tie_word_embeddings": false, "torch_dtype": "bfloat16", "transformers_version": "4.41.1", "use_cache": true, "vocab_size": 128256 } ``` ================================================ FILE: benchmark/boolq/README.md ================================================ ## Download data ``` git clone https://hf-mirror.com/datasets/google/boolq ``` ## Convert parquet to json ``` bash parquet_to_json.sh ``` ## Run benchmark ### Benchmark sglang ``` python -m sglang.launch_server --model-path ramblingpolymath/Qwen3-32B-W8A8 --port 30000 ``` ``` python3 bench_sglang.py ``` ================================================ FILE: benchmark/boolq/bench_sglang.py ================================================ import argparse import json import time import numpy as np from sglang.api import set_default_backend from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import read_jsonl def get_example(lines, i, answer): prompt = "Question: " + lines[i]["question"] + lines[i]["passage"] + "\nAnswer:" if answer: prompt += str(lines[i]["answer"]) return prompt def few_shot_examples(lines, k): prompts = "" for i in range(k): prompts += get_example(lines, i, True) + "\n\n" return prompts def main(args): # Select backend set_default_backend(select_sglang_backend(args)) # Read data train_data_path = args.train_data_path test_data_path = args.test_data_path lines_train = list(read_jsonl(train_data_path)) lines_test = list(read_jsonl(test_data_path)) # Construct prompts num_questions = args.num_questions num_shots = args.num_shots few_shots = few_shot_examples(lines_train, num_shots) questions = [] answer = [] for i in range(len(lines_test[:num_questions])): questions.append(get_example(lines_test, i, False)) answer.append(str(lines_test[i]["answer"])) arguments = [{"question": q} for q in questions] ##################################### ######### SGL Program Begin ######### ##################################### import sglang as sgl @sgl.function def few_shot_boolq(s, question): s += few_shots + question s += sgl.gen("answer", max_tokens=5, stop=["\n"]) ##################################### ########## SGL Program End ########## ##################################### # Run requests tic = time.perf_counter() states = few_shot_boolq.run_batch( arguments, temperature=0, num_threads=args.parallel, progress_bar=True, ) latency = time.perf_counter() - tic preds = [] for i in range(len(states)): preds.append(states[i]["answer"]) # Compute accuracy acc = np.mean(np.array(preds) == np.array(answer)) # Compute speed num_output_tokens = sum( s.get_meta_info("answer")["completion_tokens"] for s in states ) output_throughput = num_output_tokens / latency # Print results print(f"Accuracy: {acc:.3f}") print(f"Latency: {latency:.3f} s") print(f"Output throughput: {output_throughput:.3f} token/s") # Results with open(args.result_file, "a") as fout: value = { "task": "boolq", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "accuracy": round(acc, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--num-shots", type=int, default=5) parser.add_argument( "--train-data-path", type=str, default="./boolq/data/train-00000-of-00001.json" ) parser.add_argument( "--test-data-path", type=str, default="./boolq/data/validation-00000-of-00001.json", ) parser.add_argument("--num-questions", type=int, default=200) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/boolq/convert_parquet_to_json.py ================================================ import sys import pyarrow.parquet as pq def convert_parquet_to_json(input_file, output_file): # read parquet file table = pq.read_table(input_file) # turn parquet data to dataframe df = table.to_pandas() # turn dataframe to json form json_data = df.to_json(orient="records", lines=True) # write json to file with open(output_file, "w") as f: f.write(json_data) if __name__ == "__main__": if len(sys.argv) != 3: print("Usage:python convert_parquet_to_json.py ") input_file = sys.argv[1] output_file = sys.argv[2] convert_parquet_to_json(input_file, output_file) ================================================ FILE: benchmark/boolq/parquet_to_json.sh ================================================ #!/bin/bash #define input and output direction input_dir="./boolq/data" output_dir="./boolq/data" #define files needed to be handled files=( "train-00000-of-00001.parquet" "validation-00000-of-00001.parquet" ) #foe files above, use python script to convert the form for file in "${files[@]}"; do input_file="${input_dir}/${file}" output_file="${output_dir}/${file%.parquet}.json" echo "Converting ${input_file} to ${output_file} ..." python3 convert_parquet_to_json.py "${input_file}" "${output_file}" if [ $? -eq 0 ]; then echo "Conversion successful: ${output_file}" else echo "Conversion failed: ${input_file}" fi done ================================================ FILE: benchmark/ceval/README.md ================================================ ## Download data ``` git lfs clone https://huggingface.co/datasets/ceval/ceval-exam ``` ## Run benchmark ### Benchmark sglang ``` python -m sglang.launch_server --model-path ramblingpolymath/Qwen3-32B-W8A8 --port 30000 ``` ``` python3 bench_sglang.py ``` ================================================ FILE: benchmark/ceval/bench_sglang.py ================================================ import argparse import json import os import random import re import time import numpy as np from datasets import load_dataset from sglang.lang.api import set_default_backend from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) choices = ["A", "B", "C", "D"] def get_one_example(line, include_answer): res = line["question"] res += f"\nA. {line['A']}" res += f"\nB. {line['B']}" res += f"\nC. {line['C']}" res += f"\nD. {line['D']}" if include_answer: res += f"\nAnswer: {line['answer']} \n\n" return res def get_few_shot_examples(lines): res = "" for line in lines: res += get_one_example(line, True) + "\n\n" return res def get_answer_value(response): pattern = r"(Answer:|answer:|答案是|答案是:|正确答案是:|答案:|Assistant:)\s*([A-D])(?![\w])" match = re.search(pattern, response) if match: return match.group(2) return random.choice(choices) def main(args): # Read data && Construct prompts arguments = [] labels = [] examples = "examples:\n" data_path = args.data_path for subject in os.listdir(data_path): subject_path = os.path.join(data_path, subject) if os.path.isdir(subject_path) and subject != ".git": dataset = load_dataset(data_path, name=subject) dev_lines_temp = dataset["dev"] val_lines_temp = dataset["val"] few_shot_examples = get_few_shot_examples(dev_lines_temp) examples += f"{few_shot_examples}" for val_line in val_lines_temp: arguments.append( { "examples": few_shot_examples, "question": get_one_example(val_line, False), } ) labels.append(val_line["answer"]) ##################################### ######### SGL Program Begin ######### ##################################### import sglang as sgl @sgl.function def few_shot_ceval(s, examples, question): s += examples + question + sgl.gen("Answer") ##################################### ########## SGL Program End ########## ##################################### num_questions = args.num_questions if args.num_questions else len(arguments) # Select backend set_default_backend(select_sglang_backend(args)) # Run requests tic = time.perf_counter() states = few_shot_ceval.run_batch( arguments[:num_questions], temperature=0, num_threads=args.parallel, progress_bar=True, ) latency = time.perf_counter() - tic preds = [get_answer_value(states[i]["Answer"]) for i in range(num_questions)] # Compute accuracy acc = np.mean(np.array(preds) == np.array(labels[:num_questions])) # Compute speed num_output_tokens = sum( s.get_meta_info("Answer")["completion_tokens"] for s in states ) output_throughput = num_output_tokens / latency # Print results print(f"Accuracy: {acc:.3f}") print(f"Latency: {latency:.3f} s") print(f"Output throughput: {output_throughput:.3f} token/s") # Write results with open(args.result_file, "a") as fout: value = { "task": "ceval", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "accuracy": round(acc, 3), "num_requests": args.num_questions, "other": { "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="ceval/ceval-exam") parser.add_argument("--num-questions", type=int, default=None) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/deepseek_v3/README.md ================================================ # DeepSeek V3.1/V3/R1 Support The SGLang and DeepSeek teams collaborated to get DeepSeek V3 FP8 running on NVIDIA and AMD GPUs **from day one**. SGLang also supports [MLA optimization](https://lmsys.org/blog/2024-09-04-sglang-v0-3/#deepseek-multi-head-latent-attention-mla-throughput-optimizations) and [DP attention](https://lmsys.org/blog/2024-12-04-sglang-v0-4/#data-parallelism-attention-for-deepseek-models), making SGLang one of the best open-source LLM engines for running DeepSeek models. SGLang is the inference engine recommended by the official [DeepSeek team](https://github.com/deepseek-ai/DeepSeek-V3/tree/main?tab=readme-ov-file#62-inference-with-sglang-recommended). Special thanks to Meituan's Search & Recommend Platform Team and Baseten's Model Performance Team for implementing the model, and DataCrunch for providing GPU resources. For optimizations made on the DeepSeek series models regarding SGLang, please refer to [DeepSeek V3/V3.1/R1 Model Optimizations in SGLang](https://docs.sglang.io/basic_usage/deepseek_v3.html#optimizations). ## Installation & Launch If you encounter errors when starting the server, ensure the weights have finished downloading. It's recommended to download them beforehand or restart multiple times until all weights are downloaded. ### Using Docker (Recommended) ```bash # Pull latest image # https://hub.docker.com/r/lmsysorg/sglang/tags docker pull lmsysorg/sglang:latest # Launch docker run --gpus all --shm-size 32g -p 30000:30000 -v ~/.cache/huggingface:/root/.cache/huggingface --ipc=host --network=host --privileged lmsysorg/sglang:latest \ python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code --port 30000 ``` If you are using RDMA, please note that: 1. `--network host` and `--privileged` are required by RDMA. If you don't need RDMA, you can remove them. 2. You may need to set `NCCL_IB_GID_INDEX` if you are using RoCE, for example: `export NCCL_IB_GID_INDEX=3`. Add [performance optimization options](#performance-optimization-options) as needed. ### Using pip ```bash # Installation pip install sglang # Launch python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code ``` Add [performance optimization options](#performance-optimization-options) as needed. ### Performance Optimization Options [MLA optimizations](https://lmsys.org/blog/2024-09-04-sglang-v0-3/#deepseek-multi-head-latent-attention-mla-throughput-optimizations) are enabled by default. Here are some optional optimizations can be enabled as needed. - [Data Parallelism Attention](https://lmsys.org/blog/2024-12-04-sglang-v0-4/#data-parallelism-attention-for-deepseek-models): For high QPS scenarios, add the `--enable-dp-attention` argument to boost throughput. - [Torch.compile Optimization](https://lmsys.org/blog/2024-09-04-sglang-v0-3/#torchcompile-latency-optimizations): Add `--enable-torch-compile` argument to enable it. This will take some time while server starts. The maximum batch size for torch.compile optimization can be controlled with `--torch-compile-max-bs`. It's recommended to set it between `1` and `8`. (e.g., `--torch-compile-max-bs 8`) ### Usage: Chat with DeepSeek #### DeepSeek V3/R1 ```python3 import openai client = openai.Client( base_url="http://127.0.0.1:30000/v1", api_key="EMPTY") # Chat completion response = client.chat.completions.create( model="default", messages=[ {"role": "system", "content": "You are a helpful AI assistant"}, {"role": "user", "content": "List 3 countries and their capitals."}, ], temperature=0, max_tokens=64, ) print(response) ``` #### DeepSeek V3.1 On top of the basic usage similar to the DeepSeek V3/R1 example, DeepSeek V3.1 supports a request-level thinking/non-thinking toggle. Simply switch the `"thinking"` field in `extra_body={"chat_template_kwargs": {"thinking": True}}` to enable/disable the thinking mode. ##### Non Thinking ```python3 import openai client = openai.Client( base_url="http://127.0.0.1:30000/v1", api_key="EMPTY") # Chat completion response = client.chat.completions.create( model="default", messages=[ {"role": "system", "content": "You are a helpful AI assistant"}, {"role": "user", "content": "Answer the following with the second letter of the correct answer only: What is the capital of France?"}, ], temperature=0, max_tokens=1024, extra_body = {"chat_template_kwargs": {"thinking": False}} ) print(response.choices[0].message.content) ``` Answer: ``` h ``` * The correct response should be 'A', as the correct answer to the question is 'Paris'. ##### Thinking ```python3 import openai client = openai.Client( base_url="http://127.0.0.1:30000/v1", api_key="EMPTY") # Chat completion response = client.chat.completions.create( model="default", messages=[ {"role": "system", "content": "You are a helpful AI assistant"}, {"role": "user", "content": "Answer the following with the second letter of the correct answer only: What is the capital of France?"}, ], temperature=0, max_tokens=1024, extra_body = {"chat_template_kwargs": {"thinking": True}} ) print(response) ``` Answer: ``` First, the question is: "What is the capital of France?" I know that the capital of France is Paris. The user says: "Answer the following with the second letter of the correct answer only." So, I need to provide only the second letter of the correct answer. The correct answer is "Paris". Now, I need to find the second letter of "Paris". Let's spell it out: P-A-R-I-S. - First letter: P - Second letter: A - Third letter: R - Fourth letter: I - Fifth letter: S So, the second letter is "A". I should only output the second letter, which is "A". No additional text or explanation, just the letter. The user emphasized "the second letter of the correct answer only", so my response should be just "A". Finally, I need to make sure that this is the correct answer. Yes, Paris is indeed the capital of France.A ``` * The response contains `` thinking trace and model was able to derive the correct answer from it. ### Example: Serving with two H20\*8 nodes For example, there are two H20 nodes, each with 8 GPUs. The first node's IP is `10.0.0.1`, and the second node's IP is `10.0.0.2`. Please **use the first node's IP** for both commands. If the command fails, try setting the `GLOO_SOCKET_IFNAME` parameter. For more information, see [Common Environment Variables](https://pytorch.org/docs/stable/distributed.html#common-environment-variables). If the multi nodes support NVIDIA InfiniBand and encounter hanging issues during startup, consider adding the parameter `export NCCL_IB_GID_INDEX=3`. For more information, see [this](https://github.com/sgl-project/sglang/issues/3516#issuecomment-2668493307). ```bash # node 1 python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --tp 16 --dist-init-addr 10.0.0.1:5000 --nnodes 2 --node-rank 0 --trust-remote-code # node 2 python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --tp 16 --dist-init-addr 10.0.0.1:5000 --nnodes 2 --node-rank 1 --trust-remote-code ``` If you have two H100 nodes, the usage is similar to the aforementioned H20. > **Note that the launch command here does not enable Data Parallelism Attention or `torch.compile` Optimization**. For optimal performance, please refer to the command options in [Performance Optimization Options](#option_args). ### Example: Serving with one B200 node There is one B200 node with 4 (for FP4) GPUs or 8 (for FP4 or FP8) GPUs. Both FP4 and FP8 models are supported for DeepSeek R1. The flags to achieve optimal performance for each are slightly different. #### FP4 If using 4 GPUs: ```bash python3 -m sglang.launch_server --model-path nvidia/DeepSeek-R1-0528-FP4-V2 --host 0.0.0.0 --port 8000 --tensor-parallel-size=4 --cuda-graph-max-bs 256 --max-running-requests 256 --mem-fraction-static 0.85 --ep-size 4 --scheduler-recv-interval 30 --enable-symm-mem --stream-interval 10 ``` If using 8 GPUs: ```bash python3 -m sglang.launch_server --model-path nvidia/DeepSeek-R1-0528-FP4-V2 --host 0.0.0.0 --port 8000 --tensor-parallel-size=8 --cuda-graph-max-bs 256 --max-running-requests 256 --mem-fraction-static 0.85 --ep-size 8 --scheduler-recv-interval 30 --enable-symm-mem --stream-interval 10 ``` #### FP8 ```bash SGLANG_ENABLE_JIT_DEEPGEMM=false python3 -m sglang.launch_server --model-path=deepseek-ai/DeepSeek-R1-0528 --host=0.0.0.0 --port=8000 --tensor-parallel-size=8 --cuda-graph-max-bs 128 --max-running-requests 128 --mem-fraction-static 0.82 --kv-cache-dtype fp8_e4m3 --chunked-prefill-size 32768 --max-prefill-tokens 32768 --scheduler-recv-interval 30 --stream-interval 30 --fp8-gemm-backend flashinfer_trtllm ``` ### Example: Serving with two H200\*8 nodes and docker There are two H200 nodes, each with 8 GPUs. The first node's IP is `192.168.114.10`, and the second node's IP is `192.168.114.11`. Configure the endpoint to expose it to another Docker container using `--host 0.0.0.0` and `--port 40000`, and set up communications with `--dist-init-addr 192.168.114.10:20000`. A single H200 with 8 devices can run DeepSeek V3, the dual H200 setup is just to demonstrate multi-node usage. ```bash # node 1 docker run --gpus all \ --shm-size 32g \ --network=host \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --name sglang_multinode1 \ -it \ --rm \ --env "HF_TOKEN=$HF_TOKEN" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --tp 16 --dist-init-addr 192.168.114.10:20000 --nnodes 2 --node-rank 0 --trust-remote-code --host 0.0.0.0 --port 40000 ``` ```bash # node 2 docker run --gpus all \ --shm-size 32g \ --network=host \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --name sglang_multinode2 \ -it \ --rm \ --env "HF_TOKEN=$HF_TOKEN" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --tp 16 --dist-init-addr 192.168.114.10:20000 --nnodes 2 --node-rank 1 --trust-remote-code --host 0.0.0.0 --port 40000 ``` To ensure functionality, we include a test from a client Docker container. ```bash docker run --gpus all \ --shm-size 32g \ --network=host \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --name sglang_multinode_client \ -it \ --rm \ --env "HF_TOKEN=$HF_TOKEN" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1 --random-output 512 --random-range-ratio 1 --num-prompts 1 --host 0.0.0.0 --port 40000 --output-file "deepseekv3_multinode.jsonl" ``` > **Note that the launch command here does not enable Data Parallelism Attention or `torch.compile` Optimization**. For optimal performance, please refer to the command options in [Performance Optimization Options](#option_args). ### Example: Serving with four A100\*8 nodes To serve DeepSeek-V3 with A100 GPUs, we need to convert the [FP8 model checkpoints](https://huggingface.co/deepseek-ai/DeepSeek-V3) to BF16 with [script](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/fp8_cast_bf16.py) mentioned [here](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/fp8_cast_bf16.py) first. Since the BF16 model is over 1.3 TB, we need to prepare four A100 nodes, each with 8 80GB GPUs. Assume the first node's IP is `10.0.0.1`, and the converted model path is `/path/to/DeepSeek-V3-BF16`, we can have following commands to launch the server. ```bash # node 1 python3 -m sglang.launch_server --model-path /path/to/DeepSeek-V3-BF16 --tp 32 --dist-init-addr 10.0.0.1:5000 --nnodes 4 --node-rank 0 --trust-remote-code --host 0.0.0.0 --port 30000 # node 2 python3 -m sglang.launch_server --model-path /path/to/DeepSeek-V3-BF16 --tp 32 --dist-init-addr 10.0.0.1:5000 --nnodes 4 --node-rank 1 --trust-remote-code # node 3 python3 -m sglang.launch_server --model-path /path/to/DeepSeek-V3-BF16 --tp 32 --dist-init-addr 10.0.0.1:5000 --nnodes 4 --node-rank 2 --trust-remote-code # node 4 python3 -m sglang.launch_server --model-path /path/to/DeepSeek-V3-BF16 --tp 32 --dist-init-addr 10.0.0.1:5000 --nnodes 4 --node-rank 3 --trust-remote-code ``` > **Note that the launch command here does not enable Data Parallelism Attention or `torch.compile` Optimization**. For optimal performance, please refer to the command options in [Performance Optimization Options](#option_args). Then we can benchmark the accuracy and latency by accessing the first node's exposed port with the following example commands. ```bash # bench accuracy python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319 --host 10.0.0.1 --port 30000 # bench latency python3 -m sglang.bench_one_batch_server --model None --base-url http://10.0.0.1:30000 --batch-size 1 --input-len 128 --output-len 128 ``` ### Example: Serving with 8 A100/A800 with AWQ Quantization **Recommended Usage** Add `--quantization moe_wna16` flag to enable moe wna16 kernel for better performance. One example is as follows: ```bash python3 -m sglang.launch_server --model cognitivecomputations/DeepSeek-R1-AWQ --tp 8 --trust-remote-code --quantization moe_wna16 ``` Alternatively, you can use `--quantization awq_marlin` as follows: ```bash python3 -m sglang.launch_server --model cognitivecomputations/DeepSeek-R1-AWQ --tp 8 --trust-remote-code --quantization awq_marlin --dtype float16 ``` Note that `awq_marlin` only supports `float16` now, which may lead to some precision loss. ### Example: Serving with 16 A100/A800 with int8 Quantization There are block-wise and per-channel quantization methods, and the quantization parameters have already been uploaded to Huggingface. One example is as follows: - [meituan/DeepSeek-R1-Block-INT8](https://huggingface.co/meituan/DeepSeek-R1-Block-INT8) - [meituan/DeepSeek-R1-Channel-INT8](https://huggingface.co/meituan/DeepSeek-R1-Channel-INT8) Assuming that master node IP is `MASTER_IP`, checkpoint path is `/path/to/DeepSeek-R1-INT8` and port=5000, we can have following commands to launch the server: ```bash #master python3 -m sglang.launch_server \ --model meituan/DeepSeek-R1-Block-INT8 --tp 16 --dist-init-addr \ MASTER_IP:5000 --nnodes 2 --node-rank 0 --trust-remote-code --enable-torch-compile --torch-compile-max-bs 8 #cluster python3 -m sglang.launch_server \ --model meituan/DeepSeek-R1-Block-INT8 --tp 16 --dist-init-addr \ MASTER_IP:5000 --nnodes 2 --node-rank 1 --trust-remote-code --enable-torch-compile --torch-compile-max-bs 8 ``` > **Note that the launch command here enables `torch.compile` Optimization**. For optimal performance, please refer to the command options in [Performance Optimization Options](#option_args). Then on the **master node**, supposing the ShareGPT data is located at `/path/to/ShareGPT_V3_unfiltered_cleaned_split.json`, you can run the following commands to benchmark the launched server: ```bash # bench accuracy python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319 # bench serving python3 -m sglang.bench_serving --dataset-path /path/to/ShareGPT_V3_unfiltered_cleaned_split.json --dataset-name random --random-input 128 --random-output 128 --num-prompts 1000 --request-rate 128 --random-range-ratio 1.0 ``` > **Note: using `--parallel 200` can accelerate accuracy benchmarking**. ### Example: Serving with 32 L40S with int8 Quantization Running with per-channel quantization model: - [meituan/DeepSeek-R1-Channel-INT8](https://huggingface.co/meituan/DeepSeek-R1-Channel-INT8) Assuming that master node IP is `MASTER_IP`, checkpoint path is `/path/to/DeepSeek-R1-Channel-INT8` and port=5000, we can have following commands to launch the server: ```bash #master python3 -m sglang.launch_server --model meituan/DeepSeek-R1-Channel-INT8 --tp 32 --quantization w8a8_int8 \ --dist-init-addr MASTER_IP:5000 --nnodes 4 --node-rank 0 --trust-remote \ --enable-torch-compile --torch-compile-max-bs 32 #cluster python3 -m sglang.launch_server --model meituan/DeepSeek-R1-Channel-INT8 --tp 32 --quantization w8a8_int8 \ --dist-init-addr MASTER_IP:5000 --nnodes 4 --node-rank 1 --trust-remote \ --enable-torch-compile --torch-compile-max-bs 32 python3 -m sglang.launch_server --model meituan/DeepSeek-R1-Channel-INT8 --tp 32 --quantization w8a8_int8 \ --dist-init-addr MASTER_IP:5000 --nnodes 4 --node-rank 2 --trust-remote \ --enable-torch-compile --torch-compile-max-bs 32 python3 -m sglang.launch_server --model meituan/DeepSeek-R1-Channel-INT8 --tp 32 --quantization w8a8_int8 \ --dist-init-addr MASTER_IP:5000 --nnodes 4 --node-rank 3 --trust-remote \ --enable-torch-compile --torch-compile-max-bs 32 ``` The benchmarking method is the same as describted in the previous [16 x A100](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-16-a100a800-with-int8-quantization) example. ### Example: Serving on any cloud or Kubernetes with SkyPilot SkyPilot helps find cheapest available GPUs across any cloud or existing Kubernetes clusters and launch distributed serving with a single command. See details [here](https://github.com/skypilot-org/skypilot/tree/master/llm/deepseek-r1). To serve on multiple nodes: ```bash git clone https://github.com/skypilot-org/skypilot.git # Serve on 2 H100/H200x8 nodes sky launch -c r1 llm/deepseek-r1/deepseek-r1-671B.yaml --retry-until-up # Serve on 4 A100x8 nodes sky launch -c r1 llm/deepseek-r1/deepseek-r1-671B-A100.yaml --retry-until-up ``` #### Troubleshooting If you encounter the following error with fp16/bf16 checkpoint: ```bash ValueError: Weight output_partition_size = 576 is not divisible by weight quantization block_n = 128. ``` edit your `config.json` and remove the `quantization_config` block. For example: ```json "quantization_config": { "activation_scheme": "dynamic", "fmt": "e4m3", "quant_method": "fp8", "weight_block_size": [128, 128] }, ``` Removing this block typically resolves the error. For more details, see the discussion in [sgl-project/sglang#3491](https://github.com/sgl-project/sglang/issues/3491#issuecomment-2650779851). # Example: Serving with 4 H200 with w4fp8 Quantization There are mixed-precision quantization methods where MoE layers are computed using W4(int)A(FP)8 quantization while the dense layers remain in FP8 precision. Users can run these models efficiently on 4xH200 GPUs (or potentially 8xH100 GPUs), as the pre-quantized weights are already available on Hugging Face. Here's an example: ```bash python -m sglang.launch_server --model novita/Deepseek-V3-0324-W4AFP8 --mem-fraction-static 0.85 --disable-shared-experts-fusion --tp-size 4 ``` Other variants of pre-quantized DeepSeek models are also available: - [novita/Deepseek-V3.1-W4AFP8](https://huggingface.co/novita/Deepseek-V3.1-W4AFP8) - [novita/Deepseek-R1-0528-W4AFP8](https://huggingface.co/novita/Deepseek-R1-0528-W4AFP8) - [novita/Deepseek-R1-W4AFP8](https://huggingface.co/novita/Deepseek-R1-W4AFP8) - [novita/Deepseek-V3-0324-W4AFP8](https://huggingface.co/novita/Deepseek-V3-0324-W4AFP8) ## DeepSeek V3 Optimization Plan https://github.com/sgl-project/sglang/issues/2591 ================================================ FILE: benchmark/dspy/README.md ================================================ ## Install ``` pip3 install dspy-ai ``` Turn off cache at https://github.com/stanfordnlp/dspy/blob/34d8420383ec752037aa271825c1d3bf391e1277/dsp/modules/cache_utils.py#L10. ``` cache_turn_on = False ``` or set the environment variable ``` export DSP_CACHEBOOL=false ``` ## Benchmark SGLang ``` python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` ``` python3 bench_dspy_intro.py --backend sglang ``` ## Benchmark TGI ``` docker run --name tgi --rm -ti --gpus all --network host \ -v /home/ubuntu/model_weights/Llama-2-7b-chat-hf:/Llama-2-7b-chat-hf \ ghcr.io/huggingface/text-generation-inference:1.3.0 \ --model-id /Llama-2-7b-chat-hf --num-shard 1 --trust-remote-code \ --max-input-length 2048 --max-total-tokens 4096 \ --port 24000 ``` ``` python3 bench_dspy_intro.py --backend tgi ``` ## Benchmark vLLM ``` python3 -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` ``` python3 bench_dspy_intro.py --backend vllm ``` ================================================ FILE: benchmark/dspy/bench_dspy_intro.py ================================================ """ Adapted from https://github.com/stanfordnlp/dspy/blob/34d8420383ec752037aa271825c1d3bf391e1277/intro.ipynb#L9 """ import argparse import dspy from dspy.datasets import HotPotQA class BasicQA(dspy.Signature): """Answer questions with short factoid answers.""" question = dspy.InputField() answer = dspy.OutputField(desc="often between 1 and 5 words") class GenerateAnswer(dspy.Signature): """Answer questions with short factoid answers.""" context = dspy.InputField(desc="may contain relevant facts") question = dspy.InputField() answer = dspy.OutputField(desc="often between 1 and 5 words") class RAG(dspy.Module): def __init__(self, num_passages=3): super().__init__() self.retrieve = dspy.Retrieve(k=num_passages) self.generate_answer = dspy.ChainOfThought(GenerateAnswer) def forward(self, question): context = self.retrieve(question).passages prediction = self.generate_answer(context=context, question=question) return dspy.Prediction(context=context, answer=prediction.answer) def main(args): # lm = dspy.OpenAI(model='gpt-3.5-turbo') if args.backend == "tgi": lm = dspy.HFClientTGI( model="meta-llama/Llama-2-7b-chat-hf", port=args.port, url="http://localhost", ) elif args.backend == "sglang": lm = dspy.HFClientSGLang( model="meta-llama/Llama-2-7b-chat-hf", port=args.port, url="http://localhost", ) elif args.backend == "vllm": lm = dspy.HFClientVLLM( model="meta-llama/Llama-2-7b-chat-hf", port=args.port, url="http://localhost", ) else: raise ValueError(f"Invalid backend: {args.backend}") colbertv2_wiki17_abstracts = dspy.ColBERTv2( url="http://20.102.90.50:2017/wiki17_abstracts" ) dspy.settings.configure(lm=lm, rm=colbertv2_wiki17_abstracts) # Load the dataset. dataset = HotPotQA( train_seed=1, train_size=20, eval_seed=2023, dev_size=args.dev_size, test_size=0 ) # Tell DSPy that the 'question' field is the input. Any other fields are labels and/or metadata. trainset = [x.with_inputs("question") for x in dataset.train] devset = [x.with_inputs("question") for x in dataset.dev] print(len(trainset), len(devset)) train_example = trainset[0] print(f"Question: {train_example.question}") print(f"Answer: {train_example.answer}") dev_example = devset[18] print(f"Question: {dev_example.question}") print(f"Answer: {dev_example.answer}") print(f"Relevant Wikipedia Titles: {dev_example.gold_titles}") print( f"For this dataset, training examples have input keys {train_example.inputs().keys()} and label keys {train_example.labels().keys()}" ) print( f"For this dataset, dev examples have input keys {dev_example.inputs().keys()} and label keys {dev_example.labels().keys()}" ) # Define the predictor. generate_answer = dspy.Predict(BasicQA) # Call the predictor on a particular input. pred = generate_answer(question=dev_example.question) # Print the input and the prediction. print(f"Question: {dev_example.question}") print(f"Predicted Answer: {pred.answer}") lm.inspect_history(n=1) # Define the predictor. Notice we're just changing the class. The signature BasicQA is unchanged. generate_answer_with_chain_of_thought = dspy.ChainOfThought(BasicQA) # Call the predictor on the same input. pred = generate_answer_with_chain_of_thought(question=dev_example.question) # Print the input, the chain of thought, and the prediction. print(f"Question: {dev_example.question}") print(f"Thought: {pred.rationale.split('.', 1)[1].strip()}") print(f"Predicted Answer: {pred.answer}") retrieve = dspy.Retrieve(k=3) topK_passages = retrieve(dev_example.question).passages print( f"Top {retrieve.k} passages for question: {dev_example.question} \n", "-" * 30, "\n", ) for idx, passage in enumerate(topK_passages): print(f"{idx+1}]", passage, "\n") retrieve("When was the first FIFA World Cup held?").passages[0] from dspy.teleprompt import BootstrapFewShot # Validation logic: check that the predicted answer is correct. # Also check that the retrieved context does actually contain that answer. def validate_context_and_answer(example, pred, trace=None): answer_EM = dspy.evaluate.answer_exact_match(example, pred) answer_PM = dspy.evaluate.answer_passage_match(example, pred) return answer_EM and answer_PM # Set up a basic teleprompter, which will compile our RAG program. teleprompter = BootstrapFewShot(metric=validate_context_and_answer) # Compile! compiled_rag = teleprompter.compile(RAG(), trainset=trainset) # Ask any question you like to this simple RAG program. my_question = "What castle did David Gregory inherit?" # Get the prediction. This contains `pred.context` and `pred.answer`. pred = compiled_rag(my_question) # Print the contexts and the answer. print(f"Question: {my_question}") print(f"Predicted Answer: {pred.answer}") print(f"Retrieved Contexts (truncated): {[c[:200] + '...' for c in pred.context]}") from dspy.evaluate.evaluate import Evaluate # Set up the `evaluate_on_hotpotqa` function. We'll use this many times below. evaluate_on_hotpotqa = Evaluate( devset=devset, num_threads=args.num_threads, display_progress=True, display_table=5, ) # Evaluate the `compiled_rag` program with the `answer_exact_match` metric. metric = dspy.evaluate.answer_exact_match evaluate_on_hotpotqa(compiled_rag, metric=metric) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--port", type=int) parser.add_argument("--num-threads", type=int, default=32) parser.add_argument("--dev-size", type=int, default=150) parser.add_argument( "--backend", type=str, choices=["sglang", "tgi", "vllm"], default="sglang" ) args = parser.parse_args() if args.port is None: default_port = { "vllm": 21000, "lightllm": 22000, "tgi": 24000, "sglang": 30000, } args.port = default_port.get(args.backend, None) main(args) ================================================ FILE: benchmark/fla/benchmark_layernorm_gated.py ================================================ from typing import Optional import numpy as np import torch # Import the function to benchmark from sglang.srt.layers.attention.fla.layernorm_gated import ( _layer_norm_fwd as layer_norm_fwd, ) from sglang.srt.layers.attention.fla.layernorm_gated import ( rms_norm_ref, ) def benchmark_layer_norm_fwd( M: int = 65536, N: int = 128, eps: float = 1e-6, has_z: bool = True, has_bias: bool = False, group_size: Optional[int] = None, norm_before_gate: bool = True, is_rms_norm: bool = True, dtype: torch.dtype = torch.float16, warmup_iters: int = 10, benchmark_iters: int = 100, device: str = "cuda", verbose: bool = True, ): """ Benchmark layer_norm_fwd with specified parameters. Args: M: Number of rows (batch size) N: Number of columns (hidden dimension) eps: Epsilon for numerical stability has_z: Whether to use gating tensor z has_bias: Whether to use bias group_size: Group size for group normalization (None = full dimension) norm_before_gate: Whether to normalize before gating is_rms_norm: Whether to use RMS normalization (vs LayerNorm) dtype: Data type for tensors warmup_iters: Number of warmup iterations benchmark_iters: Number of benchmark iterations device: Device to run on """ if verbose: print("=" * 80) print("LayerNorm Forward Pass Benchmark") print("=" * 80) print(f"\nConfiguration:") print(f" x.shape: torch.Size([{M}, {N}])") print(f" weight.shape: torch.Size([{N}])") print(f" bias: {'torch.Size([{}])'.format(N) if has_bias else None}") print(f" eps: {eps}") print(f" z: {'torch.Size([{}, {}])'.format(M, N) if has_z else None}") print(f" out: None") print(f" group_size: {group_size}") print(f" norm_before_gate: {norm_before_gate}") print(f" is_rms_norm: {is_rms_norm}") print(f" dtype: {dtype}") print(f" device: {device}") print() # Create input tensors torch.manual_seed(42) x = torch.randn(M, N, dtype=dtype, device=device) weight = torch.randn(N, dtype=dtype, device=device) bias = torch.randn(N, dtype=dtype, device=device) if has_bias else None z = torch.randn(M, N, dtype=dtype, device=device) if has_z else None # Ensure contiguous memory layout x = x.contiguous() weight = weight.contiguous() if bias is not None: bias = bias.contiguous() if z is not None: z = z.contiguous() if verbose: print("Warming up...") # Warmup for _ in range(warmup_iters): out, mean, rstd = layer_norm_fwd( x=x, weight=weight, bias=bias, eps=eps, z=z, out=None, group_size=group_size, norm_before_gate=norm_before_gate, is_rms_norm=is_rms_norm, ) torch.cuda.synchronize() if verbose: print(f"Capturing CUDA graph...") # Capture the kernel execution in a CUDA graph runs_per_measurement = 100 # Create output tensor for graph capture out_graph = torch.empty_like(x) mean_graph = ( torch.empty((x.shape[0],), dtype=torch.float32, device=x.device) if not is_rms_norm else None ) rstd_graph = torch.empty((x.shape[0],), dtype=torch.float32, device=x.device) # Capture the graph graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): for _ in range(runs_per_measurement): out, mean, rstd = layer_norm_fwd( x=x, weight=weight, bias=bias, eps=eps, z=z, out=out_graph, group_size=group_size, norm_before_gate=norm_before_gate, is_rms_norm=is_rms_norm, ) if verbose: print( f"Running benchmark with {benchmark_iters} iterations using CUDA graph..." ) # Benchmark by replaying the graph times = [] for i in range(benchmark_iters): start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) start_event.record() graph.replay() end_event.record() torch.cuda.synchronize() # elapsed_time_ms returns milliseconds, divide by runs_per_measurement elapsed_ms = start_event.elapsed_time(end_event) times.append( elapsed_ms / 1000.0 / runs_per_measurement ) # Convert to seconds per run # Compute statistics times = np.array(times) * 1_000_000 # Convert to microseconds mean_time = np.mean(times) std_time = np.std(times) min_time = np.min(times) max_time = np.max(times) median_time = np.median(times) p95_time = np.percentile(times, 95) p99_time = np.percentile(times, 99) # Calculate throughput num_elements = M * N throughput_gelements_per_sec = (num_elements / mean_time) * 1_000_000 / 1e9 # Calculate memory bandwidth # Read: x, weight, z (if has_z) # Write: out, rstd, mean (if not rms_norm) bytes_per_element = 2 if dtype == torch.float16 else 4 # fp16 or fp32 read_bytes = (M * N + N) * bytes_per_element # x + weight if has_z: read_bytes += M * N * bytes_per_element # z write_bytes = M * N * bytes_per_element # out write_bytes += M * 4 # rstd (float32) if not is_rms_norm: write_bytes += M * 4 # mean (float32) total_bytes = read_bytes + write_bytes bandwidth_gb_per_sec = (total_bytes / mean_time) * 1_000_000 / 1e9 if verbose: print("\n" + "=" * 80) print("Benchmark Results") print("=" * 80) print(f"\nTiming Statistics (microseconds):") print(f" Mean: {mean_time:.2f} us") print(f" Std Dev: {std_time:.2f} us") print(f" Min: {min_time:.2f} us") print(f" Max: {max_time:.2f} us") print(f" Median: {median_time:.2f} us") print(f" P95: {p95_time:.2f} us") print(f" P99: {p99_time:.2f} us") print(f"\nThroughput:") print(f" {throughput_gelements_per_sec:.2f} GElements/sec") print(f" {bandwidth_gb_per_sec:.2f} GB/sec") print(f"\nMemory Usage:") print(f" Input size: {read_bytes / 1e6:.2f} MB") print(f" Output size: {write_bytes / 1e6:.2f} MB") print(f" Total: {total_bytes / 1e6:.2f} MB") # Verify correctness against reference implementation if verbose: print("\nVerifying correctness...") out_triton, mean_triton, rstd_triton = layer_norm_fwd( x=x, weight=weight, bias=bias, eps=eps, z=z, out=None, group_size=group_size, norm_before_gate=norm_before_gate, is_rms_norm=is_rms_norm, ) # Compute reference output out_ref = rms_norm_ref( x=x, weight=weight, bias=bias, z=z, eps=eps, group_size=group_size, norm_before_gate=norm_before_gate, upcast=True, ) # Compare outputs max_diff = torch.max(torch.abs(out_triton - out_ref)).item() mean_diff = torch.mean(torch.abs(out_triton - out_ref)).item() rel_diff = torch.mean( torch.abs(out_triton - out_ref) / (torch.abs(out_ref) + 1e-5) ).item() if verbose: print(f"\nCorrectness Check (vs Reference Implementation):") print(f" Max absolute difference: {max_diff:.6e}") print(f" Mean absolute difference: {mean_diff:.6e}") print(f" Mean relative difference: {rel_diff:.6e}") if max_diff < 1e-2: print(" ✓ PASS: Results match reference implementation") else: print(" ✗ FAIL: Results do not match reference implementation") print("\n" + "=" * 80) return { "mean_time_us": mean_time, "std_time_us": std_time, "min_time_us": min_time, "max_time_us": max_time, "median_time_us": median_time, "p95_time_us": p95_time, "p99_time_us": p99_time, "throughput_gelements_per_sec": throughput_gelements_per_sec, "bandwidth_gb_per_sec": bandwidth_gb_per_sec, "max_diff": max_diff, "mean_diff": mean_diff, "rel_diff": rel_diff, } def main(): """Run the benchmark with the specified configuration.""" # Configuration from user config = { "M": 65536, "N": 128, "eps": 1e-6, "has_z": True, "has_bias": False, "group_size": None, "norm_before_gate": True, "is_rms_norm": True, "dtype": torch.float16, "warmup_iters": 10, "benchmark_iters": 100, "device": "cuda", } if not torch.cuda.is_available(): print("CUDA is not available. This benchmark requires a CUDA-enabled GPU.") return results = benchmark_layer_norm_fwd(**config) # Collect all results all_results = [] # Test with different batch sizes print("\nRunning benchmarks for varying batch sizes...") for M in [256, 512, 1024, 4096, 16384, 65536, 2**17, 2**18]: config_var = config.copy() config_var["M"] = M config_var["warmup_iters"] = 5 config_var["benchmark_iters"] = 50 config_var["verbose"] = False result = benchmark_layer_norm_fwd(**config_var) all_results.append({"M": M, "N": config_var["N"], **result}) print(f" M={M:>5}: {result['mean_time_us']:>7.2f} us") # Print summary table print("\n\n") print("=" * 30) print("SUMMARY TABLE - Varying Batch Size (M) with N=128") print("=" * 30) print(f"{'M':>8} | {'Median (us)':>12}") print("-" * 30) for r in all_results: print(f"{r['M']:>8} | {r['median_time_us']:>12.2f}") print("=" * 30) if __name__ == "__main__": main() ================================================ FILE: benchmark/generative_agents/README.md ================================================ ## Download the dataset ``` wget -O agent_calls.jsonl https://drive.google.com/uc?export=download&id=19qLpD45e9JGTKF2cUjJJegwzSUEZEKht ``` ## Run benchmark Ensure that this benchmark is run in a serial manner (using --parallel 1) to preserve any potential dependencies between requests. ### Benchmark sglang ``` python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` ``` python3 bench_sglang.py --num-events 1000 --parallel 1 ``` ### Benchmark vllm ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` ``` python3 bench_other.py --num-events 1000 --backend vllm --parallel 1 ``` ### Benchmark guidance ``` python3 bench_other.py --num-events 1000 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` ### Benchmark lmql ``` python3 bench_other.py --num-events 1000 --backend lmql --parallel 1 ``` ================================================ FILE: benchmark/generative_agents/agent_functions.py ================================================ import sglang as sgl # here are the top five agent functions contributing ~70% LLM calls # reference: https://github.com/joonspk-research/generative_agents/ @sgl.function def poignancy_event(s, persona_name, persona_iss, event): s += "Here is a brief description of " + persona_name + ".\n" s += persona_iss + "\n" s += "On the scale of 1 to 10, where 1 is purely mundane (e.g., brushing teeth, making bed) and 10 is extremely poignant (e.g., a break up, college acceptance), rate the likely poignancy of the following event for" s += persona_name + ".\n\n" s += "Event: " + event s += "Rate (return a number between 1 to 10):" s += sgl.gen(name="Rate", max_tokens=2) def poignancy_event_prompt(persona_name, persona_iss, event): # return prompt and max_tokens s = "" s += "Here is a brief description of " + persona_name + ".\n" s += persona_iss + "\n" s += "On the scale of 1 to 10, where 1 is purely mundane (e.g., brushing teeth, making bed) and 10 is extremely poignant (e.g., a break up, college acceptance), rate the likely poignancy of the following event for" s += persona_name + ".\n\n" s += "Event: " + event s += "Rate (return a number between 1 to 10):" return {"prompt": s, "max_tokens": 2, "stop": None} @sgl.function def generate_event_triple(s, persona_name, action): s += """Task: Turn the input into (subject, predicate, object). Input: Sam Johnson is eating breakfast. Output: (Dolores Murphy, eat, breakfast) --- Input: Joon Park is brewing coffee. Output: (Joon Park, brew, coffee) --- Input: Jane Cook is sleeping. Output: (Jane Cook, is, sleep) --- Input: Michael Bernstein is writing email on a computer. Output: (Michael Bernstein, write, email) --- Input: Percy Liang is teaching students in a classroom. Output: (Percy Liang, teach, students) --- Input: Merrie Morris is running on a treadmill. Output: (Merrie Morris, run, treadmill) ---""" s += persona_name + "is" + action + ".\n" s += "(" + persona_name + "," s += sgl.gen(name="Triple", max_tokens=20, stop=")") def generate_event_triple_prompt(persona_name, action): s = "" s += """Task: Turn the input into (subject, predicate, object). Input: Sam Johnson is eating breakfast. Output: (Dolores Murphy, eat, breakfast) --- Input: Joon Park is brewing coffee. Output: (Joon Park, brew, coffee) --- Input: Jane Cook is sleeping. Output: (Jane Cook, is, sleep) --- Input: Michael Bernstein is writing email on a computer. Output: (Michael Bernstein, write, email) --- Input: Percy Liang is teaching students in a classroom. Output: (Percy Liang, teach, students) --- Input: Merrie Morris is running on a treadmill. Output: (Merrie Morris, run, treadmill) ---""" s += persona_name + "is" + action + ".\n" s += "(" + persona_name + "," return {"prompt": s, "max_tokens": 20, "stop": ")"} @sgl.function def generate_pronunciatio(s, action): s += "Convert an action description to an emoji (important: use two or less emojis).\n" s += "Action description: " + action + ".\n" s += "Emoji:" + sgl.gen(name="Emoji", max_tokens=6) def generate_pronunciatio_prompt(action): s = "" s += "Convert an action description to an emoji (important: use two or less emojis).\n" s += "Action description: " + action + ".\n" s += "Emoji:" return {"prompt": s, "max_tokens": 6, "stop": None} @sgl.function def action_location_sector( s, persona_name, living_sector, living_sector_areas, current_sector, current_sector_areas, daily_plan, sector_options, current_action, next_action, ): s += """Task -- choose an appropriate area from the area options for a task at hand. Sam Kim lives in {Sam Kim's house} that has Sam Kim's room, bathroom, kitchen. Sam Kim is currently in {Sam Kim's house} that has Sam Kim's room, bathroom, kitchen. Area options: {Sam Kim's house, The Rose and Crown Pub, Hobbs Cafe, Oak Hill College, Johnson Park, Harvey Oak Supply Store, The Willows Market and Pharmacy}. * Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. * Must be one of the "Area options," verbatim. For taking a walk, Sam Kim should go to the following area: {Johnson Park} --- Jane Anderson lives in {Oak Hill College Student Dormatory} that has Jane Anderson's room. Jane Anderson is currently in {Oak Hill College} that has a classroom, library Area options: {Oak Hill College Student Dormatory, The Rose and Crown Pub, Hobbs Cafe, Oak Hill College, Johnson Park, Harvey Oak Supply Store, The Willows Market and Pharmacy}. * Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. * Must be one of the "Area options," verbatim. For eating dinner, Jane Anderson should go to the following area: {Hobbs Cafe} ---""" s += ( persona_name + " lives in " + living_sector + " that has " + living_sector_areas + ".\n" ) s += ( persona_name + " is currently in " + current_sector + " that has " + current_sector_areas + ".\n" ) s += daily_plan + ".\n" s += "Area options: " + sector_options + ".\n" s += """* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. * Must be one of the "Area options," verbatim.\n""" s += ( persona_name + " is " + current_action + ". For " + next_action + ", " + persona_name + " should go to the following area: {" ) s += sgl.gen(name="Location", max_tokens=10, stop="}") def action_location_sector_prompt( persona_name, living_sector, living_sector_areas, current_sector, current_sector_areas, daily_plan, sector_options, current_action, next_action, ): s = "" s += """Task -- choose an appropriate area from the area options for a task at hand. Sam Kim lives in {Sam Kim's house} that has Sam Kim's room, bathroom, kitchen. Sam Kim is currently in {Sam Kim's house} that has Sam Kim's room, bathroom, kitchen. Area options: {Sam Kim's house, The Rose and Crown Pub, Hobbs Cafe, Oak Hill College, Johnson Park, Harvey Oak Supply Store, The Willows Market and Pharmacy}. * Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. * Must be one of the "Area options," verbatim. For taking a walk, Sam Kim should go to the following area: {Johnson Park} --- Jane Anderson lives in {Oak Hill College Student Dormatory} that has Jane Anderson's room. Jane Anderson is currently in {Oak Hill College} that has a classroom, library Area options: {Oak Hill College Student Dormatory, The Rose and Crown Pub, Hobbs Cafe, Oak Hill College, Johnson Park, Harvey Oak Supply Store, The Willows Market and Pharmacy}. * Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. * Must be one of the "Area options," verbatim. For eating dinner, Jane Anderson should go to the following area: {Hobbs Cafe} ---""" s += ( persona_name + " lives in " + living_sector + " that has " + living_sector_areas + ".\n" ) s += ( persona_name + " is currently in " + current_sector + " that has " + current_sector_areas + ".\n" ) s += daily_plan + ".\n" s += "Area options: " + sector_options + ".\n" s += """* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. * Must be one of the "Area options," verbatim.\n""" s += ( persona_name + " is " + current_action + ". For " + next_action + ", " + persona_name + " should go to the following area: {" ) return {"prompt": s, "max_tokens": 10, "stop": "}"} @sgl.function def action_location_object( s, persona_name, target_sector, target_sector_areas, current_action, next_action ): s += """ Jane Anderson is in kitchen in Jane Anderson's house. Jane Anderson is going to Jane Anderson's house that has the following areas: {kitchen, bedroom, bathroom} Stay in the current area if the activity can be done there. Never go into other people's rooms unless necessary. For cooking, Jane Anderson should go to the following area in Jane Anderson's house: Answer: {kitchen} --- Tom Watson is in common room in Tom Watson's apartment. Tom Watson is going to Hobbs Cafe that has the following areas: {cafe} Stay in the current area if the activity can be done there. Never go into other people's rooms unless necessary. For getting coffee, Tom Watson should go to the following area in Hobbs Cafe: Answer: {cafe} ---""" s += ( persona_name + " is going to " + target_sector + " that has the following areas: {" + target_sector_areas + "}\n" ) s += """* Stay in the current area if the activity can be done there. * NEVER go into other people's rooms unless necessary.""" s += ( persona_name + " is " + current_action + ". For " + next_action + ", " + persona_name + "should go to the following area in " + target_sector ) s += " (MUST pick one of {" + target_sector_areas + "}):\n" s += "Answer: {" + sgl.gen(name="Area", max_tokens=5, stop="}") def action_location_object_prompt( persona_name, target_sector, target_sector_areas, current_action, next_action ): s = "" s += """ Jane Anderson is in kitchen in Jane Anderson's house. Jane Anderson is going to Jane Anderson's house that has the following areas: {kitchen, bedroom, bathroom} Stay in the current area if the activity can be done there. Never go into other people's rooms unless necessary. For cooking, Jane Anderson should go to the following area in Jane Anderson's house: Answer: {kitchen} --- Tom Watson is in common room in Tom Watson's apartment. Tom Watson is going to Hobbs Cafe that has the following areas: {cafe} Stay in the current area if the activity can be done there. Never go into other people's rooms unless necessary. For getting coffee, Tom Watson should go to the following area in Hobbs Cafe: Answer: {cafe} ---""" s += ( persona_name + " is going to " + target_sector + " that has the following areas: {" + target_sector_areas + "}\n" ) s += """* Stay in the current area if the activity can be done there. * NEVER go into other people's rooms unless necessary.""" s += ( persona_name + " is " + current_action + ". For " + next_action + ", " + persona_name + "should go to the following area in " + target_sector ) s += " (MUST pick one of {" + target_sector_areas + "}):\n" s += "Answer: {" return {"prompt": s, "max_tokens": 5, "stop": "}"} ================================================ FILE: benchmark/generative_agents/bench_other.py ================================================ import argparse import json import time from agent_functions import ( action_location_object_prompt, action_location_sector_prompt, generate_event_triple_prompt, generate_pronunciatio_prompt, poignancy_event_prompt, ) from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import dump_state_text, read_jsonl def main(args): lines = read_jsonl(args.data_path)[: args.num_events] mapping = { "poignancy_event": poignancy_event_prompt, "generate_event_triple": generate_event_triple_prompt, "generate_pronunciatio": generate_pronunciatio_prompt, "action_location_sector": action_location_sector_prompt, "action_location_object": action_location_object_prompt, } arguments = [mapping[k](**v) for l in lines for k, v in l.items()] states = [] # Select backend call_generate = get_call_generate(args) def get_one_answer(arg): answer = call_generate(**arg, temperature=0) states.append(answer) async def get_one_answer_async(arg): answer = await call_generate(**arg, temperature=0) states.append(answer) tic = time.perf_counter() # we always sequentially execute agent calls to maintain its dependency if args.backend != "lmql": for arg in tqdm(arguments): get_one_answer(arg) else: import asyncio loop = asyncio.get_event_loop() for arg in tqdm(arguments): loop.run_until_complete(get_one_answer_async(arg)) latency = time.perf_counter() - tic print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "Generative Agents", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), # to pack weighted functions as a single agent "num_requests": len(arguments) / len(mapping), "other": { "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="agent_calls.jsonl") parser.add_argument("--num-events", type=int, default=10) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/generative_agents/bench_sglang.py ================================================ import argparse import json import time from agent_functions import ( action_location_object, action_location_sector, generate_event_triple, generate_pronunciatio, poignancy_event, ) import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text, read_jsonl def main(args): lines = read_jsonl(args.data_path)[: args.num_events] mapping = { "poignancy_event": poignancy_event, "generate_event_triple": generate_event_triple, "generate_pronunciatio": generate_pronunciatio, "action_location_sector": action_location_sector, "action_location_object": action_location_object, } arguments = [{mapping[k]: v for k, v in l.items()} for l in lines] # Select backend backend = select_sglang_backend(args) sgl.set_default_backend(backend) states = [] # Run requests tic = time.perf_counter() for a in arguments: # only a single key in the dict for func, arg in a.items(): result = func.run(**arg) result.sync() states.append(result) latency = time.perf_counter() - tic # Compute accuracy print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "Generative Agents", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), # to pack weighted functions as a single agent "num_requests": len(arguments) / len(mapping), "other": { "num_events": args.num_events, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="agent_calls.jsonl") parser.add_argument("--num-events", type=int, default=10) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/gpt_oss/README.md ================================================ # How to reproduce the result of GPT-OSS with SGLang ### Install the latest SGLang ```bash git clone https://github.com/sgl-project/sglang.git cd sglang git checkout v0.5.1.post3 pip install --upgrade pip pip install -e "python[all]" ``` ### Reproduce the benchmark throughput result (Batch Size 1) Launch Command ```bash # MXFP4 120B on H100 python3 -m sglang.launch_server --model openai/gpt-oss-120b --tp 8 --attention-backend triton # BF16 120B on H100 python3 -m sglang.launch_server --model lmsys/gpt-oss-120b-bf16 --tp 8 --attention-backend triton # MXFP4 120B on B200 python3 -m sglang.launch_server --model openai/gpt-oss-120b --tp 4 # BF16 120B on B200 python3 -m sglang.launch_server --model lmsys/gpt-oss-120b-bf16 --tp 4 ``` Benchmark Command ```bash # MXFP4 120B on H100 python3 -m sglang.bench_one_batch_server --model openai/gpt-oss-120b --base-url http://localhost:30000 --batch-size 1 --input-len 1024 --output-len 512 --show-report ``` ### Reproduce the benchmark throughput result (Batch Size 32) Launch Command ```bash # MXFP4 120B on H100 python3 -m sglang.launch_server --model openai/gpt-oss-120b --tp 8 # BF16 120B on H100 python3 -m sglang.launch_server --model lmsys/gpt-oss-120b-bf16 --tp 8 # MXFP4 120B on B200 python3 -m sglang.launch_server --model openai/gpt-oss-120b --tp 4 # BF16 120B on B200 python3 -m sglang.launch_server --model lmsys/gpt-oss-120b-bf16 --tp 4 ``` Benchmark Command ```bash python3 -m sglang.bench_one_batch_server --model openai/gpt-oss-120b --base-url http://localhost:30000 --batch-size 32 --input-len 1024 8192 --output-len 512 --show-report ``` ### Reproduce the evaluation result Install gpt-oss ```bash git clone https://github.com/openai/gpt-oss.git cd gpt-oss pip install -e . ``` Evaluation Command ```bash DATASET=gpqa BASE_URL=YOUR_BASE_URL OPENAI_API_KEY=dummy python -m gpt_oss.evals \ --base-url ${BASE_URL}/v1 \ --model dummy \ --reasoning-effort low,medium,high \ --eval $DATASET \ --n-threads 1000 ``` ### Reproduce the benchmark result of acceptance length > Note: On B200, if top k is 1, set `--attention-backend trtllm_mha` ```bash git clone https://github.com/sgl-project/SpecForge.git cd SpecForge/benchmarks config_list=( "1,0,0,0" "1,3,1,4" "1,5,4,8" ) python3 bench_model_speedup.py \ --model-path openai/gpt-oss-120b \ --speculative-draft-model-path lmsys/EAGLE3-gpt-oss-120b-bf16 \ --port 20001 \ --trust-remote-code \ --mem-fraction-static 0.8 \ --tp-size 4 \ --attention-backend fa3 \ --config-list "${config_list[@]}" \ --benchmark-list mtbench:80 gsm8k:200 humaneval:200 math500:200 \ --output lmsys_gpt-oss-120b_Eagle3_result.jsonl python3 bench_model_speedup.py \ --model-path openai/gpt-oss-120b \ --speculative-draft-model-path nvidia/gpt-oss-120b-Eagle3 \ --port 20001 \ --trust-remote-code \ --mem-fraction-static 0.8 \ --tp-size 4 \ --attention-backend fa3 \ --config-list "${config_list[@]}" \ --benchmark-list mtbench:80 gsm8k:200 humaneval:200 math500:200 \ --output nv_gpt-oss-120b_Eagle3_result.jsonl ``` ### Reproduce the result of speculative decoding speedup Launch Command ```bash # On Hopper: # - Tree decoding (topk > 1) and chain decoding (topk = 1) are supported on both FA3 and Triton backends. python3 -m sglang.launch_server --model openai/gpt-oss-120b --speculative-algorithm EAGLE3 --speculative-draft-model-path lmsys/EAGLE3-gpt-oss-120b-bf16 --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --tp 4 python3 -m sglang.launch_server --model openai/gpt-oss-120b --speculative-algorithm EAGLE3 --speculative-draft-model-path lmsys/EAGLE3-gpt-oss-120b-bf16 --speculative-num-steps 5 --speculative-eagle-topk 4 --speculative-num-draft-tokens 8 --tp 4 # On Blackwell: # - Chain decoding (topk = 1) is supported on TRTLLM-MHA backend. Tree decoding (topk > 1) is in progress, stay tuned! # - Both tree decoding (topk > 1) and chain decoding (topk = 1) are supported on the Triton backend. python3 -m sglang.launch_server --model openai/gpt-oss-120b --speculative-algo EAGLE3 --speculative-draft-model-path lmsys/EAGLE3-gpt-oss-120b-bf16 --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --tp 4 python3 -m sglang.launch_server --model openai/gpt-oss-120b --speculative-algo EAGLE3 --speculative-draft-model-path lmsys/EAGLE3-gpt-oss-120b-bf16 --speculative-num-steps 5 --speculative-eagle-topk 4 --speculative-num-draft-tokens 8 --attention-backend triton --tp 4 ``` Benchmark Command ```bash config_list=( "1,0,0,0" "1,3,1,4" "1,5,4,8" ) python3 bench_model_speedup.py \ --model-path openai/gpt-oss-120b \ --speculative-draft-model-path lmsys/EAGLE3-gpt-oss-120b-bf16 \ --port 20001 \ --trust-remote-code \ --mem-fraction-static 0.8 \ --tp-size 4 \ --attention-backend fa3 \ --config-list "${config_list[@]}" \ --benchmark-list gsm8k:200 humaneval:200 math500:200 \ --output lmsys_gpt-oss-120b_Eagle3_result.jsonl ``` We can gain the best speedup with the following settings: - **1.39x** speedup with the `--speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4` setting. - **1.52x** speedup with the `--speculative-num-steps 5 --speculative-eagle-topk 4 --speculative-num-draft-tokens 8` setting. ================================================ FILE: benchmark/gsm8k/README.md ================================================ ## Run benchmark ### Using GSM8K Platinum GSM8K Platinum is a revised version of the GSM8K test set with corrected labels and removed ambiguous questions. It can be more stable than the original GSM8K dataset. It's a drop-in replacement that can be used by adding the `--platinum` flag: ``` python3 bench_sglang.py --num-shots 8 --num-questions 1209 --parallel 1209 --platinum ``` For more information, see: https://huggingface.co/datasets/madrylab/gsm8k-platinum ### Benchmark sglang ``` python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` ``` python3 bench_sglang.py --num-questions 200 ``` ### Benchmark vllm ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` ``` python3 bench_other.py --num-questions 200 --backend vllm ``` ### Benchmark lightllm ``` # A10G python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 ``` ``` python3 bench_other.py --num-questions 200 --backend lightllm ``` ### Benchmark guidance ``` python3 bench_other.py --num-questions 200 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` ### Benchmark lmql ``` CUDA_VISIBLE_DEVICES=0,1 lmql serve-model meta-llama/Llama-2-7b-chat-hf --cuda --port 23000 ``` ``` python3 bench_other.py --num-questions 100 --backend lmql --parallel 2 ``` ================================================ FILE: benchmark/gsm8k/bench_other.py ================================================ import argparse import ast import asyncio import json import re import time from concurrent.futures import ThreadPoolExecutor import numpy as np from datasets import load_dataset from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import download_and_cache_file, dump_state_text, read_jsonl INVALID = -9999999 def get_one_example(lines, i, include_answer): ret = "Question: " + lines[i]["question"] + "\nAnswer:" if include_answer: ret += " " + lines[i]["answer"] return ret def get_few_shot_examples(lines, k): ret = "" for i in range(k): ret += get_one_example(lines, i, True) + "\n\n" return ret def get_answer_value(answer_str): answer_str = answer_str.replace(",", "") numbers = re.findall(r"\d+", answer_str) if len(numbers) < 1: return INVALID try: return ast.literal_eval(numbers[-1]) except SyntaxError: return INVALID def main(args): # Select backend call_generate = get_call_generate(args) # Read data if args.platinum: print("Loading GSM8K Platinum dataset from HuggingFace...") dataset = load_dataset("madrylab/gsm8k-platinum", "main", split="test") lines = [ {"question": item["question"], "answer": item["answer"]} for item in dataset ] else: url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl" filename = download_and_cache_file(url) lines = list(read_jsonl(filename)) # Construct prompts num_questions = args.num_questions num_shots = args.num_shots few_shot_examples = get_few_shot_examples(lines, num_shots) questions = [] labels = [] for i in range(len(lines[:num_questions])): questions.append(get_one_example(lines, i, False)) labels.append(get_answer_value(lines[i]["answer"])) assert all(l != INVALID for l in labels) states = [None] * len(labels) # Run requests if args.backend != "lmql": # Use thread pool def get_one_answer(i): answer = call_generate( prompt=few_shot_examples + questions[i], temperature=0, max_tokens=256, stop=["Question", "Assistant:", "<|separator|>"], ) states[i] = answer tic = time.perf_counter() if args.parallel == 1: for i in tqdm(range(len(questions))): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: list( tqdm( executor.map(get_one_answer, list(range(len(questions)))), total=len(questions), ) ) else: # Use asyncio async def batched_call(batch_size): for i in range(0, len(questions), batch_size): tasks = [] for q in questions[i : i + batch_size]: tasks.append( call_generate( few_shot_examples + q, temperature=0, max_tokens=256, stop="Question", ) ) rets = await asyncio.gather(*tasks) for j in range(len(rets)): states[i + j] = rets[j] tic = time.perf_counter() asyncio.run(batched_call(batch_size=args.parallel)) latency = time.perf_counter() - tic preds = [] for i in range(len(states)): preds.append(get_answer_value(states[i])) # Compute accuracy acc = np.mean(np.array(preds) == np.array(labels)) invalid = np.mean(np.array(preds) == INVALID) # Print results print(f"Accuracy: {acc:.3f}") print(f"Invalid: {invalid:.3f}") print(f"Latency: {latency:.3f} s") # Dump results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "gsm8k-platinum" if args.platinum else "gsm8k", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "accuracy": round(acc, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--num-shots", type=int, default=5) parser.add_argument("--data-path", type=str, default="test.jsonl") parser.add_argument("--num-questions", type=int, default=200) parser.add_argument( "--platinum", action="store_true", help="Use GSM8K Platinum dataset (drop-in replacement with corrected labels)", ) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/gsm8k/bench_sglang.py ================================================ import argparse import ast import json import os import re import time import numpy as np from datasets import load_dataset from sglang.lang.api import set_default_backend from sglang.test.test_utils import ( add_common_sglang_args_and_parse, dump_bench_raw_result, select_sglang_backend, ) from sglang.utils import download_and_cache_file, dump_state_text, read_jsonl INVALID = -9999999 def get_one_example(lines, i, include_answer): ret = "Question: " + lines[i]["question"] + "\nAnswer:" if include_answer: ret += " " + lines[i]["answer"] return ret def get_few_shot_examples(lines, k): ret = "" for i in range(k): ret += get_one_example(lines, i, True) + "\n\n" return ret def get_answer_value(answer_str): answer_str = answer_str.replace(",", "") numbers = re.findall(r"\d+", answer_str) if len(numbers) < 1: return INVALID try: return ast.literal_eval(numbers[-1]) except SyntaxError: return INVALID def main(args): # Select backend set_default_backend(select_sglang_backend(args)) # Load tokenizer if enable_thinking is set tokenizer = None if args.enable_thinking: from transformers import AutoTokenizer assert ( args.tokenizer_path is not None ), "--tokenizer-path is required when --enable-thinking is set" tokenizer = AutoTokenizer.from_pretrained( args.tokenizer_path, trust_remote_code=True ) # Read data if args.platinum: print("Loading GSM8K Platinum dataset from HuggingFace...") dataset = load_dataset("madrylab/gsm8k-platinum", "main", split="test") lines = [ {"question": item["question"], "answer": item["answer"]} for item in dataset ] else: data_path = args.data_path url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl" if not os.path.isfile(data_path): data_path = download_and_cache_file(url) lines = list(read_jsonl(data_path)) # Construct prompts num_questions = args.num_questions num_shots = args.num_shots few_shot_examples = get_few_shot_examples(lines, num_shots) questions = [] labels = [] for i in range(len(lines[:num_questions])): raw_question = few_shot_examples + get_one_example(lines, i, False) if tokenizer is not None: messages = [{"role": "user", "content": raw_question}] raw_question = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, enable_thinking=True, ) questions.append(raw_question) labels.append(get_answer_value(lines[i]["answer"])) assert all(l != INVALID for l in labels) arguments = [{"question": q} for q in questions] ##################################### ######### SGL Program Begin ######### ##################################### import sglang as sgl @sgl.function def few_shot_gsm8k(s, question): s += question s += sgl.gen( "answer", max_tokens=args.max_new_tokens, stop=["Question", "Assistant:", "<|separator|>"], ) ##################################### ########## SGL Program End ########## ##################################### # Run requests tic = time.perf_counter() states = few_shot_gsm8k.run_batch( arguments, temperature=args.temperature, top_p=args.top_p, num_threads=args.parallel, progress_bar=True, ) latency = time.perf_counter() - tic preds = [] for i in range(len(states)): preds.append(get_answer_value(states[i]["answer"])) # Compute accuracy acc = np.mean(np.array(preds) == np.array(labels)) invalid = np.mean(np.array(preds) == INVALID) # Compute speed num_output_tokens = sum( s.get_meta_info("answer")["completion_tokens"] for s in states ) output_throughput = num_output_tokens / latency # Print results print(f"Accuracy: {acc:.3f}") print(f"Invalid: {invalid:.3f}") print(f"Latency: {latency:.3f} s") print(f"Output throughput: {output_throughput:.3f} token/s") # Dump results dump_state_text(f"tmp_output_{args.backend}.txt", states) dump_bench_raw_result( path=args.raw_result_file, states=states, preds=preds, labels=labels, ) with open(args.result_file, "a") as fout: value = { "task": "gsm8k-platinum" if args.platinum else "gsm8k", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "accuracy": round(acc, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--num-shots", type=int, default=5) parser.add_argument("--data-path", type=str, default="test.jsonl") parser.add_argument("--num-questions", type=int, default=200) parser.add_argument("--max-new-tokens", type=int, default=512) parser.add_argument("--temperature", type=float, default=0.0) parser.add_argument("--top-p", type=float, default=1.0) parser.add_argument( "--enable-thinking", action="store_true", help="Enable thinking mode by wrapping prompts with chat template", ) parser.add_argument( "--tokenizer-path", type=str, default=None, help="Path to tokenizer (required when --enable-thinking is set)", ) parser.add_argument( "--platinum", action="store_true", help="Use GSM8K Platinum dataset (drop-in replacement with corrected labels)", ) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/hellaswag/README.md ================================================ ## Run benchmark ### Benchmark sglang ``` python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` ``` python3 bench_sglang.py --num-questions 200 ``` ### Benchmark vllm ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` ``` python3 bench_other.py --num-questions 200 --backend vllm ``` ### Benchmark lightllm ``` # A10G python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 ``` ``` python3 bench_other.py --num-questions 200 --backend lightllm ``` ### Benchmark guidance ``` CUDA_VISIBLE_DEVICES=0,1 python3 bench_other.py --num-questions 200 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` ### Benchmark lmql ``` lmql serve-model meta-llama/Llama-2-7b-chat-hf --cuda --port 23000 ``` ``` python3 bench_other.py --num-questions 200 --backend lmql --port 23000 --parallel 1 ``` ================================================ FILE: benchmark/hellaswag/bench_other.py ================================================ import argparse import asyncio import json import time from concurrent.futures import ThreadPoolExecutor import numpy as np from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_select from sglang.utils import download_and_cache_file, read_jsonl def get_one_example(lines, i, include_answer): ret = lines[i]["activity_label"] + ": " + lines[i]["ctx"] + " " if include_answer: ret += lines[i]["endings"][lines[i]["label"]] return ret def get_few_shot_examples(lines, k): ret = "" for i in range(k): ret += get_one_example(lines, i, True) + "\n\n" return ret def main(args): # Select backend call_select = get_call_select(args) # Read data url = "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl" filename = download_and_cache_file(url) lines = list(read_jsonl(filename)) # Construct prompts num_questions = args.num_questions num_shots = args.num_shots few_shot_examples = get_few_shot_examples(lines, num_shots) questions = [] choices = [] labels = [] for i in range(len(lines[:num_questions])): questions.append(get_one_example(lines, i, False)) choices.append(lines[i]["endings"]) labels.append(lines[i]["label"]) preds = [None] * len(labels) # Run requests if args.backend != "lmql": # Use thread pool def get_one_answer(i): preds[i] = call_select( context=few_shot_examples + questions[i], choices=choices[i] ) tic = time.perf_counter() if args.parallel == 1: for i in tqdm(range(len(questions))): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: list( tqdm( executor.map(get_one_answer, list(range(len(questions)))), total=len(questions), ) ) else: # Use asyncio async def batched_call(batch_size): for i in range(0, len(questions), batch_size): tasks = [] for q, c in zip( questions[i : i + batch_size], choices[i : i + batch_size] ): tasks.append(call_select(context=few_shot_examples + q, choices=c)) rets = await asyncio.gather(*tasks) for j in range(len(rets)): preds[i + j] = rets[j] tic = time.perf_counter() asyncio.run(batched_call(batch_size=args.parallel)) latency = time.perf_counter() - tic # Compute accuracy acc = np.mean(np.array(preds) == np.array(labels)) print(f"Latency: {latency:.3f}") print(f"Accuracy: {acc:.3f}") # Write results with open(args.result_file, "a") as fout: value = { "task": "hellaswag", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "accuracy": round(acc, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--num-shots", type=int, default=20) parser.add_argument("--data-path", type=str, default="hellaswag_val.jsonl") parser.add_argument("--num-questions", type=int, default=200) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/hellaswag/bench_sglang.py ================================================ import argparse import json import os import time import numpy as np from sglang.lang.api import set_default_backend from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import download_and_cache_file, read_jsonl def get_one_example(lines, i, include_answer): ret = lines[i]["activity_label"] + ": " + lines[i]["ctx"] + " " if include_answer: ret += lines[i]["endings"][lines[i]["label"]] return ret def get_few_shot_examples(lines, k): ret = "" for i in range(k): ret += get_one_example(lines, i, True) + "\n\n" return ret def main(args): # Select backend set_default_backend(select_sglang_backend(args)) # Read data data_path = args.data_path url = "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl" if not os.path.isfile(data_path): data_path = download_and_cache_file(url) lines = list(read_jsonl(data_path)) # Construct prompts num_questions = args.num_questions num_shots = args.num_shots few_shot_examples = get_few_shot_examples(lines, num_shots) questions = [] choices = [] labels = [] for i in range(len(lines[:num_questions])): questions.append(get_one_example(lines, i, False)) choices.append(lines[i]["endings"]) labels.append(lines[i]["label"]) arguments = [{"question": q, "choices": c} for q, c in zip(questions, choices)] ##################################### ######### SGL Program Begin ######### ##################################### import sglang as sgl @sgl.function def few_shot_hellaswag(s, question, choices): s += few_shot_examples + question s += sgl.select("answer", choices=choices) ##################################### ########## SGL Program End ########## ##################################### # Run requests tic = time.perf_counter() rets = few_shot_hellaswag.run_batch( arguments, temperature=0, num_threads=args.parallel, progress_bar=True, ) preds = [choices[i].index(rets[i]["answer"]) for i in range(len(rets))] latency = time.perf_counter() - tic # Compute accuracy acc = np.mean(np.array(preds) == np.array(labels)) print(f"Latency: {latency:.3f}") print(f"Accuracy: {acc:.3f}") # Write results with open(args.result_file, "a") as fout: value = { "task": "hellaswag", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "accuracy": round(acc, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--num-shots", type=int, default=20) parser.add_argument("--data-path", type=str, default="hellaswag_val.jsonl") parser.add_argument("--num-questions", type=int, default=200) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/hf3fs/bench.sh ================================================ export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages:/usr/local/lib/python3.12/dist-packages/torch/lib python3 benchmark/hf3fs/bench_client.py export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages:/usr/local/lib/python3.12/dist-packages/torch/lib SGLANG_HICACHE_HF3FS_CONFIG_PATH=/sgl-workspace/sglang/benchmark/hf3fs/hf3fs.json \ python3 benchmark/hf3fs/bench_storage.py export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages:/usr/local/lib/python3.12/dist-packages/torch/lib export SGLANG_HICACHE_HF3FS_CONFIG_PATH=/sgl-workspace/sglang/benchmark/hf3fs/hf3fs.json echo '{"file_path_prefix": "/data/hf3fs-test-0", "file_size": 1099511627776, "numjobs": 16, "entries": 8}' > \ ${SGLANG_HICACHE_HF3FS_CONFIG_PATH} python3 benchmark/hf3fs/bench_zerocopy.py #################################################################################################### rm -rf nohup.out && \ nohup python3 -m sglang.launch_server \ --model-path /code/models/Qwen3-32B/ \ --host 0.0.0.0 --port 33301 \ --page-size 64 \ --enable-hierarchical-cache \ --hicache-ratio 2 --hicache-size 0 \ --hicache-write-policy write_through \ --hicache-storage-backend hf3fs & rm -rf bench_multiturn.out && \ nohup python3 benchmark/hicache/bench_multiturn.py \ --model-path /code/models/Qwen3-32B \ --dataset-path /code/models/ShareGPT_V3_unfiltered_cleaned_split.json \ --port 33301 \ --request-length 2048 --num-clients 512 --num-rounds 3 --max-parallel 8 \ > bench_multiturn.out & #################################################################################################### rm -rf nohup.out && \ nohup python3 -m sglang.launch_server \ --model-path /code/models/DeepSeek-R1/ \ --tp 16 --nnodes 2 --node-rank 0 \ --dist-init-addr 10.74.249.153:5000 \ --host 0.0.0.0 --port 33301 \ --page-size 64 \ --enable-hierarchical-cache \ --hicache-ratio 2 --hicache-size 60 \ --hicache-write-policy write_through \ --hicache-storage-backend hf3fs & rm -rf bench_multiturn.out && \ nohup python3 benchmark/hicache/bench_multiturn.py \ --model-path /code/models/Qwen3-32B \ --dataset-path /code/models/ShareGPT_V3_unfiltered_cleaned_split.json \ --port 33301 \ --request-length 2048 --num-clients 1024 --num-rounds 3 --max-parallel 8 \ > bench_multiturn.out & #################################################################################################### ps aux | grep "sglang.launch_server" | grep -v grep | awk '{print $2}' | xargs kill -9 ps aux | grep "bench_multiturn.py" | grep -v grep | awk '{print $2}' | xargs kill -9 ================================================ FILE: benchmark/hf3fs/bench_client.py ================================================ import concurrent.futures import logging import random import time from typing import List import torch from tqdm import tqdm from sglang.srt.mem_cache.storage.hf3fs.hf3fs_usrbio_client import Hf3fsUsrBioClient def print_stats(x: List[int]): x = sorted(x) lenx = len(x) print( f"mean = {sum(x)/len(x):.2f}, " f"min = {min(x):.2f}, " f"p25 = {x[int(lenx*0.25)]:.2f}, " f"p50 = {x[int(lenx*0.5)]:.2f}, " f"p75 = {x[int(lenx*0.75)]:.2f}, " f"max = {max(x):.2f}" ) def test(): # /path/to/hf3fs file_path = "/data/bench.bin" file_size = 1 << 40 bytes_per_page = 16 << 20 entries = 32 file_ops = Hf3fsUsrBioClient(file_path, file_size, bytes_per_page, entries, 5) print("test batch_read / batch_write") num_pages = 128 dtype = torch.bfloat16 numel = bytes_per_page // dtype.itemsize offsets = list(range(file_size // bytes_per_page)) random.shuffle(offsets) offsets = offsets[:num_pages] offsets = [i * bytes_per_page for i in offsets] tensor_writes = [ torch.randn(numel, dtype=dtype) for _ in tqdm(range(num_pages), desc="prepare tensor") ] for i in tqdm(range(0, num_pages, file_ops.entries), desc="batch_write"): results = file_ops.batch_write( offsets[i : i + file_ops.entries], tensor_writes[i : i + file_ops.entries] ) assert all([result == numel * dtype.itemsize for result in results]) tensor_reads = [ torch.empty(numel, dtype=dtype) for _ in tqdm(range(num_pages), desc="prepare tensor") ] for i in tqdm(range(0, num_pages, file_ops.entries), desc="batch_read"): results = file_ops.batch_read( offsets[i : i + file_ops.entries], tensor_reads[i : i + file_ops.entries] ) assert all([result == numel * dtype.itemsize for result in results]) assert all([torch.allclose(r, w) for r, w in zip(tensor_reads, tensor_writes)]) file_ops.close() print("test done") def bench(): file_path = "/data/bench.bin" file_size = 1 << 40 bytes_per_page = 16 << 20 entries = 8 numjobs = 16 dtype = torch.bfloat16 numel = bytes_per_page // dtype.itemsize file_ops = [ Hf3fsUsrBioClient(file_path, file_size, bytes_per_page, entries, 5) for _ in range(numjobs) ] num_page = entries offsets = list(range(file_size // bytes_per_page)) tensors_write = [torch.randn(numel, dtype=dtype)] * num_page tensors_read = [torch.empty(numel, dtype=dtype)] * num_page random.shuffle(offsets) warmup = 50 iteration = 100 executor = concurrent.futures.ThreadPoolExecutor(max_workers=numjobs) w_bw = [] w_size = num_page * numjobs * bytes_per_page / (1 << 30) for i in tqdm(range(warmup + iteration), desc="Benchmarking write (GB/s)"): _offsets = [ [ offset * bytes_per_page for offset in offsets[ (i * numjobs + j) * num_page : (i * numjobs + j + 1) * num_page ] ] for j in range(numjobs) ] tik = time.perf_counter() futures = [ executor.submit(file_ops[j].batch_write, offset, tensors_write) for j, offset in enumerate(_offsets) ] results = [future.result() for future in futures] tok = time.perf_counter() if i < warmup: continue w_bw.append(w_size / (tok - tik)) results = [ _result == bytes_per_page for result in results for _result in result ] assert all(results) print_stats(w_bw) r_bw = [] r_size = w_size for i in tqdm(range(warmup + iteration), desc="Benchmarking read (GB/s)"): _offsets = [ [ offset * bytes_per_page for offset in offsets[ (i * numjobs + j) * num_page : (i * numjobs + j + 1) * num_page ] ] for j in range(numjobs) ] tik = time.perf_counter() futures = [ executor.submit(file_ops[j].batch_read, offset, tensors_read) for j, offset in enumerate(_offsets) ] results = [future.result() for future in futures] tok = time.perf_counter() if i < warmup: continue r_bw.append(r_size / (tok - tik)) results = [ _result == bytes_per_page for result in results for _result in result ] assert all(results) print_stats(r_bw) executor.shutdown(wait=True) for _file_ops in file_ops: _file_ops.close() print("bench done") def main(): logging.basicConfig(level=logging.INFO) test() bench() if __name__ == "__main__": main() ================================================ FILE: benchmark/hf3fs/bench_storage.py ================================================ import json import logging import os import random import time from typing import List import torch from tqdm import tqdm from sglang.srt.mem_cache.storage.hf3fs.mini_3fs_metadata_server import ( Hf3fsLocalMetadataClient, ) from sglang.srt.mem_cache.storage.hf3fs.storage_hf3fs import HiCacheHF3FS def print_stats(x: List[int]): x = sorted(x) lenx = len(x) print( f"mean = {sum(x)/len(x):.2f}, " f"min = {min(x):.2f}, " f"p25 = {x[int(lenx*0.25)]:.2f}, " f"p50 = {x[int(lenx*0.5)]:.2f}, " f"p75 = {x[int(lenx*0.75)]:.2f}, " f"max = {max(x):.2f}" ) def test(): # Qwen3-32B layer_num = 64 head_num, head_dim = 8, 128 kv_lora_rank, qk_rope_head_dim = 0, 0 store_dtype = torch.bfloat16 tokens_per_page = 64 file_path_prefix = "/data/test" file_size = 128 << 20 numjobs = 16 bytes_per_page = 16 << 20 entries = 2 dtype = store_dtype config_path = os.getenv(HiCacheHF3FS.default_env_var) assert config_path try: with open(config_path, "w") as f: json.dump( { "file_path_prefix": file_path_prefix, "file_size": file_size, "numjobs": numjobs, "entries": entries, }, f, ) except Exception as e: raise RuntimeError(f"Failed to dump config to {config_path}: {str(e)}") hicache_hf3fs = HiCacheHF3FS.from_env_config(bytes_per_page, dtype) numel = 2 * tokens_per_page * layer_num * head_num * head_dim assert numel * dtype.itemsize == bytes_per_page num_pages = 10 tensors = {} for i in range(num_pages): k = f"key_{i}" v = torch.randn((numel,)).to(dtype=dtype) ok = hicache_hf3fs.set(k, v) if i < (file_size // bytes_per_page): assert ok, f"Failed to insert {k}" else: assert not ok tensors[k] = v assert hicache_hf3fs.get("key_8") is None assert hicache_hf3fs.get("key_9") is None start = 0 for i in range(start, start + hicache_hf3fs.num_pages): k = f"key_{i}" assert hicache_hf3fs.exists(k) out = hicache_hf3fs.get(k) assert out is not None v = tensors[k] assert torch.allclose(v, out, atol=1e-3), f"Tensor mismatch for {k}" assert not hicache_hf3fs.exists("not_exists") hicache_hf3fs.delete("key_7") v2 = torch.randn((numel,)).to(dtype=dtype) assert hicache_hf3fs.set("key_new", v2) assert torch.allclose(hicache_hf3fs.get("key_new"), v2, atol=1e-3) hicache_hf3fs.clear() assert ( len(hicache_hf3fs.metadata_client.rank_metadata.free_pages) == hicache_hf3fs.metadata_client.rank_metadata.num_pages ) # batch num_pages = 10 tensors = {} keys = [] values = [] for i in range(num_pages): k = f"key_{i}" keys.append(k) v = torch.randn((numel,)).to(dtype=dtype) values.append(v) ok = hicache_hf3fs.batch_set(keys, values) assert not ok assert hicache_hf3fs.get("key_8") is None assert hicache_hf3fs.get("key_9") is None results = hicache_hf3fs.batch_get(keys[: hicache_hf3fs.num_pages]) for result, key, value in zip( results, keys[: hicache_hf3fs.num_pages], values[: hicache_hf3fs.num_pages] ): assert torch.allclose(value, result, atol=1e-3), f"Tensor mismatch for {key}" hicache_hf3fs.close() os.remove(hicache_hf3fs.file_path) print("All test cases passed.") def bench(): # Qwen3-32B layer_num = 64 head_num, head_dim = 8, 128 kv_lora_rank, qk_rope_head_dim = 0, 0 store_dtype = torch.bfloat16 tokens_per_page = 64 file_path = "/data/test.bin" file_size = 1 << 40 numjobs = 16 bytes_per_page = 16 << 20 entries = 8 dtype = store_dtype hicache_hf3fs = HiCacheHF3FS( rank=0, file_path=file_path, file_size=file_size, numjobs=numjobs, bytes_per_page=bytes_per_page, entries=entries, dtype=dtype, metadata_client=Hf3fsLocalMetadataClient(), ) numel = 2 * tokens_per_page * layer_num * head_num * head_dim assert numel * dtype.itemsize == bytes_per_page num_page = 128 values = [torch.randn((numel,)).to(dtype=dtype) for _ in tqdm(range(num_page))] warmup = 50 iteration = 100 w_bw = [] w_size = num_page * bytes_per_page / (1 << 30) for i in tqdm(range(warmup + iteration), desc="Benchmarking write (GB/s)"): keys = [f"{j}" for j in range(i * num_page, (i + 1) * num_page)] tik = time.perf_counter() ok = hicache_hf3fs.batch_set(keys, values) tok = time.perf_counter() if i < warmup: continue w_bw.append(w_size / (tok - tik)) assert ok print_stats(w_bw) r_bw = [] r_size = num_page * bytes_per_page / (1 << 30) for i in tqdm(range(warmup + iteration), desc="Benchmarking read (GB/s)"): keys = random.sample( list(hicache_hf3fs.metadata_client.rank_metadata.key_to_index.keys()), num_page, ) tik = time.perf_counter() results = hicache_hf3fs.batch_get(keys) tok = time.perf_counter() if i < warmup: continue r_bw.append(r_size / (tok - tik)) assert all([r is not None for r in results]) print_stats(r_bw) hicache_hf3fs.close() def allclose(): # Qwen3-32B layer_num = 64 head_num, head_dim = 8, 128 kv_lora_rank, qk_rope_head_dim = 0, 0 store_dtype = torch.bfloat16 tokens_per_page = 64 file_path = "/data/test.bin" file_size = 1 << 40 numjobs = 16 bytes_per_page = 16 << 20 entries = 8 dtype = store_dtype hicache_hf3fs = HiCacheHF3FS( rank=0, file_path=file_path, file_size=file_size, numjobs=numjobs, bytes_per_page=bytes_per_page, entries=entries, dtype=dtype, metadata_client=Hf3fsLocalMetadataClient(), ) numel = 2 * tokens_per_page * layer_num * head_num * head_dim assert numel * dtype.itemsize == bytes_per_page num_page = 128 values = [torch.randn((numel,)).to(dtype=dtype) for _ in tqdm(range(num_page))] iteration = 100 for i in tqdm(range(iteration), desc="Benchmarking write (GB/s)"): keys = [f"{j}" for j in range(i * num_page, (i + 1) * num_page)] ok = hicache_hf3fs.batch_set(keys, values) assert ok read_keys, read_results = [], [] for i in tqdm(range(iteration), desc="Benchmarking read (GB/s)"): keys = random.sample( list(hicache_hf3fs.metadata_client.rank_metadata.key_to_index.keys()), num_page, ) results = hicache_hf3fs.batch_get(keys) read_keys.extend(keys) read_results.extend(results) assert all([r is not None for r in results]) for key, result in tqdm(zip(read_keys, read_results)): assert torch.allclose(values[int(key) % num_page], result, atol=1e-3) hicache_hf3fs.close() def main(): logging.basicConfig(level=logging.INFO) test() bench() allclose() if __name__ == "__main__": main() ================================================ FILE: benchmark/hf3fs/bench_zerocopy.py ================================================ import threading import time import torch from tqdm import tqdm from sglang.srt.distributed import ( get_world_group, init_distributed_environment, initialize_model_parallel, ) from sglang.srt.managers.cache_controller import ( HiCacheController, PrefetchOperation, StorageOperation, ) from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool from sglang.srt.mem_cache.memory_pool_host import MHATokenToKVPoolHost init_distributed_environment( world_size=1, rank=0, distributed_init_method="tcp://127.0.0.1:23456", local_rank=0, backend="gloo", ) initialize_model_parallel( tensor_model_parallel_size=1, pipeline_model_parallel_size=1, ) group = get_world_group().cpu_group max_total_num_tokens = 524288 page_size = 64 kv_cache_dtype = torch.bfloat16 layer_num = 64 head_num, head_dim = 8, 128 device = "cuda" hicache_ratio = 2 hicache_size = 0 hicache_mem_layout = "page_first" # hicache_mem_layout = "layer_first" hicache_write_policy = "write_through" hicache_io_backend = "kernel" hicache_storage_backend = "hf3fs" prefetch_threshold = 256 op_size = 1024 op_num = 16 token_to_kv_pool = MHATokenToKVPool( max_total_num_tokens, page_size=page_size, dtype=kv_cache_dtype, head_num=head_num, head_dim=head_dim, layer_num=layer_num, device=device, enable_memory_saver=True, ) token_to_kv_pool_allocator = TokenToKVPoolAllocator( max_total_num_tokens, dtype=kv_cache_dtype, device=device, kvcache=token_to_kv_pool, need_sort=False, ) kv_cache = token_to_kv_pool_allocator.get_kvcache() token_to_kv_pool_host = MHATokenToKVPoolHost( kv_cache, hicache_ratio, hicache_size, page_size, hicache_mem_layout, ) load_cache_event = threading.Event() cache_controller = HiCacheController( token_to_kv_pool_allocator, token_to_kv_pool_host, page_size, group, load_cache_event=load_cache_event, write_policy=hicache_write_policy, io_backend=hicache_io_backend, storage_backend=hicache_storage_backend, prefetch_threshold=prefetch_threshold, ) operations = [ StorageOperation( torch.tensor(list(range(i, i + op_size))), list(range(i, i + op_size)), hash_value=[f"{j}" for j in range(i, i + op_size, page_size)], ) for i in tqdm(range(0, op_num * op_size, op_size)) ] tik = time.monotonic() if hicache_mem_layout == "page_first": for operation in operations: cache_controller.zerocopy_page_backup(operation, batch_size=128) elif hicache_mem_layout == "layer_first": for operation in operations: cache_controller.generic_page_backup(operation, batch_size=128) tok = time.monotonic() print(f"{tok-tik:.6f} s") operations = [ PrefetchOperation( f"{i}", torch.tensor(list(range(i, i + op_size))), list(range(i, i + op_size)), f"{i}", ) for i in tqdm(range(0, op_num * op_size, op_size)) ] for operation in operations: operation.hash_value = [ f"{j}" for j in range( int(operation.last_hash), int(operation.last_hash) + op_size, page_size ) ] tik = time.monotonic() if hicache_mem_layout == "page_first": for operation in operations: cache_controller.zerocopy_page_transfer(operation, batch_size=128) elif hicache_mem_layout == "layer_first": for operation in operations: cache_controller.generic_page_transfer(operation, batch_size=128) tok = time.monotonic() print(f"{tok-tik:.6f} s") ================================================ FILE: benchmark/hicache/README.md ================================================ ## Run synthetic multi-turn benchmark ``` # SGLang server with radix cache disabled python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 --disable-radix-cache # SGLang server with radix cache on and first-come-first-serve policy python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 --schedule-policy fcfs # The default SGLang server with radix cache on and long-prefix-match policy python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 # SGLang server with hierarchical radix cache enabled python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 --enable-hierarchical-cache ``` ``` python bench_multiturn.py --model-path Qwen/Qwen2.5-14B-Instruct ``` Note: The performance gain of hierarchical caching depends on the ratio of reusable tokens to GPU memory capacity. The more tokens to be reused, the larger the model, and the more constrained the GPU memory size, the greater the benefit one can expect from hierarchical caching. # Benchmark with more datasets ## Download Dataset ```bash ./download.sh {sharegpt|ultragpt|loogle|nextqa|all} ``` This script will automatically download the required dataset to the current working directory ## Multiturn Benchmark ### Supported Datasets - sharegpt - ultrachat - loogle ### Example Usage: ```bash python3 bench_serving.py --model mistralai/Mistral-7B-Instruct-v0.3 --backend sglang \ --dataset-path longdep_qa.json --dataset-name loogle --request-rate 10 --num-prompts 10 \ --port 8001 --enable-multiturn --disable-shuffle ``` This uses `mistralai/Mistral-7B-Instruct-v0.3` model with `sglang` as backend. The dataset is `longdep_qa.json`. We send `10 conversations` with `10 req/s` to port 8001. We enable multiturn chat without shuffling the order of conversations (i.e. following the original order in the dataset file). ### Note: The requests of multiple conversations are sent in a round robin fashion. For example, if we have 3 conversations A, B, C whose rounds are `[2, 3, 4]` correspondingly, multiturn chat will send the requests to the backend in the following order: `[A1, B1, C1, A2, B2, C2, B3, C3, C4]` This has implications on the cache reuse patterns: the cache reuse distance is the largest under this request pattern (which means a prefix-aware local scheduler in the backend can yield the most benefit compared to a FIFO scheduler) ## Shared Prefix Benchmark ### Supported Datasets - loogle ### Example Usage: ```bash python3 bench_serving.py --model mistralai/Mistral-7B-Instruct-v0.3 --backend sglang \ --dataset-path longdep_qa.json --dataset-name loogle --request-rate 10 --num-prompts 10 \ --port 8001 --enable-shared-prefix --disable-shuffle ``` ### Note: Shared Prefix benchmark sends the questions for the same prompt together. For example, if we have 3 shared prefix A, B, C, which have [2, 3, 4] questions correspondingly, the shared prefix benchmark will send the requests to the backend in the following order: `[A+Q1, A+Q2, B+Q1, B+Q2, B+Q3, C+Q1, C+Q2, C+Q3]`. ## Multi Modality Benchmark (WIP) ### Supported Datasets: - nextqa ### Example Usage: ```bash Server: python3 -m sglang.launch_server --model-path lmms-lab/LLaVA-NeXT-Video-7B --tp 2 --dp 1 --port 8001 \ --host 0.0.0.0 --mem-fraction-static 0.9 --tokenizer-path llava-hf/llava-1.5-7b-hf \ --json-model-override-args "{\"architectures\": [\"LlavaVidForCausalLM\"], \"model_type\":\"llava\", \"mm_spatial_pool_stride\":2}" Client: python3 bench_serving.py --model lmms-lab/LLaVA-NeXT-Video-7B --backend sglang --dataset-path \ NExTVideo --dataset-name nextqa --request-rate 10 --num-prompts 1 --disable-shuffle --port 8001 \ --enable-multiturn --max-frames 16 --tokenizer llava-hf/llava-1.5-7b-hf --fixed-output-len 2048 ``` Note: for the server args, `tokenizer-path`, overriding architecture are necessary. ## Supported Backend - sglang (oai) - vllm (oai) - lmdeploy (oai) ================================================ FILE: benchmark/hicache/bench_long_context.py ================================================ import json import queue import time import requests from bench_multiturn import ( ReadyQueue, WorkloadGenerator, gen_payload, log_to_jsonl_file, parse_args, ) from tqdm.asyncio import tqdm from sglang.benchmark.utils import get_tokenizer class ContextWorkloadGenerator(WorkloadGenerator): def __init__(self, args): # Construct the base URL for requests self.baseurl = f"http://{args.host}:{args.port}/" self.url = self.baseurl + "generate" self.tokenizer = get_tokenizer(args.model_path) self.distribution = args.distribution self.request_rate = args.request_rate self.start_time = None self.finished_time = None self.sent_requests = 0 self.completed_requests = 0 self.dataset = json.load(open(args.dataset_path)) num_requests = min(args.num_clients, len(self.dataset["queries"])) init_requests = [] for i in range(num_requests): context_id = self.dataset["queries"][i]["context"] # Tokenize the context + question to get input_ids prompt_text = ( self.dataset["contexts"][context_id] + self.dataset["queries"][i]["question"] ) input_ids = self.tokenizer.encode(prompt_text) output_len = len( self.tokenizer(self.dataset["queries"][i]["reference_answer"])[ "input_ids" ] ) init_requests.append((i, gen_payload(input_ids, output_len))) self.ready_queue = ReadyQueue(init_requests=init_requests) self.response_queue = queue.Queue() self.pbar = tqdm(total=num_requests) self.performance_metrics = { "ttft": [], "latency": [], "itl": [], "prompt_len": [], "cached_tokens": [], "generated_len": [], } self.max_parallel = args.max_parallel self.logfile = args.log_file self.enable_round_barrier = False def response_handler(self): while True: try: client_id, response = self.response_queue.get( timeout=10 ) # Block until response is available if not response.success: raise ValueError(f"Request failed with error: {response.error}") self.performance_metrics["ttft"].append(response.ttft) self.performance_metrics["itl"].extend(response.itl) self.performance_metrics["latency"].append(response.latency) self.performance_metrics["prompt_len"].append(response.prompt_len) self.performance_metrics["cached_tokens"].append(response.cached_tokens) self.performance_metrics["generated_len"].append(response.generated_len) self.completed_requests += 1 except queue.Empty: if self.pbar.n == self.pbar.total: break if __name__ == "__main__": args = parse_args() args.num_rounds = 1 args.max_parallel = 24 flush_cache_url = f"http://{args.host}:{args.port}/flush_cache" for request_rate in [24, 16, 12, 8, 4, 2, 1]: args.request_rate = request_rate requests.post(flush_cache_url) time.sleep(1) performance_data = ContextWorkloadGenerator(args).run() log_to_jsonl_file(performance_data, args.log_file, args.tag) ================================================ FILE: benchmark/hicache/bench_mix.py ================================================ import argparse import asyncio import json import logging import os import queue import random import threading import time from dataclasses import dataclass from functools import wraps import aiohttp from sglang.bench_serving import RequestFuncOutput from sglang.benchmark.datasets.random import sample_random_requests from sglang.benchmark.utils import get_tokenizer, remove_prefix # Set up logger logger = logging.getLogger(__name__) # Set up JSONL file for debug logging debug_log_file = None # Create a lock for thread-safe debug log writing debug_log_lock = threading.Lock() def write_debug_log(data): global debug_log_file """Write debug information to a JSONL file""" if debug_log_file is None: return # Acquire lock for thread-safe writing with debug_log_lock: # Write as JSONL (JSON Line format) debug_log_file.write(json.dumps(data) + "\n") debug_log_file.flush() def parse_args(): parser = argparse.ArgumentParser( description="Script to benchmark concurrent requests to a server." ) parser.add_argument( "--model-path", type=str, default="/data/models/Qwen3-0.6B", help="model path compatible with Hugging Face Transformers", ) parser.add_argument( "--dataset-path", type=str, default="/data/models/ShareGPT_V3_unfiltered_cleaned_split/ShareGPT_V3_unfiltered_cleaned_split.json", help="local dataset to sample tokens from", ) parser.add_argument( "--host", type=str, default="localhost", help="Server hostname or IP (default: localhost)", ) parser.add_argument( "--port", type=int, default=30000, help="Server port (default: 30000)", ) parser.add_argument( "--duration", type=int, default=600, help="Duration to run the benchmark in seconds (default: 300 seconds)", ) parser.add_argument( "--log-level", type=str, default="info", choices=["debug", "info"], help="Set the logging level (default: info)", ) parser.add_argument( "--debug-log-file", type=str, default="debug.log.jsonl", help="File to write debug logs in JSONL format", ) return parser.parse_args() def load_config(): config_path = os.getenv("CONFIG_PATH") if not config_path: raise ValueError("Environment variable 'CONFIG_PATH' is not set.") with open(config_path, "r") as f: config = json.load(f) required_keys = [ "num_rounds", "num_clients", "round_ratios", "mean_new_tokens_per_round", "mean_return_tokens_per_round", "mean_inter_round_interval", ] for key in required_keys: if key not in config: raise KeyError(f"Missing required configuration key: {key}") num_rounds = config["num_rounds"] assert len(config["round_ratios"]) == num_rounds assert len(config["mean_new_tokens_per_round"]) == num_rounds assert len(config["mean_return_tokens_per_round"]) == num_rounds assert len(config["mean_inter_round_interval"]) == num_rounds print(config) return config @dataclass class UserData: user_id: int current_round: int total_rounds: int prompt: str return_tokens: int start: int def synchronized(): def _decorator(func): @wraps(func) def wrapper(self, *args, **kwargs): with self.lock: return func(self, *args, **kwargs) return wrapper return _decorator class UserGenerator: def __init__(self, config, model_path, dataset_path): self.tokenizer_path = model_path self.tokenizer = get_tokenizer(self.tokenizer_path) self.dataset_path = dataset_path self.user_id = 0 self.lock = threading.Lock() self.num_rounds = config["num_rounds"] self.cumulative_ratios = [ sum(config["round_ratios"][: i + 1]) for i in range(len(config["round_ratios"])) ] self.mean_new_tokens_per_round = config["mean_new_tokens_per_round"] self.mean_return_tokens_per_round = config["mean_return_tokens_per_round"] self.mean_inter_round_interval = config["mean_inter_round_interval"] self.sigma = 100 self.range_ratio = 0.8 assert self.range_ratio <= 1 self.candidate_inputs = [ [ r for r in sample_random_requests( input_len=( self.mean_new_tokens_per_round[i] * (2 - self.range_ratio) ), output_len=( self.mean_return_tokens_per_round[i] * (2 - self.range_ratio) ), num_prompts=config["num_clients"], range_ratio=self.range_ratio / (2 - self.range_ratio), tokenizer=self.tokenizer, dataset_path=self.dataset_path, random_sample=False, ) ] for i in range(self.num_rounds) ] self.multiturn_queue = [] self.user_stats = [0 for _ in range(self.num_rounds)] self.input_stats = [[0, 0] for _ in range(self.num_rounds)] self.output_stats = [[0, 0] for _ in range(self.num_rounds)] def gen(self): user_id = self.user_id self.user_id += 1 rand_ratio = random.randint(0, self.cumulative_ratios[-1]) i = len(self.cumulative_ratios) for idx, cumulative_ratio in enumerate(self.cumulative_ratios): if rand_ratio >= cumulative_ratio: continue else: i = idx + 1 break total_rounds = i current_round = 0 candidate_input = random.sample(self.candidate_inputs[current_round], 1)[0] self.input_stats[0][0] += candidate_input.prompt_len self.input_stats[0][1] += 1 prompt = f"{user_id} " + candidate_input.prompt return_tokens = int( random.gauss(self.mean_return_tokens_per_round[current_round], self.sigma) ) if return_tokens <= 0: return_tokens = self.mean_return_tokens_per_round[current_round] start = 0 user_data = UserData( user_id, current_round, total_rounds, prompt, return_tokens, start ) self.user_stats[total_rounds - 1] += 1 return user_data @synchronized() def push(self, user_data, generated_text, len_itl): self.output_stats[user_data.current_round][0] += len_itl + 1 self.output_stats[user_data.current_round][1] += 1 user_data.current_round += 1 if user_data.current_round >= user_data.total_rounds: return candidate_input = random.sample( self.candidate_inputs[user_data.current_round], 1 )[0] self.input_stats[user_data.current_round][0] += candidate_input.prompt_len self.input_stats[user_data.current_round][1] += 1 user_data.prompt += generated_text + candidate_input.prompt user_data.return_tokens = int( random.gauss( self.mean_return_tokens_per_round[user_data.current_round], self.sigma ) ) if user_data.return_tokens <= 0: user_data.return_tokens = self.mean_return_tokens_per_round[ user_data.current_round ] interval = random.gauss( self.mean_inter_round_interval[user_data.current_round], self.sigma ) if interval <= 0: interval = self.mean_inter_round_interval[user_data.current_round] user_data.start = time.perf_counter() + interval if len(self.multiturn_queue) == 0: self.multiturn_queue.append(user_data) else: i = len(self.multiturn_queue) for idx, d in enumerate(self.multiturn_queue): if user_data.start < d.start: i = idx break self.multiturn_queue.insert(idx, user_data) @synchronized() def pop(self): if ( len(self.multiturn_queue) and time.perf_counter() > self.multiturn_queue[0].start ): return self.multiturn_queue.pop(0) return self.gen() def gen_payload(prompt, output_len): payload = { "text": prompt, "sampling_params": { "temperature": 0.0, "max_new_tokens": output_len, "ignore_eos": True, }, "stream": True, "stream_options": {"include_usage": True}, "lora_path": "", "return_logprob": False, "logprob_start_len": -1, } return payload AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60) async def async_request_sglang_generate( user_data, url, atomic_counter, ): """ Sends a streaming request to the server. Gathers text token-by-token. """ async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: headers = {} generated_text = "" ttft = 0.0 st = time.perf_counter() most_recent_timestamp = st output = RequestFuncOutput() payload = gen_payload(user_data.prompt, user_data.return_tokens) write_debug_log({"timestamp": st, "user_data": user_data.__dict__}) try: async with session.post(url=url, json=payload, headers=headers) as response: if response.status == 200: prompt_tokens = 0 cached_tokens = 0 async for chunk_bytes in response.content: chunk_bytes = chunk_bytes.strip() if not chunk_bytes: continue chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") latency = time.perf_counter() - st if chunk == "[DONE]": pass else: data = json.loads(chunk) if data.get("text"): timestamp = time.perf_counter() # First token if ttft == 0.0: ttft = time.perf_counter() - st output.ttft = ttft prompt_tokens = (data.get("meta_info") or {}).get( "prompt_tokens", 0 ) cached_tokens = (data.get("meta_info") or {}).get( "cached_tokens", 0 ) # Decoding phase else: output.itl.append(timestamp - most_recent_timestamp) most_recent_timestamp = timestamp generated_text = data["text"] output.generated_text = generated_text output.success = True output.latency = latency output.prompt_len = prompt_tokens output.cached_tokens = cached_tokens else: output.error = response.reason or "" output.success = False except Exception as e: output.success = False output.error = str(e) print(f"Request failed: {e}") atomic_counter.increment(1) return output class AtomicCounter: def __init__(self, initial_value=0): self._value = initial_value self.lock = threading.Lock() @synchronized() def increment(self, amount=1): self._value += amount @synchronized() def get(self): return self._value class WorkloadGenerator: def __init__(self, args): config = load_config() user_generator = UserGenerator( config, args.model_path, args.dataset_path, ) self.url = f"http://{args.host}:{args.port}/generate" self.tokenizer = user_generator.tokenizer self.start_time = None self.finished_time = None self.duration = args.duration self.done = False self.sent_requests = 0 self.completed_requests = 0 self.user_generator = user_generator self.response_queue = queue.Queue() self.performance_metrics = { "ttft": [], "latency": [], "prompt_len": [], "cached_tokens": [], } self.max_parallel = config["num_clients"] self.atomic_counter = AtomicCounter() async def handle_request(self, user_data): try: response = await async_request_sglang_generate( user_data, self.url, self.atomic_counter ) self.response_queue.put((user_data, response)) except Exception as e: print(f"Request failed: {e}") self.completed_requests += 1 def request_sender(self): async def request_loop(): while True: if self.sent_requests - self.completed_requests < self.max_parallel: new_request = self.user_generator.pop() if new_request: asyncio.create_task(self.handle_request(new_request)) self.sent_requests += 1 else: await asyncio.sleep(0.05) continue if time.perf_counter() - self.start_time > self.duration: self.done = True break loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) loop.run_until_complete(request_loop()) loop.close() def response_handler(self): while True: try: user_data, response = self.response_queue.get(timeout=10) logger.info( f"{((time.perf_counter()-self.start_time)/self.duration*100):.2f}%" ) if not response.success: raise ValueError(f"Request failed with error: {response.error}") self.user_generator.push( user_data, response.generated_text, len(response.itl) ) self.performance_metrics["ttft"].append(response.ttft) self.performance_metrics["latency"].append(response.latency) self.performance_metrics["prompt_len"].append(response.prompt_len) self.performance_metrics["cached_tokens"].append(response.cached_tokens) self.completed_requests += 1 self.finished_time = time.perf_counter() except queue.Empty: if self.done: break except ValueError as e: print(f"Error processing response for client {user_data}: {e}") continue def run(self): request_thread = threading.Thread(target=self.request_sender, daemon=True) response_thread = threading.Thread(target=self.response_handler, daemon=True) self.start_time = time.perf_counter() request_thread.start() response_thread.start() request_thread.join() response_thread.join() performance_data = { "summary": { "total_requests": len(self.performance_metrics["ttft"]), "average_ttft": sum(self.performance_metrics["ttft"]) / len(self.performance_metrics["ttft"]), "p90_ttft": sorted(self.performance_metrics["ttft"])[ int(0.9 * len(self.performance_metrics["ttft"])) ], "median_ttft": sorted(self.performance_metrics["ttft"])[ len(self.performance_metrics["ttft"]) // 2 ], "average_latency": sum(self.performance_metrics["latency"]) / len(self.performance_metrics["latency"]), "p90_latency": sorted(self.performance_metrics["latency"])[ int(0.9 * len(self.performance_metrics["latency"])) ], "median_latency": sorted(self.performance_metrics["latency"])[ len(self.performance_metrics["latency"]) // 2 ], "throughput": self.atomic_counter.get() / (self.finished_time - self.start_time), "cache_hit_rate": ( 0 if sum(self.performance_metrics["prompt_len"]) == 0 else sum(self.performance_metrics["cached_tokens"]) / sum(self.performance_metrics["prompt_len"]) ), }, } print("All requests completed") print("Performance metrics summary:") print(f" Total requests: {performance_data['summary']['total_requests']}") print(f" Average TTFT: {performance_data['summary']['average_ttft']:.2f}") print(f" P90 TTFT: {performance_data['summary']['p90_ttft']:.2f}") print(f" Median TTFT: {performance_data['summary']['median_ttft']:.2f}") print( f" Average latency: {performance_data['summary']['average_latency']:.2f}" ) print(f" P90 latency: {performance_data['summary']['p90_latency']:.2f}") print(f" Median latency: {performance_data['summary']['median_latency']:.2f}") print( f" Throughput: {performance_data['summary']['throughput']:.2f} requests per second" ) print(f" Cache Hit Rate: {performance_data['summary']['cache_hit_rate']:.6f}") user_stats = self.user_generator.user_stats input_stats = self.user_generator.input_stats output_stats = self.user_generator.output_stats print(f"round_ratios: {user_stats}") print( f"mean_new_tokens_per_round: {[int(a/b) if b > 0 else 0 for a, b in input_stats]}" ) print( f"mean_return_tokens_per_round: {[int(a/b) if b > 0 else 0 for a, b in output_stats]}" ) return performance_data def main(): global debug_log_file args = parse_args() if args.log_level == "debug": logging.basicConfig(level=logging.DEBUG) logger.info("use log_level debug") # Initialize debug log file debug_log_file = open(args.debug_log_file, "w") else: logging.basicConfig(level=logging.INFO) logger.info("use log_level info") performance_data = WorkloadGenerator(args).run() # Close debug log file if it was opened if debug_log_file: debug_log_file.close() if __name__ == "__main__": main() ================================================ FILE: benchmark/hicache/bench_mix.sh ================================================ #!/bin/bash export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.12/dist-packages:/usr/local/lib/python3.12/dist-packages/torch/lib rm -rf nohup.out && \ nohup python3 -m sglang.launch_server \ --attention-backend triton \ --model-path /code/models/Qwen3-32B/ \ --log-level info \ --tp 4 --mem-frac 0.25 \ --host 0.0.0.0 --port 33301 \ --enable-metrics --enable-cache-report \ --page-size 64 \ --enable-hierarchical-cache \ --hicache-ratio 2.5 --hicache-size 0 \ --hicache-io-backend kernel \ --hicache-mem-layout layer_first \ --hicache-write-policy write_through \ & ################################################## export CONFIG_PATH=/tmp/bench_mix_config.json # num_clients: Maximum number of concurrent client requests to be simulated # round_ratios: Distribution of requests across rounds. Given sum(round_ratios) total requests, # round_ratios[i] denotes the number of requests that will execute for (i+1) rounds echo '{ "num_rounds": 10, "num_clients": 60, "round_ratios": [50, 25, 15, 15, 10, 10, 9, 8, 7, 6], "mean_new_tokens_per_round": [1000, 400, 350, 300, 280, 260, 240, 220, 210, 200], "mean_return_tokens_per_round": [100, 100, 100, 100, 100, 100, 100, 100, 100, 100], "mean_inter_round_interval": [30, 30, 30, 30, 30, 30, 30, 30, 30, 30] }' > ${CONFIG_PATH} rm -rf bench_mix.out && \ nohup python3 /sgl-workspace/sglang/benchmark/hicache/bench_mix.py \ --model-path /code/models/Qwen3-32B/ \ --dataset-path /code/models/ShareGPT_V3_unfiltered_cleaned_split.json \ --port 33301 \ --duration 600 \ > bench_mix.out & ================================================ FILE: benchmark/hicache/bench_multiturn.py ================================================ import argparse import asyncio import json import queue import random import threading import time from datetime import datetime import numpy as np import requests from tqdm.asyncio import tqdm from sglang.bench_serving import RequestFuncOutput from sglang.benchmark.datasets.random import sample_random_requests from sglang.benchmark.utils import get_tokenizer from sglang.test.kits.cache_hit_kit import ( async_request_openai_chat_completions, async_request_sglang_generate, gen_payload, gen_payload_openai, ) def parse_args(): parser = argparse.ArgumentParser( description="Script to benchmark concurrent requests to a server." ) parser.add_argument( "--num-clients", type=int, default=256, help="Number of concurrent clients", ) parser.add_argument( "--max-parallel", type=int, default=128, help="Maximum number of parallel requests", ) parser.add_argument( "--request-length", type=int, default=512, help="Length of each new request", ) parser.add_argument( "--output-length", type=int, default=64, help="Length of each output", ) parser.add_argument( "--num-rounds", type=int, default=5, help="Number of rounds per client", ) parser.add_argument( "--distribution", type=str, default="poisson", choices=["poisson", "uniform"], help="Distribution type for request intervals (poisson or uniform)", ) parser.add_argument( "--request-rate", type=float, default=1.0, help="Average number of requests per second", ) parser.add_argument( "--host", type=str, default="localhost", help="Server hostname or IP (default: localhost)", ) parser.add_argument( "--port", type=int, default=30000, help="Server port (default: 30000)", ) parser.add_argument( "--model-path", type=str, default="meta-llama/Llama-3.1-8B-Instruct", help="model path compatible with Hugging Face Transformers", ) parser.add_argument( "--dataset-path", type=str, default="", help="local dataset to sample tokens from", ) parser.add_argument( "--log-file", type=str, default="performance_metrics.jsonl", help="File to log performance metrics", ) parser.add_argument( "--disable-auto-run", action="store_true", help="If set, disable automatically testing with a range of request rates.", ) parser.add_argument( "--disable-random-sample", action="store_true", help="If set, disable random sampling of requests from the ShareGPT dataset.", ) parser.add_argument( "--enable-round-barrier", action="store_true", help="If set, only send i-th turn requests after all (i-1)-th turn requests finished.", ) parser.add_argument( "--sub-question-input-length", type=int, default=0, help="Length of the sub question input for each request, if set 0 use request_length", ) parser.add_argument( "--ready-queue-policy", type=str, default="random", help="Policy for popping requests from the ready queue (random or fifo)", ) parser.add_argument( "--tag", type=str, default="", help="Tag of a certain run in the log file", ) parser.add_argument( "--min-rounds", type=int, default=0, help="Min rounds per client (0 = use --num-rounds)", ) parser.add_argument( "--max-rounds", type=int, default=0, help="Max rounds per client (0 = use --num-rounds)", ) parser.add_argument( "--range-ratio", type=float, default=1.0, help="Length variation ratio for prompts and outputs (1.0 = no variation, 0.5 = 50%% variation)", ) parser.add_argument("--seed", type=int, default=1, help="The random seed.") parser.add_argument( "--lora-path", type=str, default="", help="String of LoRA path. Currently we only support benchmarking on a single LoRA adaptor.", ) parser.add_argument( "--api-format", type=str, default="sglang", choices=["sglang", "openai"], help="API format to use: 'sglang' for native /generate endpoint, " "'openai' for OpenAI-compatible /v1/chat/completions endpoint.", ) return parser.parse_args() def log_to_jsonl_file(data, file_path="performance_metrics.jsonl", tag=""): """Append the data with a timestamp and tag to the specified JSONL file.""" timestamped_data = {"timestamp": datetime.now().isoformat(), "tag": tag, **data} try: with open(file_path, "a") as file: file.write( json.dumps(timestamped_data) + "\n" ) # Write as a single line in JSONL format except IOError as e: print(f"Error writing to JSONL file: {e}") class ReadyQueue: """ Thread-safe queue that can pop requests in different orders based on given policy. """ def __init__(self, init_requests=None, policy="random"): self.lock = threading.Lock() self.requests = init_requests or [] self.policy = policy def append(self, item): with self.lock: self.requests.append(item) def pop(self): with self.lock: if not self.requests: return None if self.policy == "random": index = random.randrange(len(self.requests)) return self.requests.pop(index) elif self.policy == "fifo": return self.requests.pop(0) else: # todo, varying thinking time of clients raise ValueError(f"{self.policy} not implemented") class WorkloadGenerator: def __init__(self, args): self.api_format = args.api_format self.model_path = args.model_path # Construct the base URL and select request/payload functions if self.api_format == "openai": self.url = f"http://{args.host}:{args.port}/v1/chat/completions" self.request_func = async_request_openai_chat_completions else: self.url = f"http://{args.host}:{args.port}/generate" self.request_func = async_request_sglang_generate self.tokenizer = get_tokenizer(args.model_path) self.distribution = args.distribution self.request_rate = args.request_rate self.start_time = None self.finished_time = None self.lora_path = args.lora_path self.sent_requests = 0 self.completed_requests = 0 # Resolve per-client round counts min_rounds = args.min_rounds max_rounds = args.max_rounds if min_rounds == 0 and max_rounds == 0: # Backward compat: all clients use --num-rounds min_rounds = args.num_rounds max_rounds = args.num_rounds elif min_rounds == 0: min_rounds = max_rounds elif max_rounds == 0: max_rounds = min_rounds if min_rounds < 1: raise ValueError(f"--min-rounds must be >= 1, got {min_rounds}") if min_rounds > max_rounds: raise ValueError( f"--min-rounds ({min_rounds}) must be <= --max-rounds ({max_rounds})" ) self.min_rounds = min_rounds self.max_rounds = max_rounds if min_rounds == max_rounds: # All clients have the same round count; skip randint to preserve random state self.client_total_rounds = [min_rounds] * args.num_clients else: self.client_total_rounds = [ random.randint(min_rounds, max_rounds) for _ in range(args.num_clients) ] # clients_per_round[r] = number of clients participating in round r self.clients_per_round = [ sum(1 for t in self.client_total_rounds if t > r) for r in range(max_rounds) ] self.total_requests = sum(self.client_total_rounds) range_ratio = args.range_ratio # Use return_text=False to get token ids instead of text first_round_samples = sample_random_requests( input_len=args.request_length, output_len=args.output_length, num_prompts=args.num_clients, range_ratio=range_ratio, tokenizer=self.tokenizer, dataset_path=args.dataset_path, random_sample=not args.disable_random_sample, return_text=False, ) # Store per-sample output_len for first round first_round_output_lens = [row.output_len for row in first_round_samples] # r.prompt is now List[int] when return_text=False self.candidate_inputs = [list(i.prompt) for i in first_round_samples] if args.sub_question_input_length != 0: sub_question_input_length = args.sub_question_input_length else: sub_question_input_length = args.request_length num_sub_questions = sum(max(t - 1, 0) for t in self.client_total_rounds) self.sub_question_inputs = sample_random_requests( input_len=sub_question_input_length, output_len=args.output_length, num_prompts=max(num_sub_questions, 1), range_ratio=range_ratio, tokenizer=self.tokenizer, dataset_path=args.dataset_path, random_sample=not args.disable_random_sample, return_text=False, ) if self.api_format == "openai": # OpenAI mode: history is a messages list for /v1/chat/completions initial_messages = { i: [ { "role": "user", "content": self.tokenizer.decode(self.candidate_inputs[i]), } ] for i in range(args.num_clients) } init_requests = [ ( i, gen_payload_openai( initial_messages[i], first_round_output_lens[i], self.model_path, ), ) for i in range(args.num_clients) ] self.client_records = { i: { "round": 0, "history": initial_messages[i], "total_rounds": self.client_total_rounds[i], } for i in range(args.num_clients) } else: # SGLang mode: history is List[int] (token ids) init_requests = [ ( i, gen_payload( self.candidate_inputs[i], first_round_output_lens[i], args.lora_path, ), ) for i in range(args.num_clients) ] self.client_records = { i: { "round": 0, "history": list(self.candidate_inputs[i]), "total_rounds": self.client_total_rounds[i], } for i in range(args.num_clients) } self.ready_queue = ReadyQueue( init_requests=init_requests, policy=args.ready_queue_policy ) self.candidate_inputs = self.candidate_inputs[args.num_clients :] self.response_queue = queue.Queue() self.pbar = tqdm(total=self.total_requests) self.performance_metrics = { "ttft": [], "itl": [], "latency": [], "prompt_len": [], "cached_tokens": [], "generated_len": [], } self.enable_round_barrier = args.enable_round_barrier if self.enable_round_barrier: # Add round-specific metrics while preserving the original structure for i in range(self.max_rounds): self.performance_metrics[f"round_{i}"] = { "ttft": [], "latency": [], "prompt_len": [], "cached_tokens": [], "generated_len": [], } self.num_clients = args.num_clients self.num_rounds = self.max_rounds self.max_parallel = args.max_parallel self.output_length = args.output_length async def handle_request(self, item): client_id, payload = item try: response = await self.request_func(payload, self.url, self.pbar) if self.pbar.n == self.pbar.total: self.finished_time = time.perf_counter() self.response_queue.put((client_id, response)) except Exception as e: print(f"Request failed for client {client_id}: {e}") failed_response = RequestFuncOutput() failed_response.success = False failed_response.error = str(e) self.response_queue.put((client_id, failed_response)) def request_sender(self): async def request_loop(): while True: if self.sent_requests - self.completed_requests < self.max_parallel: new_request = self.ready_queue.pop() if new_request: asyncio.create_task(self.handle_request(new_request)) self.sent_requests += 1 else: await asyncio.sleep(0.05) continue if self.pbar.n == self.pbar.total: break # Calculate Poisson-distributed wait time if self.distribution == "poisson": sleep_time = random.expovariate(self.request_rate) elif self.distribution == "uniform": avg_interval = ( 1.0 / self.request_rate if self.request_rate > 0 else 1.0 ) sleep_time = random.uniform(0, 2 * avg_interval) else: raise ValueError("Invalid distribution type") await asyncio.sleep(sleep_time) # Wait before sending the next request # Create and run the event loop for asynchronous requests loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) loop.run_until_complete(request_loop()) loop.close() def response_handler(self): next_round_reqs = [] current_barrier_round = 0 barrier_round_completed = 0 while True: try: client_id, response = self.response_queue.get( timeout=10 ) # Block until response is available if not response.success: print(f"Request failed for client {client_id}: {response.error}") self.completed_requests += 1 continue # Extend history with response if self.api_format == "openai": if response.generated_text: self.client_records[client_id]["history"].append( {"role": "assistant", "content": response.generated_text} ) else: self.client_records[client_id]["history"].extend( response.output_ids ) current_round = self.client_records[client_id]["round"] self.client_records[client_id]["round"] += 1 self.performance_metrics["ttft"].append(response.ttft) self.performance_metrics["itl"].extend(response.itl) self.performance_metrics["latency"].append(response.latency) self.performance_metrics["prompt_len"].append(response.prompt_len) self.performance_metrics["cached_tokens"].append(response.cached_tokens) self.performance_metrics["generated_len"].append(response.generated_len) if self.enable_round_barrier: self.performance_metrics[f"round_{current_round}"]["ttft"].append( response.ttft ) self.performance_metrics[f"round_{current_round}"][ "latency" ].append(response.latency) self.performance_metrics[f"round_{current_round}"][ "prompt_len" ].append(response.prompt_len) self.performance_metrics[f"round_{current_round}"][ "cached_tokens" ].append(response.cached_tokens) self.performance_metrics[f"round_{current_round}"][ "generated_len" ].append(response.generated_len) self.completed_requests += 1 client_total = self.client_records[client_id]["total_rounds"] if self.client_records[client_id]["round"] < client_total: sub_q = self.sub_question_inputs.pop() if self.api_format == "openai": # Append sub-question as a new user message sub_q_text = self.tokenizer.decode(list(sub_q.prompt)) self.client_records[client_id]["history"].append( {"role": "user", "content": sub_q_text} ) new_req = ( client_id, gen_payload_openai( self.client_records[client_id]["history"], sub_q.output_len, self.model_path, ), ) else: # Append sub-question token ids to client's history sub_q_ids = list(sub_q.prompt) self.client_records[client_id]["history"].extend(sub_q_ids) new_req = ( client_id, gen_payload( self.client_records[client_id]["history"], sub_q.output_len, self.lora_path, ), ) if self.enable_round_barrier: next_round_reqs.append(new_req) else: self.ready_queue.append(new_req) # Barrier logic: release next round when all clients for # current barrier round have completed if ( self.enable_round_barrier and current_barrier_round < self.max_rounds ): barrier_round_completed += 1 expected = self.clients_per_round[current_barrier_round] if barrier_round_completed == expected: print( f"\n Barrier: round {current_barrier_round} complete " f"({expected} clients), releasing {len(next_round_reqs)} " f"requests for round {current_barrier_round + 1}" ) for req in next_round_reqs: self.ready_queue.append(req) next_round_reqs = [] current_barrier_round += 1 barrier_round_completed = 0 except queue.Empty: if self.pbar.n == self.pbar.total: break except ValueError as e: print(f"Error processing response for client {client_id}: {e}") continue def run(self): request_thread = threading.Thread(target=self.request_sender, daemon=True) response_thread = threading.Thread(target=self.response_handler, daemon=True) self.start_time = time.perf_counter() request_thread.start() response_thread.start() request_thread.join() response_thread.join() self.pbar.close() duration = self.finished_time - self.start_time sorted_ttft = sorted(self.performance_metrics["ttft"]) sorted_latency = sorted(self.performance_metrics["latency"]) sorted_itl = sorted(self.performance_metrics["itl"]) sorted_prompt_len = sorted(self.performance_metrics["prompt_len"]) sorted_output_len = sorted(self.performance_metrics["generated_len"]) def percentile(sorted_vals, q): if not sorted_vals: return 0.0 idx = int(q * len(sorted_vals)) if idx >= len(sorted_vals): idx = len(sorted_vals) - 1 return sorted_vals[idx] def max_or_zero(sorted_vals): return sorted_vals[-1] if sorted_vals else 0.0 performance_data = { "summary": { "total_requests": len(self.performance_metrics["ttft"]), "request_rate": self.request_rate, "average_prompt_len": ( sum(self.performance_metrics["prompt_len"]) / len(self.performance_metrics["prompt_len"]) if self.performance_metrics["prompt_len"] else 0.0 ), "average_output_len": ( sum(self.performance_metrics["generated_len"]) / len(self.performance_metrics["generated_len"]) if self.performance_metrics["generated_len"] else 0.0 ), "p90_prompt_len": percentile(sorted_prompt_len, 0.9), "p99_prompt_len": percentile(sorted_prompt_len, 0.99), "p90_output_len": percentile(sorted_output_len, 0.9), "p99_output_len": percentile(sorted_output_len, 0.99), "average_ttft": sum(self.performance_metrics["ttft"]) / len(self.performance_metrics["ttft"]), "p90_ttft": percentile(sorted_ttft, 0.9), "p99_ttft": percentile(sorted_ttft, 0.99), "median_ttft": percentile(sorted_ttft, 0.5), "max_ttft": max_or_zero(sorted_ttft), "average_itl": ( sum(self.performance_metrics["itl"]) / len(self.performance_metrics["itl"]) if self.performance_metrics["itl"] else 0.0 ), "p90_itl": percentile(sorted_itl, 0.9), "p99_itl": percentile(sorted_itl, 0.99), "median_itl": percentile(sorted_itl, 0.5), "max_itl": max_or_zero(sorted_itl), "average_latency": sum(self.performance_metrics["latency"]) / len(self.performance_metrics["latency"]), "p90_latency": percentile(sorted_latency, 0.9), "p99_latency": percentile(sorted_latency, 0.99), "median_latency": percentile(sorted_latency, 0.5), "max_latency": max_or_zero(sorted_latency), "input_token_throughput": sum(self.performance_metrics["prompt_len"]) / duration, "output_token_throughput": sum( self.performance_metrics["generated_len"] ) / duration, "throughput": self.pbar.total / duration, "cache_hit_rate": ( 0 if sum(self.performance_metrics["prompt_len"]) == 0 else sum(self.performance_metrics["cached_tokens"]) / sum(self.performance_metrics["prompt_len"]) ), }, } if self.enable_round_barrier: performance_data["round"] = {} for round_num in range(self.num_rounds): round_key = f"round_{round_num}" round_metrics = self.performance_metrics[round_key] performance_data["round"][round_key] = { "average_ttft": ( sum(round_metrics["ttft"]) / len(round_metrics["ttft"]) if round_metrics["ttft"] else 0 ), "cache_hit_rate": ( 0 if sum(round_metrics["prompt_len"]) == 0 else sum(round_metrics["cached_tokens"]) / sum(round_metrics["prompt_len"]) ), "request_count": len(round_metrics["ttft"]), } print("All requests completed") print("Performance metrics summary:") print( f" Total requests: {performance_data['summary']['total_requests']} at {performance_data['summary']['request_rate']} requests per second" ) print( f" Average Prompt Length: {performance_data['summary']['average_prompt_len']:.2f} tokens" ) print( f" Average Output Length: {performance_data['summary']['average_output_len']:.2f} tokens" ) print( f" P90 Prompt Length: {performance_data['summary']['p90_prompt_len']:.0f} tokens" ) print( f" P99 Prompt Length: {performance_data['summary']['p99_prompt_len']:.0f} tokens" ) print( f" P90 Output Length: {performance_data['summary']['p90_output_len']:.0f} tokens" ) print( f" P99 Output Length: {performance_data['summary']['p99_output_len']:.0f} tokens" ) print(f" Average TTFT: {performance_data['summary']['average_ttft']:.2f}") print(f" P90 TTFT: {performance_data['summary']['p90_ttft']:.2f}") print(f" P99 TTFT: {performance_data['summary']['p99_ttft']:.2f}") print(f" Median TTFT: {performance_data['summary']['median_ttft']:.2f}") print(f" Max TTFT: {performance_data['summary']['max_ttft']:.2f}") print(f" Average ITL: {performance_data['summary']['average_itl']:.4f}") print(f" P90 ITL: {performance_data['summary']['p90_itl']:.4f}") print(f" P99 ITL: {performance_data['summary']['p99_itl']:.4f}") print(f" Median ITL: {performance_data['summary']['median_itl']:.4f}") print(f" Max ITL: {performance_data['summary']['max_itl']:.4f}") print( f" Average latency: {performance_data['summary']['average_latency']:.2f}" ) print(f" P90 latency: {performance_data['summary']['p90_latency']:.2f}") print(f" P99 latency: {performance_data['summary']['p99_latency']:.2f}") print(f" Median latency: {performance_data['summary']['median_latency']:.2f}") print(f" Max latency: {performance_data['summary']['max_latency']:.2f}") print( f" Input token throughput: {performance_data['summary']['input_token_throughput']:.2f} tokens per second" ) print( f" Output token throughput: {performance_data['summary']['output_token_throughput']:.2f} tokens per second" ) print( f" Request Throughput: {performance_data['summary']['throughput']:.2f} requests per second" ) print(f" Cache Hit Rate: {performance_data['summary']['cache_hit_rate']:.6f}") if self.enable_round_barrier: # Print round-basedsummary print("Per-round metrics:") if "round" in performance_data: for round_num in range(self.num_rounds): round_key = f"round_{round_num}" if round_key in performance_data["round"]: round_data = performance_data["round"][round_key] avg_ttft = round_data["average_ttft"] cache_hit_rate = round_data["cache_hit_rate"] request_count = round_data["request_count"] clients_in_round = self.clients_per_round[round_num] print( f" Round {round_num}: Average TTFT = {avg_ttft:.2f}s, " f"Cache Hit Rate = {cache_hit_rate:.6f} " f"({request_count} requests, " f"{clients_in_round} clients)" ) else: print(f" Round {round_num}: No requests completed") return performance_data if __name__ == "__main__": args = parse_args() flush_cache_url = f"http://{args.host}:{args.port}/flush_cache" random.seed(args.seed) np.random.seed(args.seed) if args.disable_auto_run: print("Running with specified request rate...") request_rates = [args.request_rate] else: print("Auto-running with different request rates...") request_rates = [16, 14, 12, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1] for rate in request_rates: args.request_rate = rate requests.post(flush_cache_url) time.sleep(1) performance_data = WorkloadGenerator(args).run() log_to_jsonl_file(performance_data, args.log_file, tag=args.tag) ================================================ FILE: benchmark/hicache/bench_serving.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/backend_request_func.py # Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/benchmark_serving.py """ Benchmark online serving with dynamic requests. Usage: python3 -m sglang.bench_serving --backend sglang --num-prompt 10 python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 3000 --random-input 1024 --random-output 1024 --random-range-ratio 0.5 python3 -m sglang.bench_serving --backend sglang --dataset-name random --request-rate-range 1,2,4,8,16,32 --random-input 4096 --random-output 1024 --random-range-ratio 0.125 --multi """ import argparse import asyncio import json import os import random import sys import time import traceback import warnings from argparse import ArgumentParser from dataclasses import dataclass, field from datetime import datetime from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple import aiohttp import numpy as np import requests from data_processing import MsgContent, SampleOutput, get_dataset from tqdm.asyncio import tqdm from transformers import PreTrainedTokenizerBase from sglang.benchmark.utils import get_tokenizer, remove_prefix, set_ulimit AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60) global args @dataclass class RequestFuncInput: prompts: List[Tuple[MsgContent, int, int]] api_url: str model: str lora_name: str extra_request_body: Dict[str, Any] # For multiturn chat, store the context prev_messages: List = field(default_factory=list) finished_prompts: int = 0 @dataclass class RequestFuncOutput: generated_text: List[str] = field(default_factory=list) prompt_len: List[int] = field(default_factory=list) output_len: List[int] = field(default_factory=list) latency: List[float] = field(default_factory=list) ttft: List[float] = field(default_factory=list) itl: List[float] = field(default_factory=list) # List of inter-token latencies success: bool = False error: str = "" # set ignore_eos True by default async def async_request_openai_completions( request_func_input: RequestFuncInput, queue: asyncio.Queue, tokenizer: PreTrainedTokenizerBase, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url assert api_url.endswith( "completions" ), "OpenAI Completions API URL must end with 'completions'." async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: payload = { "model": request_func_input.model, "temperature": 0.0, "best_of": 1, "stream": not args.disable_stream, "stream_options": {"include_usage": True}, "ignore_eos": not args.disable_ignore_eos, **request_func_input.extra_request_body, } headers = { "Content-Type": "application/json", "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", } output = RequestFuncOutput() prompt_idx = request_func_input.finished_prompts messages = request_func_input.prev_messages prompt, input_len, max_tokens = request_func_input.prompts[prompt_idx] prompt_len = sum( prompt[1] + prompt[2] # input_len + output_len for prompt in request_func_input.prompts[:prompt_idx] ) prompt_len += input_len # Messages messages.append( { "role": "user", "content": prompt, } ) payload["messages"] = messages payload["max_tokens"] = max_tokens # output.prompt_len = request_func_input.prompt_len # print(payload) generated_text = "" ttft = 0.0 st = time.perf_counter() most_recent_timestamp = st try: async with session.post( url=api_url, json=payload, headers=headers ) as response: if response.status == 200: actual_prompt_len = prompt_len - 1 actual_output_len = 0 async for chunk_bytes in response.content: chunk_bytes = chunk_bytes.strip() if not chunk_bytes: continue chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") latency = time.perf_counter() - st if chunk == "[DONE]": pass else: data = json.loads(chunk) timestamp = time.perf_counter() # NOTE: Some completion API might have a last # usage summary response without a token so we # want to check a token was generated if data["usage"] is not None and len(data["usage"]) > 0: actual_prompt_len = data["usage"]["prompt_tokens"] actual_output_len = data["usage"]["completion_tokens"] continue delta = data["choices"][0]["delta"] if delta.get("content", None): # First token if ttft == 0.0: ttft = time.perf_counter() - st output.ttft.append(ttft) # Decoding phase else: output.itl.append(timestamp - most_recent_timestamp) generated_text += delta["content"] most_recent_timestamp = timestamp output.prompt_len.append(actual_prompt_len) # truncate output.output_len.append(actual_output_len) output.generated_text.append(generated_text) output.success = True output.latency.append(latency) # Prepare for the new request request_func_input.prompts[prompt_idx] = ( prompt, input_len, actual_output_len, # changes from max_tokens to output_len ) prompt_idx += 1 messages.append( { "role": "assistant", "content": generated_text, } ) # Move the new request to the end of the queue if prompt_idx < len(request_func_input.prompts): request_func_input.finished_prompts = prompt_idx request_func_input.prev_messages = messages await queue.put(request_func_input) else: output.error = response.reason or "" output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) if pbar: pbar.update(1) return output async def async_request_profile(api_url: str) -> RequestFuncOutput: async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: output = RequestFuncOutput() try: async with session.post(url=api_url) as response: if response.status == 200: output.success = True else: output.error = response.reason or "" output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) return output ASYNC_REQUEST_FUNCS = { "sglang": async_request_openai_completions, "vllm": async_request_openai_completions, "lmdeploy": async_request_openai_completions, } @dataclass class BenchmarkMetrics: completed: int total_input: int total_output: int total_output_retokenized: int request_throughput: float input_throughput: float output_throughput: float output_throughput_retokenized: float total_throughput: float total_throughput_retokenized: float mean_ttft_ms: float median_ttft_ms: float std_ttft_ms: float p90_ttft_ms: float p99_ttft_ms: float mean_tpot_ms: float median_tpot_ms: float std_tpot_ms: float p90_tpot_ms: float p99_tpot_ms: float mean_itl_ms: float median_itl_ms: float std_itl_ms: float p90_itl_ms: float p99_itl_ms: float mean_e2e_latency_ms: float median_e2e_latency_ms: float std_e2e_latency_ms: float p99_e2e_latency_ms: float concurrency: float async def get_requests( input_requests_queue: asyncio.Queue, request_rate: float, num_actual_requests: int, ) -> AsyncGenerator[RequestFuncInput, None]: for _ in range(num_actual_requests): try: request = await asyncio.wait_for( input_requests_queue.get(), timeout=300 ) # Wait for 5 minutes then abort except Exception as e: print(f"exception: {e}") break yield request if request_rate == float("inf"): continue interval = np.random.exponential(1.0 / request_rate) await asyncio.sleep(interval) def calculate_metrics( outputs: List[RequestFuncOutput], dur_s: float, tokenizer: PreTrainedTokenizerBase, backend: str, ) -> Tuple[BenchmarkMetrics, List[int]]: output_lens: List[int] = [] retokenized_output_lens: List[int] = [] total_input = 0 completed = 0 itls: List[float] = [] tpots: List[float] = [] ttfts: List[float] = [] e2e_latencies: List[float] = [] output_success = 0 for i in range(len(outputs)): if outputs[i].success: output_success += 1 assert len(outputs[i].generated_text) == len(outputs[i].latency) assert len(outputs[i].generated_text) == len(outputs[i].ttft) for j in range(len(outputs[i].generated_text)): output_len = outputs[i].output_len[j] output_lens.append(output_len) retokenized_output_len = len( tokenizer.encode( outputs[i].generated_text[j], add_special_tokens=False ) ) retokenized_output_lens.append(retokenized_output_len) total_input += outputs[i].prompt_len[j] if output_len > 1: tpots.append( (outputs[i].latency[j] - outputs[i].ttft[j]) / (output_len - 1) ) completed += 1 itls += outputs[i].itl ttfts += outputs[i].ttft e2e_latencies += outputs[i].latency else: output_lens.append(0) retokenized_output_lens.append(0) if completed == 0: warnings.warn( "All requests failed. This is likely due to a misconfiguration " "on the benchmark arguments.", stacklevel=2, ) metrics = BenchmarkMetrics( completed=completed, total_input=total_input, total_output=sum(output_lens), total_output_retokenized=sum(retokenized_output_lens), request_throughput=completed / dur_s, input_throughput=total_input / dur_s, output_throughput=sum(output_lens) / dur_s, output_throughput_retokenized=sum(retokenized_output_lens) / dur_s, total_throughput=(total_input + sum(output_lens)) / dur_s, total_throughput_retokenized=(total_input + sum(retokenized_output_lens)) / dur_s, mean_ttft_ms=np.mean(ttfts or 0) * 1000, # ttfts is empty if streaming is not supported by backend median_ttft_ms=np.median(ttfts or 0) * 1000, std_ttft_ms=np.std(ttfts or 0) * 1000, p90_ttft_ms=np.percentile(ttfts or 0, 90) * 1000, p99_ttft_ms=np.percentile(ttfts or 0, 99) * 1000, mean_tpot_ms=np.mean(tpots or 0) * 1000, median_tpot_ms=np.median(tpots or 0) * 1000, std_tpot_ms=np.std(tpots or 0) * 1000, p90_tpot_ms=np.percentile(tpots or 0, 90) * 1000, p99_tpot_ms=np.percentile(tpots or 0, 99) * 1000, mean_itl_ms=np.mean(itls or 0) * 1000, median_itl_ms=np.median(itls or 0) * 1000, std_itl_ms=np.std(itls or 0) * 1000, p90_itl_ms=np.percentile(itls or 0, 90) * 1000, p99_itl_ms=np.percentile(itls or 0, 99) * 1000, mean_e2e_latency_ms=np.mean(e2e_latencies) * 1000, median_e2e_latency_ms=np.median(e2e_latencies) * 1000, std_e2e_latency_ms=np.std(e2e_latencies) * 1000, p99_e2e_latency_ms=np.percentile(e2e_latencies, 99) * 1000, concurrency=np.sum(e2e_latencies) / dur_s, ) return metrics, output_lens async def benchmark( backend: str, api_url: str, base_url: str, model_id: str, tokenizer: PreTrainedTokenizerBase, input_requests: SampleOutput, request_rate: float, max_concurrency: Optional[int], disable_tqdm: bool, lora_name: str, extra_request_body: Dict[str, Any], profile: bool, enable_shared_prefix: bool, ): if backend in ASYNC_REQUEST_FUNCS: request_func = ASYNC_REQUEST_FUNCS[backend] else: raise ValueError(f"Unknown backend: {backend}") # Limit concurrency # From https://github.com/vllm-project/vllm/pull/9390 semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None async def limited_request_func(request_func_input, queue, tokenizer, pbar): if semaphore is None: return await request_func( request_func_input=request_func_input, queue=queue, tokenizer=tokenizer, pbar=pbar, ) async with semaphore: return await request_func( request_func_input=request_func_input, queue=queue, tokenizer=tokenizer, pbar=pbar, ) num_actual_requests = sum(len(r) for r in input_requests) print(f"Num of shared prefixes or conversations: {len(input_requests)}") print(f"Num of total requests: {num_actual_requests}") # flatten the requests for shared prefix if enable_shared_prefix: input_requests = [[r] for requests in input_requests for r in requests] inputs_requests_queue = asyncio.Queue(maxsize=len(input_requests)) print("Starting initial single prompt test run...") # NOTE: Just use the first request of the first conversation for warmup test_input = RequestFuncInput( model=model_id, prompts=input_requests[0][:1], api_url=api_url, lora_name=lora_name, extra_request_body=extra_request_body, ) test_output = await request_func( request_func_input=test_input, queue=inputs_requests_queue, tokenizer=tokenizer ) if not test_output.success: raise ValueError( "Initial test run failed - Please make sure benchmark arguments " f"are correctly specified. Error: {test_output.error}" ) else: print("Initial test run completed. Starting main benchmark run...") # Check the states assert inputs_requests_queue.empty() # Flush cache if "sglang" in backend: requests.post(base_url + "/flush_cache") time.sleep(1.0) # Start profiler if profile: print("Starting profiler...") profile_output = await async_request_profile( api_url=base_url + "/start_profile" ) if profile_output.success: print("Profiler started") for request in input_requests: request_func_input = RequestFuncInput( model=model_id, prompts=request, api_url=api_url, lora_name=lora_name, extra_request_body=extra_request_body, ) inputs_requests_queue.put_nowait(request_func_input) if ( not args.enable_multiturn and not args.enable_shared_prefix and not args.dataset_name == "generated-shared-prefix" ): assert len(input_requests) == num_actual_requests pbar = None if disable_tqdm else tqdm(total=num_actual_requests) benchmark_start_time = time.perf_counter() tasks: List[asyncio.Task] = [] async for request in get_requests( inputs_requests_queue, request_rate, num_actual_requests ): tasks.append( asyncio.create_task( limited_request_func( request_func_input=request, queue=inputs_requests_queue, tokenizer=tokenizer, pbar=pbar, ) ) ) outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) # Stop profiler if profile: print("Stopping profiler...") profile_output = await async_request_profile(api_url=base_url + "/stop_profile") if profile_output.success: print("Profiler stopped") if pbar is not None: pbar.close() # Compute metrics and print results benchmark_duration = time.perf_counter() - benchmark_start_time metrics, output_lens = calculate_metrics( outputs=outputs, dur_s=benchmark_duration, tokenizer=tokenizer, backend=backend, ) print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) print("{:<40} {:<10}".format("Backend:", backend)) print("{:<40} {:<10}".format("Traffic request rate:", request_rate)) print( "{:<40} {:<10}".format( "Max request concurrency:", max_concurrency if max_concurrency else "not set", ) ) print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration)) print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output)) print( "{:<40} {:<10}".format( "Total generated tokens (retokenized):", metrics.total_output_retokenized ) ) print( "{:<40} {:<10.2f}".format( "Request throughput (req/s):", metrics.request_throughput ) ) print( "{:<40} {:<10.2f}".format( "Input token throughput (tok/s):", metrics.input_throughput ) ) print( "{:<40} {:<10.2f}".format( "Output token throughput (tok/s):", metrics.output_throughput ) ) print( "{:<40} {:<10.2f}".format( "Total token throughput (tok/s):", metrics.total_throughput ) ) print("{:<40} {:<10.2f}".format("Concurrency:", metrics.concurrency)) print("{s:{c}^{n}}".format(s="End-to-End Latency", n=50, c="-")) print( "{:<40} {:<10.2f}".format("Mean E2E Latency (ms):", metrics.mean_e2e_latency_ms) ) print( "{:<40} {:<10.2f}".format( "Median E2E Latency (ms):", metrics.median_e2e_latency_ms ) ) print("{s:{c}^{n}}".format(s="Time to First Token", n=50, c="-")) print("{:<40} {:<10.2f}".format("Mean TTFT (ms):", metrics.mean_ttft_ms)) print("{:<40} {:<10.2f}".format("Median TTFT (ms):", metrics.median_ttft_ms)) print("{:<40} {:<10.2f}".format("P90 TTFT (ms):", metrics.p90_ttft_ms)) print("{:<40} {:<10.2f}".format("P99 TTFT (ms):", metrics.p99_ttft_ms)) print( "{s:{c}^{n}}".format(s="Time per Output Token (excl. 1st token)", n=50, c="-") ) print("{:<40} {:<10.2f}".format("Mean TPOT (ms):", metrics.mean_tpot_ms)) print("{:<40} {:<10.2f}".format("Median TPOT (ms):", metrics.median_tpot_ms)) print("{:<40} {:<10.2f}".format("P90 TPOT (ms):", metrics.p90_tpot_ms)) print("{:<40} {:<10.2f}".format("P99 TPOT (ms):", metrics.p99_tpot_ms)) print("{s:{c}^{n}}".format(s="Inter-token Latency", n=50, c="-")) print("{:<40} {:<10.2f}".format("Mean ITL (ms):", metrics.mean_itl_ms)) print("{:<40} {:<10.2f}".format("Median ITL (ms):", metrics.median_itl_ms)) print("{:<40} {:<10.2f}".format("P90 ITL (ms):", metrics.p90_itl_ms)) print("{:<40} {:<10.2f}".format("P99 ITL (ms):", metrics.p99_itl_ms)) print("=" * 50) if ( metrics.median_ttft_ms is not None and metrics.mean_itl_ms is not None and metrics.output_throughput is not None ): result = { # Arguments "backend": args.backend, "dataset_name": args.dataset_name, "request_rate": request_rate, "max_concurrency": max_concurrency, "fixed_output_len": args.fixed_output_len, "random_input_len": args.random_input_len, "random_output_len": args.random_output_len, "random_range_ratio": args.random_range_ratio, # Results "duration": benchmark_duration, "completed": metrics.completed, "total_input_tokens": metrics.total_input, "total_output_tokens": metrics.total_output, "total_output_tokens_retokenized": metrics.total_output_retokenized, "request_throughput": metrics.request_throughput, "input_throughput": metrics.input_throughput, "output_throughput": metrics.output_throughput, "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, "median_e2e_latency_ms": metrics.median_e2e_latency_ms, "std_e2e_latency_ms": metrics.std_e2e_latency_ms, "p99_e2e_latency_ms": metrics.p99_e2e_latency_ms, "mean_ttft_ms": metrics.mean_ttft_ms, "median_ttft_ms": metrics.median_ttft_ms, "std_ttft_ms": metrics.std_ttft_ms, "p99_ttft_ms": metrics.p99_ttft_ms, "mean_tpot_ms": metrics.mean_tpot_ms, "median_tpot_ms": metrics.median_tpot_ms, "std_tpot_ms": metrics.std_tpot_ms, "p99_tpot_ms": metrics.p99_tpot_ms, "mean_itl_ms": metrics.mean_itl_ms, "median_itl_ms": metrics.median_itl_ms, "std_itl_ms": metrics.std_itl_ms, "p99_itl_ms": metrics.p99_itl_ms, "concurrency": metrics.concurrency, "input_throughput": metrics.input_throughput, "output_throughput": metrics.output_throughput, "fixed_output_len": args.fixed_output_len, "random_input_len": args.random_input_len, "random_output_len": args.random_output_len, "random_range_ratio": args.random_range_ratio, "duration": benchmark_duration, "completed": metrics.completed, } else: print(f"Error running benchmark for request rate: {request_rate}") print("-" * 30) # Determine output file name if args.output_file: output_file_name = args.output_file else: now = datetime.now().strftime("%m%d") if args.dataset_name == "random": output_file_name = f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_{args.random_output_len}.jsonl" else: output_file_name = ( f"{args.backend}_{now}_{args.num_prompts}_{args.dataset_name}.jsonl" ) # Append results to a JSONL file with open(output_file_name, "a") as file: file.write(json.dumps(result) + "\n") result = { "duration": benchmark_duration, "completed": metrics.completed, "total_input_tokens": metrics.total_input, "total_output_tokens": metrics.total_output, "total_output_tokens_retokenized": metrics.total_output_retokenized, "request_throughput": metrics.request_throughput, "input_throughput": metrics.input_throughput, "output_throughput": metrics.output_throughput, "mean_ttft_ms": metrics.mean_ttft_ms, "median_ttft_ms": metrics.median_ttft_ms, "std_ttft_ms": metrics.std_ttft_ms, "p90_ttft_ms": metrics.p90_ttft_ms, "p99_ttft_ms": metrics.p99_ttft_ms, "mean_tpot_ms": metrics.mean_tpot_ms, "median_tpot_ms": metrics.median_tpot_ms, "std_tpot_ms": metrics.std_tpot_ms, "p90_tpot_ms": metrics.p90_tpot_ms, "p99_tpot_ms": metrics.p99_tpot_ms, "mean_itl_ms": metrics.mean_itl_ms, "median_itl_ms": metrics.median_itl_ms, "std_itl_ms": metrics.std_itl_ms, "p90_itl_ms": metrics.p90_itl_ms, "p99_itl_ms": metrics.p99_itl_ms, "input_lens": [output.prompt_len for output in outputs], "output_lens": output_lens, "ttfts": [output.ttft for output in outputs], "itls": [output.itl for output in outputs], "generated_texts": [output.generated_text for output in outputs], "errors": [output.error for output in outputs], "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, "median_e2e_latency_ms": metrics.median_e2e_latency_ms, } return result def run_benchmark(args_: argparse.Namespace): global args args = args_ # Set default value for max_concurrency if not present if not hasattr(args, "max_concurrency"): args.max_concurrency = None # Set global environments set_ulimit() random.seed(args.seed) np.random.seed(args.seed) extra_request_body = {} if args.extra_request_body: extra_request_body = json.loads(args.extra_request_body) # Set url if args.port is None: args.port = { "sglang": 30000, "lmdeploy": 23333, "vllm": 8000, }.get(args.backend, 30000) model_url = ( f"{args.base_url}/v1/models" if args.base_url else f"http://{args.host}:{args.port}/v1/models" ) if args.backend in ["sglang", "vllm", "lmdeploy"]: api_url = ( f"{args.base_url}/v1/chat/completions" if args.base_url else f"http://{args.host}:{args.port}/v1/chat/completions" ) base_url = ( f"http://{args.host}:{args.port}" if args.base_url is None else args.base_url ) # Get model name if args.model is None: if args.backend == "truss": print( "Please provide a model with `--model` when using truss backend. e.g. --model meta-llama/Llama-3.1-8B-Instruct" ) sys.exit(1) try: response = requests.get(model_url) model_list = response.json().get("data", []) args.model = model_list[0]["id"] if model_list else None except Exception as e: print(f"Failed to fetch model from {model_url}. Error: {e}") print( "Please specify the correct host and port using `--host` and `--port`." ) sys.exit(1) if args.model is None: print("No model specified or found. Please provide a model using `--model`.") sys.exit(1) # Dataset compatibility check if args.enable_multiturn: # TODO: Support multiturn for random if args.dataset_name not in ["sharegpt", "ultrachat", "loogle", "nextqa"]: print( "Multiturn conversation is only supported for sharegpt, ultrachat, loogle, and nextqa datasets." ) sys.exit(1) if args.enable_shared_prefix: if args.dataset_name not in ["loogle", "nextqa"]: print("Shared prefix is only supported for loogle and nextqa datasets.") sys.exit(1) print(f"{args}\n") # Read dataset backend = args.backend model_id = args.model tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model tokenizer = get_tokenizer(tokenizer_id) input_requests = get_dataset(args, tokenizer) return asyncio.run( benchmark( backend=backend, api_url=api_url, base_url=base_url, model_id=model_id, tokenizer=tokenizer, input_requests=input_requests, request_rate=args.request_rate, max_concurrency=args.max_concurrency, disable_tqdm=args.disable_tqdm, lora_name=args.lora_name, extra_request_body=extra_request_body, profile=args.profile, enable_shared_prefix=args.enable_shared_prefix, ) ) if __name__ == "__main__": parser = ArgumentParser(description="Benchmark the online serving throughput.") parser.add_argument( "--backend", type=str, choices=list(ASYNC_REQUEST_FUNCS.keys()), default="sglang", help="Must specify a backend, depending on the LLM Inference Engine.", ) parser.add_argument( "--base-url", type=str, default=None, help="Server or API base url if not using http host and port.", ) parser.add_argument( "--host", type=str, default="0.0.0.0", help="Default host is 0.0.0.0." ) parser.add_argument( "--port", type=int, help="If not set, the default port is configured according to its default value for different LLM Inference Engines.", ) parser.add_argument( "--dataset-name", type=str, default="sharegpt", choices=[ "sharegpt", "random", "generated-shared-prefix", "ultrachat", "loogle", "nextqa", ], help="Name of the dataset to benchmark on.", ) parser.add_argument( "--dataset-path", type=str, default="", help="Path to the dataset." ) parser.add_argument( "--model", type=str, help="Name or path of the model. If not set, the default model will request /v1/models for conf.", ) parser.add_argument( "--tokenizer", type=str, help="Name or path of the tokenizer. If not set, using the model conf.", ) parser.add_argument( "--chat-template", type=str, help="The buliltin chat template name or the path of the chat template file. This is only used for OpenAI-compatible API server.", ) parser.add_argument( "--num-prompts", type=int, default=1000, help="Number of prompts to process. Default is 1000.", ) parser.add_argument( "--fixed-output-len", type=int, default=None, help="Output length for each request. Overrides the output length from the dataset.", ) parser.add_argument( "--sharegpt-context-len", type=int, default=None, help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.", ) parser.add_argument( "--random-input-len", type=int, default=1024, help="Number of input tokens per request, used only for random dataset.", ) parser.add_argument( "--random-output-len", default=1024, type=int, help="Number of output tokens per request, used only for random dataset.", ) parser.add_argument( "--random-range-ratio", type=float, default=0.0, help="Range of sampled ratio of input/output length, " "used only for random dataset.", ) parser.add_argument( "--request-rate", type=float, default=float("inf"), help="Number of requests per second. If this is inf, then all the requests are sent at time 0. " "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.", ) parser.add_argument( "--max-concurrency", type=int, default=None, help="Maximum number of concurrent requests. This can be used " "to help simulate an environment where a higher level component " "is enforcing a maximum number of concurrent requests. While the " "--request-rate argument controls the rate at which requests are " "initiated, this argument will control how many are actually allowed " "to execute at a time. This means that when used in combination, the " "actual request rate may be lower than specified with --request-rate, " "if the server is not processing requests fast enough to keep up.", ) parser.add_argument( "--multi", action="store_true", help="Use request rate range rather than single value.", ) parser.add_argument( "--request-rate-range", type=str, default="2,34,2", help="Range of request rates in the format start,stop,step. Default is 2,34,2. It also supports a list of request rates, requiring the parameters to not equal three.", ) parser.add_argument("--output-file", type=str, help="Output JSONL file name.") parser.add_argument( "--enable-multiturn", action="store_true", help="Enable multiturn chat for online serving benchmarking. " "This option is effective on the following datasets: " "sharegpt, ultrachat, loogle, nextqa", ) parser.add_argument( "--enable-shared-prefix", action="store_true", help="Enable shared prefix for online serving benchmarking. " "This option is effective on the following datasets: " "loogle, nextqa", ) parser.add_argument( "--disable-shuffle", action="store_true", help="Disable shuffling datasets. This is useful to generate stable output " "in benchmarking", ) parser.add_argument( "--disable-tqdm", action="store_true", help="Specify to disable tqdm progress bar.", ) parser.add_argument( "--disable-stream", action="store_true", help="Disable streaming mode.", ) parser.add_argument( "--return-logprob", action="store_true", help="Return logprob.", ) parser.add_argument("--seed", type=int, default=1, help="The random seed.") parser.add_argument( "--disable-ignore-eos", action="store_true", help="Disable ignoring EOS.", ) parser.add_argument( "--extra-request-body", metavar='{"key1": "value1", "key2": "value2"}', type=str, help="Append given JSON object to the request payload. You can use this to specify" "additional generate params like sampling params.", ) parser.add_argument( "--apply-chat-template", action="store_true", help="Apply chat template", ) parser.add_argument( "--profile", action="store_true", help="Use Torch Profiler. The endpoint must be launched with " "SGLANG_TORCH_PROFILER_DIR to enable profiler.", ) parser.add_argument( "--lora-name", type=str, default=None, help="The name of LoRA adapter", ) group = parser.add_argument_group("generated-shared-prefix dataset arguments") group.add_argument( "--gsp-num-groups", type=int, default=64, help="Number of system prompt groups for generated-shared-prefix dataset", ) group.add_argument( "--gsp-prompts-per-group", type=int, default=16, help="Number of prompts per system prompt group for generated-shared-prefix dataset", ) group.add_argument( "--gsp-system-prompt-len", type=int, default=2048, help="Target length in tokens for system prompts in generated-shared-prefix dataset", ) group.add_argument( "--gsp-question-len", type=int, default=128, help="Target length in tokens for questions in generated-shared-prefix dataset", ) group.add_argument( "--gsp-output-len", type=int, default=256, help="Target length in tokens for outputs in generated-shared-prefix dataset", ) # videos specific parser.add_argument( "--max-frames", type=int, default=sys.maxsize, help="The maximum number of frames to extract from each video. " "This option is specific to the nextqa dataset (video benchmark). ", ) args = parser.parse_args() if args.enable_multiturn and args.enable_shared_prefix: parser.error( "--enable-multiturn and --enable-shared-prefix cannot be set at the same time." ) run_benchmark(args) ================================================ FILE: benchmark/hicache/data_processing.py ================================================ import json import os import pickle import random from typing import List, Optional, Tuple, Union import numpy as np from nextqa import NExTQALoader # from nextqa.video import , VideoPrompt from tqdm.asyncio import tqdm from transformers import PreTrainedTokenizerBase from sglang.benchmark.datasets.common import ( SHAREGPT_FILENAME, SHAREGPT_REPO_ID, gen_prompt, ) from sglang.benchmark.datasets.generated_shared_prefix import get_gen_prefix_cache_path from sglang.benchmark.utils import download_and_cache_hf_file from sglang.lang.chat_template import get_chat_template, get_chat_template_by_model_path from sglang.srt.entrypoints.openai.protocol import ChatCompletionMessageContentPart from sglang.utils import encode_video_base64 # type of content fields, can be only prompts or with images/videos MsgContent = Union[str, List[ChatCompletionMessageContentPart]] # A list of all the conversations. Each conversation is a list of # tuples. If multiturn is not enabled, the length of list is 1, # containing only the first Q&A pair. # For the shared prefix workload (synthetic, loogle, nextqa), it # is a list of conversations sharing the same prefix (synthetic, # doc, video) SampleOutput = List[List[Tuple[MsgContent, int, int]]] def common_filter_chat( num_requests: int, new_dataset: List, tokenizer: PreTrainedTokenizerBase, min_prompt_len: Optional[int], min_output_len: Optional[int], max_prompt_len: Optional[int], max_output_len: Optional[int], fixed_output_len: Optional[int], ) -> SampleOutput: # Filter out sequences that are too long or too short filtered_dataset: SampleOutput = [] l = 0 input_tokens = 0 output_tokens = 0 while l < num_requests: for i in range(len(new_dataset)): if l == num_requests: break processed = [] for j in new_dataset[i]: # Tokenize the prompts and completions. prompt = j[0] prompt_token_ids = tokenizer.encode(prompt) prompt_len = len(prompt_token_ids) completion = j[1] completion_token_ids = tokenizer.encode(completion) output_len = ( len(completion_token_ids) if fixed_output_len is None else fixed_output_len ) if ( min_prompt_len is not None and prompt_len < min_prompt_len or min_output_len is not None and output_len < min_output_len or max_prompt_len is not None and prompt_len > max_prompt_len or max_output_len is not None and output_len > max_output_len ): # Prune too short sequences. continue input_tokens += prompt_len output_tokens += output_len processed.append((prompt, prompt_len, output_len)) if len(processed) != 0: filtered_dataset.append(processed) l += 1 print(f"#Input tokens: {input_tokens}") print(f"#Output tokens: {output_tokens}") return filtered_dataset def sample_sharegpt_requests( dataset_path: str, num_requests: int, tokenizer: PreTrainedTokenizerBase, disable_shuffle: bool = False, enable_multiturn: bool = True, fixed_output_len: Optional[int] = None, ) -> SampleOutput: if fixed_output_len is not None and fixed_output_len < 4: raise ValueError("output_len too small") # Download sharegpt if necessary if not os.path.isfile(dataset_path): dataset_path = download_and_cache_hf_file( repo_id=SHAREGPT_REPO_ID, filename=SHAREGPT_FILENAME, ) # Load the dataset. with open(dataset_path) as f: dataset = json.load(f) # Filter out the conversations with less than 2 turns. dataset = [data for data in dataset if len(data["conversations"]) >= 2] # Keep one conversation in one list new_dataset = [] for data in dataset: if len(data["conversations"]) % 2 != 0: continue if data["conversations"][0]["from"] != "human": continue chat = [] total_len = 2 if enable_multiturn: total_len = len(data["conversations"]) for i in range(0, total_len, 2): # One user One Assistant chat.append( ( data["conversations"][i]["value"], data["conversations"][i + 1]["value"], ) ) new_dataset.append(chat) if not disable_shuffle: # Shuffle the dataset. random.shuffle(new_dataset) # Filter out sequences that are too long or too short filtered_dataset: SampleOutput = common_filter_chat( num_requests, new_dataset, tokenizer, 4, 4, None, None, fixed_output_len ) return filtered_dataset def sample_ultrachat_requests( dataset_path: str, num_requests: int, tokenizer: PreTrainedTokenizerBase, disable_shuffle: bool = False, enable_multiturn: bool = True, fixed_output_len: Optional[int] = None, ) -> SampleOutput: if fixed_output_len is not None and fixed_output_len < 4: raise ValueError("output_len too small") # Load the dataset dataset = [] with open(dataset_path) as f: while True: line = f.readline() if not line: break dataset.append(json.loads(line)) # Filter out the conversations with less than 2 turns. dataset = [data for data in dataset if len(data["data"]) >= 2] # Keep one conversation in one list new_dataset = [] for data in dataset: if len(data["data"]) % 2 != 0: continue chat = [] total_len = 2 if enable_multiturn: total_len = len(data["data"]) for i in range(0, total_len, 2): # One user One Assistant chat.append((data["data"][i], data["data"][i + 1])) new_dataset.append(chat) # Shuffle the dataset. if not disable_shuffle: random.shuffle(new_dataset) # Filter out sequences that are too long or too short filtered_dataset: SampleOutput = common_filter_chat( num_requests, new_dataset, tokenizer, 4, 4, None, None, fixed_output_len ) return filtered_dataset def sample_loogle_requests( dataset_path: str, num_requests: int, tokenizer: PreTrainedTokenizerBase, disable_shuffle: bool = False, enable_multiturn: bool = True, enable_shared_prefix: bool = False, fixed_output_len: Optional[int] = None, ) -> SampleOutput: if fixed_output_len is not None and fixed_output_len < 4: raise ValueError("output_len too small") # Load the dataset dataset = [] with open(dataset_path) as f: while True: line = f.readline() if not line: break dataset.append(json.loads(line)) # Keep one conversation in one list new_dataset = [] # TODO: Add shared prefix support for loogle # NOTE: Now we preprocess it only for chat for data in dataset: chat = [] if ( "qa_pairs" not in data or data["qa_pairs"] == "none" or len(data["qa_pairs"]) == 0 ): # If Q is none (for summarization), # We add a question for summarization # And keep the summary up to 1024 words chat.append( ( "Input: " + data["input"] + " Question: " + "Please summarize the input", data["input"][:1024], ) ) new_dataset.append(chat) else: qa_pairs = eval(data["qa_pairs"]) for i, qa in enumerate(qa_pairs): if i == 0 or enable_shared_prefix: # Combine input with the first Q chat.append( ("Input: " + data["input"] + " Question: " + qa["Q"], qa["A"]) ) elif enable_multiturn: chat.append((qa["Q"], qa["A"])) new_dataset.append(chat) # Shuffle the dataset. if not disable_shuffle: random.shuffle(new_dataset) # Filter out sequences that are too long or too short filtered_dataset: SampleOutput = common_filter_chat( num_requests, new_dataset, tokenizer, 4, None, None, None, fixed_output_len ) return filtered_dataset def sample_nextqa_requests( dataset_path: str, num_requests: int, tokenizer: PreTrainedTokenizerBase, max_frames: int, # Specific for video model_path: str, disable_shuffle: bool = False, enable_multiturn: bool = True, # No multiturn support for now backend: str = "sglang-oai", chat_template_name: Optional[str] = None, fixed_output_len: Optional[int] = None, ) -> SampleOutput: """ Example of messages: message = { "role": "user", "content": [ {"type": "image_url", "image_url": {"url": base64_data}}, {"type": "text", "text": video.prompt}, ], } """ if fixed_output_len is None: fixed_output_len = 4096 # TODO: Check for multiturn dataset = NExTQALoader(video_dir=dataset_path, max_frames=max_frames) new_dataset = [] for v in dataset: new_dataset.append(v) if not disable_shuffle: random.shuffle(new_dataset) # TODO: prompt len can get from server side filtered_dataset = [] l = 0 while l < num_requests: for i in range(len(new_dataset)): if l == num_requests: break video = new_dataset[i] # text prompt prompt = video.prompt # NOTE: Chat Template is a must for video benchmark because we have to # add special image token for later expansion if backend == "sglang" or backend == "sglang-native": if "chat_template" in tokenizer.init_kwargs: chat_template = get_chat_template(tokenizer.get_chat_template()) elif chat_template_name is not None: chat_template = get_chat_template(chat_template_name) else: chat_template = get_chat_template_by_model_path(model_path) prompt = chat_template.image_token + prompt prompt_token_ids = tokenizer(prompt).input_ids prompt_len = len(prompt_token_ids) output_len = fixed_output_len # max output len, not real output len # video input base64_data = encode_video_base64(video.path, video.num_frames) # NOTE: This will be replaced by the expanded length from the server prompt_len += video.num_frames # add to content content = [ {"type": "image_url", "image_url": {"url": base64_data}}, {"type": "text", "text": prompt}, ] filtered_dataset.append([(content, prompt_len, output_len)]) l += 1 return filtered_dataset def sample_random_requests( input_len: int, output_len: int, num_prompts: int, range_ratio: float, tokenizer: PreTrainedTokenizerBase, dataset_path: str, disable_shuffle: bool = False, ) -> SampleOutput: input_lens = np.random.randint( max(int(input_len * range_ratio), 1), input_len + 1, size=num_prompts, ) output_lens = np.random.randint( int(output_len * range_ratio), output_len + 1, size=num_prompts, ) if True: # Sample token ids from ShareGPT and repeat/truncate them to satisfy the input_lens # Download sharegpt if necessary if not os.path.isfile(dataset_path): dataset_path = download_and_cache_hf_file( repo_id=SHAREGPT_REPO_ID, filename=SHAREGPT_FILENAME, ) # Load the dataset. with open(dataset_path) as f: dataset = json.load(f) # Filter out the conversations with less than 2 turns. dataset = [data for data in dataset if len(data["conversations"]) >= 2] # Only keep the first two turns of each conversation. dataset = [ (data["conversations"][0]["value"], data["conversations"][1]["value"]) for data in dataset ] if not disable_shuffle: # Shuffle the dataset. random.shuffle(dataset) # Filter out sequences that are too long or too short input_requests: SampleOutput = [] for data in dataset: i = len(input_requests) if i == num_prompts: break # Tokenize the prompts and completions. prompt = data[0] prompt_token_ids = tokenizer.encode(prompt) prompt_len = len(prompt_token_ids) # Skip empty prompt if prompt_len == 0: continue if prompt_len > input_lens[i]: input_ids = prompt_token_ids[: input_lens[i]] else: ratio = (input_lens[i] + prompt_len - 1) // prompt_len input_ids = (prompt_token_ids * ratio)[: input_lens[i]] prompt = tokenizer.decode(input_ids) input_requests.append([(prompt, int(input_lens[i]), int(output_lens[i]))]) else: # Sample token ids from random integers. This can cause some NaN issues. offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts) input_requests = [] for i in range(num_prompts): prompt = tokenizer.decode( [ (offsets[i] + i + j) % tokenizer.vocab_size for j in range(input_lens[i]) ] ) input_requests.append([(prompt, int(input_lens[i]), int(output_lens[i]))]) print(f"#Input tokens: {np.sum(input_lens)}") print(f"#Output tokens: {np.sum(output_lens)}") return input_requests def sample_generated_shared_prefix_requests( num_groups: int, prompts_per_group: int, system_prompt_len: int, question_len: int, output_len: int, tokenizer: PreTrainedTokenizerBase, args, disable_shuffle: bool = False, ) -> SampleOutput: """Generate benchmark requests with shared system prompts using random tokens and caching.""" cache_path = get_gen_prefix_cache_path( args.seed, num_groups, prompts_per_group, system_prompt_len, question_len, output_len, tokenizer, ) # Try to load from cache first if cache_path.exists(): print(f"\nLoading cached generated input data from {cache_path}") with open(cache_path, "rb") as f: return pickle.load(f) print("\nGenerating new input data...") # Generate system prompts for each group system_prompts = [] for _ in range(num_groups): system_prompt = gen_prompt(tokenizer, system_prompt_len) system_prompts.append(system_prompt) # Generate questions questions = [] for _ in range(num_groups * prompts_per_group): question = gen_prompt(tokenizer, question_len) questions.append(question) # Combine system prompts with questions input_requests = [] total_input_tokens = 0 total_output_tokens = 0 for group_idx in tqdm(range(num_groups), desc="Generating system prompt"): system_prompt = system_prompts[group_idx] input_requests.append([]) for prompt_idx in tqdm( range(prompts_per_group), desc="Generating questions", leave=False ): question = questions[group_idx * prompts_per_group + prompt_idx] full_prompt = f"{system_prompt}\n\n{question}" prompt_len = len(tokenizer.encode(full_prompt)) input_requests[-1].append((full_prompt, prompt_len, output_len)) total_input_tokens += prompt_len total_output_tokens += output_len if not disable_shuffle: # Shuffle questions random.shuffle(input_requests) # Print statistics print(f"\nGenerated shared prefix dataset statistics:") print(f"Number of groups: {num_groups}") print(f"Prompts per group: {prompts_per_group}") print(f"Total prompts: {len(input_requests) * prompts_per_group}") print(f"Total input tokens: {total_input_tokens}") print(f"Total output tokens: {total_output_tokens}") print( f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens" ) print( f"Average question length: {sum(len(tokenizer.encode(q)) for q in questions) / len(questions):.1f} tokens\n" ) # Save to cache cache_path.parent.mkdir(parents=True, exist_ok=True) print(f"Caching generated input data to {cache_path}") with open(cache_path, "wb") as f: pickle.dump(input_requests, f) return input_requests def get_dataset(args, tokenizer): if args.dataset_name == "sharegpt": input_requests = sample_sharegpt_requests( dataset_path=args.dataset_path, num_requests=args.num_prompts, tokenizer=tokenizer, disable_shuffle=args.disable_shuffle, enable_multiturn=args.enable_multiturn, fixed_output_len=args.fixed_output_len, ) elif args.dataset_name == "ultrachat": input_requests = sample_ultrachat_requests( dataset_path=args.dataset_path, num_requests=args.num_prompts, tokenizer=tokenizer, disable_shuffle=args.disable_shuffle, enable_multiturn=args.enable_multiturn, fixed_output_len=args.fixed_output_len, ) elif args.dataset_name == "loogle": input_requests = sample_loogle_requests( dataset_path=args.dataset_path, num_requests=args.num_prompts, tokenizer=tokenizer, disable_shuffle=args.disable_shuffle, enable_multiturn=args.enable_multiturn, enable_shared_prefix=args.enable_shared_prefix, fixed_output_len=args.fixed_output_len, ) elif args.dataset_name == "nextqa": input_requests = sample_nextqa_requests( dataset_path=args.dataset_path, num_requests=args.num_prompts, tokenizer=tokenizer, max_frames=args.max_frames, model_path=args.model, disable_shuffle=args.disable_shuffle, enable_multiturn=args.enable_multiturn, backend=args.backend, chat_template_name=args.chat_template, fixed_output_len=args.fixed_output_len, ) elif args.dataset_name == "random": input_requests = sample_random_requests( input_len=args.random_input_len, output_len=args.random_output_len, num_prompts=args.num_prompts, range_ratio=args.random_range_ratio, tokenizer=tokenizer, dataset_path=args.dataset_path, ) elif args.dataset_name == "generated-shared-prefix": input_requests = sample_generated_shared_prefix_requests( num_groups=args.gsp_num_groups, prompts_per_group=args.gsp_prompts_per_group, system_prompt_len=args.gsp_system_prompt_len, question_len=args.gsp_question_len, output_len=args.gsp_output_len, args=args, tokenizer=tokenizer, ) else: raise ValueError(f"Unknown dataset: {args.dataset_name}") return input_requests ================================================ FILE: benchmark/hicache/download.sh ================================================ #!/usr/bin/bash # The usage function usage() { echo "Usage: $0 {sharegpt|ultragpt|loogle|nextqa|all}" exit 1 } # The download function download() { case "$1" in sharegpt) echo $1 wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json ;; ultragpt) echo $1 # Questions about the world wget https://cloud.tsinghua.edu.cn/seafhttp/files/be1d7b87-22ca-449e-a6a7-c61d1ea7e010/ultrachat_release_230407.json # Writing and Creation wget https://cloud.tsinghua.edu.cn/seafhttp/files/61742d2a-25e2-4d08-b2b9-15f47ae50ace/ultrachat_material_release_230417.json wget https://cloud.tsinghua.edu.cn/seafhttp/files/f71f6aa6-d346-4b16-85b7-8502efa3d608/ultrachat_material_release_230412.json # External materials wget https://cloud.tsinghua.edu.cn/seafhttp/files/42d22e28-e899-4975-a70f-5eda163e265d/ultrachat_existent_material_release_230420.json.gz gunzip ultrachat_existent_material_release_230420.json.gz ;; loogle) echo $1 git lfs install git clone git@hf.co:datasets/bigainlco/LooGLE unzip LooGLE/data.zip ;; nextqa) echo $1 git lfs install git clone https://huggingface.co/datasets/lmms-lab/NExTQA unzip NExTQA/videos.zip ;; *) usage exit 1 ;; esac } # Arg check if [ "$#" -ne 1 ]; then usage fi # Invoke case "$1" in sharegpt|ultragpt|loogle|nextqa) download "$1" ;; all) download sharegpt download ultragpt download loogle download nextqa ;; *) usage ;; esac ================================================ FILE: benchmark/hicache/nextqa.py ================================================ import os import sys from typing import List import av from datasets import load_dataset def find_video_files(video_dir) -> List[str]: if os.path.isfile(video_dir): return [video_dir] video_files = [] for root, dirs, files in os.walk(video_dir): for file in files: if file.endswith((".mp4", ".avi", ".mov")): video_files.append(os.path.join(root, file)) # if file is dir elif os.path.isdir(file): video_files.extend(find_video_files(file)) return video_files def video_frames(video_path, max_frames) -> int: container = av.open(video_path) total_frames = container.streams.video[0].frames return min(total_frames, max_frames) class Video: def __init__(self, video_path, num_frames): self.path = video_path self.num_frames = num_frames def __str__(self): return f"Video({self.path}, {self.num_frames})" def __iter__(self): return iter((self.path, self.num_frames)) class VideoPrompt(Video): def __init__(self, video_path, num_frames, prompt): super().__init__(video_path, num_frames) self.prompt = prompt def __str__(self): return f"VideoPrompt({self.path}, {self.num_frames}, {self.prompt})" def __iter__(self): return iter((self.path, self.num_frames, self.prompt)) class VideoLoader: pass class VideoFileLoader(VideoLoader): """ Load all the videos in a directory """ def __init__(self, video_dir, batch_size=1, max_frames=sys.maxsize): super().__init__() self.video_dir = video_dir self.video_files = find_video_files(video_dir) self.batch_size = batch_size self.max_frames = max_frames print(f"batch_size: {batch_size}, max_frames: {max_frames}") def __iter__(self): # (file, number of frames) if self.batch_size == 1: for video_file in self.video_files: yield Video(video_file, video_frames(video_file, self.max_frames)) else: batch = [] for video_file in self.video_files: video = Video(video_file, video_frames(video_file, self.max_frames)) batch.append(video) if len(batch) == self.batch_size: yield batch batch = [] class NExTQALoader(VideoLoader): """ Load vdideos and prompts from NExT dataset set: train, test or validation """ def __init__( self, video_dir, batch_size=1, max_frames=sys.maxsize, dset="test", task="OE" ): """ task: 'MV' or 'OE' """ super().__init__() self.task = task print(f"Loading the {dset} data of {task} from lmms-lab/NExTQA") self.ds = load_dataset("lmms-lab/NExTQA", task) self.ds = self.ds[dset] # self.n = ds.num_rows self.video_dir = video_dir self.video_files = find_video_files(video_dir) self.video_to_path = dict() for video_file in self.video_files: video_id = video_file.split("/")[-1].split(".")[0] self.video_to_path[video_id] = video_file self.batch_size = batch_size self.max_frames = max_frames def get_video_prompt(self, entry, max_frames) -> VideoPrompt: # Get video video_id = entry["video"] video_path = self.video_to_path[video_id] assert os.path.exists(video_path), f"Video not found: {video_path}" num_frames = min(entry["frame_count"], max_frames) video = Video(video_path, num_frames) prompt = entry["question"] + "?" if self.task == "MC": # add choices prompt += f' a0: {entry["a0"]}, a1: {entry["a1"]}, a2: {entry["a2"]}, a3: {entry["a3"]}' return VideoPrompt(video_path, num_frames, prompt) def __iter__(self): if self.batch_size == 1: for entry in self.ds: yield self.get_video_prompt(entry, self.max_frames) else: batch = [] for entry in self.ds: video = self.get_video_prompt(entry, self.max_frames) batch.append(video) if len(batch) == self.batch_size: yield batch batch = [] # main if __name__ == "__main__": video_dir = "./videos" # video_loader = VideoFileLoader(video_dir, batch_size=16) # for batch in video_loader: # print(f"Number of videos in batch: {len(batch)}") # for video_file, num_frames in batch: # print(f"Video: {video_file} number of frames: {num_frames}") video_loader = NExTQALoader(video_dir, batch_size=16, dset="test", task="OE") for batch in video_loader: print(f"Number of videos in batch: {len(batch)}") for video_file, num_frames, prompt in batch: print( f"Video: {video_file} number of frames: {num_frames}, prompt: {prompt}" ) # break # for video_file, prompt in batch: # print(f"Video: {video_file} prompt: {prompt}") # break ================================================ FILE: benchmark/hicache/perf.py ================================================ from __future__ import annotations from typing import Any, Callable, NamedTuple import torch def jit_hicache_impl( k_cache_dst: torch.Tensor, v_cache_dst: torch.Tensor, indices_dst: torch.Tensor, k_cache_src: torch.Tensor, v_cache_src: torch.Tensor, indices_src: torch.Tensor, item_bytes: int, block_quota: int, ) -> None: from sglang.jit_kernel.hicache import transfer_hicache_one_layer _ = item_bytes transfer_hicache_one_layer( k_cache_dst=k_cache_dst, v_cache_dst=v_cache_dst, indices_dst=indices_dst, k_cache_src=k_cache_src, v_cache_src=v_cache_src, indices_src=indices_src, block_quota=block_quota, ) def ref_hicache_impl( k_cache_dst: torch.Tensor, v_cache_dst: torch.Tensor, indices_dst: torch.Tensor, k_cache_src: torch.Tensor, v_cache_src: torch.Tensor, indices_src: torch.Tensor, item_bytes: int, block_quota: int, ) -> None: from sgl_kernel import transfer_kv_per_layer transfer_kv_per_layer( src_k=k_cache_src, src_v=v_cache_src, dst_k=k_cache_dst, dst_v=v_cache_dst, src_indices=indices_src, dst_indices=indices_dst, item_size=item_bytes, block_quota=block_quota, ) class HicacheBenchArgs(NamedTuple): cache_item_size: int dtype: torch.dtype block_quota: int def perf(f: Callable[[], Any], loop: int = 100) -> float: tic = torch.cuda.Event(enable_timing=True) toc = torch.cuda.Event(enable_timing=True) torch.cuda.synchronize() # warm up f() torch.cuda._sleep(10**8) tic.record() for _ in range(loop): f() toc.record() toc.synchronize() return tic.elapsed_time(toc) / loop @torch.inference_mode() def test_hicache_kernel(args: HicacheBenchArgs) -> None: CACHE_ITEM_SIZE, DTYPE, BLOCK_QUOTA = args CUDA_CACHE_SIZE = 1024 * 1024 HOST_CACHE_SIZE = CUDA_CACHE_SIZE * 2 cuda_cache = torch.randn( (2, CUDA_CACHE_SIZE, CACHE_ITEM_SIZE), dtype=DTYPE, device="cuda", ) host_cache = torch.empty( (2, HOST_CACHE_SIZE, CACHE_ITEM_SIZE), dtype=DTYPE, device="cpu", pin_memory=True, ) ITEM_BYTES = cuda_cache.element_size() * CACHE_ITEM_SIZE def _gen_indices(size: int, bs: int) -> torch.Tensor: assert bs <= size result = ( (torch.randperm(size, dtype=torch.int64, device="cuda")[:bs]).sort().values ) if not (torch.all(result >= 0) and torch.all(result < size)): where = (result < 0) | (result >= size) place = where.nonzero(as_tuple=False) print("Invalid indices at positions:", place) print("Invalid indices values:", result[place]) raise ValueError("Generated invalid indices") return result def _calc_tput(dur: float) -> float: return (MEM / (1024**3)) / (dur / 1000) # GB/s def _gain_str(aot_dur: float, jit_dur: float) -> str: gain = 100 * (aot_dur / jit_dur - 1) if gain >= 0: return f"+{gain:>6.2f}%" else: return f"-{-gain:>6.2f}%" print(f"{CACHE_ITEM_SIZE = }, {DTYPE = }, {BLOCK_QUOTA = }") def _fast_test_correctness(bs: int): src_indices = _gen_indices(CUDA_CACHE_SIZE, bs) dst_indices = _gen_indices(HOST_CACHE_SIZE, bs) host_cache_cuda = torch.randn_like(host_cache, device="cuda") host_cache.copy_(host_cache_cuda, non_blocking=True) # copy from cuda to host jit_hicache_impl( k_cache_dst=host_cache[0], v_cache_dst=host_cache[1], indices_dst=dst_indices, k_cache_src=cuda_cache[0], v_cache_src=cuda_cache[1], indices_src=src_indices, item_bytes=ITEM_BYTES, block_quota=BLOCK_QUOTA, ) dst_indices = dst_indices.cpu() assert torch.all( host_cache[0][dst_indices].cuda() == cuda_cache[0][src_indices] ) BS_RANGE = [2**n for n in range(8, 18)] for bs in BS_RANGE: _fast_test_correctness(bs) print("Correctness passed! Start HiCache kernel performance test...") print("=" * 70) for bs in BS_RANGE: indices_dst = _gen_indices(CUDA_CACHE_SIZE, bs) indices_src = _gen_indices(HOST_CACHE_SIZE, bs) MEM = 2 * bs * ITEM_BYTES def _run_kernel_h2d(impl): return impl( k_cache_dst=cuda_cache[0], v_cache_dst=cuda_cache[1], indices_dst=indices_dst, k_cache_src=host_cache[0], v_cache_src=host_cache[1], indices_src=indices_src, item_bytes=ITEM_BYTES, block_quota=BLOCK_QUOTA, ) our_h2d_dur = perf(lambda: _run_kernel_h2d(jit_hicache_impl)) ref_h2d_dur = perf(lambda: _run_kernel_h2d(ref_hicache_impl)) print( f"{bs = :6d}, H->D", f"| aot {_calc_tput(ref_h2d_dur):<6.2f} GB/s", f"| jit {_calc_tput(our_h2d_dur):<6.2f} GB/s", f"| {_gain_str(ref_h2d_dur, our_h2d_dur)}", ) print("=" * 70) for bs in BS_RANGE: indices_dst = _gen_indices(HOST_CACHE_SIZE, bs) indices_src = _gen_indices(CUDA_CACHE_SIZE, bs) MEM = 2 * bs * ITEM_BYTES def _run_kernel_d2h(impl): return impl( k_cache_dst=host_cache[0], v_cache_dst=host_cache[1], indices_dst=indices_dst, k_cache_src=cuda_cache[0], v_cache_src=cuda_cache[1], indices_src=indices_src, item_bytes=ITEM_BYTES, block_quota=BLOCK_QUOTA, ) our_d2h_dur = perf(lambda: _run_kernel_d2h(jit_hicache_impl)) ref_d2h_dur = perf(lambda: _run_kernel_d2h(ref_hicache_impl)) print( f"{bs = :6d}, D->H", f"| aot {_calc_tput(ref_d2h_dur):<6.2f} GB/s", f"| jit {_calc_tput(our_d2h_dur):<6.2f} GB/s", f"| {_gain_str(ref_d2h_dur, our_d2h_dur)}", ) print("=" * 70) def main() -> None: torch.cuda.set_device(0) stream = torch.cuda.Stream() torch.cuda.set_stream(stream) tic = torch.cuda.Event(enable_timing=True) toc = torch.cuda.Event(enable_timing=True) BUF_SIZE = 1024 * 1024 * 1024 cuda_mem = torch.empty(BUF_SIZE, dtype=torch.uint8, device="cuda") host_mem = torch.empty(BUF_SIZE, dtype=torch.uint8, device="cpu", pin_memory=True) # test peak bandwidth tic.record() cuda_mem.copy_(host_mem, non_blocking=True) toc.record() toc.synchronize() dur = tic.elapsed_time(toc) print(f"Peak H->D Bandwidth: {(BUF_SIZE / (1024**3)) / (dur / 1000):.2f} GB/s") tic.record() host_mem.copy_(cuda_mem, non_blocking=True) toc.record() toc.synchronize() dur = tic.elapsed_time(toc) print(f"Peak D->H Bandwidth: {(BUF_SIZE / (1024**3)) / (dur / 1000):.2f} GB/s") for block_quota in [1, 2, 3, 4]: for cache_item_size in [128, 256, 512, 1024]: args = HicacheBenchArgs( cache_item_size=cache_item_size, dtype=torch.float16, block_quota=block_quota, ) test_hicache_kernel(args) if __name__ == "__main__": main() ================================================ FILE: benchmark/json_decode_regex/README.md ================================================ ## Run benchmark ### Build dataset ``` pip install wikipedia python3 build_dataset.py ``` ### Dependencies ``` llama_cpp_python 0.2.19 guidance 0.1.10 vllm 0.2.5 outlines 0.0.22 ``` ### Benchmark sglang Run Llama-7B ``` python3 -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` Run Mixtral-8x7B ``` python3 -m sglang.launch_server --model-path mistralai/Mixtral-8x7B-Instruct-v0.1 --port 30000 --tp-size 8 ``` Benchmark ``` python3 bench_sglang.py --num-questions 10 ``` ### Benchmark Outlines + vLLM Run Llama-7B ``` python3 -m outlines.serve.serve --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` Benchmark ``` python3 bench_other.py --backend outlines --num-questions 10 ``` ### Benchmark guidance Run Llama-7B and benchmark ``` python3 bench_other.py --backend guidance --num-questions 10 --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` ================================================ FILE: benchmark/json_decode_regex/bench_other.py ================================================ import argparse import json import time from concurrent.futures import ThreadPoolExecutor from functools import partial from tqdm import tqdm from sglang.lang.ir import REGEX_FLOAT, REGEX_INT, REGEX_STR from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import dump_state_text, read_jsonl REGEX_LIST = r"\[(" + REGEX_STR + ", )*" + REGEX_STR + r"\]" # fmt: off def json_decode(document, generate): s = "Please extract the information of a city from the following wikipedia page.\n" s += "Page begin.\n" + document + "Page end.\n" s += "Here is the name, country, and symbol of the city in JSON format.\n" s += "{\n" s += ' "name": ' s += generate(s, max_tokens=8, regex=REGEX_STR + ",") + "\n" s += ' "country": ' s += generate(s, max_tokens=8, regex=REGEX_STR + ",") + "\n" s += ' "latitude": ' s += generate(s, max_tokens=8, regex=REGEX_FLOAT + ",") + "\n" s += ' "population": ' s += generate(s, max_tokens=8, regex=REGEX_INT + ",") + "\n" s += ' "top 3 landmarks": ' s += generate(s, max_tokens=24, regex=REGEX_LIST) + "\n" s += "}\n" return s # fmt: on def main(args): lines = read_jsonl(args.data_path) arguments = [] for i in range(len(lines[: args.num_questions])): arguments.append( { "document": lines[i]["document"], } ) states = [None] * len(arguments) # Select backend call_generate = partial(get_call_generate(args), temperature=0) # Run requests def get_one_answer(i): states[i] = json_decode(generate=call_generate, **arguments[i]) tic = time.perf_counter() if args.parallel == 1: for i in tqdm(range(len(arguments))): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: rets = list( tqdm( executor.map(get_one_answer, list(range(len(arguments)))), total=len(arguments), ) ) for _ in rets: pass latency = time.perf_counter() - tic # Compute accuracy print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "json_decode_regex", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_questions, "other": { "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="questions.jsonl") parser.add_argument("--num-questions", type=int, default=20) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/json_decode_regex/bench_sglang.py ================================================ import argparse import json import time import sglang as sgl from sglang.lang.ir import REGEX_FLOAT, REGEX_INT, REGEX_STR from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text, read_jsonl REGEX_LIST = r"\[(" + REGEX_STR + ", )*" + REGEX_STR + r"\]" # fmt: off @sgl.function def json_warm_up(s): s += "The information about Hogwarts is in the following JSON format.\n" with s.var_scope("json_output"): s += "{\n" s += ' "name": ' + sgl.gen("name", max_tokens=8, regex=REGEX_STR + ",") + "\n" s += ' "country": ' + sgl.gen("country", max_tokens=8, regex=REGEX_STR + ",") + "\n" s += ' "latitude": ' + sgl.gen("latitude", max_tokens=8, regex=REGEX_FLOAT + ",") + "\n" s += ' "population": ' + sgl.gen("population", max_tokens=8, regex=REGEX_INT + ",") + "\n" s += ' "top 3 landmarks": ' + sgl.gen( "landmarks", max_tokens=24, regex=REGEX_LIST) + "\n" s += "}\n" print(f'The warmp up json result is:\n{s["json_output"]}') # fmt: on # fmt: off @sgl.function def json_decode(s, document): s += "Please extract the information of a city from the following wikipedia page.\n" s += "Page begin.\n" + document + "Page end.\n" s += "Here is the name, country, and symbol of the city in JSON format.\n" with s.var_scope("json_output"): s += "{\n" s += ' "name": ' + sgl.gen("name", max_tokens=8, regex=REGEX_STR + ",") + "\n" s += ' "country": ' + sgl.gen("country", max_tokens=8, regex=REGEX_STR + ",") + "\n" s += ' "latitude": ' + sgl.gen("latitude", max_tokens=8, regex=REGEX_FLOAT + ",") + "\n" s += ' "population": ' + sgl.gen("population", max_tokens=8, regex=REGEX_INT + ",") + "\n" s += ' "top 3 landmarks": ' + sgl.gen( "landmarks", max_tokens=24, regex=REGEX_LIST) + "\n" s += "}\n" # fmt: on def main(args): lines = read_jsonl(args.data_path) lines = list(lines) arguments = [] for i in range(len(lines[: args.num_questions])): arguments.append( { "document": lines[i]["document"], } ) # Select backend backend = select_sglang_backend(args) sgl.set_default_backend(backend) # Warm up json_warm_up.run().sync() # Run requests tic = time.perf_counter() states = json_decode.run_batch( arguments, temperature=0, num_threads=args.parallel, progress_bar=True ) latency = time.perf_counter() - tic # Compute accuracy print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(f"tmp_{args.backend}_json_results.txt", "w") as fout: for state in states: fout.write(state["json_output"] + "\n") with open(args.result_file, "a") as fout: value = { "task": "json_decode_regex", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_questions, "other": { "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="questions.jsonl") parser.add_argument("--num-questions", type=int, default=20) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/json_decode_regex/build_dataset.py ================================================ import json import transformers import wikipedia model_path = "meta-llama/Llama-2-7b-chat-hf" t = transformers.AutoTokenizer.from_pretrained(model_path) city_names = [ "los angles", "london", "tokyo", "beijing", "singapore", "paris", "dubai", "sydney", "moscow", "rome", "toronto", "rio de janeiro", "istanbul", "berlin", "auckland", "buenos aires", "mexico city", "mumbai", "seoul", "bangkok", "cairo", "athens", "jerusalem", ] def get_content(city_name): content = str(wikipedia.page(city_name).content) content = content.replace("\n\n", "\n") tokens = t.encode(content) expected_tokens = 3000 truncate_len = int((expected_tokens / len(tokens)) * len(content)) truncate_content = content[:truncate_len] truncate_tokens = t.encode(truncate_content) # Count token print( f"city_name: {city_name}, #tokens: {len(tokens)}, #truncate tokens: {len(truncate_tokens)}" ) return truncate_content if __name__ == "__main__": with open("questions.jsonl", "w") as fout: for city_name in city_names: truncate_content = get_content(city_name) fout.write(json.dumps({"document": truncate_content}) + "\n") ================================================ FILE: benchmark/json_jump_forward/README.md ================================================ ## Run benchmark ### Dependencies ``` llama_cpp_python 0.2.38 guidance 0.1.10 vllm 0.2.7 outlines 0.0.25 ``` ### Build dataset When benchmarking long document information retrieval, run the following command to build the dataset: ```bash pip install wikipedia python3 build_dataset.py ``` ### Benchmark sglang Run Llama-7B ```bash python3 -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` Benchmark Character Generation ```bash python3 bench_sglang.py --mode character ``` Benchmark City Information Retrieval ```bash python3 bench_sglang.py --mode city ``` ### Benchmark Outlines + vLLM Run Llama-7B ```bash python3 -m outlines.serve.serve --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` Benchmark Character Generation ```bash python3 bench_other.py --mode character --backend outlines ``` Benchmark City Information Retrieval ```bash python3 bench_other.py --mode city --backend outlines ``` ### Benchmark guidance Run Llama-7B and benchmark character generation ```bash python3 bench_other.py --mode character --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` Run Llama-7B and benchmark city information retrieval ```bash python3 bench_other.py --mode city --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` ### Benchmark lmql Run Llama-7B and benchmark character generation ``` python3 bench_other.py --mode character --backend lmql --parallel 1 ``` Run Llama-7B and benchmark city information retrieval ``` python3 bench_other.py --mode city --backend lmql --parallel 1 ``` ================================================ FILE: benchmark/json_jump_forward/bench_other.py ================================================ import argparse import json import time from concurrent.futures import ThreadPoolExecutor from functools import partial import guidance from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import dump_state_text, read_jsonl # there are some FSM bugs with json regex converted from pydantic model # here use a string regex instead # regex_string = build_regex_from_object(HarryPoterRole) character_regex = ( r"""\{\n""" + r""" "name": "[\w\d\s]{1,16}",\n""" + r""" "house": "(Gryffindor|Slytherin|Ravenclaw|Hufflepuff)",\n""" + r""" "blood status": "(Pure-blood|Half-blood|Muggle-born)",\n""" + r""" "occupation": "(student|teacher|auror|ministry of magic|death eater|order of the phoenix)",\n""" + r""" "wand": \{\n""" + r""" "wood": "[\w\d\s]{1,16}",\n""" + r""" "core": "[\w\d\s]{1,16}",\n""" + r""" "length": [0-9]{1,2}\.[0-9]{0,2}\n""" + r""" \},\n""" + r""" "alive": "(Alive|Deceased)",\n""" + r""" "patronus": "[\w\d\s]{1,16}",\n""" + r""" "bogart": "[\w\d\s]{1,16}"\n""" + r"""\}""" ) city_regex = ( r"""\{\n""" + r""" "name": "[\w\d\s]{1,16}",\n""" + r""" "country": "[\w\d\s]{1,16}",\n""" + r""" "latitude": [-+]?[0-9]*\.?[0-9]{0,2},\n""" + r""" "population": [-+]?[0-9]{1,9},\n""" + r""" "top 3 landmarks": \["[\w\d\s]{1,16}", "[\w\d\s]{1,16}", "[\w\d\s]{1,16}"\]\n""" + r"""\}""" ) # fmt: off def character_gen(name, generate): s = name + " is a character in Harry Potter. Please fill in the following information about this character.\n" s += generate(s, max_tokens=256, regex=character_regex) return s # fmt: on # fmt: off def city_gen(document, generate): s = "Please extract the information of a city from the following wikipedia page.\n" s += "Page begin.\n" + document + "Page end.\n" s += "Here is the name, country, and symbol of the city in JSON format.\n" s += generate(s, max_tokens=256, regex=city_regex) return s # fmt: on @guidance def character_maker(lm, name): regex_str_no_quote = r"[\w\d\s]+" regex_float = r"[0-9]+\.[0-9]+" lm += f"""\ {name} is a character in Harry Potter. Please fill in the following information about this character. {{ "name": "{guidance.gen("name", max_tokens=16, regex=regex_str_no_quote)}", "house": "{guidance.select(options=['Gryffindor', 'Slytherin', 'Ravenclaw', 'Hufflepuff'], name='house')}", "blood status": "{guidance.select(options=['Pure-blood', 'Half-blood', 'Muggle-born'], name='blood status')}", "occupation": "{guidance.select(options=['student', 'teacher', 'auror', 'ministry of magic', 'death eater', 'order of the phoenix'], name='occupation')}", "wand": {{ "wood": "{guidance.gen("wood", max_tokens=16, regex=regex_str_no_quote)}", "core": "{guidance.gen('core', max_tokens=16, regex=regex_str_no_quote)}", "length": {guidance.gen('length', max_tokens=10, regex=regex_float)} }}, "alive": "{guidance.select(options=['Alive', 'Deceased'], name='alive')}", "patronus": "{guidance.gen('patronus', max_tokens=16, regex=regex_str_no_quote)}", "bogart": "{guidance.gen('bogart', max_tokens=16, regex=regex_str_no_quote)}" }} """ return lm async def call_generate_lmql( prompt, temperature, max_tokens, regex, max_len=4096, model=None, **kwargs ): assert model is not None import lmql @lmql.query(model=model) async def program(question, max_tokens, regex): '''lmql """{question}[ANSWER]""" where len(TOKENS(ANSWER)) < max_tokens and REGEX(ANSWER, regex) return ANSWER ''' return await program( question=prompt, temperature=temperature, max_tokens=max_tokens, max_len=max_len, regex=regex, **kwargs, ) @guidance def city_maker(lm, document): regex_str_no_quote = r"[\w\d\s]+" regex_float = r"[0-9]+\.[0-9]+" lm += f"""\ Please extract the information of a city from the following wikipedia page. Page begin. {document} Page end. Here is the name, country, and symbol of the city in JSON format. {{ "name": "{guidance.gen("name", max_tokens=16, regex=regex_str_no_quote)}", "country": "{guidance.gen("country", max_tokens=16, regex=regex_str_no_quote)}", "latitude": {guidance.gen("latitude", max_tokens=10, regex=regex_float)}, "population": {guidance.gen("population", max_tokens=10, regex=r"[0-9]+")}, "top 3 landmarks": [ "{guidance.gen("landmark1", max_tokens=16, regex=regex_str_no_quote)}", "{guidance.gen("landmark2", max_tokens=16, regex=regex_str_no_quote)}", "{guidance.gen("landmark3", max_tokens=16, regex=regex_str_no_quote)}" ] }} """ return lm def bench_character(args): arguments = [] with open(args.data_path, "r") as f: for line in f: arguments.append({"name": line.strip()}) arguments = arguments[: args.num_jsons] states = [None] * len(arguments) # Select backend if args.backend == "outlines": call_generate = partial(get_call_generate(args), temperature=0) def get_one_answer(i): states[i] = character_gen(**arguments[i], generate=call_generate) elif args.backend == "guidance": model = guidance.models.LlamaCpp( args.model_path, n_gpu_layers=-1, n_ctx=args.n_ctx, ) def get_one_answer(i): lm = model + character_maker(**arguments[i]) states[i] = lm elif args.backend == "lmql": import asyncio import lmql model = lmql.model(args.model_path, endpoint=f"{args.host}:{args.port}") call_generate = partial( call_generate_lmql, model=model, max_tokens=256, regex=character_regex, ) async def get_one_answer_async(i): states[i] = await call_generate(prompt=arguments[i]["name"], temperature=0) else: raise ValueError(f"Invalid backend: {args.backend}") tic = time.perf_counter() if args.backend != "lmql": if args.parallel == 1: for i in tqdm(range(len(arguments))): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: rets = list( tqdm( executor.map(get_one_answer, list(range(len(arguments)))), total=len(arguments), ) ) for _ in rets: pass else: batches = [] for i in range(0, len(arguments), args.parallel): batches.append(list(range(i, min(i + args.parallel, len(arguments))))) loop = asyncio.get_event_loop() for bt in tqdm(batches): loop.run_until_complete( asyncio.gather(*[get_one_answer_async(i) for i in bt]) ) latency = time.perf_counter() - tic return states, latency def bench_city_doc(args): arguments = [] for line in read_jsonl(args.data_path): arguments.append({"document": line["document"]}) arguments = arguments[: args.num_jsons] states = [None] * len(arguments) # Select backend if args.backend == "outlines": call_generate = partial(get_call_generate(args), temperature=0) def get_one_answer(i): states[i] = city_gen(**arguments[i], generate=call_generate) elif args.backend == "guidance": model = guidance.models.LlamaCpp( args.model_path, n_gpu_layers=-1, n_ctx=args.n_ctx, ) def get_one_answer(i): lm = model + city_maker(**arguments[i]) states[i] = lm else: raise ValueError(f"Invalid backend: {args.backend}") tic = time.perf_counter() if args.parallel == 1: for i in tqdm(range(len(arguments))): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: rets = executor.map(get_one_answer, list(range(len(arguments)))) for _ in rets: pass latency = time.perf_counter() - tic return states, latency def main(args): if args.mode == "character": args.data_path = "dataset.txt" states, latency = bench_character(args) elif args.mode == "city": args.data_path = "questions.jsonl" states, latency = bench_city_doc(args) # Compute accuracy print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}_{args.mode}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "json_jump_forward", "backend": args.backend, "latency": round(latency, 3), "num_jsons": args.num_jsons, "mode": args.mode, "parallel": args.parallel, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str) parser.add_argument("--num-jsons", type=int, default=50) parser.add_argument( "--mode", type=str, default="character", choices=["character", "city"] ) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/json_jump_forward/bench_sglang.py ================================================ import argparse import json import time import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text, read_jsonl # there are some FSM bugs with json regex converted from pydantic model # here use a string regex instead # regex_string = build_regex_from_object(HarryPoterRole) character_regex = ( r"""\{\n""" + r""" "name": "[\w\d\s]{1,16}",\n""" + r""" "house": "(Gryffindor|Slytherin|Ravenclaw|Hufflepuff)",\n""" + r""" "blood status": "(Pure-blood|Half-blood|Muggle-born)",\n""" + r""" "occupation": "(student|teacher|auror|ministry of magic|death eater|order of the phoenix)",\n""" + r""" "wand": \{\n""" + r""" "wood": "[\w\d\s]{1,16}",\n""" + r""" "core": "[\w\d\s]{1,16}",\n""" + r""" "length": [0-9]{1,2}\.[0-9]{0,2}\n""" + r""" \},\n""" + r""" "alive": "(Alive|Deceased)",\n""" + r""" "patronus": "[\w\d\s]{1,16}",\n""" + r""" "bogart": "[\w\d\s]{1,16}"\n""" + r"""\}""" ) city_regex = ( r"""\{\n""" + r""" "name": "[\w\d\s]{1,16}",\n""" + r""" "country": "[\w\d\s]{1,16}",\n""" + r""" "latitude": [-+]?[0-9]*\.?[0-9]{0,2},\n""" + r""" "population": [-+]?[0-9]{1,9},\n""" + r""" "top 3 landmarks": \["[\w\d\s]{1,16}", "[\w\d\s]{1,16}", "[\w\d\s]{1,16}"\]\n""" + r"""\}""" ) # fmt: off @sgl.function def character_gen(s, name): s += name + " is a character in Harry Potter. Please fill in the following information about this character.\n" s += sgl.gen("json_output", max_tokens=256, regex=character_regex) # fmt: on # fmt: off @sgl.function def city_gen(s, document): s += "Please extract the information of a city from the following wikipedia page.\n" s += "Page begin.\n" + document + "Page end.\n" s += "Here is the name, country, and symbol of the city in JSON format.\n" s += sgl.gen("json_output",max_tokens=256, regex=city_regex) # fmt: on def bench_city_doc(args): arguments = [] for line in read_jsonl(args.data_path): arguments.append({"document": line["document"]}) arguments = arguments[: args.num_jsons] # Select backend backend = select_sglang_backend(args) sgl.set_default_backend(backend) # Run requests tic = time.perf_counter() states = city_gen.run_batch( arguments, temperature=0, num_threads=args.parallel, progress_bar=True, ) latency = time.perf_counter() - tic return states, latency def bench_character(args): arguments = [] with open(args.data_path, "r") as f: for line in f: arguments.append({"name": line.strip()}) arguments = arguments[: args.num_jsons] # Select backend backend = select_sglang_backend(args) sgl.set_default_backend(backend) # Run requests tic = time.perf_counter() states = character_gen.run_batch( arguments, temperature=0, num_threads=args.parallel, progress_bar=True, ) latency = time.perf_counter() - tic return states, latency def main(args): if args.mode == "character": args.data_path = "dataset.txt" states, latency = bench_character(args) elif args.mode == "city": args.data_path = "questions.jsonl" states, latency = bench_city_doc(args) # Compute accuracy print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}_{args.mode}.txt", states) with open(f"{args.backend}_{args.mode}.json", "w") as fout: for state in states: fout.write(state["json_output"] + "\n") with open(args.result_file, "a") as fout: value = { "task": "json_jump_forward", "backend": args.backend, "latency": round(latency, 3), "num_jsons": args.num_jsons, "mode": args.mode, "parallel": args.parallel, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str) parser.add_argument("--num-jsons", type=int, default=50) parser.add_argument( "--mode", type=str, default="character", choices=["character", "city"] ) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/json_jump_forward/build_dataset.py ================================================ import json import transformers import wikipedia model_path = "meta-llama/Llama-2-7b-chat-hf" t = transformers.AutoTokenizer.from_pretrained(model_path) city_names = [ "los angles", "london", "tokyo", "beijing", "singapore", "paris", "dubai", "sydney", "moscow", "rome", "toronto", "rio de janeiro", "istanbul", "berlin", "auckland", "buenos aires", "mexico city", "mumbai", "seoul", "bangkok", "cairo", "athens", "jerusalem", ] def get_content(city_name): content = str(wikipedia.page(city_name).content) content = content.replace("\n\n", "\n") tokens = t.encode(content) expected_tokens = 3000 truncate_len = int((expected_tokens / len(tokens)) * len(content)) truncate_content = content[:truncate_len] truncate_tokens = t.encode(truncate_content) # Count token print( f"city_name: {city_name}, #tokens: {len(tokens)}, #truncate tokens: {len(truncate_tokens)}" ) return truncate_content if __name__ == "__main__": with open("questions.jsonl", "w") as fout: for city_name in city_names: truncate_content = get_content(city_name) fout.write(json.dumps({"document": truncate_content}) + "\n") ================================================ FILE: benchmark/json_jump_forward/dataset.txt ================================================ Harry Potter Hermione Granger Ron Weasley Albus Dumbledore Severus Snape Rubeus Hagrid Draco Malfoy Ginny Weasley Fred Weasley George Weasley Percy Weasley Sirius Black Remus Lupin Neville Longbottom Luna Lovegood Cedric Diggory Cho Chang Lord Voldemort Minerva McGonagall Filius Flitwick Dolores Umbridge Bellatrix Lestrange Lucius Malfoy Molly Weasley Arthur Weasley Nymphadora Tonks Dobby Moaning Myrtle Peter Pettigrew Alastor 'Mad-Eye' Moody Horace Slughorn Vernon Dursley Petunia Dursley Dudley Dursley Argus Filch Sybill Trelawney Gilderoy Lockhart Fleur Delacour Viktor Krum Bill Weasley Oliver Wood Cornelius Fudge Barty Crouch Sr. Barty Crouch Jr. Kingsley Shacklebolt Quirinus Quirrell Nearly Headless Nick Aunt Marge Griphook Ludo Bagman ================================================ FILE: benchmark/json_schema/README.md ================================================ ## Run benchmark ### Benchmark sglang Run Llama-8b ```bash python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --port 30000 ``` Benchmark ```bash python3 bench_sglang.py ``` ================================================ FILE: benchmark/json_schema/bench_sglang.py ================================================ import argparse import json import time from typing import List, Tuple import jsonschema from datasets import load_dataset import sglang as sgl from sglang.global_config import global_config from sglang.srt.utils.hf_transformers_utils import get_tokenizer from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text @sgl.function def schema_gen(s, message: Tuple[str, str], json_schema: str): system, user = message s += sgl.system(system) s += sgl.user(user) s += sgl.assistant( sgl.gen("json_output", temperature=0, max_tokens=256, json_schema=json_schema) ) def contains_formats(schema, formats: List[str]): if isinstance(schema, dict): if schema.get("format", None) in formats: return True for value in schema.values(): if contains_formats(value, formats): return True elif isinstance(schema, list): for item in schema: if contains_formats(item, formats): return True return False def convert_dataset(path: str): raw_dataset = load_dataset(path) dataset = [] for data in raw_dataset["train"]: messages = data["prompt"] schema = data["schema"] obj = json.loads(schema) # skip some corrupted examples if obj.get("type", None) is None: continue # skip schema with format "email" # which is not supported by outlines for now if contains_formats(obj, ["email"]): continue system = messages[0] user = messages[1] assert system["role"] == "system", "invalid role" assert user["role"] == "user", "invalid role" assert len(messages) == 2, "invalid message length" message = json.dumps(system["content"]), json.dumps(user["content"]) dataset.append( { "message": message, "json_schema": schema, } ) return dataset def bench_schema(args): arguments = convert_dataset(args.data_path) if args.num_jsons < 0 or args.num_jsons > len(arguments): args.num_jsons = len(arguments) arguments = arguments[: args.num_jsons] # Select backend backend = select_sglang_backend(args) sgl.set_default_backend(backend) # Run requests tic = time.perf_counter() states = schema_gen.run_batch( arguments, temperature=0, num_threads=args.parallel, progress_bar=True, ) latency = time.perf_counter() - tic # Check if the outputs are valid indexes = [] for i, state in enumerate(states): try: schema = json.loads(arguments[i]["json_schema"]) obj = json.loads(state["json_output"]) assert jsonschema.validate(obj, schema) is None except Exception as e: print(e) indexes.append(i) return states, latency def main(args): states, latency = bench_schema(args) # Compute accuracy tokenizer = get_tokenizer( global_config.default_backend.get_server_info()["tokenizer_path"] ) output_jsons = [state["json_output"] for state in states] num_output_tokens = sum(len(tokenizer.encode(x)) for x in output_jsons) print(f"Latency: {latency:.3f}") print(f"Output throughput: {num_output_tokens / latency:.3f} token/s") print(f"#output tokens: {num_output_tokens}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(f"{args.backend}.jsonl", "w") as fout: for state in states: fout.write(state["json_output"] + "\n") with open(args.result_file, "a") as fout: value = { "task": "json_schema", "backend": args.backend, "latency": round(latency, 3), "num_jsons": args.num_jsons, "parallel": args.parallel, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="NousResearch/json-mode-eval") parser.add_argument("--num-jsons", type=int, default=-1) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/kernels/all_reduce/benchmark_aiter.py ================================================ """ Benchmark SGLang vs Aiter custom all-reduce across message sizes. Usage: torchrun --nproc_per_node=2 benchmark_aiter.py torchrun --nproc_per_node=4 benchmark_aiter.py torchrun --nproc_per_node=8 benchmark_aiter.py """ import argparse import os import sys import time from typing import List, Optional, Tuple import torch import torch.distributed as dist def parse_args(): parser = argparse.ArgumentParser( description="Benchmark SGLang vs Aiter custom all-reduce across message sizes." ) parser.add_argument( "--backend", type=str, default="gloo", help="Process group backend for the custom-AR control path (must NOT be nccl).", ) parser.add_argument( "--warmup", type=int, default=5, help="Warmup iterations per size per implementation.", ) parser.add_argument( "--iters-small", type=int, default=50, help="Benchmark iterations for sizes <= 1MB.", ) parser.add_argument( "--iters-large", type=int, default=20, help="Benchmark iterations for sizes > 1MB.", ) parser.add_argument( "--verbose", action="store_true", help="Print per-iteration timings on rank 0 for debugging.", ) return parser.parse_args() def get_env_rank_world() -> Tuple[int, int, int]: rank = int(os.environ.get("RANK", "0")) world_size = int(os.environ.get("WORLD_SIZE", "1")) local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) return rank, world_size, local_rank def init_dist(backend: str): rank, world_size, _ = get_env_rank_world() if not dist.is_initialized(): dist.init_process_group( backend=backend, init_method="env://", rank=rank, world_size=world_size, ) def get_device(local_rank: int) -> torch.device: torch.cuda.set_device(local_rank) return torch.device(f"cuda:{local_rank}") def human_size(num_bytes: int) -> str: units = [("B", 1), ("K", 1024), ("M", 1024 * 1024), ("G", 1024 * 1024 * 1024)] for suf, base in reversed(units): if num_bytes % base == 0 and num_bytes >= base: val = num_bytes // base return f"{val}{suf}" return f"{num_bytes}B" def get_message_sizes() -> List[int]: return [ 32 * 1024, 64 * 1024, 128 * 1024, 256 * 1024, 512 * 1024, 1 * 1024 * 1024, 2 * 1024 * 1024, 4 * 1024 * 1024, 8 * 1024 * 1024, 16 * 1024 * 1024, 32 * 1024 * 1024, 64 * 1024 * 1024, ] @torch.inference_mode() def run_once(comm, inp: torch.Tensor) -> Optional[torch.Tensor]: if hasattr(comm, "all_reduce_unreg"): return comm.all_reduce_unreg(inp) if hasattr(comm, "custom_all_reduce"): return comm.custom_all_reduce(inp) raise RuntimeError("No known all-reduce method found on the communicator.") @torch.inference_mode() def bench_impl( name: str, comm, sizes: List[int], device: torch.device, warmup: int, iters_small: int, iters_large: int, verbose: bool, pg: Optional[dist.ProcessGroup] = None, ) -> List[Tuple[int, Optional[float]]]: rank = dist.get_rank() world_size = dist.get_world_size() results: List[Tuple[int, Optional[float]]] = [] for size_bytes in sizes: elems = size_bytes // 2 # float16: 2 bytes per element inp = torch.empty(elems, dtype=torch.float16, device=device) inp.uniform_(0, 1) disabled = False dist.barrier(group=pg) for _ in range(warmup): torch.cuda.synchronize() out = run_once(comm, inp) torch.cuda.synchronize() if out is None: disabled = True break dist.barrier(group=pg) if disabled: if rank == 0: print( f"[{name}] {human_size(size_bytes)}: custom AR disabled (skipped)" ) results.append((size_bytes, None)) continue num_iters = iters_small if size_bytes <= (1 * 1024 * 1024) else iters_large times_ms: List[float] = [] for it in range(num_iters): dist.barrier(group=pg) torch.cuda.synchronize() t0 = time.perf_counter() out = run_once(comm, inp) torch.cuda.synchronize() t1 = time.perf_counter() dist.barrier(group=pg) if out is None: disabled = True break dt_ms = (t1 - t0) * 1000.0 times_ms.append(dt_ms) if verbose and rank == 0: print( f"[{name}] size={human_size(size_bytes)} iter={it} time={dt_ms:.3f} ms" ) if disabled or not times_ms: if rank == 0: print( f"[{name}] {human_size(size_bytes)}: custom AR disabled (no timings)" ) results.append((size_bytes, None)) continue avg_ms_local = sum(times_ms) / len(times_ms) avg_tensor = torch.tensor([avg_ms_local], dtype=torch.float64, device=device) gather_list = [torch.zeros_like(avg_tensor) for _ in range(world_size)] dist.all_gather(gather_list, avg_tensor, group=pg) if rank == 0: avg_ms = float(torch.stack(gather_list).mean().item()) print( f"[{name}] {human_size(size_bytes)}: {avg_ms:.3f} ms (avg across ranks)" ) results.append((size_bytes, avg_ms)) else: results.append((size_bytes, None)) return results def main(): args = parse_args() rank, world_size, local_rank = get_env_rank_world() if world_size not in (2, 4, 6, 8): print( f"[rank {rank}] WARNING: world_size={world_size} not in supported set (2,4,6,8). " "Custom AR may disable itself.", file=sys.stderr, ) init_dist(args.backend) device = get_device(local_rank) # Import after dist init; some libs query torch dist state on import sgl_comm = None aiter_comm = None HAVE_SGLANG = False HAVE_AITER = False try: from sglang.srt.distributed.device_communicators.custom_all_reduce import ( CustomAllreduce as SGLCustomAllreduce, ) HAVE_SGLANG = True except Exception as e: if rank == 0: print(f"SGLang CustomAllreduce import failed: {e}", file=sys.stderr) try: from aiter.dist.device_communicators.custom_all_reduce import ( CustomAllreduce as AiterCustomAllreduce, ) HAVE_AITER = True except Exception as e: if rank == 0: print(f"Aiter CustomAllreduce import failed: {e}", file=sys.stderr) if rank == 0: print(f"Initialized PG backend={args.backend} world_size={world_size}") print(f"Device: {device.type}:{device.index}") print(f"SGLang available: {HAVE_SGLANG}, Aiter available: {HAVE_AITER}") pg = dist.group.WORLD sizes = get_message_sizes() max_size = max(sizes) if sizes else (64 * 1024 * 1024) if HAVE_SGLANG: try: sgl_comm = SGLCustomAllreduce(group=pg, device=device, max_size=max_size) except Exception as e: if rank == 0: print( f"Failed to construct SGLang CustomAllreduce: {e}", file=sys.stderr ) sgl_comm = None if HAVE_AITER: try: aiter_comm = AiterCustomAllreduce( group=pg, device=device, max_size=max_size ) except Exception as e: if rank == 0: print( f"Failed to construct Aiter CustomAllreduce: {e}", file=sys.stderr ) aiter_comm = None sgl_results: List[Tuple[int, Optional[float]]] = [] aiter_results: List[Tuple[int, Optional[float]]] = [] if sgl_comm is not None: sgl_results = bench_impl( name="SGLang", comm=sgl_comm, sizes=sizes, device=device, warmup=args.warmup, iters_small=args.iters_small, iters_large=args.iters_large, verbose=args.verbose, pg=pg, ) if aiter_comm is not None: aiter_results = bench_impl( name="Aiter", comm=aiter_comm, sizes=sizes, device=device, warmup=args.warmup, iters_small=args.iters_small, iters_large=args.iters_large, verbose=args.verbose, pg=pg, ) for comm in (sgl_comm, aiter_comm): if comm is not None and hasattr(comm, "close"): try: comm.close() except Exception: pass if dist.get_rank() == 0: print("\nResults (avg ms across ranks; None = disabled/unavailable):") header = f"{'Size':>8} {'SGLang(ms)':>12} {'Aiter(ms)':>11}" print(header) print("-" * len(header)) sgl_map = {s: v for s, v in sgl_results if v is not None} aiter_map = {s: v for s, v in aiter_results if v is not None} for s in sizes: sgl_ms = sgl_map.get(s, None) aiter_ms = aiter_map.get(s, None) print( f"{human_size(s):>8} {('%.3f' % sgl_ms) if sgl_ms is not None else 'None':>12} " f"{('%.3f' % aiter_ms) if aiter_ms is not None else 'None':>11}" ) dist.barrier() dist.destroy_process_group() if __name__ == "__main__": main() ================================================ FILE: benchmark/kernels/all_reduce/benchmark_all_reduce.py ================================================ """ Benchmark SGLang custom all-reduce vs Torch symm-mem all-reduce across message sizes. Usage: torchrun --nproc_per_node=2 benchmark_all_reduce.py torchrun --nproc_per_node=4 benchmark_all_reduce.py torchrun --nproc_per_node=8 benchmark_all_reduce.py """ import argparse import os import sys import time from typing import List, Optional, Tuple import torch import torch.distributed as dist from sglang.srt.distributed.parallel_state import ( destroy_distributed_environment, destroy_model_parallel, init_distributed_environment, initialize_model_parallel, ) def parse_args(): parser = argparse.ArgumentParser( description="Benchmark SGLang custom all-reduce vs Torch symm-mem all-reduce across message sizes." ) parser.add_argument( "--backend", type=str, default="gloo", help="Process group backend for the custom-AR control path (must NOT be nccl).", ) parser.add_argument( "--warmup", type=int, default=5, help="Warmup iterations per size per implementation.", ) parser.add_argument( "--iters-small", type=int, default=50, help="Benchmark iterations for sizes <= 1MB.", ) parser.add_argument( "--iters-large", type=int, default=20, help="Benchmark iterations for sizes > 1MB.", ) parser.add_argument( "--verbose", action="store_true", help="Print per-iteration timings on rank 0 for debugging.", ) return parser.parse_args() def get_env_rank_world() -> Tuple[int, int, int]: rank = int(os.environ.get("RANK", "0")) world_size = int(os.environ.get("WORLD_SIZE", "1")) local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) return rank, world_size, local_rank def init_dist(backend: str): rank, world_size, _ = get_env_rank_world() if not dist.is_initialized(): dist.init_process_group( backend=backend, init_method="env://", rank=rank, world_size=world_size, ) device = torch.device(f"cuda:{rank}") torch.cuda.set_device(device) distributed_init_method = f"tcp://localhost:23456" init_distributed_environment( world_size=world_size, rank=rank, distributed_init_method=distributed_init_method, local_rank=rank, ) initialize_model_parallel(tensor_model_parallel_size=world_size) return dist.group.WORLD def get_device(local_rank: int) -> torch.device: torch.cuda.set_device(local_rank) return torch.device(f"cuda:{local_rank}") def human_size(num_bytes: int) -> str: units = [("B", 1), ("K", 1024), ("M", 1024 * 1024), ("G", 1024 * 1024 * 1024)] for suf, base in reversed(units): if num_bytes % base == 0 and num_bytes >= base: val = num_bytes // base return f"{val}{suf}" return f"{num_bytes}B" def get_message_sizes() -> List[int]: return [ 32 * 1024, 64 * 1024, 128 * 1024, 256 * 1024, 512 * 1024, 1 * 1024 * 1024, 2 * 1024 * 1024, 4 * 1024 * 1024, 8 * 1024 * 1024, 16 * 1024 * 1024, 32 * 1024 * 1024, 64 * 1024 * 1024, ] @torch.inference_mode() def run_once(comm, inp: torch.Tensor) -> Optional[torch.Tensor]: if hasattr(comm, "custom_all_reduce"): return comm.custom_all_reduce(inp) if hasattr(comm, "all_reduce"): return comm.all_reduce(inp) raise RuntimeError("No known all-reduce method found on the communicator.") @torch.inference_mode() def bench_impl( name: str, comm, sizes: List[int], device: torch.device, warmup: int, iters_small: int, iters_large: int, verbose: bool, pg: Optional[dist.ProcessGroup] = None, ) -> List[Tuple[int, Optional[float]]]: rank = dist.get_rank() world_size = dist.get_world_size() results: List[Tuple[int, Optional[float]]] = [] for size_bytes in sizes: elems = size_bytes // 2 # float16: 2 bytes per element inp = torch.empty(elems, dtype=torch.float16, device=device) inp.uniform_(0, 1) disabled = False dist.barrier(group=pg) for _ in range(warmup): torch.cuda.synchronize() out = run_once(comm, inp) torch.cuda.synchronize() if out is None: disabled = True break dist.barrier(group=pg) if disabled: if rank == 0: print( f"[{name}] {human_size(size_bytes)}: custom AR disabled (skipped)" ) results.append((size_bytes, None)) continue num_iters = iters_small if size_bytes <= (1 * 1024 * 1024) else iters_large times_ms: List[float] = [] for it in range(num_iters): dist.barrier(group=pg) torch.cuda.synchronize() t0 = time.perf_counter() out = run_once(comm, inp) torch.cuda.synchronize() t1 = time.perf_counter() dist.barrier(group=pg) if out is None: disabled = True break dt_ms = (t1 - t0) * 1000.0 times_ms.append(dt_ms) if verbose and rank == 0: print( f"[{name}] size={human_size(size_bytes)} iter={it} time={dt_ms:.3f} ms" ) if disabled or not times_ms: if rank == 0: print( f"[{name}] {human_size(size_bytes)}: custom AR disabled (no timings)" ) results.append((size_bytes, None)) continue avg_ms_local = sum(times_ms) / len(times_ms) avg_tensor = torch.tensor([avg_ms_local], dtype=torch.float64, device=device) gather_list = [torch.zeros_like(avg_tensor) for _ in range(world_size)] dist.all_gather(gather_list, avg_tensor, group=pg) if rank == 0: avg_ms = float(torch.stack(gather_list).mean().item()) print( f"[{name}] {human_size(size_bytes)}: {avg_ms:.3f} ms (avg across ranks)" ) results.append((size_bytes, avg_ms)) else: results.append((size_bytes, None)) return results def main(): args = parse_args() rank, world_size, local_rank = get_env_rank_world() if world_size not in (2, 4, 6, 8): print( f"[rank {rank}] WARNING: world_size={world_size} not in supported set (2,4,6,8). " "Custom AR may disable itself.", file=sys.stderr, ) group = init_dist(args.backend) device = get_device(local_rank) # Import after dist init; some libs query torch dist state on import torch_symm_mem_comm = None HAVE_SGLANG_CUSTOM = False HAVE_TORCH_SYMM_MEM = False try: from sglang.srt.distributed.device_communicators.custom_all_reduce import ( CustomAllreduce as SGLCustomAllreduce, ) HAVE_SGLANG_CUSTOM = True except Exception as e: if rank == 0: print(f"SGLang CustomAllreduce import failed: {e}", file=sys.stderr) try: from sglang.srt.distributed.device_communicators.torch_symm_mem import ( TorchSymmMemCommunicator as TorchSymmMemAllreduce, ) HAVE_TORCH_SYMM_MEM = True except Exception as e: if rank == 0: print(f"TorchSymmMemAllreduce import failed: {e}", file=sys.stderr) if rank == 0: print(f"Initialized PG backend={args.backend} world_size={world_size}") print(f"Device: {device.type}:{device.index}") print( f"SGLang Custom available: {HAVE_SGLANG_CUSTOM}, Torch Symm-Mem available: {HAVE_TORCH_SYMM_MEM}" ) sizes = get_message_sizes() max_size = max(sizes) if sizes else (128 * 1024 * 1024) if HAVE_SGLANG_CUSTOM: try: sgl_custom_comm = SGLCustomAllreduce( group=group, device=device, max_size=max_size ) except Exception as e: if rank == 0: print( f"Failed to construct SGLangCustomAllreduce: {e}", file=sys.stderr ) sgl_custom_comm = None if HAVE_TORCH_SYMM_MEM: try: torch_symm_mem_comm = TorchSymmMemAllreduce(group=group, device=device) except Exception as e: if rank == 0: print( f"Failed to construct TorchSymmMemAllreduce: {e}", file=sys.stderr ) torch_symm_mem_comm = None sgl_custom_results: List[Tuple[int, Optional[float]]] = [] symm_mem_results: List[Tuple[int, Optional[float]]] = [] if sgl_custom_comm is not None: sgl_custom_results = bench_impl( name="SGLangCustom", comm=sgl_custom_comm, sizes=sizes, device=device, warmup=args.warmup, iters_small=args.iters_small, iters_large=args.iters_large, verbose=args.verbose, pg=group, ) if torch_symm_mem_comm is not None: symm_mem_results = bench_impl( name="TorchSymmMem", comm=torch_symm_mem_comm, sizes=sizes, device=device, warmup=args.warmup, iters_small=args.iters_small, iters_large=args.iters_large, verbose=args.verbose, pg=group, ) for comm in (sgl_custom_comm, torch_symm_mem_comm): if comm is not None and hasattr(comm, "close"): try: comm.close() except Exception: pass if dist.get_rank() == 0: print( f"\nResults (avg ms across {world_size} ranks; None = disabled/unavailable):" ) header = f"{'Size':>8} {'CustomAR(ms)':>12} {'TorchSymmMem(ms)':>11}" print(header) print("-" * len(header)) sgl_custom_map = {s: v for s, v in sgl_custom_results if v is not None} symm_mem_map = {s: v for s, v in symm_mem_results if v is not None} for s in sizes: sgl_ms = sgl_custom_map.get(s, None) symm_mem_ms = symm_mem_map.get(s, None) print( f"{human_size(s):>8} {('%.3f' % sgl_ms) if sgl_ms is not None else 'None':>12} " f"{('%.3f' % symm_mem_ms) if symm_mem_ms is not None else 'None':>11}" ) torch.distributed.barrier(group=group) destroy_model_parallel() destroy_distributed_environment() if __name__ == "__main__": main() ================================================ FILE: benchmark/kernels/all_reduce/benchmark_fused_ar_rms_amd.py ================================================ """ Benchmark fused allreduce+rmsnorm on AMD with correctness checks. This script targets the same fused op used by SGLang: `tensor_model_parallel_fused_allreduce_rmsnorm`. It reports: - eager mode latency (prefill-like) - graph mode latency (decode-like) - fused availability (whether fused path returns non-None) - correctness (fused output matches split allreduce + rmsnorm reference) Usage example: torchrun --nproc_per_node=8 \ benchmark/kernels/all_reduce/benchmark_fused_ar_rms_amd.py \ --dtype bfloat16 \ --prefill-shapes 2048x8192,8192x8192 \ --decode-shapes 1x8192,4x8192,16x8192 \ --warmup 10 --iters 30 --repeats 5 """ import argparse import csv import os import statistics from typing import Dict, List, Optional, Sequence, Tuple import torch import torch.distributed as dist import torch.nn.functional as F from sglang.srt.distributed.communication_op import ( tensor_model_parallel_all_reduce, tensor_model_parallel_fused_allreduce_rmsnorm, ) from sglang.srt.distributed.parallel_state import ( destroy_distributed_environment, destroy_model_parallel, graph_capture, init_distributed_environment, initialize_model_parallel, set_custom_all_reduce, ) Shape = Tuple[int, int] def parse_shapes(raw: str) -> List[Shape]: shapes: List[Shape] = [] for item in [x.strip() for x in raw.split(",") if x.strip()]: if "x" not in item: raise ValueError(f"Invalid shape '{item}', expected MxN format.") m_str, n_str = item.split("x", 1) m = int(m_str) n = int(n_str) if m <= 0 or n <= 0: raise ValueError(f"Invalid shape '{item}', both dims must be positive.") shapes.append((m, n)) if not shapes: raise ValueError("Empty shape list is not allowed.") return shapes def dtype_from_name(name: str) -> torch.dtype: mapping = { "float16": torch.float16, "fp16": torch.float16, "bfloat16": torch.bfloat16, "bf16": torch.bfloat16, } if name not in mapping: raise ValueError(f"Unsupported dtype: {name}") return mapping[name] def check_close( a: torch.Tensor, b: torch.Tensor, dtype: torch.dtype ) -> Tuple[bool, str]: if dtype == torch.bfloat16: rtol, atol = 2e-2, 1.25e-1 else: rtol, atol = 1e-2, 2e-2 try: torch.testing.assert_close(a, b, rtol=rtol, atol=atol) return True, "PASS" except AssertionError: max_diff = torch.max(torch.abs(a - b)).item() mean_diff = torch.mean(torch.abs(a - b)).item() return False, f"FAIL(max={max_diff:.6f},mean={mean_diff:.6f})" def _measure_us( fn, warmup: int, iters: int, repeats: int, device: torch.device, ) -> Tuple[float, Dict[str, float]]: for _ in range(warmup): fn() torch.cuda.synchronize() start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) samples_us: List[float] = [] for _ in range(max(1, repeats)): _barrier(device) torch.cuda.synchronize() start_event.record() for _ in range(iters): fn() end_event.record() end_event.synchronize() samples_us.append(start_event.elapsed_time(end_event) * 1000.0 / iters) sorted_samples = sorted(samples_us) p50 = float(statistics.median(sorted_samples)) p95 = float(sorted_samples[int((len(sorted_samples) - 1) * 0.95)]) return p50, { "p50_us": p50, "p95_us": p95, "min_us": float(sorted_samples[0]), "max_us": float(sorted_samples[-1]), } def _barrier(device: torch.device): try: dist.barrier(device_ids=[device.index]) except TypeError: dist.barrier() def _mean_across_ranks(value: float, device: torch.device) -> float: t = torch.tensor([value], dtype=torch.float64, device=device) dist.all_reduce(t, op=dist.ReduceOp.SUM) t /= dist.get_world_size() return float(t.item()) def _all_true_across_ranks(value: bool, device: torch.device) -> bool: t = torch.tensor([1 if value else 0], dtype=torch.int32, device=device) dist.all_reduce(t, op=dist.ReduceOp.MIN) return bool(int(t.item())) def _make_inputs( shape: Shape, dtype: torch.dtype, seed: int, residual_mode: str, rank: int, device: torch.device, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: m, n = shape torch.manual_seed(seed + rank * 17) x = torch.randn((m, n), dtype=torch.float32, device=device).to(dtype) if residual_mode == "self": residual = x.clone() elif residual_mode == "random": residual = torch.randn((m, n), dtype=torch.float32, device=device).to(dtype) elif residual_mode == "zero": residual = torch.zeros((m, n), dtype=dtype, device=device) else: raise ValueError(f"Unknown residual_mode: {residual_mode}") weight = torch.randn((n,), dtype=torch.float32, device=device).to(dtype) return x, residual, weight def _split_reference( x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float ) -> Tuple[torch.Tensor, torch.Tensor]: ar_out = tensor_model_parallel_all_reduce(x.clone()) residual_out = ar_out + residual out = F.rms_norm( input=residual_out, normalized_shape=(residual_out.shape[-1],), weight=weight, eps=eps, ) return out, residual_out def bench_eager( x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float, warmup: int, iters: int, repeats: int, ) -> Dict[str, object]: split_fn = lambda: _split_reference(x, residual, weight, eps) split_us, split_stats = _measure_us(split_fn, warmup, iters, repeats, x.device) fused_probe = tensor_model_parallel_fused_allreduce_rmsnorm( x.clone(), residual.clone(), weight, eps ) fused_available = fused_probe is not None fused_us: Optional[float] = None fused_stats: Optional[Dict[str, float]] = None if fused_available: fused_fn = lambda: tensor_model_parallel_fused_allreduce_rmsnorm( x, residual, weight, eps ) fused_us, fused_stats = _measure_us(fused_fn, warmup, iters, repeats, x.device) ref_out, ref_residual = _split_reference(x, residual, weight, eps) if fused_available: fused_out, fused_residual = tensor_model_parallel_fused_allreduce_rmsnorm( x.clone(), residual.clone(), weight, eps ) out_ok, out_detail = check_close(fused_out, ref_out, x.dtype) res_ok, res_detail = check_close(fused_residual, ref_residual, x.dtype) correctness_ok = out_ok and res_ok correctness_detail = f"out={out_detail}, residual={res_detail}" else: correctness_ok = True correctness_detail = "SKIP(fused_unavailable)" return { "split_us": split_us, "split_stats": split_stats, "fused_available": fused_available, "fused_us": fused_us, "fused_stats": fused_stats, "correctness_ok": correctness_ok, "correctness_detail": correctness_detail, } def bench_graph( x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float, warmup: int, iters: int, repeats: int, ) -> Dict[str, object]: split_x = x.clone() split_res = residual.clone() split_graph_out: Optional[torch.Tensor] = None with graph_capture() as gc: split_graph = torch.cuda.CUDAGraph() with torch.cuda.graph(split_graph, stream=gc.stream): split_graph_out, _ = _split_reference(split_x, split_res, weight, eps) def split_replay(): split_graph.replay() split_us, split_stats = _measure_us(split_replay, warmup, iters, repeats, x.device) fused_probe = tensor_model_parallel_fused_allreduce_rmsnorm( x.clone(), residual.clone(), weight, eps ) fused_available = fused_probe is not None fused_us: Optional[float] = None fused_stats: Optional[Dict[str, float]] = None fused_graph_out: Optional[torch.Tensor] = None fused_graph_residual: Optional[torch.Tensor] = None if fused_available: fused_x = x.clone() fused_res = residual.clone() with graph_capture() as gc: fused_graph = torch.cuda.CUDAGraph() with torch.cuda.graph(fused_graph, stream=gc.stream): fused_graph_out, fused_graph_residual = ( tensor_model_parallel_fused_allreduce_rmsnorm( fused_x, fused_res, weight, eps ) ) def fused_replay(): fused_graph.replay() fused_us, fused_stats = _measure_us( fused_replay, warmup, iters, repeats, x.device ) ref_out, ref_residual = _split_reference(x, residual, weight, eps) if ( fused_available and fused_graph_out is not None and fused_graph_residual is not None ): fused_graph.replay() torch.cuda.synchronize() out_ok, out_detail = check_close(fused_graph_out, ref_out, x.dtype) res_ok, res_detail = check_close(fused_graph_residual, ref_residual, x.dtype) correctness_ok = out_ok and res_ok correctness_detail = f"out={out_detail}, residual={res_detail}" else: correctness_ok = True correctness_detail = "SKIP(fused_unavailable)" return { "split_us": split_us, "split_stats": split_stats, "fused_available": fused_available, "fused_us": fused_us, "fused_stats": fused_stats, "correctness_ok": correctness_ok, "correctness_detail": correctness_detail, } def _shape_bytes(shape: Shape, dtype: torch.dtype) -> int: m, n = shape return m * n * torch.tensor([], dtype=dtype).element_size() def parse_args(): parser = argparse.ArgumentParser( description="Benchmark fused allreduce+rmsnorm (prefill eager + decode graph)." ) parser.add_argument( "--dtype", type=str, default="bf16", choices=["fp16", "bf16", "float16", "bfloat16"], ) parser.add_argument("--eps", type=float, default=1e-6) parser.add_argument("--seed", type=int, default=1234) parser.add_argument( "--residual-mode", type=str, default="self", choices=["self", "random", "zero"], help="Use residual=x (self) to match aiter test behavior by default.", ) parser.add_argument( "--prefill-shapes", type=str, default="2048x8192,8192x8192,16384x8192", help="Comma-separated MxN shapes for eager mode.", ) parser.add_argument( "--decode-shapes", type=str, default="1x8192,2x8192,4x8192,8x8192,16x8192", help="Comma-separated MxN shapes for graph mode.", ) parser.add_argument("--warmup", type=int, default=10) parser.add_argument("--iters", type=int, default=30) parser.add_argument("--repeats", type=int, default=5) parser.add_argument( "--mode", type=str, default="both", choices=["eager", "graph", "both"], ) parser.add_argument( "--csv-out", type=str, default=None, help="Optional output CSV path (written on rank 0 only).", ) return parser.parse_args() def main(): args = parse_args() dtype = dtype_from_name(args.dtype) rank = int(os.environ.get("RANK", "0")) world_size = int(os.environ.get("WORLD_SIZE", "1")) local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) torch.cuda.set_device(local_rank % torch.cuda.device_count()) device = torch.device(f"cuda:{local_rank % torch.cuda.device_count()}") set_custom_all_reduce(True) init_distributed_environment( world_size=world_size, rank=rank, local_rank=local_rank, distributed_init_method="env://", backend="nccl", ) initialize_model_parallel(tensor_model_parallel_size=world_size) prefill_shapes = parse_shapes(args.prefill_shapes) decode_shapes = parse_shapes(args.decode_shapes) if rank == 0: print( "Config: " f"world_size={world_size}, dtype={dtype}, residual_mode={args.residual_mode}, " f"warmup={args.warmup}, iters={args.iters}, repeats={args.repeats}" ) run_modes: Sequence[str] if args.mode == "both": run_modes = ("eager", "graph") else: run_modes = (args.mode,) csv_rows: List[Dict[str, object]] = [] for mode in run_modes: shapes = prefill_shapes if mode == "eager" else decode_shapes if rank == 0: phase_name = "prefill(eager)" if mode == "eager" else "decode(graph)" print("\n" + "=" * 120) print(f"Mode: {phase_name}") print( "| Shape | Input bytes/rank | Split p50 (us) | Fused p50 (us) | Speedup | Fused available | Correctness |" ) print( "|:------|-----------------:|---------------:|---------------:|--------:|:----------------|:------------|" ) for shape in shapes: x, residual, weight = _make_inputs( shape=shape, dtype=dtype, seed=args.seed, residual_mode=args.residual_mode, rank=rank, device=device, ) if mode == "eager": metrics = bench_eager( x=x, residual=residual, weight=weight, eps=args.eps, warmup=args.warmup, iters=args.iters, repeats=args.repeats, ) else: metrics = bench_graph( x=x, residual=residual, weight=weight, eps=args.eps, warmup=args.warmup, iters=args.iters, repeats=args.repeats, ) split_us = _mean_across_ranks(float(metrics["split_us"]), device) fused_available = _all_true_across_ranks( bool(metrics["fused_available"]), device ) correctness_ok = _all_true_across_ranks( bool(metrics["correctness_ok"]), device ) fused_us: Optional[float] = None if fused_available and metrics["fused_us"] is not None: fused_us = _mean_across_ranks(float(metrics["fused_us"]), device) if rank == 0: m, n = shape shape_str = f"{m}x{n}" bytes_per_rank = _shape_bytes(shape, dtype) if fused_us is not None and fused_us > 0: speedup = split_us / fused_us speedup_str = f"{speedup:.3f}x" fused_str = f"{fused_us:.1f}" else: speedup_str = "N/A" fused_str = "N/A" correctness_text = ( "PASS" if correctness_ok else str(metrics["correctness_detail"]) ) print( f"| {shape_str} | {bytes_per_rank} | {split_us:.1f} | {fused_str} | " f"{speedup_str} | {str(fused_available)} | {correctness_text} |" ) csv_rows.append( { "mode": mode, "shape": shape_str, "m": m, "n": n, "bytes_per_rank": bytes_per_rank, "split_p50_us": split_us, "fused_p50_us": fused_us if fused_us is not None else "", "speedup_split_over_fused": ( split_us / fused_us if fused_us is not None and fused_us > 0 else "" ), "fused_available": fused_available, "correctness_ok": correctness_ok, "correctness_detail": correctness_text, "dtype": str(dtype), "world_size": world_size, "residual_mode": args.residual_mode, "warmup": args.warmup, "iters": args.iters, "repeats": args.repeats, } ) if rank == 0 and args.csv_out: os.makedirs(os.path.dirname(args.csv_out) or ".", exist_ok=True) fieldnames = [ "mode", "shape", "m", "n", "bytes_per_rank", "split_p50_us", "fused_p50_us", "speedup_split_over_fused", "fused_available", "correctness_ok", "correctness_detail", "dtype", "world_size", "residual_mode", "warmup", "iters", "repeats", ] with open(args.csv_out, "w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() writer.writerows(csv_rows) print(f"\nSaved CSV to: {args.csv_out}") _barrier(device) destroy_model_parallel() destroy_distributed_environment() if __name__ == "__main__": main() ================================================ FILE: benchmark/kernels/all_reduce/benchmark_mscclpp.py ================================================ """For Now, MSCCL is only supported on TP16 and TP8 case export WORLD_SIZE=1 export RANK=0 export MASTER_ADDR=127.0.0.1 export MASTER_PORT=12345 torchrun --nproc_per_node gpu \ --nnodes $WORLD_SIZE \ --node_rank $RANK \ --master_addr $MASTER_ADDR \ --master_port $MASTER_PORT benchmark/kernels/all_reduce/benchmark_mscclpp.py """ import os from contextlib import nullcontext from typing import List import torch import torch.distributed as dist from torch.distributed import ProcessGroup from sglang.srt.distributed import init_distributed_environment from sglang.srt.distributed.device_communicators.pymscclpp import PyMscclppCommunicator from sglang.srt.distributed.device_communicators.pynccl import PyNcclCommunicator from sglang.srt.distributed.parallel_state import ( get_tensor_model_parallel_group, graph_capture, initialize_model_parallel, set_mscclpp_all_reduce, ) def torch_allreduce(torch_input: torch.Tensor, group: ProcessGroup) -> torch.Tensor: dist.all_reduce(torch_input, group=group) return torch_input def msccl_allreduce( msccl_input: torch.Tensor, msccl_comm: PyMscclppCommunicator ) -> torch.Tensor: return msccl_comm.all_reduce(msccl_input) def pynccl_allreduce( msccl_input: torch.Tensor, pynccl_comm: PyNcclCommunicator ) -> torch.Tensor: pynccl_comm.all_reduce(msccl_input) return msccl_input def _bench_graph_time(func, inp_randn, warmup_loop=2, graph_loop=10, test_loop=10): graph_input = inp_randn.clone() with graph_capture() as graph_capture_context: graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=graph_capture_context.stream): for _ in range(graph_loop): graph_out = func(graph_input) graph.replay() func_output = graph_out.clone() for _ in range(warmup_loop): graph.replay() torch.cuda.synchronize() start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) latencies: List[float] = [] for _ in range(test_loop): torch.cuda.synchronize() dist.barrier() start_event.record() graph.replay() end_event.record() end_event.synchronize() latencies.append(start_event.elapsed_time(end_event)) func_cost_us = sum(latencies) / len(latencies) / graph_loop * 1000 graph.reset() return func_output, func_cost_us def _bench_eager_time(func, inp_randn, warmup_loop=2, test_loop=10): eager_input = inp_randn.clone() eager_output = func(eager_input) func_output = eager_output.clone() for _ in range(warmup_loop): func(eager_input) torch.cuda.synchronize() start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) torch.cuda.synchronize() start_event.record() for _ in range(test_loop): func(eager_input) end_event.record() torch.cuda.synchronize() func_cost_us = start_event.elapsed_time(end_event) / test_loop * 1000 return func_output, func_cost_us def get_torch_prof_ctx(do_prof: bool): ctx = ( torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], record_shapes=True, with_stack=True, ) if do_prof else nullcontext() ) return ctx def human_readable_size(size, decimal_places=1): for unit in ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]: if size < 1024.0 or unit == "PiB": break size /= 1024.0 return f"{size:.{decimal_places}f} {unit}" try: from tabulate import tabulate except ImportError: print("tabulate not installed, skipping table printing") tabulate = None def print_markdown_table(data): if tabulate is not None: print(tabulate(data, headers="keys", tablefmt="github")) return headers = data[0].keys() header_row = "| " + " | ".join(headers) + " |" separator = "| " + " | ".join(["---"] * len(headers)) + " |" rows = [] for item in data: row = "| " + " | ".join(str(item[key]) for key in headers) + " |" rows.append(row) markdown_table = "\n".join([header_row, separator] + rows) print(markdown_table) if __name__ == "__main__": import logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", force=True, ) if not dist.is_initialized(): dist.init_process_group(backend="nccl") world, world_size = dist.group.WORLD, dist.get_world_size() rank = dist.get_rank() torch.cuda.set_device(rank % 8) device = torch.cuda.current_device() set_mscclpp_all_reduce(True) init_distributed_environment( world_size=world_size, rank=rank, local_rank=rank % 8, ) initialize_model_parallel(tensor_model_parallel_size=world_size) group = get_tensor_model_parallel_group().device_group cpu_group = get_tensor_model_parallel_group().cpu_group pynccl_comm = get_tensor_model_parallel_group().pynccl_comm pymscclpp_comm = get_tensor_model_parallel_group().pymscclpp_comm dist.barrier() profile = False dtype = torch.bfloat16 ctx = get_torch_prof_ctx(profile) result = [] with ctx: for i in range(10, 20): sz = 2**i if sz * dtype.itemsize > 2**20: break inp_randn = torch.randint(1, 16, (sz,), dtype=dtype, device=device) memory = torch.empty_like(inp_randn) memory_out = torch.empty_like(memory) torch_eager_output, torch_eager_time = _bench_eager_time( lambda inp: torch_allreduce(inp, group), inp_randn ) msccl_eager_output, msccl_eager_time = _bench_eager_time( lambda inp: msccl_allreduce(inp, pymscclpp_comm), inp_randn ) msccl_graph_output, msccl_graph_time = _bench_graph_time( lambda inp: msccl_allreduce(inp, pymscclpp_comm), inp_randn ) # since pynccl is inplace op, this return result is not correct if graph loop > 1 _, pynccl_graph_time = _bench_graph_time( lambda inp: pynccl_allreduce(inp, pynccl_comm), inp_randn ) torch.testing.assert_close(torch_eager_output, msccl_graph_output) torch.testing.assert_close(torch_eager_output, msccl_eager_output) result.append( { "msg_size": human_readable_size(inp_randn.nbytes), "torch eager time": torch_eager_time, "msccl eager time": msccl_eager_time, "msccl graph time": msccl_graph_time, "pynccl graph time": pynccl_graph_time, } ) if rank == 0: print(f"sz={sz}, dtype={dtype}: correctness check PASS!") if rank == 0: print_markdown_table(result) if profile: prof_dir = f"prof/msccl" os.makedirs(prof_dir, exist_ok=True) ctx.export_chrome_trace(f"{prof_dir}/trace_rank{dist.get_rank()}.json.gz") ================================================ FILE: benchmark/kernels/all_reduce/benchmark_torch_symm_mem.py ================================================ """For Now, TORCH_SYMM_MEM is only supported on following limited tp case SM90: { 2: 64 * MiB, # 64 MB 4: 64 * MiB, # 64 MB 6: 128 * MiB, # 128 MB 8: 128 * MiB, # 128 MB }, SM100: { 2: 64 * MiB, # 64 MB 4: 64 * MiB, # 64 MB 6: 128 * MiB, # 128 MB 8: 128 * MiB, # 128 MB } export WORLD_SIZE=8 export RANK=0 export MASTER_ADDR=127.0.0.1 export MASTER_PORT=12345 torchrun --nproc_per_node gpu \ --nnodes $WORLD_SIZE \ --node_rank $RANK \ --master_addr $MASTER_ADDR \ --master_port $MASTER_PORT ./benchmark/kernels/all_reduce/benchmark_torch_symm_mem.py """ import os from contextlib import nullcontext from typing import List import torch import torch.distributed as dist from torch.distributed import ProcessGroup from sglang.srt.distributed import init_distributed_environment from sglang.srt.distributed.device_communicators.pynccl import PyNcclCommunicator from sglang.srt.distributed.device_communicators.torch_symm_mem import ( TorchSymmMemCommunicator, ) from sglang.srt.distributed.parallel_state import ( get_tensor_model_parallel_group, graph_capture, initialize_model_parallel, set_torch_symm_mem_all_reduce, ) # CI environment detection IS_CI = ( os.getenv("CI", "false").lower() == "true" or os.getenv("GITHUB_ACTIONS", "false").lower() == "true" ) def torch_allreduce(torch_input: torch.Tensor, group: ProcessGroup) -> torch.Tensor: dist.all_reduce(torch_input, group=group) return torch_input def torch_symm_mem_allreduce( torch_symm_mem_input: torch.Tensor, torch_symm_mem_comm: TorchSymmMemCommunicator ) -> torch.Tensor: return torch_symm_mem_comm.all_reduce(torch_symm_mem_input) def pynccl_allreduce( pynccl_input: torch.Tensor, pynccl_comm: PyNcclCommunicator ) -> torch.Tensor: pynccl_comm.all_reduce(pynccl_input) return pynccl_input def _bench_graph_time(func, inp_randn, warmup_loop=2, graph_loop=10, test_loop=10): graph_input = inp_randn.clone() with graph_capture() as graph_capture_context: graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=graph_capture_context.stream): for _ in range(graph_loop): graph_out = func(graph_input) graph.replay() func_output = graph_out.clone() for _ in range(warmup_loop): graph.replay() torch.cuda.synchronize() start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) latencies: List[float] = [] for _ in range(test_loop): torch.cuda.synchronize() dist.barrier() start_event.record() graph.replay() end_event.record() end_event.synchronize() latencies.append(start_event.elapsed_time(end_event)) func_cost_us = sum(latencies) / len(latencies) / graph_loop * 1000 graph.reset() return func_output, func_cost_us def _bench_eager_time(func, inp_randn, warmup_loop=2, test_loop=10): eager_input = inp_randn.clone() eager_output = func(eager_input) func_output = eager_output.clone() for _ in range(warmup_loop): func(eager_input) torch.cuda.synchronize() start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) torch.cuda.synchronize() start_event.record() for _ in range(test_loop): func(eager_input) end_event.record() torch.cuda.synchronize() func_cost_us = start_event.elapsed_time(end_event) / test_loop * 1000 return func_output, func_cost_us def get_torch_prof_ctx(do_prof: bool): ctx = ( torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], record_shapes=True, with_stack=True, ) if do_prof else nullcontext() ) return ctx def human_readable_size(size, decimal_places=1): for unit in ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]: if size < 1024.0 or unit == "PiB": break size /= 1024.0 return f"{size:.{decimal_places}f} {unit}" try: from tabulate import tabulate except ImportError: print("tabulate not installed, skipping table printing") tabulate = None def print_markdown_table(data): if tabulate is not None: print(tabulate(data, headers="keys", tablefmt="github")) return headers = data[0].keys() header_row = "| " + " | ".join(headers) + " |" separator = "| " + " | ".join(["---"] * len(headers)) + " |" rows = [] for item in data: row = "| " + " | ".join(str(item[key]) for key in headers) + " |" rows.append(row) markdown_table = "\n".join([header_row, separator] + rows) print(markdown_table) if __name__ == "__main__": import logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", force=True, ) if not dist.is_initialized(): dist.init_process_group(backend="nccl") world, world_size = dist.group.WORLD, dist.get_world_size() rank = dist.get_rank() torch.cuda.set_device(rank % 8) device = torch.cuda.current_device() set_torch_symm_mem_all_reduce(True) init_distributed_environment( world_size=world_size, rank=rank, local_rank=rank % 8, ) initialize_model_parallel(tensor_model_parallel_size=world_size) group = get_tensor_model_parallel_group().device_group cpu_group = get_tensor_model_parallel_group().cpu_group pynccl_comm = get_tensor_model_parallel_group().pynccl_comm torch_symm_mem_comm = get_tensor_model_parallel_group().torch_symm_mem_comm dist.barrier() profile = False dtype = torch.bfloat16 ctx = get_torch_prof_ctx(profile) result = [] with ctx: if IS_CI: i_range = range(10, 11) else: i_range = range(10, 20) for i in i_range: sz = 2**i if sz * dtype.itemsize > 2**24: break inp_randn = torch.randint(1, 16, (sz,), dtype=dtype, device=device) memory = torch.empty_like(inp_randn) memory_out = torch.empty_like(memory) torch_eager_output, torch_eager_time = _bench_eager_time( lambda inp: torch_allreduce(inp, group), inp_randn ) symm_mem_eager_output, symm_mem_eager_time = _bench_eager_time( lambda inp: torch_symm_mem_allreduce(inp, torch_symm_mem_comm), inp_randn, ) symm_mem_graph_output, symm_mem_graph_time = _bench_graph_time( lambda inp: torch_symm_mem_allreduce(inp, torch_symm_mem_comm), inp_randn, ) # since pynccl is inplace op, this return result is not correct if graph loop > 1 _, pynccl_graph_time = _bench_graph_time( lambda inp: pynccl_allreduce(inp, pynccl_comm), inp_randn ) torch.testing.assert_close(torch_eager_output, symm_mem_graph_output) torch.testing.assert_close(torch_eager_output, symm_mem_eager_output) result.append( { "msg_size": human_readable_size(inp_randn.nbytes), "torch eager time": torch_eager_time, "symm mem eager time": symm_mem_eager_time, "symm mem graph time": symm_mem_graph_time, "pynccl graph time": pynccl_graph_time, } ) if rank == 0: print(f"sz={sz}, dtype={dtype}: correctness check PASS!") if rank == 0: print_markdown_table(result) if profile: prof_dir = f"prof/torch_symm_mem" os.makedirs(prof_dir, exist_ok=True) ctx.export_chrome_trace(f"{prof_dir}/trace_rank{dist.get_rank()}.json.gz") ================================================ FILE: benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py ================================================ import itertools import math import cudnn import torch import torch.utils.benchmark as benchmark from flashinfer import BatchDecodeWithPagedKVCacheWrapper from sglang.srt.layers.attention.flashinfer_backend import should_use_tensor_core from sglang.srt.layers.attention.triton_ops.decode_attention import decode_attention_fwd def benchmark_forward( fn, *inputs, repeats=10, amp=False, amp_dtype=torch.float16, **kwinputs, ): def amp_wrapper(*inputs, **kwinputs): with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): fn(*inputs, **kwinputs) t = benchmark.Timer( stmt="fn_amp(*inputs, **kwinputs)", globals={"fn_amp": amp_wrapper, "inputs": inputs, "kwinputs": kwinputs}, num_threads=torch.get_num_threads(), ) m = t.timeit(repeats) return t, m def time_fwd(func, *args, **kwargs): time_f = benchmark_forward(func, *args, **kwargs) return time_f[1].mean * 1e6 def decode_attention_sglang( q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, num_kv_splits, warmup=10, ): k_buffer = kv_data[0].view(-1, head_num_kv, head_dim) v_buffer = kv_data[1].view(-1, head_num_kv, head_dim) o = torch.empty_like(q) total_tokens = batch_size * kv_len req_to_token = torch.arange(0, total_tokens).to(0).int().view(batch_size, kv_len) b_req_idx = torch.arange(0, batch_size).to(0).int() b_seq_len = torch.full((batch_size,), kv_len, dtype=torch.int32, device="cuda") max_len_in_batch = kv_len sm_scale = 1.0 / (head_dim**0.5) attn_logits = torch.empty( (batch_size, head_num_q, num_kv_splits, head_dim + 1), dtype=torch.float32, device="cuda", ) for _ in range(warmup): decode_attention_fwd( q, k_buffer, v_buffer, o, req_to_token, b_req_idx, b_seq_len, attn_logits, num_kv_splits, sm_scale, ) f = time_fwd( decode_attention_fwd, q, k_buffer, v_buffer, o, req_to_token, b_req_idx, b_seq_len, attn_logits, num_kv_splits, sm_scale, ) return f, o def decode_attention_flashinfer(dtype, head_num_q, head_num_kv): workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device="cuda") use_tensor_cores = should_use_tensor_core( kv_cache_dtype=dtype, num_attention_heads=head_num_q, num_kv_heads=head_num_kv, ) flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper( workspace_buffer, "NHD", use_tensor_cores=use_tensor_cores ) class FlashinferAttention(torch.autograd.Function): @staticmethod def forward( ctx, q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype, warmup=10, ): total_tokens = batch_size * kv_len kv_indptr = torch.arange(0, batch_size + 1).to(0).int() * kv_len kv_indices = torch.arange(0, total_tokens).to(0).int() kv_last_page_len = torch.full( (batch_size,), 1, dtype=torch.int32, device="cuda" ) flashinfer_decode_wrapper.end_forward() flashinfer_decode_wrapper.begin_forward( kv_indptr, kv_indices, kv_last_page_len, head_num_q, head_num_kv, head_dim, 1, pos_encoding_mode="NONE", data_type=dtype, ) for _ in range(warmup): o = flashinfer_decode_wrapper.forward( q.contiguous().view(-1, head_num_q, head_dim), kv_data ) f = time_fwd( flashinfer_decode_wrapper.forward, q.contiguous().view(-1, head_num_q, head_dim), kv_data, ) return f, o return FlashinferAttention def convert_to_cudnn_type(torch_type): if torch_type == torch.float16: return cudnn.data_type.HALF elif torch_type == torch.bfloat16: return cudnn.data_type.BFLOAT16 elif torch_type == torch.float32: return cudnn.data_type.FLOAT elif torch_type == torch.int32: return cudnn.data_type.INT32 elif torch_type == torch.int64: return cudnn.data_type.INT64 else: raise ValueError("Unsupported tensor data type.") def decode_attention_cudnn( q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype, warmup=10 ): # Prepare data: continuous q,k,v dims_q = (batch_size, head_num_q, 1, head_dim) strides_q = (head_num_q * head_dim, head_dim, head_num_q * head_dim, 1) q_gpu = q.as_strided(dims_q, strides_q) o_gpu = ( torch.empty(batch_size * head_num_q * head_dim) .half() .cuda() .as_strided(dims_q, strides_q) ) dims_kv = (batch_size, head_num_kv, kv_len, head_dim) strides_kv = ( kv_len * head_num_kv * head_dim, head_dim, head_num_kv * head_dim, 1, ) k_gpu = kv_data[0].as_strided(dims_kv, strides_kv) v_gpu = kv_data[1].as_strided(dims_kv, strides_kv) seq_len_q_gpu = torch.full((batch_size, 1, 1, 1), 1, device="cuda") seq_len_kv_gpu = torch.full((batch_size, 1, 1, 1), kv_len, device="cuda") attn_scale = 1.0 / (head_dim**0.5) # Prepare data: paged k,v block_size = 1 blocks_per_batch = math.ceil(kv_len / block_size) # [num_blocks, head_num_kv, block_size, head_dim], num_blocks = batch_size * blocks_per_batch container_k_gpu = torch.cat(k_gpu.chunk(blocks_per_batch, dim=2), dim=0) container_v_gpu = torch.cat(v_gpu.chunk(blocks_per_batch, dim=2), dim=0) page_table_k_gpu = ( torch.linspace( 0, batch_size * blocks_per_batch - 1, batch_size * blocks_per_batch, device="cuda", dtype=torch.int32, ) .reshape(blocks_per_batch, 1, batch_size, 1) .transpose(0, 2) ) page_table_v_gpu = page_table_k_gpu.clone() graph = cudnn.pygraph( io_data_type=convert_to_cudnn_type(dtype), intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT, ) q = graph.tensor_like(q_gpu) container_k = graph.tensor_like(container_k_gpu) container_v = graph.tensor_like(container_v_gpu) page_table_k = graph.tensor_like(page_table_k_gpu) page_table_v = graph.tensor_like(page_table_v_gpu) seq_len_q = graph.tensor_like(seq_len_q_gpu) seq_len_kv = graph.tensor_like(seq_len_kv_gpu) o, _ = graph.sdpa( name="sdpa", q=q, k=container_k, # Container K: non contiguous container with K blocks v=container_v, # Container V: non contiguous container with V blocks is_inference=True, attn_scale=attn_scale, use_causal_mask=False, use_padding_mask=True, seq_len_q=seq_len_q, seq_len_kv=seq_len_kv, paged_attention_k_table=page_table_k, # Page Table K: Tensor containing offsets to the container with K blocks paged_attention_v_table=page_table_v, # Page Table V: Tensor containing offsets to the container with V blocks paged_attention_max_seq_len_kv=kv_len, # The maximum sequence length for K caches (this is optional, but recommended) ) o.set_output(True).set_dim(dims_q).set_stride(strides_q) graph.validate() graph.build_operation_graph() graph.create_execution_plans([cudnn.heur_mode.A]) graph.check_support() graph.build_plans() workspace = torch.empty( graph.get_workspace_size(), device="cuda", dtype=torch.uint8 ) variant_pack = { q: q_gpu, container_k: container_k_gpu, container_v: container_v_gpu, page_table_k: page_table_k_gpu, page_table_v: page_table_v_gpu, seq_len_q: seq_len_q_gpu, seq_len_kv: seq_len_kv_gpu, o: o_gpu, } for _ in range(warmup): graph.execute(variant_pack, workspace) f = time_fwd( graph.execute, variant_pack, workspace, ) return f, o_gpu.squeeze(dim=2) def calculate_diff(): dtype = torch.float16 batch_size = 64 kv_len = 4096 head_num_q = 64 head_num_kv = 8 head_dim = 128 q = torch.randn(batch_size, head_num_q, head_dim, dtype=dtype, device="cuda") kv_data = ( torch.randn( batch_size * kv_len, head_num_kv, head_dim, dtype=dtype, device="cuda" ), torch.randn( batch_size * kv_len, head_num_kv, head_dim, dtype=dtype, device="cuda" ), ) _, output_sglang = decode_attention_sglang( q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, num_kv_splits=8, ) attn_flashinfer = decode_attention_flashinfer(dtype, head_num_q, head_num_kv).apply _, output_flashinfer = attn_flashinfer( q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype ) _, output_cudnn = decode_attention_cudnn( q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype ) print(f"SGLang output={output_sglang}") print(f"FlashInfer output={output_flashinfer}") print(f"cuDNN output={output_cudnn}") if torch.allclose(output_sglang, output_flashinfer, atol=1e-2, rtol=1e-2): print("✅ SGLang[Triton] and FlashInfer match") else: print("❌ SGLang[Triton] and FlashInfer differ") if torch.allclose(output_sglang, output_cudnn, atol=1e-2, rtol=1e-2): print("✅ SGLang[Triton] and cuDNN match") else: print("❌ SGLang[Triton] and cuDNN differ") if __name__ == "__main__": calculate_diff() head_dim = 128 dtype = torch.float16 batch_size_range = [2**i for i in range(0, 8, 2)] kv_len_range = [2**i for i in range(6, 13, 1)] configs = list(itertools.product(batch_size_range, kv_len_range)) for head_num_q, head_num_kv in [[32, 32], [64, 8], [40, 8]]: attn_flashinfer = decode_attention_flashinfer( dtype, head_num_q, head_num_kv ).apply for batch_size, kv_len in configs: q = torch.randn( batch_size, head_num_q, head_dim, dtype=dtype, device="cuda" ) kv_data = ( torch.randn( batch_size * kv_len, head_num_kv, head_dim, dtype=dtype, device="cuda", ), torch.randn( batch_size * kv_len, head_num_kv, head_dim, dtype=dtype, device="cuda", ), ) us_cudnn, output_cudnn = decode_attention_cudnn( q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype ) us_sglang, output_sglang = decode_attention_sglang( q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, num_kv_splits=8, ) us_flashinfer, _ = attn_flashinfer( q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype ) print( head_num_q, " ", head_num_kv, " ", batch_size, " ", kv_len, " ", us_cudnn, " ", us_sglang, " ", us_flashinfer, ) ================================================ FILE: benchmark/kernels/deepep/deepep_utils.py ================================================ # ADAPTED FROM https://github.com/deepseek-ai/DeepEP/blob/main/tests/utils.py import os import sys from typing import Optional import numpy as np import torch import torch.distributed as dist def init_dist(local_rank: int, num_local_ranks: int, args): ip = args.master_addr port = args.master_port num_nodes = args.nnodes node_rank = args.node_rank assert (num_local_ranks < 8 and num_nodes == 1) or num_local_ranks == 8 dist.init_process_group( backend="nccl", init_method=f"tcp://{ip}:{port}", world_size=num_nodes * num_local_ranks, rank=node_rank * num_local_ranks + local_rank, ) torch.set_default_dtype(torch.bfloat16) torch.set_default_device("cuda") torch.cuda.set_device(local_rank) return ( dist.get_rank(), dist.get_world_size(), dist.new_group(list(range(num_local_ranks * num_nodes))), ) def calc_diff(x: torch.Tensor, y: torch.Tensor): x, y = x.double() + 1, y.double() + 1 denominator = (x * x + y * y).sum() sim = 2 * (x * y).sum() / denominator return (1 - sim).item() def per_token_cast_to_fp8(x: torch.Tensor): assert x.dim() == 2 and x.size(1) % 128 == 0 m, n = x.shape x_view = x.view(m, -1, 128) x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) return (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view( m, n ), (x_amax / 448.0).view(m, -1) def per_token_cast_back(x_fp8: torch.Tensor, x_scales: torch.Tensor): x_fp32 = x_fp8.to(torch.float32).view(x_fp8.size(0), -1, 128) x_scales = x_scales.view(x_fp8.size(0), -1, 1) return (x_fp32 * x_scales).view(x_fp8.shape).to(torch.bfloat16) def inplace_unique(x: torch.Tensor, num_slots: int): assert x.dim() == 2 mask = x < 0 x_padded = x.masked_fill(mask, num_slots) bin_count = torch.zeros((x.size(0), num_slots + 1), dtype=x.dtype, device=x.device) bin_count.scatter_add_(1, x_padded, torch.ones_like(x_padded)) bin_count = bin_count[:, :num_slots] sorted_bin_count, sorted_bin_idx = torch.sort(bin_count, dim=-1, descending=True) sorted_bin_idx.masked_fill_(sorted_bin_count == 0, -1) sorted_bin_idx = torch.sort(sorted_bin_idx, descending=True, dim=-1).values x[:, :].fill_(-1) valid_len = min(num_slots, x.size(1)) x[:, :valid_len] = sorted_bin_idx[:, :valid_len] def create_grouped_scores( scores: torch.Tensor, group_idx: torch.Tensor, num_groups: int ): num_tokens, num_experts = scores.shape scores = scores.view(num_tokens, num_groups, -1) mask = torch.zeros((num_tokens, num_groups), dtype=torch.bool, device=scores.device) mask = mask.scatter_(1, group_idx, True).unsqueeze(-1).expand_as(scores) return (scores * mask).view(num_tokens, num_experts) def bench(fn, num_warmups: int = 20, num_tests: int = 30, post_fn=None): # Flush L2 cache with 256 MB data torch.cuda.synchronize() cache = torch.empty(int(256e6 // 4), dtype=torch.int, device="cuda") # Warmup for _ in range(num_warmups): fn() # Flush L2 cache.zero_() # Testing start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] for i in range(num_tests): # Record start_events[i].record() fn() end_events[i].record() if post_fn is not None: post_fn() torch.cuda.synchronize() times = np.array( [s.elapsed_time(e) / 1e3 for s, e in zip(start_events, end_events)] )[1:] return np.average(times), np.min(times), np.max(times) class empty_suppress: def __enter__(self): return self def __exit__(self, *_): pass class suppress_stdout_stderr: def __enter__(self): self.outnull_file = open(os.devnull, "w") self.errnull_file = open(os.devnull, "w") self.old_stdout_fileno_undup = sys.stdout.fileno() self.old_stderr_fileno_undup = sys.stderr.fileno() self.old_stdout_fileno = os.dup(sys.stdout.fileno()) self.old_stderr_fileno = os.dup(sys.stderr.fileno()) self.old_stdout = sys.stdout self.old_stderr = sys.stderr os.dup2(self.outnull_file.fileno(), self.old_stdout_fileno_undup) os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup) sys.stdout = self.outnull_file sys.stderr = self.errnull_file return self def __exit__(self, *_): sys.stdout = self.old_stdout sys.stderr = self.old_stderr os.dup2(self.old_stdout_fileno, self.old_stdout_fileno_undup) os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup) os.close(self.old_stdout_fileno) os.close(self.old_stderr_fileno) self.outnull_file.close() self.errnull_file.close() def bench_kineto( fn, kernel_names, num_tests: int = 30, suppress_kineto_output: bool = False, trace_path: Optional[str] = None, barrier_comm_profiling: bool = False, ): # Profile suppress = suppress_stdout_stderr if suppress_kineto_output else empty_suppress with suppress(): schedule = torch.profiler.schedule(wait=0, warmup=1, active=1, repeat=1) with torch.profiler.profile( activities=[torch.profiler.ProfilerActivity.CUDA], schedule=schedule ) as prof: for i in range(2): # NOTES: use a large kernel and a barrier to eliminate the unbalanced CPU launch overhead if barrier_comm_profiling: lhs = torch.randn((8192, 8192), dtype=torch.float, device="cuda") rhs = torch.randn((8192, 8192), dtype=torch.float, device="cuda") lhs @ rhs dist.all_reduce(torch.ones(1, dtype=torch.float, device="cuda")) for _ in range(num_tests): fn() prof.step() # Parse the profiling table assert isinstance(kernel_names, str) or isinstance(kernel_names, tuple) is_tupled = isinstance(kernel_names, tuple) prof_lines = ( prof.key_averages() .table(sort_by="cuda_time_total", max_name_column_width=100) .split("\n") ) kernel_names = (kernel_names,) if isinstance(kernel_names, str) else kernel_names assert all([isinstance(name, str) for name in kernel_names]) for name in kernel_names: assert ( sum([name in line for line in prof_lines]) == 1 ), f"Errors of the kernel {name} in the profiling table" # Save chrome traces if trace_path is not None: prof.export_chrome_trace(trace_path) # Return average kernel times units = {"ms": 1e3, "us": 1e6} kernel_times = [] for name in kernel_names: for line in prof_lines: if name in line: time_str = line.split()[-2] for unit, scale in units.items(): if unit in time_str: kernel_times.append(float(time_str.replace(unit, "")) / scale) break break return tuple(kernel_times) if is_tupled else kernel_times[0] def hash_tensor(t: torch.Tensor): return t.view(torch.int64).sum().item() ================================================ FILE: benchmark/kernels/deepep/tuning_deepep.py ================================================ # MODIFIED FROM https://github.com/deepseek-ai/DeepEP/blob/main/tests/test_internode.py """ Example usage: python tuning_deepep.py --nnodes 4 --node-rank $MY_NODE_RANK --master-addr 1.2.3.4 Then check `deepep_tuned.json` """ import argparse import json import time from copy import deepcopy from pathlib import Path # noinspection PyUnresolvedReferences import deep_ep import torch import torch.distributed as dist from deepep_utils import ( bench, calc_diff, create_grouped_scores, init_dist, inplace_unique, per_token_cast_back, per_token_cast_to_fp8, ) def test_main( num_sms: int, local_rank: int, num_local_ranks: int, num_ranks: int, num_nodes: int, rank: int, buffer: deep_ep.Buffer, group: dist.ProcessGroup, args, ): # Settings num_tokens, hidden, num_topk_groups, num_topk, num_experts = ( args.num_tokens, args.hidden, min(num_nodes, 4), args.num_topk, (args.num_experts // num_ranks) * num_ranks, ) assert num_experts % num_ranks == 0 and num_local_ranks == 8 if local_rank == 0: print( f"[config] num_tokens={num_tokens}, hidden={hidden}, num_topk_groups={num_topk_groups}, num_topk={num_topk}", flush=True, ) # Random data x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * rank x_pure_rand = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") x_e4m3 = per_token_cast_to_fp8(x) scores = ( torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs() + 1 ) group_scores = scores.view(num_tokens, num_nodes, -1).amax(dim=-1) group_idx = torch.topk( group_scores, k=num_topk_groups, dim=-1, sorted=False ).indices masked_scores = create_grouped_scores(scores, group_idx, num_nodes) topk_idx = torch.topk(masked_scores, num_topk, dim=-1, largest=True, sorted=False)[ 1 ] topk_weights = ( torch.ones((num_tokens, num_topk), dtype=torch.float32, device="cuda") * rank ) topk_weights_pure_rand = torch.randn( (num_tokens, num_topk), dtype=torch.float32, device="cuda" ) rank_idx = topk_idx // (num_experts // num_ranks) rank_idx.masked_fill_(topk_idx == -1, -1) inplace_unique(rank_idx, num_ranks) rdma_rank_idx = rank_idx // num_local_ranks rdma_rank_idx.masked_fill_(rank_idx == -1, -1) inplace_unique(rdma_rank_idx, num_nodes) # RDMA dispatch counts rdma_idx = topk_idx // (num_experts // num_nodes) rdma_idx.masked_fill_(topk_idx == -1, -1) inplace_unique(rdma_idx, num_nodes) num_rdma_token_sent = rdma_idx.ne(-1).sum().item() # Expert meta num_tokens_per_expert = torch.zeros((num_experts,), dtype=torch.int, device="cuda") for i in range(num_experts): num_tokens_per_expert[i] = (topk_idx == i).sum() gbl_num_tokens_per_expert = num_tokens_per_expert.clone() dist.all_reduce(gbl_num_tokens_per_expert, group=group) # Rank layout meta num_tokens_per_rank = torch.empty((num_ranks,), dtype=torch.int, device="cuda") num_tokens_per_rdma_rank = torch.empty((num_nodes,), dtype=torch.int, device="cuda") token_idx_in_rank = torch.full( (num_ranks, num_tokens), -1, dtype=torch.long, device="cuda" ) for i in range(num_ranks): num_tokens_per_rank[i] = (rank_idx == i).sum() token_sel = (rank_idx == i).max(dim=-1)[0] count = token_sel.sum().item() tokens = torch.sort(token_sel.to(torch.int), descending=True)[1] tokens[:count] = torch.sort(tokens[:count])[0] token_idx_in_rank[i][tokens[:count]] = torch.arange( count, dtype=torch.long, device="cuda" ) for i in range(num_nodes): num_tokens_per_rdma_rank[i] = (rdma_rank_idx == i).sum() token_idx_in_rank = token_idx_in_rank.T.contiguous().to(torch.int) is_token_in_rank = token_idx_in_rank >= 0 gbl_num_tokens_per_rank = num_tokens_per_rank.clone() dist.all_reduce(gbl_num_tokens_per_rank, group=group) ( ref_num_tokens_per_rank, ref_num_tokens_per_rdma_rank, ref_num_tokens_per_expert, ref_is_token_in_rank, _, ) = buffer.get_dispatch_layout(topk_idx, num_experts) assert torch.allclose(ref_num_tokens_per_rank, num_tokens_per_rank) assert torch.allclose(ref_num_tokens_per_rdma_rank, num_tokens_per_rdma_rank) assert torch.allclose(ref_num_tokens_per_expert, num_tokens_per_expert) assert torch.allclose(ref_is_token_in_rank, is_token_in_rank) t = bench(lambda: buffer.get_dispatch_layout(topk_idx, num_experts))[0] if local_rank == 0: print(f"[layout] Kernel performance: {t * 1000:.3f} ms", flush=True) print("", flush=True) group.barrier() time.sleep(1) # Config rdma_buffer_size, nvl_buffer_size = 128, (720 if num_ranks in (144, 160) else 512) config = deep_ep.Config(num_sms, 8, nvl_buffer_size, 16, rdma_buffer_size) # Test dispatch # noinspection PyShadowingNames def check_data(check_x, recv_gbl_rank_prefix_sum): assert torch.allclose(check_x.amin(dim=1), check_x.amax(dim=1)) check_start = 0 for i in range(num_ranks): check_end = recv_gbl_rank_prefix_sum[i].item() assert (check_x[check_start:check_end, :].int() - i).sum().item() == 0 check_start = check_end for previous_mode in (False, True): for async_mode in (False, True): for current_x in (x_pure_rand, x, x_e4m3): for with_topk in (False, True): if local_rank == 0: print( f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...', flush=True, end="", ) dispatch_args = { "x": current_x, "num_tokens_per_rank": num_tokens_per_rank, "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, "is_token_in_rank": is_token_in_rank, "num_tokens_per_expert": num_tokens_per_expert, "config": config, "async_finish": async_mode, } if with_topk: dispatch_args.update( { "topk_idx": topk_idx, "topk_weights": ( topk_weights_pure_rand if current_x is x_pure_rand else topk_weights ), } ) if previous_mode: dispatch_args.update({"previous_event": buffer.capture()}) ( recv_x, recv_topk_idx, recv_topk_weights, recv_num_tokens_per_expert_list, handle, event, ) = buffer.dispatch(**dispatch_args) event.current_stream_wait() if async_mode else () recv_x = ( per_token_cast_back(*recv_x) if isinstance(recv_x, tuple) else recv_x ) # Checks recv_gbl_rank_prefix_sum = handle[-4] assert gbl_num_tokens_per_rank[rank].item() == recv_x.size( 0 ), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}" assert ( gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist() == recv_num_tokens_per_expert_list ) if current_x is not x_pure_rand: check_data(recv_x, recv_gbl_rank_prefix_sum) if with_topk: # Check `topk_idx` assert ( recv_topk_idx.eq(-1) | ( (recv_topk_idx >= 0) & (recv_topk_idx < (num_experts // num_ranks)) ) ).sum().item() == recv_topk_idx.numel() for i, count in enumerate(recv_num_tokens_per_expert_list): assert recv_topk_idx.eq(i).sum().item() == count # Check `topk_weights` if current_x is not x_pure_rand: recv_topk_weights[recv_topk_idx.eq(-1)] = ( recv_topk_weights.amax(dim=1, keepdim=True).expand_as( recv_topk_weights )[recv_topk_idx.eq(-1)] ) check_data(recv_topk_weights, recv_gbl_rank_prefix_sum) # Test cached dispatch (must without top-k staffs) if not with_topk: dispatch_args = { "x": current_x, "handle": handle, "config": config, "async_finish": async_mode, } if previous_mode: dispatch_args.update({"previous_event": buffer.capture()}) recv_x, _, _, _, _, event = buffer.dispatch(**dispatch_args) event.current_stream_wait() if async_mode else () recv_x = ( per_token_cast_back(*recv_x) if isinstance(recv_x, tuple) else recv_x ) if current_x is not x_pure_rand: check_data(recv_x, recv_gbl_rank_prefix_sum) # Test combine combine_args = { "x": recv_x, "handle": handle, "config": config, "async_finish": async_mode, } if with_topk: combine_args.update({"topk_weights": recv_topk_weights}) if previous_mode: combine_args.update({"previous_event": buffer.capture()}) combined_x, combined_topk_weights, event = buffer.combine( **combine_args ) event.current_stream_wait() if async_mode else () check_x = combined_x.float() / is_token_in_rank.sum( dim=1 ).unsqueeze(1) ref_x = x_pure_rand if current_x is x_pure_rand else x assert calc_diff(check_x, ref_x) < 5e-6 if with_topk: check_topk_weights = ( combined_topk_weights if (current_x is x_pure_rand) else ( combined_topk_weights / is_token_in_rank.sum(dim=1).unsqueeze(1) ) ) ref_topk_weights = ( topk_weights_pure_rand if current_x is x_pure_rand else topk_weights ) assert calc_diff(check_topk_weights, ref_topk_weights) < 1e-9 # For later tuning dispatch_bf16_rdma_send_bytes = num_rdma_token_sent * hidden * 2 dispatch_bf16_nvl_recv_bytes = recv_x.numel() * 2 combine_bf16_nvl_send_bytes = dispatch_bf16_nvl_recv_bytes combine_bf16_rdma_recv_bytes = dispatch_bf16_rdma_send_bytes if local_rank == 0: print(" passed", flush=True) if local_rank == 0: print("", flush=True) output_data = {} # Tune dispatch performance best_dispatch_results = None fp8_factor = (1 + 4 / 128) / 2 for current_x in (x_e4m3, x): best_time, best_results = 1e10, None rdma_send_bytes = ( (dispatch_bf16_rdma_send_bytes * fp8_factor) if isinstance(current_x, tuple) else dispatch_bf16_rdma_send_bytes ) nvl_recv_bytes = ( (dispatch_bf16_nvl_recv_bytes * fp8_factor) if isinstance(current_x, tuple) else dispatch_bf16_nvl_recv_bytes ) for nvl_chunk_size in range(4, 33, 4): for rdma_chunk_size in range(4, 33, 4): config_kwargs = { "num_sms": num_sms, "num_max_nvl_chunked_send_tokens": nvl_chunk_size, "num_max_nvl_chunked_recv_tokens": nvl_buffer_size, "num_max_rdma_chunked_send_tokens": rdma_chunk_size, "num_max_rdma_chunked_recv_tokens": rdma_buffer_size, } config = deep_ep.Config(**config_kwargs) tune_args = {"x": current_x, "handle": handle, "config": config} t = bench(lambda: buffer.dispatch(**tune_args))[0] if t < best_time: best_time, best_results = t, ( num_sms, nvl_chunk_size, rdma_chunk_size, config_kwargs, ) if local_rank == 0: print( f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}, RDMA chunk {rdma_chunk_size}: {rdma_send_bytes / 1e9 / t:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / t:.2f} GB/s (NVL) ", flush=True, ) if local_rank == 0: print( f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)', flush=True, ) print("", flush=True) is_fp8 = isinstance(current_x, tuple) if is_fp8: output_data["normal_dispatch"] = deepcopy(best_results[3]) if isinstance(current_x, tuple): # Gather FP8 the best config from rank 0 best_dispatch_results = torch.tensor( [best_results[0], best_results[1], best_results[2]], dtype=torch.int32, device="cuda", ) all_best_fp8_results_list = [ torch.zeros_like(best_dispatch_results) for _ in range(torch.distributed.get_world_size()) ] dist.all_gather( all_best_fp8_results_list, best_dispatch_results, group=group ) best_dispatch_results = all_best_fp8_results_list[0].tolist() dispatch_config = deep_ep.Config( best_dispatch_results[0], best_dispatch_results[1], nvl_buffer_size, best_dispatch_results[2], rdma_buffer_size, ) dispatch_args = { "x": x, "num_tokens_per_rank": num_tokens_per_rank, "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, "is_token_in_rank": is_token_in_rank, "num_tokens_per_expert": num_tokens_per_expert, "config": dispatch_config if dispatch_config is not None else config, } recv_x, _, _, _, handle, _ = buffer.dispatch(**dispatch_args) # Tune combine performance best_time, best_results = 1e10, None for nvl_chunk_size in range(1, 8, 1): for rdma_chunk_size in range(12 if num_nodes == 2 else 8, 33, 4): config_kwargs = { "num_sms": num_sms, "num_max_nvl_chunked_send_tokens": nvl_chunk_size, "num_max_nvl_chunked_recv_tokens": nvl_buffer_size, "num_max_rdma_chunked_send_tokens": rdma_chunk_size, "num_max_rdma_chunked_recv_tokens": rdma_buffer_size, } config = deep_ep.Config(**config_kwargs) tune_args = {"x": recv_x, "handle": handle, "config": config} t = bench(lambda: buffer.combine(**tune_args))[0] if local_rank == 0: print( f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}, RDMA chunk {rdma_chunk_size}: {combine_bf16_rdma_recv_bytes / 1e9 / t:.2f} GB/s (RDMA), {combine_bf16_nvl_send_bytes / 1e9 / t:.2f} GB/s (NVL) ", flush=True, ) if t < best_time: best_time, best_results = t, ( num_sms, nvl_chunk_size, rdma_chunk_size, config_kwargs, ) if local_rank == 0: print( f"[tuning] Best combine: SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {combine_bf16_rdma_recv_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {combine_bf16_nvl_send_bytes / 1e9 / best_time:.2f} GB/s (NVL)", flush=True, ) print("", flush=True) output_data["normal_combine"] = deepcopy(best_results[3]) if rank == 0 and local_rank == 0: _write_output(args, output_data) def _write_output(args, output_data): text = json.dumps(output_data, indent=4) output_path = args.output_path print(f"Write to {output_path} with {text}") Path(output_path).write_text(text) # noinspection PyUnboundLocalVariable def test_loop(local_rank: int, num_local_ranks: int, args): num_nodes = args.nnodes rank, num_ranks, group = init_dist(local_rank, num_local_ranks, args) num_sms = args.num_sms num_qps_per_rank = num_sms // 2 buffer = deep_ep.Buffer( group, int(1e9), int(1e9), low_latency_mode=False, num_qps_per_rank=num_qps_per_rank, ) assert num_local_ranks == 8 and num_ranks > 8 torch.manual_seed(rank) for i in (num_sms,): test_main( i, local_rank, num_local_ranks, num_ranks, num_nodes, rank, buffer, group, args, ) if local_rank == 0: print("", flush=True) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--num-sms", type=int, default=24) parser.add_argument("--num-tokens", type=int, default=4096) parser.add_argument("--hidden", type=int, default=7168) parser.add_argument("--num-topk", type=int, default=8) parser.add_argument("--num-experts", type=int, default=256) parser.add_argument("--output-path", type=str, default="deepep_tuned.json") parser.add_argument("--nnodes", type=int, default=1) parser.add_argument("--node-rank", type=int, default=0) parser.add_argument("--master-addr", type=str, default="127.0.0.1") parser.add_argument("--master-port", type=int, default=8361) args = parser.parse_args() print(f"Start system with {args=}") num_processes = 8 torch.multiprocessing.spawn( test_loop, args=(num_processes, args), nprocs=num_processes ) ================================================ FILE: benchmark/kernels/deepseek/README.md ================================================ ## DeepSeek kernels benchmark ### Prerequisites - You should install [DeepGemm](https://github.com/deepseek-ai/DeepGEMM) from source before run `benchmark_deepgemm_fp8_gemm.py` and `benchmark_deepgemm_fp8_group_gemm.py`. ### Benchmark - `benchmark_deepgemm_fp8_gemm.py` ```bash python benchmark_deepgemm_fp8_gemm.py --run_correctness --tp_size 1 ``` - `benchmark_deepgemm_fp8_group_gemm.py` ```bash python benchmark_deepgemm_fp8_group_gemm.py --run_correctness --tp_size 1 ``` - You can use the `--run_correctness` parameter to verify all kernels results's correctness. - You can use the `--tp_size` parameter to benchmark all FP8 w8a8 block-wise matrix multiplications involved in DeepSeek V3/R1 under the current tensor parallelism (TP) setting. This benchmark compares DeepSeek's open-source [DeepGemm](https://github.com/deepseek-ai/DeepGEMM) implementation with SGLang's and VLLM Triton implementation. ================================================ FILE: benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py ================================================ from typing import Tuple import deep_gemm import tilelang import tilelang.language as T import torch import triton from deep_gemm import ceil_div from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor from vllm.model_executor.layers.quantization.utils.fp8_utils import ( w8a8_block_fp8_matmul as vllm_w8a8_block_fp8_matmul, ) from sglang.benchmark.bench_utils import run_bench from sglang.srt.layers.quantization.fp8_kernel import ( w8a8_block_fp8_matmul_deepgemm as w8a8_block_fp8_matmul, ) # Adapted from https://github.com/tile-ai/tilelang/blob/a8cfdce92795cb861c9033573534653ee040b5ed/examples/deepseek_deepgemm/example_deepgemm_fp8_2xAcc.py#L1 def tl_gemm( M, N, K, in_dtype, out_dtype, accum_dtype, ): assert in_dtype in [ "e4m3_float8", ], "Currently only e4m3_float8 is supported" assert out_dtype in [ "bfloat16", "float16", ], "Currently only bfloat16 and float16 are supported" TILE_SIZE = (128, 128, 128) block_M = TILE_SIZE[0] block_N = TILE_SIZE[1] block_K = TILE_SIZE[2] A_shape = (M, K) Scales_A_shape = (M, T.ceildiv(K, block_K)) B_shape = (N, K) Scales_B_shape = (T.ceildiv(N, block_N), T.ceildiv(K, block_K)) A_shared_shape = (block_M, block_K) B_shared_shape = (block_N, block_K) C_shared_shape = (block_M, block_N) @T.prim_func def main( A: T.Buffer(A_shape, in_dtype), scales_a: T.Buffer(Scales_A_shape, "float32"), B: T.Buffer(B_shape, in_dtype), scales_b: T.Buffer(Scales_B_shape, "float32"), C: T.Buffer((M, N), out_dtype), ): with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as ( bx, by, ): A_shared = T.alloc_shared(A_shared_shape, in_dtype) B_shared = T.alloc_shared(B_shared_shape, in_dtype) C_shared = T.alloc_shared(C_shared_shape, out_dtype) Scale_C_shared = T.alloc_shared((block_M), "float32") C_local = T.alloc_fragment(C_shared_shape, accum_dtype) C_local_accum = T.alloc_fragment(C_shared_shape, accum_dtype) # Improve L2 Cache T.use_swizzle(panel_size=10) T.clear(C_local) T.clear(C_local_accum) K_iters = T.ceildiv(K, block_K) for k in T.Pipelined(K_iters, num_stages=4): # Load A into shared memory T.copy(A[by * block_M, k * block_K], A_shared) # Load B into shared memory T.copy(B[bx * block_N, k * block_K], B_shared) # Load scale into shared memory Scale_B = scales_b[bx, k] for i in T.Parallel(block_M): Scale_C_shared[i] = scales_a[by * block_M + i, k] * Scale_B T.gemm(A_shared, B_shared, C_local, transpose_B=True) # Promote to enable 2xAcc for i, j in T.Parallel(block_M, block_N): C_local_accum[i, j] += C_local[i, j] * Scale_C_shared[i] T.clear(C_local) # TMA store T.copy(C_local_accum, C_shared) T.copy(C_shared, C[by * block_M, bx * block_N]) return main def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: assert x.dim() == 2 and x.size(1) % 128 == 0 m, n = x.shape x_view = x.view(m, -1, 128) x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) return (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view( m, n ), (x_amax / 448.0).view(m, -1) def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: assert x.dim() == 2 m, n = x.shape x_padded = torch.zeros( (ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device ) x_padded[:m, :n] = x x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128) x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4) x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn) return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view( x_view.size(0), x_view.size(2) ) def fp8_gemm_deepgemm( x_fp8: torch.Tensor, x_scale: torch.Tensor, y_fp8: torch.Tensor, y_scale: torch.Tensor, m: int, n: int, k: int, ): """DeepGEMM implementation of FP8 GEMM""" out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) # Run DeepGEMM kernel deep_gemm.fp8_gemm_nt((x_fp8, x_scale), (y_fp8, y_scale), out) return out def fp8_gemm_sglang( x_fp8: torch.Tensor, x_scale: torch.Tensor, y_fp8: torch.Tensor, y_scale: torch.Tensor, m: int, n: int, k: int, ): """SGLang implementation of FP8 GEMM""" block_size = [128, 128] # Matches the block size in per_block_cast_to_fp8 # Run SGLang kernel out = w8a8_block_fp8_matmul( x_fp8, y_fp8, x_scale, y_scale, block_size, torch.bfloat16 ) return out def fp8_gemm_vllm( x_fp8: torch.Tensor, x_scale: torch.Tensor, y_fp8: torch.Tensor, y_scale: torch.Tensor, m: int, n: int, k: int, ): """vLLM implementation of FP8 GEMM""" block_size = [128, 128] # Matches the block size in per_block_cast_to_fp8 # Run vLLM kernel out = vllm_w8a8_block_fp8_matmul( x_fp8, y_fp8, x_scale, y_scale, block_size, torch.bfloat16 ) return out def calculate_diff(m: int, n: int, k: int): x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) x_fp8, x_scale = per_token_cast_to_fp8(x.clone()) y_fp8, y_scale = per_block_cast_to_fp8(y.clone()) x_scale_col_major = get_mn_major_tma_aligned_tensor(x_scale.clone()) out_deepgemm = fp8_gemm_deepgemm( x_fp8.clone(), x_scale_col_major.clone(), y_fp8.clone(), y_scale.clone(), m, n, k, ) out_sglang = fp8_gemm_sglang( x_fp8.clone(), x_scale.clone(), y_fp8.clone(), y_scale.clone(), m, n, k ) tilelang_func = tl_gemm(m, n, k, "e4m3_float8", "bfloat16", "float32") tilelang_kernel = tilelang.compile(tilelang_func, out_idx=[-1]) out_tilelang = tilelang_kernel( x_fp8.clone(), x_scale.clone(), y_fp8.clone(), y_scale.clone() ) diff_sglang_deepgemm = torch.abs(out_deepgemm - out_sglang).mean().item() diff_tilelang_deepgemm = torch.abs(out_deepgemm - out_tilelang).mean().item() diff_tilelang_sglang = torch.abs(out_tilelang - out_sglang).mean().item() print(f"Shape m={m}, n={n}, k={k}:") print(f"DeepGEMM output: {out_deepgemm[0, 0:5]}") print(f"SGLang output: {out_sglang[0, 0:5]}") print(f"TileLang output: {out_tilelang[0, 0:5]}") print(f"Mean absolute difference (SGLang-DeepGEMM): {diff_sglang_deepgemm}") print(f"Mean absolute difference (TileLang-DeepGEMM): {diff_tilelang_deepgemm}") print(f"Mean absolute difference (TileLang-SGLang): {diff_tilelang_sglang}") sglang_deepgemm_match = torch.allclose( out_deepgemm, out_sglang, atol=1e-2, rtol=1e-2 ) tilelang_deepgemm_match = torch.allclose( out_deepgemm, out_tilelang, atol=1e-2, rtol=1e-2 ) tilelang_sglang_match = torch.allclose( out_tilelang, out_sglang, atol=1e-2, rtol=1e-2 ) if sglang_deepgemm_match and tilelang_deepgemm_match and tilelang_sglang_match: print("✅ All implementations match\n") else: print("❌ Some implementations differ:") print(f" - SGLang vs DeepGEMM: {'✅' if sglang_deepgemm_match else '❌'}") print(f" - TileLang vs DeepGEMM: {'✅' if tilelang_deepgemm_match else '❌'}") print(f" - TileLang vs SGLang: {'✅' if tilelang_sglang_match else '❌'}\n") def get_weight_shapes(tp_size): # cannot TP total = [ (512 + 64, 7168), ((128 + 64) * 128, 7168), (128 * (128 + 128), 512), (7168, 16384), (7168, 18432), ] # N can TP n_tp = [ (18432 * 2, 7168), ((128 + 64) * 128, 7168), (128 * (128 + 128), 512), (24576, 1536), (4096, 7168), ] # K can TP k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] weight_shapes = [] for t in total: weight_shapes.append(t) for n_t in n_tp: new_t = (n_t[0] // tp_size, n_t[1]) weight_shapes.append(new_t) for k_t in k_tp: new_t = (k_t[0], k_t[1] // tp_size) weight_shapes.append(new_t) return weight_shapes def create_benchmark_configs(tp_size): configs = [] weight_shapes = get_weight_shapes(tp_size) batch_sizes = [8, 16, 32, 64, 128, 256, 1024, 2048, 4096] for n, k in weight_shapes: for m in batch_sizes: configs.append((m, n, k, tp_size)) return configs def get_benchmark(tp_size): all_configs = create_benchmark_configs(tp_size) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["m", "n", "k", "tp_size"], x_vals=[list(config) for config in all_configs], line_arg="provider", line_vals=["deepgemm", "sglang", "tilelang"], line_names=["DeepGEMM", "SGLang", "TileLang"], styles=[("blue", "-"), ("red", "-"), ("green", "-")], ylabel="ms", plot_name=f"fp8-gemm-performance-comparison-tp{tp_size}", args={}, ) ) def benchmark(m, n, k, tp_size, provider): print(f"Shape (m={m}, n={n}, k={k}, tp={tp_size}), Provider: {provider}") x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) # Preprocess data before benchmarking x_fp8, x_scale = per_token_cast_to_fp8(x) y_fp8, y_scale = per_block_cast_to_fp8(y) x_scale_col_major = get_mn_major_tma_aligned_tensor(x_scale.clone()) quantiles = (0.5, 0.2, 0.8) if provider == "deepgemm": ms, min_ms, max_ms = run_bench( lambda: fp8_gemm_deepgemm( x_fp8.clone(), x_scale_col_major.clone(), y_fp8.clone(), y_scale.clone(), m, n, k, ), quantiles=quantiles, ) elif provider == "sglang": ms, min_ms, max_ms = run_bench( lambda: fp8_gemm_sglang( x_fp8.clone(), x_scale.clone(), y_fp8.clone(), y_scale.clone(), m, n, k, ), quantiles=quantiles, ) else: # tilelang tilelang_func = tl_gemm(m, n, k, "e4m3_float8", "bfloat16", "float32") tilelang_kernel = tilelang.compile(tilelang_func, out_idx=[-1]) ms, min_ms, max_ms = run_bench( lambda: tilelang_kernel( x_fp8.clone(), x_scale.clone(), y_fp8.clone(), y_scale.clone(), ), quantiles=quantiles, ) # Calculate TFLOPS flops = 2 * m * n * k # multiply-adds tflops = flops / (ms * 1e-3) / 1e12 # Print shape-specific results with TFLOPS print(f"Time: {ms*1000:.2f} ms, TFLOPS: {tflops:.2f}") return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms return benchmark if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() parser.add_argument( "--save_path", type=str, default="./configs/benchmark_ops/fp8_gemm/", help="Path to save fp8 gemm benchmark results", ) parser.add_argument( "--run_correctness", action="store_true", default=True, help="Whether to run correctness test", ) parser.add_argument( "--tp_size", type=int, default=1, help="Tensor parallelism size to benchmark (default: 1)", ) args = parser.parse_args() # Set random seed for reproducibility torch.manual_seed(0) torch.cuda.manual_seed(0) # Enable TF32, adapted from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_core.py#L148 torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True # Run correctness tests on a few examples if args.run_correctness: print("Running correctness tests...") calculate_diff(64, 512, 7168) # Small test calculate_diff(64, 7168, 16384) # Medium test calculate_diff(64, 18432, 7168) # Large test # Get the benchmark function with the specified tp_size benchmark = get_benchmark(args.tp_size) print(f"Running performance benchmark for TP size = {args.tp_size}...") benchmark.run(print_data=True, save_path=args.save_path) ================================================ FILE: benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py ================================================ import argparse from typing import Tuple import torch import triton from deep_gemm import ceil_div from flashinfer.gemm import gemm_fp8_nt_groupwise from sglang.benchmark.bench_utils import run_bench from sglang.srt.layers.quantization.fp8_kernel import ( sglang_per_token_group_quant_fp8, w8a8_block_fp8_matmul_deepgemm, ) from sglang.srt.layers.quantization.fp8_utils import requant_weight_ue8m0 BLOCK_SIZE = 128 def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: assert x.dim() == 2 assert BLOCK_SIZE == 128 m, n = x.shape x_padded = torch.zeros( (ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device ) x_padded[:m, :n] = x x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128) x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4) x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn) return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view( x_view.size(0), x_view.size(2) ) def get_weight_shapes(tp_size): # cannot TP total = [ (512 + 64, 7168), ((128 + 64) * 128, 7168), (128 * (128 + 128), 512), (7168, 16384), (7168, 18432), ] # N can TP n_tp = [ (18432 * 2, 7168), ((128 + 64) * 128, 7168), (128 * (128 + 128), 512), (24576, 1536), (4096, 7168), ] # K can TP k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] weight_shapes = [] for t in total: weight_shapes.append(t) for n_t in n_tp: new_t = (n_t[0] // tp_size, n_t[1]) weight_shapes.append(new_t) for k_t in k_tp: new_t = (k_t[0], k_t[1] // tp_size) weight_shapes.append(new_t) return weight_shapes def create_benchmark_configs(tp_size): configs = [] weight_shapes = get_weight_shapes(tp_size) batch_sizes = [8, 16, 32, 64, 128, 256, 1024, 2048, 4096] for n, k in weight_shapes: for m in batch_sizes: configs.append((m, n, k, tp_size)) return configs def fp8_gemm_flashinfer( x_fp8: torch.Tensor, x_scale: torch.Tensor, y_fp8: torch.Tensor, y_scale: torch.Tensor, ): """Flashinfer implementation of FP8 GEMM""" output = gemm_fp8_nt_groupwise( x_fp8, y_fp8, x_scale, y_scale, out_dtype=torch.bfloat16, backend="trtllm", ) return output def fp8_gemm_deepgemm_blackwell( x_fp8: torch.Tensor, x_scale: torch.Tensor, y_fp8: torch.Tensor, y_scale: torch.Tensor, ): """DeepGEMM implementation of FP8 GEMM""" block_size = [BLOCK_SIZE, BLOCK_SIZE] output = w8a8_block_fp8_matmul_deepgemm( x_fp8, y_fp8, x_scale, y_scale, block_size, output_dtype=torch.bfloat16 ) return output def check_accuracy(a, b, atol, rtol, percent): """Unified accuracy checking function with detailed error reporting.""" if not torch.isfinite(a).all(): print("Non-finite values in reference output") return False if not torch.isfinite(b).all(): print("Non-finite values in actual output") return False assert a.shape == b.shape, f"Shape mismatch: {a.shape} vs {b.shape}" close = torch.isclose(a, b, atol=atol, rtol=rtol) match_ratio = close.float().mean() if match_ratio >= percent: return True mismatch_percent = 1.0 - match_ratio.item() if mismatch_percent > 1 - percent: print( f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} " f"(threshold: {1 - percent:.4f})" ) return False def calculate_diff(m: int, n: int, k: int): x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) y_fp8, y_scale = per_block_cast_to_fp8(y) x_fp8, x_scale = sglang_per_token_group_quant_fp8( x, BLOCK_SIZE, column_major_scales=True ) out_flashinfer = fp8_gemm_flashinfer( x_fp8, x_scale, y_fp8, y_scale, ) dg_x_fp8, dg_x_scale = sglang_per_token_group_quant_fp8( x, BLOCK_SIZE, column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=True, ) # We can directly quantize y here, but to mimic the behavior of the actual # implementations, we requant it here. dg_y_fp8, dg_y_scale = requant_weight_ue8m0( y_fp8, y_scale, [BLOCK_SIZE, BLOCK_SIZE] ) out_deepgemm = fp8_gemm_deepgemm_blackwell( dg_x_fp8, dg_x_scale, dg_y_fp8, dg_y_scale ) print(f"Shape m={m}, n={n}, k={k}:") print(f"Flashinfer output: {out_flashinfer[0, 0:5]}") print(f"DeepGEMM output: {out_deepgemm[0, 0:5]}") flashinfer_deepgemm_match = check_accuracy( out_flashinfer, out_deepgemm, 0.1, 0.6, 0.95 ) print("Correctness check:") print(f" - Flashinfer vs DeepGEMM: {'✅' if flashinfer_deepgemm_match else '❌'}") def _benchmark(m, n, k, tp_size, provider): print(f"Shape (m={m}, n={n}, k={k}, tp={tp_size}), Provider: {provider}") x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) # Preprocess data before benchmarking y_fp8, y_scale = per_block_cast_to_fp8(y) x_fp8, x_scale = sglang_per_token_group_quant_fp8( x, BLOCK_SIZE, column_major_scales=True ) dg_x_fp8, dg_x_scale = sglang_per_token_group_quant_fp8( x, BLOCK_SIZE, column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=True, ) dg_y_fp8, dg_y_scale = requant_weight_ue8m0( y_fp8, y_scale, [BLOCK_SIZE, BLOCK_SIZE] ) quantiles = (0.5, 0.2, 0.8) if provider == "deepgemm": ms, min_ms, max_ms = run_bench( lambda: fp8_gemm_deepgemm_blackwell( dg_x_fp8, dg_x_scale, dg_y_fp8, dg_y_scale, ), quantiles=quantiles, ) elif provider == "flashinfer": ms, min_ms, max_ms = run_bench( lambda: fp8_gemm_flashinfer( x_fp8, x_scale, y_fp8, y_scale, ), quantiles=quantiles, ) # Calculate TFLOPS flops = 2 * m * n * k # multiply-adds tflops = flops / (ms * 1e-3) / 1e12 # Print shape-specific results with TFLOPS print(f"Time: {ms*1000:.2f} us, TFLOPS: {tflops:.2f}") return ms, max_ms, min_ms def get_benchmark_plot_friendly(tp_size): all_configs = create_benchmark_configs(tp_size) x_vals = list(range(len(all_configs))) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["cfg_id"], x_vals=x_vals, line_arg="provider", line_vals=["deepgemm", "flashinfer"], line_names=["DeepGEMM", "Flashinfer"], styles=[("blue", "-"), ("red", "-")], ylabel="us", plot_name=f"fp8-gemm-performance-comparison-tp{tp_size}", args={}, ) ) def benchmark(cfg_id, provider): m, n, k, tp_size = all_configs[cfg_id] ms, min_ms, max_ms = _benchmark(m, n, k, tp_size, provider) return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms return benchmark def get_benchmark(tp_size): all_configs = create_benchmark_configs(tp_size) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["m", "n", "k", "tp_size"], x_vals=[list(config) for config in all_configs], line_arg="provider", line_vals=["deepgemm", "flashinfer"], line_names=["DeepGEMM", "Flashinfer"], styles=[("blue", "-"), ("red", "-")], ylabel="us", plot_name=f"fp8-gemm-performance-comparison-tp{tp_size}", args={}, ) ) def benchmark(m, n, k, tp_size, provider): ms, min_ms, max_ms = _benchmark(m, n, k, tp_size, provider) return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms return benchmark if __name__ == "__main__": if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10: print("Skipping benchmark because the device is not supported") exit(0) parser = argparse.ArgumentParser() parser.add_argument( "--save-path", type=str, default="./configs/benchmark_ops/fp8_gemm/", help="Path to save fp8 gemm benchmark results", ) parser.add_argument( "--run-correctness", action="store_true", default=True, help="Whether to run correctness test", ) parser.add_argument( "--tp-size", type=int, default=1, help="Tensor parallelism size to benchmark (default: 1)", ) parser.add_argument( "--plot-friendly", action="store_true", default=False, help="Plot x axis as the config index instead of the m", ) args = parser.parse_args() # Set random seed for reproducibility torch.manual_seed(0) torch.cuda.manual_seed(0) # Run correctness tests on a few examples if args.run_correctness: print("Running correctness tests...") calculate_diff(64, 512, 7168) # Small test calculate_diff(64, 7168, 16384) # Medium test calculate_diff(64, 18432, 7168) # Large test # Get the benchmark function with the specified tp_size benchmark = ( get_benchmark_plot_friendly(args.tp_size) if args.plot_friendly else get_benchmark(args.tp_size) ) print(f"Running performance benchmark for TP size = {args.tp_size}...") benchmark.run(print_data=True, save_path=args.save_path) ================================================ FILE: benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py ================================================ from typing import Tuple import deep_gemm import torch import triton import triton.language as tl from deep_gemm import calc_diff from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor # Import shared functionality from the regular GEMM benchmark from sglang.benchmark.bench_utils import run_bench from sglang.benchmark.kernels.deepseek.benchmark_deepgemm_fp8_gemm import ( per_block_cast_to_fp8, per_token_cast_to_fp8, ) def construct_grouped_and_flat_fp8( x: torch.Tensor, y: torch.Tensor, num_groups: int, is_masked: bool ) -> Tuple[ Tuple[torch.Tensor, torch.Tensor], # grouped x_fp8 Tuple[torch.Tensor, torch.Tensor], # grouped y_fp8 Tuple[torch.Tensor, torch.Tensor], # flat x_fp8 Tuple[torch.Tensor, torch.Tensor], # flat y_fp8 torch.Tensor, # output torch.Tensor, # reference output ]: # Verify input shapes m, k = x.shape n, k_y = y.shape assert k == k_y, f"Incompatible shapes: x({m}, {k}), y({n}, {k_y})" assert m % num_groups == 0, f"m({m}) must be divisible by num_groups({num_groups})" assert m % 4 == 0, f"TMA alignment error: {m}" # Reshape inputs for grouped processing m_per_group = m // num_groups x_grouped = x.view(num_groups, m_per_group, k) y_grouped = y.unsqueeze(0).expand(num_groups, n, k) # Initialize output tensors out = torch.empty((num_groups, m_per_group, n), device="cuda", dtype=torch.bfloat16) ref_out = torch.einsum("gmk,gnk->gmn", x_grouped, y_grouped) # Quantize grouped tensors x_fp8_grouped = ( torch.empty_like(x_grouped, dtype=torch.float8_e4m3fn), torch.empty( (num_groups, m_per_group, k // 128), device="cuda", dtype=torch.float ), ) y_fp8_grouped = ( torch.empty_like(y_grouped, dtype=torch.float8_e4m3fn), torch.empty( (num_groups, (n + 127) // 128, k // 128), device="cuda", dtype=torch.float ), ) for i in range(num_groups): x_fp8_grouped[0][i], x_fp8_grouped[1][i] = per_token_cast_to_fp8(x_grouped[i]) y_fp8_grouped[0][i], y_fp8_grouped[1][i] = per_block_cast_to_fp8(y_grouped[i]) # Quantize flat tensors x_fp8_flat = per_token_cast_to_fp8(x) y_fp8_flat = per_block_cast_to_fp8(y) # For non-masked input, merge the group and M dims in output if not is_masked: x_fp8_grouped = ( x_fp8_grouped[0].view(-1, k), per_token_cast_to_fp8(x_grouped.view(-1, k))[1], ) out, ref_out = out.view(-1, n), ref_out.view(-1, n) # Transpose earlier for testing x_fp8_grouped = ( x_fp8_grouped[0], get_mn_major_tma_aligned_tensor(x_fp8_grouped[1]), ) x_fp8_flat = (x_fp8_flat[0], get_mn_major_tma_aligned_tensor(x_fp8_flat[1])) return x_fp8_grouped, y_fp8_grouped, x_fp8_flat, y_fp8_flat, out, ref_out # Since we don't have a group gemm kernel in SGLang/vLLM, we implemented a # custom kernel based on the Triton tutorial. # https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html @triton.jit def fp8_gemm_group_triton_kernel( # Pointers to matrices a_ptr, b_ptr, c_ptr, # Pointers to scaling factors a_scale_ptr, b_scale_ptr, # Matrix dimensions M, N, K, # The stride variables represent how much to increase the ptr by when moving by 1 # element in a particular dimension. stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, # Strides for scaling factors stride_a_scale_m, stride_a_scale_k, stride_b_scale_n, stride_b_scale_k, # Meta-parameters BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, ): """Kernel for computing the matmul C = A x B with FP8 inputs and scaling factors. A has shape (M, K), B has shape (K, N) and C has shape (M, N) Note: Block sizes must be multiples of 32 for optimal TMA performance. """ # Map program ids to the block of C it should compute pid_group = tl.program_id(axis=0) # Group ID pid_n = tl.program_id(axis=1) # N dimension ID # Compute the M block ID within this group group_size_m = min(M - pid_group * GROUP_SIZE_M, GROUP_SIZE_M) pid_m_within_group = tl.program_id(axis=2) % group_size_m pid_m = pid_group * GROUP_SIZE_M + pid_m_within_group # Create pointers for the first blocks of A and B offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N offs_k = tl.arange(0, BLOCK_SIZE_K) a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) # Initialize accumulator accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) # Main loop for k_block in range(0, tl.cdiv(K, BLOCK_SIZE_K)): k_offset = k_block * BLOCK_SIZE_K # Load the next block of A and B, with masks a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k_offset, other=0.0) b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k_offset, other=0.0) # Calculate indices for scaling factors for this K block a_scale_ptrs = a_scale_ptr + ( offs_am * stride_a_scale_m + k_block * stride_a_scale_k ) b_scale_ptrs = b_scale_ptr + ( pid_n * stride_b_scale_n + k_block * stride_b_scale_k ) # Perform matrix multiplication in FP8 res = tl.dot(a, b) # Load scaling factors for the current block a_scale = tl.load(a_scale_ptrs)[:, None] # [BLOCK_SIZE_M, 1] b_scale = tl.load(b_scale_ptrs) # Apply scaling factors to the accumulated result accumulator += res * a_scale * b_scale # Advance pointers a_ptrs += BLOCK_SIZE_K * stride_ak b_ptrs += BLOCK_SIZE_K * stride_bk # Convert to bfloat16 for output c = accumulator.to(tl.bfloat16) # Write back the result offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) tl.store(c_ptrs, c, mask=c_mask) def fp8_gemm_group_triton(a_tuple, b_tuple, c, num_groups): """ Perform matrix multiplication with FP8 inputs and proper scaling. Args: a_tuple: Tuple of (quantized_tensor, scale_factors) for input A b_tuple: Tuple of (quantized_tensor, scale_factors) for input B c: Output tensor in BF16 format num_groups: Number of groups for grouped GEMM Returns: Result tensor in BF16 format """ # Unpack the tuples a, a_scale = a_tuple b, b_scale = b_tuple M, K = a.shape _, N = b.shape # Configure block sizes - must be multiples of 32 for TMA alignment BLOCK_SIZE_M = 128 BLOCK_SIZE_N = 128 BLOCK_SIZE_K = 128 # Calculate grid dimensions num_pid_m = triton.cdiv(M, BLOCK_SIZE_M) num_pid_n = triton.cdiv(N, BLOCK_SIZE_N) num_groups_grid = triton.cdiv(num_pid_m, num_groups) # 3D grid launch - (group, n_blocks, m_blocks_per_group) grid = (num_groups_grid, num_pid_n, min(num_groups, num_pid_m)) fp8_gemm_group_triton_kernel[grid]( a, b, c, a_scale, b_scale, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1), a_scale.stride(0), 1, # Stride in the K dimension may be 1 b_scale.stride(0), 1 if b_scale.dim() > 1 else 0, BLOCK_SIZE_M=BLOCK_SIZE_M, BLOCK_SIZE_N=BLOCK_SIZE_N, BLOCK_SIZE_K=BLOCK_SIZE_K, GROUP_SIZE_M=num_groups, ) return c def fp8_gemm_group_deepgemm(x_fp8_grouped, y_fp8_grouped, out, m_indices): deep_gemm.m_grouped_fp8_gemm_nt_contiguous( x_fp8_grouped, y_fp8_grouped, out, m_indices, ) return out def calculate_diff(m: int, n: int, k: int, num_groups: int): print(f"Shape (m={m}, n={n}, k={k}") x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) x_fp8_grouped, y_fp8_grouped, x_fp8_flat, y_fp8_flat, out, out_torch = ( construct_grouped_and_flat_fp8(x, y, num_groups, is_masked=False) ) m_per_group = m // num_groups out_deepgemm = out.clone() m_indices = torch.arange(0, num_groups, device="cuda", dtype=torch.int) m_indices = ( m_indices.unsqueeze(-1).expand(num_groups, m_per_group).contiguous().view(-1) ) fp8_gemm_group_deepgemm( x_fp8_grouped, y_fp8_grouped, out_deepgemm, m_indices, ) torch.cuda.synchronize() # Prepare inputs for Triton a, a_scale = x_fp8_flat b, b_scale = y_fp8_flat b = b.T.contiguous() # Ensure scales are in the right format and contiguous a_scale, b_scale = a_scale.contiguous(), b_scale.contiguous() M, _ = a.shape _, N = b.shape c = torch.empty((M, N), device=a.device, dtype=torch.bfloat16) out_triton = fp8_gemm_group_triton((a, a_scale), (b, b_scale), c, num_groups) torch.cuda.synchronize() diff_torch_deepgemm = torch.abs(out_torch - out_deepgemm).mean().item() diff_torch_triton = torch.abs(out_torch - out_triton).mean().item() diff_deepgemm_triton = torch.abs(out_deepgemm - out_triton).mean().item() print(f"Shape m={m}, n={n}, k={k}:") print(f"Torch output: {out_torch[0, 0:5]}") print(f"DeepGEMM output: {out_deepgemm[0, 0:5]}") print(f"Triton output: {out_triton[0, 0:5]}") print(f"Mean absolute difference (Torch-DeepGEMM): {diff_torch_deepgemm}") print(f"Mean absolute difference (Torch-Triton): {diff_torch_triton}") print(f"Mean absolute difference (DeepGEMM-Triton): {diff_deepgemm_triton}") deepgemm_torch_diff = calc_diff(out_deepgemm, out_torch) triton_torch_diff = calc_diff(out_triton, out_torch) deepgemm_triton_diff = calc_diff(out_deepgemm, out_triton) DIFF_THRESHOLD = 0.001 all_match = ( deepgemm_torch_diff < DIFF_THRESHOLD and triton_torch_diff < DIFF_THRESHOLD and deepgemm_triton_diff < DIFF_THRESHOLD ) if all_match: print("✅ All implementations match\n") else: print("❌ Some implementations differ:") print( f" - Torch vs DeepGEMM: {'✅' if deepgemm_torch_diff < DIFF_THRESHOLD else '❌'}" f" - Torch vs Triton: {'✅' if triton_torch_diff < DIFF_THRESHOLD else '❌'}" f" - DeepGEMM vs Triton: {'✅' if deepgemm_triton_diff < DIFF_THRESHOLD else '❌'}" ) def get_weight_shapes(tp_size): # cannot TP total = [ (512 + 64, 7168), ((128 + 64) * 128, 7168), (128 * (128 + 128), 512), (7168, 16384), (7168, 18432), ] # N can TP n_tp = [ (18432 * 2, 7168), ((128 + 64) * 128, 7168), (128 * (128 + 128), 512), (24576, 1536), (4096, 7168), ] # K can TP k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] weight_shapes = [] for t in total: weight_shapes.append(t) for n_t in n_tp: new_t = (n_t[0] // tp_size, n_t[1]) weight_shapes.append(new_t) for k_t in k_tp: new_t = (k_t[0], k_t[1] // tp_size) weight_shapes.append(new_t) return weight_shapes def create_benchmark_configs(tp_size): configs = [] weight_shapes = get_weight_shapes(tp_size) batch_sizes = [2048, 4096] group_sizes = [4, 8] for n, k in weight_shapes: for m in batch_sizes: for num_groups in group_sizes: configs.append((m, n, k, num_groups, tp_size)) return configs def get_benchmark(tp_size): all_configs = create_benchmark_configs(tp_size) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["m", "n", "k", "num_groups", "tp_size"], x_vals=[config for config in all_configs], line_arg="provider", line_vals=["deepgemm", "triton"], line_names=["DeepGEMM", "Triton"], styles=[("blue", "-"), ("red", "-")], ylabel="ms", plot_name=f"fp8-group-gemm-performance-comparison-tp{tp_size}", args={}, ) ) def benchmark(m, n, k, num_groups, tp_size, provider): print( f"Shape (m={m}, n={n}, k={k}, tp={tp_size}, num_groups={num_groups}, Provider: {provider}" ) x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) x_fp8_grouped, y_fp8_grouped, x_fp8_flat, y_fp8_flat, out, out_torch = ( construct_grouped_and_flat_fp8(x, y, num_groups, is_masked=False) ) m_per_group = m // num_groups m_indices = torch.arange(0, num_groups, device="cuda", dtype=torch.int) m_indices = ( m_indices.unsqueeze(-1) .expand(num_groups, m_per_group) .contiguous() .view(-1) ) quantiles = (0.5, 0.2, 0.8) if provider == "deepgemm": ms, min_ms, max_ms = run_bench( lambda: fp8_gemm_group_deepgemm( x_fp8_grouped, y_fp8_grouped, out, m_indices, ), quantiles=quantiles, ) elif provider == "triton": # Prepare inputs for Triton # We did it outside of the lambda function to make it fair comparison like deepgemm a, a_scale = x_fp8_flat b, b_scale = y_fp8_flat b = b.T.contiguous() # Ensure scales are in the right format and contiguous a_scale, b_scale = a_scale.contiguous(), b_scale.contiguous() M, _ = a.shape _, N = b.shape c = torch.empty((M, N), device=a.device, dtype=torch.bfloat16) ms, min_ms, max_ms = run_bench( lambda: fp8_gemm_group_triton( (a, a_scale), (b, b_scale), c, num_groups, ), quantiles=quantiles, ) # Calculate TFLOPS flops = 2 * m * n * k # multiply-adds tflops = flops / (ms * 1e-3) / 1e12 print(f"Time: {ms*1000:.2f} ms, TFLOPS: {tflops:.2f}") return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms return benchmark if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() parser.add_argument( "--save_path", type=str, default="./configs/benchmark_ops/fp8_group_gemm/", help="Path to save deepgemm fp8 group gemm benchmark results", ) parser.add_argument( "--run_correctness", action="store_true", help="Whether to run correctness test", ) parser.add_argument( "--tp_size", type=int, default=1, help="Tensor parallelism size to benchmark (default: 1)", ) args = parser.parse_args() # Set random seed for reproducibility torch.manual_seed(0) torch.cuda.manual_seed(0) # Enable TF32, adapted from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_core.py#L148 torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True # Run correctness tests on a few examples if args.run_correctness: print("Running correctness tests...") calculate_diff(8192, 7168, 4096, 4) calculate_diff(8192, 2048, 7168, 4) calculate_diff(4096, 7168, 4096, 8) calculate_diff(4096, 2048, 7168, 8) calculate_diff(4096, 576, 7168, 8) # Get the benchmark function with the specified tp_size benchmark = get_benchmark(args.tp_size) print(f"Running performance benchmark for TP size = {args.tp_size}...") benchmark.run(print_data=True, save_path=args.save_path) ================================================ FILE: benchmark/kernels/elementwise/benchmark_concat_mla.py ================================================ import torch import triton import triton.language as tl from sgl_kernel import concat_mla_k as concat_mla_k_cuda from sglang.benchmark.bench_utils import run_bench DEVICE = triton.runtime.driver.active.get_active_torch_device() num_local_heads = 128 qk_nope_head_dim = 128 qk_rope_head_dim = 64 def create_data(num_tokens): k_nope_container = torch.randn( (num_tokens, num_local_heads, qk_nope_head_dim + 128), dtype=torch.bfloat16, device="cuda", ) k_nope = k_nope_container[:, :, :qk_nope_head_dim] k_rope_container = torch.randn( (num_tokens, 1, 128 + qk_rope_head_dim), dtype=torch.bfloat16, device="cuda" ) k_rope = k_rope_container[:, :, -qk_rope_head_dim:] k = torch.empty( (num_tokens, num_local_heads, qk_nope_head_dim + qk_rope_head_dim), dtype=torch.bfloat16, device="cuda", ) return dict(k=k, k_nope=k_nope, k_rope=k_rope) def fn_torch(k, k_nope, k_rope): k[..., :qk_nope_head_dim] = k_nope k[..., qk_nope_head_dim:] = k_rope def fn_hack_non_strided(k, k_nope, k_rope): k_flatten_view = k.flatten() k_flatten_view[: k_nope.numel()] = k_nope.flatten() k2 = k_flatten_view[k_nope.numel() :].view(k_rope.numel(), -1) k2 = k_rope.flatten()[:, None] @torch.compile(dynamic=True) def fn_torch_compiled(k, k_nope, k_rope): return fn_torch(k, k_nope, k_rope) def fn_cuda(k, k_nope, k_rope): concat_mla_k_cuda(k, k_nope, k_rope) @triton.jit def fn_triton_kernel( k_ptr, k_nope_ptr, k_rope_ptr, num_tokens, QK_NOPE_HEAD_DIM: tl.constexpr, QK_ROPE_HEAD_DIM: tl.constexpr, NUM_LOCAL_HEADS: tl.constexpr, K_NOPE_STRIDE_0: tl.constexpr, K_NOPE_STRIDE_1: tl.constexpr, K_STRIDE_0: tl.constexpr, K_STRIDE_1: tl.constexpr, K_ROPE_STRIDE_0: tl.constexpr, BLOCK_ROWS: tl.constexpr, ): pid = tl.program_id(axis=0) token_id = pid * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) token_mask = token_id < num_tokens head_id = tl.arange(0, NUM_LOCAL_HEADS) # nope nope_sub_id = tl.arange(0, QK_NOPE_HEAD_DIM) offs_nope = ( token_id[:, None, None] * K_NOPE_STRIDE_0 + head_id[None, :, None] * K_NOPE_STRIDE_1 + nope_sub_id[None, None, :] ) offs_k = ( token_id[:, None, None] * K_STRIDE_0 + head_id[None, :, None] * K_STRIDE_1 + nope_sub_id[None, None, :] ) vals_nope = tl.load(k_nope_ptr + offs_nope, mask=token_mask[:, None, None]) tl.store(k_ptr + offs_k, vals_nope, mask=token_mask[:, None, None]) # rope rope_sub_id = tl.arange(0, QK_ROPE_HEAD_DIM) offs_rope = token_id[:, None, None] * K_ROPE_STRIDE_0 + rope_sub_id[None, None, :] offs_k = ( token_id[:, None, None] * K_STRIDE_0 + head_id[None, :, None] * K_STRIDE_1 + rope_sub_id[None, None, :] + QK_NOPE_HEAD_DIM ) vals_rope = tl.load(k_rope_ptr + offs_rope, mask=token_mask[:, None, None]) tl.store(k_ptr + offs_k, vals_rope, mask=token_mask[:, None, None]) def fn_triton(k, k_nope, k_rope): assert k.device == DEVICE and k_nope.device == DEVICE and k_rope.device == DEVICE num_tokens, _, _ = k.shape grid = lambda meta: (triton.cdiv(num_tokens, meta["BLOCK_ROWS"]),) fn_triton_kernel[grid]( k, k_nope, k_rope, num_tokens, QK_NOPE_HEAD_DIM=qk_nope_head_dim, QK_ROPE_HEAD_DIM=qk_rope_head_dim, NUM_LOCAL_HEADS=num_local_heads, K_NOPE_STRIDE_0=k_nope.stride(0), K_NOPE_STRIDE_1=k_nope.stride(1), K_STRIDE_0=k.stride(0), K_STRIDE_1=k.stride(1), K_ROPE_STRIDE_0=k_rope.stride(0), BLOCK_ROWS=16, ) def execute_and_get_output(f, data): data["k"].zero_() f(**data) assert data["k"].sum().item() != 0 return data["k"].clone() torch.manual_seed(0) data = create_data(num_tokens=32768) output_ref = execute_and_get_output(fn_torch, data) output_exp = execute_and_get_output(fn_cuda, data) # print(output_ref) # print(output_exp) if not torch.all(output_ref == output_exp): abs_delta = torch.abs(output_ref - output_exp) raise AssertionError( f"{output_ref=} {output_exp=} " f"{abs_delta=} " f"{torch.argwhere(abs_delta != 0.0)=} " ) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["num_tokens"], # Argument names to use as an x-axis for the plot. x_vals=[ 2048, 4096, 8192, 16384, 32768, ], # Different possible values for `x_name`. x_log=False, # x axis is logarithmic. line_arg="provider", # Argument name whose value corresponds to a different line in the plot. line_vals=[ "torch", "torch_compiled", "triton", "hack_non_strided", "cuda", ], # Possible values for `line_arg`. line_names=[ "torch", "torch_compiled", "triton", "hack_non_strided", "cuda", ], # Label name for the lines. plot_name="vector-add-performance", # Name for the plot. Used also as a file name for saving the plot. args={}, # Values for function arguments not in `x_names` and `y_name`. ) ) def benchmark(num_tokens, provider): data = create_data(num_tokens=num_tokens) quantiles = (0.5, 0.2, 0.8) fn = { "torch": fn_torch, "torch_compiled": fn_torch_compiled, "triton": fn_triton, "hack_non_strided": fn_hack_non_strided, "cuda": fn_cuda, }[provider] ms, min_ms, max_ms = run_bench(lambda: fn(**data), quantiles=quantiles) return ms, min_ms, max_ms torch.cuda.cudart().cudaProfilerStart() benchmark.run(print_data=True, show_plots=True) torch.cuda.cudart().cudaProfilerStop() ================================================ FILE: benchmark/kernels/flashinfer_allreduce_fusion/README.md ================================================ # FlashInfer Fused AllReduce + RMSNorm Benchmark This benchmark script is modified from the [original implementation](https://github.com/vllm-project/vllm/blob/237e1fb887c7f5a579420fa0295097f24b006594/benchmarks/kernels/benchmark_fused_collective.py) by the vLLM community. It aims to compare the performance differences between FlashInfer fused operators in SGLang (trtllm_allreduce_fusion: AllReduce + Residual Add + RMSNorm + optional quantization) and conventional implementations (standard `tensor_model_parallel_all_reduce` + separate RMSNorm/quantization). Specifically, this script tests the timing performance of two implementation paths: 1) Standard AllReduce and RMSNorm executed separately; 2) FlashInfer's fused operator combining AllReduce, Residual Add, RMSNorm, and optional quantization operations. This benchmark script helps us tune the ipc workspace size of the `flashinfer_allreduce_residual_rmsnorm` operator in SGLang and prepare for applications with FP8/FP4 quantized fused operators. Script path: `benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py` ## Feature Overview - Compare average execution time (ms) and calculate speedup ratios for the following paths: - standard_allreduce_rmsnorm (Standard AllReduce + RMSNorm) - flashinfer_fused_allreduce_rmsnorm (Fused AllReduce + RMSNorm), including oneshot and twoshot modes - Optionally compare FP8/FP4 quantized fused paths with standard paths - Use CUDA Graph capture and batch replay to reduce measurement noise - Automatically select the faster "standard baseline" (native/compiled version) as the denominator for speedup calculation - Optionally export results in Markdown format ## Runtime Environment and Prerequisites - At least 2 GPUs, and launch multi-process distributed training using `torchrun` (NCCL backend) - Properly install/compile sglang along with sgl-kernel and custom operators ## Quick Start (Command Examples) The following examples use world_size=2. You can modify `--nproc_per_node` and parameters according to your machine: - Regular paths only (no quantization): ``` torchrun --nproc_per_node=2 \ benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py \ --no-quant --hidden-dim 1024 --seq-lens 512 1024 2048 4096 --trials 100 ``` - FP8 quantization paths only: ``` torchrun --nproc_per_node=2 \ benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py \ --quant-fp8 --hidden-dim 1024 --seq-lens 512 1024 2048 4096 --trials 100 ``` - FP4 quantization paths only: ``` torchrun --nproc_per_node=2 \ benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py \ --quant-fp4 --hidden-dim 1024 --seq-lens 512 1024 2048 4096 --trials 100 ``` - Larger hidden dimensions: ``` torchrun --nproc_per_node=2 \ benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py \ --no-quant --hidden-dim 4096 --seq-lens 512 1024 2048 4096 --trials 100 ``` ## Parameter Description - `--seq-lens`: List of sequence lengths to test (default: 128 512 1024 2048) - `--hidden-dim`: Hidden dimension (default: 8192) - `--dtypes`: Data type list, `float16|bfloat16|float32` (default: bfloat16) - `--no-residual`: Only test "no residual" scenarios (default tests both "with/without residual") - Mutually exclusive quantization options: - `--no-quant`: No quantization testing - `--quant-fp8`: Only FP8 quantization testing - `--quant-fp4`: Only FP4 quantization testing - `--quant-all`: Test all (default) - FlashInfer related: - `--disable-oneshot`: Disable oneshot mode (default enables oneshot and tests twoshot simultaneously) - Runtime configuration: - `--warmup`: Warmup count before graph capture and before graph replay (default 5) - `--trials`: Benchmark iteration count (default 20; internally each `graph.replay()` will batch replay multiple times) - `--output-file`: Save results as Markdown file (only rank0 takes effect) ## Output Example Each configuration group prints a table showing average execution time and relative speedup ratios (baseline is the faster standard implementation). For example: ``` ================================================================================ Results: seq_len=1024, hidden_dim=1024 dtype=torch.bfloat16, residual=yes, quant_mode=none ================================================================================ Operation Time (ms) Speedup -------------------------------------------------------------------------------- standard_allreduce_rmsnorm 0.024 0.98x standard_allreduce_rmsnorm_native_compiled 0.023 baseline flashinfer_fused_allreduce_rmsnorm_oneshot 0.011 2.19x flashinfer_fused_allreduce_rmsnorm_twoshot 0.041 0.57x ``` If `--output-file` is specified, all configurations will be summarized in Markdown tables in that file. ## Important Notes and Recommendations - Distributed: The script uses `torchrun` environment variables to initialize distributed training and binds tensors/communication groups to the current rank's corresponding device. - World size: Requires `WORLD_SIZE > 1` to perform communication operator benchmarks. Otherwise, the script will error and prompt. - FlashInfer: - If not installed or interfaces are missing, the script will only run standard paths and provide prompts in the logs. - The fused operator internally uses "oneshot"/"twoshot" two trigger methods; oneshot is enabled by default and twoshot is tested simultaneously. - FP8/FP4: - FP8 uses sglang's FP8 tools and dtype, with underlying platform selection of `e4m3`/`e4m3fnuz` etc. - FP4 uses sgl-kernel's `scaled_fp4_quant`, requiring corresponding platform support. - CUDA Graph: - Uses sglang's `graph_capture()` to prepare capture-ready state for communication, then uses `torch.cuda.graph` to capture kernels, reducing measurement jitter. ================================================ FILE: benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py ================================================ # Modified from https://github.com/vllm-project/vllm/blob/237e1fb887c7f5a579420fa0295097f24b006594/benchmarks/kernels/benchmark_fused_collective.py """ Benchmark for FlashInfer fused collective operations vs standard operations. This benchmark compares: 1. FlashInfer's trtllm_allreduce_fusion (fused allreduce + rmsnorm + optional quant) 2. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations Usage with torchrun: torchrun --nproc_per_node=2 benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py --no-quant --hidden-dim 1024 --seq-len 512 1024 2048 4096 --trials 100 torchrun --nproc_per_node=2 benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py --quant-fp8 --hidden-dim 1024 --seq-len 512 1024 2048 4096 --trials 100 torchrun --nproc_per_node=2 benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py --quant-fp4 --hidden-dim 1024 --seq-len 512 1024 2048 4096 --trials 100 torchrun --nproc_per_node=2 benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py --no-quant --hidden-dim 4096 --seq-len 512 1024 2048 4096 --trials 100 torchrun --nproc_per_node=2 benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py --quant-fp8 --hidden-dim 4096 --seq-len 512 1024 2048 4096 --trials 100 torchrun --nproc_per_node=2 benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py --quant-fp4 --hidden-dim 4096 --seq-len 512 1024 2048 4096 --trials 100 """ import argparse import contextlib import itertools import logging import os import time from typing import Optional import torch # type: ignore import torch.distributed as dist # type: ignore from sglang.srt.distributed import get_tp_group, tensor_model_parallel_all_reduce from sglang.srt.distributed.parallel_state import ( cleanup_dist_env_and_memory, graph_capture, init_distributed_environment, initialize_model_parallel, ) from sglang.srt.layers.layernorm import RMSNorm # noqa from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype as SGLANG_FP8_DTYPE from sglang.srt.layers.quantization.fp8_kernel import static_quant_fp8 try: from sgl_kernel import fused_add_rmsnorm as SGL_FUSED_ADD_RMS_NORM from sgl_kernel import rmsnorm as SGL_RMS_NORM from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as SGL_SCALED_FP4_QUANT except Exception: # pragma: no cover - fallback on non-supported platforms SGL_FUSED_ADD_RMS_NORM = None SGL_RMS_NORM = None SGL_SCALED_FP4_QUANT = None FP8_DTYPE = SGLANG_FP8_DTYPE logger = logging.getLogger(__name__) # Try to import FlashInfer try: import flashinfer.comm as flashinfer_comm # type: ignore if not hasattr(flashinfer_comm, "trtllm_allreduce_fusion"): flashinfer_comm = None logger.warning( "FlashInfer comm module found but missing trtllm_allreduce_fusion" ) except ImportError: flashinfer_comm = None logger.warning("FlashInfer not found, only benchmarking standard operations") # Constants MiB = 1024 * 1024 # FlashInfer max sizes per world size # Enable 64MB for 2, 4, 8 world sizes to verify large input sizes # use --disable-oneshot to disable oneshot mode for very large input sizes _FI_MAX_SIZES = { 2: 64 * MiB, # 64MB 4: 64 * MiB, # 64MB 8: 64 * MiB, # 64MB } # Global workspace tensor for FlashInfer _FI_WORKSPACE_TENSOR = None def setup_flashinfer_workspace( world_size: int, rank: int, hidden_dim: int, max_token_num: int, use_fp32_lamport: bool = False, ): """Setup FlashInfer workspace for fused allreduce operations.""" global _FI_WORKSPACE_TENSOR if flashinfer_comm is None: return None, None if world_size not in _FI_MAX_SIZES: logger.warning("FlashInfer not supported for world size %s", world_size) return None, None try: # Create IPC workspace ipc_handles, workspace_tensor = ( flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion( tp_rank=rank, tp_size=world_size, max_token_num=max_token_num, hidden_dim=hidden_dim, group=get_tp_group().device_group, use_fp32_lamport=use_fp32_lamport, ) ) _FI_WORKSPACE_TENSOR = workspace_tensor return ipc_handles, workspace_tensor except Exception as e: logger.error("Failed to setup FlashInfer workspace: %s", e) return None, None def cleanup_flashinfer_workspace(ipc_handles): """Cleanup FlashInfer workspace.""" if flashinfer_comm is None or ipc_handles is None: return try: group = get_tp_group().device_group flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce(ipc_handles, group) except Exception as e: logger.error("Failed to cleanup FlashInfer workspace: %s", e) class FlashInferFusedAllReduceParams: """Parameters for FlashInfer fused allreduce operations.""" def __init__( self, rank: int, world_size: int, use_fp32_lamport: bool = False, max_token_num: int = 1024, ): self.rank = rank self.world_size = world_size self.use_fp32_lamport = use_fp32_lamport self.trigger_completion_at_end = True self.launch_with_pdl = True self.fp32_acc = True self.max_token_num = max_token_num def get_trtllm_fused_allreduce_kwargs(self): return { "world_rank": self.rank, "world_size": self.world_size, "launch_with_pdl": self.launch_with_pdl, "trigger_completion_at_end": self.trigger_completion_at_end, "fp32_acc": self.fp32_acc, } def flashinfer_fused_allreduce_rmsnorm( input_tensor: torch.Tensor, residual: Optional[torch.Tensor], rms_gamma: torch.Tensor, rms_eps: float, allreduce_params: "FlashInferFusedAllReduceParams", use_oneshot: bool, norm_out: Optional[torch.Tensor] = None, ): """FlashInfer fused allreduce + rmsnorm operation.""" if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None: raise RuntimeError("FlashInfer not available or workspace not initialized") if norm_out is None: norm_out = input_tensor residual_out = residual else: residual_out = input_tensor flashinfer_comm.trtllm_allreduce_fusion( allreduce_in=input_tensor, token_num=input_tensor.shape[0], residual_in=residual, residual_out=residual_out, norm_out=norm_out, rms_gamma=rms_gamma, rms_eps=rms_eps, hidden_dim=input_tensor.shape[-1], workspace_ptrs=_FI_WORKSPACE_TENSOR, pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm, allreduce_out=None, quant_out=None, scale_out=None, layout_code=None, scale_factor=None, use_oneshot=use_oneshot, **allreduce_params.get_trtllm_fused_allreduce_kwargs(), ) def flashinfer_fused_allreduce_rmsnorm_fp8_quant( input_tensor: torch.Tensor, residual: Optional[torch.Tensor], rms_gamma: torch.Tensor, rms_eps: float, scale_factor: torch.Tensor, allreduce_params: FlashInferFusedAllReduceParams, use_oneshot: bool = True, norm_out: Optional[torch.Tensor] = None, quant_out: Optional[torch.Tensor] = None, ): """FlashInfer fused allreduce + rmsnorm + FP8 quantization.""" if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None: raise RuntimeError("FlashInfer not available or workspace not initialized") if norm_out is None: norm_out = input_tensor residual_out = residual else: residual_out = input_tensor flashinfer_comm.trtllm_allreduce_fusion( allreduce_in=input_tensor, token_num=input_tensor.shape[0], residual_in=residual, residual_out=residual_out, norm_out=norm_out, rms_gamma=rms_gamma, rms_eps=rms_eps, hidden_dim=input_tensor.shape[-1], workspace_ptrs=_FI_WORKSPACE_TENSOR, pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, allreduce_out=None, quant_out=quant_out, scale_out=None, layout_code=None, scale_factor=scale_factor, use_oneshot=use_oneshot, **allreduce_params.get_trtllm_fused_allreduce_kwargs(), ) def flashinfer_fused_allreduce_rmsnorm_fp4_quant( input_tensor: torch.Tensor, residual: Optional[torch.Tensor], rms_gamma: torch.Tensor, rms_eps: float, input_global_scale: torch.Tensor, allreduce_params: FlashInferFusedAllReduceParams, quant_out: torch.Tensor, use_oneshot: bool, output_scale: torch.Tensor, norm_out: Optional[torch.Tensor] = None, ): """FlashInfer fused allreduce + rmsnorm + FP4 quantization.""" if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None: raise RuntimeError("FlashInfer not available or workspace not initialized") if norm_out is None: norm_out = input_tensor residual_out = residual else: residual_out = input_tensor flashinfer_comm.trtllm_allreduce_fusion( allreduce_in=input_tensor, token_num=input_tensor.shape[0], residual_in=residual, residual_out=residual_out, norm_out=norm_out, rms_gamma=rms_gamma, rms_eps=rms_eps, hidden_dim=input_tensor.shape[-1], workspace_ptrs=_FI_WORKSPACE_TENSOR, pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, allreduce_out=None, quant_out=quant_out, scale_out=output_scale, layout_code=None, scale_factor=input_global_scale, use_oneshot=use_oneshot, **allreduce_params.get_trtllm_fused_allreduce_kwargs(), ) def standard_allreduce_rmsnorm( input_tensor: torch.Tensor, residual: Optional[torch.Tensor], rms_gamma: torch.Tensor, rms_eps: float, norm_out: Optional[torch.Tensor] = None, ): """Standard allreduce + rmsnorm operations.""" # All-reduce first allreduce_out = tensor_model_parallel_all_reduce(input_tensor) # Then RMS norm if residual is not None: # Fused add + RMS norm (in-place on allreduce_out) if SGL_FUSED_ADD_RMS_NORM is not None: SGL_FUSED_ADD_RMS_NORM(allreduce_out, residual, rms_gamma, rms_eps) else: rms = RMSNorm(allreduce_out.shape[-1], eps=rms_eps) rms.weight.data = rms_gamma rms.forward_native(allreduce_out, residual) else: # Just RMS norm if SGL_RMS_NORM is not None: _ = SGL_RMS_NORM(allreduce_out, rms_gamma, rms_eps) else: rms = RMSNorm(allreduce_out.shape[-1], eps=rms_eps) rms.weight.data = rms_gamma _ = rms.forward_native(allreduce_out) def standard_allreduce_rmsnorm_fp8_quant( input_tensor: torch.Tensor, residual: Optional[torch.Tensor], rms_gamma: torch.Tensor, rms_eps: float, scale_factor: torch.Tensor, norm_out: Optional[torch.Tensor] = None, quant_out: Optional[torch.Tensor] = None, ): """Standard allreduce + rmsnorm + FP8 quantization.""" # All-reduce first allreduce_out = tensor_model_parallel_all_reduce(input_tensor) # Then RMS norm + static FP8 quantization if residual is not None: if SGL_FUSED_ADD_RMS_NORM is not None: SGL_FUSED_ADD_RMS_NORM(allreduce_out, residual, rms_gamma, rms_eps) quant_out, _ = static_quant_fp8( allreduce_out, scale_factor, repeat_scale=False ) else: rms = RMSNorm(allreduce_out.shape[-1], eps=rms_eps) rms.weight.data = rms_gamma normed, _ = rms.forward_native(allreduce_out, residual) quant_out, _ = static_quant_fp8(normed, scale_factor, repeat_scale=False) return quant_out, residual else: if SGL_RMS_NORM is not None: normed = SGL_RMS_NORM(allreduce_out, rms_gamma, rms_eps) else: rms = RMSNorm(allreduce_out.shape[-1], eps=rms_eps) rms.weight.data = rms_gamma normed = rms.forward_native(allreduce_out) quant_out, _ = static_quant_fp8(normed, scale_factor, repeat_scale=False) return quant_out def standard_allreduce_rmsnorm_fp4_quant( input_tensor: torch.Tensor, residual: Optional[torch.Tensor], rms_gamma: torch.Tensor, rms_eps: float, input_global_scale: torch.Tensor, quant_out: torch.Tensor, output_scale: torch.Tensor, norm_out: Optional[torch.Tensor] = None, ): """Standard allreduce + rmsnorm + FP4 quantization.""" # All-reduce first allreduce_out = tensor_model_parallel_all_reduce(input_tensor) # Then RMS norm if residual is not None: if SGL_FUSED_ADD_RMS_NORM is not None: SGL_FUSED_ADD_RMS_NORM(allreduce_out, residual, rms_gamma, rms_eps) quant_input = allreduce_out else: rms = RMSNorm(allreduce_out.shape[-1], eps=rms_eps) rms.weight.data = rms_gamma quant_input, _ = rms.forward_native(allreduce_out, residual) residual_out = residual else: if SGL_RMS_NORM is not None: quant_input = SGL_RMS_NORM(allreduce_out, rms_gamma, rms_eps) else: rms = RMSNorm(allreduce_out.shape[-1], eps=rms_eps) rms.weight.data = rms_gamma quant_input = rms.forward_native(allreduce_out) residual_out = allreduce_out # Finally FP4 quantization if SGL_SCALED_FP4_QUANT is None: raise RuntimeError("scaled_fp4_quant is not available on this platform") quant_res, output_scale_res = SGL_SCALED_FP4_QUANT(quant_input, input_global_scale) if residual is not None: return quant_res, residual_out, output_scale_res else: return quant_res, quant_input def standard_allreduce_rmsnorm_native( input_tensor: torch.Tensor, residual: Optional[torch.Tensor], rmsnorm_layer: RMSNorm, norm_out: Optional[torch.Tensor] = None, ): """Standard allreduce + rmsnorm operations using native RMSNorm forward.""" # All-reduce first allreduce_out = tensor_model_parallel_all_reduce(input_tensor) # Apply native RMSNorm if residual is not None: result = rmsnorm_layer.forward_native(allreduce_out, residual) return result # Returns (norm_out, residual_out) else: result = rmsnorm_layer.forward_native(allreduce_out) return result # Returns norm_out def standard_allreduce_rmsnorm_fp8_quant_native( input_tensor: torch.Tensor, residual: Optional[torch.Tensor], rmsnorm_layer: RMSNorm, scale_factor: torch.Tensor, norm_out: Optional[torch.Tensor] = None, quant_out: Optional[torch.Tensor] = None, ): """Standard allreduce + rmsnorm + FP8 quantization using native implementations.""" # All-reduce first allreduce_out = tensor_model_parallel_all_reduce(input_tensor) # Apply native RMSNorm if residual is not None: norm_out, residual_out = rmsnorm_layer.forward_native(allreduce_out, residual) else: norm_out = rmsnorm_layer.forward_native(allreduce_out) residual_out = allreduce_out # Apply native FP8 quantization quant_out, _ = static_quant_fp8(norm_out, scale_factor, repeat_scale=False) if residual is not None: return quant_out, residual_out else: return quant_out def standard_allreduce_rmsnorm_fp4_quant_native( input_tensor: torch.Tensor, residual: Optional[torch.Tensor], rmsnorm_layer: RMSNorm, input_global_scale: torch.Tensor, quant_out: torch.Tensor, output_scale: torch.Tensor, norm_out: Optional[torch.Tensor] = None, ): """Standard allreduce + rmsnorm + FP4 quantization using native RMSNorm.""" # All-reduce first allreduce_out = tensor_model_parallel_all_reduce(input_tensor) # Apply native RMSNorm if residual is not None: norm_out, residual_out = rmsnorm_layer.forward_native(allreduce_out, residual) quant_input = norm_out else: norm_out = rmsnorm_layer.forward_native(allreduce_out) quant_input = norm_out residual_out = allreduce_out # Apply FP4 quantization (still using fused CUDA op as there's no native FP4) if SGL_SCALED_FP4_QUANT is None: raise RuntimeError("scaled_fp4_quant is not available on this platform") quant_res, output_scale_res = SGL_SCALED_FP4_QUANT(quant_input, input_global_scale) if residual is not None: return quant_res, residual_out, output_scale_res else: return quant_res, norm_out # Compiled versions of native functions @torch.compile def standard_allreduce_rmsnorm_native_compiled( input_tensor: torch.Tensor, residual: Optional[torch.Tensor], rmsnorm_layer: RMSNorm, norm_out: Optional[torch.Tensor] = None, ): """Compiled version of standard allreduce + rmsnorm.""" return standard_allreduce_rmsnorm_native( input_tensor, residual, rmsnorm_layer, norm_out ) @torch.compile def standard_allreduce_rmsnorm_fp8_quant_native_compiled( input_tensor: torch.Tensor, residual: Optional[torch.Tensor], rmsnorm_layer: RMSNorm, scale_factor: torch.Tensor, norm_out: Optional[torch.Tensor] = None, quant_out: Optional[torch.Tensor] = None, ): """Compiled version of standard allreduce + rmsnorm + FP8 quantization.""" return standard_allreduce_rmsnorm_fp8_quant_native( input_tensor, residual, rmsnorm_layer, scale_factor, norm_out, quant_out, ) @torch.compile def standard_allreduce_rmsnorm_fp4_quant_native_compiled( input_tensor: torch.Tensor, residual: Optional[torch.Tensor], rmsnorm_layer: RMSNorm, input_global_scale: torch.Tensor, quant_out: torch.Tensor, output_scale: torch.Tensor, norm_out: Optional[torch.Tensor] = None, ): """Compiled version of standard allreduce + rmsnorm + FP4 quantization.""" return standard_allreduce_rmsnorm_fp4_quant_native( input_tensor, residual, rmsnorm_layer, input_global_scale, quant_out, output_scale, norm_out, ) def create_test_tensors( seq_len: int, hidden_dim: int, dtype: torch.dtype, use_residual: bool = True ): """Create test tensors for benchmarking.""" input_tensor = torch.randn(seq_len, hidden_dim, dtype=dtype) residual = ( torch.randn_like(input_tensor) if use_residual else torch.zeros_like(input_tensor) ) rms_gamma = torch.ones(hidden_dim, dtype=dtype) norm_out = None if use_residual else torch.empty_like(input_tensor) # Quantization scales scale_fp8 = torch.tensor(1.0, dtype=torch.float32) scale_fp4 = torch.tensor(1.0, dtype=torch.float32) quant_out_fp8 = torch.empty_like(input_tensor, dtype=FP8_DTYPE) # Pre-allocate FP4 output tensors (to avoid allocation overhead in benchmarks) fp4_quant_out = torch.empty((seq_len, hidden_dim // 2), dtype=torch.uint8) fp4_output_scale = torch.empty((128, 4), dtype=torch.int32) return ( input_tensor, norm_out, residual, rms_gamma, scale_fp8, quant_out_fp8, scale_fp4, fp4_quant_out, fp4_output_scale, ) def benchmark_operation( operation_func, *args, warmup: int = 5, trials: int = 20, **kwargs ): """Benchmark a single operation using CUDA graphs.""" # Warmup before graph capture for _ in range(warmup): operation_func(*args, **kwargs) torch.cuda.synchronize() # Create CUDA graph graph = torch.cuda.CUDAGraph() num_op_per_cudagraph = 10 # Use sglang's graph_capture to make tensor_model_parallel_all_reduce graph-safe with graph_capture() as graph_capture_context: with torch.cuda.graph(graph, stream=graph_capture_context.stream): for _ in range(num_op_per_cudagraph): operation_func(*args, **kwargs) # Graph warmup torch.cuda.synchronize() for _ in range(warmup): graph.replay() # Benchmark with CUDA graph torch.cuda.synchronize() start_time = time.perf_counter() for _ in range(trials // num_op_per_cudagraph): # operation_func(*args, **kwargs) graph.replay() torch.cuda.synchronize() end_time = time.perf_counter() avg_time_ms = ((end_time - start_time) / trials) * 1000 return avg_time_ms def run_benchmarks( seq_len: int, hidden_dim: int, dtype: torch.dtype, use_residual: bool, allreduce_params: Optional[FlashInferFusedAllReduceParams], quant_mode: str = "all", disable_oneshot: bool = False, ): """Run all benchmarks for given configuration. Args: quant_mode: "none", "fp8_only", "fp4_only", or "all" """ ( input_tensor, norm_out, residual, rms_gamma, scale_fp8, quant_out_fp8, scale_fp4, fp4_quant_out, fp4_output_scale, ) = create_test_tensors(seq_len, hidden_dim, dtype, use_residual) rms_eps = 1e-6 results = {} # Create RMSNorm once for native benchmarks rmsnorm_layer = RMSNorm(hidden_dim, eps=rms_eps) rmsnorm_layer.weight.data = rms_gamma if quant_mode in ["all", "none"]: # Standard AllReduce + RMSNorm try: time_ms = benchmark_operation( standard_allreduce_rmsnorm, input_tensor, norm_out=norm_out, residual=residual, rms_gamma=rms_gamma, rms_eps=rms_eps, ) results["standard_allreduce_rmsnorm"] = time_ms except Exception as e: logger.error("Standard AllReduce+RMSNorm failed: %s", e) results["standard_allreduce_rmsnorm"] = float("inf") # Standard AllReduce + RMSNorm Native Compiled try: time_ms = benchmark_operation( standard_allreduce_rmsnorm_native_compiled, input_tensor, residual=residual, rmsnorm_layer=rmsnorm_layer, norm_out=norm_out, ) results["standard_allreduce_rmsnorm_native_compiled"] = time_ms except Exception as e: logger.error("Standard AllReduce+RMSNorm Native Compiled failed: %s", e) results["standard_allreduce_rmsnorm_native_compiled"] = float("inf") # FlashInfer Fused AllReduce + RMSNorm Oneshot if flashinfer_comm is not None and allreduce_params is not None: try: if not disable_oneshot: time_ms = benchmark_operation( flashinfer_fused_allreduce_rmsnorm, input_tensor, residual=residual, norm_out=norm_out, rms_gamma=rms_gamma, rms_eps=rms_eps, allreduce_params=allreduce_params, use_oneshot=True, ) results["flashinfer_fused_allreduce_rmsnorm_oneshot"] = time_ms except Exception as e: logger.error("FlashInfer Fused AllReduce+RMSNorm Oneshot failed: %s", e) results["flashinfer_fused_allreduce_rmsnorm_oneshot"] = float("inf") # FlashInfer Fused AllReduce + RMSNorm Two-shot try: time_ms = benchmark_operation( flashinfer_fused_allreduce_rmsnorm, input_tensor, residual=residual, norm_out=norm_out, rms_gamma=rms_gamma, rms_eps=rms_eps, allreduce_params=allreduce_params, use_oneshot=False, ) results["flashinfer_fused_allreduce_rmsnorm_twoshot"] = time_ms except Exception as e: logger.error( "FlashInfer Fused AllReduce+RMSNorm Two-shot failed: %s", e ) results["flashinfer_fused_allreduce_rmsnorm_twoshot"] = float("inf") if quant_mode in ["all", "fp8_only"]: # Standard AllReduce + RMSNorm + FP8 Quant try: time_ms = benchmark_operation( standard_allreduce_rmsnorm_fp8_quant, input_tensor, norm_out=norm_out, residual=residual, rms_gamma=rms_gamma, rms_eps=rms_eps, scale_factor=scale_fp8, quant_out=quant_out_fp8, ) results["standard_allreduce_rmsnorm_fp8_quant"] = time_ms except Exception as e: logger.error("Standard AllReduce+RMSNorm+FP8 failed: %s", e) results["standard_allreduce_rmsnorm_fp8_quant"] = float("inf") # Standard AllReduce + RMSNorm + FP8 Quant Native Compiled try: time_ms = benchmark_operation( standard_allreduce_rmsnorm_fp8_quant_native_compiled, input_tensor, residual=residual, rmsnorm_layer=rmsnorm_layer, # quant_fp8_layer removed in sglang version; static_quant_fp8 is used within the function scale_factor=scale_fp8, norm_out=norm_out, quant_out=quant_out_fp8, ) results["standard_allreduce_rmsnorm_fp8_quant_native_compiled"] = time_ms except Exception as e: logger.error("Standard AllReduce+RMSNorm+FP8 Native Compiled failed: %s", e) results["standard_allreduce_rmsnorm_fp8_quant_native_compiled"] = float( "inf" ) # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant Oneshot if flashinfer_comm is not None and allreduce_params is not None: try: if not disable_oneshot: time_ms = benchmark_operation( flashinfer_fused_allreduce_rmsnorm_fp8_quant, input_tensor, norm_out=norm_out, residual=residual, rms_gamma=rms_gamma, rms_eps=rms_eps, scale_factor=scale_fp8, quant_out=quant_out_fp8, allreduce_params=allreduce_params, use_oneshot=True, ) results["flashinfer_fused_allreduce_rmsnorm_fp8_quant_oneshot"] = ( time_ms ) except Exception as e: logger.error( "FlashInfer Fused AllReduce+RMSNorm+FP8 Oneshot failed: %s", e, ) results["flashinfer_fused_allreduce_rmsnorm_fp8_quant_oneshot"] = float( "inf" ) # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant Two-shot try: time_ms = benchmark_operation( flashinfer_fused_allreduce_rmsnorm_fp8_quant, input_tensor, norm_out=norm_out, residual=residual, rms_gamma=rms_gamma, rms_eps=rms_eps, scale_factor=scale_fp8, quant_out=quant_out_fp8, allreduce_params=allreduce_params, use_oneshot=False, ) results["flashinfer_fused_allreduce_rmsnorm_fp8_quant_twoshot"] = ( time_ms ) except Exception as e: logger.error( "FlashInfer Fused AllReduce+RMSNorm+FP8 Two-shot failed: %s", e, ) results["flashinfer_fused_allreduce_rmsnorm_fp8_quant_twoshot"] = float( "inf" ) if quant_mode in ["all", "fp4_only"]: # Standard AllReduce + RMSNorm + FP4 Quant try: time_ms = benchmark_operation( standard_allreduce_rmsnorm_fp4_quant, input_tensor, norm_out=norm_out, residual=residual, rms_gamma=rms_gamma, rms_eps=rms_eps, input_global_scale=scale_fp4, quant_out=fp4_quant_out, output_scale=fp4_output_scale, ) results["standard_allreduce_rmsnorm_fp4_quant"] = time_ms except Exception as e: logger.error("Standard AllReduce+RMSNorm+FP4 failed: %s", e) results["standard_allreduce_rmsnorm_fp4_quant"] = float("inf") # Standard AllReduce + RMSNorm + FP4 Quant Native Compiled try: time_ms = benchmark_operation( standard_allreduce_rmsnorm_fp4_quant_native_compiled, input_tensor, residual=residual, rmsnorm_layer=rmsnorm_layer, input_global_scale=scale_fp4, quant_out=fp4_quant_out, output_scale=fp4_output_scale, norm_out=norm_out, ) results["standard_allreduce_rmsnorm_fp4_quant_native_compiled"] = time_ms except Exception as e: logger.error("Standard AllReduce+RMSNorm+FP4 Native Compiled failed: %s", e) results["standard_allreduce_rmsnorm_fp4_quant_native_compiled"] = float( "inf" ) # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Oneshot if flashinfer_comm is not None and allreduce_params is not None: try: if not disable_oneshot: time_ms = benchmark_operation( flashinfer_fused_allreduce_rmsnorm_fp4_quant, input_tensor, residual=residual, norm_out=norm_out, rms_gamma=rms_gamma, rms_eps=rms_eps, input_global_scale=scale_fp4, allreduce_params=allreduce_params, quant_out=fp4_quant_out, output_scale=fp4_output_scale, use_oneshot=True, ) results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_oneshot"] = ( time_ms ) except Exception as e: logger.error( "FlashInfer Fused AllReduce+RMSNorm+FP4 Oneshot failed: %s", e, ) results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_oneshot"] = float( "inf" ) # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Two-shot if flashinfer_comm is not None and allreduce_params is not None: try: time_ms = benchmark_operation( flashinfer_fused_allreduce_rmsnorm_fp4_quant, input_tensor, residual=residual, norm_out=norm_out, rms_gamma=rms_gamma, rms_eps=rms_eps, input_global_scale=scale_fp4, allreduce_params=allreduce_params, quant_out=fp4_quant_out, output_scale=fp4_output_scale, use_oneshot=False, ) results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = ( time_ms ) except Exception as e: logger.error( "FlashInfer Fused AllReduce+RMSNorm+FP4 Two-shot failed: %s", e, ) results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = float( "inf" ) return results def prepare_results_with_speedups(results_dict): """Prepare results with speedup calculations based on dynamic baseline selection.""" prepared_results = [] # Determine the fastest baseline for each operation type def get_fastest_baseline(op_name, results_dict): """Get the fastest baseline between standard and native_compiled versions.""" if "fp8_quant" in op_name: candidates = [ "standard_allreduce_rmsnorm_fp8_quant", "standard_allreduce_rmsnorm_fp8_quant_native_compiled", ] elif "fp4_quant" in op_name: candidates = [ "standard_allreduce_rmsnorm_fp4_quant", "standard_allreduce_rmsnorm_fp4_quant_native_compiled", ] else: candidates = [ "standard_allreduce_rmsnorm", "standard_allreduce_rmsnorm_native_compiled", ] # Find the fastest among available candidates fastest_time = float("inf") fastest_baseline = None for candidate in candidates: if ( candidate in results_dict and results_dict[candidate] != float("inf") and results_dict[candidate] < fastest_time ): fastest_time = results_dict[candidate] fastest_baseline = candidate return fastest_baseline # Create dynamic baseline mapping dynamic_baseline_mapping = {} for op_name in results_dict: if ( op_name.startswith("flashinfer_") or op_name.startswith("standard_") and not op_name.endswith("_native_compiled") ): dynamic_baseline_mapping[op_name] = get_fastest_baseline( op_name, results_dict ) for op_name, time_ms in results_dict.items(): if time_ms == float("inf"): speedup_str = "FAILED" time_str = "FAILED" else: time_str = f"{time_ms:.3f}" # Find the appropriate baseline for this operation baseline_op = dynamic_baseline_mapping.get(op_name) if baseline_op and baseline_op in results_dict: baseline_time = results_dict[baseline_op] if baseline_time != float("inf") and baseline_time > 0: speedup = baseline_time / time_ms speedup_str = f"{speedup:.2f}x" else: speedup_str = "N/A" else: # For baseline operations, determine if this is the fastest baseline if op_name.endswith("_native_compiled") or ( op_name.startswith("standard_") and not op_name.endswith("_native_compiled") ): fastest_baseline = get_fastest_baseline(op_name, results_dict) if fastest_baseline == op_name: speedup_str = "baseline" else: if fastest_baseline and fastest_baseline in results_dict: baseline_time = results_dict[fastest_baseline] if baseline_time != float("inf") and baseline_time > 0: speedup = baseline_time / time_ms speedup_str = f"{speedup:.2f}x" else: speedup_str = "N/A" else: speedup_str = "N/A" else: speedup_str = "N/A" prepared_results.append( { "operation": op_name, "time_ms": time_ms, "time_str": time_str, "speedup_str": speedup_str, } ) return prepared_results def print_results(results_dict, seq_len, hidden_dim, dtype, use_residual, quant_mode): """Print benchmark results in a formatted table.""" print(f"\n{'=' * 80}") print(f"Results: seq_len={seq_len}, hidden_dim={hidden_dim}") print( f"dtype={dtype}, residual={'yes' if use_residual else 'no'}, " f"quant_mode={quant_mode}" ) print(f"{'=' * 80}") print(f"{'Operation':<50} {'Time (ms)':<12} {'Speedup':<10}") print(f"{'-' * 80}") # Prepare results with speedup calculations prepared_results = prepare_results_with_speedups(results_dict) for result in prepared_results: if result["time_ms"] == float("inf"): time_display = result["time_str"] else: time_display = f"{result['time_ms']:.3f}" print( f"{result['operation']:<50} {time_display:<12} {result['speedup_str']:<10}" ) def format_results_markdown( all_results: list[dict], world_size: int, args: argparse.Namespace ) -> str: """Format all benchmark results as markdown.""" markdown = f"""# FlashInfer Fused Collective Operations Benchmark Results **World Size:** {world_size} **Hidden Dimension:** {args.hidden_dim} **Warmup Iterations:** {args.warmup} **Benchmark Trials:** {args.trials} **Quantization Mode:** {all_results[0]["quant_mode"] if all_results else "N/A"} --- """ for result in all_results: seq_len = result["seq_len"] dtype = result["dtype"] use_residual = result["use_residual"] results_dict = result["results"] residual_str = "with residual" if use_residual else "no residual" markdown += f""" ## Configuration: seq_len={seq_len}, dtype={dtype}, {residual_str} | Operation | Time (ms) | Speedup | |-----------|-----------|---------| """ # Prepare results with speedup calculations prepared_results = prepare_results_with_speedups(results_dict) for result in prepared_results: # Format operation name for better readability formatted_op_name = result["operation"].replace("_", " ").title() markdown += f"| {formatted_op_name} | {result['time_str']} |" markdown += f"{result['speedup_str']} |\n" markdown += "\n" return markdown def save_results_to_file( all_results: list[dict], world_size: int, args: argparse.Namespace, rank: int ): """Save benchmark results to markdown file (only on rank 0).""" if rank != 0: return if not all_results: logger.warning("No results to save") return output_path = args.output_file try: markdown_content = format_results_markdown(all_results, world_size, args) with open(output_path, "w") as f: f.write(markdown_content) except Exception as e: logger.error("Failed to save results to file: %s", e) def main(): parser = argparse.ArgumentParser( description="Benchmark fused collective operations" ) parser.add_argument( "--seq-lens", type=int, nargs="+", default=[128, 512, 1024, 2048], help="Sequence lengths to test", ) parser.add_argument( "--hidden-dim", type=int, default=8192, help="Hidden dimension size" ) parser.add_argument( "--dtypes", type=str, nargs="+", default=["bfloat16"], choices=["float16", "bfloat16", "float32"], help="Data types to test", ) parser.add_argument( "--no-residual", action="store_true", help="Skip residual connection tests", ) # Quantization mode options (mutually exclusive with --no-quant) quant_group = parser.add_mutually_exclusive_group() quant_group.add_argument( "--no-quant", action="store_true", help="Skip all quantization tests" ) quant_group.add_argument( "--quant-fp8", action="store_true", help="Only run FP8 quantization tests" ) quant_group.add_argument( "--quant-fp4", action="store_true", help="Only run FP4 quantization tests" ) quant_group.add_argument( "--quant-all", action="store_true", help="Run all quantization tests (default)", ) parser.add_argument( "--disable-oneshot", action="store_true", help="Disable oneshot mode for FlashInfer operations", ) parser.add_argument( "--warmup", type=int, default=5, help="Number of warmup iterations" ) parser.add_argument( "--trials", type=int, default=20, help="Number of benchmark trials" ) parser.add_argument( "--output-file", type=str, help="""Output file path for markdown results (default: benchmark_results_.md) """, ) args = parser.parse_args() # Check if running with torchrun (required for collective operations) if "RANK" not in os.environ or "WORLD_SIZE" not in os.environ: raise RuntimeError( "Must run with torchrun for distributed benchmarking. " "Example: torchrun --nproc_per_node=2 benchmark_fused_collective.py" ) # Initialize distributed environment rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) device = torch.device(f"cuda:{rank}") torch.cuda.set_device(device) torch.set_default_device(device) init_distributed_environment( world_size=world_size, rank=rank, local_rank=rank, backend="nccl", ) initialize_model_parallel(tensor_model_parallel_size=world_size) # Validate world size (must be > 1 for collective operations) if world_size <= 1: raise ValueError( "World size must be > 1 for collective operations benchmarking. " f"Current world size: {world_size}. Use torchrun with --nproc_per_node > 1." ) # Determine quantization mode if args.no_quant: quant_mode = "none" elif args.quant_fp8: quant_mode = "fp8_only" elif args.quant_fp4: quant_mode = "fp4_only" else: # args.quant_all or default quant_mode = "all" if rank == 0: logger.info("Running benchmark with world_size=%s, rank=%s", world_size, rank) logger.info("Quantization mode: %s", quant_mode) if flashinfer_comm is not None: oneshot_status = "enabled" if not args.disable_oneshot else "disabled" logger.info( "FlashInfer available - will benchmark fused operations (oneshot: %s)", oneshot_status, ) else: logger.info( "FlashInfer not available - only benchmarking standard operations" ) # Convert dtype strings to torch dtypes dtype_map = { "float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32, } dtypes = [dtype_map[dt] for dt in args.dtypes] # Test configurations residual_options = [True] if not args.no_residual else [False] if not args.no_residual: residual_options.append(False) configs = list(itertools.product(args.seq_lens, dtypes, residual_options)) # Setup FlashInfer workspace if available ipc_handles = None allreduce_params = None if flashinfer_comm is not None: # Use the largest hidden dimension for workspace setup max_num_token = _FI_MAX_SIZES.get(world_size) // ( args.hidden_dim * world_size * 2 ) ipc_handles, workspace_tensor = setup_flashinfer_workspace( world_size, rank, args.hidden_dim, max_num_token ) if workspace_tensor is not None: allreduce_params = FlashInferFusedAllReduceParams( rank=rank, world_size=world_size, max_token_num=max_num_token, ) # Collect all results for markdown export all_results = [] try: # Run benchmarks for seq_len, dtype, use_residual in configs: if rank == 0: logger.info( "\nTesting: seq_len=%s, hidden_dim=%s, dtype=%s, residual=%s", seq_len, args.hidden_dim, dtype, use_residual, ) results = run_benchmarks( seq_len, args.hidden_dim, dtype, use_residual, allreduce_params, quant_mode=quant_mode, disable_oneshot=args.disable_oneshot, ) # Store results for markdown export if rank == 0: all_results.append( { "seq_len": seq_len, "hidden_dim": args.hidden_dim, "dtype": str(dtype).replace("torch.", ""), "use_residual": use_residual, "quant_mode": quant_mode, "results": results, } ) print_results( results, seq_len, args.hidden_dim, dtype, use_residual, quant_mode, ) # Save results to markdown file if args.output_file and rank == 0: save_results_to_file(all_results, world_size, args, rank) finally: # Cleanup if ipc_handles is not None: cleanup_flashinfer_workspace(ipc_handles) with contextlib.suppress(Exception): dist.barrier() cleanup_dist_env_and_memory(shutdown_ray=False) if __name__ == "__main__": main() ================================================ FILE: benchmark/kernels/fused_moe_triton/README.md ================================================ ## Tuning Triton MoE Kernels This directory contains benchmarking tools for MoE (Mixture of Experts) kernels. ### Overview The tuning tools support both **Tensor Parallelism (TP)** and **Expert Parallelism (EP)** modes: - **TP Mode**: Traditional tensor parallelism where intermediate layers are sharded across GPUs - **EP Mode**: Expert parallelism where experts are distributed across GPUs. Can be combined with TP mode (e.g., `--tp-size 8 --ep-size 2`) - **MLLM Support**: Multi-modal Large Language Models with text encoders (e.g., Llama4, Qwen3VL) ### Tuning Tools #### 1. `tuning_fused_moe_triton.py` A unified tool for tuning the `fused_moe_triton` kernel. Adapted from [vllm's benchmark_moe.py](https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/benchmark_moe.py), with support for EP mode and various model architectures. #### 2. `tuning_fused_moe_triton_sep.py` A specialized tool for separate kernel tuning, optimizing the first and second MoE kernels independently with TMA (Tensor Memory Accelerator) support. ### Usage Examples #### Basic TP Mode Tuning ```bash # Tune Mixtral-8x7B with default TP settings python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ --model mistralai/Mixtral-8x7B-Instruct-v0.1 \ --tune # Tune Qwen2-57B with FP8 and TP=4 python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ --model Qwen/Qwen2-57B-A14B-Instruct \ --tp-size 4 \ --dtype fp8_w8a8 \ --tune # Tune DeepSeek-V3 with FP8 and TP=8 python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ --model deepseek-ai/DeepSeek-V3-0324 \ --tp-size 8 \ --dtype fp8_w8a8 \ --tune ``` #### EP Mode Tuning (Expert Parallelism) **Note**: EP mode can be used alone or combined with TP mode. When using both, ensure `tp_size` is divisible by `ep_size`. ```bash # Tune Mixtral-8x7B with EP=2 only python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ --model mistralai/Mixtral-8x7B-Instruct-v0.1 \ --tp-size 2 \ --ep-size 2 \ --tune # Tune Qwen2-57B with TP=8 and EP=4 (combined mode) python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ --model Qwen/Qwen2-57B-A14B-Instruct \ --tp-size 8 \ --ep-size 4 \ --dtype fp8_w8a8 \ --tune ``` #### MLLM Model Tuning (Multi-modal) ```bash python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ --model Qwen/Qwen3-VL-30B-A3B-Instruct \ --tp-size 2 \ --tune ``` #### Separate Kernel Tuning with `tuning_fused_moe_triton_sep.py` This tool requires pre-generated topk_ids files and supports both TP and EP modes: Edit the code file (such as srt/models/deepseek_v2.py) in the Python site package and add the logic for saving topk_ids: ```python # import get_tensor_model_parallel_rank # DeepseekV2MoE::forward_normal if hidden_states.shape[0] >= 4096 and get_tensor_model_parallel_rank() == 0: topk_ids_dir = xxxx if not hasattr(self, "save_idx"): self.save_idx = 0 if self.save_idx <= 1: torch.save(topk_output.topk_ids, f"{topk_ids_dir}/topk_ids_layer{self.layer_id}_idx{self.save_idx}.pt") self.save_idx += 1 ``` Launch sglang server and send request using `benchmark/kernels/fused_moe_triton/tuning_client.py` ```bash python benchmark/kernels/fused_moe_triton/tuning_client.py --port 8000 ``` ```bash # TP Mode: Tune separate kernels with TP=4 python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py \ --model Qwen/Qwen2-57B-A14B-Instruct \ --tp-size 4 \ --topk-ids-dir /path/to/topk_ids \ --tune # EP Mode: Tune separate kernels with TP=4 and EP=2 python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py \ --model mistralai/Mixtral-8x7B-Instruct-v0.1 \ --tp-size 4 \ --ep-size 2 \ --topk-ids-dir /path/to/topk_ids \ --tune # MLLM: Tune DeepSeek-V3 with separate kernels, TP=8 and EP=4 python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py \ --model deepseek-ai/DeepSeek-V3-0324 \ --tp-size 8 \ --ep-size 4 \ --dtype fp8_w8a8 \ --topk-ids-dir /path/to/topk_ids \ --tune # Benchmark specific config without tuning python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py \ --model deepseek-ai/DeepSeek-V3-0324 \ --tp-size 4 \ --batch-size 1024 \ --dtype fp8_w8a8 \ --configs 128 256 128 16 8 4 \ --topk-ids-dir /path/to/topk_ids ``` #### Advanced Options ```bash # Channel-wise quantization python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ --model meituan/DeepSeek-R1-Channel-INT8 \ --tp-size 16 \ --dtype int8_w8a8 \ --per-channel-quant \ --tune # Specific batch size tuning python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ --model mistralai/Mixtral-8x7B-Instruct-v0.1 \ --batch-size 2048 \ --tune ``` ### Configuration Files After tuning, configuration files will be generated: - **Standard tuning**: `E=64,N=640,device_name=NVIDIA_GeForce_RTX_4090,dtype=fp8_w8a8.json` - **Separate kernel tuning**: Two files for up/down kernels with TMA optimization flags Move these files to `sglang/srt/layers/moe/fused_moe_triton/configs/triton_version/` directory to use them in SGLang. ### Supported Models - **Mixtral**: mistralai/Mixtral-8x7B-Instruct-v0.1, mixtral-8x22b - **Qwen**: Qwen2-57B, Qwen3-235B, Qwen3VL (MLLM) - **DeepSeek**: DeepSeek-V2, DeepSeek-V3, DeepSeek-R1 - **Llama**: Llama4-Vision (MLLM) - **DBRX**: databricks/dbrx-instruct - **Jamba**: ai21labs/AI21-Jamba - **Grok**: xai-org/grok-1 - **GLM**: THUDM/glm-4-9b-chat - **Bailing**: Custom MoE models ### Parameters Reference - `--model`: HuggingFace model name or local path - `--tp-size`: Tensor parallelism size (default: 2) - `--ep-size`: Expert parallelism size (default: 1, can be combined with TP mode, ensure tp_size is divisible by ep_size) - `--dtype`: Data type (`auto`, `fp8_w8a8`, `int8_w8a16`, `int8_w8a8`) - `--batch-size`: Specific batch size for tuning (optional) - `--tune`: Enable tuning mode - `--per-channel-quant`: Enable per-channel quantization - `--disable-shared-experts-fusion`: Disable shared expert fusion for some models - `--topk-ids-dir`: Directory containing pre-generated topk_ids (for sep tool only) - `--configs`: Manual config specification [BLOCK_M, BLOCK_N, BLOCK_K, GROUP_M, warps, stages] ### Performance Comparison Tool - `benchmark_vllm_vs_sglang_fused_moe_triton.py`: A tool for comparing the performance of fused MoE kernels between vllm and sglang implementations. Supports various model architectures and data types. Example usage: ```bash # Compare with default settings (Mixtral model) python benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py # Compare with FP8 mode for Qwen2-57B python benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py \ --model Qwen/Qwen2-57B-A14B-Instruct \ --use-fp8-w8a8 # Compare with custom TP size python benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py \ --model deepseek-ai/DeepSeek-V3-0324 \ --tp-size 8 # Compare with custom TP size python benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py \ --model deepseek-ai/DeepSeek-V3-0324 \ --tp-size 8 ``` The benchmark results will be saved as plots and data files in the specified output directory (default: `./configs/benchmark_ops/vllm_sglang_fused_moe/`). - `benchmark_torch_compile_fused_moe.py`: A tool for benchmarking the performance of the fused MoE kernel with `torch.compile` and original fused MoE kernel. Usage is similar to `benchmark_vllm_vs_sglang_fused_moe_triton.py`, note that `torch.compile` does not support `fp8_w8a8` and `int8_w8a8` fused_moe_kernel. Both tools now support EP mode with `--ep-size` parameter. ================================================ FILE: benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py ================================================ # python3 benchmark/kernels/fused_moe_triton/sglang_fused_moe_triton.py --model /DeepSeek-V3/ --tp-size 8 import argparse import torch import triton from common_utils import get_model_config from sglang.benchmark.bench_utils import run_bench from sglang.srt.distributed.parallel_state import ( destroy_distributed_environment, destroy_model_parallel, init_distributed_environment, initialize_model_parallel, ) from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( fused_moe as fused_moe_sglang, ) from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import ( triton_kernel_moe_forward, ) from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig from sglang.srt.layers.moe.topk import ( TopK, TopKConfig, TopKOutputFormat, select_experts, ) from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler def fused_moe_triton_api( x, w1, w2, input_gating, topk, ): topk_op = TopK( top_k=topk, renormalize=False, use_grouped_topk=False, output_format=TopKOutputFormat.TRITON_KERNEL, ) triton_topk_output = topk_op.forward_cuda( hidden_states=x, router_logits=input_gating, ) moe_runner_config = MoeRunnerConfig( inplace=False, ) return triton_kernel_moe_forward( x, w1, w2, triton_topk_output, moe_runner_config, ) def fused_moe_sglang_api( x, w1, w2, input_gating, topk, use_fp8_w8a8=False, w1_scale=None, w2_scale=None, a1_scale=None, a2_scale=None, block_shape=None, ): topk_output = select_experts( hidden_states=x, router_logits=input_gating, topk_config=TopKConfig(top_k=topk, renormalize=False), ) return fused_moe_sglang( x, w1, w2, topk_output, use_fp8_w8a8=use_fp8_w8a8, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, block_shape=block_shape, ) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["batch_size"], x_vals=list([128, 256, 512, 1024, 2048, 4096, 8192]), line_arg="provider", line_vals=[ "sglang_fused_moe_triton_v340", "sglang_fused_moe_triton", ], line_names=[ "sglang_fused_moe_triton_v340", "sglang_fused_moe_triton", ], styles=[ ("blue", "-"), ("green", "-"), ], ylabel="Time (ms)", plot_name="fused-moe-performance", args={}, ) ) def benchmark( batch_size, provider, model_config, use_fp8_w8a8=False, use_cuda_graph: bool = False, ): print(f"benchmark {provider} with batch_size={batch_size}") torch.set_default_device("cuda") torch.cuda.manual_seed_all(0) num_tokens = batch_size num_experts = model_config["num_experts"] hidden_size = model_config["hidden_size"] shard_intermediate_size = model_config["shard_intermediate_size"] topk = model_config["topk"] dtype = model_config["dtype"] block_shape = model_config["block_shape"] x = torch.randn(num_tokens, hidden_size, dtype=dtype) w1 = torch.randn(num_experts, shard_intermediate_size, hidden_size, dtype=dtype) w2 = torch.randn( num_experts, hidden_size, shard_intermediate_size // 2, dtype=dtype ) w1_tri = w1.clone() w2_tri = w2.clone() w1_tri = w1_tri.transpose(-2, -1).contiguous() w2_tri = w2_tri.transpose(-2, -1).contiguous() input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32) if provider == "sglang_fused_moe_triton_v340": api_func = fused_moe_triton_api api_kwargs = { "x": x, "w1": w1_tri, "w2": w2_tri, "input_gating": input_gating, "topk": topk, } else: api_func = fused_moe_sglang_api api_kwargs = { "x": x, "w1": w1, "w2": w2, "input_gating": input_gating, "topk": topk, "use_fp8_w8a8": use_fp8_w8a8, "block_shape": block_shape, } # Warmup for _ in range(10): _ = api_func(**api_kwargs) torch.cuda.synchronize() if use_cuda_graph: stream = torch.cuda.Stream() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=stream): api_func(**api_kwargs) torch.cuda.synchronize() bench_lambda = lambda: graph.replay() else: bench_lambda = lambda: api_func(**api_kwargs) quantiles = (0.5, 0.2, 0.8) ms, min_ms, max_ms = run_bench(bench_lambda, quantiles=quantiles) return ms, min_ms, max_ms def main(): parser = argparse.ArgumentParser() parser.add_argument( "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1" ) parser.add_argument("--tp-size", "--tp", type=int, default=2) parser.add_argument("--ep-size", "--ep", type=int, default=1) parser.add_argument("--use-fp8-w8a8", action="store_true") parser.add_argument( "--use-cuda-graph", action="store_true", help="Enable CUDA Graph capture/replay" ) parser.add_argument( "--save-path", type=str, default="./configs/benchmark_ops/sglang_fused_moe/", ) parser.add_argument("--trust-remote-code", action="store_true") args = parser.parse_args() # Initialize global server args (required by SGLang MoE kernels) server_args = ServerArgs(model_path=args.model) set_global_server_args_for_scheduler(server_args) try: if not torch.distributed.is_initialized(): torch.distributed.init_process_group( backend="nccl" if torch.cuda.is_available() else "gloo", init_method="tcp://127.0.0.1:23456", world_size=1, rank=0, ) init_distributed_environment( world_size=1, rank=0, distributed_init_method="tcp://127.0.0.1:23456", local_rank=0, backend="nccl" if torch.cuda.is_available() else "gloo", ) initialize_model_parallel( tensor_model_parallel_size=1, expert_model_parallel_size=1, ) model_config = get_model_config(args.model, args.tp_size, args.ep_size) benchmark.run( show_plots=True, print_data=True, save_path=args.save_path, model_config=model_config, use_fp8_w8a8=args.use_fp8_w8a8, use_cuda_graph=args.use_cuda_graph, ) finally: destroy_model_parallel() destroy_distributed_environment() if __name__ == "__main__": main() ================================================ FILE: benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py ================================================ # python3 benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py --model /DeepSeek-V3/ --tp-size 8 --use-fp8-w8a8 import argparse import torch import triton from torch.nn import functional as F from transformers import AutoConfig from sglang.benchmark.bench_utils import run_bench from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( fused_moe as fused_moe_triton, ) from sglang.srt.model_executor.cuda_graph_runner import set_torch_compile_config def get_model_config(model_name: str, tp_size: int): """Get model configuration parameters""" config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) if config.architectures[0] == "DbrxForCausalLM": E = config.ffn_config.moe_num_experts topk = config.ffn_config.moe_top_k intermediate_size = config.ffn_config.ffn_hidden_size shard_intermediate_size = 2 * intermediate_size // tp_size elif config.architectures[0] == "JambaForCausalLM": E = config.num_experts topk = config.num_experts_per_tok intermediate_size = config.intermediate_size shard_intermediate_size = 2 * intermediate_size // tp_size elif config.architectures[0] == "Qwen2MoeForCausalLM": E = config.num_experts topk = config.num_experts_per_tok intermediate_size = config.moe_intermediate_size shard_intermediate_size = 2 * intermediate_size // tp_size elif config.architectures[0] == "Qwen3MoeForCausalLM": E = config.n_routed_experts topk = config.num_experts_per_tok intermediate_size = config.moe_intermediate_size shard_intermediate_size = 2 * intermediate_size // tp_size elif config.architectures[0] in ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]: E = config.n_routed_experts topk = config.num_experts_per_tok intermediate_size = config.moe_intermediate_size shard_intermediate_size = 2 * intermediate_size // tp_size elif config.architectures[0] == "Llama4ForConditionalGeneration": E = config.text_config.num_local_experts topk = config.text_config.num_experts_per_tok intermediate_size = config.text_config.intermediate_size shard_intermediate_size = 2 * intermediate_size // tp_size elif config.architectures[0] in [ "Grok1ForCausalLM", "Grok1ImgGen", "Grok1AForCausalLM", ]: E = config.num_local_experts topk = config.num_experts_per_tok intermediate_size = config.moe_intermediate_size shard_intermediate_size = 2 * intermediate_size // tp_size else: # Default: Mixtral E = config.num_local_experts topk = config.num_experts_per_tok intermediate_size = config.intermediate_size shard_intermediate_size = 2 * intermediate_size // tp_size shape_configs = { "num_experts": E, "topk": topk, "hidden_size": config.hidden_size, "shard_intermediate_size": shard_intermediate_size, "dtype": config.torch_dtype, } print(f"{shape_configs=}") return shape_configs def fused_topk_native( hidden_states: torch.Tensor, gating_output: torch.Tensor, topk: int, renormalize: bool, ): assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" M, _ = hidden_states.shape topk_weights = torch.empty( M, topk, dtype=torch.float32, device=hidden_states.device ) topk_ids = torch.empty(M, topk, dtype=torch.int32, device=hidden_states.device) topk_weights = F.softmax(gating_output.float(), dim=-1) topk_weights, topk_ids = torch.topk(topk_weights, topk, dim=-1) if renormalize: topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) return topk_weights, topk_ids @torch.compile(dynamic=False) def fused_moe_torch( x, w1, w2, input_gating, topk, use_fp8_w8a8=False, w1_scale=None, w2_scale=None, a1_scale=None, a2_scale=None, ) -> torch.Tensor: assert not use_fp8_w8a8, "Fp8_w8a8 fused_moe is not supported for torch compile" topk_weights, topk_ids = fused_topk_native( hidden_states=x, gating_output=input_gating, topk=topk, renormalize=True, ) w13_weights = w1[topk_ids] w1_weights, w3_weights = torch.chunk(w13_weights, 2, dim=2) w2_weights = w2[topk_ids] x1 = torch.einsum("ti,taoi -> tao", x, w1_weights) x1 = F.silu(x1) x3 = torch.einsum("ti, taoi -> tao", x, w3_weights) expert_outs = torch.einsum("tao, taio -> tai", (x1 * x3), w2_weights) return torch.einsum("tai,ta -> ti", expert_outs, topk_weights.to(expert_outs.dtype)) def fused_moe_torch_compile( x, w1, w2, input_gating, topk, use_fp8_w8a8=False, w1_scale=None, w2_scale=None, a1_scale=None, a2_scale=None, ): return fused_moe_torch( x, w1, w2, input_gating, topk, use_fp8_w8a8=use_fp8_w8a8, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, ) def fused_moe_sglang_api( x, w1, w2, input_gating, topk, use_fp8_w8a8=False, w1_scale=None, w2_scale=None, a1_scale=None, a2_scale=None, ): return fused_moe_triton( x, w1, w2, input_gating, topk, renormalize=True, inplace=True, use_fp8_w8a8=use_fp8_w8a8, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, ) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["batch_size"], x_vals=list(range(1, 5)), line_arg="provider", line_vals=[ "fused_moe_triton", "fused_moe_torch_compile", ], line_names=[ "fused_moe_triton", "fused_moe_torch_compile", ], styles=[ ("blue", "-"), ("green", "-"), ], ylabel="Time (ms)", plot_name="fused-moe-performance", args={}, ) ) def benchmark(batch_size, provider, model_config, use_fp8_w8a8=False): print(f"benchmark {provider} with batch_size={batch_size}") torch.set_default_device("cuda") torch.cuda.manual_seed_all(0) set_torch_compile_config() num_tokens = batch_size num_experts = model_config["num_experts"] hidden_size = model_config["hidden_size"] shard_intermediate_size = model_config["shard_intermediate_size"] topk = model_config["topk"] dtype = model_config["dtype"] x = torch.randn(num_tokens, hidden_size, dtype=dtype) if use_fp8_w8a8: init_dtype = dtype w1 = torch.randn( num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype ) w2 = torch.randn( num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype ) w1 = w1.to(torch.float8_e4m3fn) w2 = w2.to(torch.float8_e4m3fn) w1_scale = torch.randn(num_experts, dtype=torch.float32) w2_scale = torch.randn(num_experts, dtype=torch.float32) a1_scale = torch.randn(1, dtype=torch.float32) a2_scale = torch.randn(1, dtype=torch.float32) else: w1 = torch.randn(num_experts, shard_intermediate_size, hidden_size, dtype=dtype) w2 = torch.randn( num_experts, hidden_size, shard_intermediate_size // 2, dtype=dtype ) w1_scale = w2_scale = a1_scale = a2_scale = None input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32) # Warmup api_func = ( fused_moe_torch_compile if provider == "fused_moe_torch_compile" else fused_moe_sglang_api ) for _ in range(10): y = api_func( x, w1, w2, input_gating, topk, use_fp8_w8a8=use_fp8_w8a8, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, ) torch.cuda.synchronize() quantiles = (0.5, 0.2, 0.8) ms, min_ms, max_ms = run_bench( lambda: api_func( x, w1, w2, input_gating, topk, use_fp8_w8a8=use_fp8_w8a8, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, )[0], quantiles=quantiles, ) return ms, min_ms, max_ms def main(): parser = argparse.ArgumentParser() parser.add_argument( "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1" ) parser.add_argument("--tp-size", type=int, default=2) parser.add_argument("--use-fp8-w8a8", action="store_true") parser.add_argument( "--save-path", type=str, default="./configs/benchmark_ops/fused_moe_torch_compile/", ) args = parser.parse_args() model_config = get_model_config(args.model, args.tp_size) benchmark.run( show_plots=True, print_data=True, save_path=args.save_path, model_config=model_config, use_fp8_w8a8=args.use_fp8_w8a8, ) if __name__ == "__main__": main() ================================================ FILE: benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py ================================================ # python3 benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py --model /DeepSeek-V3/ --tp-size 8 --use-fp8-w8a8 import argparse import torch import triton from vllm.model_executor.layers.fused_moe.fused_moe import fused_moe as fused_moe_vllm from sglang.benchmark.bench_utils import run_bench from sglang.srt.distributed.parallel_state import ( destroy_distributed_environment, destroy_model_parallel, init_distributed_environment, initialize_model_parallel, ) from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( fused_moe as fused_moe_sglang, ) from .common_utils import get_model_config def fused_moe_vllm_api( x, w1, w2, input_gating, topk, use_fp8_w8a8=False, w1_scale=None, w2_scale=None, a1_scale=None, a2_scale=None, block_shape=None, ): if block_shape is not None: return fused_moe_vllm( x, w1, w2, input_gating, topk, renormalize=True, inplace=True, use_fp8_w8a8=use_fp8_w8a8, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, block_shape=block_shape, ) else: return fused_moe_vllm( x, w1, w2, input_gating, topk, renormalize=True, inplace=True, use_fp8_w8a8=use_fp8_w8a8, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, ) def fused_moe_sglang_api( x, w1, w2, input_gating, topk, use_fp8_w8a8=False, w1_scale=None, w2_scale=None, a1_scale=None, a2_scale=None, block_shape=None, ): return fused_moe_sglang( x, w1, w2, input_gating, topk, renormalize=True, inplace=True, use_fp8_w8a8=use_fp8_w8a8, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, block_shape=block_shape, ) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["batch_size"], x_vals=list(range(1, 513)), line_arg="provider", line_vals=[ "vllm_fused_moe_triton", "sglang_fused_moe_triton", ], line_names=[ "vllm_fused_moe_triton", "sglang_fused_moe_triton", ], styles=[ ("blue", "-"), ("green", "-"), ], ylabel="Time (ms)", plot_name="fused-moe-performance", args={}, ) ) def benchmark(batch_size, provider, model_config, use_fp8_w8a8=False): print(f"benchmark {provider} with batch_size={batch_size}") torch.set_default_device("cuda") torch.cuda.manual_seed_all(0) num_tokens = batch_size num_experts = model_config["num_experts"] hidden_size = model_config["hidden_size"] shard_intermediate_size = model_config["shard_intermediate_size"] topk = model_config["topk"] dtype = model_config["dtype"] block_shape = model_config["block_shape"] x = torch.randn(num_tokens, hidden_size, dtype=dtype) w1_scale = w2_scale = a1_scale = a2_scale = None if use_fp8_w8a8: init_dtype = dtype w1 = torch.randn( num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype ) w2 = torch.randn( num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype ) w1 = w1.to(torch.float8_e4m3fn) w2 = w2.to(torch.float8_e4m3fn) if block_shape is None: w1_scale = torch.randn(num_experts, dtype=torch.float32) w2_scale = torch.randn(num_experts, dtype=torch.float32) a1_scale = torch.randn(1, dtype=torch.float32) a2_scale = torch.randn(1, dtype=torch.float32) else: block_n, block_k = block_shape[0], block_shape[1] n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n n_tiles_w2 = (hidden_size + block_n - 1) // block_n k_tiles_w1 = (hidden_size + block_k - 1) // block_k k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k w1_scale = torch.rand( (num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.float32 ) w2_scale = torch.rand( (num_experts, n_tiles_w2, k_tiles_w2), dtype=torch.float32 ) else: w1 = torch.randn(num_experts, shard_intermediate_size, hidden_size, dtype=dtype) w2 = torch.randn( num_experts, hidden_size, shard_intermediate_size // 2, dtype=dtype ) input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32) # Warmup api_func = ( fused_moe_vllm_api if provider == "vllm_fused_moe_triton" else fused_moe_sglang_api ) for _ in range(10): y = api_func( x, w1, w2, input_gating, topk, use_fp8_w8a8=use_fp8_w8a8, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, block_shape=block_shape, ) torch.cuda.synchronize() quantiles = (0.5, 0.2, 0.8) ms, min_ms, max_ms = run_bench( lambda: api_func( x, w1, w2, input_gating, topk, use_fp8_w8a8=use_fp8_w8a8, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, block_shape=block_shape, )[0], quantiles=quantiles, ) return ms, min_ms, max_ms def main(): parser = argparse.ArgumentParser() parser.add_argument( "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1" ) parser.add_argument("--tp-size", "--tp", type=int, default=2) parser.add_argument("--ep-size", "--ep", type=int, default=1) parser.add_argument("--use-fp8-w8a8", action="store_true") parser.add_argument( "--save-path", type=str, default="./configs/benchmark_ops/vllm_sglang_fused_moe/", ) args = parser.parse_args() try: if not torch.distributed.is_initialized(): torch.distributed.init_process_group( backend="nccl" if torch.cuda.is_available() else "gloo", init_method="tcp://127.0.0.1:23456", world_size=1, rank=0, ) init_distributed_environment( world_size=1, rank=0, distributed_init_method="tcp://127.0.0.1:23456", local_rank=0, backend="nccl" if torch.cuda.is_available() else "gloo", ) initialize_model_parallel( tensor_model_parallel_size=1, pipeline_model_parallel_size=1, ) shape_configs = get_model_config(args.model, args.tp_size, args.ep_size) benchmark.run( show_plots=True, print_data=True, save_path=args.save_path, model_config=shape_configs, use_fp8_w8a8=args.use_fp8_w8a8, ) finally: destroy_model_parallel() destroy_distributed_environment() if __name__ == "__main__": main() ================================================ FILE: benchmark/kernels/fused_moe_triton/common_utils.py ================================================ import json from typing import Dict, List, TypedDict import torch from sglang.srt.layers.moe.fused_moe_triton.fused_moe import get_config_dtype_str from sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_config import ( get_config_file_name, ) from sglang.srt.utils import is_hip from sglang.srt.utils.hf_transformers_utils import get_config class BenchmarkConfig(TypedDict): BLOCK_SIZE_M: int BLOCK_SIZE_N: int BLOCK_SIZE_K: int GROUP_SIZE_M: int num_warps: int num_stages: int def calculate_shard_intermediate_size( intermediate_size: int, tp_size: int, ep_size: int = 1 ) -> int: assert tp_size % ep_size == 0 moe_tp_size = tp_size // ep_size assert intermediate_size % moe_tp_size == 0 return 2 * intermediate_size // moe_tp_size def get_model_config( model_name: str, tp_size: int, ep_size: int = 1, disable_shared_experts_fusion: bool = False, topk_ids_dir: str = None, ) -> Dict: config = get_config(model_name, trust_remote_code=True) # Replace config with text_config for encoder-decoder models after getting block_shape and architecture if hasattr(config, "text_config"): config = config.get_text_config() block_shape = None if ( hasattr(config, "quantization_config") and "weight_block_size" in config.quantization_config ): block_shape = config.quantization_config["weight_block_size"] assert len(block_shape) == 2 if ( hasattr(config, "quantization_config") and "config_groups" in config.quantization_config ): config_groups = config.quantization_config["config_groups"] # Get group_size from the first group's weights config first_group = next(iter(config_groups.values()), {}) weights_config = first_group.get("weights", {}) group_size = weights_config.get("group_size") block_shape = [0, group_size] assert len(block_shape) == 2 architecture = config.architectures[0] hidden_size = config.hidden_size if architecture == "DbrxForCausalLM": E = config.ffn_config.moe_num_experts // ep_size topk = config.ffn_config.moe_top_k intermediate_size = config.ffn_config.ffn_hidden_size elif architecture == "JambaForCausalLM": E = config.num_experts // ep_size topk = config.num_experts_per_tok intermediate_size = config.intermediate_size elif architecture in [ "Qwen2MoeForCausalLM", "Qwen3MoeForCausalLM", "Qwen3NextForCausalLM", "Qwen3VLMoeForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration", ]: E = config.num_experts // ep_size topk = config.num_experts_per_tok intermediate_size = config.moe_intermediate_size elif architecture in [ "DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", "Glm4MoeForCausalLM", "GlmMoeDsaForCausalLM", "MistralLarge3ForCausalLM", ]: E = (config.n_routed_experts // ep_size) + ( 0 if disable_shared_experts_fusion or architecture not in [ "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", "Glm4MoeForCausalLM", "GlmMoeDsaForCausalLM", "MistralLarge3ForCausalLM", ] else 1 ) topk = config.num_experts_per_tok + ( 0 if disable_shared_experts_fusion or topk_ids_dir is None else 1 ) intermediate_size = config.moe_intermediate_size elif architecture == "Llama4ForConditionalGeneration": E = config.num_local_experts // ep_size + ( 0 if disable_shared_experts_fusion else 1 ) topk = config.num_experts_per_tok + ( 0 if disable_shared_experts_fusion or topk_ids_dir is None else 1 ) intermediate_size = config.intermediate_size elif architecture in [ "Grok1ForCausalLM", "Grok1ImgGen", "Grok1AForCausalLM", ]: E = config.num_local_experts // ep_size topk = config.num_experts_per_tok intermediate_size = config.moe_intermediate_size elif architecture in [ "BailingMoEForCausalLM", "BailingMoeForCausalLM", "BailingMoeV2ForCausalLM", ]: E = config.num_experts // ep_size topk = config.num_experts_per_tok intermediate_size = config.moe_intermediate_size elif architecture == "NemotronHForCausalLM": E = config.n_routed_experts // ep_size topk = config.num_experts_per_tok intermediate_size = config.moe_intermediate_size hidden_size = getattr(config, "moe_latent_size", None) or hidden_size else: # Default: Mixtral E = config.num_local_experts // ep_size topk = config.num_experts_per_tok intermediate_size = config.intermediate_size shard_intermediate_size = calculate_shard_intermediate_size( intermediate_size, tp_size, ep_size ) return { "num_experts": E, "topk": topk, "hidden_size": hidden_size, "shard_intermediate_size": shard_intermediate_size, "dtype": config.torch_dtype, "block_shape": block_shape, "architecture": architecture, } def get_rocm_configs_compute_bound() -> List[Dict[str, int]]: configs: List[BenchmarkConfig] = [] waves_per_eu_range = 0 for num_stages in [2]: for block_m in [32, 64, 128, 256]: for block_k in [32, 64, 128, 256]: for block_n in [16, 32, 64, 128, 256]: for num_warps in [1, 2, 4, 8]: for group_size in [1, 4, 8, 16, 32]: configs.append( { "BLOCK_SIZE_M": block_m, "BLOCK_SIZE_N": block_n, "BLOCK_SIZE_K": block_k, "GROUP_SIZE_M": group_size, "num_warps": num_warps, "num_stages": num_stages, "waves_per_eu": waves_per_eu_range, } ) return configs def get_configs_compute_bound() -> List[Dict[str, int]]: configs: List[BenchmarkConfig] = [] if is_hip(): configs = get_rocm_configs_compute_bound() else: for num_stages in [2, 3, 4, 5]: for block_m in [16, 32, 64, 128, 256]: for block_k in [64, 128, 256]: for block_n in [32, 64, 128, 256]: for num_warps in [4, 8]: for group_size in [1, 16, 32, 64]: configs.append( { "BLOCK_SIZE_M": block_m, "BLOCK_SIZE_N": block_n, "BLOCK_SIZE_K": block_k, "GROUP_SIZE_M": group_size, "num_warps": num_warps, "num_stages": num_stages, } ) return configs def sort_config(config: BenchmarkConfig) -> BenchmarkConfig: return { "BLOCK_SIZE_M": config["BLOCK_SIZE_M"], "BLOCK_SIZE_N": config["BLOCK_SIZE_N"], "BLOCK_SIZE_K": config["BLOCK_SIZE_K"], "GROUP_SIZE_M": config["GROUP_SIZE_M"], "num_warps": config["num_warps"], "num_stages": config["num_stages"], **( {"waves_per_eu": config["waves_per_eu"]} if "waves_per_eu" in config else {} ), **({"USE_TMA": config["USE_TMA"]} if "USE_TMA" in config else {}), } def save_configs( configs: Dict[int, BenchmarkConfig], filename: str, ) -> None: print(f"Writing best config to {filename}...") with open(filename, "w") as f: json.dump(configs, f, indent=4) f.write("\n") def get_config_filename( num_experts: int, shard_intermediate_size: int, hidden_size: int, topk: int, dtype: torch.dtype, use_fp8_w8a8: bool, use_int8_w8a8: bool, use_int8_w8a16: bool, use_int4_w4a16: bool, per_channel_quant: bool, block_shape: List[int], ) -> str: dtype_str = get_config_dtype_str( dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8, use_int8_w8a8=use_int8_w8a8, use_int4_w4a16=use_int4_w4a16, ) # NOTE(woosuk): The current naming convention uses w2.shape[2], which # is the intermediate size after silu_and_mul. N = shard_intermediate_size // 2 if use_int4_w4a16: N = N // 2 filename = get_config_file_name( num_experts, N, dtype_str, block_shape, per_channel_quant, ) return filename def get_default_batch_sizes() -> List[int]: return [ 1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 256, 512, 1024, 1536, 2048, 3072, 4096, ] ================================================ FILE: benchmark/kernels/fused_moe_triton/tuning_client.py ================================================ import argparse import os import time import openai """ # Edit the code file srt/models/deepseek_v2.py in the Python site package and add the logic for saving topk_ids: # import get_tensor_model_parallel_rank # DeepseekV2MoE::forward_normal if hidden_states.shape[0] >= 4096 and get_tensor_model_parallel_rank() == 0: topk_ids_dir = xxxx if not hasattr(self, "save_idx"): self.save_idx = 0 if self.save_idx <= 1: torch.save(topk_output.topk_ids, f"{topk_ids_dir}/topk_ids_layer{self.layer_id}_idx{self.save_idx}.pt") self.save_idx += 1 """ def read_long_prompt(): import json current_dir = os.path.dirname(os.path.abspath(__file__)) with open(f"{current_dir}/tuning_text.json", "r") as fp: text = fp.read() rst = json.loads(text) return rst["prompt"] def openai_stream_test(model, ip, port): client = openai.Client(base_url=f"http://{ip}:{port}/v1", api_key="None") qst = read_long_prompt() messages = [ {"role": "user", "content": qst}, ] msg2 = dict( model=model, messages=messages, temperature=0.6, top_p=0.75, max_tokens=100, ) response = client.chat.completions.create(**msg2, stream=True) time_start = time.time() time_cost = [] for chunk in response: time_end = time.time() # if chunk.choices[0].delta.content: # print(chunk.choices[0].delta.content, end="", flush=True) time_cost.append(time_end - time_start) time_start = time.time() ttft = time_cost[0] + time_cost[1] tpot = sum(time_cost[2:]) / len(time_cost[2:]) print(f"\nTTFT {ttft}, TPOT {tpot}") return ttft, tpot if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--model", type=str, default="auto") parser.add_argument( "--ip", type=str, default="127.0.0.1", ) parser.add_argument("--port", type=int, default=8188) args = parser.parse_args() openai_stream_test(args.model, args.ip, args.port) ================================================ FILE: benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/benchmark_moe.py import argparse import time from contextlib import nullcontext from datetime import datetime from typing import Any, Dict, List, Tuple import ray import torch import triton from common_utils import ( BenchmarkConfig, get_config_filename, get_configs_compute_bound, get_default_batch_sizes, get_model_config, save_configs, sort_config, ) from ray.experimental.tqdm_ray import tqdm from sglang.srt.layers.moe.fused_moe_triton import override_config from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_moe from sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_config import ( get_config_dtype_str, get_default_config, get_moe_configs, ) from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig from sglang.srt.layers.moe.topk import TopKConfig, select_experts from sglang.srt.server_args import ( ServerArgs, set_global_server_args_for_scheduler, ) from sglang.srt.utils import is_hip _is_hip = is_hip() def benchmark_config( config: BenchmarkConfig, num_tokens: int, num_experts: int, shard_intermediate_size: int, hidden_size: int, topk: int, dtype: torch.dtype, use_fp8_w8a8: bool, use_int8_w8a8: bool, use_int8_w8a16: bool, use_int4_w4a16: bool, per_channel_quant: bool, block_shape: List[int] = None, num_iters: int = 100, ) -> float: init_dtype = torch.float16 if use_fp8_w8a8 else dtype x = torch.randn(num_tokens, hidden_size, dtype=dtype) if use_int8_w8a16 or use_int8_w8a8: w1 = torch.randint( -127, 127, ( num_experts, shard_intermediate_size, hidden_size, ), dtype=torch.int8, ) w2 = torch.randint( -127, 127, ( num_experts, hidden_size, shard_intermediate_size // 2, ), dtype=torch.int8, ) elif use_int4_w4a16: w1 = torch.randint( 0, 255, ( num_experts, shard_intermediate_size, hidden_size // 2, ), dtype=torch.uint8, ) w2 = torch.randint( 0, 255, ( num_experts, hidden_size, shard_intermediate_size // 4, ), dtype=torch.uint8, ) else: w1 = torch.randn( num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype ) w2 = torch.randn( num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype ) gating_output = torch.randn(num_iters, num_tokens, num_experts, dtype=torch.float32) w1_scale = None w2_scale = None a1_scale = None a2_scale = None if use_int8_w8a16: w1_scale = torch.randn( (num_experts, 2 * shard_intermediate_size), dtype=torch.float32 ) w2_scale = torch.randn((hidden_size, num_experts), dtype=torch.float32) if use_int4_w4a16: block_n = 1 if (block_shape[0] == 0) else block_shape[0] block_k = block_shape[1] n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n n_tiles_w2 = (hidden_size + block_n - 1) // block_n k_tiles_w1 = (hidden_size + block_k - 1) // block_k k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k w1_scale = torch.randn( (num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.bfloat16 ) w2_scale = torch.randn( (num_experts, n_tiles_w2, k_tiles_w2), dtype=torch.bfloat16 ) if use_fp8_w8a8 or use_int8_w8a8: if use_int8_w8a8 and block_shape is None: w1_scale = torch.randn( num_experts, shard_intermediate_size, dtype=torch.float32 ) w2_scale = torch.randn(num_experts, hidden_size, dtype=torch.float32) elif block_shape is None: w1_scale = torch.randn(num_experts, dtype=torch.float32) w2_scale = torch.randn(num_experts, dtype=torch.float32) a1_scale = torch.randn(1, dtype=torch.float32) a2_scale = torch.randn(1, dtype=torch.float32) else: block_n, block_k = block_shape[0], block_shape[1] n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n n_tiles_w2 = (hidden_size + block_n - 1) // block_n k_tiles_w1 = (hidden_size + block_k - 1) // block_k k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k w1_scale = torch.rand( (num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.float32 ) w2_scale = torch.rand( (num_experts, n_tiles_w2, k_tiles_w2), dtype=torch.float32 ) if use_fp8_w8a8: w1 = w1.to(torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn) w2 = w2.to(torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn) input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32) topk_config = TopKConfig( top_k=topk, renormalize=True, ) topk_output = select_experts(x, input_gating, topk_config) def prepare(i: int): input_gating = gating_output[i] new_topk_output = select_experts(x, input_gating, topk_config) topk_output.topk_weights.copy_(new_topk_output.topk_weights) topk_output.topk_ids.copy_(new_topk_output.topk_ids) topk_output.router_logits.copy_(new_topk_output.router_logits) def run(): moe_runner_config = MoeRunnerConfig( inplace=True, ) with override_config(config): fused_moe( x, w1, w2, topk_output, moe_runner_config=moe_runner_config, use_fp8_w8a8=use_fp8_w8a8, use_int8_w8a8=use_int8_w8a8, use_int8_w8a16=use_int8_w8a16, use_int4_w4a16=use_int4_w4a16, w1_scale=w1_scale, w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, per_channel_quant=per_channel_quant, block_shape=block_shape, ) # JIT compilation & warmup run() torch.cuda.synchronize() # Capture 10 invocations with CUDA graph graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): for _ in range(10): run() torch.cuda.synchronize() # Warmup for _ in range(5): graph.replay() torch.cuda.synchronize() # Flush L2 cache with 256 MB data cache_flush = torch.empty(int(256e6 // 4), dtype=torch.int, device="cuda") cache_flush.zero_() start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_iters)] end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_iters)] for i in range(num_iters): prepare(i) start_events[i].record() graph.replay() end_events[i].record() torch.cuda.synchronize() latencies: List[float] = [] for i in range(num_iters): latencies.append(start_events[i].elapsed_time(end_events[i])) avg = sum(latencies) / (num_iters * 10) * 1000 # us graph.reset() return avg @ray.remote(num_gpus=1) class BenchmarkWorker: def __init__(self, seed: int, server_args: ServerArgs) -> None: torch.set_default_device("cuda") torch.cuda.manual_seed_all(0) self.seed = seed # Get the device ID to allocate tensors and kernels # on the respective GPU. self.device_id = int(ray.get_gpu_ids()[0]) set_global_server_args_for_scheduler(server_args) def benchmark( self, num_tokens: int, num_experts: int, shard_intermediate_size: int, hidden_size: int, topk: int, dtype: torch.dtype, use_fp8_w8a8: bool, use_int8_w8a8: bool, use_int8_w8a16: bool, use_int4_w4a16: bool, per_channel_quant: bool, block_shape: List[int], ) -> Tuple[Dict[str, int], float]: torch.cuda.manual_seed_all(0) dtype_str = get_config_dtype_str( dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8, use_int4_w4a16=use_int4_w4a16, ) # NOTE(woosuk): The current naming convention uses w2.shape[2], which # is the intermediate size after silu_and_mul. block_n = block_shape[0] if block_shape else 0 block_k = block_shape[1] if block_shape else 0 N = shard_intermediate_size // 2 if use_int4_w4a16: N = N // 2 op_config = get_moe_configs( num_experts, N, dtype_str, block_n, block_k, per_channel_quant, ) if op_config is None: config = get_default_config( num_tokens, num_experts, shard_intermediate_size, hidden_size, topk, dtype_str, False, block_shape, ) else: config = op_config[min(op_config.keys(), key=lambda x: abs(x - num_tokens))] with torch.cuda.device(self.device_id) if is_hip() else nullcontext(): kernel_time = benchmark_config( config, num_tokens, num_experts, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, per_channel_quant, block_shape, ) return config, kernel_time def tune( self, num_tokens: int, num_experts: int, shard_intermediate_size: int, hidden_size: int, topk: int, dtype: torch.dtype, use_fp8_w8a8: bool, use_int8_w8a8: bool, use_int8_w8a16: bool, use_int4_w4a16: bool, per_channel_quant: bool, block_shape: List[int], search_space: List[Dict[str, int]], ) -> Dict[str, int]: best_config = None best_time = float("inf") with torch.cuda.device(self.device_id) if is_hip() else nullcontext(): for config in tqdm(search_space): try: kernel_time = benchmark_config( config, num_tokens, num_experts, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, per_channel_quant, block_shape, num_iters=10, ) except (triton.runtime.autotuner.OutOfResources, RuntimeError): # Some configurations may be invalid and fail to compile. continue if kernel_time < best_time: best_time = kernel_time best_config = config now = datetime.now() print(f"{now.ctime()}] Completed tuning for batch_size={num_tokens}") assert best_config is not None return best_config def main(args: argparse.Namespace): server_args = ServerArgs( model_path=args.model, tp_size=args.tp_size, ep_size=args.ep_size ) model_config = get_model_config( args.model, args.tp_size, args.ep_size, args.disable_shared_experts_fusion ) E = model_config["num_experts"] topk = model_config["topk"] hidden_size = model_config["hidden_size"] shard_intermediate_size = model_config["shard_intermediate_size"] dtype = model_config["dtype"] block_shape = model_config["block_shape"] use_fp8_w8a8 = args.dtype == "fp8_w8a8" use_int8_w8a8 = args.dtype == "int8_w8a8" use_int8_w8a16 = args.dtype == "int8_w8a16" use_int4_w4a16 = args.dtype == "int4_w4a16" per_channel_quant = args.per_channel_quant if args.batch_size is None: batch_sizes = get_default_batch_sizes() else: batch_sizes = [args.batch_size] ray.init() num_gpus = int(ray.available_resources()["GPU"]) workers = [BenchmarkWorker.remote(args.seed, server_args) for _ in range(num_gpus)] def _distribute(method: str, inputs: List[Any]) -> List[Any]: outputs = [] worker_idx = 0 for input_args in inputs: worker = workers[worker_idx] worker_method = getattr(worker, method) output = worker_method.remote(*input_args) outputs.append(output) worker_idx = (worker_idx + 1) % num_gpus return ray.get(outputs) if args.tune: search_space = get_configs_compute_bound() if block_shape is not None: block_n, block_k = block_shape[0], block_shape[1] search_space = [ config for config in search_space if block_k % config["BLOCK_SIZE_K"] == 0 ] filename = get_config_filename( E, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, per_channel_quant, block_shape, ) print( f"Start tuning over {len(search_space)} configurations to create {filename}..." ) start = time.perf_counter() configs = _distribute( "tune", [ ( batch_size, E, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, per_channel_quant, block_shape, search_space, ) for batch_size in batch_sizes ], ) best_configs = { M: sort_config(config) for M, config in zip(batch_sizes, configs) } save_configs( best_configs, filename, ) end = time.perf_counter() print(f"Tuning took {end - start:.2f} seconds") else: outputs = _distribute( "benchmark", [ ( batch_size, E, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, per_channel_quant, block_shape, ) for batch_size in batch_sizes ], ) for batch_size, (config, kernel_time) in zip(batch_sizes, outputs): print(f"Batch size: {batch_size}, config: {config}") print(f"Kernel time: {kernel_time:.2f} us") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1" ) parser.add_argument("--tp-size", "--tp", type=int, default=2) parser.add_argument("--ep-size", "--ep", type=int, default=1) parser.add_argument( "--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16", "int8_w8a8", "int4_w4a16"], default="auto", ) parser.add_argument( "--per-channel-quant", action="store_true", ) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--batch-size", type=int, required=False) parser.add_argument("--tune", action="store_true") parser.add_argument("--disable-shared-experts-fusion", action="store_true") args = parser.parse_args() main(args) ================================================ FILE: benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/benchmark_moe.py import argparse import dataclasses import json import os import time from contextlib import nullcontext from datetime import datetime from typing import Any, Dict, List, Tuple import ray import torch import triton import triton.language as tl from common_utils import ( BenchmarkConfig, get_config_filename, get_configs_compute_bound, get_default_batch_sizes, get_model_config, sort_config, ) from ray.experimental.tqdm_ray import tqdm from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( get_config_dtype_str, invoke_fused_moe_kernel, moe_align_block_size, ) from sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_config import ( get_config_file_name, ) from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig from sglang.srt.layers.moe.topk import TopKConfig, select_experts from sglang.srt.server_args import ( ServerArgs, set_global_server_args_for_scheduler, ) from sglang.srt.utils import is_hip _is_hip = is_hip() @dataclasses.dataclass class MoeInputs: topk_ids: torch.Tensor sorted_token_ids: torch.Tensor expert_ids: torch.Tensor num_tokens_post_padded: torch.Tensor class KernelWrapper: def __init__(self, moe_inputs, use_cuda_graph=True, inner_iter=10, **kwargs): self.func = invoke_fused_moe_kernel self.use_cuda_graph = use_cuda_graph self.moe_inputs = moe_inputs self.inner_iter = inner_iter self.kwargs = kwargs if use_cuda_graph: self.graph = self.cuda_graph_wrapper() else: self.graph = None def cuda_graph_wrapper(self): moe_input = self.moe_inputs[0] self.func( **self.kwargs, topk_ids=moe_input.topk_ids, sorted_token_ids=moe_input.sorted_token_ids, expert_ids=moe_input.expert_ids, num_tokens_post_padded=moe_input.num_tokens_post_padded, ) torch.cuda.synchronize() # Capture 10 invocations with CUDA graph graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): for k in range(self.inner_iter): moe_input = self.moe_inputs[k] self.func( **self.kwargs, topk_ids=moe_input.topk_ids, sorted_token_ids=moe_input.sorted_token_ids, expert_ids=moe_input.expert_ids, num_tokens_post_padded=moe_input.num_tokens_post_padded, ) torch.cuda.synchronize() # Warmup for _ in range(5): graph.replay() torch.cuda.synchronize() return graph def forward_cost(self, try_cnt=2): time_cost = float("inf") for _ in range(try_cnt): start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) start_event.record() if self.use_cuda_graph: self.graph.replay() else: for k in range(self.inner_iter): moe_input = self.moe_inputs[k] self.func( **self.kwargs, topk_ids=moe_input.topk_ids, sorted_token_ids=moe_input.sorted_token_ids, expert_ids=moe_input.expert_ids, num_tokens_post_padded=moe_input.num_tokens_post_padded, ) end_event.record() torch.cuda.synchronize() time_cost = min(time_cost, start_event.elapsed_time(end_event)) return time_cost def load_topk_ids(topk_ids_dir, i: int): num_layers = 61 dense_layers = 3 moe_layers = num_layers - dense_layers return torch.load( f"{topk_ids_dir}/topk_ids_layer{i % moe_layers + dense_layers}_idx{i // moe_layers}.pt" ) def benchmark_config( config: BenchmarkConfig, num_tokens: int, num_experts: int, shard_intermediate_size: int, hidden_size: int, topk: int, dtype: torch.dtype, use_fp8_w8a8: bool, use_int8_w8a8: bool, use_int8_w8a16: bool, use_int4_w4a16: bool, topk_ids_list, block_shape: List[int] = None, ep_size: int = 1, num_iters: int = 100, ) -> float: ncu_enable = os.getenv("NCU_ENABLE", "0") == "1" if ncu_enable: num_iters = 1 init_dtype = torch.float16 if use_fp8_w8a8 else dtype hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype) if use_int8_w8a16 or use_int8_w8a8: w1 = torch.randint( -127, 127, ( num_experts, shard_intermediate_size, hidden_size, ), dtype=torch.int8, ) w2 = torch.randint( -127, 127, ( num_experts, hidden_size, shard_intermediate_size // 2, ), dtype=torch.int8, ) elif use_int4_w4a16: w1 = torch.randint( 0, 255, ( num_experts, shard_intermediate_size, hidden_size // 2, ), dtype=torch.uint8, ) w2 = torch.randint( 0, 255, ( num_experts, hidden_size, shard_intermediate_size // 4, ), dtype=torch.uint8, ) else: w1 = torch.randn( num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype ) w2 = torch.randn( num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype ) w1_scale = None w2_scale = None a1_scale = None a2_scale = None if use_int8_w8a16: w1_scale = torch.randn( (num_experts, 2 * shard_intermediate_size), dtype=torch.float32 ) w2_scale = torch.randn((hidden_size, num_experts), dtype=torch.float32) if use_int4_w4a16: block_n = 1 if (block_shape[0] == 0) else block_shape[0] block_k = block_shape[1] n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n n_tiles_w2 = (hidden_size + block_n - 1) // block_n k_tiles_w1 = (hidden_size + block_k - 1) // block_k k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k w1_scale = torch.randn( (num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.bfloat16 ) w2_scale = torch.randn( (num_experts, n_tiles_w2, k_tiles_w2), dtype=torch.bfloat16 ) if use_fp8_w8a8 or use_int8_w8a8: if use_int8_w8a8 and block_shape is None: w1_scale = torch.randn( num_experts, shard_intermediate_size, dtype=torch.float32 ) w2_scale = torch.randn(num_experts, hidden_size, dtype=torch.float32) elif block_shape is None: w1_scale = torch.randn(num_experts, dtype=torch.float32) w2_scale = torch.randn(num_experts, dtype=torch.float32) a1_scale = torch.randn(1, dtype=torch.float32) a2_scale = torch.randn(1, dtype=torch.float32) else: block_n, block_k = block_shape[0], block_shape[1] n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n n_tiles_w2 = (hidden_size + block_n - 1) // block_n k_tiles_w1 = (hidden_size + block_k - 1) // block_k k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k w1_scale = torch.rand( (num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.float32 ) w2_scale = torch.rand( (num_experts, n_tiles_w2, k_tiles_w2), dtype=torch.float32 ) if use_fp8_w8a8: w1 = w1.to(torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn) w2 = w2.to(torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn) input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32) topk_config = TopKConfig( top_k=topk, renormalize=True, ) topk_output_ = select_experts(hidden_states, input_gating, topk_config) sorted_token_ids_, expert_ids_, num_tokens_post_padded_ = moe_align_block_size( topk_output_.topk_ids, config["BLOCK_SIZE_M"], num_experts ) inner_iter = 10 if not ncu_enable else 1 moe_inputs = [ MoeInputs( topk_output_.topk_ids.clone(), sorted_token_ids_.clone(), expert_ids_.clone(), num_tokens_post_padded_.clone(), ) for _ in range(inner_iter) ] M = hidden_states.shape[0] E, N, _ = w1.shape padded_tokens = min(M * topk, E + 1) * ( config["BLOCK_SIZE_M"] - 1 ) # if moe_use_tma else 0 total_tokens = M * topk + padded_tokens cache = torch.empty( total_tokens * max(N, w2.shape[1]), device=hidden_states.device, dtype=hidden_states.dtype, ) intermediate_cache1 = cache[: total_tokens * N].view( (total_tokens, N), ) intermediate_cache2 = torch.empty( (total_tokens, N // 2), device=hidden_states.device, dtype=hidden_states.dtype, ) intermediate_cache3 = cache[: M * topk * w2.shape[1]].view( (M, topk, w2.shape[1]), ) def prepare(i: int, inner_iter): # update inputs according to topk_ids for k in range(inner_iter): topk_ids = topk_ids_list[i * inner_iter + k] # With EP, saved topk_ids are global expert indices; remap to local. if ep_size > 1: topk_ids = (topk_ids // ep_size).to( device=moe_inputs[k].topk_ids.device, dtype=moe_inputs[k].topk_ids.dtype, ) tokens, _topk = moe_inputs[k].topk_ids.shape moe_inputs[k].topk_ids.copy_(topk_ids[:tokens, :_topk]) sorted_token_ids_, expert_ids_, num_tokens_post_padded_ = ( moe_align_block_size( moe_inputs[k].topk_ids, config["BLOCK_SIZE_M"], num_experts ) ) moe_inputs[k].sorted_token_ids.copy_(sorted_token_ids_) moe_inputs[k].expert_ids.copy_(expert_ids_) moe_inputs[k].num_tokens_post_padded.copy_(num_tokens_post_padded_) def get_kernel_wrapper(moe_use_tma, inner_iter, use_cuda_graph): compute_type = ( tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16 ) moe_runner_config = MoeRunnerConfig( inplace=True, ) apply_router_weight_on_input = moe_runner_config.apply_router_weight_on_input kernel0 = KernelWrapper( A=hidden_states, B=w1, bias=None, C=intermediate_cache1, A_scale=a1_scale, B_scale=w1_scale, B_zp=None, topk_weights=topk_output_.topk_weights, moe_inputs=moe_inputs, mul_routed_weight=apply_router_weight_on_input, top_k=topk, config=config, compute_type=compute_type, use_fp8_w8a8=use_fp8_w8a8, use_int8_w8a8=use_int8_w8a8, use_int8_w8a16=use_int8_w8a16, use_int4_w4a16=use_int4_w4a16, per_channel_quant=False, block_shape=block_shape, b_use_tma=moe_use_tma, c_sorted=moe_use_tma, filter_expert=False, use_cuda_graph=use_cuda_graph, inner_iter=inner_iter, ) kernel1 = KernelWrapper( A=intermediate_cache2, B=w2, bias=None, C=intermediate_cache3, A_scale=a2_scale, B_scale=w2_scale, B_zp=None, topk_weights=topk_output_.topk_weights, moe_inputs=moe_inputs, mul_routed_weight=not apply_router_weight_on_input, top_k=1, config=config, compute_type=compute_type, use_fp8_w8a8=use_fp8_w8a8, use_int8_w8a8=use_int8_w8a8, use_int8_w8a16=use_int8_w8a16, use_int4_w4a16=use_int4_w4a16, per_channel_quant=False, block_shape=block_shape, a_use_tma=moe_use_tma, b_use_tma=moe_use_tma, filter_expert=False, use_cuda_graph=use_cuda_graph, inner_iter=inner_iter, ) return kernel0, kernel1 use_cuda_graph = True if not ncu_enable else False kernel0, kernel1 = get_kernel_wrapper(False, inner_iter, use_cuda_graph) kernel_tma0, kernel_tma1 = get_kernel_wrapper(True, inner_iter, use_cuda_graph) # JIT compilation & warmup if not ncu_enable: kernel0.forward_cost() kernel1.forward_cost() kernel_tma0.forward_cost() kernel_tma1.forward_cost() ts0 = [] ts1 = [] ts_tma0 = [] ts_tma1 = [] for i in range(num_iters // inner_iter): prepare(i, inner_iter) ts0.append(kernel0.forward_cost()) ts1.append(kernel1.forward_cost()) ts_tma0.append(kernel_tma0.forward_cost()) ts_tma1.append(kernel_tma1.forward_cost()) torch.cuda.synchronize() avg = sum(ts0) / (num_iters) * 1000 # us avg1 = sum(ts1) / (num_iters) * 1000 # us avg_tma = sum(ts_tma0) / (num_iters) * 1000 # us avg1_tma = sum(ts_tma1) / (num_iters) * 1000 # us return avg, avg_tma, avg1, avg1_tma class BestConfigTrace: def __init__(self, name, down_moe=False): self.name = name self.down_moe = down_moe self.best_costs_m = {} # block_m: best_cost def update(self, config, time_cost_all): block_m = config["BLOCK_SIZE_M"] if not self.down_moe: time_cost = time_cost_all[0] else: time_cost = min(time_cost_all[2], time_cost_all[3]) if ( block_m not in self.best_costs_m or time_cost < self.best_costs_m[block_m][1] ): self.best_costs_m[block_m] = config, time_cost, time_cost_all def time_cost(self, block_m): if block_m not in self.best_costs_m: return float("inf") time_cost = self.best_costs_m[block_m][1] return time_cost def config_dict(self, block_m): if block_m not in self.best_costs_m: return {} config, _, time_cost_all = self.best_costs_m[block_m] if not self.down_moe: return config else: return { **config, "USE_TMA": time_cost_all[2] > time_cost_all[3], } class BenchmarkWorker: def __init__(self, seed: int, server_args: ServerArgs) -> None: torch.set_default_device("cuda") torch.cuda.manual_seed_all(0) self.seed = seed # Get the device ID to allocate tensors and kernels # on the respective GPU. self.device_id = 0 # int(ray.get_gpu_ids()[0]) set_global_server_args_for_scheduler(server_args) def benchmark( self, num_tokens: int, num_experts: int, shard_intermediate_size: int, hidden_size: int, topk: int, dtype: torch.dtype, use_fp8_w8a8: bool, use_int8_w8a8: bool, use_int8_w8a16: bool, use_int4_w4a16: bool, block_shape: List[int], cfg: Dict[str, int], topk_ids_dir: str, ep_size: int = 1, ) -> Tuple[Dict[str, int], float]: torch.cuda.manual_seed_all(0) topk_ids_list = [load_topk_ids(topk_ids_dir, i) for i in range(100)] with torch.cuda.device(self.device_id) if is_hip() else nullcontext(): kernel_time = benchmark_config( cfg, num_tokens, num_experts, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, topk_ids_list, block_shape, ep_size=ep_size, ) return cfg, kernel_time def tune( self, num_tokens: int, num_experts: int, shard_intermediate_size: int, hidden_size: int, topk: int, dtype: torch.dtype, use_fp8_w8a8: bool, use_int8_w8a8: bool, use_int8_w8a16: bool, use_int4_w4a16: bool, block_shape: List[int], search_space: List[Dict[str, int]], topk_ids_dir: str, ep_size: int = 1, ) -> Dict[str, int]: trace0 = BestConfigTrace("kernel0", down_moe=False) trace1 = BestConfigTrace("kernel1", down_moe=True) topk_ids_list = [load_topk_ids(topk_ids_dir, i) for i in range(100)] with torch.cuda.device(self.device_id) if is_hip() else nullcontext(): for config in tqdm(search_space): try: kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma = benchmark_config( config, num_tokens, num_experts, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, topk_ids_list, block_shape, ep_size=ep_size, num_iters=100, ) except triton.runtime.autotuner.OutOfResources: # Some configurations may be invalid and fail to compile. continue trace0.update( config, (kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma), ) trace1.update( config, (kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma), ) now = datetime.now() print(f"{now.ctime()}] Completed tuning for batch_size={num_tokens}") best_block_m = 16 for block_m in (32, 64, 128, 256): if trace0.time_cost(block_m) + trace1.time_cost(block_m) < trace0.time_cost( best_block_m ) + trace1.time_cost(best_block_m): best_block_m = block_m return ( trace0.config_dict(best_block_m), trace1.config_dict(best_block_m), trace0.time_cost(best_block_m), trace1.time_cost(best_block_m), ) def cmp_configs( self, num_tokens: List[int], num_experts: int, shard_intermediate_size: int, hidden_size: int, topk: int, dtype: torch.dtype, use_fp8_w8a8: bool, use_int8_w8a8: bool, use_int8_w8a16: bool, use_int4_w4a16: bool, block_shape: List[int], cmp_config_files: List[str], topk_ids_dir: str, ep_size: int = 1, ): # compare performance of different configs cmp_configs = [] for file in cmp_config_files: with open(file) as f: cmp_configs.append({int(key): val for key, val in json.load(f).items()}) for i, file in enumerate(cmp_config_files): print(f"config {i}: {file}") topk_ids_list = [load_topk_ids(topk_ids_dir, i) for i in range(100)] torch.cuda.manual_seed_all(0) with torch.cuda.device(self.device_id) if is_hip() else nullcontext(): for bs in num_tokens: kernel_times = [] cfgs = [] for configs in cmp_configs: cfg_org = configs[min(configs.keys(), key=lambda x: abs(x - bs))] cfgs.append(cfg_org) cfg = cfg_org.copy() cfg.pop("USE_TMA", None) kernel_time = benchmark_config( cfg, bs, num_experts, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, topk_ids_list, block_shape, ep_size=ep_size, ) kernel_times.append(kernel_time) print(f"batch_size={bs=}:") for i, cfg in enumerate(cfgs): print(f" config {i} {cfg}: {kernel_times[i]}") def save_configs_sep( configs: Dict[int, BenchmarkConfig], num_experts: int, shard_intermediate_size: int, hidden_size: int, topk: int, dtype: torch.dtype, use_fp8_w8a8: bool, use_int8_w8a8: bool, use_int8_w8a16: bool, use_int4_w4a16: bool, block_shape: List[int], down_moe: bool = False, ) -> None: dtype_str = get_config_dtype_str( dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8, use_int8_w8a8=use_int8_w8a8, use_int4_w4a16=use_int4_w4a16, ) # NOTE(woosuk): The current naming convention uses w2.shape[2], which # is the intermediate size after silu_and_mul. filename = get_config_file_name( num_experts, shard_intermediate_size // 2, dtype_str, block_shape, down_moe=down_moe, ) print(f"Writing best config to {filename}...") with open(filename, "w") as f: json.dump(configs, f, indent=4) f.write("\n") def main(args: argparse.Namespace): print(args) server_args = ServerArgs( model_path=args.model, tp_size=args.tp_size, ep_size=args.ep_size ) model_config = get_model_config( args.model, args.tp_size, args.ep_size, args.disable_shared_experts_fusion, args.topk_ids_dir, ) E = model_config["num_experts"] topk = model_config["topk"] hidden_size = model_config["hidden_size"] shard_intermediate_size = model_config["shard_intermediate_size"] dtype = model_config["dtype"] block_shape = model_config["block_shape"] use_fp8_w8a8 = args.dtype == "fp8_w8a8" use_int8_w8a8 = args.dtype == "int8_w8a8" use_int8_w8a16 = args.dtype == "int8_w8a16" use_int4_w4a16 = args.dtype == "int4_w4a16" topk_ids_dir = args.topk_ids_dir if args.batch_size is None: batch_sizes = get_default_batch_sizes() batch_sizes.reverse() else: batch_sizes = [args.batch_size] if args.cmp_configs is not None: worker = BenchmarkWorker(args.seed, server_args) worker.cmp_configs( batch_sizes, E, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, block_shape, args.cmp_configs, topk_ids_dir, args.ep_size, ) return if len(batch_sizes) == 1: worker = BenchmarkWorker(args.seed, server_args) if args.tune: search_space = get_configs_compute_bound() worker.tune( batch_sizes[0], E, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, block_shape, search_space, topk_ids_dir, args.ep_size, ) else: cfg = { "BLOCK_SIZE_M": args.configs[0], "BLOCK_SIZE_N": args.configs[1], "BLOCK_SIZE_K": args.configs[2], "GROUP_SIZE_M": args.configs[3], "num_warps": args.configs[4], "num_stages": args.configs[5], } _, (t0, t0_tma, t1, t1_tma) = worker.benchmark( args.batch_size, E, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, block_shape, cfg, topk_ids_dir, args.ep_size, ) print(f"{t0=}, {t0_tma=}, {t1=}, {t1_tma=}") return assert args.tune ray.init() num_gpus = int(ray.available_resources()["GPU"]) workers = [ ray.remote(num_gpus=1)(BenchmarkWorker).remote(args.seed, server_args) for _ in range(num_gpus) ] def _distribute(method: str, inputs: List[Any]) -> List[Any]: outputs = [] worker_idx = 0 for input_args in inputs: worker = workers[worker_idx] worker_method = getattr(worker, method) output = worker_method.remote(*input_args) outputs.append(output) worker_idx = (worker_idx + 1) % num_gpus return ray.get(outputs) search_space = get_configs_compute_bound() if block_shape is not None: block_n, block_k = block_shape[0], block_shape[1] search_space = [ config for config in search_space if block_k % config["BLOCK_SIZE_K"] == 0 ] filename = get_config_filename( E, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, False, block_shape, ) print( f"Start tuning over {len(search_space)} configurations to create {filename}..." ) start = time.perf_counter() configs = _distribute( "tune", [ ( batch_size, E, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, block_shape, search_space, topk_ids_dir, args.ep_size, ) for batch_size in batch_sizes ], ) print(f"{configs=}", flush=True) cur_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) with open(f"tuning_result_{cur_time}.txt", "w") as f: print(configs, file=f) batch_sizes.reverse() configs0 = [config[0] for config in configs] configs1 = [config[1] for config in configs] configs0.reverse() configs1.reverse() best_configs0 = {M: sort_config(config) for M, config in zip(batch_sizes, configs0)} save_configs_sep( best_configs0, E, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, block_shape, ) best_configs1 = {M: sort_config(config) for M, config in zip(batch_sizes, configs1)} save_configs_sep( best_configs1, E, shard_intermediate_size, hidden_size, topk, dtype, use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, block_shape, down_moe=True, ) end = time.perf_counter() print(f"Tuning took {end - start:.2f} seconds") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1" ) parser.add_argument("--tp-size", "--tp", type=int, default=2) parser.add_argument("--ep-size", "--ep", type=int, default=1) parser.add_argument( "--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16", "int8_w8a8", "int8_w4a16"], default="auto", ) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--batch-size", type=int, required=False) parser.add_argument("--tune", action="store_true") parser.add_argument("--disable-shared-experts-fusion", action="store_true") parser.add_argument("--configs", type=int, nargs="+", required=False) parser.add_argument("--topk-ids-dir", type=str, required=True) parser.add_argument("--cmp-configs", type=str, nargs="+", required=False) args = parser.parse_args() main(args) ================================================ FILE: benchmark/kernels/fused_moe_triton/tuning_text.json ================================================ {"prompt": "Here are the relevant Wikipedia articles:\nThe president of the United States (POTUS) is the head of state and head of government of the United States of America. The president directs the executive branch of the federal government and is the commander-in-chief of the United States Armed Forces.\nThe power of the presidency has grown substantially since the first president, George Washington, took office in 1789. While presidential power has ebbed and flowed over time, the presidency has played an increasingly significant role in American political life since the beginning of the 20th century, carrying over into the 21st century with notable expansions during the presidencies of Franklin D. Roosevelt and George W. Bush. In modern times, the president is one of the world's most powerful political figures and the leader of the world's only remaining superpower. As the leader of the nation with the largest economy by nominal GDP, the president possesses significant domestic and international hard and soft power. For much of the 20th century, especially during the Cold War, the U.S. president was often called \"the leader of the free world\".\nArticle II of the Constitution establishes the executive branch of the federal government and vests executive power in the president. The power includes the execution and enforcement of federal law and the responsibility to appoint federal executive, diplomatic, regulatory, and judicial officers. Based on constitutional provisions empowering the president to appoint and receive ambassadors and conclude treaties with foreign powers, and on subsequent laws enacted by Congress, the modern presidency has primary responsibility for conducting U.S. foreign policy. The role includes responsibility for directing the world's most expensive military, which has the second-largest nuclear arsenal.\nThe president also plays a leading role in federal legislation and domestic policymaking. As part of the system of separation of powers, Article I, Section 7 of the Constitution gives the president the power to sign or veto federal legislation. Since modern presidents are typically viewed as leaders of their political parties, major policymaking is significantly shaped by the outcome of presidential elections, with presidents taking an active role in promoting their policy priorities to members of Congress who are often electorally dependent on the president. In recent decades, presidents have also made increasing use of executive orders, agency regulations, and judicial appointments to shape domestic policy.\nThe president is elected indirectly through the Electoral College to a four-year term, along with the vice president. Under the Twenty-second Amendment, ratified in 1951, no person who has been elected to two presidential terms may be elected to a third. In addition, nine vice presidents have become president by virtue of a president's intra-term death or resignation. In all, 45 individuals have served 46 presidencies spanning 58 four-year terms. Joe Biden is the 46th and current president, having assumed office on January 20, 2021.\n\nHistory and development\nOrigins\nDuring the American Revolutionary War, the Thirteen Colonies, represented by the Second Continental Congress in Philadelphia, declared themselves to be independent sovereign states and no longer under British rule. The affirmation was made in the Declaration of Independence, which was written predominantly by Thomas Jefferson and adopted unanimously on July 4, 1776, by the Second Continental Congress. Recognizing the necessity of closely coordinating their efforts against the British, the Continental Congress simultaneously began the process of drafting a constitution that would bind the states together. There were long debates on a number of issues, including representation and voting, and the exact powers to be given the central government. Congress finished work on the Articles of Confederation to establish a perpetual union between the states in November 1777 and sent it to the states for ratification.\nUnder the Articles, which took effect on March 1, 1781, the Congress of the Confederation was a central political authority without any legislative power. It could make its own resolutions, determinations, and regulations, but not any laws, and could not impose any taxes or enforce local commercial regulations upon its citizens. This institutional design reflected how Americans believed the deposed British system of Crown and Parliament ought to have functioned with respect to the royal dominion: a superintending body for matters that concerned the entire empire. The states were out from under any monarchy and assigned some formerly royal prerogatives (e.g., making war, receiving ambassadors, etc.) to Congress; the remaining prerogatives were lodged within their own respective state governments. The members of Congress elected a president of the United States in Congress Assembled to preside over its deliberation as a neutral discussion moderator. Unrelated to and quite dissimilar from the later office of president of the United States, it was a largely ceremonial position without much influence.\nIn 1783, the Treaty of Paris secured independence for each of the former colonies. With peace at hand, the states each turned toward their own internal affairs. By 1786, Americans found their continental borders besieged and weak and their respective economies in crises as neighboring states agitated trade rivalries with one another. They witnessed their hard currency pouring into foreign markets to pay for imports, their Mediterranean commerce preyed upon by North African pirates, and their foreign-financed Revolutionary War debts unpaid and accruing interest. Civil and political unrest loomed. Events such as the Newburgh Conspiracy and Shays' Rebellion demonstrated that the Articles of Confederation were not working.\nFollowing the successful resolution of commercial and fishing disputes between Virginia and Maryland at the Mount Vernon Conference in 1785, Virginia called for a trade conference between all the states, set for September 1786 in Annapolis, Maryland, with an aim toward resolving further-reaching interstate commercial antagonisms. When the convention failed for lack of attendance due to suspicions among most of the other states, Alexander Hamilton of New York led the Annapolis delegates in a call for a convention to offer revisions to the Articles, to be held the next spring in Philadelphia. Prospects for the next convention appeared bleak until James Madison and Edmund Randolph succeeded in securing George Washington's attendance to Philadelphia as a delegate for Virginia.\nWhen the Constitutional Convention convened in May 1787, the 12 state delegations in attendance (Rhode Island did not send delegates) brought with them an accumulated experience over a diverse set of institutional arrangements between legislative and executive branches from within their respective state governments. Most states maintained a weak executive without veto or appointment powers, elected annually by the legislature to a single term only, sharing power with an executive council, and countered by a strong legislature. New York offered the greatest exception, having a strong, unitary governor with veto and appointment power elected to a three-year term, and eligible for reelection to an indefinite number of terms thereafter. It was through the closed-door negotiations at Philadelphia that the presidency framed in the U.S. Constitution emerged.\n\n1789–1933\nAs the nation's first president, George Washington established many norms that would come to define the office. His decision to retire after two terms helped address fears that the nation would devolve into monarchy, and established a precedent that would not be broken until 1940 and would eventually be made permanent by the Twenty-Second Amendment. By the end of his presidency, political parties had developed, with John Adams defeating Thomas Jefferson in 1796, the first truly contested presidential election. After Jefferson defeated Adams in 1800, he and his fellow Virginians James Madison and James Monroe would each serve two terms, eventually dominating the nation's politics during the Era of Good Feelings until Adams' son John Quincy Adams won election in 1824 after the Democratic-Republican Party split.\nThe election of Andrew Jackson in 1828 was a significant milestone, as Jackson was not part of the Virginia and Massachusetts elite that had held the presidency for its first 40 years. Jacksonian democracy sought to strengthen the presidency at the expense of Congress, while broadening public participation as the nation rapidly expanded westward. However, his successor, Martin Van Buren, became unpopular after the Panic of 1837, and the death of William Henry Harrison and subsequent poor relations between John Tyler and Congress led to further weakening of the office. Including Van Buren, in the 24 years between 1837 and 1861, six presidential terms would be filled by eight different men, with none serving two terms. The Senate played an important role during this period, with the Great Triumvirate of Henry Clay, Daniel Webster, and John C. Calhoun playing key roles in shaping national policy in the 1830s and 1840s until debates over slavery began pulling the nation apart in the 1850s.\nAbraham Lincoln's leadership during the Civil War has led historians to regard him as one of the nation's greatest presidents. The circumstances of the war and Republican domination of Congress made the office very powerful, and Lincoln's re-election in 1864 was the first time a president had been re-elected since Jackson in 1832. After Lincoln's assassination, his successor Andrew Johnson lost all political support and was nearly removed from office, with Congress remaining powerful during the two-term presidency of Civil War general Ulysses S. Grant. After the end of Reconstruction, Grover Cleveland would eventually become the first Democratic president elected since before the war, running in three consecutive elections (1884, 1888, 1892) and winning twice. In 1900, William McKinley became the first incumbent to win re-election since Grant in 1872.\nAfter McKinley's assassination by Leon Czolgosz in 1901, Theodore Roosevelt became a dominant figure in American politics. Historians believe Roosevelt permanently changed the political system by strengthening the presidency, with some key accomplishments including breaking up trusts, conservationism, labor reforms, making personal character as important as the issues, and hand-picking his successor, William Howard Taft. The following decade, Woodrow Wilson led the nation to victory during World War I, although Wilson's proposal for the League of Nations was rejected by the Senate. Warren Harding, while popular in office, would see his legacy tarnished by scandals, especially Teapot Dome, and Herbert Hoover quickly became very unpopular after failing to alleviate the Great Depression.\n\nImperial presidency\nThe ascendancy of Franklin D. Roosevelt in 1933 led further toward what historians now describe as the Imperial presidency. Backed by enormous Democratic majorities in Congress and public support for major change, Roosevelt's New Deal dramatically increased the size and scope of the federal government, including more executive agencies.: 211–12  The traditionally small presidential staff was greatly expanded, with the Executive Office of the President being created in 1939, none of whom require Senate confirmation.: 229–231  Roosevelt's unprecedented re-election to a third and fourth term, the victory of the United States in World War II, and the nation's growing economy all helped established the office as a position of global leadership.: 269  His successors, Harry Truman and Dwight D. Eisenhower, each served two terms as the Cold War led the presidency to be viewed as the \"leader of the free world\", while John F. Kennedy was a youthful and popular leader who benefited from the rise of television in the 1960s.\nAfter Lyndon B. Johnson lost popular support due to the Vietnam War and Richard Nixon's presidency collapsed in the Watergate scandal, Congress enacted a series of reforms intended to reassert itself. These included the War Powers Resolution, enacted over Nixon's veto in 1973, and the Congressional Budget and Impoundment Control Act of 1974 that sought to strengthen congressional fiscal powers. By 1976, Gerald Ford conceded that \"the historic pendulum\" had swung toward Congress, raising the possibility of a \"disruptive\" erosion of his ability to govern. Ford failed to win election to a full term and his successor, Jimmy Carter, failed to win re-election. Ronald Reagan, who had been an actor before beginning his political career, used his talent as a communicator to help reshape the American agenda away from New Deal policies toward more conservative ideology.\nWith the Cold War ending and the United States becoming the world's undisputed leading power, Bill Clinton, George W. Bush, and Barack Obama each served two terms as president. Meanwhile, Congress and the nation gradually became more politically polarized, especially following the 1994 mid-term elections that saw Republicans control the House for the first time in 40 years, and the rise of routine filibusters in the Senate in recent decades. Recent presidents have thus increasingly focused on executive orders, agency regulations, and judicial appointments to implement major policies, at the expense of legislation and congressional power. Presidential elections in the 21st century have reflected this continuing polarization, with no candidate except Obama in 2008 winning by more than five percent of the popular vote and two, George W. Bush and Donald Trump, winning in the Electoral College while losing the popular vote.\n\nCritics of presidency's evolution\nThe nation's Founding Fathers expected the Congress, which was the first branch of government described in the Constitution, to be the dominant branch of government; however, they did not expect a strong executive department. However, presidential power has shifted over time, which has resulted in claims that the modern presidency has become too powerful, unchecked, unbalanced, and \"monarchist\" in nature. In 2008 professor Dana D. Nelson expressed belief that presidents over the previous thirty years worked towards \"undivided presidential control of the executive branch and its agencies\". She criticized proponents of the unitary executive theory for expanding \"the many existing uncheckable executive powers—such as executive orders, decrees, memorandums, proclamations, national security directives and legislative signing statements—that already allow presidents to enact a good deal of foreign and domestic policy without aid, interference or consent from Congress\". Bill Wilson, board member of Americans for Limited Government, opined that the expanded presidency was \"the greatest threat ever to individual freedom and democratic rule\".\n\nLegislative powers\nArticle I, Section 1 of the Constitution vests all lawmaking power in Congress's hands, and Article 1, Section 6, Clause 2 prevents the president (and all other executive branch officers) from simultaneously being a member of Congress. Nevertheless, the modern presidency exerts significant power over legislation, both due to constitutional provisions and historical developments over time.\n\nSigning and vetoing bills\nThe president's most significant legislative power derives from the Presentment Clause, which gives the president the power to veto any bill passed by Congress. While Congress can override a presidential veto, it requires a two-thirds vote of both houses, which is usually very difficult to achieve except for widely supported bipartisan legislation. The framers of the Constitution feared that Congress would seek to increase its power and enable a \"tyranny of the majority\", so giving the indirectly elected president a veto was viewed as an important check on the legislative power. While George Washington believed the veto should only be used in cases where a bill was unconstitutional, it is now routinely used in cases where presidents have policy disagreements with a bill. The veto – or threat of a veto – has thus evolved to make the modern presidency a central part of the American legislative process.\nSpecifically, under the Presentment Clause, once a bill has been presented by Congress, the president has three options:\n\nSign the legislation within ten days, excluding Sundays, the bill becomes law.\nVeto the legislation within the above timeframe and return it to the house of Congress from which it originated, expressing any objections, the bill does not become law, unless both houses of Congress vote to override the veto by a two-thirds vote.\nTake no action on the legislation within the above timeframe—the bill becomes law, as if the president had signed it, unless Congress is adjourned at the time, in which case it does not become law, which is known as a pocket veto.\nIn 1996, Congress attempted to enhance the president's veto power with the Line Item Veto Act. The legislation empowered the president to sign any spending bill into law while simultaneously striking certain spending items within the bill, particularly any new spending, any amount of discretionary spending, or any new limited tax benefit. Congress could then repass that particular item. If the president then vetoed the new legislation, Congress could override the veto by its ordinary means, a two-thirds vote in both houses. In Clinton v. City of New York, 524 U.S. 417 (1998), the U.S. Supreme Court ruled such a legislative alteration of the veto power to be unconstitutional.\n\nSetting the agenda\nFor most of American history, candidates for president have sought election on the basis of a promised legislative agenda. Article II, Section 3, Clause 2 requires the president to recommend such measures to Congress which the president deems \"necessary and expedient\". This is done through the constitutionally-based State of the Union address, which usually outlines the president's legislative proposals for the coming year, and through other formal and informal communications with Congress.\nThe president can be involved in crafting legislation by suggesting, requesting, or even insisting that Congress enact laws that the president believes are needed. Additionally, the president can attempt to shape legislation during the legislative process by exerting influence on individual members of Congress. Presidents possess this power because the Constitution is silent about who can write legislation, but the power is limited because only members of Congress can introduce legislation.\nThe president or other officials of the executive branch may draft legislation and then ask senators or representatives to introduce these drafts into Congress. Additionally, the president may attempt to have Congress alter proposed legislation by threatening to veto that legislation unless requested changes are made.\n\nPromulgating regulations\nMany laws enacted by Congress do not address every possible detail, and either explicitly or implicitly delegate powers of implementation to an appropriate federal agency. As the head of the executive branch, presidents control a vast array of agencies that can issue regulations with little oversight from Congress.\nIn the 20th century, critics charged that too many legislative and budgetary powers that should have belonged to Congress had slid into the hands of presidents. One critic charged that presidents could appoint a \"virtual army of 'czars'—each wholly unaccountable to Congress yet tasked with spearheading major policy efforts for the White House\". Presidents have been criticized for making signing statements when signing congressional legislation about how they understand a bill or plan to execute it. This practice has been criticized by the American Bar Association as unconstitutional. Conservative commentator George Will wrote of an \"increasingly swollen executive branch\" and \"the eclipse of Congress\".\n\nConvening and adjourning Congress\nTo allow the government to act quickly in case of a major domestic or international crisis arising when Congress is not in session, the president is empowered by Article II, Section 3 of the Constitution to call a special session of one or both houses of Congress. Since John Adams first did so in 1797, the president has called the full Congress to convene for a special session on 27 occasions. Harry S. Truman was the most recent to do so in July 1948, known as the Turnip Day Session. In addition, prior to ratification of the Twentieth Amendment in 1933, which brought forward the date on which Congress convenes from December to January, newly inaugurated presidents would routinely call the Senate to meet to confirm nominations or ratify treaties. In practice, the power has fallen into disuse in the modern era as Congress now formally remains in session year-round, convening pro forma sessions every three days even when ostensibly in recess. Correspondingly, the president is authorized to adjourn Congress if the House and Senate cannot agree on the time of adjournment; no president has ever had to exercise this power.\n\nExecutive powers\nThe president is head of the executive branch of the federal government and is constitutionally obligated to \"take care that the laws be faithfully executed\". The executive branch has over four million employees, including the military.\n\nAdministrative powers\nPresidents make political appointments. An incoming president may make up to 4,000 upon taking office, 1200 of which must be confirmed by the U.S. Senate. Ambassadors, members of the Cabinet, and various officers, are among the positions filled by presidential appointment with Senate confirmation.\nThe power of a president to fire executive officials has long been a contentious political issue. Generally, a president may remove executive officials at will. However, Congress can curtail and constrain a president's authority to fire commissioners of independent regulatory agencies and certain inferior executive officers by statute.\nTo manage the growing federal bureaucracy, presidents have gradually surrounded themselves with many layers of staff, who were eventually organized into the Executive Office of the President of the United States. Within the Executive Office, the president's innermost layer of aides, and their assistants, are located in the White House Office.\nThe president also possesses the power to manage operations of the federal government by issuing various types of directives, such as presidential proclamation and executive orders. When the president is lawfully exercising one of the constitutionally conferred presidential responsibilities, the scope of this power is broad. Even so, these directives are subject to judicial review by U.S. federal courts, which can find them to be unconstitutional. Congress can overturn an executive order through legislation.\n\nForeign affairs\nArticle II, Section 3, Clause 4 requires the president to \"receive Ambassadors.\" This clause, known as the Reception Clause, has been interpreted to imply that the president possesses broad power over matters of foreign policy, and to provide support for the president's exclusive authority to grant recognition to a foreign government. The Constitution also empowers the president to appoint United States ambassadors, and to propose and chiefly negotiate agreements between the United States and other countries. Such agreements, upon receiving the advice and consent of the U.S. Senate (by a two-thirds majority vote), become binding with the force of federal law.\nWhile foreign affairs has always been a significant element of presidential responsibilities, advances in technology since the Constitution's adoption have increased presidential power. Where formerly ambassadors were vested with significant power to independently negotiate on behalf of the United States, presidents now routinely meet directly with leaders of foreign countries.\n\nCommander-in-chief\nOne of the most important of executive powers is the president's role as commander-in-chief of the United States Armed Forces. The power to declare war is constitutionally vested in Congress, but the president has ultimate responsibility for the direction and disposition of the military. The exact degree of authority that the Constitution grants to the president as commander-in-chief has been the subject of much debate throughout history, with Congress at various times granting the president wide authority and at others attempting to restrict that authority. The framers of the Constitution took care to limit the president's powers regarding the military; Alexander Hamilton explained this in Federalist No. 69:The President is to be commander-in-chief of the army and navy of the United States. ... It would amount to nothing more than the supreme command and direction of the military and naval forces ... while that [the power] of the British king extends to the DECLARING of war and to the RAISING and REGULATING of fleets and armies, all [of] which ... would appertain to the legislature. [Emphasis in the original.]\nIn the modern era, pursuant to the War Powers Resolution, Congress must authorize any troop deployments longer than 60 days, although that process relies on triggering mechanisms that have never been employed, rendering it ineffectual. Additionally, Congress provides a check to presidential military power through its control over military spending and regulation. Presidents have historically initiated the process for going to war, but critics have charged that there have been several conflicts in which presidents did not get official declarations, including Theodore Roosevelt's military move into Panama in 1903, the Korean War, the Vietnam War, and the invasions of Grenada in 1983 and Panama in 1989.\nThe amount of military detail handled personally by the president in wartime has varied greatly. George Washington, the first U.S. president, firmly established military subordination under civilian authority. In 1794, Washington used his constitutional powers to assemble 12,000 militia to quell the Whiskey Rebellion, a conflict in Western Pennsylvania involving armed farmers and distillers who refused to pay an excise tax on spirits. According to historian Joseph Ellis, this was the \"first and only time a sitting American president led troops in the field\", though James Madison briefly took control of artillery units in defense of Washington, D.C., during the War of 1812. Abraham Lincoln was deeply involved in overall strategy and in day-to-day operations during the American Civil War, 1861–1865; historians have given Lincoln high praise for his strategic sense and his ability to select and encourage commanders such as Ulysses S. Grant.\nThe present-day operational command of the Armed Forces is delegated to the Department of Defense and is normally exercised through the secretary of defense. The chairman of the Joint Chiefs of Staff and the Combatant Commands assist with the operation as outlined in the presidentially approved Unified Command Plan (UCP).\n\nJuridical powers and privileges\nThe president has the power to nominate federal judges, including members of the United States courts of appeals and the Supreme Court of the United States. However, these nominations require Senate confirmation before they may take office. Securing Senate approval can provide a major obstacle for presidents who wish to orient the federal judiciary toward a particular ideological stance. When nominating judges to U.S. district courts, presidents often respect the long-standing tradition of senatorial courtesy. Presidents may also grant pardons and reprieves. Gerald Ford pardoned Richard Nixon a month after taking office. Presidents often grant pardons shortly before leaving office, like when Bill Clinton pardoned Patty Hearst on his last day in office; this is often controversial.\nTwo doctrines concerning executive power have developed that enable the president to exercise executive power with a degree of autonomy. The first is executive privilege, which allows the president to withhold from disclosure any communications made directly to the president in the performance of executive duties. George Washington first claimed the privilege when Congress requested to see Chief Justice John Jay's notes from an unpopular treaty negotiation with Great Britain. While not enshrined in the Constitution or any other law, Washington's action created the precedent for the privilege. When Nixon tried to use executive privilege as a reason for not turning over subpoenaed evidence to Congress during the Watergate scandal, the Supreme Court ruled in United States v. Nixon, 418 U.S. 683 (1974), that executive privilege did not apply in cases where a president was attempting to avoid criminal prosecution. When Bill Clinton attempted to use executive privilege regarding the Lewinsky scandal, the Supreme Court ruled in Clinton v. Jones, 520 U.S. 681 (1997), that the privilege also could not be used in civil suits. These cases established the legal precedent that executive privilege is valid, although the exact extent of the privilege has yet to be clearly defined. Additionally, federal courts have allowed this privilege to radiate outward and protect other executive branch employees but have weakened that protection for those executive branch communications that do not involve the president.\nThe state secrets privilege allows the president and the executive branch to withhold information or documents from discovery in legal proceedings if such release would harm national security. Precedent for the privilege arose early in the 19th century when Thomas Jefferson refused to release military documents in the treason trial of Aaron Burr and again in Totten v. United States 92 U.S. 105 (1876), when the Supreme Court dismissed a case brought by a former Union spy. However, the privilege was not formally recognized by the U.S. Supreme Court until United States v. Reynolds 345 U.S. 1 (1953), where it was held to be a common law evidentiary privilege. Before the September 11 attacks, use of the privilege had been rare, but increasing in frequency. Since 2001, the government has asserted the privilege in more cases and at earlier stages of the litigation, thus in some instances causing dismissal of the suits before reaching the merits of the claims, as in the Ninth Circuit's ruling in Mohamed v. Jeppesen Dataplan, Inc. Critics of the privilege claim its use has become a tool for the government to cover up illegal or embarrassing government actions.\nThe degree to which the president personally has absolute immunity from court cases is contested and has been the subject of several Supreme Court decisions. Nixon v. Fitzgerald (1982) dismissed a civil lawsuit against by-then former president Richard Nixon based on his official actions. Clinton v. Jones (1997) decided that a president has no immunity against civil suits for actions taken before becoming president and ruled that a sexual harassment suit could proceed without delay, even against a sitting president. The 2019 Mueller report on Russian interference in the 2016 presidential election detailed evidence of possible obstruction of justice, but investigators declined to refer Donald Trump for prosecution based on a United States Department of Justice policy against indicting an incumbent president. The report noted that impeachment by Congress was available as a remedy. As of October 2019, a case was pending in the federal courts regarding access to personal tax returns in a criminal case brought against Donald Trump by the New York County District Attorney alleging violations of New York state law.\n\nLeadership roles\nHead of state\nAs head of state, the president represents the United States government to its own people and represents the nation to the rest of the world. For example, during a state visit by a foreign head of state, the president typically hosts a State Arrival Ceremony held on the South Lawn, a custom begun by John F. Kennedy in 1961. This is followed by a state dinner given by the president which is held in the State Dining Room later in the evening.\n\nAs a national leader, the president also fulfills many less formal ceremonial duties. For example, William Howard Taft started the tradition of throwing out the ceremonial first pitch in 1910 at Griffith Stadium, Washington, D.C., on the Washington Senators's Opening Day. Every president since Taft, except for Jimmy Carter, threw out at least one ceremonial first ball or pitch for Opening Day, the All-Star Game, or the World Series, usually with much fanfare. Every president since Theodore Roosevelt has served as honorary president of the Boy Scouts of America.\nOther presidential traditions are associated with American holidays. Rutherford B. Hayes began in 1878 the first White House egg rolling for local children. Beginning in 1947, during the Harry S. Truman administration, every Thanksgiving the president is presented with a live domestic turkey during the annual National Thanksgiving Turkey Presentation held at the White House. Since 1989, when the custom of \"pardoning\" the turkey was formalized by George H. W. Bush, the turkey has been taken to a farm where it will live out the rest of its natural life.\nPresidential traditions also involve the president's role as head of government. Many outgoing presidents since James Buchanan traditionally give advice to their successor during the presidential transition. Ronald Reagan and his successors have also left a private message on the desk of the Oval Office on Inauguration Day for the incoming president.\nThe modern presidency holds the president as one of the nation's premier celebrities. Some argue that images of the presidency have a tendency to be manipulated by administration public relations officials as well as by presidents themselves. One critic described the presidency as \"propagandized leadership\" which has a \"mesmerizing power surrounding the office\". Administration public relations managers staged carefully crafted photo-ops of smiling presidents with smiling crowds for television cameras. One critic wrote the image of John F. Kennedy was described as carefully framed \"in rich detail\" which \"drew on the power of myth\" regarding the incident of PT 109 and wrote that Kennedy understood how to use images to further his presidential ambitions. As a result, some political commentators have opined that American voters have unrealistic expectations of presidents: voters expect a president to \"drive the economy, vanquish enemies, lead the free world, comfort tornado victims, heal the national soul and protect borrowers from hidden credit-card fees\".\n\nHead of party\nThe president is typically considered to be the head of their political party. Since the entire House of Representatives and at least one-third of the Senate is elected simultaneously with the president, candidates from a political party inevitably have their electoral success intertwined with the performance of the party's presidential candidate. The coattail effect, or lack thereof, will also often impact a party's candidates at state and local levels of government as well. However, there are often tensions between a president and others in the party, with presidents who lose significant support from their party's caucus in Congress generally viewed to be weaker and less effective.\n\nGlobal leader\nWith the rise of the United States as a superpower in the 20th century, and the United States having the world's largest economy into the 21st century, the president is typically viewed as a global leader, and at times the world's most powerful political figure. The position of the United States as the leading member of NATO, and the country's strong relationships with other wealthy or democratic nations like those comprising the European Union, have led to the moniker that the president is the \"leader of the free world\".\n\nSelection process\nEligibility\nArticle II, Section 1, Clause 5 of the Constitution sets three qualifications for holding the presidency. To serve as president, one must:\n\nbe a natural-born citizen of the United States;\nbe at least 35 years old;\nbe a resident in the United States for at least 14 years.\nA person who meets the above qualifications would, however, still be disqualified from holding the office of president under any of the following conditions:\n\nUnder Article I, Section 3, Clause 7, having been impeached, convicted and disqualified from holding further public office, although there is some legal debate as to whether the disqualification clause also includes the presidential office: the only previous persons disqualified under this clause were three federal judges.\nUnder Section 3 of the Fourteenth Amendment, no person who swore an oath to support the Constitution, and later rebelled against the United States, is eligible to hold any office. However, this disqualification can be lifted by a two-thirds vote of each house of Congress. There is, again, some debate as to whether the clause as written allows disqualification from the presidential position, or whether it would first require litigation outside of Congress, although there is precedent for use of this amendment outside of the original intended purpose of excluding Confederates from public office after the Civil War.\nUnder the Twenty-second Amendment, no person can be elected president more than twice. The amendment also specifies that if any eligible person serves as president or acting president for more than two years of a term for which some other eligible person was elected president, the former can only be elected president once.\n\nCampaigns and nomination\nThe modern presidential campaign begins before the primary elections, which the two major political parties use to clear the field of candidates before their national nominating conventions, where the most successful candidate is made the party's presidential nominee. Typically, the party's presidential candidate chooses a vice presidential nominee, and this choice is rubber-stamped by the convention. The most common previous profession of presidents is lawyer.\nNominees participate in nationally televised debates, and while the debates are usually restricted to the Democratic and Republican nominees, third party candidates may be invited, such as Ross Perot in the 1992 debates. Nominees campaign across the country to explain their views, convince voters and solicit contributions. Much of the modern electoral process is concerned with winning swing states through frequent visits and mass media advertising drives.\n\nElection\nThe president is elected indirectly by the voters of each state and the District of Columbia through the Electoral College, a body of electors formed every four years for the sole purpose of electing the president and vice president to concurrent four-year terms. As prescribed by Article II, Section 1, Clause 2, each state is entitled to a number of electors equal to the size of its total delegation in both houses of Congress. Additionally, the Twenty-third Amendment provides that the District of Columbia is entitled to the number it would have if it were a state, but in no case more than that of the least populous state. Currently, all states and the District of Columbia select their electors based on a popular election. In all but two states, the party whose presidential–vice presidential ticket receives a plurality of popular votes in the state has its entire slate of elector nominees chosen as the state's electors. Maine and Nebraska deviate from this winner-take-all practice, awarding two electors to the statewide winner and one to the winner in each congressional district.\nOn the first Monday after the second Wednesday in December, about six weeks after the election, the electors convene in their respective state capitals (and in Washington, D.C.) to vote for president and, on a separate ballot, for vice president. They typically vote for the candidates of the party that nominated them. While there is no constitutional mandate or federal law requiring them to do so, the District of Columbia and 32 states have laws requiring that their electors vote for the candidates to whom they are pledged. The constitutionality of these laws was upheld in Chiafalo v. Washington (2020). Following the vote, each state then sends a certified record of their electoral votes to Congress. The votes of the electors are opened and counted during a joint session of Congress, held in the first week of January. If a candidate has received an absolute majority of electoral votes for president (currently 270 of 538), that person is declared the winner. Otherwise, the House of Representatives must meet to elect a president using a contingent election procedure in which representatives, voting by state delegation, with each state casting a single vote, choose between the top three electoral vote-getters for president. To win the presidency, a candidate must receive the votes of an absolute majority of states (currently 26 of 50).\nThere have been two contingent presidential elections in the nation's history. A 73–73 electoral vote tie between Thomas Jefferson and fellow Democratic-Republican Aaron Burr in the election of 1800 necessitated the first. Conducted under the original procedure established by Article II, Section 1, Clause 3 of the Constitution, which stipulates that if two or three persons received a majority vote and an equal vote, the House of Representatives would choose one of them for president; the runner-up would become vice president. On February 17, 1801, Jefferson was elected president on the 36th ballot, and Burr elected vice president. Afterward, the system was overhauled through the Twelfth Amendment in time to be used in the 1804 election. A quarter-century later, the choice for president again devolved to the House when no candidate won an absolute majority of electoral votes (131 of 261) in the election of 1824. Under the Twelfth Amendment, the House was required to choose a president from among the top three electoral vote recipients: Andrew Jackson, John Quincy Adams, and William H. Crawford. Held February 9, 1825, this second and most recent contingent election resulted in John Quincy Adams being elected president on the first ballot.\n\nInauguration\nPursuant to the Twentieth Amendment, the four-year term of office for both the president and the vice president begins at noon on January 20, in the year following the preceding presidential election. The first presidential and vice presidential terms to begin on this date, known as Inauguration Day, were the second terms of President Franklin D. Roosevelt and Vice President John Nance Garner in 1937. Previously, Inauguration Day was on March 4. As a result of the date change, the first term (1933–37) of both men had been shortened by 43 days.\nBefore executing the powers of the office, a president is required to recite the presidential Oath of Office, found in Article II, Section 1, Clause 8 of the Constitution. This is the only component in the inauguration ceremony mandated by the Constitution:\n\nI do solemnly swear (or affirm) that I will faithfully execute the Office of President of the United States, and will to the best of my ability, preserve, protect, and defend the Constitution of the United States.\nPresidents have traditionally placed one hand upon a Bible while taking the oath, and have added \"So help me God\" to the end of the oath. Although the oath may be administered by any person authorized by law to administer oaths, presidents are traditionally sworn in by the chief justice of the United States.\n\nIncumbency\nTerm limit\nWhen the first president, George Washington, announced in his Farewell Address that he was not running for a third term, he established a \"two terms then out\" precedent. Precedent became tradition after Thomas Jefferson publicly embraced the principle a decade later during his second term, as did his two immediate successors, James Madison and James Monroe. In spite of the strong two-term tradition, Ulysses S. Grant sought nomination at the 1880 Republican National Convention for a non-consecutive third term, but was unsuccessful.\nIn 1940, after leading the nation through the Great Depression and focused on supporting U.S. allied nations at war with the Axis powers, Franklin Roosevelt was elected to a third term, breaking the long-standing precedent. Four years later, with the U.S. engaged in World War II, he was re-elected again despite his declining physical health; he died 82 days into his fourth term on April 12, 1945.\nIn response to the unprecedented length of Roosevelt's presidency, the Twenty-second Amendment was adopted in 1951. The amendment bars anyone from being elected president more than twice, or once if that person served more than two years (24 months) of another president's four-year term. Harry S. Truman, the president at the time it was submitted to the states by the Congress, was exempted from its limitations. Without the exemption, he would not have been eligible to run for a second full term in 1952 (which he briefly sought), as he had served nearly all of Franklin Roosevelt's unexpired 1945–1949 term and had been elected to a full four-year term beginning in 1949. Since becoming operative in 1951, the amendment has been applicable to six twice-elected presidents: Dwight D. Eisenhower, Richard Nixon, Ronald Reagan, Bill Clinton, George W. Bush, and Barack Obama.\n\nVacancies and succession\nUnder Section 1 of the Twenty-fifth Amendment, ratified in 1967, the vice president becomes president upon the removal from office, death, or resignation of the president. Deaths have occurred a number of times, resignation has occurred only once, and removal from office has never occurred.\nBefore the ratification of the Twenty-fifth amendment (which clarified the matter of succession), Article II, Section 1, Clause 6, stated only that the vice president assumes the \"powers and duties\" of the presidency in the event of a president's removal, death, resignation, or inability. Under this clause, there was ambiguity about whether the vice president would actually become president in the event of a vacancy, or simply act as president, potentially resulting in a special election. Upon the death of President William Henry Harrison in 1841, Vice President John Tyler declared that he had succeeded to the office itself, refusing to accept any papers addressed to the \"Acting President\", and Congress ultimately accepted it.\nIn the event of a double vacancy, Article II, Section 1, Clause 6 also authorizes Congress to declare who shall become acting president in the \"Case of Removal, Death, Resignation or Inability, both of the president and vice president\". The Presidential Succession Act of 1947 (codified as 3 U.S.C. § 19) provides that if both the president and vice president have left office or are both otherwise unavailable to serve during their terms of office, the presidential line of succession follows the order of: speaker of the House, then, if necessary, the president pro tempore of the Senate, and then if necessary, the eligible heads of federal executive departments who form the president's cabinet. The cabinet currently has 15 members, of which the secretary of state is first in line; the other Cabinet secretaries follow in the order in which their department (or the department of which their department is the successor) was created. Those individuals who are constitutionally ineligible to be elected to the presidency are also disqualified from assuming the powers and duties of the presidency through succession. No statutory successor has yet been called upon to act as president.\n\nDeclarations of inability\nUnder the Twenty-fifth Amendment, the president may temporarily transfer the presidential powers and duties to the vice president, who then becomes acting president, by transmitting to the speaker of the House and the president pro tempore of the Senate a statement that he is unable to discharge his duties. The president resumes his or her powers upon transmitting a second declaration stating that he is again able. The mechanism has been used by Ronald Reagan (once), George W. Bush (twice), and Joe Biden (once), each in anticipation of surgery.\nThe Twenty-fifth Amendment also provides that the vice president, together with a majority of certain members of the Cabinet, may transfer the presidential powers and duties to the vice president by transmitting a written declaration, to the speaker of the House and the president pro tempore of the Senate, to the effect that the president is unable to discharge his or her powers and duties. If the president then declares that no such inability exist, he or she resumes the presidential powers unless the vice president and Cabinet make a second declaration of presidential inability, in which case Congress decides the question.\n\nRemoval\nArticle II, Section 4 of the Constitution allows for the removal of high federal officials, including the president, from office for \"treason, bribery, or other high crimes and misdemeanors\". Article I, Section 2, Clause 5 authorizes the House of Representatives to serve as a \"grand jury\" with the power to impeach said officials by a majority vote. Article I, Section 3, Clause 6 authorizes the Senate to serve as a court with the power to remove impeached officials from office, by a two-thirds vote to convict.\nThree presidents have been impeached by the House of Representatives: Andrew Johnson in 1868, Bill Clinton in 1998, and Donald Trump in 2019 and 2021; none have been convicted by the Senate. Additionally, the House Judiciary Committee conducted an impeachment inquiry against Richard Nixon in 1973–74 and reported three articles of impeachment to the House of Representatives for final action; however, he resigned from office before the House voted on them.\n\nCircumvention of authority\nControversial measures have sometimes been taken short of removal to deal with perceived recklessness on the part of the president, or with a long-term disability. In some cases, staff have intentionally failed to deliver messages to or from the president, typically to avoid executing or promoting the president to write certain orders. This has ranged from Richard Nixon's Chief of Staff not transmitting orders to the Cabinet due to the president's heavy drinking, to staff removing memos from Donald Trump's desk. Decades before the Twenty-fifth Amendment, in 1919, President Woodrow Wilson had a stroke that left him partly incapacitated. First lady Edith Wilson kept this condition a secret from the public for a while, and controversially became the sole gatekeeper for access to the president (aside from his doctor), assisting him with paperwork and deciding which information was \"important\" enough to share with him.\n\nCompensation\nSince 2001, the president's annual salary has been $400,000, along with a: $50,000 expense allowance; $100,000 nontaxable travel account, and $19,000 entertainment account. The president's salary is set by Congress, and under Article II, Section 1, Clause 7 of the Constitution, any increase or reduction in presidential salary cannot take effect before the next presidential term of office.\n\nResidence\nThe Executive Residence of the White House in Washington, D.C. is the official residence of the president. The site was selected by George Washington, and the cornerstone was laid in 1792. Every president since John Adams (in 1800) has lived there. At various times in U.S. history, it has been known as the \"President's Palace\", the \"President's House\", and the \"Executive Mansion\". Theodore Roosevelt officially gave the White House its current name in 1901. The federal government pays for state dinners and other official functions, but the president pays for personal, family, and guest dry cleaning and food.\nCamp David, officially titled Naval Support Facility Thurmont, a mountain-based military camp in Frederick County, Maryland, is the president's country residence. A place of solitude and tranquility, the site has been used extensively to host foreign dignitaries since the 1940s.\nPresident's Guest House, located next to the Eisenhower Executive Office Building at the White House Complex and Lafayette Park, serves as the president's official guest house and as a secondary residence for the president if needed. Four interconnected, 19th-century houses—Blair House, Lee House, and 700 and 704 Jackson Place—with a combined floor space exceeding 70,000 square feet (6,500 m2) comprise the property.\n\n\tPresidential residences\n\nTravel\nThe primary means of long-distance air travel for the president is one of two identical Boeing VC-25 aircraft, which are extensively modified Boeing 747 airliners and are referred to as Air Force One while the president is on board (although any U.S. Air Force aircraft the president is aboard is designated as \"Air Force One\" for the duration of the flight). In-country trips are typically handled with just one of the two planes, while overseas trips are handled with both, one primary and one backup. The president also has access to smaller Air Force aircraft, most notably the Boeing C-32, which are used when the president must travel to airports that cannot support a jumbo jet. Any civilian aircraft the president is aboard is designated Executive One for the flight.\nFor short-distance air travel, the president has access to a fleet of U.S. Marine Corps helicopters of varying models, designated Marine One when the president is aboard any particular one in the fleet. Flights are typically handled with as many as five helicopters all flying together and frequently swapping positions as to disguise which helicopter the president is actually aboard to any would-be threats.\nFor ground travel, the president uses the presidential state car, which is an armored limousine designed to look like a Cadillac sedan, but built on a truck chassis. The U.S. Secret Service operates and maintains the fleet of several limousines. The president also has access to two armored motorcoaches, which are primarily used for touring trips.\n\n\tPresidential transportation\n\nProtection\nThe U.S. Secret Service is charged with protecting the president and the first family. As part of their protection, presidents, first ladies, their children and other immediate family members, and other prominent persons and locations are assigned Secret Service codenames. The use of such names was originally for security purposes and dates to a time when sensitive electronic communications were not routinely encrypted; today, the names simply serve for purposes of brevity, clarity, and tradition.\n\nPost-presidency\nActivities\nSome former presidents have had significant careers after leaving office. Prominent examples include William Howard Taft's tenure as chief justice of the United States and Herbert Hoover's work on government reorganization after World War II. Grover Cleveland, whose bid for reelection failed in 1888, was elected president again four years later in 1892. Two former presidents served in Congress after leaving the White House: John Quincy Adams was elected to the House of Representatives, serving there for 17 years, and Andrew Johnson returned to the Senate in 1875, though he died soon after. Some ex-presidents were very active, especially in international affairs, most notably Theodore Roosevelt; Herbert Hoover; Richard Nixon; and Jimmy Carter.\nPresidents may use their predecessors as emissaries to deliver private messages to other nations or as official representatives of the United States to state funerals and other important foreign events. Richard Nixon made multiple foreign trips to countries including China and Russia and was lauded as an elder statesman. Jimmy Carter has become a global human rights campaigner, international arbiter, and election monitor, as well as a recipient of the Nobel Peace Prize. Bill Clinton has also worked as an informal ambassador, most recently in the negotiations that led to the release of two American journalists, Laura Ling and Euna Lee, from North Korea. During his presidency, George W. Bush called on former Presidents Bush and Clinton to assist with humanitarian efforts after the 2004 Indian Ocean earthquake and tsunami. President Obama followed suit by asking Presidents Clinton and Bush to lead efforts to aid Haiti after an earthquake devastated that country in 2010.\nClinton was active politically since his presidential term ended, working with his wife Hillary on her 2008 and 2016 presidential bids and President Obama on his 2012 reelection campaign. Obama was also active politically since his presidential term ended, having worked with his former vice president Joe Biden on his 2020 election campaign. Trump has continued to make appearances in the media and at conferences and rallies since leaving office in 2021. He is currently running for a non-consecutive second term in the upcoming 2024 presidential election.\n\nPension and other benefits\nThe Former Presidents Act (FPA), enacted in 1958, grants lifetime benefits to former presidents and their widows, including a monthly pension, medical care in military facilities, health insurance, and Secret Service protection; also provided is funding for a certain number of staff and for office expenses. The act has been amended several times to provide increases in presidential pensions and in the allowances for office staff. The FPA excludes any president who was removed from office by impeachment.\nAccording to a 2008 report by the Congressional Research Service:\n\nChief executives leaving office prior to 1958 often entered retirement pursuing various occupations and received no federal assistance. When industrialist Andrew Carnegie announced a plan in 1912 to offer $25,000 annual pensions to former Presidents, many Members of Congress deemed it inappropriate that such a pension would be provided by a private corporation executive. That same year, legislation was first introduced to create presidential pensions, but it was not enacted. In 1955, such legislation was considered by Congress because of former President Harry S. Truman's financial limitations in hiring an office staff\nThe pension has increased numerous times with congressional approval. Retired presidents receive a pension based on the salary of the current administration's cabinet secretaries, which was $199,700 per year in 2012. Former presidents who served in Congress may also collect congressional pensions. The act also provides former presidents with travel funds and franking privileges.\nPrior to 1997, all former presidents, their spouses, and their children until age 16 were protected by the Secret Service until the president's death. In 1997, Congress passed legislation limiting Secret Service protection to no more than 10 years from the date a president leaves office. On January 10, 2013, President Obama signed legislation reinstating lifetime Secret Service protection for him, George W. Bush, and all subsequent presidents. A first spouse who remarries is no longer eligible for Secret Service protection.\n\nPresidential libraries\nEvery president since Herbert Hoover has created a repository known as a presidential library for preserving and making available his papers, records, and other documents and materials. Completed libraries are deeded to and maintained by the National Archives and Records Administration (NARA); the initial funding for building and equipping each library must come from private, non-federal sources. There are currently thirteen presidential libraries in the NARA system. There are also presidential libraries maintained by state governments and private foundations and Universities of Higher Education, including:\n\nThe Abraham Lincoln Presidential Library and Museum, which is run by the State of Illinois;\nThe George W. Bush Presidential Library and Museum, which is run by Southern Methodist University;\nThe George H. W. Bush Presidential Library and Museum, which is run by Texas A&M University; and\nThe Lyndon Baines Johnson Presidential Library and Museum, which is run by the University of Texas at Austin.\nSeveral former presidents have overseen the building and opening of their own presidential libraries. Some even made arrangements for their own burial at the site. Several presidential libraries contain the graves of the president they document: \n\nThe Harry S. Truman Presidential Library and Museum in Independence, Missouri;\nThe Dwight D. Eisenhower Presidential Library, Museum and Boyhood Home in Abilene, Kansas;\nThe Richard Nixon Presidential Library and Museum in Yorba Linda, California; and\nThe Ronald Reagan Presidential Library and Museum in Simi Valley, California.\nThese gravesites are open to the general public.\n\nPolitical affiliation\nPolitical parties have dominated American politics for most of the nation's history. Though the Founding Fathers generally spurned political parties as divisive and disruptive, and their rise had not been anticipated when the U.S. Constitution was drafted in 1787, organized political parties developed in the U.S. in the mid-1790s nonetheless. They evolved from political factions, which began to appear almost immediately after the Federal government came into existence. Those who supported the Washington administration were referred to as \"pro-administration\" and would eventually form the Federalist Party, while those in opposition largely joined the emerging Democratic-Republican Party.\nGreatly concerned about the very real capacity of political parties to destroy the fragile unity holding the nation together, Washington remained unaffiliated with any political faction or party throughout his eight-year presidency. He was, and remains, the only U.S. president never to be affiliated with a political party. Since Washington, every U.S. president has been affiliated with a political party at the time of assuming office.\nThe number of presidents per political party by their affiliation at the time they were first sworn into office (alphabetical, by last name) are:\n\nTimeline of presidents\nThe following timeline depicts the progression of the presidents and their political affiliation at the time of assuming office.\n\nSee also\nOutline of American politics\n\nNotes\nReferences\nFurther reading\nExternal links\n\nWhite House homepage\nUnited States Presidents Collection. General Collection, Beinecke Rare Book and Manuscript Library, Yale University\n\nJames Buchanan Jr. ( bew-KAN-ən; April 23, 1791 – June 1, 1868) was the 15th president of the United States, serving from 1857 to 1861. Buchanan also served as the secretary of state from 1845 to 1849 and represented Pennsylvania in both houses of the U.S. Congress. He was an advocate for states' rights, particularly regarding slavery, and minimized the role of the federal government preceding the Civil War.\nBuchanan was a lawyer in Pennsylvania and won his first election to the state's House of Representatives as a Federalist. He was elected to the U.S. House of Representatives in 1820 and retained that post for five terms, aligning with Andrew Jackson's Democratic Party. Buchanan served as Jackson's minister to Russia in 1832. He won the election in 1834 as a U.S. senator from Pennsylvania and continued in that position for 11 years. He was appointed to serve as President James K. Polk's secretary of state in 1845, and eight years later was named as President Franklin Pierce's minister to the United Kingdom.\nBeginning in 1844, Buchanan became a regular contender for the Democratic Party's presidential nomination. He was nominated and won the 1856 presidential election. As President, Buchanan intervened to assure the Supreme Court's majority ruling in the pro-slavery decision in the Dred Scott case. He acceded to Southern attempts to engineer Kansas' entry into the Union as a slave state under the Lecompton Constitution, and angered not only Republicans but also Northern Democrats. Buchanan honored his pledge to serve only one term and supported Breckinridge's unsuccessful candidacy in the 1860 presidential election. He failed to reconcile the fractured Democratic Party amid the grudge against Stephen Douglas, leading to the election of Republican and former Congressman Abraham Lincoln.\nBuchanan's leadership during his lame duck period, before the American Civil War, has been widely criticized. He simultaneously angered the North by not stopping secession and the South by not yielding to their demands. He supported the Corwin Amendment in an effort to reconcile the country. He made an unsuccessful attempt to reinforce Fort Sumter, but otherwise refrained from preparing the military. In his personal life, Buchanan never married and was the only U.S. president to remain a lifelong bachelor, leading some historians and authors to question his sexual orientation. His failure to forestall the Civil War has been described as incompetence, and he spent his last years defending his reputation. Historians and scholars rank Buchanan as among the worst presidents in American history.\n\nEarly life\nChildhood and education\nJames Buchanan Jr. was born into a Scottish-Irish family on April 23, 1791, in a log cabin on a farm called Stony Batter, near Cove Gap, Peters Township, in the Allegheny Mountains of southern Pennsylvania. He was the last president born in the 18th century and, until the election of Joe Biden in 2020, the only one born in Pennsylvania. Buchanan was the second of eleven children with six sisters and four brothers, and the eldest son of James Buchanan Sr. (1761–1821) and his wife Elizabeth Speer (1767–1833). James Buchanan Sr., was an Ulster-Scot from just outside Ramelton, a small town in the north-east of County Donegal in the north-west of Ulster, the northern province in Ireland, who emigrated to the newly formed United States in 1783, having sailed from Derry. He belonged to the Clan Buchanan, whose members had emigrated in large numbers from the Scottish Highlands to Ulster in the north of Ireland during the Plantation of Ulster in the seventeenth century and, later, largely because of poverty and persecution by the Crown due to their Presbyterian faith, had further emigrated in large numbers from Ulster to America from the early eighteenth century onwards. Shortly after Buchanan's birth, the family relocated to a farm near Mercersburg, Pennsylvania, and later settled in the town in 1794. His father became the area's wealthiest resident, working as a merchant, farmer, and real estate investor. Buchanan attributed his early education primarily to his mother, whereas his father had a greater influence on his character. His mother had discussed politics with him as a child and had an interest in poetry, quoting John Milton and William Shakespeare to Buchanan.\nBuchanan attended the Old Stone Academy in Mercersburg and then Dickinson College in Carlisle, Pennsylvania. In 1808, he was nearly expelled for disorderly conduct; he and his fellow students had attracted negative attention for drinking in local taverns, disturbing the peace at night and committing acts of vandalism, but he pleaded for a second chance and ultimately graduated with honors in 1809. Later that year, he moved to the state capital at Lancaster, to train as a lawyer for two and a half years with the well-known James Hopkins. Following the fashion of the time, Buchanan studied the United States Code and the Constitution of the United States as well as legal authorities such as William Blackstone during his education.\n\nEarly law practice and Pennsylvania House of Representatives\nIn 1812, Buchanan passed the bar exam and after being admitted to the bar, he remained in Lancaster, even when Harrisburg became the new capital of Pennsylvania. Buchanan quickly established himself as a prominent legal representative in the city. His income rapidly rose after he established his practice, and by 1821 he was earning over $11,000 per year (equivalent to $250,000 in 2023). At this time, Buchanan became a Freemason, and served as the Worshipful Master of Masonic Lodge No. 43 in Lancaster and as a District Deputy Grand Master of the Grand Lodge of Pennsylvania.\nBuchanan also served as chairman of the Lancaster chapter of the Federalist Party. Like his father, he supported their political program, which provided federal funds for building projects and import duties as well as the re-establishment of a central bank after the First Bank of the United States' license expired in 1811. He became a strong critic of Democratic-Republican President James Madison during the War of 1812. Although he did not himself serve in a militia during the War of 1812, during the British occupation he joined a group of young men who stole horses for the United States Army in the Baltimore area. He was the last president involved in the War of 1812.\nIn 1814, he was elected for the Federalists to the Pennsylvania House of Representatives, where he was the youngest member, and held this seat until 1816. Since the sessions in the Pennsylvania General Assembly lasted only three months, Buchanan continued practicing law at a profit by charging higher fees, and his service helped him acquire more clients. In 1815, Buchanan defended District Judge Walter Franklin in an impeachment trial before the Pennsylvania Senate, over alleged judicial misconduct. Impeachments were more common at the time because the line between abuse of office and a wrong legal decision was determined by the ruling parties' preferences and the popularity of the judge's decision. Buchanan persuaded the senators that only judicial crimes and clear violations of the law justified impeachment.\n\nCongressional career\nU.S. House of Representatives\nIn the congressional elections of 1820, Buchanan ran for a seat in the House of Representatives. Shortly after his election victory, his father died in a carriage accident. As a young Representative, Buchanan was one of the most prominent leaders of the \"Amalgamator party\" faction of Pennsylvanian politics, named that because it was made up of both Democratic-Republicans and former Federalists, which transitioned from the First Party System to the Era of Good Feelings. During this era, the Democratic-Republicans became the most influential party. Buchanan's Federalist convictions were weak, and he switched parties after opposing a nativist Federalist bill. During the 1824 presidential election, Buchanan initially supported Henry Clay, but switched to Andrew Jackson (with Clay as a second choice) when it became clear that the Pennsylvanian public overwhelmingly preferred Jackson. After Jackson lost the 1824 election, he joined his faction, but Jackson had contempt for Buchanan due to his misinterpretation of his efforts to mediate between the Clay and Jackson camps.\nIn Washington, Buchanan became an avid defender of states' rights, and was close with many southern Congressmen, viewing some New England Congressmen as dangerous radicals. Buchanan's close proximity to his constituency allowed him to establish a Democratic coalition in Pennsylvania, consisting of former Federalist farmers, Philadelphia artisans, and Ulster-Scots-Americans. In the 1828 presidential election, he secured Pennsylvania, while the \"Jacksonian Democrats\", an independent party after splitting from the National Republican Party, won an easy victory in the parallel congressional election.\n\nBuchanan gained most attention during an impeachment trial where he acted as prosecutor for federal district judge James H. Peck; however, the Senate rejected Buchanan's plea and acquitted Peck by a majority vote. He was appointed to the Agriculture Committee in his first year, and he eventually became chairman of the Judiciary Committee. In 1831, Buchanan declined a nomination for the 22nd United States Congress from his constituency consisting of Dauphin, Lebanon, and Lancaster counties. He still had political ambitions and some Pennsylvania Democrats put him forward as a candidate for the vice presidency in the 1832 election.\n\nMinister to Russia\nAfter Jackson was re-elected in 1832, he offered Buchanan the position of United States Ambassador to Russia. Buchanan was reluctant to leave the country, as the distant St. Petersburg was a kind of political exile, which was the intention of Jackson, who considered Buchanan to be an \"incompetent busybody\" and untrustworthy, but he ultimately agreed. His work focused on concluding a trade and shipping treaty with Russia. While Buchanan was successful with the former, negotiating an agreement on free merchant shipping with Foreign Minister Karl Nesselrode proved difficult. He had denounced Tsar Nicholas I as a despot merely a year prior during his tenure in Congress; many Americans had reacted negatively to Russia's reaction to the 1830 Polish uprising.\n\nU.S. Senator\nBuchanan returned home and lost the election in the State Legislature for a full six-year term in the 23rd Congress, but was appointed by the Pennsylvania state legislature to succeed William Wilkins in the U.S. Senate. Wilkins, in turn, replaced Buchanan as the ambassador to Russia. The Jacksonian Buchanan, who was re-elected in 1836 and 1842, opposed the re-chartering of the Second Bank of the United States and sought to expunge a congressional censure of Jackson stemming from the Bank War. Buchanan served in the Senate until March 1845 and was twice confirmed in office. To unite Pennsylvania Democrats at the State Convention, he was chosen as their candidate for the National Convention. Buchanan maintained a strict adherence to the Pennsylvania State Legislature's guidelines and sometimes voted against positions in Congress which he promoted in his own speeches, despite open ambitions for the White House.\nBuchanan was known for his commitment to states' rights and the Manifest Destiny ideology. He rejected President Martin Van Buren's offer to become United States Attorney General and chaired prestigious Senate committees such as the Committee on the Judiciary and the Committee on Foreign Relations. Buchanan was one of only a few senators to vote against the Webster–Ashburton Treaty for its \"surrender\" of lands to the United Kingdom, as he demanded the entire Aroostook River Valley for the United States. In the Oregon Boundary Dispute, Buchanan adopted the maximum demand of 54°40′ as the northern border and spoke out in favor of annexing the Republic of Texas. During the contentious 1838 Pennsylvania gubernatorial election, Buchanan chose to support the Democratic challenger, David Rittenhouse Porter, who was elected by fewer than 5,500 votes as Pennsylvania's first governor under the state's revised Constitution of 1838.\nBuchanan also opposed a gag rule sponsored by John C. Calhoun that would have suppressed anti-slavery petitions. He joined the majority in blocking the rule, with most senators of the belief that it would have the reverse effect of strengthening the abolitionists. He said, \"We have just as little right to interfere with slavery in the South, as we have to touch the right of petition.\" Buchanan thought that the issue of slavery was the domain of the states, and he faulted abolitionists for exciting passions over the issue. In the lead-up to the 1844 Democratic National Convention, Buchanan positioned himself as a potential alternative to former President Martin Van Buren, but the nomination went to James K. Polk, who won the election.\n\nDiplomatic career\nSecretary of State\nBuchanan was offered the position of Secretary of State in the Polk administration or, as the alternative, a seat on the Supreme Court, to compensate him for his support in the election campaign but also in order to eliminate him as an internal party rival. He accepted the State Department post and served for the duration of Polk's single term in office. During his tenure, the United States recorded its largest territorial gain in history through the Oregon Treaty and the Treaty of Guadalupe Hidalgo, which included territory that is now Texas, California, Nevada, New Mexico, Arizona, Utah, and Colorado. In negotiations with Britain over Oregon, Buchanan initially favored the 49th parallel as the boundary of Oregon Territory, while Polk called for a more northerly boundary line. When Northern Democrats rallied around the popular slogan Fifty-Four Forty or Fight (\"54°40′ or war\") in the 1844 election campaign, Buchanan adopted this position, but later followed Polk's direction, leading to the Oregon Compromise of 1846, which established the 49th parallel as the boundary in the Pacific Northwest.\nIn regards to Mexico, Buchanan maintained a dubious view that its attack on American troops on the other side of the Rio Grande in April 1846 constituted a border violation and a legitimate reason for war. During the Mexican-American War, Buchanan initially advised against claiming territory south of the Rio Grande, fearing war with Britain and France. However, as the war came to an end, Buchanan changed his mind and argued for the annexation of further territory, arguing that Mexico was to blame for the war and that the compensation negotiated for the American losses was too low. Buchanan sought the nomination at the 1848 Democratic National Convention, as Polk had promised to serve only one term, but he only won the support of the Pennsylvania and Virginia delegations, so Senator Lewis Cass of Michigan was nominated.\n\nCivilian life and 1852 presidential election\nWith the 1848 election of Whig Zachary Taylor, Buchanan returned to private life. Buchanan was getting on in years and still dressed in the old-fashioned style of his adolescence, earning him the nickname \"Old Public Functionary\" from the press. Slavery opponents in the North mocked him as a relic of prehistoric man because of his moral values. He bought the house of Wheatland on the outskirts of Lancaster and entertained various visitors while monitoring political events. During this period, Buchanan became the center of a family network consisting of 22 nieces, nephews and their descendants, seven of whom were orphans. He found public service jobs for some through patronage, and for those in his favor, he took on the role of surrogate father. He formed the strongest emotional bond with his niece Harriet Lane, who later became First Lady for Buchanan in the White House.\nIn 1852, he was named president of the Board of Trustees of Franklin and Marshall College in Lancaster, and he served in this capacity until 1866. Buchanan did not completely leave politics. He intended to publish a collection of speeches and an autobiography, but his political comeback was thwarted by the 1852 presidential election. Buchanan traveled to Washington to discuss Pennsylvania Democratic Party politics, which were divided into two camps led by Simon Cameron and George Dallas. He quietly campaigned for the 1852 Democratic presidential nomination. In light of the Compromise of 1850, which had led to the admission of California into the Union as a free state and a stricter Fugitive Slave Act, Buchanan now rejected the Missouri Compromise and welcomed Congress's rejection of the Wilmot Proviso, which prohibited slavery in all territories gained in the Mexican-American War. Buchanan criticized abolitionism as a fanatical attitude and believed that slavery should be decided by state legislatures, not Congress. He disliked abolitionist Northerners due to his party affiliation, and became known as a \"doughface\" due to his sympathy toward the South. Buchanan emerged as a promising candidate for the Democratic presidential nomination, alongside Lewis Cass, Stephen Douglas, and William L. Marcy; however, the Pennsylvania convention did not vote unanimously in his favor, with over 30 delegates protesting against him. At the 1852 Democratic National Convention, he won the support of many southern delegates but failed to win the two-thirds support needed for the presidential nomination, which went to Franklin Pierce. Buchanan declined to serve as the vice presidential nominee, and the convention instead nominated his close friend, William R. King.\n\nMinister to the United Kingdom\nPierce won the election in 1852, and six months later, Buchanan accepted the position of United States Minister to the United Kingdom, a position that represented a step backward in his career and that he had twice previously rejected. Buchanan sailed for England in the summer of 1853, and he remained abroad for the next three years. In 1850, the United States and Great Britain signed the Clayton–Bulwer Treaty, which committed both countries to joint control of any future canal that would connect the Atlantic and Pacific Oceans through Central America. Buchanan met repeatedly with Lord Clarendon, the British foreign minister, in hopes of pressuring the British to withdraw from Central America. He was able to reduce British influence in Honduras and Nicaragua while also raising the kingdom's awareness of American interests in the region. He also focused on the potential annexation of Cuba, which had long interested him.\nAt Pierce's prompting, Buchanan met in Ostend, Belgium, with U.S. Ambassador to Spain Pierre Soulé and U.S. Ambassador to France John Mason, to work out a plan for the acquisition of Cuba. A memorandum draft resulted, called the Ostend Manifesto, which proposed the purchase of Cuba from Spain, then in the midst of revolution and near bankruptcy. The document declared the island \"as necessary to the North American republic as any of its present ... family of states\". Against Buchanan's recommendation, the final draft of the manifesto suggested that \"wresting it from Spain\", if Spain refused to sell, would be justified \"by every law, human and Divine\". The manifesto was met with a divided response and was never acted upon. It weakened the Pierce administration and reduced support for Manifest Destiny. In 1855, as Buchanan's desire to return home grew, Pierce asked him to hold the fort in London in light of the relocation of a British fleet to the Caribbean.\n\nElection of 1856\nBuchanan's service abroad allowed him to conveniently avoid the debate over the Kansas–Nebraska Act then roiling the country in the slavery dispute. While he did not overtly seek the presidency, he assented to the movement on his behalf. While still in England, he campaigned by praising John Joseph Hughes, who was Archbishop of New York, to a Catholic archbishop. The latter campaigned for Buchanan among high-ranking Catholics as soon as he heard about it. When Buchanan arrived home at the end of April 1856, he led on the first ballot, supported by powerful Senators John Slidell, Jesse Bright, and Thomas F. Bayard, who presented Buchanan as an experienced leader appealing to the North and South. The 1856 Democratic National Convention met in June 1856, producing a platform that reflected Buchanan's views, including support for the Fugitive Slave Law, which required the return of escaped slaves. The platform also called for an end to anti-slavery agitation and U.S. \"ascendancy in the Gulf of Mexico\". President Pierce hoped for re-nomination, while Senator Stephen A. Douglas also loomed as a strong candidate. He won the nomination after seventeen ballots after Douglas' resignation. He was joined on the ticket by John C. Breckinridge of Kentucky in order to maintain regional proportional representation, placating supporters of Pierce and Douglas, also allies of Breckinridge.\nBuchanan faced two candidates in the general election: former Whig President Millard Fillmore ran as the candidate for the anti-Catholic, anti-immigrant American Party (or \"Know-Nothing\"), while John C. Frémont ran as the Republican nominee. The contrast between Buchanan and Frémont was particularly stark, with opposing caricaturists drawing the Democratic candidate as a fussy old man in drag. Buchanan did not actively campaign, but he wrote letters and pledged to uphold the Democratic platform. In the election, he carried every slave state except for Maryland, as well as five slavery-free states, including his home state of Pennsylvania. He won 45 percent of the popular vote and decisively won the electoral vote, taking 174 of 296 votes. His election made him the first president from Pennsylvania. In a combative victory speech, Buchanan denounced Republicans, calling them a \"dangerous\" and \"geographical\" party that had unfairly attacked the South. He also declared, \"the object of my administration will be to destroy sectional party, North or South, and to restore harmony to the Union under a national and conservative government.\" He set about this initially by feigning a sectional balance in his cabinet appointments.\n\nPresidency (1857–1861)\nInauguration\nBuchanan was inaugurated on March 4, 1857, taking the oath of office from Chief Justice Roger B. Taney. In his lengthy inaugural address, Buchanan committed himself to serving only one term, as his predecessor had done. He abhorred the growing divisions over slavery and its status in the territories, saying that Congress should play no role in determining the status of slavery in the states or territories. He proposed a solution based on the Kansas-Nebraska Act, which stated that the principle of popular sovereignty was decisive, and Congress had no say in the matter. Buchanan recommended that a federal slave code be enacted to protect the rights of slaveowners in federal territories. He alluded to a then-pending Supreme Court case, Dred Scott v. Sandford, which he said would permanently settle the issue of slavery. Dred Scott was a slave who was temporarily taken from a slave state to a free territory by his owner, John Sanford. After Scott returned to the slave state, he filed a petition for his freedom based on his time in the free territory.\n\nAssociate Justice Robert C. Grier leaked the decision in the \"Dred Scott\" case early to Buchanan. In his inaugural address, Buchanan declared that the issue of slavery in the territories would be \"speedily and finally settled\" by the Supreme Court. According to historian Paul Finkelman: Buchanan already knew what the Court was going to decide. In a major breach of Court etiquette, Justice Grier, who, like Buchanan, was from Pennsylvania, had kept the President-elect fully informed about the progress of the case and the internal debates within the Court. When Buchanan urged the nation to support the decision, he already knew what Taney would say. Republican suspicions of impropriety turned out to be fully justified.\nHistorians agree that the court decision was a major disaster because it dramatically inflamed tensions, leading to the Civil War. In 2022, historian David W. Blight argued that the year 1857 was, \"the great pivot on the road to disunion...largely because of the Dred Scott case, which stoked the fear, distrust and conspiratorial hatred already common in both the North and the South to new levels of intensity.\"\n\nPersonnel\nCabinet and administration\nAs his inauguration approached, Buchanan sought to establish an obedient, harmonious cabinet to avoid the in-fighting that had plagued Andrew Jackson's administration. The cabinet's composition had to do justice to the proportional representation within the party and between the regions of the country. Buchanan first worked on this task in Wheatland until he traveled to the capital in January 1857. There, like many other guests at the National Hotel, he contracted severe dysentery, from which he did not fully recover until several months later. Dozens of those who fell ill died, including Buchanan's nephew and private secretary Eskridge Lane.\nThe cabinet selection was disastrous, with four Southern ministers being large-scale slaveholders who later became loyal to the Confederate States of America. Secretary of the Treasury Howell Cobb was considered the greatest political talent in the Cabinet, while the three department heads from the northern states were all considered to be doughfaces. His objective was to dominate the cabinet, and he chose men who would agree with his views. Buchanan had a troubled relationship with his vice president from the beginning, when he did not receive him during his inaugural visit but referred him to his niece and First Lady, which Breckinridge never forgave him for and saw as disrespectful. He left out the influential Stephen A. Douglas, who had made Buchanan's nomination possible by resigning at the National Convention the previous year, when filling the post. Concentrating on foreign policy, he appointed the aging Lewis Cass as Secretary of State. Buchanan's appointment of Southerners and their allies alienated many in the North, and his failure to appoint any followers of Douglas divided the party. Outside of the cabinet, he left in place many of Pierce's appointments but removed a disproportionate number of Northerners who had ties to Democratic opponents Pierce or Douglas.\n\nJudicial appointments\nBuchanan appointed one Justice, Nathan Clifford, to the Supreme Court of the United States. He appointed seven other federal judges to United States district courts. He also appointed two judges to the United States Court of Claims.\n\nIntervention in the Dred Scott case\nThe case of Dred Scott v. Sandford, to which Buchanan referred to in his inaugural address, dated back to 1846. Scott sued for his release in Missouri, claiming he lived in service to the proprietor in Illinois and Wisconsin Territory. The case reached the Supreme Court and gained national attention by 1856. Buchanan consulted with Judge John Catron in January 1857, inquiring about the outcome of the case and suggesting that a broader decision, beyond the specifics of the case, would be more prudent. Buchanan hoped that a broad decision protecting slavery in the territories could lay the issue to rest, allowing him to focus on other issues.\nCatron replied on February 10, saying that the Supreme Court's Southern majority would decide against Scott, but would likely have to publish the decision on narrow grounds unless Buchanan could convince his fellow Pennsylvanian, Justice Robert Cooper Grier, to join the majority of the court. Buchanan then wrote to Grier and prevailed upon him, providing the majority leverage to issue a broad-ranging decision sufficient to render the Missouri Compromise of 1820 unconstitutional.\nTwo days after Buchanan was sworn in as president, Chief Justice Taney delivered the Dred Scott decision, which denied the petitioner's request to be set free from slavery. The ruling broadly asserted that Congress had no constitutional power to exclude slavery in the territories. According to this decision, slaves were forever the property of their owners without rights and no African American could ever be a full citizen of the United States, even if they had full civil rights in a state. Buchanan's letters were not made public at the time, but he was seen conversing quietly with the Chief Justice during his inauguration. When the decision was issued, Republicans began spreading the word that Taney had informed Buchanan of the impending outcome. Rather than destroying the Republican platform as Buchanan had hoped, the decision infuriated Northerners, who condemned it.\n\nPanic of 1857\nThe Panic of 1857 began in the summer of that year, when the New York branch of Ohio Life Insurance and Trust Company announced its insolvency. The crisis spread rapidly, and by the fall, 1,400 state banks and 5,000 businesses had gone bankrupt. Unemployment and hunger became common in northern cities, but the agricultural south was more resilient. Buchanan agreed with the southerners who attributed the economic collapse to over-speculation.\nBuchanan acted in accordance with Jacksonian Democracy principles, which restricted paper money issuance, and froze federal funds for public works projects, causing resentment among some of the population due to his refusal to implement an economic stimulus program. While the government was \"without the power to extend relief\", it would continue to pay its debts in specie, and while it would not curtail public works, none would be added. In hopes of reducing paper money supplies and inflation, he urged the states to restrict the banks to a credit level of $3 to $1 of specie and discouraged the use of federal or state bonds as security for bank note issues. The economy recovered in several years, though many Americans suffered as a result of the panic. Buchanan had hoped to reduce the deficit, but by the time he left office the federal budget grew by 15%.\n\nUtah War\nIn the spring of 1857, the Latter-day Saints and their leader Brigham Young had been challenging federal representatives in Utah Territory, causing harassment and violence against non-Mormons. Young harassed federal officers and discouraged outsiders from settling in the Salt Lake City area. In September 1857, the Utah Territorial Militia, associated with the Latter-day Saints, perpetrated the Mountain Meadows massacre, in which Young's militia attacked a wagon train and killed 125 settlers. Buchanan was offended by the militarism and polygamous behavior of Young. With reports of violence against non-Mormons, Buchanan authorized a military expedition into Utah Territory in late March 1857 to replace Young as governor. The force consisted of 2,500 men, including Alfred Cumming and his staff, and was commanded by General William S. Harney. Complicating matters, Young's notice of his replacement was not delivered because the Pierce administration had annulled the Utah mail contract, and Young portrayed the approaching forces as an unauthorized overthrow.\nBuchanan's personnel decision incited resistance from the Mormons around Young, as Harney was known for his volatility and brutality. In August 1857, Albert S. Johnston replaced him for organizational reasons. Young reacted to the military action by mustering a two-week expedition, destroying wagon trains, oxen, and other Army property. Buchanan then dispatched Thomas L. Kane as a private agent to negotiate peace. The mission was successful, a peaceful agreement to replace Governor Young with Cumming was reached, and the Utah War ended. The President granted amnesty to inhabitants affirming loyalty to the government, and placed the federal troops at a peaceable distance for the balance of his administration.\nBuchanan did not comment on the conflict again until his State of the Union Address in December 1857, leaving open the question of whether it was a rebellion in Utah. One of Buchanan's last official acts in March 1861 was to reduce the size of Utah Territory in favor of Nevada, Colorado, and Nebraska. While the Latter-day Saints had frequently defied federal authority, some historians consider Buchanan's action was an inappropriate response to uncorroborated reports.\n\nTransatlantic telegraph cable\nBuchanan was the first recipient of an official telegram transmitted across the Atlantic. Following the dispatch of test and configuration telegrams, on August 16, 1858 Queen Victoria sent a 98-word message to Buchanan at his summer residence in the Bedford Springs Hotel in Pennsylvania, expressing hope that the newly laid cable would prove \"an additional link between the nations whose friendship is founded on their common interest and reciprocal esteem\". Queen Victoria's message took 16 hours to send.\nBuchanan responded: \"It is a triumph more glorious, because far more useful to mankind, than was ever won by conqueror on the field of battle. May the Atlantic telegraph, under the blessing of Heaven, prove to be a bond of perpetual peace and friendship between the kindred nations, and an instrument destined by Divine Providence to diffuse religion, civilization, liberty, and law throughout the world.\"\n\nBleeding Kansas and constitutional dispute\nThe Kansas–Nebraska Act of 1854 created the Kansas Territory and allowed the settlers there to decide whether to allow slavery. This resulted in violence between \"Free-Soil\" (antislavery) and pro-slavery settlers, which developed into the \"Bleeding Kansas\" period. The antislavery settlers, with the help of Northern abolitionists, organized their own territorial government in Topeka. The more numerous proslavery settlers, many from the neighboring slave state Missouri, established a government in Lecompton, giving the Territory two different governments for a time, with two distinct constitutions, each claiming legitimacy. The admission of Kansas as a state required a constitution be submitted to Congress with the approval of a majority of its residents. Under President Pierce, a series of violent confrontations escalated over who had the right to vote in Kansas. The situation drew national attention, and some in Georgia and Mississippi advocated secession should Kansas be admitted as a free state. Buchanan chose to endorse the pro-slavery Lecompton government.\nBuchanan appointed Robert J. Walker to replace John W. Geary as Territorial Governor, and there ensued conflicting referendums from Topeka and Lecompton, where election fraud occurred. In October 1857, the Lecompton government framed the pro-slavery Lecompton Constitution that agreed to a referendum limited solely to the slavery question. However, the vote against slavery, as provided by the Lecompton Convention, would still permit existing slaves, and all their issue, to be enslaved, so there was no referendum that permitted the majority anti-slavery residents to prohibit slavery in Kansas. As a result, anti-slavery residents boycotted the referendum since it did not provide a meaningful choice.\nDespite the protests of Walker and two former Kansas governors, Buchanan decided to accept the Lecompton Constitution. In a December 1857 meeting with Stephen A. Douglas, the chairman of the Senate Committee on Territories, Buchanan demanded that all Democrats support the administration's position of admitting Kansas under the Lecompton Constitution. On February 2, he transmitted the Lecompton Constitution to Congress. He also transmitted a message that attacked the \"revolutionary government\" in Topeka, conflating them with the Mormons in Utah. Buchanan made every effort to secure congressional approval, offering favors, patronage appointments, and even cash for votes. The Lecompton Constitution won the approval of the Senate in March, but a combination of Know-Nothings, Republicans, and Northern Democrats defeated the bill in the House.\nBuchanan never forgave Douglas, as the Northern Democrats' rejection was the deciding factor in the House's decision, and he removed all Douglas supporters from his patronage in Illinois and Washington, D.C., installing pro-administration Democrats, including postmasters. Rather than accepting defeat, Buchanan backed the 1858 English Bill, which offered Kansas immediate statehood and vast public lands in exchange for accepting the Lecompton Constitution. In August 1858, Kansans by referendum strongly rejected the Lecompton Constitution. The territory received an abolitionist constitution, which was bitterly opposed in Congress by representatives and senators from the southern states until Kansas was admitted to the Union in January 1861.\nThe dispute over Kansas became the battlefront for control of the Democratic Party. On one side were Buchanan, the majority of Southern Democrats, and the \"doughfaces\". On the other side were Douglas and the majority of northern Democrats, as well as a few Southerners. Douglas's faction continued to support the doctrine of popular sovereignty, while Buchanan insisted that Democrats respect the Dred Scott decision and its repudiation of federal interference with slavery in the territories.\n\n1858 mid-term elections\nDouglas's Senate term was coming to an end in 1859, with the Illinois legislature, elected in 1858, determining whether Douglas would win re-election. The Senate seat was the primary issue of the legislative election, marked by the famous debates between Douglas and his Republican opponent for the seat, Abraham Lincoln. Buchanan, working through federal patronage appointees in Illinois, ran candidates for the legislature in competition with both the Republicans and the Douglas Democrats. This could easily have thrown the election to the Republicans, and showed the depth of Buchanan's animosity toward Douglas. In the end, Douglas Democrats won the legislative election and Douglas was re-elected to the Senate. In that year's elections, Douglas forces took control throughout the North, except in Buchanan's home state of Pennsylvania. Buchanan's support was otherwise reduced to a narrow base of southerners.\nThe division between northern and southern Democrats allowed the Republicans to win a plurality of the House in the 1858 elections, and allowed them to block most of Buchanan's agenda. Buchanan, in turn, added to the hostility with his veto of six substantial pieces of Republican legislation. Among these measures were the Homestead Act, which would have given 160 acres of public land to settlers who remained on the land for five years, and the Morrill Act, which would have granted public lands to establish land-grant colleges. Buchanan argued that these acts were unconstitutional. In the western and northwestern United States, where the Homestead Act was very popular, even many Democrats condemned the president's policies, while many Americans who considered education an important asset resented Buchanan's veto of agricultural colleges.\n\nForeign policy\nBuchanan took office with an ambitious foreign policy, designed to establish U.S. hegemony over Central America at the expense of Great Britain. Buchanan sought to revitalize Manifest Destiny and to enforce the Monroe Doctrine, which had been under attack from the Spanish, French, and especially the British in the 1850s. He hoped to re-negotiate the Clayton–Bulwer Treaty to counter European imperialism in the Western Hemisphere, which he thought limited U.S. influence in the region. He also sought to establish American protectorates over the Mexican states of Chihuahua and Sonora to secure American citizens and investments, and most importantly, he hoped to achieve his long-term goal of acquiring Cuba. However, Buchanan's ambitions in Cuba and Mexico were largely blocked by the House of Representatives. After long negotiations with the British, he convinced them to cede the Bay Islands to Honduras and the Mosquito Coast to Nicaragua.\nIn 1858, Buchanan ordered the Paraguay expedition to punish Paraguay for firing on the USS Water Witch, ordering 2,500 marines and 19 warships there. This costly expedition took months to reach Asunción, which successfully resulted in a Paraguayan apology and payment of an indemnity. The chiefs of Raiatea and Tahaa in the South Pacific, refusing to accept the rule of King Tamatoa V, unsuccessfully petitioned the United States to accept the islands under a protectorate in June 1858. Buchanan also considered buying Alaska from the Russian Empire, as whaling in the waters there had become of great economic importance to the United States. Buchanan fueled this by spreading the rumor to the Russian ambassador Eduard de Stoeckl in December 1857 that a large amount of Mormons intended to emigrate to Russian Alaska. In the winter of 1859, an initial purchase offer of $5,000,000 (equivalent to $169,560,000 in 2023) was made. Although the project ultimately failed due to the reservations of Foreign Minister Alexander Gorchakov, the talks formed the basis for the later negotiations to purchase Alaska.\nBuchanan sought trade agreements with the Qing Dynasty and Japan. In China, his envoy William Bradford Reed succeeded in having the United States included as a party to the Treaty of Tianjin. In May 1860, Buchanan received a Japanese delegation consisting of several princes who carried the Harris Treaty negotiated by Townsend Harris for mutual ratification. Buchanan was offered a herd of elephants by King Rama IV of Siam, though the letter arrived after Buchanan's departure from office and Buchanan's successor Abraham Lincoln declined the offer stating that the U.S. had an unsuitable climate. Other presidential pets included a pair of bald eagles and a Newfoundland dog.\n\nCovode Committee\nIn March 1860, the House impaneled the Covode Committee to investigate the Buchanan administration's patronage system for alleged impeachable offenses, such as bribery and extortion of representatives. Buchanan supporters accused the committee, consisting of three Republicans and two Democrats, of being blatantly partisan, and claimed its chairman, Republican Rep. John Covode, was acting on a personal grudge stemming from a disputed land grant designed to benefit Covode's railroad company. The Democratic committee members, as well as Democratic witnesses, were enthusiastic in their condemnation of Buchanan.\nThe committee was unable to establish grounds for impeaching Buchanan; however, the majority report issued on June 17 alleged corruption and abuse of power among members of his cabinet. The committee gathered evidence that Buchanan had tried to bribe members of Congress in his favor through intermediaries in the spring of 1858 in connection with the pro-slavery Lecompton Constitution of Kansas, and threatened their relatives with losing their posts if they did not vote in favor of the Lecompton Constitution. Witnesses also testified that the federal government used public funds to strengthen the intra-party faction of Douglas's opponents in Illinois. The Democrats pointed out that evidence was scarce, but did not refute the allegations; one of the Democratic members, Rep. James Robinson, stated that he agreed with the Republicans, though he did not sign it.\nThe public was shocked by the extent of the bribery, which affected all levels and agencies of government. Buchanan claimed to have \"passed triumphantly through this ordeal\" with complete vindication. Republican operatives distributed thousands of copies of the Covode Committee report throughout the nation as campaign material in that year's presidential election.\n\nElection of 1860\nAs he had promised in his inaugural address, Buchanan did not seek re-election. He went so far as to tell his ultimate successor, \"If you are as happy in entering the White House as I shall feel on returning to Wheatland, you are a happy man.\"\nAt the 1860 Democratic National Convention in Charleston, the party split over the issue of slavery in the territories, damaging Buchanan's reputation as the main person responsible for this issue. Though Douglas led after every ballot, he was unable to win the two-thirds majority required. The convention adjourned after 53 ballots, and re-convened in Baltimore in June. After Douglas finally won the nomination, several Southerners refused to accept the outcome, and nominated Vice President Breckinridge as their own candidate. Douglas and Breckinridge agreed on most issues except the protection of slavery. Buchanan, nursing a grudge against Douglas, failed to reconcile the party, and tepidly supported Breckinridge. With the splintering of the Democratic Party, Republican nominee Abraham Lincoln won a four-way election that also included John Bell of the Constitutional Union Party. Lincoln's support in the North was enough to give him an Electoral College majority. Buchanan became the last Democrat to win a presidential election until Grover Cleveland in 1884.\nAs early as October, the army's Commanding General, Winfield Scott, an opponent of Buchanan, warned him that Lincoln's election would likely cause at least seven states to secede from the union. He recommended that massive amounts of federal troops and artillery be deployed to those states to protect federal property, although he also warned that few reinforcements were available. Since 1857, Congress had failed to heed calls for a stronger militia and allowed the army to fall into deplorable condition. Buchanan distrusted Scott and ignored his recommendations. After Lincoln's election, Buchanan directed Secretary of War John B. Floyd to reinforce southern forts with such provisions, arms, and men as were available; however, Floyd persuaded him to revoke the order.\n\nSecession\nWith Lincoln's victory, talk of secession and disunion reached a boiling point, putting the burden on Buchanan to address it in his final speech to Congress on December 10. In his message, which was anticipated by both factions, Buchanan denied the right of states to secede but maintained the federal government was without power to prevent them. He placed the blame for the crisis solely on \"intemperate interference of the Northern people with the question of slavery in the Southern States,\" and suggested that if they did not \"repeal their unconstitutional and obnoxious enactments ... the injured States, after having first used all peaceful and constitutional means to obtain redress, would be justified in revolutionary resistance to the Government of the Union.\" Buchanan's only suggestion to solve the crisis was \"an explanatory amendment\" affirming the constitutionality of slavery in the states, the fugitive slave laws, and popular sovereignty in the territories. His address was sharply criticized both by the North, for its refusal to stop secession, and the South, for denying its right to secede. Five days after the address was delivered, Treasury Secretary Howell Cobb resigned, as his views had become irreconcilable with the President's. Even as the formation of the Confederacy by the secessionist states became increasingly apparent in the winter of 1860, the president continued to surround himself with Southerners and ignore the Republicans.\n\nSouth Carolina, long the most radical Southern state, seceded from the Union on December 20, 1860. However, Unionist sentiment remained strong among many in the South, and Buchanan sought to appeal to the Southern moderates who might prevent secession in other states. He met with South Carolinian commissioners in an attempt to resolve the situation at Fort Sumter, which federal forces remained in control of despite its location in Charleston, South Carolina. Buchanan saw Congress, not himself, as responsible for finding a solution to the secession crisis. As a compromise for the southern states, Buchanan envisioned the adoption of amendments to the United States Constitution that would guarantee the right to slavery in the southern states and territories and strengthen the right of slave owners to reclaim escaped slaves as property in the northern states.\nHe refused to dismiss Interior Secretary Jacob Thompson after the latter was chosen as Mississippi's agent to discuss secession, and he refused to fire Secretary of War John B. Floyd despite an embezzlement scandal. Floyd ended up resigning, but not before sending numerous firearms to Southern states, where they eventually fell into the hands of the Confederacy. Despite Floyd's resignation, Buchanan continued to seek the advice of counselors from the Deep South, including Jefferson Davis and William Henry Trescot. Buchanan's friend Rose O'Neal Greenhow took advantage of the proximity to the president and spied for the Confederacy, which had already established a sophisticated network for gathering information from its eventual opponent before its formation.\nEfforts were made in vain by Sen. John J. Crittenden, Rep. Thomas Corwin, and former president John Tyler to negotiate a compromise to stop secession, with Buchanan's support. Failed attempts were also made by a group of governors meeting in New York. Buchanan secretly asked President-elect Lincoln to call for a national referendum on the issue of slavery, but Lincoln declined. In December 1860, when the second session of the 36th Congress was convened, The Committee of Thirty-Three was established by the House of Representatives to prevent further states from seceding. They proposed the Corwin Amendment, which would bar Congress from interfering with slavery in states. Despite opposition from Republicans, it passed both houses of Congress and was proposed to states for ratification, but it was never ratified by the requisite number of states.\nDespite the efforts of Buchanan and others, six more slave states seceded by the end of January 1861. Buchanan replaced the departed Southern cabinet members with John Adams Dix, Edwin M. Stanton, and Joseph Holt, all of whom were committed to preserving the Union. When Buchanan considered surrendering Fort Sumter, the new cabinet members threatened to resign, and Buchanan relented. On January 5, Buchanan decided to reinforce Fort Sumter, sending the Star of the West with 250 men and supplies. However, he failed to ask Major Robert Anderson to provide covering fire for the ship, and it was forced to return North without delivering troops or supplies. Buchanan chose not to respond to this act of war, and instead sought to find a compromise to avoid secession. He received a March 3 message from Anderson, that supplies were running low, but the response became Lincoln's to make, as the latter succeeded to the presidency the next day.\n\nStates admitted to the Union\nThree new states were admitted to the Union while Buchanan was in office:\n\nMinnesota – May 11, 1858\nOregon – February 14, 1859\nKansas – January 29, 1861\n\nFinal years and death (1861–1868)\nAfter leaving office, Buchanan retired to private life in Wheatland, where he spent most of his time in his study, reading books and writing letters. The Civil War erupted within two months of Buchanan's retirement. He supported the Union and the war effort, writing to former colleagues that, \"the assault upon Sumter was the commencement of war by the Confederate states, and no alternative was left but to prosecute it with vigor on our part.\" Buchanan supported Lincoln's introduction of universal conscription in the northern states, but was an opponent of his Emancipation Proclamation. Although he recognized constitutional violations in some of the president's executive orders, he never criticized them in public. He also wrote a letter to his fellow Pennsylvania Democrats in Harrisburg, urging them and all young men to enlist in the Union army and \"join the many thousands of brave & patriotic volunteers who are already in the field.\"\nBuchanan was dedicated to defending his actions prior to the Civil War, which was referred to by some as \"Buchanan's War\". He received hate mail and threatening letters daily, and stores in Lancaster displayed Buchanan's likeness with the eyes inked red, a noose drawn around his neck and the word \"TRAITOR\" written across his forehead. The Senate proposed a resolution of condemnation which ultimately failed, and newspapers accused him of colluding with the Confederacy. His former cabinet members, five of whom had been given jobs in the Lincoln administration, refused to defend Buchanan publicly.\nBuchanan became distraught by the vitriolic attacks levied against him, and fell sick and depressed. In October 1862, he defended himself in an exchange of letters with Winfield Scott, published in the National Intelligencer. He soon began writing his fullest public defense, in the form of his memoir Mr. Buchanan's Administration on the Eve of Rebellion, which was published in 1866, one year after the Civil War ended. Buchanan attributed secession to the \"malign influence\" of Republicans and the abolitionist movement. He discussed his foreign policy successes and expressed satisfaction with his decisions, even during the secession crisis. He blamed Robert Anderson, Winfield Scott, and Congress for the unresolved issue. Two years after the publication of the memoir, Buchanan caught a cold in May 1868, which quickly worsened due to his advanced age. He died on June 1, 1868, of respiratory failure at the age of 77 at his home at Wheatland. He was interred in Woodward Hill Cemetery in Lancaster.\n\nPolitical views\nBuchanan was often considered by anti-slavery northerners a \"doughface\", a northerner with pro-southern principles. Buchanan's sympathies for the Southern states went beyond political expediency for his path to the White House. He identified with cultural and social values that he found reflected in the honor code and lifestyle of the planter class and with which he increasingly came into contact in his retirement community beginning in 1834. Shortly after his election, he said that the \"great object\" of his administration was \"to arrest, if possible, the agitation of the Slavery question in the North and to destroy sectional parties\". Although Buchanan was personally opposed to slavery, he believed that the abolitionists were preventing the solution to the slavery problem. He stated, \"Before [the abolitionists] commenced this agitation, a very large and growing party existed in several of the slave states in favor of the gradual abolition of slavery; and now not a voice is heard there in support of such a measure. The abolitionists have postponed the emancipation of the slaves in three or four states for at least half a century.\" In deference to the intentions of the typical slaveholder, he was willing to provide the benefit of the doubt. In his third annual message to Congress, the president claimed that the slaves were \"treated with kindness and humanity. ... Both the philanthropy and the self-interest of the master have combined to produce this humane result.\"\n\nBuchanan thought restraint was the essence of good self-government. He believed the constitution comprised \"... restraints, imposed not by arbitrary authority, but by the people upon themselves and their representatives. ... In an enlarged view, the people's interests may seem identical, but to the eye of local and sectional prejudice, they always appear to be conflicting ... and the jealousies that will perpetually arise can be repressed only by the mutual forbearance which pervades the constitution.\" Regarding slavery and the Constitution, he stated: \"Although in Pennsylvania we are all opposed to slavery in the abstract, we can never violate the constitutional compact we have with our sister states. Their rights will be held sacred by us. Under the constitution it is their own question; and there let it remain.\"\nOne of the prominent issues of the day was tariffs. Buchanan was conflicted by free trade as well as prohibitive tariffs, since either would benefit one section of the country to the detriment of the other. As a senator from Pennsylvania, he said: \"I am viewed as the strongest advocate of protection in other states, whilst I am denounced as its enemy in Pennsylvania.\"\nBuchanan was also torn between his desire to expand the country for the general welfare of the nation, and to guarantee the rights of the people settling particular areas. On territorial expansion, he said, \"What, sir? Prevent the people from crossing the Rocky Mountains? You might just as well command the Niagara not to flow. We must fulfill our destiny.\" On the resulting spread of slavery, through unconditional expansion, he stated: \"I feel a strong repugnance by any act of mine to extend the present limits of the Union over a new slave-holding territory.\" For instance, he hoped the acquisition of Texas would \"be the means of limiting, not enlarging, the dominion of slavery.\"\n\nPersonal life\nBuchanan suffered from esotropia. In addition, one eye was short-sighted and the other far-sighted. To cover this, he bent his head forward and leaned it to one side during social interactions. This led to ridicule, which Henry Clay, among others, used ruthlessly during a congressional debate.\nIn 1818, Buchanan met Anne Caroline Coleman at a grand ball in Lancaster, and the two began courting. Anne was the daughter of the wealthy iron manufacturer Robert Coleman; Robert, like Buchanan's father, was from County Donegal in Ulster. Anne was also the sister-in-law of Philadelphia judge Joseph Hemphill, one of Buchanan's colleagues. By 1819, the two were engaged, but spent little time together. Buchanan was busy with his law firm and political projects during the Panic of 1819, which took him away from Coleman for weeks at a time. Rumors abounded, as some suggested that he was involved with other (unidentified) women. Letters from Coleman revealed she was aware of several rumors, and she accused him of only being interested in her money. She broke off the engagement, and soon afterward, on December 9, 1819, inexplicably died of \"hysterical convulsions\" resulting from an overdose of laudanum, at the age of 23. It was never established if the drug was taken by instruction, by accident, or by intent. Buchanan wrote to her father for permission to attend the funeral, which was refused. At the time of her funeral, he said that, \"I feel happiness has fled from me forever.\" Afterwards, Buchanan claimed that he remained unmarried out of devotion to his only love, who had died young.\n\nIn 1833 and the 1840s, he spoke of plans to marry, but these came to nothing and may merely have been due to his ambitions for a seat in the federal Senate or the White House. In the latter case, the aspirant was 19-year-old Anna Payne, the niece of former First Lady Dolley Madison. During his presidency, an orphaned niece, Harriet Lane, whom he had adopted, served as official White House hostess. There was an unfounded rumor that he had an affair with President Polk's widow, Sarah Childress Polk.\nBuchanan had a close relationship with William Rufus King, which became a popular target of gossip. King was an Alabama politician who briefly served as vice president under Franklin Pierce. Buchanan and King lived together in a Washington boardinghouse and attended social functions together from 1834 until 1844. Such a living arrangement was then common, though Buchanan once referred to the relationship as a \"communion\". Andrew Jackson mockingly called them \"Miss Nancy\" and \"Aunt Fancy\", the former being a 19th-century euphemism for an effeminate man. Buchanan's Postmaster General, Aaron V. Brown, also referred to King as \"Aunt Fancy\", as well as Buchanan's \"better half\", and \"wife\". King died of tuberculosis shortly after Pierce's inauguration, four years before Buchanan became president. Buchanan described him as \"among the best, the purest and most consistent public men I have known\". Biographer Baker opines that both men's nieces may have destroyed correspondence between the two men. However, she believes that their surviving letters illustrate only \"the affection of a special friendship\".\nBuchanan's lifelong bachelorhood after Anne Coleman's death has drawn interest and speculation. Some conjecture that Anne's death merely served to deflect questions about Buchanan's sexuality and bachelorhood. One of his biographers, Jean Baker, suggests that Buchanan was celibate, if not asexual. Several writers have surmised that he was homosexual, including James W. Loewen, Robert P. Watson, and Shelley Ross. Loewen indicated that Buchanan, late in life, wrote a letter acknowledging that he might marry a woman who could accept his \"lack of ardent or romantic affection\".\n\nLegacy\nHistorical reputation\nThough Buchanan predicted that \"history will vindicate my memory,\" historians have criticized Buchanan for his unwillingness or inability to act in the face of secession. Historical rankings of presidents of the United States without exception place Buchanan among the least successful presidents. When scholars are surveyed, he ranks at or near the bottom in terms of vision/agenda-setting, domestic leadership, foreign policy leadership, moral authority, and positive historical significance of their legacy. According to surveys taken by American scholars and political scientists between 1948 and 1982, Buchanan ranks every time among the worst presidents of the United States, alongside Harding, Fillmore and Nixon.\nBuchanan biographer Philip S. Klein focused in 1962, during the Civil Rights movement, upon challenges Buchanan faced:\n\nBuchanan assumed leadership ... when an unprecedented wave of angry passion was sweeping over the nation. That he held the hostile sections in check during these revolutionary times was in itself a remarkable achievement. His weaknesses in the stormy years of his presidency were magnified by enraged partisans of the North and South. His many talents, which in a quieter era might have gained for him a place among the great presidents, were quickly overshadowed by the cataclysmic events of civil war and by the towering Abraham Lincoln.\nBiographer Jean Baker is less charitable to Buchanan, saying in 2004:\n\nAmericans have conveniently misled themselves about the presidency of James Buchanan, preferring to classify him as indecisive and inactive ... In fact Buchanan's failing during the crisis over the Union was not inactivity, but rather his partiality for the South, a favoritism that bordered on disloyalty in an officer pledged to defend all the United States. He was that most dangerous of chief executives, a stubborn, mistaken ideologue whose principles held no room for compromise. His experience in government had only rendered him too self-confident to consider other views. In his betrayal of the national trust, Buchanan came closer to committing treason than any other president in American history.Other historians, such as Robert May, argued that his politics were \"anything but pro-slavery\", nevertheless, a very negative view is to be found in Michael Birkner's works about Buchanan. For Lori Cox Han, he ranks among scholars \"as either the worst president in [American] history or as part of a lowest ranking failure category\".\n\nMemorials\nA bronze and granite memorial near the southeast corner of Washington, D.C.'s Meridian Hill Park was designed by architect William Gorden Beecher and sculpted by Maryland artist Hans Schuler. It was commissioned in 1916 but not approved by the U.S. Congress until 1918, and not completed and unveiled until June 26, 1930. The memorial features a statue of Buchanan, bookended by male and female classical figures representing law and diplomacy, with engraved text reading: \"The incorruptible statesman whose walk was upon the mountain ranges of the law,\" a quote from a member of Buchanan's cabinet, Jeremiah S. Black.\n\nAn earlier monument was constructed in 1907–1908 and dedicated in 1911, on the site of Buchanan's birthplace in Stony Batter, Pennsylvania. Part of the original 18.5-acre (75,000 m2) memorial site is a 250-ton pyramid structure that stands on the site of the original cabin where Buchanan was born. The monument was designed to show the original weathered surface of the native rubble and mortar.\nThree counties are named in his honor, in Iowa, Missouri, and Virginia. Another in Texas was christened in 1858 but renamed Stephens County, after the newly elected vice president of the Confederate States of America, Alexander Stephens, in 1861. The city of Buchanan, Michigan, was also named after him. Several other communities are named after him: the unincorporated community of Buchanan, Indiana, the city of Buchanan, Georgia, the town of Buchanan, Wisconsin, and the townships of Buchanan Township, Michigan, and Buchanan, Missouri.\nJames Buchanan High School is a small, rural high school located on the outskirts of his childhood hometown, Mercersburg, Pennsylvania.\n\nPopular culture depictions\nBuchanan and his legacy are central to the film Raising Buchanan (2019). He is portrayed by René Auberjonois.\n\nSee also\nHistorical rankings of presidents of the United States\nList of presidents of the United States\nList of presidents of the United States by previous experience\nPresidents of the United States on U.S. postage stamps\nList of federal political sex scandals in the United States\n\nReferences\nWorks cited\nFurther reading\nExternal links\n\nUnited States Congress. \"James Buchanan (id: B001005)\". Biographical Directory of the United States Congress.\nJames Buchanan: A Resource Guide from the Library of Congress\nThe James Buchanan papers, spanning the entirety of his legal, political and diplomatic career, are available for research use at the Historical Society of Pennsylvania.\nUniversity of Virginia article: Buchanan biography\nWheatland\nJames Buchanan at Tulane University\nEssay on James Buchanan and his presidency from the Miller Center of Public Affairs\nBuchanan's Birthplace State Park, Franklin County, Pennsylvania\n\"Life Portrait of James Buchanan\", from C-SPAN's American Presidents: Life Portraits, June 21, 1999\nPrimary sources\n\nWorks by James Buchanan at Project Gutenberg\nWorks by James Buchanan at LibriVox (public domain audiobooks) \nWorks by or about James Buchanan at the Internet Archive\nJames Buchanan Ill with Dysentery Before Inauguration: Original Letters Shapell Manuscript Foundation\nMr. Buchanans Administration on the Eve of the Rebellion. President Buchanans memoirs.\nInaugural Address Archived August 9, 2020, at the Wayback Machine\nFourth Annual Message to Congress, December 3, 1860\n\nHarriet Rebecca Lane Johnston (May 9, 1830 – July 3, 1903) acted as first lady of the United States during the administration of her uncle, lifelong bachelor president James Buchanan, from 1857 to 1861. She has been described as the first of the modern first ladies, being a notably charming and diplomatic hostess, whose dress-styles were copied, and who promoted deserving causes. In her will, she left funds for a new school on the grounds of Washington National Cathedral. Several ships have been named in her honor, including the cutter USCGC Harriet Lane, still in service.\n\nStatus\nLane is the only person to have served as First Lady to a bachelor president, Buchanan being the only U.S. president never to have married. She is among 11 women who have served as First Lady, but were not married to the president, with most of the other women being relatives of widowed presidents.\n\nEarly life\nHarriet Lane's family was from Franklin County, Pennsylvania. She was the youngest child of Elliott Tole Lane, a merchant, and Jane Ann Buchanan Lane. She lost her mother when she was nine; when her father's death two years later made her an orphan, she requested that her favorite uncle, James Buchanan, be appointed as her legal guardian. Buchanan, an unmarried Democratic senator from Pennsylvania, indulged his niece and her sister, enrolling them in boarding schools in Charles Town, Virginia (later for two years at the Georgetown Visitation Monastery in the Georgetown section of Washington, D.C.) By this time, Buchanan was Secretary of State, and, as he had promised, he introduced her to fashionable and political circles.\nIn 1854, she joined him in London, where he was minister to the Court of St. James's. Queen Victoria gave \"dear Miss Lane\" the rank of ambassador's wife; admiring suitors gave her the fame of a beauty. In appearance \"Hal\" Lane was of medium height, with masses of light, almost golden-colored hair. She had eyes that were described as \"violet colored\".\n\nActing First Lady of the United States\nThe capital welcomed its new \"Democratic Queen\" to the White House in 1857. Harriet was a popular hostess during the four years of the Buchanan presidency. Women copied her hair and clothing styles (especially when she lowered the neckline on her inaugural gown by 2.5 inches), parents named their daughters for her, and a popular song (\"Listen to the Mockingbird\") was dedicated to her. While in the White House, she used her position to promote social causes, such as improving the living conditions of Native Americans in reservations. She also made a point of inviting artists and musicians to White House functions. For both her popularity and her advocacy work, she has been described as the first of the modern first ladies, and her popularity at the time is compared to that of Jacqueline Kennedy in the 1960s. The presidential yacht was named for her—the first of several ships to be named after her, one of which remains in service.\n\nAs sectional tensions increased, she worked out seating arrangements for her weekly formal dinner parties with special care, to give dignitaries their proper precedence and still keep political foes apart. Her tact did not falter, but her task became impossible—as did her uncle's. Seven states had seceded by the time Buchanan retired from office and returned with his niece to his spacious country home, Wheatland, near Lancaster, Pennsylvania.\nIn the 1982 Siena College Research Institute survey asking historians to assess American first ladies, Lane and several other \"acting\" first ladies were included. The first ladies survey, which has been conducted periodically since, ranks first ladies according to a cumulative score on the independent criteria of their background, value to the country, intelligence, courage, accomplishments, integrity, leadership, being their own women, public image, and value to the president. In the 1982 survey, out of 42 first ladies and acting first ladies, Lane was assessed as the 29th most highly regarded among historians. Acting first ladies such as Lane have been excluded from subsequent iterations of this survey.\n\nRomance and marriage\nDuring her time in England, Sir Fitzroy Kelly, then Prime Minister Palmerston's attorney general, proposed marriage to her; Queen Victoria was strongly in favor of this match, as it would keep Lane in England.\nLane considered the advantages of a number of bachelors. Her uncle cautioned Lane against \"rushing precipitately into matrimonial connections\" as his ward found her potential suitors \"pleasant but dreadfully troublesome\". Lane eventually married Baltimore banker Henry Elliott Johnston at the age of 36. They had two sons: James Buchanan Johnston (1866–1881) and Henry Elliot Johnston (1869–1882), but within the 18 years from 1867 to 1885, her uncle, her husband, and her children all died.\n\nLater life and death\nHarriet wrote her will in 1895 and lived another eight years, during which the country's general prosperity greatly increased the value of her estate. She added a codicil in 1899 directing that a school building be constructed on the grounds of the Washington National Cathedral property and asked that it be called the Lane-Johnston Building \"to the end that the family names of my husband and myself may be associated with the bequest made in loving memory of our sons.\" A codicil of 1903 increased her gift by one third but said that only half the total was to be spent on the building. The remainder was \"specially to provide for the free maintenance, education and training of choirboys, primarily those in service of the Cathedral.\" This bequest founded the prestigious boys' school that today is called St. Albans School, which opened in October 1909. \nAt Harriet Lane Johnston's funeral, services were conducted by Bishop Satterlee and Canon DeVries of the Washington National Cathedral. She was buried in Green Mount Cemetery, Baltimore, Maryland, her grave marked with a Celtic cross like the Peace Cross on the cathedral close. In 1905, guests were invited to see the cornerstone of the first St. Albans School building, laid for what the invitation referred to as \"The Lane Johnston Choir School for Boys of the Washington Cathedral\".\n\nLegacy\nLane left bequests in her will that established a children's hospital and a boys' school, and she donated her collection of artwork to the Smithsonian. Several Navy and Coast Guard ships have been named in her honor.\nHer birthplace, the Lane House, was listed on the National Register of Historic Places in 1972.\n\nHospital and school\nShe dedicated $400,000 (equivalent to $13,600,000 in 2023) to establish the Harriet Lane Home for Invalid Children at the Johns Hopkins Hospital in Baltimore, Maryland as a memorial to two sons who had died in childhood. In October 1912 the Harriet Lane Home officially opened. It was the first children's clinic in the United States that was associated with a medical school. Eventually treating over 60,000 children a year, the Harriet Lane Home became a pioneer treatment, teaching, and research clinic.\nFrom 1930 to 1963 Helen Taussig, who helped to develop the blue baby operation, headed the pediatric cardiac clinic. Child psychiatrist Leo Kanner did studies of autistic children. Lawson Wilkins established an endocrine clinic that developed procedures used universally to treat children with certain glandular disorders, including dwarfism. John E. Bordley and William G. Hardy broke ground in detecting hearing impairments in very young children. It became a renowned pediatric facility; the Harriet Lane Outpatient Clinics serve thousands of children today, and the widely used manual for pediatric house officers, The Harriet Lane Handbook, bears her name.\nThe Harriet Lane Outpatient Clinics continue to operate in countries throughout the world.\nThe pediatric medicine Harriet Lane Handbook series continues in print and online, with multiple titles. The original title (subtitled A Manual for Pediatric House Officers) is in its 22nd edition, published by Mosby.\n\nArt collection\nShe had an art collection based on European works which she left to the U.S. government. The Smithsonian Institution called her the \"First Lady of the National Collection of Fine Arts\" after her collection was accepted into public ownership.\n\nNamesake ships\nThe United States Coast Guard has had three cutters named in her honor. The first was the USRC Harriet Lane, commissioned into the United States Revenue Cutter Service (predecessor of the USCG) in 1857. This cutter was transferred to the United States Navy in 1861 because of the American Civil War.\nThe second cutter named for Harriet Lane was the 125 foot USCGC Harriet Lane (WSC-141), commissioned in 1926 and decommissioned in 1946.\nThe third cutter named for Harriet Lane is the USCGC Harriet Lane (WMEC-903). The cutter was commissioned in May 1984, and as of 2021, is still in active service.\n\nFootnotes\nReferences\nFurther reading\nBalcerski, Thomas J. \"Harriet Rebecca Lane Johnston.\" in A Companion to First Ladies (2016): 197-213.\nRosenberger, Homer Tope. \"To what Extent Did Harriet Lance Influence the Public Policies of James Buchanan?\" Lancaster County Historical Society, 1970. online\nUpdike, John (1974). Buchanan Dying (play). (Ms. Johnston is a character in Updike's fictional play about President Buchanan.)\n\nExternal links\nWorks by or about Harriet Lane at the Internet Archive\n\"Harriet Lane\". First Ladies: Influence & Image. firstladies.org. CNN.\n\nSince the office was established in 1789, 45 persons have served as president of the United States. Of these, eight have died in office: four were assassinated, and four died of natural causes. In each of these instances, the vice president has succeeded to the presidency. This practice is now governed by Section One of the Twenty-fifth Amendment to the United States Constitution, ratified in 1967, which declares that, \"the Vice President shall become President\" if the president is removed from office, dies, or resigns. The initial authorization for this practice was provided by Article II, Section 1, Clause 6, of the U.S. Constitution.\nThe first incumbent U.S. president to die was William Henry Harrison, on April 4, 1841, only one month after Inauguration Day. He died from complications of what at the time was believed to be pneumonia. The second American president to die in office, Zachary Taylor, died on July 9, 1850, from acute gastroenteritis. Abraham Lincoln was the first U.S. president to be killed while in office. He was shot by John Wilkes Booth on the night of April 14, 1865, and died the following morning. Sixteen years later, on July 2, 1881, James A. Garfield was shot by Charles J. Guiteau, surviving for over two months before dying on September 19, 1881.\nOn September 14, 1901, William McKinley died, eight days after being shot by Leon Czolgosz. Next, Warren G. Harding suffered a heart attack, and died on August 2, 1923. On April 12, 1945, Franklin D. Roosevelt (who had just begun his fourth term in office) collapsed and died as a result of a cerebral hemorrhage. The most recent U.S. president to die in office was John F. Kennedy, who was shot by Lee Harvey Oswald on November 22, 1963, in Dallas, Texas.\n\n1841: William Henry Harrison\nOn March 26, 1841, William Henry Harrison became ill with a cold after being caught in a torrential downpour without cover. His symptoms grew progressively worse over the ensuing two days, at which time a team of doctors was called in to treat him. After making a diagnosis of right lower lobe pneumonia, they proceeded to place heated suction cups on his bare torso and to administer a series of bloodlettings, to supposedly draw out the disease. When those procedures failed to bring about improvement, the doctors treated him with ipecac, Castor oil, calomel, and finally with a boiled mixture of crude petroleum and Virginia snakeroot. All this only weakened Harrison further.\nInitially, no official announcement was made concerning Harrison's illness, which, the longer he remained out of public view, fueled public speculation and concern. By the end of the month large crowds were gathering outside the White House, holding vigil while awaiting any news about the president's condition. On the evening of April 4, 1841, nine days after becoming ill, and exactly one month after taking the oath of office, Harrison died at age 68. His last words were to his attending doctor, though assumed to be directed at Vice President John Tyler:\n\nSir, I wish you to understand the true principles of the government. I wish them carried out. I ask nothing more.\nA 30-day period of mourning commenced following the president's death. Various public ceremonies, modeled after European royal funeral practices, were held. An invitation-only funeral service was also held, on April 7 in the East Room of the White House, after which Harrison's coffin was brought to Congressional Cemetery in Washington, D.C., where it was placed in a temporary receiving vault.\nThat June, Harrison's body was transported by train and river barge to North Bend, Ohio. Then, on July 7, 1841, the nation's 9th president was buried in a family tomb at the summit of Mt. Nebo, overlooking the Ohio River – the William Henry Harrison Tomb State Memorial.\nHarrison's death sparked a brief constitutional crisis regarding succession to the presidency, as the U.S. Constitution was unclear as to whether Vice President John Tyler should assume the office of president or merely execute the duties of the vacant office. Tyler claimed a constitutional mandate to carry out the full powers and duties of the presidency and took the presidential oath of office, setting an important precedent for an orderly transfer of presidential power when a president leaves office intra-term.\nCoincidentally, all but one of the presidents who later died in office had, like Harrison, won a presidential election in a year ending in a zero (1840 through 1960). This pattern of tragedies came to be known as the Curse of Tippecanoe, or the Curse of Tecumseh, the name of the Shawnee leader against whom Harrison fought in the 1811 Battle of Tippecanoe. Also sometimes referred to as the Zero Factor legend, the pattern was disrupted by Ronald Reagan, who survived an assassination attempt in 1981 (69 days after taking office) and lived to complete two full terms.\n\n1850: Zachary Taylor\nZachary Taylor was known to have consumed copious amounts of ice water, cold milk, green apples, and cherries on July 4, 1850, after attending holiday celebrations and the laying of the cornerstone of the Washington Monument. That same evening, he became severely ill with an unknown digestive ailment. Doctors used popular treatments of the time. On the morning of July 9, the president asked his wife Margaret not to grieve saying:\n\nI have always done my duty, I am ready to die. My only regret is for the friends I leave behind me.\nTaylor died late that evening, five days after becoming ill, at age 65. Contemporary reports listed the cause of death as \"bilious diarrhea or a bilious cholera.\" He was succeeded by Vice President Millard Fillmore.\nTaylor's funeral took place on July 13, and like Harrison's nine years earlier, was held in the East Room of the White House. Afterward, an estimated 100,000 people gathered along the funeral route to Congressional Cemetery where his coffin was placed temporarily in the Public Vault; that October it was transported to Louisville, Kentucky. On November 1, 1850, Taylor was buried in his family's burial ground on the Taylor estate, Springfield, which became the Zachary Taylor National Cemetery.\nAlmost immediately after his death, rumors began to circulate that Taylor had been poisoned by pro-slavery Southerners, and various conspiracy theories persisted into the late-20th century. The cause of Taylor's death was definitively established in 1991, when his remains were exhumed and an autopsy conducted by Kentucky's chief medical examiner. Subsequent neutron activation analysis conducted at Oak Ridge National Laboratory revealed no evidence of poisoning, as arsenic levels were too low. The analysis concluded Taylor had contracted cholera morbus (acute gastroenteritis), as Washington had open sewers, and his food or drink may have been contaminated.\n\n1865: Abraham Lincoln\nThe assassination of Abraham Lincoln took place on Good Friday, April 14, 1865, as the Civil War was drawing to a close. He died the following morning at the age of 56. The assassination occurred five days after General Robert E. Lee and the Army of Northern Virginia surrendered to General Ulysses S. Grant and the Army of the Potomac following the Battle of Appomattox Court House. Lincoln was the first American president to be killed by an assassin. (The first U.S. president to be confronted by a would-be assassin was Andrew Jackson 30 years earlier, in January 1835.)\nThe assassination of President Lincoln was planned and carried out by the well-known stage actor John Wilkes Booth, a Confederate sympathizer, vehement in his denunciation of Lincoln, and a strong opponent of the abolition of slavery in the United States. Booth and a group of co-conspirators originally plotted to kidnap Lincoln, but later planned to kill him, Vice President Andrew Johnson, and Secretary of State William H. Seward in a bid to help the Confederacy's cause. Johnson's would-be-assassin, George Atzerodt did not carry out his part of the plan, and Johnson succeeded Lincoln as president while Lewis Powell only managed to wound Seward.\nLincoln was shot once in the back of his head while watching the play Our American Cousin with his wife Mary Todd Lincoln at Ford's Theatre in Washington, D.C., on the night of April 14, 1865. An army surgeon who happened to be at Ford's, Doctor Charles Leale, assessed Lincoln's wound as mortal. The unconscious president was then carried across the street from the theater to the Petersen House, where he remained in a coma for eight hours before dying the following morning.\nWithin two weeks of the manhunt for Lincoln's killers, on April 26, 1865, Booth and David Herold were caught in a tobacco barn in Port Conway, Virginia. While Herold surrendered, Booth was shot to death by Boston Corbett, a Union Corporal.\nA three-week series of official functions were held following the president's death. He lay in state in the East Room of the White House which was open to the public on April 18. A funeral service was held the next day, and then the coffin was transported in a procession down Pennsylvania Avenue to the United States Capitol, where a ceremonial burial service was held in the rotunda. After lying in state at the Capitol, Lincoln's remains were transported by train to Springfield, Illinois, for burial. He was interred on May 4, 1865, at Oak Ridge Cemetery in Springfield – the Lincoln Tomb State Historic Site since 1895.\n\n1881: James A. Garfield\nThe assassination of James A. Garfield happened in Washington, D.C., on July 2, 1881. Garfield was shot by Charles J. Guiteau at 9:30 a.m., less than four months into his term as the nation's 20th president. He died 11 weeks later on September 19, 1881, at the age of 49. Vice President Chester A. Arthur succeeded him as president. Garfield was scheduled to leave Washington on July 2, 1881, for his summer vacation. On that day, Guiteau lay in wait for the president at the Baltimore and Potomac Railroad station, on the southwest corner of present-day Sixth Street and Constitution Avenue NW, Washington, D.C.\nPresident Garfield came to the Sixth Street Station on his way to his alma mater, Williams College, where he was scheduled to deliver a speech. Garfield was accompanied by two of his sons, James and Harry, and Secretary of State James G. Blaine. Secretary of War Robert Todd Lincoln waited at the station to see the president off. Garfield had no bodyguard or security detail; with the exception of Abraham Lincoln during the Civil War, early U.S. presidents never used any guards.\nAs President Garfield entered the waiting room of the station, Guiteau stepped forward and pulled the trigger from behind at point-blank range. \"My God, what is that?!\" Garfield cried out, flinging up his arms. Guiteau fired again and Garfield collapsed. One bullet grazed Garfield's shoulder; the other hit him in the back, passing the first lumbar vertebra but missing the spinal cord before coming to rest behind his pancreas.\nGarfield, conscious but in shock, was carried to an upstairs floor of the train station. Lincoln sent for D.C. Bliss, a prominent Washington physician, who soon arrived and examined Garfield's wounds several times, probing for the bullet that remained lodged in the president's body with his fingers and metal probes. Two additional doctors were summoned, and they also probed the entry wound. Eventually there were about twenty people in the room, including at least ten physicians. As Garfield was being cared for, Lincoln, thinking back to the death of his father, said \"How many hours of sorrow I have passed in this town.\"\nGarfield was carried back to the White House. Although doctors told him that he would not survive the night, the president remained conscious and alert. The next morning his vital signs were good and doctors began to hope for recovery. A long vigil began, with Garfield's doctors issuing regular bulletins that the American public followed closely throughout the summer of 1881. His condition fluctuated. Fevers came and went. Garfield struggled to keep down solid food and spent most of the summer eating little, and that only liquids.\nGarfield had been a regular visitor to the shore town of Long Branch, New Jersey, one of the nation's premier summer vacation spots until World War I. In early September, it was decided to bring him to Elberon, a quiet beach town just to the south of Long Branch, in hopes that the beach air would help him recover. When they heard that the president was being brought to their town, local citizens built more than half a mile of tracks in less than 24 hours, enabling Garfield to be brought directly to the door of the oceanfront Franklyn cottage, rather than being moved by carriage from the local Elberon train station. However, Garfield died 12 days later. A granite marker on Garfield Road identifies the former site of the cottage, which was demolished in 1950. Throughout the five-month drama, anxious Americans across the country were kept informed of developments by the news media. The publisher of Frank Leslie's Illustrated Newspaper, Miriam Leslie, was especially quick to publish fully illustrated accounts of key moments, from Garfield's shooting to the embalming of his body.\nChester Arthur was at his home in New York City on the night of September 19, when word came that Garfield had died. After first getting the news, Arthur said \"I hope—my God, I do hope it is a mistake.\" But confirmation by telegram came soon after. Arthur took the presidential oath of office, administered by a New York Supreme Court judge, then left for Long Branch to pay his respects before traveling on to Washington. Garfield's body was taken to Washington, where it lay in state for two days in the Capitol Rotunda before being taken to Cleveland, where the funeral was held on September 26.\nWhen the tracks that had been hastily built to the Franklyn cottage were later torn up, actor Oliver Byron bought the wooden ties, and had local carpenter William Presley build them into a small tea house, in commemoration of the president. The red & white (originally red, white & blue) \"Garfield Tea House\" still survives, resting a couple of blocks away from the site of the cottage on the grounds of the Long Branch Historical Museum, a former Episcopal Church. The church is nicknamed \"The Church of the Presidents,\" as it had been attended by, in addition to Garfield, presidents Chester A. Arthur, Ulysses S. Grant, Benjamin Harrison, Rutherford Hayes, William McKinley, and Woodrow Wilson, during their own visits to Long Branch.\n\n1901: William McKinley\nWilliam McKinley was assassinated on September 6, 1901, inside the Temple of Music on the grounds of the Pan-American Exposition in Buffalo, New York. McKinley was shaking hands with the public when Leon Czolgosz, a Polish-American anarchist, shot him. The 58-year-old president died eight days later on September 14 from gangrene caused by the bullet wounds.\nMcKinley had been elected for a second term in 1900. He enjoyed meeting the public, and was reluctant to accept the security available to his office. The secretary to the president, George B. Cortelyou, feared an assassination attempt would take place during a visit to the Temple of Music, and twice took it off the schedule. McKinley restored it each time.\nCzolgosz had lost his job during the economic Panic of 1893 and turned to anarchism, a political philosophy whose adherents had previously killed foreign leaders. Regarding McKinley as a symbol of oppression, Czolgosz felt it was his duty as an anarchist to kill him. Unable to get near McKinley during the earlier part of the presidential visit, Czolgosz shot McKinley twice as the President reached to shake his hand in the reception line at the temple. One bullet grazed McKinley; the other entered his abdomen and was never found.\nMcKinley initially appeared to be recovering, but took a turn for the worse on September 13 as his wounds became gangrenous, and died early the next morning; Vice President Theodore Roosevelt succeeded him. Roosevelt was hiking near the top of Mt. Marcy, in New York's Adirondack region, when a runner located him to convey the news. After McKinley's murder, for which Czolgosz was put to death in the electric chair, the United States Congress passed legislation to officially charge the Secret Service with the responsibility for protecting the president.\n\n1923: Warren G. Harding\nWarren G. Harding died from a sudden heart attack in his hotel suite while visiting San Francisco on the evening of August 2, 1923, at the age of 57. His death quickly led to theories that he had been poisoned or committed suicide. Rumors of poisoning were fueled, in part, by a book called The Strange Death of President Harding by private detective and former Ohio Gang member Gaston Means, who suggested First Lady Florence Harding had poisoned her husband after learning of his infidelity. Mrs. Harding's refusal to allow an autopsy on President Harding only added to the speculation. According to the physicians attending Harding, however, the symptoms in the days prior to his death all pointed to congestive heart failure. Harding's biographer, Samuel H. Adams, concluded that \"Warren G. Harding died a natural death which, in any case, could not have been long postponed.\"\nImmediately after President Harding's death, Mrs. Harding returned to Washington, D.C., and briefly stayed in the White House with the new president Calvin Coolidge and first lady. For a month, former first lady Harding gathered and destroyed by fire President Harding's correspondence and documents, both official and unofficial. Upon her return to Marion, Ohio, Mrs. Harding hired a number of secretaries to collect and burn President Harding's personal papers. According to Mrs. Harding, she took these actions to protect her husband's legacy. The remaining papers were held and kept from public view by the Harding Memorial Association in Marion.\n\n1945: Franklin D. Roosevelt\nOn March 29, 1945, Franklin D. Roosevelt went to the Little White House in Warm Springs, Georgia, to rest before his anticipated appearance at the founding conference of the United Nations in late April in San Francisco. At around 1:00 pm on April 12, Roosevelt said, \"I have a terrific pain in the back of my head,\" which were his last words. He then slumped forward in his chair, unconscious, and was carried into his bedroom. The president's attending cardiologist, Howard Bruenn, diagnosed a massive cerebral hemorrhage (stroke). The 63-year-old Roosevelt died a few hours later, without regaining consciousness. As Allen Drury later said, \"so ended an era, and so began another.\" After Roosevelt's death, an editorial in The New York Times declared, \"Men will thank God on their knees a hundred years from now that Franklin D. Roosevelt was in the White House.\"\nIn his later years at the White House, when Roosevelt was increasingly overworked, his daughter Anna Roosevelt Boettiger had moved in to provide her father companionship and support. Anna had also arranged for her father to meet with his former mistress, the then widowed Lucy Mercer Rutherfurd. A close friend of both Roosevelt and Mercer who was present, Elizabeth Shoumatoff, rushed Mercer away to avoid negative publicity and implications of infidelity. When Eleanor heard about her husband's death, she was also faced with the news that Anna had been arranging these meetings with Mercer and that Mercer had been with Franklin when he died.\nOn the morning of April 13, Roosevelt's body was placed in a flag-draped coffin and loaded onto the presidential train. After a White House funeral on April 14, Roosevelt was transported back to Hyde Park by train, guarded by four servicemen, one each from the Army, Navy, Marines, and Coast Guard. As was his wish, Roosevelt was buried in the Rose Garden of the Springwood estate, the Roosevelt family home in Hyde Park on April 15. Eleanor died in November 1962 and was buried next to him.\nRoosevelt's death was met with shock and grief across the U.S. and around the world. His declining health had not been known to the general public. Roosevelt had been president for more than 12 years, longer than any other person, and had led the country through some of its greatest crises to the impending defeat of Nazi Germany and within sight of the defeat of Japan as well.\nLess than a month after his death, on May 8, the war in Europe ended. President Harry S. Truman dedicated Victory in Europe Day and its celebrations to Roosevelt's memory, and kept the flags across the U.S. at half-staff for the remainder of the 30-day mourning period. In doing so, Truman said that his only wish was \"that Franklin D. Roosevelt had lived to witness this day.\"\n\n1963: John F. Kennedy\nThe most recent U.S. president to die in office is John F. Kennedy, who was assassinated on November 22, 1963, in Dallas, Texas. He was fatally shot by Lee Harvey Oswald, who fired three shots from a sixth floor window of the Texas School Book Depository at 12:30 p.m. as the presidential motorcade passed through Dealey Plaza. Riding in the vehicle with the president were First Lady Jackie Kennedy, Texas governor John Connally, and Connally's wife Nellie; Governor Connally was also seriously wounded in the attack. The motorcade rushed to Parkland Memorial Hospital, where Kennedy was pronounced dead about 30 minutes later, at the age of 46. Connally recovered from his injuries.\nVice President Lyndon B. Johnson, who was a few cars behind the president in the motorcade, became U.S. president upon Kennedy's death. He took the presidential oath of office onboard Air Force One as it sat on the runway at Dallas Love Field. Oswald was arrested by the Dallas Police Department that afternoon, and was charged under Texas state law with the murder of Kennedy, as well as that of Dallas policeman J. D. Tippit, who had been fatally shot a short time after the assassination. Two days later, on November 24, 1963, as live television cameras were covering his transfer from the city jail to the county jail, Oswald was fatally shot in the basement of Dallas Police Headquarters by Dallas nightclub operator Jack Ruby. Ruby was convicted of Oswald's murder, though it was later overturned on appeal, and Ruby died in prison in 1967 while awaiting a new trial.\nIn 1964, after a 10-month investigation into the assassination, the Warren Commission concluded that President Kennedy was assassinated by Lee Harvey Oswald and that Oswald had acted entirely alone. It also concluded that Jack Ruby acted alone when he killed Oswald in police custody. Nonetheless, speculation over \"what really happened\" on November 22, 1963, in Dallas captured the public imagination during the decades that followed. Polls conducted from 1966 to 2004 found that as many as 80 percent of Americans have suspected that there was a criminal conspiracy or cover-up. Numerous books, films, television specials and websites have examined the assassination in minute detail, and numerous conspiracy theories have been advanced. Parties as varied as the FBI, the CIA, the Mafia, the Cuban and the Soviet governments, along with Kennedy's successor, Lyndon Johnson, have been identified as Suspect. In an article published prior to the 50th anniversary of Kennedy's assassination, author Vincent Bugliosi estimates that a total of 42 groups, 82 assassins, and 214 people have been accused in conspiracy theories challenging the \"lone gunman\" theory.\n\nSee also\nList of United States presidential assassination attempts\nCurse of Tippecanoe\n\nNotes\nReferences\nBibliography\nBauer, K. Jack (1985). Zachary Taylor: Soldier, Planter, Statesman of the Old Southwest. Louisiana State University Press. ISBN 0-8071-1237-2.\nCleaves, Freeman (1939). Old Tippecanoe: William Henry Harrison and His Time. New York, NY: C. Scribner's Sons.\nLeech, Margaret (1959). In the Days of McKinley. New York: Harper and Brothers. pp. 594–600. OCLC 456809.\nMcCullough, David (1992). Truman. Simon & Schuster. ISBN 0-671-86920-5.\nMillard, Candice (2011). Destiny of the Republic. Doubleday. ISBN 978-0-385-53500-7.\nMiller, Scott (2011). The President and the Assassin. New York: Random House. pp. 56–60. ISBN 978-1-4000-6752-7.\nPeskin, Allan (1978). Garfield. Kent State University Press. ISBN 0-87338-210-2.\nVowell, Sarah (2005). Assassination Vacation. Simon and Schuster. ISBN 0-7432-6003-1.\n\nExternal links\nThe Mortal Presidency Archived June 3, 2015, at the Wayback Machine (Shapell Manuscript Foundation)\n\nJames Abram Garfield (November 19, 1831 – September 19, 1881) was the 20th president of the United States, serving from March 1881 until his assassination in September that year. A preacher, lawyer, and Civil War general, Garfield served nine terms in the United States House of Representatives and is the only sitting member of the House to be elected president. Before his candidacy for the presidency, he had been elected to the U.S. Senate by the Ohio General Assembly—a position he declined when he became president-elect.\nGarfield was born into poverty in a log cabin and grew up in northeastern Ohio. After graduating from Williams College, he studied law and became an attorney. He was a preacher in the Stone–Campbell Movement and president of the Western Reserve Eclectic Institute, affiliated with the Disciples. Garfield was elected as a Republican member of the Ohio State Senate in 1859, serving until 1861. He opposed Confederate secession, was a major general in the Union Army during the American Civil War, and fought in the battles of Middle Creek, Shiloh, and Chickamauga. He was elected to Congress in 1862 to represent Ohio's 19th district. Throughout his congressional service, he firmly supported the gold standard and gained a reputation as a skilled orator. He initially agreed with Radical Republican views on Reconstruction but later favored a Moderate Republican–aligned approach to civil rights enforcement for freedmen. Garfield's aptitude for mathematics extended to his own proof of the Pythagorean theorem, which he published in 1876.\nAt the 1880 Republican National Convention, delegates chose Garfield, who had not sought the White House, as a compromise presidential nominee on the 36th ballot. In the 1880 presidential election, he conducted a low-key front porch campaign and narrowly defeated the Democratic nominee, Winfield Scott Hancock. Garfield's accomplishments as president included his assertion of presidential authority against senatorial courtesy in executive appointments, a purge of corruption in the Post Office, and his appointment of a Supreme Court justice. He advocated for agricultural technology, an educated electorate, and civil rights for African Americans. He also proposed substantial civil service reforms, which were passed by Congress in 1883 as the Pendleton Civil Service Reform Act and signed into law by his successor, Chester A. Arthur.\nGarfield was a member of the intraparty \"Half-Breed\" faction who used the powers of the presidency to defy the powerful \"Stalwart\" Senator Roscoe Conkling from New York. He did this by appointing Blaine faction leader William H. Robertson to the lucrative post of Collector of the Port of New York. The ensuing political battle resulted in Robertson's confirmation and the resignations of Conkling and Thomas C. Platt from the Senate.\nOn July 2, 1881, Charles J. Guiteau, a disappointed and delusional office seeker, shot Garfield at the Baltimore and Potomac Railroad Station in Washington. The wound was not immediately fatal, but an infection caused by his doctors' unsanitary methods in treating the wound killed Garfield on September 19. Due to his brief tenure in office, historians tend to rank Garfield as a below-average president, though he has earned praise for anti-corruption and pro-civil rights stances.\n\nChildhood and early life\nJames Abram Garfield was born the youngest of five children on November 19, 1831, in a log cabin in Orange Township, now Moreland Hills, Ohio. Garfield's ancestor Edward Garfield migrated from Hillmorton, Warwickshire, England, to Massachusetts around 1630. James's father Abram was born in Worcester, New York, and came to Ohio to woo his childhood sweetheart, Mehitabel Ballou, only to find her married. He instead wed her sister Eliza, who was born in New Hampshire. James was named after an earlier son of Eliza and Abram who had died in infancy.\nIn early 1833, Abram and Eliza Garfield joined a Stone-Campbell church, a decision that influenced their youngest son's life. Abram died later that year, and James was raised in poverty in a household led by his strong-willed mother. He was her favorite child and the two remained close for the rest of his life. Eliza remarried in 1842, but soon left her second husband, Warren (or Alfred) Belden, and a scandalous divorce was awarded in 1850. James took his mother's side in the matter and noted Belden's 1880 death with satisfaction in his diary. Garfield also enjoyed his mother's stories about his ancestry, especially those about his Welsh great-great-grandfathers and an ancestor who served as a knight of Caerphilly Castle.\nPoor and fatherless, Garfield was mocked by his peers and became sensitive to slights throughout his life; he sought escape through voracious reading. He left home at age 16 in 1847 and was rejected for work on the only ship in port in Cleveland. Garfield instead found work on a canal boat, managing the mules that pulled it. Horatio Alger later used this labor to good effect when he wrote Garfield's campaign biography in 1880.\nAfter six weeks, illness forced Garfield to return home, and during his recuperation, his mother and a local school official secured his promise to forgo canal work for a year of school. In 1848, he began at Geauga Seminary, in nearby Chester Township, Geauga County, Ohio. Garfield later said of his childhood, \"I lament that I was born to poverty, and in this chaos of childhood, seventeen years passed before I caught any inspiration ... a precious 17 years when a boy with a father and some wealth might have become fixed in manly ways.\"\n\nEducation, marriage and early career\nGarfield attended Geauga Seminary from 1848 to 1850 and learned academic subjects for which he had not previously had time. He excelled as a student and was especially interested in languages and elocution. He began to appreciate the power a speaker had over an audience, writing that the speaker's platform \"creates some excitement. I love agitation and investigation and glory in defending unpopular truth against popular error.\" Geauga was coeducational, and Garfield was attracted to one of his classmates, Lucretia Rudolph, whom he later married. To support himself at Geauga, he worked as a carpenter's assistant and teacher. The need to go from town to town to find work as a teacher aggravated Garfield, and he developed a dislike of what he called \"place-seeking\", which became, he said, \"the law of my life.\" In later years, he astounded his friends by disregarding positions that could have been his with little politicking. Garfield had attended church more to please his mother than to worship God, but in his late teens he underwent a religious awakening. He attended many camp meetings, which led to his being born again on March 4, 1850, when he was baptized into Christ by being submerged in the icy waters of the Chagrin River.\nAfter he left Geauga, Garfield worked for a year at various jobs, including teaching jobs. Finding that some New Englanders worked their way through college, Garfield determined to do the same and sought a school that could prepare him for the entrance examinations. From 1851 to 1854, he attended the Western Reserve Eclectic Institute (later named Hiram College) in Hiram, Ohio, a school founded by and still affiliated with the Christian Church (Disciples of Christ). While there, he was most interested in the study of Greek and Latin but was inclined to learn about and discuss any new thing he encountered. Securing a position on entry as janitor, he obtained a teaching position while he was still a student there. Lucretia Rudolph also enrolled at the Institute and Garfield wooed her while teaching her Greek. He developed a regular preaching circuit at neighboring churches and, in some cases, earned one gold dollar per service. By 1854, Garfield had learned all the Institute could teach him and was a full-time teacher. Garfield then enrolled at Williams College in Williamstown, Massachusetts, as a third-year student; he received credit for two years' study at the Institute after passing a cursory examination. Garfield was also impressed with the college president, Mark Hopkins, who had responded warmly to Garfield's letter inquiring about admission. He said of Hopkins, \"The ideal college is Mark Hopkins on one end of a log with a student on the other.\" Hopkins later said of Garfield in his student days, \"There was a large general capacity applicable to any subject. There was no pretense of genius, or alternation of spasmodic effort, but a satisfactory accomplishment in all directions.\" After his first term, Garfield was hired to teach penmanship to the students of nearby Pownal, Vermont, a post Chester A. Arthur previously held.\n\nGarfield graduated Phi Beta Kappa from Williams in August 1856, was named salutatorian, and spoke at the commencement. His biographer Ira Rutkow writes that Garfield's years at Williams gave him the opportunity to know and respect those of different social backgrounds, and that, despite his origin as an unsophisticated Westerner, socially conscious New Englanders liked and respected him. \"In short,\" Rutkow writes, \"Garfield had an extensive and positive first experience with the world outside the Western Reserve of Ohio.\"\nUpon his return to Ohio, the degree from a prestigious Eastern college made Garfield a man of distinction. He returned to Hiram to teach at the Institute and in 1857 was made its principal, though he did not see education as a field that would realize his full potential. The abolitionist atmosphere at Williams had enlightened him politically, after which he began to consider politics as a career. He campaigned for Republican presidential candidate John C. Frémont in 1856. In 1858, he married Lucretia, and they had seven children, five of whom survived infancy. Soon after the wedding, he registered to read law at the office of attorney Albert Gallatin Riddle in Cleveland, though he did his studying in Hiram. He was admitted to the bar in 1861.\nLocal Republican leaders invited Garfield to enter politics upon the death of Cyrus Prentiss, the presumptive nominee for the local state senate seat. He was nominated at the party convention on the sixth ballot and was elected, serving from 1860 to 1861. Garfield's major effort in the state senate was an unsuccessful bill providing for Ohio's first geological survey to measure its mineral resources.\n\nCivil War\nAfter Abraham Lincoln's election as president, several Southern states announced their secession from the Union to form a new government, the Confederate States of America. Garfield read military texts while anxiously awaiting the war effort, which he regarded as a holy crusade against the Slave Power. In April 1861, the rebels bombarded Fort Sumter, one of the South's last federal outposts, beginning the Civil War. Although he had no military training, Garfield knew his place was in the Union Army.\nAt Governor William Dennison's request, Garfield deferred his military ambitions to remain in the legislature, where he helped appropriate the funds to raise and equip Ohio's volunteer regiments. When the legislature adjourned Garfield spent the spring and early summer on a speaking tour of northeastern Ohio, encouraging enlistment in the new regiments. Following a trip to Illinois to purchase muskets, Garfield returned to Ohio and, in August 1861, received a commission as a colonel in the 42nd Ohio Infantry regiment. The 42nd Ohio existed only on paper, so Garfield's first task was to fill its ranks. He did so quickly, recruiting many of his neighbors and former students. The regiment traveled to Camp Chase, outside Columbus, Ohio, to complete training. In December, Garfield was ordered to bring the 42nd to Kentucky, where they joined the Army of the Ohio under Brigadier General Don Carlos Buell.\n\nBuell's command\nBuell quickly assigned Garfield the task of driving Confederate forces out of eastern Kentucky, giving him the 18th Brigade for the campaign, which, besides his own 42nd, included the 40th Ohio Infantry, two Kentucky infantry regiments and two cavalry units. They departed Catlettsburg, Kentucky, in mid-December, advancing through the valley of the Big Sandy River. The march was uneventful until Union forces reached Paintsville, Kentucky, on January 6, 1862, where Garfield's cavalry engaged the rebels at Jenny's Creek. Confederate troops under Brigadier General Humphrey Marshall held the town in numbers roughly equal to Garfield's own, but Garfield positioned his troops so as to deceive Marshall into believing the rebels were outnumbered. Marshall ordered his troops to withdraw to the forks of Middle Creek, on the road to Virginia, and Garfield ordered his troops to take up the pursuit. They attacked the rebel positions on January 9, 1862, in the Battle of Middle Creek, the only pitched battle Garfield commanded personally. At the fighting's end, the Confederates withdrew from the field and Garfield sent his troops to Prestonsburg to reprovision.\n\nIn recognition of his success, Garfield was promoted to brigadier general. After Marshall's retreat, Garfield's command was the sole remaining Union force in eastern Kentucky and he announced that any men who had fought for the Confederacy would be granted amnesty if they returned to their homes, lived peaceably, and remained loyal to the Union. The proclamation was surprisingly lenient, as Garfield now believed the war was a crusade for eradication of slavery. Following a brief skirmish at Pound Gap, the last rebel units in the area were outflanked and retreated to Virginia.\nGarfield's promotion gave him command of the 20th Brigade of the Army of the Ohio, which received orders to join Major General Ulysses S. Grant's forces as they advanced on Corinth, Mississippi, in early 1862. Before the 20th Brigade arrived, however, Confederate forces under General Albert Sidney Johnston surprised Grant's men in their camps, driving them back. Garfield's troops received word of the battle and advanced quickly, joining the rest of the army on the second day to drive the Confederates back across the field and into retreat. The action, later known as the Battle of Shiloh, was the bloodiest of the war to date; Garfield was exposed to fire for much of the day, but emerged uninjured. Major General Henry W. Halleck, Grant's superior, took charge of the combined armies and advanced ponderously toward Corinth; when they arrived, the Confederates had fled.\nThat summer, Garfield suffered from jaundice and significant weight loss. He was forced to return home, where his wife nursed him back to health. While he was home, Garfield's friends worked to gain him the Republican nomination for Congress, but he refused to campaign with the delegates. He returned to military duty that autumn and went to Washington to await his next assignment. During this period of idleness, a rumor of an extramarital affair caused friction in the Garfields' marriage until Lucretia eventually chose to overlook it. Garfield repeatedly received tentative assignments that were quickly withdrawn, to his frustration. In the meantime, he served on the court-martial of Fitz John Porter for his tardiness at the Second Battle of Bull Run. He was convinced of Porter's guilt and voted with his fellow generals to convict Porter. The trial lasted almost two months, from November 1862 to January 1863, and, by its end, Garfield had procured an assignment as chief of staff to Major General William S. Rosecrans.\n\nChief of staff for Rosecrans\nGenerals' chiefs of staff were usually more junior officers, but Garfield's influence with Rosecrans was greater than usual, with duties extending beyond communication of orders to actual management of his Army of the Cumberland. Rosecrans had a voracious appetite for conversation, especially when unable to sleep; in Garfield, he found \"the first well read person in the Army\" and the ideal candidate for discussions that ran deep into the night. They discussed everything, especially religion, and the two became close despite Garfield's being 12 years his junior. Rosecrans, who had converted from Methodism to Roman Catholicism, softened Garfield's view of his faith.\nGarfield recommended that Rosecrans replace wing commanders Alexander McCook and Thomas Crittenden, as he believed they were ineffective, but Rosecrans ignored the suggestion. With Rosecrans, Garfield devised the Tullahoma Campaign to pursue and trap Confederate General Braxton Bragg in Tullahoma. After initial Union success, Bragg retreated toward Chattanooga, where Rosecrans stalled and requested more troops and supplies. Garfield argued for an immediate advance, in line with demands from Halleck and Lincoln. After a council of war and lengthy deliberations, Rosecrans agreed to attack.\nAt the ensuing Battle of Chickamauga on September 19 and 20, 1863, confusion among the wing commanders over Rosecrans's orders created a gap in the lines, resulting in a rout of the right flank. Rosecrans concluded that the battle was lost and fell back on Chattanooga to establish a defensive line. Garfield, however, thought part of the army had held and, with Rosecrans's approval, headed across Missionary Ridge to survey the scene. Garfield's hunch was correct. Consequently, his ride became legendary and Rosecrans's error reignited criticism about the latter's leadership. While Rosecrans's army had avoided disaster, they were stranded in Chattanooga, surrounded by Bragg's army. Garfield sent a telegram to Secretary of War Edwin M. Stanton alerting Washington to the need for reinforcements to avoid annihilation. Lincoln and Halleck responded to the request for reinforcements by sending 20,000 troops to Garfield by rail within nine days. In the meantime, Grant was promoted to command of the western armies and quickly replaced Rosecrans with George H. Thomas. Garfield was ordered to report to Washington, where he was promoted to major general. According to historian Jean Edward Smith, Grant and Garfield had a \"guarded relationship\" since Grant promoted Thomas, rather than Garfield, to command of the Army of the Cumberland after Rosecrans's dismissal.\n\nCongressional career\nElection in 1862; Civil War years\nWhile he served in the Army in early 1862, friends of Garfield approached him about running for Congress from Ohio's newly redrawn and heavily Republican 19th district. He worried that he and other state-appointed generals would receive obscure assignments, and running for Congress would allow him to resume his political career. That the new Congress would not hold its first regular session until December 1863 allowed him to continue his war service for a time. Home on medical leave, he refused to campaign for the nomination, leaving that to political managers who secured it at the local convention in September 1862 on the eighth ballot. In the October general election, he defeated D.B. Woods by a two-to-one margin for a seat in the 38th Congress.\nDays before his Congressional term began, Garfield lost his eldest daughter, three-year-old Eliza, and became anxious and conflicted, saying his \"desolation of heart\" might require his return to \"the wild life of the army.\" He also assumed that the war would end before his joining the House, but it had not, and he felt strongly that he belonged in the field, rather than in Congress. He also thought he could expect a favorable command, so he decided to see President Lincoln. During their meeting, Lincoln recommended he take his House seat, as there was an excess of generals and a shortage of administration congressmen, especially those with knowledge of military affairs. Garfield accepted this recommendation and resigned his military commission to do so.\nGarfield met and befriended Treasury Secretary Salmon P. Chase, who saw Garfield as a younger version of himself. The two agreed politically and both were part of the Radical wing of the Republican Party. Once he took his seat in December 1863, Garfield was frustrated at Lincoln's reluctance to press the South hard. Many radicals, led in the House by Pennsylvania's Thaddeus Stevens, wanted rebel-owned lands confiscated, but Lincoln threatened to veto any bill that proposed to do so on a widespread basis. In debate on the House floor, Garfield supported such legislation and, discussing England's Glorious Revolution, hinted that Lincoln might be thrown out of office for resisting it. Garfield had supported Lincoln's Emancipation Proclamation and marveled at the \"strange phenomenon in the world's history, when a second-rate Illinois lawyer is the instrument to utter words which shall form an epoch memorable in all future ages.\"\nGarfield not only favored the abolition of slavery, but also believed the leaders of the rebellion had forfeited their constitutional rights. He supported the confiscation of Southern plantations and even exile or execution of rebellion leaders as a means to ensure a permanent end to slavery. Garfield felt Congress had an obligation \"to determine what legislation is necessary to secure equal justice to all loyal persons, without regard to color.\" He was more supportive of Lincoln when he took action against slavery.\nGarfield showed leadership early in his congressional career; he was initially the only Republican vote to terminate the use of bounties in military recruiting. Some financially able recruits had used the bounty system to buy their way out of service (called commutation), which Garfield considered reprehensible. He gave a speech pointing out the flaws in the existing conscription law: 300,000 recruits had been called upon to enlist, but barely 10,000 had done so, with the remainder claiming exemption, providing money, or recruiting a substitute. Lincoln appeared before the Military Affairs committee on which Garfield served, demanding a more effective bill; even if it cost him reelection, Lincoln was confident he could win the war before his term expired. After many false starts, Garfield, with Lincoln's support, procured the passage of a conscription bill that excluded commutation.\nUnder Chase's influence, Garfield became a staunch proponent of a dollar backed by a gold standard, and strongly opposed the \"greenback\". He also accepted the necessity of suspension of payment in gold or silver during the Civil War with strong reluctance. He voted with the Radical Republicans in passing the Wade–Davis Bill, designed to give Congress more authority over Reconstruction, but Lincoln defeated it with a pocket veto.\nGarfield did not consider Lincoln very worthy of reelection, but there seemed to be no viable alternative. \"He will probably be the man, though I think we could do better\", he said. Garfield attended the party convention and promoted Rosecrans as Lincoln's running mate, but delegates chose Military Governor of Tennessee Andrew Johnson. Lincoln was reelected, as was Garfield. By then, Chase had left the Cabinet and been appointed Chief Justice, and his relations with Garfield became more distant.\nGarfield took up the practice of law in 1865 to improve his personal finances. His efforts took him to Wall Street where, the day after Lincoln's assassination, a riotous crowd drew him into an impromptu speech to calm their passions: \"Fellow citizens! Clouds and darkness are round about Him! His pavilion is dark waters and thick clouds of the skies! Justice and judgment are the establishment of His throne! Mercy and truth shall go before His face! Fellow citizens! God reigns, and the Government at Washington still lives!\" The speech, with no mention or praise of Lincoln, was, according to Garfield biographer Robert G. Caldwell, \"quite as significant for what it did not contain as for what it did.\" In the following years, Garfield had more praise for Lincoln; a year after Lincoln's death, Garfield said, \"Greatest among all these developments were the character and fame of Abraham Lincoln,\" and in 1878 he called Lincoln \"one of the few great rulers whose wisdom increased with his power\".\nWhen in Washington, Garfield attended Vermont Avenue Christian Church, which later became National City Christian Church, a building constructed and funded by the Disciples.\n\nReconstruction\nIn 1864, the U.S. Senate passed the 13th Amendment, which abolished slavery throughout the Union. The bill failed to pass the House by a two-thirds majority until January 31, 1865, when it was then sent to the states for ratification. The Amendment opened other issues concerning African American civil rights. Garfield asked, \"[What] is freedom? Is it the bare privilege of not being chained?...If this is all, then freedom is a bitter mockery, a cruel delusion.\"\nGarfield supported black suffrage as firmly as he supported abolition. President Johnson sought the rapid restoration of the Southern states during the months between his accession and the meeting of Congress in December 1865; Garfield hesitantly supported this policy as an experiment. Johnson, an old friend, sought Garfield's backing and their conversations led Garfield to assume Johnson's differences with Congress were not large. When Congress assembled in December (to Johnson's chagrin, without the elected representatives of the Southern states, who were excluded), Garfield urged conciliation on his colleagues, although he feared that Johnson, a former Democrat, might join other Democrats to gain political control. Garfield foresaw conflict even before February 1866, when Johnson vetoed a bill to extend the life of the Freedmen's Bureau, charged with aiding the former slaves. By April, Garfield had concluded that Johnson was either \"crazy or drunk with opium.\"\n\nThe conflict between Congress and President Johnson was the major issue of the 1866 campaign, with Johnson taking to the campaign trail in a Swing Around the Circle and Garfield facing opposition within the Republican party in his home district. With the South still disenfranchised and Northern public opinion behind the Republicans, they gained a two-thirds majority in both houses of Congress. Garfield, having overcome his challengers at the district nominating convention, won reelection easily.\nGarfield opposed the proposed impeachment of Johnson initially when Congress convened in December 1866, but supported legislation to limit Johnson's powers, such as the Tenure of Office Act, which restricted Johnson's ability to remove presidential appointees. Distracted by committee duties, Garfield spoke about these bills rarely, but was a loyal Republican vote against Johnson.\nOn January 7, 1867, Garfield voted in support of the resolution that launched the first impeachment inquiry against Johnson (run by the House Committee on the Judiciary). On December 7, 1867, he voted against the unsuccessful resolution to impeach Johnson that the House Committee on the Judiciary had sent the full House. On January 27, 1868, he voted to pass the resolution that authorized the second impeachment inquiry against Johnson (run by the House Select Committee on Reconstruction). Due to a court case, he was absent on February 24, 1868, when the House impeached Johnson, but gave a speech aligning himself with Thaddeus Stevens and others who sought Johnson's removal shortly thereafter. Garfield was present on March 2 and 3, 1868, when the House voted on specific articles of impeachment, and voted in support of all 11 articles. During the March 2 debate on the articles, Garfield argued that what he characterized as Johnson's attempts to render Ulysses S. Grant, William Tecumseh Sherman, and William H. Emory personal tools of his demonstrated Johnson's intent to disregard the law and override the Constitution, suggesting that Johnson's trial perhaps could be expedited to last only a day in order to hasten his removal. When Johnson was acquitted in his trial before the Senate, Garfield was shocked and blamed the outcome on the trial's presiding officer, Chief Justice Chase, his onetime mentor.\nBy the time Grant succeeded Johnson in 1869, Garfield had moved away from the remaining radicals (Stevens, their leader, had died in 1868). By this time, many in the Republican Party wanted to remove the \"Negro question\" from national affairs. Garfield hailed the ratification of the 15th Amendment in 1870 as a triumph and favored Georgia's readmission to the Union as a matter of right, not politics. An influential Republican, Garfield said, \"[The] Fifteen Amendment confers on the African race the care of its own destiny. It places their fortunes in their own hands.\" In 1871, Congress took up the Ku Klux Klan Act, which was designed to combat attacks on African Americans' suffrage rights. Garfield opposed the act, saying, \"I have never been more perplexed by a piece of legislation.\" He was torn between his indignation at the Klan, whom he called \"terrorists\", and his concern for the power given the president to enforce the act through suspension of habeas corpus.\n\nTariffs and finance\nThroughout his political career, Garfield favored the gold standard and decried attempts to increase the money supply through the issuance of paper money not backed by gold, and later, through the free and unlimited coinage of silver. In 1865, he was put on the House Ways and Means Committee, a long-awaited opportunity to focus on financial and economic issues. He reprised his opposition to the greenback, saying, \"Any party which commits itself to paper money will go down amid the general disaster, covered with the curses of a ruined people.\" In 1868 Garfield gave a two-hour speech on currency in the House, which was widely applauded as his best oratory to that point; in it, he advocated a gradual resumption of specie payments, that is, the government paying out silver and gold, rather than paper money that could not be redeemed.\nTariffs had been raised to high levels during the Civil War. Afterward, Garfield, who made a close study of financial affairs, advocated moving toward free trade, though the standard Republican position was a protective tariff that would allow American industries to grow. This break with his party likely cost him his place on the Ways and Means Committee in 1867, and though Republicans held the majority in the House until 1875, Garfield remained off that committee. Garfield came to chair the powerful House Appropriations Committee, but it was Ways and Means, with its influence over fiscal policy, that he really wanted to lead. One reason he was denied a place on Ways and Means was the opposition of the influential Republican editor Horace Greeley.\n\nStarting in January 1870, Garfield, then chairman of the House Banking Committee, led an investigation into the Black Friday Gold Panic scandal. In 1869, during Grant's first term in office, two New York conspirators, Jay Gould and James Fisk, launched a scheme to corner the gold market. The conspiracy was broken on Friday, September 24, 1869, when Grant and Treasury Secretary George Boutwell released gold into the market, causing widespread financial panic. During the investigation, rumors spread that Grant's family might have been involved. In order not to force Grant's wife to testify, Garfield had a private meeting with Grant at the White House. When Garfield showed Grant testimony about him and his family, Grant thanked Garfield but refused to read it or give a response. Grant personally resented Garfield for investigating Black Friday and his wife Julia concerning possible involvement in the scandal.\nGarfield's investigation and final majority report, released on September 12, 1870, were thorough but found no indictable offenses and exonerated Grant and Julia of wrongdoing. Garfield thought the scandal was enabled by the greenbacks that financed the speculation. Garfield was not at all enthused about President Grant's reelection in 1872—until Greeley, who emerged as the candidate of the Democrats and Liberal Republicans, became the only serious alternative. Garfield said, \"I would say Grant was not fit to be nominated and Greeley is not fit to be elected.\" Both Grant and Garfield were overwhelmingly reelected.\n\nCrédit Mobilier scandal; salary grab\nThe Crédit Mobilier of America scandal involved corruption in the financing of the Union Pacific Railroad, part of the transcontinental railroad which was completed in 1869. Union Pacific officers and directors secretly purchased control of the Crédit Mobilier of America company, then contracted with it to undertake construction of the railroad. The railroad paid the company's grossly inflated invoices with federal funds appropriated to subsidize the project, and the company was allowed to purchase Union Pacific securities at par value, well below the market rate. Crédit Mobilier showed large profits and stock gains, and distributed substantial dividends. The high expenses meant Congress was called upon to appropriate more funds. One of the railroad officials who controlled Crédit Mobilier was also a congressman, Oakes Ames of Massachusetts. He offered some of his colleagues the opportunity to buy Crédit Mobilier stock at par value, well below what it sold for on the market, and the railroad got its additional appropriations.\n\nThe story broke in July 1872, in the middle of the presidential campaign. Among those named were Vice President Schuyler Colfax, Massachusetts Senator Henry Wilson (the Republican candidate for vice president), Speaker James G. Blaine of Maine, and Garfield. Greeley had little luck taking advantage of the scandal. When Congress reconvened after the election, Blaine, seeking to clear his name, demanded a House investigation. Evidence before the special committee exonerated Blaine. Garfield had said in September 1872 that Ames had offered him stock but he had repeatedly refused it. Testifying before the committee in January, Ames said he had offered Garfield ten shares of stock at par value, but that Garfield had never taken them or paid for them, though a year passed, from 1867 to 1868, before Garfield had finally refused. Appearing before the committee on January 14, 1873, Garfield confirmed much of this. Ames testified several weeks later that Garfield agreed to take the stock on credit, and that it was paid for by the company's huge dividends. The two men differed over $300 that Garfield received and later paid back, with Garfield deeming it a loan and Ames a dividend.\nGarfield's biographers have been unwilling to exonerate him in the scandal. Allan Peskin writes, \"Did Garfield lie? Not exactly. Did he tell the truth? Not completely. Was he corrupted? Not really. Even Garfield's enemies never claimed that his involvement in the affair influenced his behavior.\" Rutkow writes, \"Garfield's real offense was that he knowingly denied to the House investigating committee that he had agreed to accept the stock and that he had also received a dividend of $329.\" Caldwell suggests Garfield \"told the truth [before the committee, but] certainly failed to tell the whole truth, clearly evading an answer to certain vital questions and thus giving the impression of worse faults than those of which he was guilty.\" That Crédit Mobilier was a corrupt organization had been a badly kept secret, even mentioned on the floor of Congress, and editor Sam Bowles wrote at the time that Garfield, in his positions on committees dealing with finance, \"had no more right to be ignorant in a matter of such grave importance as this, than the sentinel has to snore on his post.\"\nAnother issue that caused Garfield trouble in his 1874 reelection bid was the so-called \"Salary Grab\" of 1873, which increased the compensation for members of Congress by 50%, retroactive to 1871. As chairman of the Appropriations Committee, Garfield was responsible for shepherding the appropriations bill through the House; during the debate in February 1873, Massachusetts Representative Benjamin Butler offered the increase as an amendment, and despite Garfield's opposition, it passed the House and eventually became law. The law was very popular in the House, as almost half the members were lame ducks, but the public was outraged, and many of Garfield's constituents blamed him, though he personally refused to accept the increase. In a bad year for Republicans, who lost control of the House for the first time since the Civil War, Garfield had his closest congressional election, winning with only 57% of the vote.\n\nFloor leader; Hayes administration\nThe Democratic takeover of the House of Representatives in 1875 meant the loss of Garfield's chairmanship of the Appropriations Committee, though the Democrats did put him on the Ways and Means Committee. With many of his leadership rivals defeated in the 1874 Democratic landslide, and Blaine elected to the Senate, Garfield was seen as the Republican floor leader, and the likely Speaker, should the party regain control of the chamber.\nGarfield thought the land grants given to expanding railroads was an unjust practice. He also opposed monopolistic practices by corporations, as well as the power sought by workers' unions. He supported the proposed establishment of the United States civil service as a means of ridding officials of the annoyance of aggressive office seekers. He especially wished to eliminate the practice of forcing government workers, in exchange for their positions, to kick back a percentage of their wages as political contributions.\nAs the 1876 presidential election approached, Garfield was loyal to the candidacy of Senator Blaine, and fought for the former Speaker's nomination at the 1876 Republican National Convention in Cincinnati. When it became clear, after six ballots, that Blaine could not prevail, the convention nominated Ohio Governor Rutherford B. Hayes. Although Garfield had supported Blaine, he had kept good relations with Hayes, and wholeheartedly supported the governor. Garfield had hoped to retire from politics after his term expired to devote himself full-time to the practice of law, but to help his party, he sought re-election, and won it easily that October. Any celebration was short-lived, as Garfield's youngest son, Neddie, fell ill with whooping cough shortly after the congressional election, and soon died.\n\nWhen Hayes appeared to have lost the presidential election the following month to Democrat Samuel Tilden, the Republicans launched efforts to reverse the results in South Carolina, Louisiana, and Florida, where they held the governorship. If Hayes won all three states, he would take the election by a single electoral vote. Grant asked Garfield to serve as a \"neutral observer\" of the recount in Louisiana. The observers soon recommended to the state electoral commissions that Hayes be declared the winner—Garfield recommended the entire vote of West Feliciana Parish, which had given Tilden a sizable majority, be thrown out. The Republican governors of the three states certified that Hayes had won their states, to the outrage of Democrats, who had the state legislatures submit rival returns, and threatened to prevent the counting of the electoral vote—under the Constitution, Congress is the final arbiter of the election. Congress then established an Electoral Commission, consisting of eight Republicans and seven Democrats, to determine the winner. Despite his objection to the Commission, Garfield was appointed to it. He felt Congress should count the vote and proclaim Hayes victorious. Hayes emerged the victor by a party line vote of 8–7. In exchange for recognizing Hayes as president, Southern Democrats secured the removal of federal troops from the South, ending Reconstruction.\nAlthough an Ohio Senate seat would be vacated by the resignation of John Sherman to become Treasury Secretary, Hayes needed Garfield's expertise to protect him from the agenda of a hostile Congress, and asked him not to seek it. Garfield agreed. As Hayes's key legislator in the House, he gained considerable prestige and respect for his role there. When Congress debated the Bland–Allison Act, to have the government purchase large quantities of silver and strike it into legal tender dollar coins, Garfield opposed it as a deviation from the gold standard; it was enacted over Hayes's veto in February 1878.\nIn 1876, Garfield purchased the property in Mentor that reporters later dubbed Lawnfield, where he conducted the first successful front porch campaign for the presidency. Hayes suggested that Garfield run for governor in 1879, seeing that as a road likely to take Garfield to the White House. Garfield preferred to seek election as a U.S. senator. Rivals were spoken of for the seat, such as Secretary Sherman, but he had presidential ambitions (for which he sought Garfield's support), and other candidates fell by the wayside. The General Assembly elected Garfield to the Senate in January 1880, though his term was not scheduled to commence until March 4, 1881.\n\nLegal career and other activities\nIn 1865, Garfield became a partner in the law firm of a fellow Disciple of Christ, Jeremiah Black. They had much in common, except politics: Black was an avid Democrat, having served in the cabinet of President James Buchanan. The next year, Black was retained by some pro-Confederate northern civilians who had been found guilty of treason in a military court and sentenced to death. Black saw an opportunity to strike a blow against military courts and the Republicans. He had heard Garfield's military speeches, and learned of not only his oratory skills but also his resistance to expansive powers of military commissions. Black assigned the case to Garfield one week before arguments were to be made before the U. S. Supreme Court. When Black warned him of the political peril, Garfield responded, \"It don't make any difference. I believe in English liberty and English law.\" In this landmark case, Ex parte Milligan, Garfield successfully argued that civilians could not be tried before military tribunals, despite a declaration of martial law, as long as civil courts were still operating. In his first court appearance, Garfield's oral argument lasted over two hours, and though his wealthy clients refused to pay him, he had established himself as a preeminent lawyer.\nDuring Grant's first term, Garfield was discontented with public service and in 1872 again pursued opportunities in the law. But he declined a partnership offer from a Cleveland law firm when told his prospective partner was of \"intemperate and licentious\" reputation. In 1873, after Chase's death, Garfield appealed to Grant to appoint Justice Noah H. Swayne Chief Justice, but Grant appointed Morrison R. Waite.\n\nIn 1871, Garfield traveled to Montana Territory to negotiate the removal of the Bitterroot Salish tribe to the Flathead Indian Reservation. Having been told that the people would happily move, Garfield expected an easy task. Instead, he found the Salish determined to stay in their Bitterroot Valley homeland. His attempts to coerce Chief Charlo to sign the agreement nearly brought about a military clash. In the end, he convinced two subchiefs to sign and move to the reservation with a few of the Salish people. Garfield never convinced Charlo to sign, although the official treaty document voted on by Congress bore his forged mark.\nIn 1876, Garfield developed a trapezoid proof of the Pythagorean theorem, which was published in the New England Journal of Education. Mathematics historian William Dunham wrote that Garfield's trapezoid work was \"really a very clever proof.\" According to the Journal, Garfield arrived at the proof \"in mathematical amusements and discussions with other members of congress.\"\nAfter his conversion experience in 1850, religious inquiry was a high priority for Garfield. He read widely and moved beyond the confines of his early experience as a member of the Disciples of Christ. His new, broader perspective was rooted in his devotion to freedom of inquiry and his study of history. The intensity of Garfield's religious thought was also influenced by his experience in combat and his interaction with voters.\n\nPresidential election of 1880\nRepublican nomination\nHaving just been elected to the Senate with John Sherman's support, Garfield was committed to Sherman for the 1880 Republican presidential nomination. Before the convention began, however, a few Republicans, including Wharton Barker of Philadelphia, thought Garfield the best choice for the nomination. Garfield denied any interest in the position, but the attention was enough to make Sherman suspicious of his lieutenant's ambitions. Besides Sherman, the early favorites for the nomination were Blaine, former President Grant; several other candidates attracted delegates as well.\nThe Republican Party at the time was split into two factions: the \"Stalwarts\", who supported the existing federal government patronage system, and the \"Half-Breeds\", who wanted civil service reform. As the convention began, New York Senator Roscoe Conkling, floor leader for the Stalwarts, who supported former President Ulysses S. Grant, proposed that the delegates pledge to back the eventual nominee in the general election. When three West Virginia delegates declined to be so bound, Conkling sought to expel them from the convention. Garfield rose to defend the men, giving a passionate speech in defense of their right to reserve judgment. The crowd turned against Conkling, and he withdrew the motion. The performance delighted Garfield's boosters, who were then convinced he was the only one who could attract a majority of the delegates' votes.\nAfter speeches in favor of the other front-runners, Garfield rose to place Sherman's name in nomination; his speech was well-received, but the delegates mustered little excitement for Sherman as the next president. The first ballot showed Grant leading with 304 votes to Blaine's 284, and Sherman's 93 votes placed him in a distant third. Subsequent ballots demonstrated a deadlock between Grant and Blaine, with neither having the 379 votes needed for nomination. Jeremiah McLain Rusk, a member of the Wisconsin delegation, and Benjamin Harrison, an Indiana delegate, sought to break the deadlock by shifting a few of the anti-Grant votes to a dark horse candidate—Garfield. Garfield gained 50 votes on the 35th ballot, and a stampede began. Garfield protested to the Ohio delegation that he did not seek the nomination and would not betray Sherman, but they overruled his objections and cast their ballots for him. In the next round of voting, nearly all the Sherman and Blaine delegates shifted their support to Garfield, giving him 399 votes, and the Republican nomination. Most of the Grant forces backed the former president to the end, creating a disgruntled Stalwart minority in the party. To obtain that faction's support for the ticket, Chester A. Arthur, a former New York customs collector and member of Conkling's political machine, was chosen as the vice presidential nominee.\n\nCampaign against Hancock\nEven with a Stalwart on the ticket, animosity between the Republican factions carried over from the convention, so Garfield traveled to New York to meet with party leaders. After convincing the Stalwart crowd to put aside their differences and unite for the coming campaign, Garfield returned to Ohio, leaving the active campaigning to others, as was traditional at the time. Meanwhile, the Democrats settled on their nominee, Major General Winfield Scott Hancock of Pennsylvania, a career military officer. Hancock and the Democrats expected to carry the Solid South, while much of the North was considered safe territory for Garfield and the Republicans; most of the campaign focused on a few close states, including New York and Indiana.\nPractical differences between the candidates were few, but Republicans began the campaign with the familiar theme of waving the bloody shirt. They reminded Northern voters the Democratic Party was responsible for secession and four years of civil war, and Democrats would reverse the gains of that war, dishonor Union veterans, and pay Confederate veterans pensions out of the federal treasury. Fifteen years had passed since the end of the war, and with Union generals at the head of both tickets, the bloody shirt was of diminishing value in exciting the voters. With a few months to go before the election, the Republicans switched tactics to emphasize the tariff. Seizing on the Democratic platform's call for a \"tariff for revenue only\", Republicans told Northern workers a Hancock presidency would weaken the tariff protection that kept them in good jobs. Hancock made the situation worse when, attempting to strike a moderate stance, he said, \"The tariff question is a local question.\" The Republican ploy proved effective in uniting the North behind Garfield. Ultimately, of the more than 9.2 million popular votes cast, fewer than 2,000 separated the two candidates. But in the Electoral College, Garfield had an easy victory over Hancock, 214 to 155. The election made Garfield the only sitting member of the House ever to be elected to the presidency.\n\nPresidency (1881)\nCabinet and inauguration\nBefore his inauguration, Garfield was occupied with assembling a cabinet that might engender peace between the party's Conkling and Blaine factions. Blaine's delegates had provided much of the support for Garfield's nomination, so the Maine senator received the place of honor as Secretary of State. Blaine was not only the president's closest advisor, but he was also obsessed with knowing all that took place in the White House, and allegedly posted spies there in his absence. Garfield nominated William Windom of Minnesota as Secretary of the Treasury, William H. Hunt of Louisiana as Secretary of the Navy, Robert Todd Lincoln as Secretary of War, and Samuel J. Kirkwood of Iowa as Secretary of the Interior. New York was represented by Thomas Lemuel James as Postmaster General. Garfield appointed Pennsylvania's Wayne MacVeagh, an adversary of Blaine's, as Attorney General. Blaine tried to sabotage the appointment by convincing Garfield to name an opponent of MacVeagh, William E. Chandler, as Solicitor General under MacVeagh. Only Chandler's rejection by the Senate forestalled MacVeagh's resignation over the matter.\nBecause Garfield was distracted by cabinet maneuvering, his inaugural address was a \"compendium of platitudes\" and fell below expectations. At one high point, however, Garfield emphasized the civil rights of African-Americans, saying \"Freedom can never yield its fullness of blessings so long as the law or its administration places the smallest obstacle in the pathway of any virtuous citizen.\" After discussing the gold standard, the need for education, and an unexpected denunciation of Mormon polygamy, the speech ended. The crowd applauded, but the speech, according to Peskin, \"however sincerely intended, betrayed its hasty composition by the flatness of its tone and the conventionality of its subject matter.\"\nGarfield's appointment of James infuriated Conkling, a factional opponent of the Postmaster General, who demanded a compensatory appointment for his faction, such as the position of Secretary of the Treasury. The resulting squabble occupied much of Garfield's brief presidency. The feud with Conkling reached a climax when the president, at Blaine's instigation, nominated Conkling's enemy, Judge William H. Robertson, to be Collector of the Port of New York. This was one of the prize patronage positions below cabinet level and was then held by Edwin A. Merritt. Conkling raised the time-honored principle of senatorial courtesy in an attempt to defeat the nomination, to no avail. Garfield, who believed the practice was corrupt, would not back down and threatened to withdraw all nominations unless Robertson was confirmed, intending to \"settle the question whether the president is registering clerk of the Senate or the Executive of the United States.\" Ultimately, Conkling and his New York colleague, Senator Thomas C. Platt, resigned their Senate seats to seek vindication but found only further humiliation when the New York legislature elected others in their places. Robertson was confirmed as Collector and Garfield's victory was clear. To Blaine's chagrin, the victorious Garfield returned to his goal of balancing the interests of party factions and nominated a number of Conkling's Stalwart friends to offices.\nWith his cabinet complete, Garfield had to contend with myriad office seekers. He exclaimed, \"My God! What is there in this place that a man should ever get into it.\" Garfield's family happily settled into the White House, but he found presidential duties exasperating.\n\nRefinance of national debt\nGarfield ordered the Secretary of the Treasury William Windom to refund (refinance) the national debt by calling in outstanding U.S. bonds paying 6% interest. Holders would have the option of accepting cash or new bonds at 3%, closer to the interest rates of the time. Taxpayers were saved an estimated $10 million. By comparison, federal expenditures in 1881 were below $261 million (~$7.09 billion in 2023).\n\nSupreme Court nomination\nIn 1880, President Hayes had nominated Stanley Matthews to the Supreme Court but the Senate declined to act on the nomination. In March 1881, Garfield re-nominated Matthews to the Court and the Senate confirmed Matthews by a vote of 24–23. According to The New York Times, \"opposition to Matthews's Supreme Court appointment ... stemmed from his prosecution in 1859 of a newspaper editor who had assisted two runaway slaves.\" Because Matthews was \"a professed abolitionist at the time, the matter was later framed as political expediency triumphing over moral principle.\" Matthews served on the Court until his death in 1889.\n\nReforms\nGrant and Hayes had both advocated civil service reform, and by 1881 such reform associations had organized with renewed energy across the nation. Garfield sympathized with them, believing the spoils system damaged the presidency and often eclipsed more important concerns. Some reformers became disappointed when Garfield promoted limited tenure only to minor office seekers and gave appointments to his old friends.\nCorruption in the post office also cried out for reform. In April 1880, there had been a congressional investigation of corruption in the Post Office Department, where profiteering rings allegedly stole millions of dollars, securing bogus mail contracts on star routes. After obtaining contracts with the lowest bid, costs to run the mail routes would be escalated and profits would be divided among ring members. Shortly after taking office, Garfield received word of postal corruption by an alleged star route ringleader, Assistant Postmaster General Thomas J. Brady. Garfield demanded Brady's resignation and ordered prosecutions that ended in trials for conspiracy. When told that his party, including his campaign manager, Stephen W. Dorsey, was involved, Garfield directed that the corruption in the Post Office be rooted out \"to the bone\", regardless of where it might lead. Brady resigned and was indicted for conspiracy, though jury trials in 1882 and 1883 found Brady not guilty.\n\nCivil rights and education\nGarfield believed the key to improving the state of African American civil rights was government supported education. During Reconstruction, freedmen had gained citizenship and suffrage, which enabled them to participate in government, but Garfield believed their rights were being eroded by Southern white resistance and illiteracy, and he was concerned that blacks would become America's permanent \"peasantry\". He proposed a \"universal\" education system funded by the federal government. In February 1866, as a congressman from Ohio, Garfield and Ohio School Commissioner Emerson Edward White had drafted a bill for the National Department of Education. They believed that through the use of statistics they could push the US Congress to establish a federal agency for school reform. But by the time of Garfield's presidency, Congress and the northern white public had lost interest in African-American rights, and Congress did not pass federal funding for universal education during his term. Garfield also worked to appoint several African Americans to prominent positions: Frederick Douglass, recorder of deeds in Washington; Robert Elliot, special agent to the Treasury; John M. Langston, Haitian minister; and Blanche K. Bruce, register to the Treasury. Garfield believed Southern support for the Republican Party could be gained by \"commercial and industrial\" interests rather than race issues and began to reverse Hayes's policy of conciliating Southern Democrats. He appointed William H. Hunt, a Republican from Louisiana, as Secretary of the Navy. To break the hold of the resurgent Democratic Party in the Solid South, Garfield took patronage advice from Virginia Senator William Mahone of the biracial independent Readjuster Party, hoping to add the independents' strength to the Republicans' there.\n\nForeign policy and naval reform\nGarfield had little foreign policy experience, so he leaned heavily on Blaine. They agreed on the need to promote freer trade, especially within the Western Hemisphere. Garfield and Blaine believed increasing trade with Latin America would be the best way to keep the United Kingdom of Great Britain and Ireland from dominating the region. And by encouraging exports, they believed they could increase American prosperity. Garfield authorized Blaine to call for a Pan-American conference in 1882 to mediate disputes among the Latin American nations and to serve as a forum for talks on increasing trade.\nAt the same time, they hoped to negotiate a peace in the War of the Pacific then being fought by Bolivia, Chile, and Peru. Blaine favored a resolution that would result in Peru yielding no territory, but Chile by 1881 had occupied the Peruvian capital of Lima, and rejected any settlement that restored the previous status quo.\nGarfield sought to expand American influence in other areas, calling for renegotiation of the Clayton–Bulwer Treaty to allow the United States to construct a canal through Panama without British involvement and attempting to reduce British influence in the strategically located Kingdom of Hawaii. Garfield's and Blaine's plans for the United States' involvement in the world stretched even beyond the Western Hemisphere, as he sought commercial treaties with Korea and Madagascar. Garfield also considered enhancing U.S. military strength abroad, asking Navy Secretary Hunt to investigate the navy's condition with an eye toward expansion and modernization. In the end, these ambitious plans came to nothing after Garfield was assassinated. Nine countries had accepted invitations to the Pan-American conference, but the invitations were withdrawn in April 1882 after Blaine resigned from the cabinet and Arthur, Garfield's successor, cancelled the conference. Naval reform continued under Arthur, on a more modest scale than Garfield and Hunt had envisioned, ultimately ending in the construction of the Squadron of Evolution.\n\nAssassination\nGuiteau and shooting\nCharles J. Guiteau had followed various professions in his life, but in 1880 had determined to gain federal office by supporting what he expected would be the winning Republican ticket. He composed a speech, \"Garfield vs. Hancock\", and got it printed by the Republican National Committee. One means of persuading the voters in that era was through orators expounding on the candidate's merits, but with the Republicans seeking more famous men, Guiteau received few opportunities to speak. On one occasion, according to Kenneth D. Ackerman, Guiteau was unable to finish his speech due to nerves. Guiteau, who considered himself a Stalwart, deemed his contribution to Garfield's victory sufficient to justify his appointment to the position of consul in Paris, despite the fact that he spoke no French, nor any foreign language. One medical expert has since described Guiteau as possibly a narcissistic schizophrenic; neuroscientist Kent Kiehl assessed him as a clinical psychopath.\n\nOne of Garfield's more wearying duties was seeing office-seekers, and he saw Guiteau at least once. White House officials suggested to Guiteau that he approach Blaine, as the consulship was within the Department of State. Blaine also saw the public regularly, and Guiteau became a regular at these sessions. Blaine, who had no intention of giving Guiteau a position he was unqualified for and had not earned, simply said the deadlock in the Senate over Robertson's nomination made it impossible to consider the Paris consulship, which required Senate confirmation. Once the New York senators had resigned, and Robertson had been confirmed as Collector, Guiteau pressed his claim, and Blaine told him he would not receive the position.\nGuiteau came to believe he had lost the position because he was a Stalwart. He decided the only way to end the Republican Party's internecine warfare was for Garfield to die—though he had nothing personal against the president. Arthur's succession would restore peace, he felt, and lead to rewards for fellow Stalwarts, including Guiteau.\nThe assassination of Abraham Lincoln was deemed a fluke due to the Civil War, and Garfield, like most people, saw no reason the president should be guarded; his movements and plans were often printed in the newspapers. Guiteau knew Garfield would leave Washington for a cooler climate on July 2, 1881, and made plans to kill him before then. He purchased a gun he thought would look good in a museum, and followed Garfield several times, but each time his plans were frustrated, or he lost his nerve. His opportunities dwindled to one—Garfield's departure by train for New Jersey on the morning of July 2.\nGuiteau concealed himself by the ladies' waiting room at the Sixth Street Station of the Baltimore and Potomac Railroad, from where Garfield was scheduled to depart. Most of Garfield's cabinet planned to accompany him at least part of the way. Blaine, who was to remain in Washington, came to the station to see him off. The two men were deep in conversation and did not notice Guiteau before he took out his revolver and shot Garfield twice, once in the back and once in the arm. Guiteau attempted to leave the station but was quickly captured. As Blaine recognized him, Guiteau was led away, and said, \"I did it. I will go to jail for it. I am a Stalwart and Arthur will be President.\" News of his motivation to benefit the Stalwarts reached many with the news of the shooting, causing rage against that faction.\n\nTreatment and death\nGarfield was struck by two shots: one glanced off his arm while the other pierced his back, shattering a rib and embedding itself in his abdomen. \"My God, what is this?\" he exclaimed. Among those at the station was Robert Todd Lincoln, who was deeply upset, thinking back to when his father Abraham Lincoln was assassinated 16 years earlier. Garfield was taken on a mattress upstairs to a private office, where several doctors examined him. At his request, Garfield was taken back to the White House, and his wife, then in New Jersey, was sent for. Blaine sent word to Vice President Arthur in New York City, who received threats against his life because of his animosity toward Garfield and Guiteau's statements.\nAlthough Joseph Lister's pioneering work in antisepsis was known to American doctors, few of them had confidence in it, and none of his advocates were among Garfield's treating physicians. The physician who took charge at the depot and then at the White House was Doctor Willard Bliss. A noted physician and surgeon, Bliss was an old friend of Garfield, and about a dozen doctors, led by Bliss, were soon probing the wound with unsterilized fingers and instruments. Garfield was given morphine for the pain, and asked Bliss to frankly tell him his chances, which Bliss put at one in a hundred. \"Well, Doctor, we'll take that chance.\"\nOver the next few days, Garfield made some improvement, as the nation viewed the news from the capital and prayed. Although he never stood again, he was able to sit up and write several times, and his recovery was viewed so positively that a steamer was fitted out as a seagoing hospital to aid with his convalescence. He was nourished on oatmeal porridge (which he detested) and milk from a cow on the White House lawn. When told that Indian chief Sitting Bull, a prisoner of the army, was starving, Garfield said, \"Let him starve...\" initially, but a few moments later said, \"No, send him my oatmeal.\"\n\nX-ray imaging, which could have assisted physicians in precisely locating the bullet in Garfield's body, would not be invented for another 14 years. Alexander Graham Bell tried to locate the bullet with a primitive metal detector, but was unsuccessful, though the device had been effective when tested on others. But Bliss limited its use on Garfield, ensuring he remained in charge. Because Bliss insisted the bullet rested someplace it did not, the detector could not locate it. Bell shortly returned after adjusting his device, which emitted an unusual tone in the area where Bliss believed the bullet was lodged. Bliss took this as confirmation that the bullet was where he declared it to be. Bliss recorded the test as a success, saying it was: now unanimously agreed that the location of the ball has been ascertained with reasonable certainty, and that it lies, as heretofore stated, in the front wall of the abdomen, immediately over the groin, about five inches [130 mm] below and to the right of the navel.\nOne means of keeping Garfield comfortable in Washington's summer heat was one of the first successful air conditioning units: air propelled by fans over ice and then dried reduced the temperature in the sickroom by 20 °F (11 °C). Engineers from the navy, and other scientists, worked together to develop it, though there were problems to solve, such as excessive noise and increased humidity.\nOn July 23, Garfield took a turn for the worse when his temperature increased to 104 °F (40 °C); doctors, concerned by an abscess at the wound, inserted a drainage tube. This initially helped, and the bedridden Garfield held a brief cabinet meeting on July 29; members were under orders from Bliss to discuss nothing that might excite Garfield. Doctors probed the abscess, hoping to find the bullet; they likely made the infections worse. Garfield performed only one official act in August, signing an extradition paper. By the end of the month, he was much feebler than he had been, and his weight had decreased from 210 pounds (95 kg) to 130 pounds (59 kg).\nGarfield had long been anxious to escape hot, unhealthy Washington, and in early September the doctors agreed to move him to Elberon, part of Long Branch, New Jersey, where his wife had recovered earlier in the summer. He left the White House for the last time on September 5, traveling in a specially cushioned railway car; a spur line to the Francklyn Cottage, a seaside mansion given over to his use, was built in a night by volunteers. After arriving in Elberon the next day, Garfield was moved from the train car to a bedroom where he could see the ocean as officials and reporters maintained what became (after an initial rally) a death watch. Garfield's personal secretary, Joe Stanley Brown, wrote forty years later, \"to this day I cannot hear the sound of the low slow roll of the Atlantic on the shore, the sound which filled my ears as I walked from my cottage to his bedside, without recalling again that ghastly tragedy.\"\n\nOn September 18, Garfield asked Colonel A.F. Rockwell, a friend, if he would have a place in history. Rockwell assured him he would and told Garfield he had much work still before him. But his response was, \"No, my work is done.\" The following day, Garfield, then suffering also from pneumonia and hypertension, marveled that he could not pick up a glass despite feeling well and went to sleep without discomfort. He awoke that evening around 10:15 p.m. complaining of great pain in his chest to his chief of staff General David Swaim, who was watching him, as he placed his hand over his heart. The president then requested a drink of water from Swaim. After finishing his glass, Garfield said, \"Oh Swaim, this terrible pain—press your hand on it.\" As Swaim put his hand on Garfield's chest, Garfield's hands went up reflexively. Clutching his heart, he exclaimed, \"Oh, Swaim, can't you stop this? Oh, oh, Swaim!\" Those were Garfield's last words. Swaim ordered another attendant to send for Bliss, who found Garfield unconscious. Despite efforts to revive him, Garfield never awoke, and he was pronounced dead at about 10:30 p.m. Learning from a reporter of Garfield's death the following day, Chester A. Arthur took the presidential oath of office administered by New York Supreme Court Justice John R. Brady.\nAccording to some historians and medical experts, Garfield might have survived his wounds had the doctors attending him had at their disposal today's medical research, knowledge, techniques, and equipment. Standard medical practice at the time dictated that priority be given to locating the path of the bullet. Several of his doctors inserted their unsterilized fingers into the wound to probe for the bullet, a common practice in the 1880s. Historians agree that massive infection was a significant factor in Garfield's demise. Biographer Peskin said medical malpractice did not contribute to Garfield's death; the inevitable infection and blood poisoning that would ensue from a deep bullet wound resulted in damage to multiple organs and spinal fragmentation. Rutkow, a professor of surgery at the University of Medicine and Dentistry of New Jersey, has argued that starvation also played a role. Rutkow suggests \"Garfield had such a nonlethal wound. In today's world, he would have gone home in a matter of two or three days.\" The conventional narrative regarding Garfield's post-shooting medical condition was challenged by Theodore Pappas and Shahrzad Joharifard in a 2013 article in The American Journal of Surgery. They argued that Garfield died from a late rupture of a splenic artery pseudoaneurysm, which developed secondary to the path of the bullet adjacent to the splenic artery. They also argued that his sepsis was actually caused by post-traumatic acute acalculous cholecystitis. Based on the autopsy report, the authors speculate that his gallbladder subsequently ruptured, leading to the development of a large bile-containing abscess adjacent to the gallbladder. Pappas and Joharifard say this caused the septic decline in Garfield's condition that was visible starting from July 23, 1881. Pappas and Joharifard also state that they don't believe that Garfield's doctors could have saved him even if they had been aware of his cholecystitis, since the first successful cholecystectomy (surgical removal of the gallbladder) was performed a year after Garfield's death.\nGuiteau was indicted on October 14, 1881, for the murder of the president. During his trial, Guiteau declared that he was not responsible for Garfield's death, admitting to the shooting but not the killing. In his defense, Guiteau wrote: \"General Garfield died from malpractice. According to his own physicians, he was not fatally shot. The doctors who mistreated him ought to bear the odium of his death, and not his assailant. They ought to be indicted for murdering James A. Garfield, and not me.\" After a chaotic trial in which Guiteau often interrupted and argued, and in which his counsel used the insanity defense, the jury found him guilty on January 25, 1882, and he was sentenced to death by hanging. Guiteau may have had neurosyphilis, a disease that causes physiological mental impairment. He was executed on June 30, 1882.\n\nFuneral, memorials and commemorations\nGarfield's funeral train left Long Branch on the same special track that had brought him there, traveling over tracks blanketed with flowers and past houses adorned with flags. His body was transported to the Capitol and then continued on to Cleveland for burial. Shocked by his death, Marine Band leader John Philip Sousa composed the march \"In Memoriam\", which was played when Garfield's body was received in Washington, D.C. More than 70,000 citizens, some waiting over three hours, passed by Garfield's coffin as his body lay in state from September 21 to 23, 1881, at the United States Capitol rotunda; on September 25, in Cleveland, Garfield's casket was paraded down Euclid Avenue from Wilson Avenue to Public Square, with those in attendance including former presidents Grant and Hayes, and Generals William Sherman, Sheridan and Hancock. More than 150,000—a number equal to the city's population—likewise paid their respects, and Sousa's march was again played. Garfield's body was temporarily interred in the Schofield family vault in Cleveland's Lake View Cemetery until his permanent memorial was built.\nMemorials to Garfield were erected across the country. On April 10, 1882, seven months after Garfield's death, the U.S. Post Office Department issued a postage stamp in his honor. In 1884, sculptor Frank Happersberger completed a monument on the grounds of the San Francisco Conservatory of Flowers. In 1887, the James A. Garfield Monument was dedicated in Washington. Another monument, in Philadelphia's Fairmount Park, was erected in 1896. In Victoria, Australia, Cannibal Creek was renamed Garfield in his honor.\n\nOn May 19, 1890, Garfield's body was permanently interred, with great solemnity and fanfare, in a mausoleum in Lake View Cemetery. Attending the dedication ceremonies were former President Hayes, President Benjamin Harrison, and future president William McKinley. Garfield's Treasury Secretary, William Windom, also attended. Harrison said Garfield was always a \"student and instructor\" and that his life works and death would \"continue to be instructive and inspiring incidents in American history\". Three panels on the monument display Garfield as a teacher, Union major general, and orator; another shows him taking the presidential oath, and a fifth shows his body lying in state at the Capitol rotunda in Washington, D.C.\nGarfield's murder by a deranged office-seeker awakened public awareness of the need for civil service reform legislation. Senator George H. Pendleton, a Democrat from Ohio, launched a reform effort that resulted in the Pendleton Act in January 1883. This act reversed the \"spoils system\" where office seekers paid up or gave political service to obtain or keep federally appointed positions. Under the act, appointments were awarded on merit and competitive examination. To ensure the reform was implemented, Congress and Arthur established and funded the Civil Service Commission. The Pendleton Act, however, covered only 10% of federal government workers. For Arthur, previously known for having been a \"veteran spoilsman\", civil service reform became his most noteworthy achievement.\nA marble statue of Garfield by Charles Niehaus was added to the National Statuary Hall Collection in the Capitol in Washington D.C., a gift from the State of Ohio in 1886.\nGarfield is honored with a life-size bronze sculpture inside the Cuyahoga County Soldiers' and Sailors' Monument in Cleveland, Ohio.\nOn March 2, 2019, the National Park Service erected exhibit panels in Washington to mark the site of his assassination.\n\nLegacy and historical view\nFor a few years after his assassination, Garfield's life story was seen as an exemplar of the American success story—that even the poorest boy might someday become President of the United States. Peskin wrote: \"In mourning Garfield, Americans were not only honoring a president; they were paying tribute to a man whose life story embodied their own most cherished aspirations.\" As the rivalry between Stalwarts and Half-Breeds faded from the scene in the late 1880s and after, so too did memories of Garfield. In the 1890s, Americans became disillusioned with politicians, and looked elsewhere for inspiration, focusing on industrialists, labor leaders, scientists, and others as their heroes. Increasingly, Garfield's short time as president was forgotten.\n\nThe 20th century saw no revival for Garfield. Thomas Wolfe deemed the presidents of the Gilded Age, including Garfield, \"lost Americans\" whose \"gravely vacant and bewhiskered faces mixed, melted, swam together\". The politicians of the Gilded Age faded from the public eye, their luster eclipsed by those who had influenced America outside of political office during that time; the robber barons, the inventors, those who had sought social reform, and others who had lived as America rapidly changed. Current events and more recent figures occupied America's attention. According to Ackerman, \"the busy Twentieth Century has made Garfield's era seem remote and irrelevant, its leaders ridiculed for their very obscurity.\"\nGarfield's biographers, and those who have studied his presidency, tend to think well of him, and that his presidency saw a promising start before its untimely end. Historian Justus D. Doenecke, while deeming Garfield a bit of an enigma, chronicles his achievements: \"by winning a victory over the Stalwarts, he enhanced both the power and prestige of his office. As a man, he was intelligent, sensitive, and alert, and his knowledge of how government worked was unmatched.\" Doenecke criticizes Garfield's dismissal of Merritt in Robertson's favor, and wonders if the president was truly in command of the situation even after the latter's confirmation. In 1931, Caldwell wrote: \"If Garfield lives in history, it will be partly on account of the charm of his personality—but also because in life and in death, he struck the first shrewd blows against a dangerous system of boss rule which seemed for a time about to engulf the politics of the nation. Perhaps if he had lived he could have done no more.\" Rutkow writes that \"James Abram Garfield's presidency is reduced to a tantalizing 'what if.'\"\nIn 2002, historian Bernard A. Weisberger said, \"[Garfield] was, to some extent, a perfect moderate. He read widely (and unobtrusively) without its visibly affecting his Christianity, his Republicanism, or his general laissez-faire orthodoxy. He was not so much a scholar in politics as a politic scholar.\" Peskin believes Garfield deserves more credit for his political career than he has received: \"True, his accomplishments were neither bold nor heroic, but his was not an age that called for heroism. His stormy presidency was brief, and in some respects, unfortunate, but he did leave the office stronger than he found it. As a public man he had a hand in almost every issue of national importance for almost two decades, while as a party leader he, along with Blaine, forged the Republican Party into the instrument that would lead the United States into the twentieth century.\"\n\nNotes\nReferences\nWorks cited\nFurther reading\nFuller, Corydon E. (2022) [1887]. Reminiscences of James A. Garfield. Hansebooks. ISBN 978-3-34807-944-0.\nGoodyear, C. W. (2023). President Garfield: From Radical to Unifier. New York, New York: Simon & Schuster.\nGraff Henry F., ed. The Presidents: A Reference History (3rd ed. 2002) online\nHammond, William A.; Ashhurst, Jr., John; Sims, J. Marion; Hodgen, John T. (December 1881). \"The Surgical Treatment of President Garfield\". The North American Review. 133 (301): 578–610. JSTOR 25101018.\nHoudek, John Thomas. \"James A. Garfield and Rutherford B. Hayes: A Study in State and National Politics\" (PhD dissertation, Michigan State University; Proquest Dissertations Publishing, 1970. 7111871).\nMenke, Richard. \"Media in America, 1881: Garfield, Guiteau, Bell, Whitman.\" Critical Inquiry 31.3 (2005): 638–664.\nMillard, Candice (2012). Destiny of the Republic: A Tale of Madness, Medicine and the Murder of a President. New York, New York: Anchor Books. ISBN 978-0-7679-2971-4.\nNorth, Ira Lutts. \"A rhetorical criticism of the speaking of James Abram Garfield, 1876-1880\" (PhD dissertation, Louisiana State University; ProQuest Dissertations Publishing, 1953. DP69446).\nRushford, Jerry Bryant. \"Political Disciple: The Relationship Between James A. Garfield And The Disciples Of Christ\" (PhD dissertation, University of California, Santa Barbara; ProQuest Dissertations Publishing, 1977. 7807029).\nSkidmore, Max J. \"James A. Garfield and Chester A. Arthur.\" in Maligned Presidents: The Late 19th Century (Palgrave Macmillan, New York, 2014) pp. 63–79.\nSutton, Thomas C. \"James A. Garfield.\" in The Presidents and the Constitution (Volume One. New York University Press, 2020) pp. 266–275.\nUhler, Kevin A. \"The demise of patronage: Garfield, the midterm election, and the passage of the Pendleton Civil Service Act\" (PhD. Diss. The Florida State University, 2011) online.\nVermilya, Daniel J. James Garfield and the Civil War: For Ohio and the Union (Arcadia Publishing, 2015).\n\nExternal links\n\nGarfield, James Abram, (1831–1881) Congressional Biography\nJames Garfield: A Resource Guide from the Library of Congress\nJames A. Garfield at the Database of Classical Scholars\n[http://millercenter.org/president/garfield Brief essays on James A. Garfield and his administration from the Miller Center of Public Affairs\n\"Life Portrait of James Garfield\", from C-SPAN's American Presidents: Life Portraits, July 26, 1999\nWorks by or about James A. Garfield at the Internet Archive\nWorks by James A. Garfield at LibriVox (public domain audiobooks) \nNotable alumni of Delta Upsilon fraternity, including Garfield\nJames A. Garfield Personal Manuscripts\nJames A. Garfield Collection at Williams College Chapin Library\nJames A. Garfield Collection at Williams College Archives and Special Collections\nOfficial medical bulletins relating to the health of U.S. President James Garfield from the U.S. National Library of Medicine. Contains medical bulletins issued by attending physicians D. Hayes Agnes, J.K. Barnes, D. W. Bliss, Frank H. Hamilton, Robert Reyburn, and J.J. Woodward between July 6 – September 19, 1881.\n\nBased on all the information, answer the query. \n\nQuery: If my future wife has the same first name as the 15th first lady of the United States' mother and her surname is the same as the second assassinated president's mother's maiden name, what is my future wife's name? \n\n"} ================================================ FILE: benchmark/kernels/quantization/README.md ================================================ # W8A8 Block-wise Quantization Kernel Tuning Auto-tune Triton FP8/INT8 block-wise quantization kernels for optimal performance. ## When to Use Triton FP8 Block-wise Quantization Kernel vs DeepGEMM **Use Triton FP8 Block-wise Quantization Kernel when:** - Output dtype is NOT `bfloat16` (e.g., `float16`, `float32`) - DeepGEMM is disabled (environment variable `SGLANG_ENABLE_JIT_DEEPGEMM=0`) - Running on GPUs with compute capability < SM90 (DeepGEMM requires SM90+) - You need cross-platform compatibility (Triton works on both NVIDIA and AMD GPUs) **Use DeepGEMM when:** - Output dtype is `bfloat16` AND DeepGEMM is enabled - Running on NVIDIA GPUs with compute capability >= SM90 (e.g., H100, H200) - Need maximum performance for production workloads (DeepGEMM is highly optimized for Hopper architecture) **Note:** DeepGEMM requires CUDA compute capability >= 9.0 (SM90+). It is specifically optimized for NVIDIA Hopper GPUs (H100/H200). The kernel selection logic in SGLang automatically chooses DeepGEMM when conditions are met (see `w8a8_block_fp8_matmul` function in `fp8_kernel.py`), otherwise falls back to Triton implementation. ## Quick Start **Default (DeepSeek-V3):** ```bash python benchmark/kernels/quantization/tuning_block_wise_kernel.py --tp-size 8 ``` **Custom Model (specify N and K):** ```bash python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 25600 ``` ## Parameters - `--N`, `--K`: Weight matrix dimensions (N=output_dim, K=input_dim). If not specified, uses `--tp-size` for DeepSeek-V3 - `--tp-size`: Tensor parallelism size for DeepSeek-V3 (default: 8) - `--input-type`: `fp8` or `int8` (default: fp8) - `--block-n`, `--block-k`: Block quantization granularity (default: 128) - `--batch-size`: Test single batch size (optional) ## How to Calculate N and K For a linear layer `y = xW^T` where `x` is (M, K) and `W` is (N, K): - **N**: Output features (weight matrix output dimension) - **K**: Input features (weight matrix input dimension) **Example: Qwen3-VL-32B** (hidden_size=5120, intermediate_size=25600, num_heads=64, num_kv_heads=8, head_dim=128) and TP=1 ```bash # QKV projection: Q(8192) + K(1024) + V(1024) = 10240 python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 10240 --K 5120 # MLP gate+up (SwiGLU): 2 * intermediate_size = 51200 python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 51200 --K 5120 # MLP down projection python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 25600 # O projection (if separate from QKV) python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 8192 ``` If TP=8: ```bash # QKV projection: Q(8192) + K(1024) + V(1024) = 10240 / TP=8 python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 1280 --K 5120 # MLP gate+up (SwiGLU): 2 * intermediate_size = 51200 / TP=8 python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 6400 --K 5120 # MLP down projection python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 3200 # O projection (if separate from QKV) python benchmark/kernels/quantization/tuning_block_wise_kernel.py --N 5120 --K 1024 ``` ## Output Generates JSON config files saved to `python/sglang/srt/layers/quantization/configs/`: ``` N={N},K={K},device_name={DEVICE},dtype=fp8_w8a8,block_shape=[128,128].json ``` Config maps batch size to optimal kernel parameters: ```json { "1": {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, ...}, "2048": {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128, ...} } ``` ================================================ FILE: benchmark/kernels/quantization/bench_fp4_quant.py ================================================ import argparse import itertools import torch import triton from flashinfer import ( scaled_fp4_grouped_quantize, silu_and_mul_scaled_nvfp4_experts_quantize, ) from sgl_kernel.elementwise import silu_and_mul from sglang.benchmark.bench_utils import run_bench from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.moe.ep_moe.kernels import silu_and_mul_masked_post_quant_fwd def _test_accuracy_once(E, M, K, input_dtype, device): x = torch.randn(E, M, K, device=device, dtype=input_dtype) glb_scales = torch.ones((E,), dtype=torch.float32, device=device) masks = torch.full((E,), M, dtype=torch.int32, device=device) out, blk_scales = silu_and_mul_scaled_nvfp4_experts_quantize(x, masks, glb_scales) out1, blk_scales1 = scaled_fp4_grouped_quantize( silu_and_mul(x), masks, glb_scales, ) torch.testing.assert_close(out, out1) torch.testing.assert_close(blk_scales, blk_scales1) print(f"E: {E}, M: {M}, K: {K}, type: {input_dtype} OK") NUM_RANKS = 48 M_PER_RANKs = [128, 256, 512, 1024] Ms = [M_PER_RANK * NUM_RANKS for M_PER_RANK in M_PER_RANKs] Ks = [2048, 4096, 7168] @triton.testing.perf_report( triton.testing.Benchmark( x_names=["M", "K"], x_vals=list(itertools.product(Ms, Ks)), x_log=False, line_arg="provider", line_vals=["triton_fp8", "cuda_unfused_fp4", "cuda_fused_fp4"], line_names=["triton_fp8", "cuda_unfused_fp4", "cuda_fused_fp4"], styles=[("blue", "-"), ("orange", "-"), ("green", "-")], ylabel="ms", plot_name="fp4 quant", args={}, ) ) def benchmark(M, K, provider): E = 6 device = "cuda" x = torch.randn(E, M, K, device=device, dtype=torch.bfloat16) glb_scales = torch.ones((E,), dtype=torch.float32, device=device) masks = torch.randint(1, 4096, (E,), dtype=torch.int32, device=device) fp8_out = torch.empty( ( x.shape[0], x.shape[1], x.shape[2] // 2, ), device=x.device, dtype=torch.float8_e4m3fn, ) scale_block_size = 128 fp8_scales = torch.empty( ( x.shape[0], x.shape[1], x.shape[2] // 2 // scale_block_size, ), device=x.device, dtype=torch.float32, ) quantiles = (0.5, 0.2, 0.8) if provider == "triton_fp8": ms, min_ms, max_ms = run_bench( lambda: silu_and_mul_masked_post_quant_fwd( x, fp8_out, fp8_scales, scale_block_size, masks, scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, ), quantiles=quantiles, ) if provider == "cuda_unfused_fp4": ms, min_ms, max_ms = run_bench( lambda: scaled_fp4_grouped_quantize( silu_and_mul(x), masks, glb_scales, ), quantiles=quantiles, ) if provider == "cuda_fused_fp4": ms, min_ms, max_ms = run_bench( lambda: silu_and_mul_scaled_nvfp4_experts_quantize( x, masks, glb_scales, ), quantiles=quantiles, ) return ms, min_ms, max_ms def test_accuracy(): E = 6 N_RANKS = 48 Ms = [128, 256, 512, 1024] Ks = [2048, 4096, 7168] input_dtype = torch.bfloat16 for M in Ms: for K in Ks: _test_accuracy_once(E, N_RANKS * M, K, input_dtype, "cuda") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--save_path", type=str, default="./bench_fp4_quant_res", help="Path to save fp4 quant benchmark results", ) args = parser.parse_args() test_accuracy() benchmark.run(print_data=True, show_plots=True, save_path=args.save_path) ================================================ FILE: benchmark/kernels/quantization/bench_int8_quant.py ================================================ import argparse import torch import triton from vllm._custom_ops import scaled_int8_quant as vllm_scaled_int8_quant from sglang.benchmark.bench_utils import run_bench from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8 @torch.compile(backend="inductor") def torch_int8_quant(x): int8_max = torch.iinfo(torch.int8).max abs_max = x.abs().max(dim=-1, keepdim=True).values scales = abs_max.to(torch.float32) / float(int8_max) q_x = (x / scales).round().to(torch.int8) return q_x, scales def _test_accuracy_once(M, K, input_dtype, device): x = torch.randn(M, K, dtype=input_dtype, device=device) * 5000 out, scales, _ = vllm_scaled_int8_quant(x, symmetric=True) out1, scales1 = per_token_quant_int8(x) out2, scales2 = torch_int8_quant(x) torch.testing.assert_close(out, out2, atol=1, rtol=0) torch.testing.assert_close(out, out1, atol=1, rtol=0) torch.testing.assert_close(scales, scales2) torch.testing.assert_close(scales1, scales2) print(f"M: {M}, K: {K}, type: {input_dtype} OK") def test_accuracy(): Ms = [1, 13, 128, 1024, 2048, 4096] Ks = [512, 1024, 2048, 8192] input_dtypes = [torch.float16, torch.bfloat16] for M in Ms: for K in Ks: for input_dtype in input_dtypes: _test_accuracy_once(M, K, input_dtype, "cuda") @triton.testing.perf_report( triton.testing.Benchmark( x_names=["batch_size"], x_vals=[1, 16, 32, 64, 128, 256, 512, 1024, 2048], x_log=False, line_arg="provider", line_vals=["vllm op", "triton", "torch.compile"], line_names=["vllm op", "triton", "torch.compile"], styles=[("blue", "-"), ("orange", "-"), ("red", "-")], ylabel="ms", plot_name="int8 per token quant", args={}, ) ) def benchmark(batch_size, provider): M, K = batch_size, 16384 x = torch.randn(M, K, dtype=torch.float16, device="cuda") * 1000 quantiles = (0.5, 0.2, 0.8) if provider == "vllm op": ms, min_ms, max_ms = run_bench( lambda: vllm_scaled_int8_quant(x, symmetric=True), quantiles=quantiles, ) if provider == "triton": ms, min_ms, max_ms = run_bench( lambda: per_token_quant_int8(x), quantiles=quantiles, ) if provider == "torch.compile": ms, min_ms, max_ms = run_bench( lambda: torch_int8_quant(x), quantiles=quantiles, ) return ms, min_ms, max_ms if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--save_path", type=str, default="./bench_int8_quant_res", help="Path to save int8 quant benchmark results", ) args = parser.parse_args() test_accuracy() benchmark.run(print_data=True, show_plots=True, save_path=args.save_path) ================================================ FILE: benchmark/kernels/quantization/tuning_block_wise_kernel.py ================================================ # Copyright 2025 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== import argparse import json import multiprocessing as mp import os import time from datetime import datetime from typing import Any, Dict, List import torch import triton from tqdm import tqdm mp.set_start_method("spawn", force=True) from sglang.srt.layers.quantization.fp8_kernel import ( _w8a8_block_fp8_matmul, _w8a8_block_fp8_matmul_unrolledx4, ) from sglang.srt.layers.quantization.int8_kernel import _w8a8_block_int8_matmul from sglang.srt.utils import get_device_core_count, get_device_name, is_hip _is_hip = is_hip() DTYPE_MAP = { "float32": torch.float32, "float16": torch.float16, "half": torch.half, "bfloat16": torch.bfloat16, } def w8a8_block_matmul( A: torch.Tensor, B: torch.Tensor, As: torch.Tensor, Bs: torch.Tensor, block_size: List[int], config: Dict[str, Any], output_dtype: torch.dtype = torch.float16, ) -> torch.Tensor: """This function performs matrix multiplication with block-wise quantization. It takes two input tensors `A` and `B` with scales `As` and `Bs`. The output is returned in the specified `output_dtype`. Args: A: The input tensor, e.g., activation. B: The input tensor, e.g., weight. As: The per-token-group quantization scale for `A`. Bs: The per-block quantization scale for `B`. block_size: The block size for per-block quantization. It should be 2-dim, e.g., [128, 128]. output_dytpe: The dtype of the returned tensor. Returns: torch.Tensor: The result of matmul. """ assert len(block_size) == 2 block_n, block_k = block_size[0], block_size[1] assert A.shape[-1] == B.shape[-1] assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous() assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1] M = A.numel() // A.shape[-1] assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2 N, K = B.shape assert triton.cdiv(N, block_n) == Bs.shape[0] assert triton.cdiv(K, block_k) == Bs.shape[1] C_shape = A.shape[:-1] + (N,) C = A.new_empty(C_shape, dtype=output_dtype) needs_masking = bool(K % config["BLOCK_SIZE_K"] != 0) def grid(META): return ( triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), ) # Use manually unrolledx4 kernel on AMD GPU when the grid size is small. # Empirical testing shows the sweet spot lies when it's less than the # of # compute units available on the device. num_workgroups = triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv( N, config["BLOCK_SIZE_N"] ) if A.dtype == torch.float8_e4m3fnuz or A.dtype == torch.float8_e4m3fn: kernel = ( _w8a8_block_fp8_matmul_unrolledx4 if (_is_hip == True and num_workgroups <= get_device_core_count()) else _w8a8_block_fp8_matmul ) else: kernel = _w8a8_block_int8_matmul kernel[grid]( A, B, C, As, Bs, M, N, K, block_n, block_k, A.stride(-2), A.stride(-1), B.stride(1), B.stride(0), C.stride(-2), C.stride(-1), As.stride(-2), As.stride(-1), Bs.stride(1), Bs.stride(0), **config, needs_masking=needs_masking, ) return C def get_rocm_configs_compute_bound(): configs = [] waves_per_eu_range = 0 for num_stages in [2]: for block_m in [32, 64, 128, 256]: for block_k in [32, 64, 128, 256]: for block_n in [16, 32, 64, 128, 256]: for num_warps in [4, 8]: for group_size in [1, 4, 8, 16, 32]: configs.append( { "BLOCK_SIZE_M": block_m, "BLOCK_SIZE_N": block_n, "BLOCK_SIZE_K": block_k, "GROUP_SIZE_M": group_size, "num_warps": num_warps, "num_stages": num_stages, "waves_per_eu": waves_per_eu_range, } ) return configs def get_configs_compute_bound(): configs = [] if _is_hip: configs = get_rocm_configs_compute_bound() else: for num_stages in [2, 3, 4, 5]: for block_m in [16, 32, 64, 128, 256]: for block_k in [64, 128]: for block_n in [32, 64, 128, 256]: for num_warps in [4, 8]: for group_size in [1, 16, 32, 64]: configs.append( { "BLOCK_SIZE_M": block_m, "BLOCK_SIZE_N": block_n, "BLOCK_SIZE_K": block_k, "GROUP_SIZE_M": group_size, "num_warps": num_warps, "num_stages": num_stages, } ) return configs def get_weight_shapes(tp_size): # NOTE(HandH1998): The weight shapes only works for DeepSeek-V3. Modify them, if you tune for another different model. # cannot TP total = [ (512 + 64, 7168), ((128 + 64) * 128, 7168), (128 * (128 + 128), 512), (7168, 16384), (7168, 18432), ] # N can TP n_tp = [ (18432 * 2, 7168), ((128 + 64) * 128, 7168), (128 * (128 + 128), 512), (24576, 1536), (4096, 7168), ] # K can TP k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] weight_shapes = [] for t in total: weight_shapes.append(t) for n_t in n_tp: new_t = (n_t[0] // tp_size, n_t[1]) weight_shapes.append(new_t) for k_t in k_tp: new_t = (k_t[0], k_t[1] // tp_size) weight_shapes.append(new_t) return weight_shapes def benchmark_config( A, B, As, Bs, block_size, config, out_dtype=torch.float16, num_iters=10 ): def run(): w8a8_block_matmul(A, B, As, Bs, block_size, config, out_dtype) torch.cuda.synchronize() # JIT complication & warmup for _ in range(5): run() torch.cuda.synchronize() start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) latencies: List[float] = [] for i in range(num_iters): torch.cuda.synchronize() start_event.record() run() end_event.record() end_event.synchronize() latencies.append(start_event.elapsed_time(end_event)) avg = sum(latencies) / (num_iters * 10) * 1000 # us return avg def tune(M, N, K, block_size, out_dtype, search_space, input_type): factor_for_scale = 1e-2 if input_type == "fp8": fp8_info = torch.finfo( torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn ) fp8_max, fp8_min = fp8_info.max, fp8_info.min A_fp32 = ( (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max ) A = A_fp32.clamp(min=fp8_min, max=fp8_max).to( torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn ) B_fp32 = ( (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max ) B = B_fp32.clamp(min=fp8_min, max=fp8_max).to( torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn ) else: int8_info = torch.iinfo(torch.int8) int8_max, int8_min = int8_info.max, int8_info.min A_fp32 = ( (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * int8_max ) A = A_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8) B_fp32 = ( (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * int8_max ) B = B_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8) block_n, block_k = block_size[0], block_size[1] n_tiles = (N + block_n - 1) // block_n k_tiles = (K + block_k - 1) // block_k As = torch.rand(M, k_tiles, dtype=torch.float32, device="cuda") * factor_for_scale Bs = ( torch.rand(n_tiles, k_tiles, dtype=torch.float32, device="cuda") * factor_for_scale ) best_config = None best_time = float("inf") for config in tqdm(search_space): try: kernel_time = benchmark_config( A, B, As, Bs, block_size, config, out_dtype, num_iters=10, ) except triton.runtime.autotuner.OutOfResources: # Some configurations may be invalid and fail to compile. continue if kernel_time < best_time: best_time = kernel_time best_config = config now = datetime.now() print(f"{now.ctime()}] Completed tuning for batch_size={M}") assert best_config is not None return best_config def save_configs( N, K, block_n, block_k, configs, save_path, input_type="fp8", lock=None, ) -> None: os.makedirs(save_path, exist_ok=True) device_name = get_device_name().replace(" ", "_") json_file_name = f"N={N},K={K},device_name={device_name},dtype={input_type}_w8a8,block_shape=[{block_n}, {block_k}].json" config_file_path = os.path.join(save_path, json_file_name) print(f"Writing best config to {config_file_path}...") if lock is not None: lock.acquire() try: existing_configs = {} if os.path.exists(config_file_path): with open(config_file_path, "r") as f: existing_configs = json.load(f) existing_configs = {int(k): v for k, v in existing_configs.items()} existing_configs.update(configs) with open(config_file_path, "w") as f: json.dump(existing_configs, f, indent=4) f.write("\n") finally: if lock is not None: lock.release() def get_available_gpu_count(): """Get the number of available GPUs.""" return torch.cuda.device_count() def tune_on_gpu(args_dict): """Run tuning on a specific GPU.""" gpu_id = args_dict["gpu_id"] batch_sizes = args_dict["batch_sizes"] weight_shapes = args_dict["weight_shapes"] args = args_dict["args"] lock = args_dict["lock"] torch.cuda.set_device(gpu_id) print(f"Starting tuning on GPU {gpu_id} with batch sizes {batch_sizes}") block_n = args.block_n block_k = args.block_k out_dtype = DTYPE_MAP[args.out_dtype] save_path = args.save_path input_type = args.input_type search_space = get_configs_compute_bound() search_space = [ config for config in search_space if block_k % config["BLOCK_SIZE_K"] == 0 ] start = time.perf_counter() results = {} for shape in tqdm(weight_shapes, desc=f"GPU {gpu_id} - Shapes"): N, K = shape[0], shape[1] print(f"[GPU {gpu_id}] Tune for weight shape of `N: {N}, K: {K}`") benchmark_results = [ tune( batch_size, N, K, [block_n, block_k], out_dtype, search_space, input_type, ) for batch_size in tqdm(batch_sizes, desc=f"GPU {gpu_id} - Batch sizes") ] best_configs = {M: config for M, config in zip(batch_sizes, benchmark_results)} save_configs(N, K, block_n, block_k, best_configs, save_path, input_type, lock) end = time.perf_counter() print(f"Tuning on GPU {gpu_id} took {end - start:.2f} seconds") def distribute_batch_sizes(batch_sizes, num_gpus): """Distribute batch sizes across available GPUs.""" batches_per_gpu = [] for i in range(num_gpus): start_idx = i * len(batch_sizes) // num_gpus end_idx = (i + 1) * len(batch_sizes) // num_gpus batches_per_gpu.append(batch_sizes[start_idx:end_idx]) return batches_per_gpu def main(args): print(args) num_gpus = get_available_gpu_count() if num_gpus == 0: raise RuntimeError("No GPU available for tuning") print(f"Found {num_gpus} GPUs for parallel tuning") torch.cuda.init() if args.batch_size is None: batch_sizes = [ 1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 256, 512, 1024, 1536, 2048, 3072, 4096, ] else: batch_sizes = [args.batch_size] num_gpus = 1 # If only one batch size, use only one GPU # Support manual N and K specification if args.N is not None and args.K is not None: weight_shapes = [(args.N, args.K)] print(f"Using manually specified weight shape: N={args.N}, K={args.K}") else: weight_shapes = get_weight_shapes(args.tp_size) print(f"Using predefined weight shapes for TP size {args.tp_size}") batches_per_gpu = distribute_batch_sizes(batch_sizes, num_gpus) ctx = mp.get_context("spawn") manager = ctx.Manager() lock = manager.Lock() process_args = [] for gpu_id in range(num_gpus): process_args.append( { "gpu_id": gpu_id, "batch_sizes": batches_per_gpu[gpu_id], "weight_shapes": weight_shapes, # Each GPU processes all weight shapes "args": args, "lock": lock, } ) with ctx.Pool(num_gpus) as pool: pool.map(tune_on_gpu, process_args) print("Multi-GPU tuning completed") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--tp-size", "-tp", type=int, default=8, help="Tensor parallelism size (ignored if --N and --K are specified)", ) parser.add_argument( "--N", type=int, default=None, help="Output dimension of weight matrix (number of columns)", ) parser.add_argument( "--K", type=int, default=None, help="Input dimension of weight matrix (number of rows)", ) parser.add_argument( "--input-type", type=str, choices=["fp8", "int8"], default="fp8" ) parser.add_argument( "--out-dtype", type=str, choices=["float32", "float16", "bfloat16", "half"], default="float16", ) parser.add_argument("--block-n", type=int, default=128) parser.add_argument("--block-k", type=int, default=128) parser.add_argument("--batch-size", type=int, required=False) parser.add_argument( "--save-path", type=str, default="python/sglang/srt/layers/quantization/configs" ) args = parser.parse_args() # Validate arguments if (args.N is None) != (args.K is None): parser.error("--N and --K must be specified together or not at all") main(args) ================================================ FILE: benchmark/kernels/scheduler_batch/benchmark_get_last_loc_triton.py ================================================ import os import torch import triton import triton.language as tl from sglang.benchmark.bench_utils import run_bench @torch.compile(dynamic=True) def get_last_loc_torch( req_to_token: torch.Tensor, req_pool_indices_tensor: torch.Tensor, prefix_lens_tensor: torch.Tensor, ) -> torch.Tensor: return torch.where( prefix_lens_tensor > 0, req_to_token[req_pool_indices_tensor, prefix_lens_tensor - 1], torch.full_like(prefix_lens_tensor, -1), ) @triton.jit def get_last_loc_kernel( req_to_token, req_pool_indices_tensor, prefix_lens_tensor, result, num_tokens, req_to_token_stride, BLOCK_SIZE: tl.constexpr, ): pid = tl.program_id(0) offset = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE mask = offset < num_tokens prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0) req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0) token_mask = prefix_lens > 0 token_index = req_pool_indices * req_to_token_stride + (prefix_lens - 1) tokens = tl.load(req_to_token + token_index, mask=token_mask, other=-1) tl.store(result + offset, tokens, mask=mask) def get_last_loc_triton( req_to_token: torch.Tensor, req_pool_indices_tensor: torch.Tensor, prefix_lens_tensor: torch.Tensor, ) -> torch.Tensor: BLOCK_SIZE = 256 num_tokens = prefix_lens_tensor.shape[0] result = torch.empty_like(prefix_lens_tensor) grid = (triton.cdiv(num_tokens, BLOCK_SIZE),) get_last_loc_kernel[grid]( req_to_token, req_pool_indices_tensor, prefix_lens_tensor, result, num_tokens, req_to_token.stride(0), BLOCK_SIZE, ) return result def test_get_last_loc(): max_batch = 4097 max_context_len = 6148 batch_size = 20 # Initialize input tensors req_to_token = torch.zeros( (max_batch, max_context_len), dtype=torch.int32, device="cuda" ) req_pool_indices = torch.arange(batch_size, dtype=torch.int64, device="cuda") pre_lens = torch.randint( -max_context_len // 2, max_context_len, (batch_size,), dtype=torch.int64, device="cuda", ) last_loc_res = get_last_loc_triton(req_to_token, req_pool_indices, pre_lens) last_loc_ref = get_last_loc_torch(req_to_token, req_pool_indices, pre_lens) # Compare results torch.testing.assert_close(last_loc_res, last_loc_ref) def get_benchmark(): batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024] @triton.testing.perf_report( triton.testing.Benchmark( x_names=["batch_size"], x_vals=batch_sizes, line_arg="provider", line_vals=["reference", "triton"], line_names=["PyTorch", "Triton"], styles=[("blue", "-"), ("green", "-")], ylabel="us", plot_name="get-last-loc-performance", args={}, ) ) def benchmark(batch_size, provider): max_batch = 2048 max_context_len = 16384 req_to_token = torch.zeros( (max_batch, max_context_len), dtype=torch.int32, device="cuda" ) req_pool_indices = torch.arange(batch_size, dtype=torch.int64, device="cuda") pre_lens = torch.randint( -max_context_len // 2, max_context_len, (batch_size,), dtype=torch.int64, device="cuda", ) quantiles = [0.5, 0.2, 0.8] if provider == "reference": ms, min_ms, max_ms = run_bench( lambda: get_last_loc_torch(req_to_token, req_pool_indices, pre_lens), quantiles=tuple(quantiles), ) elif provider == "triton": ms, min_ms, max_ms = run_bench( lambda: get_last_loc_triton(req_to_token, req_pool_indices, pre_lens), quantiles=tuple(quantiles), ) return 1000 * ms, 1000 * max_ms, 1000 * min_ms return benchmark def run_benchmark(save_path: str = "./configs/benchmark_ops/get_last_loc/"): """Run benchmark and save results""" # Ensure save path exists os.makedirs(save_path, exist_ok=True) # Run correctness test test_get_last_loc() print("Correctness test passed!") # Run performance test benchmark = get_benchmark() benchmark.run(print_data=True, save_path=save_path) if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() parser.add_argument( "--save_path", type=str, default="./configs/benchmark_ops/get_last_loc/", help="Path to save benchmark results", ) args = parser.parse_args() run_benchmark(args.save_path) ================================================ FILE: benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py ================================================ import itertools import os import torch import triton import triton.language as tl from sglang.benchmark.bench_utils import run_bench @triton.jit def write_req_to_token_pool_triton( req_to_token_ptr, # [max_batch, max_context_len] req_pool_indices, pre_lens, seq_lens, extend_lens, out_cache_loc, req_to_token_ptr_stride: tl.constexpr, ): BLOCK_SIZE: tl.constexpr = 512 pid = tl.program_id(0) req_pool_index = tl.load(req_pool_indices + pid) pre_len = tl.load(pre_lens + pid) seq_len = tl.load(seq_lens + pid) # TODO: optimize this? cumsum_start = 0 for i in range(pid): cumsum_start += tl.load(extend_lens + i) num_loop = tl.cdiv(seq_len - pre_len, BLOCK_SIZE) for i in range(num_loop): offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE mask = offset < (seq_len - pre_len) value = tl.load(out_cache_loc + cumsum_start + offset, mask=mask) tl.store( req_to_token_ptr + req_pool_index * req_to_token_ptr_stride + offset + pre_len, value, mask=mask, ) @triton.jit def write_req_to_token_pool_triton_optimize( req_to_token_ptr, # [max_batch, max_context_len] req_pool_indices, pre_lens, seq_lens, extend_lens, out_cache_loc, req_to_token_ptr_stride: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): pid_batch = tl.program_id(0) pid_token = tl.program_id(1) req_pool_index = tl.load(req_pool_indices + pid_batch) pre_len = tl.load(pre_lens + pid_batch) seq_len = tl.load(seq_lens + pid_batch) extend_len = seq_len - pre_len cumsum_start = 0 for i in range(pid_batch): cumsum_start += tl.load(extend_lens + i) token_start = pid_token * BLOCK_SIZE offset = tl.arange(0, BLOCK_SIZE) actual_offset = token_start + offset mask = actual_offset < extend_len src_ptr = out_cache_loc + cumsum_start + actual_offset src_ptr = tl.max_contiguous(tl.multiple_of(src_ptr, BLOCK_SIZE), BLOCK_SIZE) value = tl.load(src_ptr, mask=mask) dst_ptr = ( req_to_token_ptr + req_pool_index * req_to_token_ptr_stride + actual_offset + pre_len ) dst_ptr = tl.max_contiguous(tl.multiple_of(dst_ptr, BLOCK_SIZE), BLOCK_SIZE) tl.store(dst_ptr, value, mask=mask) def write_req_to_token_pool_reference( req_to_token: torch.Tensor, req_pool_indices: torch.Tensor, pre_lens: torch.Tensor, seq_lens: torch.Tensor, extend_lens: torch.Tensor, out_cache_loc: torch.Tensor, ) -> None: """Reference implementation using PyTorch""" for i in range(len(req_pool_indices)): req_pool_idx = req_pool_indices[i].item() pre_len = pre_lens[i].item() seq_len = seq_lens[i].item() extend_len = extend_lens[i].item() cumsum_start = sum(extend_lens[:i].tolist()) # Copy values from out_cache_loc to req_to_token req_to_token[req_pool_idx, pre_len:seq_len] = out_cache_loc[ cumsum_start : cumsum_start + extend_len ] def test_write_req_to_token_pool(): max_batch = 4097 max_context_len = 6148 batch_size = 1 extend_len = 14 # Initialize input tensors req_to_token = torch.zeros( (max_batch, max_context_len), dtype=torch.int32, device="cuda" ) req_pool_indices = torch.tensor([42], dtype=torch.int32, device="cuda") pre_lens = torch.tensor([8], dtype=torch.int32, device="cuda") seq_lens = torch.tensor([22], dtype=torch.int32, device="cuda") extend_lens = torch.tensor([extend_len], dtype=torch.int32, device="cuda") out_cache_loc = torch.arange(extend_len, dtype=torch.int32, device="cuda") # Create copies for reference implementation req_to_token_ref = req_to_token.clone() req_to_token_opt = req_to_token.clone() # Run original triton kernel write_req_to_token_pool_triton[(batch_size,)]( req_to_token, req_pool_indices, pre_lens, seq_lens, extend_lens, out_cache_loc, max_context_len, ) # Run optimized triton kernel def grid(batch_size, extend_len): num_token_blocks = triton.cdiv(extend_len, 512) return (batch_size, num_token_blocks) write_req_to_token_pool_triton_optimize[grid(batch_size, extend_len)]( req_to_token_opt, req_pool_indices, pre_lens, seq_lens, extend_lens, out_cache_loc, max_context_len, BLOCK_SIZE=512, ) # Run reference implementation write_req_to_token_pool_reference( req_to_token_ref, req_pool_indices, pre_lens, seq_lens, extend_lens, out_cache_loc, ) # Compare results torch.testing.assert_close(req_to_token, req_to_token_ref) torch.testing.assert_close(req_to_token_opt, req_to_token_ref) # Test case 2: batch size > 1 batch_size = 3 extend_lens_list = [14, 20, 30] total_extend_len = sum(extend_lens_list) req_to_token = torch.zeros( (max_batch, max_context_len), dtype=torch.int32, device="cuda" ) req_pool_indices = torch.tensor([42, 100, 200], dtype=torch.int32, device="cuda") pre_lens = torch.tensor([8, 10, 15], dtype=torch.int32, device="cuda") seq_lens = torch.tensor([22, 30, 45], dtype=torch.int32, device="cuda") extend_lens = torch.tensor(extend_lens_list, dtype=torch.int32, device="cuda") out_cache_loc = torch.arange(total_extend_len, dtype=torch.int32, device="cuda") req_to_token_ref = req_to_token.clone() req_to_token_opt = req_to_token.clone() # Run original triton kernel write_req_to_token_pool_triton[(batch_size,)]( req_to_token, req_pool_indices, pre_lens, seq_lens, extend_lens, out_cache_loc, max_context_len, ) # Run optimized triton kernel max_extend_len = max(extend_lens_list) write_req_to_token_pool_triton_optimize[grid(batch_size, max_extend_len)]( req_to_token_opt, req_pool_indices, pre_lens, seq_lens, extend_lens, out_cache_loc, max_context_len, BLOCK_SIZE=512, ) # Run reference implementation write_req_to_token_pool_reference( req_to_token_ref, req_pool_indices, pre_lens, seq_lens, extend_lens, out_cache_loc, ) # Compare results torch.testing.assert_close(req_to_token, req_to_token_ref) torch.testing.assert_close(req_to_token_opt, req_to_token_ref) def get_benchmark(): batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128] extend_lens = [32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] configs = list(itertools.product(batch_sizes, extend_lens)) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["batch_size", "extend_len"], x_vals=configs, line_arg="provider", line_vals=["reference", "triton", "triton_optimize"], line_names=["PyTorch", "Triton", "Triton Optimized"], styles=[("blue", "-"), ("green", "-"), ("red", "-")], ylabel="us", plot_name="write-req-to-token-pool-performance", args={}, ) ) def benchmark(batch_size, extend_len, provider): max_batch = 256 max_context_len = 16384 extend_lens_list = [extend_len] * batch_size total_extend_len = sum(extend_lens_list) req_to_token = torch.zeros( (max_batch, max_context_len), dtype=torch.int32, device="cuda" ) req_pool_indices = torch.arange(batch_size, dtype=torch.int32, device="cuda") pre_lens = torch.ones(batch_size, dtype=torch.int32, device="cuda") * 8 seq_lens = pre_lens + extend_len extend_lens = torch.tensor(extend_lens_list, dtype=torch.int32, device="cuda") out_cache_loc = torch.arange(total_extend_len, dtype=torch.int32, device="cuda") quantiles = [0.5, 0.2, 0.8] if provider == "reference": ms, min_ms, max_ms = run_bench( lambda: write_req_to_token_pool_reference( req_to_token.clone(), req_pool_indices, pre_lens, seq_lens, extend_lens, out_cache_loc, ), quantiles=tuple(quantiles), ) elif provider == "triton": ms, min_ms, max_ms = run_bench( lambda: write_req_to_token_pool_triton[(batch_size,)]( req_to_token.clone(), req_pool_indices, pre_lens, seq_lens, extend_lens, out_cache_loc, max_context_len, ), quantiles=tuple(quantiles), ) else: def run_optimized(): block_size = 128 if extend_len <= 1024 else 512 grid_config = (batch_size, triton.cdiv(extend_len, block_size)) write_req_to_token_pool_triton_optimize[grid_config]( req_to_token.clone(), req_pool_indices, pre_lens, seq_lens, extend_lens, out_cache_loc, max_context_len, BLOCK_SIZE=block_size, ) ms, min_ms, max_ms = run_bench(run_optimized, quantiles=tuple(quantiles)) return 1000 * ms, 1000 * max_ms, 1000 * min_ms return benchmark def run_benchmark(save_path: str = "./configs/benchmark_ops/write_req_to_token_pool/"): """Run benchmark and save results""" # Ensure save path exists os.makedirs(save_path, exist_ok=True) # Run correctness test test_write_req_to_token_pool() print("Correctness test passed!") # Run performance test benchmark = get_benchmark() benchmark.run(print_data=True, save_path=save_path) if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() parser.add_argument( "--save_path", type=str, default="./configs/benchmark_ops/write_req_to_token_pool/", help="Path to save benchmark results", ) args = parser.parse_args() run_benchmark(args.save_path) ================================================ FILE: benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py ================================================ import itertools import torch import torch.nn.functional as F import triton.testing as tt from sglang.benchmark.bench_utils import run_bench from sglang.srt.layers.attention.triton_ops.extend_attention import extend_attention_fwd def extend_attention_fwd_torch( q: torch.Tensor, # [extend_tokens, H_Q, D] k: torch.Tensor, # [extend_tokens, H_KV, D] v: torch.Tensor, # [extend_tokens, H_KV, D] o: torch.Tensor, # [extend_tokens, H_Q, D] k_cache: torch.Tensor, # [total_tokens, H_KV, D] v_cache: torch.Tensor, # [total_tokens, H_KV, D] qo_indptr: torch.Tensor, # [B+1] kv_indptr: torch.Tensor, # [B+1] kv_indices: torch.Tensor, # [prefix_tokens] sliding_window_size: int, ): B = qo_indptr.size(0) - 1 _, H_Q, D = q.shape _, H_KV, _ = k.shape group_size = H_Q // H_KV scale = 1.0 / D**0.5 for i in range(B): q_start = int(qo_indptr[i].item()) q_end = int(qo_indptr[i + 1].item()) kv_start = int(kv_indptr[i].item()) kv_end = int(kv_indptr[i + 1].item()) prefix_indices = kv_indices[kv_start:kv_end] k_prefix = k_cache[prefix_indices] # [prefix_len, H_KV, D] v_prefix = v_cache[prefix_indices] # [prefix_len, H_KV, D] k_extend = k[q_start:q_end] # [extend_len, H_KV, D] v_extend = v[q_start:q_end] # [extend_len, H_KV, D] q_extend = q[q_start:q_end] # [extend_len, H_Q, D] k_full = torch.cat([k_prefix, k_extend], dim=0) # [total_len, H_KV, D] v_full = torch.cat([v_prefix, v_extend], dim=0) # [total_len, H_KV, D] if group_size != 1: k_full_hq = k_full.repeat_interleave( group_size, dim=1 ) # [total_len, H_Q, D] v_full_hq = v_full.repeat_interleave( group_size, dim=1 ) # [total_len, H_Q, D] else: k_full_hq = k_full v_full_hq = v_full prefix_len = k_prefix.size(0) extend_len = k_extend.size(0) total_len = prefix_len + extend_len # causal pos_keys = torch.arange(total_len, device=q.device) t = prefix_len + torch.arange(extend_len, device=q.device) # [extend_len] causal_mask = pos_keys.unsqueeze(0) <= t.unsqueeze(1) # sliding window if sliding_window_size is not None and sliding_window_size > 0: start = (t - (sliding_window_size)).clamp_min(0) # [extend_len] else: start = torch.zeros_like(t) window_mask = pos_keys.unsqueeze(0) >= start.unsqueeze(1) final_mask = causal_mask & window_mask attn_scores = ( torch.einsum("qhd,khd->qhk", q_extend, k_full_hq) * scale ) # [extend_len, H_Q, total_len] attn_scores = attn_scores.masked_fill(~final_mask.unsqueeze(1), float("-inf")) attn_weights = F.softmax(attn_scores, dim=-1) o[q_start:q_end] = torch.einsum("qhk,khd->qhd", attn_weights, v_full_hq) def _build_batch( B, N_CTX, H_Q, H_KV, D, WINDOW_SIZE, dtype=torch.bfloat16, device="cuda" ): b_seq_len_prefix = torch.randint( 1, max(2, N_CTX // 2), (B,), dtype=torch.int32, device=device ) b_seq_len_extend = torch.randint( 1, max(2, N_CTX // 2), (B,), dtype=torch.int32, device=device ) b_seq_len = b_seq_len_prefix + b_seq_len_extend b_start_loc = torch.zeros((B,), dtype=torch.int32, device=device) b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0) b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device=device) b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0) kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0) kv_indices = torch.zeros( (int(b_seq_len_prefix.sum().item()),), dtype=torch.int32, device=device ) for i in range(B): s = kv_indptr[i].item() e = kv_indptr[i + 1].item() kv_indices[s:e] = torch.arange( b_start_loc[i], b_start_loc[i] + b_seq_len_prefix[i], dtype=torch.int32, device=device, ) total_token_num = int(torch.sum(b_seq_len).item()) extend_token_num = int(torch.sum(b_seq_len_extend).item()) k_buffer = torch.empty( (total_token_num, H_KV, D), dtype=dtype, device=device ).normal_(mean=0.1, std=0.2) v_buffer = torch.empty( (total_token_num, H_KV, D), dtype=dtype, device=device ).normal_(mean=0.1, std=0.2) k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device) v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device) q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device) for i in range(B): extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i] extend_end_in_buffer = b_start_loc[i] + b_seq_len[i] extend_start = b_start_loc_extend[i] extend_end = b_start_loc_extend[i] + b_seq_len_extend[i] k_extend[extend_start:extend_end] = k_buffer[ extend_start_in_buffer:extend_end_in_buffer ] v_extend[extend_start:extend_end] = v_buffer[ extend_start_in_buffer:extend_end_in_buffer ] q_extend[extend_start:extend_end] = torch.empty( (int(b_seq_len_extend[i].item()), H_Q, D), dtype=dtype, device=device ).normal_(mean=0.1, std=0.2) o_extend_triton = torch.empty( (extend_token_num, H_Q, D), dtype=dtype, device=device ) o_extend_torch = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device) b_seq_len_extend = b_seq_len - b_seq_len_prefix max_len_extend = int(torch.max(b_seq_len_extend, 0)[0].item()) qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0) inputs = dict( q_extend=q_extend, k_extend=k_extend, v_extend=v_extend, k_buffer=k_buffer, v_buffer=v_buffer, o_extend_triton=o_extend_triton, o_extend_torch=o_extend_torch, qo_indptr=qo_indptr, kv_indptr=kv_indptr, kv_indices=kv_indices, max_len_extend=max_len_extend, WINDOW_SIZE=WINDOW_SIZE, ) meta = dict( B=B, N_CTX=N_CTX, H_Q=H_Q, H_KV=H_KV, D=D, extend_token_num=extend_token_num ) return inputs, meta def _run_triton(inputs): extend_attention_fwd( inputs["q_extend"], inputs["k_extend"], inputs["v_extend"], inputs["o_extend_triton"], inputs["k_buffer"], inputs["v_buffer"], inputs["qo_indptr"], inputs["kv_indptr"], inputs["kv_indices"], custom_mask=None, is_causal=True, mask_indptr=None, max_len_extend=inputs["max_len_extend"], sliding_window_size=inputs["WINDOW_SIZE"], ) def _run_torch_ref(inputs): extend_attention_fwd_torch( inputs["q_extend"], inputs["k_extend"], inputs["v_extend"], inputs["o_extend_torch"], inputs["k_buffer"], inputs["v_buffer"], inputs["qo_indptr"], inputs["kv_indptr"], inputs["kv_indices"], inputs["WINDOW_SIZE"], ) N_CTXS = [1024, 2048, 4096, 8192] WINDOW_SIZES = [-1, 127, 256, 512] CONFIGS = list(itertools.product(N_CTXS, WINDOW_SIZES)) PROVIDERS = ["torch", "triton"] @tt.perf_report( tt.Benchmark( x_names=["N_CTX", "WINDOW_SIZE"], x_vals=CONFIGS, line_arg="provider", line_vals=PROVIDERS, line_names=PROVIDERS, ylabel="Runtime (ms)", plot_name="extend_attention_triton_vs_torch", args={ "B": 32, "H_Q": 64, "H_KV": 8, "D": 128, "dtype": "bf16", "device": "cuda", "check_correctness": False, "warmup": 25, "rep": 100, }, ) ) def bench( N_CTX, provider, B, H_Q, H_KV, D, dtype, device, WINDOW_SIZE, check_correctness, warmup, rep, ): torch.manual_seed(0) torch.cuda.manual_seed(0) dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} dt = dtype_map[dtype] inputs, _ = _build_batch( B, N_CTX, H_Q, H_KV, D, WINDOW_SIZE, dtype=dt, device=device ) if check_correctness and provider == "triton": _run_triton(inputs) _run_torch_ref(inputs) torch.cuda.synchronize() if not torch.allclose( inputs["o_extend_triton"], inputs["o_extend_torch"], rtol=1e-3, atol=1e-3 ): raise AssertionError("Mismatch between triton and torch reference.") if provider == "triton": ms = run_bench( lambda: _run_triton(inputs), quantiles=None, warmup_ms=warmup, rep_ms=rep, )[0] elif provider == "torch": ms = run_bench( lambda: _run_torch_ref(inputs), quantiles=None, warmup_ms=warmup, rep_ms=rep, )[0] else: raise ValueError(provider) return ms if __name__ == "__main__": bench.run(print_data=True, show_plots=False) ================================================ FILE: benchmark/line_retrieval/README.md ================================================ ## Download data ``` wget https://raw.githubusercontent.com/merrymercy/merrymercy.github.io/master/files/random_words.json python3 gen_data.py --number 1000 ``` ## Run benchmark ### Benchmark sglang ``` python3 -m sglang.launch_server --model-path codellama/CodeLlama-7b-hf --port 30000 ``` ``` python3 bench_sglang.py --src-index 600 --num-q 50 --parallel 1 ``` ### ``` # original Accuracy: 0.940, latency: 332.83 s # parallel encoding (no_adjust, offset = 1000) Accuracy: 0.760, latency: 238.46 s # parallel encoding (no_adjust, offset = 3000) Accuracy: 0.760, latency: 238.46 s # parallel encoding (no_adjust, offset = 0) Accuracy: 0.520, latency: 238.46 s # parallel encoding (adjust_cache) Accuracy: 0.460, latency: 257.66 s ``` ================================================ FILE: benchmark/line_retrieval/bench_sglang.py ================================================ import argparse import json import re import time import numpy as np import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text @sgl.function def line_retrieval(s, prefix, suffix, body_0, body_1, body_2, body_3): s += prefix + "\n" contexts = [body_0, body_1, body_2, body_3] position_ids_offset = [i * 1000 for i in range(len(contexts))] forks = s.fork(len(contexts), position_ids_offset) forks += lambda i: contexts[i] + "\n" forks.join(mode="concate_and_append") s += "\n" + suffix s += sgl.gen("answer", max_tokens=16) def eval_model(args, line_obj, num_hoops, src_indices, dst_percents): arguments = [] labels = [] sum_src_indices = [] sum_dst_indices = [] for i in range(len(src_indices)): for j in range(len(dst_percents)): src_index = src_indices[i] dst_percent = dst_percents[j] query_indices = line_obj["group_by_num_hoops"][str(num_hoops)] query_indices = [ q for q in query_indices if all(l <= src_index for l in line_obj["links"][q]) and q < src_index ] dst_index = query_indices[ min(int(len(query_indices) * dst_percent), len(query_indices) - 1) ] label = line_obj["values"][dst_index] body = line_obj["lines"][: src_index + 1] suffix = line_obj["suffix"].replace("???", line_obj["indices"][dst_index]) body_part_len = len(body) // 4 arguments.append( { "prefix": line_obj["prefix"], "body_0": "\n".join(body[:body_part_len]), "body_1": "\n".join(body[body_part_len : 2 * body_part_len]), "body_2": "\n".join(body[2 * body_part_len : 3 * body_part_len]), "body_3": "\n".join(body[3 * body_part_len :]), "suffix": suffix, } ) labels.append(label) sum_src_indices.append(src_index) sum_dst_indices.append(dst_index) # Select backend backend = select_sglang_backend(args) tic = time.perf_counter() states = line_retrieval.run_batch( arguments, temperature=0, backend=backend, num_threads=args.parallel, progress_bar=True, ) latency = time.perf_counter() - tic corrects = [] for i in range(len(arguments)): output = states[i]["answer"] prompt_len = states[i].get_meta_info("answer").get("prompt_length", -1) label = labels[i] # Try all numbers findall = re.findall("\d+", output) if not findall: response_number = output else: for response_number in findall: if response_number == label: break correct = response_number == label corrects.append(correct) # Log results summary = ( f"Line index: {sum_src_indices[i]} -> {sum_dst_indices[i]}, " f"Prompt len: {prompt_len}, " f"Correct: {correct}, " f"Label: {label}, Predicted: {response_number}, " ) print(summary) accuracy = np.mean(corrects) print(f"Accuracy: {accuracy:.3f}, latency: {latency:.2f} s") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "line_retrieval", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": len(arguments), "other": { "num_questions": len(arguments), "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") def main(args): line_obj = json.load(open(args.data_path, "r")) num_hoops = args.num_hoops for src_index in args.src_index: src_indices = [src_index] num_queries = args.num_queries_per_src dst_percents = [i * (1 / (num_queries)) for i in range(num_queries)] eval_model(args, line_obj, num_hoops, src_indices, dst_percents) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="lines_1000_0.0.json") parser.add_argument("--src-index", type=int, nargs="+", default=[100]) parser.add_argument("--num-queries-per-src", type=int, default=10) parser.add_argument("--num-hoops", type=int, default=1) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/line_retrieval/gen_data.py ================================================ """ Generate line data for line retrieval task. Usage: python3 gen_data.py --number 1000 """ import argparse import json from collections import defaultdict import numpy as np from tqdm import tqdm def generate_lines(random_words, num_lines, redirect_ratio): prefix = "Here is a list of lines, each with its corresponding REGISTER_CONTENT value. Please memorize them. Be prepared to provide the REGISTER_CONTENT value for a specific line index when I ask." suffix = "The list has ended. Please give the final REGISTER_CONTENT value for a specific line after resolving the redirections and references. For example, the REGISTER_CONTENT of Line __idx0__ is __val0__. The REGISTER_CONTENT of Line __idx1__ is __val1__. The REGISTER_CONTENT of Line __idx2__ is __val2__. The REGISTER_CONTENT of Line ??? is" # Raw lines visited_indices = set([None]) visited_values = set([None]) lines = [] redirects = [] indices = [] values = [] for i in tqdm(range(num_lines)): line_index = None while line_index in visited_indices: line_index = "-".join(np.random.choice(random_words, size=(2,))) visited_indices.add(line_index) line_value = np.random.randint(low=0, high=999999) line_value = f"{line_value:06}" line = f"Line {line_index}: The REGISTER_CONTENT is {line_value}." lines.append(line) redirects.append(None) indices.append(line_index) values.append(line_value) # Add redirect if redirect_ratio > 0: num_redirect_lines = int(len(lines) * redirect_ratio) redirect_indices = np.random.choice( np.arange(len(lines)), size=(num_redirect_lines,), replace=False ) for i in redirect_indices: target_idx = np.random.choice(min(i * 2 + 100, num_lines)) lines[i] = ( f"Line {indices[i]}: The REGISTER_CONTENT is the same as Line {indices[target_idx]}." ) redirects[i] = target_idx # Build links and find sources links = [[] for _ in range(num_lines)] contains_ring = set() for i in range(num_lines): if redirects[i] is None: continue tmp_link = [] cur = i visited = set() while redirects[cur] is not None: visited.add(cur) tmp_link.append(redirects[cur]) cur = redirects[cur] if cur in visited: contains_ring.add(i) tmp_link = None break values[i] = values[cur] links[i] = tmp_link # Group by num_links group_by_num_hoops = defaultdict(list) for i in range(num_lines): if i in contains_ring: continue group_by_num_hoops[len(links[i]) + 1].append(i) keys = sorted(list(group_by_num_hoops.keys())) for num_links in keys: print(f"#links: {num_links}, #lines: {len(group_by_num_hoops[num_links])}") # Append few-shot examples hoop1_candidates = list(group_by_num_hoops[1]) hoop1_candidate_keys = {c: max([c] + links[c]) for c in hoop1_candidates} hoop1_candidates.sort(key=lambda c: hoop1_candidate_keys[c]) hoop2_candidates = list(group_by_num_hoops[2]) hoop2_candidate_keys = {c: max([c] + links[c]) for c in hoop2_candidates} hoop2_candidates.sort(key=lambda c: hoop2_candidate_keys[c]) i = hoop1_candidates[5] suffix = suffix.replace("__idx0__", indices[i]).replace("__val0__", values[i]) if len(hoop2_candidates): i = hoop2_candidates[0] suffix = suffix.replace("__idx1__", indices[i]).replace("__val1__", values[i]) i = hoop2_candidates[1] suffix = suffix.replace("__idx2__", indices[i]).replace("__val2__", values[i]) else: i = hoop1_candidates[1] suffix = suffix.replace("__idx1__", indices[i]).replace("__val1__", values[i]) i = hoop1_candidates[10] suffix = suffix.replace("__idx2__", indices[i]).replace("__val2__", values[i]) obj = { "prefix": prefix, "suffix": suffix, "lines": lines, "indices": indices, "values": values, "links": links, "group_by_num_hoops": group_by_num_hoops, "contains_ring": sorted(list(contains_ring)), } return obj if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--number", type=int) parser.add_argument("--redirect-ratio", type=float, default=0.0) args = parser.parse_args() num_lines = args.number random_words_filename = "random_words.json" random_words = json.load(open(random_words_filename, "r")) np.random.seed(42) obj = generate_lines(random_words, num_lines, args.redirect_ratio) fout = f"lines_{num_lines}_{args.redirect_ratio:.1f}.json" with open(fout, "w") as fout: json.dump(obj, fout, indent=2) ================================================ FILE: benchmark/llava_bench/README.md ================================================ ## Download benchmark images ``` python3 download_images.py ``` image benchmark source: https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild ### Other Dependency ``` pip3 install "sglang[all]" pip3 install "torch>=2.1.2" "transformers>=4.36" pillow ``` ## Run benchmark ### Benchmark sglang Launch a server ``` python3 -m sglang.launch_server --model-path liuhaotian/llava-v1.6-vicuna-7b --tokenizer-path llava-hf/llava-1.5-7b-hf --port 30000 ``` Run benchmark ``` # Run with local models python3 bench_sglang.py --num-questions 60 # Run with OpenAI models python3 bench_sglang.py --num-questions 60 --backend gpt-4-vision-preview ``` ### Bench LLaVA original code ``` git clone git@github.com:haotian-liu/LLaVA.git cd LLaVA git reset --hard 9a26bd1435b4ac42c282757f2c16d34226575e96 pip3 install -e . cd ~/sglang/benchmark/llava_bench CUDA_VISIBLE_DEVICES=0 bash bench_hf_llava_bench.sh ``` ### Benchmark llama.cpp ``` # Install CMAKE_ARGS="-DLLAMA_CUBLAS=on" pip install llama-cpp-python pip install sse_starlette starlette_context pydantic_settings # Download weights mkdir -p ~/model_weights/llava-v1.5-7b/ wget https://huggingface.co/mys/ggml_llava-v1.5-7b/resolve/main/ggml-model-f16.gguf -O ~/model_weights/llava-v1.5-7b/ggml-model-f16.gguf wget https://huggingface.co/mys/ggml_llava-v1.5-7b/resolve/main/mmproj-model-f16.gguf -O ~/model_weights/llava-v1.5-7b/mmproj-model-f16.gguf ``` ``` python3 -m llama_cpp.server --model ~/model_weights/llava-v1.5-7b/ggml-model-f16.gguf --clip_model_path ~/model_weights/llava-v1.5-7b/mmproj-model-f16.gguf --chat_format llava-1-5 --port 23000 OPENAI_BASE_URL=http://localhost:23000/v1 python3 bench_sglang.py --backend gpt-4-vision-preview --num-q 1 ``` ================================================ FILE: benchmark/llava_bench/bench_hf_llava_bench.sh ================================================ #!/bin/bash python -m llava.eval.model_vqa \ --model-path liuhaotian/llava-v1.5-7b \ --question-file ./questions.jsonl \ --image-folder ./images \ --answers-file ./answers_hf.jsonl \ --temperature 0 \ --conv-mode vicuna_v1 ================================================ FILE: benchmark/llava_bench/bench_hf_mme.sh ================================================ #!/bin/bash python -m llava.eval.model_vqa_loader \ --model-path liuhaotian/llava-v1.5-7b \ --question-file ./mme_pack/llava_mme_bench_replace.jsonl \ --image-folder ./mme_pack/MME_Benchmark_release_version \ --answers-file ./answers_hf_mme.jsonl \ --temperature 0 \ --conv-mode vicuna_v1 ================================================ FILE: benchmark/llava_bench/bench_sglang.py ================================================ import argparse import json import os import time import tqdm import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text, read_jsonl @sgl.function def image_qa(s, image_file, question): s += sgl.user(sgl.image(image_file) + question) s += sgl.assistant(sgl.gen("answer", max_tokens=args.max_tokens)) def main(args): lines = list(read_jsonl(args.question_file))[: args.num_questions] arguments = [ { "image_file": os.path.abspath(args.image_folder + "/" + l["image"]), "question": l["text"], } for l in lines ] # arguments = [ # {"image_file": # Image.open(os.path.abspath(args.image_folder + "/" + l["image"])), # "question": l["text"]} for l in lines # ] states = [None] * len(lines) # Select backend backend = select_sglang_backend(args) sgl.set_default_backend(backend) # Run requests tic = time.perf_counter() if args.parallel == 1: for i in tqdm.tqdm(range(len(lines))): image_file = arguments[i]["image_file"] question = arguments[i]["question"] ret = image_qa.run(image_file=image_file, question=question, temperature=0) states[i] = ret else: states = image_qa.run_batch( arguments, temperature=0, num_threads=args.parallel, progress_bar=True ) latency = time.perf_counter() - tic print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) print(f"Write output to {args.answer_file}") with open(args.answer_file, "w") as fout: for i in range(len(lines)): value = { "question_id": lines[i]["question_id"], "prompt": lines[i]["text"], "text": states[i]["answer"].strip(), "model_id": backend.model_info["model_path"], "answer_id": i, "metadata": {}, } fout.write(json.dumps(value) + "\n") with open(args.result_file, "a") as fout: value = { "task": "llava_bench", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": len(lines), "parallel": args.parallel, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--question-file", type=str, default="questions.jsonl") parser.add_argument("--answer-file", type=str, default="answers.jsonl") parser.add_argument("--image-folder", type=str, default="./images") parser.add_argument("--temperature", type=float, default=0.0) parser.add_argument("--num-questions", type=int, default=None) parser.add_argument("--max-tokens", type=int, default=768) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/llava_bench/bench_sglang_mme.sh ================================================ MME_FOLDER=./mme_pack python3 bench_sglang.py --num-questions 5000 --question-file $MME_FOLDER/llava_mme_bench_replace.jsonl --answer-file answer_mme.jsonl --image-folder $MME_FOLDER/MME_Benchmark_release_version --max-tokens 4 ================================================ FILE: benchmark/llava_bench/download_images.py ================================================ import os # Create the 'images' directory if it doesn't exist if not os.path.exists("images"): os.makedirs("images") # Base URL base_url = "https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild/resolve/main/images/" # Loop through image numbers for i in range(1, 25): # Format the image number with leading zeros image_number = str(i).zfill(3) image_url = base_url + image_number + ".jpg" image_path = "images/" + image_number + ".jpg" # Download the image using wget os.system(f"wget -O {image_path} {image_url}") print("Download complete.") ================================================ FILE: benchmark/llm_judge/README.md ================================================ ## Run benchmark ### Benchmark sglang ``` python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` ``` python3 bench_sglang.py --num-questions 25 --parallel 8 python3 bench_sglang.py --num-questions 16 --parallel 1 ``` ### Benchmark vllm ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` ``` python3 bench_other.py --backend vllm --num-questions 25 ``` ### Benchmark guidance ``` python3 bench_other.py --backend guidance --num-questions 25 --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` ### Benchmark lmql ``` python3 bench_other.py --backend lmql --num-questions 25 --parallel 1 ``` ================================================ FILE: benchmark/llm_judge/bench_other.py ================================================ import argparse import json import time from concurrent.futures import ThreadPoolExecutor from functools import partial from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import dump_state_text, read_jsonl system_prompt = "Please serve as an impartial judge and rigorously evaluate the quality of the following article. Apply the most stringent standards possible, showing no leniency." dimension_prompts = [ "Content: This refers to the essences of the essay. The substance should be well researched, accurate, relevant to the topic and should show a thorough understanding of the subject. The essay should also reflect a clear goal or purpose.", "Organization and Structure: An essay needs to be properly structured with a clear introduction, body, and conclusion. The essay should flow naturally, with one paragraph leading seamlessly into the next.", "Argument and Analysis: The argument made in the essay should be logical, coherent and clearly articulated. Each point made should be backed up by solid evidence and thorough analysis.", "Clarity and Precision: The essay should be written in a clear and concise manner. The points made should be easily understood by the reader. The language used should also be precise and unambiguous.", "Grammar and Punctuation: Proper use of grammar and punctuation is vital in an academic essay. Errors in grammar and punctuation not only distract the reader but can also negatively impact the meaning and interpretation of the content.", "Referencing and Citation: An essay should contain proper citations and references for all sources used. This not only prevents accusations of plagiarism but also gives credit to the authors of the works that have contributed to the essay. The citation should adhere to a specific format as required by the academic institution or specified by the professor.", ] def multi_dimension_judge(article, generate): s = system_prompt s += "\n```\n" + article + "\n```\n\n" judges = [] for i in range(len(dimension_prompts)): comp = generate( s + "USER: Please judge the quality based on the following metric. " + dimension_prompts[i] + " Please provide a single-paragraph judgement. " + "Focus on the provided metric and do not say other things. " 'End your judgement paragraph with the word "END"\nJUDGE:', max_tokens=256, stop="END", ) judges.append(comp) s += "I will judge the quality based on the following metrics.\n" for i in range(len(dimension_prompts)): s += dimension_prompts[i].split(":")[0] + ": " + judges[i].strip() + "\n" s += "In summary, on a scale of 1 to 10, I would give the article a score of" s += generate(s, max_tokens=2, stop=None) return s async def multi_dimension_judge_async(article, generate): s = system_prompt s += "\n```\n" + article + "\n```\n\n" judges = [] for i in range(len(dimension_prompts)): comp = await generate( s + "USER: Please judge the quality based on the following metric. " + dimension_prompts[i] + " Please provide a single-paragraph judgement. " + "Focus on the provided metric and do not say other things. " 'End your judgement paragraph with the word "END"\nJUDGE:', max_tokens=256, stop="END", ) judges.append(comp) s += "I will judge the quality based on the following metrics.\n" for i in range(len(dimension_prompts)): s += dimension_prompts[i].split(":")[0] + ": " + judges[i].strip() + "\n" s += "In summary, on a scale of 1 to 10, I would give the article a score of" s += await generate(s, max_tokens=2, stop=None) return s def main(args): lines = read_jsonl(args.data_path)[: args.num_questions] states = [None] * len(lines) # Select backend call_generate = partial(get_call_generate(args), temperature=0) # Run requests tic = time.perf_counter() if args.backend != "lmql": def get_one_answer(i): states[i] = multi_dimension_judge(lines[i], call_generate) if args.parallel == 1: for i in tqdm(range(len(lines))): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: list( tqdm( executor.map(get_one_answer, list(range(len(lines)))), total=len(lines), ) ) else: import asyncio async def get_one_answer_async(i): states[i] = await multi_dimension_judge_async(lines[i], call_generate) batches = [] for i in range(0, len(lines), args.parallel): batches.append(list(range(i, min(i + args.parallel, len(lines))))) loop = asyncio.get_event_loop() for bt in tqdm(batches): loop.run_until_complete( asyncio.gather(*[get_one_answer_async(i) for i in bt]) ) latency = time.perf_counter() - tic # Compute accuracy print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "llm_judge", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="articles.jsonl") parser.add_argument("--num-questions", type=int, default=20) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/llm_judge/bench_sglang.py ================================================ import argparse import json import time import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text, read_jsonl system_prompt = "Please serve as an impartial judge and rigorously evaluate the quality of the following article. Apply the most stringent standards possible, showing no leniency." dimension_prompts = [ "Content: This refers to the essences of the essay. The substance should be well researched, accurate, relevant to the topic and should show a thorough understanding of the subject. The essay should also reflect a clear goal or purpose.", "Organization and Structure: An essay needs to be properly structured with a clear introduction, body, and conclusion. The essay should flow naturally, with one paragraph leading seamlessly into the next.", "Argument and Analysis: The argument made in the essay should be logical, coherent and clearly articulated. Each point made should be backed up by solid evidence and thorough analysis.", "Clarity and Precision: The essay should be written in a clear and concise manner. The points made should be easily understood by the reader. The language used should also be precise and unambiguous.", "Grammar and Punctuation: Proper use of grammar and punctuation is vital in an academic essay. Errors in grammar and punctuation not only distract the reader but can also negatively impact the meaning and interpretation of the content.", "Referencing and Citation: An essay should contain proper citations and references for all sources used. This not only prevents accusations of plagiarism but also gives credit to the authors of the works that have contributed to the essay. The citation should adhere to a specific format as required by the academic institution or specified by the professor.", ] @sgl.function def multi_dimension_judge(s, article): s += system_prompt s += "\n```\n" + article + "\n```\n\n" forks = s.fork(len(dimension_prompts)) for i in range(len(dimension_prompts)): forks[i] += ( "USER: Please judge the quality based on the following metric. " + dimension_prompts[i] + " Please provide a single-paragraph judgement. " + "Focus on the provided metric and do not say other things. " 'End your judgement paragraph with the word "END"\nJUDGE:' ) forks[i] += sgl.gen("judgement", max_tokens=256, stop="END") forks.join() s += "I will judge the quality based on the following metrics.\n" for i in range(len(dimension_prompts)): s += ( dimension_prompts[i].split(":")[0] + ": " + forks[i]["judgement"].strip() + "\n" ) s += "In summary, on a scale of 1 to 10, I would give the article a score of" s += sgl.gen("score", max_tokens=2) def main(args): lines = read_jsonl(args.data_path)[: args.num_questions] arguments = [{"article": l} for l in lines] # Select backend backend = select_sglang_backend(args) # Run requests tic = time.perf_counter() states = multi_dimension_judge.run_batch( arguments, temperature=0, backend=backend, num_threads=args.parallel, progress_bar=True, ) latency = time.perf_counter() - tic print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "llm_judge", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="articles.jsonl") parser.add_argument("--num-questions", type=int, default=20) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/long_json_decode/README.md ================================================ ## Run benchmark ### Benchmark sglang ``` python3 -m sglang.launch_server --model-path codellama/CodeLlama-7b-instruct-hf --port 30000 ``` ``` python3 bench_sglang.py --num-questions 5 --parallel 1 ``` ### Benchmark vllm ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model codellama/CodeLlama-7b-instruct-hf --disable-log-requests --port 21000 --gpu 0.97 ``` ``` python3 bench_other.py --backend vllm --num-questions 5 ``` ### Benchmark guidance ``` python3 bench_other.py --backend guidance --num-questions 5 --parallel 1 --n-ctx 11000 --model-path path/to/code-llama/gguf ``` ### Build dataset ``` pip install wikipedia python3 build_dataset.py ``` ================================================ FILE: benchmark/long_json_decode/bench_other.py ================================================ import argparse import json import time from concurrent.futures import ThreadPoolExecutor from functools import partial from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import dump_state_text, read_jsonl def json_decode(document, generate): s = "Please extract the information of a city from the following wikipedia page.\n" s += "Page begin.\n" + document + "Page end.\n" s += "Here is the name, country, and symbol of the city in JSON format.\n" s += "{\n" s += ' "name": "' s += generate(s, max_tokens=8, stop='"') + '",\n' s += ' "country": "' s += generate(s, max_tokens=8, stop='"') + '",\n' s += ' "air port code": "' s += generate(s, max_tokens=8, stop='"') + '",\n' s += ' "top 3 landmarks": "' s += generate(s, max_tokens=24, stop='"') + '",\n' s += "}\n" return s def main(args): lines = read_jsonl(args.data_path) arguments = [] for i in range(len(lines[: args.num_questions])): arguments.append( { "document": lines[i]["document"], } ) states = [None] * len(arguments) # Select backend call_generate = partial(get_call_generate(args), temperature=0) # Run requests def get_one_answer(i): states[i] = json_decode(generate=call_generate, **arguments[i]) tic = time.perf_counter() if args.parallel == 1: for i in tqdm(range(len(arguments))): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: list( tqdm( executor.map(get_one_answer, list(range(len(arguments)))), total=len(arguments), ) ) latency = time.perf_counter() - tic # Compute accuracy print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "long_json_decode", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="questions.jsonl") parser.add_argument("--num-questions", type=int, default=100) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/long_json_decode/bench_sglang.py ================================================ import argparse import json import time import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text, read_jsonl @sgl.function def json_decode(s, document): s += "Please extract the information of a city from the following wikipedia page.\n" s += "Page begin.\n" + document + "Page end.\n" s += "Here is the name, country, and symbol of the city in JSON format.\n" s += "{\n" s += ' "name": "' + sgl.gen("name", max_tokens=8, stop='"') + '",\n' s += ' "country": "' + sgl.gen("country", max_tokens=8, stop='"') + '",\n' s += ( ' "air port code": "' + sgl.gen("air port code", max_tokens=8, stop='"') + '",\n' ) s += ( ' "top 3 landmarks": "' + sgl.gen("landmarks", max_tokens=24, stop='"') + '",\n' ) s += "}\n" def main(args): lines = read_jsonl(args.data_path) arguments = [] for i in range(len(lines[: args.num_questions])): arguments.append( { "document": lines[i]["document"], } ) # Select backend backend = select_sglang_backend(args) sgl.set_default_backend(backend) # Run requests tic = time.perf_counter() states = json_decode.run_batch( arguments, temperature=0, num_threads=args.parallel, progress_bar=True ) latency = time.perf_counter() - tic # Compute accuracy print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "long_json_decode", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="questions.jsonl") parser.add_argument("--num-questions", type=int, default=10) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/long_json_decode/build_dataset.py ================================================ import json import transformers import wikipedia name = "meta-llama/Llama-2-7b-chat-hf" t = transformers.AutoTokenizer.from_pretrained(name) city_names = ["los angles", "london", "tokyo", "beijing", "singapore"] for city_name in city_names: content = str(wikipedia.page(city_name).content) content = content.replace("\n\n", "\n") tokens = t.encode(content) truncate_len = int((10000 / len(tokens)) * len(content)) truncate_content = content[:truncate_len] truncate_tokens = t.encode(truncate_content) # Count token print( f"city_name: {city_name}, #tokens: {len(tokens)}, #truncate tokens: {len(truncate_tokens)}" ) with open("questions.jsonl", "a") as fout: fout.write(json.dumps({"document": truncate_content}) + "\n") ================================================ FILE: benchmark/lora/launch_server.py ================================================ import argparse import os NUM_LORAS = 4 LORA_PATH = { "base": "meta-llama/Llama-2-7b-hf", "lora": "winddude/wizardLM-LlaMA-LoRA-7B", } def launch_server(args): base_path = LORA_PATH["base"] lora_path = LORA_PATH["lora"] if args.base_only: cmd = f"python3 -m sglang.launch_server --model {base_path} " else: cmd = f"python3 -m sglang.launch_server --model {base_path} --lora-paths " for i in range(NUM_LORAS): lora_name = f"lora{i}" cmd += f"{lora_name}={lora_path} " cmd += f"--disable-radix " cmd += f"--max-loras-per-batch {args.max_loras_per_batch} " cmd += f"--max-running-requests {args.max_running_requests} " cmd += f"--lora-backend {args.lora_backend} " cmd += f"--tp-size {args.tp_size} " if args.disable_custom_all_reduce: cmd += "--disable-custom-all-reduce" if args.enable_mscclpp: cmd += "--enable-mscclpp" if args.enable_torch_symm_mem: cmd += "--enable-torch-symm-mem" print(cmd) os.system(cmd) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--base-only", action="store_true", ) parser.add_argument( "--max-loras-per-batch", type=int, default=8, ) parser.add_argument( "--max-running-requests", type=int, default=8, ) parser.add_argument( "--lora-backend", type=str, default="csgmv", ) parser.add_argument( "--tp-size", type=int, default=1, help="Tensor parallel size for distributed inference", ) # disable_custom_all_reduce parser.add_argument( "--disable-custom-all-reduce", action="store_true", help="Disable custom all reduce when device does not support p2p communication", ) parser.add_argument( "--enable-mscclpp", action="store_true", help="Enable using mscclpp for small messages for all-reduce kernel and fall back to NCCL.", ) parser.add_argument( "--enable-torch-symm-mem", action="store_true", help="Enable using torch symm mem for all-reduce kernel and fall back to NCCL.", ) args = parser.parse_args() launch_server(args) ================================================ FILE: benchmark/lora/lora_bench.py ================================================ # Copyright 2023-2024 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== import argparse import asyncio import json import random import resource import sys import time import traceback from argparse import ArgumentParser from datetime import datetime from typing import Any, Dict, List, Optional, Tuple import numpy as np from launch_server import LORA_PATH, NUM_LORAS from tqdm.asyncio import tqdm from transformers import PreTrainedTokenizerBase from sglang.bench_serving import ( RequestFuncInput, RequestFuncOutput, _create_bench_client_session, calculate_metrics, get_request, ) from sglang.benchmark.datasets.random import sample_random_requests from sglang.benchmark.utils import get_tokenizer, remove_prefix global args # set ignore_eos True by default async def async_request_openai_completions( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url # assert api_url.endswith( # "completions" # ), "OpenAI Completions API URL must end with 'completions'." prompt = request_func_input.prompt async with _create_bench_client_session() as session: # payload = { # "model": request_func_input.model, # "prompt": prompt, # "temperature": 0.0, # "best_of": 1, # "max_tokens": request_func_input.output_len, # "stream": not args.disable_stream, # "ignore_eos": not args.disable_ignore_eos, # **request_func_input.extra_request_body, # } # headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"} if args.base_only: payload = { "text": prompt, "sampling_params": {"max_new_tokens": request_func_input.output_len}, } else: payload = { "text": prompt, "sampling_params": {"max_new_tokens": request_func_input.output_len}, "lora_path": f"lora{random.randint(0, NUM_LORAS - 1)}", } headers = {"Authorization": ""} output = RequestFuncOutput() output.prompt_len = request_func_input.prompt_len generated_text = "" ttft = 0.0 st = time.perf_counter() most_recent_timestamp = st try: async with session.post( url=api_url, json=payload, headers=headers ) as response: if response.status == 200: async for chunk_bytes in response.content: chunk_bytes = chunk_bytes.strip() if not chunk_bytes: continue chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") latency = time.perf_counter() - st if chunk == "[DONE]": pass else: data = json.loads(chunk) # NOTE: Some completion API might have a last # usage summary response without a token so we # want to check a token was generated if data["text"]: # if data["choices"][0]["text"]: timestamp = time.perf_counter() # First token if ttft == 0.0: ttft = time.perf_counter() - st output.ttft = ttft # Decoding phase else: output.itl.append(timestamp - most_recent_timestamp) most_recent_timestamp = timestamp # generated_text += data["choices"][0]["text"] generated_text += data["text"] output.generated_text = generated_text output.success = True output.latency = latency output.output_len = request_func_input.output_len else: output.error = response.reason or "" output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) if pbar: pbar.update(1) return output ASYNC_REQUEST_FUNCS = { "sglang": async_request_openai_completions, } async def benchmark( backend: str, api_url: str, model_id: str, tokenizer: PreTrainedTokenizerBase, input_requests: List[Tuple[str, int, int]], request_rate: float, disable_tqdm: bool, extra_request_body: Dict[str, Any], ): if backend in ASYNC_REQUEST_FUNCS: request_func = ASYNC_REQUEST_FUNCS[backend] else: raise ValueError(f"Unknown backend: {backend}") print("Starting initial single prompt test run...") test_request = input_requests[0] test_input = RequestFuncInput( model=model_id, prompt=test_request.prompt, api_url=api_url, prompt_len=test_request.prompt_len, output_len=test_request.output_len, lora_name="dummy", # the lora_name argument will not be used image_data=None, extra_request_body=extra_request_body, ) test_output = await request_func(request_func_input=test_input) if not test_output.success: raise ValueError( "Initial test run failed - Please make sure benchmark arguments " f"are correctly specified. Error: {test_output.error}" ) else: print("Initial test run completed. Starting main benchmark run...") pbar = None if disable_tqdm else tqdm(total=len(input_requests)) benchmark_start_time = time.perf_counter() tasks: List[asyncio.Task] = [] async for request in get_request(input_requests, request_rate): request_func_input = RequestFuncInput( model=model_id, prompt=request.prompt, api_url=api_url, prompt_len=request.prompt_len, output_len=request.output_len, lora_name="dummy", image_data=None, extra_request_body=extra_request_body, ) tasks.append( asyncio.create_task( request_func(request_func_input=request_func_input, pbar=pbar) ) ) outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) if pbar is not None: pbar.close() benchmark_duration = time.perf_counter() - benchmark_start_time metrics, output_lens = calculate_metrics( input_requests=input_requests, outputs=outputs, dur_s=benchmark_duration, tokenizer=tokenizer, backend=backend, ) print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) print("{:<40} {:<10}".format("Backend:", backend)) print("{:<40} {:<10}".format("Traffic request rate:", request_rate)) print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration)) print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output)) print( "{:<40} {:<10}".format( "Total generated tokens (retokenized):", metrics.total_output_retokenized ) ) print( "{:<40} {:<10.2f}".format( "Request throughput (req/s):", metrics.request_throughput ) ) print( "{:<40} {:<10.2f}".format( "Input token throughput (tok/s):", metrics.input_throughput ) ) print( "{:<40} {:<10.2f}".format( "Output token throughput (tok/s):", metrics.output_throughput ) ) print( "{:<40} {:<10.2f}".format("Total throughput (tok/s):", metrics.total_throughput) ) print("{s:{c}^{n}}".format(s="End-to-End Latency", n=50, c="-")) print( "{:<40} {:<10.2f}".format("Mean E2E Latency (ms):", metrics.mean_e2e_latency_ms) ) print( "{:<40} {:<10.2f}".format( "Median E2E Latency (ms):", metrics.median_e2e_latency_ms ) ) print("{s:{c}^{n}}".format(s="Time to First Token", n=50, c="-")) print("{:<40} {:<10.2f}".format("Mean TTFT (ms):", metrics.mean_ttft_ms)) print("{:<40} {:<10.2f}".format("Median TTFT (ms):", metrics.median_ttft_ms)) print("{:<40} {:<10.2f}".format("P99 TTFT (ms):", metrics.p99_ttft_ms)) print( "{s:{c}^{n}}".format(s="Time per Output Token (excl. 1st token)", n=50, c="-") ) print("{:<40} {:<10.2f}".format("Mean TPOT (ms):", metrics.mean_tpot_ms)) print("{:<40} {:<10.2f}".format("Median TPOT (ms):", metrics.median_tpot_ms)) print("{:<40} {:<10.2f}".format("P99 TPOT (ms):", metrics.p99_tpot_ms)) print("{s:{c}^{n}}".format(s="Inter-token Latency", n=50, c="-")) print("{:<40} {:<10.2f}".format("Mean ITL (ms):", metrics.mean_itl_ms)) print("{:<40} {:<10.2f}".format("Median ITL (ms):", metrics.median_itl_ms)) print("{:<40} {:<10.2f}".format("P99 ITL (ms):", metrics.p99_itl_ms)) print("=" * 50) if ( metrics.median_ttft_ms is not None and metrics.mean_itl_ms is not None and metrics.output_throughput is not None ): result = { "backend": args.backend, "request_rate": request_rate, "total_input_tokens": metrics.total_input, "total_output_tokens": metrics.total_output, "total_output_tokens_retokenized": metrics.total_output_retokenized, "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, "median_e2e_latency_ms": metrics.median_e2e_latency_ms, "median_ttft_ms": metrics.median_ttft_ms, "median_itl_ms": metrics.median_itl_ms, "output_throughput": metrics.output_throughput, "random_input_len": args.random_input_len, "random_output_len": args.random_output_len, "random_range_ratio": args.random_range_ratio, "duration": benchmark_duration, "completed": metrics.completed, } else: print(f"Error running benchmark for request rate: {request_rate}") print("-" * 30) # Determine output file name if args.output_file: output_file_name = args.output_file else: now = datetime.now().strftime("%m%d") output_file_name = f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_{args.random_output_len}.jsonl" # Append results to a JSONL file with open(output_file_name, "a") as file: file.write(json.dumps(result) + "\n") result = { "duration": benchmark_duration, "completed": metrics.completed, "total_input_tokens": metrics.total_input, "total_output_tokens": metrics.total_output, "total_output_tokens_retokenized": metrics.total_output_retokenized, "request_throughput": metrics.request_throughput, "input_throughput": metrics.input_throughput, "output_throughput": metrics.output_throughput, "mean_ttft_ms": metrics.mean_ttft_ms, "median_ttft_ms": metrics.median_ttft_ms, "std_ttft_ms": metrics.std_ttft_ms, "p99_ttft_ms": metrics.p99_ttft_ms, "mean_tpot_ms": metrics.mean_tpot_ms, "median_tpot_ms": metrics.median_tpot_ms, "std_tpot_ms": metrics.std_tpot_ms, "p99_tpot_ms": metrics.p99_tpot_ms, "mean_itl_ms": metrics.mean_itl_ms, "median_itl_ms": metrics.median_itl_ms, "std_itl_ms": metrics.std_itl_ms, "p99_itl_ms": metrics.p99_itl_ms, "input_lens": [output.prompt_len for output in outputs], "output_lens": output_lens, "ttfts": [output.ttft for output in outputs], "itls": [output.itl for output in outputs], "generated_texts": [output.generated_text for output in outputs], "errors": [output.error for output in outputs], "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, "median_e2e_latency_ms": metrics.median_e2e_latency_ms, } return result def run_benchmark(args_: argparse.Namespace): global args args = args_ # Set global environments set_ulimit() random.seed(args.seed) np.random.seed(args.seed) # Set url if args.port is None: args.port = { "sglang": 30000, }.get(args.backend, 30000) # api_url = ( # f"{args.base_url}/v1/completions" # if args.base_url # else f"http://{args.host}:{args.port}/v1/completions" # ) api_url = ( f"{args.base_url}/generate" if args.base_url else f"http://{args.host}:{args.port}/generate" ) print(f"{args}\n") # Read dataset backend = args.backend model_id = args.model = LORA_PATH["base"] tokenizer_id = args.model tokenizer = get_tokenizer(tokenizer_id) input_requests = sample_random_requests( input_len=args.random_input_len, output_len=args.random_output_len, num_prompts=args.num_prompts, range_ratio=args.random_range_ratio, tokenizer=tokenizer, dataset_path="", ) return asyncio.run( benchmark( backend=backend, api_url=api_url, model_id=model_id, tokenizer=tokenizer, input_requests=input_requests, request_rate=args.request_rate, disable_tqdm=False, extra_request_body={}, ) ) def set_ulimit(target_soft_limit=65535): resource_type = resource.RLIMIT_NOFILE current_soft, current_hard = resource.getrlimit(resource_type) if current_soft < target_soft_limit: try: resource.setrlimit(resource_type, (target_soft_limit, current_hard)) except ValueError as e: print(f"Fail to set RLIMIT_NOFILE: {e}") if __name__ == "__main__": parser = ArgumentParser(description="Benchmark the online lora serving throughput.") parser.add_argument( "--backend", type=str, choices=list(ASYNC_REQUEST_FUNCS.keys()), default="sglang", help="Must specify a backend, depending on the LLM Inference Engine.", ) parser.add_argument( "--base-url", type=str, default=None, help="Server or API base url if not using http host and port.", ) parser.add_argument( "--host", type=str, default="0.0.0.0", help="Default host is 0.0.0.0." ) parser.add_argument( "--port", type=int, help="If not set, the default port is configured according to its default value for different LLM Inference Engines.", ) parser.add_argument( "--num-prompts", type=int, default=50, help="Number of prompts to process. Default is 1000.", ) parser.add_argument( "--random-input-len", type=int, default=1024, help="Number of input tokens per request, used only for random dataset.", ) parser.add_argument( "--random-output-len", type=int, default=128, help="Number of output tokens per request, used only for random dataset.", ) parser.add_argument( "--random-range-ratio", type=float, default=0.0, help="Range of sampled ratio of input/output length, " "used only for random dataset.", ) parser.add_argument( "--request-rate", type=float, default=float("inf"), help="Number of requests per second. If this is inf, then all the requests are sent at time 0. " "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.", ) parser.add_argument( "--base-only", action="store_true", ) parser.add_argument("--output-file", type=str, help="Output JSONL file name.") parser.add_argument("--seed", type=int, default=1, help="The random seed.") args = parser.parse_args() run_benchmark(args) ================================================ FILE: benchmark/mmlu/README.md ================================================ ## Download data ``` bash download_data.sh ``` ## Run benchmark ### Benchmark sglang ``` python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` ``` python3 bench_sglang.py --nsub 10 ``` ``` # OpenAI models python3 bench_sglang.py --backend gpt-3.5-turbo --parallel 8 ``` ### Benchmark vllm ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` ``` python3 bench_other.py --nsub 10 --backend vllm ``` ### Benchmark lightllm ``` # A10G python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 # V100 python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 4500 --port 22000 ``` ``` python3 bench_other.py --nsub 10 --backend lightllm ``` ### Benchmark guidance ``` python3 bench_other.py --nsub 10 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` ### Benchmark lmql ``` CUDA_VISIBLE_DEVICES=0,1 lmql serve-model meta-llama/Llama-2-7b-chat-hf --cuda --port 23000 ``` ``` python3 bench_other.py --nsub 10 --backend lmql --parallel 2 ``` ================================================ FILE: benchmark/mmlu/bench_other.py ================================================ import argparse import asyncio import json import os import time from concurrent.futures import ThreadPoolExecutor import numpy as np import pandas as pd import tiktoken from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate choices = ["A", "B", "C", "D"] tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo") def format_subject(subject): l = subject.split("_") s = "" for entry in l: s += " " + entry return s def format_example(df, idx, include_answer=True): prompt = df.iloc[idx, 0] k = df.shape[1] - 2 for j in range(k): prompt += "\n{}. {}".format(choices[j], df.iloc[idx, j + 1]) prompt += "\nAnswer:" if include_answer: prompt += " {}\n\n".format(df.iloc[idx, k + 1]) return prompt def gen_prompt(train_df, subject, k=-1): prompt = "The following are multiple choice questions (with answers) about{}.\n\n".format( format_subject(subject) ) if k == -1: k = train_df.shape[0] for i in range(k): prompt += format_example(train_df, i) return prompt def evaluate(args, subject, dev_df, test_df, call_generate): prompts = [] labels = [] # Construct prompts k = args.ntrain train_prompt = gen_prompt(dev_df, subject, k) while len(tokenizer.encode(train_prompt)) > 1536: k -= 1 train_prompt = gen_prompt(dev_df, subject, k) for i in range(test_df.shape[0]): prompt_end = format_example(test_df, i, include_answer=False) prompt = train_prompt + prompt_end prompts.append(prompt) label = test_df.iloc[i, test_df.shape[1] - 1] labels.append(label) preds = [None] * len(prompts) max_tokens = 1 # Run requests if args.backend != "lmql": # Use thread pool def get_one_answer(i): pred = call_generate(prompts[i], temperature=0, max_tokens=max_tokens) preds[i] = pred.strip()[0] tic = time.perf_counter() if args.parallel == 1: for i in range(len(prompts)): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: executor.map(get_one_answer, list(range(len(prompts)))) else: # Use asyncio async def batched_call(batch_size): for i in range(0, len(prompts), batch_size): tasks = [] for p in prompts[i : i + batch_size]: tasks.append(call_generate(p, temperature=0, max_tokens=max_tokens)) rets = await asyncio.gather(*tasks) for j in range(len(rets)): preds[i + j] = rets[j].strip()[0] tic = time.perf_counter() asyncio.run(batched_call(batch_size=args.parallel)) latency = time.perf_counter() - tic # Compute accuracy cors = [pred == label for pred, label in zip(preds, labels)] acc = np.mean(cors) cors = np.array(cors) print( "Average accuracy {:.3f}, latency {:.2f}, #q: {} - {}".format( acc, latency, len(prompts), subject ) ) return cors, acc, latency def main(args): subjects = sorted( [ f.split("_test.csv")[0] for f in os.listdir(os.path.join(args.data_dir, "test")) if "_test.csv" in f ] ) all_cors = [] all_latencies = [] num_requests = 0 # Select backend call_generate = get_call_generate(args) for subject in tqdm(subjects[: args.nsub]): dev_df = pd.read_csv( os.path.join(args.data_dir, "dev", subject + "_dev.csv"), header=None )[: args.ntrain] test_df = pd.read_csv( os.path.join(args.data_dir, "test", subject + "_test.csv"), header=None ) cors, acc, latency = evaluate(args, subject, dev_df, test_df, call_generate) all_cors.append(cors) all_latencies.append(latency) num_requests += len(test_df) total_latency = np.sum(all_latencies) print("Total latency: {:.3f}".format(total_latency)) weighted_acc = np.mean(np.concatenate(all_cors)) print("Average accuracy: {:.3f}".format(weighted_acc)) # Write results with open(args.result_file, "a") as fout: value = { "task": "mmlu", "backend": args.backend, "num_gpus": 1, "latency": round(total_latency, 3), "accuracy": round(weighted_acc, 3), "num_requests": num_requests, "other": { "nsub": args.nsub, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--ntrain", type=int, default=5) parser.add_argument("--data_dir", type=str, default="data") parser.add_argument("--nsub", type=int, default=60) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/mmlu/bench_sglang.py ================================================ import argparse import json import os import subprocess import tarfile import time import numpy as np import pandas as pd import tiktoken from sglang.test.test_utils import ( add_common_sglang_args_and_parse, dump_bench_raw_result, select_sglang_backend, ) SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) choices = ["A", "B", "C", "D"] tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo") def format_subject(subject): l = subject.split("_") s = "" for entry in l: s += " " + entry return s def format_example(df, idx, include_answer=True): prompt = df.iloc[idx, 0] k = df.shape[1] - 2 for j in range(k): prompt += "\n{}. {}".format(choices[j], df.iloc[idx, j + 1]) prompt += "\nAnswer:" if include_answer: prompt += " {}\n\n".format(df.iloc[idx, k + 1]) return prompt def gen_prompt(train_df, subject, k=-1): prompt = "The following are multiple choice questions (with answers) about{}.\n\n".format( format_subject(subject) ) if k == -1: k = train_df.shape[0] for i in range(k): prompt += format_example(train_df, i) return prompt def download_data(data_dir): """Download and extract MMLU data if it doesn't exist.""" if os.path.isdir(os.path.join(data_dir, "test")): return print(f"Data not found at {data_dir}. Downloading...") os.makedirs(data_dir, exist_ok=True) tar_path = os.path.join(data_dir, "data.tar") subprocess.check_call( ["wget", "-O", tar_path, "https://people.eecs.berkeley.edu/~hendrycks/data.tar"] ) with tarfile.open(tar_path) as tar: tar.extractall(path=data_dir, filter="data") # The tarball extracts into a "data/" subdirectory; move contents up if needed nested = os.path.join(data_dir, "data") if os.path.isdir(nested): for item in os.listdir(nested): os.rename(os.path.join(nested, item), os.path.join(data_dir, item)) os.rmdir(nested) os.remove(tar_path) print("Download complete.") def main(args): subjects = sorted( [ f.split("_test.csv")[0] for f in os.listdir(os.path.join(args.data_dir, "test")) if "_test.csv" in f ] ) # Build prompts arguments = [] labels = [] num_questions = [] for subject in subjects[: args.nsub]: dev_df = pd.read_csv( os.path.join(args.data_dir, "dev", subject + "_dev.csv"), header=None )[: args.ntrain] test_df = pd.read_csv( os.path.join(args.data_dir, "test", subject + "_test.csv"), header=None ) num_questions.append(test_df.shape[0]) k = args.ntrain few_shot_examples = gen_prompt(dev_df, subject, k) while len(tokenizer.encode(few_shot_examples)) > 1536: k -= 1 few_shot_examples = gen_prompt(dev_df, subject, k) for i in range(test_df.shape[0]): prompt_end = format_example(test_df, i, include_answer=False) arguments.append( { "examples": few_shot_examples, "question": prompt_end, } ) label = test_df.iloc[i, test_df.shape[1] - 1] labels.append(label) ##################################### ######### SGL Program Begin ######### ##################################### import sglang as sgl if args.backend.startswith("gpt-"): @sgl.function def few_shot_mmlu(s, examples, question): s += sgl.user(examples + question) s += sgl.assistant(sgl.gen("answer")) else: @sgl.function def few_shot_mmlu(s, examples, question): s += examples + question + sgl.gen("answer") ##################################### ########## SGL Program End ########## ##################################### # Select backend backend = select_sglang_backend(args) # Run tic = time.perf_counter() states = few_shot_mmlu.run_batch( arguments, temperature=0, max_new_tokens=1, backend=backend, num_threads=args.parallel, progress_bar=True, ) preds = [ s["answer"].strip()[0] if len(s["answer"].strip()) > 0 else "" for s in states ] latency = time.perf_counter() - tic # Compute accuracy cors = [pred == label for pred, label in zip(preds, labels)] pt = 0 for subject, num_qs in zip(subjects[: args.nsub], num_questions): print( f"subject: {subject}, #q:{num_qs}, acc: {np.mean(cors[pt: pt + num_qs]):.3f}" ) pt += num_qs assert pt == len(cors) weighted_acc = np.mean(cors) dump_bench_raw_result( path=args.raw_result_file, states=states, preds=preds, labels=labels, ) # Print results print("Total latency: {:.3f}".format(latency)) print("Average accuracy: {:.3f}".format(weighted_acc)) # Write results with open(args.result_file, "a") as fout: value = { "task": "mmlu", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "accuracy": round(weighted_acc, 3), "num_requests": len(arguments), "other": { "nsub": args.nsub, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--ntrain", "-k", type=int, default=5) parser.add_argument( "--data_dir", "-d", type=str, default=os.path.join(SCRIPT_DIR, "data") ) parser.add_argument("--save_dir", "-s", type=str, default="results") parser.add_argument("--nsub", type=int, default=60) args = add_common_sglang_args_and_parse(parser) download_data(args.data_dir) main(args) ================================================ FILE: benchmark/mmlu/download_data.sh ================================================ wget https://people.eecs.berkeley.edu/~hendrycks/data.tar tar xf data.tar ================================================ FILE: benchmark/mmmu/README.md ================================================ ## Run evaluation ### Evaluate sglang Host the VLM: ``` python -m sglang.launch_server --model-path Qwen/Qwen2-VL-7B-Instruct --port 30000 ``` It's recommended to reduce the memory usage by appending something like `--mem-fraction-static 0.6` to the command above. Benchmark: ``` python benchmark/mmmu/bench_sglang.py --port 30000 --concurrency 16 ``` You can adjust the `--concurrency` to control the number of concurrent OpenAI calls. You can use `--lora-path` to specify the LoRA adapter to apply during benchmarking. E.g., ``` # Launch server with LoRA enabled python -m sglang.launch_server --model-path microsoft/Phi-4-multimodal-instruct --port 30000 --trust-remote-code --disable-radix-cache --lora-paths vision= # Apply LoRA adapter during inferencing python -m benchmark/mmmu/bench_sglang.py --concurrency 8 --lora-path vision ``` You can use `--response-answer-regex` to specify how to extract the answer from the response string. E.g., ``` python3 -m sglang.launch_server --model-path zai-org/GLM-4.1V-9B-Thinking --reasoning-parser glm45 python3 bench_sglang.py --response-answer-regex "<\|begin_of_box\|>(.*)<\|end_of_box\|>" --concurrency 64 ``` You can use `--extra-request-body` to specify additional OpenAI request parameters. E.g., ``` python3 bench_sglang.py --extra-request-body '{"max_new_tokens": 128, "temperature": 0.01}' ``` ### Evaluate HF ``` python benchmark/mmmu/bench_hf.py --model-path Qwen/Qwen2-VL-7B-Instruct ``` # Profiling MMMU You should use the standard instructions found in the [dedicated profiling doc](../../docs/developer_guide/benchmark_and_profiling.md) if running this benchmark with the profile option. We recommend using `--concurrency 1` for consistency, which makes profiling and debugging easier. ================================================ FILE: benchmark/mmmu/bench_hf.py ================================================ import argparse import PIL import torch from data_utils import save_json from eval_utils import ( EvalArgs, eval_result, get_sampling_params, prepare_samples, process_result, ) from tqdm import tqdm from transformers import AutoModel, AutoProcessor, GenerationConfig @torch.no_grad() def eval_mmmu(args): eval_args = EvalArgs.from_cli_args(args) sampling_params = get_sampling_params(eval_args) generation_config = GenerationConfig( max_new_tokens=sampling_params["max_new_tokens"], do_sample=False, ) try: from transformers import AutoModelForImageTextToText model = AutoModelForImageTextToText.from_pretrained( args.model_path, torch_dtype="auto", trust_remote_code=True, ) except Exception as first_exception: try: # check if the model is belongs to internvl if "InternVL" in args.model_path: from transformers import AutoTokenizer from sglang.srt.multimodal.internvl_utils import image_to_pixel_values tokenizer = AutoTokenizer.from_pretrained(args.model_path) model = AutoModel.from_pretrained( args.model_path, torch_dtype="auto", trust_remote_code=True, ) generation_config_internvl = dict( max_new_tokens=sampling_params["max_new_tokens"], do_sample=False ) else: model = AutoModel.from_pretrained( args.model_path, torch_dtype="auto", trust_remote_code=True, init_tts=False, ) except Exception as second_exception: raise RuntimeError( f"Failed to load model: First attempt failed with {first_exception}, " f"second attempt failed with {second_exception}" ) from second_exception model = model.eval().cuda() processor = AutoProcessor.from_pretrained( args.model_path, torch_dtype="auto", device_map="auto", trust_remote_code=True ) samples = prepare_samples(eval_args) out_samples = dict() answer_dict = {} for sample in tqdm(samples): prompt = sample["final_input_prompt"] image = sample["image"] prefix = prompt.split("<")[0] suffix = prompt.split(">")[1] assert image is not None if "InternVL" in args.model_path: image = PIL.Image.open(sample["image_path"]).convert("RGB") pixel_values = image_to_pixel_values( image, input_size=448, max_num=12, use_thumbnail=True ) pixel_values = pixel_values.to(device="cuda", dtype=torch.bfloat16) contents = "" if prefix: contents += prefix contents += "\n" if suffix: contents += suffix response = model.chat( tokenizer, pixel_values, contents, generation_config_internvl ) print(f"response: {response}") process_result(response, sample, answer_dict, out_samples) continue contents = [] if prefix: contents += [{"type": "text", "text": prefix}] contents += [ { "type": "image", "image": sample["image_path"], } ] if suffix: contents += [{"type": "text", "text": suffix}] messages = [{"role": "user", "content": contents}] try: model_inputs = processor.tokenizer.apply_chat_template( messages, tokenize=True, return_dict=True, add_generation_prompt=True, return_tensors="pt", ).to(model.device) input_len = model_inputs["input_ids"].shape[-1] generation = model.generate( **model_inputs, generation_config=generation_config ) generation = generation[0][input_len:] response = processor.decode(generation, skip_special_tokens=True) except: contents = [] if prefix: contents += [prefix] image = PIL.Image.open(sample["image_path"]) contents += [image] if suffix: contents += [suffix] messages = [{"role": "user", "content": contents}] response = model.chat( msgs=messages, tokenizer=processor.tokenizer, sampling=False, max_new_tokens=sampling_params["max_new_tokens"], use_tts_template=False, generate_audio=False, temperature=0.0, ) print(f"response: {response}") process_result(response, sample, answer_dict, out_samples) args.output_path = f"{args.model_path}_answer_hf.json" save_json(args.output_path, out_samples) eval_result( model_answer_path=args.output_path, answer_dict=answer_dict, eval_output_path=f"{args.model_path}_val_hf.json", ) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--model-path", type=str, help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.", required=True, ) EvalArgs.add_cli_args(parser) args = parser.parse_args() eval_mmmu(args) ================================================ FILE: benchmark/mmmu/bench_sglang.py ================================================ """ Bench the sglang-hosted vLM with benchmark MMMU Usage: Host the VLM: python -m sglang.launch_server --model-path Qwen/Qwen2-VL-7B-Instruct --port 30000 Benchmark: python benchmark/mmmu/bench_sglang.py --port 30000 --concurrency 16 The eval output will be logged """ import argparse import asyncio import base64 import mimetypes import re import sys import time import traceback from dataclasses import dataclass, field from pathlib import Path from typing import Any, List, Optional, Tuple import aiohttp import openai from data_utils import save_json from eval_utils import ( EvalArgs, eval_result, get_sampling_params, prepare_samples, process_result, ) from tqdm import tqdm from sglang.test.test_utils import add_common_sglang_args_and_parse AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60) @dataclass class RequestFuncOutput: generated_text: List[str] = field(default_factory=list) prompt_len: List[int] = field(default_factory=list) output_len: List[int] = field(default_factory=list) latency: List[float] = field(default_factory=list) ttft: List[float] = field(default_factory=list) itl: List[float] = field(default_factory=list) # List of inter-token latencies success: bool = False error: str = "" async def async_request_profile(api_url: str) -> RequestFuncOutput: async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: output = RequestFuncOutput() try: async with session.post(url=api_url) as response: if response.status == 200: output.success = True else: output.error = response.reason or "" output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) return output def _get_prefix_suffix(prompt: str) -> Tuple[str, str]: """Split the prompt into prefix and suffix.""" prefix = prompt.split("<")[0] suffix = prompt.split(">", 1)[1] return prefix, suffix async def process_sample( client: Any, sample: dict, sampling_params: dict, model: str, reasoning_effort: Optional[str] = None, lora_path: Optional[str] = None, ) -> Tuple[dict, str]: """Send a single sample to the LLM and return (sample, response).""" prompt = sample["final_input_prompt"] prefix, suffix = _get_prefix_suffix(prompt) image = sample["image"] assert image is not None image_path = sample["image_path"] if image_path and not image_path.startswith(("http://", "https://", "data:")): p = Path(image_path) mime = mimetypes.guess_type(str(p))[0] or "image/png" with open(p, "rb") as f: b64 = base64.b64encode(f.read()).decode() image_url = f"data:{mime};base64,{b64}" else: image_url = image_path extra_body = {"lora_path": lora_path} if lora_path else None payload = { "model": model, "messages": [ { "role": "user", "content": [ {"type": "text", "text": prefix}, {"type": "image_url", "image_url": {"url": image_url}}, {"type": "text", "text": suffix}, ], } ], "extra_body": extra_body, **sampling_params, } if reasoning_effort: payload["reasoning_effort"] = reasoning_effort response = await client.chat.completions.create(**payload) msg = response.choices[0].message content = msg.content if content is None: content = getattr(msg, "reasoning_content", None) return sample, content async def process_sample_with_semaphore( semaphore: asyncio.Semaphore, client: Any, sample: dict, sampling_params: dict, model: str, reasoning_effort: Optional[str] = None, lora_path: Optional[str] = None, ) -> Tuple[dict, str]: """Wrap process_sample with a semaphore for concurrency control.""" async with semaphore: return await process_sample( client, sample, sampling_params, model, reasoning_effort, lora_path ) async def eval_mmmu(args) -> None: """Main evaluation loop with concurrency control.""" eval_args = EvalArgs.from_cli_args(args) sampling_params = get_sampling_params(eval_args) samples = prepare_samples(eval_args) model = args.model reasoning_effort = eval_args.reasoning_effort lora_path = eval_args.lora_path answer_dict = {} out_samples = {} client = openai.AsyncOpenAI( api_key="sk", base_url=f"http://127.0.0.1:{args.port}/v1", timeout=20 * 60 * 60, ) start = time.perf_counter() base_url = f"http://127.0.0.1:{args.port}" if args.profile: print("Starting profiler...") profile_output = await async_request_profile( api_url=f"{base_url}/start_profile" ) if profile_output.success: print("Profiler started") samples = samples[: args.profile_number] if args.concurrency == 1: # For concurrency == 1, run in sequential mode to ensure consistent order # this is mainly for profiling for sample in tqdm(samples): _, response = await process_sample( client, sample, sampling_params, model, reasoning_effort, lora_path ) sample["original_response"] = response answer = ( re.search(args.response_answer_regex, response) if response is not None else None ) process_result( answer.group(1).strip() if answer else response, sample, answer_dict, out_samples, ) else: semaphore = asyncio.Semaphore(args.concurrency) tasks = [ process_sample_with_semaphore( semaphore, client, sample, sampling_params, model, reasoning_effort, lora_path, ) for sample in samples ] for coro in tqdm(asyncio.as_completed(tasks), total=len(tasks)): sample, response = await coro sample["original_response"] = response answer = ( re.search(args.response_answer_regex, response) if response is not None else None ) process_result( answer.group(1).strip() if answer else response, sample, answer_dict, out_samples, ) if args.profile: print("Stopping profiler...") profile_output = await async_request_profile(api_url=f"{base_url}/stop_profile") if profile_output.success: print("Profiler stopped") print(f"Benchmark time: {time.perf_counter() - start}") args.output_path = "./answer_sglang.json" save_json(args.output_path, out_samples) eval_result( model_answer_path=args.output_path, answer_dict=answer_dict, eval_output_path="./val_sglang.json", ) def parse_args(): parser = argparse.ArgumentParser() parser.add_argument( "--model", type=str, default="default", help="Model name to use in API requests.", ) EvalArgs.add_cli_args(parser) args = add_common_sglang_args_and_parse(parser) return args def main(): args = parse_args() asyncio.run(eval_mmmu(args)) if __name__ == "__main__": main() ================================================ FILE: benchmark/mmmu/data_utils.py ================================================ """Utils for data load, save, and process (e.g., prompt construction)""" import json import os import re import yaml DOMAIN_CAT2SUB_CAT = { "Art and Design": ["Art", "Art_Theory", "Design", "Music"], "Business": ["Accounting", "Economics", "Finance", "Manage", "Marketing"], "Science": [ "Biology", "Chemistry", "Geography", "Math", "Physics", ], "Health and Medicine": [ "Basic_Medical_Science", "Clinical_Medicine", "Diagnostics_and_Laboratory_Medicine", "Pharmacy", "Public_Health", ], "Humanities and Social Science": [ "History", "Literature", "Sociology", "Psychology", ], "Tech and Engineering": [ "Agriculture", "Architecture_and_Engineering", "Computer_Science", "Electronics", "Energy_and_Power", "Materials", "Mechanical_Engineering", ], } CAT_SHORT2LONG = { "acc": "Accounting", "agri": "Agriculture", "arch": "Architecture_and_Engineering", "art": "Art", "art_theory": "Art_Theory", "bas_med": "Basic_Medical_Science", "bio": "Biology", "chem": "Chemistry", "cli_med": "Clinical_Medicine", "cs": "Computer_Science", "design": "Design", "diag_med": "Diagnostics_and_Laboratory_Medicine", "econ": "Economics", "elec": "Electronics", "ep": "Energy_and_Power", "fin": "Finance", "geo": "Geography", "his": "History", "liter": "Literature", "manage": "Manage", "mark": "Marketing", "mate": "Materials", "math": "Math", "mech": "Mechanical_Engineering", "music": "Music", "phar": "Pharmacy", "phys": "Physics", "psy": "Psychology", "pub_health": "Public_Health", "socio": "Sociology", } def get_multi_choice_info(options): """ Given the list of options for multiple choice question Return the index2ans and all_choices """ start_chr = "A" all_choices = [] index2ans = {} for i, option in enumerate(options): index2ans[chr(ord(start_chr) + i)] = option all_choices.append(chr(ord(start_chr) + i)) return index2ans, all_choices def load_yaml(file_path): with open(file_path, "r") as stream: try: yaml_dict = yaml.safe_load(stream) except yaml.YAMLError as exc: print(exc) return yaml_dict def parse_img_path(text): matches = re.findall("", text) return matches def process_single_sample(data): question = data["question"] o_imgs_paths = [] for option in data["options"]: current_o_imgs_paths = parse_img_path(option) for img_path in current_o_imgs_paths: o_imgs_paths.append(img_path) if len(o_imgs_paths) > 1: # multiple images in options, used for random selection return { "id": data["id"], "question": question, "options": data["options"], "answer": data["answer"], "image": None, "question_type": data["question_type"], } else: return { "id": data["id"], "question": question, "options": data["options"], "answer": data["answer"], "image": data["image_1"], "question_type": data["question_type"], } # DATA SAVING def save_json(filename, ds): print(f"answers saved to: {filename}") os.makedirs(os.path.dirname(filename), exist_ok=True) with open(filename, "w") as f: json.dump(ds, f, indent=4) def save_jsonl(filename, data): """ Save a dictionary of data to a JSON Lines file with the filename as key and caption as value. Args: filename (str): The path to the file where the data should be saved. data (dict): The dictionary containing the data to save where key is the image path and value is the caption. """ with open(filename, "w", encoding="utf-8") as f: for img_path, caption in data.items(): # Extract the base filename without the extension base_filename = os.path.basename(img_path) # Create a JSON object with the filename as the key and caption as the value json_record = json.dumps({base_filename: caption}, ensure_ascii=False) # Write the JSON object to the file, one per line f.write(json_record + "\n") def save_args(args, path_dir): argsDict = args.__dict__ with open(path_dir + "setting.txt", "w") as f: f.writelines("------------------ start ------------------" + "\n") for eachArg, value in argsDict.items(): f.writelines(eachArg + " : " + str(value) + "\n") f.writelines("------------------- end -------------------") # DATA PROCESSING def construct_prompt(sample, config): question = sample["question"] options = eval(sample["options"]) example = "" if sample["question_type"] == "multiple-choice": start_chr = "A" prediction_range = [] index2ans = {} for option in options: prediction_range.append(start_chr) example += f"({start_chr}) {option}\n" index2ans[start_chr] = option start_chr = chr(ord(start_chr) + 1) empty_prompt_sample_structure = config["multi_choice_example_format"] empty_prompt = empty_prompt_sample_structure.format(question, example) res_dict = {} res_dict["index2ans"] = index2ans res_dict["correct_choice"] = sample["answer"] res_dict["all_choices"] = prediction_range res_dict["empty_prompt"] = empty_prompt if config["task_instructions"]: res_dict["final_input_prompt"] = ( config["task_instructions"].strip() + "\n\n" + empty_prompt ) else: res_dict["final_input_prompt"] = empty_prompt res_dict["gt_content"] = options[ord(sample["answer"].upper()) - ord("A")] else: empty_prompt_sample_structure = config["short_ans_example_format"] empty_prompt = empty_prompt_sample_structure.format(question) res_dict = {} res_dict["empty_prompt"] = empty_prompt if config["task_instructions"]: res_dict["final_input_prompt"] = ( config["task_instructions"].strip() + "\n\n" + empty_prompt ) else: res_dict["final_input_prompt"] = empty_prompt res_dict["gt_content"] = sample["answer"] res_dict.update(sample) return res_dict ================================================ FILE: benchmark/mmmu/eval_utils.py ================================================ """Response Parsing and Evaluation for various models""" import argparse import dataclasses import json import os import pprint import random import re from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Dict, Optional import numpy as np import torch from data_utils import ( CAT_SHORT2LONG, DOMAIN_CAT2SUB_CAT, construct_prompt, load_yaml, process_single_sample, save_json, ) from datasets import concatenate_datasets, load_dataset from tqdm import tqdm @dataclasses.dataclass class EvalArgs: seed: int = 42 split: str = "validation" image_pixels_limit: int = -1 result_filename: str = f"./val_sglang.json" prompt_format_file: str = "prompt_format.yaml" dataset_path: str = "MMMU/MMMU" extra_request_body: Optional[str] = None profile: bool = False profile_number: int = 5 concurrency: int = 1 max_new_tokens: Optional[int] = None temperature: Optional[float] = None response_answer_regex: str = "(.*)" lora_path: Optional[str] = None reasoning_effort: Optional[str] = None @staticmethod def add_cli_args(parser: argparse.ArgumentParser): parser.add_argument( "--result-filename", type=str, default=EvalArgs.result_filename, help="The filename to save the evaluation results.", ) parser.add_argument( "--image-pixels-limit", type=int, default=EvalArgs.image_pixels_limit, help="The maximum number of pixels allowed for an image. If an image exceeds this limit, it will be skipped during evaluation.", ) parser.add_argument( "--dataset-path", type=str, default=EvalArgs.dataset_path, help="path to the dataset", ) parser.add_argument("--seed", type=int, default=1, help="The random seed.") parser.add_argument( "--prompt-format-file", type=str, help="The path to the prompt format of mmmu. If not, a default format llava_config.yaml will be used", ) parser.add_argument( "--split", type=str, default=EvalArgs.split, help='Split of the dataset to use for evaluation. Default is "validation".', ) parser.add_argument( "--extra-request-body", metavar='{"key1": "value1", "key2": "value2"}', type=str, default=EvalArgs.extra_request_body, help="Append given JSON object to the request payload. You can use this to specify" "additional generate params like sampling params.", ) parser.add_argument( "--profile", action="store_true", help="enable mmmu profile" ) parser.add_argument( "--profile-number", type=int, default=EvalArgs.profile_number, help="Number of samples to profile. If not set, will profile all samples.", ) parser.add_argument( "--concurrency", type=int, default=EvalArgs.concurrency, help="Number of concurrent requests to make during evaluation. Default is 1, which means no concurrency.", ) parser.add_argument( "--max-new-tokens", type=int, default=EvalArgs.max_new_tokens, help="Maximum number of new tokens to generate per sample.", ) parser.add_argument( "--temperature", type=float, default=EvalArgs.temperature, help="Sampling temperature for generation.", ) parser.add_argument( "--response-answer-regex", type=str, default=EvalArgs.response_answer_regex, help="Specific regex to capture the answer from the response, string", ) parser.add_argument( "--lora-path", type=str, default=EvalArgs.lora_path, help="Specify the LoRA path to use for evaluation. If specified, the value will be specified in the body of every request as `lora-path`.", ) parser.add_argument( "--reasoning-effort", type=str, default=EvalArgs.reasoning_effort, choices=["none", "high"], help="Reasoning effort for the model (none or high).", ) @classmethod def from_cli_args(cls, args: argparse.Namespace): attrs = [attr.name for attr in dataclasses.fields(cls)] return cls(**{attr: getattr(args, attr) for attr in attrs}) def set_seed(seed_value): """ Set the seed for PyTorch (both CPU and CUDA), Python, and NumPy for reproducible results. :param seed_value: An integer value to be used as the seed. """ torch.manual_seed(seed_value) if torch.cuda.is_available(): torch.cuda.manual_seed(seed_value) torch.cuda.manual_seed_all(seed_value) # For multi-GPU setups random.seed(seed_value) np.random.seed(seed_value) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False def prepare_samples(eval_args: EvalArgs): print("Preparing samples...") # Build prompts set_seed(eval_args.seed) prompt_format_file = ( eval_args.prompt_format_file if eval_args.prompt_format_file is not None else os.path.join(os.path.dirname(__file__), "prompt_format.yaml") ) # load config and process to one value eval_args.config = load_yaml(prompt_format_file) for key, value in eval_args.config.items(): if key != "eval_params" and type(value) == list: assert len(value) == 1, "key {} has more than one value".format(key) eval_args.config[key] = value[0] # run for each subject in parallel sub_dataset_list = [] subjects = list(CAT_SHORT2LONG.values()) # Get a fixed list of subjects print(f"Loading datasets for {len(subjects)} subjects...") with ThreadPoolExecutor() as executor: # Submit all load_dataset tasks future_to_subject = { executor.submit( load_dataset, eval_args.dataset_path, subject, split=eval_args.split ): subject for subject in subjects } # Collect results as they complete results = {} for future in tqdm( as_completed(future_to_subject), total=len(subjects), desc="Loading datasets", ): subject = future_to_subject[future] try: results[subject] = future.result() except Exception as exc: print(f"{subject} generated an exception: {exc}") # Ensure datasets are added in the original order for consistency for subject in subjects: if subject in results: sub_dataset_list.append(results[subject]) else: # Handle cases where a dataset failed to load (optional, depends on desired behavior) print(f"Warning: Dataset for subject '{subject}' could not be loaded.") # merge all dataset dataset = concatenate_datasets(sub_dataset_list) # Prepare images in parallel images_path = os.path.expanduser("~/.cache/mmmu/images") os.makedirs(images_path, exist_ok=True) print(f"Saving images to: {images_path}") samples = [] skip_count = 0 def process_sample(i, sample): sample = process_single_sample(sample) sample = construct_prompt(sample, eval_args.config) image = sample["image"] width, height = image.size if 0 < eval_args.image_pixels_limit <= width * height: return None, True # Use a unique identifier for the image path to avoid potential collisions if indices reset image_path = f"{images_path}/image_{sample['id']}.png" if not os.path.exists(image_path): image.save(image_path) sample["image_path"] = image_path return sample, False print("Processing samples...") with ThreadPoolExecutor() as executor: # Pass the sample itself to process_sample, index is less reliable now futures = [ executor.submit( process_sample, i, sample ) # Keep index i for tqdm maybe? Or remove it. Let's keep it for now. for i, sample in enumerate(dataset) ] for future in tqdm( as_completed(futures), total=len(dataset), desc="Processing samples" ): sample, skipped = future.result() if skipped: skip_count += 1 elif sample: samples.append(sample) samples.sort(key=lambda x: x["final_input_prompt"]) print( f"Skipping {skip_count} samples with large images, {round((float(skip_count) / len(dataset)) * 100, 2)}% of dataset" ) print("Samples have been prepared") return samples def get_sampling_params(eval_args): extra_request_body = {} if eval_args.extra_request_body: extra_request_body = json.loads(eval_args.extra_request_body) sampling_params = { **extra_request_body, } if eval_args.max_new_tokens is not None and eval_args.max_new_tokens > 0: sampling_params.update({"max_completion_tokens": eval_args.max_new_tokens}) if eval_args.temperature is not None: sampling_params.update({"temperature": eval_args.temperature}) return sampling_params # ----------- Process Multi-choice ------------- def parse_multi_choice_response(response, all_choices, index2ans): """ Parse the prediction from the generated response. Return the predicted index e.g., A, B, C, D. """ for char in [",", ".", "!", "?", ";", ":", "'"]: response = response.strip(char) response = " " + response + " " # add space to avoid partial match index_ans = True ans_with_brack = False candidates = [] for choice in all_choices: # e.g., (A) (B) (C) (D) if f"({choice})" in response: candidates.append(choice) ans_with_brack = True if len(candidates) == 0: for choice in all_choices: # e.g., A B C D if f" {choice} " in response: candidates.append(choice) # if all above doesn't get candidates, check if the content is larger than 5 tokens and try to parse the example if len(candidates) == 0 and len(response.split()) > 5: for index, ans in index2ans.items(): if ans.lower() in response.lower(): candidates.append(index) index_ans = False # it's content ans. if len(candidates) == 0: # still not get answer, randomly choose one. pred_index = random.choice(all_choices) elif len(candidates) > 1: start_indexes = [] if index_ans: if ans_with_brack: for can in candidates: index = response.rfind(f"({can})") start_indexes.append(index) # -1 will be ignored anyway # start_indexes = [generated_response.index(f'({can})') for can in candidates] else: for can in candidates: index = response.rfind(f" {can} ") start_indexes.append(index) else: for can in candidates: index = response.lower().rfind(index2ans[can].lower()) start_indexes.append(index) # get the last one pred_index = candidates[np.argmax(start_indexes)] else: # if only one candidate, use it. pred_index = candidates[0] return pred_index # ----------- Process Open ------------- def check_is_number(string): """ Check if the given string a number. """ try: float(string.replace(",", "")) return True except ValueError: # check if there's comma inside return False def normalize_str(string): """ Normalize the str to lower case and make them float numbers if possible. """ # check if characters in the string # if number, numerize it. string = string.strip() is_number = check_is_number(string) if is_number: string = string.replace(",", "") string = float(string) # leave 2 decimal string = round(string, 2) return [string] else: # it's likely to be a string # lower it string = string.lower() if len(string) == 1: return [" " + string, string + " "] # avoid trivial matches return [string] def extract_numbers(string): """ Exact all forms of numbers from a string with regex. """ # Pattern for numbers with commas pattern_commas = r"-?\b\d{1,3}(?:,\d{3})+\b" # Pattern for scientific notation pattern_scientific = r"-?\d+(?:\.\d+)?[eE][+-]?\d+" # Pattern for simple numbers without commas pattern_simple = r"-?(?:\d+\.\d+|\.\d+|\d+\b)(?![eE][+-]?\d+)(?![,\d])" # Extract numbers with commas numbers_with_commas = re.findall(pattern_commas, string) # Extract numbers in scientific notation numbers_scientific = re.findall(pattern_scientific, string) # Extract simple numbers without commas numbers_simple = re.findall(pattern_simple, string) # Combine all extracted numbers all_numbers = numbers_with_commas + numbers_scientific + numbers_simple return all_numbers def parse_open_response(response): """ Parse the prediction from the generated response. Return a list of predicted strings or numbers. """ # content = content.strip("\n").strip(".").strip(" ") def get_key_subresponses(response): key_responses = [] response = response.strip().strip(".").lower() sub_responses = re.split(r"\.\s(?=[A-Z])|\n", response) indicators_of_keys = [ "could be ", "so ", "is ", "thus ", "therefore ", "final ", "answer ", "result ", ] key_responses = [] for index, resp in enumerate(sub_responses): # if last one, accept it's an equation (the entire response can be just one sentence with equation) if index == len(sub_responses) - 1: indicators_of_keys.extend(["="]) shortest_key_response = None # the shortest response that may contain the answer (tail part of the response) for indicator in indicators_of_keys: if indicator in resp: if not shortest_key_response: shortest_key_response = resp.split(indicator)[-1].strip() else: if len(resp.split(indicator)[-1].strip()) < len( shortest_key_response ): shortest_key_response = resp.split(indicator)[-1].strip() # key_responses.append(resp.split(indicator)[1].strip()) if shortest_key_response: # and it's not trivial if shortest_key_response.strip() not in [ ":", ",", ".", "!", "?", ";", ":", "'", ]: key_responses.append(shortest_key_response) if len(key_responses) == 0: # did not found any return [response] return key_responses # pdb.set_trace() key_responses = get_key_subresponses(response) pred_list = key_responses.copy() # keep the original string response for resp in key_responses: pred_list.extend(extract_numbers(resp)) tmp_pred_list = [] for i in range(len(pred_list)): tmp_pred_list.extend(normalize_str(pred_list[i])) pred_list = tmp_pred_list # remove duplicates pred_list = list(set(pred_list)) return pred_list # ----------- Evaluation ------------- def eval_multi_choice(gold_i, pred_i): """ Evaluate a multiple choice instance. """ correct = False # for case like Answer: A, Answer is A, answer is A, answer: A for _exp in ["Answer:", "Answer is ", "answer is ", "answer: "]: if _exp in pred_i: pred_i = pred_i.split(_exp)[1].strip() break # for case like (A), (B), (C), (D) ...... if "(" in pred_i and ")" in pred_i: try: pred_i = re.search(r"\(([A-Z])\)", pred_i).group(1) except: print(f"Error to extract answer from: {pred_i}") pass # only they are exactly the same, we consider it as correct if isinstance(gold_i, list): for answer in gold_i: if answer == pred_i: correct = True break else: # gold_i is a string if gold_i == pred_i: correct = True return correct def eval_open(gold_i, pred_i): """ Evaluate an open question instance """ correct = False if isinstance(gold_i, list): # use float to avoid trivial matches norm_answers = [] for answer in gold_i: norm_answers.extend(normalize_str(answer)) else: norm_answers = normalize_str(gold_i) for pred in pred_i: # pred is already normalized in parse response phase if isinstance(pred, str): # if it's a string, then find if ans in the pred_i for norm_ans in norm_answers: # only see if the string answer in the string pred if isinstance(norm_ans, str) and norm_ans in pred: if not correct: correct = True break else: # it's a float number if pred in norm_answers: if not correct: correct = True break return correct # ----------- Batch Evaluation ------------- def evaluate(samples): """ Batch evaluation for multiple choice and open questions. """ pred_correct = 0 judge_dict = dict() for sample in samples: gold_i = sample["answer"] pred_i = sample["parsed_pred"] if sample["question_type"] == "multiple-choice": correct = eval_multi_choice(gold_i, pred_i) else: # open question correct = eval_open(gold_i, pred_i) if correct: judge_dict[sample["id"]] = "Correct" pred_correct += 1 else: # print(f"Wrong! expected {pred_i}, answered with {gold_i}") judge_dict[sample["id"]] = "Wrong" if len(samples) == 0: return {"acc": 0} return judge_dict, {"acc": pred_correct / len(samples)} # ----------- Calculate Accuracy ------------- def calculate_ins_level_acc(results: Dict): """Calculate the instruction level accuracy for given Subject results""" acc = 0 ins_num = 0 for cat_results in results.values(): acc += cat_results["acc"] * cat_results["num_example"] ins_num += cat_results["num_example"] if ins_num == 0: return 0 return acc / ins_num def process_result(response, sample, answer_dict, out_samples): if response is None: return if sample["question_type"] == "multiple-choice": pred_ans = parse_multi_choice_response( response, sample["all_choices"], sample["index2ans"] ) else: # open question pred_ans = response out_samples[sample["id"]] = { "pred_ans": pred_ans, "original_response": sample["original_response"], "ground_truth": sample["answer"], "question_type": sample["question_type"], } # set ground truth answer answer_dict[sample["id"]] = { "question_type": sample["question_type"], "ground_truth": sample["answer"], } def eval_result(model_answer_path, answer_dict, eval_output_path=None): if eval_output_path is None: eval_output_path = model_answer_path print("Evaluating...") output_dict = json.load(open(model_answer_path)) # answer_dict = json.load(open(answer_path)) # group by category output_dict_w_cat = {} for data_id, parsed_pred in output_dict.items(): if isinstance(parsed_pred, str): parsed_pred = parsed_pred elif isinstance(parsed_pred, dict): parsed_pred = parsed_pred["pred_ans"] else: raise ValueError(f"Unknown type of parsed_pred: {type(parsed_pred)}") category = "_".join(data_id.split("_")[1:-1]) if category not in output_dict_w_cat: output_dict_w_cat.update({category: {}}) output_dict_w_cat[category].update({data_id: parsed_pred}) # group by category answer_dict_w_cat = {} for data_id, parsed_pred in answer_dict.items(): category = "_".join(data_id.split("_")[1:-1]) if category not in answer_dict_w_cat: answer_dict_w_cat.update({category: {}}) answer_dict_w_cat[category].update({data_id: parsed_pred}) evaluation_result = {} for category in CAT_SHORT2LONG.values(): # print("Evaluating: {}".format(category)) # get cat_outputs and cat_answers try: cat_outputs = output_dict_w_cat[category] cat_answers = answer_dict_w_cat[category] except KeyError: # print("Skipping {} for not found".format(category)) continue exampels_to_eval = [] for data_id, parsed_pred in cat_outputs.items(): question_type = cat_answers[data_id]["question_type"] if question_type != "multiple-choice": parsed_pred = parse_open_response( parsed_pred ) # mainly for type consistency (make it number, etc.) else: parsed_pred = parsed_pred exampels_to_eval.append( { "id": data_id, "question_type": question_type, "answer": cat_answers[data_id]["ground_truth"], "parsed_pred": parsed_pred, } ) judge_dict, metric_dict = evaluate(exampels_to_eval) metric_dict.update({"num_example": len(exampels_to_eval)}) for key, value in judge_dict.items(): output_dict[key]["judge"] = value evaluation_result[category] = metric_dict save_json(model_answer_path, output_dict) printable_results = {} # pdb.set_trace() # add domain Subject for domain, in_domain_cats in DOMAIN_CAT2SUB_CAT.items(): in_domain_cat_results = {} for cat_name in in_domain_cats: # use the order in DOMAIN_CAT2SUB_CAT if cat_name in evaluation_result.keys(): in_domain_cat_results[cat_name] = evaluation_result[cat_name] else: pass in_domain_ins_acc = calculate_ins_level_acc(in_domain_cat_results) in_domain_data_num = sum( [ cat_results["num_example"] for cat_results in in_domain_cat_results.values() ] ) printable_results["Overall-" + domain] = { "num": int(in_domain_data_num), "acc": round(in_domain_ins_acc, 3), } # add sub category for cat_name, cat_results in in_domain_cat_results.items(): printable_results[cat_name] = { "num": int(cat_results["num_example"]), "acc": round(cat_results["acc"], 3), } # table.append(["-----------------------------", "-----", "----"]) all_ins_acc = calculate_ins_level_acc(evaluation_result) overall_acc = round(all_ins_acc, 3) printable_results["Overall"] = { "num": sum( [cat_results["num_example"] for cat_results in evaluation_result.values()] ), "acc": overall_acc, } pprint.pprint(printable_results) out = eval_output_path with open(out, "w", encoding="utf-8") as outfile: json.dump(printable_results, outfile) print(f"eval out saved to {out}") print(f"Overall accuracy: {overall_acc}") ================================================ FILE: benchmark/mmmu/prompt_format.yaml ================================================ task_instructions: - "" multi_choice_example_format: - "{} {} Answer with the option's letter from the given choices directly." short_ans_example_format: - "{} Answer the question using a single word or phrase." temperature: - 0 ================================================ FILE: benchmark/mtbench/README.md ================================================ ## Download Dataset ```sh wget -O question.jsonl https://raw.githubusercontent.com/lm-sys/FastChat/main/fastchat/llm_judge/data/mt_bench/question.jsonl ``` ## Run benchmark ### Benchmark sglang ``` python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` ``` python3 bench_sglang.py --num-questions 80 ``` ### Benchmark sglang EAGLE ``` python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3-8B-Instruct --speculative-algo EAGLE \ --speculative-draft-model-path lmsys/sglang-EAGLE-LLaMA3-Instruct-8B --speculative-num-steps 5 \ --speculative-eagle-topk 8 --speculative-num-draft-tokens 64 --dtype float16 --port 30000 ``` ``` python3 bench_sglang_eagle.py --num-questions 80 --parallel 1 ``` ### Benchmark vllm ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` ``` python3 bench_other.py --num-questions 80 --backend vllm ``` ### Benchmark lightllm ``` # A10G python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 ``` ``` python3 bench_other.py --num-questions 80 --backend lightllm ``` ================================================ FILE: benchmark/mtbench/bench_other.py ================================================ import argparse import json import os import time import uuid from concurrent.futures import ThreadPoolExecutor from fastchat.model import get_conversation_template from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import download_and_cache_file def load_questions(filename): questions = [] with open(filename, "r") as fin: for line in fin: obj = json.loads(line) questions.append(obj) return questions def write_answers(filename, model_id, questions, answers): with open(os.path.expanduser(filename), "w") as fout: for i in range(len(answers)): ans_json = { "question_id": questions[i]["question_id"], "answer_id": uuid.uuid4().hex, "model_id": model_id, "choices": { "index": 0, "turns": [answers[i][0], answers[i][1]], }, "tstamp": time.time(), } fout.write(json.dumps(ans_json) + "\n") def main(args): # Download question file if not exist question_file = args.question_file url = "https://raw.githubusercontent.com/lm-sys/FastChat/main/fastchat/llm_judge/data/mt_bench/question.jsonl" if not os.path.isfile(question_file): question_file = download_and_cache_file(url) questions = load_questions(question_file) questions = (questions * 10)[: args.num_questions] max_tokens = 256 model_id = "llama-2-chat" conv_main = get_conversation_template(model_id) # Select backend call_generate = get_call_generate(args) answers = [None] * len(questions) def get_answer(i): conv = conv_main.copy() cur_answers = [] for j in range(2): q = questions[i]["turns"][j] conv.append_message(conv.roles[0], q) conv.append_message(conv.roles[1], None) prompt = conv.get_prompt() output = call_generate(prompt, temperature=0, max_tokens=max_tokens).strip() cur_answers.append(output) conv.update_last_message(output) answers[i] = cur_answers # Run requests tic = time.perf_counter() if args.parallel == 1: for i in tqdm(range(len(questions))): get_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: list( tqdm( executor.map(get_answer, list(range(len(questions)))), total=len(questions), ) ) latency = time.perf_counter() - tic print(f"#questions: {len(questions)}, Latency: {latency:.2f}") # Write results answer_file = args.answer_file or f"tmp_output_{args.backend}.txt" write_answers(answer_file, model_id, questions, answers) with open(args.result_file, "a") as fout: value = { "task": "mtbench", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--question-file", type=str, default="question.jsonl") parser.add_argument("--answer-file", type=str, default=None) parser.add_argument("--num-questions", type=int, default=80) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/mtbench/bench_sglang.py ================================================ import argparse import json import os import time import uuid import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import download_and_cache_file def load_questions(filename): questions = [] with open(filename, "r") as fin: for line in fin: obj = json.loads(line) questions.append(obj) return questions def write_answers(filename, model_id, questions, answers): with open(os.path.expanduser(filename), "w") as fout: for i in range(len(answers)): ans_json = { "question_id": questions[i]["question_id"], "answer_id": uuid.uuid4().hex, "model_id": model_id, "choices": { "index": 0, "turns": [answers[i][0], answers[i][1]], }, "tstamp": time.time(), } fout.write(json.dumps(ans_json) + "\n") @sgl.function def answer_mt_bench(s, question_1, question_2): s += sgl.system() s += sgl.user(question_1) s += sgl.assistant(sgl.gen("answer_1")) s += sgl.user(question_2) s += sgl.assistant(sgl.gen("answer_2")) def main(args): # Download question file if not exist question_file = args.question_file url = "https://raw.githubusercontent.com/lm-sys/FastChat/main/fastchat/llm_judge/data/mt_bench/question.jsonl" if not os.path.isfile(question_file): question_file = download_and_cache_file(url) # Construct prompts questions = load_questions(question_file)[: args.num_questions] arguments = [ {"question_1": q["turns"][0], "question_2": q["turns"][1]} for q in questions ] # Select backend backend = select_sglang_backend(args) sgl.set_default_backend(backend) # Run requests tic = time.perf_counter() rets = answer_mt_bench.run_batch( arguments, temperature=0, max_new_tokens=256, num_threads=args.parallel, progress_bar=True, ) answers = [[s["answer_1"], s["answer_2"]] for s in rets] latency = time.perf_counter() - tic print(f"#questions: {len(questions)}, Latency: {latency:.2f}") # Write results model_id = backend.model_info["model_path"] answer_file = args.answer_file or f"tmp_output_{args.backend}.txt" write_answers(answer_file, model_id, questions, answers) with open(args.result_file, "a") as fout: value = { "task": "mtbench", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--question-file", type=str, default="question.jsonl") parser.add_argument("--answer-file", type=str, default=None) parser.add_argument("--num-questions", type=int, default=80) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/mtbench/bench_sglang_eagle.py ================================================ """ Adapted from https://github.com/chromecast56/sglang/blob/6f145d2eadb93a116134f703358ce76f15381045/benchmark/mtbench/bench_sglang.py Benchmark SGLang EAGLE/EAGLE3 Speculative Decoding Usage: python3 benchmark/mtbench/bench_sglang_eagle.py --num-questions 80 --parallel 1 """ import argparse import json import os import time import uuid import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import download_and_cache_file def load_questions(filename): questions = [] with open(filename, "r") as fin: for line in fin: obj = json.loads(line) questions.append(obj) return questions def write_answers(filename, model_id, questions, answers): with open(os.path.expanduser(filename), "w") as fout: for i in range(len(answers)): ans_json = { "question_id": questions[i]["question_id"], "answer_id": uuid.uuid4().hex, "model_id": model_id, "choices": { "index": 0, "turns": [answers[i][0], answers[i][1]], }, "tstamp": time.time(), } fout.write(json.dumps(ans_json) + "\n") @sgl.function def answer_mt_bench(s, question_1, question_2): s += sgl.system( "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information." ) s += sgl.user(question_1) s += sgl.assistant(sgl.gen("answer_1")) s += sgl.user(question_2) s += sgl.assistant(sgl.gen("answer_2")) def main(args): # Download question file if not exist question_file = args.question_file url = "https://raw.githubusercontent.com/lm-sys/FastChat/main/fastchat/llm_judge/data/mt_bench/question.jsonl" if not os.path.isfile(question_file): question_file = download_and_cache_file(url) # Construct prompts questions = load_questions(question_file)[: args.num_questions] arguments = [ {"question_1": q["turns"][0], "question_2": q["turns"][1]} for q in questions ] # Select backend backend = select_sglang_backend(args) sgl.set_default_backend(backend) # Run requests tic = time.perf_counter() rets = answer_mt_bench.run_batch( arguments, temperature=0, max_new_tokens=2048, num_threads=args.parallel, progress_bar=True, ) answers = [[s["answer_1"], s["answer_2"]] for s in rets] latency = time.perf_counter() - tic num_output_tokens = sum( s.get_meta_info("answer_1")["completion_tokens"] + s.get_meta_info("answer_2")["completion_tokens"] for s in rets ) # NOTE: acceptance length is just completion_tokens / spec_verify_ct # {'id': '3bb9c5ead109488d8ed5ee9cbecaec29', 'finish_reason': {'type': 'length', 'length': 256}, 'prompt_tokens': 37, 'spec_verify_ct': 101, 'completion_tokens': 256, 'cached_tokens': 0} output_throughput = num_output_tokens / latency has_verify = "spec_verify_ct" in rets[0].get_meta_info("answer_1") if has_verify: num_verify_tokens = sum( s.get_meta_info("answer_1")["spec_verify_ct"] + s.get_meta_info("answer_2")["spec_verify_ct"] for s in rets ) accept_length = num_output_tokens / num_verify_tokens else: accept_length = 1.0 print( f"#questions: {len(questions)}, Throughput: {output_throughput:.2f} token/s, Acceptance length: {accept_length:.2f}" ) # Write results model_id = backend.model_info["model_path"] answer_file = args.answer_file or f"tmp_output_{args.backend}.txt" write_answers(answer_file, model_id, questions, answers) with open(args.result_file, "a") as fout: value = { "task": "mtbench", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "throughput": round(output_throughput, 3), "accept_length": round(accept_length, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--question-file", type=str, default="question.jsonl") parser.add_argument("--answer-file", type=str, default=None) parser.add_argument("--num-questions", type=int, default=80) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/multi_chain_reasoning/README.md ================================================ ## Download data ``` wget https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl ``` ## Run benchmark ### Benchmark sglang ``` python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --schedule-conservativeness 1.3 ``` ``` python3 bench_sglang.py --num-questions 64 python3 bench_sglang.py --num-questions 32 --parallel 1 ``` ### Benchmark vllm ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` ``` python3 bench_other.py --num-questions 64 --backend vllm ``` ### Benchmark lightllm ``` # A10G python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 ``` ``` python3 bench_other.py --num-questions 64 --backend lightllm ``` ### Benchmark guidance ``` python3 bench_other.py --num-questions 8 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` ### Benchmark lmql ``` python3 bench_other.py --num-questions 64 --backend lmql --parallel 1 ``` ================================================ FILE: benchmark/multi_chain_reasoning/bench_other.py ================================================ import argparse import ast import asyncio import json import re import time from concurrent.futures import ThreadPoolExecutor import numpy as np from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import dump_state_text, read_jsonl INVALID = -9999999 def get_answer_value(answer_str): answer_str = answer_str.replace(",", "") numbers = re.findall(r"\d+", answer_str) if len(numbers) < 1: return INVALID try: return ast.literal_eval(numbers[-1]) except SyntaxError: return INVALID prompt_lib = [ "Let us think step by step.", "Approach this methodically. Let's dissect the problem into smaller, more manageable parts.", "It's important to proceed step by step, ensuring accuracy at each stage.", "Take a deep breath and break this down.", "A little bit of arithmetic and a logical approach will help us quickly arrive at the solution to this problem.", "I am extremely good at math.", ] def multi_chain_gsm8k(question, num_chains, call_generate): s = "Question: " + question + "\n" # s += call_generate(s + "Answer: " + prompt_lib[0], max_tokens=256, # stop="Question", temperature=0) # return s comps = [] for i in range(num_chains): comps.append( call_generate( s + "Answer: " + prompt_lib[i % num_chains], max_tokens=256, temperature=0.3, stop="Question", ) ) s += "Answer: To answer this question, here are some possible solutions. " s += "After considering all of them, I will do a majority vote.\n\n" for i in range(num_chains): s += f"Solution {i+1}: " + comps[i].strip() + "\n\n" s += "\nBy considering the above solutions and doing a majority vote, I think the final answer (a single integer number) is " s += call_generate(s, max_tokens=16, temperature=0, stop=None) return s async def multi_chain_gsm8k_async(question, num_chains, call_generate): s = "Question: " + question + "\n" # s += call_generate(s + "Answer: " + prompt_lib[0], max_tokens=256, # stop="Question", temperature=0) # return s comps = [] for i in range(num_chains): comps.append( await call_generate( s + "Answer: " + prompt_lib[i % num_chains], max_tokens=256, temperature=0.3, stop="Question", ) ) s += "Answer: To answer this question, here are some possible solutions. " s += "After considering all of them, I will do a majority vote.\n\n" for i in range(num_chains): s += f"Solution {i+1}: " + comps[i].strip() + "\n\n" s += "\nBy considering the above solutions and doing a majority vote, I think the final answer (a single integer number) is " s += await call_generate(s, max_tokens=16, temperature=0, stop=None) return s def main(args): lines = list(read_jsonl(args.data_path)) # Construct prompts k = args.num_shot questions = [] labels = [] for i in range(len(lines[: args.num_questions])): questions.append(lines[i]["question"]) labels.append(get_answer_value(lines[i]["answer"])) assert all(l != INVALID for l in labels) states = [None] * len(labels) # Select backend call_generate = get_call_generate(args) # Run requests if args.backend != "lmql": # Use thread pool def get_one_answer(i): answer = multi_chain_gsm8k(questions[i], args.num_chains, call_generate) states[i] = answer tic = time.perf_counter() if args.parallel == 1: for i in tqdm(range(len(questions))): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: list( tqdm( executor.map(get_one_answer, list(range(len(questions)))), total=len(questions), ) ) else: # Use asyncio async def get_one_answer_asyncio(i): answer = await multi_chain_gsm8k_async( questions[i], args.num_chains, call_generate ) states[i] = answer tic = time.perf_counter() loop = asyncio.get_event_loop() batches = [ list(range(i, min(i + args.parallel, len(questions)))) for i in range(0, len(questions), args.parallel) ] for bt in tqdm(batches): tasks = [get_one_answer_asyncio(k) for k in bt] loop.run_until_complete(asyncio.gather(*tasks)) latency = time.perf_counter() - tic preds = [] for i in range(len(states)): preds.append(get_answer_value(states[i])) # Compute accuracy acc = np.mean(np.array(preds) == np.array(labels)) invalid = np.mean(np.array(preds) == INVALID) print(f"Latency: {latency:.3f}") print(f"Invalid: {invalid:.3f}") print(f"Accuracy: {acc:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "multi_chain_gsm8k", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "accuracy": round(acc, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--num-shot", type=int, default=0) parser.add_argument("--num-chains", type=int, default=5) parser.add_argument("--data-path", type=str, default="test.jsonl") parser.add_argument("--num-questions", type=int, default=50) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/multi_chain_reasoning/bench_sglang.py ================================================ import argparse import ast import json import re import time import numpy as np from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text, read_jsonl INVALID = -9999999 def get_answer_value(answer_str): answer_str = answer_str.replace(",", "") numbers = re.findall(r"\d+", answer_str) if len(numbers) < 1: return INVALID try: return ast.literal_eval(numbers[-1]) except SyntaxError: return INVALID prompt_lib = [ "Let us think step by step.", "Approach this methodically. Let's dissect the problem into smaller, more manageable parts.", "It's important to proceed step by step, ensuring accuracy at each stage.", "Take a deep breath and break this down.", "A little bit of arithmetic and a logical approach will help us quickly arrive at the solution to this problem.", "I am extremely good at math.", ] def main(args): lines = list(read_jsonl(args.data_path)) # Construct prompts # k = args.num_shot # few_shot_examples = get_few_shot_examples(lines, k) questions = [] labels = [] for i in range(len(lines[: args.num_questions])): questions.append(lines[i]["question"]) labels.append(get_answer_value(lines[i]["answer"])) assert all(l != INVALID for l in labels) arguments = [{"question": q} for q in questions] num_chains = args.num_chains ##################################### ######### SGL Program Begin ######### ##################################### import sglang as sgl @sgl.function def multi_chain_gsm8k(s, question): s += "Question: " + question + "\n" # s += "Answer: " + prompt_lib[0] + sgl.gen("answer", max_tokens=256, stop="Question", # temperature=0) # return forks = s.fork(num_chains) for i in range(num_chains): forks[i] += ( "Answer: " + prompt_lib[i % num_chains] + sgl.gen("chain", max_tokens=256, temperature=0.3, stop="Question") ) forks.join() s += "Answer: To answer this question, here are some possible solutions. " s += "After considering all of them, I will do a majority vote.\n\n" for i in range(num_chains): s += f"Solution {i+1}: " + forks[i]["chain"].strip() + "\n\n" s += "\nBy considering the above solutions and doing a majority vote, I think the final answer (a single integer number) is " s += sgl.gen("answer", max_tokens=16) ##################################### ########## SGL Program End ########## ##################################### # Select backend backend = select_sglang_backend(args) # Run requests tic = time.perf_counter() states = multi_chain_gsm8k.run_batch( arguments, temperature=0, backend=backend, num_threads=args.parallel, progress_bar=True, ) latency = time.perf_counter() - tic preds = [] for i in range(len(states)): preds.append(get_answer_value(states[i]["answer"])) # Compute accuracy acc = np.mean(np.array(preds) == np.array(labels)) invalid = np.mean(np.array(preds) == INVALID) print(f"Latency: {latency:.3f}") print(f"Invalid: {invalid:.3f}") print(f"Accuracy: {acc:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "multi_chain_gsm8k", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "accuracy": round(acc, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--num-shot", type=int, default=0) parser.add_argument("--num-chains", type=int, default=5) parser.add_argument("--data-path", type=str, default="test.jsonl") parser.add_argument("--num-questions", type=int, default=50) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/multi_document_qa/README.md ================================================ ## Run benchmark ### Benchmark sglang ``` python3 -m sglang.launch_server --model-path codellama/CodeLlama-7b-instruct-hf --port 30000 ``` ``` python3 bench_sglang.py --num-questions 10 --parallel 1 ``` ### Benchmark vllm ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model codellama/CodeLlama-7b-instruct-hf --disable-log-requests --port 21000 --gpu 0.97 ``` ``` python3 bench_other.py --backend vllm --num-questions 64 ``` ### Benchmark guidance ``` python3 bench_other.py --backend guidance --num-questions 32 --parallel 1 --n-ctx 11000 --model-path path/to/code-llama/gguf ``` ### Build dataset ``` pip install PyPDF2 python3 build_dataset.py ``` ```python import PyPDF2 with open('llama2.pdf', 'rb') as file: reader = PyPDF2.PdfReader(file) text = '' for page_num in range(len(reader.pages)): text += reader.pages[page_num].extract_text() with open('output.txt', 'w') as text_file: text_file.write(text) ``` ================================================ FILE: benchmark/multi_document_qa/bench_other.py ================================================ import argparse import json import time from concurrent.futures import ThreadPoolExecutor from functools import partial from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import dump_state_text, read_jsonl USER_PREFIX = "[INST] " USER_SUFFIX = " [/INST]" ASSISTANT_PREFIX = "" ASSISTANT_SUFFIX = " " def multi_document_qa(docs, question, generate): s = USER_PREFIX s += "Please answer a question according to given documents.\n" s += "Question:" + question + "Documents begin.\n" s += "".join(docs) s += "\nDocuments end." s += ( "\n\nBased on the above documents, please answer this question:\n" + question + "\nAnswer in three words or fewer." ) s += USER_SUFFIX s += ASSISTANT_PREFIX answer = generate(s, max_tokens=16, stop=None) return answer def main(args): lines = read_jsonl(args.data_path) l = lines[0] arguments = [] labels = [] num_docs = 10 if args.backend == "guidance": num_docs = 7 # due to OOM for i in range(len(l["questions"][: args.num_questions])): arguments.append( { "docs": l["documents"][:num_docs], "question": l["questions"][i], } ) labels.append(l["answers"][i]) states = [None] * len(arguments) # Select backend call_generate = partial(get_call_generate(args), temperature=0) # Run requests def get_one_answer(i): states[i] = multi_document_qa(generate=call_generate, **arguments[i]) tic = time.perf_counter() if args.parallel == 1: for i in tqdm(range(len(labels))): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: list( tqdm( executor.map(get_one_answer, list(range(len(labels)))), total=len(labels), ) ) latency = time.perf_counter() - tic # Compute accuracy print(states) correct = 0 for s, label in zip(states, labels): answer = s.lower() if all(x in answer for x in label.lower().split(" ")): correct += 1 accuracy = correct / len(labels) print(f"Accuracy: {accuracy:.3f}") print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "multi_document_qa", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_questions, "accuracy": accuracy, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="questions.jsonl") parser.add_argument("--num-questions", type=int, default=100) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/multi_document_qa/bench_sglang.py ================================================ import argparse import json import time import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text, read_jsonl @sgl.function def multi_document_qa(s, docs, question): s += sgl.user_begin() s += "Please answer a question according to given documents.\n" s += "Question:" + question + "Documents begin.\n" forks = s.fork(len(docs)) forks += lambda i: docs[i] forks.join("concate_and_append") s += "\nDocuments end." s += ( "\n\nBased on the above documents, please answer this question:\n" + question + "\nAnswer in three words or fewer." ) s += sgl.user_end() s += sgl.assistant(sgl.gen("answer", max_tokens=16)) def main(args): lines = read_jsonl(args.data_path) l = lines[0] arguments = [] labels = [] for i in range(len(l["questions"][: args.num_questions])): arguments.append( { "docs": l["documents"][:10], "question": l["questions"][i], } ) labels.append(l["answers"][i]) # Select backend backend = select_sglang_backend(args) sgl.set_default_backend(backend) # Run requests tic = time.perf_counter() states = multi_document_qa.run_batch( arguments, temperature=0, num_threads=args.parallel, progress_bar=True ) latency = time.perf_counter() - tic # Compute accuracy print([s["answer"] for s in states]) correct = 0 for s, label in zip(states, labels): answer = s["answer"].lower() if all(x in answer for x in label.lower().split(" ")): correct += 1 accuracy = correct / len(labels) print(f"Accuracy: {accuracy:.3f}") print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "multi_document_qa", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_questions, "accuracy": accuracy, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="questions.jsonl") parser.add_argument("--num-questions", type=int, default=100) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/multi_document_qa/build_dataset.py ================================================ import json import transformers content = "\n".join( open("llama2.txt", "r", encoding="utf-8", errors="ignore").readlines() ) content = content.replace("\n\n", "\n") # Count token name = "meta-llama/Llama-2-7b-chat-hf" t = transformers.AutoTokenizer.from_pretrained(name) print(f"num tokens: {len(t.encode(content))}") # Segment SEP = "\n\n" parts = content.split(SEP) print(f"num segments: {len(parts)}") segment_len = 1100 segments = [] tmp = [] tmp_len = 0 for i in range(len(parts)): tmp.append(parts[i]) tmp_len += len(t.encode(parts[i])) if tmp_len > segment_len: segments.append(SEP.join(tmp)) tmp = [] tmp_len = 0 for i, s in enumerate(segments): print(i, len(t.encode(segments[i]))) # Dump with open("questions.jsonl", "w") as fout: fout.write( json.dumps( { "documents": segments[:30], "questions": [ "What is the name of the fine-tuned LLMs?", "Which figure shows the helpfulness human evaluation results for Llama 2-Chat?", "What is the number of parameters in the largest Llama 2 model?", "What is the batch size of fine-tuning?", "Where can we find the details of potential data contamination?", "What is the full name of MPT?", "What is the power consumption of RSC in Watt?", "How many tokens of data do they train on?", "Which model's release is delayed due to a lack of time to sufficiently red team?", "Which activation function is used in Llama?", ], "answers": [ "Llama 2 Chat", "1", "70 B", "64", "A 6", "MosaicML", "400", "2 trillion", "34 B", "SwiGLU", ], } ) + "\n" ) ================================================ FILE: benchmark/multi_turn_chat/README.md ================================================ ### Benchmark sglang Run Llama-7B ``` python3 -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` Run Mixtral-8x7B (When there is a CUDA out-of-memory error, try to reduce the `--mem-fraction-static`) ``` python3 -m sglang.launch_server --model-path mistralai/Mixtral-8x7B-Instruct-v0.1 --port 30000 --tp-size 8 ``` Benchmark(short output) ``` python3 bench_sglang.py --tokenizer meta-llama/Llama-2-7b-chat-hf ``` Benchmark(long output) ``` python3 bench_sglang.py --tokenizer meta-llama/Llama-2-7b-chat-hf --long ``` ### Benchmark vLLM Run Llama-7B ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` Run Mixtral-8x7B ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model mistralai/Mixtral-8x7B-Instruct-v0.1 --disable-log-requests --port 21000 --tensor-parallel-size 8 ``` Benchmark(short output) ``` python3 bench_other.py --tokenizer meta-llama/Llama-2-7b-chat-hf --backend vllm ``` Benchmark(long output) ``` python3 bench_other.py --tokenizer meta-llama/Llama-2-7b-chat-hf --backend vllm --long ``` ### Benchmark guidance Benchmark Llama-7B (short output) ``` python3 bench_other.py --tokenizer meta-llama/Llama-2-7b-chat-hf --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` Benchmark Llama-7B (long output) ``` python3 bench_other.py --tokenizer meta-llama/Llama-2-7b-chat-hf --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf --long ``` ================================================ FILE: benchmark/multi_turn_chat/bench_other.py ================================================ import json import time from argparse import ArgumentParser from concurrent.futures import ThreadPoolExecutor from functools import partial from data_gen import gen_arguments from tqdm import tqdm from vllm.transformers_utils.tokenizer import get_tokenizer from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import dump_state_text def multi_turns(generate, qas): s = "" for qa in qas: s += qa["prompt"] s += generate(s, max_tokens=qa["new_tokens"]) return s def main(args): print(args) tokenizer = get_tokenizer(args.tokenizer, trust_remote_code=args.trust_remote_code) multi_qas = gen_arguments(args, tokenizer) states = [None] * args.num_qa call_generate = partial(get_call_generate(args), temperature=0) def get_one_answer(i): states[i] = multi_turns(generate=call_generate, **multi_qas[i]) tic = time.perf_counter() if args.parallel == 1: for i in tqdm(range(len(multi_qas))): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: rets = list( tqdm( executor.map(get_one_answer, list(range(len(multi_qas)))), total=len(multi_qas), ) ) for _ in rets: pass latency = time.perf_counter() - tic # Compute accuracy print(f"Latency: {latency:.3f}") dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "multi_turn_chat", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_qa, "num_turns": args.turns, "other": { "parallel": args.parallel, "output_mode": "long" if args.long else "short", }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = ArgumentParser() parser.add_argument("--turns", type=int, default=4) parser.add_argument("--num-qa", type=int, default=20) parser.add_argument("--min-len-q", type=int, default=256) parser.add_argument("--max-len-q", type=int, default=512) parser.add_argument("--min-len-a", type=int, default=4) parser.add_argument("--max-len-a", type=int, default=8) parser.add_argument("--tokenizer", type=str, required=True) parser.add_argument("--trust-remote-code", action="store_true") parser.add_argument("--long", action="store_true") args = add_common_other_args_and_parse(parser) if args.long: args.min_len_a = 256 args.max_len_a = 512 args.num_qa = 20 main(args) ================================================ FILE: benchmark/multi_turn_chat/bench_sglang.py ================================================ import json import time from argparse import ArgumentParser from data_gen import gen_arguments from vllm.transformers_utils.tokenizer import get_tokenizer import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text @sgl.function def multi_turns(s, qas): for qa in qas: s += qa["prompt"] s += sgl.gen(max_tokens=qa["new_tokens"], ignore_eos=True) def main(args): tokenizer = get_tokenizer(args.tokenizer, trust_remote_code=args.trust_remote_code) multi_qas = gen_arguments(args, tokenizer) backend = select_sglang_backend(args) tic = time.perf_counter() states = multi_turns.run_batch( multi_qas, temperature=0, backend=backend, num_threads=args.parallel, progress_bar=True, ) latency = time.perf_counter() - tic print(f"Latency: {latency:.3f}") dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "multi_turn_chat", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_qa, "num_turns": args.turns, "other": { "parallel": args.parallel, "output_mode": "long" if args.long else "short", }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = ArgumentParser() parser.add_argument("--turns", type=int, default=4) parser.add_argument("--num-qa", type=int, default=20) parser.add_argument("--min-len-q", type=int, default=256) parser.add_argument("--max-len-q", type=int, default=512) parser.add_argument("--min-len-a", type=int, default=4) parser.add_argument("--max-len-a", type=int, default=8) parser.add_argument("--tokenizer", type=str, required=True) parser.add_argument("--trust-remote-code", action="store_true") parser.add_argument("--long", action="store_true") args = add_common_sglang_args_and_parse(parser) if args.long: args.min_len_a = 256 args.max_len_a = 512 args.num_qa = 20 print(args) main(args) ================================================ FILE: benchmark/multi_turn_chat/data_gen.py ================================================ import random import string random.seed(42) def gen_prompt(tokenizer, token_num): cha_set = string.ascii_letters + string.digits ret = "".join(random.choices(cha_set, k=token_num)) while len(tokenizer(ret).input_ids) < token_num: ret += random.choice(cha_set) return ret def gen_arguments(args, tokenizer): multi_qas = [{"qas": []} for _ in range(args.num_qa)] for i in range(args.num_qa): qas = multi_qas[i]["qas"] for _ in range(args.turns): prompt_len = random.randint(args.min_len_q, args.max_len_q) new_tokens = random.randint(args.min_len_a, args.max_len_a) qas.append( { "prompt": gen_prompt(tokenizer, prompt_len), "new_tokens": new_tokens, } ) return multi_qas ================================================ FILE: benchmark/multi_turn_chat/long_prompt_multi_turn.py ================================================ import json import random import time from argparse import ArgumentParser from pathlib import Path from tqdm import tqdm import sglang as sgl from sglang.srt.utils.hf_transformers_utils import get_tokenizer from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text def gen_prompt(tokenizer, token_num): all_available_tokens = list(tokenizer.get_vocab().values()) selected_tokens = random.choices(all_available_tokens, k=token_num) ret = tokenizer.decode(selected_tokens) return ret def get_cache_path(args): # Create cache directory under ~/.cache/sglang cache_dir = Path.home() / ".cache" / "sglang" # Create a unique cache filename based on the arguments that affect generation cache_key = f"qa_{args.num_qa}_{args.turns}_{args.system_prompt_len}_{args.len_q}_{args.len_a}_{args.tokenizer.replace('/', '_')}.json" return cache_dir / cache_key def gen_arguments(args, tokenizer): cache_path = get_cache_path(args) # Try to load from cache first if cache_path.exists(): print(f"Loading cached arguments from {cache_path}") with open(cache_path, "r") as f: return json.load(f) print("Generating new arguments...") # First progress bar for system prompts multi_qas = [] for _ in tqdm(range(args.num_qa), desc="Generating system prompts"): multi_qas.append( {"system_prompt": gen_prompt(tokenizer, args.system_prompt_len), "qas": []} ) # Nested progress bars for QA pairs for i in tqdm(range(args.num_qa), desc="Generating QA pairs"): qas = multi_qas[i]["qas"] for j in range(args.turns): qas.append( { "prompt": gen_prompt(tokenizer, args.len_q), "new_tokens": args.len_a, } ) # Save to cache cache_path.parent.mkdir(parents=True, exist_ok=True) with open(cache_path, "w") as f: json.dump(multi_qas, f) print(f"Cached arguments saved to {cache_path}") return multi_qas @sgl.function def multi_turns(s, system_prompt, qas): s += system_prompt for i, qa in enumerate(qas): s += qa["prompt"] s += sgl.gen(max_tokens=qa["new_tokens"], ignore_eos=True) def main(args): tokenizer = get_tokenizer(args.tokenizer, trust_remote_code=args.trust_remote_code) multi_qas = gen_arguments(args, tokenizer) backend = select_sglang_backend(args) tic = time.perf_counter() states = multi_turns.run_batch( multi_qas, temperature=0, backend=backend, num_threads="auto", progress_bar=True, ) latency = time.perf_counter() - tic print(f"Latency: {latency:.3f}") dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "multi_turn_system_prompt_chat", "backend": args.backend, "latency": round(latency, 3), "num_requests": args.num_qa, "num_turns": args.turns, "other": { "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = ArgumentParser() parser.add_argument("--turns", type=int, default=8) parser.add_argument("--num-qa", type=int, default=128) parser.add_argument("--system-prompt-len", type=int, default=2048) parser.add_argument("--len-q", type=int, default=32) parser.add_argument("--len-a", type=int, default=128) parser.add_argument( "--tokenizer", type=str, default="meta-llama/Meta-Llama-3-8B-Instruct" ) parser.add_argument("--trust-remote-code", action="store_true") args = add_common_sglang_args_and_parse(parser) print(args) main(args) ================================================ FILE: benchmark/prefill_only/bench_embeddings.py ================================================ """ SGLang Embeddings Benchmark Script This script benchmarks SGLang's /v1/embeddings API performance using HTTP requests. Features: - HTTP-only implementation - Uses /v1/embeddings API endpoint directly - Configurable RPS, duration, and batch sizes - Progress tracking and detailed metrics - Poisson and constant request distributions Usage: - Update configuration variables at the top of the file - Ensure SGLang server is running on the configured HTTP_URL - Run: python bench_embeddings.py """ import asyncio import logging from typing import Optional from transformers import AutoTokenizer from util import ( BenchmarkConfig, generate_text_with_token_count, run_benchmark_main, run_generic_benchmark, ) # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) ############################################################################### # CONFIG ############################################################################### # Create benchmark configuration config = BenchmarkConfig() config.rps_values = [500] config.duration_secs_values = [60] config.num_unique_requests = 100 config.distribution = "POISSON" config.profile = False config.freeze_gc = True # Enable GC freeze functionality # Profiler output directory - by default uses present working directory (pwd) # Uncomment and customize the line below to override the default location: # config.profiler_dir = "/sglang-oss-trace" # HTTP Configuration HTTP_URL = "http://localhost:30000/v1/embeddings" # Embeddings API Config EMBEDDINGS_MODEL_PATH = "Qwen/Qwen3-Embedding-0.6B" BATCH_SIZE = [1] # Number of items per request (batch size) # Configurable input token length EMBEDDINGS_INPUT_TOKENS = 500 # Default token length MATRYOSHKA_DIMENSIONS: Optional[int] = ( None # Set to None to disable matryoshka embeddings ) # Load tokenizer once for embeddings text generation print("Loading tokenizer for embeddings input generation...") embeddings_tokenizer = AutoTokenizer.from_pretrained(EMBEDDINGS_MODEL_PATH) # Generate input text with the specified token length using pre-loaded tokenizer EMBEDDINGS_INPUT_TEXT = generate_text_with_token_count( EMBEDDINGS_MODEL_PATH, EMBEDDINGS_INPUT_TOKENS, config.special_replicated_token, tokenizer=embeddings_tokenizer, ) ############################################################################### # REQUEST GENERATION (in parallel) ############################################################################### def build_embeddings_request(index: int, item_count: int) -> tuple: """Build a single embeddings request.""" try: # For embeddings, input can be a string or list of strings if item_count == 1: input_data = EMBEDDINGS_INPUT_TEXT else: input_data = [EMBEDDINGS_INPUT_TEXT for _ in range(item_count)] req = { "input": input_data, "model": EMBEDDINGS_MODEL_PATH, "dimensions": MATRYOSHKA_DIMENSIONS, } return (index, req) except Exception as e: logger.error(f"Error building request {index}: {e}") return (index, None) def validate_embeddings_response(response_data: dict) -> bool: """Validate embeddings API response.""" return ( "data" in response_data and len(response_data["data"][0]["embedding"]) == MATRYOSHKA_DIMENSIONS if MATRYOSHKA_DIMENSIONS else True ) def build_warmup_embeddings_request() -> dict: """Build a warmup request for the embeddings API.""" return { "input": EMBEDDINGS_INPUT_TEXT, "model": EMBEDDINGS_MODEL_PATH, "dimensions": MATRYOSHKA_DIMENSIONS, } ############################################################################### # MAIN ############################################################################### async def run_benchmark(rps, duration_secs, item_count): """Run a single embeddings benchmark with the given RPS value.""" return await run_generic_benchmark( rps=rps, duration_secs=duration_secs, item_count=item_count, config=config, http_url=HTTP_URL, build_request_func=build_embeddings_request, response_validator=validate_embeddings_response, api_name="EMBEDDINGS", request_description="embeddings requests", ) async def main(): additional_info = { "Input text length": f"{EMBEDDINGS_INPUT_TOKENS} tokens", "Input text preview": ( EMBEDDINGS_INPUT_TEXT[:100] + "..." if len(EMBEDDINGS_INPUT_TEXT) > 100 else EMBEDDINGS_INPUT_TEXT ), } await run_benchmark_main( config, run_benchmark, "EMBEDDINGS", HTTP_URL, BATCH_SIZE, additional_info, build_warmup_embeddings_request, ) if __name__ == "__main__": asyncio.run(main()) ================================================ FILE: benchmark/prefill_only/bench_score.py ================================================ """ SGLang Scoring Benchmark Script This script benchmarks SGLang's scoring API performance using HTTP requests. Current Features: - HTTP-only implementation (open source compatible) - Uses /v1/score API endpoint directly - Single item scoring with batching support - Configurable RPS, duration, and batch sizes - Progress tracking and detailed metrics - Poisson and constant request distributions Usage: - Update configuration variables at the top of the file - Ensure SGLang server is running on the configured HTTP_URL - Run: python bench_score.py - Each request will contain ITEM_COUNT_VALUES items for batch scoring """ import asyncio from transformers import AutoTokenizer from util import ( BenchmarkConfig, generate_text_with_token_count, run_benchmark_main, run_generic_benchmark, ) ############################################################################### # CONFIG ############################################################################### # Create benchmark configuration config = BenchmarkConfig() config.rps_values = [160] config.duration_secs_values = [60] config.num_unique_requests = 100 config.distribution = "POISSON" config.profile = False config.freeze_gc = True # Enable GC freeze functionality # Profiler output directory - by default uses present working directory (pwd) # Uncomment and customize the line below to override the default location: # config.profiler_dir = "/sglang-oss-trace" # HTTP Configuration HTTP_URL = "http://localhost:30000/v1/score" # Use score API directly # Score API Config # ITEM_COUNT_VALUES determines number of items per score request (batch size) SCORE_QUERY_TOKENS = 120 SCORE_ITEM_TOKENS = 180 SCORE_MODEL_PATH = "Qwen/Qwen3-0.6B" SCORE_LABEL_TOKEN_IDS = [9454, 2753] # Yes/No token IDs ITEM_COUNT_VALUES = [10] # Number of items per request # Special token to replicate for precise token counting SPECIAL_REPLICATED_TOKEN = "<|im_start|>" ############################################################################### # REQUEST GENERATION (in parallel) ############################################################################### def create_score_request_builder(): """Create a score request builder function with shared tokenizer.""" # Load tokenizer once here to verify special token and get precise counts print("Loading tokenizer...") tokenizer = AutoTokenizer.from_pretrained(SCORE_MODEL_PATH) # Verify that our special token produces exactly 1 token special_token_count = len( tokenizer.encode(config.special_replicated_token, add_special_tokens=False) ) print( f"Special token '{config.special_replicated_token}' produces " f"{special_token_count} token(s)" ) def generate_text_with_token_count_local(num_toks): """Generate text with precise token count using replicated token.""" return generate_text_with_token_count( SCORE_MODEL_PATH, num_toks, config.special_replicated_token, tokenizer=tokenizer, ) def build_score_request(index: int, item_count: int) -> tuple: """Build a single score request.""" try: # Generate query and items for score API query = generate_text_with_token_count_local(SCORE_QUERY_TOKENS) items = [ generate_text_with_token_count_local(SCORE_ITEM_TOKENS) for _ in range(item_count) ] # Return as dict for score API format score_data = { "query": query, "items": items, "label_token_ids": SCORE_LABEL_TOKEN_IDS, "model": SCORE_MODEL_PATH, } return (index, score_data) except Exception as e: print(f"Error building request {index}: {e}") return (index, None) return build_score_request def validate_score_response(response_data: dict) -> bool: """Validate score API response.""" return "scores" in response_data or "logprobs" in response_data def build_warmup_score_request() -> dict: """Build a warmup request for the score API.""" # Load tokenizer once for warmup generation tokenizer = AutoTokenizer.from_pretrained(SCORE_MODEL_PATH) warmup_query = generate_text_with_token_count( SCORE_MODEL_PATH, SCORE_QUERY_TOKENS, config.special_replicated_token, tokenizer=tokenizer, ) warmup_items = [ generate_text_with_token_count( SCORE_MODEL_PATH, SCORE_ITEM_TOKENS, config.special_replicated_token, tokenizer=tokenizer, ) for _ in range(3) ] return { "query": warmup_query, "items": warmup_items, "label_token_ids": SCORE_LABEL_TOKEN_IDS, "model": SCORE_MODEL_PATH, # Add missing parameters for consistency with the original warmup "apply_softmax": True, "item_first": False, } ############################################################################### # MAIN ############################################################################### async def run_benchmark(rps, duration_secs, item_count): """Run a single benchmark with the given RPS value.""" # Create the request builder function with shared tokenizer build_request_func = create_score_request_builder() return await run_generic_benchmark( rps=rps, duration_secs=duration_secs, item_count=item_count, config=config, http_url=HTTP_URL, build_request_func=build_request_func, response_validator=validate_score_response, api_name="SINGLE_ITEM_SCORING", request_description="score requests", ) async def main(): """Main function that runs benchmarks for all RPS values.""" additional_info = { "Query tokens per request": SCORE_QUERY_TOKENS, "Item tokens per item": SCORE_ITEM_TOKENS, } await run_benchmark_main( config, run_benchmark, "SINGLE_ITEM_SCORING", HTTP_URL, ITEM_COUNT_VALUES, additional_info, build_warmup_score_request, ) if __name__ == "__main__": asyncio.run(main()) ================================================ FILE: benchmark/prefill_only/util.py ================================================ """ Common utilities for SGLang benchmark scripts. This module contains shared code for benchmarking different SGLang APIs including scoring, embeddings, and other endpoints. """ import asyncio import concurrent.futures import json import os import random from statistics import mean from typing import Any, Callable, Dict, List, Optional, Tuple import aiohttp import numpy as np from tqdm import tqdm from transformers import AutoTokenizer class BenchmarkConfig: """Configuration for benchmark parameters.""" def __init__(self): # Common benchmark settings self.server_type = "HTTP" self.rps_values = [70] self.duration_secs_values = [60] self.num_unique_requests = 100 self.distribution = "POISSON" # Options: "CONSTANT", "POISSON" self.profile = False # Garbage Collection Control self.freeze_gc = True # Enable/disable garbage collection freezing # Profiler configuration self.profiler_dir = ( os.getcwd() ) # Default profiler output directory (current working directory) # Special token for text generation self.special_replicated_token = "<|im_start|>" def generate_text_with_token_count( model_path: str, num_tokens: int, special_token: str = "<|im_start|>", tokenizer: Optional[Any] = None, ) -> str: """ Generate text with precise token count using a replicated token. Args: model_path: Path to the model for tokenizer num_tokens: Target number of tokens special_token: Token to replicate tokenizer: Optional pre-loaded tokenizer to avoid repeated loading Returns: Generated text with approximately the target token count """ if tokenizer is None: tokenizer = AutoTokenizer.from_pretrained(model_path) # Verify token count special_token_count = len(tokenizer.encode(special_token, add_special_tokens=False)) if special_token_count == 1: # Simple case: token maps to exactly 1 token return special_token * num_tokens else: print(f"Special token '{special_token}' produces {special_token_count} tokens") # Handle case where special token produces multiple tokens repetitions = (num_tokens + special_token_count - 1) // special_token_count text = special_token * repetitions # Verify we got the expected token count actual_tokens = len(tokenizer.encode(text, add_special_tokens=False)) if actual_tokens < num_tokens: print(f"Warning: Generated {actual_tokens} tokens, expected {num_tokens}") return text def setup_profiler(config: BenchmarkConfig, benchmark_name: str) -> None: """ Set up profiler environment if profiling is enabled. Args: config: Benchmark configuration benchmark_name: Name of the benchmark (used in directory path) """ if config.profile: # Create benchmark-specific subdirectory profiler_path = os.path.join( config.profiler_dir, benchmark_name.lower().replace("_", "-") ) os.environ["SGLANG_TORCH_PROFILER_DIR"] = profiler_path print(f"Profiler enabled. Output directory: {profiler_path}") else: print("Profiler disabled") def prepare_all_requests_parallel( num_requests: int, item_count: int, build_request_func: Callable[[int, int], Tuple[int, Any]], config: BenchmarkConfig, description: str = "requests", ) -> List[Any]: """ Generic function to generate unique requests in parallel, then reuse them. Args: num_requests: Total number of requests needed item_count: Number of items per request (batch size) build_request_func: Function that takes (index, item_count) and returns (index, request_data) config: Benchmark configuration description: Description for progress bars Returns: List of request data objects """ def build_request_wrapper(index): """Wrapper to call the provided build_request_func.""" try: return build_request_func(index, item_count) except Exception as e: print(f"Error building request {index}: {e}") return (index, None) # Generate only the unique requests unique_requests = [None] * config.num_unique_requests max_workers = min(8, os.cpu_count() or 1) # Limit to 8 threads max with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [] for i in tqdm( range(config.num_unique_requests), desc=f"Submitting {description} generation tasks", ): future = executor.submit(build_request_wrapper, i) futures.append(future) # Collect results as they complete for f in tqdm( concurrent.futures.as_completed(futures), desc=f"Building unique {description}", total=config.num_unique_requests, ): try: index, req_data = f.result() if req_data is not None: unique_requests[index] = req_data else: print(f"Failed to build request {index}") except Exception as e: print(f"Error processing request result: {e}") # Check if we have any valid requests valid_requests = [req for req in unique_requests if req is not None] if not valid_requests: raise RuntimeError("Failed to generate any valid requests") print( f"Successfully generated {len(valid_requests)} out of " f"{config.num_unique_requests} unique {description}" ) # Create the full request list by cycling through unique requests print( f"Reusing {len(valid_requests)} unique {description} to create " f"{num_requests} total requests..." ) all_requests = [] for i in tqdm(range(num_requests), desc=f"Reusing {description}"): unique_index = i % len(valid_requests) all_requests.append(valid_requests[unique_index]) print(f"All {description} prepared.\n") return all_requests async def sleep_with_distribution(distribution: str, rps: float) -> None: """ Sleep according to the specified distribution pattern. Args: distribution: "CONSTANT" or "POISSON" rps: Requests per second rate """ if distribution == "CONSTANT": interval = 1 / rps await asyncio.sleep(interval) elif distribution == "POISSON": # For Poisson process, inter-arrival times follow exponential distribution interval = random.expovariate(rps) await asyncio.sleep(interval) else: raise ValueError( f"Unknown distribution: {distribution}. Use 'CONSTANT' or 'POISSON'." ) def build_http_request_json(request_data: Any) -> str: """ Generic function to build HTTP request JSON. Args: request_data: The data to serialize to JSON Returns: JSON string representation of the request data """ return json.dumps(request_data) async def make_http_call( session: aiohttp.ClientSession, request_data: Any, request_id: int, results_queue: asyncio.Queue, http_url: str, response_validator: Callable[[Dict[str, Any]], bool], api_name: str = "API", ) -> None: """ Generic HTTP call function for API requests. Args: session: aiohttp client session request_data: Data to send in the request request_id: Unique identifier for this request results_queue: Queue to put results http_url: URL to send the request to response_validator: Function to validate the response JSON api_name: Name of the API for error messages """ try: start_time = asyncio.get_running_loop().time() request_json = build_http_request_json(request_data) headers = {"Content-Type": "application/json"} async with session.post(http_url, data=request_json, headers=headers) as resp: resp_text = await resp.text() if resp.status != 200: print( f"[HTTP] {api_name} Request {request_id} failed with status " f"{resp.status}: {resp_text}" ) completion_time = asyncio.get_running_loop().time() await results_queue.put((request_id, 0, False, completion_time)) return # Parse and validate response try: response_data = json.loads(resp_text) success = response_validator(response_data) if not success: print( f"[HTTP] {api_name} Request {request_id} failed response validation" ) except json.JSONDecodeError: print( f"[HTTP] {api_name} Request {request_id} failed to parse JSON response" ) success = False completion_time = asyncio.get_running_loop().time() elapsed_time = (completion_time - start_time) * 1000 await results_queue.put((request_id, elapsed_time, success, completion_time)) except Exception as e: print(f"[HTTP] {api_name} Error for request {request_id}: {e}") completion_time = asyncio.get_running_loop().time() await results_queue.put((request_id, 0, False, completion_time)) async def send_profile_request( profile_text: str, http_url: str, session: Optional[aiohttp.ClientSession] = None ) -> None: """ Send a profile request (START_PROFILE or STOP_PROFILE) and wait for completion. Args: profile_text: "START_PROFILE" or "STOP_PROFILE" http_url: Base HTTP URL (will derive profile endpoints from this) session: Optional aiohttp session to use """ try: if session: print(f"Sending {profile_text} request via HTTP...") # Determine the correct endpoint if "/v1/" in http_url: base_url = http_url.rsplit("/v1/", 1)[0] # Remove /v1/xxx else: base_url = http_url.rsplit("/", 1)[0] # Remove last path component if profile_text == "START_PROFILE": endpoint_url = f"{base_url}/start_profile" elif profile_text == "STOP_PROFILE": endpoint_url = f"{base_url}/stop_profile" else: print(f"Unknown profile request: {profile_text}") return headers = {"Content-Type": "application/json"} async with session.post(endpoint_url, headers=headers) as resp: resp_text = await resp.text() if resp.status == 200: print(f"{profile_text} request completed") else: print( f"{profile_text} request failed with status " f"{resp.status}: {resp_text}" ) else: print(f"Cannot send {profile_text} request - missing session") except Exception as e: print(f"Error sending {profile_text} request: {e}") async def call_freeze_gc_http(session: aiohttp.ClientSession, http_url: str) -> None: """ Call the /freeze_gc HTTP endpoint. Args: session: aiohttp client session http_url: Base HTTP URL to derive the freeze_gc endpoint from """ try: # Derive freeze_gc endpoint from the API URL if "/v1/" in http_url: freeze_gc_url = http_url.rsplit("/v1/", 1)[0] + "/freeze_gc" else: freeze_gc_url = http_url.rsplit("/", 1)[0] + "/freeze_gc" print(f"Calling freeze_gc endpoint: {freeze_gc_url}") async with session.post(freeze_gc_url) as resp: if resp.status == 200: print("freeze_gc called successfully") else: resp_text = await resp.text() print(f"freeze_gc failed with status {resp.status}: {resp_text}") except Exception as e: print(f"Failed to call freeze_gc: {e}") async def send_warmup_requests( session: aiohttp.ClientSession, http_url: str, build_warmup_request_func: Callable[[], Any], num_warmup: int = 3, ) -> None: """ Send warmup requests to HTTP server. Args: session: aiohttp client session http_url: URL to send warmup requests to build_warmup_request_func: Function that returns a warmup request object num_warmup: Number of warmup requests to send """ print(f"Sending {num_warmup} HTTP warmup requests...") for i in range(num_warmup): try: warmup_data = build_warmup_request_func() request_json = build_http_request_json(warmup_data) headers = {"Content-Type": "application/json"} async with session.post( http_url, data=request_json, headers=headers ) as resp: if resp.status == 200: print(f"Warmup request {i+1}/{num_warmup} completed successfully") else: print( f"Warmup request {i+1}/{num_warmup} failed with status {resp.status}" ) except Exception as e: print(f"Warmup request {i+1}/{num_warmup} failed with error: {e}") print("HTTP warmup requests completed") async def perform_global_warmup_and_freeze( config: BenchmarkConfig, http_url: str, build_warmup_request_func: Callable[[], Any], ) -> None: """ Perform warmup and optionally GC freeze operations once before all benchmark runs. Args: config: Benchmark configuration http_url: URL for API requests build_warmup_request_func: Function that returns a warmup request object """ print("=" * 80) print(f"PERFORMING GLOBAL WARMUP{' AND GC FREEZE' if config.freeze_gc else ''}") print("=" * 80) print(f"Performing HTTP warmup{' and GC freeze' if config.freeze_gc else ''}...") async with aiohttp.ClientSession() as session: await send_warmup_requests(session, http_url, build_warmup_request_func) if config.freeze_gc: await call_freeze_gc_http(session, http_url) print( f"HTTP warmup{' and GC freeze' if config.freeze_gc else ''} completed successfully." ) print( f"Global warmup{' and GC freeze' if config.freeze_gc else ''} operations completed." ) print("=" * 80) async def process_results( results_queue: asyncio.Queue, num_requests: int, send_duration: float, total_duration: float, rps: int, duration_secs: int, item_count: int, test_start_time: float, config: BenchmarkConfig, http_mode: str = "UNKNOWN", ) -> List[Dict[str, Any]]: """ Process benchmark results and group them by minute intervals. Args: results_queue: Queue containing result tuples num_requests: Total number of requests sent send_duration: Time taken to send all requests total_duration: Total time for all requests to complete rps: Target requests per second duration_secs: Test duration in seconds item_count: Number of items per request test_start_time: Start time of the test config: Benchmark configuration http_mode: Description of the HTTP mode/API being tested Returns: List of dictionaries containing minute-by-minute results """ all_results = [] # Collect all results for _ in range(num_requests): result = await results_queue.get() request_id, elapsed_time, success, completion_time = result all_results.append( { "request_id": request_id, "elapsed_time": elapsed_time, "success": success, "completion_time": completion_time, } ) # Group results by minute intervals minute_results = [] num_minutes = int(duration_secs // 60) + (1 if duration_secs % 60 > 0 else 0) for minute in range(num_minutes): minute_start = test_start_time + (minute * 60) minute_end = test_start_time + ((minute + 1) * 60) # Filter results that completed in this minute minute_data = [ r for r in all_results if minute_start <= r["completion_time"] < minute_end ] response_times = [r["elapsed_time"] for r in minute_data if r["success"]] successful_requests = len([r for r in minute_data if r["success"]]) failed_requests = len([r for r in minute_data if not r["success"]]) avg_response_time = mean(response_times) if response_times else 0 # Calculate percentiles using numpy if response_times: p50 = np.percentile(response_times, 50) p90 = np.percentile(response_times, 90) p99 = np.percentile(response_times, 99) else: p50 = p90 = p99 = 0 minute_result = { "test_duration_secs": duration_secs, "minute_interval": minute + 1, "target_rps": rps, "item_count": item_count, "server_type": config.server_type, "distribution": config.distribution, "unique_requests": config.num_unique_requests, "total_requests": len(minute_data), "successful_requests": successful_requests, "failed_requests": failed_requests, "send_duration_secs": send_duration, "total_duration_secs": total_duration, "avg_response_time_ms": avg_response_time, "p50_response_time_ms": p50, "p90_response_time_ms": p90, "p99_response_time_ms": p99, } minute_results.append(minute_result) print( f"\nMinute {minute + 1} Summary for RPS {rps}, " f"Duration {duration_secs}s, Item Count {item_count}:" ) print(f" Requests completed in minute: {len(minute_data)}") print(f" Successful requests: {successful_requests}") print(f" Failed requests: {failed_requests}") print(f" Average response time: {avg_response_time:.2f} ms") print(f" P50 response time: {p50:.2f} ms") print(f" P90 response time: {p90:.2f} ms") print(f" P99 response time: {p99:.2f} ms") # Print overall summary all_response_times = [r["elapsed_time"] for r in all_results if r["success"]] total_successful = len([r for r in all_results if r["success"]]) total_failed = len([r for r in all_results if not r["success"]]) overall_avg = mean(all_response_times) if all_response_times else 0 if all_response_times: overall_p50 = np.percentile(all_response_times, 50) overall_p90 = np.percentile(all_response_times, 90) overall_p99 = np.percentile(all_response_times, 99) else: overall_p50 = overall_p90 = overall_p99 = 0 print( f"\nOverall Summary for RPS {rps}, Duration {duration_secs}s, " f"Item Count {item_count}:" ) print(f" Test duration: {duration_secs} seconds") print(f" Server type: {config.server_type}") print(f" HTTP mode: {http_mode}") print(f" Target RPS: {rps}") print(f" Achieved RPS: {len(all_results) / total_duration:.2f}") print(f" Item count: {item_count}") print(f" Distribution: {config.distribution}") print(f" Unique requests generated: {config.num_unique_requests}") print(f" Total requests sent: {num_requests}") print(f" Successful requests: {total_successful}") print(f" Failed requests: {total_failed}") print(f" Time to send all requests: {send_duration:.2f} seconds") print(f" Time for all requests to complete: {total_duration:.2f} seconds") print(f" Average response time: {overall_avg:.2f} ms") print(f" P50 response time: {overall_p50:.2f} ms") print(f" P90 response time: {overall_p90:.2f} ms") print(f" P99 response time: {overall_p99:.2f} ms\n") return minute_results def print_csv_results(all_results: List[Dict[str, Any]]) -> None: """ Print benchmark results in CSV format. Args: all_results: List of result dictionaries from process_results """ print("\n" + "=" * 80) print("FINAL CSV RESULTS:") print("=" * 80) # CSV Header headers = [ "test_duration_secs", "minute_interval", "target_rps", "item_count", "server_type", "distribution", "unique_requests", "total_requests", "successful_requests", "failed_requests", "send_duration_secs", "total_duration_secs", "avg_response_time_ms", "p50_response_time_ms", "p90_response_time_ms", "p99_response_time_ms", ] print(",".join(headers)) # CSV Data for result in all_results: row = [ result["test_duration_secs"], result["minute_interval"], result["target_rps"], result["item_count"], result["server_type"], result["distribution"], result["unique_requests"], result["total_requests"], result["successful_requests"], result["failed_requests"], f"{result['send_duration_secs']:.2f}", f"{result['total_duration_secs']:.2f}", f"{result['avg_response_time_ms']:.2f}", f"{result['p50_response_time_ms']:.2f}", f"{result['p90_response_time_ms']:.2f}", f"{result['p99_response_time_ms']:.2f}", ] print(",".join(map(str, row))) async def run_benchmark_main( config: BenchmarkConfig, run_single_benchmark_func, benchmark_name: str, http_url: str, item_count_values: List[int], additional_info: Optional[Dict[str, Any]] = None, build_warmup_request_func: Optional[Callable[[], Any]] = None, ) -> None: """ Main benchmark orchestration function. Args: config: Benchmark configuration run_single_benchmark_func: Async function to run a single benchmark benchmark_name: Name of the benchmark (e.g., "SCORING", "EMBEDDINGS") http_url: URL of the API endpoint item_count_values: List of item counts to test additional_info: Additional information to print in the header build_warmup_request_func: Optional function to build warmup requests """ total_combinations = ( len(config.duration_secs_values) * len(config.rps_values) * len(item_count_values) ) print( f"Running benchmarks for {len(config.duration_secs_values)} duration " f"values, {len(config.rps_values)} RPS values, and " f"{len(item_count_values)} item count values = " f"{total_combinations} total combinations" ) print(f"Server Type: {config.server_type}") print(f"HTTP Mode: {benchmark_name}") print(f"API URL: {http_url}") if additional_info: for key, value in additional_info.items(): print(f"{key}: {value}") print(f"Items per request (batch size): {item_count_values}") print(f"Profiling Enabled: {config.profile}") print(f"Duration values: {config.duration_secs_values}") print(f"RPS values: {config.rps_values}") print(f"Item count values: {item_count_values}") print("=" * 80) # Set up profiler environment setup_profiler(config, benchmark_name) # Perform global warmup and GC freeze operations if warmup function is provided if build_warmup_request_func is not None: await perform_global_warmup_and_freeze( config, http_url, build_warmup_request_func ) all_results = [] for duration_secs in config.duration_secs_values: for rps in config.rps_values: for item_count in item_count_values: result = await run_single_benchmark_func(rps, duration_secs, item_count) all_results.extend(result) # Extend with minute results print_csv_results(all_results) async def run_generic_benchmark( rps: int, duration_secs: int, item_count: int, config: BenchmarkConfig, http_url: str, build_request_func: Callable[[int, int], Tuple[int, Any]], response_validator: Callable[[Dict[str, Any]], bool], api_name: str, request_description: str = "requests", ) -> List[Dict[str, Any]]: """ Generic benchmark runner that can be used for different APIs. Args: rps: Requests per second duration_secs: Duration of the test in seconds item_count: Number of items per request (batch size) config: Benchmark configuration http_url: URL of the API endpoint build_request_func: Function to build individual requests response_validator: Function to validate API responses api_name: Name of the API for logging request_description: Description for progress bars Returns: List of dictionaries containing minute-by-minute results """ num_requests = int(rps * duration_secs) print( f"Starting benchmark with RPS={rps}, Duration={duration_secs}s, " f"Item Count={item_count}, num_requests={num_requests}" ) print(f"Server Type: {config.server_type}") print(f"HTTP Mode: {api_name}") print(f"Profiling Enabled: {config.profile}") # Build requests in parallel (unmeasured) all_requests = prepare_all_requests_parallel( num_requests, item_count, build_request_func, config, request_description ) results_queue = asyncio.Queue() tasks = [] # Track timing for sending requests send_start_time = asyncio.get_running_loop().time() # HTTP implementation async with aiohttp.ClientSession( timeout=aiohttp.ClientTimeout(total=300) ) as session: # Send START_PROFILE if profiling is enabled if config.profile: await send_profile_request("START_PROFILE", http_url, session=session) # Add progress bar for sending requests with tqdm( total=len(all_requests), desc=f"Sending HTTP {request_description} at {rps} RPS", unit="req", ) as pbar: for i, request_data in enumerate(all_requests): request_id = i + 1 tasks.append( asyncio.create_task( make_http_call( session, request_data, request_id, results_queue, http_url, response_validator, api_name, ) ) ) # Update progress bar pbar.update(1) # Throttle based on distribution if i < len(all_requests) - 1: await sleep_with_distribution(config.distribution, rps) send_end_time = asyncio.get_running_loop().time() send_duration = send_end_time - send_start_time # Wait for all requests to complete with progress tracking print(f"Waiting for {len(tasks)} HTTP {request_description} to complete...") with tqdm( total=len(tasks), desc=f"Completing HTTP {request_description}", unit="req" ) as completion_pbar: completed_tasks = [] for task in asyncio.as_completed(tasks): await task completed_tasks.append(task) completion_pbar.update(1) # Send STOP_PROFILE if profiling is enabled if config.profile: await send_profile_request("STOP_PROFILE", http_url, session=session) completion_end_time = asyncio.get_running_loop().time() total_duration = completion_end_time - send_start_time return await process_results( results_queue, num_requests, send_duration, total_duration, rps, duration_secs, item_count, send_start_time, config, api_name, ) ================================================ FILE: benchmark/react/README.md ================================================ ## Run benchmark NOTE: This is an implementation for replaying a given trace for throughput/latency benchmark purposes. It is not an actual ReAct agent implementation. ### Benchmark sglang ``` python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` ``` python3 bench_sglang.py --num-questions 100 ``` ### Benchmark vllm ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` ``` python3 bench_other.py --num-questions 100 --backend vllm ``` ### Benchmark guidance ``` python3 bench_other.py --num-questions 100 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` ### Benchmark lmql ``` python3 bench_other.py --num-questions 100 --backend lmql --parallel 1 ``` ================================================ FILE: benchmark/react/bench_other.py ================================================ import argparse import json import time from concurrent.futures import ThreadPoolExecutor from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import dump_state_text, read_jsonl def get_prompt(question): prompt = ( """Solve a question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types: (1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search. (2) Lookup[keyword], which returns the next sentence containing keyword in the current passage. (3) Finish[answer], which returns the answer and finishes the task. Here are some examples. Question: What is the elevation range for the area that the eastern sector of the Colorado orogeny extends into? Thought 1: I need to search Colorado orogeny, find the area that the eastern sector of the Colorado orogeny extends into, then find the elevation range of the area. Action 1: Search[Colorado orogeny] Observation 1: The Colorado orogeny was an episode of mountain building (an orogeny) in Colorado and surrounding areas. Thought 2: It does not mention the eastern sector. So I need to look up eastern sector. Action 2: Lookup[eastern sector] Observation 2: (Result 1 / 1) The eastern sector extends into the High Plains and is called the Central Plains orogeny. Thought 3: The eastern sector of Colorado orogeny extends into the High Plains. So I need to search High Plains and find its elevation range. Action 3: Search[High Plains] Observation 3: High Plains refers to one of two distinct land regions: Thought 4: I need to instead search High Plains (United States). Action 4: Search[High Plains (United States)] Observation 4: The High Plains are a subregion of the Great Plains. From east to west, the High Plains rise in elevation from around 1,800 to 7,000 ft (550 to 2,130 m).[3] Thought 5: High Plains rise in elevation from around 1,800 to 7,000 ft, so the answer is 1,800 to 7,000 ft. Action 5: Finish[1,800 to 7,000 ft] Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who? Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after. Action 1: Search[Milhouse] Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening. Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after". Action 2: Lookup[named after] Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous. Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon. Action 3: Finish[Richard Nixon] Question: Which documentary is about Finnish rock groups, Adam Clayton Powell or The Saimaa Gesture? Thought 1: I need to search Adam Clayton Powell and The Saimaa Gesture, and find which documentary is about Finnish rock groups. Action 1: Search[Adam Clayton Powell] Observation 1: Could not find [Adam Clayton Powell]. Similar: ['Adam Clayton Powell III', 'Seventh Avenue (Manhattan)', 'Adam Clayton Powell Jr. State Office Building', 'Isabel Washington Powell', 'Adam Powell', 'Adam Clayton Powell (film)', 'Giancarlo Esposito']. Thought 2: To find the documentary, I can search Adam Clayton Powell (film). Action 2: Search[Adam Clayton Powell (film)] Observation 2: Adam Clayton Powell is a 1989 American documentary film directed by Richard Kilberg. The film is about the rise and fall of influential African-American politician Adam Clayton Powell Jr.[3][4] It was later aired as part of the PBS series The American Experience. Thought 3: Adam Clayton Powell (film) is a documentary about an African-American politician, not Finnish rock groups. So the documentary about Finnish rock groups must instead be The Saimaa Gesture. Action 3: Finish[The Saimaa Gesture] Question: What profession does Nicholas Ray and Elia Kazan have in common? Thought 1: I need to search Nicholas Ray and Elia Kazan, find their professions, then find the profession they have in common. Action 1: Search[Nicholas Ray] Observation 1: Nicholas Ray (born Raymond Nicholas Kienzle Jr., August 7, 1911 – June 16, 1979) was an American film director, screenwriter, and actor best known for the 1955 film Rebel Without a Cause. Thought 2: Professions of Nicholas Ray are director, screenwriter, and actor. I need to search Elia Kazan next and find his professions. Action 2: Search[Elia Kazan] Observation 2: Elia Kazan was an American film and theatre director, producer, screenwriter and actor. Thought 3: Professions of Elia Kazan are director, producer, screenwriter, and actor. So profession Nicholas Ray and Elia Kazan have in common is director, screenwriter, and actor. Action 3: Finish[director, screenwriter, actor] Question: Which magazine was started first Arthur's Magazine or First for Women? Thought 1: I need to search Arthur's Magazine and First for Women, and find which was started first. Action 1: Search[Arthur's Magazine] Observation 1: Arthur's Magazine (1844-1846) was an American literary periodical published in Philadelphia in the 19th century. Thought 2: Arthur's Magazine was started in 1844. I need to search First for Women next. Action 2: Search[First for Women] Observation 2: First for Women is a woman's magazine published by Bauer Media Group in the USA.[1] The magazine was started in 1989. Thought 3: First for Women was started in 1989. 1844 (Arthur's Magazine) < 1989 (First for Women), so Arthur's Magazine was started first. Action 3: Finish[Arthur's Magazine] Question: Were Pavel Urysohn and Leonid Levin known for the same type of work? Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same. Action 1: Search[Pavel Urysohn] Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory. Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work. Action 2: Search[Leonid Levin] Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist. Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work. Action 3: Finish[yes] """ + question ) return prompt def main(args): lines = read_jsonl(args.data_path)[: args.num_questions] arguments = [{"question": k, "triplets": v} for l in lines for k, v in l.items()] states = [] # Select backend call_generate = get_call_generate(args) def run_single_agent(argument): question = argument["question"] triplets = argument["triplets"] prompt = get_prompt(question) for i in range(1, len(triplets) + 2): prompt += "Thought " + str(i) + ":" states.append(prompt) answer = call_generate( prompt, max_tokens=200, temperature=0, stop="Observation" ) if i > len(triplets): break prompt += ( triplets[i - 1]["thought"] + "\nAction " + str(i) + ":" + triplets[i - 1]["action"] + "\nObservation " + str(i) + ":" + triplets[i - 1]["observation"] + "\n" ) states.append(answer) async def run_single_agent_async(argument): question = argument["question"] triplets = argument["triplets"] prompt = get_prompt(question) for i in range(1, len(triplets) + 2): prompt += "Thought " + str(i) + ":" states.append(prompt) answer = await call_generate( prompt, max_tokens=200, temperature=0, stop="Observation", max_len=4096 ) if i > len(triplets): break prompt += ( triplets[i - 1]["thought"] + "\nAction " + str(i) + ":" + triplets[i - 1]["action"] + "\nObservation " + str(i) + ":" + triplets[i - 1]["observation"] + "\n" ) states.append(answer) tic = time.perf_counter() if args.backend != "lmql": if args.parallel == 1: for arg in tqdm(arguments): run_single_agent(arg) else: with ThreadPoolExecutor(args.parallel) as executor: list( tqdm( executor.map(run_single_agent, arguments), total=len(arguments) ) ) else: import asyncio loop = asyncio.get_event_loop() batches = [ [] for _ in range((len(arguments) + args.parallel - 1) // args.parallel) ] for i, arg in enumerate(arguments): batches[i // args.parallel].append(arg) for bt in tqdm(batches): tasks = [run_single_agent_async(arg) for arg in bt] loop.run_until_complete(asyncio.gather(*tasks)) latency = time.perf_counter() - tic print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "ReAct Agents", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": len(arguments), "other": { "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="hotpotqa_100.jsonl") parser.add_argument("--num-questions", type=int, default=10) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/react/bench_sglang.py ================================================ import argparse import json import time import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text, read_jsonl @sgl.function def webthink(s, question, triplets): s += ( """Solve a question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types: (1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search. (2) Lookup[keyword], which returns the next sentence containing keyword in the current passage. (3) Finish[answer], which returns the answer and finishes the task. Here are some examples. Question: What is the elevation range for the area that the eastern sector of the Colorado orogeny extends into? Thought 1: I need to search Colorado orogeny, find the area that the eastern sector of the Colorado orogeny extends into, then find the elevation range of the area. Action 1: Search[Colorado orogeny] Observation 1: The Colorado orogeny was an episode of mountain building (an orogeny) in Colorado and surrounding areas. Thought 2: It does not mention the eastern sector. So I need to look up eastern sector. Action 2: Lookup[eastern sector] Observation 2: (Result 1 / 1) The eastern sector extends into the High Plains and is called the Central Plains orogeny. Thought 3: The eastern sector of Colorado orogeny extends into the High Plains. So I need to search High Plains and find its elevation range. Action 3: Search[High Plains] Observation 3: High Plains refers to one of two distinct land regions: Thought 4: I need to instead search High Plains (United States). Action 4: Search[High Plains (United States)] Observation 4: The High Plains are a subregion of the Great Plains. From east to west, the High Plains rise in elevation from around 1,800 to 7,000 ft (550 to 2,130 m).[3] Thought 5: High Plains rise in elevation from around 1,800 to 7,000 ft, so the answer is 1,800 to 7,000 ft. Action 5: Finish[1,800 to 7,000 ft] Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who? Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after. Action 1: Search[Milhouse] Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening. Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after". Action 2: Lookup[named after] Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous. Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon. Action 3: Finish[Richard Nixon] Question: Which documentary is about Finnish rock groups, Adam Clayton Powell or The Saimaa Gesture? Thought 1: I need to search Adam Clayton Powell and The Saimaa Gesture, and find which documentary is about Finnish rock groups. Action 1: Search[Adam Clayton Powell] Observation 1: Could not find [Adam Clayton Powell]. Similar: ['Adam Clayton Powell III', 'Seventh Avenue (Manhattan)', 'Adam Clayton Powell Jr. State Office Building', 'Isabel Washington Powell', 'Adam Powell', 'Adam Clayton Powell (film)', 'Giancarlo Esposito']. Thought 2: To find the documentary, I can search Adam Clayton Powell (film). Action 2: Search[Adam Clayton Powell (film)] Observation 2: Adam Clayton Powell is a 1989 American documentary film directed by Richard Kilberg. The film is about the rise and fall of influential African-American politician Adam Clayton Powell Jr.[3][4] It was later aired as part of the PBS series The American Experience. Thought 3: Adam Clayton Powell (film) is a documentary about an African-American politician, not Finnish rock groups. So the documentary about Finnish rock groups must instead be The Saimaa Gesture. Action 3: Finish[The Saimaa Gesture] Question: What profession does Nicholas Ray and Elia Kazan have in common? Thought 1: I need to search Nicholas Ray and Elia Kazan, find their professions, then find the profession they have in common. Action 1: Search[Nicholas Ray] Observation 1: Nicholas Ray (born Raymond Nicholas Kienzle Jr., August 7, 1911 – June 16, 1979) was an American film director, screenwriter, and actor best known for the 1955 film Rebel Without a Cause. Thought 2: Professions of Nicholas Ray are director, screenwriter, and actor. I need to search Elia Kazan next and find his professions. Action 2: Search[Elia Kazan] Observation 2: Elia Kazan was an American film and theatre director, producer, screenwriter and actor. Thought 3: Professions of Elia Kazan are director, producer, screenwriter, and actor. So profession Nicholas Ray and Elia Kazan have in common is director, screenwriter, and actor. Action 3: Finish[director, screenwriter, actor] Question: Which magazine was started first Arthur's Magazine or First for Women? Thought 1: I need to search Arthur's Magazine and First for Women, and find which was started first. Action 1: Search[Arthur's Magazine] Observation 1: Arthur's Magazine (1844-1846) was an American literary periodical published in Philadelphia in the 19th century. Thought 2: Arthur's Magazine was started in 1844. I need to search First for Women next. Action 2: Search[First for Women] Observation 2: First for Women is a woman's magazine published by Bauer Media Group in the USA.[1] The magazine was started in 1989. Thought 3: First for Women was started in 1989. 1844 (Arthur's Magazine) < 1989 (First for Women), so Arthur's Magazine was started first. Action 3: Finish[Arthur's Magazine] Question: Were Pavel Urysohn and Leonid Levin known for the same type of work? Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same. Action 1: Search[Pavel Urysohn] Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory. Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work. Action 2: Search[Leonid Levin] Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist. Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work. Action 3: Finish[yes] """ + question ) for i in range(1, len(triplets) + 2): s += "Thought " + str(i) + ":" # NOTE: This is an implementation for replaying a given trace for benchmark purposes. It is not an actual ReAct agent implementation. ss = s.fork(1) ss[0] += sgl.gen(name="thought_action", max_tokens=200, stop="Observation") ss.join() # to verify the correctness of output, this should be collected # print(ss[0]["thought_action"]) if i > len(triplets): break s += ( triplets[i - 1]["thought"] + "\nAction " + str(i) + ":" + triplets[i - 1]["action"] + "\nObservation " + str(i) + ":" + triplets[i - 1]["observation"] + "\n" ) def main(args): lines = read_jsonl(args.data_path)[: args.num_questions] arguments = [{"question": k, "triplets": v} for l in lines for k, v in l.items()] # Select backend backend = select_sglang_backend(args) sgl.set_default_backend(backend) states = [] tic = time.perf_counter() states = webthink.run_batch( arguments, temperature=0, num_threads=args.parallel, progress_bar=True, ) latency = time.perf_counter() - tic # Compute accuracy print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "ReAct Agents", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": len(arguments), "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="hotpotqa_100.jsonl") parser.add_argument("--num-questions", type=int, default=10) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/reasoning_benchmark/README.md ================================================ # Run benchmark This benchmark is primarily intended to be used with reasoning models like `DeepSeek-R1` and its distilled models like `DeepSeek-R1-Distill-Qwen-1.5B`. Please use ```bash pip install antlr4-python3-runtime ``` for `parse_latex` which we use for symbolic equality check. ## Benchmark sglang 1. Launch the Server ```bash python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --port 30000 ``` Note that depending on the GPU this benchmark will take quiet some time. To employ data parallelism please use: ```bash python3 -m sglang_router.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --port 30000 --dp-size 4 ``` 2. Benchmarking We use [suggested](https://github.com/deepseek-ai/DeepSeek-R1) parameters of `temperature=0.6`, `top_p=.95`, `max_new_tokens=32768`. The command line argument `num-tries` can be used to evaluate the model multiple times on the same question. We use the suggested `64` from the repo for AIME 2024. For LIMO, we use `8` as the number of tries due to the size of the dataset. By default evaluate on LIMO dataset. ```bash python3 bench_sglang.py --parallel 256 --num-tries 64 --port 30000 ``` Evaluate on AIME 2024 dataset. ```bash python3 bench_sglang.py --parallel 256 --port 30000 --data-path Maxwell-Jia/AIME_2024 --question-key Problem --answer-key Answer --num-tries 64 ``` Evaluate on [AIME 2025 I dataset](https://huggingface.co/datasets/opencompass/AIME2025). For benchmark result see [here](https://matharena.ai/). ```bash python3 bench_sglang.py --parallel 256 --port 30000 --data-path opencompass/AIME2025 --question-key question --answer-key answer --num-tries 64 ``` ## Results ### Evaluation Results | Dataset | Num Tries | Accuracy | Reference | Standard Error | |------------|-----------|----------|-----------|-----------| | LIMO | 8 | 47.7% | ? | ? | | AIME 2024 | 64 | 33.2% | 28.9% | 3.4% | | AIME 2025 I| 64 | 29.9% | 25.0% | ? | ### Statistic Analysis Results Set up SGLang engine for statistic analysis, for high efficiency we use `--dp-size 8` for data parallelism: ```bash python3 -m sglang_router.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --port 30000 --dp-size 8 ``` **Experiment 1**: We fixed the number of attempts (num_tries) and conducted multiple runs to assess the consistency of the model's performance. The results show that all recorded accuracies lie within ± one standard error deviation from the mean. This suggests that **our metric serves as an effective upper bound for the deviation of reported accuracy**. To collect the accuracy, run the following command 30 times: ```bash python3 bench_sglang.py --parallel 64 --port 30000 --data-path Maxwell-Jia/AIME_2024 --question-key Problem --answer-key Answer --num-tries 64 ``` ![acc_hist](figure/Acc_histplot.png) **Experiment 2**: We explored the relationship between the number of attempts (num_tries) and the standard error (SE) by varying num_tries across a range (e.g., 8, 16, 32, ..., 256) and performing a single run for each value. The results demonstrate that as the number of attempts increases, the standard error decreases, leading to **greater stability in answer accuracy**. To reveal the relationship, run the command 6 times and adjust the parameter `--num-tries` for each run: ```bash python3 bench_sglang.py --parallel 64 --port 30000 --data-path Maxwell-Jia/AIME_2024 --question-key Problem --answer-key Answer --num-tries ``` ![SE_num_tries](figure/SE_numtries.png) ================================================ FILE: benchmark/reasoning_benchmark/answer_extraction.py ================================================ # Adapted from https://github.com/deepseek-ai/DeepSeek-Math/blob/main/evaluation/data_processing/answer_extraction.py import re import regex def _fix_fracs(string): substrs = string.split("\\frac") new_str = substrs[0] if len(substrs) > 1: substrs = substrs[1:] for substr in substrs: new_str += "\\frac" if len(substr) > 0 and substr[0] == "{": new_str += substr else: try: assert len(substr) >= 2 except: return string a = substr[0] b = substr[1] if b != "{": if len(substr) > 2: post_substr = substr[2:] new_str += "{" + a + "}{" + b + "}" + post_substr else: new_str += "{" + a + "}{" + b + "}" else: if len(substr) > 2: post_substr = substr[2:] new_str += "{" + a + "}" + b + post_substr else: new_str += "{" + a + "}" + b string = new_str return string def _fix_a_slash_b(string): if len(string.split("/")) != 2: return string a = string.split("/")[0] b = string.split("/")[1] try: if "sqrt" not in a: a = int(a) if "sqrt" not in b: b = int(b) assert string == "{}/{}".format(a, b) new_string = "\\frac{" + str(a) + "}{" + str(b) + "}" return new_string except: return string def _fix_sqrt(string): _string = re.sub(r"\\sqrt(-?[0-9.a-zA-Z]+)", r"\\sqrt{\1}", string) _string = re.sub(r"\\sqrt\s+(\w+)$", r"\\sqrt{\1}", _string) return _string def _fix_tan(string): _string = re.sub(r"\\tan(-?[0-9.a-zA-Z]+)", r"\\tan{\1}", string) _string = re.sub(r"\\tan\s+(\w+)$", r"\\tan{\1}", _string) return _string def strip_string(string): string = str(string).strip() # linebreaks string = string.replace("\n", "") # right "." string = string.rstrip(".") # remove inverse spaces string = string.replace("\\!", "") # string = string.replace("\\ ", "") # replace \\ with \ # string = string.replace("\\\\", "\\") # string = string.replace("\\\\", "\\") if string.startswith("\\text{") and string.endswith("}"): string = string.split("{", 1)[1][:-1] # replace tfrac and dfrac with frac string = string.replace("tfrac", "frac") string = string.replace("dfrac", "frac") string = string.replace("cfrac", "frac") # remove \left and \right string = string.replace("\\left", "") string = string.replace("\\right", "") # Remove unit: miles, dollars if after is not none _string = re.sub(r"\\text{.*?}$", "", string).strip() if _string != "" and _string != string: # print("Warning: unit not removed: '{}' -> '{}'".format(string, _string)) string = _string # Remove circ (degrees) string = string.replace("^{\\circ}", "").strip() string = string.replace("^\\circ", "").strip() string = regex.sub(r"\{(c|m)?m\}(\^(2|3))?", "", string).strip() string = regex.sub(r"p\.m\.$", "", string).strip() string = regex.sub(r"(\d)\s*t$", r"\1", string).strip() # remove dollar signs string = string.replace("\\$", "") string = string.replace("$", "") # string = string.replace("\\text", "") string = string.replace("x\\in", "") # remove percentage string = string.replace("\\%", "%") string = string.replace("\%", "%") # string = string.replace("%", "") # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string string = string.replace(" .", " 0.") string = string.replace("{.", "{0.") # cdot string = string.replace("\\cdot", "") # inf string = string.replace("infinity", "\\infty") if "\\infty" not in string: string = string.replace("inf", "\\infty") string = string.replace("+\\inity", "\\infty") # and # string = string.replace("and", "") string = string.replace("\\mathbf", "") string = string.replace("\\mathrm", "") # use regex to remove \mbox{...} string = re.sub(r"\\mbox{.*?}", "", string) # quote string.replace("'", "") string.replace('"', "") # i, j if "j" in string and "i" not in string: string = string.replace("j", "i") # replace a.000b where b is not number or b is end, with ab, use regex string = re.sub(r"(\d+)\.0+([^\d])", r"\1\2", string) string = re.sub(r"(\d+)\.0+$", r"\1", string) # if empty, return empty string if len(string) == 0: return string if string[0] == ".": string = "0" + string # to consider: get rid of e.g. "k = " or "q = " at beginning # if len(string.split("=")) == 2: # if len(string.split("=")[0]) <= 2: # string = string.split("=")[1] string = _fix_sqrt(string) string = _fix_tan(string) string = string.replace(" ", "") # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b} string = _fix_fracs(string) # NOTE: X/Y changed to \frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y string = _fix_a_slash_b(string) string = regex.sub(r"(\\|,|\.)+$", "", string) return string def extract_boxed_answers(text): answers = [] for piece in text.split("boxed{")[1:]: n = 0 for i in range(len(piece)): if piece[i] == "{": n += 1 elif piece[i] == "}": n -= 1 if n < 0: if i + 1 < len(piece) and piece[i + 1] == "%": answers.append(piece[: i + 1]) else: answers.append(piece[:i]) break return answers def extract_program_output(pred_str): """ extract output between the last ```output\n...\n``` """ if "```output" not in pred_str: return "" if "```output" in pred_str: pred_str = pred_str.split("```output")[-1] if "```" in pred_str: pred_str = pred_str.split("```")[0] output = pred_str.strip() return output def extract_answer(pred_str, exhaust=False): pred = [] if "final answer is $" in pred_str and "$. I hope" in pred_str: tmp = pred_str.split("final answer is $", 1)[1] pred = [tmp.split("$. I hope", 1)[0].strip()] elif "boxed" in pred_str: pred = extract_boxed_answers(pred_str) elif "he answer is" in pred_str: pred = [pred_str.split("he answer is")[-1].strip()] else: program_output = extract_program_output(pred_str) if program_output != "": # fall back to program pred.append(program_output) else: # use the last number pattern = "-?\d*\.?\d+" answers = re.findall(pattern, pred_str.replace(",", "")) if len(answers) >= 1: last_ans = answers[-1] else: last_ans = "" if last_ans: pred.append(last_ans) # multiple line _pred = [] for each_ans in pred: each_ans = each_ans.strip().split("\n")[0] each_ans = each_ans.lstrip(":") each_ans = each_ans.rstrip(".") each_ans = each_ans.rstrip("/") each_ans = strip_string(each_ans) _pred.append(each_ans) if exhaust: return _pred else: return _pred[-1] if _pred else "" def extract_math_answer(question, reasoning, task): answer = [] for ans in extract_answer(reasoning, exhaust=True): if "separated by commas" in question and all(ch not in ans for ch in "()[]"): answer.extend([a.strip() for a in ans.split(",")]) elif regex.search(r"\\text\{\s*and\s*\}", ans): answer.extend( [ a.strip() for a in regex.sub(r"\\text\{\s*and\s*\}", "[SEP]", ans).split( "[SEP]" ) ] ) else: answer.append(ans.strip()) return answer ================================================ FILE: benchmark/reasoning_benchmark/bench_sglang.py ================================================ import argparse import json import time import answer_extraction import eval_utils import numpy as np from datasets import load_dataset import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text @sgl.function def reasoning_gen(s, question: str): s += sgl.user( question + "\nPlease reason step by step, and put your final answer within \boxed{}." ) s += sgl.assistant( sgl.gen( "answer", ) ) def convert_dataset(path: str, question_key: str, answer_key: str, num_tries: int): raw_dataset = load_dataset(path) questions = [] answers = [] for data in raw_dataset["train"]: question = data[question_key] answer = data[answer_key] for _ in range(num_tries): questions.append({"question": question}) answers.append({"answer": answer}) return questions, answers def main(args): # Select backend sgl.set_default_backend(select_sglang_backend(args)) # Get dataset questions, answers = convert_dataset( args.data_path, args.question_key, args.answer_key, args.num_tries ) # Run requests tic = time.perf_counter() states = reasoning_gen.run_batch( questions, num_threads=args.parallel, progress_bar=True, temperature=0.6, max_new_tokens=32768, top_p=0.95, ) latency = time.perf_counter() - tic # Extract results and record outcomes in a list. outcomes = [] for i, state in enumerate(states): try: pred_answer = answer_extraction.extract_math_answer( questions[i]["question"], state["answer"], "limo" ) gt_answer = str(answers[i]["answer"]) pred_answer = ( pred_answer[-1] if isinstance(pred_answer, list) else pred_answer ) is_correct = 1 if eval_utils.math_equal(pred_answer, gt_answer) else 0 except Exception as e: print(f"Error extracting answer: {e}") is_correct = 0 outcomes.append(is_correct) # Calculate overall accuracy using numpy overall_accuracy = np.mean(outcomes) print(f"Overall Accuracy: {overall_accuracy}") # Calculate mean standard error over questions if num_tries >= 2 if args.num_tries > 1: outcomes_np = np.array(outcomes).reshape(-1, args.num_tries) # Using sample standard deviation with ddof=1 std_per_question = np.std(outcomes_np, axis=1, ddof=1) # Compute the standard error for each question: std / sqrt(num_tries) se_per_question = std_per_question / np.sqrt(args.num_tries) mean_se = se_per_question.mean() print(f"Mean Standard Error of Accuracy across questions: {mean_se}") else: mean_se = None print("Not enough samples per question to compute standard error.") # Calculate output throughput num_output_tokens = sum( s.get_meta_info("answer")["completion_tokens"] for s in states ) output_throughput = num_output_tokens / latency print(f"Output throughput: {output_throughput} token/s") # Dump results dump_state_text(f"tmp_output_{args.backend}.txt", states) # Write results with open(args.result_file, "a") as fout: value = { "task": "limo", "backend": args.backend, "latency": round(latency, 3), "overall_accuracy": round(overall_accuracy, 3), "mean_se_accuracy": round(mean_se, 3) if mean_se is not None else None, "num_requests": len(questions), "other": { "num_questions": len(questions), "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="GAIR/LIMO") parser.add_argument("--question-key", type=str, default="question") parser.add_argument("--answer-key", type=str, default="answer") parser.add_argument("--num-tries", type=int, default=1) add_common_sglang_args_and_parse(parser) args = parser.parse_args() main(args) ================================================ FILE: benchmark/reasoning_benchmark/eval_utils.py ================================================ # Adapted from https://github.com/deepseek-ai/DeepSeek-Math/blob/main/evaluation/eval/eval_utils.py from math import isclose import regex from sympy import N, simplify from sympy.parsing.latex import parse_latex from sympy.parsing.sympy_parser import parse_expr def parse_digits(num): # format: 234.23 || 23% num = regex.sub(",", "", str(num)) try: return float(num) except: if num.endswith("%"): num = num[:-1] if num.endswith("\\"): num = num[:-1] try: return float(num) / 100 except: pass return None def is_digit(num): # paired with parse_digits return parse_digits(num) is not None def symbolic_equal(a, b): def _parse(s): for f in [parse_latex, parse_expr]: try: return f(s) except: pass return s a = _parse(a) b = _parse(b) try: if simplify(a - b) == 0: return True except: pass try: if isclose(N(a), N(b), abs_tol=1e-3): return True except: pass return False def math_equal(prediction, reference, include_percentage=True, is_close=True): """ Exact match of math if and only if: 1. numerical equal: both can convert to float and are equal 2. symbolic equal: both can convert to sympy expression and are equal """ if str(prediction) == str(reference): return True try: # 1. numerical equal if is_digit(prediction) and is_digit(reference): prediction = parse_digits(prediction) reference = parse_digits(reference) # number questions if include_percentage: gt_result = [reference / 100, reference, reference * 100] else: gt_result = [reference] for item in gt_result: try: if is_close: if isclose(item, prediction, abs_tol=1e-3): return True else: if item == prediction: return True except Exception: continue return False except: pass if not prediction and prediction not in [0, False]: return False # 2. symbolic equal reference = str(reference).strip() prediction = str(prediction).strip() if ( regex.match(r"(\(|\[).+(\)|\])", prediction) is not None and regex.match(r"(\(|\[).+(\)|\])", reference) is not None ): pred_parts = prediction[1:-1].split(",") ref_parts = reference[1:-1].split(",") if len(pred_parts) == len(ref_parts): if all( [ math_equal( pred_parts[i], ref_parts[i], include_percentage, is_close ) for i in range(len(pred_parts)) ] ): return True # Add back matrix comparison if ( ( prediction.startswith("\\begin{pmatrix}") or prediction.startswith("\\begin{bmatrix}") ) and ( prediction.endswith("\\end{pmatrix}") or prediction.endswith("\\end{bmatrix}") ) and ( reference.startswith("\\begin{pmatrix}") or reference.startswith("\\begin{bmatrix}") ) and ( reference.endswith("\\end{pmatrix}") or reference.endswith("\\end{bmatrix}") ) ): pred_lines = [ line.strip() for line in prediction[ len("\\begin{pmatrix}") : -len("\\end{pmatrix}") ].split("\\\\") if line.strip() ] ref_lines = [ line.strip() for line in reference[ len("\\begin{pmatrix}") : -len("\\end{pmatrix}") ].split("\\\\") if line.strip() ] matched = True if len(pred_lines) == len(ref_lines): for pred_line, ref_line in zip(pred_lines, ref_lines): pred_parts = pred_line.split("&") ref_parts = ref_line.split("&") if len(pred_parts) == len(ref_parts): if not all( [ math_equal( pred_parts[i], ref_parts[i], include_percentage, is_close, ) for i in range(len(pred_parts)) ] ): matched = False break else: matched = False if not matched: break else: matched = False if matched: return True # Add back equation comparison if prediction.count("=") == 1 and reference.count("=") == 1: pred = prediction.split("=") pred = f"{pred[0].strip()} - ({pred[1].strip()})" ref = reference.split("=") ref = f"{ref[0].strip()} - ({ref[1].strip()})" if symbolic_equal(pred, ref) or symbolic_equal(f"-({pred})", ref): return True elif ( prediction.count("=") == 1 and len(prediction.split("=")[0].strip()) <= 2 and "=" not in reference ): if math_equal( prediction.split("=")[1], reference, include_percentage, is_close ): return True elif ( reference.count("=") == 1 and len(reference.split("=")[0].strip()) <= 2 and "=" not in prediction ): if math_equal( prediction, reference.split("=")[1], include_percentage, is_close ): return True # symbolic equal with sympy if symbolic_equal(prediction, reference): return True return False ================================================ FILE: benchmark/tip_suggestion/.gitignore ================================================ !topic.jsonl ================================================ FILE: benchmark/tip_suggestion/README.md ================================================ ## Run benchmark ### Benchmark sglang ``` python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` ``` python3 bench_sglang.py --num-questions 64 python3 bench_sglang.py --num-questions 32 --parallel 1 ``` ### Benchmark vllm ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` ``` python3 bench_other.py --backend vllm --num-questions 64 ``` ### Benchmark guidance ``` python3 bench_other.py --backend guidance --num-questions 32 --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` ### Benchmark lmql ``` python3 bench_other.py --backend lmql --num-questions 32 --parallel 1 ``` ================================================ FILE: benchmark/tip_suggestion/bench_other.py ================================================ import argparse import json import time from concurrent.futures import ThreadPoolExecutor from functools import partial from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import dump_state_text, read_jsonl number = 5 def expand_tip(topic, tip, generate): s = """Please expand a tip for a topic into a detailed paragraph. Topic: staying healthy Tip: Regular Exercise Paragraph: Incorporate physical activity into your daily routine. This doesn't necessarily mean intense gym workouts; it can be as simple as walking, cycling, or yoga. Regular exercise helps in maintaining a healthy weight, improves cardiovascular health, boosts mental health, and can enhance cognitive function, which is crucial for fields that require intense intellectual engagement. Topic: building a campfire Tip: Choose the Right Location Paragraph: Always build your campfire in a safe spot. This means selecting a location that's away from trees, bushes, and other flammable materials. Ideally, use a fire ring if available. If you're building a fire pit, it should be on bare soil or on a bed of stones, not on grass or near roots which can catch fire underground. Make sure the area above is clear of low-hanging branches. Topic: writing a blog post Tip: structure your content effectively Paragraph: A well-structured post is easier to read and more enjoyable. Start with an engaging introduction that hooks the reader and clearly states the purpose of your post. Use headings and subheadings to break up the text and guide readers through your content. Bullet points and numbered lists can make information more digestible. Ensure each paragraph flows logically into the next, and conclude with a summary or call-to-action that encourages reader engagement. Topic: """ + topic + "\nTip: " + tip + "\nParagraph:" return generate(s, max_tokens=128, stop=["\n\n"]) def suggest_tips(topic, generate): s = "Please act as a helpful assistant. Your job is to provide users with useful tips on a specific topic.\n" s += "USER: Give some tips for " + topic + ".\n" s += ( "ASSISTANT: Okay. Here are " + str(number) + " concise tips, each under 8 words:\n" ) tips = [] for i in range(1, 1 + number): s += f"{i}." tip = generate(s, max_tokens=24, stop=[".", "\n"]) s += tip + ".\n" tips.append(tip) paragraphs = [expand_tip(topic, tip, generate=generate) for tip in tips] for i in range(1, 1 + number): s += f"Tip {i}:" + paragraphs[i - 1] + "\n" return s def main(args): lines = read_jsonl(args.data_path)[: args.num_questions] states = [None] * len(lines) # Select backend call_generate = partial(get_call_generate(args), temperature=0) # Run requests tic = time.perf_counter() if args.backend != "lmql": def get_one_answer(i): states[i] = suggest_tips(lines[i]["topic"], call_generate) if args.parallel == 1: for i in tqdm(range(len(lines))): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: list( tqdm( executor.map(get_one_answer, list(range(len(lines)))), total=len(lines), ) ) else: import asyncio from lmql_funcs import suggest_tips_async async def get_one_answer_async(i): states[i] = await suggest_tips_async(lines[i]["topic"], call_generate) batches = [] for i in range(0, len(lines), args.parallel): batches.append(list(range(i, min(i + args.parallel, len(lines))))) loop = asyncio.get_event_loop() for batch in tqdm(batches): loop.run_until_complete( asyncio.gather(*[get_one_answer_async(i) for i in batch]) ) latency = time.perf_counter() - tic # Compute accuracy print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "tip_suggestion", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="topic.jsonl") parser.add_argument("--num-questions", type=int, default=100) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/tip_suggestion/bench_sglang.py ================================================ import argparse import json import time import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text, read_jsonl number = 5 @sgl.function def expand_tip(s, topic, tip): s += """Please expand a tip for a topic into a detailed paragraph. Topic: staying healthy Tip: Regular Exercise Paragraph: Incorporate physical activity into your daily routine. This doesn't necessarily mean intense gym workouts; it can be as simple as walking, cycling, or yoga. Regular exercise helps in maintaining a healthy weight, improves cardiovascular health, boosts mental health, and can enhance cognitive function, which is crucial for fields that require intense intellectual engagement. Topic: building a campfire Tip: Choose the Right Location Paragraph: Always build your campfire in a safe spot. This means selecting a location that's away from trees, bushes, and other flammable materials. Ideally, use a fire ring if available. If you're building a fire pit, it should be on bare soil or on a bed of stones, not on grass or near roots which can catch fire underground. Make sure the area above is clear of low-hanging branches. Topic: writing a blog post Tip: structure your content effectively Paragraph: A well-structured post is easier to read and more enjoyable. Start with an engaging introduction that hooks the reader and clearly states the purpose of your post. Use headings and subheadings to break up the text and guide readers through your content. Bullet points and numbered lists can make information more digestible. Ensure each paragraph flows logically into the next, and conclude with a summary or call-to-action that encourages reader engagement. Topic: """ + topic + "\nTip: " + tip + "\nParagraph:" s += sgl.gen("paragraph", max_tokens=128, stop=["\n\n"], temperature=0) @sgl.function def suggest_tips(s, topic): s += "Please act as a helpful assistant. Your job is to provide users with useful tips on a specific topic.\n" s += "USER: Give some tips for " + topic + ".\n" s += ( "ASSISTANT: Okay. Here are " + str(number) + " concise tips, each under 8 words:\n" ) paragraphs = [] for i in range(1, 1 + number): s += f"{i}." + sgl.gen(f"tip_{i}", max_tokens=24, stop=[".", "\n"]) + ".\n" paragraphs.append(expand_tip(topic=topic, tip=s[f"tip_{i}"])) for i in range(1, 1 + number): s += f"Tip {i}:" + paragraphs[i - 1]["paragraph"] + "\n" def main(args): lines = read_jsonl(args.data_path)[: args.num_questions] arguments = [{"topic": l["topic"]} for l in lines] # Select backend sgl.set_default_backend(select_sglang_backend(args)) # Run requests tic = time.perf_counter() states = suggest_tips.run_batch( arguments, temperature=0, num_threads=args.parallel, progress_bar=True ) latency = time.perf_counter() - tic # Compute accuracy print(f"Latency: {latency:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", states) with open(args.result_file, "a") as fout: value = { "task": "tip_suggestion", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="topic.jsonl") parser.add_argument("--num-questions", type=int, default=100) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/tip_suggestion/lmql_funcs.py ================================================ number = 5 async def expand_tip_async(topic, tip, generate): s = """Please expand a tip for a topic into a detailed paragraph. Topic: staying healthy Tip: Regular Exercise Paragraph: Incorporate physical activity into your daily routine. This doesn't necessarily mean intense gym workouts; it can be as simple as walking, cycling, or yoga. Regular exercise helps in maintaining a healthy weight, improves cardiovascular health, boosts mental health, and can enhance cognitive function, which is crucial for fields that require intense intellectual engagement. Topic: building a campfire Tip: Choose the Right Location Paragraph: Always build your campfire in a safe spot. This means selecting a location that's away from trees, bushes, and other flammable materials. Ideally, use a fire ring if available. If you're building a fire pit, it should be on bare soil or on a bed of stones, not on grass or near roots which can catch fire underground. Make sure the area above is clear of low-hanging branches. Topic: writing a blog post Tip: structure your content effectively Paragraph: A well-structured post is easier to read and more enjoyable. Start with an engaging introduction that hooks the reader and clearly states the purpose of your post. Use headings and subheadings to break up the text and guide readers through your content. Bullet points and numbered lists can make information more digestible. Ensure each paragraph flows logically into the next, and conclude with a summary or call-to-action that encourages reader engagement. Topic: """ + topic + "\nTip: " + tip + "\nParagraph:" return await generate(s, max_tokens=128, stop="\n\n") async def suggest_tips_async(topic, generate): s = "Please act as a helpful assistant. Your job is to provide users with useful tips on a specific topic.\n" s += "USER: Give some tips for " + topic + ".\n" s += ( "ASSISTANT: Okay. Here are " + str(number) + " concise tips, each under 8 words:\n" ) tips = [] for i in range(1, 1 + number): s += f"{i}." # NOTE: stop is different due to lmql does not support a list of stop tokens tip = await generate(s, max_tokens=24, stop=".\n") s += tip + ".\n" tips.append(tip) paragraphs = [await expand_tip_async(topic, tip, generate=generate) for tip in tips] for i in range(1, 1 + number): s += f"Tip {i}:" + paragraphs[i - 1] + "\n" return s ================================================ FILE: benchmark/tip_suggestion/topic.jsonl ================================================ {"topic": "organizing a successful charity event", "number": 6} {"topic": "improving personal credit scores", "number": 7} {"topic": "staying motivated during job searches", "number": 5} {"topic": "maintaining a work-life balance", "number": 9} {"topic": "reducing carbon footprint at home", "number": 8} {"topic": "starting a book club", "number": 5} {"topic": "learning to play a musical instrument", "number": 7} {"topic": "getting into freelance writing", "number": 6} {"topic": "beginner yoga poses", "number": 8} {"topic": "preparing for graduate school exams", "number": 5} {"topic": "exploring minimalist living", "number": 9} {"topic": "effective grocery shopping", "number": 7} {"topic": "winter camping", "number": 5} {"topic": "starting a podcast on a budget", "number": 8} {"topic": "creating a capsule wardrobe", "number": 6} {"topic": "improving your writing skills", "number": 7} {"topic": "learning a new software quickly", "number": 9} {"topic": "reducing anxiety before public speaking", "number": 5} {"topic": "planning a solo travel adventure", "number": 8} {"topic": "beginner skateboarders", "number": 6} {"topic": "studying abroad", "number": 7} {"topic": "planting a vegetable garden", "number": 5} {"topic": "adopting a shelter pet", "number": 9} {"topic": "learning to cook ethnic cuisines", "number": 8} {"topic": "effective conflict resolution", "number": 5} {"topic": "starting a vlog", "number": 7} {"topic": "keeping a daily journal", "number": 6} {"topic": "improving sleep hygiene", "number": 8} {"topic": "beginner mountain climbers", "number": 5} {"topic": "creating a mobile app", "number": 9} {"topic": "maintaining a saltwater aquarium", "number": 7} {"topic": "preparing for a baby's arrival", "number": 6} {"topic": "writing a fantasy novel", "number": 5} {"topic": "effective team leadership", "number": 8} {"topic": "making a documentary film", "number": 9} {"topic": "learning about historical events", "number": 7} {"topic": "baking gluten-free treats", "number": 6} {"topic": "improving mental arithmetic skills", "number": 5} {"topic": "building a treehouse", "number": 8} {"topic": "getting started with watercolor painting", "number": 9} {"topic": "creating a YouTube tutorial series", "number": 7} {"topic": "landscape photography", "number": 5} {"topic": "navigating cultural differences", "number": 6} {"topic": "preparing for a marathon", "number": 8} {"topic": "building an online business", "number": 9} {"topic": "learning to dance at home", "number": 5} {"topic": "self-publishing a book", "number": 7} {"topic": "starting an urban farm", "number": 6} {"topic": "improving your memory", "number": 8} {"topic": "creating a personal brand online", "number": 9} ================================================ FILE: benchmark/tree_of_thought_deep/README.md ================================================ ## Download data ``` wget https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl ``` ## Run benchmark NOTE: This is an implementation for throughput/latency benchmark purposes. The prompts are not tuned to achieve good accuracy on the GSM-8K tasks. ### Benchmark sglang ``` python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` ``` python3 bench_sglang.py --num-questions 32 python3 bench_sglang.py --num-questions 16 --parallel 1 ``` ### Benchmark vllm ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` ``` python3 bench_other.py --num-questions 32 --backend vllm ``` ### Benchmark lightllm ``` # A10G python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 ``` ``` python3 bench_other.py --num-questions 32 --backend lightllm ``` ### Benchmark guidance ``` python3 bench_other.py --num-questions 8 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` ### Benchmark lmql ``` python3 bench_other.py --num-questions 8 --backend lmql --parallel 1 ``` ================================================ FILE: benchmark/tree_of_thought_deep/bench_other.py ================================================ import argparse import ast import json import re import time from collections import Counter from concurrent.futures import ThreadPoolExecutor import numpy as np from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import dump_state_text, read_jsonl INVALID = -9999999 def get_answer_value(answer_str): answer_str = answer_str.replace(",", "") numbers = re.findall(r"\d+", answer_str) if len(numbers) < 1: return INVALID try: return ast.literal_eval(numbers[-1]) except SyntaxError: return INVALID def most_frequent_number(numbers): if not numbers: return None frequency = Counter(numbers) most_frequent = max(frequency, key=frequency.get) return most_frequent USER_PREFIX = "[INST] " USER_SUFFIX = " [/INST]" ASSISTANT_PREFIX = "" ASSISTANT_SUFFIX = " " # Use a low temp to make the results more deterministic and the comparison more fair. temp = 0.001 def propose_plan(s, question, num_branches, call_generate): s += ( USER_PREFIX + """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + question + USER_SUFFIX ) s += ASSISTANT_PREFIX comps = call_generate( s, max_tokens=256, temperature=temp, stop=None, n=num_branches ) return [s + comp + ASSISTANT_SUFFIX for comp in comps] def execute_plan(s, num_branches, call_generate): s += ( USER_PREFIX + """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + USER_SUFFIX ) s += ASSISTANT_PREFIX comps = call_generate( s, max_tokens=256, temperature=temp, stop=None, n=num_branches ) return [s + comp + ASSISTANT_SUFFIX for comp in comps] def reflect_solution(s, num_branches, call_generate): s += ( USER_PREFIX + """Okay. Now, evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + USER_SUFFIX ) s += ASSISTANT_PREFIX comps = call_generate( s, max_tokens=256, temperature=temp, stop=None, n=num_branches ) return [s + comp + ASSISTANT_SUFFIX for comp in comps] def get_final_answer(s, num_branches, call_generate): s += ( USER_PREFIX + """Based on your reflection, do you change your mind? Now, give me the final answer after careful consideration.""" + USER_SUFFIX ) s += ASSISTANT_PREFIX comps = call_generate( s, max_tokens=256, temperature=temp, stop=None, n=num_branches ) return [s + comp + ASSISTANT_SUFFIX for comp in comps] def tree_search(question, num_branches, call_generate): plan_forks = propose_plan("", question, num_branches, call_generate) sol_states = [] for plan in plan_forks: forks = execute_plan(plan, num_branches, call_generate) sol_states.extend(forks) ref_states = [] for sol in sol_states: forks = reflect_solution(sol, num_branches, call_generate) ref_states.extend(forks) solutions = [] for sol in ref_states: ans = get_final_answer(sol, num_branches, call_generate) solutions.append(ans) return solutions def main(args): lines = read_jsonl(args.data_path) # Construct prompts num_branches = 2 questions = [] labels = [] for i in range(len(lines[: args.num_questions])): questions.append(lines[i]["question"]) labels.append(get_answer_value(lines[i]["answer"])) assert all(l != INVALID for l in labels) arguments = [{"question": q, "num_branches": num_branches} for q in questions] # Select backend call_generate = get_call_generate(args) # Run requests states = [None] * len(questions) tic = time.perf_counter() if args.backend != "lmql": def get_one_answer(i): states[i] = tree_search(**arguments[i], call_generate=call_generate) if args.parallel == 1: for i in tqdm(range(len(questions))): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: list( tqdm( executor.map(get_one_answer, list(range(len(questions)))), total=len(questions), ) ) else: import asyncio from lmql_funcs import tree_search_async async def get_one_answer_async(i): states[i] = await tree_search_async( **arguments[i], call_generate=call_generate ) batches = [ [] for _ in range((len(questions) + args.parallel - 1) // args.parallel) ] for i in range(len(questions)): batches[i // args.parallel].append(i) loop = asyncio.get_event_loop() for bt in tqdm(batches): tasks = [get_one_answer_async(k) for k in bt] loop.run_until_complete(asyncio.gather(*tasks)) latency = time.perf_counter() - tic answers_text = [] for s in states: answers_text.append([x for xs in s for x in xs]) preds = [] for i in range(len(states)): answers = [get_answer_value(v) for v in answers_text[i]] preds.append(most_frequent_number(answers)) # Compute accuracy acc = np.mean(np.array(preds) == np.array(labels)) invalid = np.mean(np.array(preds) == INVALID) print(f"Latency: {latency:.3f}") print(f"Invalid: {invalid:.3f}") print(f"Accuracy: {acc:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", answers_text) with open(args.result_file, "a") as fout: value = { "task": "tree_of_thought_gsm8k", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "accuracy": round(acc, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="test.jsonl") parser.add_argument("--num-questions", type=int, default=200) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/tree_of_thought_deep/bench_sglang.py ================================================ import argparse import ast import json import re import time from collections import Counter import numpy as np import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text, read_jsonl INVALID = -9999999 def get_answer_value(answer_str): answer_str = answer_str.replace(",", "") numbers = re.findall(r"\d+", answer_str) if len(numbers) < 1: return INVALID try: return ast.literal_eval(numbers[-1]) except SyntaxError: return INVALID def most_frequent_number(numbers): if not numbers: return None frequency = Counter(numbers) most_frequent = max(frequency, key=frequency.get) return most_frequent # Use a low temp to make the results more deterministic and the comparison more fair. temp = 0.001 def propose_plan(s, question, num_branches): s += sgl.user( """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + question ) forks = s.fork(num_branches) forks += sgl.assistant(sgl.gen("plan", max_tokens=256, temperature=temp)) return forks def execute_plan(s, num_branches): s += sgl.user( """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" ) forks = s.fork(num_branches) forks += sgl.assistant(sgl.gen("answer", max_tokens=256, temperature=temp)) return forks def reflect_solution(s, num_branches): s += sgl.user( """Okay. Now, evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" ) forks = s.fork(num_branches) forks += sgl.assistant(sgl.gen("score", max_tokens=256, temperature=temp)) return forks def get_final_answer(s, num_branches): s += sgl.user( """Based on your reflection, do you change your mind? Now, give me the final answer after careful consideration.""" ) forks = s.fork(num_branches) forks += sgl.assistant(sgl.gen("final_answer", max_tokens=256, temperature=temp)) return forks @sgl.function def tree_search(s, question, num_branches): plan_forks = propose_plan(s, question, num_branches) sol_states = [] for plan in plan_forks: forks = execute_plan(plan, num_branches) sol_states.extend(forks) ref_states = [] for sol in sol_states: forks = reflect_solution(sol, num_branches) ref_states.extend(forks) solutions = [] for sol in ref_states: forks = get_final_answer(sol, num_branches) solutions.append(forks) solutions = [[s.text() for s in forks] for forks in solutions] return solutions def main(args): lines = read_jsonl(args.data_path) lines = list(lines) # Construct prompts num_branches = 2 questions = [] labels = [] for i in range(len(lines[: args.num_questions])): questions.append(lines[i]["question"]) labels.append(get_answer_value(lines[i]["answer"])) assert all(l != INVALID for l in labels) arguments = [{"question": q, "num_branches": num_branches} for q in questions] # Select backend backend = select_sglang_backend(args) # Run requests tic = time.perf_counter() states = tree_search.run_batch( arguments, temperature=0, backend=backend, num_threads=args.parallel, progress_bar=True, ) latency = time.perf_counter() - tic answers_text = [] for s in states: answers_text.append([x for xs in s.ret_value for x in xs]) preds = [] for i in range(len(states)): answers = [get_answer_value(v) for v in answers_text[i]] preds.append(most_frequent_number(answers)) # Compute accuracy acc = np.mean(np.array(preds) == np.array(labels)) invalid = np.mean(np.array(preds) == INVALID) print(f"Latency: {latency:.3f}") print(f"Invalid: {invalid:.3f}") print(f"Accuracy: {acc:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", answers_text) with open(args.result_file, "a") as fout: value = { "task": "tree_of_thought_gsm8k", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "accuracy": round(acc, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="test.jsonl") parser.add_argument("--num-questions", type=int, default=200) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: benchmark/tree_of_thought_deep/lmql_funcs.py ================================================ from bench_other import ( ASSISTANT_PREFIX, ASSISTANT_SUFFIX, USER_PREFIX, USER_SUFFIX, temp, ) async def propose_plan_async(s, question, num_branches, call_generate): s += ( USER_PREFIX + """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + question + USER_SUFFIX ) s += ASSISTANT_PREFIX comps = await call_generate( s, max_tokens=256, temperature=temp, stop=None, n=num_branches ) return [s + comp + ASSISTANT_SUFFIX for comp in comps] async def execute_plan_async(s, num_branches, call_generate): s += ( USER_PREFIX + """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + USER_SUFFIX ) s += ASSISTANT_PREFIX comps = await call_generate( s, max_tokens=256, temperature=temp, stop=None, n=num_branches ) return [s + comp + ASSISTANT_SUFFIX for comp in comps] async def reflect_solution_async(s, num_branches, call_generate): s += ( USER_PREFIX + """Okay. Now, evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + USER_SUFFIX ) s += ASSISTANT_PREFIX comps = await call_generate( s, max_tokens=256, temperature=temp, stop=None, n=num_branches ) return [s + comp + ASSISTANT_SUFFIX for comp in comps] async def get_final_answer_async(s, num_branches, call_generate): s += ( USER_PREFIX + """Based on your reflection, do you change your mind? Now, give me the final answer after careful consideration.""" + USER_SUFFIX ) s += ASSISTANT_PREFIX comps = await call_generate( s, max_tokens=256, temperature=temp, stop=None, n=num_branches ) return [s + comp + ASSISTANT_SUFFIX for comp in comps] async def tree_search_async(question, num_branches, call_generate): plan_forks = await propose_plan_async("", question, num_branches, call_generate) sol_states = [] for plan in plan_forks: forks = await execute_plan_async(plan, num_branches, call_generate) sol_states.extend(forks) ref_states = [] for sol in sol_states: forks = await reflect_solution_async(sol, num_branches, call_generate) ref_states.extend(forks) solutions = [] for sol in ref_states: ans = await get_final_answer_async(sol, num_branches, call_generate) solutions.append(ans) return solutions ================================================ FILE: benchmark/tree_of_thought_v0/README.md ================================================ ## Download data ``` wget https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl ``` ## Run benchmark ### Benchmark sglang ``` python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 ``` ``` python3 bench_sglang.py --num-questions 32 --parallel 16 python3 bench_sglang.py --num-questions 10 --parallel 1 ``` ### Benchmark vllm ``` python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 ``` ``` python3 bench_other.py --num-questions 32 --backend vllm ``` ### Benchmark lightllm ``` # A10G python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 ``` ``` python3 bench_other.py --num-questions 32 --backend lightllm ``` ### Benchmark guidance ``` python3 bench_other.py --num-questions 32 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf ``` ================================================ FILE: benchmark/tree_of_thought_v0/bench_other.py ================================================ import argparse import ast import json import re import time from collections import Counter from concurrent.futures import ThreadPoolExecutor import numpy as np from tqdm import tqdm from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate from sglang.utils import dump_state_text, read_jsonl INVALID = -9999999 def get_answer_value(answer_str): answer_str = answer_str.replace(",", "") numbers = re.findall(r"\d+", answer_str) if len(numbers) < 1: return INVALID try: return ast.literal_eval(numbers[-1]) except SyntaxError: return INVALID def most_frequent_number(numbers): if not numbers: return None frequency = Counter(numbers) most_frequent = max(frequency, key=frequency.get) return most_frequent USER_PREFIX = "[INST] " USER_SUFFIX = " [/INST]" ASSISTANT_PREFIX = "" ASSISTANT_SUFFIX = " " # Use a low temp to make the results more deterministic and the comparison more fair. temp = 0.3 def propose_plan(s, question, num_branches, call_generate): s += ( USER_PREFIX + """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + question + USER_SUFFIX ) s += ASSISTANT_PREFIX comps = call_generate( s, max_tokens=256, temperature=temp, stop=None, n=num_branches ) return [s + comp + ASSISTANT_SUFFIX for comp in comps] def execute_plan(s, num_branches, call_generate): s += ( USER_PREFIX + """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + USER_SUFFIX ) s += ASSISTANT_PREFIX comps = call_generate( s, max_tokens=256, temperature=temp, stop=None, n=num_branches ) return [s + comp + ASSISTANT_SUFFIX for comp in comps] def reflect_solution(s, num_branches, call_generate): s += ( USER_PREFIX + """Okay. Now you evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + USER_SUFFIX ) s += ASSISTANT_PREFIX comps = call_generate( s, max_tokens=256, temperature=temp, stop=None, n=num_branches ) return [s + comp + ASSISTANT_SUFFIX for comp in comps] def tree_search(question, num_branches, call_generate): s = "" solutions = [] plan_forks = propose_plan(s, question, num_branches, call_generate) for plan in plan_forks: sol_forks = execute_plan(plan, num_branches, call_generate) for sol in sol_forks: score_forks = reflect_solution(sol, num_branches, call_generate) solutions.append(sol_forks) return solutions def main(args): lines = read_jsonl(args.data_path) # Construct prompts num_branches = 3 questions = [] labels = [] for i in range(len(lines[: args.num_questions])): questions.append(lines[i]["question"]) labels.append(get_answer_value(lines[i]["answer"])) assert all(l != INVALID for l in labels) arguments = [{"question": q, "num_branches": num_branches} for q in questions] # Select backend call_generate = get_call_generate(args) # Run requests states = [None] * len(questions) def get_one_answer(i): states[i] = tree_search(**arguments[i], call_generate=call_generate) tic = time.perf_counter() if args.parallel == 1: for i in tqdm(range(len(questions))): get_one_answer(i) else: with ThreadPoolExecutor(args.parallel) as executor: list( tqdm( executor.map(get_one_answer, list(range(len(questions)))), total=len(questions), ) ) latency = time.perf_counter() - tic answers_text = [] for s in states: answers_text.append([x for xs in s for x in xs]) preds = [] for i in range(len(states)): answers = [get_answer_value(v) for v in answers_text[i]] preds.append(most_frequent_number(answers)) # Compute accuracy acc = np.mean(np.array(preds) == np.array(labels)) invalid = np.mean(np.array(preds) == INVALID) print(f"Latency: {latency:.3f}") print(f"Invalid: {invalid:.3f}") print(f"Accuracy: {acc:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", answers_text) with open(args.result_file, "a") as fout: value = { "task": "tree_of_thought_gsm8k", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "accuracy": round(acc, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="test.jsonl") parser.add_argument("--num-questions", type=int, default=200) args = add_common_other_args_and_parse(parser) main(args) ================================================ FILE: benchmark/tree_of_thought_v0/bench_sglang.py ================================================ import argparse import ast import json import re import time from collections import Counter import numpy as np import sglang as sgl from sglang.test.test_utils import ( add_common_sglang_args_and_parse, select_sglang_backend, ) from sglang.utils import dump_state_text, read_jsonl INVALID = -9999999 def get_answer_value(answer_str): answer_str = answer_str.replace(",", "") numbers = re.findall(r"\d+", answer_str) if len(numbers) < 1: return INVALID try: return ast.literal_eval(numbers[-1]) except SyntaxError: return INVALID def most_frequent_number(numbers): if not numbers: return None frequency = Counter(numbers) most_frequent = max(frequency, key=frequency.get) return most_frequent # Use a low temp to make the results more deterministic and the comparison more fair. temp = 0.3 def propose_plan(s, question, num_branches): s += sgl.user( """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + question ) forks = s.fork(num_branches) forks += sgl.assistant(sgl.gen("plan", max_tokens=256, temperature=temp)) return forks def execute_plan(s, num_branches): s += sgl.user( """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" ) forks = s.fork(num_branches) forks += sgl.assistant(sgl.gen("answer", max_tokens=256, temperature=temp)) return forks def reflect_solution(s, num_branches): s += sgl.user( """Okay. Now you evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" ) forks = s.fork(num_branches) forks += sgl.assistant(sgl.gen("score", max_tokens=256, temperature=temp)) return forks @sgl.function def tree_search(s, question, num_branches): forks_to_join = [] plan_forks = propose_plan(s, question, num_branches) forks_to_join.append(plan_forks) sol_states = [] for plan in plan_forks: forks = execute_plan(plan, num_branches) forks_to_join.append(forks) sol_states.extend(forks) for sol in sol_states: forks = reflect_solution(sol, num_branches) forks_to_join.append(forks) for f in reversed(forks_to_join): f.join() def main(args): lines = read_jsonl(args.data_path) # Construct prompts num_branches = 3 questions = [] labels = [] for i in range(len(lines[: args.num_questions])): questions.append(lines[i]["question"]) labels.append(get_answer_value(lines[i]["answer"])) assert all(l != INVALID for l in labels) arguments = [{"question": q, "num_branches": num_branches} for q in questions] # Select backend backend = select_sglang_backend(args) # Run requests tic = time.perf_counter() states = tree_search.run_batch( arguments, temperature=0, backend=backend, num_threads=args.parallel, progress_bar=True, ) latency = time.perf_counter() - tic answers_text = [] for s in states: answers_text.append([x for xs in s["answer"] for x in xs]) preds = [] for i in range(len(states)): answers = [get_answer_value(v) for v in answers_text[i]] preds.append(most_frequent_number(answers)) # Compute accuracy acc = np.mean(np.array(preds) == np.array(labels)) invalid = np.mean(np.array(preds) == INVALID) print(f"Latency: {latency:.3f}") print(f"Invalid: {invalid:.3f}") print(f"Accuracy: {acc:.3f}") # Write results dump_state_text(f"tmp_output_{args.backend}.txt", answers_text) with open(args.result_file, "a") as fout: value = { "task": "tree_of_thought_gsm8k", "backend": args.backend, "num_gpus": 1, "latency": round(latency, 3), "accuracy": round(acc, 3), "num_requests": args.num_questions, "other": { "num_questions": args.num_questions, "parallel": args.parallel, }, } fout.write(json.dumps(value) + "\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data-path", type=str, default="test.jsonl") parser.add_argument("--num-questions", type=int, default=200) args = add_common_sglang_args_and_parse(parser) main(args) ================================================ FILE: docker/Dockerfile ================================================ ARG CUDA_VERSION=12.9.1 FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu24.04 AS base ARG TARGETARCH ARG BUILD_TYPE=all ARG BRANCH_TYPE=remote ARG GRACE_BLACKWELL=0 ARG HOPPER_SBO=0 ARG GRACE_BLACKWELL_DEEPEP_BRANCH=gb200_blog_part_2 ARG HOPPER_SBO_DEEPEP_COMMIT=9f2fc4b3182a51044ae7ecb6610f7c9c3258c4d6 ARG DEEPEP_COMMIT=9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee ARG BUILD_AND_DOWNLOAD_PARALLEL=8 ARG SGL_KERNEL_VERSION=0.4.0 ARG SGL_VERSION ARG USE_LATEST_SGLANG=0 ARG GDRCOPY_VERSION=2.5.1 ARG PIP_DEFAULT_INDEX ARG UBUNTU_MIRROR ARG GITHUB_ARTIFACTORY=github.com ARG INSTALL_FLASHINFER_JIT_CACHE=0 ARG FLASHINFER_VERSION=0.6.6 ARG MOONCAKE_VERSION=0.3.9 #if need other arg please add in MOONCAKE_COMPILE_ARG ARG MOONCAKE_COMPILE_ARG="-DUSE_HTTP=ON -DUSE_MNNVL=ON -DUSE_CUDA=ON -DWITH_EP=ON" ENV DEBIAN_FRONTEND=noninteractive \ CUDA_HOME=/usr/local/cuda \ GDRCOPY_HOME=/usr/src/gdrdrv-${GDRCOPY_VERSION}/ \ FLASHINFER_VERSION=${FLASHINFER_VERSION} # Add GKE default lib and bin locations ENV PATH="${PATH}:/usr/local/nvidia/bin" \ LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/nvidia/lib:/usr/local/nvidia/lib64" # Replace Ubuntu sources if specified RUN if [ -n "$UBUNTU_MIRROR" ]; then \ sed -i "s|http://.*archive.ubuntu.com|$UBUNTU_MIRROR|g" /etc/apt/sources.list && \ sed -i "s|http://.*security.ubuntu.com|$UBUNTU_MIRROR|g" /etc/apt/sources.list; \ fi # Python setup (combined with apt update to reduce layers) RUN --mount=type=cache,target=/var/cache/apt,id=base-apt \ apt update && apt install -y --no-install-recommends wget software-properties-common \ && add-apt-repository ppa:deadsnakes/ppa -y \ && apt install -y --no-install-recommends python3.12-full python3.12-dev python3.10-venv \ && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 \ && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 2 \ && update-alternatives --set python3 /usr/bin/python3.12 \ && wget -q https://bootstrap.pypa.io/get-pip.py \ && python3 get-pip.py --break-system-packages \ && rm get-pip.py \ # Allow pip to install packages globally (PEP 668 workaround for Ubuntu 24.04) && python3 -m pip config set global.break-system-packages true \ # Fix for apt-add-repository && cd /usr/lib/python3/dist-packages/ \ && ln -s apt_pkg.cpython-310-*-linux-gnu.so apt_pkg.so # Install system dependencies (organized by category for better caching) RUN --mount=type=cache,target=/var/cache/apt,id=base-apt \ apt-get update && apt-get install -y --no-install-recommends \ # Core system utilities ca-certificates \ software-properties-common \ netcat-openbsd \ kmod \ unzip \ openssh-server \ curl \ wget \ lsof \ locales \ # Build essentials (needed for framework stage) build-essential \ cmake \ perl \ patchelf \ ccache \ git-lfs \ # MPI and NUMA libopenmpi-dev \ libnuma1 \ libnuma-dev \ numactl \ # transformers multimodal VLM ffmpeg \ # InfiniBand/RDMA libibverbs-dev \ libibverbs1 \ libibumad3 \ librdmacm1 \ libnl-3-200 \ libnl-route-3-200 \ libnl-route-3-dev \ libnl-3-dev \ ibverbs-providers \ infiniband-diags \ perftest \ # Development libraries libgoogle-glog-dev \ libgtest-dev \ libjsoncpp-dev \ libunwind-dev \ libboost-all-dev \ libssl-dev \ libgrpc-dev \ libgrpc++-dev \ libprotobuf-dev \ protobuf-compiler \ protobuf-compiler-grpc \ pybind11-dev \ libhiredis-dev \ libcurl4-openssl-dev \ libczmq4 \ libczmq-dev \ libfabric-dev \ # Package building tools devscripts \ debhelper \ fakeroot \ dkms \ check \ libsubunit0 \ libsubunit-dev \ && ln -sf /usr/bin/python3.12 /usr/bin/python \ && rm -rf /var/lib/apt/lists/* \ && apt-get clean # Replace pip global cache if specified RUN if [ -n "${PIP_DEFAULT_INDEX}" ]; then \ python3 -m pip config set global.index-url ${PIP_DEFAULT_INDEX}; \ fi # GDRCopy installation RUN mkdir -p /tmp/gdrcopy && cd /tmp \ && curl --retry 3 --retry-delay 2 -fsSL -o v${GDRCOPY_VERSION}.tar.gz \ https://${GITHUB_ARTIFACTORY}/NVIDIA/gdrcopy/archive/refs/tags/v${GDRCOPY_VERSION}.tar.gz \ && tar -xzf v${GDRCOPY_VERSION}.tar.gz && rm v${GDRCOPY_VERSION}.tar.gz \ && cd gdrcopy-${GDRCOPY_VERSION}/packages \ && CUDA=/usr/local/cuda ./build-deb-packages.sh \ && dpkg -i gdrdrv-dkms_*.deb libgdrapi_*.deb gdrcopy-tests_*.deb gdrcopy_*.deb \ && cd / && rm -rf /tmp/gdrcopy # Fix DeepEP IBGDA symlink RUN ln -sf /usr/lib/$(uname -m)-linux-gnu/libmlx5.so.1 /usr/lib/$(uname -m)-linux-gnu/libmlx5.so # Set up locale RUN locale-gen en_US.UTF-8 ENV LANG=en_US.UTF-8 \ LANGUAGE=en_US:en \ LC_ALL=en_US.UTF-8 ######################################################## ########## Framework Development Image ################ ######################################################## # Copy local source if building from local FROM scratch AS local_src COPY . /src FROM base AS framework ARG BRANCH_TYPE ARG BUILD_TYPE ARG CUDA_VERSION ARG BUILD_AND_DOWNLOAD_PARALLEL ARG SGL_KERNEL_VERSION ARG SGL_VERSION ARG USE_LATEST_SGLANG ARG INSTALL_FLASHINFER_JIT_CACHE ARG FLASHINFER_VERSION ARG GRACE_BLACKWELL ARG GRACE_BLACKWELL_DEEPEP_BRANCH ARG DEEPEP_COMMIT ARG TRITON_LANG_COMMIT ARG GITHUB_ARTIFACTORY WORKDIR /sgl-workspace # Install SGLang COPY --from=local_src /src /tmp/local_src RUN if [ "$BRANCH_TYPE" = "local" ]; then \ cp -r /tmp/local_src /sgl-workspace/sglang; \ elif [ "$USE_LATEST_SGLANG" = "1" ]; then \ git clone --depth=1 https://github.com/sgl-project/sglang.git /sgl-workspace/sglang; \ elif [ -z "$SGL_VERSION" ]; then \ echo "ERROR: SGL_VERSION must be set when USE_LATEST_SGLANG=0 and BRANCH_TYPE!=local" && exit 1; \ else \ git clone --depth=1 --branch v${SGL_VERSION} https://github.com/sgl-project/sglang.git /sgl-workspace/sglang; \ fi \ && rm -rf /tmp/local_src RUN --mount=type=cache,target=/root/.cache/pip \ python3 -m pip install --upgrade pip setuptools wheel html5lib six \ && cd sglang \ && case "$CUDA_VERSION" in \ 12.6.1) CUINDEX=126 ;; \ 12.8.1) CUINDEX=128 ;; \ 12.9.1) CUINDEX=129 ;; \ 13.0.1) CUINDEX=130 ;; \ *) echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 ;; \ esac \ && if [ "$CUDA_VERSION" = "12.6.1" ]; then \ python3 -m pip install https://${GITHUB_ARTIFACTORY}/sgl-project/whl/releases/download/v${SGL_KERNEL_VERSION}/sglang_kernel-${SGL_KERNEL_VERSION}+cu124-cp310-abi3-manylinux2014_$(uname -m).whl --force-reinstall --no-deps \ ; \ elif [ "$CUDA_VERSION" = "12.8.1" ] || [ "$CUDA_VERSION" = "12.9.1" ]; then \ python3 -m pip install sglang-kernel==${SGL_KERNEL_VERSION} \ ; \ elif [ "$CUDA_VERSION" = "13.0.1" ]; then \ python3 -m pip install https://github.com/sgl-project/whl/releases/download/v${SGL_KERNEL_VERSION}/sglang_kernel-${SGL_KERNEL_VERSION}+cu130-cp310-abi3-manylinux2014_$(uname -m).whl --force-reinstall --no-deps \ ; \ else \ echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 \ ; \ fi \ && python3 -m pip install -e "python[${BUILD_TYPE}]" --extra-index-url https://download.pytorch.org/whl/cu${CUINDEX} \ && if [ "$INSTALL_FLASHINFER_JIT_CACHE" = "1" ]; then \ python3 -m pip install flashinfer-jit-cache==${FLASHINFER_VERSION} --index-url https://flashinfer.ai/whl/cu${CUINDEX} ; \ fi \ && FLASHINFER_CUBIN_DOWNLOAD_THREADS=${BUILD_AND_DOWNLOAD_PARALLEL} FLASHINFER_LOGGING_LEVEL=warning python3 -m flashinfer --download-cubin # DeepEP # We use Tom's DeepEP fork for GB200 for now; the 1fd57b0276311d035d16176bb0076426166e52f3 commit is https://github.com/fzyzcjy/DeepEP/tree/gb200_blog_part_2 # TODO: move from Tom's branch to DeepEP hybrid-ep branch # We use the nvshmem version that ships with torch 2.9.1 # CU12 uses 3.3.20 and CU13 uses 3.3.24 RUN set -eux; \ if [ "$GRACE_BLACKWELL" = "1" ]; then \ git clone https://github.com/fzyzcjy/DeepEP.git && \ cd DeepEP && \ git checkout ${GRACE_BLACKWELL_DEEPEP_BRANCH} && \ sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh && \ sed -i 's/#define NUM_TIMEOUT_CYCLES 200000000000ull/#define NUM_TIMEOUT_CYCLES 2000000000000ull/' csrc/kernels/configs.cuh && \ cd .. ; \ elif [ "$HOPPER_SBO" = "1" ]; then \ git clone https://github.com/deepseek-ai/DeepEP.git -b antgroup-opt && \ cd DeepEP && \ git checkout ${HOPPER_SBO_DEEPEP_COMMIT} && \ sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh && \ sed -i 's/#define NUM_TIMEOUT_CYCLES 200000000000ull/#define NUM_TIMEOUT_CYCLES 2000000000000ull/' csrc/kernels/configs.cuh && \ cd .. ; \ else \ curl --retry 3 --retry-delay 2 -fsSL -o ${DEEPEP_COMMIT}.zip \ https://${GITHUB_ARTIFACTORY}/deepseek-ai/DeepEP/archive/${DEEPEP_COMMIT}.zip && \ unzip -q ${DEEPEP_COMMIT}.zip && rm ${DEEPEP_COMMIT}.zip && mv DeepEP-${DEEPEP_COMMIT} DeepEP && cd DeepEP && \ sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh && \ sed -i 's/#define NUM_TIMEOUT_CYCLES 200000000000ull/#define NUM_TIMEOUT_CYCLES 2000000000000ull/' csrc/kernels/configs.cuh && \ cd .. ; \ fi # Install DeepEP RUN --mount=type=cache,target=/root/.cache/pip \ cd /sgl-workspace/DeepEP && \ case "$CUDA_VERSION" in \ 12.6.1) \ CHOSEN_TORCH_CUDA_ARCH_LIST='9.0' \ ;; \ 12.8.1) \ # FIXED: 12.8.1 does NOT support Blackwell 10.3 \ CHOSEN_TORCH_CUDA_ARCH_LIST='9.0;10.0' \ ;; \ 12.9.1|13.0.1) \ # 12.9.1+ properly supports Blackwell 10.3 \ CHOSEN_TORCH_CUDA_ARCH_LIST='9.0;10.0;10.3' \ ;; \ *) \ echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 \ ;; \ esac && \ if [ "${CUDA_VERSION%%.*}" = "13" ]; then \ sed -i "/^ include_dirs = \['csrc\/'\]/a\ include_dirs.append('${CUDA_HOME}/include/cccl')" setup.py; \ fi && \ TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" MAX_JOBS=${BUILD_AND_DOWNLOAD_PARALLEL} pip install --no-build-isolation . # Install Mooncake RUN --mount=type=cache,target=/root/.cache/pip \ CUDA_MAJOR="${CUDA_VERSION%%.*}" && \ if [ "$CUDA_MAJOR" -ge 13 ]; then \ echo "CUDA >= 13, installing mooncake-transfer-engine from source code"; \ git clone --branch v${MOONCAKE_VERSION} --depth 1 https://github.com/kvcache-ai/Mooncake.git && \ cd Mooncake && \ bash dependencies.sh && \ mkdir -p build && \ cd build && \ cmake .. ${MOONCAKE_COMPILE_ARG} && \ make -j$(nproc) && \ make install; \ else \ echo "CUDA < 13, installing mooncake-transfer-engine from pip"; \ python3 -m pip install mooncake-transfer-engine==${MOONCAKE_VERSION}; \ fi # Install essential Python packages RUN --mount=type=cache,target=/root/.cache/pip \ python3 -m pip install \ datamodel_code_generator \ pre-commit \ pytest \ black \ isort \ icdiff \ uv \ wheel \ scikit-build-core \ nixl \ py-spy \ cubloaty \ google-cloud-storage # Build and install sgl-model-gateway (install Rust, build, then remove to save space) RUN --mount=type=cache,target=/root/.cache/pip \ curl --proto '=https' --tlsv1.2 --retry 3 --retry-delay 2 -sSf https://sh.rustup.rs | sh -s -- -y \ && export PATH="/root/.cargo/bin:${PATH}" \ && rustc --version && cargo --version \ && python3 -m pip install maturin \ && cd /sgl-workspace/sglang/sgl-model-gateway/bindings/python \ && ulimit -n 65536 && maturin build --release --features vendored-openssl --out dist \ && python3 -m pip install --force-reinstall dist/*.whl \ && cd /sgl-workspace/sglang/sgl-model-gateway \ && cargo build --release --bin sglang-router --features vendored-openssl \ && cp target/release/sglang-router /usr/local/bin/sglang-router \ && rm -rf /root/.cargo /root/.rustup target dist ~/.cargo \ && sed -i '/\.cargo\/env/d' /root/.profile /root/.bashrc 2>/dev/null || true RUN --mount=type=cache,target=/root/.cache/pip \ python3 -m pip install "nvidia-cutlass-dsl>=4.4.1" "nvidia-cutlass-dsl-libs-base>=4.4.1" --force-reinstall --no-deps; # Patching packages for CUDA 12/13 compatibility # TODO: Remove when torch version covers these packages RUN --mount=type=cache,target=/root/.cache/pip if [ "${CUDA_VERSION%%.*}" = "12" ]; then \ python3 -m pip install nvidia-nccl-cu12==2.28.3 --force-reinstall --no-deps ; \ python3 -m pip install nvidia-cudnn-cu12==9.16.0.29 --force-reinstall --no-deps ; \ python3 -m pip install cuda-python==12.9 ; \ elif [ "${CUDA_VERSION%%.*}" = "13" ]; then \ python3 -m pip install nvidia-nccl-cu13==2.28.3 --force-reinstall --no-deps ; \ python3 -m pip install nvidia-cudnn-cu13==9.16.0.29 --force-reinstall --no-deps ; \ python3 -m pip install nvidia-cublas==13.1.0.3 --force-reinstall --no-deps ; \ python3 -m pip install nixl-cu13 --no-deps ; \ python3 -m pip install cuda-python==13.2.0 ; \ fi # Install development tools RUN --mount=type=cache,target=/var/cache/apt,id=framework-apt \ apt-get update && apt-get install -y --no-install-recommends \ gdb \ ninja-build \ vim \ tmux \ htop \ zsh \ tree \ silversearcher-ag \ cloc \ pkg-config \ bear \ less \ rdma-core \ openssh-server \ gnuplot \ infiniband-diags \ perftest \ ibverbs-providers \ libibumad3 \ libibverbs1 \ libnl-3-200 \ libnl-route-3-200 \ librdmacm1 \ && rm -rf /var/lib/apt/lists/* \ && apt-get clean # Install NVIDIA development tools RUN --mount=type=cache,target=/var/cache/apt,id=framework-apt \ apt update -y \ && apt install -y --no-install-recommends gnupg \ && echo "deb http://developer.download.nvidia.com/devtools/repos/ubuntu2004/$(if [ "$(uname -m)" = "aarch64" ]; then echo "arm64"; else echo "amd64"; fi) /" | tee /etc/apt/sources.list.d/nvidia-devtools.list \ && apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/$(if [ "$(uname -m)" = "aarch64" ]; then echo "arm64"; else echo "x86_64"; fi)/7fa2af80.pub \ && apt update -y \ && apt install -y --no-install-recommends nsight-systems-cli \ && rm -rf /var/lib/apt/lists/* # Install minimal Python dev packages RUN --mount=type=cache,target=/root/.cache/pip \ python3 -m pip install --break-system-packages \ pytest \ black \ isort \ icdiff \ scikit-build-core \ uv \ pre-commit \ pandas \ matplotlib \ tabulate \ termplotlib # diff-so-fancy RUN curl --retry 3 --retry-delay 2 -LSso /usr/local/bin/diff-so-fancy \ https://${GITHUB_ARTIFACTORY}/so-fancy/diff-so-fancy/releases/download/v1.4.4/diff-so-fancy \ && chmod +x /usr/local/bin/diff-so-fancy # clang-format RUN curl --retry 3 --retry-delay 2 -LSso /usr/local/bin/clang-format \ https://${GITHUB_ARTIFACTORY}/muttleyxd/clang-tools-static-binaries/releases/download/master-32d3ac78/clang-format-16_linux-amd64 \ && chmod +x /usr/local/bin/clang-format # clangd RUN curl --retry 3 --retry-delay 2 -fsSL -o clangd.zip \ https://${GITHUB_ARTIFACTORY}/clangd/clangd/releases/download/18.1.3/clangd-linux-18.1.3.zip \ && unzip -q clangd.zip \ && cp -r clangd_18.1.3/bin/* /usr/local/bin/ \ && cp -r clangd_18.1.3/lib/* /usr/local/lib/ \ && rm -rf clangd_18.1.3 clangd.zip # CMake RUN CMAKE_VERSION=3.31.1 \ && ARCH=$(uname -m) \ && CMAKE_INSTALLER="cmake-${CMAKE_VERSION}-linux-${ARCH}" \ && curl --retry 3 --retry-delay 2 -fsSL -o "${CMAKE_INSTALLER}.tar.gz" \ "https://${GITHUB_ARTIFACTORY}/Kitware/CMake/releases/download/v${CMAKE_VERSION}/${CMAKE_INSTALLER}.tar.gz" \ && tar -xzf "${CMAKE_INSTALLER}.tar.gz" \ && cp -r "${CMAKE_INSTALLER}/bin/"* /usr/local/bin/ \ && cp -r "${CMAKE_INSTALLER}/share/"* /usr/local/share/ \ && rm -rf "${CMAKE_INSTALLER}" "${CMAKE_INSTALLER}.tar.gz" # Install just RUN curl --proto '=https' --tlsv1.2 --retry 3 --retry-delay 2 -sSf https://just.systems/install.sh | \ sed "s|https://github.com|https://${GITHUB_ARTIFACTORY}|g" | \ bash -s -- --tag 1.42.4 --to /usr/local/bin # Add yank script COPY --chown=root:root --chmod=755 docker/configs/yank /usr/local/bin/yank # Install oh-my-zsh and plugins RUN sh -c "$(curl --retry 3 --retry-delay 2 -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended \ && git clone --depth 1 https://github.com/zsh-users/zsh-autosuggestions ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-autosuggestions \ && git clone --depth 1 https://github.com/zsh-users/zsh-syntax-highlighting.git ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-syntax-highlighting # These configs are optional; users can override them by mounting their own files COPY docker/configs/opt/.vimrc /opt/sglang/.vimrc COPY docker/configs/opt/.tmux.conf /opt/sglang/.tmux.conf COPY docker/configs/opt/.gitconfig /opt/sglang/.gitconfig # Configure development environment COPY docker/configs/.zshrc /root/.zshrc # Fix Triton to use system ptxas for Blackwell (sm_103a) support (CUDA 13+ only) RUN if [ "${CUDA_VERSION%%.*}" = "13" ] && [ -d /usr/local/lib/python3.12/dist-packages/triton/backends/nvidia/bin ]; then \ rm -f /usr/local/lib/python3.12/dist-packages/triton/backends/nvidia/bin/ptxas && \ ln -s /usr/local/cuda/bin/ptxas /usr/local/lib/python3.12/dist-packages/triton/backends/nvidia/bin/ptxas; \ fi RUN python3 -m pip install --upgrade "urllib3>=2.6.3" # Set workspace directory WORKDIR /sgl-workspace/sglang ######################################################## ########## Runtime Image ############################## ######################################################## # # PURPOSE: Production runtime environment with JIT support # # This stage creates a production-ready image containing: # - Pre-compiled SGLang and DeepEP components # - Full CUDA toolchain for JIT compilation (DeepGEMM, Triton, FlashInfer) # - Optimized for inference workloads and deployment # - Smaller than framework (no dev tools like vim, tmux, nsight, etc.) # # Use this stage when you need: # - Production deployment of SGLang # - JIT compilation support for FP8/microscaling kernels # - Ready-to-run inference server environment # # Note: Uses devel base for complete NVCC toolchain required by DeepGEMM JIT FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu24.04 AS runtime ARG CUDA_VERSION ARG TARGETARCH ARG GDRCOPY_VERSION=2.5.1 ENV DEBIAN_FRONTEND=noninteractive \ CUDA_HOME=/usr/local/cuda \ GDRCOPY_HOME=/usr/src/gdrdrv-${GDRCOPY_VERSION}/ # Add GKE default lib and bin locations + CUDA compiler paths for FlashInfer JIT ENV PATH="${PATH}:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/cuda/nvvm/bin" \ LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/nvidia/lib:/usr/local/nvidia/lib64" # Install runtime dependencies (devel base provides gcc/g++/build tools) RUN --mount=type=cache,target=/var/cache/apt,id=runtime-apt \ apt-get update && apt-get install -y --no-install-recommends \ # Python runtime software-properties-common \ && add-apt-repository ppa:deadsnakes/ppa -y \ && apt-get update && apt-get install -y --no-install-recommends --allow-change-held-packages \ python3.12-full \ python3.12-dev \ wget \ # Core system utilities ca-certificates \ netcat-openbsd \ curl \ git \ # Runtime libraries libopenmpi3 \ libnuma1 \ libibverbs1 \ libibumad3 \ librdmacm1 \ libnl-3-200 \ libnl-route-3-200 \ ibverbs-providers \ libgoogle-glog0v6t64 \ libunwind8 \ libboost-system1.83.0 \ libboost-thread1.83.0 \ libboost-filesystem1.83.0 \ libgrpc++1.51t64 \ libprotobuf32t64 \ libhiredis1.1.0 \ libcurl4 \ libczmq4 \ libfabric1 \ libssl3 \ # RDMA runtime rdma-core \ infiniband-diags \ perftest \ # Build tools for JIT compilation ninja-build \ # NCCL packages needed for pynccl_allocator JIT compilation (-lnccl) libnccl2 \ libnccl-dev \ # GPG key verification gnupg2 \ && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 2 \ && update-alternatives --set python3 /usr/bin/python3.12 \ && ln -sf /usr/bin/python3.12 /usr/bin/python \ && wget -q https://bootstrap.pypa.io/get-pip.py \ && python3 get-pip.py --break-system-packages \ && rm get-pip.py \ # Allow pip to install packages globally (PEP 668 workaround for Ubuntu 24.04) && python3 -m pip config set global.break-system-packages true \ && rm -rf /var/lib/apt/lists/* \ && apt-get clean # Set up locale RUN apt-get update && apt-get install -y --no-install-recommends locales \ && locale-gen en_US.UTF-8 \ && rm -rf /var/lib/apt/lists/* ENV LANG=en_US.UTF-8 \ LANGUAGE=en_US:en \ LC_ALL=en_US.UTF-8 # Copy Python site-packages from framework (contains all built packages) COPY --from=framework /usr/local/lib/python3.12/dist-packages /usr/local/lib/python3.12/dist-packages # Copy SGLang workspace COPY --from=framework /sgl-workspace /sgl-workspace # Fix Triton to use system ptxas for Blackwell (sm_103a) support (CUDA 13+ only) RUN if [ "${CUDA_VERSION%%.*}" = "13" ] && [ -d /usr/local/lib/python3.12/dist-packages/triton/backends/nvidia/bin ]; then \ rm -f /usr/local/lib/python3.12/dist-packages/triton/backends/nvidia/bin/ptxas && \ ln -s /usr/local/cuda/bin/ptxas /usr/local/lib/python3.12/dist-packages/triton/backends/nvidia/bin/ptxas; \ fi # Copy GDRCopy runtime libraries (but not the build artifacts) COPY --from=framework /usr/lib/libgdrapi.so* /usr/lib/ COPY --from=framework /usr/bin/gdrcopy_* /usr/bin/ COPY --from=framework /usr/src/gdrdrv-2.5.1 /usr/src/gdrdrv-2.5.1 # Fix DeepEP IBGDA symlink in runtime RUN ln -sf /usr/lib/$(uname -m)-linux-gnu/libmlx5.so.1 /usr/lib/$(uname -m)-linux-gnu/libmlx5.so WORKDIR /sgl-workspace/sglang # Default command CMD ["/bin/bash"] ================================================ FILE: docker/compose.yaml ================================================ services: sglang: image: lmsysorg/sglang:latest container_name: sglang volumes: - ${HOME}/.cache/huggingface:/root/.cache/huggingface # If you use modelscope, you need mount this directory # - ${HOME}/.cache/modelscope:/root/.cache/modelscope restart: always network_mode: host # required by RDMA privileged: true # required by RDMA # Or you can only publish port 30000 # ports: # - 30000:30000 environment: - HF_TOKEN= # if you use modelscope to download model, you need set this environment # - SGLANG_USE_MODELSCOPE=true entrypoint: python3 -m sglang.launch_server command: --model-path meta-llama/Llama-3.1-8B-Instruct --host 0.0.0.0 --port 30000 ulimits: memlock: -1 stack: 67108864 ipc: host healthcheck: test: ["CMD-SHELL", "curl -f http://localhost:30000/health || exit 1"] deploy: resources: reservations: devices: - driver: nvidia count: 1 capabilities: [gpu] ================================================ FILE: docker/configs/.zshrc ================================================ export ZSH="/root/.oh-my-zsh" # Theme ZSH_THEME="robbyrussell" # Plugins plugins=( git z zsh-autosuggestions zsh-syntax-highlighting ) source $ZSH/oh-my-zsh.sh # Aliases alias ll='ls -alF' alias la='ls -A' alias l='ls -CF' alias vi='vim' # Enhanced history HISTSIZE=10000 SAVEHIST=10000 setopt HIST_IGNORE_ALL_DUPS setopt HIST_FIND_NO_DUPS setopt INC_APPEND_HISTORY ================================================ FILE: docker/configs/opt/.gitconfig ================================================ [core] editor = vim whitespace = fix,-indent-with-non-tab,trailing-space,cr-at-eol pager = diff-so-fancy | less --tabs=4 -RFX [color] ui = true [color "diff-highlight"] oldNormal = red bold oldHighlight = red bold 52 newNormal = green bold newHighlight = green bold 22 [color "diff"] meta = 11 frag = magenta bold commit = yellow bold old = red bold new = green bold whitespace = red reverse [alias] lg = log --color --graph --pretty=format:'%Cred%h%Creset - %s %Cgreen(%cr) %C(bold blue)<%an>%Creset%C(auto)%d%Creset' --abbrev-commit -- [http] sslVerify = false [pull] rebase = true ================================================ FILE: docker/configs/opt/.tmux.conf ================================================ # Pane border styling set -g pane-border-style fg='#742727',bg=black set -g pane-active-border-style fg=red,bg=black # Status bar styling set -g status-style bg='#0C8A92',fg=black # Change prefix key to backtick set-option -g prefix ` unbind C-b bind-key ` send-prefix # Split panes using - and = with current path unbind '"' bind - splitw -v -c '#{pane_current_path}' unbind '%' bind = splitw -h -c '#{pane_current_path}' # Vi mode settings bind-key -T copy-mode-vi Y send-keys -X copy-pipe 'yank > #{pane_tty}' set-window-option -g mode-keys vi # Other settings set-option -g escape-time 0 set-option -g base-index 1 set-window-option -g mouse on set -g history-limit 100000 ================================================ FILE: docker/configs/opt/.vimrc ================================================ function! Yank(text) abort let escape = system('yank', a:text) if v:shell_error echoerr escape else call writefile([escape], '/dev/tty', 'b') endif endfunction noremap y y:call Yank(@0) " automatically run yank(1) whenever yanking in Vim function! CopyYank() abort call Yank(join(v:event.regcontents, "\n")) endfunction autocmd TextYankPost * call CopyYank() " Basic settings set number syntax on set mouse=a filetype indent on " Indentation set autoindent nosmartindent set smarttab set expandtab set shiftwidth=4 set softtabstop=4 " Visual guides set colorcolumn=120 highlight ColorColumn ctermbg=5 " Status line set laststatus=2 set statusline=%<%f\ %h%m%r%=%{\"[\".(&fenc==\"\"?&enc:&fenc).((exists(\"+bomb\")\ &&\ &bomb)?\",B\":\"\").\"]\ \"}%k\ %-14.(%l,%c%V%)\ %P " Backspace behavior set backspace=2 " Encoding set encoding=utf-8 set fileencoding=utf-8 ================================================ FILE: docker/configs/yank ================================================ #!/bin/bash put() { esc=$1 test -n "$TMUX" -o -z "${TERM##screen*}" && esc="\033Ptmux;\033$esc\033\\" printf "$esc" } put "\033]52;c;!\a" buf=$( cat "$@" ) len=$( printf %s "$buf" | wc -c ) max=74994 test $len -gt $max && echo "$0: input is $(( len - max )) bytes too long" >&2 put "\033]52;c;$( printf %s "$buf" | head -c $max | base64 | tr -d '\r\n' )\a" test -n "$TMUX" && tmux set-buffer "$buf" ||: ================================================ FILE: docker/diffusion.Dockerfile ================================================ FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04 ENV DEBIAN_FRONTEND=noninteractive SHELL ["/bin/bash", "-c"] WORKDIR /sgl-workspace/sglang RUN apt-get update && apt-get install -y --no-install-recommends \ wget \ git \ ca-certificates \ openssh-server \ zsh \ vim \ curl \ gcc-11 \ g++-11 \ clang-11 \ libnuma1 libnuma-dev \ && rm -rf /var/lib/apt/lists/* # Install oh-my-zsh and plugins RUN sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended \ && git clone https://github.com/zsh-users/zsh-autosuggestions ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-autosuggestions \ && git clone https://github.com/zsh-users/zsh-syntax-highlighting.git ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-syntax-highlighting # Set up C++20 compilers for ThunderKittens RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11 # Set CUDA environment variables ENV CUDA_HOME=/usr/local/cuda-12.8 ENV PATH=${CUDA_HOME}/bin:${PATH} ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH # Install uv and source its environment RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \ echo 'source $HOME/.local/bin/env' >> /root/.zshrc # Copy just the pyproject.toml first to leverage Docker cache COPY python/pyproject.toml python/ # Create a dummy README to satisfy the installation RUN mkdir -p python && echo "# Placeholder" > python/README.md # Create and activate virtual environment with specific Python version and seed RUN source $HOME/.local/bin/env && \ uv venv --python 3.12 --seed /opt/venv && \ source /opt/venv/bin/activate && \ uv pip install nvitop && \ uv pip install --no-cache-dir --upgrade pip && \ uv pip install --no-cache-dir --prerelease=allow ./python[diffusion] COPY . . # Install dependencies using uv and set up shell configuration RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ git config --unset-all http.https://github.com/.extraheader || true && \ echo 'source /opt/venv/bin/activate' >> /root/.zshrc && \ echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile # Set PATH to include venv bin ENV PATH=/opt/venv/bin:$PATH # Configure zsh COPY --chown=root:root <<-"EOF" /root/.zshrc export ZSH="/root/.oh-my-zsh" source $HOME/.local/bin/env source /opt/venv/bin/activate ## Theme ZSH_THEME="robbyrussell" ## Plugins plugins=( git z zsh-autosuggestions zsh-syntax-highlighting ) source $ZSH/oh-my-zsh.sh ## Aliases alias ll='ls -alF' alias la='ls -A' alias l='ls -CF' alias vi='vim' ## Enhanced history HISTSIZE=10000 SAVEHIST=10000 setopt HIST_IGNORE_ALL_DUPS setopt HIST_FIND_NO_DUPS setopt INC_APPEND_HISTORY EOF EXPOSE 22 CMD ["/bin/zsh"] ================================================ FILE: docker/gateway.Dockerfile ================================================ ######################## BASE IMAGE ########################## FROM ubuntu:24.04 AS base ARG PYTHON_VERSION=3.12 # set the environment variables ENV PATH="/root/.local/bin:${PATH}" ENV DEBIAN_FRONTEND=noninteractive # uv environment variables ENV UV_HTTP_TIMEOUT=500 ENV VIRTUAL_ENV="/opt/venv" ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python ENV UV_LINK_MODE="copy" ENV PATH="$VIRTUAL_ENV/bin:$PATH" # install dependencies RUN apt update -y \ && apt install -y curl \ && rm -rf /var/lib/apt/lists/* \ && apt clean # install uv RUN curl -LsSf https://astral.sh/uv/install.sh | sh # install python RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV} FROM scratch AS local_src COPY . /src ######################### BUILD IMAGE ######################### FROM base AS build-image # set the environment variables ENV PATH="/root/.cargo/bin:${PATH}" # install dependencies RUN apt update -y \ && apt install -y git build-essential libssl-dev pkg-config protobuf-compiler \ && rm -rf /var/lib/apt/lists/* \ && apt clean # install rustup from rustup.rs RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ && rustc --version && cargo --version && protoc --version # copy source code COPY --from=local_src /src /opt/sglang # working directory WORKDIR /opt/sglang/sgl-model-gateway # install maturin and build the wheel with vendored OpenSSL RUN uv pip install maturin \ && cargo clean \ && rm -rf bindings/python/dist/ \ && cd bindings/python \ && ulimit -n 65536 && maturin build --release --features vendored-openssl --out dist \ && rm -rf /root/.cache ######################### ROUTER IMAGE ######################### FROM base AS router-image # Copy the built package from the build image COPY --from=build-image /opt/sglang/sgl-model-gateway/bindings/python/dist/*.whl dist/ # Build the package and install RUN uv pip install --force-reinstall dist/*.whl # Clean up unnecessary files to reduce the image size RUN rm -rf /root/.cache dist/ \ && apt purge -y --auto-remove curl # Set the entrypoint to the main command ENTRYPOINT ["python3", "-m", "sglang_router.launch_router"] ================================================ FILE: docker/k8s-sglang-distributed-sts.yaml ================================================ # Two Nodes Sglang example apiVersion: apps/v1 kind: StatefulSet metadata: name: distributed-sglang spec: replicas: 2 # number of nodes/pods to run distributed sglang selector: matchLabels: app: distributed-sglang serviceName: "" template: metadata: labels: app: distributed-sglang spec: containers: - name: sglang-container image: docker.io/lmsysorg/sglang:latest imagePullPolicy: Always # image may be replaced by official CI versioned image command: - /bin/bash - -c # please modify the sglang serving arguments below, as necessary. # NOTE: the --expert-parallel-size is for MoE model like DeepSeek-R1 args: - | python3 -m sglang.launch_server \ --model /llm-folder \ --dist-init-addr sglang-master-pod:5000 \ --tensor-parallel-size 16 \ --nnodes 2 \ --node-rank $POD_INDEX \ --trust-remote-code \ --host 0.0.0.0 \ --port 8000 \ --enable-metrics \ --expert-parallel-size 16 env: - name: POD_INDEX # reflects the node-rank valueFrom: fieldRef: apiVersion: v1 fieldPath: metadata.labels['apps.kubernetes.io/pod-index'] - name: NCCL_DEBUG value: INFO resources: limits: nvidia.com/gpu: "8" requests: volumeMounts: - mountPath: /dev/shm name: dshm - mountPath: /llm-folder name: llm securityContext: privileged: true # to leverage RDMA/InfiniBand device, co-work with HostNetwork=true hostNetwork: true volumes: - emptyDir: medium: Memory sizeLimit: 10Gi name: dshm - hostPath: path: /llm-folder # replace with PVC or hostPath with your model weights type: DirectoryOrCreate name: llm #- persistentVolumeClaim: # claimName: llm-pvc # name: llm --- apiVersion: v1 kind: Service metadata: name: sglang-master-pod spec: type: ClusterIP selector: app: distributed-sglang apps.kubernetes.io/pod-index: "0" ports: - name: dist-port port: 5000 targetPort: 5000 --- # the serving service apiVersion: v1 kind: Service metadata: name: sglang-serving-on-master spec: type: NodePort selector: app: distributed-sglang apps.kubernetes.io/pod-index: "0" ports: - name: serving port: 8000 targetPort: 8000 - name: metrics port: 8080 targetPort: 8080 ================================================ FILE: docker/k8s-sglang-service.yaml ================================================ apiVersion: v1 kind: PersistentVolumeClaim metadata: name: llama-31-8b-sglang spec: accessModes: - ReadWriteMany resources: requests: storage: 30Gi storageClassName: default # change this to your preferred storage class volumeMode: Filesystem --- apiVersion: node.k8s.io/v1 kind: RuntimeClass metadata: name: nvidia handler: nvidia --- apiVersion: apps/v1 kind: Deployment metadata: name: meta-llama-31-8b-instruct-sglang spec: replicas: 1 strategy: type: Recreate selector: matchLabels: app: meta-llama-31-8b-instruct-sglang template: metadata: labels: app: meta-llama-31-8b-instruct-sglang model: meta-llama-31-8b-instruct engine: sglang spec: restartPolicy: Always runtimeClassName: nvidia containers: - name: meta-llama-31-8b-instruct-sglang image: docker.io/lmsysorg/sglang:latest imagePullPolicy: Always # IfNotPresent or Never ports: - containerPort: 30000 command: ["python3", "-m", "sglang.launch_server"] args: [ "--model-path", "meta-llama/Llama-3.1-8B-Instruct", "--host", "0.0.0.0", "--port", "30000", ] env: - name: HF_TOKEN value: resources: limits: nvidia.com/gpu: 1 cpu: 8 memory: 40Gi requests: cpu: 2 memory: 16Gi nvidia.com/gpu: 1 volumeMounts: - name: shm mountPath: /dev/shm - name: hf-cache mountPath: /root/.cache/huggingface - name: localtime mountPath: /etc/localtime readOnly: true livenessProbe: httpGet: path: /health port: 30000 initialDelaySeconds: 120 periodSeconds: 15 timeoutSeconds: 10 failureThreshold: 3 readinessProbe: httpGet: path: /health_generate port: 30000 initialDelaySeconds: 120 periodSeconds: 15 timeoutSeconds: 10 failureThreshold: 3 successThreshold: 1 volumes: - name: shm emptyDir: medium: Memory sizeLimit: 10Gi - name: hf-cache persistentVolumeClaim: claimName: llama-31-8b-sglang - name: localtime hostPath: path: /etc/localtime type: File --- apiVersion: v1 kind: Service metadata: name: meta-llama-31-8b-instruct-sglang spec: selector: app: meta-llama-31-8b-instruct-sglang ports: - protocol: TCP port: 80 # port on host targetPort: 30000 # port in container type: LoadBalancer # change to ClusterIP if needed ================================================ FILE: docker/npu.Dockerfile ================================================ ARG CANN_VERSION=8.5.0 ARG DEVICE_TYPE=a3 ARG OS=ubuntu22.04 ARG PYTHON_VERSION=py3.11 FROM quay.io/ascend/cann:$CANN_VERSION-$DEVICE_TYPE-$OS-$PYTHON_VERSION # Update pip & apt sources ARG TARGETARCH ARG CANN_VERSION ARG DEVICE_TYPE ARG PIP_INDEX_URL="https://pypi.org/simple/" ARG APTMIRROR="" ARG PYTORCH_VERSION="2.8.0" ARG TORCHVISION_VERSION="0.23.0" ARG PTA_URL_ARM64="https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.8.0/torch_npu-2.8.0.post2-cp311-cp311-manylinux_2_28_aarch64.whl" ARG PTA_URL_AMD64="https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.8.0/torch_npu-2.8.0.post2-cp311-cp311-manylinux_2_28_x86_64.whl" ARG SGLANG_TAG=main ARG ASCEND_CANN_PATH=/usr/local/Ascend/ascend-toolkit ARG SGLANG_KERNEL_NPU_TAG=main ARG PIP_INSTALL="python3 -m pip install --no-cache-dir" ARG DEVICE_TYPE RUN if [ "$TARGETARCH" = "amd64" ]; then \ echo "Using x86_64 dependencies"; \ echo "PTA_URL=$PTA_URL_AMD64" >> /etc/environment_new; \ elif [ "$TARGETARCH" = "arm64" ]; then \ echo "Using aarch64 dependencies"; \ echo "PTA_URL=$PTA_URL_ARM64" >> /etc/environment_new; \ else \ echo "Unsupported TARGETARCH: $TARGETARCH"; exit 1; \ fi WORKDIR /workspace # Define environments ENV DEBIAN_FRONTEND=noninteractive RUN pip config set global.index-url $PIP_INDEX_URL RUN if [ -n "$APTMIRROR" ];then sed -i "s|.*.ubuntu.com|$APTMIRROR|g" /etc/apt/sources.list ;fi # Install development tools and utilities RUN apt-get update -y && apt upgrade -y && apt-get install -y \ unzip \ build-essential \ cmake \ vim \ wget \ curl \ net-tools \ zlib1g-dev \ lld \ clang \ locales \ ccache \ openssl \ libssl-dev \ pkg-config \ ca-certificates \ && rm -rf /var/cache/apt/* \ && rm -rf /var/lib/apt/lists/* \ && update-ca-certificates \ && locale-gen en_US.UTF-8 ENV LANG=en_US.UTF-8 ENV LANGUAGE=en_US:en ENV LC_ALL=en_US.UTF-8 ### Install MemFabric RUN ${PIP_INSTALL} memfabric-hybrid==1.0.5 ### Install SGLang Model Gateway RUN ${PIP_INSTALL} sglang-router ### Install PyTorch and PTA RUN . /etc/environment_new && \ (${PIP_INSTALL} torch==${PYTORCH_VERSION} torchvision==${TORCHVISION_VERSION} --index-url https://download.pytorch.org/whl/cpu) \ && (${PIP_INSTALL} ${PTA_URL}) ## Install triton-ascend RUN (${PIP_INSTALL} pybind11 triton-ascend) # Install SGLang RUN git clone https://github.com/sgl-project/sglang --branch $SGLANG_TAG && \ (cd sglang/python && rm -rf pyproject.toml && mv pyproject_npu.toml pyproject.toml && ${PIP_INSTALL} -v .[all_npu]) && \ rm -rf sglang # Install Deep-ep # pin wheel to 0.45.1 ref: https://github.com/pypa/wheel/issues/662 RUN ${PIP_INSTALL} wheel==0.45.1 pybind11 pyyaml decorator scipy attrs psutil \ && mkdir sgl-kernel-npu \ && cd sgl-kernel-npu \ && wget https://github.com/sgl-project/sgl-kernel-npu/releases/download/${SGLANG_KERNEL_NPU_TAG}/sgl-kernel-npu-${SGLANG_KERNEL_NPU_TAG}-torch2.8.0-py311-cann${CANN_VERSION}-${DEVICE_TYPE}-$(arch).zip \ && unzip sgl-kernel-npu-${SGLANG_KERNEL_NPU_TAG}-torch2.8.0-py311-cann${CANN_VERSION}-${DEVICE_TYPE}-$(arch).zip \ && ${PIP_INSTALL} deep_ep*.whl sgl_kernel_npu*.whl \ && cd .. && rm -rf sgl-kernel-npu \ && cd "$(python3 -m pip show deep-ep | awk '/^Location:/ {print $2}')" && ln -sf deep_ep/deep_ep_cpp*.so CMD ["/bin/bash"] ================================================ FILE: docker/rocm.Dockerfile ================================================ # Usage (to build SGLang ROCm docker image): # docker build --build-arg SGL_BRANCH=v0.5.9 --build-arg GPU_ARCH=gfx942 -t v0.5.9-rocm700-mi30x -f rocm.Dockerfile . # docker build --build-arg SGL_BRANCH=v0.5.9 --build-arg GPU_ARCH=gfx942-rocm720 -t v0.5.9-rocm720-mi30x -f rocm.Dockerfile . # docker build --build-arg SGL_BRANCH=v0.5.9 --build-arg GPU_ARCH=gfx950 -t v0.5.9-rocm700-mi35x -f rocm.Dockerfile . # docker build --build-arg SGL_BRANCH=v0.5.9 --build-arg GPU_ARCH=gfx950-rocm720 -t v0.5.9-rocm720-mi35x -f rocm.Dockerfile . # Usage (to build SGLang ROCm + Mori docker image): # docker build --build-arg SGL_BRANCH=v0.5.9 --build-arg GPU_ARCH=gfx942 --build-arg ENABLE_MORI=1 --build-arg NIC_BACKEND=ainic -t v0.5.9-rocm700-mi30x -f rocm.Dockerfile . # docker build --build-arg SGL_BRANCH=v0.5.9 --build-arg GPU_ARCH=gfx942-rocm720 --build-arg ENABLE_MORI=1 --build-arg NIC_BACKEND=ainic -t v0.5.9-rocm720-mi30x -f rocm.Dockerfile . # docker build --build-arg SGL_BRANCH=v0.5.9 --build-arg GPU_ARCH=gfx950 --build-arg ENABLE_MORI=1 --build-arg NIC_BACKEND=ainic -t v0.5.9-rocm700-mi35x -f rocm.Dockerfile . # docker build --build-arg SGL_BRANCH=v0.5.9 --build-arg GPU_ARCH=gfx950-rocm720 --build-arg ENABLE_MORI=1 --build-arg NIC_BACKEND=ainic -t v0.5.9-rocm720-mi35x -f rocm.Dockerfile . # Default base images ARG BASE_IMAGE_942="rocm/sgl-dev:rocm7-vllm-20250904" ARG BASE_IMAGE_942_ROCM720="rocm/pytorch:rocm7.2_ubuntu22.04_py3.10_pytorch_release_2.9.1" ARG BASE_IMAGE_950="rocm/sgl-dev:rocm7-vllm-20250904" ARG BASE_IMAGE_950_ROCM720="rocm/pytorch:rocm7.2_ubuntu22.04_py3.10_pytorch_release_2.9.1" # This is necessary for scope purpose ARG GPU_ARCH=gfx950 # =============================== # Base image 942 with rocm700 and args FROM $BASE_IMAGE_942 AS gfx942 ENV BUILD_VLLM="0" ENV BUILD_TRITON="0" ENV BUILD_LLVM="0" ENV BUILD_AITER_ALL="1" ENV BUILD_MOONCAKE="1" ENV AITER_COMMIT="v0.1.11.post1" # =============================== # Base image 942 with rocm720 and args FROM $BASE_IMAGE_942_ROCM720 AS gfx942-rocm720 ENV BUILD_VLLM="0" ENV BUILD_TRITON="1" ENV BUILD_LLVM="0" ENV BUILD_AITER_ALL="1" ENV BUILD_MOONCAKE="1" ENV AITER_COMMIT="v0.1.11.post1" # =============================== # Base image 950 and args FROM $BASE_IMAGE_950 AS gfx950 ENV BUILD_VLLM="0" ENV BUILD_TRITON="0" ENV BUILD_LLVM="0" ENV BUILD_AITER_ALL="1" ENV BUILD_MOONCAKE="1" ENV AITER_COMMIT="v0.1.11.post1" # =============================== # Base image 950 with rocm720 and args FROM $BASE_IMAGE_950_ROCM720 AS gfx950-rocm720 ENV BUILD_VLLM="0" ENV BUILD_TRITON="1" ENV BUILD_LLVM="0" ENV BUILD_AITER_ALL="1" ENV BUILD_MOONCAKE="1" ENV AITER_COMMIT="v0.1.11.post1" # =============================== # Chosen arch and args FROM ${GPU_ARCH} # This is necessary for scope purpose, again ARG GPU_ARCH=gfx950 ENV GPU_ARCH_LIST=${GPU_ARCH%-*} ENV PYTORCH_ROCM_ARCH=gfx942;gfx950 ARG SGL_REPO="https://github.com/sgl-project/sglang.git" ARG SGL_DEFAULT="main" ARG SGL_BRANCH=${SGL_DEFAULT} # Version override for setuptools_scm (used in nightly builds) ARG SETUPTOOLS_SCM_PRETEND_VERSION="" ARG TRITON_REPO="https://github.com/triton-lang/triton.git" ARG TRITON_COMMIT="42270451990532c67e69d753fbd026f28fcc4840" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG LLVM_REPO="https://github.com/jrbyrnes/llvm-project.git" ARG LLVM_BRANCH="MainOpSelV2" ARG LLVM_COMMIT="6520ace8227ffe2728148d5f3b9872a870b0a560" ARG MOONCAKE_REPO="https://github.com/kvcache-ai/Mooncake.git" ARG MOONCAKE_COMMIT="b6a841dc78c707ec655a563453277d969fb8f38d" ARG TILELANG_REPO="https://github.com/tile-ai/tilelang.git" ARG TILELANG_COMMIT="ebf4a7cb8881432165ae8760e99d209d905c704a" ARG FHT_REPO="https://github.com/jeffdaily/fast-hadamard-transform.git" ARG FHT_BRANCH="rocm" ARG FHT_COMMIT="46efb7d776d38638fc39f3c803eaee3dd7016bd1" ARG ENABLE_MORI=0 ARG NIC_BACKEND=none ARG MORI_REPO="https://github.com/ROCm/mori.git" ARG MORI_COMMIT="2f88d06aba75400262ca5c1ca5986cf1fdf4cd82" # AMD AINIC apt repo settings ARG AINIC_VERSION=1.117.5 ARG UBUNTU_CODENAME=jammy USER root # Fix hipDeviceGetName returning empty string in ROCm 7.0 docker images. # The ROCm 7.0 base image is missing libdrm-amdgpu-common which provides the # amdgpu.ids device-ID-to-marketing-name mapping file. # ROCm 7.2 base images already ship these packages, so this step is skipped. # See https://github.com/ROCm/ROCm/issues/5992 RUN set -eux; \ case "${GPU_ARCH}" in \ *rocm720*) \ echo "ROCm 7.2 (GPU_ARCH=${GPU_ARCH}): libdrm-amdgpu packages already present, skipping"; \ ;; \ *) \ echo "ROCm 7.0 (GPU_ARCH=${GPU_ARCH}): installing libdrm-amdgpu packages"; \ curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key \ | gpg --dearmor -o /etc/apt/keyrings/amdgpu-graphics.gpg \ && echo 'deb [arch=amd64,i386 signed-by=/etc/apt/keyrings/amdgpu-graphics.gpg] https://repo.radeon.com/graphics/7.0/ubuntu jammy main' \ > /etc/apt/sources.list.d/amdgpu-graphics.list \ && apt-get update \ && apt-get install -y --no-install-recommends \ libdrm-amdgpu-common \ libdrm-amdgpu-amdgpu1 \ libdrm2-amdgpu \ && rm -rf /var/lib/apt/lists/* \ && cp /opt/amdgpu/share/libdrm/amdgpu.ids /usr/share/libdrm/amdgpu.ids; \ ;; \ esac # Install some basic utilities RUN python -m pip install --upgrade pip && pip install setuptools_scm RUN apt-get purge -y sccache; python -m pip uninstall -y sccache; rm -f "$(which sccache)" # Install AMD SMI Python package from ROCm distribution. # The ROCm 7.2 base image (rocm/pytorch) does not pre-install this package. RUN set -eux; \ case "${GPU_ARCH}" in \ *rocm720*) \ echo "ROCm 7.2 flavor detected from GPU_ARCH=${GPU_ARCH}"; \ cd /opt/rocm/share/amd_smi \ && python3 -m pip install --no-cache-dir . \ ;; \ *) \ echo "Not rocm720 (GPU_ARCH=${GPU_ARCH}), skip amdsmi installation"; \ ;; \ esac WORKDIR /sgl-workspace # ----------------------- # llvm RUN if [ "$BUILD_LLVM" = "1" ]; then \ ENV HIP_CLANG_PATH="/sgl-workspace/llvm-project/build/bin/" \ git clone --single-branch ${LLVM_REPO} -b ${LLVM_BRANCH} \ && cd llvm-project \ && git checkout ${LLVM_COMMIT} \ && mkdir build \ && cd build \ && cmake -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=1 -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" -DLLVM_ENABLE_PROJECTS="clang;lld;" -DLLVM_ENABLE_RUNTIMES="compiler-rt" ../llvm \ && make -j$(nproc); \ fi # ----------------------- # AITER # Unset setuptools_scm override so AITER gets its own version (AITER_COMMIT), not SGLang's # (SETUPTOOLS_SCM_PRETEND_VERSION is set later for SGLang nightly builds and would otherwise # leak into AITER's version when AITER uses setuptools_scm) ENV SETUPTOOLS_SCM_PRETEND_VERSION= RUN pip uninstall -y aiter \ && pip install flydsl==0.0.1.dev95158637 \ && pip install psutil pybind11 # Required by AITER setup.py RUN git clone ${AITER_REPO} \ && cd aiter \ && git checkout ${AITER_COMMIT} \ && git submodule update --init --recursive # Hot patches for AITER in v0.1.10.post3 # This is for ROCm 7.2 only, because of the image rebase from vllm # to rocm/pytorch. RUN set -eux; \ case "${GPU_ARCH}" in \ *rocm720*) \ echo "ROCm 7.2 flavor detected from GPU_ARCH=${GPU_ARCH}"; \ cd aiter \ && sed -i '459 s/if.*:/if False:/' aiter/ops/triton/attention/pa_mqa_logits.py; \ ;; \ *) \ echo "Not rocm720 (GPU_ARCH=${GPU_ARCH}), skip patch"; \ ;; \ esac # [WA] from kk-huang # add sed -i '/c1 = torch.empty((M, D, S1 + S3) for aiter triton gemm config issue # the corresponding pr is https://github.com/ROCm/aiter/pull/2173 # it will be removed when server launched issue is fixed by aiter RUN cd aiter \ && echo "[AITER] GPU_ARCH=${GPU_ARCH}" \ && sed -i '/c1 = torch.empty((M, D, S1 + S3), dtype=dtype, device=x.device)/i\ config = dict(config)' aiter/ops/triton/gemm/fused/fused_gemm_afp4wfp4_split_cat.py \ && if [ "$BUILD_AITER_ALL" = "1" ] && [ "$BUILD_LLVM" = "1" ]; then \ sh -c "HIP_CLANG_PATH=/sgl-workspace/llvm-project/build/bin/ PREBUILD_KERNELS=1 GPU_ARCHS=$GPU_ARCH_LIST python setup.py build_ext --inplace" \ && sh -c "HIP_CLANG_PATH=/sgl-workspace/llvm-project/build/bin/ GPU_ARCHS=$GPU_ARCH_LIST pip install -e ."; \ elif [ "$BUILD_AITER_ALL" = "1" ]; then \ sh -c "PREBUILD_KERNELS=1 GPU_ARCHS=$GPU_ARCH_LIST python setup.py build_ext --inplace" \ && sh -c "GPU_ARCHS=$GPU_ARCH_LIST pip install -e ."; \ else \ sh -c "GPU_ARCHS=$GPU_ARCH_LIST pip install -e ."; \ fi \ && echo "export PYTHONPATH=/sgl-workspace/aiter:\${PYTHONPATH}" >> /etc/bash.bashrc # ----------------------- # Build Mooncake ENV PATH=$PATH:/usr/local/go/bin RUN if [ "$BUILD_MOONCAKE" = "1" ]; then \ apt update && apt install -y zip unzip wget && \ apt install -y gcc make libtool autoconf librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool libibverbs-dev rdma-core && \ apt install -y openssh-server openmpi-bin openmpi-common libopenmpi-dev && \ git clone ${MOONCAKE_REPO} && \ cd Mooncake && \ git checkout ${MOONCAKE_COMMIT} && \ git submodule update --init --recursive && \ bash dependencies.sh -y && \ rm -rf /usr/local/go && \ wget https://go.dev/dl/go1.22.2.linux-amd64.tar.gz && \ tar -C /usr/local -xzf go1.22.2.linux-amd64.tar.gz && \ rm go1.22.2.linux-amd64.tar.gz && \ mkdir -p build && \ cd build && \ cmake .. -DUSE_HIP=ON -DUSE_ETCD=ON && \ make -j "$(nproc)" && make install; \ fi # ----------------------- # Build SGLang ARG BUILD_TYPE=all # Set version for setuptools_scm if provided (for nightly builds). Only pass in the SGLang # pip install RUN so it does not affect AITER, sgl-model-gateway, TileLang, FHT, MORI, etc. ARG SETUPTOOLS_SCM_PRETEND_VERSION RUN pip install IPython \ && pip install orjson \ && pip install python-multipart \ && pip install torchao==0.9.0 \ && pip install pybind11 RUN pip uninstall -y sgl_kernel sglang RUN git clone ${SGL_REPO} \ && cd sglang \ && if [ "${SGL_BRANCH}" = ${SGL_DEFAULT} ]; then \ echo "Using ${SGL_DEFAULT}, default branch."; \ git checkout ${SGL_DEFAULT}; \ else \ echo "Using ${SGL_BRANCH} branch."; \ git checkout ${SGL_BRANCH}; \ fi \ && cd sgl-kernel \ && rm -f pyproject.toml \ && mv pyproject_rocm.toml pyproject.toml \ && AMDGPU_TARGET=$GPU_ARCH_LIST python setup_rocm.py install \ && cd .. \ && rm -rf python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml \ && if [ "$BUILD_TYPE" = "srt" ]; then \ export SETUPTOOLS_SCM_PRETEND_VERSION="${SETUPTOOLS_SCM_PRETEND_VERSION}" && python -m pip --no-cache-dir install -e "python[srt_hip,diffusion_hip]"; \ else \ export SETUPTOOLS_SCM_PRETEND_VERSION="${SETUPTOOLS_SCM_PRETEND_VERSION}" && python -m pip --no-cache-dir install -e "python[all_hip]"; \ fi RUN python -m pip cache purge # Copy config files to support MI300X in virtualized environments (MI300X_VF). Symlinks will not be created in image build. RUN find /sgl-workspace/sglang/python/sglang/srt/layers/quantization/configs/ \ /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs/ \ -type f -name '*MI300X*' | xargs -I {} sh -c 'vf_config=$(echo "$1" | sed "s/MI300X/MI300X_VF/"); cp "$1" "$vf_config"' -- {} # Install Rust toolchain for sgl-model-gateway ENV PATH="/root/.cargo/bin:${PATH}" RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ && rustc --version && cargo --version ENV CARGO_BUILD_JOBS=4 # Build and install sgl-model-gateway RUN python3 -m pip install --no-cache-dir maturin \ && cd /sgl-workspace/sglang/sgl-model-gateway/bindings/python \ && ulimit -n 65536 && maturin build --release --features vendored-openssl --out dist \ && python3 -m pip install --force-reinstall dist/*.whl \ && rm -rf /root/.cache # ----------------------- # TileLang ENV DEBIAN_FRONTEND=noninteractive ENV LIBGL_ALWAYS_INDIRECT=1 RUN echo "LC_ALL=en_US.UTF-8" >> /etc/environment RUN /bin/bash -lc 'set -euo pipefail; \ echo "[TileLang] Building TileLang for ${GPU_ARCH}"; \ # System dependencies (NO llvm-dev to avoid llvm-config-16 shadowing) apt-get update && apt-get install -y --no-install-recommends \ build-essential git wget curl ca-certificates gnupg \ libgtest-dev libgmock-dev \ libprotobuf-dev protobuf-compiler libgflags-dev libsqlite3-dev \ python3 python3-dev python3-setuptools python3-pip python3-apt \ gcc libtinfo-dev zlib1g-dev libedit-dev libxml2-dev vim \ cmake ninja-build pkg-config libstdc++6 software-properties-common \ && rm -rf /var/lib/apt/lists/*; \ \ # Prefer the container venv VENV_PY="/opt/venv/bin/python"; \ VENV_PIP="/opt/venv/bin/pip"; \ if [ ! -x "$VENV_PY" ]; then VENV_PY="python3"; fi; \ if [ ! -x "$VENV_PIP" ]; then VENV_PIP="pip3"; fi; \ \ # Build GoogleTest static libs (Ubuntu package ships sources only) cmake -S /usr/src/googletest -B /tmp/build-gtest -DBUILD_GTEST=ON -DBUILD_GMOCK=ON -DCMAKE_BUILD_TYPE=Release && \ cmake --build /tmp/build-gtest -j"$(nproc)" && \ cp -v /tmp/build-gtest/lib/*.a /usr/lib/x86_64-linux-gnu/ && \ rm -rf /tmp/build-gtest; \ \ # Keep setuptools < 80 (compat with base image) "$VENV_PIP" install --upgrade "setuptools>=77.0.3,<80" wheel cmake ninja scikit-build-core && \ "$VENV_PIP" cache purge || true; \ \ # Locate ROCm llvm-config; fallback to installing LLVM 18 if missing LLVM_CONFIG_PATH=""; \ for p in /opt/rocm/llvm/bin/llvm-config /opt/rocm/llvm-*/bin/llvm-config /opt/rocm-*/llvm*/bin/llvm-config; do \ if [ -x "$p" ]; then LLVM_CONFIG_PATH="$p"; break; fi; \ done; \ if [ -z "$LLVM_CONFIG_PATH" ]; then \ echo "[TileLang] ROCm llvm-config not found; installing LLVM 18..."; \ curl -fsSL https://apt.llvm.org/llvm-snapshot.gpg.key | gpg --dearmor -o /etc/apt/keyrings/llvm.gpg; \ echo "deb [signed-by=/etc/apt/keyrings/llvm.gpg] http://apt.llvm.org/jammy/ llvm-toolchain-jammy-18 main" > /etc/apt/sources.list.d/llvm.list; \ apt-get update; \ apt-get install -y --no-install-recommends llvm-18; \ rm -rf /var/lib/apt/lists/*; \ LLVM_CONFIG_PATH="$(command -v llvm-config-18)"; \ if [ -z "$LLVM_CONFIG_PATH" ]; then echo "ERROR: llvm-config-18 not found after install"; exit 1; fi; \ fi; \ echo "[TileLang] Using LLVM_CONFIG at: $LLVM_CONFIG_PATH"; \ export PATH="$(dirname "$LLVM_CONFIG_PATH"):/usr/local/bin:${PATH}"; \ export LLVM_CONFIG="$LLVM_CONFIG_PATH"; \ \ # Optional shim for tools that expect llvm-config-16 mkdir -p /usr/local/bin && \ printf "#!/usr/bin/env bash\nexec \"%s\" \"\$@\"\n" "$LLVM_CONFIG_PATH" > /usr/local/bin/llvm-config-16 && \ chmod +x /usr/local/bin/llvm-config-16; \ \ # TVM Python bits need Cython + z3 before configure. # Pin z3-solver==4.15.4.0: 4.15.4.0 has a manylinux wheel; 4.15.5.0 has no wheel and builds from source (fails: C++20 needs GCC 14+, image has GCC 11). "$VENV_PIP" install --no-cache-dir "cython>=0.29.36,<3.0" "apache-tvm-ffi @ git+https://github.com/apache/tvm-ffi.git@37d0485b2058885bf4e7a486f7d7b2174a8ac1ce" "z3-solver==4.15.4.0"; \ \ # Clone + pin TileLang (bundled TVM), then build git clone --recursive "${TILELANG_REPO}" /opt/tilelang && \ cd /opt/tilelang && \ git fetch --depth=1 origin "${TILELANG_COMMIT}" || true && \ git checkout -f "${TILELANG_COMMIT}" && \ git submodule update --init --recursive && \ export CMAKE_ARGS="-DUSE_CUDA=OFF -DUSE_ROCM=ON -DROCM_PATH=/opt/rocm -DLLVM_CONFIG=${LLVM_CONFIG} -DSKBUILD_SABI_VERSION= ${CMAKE_ARGS:-}" && \ "$VENV_PIP" install -e . -v --no-build-isolation --no-deps; \ if [ -f pyproject.toml ]; then sed -i "/^[[:space:]]*\"torch/d" pyproject.toml || true; fi; \ "$VENV_PIP" cache purge || true; \ "$VENV_PY" -c "import tilelang; print(tilelang.__version__)"' # ----------------------- # Hadamard-transform (HIP build) RUN /bin/bash -lc 'set -euo pipefail; \ git clone --branch "${FHT_BRANCH}" "${FHT_REPO}" fast-hadamard-transform; \ cd fast-hadamard-transform; \ git checkout -f "${FHT_COMMIT}"; \ python setup.py install' # ----------------------- # Python tools RUN python3 -m pip install --no-cache-dir \ py-spy \ pre-commit \ tabulate # ----------------------- # MORI (optional) RUN /bin/bash -lc 'set -euo pipefail; \ if [ "${ENABLE_MORI}" != "1" ]; then \ echo "[MORI] Skipping (ENABLE_MORI=${ENABLE_MORI})"; \ exit 0; \ fi; \ echo "[MORI] Enabling MORI (NIC_BACKEND=${NIC_BACKEND})"; \ \ # Base deps for MORI build apt-get update && apt-get install -y --no-install-recommends \ build-essential \ g++ \ jq \ libopenmpi-dev \ libpci-dev \ initramfs-tools \ && rm -rf /var/lib/apt/lists/*; \ \ # NIC backend deps case "${NIC_BACKEND}" in \ # default: mlx5 none) \ export USE_IONIC="OFF"; \ export USE_BNXT="OFF"; \ ;; \ # AMD NIC ainic) \ export USE_IONIC="ON"; \ export USE_BNXT="OFF"; \ apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https && \ rm -rf /var/lib/apt/lists/* && mkdir -p /etc/apt/keyrings; \ curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \ echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \ > /etc/apt/sources.list.d/amdainic.list; \ apt-get update && apt-get install -y --no-install-recommends \ libionic-dev \ ionic-common \ ; \ rm -rf /var/lib/apt/lists/*; \ ;; \ # TODO: Add Broadcom bnxt packages/repos here later. # bnxt) \ # export USE_IONIC="OFF"; \ # export USE_BNXT="ON"; \ # echo "[MORI] NIC_BACKEND=bnxt: USE_BNXT=ON. Add Broadcom bnxt packages/repos here later."; \ # ;; \ *) \ echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic"; \ exit 2; \ ;; \ esac; \ \ # Build/install MORI export MORI_GPU_ARCHS="${GPU_ARCH_LIST}"; \ echo "[MORI] MORI_GPU_ARCHS=${MORI_GPU_ARCHS} USE_IONIC=${USE_IONIC} USE_BNXT=${USE_BNXT}"; \ rm -rf /sgl-workspace/mori; \ git clone "${MORI_REPO}" /sgl-workspace/mori; \ cd /sgl-workspace/mori; \ git checkout "${MORI_COMMIT}"; \ git submodule update --init --recursive; \ python3 setup.py develop; \ python3 -c "import os, torch; print(os.path.join(os.path.dirname(torch.__file__), \"lib\"))" > /etc/ld.so.conf.d/torch.conf; \ ldconfig; \ echo "export PYTHONPATH=/sgl-workspace/mori:\${PYTHONPATH}" >> /etc/bash.bashrc; \ echo "[MORI] Done."' # ----------------------- # Hot patch: torch-ROCm # The artifact hardcoded the supported triton version to be 3.5.1. # Rewrite the restriction directly. ARG TORCH_ROCM_FILE="torch-2.9.1+rocm7.2.0.lw.git7e1940d4-cp310-cp310-linux_x86_64.whl" RUN mkdir /tmp/whl && cd /tmp/whl \ && export TORCH_ROCM_FILE="${TORCH_ROCM_FILE}" \ && cat > hack.py <<"PY" import zipfile, csv, os, re from pathlib import Path fname = os.environ["TORCH_ROCM_FILE"] in_whl = Path("/") / fname out_whl = Path("/tmp")/ fname work = Path("/tmp/whl") # 1) Extract with zipfile.ZipFile(in_whl, "r") as z: z.extractall(work) # 2) Locate dist-info and patch METADATA (edit this logic to match your exact line) dist_info = next(work.glob("*.dist-info")) meta = dist_info / "METADATA" txt = meta.read_text(encoding="utf-8") # Example: replace one exact requirement form. # Adjust the string to match what you actually see. pat = r"^Requires-Dist:\s*triton==3.5.1[^\s]*;" txt2, n = re.subn(pat, r"triton>=3.5.1;", txt, flags=re.MULTILINE) if txt2 == txt: raise SystemExit("Did not find expected Requires-Dist line to replace in METADATA") meta.write_text(txt2, encoding="utf-8") # 3) Hacky step: blank hash/size columns in RECORD record = dist_info / "RECORD" rows = [] with record.open(newline="", encoding="utf-8") as f: for r in csv.reader(f): if not r: continue # keep filename, blank out hash and size rows.append([r[0], "", ""]) with record.open("w", newline="", encoding="utf-8") as f: csv.writer(f).writerows(rows) # 4) Re-zip as a wheel with zipfile.ZipFile(out_whl, "w", compression=zipfile.ZIP_DEFLATED) as z: for p in work.rglob("*"): if p.is_file(): z.write(p, p.relative_to(work).as_posix()) print("Wrote", out_whl) PY RUN cd /tmp/whl \ && case "${GPU_ARCH}" in \ *rocm720*) \ echo "ROCm 7.2 flavor detected from GPU_ARCH=${GPU_ARCH}"; \ python hack.py \ && python3 -m pip install --force --no-deps /tmp/${TORCH_ROCM_FILE} \ && rm -fr /tmp/whl /tmp/${TORCH_ROCM_FILE} \ ;; \ *) \ echo "Not rocm720 (GPU_ARCH=${GPU_ARCH}), skip patch"; \ ;; \ esac # ----------------------- # Hot patch: Triton # For ROCm 7.2, this custom build breaks pip dependency management, # so future `pip install` will break the ROCm stack. # A workaround for this is to reinstall the default triton # wheel with the `rocm/pytorch` image in the root directory. RUN if [ "$BUILD_TRITON" = "1" ]; then \ pip uninstall -y triton \ && apt install -y cmake \ && git clone ${TRITON_REPO} triton-custom \ && cd triton-custom \ && git checkout ${TRITON_COMMIT} \ && pip install -r python/requirements.txt \ && pip install -e .; \ fi # ----------------------- # Performance environment variable. # Skip CuDNN compatibility check - not applicable for ROCm (uses MIOpen instead) ENV SGLANG_DISABLE_CUDNN_CHECK=1 ENV HIP_FORCE_DEV_KERNARG=1 ENV HSA_NO_SCRATCH_RECLAIM=1 ENV SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 ENV SGLANG_INT4_WEIGHT=0 ENV SGLANG_MOE_PADDING=1 ENV SGLANG_ROCM_DISABLE_LINEARQUANT=0 ENV SGLANG_ROCM_FUSED_DECODE_MLA=1 ENV SGLANG_SET_CPU_AFFINITY=1 ENV SGLANG_USE_AITER=1 ENV SGLANG_USE_ROCM700A=1 ENV NCCL_MIN_NCHANNELS=112 ENV ROCM_QUICK_REDUCE_QUANTIZATION=INT8 ENV TORCHINDUCTOR_MAX_AUTOTUNE=1 ENV TORCHINDUCTOR_MAX_AUTOTUNE_POINTWISE=1 CMD ["/bin/bash"] ================================================ FILE: docker/sagemaker.Dockerfile ================================================ FROM lmsysorg/sglang:latest COPY serve /usr/bin/serve RUN chmod 777 /usr/bin/serve ENTRYPOINT [ "/usr/bin/serve" ] ================================================ FILE: docker/serve ================================================ #!/bin/bash echo "Starting server" PREFIX="SM_SGLANG_" ARG_PREFIX="--" ARGS=() while IFS='=' read -r key value; do arg_name=$(echo "${key#"${PREFIX}"}" | tr '[:upper:]' '[:lower:]' | tr '_' '-') ARGS+=("${ARG_PREFIX}${arg_name}") if [ -n "$value" ]; then ARGS+=("$value") fi done < <(env | grep "^${PREFIX}") # Add default port only if not already set if ! [[ " ${ARGS[@]} " =~ " --port " ]]; then ARGS+=(--port "${SM_SGLANG_PORT:-8080}") fi # Add default host only if not already set if ! [[ " ${ARGS[@]} " =~ " --host " ]]; then ARGS+=(--host "${SM_SGLANG_HOST:-0.0.0.0}") fi # Add default model-path only if not already set if ! [[ " ${ARGS[@]} " =~ " --model-path " ]]; then ARGS+=(--model-path "${SM_SGLANG_MODEL_PATH:-/opt/ml/model}") fi echo "Running command: exec python3 -m sglang.launch_server ${ARGS[@]}" exec python3 -m sglang.launch_server "${ARGS[@]}" ================================================ FILE: docker/xeon.Dockerfile ================================================ FROM ubuntu:24.04 SHELL ["/bin/bash", "-c"] ARG SGLANG_REPO=https://github.com/sgl-project/sglang.git ARG VER_SGLANG=main RUN apt-get update && \ apt-get full-upgrade -y && \ DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \ ca-certificates \ git \ curl \ wget \ vim \ gcc \ g++ \ make \ libsqlite3-dev \ google-perftools \ libtbb-dev \ libnuma-dev \ numactl WORKDIR /opt RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \ source $HOME/.local/bin/env && \ uv venv --python 3.12 RUN echo -e '[[index]]\nname = "torch"\nurl = "https://download.pytorch.org/whl/cpu"\n\n[[index]]\nname = "torchvision"\nurl = "https://download.pytorch.org/whl/cpu"\n\n[[index]]\nname = "torchaudio"\nurl = "https://download.pytorch.org/whl/cpu"\n\n[[index]]\nname = "triton"\nurl = "https://download.pytorch.org/whl/cpu"' > .venv/uv.toml ENV UV_CONFIG_FILE=/opt/.venv/uv.toml WORKDIR /sgl-workspace RUN source $HOME/.local/bin/env && \ source /opt/.venv/bin/activate && \ git clone ${SGLANG_REPO} sglang && \ cd sglang && \ git checkout ${VER_SGLANG} && \ cd python && \ cp pyproject_cpu.toml pyproject.toml && \ uv pip install . && \ cd ../sgl-kernel && \ cp pyproject_cpu.toml pyproject.toml && \ uv pip install . ENV SGLANG_USE_CPU_ENGINE=1 ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4:/usr/lib/x86_64-linux-gnu/libtbbmalloc.so:/opt/.venv/lib/libiomp5.so RUN echo 'source /opt/.venv/bin/activate' >> /root/.bashrc WORKDIR /sgl-workspace/sglang ================================================ FILE: docker/xpu.Dockerfile ================================================ # If the device is Battlemage, we need to set UBUNTU_VERSION to 24.10 # Usage: docker build --build-arg UBUNTU_VERSION=24.04 --build-arg PYTHON_VERSION=3.10 -t sglang:xpu_kernel -f xpu.Dockerfile --no-cache . # Use Intel deep learning essentials base image with Ubuntu 24.04 FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 # Avoid interactive prompts during package install ENV DEBIAN_FRONTEND=noninteractive # Define build arguments ARG PYTHON_VERSION=3.10 ARG SG_LANG_REPO=https://github.com/sgl-project/sglang.git ARG SG_LANG_BRANCH=main ARG SG_LANG_KERNEL_REPO=https://github.com/sgl-project/sgl-kernel-xpu.git ARG SG_LANG_KERNEL_BRANCH=main RUN useradd -m -d /home/sdp -s /bin/bash sdp && \ chown -R sdp:sdp /home/sdp # Switch to non-root user 'sdp' USER sdp # Set HOME and WORKDIR to user's home directory ENV HOME=/home/sdp WORKDIR /home/sdp RUN curl -fsSL -v -o miniforge.sh -O https://github.com/conda-forge/miniforge/releases/download/25.1.1-0/Miniforge3-Linux-x86_64.sh && \ bash miniforge.sh -b -p ./miniforge3 && \ rm miniforge.sh && \ # Initialize conda environment and install pip . ./miniforge3/bin/activate && \ conda create -y -n py${PYTHON_VERSION} python=${PYTHON_VERSION} && \ conda activate py${PYTHON_VERSION} && \ conda install pip && \ # Append environment activation to .bashrc for interactive shells echo ". /home/sdp/miniforge3/bin/activate; conda activate py${PYTHON_VERSION}; . /opt/intel/oneapi/setvars.sh; cd /home/sdp" >> /home/sdp/.bashrc USER root RUN apt-get update && apt install -y intel-ocloc # Switch back to user sdp USER sdp RUN --mount=type=secret,id=github_token \ cd /home/sdp && \ . /home/sdp/miniforge3/bin/activate && \ conda activate py${PYTHON_VERSION} && \ pip3 install torch==2.10.0+xpu torchao torchvision torchaudio triton-xpu==3.6.0 --index-url https://download.pytorch.org/whl/xpu RUN --mount=type=secret,id=github_token \ cd /home/sdp && \ . /home/sdp/miniforge3/bin/activate && \ conda activate py${PYTHON_VERSION} && \ echo "Cloning ${SG_LANG_BRANCH} from ${SG_LANG_REPO}" && \ git clone --branch ${SG_LANG_BRANCH} --single-branch ${SG_LANG_REPO} && \ cd sglang && cd python && \ cp pyproject_xpu.toml pyproject.toml && \ pip install . && \ pip install xgrammar --no-deps && \ pip install msgspec blake3 py-cpuinfo compressed_tensors gguf partial_json_parser einops tabulate --root-user-action=ignore && \ conda install libsqlite=3.48.0 -y && \ # Add environment setup commands to .bashrc again (in case it was overwritten) echo ". /home/sdp/miniforge3/bin/activate; conda activate py${PYTHON_VERSION}; cd /home/sdp" >> /home/sdp/.bashrc # Use bash as default shell with initialization from .bashrc SHELL ["bash", "-c"] # Start an interactive bash shell with all environment set up USER sdp CMD ["bash", "-c", "source /home/sdp/.bashrc && exec bash"] ================================================ FILE: docs/Makefile ================================================ # Minimal Makefile for Sphinx documentation SPHINXOPTS ?= SPHINXBUILD ?= sphinx-build SPHINXAUTOBUILD ?= sphinx-autobuild SOURCEDIR = . BUILDDIR = _build PORT ?= 8003 help: @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) @echo "" @echo "Additional targets:" @echo " serve to build and serve documentation with auto-build and live reload" # Compile Notebook files and record execution time compile: @set -e; \ echo "Starting Notebook compilation..."; \ mkdir -p logs; \ echo "Notebook execution timings:" > logs/timing.log; \ START_TOTAL=$$(date +%s); \ find $(SOURCEDIR) -path "*/_build/*" -prune -o -name "*.ipynb" -print0 | \ parallel -0 -j3 --halt soon,fail=1 ' \ NB_NAME=$$(basename {}); \ START_TIME=$$(date +%s); \ retry --delay=0 --times=2 -- \ jupyter nbconvert --to notebook --execute --inplace "{}" \ --ExecutePreprocessor.timeout=600 \ --ExecutePreprocessor.kernel_name=python3; \ RET_CODE=$$?; \ END_TIME=$$(date +%s); \ ELAPSED_TIME=$$((END_TIME - START_TIME)); \ echo "$${NB_NAME}: $${ELAPSED_TIME}s" >> logs/timing.log; \ exit $$RET_CODE' || exit 1; \ END_TOTAL=$$(date +%s); \ TOTAL_ELAPSED=$$((END_TOTAL - START_TOTAL)); \ echo "---------------------------------" >> logs/timing.log; \ echo "Total execution time: $${TOTAL_ELAPSED}s" >> logs/timing.log; \ echo "All Notebook execution timings:" && cat logs/timing.log # Convert Notebook files to Markdown artifacts (no execution) markdown: @set -e; \ echo "Exporting docs to Markdown..."; \ mkdir -p "$(BUILDDIR)/html/markdown"; \ \ # 1) Copy .md and .rst files as-is; additionally convert .rst -> .md \ find $(SOURCEDIR) -path "*/_build/*" -prune -o \( -name "*.md" -o -name "*.rst" \) -print0 | \ parallel -0 -j3 --halt soon,fail=1 ' \ SRC="{}"; \ REL_DIR=$$(dirname "$$SRC"); \ OUT_DIR="$(BUILDDIR)/html/markdown/$$REL_DIR"; \ mkdir -p "$$OUT_DIR"; \ cp -f "$$SRC" "$$OUT_DIR/"; \ case "$$SRC" in \ *.rst) \ BASE=$$(basename "$$SRC" .rst); \ pandoc -f rst -t gfm "$$SRC" -o "$$OUT_DIR/$$BASE.md" ;; \ esac \ ' || exit 1; \ \ # 2) Convert .ipynb -> .md \ find $(SOURCEDIR) -path "*/_build/*" -prune -o -name "*.ipynb" -print0 | \ parallel -0 -j3 --halt soon,fail=1 ' \ NB_SRC="{}"; \ REL_DIR=$$(dirname "$$NB_SRC"); \ NB_NAME=$$(basename "$$NB_SRC"); \ NB_BASE=$${NB_NAME%.ipynb}; \ OUT_DIR="$(BUILDDIR)/html/markdown/$$REL_DIR"; \ mkdir -p "$$OUT_DIR"; \ jupyter nbconvert --to markdown "$$NB_SRC" \ --output "$$NB_BASE.md" \ --output-dir "$$OUT_DIR" \ >/dev/null; \ ' || exit 1; \ \ echo "Markdown artifacts written to: $(BUILDDIR)/html/markdown" # Serve documentation with auto-build and live reload serve: @echo "Starting auto-build server at http://0.0.0.0:$(PORT)" @$(SPHINXAUTOBUILD) "$(SOURCEDIR)" "$(BUILDDIR)/html" \ --host 0.0.0.0 \ --port $(PORT) \ --watch $(SOURCEDIR) \ --re-ignore ".*\.(ipynb_checkpoints|pyc|pyo|pyd|git)" .PHONY: help Makefile compile clean serve %: Makefile @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) clean: find . -name "*.ipynb" -exec nbstripout {} \; rm -rf $(BUILDDIR) rm -rf logs ================================================ FILE: docs/README.md ================================================ # SGLang Documentation This is the documentation website for the SGLang project (https://github.com/sgl-project/sglang). We recommend new contributors start from writing documentation, which helps you quickly understand SGLang codebase. Most documentation files are located under the `docs/` folder. ## Docs Workflow ### Install Dependency **Linux:** ```bash apt-get update && apt-get install -y pandoc parallel retry pip install -r requirements.txt ``` **macOS:** ```bash brew install pandoc parallel retry pip install -r requirements.txt ``` ### Update Documentation Update your Jupyter notebooks in the appropriate subdirectories under `docs/`. If you add new files, remember to update `index.rst` (or relevant `.rst` files) accordingly. - **`pre-commit run --all-files`** manually runs all configured checks, applying fixes if possible. If it fails the first time, re-run it to ensure lint errors are fully resolved. Make sure your code passes all checks **before** creating a Pull Request. ```bash # 1) Compile all Jupyter notebooks make compile # This step can take a long time (10+ mins). You can consider skipping this step if you can make sure your added files are correct. make html # 2) Compile and Preview documentation locally with auto-build # This will automatically rebuild docs when files change # Open your browser at the displayed port to view the docs bash serve.sh # 2a) Alternative ways to serve documentation # Directly use make serve make serve # With custom port PORT=8080 make serve # 3) Clean notebook outputs # nbstripout removes notebook outputs so your PR stays clean pip install nbstripout find . -name '*.ipynb' -exec nbstripout {} \; # 4) Pre-commit checks and create a PR # After these checks pass, push your changes and open a PR on your branch pre-commit run --all-files ``` ## Documentation Style Guidelines - For common functionalities, we prefer **Jupyter Notebooks** over Markdown so that all examples can be executed and validated by our docs CI pipeline. For complex features (e.g., distributed serving), Markdown is preferred. - Keep in mind the documentation execution time when writing interactive Jupyter notebooks. Each interactive notebook will be run and compiled against every commit to ensure they are runnable, so it is important to apply some tips to reduce the documentation compilation time: - Use small models (e.g., `qwen/qwen2.5-0.5b-instruct`) for most cases to reduce server launch time. - Reuse the launched server as much as possible to reduce server launch time. - Do not use absolute links (e.g., `https://docs.sglang.io/get_started/install.html`). Always prefer relative links (e.g., `../get_started/install.md`). - Follow the existing examples to learn how to launch a server, send a query and other common styles. ## Documentation Build, Deployment, and CI The SGLang documentation pipeline is based on **Sphinx** and supports rendering Jupyter notebooks (`.ipynb`) into HTML/Markdown for web display. Detailed logits can be found in the [Makefile](./Makefile). ### Notebook Execution (`make compile`) The `make compile` target is responsible for executing notebooks before rendering: * Finds all `.ipynb` files under `docs/` (excluding `_build/`) * Executes notebooks in parallel using GNU Parallel, with a relatively small `--mem-fraction-static` * Wraps execution with `retry` to reduce flaky failures * Executes notebooks via `jupyter nbconvert --execute --inplace` * Records execution timing in `logs/timing.log` This step ensures notebooks contain up-to-date outputs with each commit in the main branch before rendering. ### Web Rendering (`make html`) After compilation, Sphinx builds the website: * Reads Markdown, reStructuredText, and Jupyter notebooks * Renders them into HTML pages * Outputs the website into: ``` docs/_build/html/ ``` This directory is the source for online documentation hosting. ### Markdown Export (`make markdown`) To support downstream consumers, we add a **new Makefile target**: ```bash make markdown ``` This target: * Does **not modify** `make compile` * Scans all `.ipynb` files (excluding `_build/`) * Converts notebooks directly to Markdown using `jupyter nbconvert --to markdown` * Writes Markdown artifacts into the existing build directory: ``` docs/_build/html/markdown/.md ``` Example: ``` docs/advanced_features/lora.ipynb → docs/_build/html/markdown/advanced_features/lora.md ``` ### CI Execution In our [CI](https://github.com/sgl-project/sglang/blob/main/.github/workflows/release-docs.yml), the documentation pipeline first gets all the executed results and renders HTML and Markdown by: ```bash make compile # execute notebooks (ensure outputs are up to date) make html # build website as usual make markdown # export markdown artifacts into _build/html/markdown ``` Then, the compiled results are forced pushed to [sgl-project.io](https://github.com/sgl-project/sgl-project.github.io) for rendering. In other words, sgl-project.io is push-only. All the changes of SGLang docs should be made directly in SGLang main repo, then push to the sgl-project.io. ================================================ FILE: docs/_static/css/custom_log.css ================================================ .output_area { color: #615656; } table.autosummary td { width: 50% } img.align-center { display: block; margin-left: auto; margin-right: auto; } .output_area.stderr { color: #d3d3d3 !important; } .output_area.stdout { color: #d3d3d3 !important; } div.output_area.stderr { color: #d3d3d3 !important; } div.output_area.stdout { color: #d3d3d3 !important; } ================================================ FILE: docs/_static/css/readthedocs.css ================================================ table.autosummary td { width: 50% } img.align-center { display: block; margin-left: auto; margin-right: auto; } ================================================ FILE: docs/advanced_features/attention_backend.md ================================================ # Attention Backend SGLang supports a large variety of attention backends. Each of them has different pros and cons. You can test them according to your needs. ```{important} Selecting an optimal attention backend is crucial for maximizing your performance. Different backends excel in various scenarios, so choose based on your model, hardware, and use case. Not all backends are supported on all platforms and model architectures. If you don't specify `--attention-backend`, SGLang makes a best effort to automatically select the most performant backend based on your hardware and model architecture. ``` ## Support Matrix The support matrix is split into two parts: MHA (standard attention) and MLA (multi-head latent attention). For an explanation of the key differences between MHA and MLA, please see the [SGLang documentation on DeepSeek MLA](../basic_usage/deepseek_v3.md#multi-head-latent-attention-mla-throughput-optimizations) and the original [DeepSeek MLA paper](https://arxiv.org/pdf/2405.04434). ### MHA Backends | **Backend** | **Page Size > 1 (native)** | **FP8 KV Cache** | **FP4 KV Cache** | **Spec topk=1** | **Spec topk>1** | **Sliding Window** | **MultiModal** | |---------------------------------|-----------------------------|------------------|-----------------|-----------------|-----------------|--------------------|----------------| | **FlashInfer** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | | **FA3 (FlashAttention 3)** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | | **FA4 (FlashAttention 4)** | 128 | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | | **Triton** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | **Torch Native (SDPA)** | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | | **FlexAttention (PyTorch)** | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | | **TRTLLM MHA** | 16, 32 or 64 | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | | **Dual Chunk FlashAttention** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **AITER (ROCm)** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | | **Wave (ROCm)** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Ascend (NPU)** | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | | **Intel XPU** | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | | **Intel AMX (CPU)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ### MLA Backends | **Backend** | **Native Page Sizes** | **FP8 KV Cache** | **FP4 KV Cache** | **Chunked Prefix Cache** | **Spec topk=1** | **Spec topk>1** | |----------------------------|---------------------------|------------------|------------------|--------------------------|-----------------|-----------------| | **FlashInfer MLA** | 1 | ❌ | ✅ | ✅ | ✅ | ❌ | | **FlashMLA** | 64 | ✅ | ✅ | ✅ | ✅ | ❌ | | **Cutlass MLA** | 128 | ✅ | ✅ | ✅ | ✅ | ❌ | | **TRTLLM MLA (Blackwell)** | 32 or 64 | ✅ | ✅ | ✅ | ✅ | ❌ | | **FA3 (FlashAttention 3)** | n/a | ❌ | ❌ | ✅ | ✅ | ⚠️ (page_size=1 only) | | **Triton** | n/a | ❌ | ❌ | ❌ | ✅ | ⚠️ (page_size=1 only) | | **FA4** | 1 | ❌ | ✅ | ✅ | ❌ | ❌ | | **Ascend MLA (NPU)** | 128 | ❌ | ❌ | ❌ | ❌ | ❌ | ```{note} Multimodal attention is selected by `--mm-attention-backend`. The "MultiModal" column indicates whether a corresponding multimodal implementation exists for that backend family. ``` ```{note} - FlashAttention 4 supports both prefill and decode on SM90 (Hopper) and SM100 (Blackwell). FA4 MLA supports `page_size = 1`; FA4 MHA requires `page_size = 128`. On SM100, this is auto-enforced by the server; on SM90, users must set `--page-size 128` manually. - NSA is specifically designed for [DeepSeek V3.2 DSA](https://lmsys.org/blog/2025-09-29-deepseek-V32/). See the [DSA Attention Backend (NSA)](#dsa-attention-backend-nsa) section and [DeepSeek V3.2 deployment guide](../basic_usage/deepseek_v32.md) for details. ``` ```{warning} **FA4 on Hopper (SM90):** FA4 decode speed decreases as sequence length grows due to lack of SplitKV support. At batch=1 compared to FA3 on H100: ~-10% at 2K tokens, ~-18% at 4K, ~-31% at 8K, ~-49% at 16K. Larger batch sizes reduce the gap (e.g., batch=8: ~-2% at 2K, ~-8% at 4K). Blackwell (SM100) is not affected. ``` ```{note} For the KV4 FA4 scenario, FA4 requires using a different --decode-attention-backend to run. Except for trtllm_mha being incompatible with FA4, all other decode backends behave as shown in the table. ``` ```{tip} Speculative decoding topk: `topk` is the number of draft tokens sampled per step from the draft model. `topk = 1` follows classic EAGLE; `topk > 1` explores multiple branches and requires backend support in both draft and verification paths. ``` ```{note} **Speculative Decoding V2 (Spec V2):** Spec V2 uses overlap scheduling (`SGLANG_ENABLE_SPEC_V2=True`) that benefits various attention backends. Requires `--speculative-eagle-topk 1` and currently applies to EAGLE and EAGLE3. **Verified backends:** TRTLLM MLA, TRTLLM MHA, FA3, Ascend (NPU), Triton. **Limited support:** FlashInfer can run under Spec V2, but its plan stream (used for split-KV optimization) introduces a synchronization point that limits overlap benefits. ``` ```{tip} Page size controls how many tokens are grouped into a KV cache block. For the prefix cache to take effect, the number of tokens must fill at least one complete page. For example, if your prompt is only 32 tokens and `page_size = 64`, it won't fill a complete page and cannot be matched in the prefix cache (pages cannot be padded). With 65 tokens and `page_size = 64`, only the first page of 64 tokens will be cached and matched; the remaining 1 token is discarded. Use `page_size = 1` for maximum prefix reuse (token-level matching). Note that higher page sizes generally improve attention kernel performance, so prefer `page_size > 1` when prefix cache reuse is not critical. ``` Many backends that do not natively operate on pages can emulate `page_size > 1` at the wrapper layer by expanding page tables to per-token indices. The "Page Size > 1 (native)" column indicates true in-kernel paging. Some backends require fixed native page sizes and cannot be reduced/emulated differently: TRTLLM MHA (16/32/64), TRTLLM MLA (32/64), FlashMLA (64), Cutlass MLA (128), Ascend (128). MLA page-size constraints: - FlashInfer MLA: page_size = 1. - FlashMLA: page_size = 64. - Cutlass MLA: page_size = 128. - TRTLLM MLA: page_size ∈ {32, 64}. ### GDN Attention Backends GDN (Gated Delta Network) is a linear attention mechanism with O(n) complexity, used in hybrid models that alternate GDN linear attention layers with standard full attention layers. GDN is **not** selected via `--attention-backend`; it is automatically activated when the model architecture requires it (e.g., Qwen 3.5, Qwen 3 Next, Jet Nemotron, Jet VLM). The GDN linear attention layers have their own kernel backends, selected via `--linear-attn-backend` (default: `triton`). You can override the kernel per phase with `--linear-attn-decode-backend` and `--linear-attn-prefill-backend`. | **Backend** | **Decode** | **Prefill / Extend** | **Spec Decoding (Target Verify)** | |--------------------------|------------|----------------------|-----------------------------------| | **Triton (CUDA)** | ✅ | ✅ | ✅ | | **Triton (AMD/ROCm)** | ✅ | ✅ | ✅ | | **Triton (NPU)** | ✅ | ✅ | ❌ | | **Triton (CPU)** | ✅ | ✅ | ❌ | | **CuTe DSL (CUDA only)**| ✅ | ❌ | ❌ | ```{important} GDN models are hybrid: the full-attention layers still require a standard `--attention-backend`. Platform constraints for the full-attention backend on hybrid GDN models: - **Blackwell (e.g., B200)**: `triton`, `trtllm_mha`, or `fa4` only. - **NPU (Ascend)**: `ascend` only. - **AMD (ROCm)**: `triton` recommended. - **Other CUDA (Hopper, Ampere, etc.)**: auto-selection works; no special constraints. ``` ### DSA Attention Backend (NSA) DSA (Deepseek Sparse Attention) is a native sparse attention mechanism used by [DeepSeek V3.2](https://lmsys.org/blog/2025-09-29-deepseek-V32/). It is activated automatically when the model architecture requires it and is selected via `--attention-backend nsa`. Internally, the NSA backend dispatches to different sub-backends for prefill and decode phases. You can override these with `--nsa-prefill-backend` and `--nsa-decode-backend`: | **Sub-backend** | **Prefill** | **Decode** | **Notes** | |-----------------------|-------------|------------|-----------------------------------------------| | **flashmla_sparse** | ✅ | ✅ | Default prefill on Hopper and Blackwell (bf16) | | **flashmla_kv** | ✅ | ✅ | Default decode for FP8 on Blackwell with DP | | **flashmla_auto** | ✅ | ❌ | Auto-selects flashmla_sparse or flashmla_kv based on kv_cache_dtype | | **fa3** | ✅ | ✅ | Default decode on Hopper (bf16) | | **trtllm** | ✅ | ✅ | Default decode on Blackwell (bf16); default for both on Blackwell without DP | | **tilelang** | ✅ | ✅ | Default on AMD (ROCm) | | **aiter** | ✅ | ✅ | AMD-specific kernel library (requires aiter package) | For deployment examples, see the [DeepSeek V3.2 deployment guide](../basic_usage/deepseek_v32.md). ### Hybrid attention (different backends for prefill vs decode) (Experimental) ```{warning} Hybrid attention is an experimental feature. ``` You can mix-and-match attention backends for prefill and decode. This is useful when one backend excels at prefill and another excels at decode. For the implementation details, please see `python/sglang/srt/layers/attention/hybrid_attn_backend.py`. ```bash # Example: Prefill with FA4, Decode with TRTLLM MLA (Blackwell) python3 -m sglang.launch_server \ --model-path nvidia/DeepSeek-R1-FP4 \ --tp 8 \ --attention-backend trtllm_mla \ --moe-runner-backend flashinfer_trtllm \ --quantization modelopt_fp4 \ --prefill-attention-backend fa4 ``` #### Speculative decoding with hybrid attention Hybrid attention also works with speculative decoding. The backend used for draft decoding and target verification depends on `--speculative-attention-mode`: - `--speculative-attention-mode decode` (recommended): draft/verify use the decode backend. - `--speculative-attention-mode prefill` (default): draft/verify use the prefill backend. Constraints when combining hybrid attention with speculative decoding: - If any attention backend is `trtllm_mha`, speculative decoding supports only `--speculative-eagle-topk 1`. - For paged MHA backends with `--page-size > 1` and `--speculative-eagle-topk > 1`, only `flashinfer` is supported. - CUDA Graph: the decode backend is always captured; the prefill backend is captured only when `--speculative-attention-mode prefill`. ```{tip} If you set only one of `--prefill-attention-backend` or `--decode-attention-backend`, the unspecified phase inherits `--attention-backend`. If both are specified and differ, SGLang automatically enables a hybrid wrapper to dispatch to the chosen backend per phase. ``` ## Attention Backend Selection Guide (CUDA) If the `--attention-backend` argument is not specified, SGLang automatically selects the best backend based on the hardware (CUDA) and model architecture. ### Automatic Selection Logic **1. MHA Models (e.g., Llama, Qwen)** - **Hopper (e.g., H100, H200)**: Defaults to `fa3` if using CUDA 12.3+ and the model configuration is supported. - **Blackwell (e.g., B200)**: Defaults to `trtllm_mha`, unless using speculative decoding with `topk > 1`. - **Other Architectures (Ampere, Ada, etc.)**: Defaults to `flashinfer` if available; otherwise falls back to `triton`. **2. MLA Models (e.g., DeepSeek V3)** - **Hopper**: Defaults to `fa3` (requires CUDA 12.3+). - **Blackwell**: Defaults to `flashinfer`; `trtllm_mla` is auto-selected for DeepSeek V3 models specifically. - **Other Architectures**: Defaults to `triton`. ## User Guide ### Launch Command for Different Attention Backends - FlashInfer (Default for Non-Hopper Machines, e.g., A100, A40) ```bash python3 -m sglang.launch_server \ --model meta-llama/Meta-Llama-3.1-8B-Instruct \ --attention-backend flashinfer python3 -m sglang.launch_server \ --tp 8 \ --model deepseek-ai/DeepSeek-V3 \ --attention-backend flashinfer \ --trust-remote-code ``` - FlashAttention 3 (Default for Hopper Machines, e.g., H100, H200, H20) ```bash python3 -m sglang.launch_server \ --model meta-llama/Meta-Llama-3.1-8B-Instruct \ --attention-backend fa3 python3 -m sglang.launch_server \ --tp 8 \ --model deepseek-ai/DeepSeek-V3 \ --trust-remote-code \ --attention-backend fa3 ``` - Triton ```bash python3 -m sglang.launch_server \ --model meta-llama/Meta-Llama-3.1-8B-Instruct \ --attention-backend triton python3 -m sglang.launch_server \ --tp 8 \ --model deepseek-ai/DeepSeek-V3 \ --attention-backend triton \ --trust-remote-code ``` - FlashMLA ```bash python3 -m sglang.launch_server \ --tp 8 \ --model deepseek-ai/DeepSeek-R1 \ --attention-backend flashmla \ --trust-remote-code python3 -m sglang.launch_server \ --tp 8 \ --model deepseek-ai/DeepSeek-R1 \ --attention-backend flashmla \ --kv-cache-dtype fp8_e4m3 \ --trust-remote-code ``` - TRTLLM MLA (Optimized for Blackwell Architecture, e.g., B200) ```bash python3 -m sglang.launch_server \ --tp 8 \ --model deepseek-ai/DeepSeek-R1 \ --attention-backend trtllm_mla \ --trust-remote-code ``` - TRTLLM MLA with FP8 KV Cache (Higher concurrency, lower memory footprint) ```bash python3 -m sglang.launch_server \ --tp 8 \ --model deepseek-ai/DeepSeek-R1 \ --attention-backend trtllm_mla \ --kv-cache-dtype fp8_e4m3 \ --trust-remote-code ``` - TRTLLM MHA (Optimized for Blackwell Architecture, e.g., B200) ```bash python3 -m sglang.launch_server \ --tp 4 \ --model Qwen/Qwen3.5-35B-A3B-FP8 \ --attention-backend trtllm_mha \ --trust-remote-code ``` - TRTLLM MHA (XQA backend) (Optimized for SM90 and SM120, e.g., H20, H200, 5090) Note that TRTLLM XQA backend only works well for pagesize 64. ```bash python3 -m sglang.launch_server \ --tp 4 \ --model Qwen/Qwen3.5-35B-A3B-FP8 \ --decode-attention-backend trtllm_mha \ --trust-remote-code ``` - FlashAttention 4 (MHA & MLA) ```bash # FA4 for both prefill and decode on SM90/SM100 python3 -m sglang.launch_server \ --model-path Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 \ --attention-backend fa4 \ --page-size 128 \ --trust-remote-code python3 -m sglang.launch_server \ --tp 8 \ --model deepseek-ai/DeepSeek-R1 \ --prefill-attention-backend fa4 \ --trust-remote-code ``` - Cutlass MLA ```bash python3 -m sglang.launch_server \ --tp 8 \ --model deepseek-ai/DeepSeek-R1 \ --attention-backend cutlass_mla \ --trust-remote-code ``` - Ascend ```bash python3 -m sglang.launch_server \ --model meta-llama/Meta-Llama-3.1-8B-Instruct \ --attention-backend ascend ``` - Intel XPU ```bash python3 -m sglang.launch_server \ --model meta-llama/Meta-Llama-3.1-8B-Instruct \ --attention-backend intel_xpu ``` - Wave ```bash python3 -m sglang.launch_server \ --model meta-llama/Meta-Llama-3.1-8B-Instruct \ --attention-backend wave ``` - FlexAttention ```bash python3 -m sglang.launch_server \ --model meta-llama/Meta-Llama-3.1-8B-Instruct \ --attention-backend flex_attention ``` - Dual Chunk FlashAttention ```bash python3 -m sglang.launch_server \ --model Qwen/Qwen2.5-14B-Instruct-1M \ --attention-backend dual_chunk_flash_attn ``` - Torch Native ```bash python3 -m sglang.launch_server \ --model meta-llama/Meta-Llama-3.1-8B-Instruct \ --attention-backend torch_native ``` ## Steps to add a new attention backend To add a new attention backend, you can learn from the existing backends (`python/sglang/srt/layers/attention/triton_backend.py`, `python/sglang/srt/layers/attention/flashattention_backend.py`) and follow the steps below. ```{note} Linear attention kernel backends (GDN, KDA) follow a different pattern. They implement `LinearAttnKernelBase` in `python/sglang/srt/layers/attention/linear/kernels/` and are dispatched by `GDNKernelDispatcher` / `KDAKernelDispatcher` rather than registered via `@register_attention_backend`. ``` 1. Run without cuda graph. Support the two forward functions - forward_extend - Will be used for prefill, prefill with KV cache, and target verification - It will be called once per layer - forward_decode - Will be used for normal decode, and draft decode - It will be called once per layer - init_forward_metadata - Initialize the class and common metadata shared by all layers - Call the plan function for optimizations like split_kv - It will be called once per forward 2. Run with cuda graph. It has two phases (capture and replay) and you need to implement three functions - init_cuda_graph_state - It will be called once during life time - Create all common shared buffers - init_forward_metadata_capture_cuda_graph - It will be called before capturing a cuda graph - It is similar to init_forward_metadata but write the medatada to some pre-defined buffers - init_forward_metadata_replay_cuda_graph - It will be called before replaying a cuda graph - This function is in the critical path and needs to be fast ================================================ FILE: docs/advanced_features/checkpoint_engine.md ================================================ # Checkpoint Engine Integration The SGLang checkpoint engine integration provides an efficient way to load model weights using a distributed checkpoint loading system. This feature significantly reduces model loading time, especially for large models and multi-node setups, by parallelizing the weight loading process across multiple processes and nodes. ## Overview The checkpoint engine integration allows SGLang to: - Load model weights in parallel using multiple processes - Distribute weight loading across multiple nodes to increase effective disk bandwidth - Overlap weight loading with other initialization tasks like CUDA graph capture - Support both single-node and multi-node deployments ## Installation First, install the checkpoint engine package: ```bash pip install 'checkpoint-engine[p2p]' ``` ## Architecture The system consists of two main components: 1. **SGLang Server**: Runs with `--wait-for-initial-weights` flag to wait for weights before becoming ready 2. **Checkpoint Engine Workers**: Separate processes (managed by torchrun) that load and distribute model weights The checkpoint engine uses a parameter server architecture with support for: - **Broadcast mode**: Weights are broadcast from loading processes to inference processes - **P2P mode**: Direct peer-to-peer weight transfer between processes - **All mode**: Combination of both broadcast and P2P methods ## Usage Examples ### Single Node Setup **Terminal 1 - Launch SGLang Server:** ```bash python -m sglang.launch_server \ --model-path Qwen/Qwen3-8B \ --tp 8 \ --load-format dummy \ --wait-for-initial-weights ``` **Terminal 2 - Run Checkpoint Engine:** Using sglang entrypoint: ```bash python -m sglang.srt.checkpoint_engine.update \ --update-method broadcast \ --checkpoint-path /path/to/Qwen/Qwen3-8B/ \ --inference-parallel-size 8 ``` Using torchrun directly: ```bash torchrun --nproc-per-node 8 \ examples/checkpoint_engine/update.py \ --update-method broadcast \ --checkpoint-path /path/to/Qwen/Qwen3-8B/ \ --inference-parallel-size 8 ``` ### Multi-Node Setup (2 Nodes) **Node 0:** Launch SGLang server: ```bash python -m sglang.launch_server \ --model-path Qwen/Qwen3-8B \ --tp 8 \ --load-format dummy \ --wait-for-initial-weights \ --host [IP] ``` Run checkpoint engine: Using sglang entrypoint (recommended): ```bash python -m sglang.srt.checkpoint_engine.update \ --update-method broadcast \ --checkpoint-path /path/to/Qwen/Qwen3-8B/ \ --inference-parallel-size 8 ``` Using torchrun directly: ```bash torchrun --nproc-per-node 8 \ --nnodes 2 \ --node-rank 0 \ --master-addr [IP] \ --master-port 29500 \ examples/checkpoint_engine/update.py \ --update-method broadcast \ --checkpoint-path /path/to/Qwen/Qwen3-8B/ \ --inference-parallel-size 8 ``` **Node 1:** Launch SGLang server: ```bash python -m sglang.launch_server \ --model-path Qwen/Qwen3-8B \ --tp 8 \ --load-format dummy \ --wait-for-initial-weights \ --host [IP] ``` Run checkpoint engine: Using sglang entrypoint (recommended): ```bash python -m sglang.srt.checkpoint_engine.update \ --update-method broadcast \ --checkpoint-path /path/to/Qwen/Qwen3-8B/ \ --inference-parallel-size 8 ``` Using torchrun directly: ```bash torchrun --nproc-per-node 8 \ --nnodes 2 \ --node-rank 1 \ --master-addr [IP] \ --master-port 29500 \ examples/checkpoint_engine/update.py \ --update-method broadcast \ --checkpoint-path /path/to/Qwen/Qwen3-8B/ \ --inference-parallel-size 8 ``` ### Multi-Node Setup with Tensor Parallelism (TP=16) **Node 0:** Launch SGLang server: ```bash python -m sglang.launch_server \ --model-path Qwen/Qwen3-8B \ --tp 8 \ --load-format dummy \ --wait-for-initial-weights \ --host [IP] \ --dist-init-addr [IP]:9120 \ --nnodes 2 \ --node-rank 0 ``` Run checkpoint engine: Using sglang entrypoint (recommended): ```bash python -m sglang.srt.checkpoint_engine.update \ --update-method broadcast \ --checkpoint-path /path/to/Qwen/Qwen3-8B/ \ --inference-parallel-size 16 ``` Using torchrun directly: ```bash torchrun --nproc-per-node 8 \ --nnodes 2 \ --node-rank 0 \ --master-addr [IP] \ --master-port 29500 \ examples/checkpoint_engine/update.py \ --update-method broadcast \ --checkpoint-path /path/to/Qwen/Qwen3-8B/ \ --inference-parallel-size 16 ``` **Node 1:** Launch SGLang server: ```bash python -m sglang.launch_server \ --model-path Qwen/Qwen3-8B \ --tp 8 \ --load-format dummy \ --wait-for-initial-weights \ --host [IP] \ --dist-init-addr [IP]:9120 \ --nnodes 2 \ --node-rank 1 ``` Run checkpoint engine: Using sglang entrypoint (recommended): ```bash python -m sglang.srt.checkpoint_engine.update \ --update-method broadcast \ --checkpoint-path /path/to/Qwen/Qwen3-8B/ \ --inference-parallel-size 16 ``` Using torchrun directly: ```bash torchrun --nproc-per-node 8 \ --nnodes 2 \ --node-rank 1 \ --master-addr [IP] \ --master-port 29500 \ examples/checkpoint_engine/update.py \ --update-method broadcast \ --checkpoint-path /path/to/Qwen/Qwen3-8B/ \ --inference-parallel-size 16 ``` ## Configuration Options ### SGLang Server Options - `--load-format dummy`: Use dummy format for initial loading (allows overlapping with other tasks) - `--wait-for-initial-weights`: Wait for checkpoint engine to provide weights before becoming ready - `--host`: Host address for multi-node setups - `--dist-init-addr`: Distributed initialization address for tensor parallelism ### Checkpoint Engine Options - `--update-method`: Weight update method (`broadcast`, `p2p`, or `all`) - `--checkpoint-path`: Path to model checkpoint directory - `--inference-parallel-size`: Number of inference parallel processes - `--endpoint`: SGLang server endpoint (default: `http://localhost:19730`) - `--checkpoint-name`: Name for the checkpoint (default: `my-checkpoint-iter-0`) - `--save-metas-file`: File to save checkpoint metadata - `--load-metas-file`: File to load checkpoint metadata from - `--uds`: Unix domain socket path for communication - `--weight-version`: Version identifier for weights ## Performance Benefits The checkpoint engine provides significant time savings in two main aspects: 1. **Multi-node Loading**: Each node only loads a portion of weights from disk, effectively increasing disk bandwidth. More participating nodes provide greater acceleration. Preliminary tests show 20-second acceleration when loading DeepSeek-R1 on H20-3e with two nodes. 2. **Single Process Optimization**: Using dummy format allows overlapping disk-to-CPU transfer with CUDA graph capture and other initialization tasks, providing additional time savings. ## Troubleshooting - Ensure checkpoint engine package is installed: `pip install 'checkpoint-engine[p2p]'` - Verify network connectivity between nodes in multi-node setups - Check that the checkpoint path contains valid model files - Monitor logs for connection errors between SGLang server and checkpoint engine - Use `--sleep-time` parameter to add delays if needed for debugging ## References - [Checkpoint Engine Repository](https://github.com/MoonshotAI/checkpoint-engine) ================================================ FILE: docs/advanced_features/cuda_graph_for_multi_modal_encoder.md ================================================ # Cuda Graph for Multi-Modal Encoder in SGLang ## Motivation In multimodal reasoning services, the visual encoder (ViT / Vision Transformer) typically has a few characteristic traits: Many layers, fragmented operators: Each layer includes LN, QKV projections, attention, MLP, residual connections, etc., resulting in extremely frequent kernel launches. Server-side “small batch / low latency” is common: The batch size is very small (sometimes it looks like 1 after “flattening” the batch), so kernel launch overhead accounts for a large portion of end-to-end latency. Input token count (number of patches) varies frequently: Different image/video resolutions and different batch composition lead to different sequence lengths S — and this is precisely the biggest obstacle for CUDA Graph (unstable shapes). The value of CUDA Graph: It captures a long sequence of GPU kernels with fixed shapes and fixed memory addresses into a graph; later, for the same shapes, it can replay the graph directly, dramatically reducing launch overhead and making GPU scheduling more compact. This led us to seek a CUDA Graph enabled feature for ViT in order to improve ViT performance. ## Design and Restrictions The new CUDA Graph enabled ViT logic is built on ViTCudaGraphRunner. This runner captures the "blocks + merger + deepstack merger (optional)" part of a vision transformer into a CUDA graph and replays it for identical shapes. See the following design consideration and restrictions for more details. ### Dynamic inputs to fit static constraints of CUDA Graph Variable sequence length S is very common in ViT. While CUDA Graph requires fixed shapes. The solution is to build a graph cache by S(e.g., graph_key = S). The first time create a new S, and then capture a graph; afterwards, replay it. If there are many distinct S values, we need to increase VRAM usage which is graph-private memory pools for many graphs. ### Stable addresses Everything "parameter-like" becomes a static buffer: - block_input / block_ws / block_output - cu_full_len / cu_window_len and their kk variants - sin_cos_ws In this way to solve the underlying requirement: during replay, not allowed to swap tensors, can only modify tensor contents. ### Attention backend arguments Attention backend arguments are fixed inside the graph: TritonAttn expects [cu_seqlens, cu_seqlens_kk, max_len] FA3 expects [cu_seqlens, max_len] max_len is frozen as an int constant. cu_seqlens is cached into a dict during create_graph(), and its contents are not updated during subsequent replays. For the same graph_key = S, you not only require the input shape to match, but also require the segmentation pattern in cu_seqlens (and window seqlens) to be identical. Otherwise, attention will segment the sequence incorrectly. ### Rotary buffer management The feature reallocates a larger sin_cos_ws when seq_len increases. The max_content_len is used to make sure the maximum size of the allocated rotary buffer. ## Command Example You can enable CUDA Graph for ViT by setting env variable `SGLANG_VIT_ENABLE_CUDA_GRAPH=1`, for example: ``` SGLANG_VIT_ENABLE_CUDA_GRAPH=1 \ python3 -m sglang.launch_server \ --model Qwen/Qwen3-VL-8B-Instruct ``` Or you can run CUDA Graph for ViT together with Piecewise CUDA Graph feature by both setting env variable `SGLANG_VIT_ENABLE_CUDA_GRAPH=1` and setting `--enable-piecewise-cuda-graph`, for example: ``` SGLANG_VIT_ENABLE_CUDA_GRAPH=1 \ python3 -m sglang.launch_server \ --model Qwen/Qwen3-VL-8B-Instruct \ --piecewise-cuda-graph-max-tokens 4096 \ --enable-piecewise-cuda-graph \ --piecewise-cuda-graph-compiler eager ``` ## Known supported models - Qwen2.5-VL (https://github.com/sgl-project/sglang/pull/14422) - Qwen3-VL (https://github.com/sgl-project/sglang/pull/15320) ================================================ FILE: docs/advanced_features/deterministic_inference.md ================================================ # Deterministic Inference ## Why Deterministic Inference Matters Deterministic inference ensures consistent LLM outputs across runs, which is critical for: - **Reinforcement Learning**: Ensures consistent logprobs across runs, reducing stochastic noise and making RL training more stable, reproducible, and debuggable. - **Testing & Debugging**: Enables reproducible validation - **Production**: Improves reliability and user experience Even with `temperature=0`, standard LLM inference can produce different outputs due to dynamic batching and varying reduction orders in GPU kernels. ## The Root Cause of Non-Determinism The main source is **varying batch sizes**. Different batch sizes cause GPU kernels to split reduction operations differently, leading to different addition orders. Due to floating-point non-associativity (`(a + b) + c ≠ a + (b + c)`), this produces different results even for identical inputs. ## SGLang's Solution Building on [Thinking Machines Lab's batch-invariant operators](https://github.com/thinking-machines-lab/batch_invariant_ops), SGLang achieves fully deterministic inference while maintaining compatibility with chunked prefill, CUDA graphs, radix cache, and non-greedy sampling. The development roadmap for deterministic inference features can be found in this [issue](https://github.com/sgl-project/sglang/issues/10278). ### Supported Backends Deterministic inference is only supported with the following three attention backends: **FlashInfer**, **FlashAttention 3 (FA3)**, and **Triton**. The following table shows feature compatibility for deterministic inference across different attention backends: | Attention Backend | CUDA Graph | Chunked Prefill | Radix Cache | Non-greedy Sampling (Temp > 0) | |-------------------|------------|-----------------|-------------|---------------------| | **FlashInfer** | ✅ Yes | ✅ Yes | ❌ No | ✅ Yes | | **FlashAttention 3 (FA3)** | ✅ Yes | ✅ Yes | ✅ Yes | ✅ Yes | | **Triton** | ✅ Yes | ✅ Yes | ✅ Yes | ✅ Yes | ## Usage ### Basic Usage Enable deterministic inference by adding the `--enable-deterministic-inference` flag: ```bash python3 -m sglang.launch_server \ --model-path Qwen/Qwen3-8B \ --attention-backend fa3 \ --enable-deterministic-inference ``` ### Server Arguments | Argument | Type/Default | Description | |----------|--------------|-------------| | `--enable-deterministic-inference` | flag; default: disabled | Enable deterministic inference with batch-invariant operations | | `--attention-backend` | string; default: fa3 | Choose attention backend (flashinfer, fa3, or triton) | ### Example Configurations #### Qwen3-8B ```bash python3 -m sglang.launch_server \ --model-path Qwen/Qwen3-8B \ --attention-backend flashinfer \ --enable-deterministic-inference ``` #### Llama Models ```bash python3 -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --attention-backend fa3 \ --enable-deterministic-inference ``` #### Qwen3-30B-A3B (MoE Model) ```bash python3 -m sglang.launch_server \ --model-path Qwen/Qwen3-30B-A3B \ --attention-backend fa3 \ --enable-deterministic-inference ``` ### Deterministic Inference with Non-Greedy Sampling (Temperature > 0) SGLang supports deterministic inference even with non-greedy sampling by using sampling seeds. This is particularly useful for reinforcement learning scenarios like GRPO (Group Relative Policy Optimization) where you need multiple diverse but reproducible responses. #### Default Behavior By default, SGLang uses a sampling seed of `42` for reproducible sampling: ```python import requests response = requests.post( "http://localhost:30000/generate", json={ "text": "Tell me a joke", "sampling_params": { "temperature": 0.8, # Non-greedy sampling "max_new_tokens": 128, }, }, ) print(response.json()) # This will always produce the same response across runs ``` #### Generating Multiple Reproducible Responses To sample different responses from the same prompt while maintaining reproducibility (e.g., for GRPO training), provide different sampling seeds in your requests: ```python import requests # Prepare a list of sampling seeds for different responses sampling_seeds = [42, 43, 44, 45, 46] responses = [] for seed in sampling_seeds: response = requests.post( "http://localhost:30000/generate", json={ "text": "Tell me a joke", "sampling_params": { "temperature": 0.8, "max_new_tokens": 128, "sampling_seed": seed, # Specify sampling seed }, }, ) responses.append(response.json()) # Each seed will produce a different but reproducible response # Using the same seed will always produce the same response ``` This approach ensures that: - Different seeds produce diverse responses - The same seed always produces the same response across different runs - Results are reproducible for debugging and evaluation ## Verification Run deterministic tests to verify consistent outputs: ```bash # Single test: same prompt, varying batch sizes python3 -m sglang.test.test_deterministic --test-mode single --n-trials 50 # Prefix test: prompts with different prefix lengths python3 -m sglang.test.test_deterministic --test-mode prefix --n-trials 50 # Radix Cache Consistency mode: test radix cache determinism (cached vs uncached prefill) python3 -m sglang.test.test_deterministic --test-mode radix_cache ``` Expected result: All tests should show `Unique samples: 1` (perfectly deterministic). ================================================ FILE: docs/advanced_features/dp_dpa_smg_guide.md ================================================ # DP, DPA and SGLang DP Router This guide explains the difference between Data Parallelism (DP) and Data Parallelism Attention (DPA), how to enable each mode correctly, and how to use the SGLang Model Gateway (SMG) for production-grade DP deployments. ## Data Parallelism (DP) **Data Parallelism (DP)** is the most common parallelism strategy that replicates the entire model across multiple GPU sets and processes different batches of requests in parallel. Each GPU set handles independent requests. With dedicated routing strategies, as we will introduce later, with those proper routing algorithms in SGLang Model Gateway, the throughput of your serving system could be multiplied nearly linearly. ### Key characteristics - Each replica has a full copy of the model - Requests are distributed/scattered across replicas - No inter-replica communication during one request's inference (for simple DP) ## Data Parallelism Attention (DPA) **Data Parallelism Attention (DPA)**, also known as DP Attention, is an advanced parallelism strategy. While DPA provides the most significant benefits for **Multi-Head Latent Attention (MLA)** models (such as DeepSeek, MiniMax, Kimi-K2), it also supports **standard attention models** like Qwen. ### The Problem with Tensor Parallelism for MLA Models The most common parallelism strategy for inference is **Tensor Parallelism (TP)**. However, TP might not be the most efficient strategy for certain models. For example, DeepSeek models use MLA and only have **one KV head**. If we use tensor parallelism on 8 GPUs, it will lead to: - **Duplicated KV cache** across all GPUs - **Unwanted memory usage** that limits batch size - **Reduced throughput** due to memory constraints ### How DPA Works DPA addresses these limitations by applying **data parallelism specifically to the attention component**.
DPA + EP Architecture **Each DP replica:** - Processes different batches independently (can be in different forward modes: prefill, decode, or idle) - Maintains its own KV cache (no duplication) - Enables significantly larger batch sizes due to memory savings **Communication patterns in DPA + EP:** - - **All2All (Dispatch)**: Routes tokens to expert sub-groups based on gating decisions - **All2All (Combine)**: Gathers computed results from experts back to original token positions
### Key benefits of DPA 1. **Significantly reduced KV cache memory**: Each DP replica only stores KV cache for its own batches 2. **Larger batch sizes**: Memory savings enable larger batch sizes 3. **Improved decoding throughput**: Significant throughput gains for MLA-based models 4. **Independent forward modes**: Each DP replica can be in different forward modes (prefill, decode, or idle) and handles its assigned batches independently during attention computation ### DPA with Expert Parallelism for MoE For MoE models like DeepSeek, DPA is **often** paired with Expert Parallelism (EP) for best throughput at scale. However, **DPA does not require EP**: you can enable DPA without EP if your deployment does not need expert sharding. - Distribute 256+ expert weights across GPUs (cannot fit on a single GPU) - Enable efficient all-to-all token routing via DeepEP - Scale to large clusters (up to 5x throughput improvement over vanilla TP) ### Recommended setup for DeepSeek ```bash python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3 \ --tp 8 \ --dp-size 8 \ --ep 8 \ --enable-dp-attention \ --moe-a2a-backend deepep \ --moe-runner-backend deep_gemm ``` > **Note**: `--dp-size` must be explicitly set when using `--enable-dp-attention`. If `dp_size` is 1 (default), DPA will be disabled. For detailed EP configuration (DeepEP, Two-Batch Overlap, EPLB), see [Expert Parallelism](expert_parallelism.md). ### Target Models DPA supports the following model architectures: - **MLA (Multi-Head Latent Attention) models** - where DPA provides the most significant benefits: - DeepSeek family (DeepSeek-V2, DeepSeek-V3, DeepSeek-R1) - MiniMax models - Kimi-K2 - Other models using MLA architecture - **Standard attention models** - also supported: - Qwen models (see [PR #6121](https://github.com/sgl-project/sglang/pull/6121)) For models like Llama, with standard GQA, standard DP, or TP is typically recommended. To enable DPA, add `--enable-dp-attention` to your server launch command. ### Activation Logic DPA is enabled explicitly via server arguments (CLI or config). You must set both `--dp-size` and `--enable-dp-attention`: ```bash python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3 \ --tp 8 \ --dp-size 8 \ --enable-dp-attention ``` **Important**: `--dp-size` must be greater than 1 for DPA to work. When `dp_size == 1` (default), `--enable-dp-attention` is automatically disabled. The constraint `tp_size % dp_size == 0` must also be satisfied. ### Standard DP for MLA models Note that MLA models, of course, also support DP. Suppose you want to enable standard DP for MLA models. First, launch each MLA model's replica independently. You may launch these replicas one by one with DPA enabled. After launching each MLA model's replica, launch an SMG and connect all the replicas to the SMG. A detailed explanation of SMG is as follows. ## Modern Data Parallelism SGLang Model Gateway (SMG) ### Native DP Mode Native DP (built-in Data Parallelism) in SGLang creates multiple worker processes within a single SGLang instance, under the control of `DataParallelController` with the launching parameter of `dp-size`. ```bash # Native DP mode python -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ --dp-size 4 ``` **Limitations:** - Built-in in-process load balancing only (e.g., `round_robin`, `total_requests`, `total_tokens`) - No cache-aware routing - Limited observability and metrics - No fault tolerance or circuit breakers - Not suitable for production workloads ⚠️ Native DP is **highly not recommended for use right now**. It is only used in some ancient/outdated RL frameworks. You can use SGLang Model Gateway (SMG) to power up your data parallelism in any use case. ### SMG-Based DP (Recommended) Starting from September 2024, SGLang Model Gateway, i.e., SMG, formerly named as SGLang DP Router, was built especially as a production-ready DP routing system with Rust. It starts from DP routing, but later we further expanded its scope to coordinate RL, PD Disaggregation, and other scenarios. This doc only discusses SMG's usage in DP routing. For other usage, please refer to [SGLang Model Gateway Documentation](sgl_model_gateway.md). > To achieve the best production-level routing performance and reduce the overhead to an extreme extent, we use Rust to build SMG, but not Python, since Python is never FAST enough. **We strongly recommend using the SGLang Model Gateway (SMG) for production-grade Data Parallelism.** SMG provides significant advantages over native DP mode. ```bash # SMG-based DP mode (Recommended) python -m sglang_router.launch_server \ --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ --dp-size 4 ``` ⚠️ Note that **SMG and Naive DP share the same launching parameter, `--dp-size`**. But the entrypoint of Naive DP is `python -m sglang.launch_server`, and SMG's entrypoint is `python -m sglang_router.launch_server`. **Advantages of SMG-Based DP:** | Feature | Native DP | SMG-Based DP | |---------|-----------|--------------| | **Load Balancing** | Built-in in-process methods | Advanced policies (cache-aware, power-of-two, etc.) | | **Cache Awareness** | ❌ No | ✅ Yes - significantly higher cache hit rate | | **Throughput** | Baseline | Significant improvement | | **Multi-Node Support** | Limited | ✅ Full support | | **Worker Health Monitoring** | Basic | ✅ Circuit breakers, health checks | | **Reliability** | Basic | ✅ Retries, rate limiting, queuing | | **Observability** | Basic metrics | ✅ 40+ Prometheus metrics, OpenTelemetry | | **Hot Worker Add/Remove** | ❌ No | ✅ Yes | ### SMG's Performance The cache-aware routing policy in SMG significantly improves performance for workloads with shared prefixes: | Metric | Without Cache-Aware | With Cache-Aware SMG | |--------|---------------------|----------------------| | Throughput (token/s) | 82,665 | 158,596 (+92%) | | Cache Hit Rate | 20% | 75% (+275%) | *Benchmark from [SGLang v0.4 blog](https://lmsys.org/blog/2024-12-04-sglang-v0-4/), workload with multiple long prefix groups, 8x A100 80GB GPUs, dp-size=8* ### When to Use Each **Use Native DP when:** - ~Never use Native/Naive DP~ - Learning material of DP routing **Use SMG-Based DP when:** - In any case, when you think DP is needed - Production deployments - Multi-node distributed setups - Workloads with shared prefixes (high cache reuse potential) - You need high availability and reliability features - You require detailed observability and metrics - You want to have highly efficient RL rollout systems Note that for RL rollout systems, **there are four crucial reasons that SMG-Based DP is far better than naive DP routing**. Details can be found at [Load Balancing Router in RL](./sglang_for_rl.md#load-balancing-router). ### Quick Start For SMG **Installation** ```bash pip install sglang-router # or pip install "sglang[all]" ``` **Option A: Co-launch Workers and SMG (Simplest)** This is the easiest way to get started - SMG and workers are launched together: ```bash python -m sglang_router.launch_server \ --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ --dp-size 4 \ --host 0.0.0.0 \ --port 30000 ``` **Option B: Separate Launch (Multi-Node)** For distributed deployments across multiple machines: 1. Launch workers on each node ```bash # Node 1 python -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ --port 8000 # Node 2 python -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ --port 8000 ``` 2. Launch SMG pointing to workers ```bash python -m sglang_router.launch_router \ --worker-urls http://node1:8000 http://node2:8000 \ --policy cache_aware \ --host 0.0.0.0 \ --port 30000 ``` **Option C: Dynamic Worker Registration** For elastic deployments where workers can be added/removed dynamically: ```bash # Launch SMG first python -m sglang_router.launch_router \ --policy cache_aware \ --host 0.0.0.0 \ --port 30000 # Register workers dynamically curl -X POST http://localhost:30000/workers \ -H "Content-Type: application/json" \ -d '{"url": "http://worker1:8000"}' curl -X POST http://localhost:30000/workers \ -H "Content-Type: application/json" \ -d '{"url": "http://worker2:8000"}' ``` ### Load Balancing Policies SMG supports multiple load balancing policies: | Policy | Description | Best For | |--------|-------------|----------| | `cache_aware` | Combines cache locality with load balancing | **Recommended for most workloads** | | `round_robin` | Cycles through workers in order | Simple, predictable distribution | | `random` | Random worker selection | Baseline, testing | | `power_of_two` | Samples two workers, picks lighter one | Low latency requirements | **Cache-Aware Policy (Default, Recommended)** The cache-aware policy provides the best performance for most workloads: ```bash python -m sglang_router.launch_router \ --worker-urls http://worker1:8000 http://worker2:8000 \ --policy cache_aware \ --cache-threshold 0.5 \ --balance-abs-threshold 32 \ --balance-rel-threshold 1.5 \ --eviction-interval-secs 120 \ --max-tree-size 67108864 ``` **How it works:** 1. Maintains an approximate radix tree for each worker based on request history 2. Routes requests to workers with the highest prefix match (cache hit) 3. Falls back to shortest-queue routing when load is imbalanced 4. Automatically evicts old entries to prevent memory overflow ### Best Practices 1. **Start with `cache_aware` policy** - It provides the best balance between cache locality and load distribution for most workloads 2. **Use SMG for production** - Prefer `sglang_router.launch_server` over `sglang.launch_server` for better reliability and observability 3. **Enable health checks** - Configure `--router-health-check-interval-secs` to detect and remove unhealthy workers automatically **Recommended command with best practices applied:** ```bash python -m sglang_router.launch_server \ --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ --dp-size 4 \ --router-policy cache_aware \ --router-health-check-interval-secs 30 \ --router-prometheus-port 10001 \ --host 0.0.0.0 \ --port 30000 ``` For advanced configuration (circuit breakers, retries, Prometheus metrics, K8s integration), see [SGLang Model Gateway Documentation](sgl_model_gateway.md). ### Verifying Traffic Distribution After launching SMG, verify that traffic is being distributed correctly: **1. Check worker status:** ```bash curl http://localhost:30000/workers ``` **2. Check load distribution:** ```bash curl http://localhost:30000/get_loads ``` **3. Monitor metrics (if Prometheus enabled):** ```bash # Key metrics to check smg_router_requests_total{model="..."} smg_worker_requests_active{worker="..."} sglang_cache_hit_rate{source="..."} ``` For detailed metrics and monitoring setup, see [SGLang Model Gateway Documentation](sgl_model_gateway.md). ## Reference | Strategy | Use Case | Key Benefit | |----------|----------|-------------| | **Native DP** (`--dp-size`) | Never | Easy to understand, not rust based | | **SMG-Based DP** | **Production (recommended)** | Cache-aware routing, high availability | | **DPA** (`--dp-size N --enable-dp-attention`) | DeepSeek/MLA models | Eliminates KV cache duplication, improved throughput | | **DPA + EP** | DeepSeek MoE models | Significant throughput improvement vs vanilla TP | **Recommended production setup for DeepSeek:** 1. Enable **DPA** for attention layers (`--dp-size 8 --enable-dp-attention`) 2. Enable **EP** for MoE layers (`--ep 8 --moe-a2a-backend deepep`) 3. Use **SMG** with **cache_aware** policy **Related documentation:** - [Expert Parallelism](expert_parallelism.md) - DeepEP, Two-Batch Overlap, EPLB - [SGLang Model Gateway Documentation](sgl_model_gateway.md) - SMG configuration & troubleshooting - [Large-Scale EP Blog](https://lmsys.org/blog/2025-05-05-large-scale-ep/) - 96 GPU deployment guide ================================================ FILE: docs/advanced_features/dp_for_multi_modal_encoder.md ================================================ # DP for Multi-Modal Encoder in SGLang A typical VLM architecture involves two main components: an multi-modal encoder and a text decoder. Most VLMs utilize a Vision Transformer (ViT) as their multi-modal encoder, it is responsible for processing visual data, extracting features (objects, colors, textures, etc.), and transforming them into a format that can be understood by the model. The text decoder is based on LLM. It processes textual data and generates output based on the encoded visual features. However, since the size of ViT is very small compared to language decoders, there is relatively little gain from TP. On the other hand, TP incurs significant communication overhead because of all-reduce being performed after every layer. Placing the ViT in data parallel while keeping the LLM in tensor parallel consistently lowers TTFT and boosts end-to-end throughput. In this hybrid layout, the vision front-end becomes parallel and lightweight, while scarce interconnect bandwidth and collective ops are reserved for the LLM. Data parallelism replicates the entire model across multiple GPU sets and processes different batches of requests in parallel. ## Command Example You can enable batch-level DP by setting `mm-enable-dp-encoder`, for example: ``` python3 -m sglang.launch_server \ --model-path Qwen/Qwen2.5-VL-7B-Instruct \ --tp 2 \ --mm-enable-dp-encoder ``` ## Known supported models - Qwen2.5-VL () - Qwen3-VL () - InternVL () - GLM-4.5V & GLM-4.6V () ================================================ FILE: docs/advanced_features/epd_disaggregation.md ================================================ # EPD Disaggregation ## Why and What is EPD Disaggregation? In modern Vision-Language Model (VLM) inference, request execution naturally decomposes into three distinct stages: Encoder, Prefill, and Decode. The Encoder stage performs vision preprocessing and ViT-based image encoding, which is highly compute-intensive but only required during request initialization. The Prefill stage processes the full multimodal input sequence to initialize the language model’s Key-Value (KV) cache, while the Decode stage is dominated by memory bandwidth and KV cache access for autoregressive token generation. Existing deployments typically colocate these stages within a unified execution engine, or at best apply Prefill–Decode (PD) disaggregation. However, such designs still tightly couple vision encoding with language prefill, leading to inefficient resource utilization, limited scalability for image-heavy workloads, and suboptimal scheduling under load. To address these challenges, we introduce Encoder–Prefill–Decode (EPD) Disaggregation in SGLang. EPD further separates vision encoding from language processing, enabling independent horizontal scaling of encoder servers, improved load balancing for multimodal requests, and seamless integration with existing PD disaggregation to form a fully decoupled three-tier inference architecture. ### Usage You can launch a language-only model using `--language-only`, or an encoder-only model using `--encoder-only`. When launching a language-only model, you must additionally specify the encoder service endpoints via `--encoder-urls`. We support multiple encoder transfer backends, including zmq_to_scheduler, zmq_to_tokenizer, and mooncake (the default is zmq_to_scheduler). The backend can be selected using `--encoder-transfer-backend`. ### Encoder transfer with Mooncake `--encoder-transfer-backend mooncake` controls **how encoder outputs are transferred** between encoder and language/prefill services. It is an encoder transfer option and can be used independently of the global multimodal embedding cache. Example: ```bash # encoder python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --encoder-only \ --encoder-transfer-backend mooncake \ --port 30000 # language-only server python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --language-only \ --encoder-urls http://127.0.0.1:30000 \ --encoder-transfer-backend mooncake \ --port 30002 ``` ### Global multimodal embedding cache with Mooncake SGLang also supports a Mooncake-backed **global multimodal embedding cache** for EPD workloads. When enabled on encoder servers, repeated image inputs can reuse previously computed ViT embeddings across instances instead of running the vision encoder again. This feature is useful when: - the deployment serves repeated or overlapping image inputs, - encoder compute is the bottleneck, and - Mooncake is already available in the cluster. At a high level, the encoder checks whether the image embedding already exists in Mooncake. Cache hits are prefetched from the global store, while misses are encoded normally and inserted into the cache in the background. To enable it: - install and configure Mooncake in the same way as other SGLang Mooncake integrations, - add `--enable-mm-global-cache` on the encoder server. `--enable-mm-global-cache` controls **whether multimodal embeddings are looked up and stored in the global Mooncake cache**. It is separate from `--encoder-transfer-backend`, which only controls encoder output transport. For Mooncake deployment and configuration details, see [HiCache best practices](hicache_best_practices.md#deployment-with-mooncake) and the [Mooncake backend README](../../python/sglang/srt/mem_cache/storage/mooncake_store/README.md). Example: ```bash # Shared Mooncake configuration export MOONCAKE_TE_META_DATA_SERVER="http://127.0.0.1:8080/metadata" export MOONCAKE_MASTER="127.0.0.1:50051" export MOONCAKE_PROTOCOL="rdma" export MOONCAKE_GLOBAL_SEGMENT_SIZE="4gb" # encoder with global multimodal cache enabled python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --encoder-only \ --enable-mm-global-cache \ --port 30000 # language-only server python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --language-only \ --encoder-urls http://127.0.0.1:30000 \ --port 30002 ``` Notes: - This cache is for **multimodal encoder embeddings**, not the language model KV cache. - The feature currently uses Mooncake as the shared backing store. - It can be enabled regardless of which `--encoder-transfer-backend` you use. - It is most relevant for EPD or encoder-disaggregated VLM deployments where the same images are likely to appear across requests or instances. #### Qwen VL - EP Disaggregation ```bash # encoder 0 python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --encoder-only \ --encoder-transfer-backend zmq_to_scheduler \ --port 30000 # encoder 1 python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --encoder-only \ --encoder-transfer-backend zmq_to_scheduler \ --port 30001 # language-only server python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --language-only \ --encoder-urls http://127.0.0.1:30000 http://127.0.0.1:30001 \ --encoder-transfer-backend zmq_to_scheduler \ --port 30002 ``` - EPD Disaggregation ```bash # encoder 0 python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --encoder-only \ --encoder-transfer-backend zmq_to_scheduler \ --port 30000 # encoder 1 python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --encoder-only \ --encoder-transfer-backend zmq_to_scheduler \ --port 30001 # prefill 0 python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --disaggregation-mode prefill \ --language-only \ --encoder-urls http://127.0.0.1:30000 http://127.0.0.1:30001 \ --encoder-transfer-backend zmq_to_scheduler \ --port 30002 # decode 0 python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --disaggregation-mode decode \ --port 30003 # router python -m sglang_router.launch_router \ --pd-disaggregation \ --prefill http://$PREFILL_HOST:30002 \ --decode http://$DECODE_HOST:30003 \ --port 8000 ``` #### gRPC Encoder (EPD) You can run the encoder as a gRPC server while keeping prefill/decode as HTTP. When using gRPC encoders, set `SGLANG_ENCODER_MM_RECEIVER_MODE=grpc` for the prefill process so it uses the gRPC receiver. ```bash # gRPC encoder python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --encoder-only \ --grpc-mode \ --encoder-transfer-backend zmq_to_scheduler \ --port 30000 # prefill (HTTP) - tell it to use gRPC receiver SGLANG_ENCODER_MM_RECEIVER_MODE=grpc \ python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --disaggregation-mode prefill \ --language-only \ --encoder-urls grpc://127.0.0.1:30000 \ --encoder-transfer-backend zmq_to_scheduler \ --port 30002 # decode (HTTP) python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --disaggregation-mode decode \ --port 30003 # router python -m sglang_router.launch_router \ --pd-disaggregation \ --prefill http://$PREFILL_HOST:30002 \ --decode http://$DECODE_HOST:30003 \ --port 8000 ``` ================================================ FILE: docs/advanced_features/expert_parallelism.md ================================================ # Expert Parallelism Expert Parallelism (EP) in SGLang distributes expert weights across multiple devices in Mixture-of-Experts (MoE) models, addressing memory bottlenecks and enabling efficient scaling for high-performance inference. It is particularly vital for serving large-scale MoE models where tokens are dynamically routed to specialized experts across GPUs. By leveraging optimized all-to-all communication and grouped matrix multiplications (GEMMs), EP reduces latency, boosts throughput, and minimizes idle GPU time. SGLang's EP offers strong extensibility through its modular framework, allowing seamless integration of custom kernels, backends, and optimizations without refactoring core logic, supporting diverse hardware and quantization schemes. ## Supported Backends and Selection Guidance SGLang's EP integrates diverse, highly efficient backends for different use cases, allowing fine-grained control over performance trade-offs. Users specify backends via command-line flags: - `--moe-a2a-backend`: Selects the backend for all-to-all communication. - `--moe-runner-backend`: Selects the backend for MoE computation. ### Backends for All-to-All Communication | Backend | Description | Use Cases | |--------------|-----------------------------------------------------------------------------|------------------------------------| | **`none` (default)** | Disables all-to-all for EP. Uses All-Reduce or All-Gather for token dispatch. | Hybrid EP and TP setups. | | `deepep` | DeepEP, a communication library for efficient token shuffling in MoE models. | Large-scale EP deployments. | | `mooncake` | An extension of DeepEP for elastic inference, leveraging RDMA for high-performance data transfers. | Elastic EP serving. | | `nixl` | [NIXL-EP](https://github.com/ai-dynamo/nixl/tree/main/examples/device/ep), an elastic EP communication library built on NVIDIA's [NIXL](https://github.com/ai-dynamo/nixl) framework with native RDMA and NVLink support. | Elastic EP serving with fault tolerance and dynamic scaling. | | `mori` | MORI-EP, AMD's native all-to-all communication implementation optimized for ROCm. | AMD GPU deployments. | | `flashinfer` | Flashinfer implementation of all-to-all. | Large-scale EP deployments. | | `ascend_fuseep` | Ascend NPU native fused all-to-all communication. | Ascend NPU deployments. | DeepEP and Mooncake backends support two modes for token dispatch: `normal` mode (optimized for prefill workloads with high throughput) and `low_latency` mode (optimized for decode workloads with low latency and CUDA Graph compatibility). MORI backend only supports `normal` mode now. NIXL-EP currently operates in low-latency mode with CUDA Graph support. Users are recommended to set `--deepep-mode auto` to enable automatic dispatch mode switching during runtime. Setting `--deepep-mode normal` or `--deepep-mode low_latency` is useful for debugging or development purposes. Currently, DeepEP, Mooncake, NIXL-EP, `ascend_fuseep` and MORI only support cases where `ep_size = tp_size`. For hybrid EP and TP (i.e., `ep_size < tp_size`), only the `none` backend (All-Reduce or All-Gather-based dispatching) is supported. ### Backends for MoE Computation | Backend | Description | Use Cases | |--------------------------|-----------------------------------------------------------------------------|------------------------------------| | **`auto` (default)** | Automatically selects the optimal backend based on model architecture, hardware (e.g., NVIDIA architecture like Ampere, Hopper, Blackwell), quantization scheme (e.g., FP8, FP4), and runtime conditions. | General-purpose deployments; ensures compatibility and performance without user intervention. | | `triton` | Triton-based implementation for grouped GEMMs. To achieve higher performance, it's highly recommended to create [tuned configurations](https://github.com/sgl-project/sglang/blob/main/benchmark/kernels/fused_moe_triton/README.md). | Custom kernel development or scenarios requiring high extensibility with Torch compilation support. | | `deep_gemm` | DeepGEMM backend optimized for MoE matrix multiplications, supporting contiguous layouts for prefill and masked layouts for decode; often JIT-compiled for performance. | Large-scale EP deployments with FP8 block-wise quantization. | | `cutlass` | CUTLASS-based backend for efficient GEMMs. | NVIDIA architectures with CUTLASS support. | | `flashinfer_trtllm` | FlashInfer integrated with TensorRT-LLM for accelerated MoE computations, supporting FP4 communication operators and high-performance GEMMs. | Blackwell with TRT-LLM. | | `flashinfer_trtllm_routed` | FlashInfer integrated with TensorRT-LLM for accelerated routed MoE computations, consuming SGLang-computed top-k expert assignments and weights. | Blackwell with TRT-LLM. | | `flashinfer_cutlass` | FlashInfer combined with CUTLASS for high-performance grouped GEMMs in MoE layers, handling FP4/FP8 quantization efficiently. | Blackwell with FP4/FP8 models. | | `flashinfer_mxfp4` | FlashInfer variant optimized for MXFP4 (mixed FP4) quantization in MoE runners, focusing on memory-efficient low-precision inference. | Low-precision models with MXFP4. | | `flashinfer_cutedsl` | FlashInfer with a custom DSL for flexible and efficient MoE kernel generation, integrated with ModelOpt FP4 quantization. | Low-precision models with NVFP4. | ### Examples Launch with DeepEP and DeepGEMM for DeepSeek-V3: ```bash python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --moe-a2a-backend deepep --moe-runner-backend deep_gemm --tp 8 --ep 8 ``` ## Extensible EP Framework SGLang's EP framework provides modular abstractions for easy integration of custom kernels, backends, and optimizations. It decouples the MoE forward pass into stages (dispatch → pre-permute → core runner → post-permute → combine), enabling seamless extensions without refactoring core logic. ### Framework Overview The framework centers on `FusedMoE` as the unified entry point for a single, extensible structure. Key components include: - **Dispatcher**: Manages dispatch/combine for backends like DeepEP (implements `BaseDispatcher` subclasses). - **MoeRunner**: Orchestrates grouped-GEMM execution via `MoeRunnerCore` implementations (e.g., `TritonRunnerCore`). - **PermuteMethodPool**: Auto-registers layout conversions (e.g., pre/post-permute via `register_pre_permute` and `register_post_permute` for dynamic modes, or `register_fused_func` for static, torch.compile-compatible fused operations). - **TopK Router**: Backend-agnostic expert selection. This design supports multiple backends via `--moe-a2a-backend` and `--moe-runner-backend`, with quantization integrated through a standardized `apply()` method. The computation flow ensures modularity: ``` [input_hidden_states] | v TopK.forward -> select_experts / triton_kernels.routing / bypass | v [TopKOutput] | v FusedMoE.forward -> Dispatcher.dispatch -> DeepEP / bypass | | | v | [DispatchOutput] | | | v | quant_method.apply -> MoeRunner.forward | | | | | v | | pre-permute + grouped_gemm + post-permute | | | | |-------------- | v | [CombineInput] | | | v | Dispatcher.combine -> DeepEP / bypass | | |--------------------- v [final_hidden_states] ``` For details, see the [MoE Refactor Roadmap](https://github.com/sgl-project/sglang/issues/8715). ### Implementing New Backends To add a new backend: 1. For a new all-to-all dispatcher, implement a `BaseDispatcher` subclass with `dispatch` and `combine` methods. 2. For a new MoE runner backend, define a `MoeRunnerCore` subclass for core operations (e.g., grouped GEMMs). 3. Define new input/output formats for the dispatcher or model runner (e.g., `RunnerInput`, `RunnerOutput`). 4. Register permute/unpermute methods to ensure compatibility: - **Fused Mode** (static, torch.compile-compatible): Use `register_fused_func` for end-to-end operations. - **Permute Mode** (dynamic): Register `register_pre_permute` and `register_post_permute` for flexible layouts. See the [MoE Refactor Implementation PR](https://github.com/sgl-project/sglang/pull/9269) for full changes, including type hints and config expansions. ### Examples For an example implementation, see [moe_runner/triton.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/moe/moe_runner/triton.py), which demonstrates Triton-based grouped GEMMs with registered fused and permutation functions. ## Computation and Communication Overlap SGLang's EP employs advanced overlap techniques to hide communication latency behind computation, maximizing GPU utilization in MoE layers. ### Two-Batch Overlap (TBO) TBO splits requests into micro-batches, interleaving attention computation with dispatch/combine operations. Yield points in the execution graph allow pausing for overlaps, increasing overall throughput without peak memory spikes: ```python operations = [ self._forward_attn, YieldOperation(), # Overlap with dispatch of prior micro-batch self._forward_dispatch, self._forward_mlp, YieldOperation(), # Overlap with combine self._forward_combine, ] ``` Users need to specify `--enable-two-batch-overlap` to unlock up to 2x throughput. For details, see the [Large-Scale EP Blog](https://lmsys.org/blog/2025-05-05-large-scale-ep/#two-batch-overlap). ### Single-Batch Overlap (SBO) SGLang introduces a dispatcher-hook system for Single-Batch Overlap (SBO), enabling the overlap of operations within a single batch—such as shared experts computation with communication—while decentralizing logic to enhance modularity. These hooks execute before and after the `dispatch` and `combine` operations without modifying core MoE modules. This design simplifies interfaces, reduces coupling, and improves extensibility. For implementation details and an example of overlapping shared experts with DeepEP's combine operation, refer to [PR #13327](https://github.com/sgl-project/sglang/pull/13327). Users can set `--enable-single-batch-overlap` to enable this feature. ## Workload Balancer SGLang integrates the [Expert Parallelism Load Balancer (EPLB)](https://github.com/deepseek-ai/EPLB) from DeepSeek to address routing imbalances in MoE models. By analyzing expert activation statistics, EPLB computes an optimal expert arrangement, strategically placing or replicating experts to minimize GPU utilization variance, reduce idle cycles, and enhance scalability. To enable EPLB, use the flags `--enable-eplb`. For optimal performance, increase batch sizes to stabilize activation statistics and configure periodic rebalancing (e.g., every 1000 requests) to adapt to evolving workloads. Simulations demonstrate significant improvements in load balancedness (ratio of mean to max computation time), correlating strongly with throughput gains. For more details, refer to the [EPLB Section in the Large-Scale EP Blog](https://lmsys.org/blog/2025-05-05-large-scale-ep/#expert-parallelism-load-balancer) and the [EPLB Repository](https://github.com/deepseek-ai/eplb). ## EP with Spectulative Decoding When utilizing speculative decoding with MTP on MoE architectures, use the `--speculative-moe-runner-backend` and `--speculative-moe-a2a-backend` arguments to customize the MoE layer behavior for the draft model. While they default to the target model’s settings, users can differentiate them for varying precisions between target and draft models. For model like `nvidia/DeepSeek-R1-0528-NVFP4-v2`, the target model uses NVFP4 precision while the draft model uses BF16. To apply `flashinfer_trtllm` kernel for target MoE layer while falling back to triton fused MoE kernel for draft MoE layer, users can set the arguments as follows: ``` ... --moe-runner-backend flashinfer_trtllm \ --speculative-moe-runner-backend triton \ ... ``` ## Ascend NPU Guidance ### Guidance on SGLang configuration in Ascend NPU - `--moe-a2a-backend` only supports `deepep` and `ascend_fuseep` backends, - `deepep`: The mechanism is consistent with the above description. - `ascend_fuseep`: Offer a large fused operator which integrates all operations between dispatch and combine to boost MoE computation. Only used for decode stage in PD Disaggregation Mode. - `--moe-runner-backend` parameter does not need to be configured. - `--deepep-mode`: - In PD mixed mode, please set `--deepep-mode auto`. - In PD Disaggregation Mode, prefill instance sets `--deepep-mode normal`, and decode instance sets `--deepep-mode low_latency`. ### DeepEP Ascend Introduction DeepEP Ascend is the adapted version of the DeepEP communication library for Huawei Ascend NPUs, specifically designed for Mixture-of-Experts (MoE) model Expert Parallelism (EP). It supports the Ant-moving Function (Split the sequence length into rounds for streaming batch transmission) to optimize the buffer size occupied during collective communication in prefill stage, especially for long sequences. Ant-moving Function can be enabled for both the dispatch and combine phases via the following environment variables: - `DEEPEP_NORMAL_LONG_SEQ_PER_ROUND_TOKENS`: Enable ant-moving function in dispatch stage. Indicates the number of tokens transmitted per round on each rank, default 8192. - `DEEPEP_NORMAL_LONG_SEQ_ROUND`: Enable ant-moving function in dispatch stage. Indicates the number of rounds transmitted on each rank, default 1. - `DEEPEP_NORMAL_COMBINE_ENABLE_LONG_SEQ`: Enable ant-moving function in combine stage, default 0 (means disabled). `DEEPEP_NORMAL_LONG_SEQ_PER_ROUND_TOKENS * DEEPEP_NORMAL_LONG_SEQ_ROUND` means input sequence length. When the input sequence length exceeds 8192, it is recommended to enable the ant-moving function in both dispatch and combine phase. The environment variable `HCCL_BUFFSIZE` is used to configure the buffer size (MB) actually allocated. Its calculation formula is as follows: ```angular2html # Enable Ant-moving Function HCCL_BUFFSIZE >= 2 * (102MB + 4MB + DEEPEP_NORMAL_LONG_SEQ_PER_ROUND_TOKENS * (hidden_size + hidden_size + hidden_size) * topk) + PADDING_BUFFSIZE # Disable Ant-moving Function HCCL_BUFFSIZE >= 2 * (102MB + 4MB + TOTAL_SEQ_LEN * (hidden_size + hidden_size) * topk) + PADDING_BUFFSIZE ``` Wherein the parameters are described as follows: - `hidden_size`: hidden size in model config. - `topk`: The number of selected routing experts. - `TOTAL_SEQ_LEN`: input sequence length. - `PADDING_BUFFSIZE`: A value of 20 or greater is recommended. ================================================ FILE: docs/advanced_features/forward_hooks.md ================================================ ## Model Hooks SGLang supports attaching PyTorch forward hooks to specific submodules in the loaded model, configured entirely via `server_args` JSON. This is useful for: * Logging intermediate activations * Debugging model internals * Exporting hidden states to external tooling Hooks are attached once during `ModelRunner.initialize` and run on every forward pass. --- ### Configuration overview Hooks are configured via a `ServerArgs` field: ```python class ServerArgs: ... # For forward hooks forward_hooks: Optional[List[dict[str, Any]]] = None ```` In JSON form, a minimal configuration looks like: ```jsonc { "forward_hooks": [ { "name": "outer_linear_hooks", "target_modules": ["outer.0", "outer.1"], "hook_factory": "my_project.hooks:dummy_hook_factory", "config": { "tag": "outer-layer" } } ] } ``` #### Top-level fields * `forward_hooks` (optional list of objects) Each element is a hook spec describing: * Which modules to target * Which Python factory to call * What configuration to pass into that factory --- ### Hook spec schema Each entry in `forward_hooks` is a JSON object with the following shape: ```jsonc { "name": "optional-descriptive-name", "target_modules": ["pattern1", "pattern2", "..."], "hook_factory": "module.submodule:factory_name", "config": { "...": "arbitrary JSON" } } ``` #### `name` (optional) * Human-readable name for logging. * Used only in log messages such as: ```text Registered forward hook 'outer_linear_hooks' on outer.0 ``` #### `target_modules` (required) * List of **module name patterns** used to match entries in `model.named_modules()`. * Patterns are matched using `fnmatch.fnmatch`, so: * `"outer.0"` matches exactly `"outer.0"`. * `"outer.*"` matches `"outer.0"`, `"outer.1"`, `"outer.inner"`, etc. * `"outer.inner.*"` matches children under `outer.inner`. > If no modules match the given patterns, hook registration does **not** fail. > Instead, SGLang logs a warning and continues: > > ```text > No modules matched hook spec 'name' patterns=['...'] > ``` #### `hook_factory` (required) * String path to the Python factory function that creates the hook. * Supported formats: * `"package.module:factory_name"` * `"package.module.submodule.factory_name"` The path is resolved via: ```python def resolve_callable(path: Optional[str]) -> Optional[Callable]: if path is None: return None if ":" in path: module_name, fn_name = path.split(":", 1) else: parts = path.split(".") if len(parts) < 2: raise ValueError( f"Invalid hook callable path '{path}'. " "Expected 'module.submodule:factory' or 'module.submodule.factory'." ) *mod_parts, fn_name = parts module_name = ".".join(mod_parts) module = importlib.import_module(module_name) try: return getattr(module, fn_name) except AttributeError as e: raise AttributeError( f"Module '{module_name}' has no attribute '{fn_name}' " f"(from hook path '{path}')" ) from e ``` **Failure modes**: * If the path is malformed (not enough dots and no `:`), a `ValueError` is raised at startup. * If the module imports but the attribute is missing, an `AttributeError` is raised with a clear error message. * If the hook factory returns `None`, a warning is logged and no hook is registered for that spec (initialization continues). The first two cause initialization to fail fast with a descriptive error; the last one is non-fatal. #### `config` (optional) * Arbitrary JSON object. * Passed directly to the hook factory as a Python `dict`. * This lets you parameterize hook behavior from config (e.g. tags, log levels, sampling rates, etc.). --- ### Hook lifecycle and behavior Hooks are registered in `ModelRunner.initialize()`: ```python if server_args.forward_hooks: register_forward_hooks(self.model, server_args.forward_hooks) ``` The actual registration logic is implemented by `register_forward_hooks`: ```python def register_forward_hooks(model: nn.Module, hook_specs: List[dict[str, Any]]) -> None: """ hook_specs is a list of dicts from server_args.forward_hooks. Attaches forward hooks to the matching modules. """ name_to_module = dict(model.named_modules()) for spec in hook_specs: spec_name = spec.get("name", "") target_patterns = spec.get("target_modules", []) if not target_patterns: logger.warning( f"Hook spec '{spec_name}' has no 'target_modules', skipping" ) continue hook_factory_path = spec.get("hook_factory") if not hook_factory_path: logger.warning( f"Hook spec '{spec_name}' has no 'hook_factory', skipping" ) continue config = spec.get("config") or {} hook_factory = resolve_callable(hook_factory_path) hook = hook_factory(config) if hook_factory else None if hook is None: logger.warning( f"Hook factory '{hook_factory_path}' for spec '{spec_name}' " "returned None, not registering any hook" ) continue # Resolve patterns like "model.layers.*.mlp" matched = [] for name, module in name_to_module.items(): if any(fnmatch.fnmatch(name, pattern) for pattern in target_patterns): matched.append((name, module)) if not matched: logger.warning( f"No modules matched hook spec '{spec_name}' " f"patterns={target_patterns}" ) continue for module_name, module in matched: if hook: _ = module.register_forward_hook(hook) logger.info( f"Registered forward hook '{spec_name}' " f"on {module_name}" ) ``` Key points: * Hooks are **forward hooks only** (via `module.register_forward_hook`). * They are attached once at initialization. * Hook handles are currently not stored on `ModelRunner` (they cannot be removed later via this API). * Failure to match any modules is non-fatal; a warning is logged instead. * If a hook factory returns `None`, a warning is logged and that spec is skipped. --- ### Writing a hook factory A hook factory is a regular Python function: * Takes a `config: dict` (from JSON) * Returns a forward hook function with signature `(module, inputs, output)` Example: ```python HOOK_CALLS = [] def dummy_hook_factory(config): """Factory that returns a forward hook capturing a tag from config.""" tag = config.get("tag", "default") def hook(module, inputs, output): HOOK_CALLS.append( { "module_type": type(module).__name__, "tag": tag, "shape": tuple(output.shape), } ) return output # must return output if you don’t want to modify the tensor return hook ``` In JSON: ```jsonc { "forward_hooks": [ { "name": "capture_outer", "target_modules": ["outer.0", "outer.1"], "hook_factory": "my_project.hooks:dummy_hook_factory", "config": { "tag": "outer" } } ] } ``` This will: * Resolve `my_project.hooks:dummy_hook_factory` to a Python callable. * Call it with `config = {"tag": "outer"}`. * Use the returned hook for all modules matching `outer.0` and `outer.1`. * Append metadata about each call to `HOOK_CALLS`. --- ### Summary * Define `forward_hooks` as a list of specs in `ServerArgs` to turn on the feature. * Each spec: * selects modules via `target_modules` (glob patterns over `model.named_modules()`), * points to a hook factory via `hook_factory`, * passes arbitrary `config` into that factory. * Hook factories are resolved via `resolve_callable`, which supports `module:factory` and `module.submodule.factory`. * Hooks are standard PyTorch forward hooks, attached once at startup and invoked on every forward pass. * Misconfiguration is either: * **fatal and explicit** (bad path / missing attribute), or * **non-fatal with clear warnings** (no targets matched, or factory returned `None`). ================================================ FILE: docs/advanced_features/hicache.rst ================================================ Hierarchical KV Caching (HiCache) ================================= .. toctree:: :maxdepth: 1 hicache_best_practices.md hicache_design.md hicache_storage_runtime_attach_detach.md ================================================ FILE: docs/advanced_features/hicache_best_practices.md ================================================ # SGLang HiCache Best Practices ## Why HiCache Matters SGLang HiCache extends the traditional RadixAttention with a three-tier hierarchical KV caching system that dramatically improves performance for long-context and multi-turn conversation scenarios. By intelligently managing KV caches across GPU memory, host memory, and external storage backends, HiCache addresses the fundamental capacity bottleneck that limits cache hit rates in conventional systems. ## Configuration Guidelines ## Core HiCache Parameters ```bash # Essential HiCache flags --page-size 64 # Page size for cache management --enable-hierarchical-cache # Enable HiCache --hicache-ratio 2 # Host memory ratio (2x GPU memory) --hicache-size 100 # Host memory size in GBs, will override the above ratio --hicache-io-backend kernel # The I/O backend of moving data between CPU and GPU --hicache-write-policy write_through # Cache write policy from GPU to CPU --hicache-storage-backend # Optional storage backend (e.g., hf3fs, mooncake, etc.) ``` Notes: - Besides configuring `--hicache-storage-backend` at startup, SGLang also supports **runtime attach/detach** of the HiCache storage backend (no restart required) via HTTP admin endpoints. See [Runtime Attach/Detach HiCache Storage Backend](hicache_storage_runtime_attach_detach.md). ## Key Configurations with Storage Backends Enabled ### Memory Layout Optimization ```bash # Page-first: Optimized for I/O efficiency with zero-copy (recommended with kernel backend) --hicache-mem-layout page_first # Page-first-direct: Optimized for direct I/O operations (Compatible with fa3 and same zero-copy performance as page_first) --hicache-mem-layout page_first_direct # Layer-first --hicache-mem-layout layer_first ``` **Layout Compatibility:** - `page_first`: Only compatible with `kernel` I/O backend, automatically switches to `layer_first` with `direct` backend - `page_first_direct`: Specifically designed for `direct` I/O backend with optimized memory organization ### Heterogeneous TP Support (GQA/MHA models) HiCache storage supports cross-cluster KV reuse when different deployments use different TP sizes (for example, `tp=4` and `tp=8`) and share the same storage backend namespace. Use `tp_lcm_size` in `--hicache-storage-backend-extra-config`: ```bash # Example: heterogeneous TP = {4, 8}, so lcm = 8 --hicache-storage-backend-extra-config '{"tp_lcm_size": 8}' ``` Guidelines: - Set `tp_lcm_size` to the least common multiple (LCM) of all TP sizes that will share the same HiCache storage. - For MHA models with Mooncake and `page_head` layout, HiCache will split head shards based on `tp_lcm_size` to make keys reusable across heterogeneous TP deployments. - If all clusters use the same TP size, this option is not needed. ### Prefetch Policies ```bash # Best-effort: Terminate prefetch when needed --hicache-storage-prefetch-policy best_effort # Wait-complete: Ensure complete prefetch, higher cache reuse --hicache-storage-prefetch-policy wait_complete # Timeout: Balance between completion and best-effort --hicache-storage-prefetch-policy timeout ``` ### Integration with PD Disaggregation HiCache works seamlessly with PD Disaggregation. You can choose between two configurations: 1. **Prefill-only HiCache**: Enable HiCache only on Prefill nodes, allowing KV cache sharing among Prefill instances 2. **Full HiCache with async offloading**: Enable HiCache on Prefill nodes and async KV cache offloading on Decode nodes, allowing Prefill nodes to reuse KV caches from Decode nodes in multi-turn dialogue scenarios ```bash # Prefill node with HiCache enabled for cross-prefill sharing (ideal for SystemPrompt scenarios) python3 -m sglang.launch_server \ --model-path /xxx/DeepSeek-R1/ \ --tp 8 \ --host 0.0.0.0 \ --port 10000 \ --enable-metrics \ --enable-cache-report \ --mem-fraction-static 0.85 \ --page-size 64 \ --enable-hierarchical-cache \ --hicache-ratio 2 \ --hicache-size 0 \ --hicache-mem-layout page_first_direct \ --hicache-io-backend direct \ --hicache-write-policy write_through \ --hicache-storage-backend hf3fs \ --hicache-storage-prefetch-policy wait_complete \ --disaggregation-ib-device mlx5_0 \ --disaggregation-mode prefill \ --disaggregation-transfer-backend mooncake # Decode node with async offloading enabled for KV cache reuse by Prefill (ideal for multi-turn conversations) python3 -m sglang.launch_server \ --model-path /xxx/DeepSeek-R1/ \ --tp 8 \ --host 0.0.0.0 \ --port 10000 \ --enable-metrics \ --enable-cache-report \ --page-size 64 \ --hicache-ratio 2 \ --hicache-size 0 \ --hicache-mem-layout page_first_direct \ --hicache-io-backend direct \ --hicache-write-policy write_through \ --hicache-storage-backend hf3fs \ --hicache-storage-prefetch-policy wait_complete \ --disaggregation-decode-enable-offload-kvcache \ # Enable async KV cache offloading in decode node --disaggregation-ib-device mlx5_0 \ --disaggregation-mode decode \ --disaggregation-transfer-backend mooncake ``` ### Deployment with HF3FS Here is an example of deploying DeepSeek-R1 with HiCache-HF3FS. For more details, see the [HF3FS Documentation](../../python/sglang/srt/mem_cache/storage/hf3fs/docs/README.md). ```bash python3 -m sglang.launch_server \ --model-path /xxx/DeepSeek-R1/ \ --log-level info \ --tp 8 \ --host 0.0.0.0 \ --port 10000 \ --enable-metrics \ --enable-cache-report \ --page-size 64 \ --mem-fraction-static 0.85 \ --enable-hierarchical-cache \ --hicache-ratio 2 \ --hicache-size 0 \ --hicache-mem-layout page_first_direct \ --hicache-io-backend direct \ --hicache-write-policy write_through \ --hicache-storage-backend hf3fs \ --hicache-storage-prefetch-policy wait_complete \ ``` ### Deployment with Mooncake Here is an example of deploying Qwen3-235B-A22B-Instruct-2507 with Mooncake. For more details, see the [Mooncake Documentation](../../python/sglang/srt/mem_cache/storage/mooncake_store/README.md). ```bash # Set Mooncake environment variables export MOONCAKE_TE_META_DATA_SERVER="http://127.0.0.1:8080/metadata" export MOONCAKE_GLOBAL_SEGMENT_SIZE=816043786240 export MOONCAKE_PROTOCOL="rdma" export MOONCAKE_DEVICE="$DEVICE_LIST" export MOONCAKE_MASTER=127.0.0.1:50051 # Launch SGLang server with Mooncake backend python3 -m sglang.launch_server \ --model-path $MODEL_PATH \ --tp 8 \ --page-size 64 \ --enable-hierarchical-cache \ --hicache-ratio 2 \ --hicache-mem-layout page_first_direct \ --hicache-io-backend direct \ --hicache-storage-backend mooncake \ --hicache-write-policy write_through \ --hicache-storage-prefetch-policy timeout ``` ## Custom Storage Backend Integration To integrate a new storage backend: 1. **Implement three core methods:** - `get(key)`: Retrieve value by key - `exists(key)`: Check key existence - `set(key, value)`: Store key-value pair 2. **Register your backend:** Add your storage backend to the HiCache [BackendFactory](../../python/sglang/srt/mem_cache/storage/backend_factory.py#L188) The HiCache controller handles all scheduling and synchronization automatically. ### Dynamic Backend Loading Alternatively, you can use dynamic loading to avoid hard-coding your backend in the repository: ```bash python3 -m sglang.launch_server \ --model-path your-model \ --enable-hierarchical-cache \ --hicache-storage-backend dynamic \ --hicache-storage-backend-extra-config '{"backend_name":"custom_backend_name", "module_path": "your_module_path", "class_name": "YourHiCacheClassName"}' ``` **Configuration Parameters:** - `--hicache-storage-backend`: Set to `dynamic` - `--hicache-storage-backend-extra-config`: JSON configuration with: - `backend_name`: Custom backend identifier - `module_path`: Python module path to your implementation - `class_name`: Your HiCache implementation class name - `interface_v1`: 0 (disable) or 1 (enable) to control usage of batch_get_v1 and batch_set_v1 methods ## Community and Support - **GitHub Issues**: Report bugs and feature requests - **Slack Channel**: Join community discussions in #sgl-kv-cache-store - **Documentation**: Refer to storage backend-specific guides --- *This document will be continuously updated based on community feedback and new features. Contributions and suggestions are welcome!* ================================================ FILE: docs/advanced_features/hicache_design.md ================================================ # HiCache System Design and Optimization This document provides a comprehensive overview of SGLang HiCache, covering its system architecture, workflow and key components. It also details configuration parameters, optimization techniques, and integration with various L3 storage backends, serving as a complete reference for users and developers to understand and tune HiCache for efficient LLM inference. ## Why and What is HiCache? In large language model inference, the prefill phase is often time-consuming: input sequences need to be first converted into Key-Value cache (KV cache) for subsequent decoding. When multiple requests share the same prefix, the KV cache for that prefix is identical. By caching and reusing these shared KV caches, redundant computation can be avoided. To address this, SGLang introduced RadixAttention, which leverages idle GPU memory to cache and reuse prefix KV caches, and **HiCache**, which extends this idea to host memory and distributed storage. Inspired by the classic three-level cache design of modern CPUs, HiCache organizes GPU memory as L1, host memory as L2, and distributed storage as L3. This hierarchy enables HiCache to fully exploit the "idle" storage space of GPUs and CPUs, while integrating distributed cache systems such as Mooncake, 3FS, NIXL, and AIBrix KVCache for global KV cache storage and scheduling. As a result, HiCache significantly expands KV cache capacity while maintaining strong read performance—especially in workloads such as multi-QA and long-context inference, where KV cache reuse is frequent. For detailed benchmark results, see [this blog](https://lmsys.org/blog/2025-09-10-sglang-hicache/). ## System Design ### Overall Architecture In many modern CPU architectures, the small but fast L1 and L2 caches are private to each core, enabling rapid access to the hottest data, while the larger L3 cache is shared across all cores to significantly reduce redundancy within the cache. Similarly, in HiCache, the L1 and L2 KV caches are private to each inference instance, whereas the L3 KV cache is shared among all inference instances within the cluster. ### HiRadixTree: Metadata Organization in HiCache For KV cache data organization, HiCache builds upon the RadixTree structure introduced in RadixAttention and proposes HiRadixTree. In RadixAttention, each node of the RadixTree corresponds to the KV cache of a consecutive span of tokens in GPU memory. A path from the root to a leaf node represents the prefix of a request, and shared prefixes across multiple requests can reuse the same nodes, thereby avoiding redundant storage. HiRadixTree extends this idea: each node corresponds to the KV cache of a span of consecutive tokens and records where that KV cache is stored—whether in local GPU memory, CPU memory, L3 storage, or multiple of these tiers. If stored locally, HiRadixTree maintains precise metadata, including the exact storage address. However, to reduce overhead, HiRadixTree does not store or continuously synchronize metadata for L3 KV cache. Instead, when accessing L3 data, it queries the backend in real time to retrieve the necessary metadata, such as whether the data exists and on which server and location it resides. ### Overall Workflow The workflow of HiCache mainly involves three key operations: **local match**, **prefetch** and **write-back**. When the system receives a new request, it first searches the local L1 and L2 caches for matching KV caches. For parts not found locally, it attempts to prefetch from L3. After prefetching, all required KV caches are loaded into the GPU for computation. Once the prefill computation is complete, the system considers storing the newly generated data into L2 or L3. ![HiCache Workflow](https://lmsys.org/images/blog/hicache/hicache_overview.png) ### Local Match Local matching is the first step in HiCache's workflow, where incoming request tokens are matched against the HiRadixTree to locate cached KV data in local memory tiers (L1 GPU memory and L2 host memory). The matching algorithm traverses the HiRadixTree from the root node, following child nodes that match the token sequence prefix. At each node, the incoming token sequence is compared with the node’s stored token sequence. When `page_size > 1`, matching is performed at the page granularity to optimize memory access patterns. If a match terminates within a node’s stored sequence, the node is automatically split to create an exact boundary, improving the efficiency of future matches. The algorithm returns a continuous prefix of the request, with the first part residing in L1 and the latter part in L2. Since the process only requires traversing the local HiRadixTree and does not involve any actual data copying, local matching is extremely fast. ### Prefetch from L3 Data prefetching is one of HiCache’s core optimization techniques, designed to proactively load KV caches from L3 storage into local L2 memory, thereby reducing access latency during subsequent operations. **Prefetch Trigger Conditions**: After local matching, for the parts not found in L1 or L2, the system queries L3 to retrieve metadata for the next continuous matching KV caches. If the length of hit cache in L3 exceeds a threshold (default: 256 tokens, configurable), a prefetch operation is triggered. **Prefetch Strategies**: HiCache provides three different prefetch termination strategies to address different scenario needs: - **best_effort**: Terminates immediately when GPU can execute prefill computation, with no waiting time, suitable for scenarios extremely sensitive to latency. - **wait_complete**: Must wait for all prefetch operations to complete, suitable for scenarios requiring high cache hit rates. - **timeout**: Terminates after specified time or when complete, balancing latency and cache hit rate needs. After prefetching stops, the data already fetched is used together with the local data for the prefill computation. For **timeout** strategy, HiCache introduces two configuration parameters to support fine-grained control over prefetch timeout conditions: * `prefetch_timeout_base`: the base timeout, representing overhead unrelated to the number of tokens (e.g., scheduling and synchronization). * `prefetch_timeout_per_ki_token`: the incremental timeout per thousand tokens. The timeout is computed as: ``` timeout = prefetch_timeout_base + prefetch_timeout_per_ki_token * num_token_to_fetch / 1024 ``` ### Data Write-back The write-back mechanism is responsible for moving frequently accessed KV caches from L1 to L2 and L3, enabling larger and longer-term storage as well as cache sharing across instances. **Configurable Write-back Policies**: HiCache supports three write-back strategies: * **write_through**: Every access is immediately written back to the next level. When bandwidth is sufficient, this strategy provides the strongest caching benefit. * **write_through_selective**: Data is written back only after the access frequency exceeds a threshold. This strategy backs up only hot data, reducing I/O overhead. * **write_back**: Data is written back to the next level only when it is evicted from the upper level. This strategy alleviates storage pressure and is suitable for scenarios where storage capacity is limited but memory utilization must be maximized. **Cross-instance Sharing**: When data is written back from L2 to L3, only data not already present in L3 is transferred. KV caches stored in L3 can then be shared across all SGLang instances in the cluster (depending on the L3 backend implementation), significantly improving cache hit rates within the same memory budget. ### Multi-Rank Synchronization During multi-GPU parallel computation, such as tensor parallelism (TP), HiCache must ensure consistent states across different ranks. Therefore, critical computation steps require the use of `all_reduce` for state synchronization. For example, during prefetching, `all_reduce(op=min)` is used to ensure that all ranks obtain the same number of L3 hits, preventing inconsistent judgments about whether the prefetch threshold has been reached. Similarly, after prefetching completes or terminates, `all_reduce(op=min)` is again required to guarantee consensus among ranks on the prefix length of the successfully retrieved KV cache. ### Data Transfer Optimization **Zero-Copy Data Transfers**: Both prefetching and write-back involve substantial data movement. Minimizing the number of data copies can significantly improve system performance. HiCache supports passing memory addresses and sizes directly when transferring data from L2 memory to an L3 backend. **“Batch-Oriented” Data Organization**: The granularity of data reads and writes has a major impact on performance. To address this, HiCache L3 stores and transfers KV cache data at the granularity of **pages** and supports different data layouts beyond the existing `layer first` scheme, including `page first` and `page first direct`. Under the `page first` and `page first direct` layouts, all KV cache data belonging to the same page is placed in contiguous memory, allowing it to be passed as a single object to L3 using zero-copy transfers. ![HiCache L2 MEM layout](https://lmsys.org/images/blog/hicache/hicache_layout.png) However, because GPU KV computation is naturally performed layer by layer, the GPU inherently operates in a `layer first` layout. When transferring `page first` data from L2 to the GPU, data must be transferred at the granularity of one token per layer. The `page first direct` layout mitigates this issue by grouping together all tokens of a given layer within a page, allowing transfers from L2 to GPU to be aggregated at the page-layer level. **CPU-to-GPU Transfer Optimizations**: In HiCache, moving data from CPU memory to GPU is as performance-critical as prefetching data from L3 to L2. HiCache employs several optimizations for this process: * **Compute-Transfer Overlap**: During the prefill phase, when transferring data from CPU to GPU, HiCache overlaps layers by concurrently loading the KV cache of layer N+1 while computing layer N. This effectively hides data transfer latency. * **GPU-assisted I/O Kernels**: On top of `cudaMemcpyAsync`, HiCache implements a set of GPU-assisted I/O kernels specifically optimized for KV cache transfers between CPU and GPU. Compared to the baseline approach, these kernels achieve up to 3x higher transfer speed. **Write-back Optimization for MLA**: For MHA (Multi-Head Attention) models under multi-TP, each rank holds `1/tp_size` of a token’s KV data. In contrast, for MLA (Multi-Layer Attention) models, all ranks hold the complete and identical KV data for each token. HiCache includes a dedicated optimization for MLA: only one rank initiates the write-back operation, ensuring that data is not redundantly stored across ranks. ### Integration with PD-Disaggregation Deployment Mode SGLang supports a PD (Prefill-Decode) disaggregation deployment mode through the Mooncake TransferEngine (for details, see [this doc](https://docs.sglang.io/advanced_features/pd_disaggregation.html)). In the PD-disaggregation deployment mode, HiCache can be enabled on both the prefill nodes and decode nodes to optimize prefill performance. If enabled on decode nodes, the decode output will also be written back to L3. ### Unified Interfaces and Rich L3 Storage Backends HiCache encapsulates all read, write, and query operations on L3 backends within the `class HiCacheStorage(ABC)`, exposing a set of simple and consistent interfaces. This design supports a wide range of L3 storage backends and allows users to select the one that best fits their specific use cases. - **Mooncake**: Mooncake is a high-performance caching system for LLM inference that leverages RDMA and multi-NIC resources to enable zero-copy, ultra-fast data transfers. Try Mooncake [here](https://github.com/sgl-project/sglang/tree/main/python/sglang/srt/mem_cache/storage/mooncake_store). - **DeepSeek 3FS (HF3FS)**: HF3FS is a Kubernetes-native distributed storage solution with operator-based deployment. Try HF3FS [here](https://github.com/sgl-project/sglang/tree/main/python/sglang/srt/mem_cache/storage/hf3fs). - **NIXL**: NIXL provides a unified API for accessing various storage plugins, including but not limited to DeepSeek's 3FS, GPU Direct Storage (GDS) and Amazon S3-compatible object storage. Try NIXL [here](https://github.com/sgl-project/sglang/tree/main/python/sglang/srt/mem_cache/storage/nixl). - **AIBrix KVCache**: AIBrix KVCache is a production-ready KVCache Offloading Framework, which enables efficient memory tiering and low-overhead cross-engine reuse. Try AIBrix KVCache [here](https://github.com/sgl-project/sglang/tree/main/python/sglang/srt/mem_cache/storage/aibrix_kvcache). - **HiCacheFile**: A simple file-based storage backend for demonstration purposes. Specifically, **LMCache**, an efficient KV cache layer for enterprise-scale LLM inference, provides an alternative solution to HiCache. Try LMCache [here](https://github.com/sgl-project/sglang/tree/main/python/sglang/srt/mem_cache/storage/lmcache). ## Related Parameters - **`--enable-hierarchical-cache`**: Enable hierarchical cache functionality. This is required to use HiCache. - **`--hicache-ratio HICACHE_RATIO`**: The ratio of the size of host KV cache memory pool to the size of device pool. For example, a value of 2 means the host memory pool is twice as large as the device memory pool. The value of this parameter must be greater than 1, as the current implementation requires the host memory allocated for the KV cache to be larger than the device memory allocated for the KV cache. - **`--hicache-size HICACHE_SIZE`**: The size of host KV cache memory pool in gigabytes. This parameter overrides `hicache-ratio` if set. For example, `--hicache-size 30` allocates 30GB (1GB = 1e9 bytes) for the host memory pool **for each rank**. If there are 8 ranks, then the total memory size is 240GB. Just like `hicache-ratio`, the value of this parameter must be larger than the size of device memory allocated for KV cache. **Note**: `--hicache-ratio` and `--hicache-size` are two critical parameters. In general, a larger HiCache size leads to a higher cache hit rate, which improves prefill performance. However, the relationship between cache size and hit rate is not linear. Once most reusable KV data—especially hot tokens—are already cached, further increasing the size may yield only marginal performance gains. Users can set these parameters based on their workload characteristics and performance requirements. - **`--page-size PAGE_SIZE`**: The number of tokens per page. This parameter determines the granularity of KV cache storage and retrieval. Larger page sizes reduce metadata overhead and improve I/O efficiency for storage backends, but may lower the cache hit rate when only part of a page matches the stored KV cache. For workloads with long common prefixes, larger pages can improve performance, while workloads with more diverse prefixes may benefit from smaller pages. See [Data Transfer Optimization](#data-transfer-optimization) for how page granularity affects I/O performance. - **`--hicache-storage-prefetch-policy {best_effort,wait_complete,timeout}`**: Controls when prefetching from storage should stop. See [Prefetch from L3](#prefetch-from-l3) for details. - `best_effort`: Prefetch as much as possible without blocking - `wait_complete`: Wait for prefetch to complete before proceeding - `timeout`: Terminates after specified time or when complete (Recommended for production environments, as setting an appropriate timeout helps the system meet required SLOs) - **`--hicache-write-policy {write_back,write_through,write_through_selective}`**: Controls how data is written from faster to slower memory tiers. See [Data Write-back](#data-write-back) for details. - `write_through`: Immediately writes data to all tiers (strongest caching benefits) - `write_through_selective`: Uses hit-count tracking to back up only frequently accessed data - `write_back`: Writes data back to slower tiers only when eviction is needed (reduces I/O load) - **`--hicache-io-backend {direct,kernel}`**: Choose the I/O backend for KV cache transfer between CPU and GPU. See [Data Transfer Optimization](#data-transfer-optimization) for details. - `direct`: Standard CUDA memory copy operations - `kernel`: GPU-assisted I/O kernels (recommended for better performance) - **`--hicache-mem-layout {layer_first,page_first,page_first_direct}`**: Memory layout for the host memory pool. See [Data Transfer Optimization](#data-transfer-optimization) for details. - `layer_first`: Compatible with GPU computation kernels (default for GPU memory) - `page_first`: Optimized for I/O efficiency - `page_first_direct`: Groups all tokens of a given layer within a page, allowing transfers from L2 to GPU to be aggregated at the page-layer level - **`--hicache-storage-backend {file,mooncake,hf3fs,nixl,aibrix,dynamic}`**: Choose the storage backend for the L3 tier. Built-in backends: file, mooncake, hf3fs, nixl, aibrix. For dynamic backend, use --hicache-storage-backend-extra-config to specify: `backend_name` (custom name), `module_path` (Python module path), `class_name` (backend class name). See [Unified Interfaces and Rich L3 Storage Backends](#unified-interfaces-and-rich-l3-storage-backends) for available backends. - **`--enable-lmcache`**: Using LMCache as an alternative hierarchical cache solution. - **`--hicache-storage-backend-extra-config HICACHE_STORAGE_BACKEND_EXTRA_CONFIG`**: the extra config can be either - a JSON string containing extra configuration for the storage backend, e.g., `--hicache-storage-backend-extra-config '{"prefetch_threshold":512, "prefetch_timeout_base": 0.5, "prefetch_timeout_per_ki_token": 0.25}' `, or - a TOML or JSON or YAML file specifying the extra configuration for the storage backend (to differentiate from the JSON string input, prepend a `@` in front of the file name), e.g., `--hicache-storage-backend-extra-config "@config.toml"` where `config.toml` is the config file containing the complex configurations. This can be useful when the configuration consists of many or complex key-value pairs (for instance, it is preferred to use a config file for NIXL backend as its configurations can be complex). ================================================ FILE: docs/advanced_features/hicache_storage_runtime_attach_detach.md ================================================ # Runtime Attach/Detach HiCache Storage Backend (No Restart) This document explains how to **dynamically attach/detach the HiCache L3 storage backend at runtime** (e.g., `mooncake` / `hf3fs` / `nixl` / `file` / `aibrix` / `eic`) while **SGLang is already running and serving traffic**, without restarting the process. For safety and consistency, the current implementation **strictly requires** these operations to happen only when the service is **idle**: - **No running requests** - **No waiting/queued requests** If the idle condition is not met, the API will fail fast (HTTP 400) and **will not modify** the current service state. --- ## 1. Background and implementation overview ### 1.1 Architecture / control path The control path is: 1. **HTTP Server** (`python/sglang/srt/entrypoints/http_server.py`) - Exposes `PUT /hicache/storage-backend`, `DELETE /hicache/storage-backend`, `GET /hicache/storage-backend` 2. **TokenizerManager** (`python/sglang/srt/managers/tokenizer_communicator_mixin.py`) - Sends the request to the Scheduler via `_Communicator` 3. **Scheduler** (`python/sglang/srt/managers/scheduler.py`) - Performs a **strict idle check** - Calls `tree_cache.attach_storage_backend(...)` / `detach_storage_backend(...)` 4. **HiRadixCache** (`python/sglang/srt/mem_cache/hiradix_cache.py`) - Parses `hicache_storage_backend_extra_config_json` (supports both backend config and prefetch knobs) - Calls `cache_controller.attach_storage_backend(...)` / `detach_storage_backend(...)` 5. **HiCacheController** (`python/sglang/srt/managers/cache_controller.py`) - Creates/destroys the storage backend instance (via `StorageBackendFactory`) - Starts/stops backend background threads at runtime (prefetch/backup) --- ## 2. Idle-state requirement (strict) The Scheduler uses `is_fully_idle()` which checks: - No running batches (including chunked prefill, overlap, pipeline-parallel, and disaggregation paths) - No waiting requests in any queue (waiting, grammar, disagg bootstrap/prealloc/transfer/inflight) - No DLLM staging requests If the condition is not met, attach/detach returns an error like: - `Reject attach: scheduler is not idle. #queue-req=... #running-req=...` > Tip: before switching, drain upstream traffic and wait for the server to become idle, then call attach/detach. ### 2.1 DP (data parallel) semantics When `dp_size > 1`, the tokenizer dispatches the request to **all DP scheduler instances** and aggregates their responses: - The final `success` is **true only if all DP ranks return success** - The final `message` concatenates messages from all DP ranks This is intended to prevent “silent partial success”, but it also means you may see: - Overall **failure** even though **some ranks already succeeded** Currently there is **no automatic partial rollback** across DP ranks (see TODO in code). Operationally: - Prefer to keep backend config identical across ranks - If attach fails, immediately call detach (best-effort/idempotent), fix config, then retry attach --- ## 3. How to use (HTTP Admin API) The examples below assume your SGLang HTTP server is at `http://127.0.0.1:30000`. ### 3.1 Query current storage backend status ```bash curl -s http://127.0.0.1:30000/hicache/storage-backend ``` Example response: ```json { "hicache_storage_backend": "mooncake", "hicache_storage_backend_extra_config": "{\"master_server_address\":\"127.0.0.1:50051\", ...}" } ``` ### 3.2 Attach (enable) a storage backend ```bash curl -s -X PUT http://127.0.0.1:30000/hicache/storage-backend \ -H 'Content-Type: application/json' \ -d '{ "hicache_storage_backend": "mooncake" }' ``` ```bash curl -s -X PUT http://127.0.0.1:30000/hicache/storage-backend \ -H 'Content-Type: application/json' \ -d '{ "hicache_storage_backend": "mooncake", "hicache_storage_backend_extra_config_json": "{\"master_server_address\":\"127.0.0.1:50051\",\"protocol\":\"tcp\",\"global_segment_size\":\"4gb\",\"prefetch_threshold\":256}", "hicache_storage_prefetch_policy": "timeout" }' ``` Notes: - `hicache_storage_backend_extra_config_json` can include both: - **Backend configuration** (e.g., Mooncake master/metadata/protocol, etc.) - **Prefetch configuration** (`prefetch_threshold`, `prefetch_timeout_base`, `prefetch_timeout_per_ki_token`, `hicache_storage_pass_prefix_keys`) ### 3.3 Detach (disable) the storage backend ```bash curl -s -X DELETE http://127.0.0.1:30000/hicache/storage-backend ``` Notes: - Detach only makes SGLang **stop using** the L3 storage backend and stops prefetch/backup threads - It **does not automatically delete** data stored in Mooncake/HF3FS (or other remote backends) --- ## 4. Behavior and caveats - **No restart required**: attach/detach switches in-process at runtime - **Must be idle**: otherwise the request is rejected to avoid consistency issues - **Host KV layout constraints still apply**: for example, Mooncake still requires layouts like `page_first/page_first_direct/page_head`; if the server's HiCache host-memory layout does not satisfy the backend requirements, attach will fail with an error - **Observability**: - After attach, `server_args.hicache_storage_backend*` is updated on both the tokenizer and scheduler sides - If metrics are enabled, attach will create a storage metrics collector in `HiRadixCache` on demand ================================================ FILE: docs/advanced_features/hyperparameter_tuning.md ================================================ # Hyperparameter Tuning ## Achieving high throughput for offline batch inference Achieving a large batch size is the most important thing for attaining high throughput in offline batch inference. When the server is running at full load in a steady state, look for the following in the log: ```Decode batch. #running-req: 233, #token: 370959, token usage: 0.82, cuda graph: True, gen throughput (token/s): 4594.01, #queue-req: 317``` ### Adjust the request submission speed to control `#queue-req` `#queue-req` indicates the number of requests in the queue. If you frequently see `#queue-req: 0`, it suggests that your client code is submitting requests too slowly. A healthy range for `#queue-req` is `100 - 2000`. However, avoid making `#queue-req` too large, as this will increase the scheduling overhead on the server. ### Achieve a high `token usage` `token usage` indicates the KV cache memory utilization of the server. `token usage > 0.9` means good utilization. If you frequently see `token usage < 0.9` and `#queue-req > 0`, it means the server is too conservative about taking in new requests. You can decrease `--schedule-conservativeness` to a value like 0.3. The case of a server being too conservative can happen when users send many requests with a large `max_new_tokens` but the requests stop very early due to EOS or stop strings. On the other hand, if you see `token usage` very high and you frequently see warnings like `KV cache pool is full. Retract requests. #retracted_reqs: 1, #new_token_ratio: 0.9998 -> 1.0000`, you can increase `--schedule-conservativeness` to a value like 1.3. If you see `KV cache pool is full. Retract requests.` occasionally but not frequently (~1 time per minute), it is okay. ### Tune `--mem-fraction-static` to increase KV cache pool capacity SGLang allocates memory as follows: Total memory usage = model weights + KV cache pool + CUDA graph buffers + activations The `--mem-fraction-static` parameter determines how much memory is allocated to the first two components: mem_fraction_static = (model weights + KV cache pool) / GPU memory capacity To support higher concurrency, you should maximize the KV cache pool capacity by setting `--mem-fraction-static` as high as possible while still reserving enough memory for activations and CUDA graph buffers. SGLang uses simple heuristics to set the default value of `--mem-fraction-static`, but you can optimize it for your use cases. As a rule of thumb, reserving 5–8 GB of memory for activations is typically sufficient. You can check this by inspecting the logs just before the server is ready. Look for log entries like this: ``` [2025-08-11 17:17:03] max_total_num_tokens=665690, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=4096, context_len=65536, available_gpu_mem=13.50 GB ``` Check the `available_gpu_mem` value. - If it is between 5–8 GB, the setting is good. - If it is too high (e.g., 10 - 20 GB), increase `--mem-fraction-static` to allocate more memory to the KV cache. - If it is too low, you risk out-of-memory (OOM) errors later, so decrease `--mem-fraction-static`. Another straightforward approach is to increase `--mem-fraction-static` in increments of 0.01 until you encounter OOM errors for your workloads. ### Avoid out-of-memory errors by tuning `--chunked-prefill-size`, `--mem-fraction-static`, and `--max-running-requests` If you encounter out-of-memory (OOM) errors, you can adjust the following parameters: - If OOM occurs during prefill, try reducing `--chunked-prefill-size` to `4096` or `2048`. This saves memory but slows down the prefill speed for long prompts. - If OOM occurs during decoding, try lowering `--max-running-requests`. - You can also reduce `--mem-fraction-static` to a smaller value, such as 0.8 or 0.7. This decreases the memory usage of the KV cache memory pool and helps prevent OOM errors during both prefill and decoding. However, it limits maximum concurrency and reduces peak throughput. ### Tune `--cuda-graph-max-bs` By default, CUDA graph is enabled only for small batch sizes (e.g., less than 160 or 256). However, for some models, especially at large tensor parallelism sizes, CUDA graph can be useful for batch sizes up to 512 or 768. Therefore, it may be beneficial to increase `--cuda-graph-max-bs` to a larger value. Note that CUDA graph consumes more memory, so you may need to reduce `--mem-fraction-static` at the same time. ### Tune `--dp-size` and `--tp-size` Data parallelism is better for throughput. When there is enough GPU memory, always favor data parallelism for throughput. Refer to [SGLang Model Gateway (former Router)](../advanced_features/sgl_model_gateway.md) for a better data parallelism rather than using `dp_size` parameter. ### Try other options - `torch.compile` accelerates small models on small batch sizes. You can enable it with `--enable-torch-compile`. - Try other quantization (e.g. FP8 quantization with `--quantization fp8`) - Try other parallelism strategies (e.g. [expert parallelism](https://lmsys.org/blog/2025-05-05-large-scale-ep/)) or DP attention for deepseek models (with `--enable-dp-attention --dp-size 8`). - If the workload has many shared prefixes, try `--schedule-policy lpm`. Here, `lpm` stands for longest prefix match. It reorders requests to encourage more cache hits but introduces more scheduling overhead. ================================================ FILE: docs/advanced_features/lora.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# LoRA Serving" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "SGLang enables the use of [LoRA adapters](https://arxiv.org/abs/2106.09685) with a base model. By incorporating techniques from [S-LoRA](https://arxiv.org/pdf/2311.03285) and [Punica](https://arxiv.org/pdf/2310.18547), SGLang can efficiently support multiple LoRA adapters for different sequences within a single batch of inputs." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Arguments for LoRA Serving" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The following server arguments are relevant for multi-LoRA serving:\n", "\n", "* `enable_lora`: Enable LoRA support for the model. This argument is automatically set to True if `--lora-paths` is provided for backward compatibility.\n", "\n", "* `enable_lora_overlap_loading`: Enable asynchronous LoRA weight loading in order to overlap H2D transfers with GPU compute. This should be enabled if you find that your LoRA workloads are bottlenecked by adapter weight loading, for example when frequently loading large LoRA adapters.\n", "\n", "* `lora_paths`: The list of LoRA adapters to load. Each adapter must be specified in one of the following formats: | = | JSON with schema {\"lora_name\":str,\"lora_path\":str,\"pinned\":bool}.\n", "\n", "* `max_loras_per_batch`: Maximum number of adaptors used by each batch. This argument can affect the amount of GPU memory reserved for multi-LoRA serving, so it should be set to a smaller value when memory is scarce. Defaults to be 8.\n", "\n", "* `max_loaded_loras`: If specified, it limits the maximum number of LoRA adapters loaded in CPU memory at a time. The value must be greater than or equal to `max-loras-per-batch`.\n", "\n", "* `lora_eviction_policy`: LoRA adapter eviction policy when GPU memory pool is full. `lru`: Least Recently Used (default, better cache efficiency). `fifo`: First-In-First-Out.\n", "\n", "* `lora_backend`: The backend of running GEMM kernels for Lora modules. Currently we support Triton LoRA backend (`triton`) and Chunked SGMV backend (`csgmv`). In the future, faster backend built upon Cutlass or Cuda kernels will be added.\n", "\n", "* `max_lora_rank`: The maximum LoRA rank that should be supported. If not specified, it will be automatically inferred from the adapters provided in `--lora-paths`. This argument is needed when you expect to dynamically load adapters of larger LoRA rank after server startup.\n", "\n", "* `lora_target_modules`: The union set of all target modules where LoRA should be applied (e.g., `q_proj`, `k_proj`, `gate_proj`). If not specified, it will be automatically inferred from the adapters provided in `--lora-paths`. This argument is needed when you expect to dynamically load adapters of different target modules after server startup. You can also set it to `all` to enable LoRA for all supported modules. However, enabling LoRA on additional modules introduces a minor performance overhead. If your application is performance-sensitive, we recommend only specifying the modules for which you plan to load adapters.\n", "\n", "* `--max-lora-chunk-size`: Maximum chunk size for the ChunkedSGMV LoRA backend. Only used when --lora-backend is 'csgmv'. Choosing a larger value might improve performance. Please tune this value based on your hardware and workload as needed. Defaults to 16.\n", "\n", "* `tp_size`: LoRA serving along with Tensor Parallelism is supported by SGLang. `tp_size` controls the number of GPUs for tensor parallelism. More details on the tensor sharding strategy can be found in [S-Lora](https://arxiv.org/pdf/2311.03285) paper.\n", "\n", "From client side, the user needs to provide a list of strings as input batch, and a list of adaptor names that each input sequence corresponds to." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Usage\n", "\n", "### Serving Single Adaptor" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Note:** SGLang supports LoRA adapters through two APIs:\n", "\n", "1. **OpenAI-Compatible API** (`/v1/chat/completions`, `/v1/completions`): Use the `model:adapter-name` syntax. See [OpenAI API with LoRA](../basic_usage/openai_api_completions.ipynb#Using-LoRA-Adapters) for examples.\n", "\n", "2. **Native API** (`/generate`): Pass `lora_path` in the request body (shown below)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\n", "import requests\n", "\n", "from sglang.test.doc_patch import launch_server_cmd\n", "from sglang.utils import wait_for_server, terminate_process" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "server_process, port = launch_server_cmd(\n", " # Here we set max-loras-per-batch to 2: one slot for adaptor and another one for base model\n", " \"\"\"\n", "python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n", " --enable-lora \\\n", " --lora-paths lora0=algoprog/fact-generation-llama-3.1-8b-instruct-lora \\\n", " --max-loras-per-batch 2 \\\n", " --log-level warning \\\n", "\"\"\"\n", ")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=server_process)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "url = f\"http://127.0.0.1:{port}\"\n", "json_data = {\n", " \"text\": [\n", " \"List 3 countries and their capitals.\",\n", " \"List 3 countries and their capitals.\",\n", " ],\n", " \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n", " # The first input uses lora0, and the second input uses the base model\n", " \"lora_path\": [\"lora0\", None],\n", "}\n", "response = requests.post(\n", " url + \"/generate\",\n", " json=json_data,\n", ")\n", "print(f\"Output 0: {response.json()[0]['text']}\")\n", "print(f\"Output 1: {response.json()[1]['text']}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Serving Multiple Adaptors" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "server_process, port = launch_server_cmd(\"\"\"\n", "python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n", " --enable-lora \\\n", " --lora-paths lora0=algoprog/fact-generation-llama-3.1-8b-instruct-lora \\\n", " lora1=Nutanix/Meta-Llama-3.1-8B-Instruct_SFT_lora_4_alpha_16_humaneval_raw_json \\\n", " --max-loras-per-batch 2 \\\n", " --log-level warning \\\n", "\"\"\")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=server_process)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "url = f\"http://127.0.0.1:{port}\"\n", "json_data = {\n", " \"text\": [\n", " \"List 3 countries and their capitals.\",\n", " \"List 3 countries and their capitals.\",\n", " ],\n", " \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n", " # The first input uses lora0, and the second input uses lora1\n", " \"lora_path\": [\"lora0\", \"lora1\"],\n", "}\n", "response = requests.post(\n", " url + \"/generate\",\n", " json=json_data,\n", ")\n", "print(f\"Output 0: {response.json()[0]['text']}\")\n", "print(f\"Output 1: {response.json()[1]['text']}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Dynamic LoRA loading" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Instead of specifying all adapters during server startup via `--lora-paths`. You can also load & unload LoRA adapters dynamically via the `/load_lora_adapter` and `/unload_lora_adapter` API.\n", "\n", "When using dynamic LoRA loading, it's recommended to explicitly specify both `--max-lora-rank` and `--lora-target-modules` at startup. For backward compatibility, SGLang will infer these values from `--lora-paths` if they are not explicitly provided. However, in that case, you would have to ensure that all dynamically loaded adapters share the same shape (rank and target modules) as those in the initial `--lora-paths` or are strictly \"smaller\"." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "lora0 = \"Nutanix/Meta-Llama-3.1-8B-Instruct_SFT_lora_4_alpha_16_humaneval_raw_json\" # rank - 4, target modules - q_proj, k_proj, v_proj, o_proj, gate_proj\n", "lora1 = \"algoprog/fact-generation-llama-3.1-8b-instruct-lora\" # rank - 64, target modules - q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj\n", "lora0_new = \"philschmid/code-llama-3-1-8b-text-to-sql-lora\" # rank - 256, target modules - q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj\n", "\n", "\n", "# The `--target-lora-modules` param below is technically not needed, as the server will infer it from lora0 which already has all the target modules specified.\n", "# We are adding it here just to demonstrate usage.\n", "server_process, port = launch_server_cmd(\"\"\"\n", " python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n", " --enable-lora \\\n", " --cuda-graph-max-bs 2 \\\n", " --max-loras-per-batch 2 \\\n", " --max-lora-rank 256\n", " --lora-target-modules all\n", " --log-level warning\n", " \"\"\")\n", "\n", "url = f\"http://127.0.0.1:{port}\"\n", "wait_for_server(url, process=server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Load adapter lora0" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = requests.post(\n", " url + \"/load_lora_adapter\",\n", " json={\n", " \"lora_name\": \"lora0\",\n", " \"lora_path\": lora0,\n", " },\n", ")\n", "\n", "if response.status_code == 200:\n", " print(\"LoRA adapter loaded successfully.\", response.json())\n", "else:\n", " print(\"Failed to load LoRA adapter.\", response.json())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Load adapter lora1:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = requests.post(\n", " url + \"/load_lora_adapter\",\n", " json={\n", " \"lora_name\": \"lora1\",\n", " \"lora_path\": lora1,\n", " },\n", ")\n", "\n", "if response.status_code == 200:\n", " print(\"LoRA adapter loaded successfully.\", response.json())\n", "else:\n", " print(\"Failed to load LoRA adapter.\", response.json())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Check inference output:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "url = f\"http://127.0.0.1:{port}\"\n", "json_data = {\n", " \"text\": [\n", " \"List 3 countries and their capitals.\",\n", " \"List 3 countries and their capitals.\",\n", " ],\n", " \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n", " # The first input uses lora0, and the second input uses lora1\n", " \"lora_path\": [\"lora0\", \"lora1\"],\n", "}\n", "response = requests.post(\n", " url + \"/generate\",\n", " json=json_data,\n", ")\n", "print(f\"Output from lora0: \\n{response.json()[0]['text']}\\n\")\n", "print(f\"Output from lora1 (updated): \\n{response.json()[1]['text']}\\n\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Unload lora0 and replace it with a different adapter:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = requests.post(\n", " url + \"/unload_lora_adapter\",\n", " json={\n", " \"lora_name\": \"lora0\",\n", " },\n", ")\n", "\n", "response = requests.post(\n", " url + \"/load_lora_adapter\",\n", " json={\n", " \"lora_name\": \"lora0\",\n", " \"lora_path\": lora0_new,\n", " },\n", ")\n", "\n", "if response.status_code == 200:\n", " print(\"LoRA adapter loaded successfully.\", response.json())\n", "else:\n", " print(\"Failed to load LoRA adapter.\", response.json())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Check output again:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "url = f\"http://127.0.0.1:{port}\"\n", "json_data = {\n", " \"text\": [\n", " \"List 3 countries and their capitals.\",\n", " \"List 3 countries and their capitals.\",\n", " ],\n", " \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n", " # The first input uses lora0, and the second input uses lora1\n", " \"lora_path\": [\"lora0\", \"lora1\"],\n", "}\n", "response = requests.post(\n", " url + \"/generate\",\n", " json=json_data,\n", ")\n", "print(f\"Output from lora0: \\n{response.json()[0]['text']}\\n\")\n", "print(f\"Output from lora1 (updated): \\n{response.json()[1]['text']}\\n\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### OpenAI-compatible API usage\n", "\n", "You can use LoRA adapters via the OpenAI-compatible APIs by specifying the adapter in the `model` field using the `base-model:adapter-name` syntax (for example, `qwen/qwen2.5-0.5b-instruct:adapter_a`). For more details and examples, see the “Using LoRA Adapters” section in the OpenAI API documentation: [openai_api_completions.ipynb](../basic_usage/openai_api_completions.ipynb).\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### LoRA GPU Pinning" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Another advanced option is to specify adapters as `pinned` during loading. When an adapter is pinned, it is permanently assigned to one of the available GPU pool slots (as configured by `--max-loras-per-batch`) and will not be evicted from GPU memory during runtime. Instead, it remains resident until it is explicitly unloaded.\n", "\n", "This can improve performance in scenarios where the same adapter is frequently used across requests, by avoiding repeated memory transfers and reinitialization overhead. However, since GPU pool slots are limited, pinning adapters reduces the flexibility of the system to dynamically load other adapters on demand. If too many adapters are pinned, it may lead to degraded performance, or in the most extreme case (`Number of pinned adapters == max-loras-per-batch`), halt all unpinned requests. Therefore, currently SGLang limits maximal number of pinned adapters to `max-loras-per-batch - 1` to prevent unexpected starvations. \n", "\n", "In the example below, we start a server with `lora1` loaded as pinned, `lora2` and `lora3` loaded as regular (unpinned) adapters. Please note that, we intentionally specify `lora2` and `lora3` in two different formats to demonstrate that both are supported." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "server_process, port = launch_server_cmd(\"\"\"\n", " python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n", " --enable-lora \\\n", " --cuda-graph-max-bs 8 \\\n", " --max-loras-per-batch 3 \\\n", " --max-lora-rank 256 \\\n", " --lora-target-modules all \\\n", " --lora-paths \\\n", " {\"lora_name\":\"lora0\",\"lora_path\":\"Nutanix/Meta-Llama-3.1-8B-Instruct_SFT_lora_4_alpha_16_humaneval_raw_json\",\"pinned\":true} \\\n", " {\"lora_name\":\"lora1\",\"lora_path\":\"algoprog/fact-generation-llama-3.1-8b-instruct-lora\"} \\\n", " lora2=philschmid/code-llama-3-1-8b-text-to-sql-lora\n", " --log-level warning\n", " \"\"\")\n", "\n", "\n", "url = f\"http://127.0.0.1:{port}\"\n", "wait_for_server(url, process=server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can also specify adapter as pinned during dynamic adapter loading. In the example below, we reload `lora2` as pinned adapter:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = requests.post(\n", " url + \"/unload_lora_adapter\",\n", " json={\n", " \"lora_name\": \"lora1\",\n", " },\n", ")\n", "\n", "response = requests.post(\n", " url + \"/load_lora_adapter\",\n", " json={\n", " \"lora_name\": \"lora1\",\n", " \"lora_path\": \"algoprog/fact-generation-llama-3.1-8b-instruct-lora\",\n", " \"pinned\": True, # Pin the adapter to GPU\n", " },\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Verify that the results are expected:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "url = f\"http://127.0.0.1:{port}\"\n", "json_data = {\n", " \"text\": [\n", " \"List 3 countries and their capitals.\",\n", " \"List 3 countries and their capitals.\",\n", " \"List 3 countries and their capitals.\",\n", " ],\n", " \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n", " # The first input uses lora0, and the second input uses lora1\n", " \"lora_path\": [\"lora0\", \"lora1\", \"lora2\"],\n", "}\n", "response = requests.post(\n", " url + \"/generate\",\n", " json=json_data,\n", ")\n", "print(f\"Output from lora0 (pinned): \\n{response.json()[0]['text']}\\n\")\n", "print(f\"Output from lora1 (pinned): \\n{response.json()[1]['text']}\\n\")\n", "print(f\"Output from lora2 (not pinned): \\n{response.json()[2]['text']}\\n\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Choosing LoRA Backend\n", "\n", "SGLang supports two LoRA backends that you can choose from using the `--lora-backend` argument:\n", "\n", "- `triton`: Basic Triton-based backend.\n", "- `csgmv`: Default chunked SGMV backend optimized for high concurrency scenarios.\n", "\n", "The `csgmv` backend was recently introduced to improve performance especially at high-concurrency scenarios. Our benchmark shows that it achieves 20% to 80% latency improvements over the basic triton backend." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "server_process, port = launch_server_cmd(\"\"\"\n", " python3 -m sglang.launch_server \\\n", " --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n", " --enable-lora \\\n", " --lora-backend csgmv \\\n", " --max-loras-per-batch 16 \\\n", " --lora-paths lora1=path/to/lora1 lora2=path/to/lora2\n", " \"\"\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## LoRA Overlap Loading" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "By using the `--enable-lora-overlap-loading` server argument, the SGLang engine is able to overlap the loading of LoRA weights with prefill and decode compute, essentially hiding the data movement for LoRA weights behind GPU computation. Our benchmarks show that under adversarial conditions, enabling this feature can result in a ~35% reduction in median TTFT - (see the [LoRA overlap loading PR](https://github.com/sgl-project/sglang/pull/15512) for detailed benchmarks)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "lora0 = \"Nutanix/Meta-Llama-3.1-8B-Instruct_SFT_lora_4_alpha_16_humaneval_raw_json\"\n", "lora1 = \"algoprog/fact-generation-llama-3.1-8b-instruct-lora\"\n", "lora2 = \"philschmid/code-llama-3-1-8b-text-to-sql-lora\"\n", "\n", "\n", "server_process, port = launch_server_cmd(\"\"\"\n", " python3 -m sglang.launch_server \\\n", " --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n", " --enable-lora \\\n", " --enable-lora-overlap-loading \\\n", " --lora-paths lora0=Nutanix/Meta-Llama-3.1-8B-Instruct_SFT_lora_4_alpha_16_humaneval_raw_json \\\n", " lora1=algoprog/fact-generation-llama-3.1-8b-instruct-lora \\\n", " lora2=philschmid/code-llama-3-1-8b-text-to-sql-lora \\\n", " --max-lora-rank 256 \\\n", " --max-loras-per-batch 2 \\\n", " --max-loaded-loras 4\n", " \"\"\")\n", "\n", "url = f\"http://127.0.0.1:{port}\"\n", "wait_for_server(url, process=server_process)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "json_data = {\n", " \"text\": [\n", " \"Write a very long fairy-tale.\",\n", " \"List 3 countries and their capitals.\",\n", " \"List 3 countries and their capitals.\",\n", " ],\n", " \"sampling_params\": [\n", " {\"max_new_tokens\": 1024, \"temperature\": 0},\n", " {\"max_new_tokens\": 64, \"temperature\": 0},\n", " {\"max_new_tokens\": 64, \"temperature\": 0},\n", " ],\n", " \"lora_path\": [\"lora0\", \"lora1\", \"lora2\"],\n", "}\n", "\n", "# lora0 and lora1 will be loaded into the memory pool first, and because max_loras_per_batch = 2, lora2's request will remain in the queue.\n", "# lora1's request will likely finish first, and once it does, lora2 will be loaded. With --enable-lora-overlap-loading, this loading will\n", "# occur asynchronously and thus decoding for lora0's request won't be blocked.\n", "response = requests.post(\n", " url + \"/generate\",\n", " json=json_data,\n", ")\n", "\n", "for i in range(3):\n", " print(f\"Output from lora{i}: \\n{response.json()[i]['text']}\\n\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Limitations of LoRA Overlap Loading" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "However, LoRA overlap loading is not free and comes with two important caveats:\n", "\n", "1. **Pinned CPU memory requirement**:\n", " Asynchronous H2D memory copies require LoRA weights to be pinned in CPU memory, which is a finite system resource. To mitigate excessive pinned-memory usage, SGLang currently restricts `max_loaded_loras` to be at most 2× `max_loras_per_batch` when LoRA overlap loading is enabled.\n", "\n", "2. **Reduced multi-adapter prefill batching**:\n", " With overlap loading, adapters become available on the GPU at different times because each adapter is loaded asynchronously. This can reduce the scheduler’s ability to form multi-adapter prefill batches, since only requests whose adapters are currently loaded can be grouped together. As a result, requests for different adapters will be scheduled in separate (or smaller) prefill batches, which can increase TTFT when adapter load time is small compared to prefill compute time. This is why LoRA overlap loading is disabled by default: it should only be enabled when users have determined that LoRA weight loading is a bottleneck (EG high adapter churn, heavy adapter weights, or PCIe-bottlenecked workloads).\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Example When Overlap Loading Results in Higher Latency" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "For instance, suppose we have four LoRA adapters: `lora0`, `lora1`, `lora2`, and `lora3`. Loading any adapter takes 2ms, while the prefill step for requests for that adapter takes 20ms.\n", "\n", "1. **Baseline**:\n", " The engine loads all four adapters synchronously, then runs one combined prefill batch, giving us a total time of ≈ `2 * 4 + 20 = 28ms`\n", "\n", "2. **With LoRA overlap loading enabled**:\n", " The engine begins loading `lora0` and, once it is ready, schedules a prefill batch containing only `lora0` while `lora1` loads in the background. Then it schedules `lora1`’s prefill while `lora2` loads, and so on. In the worst case where prefill cannot be batched across adapters, total time is ≈ `2 + 4 * 20 = 82ms`\n", "\n", "In this scenario, overlap loading reduces adapter-load overhead, but the loss of multi-adapter prefill batching dominates and leads to higher TTFT." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Future Works\n", "\n", "The development roadmap for LoRA-related features can be found in this [issue](https://github.com/sgl-project/sglang/issues/2929). Other features, including Embedding Layer, Unified Paging, Cutlass backend are still under development." ] } ], "metadata": { "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 2 } ================================================ FILE: docs/advanced_features/observability.md ================================================ # Observability ## Production Metrics SGLang exposes the following metrics via Prometheus. You can enable them by adding `--enable-metrics` when launching the server. You can query them by: ``` curl http://localhost:30000/metrics ``` See [Production Metrics](../references/production_metrics.md) and [Production Request Tracing](../references/production_request_trace.md) for more details. ## Logging By default, SGLang does not log any request contents. You can log them by using `--log-requests`. You can control the verbosity by using `--log-request-level`. See [Logging](server_arguments.md#logging) for more details. ## Request Dump and Replay You can dump all requests and replay them later for benchmarking or other purposes. To start dumping, use the following command to send a request to a server: ``` python3 -m sglang.srt.managers.configure_logging --url http://localhost:30000 --dump-requests-folder /tmp/sglang_request_dump --dump-requests-threshold 100 ``` The server will dump the requests into a pickle file for every 100 requests. To replay the request dump, use `scripts/playground/replay_request_dump.py`. ## Crash Dump and Replay Sometimes the server might crash, and you may want to debug the cause of the crash. SGLang supports crash dumping, which will dump all requests from the 5 minutes before the crash, allowing you to replay the requests and debug the reason later. To enable crash dumping, use `--crash-dump-folder /tmp/crash_dump`. To replay the crash dump, use `scripts/playground/replay_request_dump.py`. ================================================ FILE: docs/advanced_features/pd_disaggregation.md ================================================ # PD Disaggregation ## Why and What is PD Disaggregation? Large Language Model (LLM) inference comprises two distinct phases: **Prefill** and **Decode**. The Prefill phase is computation-intensive, processing the entire input sequence, while the Decode phase is memory-intensive, managing the Key-Value (KV) cache for token generation. Traditionally, these phases are handled within a unified engine, where combined scheduling of prefill and decode batches introduces inefficiencies. To address these challenges, we introduce **Prefill and Decoding (PD) Disaggregation** in SGLang. ### Issues with Unified Scheduling The conventional unified engine, which processes prefill and decode batches together, results in two significant problems: 1. **Prefill Interruption**: Incoming prefill batches frequently interrupt ongoing decode batches, causing substantial delays in token generation. 2. **DP Attention Imbalance**: In data-parallel (DP) attention, one DP worker may process a prefill batch while another handles a decode batch simultaneously, leading to increased decode latency. PD Disaggregation resolves these by separating the two stages, enabling tailored optimizations for each. For the design details, please refer to [link](https://docs.google.com/document/d/1rQXJwKd5b9b1aOzLh98mnyMhBMhlxXA5ATZTHoQrwvc/edit?tab=t.0). Currently, we support Mooncake and NIXL as the transfer engine. ## Profiling in PD Disaggregation Mode When you need to profile prefill or decode workers in PD disaggregation mode, please refer to the [Profile In PD Disaggregation Mode](https://docs.sglang.io/developer_guide/benchmark_and_profiling.html#profile-in-pd-disaggregation-mode) section in the Benchmark and Profiling guide. Due to torch profiler limitations, prefill and decode workers must be profiled separately using dedicated command-line options. ## Router Integration For deploying PD disaggregation at scale with load balancing and fault tolerance, SGLang provides a router. The router can distribute requests between prefill and decode instances using various routing policies. For detailed information on setting up routing with PD disaggregation, including configuration options and deployment patterns, see the [SGLang Model Gateway (former Router)](../advanced_features/sgl_model_gateway.md#prefill-decode-disaggregation). ## Mooncake ### Requirements ```bash uv pip install mooncake-transfer-engine ``` ### Usage ### Llama Single Node ```bash python -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --disaggregation-mode prefill \ --port 30000 \ --disaggregation-ib-device mlx5_roce0 python -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --disaggregation-mode decode \ --port 30001 \ --base-gpu-id 1 \ --disaggregation-ib-device mlx5_roce0 python -m sglang_router.launch_router --pd-disaggregation --prefill http://127.0.0.1:30000 --decode http://127.0.0.1:30001 --host 0.0.0.0 --port 8000 ``` ### DeepSeek Multi-Node ```bash # prefill 0 python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3-0324 \ --disaggregation-ib-device ${device_name} \ --disaggregation-mode prefill \ --host ${local_ip} \ --port 30000 \ --trust-remote-code \ --dist-init-addr ${prefill_master_ip}:5000 \ --nnodes 2 \ --node-rank 0 \ --tp-size 16 \ --dp-size 8 \ --enable-dp-attention \ --moe-a2a-backend deepep \ --mem-fraction-static 0.8 # prefill 1 python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3-0324 \ --disaggregation-ib-device ${device_name} \ --disaggregation-mode prefill \ --host ${local_ip} \ --port 30000 \ --trust-remote-code \ --dist-init-addr ${prefill_master_ip}:5000 \ --nnodes 2 \ --node-rank 1 \ --tp-size 16 \ --dp-size 8 \ --enable-dp-attention \ --moe-a2a-backend deepep \ --mem-fraction-static 0.8 # decode 0 python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3-0324 \ --disaggregation-ib-device ${device_name} \ --disaggregation-mode decode \ --host ${local_ip} \ --port 30001 \ --trust-remote-code \ --dist-init-addr ${decode_master_ip}:5000 \ --nnodes 2 \ --node-rank 0 \ --tp-size 16 \ --dp-size 8 \ --enable-dp-attention \ --moe-a2a-backend deepep \ --mem-fraction-static 0.8 \ --max-running-requests 128 # decode 1 python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3-0324 \ --disaggregation-ib-device ${device_name} \ --disaggregation-mode decode \ --host ${local_ip} \ --port 30001 \ --trust-remote-code \ --dist-init-addr ${decode_master_ip}:5000 \ --nnodes 2 \ --node-rank 1 \ --tp-size 16 \ --dp-size 8 \ --enable-dp-attention \ --moe-a2a-backend deepep \ --mem-fraction-static 0.8 \ --max-running-requests 128 ``` ### Advanced Configuration PD Disaggregation with Mooncake supports the following environment variables for fine-grained control over system behavior. #### NVLink Transport Configuration To enable NVLink transport for KV cache transfers with the mooncake backend (recommended for NVL72 deployments), set the following environment variables. Note that auxiliary data transfer will still use TCP as a temporary workaround. ```bash export SGLANG_MOONCAKE_CUSTOM_MEM_POOL=NVLINK export MC_FORCE_MNNVL=True ``` The `SGLANG_MOONCAKE_CUSTOM_MEM_POOL` environment variable enables the custom memory pool. Supported values are `NVLINK` (or `True`), `BAREX`, and `INTRA_NODE_NVLINK`. #### Prefill Server Configuration | Variable | Description | Default | |:--------:|:-----------:|:--------: | **`SGLANG_DISAGGREGATION_THREAD_POOL_SIZE`** | Controls the total number of worker threads for KVCache transfer operations per TP rank | A dynamic value calculated by `int(0.75 * os.cpu_count()) // 8)`, which is limited to be larger than 4 and less than 12 to ensure efficiency and prevent thread race conditions | | **`SGLANG_DISAGGREGATION_QUEUE_SIZE`** | Sets the number of parallel transfer queues. KVCache transfer requests from multiple decode instances will be sharded into these queues so that they can share the threads and the transfer bandwidth at the same time. If it is set to `1`, then we transfer requests one by one according to fcfs strategy | `4` | | **`SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT`** | Timeout (seconds) for receiving destination KV indices during request initialization | `300` | | **`SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL`** | Interval (seconds) between cleanups of bootstrap entries | `120` | If a greater mean TTFT is acceptable, you can `export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600` (10 minutes) to relax the timeout condition. Please be aware that this setting will cause prefill instances to take a longer time to clean up the affected memory resources when a running decode node loses connection. #### Decode Server Configuration | Variable | Description | Default | |:--------:|:-----------:|:--------: | **`SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL`** | Interval (seconds) between health checks to prefill bootstrap servers | `5.0` | | **`SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE`** | Consecutive heartbeat failures before marking prefill server offline | `2` | | **`SGLANG_DISAGGREGATION_WAITING_TIMEOUT`** | Timeout (seconds) for receiving KV Cache after request initialization | `300` | If a greater mean TTFT is acceptable, you can `export SGLANG_DISAGGREGATION_WAITING_TIMEOUT=600` (10 minutes) to relax the timeout condition. ## NIXL ### Requirements Install via pip. ```bash pip install nixl ``` Or build from source - may be required if you already have UCX installed. ```bash git clone https://github.com/ai-dynamo/nixl.git cd nixl pip install . --config-settings=setup-args="-Ducx_path=/path/to/ucx" ``` ### Usage ### Llama Single Node ```bash python -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --disaggregation-mode prefill \ --port 30000 \ --disaggregation-transfer-backend nixl python -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --disaggregation-mode decode \ --port 30001 \ --base-gpu-id 1 \ --disaggregation-transfer-backend nixl python -m sglang_router.launch_router --pd-disaggregation --prefill http://127.0.0.1:30000 --decode http://127.0.0.1:30001 --host 0.0.0.0 --port 8000 ``` ### DeepSeek Multi-Node ```bash # prefill 0 python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3-0324 \ --disaggregation-transfer-backend nixl \ --disaggregation-mode prefill \ --host ${local_ip} \ --port 30000 \ --trust-remote-code \ --dist-init-addr ${prefill_master_ip}:5000 \ --nnodes 2 \ --node-rank 0 \ --tp-size 16 \ --dp-size 8 \ --enable-dp-attention \ --moe-a2a-backend deepep \ --mem-fraction-static 0.8 # prefill 1 python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3-0324 \ --disaggregation-transfer-backend nixl \ --disaggregation-mode prefill \ --host ${local_ip} \ --port 30000 \ --trust-remote-code \ --dist-init-addr ${prefill_master_ip}:5000 \ --nnodes 2 \ --node-rank 1 \ --tp-size 16 \ --dp-size 8 \ --enable-dp-attention \ --moe-a2a-backend deepep \ --mem-fraction-static 0.8 # decode 0 python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3-0324 \ --disaggregation-transfer-backend nixl \ --disaggregation-mode decode \ --host ${local_ip} \ --port 30001 \ --trust-remote-code \ --dist-init-addr ${decode_master_ip}:5000 \ --nnodes 2 \ --node-rank 0 \ --tp-size 16 \ --dp-size 8 \ --enable-dp-attention \ --moe-a2a-backend deepep \ --mem-fraction-static 0.8 \ --max-running-requests 128 # decode 1 python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3-0324 \ --disaggregation-transfer-backend nixl \ --disaggregation-mode decode \ --host ${local_ip} \ --port 30001 \ --trust-remote-code \ --dist-init-addr ${decode_master_ip}:5000 \ --nnodes 2 \ --node-rank 1 \ --tp-size 16 \ --dp-size 8 \ --enable-dp-attention \ --moe-a2a-backend deepep \ --mem-fraction-static 0.8 \ --max-running-requests 128 ``` ### Advanced Configuration #### NIXL Backend Selection By default, NIXL uses the **UCX** backend for KV cache transfers. You can select a different NIXL plugin backend depending on your infrastructure using the environment variable `SGLANG_DISAGGREGATION_NIXL_BACKEND`. Example: `export SGLANG_DISAGGREGATION_NIXL_BACKEND=LIBFABRIC` **Available backends:** UCX (default), LIBFABRIC, or any installed NIXL plugin. Example usage: ```bash export SGLANG_DISAGGREGATION_NIXL_BACKEND=LIBFABRIC python -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --disaggregation-mode prefill \ --disaggregation-transfer-backend nixl \ --port 30000 ``` ## ASCEND ### Usage Use ascend backend with [memfabric_hybrid](https://gitcode.com/Ascend/memfabric_hybrid) and ASCEND_MF_STORE_URL being set ```bash pip install memfabric-hybrid==1.0.0 export ASCEND_MF_STORE_URL="tcp://xxx.xx.xxx.xxx:xxxx" ``` Use mooncake backend, more details can be found in mooncake section. ```bash export ENABLE_ASCEND_TRANSFER_WITH_MOONCAKE=true ``` ASCEND_NPU_PHY_ID need to be set in container env ```bash export ASCEND_NPU_PHY_ID=xxx ``` ### Llama Single Node ```bash python -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --disaggregation-mode prefill \ --port 30000 \ --disaggregation-transfer-backend ascend python -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --disaggregation-mode decode \ --port 30001 \ --base-gpu-id 1 \ --disaggregation-transfer-backend ascend python -m sglang_router.launch_router --pd-disaggregation --prefill http://127.0.0.1:30000 --decode http://127.0.0.1:30001 --host 0.0.0.0 --port 8000 ``` ### DeepSeek Multi-Node ```bash # prefill 0 python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3-0324 \ --disaggregation-transfer-backend ascend \ --disaggregation-mode prefill \ --host ${local_ip} \ --port 30000 \ --trust-remote-code \ --dist-init-addr ${prefill_master_ip}:5000 \ --nnodes 1 \ --node-rank 0 \ --tp-size 16 # decode 0 python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3-0324 \ --disaggregation-transfer-backend ascend \ --disaggregation-mode decode \ --host ${local_ip} \ --port 30001 \ --trust-remote-code \ --dist-init-addr ${decode_master_ip}:5000 \ --nnodes 1 \ --node-rank 0 \ --tp-size 16 ``` ================================================ FILE: docs/advanced_features/piecewise_cuda_graph.md ================================================ # Piecewise CUDA Graph ## Motivation Standard CUDA graphs capture the entire model forward pass as a single graph. This works well for decode (fixed batch size), but not for extend/prefill where the number of tokens varies across iterations. Piecewise CUDA Graph (PCG) solves this by splitting the model's computation graph into pieces (roughly one per layer) at "split points" (e.g., MoE dispatch ops). Each piece is captured as a separate CUDA graph for a set of pre-defined token lengths. At runtime, the input is padded to the nearest captured size, and each piece is replayed. This eliminates kernel launch overhead for prefill/extend while still supporting dynamic shapes. Recently we **enabled PCG by default**, which means that the old `--enable-piecewise-cuda-graph` flag is deprecated. Use `--disable-piecewise-cuda-graph` to turn it off. ## Usage PCG is enabled by default for supported configurations. No extra flags needed: ```bash python3 -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct ``` ### Disable PCG ```bash python3 -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --disable-piecewise-cuda-graph ``` ### Custom capture sizes ```bash python3 -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --piecewise-cuda-graph-max-tokens 2048 ``` ### Server Args | Argument | Default | Description | |---|---|---| | `--disable-piecewise-cuda-graph` | `False` | Disable PCG for extend/prefill. | | `--enforce-piecewise-cuda-graph` | `False` | Force-enable PCG, skipping all auto-disable conditions. For testing only. | | `--piecewise-cuda-graph-max-tokens` | `None` (auto) | Maximum token count to capture. Defaults to `chunked_prefill_size` (non-MLA) or `2048` (MLA). | | `--piecewise-cuda-graph-tokens` | `None` (auto) | Explicit list of token lengths to capture. Auto-generated if not set. | | `--piecewise-cuda-graph-compiler` | `"eager"` | Compiler backend for the captured subgraphs. Choices: `eager`, `inductor`. | | ~~`--enable-piecewise-cuda-graph`~~ | — | **Deprecated.** PCG is now enabled by default. Use `--enforce-piecewise-cuda-graph` to skip auto-disable conditions. | ## Bug Report PCG is enabled by default but is still in an experimental stage. Since PCG relies on `torch.compile` to trace the model's forward pass, most bugs are introduced by torch compile tracing failures (e.g., untraceable ops, dynamic control flow, or graph breaks). If you encounter any issues related to PCG, please disable it by adding `--disable-piecewise-cuda-graph` to your launch command and report the bug at [GitHub Issues](https://github.com/sgl-project/sglang/issues/new/choose). We greatly appreciate your help in improving this feature. ### For Users If you see an error message like the following during server startup, it is a PCG bug: ``` Piecewise CUDA Graph is enabled by default as an experimental feature. To work around this error, add --disable-piecewise-cuda-graph to your launch command. Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose ``` To work around it, add `--disable-piecewise-cuda-graph` to your launch command. When filing a bug report, please include: 1. The full error traceback 2. Model name and quantization method 3. Launch command with all arguments 4. GPU type and driver version ### For Developers Since PCG relies on `torch.compile` to trace the model's forward pass, newly developed CUDA kernels (both JIT kernels and sgl-kernels) are typically not compatible with `torch.compile` out of the box. The tracing will fail on untraceable operations such as JIT compilation, file I/O, or dynamic module loading inside the kernel. To make a kernel compatible with PCG, you need to register it as a custom op using `register_custom_op` from `sglang.srt.utils.custom_op`. This wraps the kernel as an opaque node in the compiled graph so that `torch.compile` will not trace inside it. **Example usage (JIT kernel):** ```python from sglang.srt.utils.custom_op import register_custom_op # Inplace operator (no return value) @register_custom_op(mutates_args=["output_q", "output_s"]) def per_token_group_quant_8bit( input: torch.Tensor, output_q: torch.Tensor, output_s: torch.Tensor, ) -> None: # kernel implementation ... ``` **Example usage (operator with output):** ```python # out_shape indicates which argument has the same shape as the output @register_custom_op(mutates_args=["x"], out_shape=0) def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return x.add_(y) ``` For wrapping external library functions (e.g., FlashInfer kernels), use `register_custom_op_from_extern` instead. See `python/sglang/srt/utils/custom_op.py` for full API documentation. ## How it works ### Torch compile backend PCG uses `torch.compile` with a custom backend (`SGLangBackend`) to split and compile the model's forward pass. The flow is: ``` model.forward wrapper → torch.compile(..., backend=SGLangBackend) → FX graph → split_graph() at registered split ops → split_gm (top-level graph that chains the pieces) → replace capturable submodules with CUDAPiecewiseBackend → runtime dispatch: eager split ops + per-piece capture/replay ``` - **Install**: `install_torch_compiled()` replaces `model.forward` with a wrapper function. When `is_in_piecewise_cuda_graph()` returns True, the wrapper dispatches to the compiled callable; otherwise it falls back to the original forward. The first invocation through this path triggers Dynamo tracing and graph compilation — CUDA graph replay only happens after the capture phase completes. - **Split**: When `torch.compile` traces the model, `SGLangBackend` receives the FX graph and calls `split_graph()`. Ops listed in `CompilationConfig.split_ops` are treated as split points, so the graph is cut at each one. These split-op submodules are left to run eagerly at runtime, while the surrounding submodules are compiled and wrapped by `CUDAPiecewiseBackend`. The result is a top-level "stitching graph" (`split_gm`) with children such as `submod_0`, `submod_1`, … interleaving capturable subgraphs and eager split-op submodules. - **Replace**: `PiecewiseCompileInterpreter` iterates over each capturable submodule in `split_gm`, compiles it for general (dynamic) shapes, and replaces it in-place with a `CUDAPiecewiseBackend` instance. Split-op submodules (e.g., attention, all-reduce) are left as-is and run eagerly at runtime. - **Dispatch**: At runtime, calling `split_gm` executes the stitching graph, which calls each submodule in order. Split-op submodules run eagerly. Each `CUDAPiecewiseBackend` submodule goes through three phases: - **Compile warmup** — runs the general-shape compiled path. - **Capture** — for each capture size, runs one warmup pass then records a CUDA graph. - **Steady-state replay** — replays the captured CUDA graph for each forward pass. ### Piecewise cuda graph runner `PiecewiseCudaGraphRunner` orchestrates the full lifecycle through three phases: - **Compile** — Warms up JIT kernels with a dummy forward pass, then wraps the model with `torch.compile`, triggering Dynamo tracing to split the FX graph and create `CUDAPiecewiseBackend` instances for each subgraph piece. - **Capture** — Iterates over capture sizes in reverse order (largest first). For each size, runs the forward pass twice (one warmup, one CUDA graph capture). - **Replay** — At runtime, finds the smallest captured size >= actual token count via binary search, copies inputs into static buffers with zero-padding, replays the captured CUDA graphs, and slices outputs back to the actual token count. ### Memory optimization The memory cost of PCG comes from two parts: **torch memory allocator** and **non-torch memory**. The torch memory allocator overhead is trivial thanks to several optimizations: a global shared memory pool is reused across all CUDA graph runners and capture sizes, capture is done in reverse order (large to small) so smaller graphs reuse memory allocated by larger ones, and output tensors of the last subgraph are stored as weak references to maximize memory reuse. The main memory overhead comes from non-torch memory — the CUDA graph objects themselves require GPU memory to store the recorded kernel launch parameters and internal state. This overhead scales with the number of captured sizes, which is why `piecewise_cuda_graph_max_tokens` is capped conservatively by default. ### Shape configuration Piecewise CUDA graph pre-captures graphs for a set of token counts. At runtime, the actual token count is rounded up to the nearest captured size (via binary search), and the corresponding graph is replayed. If the token count exceeds the largest captured size, the runtime falls back to the normal (non-graph) forward path. The default capture schedule is auto-generated with increasing granularity: | Token range | Step size | |-------------|-----------| | 4 – 32 | 4 | | 48 – 256 | 16 | | 288 – 512 | 32 | | 576 – 1024 | 64 | | 1280 – 4096 | 256 | | 4096+ | 512 | For the auto-generated schedule, sizes are capped at `--piecewise-cuda-graph-max-tokens`. The default cap is `chunked_prefill_size` for non-MLA models and `2048` for MLA backend models. If `--max-total-tokens` is set, the cap is further limited to not exceed it. Additionally, Llama-2 models are auto-capped at 4096 tokens as a temporary workaround. ## Compatibility PCG is auto-disabled in the following scenarios. We are actively working on expanding compatibility — support for many of these will be coming soon. - Disabled model architectures (e.g., `DeepseekV32ForCausalLM`) - Speculative decoding - DP attention - Pipeline parallelism (`pp_size > 1`) - Non-CUDA hardware (AMD ROCm, Ascend NPU) - MoE A2A backend - LoRA - Multimodal / VLM models - DLLM (diffusion LLM) - Deterministic inference - PD disaggregation - Expert distribution recorder / EPLB Use `--enforce-piecewise-cuda-graph` to skip all auto-disable checks (for testing/debugging only). ## Code Reference | File | Description | |---|---| | `python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py` | Main runner: init, capture, replay | | `python/sglang/srt/compilation/compile.py` | `install_torch_compiled` trampoline | | `python/sglang/srt/compilation/backend.py` | `SGLangBackend`, graph splitting, piecewise compilation | | `python/sglang/srt/compilation/cuda_piecewise_backend.py` | Per-subgraph CUDA graph capture/replay | | `python/sglang/srt/compilation/piecewise_context_manager.py` | Global context flags and `ForwardContext` | | `python/sglang/srt/compilation/compilation_config.py` | Capture sizes, split ops, compiler config | | `python/sglang/srt/utils/custom_op.py` | `register_custom_op` for torch.compile compatibility | | `python/sglang/srt/server_args.py` | Server arguments and auto-disable logic | ================================================ FILE: docs/advanced_features/pipeline_parallelism.md ================================================ # Pipeline Parallelism for Long Context ## Why Pipeline Parallelism? As Large Language Models (LLMs) scale toward trillion-parameter architectures and "infinite" context windows, the underlying serving infrastructure must evolve toward more granular, cross-node parallelization strategies. While KV cache techniques effectively mitigate redundant computation, they cannot circumvent the prohibitive Time to First Token (TTFT) inherent in ultra-long sequences with extremely large initial Input Token Length (ITL). Although Tensor Parallelism (TP) remains the conventional approach for intra-node scaling, it frequently encounters communication bottlenecks during multi-node deployments. On the other hand, pipeline parallelism only requires cross-node communication at the boundaries of each pipeline stage, which can achieve better computation-communication overlap compared to a large TP. Therefore, it is also a promising parallelization strategy for improving throughput. Detailed analysis can be found in this [blog](https://lmsys.org/blog/2026-01-15-chunked-pipeline/). ## Implementation Refactoring based on Async Communication With Dynamic Chunked Prefill, pipeline parallelism has the potential to reduce the TTFT of long-context inputs. For each request, its input tokens can be partitioned into multiple chunks, each no longer than the chunked prefill size. Different chunks of the same request can be processed simultaneously by different nodes, thus parallelizing the processing and reducing TTFT. SGLang has supported Pipeline Parallelism (#5724) for some time and made it compatible with the PD Disaggregation feature (#8846), but the implementation was not perfect and had significant room for performance improvements. To eliminate this performance hazard, SGLang implements a Micro-batching Event Loop with non-blocking asynchronous peer-to-peer (P2P) communication to overlap GPU computation with CPU metadata processing and PP communication. This ensures that while one micro-batch is being computed on the GPU, the next one is already being prepared and moved into position effectively, ensuring the pipeline remains as saturated as possible. This approach was first proposed in #7979 and has been redesigned and included in #11852. The key mechanisms of the implementation include: * **Decoupled Sync/Async Logic in the Event Loop:** The scheduler uses `async_send` in `_pp_send_pyobj_to_next_stage`. Instead of waiting for a transfer to complete, it returns a `P2PWork` handle. The actual synchronization (`P2PWork.work.wait()`) is deferred until `_pp_commit_comm_work` is called, allowing the CPU to perform other work—like scheduling the next batch or processing metadata—while data is in flight. * **Multi-Stream Execution:** In addition to the main `default_stream`, which serves as the synchronization stream, SGLang utilizes dedicated `forward_stream` and `copy_stream` to execute forward pass GPU computation and Data-to-Host (D2H) memory transfers separately for better overlapping. While `_pp_launch_batch` is executing the current micro-batch on the GPU for the current stage, the CPU processes the previous micro-batch's results using `_pp_process_batch_result`. ## Guidance about Dynamic Chunking ### Why Dynamic Chunking Chunked prefill with a fixed size can cause bubbles in the pipeline, especially when the pp size is large. The main reason behind this phenomenon is that the model has a non-uniform running time, even though each chunk size is identical (brought by the Transformer structure). The larger the prefix sequence length, the longer the running time of the chunk. And these bubbles will be propagated to the next stage, and will significantly degrade the scale efficiency of larger pp ranks. To address this issue, SGLang introduces a dynamic chunking mechanism to predict the optimal size for the next chunk such that it satisfies this condition: Runtime(L + Next Chunk Size) - Runtime(L) = Runtime(Initial Chunk Size) where ***L*** denotes the Prefix Sequence Length. By profiling a series of requests with different ITLs, we model the cumulative runtime as a quadratic function of sequence length. Using this model, we solve the optimal next chunk size for any given prefix length ***L***. Since the computation complexity of the Attention mechanism scales with ***L***, the next chunk size will be progressively reduced as ***L*** grows to maintain an aligned chunk execution time across pipeline stages. Based on this method, the scheduler can predict and dynamically reduce the chunk size during runtime to minimize the bubbles caused by the stage misalignment. To be noticed, the scheduler does not use the raw predicted value. To facilitate efficient KVCache memory management and ensure affinity with hardware execution efficiency, the value is aligned downward to the nearest multiple of max(`--page-size`, 64). ### Chunked Prefill Size and Smoothing Factor When `--enable-dynamic-chunking` is enabled, each chunk size of a sequence is determined dynamically based on the quadratic model that predicts the next chunk size based on the estimated runtime of the initial chunk length. In this case, we use `--chunked-prefill-size` to set up the initial chunk size. When switching to the dynamic chunking mode, the initial chunk size (`--chunked-prefill-size`) should be set to a larger value comparable to the original chunked prefill size, so that there won't be too many chunks. **`SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR`** is an environmental variable that controls the smoothing factor for the dynamic chunking algorithm, defaulting to 0.75. It determines how much the chunk size can change during the prefill phase. A larger value means a more aggressive chunk size change, which may lead to better performance but also to greater chunk size changes (the chunk size at the end may become very small, which could lead to performance degradation) and more total chunks. When it is set to 1, the chunk size will be adjusted strictly based on the aforementioned quadratic model that predicts the next chunk size. A smaller value means a more conservative chunk size change, which may lead to smaller chunk size changes and fewer total chunks. When it is set to 0, the chunk size will not be adjusted dynamically, so it is identical to the traditional way with a fixed chunked prefill size. Due to the variation in hardware, models, and target workloads, a static configuration is seldom optimal across all scenarios. Consequently, achieving peak performance necessitates a degree of hyperparameter tuning when switching to the dynamic chunking mode. **Tuning Guidance for Dynamic Chunked Prefill** * **Step 1 \- Iterate to find the optimal fixed chunked prefill size for the targeted PP size**: Different PP sizes for targeted ITL may have different optimal chunked prefill sizes. Therefore, users should iterate to obtain the baseline according to the available resources for scaling. * **Step 2 \- Initial Chunk Size Selection for Dynamic Chunking**: Set the initial size to 2× or 3× the optimal fixed chunked prefill size. This reduces the total number of chunks and prevents "tail chunks" from underutilizing hardware. To maintain efficiency for extremely large Input Token Lengths (ITL), the dynamic predictor automatically ensures subsequent chunks are at least 1/4 of this initial size. In addition, it is recommended to use a larger initial chunk size (e.g., 4× the optimal fixed chunked prefill size) for such cases as well. * **Step 3 \- Smooth Factor Adjustment**: This factor controls how strictly the chunk size adjusts the prediction given by the quadratic performance fitting model. * 1.0: Follows the model strictly. * **0.6 – 0.85 (Recommended)**: Typical range for the best balance between dynamic scaling and hardware stability. Through experiments, we find that a range between 0.6 and 0.85 typically yields the best performance for dynamic chunking. * 0: Disables dynamic adjustment, reverting to traditional fixed-size chunking. * **Another small optimization tip:** Put the larger partition in the higher PP rank when the layers are not evenly divisible across ranks. It can increase the GPU utilization when a larger PP rank is waiting for the previous stage’s result, hence reducing the bubbles on higher PP ranks. If we take DeepSeek-V3.1 as an example, `SGLANG_PP_LAYER_PARTITION=15,15,15,16` usually performs better than `16,15,15,15`. ## Best Practice for Long Context ### Tuning the Chunked Prefill Size Optimizing the chunked prefill size is crucial for balancing pipeline efficiency and resource utilization. The ideal size depends on factors including model architecture, hardware configuration, and typical input lengths. We recommend starting with a small chunk size, such as 4K, and gradually increasing it until you find the optimal size for your specific use case (Different targeted ITL and PP Sizes may have different optimal chunked prefill sizes. Therefore, users should iterate to obtain the baseline according to the available resources for scaling). Alternatively, you can analyze the hardware capacity and determine the optimal chunk size based on the roofline model. ### Enable Dynamic Chunking and Adjust Smoothing Factor for Ultra-long ITL SGLang also offers a dynamic chunking solution that could further improve performance. This feature is currently an experimental feature that requires a certain amount of tuning experimentation and may not be suitable for all workloads. In addition, fine-tuning the smoothing factor can help optimize performance for specific workloads and model characteristics. ### Case Study on NVIDIA H20 When evaluating pipeline parallelism with fixed chunked prefill sizes from 2K to 16K, experiment results show that a 4K chunk size delivered optimal prefill TTFT performance for the DeepSeek-V3.1, and a 6K chunk size delivered optimal prefill TTFT performance for the Qwen3-235B-A22B-FP8. When enabling dynamic chunking, we first scale the optimal fixed chunked prefill size by a factor of 3 as the initial chunk size. Through experimentation, we found that a multiplier of 2-3 provides an appropriate balance—avoiding excessive initial pipeline bubbles while ensuring that subsequent chunks don't become too small as context length increases. With the default dynamic chunking smoothing factor of 0.75, we performed parameter tuning and determined that a value of 0.65 works optimally with the 12K initial chunk size for the DeepSeek-V3.1, while a value of 0.8 works optimally with the 18K initial chunk size for the Qwen3-235B-A22B-FP8. #### DeepSeek-V3.1 with 128K Input Token Length ```bash # prefill node 0 (fixed chunked prefill size) python3 -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3.1 --trust-remote-code \ --nnodes 4 --node-rank 0 --tp 8 --pp-size 4 \ --port 30000 --dist-init-addr \ --disable-radix-cache --mem-fraction-static 0.8 \ --attention-backend fa3 --host 0.0.0.0 --watchdog-timeout 3600 \ --max-running-requests 128 --chunked-prefill-size 4096 ``` ```bash # prefill node 0 (with dynamic chunking) export SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR=0.65 python3 -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3.1 --trust-remote-code \ --nnodes 4 --node-rank 0 --tp 8 --pp-size 4 \ --port 30000 --dist-init-addr \ --disable-radix-cache --mem-fraction-static 0.8 \ --attention-backend fa3 --host 0.0.0.0 --watchdog-timeout 3600 \ --max-running-requests 128 --chunked-prefill-size 12288 --enable-dynamic-chunking ``` #### Qwen3-235B-A22B-FP8 with 128K Input Token Length ```bash # prefill node 0 (fixed chunked prefill size) python3 -m sglang.launch_server \ --model-path Qwen/Qwen3-235B-A22B-FP8 --trust-remote-code \ --nnodes 4 --node-rank 0 --tp 4 --pp-size 8 \ --port 30000 --dist-init-addr \ --disable-radix-cache --mem-fraction-static 0.8 \ --attention-backend fa3 --host 0.0.0.0 --watchdog-timeout 3600 \ --max-running-requests 128 --chunked-prefill-size 6144 ``` ```bash # prefill node 0 (with dynamic chunking) export SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR=0.8 python3 -m sglang.launch_server \ --model-path Qwen/Qwen3-235B-A22B-FP8 --trust-remote-code \ --nnodes 4 --node-rank 0 --tp 4 --pp-size 8 \ --port 30000 --dist-init-addr \ --disable-radix-cache --mem-fraction-static 0.8 \ --attention-backend fa3 --host 0.0.0.0 --watchdog-timeout 3600 \ --max-running-requests 128 --chunked-prefill-size 18432 --enable-dynamic-chunking ``` Note: `--disable-radix-cache` is enabled only for reproducible benchmarking purposes. It is not recommended to use it in production. ## Best Practice for Pipeline Parallelism with PD Disaggregation To be added. Stay tuned for the latest updates on Pipeline Parallelism with PD Disaggregation. ================================================ FILE: docs/advanced_features/quantization.md ================================================ # Quantization SGLang supports various quantization methods, including offline quantization and online dynamic quantization. Offline quantization loads pre-quantized model weights directly during inference. This is required for quantization methods such as GPTQ and AWQ, which collect and pre-compute various statistics from the original weights using the calibration dataset. Online quantization dynamically computes scaling parameters—such as the maximum/minimum values of model weights—during runtime. Like NVIDIA FP8 training's [delayed scaling](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html#Mixed-precision-training-with-FP8) mechanism, online quantization calculates the appropriate scaling factors on-the-fly to convert high-precision weights into a lower-precision format. **Note: For better performance, usability and convenience, offline quantization is recommended over online quantization.** If you use a pre-quantized model, do not add `--quantization` to enable online quantization at the same time. For popular pre-quantized models, please visit [Unsloth](https://huggingface.co/unsloth), [NVIDIA ModelOpt](https://huggingface.co/collections/nvidia/inference-optimized-checkpoints-with-model-optimizer) or [NeuralMagic](https://huggingface.co/collections/neuralmagic) collections on HF for some popular quality validated quantized models. Quantized models must be validated via benchmarks post-quantization to guard against abnormal quantization loss regressions. ## Platform Compatibility The following table summarizes quantization method support across NVIDIA and AMD GPUs. | Method | NVIDIA GPUs | AMD GPUs (MI300X/MI325X/MI350X) | Notes | |--------|:-----------:|:-------------------------------:|-------| | `fp8` | Yes | Yes | Aiter or Triton backend on AMD | | `mxfp4` | Yes | Yes | Requires CDNA3/CDNA4 with MXFP support; uses Aiter | | `blockwise_int8` | Yes | Yes | Triton-based, works on both platforms | | `w8a8_int8` | Yes | Yes | | | `w8a8_fp8` | Yes | Yes | Aiter or Triton FP8 on AMD | | `awq` | Yes | Yes | Uses Triton dequantize on AMD (vs. optimized CUDA kernels on NVIDIA) | | `gptq` | Yes | Yes | Uses Triton or vLLM kernels on AMD | | `compressed-tensors` | Yes | Yes | Aiter paths for FP8/MoE on AMD | | `quark` | Yes | Yes | AMD Quark quantization; Aiter GEMM paths on AMD | | `auto-round` | Yes | Yes | Platform-agnostic (Intel auto-round) | | `quark_int4fp8_moe` | No | Yes | AMD-only; online INT4-to-FP8 MoE quantization (CDNA3/CDNA4) | | `awq_marlin` | Yes | No | Marlin kernels are CUDA-only | | `gptq_marlin` | Yes | No | Marlin kernels are CUDA-only | | `gguf` | Yes | No | CUDA-only kernels in sgl-kernel | | `modelopt` / `modelopt_fp8` | Yes (Hopper/SM90+) | No | [NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer); requires NVIDIA hardware | | `modelopt_fp4` | Yes (Blackwell/SM100+) | No | [NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer); native FP4 on Blackwell (B200, GB200) | | `petit_nvfp4` | No | Yes (MI250/MI300X/MI325X) | Enables NVFP4 on ROCm via [Petit](https://github.com/causalflow-ai/petit-kernel); use `modelopt_fp4` on NVIDIA Blackwell. Auto-selected when loading NVFP4 models on AMD. See [LMSYS blog](https://lmsys.org/blog/2025-09-21-petit-amdgpu/) and [AMD ROCm blog](https://rocm.blogs.amd.com/artificial-intelligence/fp4-mixed-precision/README.html). | | `bitsandbytes` | Yes | Experimental | Depends on bitsandbytes ROCm support | | `torchao` (`int4wo`, etc.) | Yes | Partial | `int4wo` not supported on AMD; other methods may work | On AMD, several of these methods use [Aiter](https://github.com/ROCm/aiter) for acceleration -- set `SGLANG_USE_AITER=1` where noted. See [AMD GPU setup](../platforms/amd_gpu.md) for installation and configuration details. ## GEMM Backends for FP4/FP8 Quantization :::{note} Backend selection is supported only for **blockwise FP8** and **NVFP4** GEMM. When running FP8 or FP4 quantized models, you can select the GEMM backend via `--fp8-gemm-backend` and `--fp4-gemm-backend`. ::: ### `--fp8-gemm-backend` (Blockwise FP8 GEMM) | Backend | Hardware | Description | |---------|----------|-------------| | `auto` | All | Auto-selects based on hardware | | `deep_gemm` | SM90, SM100 | JIT-compiled; enabled when DeepGEMM is installed | | `flashinfer_trtllm` | SM100 | FlashInfer TensorRT-LLM backend; optimal for low-latency | | `flashinfer_cutlass` | SM100/120 | FlashInfer CUTLASS groupwise FP8 GEMM | | `flashinfer_deepgemm` | SM90 | Uses swapAB optimization for small M dimensions in decoding | | `cutlass` | SM90, SM100/120 | sgl-kernel CUTLASS | | `triton` | All | Fallback; widely compatible | | `aiter` | ROCm | AMD AITER backend | **`auto` selection order:** 1) DeepGEMM (SM90/SM100, installed); 2) FlashInfer TRTLLM (SM100, FlashInfer available); 3) CUTLASS (SM90/SM100/120); 4) AITER (AMD); 5) Triton. **Exception:** SM120 always resolves to Triton. ### `--fp4-gemm-backend` (NVFP4 GEMM) | Backend | Hardware | Description | |---------|----------|-------------| | `auto` | SM100/120 | Auto-selects: `flashinfer_cudnn` on SM120; `flashinfer_cutlass` on SM100 | | `flashinfer_cutlass` | SM100/120 | FlashInfer CUTLASS backend | | `flashinfer_cudnn` | SM100/120 (CUDA 13+, cuDNN 9.15+) | FlashInfer cuDNN backend; used on SM120 for performance | | `flashinfer_trtllm` | SM100 | FlashInfer TensorRT-LLM backend | When FlashInfer is unavailable for NVFP4, sgl-kernel CUTLASS is used as an automatic fallback. ## Offline Quantization To load already quantized models, simply load the model weights and config. **Again, if the model has been quantized offline, there's no need to add `--quantization` argument when starting the engine. The quantization method will be parsed from the downloaded Hugging Face config. For example, DeepSeek V3/R1 models are already in FP8, so do not add redundant parameters.** ```bash python3 -m sglang.launch_server \ --model-path hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4 \ --port 30000 --host 0.0.0.0 ``` Take note, if your model is **per-channel quantized (INT8 or FP8) with per-token dynamic quantization activation**, you can opt to include `--quantization w8a8_int8` or `--quantization w8a8_fp8` to invoke the corresponding CUTLASS int8_kernel or fp8_kernel in sgl-kernel. This action will ignore the Hugging Face config's quantization settings. For instance, with `neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8-dynamic`, if you execute with `--quantization w8a8_fp8`, the system will use the `W8A8Fp8Config` from SGLang to invoke the sgl-kernel, rather than the `CompressedTensorsConfig` for vLLM kernels. ```bash python3 -m sglang.launch_server \ --model-path neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8-dynamic \ --quantization w8a8_fp8 \ --port 30000 --host 0.0.0.0 ``` ### Examples of Offline Model Quantization #### Using [Unsloth](https://docs.unsloth.ai/basics/inference-and-deployment/sglang-guide) We strongly suggest the use of Unsloth to quantize and load the model. Please refer to [SGLang Deployment & Inference Guide with Unsloth](https://docs.unsloth.ai/basics/inference-and-deployment/sglang-guide). #### Using [auto-round](https://github.com/intel/auto-round) ```bash # Install pip install auto-round ``` - LLM quantization ```py # for LLM from auto_round import AutoRound model_id = "meta-llama/Llama-3.2-1B-Instruct" quant_path = "Llama-3.2-1B-Instruct-autoround-4bit" # Scheme examples: "W2A16", "W3A16", "W4A16", "W8A16", "NVFP4", "MXFP4" (no real kernels), "GGUF:Q4_K_M", etc. scheme = "W4A16" format = "auto_round" autoround = AutoRound(model_id, scheme=scheme) autoround.quantize_and_save(quant_path, format=format) # quantize and save ``` - VLM quantization ```py # for VLMs from auto_round import AutoRoundMLLM model_name = "Qwen/Qwen2-VL-2B-Instruct" quant_path = "Qwen2-VL-2B-Instruct-autoround-4bit" scheme = "W4A16" format = "auto_round" autoround = AutoRoundMLLM(model_name, scheme) autoround.quantize_and_save(quant_path, format=format) # quantize and save ``` - Command Line Usage (Gaudi/CPU/Intel GPU/CUDA) ```bash auto-round \ --model meta-llama/Llama-3.2-1B-Instruct \ --bits 4 \ --group_size 128 \ --format "auto_round" \ --output_dir ./tmp_autoround ``` - known issues Several limitations currently affect offline quantized model loading in sglang, These issues might be resolved in future updates of sglang. If you experience any problems, consider using Hugging Face Transformers as an alternative. 1. Mixed-bit Quantization Limitations Mixed-bit quantization is not fully supported. Due to vLLM's layer fusion (e.g., QKV fusion), applying different bit-widths to components within the same fused layer can lead to compatibility issues. 2. Limited Support for Quantized MoE Models Quantized MoE models may encounter inference issues due to kernel limitations (e.g., lack of support for mlp.gate layer quantization). please try to skip quantizing these layers to avoid such errors. 3. Limited Support for Quantized VLMs
VLM failure cases Qwen2.5-VL-7B auto_round:auto_gptq format: Accuracy is close to zero. GPTQ format: Fails with: ``` The output size is not aligned with the quantized weight shape ``` auto_round:auto_awq and AWQ format: These work as expected.
#### Using [GPTQModel](https://github.com/ModelCloud/GPTQModel) ```bash # install pip install gptqmodel --no-build-isolation -v ``` ```py from datasets import load_dataset from gptqmodel import GPTQModel, QuantizeConfig model_id = "meta-llama/Llama-3.2-1B-Instruct" quant_path = "Llama-3.2-1B-Instruct-gptqmodel-4bit" calibration_dataset = load_dataset( "allenai/c4", data_files="en/c4-train.00001-of-01024.json.gz", split="train" ).select(range(1024))["text"] quant_config = QuantizeConfig(bits=4, group_size=128) # quantization config model = GPTQModel.load(model_id, quant_config) # load model model.quantize(calibration_dataset, batch_size=2) # quantize model.save(quant_path) # save model ``` #### Using [LLM Compressor](https://github.com/vllm-project/llm-compressor/) ```bash # install pip install llmcompressor ``` Here, we take quantize `meta-llama/Meta-Llama-3-8B-Instruct` to `FP8` as an example to elaborate on how to do offline quantization. ```python from transformers import AutoTokenizer from llmcompressor.transformers import SparseAutoModelForCausalLM from llmcompressor.transformers import oneshot from llmcompressor.modifiers.quantization import QuantizationModifier # Step 1: Load the original model. MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct" model = SparseAutoModelForCausalLM.from_pretrained( MODEL_ID, device_map="auto", torch_dtype="auto") tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) # Step 2: Perform offline quantization. # Step 2.1: Configure the simple PTQ quantization. recipe = QuantizationModifier( targets="Linear", scheme="FP8_DYNAMIC", ignore=["lm_head"]) # Step 2.2: Apply the quantization algorithm. oneshot(model=model, recipe=recipe) # Step 3: Save the model. SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-Dynamic" model.save_pretrained(SAVE_DIR) tokenizer.save_pretrained(SAVE_DIR) ``` Then, you can directly use the quantized model with `SGLang`, by using the following command: ```bash python3 -m sglang.launch_server \ --model-path $PWD/Meta-Llama-3-8B-Instruct-FP8-Dynamic \ --port 30000 --host 0.0.0.0 ``` #### Using [NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer) NVIDIA Model Optimizer (ModelOpt) provides advanced quantization techniques optimized for NVIDIA hardware. **Offline vs. Online Quantization:** SGLang supports two modes for ModelOpt. * **Offline Quantization (pre-quantized):** * **Usage:** Download a pre-quantized model from Hugging Face or run `hf_ptq.py` once to create a new quantized checkpoint. Then load this quantized checkpoint. * **Pros:** Fast server startup, quantization can be validated before deployment, efficient resource usage. * **Cons:** Requires an extra preparation step. * **Online Quantization (quant and serve):** * **Usage:** Load a standard BF16/FP16 model and add a flag. The engine applies quantization *on startup*. * **Pros:** Convenient (no new checkpoint needed). * **Cons:** **High startup time**, increases VRAM usage during initialization (risk of OOM). The following sections guide you through using the Offline path: loading pre-quantized models or creating your own checkpoints. ##### Using Pre-Quantized Checkpoints If a model is already quantized (e.g., from Hugging Face), you can load it directly. * **FP8 Models:** Use `--quantization modelopt_fp8`. ```bash python3 -m sglang.launch_server \ --model-path nvidia/Llama-3.1-8B-Instruct-FP8 \ --quantization modelopt_fp8 \ --port 30000 ``` * **FP4 Models:** Use `--quantization modelopt_fp4`. ```bash python3 -m sglang.launch_server \ --model-path nvidia/Llama-3.3-70B-Instruct-NVFP4 \ --quantization modelopt_fp4 \ --port 30000 ``` ##### Creating Your Own Quantized Checkpoints If a pre-quantized checkpoint is not available for your model, you can create one using NVIDIA Model Optimizer's `hf_ptq.py` script. **Why quantize?** - Reduce VRAM usage - Higher throughput and lower latency - More flexible deployment (on smaller GPUs) **What can be quantized?** - The entire model - MLP layers only - KV cache **Key options in `hf_ptq.py`:** `--qformat`: Quantization formats `fp8`, `nvfp4`, `nvfp4_mlp_only` `--kv_cache_qformat`: KV cache quantization format (default: `fp8`) **Note:** The default `kv_cache_qformat` may not be optimal for all use cases. Consider setting this explicitly. **Hardware requirements:** Hopper and higher are recommended. Insufficient GPU memory may cause weight offloading, resulting in extremely long quantization time. For detailed usage and supported model architectures, see [NVIDIA Model Optimizer LLM PTQ](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_ptq). SGLang includes a streamlined workflow for quantizing models with ModelOpt and automatically exporting them for deployment. ##### Installation First, install ModelOpt: ```bash pip install nvidia-modelopt ``` ##### Quantization and Export Workflow SGLang provides an example script that demonstrates the complete ModelOpt quantization and export workflow. Run from the SGLang repository root (see [modelopt_quantize_and_export.py](https://github.com/sgl-project/sglang/blob/main/examples/usage/modelopt_quantize_and_export.py)): ```bash # Quantize and export a model using ModelOpt FP8 quantization python examples/usage/modelopt_quantize_and_export.py quantize \ --model-path TinyLlama/TinyLlama-1.1B-Chat-v1.0 \ --export-dir ./quantized_tinyllama_fp8 \ --quantization-method modelopt_fp8 # For FP4 quantization (requires Blackwell GPU) python examples/usage/modelopt_quantize_and_export.py quantize \ --model-path TinyLlama/TinyLlama-1.1B-Chat-v1.0 \ --export-dir ./quantized_tinyllama_fp4 \ --quantization-method modelopt_fp4 ``` ##### Available Quantization Methods - `modelopt_fp8`: FP8 quantization with optimal performance on NVIDIA Hopper and Blackwell GPUs - `modelopt_fp4`: FP4 quantization with optimal performance on Nvidia Blackwell GPUs ##### Python API Usage You can also use ModelOpt quantization programmatically: ```python import sglang as sgl from sglang.srt.configs.device_config import DeviceConfig from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.model_config import ModelConfig from sglang.srt.model_loader.loader import get_model_loader # Configure model with ModelOpt quantization and export model_config = ModelConfig( model_path="TinyLlama/TinyLlama-1.1B-Chat-v1.0", quantization="modelopt_fp8", # or "modelopt_fp4" trust_remote_code=True, ) load_config = LoadConfig( modelopt_export_path="./exported_model", modelopt_checkpoint_save_path="./checkpoint.pth", # optional, fake quantized checkpoint ) device_config = DeviceConfig(device="cuda") # Load and quantize the model (export happens automatically) model_loader = get_model_loader(load_config, model_config) quantized_model = model_loader.load_model( model_config=model_config, device_config=device_config, ) ``` ##### Deploying Quantized Models After quantization and export, you can deploy the model with SGLang: ```bash # Deploy the exported quantized model python -m sglang.launch_server \ --model-path ./quantized_tinyllama_fp8 \ --quantization modelopt \ --port 30000 --host 0.0.0.0 ``` Or using the Python API (use the same path as `modelopt_export_path` from the quantize step): ```python import sglang as sgl def main(): # Deploy exported ModelOpt quantized model # Path must match modelopt_export_path from quantize step (e.g., ./exported_model) llm = sgl.Engine( model_path="./exported_model", quantization="modelopt", ) # Run inference prompts = [ "Hello, how are you?", "What is the capital of France?", ] sampling_params = { "temperature": 0.8, "top_p": 0.95, "max_new_tokens": 100, } outputs = llm.generate(prompts, sampling_params) for i, output in enumerate(outputs): print(f"Prompt: {prompts[i]}") print(f"Output: {output['text']}") if __name__ == "__main__": main() ``` ##### Advanced Features **Checkpoint Management**: Save and restore fake quantized checkpoints for reuse: ```bash # Save the fake quantized checkpoint during quantization python examples/usage/modelopt_quantize_and_export.py quantize \ --model-path meta-llama/Llama-3.2-1B-Instruct \ --export-dir ./quantized_model \ --quantization-method modelopt_fp8 \ --checkpoint-save-path ./my_checkpoint.pth # The checkpoint can be reused for future quantization runs and skip calibration ``` **Export-only Workflow**: If you have a pre-existing fake quantized ModelOpt checkpoint, you can export it directly. See [LoadConfig](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/configs/load_config.py) for the full API: ```python from sglang.srt.configs.device_config import DeviceConfig from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.model_config import ModelConfig from sglang.srt.model_loader.loader import get_model_loader model_config = ModelConfig( model_path="meta-llama/Llama-3.2-1B-Instruct", quantization="modelopt_fp8", trust_remote_code=True, ) load_config = LoadConfig( modelopt_checkpoint_restore_path="./my_checkpoint.pth", modelopt_export_path="./exported_model", ) # Load and export the model (DeviceConfig defaults to device="cuda") model_loader = get_model_loader(load_config, model_config) model_loader.load_model(model_config=model_config, device_config=DeviceConfig()) ``` ##### Benefits of ModelOpt - **Hardware Optimization**: Specifically optimized for NVIDIA GPU architectures - **Advanced Quantization**: Supports cutting-edge FP8 and FP4 quantization techniques - **Seamless Integration**: Automatic export to HuggingFace format for easy deployment - **Calibration-based**: Uses calibration datasets for optimal quantization quality - **Production Ready**: Enterprise-grade quantization with NVIDIA support ## Online Quantization To enable online quantization, you can simply specify `--quantization` in the command line. For example, you can launch the server with the following command to enable `FP8` quantization for model `meta-llama/Meta-Llama-3.1-8B-Instruct`: ```bash python3 -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ --quantization fp8 \ --port 30000 --host 0.0.0.0 ``` Our team is working on supporting more online quantization methods. SGLang will soon support methods including but not limited to `["awq", "gptq", "marlin", "gptq_marlin", "awq_marlin", "bitsandbytes", "gguf"]`. ### torchao online quantization method SGLang also supports quantization methods based on [torchao](https://github.com/pytorch/ao). You can simply specify `--torchao-config` in the command line to support this feature. For example, if you want to enable `int4wo-128` for model `meta-llama/Meta-Llama-3.1-8B-Instruct`, you can launch the server with the following command: ```bash python3 -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ --torchao-config int4wo-128 \ --port 30000 --host 0.0.0.0 ``` SGLang supports the following quantization methods based on torchao `["int8dq", "int8wo", "fp8wo", "fp8dq-per_tensor", "fp8dq-per_row", "int4wo-32", "int4wo-64", "int4wo-128", "int4wo-256"]`. Note: According to [this issue](https://github.com/sgl-project/sglang/issues/2219#issuecomment-2561890230), `"int8dq"` method currently has some bugs when using together with cuda graph capture. So we suggest to disable cuda graph capture when using `"int8dq"` method. Namely, please use the following command: ```bash python3 -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ --torchao-config int8dq \ --disable-cuda-graph \ --port 30000 --host 0.0.0.0 ``` ### `quark_int4fp8_moe` online quantization method SGLang running on AMD GPUs (CDNA3 or CDNA4 architecture) supports the quantization method `--quantization quark_int4fp8_moe`, that will replace [MoE layers](https://github.com/sgl-project/sglang/blob/v0.4.8/python/sglang/srt/layers/moe/fused_moe_triton/layer.py#L271) originally in high precision (bfloat16, float16 or float32) to use weights dynamically quantized to int4, that are upcasted to float8 during inference to run compute in float8 precision with activations dynamically quantized on the fly to float8. Other layers (e.g. projections in the attention layers) have their weights quantized online to float8 directly. ## Reference - [GPTQModel](https://github.com/ModelCloud/GPTQModel) - [LLM Compressor](https://github.com/vllm-project/llm-compressor/) - [NVIDIA Model Optimizer (ModelOpt)](https://github.com/NVIDIA/Model-Optimizer) - [NVIDIA Model Optimizer LLM PTQ](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_ptq) - [Petit: NVFP4 on ROCm](https://github.com/causalflow-ai/petit-kernel) — [LMSYS blog](https://lmsys.org/blog/2025-09-21-petit-amdgpu/), [AMD ROCm blog](https://rocm.blogs.amd.com/artificial-intelligence/fp4-mixed-precision/README.html) - [Torchao: PyTorch Architecture Optimization](https://github.com/pytorch/ao) - [vLLM Quantization](https://docs.vllm.ai/en/latest/quantization/) - [auto-round](https://github.com/intel/auto-round) ================================================ FILE: docs/advanced_features/quantized_kv_cache.md ================================================ # Quantized KV Cache Quantized KV cache reduces the memory footprint of key-value cache storage by using lower-precision data types (FP8 or FP4) instead of the default model precision in BF16. During autoregressive generation, LLMs cache previously computed key-value pairs to avoid redundant calculations. The KV cache typically consumes a significant portion of GPU memory, especially for long sequences. Quantized KV cache is a memory optimization technique that primarily benefits throughput by allowing more tokens to be cached, but may introduce minimal accuracy degradation depending on the quantization format used. ```{warning} **Performance Warning**: When quantized KV cache must be dequantized before use in attention operations, performance can be extremely slow if dequantization is not fused with the attention kernel. Always verify that your chosen attention backend supports quantized KV cache. Backends without fused support may experience significant throughput degradation, potentially negating the memory benefits. **Backend Support**: Not all attention backends support quantized KV cache. Refer to [Attention Backend](attention_backend.md) for which backends support it. ``` ## Supported Formats SGLang supports the following quantized KV cache formats: ### FP8 Format [OCP (Open Compute Project)](https://www.opencompute.org) specifies two common 8-bit floating point formats: - **E5M2** (5 exponent bits, 2 mantissa bits): Larger dynamic range (±57344.0), lower precision - **E4M3** (4 exponent bits, 3 mantissa bits): Higher precision, smaller dynamic range (±240.0) ### FP4 Format ```{warning} FP4 quantization is currently experimental. ``` [OCP (Open Compute Project)](https://www.opencompute.org) specifies MXFP4 (Microscaling FP4), a 4-bit floating-point format: - **E2M1** (1 sign bit, 2 exponent bits, 1 mantissa bit): Uses block-based microscaling where tensors are divided into blocks of consecutive elements, with each block sharing a single 8-bit exponential scaling factor. While OCP specifies blocks of 32 elements, SGLang's current implementation uses blocks of 16 elements for KV cache quantization. ## Usage ### Enabling Quantized KV Cache To enable quantized KV cache, use the `--kv-cache-dtype` argument when launching the server: ```bash # Enable FP8 E5M2 KV cache python3 -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-R1-0528 \ --kv-cache-dtype fp8_e5m2 \ # Enable FP8 E4M3 KV cache python3 -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-R1-0528 \ --kv-cache-dtype fp8_e4m3 \ # Enable FP4 E2M1 KV cache python3 -m sglang.launch_server \ --model-path nvidia/DeepSeek-R1-0528-NVFP4 \ --kv-cache-dtype fp4_e2m1 \ ``` ### Scaling Factors FP8 quantization requires scaling factors to properly quantize and dequantize the KV cache. ```{note} Currently, only per-tensor (scalar) scaling factors are supported. ``` Scaling factors can be: - **Loaded from checkpoints**: Pre-quantized models (e.g., ModelOpt) may include `k_scale` and `v_scale` parameters that are automatically loaded - **Provided via JSON**: Supply scaling factors via `--quantization-param-path`. The JSON file should follow this format: ```json { "kv_cache": { "dtype": "float8_e4m3fn", "scaling_factor": { "0": { "0": 1.0, "1": 1.0 } } } } ``` Where the outer keys in `scaling_factor` are tensor parallel ranks and inner keys are layer indices. ```{warning} If scaling factors are not provided and not found in the checkpoint, it will default to 1.0, which may cause accuracy issues. ``` ```{tip} **FP4 (MXFP4)**: Unlike FP8, FP4 quantization handles scaling factors automatically on-the-fly during quantization and dequantization. No pre-quantized models or external scaling factor files are required—the block-based scaling factors are computed dynamically as needed. ``` ## Performance Considerations ### Memory Savings Quantized KV cache provides significant memory savings: - **BF16 → FP4**: Supports approximately 3.56× more tokens than BF16 (accounting for scaling factor overhead) ```{note} FP4 and FP8 quantization require additional memory for block-based scaling factors, which reduces the effective memory savings compared to the raw bit-width reduction. FP4 with block size 16 supports approximately 1.78× more tokens than FP8, and approximately 3.56× more tokens than BF16. The relative token capacity between FP8 and BF16 can be derived from these ratios. ``` This enables longer context lengths or more concurrent requests within the same memory budget. ### Accuracy Impact #### FP8 Accuracy FP8 E4M3 quantization typically introduces minimal accuracy degradation. The impact depends on model architecture, sequence length, and quantization format (generally, E4M3 has better accuracy than E5M2). #### FP4 Accuracy FP4 (MXFP4) quantization provides significant memory savings with varying accuracy impact depending on model size and dataset complexity. Preliminary accuracy test results from [PR #10078](https://github.com/sgl-project/sglang/pull/10078) (MLA) and [PR #12612](https://github.com/sgl-project/sglang/pull/12612) (MHA) show: **Large Models (e.g., Qwen3-235B-A22B, DeepSeek-R1-0528)** On large-scale models, FP4 maintains accuracy close to FP8/BF16, especially on simpler datasets: | Model | Dataset | KV16 | KV8 (FP8 E4M3) | KV4 (FP4 E2M1) | |-------|---------|------|----------------|----------------| | Qwen3-235B-A22B | gsm8k | 0.9168 | 0.9181 | 0.9186 | | Qwen3-235B-A22B | aime25 | 0.7733 | 0.7333 | 0.6000 | | Qwen3-235B-A22B | gpqa_diamond | 0.7010 | 0.6899 | 0.6778 | | DeepSeek-R1-0528 | gsm8k | 0.9157 | 0.9154 | 0.9124 | | DeepSeek-R1-0528 | aime25 | 0.5067 | 0.4934 | 0.4000 | | DeepSeek-R1-0528 | gpqa_diamond | 0.7707 | 0.7697 | 0.7273 | **Smaller Models (e.g., GPT-OSS-120B)** On smaller models, FP4 shows more pronounced accuracy drops, particularly on challenging datasets: | Model | Dataset | KV16 | KV8 (FP8 E4M3) | KV4 (FP4 E2M1) | |-------|---------|------|----------------|----------------| | GPT-OSS-120B | gsm8k | 0.9161 | 0.9163 | 0.9152 | | GPT-OSS-120B | aime25 | 0.7533 | 0.7667 | 0.3533 | | GPT-OSS-120B | gpqa_diamond | 0.5081 | 0.5434 | 0.3202 | **Key Observations:** - **Simple datasets (e.g., gsm8k)**: FP4 maintains accuracy close to FP8/BF16 across model sizes - **Model size matters**: Large models (200B+ parameters) generally tolerate FP4 quantization better than smaller models - **Context length**: Accuracy degradation may be more pronounced in long-context scenarios, as the accumulation of the quantization error may become significant. ```{tip} Evaluate FP4 accuracy on your specific model and workload. Large models on simpler tasks typically show minimal degradation, while smaller models or complex reasoning tasks may require FP8 or BF16 for acceptable accuracy. ``` ## Best Practices - **Use pre-quantized models**: Prefer models quantized offline with scaling factors included in the checkpoint. - **Choose the right format**: Use `fp8_e4m3` for better accuracy (recommended), `fp8_e5m2` for larger dynamic range, or `fp4_e2m1` for maximum memory savings (experimental) - **Check backend compatibility**: Verify that your chosen attention backend supports quantized KV cache ```{seealso} - [Quantization](quantization.md) - [Attention Backend](attention_backend.md) - [Server Arguments](server_arguments.md) ``` ================================================ FILE: docs/advanced_features/rfork.md ================================================ # R-Fork R-Fork (Tensor Remote Fork) is a novel weight loading methodology that leverages efficient inter-node GPU-to-GPU data transfer path to load tensors from a running SGLang instance to a new instance with zero-copy. It can significantly optimize the SGLang instance boot-up time by reducing model weights loading from several minutes to mere seconds. To learn more details about R-Fork, please check ** R-Fork blog ** ## Usage | Argument | Usage | |--------------|--------------------------------------------| | load-format | set to `remote_instance` to enable R-Fork. | | remote-instance-weight-loader-backend | `nccl`, `transfer_engine`, or `modelexpress`. Default is `nccl`. | | remote-instance-weight-loader-seed-instance-ip | IP address of the seed instance who will provide the model weight. Used by `nccl` and `transfer_engine` backends. | | remote-instance-weight-loader-seed-instance-service-port | the port that the seed instance's HTTP server is listening on. Used by `nccl` and `transfer_engine` backends. | | remote-instance-weight-loader-send-weights-group-ports | the list of available ports on the seed instance that will be used to build NCCL communication groups between seed and client instance. Only needed by `nccl` backend. | | remote-instance-weight-loader-start-seed-via-transfer-engine | set to start seed service that supports TransferEngine as backend. Needed for seed instances when using `transfer_engine` as backend. | | modelexpress-config | JSON config for `modelexpress` backend. Keys: `"url"` (required, gRPC host:port of ModelExpress server), `"model_name"` (optional, defaults to `--model-path`), `"source"` (optional bool, `true` for seed mode). | ### NCCL as backend seed instance: ```shell python -m sglang.launch_server [args] ``` client instance: ```shell python -m sglang.launch_server [args] \ --load-format remote_instance \ --remote-instance-weight-loader-seed-instance-ip [seed_instance_ip] \ --remote-instance-weight-loader-seed-instance-service-port [seed_instance_service_port] \ --remote-instance-weight-loader-send-weights-group-ports [send_weights_nccl_group_ports_list] \ --remote-instance-weight-loader-backend nccl ``` ### TransferEngine as backend seed instance: ```shell python -m sglang.launch_server [args] \ --remote-instance-weight-loader-start-seed-via-transfer-engine ``` ```shell python -m sglang.launch_server [args] \ --load-format remote_instance \ --remote-instance-weight-loader-seed-instance-ip [seed_instance_ip] \ --remote-instance-weight-loader-seed-instance-service-port [seed_instance_service_port] \ --remote-instance-weight-loader-backend transfer_engine ``` ### ModelExpress as backend [ModelExpress](https://github.com/ai-dynamo/modelexpress) is a coordination service that manages P2P weight transfer metadata. It removes the need for direct seed IP/port configuration by providing a centralized registry that seeds publish to and clients discover from. Under the hood it uses TransferEngine (Mooncake) for the actual RDMA data transfer. A running ModelExpress server is required. See the [ModelExpress documentation](https://github.com/ai-dynamo/modelexpress) for setup instructions. seed instance: ```shell python -m sglang.launch_server [args] \ --modelexpress-config '{"url": "[modelexpress_grpc_host:port]", "model_name": "[model_name]", "source": true}' ``` client instance: ```shell python -m sglang.launch_server [args] \ --load-format remote_instance \ --remote-instance-weight-loader-backend modelexpress \ --modelexpress-config '{"url": "[modelexpress_grpc_host:port]", "model_name": "[model_name]"}' ``` The seed publishes its TransferEngine session ID and tensor layout to ModelExpress. The client queries ModelExpress to discover the seed, then pulls weights directly via RDMA. This enables dynamic seed discovery without hardcoding IPs, and supports multiple models through a single ModelExpress instance. ================================================ FILE: docs/advanced_features/separate_reasoning.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Reasoning Parser\n", "\n", "SGLang supports parsing reasoning content out from \"normal\" content for reasoning models such as [DeepSeek R1](https://huggingface.co/deepseek-ai/DeepSeek-R1).\n", "\n", "## Supported Models & Parsers\n", "\n", "| Model | Reasoning tags | Parser | Notes |\n", "|---------|-----------------------------|------------------|-------|\n", "| [DeepSeek‑R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `` … `` | `deepseek-r1` | Supports all variants (R1, R1-0528, R1-Distill) |\n", "| [DeepSeek‑V3 series](https://huggingface.co/deepseek-ai/DeepSeek-V3.1) | `` … `` | `deepseek-v3` | Including [DeepSeek‑V3.2](https://huggingface.co/deepseek-ai/DeepSeek-V3.2-Exp). Supports `thinking` parameter |\n", "| [Standard Qwen3 models](https://huggingface.co/collections/Qwen/qwen3-67dd247413f0e2e4f653967f) | `` … `` | `qwen3` | Supports `enable_thinking` parameter |\n", "| [Qwen3-Thinking models](https://huggingface.co/Qwen/Qwen3-235B-A22B-Thinking-2507) | `` … `` | `qwen3` or `qwen3-thinking` | Always generates thinking content |\n", "| [Kimi K2 Thinking](https://huggingface.co/moonshotai/Kimi-K2-Thinking) | `◁think▷` … `◁/think▷` | `kimi_k2` | Uses special thinking delimiters. Also requires `--tool-call-parser kimi_k2` for tool use. |\n", "| [GPT OSS](https://huggingface.co/openai/gpt-oss-120b) | `<\\|channel\\|>analysis<\\|message\\|>` … `<\\|end\\|>` | `gpt-oss` | N/A |\n", "### Model-Specific Behaviors\n", "\n", "**DeepSeek-R1 Family:**\n", "- DeepSeek-R1: No `` start tag, jumps directly to thinking content\n", "- DeepSeek-R1-0528: Generates both `` start and `` end tags\n", "- Both are handled by the same `deepseek-r1` parser\n", "\n", "**DeepSeek-V3 Family:**\n", "- DeepSeek-V3.1/V3.2: Hybrid model supporting both thinking and non-thinking modes, use the `deepseek-v3` parser and `thinking` parameter (NOTE: not `enable_thinking`)\n", "\n", "**Qwen3 Family:**\n", "- Standard Qwen3 (e.g., Qwen3-2507): Use `qwen3` parser, supports `enable_thinking` in chat templates\n", "- Qwen3-Thinking (e.g., Qwen3-235B-A22B-Thinking-2507): Use `qwen3` or `qwen3-thinking` parser, always thinks\n", "\n", "**Kimi K2:**\n", "- Kimi K2 Thinking: Uses special `◁think▷` and `◁/think▷` tags. For agentic tool use, also specify `--tool-call-parser kimi_k2`.\n", "\n", "**GPT OSS:**\n", "- GPT OSS: Uses special `<|channel|>analysis<|message|>` and `<|end|>` tags" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Usage\n", "\n", "### Launching the Server" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Specify the `--reasoning-parser` option." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import requests\n", "from openai import OpenAI\n", "from sglang.test.doc_patch import launch_server_cmd\n", "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", "\n", "server_process, port = launch_server_cmd(\n", " \"python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-7B --host 0.0.0.0 --reasoning-parser deepseek-r1 --log-level warning\"\n", ")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Note that `--reasoning-parser` defines the parser used to interpret responses." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### OpenAI Compatible API\n", "\n", "Using the OpenAI compatible API, the contract follows the [DeepSeek API design](https://api-docs.deepseek.com/guides/reasoning_model) established with the release of DeepSeek-R1:\n", "\n", "- `reasoning_content`: The content of the CoT.\n", "- `content`: The content of the final answer." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Initialize OpenAI-like client\n", "client = OpenAI(api_key=\"None\", base_url=f\"http://0.0.0.0:{port}/v1\")\n", "model_name = client.models.list().data[0].id\n", "\n", "messages = [\n", " {\n", " \"role\": \"user\",\n", " \"content\": \"What is 1+3?\",\n", " }\n", "]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Non-Streaming Request" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response_non_stream = client.chat.completions.create(\n", " model=model_name,\n", " messages=messages,\n", " temperature=0.6,\n", " top_p=0.95,\n", " stream=False, # Non-streaming\n", " extra_body={\"separate_reasoning\": True},\n", ")\n", "print_highlight(\"==== Reasoning ====\")\n", "print_highlight(response_non_stream.choices[0].message.reasoning_content)\n", "\n", "print_highlight(\"==== Text ====\")\n", "print_highlight(response_non_stream.choices[0].message.content)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Streaming Request" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response_stream = client.chat.completions.create(\n", " model=model_name,\n", " messages=messages,\n", " temperature=0.6,\n", " top_p=0.95,\n", " stream=True, # Non-streaming\n", " extra_body={\"separate_reasoning\": True},\n", ")\n", "\n", "reasoning_content = \"\"\n", "content = \"\"\n", "for chunk in response_stream:\n", " if chunk.choices[0].delta.content:\n", " content += chunk.choices[0].delta.content\n", " if chunk.choices[0].delta.reasoning_content:\n", " reasoning_content += chunk.choices[0].delta.reasoning_content\n", "\n", "print_highlight(\"==== Reasoning ====\")\n", "print_highlight(reasoning_content)\n", "\n", "print_highlight(\"==== Text ====\")\n", "print_highlight(content)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Optionally, you can buffer the reasoning content to the last reasoning chunk (or the first chunk after the reasoning content)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response_stream = client.chat.completions.create(\n", " model=model_name,\n", " messages=messages,\n", " temperature=0.6,\n", " top_p=0.95,\n", " stream=True, # Non-streaming\n", " extra_body={\"separate_reasoning\": True, \"stream_reasoning\": False},\n", ")\n", "\n", "reasoning_content = \"\"\n", "content = \"\"\n", "for chunk in response_stream:\n", " if chunk.choices[0].delta.content:\n", " content += chunk.choices[0].delta.content\n", " if chunk.choices[0].delta.reasoning_content:\n", " reasoning_content += chunk.choices[0].delta.reasoning_content\n", "\n", "print_highlight(\"==== Reasoning ====\")\n", "print_highlight(reasoning_content)\n", "\n", "print_highlight(\"==== Text ====\")\n", "print_highlight(content)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The reasoning separation is enable by default when specify . \n", "**To disable it, set the `separate_reasoning` option to `False` in request.**" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response_non_stream = client.chat.completions.create(\n", " model=model_name,\n", " messages=messages,\n", " temperature=0.6,\n", " top_p=0.95,\n", " stream=False, # Non-streaming\n", " extra_body={\"separate_reasoning\": False},\n", ")\n", "\n", "print_highlight(\"==== Original Output ====\")\n", "print_highlight(response_non_stream.choices[0].message.content)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### SGLang Native API " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from transformers import AutoTokenizer\n", "\n", "tokenizer = AutoTokenizer.from_pretrained(\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\")\n", "input = tokenizer.apply_chat_template(\n", " messages, tokenize=False, add_generation_prompt=True, return_dict=False\n", ")\n", "\n", "gen_url = f\"http://localhost:{port}/generate\"\n", "gen_data = {\n", " \"text\": input,\n", " \"sampling_params\": {\n", " \"skip_special_tokens\": False,\n", " \"max_new_tokens\": 1024,\n", " \"temperature\": 0.6,\n", " \"top_p\": 0.95,\n", " },\n", "}\n", "gen_response = requests.post(gen_url, json=gen_data).json()[\"text\"]\n", "\n", "print_highlight(\"==== Original Output ====\")\n", "print_highlight(gen_response)\n", "\n", "parse_url = f\"http://localhost:{port}/separate_reasoning\"\n", "separate_reasoning_data = {\n", " \"text\": gen_response,\n", " \"reasoning_parser\": \"deepseek-r1\",\n", "}\n", "separate_reasoning_response_json = requests.post(\n", " parse_url, json=separate_reasoning_data\n", ").json()\n", "print_highlight(\"==== Reasoning ====\")\n", "print_highlight(separate_reasoning_response_json[\"reasoning_text\"])\n", "print_highlight(\"==== Text ====\")\n", "print_highlight(separate_reasoning_response_json[\"text\"])" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Offline Engine API" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import sglang as sgl\n", "from sglang.srt.parser.reasoning_parser import ReasoningParser\n", "from sglang.utils import print_highlight\n", "\n", "llm = sgl.Engine(model_path=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\")\n", "tokenizer = AutoTokenizer.from_pretrained(\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\")\n", "input = tokenizer.apply_chat_template(\n", " messages, tokenize=False, add_generation_prompt=True, return_dict=False\n", ")\n", "sampling_params = {\n", " \"max_new_tokens\": 1024,\n", " \"skip_special_tokens\": False,\n", " \"temperature\": 0.6,\n", " \"top_p\": 0.95,\n", "}\n", "result = llm.generate(prompt=input, sampling_params=sampling_params)\n", "\n", "generated_text = result[\"text\"] # Assume there is only one prompt\n", "\n", "print_highlight(\"==== Original Output ====\")\n", "print_highlight(generated_text)\n", "\n", "parser = ReasoningParser(\"deepseek-r1\")\n", "reasoning_text, text = parser.parse_non_stream(generated_text)\n", "print_highlight(\"==== Reasoning ====\")\n", "print_highlight(reasoning_text)\n", "print_highlight(\"==== Text ====\")\n", "print_highlight(text)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "llm.shutdown()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Supporting New Reasoning Model Schemas\n", "\n", "For future reasoning models, you can implement the reasoning parser as a subclass of `BaseReasoningFormatDetector` in `python/sglang/srt/reasoning_parser.py` and specify the reasoning parser for new reasoning model schemas accordingly." ] } ], "metadata": { "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 4 } ================================================ FILE: docs/advanced_features/server_arguments.md ================================================ # Server Arguments This page provides a list of server arguments used in the command line to configure the behavior and performance of the language model server during deployment. These arguments enable users to customize key aspects of the server, including model selection, parallelism policies, memory management, and optimization techniques. You can find all arguments by `python3 -m sglang.launch_server --help` ## Common launch commands - To use a configuration file, create a YAML file with your server arguments and specify it with `--config`. CLI arguments will override config file values. ```bash # Create config.yaml cat > config.yaml << EOF model-path: meta-llama/Meta-Llama-3-8B-Instruct host: 0.0.0.0 port: 30000 tensor-parallel-size: 2 enable-metrics: true log-requests: true EOF # Launch server with config file python -m sglang.launch_server --config config.yaml ``` - To enable multi-GPU tensor parallelism, add `--tp 2`. If it reports the error "peer access is not supported between these two devices", add `--enable-p2p-check` to the server launch command. ```bash python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --tp 2 ``` - To enable multi-GPU data parallelism, add `--dp 2`. Data parallelism is better for throughput if there is enough memory. It can also be used together with tensor parallelism. The following command uses 4 GPUs in total. We recommend [SGLang Model Gateway (former Router)](../advanced_features/sgl_model_gateway.md) for data parallelism. ```bash python -m sglang_router.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --dp 2 --tp 2 ``` - If you see out-of-memory errors during serving, try to reduce the memory usage of the KV cache pool by setting a smaller value of `--mem-fraction-static`. The default value is `0.9`. ```bash python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --mem-fraction-static 0.7 ``` - See [hyperparameter tuning](hyperparameter_tuning.md) on tuning hyperparameters for better performance. - For docker and Kubernetes runs, you need to set up shared memory which is used for communication between processes. See `--shm-size` for docker and `/dev/shm` size update for Kubernetes manifests. - If you see out-of-memory errors during prefill for long prompts, try to set a smaller chunked prefill size. ```bash python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --chunked-prefill-size 4096 ``` - To enable fp8 weight quantization, add `--quantization fp8` on a fp16 checkpoint or directly load a fp8 checkpoint without specifying any arguments. - To enable fp8 kv cache quantization, add `--kv-cache-dtype fp8_e4m3` or `--kv-cache-dtype fp8_e5m2`. - To enable deterministic inference and batch invariant operations, add `--enable-deterministic-inference`. More details can be found in [deterministic inference document](../advanced_features/deterministic_inference.md). - If the model does not have a chat template in the Hugging Face tokenizer, you can specify a [custom chat template](../references/custom_chat_template.md). If the tokenizer has multiple named templates (e.g., 'default', 'tool_use'), you can select one using `--hf-chat-template-name tool_use`. - To run tensor parallelism on multiple nodes, add `--nnodes 2`. If you have two nodes with two GPUs on each node and want to run TP=4, let `sgl-dev-0` be the hostname of the first node and `50000` be an available port, you can use the following commands. If you meet deadlock, please try to add `--disable-cuda-graph` - (Note: This feature is out of maintenance and might cause error) To enable `torch.compile` acceleration, add `--enable-torch-compile`. It accelerates small models on small batch sizes. By default, the cache path is located at `/tmp/torchinductor_root`, you can customize it using environment variable `TORCHINDUCTOR_CACHE_DIR`. For more details, please refer to [PyTorch official documentation](https://pytorch.org/tutorials/recipes/torch_compile_caching_tutorial.html) and [Enabling cache for torch.compile](https://docs.sglang.io/references/torch_compile_cache.html). ```bash # Node 0 python -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3-8B-Instruct \ --tp 4 \ --dist-init-addr sgl-dev-0:50000 \ --nnodes 2 \ --node-rank 0 # Node 1 python -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3-8B-Instruct \ --tp 4 \ --dist-init-addr sgl-dev-0:50000 \ --nnodes 2 \ --node-rank 1 ``` Please consult the documentation below and [server_args.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py) to learn more about the arguments you may provide when launching a server. ## Model and tokenizer | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--model-path`
`--model` | The path of the model weights. This can be a local folder or a Hugging Face repo ID. | `None` | Type: str | | `--tokenizer-path` | The path of the tokenizer. | `None` | Type: str | | `--tokenizer-mode` | Tokenizer mode. 'auto' will use the fast tokenizer if available, and 'slow' will always use the slow tokenizer. | `auto` | `auto`, `slow` | | `--tokenizer-worker-num` | The worker num of the tokenizer manager. | `1` | Type: int | | `--skip-tokenizer-init` | If set, skip init tokenizer and pass input_ids in generate request. | `False` | bool flag (set to enable) | | `--load-format` | The format of the model weights to load. "auto" will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available. "pt" will load the weights in the pytorch bin format. "safetensors" will load the weights in the safetensors format. "npcache" will load the weights in pytorch format and store a numpy cache to speed up the loading. "dummy" will initialize the weights with random values, which is mainly for profiling."gguf" will load the weights in the gguf format. "bitsandbytes" will load the weights using bitsandbytes quantization."layered" loads weights layer by layer so that one can quantize a layer before loading another to make the peak memory envelope smaller. "flash_rl" will load the weights in flash_rl format. "fastsafetensors" and "private" are also supported. | `auto` | `auto`, `pt`, `safetensors`, `npcache`, `dummy`, `sharded_state`, `gguf`, `bitsandbytes`, `layered`, `flash_rl`, `remote`, `remote_instance`, `fastsafetensors`, `private` | | `--model-loader-extra-config` | Extra config for model loader. This will be passed to the model loader corresponding to the chosen load_format. | `{}` | Type: str | | `--trust-remote-code` | Whether or not to allow for custom models defined on the Hub in their own modeling files. | `False` | bool flag (set to enable) | | `--context-length` | The model's maximum context length. Defaults to None (will use the value from the model's config.json instead). | `None` | Type: int | | `--is-embedding` | Whether to use a CausalLM as an embedding model. | `False` | bool flag (set to enable) | | `--enable-multimodal` | Enable the multimodal functionality for the served model. If the model being served is not multimodal, nothing will happen | `None` | bool flag (set to enable) | | `--revision` | The specific model version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. | `None` | Type: str | | `--model-impl` | Which implementation of the model to use. * "auto" will try to use the SGLang implementation if it exists and fall back to the Transformers implementation if no SGLang implementation is available. * "sglang" will use the SGLang model implementation. * "transformers" will use the Transformers model implementation. | `auto` | Type: str | ## HTTP server | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--host` | The host of the HTTP server. | `127.0.0.1` | Type: str | | `--port` | The port of the HTTP server. | `30000` | Type: int | | `--fastapi-root-path` | App is behind a path based routing proxy. | `""` | Type: str | | `--grpc-mode` | If set, use gRPC server instead of HTTP server. | `False` | bool flag (set to enable) | | `--skip-server-warmup` | If set, skip warmup. | `False` | bool flag (set to enable) | | `--warmups` | Specify custom warmup functions (csv) to run before server starts eg. --warmups=warmup_name1,warmup_name2 will run the functions `warmup_name1` and `warmup_name2` specified in warmup.py before the server starts listening for requests | `None` | Type: str | | `--nccl-port` | The port for NCCL distributed environment setup. Defaults to a random port. | `None` | Type: int | | `--checkpoint-engine-wait-weights-before-ready` | If set, the server will wait for initial weights to be loaded via checkpoint-engine or other update methods before serving inference requests. | `False` | bool flag (set to enable) | ## Quantization and data type | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--dtype` | Data type for model weights and activations. * "auto" will use FP16 precision for FP32 and FP16 models, and BF16 precision for BF16 models. * "half" for FP16. Recommended for AWQ quantization. * "float16" is the same as "half". * "bfloat16" for a balance between precision and range. * "float" is shorthand for FP32 precision. * "float32" for FP32 precision. | `auto` | `auto`, `half`, `float16`, `bfloat16`, `float`, `float32` | | `--quantization` | The quantization method. | `None` | `awq`, `fp8`, `gptq`, `marlin`, `gptq_marlin`, `awq_marlin`, `bitsandbytes`, `gguf`, `modelopt`, `modelopt_fp8`, `modelopt_fp4`, `petit_nvfp4`, `w8a8_int8`, `w8a8_fp8`, `moe_wna16`, `qoq`, `w4afp8`, `mxfp4`, `mxfp8`, `auto-round`, `compressed-tensors`, `modelslim`, `quark_int4fp8_moe` | | `--quantization-param-path` | Path to the JSON file containing the KV cache scaling factors. This should generally be supplied, when KV cache dtype is FP8. Otherwise, KV cache scaling factors default to 1.0, which may cause accuracy issues. | `None` | Type: Optional[str] | | `--kv-cache-dtype` | Data type for kv cache storage. "auto" will use model data type. "bf16" or "bfloat16" for BF16 KV cache. "fp8_e5m2" and "fp8_e4m3" are supported for CUDA 11.8+. "fp4_e2m1" (only mxfp4) is supported for CUDA 12.8+ and PyTorch 2.8.0+ | `auto` | `auto`, `fp8_e5m2`, `fp8_e4m3`, `bf16`, `bfloat16`, `fp4_e2m1` | | `--enable-fp32-lm-head` | If set, the LM head outputs (logits) are in FP32. | `False` | bool flag (set to enable) | | `--modelopt-quant` | The ModelOpt quantization configuration. Supported values: 'fp8', 'int4_awq', 'w4a8_awq', 'nvfp4', 'nvfp4_awq'. This requires the NVIDIA Model Optimizer library to be installed: pip install nvidia-modelopt | `None` | Type: str | | `--modelopt-checkpoint-restore-path` | Path to restore a previously saved ModelOpt quantized checkpoint. If provided, the quantization process will be skipped and the model will be loaded from this checkpoint. | `None` | Type: str | | `--modelopt-checkpoint-save-path` | Path to save the ModelOpt quantized checkpoint after quantization. This allows reusing the quantized model in future runs. | `None` | Type: str | | `--modelopt-export-path` | Path to export the quantized model in HuggingFace format after ModelOpt quantization. The exported model can then be used directly with SGLang for inference. If not provided, the model will not be exported. | `None` | Type: str | | `--quantize-and-serve` | Quantize the model with ModelOpt and immediately serve it without exporting. This is useful for development and prototyping. For production, it's recommended to use separate quantization and deployment steps. | `False` | bool flag (set to enable) | | `--rl-quant-profile` | Path to the FlashRL quantization profile. Required when using --load-format flash_rl. | `None` | Type: str | ## Memory and scheduling | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--mem-fraction-static` | The fraction of the memory used for static allocation (model weights and KV cache memory pool). Use a smaller value if you see out-of-memory errors. | `None` | Type: float | | `--max-running-requests` | The maximum number of running requests. | `None` | Type: int | | `--max-queued-requests` | The maximum number of queued requests. This option is ignored when using disaggregation-mode. | `None` | Type: int | | `--max-total-tokens` | The maximum number of tokens in the memory pool. If not specified, it will be automatically calculated based on the memory usage fraction. This option is typically used for development and debugging purposes. | `None` | Type: int | | `--chunked-prefill-size` | The maximum number of tokens in a chunk for the chunked prefill. Setting this to -1 means disabling chunked prefill. | `None` | Type: int | | `--prefill-max-requests` | The maximum number of requests in a prefill batch. If not specified, there is no limit. | `None` | Type: int | | `--enable-dynamic-chunking` | Enable dynamic chunk size adjustment for pipeline parallelism. When enabled, chunk sizes are dynamically calculated based on fitted function to maintain consistent execution time across chunks. | `False` | bool flag (set to enable) | | `--max-prefill-tokens` | The maximum number of tokens in a prefill batch. The real bound will be the maximum of this value and the model's maximum context length. | `16384` | Type: int | | `--schedule-policy` | The scheduling policy of the requests. | `fcfs` | `lpm`, `random`, `fcfs`, `dfs-weight`, `lof`, `priority`, `routing-key` | | `--enable-priority-scheduling` | Enable priority scheduling. Requests with higher priority integer values will be scheduled first by default. | `False` | bool flag (set to enable) | | `--abort-on-priority-when-disabled` | If set, abort requests that specify a priority when priority scheduling is disabled. | `False` | bool flag (set to enable) | | `--schedule-low-priority-values-first` | If specified with --enable-priority-scheduling, the scheduler will schedule requests with lower priority integer values first. | `False` | bool flag (set to enable) | | `--priority-scheduling-preemption-threshold` | Minimum difference in priorities for an incoming request to have to preempt running request(s). | `10` | Type: int | | `--schedule-conservativeness` | How conservative the schedule policy is. A larger value means more conservative scheduling. Use a larger value if you see requests being retracted frequently. | `1.0` | Type: float | | `--page-size` | The number of tokens in a page. | `1` | Type: int | | `--swa-full-tokens-ratio` | The ratio of SWA layer KV tokens / full layer KV tokens, regardless of the number of swa:full layers. It should be between 0 and 1. E.g. 0.5 means if each swa layer has 50 tokens, then each full layer has 100 tokens. | `0.8` | Type: float | | `--disable-hybrid-swa-memory` | Disable the hybrid SWA memory. | `False` | bool flag (set to enable) | | `--radix-eviction-policy` | The eviction policy of radix trees. 'lru' stands for Least Recently Used, 'lfu' stands for Least Frequently Used. | `lru` | `lru`, `lfu` | | `--enable-prefill-delayer` | Enable prefill delayer for DP attention to reduce idle time. | `False` | bool flag (set to enable) | | `--prefill-delayer-max-delay-passes` | Maximum forward passes to delay prefill. | `30` | Type: int | | `--prefill-delayer-token-usage-low-watermark` | Token usage low watermark for prefill delayer. | `None` | Type: float | | `--prefill-delayer-forward-passes-buckets` | Custom buckets for prefill delayer forward passes histogram. 0 and max_delay_passes-1 will be auto-added. | `None` | List[float] | | `--prefill-delayer-wait-seconds-buckets` | Custom buckets for prefill delayer wait seconds histogram. 0 will be auto-added. | `None` | List[float] | ## Runtime options | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--device` | The device to use ('cuda', 'xpu', 'hpu', 'npu', 'cpu'). Defaults to auto-detection if not specified. | `None` | Type: str | | `--tensor-parallel-size`
`--tp-size` | The tensor parallelism size. | `1` | Type: int | | `--pipeline-parallel-size`
`--pp-size` | The pipeline parallelism size. | `1` | Type: int | | `--attention-context-parallel-size`
`--attn-cp-size`| The attention context parallelism size. | `1` | Type: int| | `--moe-data-parallel-size`
`--moe-dp-size`| The moe data parallelism size. | `1` | Type: int| | `--pp-max-micro-batch-size` | The maximum micro batch size in pipeline parallelism. | `None` | Type: int | | `--pp-async-batch-depth` | The async batch depth of pipeline parallelism. | `0` | Type: int | | `--stream-interval` | The interval (or buffer size) for streaming in terms of the token length. A smaller value makes streaming smoother, while a larger value makes the throughput higher | `1` | Type: int | | `--incremental-streaming-output` | Whether to output as a sequence of disjoint segments. | `False` | bool flag (set to enable) | | `--random-seed` | The random seed. | `None` | Type: int | | `--constrained-json-whitespace-pattern` | (outlines and llguidance backends only) Regex pattern for syntactic whitespaces allowed in JSON constrained output. For example, to allow the model to generate consecutive whitespaces, set the pattern to [\n\t ]* | `None` | Type: str | | `--constrained-json-disable-any-whitespace` | (xgrammar and llguidance backends only) Enforce compact representation in JSON constrained output. | `False` | bool flag (set to enable) | | `--watchdog-timeout` | Set watchdog timeout in seconds. If a forward batch takes longer than this, the server will crash to prevent hanging. | `300` | Type: float | | `--soft-watchdog-timeout` | Set soft watchdog timeout in seconds. If a forward batch takes longer than this, the server will dump information for debugging. | `None` | Type: float | | `--dist-timeout` | Set timeout for torch.distributed initialization. | `None` | Type: int | | `--download-dir` | Model download directory for huggingface. | `None` | Type: str | | `--model-checksum` | Model file integrity verification. If provided without value, uses model-path as HF repo ID. Otherwise, provide checksums JSON file path or HuggingFace repo ID. | `None` | Type: str | | `--base-gpu-id` | The base GPU ID to start allocating GPUs from. Useful when running multiple instances on the same machine. | `0` | Type: int | | `--gpu-id-step` | The delta between consecutive GPU IDs that are used. For example, setting it to 2 will use GPU 0,2,4,... | `1` | Type: int | | `--sleep-on-idle` | Reduce CPU usage when sglang is idle. | `False` | bool flag (set to enable) | | `--custom-sigquit-handler` | Register a custom sigquit handler so you can do additional cleanup after the server is shutdown. This is only available for Engine, not for CLI. | `None` | Type: str | ## Logging | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--log-level` | The logging level of all loggers. | `info` | Type: str | | `--log-level-http` | The logging level of HTTP server. If not set, reuse --log-level by default. | `None` | Type: str | | `--log-requests` | Log metadata, inputs, outputs of all requests. The verbosity is decided by --log-requests-level | `False` | bool flag (set to enable) | | `--log-requests-level` | 0: Log metadata (no sampling parameters). 1: Log metadata and sampling parameters. 2: Log metadata, sampling parameters and partial input/output. 3: Log every input/output. | `2` | `0`, `1`, `2`, `3` | | `--log-requests-format` | Format for request logging: 'text' (human-readable) or 'json' (structured) | `text` | `text`, `json` | | `--log-requests-target` | Target(s) for request logging: 'stdout' and/or directory path(s) for file output. Can specify multiple targets, e.g., '--log-requests-target stdout /my/path'. | `None` | List[str] | | `--uvicorn-access-log-exclude-prefixes` | Exclude uvicorn access logs whose request path starts with any of these prefixes. Defaults to empty (disabled). | `[]` | List[str] | | `--crash-dump-folder` | Folder path to dump requests from the last 5 min before a crash (if any). If not specified, crash dumping is disabled. | `None` | Type: str | | `--show-time-cost` | Show time cost of custom marks. | `False` | bool flag (set to enable) | | `--enable-metrics` | Enable log prometheus metrics. | `False` | bool flag (set to enable) | | `--enable-metrics-for-all-schedulers` | Enable --enable-metrics-for-all-schedulers when you want schedulers on all TP ranks (not just TP 0) to record request metrics separately. This is especially useful when dp_attention is enabled, as otherwise all metrics appear to come from TP 0. | `False` | bool flag (set to enable) | | `--tokenizer-metrics-custom-labels-header` | Specify the HTTP header for passing custom labels for tokenizer metrics. | `x-custom-labels` | Type: str | | `--tokenizer-metrics-allowed-custom-labels` | The custom labels allowed for tokenizer metrics. The labels are specified via a dict in '--tokenizer-metrics-custom-labels-header' field in HTTP requests, e.g., {'label1': 'value1', 'label2': 'value2'} is allowed if '--tokenizer-metrics-allowed-custom-labels label1 label2' is set. | `None` | List[str] | | `--bucket-time-to-first-token` | The buckets of time to first token, specified as a list of floats. | `None` | List[float] | | `--bucket-inter-token-latency` | The buckets of inter-token latency, specified as a list of floats. | `None` | List[float] | | `--bucket-e2e-request-latency` | The buckets of end-to-end request latency, specified as a list of floats. | `None` | List[float] | | `--collect-tokens-histogram` | Collect prompt/generation tokens histogram. | `False` | bool flag (set to enable) | | `--prompt-tokens-buckets` | The buckets rule of prompt tokens. Supports 3 rule types: 'default' uses predefined buckets; 'tse ' generates two sides exponential distributed buckets (e.g., 'tse 1000 2 8' generates buckets [984.0, 992.0, 996.0, 998.0, 1000.0, 1002.0, 1004.0, 1008.0, 1016.0]).); 'custom ...' uses custom bucket values (e.g., 'custom 10 50 100 500'). | `None` | List[str] | | `--generation-tokens-buckets` | The buckets rule for generation tokens histogram. Supports 3 rule types: 'default' uses predefined buckets; 'tse ' generates two sides exponential distributed buckets (e.g., 'tse 1000 2 8' generates buckets [984.0, 992.0, 996.0, 998.0, 1000.0, 1002.0, 1004.0, 1008.0, 1016.0]).); 'custom ...' uses custom bucket values (e.g., 'custom 10 50 100 500'). | `None` | List[str] | | `--gc-warning-threshold-secs` | The threshold for long GC warning. If a GC takes longer than this, a warning will be logged. Set to 0 to disable. | `0.0` | Type: float | | `--decode-log-interval` | The log interval of decode batch. | `40` | Type: int | | `--enable-request-time-stats-logging` | Enable per request time stats logging | `False` | bool flag (set to enable) | | `--kv-events-config` | Config in json format for NVIDIA dynamo KV event publishing. Publishing will be enabled if this flag is used. | `None` | Type: str | | `--enable-trace` | Enable opentelemetry trace | `False` | bool flag (set to enable) | | `--otlp-traces-endpoint` | Config opentelemetry collector endpoint if --enable-trace is set. format: : | `localhost:4317` | Type: str | ## RequestMetricsExporter configuration | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--export-metrics-to-file` | Export performance metrics for each request to local file (e.g. for forwarding to external systems). | `False` | bool flag (set to enable) | | `--export-metrics-to-file-dir` | Directory path for writing performance metrics files (required when --export-metrics-to-file is enabled). | `None` | Type: str | ## API related | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--api-key` | Set API key of the server. It is also used in the OpenAI API compatible server. | `None` | Type: str | | `--admin-api-key` | Set **admin API key** for administrative/control endpoints (e.g., weights update, cache flush, `/get_server_info`). Endpoints marked as admin-only require `Authorization: Bearer ` when this is set. | `None` | Type: str | | `--served-model-name` | Override the model name returned by the v1/models endpoint in OpenAI API server. | `None` | Type: str | | `--weight-version` | Version identifier for the model weights. Defaults to 'default' if not specified. | `default` | Type: str | | `--chat-template` | The builtin chat template name or the path of the chat template file. This is only used for OpenAI-compatible API server. | `None` | Type: str | | `--hf-chat-template-name` | When the HuggingFace tokenizer has multiple chat templates (e.g., 'default', 'tool_use', 'rag'), specify which named template to use. If not set, the first available template is used. | `None` | Type: str | | `--completion-template` | The builtin completion template name or the path of the completion template file. This is only used for OpenAI-compatible API server. only for code completion currently. | `None` | Type: str | | `--file-storage-path` | The path of the file storage in backend. | `sglang_storage` | Type: str | | `--enable-cache-report` | Return number of cached tokens in usage.prompt_tokens_details for each openai request. | `False` | bool flag (set to enable) | | `--reasoning-parser` | Specify the parser for reasoning models. Supported parsers: [deepseek-r1, deepseek-v3, glm45, gpt-oss, kimi, qwen3, qwen3-thinking, step3]. | `None` | `deepseek-r1`, `deepseek-v3`, `glm45`, `gpt-oss`, `kimi`, `qwen3`, `qwen3-thinking`, `step3` | | `--tool-call-parser` | Specify the parser for handling tool-call interactions. Supported parsers: [deepseekv3, deepseekv31, glm, glm45, glm47, gpt-oss, kimi_k2, llama3, mistral, pythonic, qwen, qwen25, qwen3_coder, step3]. | `None` | `deepseekv3`, `deepseekv31`, `glm`, `glm45`, `glm47`, `gpt-oss`, `kimi_k2`, `llama3`, `mistral`, `pythonic`, `qwen`, `qwen25`, `qwen3_coder`, `step3`, `gigachat3` | | `--tool-server` | Either 'demo' or a comma-separated list of tool server urls to use for the model. If not specified, no tool server will be used. | `None` | Type: str | | `--sampling-defaults` | Where to get default sampling parameters. 'openai' uses SGLang/OpenAI defaults (temperature=1.0, top_p=1.0, etc.). 'model' uses the model's generation_config.json to get the recommended sampling parameters if available. Default is 'model'. | `model` | `openai`, `model` | ## Data parallelism | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--data-parallel-size`
`--dp-size` | The data parallelism size. | `1` | Type: int | | `--load-balance-method` | The load balancing strategy for data parallelism. The `total_tokens` algorithm can only be used when DP attention is applied. This algorithm performs load balancing based on the real-time token load of the DP workers. | `auto` | `auto`, `round_robin`, `follow_bootstrap_room`, `total_requests`, `total_tokens` | ## Multi-node distributed serving | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--dist-init-addr`
`--nccl-init-addr` | The host address for initializing distributed backend (e.g., `192.168.0.2:25000`). | `None` | Type: str | | `--nnodes` | The number of nodes. | `1` | Type: int | | `--node-rank` | The node rank. | `0` | Type: int | ## Model override args | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--json-model-override-args` | A dictionary in JSON string format used to override default model configurations. | `{}` | Type: str | | `--preferred-sampling-params` | json-formatted sampling settings that will be returned in /get_model_info | `None` | Type: str | ## LoRA | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--enable-lora` | Enable LoRA support for the model. This argument is automatically set to `True` if `--lora-paths` is provided for backward compatibility. | `False` | Bool flag (set to enable) | | `--enable-lora-overlap-loading` | Enable asynchronous LoRA weight loading in order to overlap H2D transfers with GPU compute. This should be enabled if you find that your LoRA workloads are bottlenecked by adapter weight loading, for example when frequently loading large LoRA adapters. | `False` | Bool flag (set to enable) | `--max-lora-rank` | The maximum LoRA rank that should be supported. If not specified, it will be automatically inferred from the adapters provided in `--lora-paths`. This argument is needed when you expect to dynamically load adapters of larger LoRA rank after server startup. | `None` | Type: int | | `--lora-target-modules` | The union set of all target modules where LoRA should be applied (e.g., `q_proj`, `k_proj`, `gate_proj`). If not specified, it will be automatically inferred from the adapters provided in `--lora-paths`. You can also set it to `all` to enable LoRA for all supported modules; note this may introduce minor performance overhead. | `None` | `q_proj`, `k_proj`, `v_proj`, `o_proj`, `gate_proj`, `up_proj`, `down_proj`, `qkv_proj`, `gate_up_proj`, `all` | | `--lora-paths` | The list of LoRA adapters to load. Each adapter must be specified in one of the following formats: `` \| `=` \| JSON with schema `{"lora_name": str, "lora_path": str, "pinned": bool}`. | `None` | Type: List[str] / JSON objects | | `--max-loras-per-batch` | Maximum number of adapters for a running batch, including base-only requests. | `8` | Type: int | | `--max-loaded-loras` | If specified, limits the maximum number of LoRA adapters loaded in CPU memory at a time. Must be ≥ `--max-loras-per-batch`. | `None` | Type: int | | `--lora-eviction-policy` | LoRA adapter eviction policy when the GPU memory pool is full. | `lru` | `lru`, `fifo` | | `--lora-backend` | Choose the kernel backend for multi-LoRA serving. | `csgmv` | `triton`, `csgmv`, `ascend`, `torch_native` | | `--max-lora-chunk-size` | Maximum chunk size for the ChunkedSGMV LoRA backend. Only used when `--lora-backend` is `csgmv`. Larger values may improve performance. | `16` | `16`, `32`, `64`, `128` | ## Kernel Backends (Attention, Sampling, Grammar, GEMM) | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--attention-backend` | Choose the kernels for attention layers. | `None` | `triton`, `torch_native`, `flex_attention`, `nsa`, `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` | | `--prefill-attention-backend` | Choose the kernels for prefill attention layers (have priority over --attention-backend). | `None` | `triton`, `torch_native`, `flex_attention`, `nsa`, `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` | | `--decode-attention-backend` | Choose the kernels for decode attention layers (have priority over --attention-backend). | `None` | `triton`, `torch_native`, `flex_attention`, `nsa`, `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` | | `--sampling-backend` | Choose the kernels for sampling layers. | `None` | `flashinfer`, `pytorch`, `ascend` | | `--grammar-backend` | Choose the backend for grammar-guided decoding. | `None` | `xgrammar`, `outlines`, `llguidance`, `none` | | `--mm-attention-backend` | Set multimodal attention backend. | `None` | `sdpa`, `fa3`, `fa4`, `triton_attn`, `ascend_attn`, `aiter_attn` | | `--nsa-prefill-backend` | Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention). | `flashmla_sparse` | `flashmla_sparse`, `flashmla_kv`, `flashmla_auto`, `fa3`, `tilelang`, `aiter`, `trtllm` | | `--nsa-decode-backend` | Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding. | `fa3` | `flashmla_sparse`, `flashmla_kv`, `fa3`, `tilelang`, `aiter`, `trtllm` | | `--fp8-gemm-backend` | Choose the runner backend for Blockwise FP8 GEMM operations. Options: 'auto' (default, auto-selects based on hardware), 'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), 'flashinfer_trtllm' (FlashInfer TRTLLM backend; SM100/SM103 only), 'flashinfer_cutlass' (FlashInfer CUTLASS backend, SM120 only), 'flashinfer_deepgemm' (Hopper SM90 only, uses swapAB optimization for small M dimensions in decoding), 'cutlass' (optimal for Hopper/Blackwell GPUs and high-throughput), 'triton' (fallback, widely compatible), 'aiter' (ROCm only). **NOTE**: This replaces the deprecated environment variables SGLANG_ENABLE_FLASHINFER_FP8_GEMM and SGLANG_SUPPORT_CUTLASS_BLOCK_FP8. | `auto` | `auto`, `deep_gemm`, `flashinfer_trtllm`, `flashinfer_cutlass`, `flashinfer_deepgemm`, `cutlass`, `triton`, `aiter` | | `--fp4-gemm-backend` | Choose the runner backend for NVFP4 GEMM operations. Options: 'flashinfer_cutlass' (default), 'auto' (auto-selects between flashinfer_cudnn/flashinfer_cutlass based on CUDA/cuDNN version), 'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), 'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling). All backends are from FlashInfer; when FlashInfer is unavailable, sgl-kernel CUTLASS is used as an automatic fallback. **NOTE**: This replaces the deprecated environment variable SGLANG_FLASHINFER_FP4_GEMM_BACKEND. | `flashinfer_cutlass` | `auto`, `flashinfer_cudnn`, `flashinfer_cutlass`, `flashinfer_trtllm` | | `--disable-flashinfer-autotune` | Flashinfer autotune is enabled by default. Set this flag to disable the autotune. | `False` | bool flag (set to enable) | ## Speculative decoding | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--speculative-algorithm` | Speculative algorithm. | `None` | `EAGLE`, `EAGLE3`, `NEXTN`, `STANDALONE`, `NGRAM` | | `--speculative-draft-model-path`
`--speculative-draft-model` | The path of the draft model weights. This can be a local folder or a Hugging Face repo ID. | `None` | Type: str | | `--speculative-draft-model-revision` | The specific draft model version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. | `None` | Type: str | | `--speculative-draft-load-format` | The format of the draft model weights to load. If not specified, will use the same format as --load-format. Use 'dummy' to initialize draft model weights with random values for profiling. | `None` | Same as --load-format options | | `--speculative-num-steps` | The number of steps sampled from draft model in Speculative Decoding. | `None` | Type: int | | `--speculative-eagle-topk` | The number of tokens sampled from the draft model in eagle2 each step. | `None` | Type: int | | `--speculative-num-draft-tokens` | The number of tokens sampled from the draft model in Speculative Decoding. | `None` | Type: int | | `--speculative-accept-threshold-single` | Accept a draft token if its probability in the target model is greater than this threshold. | `1.0` | Type: float | | `--speculative-accept-threshold-acc` | The accept probability of a draft token is raised from its target probability p to min(1, p / threshold_acc). | `1.0` | Type: float | | `--speculative-token-map` | The path of the draft model's small vocab table. | `None` | Type: str | | `--speculative-attention-mode` | Attention backend for speculative decoding operations (both target verify and draft extend). Can be one of 'prefill' (default) or 'decode'. | `prefill` | `prefill`, `decode` | | `--speculative-draft-attention-backend` | Attention backend for speculative decoding drafting. | `None` | Same as attention backend options | | `--speculative-moe-runner-backend` | MOE backend for EAGLE speculative decoding, see --moe-runner-backend for options. Same as moe runner backend if unset. | `None` | Same as --moe-runner-backend options | | `--speculative-moe-a2a-backend` | MOE A2A backend for EAGLE speculative decoding, see --moe-a2a-backend for options. Same as moe a2a backend if unset. | `None` | Same as --moe-a2a-backend options | | `--speculative-draft-model-quantization` | The quantization method for speculative model. | `None` | Same as --quantization options | ## Ngram speculative decoding | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--speculative-ngram-min-match-window-size` | The minimum window size for pattern matching in ngram speculative decoding. | `1` | Type: int | | `--speculative-ngram-max-match-window-size` | The maximum window size for pattern matching in ngram speculative decoding. | `12` | Type: int | | `--speculative-ngram-min-bfs-breadth` | The minimum breadth for BFS (Breadth-First Search) in ngram speculative decoding. | `1` | Type: int | | `--speculative-ngram-max-bfs-breadth` | The maximum breadth for BFS (Breadth-First Search) in ngram speculative decoding. | `10` | Type: int | | `--speculative-ngram-match-type` | The match type for cache tree. | `BFS` | `BFS`, `PROB` | | `--speculative-ngram-branch-length` | The branch length for ngram speculative decoding. | `18` | Type: int | | `--speculative-ngram-capacity` | The cache capacity for ngram speculative decoding. | `10000000` | Type: int | ## Multi-layer Eagle speculative decoding | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--enable-multi-layer-eagle` | Enable multi-layer Eagle speculative decoding. | `False` | bool flag (set to enable) | ## MoE | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--expert-parallel-size`
`--ep-size`
`--ep` | The expert parallelism size. | `1` | Type: int | | `--moe-a2a-backend` | Select the backend for all-to-all communication for expert parallelism. | `none` | `none`, `deepep`, `mooncake`, `mori`, `nixl`, `ascend_fuseep`| | `--moe-runner-backend` | Choose the runner backend for MoE. | `auto` | `auto`, `deep_gemm`, `triton`, `triton_kernel`, `flashinfer_trtllm`, `flashinfer_trtllm_routed`, `flashinfer_cutlass`, `flashinfer_mxfp4`, `flashinfer_cutedsl`, `cutlass` | | `--flashinfer-mxfp4-moe-precision` | Choose the computation precision of flashinfer mxfp4 moe | `default` | `default`, `bf16` | | `--enable-flashinfer-allreduce-fusion` | Enable FlashInfer allreduce fusion with Residual RMSNorm. | `False` | bool flag (set to enable) | | `--enable-aiter-allreduce-fusion` | Enable aiter allreduce fusion with Residual RMSNorm. | `False` | bool flag (set to enable) | | `--deepep-mode` | Select the mode when enable DeepEP MoE, could be `normal`, `low_latency` or `auto`. Default is `auto`, which means `low_latency` for decode batch and `normal` for prefill batch. | `auto` | `normal`, `low_latency`, `auto` | | `--ep-num-redundant-experts` | Allocate this number of redundant experts in expert parallel. | `0` | Type: int | | `--ep-dispatch-algorithm` | The algorithm to choose ranks for redundant experts in expert parallel. | `None` | Type: str | | `--init-expert-location` | Initial location of EP experts. | `trivial` | Type: str | | `--enable-eplb` | Enable EPLB algorithm | `False` | bool flag (set to enable) | | `--eplb-algorithm` | Chosen EPLB algorithm | `auto` | Type: str | | `--eplb-rebalance-num-iterations` | Number of iterations to automatically trigger a EPLB re-balance. | `1000` | Type: int | | `--eplb-rebalance-layers-per-chunk` | Number of layers to rebalance per forward pass. | `None` | Type: int | | `--eplb-min-rebalancing-utilization-threshold` | Minimum threshold for GPU average utilization to trigger EPLB rebalancing. Must be in the range [0.0, 1.0]. | `1.0` | Type: float | | `--expert-distribution-recorder-mode` | Mode of expert distribution recorder. | `None` | Type: str | | `--expert-distribution-recorder-buffer-size` | Circular buffer size of expert distribution recorder. Set to -1 to denote infinite buffer. | `None` | Type: int | | `--enable-expert-distribution-metrics` | Enable logging metrics for expert balancedness | `False` | bool flag (set to enable) | | `--deepep-config` | Tuned DeepEP config suitable for your own cluster. It can be either a string with JSON content or a file path. | `None` | Type: str | | `--moe-dense-tp-size` | TP size for MoE dense MLP layers. This flag is useful when, with large TP size, there are errors caused by weights in MLP layers having dimension smaller than the min dimension GEMM supports. | `None` | Type: int | | `--elastic-ep-backend` | Specify the collective communication backend for elastic EP. Currently supports 'mooncake'. | `none` | `none`, `mooncake` | | `--enable-elastic-expert-backup` | Enable elastic EP backend to backup expert weights in DRAM feature. Currently supports 'mooncake'.| `False` | bool flag (set to enable) | | `--mooncake-ib-device` | The InfiniBand devices for Mooncake Backend transfer, accepts multiple comma-separated devices (e.g., --mooncake-ib-device mlx5_0,mlx5_1). Default is None, which triggers automatic device detection when Mooncake Backend is enabled. | `None` | Type: str | ## Mamba Cache | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--max-mamba-cache-size` | The maximum size of the mamba cache. | `None` | Type: int | | `--mamba-ssm-dtype` | The data type of the SSM states in mamba cache. | `float32` | `float32`, `bfloat16`, `float16` | | `--mamba-full-memory-ratio` | The ratio of mamba state memory to full kv cache memory. | `0.9` | Type: float | | `--mamba-scheduler-strategy` | The strategy to use for mamba scheduler. `auto` currently defaults to `no_buffer`. 1. `no_buffer` does not support overlap scheduler due to not allocating extra mamba state buffers. Branching point caching support is feasible but not implemented. 2. `extra_buffer` supports overlap schedule by allocating extra mamba state buffers to track mamba state for caching (mamba state usage per running req becomes `2x` for non-spec; `1+(1/(2+speculative_num_draft_tokens))x` for spec dec (e.g. 1.16x if speculative_num_draft_tokens==4)). 2a. `extra_buffer` is strictly better for non-KV-cache-bound cases; for KV-cache-bound cases, the tradeoff depends on whether enabling overlap outweighs reduced max running requests. 2b. mamba caching at radix cache branching point is strictly better than non-branch but requires kernel support (currently only FLA backend), currently only extra_buffer supports branching. | `auto` | `auto`, `no_buffer`, `extra_buffer` | | `--mamba-track-interval` | The interval (in tokens) to track the mamba state during decode. Only used when `--mamba-scheduler-strategy` is `extra_buffer`. Must be divisible by page_size if set, and must be >= speculative_num_draft_tokens when using speculative decoding. | `256` | Type: int | ## Hierarchical cache | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--enable-hierarchical-cache` | Enable hierarchical cache | `False` | bool flag (set to enable) | | `--hicache-ratio` | The ratio of the size of host KV cache memory pool to the size of device pool. | `2.0` | Type: float | | `--hicache-size` | The size of host KV cache memory pool in gigabytes, which will override the hicache_ratio if set. | `0` | Type: int | | `--hicache-write-policy` | The write policy of hierarchical cache. | `write_through` | `write_back`, `write_through`, `write_through_selective` | | `--hicache-io-backend` | The IO backend for KV cache transfer between CPU and GPU | `kernel` | `direct`, `kernel`, `kernel_ascend` | | `--hicache-mem-layout` | The layout of host memory pool for hierarchical cache. | `layer_first` | `layer_first`, `page_first`, `page_first_direct`, `page_first_kv_split`, `page_head` | | `--hicache-storage-backend` | The storage backend for hierarchical KV cache. Built-in backends: file, mooncake, hf3fs, nixl, aibrix. For dynamic backend, use --hicache-storage-backend-extra-config to specify: backend_name (custom name), module_path (Python module path), class_name (backend class name). | `None` | `file`, `mooncake`, `hf3fs`, `nixl`, `aibrix`, `dynamic`, `eic` | | `--hicache-storage-prefetch-policy` | Control when prefetching from the storage backend should stop. | `best_effort` | `best_effort`, `wait_complete`, `timeout` | | `--hicache-storage-backend-extra-config` | A dictionary in JSON string format, or a string starting with a `@` followed by a config file in JSON/YAML/TOML format, containing extra configuration for the storage backend. | `None` | Type: str | ## Hierarchical sparse attention | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--hierarchical-sparse-attention-extra-config` | A dictionary in JSON string format for hierarchical sparse attention configuration. Required fields: `algorithm` (str), `backend` (str). All other fields are algorithm-specific and passed to the algorithm constructor. | `None` | Type: str | ## LMCache | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--enable-lmcache` | Using LMCache as an alternative hierarchical cache solution | `False` | bool flag (set to enable) | ## Ktransformers | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--kt-weight-path` | [ktransformers parameter] The path of the quantized expert weights for amx kernel. A local folder. | `None` | Type: str | | `--kt-method` | [ktransformers parameter] Quantization formats for CPU execution. | `AMXINT4` | Type: str | | `--kt-cpuinfer` | [ktransformers parameter] The number of CPUInfer threads. | `None` | Type: int | | `--kt-threadpool-count` | [ktransformers parameter] One-to-one with the number of NUMA nodes (one thread pool per NUMA). | `2` | Type: int | | `--kt-num-gpu-experts` | [ktransformers parameter] The number of GPU experts. | `None` | Type: int | | `--kt-max-deferred-experts-per-token` | [ktransformers parameter] Maximum number of experts deferred to CPU per token. All MoE layers except the final one use this value; the final layer always uses 0. | `None` | Type: int | ## Diffusion LLM | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--dllm-algorithm` | The diffusion LLM algorithm, such as LowConfidence. | `None` | Type: str | | `--dllm-algorithm-config` | The diffusion LLM algorithm configurations. Must be a YAML file. | `None` | Type: str | ## Double Sparsity | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--enable-double-sparsity` | Enable double sparsity attention | `False` | bool flag (set to enable) | | `--ds-channel-config-path` | The path of the double sparsity channel config | `None` | Type: str | | `--ds-heavy-channel-num` | The number of heavy channels in double sparsity attention | `32` | Type: int | | `--ds-heavy-token-num` | The number of heavy tokens in double sparsity attention | `256` | Type: int | | `--ds-heavy-channel-type` | The type of heavy channels in double sparsity attention | `qk` | Type: str | | `--ds-sparse-decode-threshold` | The minimum decode sequence length required before the double-sparsity backend switches from the dense fallback to the sparse decode kernel. | `4096` | Type: int | ## Offloading | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--cpu-offload-gb` | How many GBs of RAM to reserve for CPU offloading. | `0` | Type: int | | `--offload-group-size` | Number of layers per group in offloading. | `-1` | Type: int | | `--offload-num-in-group` | Number of layers to be offloaded within a group. | `1` | Type: int | | `--offload-prefetch-step` | Steps to prefetch in offloading. | `1` | Type: int | | `--offload-mode` | Mode of offloading. | `cpu` | Type: str | ## Args for multi-item scoring | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--multi-item-scoring-delimiter` | Delimiter token ID for multi-item scoring. Used to combine Query and Items into a single sequence: QueryItem1Item2... This enables efficient batch processing of multiple items against a single query. | `None` | Type: int | ## Optimization/debug options | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--disable-radix-cache` | Disable RadixAttention for prefix caching. | `False` | bool flag (set to enable) | | `--cuda-graph-max-bs` | Set the maximum batch size for cuda graph. It will extend the cuda graph capture batch size to this value. | `None` | Type: int | | `--cuda-graph-bs` | Set the list of batch sizes for cuda graph. | `None` | List[int] | | `--disable-cuda-graph` | Disable cuda graph. | `False` | bool flag (set to enable) | | `--disable-cuda-graph-padding` | Disable cuda graph when padding is needed. Still uses cuda graph when padding is not needed. | `False` | bool flag (set to enable) | | `--enable-profile-cuda-graph` | Enable profiling of cuda graph capture. | `False` | bool flag (set to enable) | | `--enable-cudagraph-gc` | Enable garbage collection during CUDA graph capture. If disabled (default), GC is frozen during capture to speed up the process. | `False` | bool flag (set to enable) | | `--enable-layerwise-nvtx-marker` | Enable layerwise NVTX profiling annotations for the model. This adds NVTX markers to every layer for detailed per-layer performance analysis with Nsight Systems. | `False` | bool flag (set to enable) | | `--enable-nccl-nvls` | Enable NCCL NVLS for prefill heavy requests when available. | `False` | bool flag (set to enable) | | `--enable-symm-mem` | Enable NCCL symmetric memory for fast collectives. | `False` | bool flag (set to enable) | | `--disable-flashinfer-cutlass-moe-fp4-allgather` | Disables quantize before all-gather for flashinfer cutlass moe. | `False` | bool flag (set to enable) | | `--enable-tokenizer-batch-encode` | Enable batch tokenization for improved performance when processing multiple text inputs. Do not use with image inputs, pre-tokenized input_ids, or input_embeds. | `False` | bool flag (set to enable) | | `--disable-tokenizer-batch-decode` | Disable batch decoding when decoding multiple completions. | `False` | bool flag (set to enable) | | `--disable-outlines-disk-cache` | Disable disk cache of outlines to avoid possible crashes related to file system or high concurrency. | `False` | bool flag (set to enable) | | `--disable-custom-all-reduce` | Disable the custom all-reduce kernel and fall back to NCCL. | `False` | bool flag (set to enable) | | `--enable-mscclpp` | Enable using mscclpp for small messages for all-reduce kernel and fall back to NCCL. | `False` | bool flag (set to enable) | | `--enable-torch-symm-mem` | Enable using torch symm mem for all-reduce kernel and fall back to NCCL. Only supports CUDA device SM90 and above. SM90 supports world size 4, 6, 8. SM10 supports world size 6, 8. | `False` | bool flag (set to enable) | | `--disable-overlap-schedule` | Disable the overlap scheduler, which overlaps the CPU scheduler with GPU model worker. | `False` | bool flag (set to enable) | | `--enable-mixed-chunk` | Enabling mixing prefill and decode in a batch when using chunked prefill. | `False` | bool flag (set to enable) | | `--enable-dp-attention` | Enabling data parallelism for attention and tensor parallelism for FFN. The dp size should be equal to the tp size. Currently DeepSeek-V2 and Qwen 2/3 MoE models are supported. | `False` | bool flag (set to enable) | | `--enable-dp-lm-head` | Enable vocabulary parallel across the attention TP group to avoid all-gather across DP groups, optimizing performance under DP attention. | `False` | bool flag (set to enable) | | `--enable-two-batch-overlap` | Enabling two micro batches to overlap. | `False` | bool flag (set to enable) | | `--enable-single-batch-overlap` | Let computation and communication overlap within one micro batch. | `False` | bool flag (set to enable) | | `--tbo-token-distribution-threshold` | The threshold of token distribution between two batches in micro-batch-overlap, determines whether to two-batch-overlap or two-chunk-overlap. Set to 0 denote disable two-chunk-overlap. | `0.48` | Type: float | | `--enable-torch-compile` | Optimize the model with torch.compile. Experimental feature. | `False` | bool flag (set to enable) | | `--enable-torch-compile-debug-mode` | Enable debug mode for torch compile. | `False` | bool flag (set to enable) | | `--disable-piecewise-cuda-graph` | Disable piecewise cuda graph for extend/prefill. PCG is enabled by default. | `False` | bool flag (set to disable) | | `--enforce-piecewise-cuda-graph` | Enforce piecewise cuda graph, skipping all auto-disable conditions. For testing only. | `False` | bool flag (set to enable) | | `--piecewise-cuda-graph-tokens` | Set the list of tokens when using piecewise cuda graph. | `None` | Type: JSON list | | `--piecewise-cuda-graph-compiler` | Set the compiler for piecewise cuda graph. Choices are: eager, inductor. | `eager` | `eager`, `inductor` | | `--torch-compile-max-bs` | Set the maximum batch size when using torch compile. | `32` | Type: int | | `--piecewise-cuda-graph-max-tokens` | Set the maximum tokens when using piecewise cuda graph. | `4096` | Type: int | | `--torchao-config` | Optimize the model with torchao. Experimental feature. Current choices are: int8dq, int8wo, int4wo-, fp8wo, fp8dq-per_tensor, fp8dq-per_row | `` | Type: str | | `--enable-nan-detection` | Enable the NaN detection for debugging purposes. | `False` | bool flag (set to enable) | | `--enable-p2p-check` | Enable P2P check for GPU access, otherwise the p2p access is allowed by default. | `False` | bool flag (set to enable) | | `--triton-attention-reduce-in-fp32` | Cast the intermediate attention results to fp32 to avoid possible crashes related to fp16. This only affects Triton attention kernels. | `False` | bool flag (set to enable) | | `--triton-attention-num-kv-splits` | The number of KV splits in flash decoding Triton kernel. Larger value is better in longer context scenarios. The default value is 8. | `8` | Type: int | | `--triton-attention-split-tile-size` | The size of split KV tile in flash decoding Triton kernel. Used for deterministic inference. | `None` | Type: int | | `--num-continuous-decode-steps` | Run multiple continuous decoding steps to reduce scheduling overhead. This can potentially increase throughput but may also increase time-to-first-token latency. The default value is 1, meaning only run one decoding step at a time. | `1` | Type: int | | `--delete-ckpt-after-loading` | Delete the model checkpoint after loading the model. | `False` | bool flag (set to enable) | | `--enable-memory-saver` | Allow saving memory using release_memory_occupation and resume_memory_occupation | `False` | bool flag (set to enable) | | `--enable-weights-cpu-backup` | Save model weights to CPU memory during release_weights_occupation and resume_weights_occupation | `False` | bool flag (set to enable) | | `--enable-draft-weights-cpu-backup` | Save draft model weights to CPU memory during release_weights_occupation and resume_weights_occupation | `False` | bool flag (set to enable) | | `--allow-auto-truncate` | Allow automatically truncating requests that exceed the maximum input length instead of returning an error. | `False` | bool flag (set to enable) | | `--enable-custom-logit-processor` | Enable users to pass custom logit processors to the server (disabled by default for security) | `False` | bool flag (set to enable) | | `--flashinfer-mla-disable-ragged` | Not using ragged prefill wrapper when running flashinfer mla | `False` | bool flag (set to enable) | | `--disable-shared-experts-fusion` | Disable shared experts fusion optimization for deepseek v3/r1. | `False` | bool flag (set to enable) | | `--disable-chunked-prefix-cache` | Disable chunked prefix cache feature for deepseek, which should save overhead for short sequences. | `False` | bool flag (set to enable) | | `--disable-fast-image-processor` | Adopt base image processor instead of fast image processor. | `False` | bool flag (set to enable) | | `--keep-mm-feature-on-device` | Keep multimodal feature tensors on device after processing to save D2H copy. | `False` | bool flag (set to enable) | | `--enable-return-hidden-states` | Enable returning hidden states with responses. | `False` | bool flag (set to enable) | | `--enable-return-routed-experts` | Enable returning routed experts of each layer with responses. | `False` | bool flag (set to enable) | | `--scheduler-recv-interval` | The interval to poll requests in scheduler. Can be set to >1 to reduce the overhead of this. | `1` | Type: int | | `--numa-node` | Sets the numa node for the subprocesses. i-th element corresponds to i-th subprocess. | `None` | List[int] | | `--enable-deterministic-inference` | Enable deterministic inference mode with batch invariant ops. | `False` | bool flag (set to enable) | | `--rl-on-policy-target` | The training system that SGLang needs to match for true on-policy. | `None` | `fsdp` | | `--enable-attn-tp-input-scattered` | Allow input of attention to be scattered when only using tensor parallelism, to reduce the computational load of operations such as qkv latent. | `False` | bool flag (set to enable) | | `--enable-nsa-prefill-context-parallel` | Enable context parallelism used in the long sequence prefill phase of DeepSeek v3.2. | `False` | bool flag (set to enable) | | `--nsa-prefill-cp-mode` | Token splitting mode for the prefill phase of DeepSeek v3.2 under context parallelism. Optional values: `round-robin-split`(default),`in-seq-split`. `round-robin-split` distributes tokens across ranks based on `token_idx % cp_size`. It supports multi-batch prefill, fused MoE, and FP8 KV cache. | `in-seq-split` | `in-seq-split`, `round-robin-split` | | `--enable-fused-qk-norm-rope` | Enable fused qk normalization and rope rotary embedding. | `False` | bool flag (set to enable) | | `--enable-precise-embedding-interpolation` | Enable corner alignment for resize of embeddings grid to ensure more accurate(but slower) evaluation of interpolated embedding values. | `False` | bool flag (set to enable) | ## Dynamic batch tokenizer | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--enable-dynamic-batch-tokenizer` | Enable async dynamic batch tokenizer for improved performance when multiple requests arrive concurrently. | `False` | bool flag (set to enable) | | `--dynamic-batch-tokenizer-batch-size` | [Only used if --enable-dynamic-batch-tokenizer is set] Maximum batch size for dynamic batch tokenizer. | `32` | Type: int | | `--dynamic-batch-tokenizer-batch-timeout` | [Only used if --enable-dynamic-batch-tokenizer is set] Timeout in seconds for batching tokenization requests. | `0.002` | Type: float | ## Debug tensor dumps | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--debug-tensor-dump-output-folder` | The output folder for dumping tensors. | `None` | Type: str | | `--debug-tensor-dump-layers` | The layer ids to dump. Dump all layers if not specified. | `None` | Type: JSON list | | `--debug-tensor-dump-input-file` | The input filename for dumping tensors | `None` | Type: str | | `--debug-tensor-dump-inject` | Inject the outputs from jax as the input of every layer. | `False` | Type: str | ## PD disaggregation | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--disaggregation-mode` | Only used for PD disaggregation. "prefill" for prefill-only server, and "decode" for decode-only server. If not specified, it is not PD disaggregated | `null` | `null`, `prefill`, `decode` | | `--disaggregation-transfer-backend` | The backend for disaggregation transfer. Default is mooncake. | `mooncake` | `mooncake`, `nixl`, `ascend`, `fake` | | `--disaggregation-bootstrap-port` | Bootstrap server port on the prefill server. Default is 8998. | `8998` | Type: int | | `--disaggregation-ib-device` | The InfiniBand devices for disaggregation transfer, accepts single device (e.g., --disaggregation-ib-device mlx5_0) or multiple comma-separated devices (e.g., --disaggregation-ib-device mlx5_0,mlx5_1). Default is None, which triggers automatic device detection when mooncake backend is enabled. | `None` | Type: str | | `--disaggregation-decode-enable-offload-kvcache` | Enable async KV cache offloading on decode server (PD mode). | `False` | bool flag (set to enable) | | `--num-reserved-decode-tokens` | Number of decode tokens that will have memory reserved when adding new request to the running batch. | `512` | Type: int | | `--disaggregation-decode-polling-interval` | The interval to poll requests in decode server. Can be set to >1 to reduce the overhead of this. | `1` | Type: int | ## Encode prefill disaggregation | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--encoder-only` | For MLLM with an encoder, launch an encoder-only server | `False` | bool flag (set to enable) | | `--language-only` | For VLM, load weights for the language model only. | `False` | bool flag (set to enable) | | `--encoder-transfer-backend` | The backend for encoder disaggregation transfer. Default is zmq_to_scheduler. | `zmq_to_scheduler` | `zmq_to_scheduler`, `zmq_to_tokenizer`, `mooncake` | | `--encoder-urls` | List of encoder server urls. | `[]` | Type: JSON list | ## Custom weight loader | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--custom-weight-loader` | The custom dataloader which used to update the model. Should be set with a valid import path, such as my_package.weight_load_func | `None` | List[str] | | `--weight-loader-disable-mmap` | Disable mmap while loading weight using safetensors. | `False` | bool flag (set to enable) | | `--remote-instance-weight-loader-seed-instance-ip` | The ip of the seed instance for loading weights from remote instance. | `None` | Type: str | | `--remote-instance-weight-loader-seed-instance-service-port` | The service port of the seed instance for loading weights from remote instance. | `None` | Type: int | | `--remote-instance-weight-loader-send-weights-group-ports` | The communication group ports for loading weights from remote instance. | `None` | Type: JSON list | | `--remote-instance-weight-loader-backend` | The backend for loading weights from remote instance. Can be 'transfer_engine' or 'nccl'. Default is 'nccl'. | `nccl` | `transfer_engine`, `nccl` | | `--remote-instance-weight-loader-start-seed-via-transfer-engine` | Start seed server via transfer engine backend for remote instance weight loader. | `False` | bool flag (set to enable) | ## For PD-Multiplexing | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--enable-pdmux` | Enable PD-Multiplexing, PD running on greenctx stream. | `False` | bool flag (set to enable) | | `--pdmux-config-path` | The path of the PD-Multiplexing config file. | `None` | Type: str | | `--sm-group-num` | Number of sm partition groups. | `8` | Type: int | ## Configuration file support | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--config` | Read CLI options from a config file. Must be a YAML file with configuration options. | `None` | Type: str | ## For Multi-Modal | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--mm-max-concurrent-calls` | The max concurrent calls for async mm data processing. | `32` | Type: int | | `--mm-per-request-timeout` | The timeout for each multi-modal request in seconds. | `10.0` | Type: int | | `--enable-broadcast-mm-inputs-process` | Enable broadcast mm-inputs process in scheduler. | `False` | bool flag (set to enable) | | `--mm-process-config` | Multimodal preprocessing config, a json config contains keys: `image`, `video`, `audio`. | `{}` | Type: JSON / Dict | | `--mm-enable-dp-encoder` | Enabling data parallelism for mm encoder. The dp size will be set to the tp size automatically. | `False` | bool flag (set to enable) | | `--limit-mm-data-per-request` | Limit the number of multimodal inputs per request. e.g. '{"image": 1, "video": 1, "audio": 1}' | `None` | Type: JSON / Dict | | `--enable-mm-global-cache` | Enable Mooncake-backed global multimodal embedding cache on encoder servers so repeated images can reuse cached ViT embeddings instead of recomputing them. | `False` | bool flag (set to enable) | ## For checkpoint decryption | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--decrypted-config-file` | The path of the decrypted config file. | `None` | Type: str | | `--decrypted-draft-config-file` | The path of the decrypted draft config file. | `None` | Type: str | | `--enable-prefix-mm-cache` | Enable prefix multimodal cache. Currently only supports mm-only. | `False` | bool flag (set to enable) | ## Forward hooks | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--forward-hooks` | JSON-formatted list of forward hook specifications. Each element must include `target_modules` (list of glob patterns matched against `model.named_modules()` names) and `hook_factory` (Python import path to a factory, e.g. `my_package.hooks:make_hook`). An optional `name` field is used for logging, and an optional `config` object is passed as a `dict` to the factory. | `None` | Type: JSON list | ## Deprecated arguments | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--enable-ep-moe` | NOTE: --enable-ep-moe is deprecated. Please set `--ep-size` to the same value as `--tp-size` instead. | `None` | N/A | | `--enable-deepep-moe` | NOTE: --enable-deepep-moe is deprecated. Please set `--moe-a2a-backend` to 'deepep' instead. | `None` | N/A | | `--prefill-round-robin-balance` | Note: Note: --prefill-round-robin-balance is deprecated now. | `None` | N/A | | `--enable-flashinfer-cutlass-moe` | NOTE: --enable-flashinfer-cutlass-moe is deprecated. Please set `--moe-runner-backend` to 'flashinfer_cutlass' instead. | `None` | N/A | | `--enable-flashinfer-cutedsl-moe` | NOTE: --enable-flashinfer-cutedsl-moe is deprecated. Please set `--moe-runner-backend` to 'flashinfer_cutedsl' instead. | `None` | N/A | | `--enable-flashinfer-trtllm-moe` | NOTE: --enable-flashinfer-trtllm-moe is deprecated. Please set `--moe-runner-backend` to 'flashinfer_trtllm' instead. | `None` | N/A | | `--enable-triton-kernel-moe` | NOTE: --enable-triton-kernel-moe is deprecated. Please set `--moe-runner-backend` to 'triton_kernel' instead. | `None` | N/A | | `--enable-flashinfer-mxfp4-moe` | NOTE: --enable-flashinfer-mxfp4-moe is deprecated. Please set `--moe-runner-backend` to 'flashinfer_mxfp4' instead. | `None` | N/A | | `--crash-on-nan` | Crash the server on nan logprobs. | `False` | Type: str | | `--hybrid-kvcache-ratio` | Mix ratio in [0,1] between uniform and hybrid kv buffers (0.0 = pure uniform: swa_size / full_size = 1)(1.0 = pure hybrid: swa_size / full_size = local_attention_size / context_length) | `None` | Optional[float] | | `--load-watch-interval` | The interval of load watching in seconds. | `0.1` | Type: float | | `--nsa-prefill` | Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention). | `flashmla_sparse` | `flashmla_sparse`, `flashmla_decode`, `fa3`, `tilelang`, `aiter` | | `--nsa-decode` | Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding. | `flashmla_kv` | `flashmla_prefill`, `flashmla_kv`, `fa3`, `tilelang`, `aiter` | ================================================ FILE: docs/advanced_features/sgl_model_gateway.md ================================================ # SGLang Model Gateway SGLang Model Gateway is a high-performance model-routing gateway for large-scale LLM deployments. It centralizes worker lifecycle management, balances traffic across heterogeneous protocols (HTTP, gRPC, OpenAI-compatible), and provides enterprise-ready control over history storage, MCP tooling, and privacy-sensitive workflows. The gateway is deeply optimized for the SGLang serving runtime, but can route to any OpenAI-compatible backend. --- ## Table of Contents 1. [Overview](#overview) 2. [Architecture](#architecture) - [Control Plane](#control-plane) - [Data Plane](#data-plane) - [Storage and Privacy](#storage-and-privacy) 3. [Installation](#installation) 4. [Quick Start](#quick-start) 5. [Deployment Modes](#deployment-modes) - [Co-launch Router and Workers](#co-launch-router-and-workers) - [Separate Launch (HTTP)](#separate-launch-http) - [gRPC Launch](#grpc-launch) - [Prefill-Decode Disaggregation](#prefill-decode-disaggregation) - [OpenAI Backend Proxy](#openai-backend-proxy) - [Multi-Model Inference Gateway](#multi-model-inference-gateway) 6. [API Reference](#api-reference) - [Inference Endpoints](#inference-endpoints) - [Tokenization Endpoints](#tokenization-endpoints) - [Parser Endpoints](#parser-endpoints) - [Classification API](#classification-api) - [Conversation and Response APIs](#conversation-and-response-apis) - [Worker Management APIs](#worker-management-apis) - [Admin and Health Endpoints](#admin-and-health-endpoints) 7. [Load Balancing Policies](#load-balancing-policies) 8. [Reliability and Flow Control](#reliability-and-flow-control) - [Retries](#retries) - [Circuit Breaker](#circuit-breaker) - [Rate Limiting and Queuing](#rate-limiting-and-queuing) - [Health Checks](#health-checks) 9. [Reasoning Parser Integration](#reasoning-parser-integration) 10. [Tool Call Parsing](#tool-call-parsing) 11. [Tokenizer Management](#tokenizer-management) 12. [MCP Integration](#mcp-integration) 13. [Service Discovery (Kubernetes)](#service-discovery-kubernetes) 14. [History and Data Connectors](#history-and-data-connectors) 15. [WASM Middleware](#wasm-middleware) 16. [Language Bindings](#language-bindings) 17. [Security and Authentication](#security-and-authentication) - [TLS (HTTPS) for Gateway Server](#tls-https-for-gateway-server) - [mTLS for Worker Communication](#mtls-for-worker-communication) 18. [Observability](#observability) - [Prometheus Metrics](#prometheus-metrics) - [OpenTelemetry Tracing](#opentelemetry-tracing) - [Logging](#logging) 19. [Production Recommendations](#production-recommendations) - [Security Best Practices](#security-best-practices) - [High Availability](#high-availability) - [Performance](#performance) - [Kubernetes Deployment](#kubernetes-deployment) - [Monitoring with PromQL](#monitoring-with-promql) 20. [Configuration Reference](#configuration-reference) 21. [Troubleshooting](#troubleshooting) --- ## Overview - **Unified control plane** for registering, monitoring, and orchestrating regular, prefill, and decode workers across heterogeneous model fleets. - **Multi-protocol data plane** that routes traffic across HTTP, PD (prefill/decode), gRPC, and OpenAI-compatible backends with shared reliability primitives. - **Industry-first gRPC pipeline** with native Rust tokenization, reasoning parsers, and tool-call execution for high-throughput, OpenAI-compatible serving; supports both single-stage and PD topologies. - **Inference Gateway Mode (`--enable-igw`)** dynamically instantiates multiple router stacks (HTTP regular/PD, gRPC) and applies per-model policies for multi-tenant deployments. - **Conversation & responses connectors** centralize chat history inside the router so the same context can be reused across models and MCP loops without leaking data to upstream vendors (memory, none, Oracle ATP, PostgreSQL). - **Enterprise privacy**: agentic multi-turn `/v1/responses`, native MCP client (STDIO/HTTP/SSE/Streamable), and history storage all operate within the router boundary. - **Reliability core**: retries with jitter, worker-scoped circuit breakers, token-bucket rate limiting with queuing, background health checks, and cache-aware load monitoring. - **Comprehensive observability**: 40+ Prometheus metrics, OpenTelemetry distributed tracing, structured logging, and request ID propagation. --- ## Architecture ### Control Plane - **Worker Manager** discovers capabilities (`/get_server_info`, `/get_model_info`), tracks load, and registers/removes workers in the shared registry. - **Job Queue** serializes add/remove requests and exposes status (`/workers/{worker_id}`) so clients can track onboarding progress. - **Load Monitor** feeds cache-aware and power-of-two policies with live worker load statistics. - **Health Checker** continuously probes workers and updates readiness, circuit breaker state, and router metrics. - **Tokenizer Registry** manages dynamically registered tokenizers with async loading from HuggingFace or local paths. ### Data Plane - **HTTP routers** (regular & PD) implement `/generate`, `/v1/chat/completions`, `/v1/completions`, `/v1/responses`, `/v1/embeddings`, `/v1/rerank`, `/v1/classify`, `/v1/tokenize`, `/v1/detokenize`, and associated admin endpoints. - **gRPC router** streams tokenized requests directly to SRT gRPC workers, running fully in Rust—tokenizer, reasoning parser, and tool parser all reside in-process. Supports both single-stage and PD routing, including embeddings and classification. - **OpenAI router** proxies OpenAI-compatible endpoints to external vendors (OpenAI, xAI, etc.) while keeping chat history and multi-turn orchestration local. ### Storage and Privacy - Conversation and response history is stored at the router tier (memory, none, Oracle ATP, or PostgreSQL). The same history can power multiple models or MCP loops without sending data to upstream vendors. - `/v1/responses` agentic flows, MCP sessions, and conversation APIs share the same storage layer, enabling compliance for regulated workloads. --- ## Installation ### Docker Pre-built Docker images are available on Docker Hub with multi-architecture support (x86_64 and ARM64): ```bash docker pull lmsysorg/sgl-model-gateway:latest ``` ### Prerequisites - **Rust and Cargo** ```bash curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh source "$HOME/.cargo/env" rustc --version cargo --version ``` - **Python** with `pip` and virtualenv tooling available. ### Rust Binary ```bash cd sgl-model-gateway cargo build --release ``` ### Python Package ```bash pip install maturin # Fast development mode cd sgl-model-gateway/bindings/python maturin develop # Production build maturin build --release --out dist --features vendored-openssl pip install --force-reinstall dist/*.whl ``` --- ## Quick Start ### Regular HTTP Routing ```bash # Rust binary ./target/release/sgl-model-gateway \ --worker-urls http://worker1:8000 http://worker2:8000 \ --policy cache_aware # Python launcher python -m sglang_router.launch_router \ --worker-urls http://worker1:8000 http://worker2:8000 \ --policy cache_aware ``` ### gRPC Routing ```bash python -m sglang_router.launch_router \ --worker-urls grpc://127.0.0.1:20000 \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --reasoning-parser deepseek-r1 \ --tool-call-parser json \ --host 0.0.0.0 --port 8080 ``` --- ## Deployment Modes ### Co-launch Router and Workers Launch the router and a fleet of SGLang workers in one process: ```bash python -m sglang_router.launch_server \ --model meta-llama/Meta-Llama-3.1-8B-Instruct \ --dp-size 4 \ --host 0.0.0.0 \ --port 30000 ``` Comprehensive example with router arguments (prefixed with `--router-`): ```bash python -m sglang_router.launch_server \ --host 0.0.0.0 \ --port 8080 \ --model meta-llama/Llama-3.1-8B-Instruct \ --tp-size 1 \ --dp-size 8 \ --grpc-mode \ --log-level debug \ --router-prometheus-port 10001 \ --router-tool-call-parser llama \ --router-model-path meta-llama/Llama-3.1-8B-Instruct \ --router-policy round_robin \ --router-log-level debug ``` ### Separate Launch (HTTP) Run workers independently and point the router at their HTTP endpoints: ```bash # Worker nodes python -m sglang.launch_server --model meta-llama/Meta-Llama-3.1-8B-Instruct --port 8000 python -m sglang.launch_server --model meta-llama/Meta-Llama-3.1-8B-Instruct --port 8001 # Router node python -m sglang_router.launch_router \ --worker-urls http://worker1:8000 http://worker2:8001 \ --policy cache_aware \ --host 0.0.0.0 --port 30000 ``` ### gRPC Launch Use SRT gRPC workers to unlock the highest throughput and access native reasoning/tool pipelines: ```bash # Workers expose gRPC endpoints python -m sglang.launch_server \ --model meta-llama/Llama-3.1-8B-Instruct \ --grpc-mode \ --port 20000 # Router python -m sglang_router.launch_router \ --worker-urls grpc://127.0.0.1:20000 \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --reasoning-parser deepseek-r1 \ --tool-call-parser json \ --host 0.0.0.0 --port 8080 ``` The gRPC router supports both regular HTTP-equivalent serving and PD (prefill/decode) serving. Provide `--tokenizer-path` or `--model-path` (HuggingFace ID or local directory) whenever connection mode resolves to gRPC. ### Prefill-Decode Disaggregation Split prefill and decode workers for PD-aware caching and balancing: ```bash python -m sglang_router.launch_router \ --pd-disaggregation \ --prefill http://prefill1:30001 9001 \ --decode http://decode1:30011 \ --prefill-policy cache_aware \ --decode-policy power_of_two ``` Prefill entries accept an optional bootstrap port. PD mode merges prefill metadata with decode outputs and streams results back to the client. ### OpenAI Backend Proxy Proxy OpenAI-compatible endpoints while keeping history and MCP sessions local: ```bash python -m sglang_router.launch_router \ --backend openai \ --worker-urls https://api.openai.com \ --history-backend memory ``` OpenAI backend mode expects exactly one `--worker-urls` entry per router instance. ### Multi-Model Inference Gateway Enable IGW mode to route multiple models through a single router: ```bash ./target/release/sgl-model-gateway \ --enable-igw \ --policy cache_aware \ --max-concurrent-requests 512 # Register workers dynamically curl -X POST http://localhost:30000/workers \ -H "Content-Type: application/json" \ -d '{ "url": "http://worker-a:8000", "model_id": "mistral", "priority": 10, "labels": {"tier": "gold"} }' ``` --- ## API Reference ### Inference Endpoints | Method | Path | Description | |--------|------|-------------| | `POST` | `/generate` | SGLang generate API | | `POST` | `/v1/chat/completions` | OpenAI-compatible chat completions (streaming/tool calls) | | `POST` | `/v1/completions` | OpenAI-compatible text completions | | `POST` | `/v1/embeddings` | Embedding generation (HTTP and gRPC) | | `POST` | `/v1/rerank`, `/rerank` | Reranking requests | | `POST` | `/v1/classify` | Text classification | ### Tokenization Endpoints The gateway provides HTTP endpoints for text tokenization with batch support, designed to mirror the SGLang Python tokenization API. | Method | Path | Description | |--------|------|-------------| | `POST` | `/v1/tokenize` | Tokenize text to token IDs (single or batch) | | `POST` | `/v1/detokenize` | Convert token IDs back to text (single or batch) | | `POST` | `/v1/tokenizers` | Register a new tokenizer (async, returns job status) | | `GET` | `/v1/tokenizers` | List all registered tokenizers | | `GET` | `/v1/tokenizers/{id}` | Get tokenizer info by UUID | | `GET` | `/v1/tokenizers/{id}/status` | Check async tokenizer loading status | | `DELETE` | `/v1/tokenizers/{id}` | Remove a tokenizer from the registry | #### Tokenize Request ```json { "model": "meta-llama/Llama-3.1-8B-Instruct", "prompt": "Hello, world!" } ``` #### Batch Tokenize Request ```json { "model": "meta-llama/Llama-3.1-8B-Instruct", "prompt": ["Hello", "World", "How are you?"] } ``` #### Tokenize Response ```json { "tokens": [15339, 11, 1917, 0], "count": 4, "char_count": 13 } ``` #### Detokenize Request ```json { "model": "meta-llama/Llama-3.1-8B-Instruct", "tokens": [15339, 11, 1917, 0], "skip_special_tokens": true } ``` #### Detokenize Response ```json { "text": "Hello, world!" } ``` #### Add Tokenizer (Async) ```bash curl -X POST http://localhost:30000/v1/tokenizers \ -H "Content-Type: application/json" \ -d '{"name": "llama3", "source": "meta-llama/Llama-3.1-8B-Instruct"}' ``` Response: ```json { "id": "550e8400-e29b-41d4-a716-446655440000", "status": "pending", "message": "Tokenizer registration queued" } ``` Check status: ```bash curl http://localhost:30000/v1/tokenizers/550e8400-e29b-41d4-a716-446655440000/status ``` ### Parser Endpoints The gateway provides admin endpoints for parsing reasoning content and function calls from LLM outputs. | Method | Path | Description | |--------|------|-------------| | `POST` | `/parse/reasoning` | Separate reasoning (``) from normal text | | `POST` | `/parse/function_call` | Parse function/tool calls from text | #### Separate Reasoning Request ```json { "text": "Let me analyze this step by step...The answer is 42.", "parser": "deepseek-r1" } ``` #### Response ```json { "normal_text": "The answer is 42.", "reasoning_text": "Let me analyze this step by step..." } ``` #### Function Call Parsing ```json { "text": "{\"name\": \"get_weather\", \"arguments\": {\"city\": \"NYC\"}}", "parser": "json" } ``` ### Classification API The `/v1/classify` endpoint provides text classification using sequence classification models (e.g., `Qwen2ForSequenceClassification`, `BertForSequenceClassification`). #### Request ```bash curl http://localhost:30000/v1/classify \ -H "Content-Type: application/json" \ -d '{ "model": "jason9693/Qwen2.5-1.5B-apeach", "input": "I love this product!" }' ``` #### Response ```json { "id": "classify-a1b2c3d4-5678-90ab-cdef-1234567890ab", "object": "list", "created": 1767034308, "model": "jason9693/Qwen2.5-1.5B-apeach", "data": [ { "index": 0, "label": "positive", "probs": [0.12, 0.88], "num_classes": 2 } ], "usage": { "prompt_tokens": 6, "completion_tokens": 0, "total_tokens": 6 } } ``` #### Response Fields | Field | Description | |-------|-------------| | `label` | Predicted class label (from model's `id2label` config, or `LABEL_N` fallback) | | `probs` | Probability distribution over all classes (softmax of logits) | | `num_classes` | Number of classification classes | #### Notes - Classification reuses the embedding backend—the scheduler returns logits which are converted to probabilities via softmax - Labels come from the model's HuggingFace config (`id2label` field); models without this mapping use generic labels (`LABEL_0`, `LABEL_1`, etc.) - Both HTTP and gRPC routers support classification ### Conversation and Response APIs | Method | Path | Description | |--------|------|-------------| | `POST` | `/v1/responses` | Create background responses (agentic loops) | | `GET` | `/v1/responses/{id}` | Retrieve stored response | | `POST` | `/v1/responses/{id}/cancel` | Cancel background response | | `DELETE` | `/v1/responses/{id}` | Delete response | | `GET` | `/v1/responses/{id}/input_items` | List response input items | | `POST` | `/v1/conversations` | Create conversation | | `GET` | `/v1/conversations/{id}` | Get conversation | | `POST` | `/v1/conversations/{id}` | Update conversation | | `DELETE` | `/v1/conversations/{id}` | Delete conversation | | `GET` | `/v1/conversations/{id}/items` | List conversation items | | `POST` | `/v1/conversations/{id}/items` | Add items to conversation | | `GET` | `/v1/conversations/{id}/items/{item_id}` | Get conversation item | | `DELETE` | `/v1/conversations/{id}/items/{item_id}` | Delete conversation item | ### Worker Management APIs | Method | Path | Description | |--------|------|-------------| | `POST` | `/workers` | Queue worker registration (returns 202 Accepted) | | `GET` | `/workers` | List workers with health, load, and policy metadata | | `GET` | `/workers/{worker_id}` | Inspect specific worker or job queue entry | | `PUT` | `/workers/{worker_id}` | Queue worker update | | `DELETE` | `/workers/{worker_id}` | Queue worker removal | #### Add Worker ```bash curl -X POST http://localhost:30000/workers \ -H "Content-Type: application/json" \ -d '{"url":"grpc://0.0.0.0:31000","worker_type":"regular"}' ``` #### List Workers ```bash curl http://localhost:30000/workers ``` Response: ```json { "workers": [ { "id": "2f3a0c3e-3a7b-4c3f-8c70-1b7d4c3a6e1f", "url": "http://0.0.0.0:31378", "model_id": "mistral", "priority": 50, "cost": 1.0, "worker_type": "regular", "is_healthy": true, "load": 0, "connection_mode": "Http" } ], "total": 1, "stats": { "prefill_count": 0, "decode_count": 0, "regular_count": 1 } } ``` ### Admin and Health Endpoints | Method | Path | Description | |--------|------|-------------| | `GET` | `/liveness` | Health check (always returns OK) | | `GET` | `/readiness` | Readiness check (checks healthy worker availability) | | `GET` | `/health` | Alias for liveness | | `GET` | `/health_generate` | Health generate test | | `GET` | `/engine_metrics` | Engine-level metrics from workers | | `GET` | `/v1/models` | List available models | | `GET` | `/get_model_info` | Get model information | | `GET` | `/get_server_info` | Get server information | | `POST` | `/flush_cache` | Clear all caches | | `GET` | `/get_loads` | Get all worker loads | | `POST` | `/wasm` | Upload WASM module | | `GET` | `/wasm` | List WASM modules | | `DELETE` | `/wasm/{module_uuid}` | Remove WASM module | --- ## Load Balancing Policies | Policy | Description | Usage | |--------|-------------|-------| | `random` | Uniform random selection | `--policy random` | | `round_robin` | Cycles through workers in order | `--policy round_robin` | | `power_of_two` | Samples two workers and picks the lighter one | `--policy power_of_two` | | `cache_aware` | Combines cache locality with load balancing (default) | `--policy cache_aware` | | `bucket` | Divides workers into load buckets with dynamic boundaries | `--policy bucket` | ### Cache-Aware Policy Tuning ```bash --cache-threshold 0.5 \ --balance-abs-threshold 32 \ --balance-rel-threshold 1.5 \ --eviction-interval-secs 120 \ --max-tree-size 67108864 ``` | Parameter | Default | Description | |-----------|---------|-------------| | `--cache-threshold` | 0.3 | Minimum prefix match ratio for cache hit | | `--balance-abs-threshold` | 64 | Absolute load difference before rebalancing | | `--balance-rel-threshold` | 1.5 | Relative load ratio before rebalancing | | `--eviction-interval-secs` | 120 | Cache eviction cadence in seconds | | `--max-tree-size` | 67108864 | Maximum nodes in cache tree | --- ## Reliability and Flow Control ### Retries Configure exponential backoff retries: ```bash python -m sglang_router.launch_router \ --worker-urls http://worker1:8000 http://worker2:8001 \ --retry-max-retries 5 \ --retry-initial-backoff-ms 50 \ --retry-max-backoff-ms 30000 \ --retry-backoff-multiplier 1.5 \ --retry-jitter-factor 0.2 ``` | Parameter | Default | Description | |-----------|---------|-------------| | `--retry-max-retries` | 5 | Maximum retry attempts | | `--retry-initial-backoff-ms` | 50 | Initial backoff duration (ms) | | `--retry-max-backoff-ms` | 5000 | Maximum backoff duration (ms) | | `--retry-backoff-multiplier` | 2.0 | Exponential backoff multiplier | | `--retry-jitter-factor` | 0.1 | Random jitter factor (0.0-1.0) | | `--disable-retries` | false | Disable retries entirely | **Retryable Status Codes:** 408, 429, 500, 502, 503, 504 ### Circuit Breaker Per-worker circuit breakers prevent cascading failures: ```bash python -m sglang_router.launch_router \ --worker-urls http://worker1:8000 http://worker2:8001 \ --cb-failure-threshold 5 \ --cb-success-threshold 2 \ --cb-timeout-duration-secs 30 \ --cb-window-duration-secs 60 ``` | Parameter | Default | Description | |-----------|---------|-------------| | `--cb-failure-threshold` | 5 | Consecutive failures to open circuit | | `--cb-success-threshold` | 2 | Successes to close from half-open | | `--cb-timeout-duration-secs` | 30 | Time before half-open attempt | | `--cb-window-duration-secs` | 60 | Failure counting window | | `--disable-circuit-breaker` | false | Disable circuit breaker | **Circuit Breaker States:** - **Closed**: Normal operation, requests allowed - **Open**: Failing, requests rejected immediately - **Half-Open**: Testing recovery, limited requests allowed ### Rate Limiting and Queuing ```bash python -m sglang_router.launch_router \ --worker-urls http://worker1:8000 http://worker2:8001 \ --max-concurrent-requests 256 \ --rate-limit-tokens-per-second 512 \ --queue-size 128 \ --queue-timeout-secs 30 ``` Requests beyond the concurrency limit wait in a FIFO queue. Returns: - `429 Too Many Requests` when queue is full - `408 Request Timeout` when queue timeout expires ### Health Checks ```bash --health-check-interval-secs 30 \ --health-check-timeout-secs 10 \ --health-success-threshold 2 \ --health-failure-threshold 3 \ --health-check-endpoint /health ``` --- ## Reasoning Parser Integration The gateway includes built-in reasoning parsers for models that use Chain-of-Thought (CoT) reasoning with explicit thinking blocks. ### Supported Parsers | Parser ID | Model Family | Think Tokens | |-----------|--------------|--------------| | `deepseek-r1` | DeepSeek-R1 | `...` (initial reasoning) | | `qwen3` | Qwen-3 | `...` | | `qwen3-thinking` | Qwen-3 Thinking | `...` (initial reasoning) | | `kimi` | Kimi K2 | Unicode think tokens | | `glm45` | GLM-4.5/4.6/4.7 | `...` | | `step3` | Step-3 | `...` | | `minimax` | MiniMax | `...` | ### Usage ```bash python -m sglang_router.launch_router \ --worker-urls grpc://127.0.0.1:20000 \ --model-path deepseek-ai/DeepSeek-R1 \ --reasoning-parser deepseek-r1 ``` The gRPC router automatically: 1. Detects reasoning blocks in streaming output 2. Separates reasoning content from normal text 3. Applies incremental streaming parsing with buffer management 4. Handles partial token detection for correct streaming behavior --- ## Tool Call Parsing The gateway supports parsing function/tool calls from LLM outputs in multiple formats. ### Supported Formats | Parser | Format | Description | |--------|--------|-------------| | `json` | JSON | Standard JSON tool calls | | `python` | Pythonic | Python function call syntax | | `xml` | XML | XML-formatted tool calls | ### Usage ```bash python -m sglang_router.launch_router \ --worker-urls grpc://127.0.0.1:20000 \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --tool-call-parser json ``` --- ## Tokenizer Management ### Tokenizer Sources The gateway supports multiple tokenizer backends: - **HuggingFace**: Load from HuggingFace Hub by model ID - **Local**: Load from local `tokenizer.json` or directory - **Tiktoken**: Auto-detect OpenAI GPT models (gpt-4, davinci, etc.) ### Configuration ```bash # HuggingFace model --model-path meta-llama/Llama-3.1-8B-Instruct # Local tokenizer --tokenizer-path /path/to/tokenizer.json # With chat template override --chat-template /path/to/template.jinja ``` ### Tokenizer Caching Two-level caching for optimal performance: | Cache | Type | Description | |-------|------|-------------| | L0 | Exact match | Whole-string caching for repeated prompts | | L1 | Prefix match | Prefix boundary matching for incremental prompts | ```bash --enable-l0-cache \ --l0-max-entries 10000 \ --enable-l1-cache \ --l1-max-memory 52428800 # 50MB ``` --- ## MCP Integration The gateway provides native Model Context Protocol (MCP) client integration for tool execution. ### Supported Transports | Transport | Description | |-----------|-------------| | STDIO | Local process execution | | SSE | Server-Sent Events (HTTP) | | Streamable | Bidirectional streaming | ### Configuration ```bash python -m sglang_router.launch_router \ --mcp-config-path /path/to/mcp-config.yaml \ --worker-urls http://worker1:8000 ``` ### MCP Configuration File ```yaml servers: - name: "filesystem" command: "npx" args: ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"] protocol: "stdio" required: false - name: "github" url: "https://api.github.com/mcp" token: "ghp_xxxxx" protocol: "sse" required: false - name: "custom-tools" url: "https://tools.example.com/mcp" protocol: "streamable" required: true pool: max_connections: 100 idle_timeout: 300 proxy: http: "http://proxy.internal:8080" https: "https://proxy.internal:8443" no_proxy: "localhost,127.0.0.1,*.internal" inventory: enable_refresh: true tool_ttl: 300 refresh_interval: 300 ``` --- ## Service Discovery (Kubernetes) Enable automatic worker discovery via Kubernetes pod selectors: ```bash python -m sglang_router.launch_router \ --service-discovery \ --selector app=sglang-worker role=inference \ --service-discovery-namespace production \ --service-discovery-port 8000 ``` ### PD Mode Discovery ```bash --pd-disaggregation \ --prefill-selector app=sglang component=prefill \ --decode-selector app=sglang component=decode \ --service-discovery ``` Prefill pods can expose bootstrap ports via the `sglang.ai/bootstrap-port` annotation. RBAC must allow `get`, `list`, and `watch` on pods. --- ## History and Data Connectors | Backend | Description | Usage | |---------|-------------|-------| | `memory` | In-memory storage (default) | `--history-backend memory` | | `none` | No persistence | `--history-backend none` | | `oracle` | Oracle Autonomous Database | `--history-backend oracle` | | `postgres` | PostgreSQL Database | `--history-backend postgres` | | `redis` | Redis | `--history-backend redis` | ### Oracle Configuration ```bash # Connection descriptor export ATP_DSN="(description=(address=(protocol=tcps)(port=1522)(host=adb.region.oraclecloud.com))(connect_data=(service_name=service_name)))" # Or TNS alias (requires wallet) export ATP_TNS_ALIAS="sglroutertestatp_high" export ATP_WALLET_PATH="/path/to/wallet" # Credentials export ATP_USER="admin" export ATP_PASSWORD="secret" export ATP_POOL_MIN=4 export ATP_POOL_MAX=32 python -m sglang_router.launch_router \ --backend openai \ --worker-urls https://api.openai.com \ --history-backend oracle ``` ### PostgreSQL Configuration ```bash export POSTGRES_DB_URL="postgres://user:password@host:5432/dbname" python -m sglang_router.launch_router \ --backend openai \ --worker-urls https://api.openai.com \ --history-backend postgres ``` ### Redis Configuration ```bash export REDIS_URL="redis://localhost:6379" export REDIS_POOL_MAX=16 export REDIS_RETENTION_DAYS=30 python -m sglang_router.launch_router \ --backend openai \ --worker-urls https://api.openai.com \ --history-backend redis \ --redis-retention-days 30 ``` Use `--redis-retention-days -1` for persistent storage (default is 30 days). --- ## WASM Middleware The gateway supports WebAssembly (WASM) middleware modules for custom request/response processing. This enables organization-specific logic for authentication, rate limiting, billing, logging, and more—without modifying or recompiling the gateway. ### Overview WASM middleware runs in a sandboxed environment with memory isolation, no network/filesystem access, and configurable resource limits. | Attach Point | When Executed | Use Cases | |--------------|---------------|-----------| | `OnRequest` | Before forwarding to workers | Auth, rate limiting, request modification | | `OnResponse` | After receiving worker response | Logging, response modification, error handling | | Action | Description | |--------|-------------| | `Continue` | Proceed without modification | | `Reject(status)` | Reject request with HTTP status code | | `Modify(...)` | Modify headers, body, or status | ### Examples Complete working examples are available in `examples/wasm/`: | Example | Description | |---------|-------------| | `auth/` | API key authentication for protected routes | | `rate_limit/` | Per-client rate limiting (requests/minute) | | `logging/` | Request tracking headers and response modification | The interface definition is located at `src/wasm/interface`. ### Building Modules ```bash # Prerequisites rustup target add wasm32-wasip2 cargo install wasm-tools # Build cargo build --target wasm32-wasip2 --release # Convert to component format wasm-tools component new \ target/wasm32-wasip2/release/my_middleware.wasm \ -o my_middleware.component.wasm ``` ### Deploying Modules ```bash # Enable WASM support python -m sglang_router.launch_router \ --worker-urls http://worker1:8000 \ --enable-wasm # Upload module curl -X POST http://localhost:30000/wasm \ -H "Content-Type: application/json" \ -d '{ "modules": [{ "name": "auth-middleware", "file_path": "/absolute/path/to/auth.component.wasm", "module_type": "Middleware", "attach_points": [{"Middleware": "OnRequest"}] }] }' # List modules curl http://localhost:30000/wasm # Remove module curl -X DELETE http://localhost:30000/wasm/{module_uuid} ``` ### Runtime Configuration | Parameter | Default | Description | |-----------|---------|-------------| | `max_memory_pages` | 1024 (64MB) | Maximum WASM memory | | `max_execution_time_ms` | 1000 | Execution timeout | | `max_stack_size` | 1MB | Stack size limit | | `module_cache_size` | 10 | Cached modules per worker | **Note:** Rate limiting state is per-worker thread and not shared across gateway replicas. For production, consider implementing rate limiting at a shared layer (e.g., Redis) --- ## Language Bindings SGLang Model Gateway provides official language bindings for Python and Go, enabling integration with different technology stacks and organizational requirements. ### Python Bindings The Python bindings provide a PyO3-based wrapper around the Rust gateway library. This is a straightforward binding that calls the gateway server startup from Python. #### Installation ```bash # From PyPI pip install sglang-router # Development build cd sgl-model-gateway/bindings/python pip install maturin && maturin develop --features vendored-openssl ``` #### Usage The Python bindings are used throughout this documentation. See the [Quick Start](#quick-start) and [Deployment Modes](#deployment-modes) sections for detailed examples. Key components: - `RouterArgs` dataclass with 50+ configuration options - `Router.from_args()` for programmatic startup - CLI commands: `smg launch`, `smg server`, `python -m sglang_router.launch_router` ### Go Bindings The Go bindings provide a high-performance gRPC client library for organizations with Go-based infrastructure. This is ideal for: - Integration with internal Go services and tooling - High-performance client applications - Building custom OpenAI-compatible proxy servers #### Architecture ``` ┌─────────────────────────────────────────┐ │ High-Level Go API │ │ (client.go - OpenAI-style interface) │ ├─────────────────────────────────────────┤ │ gRPC Layer │ ├─────────────────────────────────────────┤ │ Rust FFI Layer │ │ (Tokenization, Parsing, Conversion) │ └─────────────────────────────────────────┘ ``` **Key Features:** - Native Rust tokenization via FFI (thread-safe, lock-free) - Full streaming support with context cancellation - Configurable channel buffer sizes for high concurrency - Built-in tool call parsing and chat template application #### Installation ```bash # Build the FFI library first cd sgl-model-gateway/bindings/golang make build && make lib # Then use in your Go project go get github.com/sgl-project/sgl-go-sdk ``` **Requirements:** Go 1.24+, Rust toolchain #### Examples Complete working examples are available in `bindings/golang/examples/`: | Example | Description | |---------|-------------| | `simple/` | Non-streaming chat completion | | `streaming/` | Streaming chat completion with SSE | | `oai_server/` | Full OpenAI-compatible HTTP server | ```bash # Run examples cd sgl-model-gateway/bindings/golang/examples/simple && ./run.sh cd sgl-model-gateway/bindings/golang/examples/streaming && ./run.sh cd sgl-model-gateway/bindings/golang/examples/oai_server && ./run.sh ``` #### Testing ```bash cd sgl-model-gateway/bindings/golang # Unit tests go test -v ./... # Integration tests (requires running SGLang server) export SGL_GRPC_ENDPOINT=grpc://localhost:20000 export SGL_TOKENIZER_PATH=/path/to/tokenizer go test -tags=integration -v ./... ``` ### Comparison | Feature | Python | Go | |---------|--------|-----| | **Primary Use** | Gateway server launcher | gRPC client library | | **CLI Support** | Full CLI (smg, sglang-router) | Library only | | **K8s Discovery** | Native support | N/A (client library) | | **PD Mode** | Built-in | N/A (client library) | **When to Use Python:** Launching and managing the gateway server, service discovery, PD disaggregation. **When to Use Go:** Building custom client applications, integration with Go microservices, OpenAI-compatible proxy servers --- ## Security and Authentication ### Router API Key ```bash python -m sglang_router.launch_router \ --api-key "your-router-api-key" \ --worker-urls http://worker1:8000 ``` Clients must supply `Authorization: Bearer ` for protected endpoints. ### Worker API Keys ```bash # Add worker with explicit key curl -H "Authorization: Bearer router-key" \ -X POST http://localhost:8080/workers \ -H "Content-Type: application/json" \ -d '{"url":"http://worker:8000","api_key":"worker-key"}' ``` ### Security Configurations 1. **No Authentication** (default): Use only in trusted environments 2. **Router-only Authentication**: Clients authenticate to router 3. **Worker-only Authentication**: Router open, workers require keys 4. **Full Authentication**: Both router and workers protected ### TLS (HTTPS) for Gateway Server Enable TLS to serve the gateway over HTTPS: ```bash python -m sglang_router.launch_router \ --worker-urls http://worker1:8000 \ --tls-cert-path /path/to/server.crt \ --tls-key-path /path/to/server.key ``` | Parameter | Description | |-----------|-------------| | `--tls-cert-path` | Path to server certificate (PEM format) | | `--tls-key-path` | Path to server private key (PEM format) | Both parameters must be provided together. The gateway uses rustls with the ring crypto provider for TLS termination. If TLS is not configured, the gateway falls back to plain HTTP. ### mTLS for Worker Communication Enable mutual TLS (mTLS) for secure communication with workers in HTTP mode: ```bash python -m sglang_router.launch_router \ --worker-urls https://worker1:8443 https://worker2:8443 \ --client-cert-path /path/to/client.crt \ --client-key-path /path/to/client.key \ --ca-cert-path /path/to/ca.crt ``` | Parameter | Description | |-----------|-------------| | `--client-cert-path` | Path to client certificate for mTLS (PEM format) | | `--client-key-path` | Path to client private key for mTLS (PEM format) | | `--ca-cert-path` | Path to CA certificate for verifying worker TLS (PEM format, repeatable) | **Key Points:** - Client certificate and key must be provided together - Multiple CA certificates can be added with multiple `--ca-cert-path` flags - Uses rustls backend when TLS is configured - Single HTTP client is created for all workers (assumes single security domain) - TCP keepalive (30 seconds) is enabled for long-lived connections ### Full TLS Configuration Example Gateway HTTPS + Worker mTLS + API Key authentication: ```bash python -m sglang_router.launch_router \ --worker-urls https://worker1:8443 https://worker2:8443 \ --tls-cert-path /etc/certs/server.crt \ --tls-key-path /etc/certs/server.key \ --client-cert-path /etc/certs/client.crt \ --client-key-path /etc/certs/client.key \ --ca-cert-path /etc/certs/ca.crt \ --api-key "secure-api-key" \ --policy cache_aware ``` --- ## Observability ### Prometheus Metrics Enable with `--prometheus-host`/`--prometheus-port` (defaults to `0.0.0.0:29000`). #### Metric Categories (40+ metrics) | Layer | Prefix | Metrics | |-------|--------|---------| | HTTP | `smg_http_*` | `requests_total`, `request_duration_seconds`, `responses_total`, `connections_active`, `rate_limit_total` | | Router | `smg_router_*` | `requests_total`, `request_duration_seconds`, `request_errors_total`, `stage_duration_seconds`, `upstream_responses_total` | | Inference | `smg_router_*` | `ttft_seconds`, `tpot_seconds`, `tokens_total`, `generation_duration_seconds` | | Worker | `smg_worker_*` | `pool_size`, `connections_active`, `requests_active`, `health_checks_total`, `selection_total`, `errors_total` | | Circuit Breaker | `smg_worker_cb_*` | `state`, `transitions_total`, `outcomes_total`, `consecutive_failures`, `consecutive_successes` | | Retry | `smg_worker_*` | `retries_total`, `retries_exhausted_total`, `retry_backoff_seconds` | | Discovery | `smg_discovery_*` | `registrations_total`, `deregistrations_total`, `sync_duration_seconds`, `workers_discovered` | | MCP | `smg_mcp_*` | `tool_calls_total`, `tool_duration_seconds`, `servers_active`, `tool_iterations_total` | | Database | `smg_db_*` | `operations_total`, `operation_duration_seconds`, `connections_active`, `items_stored` | #### Key Inference Metrics (gRPC mode) | Metric | Type | Description | |--------|------|-------------| | `smg_router_ttft_seconds` | Histogram | Time to first token | | `smg_router_tpot_seconds` | Histogram | Time per output token | | `smg_router_tokens_total` | Counter | Total tokens (input/output) | | `smg_router_generation_duration_seconds` | Histogram | End-to-end generation time | #### Duration Buckets 1ms, 5ms, 10ms, 25ms, 50ms, 100ms, 250ms, 500ms, 1s, 2.5s, 5s, 10s, 15s, 30s, 45s, 60s, 90s, 120s, 180s, 240s ### OpenTelemetry Tracing Enable distributed tracing with OTLP export: ```bash python -m sglang_router.launch_router \ --worker-urls http://worker1:8000 \ --enable-trace \ --otlp-traces-endpoint localhost:4317 ``` #### Features - OTLP/gRPC exporter (default port 4317) - W3C Trace Context propagation for HTTP and gRPC - Batch span processing (500ms delay, 64 span batch size) - Custom filtering to reduce noise - Trace context injection into upstream worker requests - Service name: `sgl-router` ### Logging ```bash python -m sglang_router.launch_router \ --worker-urls http://worker1:8000 \ --log-level debug \ --log-dir ./router_logs ``` Structured tracing with optional file sink. Log levels: `debug`, `info`, `warn`, `error`. ### Request ID Propagation ```bash --request-id-headers x-request-id x-trace-id x-correlation-id ``` Responses include `x-request-id` header for correlation. --- ## Production Recommendations This section provides guidance for deploying SGLang Model Gateway in production environments. ### Security Best Practices **Always enable TLS in production:** ```bash python -m sglang_router.launch_router \ --worker-urls https://worker1:8443 https://worker2:8443 \ --tls-cert-path /etc/certs/server.crt \ --tls-key-path /etc/certs/server.key \ --client-cert-path /etc/certs/client.crt \ --client-key-path /etc/certs/client.key \ --ca-cert-path /etc/certs/ca.crt \ --api-key "${ROUTER_API_KEY}" ``` **Security Checklist:** - Enable TLS for gateway HTTPS termination - Enable mTLS for worker communication when workers are on untrusted networks - Set `--api-key` to protect router endpoints - Use Kubernetes Secrets or a secrets manager for credentials - Rotate certificates and API keys periodically - Restrict network access with firewalls or network policies ### High Availability **Scaling Strategy:** The gateway supports running multiple replicas behind a load balancer for high availability. However, there are important considerations: | Component | Shared Across Replicas | Impact | |-----------|----------------------|--------| | Worker Registry | No (independent) | Each replica discovers workers independently | | Radix Cache Tree | No (independent) | Cache hits may decrease by 10-20% | | Circuit Breaker State | No (independent) | Each replica tracks failures independently | | Rate Limiting | No (independent) | Limits apply per-replica, not globally | **Recommendations:** 1. **Prefer horizontal scaling over vertical scaling**: Deploy multiple smaller gateway replicas rather than one large instance with excessive CPU and memory. This provides: - Better fault tolerance (single replica failure doesn't take down the gateway) - More predictable resource usage - Easier capacity planning 2. **Use Kubernetes Service Discovery**: Let the gateway automatically discover and manage workers: ```bash python -m sglang_router.launch_router \ --service-discovery \ --selector app=sglang-worker \ --service-discovery-namespace production ``` 3. **Accept cache efficiency trade-off**: With multiple replicas, the cache-aware routing policy's radix tree is not synchronized across replicas. This means: - Each replica builds its own cache tree - Requests from the same user may hit different replicas - Expected cache hit rate reduction: **10-20%** - This is often acceptable given the HA benefits 4. **Configure session affinity (optional)**: If cache efficiency is critical, configure your load balancer for session affinity based on a consistent hash of the request (e.g., user ID or API key). **Example HA Architecture:** ``` ┌─────────────────┐ │ Load Balancer │ │ (L4/L7) │ └────────┬────────┘ ┌──────────────┼──────────────┐ │ │ │ ┌─────▼─────┐ ┌─────▼─────┐ ┌─────▼─────┐ │ Gateway │ │ Gateway │ │ Gateway │ │ Replica 1 │ │ Replica 2 │ │ Replica 3 │ └─────┬─────┘ └─────┬─────┘ └─────┬─────┘ │ │ │ └──────────────┼──────────────┘ │ ┌──────────────┼──────────────┐ │ │ │ ┌─────▼─────┐ ┌─────▼─────┐ ┌─────▼─────┐ │ Worker │ │ Worker │ │ Worker │ │ Pod 1 │ │ Pod 2 │ │ Pod N │ └───────────┘ └───────────┘ └───────────┘ ``` ### Performance **Use gRPC mode for high throughput:** gRPC mode provides the highest performance for SGLang workers: ```bash # Start workers in gRPC mode python -m sglang.launch_server \ --model meta-llama/Llama-3.1-8B-Instruct \ --grpc-mode \ --port 20000 # Configure gateway for gRPC python -m sglang_router.launch_router \ --worker-urls grpc://worker1:20000 grpc://worker2:20000 \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --policy cache_aware ``` **Performance Benefits of gRPC:** - Native Rust tokenization (no Python overhead) - Streaming with lower latency - Built-in reasoning parser execution - Tool call parsing in the gateway - Reduced serialization overhead **Tuning Recommendations:** | Parameter | Recommendation | Reason | |-----------|---------------|--------| | `--policy` | `cache_aware` | Best for repeated prompts, ~30% latency reduction | | `--max-concurrent-requests` | 2-4x worker count | Prevent overload while maximizing throughput | | `--queue-size` | 2x max-concurrent | Buffer for burst traffic | | `--request-timeout-secs` | Based on max generation length | Prevent stuck requests | ### Kubernetes Deployment **Pod Labeling for Service Discovery:** For the gateway to discover workers automatically, label your worker pods consistently: ```yaml # Worker Deployment (Regular Mode) apiVersion: apps/v1 kind: Deployment metadata: name: sglang-worker namespace: production spec: replicas: 4 selector: matchLabels: app: sglang-worker component: inference template: metadata: labels: app: sglang-worker component: inference model: llama-3-8b spec: containers: - name: worker image: lmsysorg/sglang:latest ports: - containerPort: 8000 name: http - containerPort: 20000 name: grpc ``` **Gateway configuration for discovery:** ```bash python -m sglang_router.launch_router \ --service-discovery \ --selector app=sglang-worker component=inference \ --service-discovery-namespace production \ --service-discovery-port 8000 ``` **PD (Prefill/Decode) Mode Labeling:** ```yaml # Prefill Worker metadata: labels: app: sglang-worker component: prefill annotations: sglang.ai/bootstrap-port: "9001" # Decode Worker metadata: labels: app: sglang-worker component: decode ``` **Gateway configuration for PD discovery:** ```bash python -m sglang_router.launch_router \ --service-discovery \ --pd-disaggregation \ --prefill-selector app=sglang-worker component=prefill \ --decode-selector app=sglang-worker component=decode \ --service-discovery-namespace production ``` **RBAC Requirements:** The gateway needs permissions to watch pods: ```yaml apiVersion: rbac.authorization.k8s.io/v1 kind: Role metadata: name: sglang-gateway namespace: production rules: - apiGroups: [""] resources: ["pods"] verbs: ["get", "list", "watch"] --- apiVersion: rbac.authorization.k8s.io/v1 kind: RoleBinding metadata: name: sglang-gateway namespace: production subjects: - kind: ServiceAccount name: sglang-gateway namespace: production roleRef: kind: Role name: sglang-gateway apiGroup: rbac.authorization.k8s.io ``` ### Monitoring with PromQL Configure Prometheus to scrape the gateway metrics endpoint (default: `:29000/metrics`). **Essential Dashboards:** **1. Request Rate and Latency:** ```promql # Request rate by endpoint sum(rate(smg_http_requests_total[5m])) by (path, method) # P50 latency histogram_quantile(0.50, sum(rate(smg_http_request_duration_seconds_bucket[5m])) by (le)) # P99 latency histogram_quantile(0.99, sum(rate(smg_http_request_duration_seconds_bucket[5m])) by (le)) # Error rate sum(rate(smg_http_responses_total{status=~"5.."}[5m])) / sum(rate(smg_http_responses_total[5m])) ``` **2. Worker Health:** ```promql # Healthy workers sum(smg_worker_pool_size) # Active connections per worker smg_worker_connections_active # Worker health check failures sum(rate(smg_worker_health_checks_total{result="failure"}[5m])) by (worker_id) ``` **3. Circuit Breaker Status:** ```promql # Circuit breaker states (0=closed, 1=open, 2=half-open) smg_worker_cb_state # Circuit breaker transitions sum(rate(smg_worker_cb_transitions_total[5m])) by (worker_id, from_state, to_state) # Workers with open circuits count(smg_worker_cb_state == 1) ``` **4. Inference Performance (gRPC mode):** ```promql # Time to first token (P50) histogram_quantile(0.50, sum(rate(smg_router_ttft_seconds_bucket[5m])) by (le, model)) # Time per output token (P99) histogram_quantile(0.99, sum(rate(smg_router_tpot_seconds_bucket[5m])) by (le, model)) # Token throughput sum(rate(smg_router_tokens_total[5m])) by (model, direction) # Generation duration P95 histogram_quantile(0.95, sum(rate(smg_router_generation_duration_seconds_bucket[5m])) by (le)) ``` **5. Rate Limiting and Queuing:** ```promql # Rate limit rejections sum(rate(smg_http_rate_limit_total{decision="rejected"}[5m])) # Queue depth (if using concurrency limiting) smg_worker_requests_active # Retry attempts sum(rate(smg_worker_retries_total[5m])) by (worker_id) # Exhausted retries (failures after all retries) sum(rate(smg_worker_retries_exhausted_total[5m])) ``` **6. MCP Tool Execution:** ```promql # Tool call rate sum(rate(smg_mcp_tool_calls_total[5m])) by (server, tool) # Tool latency P95 histogram_quantile(0.95, sum(rate(smg_mcp_tool_duration_seconds_bucket[5m])) by (le, tool)) # Active MCP server connections smg_mcp_servers_active ``` **Alerting Rules Example:** ```yaml groups: - name: sglang-gateway rules: - alert: HighErrorRate expr: | sum(rate(smg_http_responses_total{status=~"5.."}[5m])) / sum(rate(smg_http_responses_total[5m])) > 0.05 for: 5m labels: severity: critical annotations: summary: "High error rate on SGLang Gateway" - alert: CircuitBreakerOpen expr: count(smg_worker_cb_state == 1) > 0 for: 2m labels: severity: warning annotations: summary: "Worker circuit breaker is open" - alert: HighLatency expr: | histogram_quantile(0.99, sum(rate(smg_http_request_duration_seconds_bucket[5m])) by (le)) > 30 for: 5m labels: severity: warning annotations: summary: "P99 latency exceeds 30 seconds" - alert: NoHealthyWorkers expr: sum(smg_worker_pool_size) == 0 for: 1m labels: severity: critical annotations: summary: "No healthy workers available" ``` --- ## Configuration Reference ### Core Settings | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `--host` | str | 127.0.0.1 | Router host | | `--port` | int | 30000 | Router port | | `--worker-urls` | list | [] | Worker URLs (HTTP or gRPC) | | `--policy` | str | cache_aware | Routing policy | | `--max-concurrent-requests` | int | -1 | Concurrency limit (-1 disables) | | `--request-timeout-secs` | int | 600 | Request timeout | | `--max-payload-size` | int | 256MB | Maximum request payload | ### Prefill/Decode | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `--pd-disaggregation` | flag | false | Enable PD mode | | `--prefill` | list | [] | Prefill URLs + optional bootstrap ports | | `--decode` | list | [] | Decode URLs | | `--prefill-policy` | str | None | Override policy for prefill nodes | | `--decode-policy` | str | None | Override policy for decode nodes | | `--worker-startup-timeout-secs` | int | 600 | Worker init timeout | ### Kubernetes Discovery | Parameter | Type | Description | |-----------|------|-------------| | `--service-discovery` | flag | Enable discovery | | `--selector` | list | Label selectors (key=value) | | `--prefill-selector` / `--decode-selector` | list | PD mode selectors | | `--service-discovery-namespace` | str | Namespace to watch | | `--service-discovery-port` | int | Worker port (default 80) | | `--bootstrap-port-annotation` | str | Annotation for bootstrap ports | ### TLS Configuration | Parameter | Type | Description | |-----------|------|-------------| | `--tls-cert-path` | str | Server certificate for gateway HTTPS (PEM) | | `--tls-key-path` | str | Server private key for gateway HTTPS (PEM) | | `--client-cert-path` | str | Client certificate for worker mTLS (PEM) | | `--client-key-path` | str | Client private key for worker mTLS (PEM) | | `--ca-cert-path` | str | CA certificate for verifying workers (PEM, repeatable) | --- ## Troubleshooting ### Workers Never Ready Increase `--worker-startup-timeout-secs` or ensure health probes respond before router startup. ### Load Imbalance / Hot Workers Inspect `smg_router_requests_total` by worker and tune cache-aware thresholds (`--balance-*`, `--cache-threshold`). ### Circuit Breaker Flapping Increase `--cb-failure-threshold` or extend the timeout/window durations. Consider temporarily disabling retries. ### Queue Overflow (429) Increase `--queue-size` or reduce client concurrency. Ensure `--max-concurrent-requests` matches downstream capacity. ### Memory Growth Reduce `--max-tree-size` or lower `--eviction-interval-secs` for more aggressive cache pruning. ### Debugging ```bash python -m sglang_router.launch_router \ --worker-urls http://worker1:8000 \ --log-level debug \ --log-dir ./router_logs ``` ### gRPC Connection Issues Ensure workers are started with `--grpc-mode` and verify `--model-path` or `--tokenizer-path` is provided to the router. ### Tokenizer Loading Failures Check HuggingFace Hub credentials (`HF_TOKEN` environment variable) for private models. Verify local paths are accessible. --- SGLang Model Gateway continues to evolve alongside the SGLang runtime. Keep CLI flags, integrations, and documentation aligned when adopting new features or contributing improvements. ================================================ FILE: docs/advanced_features/sglang_for_rl.md ================================================ # SGLang for RL Systems This document is a practical guide for infrastructure teams integrating SGLang into RL and post-training systems. It focuses on the operational pain points in the loop (rollout, evaluation, training, weight sync) and maps them to concrete SGLang APIs, flags, and integration patterns. The focus is on maximizing rollout efficiency, accuracy and stability while keeping rollout-serving behavior aligned in production environments. ## Why SGLang for RL Lifecycle? Let's embrace a guiding principle from early DeepMind's RL engineering: **Be a library, not a framework.** This philosophy empowers innovation by providing SGLang as flexible tools, not rigid structures. Here are five reasons to use SGLang for your RL lifecycle: * **Fine-Grained Engine Sleep and Wake Up**: facilitate maximum-powered rollout and training * **Open-To-Use Refit Functionality**: diverse methods for co-location or disaggregation * **Easy To Postpone Generation**: enable partial rollout and dedicated rollout control * **Deterministic Inference**: achieve deterministic inference to enable zero training-inference mismatch * **Load Balancing Router**: cache-aware load-balancing for high-throughput rollout The following sections cover these aspects in detail. ## Fine-Grained Engine Sleep and Wake Up Rollout and training are both memory-intensive, and co-locating them on the same GPUs often leads to memory pressure and slow handoffs. SGLang provides a memory-aware sleep/wake mechanism that releases KV cache and weights while keeping the server process alive, then resumes them for rollout without a full restart. This avoids repeated disk I/O and CUDA graph recapture during each RL step. Under the hood, the RL team uses CUDA-graph-aware weight offload via [torch_memory_saver](https://github.com/fzyzcjy/torch_memory_saver) to preserve virtual memory addresses for graph replay. For details, see: [Efficient RL Training - Optimizing Memory Usage in verl](https://hebiao064.github.io/rl-memory-management). ### Server flag Enable memory saver support when launching the server: ``` --enable-memory-saver ``` ### Release Memory **Endpoint:** `POST /release_memory_occupation` **Request body:** | Field | Description | Defaults | Options | | --- | --- | --- | --- | | `tags` | Which memory regions to release. If omitted, all are released. | `None` | Type: list[str], values: `kv_cache`, `weights` | **Behavior notes:** - This call asserts there are no ongoing requests. Ensure the engine is idle before calling it. - If `kv_cache` is released, SGLang flushes cache; subsequent requests will rebuild KV cache as needed. ### Resume Memory **Endpoint:** `POST /resume_memory_occupation` **Request body:** | Field | Description | Defaults | Options | | --- | --- | --- | --- | | `tags` | Which memory regions to resume. If omitted, all are resumed. | `None` | Type: list[str], values: `kv_cache`, `weights` | ## Open-To-Use Refit Functionality After training completes each step, rollout engines must be refit with new weights. SGLang supports three refit strategies so you can match your infrastructure style (co-located vs disaggregated) and scaling needs. Each strategy maps to a concrete API with clear request schemas. For a deeper dive into SGLang's weight update utilities, see [RL System Deep Thinking: Weight Update Mechanisms](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/sys-design/readme-1-EN.md). **How to choose:** - **From disk** is simplest and best for elastic rollout scaling and checkpointing. - **From tensor** is best for co-located training/rollout when you can pass in-memory tensors. - **From distributed** is best for disaggregated training/rollout with dedicated communication groups (NCCL/IB). ### Update Weights from Disk **When to use:** - Save checkpoint to disk and update weights from disk - Dynamic scaling (new rollout instances can load from the same checkpoint) **Why it works well:** This path trades some I/O overhead for simplicity and flexibility. It integrates naturally with checkpointing and makes it trivial to add new rollout engines: point them at the same checkpoint and call the API. It is also the safest option for high availability because the checkpoint itself is the source of truth. **Endpoint:** `POST /update_weights_from_disk` **Request body:** | Field | Description | Defaults | Options | | --- | --- | --- | --- | | `model_path` | The model path with the new weights. | Required | Type: str | | `load_format` | The format to load the weights. | `None` | Type: str | | `abort_all_requests` | Abort all running requests before update. | `False` | Type: bool | | `weight_version` | Optional weight version label tracked by the server. | `None` | Type: str | | `is_async` | Perform weight load asynchronously. | `False` | Type: bool | | `torch_empty_cache` | Empty torch cache. | `False` | Type: bool | | `keep_pause` | Keep scheduler paused after update. | `False` | Type: bool | | `recapture_cuda_graph` | Recapture CUDA graphs after update. | `False` | Type: bool | | `token_step` | Trainer step id for rollout bookkeeping. | `0` | Type: int | | `flush_cache` | Flush KV cache after update. | `True` | Type: bool | **Response body:** | Field | Description | Defaults | Options | | --- | --- | --- | --- | | `success` | Whether the update succeeded. | - | Type: bool | | `message` | Status / error message. | - | Type: str | | `num_paused_requests` | Number of paused requests during update. | `0` | Type: int | **Python Engine API:** `engine.update_weights_from_disk(model_path, load_format=None)` **Diffusion engine (SGLang-Diffusion):** The diffusion engine exposes the same `POST /update_weights_from_disk` endpoint with the following behavior: - **All-or-nothing with rollback:** if any module fails to load, all previously updated modules are rolled back to the original weights by reloading from the original model path. No partial updates are left behind. If rollback itself fails, the exception propagates so the caller knows the model is in an inconsistent state. - **Offload-aware:** when layerwise offload (`--dit-layerwise-offload`) is enabled, the diffusion offload manager replaces GPU parameters with small `torch.empty((1,))` placeholders while real weights live in consolidated pinned CPU buffers. A naive `param.data.copy_()` would fail with a shape mismatch. Instead, the updater dynamically detects active offload managers and writes new weights directly into their CPU buffers, bypassing the placeholders entirely. For any layer that happens to be prefetched on GPU at update time, the live GPU tensor is also updated so the change takes effect immediately. This requires no extra GPU memory and does not disturb the offload state. - **DTensor-aware:** parameters distributed via `torch.distributed.tensor` (tensor parallelism) are updated through `distribute_tensor` so that each shard is correctly placed on the right device mesh. **Request body:** | Field | Description | Defaults | Options | | --- | --- | --- | --- | | `model_path` | The model path with the new weights. | Required | Type: str | | `flush_cache` | Flush TeaCache state after update. | `True` | Type: bool | | `target_modules` | List of module names to update (e.g. `["transformer"]`). If omitted, all `nn.Module` components are updated. | `None` | Type: list[str] | **Response body:** | Field | Description | Defaults | Options | | --- | --- | --- | --- | | `success` | Whether the update succeeded. | - | Type: bool | | `message` | Status / error message. | - | Type: str | > **Note:** The diffusion engine (SGLang-Diffusion) does not currently support hot refit (updating weights while inference is in progress). The diffusion scheduler processes one request at a time and completes the entire inference before handling the next request, so weight updates and inference never run concurrently. ### Update Weights from Tensor **When to use:** - Co-located training and rollout, where training can provide tensors directly - Fast in-memory updates **Important constraints:** This strategy requires the training process and rollout engine to share access to the tensors. Co-located setups must keep the model on GPU; moving tensors to CPU will break the update path. For high-performance MoE or specialized attention kernels, co-location may limit some optimizations compared to disaggregated rollouts. **Endpoint:** `POST /update_weights_from_tensor` **Request body:** | Field | Description | Defaults | Options | | --- | --- | --- | --- | | `serialized_named_tensors` | Per-TP serialized tensor payloads. | Required | Type: list[str|bytes] | | `load_format` | Optional load format selector. | `None` | `None`, `direct`, `flattened_bucket`, or a custom loader path string | | `flush_cache` | Flush KV cache after update. | `True` | Type: bool | | `abort_all_requests` | Abort all running requests before update. | `False` | Type: bool | | `weight_version` | Optional version label tracked by the server. | `None` | Type: str | **Note:** The serialized tensor payloads must be created with `MultiprocessingSerializer.serialize(...)` and should be base64-safe strings. **Python Engine API:** `engine.update_weights_from_tensor(named_tensors, load_format=None, flush_cache=True)` ### Update Weights from Distributed Group **When to use:** - Disaggregated training and rollout - NCCL or IB-backed weight broadcast from training workers to rollout workers **How it works:** Training workers gather weights (typically on TP rank 0), broadcast them to the rollout group, and each rollout TP shard loads the parameters it needs. This avoids disk I/O and keeps training and rollout decoupled, at the cost of managing a dedicated communication group. **Initialize weight update group** **Endpoint:** `POST /init_weights_update_group` **Request body:** | Field | Description | Defaults | Options | | --- | --- | --- | --- | | `master_address` | Group master address. | Required | Type: str | | `master_port` | Group master port. | Required | Type: int | | `rank_offset` | Offset for local rank mapping. | Required | Type: int | | `world_size` | Total world size. | Required | Type: int | | `group_name` | Group name. | `weight_update_group` | Type: str | | `backend` | Communication backend. | `nccl` | Type: str | **Update weight** **Endpoint:** `POST /update_weights_from_distributed` **Request body:** | Field | Description | Defaults | Options | | --- | --- | --- | --- | | `names` | Parameter names to update. | Required | Type: list[str] | | `dtypes` | Dtype strings for each parameter. | Required | Type: list[str] | | `shapes` | Tensor shapes. | Required | Type: list[list[int]] | | `group_name` | Group name. | `weight_update_group` | Type: str | | `flush_cache` | Flush KV cache after update. | `True` | Type: bool | | `abort_all_requests` | Abort all running requests before update. | `False` | Type: bool | | `weight_version` | Optional version label. | `None` | Type: str | | `load_format` | Optional format selector. | `None` | `None` or `flattened_bucket` | **Destroy weights update group** **Endpoint:** `POST /destroy_weights_update_group` **Request body:** | Field | Description | Defaults | Options | | --- | --- | --- | --- | | `group_name` | Group name. | `weight_update_group` | Type: str | **Python Engine APIs:** - `engine.init_weights_update_group(...)` - `engine.update_weights_from_distributed(names, dtypes, shapes, ...)` - `engine.destroy_weights_update_group(group_name)` ## Easy To Postpone Generation Multi-turn RL rollouts often suffer from long-tail requests that block the entire batch. A small number of slow interactions can stall all GPUs, and the long-tail behavior makes profiling and monitoring difficult. SGLang exposes explicit pause/resume APIs so you can pause slow requests and continue them later. This pattern matches systems like [APRIL](https://arxiv.org/abs/2509.18521), terminate once enough responses are collected, and recycle incomplete responses in the next step. The result is higher GPU utilization without discarding partial work. `pause_generation` --- update weights --- `continue_generation` is the correct execution flow when updating weights from training. An update can only happen when SGLang is not actively processing inference tasks. ### Pause Generation **Endpoint:** `POST /pause_generation` **Request body:** | Field | Description | Defaults | Options | | --- | --- | --- | --- | | `mode` | Pause mode. | `abort` | `abort`, `retract`, `in_place` | **Modes:** - `abort`: Default behavior, identical to `abort` endpoint with `abort_all` set. Pending requests from `waiting_queue` and `running_queue` will be returned immediately to the caller. - `retract`: Put engine in "paused" state. Move running requests back to waiting queue. KV cache can be flushed and recomputed later. - `in_place`: Put engine in "paused" state without changing states of the requests. Running requests rely on availability of KV caches to continue, so any subsequent `flush_cache` call will be unsuccessful. ### Continue Generation **Endpoint:** `POST /continue_generation` ## Deterministic Inference In many RL stacks, rollout and training are implemented with different kernels or batching behavior. Even when weights are identical, token probabilities can drift, silently breaking the on-policy assumption. This is the training–inference mismatch problem. SGLang supports a deterministic inference mode that reduces non-determinism across batch shapes. This mitigates variance introduced by runtime batching and kernel selection. To further achieve true on-policy training, you need to modify the training engine to use the same deterministic kernels. For implementation details, see these miles examples: [True On-Policy](https://github.com/radixark/miles/tree/main/examples/true_on_policy) and [True On-Policy for VLM](https://github.com/radixark/miles/tree/main/examples/true_on_policy_vlm). For additional context, see the blog post [Let Speed Be With Stability: All-In-One Solution to Training-Inference Mismatch with Miles](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/slime/mismatch/blog-en.md). **Server flag:** ``` --enable-deterministic-inference ``` For more details, see [Deterministic Inference](deterministic_inference.md) ## Load Balancing Router SGLang Model Gateway is the recommended control plane for large‑scale RL rollouts. It provides async, non‑blocking request handling, cache‑aware load balancing, and fault‑tolerant routing across rollout and reward servers. This lets you keep GPUs saturated while avoiding long‑tail stalls and brittle, engine‑local concurrency logic. It has been deployed in the training of GLM 4.5+ models and proven to be highly efficient in production-level large-scale RL workloads. Key benefits for RL infrastructure: - **Async non-blocking efficiency**: SGLang’s native async server/router architecture (HTTPS/gRPC) manages concurrency automatically. This guarantees maximum GPU saturation and effective continuous batching without requiring complex, manual implementation by engineers. - **Elasticity and fault tolerance**: By encapsulating the reward model and rollout as independent servers, SGLang decouples them logically and physically. This architecture provides robust disaster recovery for large-scale distributed training; if a server fails, the router automatically redirects traffic to healthy nodes, ensuring the training process continues without interruption. - **Training–Inference alignment**: Using the SGLang Model Gateway for both training and inference ensures "What You See Is What You Get." This eliminates score discrepancies and the painful backend alignment issues often caused by using different engines for training versus deployment. - **Dynamic load balancing and long-tail mitigation**: Unlike static partitioning, the SGLang Model Gateway enables request-level dynamic dispatching for multi-turn RL. It can distribute different turns of a conversation across different servers to balance workloads and eliminate long-tail latency caused by varying sequence lengths. For deployment and configuration, see: [SGLang Model Gateway](sgl_model_gateway.md) ================================================ FILE: docs/advanced_features/speculative_decoding.md ================================================ # Speculative Decoding SGLang provides several speculative decoding options, including EAGLE-2/EAGLE-3, MTP, classic draft-model decoding, and an NGRAM-based variant. Our implementation aims to maximize speed and efficiency and is considered to be among the fastest in open-source LLM engines. ## Summary ### Jump to sections - [EAGLE Decoding](#eagle-decoding) - [EAGLE-2 Decoding](#eagle-2-decoding) - [EAGLE-2 Decoding with torch.compile](#eagle-2-decoding-with-torchcompile) - [EAGLE-2 Decoding via Frequency-Ranked Speculative Sampling](#eagle-2-decoding-via-frequency-ranked-speculative-sampling) - [EAGLE-3 Decoding](#eagle-3-decoding) - [Multi Token Prediction](#multi-token-prediction) - [Standalone Speculative Decoding (Small Draft Model)](#standalone-speculative-decoding-small-draft-model) - [Speculative Decoding V2 (Overlap Scheduler)](#speculative-decoding-v2-overlap-scheduler) - [Ngram Speculative Decoding](#ngram-speculative-decoding) - [Full Parameter Reference](#full-parameter-reference) - [OOM Troubleshooting](#oom-troubleshooting) - [References](#references) ### Quick guidance - **Best speed/quality (recommended)**: Use **EAGLE-3** with `--speculative-algorithm EAGLE3`. - **Strong default / broad compatibility**: Use **EAGLE-2** with `--speculative-algorithm EAGLE`. - **Lower `lm_head` overhead for EAGLE-2**: Enable **FR-Spec** with `--speculative-token-map`. - **Model is MTP-enabled**: Use **MTP via speculative decoding** (often with small `speculative_num_steps/topk/num_draft_tokens`, see the example section). - **You have a smaller draft LLM**: Use **STANDALONE** (`--speculative-algorithm STANDALONE`). - **No extra model available**: Use **NGRAM** (`--speculative-algorithm NGRAM`, CUDA-only). - **Want overlap scheduler (experimental)**: Enable **SpecV2** with `SGLANG_ENABLE_SPEC_V2=True` (requires `--speculative-eagle-topk 1`). ### Method comparison (mini table) | Method | Draft source | Separate draft model? | How to enable | Notes / constraints | |---|---|---:|---|---| | EAGLE-2 | EAGLE draft model (feature drafting + tree) | Typically yes | `--speculative-algorithm EAGLE` + `--speculative-draft-model-path ...` | Tune `--speculative-num-steps`, `--speculative-eagle-topk`, `--speculative-num-draft-tokens` | | EAGLE-2 + `torch.compile` | Same as EAGLE-2 | Typically yes | Add `--enable-torch-compile` (optionally `--torch-compile-max-bs`) | Benefit varies by hardware/model; benchmark to verify | | EAGLE-2 + FR-Spec | Same as EAGLE-2 + token subset | Typically yes | Add `--speculative-token-map ...` | Reduces `lm_head` overhead with high-frequency token vocab | | EAGLE-3 | EAGLE3 draft model | Yes | `--speculative-algorithm EAGLE3` + `--speculative-draft-model-path ...` | Best throughput in the benchmark below | | MTP | Built-in multi-token heads (model-specific) | Often no | See **Multi Token Prediction** section | Uses speculative workflow; draft path may be auto-handled for some models | | STANDALONE | Smaller draft LLM (token-level) | Yes | `--speculative-algorithm STANDALONE` + `--speculative-draft-model-path ...` | Does **not** support `--enable-dp-attention` | | SpecV2 (experimental) | V2 workers + overlap scheduler | N/A | `SGLANG_ENABLE_SPEC_V2=True` | Only supports `--speculative-eagle-topk 1`; applies to `EAGLE`, `EAGLE3`, `STANDALONE` | | NGRAM | Ngram cache from previous tokens | No | `--speculative-algorithm NGRAM` | CUDA-only; no `--enable-dp-attention`; disables overlap scheduler & mixed chunked prefill | ### Performance Highlights Please see below for the huge improvements on throughput for LLaMA-Instruct 3.1 8B tested on MT bench that can be achieved via EAGLE3 decoding. For further details please see the [EAGLE3 paper](https://arxiv.org/pdf/2503.01840). | Method | Throughput (tokens/s) | |--------|----------------| | SGLang (w/o speculative, 1x H100) | 158.34 tokens/s | | SGLang + EAGLE-2 (1x H100) | 244.10 tokens/s | | SGLang + EAGLE-3 (1x H100) | 373.25 tokens/s | --- ## EAGLE Decoding To enable EAGLE speculative decoding the following parameters are relevant: | Parameter | Description | Default | |---|---|---| | `--speculative-draft-model-path` | Draft model path/weights. **Typically required** for EAGLE/EAGLE3 and STANDALONE. For some MTP-enabled models, this can be omitted. | `None` | | `--speculative-num-steps` | Depth of autoregressive drafting. Increases speculation range but risks rejection cascades. | Auto (`5` for Llama/Grok; `3` for many other models) | | `--speculative-eagle-topk` | Branching factor per step. Improves candidate diversity and acceptance rate, but increases memory/compute consumption. | Auto (`4` for Llama/Grok; `1` for many other models) | | `--speculative-num-draft-tokens` | Maximum parallel verification capacity. Allows deeper tree evaluation but increases GPU memory usage. | Auto (`8` for Llama/Grok; `4` for many other models). If `topk=1`, it is adjusted to `num_steps + 1`. | | `--speculative-accept-threshold-single` | Acceptance threshold for single-token verification. Lower values accept more aggressively. | `1.0` | | `--speculative-accept-threshold-acc` | Accumulated acceptance threshold across steps. | `1.0` | | `--speculative-attention-mode` | Attention mode for speculative operations (`prefill` or `decode`), affecting both target verification and draft extension. | `"prefill"` | | `--speculative-draft-attention-backend` | Override attention backend for the draft model. | `None` (same as target) | | `--speculative-draft-model-quantization` | Quantization method for the draft model. Use `"unquant"` to force no quantization even when the target model is quantized. | Same as target model | | `--speculative-draft-model-revision` | Specific revision/commit of the draft model to load. | `None` (auto-set to `"main"` when `--speculative-draft-model-path` is set and revision is omitted) | | `--speculative-draft-load-format` | Load format for the draft model weights. | `None` | These parameters are mostly the same for EAGLE-2 and EAGLE-3. `--speculative-token-map` is ignored for EAGLE-3 models. For `--speculative-num-steps`, `--speculative-eagle-topk`, and `--speculative-num-draft-tokens`: leave all three unset to use auto-tuning, or set all three explicitly when tuning. You can find the best combinations of these parameters with [bench_speculative.py](https://github.com/sgl-project/sglang/blob/main/scripts/playground/bench_speculative.py). ### EAGLE-2 Decoding You can enable EAGLE-2 Decoding by setting `--speculative-algorithm EAGLE` and choosing an appropriate model. **Launch the server:** ```bash python3 -m sglang.launch_server \ --model meta-llama/Llama-2-7b-chat-hf \ --speculative-algorithm EAGLE \ --speculative-draft-model-path lmsys/sglang-EAGLE-llama2-chat-7B \ --speculative-num-steps 3 \ --speculative-eagle-topk 4 \ --speculative-num-draft-tokens 16 \ --mem-fraction-static 0.7 \ --cuda-graph-max-bs 8 \ --log-level warning ``` **Send a request:** ```python import openai client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="None") response = client.chat.completions.create( model="meta-llama/Llama-2-7b-chat-hf", messages=[ {"role": "user", "content": "List 3 countries and their capitals."}, ], temperature=0, max_tokens=64, ) print(response.choices[0].message.content) ``` --- ### EAGLE-2 Decoding with `torch.compile` You can optionally enable `torch.compile` to apply kernel-level optimizations (operator fusion, autotune) to the draft model. The actual speedup depends on your hardware, model architecture, and batch size. In some configurations (e.g., small draft models on H100 where cuBLAS is already optimal and CUDA graphs are enabled), the benefit may be negligible. We recommend benchmarking with and without this flag on your specific setup to verify whether it helps. To enable it, add `--enable-torch-compile` and optionally set `--torch-compile-max-bs`: ```bash python3 -m sglang.launch_server \ --model meta-llama/Llama-2-7b-chat-hf \ --speculative-algorithm EAGLE \ --speculative-draft-model-path lmsys/sglang-EAGLE-llama2-chat-7B \ --speculative-num-steps 3 \ --speculative-eagle-topk 4 \ --speculative-num-draft-tokens 16 \ --mem-fraction-static 0.7 \ --enable-torch-compile \ --torch-compile-max-bs 8 \ --log-level warning ``` **Send a request:** ```python import openai client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="None") response = client.chat.completions.create( model="meta-llama/Llama-2-7b-chat-hf", messages=[ {"role": "user", "content": "List 3 countries and their capitals."}, ], temperature=0, max_tokens=64, ) print(response.choices[0].message.content) ``` --- ### EAGLE-2 Decoding via Frequency-Ranked Speculative Sampling By employing a truncated high-frequency token vocabulary in the draft model, EAGLE speculative decoding reduces `lm_head` computational overhead while accelerating the pipeline without quality degradation. For more details, check out [the paper](https://arxiv.org/pdf/2502.14856). In our implementation, set `--speculative-token-map` to enable the optimization. You can get the high-frequency tokens in FR-Spec from [this model](https://huggingface.co/thunlp/LLaMA3-Instruct-8B-FR-Spec). Or you can obtain high-frequency tokens by directly downloading these tokens from [this repo](https://github.com/thunlp/FR-Spec/tree/main?tab=readme-ov-file#prepare-fr-spec-vocabulary-subset). Thanks for the contribution from [Weilin Zhao](https://github.com/Achazwl) and [Zhousx](https://github.com/Zhou-sx). ```bash python3 -m sglang.launch_server \ --model meta-llama/Meta-Llama-3-8B-Instruct \ --speculative-algorithm EAGLE \ --speculative-draft-model-path lmsys/sglang-EAGLE-LLaMA3-Instruct-8B \ --speculative-num-steps 3 \ --speculative-eagle-topk 4 \ --speculative-num-draft-tokens 16 \ --speculative-token-map thunlp/LLaMA3-Instruct-8B-FR-Spec/freq_32768.pt \ --mem-fraction-static 0.7 \ --cuda-graph-max-bs 8 \ --dtype float16 \ --log-level warning ``` **Send a request:** ```python import openai client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="None") response = client.chat.completions.create( model="meta-llama/Meta-Llama-3-8B-Instruct", messages=[ {"role": "user", "content": "List 3 countries and their capitals."}, ], temperature=0, max_tokens=64, ) print(response.choices[0].message.content) ``` --- ### EAGLE-3 Decoding You can enable EAGLE-3 decoding by setting `--speculative-algorithm EAGLE3` and choosing an appropriate model. ```bash python3 -m sglang.launch_server \ --model meta-llama/Meta-Llama-3.1-8B-Instruct \ --speculative-algorithm EAGLE3 \ --speculative-draft-model-path jamesliu1/sglang-EAGLE3-Llama-3.1-Instruct-8B \ --speculative-num-steps 3 \ --speculative-eagle-topk 4 \ --speculative-num-draft-tokens 16 \ --mem-fraction-static 0.7 \ --cuda-graph-max-bs 8 \ --dtype float16 \ --log-level warning ``` **Send a request:** ```python import openai client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="None") response = client.chat.completions.create( model="meta-llama/Meta-Llama-3.1-8B-Instruct", messages=[ {"role": "user", "content": "List 3 countries and their capitals."}, ], temperature=0, max_tokens=64, ) print(response.choices[0].message.content) ``` --- ## Multi Token Prediction We support [MTP (Multi-Token Prediction)](https://arxiv.org/pdf/2404.19737) in SGLang by using speculative decoding. We use `XiaomiMiMo/MiMo-7B-RL` as an example here (for DeepSeek MTP usage, refer to [deepseek_v32 doc](../basic_usage/deepseek_v32.md#multi-token-prediction)). ```bash python3 -m sglang.launch_server \ --model XiaomiMiMo/MiMo-7B-RL \ --host 0.0.0.0 \ --trust-remote-code \ --speculative-algorithm EAGLE \ --speculative-num-steps 1 \ --speculative-eagle-topk 1 \ --speculative-num-draft-tokens 2 \ --mem-fraction-static 0.7 \ --cuda-graph-max-bs 8 \ --log-level warning ``` **Send a request:** ```python import requests url = "http://localhost:30000/v1/chat/completions" data = { "model": "XiaomiMiMo/MiMo-7B-RL", "messages": [{"role": "user", "content": "What is the capital of France?"}], } response = requests.post(url, json=data) print(response.json()) ``` --- ## Standalone Speculative Decoding (Small Draft Model) Besides EAGLE/MTP, SGLang also supports **token-level speculative decoding** using a smaller **draft model**. Enable it with `--speculative-algorithm STANDALONE` and provide a draft model via `--speculative-draft-model-path`. Relevant parameters: | Parameter | Description | Default | |---|---|---| | `--speculative-draft-model-path` | Draft model weights (smaller than the target model). | `None` | | `--speculative-num-steps` | Draft depth (how many steps the draft model runs autoregressively). | `3` (auto default for STANDALONE) | | `--speculative-eagle-topk` | Branching factor (token candidates per step). | `1` (auto default for STANDALONE) | | `--speculative-num-draft-tokens` | Verification capacity. | `4` (auto default for STANDALONE) | | `--speculative-draft-model-quantization` | Quantization for the draft model. Use `"unquant"` to disable quantization on the draft even when the target is quantized. | Same as target | > **Note:** Standalone speculative decoding currently **does not support** `--enable-dp-attention`. ```bash python3 -m sglang.launch_server \ --model Qwen/Qwen2.5-7B-Instruct \ --speculative-algorithm STANDALONE \ --speculative-draft-model-path Qwen/Qwen2.5-1.5B-Instruct \ --speculative-num-steps 4 \ --speculative-eagle-topk 2 \ --speculative-num-draft-tokens 7 \ --mem-fraction-static 0.7 \ --cuda-graph-max-bs 8 \ --log-level warning ``` **Send a request:** ```python import openai client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="None") response = client.chat.completions.create( model="Qwen/Qwen2.5-7B-Instruct", messages=[ {"role": "user", "content": "List 3 countries and their capitals."}, ], temperature=0, max_tokens=64, ) print(response.choices[0].message.content) ``` --- ## Speculative Decoding V2 (Overlap Scheduler) SGLang provides an **experimental Speculative Decoding V2** implementation that enables an overlap scheduler and uses V2 speculative workers (e.g. `StandaloneWorkerV2`, `EAGLEWorkerV2`). To enable it, set the environment variable: - `SGLANG_ENABLE_SPEC_V2=True` Notes: - SpecV2 currently only supports `--speculative-eagle-topk 1`. When SpecV2 is enabled, **set `--speculative-eagle-topk 1` explicitly**. - If you explicitly set `--speculative-eagle-topk > 1`, the server will error. - If you omit `--speculative-eagle-topk`, auto-tuning may pick `topk > 1` for some models (e.g. Llama). This is incompatible with SpecV2 and may not always trigger an immediate config error, so set `--speculative-eagle-topk 1` explicitly. - This applies to `EAGLE`, `EAGLE3`, and `STANDALONE`. ```bash SGLANG_ENABLE_SPEC_V2=True python3 -m sglang.launch_server \ --model Qwen/Qwen2.5-7B-Instruct \ --speculative-algorithm STANDALONE \ --speculative-draft-model-path Qwen/Qwen2.5-1.5B-Instruct \ --speculative-num-steps 4 \ --speculative-eagle-topk 1 \ --speculative-num-draft-tokens 5 \ --mem-fraction-static 0.7 \ --cuda-graph-max-bs 8 \ --log-level warning ``` **Send a request:** ```python import openai client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="None") response = client.chat.completions.create( model="Qwen/Qwen2.5-7B-Instruct", messages=[ {"role": "user", "content": "List 3 countries and their capitals."}, ], temperature=0, max_tokens=64, ) print(response.choices[0].message.content) ``` --- ## Ngram Speculative Decoding SGLang also supports **ngram-based speculative decoding** (no separate draft model). It retrieves draft tokens from an ngram cache built from previously generated tokens, and then verifies them with the target model. Enable it with: - `--speculative-algorithm NGRAM` ### Ngram-specific parameters | Parameter | Description | Default | |---|---|---| | `--speculative-num-draft-tokens` | Number of draft tokens verified per step. If omitted, defaults to `--speculative-ngram-max-match-window-size`. | `12` (with default ngram settings) | | `--speculative-ngram-min-match-window-size` | Minimum matching window size. | `1` | | `--speculative-ngram-max-match-window-size` | Maximum matching window size. | `12` | | `--speculative-ngram-min-bfs-breadth` | Minimum BFS breadth. | `1` | | `--speculative-ngram-max-bfs-breadth` | Maximum BFS breadth. | `10` | | `--speculative-ngram-match-type` | Match type: `"BFS"` or `"PROB"`. | `"BFS"` | | `--speculative-ngram-branch-length` | How many recent tokens to insert into the cache. | `18` | | `--speculative-ngram-capacity` | Cache capacity (number of entries). | `10,000,000` | Notes: - Ngram speculative decoding **only supports CUDA**. - It currently **does not support** `--enable-dp-attention`. - It disables the overlap scheduler and mixed chunked prefill. - If `--speculative-ngram-max-bfs-breadth > 1` (thus `speculative_eagle_topk > 1`) and `page_size > 1`, use `--attention-backend flashinfer`; otherwise the server will error. - Optional: set `SGLANG_NGRAM_FORCE_GREEDY_VERIFY=True` to force greedy verification. ```bash python3 -m sglang.launch_server \ --model Qwen/Qwen2.5-7B-Instruct \ --speculative-algorithm NGRAM \ --speculative-num-draft-tokens 16 \ --speculative-ngram-max-match-window-size 12 \ --speculative-ngram-max-bfs-breadth 10 \ --mem-fraction-static 0.7 \ --cuda-graph-max-bs 8 \ --log-level warning ``` **Send a request:** ```python import openai client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="None") response = client.chat.completions.create( model="Qwen/Qwen2.5-7B-Instruct", messages=[ {"role": "user", "content": "List 3 countries and their capitals."}, ], temperature=0, max_tokens=64, ) print(response.choices[0].message.content) ``` --- ## Full Parameter Reference Below is a comprehensive list of all speculative decoding parameters available in SGLang: ### Core parameters | Parameter | Type | Default | Description | |---|---|---|---| | `--speculative-algorithm` | `str` | `None` | Algorithm to use: `EAGLE`, `EAGLE3`, `STANDALONE`, `NGRAM`, `NEXTN` (alias of `EAGLE`) | | `--speculative-draft-model-path` | `str` | `None` | Path to the draft model weights | | `--speculative-draft-model-revision` | `str` | `None` | Specific revision/commit of the draft model (`"main"` is auto-used when draft path is set and revision is omitted) | | `--speculative-draft-load-format` | `str` | `None` | Load format for draft model weights | | `--speculative-num-steps` | `int` | `None` (auto-chosen when omitted) | Autoregressive drafting depth | | `--speculative-eagle-topk` | `int` | `None` (auto-chosen when omitted) | Branching factor per drafting step | | `--speculative-num-draft-tokens` | `int` | `None` (auto-chosen when omitted) | Maximum number of draft tokens for verification | | `--speculative-accept-threshold-single` | `float` | `1.0` | Single-token acceptance threshold | | `--speculative-accept-threshold-acc` | `float` | `1.0` | Accumulated acceptance threshold | | `--speculative-token-map` | `str` | `None` | Path to FR-Spec high-frequency token map | | `--speculative-attention-mode` | `str` | `"prefill"` | Attention mode for speculative operations (`"prefill"` or `"decode"`) | | `--speculative-draft-attention-backend` | `str` | `None` | Override attention backend for the draft model | | `--speculative-moe-runner-backend` | `str` | `None` | MoE runner backend for the draft model | | `--speculative-moe-a2a-backend` | `str` | `None` | MoE all-to-all backend for the draft model | | `--speculative-draft-model-quantization` | `str` | Same as target | Quantization for the draft model (`"unquant"` to disable) | ### Ngram-specific parameters | Parameter | Type | Default | Description | |---|---|---|---| | `--speculative-ngram-min-match-window-size` | `int` | `1` | Minimum ngram matching window | | `--speculative-ngram-max-match-window-size` | `int` | `12` | Maximum ngram matching window | | `--speculative-ngram-min-bfs-breadth` | `int` | `1` | Minimum BFS breadth | | `--speculative-ngram-max-bfs-breadth` | `int` | `10` | Maximum BFS breadth | | `--speculative-ngram-match-type` | `str` | `"BFS"` | Match type: `"BFS"` or `"PROB"` | | `--speculative-ngram-branch-length` | `int` | `18` | Recent tokens to insert into cache | | `--speculative-ngram-capacity` | `int` | `10,000,000` | Cache capacity | ### Environment variables | Variable | Default | Description | |---|---|---| | `SGLANG_ENABLE_SPEC_V2` | `False` | Enable Speculative Decoding V2 (overlap scheduler) | | `SGLANG_NGRAM_FORCE_GREEDY_VERIFY` | `False` | Force greedy verification for ngram decoding | ### Other related flags | Parameter | Description | |---|---| | `--enable-multi-layer-eagle` | Enable multi-layer EAGLE (auto-enabled for MiMoV2 and Step3p5 models) | | `--enable-torch-compile` | Enable `torch.compile` for kernel-level optimizations | | `--torch-compile-max-bs` | Maximum batch size for `torch.compile` | --- ## OOM Troubleshooting > [!WARNING] > **Out of Memory (OOM)?** Speculative decoding may increase GPU memory usage because the draft tree, CUDA graphs, and verification-related buffers consume additional VRAM. If you encounter OOM errors, try the following adjustments. ### Step 1: Lower static memory fraction (most effective) ```bash --mem-fraction-static 0.5 # when omitted, this value is auto-computed ``` - `--mem-fraction-static` controls the memory budget for model weights + KV cache pool. - Lowering it directly increases dynamic headroom for activations and CUDA graph buffers. - If omitted, SGLang auto-estimates this value from other settings, and those auto settings can still be too aggressive for some workloads. ### Step 2: Reduce CUDA graph batch size ```bash # Fewer CUDA graph captures = less memory reserved --cuda-graph-max-bs 4 # or even 2 for tight memory situations ``` - If omitted, `--cuda-graph-max-bs` is auto-selected based on GPU memory and TP size, and can be much larger on high-memory GPUs. ### Step 3: Reduce draft tree size These three parameters directly control how much memory the draft tree consumes: ```bash # Before (aggressive, high memory) --speculative-num-steps 5 --speculative-eagle-topk 8 --speculative-num-draft-tokens 64 # After (conservative, lower memory) --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 ``` ### Step 4: Limit concurrent requests ```bash # Fewer concurrent requests lowers in-flight load and can reduce OOM risk --max-running-requests 4 ``` ### Quick OOM recovery recipe If you're hitting OOM and just want something that works, start with this minimal configuration and scale up: ```bash python3 -m sglang.launch_server \ --model \ --speculative-algorithm EAGLE \ --speculative-draft-model-path \ --speculative-num-steps 3 \ --speculative-eagle-topk 1 \ --speculative-num-draft-tokens 4 \ --cuda-graph-max-bs 2 \ --mem-fraction-static 0.5 \ --max-running-requests 4 \ --log-level warning ``` Then gradually increase `--speculative-num-draft-tokens`, `--speculative-eagle-topk`, and `--cuda-graph-max-bs`. Increase `--mem-fraction-static` last, only after the run is stable. --- ## References EAGLE process is as follows: - Within EAGLE the draft model predicts the next feature vector, i.e. the last hidden state of the original LLM, using the feature sequence $(f_1, ..., f_k)$ and the token sequence $(t_2, ..., t_{k+1})$. - The next token is then sampled from $p_{k+2}=\text{LMHead}(f_{k+1})$. Afterwards, the two sequences are extended in a tree style—branching out multiple potential continuations, with the branching factor per step controlled by the `speculative_eagle_topk` parameter—to ensure a more coherent connection of context, and are given as input again. - In SGLang's EAGLE-2 implementation, the draft tree is expanded for the configured steps and then reranked to select the top `speculative_num_draft_tokens` final nodes as draft tokens. - EAGLE-3 removes the feature prediction objective, incorporates low and mid-layer features, and is trained in an on-policy manner. This enhances drafting accuracy by operating on features instead of tokens for more regular inputs and by additionally passing tokens from the next timestep to reduce sampling randomness. For more details, see the [EAGLE-2](https://arxiv.org/abs/2406.16858) and [EAGLE-3](https://arxiv.org/abs/2503.01840) papers. For guidance on how to train your own EAGLE model please see the [EAGLE repo](https://github.com/SafeAILab/EAGLE/tree/main?tab=readme-ov-file#train). For EAGLE-3 training specifically, check out [SpecForge](https://github.com/sgl-project/SpecForge), the SGLang team's training framework designed for EAGLE-3 speculative decoding models with seamless porting to SGLang serving. See the [SpecForge documentation](https://docs.sglang.ai/SpecForge/) and [blog post](https://lmsys.org/blog/2025-07-25-spec-forge) for details. ================================================ FILE: docs/advanced_features/structured_outputs.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Structured Outputs" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can specify a JSON schema, [regular expression](https://en.wikipedia.org/wiki/Regular_expression) or [EBNF](https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_form) to constrain the model output. The model output will be guaranteed to follow the given constraints. Only one constraint parameter (`json_schema`, `regex`, or `ebnf`) can be specified for a request.\n", "\n", "SGLang supports three grammar backends:\n", "\n", "- [XGrammar](https://github.com/mlc-ai/xgrammar)(default): Supports JSON schema, regular expression, and EBNF constraints.\n", "- [Outlines](https://github.com/dottxt-ai/outlines): Supports JSON schema and regular expression constraints.\n", "- [Llguidance](https://github.com/guidance-ai/llguidance): Supports JSON schema, regular expression, and EBNF constraints.\n", "\n", "We suggest using XGrammar for its better performance and utility. XGrammar currently uses the [GGML BNF format](https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md). For more details, see [XGrammar technical overview](https://blog.mlc.ai/2024/11/22/achieving-efficient-flexible-portable-structured-generation-with-xgrammar).\n", "\n", "To use Outlines, simply add `--grammar-backend outlines` when launching the server.\n", "To use llguidance, add `--grammar-backend llguidance` when launching the server.\n", "If no backend is specified, XGrammar will be used as the default.\n", "\n", "For better output quality, **It's advisable to explicitly include instructions in the prompt to guide the model to generate the desired format.** For example, you can specify, 'Please generate the output in the following JSON format: ...'.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## OpenAI Compatible API" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import openai\n", "import os\n", "\n", "from sglang.test.doc_patch import launch_server_cmd\n", "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", "\n", "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n", "\n", "\n", "server_process, port = launch_server_cmd(\n", " \"python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --host 0.0.0.0 --log-level warning\"\n", ")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=server_process)\n", "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### JSON\n", "\n", "you can directly define a JSON schema or use [Pydantic](https://docs.pydantic.dev/latest/) to define and validate the response." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Using Pydantic**" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from pydantic import BaseModel, Field\n", "\n", "\n", "# Define the schema using Pydantic\n", "class CapitalInfo(BaseModel):\n", " name: str = Field(..., pattern=r\"^\\w+$\", description=\"Name of the capital city\")\n", " population: int = Field(..., description=\"Population of the capital city\")\n", "\n", "\n", "response = client.chat.completions.create(\n", " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", " messages=[\n", " {\n", " \"role\": \"user\",\n", " \"content\": \"Please generate the information of the capital of France in the JSON format.\",\n", " },\n", " ],\n", " temperature=0,\n", " max_tokens=128,\n", " response_format={\n", " \"type\": \"json_schema\",\n", " \"json_schema\": {\n", " \"name\": \"foo\",\n", " # convert the pydantic model to json schema\n", " \"schema\": CapitalInfo.model_json_schema(),\n", " },\n", " },\n", ")\n", "\n", "response_content = response.choices[0].message.content\n", "# validate the JSON response by the pydantic model\n", "capital_info = CapitalInfo.model_validate_json(response_content)\n", "print_highlight(f\"Validated response: {capital_info.model_dump_json()}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**JSON Schema Directly**\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\n", "\n", "json_schema = json.dumps(\n", " {\n", " \"type\": \"object\",\n", " \"properties\": {\n", " \"name\": {\"type\": \"string\", \"pattern\": \"^[\\\\w]+$\"},\n", " \"population\": {\"type\": \"integer\"},\n", " },\n", " \"required\": [\"name\", \"population\"],\n", " }\n", ")\n", "\n", "response = client.chat.completions.create(\n", " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", " messages=[\n", " {\n", " \"role\": \"user\",\n", " \"content\": \"Give me the information of the capital of France in the JSON format.\",\n", " },\n", " ],\n", " temperature=0,\n", " max_tokens=128,\n", " response_format={\n", " \"type\": \"json_schema\",\n", " \"json_schema\": {\"name\": \"foo\", \"schema\": json.loads(json_schema)},\n", " },\n", ")\n", "\n", "print_highlight(response.choices[0].message.content)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### EBNF" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ebnf_grammar = \"\"\"\n", "root ::= city | description\n", "city ::= \"London\" | \"Paris\" | \"Berlin\" | \"Rome\"\n", "description ::= city \" is \" status\n", "status ::= \"the capital of \" country\n", "country ::= \"England\" | \"France\" | \"Germany\" | \"Italy\"\n", "\"\"\"\n", "\n", "response = client.chat.completions.create(\n", " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", " messages=[\n", " {\"role\": \"system\", \"content\": \"You are a helpful geography bot.\"},\n", " {\n", " \"role\": \"user\",\n", " \"content\": \"Give me the information of the capital of France.\",\n", " },\n", " ],\n", " temperature=0,\n", " max_tokens=32,\n", " extra_body={\"ebnf\": ebnf_grammar},\n", ")\n", "\n", "print_highlight(response.choices[0].message.content)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Regular expression" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = client.chat.completions.create(\n", " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", " messages=[\n", " {\"role\": \"user\", \"content\": \"What is the capital of France?\"},\n", " ],\n", " temperature=0,\n", " max_tokens=128,\n", " extra_body={\"regex\": \"(Paris|London)\"},\n", ")\n", "\n", "print_highlight(response.choices[0].message.content)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Structural Tag" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "tool_get_current_weather = {\n", " \"type\": \"function\",\n", " \"function\": {\n", " \"name\": \"get_current_weather\",\n", " \"description\": \"Get the current weather in a given location\",\n", " \"parameters\": {\n", " \"type\": \"object\",\n", " \"properties\": {\n", " \"city\": {\n", " \"type\": \"string\",\n", " \"description\": \"The city to find the weather for, e.g. 'San Francisco'\",\n", " },\n", " \"state\": {\n", " \"type\": \"string\",\n", " \"description\": \"the two-letter abbreviation for the state that the city is\"\n", " \" in, e.g. 'CA' which would mean 'California'\",\n", " },\n", " \"unit\": {\n", " \"type\": \"string\",\n", " \"description\": \"The unit to fetch the temperature in\",\n", " \"enum\": [\"celsius\", \"fahrenheit\"],\n", " },\n", " },\n", " \"required\": [\"city\", \"state\", \"unit\"],\n", " },\n", " },\n", "}\n", "\n", "tool_get_current_date = {\n", " \"type\": \"function\",\n", " \"function\": {\n", " \"name\": \"get_current_date\",\n", " \"description\": \"Get the current date and time for a given timezone\",\n", " \"parameters\": {\n", " \"type\": \"object\",\n", " \"properties\": {\n", " \"timezone\": {\n", " \"type\": \"string\",\n", " \"description\": \"The timezone to fetch the current date and time for, e.g. 'America/New_York'\",\n", " }\n", " },\n", " \"required\": [\"timezone\"],\n", " },\n", " },\n", "}\n", "\n", "schema_get_current_weather = tool_get_current_weather[\"function\"][\"parameters\"]\n", "schema_get_current_date = tool_get_current_date[\"function\"][\"parameters\"]\n", "\n", "\n", "def get_messages():\n", " return [\n", " {\n", " \"role\": \"system\",\n", " \"content\": f\"\"\"\n", "# Tool Instructions\n", "- Always execute python code in messages that you share.\n", "- When looking for real time information use relevant functions if available else fallback to brave_search\n", "You have access to the following functions:\n", "Use the function 'get_current_weather' to: Get the current weather in a given location\n", "{tool_get_current_weather[\"function\"]}\n", "Use the function 'get_current_date' to: Get the current date and time for a given timezone\n", "{tool_get_current_date[\"function\"]}\n", "If a you choose to call a function ONLY reply in the following format:\n", "<{{start_tag}}={{function_name}}>{{parameters}}{{end_tag}}\n", "where\n", "start_tag => ` a JSON dict with the function argument name as key and function argument value as value.\n", "end_tag => ``\n", "Here is an example,\n", "{{\"example_name\": \"example_value\"}}\n", "Reminder:\n", "- Function calls MUST follow the specified format\n", "- Required parameters MUST be specified\n", "- Only call one function at a time\n", "- Put the entire function call reply on one line\n", "- Always add your sources when using search results to answer the user query\n", "You are a helpful assistant.\"\"\",\n", " },\n", " {\n", " \"role\": \"user\",\n", " \"content\": \"You are in New York. Please get the current date and time, and the weather.\",\n", " },\n", " ]\n", "\n", "\n", "messages = get_messages()\n", "\n", "response = client.chat.completions.create(\n", " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", " messages=messages,\n", " response_format={\n", " \"type\": \"structural_tag\",\n", " \"structures\": [\n", " {\n", " \"begin\": \"\",\n", " \"schema\": schema_get_current_weather,\n", " \"end\": \"\",\n", " },\n", " {\n", " \"begin\": \"\",\n", " \"schema\": schema_get_current_date,\n", " \"end\": \"\",\n", " },\n", " ],\n", " \"triggers\": [\"\",\n", " \"content\": {\n", " \"type\": \"json_schema\",\n", " \"json_schema\": schema_get_current_weather,\n", " },\n", " \"end\": \"\",\n", " },\n", " {\n", " \"begin\": \"\",\n", " \"content\": {\n", " \"type\": \"json_schema\",\n", " \"json_schema\": schema_get_current_date,\n", " },\n", " \"end\": \"\",\n", " },\n", " ],\n", " \"at_least_one\": False,\n", " \"stop_after_first\": False,\n", " },\n", " },\n", ")\n", "\n", "print_highlight(response.choices[0].message.content)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Native API and SGLang Runtime (SRT)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### JSON" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Using Pydantic**" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import requests\n", "import json\n", "from pydantic import BaseModel, Field\n", "\n", "from transformers import AutoTokenizer\n", "\n", "tokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Meta-Llama-3.1-8B-Instruct\")\n", "\n", "\n", "# Define the schema using Pydantic\n", "class CapitalInfo(BaseModel):\n", " name: str = Field(..., pattern=r\"^\\w+$\", description=\"Name of the capital city\")\n", " population: int = Field(..., description=\"Population of the capital city\")\n", "\n", "\n", "# Make API request\n", "messages = [\n", " {\n", " \"role\": \"user\",\n", " \"content\": \"Here is the information of the capital of France in the JSON format.\\n\",\n", " }\n", "]\n", "text = tokenizer.apply_chat_template(\n", " messages, tokenize=False, add_generation_prompt=True, return_dict=False\n", ")\n", "response = requests.post(\n", " f\"http://localhost:{port}/generate\",\n", " json={\n", " \"text\": text,\n", " \"sampling_params\": {\n", " \"temperature\": 0,\n", " \"max_new_tokens\": 64,\n", " \"json_schema\": json.dumps(CapitalInfo.model_json_schema()),\n", " },\n", " },\n", ")\n", "print_highlight(response.json())\n", "\n", "\n", "response_data = json.loads(response.json()[\"text\"])\n", "# validate the response by the pydantic model\n", "capital_info = CapitalInfo.model_validate(response_data)\n", "print_highlight(f\"Validated response: {capital_info.model_dump_json()}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**JSON Schema Directly**" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "json_schema = json.dumps(\n", " {\n", " \"type\": \"object\",\n", " \"properties\": {\n", " \"name\": {\"type\": \"string\", \"pattern\": \"^[\\\\w]+$\"},\n", " \"population\": {\"type\": \"integer\"},\n", " },\n", " \"required\": [\"name\", \"population\"],\n", " }\n", ")\n", "\n", "# JSON\n", "response = requests.post(\n", " f\"http://localhost:{port}/generate\",\n", " json={\n", " \"text\": text,\n", " \"sampling_params\": {\n", " \"temperature\": 0,\n", " \"max_new_tokens\": 64,\n", " \"json_schema\": json_schema,\n", " },\n", " },\n", ")\n", "\n", "print_highlight(response.json())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### EBNF" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "messages = [\n", " {\n", " \"role\": \"user\",\n", " \"content\": \"Give me the information of the capital of France.\",\n", " }\n", "]\n", "text = tokenizer.apply_chat_template(\n", " messages, tokenize=False, add_generation_prompt=True, return_dict=False\n", ")\n", "response = requests.post(\n", " f\"http://localhost:{port}/generate\",\n", " json={\n", " \"text\": text,\n", " \"sampling_params\": {\n", " \"max_new_tokens\": 128,\n", " \"temperature\": 0,\n", " \"n\": 3,\n", " \"ebnf\": (\n", " \"root ::= city | description\\n\"\n", " 'city ::= \"London\" | \"Paris\" | \"Berlin\" | \"Rome\"\\n'\n", " 'description ::= city \" is \" status\\n'\n", " 'status ::= \"the capital of \" country\\n'\n", " 'country ::= \"England\" | \"France\" | \"Germany\" | \"Italy\"'\n", " ),\n", " },\n", " \"stream\": False,\n", " \"return_logprob\": False,\n", " },\n", ")\n", "\n", "print_highlight(response.json())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Regular expression" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "messages = [\n", " {\n", " \"role\": \"user\",\n", " \"content\": \"Paris is the capital of\",\n", " }\n", "]\n", "text = tokenizer.apply_chat_template(\n", " messages, tokenize=False, add_generation_prompt=True, return_dict=False\n", ")\n", "response = requests.post(\n", " f\"http://localhost:{port}/generate\",\n", " json={\n", " \"text\": text,\n", " \"sampling_params\": {\n", " \"temperature\": 0,\n", " \"max_new_tokens\": 64,\n", " \"regex\": \"(France|England)\",\n", " },\n", " },\n", ")\n", "print_highlight(response.json())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Structural Tag" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from transformers import AutoTokenizer\n", "\n", "# generate an answer\n", "tokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Meta-Llama-3.1-8B-Instruct\")\n", "\n", "text = tokenizer.apply_chat_template(\n", " messages, tokenize=False, add_generation_prompt=True, return_dict=False\n", ")\n", "payload = {\n", " \"text\": text,\n", " \"sampling_params\": {\n", " \"structural_tag\": json.dumps(\n", " {\n", " \"type\": \"structural_tag\",\n", " \"structures\": [\n", " {\n", " \"begin\": \"\",\n", " \"schema\": schema_get_current_weather,\n", " \"end\": \"\",\n", " },\n", " {\n", " \"begin\": \"\",\n", " \"schema\": schema_get_current_date,\n", " \"end\": \"\",\n", " },\n", " ],\n", " \"triggers\": [\"\",\n", " \"content\": {\n", " \"type\": \"json_schema\",\n", " \"json_schema\": schema_get_current_weather,\n", " },\n", " \"end\": \"\",\n", " },\n", " {\n", " \"begin\": \"\",\n", " \"content\": {\n", " \"type\": \"json_schema\",\n", " \"json_schema\": schema_get_current_date,\n", " },\n", " \"end\": \"\",\n", " },\n", " ],\n", " \"at_least_one\": False,\n", " \"stop_after_first\": False,\n", " },\n", " }\n", " )\n", " },\n", "}\n", "\n", "\n", "# Send POST request to the API endpoint\n", "response = requests.post(f\"http://localhost:{port}/generate\", json=payload)\n", "print_highlight(response.json())" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Offline Engine API" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import sglang as sgl\n", "\n", "llm = sgl.Engine(\n", " model_path=\"meta-llama/Meta-Llama-3.1-8B-Instruct\", grammar_backend=\"xgrammar\"\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### JSON" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Using Pydantic**" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\n", "from pydantic import BaseModel, Field\n", "\n", "prompts = [\n", " \"Give me the information of the capital of China in the JSON format.\",\n", " \"Give me the information of the capital of France in the JSON format.\",\n", " \"Give me the information of the capital of Ireland in the JSON format.\",\n", "]\n", "\n", "\n", "# Define the schema using Pydantic\n", "class CapitalInfo(BaseModel):\n", " name: str = Field(..., pattern=r\"^\\w+$\", description=\"Name of the capital city\")\n", " population: int = Field(..., description=\"Population of the capital city\")\n", "\n", "\n", "sampling_params = {\n", " \"temperature\": 0.1,\n", " \"top_p\": 0.95,\n", " \"json_schema\": json.dumps(CapitalInfo.model_json_schema()),\n", "}\n", "\n", "outputs = llm.generate(prompts, sampling_params)\n", "for prompt, output in zip(prompts, outputs):\n", " print_highlight(\"===============================\")\n", " print_highlight(f\"Prompt: {prompt}\") # validate the output by the pydantic model\n", " capital_info = CapitalInfo.model_validate_json(output[\"text\"])\n", " print_highlight(f\"Validated output: {capital_info.model_dump_json()}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**JSON Schema Directly**" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "prompts = [\n", " \"Give me the information of the capital of China in the JSON format.\",\n", " \"Give me the information of the capital of France in the JSON format.\",\n", " \"Give me the information of the capital of Ireland in the JSON format.\",\n", "]\n", "\n", "json_schema = json.dumps(\n", " {\n", " \"type\": \"object\",\n", " \"properties\": {\n", " \"name\": {\"type\": \"string\", \"pattern\": \"^[\\\\w]+$\"},\n", " \"population\": {\"type\": \"integer\"},\n", " },\n", " \"required\": [\"name\", \"population\"],\n", " }\n", ")\n", "\n", "sampling_params = {\"temperature\": 0.1, \"top_p\": 0.95, \"json_schema\": json_schema}\n", "\n", "outputs = llm.generate(prompts, sampling_params)\n", "for prompt, output in zip(prompts, outputs):\n", " print_highlight(\"===============================\")\n", " print_highlight(f\"Prompt: {prompt}\\nGenerated text: {output['text']}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### EBNF\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "prompts = [\n", " \"Give me the information of the capital of France.\",\n", " \"Give me the information of the capital of Germany.\",\n", " \"Give me the information of the capital of Italy.\",\n", "]\n", "\n", "sampling_params = {\n", " \"temperature\": 0.8,\n", " \"top_p\": 0.95,\n", " \"ebnf\": (\n", " \"root ::= city | description\\n\"\n", " 'city ::= \"London\" | \"Paris\" | \"Berlin\" | \"Rome\"\\n'\n", " 'description ::= city \" is \" status\\n'\n", " 'status ::= \"the capital of \" country\\n'\n", " 'country ::= \"England\" | \"France\" | \"Germany\" | \"Italy\"'\n", " ),\n", "}\n", "\n", "outputs = llm.generate(prompts, sampling_params)\n", "for prompt, output in zip(prompts, outputs):\n", " print_highlight(\"===============================\")\n", " print_highlight(f\"Prompt: {prompt}\\nGenerated text: {output['text']}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Regular expression" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "prompts = [\n", " \"Please provide information about London as a major global city:\",\n", " \"Please provide information about Paris as a major global city:\",\n", "]\n", "\n", "sampling_params = {\"temperature\": 0.8, \"top_p\": 0.95, \"regex\": \"(France|England)\"}\n", "\n", "outputs = llm.generate(prompts, sampling_params)\n", "for prompt, output in zip(prompts, outputs):\n", " print_highlight(\"===============================\")\n", " print_highlight(f\"Prompt: {prompt}\\nGenerated text: {output['text']}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Structural Tag" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "text = tokenizer.apply_chat_template(\n", " messages, tokenize=False, add_generation_prompt=True, return_dict=False\n", ")\n", "prompts = [text]\n", "\n", "\n", "sampling_params = {\n", " \"temperature\": 0.8,\n", " \"top_p\": 0.95,\n", " \"structural_tag\": json.dumps(\n", " {\n", " \"type\": \"structural_tag\",\n", " \"structures\": [\n", " {\n", " \"begin\": \"\",\n", " \"schema\": schema_get_current_weather,\n", " \"end\": \"\",\n", " },\n", " {\n", " \"begin\": \"\",\n", " \"schema\": schema_get_current_date,\n", " \"end\": \"\",\n", " },\n", " ],\n", " \"triggers\": [\"\",\n", " \"content\": {\n", " \"type\": \"json_schema\",\n", " \"json_schema\": schema_get_current_weather,\n", " },\n", " \"end\": \"\",\n", " },\n", " {\n", " \"begin\": \"\",\n", " \"content\": {\n", " \"type\": \"json_schema\",\n", " \"json_schema\": schema_get_current_date,\n", " },\n", " \"end\": \"\",\n", " },\n", " ],\n", " \"at_least_one\": False,\n", " \"stop_after_first\": False,\n", " },\n", " }\n", " ),\n", "}\n", "\n", "\n", "# Send POST request to the API endpoint\n", "outputs = llm.generate(prompts, sampling_params)\n", "for prompt, output in zip(prompts, outputs):\n", " print_highlight(\"===============================\")\n", " print_highlight(f\"Prompt: {prompt}\\nGenerated text: {output['text']}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "llm.shutdown()" ] } ], "metadata": { "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 2 } ================================================ FILE: docs/advanced_features/structured_outputs_for_reasoning_models.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Structured Outputs For Reasoning Models\n", "\n", "When working with reasoning models that use special tokens like `...` to denote reasoning sections, you might want to allow free-form text within these sections while still enforcing grammar constraints on the rest of the output.\n", "\n", "SGLang provides a feature to disable grammar restrictions within reasoning sections. This is particularly useful for models that need to perform complex reasoning steps before providing a structured output.\n", "\n", "To enable this feature, use the `--reasoning-parser` flag which decide the think_end_token, such as ``, when launching the server. You can also specify the reasoning parser using the `--reasoning-parser` flag.\n", "\n", "## Supported Models\n", "\n", "Currently, SGLang supports the following reasoning models:\n", "- [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d): The reasoning content is wrapped with `` and `` tags.\n", "- [QwQ](https://huggingface.co/Qwen/QwQ-32B): The reasoning content is wrapped with `` and `` tags.\n", "\n", "\n", "## Usage\n", "\n", "## OpenAI Compatible API" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Specify the `--grammar-backend`, `--reasoning-parser` option." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import openai\n", "import os\n", "\n", "from sglang.test.doc_patch import launch_server_cmd\n", "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", "\n", "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n", "\n", "\n", "server_process, port = launch_server_cmd(\n", " \"python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-7B --host 0.0.0.0 --reasoning-parser deepseek-r1 --log-level warning\"\n", ")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=server_process)\n", "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### JSON\n", "\n", "you can directly define a JSON schema or use [Pydantic](https://docs.pydantic.dev/latest/) to define and validate the response." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Using Pydantic**" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from pydantic import BaseModel, Field\n", "\n", "\n", "# Define the schema using Pydantic\n", "class CapitalInfo(BaseModel):\n", " name: str = Field(..., pattern=r\"^\\w+$\", description=\"Name of the capital city\")\n", " population: int = Field(..., description=\"Population of the capital city\")\n", "\n", "\n", "response = client.chat.completions.create(\n", " model=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\",\n", " messages=[\n", " {\n", " \"role\": \"assistant\",\n", " \"content\": \"Give me the information and population of the capital of France in the JSON format.\",\n", " },\n", " ],\n", " temperature=0,\n", " max_tokens=2048,\n", " response_format={\n", " \"type\": \"json_schema\",\n", " \"json_schema\": {\n", " \"name\": \"foo\",\n", " # convert the pydantic model to json schema\n", " \"schema\": CapitalInfo.model_json_schema(),\n", " },\n", " },\n", ")\n", "\n", "print_highlight(\n", " f\"reasoing_content: {response.choices[0].message.reasoning_content}\\n\\ncontent: {response.choices[0].message.content}\"\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**JSON Schema Directly**\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\n", "\n", "json_schema = json.dumps(\n", " {\n", " \"type\": \"object\",\n", " \"properties\": {\n", " \"name\": {\"type\": \"string\", \"pattern\": \"^[\\\\w]+$\"},\n", " \"population\": {\"type\": \"integer\"},\n", " },\n", " \"required\": [\"name\", \"population\"],\n", " }\n", ")\n", "\n", "response = client.chat.completions.create(\n", " model=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\",\n", " messages=[\n", " {\n", " \"role\": \"assistant\",\n", " \"content\": \"Give me the information and population of the capital of France in the JSON format.\",\n", " },\n", " ],\n", " temperature=0,\n", " max_tokens=2048,\n", " response_format={\n", " \"type\": \"json_schema\",\n", " \"json_schema\": {\"name\": \"foo\", \"schema\": json.loads(json_schema)},\n", " },\n", ")\n", "\n", "print_highlight(\n", " f\"reasoing_content: {response.choices[0].message.reasoning_content}\\n\\ncontent: {response.choices[0].message.content}\"\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### EBNF" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ebnf_grammar = \"\"\"\n", "root ::= city | description\n", "city ::= \"London\" | \"Paris\" | \"Berlin\" | \"Rome\"\n", "description ::= city \" is \" status\n", "status ::= \"the capital of \" country\n", "country ::= \"England\" | \"France\" | \"Germany\" | \"Italy\"\n", "\"\"\"\n", "\n", "response = client.chat.completions.create(\n", " model=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\",\n", " messages=[\n", " {\"role\": \"system\", \"content\": \"You are a helpful geography bot.\"},\n", " {\n", " \"role\": \"assistant\",\n", " \"content\": \"Give me the information and population of the capital of France in the JSON format.\",\n", " },\n", " ],\n", " temperature=0,\n", " max_tokens=2048,\n", " extra_body={\"ebnf\": ebnf_grammar},\n", ")\n", "\n", "print_highlight(\n", " f\"reasoing_content: {response.choices[0].message.reasoning_content}\\n\\ncontent: {response.choices[0].message.content}\"\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Regular expression" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = client.chat.completions.create(\n", " model=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\",\n", " messages=[\n", " {\"role\": \"assistant\", \"content\": \"What is the capital of France?\"},\n", " ],\n", " temperature=0,\n", " max_tokens=2048,\n", " extra_body={\"regex\": \"(Paris|London)\"},\n", ")\n", "\n", "print_highlight(\n", " f\"reasoing_content: {response.choices[0].message.reasoning_content}\\n\\ncontent: {response.choices[0].message.content}\"\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Structural Tag" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "tool_get_current_weather = {\n", " \"type\": \"function\",\n", " \"function\": {\n", " \"name\": \"get_current_weather\",\n", " \"description\": \"Get the current weather in a given location\",\n", " \"parameters\": {\n", " \"type\": \"object\",\n", " \"properties\": {\n", " \"city\": {\n", " \"type\": \"string\",\n", " \"description\": \"The city to find the weather for, e.g. 'San Francisco'\",\n", " },\n", " \"state\": {\n", " \"type\": \"string\",\n", " \"description\": \"the two-letter abbreviation for the state that the city is\"\n", " \" in, e.g. 'CA' which would mean 'California'\",\n", " },\n", " \"unit\": {\n", " \"type\": \"string\",\n", " \"description\": \"The unit to fetch the temperature in\",\n", " \"enum\": [\"celsius\", \"fahrenheit\"],\n", " },\n", " },\n", " \"required\": [\"city\", \"state\", \"unit\"],\n", " },\n", " },\n", "}\n", "\n", "tool_get_current_date = {\n", " \"type\": \"function\",\n", " \"function\": {\n", " \"name\": \"get_current_date\",\n", " \"description\": \"Get the current date and time for a given timezone\",\n", " \"parameters\": {\n", " \"type\": \"object\",\n", " \"properties\": {\n", " \"timezone\": {\n", " \"type\": \"string\",\n", " \"description\": \"The timezone to fetch the current date and time for, e.g. 'America/New_York'\",\n", " }\n", " },\n", " \"required\": [\"timezone\"],\n", " },\n", " },\n", "}\n", "\n", "schema_get_current_weather = tool_get_current_weather[\"function\"][\"parameters\"]\n", "schema_get_current_date = tool_get_current_date[\"function\"][\"parameters\"]\n", "\n", "\n", "def get_messages():\n", " return [\n", " {\n", " \"role\": \"system\",\n", " \"content\": f\"\"\"\n", "# Tool Instructions\n", "- Always execute python code in messages that you share.\n", "- When looking for real time information use relevant functions if available else fallback to brave_search\n", "You have access to the following functions:\n", "Use the function 'get_current_weather' to: Get the current weather in a given location\n", "{tool_get_current_weather[\"function\"]}\n", "Use the function 'get_current_date' to: Get the current date and time for a given timezone\n", "{tool_get_current_date[\"function\"]}\n", "If a you choose to call a function ONLY reply in the following format:\n", "<{{start_tag}}={{function_name}}>{{parameters}}{{end_tag}}\n", "where\n", "start_tag => ` a JSON dict with the function argument name as key and function argument value as value.\n", "end_tag => ``\n", "Here is an example,\n", "{{\"example_name\": \"example_value\"}}\n", "Reminder:\n", "- Function calls MUST follow the specified format\n", "- Required parameters MUST be specified\n", "- Only call one function at a time\n", "- Put the entire function call reply on one line\n", "- Always add your sources when using search results to answer the user query\n", "You are a helpful assistant.\"\"\",\n", " },\n", " {\n", " \"role\": \"assistant\",\n", " \"content\": \"You are in New York. Please get the current date and time, and the weather.\",\n", " },\n", " ]\n", "\n", "\n", "messages = get_messages()\n", "\n", "response = client.chat.completions.create(\n", " model=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\",\n", " messages=messages,\n", " response_format={\n", " \"type\": \"structural_tag\",\n", " \"max_new_tokens\": 2048,\n", " \"structures\": [\n", " {\n", " \"begin\": \"\",\n", " \"schema\": schema_get_current_weather,\n", " \"end\": \"\",\n", " },\n", " {\n", " \"begin\": \"\",\n", " \"schema\": schema_get_current_date,\n", " \"end\": \"\",\n", " },\n", " ],\n", " \"triggers\": [\" Note: For native API, as a work-around, you need to set `require_reasoning` argument to `True` to ensure the model will think before generating the structured output. It's not required for chat-completion API." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### JSON" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Using Pydantic**" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import requests\n", "from pydantic import BaseModel, Field\n", "from transformers import AutoTokenizer\n", "\n", "tokenizer = AutoTokenizer.from_pretrained(\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\")\n", "\n", "\n", "# Define the schema using Pydantic\n", "class CapitalInfo(BaseModel):\n", " name: str = Field(..., pattern=r\"^\\w+$\", description=\"Name of the capital city\")\n", " population: int = Field(..., description=\"Population of the capital city\")\n", "\n", "\n", "messages = [\n", " {\n", " \"role\": \"assistant\",\n", " \"content\": \"Give me the information and population of the capital of France in the JSON format.\",\n", " },\n", "]\n", "text = tokenizer.apply_chat_template(\n", " messages, tokenize=False, add_generation_prompt=True, return_dict=False\n", ")\n", "# Make API request\n", "response = requests.post(\n", " f\"http://localhost:{port}/generate\",\n", " json={\n", " \"text\": text,\n", " \"require_reasoning\": True,\n", " \"sampling_params\": {\n", " \"temperature\": 0,\n", " \"max_new_tokens\": 2048,\n", " \"json_schema\": json.dumps(CapitalInfo.model_json_schema()),\n", " },\n", " },\n", ")\n", "print(response.json())\n", "\n", "\n", "reasoing_content = response.json()[\"text\"].split(\"\")[0]\n", "content = response.json()[\"text\"].split(\"\")[1]\n", "print_highlight(f\"reasoing_content: {reasoing_content}\\n\\ncontent: {content}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**JSON Schema Directly**" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "json_schema = json.dumps(\n", " {\n", " \"type\": \"object\",\n", " \"properties\": {\n", " \"name\": {\"type\": \"string\", \"pattern\": \"^[\\\\w]+$\"},\n", " \"population\": {\"type\": \"integer\"},\n", " },\n", " \"required\": [\"name\", \"population\"],\n", " }\n", ")\n", "\n", "# JSON\n", "text = tokenizer.apply_chat_template(\n", " messages, tokenize=False, add_generation_prompt=True, return_dict=False\n", ")\n", "response = requests.post(\n", " f\"http://localhost:{port}/generate\",\n", " json={\n", " \"text\": text,\n", " \"require_reasoning\": True,\n", " \"sampling_params\": {\n", " \"temperature\": 0,\n", " \"max_new_tokens\": 2048,\n", " \"json_schema\": json_schema,\n", " },\n", " },\n", ")\n", "\n", "print_highlight(response.json())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### EBNF" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = requests.post(\n", " f\"http://localhost:{port}/generate\",\n", " json={\n", " \"text\": \"Give me the information of the capital of France.\",\n", " \"require_reasoning\": True,\n", " \"sampling_params\": {\n", " \"max_new_tokens\": 2048,\n", " \"temperature\": 0,\n", " \"n\": 3,\n", " \"ebnf\": (\n", " \"root ::= city | description\\n\"\n", " 'city ::= \"London\" | \"Paris\" | \"Berlin\" | \"Rome\"\\n'\n", " 'description ::= city \" is \" status\\n'\n", " 'status ::= \"the capital of \" country\\n'\n", " 'country ::= \"England\" | \"France\" | \"Germany\" | \"Italy\"'\n", " ),\n", " },\n", " \"stream\": False,\n", " \"return_logprob\": False,\n", " },\n", ")\n", "\n", "print(response.json())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Regular expression" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = requests.post(\n", " f\"http://localhost:{port}/generate\",\n", " json={\n", " \"text\": \"Paris is the capital of\",\n", " \"require_reasoning\": True,\n", " \"sampling_params\": {\n", " \"temperature\": 0,\n", " \"max_new_tokens\": 2048,\n", " \"regex\": \"(France|England)\",\n", " },\n", " },\n", ")\n", "print(response.json())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Structural Tag" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "text = tokenizer.apply_chat_template(\n", " messages, tokenize=False, add_generation_prompt=True, return_dict=False\n", ")\n", "payload = {\n", " \"text\": text,\n", " \"require_reasoning\": True,\n", " \"sampling_params\": {\n", " \"max_new_tokens\": 2048,\n", " \"structural_tag\": json.dumps(\n", " {\n", " \"type\": \"structural_tag\",\n", " \"structures\": [\n", " {\n", " \"begin\": \"\",\n", " \"schema\": schema_get_current_weather,\n", " \"end\": \"\",\n", " },\n", " {\n", " \"begin\": \"\",\n", " \"schema\": schema_get_current_date,\n", " \"end\": \"\",\n", " },\n", " ],\n", " \"triggers\": [\"\",\n", " \"schema\": schema_get_current_weather,\n", " \"end\": \"\",\n", " },\n", " {\n", " \"begin\": \"\",\n", " \"schema\": schema_get_current_date,\n", " \"end\": \"\",\n", " },\n", " ],\n", " \"triggers\": [\" is not trimmed.\n", "\n", "sampling_params = {\n", " \"max_new_tokens\": 1024,\n", " \"temperature\": 0,\n", " \"top_p\": 0.95,\n", " \"skip_special_tokens\": False,\n", "}\n", "\n", "# 1) Offline generation\n", "result = llm.generate(input_ids=input_ids, sampling_params=sampling_params)\n", "generated_text = result[\"text\"] # Assume there is only one prompt\n", "\n", "print_highlight(\"=== Offline Engine Output Text ===\")\n", "print_highlight(generated_text)\n", "\n", "\n", "# 2) Parse using FunctionCallParser\n", "def convert_dict_to_tool(tool_dict: dict) -> Tool:\n", " function_dict = tool_dict.get(\"function\", {})\n", " return Tool(\n", " type=tool_dict.get(\"type\", \"function\"),\n", " function=Function(\n", " name=function_dict.get(\"name\"),\n", " description=function_dict.get(\"description\"),\n", " parameters=function_dict.get(\"parameters\"),\n", " ),\n", " )\n", "\n", "\n", "tools = [convert_dict_to_tool(raw_tool) for raw_tool in tools]\n", "\n", "parser = FunctionCallParser(tools=tools, tool_call_parser=\"qwen25\")\n", "normal_text, calls = parser.parse_non_stream(generated_text)\n", "\n", "print_highlight(\"=== Parsing Result ===\")\n", "print(\"Normal text portion:\", normal_text)\n", "print_highlight(\"Function call portion:\")\n", "for call in calls:\n", " # call: ToolCallItem\n", " print_highlight(f\" - tool name: {call.name}\")\n", " print_highlight(f\" parameters: {call.parameters}\")\n", "\n", "# 3) If needed, perform additional logic on the parsed functions, such as automatically calling the corresponding function to obtain a return value, etc." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "llm.shutdown()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Tool Choice Mode\n", "\n", "SGLang supports OpenAI's `tool_choice` parameter to control when and which tools the model should call. This feature is implemented using EBNF (Extended Backus-Naur Form) grammar to ensure reliable tool calling behavior.\n", "\n", "### Supported Tool Choice Options\n", "\n", "- **`tool_choice=\"required\"`**: Forces the model to call at least one tool\n", "- **`tool_choice={\"type\": \"function\", \"function\": {\"name\": \"specific_function\"}}`**: Forces the model to call a specific function\n", "\n", "### Backend Compatibility\n", "\n", "Tool choice is fully supported with the **Xgrammar backend**, which is the default grammar backend (`--grammar-backend xgrammar`). However, it may not be fully supported with other backends such as `outlines`.\n", "\n", "### Example: Required Tool Choice" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from openai import OpenAI\n", "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", "from sglang.test.doc_patch import launch_server_cmd\n", "\n", "# Start a new server session for tool choice examples\n", "server_process_tool_choice, port_tool_choice = launch_server_cmd(\n", " \"python3 -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct --tool-call-parser qwen25 --host 0.0.0.0 --log-level warning\"\n", ")\n", "wait_for_server(\n", " f\"http://localhost:{port_tool_choice}\", process=server_process_tool_choice\n", ")\n", "\n", "# Initialize client for tool choice examples\n", "client_tool_choice = OpenAI(\n", " api_key=\"None\", base_url=f\"http://0.0.0.0:{port_tool_choice}/v1\"\n", ")\n", "model_name_tool_choice = client_tool_choice.models.list().data[0].id\n", "\n", "# Example with tool_choice=\"required\" - forces the model to call a tool\n", "messages_required = [\n", " {\"role\": \"user\", \"content\": \"Hello, what is the capital of France?\"}\n", "]\n", "\n", "# Define tools\n", "tools = [\n", " {\n", " \"type\": \"function\",\n", " \"function\": {\n", " \"name\": \"get_current_weather\",\n", " \"description\": \"Get the current weather in a given location\",\n", " \"parameters\": {\n", " \"type\": \"object\",\n", " \"properties\": {\n", " \"city\": {\n", " \"type\": \"string\",\n", " \"description\": \"The city to find the weather for, e.g. 'San Francisco'\",\n", " },\n", " \"unit\": {\n", " \"type\": \"string\",\n", " \"description\": \"The unit to fetch the temperature in\",\n", " \"enum\": [\"celsius\", \"fahrenheit\"],\n", " },\n", " },\n", " \"required\": [\"city\", \"unit\"],\n", " },\n", " },\n", " }\n", "]\n", "\n", "response_required = client_tool_choice.chat.completions.create(\n", " model=model_name_tool_choice,\n", " messages=messages_required,\n", " temperature=0,\n", " max_tokens=1024,\n", " tools=tools,\n", " tool_choice=\"required\", # Force the model to call a tool\n", ")\n", "\n", "print_highlight(\"Response with tool_choice='required':\")\n", "print(\"Content:\", response_required.choices[0].message.content)\n", "print(\"Tool calls:\", response_required.choices[0].message.tool_calls)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Example: Specific Function Choice\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Example with specific function choice - forces the model to call a specific function\n", "messages_specific = [\n", " {\"role\": \"user\", \"content\": \"What are the most attactive places in France?\"}\n", "]\n", "\n", "response_specific = client_tool_choice.chat.completions.create(\n", " model=model_name_tool_choice,\n", " messages=messages_specific,\n", " temperature=0,\n", " max_tokens=1024,\n", " tools=tools,\n", " tool_choice={\n", " \"type\": \"function\",\n", " \"function\": {\"name\": \"get_current_weather\"},\n", " }, # Force the model to call the specific get_current_weather function\n", ")\n", "\n", "print_highlight(\"Response with specific function choice:\")\n", "print(\"Content:\", response_specific.choices[0].message.content)\n", "print(\"Tool calls:\", response_specific.choices[0].message.tool_calls)\n", "\n", "if response_specific.choices[0].message.tool_calls:\n", " tool_call = response_specific.choices[0].message.tool_calls[0]\n", " print_highlight(f\"Called function: {tool_call.function.name}\")\n", " print_highlight(f\"Arguments: {tool_call.function.arguments}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process_tool_choice)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Pythonic Tool Call Format (Llama-3.2 / Llama-3.3 / Llama-4)\n", "\n", "Some Llama models (such as Llama-3.2-1B, Llama-3.2-3B, Llama-3.3-70B, and Llama-4) support a \"pythonic\" tool call format, where the model outputs function calls as Python code, e.g.:\n", "\n", "```python\n", "[get_current_weather(city=\"San Francisco\", state=\"CA\", unit=\"celsius\")]\n", "```\n", "\n", "- The output is a Python list of function calls, with arguments as Python literals (not JSON).\n", "- Multiple tool calls can be returned in the same list:\n", "```python\n", "[get_current_weather(city=\"San Francisco\", state=\"CA\", unit=\"celsius\"),\n", " get_current_weather(city=\"New York\", state=\"NY\", unit=\"fahrenheit\")]\n", "```\n", "\n", "For more information, refer to Meta’s documentation on [Zero shot function calling](https://github.com/meta-llama/llama-models/blob/main/models/llama4/prompt_format.md#zero-shot-function-calling---system-message).\n", "\n", "Note that this feature is still under development on Blackwell.\n", "\n", "### How to enable\n", "- Launch the server with `--tool-call-parser pythonic`\n", "- You may also specify --chat-template with the improved template for the model (e.g., `--chat-template=examples/chat_template/tool_chat_template_llama4_pythonic.jinja`).\n", "This is recommended because the model expects a special prompt format to reliably produce valid pythonic tool call outputs. The template ensures that the prompt structure (e.g., special tokens, message boundaries like `<|eom|>`, and function call delimiters) matches what the model was trained or fine-tuned on. If you do not use the correct chat template, tool calling may fail or produce inconsistent results.\n", "\n", "#### Forcing Pythonic Tool Call Output Without a Chat Template\n", "If you don't want to specify a chat template, you must give the model extremely explicit instructions in your messages to enforce pythonic output. For example, for `Llama-3.2-1B-Instruct`, you need:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import openai\n", "\n", "server_process, port = launch_server_cmd(\n", " \" python3 -m sglang.launch_server --model-path meta-llama/Llama-3.2-1B-Instruct --tool-call-parser pythonic --tp 1 --log-level warning\" # llama-3.2-1b-instruct\n", ")\n", "wait_for_server(f\"http://localhost:{port}\", process=server_process)\n", "\n", "tools = [\n", " {\n", " \"type\": \"function\",\n", " \"function\": {\n", " \"name\": \"get_weather\",\n", " \"description\": \"Get the current weather for a given location.\",\n", " \"parameters\": {\n", " \"type\": \"object\",\n", " \"properties\": {\n", " \"location\": {\n", " \"type\": \"string\",\n", " \"description\": \"The name of the city or location.\",\n", " }\n", " },\n", " \"required\": [\"location\"],\n", " },\n", " },\n", " },\n", " {\n", " \"type\": \"function\",\n", " \"function\": {\n", " \"name\": \"get_tourist_attractions\",\n", " \"description\": \"Get a list of top tourist attractions for a given city.\",\n", " \"parameters\": {\n", " \"type\": \"object\",\n", " \"properties\": {\n", " \"city\": {\n", " \"type\": \"string\",\n", " \"description\": \"The name of the city to find attractions for.\",\n", " }\n", " },\n", " \"required\": [\"city\"],\n", " },\n", " },\n", " },\n", "]\n", "\n", "\n", "def get_messages():\n", " return [\n", " {\n", " \"role\": \"system\",\n", " \"content\": (\n", " \"You are a travel assistant. \"\n", " \"When asked to call functions, ALWAYS respond ONLY with a python list of function calls, \"\n", " \"using this format: [func_name1(param1=value1, param2=value2), func_name2(param=value)]. \"\n", " \"Do NOT use JSON, do NOT use variables, do NOT use any other format. \"\n", " \"Here is an example:\\n\"\n", " '[get_weather(location=\"Paris\"), get_tourist_attractions(city=\"Paris\")]'\n", " ),\n", " },\n", " {\n", " \"role\": \"user\",\n", " \"content\": (\n", " \"I'm planning a trip to Tokyo next week. What's the weather like and what are some top tourist attractions? \"\n", " \"Propose parallel tool calls at once, using the python list of function calls format as shown above.\"\n", " ),\n", " },\n", " ]\n", "\n", "\n", "messages = get_messages()\n", "\n", "client = openai.Client(base_url=f\"http://localhost:{port}/v1\", api_key=\"xxxxxx\")\n", "model_name = client.models.list().data[0].id\n", "\n", "\n", "response_non_stream = client.chat.completions.create(\n", " model=model_name,\n", " messages=messages,\n", " temperature=0,\n", " top_p=0.9,\n", " stream=False, # Non-streaming\n", " tools=tools,\n", ")\n", "print_highlight(\"Non-stream response:\")\n", "print_highlight(response_non_stream)\n", "\n", "response_stream = client.chat.completions.create(\n", " model=model_name,\n", " messages=messages,\n", " temperature=0,\n", " top_p=0.9,\n", " stream=True,\n", " tools=tools,\n", ")\n", "texts = \"\"\n", "tool_calls = []\n", "name = \"\"\n", "arguments = \"\"\n", "\n", "for chunk in response_stream:\n", " if chunk.choices[0].delta.content:\n", " texts += chunk.choices[0].delta.content\n", " if chunk.choices[0].delta.tool_calls:\n", " tool_calls.append(chunk.choices[0].delta.tool_calls[0])\n", "\n", "print_highlight(\"Streaming Response:\")\n", "print_highlight(\"==== Text ====\")\n", "print_highlight(texts)\n", "\n", "print_highlight(\"==== Tool Call ====\")\n", "for tool_call in tool_calls:\n", " print_highlight(tool_call)\n", "\n", "terminate_process(server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Note:** \n", "> The model may still default to JSON if it was heavily finetuned on that format. Prompt engineering (including examples) is the only way to increase the chance of pythonic output if you are not using a chat template." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## How to support a new model?\n", "1. Update the TOOLS_TAG_LIST in sglang/srt/function_call_parser.py with the model’s tool tags. Currently supported tags include:\n", "```\n", "\tTOOLS_TAG_LIST = [\n", "\t “<|plugin|>“,\n", "\t ““,\n", "\t “<|python_tag|>“,\n", "\t “[TOOL_CALLS]”\n", "\t]\n", "```\n", "2. Create a new detector class in sglang/srt/function_call_parser.py that inherits from BaseFormatDetector. The detector should handle the model’s specific function call format. For example:\n", "```\n", " class NewModelDetector(BaseFormatDetector):\n", "```\n", "3. Add the new detector to the MultiFormatParser class that manages all the format detectors." ] } ], "metadata": { "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 4 } ================================================ FILE: docs/advanced_features/vlm_query.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# Query VLM with Offline Engine\n", "\n", "This tutorial demonstrates how to use SGLang's **offline Engine API** to query VLMs. We will demonstrate usage with Qwen2.5-VL and Llama 4. This section demonstrates three different calling approaches:\n", "\n", "1. **Basic Call**: Directly pass images and text.\n", "2. **Processor Output**: Use HuggingFace processor for data preprocessing.\n", "3. **Precomputed Embeddings**: Pre-calculate image features to improve inference efficiency." ] }, { "cell_type": "markdown", "id": "1", "metadata": {}, "source": [ "## Understanding the Three Input Formats\n", "\n", "SGLang supports three ways to pass visual data, each optimized for different scenarios:\n", "\n", "### 1. **Raw Images** - Simplest approach\n", "- Pass PIL Images, file paths, URLs, or base64 strings directly\n", "- SGLang handles all preprocessing automatically\n", "- Best for: Quick prototyping, simple applications\n", "\n", "### 2. **Processor Output** - For custom preprocessing\n", "- Pre-process images with HuggingFace processor\n", "- Pass the complete processor output dict with `format: \"processor_output\"`\n", "- Best for: Custom image transformations, integration with existing pipelines\n", "- Requirement: Must use `input_ids` instead of text prompt\n", "\n", "### 3. **Precomputed Embeddings** - For maximum performance\n", "- Pre-calculate visual embeddings using the vision encoder\n", "- Pass embeddings with `format: \"precomputed_embedding\"`\n", "- Best for: Repeated queries on same images, caching, high-throughput serving\n", "- Performance gain: Avoids redundant vision encoder computation (30-50% speedup)\n", "\n", "**Key Rule**: Within a single request, use only one format for all images. Don't mix formats.\n", "\n", "The examples below demonstrate all three approaches with both Qwen2.5-VL and Llama 4 models." ] }, { "cell_type": "markdown", "id": "2", "metadata": {}, "source": [ "## Querying Qwen2.5-VL Model" ] }, { "cell_type": "code", "execution_count": null, "id": "3", "metadata": {}, "outputs": [], "source": [ "import nest_asyncio\n", "\n", "nest_asyncio.apply()\n", "\n", "import sglang.test.doc_patch # noqa: F401\n", "\n", "model_path = \"Qwen/Qwen2.5-VL-3B-Instruct\"\n", "chat_template = \"qwen2-vl\"\n", "example_image_url = \"https://raw.githubusercontent.com/sgl-project/sglang/main/examples/assets/example_image.png\"" ] }, { "cell_type": "code", "execution_count": null, "id": "4", "metadata": {}, "outputs": [], "source": [ "from io import BytesIO\n", "import requests\n", "from PIL import Image\n", "\n", "from sglang.srt.parser.conversation import chat_templates\n", "\n", "image = Image.open(BytesIO(requests.get(example_image_url).content))\n", "\n", "conv = chat_templates[chat_template].copy()\n", "conv.append_message(conv.roles[0], f\"What's shown here: {conv.image_token}?\")\n", "conv.append_message(conv.roles[1], \"\")\n", "conv.image_data = [image]\n", "\n", "print(\"Generated prompt text:\")\n", "print(conv.get_prompt())\n", "print(f\"\\nImage size: {image.size}\")\n", "image" ] }, { "cell_type": "markdown", "id": "5", "metadata": {}, "source": [ "### Basic Offline Engine API Call" ] }, { "cell_type": "code", "execution_count": null, "id": "6", "metadata": {}, "outputs": [], "source": [ "from sglang import Engine\n", "\n", "llm = Engine(model_path=model_path, chat_template=chat_template, log_level=\"warning\")" ] }, { "cell_type": "code", "execution_count": null, "id": "7", "metadata": {}, "outputs": [], "source": [ "out = llm.generate(prompt=conv.get_prompt(), image_data=[image])\n", "print(\"Model response:\")\n", "print(out[\"text\"])" ] }, { "cell_type": "markdown", "id": "8", "metadata": {}, "source": [ "### Call with Processor Output\n", "\n", "Using a HuggingFace processor to preprocess text and images, and passing the `processor_output` directly into `Engine.generate`." ] }, { "cell_type": "code", "execution_count": null, "id": "9", "metadata": {}, "outputs": [], "source": [ "from transformers import AutoProcessor\n", "\n", "processor = AutoProcessor.from_pretrained(model_path, use_fast=True)\n", "processor_output = processor(\n", " images=[image], text=conv.get_prompt(), return_tensors=\"pt\"\n", ")\n", "\n", "out = llm.generate(\n", " input_ids=processor_output[\"input_ids\"][0].detach().cpu().tolist(),\n", " image_data=[dict(processor_output, format=\"processor_output\")],\n", ")\n", "print(\"Response using processor output:\")\n", "print(out[\"text\"])" ] }, { "cell_type": "markdown", "id": "10", "metadata": {}, "source": [ "### Call with Precomputed Embeddings\n", "\n", "You can pre-calculate image features to avoid repeated visual encoding processes." ] }, { "cell_type": "code", "execution_count": null, "id": "11", "metadata": {}, "outputs": [], "source": [ "from transformers import AutoProcessor\n", "from transformers import Qwen2_5_VLForConditionalGeneration\n", "\n", "processor = AutoProcessor.from_pretrained(model_path, use_fast=True)\n", "model = Qwen2_5_VLForConditionalGeneration.from_pretrained(model_path).eval()\n", "vision = model.model.visual.cuda()" ] }, { "cell_type": "code", "execution_count": null, "id": "12", "metadata": {}, "outputs": [], "source": [ "processor_output = processor(\n", " images=[image], text=conv.get_prompt(), return_tensors=\"pt\"\n", ")\n", "\n", "input_ids = processor_output[\"input_ids\"][0].detach().cpu().tolist()\n", "\n", "precomputed_embeddings = vision(\n", " processor_output[\"pixel_values\"].cuda(), processor_output[\"image_grid_thw\"].cuda()\n", ")\n", "precomputed_embeddings = precomputed_embeddings.pooler_output\n", "\n", "multi_modal_item = dict(\n", " processor_output,\n", " format=\"precomputed_embedding\",\n", " feature=precomputed_embeddings,\n", ")\n", "\n", "out = llm.generate(input_ids=input_ids, image_data=[multi_modal_item])\n", "print(\"Response using precomputed embeddings:\")\n", "print(out[\"text\"])\n", "\n", "llm.shutdown()" ] }, { "cell_type": "markdown", "id": "13", "metadata": {}, "source": [ "## Querying Llama 4 Vision Model\n", "\n", "```python\n", "model_path = \"meta-llama/Llama-4-Scout-17B-16E-Instruct\"\n", "chat_template = \"llama-4\"\n", "\n", "from io import BytesIO\n", "import requests\n", "from PIL import Image\n", "\n", "from sglang.srt.parser.conversation import chat_templates\n", "\n", "# Download the same example image\n", "image = Image.open(BytesIO(requests.get(example_image_url).content))\n", "\n", "conv = chat_templates[chat_template].copy()\n", "conv.append_message(conv.roles[0], f\"What's shown here: {conv.image_token}?\")\n", "conv.append_message(conv.roles[1], \"\")\n", "conv.image_data = [image]\n", "\n", "print(\"Llama 4 generated prompt text:\")\n", "print(conv.get_prompt())\n", "print(f\"Image size: {image.size}\")\n", "\n", "image\n", "```" ] }, { "cell_type": "markdown", "id": "14", "metadata": {}, "source": [ "### Llama 4 Basic Call\n", "\n", "Llama 4 requires more computational resources, so it's configured with multi-GPU parallelism (tp_size=4) and larger context length.\n", "\n", "```python\n", "llm = Engine(\n", " model_path=model_path,\n", " enable_multimodal=True,\n", " attention_backend=\"fa3\",\n", " tp_size=4,\n", " context_length=65536,\n", ")\n", "\n", "out = llm.generate(prompt=conv.get_prompt(), image_data=[image])\n", "print(\"Llama 4 response:\")\n", "print(out[\"text\"])\n", "```" ] }, { "cell_type": "markdown", "id": "15", "metadata": {}, "source": [ "### Call with Processor Output\n", "\n", "Using HuggingFace processor to preprocess data can reduce computational overhead during inference.\n", "\n", "```python\n", "from transformers import AutoProcessor\n", "\n", "processor = AutoProcessor.from_pretrained(model_path, use_fast=True)\n", "processor_output = processor(\n", " images=[image], text=conv.get_prompt(), return_tensors=\"pt\"\n", ")\n", "\n", "out = llm.generate(\n", " input_ids=processor_output[\"input_ids\"][0].detach().cpu().tolist(),\n", " image_data=[dict(processor_output, format=\"processor_output\")],\n", ")\n", "print(\"Response using processor output:\")\n", "print(out)\n", "```" ] }, { "cell_type": "markdown", "id": "16", "metadata": {}, "source": [ "### Call with Precomputed Embeddings\n", "\n", "```python\n", "from transformers import AutoProcessor\n", "from transformers import Llama4ForConditionalGeneration\n", "\n", "processor = AutoProcessor.from_pretrained(model_path, use_fast=True)\n", "model = Llama4ForConditionalGeneration.from_pretrained(\n", " model_path, torch_dtype=\"auto\"\n", ").eval()\n", "\n", "vision = model.vision_model.cuda()\n", "multi_modal_projector = model.multi_modal_projector.cuda()\n", "\n", "print(f'Image pixel values shape: {processor_output[\"pixel_values\"].shape}')\n", "input_ids = processor_output[\"input_ids\"][0].detach().cpu().tolist()\n", "\n", "# Process image through vision encoder\n", "image_outputs = vision(\n", " processor_output[\"pixel_values\"].to(\"cuda\"), \n", " aspect_ratio_ids=processor_output[\"aspect_ratio_ids\"].to(\"cuda\"),\n", " aspect_ratio_mask=processor_output[\"aspect_ratio_mask\"].to(\"cuda\"),\n", " output_hidden_states=False\n", ")\n", "image_features = image_outputs.last_hidden_state\n", "\n", "# Flatten image features and pass through multimodal projector\n", "vision_flat = image_features.view(-1, image_features.size(-1))\n", "precomputed_embeddings = multi_modal_projector(vision_flat)\n", "\n", "# Build precomputed embedding data item\n", "mm_item = dict(\n", " processor_output, \n", " format=\"precomputed_embedding\", \n", " feature=precomputed_embeddings\n", ")\n", "\n", "# Use precomputed embeddings for efficient inference\n", "out = llm.generate(input_ids=input_ids, image_data=[mm_item])\n", "print(\"Llama 4 precomputed embedding response:\")\n", "print(out[\"text\"])\n", "```" ] } ], "metadata": { "jupytext": { "cell_metadata_filter": "-all", "custom_cell_magics": "kql", "encoding": "# -*- coding: utf-8 -*-", "text_representation": { "extension": ".py", "format_name": "light", "format_version": "1.5", "jupytext_version": "1.16.1" } }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 5 } ================================================ FILE: docs/basic_usage/deepseek_ocr.md ================================================ # DeepSeek OCR (OCR-1 / OCR-2) DeepSeek OCR models are multimodal (image + text) models for OCR and document understanding. ## Launch server ```shell python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-OCR-2 \ --trust-remote-code \ --host 0.0.0.0 \ --port 30000 ``` > You can replace `deepseek-ai/DeepSeek-OCR-2` with `deepseek-ai/DeepSeek-OCR`. ## Prompt examples Recommended prompts from the model card: ``` <|grounding|>Convert the document to markdown. ``` ``` Free OCR. ``` ## OpenAI-compatible request example ```python import requests url = "http://localhost:30000/v1/chat/completions" data = { "model": "deepseek-ai/DeepSeek-OCR-2", "messages": [ { "role": "user", "content": [ {"type": "text", "text": "\n<|grounding|>Convert the document to markdown."}, {"type": "image_url", "image_url": {"url": "https://example.com/your_image.jpg"}}, ], } ], "max_tokens": 512, } response = requests.post(url, json=data) print(response.text) ``` ================================================ FILE: docs/basic_usage/deepseek_v3.md ================================================ # DeepSeek V3/V3.1/R1 Usage SGLang provides many optimizations specifically designed for the DeepSeek models, making it the inference engine recommended by the official [DeepSeek team](https://github.com/deepseek-ai/DeepSeek-V3/tree/main?tab=readme-ov-file#62-inference-with-sglang-recommended) from Day 0. This document outlines current optimizations for DeepSeek. For an overview of the implemented features see the completed [Roadmap](https://github.com/sgl-project/sglang/issues/2591). ## Launch DeepSeek V3.1/V3/R1 with SGLang To run DeepSeek V3.1/V3/R1 models, the recommended settings are as follows: | Weight Type | Configuration | |------------|-------------------| | **Full precision [FP8](https://huggingface.co/deepseek-ai/DeepSeek-R1-0528)**
*(recommended)* | 8 x H200 | | | 8 x B200 | | | 8 x MI300X | | | 2 x 8 x H100/800/20 | | | Xeon 6980P CPU | | **Full precision ([BF16](https://huggingface.co/unsloth/DeepSeek-R1-0528-BF16))** (upcast from original FP8) | 2 x 8 x H200 | | | 2 x 8 x MI300X | | | 4 x 8 x H100/800/20 | | | 4 x 8 x A100/A800 | | **Quantized weights ([INT8](https://huggingface.co/meituan/DeepSeek-R1-Channel-INT8))** | 16 x A100/800 | | | 32 x L40S | | | Xeon 6980P CPU | | | 4 x Atlas 800I A3 | | **Quantized weights ([W4A8](https://huggingface.co/novita/Deepseek-R1-0528-W4AFP8))** | 8 x H20/100, 4 x H200 | | **Quantized weights ([AWQ](https://huggingface.co/QuixiAI/DeepSeek-R1-0528-AWQ))** | 8 x H100/800/20 | | | 8 x A100/A800 | | **Quantized weights ([MXFP4](https://huggingface.co/amd/DeepSeek-R1-MXFP4-Preview))** | 8, 4 x MI355X/350X | | **Quantized weights ([NVFP4](https://huggingface.co/nvidia/DeepSeek-R1-0528-NVFP4-v2))** | 8, 4 x B200 | ```{important} The official DeepSeek V3 is already in FP8 format, so you should not run it with any quantization arguments like `--quantization fp8`. ``` Detailed commands for reference: - [8 x H200](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#using-docker-recommended) - [4 x B200, 8 x B200](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-one-b200-node) - [8 x MI300X](../platforms/amd_gpu.md#running-deepseek-v3) - [2 x 8 x H200](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-two-h2008-nodes-and-docker) - [4 x 8 x A100](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-four-a1008-nodes) - [8 x A100 (AWQ)](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-8-a100a800-with-awq-quantization) - [16 x A100 (INT8)](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-16-a100a800-with-int8-quantization) - [32 x L40S (INT8)](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-32-l40s-with-int8-quantization) - [Xeon 6980P CPU](../platforms/cpu_server.md#example-running-deepseek-r1) - [4 x Atlas 800I A3 (int8)](../platforms/ascend_npu_deepseek_example.md#running-deepseek-with-pd-disaggregation-on-4-x-atlas-800i-a3) ### Download Weights If you encounter errors when starting the server, ensure the weights have finished downloading. It's recommended to download them beforehand or restart multiple times until all weights are downloaded. Please refer to [DeepSeek V3](https://huggingface.co/deepseek-ai/DeepSeek-V3-Base#61-inference-with-deepseek-infer-demo-example-only) official guide to download the weights. ### Launch with one node of 8 x H200 Please refer to [the example](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#installation--launch). ### Running examples on Multi-Node - [Deploying DeepSeek on GB200 NVL72 with PD and Large Scale EP](https://lmsys.org/blog/2025-06-16-gb200-part-1/) ([Part I](https://lmsys.org/blog/2025-06-16-gb200-part-1/), [Part II](https://lmsys.org/blog/2025-09-25-gb200-part-2/)) - Comprehensive guide on GB200 optimizations. - [Deploying DeepSeek with PD Disaggregation and Large-Scale Expert Parallelism on 96 H100 GPUs](https://lmsys.org/blog/2025-05-05-deepseek-pd-ep/) - Guide on PD disaggregation and large-scale EP. - [Serving with two H20*8 nodes](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-two-h208-nodes). - [Best Practices for Serving DeepSeek-R1 on H20](https://lmsys.org/blog/2025-09-26-sglang-ant-group/) - Comprehensive guide on H20 optimizations, deployment and performance. - [Serving with two H200*8 nodes and docker](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-two-h2008-nodes-and-docker). - [Serving with four A100*8 nodes](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-four-a1008-nodes). ## Optimizations ### Multi-head Latent Attention (MLA) Throughput Optimizations **Description**: [MLA](https://arxiv.org/pdf/2405.04434) is an innovative attention mechanism introduced by the DeepSeek team, aimed at improving inference efficiency. SGLang has implemented specific optimizations for this, including: - **Weight Absorption**: By applying the associative law of matrix multiplication to reorder computation steps, this method balances computation and memory access and improves efficiency in the decoding phase. - **MLA Attention Backends**: Currently SGLang supports different optimized MLA attention backends, including [FlashAttention3](https://github.com/Dao-AILab/flash-attention), [Flashinfer](https://docs.flashinfer.ai/api/attention.html#flashinfer-mla), [FlashMLA](https://github.com/deepseek-ai/FlashMLA), [CutlassMLA](https://github.com/sgl-project/sglang/pull/5390), **TRTLLM MLA** (optimized for Blackwell architecture), and [Triton](https://github.com/triton-lang/triton) backends. The default FA3 provides good performance across wide workloads. - **FP8 Quantization**: W8A8 FP8 and KV Cache FP8 quantization enables efficient FP8 inference. Additionally, we have implemented Batched Matrix Multiplication (BMM) operator to facilitate FP8 inference in MLA with weight absorption. - **CUDA Graph & Torch.compile**: Both MLA and Mixture of Experts (MoE) are compatible with CUDA Graph and Torch.compile, which reduces latency and accelerates decoding speed for small batch sizes. - **Chunked Prefix Cache**: Chunked prefix cache optimization can increase throughput by cutting prefix cache into chunks, processing them with multi-head attention and merging their states. Its improvement can be significant when doing chunked prefill on long sequences. Currently this optimization is only available for FlashAttention3 backend. Overall, with these optimizations, we have achieved up to **7x** acceleration in output throughput compared to the previous version.

Multi-head Latent Attention for DeepSeek Series Models

**Usage**: MLA optimization is enabled by default. **Reference**: Check [Blog](https://lmsys.org/blog/2024-09-04-sglang-v0-3/#deepseek-multi-head-latent-attention-mla-throughput-optimizations) and [Slides](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/lmsys_1st_meetup_deepseek_mla.pdf) for more details. ### Data Parallelism Attention **Description**: This optimization involves data parallelism (DP) for the MLA attention mechanism of DeepSeek Series Models, which allows for a significant reduction in the KV cache size, enabling larger batch sizes. Each DP worker independently handles different types of batches (prefill, decode, idle), which are then synchronized before and after processing through the Mixture-of-Experts (MoE) layer. If you do not use DP attention, KV cache will be duplicated among all TP ranks.

Data Parallelism Attention for DeepSeek Series Models

With data parallelism attention enabled, we have achieved up to **1.9x** decoding throughput improvement compared to the previous version.

Data Parallelism Attention Performance Comparison

**Usage**: - Append `--enable-dp-attention --tp 8 --dp 8` to the server arguments when using 8 H200 GPUs. This optimization improves peak throughput in high batch size scenarios where the server is limited by KV cache capacity. - DP and TP attention can be flexibly combined. For example, to deploy DeepSeek-V3/R1 on 2 nodes with 8 H100 GPUs each, you can specify `--enable-dp-attention --tp 16 --dp 2`. This configuration runs attention with 2 DP groups, each containing 8 TP GPUs. ```{caution} Data parallelism attention is not recommended for low-latency, small-batch use cases. It is optimized for high-throughput scenarios with large batch sizes. ``` **Reference**: Check [Blog](https://lmsys.org/blog/2024-12-04-sglang-v0-4/#data-parallelism-attention-for-deepseek-models). ### Multi-Node Tensor Parallelism **Description**: For users with limited memory on a single node, SGLang supports serving DeepSeek Series Models, including DeepSeek V3, across multiple nodes using tensor parallelism. This approach partitions the model parameters across multiple GPUs or nodes to handle models that are too large for one node's memory. **Usage**: Check [here](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-two-h2008-nodes-and-docker) for usage examples. ### Block-wise FP8 **Description**: SGLang implements block-wise FP8 quantization with two key optimizations: - **Activation**: E4M3 format using per-token-per-128-channel sub-vector scales with online casting. - **Weight**: Per-128x128-block quantization for better numerical stability. - **DeepGEMM**: The [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) kernel library optimized for FP8 matrix multiplications. **Usage**: The activation and weight optimization above are turned on by default for DeepSeek V3 models. DeepGEMM is enabled by default on NVIDIA Hopper/Blackwell GPUs and disabled by default on other devices. DeepGEMM can also be manually turned off by setting the environment variable `SGLANG_ENABLE_JIT_DEEPGEMM=0`. ```{tip} Before serving the DeepSeek model, precompile the DeepGEMM kernels to improve first-run performance. The precompilation process typically takes around 10 minutes to complete. ``` ```bash python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code ``` ### Multi-token Prediction **Description**: SGLang implements DeepSeek V3 Multi-Token Prediction (MTP) based on [EAGLE speculative decoding](https://docs.sglang.io/advanced_features/speculative_decoding.html#EAGLE-Decoding). With this optimization, the decoding speed can be improved by **1.8x** for batch size 1 and **1.5x** for batch size 32 respectively on H200 TP8 setting. **Usage**: Add `--speculative-algorithm EAGLE`. Other flags, like `--speculative-num-steps`, `--speculative-eagle-topk` and `--speculative-num-draft-tokens` are optional. For example: ``` python3 -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3-0324 \ --speculative-algorithm EAGLE \ --trust-remote-code \ --tp 8 ``` - The default configuration for DeepSeek models is `--speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4`. The best configuration for `--speculative-num-steps`, `--speculative-eagle-topk` and `--speculative-num-draft-tokens` can be searched with [bench_speculative.py](https://github.com/sgl-project/sglang/blob/main/scripts/playground/bench_speculative.py) script for given batch size. The minimum configuration is `--speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2`, which can achieve speedup for larger batch sizes. - Most MLA attention backends fully support MTP usage. See [MLA Backends](../advanced_features/attention_backend.md#mla-backends) for details. ```{note} To enable DeepSeek MTP for large batch sizes (>48), you need to adjust some parameters (Reference [this discussion](https://github.com/sgl-project/sglang/issues/4543#issuecomment-2737413756)): - Adjust `--max-running-requests` to a larger number. The default value is `48` for MTP. For larger batch sizes, you should increase this value beyond the default value. - Set `--cuda-graph-bs`. It's a list of batch sizes for cuda graph capture. The [default captured batch sizes for speculative decoding](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py#L888-L895) is 48. You can customize this by including more batch sizes. ``` ```{tip} To enable the experimental overlap scheduler for EAGLE speculative decoding, set the environment variable `SGLANG_ENABLE_SPEC_V2=1`. This can improve performance by enabling overlap scheduling between draft and verification stages. ``` ### Reasoning Content for DeepSeek R1 & V3.1 See [Reasoning Parser](https://docs.sglang.io/advanced_features/separate_reasoning.html) and [Thinking Parameter for DeepSeek V3.1](https://docs.sglang.io/basic_usage/openai_api_completions.html#Example:-DeepSeek-V3-Models). ### Function calling for DeepSeek Models Add arguments `--tool-call-parser deepseekv3` and `--chat-template ./examples/chat_template/tool_chat_template_deepseekv3.jinja`(recommended) to enable this feature. For example (running on 1 * H20 node): ``` python3 -m sglang.launch_server \ --model deepseek-ai/DeepSeek-V3-0324 \ --tp 8 \ --port 30000 \ --host 0.0.0.0 \ --mem-fraction-static 0.9 \ --tool-call-parser deepseekv3 \ --chat-template ./examples/chat_template/tool_chat_template_deepseekv3.jinja ``` Sample Request: ``` curl "http://127.0.0.1:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ -d '{"temperature": 0, "max_tokens": 100, "model": "deepseek-ai/DeepSeek-V3-0324", "tools": [{"type": "function", "function": {"name": "query_weather", "description": "Get weather of a city, the user should supply a city first", "parameters": {"type": "object", "properties": {"city": {"type": "string", "description": "The city, e.g. Beijing"}}, "required": ["city"]}}}], "messages": [{"role": "user", "content": "How'\''s the weather like in Qingdao today"}]}' ``` Expected Response ``` {"id":"6501ef8e2d874006bf555bc80cddc7c5","object":"chat.completion","created":1745993638,"model":"deepseek-ai/DeepSeek-V3-0324","choices":[{"index":0,"message":{"role":"assistant","content":null,"reasoning_content":null,"tool_calls":[{"id":"0","index":null,"type":"function","function":{"name":"query_weather","arguments":"{\"city\": \"Qingdao\"}"}}]},"logprobs":null,"finish_reason":"tool_calls","matched_stop":null}],"usage":{"prompt_tokens":116,"total_tokens":138,"completion_tokens":22,"prompt_tokens_details":null}} ``` Sample Streaming Request: ``` curl "http://127.0.0.1:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ -d '{"temperature": 0, "max_tokens": 100, "model": "deepseek-ai/DeepSeek-V3-0324","stream":true,"tools": [{"type": "function", "function": {"name": "query_weather", "description": "Get weather of a city, the user should supply a city first", "parameters": {"type": "object", "properties": {"city": {"type": "string", "description": "The city, e.g. Beijing"}}, "required": ["city"]}}}], "messages": [{"role": "user", "content": "How'\''s the weather like in Qingdao today"}]}' ``` Expected Streamed Chunks (simplified for clarity): ``` data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"{\""}}]}}]} data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"city"}}]}}]} data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\":\""}}]}}]} data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"Q"}}]}}]} data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"ing"}}]}}]} data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"dao"}}]}}]} data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\"}"}}]}}]} data: {"choices":[{"delta":{"tool_calls":null}}], "finish_reason": "tool_calls"} data: [DONE] ``` The client needs to concatenate all arguments fragments to reconstruct the complete tool call: ``` {"city": "Qingdao"} ``` ```{important} 1. Use a lower `"temperature"` value for better results. 2. To receive more consistent tool call results, it is recommended to use `--chat-template examples/chat_template/tool_chat_template_deepseekv3.jinja`. It provides an improved unified prompt. ``` ### Thinking Budget for DeepSeek R1 In SGLang, we can implement thinking budget with `CustomLogitProcessor`. Launch a server with `--enable-custom-logit-processor` flag on. ``` python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-R1 --tp 8 --port 30000 --host 0.0.0.0 --mem-fraction-static 0.9 --disable-cuda-graph --reasoning-parser deepseek-r1 --enable-custom-logit-processor ``` Sample Request: ```python import openai from rich.pretty import pprint from sglang.srt.sampling.custom_logit_processor import DeepSeekR1ThinkingBudgetLogitProcessor client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="*") response = client.chat.completions.create( model="deepseek-ai/DeepSeek-R1", messages=[ { "role": "user", "content": "Question: Is Paris the Capital of France?", } ], max_tokens=1024, extra_body={ "custom_logit_processor": DeepSeekR1ThinkingBudgetLogitProcessor().to_str(), "custom_params": { "thinking_budget": 512, }, }, ) pprint(response) ``` ## FAQ **Q: Model loading is taking too long, and I'm encountering an NCCL timeout. What should I do?** A: If you're experiencing extended model loading times and an NCCL timeout, you can try increasing the timeout duration. Add the argument `--dist-timeout 3600` when launching your model. This will set the timeout to one hour, which often resolves the issue. ================================================ FILE: docs/basic_usage/deepseek_v32.md ================================================ # DeepSeek V3.2 Usage DeepSeek-V3.2 model family equips DeepSeek-V3.1-Terminus with DeepSeek Sparse Attention (DSA) through continued training. With DSA, a fine-grained sparse attention mechanism powered by a lightning indexer, DeepSeek-V3.2 achieves efficiency improvements in long-context scenarios. For reporting issues or tracking upcoming features, please refer to this [Roadmap](https://github.com/sgl-project/sglang/issues/11060). Note: This document is originally written for the usage of [DeepSeek-V3.2-Exp](https://huggingface.co/deepseek-ai/DeepSeek-V3.2-Exp) model. The usage of [DeepSeek-V3.2](https://huggingface.co/deepseek-ai/DeepSeek-V3.2) or [DeepSeek-V3.2-Speciale](https://huggingface.co/deepseek-ai/DeepSeek-V3.2-Speciale) is the same as DeepSeek-V3.2-Exp except for the tool call parser. ## Installation ### Docker ```bash # H200/B200 docker pull lmsysorg/sglang:latest # MI350/MI355 docker pull lmsysorg/sglang:v0.5.8-rocm700-mi35x # MI300 # v0.5.8-rocm700-mi30x does not include PR #17504. Prefer the newest MI30x ROCm # image tag from Docker Hub when available, or build from source (below). docker pull lmsysorg/sglang:v0.5.8-rocm700-mi30x # NPUs docker pull lmsysorg/sglang:dsv32-a2 docker pull lmsysorg/sglang:dsv32-a3 ``` ### Build From Source ```bash # Install SGLang git clone https://github.com/sgl-project/sglang cd sglang pip3 install pip --upgrade pip3 install -e "python" ``` ## Launch DeepSeek V3.2 with SGLang To serve [DeepSeek-V3.2-Exp](https://huggingface.co/deepseek-ai/DeepSeek-V3.2-Exp) on 8xH200/B200 GPUs: ```bash # Launch with TP + DP (Recommended) python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --dp 8 --enable-dp-attention # Launch with EP + DP python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --ep 8 --dp 8 --enable-dp-attention # Launch with Pure TP python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 # Launch with TP on MI30x/MI35x python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --nsa-prefill-backend tilelang --nsa-decode-backend tilelang ``` ### Configuration Tips - **DP Attention (Recommended)**: For DeepSeek V3.2 model, the kernels are customized for the use case of `dp_size=8`, so DP attention (`--dp 8 --enable-dp-attention`) is the recommended configuration for better stability and performance. All test cases use this configuration by default. - **Pure TP Mode**: Launching with pure TP (without `--dp` and `--enable-dp-attention`) is also supported. Note that this mode has not been fully validated in PD disaggregation scenarios. - **Short-sequence MHA prefill (adaptive)**: For short prefill sequences (default threshold: **2048 tokens**), the NSA backend uses standard MHA automatically (no extra flags). On H200 (SM90) this path uses the FlashAttention variable-length kernel; on B200 (SM100) it uses TRT-LLM ragged MHA. MHA uses `MHA_ONE_SHOT` for best performance. `MHA_ONE_SHOT` computes multi-head attention over all tokens (both cached prefix and newly extended tokens) in a single kernel invocation, avoiding the overhead of chunked KV cache processing. This achieves optimal throughput for short sequences where total sequence length fits within the chunk capacity limit. - **Choices of Attention Kernels**: The attention backend is automatically set to `nsa` attention backend for DeepSeek V3.2 model. In this backend, different kernels for sparse prefilling/decoding are implemented, which can be specified by `--nsa-prefill-backend` and `--nsa-decode-backend` server arguments. The choices of nsa prefill/decode attention kernels include: - `flashmla_sparse`: `flash_mla_sparse_fwd` kernel from `flash_mla` library. Can run on both Hopper and Blackwell GPUs. It requires bf16 q, kv inputs. - `flashmla_kv`: `flash_mla_with_kvcache` kernel from `flash_mla` library. Can run on both Hopper and Blackwell GPUs. It requires bf16 q, fp8 k_cache inputs. - `fa3`: `flash_attn_with_kvcache` kernel from `flash_attn` library. Can only run on Hopper GPUs. It requires bf16 q, kv inputs. - `tilelang`: `tilelang` implementation that can run on GPU, HPU and NPU. - `aiter`: Aiter kernel on AMD HPUs. Can only be used as decode kernel. - `trtllm`: `trtllm-mla` sparse kernel from flashinfer library. Only run on blackwell GPUs. It requires QKV bf16 or QKV fp8. - On the basis of performance benchmarks, the default configuration on H200 and B200 are set as follows : - H200: `flashmla_sparse` prefill attention (short-seq prefill uses MHA via FlashAttention varlen), `fa3` decode attention, `bf16` kv cache dtype. - B200: `flashmla_auto` prefill attention (short-seq prefill uses MHA via TRT-LLM ragged), `flashmla_kv` decode attention, `fp8_e4m3` kv cache dtype. `flashmla_auto` enables automatic selection of either `flashmla_sparse` or `flashmla_kv` kernel for prefill based on KV cache dtype, hardware, and heuristics. When FP8 KV cache is enabled and `total_kv_tokens < total_q_tokens * 512`, it uses the `flashmla_sparse` kernel; otherwise, it falls back to the `flashmla_kv` kernel. The heuristics may need to be tuned if the performance of either the `flashmla_sparse` or `flashmla_kv` kernel changes significantly. - On Blackwell platform, with slightly accuracy drop, the performance can boost up to 3x-5x - B200: by choosing `trtllm` for both `--nsa-prefill-backend` and `--nsa-decode-backend`, the prefill attention use MHA via TRT-LLM ragged for both short and long sequence (**accuracy impact**). Combine the `trtllm` with `fp8_e4m3` kv cache, the kv cache dim is `576` (kv_lora_rank + qk_rope_head_dim) (**accuracy impact**), compare to the combination of `flashmla_auto` and `fp8_e4m` kv cache dim is `656` (kv_lora_rank + scale storage (kv_lora_rank // quant_block_size * 4 bytes) + rope dimension storage). ## Multi-token Prediction SGLang implements Multi-Token Prediction (MTP) for DeepSeek V3.2 based on [EAGLE speculative decoding](https://docs.sglang.io/advanced_features/speculative_decoding.html#EAGLE-Decoding). With this optimization, the decoding speed can be improved significantly on small batch sizes. Please look at [this PR](https://github.com/sgl-project/sglang/pull/11652) for more information. Example usage with DP Attention: ```bash python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --dp 8 --enable-dp-attention --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 ``` Example usage with Pure TP: ```bash python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 ``` - The best configuration for `--speculative-num-steps`, `--speculative-eagle-topk` and `--speculative-num-draft-tokens` can be searched with [bench_speculative.py](https://github.com/sgl-project/sglang/blob/main/scripts/playground/bench_speculative.py) script for given batch size. The minimum configuration is `--speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2`, which can achieve speedup for larger batch sizes. - The default value of `--max-running-requests` is set to `48` for MTP. For larger batch sizes, this value should be increased beyond the default value. ```{tip} To enable the experimental overlap scheduler for EAGLE speculative decoding, set the environment variable `SGLANG_ENABLE_SPEC_V2=1`. This can improve performance by enabling overlap scheduling between draft and verification stages. ``` ## Function Calling and Reasoning Parser The usage of function calling and reasoning parser is the same as DeepSeek V3.1. Please refer to [Reasoning Parser](https://docs.sglang.io/advanced_features/separate_reasoning.html) and [Tool Parser](https://docs.sglang.io/advanced_features/tool_parser.html) documents. To launch `DeepSeek-V3.2-Exp` with function calling and reasoning parser: > Note: It is recommended to specify the chat-template, ensuring that you are within the sglang's root directory. ```bash python3 -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3.2-Exp \ --trust-remote-code \ --tp-size 8 --dp-size 8 --enable-dp-attention \ --tool-call-parser deepseekv31 \ --reasoning-parser deepseek-v3 \ --chat-template ./examples/chat_template/tool_chat_template_deepseekv32.jinja ``` To launch `DeepSeek-V3.2` with function calling and reasoning parser: ```bash python3 -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3.2 \ --trust-remote-code \ --tp-size 8 --dp-size 8 --enable-dp-attention \ --tool-call-parser deepseekv32 \ --reasoning-parser deepseek-v3 ``` `DeepSeek-V3.2-Speciale` doesn't support tool calling, so can only be launched with reasoning parser: ```bash python3 -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3.2-Speciale \ --trust-remote-code \ --tp-size 8 --dp-size 8 --enable-dp-attention \ --reasoning-parser deepseek-v3 ``` ## NVFP4 Checkpoint To launch deepseek v3.2 [NVFP4 checkpoint](https://huggingface.co/nvidia/DeepSeek-V3.2-NVFP4) on Blackwell devices, the user needs to specify the quantization method as `modelopt_fp4`, and moe runner backend as one of `flashinfer_trtllm`(recommended), `flashinfer_cutlass` and `flashinfer_cutedsl`. Any other usage (parallelism, reasoning parser, ...) is the same as FP8 checkpoint. An example launching command can be: ```bash python -m sglang.launch_server --model nvidia/DeepSeek-V3.2-NVFP4 --tp 4 --quantization modelopt_fp4 --moe-runner-backend flashinfer_trtllm --tool-call-parser deepseekv32 --reasoning-parser deepseek-v3 ``` ## PD Disaggregation Prefill Command: ```bash python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3.2-Exp \ --disaggregation-mode prefill \ --host $LOCAL_IP \ --port $PORT \ --tp 8 \ --dp 8 \ --enable-dp-attention \ --dist-init-addr ${HOST}:${DIST_PORT} \ --trust-remote-code \ --disaggregation-bootstrap-port 8998 \ --mem-fraction-static 0.9 \ ``` Decode command: ```bash python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3.2-Exp \ --disaggregation-mode decode \ --host $LOCAL_IP \ --port $PORT \ --tp 8 \ --dp 8 \ --enable-dp-attention \ --dist-init-addr ${HOST}:${DIST_PORT} \ --trust-remote-code \ --mem-fraction-static 0.9 \ ``` Router command: ```bash python -m sglang_router.launch_router --pd-disaggregation \ --prefill $PREFILL_ADDR 8998 \ --decode $DECODE_ADDR \ --host 127.0.0.1 \ --port 8000 \ ``` If you need more advanced deployment methods or production-ready deployment methods, such as RBG or LWS-based deployment, please refer to [references/multi_node_deployment/rbg_pd/deepseekv32_pd.md](../references/multi_node_deployment/rbg_pd/deepseekv32_pd.md). Additionally, you can also find startup commands for DeepEP-based EP parallelism in the aforementioned documentation. ## Benchmarking Results ### Accuracy Test with `gsm8k` A simple accuracy benchmark can be tested with `gsm8k` dataset: ```bash python3 benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1319 --parallel 1319 ``` The result is 0.956, which matches our expectation: ```bash Accuracy: 0.956 Invalid: 0.000 Latency: 25.109 s Output throughput: 5226.235 token/s ``` To test long-context accuracy, run gsm8k with `--num-shots 20`. The results are very close to the 8 shots results: ``` Accuracy: 0.956 Invalid: 0.000 Latency: 29.545 s Output throughput: 4418.617 token/s ``` ### Accuracy Test with `gpqa-diamond` Accuracy benchmark on long context can be tested on GPQA-diamond dataset with long output tokens and thinking enabled: ```bash python3 -m sglang.test.run_eval --port 30000 --eval-name gpqa --num-examples 198 --max-tokens 128000 --repeat 8 --thinking-mode deepseek-v3 ``` The mean accuracy over 8 runs shows 0.797, which matches the number 0.799 in official tech report. ```bash Repeat: 8, mean: 0.797 Scores: ['0.808', '0.798', '0.808', '0.798', '0.783', '0.788', '0.803', '0.793'] ``` For Deepseek V3.2, Deepseek recommends setting the sampling parameters to temperature = 1.0, top_p = 0.95: ```bash python3 -m sglang.test.run_eval --port 30000 --eval-name gpqa --num-examples 198 --max-tokens 128000 --repeat 8 --top-p 0.95 --temperature 1.0 --thinking-mode deepseek-v3 Repeat: 8, mean: 0.840 Scores: ['0.848', '0.808', '0.848', '0.838', '0.879', '0.813', '0.838', '0.848'] ``` which matches the official score, 0.824, as reported in the [Deepseek-V3.2 technical report](https://huggingface.co/deepseek-ai/DeepSeek-V3.2/blob/main/assets/paper.pdf). ### Accuracy Test with `aime 2025` Prepare the environment by installing NeMo-Skills in the docker or your own virtual environment: ``` pip install git+https://github.com/NVIDIA/NeMo-Skills.git --ignore-installed blinker ``` Then launch the SGLang server: ``` python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --dp 8 --enable-dp-attention ``` **For `DeepSeek-V3.2` and `DeepSeek-V3.2-Speciale`**: ``` python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3.2 --trust-remote-code --tp-size 8 --dp-size 8 --enable-dp-attention --tool-call-parser deepseekv32 --reasoning-parser deepseek-v3 ``` Run the following script to evaluate AIME 2025: ``` #! /bin/bash export NEMO_SKILLS_DISABLE_UNCOMMITTED_CHANGES_CHECK=1 ns prepare_data aime25 PORT=30000 BACKEND=sglang MODEL="deepseek-ai/DeepSeek-V3.2-Exp" # Should be changed to the model name MODEL_NAME="dsv32-fp8" echo "Starting AIME25 evaluation with model $MODEL on port $PORT using backend $BACKEND..." ns eval \ --benchmarks=aime25:4 \ --server_type=$BACKEND \ --model=$MODEL \ --server_address=http://localhost:${PORT}/v1 \ --output_dir=nemo_skills_aime25_${MODEL_NAME}_output_${BACKEND}_$(date +%Y%m%d_%H%M%S) \ ++chat_template_kwargs.thinking=true \ ++inference.temperature=1.0 \ ++inference.top_p=0.95 \ ++inference.tokens_to_generate=64000 # ++inference.tokens_to_generate=120000 for Speciale model ``` Test results (8*B200): DeepSeek-V3.2-Exp： | evaluation_mode | num_entries | avg_tokens | gen_seconds | symbolic_correct | no_answer | |--------------------|-------------|------------|-------------|-----------------------|-----------| | pass@1[avg-of-4] | 30 | 15040 | 1673 | 87.50% ± 1.67% | 0.00% | | majority@4 | 30 | 15040 | 1673 | 90.00% | 0.00% | | pass@4 | 30 | 15040 | 1673 | 90.00% | 0.00% | DeepSeek-V3.2: | evaluation_mode | num_entries | avg_tokens | gen_seconds | symbolic_correct | no_answer | |--------------------|-------------|------------|-------------|-----------------------|-----------| | pass@1[avg-of-4] | 30 | 13550 | 1632 | 92.50% ± 1.67% | 0.00% | | majority@4 | 30 | 13550 | 1632 | 94.71% | 0.00% | | pass@4 | 30 | 13550 | 1632 | 96.67% | 0.00% | DeepSeek-V3.2-Speciale: | evaluation_mode | num_entries | avg_tokens | gen_seconds | symbolic_correct | no_answer | |--------------------|-------------|------------|-------------|-----------------------|-----------| | pass@1[avg-of-4] | 30 | 24155 | 3583 | 95.00% ± 1.92% | 0.00% | | majority@4 | 30 | 24155 | 3583 | 95.83% | 0.00% | | pass@4 | 30 | 24155 | 3583 | 100.00% | 0.00% | ## DSA long sequence context parallel optimization(experimental) **Note: This feature is only verified on Hopper machines** For context parallel in DeepSeek V3.2 model, we provide two different modes of splitting tokens, which can be controlled with argument `--nsa-prefill-cp-mode`. ### In sequence splitting The first mode can be enabled by `--nsa-prefill-cp-mode in-seq-split`. This mode implements context parallel for DSA by splitting the sequence uniformly between context parallel ranks. At attention stage, each cp rank computes the indexer results of sharded sequence, and collects the whole kv cache through all gather operator. Add `attn_cp_size` for communication group for context parallel. Note that in sequence splitting mode has the following restrictions: - The batch size is restricted to 1 for prefill batches - `moe_dense_tp_size=1`, `moe_a2a_backend = "deepep"` - To ensure `cp_size > 1`, the passed in `tp_size` must be larger than `dp_size` For more details, please refer to PR https://github.com/sgl-project/sglang/pull/12065. Example: ```bash # In-seq splitting mode launched with EP + DP python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --ep 8 --dp 2 --enable-dp-attention --enable-nsa-prefill-context-parallel --attn-cp-size 4 --nsa-prefill-cp-mode in-seq-split --max-running-requests 32 ``` ### Round robin splitting (default setting) This mode can be enabled by specifying the parameter `--nsa-prefill-cp-mode round-robin-split`, which distributes tokens across ranks based on `token_idx % cp_size`. In this scenario, compared with the aforementioned method, it additionally supports the fused MoE backend (the fused MoE backend may deliver better performance than DeepEP in single-machine scenarios), FP8 KV-cache, and multi-batch prefill inference. But it cannot be enabled with dp attention together. For more details, please refer to PR https://github.com/sgl-project/sglang/pull/13959. Example usage: ```bash # Launch with FusedMoe + CP8 python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --enable-nsa-prefill-context-parallel --attn-cp-size 8 --nsa-prefill-cp-mode round-robin-split --max-running-requests 32 ``` ### Pipeline Parallel + Context Parallel (PP + CP) This mode combines Pipeline Parallelism (PP) and Context Parallelism (CP) to scale across multiple nodes, which can achieve better throughput and Time To First Token (TTFT). Note that this method has only been tested on H20 96G. #### Standard Usage To launch with PP=2 and CP (via `round-robin-split` mode) on 2 nodes. This configuration uses the fused MoE kernel by default, which generally provides better performance. For related development details, please refer to: - Fused MoE + CP support: [PR #13959](https://github.com/sgl-project/sglang/pull/13959) - PP + CP support: [Issue #15358](https://github.com/sgl-project/sglang/issues/15358) and [PR #16380](https://github.com/sgl-project/sglang/pull/16380) Node 0: ```bash export SGLANG_PP_LAYER_PARTITION=30,31 python3 -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3.2-Exp \ --nnodes 2 --node-rank 0 \ --dist-init-addr :62001 \ --tp 8 --pp-size 2 \ --dp-size 1 --moe-dense-tp-size 1 \ --enable-nsa-prefill-context-parallel \ --attn-cp-size 8 \ --nsa-prefill-cp-mode round-robin-split \ --trust-remote-code \ --disable-radix-cache \ --mem-fraction-static 0.8 \ --max-running-requests 128 \ --chunked-prefill-size 16384 \ --cuda-graph-max-bs 8 \ --page-size 64 \ --watchdog-timeout 3600 \ --host 0.0.0.0 --port 8000 \ --tool-call-parser deepseekv32 ``` Node 1: ```bash export SGLANG_PP_LAYER_PARTITION=30,31 python3 -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3.2-Exp \ --nnodes 2 --node-rank 1 \ --dist-init-addr :62001 \ --tp 8 --pp-size 2 \ --dp-size 1 --moe-dense-tp-size 1 \ --enable-nsa-prefill-context-parallel \ --attn-cp-size 8 \ --nsa-prefill-cp-mode round-robin-split \ --trust-remote-code \ --disable-radix-cache \ --mem-fraction-static 0.8 \ --max-running-requests 128 \ --chunked-prefill-size 16384 \ --cuda-graph-max-bs 8 \ --page-size 64 \ --watchdog-timeout 3600 \ --host 0.0.0.0 --port 8000 \ --tool-call-parser deepseekv32 ``` #### PD Disaggregation with PP + CP If using PD (Prefill-Decode) Disaggregation, the Prefill nodes can be configured with PP + CP as follows. Prefill Node 0: ```bash python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3.2-Exp \ --served-model-name deepseek-v32 \ --nnodes 2 --node-rank 0 \ --dist-init-addr :20102 \ --tp 8 --pp-size 2 \ --dp-size 1 --moe-dense-tp-size 1 \ --enable-nsa-prefill-context-parallel \ --attn-cp-size 8 \ --nsa-prefill-cp-mode round-robin-split \ --disaggregation-ib-device mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 \ --trust-remote-code \ --disable-radix-cache \ --max-running-requests 512 \ --chunked-prefill-size 4096 \ --context-length 131072 \ --mem-fraction-static 0.9 \ --page-size 64 \ --enable-metrics \ --collect-tokens-histogram \ --tokenizer-worker-num 8 \ --host 0.0.0.0 --port 30000 ``` Prefill Node 1: ```bash python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3.2-Exp \ --served-model-name deepseek-v32-prefill \ --nnodes 2 --node-rank 1 \ --dist-init-addr :20102 \ --tp 8 --pp-size 2 \ --dp-size 1 --moe-dense-tp-size 1 \ --enable-nsa-prefill-context-parallel \ --attn-cp-size 8 \ --nsa-prefill-cp-mode round-robin-split \ --disaggregation-ib-device mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 \ --trust-remote-code \ --disable-radix-cache \ --max-running-requests 512 \ --chunked-prefill-size 4096 \ --context-length 131072 \ --mem-fraction-static 0.9 \ --page-size 64 \ --enable-metrics \ --collect-tokens-histogram \ --tokenizer-worker-num 8 \ --host 0.0.0.0 --port 30000 ``` For the Decode nodes, it is recommended to use the **EP mode**. ================================================ FILE: docs/basic_usage/glm45.md ================================================ ## Launch GLM-4.5 / GLM-4.6 / GLM-4.7 with SGLang To serve GLM-4.5 / GLM-4.6 FP8 models on 8xH100/H200 GPUs: ```bash python3 -m sglang.launch_server --model zai-org/GLM-4.6-FP8 --tp 8 ``` ### EAGLE Speculative Decoding **Description**: SGLang has supported GLM-4.5 / GLM-4.6 models with [EAGLE speculative decoding](https://docs.sglang.io/advanced_features/speculative_decoding.html#EAGLE-Decoding). **Usage**: Add arguments `--speculative-algorithm`, `--speculative-num-steps`, `--speculative-eagle-topk` and `--speculative-num-draft-tokens` to enable this feature. For example: ``` bash python3 -m sglang.launch_server \ --model-path zai-org/GLM-4.6-FP8 \ --tp-size 8 \ --tool-call-parser glm45 \ --reasoning-parser glm45 \ --speculative-algorithm EAGLE \ --speculative-num-steps 3 \ --speculative-eagle-topk 1 \ --speculative-num-draft-tokens 4 \ --mem-fraction-static 0.9 \ --served-model-name glm-4.6-fp8 \ --enable-custom-logit-processor ``` ```{tip} To enable the experimental overlap scheduler for EAGLE speculative decoding, set the environment variable `SGLANG_ENABLE_SPEC_V2=1`. This can improve performance by enabling overlap scheduling between draft and verification stages. ``` ### Thinking Budget for GLM-4.5 / GLM-4.6 **Note**: For GLM-4.7, `--tool-call-parser` should be set to `glm47`, for GLM-4.5 and GLM-4.6, it should be set to `glm45`. In SGLang, we can implement thinking budget with `CustomLogitProcessor`. Launch a server with `--enable-custom-logit-processor` flag on. Sample Request: ```python import openai from rich.pretty import pprint from sglang.srt.sampling.custom_logit_processor import Glm4MoeThinkingBudgetLogitProcessor client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="*") response = client.chat.completions.create( model="zai-org/GLM-4.6", messages=[ { "role": "user", "content": "Question: Is Paris the Capital of France?", } ], max_tokens=1024, extra_body={ "custom_logit_processor": Glm4MoeThinkingBudgetLogitProcessor().to_str(), "custom_params": { "thinking_budget": 512, }, }, ) pprint(response) ``` ================================================ FILE: docs/basic_usage/glmv.md ================================================ # GLM-4.6V / GLM-4.5V Usage ## Launch commands for SGLang Below are suggested launch commands tailored for different hardware / precision modes ### FP8 (quantised) mode For high memory-efficiency and latency optimized deployments (e.g., on H100, H200) where FP8 checkpoint is supported: ```bash python3 -m sglang.launch_server \ --model-path zai-org/GLM-4.6V-FP8 \ --tp 2 \ --ep 2 \ --host 0.0.0.0 \ --port 30000 \ --keep-mm-feature-on-device ``` ### Non-FP8 (BF16 / full precision) mode For deployments on A100/H100 where BF16 is used (or FP8 snapshot not used): ```bash python3 -m sglang.launch_server \ --model-path zai-org/GLM-4.6V \ --tp 4 \ --ep 4 \ --host 0.0.0.0 \ --port 30000 ``` ## Hardware-specific notes / recommendations - On H100 with FP8: Use the FP8 checkpoint for best memory efficiency. - On A100 / H100 with BF16 (non-FP8): It’s recommended to use `--mm-max-concurrent-calls` to control parallel throughput and GPU memory usage during image/video inference. - On H200 & B200: The model can be run “out of the box”, supporting full context length plus concurrent image + video processing. ## Sending Image/Video Requests ### Image input: ```python import requests url = f"http://localhost:30000/v1/chat/completions" data = { "model": "zai-org/GLM-4.6V", "messages": [ { "role": "user", "content": [ {"type": "text", "text": "What’s in this image?"}, { "type": "image_url", "image_url": { "url": "https://github.com/sgl-project/sglang/blob/main/examples/assets/example_image.png?raw=true" }, }, ], } ], "max_tokens": 300, } response = requests.post(url, json=data) print(response.text) ``` ### Video Input: ```python import requests url = f"http://localhost:30000/v1/chat/completions" data = { "model": "zai-org/GLM-4.6V", "messages": [ { "role": "user", "content": [ {"type": "text", "text": "What’s happening in this video?"}, { "type": "video_url", "video_url": { "url": "https://github.com/sgl-project/sgl-test-files/raw/refs/heads/main/videos/jobs_presenting_ipod.mp4" }, }, ], } ], "max_tokens": 300, } response = requests.post(url, json=data) print(response.text) ``` ## Important Server Parameters and Flags When launching the model server for **multimodal support**, you can use the following command-line arguments to fine-tune performance and behavior: - `--mm-attention-backend`: Specify multimodal attention backend. Eg. `fa3`(Flash Attention 3) - `--mm-max-concurrent-calls `: Specifies the **maximum number of concurrent asynchronous multimodal data processing calls** allowed on the server. Use this to control parallel throughput and GPU memory usage during image/video inference. - `--mm-per-request-timeout `: Defines the **timeout duration (in seconds)** for each multimodal request. If a request exceeds this time limit (e.g., for very large video inputs), it will be automatically terminated. - `--keep-mm-feature-on-device`: Instructs the server to **retain multimodal feature tensors on the GPU** after processing. This avoids device-to-host (D2H) memory copies and improves performance for repeated or high-frequency inference workloads. - `--mm-enable-dp-encoder`: Placing the ViT in data parallel while keeping the LLM in tensor parallel consistently lowers TTFT and boosts end-to-end throughput. - `SGLANG_USE_CUDA_IPC_TRANSPORT=1`: Shared memory pool based CUDA IPC for multi-modal data transport. For significantly improving e2e latency. ### Example usage with the above optimizations: ```bash SGLANG_USE_CUDA_IPC_TRANSPORT=1 \ SGLANG_VLM_CACHE_SIZE_MB=0 \ python -m sglang.launch_server \ --model-path zai-org/GLM-4.6V \ --host 0.0.0.0 \ --port 30000 \ --trust-remote-code \ --tp-size 8 \ --enable-cache-report \ --log-level info \ --max-running-requests 64 \ --mem-fraction-static 0.65 \ --chunked-prefill-size 8192 \ --attention-backend fa3 \ --mm-attention-backend fa3 \ --mm-enable-dp-encoder \ --enable-metrics ``` ### Thinking Budget for GLM-4.5V / GLM-4.6V In SGLang, we can implement thinking budget with `CustomLogitProcessor`. Launch a server with the `--enable-custom-logit-processor` flag. Then, use `Glm4MoeThinkingBudgetLogitProcessor` in the request, similar to the `GLM-4.6` example in [glm45.md](./glm45.md). ================================================ FILE: docs/basic_usage/gpt_oss.md ================================================ # GPT OSS Usage Please refer to [https://github.com/sgl-project/sglang/issues/8833](https://github.com/sgl-project/sglang/issues/8833). ## Responses API & Built-in Tools ### Responses API GPT‑OSS is compatible with the OpenAI Responses API. Use `client.responses.create(...)` with `model`, `instructions`, `input`, and optional `tools` to enable built‑in tool use. You can set reasoning level via `instructions`, e.g., "Reasoning: high" (also supports "medium" and "low") — levels: low (fast), medium (balanced), high (deep). ### Built-in Tools GPT‑OSS can call built‑in tools for web search and Python execution. You can use the demo tool server or connect to external MCP tool servers. #### Python Tool - Executes short Python snippets for calculations, parsing, and quick scripts. - By default runs in a Docker-based sandbox. To run on the host, set `PYTHON_EXECUTION_BACKEND=UV` (this executes model-generated code locally; use with care). - Ensure Docker is available if you are not using the UV backend. It is recommended to run `docker pull python:3.11` in advance. #### Web Search Tool - Uses the Exa backend for web search. - Requires an Exa API key; set `EXA_API_KEY` in your environment. Create a key at `https://exa.ai`. ### Tool & Reasoning Parser - We support OpenAI Reasoning and Tool Call parser, as well as our SGLang native api for tool call and reasoning. Refer to [reasoning parser](../advanced_features/separate_reasoning.ipynb) and [tool call parser](../advanced_features/function_calling.ipynb) for more details. ## Notes - Use **Python 3.12** for the demo tools. And install the required `gpt-oss` packages. - The default demo integrates the web search tool (Exa backend) and a demo Python interpreter via Docker. - For search, set `EXA_API_KEY`. For Python execution, either have Docker available or set `PYTHON_EXECUTION_BACKEND=UV`. Examples: ```bash export EXA_API_KEY=YOUR_EXA_KEY # Optional: run Python tool locally instead of Docker (use with care) export PYTHON_EXECUTION_BACKEND=UV ``` Launch the server with the demo tool server: ```bash python3 -m sglang.launch_server \ --model-path openai/gpt-oss-120b \ --tool-server demo \ --tp 2 ``` For production usage, sglang can act as an MCP client for multiple services. An [example tool server](https://github.com/openai/gpt-oss/tree/main/gpt-oss-mcp-server) is provided. Start the servers and point sglang to them: ```bash mcp run -t sse browser_server.py:mcp mcp run -t sse python_server.py:mcp python -m sglang.launch_server ... --tool-server ip-1:port-1,ip-2:port-2 ``` The URLs should be MCP SSE servers that expose server information and well-documented tools. These tools are added to the system prompt so the model can use them. ## Speculative Decoding SGLang supports speculative decoding for GPT-OSS models using EAGLE3 algorithm. This can significantly improve decoding speed, especially for small batch sizes. **Usage**: Add `--speculative-algorithm EAGLE3` along with the draft model path. ```bash python3 -m sglang.launch_server \ --model-path openai/gpt-oss-120b \ --speculative-algorithm EAGLE3 \ --speculative-draft-model-path lmsys/EAGLE3-gpt-oss-120b-bf16 \ --tp 2 ``` ```{tip} To enable the experimental overlap scheduler for EAGLE3 speculative decoding, set the environment variable `SGLANG_ENABLE_SPEC_V2=1`. This can improve performance by enabling overlap scheduling between draft and verification stages. ``` ### Quick Demo ```python from openai import OpenAI client = OpenAI( base_url="http://localhost:30000/v1", api_key="sk-123456" ) tools = [ {"type": "code_interpreter"}, {"type": "web_search_preview"}, ] # Reasoning level example response = client.responses.create( model="openai/gpt-oss-120b", instructions="You are a helpful assistant." reasoning_effort="high" # Supports high, medium, or low input="In one sentence, explain the transformer architecture.", ) print("====== reasoning: high ======") print(response.output_text) # Test python tool response = client.responses.create( model="openai/gpt-oss-120b", instructions="You are a helpful assistant, you could use python tool to execute code.", input="Use python tool to calculate the sum of 29138749187 and 29138749187", # 58,277,498,374 tools=tools ) print("====== test python tool ======") print(response.output_text) # Test browser tool response = client.responses.create( model="openai/gpt-oss-120b", instructions="You are a helpful assistant, you could use browser to search the web", input="Search the web for the latest news about Nvidia stock price", tools=tools ) print("====== test browser tool ======") print(response.output_text) ``` Example output: ``` ====== test python tool ====== The sum of 29,138,749,187 and 29,138,749,187 is **58,277,498,374**. ====== test browser tool ====== **Recent headlines on Nvidia (NVDA) stock** | Date (2025) | Source | Key news points | Stock‑price detail | |-------------|--------|----------------|--------------------| | **May 13** | Reuters | The market data page shows Nvidia trading “higher” at **$116.61** with no change from the previous close. | **$116.61** – latest trade (delayed ≈ 15 min)【14†L34-L38】 | | **Aug 18** | CNBC | Morgan Stanley kept an **overweight** rating and lifted its price target to **$206** (up from $200), implying a 14 % upside from the Friday close. The firm notes Nvidia shares have already **jumped 34 % this year**. | No exact price quoted, but the article signals strong upside expectations【9†L27-L31】 | | **Aug 20** | The Motley Fool | Nvidia is set to release its Q2 earnings on Aug 27. The article lists the **current price of $175.36**, down 0.16 % on the day (as of 3:58 p.m. ET). | **$175.36** – current price on Aug 20【10†L12-L15】【10†L53-L57】 | **What the news tells us** * Nvidia’s share price has risen sharply this year – up roughly a third according to Morgan Stanley – and analysts are still raising targets (now $206). * The most recent market quote (Reuters, May 13) was **$116.61**, but the stock has surged since then, reaching **$175.36** by mid‑August. * Upcoming earnings on **Aug 27** are a focal point; both the Motley Fool and Morgan Stanley expect the results could keep the rally going. **Bottom line:** Nvidia’s stock is on a strong upward trajectory in 2025, with price targets climbing toward $200‑$210 and the market price already near $175 as of late August. ``` ================================================ FILE: docs/basic_usage/llama4.md ================================================ # Llama4 Usage [Llama 4](https://github.com/meta-llama/llama-models/blob/main/models/llama4/MODEL_CARD.md) is Meta's latest generation of open-source LLM model with industry-leading performance. SGLang has supported Llama 4 Scout (109B) and Llama 4 Maverick (400B) since [v0.4.5](https://github.com/sgl-project/sglang/releases/tag/v0.4.5). Ongoing optimizations are tracked in the [Roadmap](https://github.com/sgl-project/sglang/issues/5118). ## Launch Llama 4 with SGLang To serve Llama 4 models on 8xH100/H200 GPUs: ```bash python3 -m sglang.launch_server \ --model-path meta-llama/Llama-4-Scout-17B-16E-Instruct \ --tp 8 \ --context-length 1000000 ``` ### Configuration Tips - **OOM Mitigation**: Adjust `--context-length` to avoid a GPU out-of-memory issue. For the Scout model, we recommend setting this value up to 1M on 8\*H100 and up to 2.5M on 8\*H200. For the Maverick model, we don't need to set context length on 8\*H200. When hybrid kv cache is enabled, `--context-length` can be set up to 5M on 8\*H100 and up to 10M on 8\*H200 for the Scout model. - **Attention Backend Auto-Selection**: SGLang automatically selects the optimal attention backend for Llama 4 based on your hardware. You typically don't need to specify `--attention-backend` manually: - **Blackwell GPUs (B200/GB200)**: `trtllm_mha` - **Hopper GPUs (H100/H200)**: `fa3` - **AMD GPUs**: `aiter` - **Intel XPU**: `intel_xpu` - **Other platforms**: `triton` (fallback) To override the auto-selection, explicitly specify `--attention-backend` with one of the supported backends: `fa3`, `aiter`, `triton`, `trtllm_mha`, or `intel_xpu`. - **Chat Template**: Add `--chat-template llama-4` for chat completion tasks. - **Enable Multi-Modal**: Add `--enable-multimodal` for multi-modal capabilities. - **Enable Hybrid-KVCache**: Set `--swa-full-tokens-ratio` to adjust the ratio of SWA layer (for Llama4, it's local attention layer) KV tokens / full layer KV tokens. (default: 0.8, range: 0-1) ### EAGLE Speculative Decoding **Description**: SGLang has supported Llama 4 Maverick (400B) with [EAGLE speculative decoding](https://docs.sglang.io/advanced_features/speculative_decoding.html#EAGLE-Decoding). **Usage**: Add arguments `--speculative-draft-model-path`, `--speculative-algorithm`, `--speculative-num-steps`, `--speculative-eagle-topk` and `--speculative-num-draft-tokens` to enable this feature. For example: ``` python3 -m sglang.launch_server \ --model-path meta-llama/Llama-4-Maverick-17B-128E-Instruct \ --speculative-algorithm EAGLE3 \ --speculative-draft-model-path nvidia/Llama-4-Maverick-17B-128E-Eagle3 \ --speculative-num-steps 3 \ --speculative-eagle-topk 1 \ --speculative-num-draft-tokens 4 \ --trust-remote-code \ --tp 8 \ --context-length 1000000 ``` - **Note** The Llama 4 draft model *nvidia/Llama-4-Maverick-17B-128E-Eagle3* can only recognize conversations in chat mode. ## Benchmarking Results ### Accuracy Test with `lm_eval` The accuracy on SGLang for both Llama4 Scout and Llama4 Maverick can match the [official benchmark numbers](https://ai.meta.com/blog/llama-4-multimodal-intelligence/). Benchmark results on MMLU Pro dataset with 8*H100: | | Llama-4-Scout-17B-16E-Instruct | Llama-4-Maverick-17B-128E-Instruct | |--------------------|--------------------------------|-------------------------------------| | Official Benchmark | 74.3 | 80.5 | | SGLang | 75.2 | 80.7 | Commands: ```bash # Llama-4-Scout-17B-16E-Instruct model python -m sglang.launch_server \ --model-path meta-llama/Llama-4-Scout-17B-16E-Instruct \ --port 30000 \ --tp 8 \ --mem-fraction-static 0.8 \ --context-length 65536 lm_eval --model local-chat-completions --model_args model=meta-llama/Llama-4-Scout-17B-16E-Instruct,base_url=http://localhost:30000/v1/chat/completions,num_concurrent=128,timeout=999999,max_gen_toks=2048 --tasks mmlu_pro --batch_size 128 --apply_chat_template --num_fewshot 0 # Llama-4-Maverick-17B-128E-Instruct python -m sglang.launch_server \ --model-path meta-llama/Llama-4-Maverick-17B-128E-Instruct \ --port 30000 \ --tp 8 \ --mem-fraction-static 0.8 \ --context-length 65536 lm_eval --model local-chat-completions --model_args model=meta-llama/Llama-4-Maverick-17B-128E-Instruct,base_url=http://localhost:30000/v1/chat/completions,num_concurrent=128,timeout=999999,max_gen_toks=2048 --tasks mmlu_pro --batch_size 128 --apply_chat_template --num_fewshot 0 ``` Details can be seen in [this PR](https://github.com/sgl-project/sglang/pull/5092). ================================================ FILE: docs/basic_usage/minimax_m2.md ================================================ # MiniMax M2.5/M2.1/M2 Usage [MiniMax-M2.5](https://huggingface.co/MiniMaxAI/MiniMax-M2.5), [MiniMax-M2.1](https://huggingface.co/MiniMaxAI/MiniMax-M2.1), and [MiniMax-M2](https://huggingface.co/MiniMaxAI/MiniMax-M2) are advanced large language models created by [MiniMax](https://www.minimax.io/). The MiniMax-M2 series redefines efficiency for agents. These compact, fast, and cost-effective MoE models (230 billion total parameters with 10 billion active parameters) are built for elite performance in coding and agentic tasks, all while maintaining powerful general intelligence. With just 10 billion activated parameters, the MiniMax-M2 series provides sophisticated, end-to-end tool use performance expected from today's leading models, but in a streamlined form factor that makes deployment and scaling easier than ever. ## Supported Models This guide applies to the following models. You only need to update the model name during deployment. The following examples use **MiniMax-M2**: - [MiniMaxAI/MiniMax-M2.5](https://huggingface.co/MiniMaxAI/MiniMax-M2.5) - [MiniMaxAI/MiniMax-M2.1](https://huggingface.co/MiniMaxAI/MiniMax-M2.1) - [MiniMaxAI/MiniMax-M2](https://huggingface.co/MiniMaxAI/MiniMax-M2) ## System Requirements The following are recommended configurations; actual requirements should be adjusted based on your use case: - 4x 96GB GPUs: Supported context length of up to 400K tokens. - 8x 144GB GPUs: Supported context length of up to 3M tokens. ## Deployment with Python 4-GPU deployment command: ```bash python -m sglang.launch_server \ --model-path MiniMaxAI/MiniMax-M2 \ --tp-size 4 \ --tool-call-parser minimax-m2 \ --reasoning-parser minimax-append-think \ --host 0.0.0.0 \ --trust-remote-code \ --port 8000 \ --mem-fraction-static 0.85 ``` 8-GPU deployment command: ```bash python -m sglang.launch_server \ --model-path MiniMaxAI/MiniMax-M2 \ --tp-size 8 \ --ep-size 8 \ --tool-call-parser minimax-m2 \ --reasoning-parser minimax-append-think \ --host 0.0.0.0 \ --trust-remote-code \ --port 8000 \ --mem-fraction-static 0.85 ``` ### AMD GPUs (MI300X/MI325X/MI355X) 8-GPU deployment command: ```bash SGLANG_USE_AITER=1 python -m sglang.launch_server \ --model-path MiniMaxAI/MiniMax-M2.5 \ --tp-size 8 \ --ep-size 8 \ --attention-backend aiter \ --tool-call-parser minimax-m2 \ --reasoning-parser minimax-append-think \ --host 0.0.0.0 \ --trust-remote-code \ --port 8000 \ --mem-fraction-static 0.85 ``` ## Testing Deployment After startup, you can test the SGLang OpenAI-compatible API with the following command: ```bash curl http://localhost:8000/v1/chat/completions \ -H "Content-Type: application/json" \ -d '{ "model": "MiniMaxAI/MiniMax-M2", "messages": [ {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]}, {"role": "user", "content": [{"type": "text", "text": "Who won the world series in 2020?"}]} ] }' ``` ================================================ FILE: docs/basic_usage/native_api.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# SGLang Native APIs\n", "\n", "Apart from the OpenAI compatible APIs, the SGLang Runtime also provides its native server APIs. We introduce the following APIs:\n", "\n", "- `/generate` (text generation model)\n", "- `/get_model_info`\n", "- `/get_server_info`\n", "- `/health`\n", "- `/health_generate`\n", "- `/flush_cache`\n", "- `/update_weights`\n", "- `/encode`(embedding model)\n", "- `/v1/rerank`(cross encoder rerank model)\n", "- `/v1/score`(decoder-only scoring)\n", "- `/classify`(reward model)\n", "- `/start_expert_distribution_record`\n", "- `/stop_expert_distribution_record`\n", "- `/dump_expert_distribution_record`\n", "- `/tokenize`\n", "- `/detokenize`\n", "- A full list of these APIs can be found at [http_server.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/entrypoints/http_server.py)\n", "\n", "We mainly use `requests` to test these APIs in the following examples. You can also use `curl`.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Launch A Server" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from sglang.test.doc_patch import launch_server_cmd\n", "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", "\n", "server_process, port = launch_server_cmd(\n", " \"python3 -m sglang.launch_server --model-path qwen/qwen2.5-0.5b-instruct --host 0.0.0.0 --log-level warning\"\n", ")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Generate (text generation model)\n", "Generate completions. This is similar to the `/v1/completions` in OpenAI API. Detailed parameters can be found in the [sampling parameters](sampling_params.md)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import requests\n", "\n", "url = f\"http://localhost:{port}/generate\"\n", "data = {\"text\": \"What is the capital of France?\"}\n", "\n", "response = requests.post(url, json=data)\n", "print_highlight(response.json())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Get Model Info\n", "\n", "Get the information of the model.\n", "\n", "- `model_path`: The path/name of the model.\n", "- `is_generation`: Whether the model is used as generation model or embedding model.\n", "- `tokenizer_path`: The path/name of the tokenizer.\n", "- `preferred_sampling_params`: The default sampling params specified via `--preferred-sampling-params`. `None` is returned in this example as we did not explicitly configure it in server args.\n", "- `weight_version`: This field contains the version of the model weights. This is often used to track changes or updates to the model’s trained parameters.\n", "- `has_image_understanding`: Whether the model has image-understanding capability.\n", "- `has_audio_understanding`: Whether the model has audio-understanding capability.\n", "- `model_type`: The model type from the HuggingFace config (e.g., \"qwen2\", \"llama\").\n", "- `architectures`: The model architectures from the HuggingFace config (e.g., [\"Qwen2ForCausalLM\"])." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "url = f\"http://localhost:{port}/get_model_info\"\n", "\n", "response = requests.get(url)\n", "response_json = response.json()\n", "print_highlight(response_json)\n", "assert response_json[\"model_path\"] == \"qwen/qwen2.5-0.5b-instruct\"\n", "assert response_json[\"is_generation\"] is True\n", "assert response_json[\"tokenizer_path\"] == \"qwen/qwen2.5-0.5b-instruct\"\n", "assert response_json[\"preferred_sampling_params\"] is None\n", "assert response_json.keys() == {\n", " \"model_path\",\n", " \"is_generation\",\n", " \"tokenizer_path\",\n", " \"preferred_sampling_params\",\n", " \"weight_version\",\n", " \"has_image_understanding\",\n", " \"has_audio_understanding\",\n", " \"model_type\",\n", " \"architectures\",\n", "}" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Get Server Info\n", "Gets the server information including CLI arguments, token limits, and memory pool sizes.\n", "- Note: `get_server_info` merges the following deprecated endpoints:\n", " - `get_server_args`\n", " - `get_memory_pool_size`\n", " - `get_max_total_num_tokens`" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "url = f\"http://localhost:{port}/get_server_info\"\n", "\n", "response = requests.get(url)\n", "print_highlight(response.text)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Health Check\n", "- `/health`: Check the health of the server.\n", "- `/health_generate`: Check the health of the server by generating one token." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "url = f\"http://localhost:{port}/health_generate\"\n", "\n", "response = requests.get(url)\n", "print_highlight(response.text)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "url = f\"http://localhost:{port}/health\"\n", "\n", "response = requests.get(url)\n", "print_highlight(response.text)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Flush Cache\n", "\n", "Flush the radix cache. It will be automatically triggered when the model weights are updated by the `/update_weights` API." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "url = f\"http://localhost:{port}/flush_cache\"\n", "\n", "response = requests.post(url)\n", "print_highlight(response.text)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Update Weights From Disk\n", "\n", "Update model weights from disk without restarting the server. Only applicable for models with the same architecture and parameter size.\n", "\n", "SGLang support `update_weights_from_disk` API for continuous evaluation during training (save checkpoint to disk and update weights from disk).\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# successful update with same architecture and size\n", "\n", "url = f\"http://localhost:{port}/update_weights_from_disk\"\n", "data = {\"model_path\": \"qwen/qwen2.5-0.5b-instruct\"}\n", "\n", "response = requests.post(url, json=data)\n", "print_highlight(response.text)\n", "assert response.json()[\"success\"] is True\n", "assert response.json()[\"message\"] == \"Succeeded to update model weights.\"" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# failed update with different parameter size or wrong name\n", "\n", "url = f\"http://localhost:{port}/update_weights_from_disk\"\n", "data = {\"model_path\": \"qwen/qwen2.5-0.5b-instruct-wrong\"}\n", "\n", "response = requests.post(url, json=data)\n", "response_json = response.json()\n", "print_highlight(response_json)\n", "assert response_json[\"success\"] is False\n", "assert response_json[\"message\"] == (\n", " \"Failed to get weights iterator: \"\n", " \"qwen/qwen2.5-0.5b-instruct-wrong\"\n", " \" (repository not found).\"\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Encode (embedding model)\n", "\n", "Encode text into embeddings. Note that this API is only available for [embedding models](openai_api_embeddings.ipynb) and will raise an error for generation models.\n", "Therefore, we launch a new server to server an embedding model." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "embedding_process, port = launch_server_cmd(\"\"\"\n", "python3 -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-1.5B-instruct \\\n", " --host 0.0.0.0 --is-embedding --log-level warning\n", "\"\"\")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=embedding_process)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# successful encode for embedding model\n", "\n", "url = f\"http://localhost:{port}/encode\"\n", "data = {\"model\": \"Alibaba-NLP/gte-Qwen2-1.5B-instruct\", \"text\": \"Once upon a time\"}\n", "\n", "response = requests.post(url, json=data)\n", "response_json = response.json()\n", "print_highlight(f\"Text embedding (first 10): {response_json['embedding'][:10]}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(embedding_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## v1/rerank (cross encoder rerank model)\n", "Rerank a list of documents given a query using a cross-encoder model. Note that this API is only available for cross encoder model like [BAAI/bge-reranker-v2-m3](https://huggingface.co/BAAI/bge-reranker-v2-m3) with `attention-backend` `triton` and `torch_native`.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "reranker_process, port = launch_server_cmd(\"\"\"\n", "python3 -m sglang.launch_server --model-path BAAI/bge-reranker-v2-m3 \\\n", " --host 0.0.0.0 --disable-radix-cache --chunked-prefill-size -1 --attention-backend triton --is-embedding --log-level warning\n", "\"\"\")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=reranker_process)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# compute rerank scores for query and documents\n", "\n", "url = f\"http://localhost:{port}/v1/rerank\"\n", "data = {\n", " \"model\": \"BAAI/bge-reranker-v2-m3\",\n", " \"query\": \"what is panda?\",\n", " \"documents\": [\n", " \"hi\",\n", " \"The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.\",\n", " ],\n", "}\n", "\n", "response = requests.post(url, json=data)\n", "response_json = response.json()\n", "for item in response_json:\n", " print_highlight(f\"Score: {item['score']:.2f} - Document: '{item['document']}'\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(reranker_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## v1/score (decoder-only scoring)\n", "\n", "Compute token probabilities for specified tokens given a query and items. This is useful for classification tasks, scoring responses, or computing log-probabilities.\n", "\n", "Parameters:\n", "- `query`: Query text\n", "- `items`: Item text(s) to score\n", "- `label_token_ids`: Token IDs to compute probabilities for\n", "- `apply_softmax`: Whether to apply softmax to get normalized probabilities (default: False)\n", "- `item_first`: Whether items come first in concatenation order (default: False)\n", "- `model`: Model name\n", "\n", "The response contains `scores` - a list of probability lists, one per item, each in the order of `label_token_ids`." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "score_process, port = launch_server_cmd(\"\"\"\n", "python3 -m sglang.launch_server --model-path qwen/qwen2.5-0.5b-instruct \\\n", " --host 0.0.0.0 --log-level warning\n", "\"\"\")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=score_process)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Score the probability of different completions given a query\n", "query = \"The capital of France is\"\n", "items = [\"Paris\", \"London\", \"Berlin\"]\n", "\n", "url = f\"http://localhost:{port}/v1/score\"\n", "data = {\n", " \"model\": \"qwen/qwen2.5-0.5b-instruct\",\n", " \"query\": query,\n", " \"items\": items,\n", " \"label_token_ids\": [9454, 2753], # e.g. \"Yes\" and \"No\" token ids\n", " \"apply_softmax\": True, # Normalize probabilities to sum to 1\n", "}\n", "\n", "response = requests.post(url, json=data)\n", "response_json = response.json()\n", "\n", "# Display scores for each item\n", "for item, scores in zip(items, response_json[\"scores\"]):\n", " print_highlight(f\"Item '{item}': probabilities = {[f'{s:.4f}' for s in scores]}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(score_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Classify (reward model)\n", "\n", "SGLang Runtime also supports reward models. Here we use a reward model to classify the quality of pairwise generations." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Note that SGLang now treats embedding models and reward models as the same type of models.\n", "# This will be updated in the future.\n", "\n", "reward_process, port = launch_server_cmd(\"\"\"\n", "python3 -m sglang.launch_server --model-path Skywork/Skywork-Reward-Llama-3.1-8B-v0.2 --host 0.0.0.0 --is-embedding --log-level warning\n", "\"\"\")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=reward_process)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from transformers import AutoTokenizer\n", "\n", "PROMPT = (\n", " \"What is the range of the numeric output of a sigmoid node in a neural network?\"\n", ")\n", "\n", "RESPONSE1 = \"The output of a sigmoid node is bounded between -1 and 1.\"\n", "RESPONSE2 = \"The output of a sigmoid node is bounded between 0 and 1.\"\n", "\n", "CONVS = [\n", " [{\"role\": \"user\", \"content\": PROMPT}, {\"role\": \"assistant\", \"content\": RESPONSE1}],\n", " [{\"role\": \"user\", \"content\": PROMPT}, {\"role\": \"assistant\", \"content\": RESPONSE2}],\n", "]\n", "\n", "tokenizer = AutoTokenizer.from_pretrained(\"Skywork/Skywork-Reward-Llama-3.1-8B-v0.2\")\n", "prompts = tokenizer.apply_chat_template(CONVS, tokenize=False, return_dict=False)\n", "\n", "url = f\"http://localhost:{port}/classify\"\n", "data = {\"model\": \"Skywork/Skywork-Reward-Llama-3.1-8B-v0.2\", \"text\": prompts}\n", "\n", "responses = requests.post(url, json=data).json()\n", "for response in responses:\n", " print_highlight(f\"reward: {response['embedding'][0]}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(reward_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Capture expert selection distribution in MoE models\n", "\n", "SGLang Runtime supports recording the number of times an expert is selected in a MoE model run for each expert in the model. This is useful when analyzing the throughput of the model and plan for optimization.\n", "\n", "*Note: We only print out the first 10 lines of the csv below for better readability. Please adjust accordingly if you want to analyze the results more deeply.*" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "expert_record_server_process, port = launch_server_cmd(\n", " \"python3 -m sglang.launch_server --model-path Qwen/Qwen1.5-MoE-A2.7B --host 0.0.0.0 --expert-distribution-recorder-mode stat --log-level warning\"\n", ")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=expert_record_server_process)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = requests.post(f\"http://localhost:{port}/start_expert_distribution_record\")\n", "print_highlight(response)\n", "\n", "url = f\"http://localhost:{port}/generate\"\n", "data = {\"text\": \"What is the capital of France?\"}\n", "\n", "response = requests.post(url, json=data)\n", "print_highlight(response.json())\n", "\n", "response = requests.post(f\"http://localhost:{port}/stop_expert_distribution_record\")\n", "print_highlight(response)\n", "\n", "response = requests.post(f\"http://localhost:{port}/dump_expert_distribution_record\")\n", "print_highlight(response)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(expert_record_server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Tokenize/Detokenize Example (Round Trip)\n", "\n", "This example demonstrates how to use the /tokenize and /detokenize endpoints together. We first tokenize a string, then detokenize the resulting IDs to reconstruct the original text. This workflow is useful when you need to handle tokenization externally but still leverage the server for detokenization." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "tokenizer_free_server_process, port = launch_server_cmd(\"\"\"\n", "python3 -m sglang.launch_server --model-path qwen/qwen2.5-0.5b-instruct\n", "\"\"\")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=tokenizer_free_server_process)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import requests\n", "from sglang.utils import print_highlight\n", "\n", "base_url = f\"http://localhost:{port}\"\n", "tokenize_url = f\"{base_url}/tokenize\"\n", "detokenize_url = f\"{base_url}/detokenize\"\n", "\n", "model_name = \"qwen/qwen2.5-0.5b-instruct\"\n", "input_text = \"SGLang provides efficient tokenization endpoints.\"\n", "print_highlight(f\"Original Input Text:\\n'{input_text}'\")\n", "\n", "# --- tokenize the input text ---\n", "tokenize_payload = {\n", " \"model\": model_name,\n", " \"prompt\": input_text,\n", " \"add_special_tokens\": False,\n", "}\n", "try:\n", " tokenize_response = requests.post(tokenize_url, json=tokenize_payload)\n", " tokenize_response.raise_for_status()\n", " tokenization_result = tokenize_response.json()\n", " token_ids = tokenization_result.get(\"tokens\")\n", "\n", " if not token_ids:\n", " raise ValueError(\"Tokenization returned empty tokens.\")\n", "\n", " print_highlight(f\"\\nTokenized Output (IDs):\\n{token_ids}\")\n", " print_highlight(f\"Token Count: {tokenization_result.get('count')}\")\n", " print_highlight(f\"Max Model Length: {tokenization_result.get('max_model_len')}\")\n", "\n", " # --- detokenize the obtained token IDs ---\n", " detokenize_payload = {\n", " \"model\": model_name,\n", " \"tokens\": token_ids,\n", " \"skip_special_tokens\": True,\n", " }\n", "\n", " detokenize_response = requests.post(detokenize_url, json=detokenize_payload)\n", " detokenize_response.raise_for_status()\n", " detokenization_result = detokenize_response.json()\n", " reconstructed_text = detokenization_result.get(\"text\")\n", "\n", " print_highlight(f\"\\nDetokenized Output (Text):\\n'{reconstructed_text}'\")\n", "\n", " if input_text == reconstructed_text:\n", " print_highlight(\n", " \"\\nRound Trip Successful: Original and reconstructed text match.\"\n", " )\n", " else:\n", " print_highlight(\n", " \"\\nRound Trip Mismatch: Original and reconstructed text differ.\"\n", " )\n", "\n", "except requests.exceptions.RequestException as e:\n", " print_highlight(f\"\\nHTTP Request Error: {e}\")\n", "except Exception as e:\n", " print_highlight(f\"\\nAn error occurred: {e}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(tokenizer_free_server_process)" ] } ], "metadata": { "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 4 } ================================================ FILE: docs/basic_usage/offline_engine_api.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Offline Engine API\n", "\n", "SGLang provides a direct inference engine without the need for an HTTP server, especially for use cases where additional HTTP server adds unnecessary complexity or overhead. Here are two general use cases:\n", "\n", "- Offline Batch Inference\n", "- Custom Server on Top of the Engine\n", "\n", "This document focuses on the offline batch inference, demonstrating four different inference modes:\n", "\n", "- Non-streaming synchronous generation\n", "- Streaming synchronous generation\n", "- Non-streaming asynchronous generation\n", "- Streaming asynchronous generation\n", "\n", "Additionally, you can easily build a custom server on top of the SGLang offline engine. A detailed example working in a python script can be found in [custom_server](https://github.com/sgl-project/sglang/blob/main/examples/runtime/engine/custom_server.py).\n", "\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Nest Asyncio\n", "Note that if you want to use **Offline Engine** in ipython or some other nested loop code, you need to add the following code:\n", "```python\n", "import nest_asyncio\n", "\n", "nest_asyncio.apply()\n", "\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Advanced Usage\n", "\n", "The engine supports [vlm inference](https://github.com/sgl-project/sglang/blob/main/examples/runtime/engine/offline_batch_inference_vlm.py) as well as [extracting hidden states](https://github.com/sgl-project/sglang/blob/main/examples/runtime/hidden_states). \n", "\n", "Please see [the examples](https://github.com/sgl-project/sglang/tree/main/examples/runtime/engine) for further use cases." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Offline Batch Inference\n", "\n", "SGLang offline engine supports batch inference with efficient scheduling." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# launch the offline engine\n", "import asyncio\n", "\n", "import sglang as sgl\n", "import sglang.test.doc_patch # noqa: F401\n", "from sglang.utils import async_stream_and_merge, stream_and_merge\n", "\n", "llm = sgl.Engine(model_path=\"qwen/qwen2.5-0.5b-instruct\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Non-streaming Synchronous Generation" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "prompts = [\n", " \"Hello, my name is\",\n", " \"The president of the United States is\",\n", " \"The capital of France is\",\n", " \"The future of AI is\",\n", "]\n", "\n", "sampling_params = {\"temperature\": 0.8, \"top_p\": 0.95}\n", "\n", "outputs = llm.generate(prompts, sampling_params)\n", "for prompt, output in zip(prompts, outputs):\n", " print(\"===============================\")\n", " print(f\"Prompt: {prompt}\\nGenerated text: {output['text']}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Streaming Synchronous Generation" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "prompts = [\n", " \"Write a short, neutral self-introduction for a fictional character. Hello, my name is\",\n", " \"Provide a concise factual statement about France’s capital city. The capital of France is\",\n", " \"Explain possible future trends in artificial intelligence. The future of AI is\",\n", "]\n", "\n", "sampling_params = {\n", " \"temperature\": 0.2,\n", " \"top_p\": 0.9,\n", "}\n", "\n", "print(\"\\n=== Testing synchronous streaming generation with overlap removal ===\\n\")\n", "\n", "for prompt in prompts:\n", " print(f\"Prompt: {prompt}\")\n", " merged_output = stream_and_merge(llm, prompt, sampling_params)\n", " print(\"Generated text:\", merged_output)\n", " print()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Non-streaming Asynchronous Generation" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "prompts = [\n", " \"Write a short, neutral self-introduction for a fictional character. Hello, my name is\",\n", " \"Provide a concise factual statement about France’s capital city. The capital of France is\",\n", " \"Explain possible future trends in artificial intelligence. The future of AI is\",\n", "]\n", "\n", "sampling_params = {\"temperature\": 0.8, \"top_p\": 0.95}\n", "\n", "print(\"\\n=== Testing asynchronous batch generation ===\")\n", "\n", "\n", "async def main():\n", " outputs = await llm.async_generate(prompts, sampling_params)\n", "\n", " for prompt, output in zip(prompts, outputs):\n", " print(f\"\\nPrompt: {prompt}\")\n", " print(f\"Generated text: {output['text']}\")\n", "\n", "\n", "asyncio.run(main())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Streaming Asynchronous Generation" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "prompts = [\n", " \"Write a short, neutral self-introduction for a fictional character. Hello, my name is\",\n", " \"Provide a concise factual statement about France’s capital city. The capital of France is\",\n", " \"Explain possible future trends in artificial intelligence. The future of AI is\",\n", "]\n", "\n", "sampling_params = {\"temperature\": 0.8, \"top_p\": 0.95}\n", "\n", "print(\"\\n=== Testing asynchronous streaming generation (no repeats) ===\")\n", "\n", "\n", "async def main():\n", " for prompt in prompts:\n", " print(f\"\\nPrompt: {prompt}\")\n", " print(\"Generated text: \", end=\"\", flush=True)\n", "\n", " # Replace direct calls to async_generate with our custom overlap-aware version\n", " async for cleaned_chunk in async_stream_and_merge(llm, prompt, sampling_params):\n", " print(cleaned_chunk, end=\"\", flush=True)\n", "\n", " print() # New line after each prompt\n", "\n", "\n", "asyncio.run(main())" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "llm.shutdown()" ] } ], "metadata": { "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 2 } ================================================ FILE: docs/basic_usage/ollama_api.md ================================================ # Ollama-Compatible API SGLang provides Ollama API compatibility, allowing you to use the Ollama CLI and Python library with SGLang as the inference backend. ## Prerequisites ```bash # Install the Ollama Python library (for Python client usage) pip install ollama ``` > **Note**: You don't need the Ollama server installed - SGLang acts as the backend. You only need the `ollama` CLI or Python library as the client. ## Endpoints | Endpoint | Method | Description | |----------|--------|-------------| | `/` | GET, HEAD | Health check for Ollama CLI | | `/api/tags` | GET | List available models | | `/api/chat` | POST | Chat completions (streaming & non-streaming) | | `/api/generate` | POST | Text generation (streaming & non-streaming) | | `/api/show` | POST | Model information | ## Quick Start ### 1. Launch SGLang Server ```bash python -m sglang.launch_server \ --model Qwen/Qwen2.5-1.5B-Instruct \ --port 30001 \ --host 0.0.0.0 ``` > **Note**: The model name used with `ollama run` must match exactly what you passed to `--model`. ### 2. Use Ollama CLI ```bash # List available models OLLAMA_HOST=http://localhost:30001 ollama list # Interactive chat OLLAMA_HOST=http://localhost:30001 ollama run "Qwen/Qwen2.5-1.5B-Instruct" ``` If connecting to a remote server behind a firewall: ```bash # SSH tunnel ssh -L 30001:localhost:30001 user@gpu-server -N & # Then use Ollama CLI as above OLLAMA_HOST=http://localhost:30001 ollama list ``` ### 3. Use Ollama Python Library ```python import ollama client = ollama.Client(host='http://localhost:30001') # Non-streaming response = client.chat( model='Qwen/Qwen2.5-1.5B-Instruct', messages=[{'role': 'user', 'content': 'Hello!'}] ) print(response['message']['content']) # Streaming stream = client.chat( model='Qwen/Qwen2.5-1.5B-Instruct', messages=[{'role': 'user', 'content': 'Tell me a story'}], stream=True ) for chunk in stream: print(chunk['message']['content'], end='', flush=True) ``` ## Smart Router For intelligent routing between local Ollama (fast) and remote SGLang (powerful) using an LLM judge, see the [Smart Router documentation](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/entrypoints/ollama/README.md). ## Summary | Component | Purpose | |-----------|---------| | **Ollama API** | Familiar CLI/API that developers already know | | **SGLang Backend** | High-performance inference engine | | **Smart Router** | Intelligent routing - fast local for simple tasks, powerful remote for complex tasks | ================================================ FILE: docs/basic_usage/openai_api.rst ================================================ OpenAI-Compatible APIs ====================== .. toctree:: :maxdepth: 1 openai_api_completions.ipynb openai_api_vision.ipynb openai_api_embeddings.ipynb ================================================ FILE: docs/basic_usage/openai_api_completions.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# OpenAI APIs - Completions\n", "\n", "SGLang provides OpenAI-compatible APIs to enable a smooth transition from OpenAI services to self-hosted local models.\n", "A complete reference for the API is available in the [OpenAI API Reference](https://platform.openai.com/docs/api-reference).\n", "\n", "This tutorial covers the following popular APIs:\n", "\n", "- `chat/completions`\n", "- `completions`\n", "\n", "Check out other tutorials to learn about [vision APIs](openai_api_vision.ipynb) for vision-language models and [embedding APIs](openai_api_embeddings.ipynb) for embedding models." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Launch A Server\n", "\n", "Launch the server in your terminal and wait for it to initialize." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from sglang.test.doc_patch import launch_server_cmd\n", "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", "\n", "server_process, port = launch_server_cmd(\n", " \"python3 -m sglang.launch_server --model-path qwen/qwen2.5-0.5b-instruct --host 0.0.0.0 --log-level warning\"\n", ")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=server_process)\n", "print(f\"Server started on http://localhost:{port}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Chat Completions\n", "\n", "### Usage\n", "\n", "The server fully implements the OpenAI API.\n", "It will automatically apply the chat template specified in the Hugging Face tokenizer, if one is available.\n", "You can also specify a custom chat template with `--chat-template` when launching the server." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import openai\n", "\n", "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", "\n", "response = client.chat.completions.create(\n", " model=\"qwen/qwen2.5-0.5b-instruct\",\n", " messages=[\n", " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", " ],\n", " temperature=0,\n", " max_tokens=64,\n", ")\n", "\n", "print_highlight(f\"Response: {response}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Model Thinking/Reasoning Support\n", "\n", "Some models support internal reasoning or thinking processes that can be exposed in the API response. SGLang provides unified support for various reasoning models through the `chat_template_kwargs` parameter and compatible reasoning parsers.\n", "\n", "#### Supported Models and Configuration\n", "\n", "| Model Family | Chat Template Parameter | Reasoning Parser | Notes |\n", "|--------------|------------------------|------------------|--------|\n", "| DeepSeek-R1 (R1, R1-0528, R1-Distill) | `enable_thinking` | `--reasoning-parser deepseek-r1` | Standard reasoning models |\n", "| DeepSeek-V3.1 | `thinking` | `--reasoning-parser deepseek-v3` | Hybrid model (thinking/non-thinking modes) |\n", "| Qwen3 (standard) | `enable_thinking` | `--reasoning-parser qwen3` | Hybrid model (thinking/non-thinking modes) |\n", "| Qwen3-Thinking | N/A (always enabled) | `--reasoning-parser qwen3-thinking` | Always generates reasoning |\n", "| Kimi | N/A (always enabled) | `--reasoning-parser kimi` | Kimi thinking models |\n", "| Gpt-Oss | N/A (always enabled) | `--reasoning-parser gpt-oss` | Gpt-Oss thinking models |\n", "\n", "#### Basic Usage\n", "\n", "To enable reasoning output, you need to:\n", "1. Launch the server with the appropriate reasoning parser\n", "2. Set the model-specific parameter in `chat_template_kwargs`\n", "3. Optionally use `separate_reasoning: False` to not get reasoning content separately (default to `True`)\n", "\n", "**Note for Qwen3-Thinking models:** These models always generate thinking content and do not support the `enable_thinking` parameter. Use `--reasoning-parser qwen3-thinking` or `--reasoning-parser qwen3` to parse the thinking content.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Example: Qwen3 Models\n", "\n", "```python\n", "# Launch server:\n", "# python3 -m sglang.launch_server --model Qwen/Qwen3-4B --reasoning-parser qwen3\n", "\n", "from openai import OpenAI\n", "\n", "client = OpenAI(\n", " api_key=\"EMPTY\",\n", " base_url=f\"http://127.0.0.1:30000/v1\",\n", ")\n", "\n", "model = \"Qwen/Qwen3-4B\"\n", "messages = [{\"role\": \"user\", \"content\": \"How many r's are in 'strawberry'?\"}]\n", "\n", "response = client.chat.completions.create(\n", " model=model,\n", " messages=messages,\n", " extra_body={\n", " \"chat_template_kwargs\": {\"enable_thinking\": True},\n", " \"separate_reasoning\": True\n", " }\n", ")\n", "\n", "print(\"Reasoning:\", response.choices[0].message.reasoning_content)\n", "print(\"-\"*100)\n", "print(\"Answer:\", response.choices[0].message.content)\n", "```\n", "\n", "**ExampleOutput:**\n", "```\n", "Reasoning: Okay, so the user is asking how many 'r's are in the word 'strawberry'. Let me think. First, I need to make sure I have the word spelled correctly. Strawberry... S-T-R-A-W-B-E-R-R-Y. Wait, is that right? Let me break it down.\n", "\n", "Starting with 'strawberry', let's write out the letters one by one. S, T, R, A, W, B, E, R, R, Y. Hmm, wait, that's 10 letters. Let me check again. S (1), T (2), R (3), A (4), W (5), B (6), E (7), R (8), R (9), Y (10). So the letters are S-T-R-A-W-B-E-R-R-Y. \n", "...\n", "Therefore, the answer should be three R's in 'strawberry'. But I need to make sure I'm not counting any other letters as R. Let me check again. S, T, R, A, W, B, E, R, R, Y. No other R's. So three in total. Yeah, that seems right.\n", "\n", "----------------------------------------------------------------------------------------------------\n", "Answer: The word \"strawberry\" contains **three** letters 'r'. Here's the breakdown:\n", "\n", "1. **S-T-R-A-W-B-E-R-R-Y** \n", " - The **third letter** is 'R'. \n", " - The **eighth and ninth letters** are also 'R's. \n", "\n", "Thus, the total count is **3**. \n", "\n", "**Answer:** 3.\n", "```\n", "\n", "**Note:** Setting `\"enable_thinking\": False` (or omitting it) will result in `reasoning_content` being `None`. Qwen3-Thinking models always generate reasoning content and don't support the `enable_thinking` parameter.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Logit Bias Support\n", "\n", "SGLang supports the `logit_bias` parameter for both chat completions and completions APIs. This parameter allows you to modify the likelihood of specific tokens being generated by adding bias values to their logits. The bias values can range from -100 to 100, where:\n", "\n", "- **Positive values** (0 to 100) increase the likelihood of the token being selected\n", "- **Negative values** (-100 to 0) decrease the likelihood of the token being selected\n", "- **-100** effectively prevents the token from being generated\n", "\n", "The `logit_bias` parameter accepts a dictionary where keys are token IDs (as strings) and values are the bias amounts (as floats).\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Getting Token IDs\n", "\n", "To use `logit_bias` effectively, you need to know the token IDs for the words you want to bias. Here's how to get token IDs:\n", "\n", "```python\n", "# Get tokenizer to find token IDs\n", "import tiktoken\n", "\n", "# For OpenAI models, use the appropriate encoding\n", "tokenizer = tiktoken.encoding_for_model(\"gpt-3.5-turbo\") # or your model\n", "\n", "# Get token IDs for specific words\n", "word = \"sunny\"\n", "token_ids = tokenizer.encode(word)\n", "print(f\"Token IDs for '{word}': {token_ids}\")\n", "\n", "# For SGLang models, you can access the tokenizer through the client\n", "# and get token IDs for bias\n", "```\n", "\n", "**Important:** The `logit_bias` parameter uses token IDs as string keys, not the actual words.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Example: DeepSeek-V3 Models\n", "\n", "DeepSeek-V3 models support thinking mode through the `thinking` parameter:\n", "\n", "```python\n", "# Launch server:\n", "# python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.1 --tp 8 --reasoning-parser deepseek-v3\n", "\n", "from openai import OpenAI\n", "\n", "client = OpenAI(\n", " api_key=\"EMPTY\",\n", " base_url=f\"http://127.0.0.1:30000/v1\",\n", ")\n", "\n", "model = \"deepseek-ai/DeepSeek-V3.1\"\n", "messages = [{\"role\": \"user\", \"content\": \"How many r's are in 'strawberry'?\"}]\n", "\n", "response = client.chat.completions.create(\n", " model=model,\n", " messages=messages,\n", " extra_body={\n", " \"chat_template_kwargs\": {\"thinking\": True},\n", " \"separate_reasoning\": True\n", " }\n", ")\n", "\n", "print(\"Reasoning:\", response.choices[0].message.reasoning_content)\n", "print(\"-\"*100)\n", "print(\"Answer:\", response.choices[0].message.content)\n", "```\n", "\n", "**Example Output:**\n", "```\n", "Reasoning: First, the question is: \"How many r's are in 'strawberry'?\"\n", "\n", "I need to count the number of times the letter 'r' appears in the word \"strawberry\".\n", "\n", "Let me write out the word: S-T-R-A-W-B-E-R-R-Y.\n", "\n", "Now, I'll go through each letter and count the 'r's.\n", "...\n", "So, I have three 'r's in \"strawberry\".\n", "\n", "I should double-check. The word is spelled S-T-R-A-W-B-E-R-R-Y. The letters are at positions: 3, 8, and 9 are 'r's. Yes, that's correct.\n", "\n", "Therefore, the answer should be 3.\n", "----------------------------------------------------------------------------------------------------\n", "Answer: The word \"strawberry\" contains **3** instances of the letter \"r\". Here's a breakdown for clarity:\n", "\n", "- The word is spelled: S-T-R-A-W-B-E-R-R-Y\n", "- The \"r\" appears at the 3rd, 8th, and 9th positions.\n", "```\n", "\n", "**Note:** DeepSeek-V3 models use the `thinking` parameter (not `enable_thinking`) to control reasoning output.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Example with logit_bias parameter\n", "# Note: You need to get the actual token IDs from your tokenizer\n", "# For demonstration, we'll use some example token IDs\n", "response = client.chat.completions.create(\n", " model=\"qwen/qwen2.5-0.5b-instruct\",\n", " messages=[\n", " {\"role\": \"user\", \"content\": \"Complete this sentence: The weather today is\"}\n", " ],\n", " temperature=0.7,\n", " max_tokens=20,\n", " logit_bias={\n", " \"12345\": 50, # Increase likelihood of token ID 12345\n", " \"67890\": -50, # Decrease likelihood of token ID 67890\n", " \"11111\": 25, # Slightly increase likelihood of token ID 11111\n", " },\n", ")\n", "\n", "print_highlight(f\"Response with logit bias: {response.choices[0].message.content}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Parameters\n", "\n", "The chat completions API accepts OpenAI Chat Completions API's parameters. Refer to [OpenAI Chat Completions API](https://platform.openai.com/docs/api-reference/chat/create) for more details.\n", "\n", "SGLang extends the standard API with the `extra_body` parameter, allowing for additional customization. One key option within `extra_body` is `chat_template_kwargs`, which can be used to pass arguments to the chat template processor." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = client.chat.completions.create(\n", " model=\"qwen/qwen2.5-0.5b-instruct\",\n", " messages=[\n", " {\n", " \"role\": \"system\",\n", " \"content\": \"You are a knowledgeable historian who provides concise responses.\",\n", " },\n", " {\"role\": \"user\", \"content\": \"Tell me about ancient Rome\"},\n", " {\n", " \"role\": \"assistant\",\n", " \"content\": \"Ancient Rome was a civilization centered in Italy.\",\n", " },\n", " {\"role\": \"user\", \"content\": \"What were their major achievements?\"},\n", " ],\n", " temperature=0.3, # Lower temperature for more focused responses\n", " max_tokens=128, # Reasonable length for a concise response\n", " top_p=0.95, # Slightly higher for better fluency\n", " presence_penalty=0.2, # Mild penalty to avoid repetition\n", " frequency_penalty=0.2, # Mild penalty for more natural language\n", " n=1, # Single response is usually more stable\n", " seed=42, # Keep for reproducibility\n", ")\n", "\n", "print_highlight(response.choices[0].message.content)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Streaming mode is also supported." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Logit Bias Support\n", "\n", "The completions API also supports the `logit_bias` parameter with the same functionality as described in the chat completions section above.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "stream = client.chat.completions.create(\n", " model=\"qwen/qwen2.5-0.5b-instruct\",\n", " messages=[{\"role\": \"user\", \"content\": \"Say this is a test\"}],\n", " stream=True,\n", ")\n", "for chunk in stream:\n", " if chunk.choices[0].delta.content is not None:\n", " print(chunk.choices[0].delta.content, end=\"\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Returning Routed Experts (MoE Models)\n", "\n", "For MoE models, set `return_routed_experts: true` in `extra_body` to return expert routing data. Requires `--enable-return-routed-experts` server flag. The `routed_experts` field will be returned in the `sgl_ext` object on each choice, containing base64-encoded int32 expert IDs as a flattened array with logical shape `[num_tokens, num_layers, top_k]`." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Example with logit_bias parameter for completions API\n", "# Note: You need to get the actual token IDs from your tokenizer\n", "# For demonstration, we'll use some example token IDs\n", "response = client.completions.create(\n", " model=\"qwen/qwen2.5-0.5b-instruct\",\n", " prompt=\"The best programming language for AI is\",\n", " temperature=0.7,\n", " max_tokens=20,\n", " logit_bias={\n", " \"12345\": 75, # Strongly favor token ID 12345\n", " \"67890\": -100, # Completely avoid token ID 67890\n", " \"11111\": -25, # Slightly discourage token ID 11111\n", " },\n", ")\n", "\n", "print_highlight(f\"Response with logit bias: {response.choices[0].text}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Completions\n", "\n", "### Usage\n", "Completions API is similar to Chat Completions API, but without the `messages` parameter or chat templates." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = client.completions.create(\n", " model=\"qwen/qwen2.5-0.5b-instruct\",\n", " prompt=\"List 3 countries and their capitals.\",\n", " temperature=0,\n", " max_tokens=64,\n", " n=1,\n", " stop=None,\n", ")\n", "\n", "print_highlight(f\"Response: {response}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Parameters\n", "\n", "The completions API accepts OpenAI Completions API's parameters. Refer to [OpenAI Completions API](https://platform.openai.com/docs/api-reference/completions/create) for more details.\n", "\n", "Here is an example of a detailed completions request:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = client.completions.create(\n", " model=\"qwen/qwen2.5-0.5b-instruct\",\n", " prompt=\"Write a short story about a space explorer.\",\n", " temperature=0.7, # Moderate temperature for creative writing\n", " max_tokens=150, # Longer response for a story\n", " top_p=0.9, # Balanced diversity in word choice\n", " stop=[\"\\n\\n\", \"THE END\"], # Multiple stop sequences\n", " presence_penalty=0.3, # Encourage novel elements\n", " frequency_penalty=0.3, # Reduce repetitive phrases\n", " n=1, # Generate one completion\n", " seed=123, # For reproducible results\n", ")\n", "\n", "print_highlight(f\"Response: {response}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Returning Routed Experts (MoE Models)\n", "\n", "For MoE models, set `return_routed_experts: true` in `extra_body` to return expert routing data. Requires `--enable-return-routed-experts` server flag. The `routed_experts` field will be returned in the `sgl_ext` object on each choice, containing base64-encoded int32 expert IDs as a flattened array with logical shape `[num_tokens, num_layers, top_k]`." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Structured Outputs (JSON, Regex, EBNF)\n", "\n", "For OpenAI compatible structured outputs API, refer to [Structured Outputs](../advanced_features/structured_outputs.ipynb) for more details.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Using LoRA Adapters\n", "\n", "SGLang supports LoRA (Low-Rank Adaptation) adapters with OpenAI-compatible APIs. You can specify which adapter to use directly in the `model` parameter using the `base-model:adapter-name` syntax.\n", "\n", "**Server Setup:**\n", "```bash\n", "python -m sglang.launch_server \\\n", " --model-path qwen/qwen2.5-0.5b-instruct \\\n", " --enable-lora \\\n", " --lora-paths adapter_a=/path/to/adapter_a adapter_b=/path/to/adapter_b\n", "```\n", "\n", "For more details on LoRA serving configuration, see the [LoRA documentation](../advanced_features/lora.ipynb).\n", "\n", "**API Call:**\n", "\n", "(Recommended) Use the `model:adapter` syntax to specify which adapter to use:\n", "```python\n", "response = client.chat.completions.create(\n", " model=\"qwen/qwen2.5-0.5b-instruct:adapter_a\", # ← base-model:adapter-name\n", " messages=[{\"role\": \"user\", \"content\": \"Convert to SQL: show all users\"}],\n", " max_tokens=50,\n", ")\n", "```\n", "\n", "**Backward Compatible: Using `extra_body`**\n", "\n", "The old `extra_body` method is still supported for backward compatibility:\n", "```python\n", "# Backward compatible method\n", "response = client.chat.completions.create(\n", " model=\"qwen/qwen2.5-0.5b-instruct\",\n", " messages=[{\"role\": \"user\", \"content\": \"Convert to SQL: show all users\"}],\n", " extra_body={\"lora_path\": \"adapter_a\"}, # ← old method\n", " max_tokens=50,\n", ")\n", "```\n", "**Note:** When both `model:adapter` and `extra_body[\"lora_path\"]` are specified, the `model:adapter` syntax takes precedence." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process)" ] } ], "metadata": { "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 2 } ================================================ FILE: docs/basic_usage/openai_api_embeddings.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# OpenAI APIs - Embedding\n", "\n", "SGLang provides OpenAI-compatible APIs to enable a smooth transition from OpenAI services to self-hosted local models.\n", "A complete reference for the API is available in the [OpenAI API Reference](https://platform.openai.com/docs/guides/embeddings).\n", "\n", "This tutorial covers the embedding APIs for embedding models. For a list of the supported models see the [corresponding overview page](../supported_models/retrieval_ranking/embedding_models.md)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Launch A Server\n", "\n", "Launch the server in your terminal and wait for it to initialize. Remember to add `--is-embedding` to the command." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from sglang.test.doc_patch import launch_server_cmd\n", "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", "\n", "embedding_process, port = launch_server_cmd(\"\"\"\n", "python3 -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-1.5B-instruct \\\n", " --host 0.0.0.0 --is-embedding --log-level warning\n", "\"\"\")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=embedding_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Using cURL" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import subprocess, json\n", "\n", "text = \"Once upon a time\"\n", "\n", "curl_text = f\"\"\"curl -s http://localhost:{port}/v1/embeddings \\\n", " -H \"Content-Type: application/json\" \\\n", " -d '{{\"model\": \"Alibaba-NLP/gte-Qwen2-1.5B-instruct\", \"input\": \"{text}\"}}'\"\"\"\n", "\n", "result = subprocess.check_output(curl_text, shell=True)\n", "\n", "print(result)\n", "\n", "text_embedding = json.loads(result)[\"data\"][0][\"embedding\"]\n", "\n", "print_highlight(f\"Text embedding (first 10): {text_embedding[:10]}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Using Python Requests" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import requests\n", "\n", "text = \"Once upon a time\"\n", "\n", "response = requests.post(\n", " f\"http://localhost:{port}/v1/embeddings\",\n", " json={\"model\": \"Alibaba-NLP/gte-Qwen2-1.5B-instruct\", \"input\": text},\n", ")\n", "\n", "text_embedding = response.json()[\"data\"][0][\"embedding\"]\n", "\n", "print_highlight(f\"Text embedding (first 10): {text_embedding[:10]}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Using OpenAI Python Client" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import openai\n", "\n", "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", "\n", "# Text embedding example\n", "response = client.embeddings.create(\n", " model=\"Alibaba-NLP/gte-Qwen2-1.5B-instruct\",\n", " input=text,\n", ")\n", "\n", "embedding = response.data[0].embedding[:10]\n", "print_highlight(f\"Text embedding (first 10): {embedding}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Using Input IDs\n", "\n", "SGLang also supports `input_ids` as input to get the embedding." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\n", "import os\n", "from transformers import AutoTokenizer\n", "\n", "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n", "\n", "tokenizer = AutoTokenizer.from_pretrained(\"Alibaba-NLP/gte-Qwen2-1.5B-instruct\")\n", "input_ids = tokenizer.encode(text)\n", "\n", "curl_ids = f\"\"\"curl -s http://localhost:{port}/v1/embeddings \\\n", " -H \"Content-Type: application/json\" \\\n", " -d '{{\"model\": \"Alibaba-NLP/gte-Qwen2-1.5B-instruct\", \"input\": {json.dumps(input_ids)}}}'\"\"\"\n", "\n", "input_ids_embedding = json.loads(subprocess.check_output(curl_ids, shell=True))[\"data\"][\n", " 0\n", "][\"embedding\"]\n", "\n", "print_highlight(f\"Input IDs embedding (first 10): {input_ids_embedding[:10]}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(embedding_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Multi-Modal Embedding Model\n", "Please refer to [Multi-Modal Embedding Model](../supported_models/retrieval_ranking/embedding_models.md)" ] } ], "metadata": { "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 2 } ================================================ FILE: docs/basic_usage/openai_api_vision.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# OpenAI APIs - Vision\n", "\n", "SGLang provides OpenAI-compatible APIs to enable a smooth transition from OpenAI services to self-hosted local models.\n", "A complete reference for the API is available in the [OpenAI API Reference](https://platform.openai.com/docs/guides/vision).\n", "This tutorial covers the vision APIs for vision language models.\n", "\n", "SGLang supports various vision language models such as Llama 3.2, LLaVA-OneVision, Qwen2.5-VL, Gemma3 and [more](../supported_models/text_generation/multimodal_language_models.md).\n", "\n", "As an alternative to the OpenAI API, you can also use the [SGLang offline engine](https://github.com/sgl-project/sglang/blob/main/examples/runtime/engine/offline_batch_inference_vlm.py)." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Launch A Server\n", "\n", "Launch the server in your terminal and wait for it to initialize." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from sglang.test.doc_patch import launch_server_cmd\n", "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", "\n", "example_image_url = \"https://raw.githubusercontent.com/sgl-project/sglang/main/examples/assets/example_image.png\"\n", "logo_image_url = (\n", " \"https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png\"\n", ")\n", "\n", "vision_process, port = launch_server_cmd(\"\"\"\n", "python3 -m sglang.launch_server --model-path Qwen/Qwen2.5-VL-7B-Instruct --log-level warning\n", "\"\"\")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=vision_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Using cURL\n", "\n", "Once the server is up, you can send test requests using curl or requests." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import subprocess\n", "\n", "curl_command = f\"\"\"\n", "curl -s http://localhost:{port}/v1/chat/completions \\\\\n", " -H \"Content-Type: application/json\" \\\\\n", " -d '{{\n", " \"model\": \"Qwen/Qwen2.5-VL-7B-Instruct\",\n", " \"messages\": [\n", " {{\n", " \"role\": \"user\",\n", " \"content\": [\n", " {{\n", " \"type\": \"text\",\n", " \"text\": \"What’s in this image?\"\n", " }},\n", " {{\n", " \"type\": \"image_url\",\n", " \"image_url\": {{\n", " \"url\": \"{example_image_url}\"\n", " }}\n", " }}\n", " ]\n", " }}\n", " ],\n", " \"max_tokens\": 300\n", " }}'\n", "\"\"\"\n", "\n", "response = subprocess.check_output(curl_command, shell=True).decode()\n", "print_highlight(response)\n", "\n", "\n", "response = subprocess.check_output(curl_command, shell=True).decode()\n", "print_highlight(response)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Using Python Requests" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import requests\n", "\n", "url = f\"http://localhost:{port}/v1/chat/completions\"\n", "\n", "data = {\n", " \"model\": \"Qwen/Qwen2.5-VL-7B-Instruct\",\n", " \"messages\": [\n", " {\n", " \"role\": \"user\",\n", " \"content\": [\n", " {\"type\": \"text\", \"text\": \"What’s in this image?\"},\n", " {\n", " \"type\": \"image_url\",\n", " \"image_url\": {\"url\": example_image_url},\n", " },\n", " ],\n", " }\n", " ],\n", " \"max_tokens\": 300,\n", "}\n", "\n", "response = requests.post(url, json=data)\n", "print_highlight(response.text)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Using OpenAI Python Client" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from openai import OpenAI\n", "\n", "client = OpenAI(base_url=f\"http://localhost:{port}/v1\", api_key=\"None\")\n", "\n", "response = client.chat.completions.create(\n", " model=\"Qwen/Qwen2.5-VL-7B-Instruct\",\n", " messages=[\n", " {\n", " \"role\": \"user\",\n", " \"content\": [\n", " {\n", " \"type\": \"text\",\n", " \"text\": \"What is in this image?\",\n", " },\n", " {\n", " \"type\": \"image_url\",\n", " \"image_url\": {\"url\": example_image_url},\n", " },\n", " ],\n", " }\n", " ],\n", " max_tokens=300,\n", ")\n", "\n", "print_highlight(response.choices[0].message.content)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Multiple-Image Inputs\n", "\n", "The server also supports multiple images and interleaved text and images if the model supports it." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from openai import OpenAI\n", "\n", "client = OpenAI(base_url=f\"http://localhost:{port}/v1\", api_key=\"None\")\n", "\n", "response = client.chat.completions.create(\n", " model=\"Qwen/Qwen2.5-VL-7B-Instruct\",\n", " messages=[\n", " {\n", " \"role\": \"user\",\n", " \"content\": [\n", " {\n", " \"type\": \"image_url\",\n", " \"image_url\": {\n", " \"url\": example_image_url,\n", " },\n", " },\n", " {\n", " \"type\": \"image_url\",\n", " \"image_url\": {\n", " \"url\": logo_image_url,\n", " },\n", " },\n", " {\n", " \"type\": \"text\",\n", " \"text\": \"I have two very different images. They are not related at all. \"\n", " \"Please describe the first image in one sentence, and then describe the second image in another sentence.\",\n", " },\n", " ],\n", " }\n", " ],\n", " temperature=0,\n", ")\n", "\n", "print_highlight(response.choices[0].message.content)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(vision_process)" ] } ], "metadata": { "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 2 } ================================================ FILE: docs/basic_usage/popular_model_usage.rst ================================================ Popular Model Usage (DeepSeek, GPT-OSS, GLM, Llama, MiniMax, Qwen, and more) =============================================================== For more usage examples and recipes, visit the `SGLang Cookbook `_. .. toctree:: :maxdepth: 1 deepseek_v3.md deepseek_v32.md glm45.md glmv.md gpt_oss.md minimax_m2.md qwen3.md qwen3_5.md qwen3_vl.md deepseek_ocr.md llama4.md ================================================ FILE: docs/basic_usage/qwen3.md ================================================ # Qwen3-Next Usage SGLang has supported Qwen3-Next-80B-A3B-Instruct and Qwen3-Next-80B-A3B-Thinking since [this PR](https://github.com/sgl-project/sglang/pull/10233). ## Launch Qwen3-Next with SGLang To serve Qwen3-Next models on 4xH100/H200 GPUs: ```bash python3 -m sglang.launch_server --model Qwen/Qwen3-Next-80B-A3B-Instruct --tp 4 ``` ### Configuration Tips - `--max-mamba-cache-size`: Adjust `--max-mamba-cache-size` to increase mamba cache space and max running requests capability. It will decrease KV cache space as a trade-off. You can adjust it according to workload. - `--mamba-ssm-dtype`: `bfloat16` or `float32`, use `bfloat16` to save mamba cache size and `float32` to get more accurate results. The default setting is `float32`. - `--mamba-full-memory-ratio`: The ratio of mamba state memory to full kv cache memory. The default is 0.9. ### Mamba Radix Cache SGLang supports prefix caching for Qwen3-Next models named `MambaRadixCache`, which improves inference speed by reusing computation results. There are two versions of `MambaRadixCache`: - `no_buffer`: The default version, which is also other hybrid linear models' choice. When it is enabled, SGLang will automatically close overlap schedule for compatibility reasons. - `extra_buffer`: An optimized version that is compatible with features like page size > 1, overlap schedule, and speculative decoding. It also supports storing mamba state in branching positions. However, it requires two extra mamba spaces for a ping-pong buffer for each request. To enable it, add the argument `--mamba-scheduler-strategy extra_buffer` when launching the server. ### EAGLE Speculative Decoding **Description**: SGLang has supported Qwen3-Next models with [EAGLE speculative decoding](https://docs.sglang.io/advanced_features/speculative_decoding.html#EAGLE-Decoding). **Usage**: Add arguments `--speculative-algorithm`, `--speculative-num-steps`, `--speculative-eagle-topk` and `--speculative-num-draft-tokens` to enable this feature. For example: ``` bash python3 -m sglang.launch_server \ --model Qwen/Qwen3-Next-80B-A3B-Instruct \ --tp 4 \ --speculative-num-steps 3 \ --speculative-eagle-topk 1 \ --speculative-num-draft-tokens 4 \ --speculative-algo NEXTN ``` Details can be seen in [this PR](https://github.com/sgl-project/sglang/pull/10233). ================================================ FILE: docs/basic_usage/qwen3_5.md ================================================ # Qwen 3.5 Usage Qwen 3.5 is Alibaba's latest generation LLM featuring a hybrid attention architecture, advanced MoE with shared experts, and native multimodal capabilities. Key architecture features: - **Hybrid Attention**: Gated Delta Networks (linear, O(n) complexity) combined with full attention every 4th layer for high associative recall - **MoE with Shared Experts**: Top-8 active out of 64 routed experts plus a dedicated shared expert for universal features - **Multimodal**: DeepStack Vision Transformer with Conv3d for native image and video understanding ## Launch Qwen 3.5 with SGLang ### Dense Model To serve `Qwen/Qwen3.5-397B-A17B` on 8 GPUs: ```bash python3 -m sglang.launch_server \ --model-path Qwen/Qwen3.5-397B-A17B \ --tp 8 \ --trust-remote-code ``` ### AMD GPU (MI300X / MI325X / MI35X) On AMD Instinct GPUs, use the `triton` attention backend. Both the full attention layers and the Gated Delta Net (linear attention) layers use Triton-based kernels on ROCm: ```bash SGLANG_USE_AITER=1 python3 -m sglang.launch_server \ --model-path Qwen/Qwen3.5-397B-A17B \ --tp 8 \ --attention-backend triton \ --trust-remote-code ``` ```{tip} Set `SGLANG_USE_AITER=1` to enable AMD's optimized aiter kernels for MoE and GEMM operations. ``` ### Configuration Tips - `--attention-backend`: Use `triton` on AMD GPUs for Qwen 3.5. The hybrid attention architecture (Gated Delta Networks + full attention) works best with the Triton backend on ROCm. The linear attention (GDN) layers always use Triton kernels internally via the `GDNAttnBackend`. - `--watchdog-timeout`: Increase to `1200` or higher for this large model, as weight loading takes significant time. - `--model-loader-extra-config '{"enable_multithread_load": true}'`: Enables parallel weight loading for faster startup. ### Reasoning and Tool Calling Qwen 3.5 supports reasoning and tool calling via the Qwen3 parsers: ```bash python3 -m sglang.launch_server \ --model-path Qwen/Qwen3.5-397B-A17B \ --tp 8 \ --trust-remote-code \ --reasoning-parser qwen3 \ --tool-call-parser qwen3_coder ``` ## Accuracy Evaluation You can evaluate the model accuracy using `lm-eval`: ```bash pip install lm-eval[api] lm_eval --model local-completions \ --model_args '{"base_url": "http://localhost:8000/v1/completions", "model": "Qwen/Qwen3.5-397B-A17B", "num_concurrent": 256, "max_retries": 10, "max_gen_toks": 2048}' \ --tasks gsm8k \ --batch_size auto \ --num_fewshot 5 \ --trust_remote_code ``` ## Additional Resources - [AMD Day 0 Support for Qwen 3.5 on AMD Instinct GPUs](https://www.amd.com/en/developer/resources/technical-articles/2026/day-0-support-for-qwen-3-5-on-amd-instinct-gpus.html) - [HuggingFace Model Card](https://huggingface.co/Qwen/Qwen3.5-397B-A17B) ================================================ FILE: docs/basic_usage/qwen3_vl.md ================================================ # Qwen3-VL Usage [Qwen3-VL](https://huggingface.co/collections/Qwen/qwen3-vl) is Alibaba’s latest multimodal large language model with strong text, vision, and reasoning capabilities. SGLang supports Qwen3-VL Family of models with Image and Video input support. ## Launch commands for SGLang Below are suggested launch commands tailored for different hardware / precision modes ### FP8 (quantised) mode For high memory-efficiency and latency optimized deployments (e.g., on H100, H200) where FP8 checkpoint is supported: ```bash python3 -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-235B-A22B-Instruct-FP8 \ --tp 8 \ --ep 8 \ --host 0.0.0.0 \ --port 30000 \ --keep-mm-feature-on-device ``` ### Non-FP8 (BF16 / full precision) mode For deployments on A100/H100 where BF16 is used (or FP8 snapshot not used): ```bash python3 -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-235B-A22B-Instruct \ --tp 8 \ --ep 8 \ --host 0.0.0.0 \ --port 30000 \ ``` ## Hardware-specific notes / recommendations - On H100 with FP8: Use the FP8 checkpoint for best memory efficiency. - On A100 / H100 with BF16 (non-FP8): It’s recommended to use `--mm-max-concurrent-calls` to control parallel throughput and GPU memory usage during image/video inference. - On H200 & B200: The model can be run “out of the box”, supporting full context length plus concurrent image + video processing. ## Sending Image/Video Requests ### Image input: ```python import requests url = f"http://localhost:30000/v1/chat/completions" data = { "model": "Qwen/Qwen3-VL-30B-A3B-Instruct", "messages": [ { "role": "user", "content": [ {"type": "text", "text": "What’s in this image?"}, { "type": "image_url", "image_url": { "url": "https://github.com/sgl-project/sglang/blob/main/examples/assets/example_image.png?raw=true" }, }, ], } ], "max_tokens": 300, } response = requests.post(url, json=data) print(response.text) ``` ### Video Input: ```python import requests url = f"http://localhost:30000/v1/chat/completions" data = { "model": "Qwen/Qwen3-VL-30B-A3B-Instruct", "messages": [ { "role": "user", "content": [ {"type": "text", "text": "What’s happening in this video?"}, { "type": "video_url", "video_url": { "url": "https://github.com/sgl-project/sgl-test-files/raw/refs/heads/main/videos/jobs_presenting_ipod.mp4" }, }, ], } ], "max_tokens": 300, } response = requests.post(url, json=data) print(response.text) ``` ## Important Server Parameters and Flags When launching the model server for **multimodal support**, you can use the following command-line arguments to fine-tune performance and behavior: - `--mm-attention-backend`: Specify multimodal attention backend. Eg. `fa3`(Flash Attention 3) - `--mm-max-concurrent-calls `: Specifies the **maximum number of concurrent asynchronous multimodal data processing calls** allowed on the server. Use this to control parallel throughput and GPU memory usage during image/video inference. - `--mm-per-request-timeout `: Defines the **timeout duration (in seconds)** for each multimodal request. If a request exceeds this time limit (e.g., for very large video inputs), it will be automatically terminated. - `--keep-mm-feature-on-device`: Instructs the server to **retain multimodal feature tensors on the GPU** after processing. This avoids device-to-host (D2H) memory copies and improves performance for repeated or high-frequency inference workloads. - `SGLANG_USE_CUDA_IPC_TRANSPORT=1`: Shared memory pool based CUDA IPC for multi-modal data transport. For significantly improving e2e latency. ### Example usage with the above optimizations: ```bash SGLANG_USE_CUDA_IPC_TRANSPORT=1 \ SGLANG_VLM_CACHE_SIZE_MB=0 \ python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-235B-A22B-Instruct \ --host 0.0.0.0 \ --port 30000 \ --trust-remote-code \ --tp-size 8 \ --enable-cache-report \ --log-level info \ --max-running-requests 64 \ --mem-fraction-static 0.65 \ --chunked-prefill-size 8192 \ --attention-backend fa3 \ --mm-attention-backend fa3 \ --enable-metrics ``` ================================================ FILE: docs/basic_usage/sampling_params.md ================================================ # Sampling Parameters This doc describes the sampling parameters of the SGLang Runtime. It is the low-level endpoint of the runtime. If you want a high-level endpoint that can automatically handle chat templates, consider using the [OpenAI Compatible API](openai_api_completions.ipynb). ## `/generate` Endpoint The `/generate` endpoint accepts the following parameters in JSON format. For detailed usage, see the [native API doc](native_api.ipynb). The object is defined at `io_struct.py::GenerateReqInput`. You can also read the source code to find more arguments and docs. | Argument | Type/Default | Description | |----------------------------|------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------| | text | `Optional[Union[List[str], str]] = None` | The input prompt. Can be a single prompt or a batch of prompts. | | input_ids | `Optional[Union[List[List[int]], List[int]]] = None` | The token IDs for text; one can specify either text or input_ids. | | input_embeds | `Optional[Union[List[List[List[float]]], List[List[float]]]] = None` | The embeddings for input_ids; one can specify either text, input_ids, or input_embeds. | | image_data | `Optional[Union[List[List[ImageDataItem]], List[ImageDataItem], ImageDataItem]] = None` | The image input. Supports three formats: (1) **Raw images**: PIL Image, file path, URL, or base64 string; (2) **Processor output**: Dict with `format: "processor_output"` containing HuggingFace processor outputs; (3) **Precomputed embeddings**: Dict with `format: "precomputed_embedding"` and `feature` containing pre-calculated visual embeddings. Can be a single image, list of images, or list of lists of images. See [Multimodal Input Formats](#multimodal-input-formats) for details. | | audio_data | `Optional[Union[List[AudioDataItem], AudioDataItem]] = None` | The audio input. Can be a file name, URL, or base64 encoded string. | | sampling_params | `Optional[Union[List[Dict], Dict]] = None` | The sampling parameters as described in the sections below. | | rid | `Optional[Union[List[str], str]] = None` | The request ID. | | return_logprob | `Optional[Union[List[bool], bool]] = None` | Whether to return log probabilities for tokens. | | logprob_start_len | `Optional[Union[List[int], int]] = None` | If return_logprob, the start location in the prompt for returning logprobs. Default is "-1", which returns logprobs for output tokens only. | | top_logprobs_num | `Optional[Union[List[int], int]] = None` | If return_logprob, the number of top logprobs to return at each position. | | token_ids_logprob | `Optional[Union[List[List[int]], List[int]]] = None` | If return_logprob, the token IDs to return logprob for. | | return_text_in_logprobs | `bool = False` | Whether to detokenize tokens in text in the returned logprobs. | | stream | `bool = False` | Whether to stream output. | | lora_path | `Optional[Union[List[Optional[str]], Optional[str]]] = None` | The path to the LoRA. | | custom_logit_processor | `Optional[Union[List[Optional[str]], str]] = None` | Custom logit processor for advanced sampling control. Must be a serialized instance of `CustomLogitProcessor` using its `to_str()` method. For usage see below. | | return_hidden_states | `Union[List[bool], bool] = False` | Whether to return hidden states. | | return_routed_experts | `bool = False` | Whether to return routed experts for MoE models. Requires `--enable-return-routed-experts` server flag. Returns base64-encoded int32 expert IDs as a flattened array with logical shape `[num_tokens, num_layers, top_k]`. | ## Sampling parameters The object is defined at `sampling_params.py::SamplingParams`. You can also read the source code to find more arguments and docs. ### Note on defaults By default, SGLang initializes several sampling parameters from the model's `generation_config.json` (when the server is launched with `--sampling-defaults model`, which is the default). To use SGLang/OpenAI constant defaults instead, start the server with `--sampling-defaults openai`. You can always override any parameter per request via `sampling_params`. ```bash # Use model-provided defaults from generation_config.json (default behavior) python -m sglang.launch_server --model-path --sampling-defaults model # Use SGLang/OpenAI constant defaults instead python -m sglang.launch_server --model-path --sampling-defaults openai ``` ### Core parameters | Argument | Type/Default | Description | |-----------------|----------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------| | max_new_tokens | `int = 128` | The maximum output length measured in tokens. | | stop | `Optional[Union[str, List[str]]] = None` | One or multiple [stop words](https://platform.openai.com/docs/api-reference/chat/create#chat-create-stop). Generation will stop if one of these words is sampled. | | stop_token_ids | `Optional[List[int]] = None` | Provide stop words in the form of token IDs. Generation will stop if one of these token IDs is sampled. | | stop_regex | `Optional[Union[str, List[str]]] = None` | Stop when hitting any of the regex patterns in this list | | temperature | `float (model default; fallback 1.0)` | [Temperature](https://platform.openai.com/docs/api-reference/chat/create#chat-create-temperature) when sampling the next token. `temperature = 0` corresponds to greedy sampling, a higher temperature leads to more diversity. | | top_p | `float (model default; fallback 1.0)` | [Top-p](https://platform.openai.com/docs/api-reference/chat/create#chat-create-top_p) selects tokens from the smallest sorted set whose cumulative probability exceeds `top_p`. When `top_p = 1`, this reduces to unrestricted sampling from all tokens. | | top_k | `int (model default; fallback -1)` | [Top-k](https://developer.nvidia.com/blog/how-to-get-better-outputs-from-your-large-language-model/#predictability_vs_creativity) randomly selects from the `k` highest-probability tokens. | | min_p | `float (model default; fallback 0.0)` | [Min-p](https://github.com/huggingface/transformers/issues/27670) samples from tokens with probability larger than `min_p * highest_token_probability`. | ### Penalizers | Argument | Type/Default | Description | |--------------------|------------------------|------------------------------------------------------------------------------------------------------------------------------------------------| | frequency_penalty | `float = 0.0` | Penalizes tokens based on their frequency in generation so far. Must be between `-2` and `2` where negative numbers encourage repeatment of tokens and positive number encourages sampling of new tokens. The scaling of penalization grows linearly with each appearance of a token. | | presence_penalty | `float = 0.0` | Penalizes tokens if they appeared in the generation so far. Must be between `-2` and `2` where negative numbers encourage repeatment of tokens and positive number encourages sampling of new tokens. The scaling of the penalization is constant if a token occurred. | | repetition_penalty | `float = 1.0` | Scales the logits of previously generated tokens to discourage (values > 1) or encourage (values < 1) repetition. Valid range is `[0, 2]`; `1.0` leaves probabilities unchanged. | | min_new_tokens | `int = 0` | Forces the model to generate at least `min_new_tokens` until a stop word or EOS token is sampled. Note that this might lead to unintended behavior, for example, if the distribution is highly skewed towards these tokens. | ### Constrained decoding Please refer to our dedicated guide on [constrained decoding](../advanced_features/structured_outputs.ipynb) for the following parameters. | Argument | Type/Default | Description | |-----------------|---------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------| | json_schema | `Optional[str] = None` | JSON schema for structured outputs. | | regex | `Optional[str] = None` | Regex for structured outputs. | | ebnf | `Optional[str] = None` | EBNF for structured outputs. | | structural_tag | `Optional[str] = None` | The structural tag for structured outputs. | ### Other options | Argument | Type/Default | Description | |-------------------------------|---------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------| | n | `int = 1` | Specifies the number of output sequences to generate per request. (Generating multiple outputs in one request (n > 1) is discouraged; repeating the same prompts several times offers better control and efficiency.) | | ignore_eos | `bool = False` | Don't stop generation when EOS token is sampled. | | skip_special_tokens | `bool = True` | Remove special tokens during decoding. | | spaces_between_special_tokens | `bool = True` | Whether or not to add spaces between special tokens during detokenization. | | no_stop_trim | `bool = False` | Don't trim stop words or EOS token from the generated text. | | custom_params | `Optional[List[Optional[Dict[str, Any]]]] = None` | Used when employing `CustomLogitProcessor`. For usage, see below. | ## Examples ### Normal Launch a server: ```bash python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --port 30000 ``` Send a request: ```python import requests response = requests.post( "http://localhost:30000/generate", json={ "text": "The capital of France is", "sampling_params": { "temperature": 0, "max_new_tokens": 32, }, }, ) print(response.json()) ``` Detailed example in [send request](./send_request.ipynb). ### Streaming Send a request and stream the output: ```python import requests, json response = requests.post( "http://localhost:30000/generate", json={ "text": "The capital of France is", "sampling_params": { "temperature": 0, "max_new_tokens": 32, }, "stream": True, }, stream=True, ) prev = 0 for chunk in response.iter_lines(decode_unicode=False): chunk = chunk.decode("utf-8") if chunk and chunk.startswith("data:"): if chunk == "data: [DONE]": break data = json.loads(chunk[5:].strip("\n")) output = data["text"].strip() print(output[prev:], end="", flush=True) prev = len(output) print("") ``` Detailed example in [openai compatible api](openai_api_completions.ipynb). ### Multimodal Launch a server: ```bash python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-7b-ov ``` Download an image: ```bash curl -o example_image.png -L https://github.com/sgl-project/sglang/blob/main/examples/assets/example_image.png?raw=true ``` Send a request: ```python import requests response = requests.post( "http://localhost:30000/generate", json={ "text": "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" "<|im_start|>user\n\nDescribe this image in a very short sentence.<|im_end|>\n" "<|im_start|>assistant\n", "image_data": "example_image.png", "sampling_params": { "temperature": 0, "max_new_tokens": 32, }, }, ) print(response.json()) ``` The `image_data` can be a file name, a URL, or a base64 encoded string. See also `python/sglang/srt/utils.py:load_image`. Streaming is supported in a similar manner as [above](#streaming). Detailed example in [OpenAI API Vision](openai_api_vision.ipynb). ### Structured Outputs (JSON, Regex, EBNF) You can specify a JSON schema, regular expression or [EBNF](https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_form) to constrain the model output. The model output will be guaranteed to follow the given constraints. Only one constraint parameter (`json_schema`, `regex`, or `ebnf`) can be specified for a request. SGLang supports two grammar backends: - [XGrammar](https://github.com/mlc-ai/xgrammar) (default): Supports JSON schema, regular expression, and EBNF constraints. - XGrammar currently uses the [GGML BNF format](https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md). - [Outlines](https://github.com/dottxt-ai/outlines): Supports JSON schema and regular expression constraints. If instead you want to initialize the Outlines backend, you can use `--grammar-backend outlines` flag: ```bash python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ --port 30000 --host 0.0.0.0 --grammar-backend [xgrammar|outlines] # xgrammar or outlines (default: xgrammar) ``` ```python import json import requests json_schema = json.dumps({ "type": "object", "properties": { "name": {"type": "string", "pattern": "^[\\w]+$"}, "population": {"type": "integer"}, }, "required": ["name", "population"], }) # JSON (works with both Outlines and XGrammar) response = requests.post( "http://localhost:30000/generate", json={ "text": "Here is the information of the capital of France in the JSON format.\n", "sampling_params": { "temperature": 0, "max_new_tokens": 64, "json_schema": json_schema, }, }, ) print(response.json()) # Regular expression (Outlines backend only) response = requests.post( "http://localhost:30000/generate", json={ "text": "Paris is the capital of", "sampling_params": { "temperature": 0, "max_new_tokens": 64, "regex": "(France|England)", }, }, ) print(response.json()) # EBNF (XGrammar backend only) response = requests.post( "http://localhost:30000/generate", json={ "text": "Write a greeting.", "sampling_params": { "temperature": 0, "max_new_tokens": 64, "ebnf": 'root ::= "Hello" | "Hi" | "Hey"', }, }, ) print(response.json()) ``` Detailed example in [structured outputs](../advanced_features/structured_outputs.ipynb). ### Custom logit processor Launch a server with `--enable-custom-logit-processor` flag on. ```bash python -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3-8B-Instruct \ --port 30000 \ --enable-custom-logit-processor ``` Define a custom logit processor that will always sample a specific token id. ```python from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor class DeterministicLogitProcessor(CustomLogitProcessor): """A dummy logit processor that changes the logits to always sample the given token id. """ def __call__(self, logits, custom_param_list): # Check that the number of logits matches the number of custom parameters assert logits.shape[0] == len(custom_param_list) key = "token_id" for i, param_dict in enumerate(custom_param_list): # Mask all other tokens logits[i, :] = -float("inf") # Assign highest probability to the specified token logits[i, param_dict[key]] = 0.0 return logits ``` Send a request: ```python import requests response = requests.post( "http://localhost:30000/generate", json={ "text": "The capital of France is", "custom_logit_processor": DeterministicLogitProcessor().to_str(), "sampling_params": { "temperature": 0.0, "max_new_tokens": 32, "custom_params": {"token_id": 5}, }, }, ) print(response.json()) ``` Send an OpenAI chat completion request: ```python import openai from sglang.utils import print_highlight client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="None") response = client.chat.completions.create( model="meta-llama/Meta-Llama-3-8B-Instruct", messages=[ {"role": "user", "content": "List 3 countries and their capitals."}, ], temperature=0.0, max_tokens=32, extra_body={ "custom_logit_processor": DeterministicLogitProcessor().to_str(), "custom_params": {"token_id": 5}, }, ) print_highlight(f"Response: {response}") ``` ================================================ FILE: docs/basic_usage/send_request.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Sending Requests\n", "This notebook provides a quick-start guide to use SGLang in chat completions after installation. Once your server is running, API documentation is available at `http://localhost:30000/docs` (Swagger UI), `http://localhost:30000/redoc` (ReDoc), or `http://localhost:30000/openapi.json` (OpenAPI spec, useful for AI agents). Replace `30000` with your port if using a different one.\n", "\n", "- For Vision Language Models, see [OpenAI APIs - Vision](openai_api_vision.ipynb).\n", "- For Embedding Models, see [OpenAI APIs - Embedding](openai_api_embeddings.ipynb) and [Encode (embedding model)](native_api.html#Encode-(embedding-model)).\n", "- For Reward Models, see [Classify (reward model)](native_api.html#Classify-(reward-model))." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Launch A Server" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from sglang.test.doc_patch import launch_server_cmd\n", "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", "\n", "# This is equivalent to running the following command in your terminal\n", "# python3 -m sglang.launch_server --model-path qwen/qwen2.5-0.5b-instruct --host 0.0.0.0\n", "\n", "server_process, port = launch_server_cmd(\"\"\"\n", "python3 -m sglang.launch_server --model-path qwen/qwen2.5-0.5b-instruct \\\n", " --host 0.0.0.0 --log-level warning\n", "\"\"\")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Using cURL\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import subprocess, json\n", "\n", "curl_command = f\"\"\"\n", "curl -s http://localhost:{port}/v1/chat/completions \\\n", " -H \"Content-Type: application/json\" \\\n", " -d '{{\"model\": \"qwen/qwen2.5-0.5b-instruct\", \"messages\": [{{\"role\": \"user\", \"content\": \"What is the capital of France?\"}}]}}'\n", "\"\"\"\n", "\n", "response = json.loads(subprocess.check_output(curl_command, shell=True))\n", "print_highlight(response)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Using Python Requests" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import requests\n", "\n", "url = f\"http://localhost:{port}/v1/chat/completions\"\n", "\n", "data = {\n", " \"model\": \"qwen/qwen2.5-0.5b-instruct\",\n", " \"messages\": [{\"role\": \"user\", \"content\": \"What is the capital of France?\"}],\n", "}\n", "\n", "response = requests.post(url, json=data)\n", "print_highlight(response.json())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Using OpenAI Python Client" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import openai\n", "\n", "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", "\n", "response = client.chat.completions.create(\n", " model=\"qwen/qwen2.5-0.5b-instruct\",\n", " messages=[\n", " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", " ],\n", " temperature=0,\n", " max_tokens=64,\n", ")\n", "print_highlight(response)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Streaming" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import openai\n", "\n", "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", "\n", "# Use stream=True for streaming responses\n", "response = client.chat.completions.create(\n", " model=\"qwen/qwen2.5-0.5b-instruct\",\n", " messages=[\n", " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", " ],\n", " temperature=0,\n", " max_tokens=64,\n", " stream=True,\n", ")\n", "\n", "# Handle the streaming output\n", "for chunk in response:\n", " if chunk.choices[0].delta.content:\n", " print(chunk.choices[0].delta.content, end=\"\", flush=True)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Using Native Generation APIs\n", "\n", "You can also use the native `/generate` endpoint with requests, which provides more flexibility. An API reference is available at [Sampling Parameters](sampling_params.md)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import requests\n", "\n", "response = requests.post(\n", " f\"http://localhost:{port}/generate\",\n", " json={\n", " \"text\": \"The capital of France is\",\n", " \"sampling_params\": {\n", " \"temperature\": 0,\n", " \"max_new_tokens\": 32,\n", " },\n", " },\n", ")\n", "\n", "print_highlight(response.json())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Streaming" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import requests, json\n", "\n", "response = requests.post(\n", " f\"http://localhost:{port}/generate\",\n", " json={\n", " \"text\": \"The capital of France is\",\n", " \"sampling_params\": {\n", " \"temperature\": 0,\n", " \"max_new_tokens\": 32,\n", " },\n", " \"stream\": True,\n", " },\n", " stream=True,\n", ")\n", "\n", "prev = 0\n", "for chunk in response.iter_lines(decode_unicode=False):\n", " chunk = chunk.decode(\"utf-8\")\n", " if chunk and chunk.startswith(\"data:\"):\n", " if chunk == \"data: [DONE]\":\n", " break\n", " data = json.loads(chunk[5:].strip(\"\\n\"))\n", " output = data[\"text\"]\n", " print(output[prev:], end=\"\", flush=True)\n", " prev = len(output)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process)" ] } ], "metadata": { "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 2 } ================================================ FILE: docs/conf.py ================================================ import os import sys from datetime import datetime sys.path.insert(0, os.path.abspath("../..")) version_file = "../python/sglang/version.py" with open(version_file, "r") as f: exec(compile(f.read(), version_file, "exec")) __version__ = locals()["__version__"] project = "SGLang" copyright = f"2023-{datetime.now().year}, SGLang" author = "SGLang Team" version = __version__ release = __version__ extensions = [ "sphinx.ext.autodoc", "sphinx.ext.autosummary", "sphinx.ext.napoleon", "sphinx.ext.viewcode", "sphinx.ext.autosectionlabel", "sphinx.ext.intersphinx", "sphinx_tabs.tabs", "myst_parser", "sphinx_copybutton", "sphinxcontrib.mermaid", "nbsphinx", "sphinx.ext.mathjax", ] nbsphinx_allow_errors = True nbsphinx_execute = "never" autosectionlabel_prefix_document = True nbsphinx_allow_directives = True myst_enable_extensions = [ "dollarmath", "amsmath", "deflist", "colon_fence", "html_image", "linkify", "substitution", ] myst_heading_anchors = 3 nbsphinx_kernel_name = "python3" nbsphinx_execute_arguments = [ "--InlineBackend.figure_formats={'svg', 'pdf'}", "--InlineBackend.rc={'figure.dpi': 96}", ] nb_render_priority = { "html": ( "application/vnd.jupyter.widget-view+json", "application/javascript", "text/html", "image/svg+xml", "image/png", "image/jpeg", "text/markdown", "text/latex", "text/plain", ) } myst_enable_extensions = [ "dollarmath", "amsmath", "deflist", "colon_fence", "html_image", "linkify", "substitution", ] myst_heading_anchors = 3 myst_ref_domains = ["std", "py"] templates_path = ["_templates"] source_suffix = { ".rst": "restructuredtext", ".md": "markdown", } master_doc = "index" language = "en" exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] pygments_style = "sphinx" html_theme = "sphinx_book_theme" html_logo = "_static/image/logo.png" html_favicon = "_static/image/logo.ico" html_title = project html_copy_source = True html_last_updated_fmt = "" html_theme_options = { "repository_url": "https://github.com/sgl-project/sgl-project.github.io", "repository_branch": "main", "show_navbar_depth": 3, "max_navbar_depth": 4, "collapse_navbar": True, "use_edit_page_button": True, "use_source_button": True, "use_issues_button": True, "use_repository_button": True, "use_download_button": True, "use_sidenotes": True, "show_toc_level": 2, } html_context = { "display_github": True, "github_user": "sgl-project", "github_repo": "sgl-project.github.io", "github_version": "main", "conf_py_path": "/docs/", } html_static_path = ["_static"] html_css_files = ["css/custom_log.css"] def setup(app): app.add_css_file("css/custom_log.css") myst_enable_extensions = [ "dollarmath", "amsmath", "deflist", "colon_fence", ] myst_heading_anchors = 5 htmlhelp_basename = "sglangdoc" latex_elements = {} latex_documents = [ (master_doc, "sglang.tex", "sglang Documentation", "SGLang Team", "manual"), ] man_pages = [(master_doc, "sglang", "sglang Documentation", [author], 1)] texinfo_documents = [ ( master_doc, "sglang", "sglang Documentation", author, "sglang", "One line description of project.", "Miscellaneous", ), ] epub_title = project epub_exclude_files = ["search.html"] copybutton_prompt_text = r">>> |\.\.\. " copybutton_prompt_is_regexp = True autodoc_preserve_defaults = True navigation_with_keys = False autodoc_mock_imports = [ "torch", "transformers", "triton", ] intersphinx_mapping = { "python": ("https://docs.python.org/3.12", None), "typing_extensions": ("https://typing-extensions.readthedocs.io/en/latest", None), "pillow": ("https://pillow.readthedocs.io/en/stable", None), "numpy": ("https://numpy.org/doc/stable", None), "torch": ("https://pytorch.org/docs/stable", None), } html_theme = "sphinx_book_theme" nbsphinx_prolog = """ .. raw:: html """ ================================================ FILE: docs/deploy.py ================================================ # Deploy the documents import os from datetime import datetime def run_cmd(cmd): print(cmd) os.system(cmd) run_cmd("cd $DOC_SITE_PATH; git pull") # (Optional) Remove old files # run_cmd("rm -rf $ALPA_SITE_PATH/*") run_cmd("cp -r _build/html/* $DOC_SITE_PATH") cmd_message = f"Update {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" run_cmd( f"cd $DOC_SITE_PATH; git add .; git commit -m '{cmd_message}'; git push origin main" ) ================================================ FILE: docs/developer_guide/bench_serving.md ================================================ # Bench Serving Guide This guide explains how to benchmark online serving throughput and latency using `python -m sglang.bench_serving`. It supports multiple inference backends via OpenAI-compatible and native endpoints, and produces both console metrics and optional JSONL outputs. ### What it does - Generates synthetic or dataset-driven prompts and submits them to a target serving endpoint - Measures throughput, time-to-first-token (TTFT), inter-token latency (ITL), per-request end-to-end latency, and more - Supports streaming or non-streaming modes, rate control, and concurrency limits ### Supported backends and endpoints - `sglang` / `sglang-native`: `POST /generate` - `sglang-oai`, `vllm`, `lmdeploy`: `POST /v1/completions` - `sglang-oai-chat`, `vllm-chat`, `lmdeploy-chat`: `POST /v1/chat/completions` - `trt` (TensorRT-LLM): `POST /v2/models/ensemble/generate_stream` - `gserver`: Custom server (Not Implemented yet in this script) - `truss`: `POST /v1/models/model:predict` If `--base-url` is provided, requests are sent to it. Otherwise, `--host` and `--port` are used. When `--model` is not provided, the script will attempt to query `GET /v1/models` for an available model ID (OpenAI-compatible endpoints). ### Prerequisites - Python 3.8+ - Dependencies typically used by this script: `aiohttp`, `numpy`, `requests`, `tqdm`, `transformers`, and for some datasets `datasets`, `pillow`, `pybase64`. Install as needed. - An inference server running and reachable via the endpoints above - If your server requires authentication, set environment variable `OPENAI_API_KEY` (used as `Authorization: Bearer `) ### Quick start Run a basic benchmark against an sglang server exposing `/generate`: ```bash python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct ``` ```bash python3 -m sglang.bench_serving \ --backend sglang \ --host 127.0.0.1 --port 30000 \ --num-prompts 1000 \ --model meta-llama/Llama-3.1-8B-Instruct ``` Or, using an OpenAI-compatible endpoint (completions): ```bash python3 -m sglang.bench_serving \ --backend vllm \ --base-url http://127.0.0.1:8000 \ --num-prompts 1000 \ --model meta-llama/Llama-3.1-8B-Instruct ``` ### Datasets Select with `--dataset-name`: - `sharegpt` (default): loads ShareGPT-style pairs; optionally restrict with `--sharegpt-context-len` and override outputs with `--sharegpt-output-len` - `random`: random text lengths; sampled from ShareGPT token space - `random-ids`: random token ids (can lead to gibberish) - `image`: generates images and wraps them in chat messages; supports custom resolutions, multiple formats, and different content types - `generated-shared-prefix`: synthetic dataset with shared long system prompts and short questions - `mmmu`: samples from MMMU (Math split) and includes images Common dataset flags: - `--num-prompts N`: number of requests - `--random-input-len`, `--random-output-len`, `--random-range-ratio`: for random/random-ids/image - `--image-count`: Number of images per request (for `image` dataset). - `--apply-chat-template`: apply tokenizer chat template when constructing prompts - `--dataset-path PATH`: file path for ShareGPT json; if blank and missing, it will be downloaded and cached Generated Shared Prefix flags (for `generated-shared-prefix`): - `--gsp-num-groups` - `--gsp-prompts-per-group` - `--gsp-system-prompt-len` - `--gsp-question-len` - `--gsp-output-len` Image dataset flags (for `image`): - `--image-count`: Number of images per request - `--image-resolution`: Image resolution; supports presets (4k, 1080p, 720p, 360p) or custom 'heightxwidth' format (e.g., 1080x1920, 512x768) - `--image-format`: Image format (jpeg or png) - `--image-content`: Image content type (random or blank) ### Examples 1. To benchmark image dataset with 3 images per request, 500 prompts, 512 input length, and 512 output length, you can run: ```bash python -m sglang.launch_server --model-path Qwen/Qwen2.5-VL-3B-Instruct --disable-radix-cache ``` ```bash python -m sglang.bench_serving \ --backend sglang-oai-chat \ --dataset-name image \ --num-prompts 500 \ --image-count 3 \ --image-resolution 720p \ --random-input-len 512 \ --random-output-len 512 ``` 2. To benchmark random dataset with 3000 prompts, 1024 input length, and 1024 output length, you can run: ```bash python -m sglang.launch_server --model-path Qwen/Qwen2.5-3B-Instruct ``` ```bash python3 -m sglang.bench_serving \ --backend sglang \ --dataset-name random \ --num-prompts 3000 \ --random-input 1024 \ --random-output 1024 \ --random-range-ratio 0.5 ``` ### Choosing model and tokenizer - `--model` is required unless the backend exposes `GET /v1/models`, in which case the first model ID is auto-selected. - `--tokenizer` defaults to `--model`. Both can be HF model IDs or local paths. - For ModelScope workflows, setting `SGLANG_USE_MODELSCOPE=true` enables fetching via ModelScope (weights are skipped for speed). - If your tokenizer lacks a chat template, the script warns because token counting can be less robust for gibberish outputs. ### Rate, concurrency, and streaming - `--request-rate`: requests per second. `inf` sends all immediately (burst). Non-infinite rate uses a Poisson process for arrival times. - `--max-concurrency`: caps concurrent in-flight requests regardless of arrival rate. - `--disable-stream`: switch to non-streaming mode when supported; TTFT then equals total latency for chat completions. ### Other key options - `--output-file FILE.jsonl`: append JSONL results to file; auto-named if unspecified - `--output-details`: include per-request arrays (generated texts, errors, ttfts, itls, input/output lens) - `--extra-request-body '{"top_p":0.9,"temperature":0.6}'`: merged into payload (sampling params, etc.) - `--disable-ignore-eos`: pass through EOS behavior (varies by backend) - `--warmup-requests N`: run warmup requests with short output first (default 1) - `--flush-cache`: call `/flush_cache` (sglang) before main run - `--profile`: call `/start_profile` and `/stop_profile` (requires server to enable profiling, e.g., `SGLANG_TORCH_PROFILER_DIR`) - `--lora-name name1 name2 ...`: randomly pick one per request and pass to backend (e.g., `lora_path` for sglang) - `--tokenize-prompt`: send integer IDs instead of text (currently supports `--backend sglang` only) ### Authentication If your target endpoint requires OpenAI-style auth, set: ```bash export OPENAI_API_KEY=sk-...yourkey... ``` The script will add `Authorization: Bearer $OPENAI_API_KEY` automatically for OpenAI-compatible routes. ### Metrics explained Printed after each run: - Request throughput (req/s) - Input token throughput (tok/s) - includes both text and vision tokens - Output token throughput (tok/s) - Total token throughput (tok/s) - includes both text and vision tokens - Total input text tokens and Total input vision tokens - per-modality breakdown - Concurrency: aggregate time of all requests divided by wall time - End-to-End Latency (ms): mean/median/std/p99 per-request total latency - Time to First Token (TTFT, ms): mean/median/std/p99 for streaming mode - Inter-Token Latency (ITL, ms): mean/median/std/p95/p99/max between tokens - TPOT (ms): Token processing time after first token, i.e., `(latency - ttft)/(tokens-1)` - Accept length (sglang-only, if available): speculative decoding accept length The script also retokenizes generated text with the configured tokenizer and reports "retokenized" counts. ### JSONL output format When `--output-file` is set, one JSON object is appended per run. Base fields: - Arguments summary: backend, dataset, request_rate, max_concurrency, etc. - Duration and totals: completed, total_input_tokens, total_output_tokens, retokenized totals - Throughputs and latency statistics as printed in the console - `accept_length` when available (sglang) With `--output-details`, an extended object also includes arrays: - `input_lens`, `output_lens` - `ttfts`, `itls` (per request: ITL arrays) - `generated_texts`, `errors` ### End-to-end examples 1) sglang native `/generate` (streaming): ```bash python3 -m sglang.bench_serving \ --backend sglang \ --host 127.0.0.1 --port 30000 \ --model meta-llama/Llama-3.1-8B-Instruct \ --dataset-name random \ --random-input-len 1024 --random-output-len 1024 --random-range-ratio 0.5 \ --num-prompts 2000 \ --request-rate 100 \ --max-concurrency 512 \ --output-file sglang_random.jsonl --output-details ``` 2) OpenAI-compatible Completions (e.g., vLLM): ```bash python3 -m sglang.bench_serving \ --backend vllm \ --base-url http://127.0.0.1:8000 \ --model meta-llama/Llama-3.1-8B-Instruct \ --dataset-name sharegpt \ --num-prompts 1000 \ --sharegpt-output-len 256 ``` 3) OpenAI-compatible Chat Completions (streaming): ```bash python3 -m sglang.bench_serving \ --backend vllm-chat \ --base-url http://127.0.0.1:8000 \ --model meta-llama/Llama-3.1-8B-Instruct \ --dataset-name random \ --num-prompts 500 \ --apply-chat-template ``` 4) Images (VLM) with chat template: ```bash python3 -m sglang.bench_serving \ --backend sglang \ --host 127.0.0.1 --port 30000 \ --model your-vlm-model \ --dataset-name image \ --image-count 2 \ --image-resolution 720p \ --random-input-len 128 --random-output-len 256 \ --num-prompts 200 \ --apply-chat-template ``` 4a) Images with custom resolution: ```bash python3 -m sglang.bench_serving \ --backend sglang \ --host 127.0.0.1 --port 30000 \ --model your-vlm-model \ --dataset-name image \ --image-count 1 \ --image-resolution 512x768 \ --random-input-len 64 --random-output-len 128 \ --num-prompts 100 \ --apply-chat-template ``` 4b) 1080p images with PNG format and blank content: ```bash python3 -m sglang.bench_serving \ --backend sglang \ --host 127.0.0.1 --port 30000 \ --model your-vlm-model \ --dataset-name image \ --image-count 1 \ --image-resolution 1080p \ --image-format png \ --image-content blank \ --random-input-len 64 --random-output-len 128 \ --num-prompts 100 \ --apply-chat-template ``` 5) Generated shared prefix (long system prompts + short questions): ```bash python3 -m sglang.bench_serving \ --backend sglang \ --host 127.0.0.1 --port 30000 \ --model meta-llama/Llama-3.1-8B-Instruct \ --dataset-name generated-shared-prefix \ --gsp-num-groups 64 --gsp-prompts-per-group 16 \ --gsp-system-prompt-len 2048 --gsp-question-len 128 --gsp-output-len 256 \ --num-prompts 1024 ``` 6) Tokenized prompts (ids) for strict length control (sglang only): ```bash python3 -m sglang.bench_serving \ --backend sglang \ --host 127.0.0.1 --port 30000 \ --model meta-llama/Llama-3.1-8B-Instruct \ --dataset-name random \ --tokenize-prompt \ --random-input-len 2048 --random-output-len 256 --random-range-ratio 0.2 ``` 7) Profiling and cache flush (sglang): ```bash python3 -m sglang.bench_serving \ --backend sglang \ --host 127.0.0.1 --port 30000 \ --model meta-llama/Llama-3.1-8B-Instruct \ --profile \ --flush-cache ``` 8) TensorRT-LLM streaming endpoint: ```bash python3 -m sglang.bench_serving \ --backend trt \ --base-url http://127.0.0.1:8000 \ --model your-trt-llm-model \ --dataset-name random \ --num-prompts 100 \ --disable-ignore-eos ``` 9) Evaluating large-scale KVCache sharing with mooncake trace (sglang only): ```bash python3 -m sglang.bench_serving \ --backend sglang \ --host 127.0.0.1 --port 30000 \ --model model-name \ --dataset-name mooncake \ --mooncake-slowdown-factor 1.0 \ --mooncake-num-rounds 1000 \ --mooncake-workload conversation|mooncake|agent|synthetic --use-trace-timestamps true \ --random-output-len 256 ``` ### Troubleshooting - All requests failed: verify `--backend`, server URL/port, `--model`, and authentication. Check warmup errors printed by the script. - Throughput seems too low: adjust `--request-rate` and `--max-concurrency`; verify server batch size/scheduling; ensure streaming is enabled if appropriate. - Token counts look odd: prefer chat/instruct models with proper chat templates; otherwise tokenization of gibberish may be inconsistent. - Image/MMMU datasets: ensure you installed extra deps (`pillow`, `datasets`, `pybase64`). - Authentication errors (401/403): set `OPENAI_API_KEY` or disable auth on your server. ### Notes - The script raises the file descriptor soft limit (`RLIMIT_NOFILE`) to help with many concurrent connections. - For sglang, `/get_server_info` is queried post-run to report speculative decoding accept length when available. ================================================ FILE: docs/developer_guide/benchmark_and_profiling.md ================================================ # Benchmark and Profiling ## Benchmark SGLang provides four benchmark tools that operate at different levels of the stack. The table below summarizes their key differences: | Tool | HTTP Server | Scheduler | Use Case | | -------------------------- | --------------------------------------------- | --------------------------------------- | -------------------------------------------------------------------------- | | `bench_serving` | Yes (async HTTP client to a running server) | Yes (indirectly, via server) | Realistic online serving benchmarks with latency metrics (TTFT, TPOT, ITL) | | `bench_one_batch_server` | Yes (sends HTTP requests to a running server) | Yes (indirectly, via server) | End-to-end single-batch latency including HTTP and scheduler overhead | | `bench_offline_throughput` | No | Yes (directly uses `Engine` in-process) | Maximum throughput measurement without HTTP overhead | | `bench_one_batch` | No | No (directly calls `ModelRunner`) | Kernel-level latency profiling of a single static batch | Use `bench_serving` by default unless there are specific needs. **`bench_serving`** is an async HTTP load-testing client that sends requests at controlled rates with configurable concurrency to a running server. It measures realistic online serving metrics including time-to-first-token (TTFT), time-per-output-token (TPOT), inter-token latency (ITL), and throughput. Use `num-prompts >= 5 * max-concurrency` to measure steady-state performance. Launch a server with `sglang.launch_server` first. ```bash python3 -m sglang.bench_serving --backend sglang --max-concurrency 16 --num-prompts 80 --random-input-len 256 --random-output-len 32 --dataset-name random ``` **`bench_one_batch_server`** sends a single batch as one HTTP request to a running server. Due to only having a single batch, the server is never in a steady-state and metrics will be biased. Launch a server with `sglang.launch_server` first. ```bash python3 -m sglang.bench_one_batch_server --base-url http://127.0.0.1:30000 --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --batch-size 32 --input-len 256 --output-len 32 ``` **`bench_offline_throughput`** directly instantiates the `Engine` object in-process (no HTTP server) and submits all requests at once via `engine.generate()`. The engine's scheduler handles batching and execution. This measures maximum achievable throughput without any network overhead. ```bash python3 -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --num-prompts 10 ``` **`bench_one_batch`** is the lowest-level tool. It directly instantiates a `ModelRunner` and calls `extend()` / `decode()` on a fixed static batch, bypassing the scheduler entirely. The prefill and decode phases are run separately, making profiling easier but rendering the metrics unrealistic. Because there is no dynamic batching, it may run out of memory for batch sizes that a real server can handle (a real server chunks prefill into smaller batches). This is best suited for profiling individual kernel performance. ```bash python3 -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --batch-size 32 --input-len 256 --output-len 32 ``` ## Profile with PyTorch Profiler [Pytorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) is a convenient basic tool to inspect kernel execution time, call stack, and kernel overlap and occupancy. ### Profile a server with `sglang.bench_serving` ```bash # set trace path export SGLANG_TORCH_PROFILER_DIR=/root/sglang/profile_log # start server python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct # send profiling request from client python -m sglang.bench_serving --backend sglang --model meta-llama/Llama-3.1-8B-Instruct --num-prompts 10 --sharegpt-output-len 100 --profile ``` The `SGLANG_TORCH_PROFILER_DIR` environment variable must be set on both the server and client side; otherwise, the trace file will not be generated correctly. A secure way to do this is by setting it in your shell's resource file (e.g., `~/.bashrc` for bash). For more details, please refer to [Bench Serving Guide](./bench_serving.md). ### Profile In PD Disaggregation Mode When profiling in PD disaggregation mode, prefill and decode workers **must be profiled separately** due to torch profiler limitations. The `bench_serving` command provides dedicated options for this: #### Profile Prefill Workers ```bash # set trace path export SGLANG_TORCH_PROFILER_DIR=/root/sglang/profile_log # start prefill and decode servers (see PD disaggregation docs for setup) python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disaggregation-mode prefill python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disaggregation-mode decode --port 30001 --base-gpu-id 1 # start router python -m sglang_router.launch_router --pd-disaggregation --prefill http://127.0.0.1:30000 --decode http://127.0.0.1:30001 --host 0.0.0.0 --port 8000 # send profiling request targeting prefill workers python -m sglang.bench_serving --backend sglang --model meta-llama/Llama-3.1-8B-Instruct --num-prompts 10 --sharegpt-output-len 100 --profile --pd-separated --profile-prefill-url http://127.0.0.1:30000 ``` #### Profile Decode Workers ```bash # send profiling request targeting decode workers python -m sglang.bench_serving --backend sglang --model meta-llama/Llama-3.1-8B-Instruct --num-prompts 10 --sharegpt-output-len 100 --profile --pd-separated --profile-decode-url http://127.0.0.1:30001 ``` #### Important Notes - `--profile-prefill-url` and `--profile-decode-url` are **mutually exclusive** - you cannot profile both at the same time - Both options support multiple worker URLs for multi-instance setups: ```bash # Profile multiple prefill workers python -m sglang.bench_serving --backend sglang --model meta-llama/Llama-3.1-8B-Instruct --num-prompts 10 --profile --pd-separated --profile-prefill-url http://127.0.0.1:30000 http://127.0.0.1:30002 # Profile multiple decode workers python -m sglang.bench_serving --backend sglang --model meta-llama/Llama-3.1-8B-Instruct --num-prompts 10 --profile --pd-separated --profile-decode-url http://127.0.0.1:30001 http://127.0.0.1:30003 ``` - Make sure `SGLANG_TORCH_PROFILER_DIR` is set on all worker nodes before starting the servers - For more details on setting up PD disaggregation, see [PD Disaggregation Guide](../advanced_features/pd_disaggregation.md) ### Profile a server with `sglang.bench_offline_throughput` ```bash export SGLANG_TORCH_PROFILER_DIR=/root/sglang/profile_log # profile one batch with bench_one_batch.py # batch size can be controlled with --batch argument python3 -m sglang.bench_one_batch --model-path meta-llama/Llama-3.1-8B-Instruct --batch 32 --input-len 1024 --output-len 10 --profile # profile multiple batches with bench_offline_throughput.py python -m sglang.bench_offline_throughput --model-path meta-llama/Llama-3.1-8B-Instruct --dataset-name random --num-prompts 10 --profile --mem-frac=0.8 ``` ### Profile a server with `sglang.profiler` When the server is running (e.g., processing a decoding request), you can start live profiling immediately by sending a profile request to the server. You can do this by running `python3 -m sglang.profiler`. For example: ``` # Terminal 1: Send a generation request python3 -m sglang.test.send_one # Terminal 2: Before the above request finishes, quickly launch the following command in a separate terminal. # It will generate a profile of the above request for several decoding batches. python3 -m sglang.profiler ``` You can also combine the above operations into a single command ``` python3 -m sglang.test.send_one --profile ``` ### Profile a server with HTTP API endpoints SGLang provides HTTP API endpoints to control profiling on a running server. This allows you to start and stop profiling programmatically, which is useful for capturing specific workload patterns. #### Using `/start_profile` endpoint The `/start_profile` endpoint starts profiling on the server. You can control when profiling begins and how long it runs using the following parameters: **Basic usage:** ```bash # Start profiling immediately for 10 steps curl -X POST http://127.0.0.1:30000/start_profile \ -H "Content-Type: application/json" \ -d '{ "num_steps": 10 }' ``` **Parameters:** - `output_dir` (optional): Directory where profile traces will be saved. If not specified, uses `SGLANG_TORCH_PROFILER_DIR` environment variable, or `/tmp` as the default - `num_steps` (optional): Number of steps to profile. If not specified, profiling continues until manually stopped with `/end_profile` - `start_step` (optional): Step number at which to start profiling (inclusive). Useful for skipping warmup iterations - `activities` (optional): List of activities to profile, e.g., `["CPU", "GPU"]`. Default is `["CPU", "GPU"]` - `merge_profiles` (optional): Whether to merge distributed traces. Default is `false` **Note on step ranges:** Profiling starts at `start_step` (inclusive) and continues for `num_steps` iterations. For example, with `start_step=3` and `num_steps=10`, profiling captures steps 3, 4, 5, 6, 7, 8, 9, 10, 11, and 12 (10 steps total, starting from step 3). **Advanced usage with `start_step`:** ```bash # Wait 5 steps (warmup), then profile for 10 steps curl -X POST http://127.0.0.1:30000/start_profile \ -H "Content-Type: application/json" \ -d '{ "output_dir": "/tmp/profiles", "start_step": 5, "num_steps": 10, "activities": ["CPU", "GPU"] }' ``` **Continuous profiling (manual stop):** ```bash # Start profiling without num_steps - must manually stop with /end_profile curl -X POST http://127.0.0.1:30000/start_profile ``` #### Using `/end_profile` endpoint The `/end_profile` endpoint stops an ongoing profiling session and saves the trace file. ```bash # Stop profiling and save traces curl -X POST http://127.0.0.1:30000/end_profile ``` This is only needed when you start profiling without specifying `num_steps`. If `num_steps` is specified, profiling will automatically stop after that many steps. #### Example workflow ```bash # Terminal 1: Start the server export SGLANG_TORCH_PROFILER_DIR=/tmp/profiles python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct # Terminal 2: Start continuous profiling curl -X POST http://127.0.0.1:30000/start_profile \ -H "Content-Type: application/json" \ -d '{ "start_step": 3 }' # Terminal 3: Send requests to generate load python -m sglang.bench_serving --backend sglang --num-prompts 100 # Terminal 2: Stop profiling when done curl -X POST http://127.0.0.1:30000/end_profile ``` ### Profiler Trace Merger for Distributed Traces SGLang now supports automatic merging of profiling traces from distributed setups with multiple parallelism types (TP, DP, PP, EP). This feature is particularly useful for analyzing performance across distributed runs. #### Multi-Node Profiling and Shared Storage Considerations Single-node profiler output merging is completely supported. When profiling in distributed environments spanning multiple nodes, shared storage (e.g., NFS, Lustre) should be accessible by all nodes for the output directory to enable merging of trace files. If there is no shared storage accessible across nodes, automatic merging of trace files during profiling is not supported directly as of now. #### HTTP API Usage ```bash # Start profiling with automatic trace merging enabled curl -X POST /start_profile \ -H "Content-Type: application/json" \ -d '{ "output_dir": "/tmp/profiles", # where to store profile traces "num_steps": 10, "activities": ["CPU", "GPU"], "merge_profiles": true # optional argument to merge profile traces (default=False) }' ``` #### Command Line Usage ```bash # Start profiling with merge enabled python -m sglang.profiler \ --num-steps 10 \ --cpu \ --gpu \ --output-dir /tmp/profiles \ --merge-profiles # optional argument to merge profile traces (default=False) ``` #### Output Files The profile merger generates: - Individual rank trace files: `{profile_id}-TP-{tp}-DP-{dp}-PP-{pp}-EP-{ep}.trace.json.gz` - Merged trace file: `merged-{profile_id}.trace.json.gz` ### Possible PyTorch bugs If in any cases you encounter the following error (for example, using qwen 2.5 VL): ```bash RuntimeError: !stack.empty() INTERNAL ASSERT FAILED at "/pytorch/torch/csrc/autograd/profiler_python.cpp":983, please report a bug to PyTorch. Python replay stack is empty. ``` This is likely a PyTorch Bug reported in [Bug: vLLM Profiler](https://github.com/vllm-project/vllm/issues/18240) and [Bug: torch.profiler.profile](https://github.com/pytorch/pytorch/issues/101632). As a workaround, you may disable `with_stack` with an environment variable such as follows: ```bash export SGLANG_PROFILE_WITH_STACK=False python -m sglang.bench_offline_throughput --model-path meta-llama/Llama-3.1-8B-Instruct --dataset-name random --num-prompts 10 --profile --mem-frac=0.8 ``` ### View traces Trace files can be loaded and visualized from: 1. https://ui.perfetto.dev/ (any browser) 2. chrome://tracing (Chrome browser only) If browser cannot open trace file due to its large size, client can generate a small trace file (<100MB) by controlling number of prompts and lengths of prompt outputs. For example, when profiling a server, ```bash python -m sglang.bench_serving --backend sglang --model meta-llama/Llama-3.1-8B-Instruct --num-prompts 2 --sharegpt-output-len 100 --profile ``` This command sets the number of prompts to 2 with `--num-prompts` argument and limits the length of output sequences to 100 with `--sharegpt-output-len` argument, which can generate a small trace file for browser to open smoothly. Additionally, if you want to locate the SGLang Python source code through the cuda kernel in Trace, you need to disable CUDA Graph when starting the service. This can be done by using the `--disable-cuda-graph` parameter in the command to start the service. ## Profile with Nsight [Nsight systems](https://docs.nvidia.com/nsight-systems/) is an advanced tool that exposes more profiling details, such as register and shared memory usage, annotated code regions and low-level CUDA APIs and events. 1. Prerequisite: Install using apt, or run inside a [NVIDIA Docker container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags) or [SGLang Docker container](https://github.com/sgl-project/sglang/tree/main/docker). ```bash # install nsys # https://docs.nvidia.com/nsight-systems/InstallationGuide/index.html apt update apt install -y --no-install-recommends gnupg echo "deb http://developer.download.nvidia.com/devtools/repos/ubuntu$(source /etc/lsb-release; echo "$DISTRIB_RELEASE" | tr -d .)/$(dpkg --print-architecture) /" | tee /etc/apt/sources.list.d/nvidia-devtools.list apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub apt update apt install nsight-systems-cli ``` 2. To profile a single batch, use ```bash nsys profile --trace-fork-before-exec=true --cuda-graph-trace=node python3 -m sglang.bench_one_batch --model meta-llama/Meta-Llama-3-8B --batch-size 64 --input-len 512 ``` 3. To profile a server, e.g. ```bash # launch the server, set the delay and duration times according to needs # after the duration time has been used up, server will be killed by nsys nsys profile --trace-fork-before-exec=true --cuda-graph-trace=node -o sglang.out --delay 60 --duration 70 python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disable-radix-cache # client python3 -m sglang.bench_serving --backend sglang --num-prompts 1000 --dataset-name random --random-input 1024 --random-output 512 ``` In practice, we recommend users to set `--duration` argument to a large value. Whenever user wants the server to stop profiling. Firstly run: ```bash nsys sessions list ``` to get the session id in the form of `profile-XXXXX`, then run: ```bash nsys stop --session=profile-XXXXX ``` to manually kill the profiler and generate `nsys-rep` files instantly. 4. Use NVTX to annotate code regions, e.g. to see their execution time. ```bash # install nvtx pip install nvtx ``` ```python # code snippets import nvtx with nvtx.annotate("description", color="color"): # some critical code ``` ### Layer-wise NVTX Profiling with Nsight Systems SGLang provides built-in layerwise NVTX annotations that can be combined with the CUDA Profiler for detailed per-layer profiling in Nsight Systems. This is particularly useful for identifying performance bottlenecks at the layer level. #### Using `--enable-layerwise-nvtx-marker` with Nsight Systems and `/start_profile` The `--enable-layerwise-nvtx-marker` flag automatically adds NVTX markers to every layer in your model. This is particularly powerful when combined with Nsight Systems profiling to see detailed per-layer performance. **Method 1: Using `/start_profile` with CUDA_PROFILER (for programmatic control)** This method allows you to control exactly when profiling starts/stops via HTTP API while Nsight Systems is running. 1. Launch the server with layerwise NVTX enabled under Nsight Systems: ```bash # Terminal 1: Start server with nsys and capture-range option nsys profile --trace-fork-before-exec=true \ --cuda-graph-trace=node \ --capture-range=cudaProfilerApi \ --capture-range-end=stop \ -o layerwise_profile \ python -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --enable-layerwise-nvtx-marker \ --disable-cuda-graph ``` Note: NVTX markers are not emitted for kernel launches captured by CUDA graphs. Use `--disable-cuda-graph` to ensure all layerwise NVTX markers are emitted in the trace. 2. In another terminal, control profiling via `/start_profile` with `CUDA_PROFILER` activity: ```bash # Terminal 2: Wait for server to be ready, then start CUDA profiling # Wait 3 steps for warmup, then profile for 10 steps curl -X POST http://127.0.0.1:30000/start_profile \ -H "Content-Type: application/json" \ -d '{ "start_step": 3, "num_steps": 10, "activities": ["CUDA_PROFILER"] }' ``` 3. Send requests to generate load: ```bash # Terminal 3: Generate workload python -m sglang.bench_serving --backend sglang --num-prompts 100 ``` 4. Profiling will automatically stop after 10 steps (due to `num_steps: 10`). If you hadn't specified `num_steps`, you would need to manually stop it: ```bash # Terminal 2: Only needed if num_steps was not specified curl -X POST http://127.0.0.1:30000/end_profile ``` The `--capture-range=cudaProfilerApi` option tells Nsight Systems to only capture data between `cudaProfilerStart()` and `cudaProfilerStop()` calls (triggered by `/start_profile` and `/end_profile`), reducing overhead and file size. The `start_step` parameter skips the first 3 steps to avoid capturing warmup overhead. **Method 2: Simpler approach without `/start_profile` API** For simpler use cases where you don't need fine-grained control over profiling start/stop, you can profile with Nsight Systems capturing the entire workload: ```bash # Terminal 1: Start server with layerwise NVTX # Note: --disable-cuda-graph ensures all NVTX markers are emitted python -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --enable-layerwise-nvtx-marker \ --disable-cuda-graph # Terminal 2: Profile the benchmarking client nsys profile --trace-fork-before-exec=true \ --cuda-graph-trace=node \ -o layerwise_profile \ python -m sglang.bench_serving --backend sglang --num-prompts 10 ``` This approach profiles the entire client execution, including all server interactions. The layerwise NVTX markers will be visible in the Nsight Systems timeline. **Viewing the profiling results:** Open the generated `.qdrep` file with Nsight Systems: ```bash nsys-ui layerwise_profile.qdrep ``` In the Nsight Systems GUI, you'll see: - **NVTX ranges**: Each layer appears as a labeled range in the timeline with detailed information in the marker metadata - **CUDA kernels**: All GPU kernels are shown alongside the layer annotations - **Layer hierarchy**: The full module path (e.g., `meta-llama/Meta-Llama-3.1-8B-Instruct.model.layers.0.self_attn.qkv_proj`) helps identify specific layers. The prefix uses the full model path from `--model-path`. - **Tensor shapes**: Input/output dimensions and parameter shapes are included in the NVTX marker data **Benefits of layerwise NVTX profiling:** - **Granular visibility**: See exactly which layers are taking the most time - **Memory tracking**: Identify layers with large memory allocations - **Bottleneck identification**: Quickly locate inefficient operations - **Communication overhead**: In multi-GPU setups, see per-layer communication costs - **Development debugging**: Validate that model architecture changes have the expected performance impact ## Other tips 1. You can benchmark a model using dummy weights by only providing the config.json file. This allows for quick testing of model variants without training. To do so, add `--load-format dummy` to the above commands and then you only need a correct `config.json` under the checkpoint folder. 2. You can benchmark a model with modified configs (e.g., less layers) by using `--json-model-override-args`. For example, you can benchmark a model with only 2 layers and 2 kv heads using: ```bash python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --batch 32 --input-len 256 --output-len 32 --load-format dummy --json-model-override-args '{"num_hidden_layers": 1, "num_key_value_heads": 1}' ``` 3. You can use `--python-backtrace=cuda` to see python call stack for all CUDA kernels, as in PyTorch Profiler. (Caveat: this can cause inaccurately long kernel runtimes for CUDA event based timing) 4. For more arguments see [Nsight Systems User Guide](https://docs.nvidia.com/nsight-systems/UserGuide/index.html). ================================================ FILE: docs/developer_guide/contribution_guide.md ================================================ # Contribution Guide Welcome to **SGLang**! We appreciate your interest in contributing. This guide provides a concise overview of how to set up your environment, run tests, build documentation, and open a Pull Request (PR). Whether you’re fixing a small bug or developing a major feature, we encourage following these steps for a smooth contribution process. ## Install SGLang from Source ### Fork and clone the repository **Note**: New contributors do **not** have the write permission to push to the official SGLang repo. Please fork the repository under your GitHub account, then clone your fork locally. ```bash git clone https://github.com//sglang.git ``` ### Build from source Refer to [Install SGLang from Source](../get_started/install.md#method-2-from-source). ## Format code with pre-commit We use [pre-commit](https://pre-commit.com/) to maintain consistent code style checks. Before pushing your changes, please run: ```bash pip3 install pre-commit pre-commit install pre-commit run --all-files ``` - **`pre-commit run --all-files`** manually runs all configured checks, applying fixes if possible. If it fails the first time, re-run it to ensure lint errors are fully resolved. Make sure your code passes all checks **before** creating a Pull Request. - **Do not commit** directly to the `main` branch. Always create a new branch (e.g., `feature/my-new-feature`), push your changes, and open a PR from that branch. ## Run and add unit tests If you add a new feature or fix a bug, please add corresponding unit tests to ensure coverage and prevent regression. SGLang uses Python's built-in [unittest](https://docs.python.org/3/library/unittest.html) framework with [pytest](https://docs.pytest.org/) as the test runner. ### Unit tests (no server required) Unit tests live under [`test/registered/unit/`](https://github.com/sgl-project/sglang/tree/main/test/registered/unit), organized to mirror the `python/sglang/srt/` source tree. These tests validate component logic **without** launching a server or loading real model weights. **When to add a unit test:** If you modify a file under `python/sglang/srt/`, check whether a corresponding test exists in `test/registered/unit/` and add coverage for your changes. For example: ``` srt/mem_cache/radix_cache.py → unit/mem_cache/test_radix_cache.py srt/sampling/sampling_params.py → unit/sampling/test_sampling_params.py ``` **Run unit tests locally:** ```bash pytest test/registered/unit/ -v # all unit tests pytest test/registered/unit/mem_cache/ -v # one module ``` **Run with coverage:** ```bash pytest test/registered/unit/ --cov --cov-config=.coveragerc -v ``` For conventions on CI registration, test structure, and examples, see [`test/registered/unit/README.md`](https://github.com/sgl-project/sglang/tree/main/test/registered/unit/README.md). ### E2E tests (server required) For tests that require launching a server, refer to [`test/registered/README.md`](https://github.com/sgl-project/sglang/tree/main/test/registered/README.md) for guidance on where to place your test. For detailed instructions on running tests and integrating them into CI, refer to [test/README.md](https://github.com/sgl-project/sglang/tree/main/test/README.md). ## Write documentations We recommend new contributors start from writing documentation, which helps you quickly understand SGLang codebase. For more details, please refer to [docs/README.md](https://github.com/sgl-project/sglang/tree/main/docs/README.md). ## Test the accuracy If your code changes the model output, please run the accuracy tests. A quick sanity check is the few-shot GSM8K. ``` # Launch a server python3 -m sglang.launch_server --model Qwen/Qwen2-7B-Instruct # Evaluate python3 -m sglang.test.few_shot_gsm8k --num-questions 200 ``` Please note that the above script is primarily a sanity check, not a rigorous accuracy or speed test. This test can have significant variance (1%–5%) in accuracy due to batching and the non-deterministic nature of the inference engine. Also, do not rely on the "Latency/Output throughput" from this script, as it is not a proper speed test. GSM8K is too easy for state-of-the-art models nowadays. Please try your own more challenging accuracy tests. You can find additional accuracy eval examples in: - [test_eval_accuracy_large.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_eval_accuracy_large.py) - [test_gpt_oss_1gpu.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_gpt_oss_1gpu.py) ## Benchmark the speed Refer to [Benchmark and Profiling](../developer_guide/benchmark_and_profiling.md). ## Requesting a review for merge You can follow the pull request merge process described in [MAINTAINER.md](https://github.com/sgl-project/sglang/blob/main/.github/MAINTAINER.md). You will need to work with the Merge Oncall, Codeowner, and other reviewers to get their approvals. Then your PR can be merged. ## How to Trigger CI Tests We have a lot of open PRs but limited CI machines, so only top and trusted contributors have permission to trigger CI tests. Users with permission are listed in the [CI_PERMISSIONS.json](https://github.com/sgl-project/sglang/blob/main/.github/CI_PERMISSIONS.json) **PR authors** can always use `/rerun-failed-ci` on their own PRs, even if they are not listed in `CI_PERMISSIONS.json`. For CI to run on a pull request, it must have the "run-ci" label. Authorized users can add the label or rerun failed tests by commenting on the PR with one of these commands: - `/tag-run-ci-label`: Adds the "run-ci" label. Every future commit will trigger CI. - `/rerun-failed-ci`: Reruns the failed or flaky tests from the most recent commit. - `/tag-and-rerun-ci`: A single command that performs both `/tag-run-ci-label` and `/rerun-failed-ci`. - `/rerun-stage `: Reruns a specific test stage without waiting for its dependencies. This is useful when you want to quickly validate a fix for a specific test failure instead of waiting ~30 minutes for preceding stages to complete. If you have permission, the [Slash Command Handler](https://github.com/sgl-project/sglang/actions/workflows/slash-command-handler.yml) will run your command and react with a 👍 to your comment. It may take up to a few minutes for the reaction to appear. Here’s a usage [example](https://github.com/sgl-project/sglang/pull/14253#issuecomment-3599509302). To avoid spamming a PR with too many `/rerun-failed-ci` comments, you can also trigger the command by editing an existing comment and adding any suffix (e.g., `/rerun-failed-ci try again`). Example of rerunning a single test stage: `/rerun-stage unit-test-backend-4-gpu`. If you don’t have permission and you’re not the PR author, please ask maintainers to trigger CI for you. ### CI rate limits Due to CI scheduling and limited resources, higher-priority PRs may preempt running jobs. In such cases, you may need to rerun the tests. We apply CI rate limits to prevent abuse and ensure fair usage of our CI resources. Each CI workflow has a default limit defined in its workflow configuration file. For example, in [pr-gate.yml](https://github.com/sgl-project/sglang/blob/main/.github/workflows/pr-gate.yml), the default cooldown period is 120 minutes, and each workflow can override it via the `cool-down-minutes` input parameter: ```yaml cool-down-minutes: description: "Default cooldown period in minutes; 0 disables rate limiting" type: number default: 120 ``` Users listed in [CI_PERMISSIONS.json](https://github.com/sgl-project/sglang/blob/main/.github/CI_PERMISSIONS.json) may have a per-user cooldown interval. In practice, we use the minimum of the workflow’s default window and the user-specific interval. ## Code style guidance - Avoid code duplication. If the same code snippet (more than five lines) appears multiple times, extract it into a shared function. - Minimize device synchronization. Reduce expensive CPU-GPU synchronization operations, such as `tensor.item()` or `tensor.cpu()`, whenever possible. Use vectorized code. - Prioritize extreme efficiency. SGLang is a runtime, and most of your code runs on the critical path for every request. Optimize all minor overheads as much as possible, especially in the model forward code. - A common pattern is some runtime checks in the model forward pass (e.g., [this](https://github.com/sgl-project/sglang/blob/f1b0eda55c2c4838e8ab90a0fac7fb1e3d7064ab/python/sglang/srt/models/deepseek_v2.py#L486-L491)). These are very likely the same for every layer. Please cache the result as a single boolean value whenever possible. - Make functions as pure as possible. Avoid in-place modification of arguments. - Keep files concise. If a file exceeds 2,000 lines of code, split it into multiple smaller files. (e.g., `scheduler.py`, `scheduler_output_processor_mixin.py`) - Keep tests run fast. - If a single test file run longer than 500 seconds, split it into multiple smaller files (e.g., `test_eagle_infer_a.py`, `test_eagle_infer_b.py`). - If a single job in a github workflow runs longer than 30 mins, split it into smaller jobs/steps. - Reuse server launches in your unit tests to make tests run faster. - When supporting new hardware or features, follow these guidelines: - Do not drastically change existing code. - Always prefer new files to introduce specific components for your new hardware (e.g., `allocator_ascend.py`). - If you write multiple if/else blocks for new features, ensure the common path (e.g., NVIDIA hardware or the existing code path) is the first branch. ## How to update sgl-kernel Since sglang and the `sglang-kernel` (prior `sgl-kernel`) distribution are separate Python packages, our current GitHub CI infrastructure does not support updating a kernel and using it immediately within the same pull request (PR). To add a new kernel or modify an existing one in the `sgl-kernel/` source tree, you must use multiple PRs. Follow these steps: 1. Submit a PR to update the sgl-kernel source code without using it in sglang python package (e.g., [#8884](https://github.com/sgl-project/sglang/pull/8884/files)). 2. Bump the version of the kernel package (e.g., [#9220](https://github.com/sgl-project/sglang/pull/9220/files)). - Once merged, this will trigger an automatic release of the `sglang-kernel` wheel to PyPI. - If not urgent, you can wait for other people to release the wheel. A new version will typically be released within one week. 3. Apply the changes: - Update the `sglang-kernel` version in `sglang/python/pyproject.toml` to use the modified kernels. - Update the related caller code in the sglang to use the new kernel. ## Tips for newcomers If you want to contribute but don’t have a specific idea in mind, pick issues labeled [“good first issue” or “help wanted”](https://github.com/sgl-project/sglang/issues?q=is%3Aissue+label%3A%22good+first+issue%22%2C%22help+wanted%22). These tasks typically have lower complexity and provide an excellent introduction to the codebase. Also check out this [code walk-through](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/tree/main/sglang/code-walk-through) for a deeper look into SGLang’s workflow. If you have any questions or want to start a discussion, please feel free to ask in our [Slack channel](https://slack.sglang.io). Thank you for your interest in SGLang. Happy coding! ================================================ FILE: docs/developer_guide/development_guide_using_docker.md ================================================ # Development Guide Using Docker ## Setup VSCode on a Remote Host (Optional - you can skip this step if you plan to run sglang dev container locally) 1. In the remote host, download `code` from [Https://code.visualstudio.com/docs/?dv=linux64cli](https://code.visualstudio.com/download) and run `code tunnel` in a shell. Example ```bash wget https://vscode.download.prss.microsoft.com/dbazure/download/stable/fabdb6a30b49f79a7aba0f2ad9df9b399473380f/vscode_cli_alpine_x64_cli.tar.gz tar xf vscode_cli_alpine_x64_cli.tar.gz # https://code.visualstudio.com/docs/remote/tunnels ./code tunnel ``` 2. In your local machine, press F1 in VSCode and choose "Remote Tunnels: Connect to Tunnel". ## Setup Docker Container ### Option 1. Use the default dev container automatically from VSCode There is a `.devcontainer` folder in the sglang repository root folder to allow VSCode to automatically start up within dev container. You can read more about this VSCode extension in VSCode official document [Developing inside a Container](https://code.visualstudio.com/docs/devcontainers/containers). ![image](https://github.com/user-attachments/assets/6a245da8-2d4d-4ea8-8db1-5a05b3a66f6d) (*Figure 1: Diagram from VSCode official documentation [Developing inside a Container](https://code.visualstudio.com/docs/devcontainers/containers).*) To enable this, you only need to: 1. Start Visual Studio Code and install [VSCode dev container extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers). 2. Press F1, type and choose "Dev Container: Open Folder in Container. 3. Input the `sglang` local repo path in your machine and press enter. The first time you open it in dev container might take longer due to docker pull and build. Once it's successful, you should set on your status bar at the bottom left displaying that you are in a dev container: ![image](https://github.com/user-attachments/assets/650bba0b-c023-455f-91f9-ab357340106b) Now when you run `sglang.launch_server` in the VSCode terminal or start debugging using F5, sglang server will be started in the dev container with all your local changes applied automatically: ![image](https://github.com/user-attachments/assets/748c85ba-7f8c-465e-8599-2bf7a8dde895) ### Option 2. Start up containers manually (advanced) The following startup command is an example for internal development by the SGLang team. You can **modify or add directory mappings as needed**, especially for model weight downloads, to prevent repeated downloads by different Docker containers. ❗️ **Note on RDMA** 1. `--network host` and `--privileged` are required by RDMA. If you don't need RDMA, you can remove them but keeping them there does not harm. Thus, we enable these two flags by default in the commands below. 2. You may need to set `NCCL_IB_GID_INDEX` if you are using RoCE, for example: `export NCCL_IB_GID_INDEX=3`. ```bash # Change the name to yours docker run -itd --shm-size 32g --gpus all -v --ipc=host --network=host --privileged --name sglang_dev lmsysorg/sglang:dev /bin/zsh docker exec -it sglang_dev /bin/zsh ``` Some useful volumes to mount are: 1. **Huggingface model cache**: mounting model cache can avoid re-download every time docker restarts. Default location on Linux is `~/.cache/huggingface/`. 2. **SGLang repository**: code changes in the SGLang local repository will be automatically synced to the .devcontainer. Example 1: Mounting local cache folder `/opt/dlami/nvme/.cache` but not the SGLang repo. Use this when you prefer to manually transfer local code changes to the devcontainer. ```bash docker run -itd --shm-size 32g --gpus all -v /opt/dlami/nvme/.cache:/root/.cache --ipc=host --network=host --privileged --name sglang_zhyncs lmsysorg/sglang:dev /bin/zsh docker exec -it sglang_zhyncs /bin/zsh ``` Example 2: Mounting both HuggingFace cache and local SGLang repo. Local code changes are automatically synced to the devcontainer as the SGLang is installed in editable mode in the dev image. ```bash docker run -itd --shm-size 32g --gpus all -v $HOME/.cache/huggingface/:/root/.cache/huggingface -v $HOME/src/sglang:/sgl-workspace/sglang --ipc=host --network=host --privileged --name sglang_zhyncs lmsysorg/sglang:dev /bin/zsh docker exec -it sglang_zhyncs /bin/zsh ``` ## Debug SGLang with VSCode Debugger 1. (Create if not exist) open `launch.json` in VSCode. 2. Add the following config and save. Please note that you can edit the script as needed to apply different parameters or debug a different program (e.g. benchmark script). ```JSON { "version": "0.2.0", "configurations": [ { "name": "Python Debugger: launch_server", "type": "debugpy", "request": "launch", "module": "sglang.launch_server", "console": "integratedTerminal", "args": [ "--model-path", "meta-llama/Llama-3.2-1B", "--host", "0.0.0.0", "--port", "30000", "--trust-remote-code", ], "justMyCode": false } ] } ``` 3. Press "F5" to start. VSCode debugger will ensure that the program will pause at the breakpoints even if the program is running at remote SSH/Tunnel host + dev container. ## Profile ```bash # Change batch size, input, output and add `disable-cuda-graph` (for easier analysis) # e.g. DeepSeek V3 nsys profile -o deepseek_v3 python3 -m sglang.bench_one_batch --batch-size 1 --input 128 --output 256 --model deepseek-ai/DeepSeek-V3 --trust-remote-code --tp 8 --disable-cuda-graph ``` ## Evaluation ```bash # e.g. gsm8k 8 shot python3 benchmark/gsm8k/bench_sglang.py --num-questions 2000 --parallel 2000 --num-shots 8 ``` ================================================ FILE: docs/developer_guide/development_jit_kernel_guide.md ================================================ # Development Guide for JIT Kernels ## Environment Setup We strongly recommend using `clangd` as the language server for JIT kernel development. For Ubuntu/Debian, you can download clangd from [apt.llvm.org](https://apt.llvm.org/). If you are using VS Code, we recommend installing the `clangd` extension for better IDE integration. All JIT-related files are located in `python/sglang/jit_kernel`. Unlike `sgl-kernel`, which compiles CUDA/C++ binaries ahead of time (AOT), just-in-time (JIT) kernels are compiled at runtime. Consequently, a static `compile_commands.json` cannot be generated. To enable code completion with `clangd`, run `python -m sglang.jit_kernel` to generate a `.clangd` configuration file in your current directory. After generating the file, restart the clangd language server. It should now recognize all JIT kernel files. ## Code Structure ### C++ Implementation C++ source code is located in `python/sglang/jit_kernel/csrc`. Reusable functions should be placed in `python/sglang/jit_kernel/include`. We use [tvm-ffi](https://github.com/apache/tvm-ffi) for efficient foreign language bindings. Refer to the [documentation](https://tvm.apache.org/ffi/) for advanced usage, such as exporting C++ objects. Typically, `tvm::ffi::TensorView` is sufficient for passing PyTorch Tensors from Python. ### Python Interface Python interfaces are defined in `python/sglang/jit_kernel`. The `load_jit` utility function in `python/sglang/jit_kernel/utils.py` loads and returns the compiled module. To export a C++ function (e.g., `cpp_func`), pass `cuda_wrappers=[("func", "cpp_func")]` to `load_jit`. The function can then be called in Python as `module.func`. For caching compiled modules, prefer `sglang.jit_kernel.utils.cache_once` over `functools.lru_cache`. `functools.lru_cache` is not compatible with `torch.compile`. ### C++ Utilities The following C++ utilities are available: #### Integer Range Similar to PyTorch, we provide an `irange` function to represent an integer range. ```C++ #include void test() { for (auto i : host::irange(100)) { // [0, 100) // do something } for (auto i : host::irange(0, 100)) { // [0, 100) // do something } } ``` #### Runtime Checking `RuntimeCheck` validates conditions at runtime. It accepts optional arguments for error reporting. If the check fails, these arguments are output to aid debugging. `RuntimeDeviceCheck` verifies the status of the last kernel launch. ```C++ #include #include void test() { host::RuntimeCheck(1 + 1 == 2, 1 + 1, " != ", 2); host::RuntimeDeviceCheck(); // check the provided `cudaError_t` host::RuntimeDeviceCheck(cudaGetLastError()); } ``` #### Tensor Checking `TensorMatcher` provides a readable way to validate and extract tensor shape information. ```cpp #include void test(const tvm::ffi::TensorView k_cache, const tvm::ffi::TensorView v_cache) { using namespace host; auto D = SymbolicSize{"D"}; // cache dimension auto N = SymbolicSize{"N"}; // kvcache stride auto dtype = SymbolicDType{}; auto device = SymbolicDevice{}; TensorMatcher({-1, D}) // .with_strides({N, 1}) .with_dtype(dtype) .with_device(device) .verify(k_cache) .verify(v_cache); } ``` Configure the `TensorMatcher` with expected stride, dtype, and device properties before verification. - If `with_strides` is omitted, the tensor is expected to be contiguous. - Template arguments in `with_dtype` restrict the allowed data types. - Template arguments in `with_device` restrict the allowed devices. - Values passed to `with_xxx` methods enforce equality checks. - Passing `-1` for size or stride allows matching any value. A `Symbolic` variable must resolve to the same value across all verifications. Use `.unwrap()` to retrieve the matched value after verification. > Note: `TensorMatcher` is a temporary expression and should not be stored in a variable. > Tip: Add `//` at the end of the `TensorMatcher` chain to enforce proper indentation. #### Kernel Launching `LaunchKernel::resolve_device` retrieves the current `cudaStream` from PyTorch. Kernels can also be launched directly using `LaunchKernel`. ```cpp #include #include __global__ void kernel() {} void test() { const auto num_blocks = 1; const auto num_threads = 32; const auto dynamic_smem = 0; DLDevice dev; // suppose this is initialized properly host::LaunchKernel(num_blocks, num_threads, dev)(kernel); cudaStream_t stream = host::LaunchKernel::resolve_device(dev); host::LaunchKernel(num_blocks, num_threads, stream, dynamic_smem)(kernel); } ``` ## Add new kernels This section walks through a complete, end-to-end example of adding a new JIT kernel to the system. We use a simple add_constant kernel as a running example, which adds a constant integer value to every element of an input tensor. Conceptually, the Python interface looks like this: ```python def add_constant(src: torch.Tensor, c: int): return src + c ``` ### STEP 1: Write the C++ kernel Write your CUDA kernel in [jit_kernel/csrc/add_constant.cuh](../../python/sglang/jit_kernel/csrc/add_constant.cuh). For demonstration purposes, we pass the constant value as a template parameter. ```cpp #include // For TensorMatcher, SymbolicSize, SymbolicDevice #include // For LaunchKernel #include // For div_ceil, RuntimeCheck #include #include #include #include namespace { template __global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t length) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < length) { dst[idx] = src[idx] + kConstant; } } constexpr size_t kBlockSize = 256; // You can also use struct with static method as an alternative template void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { using namespace host; // 1. Validate input tensors SymbolicSize N = {"num_elements"}; SymbolicDevice device_; TensorMatcher({N}) // 1D tensor, must be contiguous .with_dtype() // must be int32 .with_device(device_) // must be on CUDA device .verify(dst) // check tensor dst .verify(src); // check tensor src // 2. Extract required parameters, prepare for kernel launch const size_t num_elements = N.unwrap(); const size_t grid_size = div_ceil(num_elements, kBlockSize); const DLDevice device = device_.unwrap(); // some extra runtime checks using host::RuntimeCheck RuntimeCheck(num_elements > 0, "We only support non-empty tensors, got num_elements = ", num_elements); // 3. Launch the kernel. Error code will be automatically checked. LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)( // kernel function add_constant_kernel, // kernel arguments static_cast(dst.data_ptr()), static_cast(src.data_ptr()), num_elements); } } // namespace ``` ### STEP 2: Create Python Interfaces Next, expose the kernel through a Python wrapper. Create a new file at [jit_kernel/add_constant.py](../../python/sglang/jit_kernel/add_constant.py) and expose the needed interfaces. ```python from __future__ import annotations from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_add_constant_module(constant: int) -> Module: args = make_cpp_args(constant) # pass all the template argument return load_jit( "add_constant", *args, cuda_files=["add_constant.cuh"], cuda_wrappers=[("add_constant", f"add_constant<{args}>")], ) def add_constant(src: torch.Tensor, constant: int) -> torch.Tensor: if not src.is_cuda: raise RuntimeError("src must be a CUDA tensor") if src.dtype != torch.int32: raise RuntimeError(f"Unsupported dtype {src.dtype}. Supported: int32") dst = torch.empty_like(src) module = _jit_add_constant_module(constant) module.add_constant(dst, src) return dst ``` Keep the Python wrapper thin, but still validate the basic invariants such as device and dtype before dispatch. In the current JIT/FFI path, invalid tensors are not always rejected safely before launch. ### STEP 3: Use your kernel Finally, import and use the kernel like a regular Python function: ```python from sglang.jit_kernel.add_constant import add_constant ``` For a complete, runnable example, refer to [test_add_constant.py](../../python/sglang/jit_kernel/tests/test_add_constant.py). ## C++ Include Library Reference The JIT kernel framework provides a set of reusable C++ headers in `python/sglang/jit_kernel/include/sgl_kernel/`. Each header is designed to be lightweight and self-contained. Below is a summary of each header and its key APIs. ### Core Utilities | Header | Namespace | Purpose | |--------|-----------|---------| | `utils.h` | `host` | Host-side essentials: `RuntimeCheck`, `Panic`, `div_ceil`, `irange` | | `utils.cuh` | `device` / `host` | Type aliases (`fp16_t`, `bf16_t`, ...), `SGL_DEVICE` macro, PDL helpers, `LaunchKernel`, `RuntimeDeviceCheck` | | `source_location.h` | (global) | Portable `std::source_location` wrapper for error reporting | | `runtime.cuh` | `host::runtime` | CUDA runtime queries: `get_blocks_per_sm`, `get_sm_count`, `get_cc_major`, `get_runtime_version`, `get_available_dynamic_smem_per_block` | ### Tensor Validation | Header | Namespace | Purpose | |--------|-----------|---------| | `tensor.h` | `host` | `TensorMatcher`, `SymbolicSize`, `SymbolicDType`, `SymbolicDevice` | ### Math & Type System | Header | Namespace | Purpose | |--------|-----------|---------| | `math.cuh` | `device::math` | `max`, `min`, `abs`, `sqrt`, `rsqrt`, `exp`, `sin`, `cos`, constants | | `type.cuh` | (global) / `device` | `dtype_trait`, `packed_t`, `device::cast(from)` | ### Memory Access | Header | Namespace | Purpose | |--------|-----------|---------| | `vec.cuh` | `device` | `AlignedVector` - vectorized load/store (up to 128-bit; 256-bit requires Blackwell GPUs) | | `tile.cuh` | `device::tile` | `Memory` - cooperative tiled memory I/O (thread/warp/CTA) | ### Parallel Primitives | Header | Namespace | Purpose | |--------|-----------|---------| | `warp.cuh` | `device::warp` | `reduce_sum`, `reduce_max` via `__shfl_xor_sync` | | `cta.cuh` | `device::cta` | `reduce_max` across warps via shared memory | | `atomic.cuh` | `device::atomic` | `max` - atomic float max (CUDA + ROCm fallback) | ### Reusable Kernel Templates | Header | Namespace | Purpose | |--------|-----------|---------| | `impl/norm.cuh` | `host::norm` / `device::norm` | RMSNorm building blocks (warp & CTA paths, `StorageType`) | ================================================ FILE: docs/developer_guide/evaluating_new_models.md ================================================ # Evaluating New Models with SGLang This document provides commands for evaluating models' accuracy and performance. Before open-sourcing new models, we strongly suggest running these commands to verify whether the score matches your internal benchmark results. **For cross verification, please submit commands for installation, server launching, and benchmark running with all the scores and hardware requirements when open-sourcing your models.** [Reference: MiniMax M2](https://github.com/sgl-project/sglang/pull/12129) ## Accuracy ### LLMs SGLang provides built-in scripts to evaluate common benchmarks. **MMLU** ```bash python -m sglang.test.run_eval \ --eval-name mmlu \ --port 30000 \ --num-examples 1000 \ --max-tokens 8192 ``` **GSM8K** ```bash python -m sglang.test.few_shot_gsm8k \ --host 127.0.0.1 \ --port 30000 \ --num-questions 200 \ --num-shots 5 ``` **HellaSwag** ```bash python benchmark/hellaswag/bench_sglang.py \ --host 127.0.0.1 \ --port 30000 \ --num-questions 200 \ --num-shots 20 ``` **GPQA** ```bash python -m sglang.test.run_eval \ --eval-name gpqa \ --port 30000 \ --num-examples 198 \ --max-tokens 120000 \ --repeat 8 ``` ```{tip} For reasoning models, add `--thinking-mode ` (e.g., `qwen3`, `deepseek-v3`). You may skip it if the model has forced thinking enabled. ``` **HumanEval** ```bash pip install human_eval python -m sglang.test.run_eval \ --eval-name humaneval \ --num-examples 10 \ --port 30000 ``` ### VLMs **MMMU** ```bash python benchmark/mmmu/bench_sglang.py \ --port 30000 \ --concurrency 64 ``` ```{tip} You can set max tokens by passing `--extra-request-body '{"max_tokens": 4096}'`. ``` For models capable of processing video, we recommend extending the evaluation to include `VideoMME`, `MVBench`, and other relevant benchmarks. ## Performance Performance benchmarks measure **Latency** (Time To First Token - TTFT) and **Throughput** (tokens/second). ### LLMs **Latency-Sensitive Benchmark** This simulates a scenario with low concurrency (e.g., single user) to measure latency. ```bash python -m sglang.bench_serving \ --backend sglang \ --host 0.0.0.0 \ --port 30000 \ --dataset-name random \ --num-prompts 10 \ --max-concurrency 1 ``` **Throughput-Sensitive Benchmark** This simulates a high-traffic scenario to measure maximum system throughput. ```bash python -m sglang.bench_serving \ --backend sglang \ --host 0.0.0.0 \ --port 30000 \ --dataset-name random \ --num-prompts 1000 \ --max-concurrency 100 ``` **Single Batch Performance** You can also benchmark the performance of processing a single batch offline. ```bash python -m sglang.bench_one_batch_server \ --model \ --batch-size 8 \ --input-len 1024 \ --output-len 1024 ``` You can run more granular benchmarks: - **Low Concurrency**: `--num-prompts 10 --max-concurrency 1` - **Medium Concurrency**: `--num-prompts 80 --max-concurrency 16` - **High Concurrency**: `--num-prompts 500 --max-concurrency 100` ## Reporting Results For each evaluation, please report: 1. **Metric Score**: Accuracy % (LLMs and VLMs); Latency (ms) and Throughput (tok/s) (LLMs only). 2. **Environment settings**: GPU type/count, SGLang commit hash. 3. **Launch configuration**: Model path, TP size, and any special flags. 4. **Evaluation parameters**: Number of shots, examples, max tokens. ================================================ FILE: docs/developer_guide/release_process.md ================================================ # PyPI Package Release Process ## Update the version in code Update the package version in `python/pyproject.toml` and `python/sglang/__init__.py`. ## Upload the PyPI package ``` pip install build twine ``` ``` cd python bash upload_pypi.sh ``` ## Make a release in GitHub Make a new release https://github.com/sgl-project/sglang/releases/new. ================================================ FILE: docs/developer_guide/setup_github_runner.md ================================================ # Set Up Self-Hosted Runners for GitHub Actions ## Add a Runner ### Step 1: Start a docker container. **You can mount a folder for the shared huggingface model weights cache. ** The command below uses `/tmp/huggingface` as an example. ``` docker pull nvidia/cuda:12.9.1-devel-ubuntu22.04 # Nvidia docker run --shm-size 128g -it -v /tmp/huggingface:/hf_home --gpus all nvidia/cuda:12.9.1-devel-ubuntu22.04 /bin/bash # AMD docker run --rm --device=/dev/kfd --device=/dev/dri --group-add video --shm-size 128g -it -v /tmp/huggingface:/hf_home lmsysorg/sglang:v0.5.8-rocm700-mi30x /bin/bash # AMD just the last 2 GPUs docker run --rm --device=/dev/kfd --device=/dev/dri/renderD176 --device=/dev/dri/renderD184 --group-add video --shm-size 128g -it -v /tmp/huggingface:/hf_home lmsysorg/sglang:v0.5.8-rocm700-mi30x /bin/bash ``` ### Step 2: Configure the runner by `config.sh` Run these commands inside the container. ``` apt update && apt install -y curl python3-pip git pip install --upgrade pip export RUNNER_ALLOW_RUNASROOT=1 ``` Then follow https://github.com/sgl-project/sglang/settings/actions/runners/new?arch=x64&os=linux to run `config.sh` **Notes** - Do not need to specify the runner group - Give it a name (e.g., `test-sgl-gpu-0`) and some labels (e.g., `1-gpu-runner`). The labels can be edited later in Github Settings. - Do not need to change the work folder. ### Step 3: Run the runner by `run.sh` - Set up environment variables ``` export HF_HOME=/hf_home export SGLANG_IS_IN_CI=true export HF_TOKEN=hf_xxx export OPENAI_API_KEY=sk-xxx export CUDA_VISIBLE_DEVICES=0 ``` - Run it forever ``` while true; do ./run.sh; echo "Restarting..."; sleep 2; done ``` ================================================ FILE: docs/diffusion/api/cli.md ================================================ # SGLang diffusion CLI Inference The SGLang-diffusion CLI provides a quick way to access the inference pipeline for image and video generation. ## Prerequisites - A working SGLang diffusion installation and the `sglang` CLI available in `$PATH`. ## Supported Arguments ### Server Arguments - `--model-path {MODEL_PATH}`: Path to the model or model ID - `--lora-path {LORA_PATH}`: Path to a LoRA adapter (local path or HuggingFace model ID). If not specified, LoRA will not be applied. - `--lora-nickname {NAME}`: Nickname for the LoRA adapter. (default: `default`). - `--num-gpus {NUM_GPUS}`: Number of GPUs to use - `--tp-size {TP_SIZE}`: Tensor parallelism size (only for the encoder; should not be larger than 1 if text encoder offload is enabled, as layer-wise offload plus prefetch is faster) - `--sp-degree {SP_SIZE}`: Sequence parallelism size (typically should match the number of GPUs) - `--ulysses-degree {ULYSSES_DEGREE}`: The degree of DeepSpeed-Ulysses-style SP in USP - `--ring-degree {RING_DEGREE}`: The degree of ring attention-style SP in USP - `--attention-backend {BACKEND}`: Attention backend to use. For SGLang-native pipelines use `fa`, `torch_sdpa`, `sage_attn`, etc. For diffusers pipelines use diffusers backend names like `flash`, `_flash_3_hub`, `sage`, `xformers`. - `--attention-backend-config {CONFIG}`: Configuration for the attention backend. Can be a JSON string (e.g., '{"k": "v"}'), a path to a JSON/YAML file, or key=value pairs (e.g., "k=v,k2=v2"). - `--cache-dit-config {PATH}`: Path to a Cache-DiT YAML/JSON config (diffusers backend only) - `--dit-precision {DTYPE}`: Precision for the DiT model (currently supports fp32, fp16, and bf16). ### Sampling Parameters - `--prompt {PROMPT}`: Text description for the video you want to generate - `--num-inference-steps {STEPS}`: Number of denoising steps - `--negative-prompt {PROMPT}`: Negative prompt to guide generation away from certain concepts - `--seed {SEED}`: Random seed for reproducible generation **Image/Video Configuration** - `--height {HEIGHT}`: Height of the generated output - `--width {WIDTH}`: Width of the generated output - `--num-frames {NUM_FRAMES}`: Number of frames to generate - `--fps {FPS}`: Frames per second for the saved output, if this is a video-generation task **Post-Processing** (frame interpolation & upscaling) SGLang diffusion supports optional post-processing steps — frame interpolation (RIFE) for smoother video and upscaling (Real-ESRGAN) for higher resolution. See the dedicated **[Post-Processing](post_processing.md)** page for full details, supported models, and examples. **Output Options** - `--output-path {PATH}`: Directory to save the generated video - `--save-output`: Whether to save the image/video to disk - `--return-frames`: Whether to return the raw frames ### Using Configuration Files Instead of specifying all parameters on the command line, you can use a configuration file: ```bash sglang generate --config {CONFIG_FILE_PATH} ``` The configuration file should be in JSON or YAML format with the same parameter names as the CLI options. Command-line arguments take precedence over settings in the configuration file, allowing you to override specific values while keeping the rest from the configuration file. Example configuration file (config.json): ```json { "model_path": "FastVideo/FastHunyuan-diffusers", "prompt": "A beautiful woman in a red dress walking down a street", "output_path": "outputs/", "num_gpus": 2, "sp_size": 2, "tp_size": 1, "num_frames": 45, "height": 720, "width": 1280, "num_inference_steps": 6, "seed": 1024, "fps": 24, "precision": "bf16", "vae_precision": "fp16", "vae_tiling": true, "vae_sp": true, "vae_config": { "load_encoder": false, "load_decoder": true, "tile_sample_min_height": 256, "tile_sample_min_width": 256 }, "text_encoder_precisions": [ "fp16", "fp16" ], "mask_strategy_file_path": null, "enable_torch_compile": false } ``` Or using YAML format (config.yaml): ```yaml model_path: "FastVideo/FastHunyuan-diffusers" prompt: "A beautiful woman in a red dress walking down a street" output_path: "outputs/" num_gpus: 2 sp_size: 2 tp_size: 1 num_frames: 45 height: 720 width: 1280 num_inference_steps: 6 seed: 1024 fps: 24 precision: "bf16" vae_precision: "fp16" vae_tiling: true vae_sp: true vae_config: load_encoder: false load_decoder: true tile_sample_min_height: 256 tile_sample_min_width: 256 text_encoder_precisions: - "fp16" - "fp16" mask_strategy_file_path: null enable_torch_compile: false ``` To see all the options, you can use the `--help` flag: ```bash sglang generate --help ``` ## Serve Launch the SGLang diffusion HTTP server and interact with it using the OpenAI SDK and curl. ### Start the server Use the following command to launch the server: ```bash SERVER_ARGS=( --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers --text-encoder-cpu-offload --pin-cpu-memory --num-gpus 4 --ulysses-degree=2 --ring-degree=2 ) sglang serve "${SERVER_ARGS[@]}" ``` - **--model-path**: Which model to load. The example uses `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`. - **--port**: HTTP port to listen on (the default here is `30010`). For detailed API usage, including Image, Video Generation and LoRA management, please refer to the [OpenAI API Documentation](openai_api.md). ### Cloud Storage Support SGLang diffusion supports automatically uploading generated images and videos to S3-compatible cloud storage (e.g., AWS S3, MinIO, Alibaba Cloud OSS, Tencent Cloud COS). When enabled, the server follows a **Generate -> Upload -> Delete** workflow: 1. The artifact is generated to a temporary local file. 2. The file is immediately uploaded to the configured S3 bucket in a background thread. 3. Upon successful upload, the local file is deleted. 4. The API response returns the public URL of the uploaded object. **Configuration** Cloud storage is enabled via environment variables. Note that `boto3` must be installed separately (`pip install boto3`) to use this feature. ```bash # Enable S3 storage export SGLANG_CLOUD_STORAGE_TYPE=s3 export SGLANG_S3_BUCKET_NAME=my-bucket export SGLANG_S3_ACCESS_KEY_ID=your-access-key export SGLANG_S3_SECRET_ACCESS_KEY=your-secret-key # Optional: Custom endpoint for MinIO/OSS/COS export SGLANG_S3_ENDPOINT_URL=https://minio.example.com ``` See [Environment Variables Documentation](../environment_variables.md) for more details. ## Generate Run a one-off generation task without launching a persistent server. To use it, pass both server arguments and sampling parameters in one command, after the `generate` subcommand, for example: ```bash SERVER_ARGS=( --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers --text-encoder-cpu-offload --pin-cpu-memory --num-gpus 4 --ulysses-degree=2 --ring-degree=2 ) SAMPLING_ARGS=( --prompt "A curious raccoon" --save-output --output-path outputs --output-file-name "A curious raccoon.mp4" ) sglang generate "${SERVER_ARGS[@]}" "${SAMPLING_ARGS[@]}" # Or, users can set `SGLANG_CACHE_DIT_ENABLED` env as `true` to enable cache acceleration SGLANG_CACHE_DIT_ENABLED=true sglang generate "${SERVER_ARGS[@]}" "${SAMPLING_ARGS[@]}" ``` Once the generation task has finished, the server will shut down automatically. > [!NOTE] > The HTTP server-related arguments are ignored in this subcommand. ## Component Path Overrides SGLang diffusion allows you to override any pipeline component (e.g., `vae`, `transformer`, `text_encoder`) by specifying a custom checkpoint path. This is useful for: ### Example: FLUX.2-dev with Tiny AutoEncoder You can override **any** component by using `---path`, where `` matches the key in the model's `model_index.json`: For example, replace the default VAE with a distilled tiny autoencoder for ~3x faster decoding: ```bash sglang serve \ --model-path=black-forest-labs/FLUX.2-dev \ # with a Huggingface Repo ID --vae-path=fal/FLUX.2-Tiny-AutoEncoder # or use a local path --vae-path=~/.cache/huggingface/hub/models--fal--FLUX.2-Tiny-AutoEncoder/snapshots/.../vae ``` **Important:** - The component key must match the one in your model's `model_index.json` (e.g., `vae`). - The path must: - either be a Huggingface Repo ID (e.g., fal/FLUX.2-Tiny-AutoEncoder) - or point to a **complete component folder**, containing `config.json` and safetensors files ## Diffusers Backend SGLang diffusion supports a **diffusers backend** that allows you to run any diffusers-compatible model through SGLang's infrastructure using vanilla diffusers pipelines. This is useful for running models without native SGLang implementations or models with custom pipeline classes. ### Arguments | Argument | Values | Description | |----------|--------|-------------| | `--backend` | `auto` (default), `sglang`, `diffusers` | `auto`: prefer native SGLang, fallback to diffusers. `sglang`: force native (fails if unavailable). `diffusers`: force vanilla diffusers pipeline. | | `--diffusers-attention-backend` | `flash`, `_flash_3_hub`, `sage`, `xformers`, `native` | Attention backend for diffusers pipelines. See [diffusers attention backends](https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends). | | `--trust-remote-code` | flag | Required for models with custom pipeline classes (e.g., Ovis). | | `--vae-tiling` | flag | Enable VAE tiling for large image support (decodes tile-by-tile). | | `--vae-slicing` | flag | Enable VAE slicing for lower memory usage (decodes slice-by-slice). | | `--dit-precision` | `fp16`, `bf16`, `fp32` | Precision for the diffusion transformer. | | `--vae-precision` | `fp16`, `bf16`, `fp32` | Precision for the VAE. | | `--enable-torch-compile` | flag | Enable `torch.compile` for diffusers pipelines. | | `--cache-dit-config` | `{PATH}` | Path to a Cache-DiT YAML/JSON config file for accelerating diffusers pipelines with Cache-DiT. | ### Example: Running Ovis-Image-7B [Ovis-Image-7B](https://huggingface.co/AIDC-AI/Ovis-Image-7B) is a 7B text-to-image model optimized for high-quality text rendering. ```bash sglang generate \ --model-path AIDC-AI/Ovis-Image-7B \ --backend diffusers \ --trust-remote-code \ --diffusers-attention-backend flash \ --prompt "A serene Japanese garden with cherry blossoms" \ --height 1024 \ --width 1024 \ --num-inference-steps 30 \ --save-output \ --output-path outputs \ --output-file-name ovis_garden.png ``` ### Extra Diffusers Arguments For pipeline-specific parameters not exposed via CLI, use `diffusers_kwargs` in a config file: ```json { "model_path": "AIDC-AI/Ovis-Image-7B", "backend": "diffusers", "prompt": "A beautiful landscape", "diffusers_kwargs": { "cross_attention_kwargs": {"scale": 0.5} } } ``` ```bash sglang generate --config config.json ``` ### Cache-DiT Acceleration Users who use the diffusers backend can also leverage Cache-DiT acceleration and load custom cache configs from a YAML file to boost performance of diffusers pipelines. See the [Cache-DiT Acceleration](https://docs.sglang.io/diffusion/performance/cache/cache_dit.html) documentation for details. ================================================ FILE: docs/diffusion/api/openai_api.md ================================================ # SGLang Diffusion OpenAI API The SGLang diffusion HTTP server implements an OpenAI-compatible API for image and video generation, as well as LoRA adapter management. ## Prerequisites - Python 3.11+ if you plan to use the OpenAI Python SDK. ## Serve Launch the server using the `sglang serve` command. ### Start the server ```bash SERVER_ARGS=( --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers --text-encoder-cpu-offload --pin-cpu-memory --num-gpus 4 --ulysses-degree=2 --ring-degree=2 --port 30010 ) sglang serve "${SERVER_ARGS[@]}" ``` - **--model-path**: Path to the model or model ID. - **--port**: HTTP port to listen on (default: `30000`). **Get Model Information** **Endpoint:** `GET /models` Returns information about the model served by this server, including model path, task type, pipeline configuration, and precision settings. **Curl Example:** ```bash curl -sS -X GET "http://localhost:30010/models" ``` **Response Example:** ```json { "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "task_type": "T2V", "pipeline_name": "wan_pipeline", "pipeline_class": "WanPipeline", "num_gpus": 4, "dit_precision": "bf16", "vae_precision": "fp16" } ``` --- ## Endpoints ### Image Generation The server implements an OpenAI-compatible Images API under the `/v1/images` namespace. **Create an image** **Endpoint:** `POST /v1/images/generations` **Python Example (b64_json response):** ```python import base64 from openai import OpenAI client = OpenAI(api_key="sk-proj-1234567890", base_url="http://localhost:30010/v1") img = client.images.generate( prompt="A calico cat playing a piano on stage", size="1024x1024", n=1, response_format="b64_json", ) image_bytes = base64.b64decode(img.data[0].b64_json) with open("output.png", "wb") as f: f.write(image_bytes) ``` **Curl Example:** ```bash curl -sS -X POST "http://localhost:30010/v1/images/generations" \ -H "Content-Type: application/json" \ -H "Authorization: Bearer sk-proj-1234567890" \ -d '{ "prompt": "A calico cat playing a piano on stage", "size": "1024x1024", "n": 1, "response_format": "b64_json" }' ``` > **Note** > If `response_format=url` is used and cloud storage is not configured, the API returns > a relative URL like `/v1/images//content`. **Edit an image** **Endpoint:** `POST /v1/images/edits` This endpoint accepts a multipart form upload with input images and a text prompt. The server can return either a base64-encoded image or a URL to download the image. **Curl Example (b64_json response):** ```bash curl -sS -X POST "http://localhost:30010/v1/images/edits" \ -H "Authorization: Bearer sk-proj-1234567890" \ -F "image=@local_input_image.png" \ -F "url=image_url.jpg" \ -F "prompt=A calico cat playing a piano on stage" \ -F "size=1024x1024" \ -F "response_format=b64_json" ``` **Curl Example (URL response):** ```bash curl -sS -X POST "http://localhost:30010/v1/images/edits" \ -H "Authorization: Bearer sk-proj-1234567890" \ -F "image=@local_input_image.png" \ -F "url=image_url.jpg" \ -F "prompt=A calico cat playing a piano on stage" \ -F "size=1024x1024" \ -F "response_format=url" ``` **Download image content** When `response_format=url` is used with `POST /v1/images/generations` or `POST /v1/images/edits`, the API returns a relative URL like `/v1/images//content`. **Endpoint:** `GET /v1/images/{image_id}/content` **Curl Example:** ```bash curl -sS -L "http://localhost:30010/v1/images//content" \ -H "Authorization: Bearer sk-proj-1234567890" \ -o output.png ``` ### Video Generation The server implements a subset of the OpenAI Videos API under the `/v1/videos` namespace. **Create a video** **Endpoint:** `POST /v1/videos` **Python Example:** ```python from openai import OpenAI client = OpenAI(api_key="sk-proj-1234567890", base_url="http://localhost:30010/v1") video = client.videos.create( prompt="A calico cat playing a piano on stage", size="1280x720" ) print(f"Video ID: {video.id}, Status: {video.status}") ``` **Curl Example:** ```bash curl -sS -X POST "http://localhost:30010/v1/videos" \ -H "Content-Type: application/json" \ -H "Authorization: Bearer sk-proj-1234567890" \ -d '{ "prompt": "A calico cat playing a piano on stage", "size": "1280x720" }' ``` **List videos** **Endpoint:** `GET /v1/videos` **Python Example:** ```python videos = client.videos.list() for item in videos.data: print(item.id, item.status) ``` **Curl Example:** ```bash curl -sS -X GET "http://localhost:30010/v1/videos" \ -H "Authorization: Bearer sk-proj-1234567890" ``` **Download video content** **Endpoint:** `GET /v1/videos/{video_id}/content` **Python Example:** ```python import time # Poll for completion while True: page = client.videos.list() item = next((v for v in page.data if v.id == video_id), None) if item and item.status == "completed": break time.sleep(5) # Download content resp = client.videos.download_content(video_id=video_id) with open("output.mp4", "wb") as f: f.write(resp.read()) ``` **Curl Example:** ```bash curl -sS -L "http://localhost:30010/v1/videos//content" \ -H "Authorization: Bearer sk-proj-1234567890" \ -o output.mp4 ``` --- ### LoRA Management The server supports dynamic loading, merging, and unmerging of LoRA adapters. **Important Notes:** - Mutual Exclusion: Only one LoRA can be *merged* (active) at a time - Switching: To switch LoRAs, you must first `unmerge` the current one, then `set` the new one - Caching: The server caches loaded LoRA weights in memory. Switching back to a previously loaded LoRA (same path) has little cost **Set LoRA Adapter** Loads one or more LoRA adapters and merges their weights into the model. Supports both single LoRA (backward compatible) and multiple LoRA adapters. **Endpoint:** `POST /v1/set_lora` **Parameters:** - `lora_nickname` (string or list of strings, required): A unique identifier for the LoRA adapter(s). Can be a single string or a list of strings for multiple LoRAs - `lora_path` (string or list of strings/None, optional): Path to the `.safetensors` file(s) or Hugging Face repo ID(s). Required for the first load; optional if re-activating a cached nickname. If a list, must match the length of `lora_nickname` - `target` (string or list of strings, optional): Which transformer(s) to apply the LoRA to. If a list, must match the length of `lora_nickname`. Valid values: - `"all"` (default): Apply to all transformers - `"transformer"`: Apply only to the primary transformer (high noise for Wan2.2) - `"transformer_2"`: Apply only to transformer_2 (low noise for Wan2.2) - `"critic"`: Apply only to the critic model - `strength` (float or list of floats, optional): LoRA strength for merge, default 1.0. If a list, must match the length of `lora_nickname`. Values < 1.0 reduce the effect, values > 1.0 amplify the effect **Single LoRA Example:** ```bash curl -X POST http://localhost:30010/v1/set_lora \ -H "Content-Type: application/json" \ -d '{ "lora_nickname": "lora_name", "lora_path": "/path/to/lora.safetensors", "target": "all", "strength": 0.8 }' ``` **Multiple LoRA Example:** ```bash curl -X POST http://localhost:30010/v1/set_lora \ -H "Content-Type: application/json" \ -d '{ "lora_nickname": ["lora_1", "lora_2"], "lora_path": ["/path/to/lora1.safetensors", "/path/to/lora2.safetensors"], "target": ["transformer", "transformer_2"], "strength": [0.8, 1.0] }' ``` **Multiple LoRA with Same Target:** ```bash curl -X POST http://localhost:30010/v1/set_lora \ -H "Content-Type: application/json" \ -d '{ "lora_nickname": ["style_lora", "character_lora"], "lora_path": ["/path/to/style.safetensors", "/path/to/character.safetensors"], "target": "all", "strength": [0.7, 0.9] }' ``` > [!NOTE] > When using multiple LoRAs: > - All list parameters (`lora_nickname`, `lora_path`, `target`, `strength`) must have the same length > - If `target` or `strength` is a single value, it will be applied to all LoRAs > - Multiple LoRAs applied to the same target will be merged in order **Merge LoRA Weights** Manually merges the currently set LoRA weights into the base model. > [!NOTE] > `set_lora` automatically performs a merge, so this is typically only needed if you have manually unmerged but want to re-apply the same LoRA without calling `set_lora` again.* **Endpoint:** `POST /v1/merge_lora_weights` **Parameters:** - `target` (string, optional): Which transformer(s) to merge. One of "all" (default), "transformer", "transformer_2", "critic" - `strength` (float, optional): LoRA strength for merge, default 1.0. Values < 1.0 reduce the effect, values > 1.0 amplify the effect **Curl Example:** ```bash curl -X POST http://localhost:30010/v1/merge_lora_weights \ -H "Content-Type: application/json" \ -d '{"strength": 0.8}' ``` **Unmerge LoRA Weights** Unmerges the currently active LoRA weights from the base model, restoring it to its original state. This **must** be called before setting a different LoRA. **Endpoint:** `POST /v1/unmerge_lora_weights` **Curl Example:** ```bash curl -X POST http://localhost:30010/v1/unmerge_lora_weights \ -H "Content-Type: application/json" ``` **List LoRA Adapters** Returns loaded LoRA adapters and current application status per module. **Endpoint:** `GET /v1/list_loras` **Curl Example:** ```bash curl -sS -X GET "http://localhost:30010/v1/list_loras" ``` **Response Example:** ```json { "loaded_adapters": [ { "nickname": "lora_a", "path": "/weights/lora_a.safetensors" }, { "nickname": "lora_b", "path": "/weights/lora_b.safetensors" } ], "active": { "transformer": [ { "nickname": "lora2", "path": "tarn59/pixel_art_style_lora_z_image_turbo", "merged": true, "strength": 1.0 } ] } } ``` Notes: - If LoRA is not enabled for the current pipeline, the server will return an error. - `num_lora_layers_with_weights` counts only layers that have LoRA weights applied for the active adapter. ### Example: Switching LoRAs 1. Set LoRA A: ```bash curl -X POST http://localhost:30010/v1/set_lora -d '{"lora_nickname": "lora_a", "lora_path": "path/to/A"}' ``` 2. Generate with LoRA A... 3. Unmerge LoRA A: ```bash curl -X POST http://localhost:30010/v1/unmerge_lora_weights ``` 4. Set LoRA B: ```bash curl -X POST http://localhost:30010/v1/set_lora -d '{"lora_nickname": "lora_b", "lora_path": "path/to/B"}' ``` 5. Generate with LoRA B... ### Adjust Output Quality The server supports adjusting output quality and compression levels for both image and video generation through the `output-quality` and `output-compression` parameters. #### Parameters - **`output-quality`** (string, optional): Preset quality level that automatically sets compression. **Default is `"default"`**. Valid values: - `"maximum"`: Highest quality (100) - `"high"`: High quality (90) - `"medium"`: Medium quality (55) - `"low"`: Lower quality (35) - `"default"`: Auto-adjust based on media type (50 for video, 75 for image) - **`output-compression`** (integer, optional): Direct compression level override (0-100). **Default is `None`**. When provided (not `None`), takes precedence over `output-quality`. - `0`: Lowest quality, smallest file size - `100`: Highest quality, largest file size #### Notes - **Precedence**: When both `output-quality` and `output-compression` are provided, `output-compression` takes precedence - **Format Support**: Quality settings apply to JPEG, and video formats. PNG uses lossless compression and ignores these settings - **File Size vs Quality**: Lower compression values (or "low" quality preset) produce smaller files but may show visible artifacts ================================================ FILE: docs/diffusion/api/post_processing.md ================================================ # Post-Processing SGLang diffusion supports optional post-processing steps that run after generation to improve temporal smoothness (frame interpolation) or spatial resolution (upscaling). These steps are independent of the diffusion model and can be combined in a single run. When both are enabled, **frame interpolation runs first** (increasing the frame count), then **upscaling runs on every frame** (increasing the spatial resolution). --- ## Frame Interpolation (video only) Frame interpolation synthesizes new frames between each pair of consecutive generated frames, producing smoother motion without re-running the diffusion model. The `--frame-interpolation-exp` flag controls how many rounds of interpolation to apply: each round inserts one new frame into every gap between adjacent frames, so the output frame count follows the formula: > **(N − 1) × 2^exp + 1** > > e.g. 5 original frames with `exp=1` → 4 gaps × 1 new frame + 5 originals = **9** frames; > with `exp=2` → **17** frames. ### CLI Arguments | Argument | Description | |----------|-------------| | `--enable-frame-interpolation` | Enable frame interpolation. Model weights are downloaded automatically on first use. | | `--frame-interpolation-exp {EXP}` | Interpolation exponent — `1` = 2× temporal resolution, `2` = 4×, etc. (default: `1`) | | `--frame-interpolation-scale {SCALE}` | RIFE inference scale; use `0.5` for high-resolution inputs to save memory (default: `1.0`) | | `--frame-interpolation-model-path {PATH}` | Local directory or HuggingFace repo ID containing RIFE `flownet.pkl` weights (default: `elfgum/RIFE-4.22.lite`, downloaded automatically) | ### Supported Models Frame interpolation uses the [RIFE](https://github.com/hzwer/Practical-RIFE) (Real-Time Intermediate Flow Estimation) architecture. Only **RIFE 4.22.lite** (`IFNet` with 4-scale `IFBlock` backbone) is supported. The network topology is hard-coded, so custom weights provided via `--frame-interpolation-model-path` must be a `flownet.pkl` checkpoint that is compatible with this architecture. Other RIFE versions (e.g., older `v4.x` variants with different block counts) or entirely different frame interpolation methods (FILM, AMT, etc.) are **not supported**. | Weight | HuggingFace Repo | Description | |--------|------------------|-------------| | RIFE 4.22.lite *(default)* | [`elfgum/RIFE-4.22.lite`](https://huggingface.co/elfgum/RIFE-4.22.lite) | Lightweight model, downloaded automatically on first use | ### Example Generate a 5-frame video and interpolate to 9 frames ((5 − 1) × 2¹ + 1 = 9): ```bash sglang generate \ --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --prompt "A dog running through a park" \ --num-frames 5 \ --enable-frame-interpolation \ --frame-interpolation-exp 1 \ --save-output ``` --- ## Upscaling (image and video) Upscaling increases the spatial resolution of generated images or video frames using [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN). The model weights are downloaded automatically on first use and cached for subsequent runs. ### CLI Arguments | Argument | Description | |----------|-------------| | `--enable-upscaling` | Enable post-generation upscaling using Real-ESRGAN. | | `--upscaling-scale {SCALE}` | Desired upscaling factor (default: `4`). The 4× model is used internally; if a different scale is requested, a bicubic resize is applied after the network output. | | `--upscaling-model-path {PATH}` | Local `.pth` file, HuggingFace repo ID, or `repo_id:filename` for Real-ESRGAN weights (default: `ai-forever/Real-ESRGAN` with `RealESRGAN_x4.pth`, downloaded automatically). Use the `repo_id:filename` format to specify a custom weight file from a HuggingFace repo (e.g. `my-org/my-esrgan:weights.pth`). | ### Supported Models Upscaling supports two Real-ESRGAN network architectures. The correct architecture is **auto-detected** from the checkpoint keys, so you only need to point `--upscaling-model-path` at a valid `.pth` file: | Architecture | Example Weights | Description | |--------------|-----------------|-------------| | **RRDBNet** | `RealESRGAN_x4plus.pth` | Heavier model with higher quality; best for photos | | **SRVGGNetCompact** | `RealESRGAN_x4.pth` *(default)*, `realesr-animevideov3.pth`, `realesr-general-x4v3.pth` | Lightweight model; faster inference, good for video | The default weight file is [`ai-forever/Real-ESRGAN`](https://huggingface.co/ai-forever/Real-ESRGAN) with `RealESRGAN_x4.pth` (SRVGGNetCompact, 4× native scale). Other super-resolution models (e.g., SwinIR, HAT, BSRGAN) are **not supported** — only Real-ESRGAN checkpoints using the two architectures above are compatible. ### Examples Generate a 1024×1024 image and upscale to 4096×4096: ```bash sglang generate \ --model-path black-forest-labs/FLUX.2-dev \ --prompt "A cat sitting on a windowsill" \ --output-size 1024x1024 \ --enable-upscaling \ --save-output ``` Generate a video and upscale each frame by 4×: ```bash sglang generate \ --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ --prompt "A curious raccoon" \ --enable-upscaling \ --upscaling-scale 4 \ --save-output ``` --- ## Combining Frame Interpolation and Upscaling Frame interpolation and upscaling can be combined in a single run. Interpolation is applied first (increasing the frame count), then upscaling is applied to every frame (increasing the spatial resolution). Example — generate 5 frames, interpolate to 9 frames, and upscale each frame by 4×: ```bash sglang generate \ --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ --prompt "A curious raccoon" \ --num-frames 5 \ --enable-frame-interpolation \ --frame-interpolation-exp 1 \ --enable-upscaling \ --upscaling-scale 4 \ --save-output ``` ================================================ FILE: docs/diffusion/ci_perf.md ================================================ ## Perf Baseline Generation Script `python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py` starts a local diffusion server, issues requests for selected test cases, aggregates stage/denoise-step/E2E timings from the perf log, and writes the results back to the `scenarios` section of `perf_baselines.json`. ### Usage Update a single case: ```bash python python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py --case qwen_image_t2i ``` Select by regex: ```bash python python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py --match 'qwen_image_.*' ``` Run all keys from the baseline file `scenarios`: ```bash python python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py --all-from-baseline ``` Specify input/output paths and timeout: ```bash python python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py --baseline python/sglang/multimodal_gen/test/server/perf_baselines.json --out /tmp/perf_baselines.json --timeout 600 ``` ================================================ FILE: docs/diffusion/compatibility_matrix.md ================================================ # Compatibility Matrix The table below shows every supported model and the optimizations supported for them. The symbols used have the following meanings: - ✅ = Full compatibility - ❌ = No compatibility - ⭕ = Does not apply to this model ## Models x Optimization The `HuggingFace Model ID` can be passed directly to `from_pretrained()` methods, and sglang-diffusion will use the optimal default parameters when initializing and generating videos. ### Video Generation Models | Model Name | Hugging Face Model ID | Resolutions | TeaCache | Sliding Tile Attn | Sage Attn | Video Sparse Attention (VSA) | Sparse Linear Attention (SLA) | Sage Sparse Linear Attention (SageSLA) | Sparse Video Gen 2 (SVG2) | |:-----------------------------|:--------------------------------------------------|:--------------------|:--------:|:-----------------:|:---------:|:----------------------------:|:----------------------------:|:-----------------------------------------------:|:----------------------------------:| | FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480p | ⭕ | ⭕ | ⭕ | ✅ | ❌ | ❌ | ❌ | | FastWan2.2 TI2V 5B Full Attn | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720p | ⭕ | ⭕ | ⭕ | ✅ | ❌ | ❌ | ❌ | | Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720p | ⭕ | ⭕ | ✅ | ⭕ | ❌ | ❌ | ❌ | | Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ | ❌ | ❌ | ❌ | | Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480p
720p | ❌ | ❌ | ✅ | ⭕ | ❌ | ❌ | ❌ | | HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | | FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720×1280
544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | | Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | | Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480p, 720p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | | Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | | Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | | TurboWan2.1 T2V 1.3B | `IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | | TurboWan2.1 T2V 14B | `IPostYellow/TurboWan2.1-T2V-14B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | | TurboWan2.1 T2V 14B 720P | `IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | | TurboWan2.2 I2V A14B | `IPostYellow/TurboWan2.2-I2V-A14B-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | **Note**: 1.Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue. 2.SageSLA Based on SpargeAttn. Install it first with `pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation` ### Image Generation Models | Model Name | HuggingFace Model ID | Resolutions | |:-----------------|:----------------------------------------|:---------------| | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | Any resolution | | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Any resolution | | FLUX.2-Klein | `black-forest-labs/FLUX.2-klein-4B` | Any resolution | | Z-Image-Turbo | `Tongyi-MAI/Z-Image-Turbo` | Any resolution | | GLM-Image | `zai-org/GLM-Image` | Any resolution | | Qwen Image | `Qwen/Qwen-Image` | Any resolution | | Qwen Image 2512 | `Qwen/Qwen-Image-2512` | Any resolution | | Qwen Image Edit | `Qwen/Qwen-Image-Edit` | Any resolution | ## Verified LoRA Examples This section lists example LoRAs that have been explicitly tested and verified with each base model in the **SGLang Diffusion** pipeline. > Important: > LoRAs that are not listed here are not necessarily incompatible. > In practice, most standard LoRAs are expected to work, especially those following common Diffusers or SD-style conventions. > The entries below simply reflect configurations that have been manually validated by the SGLang team. ### Verified LoRAs by Base Model | Base Model | Supported LoRAs | |:-----------------|:----------------| | Wan2.2 | `lightx2v/Wan2.2-Distill-Loras`
`Cseti/wan2.2-14B-Arcane_Jinx-lora-v1` | | Wan2.1 | `lightx2v/Wan2.1-Distill-Loras` | | Z-Image-Turbo | `tarn59/pixel_art_style_lora_z_image_turbo`
`wcde/Z-Image-Turbo-DeJPEG-Lora` | | Qwen-Image | `lightx2v/Qwen-Image-Lightning`
`flymy-ai/qwen-image-realism-lora`
`prithivMLmods/Qwen-Image-HeadshotX`
`starsfriday/Qwen-Image-EVA-LoRA` | | Qwen-Image-Edit | `ostris/qwen_image_edit_inpainting`
`lightx2v/Qwen-Image-Edit-2511-Lightning` | | Flux | `dvyio/flux-lora-simple-illustration`
`XLabs-AI/flux-furry-lora`
`XLabs-AI/flux-RealismLora` | ## Special requirements ### Sliding Tile Attention - Currently, only Hopper GPUs (H100s) are supported. ================================================ FILE: docs/diffusion/contributing.md ================================================ # Contributing to SGLang Diffusion This guide outlines the requirements for contributing to the SGLang Diffusion module (`sglang.multimodal_gen`). ## On AI-Assisted ("Vibe Coding") PRs Vibe-coded PRs are welcome — we judge code quality, not how it was produced. The bar is the same for all PRs: - **No over-commenting.** If the name says it all, skip the docstring. - **No over-catching.** Don't guard against errors that virtually never happen in practice. - **Test before submitting.** AI-generated code can be subtly wrong — verify correctness end-to-end. ## Commit Message Convention We follow a structured commit message format to maintain a clean history. **Format:** ```text [diffusion] : ``` **Examples:** - `[diffusion] cli: add --perf-dump-path argument` - `[diffusion] scheduler: fix deadlock in batch processing` - `[diffusion] model: support Stable Diffusion 3.5` **Rules:** - **Prefix**: Always start with `[diffusion]`. - **Scope** (Optional): `cli`, `scheduler`, `model`, `pipeline`, `docs`, etc. - **Subject**: Imperative mood, short and clear (e.g., "add feature" not "added feature"). ## Performance Reporting For PRs that impact **latency**, **throughput**, or **memory usage**, you **should** provide a performance comparison report. ### How to Generate a Report 1. **Baseline**: run the benchmark (for a single generation task) ```bash $ sglang generate --model-path --prompt "A benchmark prompt" --perf-dump-path baseline.json ``` 2. **New**: run the same benchmark, without modifying any server_args or sampling_params ```bash $ sglang generate --model-path --prompt "A benchmark prompt" --perf-dump-path new.json ``` 3. **Compare**: run the compare script, which will print a Markdown table to the console ```bash $ python python/sglang/multimodal_gen/benchmarks/compare_perf.py baseline.json new.json [new2.json ...] ### Performance Comparison Report ... ``` 4. **Paste**: paste the table into the PR description ## CI-Based Change Protection Consider adding tests to the `pr-test` or `nightly-test` suites to safeguard your changes, especially for PRs that: - support a new model - add a testcase for this new model to `testcase_configs.py` - support or fix important features - significantly improve performance Please run the according testcase, then update/add the baseline to `perf_baselines.json` by following the instruction in console if applicable. See [test](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen/test) for examples ================================================ FILE: docs/diffusion/environment_variables.md ================================================ ## Apple MPS | Environment Variable | Default | Description | |----------------------|---------|--------------------------------------------------------------| | `SGLANG_USE_MLX` | not set | Set to `1` to enable MLX fused Metal kernels for norm ops on MPS | ## Caching Acceleration These variables configure caching acceleration for Diffusion Transformer (DiT) models. SGLang supports multiple caching strategies - see [caching documentation](performance/cache/index.md) for an overview. ### Cache-DiT Configuration See [cache-dit documentation](performance/cache/cache_dit.md) for detailed configuration. | Environment Variable | Default | Description | |-------------------------------------|---------|------------------------------------------| | `SGLANG_CACHE_DIT_ENABLED` | false | Enable Cache-DiT acceleration | | `SGLANG_CACHE_DIT_FN` | 1 | First N blocks to always compute | | `SGLANG_CACHE_DIT_BN` | 0 | Last N blocks to always compute | | `SGLANG_CACHE_DIT_WARMUP` | 4 | Warmup steps before caching | | `SGLANG_CACHE_DIT_RDT` | 0.24 | Residual difference threshold | | `SGLANG_CACHE_DIT_MC` | 3 | Max continuous cached steps | | `SGLANG_CACHE_DIT_TAYLORSEER` | false | Enable TaylorSeer calibrator | | `SGLANG_CACHE_DIT_TS_ORDER` | 1 | TaylorSeer order (1 or 2) | | `SGLANG_CACHE_DIT_SCM_PRESET` | none | SCM preset (none/slow/medium/fast/ultra) | | `SGLANG_CACHE_DIT_SCM_POLICY` | dynamic | SCM caching policy | | `SGLANG_CACHE_DIT_SCM_COMPUTE_BINS` | not set | Custom SCM compute bins | | `SGLANG_CACHE_DIT_SCM_CACHE_BINS` | not set | Custom SCM cache bins | ## Cloud Storage These variables configure S3-compatible cloud storage for automatically uploading generated images and videos. | Environment Variable | Default | Description | |---------------------------------|---------|--------------------------------------------------------| | `SGLANG_CLOUD_STORAGE_TYPE` | not set | Set to `s3` to enable cloud storage | | `SGLANG_S3_BUCKET_NAME` | not set | The name of the S3 bucket | | `SGLANG_S3_ENDPOINT_URL` | not set | Custom endpoint URL (for MinIO, OSS, etc.) | | `SGLANG_S3_REGION_NAME` | us-east-1 | AWS region name | | `SGLANG_S3_ACCESS_KEY_ID` | not set | AWS Access Key ID | | `SGLANG_S3_SECRET_ACCESS_KEY` | not set | AWS Secret Access Key | ================================================ FILE: docs/diffusion/index.md ================================================ # SGLang Diffusion SGLang Diffusion is an inference framework for accelerated image and video generation using diffusion models. It provides an end-to-end unified pipeline with optimized kernels and an efficient scheduler loop. ## Key Features - **Broad Model Support**: Wan series, FastWan series, Hunyuan, Qwen-Image, Qwen-Image-Edit, Flux, Z-Image, GLM-Image, and more - **Fast Inference**: Optimized kernels, efficient scheduler loop, and Cache-DiT acceleration - **Ease of Use**: OpenAI-compatible API, CLI, and Python SDK - **Multi-Platform**: - NVIDIA GPUs (H100, H200, A100, B200, 4090) - AMD GPUs (MI300X, MI325X) - Ascend NPU (A2, A3) - Apple Silicon (M-series via MPS) - Moore Threads GPUs (MTT S5000) --- ## Quick Start ### Installation ```bash uv pip install "sglang[diffusion]" --prerelease=allow ``` See [Installation Guide](installation.md) for more installation methods and ROCm-specific instructions. ### Basic Usage Generate an image with the CLI: ```bash sglang generate --model-path Qwen/Qwen-Image \ --prompt "A beautiful sunset over the mountains" \ --save-output ``` Or start a server with the OpenAI-compatible API: ```bash sglang serve --model-path Qwen/Qwen-Image --port 30010 ``` --- ## Documentation ### Getting Started - **[Installation](installation.md)** - Install SGLang Diffusion via pip, uv, Docker, or from source - **[Compatibility Matrix](compatibility_matrix.md)** - Supported models and optimization compatibility ### Usage - **[CLI Documentation](api/cli.md)** - Command-line interface for `sglang generate` and `sglang serve` - **[OpenAI API](api/openai_api.md)** - OpenAI-compatible API for image/video generation and LoRA management - **[Post-Processing](api/post_processing.md)** - Frame interpolation (RIFE) and upscaling (Real-ESRGAN) ### Performance Optimization - **[Performance Overview](performance/index.md)** - Overview of all performance optimization strategies - **[Attention Backends](performance/attention_backends.md)** - Available attention backends (FlashAttention, SageAttention, etc.) - **[Caching Strategies](performance/cache/)** - Cache-DiT and TeaCache acceleration - **[Profiling](performance/profiling.md)** - Profiling techniques with PyTorch Profiler and Nsight Systems ### Reference - **[Environment Variables](environment_variables.md)** - Configuration via environment variables - **[Support New Models](support_new_models.md)** - Guide for adding new diffusion models - **[Contributing](contributing.md)** - Contribution guidelines and commit message conventions - **[CI Performance](ci_perf.md)** - Performance baseline generation script --- ## CLI Quick Reference ### Generate (one-off generation) ```bash sglang generate --model-path --prompt "" --save-output ``` ### Serve (HTTP server) ```bash sglang serve --model-path --port 30010 ``` ### Enable Cache-DiT acceleration ```bash SGLANG_CACHE_DIT_ENABLED=true sglang generate --model-path --prompt "" ``` --- ## References - [SGLang GitHub](https://github.com/sgl-project/sglang) - [Cache-DiT](https://github.com/vipshop/cache-dit) - [FastVideo](https://github.com/hao-ai-lab/FastVideo) - [xDiT](https://github.com/xdit-project/xDiT) - [Diffusers](https://github.com/huggingface/diffusers) ================================================ FILE: docs/diffusion/installation.md ================================================ # Install SGLang-Diffusion You can install SGLang-Diffusion using one of the methods below. ## Standard Installation (NVIDIA GPUs) ### Method 1: With pip or uv It is recommended to use uv for a faster installation: ```bash pip install --upgrade pip pip install uv uv pip install "sglang[diffusion]" --prerelease=allow ``` ### Method 2: From source ```bash # Use the latest release branch git clone https://github.com/sgl-project/sglang.git cd sglang # Install the Python packages pip install --upgrade pip pip install -e "python[diffusion]" # With uv uv pip install -e "python[diffusion]" --prerelease=allow ``` ### Method 3: Using Docker The Docker images are available on Docker Hub at [lmsysorg/sglang](https://hub.docker.com/r/lmsysorg/sglang), built from the [Dockerfile](https://github.com/sgl-project/sglang/blob/main/docker/Dockerfile). Replace `` below with your HuggingFace Hub [token](https://huggingface.co/docs/hub/en/security-tokens). ```bash docker run --gpus all \ --shm-size 32g \ -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=" \ --ipc=host \ lmsysorg/sglang:dev \ zsh -c '\ echo "Installing diffusion dependencies..." && \ pip install -e "python[diffusion]" && \ echo "Starting SGLang-Diffusion..." && \ sglang generate \ --model-path black-forest-labs/FLUX.1-dev \ --prompt "A logo With Bold Large text: SGL Diffusion" \ --save-output \ ' ``` ## Platform-Specific: ROCm (AMD GPUs) For AMD Instinct GPUs (e.g., MI300X), you can use the ROCm-enabled Docker image: ```bash docker run --device=/dev/kfd --device=/dev/dri --ipc=host \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env HF_TOKEN= \ lmsysorg/sglang:v0.5.5.post2-rocm700-mi30x \ sglang generate --model-path black-forest-labs/FLUX.1-dev --prompt "A logo With Bold Large text: SGL Diffusion" --save-output ``` For detailed ROCm system configuration and installation from source, see [AMD GPUs](../../platforms/amd_gpu.md). ## Platform-Specific: MUSA (Moore Threads GPUs) For Moore Threads GPUs (MTGPU) with the MUSA software stack, please follow the instructions below to install from source: ```bash # Clone the repository git clone https://github.com/sgl-project/sglang.git cd sglang # Install the Python packages pip install --upgrade pip rm -f python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml pip install -e "python[all_musa]" ``` ## Platform-Specific: Ascend NPU For Ascend NPU, please follow the [NPU installation guide](../platforms/ascend_npu.md). Quick test: ```bash sglang generate --model-path black-forest-labs/FLUX.1-dev \ --prompt "A logo With Bold Large text: SGL Diffusion" \ --save-output ``` ## Platform-Specific: Apple MPS For Apple MPS, please follow the instructions below to install from source: ```bash # Install ffmpeg brew install ffmpeg # Install uv brew install uv # Clone the repository git clone https://github.com/sgl-project/sglang.git cd sglang # Create and activate a virtual environment uv venv -p 3.11 sglang-diffusion source sglang-diffusion/bin/activate # Install the Python packages uv pip install --upgrade pip rm -f python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml uv pip install -e "python[all_mps]" ``` ================================================ FILE: docs/diffusion/performance/attention_backends.md ================================================ # Attention Backends This document describes the attention backends available in sglang diffusion (`sglang.multimodal_gen`) and how to select them. ## Overview Attention backends are defined by `AttentionBackendEnum` (`sglang.multimodal_gen.runtime.platforms.interface.AttentionBackendEnum`) and selected via the CLI flag `--attention-backend`. Backend selection is performed by the shared attention layers (e.g. `LocalAttention` / `USPAttention` / `UlyssesAttention` in `sglang.multimodal_gen.runtime.layers.attention.layer`) and therefore applies to any model component using these layers (e.g. diffusion transformer / DiT and encoders). When using the diffusers backend, `--attention-backend` is passed through to diffusers' `set_attention_backend` (e.g., `flash`, `_flash_3_hub`, `sage`, `xformers`, `native`). - **CUDA**: prefers FlashAttention (FA3/FA4) when supported; otherwise falls back to PyTorch SDPA. - **ROCm**: uses FlashAttention when available; otherwise falls back to PyTorch SDPA. - **MPS**: always uses PyTorch SDPA. - **NPU**: always uses PyTorch SDPA. ## Backend options For SGLang-native pipelines, the CLI accepts the lowercase names of `AttentionBackendEnum`. The table below lists the backends implemented by the built-in platforms. `fa3`/`fa4` are accepted as aliases for `fa`. | CLI value | Enum value | Notes | |---|---|---| | `fa` / `fa3` / `fa4` | `FA` | FlashAttention. `fa3/fa4` are normalized to `fa` during argument parsing (`ServerArgs.__post_init__`). | | `torch_sdpa` | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention`. | | `sliding_tile_attn` | `SLIDING_TILE_ATTN` | Sliding Tile Attention (STA). Requires `st_attn`. Configure via `--attention-backend-config`. | | `sage_attn` | `SAGE_ATTN` | Requires `sageattention`. Upstream SageAttention CUDA extensions target SM80/SM86/SM89/SM90/SM120 (compute capability 8.0/8.6/8.9/9.0/12.0); see upstream `setup.py`: https://github.com/thu-ml/SageAttention/blob/main/setup.py. | | `sage_attn_3` | `SAGE_ATTN_3` | Requires SageAttention3 installed per upstream instructions. | | `video_sparse_attn` | `VIDEO_SPARSE_ATTN` | Requires `vsa`. Configure `sparsity` via `--attention-backend-config`. | | `vmoba_attn` | `VMOBA_ATTN` | Requires `kernel.attn.vmoba_attn.vmoba`. Configure via `--attention-backend-config`. | | `aiter` | `AITER` | Requires `aiter`. | | `aiter_sage` | `AITER_SAGE` | Requires `aiter`. | | `sparse_video_gen_2_attn` | `SPARSE_VIDEO_GEN_2_ATTN` | Requires `svg`. See installation instructions at https://github.com/svg-project/Sparse-VideoGen. | ## Selection priority The selection order in `runtime/layers/attention/selector.py` is: 1. `global_force_attn_backend(...)` / `global_force_attn_backend_context_manager(...)` 2. CLI `--attention-backend` (`ServerArgs.attention_backend`) 3. Auto selection (platform capability, dtype, and installed packages) ## Configuration Some backends require additional configuration. You can pass these parameters via `--attention-backend-config`. This argument accepts: - A path to a JSON or YAML configuration file. - A JSON string (e.g., `'{"sparsity": 0.5}'`). - Key-value pairs (e.g., `"sparsity=0.5,enable_x=true"`). ### Supported Configuration Parameters **Sliding Tile Attention (`sliding_tile_attn`)** | Parameter | Type | Description | Default | | :--- | :--- | :--- | :--- | | `mask_strategy_file_path` | `str` | **Required.** Path to the mask strategy JSON file. | - | | `sta_mode` | `str` | Mode of STA. | `STA_inference` | | `skip_time_steps` | `int` | Number of steps to use full attention before switching to sparse attention. | `15` | **Video Sparse Attention (`video_sparse_attn`)** | Parameter | Type | Description | Default | | :--- | :--- | :--- | :--- | | `sparsity` | `float` | Validation sparsity (0.0 - 1.0). | `0.0` | **V-MoBA (`vmoba_attn`)** | Parameter | Type | Description | Default | | :--- | :--- | :--- | :--- | | `temporal_chunk_size` | `int` | Chunk size for temporal dimension. | - | | `temporal_topk` | `int` | Top-K tokens to select in temporal dimension. | - | | `spatial_chunk_size` | `list[int]` | Chunk size for spatial dimension (H, W). | - | | `spatial_topk` | `int` | Top-K tokens to select in spatial dimension. | - | | `st_chunk_size` | `list[int]` | Chunk size for spatiotemporal dimension (T, H, W). | - | | `st_topk` | `int` | Top-K tokens to select in spatiotemporal dimension. | - | | `moba_select_mode` | `str` | Selection mode (e.g., `threshold`). | `threshold` | | `moba_threshold` | `float` | Threshold value for selection. | `0.25` | | `moba_threshold_type` | `str` | Type of thresholding (e.g., `query_head`). | `query_head` | | `first_full_step` | `int` | Number of initial steps to use full attention. | `12` | | `first_full_layer` | `int` | Number of initial layers to use full attention. | `0` | | `temporal_layer` | `int` | Number of temporal layers. | `1` | | `spatial_layer` | `int` | Number of spatial layers. | `1` | | `st_layer` | `int` | Number of spatiotemporal layers. | `1` | ## Platform support matrix | Backend | CUDA | ROCm | MPS | NPU | Notes | |---|---:|---:|---:|---:|---| | `fa` | ✅ | ✅ | ❌ | ❌ | CUDA requires SM80+ and fp16/bf16. FlashAttention is only used when the required runtime is installed; otherwise it falls back to `torch_sdpa`. | | `torch_sdpa` | ✅ | ✅ | ✅ | ✅ | Most compatible option across platforms. | | `sliding_tile_attn` | ✅ | ❌ | ❌ | ❌ | CUDA-only. Requires `st_attn`. Configure via `--attention-backend-config`. | | `sage_attn` | ✅ | ❌ | ❌ | ❌ | CUDA-only (optional dependency). | | `sage_attn_3` | ✅ | ❌ | ❌ | ❌ | CUDA-only (optional dependency). | | `video_sparse_attn` | ✅ | ❌ | ❌ | ❌ | CUDA-only. Requires `vsa`. Configure `sparsity` via `--attention-backend-config`. | | `vmoba_attn` | ✅ | ❌ | ❌ | ❌ | CUDA-only. Requires `kernel.attn.vmoba_attn.vmoba`. Configure via `--attention-backend-config`. | | `aiter` | ❌ | ✅ | ❌ | ❌ | Requires `aiter`. | | `aiter_sage` | ❌ | ✅ | ❌ | ❌ | Requires `aiter`. | | `sparse_video_gen_2_attn` | ✅ | ❌ | ❌ | ❌ | CUDA-only. Requires `svg`. | ## Usage ### Select a backend via CLI ```bash sglang generate \ --model-path \ --prompt "..." \ --attention-backend fa ``` ```bash sglang generate \ --model-path \ --prompt "..." \ --attention-backend torch_sdpa ``` ### Using Sliding Tile Attention (STA) ```bash # Pass the mask strategy file path via config sglang generate \ --model-path \ --prompt "..." \ --attention-backend sliding_tile_attn \ --attention-backend-config "mask_strategy_file_path=/abs/path/to/mask_strategy.json" ``` ### Notes for ROCm / MPS - ROCm: use `--attention-backend torch_sdpa` or `fa` depending on what is available in your environment. - MPS: the platform implementation always uses `torch_sdpa`. ================================================ FILE: docs/diffusion/performance/cache/cache_dit.md ================================================ # Cache-DiT Acceleration SGLang integrates [Cache-DiT](https://github.com/vipshop/cache-dit), a caching acceleration engine for Diffusion Transformers (DiT), to achieve up to **1.69x inference speedup** with minimal quality loss. ## Overview **Cache-DiT** uses intelligent caching strategies to skip redundant computation in the denoising loop: - **DBCache (Dual Block Cache)**: Dynamically decides when to cache transformer blocks based on residual differences - **TaylorSeer**: Uses Taylor expansion for calibration to optimize caching decisions - **SCM (Step Computation Masking)**: Step-level caching control for additional speedup ## Basic Usage Enable Cache-DiT by exporting the environment variable and using `sglang generate` or `sglang serve` : ```bash SGLANG_CACHE_DIT_ENABLED=true \ sglang generate --model-path Qwen/Qwen-Image \ --prompt "A beautiful sunset over the mountains" ``` ## Diffusers Backend Cache-DiT supports loading acceleration configs from a custom YAML file. For diffusers pipelines (`diffusers` backend), pass the YAML/JSON path via `--cache-dit-config`. This flow requires cache-dit >= 1.2.0 (`cache_dit.load_configs`). ### Single GPU inference Define a `cache.yaml` file that contains: - DBCache + TaylorSeer ```yaml cache_config: max_warmup_steps: 8 warmup_interval: 2 max_cached_steps: -1 max_continuous_cached_steps: 2 Fn_compute_blocks: 1 Bn_compute_blocks: 0 residual_diff_threshold: 0.12 enable_taylorseer: true taylorseer_order: 1 ``` Then apply the config with: ```bash sglang generate \ --backend diffusers \ --model-path Qwen/Qwen-Image \ --cache-dit-config cache.yaml \ --prompt "A beautiful sunset over the mountains" ``` - DBCache + TaylorSeer + SCM (Step Computation Mask) ```yaml cache_config: max_warmup_steps: 8 warmup_interval: 2 max_cached_steps: -1 max_continuous_cached_steps: 2 Fn_compute_blocks: 1 Bn_compute_blocks: 0 residual_diff_threshold: 0.12 enable_taylorseer: true taylorseer_order: 1 # Must set the num_inference_steps for SCM. The SCM will automatically # generate the steps computation mask based on the num_inference_steps. # Reference: https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/#scm-steps-computation-masking num_inference_steps: 28 steps_computation_mask: fast ``` - DBCache + TaylorSeer + SCM (Step Computation Mask) + Cache CFG ```yaml cache_config: max_warmup_steps: 8 warmup_interval: 2 max_cached_steps: -1 max_continuous_cached_steps: 2 Fn_compute_blocks: 1 Bn_compute_blocks: 0 residual_diff_threshold: 0.12 enable_taylorseer: true taylorseer_order: 1 num_inference_steps: 28 steps_computation_mask: fast enable_sperate_cfg: true # e.g, Qwen-Image, Wan, Chroma, Ovis-Image, etc. ``` ### Distributed inference - 1D Parallelism Define a parallelism only config yaml `parallel.yaml` file that contains: ```yaml parallelism_config: ulysses_size: auto attention_backend: native ``` Then, apply the distributed inference acceleration config from yaml. `ulysses_size: auto` means that cache-dit will auto detect the `world_size` as the ulysses_size. Otherwise, you should manually set it as specific int number, e.g, 4. Then apply the distributed config with: (Note: please add `--num-gpus N` to specify the number of gpus for distributed inference) ```bash sglang generate \ --backend diffusers \ --num-gpus 4 \ --model-path Qwen/Qwen-Image \ --cache-dit-config parallel.yaml \ --prompt "A futuristic cityscape at sunset" ``` - 2D Parallelism You can also define a 2D parallelism config yaml `parallel_2d.yaml` file that contains: ```yaml parallelism_config: ulysses_size: auto tp_size: 2 attention_backend: native ``` Then, apply the 2D parallelism config from yaml. Here `tp_size: 2` means using tensor parallelism with size 2. The `ulysses_size: auto` means that cache-dit will auto detect the `world_size // tp_size` as the ulysses_size. - 3D Parallelism You can also define a 3D parallelism config yaml `parallel_3d.yaml` file that contains: ```yaml parallelism_config: ulysses_size: 2 ring_size: 2 tp_size: 2 attention_backend: native ``` Then, apply the 3D parallelism config from yaml. Here `ulysses_size: 2`, `ring_size: 2`, `tp_size: 2` means using ulysses parallelism with size 2, ring parallelism with size 2 and tensor parallelism with size 2. - Ulysses Anything Attention To enable Ulysses Anything Attention, you can define a parallelism config yaml `parallel_uaa.yaml` file that contains: ```yaml parallelism_config: ulysses_size: auto attention_backend: native ulysses_anything: true ``` - Ulysses FP8 Communication For device that don't have NVLink support, you can enable Ulysses FP8 Communication to further reduce the communication overhead. You can define a parallelism config yaml `parallel_fp8.yaml` file that contains: ```yaml parallelism_config: ulysses_size: auto attention_backend: native ulysses_float8: true ``` - Async Ulysses CP You can also enable async ulysses CP to overlap the communication and computation. Define a parallelism config yaml `parallel_async.yaml` file that contains: ```yaml parallelism_config: ulysses_size: auto attention_backend: native ulysses_async: true # Now, only support for FLUX.1, Qwen-Image, Ovis-Image and Z-Image. ``` Then, apply the config from yaml. Here `ulysses_async: true` means enabling async ulysses CP. - TE-P and VAE-P You can also specify the extra parallel modules in the yaml config. For example, define a parallelism config yaml `parallel_extra.yaml` file that contains: ```yaml parallelism_config: ulysses_size: auto attention_backend: native extra_parallel_modules: ["text_encoder", "vae"] ``` ### Hybrid Cache and Parallelism Define a hybrid cache and parallel acceleration config yaml `hybrid.yaml` file that contains: ```yaml cache_config: max_warmup_steps: 8 warmup_interval: 2 max_cached_steps: -1 max_continuous_cached_steps: 2 Fn_compute_blocks: 1 Bn_compute_blocks: 0 residual_diff_threshold: 0.12 enable_taylorseer: true taylorseer_order: 1 parallelism_config: ulysses_size: auto attention_backend: native extra_parallel_modules: ["text_encoder", "vae"] ``` Then, apply the hybrid cache and parallel acceleration config from yaml. ```bash sglang generate \ --backend diffusers \ --num-gpus 4 \ --model-path Qwen/Qwen-Image \ --cache-dit-config hybrid.yaml \ --prompt "A beautiful sunset over the mountains" ``` ### Attention Backend In some cases, users may want to only specify the attention backend without any other optimization configs. In this case, you can define a yaml file `attention.yaml` that only contains: ```yaml attention_backend: "flash" # '_flash_3' for Hopper ``` ### Quantization You can also specify the quantization config in the yaml file, required `torchao>=0.16.0`. For example, define a yaml file `quantize.yaml` that contains: ```yaml quantize_config: # quantization configuration for transformer modules # float8 (DQ), float8_weight_only, float8_blockwise, int8 (DQ), int8_weight_only, etc. quant_type: "float8" # layers to exclude from quantization (transformer). layers that contains any of the # keywords in the exclude_layers list will be excluded from quantization. This is useful # for some sensitive layers that are not robust to quantization, e.g., embedding layers. exclude_layers: - "embedder" - "embed" verbose: false # whether to print verbose logs during quantization ``` Then, apply the quantization config from yaml. Please also enable torch.compile for better performance if you are using quantization. For example: ```bash sglang generate \ --backend diffusers \ --model-path Qwen/Qwen-Image \ --warmup \ --cache-dit-config quantize.yaml \ --enable-torch-compile \ --dit-cpu-offload false \ --text-encoder-cpu-offload false \ --prompt "A beautiful sunset over the mountains" ``` ### Combined Configs: Cache + Parallelism + Quantization You can also combine all the above configs together in a single yaml file `combined.yaml` that contains: ```yaml cache_config: max_warmup_steps: 8 warmup_interval: 2 max_cached_steps: -1 max_continuous_cached_steps: 2 Fn_compute_blocks: 1 Bn_compute_blocks: 0 residual_diff_threshold: 0.12 enable_taylorseer: true taylorseer_order: 1 parallelism_config: ulysses_size: auto attention_backend: native extra_parallel_modules: ["text_encoder", "vae"] quantize_config: quant_type: "float8" exclude_layers: - "embedder" - "embed" verbose: false ``` Then, apply the combined cache, parallelism and quantization config from yaml. Please also enable torch.compile for better performance if you are using quantization. ## Advanced Configuration ### DBCache Parameters DBCache controls block-level caching behavior: | Parameter | Env Variable | Default | Description | |-----------|---------------------------|---------|------------------------------------------| | Fn | `SGLANG_CACHE_DIT_FN` | 1 | Number of first blocks to always compute | | Bn | `SGLANG_CACHE_DIT_BN` | 0 | Number of last blocks to always compute | | W | `SGLANG_CACHE_DIT_WARMUP` | 4 | Warmup steps before caching starts | | R | `SGLANG_CACHE_DIT_RDT` | 0.24 | Residual difference threshold | | MC | `SGLANG_CACHE_DIT_MC` | 3 | Maximum continuous cached steps | ### TaylorSeer Configuration TaylorSeer improves caching accuracy using Taylor expansion: | Parameter | Env Variable | Default | Description | |-----------|-------------------------------|---------|---------------------------------| | Enable | `SGLANG_CACHE_DIT_TAYLORSEER` | false | Enable TaylorSeer calibrator | | Order | `SGLANG_CACHE_DIT_TS_ORDER` | 1 | Taylor expansion order (1 or 2) | ### Combined Configuration Example DBCache and TaylorSeer are complementary strategies that work together, you can configure both sets of parameters simultaneously: ```bash SGLANG_CACHE_DIT_ENABLED=true \ SGLANG_CACHE_DIT_FN=2 \ SGLANG_CACHE_DIT_BN=1 \ SGLANG_CACHE_DIT_WARMUP=4 \ SGLANG_CACHE_DIT_RDT=0.4 \ SGLANG_CACHE_DIT_MC=4 \ SGLANG_CACHE_DIT_TAYLORSEER=true \ SGLANG_CACHE_DIT_TS_ORDER=2 \ sglang generate --model-path black-forest-labs/FLUX.1-dev \ --prompt "A curious raccoon in a forest" ``` ### SCM (Step Computation Masking) SCM provides step-level caching control for additional speedup. It decides which denoising steps to compute fully and which to use cached results. **SCM Presets** SCM is configured with presets: | Preset | Compute Ratio | Speed | Quality | |----------|---------------|----------|------------| | `none` | 100% | Baseline | Best | | `slow` | ~75% | ~1.3x | High | | `medium` | ~50% | ~2x | Good | | `fast` | ~35% | ~3x | Acceptable | | `ultra` | ~25% | ~4x | Lower | **Usage** ```bash SGLANG_CACHE_DIT_ENABLED=true \ SGLANG_CACHE_DIT_SCM_PRESET=medium \ sglang generate --model-path Qwen/Qwen-Image \ --prompt "A futuristic cityscape at sunset" ``` **Custom SCM Bins** For fine-grained control over which steps to compute vs cache: ```bash SGLANG_CACHE_DIT_ENABLED=true \ SGLANG_CACHE_DIT_SCM_COMPUTE_BINS="8,3,3,2,2" \ SGLANG_CACHE_DIT_SCM_CACHE_BINS="1,2,2,2,3" \ sglang generate --model-path Qwen/Qwen-Image \ --prompt "A futuristic cityscape at sunset" ``` **SCM Policy** | Policy | Env Variable | Description | |-----------|---------------------------------------|---------------------------------------------| | `dynamic` | `SGLANG_CACHE_DIT_SCM_POLICY=dynamic` | Adaptive caching based on content (default) | | `static` | `SGLANG_CACHE_DIT_SCM_POLICY=static` | Fixed caching pattern | ## Environment Variables All Cache-DiT parameters can be configured via environment variables. See [Environment Variables](../../environment_variables.md) for the complete list. ## Supported Models SGLang Diffusion x Cache-DiT supports almost all models originally supported in SGLang Diffusion: | Model Family | Example Models | |--------------|-----------------------------| | Wan | Wan2.1, Wan2.2 | | Flux | FLUX.1-dev, FLUX.2-dev | | Z-Image | Z-Image-Turbo | | Qwen | Qwen-Image, Qwen-Image-Edit | | Hunyuan | HunyuanVideo | ## Performance Tips 1. **Start with defaults**: The default parameters work well for most models 2. **Use TaylorSeer**: It typically improves both speed and quality 3. **Tune R threshold**: Lower values = better quality, higher values = faster 4. **SCM for extra speed**: Use `medium` preset for good speed/quality balance 5. **Warmup matters**: Higher warmup = more stable caching decisions ## Limitations - **SGLang-native pipelines**: Distributed support (TP/SP) is not yet validated; Cache-DiT will be automatically disabled when `world_size > 1`. - **SCM minimum steps**: SCM requires >= 8 inference steps to be effective - **Model support**: Only models registered in Cache-DiT's BlockAdapterRegister are supported ## Troubleshooting ### SCM disabled for low step count For models with < 8 inference steps (e.g., DMD distilled models), SCM will be automatically disabled. DBCache acceleration still works. ## References - [Cache-DiT](https://github.com/vipshop/cache-dit) - [SGLang Diffusion](../index.md) ================================================ FILE: docs/diffusion/performance/cache/index.md ================================================ # Caching Acceleration for Diffusion Models SGLang provides multiple caching acceleration strategies for Diffusion Transformer (DiT) models. These strategies can significantly reduce inference time by skipping redundant computation. ## Overview SGLang supports two complementary caching approaches: | Strategy | Scope | Mechanism | Best For | |----------|-------|-----------|----------| | **Cache-DiT** | Block-level | Skip individual transformer blocks dynamically | Advanced, higher speedup | | **TeaCache** | Timestep-level | Skip entire denoising steps based on L1 similarity | Simple, built-in | ## Cache-DiT [Cache-DiT](https://github.com/vipshop/cache-dit) provides block-level caching with advanced strategies like DBCache and TaylorSeer. It can achieve up to **1.69x speedup**. See [cache_dit.md](cache_dit.md) for detailed configuration. ### Quick Start ```bash SGLANG_CACHE_DIT_ENABLED=true \ sglang generate --model-path Qwen/Qwen-Image \ --prompt "A beautiful sunset over the mountains" ``` ### Key Features - **DBCache**: Dynamic block-level caching based on residual differences - **TaylorSeer**: Taylor expansion-based calibration for optimized caching - **SCM**: Step-level computation masking for additional speedup ## TeaCache TeaCache (Temporal similarity-based caching) accelerates diffusion inference by detecting when consecutive denoising steps are similar enough to skip computation entirely. See [teacache.md](teacache.md) for detailed documentation. ### Quick Overview - Tracks L1 distance between modulated inputs across timesteps - When accumulated distance is below threshold, reuses cached residual - Supports CFG with separate positive/negative caches ### Supported Models - Wan (wan2.1, wan2.2) - Hunyuan (HunyuanVideo) - Z-Image For Flux and Qwen models, TeaCache is automatically disabled when CFG is enabled. ## References - [Cache-DiT Repository](https://github.com/vipshop/cache-dit) - [TeaCache Paper](https://arxiv.org/abs/2411.14324) ================================================ FILE: docs/diffusion/performance/cache/teacache.md ================================================ # TeaCache Acceleration > **Note**: This is one of two caching strategies available in SGLang. > For an overview of all caching options, see [caching](../index.md). TeaCache (Temporal similarity-based caching) accelerates diffusion inference by detecting when consecutive denoising steps are similar enough to skip computation entirely. ## Overview TeaCache works by: 1. Tracking the L1 distance between modulated inputs across consecutive timesteps 2. Accumulating the rescaled L1 distance over steps 3. When accumulated distance is below a threshold, reusing the cached residual 4. Supporting CFG (Classifier-Free Guidance) with separate positive/negative caches ## How It Works ### L1 Distance Tracking At each denoising step, TeaCache computes the relative L1 distance between the current and previous modulated inputs: ``` rel_l1 = |current - previous|.mean() / |previous|.mean() ``` This distance is then rescaled using polynomial coefficients and accumulated: ``` accumulated += poly(coefficients)(rel_l1) ``` ### Cache Decision - If `accumulated >= threshold`: Force computation, reset accumulator - If `accumulated < threshold`: Skip computation, use cached residual ### CFG Support For models that support CFG cache separation (Wan, Hunyuan, Z-Image), TeaCache maintains separate caches for positive and negative branches: - `previous_modulated_input` / `previous_residual` for positive branch - `previous_modulated_input_negative` / `previous_residual_negative` for negative branch For models that don't support CFG separation (Flux, Qwen), TeaCache is automatically disabled when CFG is enabled. ## Configuration TeaCache is configured via `TeaCacheParams` in the sampling parameters: ```python from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams params = TeaCacheParams( teacache_thresh=0.1, # Threshold for accumulated L1 distance coefficients=[1.0, 0.0, 0.0], # Polynomial coefficients for L1 rescaling ) ``` ### Parameters | Parameter | Type | Description | |-----------|------|-------------| | `teacache_thresh` | float | Threshold for accumulated L1 distance. Lower = more caching, faster but potentially lower quality | | `coefficients` | list[float] | Polynomial coefficients for L1 rescaling. Model-specific tuning | ### Model-Specific Configurations Different models may have different optimal configurations. The coefficients are typically tuned per-model to balance speed and quality. ## Supported Models TeaCache is built into the following model families: | Model Family | CFG Cache Separation | Notes | |--------------|---------------------|-------| | Wan (wan2.1, wan2.2) | Yes | Full support | | Hunyuan (HunyuanVideo) | Yes | To be supported | | Z-Image | Yes | To be supported | | Flux | No | To be supported | | Qwen | No | To be supported | ## References - [TeaCache: Accelerating Diffusion Models with Temporal Similarity](https://arxiv.org/abs/2411.14324) ================================================ FILE: docs/diffusion/performance/index.md ================================================ # Performance Optimization SGLang-Diffusion provides multiple performance optimization strategies to accelerate inference. This section covers all available performance tuning options. ## Overview | Optimization | Type | Description | |--------------|------|-------------| | **Cache-DiT** | Caching | Block-level caching with DBCache, TaylorSeer, and SCM | | **TeaCache** | Caching | Timestep-level caching using L1 similarity | | **Attention Backends** | Kernel | Optimized attention implementations (FlashAttention, SageAttention, etc.) | | **Profiling** | Diagnostics | PyTorch Profiler and Nsight Systems guidance | ## Caching Strategies SGLang supports two complementary caching approaches: ### Cache-DiT [Cache-DiT](https://github.com/vipshop/cache-dit) provides block-level caching with advanced strategies. It can achieve up to **1.69x speedup**. **Quick Start:** ```bash SGLANG_CACHE_DIT_ENABLED=true \ sglang generate --model-path Qwen/Qwen-Image \ --prompt "A beautiful sunset over the mountains" ``` **Key Features:** - **DBCache**: Dynamic block-level caching based on residual differences - **TaylorSeer**: Taylor expansion-based calibration for optimized caching - **SCM**: Step-level computation masking for additional speedup See [Cache-DiT Documentation](cache/cache_dit.md) for detailed configuration. ### TeaCache TeaCache (Temporal similarity-based caching) accelerates diffusion inference by detecting when consecutive denoising steps are similar enough to skip computation entirely. **Quick Overview:** - Tracks L1 distance between modulated inputs across timesteps - When accumulated distance is below threshold, reuses cached residual - Supports CFG with separate positive/negative caches **Supported Models:** Wan (wan2.1, wan2.2), Hunyuan (HunyuanVideo), Z-Image See [TeaCache Documentation](cache/teacache.md) for detailed configuration. ## Attention Backends Different attention backends offer varying performance characteristics depending on your hardware and model: - **FlashAttention**: Fastest on NVIDIA GPUs with fp16/bf16 - **SageAttention**: Alternative optimized implementation - **xformers**: Memory-efficient attention - **SDPA**: PyTorch native scaled dot-product attention See [Attention Backends](attention_backends.md) for platform support and configuration options. ## Profiling To diagnose performance bottlenecks, SGLang-Diffusion supports profiling tools: - **PyTorch Profiler**: Built-in Python profiling - **Nsight Systems**: GPU kernel-level analysis See [Profiling Guide](profiling.md) for detailed instructions. ## References - [Cache-DiT Repository](https://github.com/vipshop/cache-dit) - [TeaCache Paper](https://arxiv.org/abs/2411.14324) ================================================ FILE: docs/diffusion/performance/profiling.md ================================================ # Profiling Multimodal Generation This guide covers profiling techniques for multimodal generation pipelines in SGLang. ## PyTorch Profiler PyTorch Profiler provides detailed kernel execution time, call stack, and GPU utilization metrics. ### Denoising Stage Profiling Profile the denoising stage with sampled timesteps (default: 5 steps after 1 warmup step): ```bash sglang generate \ --model-path Qwen/Qwen-Image \ --prompt "A Logo With Bold Large Text: SGL Diffusion" \ --seed 0 \ --profile ``` **Parameters:** - `--profile`: Enable profiling for the denoising stage - `--num-profiled-timesteps N`: Number of timesteps to profile after warmup (default: 5) - Smaller values reduce trace file size - Example: `--num-profiled-timesteps 10` profiles 10 steps after 1 warmup step ### Full Pipeline Profiling Profile all pipeline stages (text encoding, denoising, VAE decoding, etc.): ```bash sglang generate \ --model-path Qwen/Qwen-Image \ --prompt "A Logo With Bold Large Text: SGL Diffusion" \ --seed 0 \ --profile \ --profile-all-stages ``` **Parameters:** - `--profile-all-stages`: Used with `--profile`, profile all pipeline stages instead of just denoising ### Output Location By default, trace files are saved in the ./logs/ directory. The exact output file path will be shown in the console output, for example: ```bash [mm-dd hh:mm:ss] Saved profiler traces to: /sgl-workspace/sglang/logs/mocked_fake_id_for_offline_generate-5_steps-global-rank0.trace.json.gz ``` ### View Traces Load and visualize trace files at: - https://ui.perfetto.dev/ (recommended) - chrome://tracing (Chrome only) For large trace files, reduce `--num-profiled-timesteps` or avoid using `--profile-all-stages`. ### `--perf-dump-path` (Stage/Step Timing Dump) Besides profiler traces, you can also dump a lightweight JSON report that contains: - stage-level timing breakdown for the full pipeline - step-level timing breakdown for the denoising stage (per diffusion step) This is useful to quickly identify which stage dominates end-to-end latency, and whether denoising steps have uniform runtimes (and if not, which step has an abnormal spike). The dumped JSON contains a `denoise_steps_ms` field formatted as an array of objects, each with a `step` key (the step index) and a `duration_ms` key. Example: ```bash sglang generate \ --model-path \ --prompt "" \ --perf-dump-path perf.json ``` ## Nsight Systems Nsight Systems provides low-level CUDA profiling with kernel details, register usage, and memory access patterns. ### Installation See the [SGLang profiling guide](https://github.com/sgl-project/sglang/blob/main/docs/developer_guide/benchmark_and_profiling.md#profile-with-nsight) for installation instructions. ### Basic Profiling Profile the entire pipeline execution: ```bash nsys profile \ --trace-fork-before-exec=true \ --cuda-graph-trace=node \ --force-overwrite=true \ -o QwenImage \ sglang generate \ --model-path Qwen/Qwen-Image \ --prompt "A Logo With Bold Large Text: SGL Diffusion" \ --seed 0 ``` ### Targeted Stage Profiling Use `--delay` and `--duration` to capture specific stages and reduce file size: ```bash nsys profile \ --trace-fork-before-exec=true \ --cuda-graph-trace=node \ --force-overwrite=true \ --delay 10 \ --duration 30 \ -o QwenImage_denoising \ sglang generate \ --model-path Qwen/Qwen-Image \ --prompt "A Logo With Bold Large Text: SGL Diffusion" \ --seed 0 ``` **Parameters:** - `--delay N`: Wait N seconds before starting capture (skip initialization overhead) - `--duration N`: Capture for N seconds (focus on specific stages) - `--force-overwrite`: Overwrite existing output files ## Notes - **Reduce trace size**: Use `--num-profiled-timesteps` with smaller values or `--delay`/`--duration` with Nsight Systems - **Stage-specific analysis**: Use `--profile` alone for denoising stage, add `--profile-all-stages` for full pipeline - **Multiple runs**: Profile with different prompts and resolutions to identify bottlenecks across workloads ## FAQ - If you are profiling `sglang generate` with Nsight Systems and find that the generated profiler file did not capture any CUDA kernels, you can resolve this issue by increasing the model's inference steps to extend the execution time. ================================================ FILE: docs/diffusion/support_new_models.md ================================================ # How to Support New Diffusion Models This document explains how to add support for new diffusion models in SGLang Diffusion. ## Architecture Overview SGLang Diffusion is engineered for both performance and flexibility, built upon a pipeline architecture. This design allows developers to construct pipelines for various diffusion models while keeping the core generation loop standardized for optimization. At its core, the architecture revolves around two key concepts, as highlighted in our [blog post](https://lmsys.org/blog/2025-11-07-sglang-diffusion/#architecture): - **`ComposedPipeline`**: This class orchestrates a series of `PipelineStage`s to define the complete generation process for a specific model. It acts as the main entry point for a model and manages the data flow between the different stages of the diffusion process. - **`PipelineStage`**: Each stage is a modular component that encapsulates a function within the diffusion process. Examples include prompt encoding, the denoising loop, or VAE decoding. ### Two Pipeline Styles SGLang Diffusion supports two pipeline composition styles. Both are valid; choose the one that best fits your model. #### Style A: Hybrid Monolithic Pipeline (Recommended Default) The recommended default for most new models. Uses a three-stage structure: ``` BeforeDenoisingStage (model-specific) → DenoisingStage (standard) → DecodingStage (standard) ``` | Stage | Ownership | Responsibility | |-------|-----------|----------------| | `{Model}BeforeDenoisingStage` | Model-specific | All pre-processing: input validation, text/image encoding, latent preparation, timestep computation | | `DenoisingStage` | Framework-standard | The denoising loop (DiT/UNet forward passes), shared across all models | | `DecodingStage` | Framework-standard | VAE decoding from latent space to pixel space, shared across all models | **Why recommended?** Modern diffusion models often have highly heterogeneous pre-processing requirements — different text encoders, different latent formats, different conditioning mechanisms. The Hybrid approach keeps pre-processing isolated per model, avoids fragile shared stages with excessive conditional logic, and lets developers port Diffusers reference code quickly. #### Style B: Modular Composition Style Uses the framework's fine-grained standard stages (`TextEncodingStage`, `LatentPreparationStage`, `TimestepPreparationStage`, etc.) to build the pipeline by composition. Convenience methods like `add_standard_t2i_stages()` and `add_standard_ti2i_stages()` make this very concise. This style is appropriate when: - **The new model's pre-processing can largely reuse existing stages** — e.g., a model that uses standard CLIP/T5 text encoding + standard latent preparation with minimal customization. - **A model-specific optimization needs to be extracted as a standalone stage** — e.g., a specialized encoding or conditioning step that benefits from being a separate stage for profiling, parallelism control, or reuse across multiple pipeline variants. #### How to Choose | Situation | Recommended Style | |-----------|-------------------| | Model has unique/complex pre-processing (VLM captioning, AR token generation, custom latent packing, etc.) | **Hybrid** — consolidate into a BeforeDenoisingStage | | Model fits neatly into standard text-to-image or text+image-to-image pattern | **Modular** — use `add_standard_t2i_stages()` / `add_standard_ti2i_stages()` | | Porting a Diffusers pipeline with many custom steps | **Hybrid** — copy the `__call__` logic into a single stage | | Adding a variant of an existing model that shares most logic | **Modular** — reuse existing stages, customize via PipelineConfig callbacks | | A specific pre-processing step needs special parallelism or profiling isolation | **Modular** — extract that step as a dedicated stage | ## Key Components for Implementation To add support for a new diffusion model, you will need to define or configure the following components: 1. **`PipelineConfig`**: A dataclass holding static configurations for your model pipeline — precision settings, model architecture parameters, and callback methods used by the standard `DenoisingStage` and `DecodingStage`. Each model has its own subclass. 2. **`SamplingParams`**: A dataclass defining runtime generation parameters — `prompt`, `negative_prompt`, `guidance_scale`, `num_inference_steps`, `seed`, `height`, `width`, etc. 3. **Pre-processing stage(s)**: Either a single model-specific `{Model}BeforeDenoisingStage` (Hybrid style) or a combination of standard stages (Modular style). See [Two Pipeline Styles](#two-pipeline-styles) above. 4. **`ComposedPipeline`**: A class that wires together your pre-processing stage(s) with the standard `DenoisingStage` and `DecodingStage`. See base definitions: - [`ComposedPipelineBase`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py) - [`PipelineStage`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py) - [Central registry](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/registry.py) 5. **Modules (model components)**: Each pipeline references modules loaded from the model repository (e.g., Diffusers `model_index.json`): - `text_encoder`: Encodes text prompts into embeddings. - `tokenizer`: Tokenizes raw text input for the text encoder(s). - `processor`: Preprocesses images and extracts features; often used in image-to-image tasks. - `image_encoder`: Specialized image feature extractor. - `dit/transformer`: The core denoising network (DiT/UNet architecture) operating in latent space. - `scheduler`: Controls the timestep schedule and denoising dynamics. - `vae`: Variational Autoencoder for encoding/decoding between pixel space and latent space. ## Pipeline Stages Reference ### Core Stages (used by all pipelines) | Stage Class | Description | | -------------------------------- | ------------------------------------------------------------------------------------------------------- | | `DenoisingStage` | Executes the main denoising loop, iteratively applying the model (DiT/UNet) to refine the latents. | | `DecodingStage` | Decodes the final latent tensor back into pixel space using the VAE. | | `DmdDenoisingStage` | A specialized denoising stage for DMD model architectures. | | `CausalDMDDenoisingStage` | A specialized causal denoising stage for specific video models. | ### Pre-processing Stages (for Modular Composition Style) The following fine-grained stages can be composed to build the pre-processing portion of a pipeline. They are best suited for models whose pre-processing largely fits the standard patterns. If your model requires significant customization, consider the Hybrid style with a single `BeforeDenoisingStage` instead. | Stage Class | Description | | -------------------------------- | ------------------------------------------------------------------------------------------------------- | | `InputValidationStage` | Validates user-provided `SamplingParams`. | | `TextEncodingStage` | Encodes text prompts into embeddings using one or more text encoders. | | `ImageEncodingStage` | Encodes input images into embeddings, often used in image-to-image tasks. | | `ImageVAEEncodingStage` | Encodes an input image into latent space using the VAE. | | `TimestepPreparationStage` | Prepares the scheduler's timesteps for the diffusion process. | | `LatentPreparationStage` | Creates the initial noisy latent tensor that will be denoised. | ## Implementation Guide ### Step 1: Obtain and Study the Reference Implementation Before writing any code, obtain the model's original implementation or Diffusers pipeline code: - The model's Diffusers pipeline source (e.g., the `pipeline_*.py` file from the `diffusers` library or HuggingFace repo) - Or the model's official reference implementation (e.g., from the model author's GitHub repo) - Or the HuggingFace model ID to look up `model_index.json` and the associated pipeline class Once you have the reference code, study it thoroughly: 1. Find the model's `model_index.json` to identify required modules. 2. Read the Diffusers pipeline's `__call__` method to understand: - How text prompts are encoded - How latents are prepared (shape, dtype, scaling) - How timesteps/sigmas are computed - What conditioning kwargs the DiT expects - How the denoising loop works - How VAE decoding is done ### Step 2: Evaluate Reuse of Existing Pipelines and Stages Before creating any new files, check whether an existing pipeline or stage can be reused or extended. Only create new pipelines/stages when the existing ones would need substantial structural changes or when no architecturally similar implementation exists. - **Compare against existing pipelines** (Flux, Wan, Qwen-Image, GLM-Image, HunyuanVideo, LTX, etc.). If the new model shares most of its structure with an existing one, prefer adding a new config variant or reusing existing stages. - **Check existing stages** in `runtime/pipelines_core/stages/` and `stages/model_specific_stages/`. - **Check existing model components** — many models share VAEs (e.g., `AutoencoderKL`), text encoders (CLIP, T5), and schedulers. Reuse these directly. ### Step 3: Implement Model Components Adapt the model's core components: - **DiT/Transformer**: Implement in [`runtime/models/dits/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/dits/) - **Encoders**: Implement in [`runtime/models/encoders/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/encoders/) - **VAEs**: Implement in [`runtime/models/vaes/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/vaes/) - **Schedulers**: Implement in [`runtime/models/schedulers/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/schedulers/) if needed Use SGLang's fused kernels where possible (see `LayerNormScaleShift`, `RMSNormScaleShift`, `apply_qk_norm`, etc.). **Tensor Parallel (TP) and Sequence Parallel (SP)**: For multi-GPU deployment, it is recommended to add TP/SP support to the DiT model. This can be done incrementally after the single-GPU implementation is verified. Reference implementations: - **Wan model** (`runtime/models/dits/wanvideo.py`) — Full TP + SP: `ColumnParallelLinear`/`RowParallelLinear` for attention, sequence dimension sharding via `get_sp_world_size()` - **Qwen-Image model** (`runtime/models/dits/qwen_image.py`) — SP via `USPAttention` (Ulysses + Ring Attention) ### Step 4: Create Configs - **DiT Config**: `configs/models/dits/{model_name}.py` - **VAE Config**: `configs/models/vaes/{model_name}.py` - **SamplingParams**: `configs/sample/{model_name}.py` ### Step 5: Create PipelineConfig The `PipelineConfig` provides callbacks that the standard `DenoisingStage` and `DecodingStage` use: ```python # python/sglang/multimodal_gen/configs/pipeline_configs/my_model.py @dataclass class MyModelPipelineConfig(ImagePipelineConfig): task_type: ModelTaskType = ModelTaskType.T2I vae_precision: str = "bf16" should_use_guidance: bool = True dit_config: DiTConfig = field(default_factory=MyModelDitConfig) vae_config: VAEConfig = field(default_factory=MyModelVAEConfig) def get_freqs_cis(self, batch, device, rotary_emb, dtype): """Prepare rotary position embeddings for the DiT.""" ... def prepare_pos_cond_kwargs(self, batch, latent_model_input, t, **kwargs): """Build positive conditioning kwargs for each denoising step.""" return { "hidden_states": latent_model_input, "encoder_hidden_states": batch.prompt_embeds[0], "timestep": t, } def prepare_neg_cond_kwargs(self, batch, latent_model_input, t, **kwargs): """Build negative conditioning kwargs for CFG.""" return { "hidden_states": latent_model_input, "encoder_hidden_states": batch.negative_prompt_embeds[0], "timestep": t, } def get_decode_scale_and_shift(self): """Return (scale, shift) for latent denormalization before VAE decode.""" ... ``` ### Step 6: Implement Pre-processing Choose based on your model's needs (see [How to Choose](#how-to-choose)): #### Option A: BeforeDenoisingStage (Hybrid Style) Create a single stage that handles all pre-processing. Best when the model has custom/complex pre-processing logic. ```python # python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/my_model.py class MyModelBeforeDenoisingStage(PipelineStage): """Monolithic pre-processing stage for MyModel. Consolidates: input validation, text/image encoding, latent preparation, and timestep computation. """ def __init__(self, vae, text_encoder, tokenizer, transformer, scheduler): super().__init__() self.vae = vae self.text_encoder = text_encoder self.tokenizer = tokenizer self.transformer = transformer self.scheduler = scheduler @torch.no_grad() def forward(self, batch: Req, server_args: ServerArgs) -> Req: device = get_local_torch_device() # 1. Encode prompt (model-specific logic) prompt_embeds, negative_prompt_embeds = self._encode_prompt(...) # 2. Prepare latents latents = self._prepare_latents(...) # 3. Prepare timesteps timesteps, sigmas = self._prepare_timesteps(...) # 4. Populate batch for DenoisingStage batch.prompt_embeds = [prompt_embeds] batch.negative_prompt_embeds = [negative_prompt_embeds] batch.latents = latents batch.timesteps = timesteps batch.num_inference_steps = len(timesteps) batch.sigmas = sigmas.tolist() batch.generator = generator batch.raw_latent_shape = latents.shape return batch ``` #### Option B: Standard Stages (Modular Style) Skip creating a custom stage entirely — configure via `PipelineConfig` callbacks and use framework helpers. Best when the model fits standard patterns. (This option has no separate stage file; the pipeline class in Step 7 calls `add_standard_t2i_stages()` directly.) **Key batch fields that `DenoisingStage` expects** (regardless of which option you choose): | Field | Type | Description | |-------|------|-------------| | `batch.latents` | `torch.Tensor` | Initial noisy latent tensor | | `batch.timesteps` | `torch.Tensor` | Timestep schedule | | `batch.num_inference_steps` | `int` | Number of denoising steps | | `batch.sigmas` | `list[float]` | Sigma schedule (must be a Python list, not numpy) | | `batch.prompt_embeds` | `list[torch.Tensor]` | Positive prompt embeddings (wrapped in a list) | | `batch.negative_prompt_embeds` | `list[torch.Tensor]` | Negative prompt embeddings (wrapped in a list) | | `batch.generator` | `torch.Generator` | RNG generator for reproducibility | | `batch.raw_latent_shape` | `tuple` | Original latent shape before any packing | ### Step 7: Define the Pipeline Class #### Hybrid Style ```python # python/sglang/multimodal_gen/runtime/pipelines/my_model.py class MyModelPipeline(LoRAPipeline, ComposedPipelineBase): pipeline_name = "MyModelPipeline" # Must match model_index.json _class_name _required_config_modules = [ "text_encoder", "tokenizer", "vae", "transformer", "scheduler", ] def create_pipeline_stages(self, server_args: ServerArgs): # 1. Monolithic pre-processing (model-specific) self.add_stage( MyModelBeforeDenoisingStage( vae=self.get_module("vae"), text_encoder=self.get_module("text_encoder"), tokenizer=self.get_module("tokenizer"), transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), ), ) # 2. Standard denoising loop (framework-provided) self.add_stage( DenoisingStage( transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), ), ) # 3. Standard VAE decoding (framework-provided) self.add_standard_decoding_stage() EntryClass = [MyModelPipeline] ``` #### Modular Style ```python # python/sglang/multimodal_gen/runtime/pipelines/my_model.py class MyModelPipeline(LoRAPipeline, ComposedPipelineBase): pipeline_name = "MyModelPipeline" _required_config_modules = [ "text_encoder", "tokenizer", "vae", "transformer", "scheduler", ] def create_pipeline_stages(self, server_args: ServerArgs): # All pre-processing + denoising + decoding in one call self.add_standard_t2i_stages( prepare_extra_timestep_kwargs=[prepare_mu], # model-specific hooks ) EntryClass = [MyModelPipeline] ``` ### Step 8: Register the Model Register your configs in [`registry.py`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/registry.py): ```python register_configs( model_family="my_model", sampling_param_cls=MyModelSamplingParams, pipeline_config_cls=MyModelPipelineConfig, hf_model_paths=["org/my-model-name"], ) ``` The `EntryClass` in your pipeline file is automatically discovered by the registry — no additional registration needed for the pipeline class itself. ### Step 9: Verify Output Quality After implementation, verify that the generated output is not noise. A noisy or garbled output is the most common sign of an incorrect implementation. Common causes include: - Incorrect latent scale/shift factors - Wrong timestep/sigma schedule (order, dtype, or value range) - Mismatched conditioning kwargs - Rotary embedding style mismatch (`is_neox_style`) Debug by comparing intermediate tensor values against the Diffusers reference pipeline with the same seed. ## Reference Implementations ### Hybrid Style | Model | Pipeline | BeforeDenoisingStage | PipelineConfig | |-------|----------|---------------------|----------------| | GLM-Image | `runtime/pipelines/glm_image.py` | `stages/model_specific_stages/glm_image.py` | `configs/pipeline_configs/glm_image.py` | | Qwen-Image-Layered | `runtime/pipelines/qwen_image.py` | `stages/model_specific_stages/qwen_image_layered.py` | `configs/pipeline_configs/qwen_image.py` | ### Modular Style | Model | Pipeline | Notes | |-------|----------|-------| | Qwen-Image (T2I) | `runtime/pipelines/qwen_image.py` | Uses `add_standard_t2i_stages()` | | Qwen-Image-Edit | `runtime/pipelines/qwen_image.py` | Uses `add_standard_ti2i_stages()` | | Flux | `runtime/pipelines/flux.py` | Uses `add_standard_t2i_stages()` with custom `prepare_mu` | | Wan | `runtime/pipelines/wan_pipeline.py` | Uses `add_standard_ti2v_stages()` | ## Checklist Before submitting your implementation, verify: **Common (both styles):** - [ ] **Pipeline file** at `runtime/pipelines/{model_name}.py` with `EntryClass` - [ ] **PipelineConfig** at `configs/pipeline_configs/{model_name}.py` - [ ] **SamplingParams** at `configs/sample/{model_name}.py` - [ ] **DiT model** at `runtime/models/dits/{model_name}.py` - [ ] **Model configs** (DiT, VAE) at `configs/models/dits/` and `configs/models/vaes/` - [ ] **Registry entry** in `registry.py` via `register_configs()` - [ ] `pipeline_name` matches Diffusers `model_index.json` `_class_name` - [ ] `_required_config_modules` lists all modules from `model_index.json` - [ ] `PipelineConfig` callbacks (`prepare_pos_cond_kwargs`, etc.) match the DiT's `forward()` signature - [ ] Uses framework-standard `DenoisingStage` and `DecodingStage` (not custom denoising loops) - [ ] **TP/SP support** considered for DiT model (recommended; reference `wanvideo.py` for TP+SP, `qwen_image.py` for USPAttention) - [ ] **Output quality verified** — generated images/videos are not noise; compared against Diffusers reference output **Hybrid style only:** - [ ] **BeforeDenoisingStage** at `stages/model_specific_stages/{model_name}.py` - [ ] `BeforeDenoisingStage.forward()` populates all batch fields required by `DenoisingStage` ================================================ FILE: docs/get_started/install.md ================================================ # Install SGLang You can install SGLang using one of the methods below. This page primarily applies to common NVIDIA GPU platforms. For other or newer platforms, please refer to the dedicated pages for [AMD GPUs](../platforms/amd_gpu.md), [Intel Xeon CPUs](../platforms/cpu_server.md), [TPU](../platforms/tpu.md), [NVIDIA DGX Spark](https://lmsys.org/blog/2025-11-03-gpt-oss-on-nvidia-dgx-spark/), [NVIDIA Jetson](../platforms/nvidia_jetson.md), [Ascend NPUs](../platforms/ascend_npu.md), and [Intel XPU](../platforms/xpu.md). ## Method 1: With pip or uv It is recommended to use uv for faster installation: ```bash pip install --upgrade pip pip install uv uv pip install sglang ``` ### For CUDA 13 Docker is recommended (see Method 3 note on B300/GB300/CUDA 13). If you do not have Docker access, follow these steps: 1. Install PyTorch with CUDA 13 support first: ```bash # Replace X.Y.Z with the version by your SGLang install uv pip install torch==X.Y.Z torchvision torchaudio --index-url https://download.pytorch.org/whl/cu130 ``` 2. Install sglang: ```bash uv pip install sglang ``` 3. Install the `sglang-kernel` wheel for CUDA 13 from [the sgl-project whl releases](https://github.com/sgl-project/whl/blob/gh-pages/cu130/sglang-kernel/index.html). Replace `X.Y.Z` with the `sglang-kernel` version required by your SGLang install (you can find this by running `uv pip show sglang-kernel`). Examples: ```bash # x86_64 uv pip install "https://github.com/sgl-project/whl/releases/download/vX.Y.Z/sglang_kernel-X.Y.Z+cu130-cp310-abi3-manylinux2014_x86_64.whl" # aarch64 uv pip install "https://github.com/sgl-project/whl/releases/download/vX.Y.Z/sglang_kernel-X.Y.Z+cu130-cp310-abi3-manylinux2014_aarch64.whl" ``` ### **Quick fixes to common problems** - If you encounter `OSError: CUDA_HOME environment variable is not set`. Please set it to your CUDA install root with either of the following solutions: 1. Use `export CUDA_HOME=/usr/local/cuda-` to set the `CUDA_HOME` environment variable. 2. Install FlashInfer first following [FlashInfer installation doc](https://docs.flashinfer.ai/installation.html), then install SGLang as described above. ## Method 2: From source ```bash # Use the last release branch git clone -b v0.5.9 https://github.com/sgl-project/sglang.git cd sglang # Install the python packages pip install --upgrade pip pip install -e "python" ``` **Quick fixes to common problems** - If you want to develop SGLang, you can try the dev docker image. Please refer to [setup docker container](../developer_guide/development_guide_using_docker.md#setup-docker-container). The docker image is `lmsysorg/sglang:dev`. ## Method 3: Using docker The docker images are available on Docker Hub at [lmsysorg/sglang](https://hub.docker.com/r/lmsysorg/sglang/tags), built from [Dockerfile](https://github.com/sgl-project/sglang/tree/main/docker). Replace `` below with your huggingface hub [token](https://huggingface.co/docs/hub/en/security-tokens). ```bash docker run --gpus all \ --shm-size 32g \ -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --host 0.0.0.0 --port 30000 ``` For production deployments, use the `runtime` variant which is significantly smaller (~40% reduction) by excluding build tools and development dependencies: ```bash docker run --gpus all \ --shm-size 32g \ -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=" \ --ipc=host \ lmsysorg/sglang:latest-runtime \ python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --host 0.0.0.0 --port 30000 ``` You can also find the nightly docker images [here](https://hub.docker.com/r/lmsysorg/sglang/tags?name=nightly). Notes: - On B300/GB300 (SM103) or CUDA 13 environment, we recommend using the nightly image at `lmsysorg/sglang:dev-cu13` or stable image at `lmsysorg/sglang:latest-cu130-runtime`. Please, do not re-install the project as editable inside the docker image, since it will override the version of libraries specified by the cu13 docker image. ## Method 4: Using Kubernetes Please check out [OME](https://github.com/sgl-project/ome), a Kubernetes operator for enterprise-grade management and serving of large language models (LLMs).
More 1. Option 1: For single node serving (typically when the model size fits into GPUs on one node) Execute command `kubectl apply -f docker/k8s-sglang-service.yaml`, to create k8s deployment and service, with llama-31-8b as example. 2. Option 2: For multi-node serving (usually when a large model requires more than one GPU node, such as `DeepSeek-R1`) Modify the LLM model path and arguments as necessary, then execute command `kubectl apply -f docker/k8s-sglang-distributed-sts.yaml`, to create two nodes k8s statefulset and serving service.
## Method 5: Using docker compose
More > This method is recommended if you plan to serve it as a service. > A better approach is to use the [k8s-sglang-service.yaml](https://github.com/sgl-project/sglang/blob/main/docker/k8s-sglang-service.yaml). 1. Copy the [compose.yml](https://github.com/sgl-project/sglang/blob/main/docker/compose.yaml) to your local machine 2. Execute the command `docker compose up -d` in your terminal.
## Method 6: Run on Kubernetes or Clouds with SkyPilot
More To deploy on Kubernetes or 12+ clouds, you can use [SkyPilot](https://github.com/skypilot-org/skypilot). 1. Install SkyPilot and set up Kubernetes cluster or cloud access: see [SkyPilot's documentation](https://skypilot.readthedocs.io/en/latest/getting-started/installation.html). 2. Deploy on your own infra with a single command and get the HTTP API endpoint:
SkyPilot YAML: sglang.yaml ```yaml # sglang.yaml envs: HF_TOKEN: null resources: image_id: docker:lmsysorg/sglang:latest accelerators: A100 ports: 30000 run: | conda deactivate python3 -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --host 0.0.0.0 \ --port 30000 ```
```bash # Deploy on any cloud or Kubernetes cluster. Use --cloud to select a specific cloud provider. HF_TOKEN= sky launch -c sglang --env HF_TOKEN sglang.yaml # Get the HTTP API endpoint sky status --endpoint 30000 sglang ``` 3. To further scale up your deployment with autoscaling and failure recovery, check out the [SkyServe + SGLang guide](https://github.com/skypilot-org/skypilot/tree/master/llm/sglang#serving-llama-2-with-sglang-for-more-traffic-using-skyserve).
## Method 7: Run on AWS SageMaker
More To deploy on SGLang on AWS SageMaker, check out [AWS SageMaker Inference](https://aws.amazon.com/sagemaker/ai/deploy) Amazon Web Services provide supports for SGLang containers along with routine security patching. For available SGLang containers, check out [AWS SGLang DLCs](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#sglang-containers) To host a model with your own container, follow the following steps: 1. Build a docker container with [sagemaker.Dockerfile](https://github.com/sgl-project/sglang/blob/main/docker/sagemaker.Dockerfile) alongside the [serve](https://github.com/sgl-project/sglang/blob/main/docker/serve) script. 2. Push your container onto AWS ECR.
Dockerfile Build Script: build-and-push.sh ```bash #!/bin/bash AWS_ACCOUNT="" AWS_REGION="" REPOSITORY_NAME="" IMAGE_TAG="" ECR_REGISTRY="${AWS_ACCOUNT}.dkr.ecr.${AWS_REGION}.amazonaws.com" IMAGE_URI="${ECR_REGISTRY}/${REPOSITORY_NAME}:${IMAGE_TAG}" echo "Starting build and push process..." # Login to ECR echo "Logging into ECR..." aws ecr get-login-password --region ${AWS_REGION} | docker login --username AWS --password-stdin ${ECR_REGISTRY} # Build the image echo "Building Docker image..." docker build -t ${IMAGE_URI} -f sagemaker.Dockerfile . echo "Pushing ${IMAGE_URI}" docker push ${IMAGE_URI} echo "Build and push completed successfully!" ```
3. Deploy a model for serving on AWS Sagemaker, refer to [deploy_and_serve_endpoint.py](https://github.com/sgl-project/sglang/blob/main/examples/sagemaker/deploy_and_serve_endpoint.py). For more information, check out [sagemaker-python-sdk](https://github.com/aws/sagemaker-python-sdk). 1. By default, the model server on SageMaker will run with the following command: `python3 -m sglang.launch_server --model-path opt/ml/model --host 0.0.0.0 --port 8080`. This is optimal for hosting your own model with SageMaker. 2. To modify your model serving parameters, the [serve](https://github.com/sgl-project/sglang/blob/main/docker/serve) script allows for all available options within `python3 -m sglang.launch_server --help` cli by specifying environment variables with prefix `SM_SGLANG_`. 3. The serve script will automatically convert all environment variables with prefix `SM_SGLANG_` from `SM_SGLANG_INPUT_ARGUMENT` into `--input-argument` to be parsed into `python3 -m sglang.launch_server` cli. 4. For example, to run [Qwen/Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B) with reasoning parser, simply add additional environment variables `SM_SGLANG_MODEL_PATH=Qwen/Qwen3-0.6B` and `SM_SGLANG_REASONING_PARSER=qwen3`.
## Common Notes - [FlashInfer](https://github.com/flashinfer-ai/flashinfer) is the default attention kernel backend. It only supports sm75 and above. If you encounter any FlashInfer-related issues on sm75+ devices (e.g., T4, A10, A100, L4, L40S, H100), please switch to other kernels by adding `--attention-backend triton --sampling-backend pytorch` and open an issue on GitHub. - To reinstall flashinfer locally, use the following command: `pip3 install --upgrade flashinfer-python --force-reinstall --no-deps` and then delete the cache with `rm -rf ~/.cache/flashinfer`. - When encountering `ptxas fatal : Value 'sm_103a' is not defined for option 'gpu-name'` on B300/GB300, fix it with `export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas`. ================================================ FILE: docs/index.rst ================================================ SGLang Documentation ==================== .. raw:: html Star Fork

SGLang is a high-performance serving framework for large language models and multimodal models. It is designed to deliver low-latency and high-throughput inference across a wide range of setups, from a single GPU to large distributed clusters. Its core features include: - **Fast Runtime**: Provides efficient serving with RadixAttention for prefix caching, a zero-overhead CPU scheduler, prefill-decode disaggregation, speculative decoding, continuous batching, paged attention, tensor/pipeline/expert/data parallelism, structured outputs, chunked prefill, quantization (FP4/FP8/INT4/AWQ/GPTQ), and multi-LoRA batching. - **Broad Model Support**: Supports a wide range of language models (Llama, Qwen, DeepSeek, Kimi, GLM, GPT, Gemma, Mistral, etc.), embedding models (e5-mistral, gte, mcdse), reward models (Skywork), and diffusion models (WAN, Qwen-Image), with easy extensibility for adding new models. Compatible with most Hugging Face models and OpenAI APIs. - **Extensive Hardware Support**: Runs on NVIDIA GPUs (GB200/B300/H100/A100/Spark), AMD GPUs (MI355/MI300), Intel Xeon CPUs, Google TPUs, Ascend NPUs, and more. - **Active Community**: SGLang is open-source and supported by a vibrant community with widespread industry adoption, powering over 400,000 GPUs worldwide. - **RL & Post-Training Backbone**: SGLang is a proven rollout backend used for training many frontier models, with native RL integrations and adoption by well-known post-training frameworks such as AReaL, Miles, slime, Tunix, verl and more. .. toctree:: :maxdepth: 1 :caption: Get Started get_started/install.md .. toctree:: :maxdepth: 1 :caption: Basic Usage basic_usage/send_request.ipynb basic_usage/openai_api.rst basic_usage/ollama_api.md basic_usage/offline_engine_api.ipynb basic_usage/native_api.ipynb basic_usage/sampling_params.md basic_usage/popular_model_usage.rst .. toctree:: :maxdepth: 1 :caption: Advanced Features advanced_features/server_arguments.md advanced_features/hyperparameter_tuning.md advanced_features/attention_backend.md advanced_features/speculative_decoding.ipynb advanced_features/structured_outputs.ipynb advanced_features/structured_outputs_for_reasoning_models.ipynb advanced_features/tool_parser.ipynb advanced_features/separate_reasoning.ipynb advanced_features/quantization.md advanced_features/quantized_kv_cache.md advanced_features/expert_parallelism.md advanced_features/dp_dpa_smg_guide.md advanced_features/lora.ipynb advanced_features/pd_disaggregation.md advanced_features/epd_disaggregation.md advanced_features/pipeline_parallelism.md advanced_features/hicache.rst advanced_features/pd_multiplexing.md advanced_features/vlm_query.ipynb advanced_features/dp_for_multi_modal_encoder.md advanced_features/cuda_graph_for_multi_modal_encoder.md advanced_features/piecewise_cuda_graph.md advanced_features/sgl_model_gateway.md advanced_features/deterministic_inference.md advanced_features/observability.md advanced_features/checkpoint_engine.md advanced_features/sglang_for_rl.md .. toctree:: :maxdepth: 2 :caption: Supported Models supported_models/text_generation/index supported_models/retrieval_ranking/index supported_models/specialized/index supported_models/extending/index .. toctree:: :maxdepth: 2 :caption: SGLang Diffusion diffusion/index diffusion/installation diffusion/compatibility_matrix diffusion/api/cli diffusion/api/openai_api diffusion/performance/index diffusion/performance/attention_backends diffusion/performance/profiling diffusion/performance/cache/index diffusion/performance/cache/cache_dit diffusion/performance/cache/teacache diffusion/support_new_models diffusion/contributing diffusion/ci_perf diffusion/environment_variables .. toctree:: :maxdepth: 1 :caption: Hardware Platforms platforms/amd_gpu.md platforms/cpu_server.md platforms/tpu.md platforms/nvidia_jetson.md platforms/ascend_npu_support.rst platforms/xpu.md .. toctree:: :maxdepth: 1 :caption: Developer Guide developer_guide/contribution_guide.md developer_guide/development_guide_using_docker.md developer_guide/development_jit_kernel_guide.md developer_guide/benchmark_and_profiling.md developer_guide/bench_serving.md developer_guide/evaluating_new_models.md .. toctree:: :maxdepth: 1 :caption: References references/faq.md references/environment_variables.md references/production_metrics.md references/production_request_trace.md references/multi_node_deployment/multi_node_index.rst references/custom_chat_template.md references/frontend/frontend_index.rst references/post_training_integration.md references/release_lookup references/learn_more.md .. toctree:: :maxdepth: 1 :caption: Security Acknowledgement security/acknowledgements.md ================================================ FILE: docs/performance_dashboard/README.md ================================================ # SGLang Performance Dashboard A web-based dashboard for visualizing SGLang nightly test performance metrics. ## Features - **Performance Trends**: View throughput, latency, and TTFT trends over time - **Model Comparison**: Compare performance across different models and configurations - **Filtering**: Filter by GPU configuration, model, variant, and batch size - **Interactive Charts**: Zoom, pan, and hover for detailed metrics - **Run History**: View recent benchmark runs with links to GitHub Actions ## Quick Start ### Option 1: Run with Local Server (Recommended) For live data from GitHub Actions artifacts: ```bash # Install requirements pip install requests # Run the server python server.py --fetch-on-start # Visit http://localhost:8000 ``` The server provides: - Automatic fetching of metrics from GitHub - Caching to reduce API calls - `/api/metrics` endpoint for the frontend ### Option 2: Fetch Data Manually Use the fetch script to download metrics data: ```bash # Fetch last 30 days of metrics python fetch_metrics.py --output metrics_data.json # Fetch a specific run python fetch_metrics.py --run-id 21338741812 --output single_run.json # Fetch only scheduled (nightly) runs python fetch_metrics.py --scheduled-only --days 7 ``` ## GitHub Token To download artifacts from GitHub, you need authentication: 1. **Using `gh` CLI** (recommended): ```bash gh auth login ``` 2. **Using environment variable**: ```bash export GITHUB_TOKEN=your_token_here ``` Without a token, the dashboard will show run metadata but not detailed benchmark results. ## Data Structure The metrics JSON has this structure: ```json { "run_id": "21338741812", "run_date": "2026-01-25T22:24:02.090218+00:00", "commit_sha": "5cdb391...", "branch": "main", "results": [ { "gpu_config": "8-gpu-h200", "partition": 0, "model": "deepseek-ai/DeepSeek-V3.1", "variant": "TP8+MTP", "benchmarks": [ { "batch_size": 1, "input_len": 4096, "output_len": 512, "latency_ms": 2400.72, "input_throughput": 21408.64, "output_throughput": 231.74, "overall_throughput": 1919.43, "ttft_ms": 191.32, "acc_length": 3.19 } ] } ] } ``` ## Deployment ### GitHub Pages The dashboard can be deployed to GitHub Pages for public access: 1. Copy the dashboard files to `docs/performance_dashboard/` 2. Enable GitHub Pages in repository settings 3. Set up a GitHub Action to periodically update metrics data ### Self-Hosted For a self-hosted deployment with live data: 1. Set up a server running `server.py` 2. Configure a cron job or systemd timer to refresh data 3. Optionally put behind nginx/caddy for SSL ## Metrics Explained - **Overall Throughput**: Total tokens (input + output) processed per second - **Input Throughput**: Input tokens processed per second (prefill speed) - **Output Throughput**: Output tokens generated per second (decode speed) - **Latency**: End-to-end time to complete the request - **TTFT**: Time to First Token - time until the first output token - **Acc Length**: Acceptance length for speculative decoding (MTP variants) ## Contributing To add support for new metrics or visualizations: 1. Update `fetch_metrics.py` if data collection needs changes 2. Modify `app.js` to add new chart types or filters 3. Update `index.html` for UI changes ## Troubleshooting **No data displayed** - Check browser console for errors - Verify GitHub API is accessible - Try running with `server.py --fetch-on-start` **API rate limits** - Use a GitHub token for higher limits - The server caches data for 5 minutes **Charts not rendering** - Ensure Chart.js is loading from CDN - Check for JavaScript errors in console ================================================ FILE: docs/performance_dashboard/app.js ================================================ // SGLang Performance Dashboard Application const GITHUB_REPO = 'sgl-project/sglang'; const WORKFLOW_NAME = 'nightly-test-nvidia.yml'; const ARTIFACT_PREFIX = 'consolidated-metrics-'; // Chart instances (array for batch-separated charts) let activeCharts = []; // Data storage let allMetricsData = []; let currentModel = null; let currentMetricType = 'throughput'; // throughput, latency, ttft, inputThroughput // Metric type definitions const metricTypes = { // Text/VLM metrics throughput: { label: 'Overall Throughput', unit: 'tokens/sec', field: 'throughput', type: 'text' }, outputThroughput: { label: 'Output Throughput', unit: 'tokens/sec', field: 'outputThroughput', type: 'text' }, inputThroughput: { label: 'Input Throughput', unit: 'tokens/sec', field: 'inputThroughput', type: 'text' }, latency: { label: 'Latency', unit: 'ms', field: 'latency', type: 'text' }, ttft: { label: 'Time to First Token', unit: 'ms', field: 'ttft', type: 'text' }, accLength: { label: 'Accept Length', unit: 'tokens', field: 'accLength', filterInvalid: true, type: 'text' }, // Diffusion metrics e2eMs: { label: 'End-to-End Time', unit: 'ms', field: 'e2e_ms', type: 'diffusion' }, avgDenoiseMs: { label: 'Avg Denoise Time', unit: 'ms', field: 'avg_denoise_ms', type: 'diffusion' }, medianDenoiseMs: { label: 'Median Denoise Time', unit: 'ms', field: 'median_denoise_ms', type: 'diffusion' } }; // Chart.js default configuration for dark theme Chart.defaults.color = '#94a3b8'; Chart.defaults.borderColor = '#1e293b'; const chartColors = [ '#22d3ee', '#34d399', '#fbbf24', '#f87171', '#a78bfa', '#67e8f9', '#6ee7b7', '#fcd34d', '#fca5a5', '#c4b5fd' ]; // Initialize the dashboard async function init() { try { await loadData(); document.getElementById('loading').style.display = 'none'; document.getElementById('content').style.display = 'block'; populateFilters(); updateStats(); updateCharts(); updateRunsTable(); } catch (error) { console.error('Failed to initialize dashboard:', error); document.getElementById('loading').style.display = 'none'; document.getElementById('error').style.display = 'block'; document.getElementById('error-message').textContent = error.message; } } // Load data from local server API or GitHub async function loadData() { // Try local server API first (if running server.py) try { const response = await fetch('/api/metrics', { headers: getAuthHeaders() }); if (response.ok) { const data = await response.json(); if (data.length > 0 && data[0].results && data[0].results.length > 0) { allMetricsData = data; console.log(`Loaded ${data.length} records from local API`); allMetricsData.sort((a, b) => new Date(b.run_date) - new Date(a.run_date)); return; } } } catch (error) { console.log('Local API not available, trying GitHub API'); } // Try to load from GitHub API const runs = await fetchWorkflowRuns(); const metricsPromises = runs.map(run => fetchMetricsForRun(run)); const results = await Promise.allSettled(metricsPromises); allMetricsData = results .filter(r => r.status === 'fulfilled' && r.value !== null) .map(r => r.value); if (allMetricsData.length === 0) { throw new Error('No metrics data available. Please run the server.py with --fetch-on-start to fetch data from GitHub.'); } // Sort by date descending allMetricsData.sort((a, b) => new Date(b.run_date) - new Date(a.run_date)); } // Fetch workflow runs from GitHub API async function fetchWorkflowRuns() { const response = await fetch( `https://api.github.com/repos/${GITHUB_REPO}/actions/workflows/${WORKFLOW_NAME}/runs?status=completed&per_page=30`, { headers: { 'Accept': 'application/vnd.github.v3+json' } } ); if (!response.ok) { throw new Error(`GitHub API error: ${response.status}`); } const data = await response.json(); return data.workflow_runs || []; } // Fetch metrics artifact for a specific run async function fetchMetricsForRun(run) { try { // Get artifacts for this run const artifactsResponse = await fetch( `https://api.github.com/repos/${GITHUB_REPO}/actions/runs/${run.id}/artifacts`, { headers: { 'Accept': 'application/vnd.github.v3+json' } } ); if (!artifactsResponse.ok) return null; const artifactsData = await artifactsResponse.json(); const metricsArtifact = artifactsData.artifacts.find( a => a.name.startsWith(ARTIFACT_PREFIX) ); if (!metricsArtifact) return null; // Note: GitHub API doesn't allow direct artifact download without authentication // For public access, we would need to use a proxy or pre-process the data // For now, return run metadata - in production, use a backend to fetch artifacts return { run_id: run.id.toString(), run_date: run.created_at, commit_sha: run.head_sha, branch: run.head_branch, artifact_id: metricsArtifact.id, results: [] // Would be populated from artifact content }; } catch (error) { console.warn(`Failed to fetch metrics for run ${run.id}:`, error); return null; } } // Helper function to detect if result is diffusion type function isDiffusionResult(result) { return result.test_type === 'diffusion' || (result.tests && !result.benchmarks); } // Populate filter dropdowns function populateFilters() { const gpuConfigs = new Set(); const models = new Set(); const testNames = new Set(); // For diffusion tests const batchSizes = new Set(); const ioLengths = new Set(); allMetricsData.forEach(run => { run.results.forEach(result => { gpuConfigs.add(result.gpu_config); // Handle diffusion results if (isDiffusionResult(result)) { models.add(result.test_suite || 'diffusion'); if (result.tests) { result.tests.forEach(test => { testNames.add(test.test_name); }); } } // Handle text/VLM results else { models.add(result.model); // Try new structure first (benchmarks_by_io_len), fall back to flat benchmarks if (result.benchmarks_by_io_len) { Object.entries(result.benchmarks_by_io_len).forEach(([ioKey, ioData]) => { ioLengths.add(ioKey); ioData.benchmarks.forEach(bench => { batchSizes.add(bench.batch_size); }); }); } else if (result.benchmarks) { result.benchmarks.forEach(bench => { batchSizes.add(bench.batch_size); if (bench.input_len && bench.output_len) { ioLengths.add(`${bench.input_len}_${bench.output_len}`); } }); } } }); }); // No "all" option for GPU and Model - populate with first value selected const gpuArray = Array.from(gpuConfigs).sort(); const modelArray = Array.from(models).sort(); populateSelectNoAll('gpu-filter', gpuArray); populateSelectNoAll('model-filter', modelArray); populateSelect('batch-filter', Array.from(batchSizes).sort((a, b) => a - b)); populateSelectWithLabels('io-len-filter', sortIoLengths(Array.from(ioLengths)), formatIoLenLabel); // Set initial values (first option) if (gpuArray.length > 0) { document.getElementById('gpu-filter').value = gpuArray[0]; } if (modelArray.length > 0) { document.getElementById('model-filter').value = modelArray[0]; currentModel = modelArray[0]; } // Update variants based on selected model updateVariantFilter(); // Update IO length filter based on selected GPU/model updateIoLenFilter(); // Create metric type tabs createMetricTabs(); } // Format input/output length key for display function formatIoLenLabel(ioKey) { if (!ioKey) return 'Unknown'; const parts = ioKey.split('_'); if (parts.length === 2) { return `In: ${parts[0]}, Out: ${parts[1]}`; } return ioKey; } // Sort IO length keys numerically (by input length, then output length) function sortIoLengths(ioLengths) { return ioLengths.filter(key => key && key.includes('_')).sort((a, b) => { const [aIn, aOut] = a.split('_').map(Number); const [bIn, bOut] = b.split('_').map(Number); if (isNaN(aIn) || isNaN(bIn)) return 0; return (aIn - bIn) || (aOut - bOut); }); } // Populate select with custom label formatting function populateSelectWithLabels(selectId, options, labelFormatter) { const select = document.getElementById(selectId); options.forEach(option => { const opt = document.createElement('option'); opt.value = option; opt.textContent = labelFormatter ? labelFormatter(option) : option; select.appendChild(opt); }); } // Update IO length filter based on selected GPU and model function updateIoLenFilter() { const gpuFilterEl = document.getElementById('gpu-filter'); const modelFilterEl = document.getElementById('model-filter'); const ioLenSelect = document.getElementById('io-len-filter'); if (!gpuFilterEl || !modelFilterEl || !ioLenSelect) return; const gpuFilter = gpuFilterEl.value; const modelFilter = modelFilterEl.value; const ioLengths = new Set(); allMetricsData.forEach(run => { run.results.forEach(result => { if (result.gpu_config === gpuFilter && result.model === modelFilter) { if (result.benchmarks_by_io_len) { Object.keys(result.benchmarks_by_io_len).forEach(ioKey => { ioLengths.add(ioKey); }); } else if (result.benchmarks) { result.benchmarks.forEach(bench => { if (bench.input_len && bench.output_len) { ioLengths.add(`${bench.input_len}_${bench.output_len}`); } }); } } }); }); const ioLenArray = sortIoLengths(Array.from(ioLengths)); const currentIoLen = ioLenSelect.value; // Clear and repopulate ioLenSelect.innerHTML = ''; ioLenArray.forEach(ioLen => { const opt = document.createElement('option'); opt.value = ioLen; opt.textContent = formatIoLenLabel(ioLen); ioLenSelect.appendChild(opt); }); // Try to restore previous selection if still valid if (ioLenArray.includes(currentIoLen)) { ioLenSelect.value = currentIoLen; } else { ioLenSelect.value = 'all'; } } // Update variant filter based on selected GPU and model function updateVariantFilter() { const gpuFilter = document.getElementById('gpu-filter').value; const modelFilter = document.getElementById('model-filter').value; const variants = new Set(); allMetricsData.forEach(run => { run.results.forEach(result => { if (result.gpu_config === gpuFilter && result.model === modelFilter) { // Use 'default' for null/undefined variants variants.add(result.variant || 'default'); } }); }); const variantArray = Array.from(variants).sort(); const variantSelect = document.getElementById('variant-filter'); const currentVariant = variantSelect.value; // Clear and repopulate variantSelect.innerHTML = ''; variantArray.forEach(variant => { const opt = document.createElement('option'); opt.value = variant; opt.textContent = variant; variantSelect.appendChild(opt); }); // Try to restore previous selection if still valid if (variantArray.includes(currentVariant)) { variantSelect.value = currentVariant; } else { variantSelect.value = 'all'; } } function populateSelect(selectId, options) { const select = document.getElementById(selectId); options.forEach(option => { const opt = document.createElement('option'); opt.value = option; opt.textContent = option; select.appendChild(opt); }); } function populateSelectNoAll(selectId, options) { const select = document.getElementById(selectId); // Remove the "all" option if present while (select.options.length > 0) { select.remove(0); } options.forEach(option => { const opt = document.createElement('option'); opt.value = option; opt.textContent = option; select.appendChild(opt); }); } function createMetricTabs() { const tabsContainer = document.getElementById('metric-tabs'); tabsContainer.innerHTML = ''; // Detect if current data is diffusion or text const isDiffusion = detectCurrentDataType() === 'diffusion'; const dataType = isDiffusion ? 'diffusion' : 'text'; // Filter metrics based on data type const relevantMetrics = Object.entries(metricTypes).filter(([key, metric]) => metric.type === dataType ); relevantMetrics.forEach(([key, metric], index) => { const tab = document.createElement('div'); tab.className = index === 0 ? 'tab active' : 'tab'; tab.textContent = metric.label; tab.dataset.metric = key; tab.onclick = () => selectMetricTab(key, tab); tabsContainer.appendChild(tab); }); // Set initial metric type if (relevantMetrics.length > 0) { currentMetricType = relevantMetrics[0][0]; } } function detectCurrentDataType() { // Check if currently selected model/GPU config has diffusion data const gpuFilter = document.getElementById('gpu-filter')?.value; const modelFilter = currentModel; if (!gpuFilter || !modelFilter) return 'text'; for (const run of allMetricsData) { for (const result of run.results) { if (result.gpu_config === gpuFilter) { const resultModel = result.test_suite || result.model; if (resultModel === modelFilter && isDiffusionResult(result)) { return 'diffusion'; } } } } return 'text'; } function selectMetricTab(metricKey, tabElement) { document.querySelectorAll('.tab').forEach(t => t.classList.remove('active')); tabElement.classList.add('active'); currentMetricType = metricKey; // Update chart title const metric = metricTypes[metricKey]; document.getElementById('metric-title').textContent = `${metric.label} (${metric.unit})`; updateCharts(); } // Handle model filter dropdown change function handleModelFilterChange(model) { currentModel = model; // Update variant filter based on new model selection updateVariantFilter(); // Update IO length filter based on new model selection updateIoLenFilter(); // Recreate metric tabs in case data type changed (text vs diffusion) createMetricTabs(); updateCharts(); } // Handle GPU filter change function handleGpuFilterChange() { // Update variant filter based on new GPU selection updateVariantFilter(); // Update IO length filter based on new GPU selection updateIoLenFilter(); // Recreate metric tabs in case data type changed (text vs diffusion) createMetricTabs(); updateCharts(); } // Update summary stats function updateStats() { const statsRow = document.getElementById('stats-row'); const latestRun = allMetricsData[0]; if (!latestRun) { statsRow.innerHTML = ''; const noDataDiv = document.createElement('div'); noDataDiv.className = 'no-data'; noDataDiv.textContent = 'No data available'; statsRow.appendChild(noDataDiv); return; } const totalModels = new Set(latestRun.results.map(r => r.model)).size; const totalBenchmarks = latestRun.results.reduce((sum, r) => { // Count benchmarks from either structure if (r.benchmarks_by_io_len) { return sum + Object.values(r.benchmarks_by_io_len).reduce( (ioSum, ioData) => ioSum + ioData.benchmarks.length, 0 ); } return sum + (r.benchmarks ? r.benchmarks.length : 0); }, 0); statsRow.innerHTML = ''; // Clear previous stats const addStat = (label, value) => { const card = document.createElement('div'); card.className = 'stat-card'; const labelEl = document.createElement('div'); labelEl.className = 'label'; labelEl.textContent = label; const valueEl = document.createElement('div'); valueEl.className = 'value'; valueEl.textContent = value; card.appendChild(labelEl); card.appendChild(valueEl); statsRow.appendChild(card); }; addStat('Total Runs', allMetricsData.length); addStat('Models Tested', totalModels); addStat('Benchmarks', totalBenchmarks); } // Update charts based on current filters and selected metric type function updateCharts() { const gpuFilter = document.getElementById('gpu-filter').value; const modelFilter = currentModel; const variantFilter = document.getElementById('variant-filter').value; const ioLenFilter = document.getElementById('io-len-filter').value; const batchFilter = document.getElementById('batch-filter').value; // Prepare data for charts - grouped by batch size const chartDataByBatch = prepareChartDataByBatch(gpuFilter, modelFilter, variantFilter, ioLenFilter, batchFilter); // Update chart for the selected metric type updateMetricChart(chartDataByBatch, currentMetricType); } function prepareChartData(gpuFilter, modelFilter, variantFilter, ioLenFilter, batchFilter) { const seriesMap = new Map(); allMetricsData.forEach(run => { const runDate = new Date(run.run_date); run.results.forEach(result => { // Apply filters if (result.gpu_config !== gpuFilter) return; if (result.model !== modelFilter) return; if (variantFilter !== 'all' && result.variant !== variantFilter) return; // Helper function to process a benchmark entry const processBenchmark = (bench, ioKey) => { if (batchFilter !== 'all' && bench.batch_size !== parseInt(batchFilter)) return; const ioLabel = ioKey ? `, ${formatIoLenLabel(ioKey)}` : ''; const seriesKey = `${result.model.split('/').pop()} (${result.variant}, BS=${bench.batch_size}${ioLabel})`; if (!seriesMap.has(seriesKey)) { seriesMap.set(seriesKey, { label: seriesKey, data: [], model: result.model, variant: result.variant, batchSize: bench.batch_size, ioKey: ioKey }); } seriesMap.get(seriesKey).data.push({ x: runDate, throughput: bench.overall_throughput, outputThroughput: bench.output_throughput, latency: bench.latency_ms, ttft: bench.ttft_ms, inputThroughput: bench.input_throughput, accLength: bench.acc_length, runId: run.run_id }); }; // Use benchmarks_by_io_len if available if (result.benchmarks_by_io_len) { Object.entries(result.benchmarks_by_io_len).forEach(([ioKey, ioData]) => { if (ioLenFilter !== 'all' && ioKey !== ioLenFilter) return; ioData.benchmarks.forEach(bench => processBenchmark(bench, ioKey)); }); } else if (result.benchmarks) { result.benchmarks.forEach(bench => { const benchIoKey = bench.input_len && bench.output_len ? `${bench.input_len}_${bench.output_len}` : null; if (ioLenFilter !== 'all' && benchIoKey !== ioLenFilter) return; processBenchmark(bench, benchIoKey); }); } }); }); // Sort data points by date seriesMap.forEach(series => { series.data.sort((a, b) => a.x - b.x); }); return Array.from(seriesMap.values()); } // Prepare chart data grouped by batch size - each batch size is a separate series function prepareChartDataByBatch(gpuFilter, modelFilter, variantFilter, ioLenFilter, batchFilter) { const batchDataMap = new Map(); // batch_size -> Map of variant -> data const testDataMap = new Map(); // For diffusion: test_name -> data allMetricsData.forEach(run => { const runDate = new Date(run.run_date); run.results.forEach(result => { // Apply filters - GPU and Model are required (no "all" option) if (result.gpu_config !== gpuFilter) return; // Handle diffusion results if (isDiffusionResult(result)) { const resultModel = result.test_suite || 'diffusion'; if (resultModel !== modelFilter) return; if (result.tests) { result.tests.forEach(test => { const testName = test.test_name; if (!testDataMap.has(testName)) { testDataMap.set(testName, { label: testName, data: [], model: resultModel, testName: testName }); } testDataMap.get(testName).data.push({ x: runDate, e2e_ms: test.e2e_ms, avg_denoise_ms: test.avg_denoise_ms, median_denoise_ms: test.median_denoise_ms, runId: run.run_id }); }); } return; } // Handle text/VLM results if (result.model !== modelFilter) return; if (variantFilter !== 'all' && result.variant !== variantFilter) return; // Use benchmarks_by_io_len if available, otherwise fall back to flat benchmarks if (result.benchmarks_by_io_len) { Object.entries(result.benchmarks_by_io_len).forEach(([ioKey, ioData]) => { // Apply IO length filter if (ioLenFilter !== 'all' && ioKey !== ioLenFilter) return; ioData.benchmarks.forEach(bench => { if (batchFilter !== 'all' && bench.batch_size !== parseInt(batchFilter)) return; const batchSize = bench.batch_size; const variantLabel = result.variant || 'default'; // Include IO length in series key when showing all lengths const seriesKey = ioLenFilter === 'all' ? `${variantLabel} (${formatIoLenLabel(ioKey)})` : variantLabel; if (!batchDataMap.has(batchSize)) { batchDataMap.set(batchSize, new Map()); } const variantMap = batchDataMap.get(batchSize); if (!variantMap.has(seriesKey)) { variantMap.set(seriesKey, { label: seriesKey, data: [], model: result.model, variant: result.variant, batchSize: batchSize, ioKey: ioKey }); } variantMap.get(seriesKey).data.push({ x: runDate, throughput: bench.overall_throughput, outputThroughput: bench.output_throughput, latency: bench.latency_ms, ttft: bench.ttft_ms, inputThroughput: bench.input_throughput, accLength: bench.acc_length, runId: run.run_id }); }); }); } else if (result.benchmarks) { // Fall back to flat benchmarks for backward compatibility result.benchmarks.forEach(bench => { // Apply IO length filter using flat structure const benchIoKey = bench.input_len && bench.output_len ? `${bench.input_len}_${bench.output_len}` : null; if (ioLenFilter !== 'all' && benchIoKey !== ioLenFilter) return; if (batchFilter !== 'all' && bench.batch_size !== parseInt(batchFilter)) return; const batchSize = bench.batch_size; const variantLabel = result.variant || 'default'; // Include IO length in series key when showing all lengths const seriesKey = ioLenFilter === 'all' && benchIoKey ? `${variantLabel} (${formatIoLenLabel(benchIoKey)})` : variantLabel; if (!batchDataMap.has(batchSize)) { batchDataMap.set(batchSize, new Map()); } const variantMap = batchDataMap.get(batchSize); if (!variantMap.has(seriesKey)) { variantMap.set(seriesKey, { label: seriesKey, data: [], model: result.model, variant: result.variant, batchSize: batchSize, ioKey: benchIoKey }); } variantMap.get(seriesKey).data.push({ x: runDate, throughput: bench.overall_throughput, outputThroughput: bench.output_throughput, latency: bench.latency_ms, ttft: bench.ttft_ms, inputThroughput: bench.input_throughput, accLength: bench.acc_length, runId: run.run_id }); }); } }); }); // Sort data points by date and convert to array format const result = {}; // For diffusion data, use test names as "batch sizes" if (testDataMap.size > 0) { testDataMap.forEach((series, testName) => { series.data.sort((a, b) => a.x - b.x); result[testName] = [series]; // Each test is its own series }); return result; } // For text/VLM data, use batch sizes batchDataMap.forEach((variantMap, batchSize) => { variantMap.forEach(series => { series.data.sort((a, b) => a.x - b.x); }); result[batchSize] = Array.from(variantMap.values()); }); return result; } // Unified chart update function for any metric type function updateMetricChart(chartDataByBatch, metricType) { const container = document.getElementById('charts-container'); container.innerHTML = ''; // Destroy existing charts activeCharts.forEach(chart => chart.destroy()); activeCharts = []; const metric = metricTypes[metricType]; const isDiffusion = metric.type === 'diffusion'; // For diffusion, keys are test names; for text, keys are batch sizes const keys = Object.keys(chartDataByBatch); if (!isDiffusion) { keys.sort((a, b) => parseInt(a) - parseInt(b)); } else { keys.sort(); // Alphabetical sort for test names } const batchSizes = keys; // Keep variable name for compatibility if (batchSizes.length === 0) { container.innerHTML = '
No data available for the selected filters
'; return; } let hasAnyData = false; batchSizes.forEach(batchSize => { const chartData = chartDataByBatch[batchSize]; const ctx_datasets = chartData.map((series, index) => { // Filter data points - for metrics like accLength, exclude invalid values (-1 or null) let dataPoints = series.data.map(d => ({ x: d.x, y: d[metric.field] })); if (metric.filterInvalid) { dataPoints = dataPoints.filter(d => d.y != null && d.y !== -1 && d.y > 0); } return { label: series.label, data: dataPoints, borderColor: chartColors[index % chartColors.length], backgroundColor: chartColors[index % chartColors.length] + '20', tension: 0.1, fill: false }; }).filter(dataset => dataset.data.length > 0); // Remove empty datasets // Skip this batch size if no valid data if (ctx_datasets.length === 0) { return; } hasAnyData = true; const chartWrapper = document.createElement('div'); chartWrapper.className = 'batch-chart-wrapper'; const title = document.createElement('div'); title.className = 'batch-chart-title'; // For diffusion, show test name; for text, show batch size title.textContent = isDiffusion ? `Test: ${batchSize}` : `Batch Size: ${batchSize}`; chartWrapper.appendChild(title); const chartContainer = document.createElement('div'); chartContainer.className = 'chart-container'; const canvas = document.createElement('canvas'); chartContainer.appendChild(canvas); chartWrapper.appendChild(chartContainer); container.appendChild(chartWrapper); const ctx = canvas.getContext('2d'); const chart = new Chart(ctx, { type: 'line', data: { datasets: ctx_datasets }, options: getChartOptions(metric.unit) }); activeCharts.push(chart); }); // Show message if no valid data for this metric if (!hasAnyData) { container.innerHTML = `
No valid ${metric.label.toLowerCase()} data available for the selected filters
`; } } function getChartOptions(yAxisLabel) { return { responsive: true, maintainAspectRatio: false, interaction: { mode: 'index', intersect: false }, plugins: { legend: { position: 'bottom', labels: { boxWidth: 12, padding: 10, font: { size: 11 } } }, tooltip: { backgroundColor: '#1a2332', borderColor: 'rgba(148, 163, 184, 0.1)', borderWidth: 1, titleFont: { size: 13, family: "'DM Sans', sans-serif" }, bodyFont: { size: 12, family: "'JetBrains Mono', monospace" }, padding: 14, cornerRadius: 8 } }, scales: { x: { type: 'time', time: { unit: 'day', displayFormats: { day: 'MMM d' } }, grid: { color: 'rgba(148, 163, 184, 0.06)' } }, y: { title: { display: true, text: yAxisLabel }, grid: { color: 'rgba(148, 163, 184, 0.06)' } } } }; } // Escape HTML to prevent XSS function escapeHtml(text) { const div = document.createElement('div'); div.textContent = text; return div.innerHTML; } // Update runs table function updateRunsTable() { const tbody = document.getElementById('runs-table-body'); tbody.innerHTML = ''; allMetricsData.slice(0, 10).forEach(run => { const models = new Set(run.results.map(r => r.model.split('/').pop())); const date = new Date(run.run_date); const row = document.createElement('tr'); // Create cells safely to prevent XSS const dateCell = document.createElement('td'); dateCell.textContent = `${date.toLocaleDateString()} ${date.toLocaleTimeString()}`; const runIdCell = document.createElement('td'); const runLink = document.createElement('a'); runLink.href = `https://github.com/${GITHUB_REPO}/actions/runs/${encodeURIComponent(run.run_id)}`; runLink.target = '_blank'; runLink.className = 'run-link'; runLink.textContent = run.run_id; runIdCell.appendChild(runLink); const commitCell = document.createElement('td'); const commitCode = document.createElement('code'); commitCode.textContent = run.commit_sha.substring(0, 7); commitCell.appendChild(commitCode); const branchCell = document.createElement('td'); branchCell.textContent = run.branch; const modelsCell = document.createElement('td'); Array.from(models).forEach((model, index) => { if (index > 0) modelsCell.appendChild(document.createTextNode(' ')); const badge = document.createElement('span'); badge.className = 'model-badge'; badge.textContent = model; modelsCell.appendChild(badge); }); row.appendChild(dateCell); row.appendChild(runIdCell); row.appendChild(commitCell); row.appendChild(branchCell); row.appendChild(modelsCell); tbody.appendChild(row); }); } // Refresh data async function refreshData() { document.getElementById('content').style.display = 'none'; document.getElementById('loading').style.display = 'flex'; await init(); } // Format numbers for display function formatNumber(num) { if (num >= 1000) { return (num / 1000).toFixed(1) + 'k'; } return num.toFixed(1); } // Authentication state let authToken = sessionStorage.getItem('dashboard_auth_token') || null; // Get auth headers for API requests function getAuthHeaders() { const headers = {}; if (authToken) { headers['Authorization'] = `Bearer ${authToken}`; } return headers; } // Check if server requires authentication and show/hide login accordingly async function checkAuthAndInit() { const loginOverlay = document.getElementById('login-overlay'); const dashboardContainer = document.getElementById('dashboard-container'); try { const response = await fetch('/api/auth-check'); if (response.ok) { const data = await response.json(); if (!data.auth_required) { // No auth required - skip login, show dashboard directly loginOverlay.style.display = 'none'; dashboardContainer.style.display = 'block'; init(); return; } } } catch (e) { // Server not available (e.g. static hosting) - skip login loginOverlay.style.display = 'none'; dashboardContainer.style.display = 'block'; init(); return; } // Auth is required - check if we have a valid token from a previous session if (authToken) { try { const testResponse = await fetch('/api/metrics', { headers: getAuthHeaders() }); if (testResponse.ok) { loginOverlay.style.display = 'none'; dashboardContainer.style.display = 'block'; init(); return; } } catch (e) { // Token invalid or expired } // Clear invalid token authToken = null; sessionStorage.removeItem('dashboard_auth_token'); } // Show login form loginOverlay.style.display = 'flex'; dashboardContainer.style.display = 'none'; } // Handle login form submission async function handleLogin(event) { event.preventDefault(); const username = document.getElementById('login-username').value; const password = document.getElementById('login-password').value; const errorEl = document.getElementById('login-error'); const loginBtn = document.getElementById('login-btn'); errorEl.textContent = ''; loginBtn.disabled = true; loginBtn.textContent = 'Signing in...'; try { const response = await fetch('/api/login', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ username, password }) }); const data = await response.json(); if (response.ok && data.token) { authToken = data.token; sessionStorage.setItem('dashboard_auth_token', authToken); document.getElementById('login-overlay').style.display = 'none'; document.getElementById('dashboard-container').style.display = 'block'; init(); } else { errorEl.textContent = data.error || 'Invalid username or password'; } } catch (e) { errorEl.textContent = 'Unable to connect to server'; } finally { loginBtn.disabled = false; loginBtn.textContent = 'Sign In'; } return false; } // Initialize on page load document.addEventListener('DOMContentLoaded', checkAuthAndInit); ================================================ FILE: docs/performance_dashboard/fetch_metrics.py ================================================ #!/usr/bin/env python3 """ Fetch and process SGLang nightly test metrics from GitHub Actions artifacts. This script fetches consolidated metrics from GitHub Actions workflow runs and outputs them as JSON for the performance dashboard. Usage: python fetch_metrics.py --output metrics_data.json python fetch_metrics.py --output metrics_data.json --days 30 python fetch_metrics.py --output metrics_data.json --run-id 21338741812 """ import argparse import io import json import os import sys import zipfile from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Optional import requests GITHUB_REPO = "sgl-project/sglang" WORKFLOW_NAME = "nightly-test-nvidia.yml" ARTIFACT_PREFIX = "consolidated-metrics-" def get_github_token() -> Optional[str]: """Get GitHub token from environment or gh CLI.""" # Check environment variable first token = os.environ.get("GITHUB_TOKEN") if token: return token # Try gh CLI try: import subprocess result = subprocess.run( ["gh", "auth", "token"], capture_output=True, text=True, check=True, ) return result.stdout.strip() except (subprocess.CalledProcessError, FileNotFoundError): pass return None def get_headers(token: Optional[str]) -> dict: """Get request headers with optional authentication.""" headers = { "Accept": "application/vnd.github.v3+json", } if token: headers["Authorization"] = f"Bearer {token}" return headers def fetch_workflow_runs( token: Optional[str], days: int = 30, event: Optional[str] = None, ) -> list: """Fetch completed workflow runs from GitHub Actions.""" url = f"https://api.github.com/repos/{GITHUB_REPO}/actions/workflows/{WORKFLOW_NAME}/runs" params = { "status": "completed", "per_page": 100, } if event: params["event"] = event response = requests.get(url, headers=get_headers(token), params=params, timeout=30) response.raise_for_status() runs = response.json().get("workflow_runs", []) # Filter by date cutoff = datetime.now(timezone.utc) - timedelta(days=days) runs = [ run for run in runs if datetime.fromisoformat(run["created_at"].replace("Z", "+00:00")) > cutoff ] return runs def fetch_run_artifacts(token: Optional[str], run_id: int) -> list: """Fetch artifacts for a specific workflow run.""" url = f"https://api.github.com/repos/{GITHUB_REPO}/actions/runs/{run_id}/artifacts" response = requests.get(url, headers=get_headers(token), timeout=30) response.raise_for_status() return response.json().get("artifacts", []) def download_artifact(token: Optional[str], artifact_id: int) -> Optional[bytes]: """Download an artifact by ID.""" if not token: print(f"Warning: GitHub token required to download artifacts", file=sys.stderr) return None url = f"https://api.github.com/repos/{GITHUB_REPO}/actions/artifacts/{artifact_id}/zip" headers = get_headers(token) response = requests.get(url, headers=headers, allow_redirects=True, timeout=60) if response.status_code == 200: return response.content print( f"Failed to download artifact {artifact_id}: {response.status_code}", file=sys.stderr, ) return None def extract_metrics_from_zip(zip_content: bytes) -> Optional[dict]: """Extract metrics JSON from a zip file.""" try: with zipfile.ZipFile(io.BytesIO(zip_content)) as zf: # Find the JSON file in the archive json_files = [f for f in zf.namelist() if f.endswith(".json")] if not json_files: return None with zf.open(json_files[0]) as f: return json.load(f) except (zipfile.BadZipFile, json.JSONDecodeError) as e: print(f"Failed to extract metrics: {e}", file=sys.stderr) return None def fetch_metrics_for_run(token: Optional[str], run: dict) -> Optional[dict]: """Fetch metrics for a single workflow run.""" run_id = run["id"] print(f"Fetching metrics for run {run_id}...", file=sys.stderr) artifacts = fetch_run_artifacts(token, run_id) # Find consolidated metrics artifact metrics_artifact = None for artifact in artifacts: if artifact["name"].startswith(ARTIFACT_PREFIX): metrics_artifact = artifact break if not metrics_artifact: print(f"No consolidated metrics found for run {run_id}", file=sys.stderr) return None # Download and extract zip_content = download_artifact(token, metrics_artifact["id"]) if not zip_content: return None metrics = extract_metrics_from_zip(zip_content) if not metrics: return None # Ensure required fields are present if "run_id" not in metrics: metrics["run_id"] = str(run_id) if "run_date" not in metrics: metrics["run_date"] = run["created_at"] if "commit_sha" not in metrics: metrics["commit_sha"] = run["head_sha"] if "branch" not in metrics: metrics["branch"] = run["head_branch"] return metrics def fetch_single_run(token: Optional[str], run_id: int) -> Optional[dict]: """Fetch metrics for a single run by ID.""" url = f"https://api.github.com/repos/{GITHUB_REPO}/actions/runs/{run_id}" response = requests.get(url, headers=get_headers(token), timeout=30) response.raise_for_status() run = response.json() return fetch_metrics_for_run(token, run) def main(): parser = argparse.ArgumentParser( description="Fetch SGLang nightly test metrics from GitHub Actions" ) parser.add_argument( "--output", "-o", type=str, default="metrics_data.json", help="Output JSON file path", ) parser.add_argument( "--days", type=int, default=30, help="Number of days to fetch (default: 30)", ) parser.add_argument( "--run-id", type=int, help="Fetch a specific run by ID", ) parser.add_argument( "--event", type=str, choices=["schedule", "workflow_dispatch", "push"], help="Filter by trigger event type", ) parser.add_argument( "--scheduled-only", action="store_true", help="Only fetch scheduled (nightly) runs", ) args = parser.parse_args() token = get_github_token() if not token: print( "Warning: No GitHub token found. Some features may be limited.", file=sys.stderr, ) print( "Set GITHUB_TOKEN env var or login with 'gh auth login'", file=sys.stderr, ) all_metrics = [] if args.run_id: # Fetch single run metrics = fetch_single_run(token, args.run_id) if metrics: all_metrics.append(metrics) else: # Fetch multiple runs event = "schedule" if args.scheduled_only else args.event runs = fetch_workflow_runs(token, days=args.days, event=event) print(f"Found {len(runs)} workflow runs", file=sys.stderr) for run in runs: metrics = fetch_metrics_for_run(token, run) if metrics: all_metrics.append(metrics) # Sort by date descending all_metrics.sort(key=lambda x: x.get("run_date", ""), reverse=True) # Write output output_path = Path(args.output) with open(output_path, "w") as f: json.dump(all_metrics, f, indent=2) print(f"Wrote {len(all_metrics)} metrics records to {output_path}", file=sys.stderr) if __name__ == "__main__": main() ================================================ FILE: docs/performance_dashboard/index.html ================================================ SGLang Performance Dashboard ================================================ FILE: docs/performance_dashboard/server.py ================================================ #!/usr/bin/env python3 """ Simple development server for the SGLang Performance Dashboard. This server: 1. Serves the static HTML/JS files 2. Provides an API endpoint to fetch metrics from GitHub 3. Caches metrics data to reduce API calls Usage: python server.py python server.py --port 8080 python server.py --host 0.0.0.0 # Allow external access python server.py --fetch-on-start python server.py --username admin --password secret # Enable authentication DASHBOARD_USERNAME=admin DASHBOARD_PASSWORD=secret python server.py # Via env vars python server.py --refresh-interval 12 # Auto-refresh data every 12 hours """ import argparse import hashlib import hmac import http.server import io import json import os import secrets import socketserver import threading import time import zipfile from datetime import datetime, timedelta, timezone from pathlib import Path from urllib.parse import urlparse import requests GITHUB_REPO = "sgl-project/sglang" WORKFLOW_NAME = "nightly-test-nvidia.yml" ARTIFACT_PREFIX = "consolidated-metrics-" # Cache for metrics data with thread-safe lock cache_lock = threading.Lock() metrics_cache = { "data": [], "last_updated": None, "updating": False, } CACHE_TTL = 300 # 5 minutes REQUEST_TIMEOUT = 30 # seconds # Authentication configuration (set via CLI flags) auth_config = { "enabled": False, "username": None, "password_hash": None, # SHA-256 hash of the password "active_tokens": {}, # token -> expiry timestamp } auth_lock = threading.Lock() AUTH_TOKEN_TTL = 3600 # 1 hour def hash_password(password): """Hash a password using SHA-256 for constant-time comparison.""" return hashlib.sha256(password.encode("utf-8")).hexdigest() def create_auth_token(): """Create a new session token.""" token = secrets.token_hex(32) with auth_lock: # Clean up expired tokens now = time.time() auth_config["active_tokens"] = { t: exp for t, exp in auth_config["active_tokens"].items() if exp > now } auth_config["active_tokens"][token] = now + AUTH_TOKEN_TTL return token def verify_auth_token(token): """Verify a session token is valid and not expired.""" if not token: return False with auth_lock: expiry = auth_config["active_tokens"].get(token) if expiry and expiry > time.time(): return True # Remove expired token auth_config["active_tokens"].pop(token, None) return False def get_github_token(): """Get GitHub token from environment or gh CLI.""" token = os.environ.get("GITHUB_TOKEN") if token: return token try: import subprocess result = subprocess.run( ["gh", "auth", "token"], capture_output=True, text=True, check=True, ) return result.stdout.strip() except (subprocess.CalledProcessError, FileNotFoundError): pass return None def fetch_metrics_from_github(days=30): """Fetch metrics from GitHub Actions artifacts.""" token = get_github_token() headers = {"Accept": "application/vnd.github.v3+json"} if token: headers["Authorization"] = f"Bearer {token}" # Get workflow runs - only scheduled (nightly) runs, not workflow_dispatch url = f"https://api.github.com/repos/{GITHUB_REPO}/actions/workflows/{WORKFLOW_NAME}/runs" params = {"status": "completed", "per_page": 50, "event": "schedule"} try: response = requests.get( url, headers=headers, params=params, timeout=REQUEST_TIMEOUT ) if not response.ok: print(f"Failed to fetch workflow runs: {response.status_code}") return [] except requests.exceptions.RequestException as e: print(f"Network error fetching workflow runs: {e}") return [] runs = response.json().get("workflow_runs", []) # Filter by date cutoff = datetime.now(timezone.utc) - timedelta(days=days) runs = [ run for run in runs if datetime.fromisoformat(run["created_at"].replace("Z", "+00:00")) > cutoff ] all_metrics = [] for run in runs[:20]: # Limit to 20 most recent run_id = run["id"] # Get artifacts artifacts_url = f"https://api.github.com/repos/{GITHUB_REPO}/actions/runs/{run_id}/artifacts" try: artifacts_resp = requests.get( artifacts_url, headers=headers, timeout=REQUEST_TIMEOUT ) if not artifacts_resp.ok: continue except requests.exceptions.RequestException as e: print(f"Network error fetching artifacts for run {run_id}: {e}") continue artifacts = artifacts_resp.json().get("artifacts", []) # Find consolidated metrics for artifact in artifacts: if artifact["name"].startswith(ARTIFACT_PREFIX): if not token: # Without token, we can't download - return metadata only all_metrics.append( { "run_id": str(run_id), "run_date": run["created_at"], "commit_sha": run["head_sha"], "branch": run["head_branch"], "results": [], } ) break # Download artifact download_url = f"https://api.github.com/repos/{GITHUB_REPO}/actions/artifacts/{artifact['id']}/zip" try: download_resp = requests.get( download_url, headers=headers, allow_redirects=True, timeout=REQUEST_TIMEOUT, ) except requests.exceptions.RequestException as e: print(f"Network error downloading artifact: {e}") break if download_resp.ok: try: with zipfile.ZipFile(io.BytesIO(download_resp.content)) as zf: json_files = [ f for f in zf.namelist() if f.endswith(".json") ] if json_files: with zf.open(json_files[0]) as f: metrics = json.load(f) # Ensure required fields metrics.setdefault("run_id", str(run_id)) metrics.setdefault("run_date", run["created_at"]) metrics.setdefault("commit_sha", run["head_sha"]) metrics.setdefault("branch", run["head_branch"]) all_metrics.append(metrics) except (zipfile.BadZipFile, json.JSONDecodeError) as e: print(f"Failed to process artifact: {e}") break return all_metrics def update_cache_async(): """Update the metrics cache in background with thread safety.""" with cache_lock: if metrics_cache["updating"]: return metrics_cache["updating"] = True try: data = fetch_metrics_from_github() with cache_lock: metrics_cache["data"] = data metrics_cache["last_updated"] = time.time() print(f"Cache updated with {len(data)} metrics records") finally: with cache_lock: metrics_cache["updating"] = False def start_periodic_refresh(interval_hours): """Start a background thread that refreshes the cache periodically.""" interval_seconds = interval_hours * 3600 def refresh_loop(): while True: time.sleep(interval_seconds) print(f"Periodic refresh triggered (every {interval_hours}h)") update_cache_async() thread = threading.Thread(target=refresh_loop, daemon=True) thread.start() print(f"Periodic refresh enabled: every {interval_hours} hours") class DashboardHandler(http.server.SimpleHTTPRequestHandler): """HTTP request handler for the dashboard.""" def __init__(self, *args, directory=None, **kwargs): super().__init__(*args, directory=directory, **kwargs) def _send_json(self, data, status=200): """Send a JSON response.""" self.send_response(status) self.send_header("Content-Type", "application/json") self.send_header("Access-Control-Allow-Origin", "*") self.end_headers() self.wfile.write(json.dumps(data).encode()) def _check_auth(self): """Check if request is authenticated. Returns True if OK, sends 401 and returns False otherwise.""" if not auth_config["enabled"]: return True auth_header = self.headers.get("Authorization", "") if auth_header.startswith("Bearer "): token = auth_header[7:] if verify_auth_token(token): return True self._send_json({"error": "Unauthorized"}, status=401) return False def do_GET(self): parsed = urlparse(self.path) # Prevent directory traversal attacks if ".." in parsed.path or parsed.path.startswith("//"): self.send_error(400, "Invalid path") return if parsed.path == "/api/auth-check": self.handle_auth_check() elif parsed.path == "/api/metrics": if self._check_auth(): self.handle_metrics_api(parsed) elif parsed.path == "/api/refresh": if self._check_auth(): self.handle_refresh_api() else: super().do_GET() def do_POST(self): parsed = urlparse(self.path) if parsed.path == "/api/login": self.handle_login() else: self.send_error(404, "Not Found") def handle_auth_check(self): """Tell the frontend whether authentication is required.""" self._send_json({"auth_required": auth_config["enabled"]}) def handle_login(self): """Validate username/password and return a session token.""" content_length = int(self.headers.get("Content-Length", 0)) if content_length == 0 or content_length > 4096: self._send_json({"error": "Invalid request"}, status=400) return try: body = json.loads(self.rfile.read(content_length)) except (json.JSONDecodeError, ValueError): self._send_json({"error": "Invalid JSON"}, status=400) return username = body.get("username", "") password = body.get("password", "") if hmac.compare_digest( username, auth_config["username"] ) and hmac.compare_digest( hash_password(password), auth_config["password_hash"] ): token = create_auth_token() self._send_json({"token": token}) else: self._send_json({"error": "Invalid username or password"}, status=401) def handle_metrics_api(self, parsed): """Handle /api/metrics endpoint.""" # Check cache with thread safety with cache_lock: cache_valid = ( metrics_cache["last_updated"] and time.time() - metrics_cache["last_updated"] < CACHE_TTL ) data = metrics_cache["data"].copy() if not cache_valid: # Trigger background update threading.Thread(target=update_cache_async, daemon=True).start() self._send_json(data) def handle_refresh_api(self): """Handle /api/refresh endpoint.""" threading.Thread(target=update_cache_async, daemon=True).start() self._send_json({"status": "refreshing"}) def log_message(self, format, *args): """Custom log format.""" print(f"[{self.log_date_time_string()}] {args[0]}") def main(): parser = argparse.ArgumentParser(description="SGLang Performance Dashboard Server") parser.add_argument("--port", type=int, default=8000, help="Port to serve on") parser.add_argument( "--host", default="127.0.0.1", help="Host to bind to (use 0.0.0.0 for external access)", ) parser.add_argument( "--fetch-on-start", action="store_true", help="Fetch metrics on startup" ) parser.add_argument( "--refresh-interval", type=float, default=12, help="Auto-refresh interval in hours (default: 12, set to 0 to disable)", ) parser.add_argument( "--username", default=os.environ.get("DASHBOARD_USERNAME"), help="Username for dashboard authentication (or set DASHBOARD_USERNAME env var)", ) parser.add_argument( "--password", default=os.environ.get("DASHBOARD_PASSWORD"), help="Password for dashboard authentication (or set DASHBOARD_PASSWORD env var)", ) args = parser.parse_args() # Configure authentication if both username and password are provided if args.username and args.password: auth_config["enabled"] = True auth_config["username"] = args.username auth_config["password_hash"] = hash_password(args.password) print(f"Authentication enabled for user: {args.username}") elif args.username or args.password: parser.error("Both --username and --password must be provided together") # Change to dashboard directory dashboard_dir = Path(__file__).parent os.chdir(dashboard_dir) if args.fetch_on_start: print("Fetching initial metrics data...") update_cache_async() if args.refresh_interval > 0: start_periodic_refresh(args.refresh_interval) handler = lambda *a, **kw: DashboardHandler(*a, directory=str(dashboard_dir), **kw) with socketserver.TCPServer((args.host, args.port), handler) as httpd: print(f"Serving dashboard at http://{args.host}:{args.port}") print("Press Ctrl+C to stop") try: httpd.serve_forever() except KeyboardInterrupt: print("\nShutting down...") if __name__ == "__main__": main() ================================================ FILE: docs/platforms/amd_gpu.md ================================================ # AMD GPUs This document describes how to run SGLang on AMD GPUs. If you encounter issues or have questions, please [open an issue](https://github.com/sgl-project/sglang/issues). ## System Configuration When using AMD GPUs (such as MI300X), certain system-level optimizations help ensure stable performance. Here we take MI300X as an example. AMD provides official documentation for MI300X optimization and system tuning: - [AMD MI300X Tuning Guides](https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/index.html) - [LLM inference performance validation on AMD Instinct MI300X](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/inference/vllm-benchmark.html) - [AMD Instinct MI300X System Optimization](https://rocm.docs.amd.com/en/latest/how-to/system-optimization/mi300x.html) - [AMD Instinct MI300X Workload Optimization](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/inference-optimization/workload.html) - [Supercharge DeepSeek-R1 Inference on AMD Instinct MI300X](https://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1-Part2/README.html) **NOTE:** We strongly recommend reading these docs and guides entirely to fully utilize your system. Below are a few key settings to confirm or enable for SGLang: ### Update GRUB Settings In `/etc/default/grub`, append the following to `GRUB_CMDLINE_LINUX`: ```text pci=realloc=off iommu=pt ``` Afterward, run `sudo update-grub` (or your distro’s equivalent) and reboot. ### Disable NUMA Auto-Balancing ```bash sudo sh -c 'echo 0 > /proc/sys/kernel/numa_balancing' ``` You can automate or verify this change using [this helpful script](https://github.com/ROCm/triton/blob/rocm_env/scripts/amd/env_check.sh). Again, please go through the entire documentation to confirm your system is using the recommended configuration. ## Install SGLang You can install SGLang using one of the methods below. ### Install from Source ```bash # Use the last release branch git clone -b v0.5.9 https://github.com/sgl-project/sglang.git cd sglang # Compile sgl-kernel pip install --upgrade pip cd sgl-kernel python setup_rocm.py install # Install sglang python package along with diffusion support cd .. rm -rf python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml pip install -e "python[all_hip]" ``` ### Install Using Docker (Recommended) The docker images are available on Docker Hub at [lmsysorg/sglang](https://hub.docker.com/r/lmsysorg/sglang/tags), built from [rocm.Dockerfile](https://github.com/sgl-project/sglang/tree/main/docker). The steps below show how to build and use an image. 1. Build the docker image. If you use pre-built images, you can skip this step and replace `sglang_image` with the pre-built image names in the steps below. ```bash docker build -t sglang_image -f rocm.Dockerfile . ``` 2. Create a convenient alias. ```bash alias drun='docker run -it --rm --network=host --privileged --device=/dev/kfd --device=/dev/dri \ --ipc=host --shm-size 16G --group-add video --cap-add=SYS_PTRACE \ --security-opt seccomp=unconfined \ -v $HOME/dockerx:/dockerx \ -v /data:/data' ``` If you are using RDMA, please note that: - `--network host` and `--privileged` are required by RDMA. If you don't need RDMA, you can remove them. - You may need to set `NCCL_IB_GID_INDEX` if you are using RoCE, for example: `export NCCL_IB_GID_INDEX=3`. 3. Launch the server. **NOTE:** Replace `` below with your [huggingface hub token](https://huggingface.co/docs/hub/en/security-tokens). ```bash drun -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=" \ sglang_image \ python3 -m sglang.launch_server \ --model-path NousResearch/Meta-Llama-3.1-8B \ --host 0.0.0.0 \ --port 30000 ``` 4. To verify the utility, you can run a benchmark in another terminal or refer to [other docs](https://docs.sglang.io/basic_usage/openai_api_completions.html) to send requests to the engine. ```bash drun sglang_image \ python3 -m sglang.bench_serving \ --backend sglang \ --dataset-name random \ --num-prompts 4000 \ --random-input 128 \ --random-output 128 ``` With your AMD system properly configured and SGLang installed, you can now fully leverage AMD hardware to power SGLang’s machine learning capabilities. ## Quantization on AMD GPUs The [Quantization documentation](../advanced_features/quantization.md#platform-compatibility) has a full compatibility matrix. The short version: FP8, AWQ, MXFP4, W8A8, GPTQ, compressed-tensors, Quark, and **petit_nvfp4** (NVFP4 on ROCm via [Petit](https://github.com/causalflow-ai/petit-kernel)) all work on AMD. Methods that depend on Marlin or NVIDIA-specific kernels (`awq_marlin`, `gptq_marlin`, `gguf`, `modelopt_fp8`, `modelopt_fp4`) do not. A few things to keep in mind: - FP8 works via Aiter or Triton. Pre-quantized FP8 models like DeepSeek-V3/R1 work out of the box. - AWQ uses Triton dequantization kernels on AMD. The faster Marlin path is not available. - MXFP4 requires CDNA3/CDNA4 and `SGLANG_USE_AITER=1`. - `petit_nvfp4` enables NVFP4 models (e.g., [Llama 3.3 70B FP4](https://huggingface.co/nvidia/Llama-3.3-70B-Instruct-FP4)) on MI250/MI300X via [Petit](https://github.com/causalflow-ai/petit-kernel). Install with `pip install petit-kernel`; no `--quantization` flag needed when loading pre-quantized NVFP4 models. - `quark_int4fp8_moe` is an AMD-only online quantization method for MoE models on CDNA3/CDNA4. Several of these backends are accelerated by [Aiter](https://github.com/ROCm/aiter). Enable it with: ```bash export SGLANG_USE_AITER=1 ``` Example -- serving an AWQ model: ```bash python3 -m sglang.launch_server \ --model-path hugging-quants/Mixtral-8x7B-Instruct-v0.1-AWQ-INT4 \ --trust-remote-code \ --port 30000 --host 0.0.0.0 ``` Example -- FP8 online quantization: ```bash python3 -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ --quantization fp8 \ --port 30000 --host 0.0.0.0 ``` ## Examples ### Running DeepSeek-V3 The only difference when running DeepSeek-V3 is in how you start the server. Here's an example command: ```bash drun -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --ipc=host \ --env "HF_TOKEN=" \ sglang_image \ python3 -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-V3 \ # <- here --tp 8 \ --trust-remote-code \ --host 0.0.0.0 \ --port 30000 ``` [Running DeepSeek-R1 on a single NDv5 MI300X VM](https://techcommunity.microsoft.com/blog/azurehighperformancecomputingblog/running-deepseek-r1-on-a-single-ndv5-mi300x-vm/4372726) could also be a good reference. ### Running Llama3.1 Running Llama3.1 is nearly identical to running DeepSeek-V3. The only difference is in the model specified when starting the server, shown by the following example command: ```bash drun -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --ipc=host \ --env "HF_TOKEN=" \ sglang_image \ python3 -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ # <- here --tp 8 \ --trust-remote-code \ --host 0.0.0.0 \ --port 30000 ``` ### Warmup Step When the server displays `The server is fired up and ready to roll!`, it means the startup is successful. ================================================ FILE: docs/platforms/apple_metal.md ================================================ # Apple Silicon with Metal This document describes how run SGLang on Apple Silicon using [Metal](https://developer.apple.com/metal/). If you encounter issues or have questions, please [open an issue](https://github.com/sgl-project/sglang/issues). ## Install SGLang You can install SGLang using one of the methods below. ### Install from Source ```bash # Use the default branch git clone https://github.com/sgl-project/sglang.git cd sglang # Install sglang python package pip install --upgrade pip rm -f python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml uv pip install -e "python[all_mps]" ``` ================================================ FILE: docs/platforms/ascend_contribution_guide.md ================================================ # Contribution Guide Welcome to **SGLang**! We appreciate your interest in contributing. This guide provides a concise overview of how to set up your environment, run tests, build documentation, and open a Pull Request (PR). Whether you’re fixing a small bug or developing a major feature, we encourage following these steps for a smooth contribution process. ## Install SGLang from Source ### Prepare Environment Before contributing, please ensure that your environment is set up correctly. Follow the steps in the [Installation Guide](../platforms/ascend_npu.md) to install the necessary dependencies. We recommend [using docker](../platforms/ascend_npu.md#method-2-using-docker-image) to build the environment. ### Fork and clone the repository **Note**: New contributors do **not** have the write permission to push to the official SGLang repo. Please fork the repository under your GitHub account, then clone your fork locally. ```bash git clone https://github.com//sglang.git # if you are using docker, the environment is already set up. cd sglang export PYTHONPATH=$PWD/python:$PYTHONPATH ``` ## Format code with pre-commit We use [pre-commit](https://pre-commit.com/) to maintain consistent code style checks. Before pushing your changes, please run: ```bash pip3 install pre-commit pre-commit install pre-commit run --all-files ``` - **`pre-commit run --all-files`** manually runs all configured checks, applying fixes if possible. If it fails the first time, re-run it to ensure lint errors are fully resolved. Make sure your code passes all checks **before** creating a Pull Request. - **Do not commit** directly to the `main` branch. Always create a new branch (e.g., `feature/my-new-feature`), push your changes, and open a PR from that branch. ## Run and add unit tests If you add a new feature or fix a bug, please add corresponding unit tests to ensure coverage and prevent regression. SGLang uses Python's built-in [unittest](https://docs.python.org/3/library/unittest.html) framework. For detailed instructions on running tests and integrating them into CI, refer to [test/README.md](https://github.com/sgl-project/sglang/tree/main/test/README.md). ## Write documentations We recommend new contributors start from writing documentation, which helps you quickly understand SGLang codebase. For more details, please refer to [docs/README.md](https://github.com/sgl-project/sglang/tree/main/docs/README.md). ## Test the accuracy If your code changes the model output, please run the accuracy tests. A quick sanity check is the few-shot GSM8K. ``` # Launch a server python3 -m sglang.launch_server --model Qwen/Qwen2-7B-Instruct # Evaluate python3 -m sglang.test.few_shot_gsm8k --num-questions 200 ``` Please note that the above script is primarily a sanity check, not a rigorous accuracy or speed test. This test can have significant variance (1%–5%) in accuracy due to batching and the non-deterministic nature of the inference engine. Also, do not rely on the "Latency/Output throughput" from this script, as it is not a proper speed test. GSM8K is too easy for state-of-the-art models nowadays. Please try your own more challenging accuracy tests. You can find additional accuracy eval examples in: - [test_eval_accuracy_large.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_eval_accuracy_large.py) - [test_moe_eval_accuracy_large.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_moe_eval_accuracy_large.py) ## Benchmark the speed Refer to [Benchmark and Profiling](../developer_guide/benchmark_and_profiling.md). ## Requesting a review for merge You can follow the pull request merge process described in [MAINTAINER.md](https://github.com/sgl-project/sglang/blob/main/.github/MAINTAINER.md). You will need to work with the Merge Oncall, Codeowner, and other reviewers to get their approvals. Then your PR can be merged. ## How to Trigger CI Tests We have a lot of open PRs but limited CI machines, so only top and trusted contributors have permission to trigger CI tests. Users with permission are listed in the [CI_PERMISSIONS.json](https://github.com/sgl-project/sglang/blob/main/.github/CI_PERMISSIONS.json) For CI to run on a pull request, it must have the "run-ci" label. Authorized users can add the label or rerun failed tests by commenting on the PR with one of these commands: - `/tag-run-ci-label`: Adds the "run-ci" label. Every future commit will trigger CI. - `/rerun-failed-ci`: Reruns the failed or flaky tests from the most recent commit. - `/tag-and-rerun-ci`: A single command that performs both `/tag-run-ci-label` and `/rerun-failed-ci`. - `/rerun-stage `: Reruns a specific test stage without waiting for its dependencies. This is useful when you want to quickly validate a fix for a specific test failure instead of waiting ~30 minutes for preceding stages to complete. If you have permission, the [Slash Command Handler](https://github.com/sgl-project/sglang/actions/workflows/slash-command-handler.yml) will run your command and react with a 👍 to your comment. It may take up to a few minutes for the reaction to appear. Here’s a usage [example](https://github.com/sgl-project/sglang/pull/14253#issuecomment-3599509302). To avoid spamming a PR with too many `/rerun-failed-ci` comments, you can also trigger the command by editing an existing comment and adding any suffix (e.g., `/rerun-failed-ci try again`). Example of rerunning a single test stage: `/rerun-stage unit-test-backend-4-gpu`. If you don’t have permission, please ask maintainers to trigger CI for you. ### CI rate limits Due to CI scheduling and limited resources, higher-priority PRs may preempt running jobs. In such cases, you may need to rerun the tests. We apply CI rate limits to prevent abuse and ensure fair usage of our CI resources. Each CI workflow has a default limit defined in its workflow configuration file. For example, in [pr-gate.yml](https://github.com/sgl-project/sglang/blob/main/.github/workflows/pr-gate.yml), the default cooldown period is 120 minutes, and each workflow can override it via the `cool-down-minutes` input parameter: ```yaml cool-down-minutes: description: "Default cooldown period in minutes; 0 disables rate limiting" type: number default: 120 ``` Users listed in [CI_PERMISSIONS.json](https://github.com/sgl-project/sglang/blob/main/.github/CI_PERMISSIONS.json) may have a per-user cooldown interval. In practice, we use the minimum of the workflow’s default window and the user-specific interval. ## Code style guidance - Avoid code duplication. If the same code snippet (more than five lines) appears multiple times, extract it into a shared function. - Minimize device synchronization. Reduce expensive CPU-GPU synchronization operations, such as `tensor.item()` or `tensor.cpu()`, whenever possible. Use vectorized code. - Prioritize extreme efficiency. SGLang is a runtime, and most of your code runs on the critical path for every request. Optimize all minor overheads as much as possible, especially in the model forward code. - A common pattern is some runtime checks in the model forward pass (e.g., [this](https://github.com/sgl-project/sglang/blob/f1b0eda55c2c4838e8ab90a0fac7fb1e3d7064ab/python/sglang/srt/models/deepseek_v2.py#L486-L491)). These are very likely the same for every layer. Please cache the result as a single boolean value whenever possible. - Make functions as pure as possible. Avoid in-place modification of arguments. - Keep files concise. If a file exceeds 2,000 lines of code, split it into multiple smaller files. (e.g., `scheduler.py`, `scheduler_output_processor_mixin.py`) - Keep tests run fast. - If a single test file run longer than 500 seconds, split it into multiple smaller files (e.g., `test_eagle_infer_a.py`, `test_eagle_infer_b.py`). - If a single job in a github workflow runs longer than 30 mins, split it into smaller jobs/steps. - Reuse server launches in your unit tests to make tests run faster. - When supporting new hardware or features, follow these guidelines: - Do not drastically change existing code. - Always prefer new files to introduce specific components for your new hardware (e.g., `allocator_ascend.py`). - If you write multiple if/else blocks for new features, ensure the common path (e.g., NVIDIA hardware or the existing code path) is the first branch. ## How to update sgl-kernel Since sglang and sgl-kernel are separate Python packages, our current GitHub CI infrastructure does not support updating a kernel and using it immediately within the same pull request (PR). To add a new kernel or modify an existing one in the `sgl-kernel/` source tree, you must use multiple PRs. Follow these steps: 1. Submit a PR to update the sgl-kernel source code without using it in sglang python package (e.g., [#8884](https://github.com/sgl-project/sglang/pull/8884/files)). 2. Bump the version of the kernel package (e.g., [#9220](https://github.com/sgl-project/sglang/pull/9220/files)). - Once merged, this will trigger an automatic release of the `sglang-kernel` wheel to PyPI. - If not urgent, you can wait for other people to release the wheel. A new version will typically be released within one week. 3. Apply the changes: - Update the `sglang-kernel` version in `sglang/python/pyproject.toml` to use the modified kernels. - Update the related caller code in the sglang to use the new kernel. ## How to update sgl-kernel-npu Sgl-kernel-npu is the kernel package for Ascend NPU and is maintained in the [sgl-kernel-npu](https://github.com/sgl-project/sgl-kernel-npu) repository. if you want to add a new kernel and want to use it in sglang, please follow the steps in [Contribution Guide](https://github.com/sgl-project/sgl-kernel-npu/blob/main/docs/developer_guide/contribution_guide.md). ## Tips for newcomers If you want to contribute but don’t have a specific idea in mind, pick issues labeled [“good first issue” or “help wanted”](https://github.com/sgl-project/sglang/issues?q=is%3Aissue+label%3A%22good+first+issue%22%2C%22help+wanted%22). These tasks typically have lower complexity and provide an excellent introduction to the codebase. Also check out this [code walk-through](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/tree/main/sglang/code-walk-through) for a deeper look into SGLang’s workflow. If you have any questions or want to start a discussion, please feel free to ask in our [Slack channel](https://slack.sglang.io). Thank you for your interest in SGLang. Happy coding! ================================================ FILE: docs/platforms/ascend_npu.md ================================================ # SGLang installation with NPUs support You can install SGLang using any of the methods below. Please go through `System Settings` section to ensure the clusters are roaring at max performance. Feel free to leave an issue [here at sglang](https://github.com/sgl-project/sglang/issues) if you encounter any issues or have any problems. ## Component Version Mapping For SGLang | Component | Version | Obtain Way | |-------------------|-------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | HDK | 25.3.RC1 | [link](https://hiascend.com/hardware/firmware-drivers/commercial?product=7&model=33) | | CANN | 8.5.0 | [Obtain Images](#obtain-cann-image) | | Pytorch Adapter | 7.3.0 | [link](https://gitcode.com/Ascend/pytorch/releases) | | MemFabric | 1.0.5 | `pip install memfabric-hybrid==1.0.5` | | Triton | 3.2.0 | `pip install triton-ascend`| | SGLang NPU Kernel | NA | [link](https://github.com/sgl-project/sgl-kernel-npu/releases) | ### Obtain CANN Image You can obtain the dependency of a specified version of CANN through an image. ```shell # for Atlas 800I A3 and Ubuntu OS docker pull quay.io/ascend/cann:8.5.0-a3-ubuntu22.04-py3.11 # for Atlas 800I A2 and Ubuntu OS docker pull quay.io/ascend/cann:8.5.0-910b-ubuntu22.04-py3.11 ``` ## Preparing the Running Environment ### Method 1: Installing from source with prerequisites #### Python Version Only `python==3.11` is supported currently. If you don't want to break system pre-installed python, try installing with [conda](https://github.com/conda/conda). ```shell conda create --name sglang_npu python=3.11 conda activate sglang_npu ``` #### CANN Prior to start work with SGLang on Ascend you need to install CANN Toolkit, Kernels operator package and NNAL version 8.3.RC2 or higher, check the [installation guide](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/83RC1/softwareinst/instg/instg_0008.html?Mode=PmIns&InstallType=local&OS=openEuler&Software=cannToolKit) #### MemFabric-Hybrid If you want to use PD disaggregation mode, you need to install MemFabric-Hybrid. MemFabric-Hybrid is a drop-in replacement of Mooncake Transfer Engine that enables KV cache transfer on Ascend NPU clusters. ```shell pip install memfabric-hybrid==1.0.5 ``` #### Pytorch and Pytorch Framework Adaptor on Ascend ```shell PYTORCH_VERSION=2.8.0 TORCHVISION_VERSION=0.23.0 TORCH_NPU_VERSION=2.8.0 pip install torch==$PYTORCH_VERSION torchvision==$TORCHVISION_VERSION --index-url https://download.pytorch.org/whl/cpu pip install torch_npu==$TORCH_NPU_VERSION ``` If you are using other versions of `torch` and install `torch_npu`, check [installation guide](https://github.com/Ascend/pytorch/blob/master/README.md) #### Triton on Ascend We provide our own implementation of Triton for Ascend. ```shell pip install triton-ascend ``` For installation of Triton on Ascend nightly builds or from sources, follow [installation guide](https://gitcode.com/Ascend/triton-ascend/blob/master/docs/sources/getting-started/installation.md) #### SGLang Kernels NPU We provide SGL kernels for Ascend NPU, check [installation guide](https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/sgl_kernel_npu/README.md). #### DeepEP-compatible Library We provide a DeepEP-compatible Library as a drop-in replacement of deepseek-ai's DeepEP library, check the [installation guide](https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/deep_ep/README.md). #### Installing SGLang from source ```shell # Use the last release branch git clone https://github.com/sgl-project/sglang.git cd sglang mv python/pyproject_npu.toml python/pyproject.toml pip install -e python[all_npu] ``` ### Method 2: Using Docker Image #### Obtain Image You can download the SGLang image or build an image based on Dockerfile to obtain the Ascend NPU image. 1. Download SGLang image ```angular2html dockerhub: docker.io/lmsysorg/sglang:$tag # Main-based tag, change main to specific version like v0.5.6, # you can get image for specific version Atlas 800I A3 : {main}-cann8.5.0-a3 Atlas 800I A2: {main}-cann8.5.0-910b ``` 2. Build an image based on Dockerfile ```shell # Clone the SGLang repository git clone https://github.com/sgl-project/sglang.git cd sglang/docker # Build the docker image # If there are network errors, please modify the Dockerfile to use offline dependencies or use a proxy docker build -t -f npu.Dockerfile . ``` #### Create Docker __Notice:__ `--privileged` and `--network=host` are required by RDMA, which is typically needed by Ascend NPU clusters. __Notice:__ The following docker command is based on Atlas 800I A3 machines. If you are using Atlas 800I A2, make sure only `davinci[0-7]` are mapped into container. ```shell alias drun='docker run -it --rm --privileged --network=host --ipc=host --shm-size=16g \ --device=/dev/davinci0 --device=/dev/davinci1 --device=/dev/davinci2 --device=/dev/davinci3 \ --device=/dev/davinci4 --device=/dev/davinci5 --device=/dev/davinci6 --device=/dev/davinci7 \ --device=/dev/davinci8 --device=/dev/davinci9 --device=/dev/davinci10 --device=/dev/davinci11 \ --device=/dev/davinci12 --device=/dev/davinci13 --device=/dev/davinci14 --device=/dev/davinci15 \ --device=/dev/davinci_manager --device=/dev/hisi_hdc \ --volume /usr/local/sbin:/usr/local/sbin --volume /usr/local/Ascend/driver:/usr/local/Ascend/driver \ --volume /usr/local/Ascend/firmware:/usr/local/Ascend/firmware \ --volume /etc/ascend_install.info:/etc/ascend_install.info \ --volume /var/queue_schedule:/var/queue_schedule --volume ~/.cache/:/root/.cache/' # Add HF_TOKEN env for download model by SGLang. drun --env "HF_TOKEN=" \ \ python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --attention-backend ascend ``` ## System Settings ### CPU performance power scheme The default power scheme on Ascend hardware is `ondemand` which could affect performance, changing it to `performance` is recommended. ```shell echo performance | sudo tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor # Make sure changes are applied successfully cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor # shows performance ``` ### Disable NUMA balancing ```shell sudo sysctl -w kernel.numa_balancing=0 # Check cat /proc/sys/kernel/numa_balancing # shows 0 ``` ### Prevent swapping out system memory ```shell sudo sysctl -w vm.swappiness=10 # Check cat /proc/sys/vm/swappiness # shows 10 ``` ## Running SGLang Service ### Running Service For Large Language Models #### PD Mixed Scene ```shell # Enabling CPU Affinity export SGLANG_SET_CPU_AFFINITY=1 python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --attention-backend ascend ``` #### PD Separation Scene 1. Launch Prefill Server ```shell # Enabling CPU Affinity export SGLANG_SET_CPU_AFFINITY=1 # PIP: recommended to config first Prefill Server IP # PORT: one free port # all sglang servers need to be config the same PIP and PORT, export ASCEND_MF_STORE_URL="tcp://PIP:PORT" # if you are Atlas 800I A2 hardware and use rdma for kv cache transfer, add this parameter export ASCEND_MF_TRANSFER_PROTOCOL="device_rdma" python3 -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --disaggregation-mode prefill \ --disaggregation-transfer-backend ascend \ --disaggregation-bootstrap-port 8995 \ --attention-backend ascend \ --device npu \ --base-gpu-id 0 \ --tp-size 1 \ --host 127.0.0.1 \ --port 8000 ``` 2. Launch Decode Server ```shell # PIP: recommended to config first Prefill Server IP # PORT: one free port # all sglang servers need to be config the same PIP and PORT, export ASCEND_MF_STORE_URL="tcp://PIP:PORT" # if you are Atlas 800I A2 hardware and use rdma for kv cache transfer, add this parameter export ASCEND_MF_TRANSFER_PROTOCOL="device_rdma" python3 -m sglang.launch_server \ --model-path meta-llama/Llama-3.1-8B-Instruct \ --disaggregation-mode decode \ --disaggregation-transfer-backend ascend \ --attention-backend ascend \ --device npu \ --base-gpu-id 1 \ --tp-size 1 \ --host 127.0.0.1 \ --port 8001 ``` 3. Launch Router ```shell python3 -m sglang_router.launch_router \ --pd-disaggregation \ --policy cache_aware \ --prefill http://127.0.0.1:8000 8995 \ --decode http://127.0.0.1:8001 \ --host 127.0.0.1 \ --port 6688 ``` ### Running Service For Multimodal Language Models #### PD Mixed Scene ```shell python3 -m sglang.launch_server \ --model-path Qwen3-VL-30B-A3B-Instruct \ --host 127.0.0.1 \ --port 8000 \ --tp 4 \ --device npu \ --attention-backend ascend \ --mm-attention-backend ascend_attn \ --disable-radix-cache \ --trust-remote-code \ --enable-multimodal \ --sampling-backend ascend ``` ================================================ FILE: docs/platforms/ascend_npu_best_practice.md ================================================ # Best Practice on Ascend NPU This section describes the best practice data of mainstream LLM models such as DeepSeek and Qwen on the Ascend NPU. If you encounter issues or have any questions, please [open an issue](https://github.com/sgl-project/sglang/issues). ## DeepSeek Series Models ### Low Latency | Model | Hardware | Cards | Deploy Mode | Dataset | TPOT | Quantization | Configuration | |-------------------|---------------|-------|---------------|-----------|------|--------------|---------------------------------------------------------------------------------------| | Deepseek-R1 | Atlas 800I A3 | 32 | PD Separation | 6K+1.6K | 20ms | W8A8 INT8 | [Optimal Configuration](#deepseek-r1-6k-1_6k-20ms-on-a3-32-cards-separation-mode) | | Deepseek-R1 | Atlas 800I A3 | 32 | PD Separation | 3.9K+1K | 20ms | W8A8 INT8 | [Optimal Configuration](#deepseek-r1-3_9k-1k-20ms-on-a3-32-cards-separation-mode) | | Deepseek-R1 | Atlas 800I A3 | 32 | PD Separation | 3.5K+1.5K | 20ms | W8A8 INT8 | [Optimal Configuration](#deepseek-r1-3_5k-1_5k-20ms-on-a3-32-cards-separation-mode) | | Deepseek-R1 | Atlas 800I A3 | 32 | PD Separation | 3.5K+1K | 20ms | W8A8 INT8 | [Optimal Configuration](#deepseek-r1-3_5k-1k-20ms-on-a3-32-cards-separation-mode) | | DeepSeek-V3.2-Exp | Atlas 800I A3 | 32 | PD Separation | 64K+3K | 30ms | W8A8 INT8 | [Optimal Configuration](#deepseek-v32-exp-64k-3k-30ms-on-a3-32-cards-separation-mode) | ### High Throughput | Model | Hardware | Cards | Deploy Mode | Dataset | TPOT | Quantization | Configuration | |-------------|---------------|-------|---------------|-----------|------|--------------|-------------------------------------------------------------------------------------| | Deepseek-R1 | Atlas 800I A3 | 32 | PD Separation | 3.5K+1.5K | 50ms | W8A8 INT8 | [Optimal Configuration](#deepseek-r1-3_5k-1_5k-50ms-on-a3-32-cards-separation-mode) | | Deepseek-R1 | Atlas 800I A3 | 8 | PD Mixed | 2K+2K | 50ms | W4A8 INT8 | [Optimal Configuration](#deepseek-r1-2k-2k-50ms-on-a3-8-cards-mixed-mode) | | Deepseek-R1 | Atlas 800I A3 | 16 | PD Separation | 2K+2K | 50ms | W4A8 INT8 | [Optimal Configuration](#deepseek-r1-2k-2k-50ms-on-a3-16-cards-separation-mode) | | Deepseek-R1 | Atlas 800I A3 | 8 | PD Mixed | 3.5K+1.5K | 50ms | W4A8 INT8 | [Optimal Configuration](#deepseek-r1-3_5k-1_5k-50ms-on-a3-8-cards-mixed-mode) | | Deepseek-R1 | Atlas 800I A3 | 16 | PD Separation | 3.5K+1.5K | 50ms | W4A8 INT8 | [Optimal Configuration](#deepseek-r1-3_5k-1_5k-50ms-on-a3-16-cards-separation-mode) | ## Qwen Series Models ### Low Latency | Model | Hardware | Cards | Deploy Mode | Dataset | TPOT | Quantization | Configuration | |-----------------|---------------|-------|-------------|---------|------|--------------|--------------------------------------------------------------------------------| | Qwen3-235B-A22B | Atlas 800I A3 | 8 | PD Mixed | 11K+1K | 10ms | BF16 | [Optimal Configuration](#qwen3-235b-a22b-11k-1k-10ms-on-a3-8-cards-mixed-mode) | | Qwen3-32B | Atlas 800I A3 | 4 | PD Mixed | 6K+1.5K | 18ms | BF16 | [Optimal Configuration](#qwen3-32b-6k-1_5k-18ms-on-a3-4-cards-mixed-mode) | | Qwen3-32B | Atlas 800I A3 | 4 | PD Mixed | 4K+1.5K | 11ms | BF16 | [Optimal Configuration](#qwen3-32b-4k-1_5k-11ms-on-a3-4-cards-mixed-mode) | | Qwen3-32B | Atlas 800I A3 | 8 | PD Mixed | 18K+4K | 12ms | BF16 | [Optimal Configuration](#qwen3-32b-18k-4k-12ms-on-a3-8-cards-mixed-mode) | | Qwen3-32B | Atlas 800I A2 | 8 | PD Mixed | 6K+1.5K | 18ms | W8A8 INT8 | [Optimal Configuration](#qwen3-32b-6k-1_5k-18ms-on-a2-8-cards-mixed-mode) | | Qwen3-32B | Atlas 800I A2 | 8 | PD Mixed | 4K+1.5K | 11ms | BF16 | [Optimal Configuration](#qwen3-32b-4k-1_5k-11ms-on-a2-8-cards-mixed-mode) | ### High Throughput | Model | Hardware | Cards | Deploy Mode | Dataset | TPOT | Quantization | Configuration | |--------------------------------|---------------|-------|---------------|-----------|-------|--------------|--------------------------------------------------------------------------------------------------------| | Qwen3-235B-A22B | Atlas 800I A3 | 24 | PD Separation | 3.5K+1.5K | 50ms | W8A8 INT8 | [Optimal Configuration](#qwen3-235b-a22b-3_5k-1_5k-50ms-on-a3-24-cards-separation-mode) | | Qwen3-235B-A22B | Atlas 800I A3 | 8 | PD Mixed | 3.5K+1.5K | 50ms | W8A8 INT8 | [Optimal Configuration](#qwen3-235b-a22b-3_5k-1_5k-50ms-on-a3-8-cards-mixed-mode) | | Qwen3-235B-A22B | Atlas 800I A3 | 8 | PD Mixed | 2K+2K | 100ms | W8A8 INT8 | [Optimal Configuration](#qwen3-235b-a22b-2k-2k-100ms-on-a3-8-cards-mixed-mode) | | Qwen3-235B-A22B | Atlas 800I A3 | 8 | PD Mixed | 2K+2K | 50ms | W8A8 INT8 | [Optimal Configuration](#qwen3-235b-a22b-2k-2k-50ms-on-a3-8-cards-mixed-mode) | | Qwen3-235B-A22B | Atlas 800I A3 | 16 | PD Mixed | 2K+2K | 50ms | W8A8 INT8 | [Optimal Configuration](#qwen3-235b-a22b-2k-2k-50ms-on-a3-16-cards-mixed-mode) | | Qwen3-32B | Atlas 800I A3 | 2 | PD Mixed | 3.5K+1.5K | 50ms | W8A8 INT8 | [Optimal Configuration](#qwen3-32b-3_5k-1_5k-50ms-on-a3-2-cards-mixed-mode) | | Qwen3-32B | Atlas 800I A3 | 2 | PD Mixed | 2K+2K | 50ms | W8A8 INT8 | [Optimal Configuration](#qwen3-32b-2k-2k-50ms-on-a3-2-cards-mixed-mode) | | Qwen3-30B-A3B | Atlas 800I A3 | 1 | PD Mixed | 3.5K+1.5K | 50ms | W8A8 INT8 | [Optimal Configuration](#qwen3-30b-a3b-3_5k-1_5k-50ms-on-a3-1-card-mixed-mode) | | Qwen3-Coder-480B-A35B-Instruct | Atlas 800I A3 | 24 | PD Separation | 3.5K+1.5K | 50ms | W8A8 INT8 | [Optimal Configuration](#qwen3-coder-480b-a35b-instruct-3_5k-1_5k-50ms-on-a3-24-cards-separation-mode) | | Qwen3-Coder-480B-A35B-Instruct | Atlas 800I A3 | 16 | PD Mixed | 3.5K+1.5K | 50ms | W8A8 INT8 | [Optimal Configuration](#qwen3-coder-480b-a35b-instruct-3_5k-1_5k-50ms-on-a3-16-cards-mixed-mode) | | Qwen3-Coder-480B-A35B-Instruct | Atlas 800I A3 | 8 | PD Mixed | 3.5K+1.5K | 50ms | W8A8 INT8 | [Optimal Configuration](#qwen3-coder-480b-a35b-instruct-3_5k-1_5k-50ms-on-a3-8-cards-mixed-mode) | | Qwen3-Next-80B-A3B-Instruct | Atlas 800I A3 | 2 | PD Mixed | 3.5K+1.5K | 50ms | W8A8 INT8 | [Optimal Configuration](#qwen3-next-80B-a3b-instruct-3_5k-1_5k-50ms-on-a3-2-cards-mixed-mode) | | Qwen3-32B | Atlas 800I A2 | 8 | PD Mixed | 3.5K+1.5K | 50ms | W8A8 INT8 | [Optimal Configuration](#qwen3-32b-3_5k-1_5k-50ms-on-a2-8-cards-mixed-mode) | | Qwen3-32B | Atlas 800I A2 | 8 | PD Mixed | 2K+2K | 50ms | W8A8 INT8 | [Optimal Configuration](#qwen3-32b-2k-2k-50ms-on-a2-8-cards-mixed-mode) | ## Optimal Configuration ### DeepSeek-R1 3_5K-1_5K 50ms on A3 32 Cards Separation Mode Model: Deepseek R1 Hardware: Atlas 800I A3 32Card DeployMode: PD Separation Dataset: random Input Output Length: 3.5K+1.5K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export ASCEND_MF_STORE_URL="tcp://your prefill ip1:24669" P_IP=('your prefill ip1' 'your prefill ip2') D_IP=('your decode ip1' 'your decode ip2') MODEL_PATH=xxx export SGLANG_NPU_USE_MLAPO=1 export SGLANG_USE_FIA_NZ=1 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" # prefill for i in "${!P_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${P_IP[$i]}" || "$LOCAL_HOST2" == "${P_IP[$i]}" ]]; then echo "${P_IP[$i]}" export HCCL_BUFFSIZE=1536 export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 export TASK_QUEUE_ENABLE=2 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode prefill --host ${P_IP[$i]} \ --port 8000 --disaggregation-bootstrap-port $((8998+$i)) --trust-remote-code --nnodes 1 --node-rank 0 \ --tp-size 16 --mem-fraction-static 0.81 --attention-backend ascend --device npu --quantization modelslim \ --disaggregation-transfer-backend ascend --max-running-requests 8 --context-length 8192 --disable-radix-cache \ --chunked-prefill-size -1 --max-prefill-tokens 28680 --moe-a2a-backend deepep --deepep-mode normal \ --speculative-algorithm NEXTN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \ --dp-size 2 --enable-dp-attention --disable-shared-experts-fusion --dtype bfloat16 --enable-attn-tp-input-scattered NODE_RANK=$i break fi done # decode for i in "${!D_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${D_IP[$i]}" || "$LOCAL_HOST2" == "${D_IP[$i]}" ]]; then echo "${D_IP[$i]}" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 export HCCL_BUFFSIZE=650 export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=78 export TASK_QUEUE_ENABLE=1 export SGLANG_SCHEDULER_SKIP_ALL_GATHER=1 export HCCL_SOCKET_IFNAME=xxx export GLOO_SOCKET_IFNAME=xxx python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode decode --host ${D_IP[$i]} \ --port 8001 --trust-remote-code --dist-init-addr ${D_IP[0]}:5000 --nnodes 2 --node-rank $i --tp-size 32 --dp-size 32 \ --mem-fraction-static 0.815 --max-running-requests 832 --attention-backend ascend --device npu --quantization modelslim \ --moe-a2a-backend deepep --enable-dp-attention --deepep-mode low_latency --enable-dp-lm-head --moe-dense-tp 1 \ --cuda-graph-bs 12 14 16 18 20 22 24 26 --disaggregation-transfer-backend ascend --watchdog-timeout 9000 --context-length 8192 \ --speculative-algorithm NEXTN --speculative-num-steps 2 --speculative-eagle-topk 1 --speculative-num-draft-tokens 3 \ --tokenizer-worker-num 4 --prefill-round-robin-balance --disable-shared-experts-fusion --dtype bfloat16 \ --load-balance-method decode_round_robin NODE_RANK=$i break fi done ``` ```shell export SGLANG_DP_ROUND_ROBIN=1 python -m sglang_router.launch_router \ --pd-disaggregation \ --policy cache_aware \ --prefill http://P_IP:8000 8998 \ --prefill http://P_IP:8000 8999 \ --decode http://D_IP:8001 \ --host 127.0.0.1 \ --port 6688 \ --mini-lb ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 6688 --max-concurrency 768 --random-input-len 3500 --random-output-len 1500 --num-prompts 3072 --random-range-ratio 1 --request-rate 16 ``` ### DeepSeek-R1 6K-1_6K 20ms on A3 32 Cards Separation Mode Model: Deepseek R1 Hardware: Atlas 800I A3 32Card DeployMode: PD Separation Dataset: random Input Output Length: 6K+1.6K TPOT: 20ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export ASCEND_MF_STORE_URL="tcp://your prefill ip1:24669" P_IP=('your prefill ip1' 'your prefill ip2') D_IP=('your decode ip1' 'your decode ip2') MODEL_PATH=xxx export SGLANG_NPU_USE_MLAPO=1 export SGLANG_USE_FIA_NZ=1 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" # prefill for i in "${!P_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${P_IP[$i]}" || "$LOCAL_HOST2" == "${P_IP[$i]}" ]]; then echo "${P_IP[$i]}" export HCCL_BUFFSIZE=1536 export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 export TASK_QUEUE_ENABLE=2 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode prefill --host ${P_IP[$i]} \ --port 8000 --disaggregation-bootstrap-port $((8998+$i)) --trust-remote-code --nnodes 1 --node-rank 0 \ --tp-size 16 --mem-fraction-static 0.81 --attention-backend ascend --device npu --quantization modelslim \ --disaggregation-transfer-backend ascend --max-running-requests 4 --context-length 8192 --disable-radix-cache \ --chunked-prefill-size -1 --max-prefill-tokens 28680 --moe-a2a-backend deepep --deepep-mode normal \ --speculative-algorithm NEXTN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \ --dp-size 2 --enable-dp-attention --disable-shared-experts-fusion --dtype bfloat16 --enable-attn-tp-input-scattered NODE_RANK=$i break fi done # decode for i in "${!D_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${D_IP[$i]}" || "$LOCAL_HOST2" == "${D_IP[$i]}" ]]; then echo "${D_IP[$i]}" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 export HCCL_BUFFSIZE=650 export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=12 export TASK_QUEUE_ENABLE=1 export SGLANG_SCHEDULER_SKIP_ALL_GATHER=1 export HCCL_SOCKET_IFNAME=xxx export GLOO_SOCKET_IFNAME=xxx python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode decode --host ${D_IP[$i]} \ --port 8001 --trust-remote-code --dist-init-addr DIP1:5000 --nnodes 2 --node-rank $i --tp-size 32 --dp-size 16 \ --mem-fraction-static 0.75 --max-running-requests 32 --attention-backend ascend --device npu --quantization modelslim \ --moe-a2a-backend deepep --enable-dp-attention --deepep-mode low_latency --enable-dp-lm-head --moe-dense-tp 1 \ --cuda-graph-bs 2 4 6 --disaggregation-transfer-backend ascend --watchdog-timeout 9000 --context-length 8192 \ --speculative-algorithm NEXTN --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --tokenizer-worker-num 4 --prefill-round-robin-balance --disable-shared-experts-fusion --dtype bfloat16 \ --load-balance-method decode_round_robin NODE_RANK=$i break fi done ``` ```shell export SGLANG_DP_ROUND_ROBIN=1 python -m sglang_router.launch_router \ --pd-disaggregation \ --policy cache_aware \ --prefill http://P_IP:8000 8998 \ --prefill http://P_IP:8000 8999 \ --decode http://D_IP:8001 \ --host 127.0.0.1 \ --port 6688 \ --mini-lb ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 6688 --max-concurrency 32 --random-input-len 6000 --random-output-len 1600 --num-prompts 32 --random-range-ratio 1 ``` ### DeepSeek-R1 3_9K-1K 20ms on A3 32 Cards Separation Mode Model: Deepseek R1 Hardware: Atlas 800I A3 32Card DeployMode: PD Separation Dataset: random Input Output Length: 3.9K+1K TPOT: 20ms #### Model Deployment Please Turn to [DeepSeek-R1 6K-1_6K 20ms on A3 32 Cards Separation Mode](#deepseek-r1-6k-1_6k-20ms-on-a3-32-cards-separation-mode) #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 6688 --max-concurrency 768 --random-input-len 3900 --random-output-len 1000 --num-prompts 768 --random-range-ratio 1 --request-rate 16 ``` ### DeepSeek-R1 3_5K-1_5K 20ms on A3 32 Cards Separation Mode Model: Deepseek R1 Hardware: Atlas 800I A3 32Card DeployMode: PD Separation Dataset: random Input Output Length: 3.5K+1.5K TPOT: 20ms #### Model Deployment Please Turn to [DeepSeek-R1 6K-1_6K 20ms on A3 32 Cards Separation Mode](#deepseek-r1-6k-1_6k-20ms-on-a3-32-cards-separation-mode) #### Benchmark We tested it based on the `RANDOM` dataset. ```bash python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 6688 --max-concurrency 768 --random-input-len 3500 --random-output-len 1500 --num-prompts 768 --random-range-ratio 1 --request-rate 16 ``` ### DeepSeek-R1 3_5K-1K 20ms on A3 32 Cards Separation Mode Model: Deepseek R1 Hardware: Atlas 800I A3 32Card DeployMode: PD Separation Dataset: random Input Output Length: 3.5K+1K TPOT: 20ms #### Model Deployment Please Turn to [DeepSeek-R1 6K-1_6K 20ms on A3 32 Cards Separation Mode](#deepseek-r1-6k-1_6k-20ms-on-a3-32-cards-separation-mode) #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 6688 --max-concurrency 768 --random-input-len 3500 --random-output-len 1000 --num-prompts 768 --random-range-ratio 1 --request-rate 16 ``` ### DeepSeek-R1 2K-2K 50ms on A3 8 Cards Mixed Mode Model: Deepseek R1 Hardware: Atlas 800I A3 8Card DeployMode: PD Mixed Dataset: random Input Output Length: 2K+2K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE=1 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=64 export HCCL_BUFFSIZE=1600 export DEEPEP_NORMAL_LONG_SEQ_ROUND=10 export DEEPEP_NORMAL_LONG_SEQ_PER_ROUND_TOKENS=512 MODEL_PATH=xxx export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 export SGLANG_NPU_USE_MLAPO=1 export SGLANG_ENABLE_SPEC_V2=1 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_USE_FIA_NZ=1 export ENABLE_MOE_NZ=1 python3 -m sglang.launch_server --model-path ${MODEL_PATH} \ --tp 16 \ --trust-remote-code \ --attention-backend ascend \ --device npu \ --quantization modelslim \ --watchdog-timeout 9000 \ --host 127.0.0.1 --port 6699 \ --cuda-graph-bs 4 8 16 \ --mem-fraction-static 0.74 \ --max-running-requests 256 \ --disable-radix-cache --chunked-prefill-size -1 --max-prefill-tokens 1500 \ --moe-a2a-backend deepep --deepep-mode auto \ --enable-dp-attention --dp-size 16 --enable-dp-lm-head \ --speculative-algorithm NEXTN --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --dtype bfloat16 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 6699 --max-concurrency 256 --random-input-len 2048 --random-output-len 2048 --num-prompts 1024 --random-range-ratio 1 ``` ### DeepSeek-R1 2K-2K 50ms on A3 16 Cards Separation Mode Model: Deepseek R1 Hardware: Atlas 800I A3 16Card DeployMode: PD Separation Dataset: random Input Output Length: 2K+2K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export ASCEND_MF_STORE_URL="tcp://your prefill ip1:24667" P_IP=('your prefill ip1') D_IP=('your decode ip1') MODEL_PATH=xxx export SGLANG_NPU_USE_MLAPO=1 export SGLANG_USE_FIA_NZ=1 export ENABLE_MOE_NZ=1 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" # prefill for i in "${!P_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${P_IP[$i]}" || "$LOCAL_HOST2" == "${P_IP[$i]}" ]]; then echo "${P_IP[$i]}" export HCCL_BUFFSIZE=1536 export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 export TASK_QUEUE_ENABLE=2 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode prefill --host ${P_IP[$i]} \ --port 8000 --disaggregation-bootstrap-port $((8998+$i)) --trust-remote-code --nnodes 1 --node-rank 0 \ --tp-size 16 --mem-fraction-static 0.6 --attention-backend ascend --device npu --quantization modelslim \ --disaggregation-transfer-backend ascend --max-running-requests 8 --context-length 8192 --disable-radix-cache \ --chunked-prefill-size 32768 --max-prefill-tokens 28680 --moe-a2a-backend deepep --deepep-mode normal \ --speculative-algorithm NEXTN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \ --dp-size 2 --enable-dp-attention --disable-shared-experts-fusion --dtype bfloat16 NODE_RANK=$i break fi done # decode for i in "${!D_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${D_IP[$i]}" || "$LOCAL_HOST2" == "${D_IP[$i]}" ]]; then echo "${D_IP[$i]}" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 export HCCL_BUFFSIZE=720 export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=96 export TASK_QUEUE_ENABLE=1 export HCCL_SOCKET_IFNAME=xxx export GLOO_SOCKET_IFNAME=xxx python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode decode --host ${D_IP[$i]} \ --port 8001 --trust-remote-code --nnodes 1 --node-rank 0 --tp-size 16 --dp-size 16 \ --mem-fraction-static 0.8 --max-running-requests 384 --attention-backend ascend --device npu --quantization modelslim \ --moe-a2a-backend deepep --enable-dp-attention --deepep-mode low_latency --enable-dp-lm-head \ --cuda-graph-bs 8 10 12 14 16 18 20 22 24 --disaggregation-transfer-backend ascend --watchdog-timeout 9000 --context-length 8192 \ --speculative-algorithm NEXTN --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --prefill-round-robin-balance --disable-shared-experts-fusion --dtype bfloat16 --tokenizer-worker-num 4 \ --load-balance-method decode_round_robin NODE_RANK=$i break fi done ``` ```shell export SGLANG_DP_ROUND_ROBIN=1 python -m sglang_router.launch_router \ --pd-disaggregation \ --policy cache_aware \ --prefill http://P_IP:8000 8998 \ --decode http://D_IP:8001 \ --host 127.0.0.1 \ --port 6688 \ --mini-lb ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 6688 --max-concurrency 400 --random-input-len 2048 --random-output-len 2048 --num-prompts 3200 --random-range-ratio 1 --request-rate 8 ``` ### DeepSeek-R1 3_5K-1_5K 50ms on A3 8 Cards Mixed Mode Model: Deepseek R1 Hardware: Atlas 800I A3 8Card DeployMode: PD Mixed Dataset: random Input Output Length: 3.5K+1.5K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE=1 export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=36 export HCCL_BUFFSIZE=1600 export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 export SGLANG_NPU_USE_MLAPO=1 export SGLANG_ENABLE_SPEC_V2=1 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_USE_FIA_NZ=1 export ENABLE_MOE_NZ=1 MODEL_PATH=xxx python3 -m sglang.launch_server --model-path ${MODEL_PATH} \ --tp 16 \ --trust-remote-code \ --attention-backend ascend \ --device npu \ --quantization modelslim \ --watchdog-timeout 9000 \ --host 127.0.0.1 --port 6699 \ --cuda-graph-bs 8 16 24 28 32 36 \ --mem-fraction-static 0.71 \ --max-running-requests 144 \ --context-length 8188 --disable-radix-cache --chunked-prefill-size -1 --max-prefill-tokens 9000 \ --moe-a2a-backend deepep --deepep-mode auto \ --enable-dp-attention --dp-size 4 --enable-dp-lm-head \ --speculative-algorithm NEXTN --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --dtype bfloat16 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 6699 --max-concurrency 144 --random-input-len 3500 --random-output-len 1500 --num-prompts 576 --random-range-ratio 1 ``` ### DeepSeek-R1 3_5K-1_5K 50ms on A3 16 Cards Separation Mode Model: Deepseek R1 Hardware: Atlas 800I A3 16Card DeployMode: PD Separation Dataset: random Input Output Length: 3.5K+1.5K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export ASCEND_MF_STORE_URL="tcp://your prefill ip1:24667" P_IP=('your prefill ip1') D_IP=('your decode ip1') MODEL_PATH=xxx export SGLANG_NPU_USE_MLAPO=1 export SGLANG_USE_FIA_NZ=1 export ENABLE_MOE_NZ=1 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" # prefill for i in "${!P_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${P_IP[$i]}" || "$LOCAL_HOST2" == "${P_IP[$i]}" ]]; then echo "${P_IP[$i]}" export HCCL_BUFFSIZE=1536 export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 export TASK_QUEUE_ENABLE=2 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode prefill --host ${P_IP[$i]} \ --port 8000 --disaggregation-bootstrap-port $((8998+$i)) --trust-remote-code --nnodes 1 --node-rank 0 \ --tp-size 16 --mem-fraction-static 0.6 --attention-backend ascend --device npu --quantization modelslim \ --disaggregation-transfer-backend ascend --max-running-requests 8 --context-length 8192 --disable-radix-cache \ --chunked-prefill-size -1 --max-prefill-tokens 28680 --moe-a2a-backend deepep --deepep-mode normal \ --speculative-algorithm NEXTN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \ --dp-size 2 --enable-dp-attention --disable-shared-experts-fusion --dtype bfloat16 NODE_RANK=$i break fi done # decode for i in "${!D_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${D_IP[$i]}" || "$LOCAL_HOST2" == "${D_IP[$i]}" ]]; then echo "${D_IP[$i]}" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 export HCCL_BUFFSIZE=720 export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=96 export TASK_QUEUE_ENABLE=1 export HCCL_SOCKET_IFNAME=xxx export GLOO_SOCKET_IFNAME=xxx python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode decode --host ${D_IP[$i]} \ --port 8001 --trust-remote-code --nnodes 1 --node-rank 0 --tp-size 16 --dp-size 16 \ --mem-fraction-static 0.8 --max-running-requests 384 --attention-backend ascend --device npu --quantization modelslim \ --moe-a2a-backend deepep --enable-dp-attention --deepep-mode low_latency --enable-dp-lm-head \ --cuda-graph-bs 8 10 12 14 16 18 20 22 24 --disaggregation-transfer-backend ascend --watchdog-timeout 9000 --context-length 8192 \ --speculative-algorithm NEXTN --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --prefill-round-robin-balance --disable-shared-experts-fusion --dtype bfloat16 --tokenizer-worker-num 4 \ --load-balance-method decode_round_robin NODE_RANK=$i break fi done ``` ```shell export SGLANG_DP_ROUND_ROBIN=1 python -m sglang_router.launch_router \ --pd-disaggregation \ --policy cache_aware \ --prefill http://P_IP:8000 8998 \ --decode http://D_IP:8001 \ --host 127.0.0.1 \ --port 6688 \ --mini-lb ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 6688 --max-concurrency 384 --random-input-len 3500 --random-output-len 1500 --num-prompts 1536 --random-range-ratio 1 ``` ### DeepSeek-V3.2-Exp 64K-3K 30ms on A3 32 Cards Separation Mode Model: DeepSeek-V3.2-Exp-W8A8 Hardware: Atlas 800I A3 32Card DeployMode: PD Separation Dataset: random Input Output Length: 64K+3K TPOT: 30ms #### Model Deployment Deploy Prefill Instance ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/op_api/lib/:${LD_LIBRARY_PATH} export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export HCCL_BUFFSIZE=1024 export DEEPEP_NORMAL_LONG_SEQ_ROUND=5 export DEEPEP_NORMAL_LONG_SEQ_PER_ROUND_TOKENS=512 MODEL_PATH=xxx export SGLANG_NPU_USE_MLAPO=1 export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 export SGLANG_NPU_USE_MULTI_STREAM=1 export HCCL_OP_EXPANSION_MODE=AIV IPs=('your prefill ip1' 'your prefill ip2') # get IP in current node LOCAL_HOST=`hostname -I|awk -F " " '{print$1}'` echo "LOCAL_HOST = " ${LOCAL_HOST} # get node index for i in "${!IPs[@]}"; do echo "LOCAL_HOST=${LOCAL_HOST}, IPs[${i}]=${IPs[$i]}" if [ "$LOCAL_HOST" == "${IPs[$i]}" ]; then echo "Node Rank : ${i}" VC_TASK_INDEX=$i break fi done IFNAMES=('xxx' 'xxx') export HCCL_SOCKET_IFNAME=${IFNAMES[$VC_TASK_INDEX]} export GLOO_SOCKET_IFNAME=${HCCL_SOCKET_IFNAME} echo "HCCL_SOCKET_IFNAME : ${HCCL_SOCKET_IFNAME}" nnodes=${#IPs[@]} tp_size=`expr 16 \* ${nnodes}` export ASCEND_MF_STORE_URL=tcp://${IPs[0]}:24667 python3 -m sglang.launch_server --model-path ${MODEL_PATH} \ --tp $tp_size \ --trust-remote-code \ --attention-backend ascend \ --device npu \ --watchdog-timeout 9000 \ --host ${IPs[$VC_TASK_INDEX]} --port 8000 \ --mem-fraction-static 0.73 \ --disable-radix-cache --chunked-prefill-size -1 --max-prefill-tokens 68000 \ --max-running-requests 1 \ --moe-a2a-backend deepep --deepep-mode normal \ --quantization modelslim \ --disaggregation-transfer-backend ascend \ --disaggregation-mode prefill \ --disable-cuda-graph \ --nnodes $nnodes --node-rank $VC_TASK_INDEX \ --disaggregation-bootstrap-port 8995 \ --enable-nsa-prefill-context-parallel --moe-dense-tp-size 1 \ --speculative-algorithm NEXTN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \ --dist-init-addr ${IPs[0]}:10000 ``` Deploy Decode Instance ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/op_api/lib/:${LD_LIBRARY_PATH} export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 MODEL_PATH=xxx export SGLANG_NPU_USE_MULTI_STREAM=1 export SGLANG_NPU_USE_MLAPO=1 export HCCL_OP_EXPANSION_MODE=AIV export SGLANG_SCHEDULER_SKIP_ALL_GATHER=1 export TASK_QUEUE_ENABLE=0 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 IPs=('your decode ip1' 'your decode ip2') export prefill_ip=your prefill ip1 # get IP in current node LOCAL_HOST=`hostname -I|awk -F " " '{print$1}'` echo "LOCAL_HOST = " ${LOCAL_HOST} # get node index for i in "${!IPs[@]}"; do echo "LOCAL_HOST=${LOCAL_HOST}, IPs[${i}]=${IPs[$i]}" if [ "$LOCAL_HOST" == "${IPs[$i]}" ]; then echo "Node Rank : ${i}" VC_TASK_INDEX=$i break fi done IFNAMES=('xxx' 'xxx') export HCCL_SOCKET_IFNAME=${IFNAMES[$VC_TASK_INDEX]} export GLOO_SOCKET_IFNAME=${HCCL_SOCKET_IFNAME} nnodes=${#IPs[@]} tp_size=`expr 16 \* ${nnodes}` export ASCEND_MF_STORE_URL=tcp://${prefill_ip}:24667 CHUNKED_SIZE=65536 DP=8 export HCCL_BUFFSIZE=400 export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=8 python3 -m sglang.launch_server --model-path ${MODEL_PATH} \ --tp $tp_size \ --dp ${DP} \ --ep $tp_size \ --moe-dense-tp-size 1 \ --enable-dp-attention \ --enable-dp-lm-head \ --trust-remote-code \ --attention-backend ascend \ --device npu \ --watchdog-timeout 9000 \ --host ${IPs[$VC_TASK_INDEX]} --port 8001 \ --mem-fraction-static 0.79 \ --disable-radix-cache \ --chunked-prefill-size -1 --max-prefill-tokens 68000 \ --max-running-requests 32 \ --cuda-graph-max-bs 4 \ --moe-a2a-backend deepep \ --deepep-mode low_latency \ --quantization modelslim \ --speculative-algorithm NEXTN --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --disaggregation-transfer-backend ascend \ --disaggregation-mode decode \ --prefill-round-robin-balance \ --load-balance-method round_robin \ --nnodes $nnodes --node-rank $VC_TASK_INDEX \ --dist-init-addr ${IPs[0]}:10000 --load-balance-method decode_round_robin ``` ```shell export SGLANG_DP_ROUND_ROBIN=1 python -m sglang_router.launch_router \ --pd-disaggregation \ --policy cache_aware \ --prefill http://PIP1:8000 8995 \ --decode http://DIP1:8001 \ --host 127.0.0.1 \ --port 6688 \ --mini-lb ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 6688 --max-concurrency 32 --random-input-len 64000 --random-output-len 3000 --num-prompts 64 --random-range-ratio 1 ``` ### Qwen3-235B-A22B 3_5K-1_5K 50ms on A3 24 Cards Separation Mode Model: Qwen3-235B-A22B-W8A8 Hardware: Atlas 800I A3 24Card DeployMode: PD Separation Dataset: random Input Output Length: 3.5K+1.5K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=16 MODEL_PATH=xxx export ASCEND_MF_STORE_URL="tcp://your prefill ip1:24667" P_IP=('your prefill ip1') D_IP=('your decode ip1' 'your decode ip2') export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 export SGLANG_DP_ROUND_ROBIN=1 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" for i in "${!P_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${P_IP[$i]}" || "$LOCAL_HOST2" == "${P_IP[$i]}" ]]; then echo "${P_IP[$i]}" source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export DEEPEP_NORMAL_LONG_SEQ_PER_ROUND_TOKENS=1024 export DEEPEP_NORMAL_LONG_SEQ_ROUND=16 export HCCL_BUFFSIZE=4300 export TASK_QUEUE_ENABLE=2 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export STREAMS_PER_DEVICE=32 export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 # P节点 python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode prefill \ --host ${P_IP[$i]} --port 8000 --disaggregation-bootstrap-port 8995 --trust-remote-code \ --nnodes 1 --node-rank $i --tp-size 16 --dp-size 16 --mem-fraction-static 0.6 \ --disable-radix-cache \ --attention-backend ascend --device npu --quantization modelslim --disaggregation-transfer-backend ascend \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --speculative-draft-model-quantization unquant \ --max-running-requests 128 --chunked-prefill-size 262144 --max-prefill-tokens 262144 \ --enable-dp-attention \ --moe-a2a-backend deepep --deepep-mode normal --dtype bfloat16 NODE_RANK=$i break fi done for i in "${!D_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${D_IP[$i]}" || "$LOCAL_HOST2" == "${D_IP[$i]}" ]]; then echo "${D_IP[$i]}" source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=24 export HCCL_BUFFSIZE=512 export HCCL_SOCKET_IFNAME=data0.3001 export GLOO_SOCKET_IFNAME=data0.3001 export STREAMS_PER_DEVICE=32 python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode decode \ --host ${D_IP[$i]} --port 8001 --trust-remote-code \ --nnodes 2 --node-rank $i --tp-size 32 --dp-size 32 --mem-fraction-static 0.83 --max-running-requests 768 \ --attention-backend ascend --device npu --quantization modelslim --enable-dp-attention \ --moe-a2a-backend ascend_fuseep --cuda-graph-bs 6 8 12 15 18 20 22 24 \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-draft-model-quantization unquant \ --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --dist-init-addr xxx:5000 \ --disaggregation-transfer-backend ascend --watchdog-timeout 9000 --context-length 8192 \ --prefill-round-robin-balance --enable-dp-lm-head --dtype bfloat16 --tokenizer-worker-num 4 \ --load-balance-method decode_round_robin NODE_RANK=$i break fi done ``` ```shell export SGLANG_DP_ROUND_ROBIN=1 python -m sglang_router.launch_router \ --pd-disaggregation \ --policy cache_aware \ --prefill http://PIP:8000 8995 \ --decode http://DIP:8001 \ --host 127.0.0.1 \ --port 6688 \ --mini-lb ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang-oai --host 127.0.0.1 --port 7239 --max-concurrency 860 --random-input-len 3500 --random-output-len 1500 --num-prompts 3440 --random-range-ratio 1 ``` ### Qwen3-235B-A22B 3_5K-1_5K 50ms on A3 8 Cards Mixed Mode Model: Qwen3-235B-A22B-W8A8 Hardware: Atlas 800I A3 8Card DeployMode: PD Mixed Dataset: random Input Output Length: 3.5K+1.5K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=1600 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 export SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE=2 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7439 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu --quantization modelslim \ --max-running-requests 272 --context-length 8192 --dtype bfloat16 \ --chunked-prefill-size 32768 --max-prefill-tokens 32768 \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --disable-radix-cache --moe-a2a-backend deepep --deepep-mode auto --speculative-draft-model-quantization unquant \ --tp 16 --dp-size 16 --enable-dp-attention --enable-dp-lm-head --mem-fraction-static 0.8 --cuda-graph-bs 3 4 6 8 10 12 13 14 15 16 17 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7439 --max-concurrency 272 --random-input-len 3500 --random-output-len 1500 --num-prompts 1088 --random-range-ratio 1 ``` ### Qwen3-235B-A22B 2K-2K 100ms on A3 8 Cards Mixed Mode Model: Qwen3-235B-A22B-W8A8 Hardware: Atlas 800I A3 8Card DeployMode: PD Mixed Dataset: random Input Output Length: 2K+2K TPOT: 100ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=1200 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7439 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu --quantization modelslim \ --max-running-requests 576 --context-length 8192 --dtype bfloat16 \ --chunked-prefill-size 32768 --max-prefill-tokens 458880 \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --disable-radix-cache --moe-a2a-backend deepep --deepep-mode auto --speculative-draft-model-quantization unquant \ --tp 16 --dp-size 16 --enable-dp-attention --enable-dp-lm-head --mem-fraction-static 0.81 --cuda-graph-bs 8 16 20 24 32 36 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7439 --max-concurrency 576 --random-input-len 2000 --random-output-len 2000 --num-prompts 576 --random-range-ratio 1 ``` ### Qwen3-235B-A22B 2K-2K 50ms on A3 8 Cards Mixed Mode Model: Qwen3-235B-A22B-W8A8 Hardware: Atlas 800I A3 8Card DeployMode: PD Mixed Dataset: random Input Output Length: 2K+2K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=2100 export HCCL_SOCKET_IFNAME=xxx export GLOO_SOCKET_IFNAME=xxx export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7439 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu --quantization modelslim \ --max-running-requests 480 --context-length 8192 --dtype bfloat16 \ --chunked-prefill-size -1 --max-prefill-tokens 4096 --speculative-draft-model-quantization unquant \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --disable-radix-cache --moe-a2a-backend deepep --deepep-mode auto \ --tp 16 --dp-size 16 --enable-dp-attention --enable-dp-lm-head --mem-fraction-static 0.75 --cuda-graph-bs 6 8 10 12 15 18 28 30 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7439 --max-concurrency 480 --random-input-len 2048 --random-output-len 2048 --num-prompts 480 --random-range-ratio 1 ``` ### Qwen3-235B-A22B 2K-2K 50ms on A3 16 Cards Mixed Mode Model: Qwen3-235B-A22B-W8A8 Hardware: Atlas 800I A3 16Card DeployMode: PD Mixed Dataset: random Input Output Length: 2K+2K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=1600 export HCCL_SOCKET_IFNAME=xxx export GLOO_SOCKET_IFNAME=xxx export HCCL_OP_EXPANSION_MODE="AIV" MIX_IP=('IP1' 'IP2') for i in "${!MIX_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${MIX_IP[$i]}" || "$LOCAL_HOST2" == "${MIX_IP[$i]}" ]]; then echo "${MIX_IP[$i]}" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path ${MODEL_PATH} \ --host 127.0.0.1 --port 7439 --trust-remote-code \ --nnodes 2 --node-rank $i --tp-size 32 --dp-size 32 --mem-fraction-static 0.8 --max-running-requests 768 \ --attention-backend ascend --device npu --quantization modelslim --enable-dp-attention \ --moe-a2a-backend deepep --deepep-mode auto --cuda-graph-bs 6 8 10 12 18 24 \ --dist-init-addr ${MIX_IP[0]}:5000 --chunked-prefill-size 131072 --max-prefill-tokens 458880 \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx --speculative-draft-model-quantization= unquant \ --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --context-length 8192 --disable-radix-cache \ --enable-dp-lm-head --dtype bfloat16 NODE_RANK=$i break fi done ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7439 --max-concurrency 768 --random-input-len 2000 --random-output-len 2000 --num-prompts 768 --random-range-ratio 1 ``` ### Qwen3-235B-A22B 11K-1K 10ms on A3 8 Cards Mixed Mode Model: Qwen3-235B-A22B-W8A8 Hardware: Atlas 800I A3 8Card DeployMode: PD Mixed Dataset: random Input Output Length: 11K+1K TPOT: 10ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=1600 export HCCL_SOCKET_IFNAME=xxx export GLOO_SOCKET_IFNAME=xxx export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7439 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu --quantization modelslim \ --max-running-requests 1 --dtype bfloat16 \ --chunked-prefill-size -1 --max-prefill-tokens 16384 --speculative-draft-model-quantization unquant \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 4 --speculative-eagle-topk 1 --speculative-num-draft-tokens 5 \ --disable-radix-cache --enable-dp-lm-head \ --tp 16 --mem-fraction-static 0.78 --cuda-graph-bs 1 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7439 --max-concurrency 1 --random-input-len 11000 --random-output-len 1000 --num-prompts 1 --random-range-ratio 1 ``` ### Qwen3-32B 6K-1_5K 18ms on A3 4 Cards Mixed Mode Model: Qwen3-32B Hardware: Atlas 800I A3 4Card DeployMode: PD Mixed Dataset: random Input Output Length: 6K+1.5K TPOT: 18ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=xxx export GLOO_SOCKET_IFNAME=xxx export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7439 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu \ --max-running-requests 32 \ --disable-radix-cache \ --chunked-prefill-size 24576 --max-prefill-tokens 65536 \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 4 --speculative-eagle-topk 1 --speculative-num-draft-tokens 5 \ --tp-size 8 --mem-fraction-static 0.72 --cuda-graph-bs 8 16 24 32 --dtype bfloat16 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python3 -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7439 --max-concurrency 32 --random-output-len 1500 --random-input-len 6000 --num-prompts 32 --random-range-ratio 1 ``` ### Qwen3-32B 4K-1_5K 11ms on A3 4 Cards Mixed Mode Model: Qwen3-32B Hardware: Atlas 800I A3 4Card DeployMode: PD Mixed Dataset: random Input Output Length: 4K+1.5K TPOT: 11ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7339 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu \ --max-running-requests 1 \ --disable-radix-cache \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 4 --speculative-eagle-topk 1 --speculative-num-draft-tokens 5 \ --chunked-prefill-size 24576 --max-prefill-tokens 65536 \ --tp-size 8 --mem-fraction-static 0.72 --cuda-graph-bs 1 --dtype bfloat16 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python3 -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7239 --random-range-ratio 1 --max-concurrency 1 --random-output-len 1500 --random-input-len 4096 --num-prompts 4 ``` ### Qwen3-32B 18K-4K 12ms on A3 8 Cards Mixed Mode Model: Qwen3-32B Hardware: Atlas 800I A3 8Card DeployMode: PD Mixed Dataset: random Input Output Length: 18K+4K TPOT: 12ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7339 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu \ --max-running-requests 1 \ --disable-radix-cache --speculative-draft-model-quantization unquant \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 4 --speculative-eagle-topk 1 --speculative-num-draft-tokens 5 \ --chunked-prefill-size -1 --max-prefill-tokens 65536 \ --tp-size 16 --mem-fraction-static 0.72 --cuda-graph-bs 1 --dtype bfloat16 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python3 -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7339 --random-range-ratio 1 --max-concurrency 1 --random-output-len 18000 --random-input-len 4000 --num-prompts 1 ``` ### Qwen3-32B 3_5K-1_5K 50ms on A3 2 Cards Mixed Mode Model: Qwen3-32B Hardware: Atlas 800I A3 2Card DeployMode: PD Mixed Dataset: random Input Output Length: 3.5K+1.5K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7239 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu --quantization modelslim \ --max-running-requests 78 \ --disable-radix-cache --speculative-draft-model-quantization unquant \ --chunked-prefill-size -1 --max-prefill-tokens 49152 \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --tp-size 4 --mem-fraction-static 0.72 --cuda-graph-bs 16 32 64 68 72 78 --dtype bfloat16 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python3 -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7239 --max-concurrency 78 --random-output-len 1500 --random-input-len 3500 --num-prompts 312 --random-range-ratio 1 ``` ### Qwen3-32B 2K-2K 50ms on A3 2 Cards Mixed Mode Model: Qwen3-32B Hardware: Atlas 800I A3 2Card DeployMode: PD Mixed Dataset: random Input Output Length: 2K+2K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7239 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu --quantization modelslim \ --max-running-requests 120 \ --disable-radix-cache --speculative-draft-model-quantization unquant \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --chunked-prefill-size -1 --max-prefill-tokens 49152 \ --tp-size 4 --mem-fraction-static 0.7 --cuda-graph-bs 54 60 66 72 78 84 90 108 114 120 --dtype bfloat16 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python3 -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7239 --max-concurrency 120 --random-output-len 2000 --random-input-len 2000 --num-prompts 480 --random-range-ratio 1 ``` ### Qwen3-30B-A3B 3_5K-1_5K 50ms on A3 1 Card Mixed Mode Model: Qwen3-30B-A3B-Instruct-2507 Hardware: Atlas 800I A3 1Card DeployMode: PD Mixed Dataset: random Input Output Length: 3.5K+1.5K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7239 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu --quantization modelslim \ --max-running-requests 192 \ --disable-radix-cache \ --speculative-draft-model-quantization unquant \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --chunked-prefill-size -1 --max-prefill-tokens 32768 \ --tp-size 2 --mem-fraction-static 0.86 --cuda-graph-bs 42 88 96 132 144 156 172 178 192 --dtype bfloat16 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7239 --max-concurrency 156 --random-input-len 3500 --random-output-len 1500 --num-prompts 624 --random-range-ratio 1 ``` ### Qwen3-Coder-480B-A35B-Instruct 3_5K-1_5K 50ms on A3 24 Cards Separation Mode Model: Qwen3-Coder-480B-A35B-Instruct Hardware: Atlas 800I A3 24Card DeployMode: PD Separation Dataset: random Input Output Length: 3.5K+1.5K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=16 MODEL_PATH=xxx export ASCEND_MF_STORE_URL="tcp://PIP:24667" P_IP=('PIP') D_IP=('DIP1' 'DIP2') export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" for i in "${!P_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${P_IP[$i]}" || "$LOCAL_HOST2" == "${P_IP[$i]}" ]]; then echo "${P_IP[$i]}" source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export DEEPEP_NORMAL_LONG_SEQ_PER_ROUND_TOKENS=1024 export DEEPEP_NORMAL_LONG_SEQ_ROUND=16 export HCCL_BUFFSIZE=4300 export TASK_QUEUE_ENABLE=2 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export STREAMS_PER_DEVICE=32 export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode prefill \ --host ${P_IP[$i]} --port 8000 --disaggregation-bootstrap-port 8995 --trust-remote-code \ --nnodes 1 --node-rank $i --tp-size 16 --dp-size 2 --mem-fraction-static 0.6 \ --disable-radix-cache \ --attention-backend ascend --device npu --quantization modelslim --disaggregation-transfer-backend ascend \ --max-running-requests 128 --chunked-prefill-size 65536 --max-prefill-tokens 262144 \ --enable-dp-attention \ --moe-a2a-backend deepep --deepep-mode normal --dtype bfloat16 NODE_RANK=$i break fi done for i in "${!D_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${D_IP[$i]}" || "$LOCAL_HOST2" == "${D_IP[$i]}" ]]; then echo "${D_IP[$i]}" source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=72 export HCCL_BUFFSIZE=512 export HCCL_SOCKET_IFNAME=xxx export GLOO_SOCKET_IFNAME=xxx export STREAMS_PER_DEVICE=32 python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode decode \ --host ${D_IP[$i]} --port 8001 --trust-remote-code \ --nnodes 2 --node-rank $i --tp-size 32 --dp-size 4 --mem-fraction-static 0.73 --max-running-requests 384 \ --attention-backend ascend --device npu --quantization modelslim --enable-dp-attention \ --moe-a2a-backend ascend_fuseep --cuda-graph-bs 16 32 48 56 64 72 80 88 96 \ --dist-init-addr DIP1:5000 \ --disaggregation-transfer-backend ascend --watchdog-timeout 9000 --context-length 8192 \ --prefill-round-robin-balance --enable-dp-lm-head --dtype bfloat16 --tokenizer-worker-num 4 --load-balance-method decode_round_robin NODE_RANK=$i break fi done ``` ```shell export SGLANG_DP_ROUND_ROBIN=1 python -m sglang_router.launch_router \ --pd-disaggregation \ --policy cache_aware \ --prefill http://PIP:8000 8995 \ --decode http://DIP:8001 \ --host 127.0.0.1 \ --port 6688 \ --mini-lb ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7239 --max-concurrency 410 --random-input-len 3500 --random-output-len 1500 --num-prompts 1640 --random-range-ratio 1 --request-rate 8 ``` ### Qwen3-Coder-480B-A35B-Instruct 3_5K-1_5K 50ms on A3 16 Cards Mixed Mode Model: Qwen3-Coder-480B-A35B-Instruct Hardware: Atlas 800I A3 16Card DeployMode: PD Mixed Dataset: random Input Output Length: 3.5K+1.5K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=16 export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=1800 export HCCL_SOCKET_IFNAME=xxx export GLOO_SOCKET_IFNAME=xxx export HCCL_OP_EXPANSION_MODE="AIV" MIX_IP=('IP1' 'IP2') for i in "${!MIX_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${MIX_IP[$i]}" || "$LOCAL_HOST2" == "${MIX_IP[$i]}" ]]; then echo "${MIX_IP[$i]}" python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7439 --trust-remote-code --nnodes 2 --node-rank $i \ --dist-init-addr 141.61.133.128:5000 \ --attention-backend ascend --device npu --quantization modelslim \ --max-running-requests 288 --context-length 8192 --dtype bfloat16 \ --chunked-prefill-size 114688 --max-prefill-tokens 458880 \ --disable-radix-cache --moe-a2a-backend deepep --deepep-mode auto \ --tp 32 --dp-size 4 --enable-dp-attention --enable-dp-lm-head --mem-fraction-static 0.7 --cuda-graph-bs 56 64 72 NODE_RANK=$i break fi done ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7439 --max-concurrency 288 --random-input-len 3500 --random-output-len 1500 --num-prompts 1152 --random-range-ratio 1 --request-rate 20 ``` ### Qwen3-Coder-480B-A35B-Instruct 3_5K-1_5K 50ms on A3 8 Cards Mixed Mode Model: Qwen3-Coder-480B-A35B-Instruct Hardware: Atlas 800I A3 8Card DeployMode: PD Mixed Dataset: random Input Output Length: 3.5K+1.5K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=2100 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7439 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu --quantization modelslim \ --max-running-requests 80 --context-length 8192 --dtype bfloat16 \ --chunked-prefill-size 28672 --max-prefill-tokens 458880 \ --disable-radix-cache --moe-a2a-backend deepep --deepep-mode auto --enable-dp-attention --enable-dp-lm-head \ --tp 16 --dp-size 4 --mem-fraction-static 0.7 --cuda-graph-bs 16 20 24 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7439 --max-concurrency 80 --random-input-len 3500 --random-output-len 1500 --num-prompts 320 --random-range-ratio 1 ``` ### Qwen3-Next-80B-A3B-Instruct 3_5K-1_5K 50ms on A3 2 Cards Mixed Mode Model: Qwen3-Next-80B-A3B-Instruct Hardware: Atlas 800I A3 2Card DeployMode: PD Mixed Dataset: random Input Output Length: 3.5K+1.5K TPOT: 50ms #### Model Deployment ```shell export cann_path=/usr/local/Ascend/ascend-toolkit/latest source /usr/local/Ascend/driver/bin/setenv.bash source ${cann_path}/../set_env.sh source ${cann_path}/../../nnal/atb/set_env.sh source ${cann_path}/opp/vendors/customize/bin/set_env.bash export ASCEND_HOME_PATH=${cann_path} source /usr/local/Ascend/8.5.0/bisheng_toolkit/set_env.sh echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE=AIV export HCCL_ALGO="level0:NA;level1:ring" export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=20 export HCCL_BUFFSIZE=2000 python -m sglang.launch_server \ --model-path /path/to/Qwen3-Next-80B-A3B-Instruct-W8A8-3 \ --host 127.0.0.1 \ --port 6699 \ --tp-size 4 \ --device npu \ --attention-backend ascend \ --mem-fraction-static 0.685 \ --max-running-requests 80 \ --watchdog-timeout 3600 \ --disable-radix-cache \ --cuda-graph-bs 80 \ --max-prefill-tokens 28672 --max-total-tokens 450560 \ --moe-a2a-backend deepep --deepep-mode auto \ --quantization modelslim \ --chunked-prefill-size -1 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python3 -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 6699 --max-concurrency 80 --random-output-len 1536 --random-input-len 3584 --num-prompts 160 --random-range-ratio 1 ``` ### Qwen3-32B 6K-1_5K 18ms on A2 8 Cards Mixed Mode Model: Qwen3-32B Hardware: Atlas 800I A2 8Card DeployMode: PD Mixed Dataset: random Input Output Length: 6K+1.5K TPOT: 18ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7439 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu --quantization modelslim \ --max-running-requests 32 \ --disable-radix-cache \ --chunked-prefill-size 24576 --max-prefill-tokens 65536 \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 4 --speculative-eagle-topk 1 --speculative-num-draft-tokens 5 \ --tp-size 8 --mem-fraction-static 0.72 --cuda-graph-bs 8 16 24 32 --dtype bfloat16 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python3 -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7439 --max-concurrency 32 --random-output-len 1500 --random-input-len 6000 --num-prompts 32 --random-range-ratio 1 ``` ### Qwen3-32B 4K-1_5K 11ms on A2 8 Cards Mixed Mode Model: Qwen3-32B Hardware: Atlas 800I A2 8Card DeployMode: PD Mixed Dataset: random Input Output Length: 4K+1.5K TPOT: 11ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7339 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu \ --max-running-requests 32 \ --disable-radix-cache \ --speculative-draft-model-quantization unquant \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 4 --speculative-eagle-topk 1 --speculative-num-draft-tokens 5 \ --chunked-prefill-size -1 --max-prefill-tokens 65536 \ --tp-size 8 --mem-fraction-static 0.72 --cuda-graph-bs 1 4 6 12 18 24 30 32 --dtype bfloat16 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python3 -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7339 --random-range-ratio 1 --max-concurrency 1 --random-output-len 1500 --random-input-len 4096 --num-prompts 4 ``` ### Qwen3-32B 3_5K-1_5K 50ms on A2 8 Cards Mixed Mode Model: Qwen3-32B Hardware: Atlas 800I A2 8Card DeployMode: PD Mixed Dataset: random Input Output Length: 3.5K+1.5K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7239 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu --quantization modelslim \ --max-running-requests 78 \ --disable-radix-cache --speculative-draft-model-quantization unquant \ --chunked-prefill-size -1 --max-prefill-tokens 65536 \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --tp-size 4 --mem-fraction-static 0.72 --cuda-graph-bs 1 4 8 16 32 64 68 72 78 --dtype bfloat16 --base-gpu-id 4 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python3 -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7239 --max-concurrency 78 --random-output-len 1500 --random-input-len 3500 --num-prompts 312 --random-range-ratio 1 ``` ### Qwen3-32B 2K-2K 50ms on A2 8 Cards Mixed Mode Model: Qwen3-32B Hardware: Atlas 800I A2 8Card DeployMode: PD Mixed Dataset: random Input Output Length: 2K+2K TPOT: 50ms #### Model Deployment ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/bin/set_env.bash export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH MODEL_PATH=xxx export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server --model-path $MODEL_PATH \ --host 127.0.0.1 --port 7239 --trust-remote-code --nnodes 1 --node-rank 0 \ --attention-backend ascend --device npu --quantization modelslim \ --max-running-requests 120 \ --disable-radix-cache \ --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --speculative-draft-model-quantization unquant \ --chunked-prefill-size -1 --max-prefill-tokens 49152 --base-gpu-id 4 \ --tp-size 4 --mem-fraction-static 0.7 --cuda-graph-bs 54 60 66 72 78 84 90 108 114 120 --dtype bfloat16 ``` #### Benchmark We tested it based on the `RANDOM` dataset. ```shell python3 -m sglang.bench_serving --dataset-name random --backend sglang --host 127.0.0.1 --port 7239 --max-concurrency 120 --random-output-len 2000 --random-input-len 2000 --num-prompts 120 --random-range-ratio 1 ``` ================================================ FILE: docs/platforms/ascend_npu_deepseek_example.md ================================================ ## DeepSeek examples ### Running DeepSeek-V3 #### Running DeepSeek in PD mixed mode on 1 x Atlas 800I A3. W4A8 Model weights could be found [here](https://modelers.cn/models/Modelers_Park/DeepSeek-R1-0528-w4a8). ```shell export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 #Deepep communication settings export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=32 export HCCL_BUFFSIZE=1600 #spec overlap export SGLANG_ENABLE_SPEC_V2=1 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 #npu acceleration operator export SGLANG_NPU_USE_MLAPO=1 export SGLANG_USE_FIA_NZ=1 python3 -m sglang.launch_server \ --model-path ${MODEL_PATH} \ --tp 16 \ --trust-remote-code \ --attention-backend ascend \ --device npu \ --quantization modelslim \ --watchdog-timeout 9000 \ --cuda-graph-bs 8 16 24 28 32 \ --mem-fraction-static 0.68 \ --max-running-requests 128 \ --context-length 8188 \ --disable-radix-cache \ --chunked-prefill-size -1 \ --max-prefill-tokens 16384 \ --moe-a2a-backend deepep \ --deepep-mode auto \ --enable-dp-attention \ --dp-size 4 \ --enable-dp-lm-head \ --speculative-algorithm NEXTN \ --speculative-num-steps 3 \ --speculative-eagle-topk 1 \ --speculative-num-draft-tokens 4 \ --dtype bfloat16 ``` #### Running DeepSeek with PD disaggregation mode on 2 x Atlas 800I A3. W4A8 Model weights could be found [here](https://modelers.cn/models/Modelers_Park/DeepSeek-R1-0528-w4a8). 1. Prefill: ```shell export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 #memfabric config store export ASCEND_MF_STORE_URL="tcp://:" #Deepep communication settings export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 export HCCL_BUFFSIZE=1536 #npu acceleration operator export SGLANG_NPU_USE_MLAPO=1 export SGLANG_USE_FIA_NZ=1 export TASK_QUEUE_ENABLE=2 python -m sglang.launch_server \ --model-path ${MODEL_PATH} \ --host $PREFILL_HOST_IP \ --port 8000 \ --disaggregation-mode prefill \ --disaggregation-bootstrap-port 8996 \ --disaggregation-transfer-backend ascend \ --trust-remote-code \ --nnodes 1 \ --node-rank 0 \ --tp-size 16 \ --mem-fraction-static 0.6 \ --attention-backend ascend \ --device npu \ --quantization modelslim \ --load-balance-method round_robin \ --max-running-requests 8 \ --context-length 8192 \ --disable-radix-cache \ --chunked-prefill-size -1 \ --max-prefill-tokens 28680 \ --moe-a2a-backend deepep \ --deepep-mode normal \ --speculative-algorithm NEXTN \ --speculative-num-steps 3 \ --speculative-eagle-topk 1 \ --speculative-num-draft-tokens 4 \ --dp-size 2 \ --enable-dp-attention \ --disable-shared-experts-fusion \ --dtype bfloat16 ``` 2. Decode: ```shell export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 #memfabric config store export ASCEND_MF_STORE_URL="tcp://:" #Deepep communication settings export HCCL_BUFFSIZE=720 export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=88 #spec overlap export SGLANG_ENABLE_SPEC_V2=1 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 #npu acceleration operator unset TASK_QUEUE_ENABLE export SGLANG_NPU_USE_MLAPO=1 export SGLANG_USE_FIA_NZ=1 # suggest max-running-requests <= max-cuda-graph-bs * dp_size, Because when this value is exceeded, performance will significantly degrade. python -m sglang.launch_server \ --model-path ${MODEL_PATH} \ --disaggregation-mode decode \ --host $DECODE_HOST_IP \ --port 8001 \ --trust-remote-code \ --nnodes 1 \ --node-rank 0 \ --tp-size 16 \ --dp-size 16 \ --mem-fraction-static 0.8 \ --max-running-requests 352 \ --attention-backend ascend \ --device npu \ --quantization modelslim \ --prefill-round-robin-balance \ --moe-a2a-backend deepep \ --enable-dp-attention \ --deepep-mode low_latency \ --enable-dp-lm-head \ --cuda-graph-bs 8 10 12 14 16 18 20 22 \ --disaggregation-transfer-backend ascend \ --watchdog-timeout 9000 \ --context-length 8192 \ --speculative-algorithm NEXTN \ --speculative-num-steps 3 \ --speculative-eagle-topk 1 \ --speculative-num-draft-tokens 4 \ --disable-shared-experts-fusion \ --dtype bfloat16 \ --tokenizer-worker-num 4 ``` 3. SGLang Model Gateway (former Router) ```shell python -m sglang_router.launch_router \ --pd-disaggregation \ --policy cache_aware \ --prefill http://:8000 8996 \ --decode http://:8001 \ --host 127.0.0.1 \ --port 6688 ``` #### Running DeepSeek with PD disaggregation on 4 x Atlas 800I A3. W8A8 Model weights could be found [here](https://modelers.cn/models/State_Cloud/Deepseek-R1-bf16-hfd-w8a8). 1. Prefill & Decode: ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 export SGLANG_SET_CPU_AFFINITY=1 unset ASCEND_LAUNCH_BLOCKING source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export PATH=/usr/local/Ascend/8.5.0/compiler/bishengir/bin:$PATH export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export ASCEND_MF_STORE_URL="tcp://your prefill ip1:24669" P_IP=('your prefill ip1' 'your prefill ip2') D_IP=('your decode ip1' 'your decode ip2') MODEL_PATH=xxx export SGLANG_NPU_USE_MLAPO=1 export SGLANG_USE_FIA_NZ=1 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" # prefill for i in "${!P_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${P_IP[$i]}" || "$LOCAL_HOST2" == "${P_IP[$i]}" ]]; then echo "${P_IP[$i]}" export HCCL_BUFFSIZE=1536 export DEEP_NORMAL_MODE_USE_INT8_QUANT=1 export TASK_QUEUE_ENABLE=2 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode prefill --host ${P_IP[$i]} \ --port 8000 --disaggregation-bootstrap-port $((8998+$i)) --trust-remote-code --nnodes 1 --node-rank 0 \ --tp-size 16 --mem-fraction-static 0.81 --attention-backend ascend --device npu --quantization modelslim \ --disaggregation-transfer-backend ascend --max-running-requests 8 --context-length 8192 --disable-radix-cache \ --chunked-prefill-size -1 --max-prefill-tokens 28680 --moe-a2a-backend deepep --deepep-mode normal \ --speculative-algorithm NEXTN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \ --dp-size 2 --enable-dp-attention --disable-shared-experts-fusion --dtype bfloat16 --enable-attn-tp-input-scattered NODE_RANK=$i break fi done # decode for i in "${!D_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${D_IP[$i]}" || "$LOCAL_HOST2" == "${D_IP[$i]}" ]]; then echo "${D_IP[$i]}" export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 export HCCL_BUFFSIZE=650 export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=78 export TASK_QUEUE_ENABLE=1 export SGLANG_SCHEDULER_SKIP_ALL_GATHER=1 export HCCL_SOCKET_IFNAME=xxx export GLOO_SOCKET_IFNAME=xxx python -m sglang.launch_server --model-path ${MODEL_PATH} --disaggregation-mode decode --host ${D_IP[$i]} \ --port 8001 --trust-remote-code --dist-init-addr ${D_IP[0]}:5000 --nnodes 2 --node-rank $i --tp-size 32 --dp-size 32 \ --mem-fraction-static 0.815 --max-running-requests 832 --attention-backend ascend --device npu --quantization modelslim \ --moe-a2a-backend deepep --enable-dp-attention --deepep-mode low_latency --enable-dp-lm-head --moe-dense-tp 1 \ --cuda-graph-bs 12 14 16 18 20 22 24 26 --disaggregation-transfer-backend ascend --watchdog-timeout 9000 --context-length 8192 \ --speculative-algorithm NEXTN --speculative-num-steps 2 --speculative-eagle-topk 1 --speculative-num-draft-tokens 3 \ --tokenizer-worker-num 4 --prefill-round-robin-balance --disable-shared-experts-fusion --dtype bfloat16 \ --load-balance-method decode_round_robin NODE_RANK=$i break fi done ``` 2. SGLang Model Gateway (former Router): ```shell export SGLANG_DP_ROUND_ROBIN=1 python -m sglang_router.launch_router \ --pd-disaggregation \ --policy cache_aware \ --prefill http://P_IP:8000 8998 \ --prefill http://P_IP:8000 8999 \ --decode http://D_IP:8001 \ --host 127.0.0.1 \ --port 6688 \ --mini-lb ``` #### test gsm8k ```python from types import SimpleNamespace from sglang.test.few_shot_gsm8k import run_eval def gsm8k(): args = SimpleNamespace( num_shots=5, data_path=None, num_questions=200, max_new_tokens=512, parallel=32, host=f"http://127.0.0.1", port=6688, ) metrics = run_eval(args) print(f"{metrics=}") print(f"{metrics['accuracy']=}") if __name__ == "__main__": gsm8k() ``` ================================================ FILE: docs/platforms/ascend_npu_environment_variables.md ================================================ # Environment Variables SGLang supports various environment variables related to Ascend NPU that can be used to configure its runtime behavior. This document provides a list of commonly used environment variables and aims to stay updated over time. ## Directly Used in SGLang | Environment Variable | Description | Default Value | |--------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------| | `SGLANG_NPU_USE_MLAPO` | Adopts the `MLAPO` fusion operator in attention
preprocessing stage of the MLA model. | `false` | | `SGLANG_USE_FIA_NZ` | Reshapes KV Cache for FIA NZ format.
`SGLANG_USE_FIA_NZ` must be enabled with `SGLANG_NPU_USE_MLAPO` | `false` | | `SGLANG_NPU_USE_MULTI_STREAM` | Enable dual-stream computation of shared experts
and routing experts in DeepSeek models.
Enable dual-stream computation in DeepSeek NSA Indexer. | `false` | | `SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT` | Disable cast model weight tensor to a specific NPU
ACL format. | `false` | | `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | The maximum number of dispatched tokens on each rank. | `128` | ## Used in DeepEP Ascend | Environment Variable | Description | Default Value | |-------------------------------------------|------------------------------------------------------------------------------------------------------------------------|---------------| | `DEEPEP_NORMAL_LONG_SEQ_PER_ROUND_TOKENS` | Enable ant-moving function in dispatch stage. Indicates
the number of tokens transmitted per round on each rank. | `8192` | | `DEEPEP_NORMAL_LONG_SEQ_ROUND` | Enable ant-moving function in dispatch stage. Indicates
the number of rounds transmitted on each rank. | `1` | | `DEEPEP_NORMAL_COMBINE_ENABLE_LONG_SEQ` | Enable ant-moving function in combine stage.
The value `0` means disabled. | `0` | | `MOE_ENABLE_TOPK_NEG_ONE` | Needs to be enabled when the expert ID to be processed by
DEEPEP contains -1. | `0` | | `DEEP_NORMAL_MODE_USE_INT8_QUANT` | Quantizes x to int8 and returns (tensor, scales) in dispatch operator. | `0` | ## Others | Environment Variable | Description | Default Value | |--------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------| | `TASK_QUEUE_ENABLE` | Used to control the optimization level of the dispatch queue
about the task_queue operator. [Detail](https://www.hiascend.com/document/detail/zh/Pytorch/730/comref/Envvariables/docs/zh/environment_variable_reference/TASK_QUEUE_ENABLE.md) | `1` | | `INF_NAN_MODE_ENABLE` | Controls whether the chip uses saturation mode or INF_NAN mode. [Detail](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/800alpha001/apiref/envref/envref_07_0056.html) | `1` | | `STREAMS_PER_DEVICE` | Configures the maximum number of streams for the stream pool. [Detail](https://www.hiascend.com/document/detail/zh/Pytorch/720/comref/Envvariables/Envir_041.html) | `32` | | `PYTORCH_NPU_ALLOC_CONF` | Controls the behavior of the cache allocator.
This variable changes memory usage and may cause performance fluctuations. [Detail](https://www.hiascend.com/document/detail/zh/Pytorch/700/comref/Envvariables/Envir_012.html) | | | `ASCEND_MF_STORE_URL` | The address of config store in MemFabric during PD separation,
which is generally set to the IP address of the P primary node
with an arbitrary port number. | | | `ASCEND_LAUNCH_BLOCKING` | Controls whether synchronous mode is enabled during operator execution. [Detail](https://www.hiascend.com/document/detail/zh/Pytorch/710/comref/Envvariables/Envir_006.html) | `0` | | `HCCL_OP_EXPANSION_MODE` | Configures the expansion position for communication algorithm scheduling. [Detail](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/800alpha001/apiref/envref/envref_07_0094.html) | | | `HCCL_BUFFSIZE` | Controls the size of the buffer area for shared data between two NPUs.
The unit is MB, and the value must be greater than or equal to 1. [Detail](https://www.hiascend.com/document/detail/zh/Pytorch/60RC3/ptmoddevg/trainingmigrguide/performance_tuning_0047.html) | `200` | | `HCCL_SOCKET_IFNAME` | Configures the name of the network card used by the Host
during HCCL initialization. [Detail](https://www.hiascend.com/document/detail/zh/canncommercial/81RC1/apiref/envvar/envref_07_0075.html) | | | `GLOO_SOCKET_IFNAME` | Configures the network interface name for GLOO communication. | | ================================================ FILE: docs/platforms/ascend_npu_glm5_examples.md ================================================ # GLM-5 examples ## Introduction The GLM (General Language Model) series is an open-source bilingual large language model family jointly developed by the KEG Laboratory of Tsinghua University and Zhipu AI. This series of models has performed outstandingly in the field of Chinese NLP with its unique unified pre-training framework and bilingual capabilities. [GLM-5](https://huggingface.co/zai-org/GLM-5) adopts the DeepSeek-V3/V3.2 architecture, including the sparse attention (DSA) and multi-token prediction (MTP). Ascend supports GLM-5 with 0Day based on the SGLang inference framework, achieving low-code seamless enablement and compatibility with the mainstream distributed parallel capabilities within the current SGLang framework. We welcome developers to download and experience it. ## Environment Preparation ### Model Weight - `GLM-5.0`(BF16 version): [Download model weight](https://www.modelscope.cn/models/ZhipuAI/GLM-5). - `GLM-5.0-w4a8`(Quantized version without mtp): [Download model weight](https://modelers.cn/models/Eco-Tech/GLM-5-w4a8). - You can use [msmodelslim](https://gitcode.com/Ascend/msmodelslim) to quantify the model naively. ### Installation The dependencies required for the NPU runtime environment have been integrated into a Docker image and uploaded to the Ascend platform. You can directly pull it. ```{code-block} bash #Atlas 800 A3 docker pull swr.cn-southwest-2.myhuaweicloud.com/base_image/dockerhub/lmsysorg/sglang:cann8.5.0-a3-glm5 #Atlas 800 A2 docker pull swr.cn-southwest-2.myhuaweicloud.com/base_image/dockerhub/lmsysorg/sglang:cann8.5.0-910b-glm5 #start container docker run -itd --shm-size=16g --privileged=true --name ${NAME} \ --privileged=true --net=host \ -v /var/queue_schedule:/var/queue_schedule \ -v /etc/ascend_install.info:/etc/ascend_install.info \ -v /usr/local/sbin:/usr/local/sbin \ -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \ -v /usr/local/Ascend/firmware:/usr/local/Ascend/firmware \ --device=/dev/davinci0:/dev/davinci0 \ --device=/dev/davinci1:/dev/davinci1 \ --device=/dev/davinci2:/dev/davinci2 \ --device=/dev/davinci3:/dev/davinci3 \ --device=/dev/davinci4:/dev/davinci4 \ --device=/dev/davinci5:/dev/davinci5 \ --device=/dev/davinci6:/dev/davinci6 \ --device=/dev/davinci7:/dev/davinci7 \ --device=/dev/davinci8:/dev/davinci8 \ --device=/dev/davinci9:/dev/davinci9 \ --device=/dev/davinci10:/dev/davinci10 \ --device=/dev/davinci11:/dev/davinci11 \ --device=/dev/davinci12:/dev/davinci12 \ --device=/dev/davinci13:/dev/davinci13 \ --device=/dev/davinci14:/dev/davinci14 \ --device=/dev/davinci15:/dev/davinci15 \ --device=/dev/davinci_manager:/dev/davinci_manager \ --device=/dev/hisi_hdc:/dev/hisi_hdc \ --entrypoint=bash \ swr.cn-southwest-2.myhuaweicloud.com/base_image/dockerhub/lmsysorg/sglang:${TAG} ``` Note: Using this image, you need to update transformers to main branch ``` shell # reinstall transformers pip install git+https://github.com/huggingface/transformers.git ``` ## Deployment ### Single-node Deployment - Quantized model `glm5_w4a8` can be deployed on 1 Atlas 800 A3 (64G × 16) . Run the following script to execute online inference. ```shell # high performance cpu echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 # bind cpu export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING # cann source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export STREAMS_PER_DEVICE=32 export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 export SGLANG_ENABLE_SPEC_V2=1 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_NPU_USE_MULTI_STREAM=1 export HCCL_BUFFSIZE=1000 export HCCL_OP_EXPANSION_MODE=AIV export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo python3 -m sglang.launch_server \ --model-path $MODEL_PATH \ --attention-backend ascend \ --device npu \ --tp-size 16 --nnodes 1 --node-rank 0 \ --chunked-prefill-size 16384 --max-prefill-tokens 280000 \ --trust-remote-code \ --host 127.0.0.1 \ --mem-fraction-static 0.7 \ --port 8000 \ --served-model-name glm-5 \ --cuda-graph-bs 16 \ --quantization modelslim \ --moe-a2a-backend deepep --deepep-mode auto ``` ### Multi-node Deployment - `GLM-5-bf16`: require at least 2 Atlas 800 A3 (64G × 16). **A3 series** Modify the IP of 2 nodes, then run the same scripts on two nodes. **node 0/1** ```shell echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 # bind cpu export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING # cann source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export STREAMS_PER_DEVICE=32 export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 export SGLANG_ENABLE_SPEC_V2=1 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_NPU_USE_MULTI_STREAM=1 export HCCL_BUFFSIZE=1000 export HCCL_OP_EXPANSION_MODE=AIV # Run command ifconfig on two nodes, find out which inet addr has same IP with your node IP. That is your public interface, which should be added here export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo P_IP=('your ip1' 'your ip2') P_MASTER="${P_IP[0]}:your port" export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600 export SGLANG_ENABLE_SPEC_V2=1 export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 LOCAL_HOST1=`hostname -I|awk -F " " '{print$1}'` LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` for i in "${!P_IP[@]}"; do if [[ "$LOCAL_HOST1" == "${P_IP[$i]}" || "$LOCAL_HOST2" == "${P_IP[$i]}" ]]; then echo "${P_IP[$i]}" python3 -m sglang.launch_server \ --model-path $MODEL_PATH \ --attention-backend ascend \ --device npu \ --tp-size 32 --nnodes 2 --node-rank $i --dist-init-addr $P_MASTER \ --chunked-prefill-size 16384 --max-prefill-tokens 131072 \ --trust-remote-code \ --host 127.0.0.1 \ --mem-fraction-static 0.8\ --port 8000 \ --served-model-name glm-5 \ --cuda-graph-max-bs 16 \ --disable-radix-cache NODE_RANK=$i break fi done ``` ### Prefill-Decode Disaggregation Not test yet. ### Using Benchmark Refer to [Benchmark and Profiling](../developer_guide/benchmark_and_profiling.md) for details. ================================================ FILE: docs/platforms/ascend_npu_quantization.md ================================================ Quantization on Ascend. To load already quantized models, simply load the model weights and config. Again, if the model has been quantized offline, there's no need to add `--quantization` argument when starting the engine. The quantization method will be automatically parsed from the downloaded `quant_model_description.json` or `config.json` config. [ModelSlim on Ascend support](https://github.com/sgl-project/sglang/pull/14504): - [x] W4A4 dynamic linear - [x] W8A8 static linear - [x] W8A8 dynamic linear - [x] W4A4 dynamic MOE - [x] W4A8 dynamic MOE - [x] W8A8 dynamic MOE [AWQ on Ascend support](https://github.com/sgl-project/sglang/pull/10158): - [x] W4A16 linear - [x] W8A16 linear # Need to test - [x] W4A16 MOE # Need to test Compressed-tensors (LLM Compressor) on Ascend support: - [x] [W4A8 dynamic MOE with/without activation clip](https://github.com/sgl-project/sglang/pull/14736) # Need to test - [x] [W4A16 MOE](https://github.com/sgl-project/sglang/pull/12759) - [x] [W8A8 dynamic linear](https://github.com/sgl-project/sglang/pull/14504) - [x] [W8A8 dynamic MOE](https://github.com/sgl-project/sglang/pull/14504) Diffusion model [modelslim](https://github.com/sgl-project/sglang/pull/17996) quantization on Ascend support: - [x] W4A4 dynamic linear - [x] W8A8 static linear - [x] W8A8 dynamic linear ================================================ FILE: docs/platforms/ascend_npu_qwen3_5_examples.md ================================================ # Qwen3.5 examples ## Environment Preparation ### Installation The dependencies required for the NPU runtime environment have been integrated into a Docker image and uploaded to the quay.io platform. You can directly pull it. ```{code-block} bash #Atlas 800 A3 docker pull quay.io/ascend/sglang:main-cann8.5.0-a3 #Atlas 800 A2 docker pull quay.io/ascend/sglang:main-cann8.5.0-910b #start container docker run -itd --shm-size=16g --privileged=true --name ${NAME} \ --privileged=true --net=host \ -v /var/queue_schedule:/var/queue_schedule \ -v /etc/ascend_install.info:/etc/ascend_install.info \ -v /usr/local/sbin:/usr/local/sbin \ -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \ -v /usr/local/Ascend/firmware:/usr/local/Ascend/firmware \ --device=/dev/davinci0:/dev/davinci0 \ --device=/dev/davinci1:/dev/davinci1 \ --device=/dev/davinci2:/dev/davinci2 \ --device=/dev/davinci3:/dev/davinci3 \ --device=/dev/davinci4:/dev/davinci4 \ --device=/dev/davinci5:/dev/davinci5 \ --device=/dev/davinci6:/dev/davinci6 \ --device=/dev/davinci7:/dev/davinci7 \ --device=/dev/davinci8:/dev/davinci8 \ --device=/dev/davinci9:/dev/davinci9 \ --device=/dev/davinci10:/dev/davinci10 \ --device=/dev/davinci11:/dev/davinci11 \ --device=/dev/davinci12:/dev/davinci12 \ --device=/dev/davinci13:/dev/davinci13 \ --device=/dev/davinci14:/dev/davinci14 \ --device=/dev/davinci15:/dev/davinci15 \ --device=/dev/davinci_manager:/dev/davinci_manager \ --device=/dev/hisi_hdc:/dev/hisi_hdc \ --entrypoint=bash \ quay.io/ascend/sglang:${tag} ``` ## Deployment ### Single-node Deployment Run the following script to execute online inference. #### Qwen3.5 397B ```shell # high performance cpu echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 # bind cpu export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING # cann source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export STREAMS_PER_DEVICE=32 export HCCL_BUFFSIZE=1000 export HCCL_OP_EXPANSION_MODE=AIV export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo python3 -m sglang.launch_server \ --model-path $MODEL_PATH \ --attention-backend ascend \ --device npu \ --tp-size 16 --nnodes 1 --node-rank 0 \ --chunked-prefill-size 4096 --max-prefill-tokens 280000 \ --disable-radix-cache \ --trust-remote-code \ --host 127.0.0.1 \ --mem-fraction-static 0.7 \ --port 8000 \ --cuda-graph-bs 16 \ --quantization modelslim \ --enable-multimodal \ --mm-attention-backend ascend_attn \ --dtype bfloat16 ``` #### Qwen3.5 122B ```shell # high performance cpu echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 # bind cpu export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING # cann source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export STREAMS_PER_DEVICE=32 export HCCL_BUFFSIZE=1000 export HCCL_OP_EXPANSION_MODE=AIV export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo python3 -m sglang.launch_server \ --model-path $MODEL_PATH \ --attention-backend ascend \ --device npu \ --tp-size 8 --nnodes 1 --node-rank 0 \ --chunked-prefill-size 4096 --max-prefill-tokens 280000 \ --disable-radix-cache \ --trust-remote-code \ --host 127.0.0.1 \ --mem-fraction-static 0.7 \ --port 8000 \ --cuda-graph-bs 16 \ --quantization modelslim \ --enable-multimodal \ --mm-attention-backend ascend_attn \ --dtype bfloat16 ``` #### Qwen3.5 35B ```shell # high performance cpu echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 # bind cpu export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING # cann source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export STREAMS_PER_DEVICE=32 export HCCL_BUFFSIZE=1000 export HCCL_OP_EXPANSION_MODE=AIV export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo python3 -m sglang.launch_server \ --model-path $MODEL_PATH \ --attention-backend ascend \ --device npu \ --tp-size 2 --nnodes 1 --node-rank 0 \ --chunked-prefill-size 4096 --max-prefill-tokens 280000 \ --disable-radix-cache \ --trust-remote-code \ --host 127.0.0.1 \ --mem-fraction-static 0.7 \ --port 8000 \ --cuda-graph-bs 16 \ --quantization modelslim \ --enable-multimodal \ --mm-attention-backend ascend_attn \ --dtype bfloat16 ``` #### Qwen3.5 27B ```shell # high performance cpu echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor sysctl -w vm.swappiness=0 sysctl -w kernel.numa_balancing=0 sysctl -w kernel.sched_migration_cost_ns=50000 # bind cpu export SGLANG_SET_CPU_AFFINITY=1 unset https_proxy unset http_proxy unset HTTPS_PROXY unset HTTP_PROXY unset ASCEND_LAUNCH_BLOCKING # cann source /usr/local/Ascend/ascend-toolkit/set_env.sh source /usr/local/Ascend/nnal/atb/set_env.sh export STREAMS_PER_DEVICE=32 export HCCL_BUFFSIZE=1000 export HCCL_OP_EXPANSION_MODE=AIV export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo python3 -m sglang.launch_server \ --model-path $MODEL_PATH \ --attention-backend ascend \ --device npu \ --tp-size 2 \ --chunked-prefill-size -1 --max-prefill-tokens 120000 \ --disable-radix-cache \ --trust-remote-code \ --host 127.0.0.1 \ --mem-fraction-static 0.8 \ --port 8000 \ --cuda-graph-bs 32 \ --enable-multimodal \ --mm-attention-backend ascend_attn ``` ### Prefill-Decode Disaggregation Not test yet. ### Using Benchmark Refer to [Benchmark and Profiling](../developer_guide/benchmark_and_profiling.md) for details. ================================================ FILE: docs/platforms/ascend_npu_qwen3_examples.md ================================================ ## Qwen3 examples ### Running Qwen3 #### Running Qwen3-32B on 1 x Atlas 800I A3. Model weights could be found [here](https://huggingface.co/Qwen/Qwen3-32B) ```shell export SGLANG_SET_CPU_AFFINITY=1 export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export HCCL_BUFFSIZE=1536 export HCCL_OP_EXPANSION_MODE=AIV python -m sglang.launch_server \ --device npu \ --attention-backend ascend \ --trust-remote-code \ --tp-size 4 \ --model-path Qwen/Qwen3-32B \ --mem-fraction-static 0.8 ``` #### Running Qwen3-32B on 1 x Atlas 800I A3 with Qwen3-32B-Eagle3. Model weights could be found [here](https://huggingface.co/Qwen/Qwen3-32B) Speculative model weights could be found [here](https://huggingface.co/Zhihu-ai/Zhi-Create-Qwen3-32B-Eagle3) ```shell export SGLANG_SET_CPU_AFFINITY=1 export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export HCCL_OP_EXPANSION_MODE=AIV export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 export SGLANG_ENABLE_SPEC_V2=1 python -m sglang.launch_server \ --device npu \ --attention-backend ascend \ --trust-remote-code \ --tp-size 4 \ --model-path Qwen/Qwen3-32B \ --mem-fraction-static 0.8 \ --speculative-algorithm EAGLE3 \ --speculative-draft-model-path Qwen/Qwen3-32B-Eagle3 \ --speculative-num-steps 1 \ --speculative-eagle-topk 1 \ --speculative-num-draft-tokens 2 ``` #### Running Qwen3-30B-A3B MOE on 1 x Atlas 800I A3. Model weights could be found [here](https://huggingface.co/Qwen/Qwen3-30B-A3B) ```shell export SGLANG_SET_CPU_AFFINITY=1 export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export HCCL_BUFFSIZE=1536 export HCCL_OP_EXPANSION_MODE=AIV export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=32 export SGLANG_DEEPEP_BF16_DISPATCH=1 python -m sglang.launch_server \ --device npu \ --attention-backend ascend \ --trust-remote-code \ --tp-size 4 \ --model-path Qwen/Qwen3-30B-A3B \ --mem-fraction-static 0.8 ``` #### Running Qwen3-235B-A22B-Instruct-2507 MOE on 1 x Atlas 800I A3. Model weights could be found [here](https://huggingface.co/Qwen/Qwen3-235B-A22B-Instruct-2507) ```shell export SGLANG_SET_CPU_AFFINITY=1 export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export HCCL_BUFFSIZE=1536 export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=32 export SGLANG_DEEPEP_BF16_DISPATCH=1 python -m sglang.launch_server \ --model-path Qwen/Qwen3-235B-A22B-Instruct-2507 \ --tp-size 16 \ --trust-remote-code \ --attention-backend ascend \ --device npu \ --watchdog-timeout 9000 \ --mem-fraction-static 0.8 ``` #### Running Qwen3-VL-8B-Instruct on 1 x Atlas 800I A3. Model weights could be found [here](https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct) ```shell export SGLANG_SET_CPU_AFFINITY=1 export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export STREAMS_PER_DEVICE=32 export HCCL_BUFFSIZE=1536 export HCCL_OP_EXPANSION_MODE=AIV python -m sglang.launch_server \ --enable-multimodal \ --attention-backend ascend \ --mm-attention-backend ascend_attn \ --trust-remote-code \ --tp-size 4 \ --model-path Qwen/Qwen3-VL-8B-Instruct \ --mem-fraction-static 0.8 ``` ================================================ FILE: docs/platforms/ascend_npu_support.rst ================================================ Ascend NPUs =============================================================== .. toctree:: :maxdepth: 1 ascend_npu.md ascend_npu_support_features.md ascend_npu_support_models.md ascend_npu_deepseek_example.md ascend_npu_qwen3_examples.md mindspore_backend.md ascend_contribution_guide.md ascend_npu_best_practice.md ascend_npu_qwen3_5_examples.md ascend_npu_glm5_examples.md ascend_npu_environment_variables.md ================================================ FILE: docs/platforms/ascend_npu_support_features.md ================================================ # Support Features on Ascend NPU This section describes the basic functions and features supported by the Ascend NPU.If you encounter issues or have any questions, please [open an issue](https://github.com/sgl-project/sglang/issues). If you want to know the meaning and usage of each parameter, click [Server Arguments](https://docs.sglang.io/advanced_features/server_arguments.html). ## Model and tokenizer | Argument | Defaults | Options | Server supported | |----------------------------------------|----------|---------------------------------------|:----------------:| | `--model-path`
`--model` | `None` | Type: str | A2, A3 | | `--tokenizer-path` | `None` | Type: str | A2, A3 | | `--tokenizer-mode` | `auto` | `auto`, `slow` | A2, A3 | | `--tokenizer-worker-num` | `1` | Type: int | A2, A3 | | `--skip-tokenizer-init` | `False` | bool flag (set to enable) | A2, A3 | | `--load-format` | `auto` | `auto`, `safetensors` | A2, A3 | | `--model-loader-`
`extra-config` | `{}` | Type: str | A2, A3 | | `--trust-remote-code` | `False` | bool flag (set to enable) | A2, A3 | | `--context-length` | `None` | Type: int | A2, A3 | | `--is-embedding` | `False` | bool flag (set to enable) | A2, A3 | | `--enable-multimodal` | `None` | bool flag (set to enable) | A2, A3 | | `--revision` | `None` | Type: str | A2, A3 | | `--model-impl` | `auto` | `auto`, `sglang`,
`transformers` | A2, A3 | ## HTTP server | Argument | Defaults | Options | Server supported | |------------------------|-------------|---------------------------|:----------------:| | `--host` | `127.0.0.1` | Type: str | A2, A3 | | `--port` | `30000` | Type: int | A2, A3 | | `--skip-server-warmup` | `False` | bool flag (set to enable) | A2, A3 | | `--warmups` | `None` | Type: str | A2, A3 | | `--nccl-port` | `None` | Type: int | A2, A3 | | `--fastapi-root-path` | `None` | Type: str | A2, A3 | | `--grpc-mode` | `False` | bool flag (set to enable) | A2, A3 | ## Quantization and data type | Argument | Defaults | Options | Server supported | |---------------------------------------------|----------|-----------------------------------------|:----------------:| | `--dtype` | `auto` | `auto`,
`float16`,
`bfloat16` | A2, A3 | | `--quantization` | `None` | `modelslim` | A2, A3 | | `--quantization-param-path` | `None` | Type: str | Special For GPU | | `--kv-cache-dtype` | `auto` | `auto` | A2, A3 | | `--enable-fp32-lm-head` | `False` | bool flag
(set to enable) | A2, A3 | | `--modelopt-quant` | `None` | Type: str | Special For GPU | | `--modelopt-checkpoint-`
`restore-path` | `None` | Type: str | Special For GPU | | `--modelopt-checkpoint-`
`save-path` | `None` | Type: str | Special For GPU | | `--modelopt-export-path` | `None` | Type: str | Special For GPU | | `--quantize-and-serve` | `False` | bool flag
(set to enable) | Special For GPU | | `--rl-quant-profile` | `None` | Type: str | Special For GPU | ## Memory and scheduling | Argument | Defaults | Options | Server supported | |-----------------------------------------------------|----------|--------------------------------|:----------------:| | `--mem-fraction-static` | `None` | Type: float | A2, A3 | | `--max-running-requests` | `None` | Type: int | A2, A3 | | `--prefill-max-requests` | `None` | Type: int | A2, A3 | | `--max-queued-requests` | `None` | Type: int | A2, A3 | | `--max-total-tokens` | `None` | Type: int | A2, A3 | | `--chunked-prefill-size` | `None` | Type: int | A2, A3 | | `--max-prefill-tokens` | `16384` | Type: int | A2, A3 | | `--schedule-policy` | `fcfs` | `lpm`, `fcfs` | A2, A3 | | `--enable-priority-`
`scheduling` | `False` | bool flag
(set to enable) | A2, A3 | | `--schedule-low-priority-`
`values-first` | `False` | bool flag
(set to enable) | A2, A3 | | `--priority-scheduling-`
`preemption-threshold` | `10` | Type: int | A2, A3 | | `--schedule-conservativeness` | `1.0` | Type: float | A2, A3 | | `--page-size` | `128` | Type: int | A2, A3 | | `--swa-full-tokens-ratio` | `0.8` | Type: float | A2, A3 | | `--disable-hybrid-swa-memory` | `False` | bool flag
(set to enable) | A2, A3 | | `--radix-eviction-policy` | `lru` | `lru`,
`lfu` | A2, A3 | | `--abort-on-priority-`
`when-disabled` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-dynamic-chunking` | `False` | bool flag
(set to enable) | A2, A3 | ## Runtime options | Argument | Defaults | Options | Server supported | |----------------------------------------------------|----------|---------------------------|:----------------:| | `--device` | `None` | Type: str | A2, A3 | | `--tensor-parallel-size`
`--tp-size` | `1` | Type: int | A2, A3 | | `--pipeline-parallel-size`
`--pp-size` | `1` | Type: int | A2, A3 | | `--pp-max-micro-batch-size` | `None` | Type: int | A2, A3 | | `--pp-async-batch-depth` | `None` | Type: int | A2, A3 | | `--stream-interval` | `1` | Type: int | A2, A3 | | `--incremental-streaming-output` | `False` | bool flag (set to enable) | A2, A3 | | `--random-seed` | `None` | Type: int | A2, A3 | | `--constrained-json-`
`whitespace-pattern` | `None` | Type: str | A2, A3 | | `--constrained-json-`
`disable-any-whitespace` | `False` | bool flag (set to enable) | A2, A3 | | `--watchdog-timeout` | `300` | Type: float | A2, A3 | | `--soft-watchdog-timeout` | `300` | Type: float | A2, A3 | | `--dist-timeout` | `None` | Type: int | A2, A3 | | `--base-gpu-id` | `0` | Type: int | A2, A3 | | `--gpu-id-step` | `1` | Type: int | A2, A3 | | `--sleep-on-idle` | `False` | bool flag (set to enable) | A2, A3 | | `--custom-sigquit-handler` | `None` | Optional[Callable] | A2, A3 | ## Logging | Argument | Defaults | Options | Server supported | |----------------------------------------------------|-------------------|--------------------------------|:----------------:| | `--log-level` | `info` | Type: str | A2, A3 | | `--log-level-http` | `None` | Type: str | A2, A3 | | `--log-requests` | `False` | bool flag
(set to enable) | A2, A3 | | `--log-requests-level` | `2` | `0`, `1`, `2`, `3` | A2, A3 | | `--log-requests-format` | `text` | `text`, `json` | A2, A3 | | `--crash-dump-folder` | `None` | Type: str | A2, A3 | | `--enable-metrics` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-metrics-for-`
`all-schedulers` | `False` | bool flag
(set to enable) | A2, A3 | | `--tokenizer-metrics-`
`custom-labels-header` | `x-custom-labels` | Type: str | A2, A3 | | `--tokenizer-metrics-`
`allowed-custom-labels` | `None` | List[str] | A2, A3 | | `--bucket-time-to-`
`first-token` | `None` | List[float] | A2, A3 | | `--bucket-inter-token-`
`latency` | `None` | List[float] | A2, A3 | | `--bucket-e2e-request-`
`latency` | `None` | List[float] | A2, A3 | | `--collect-tokens-`
`histogram` | `False` | bool flag
(set to enable) | A2, A3 | | `--prompt-tokens-buckets` | `None` | List[str] | A2, A3 | | `--generation-tokens-buckets` | `None` | List[str] | A2, A3 | | `--gc-warning-threshold-secs` | `0.0` | Type: float | A2, A3 | | `--decode-log-interval` | `40` | Type: int | A2, A3 | | `--enable-request-time-`
`stats-logging` | `False` | bool flag
(set to enable) | A2, A3 | | `--kv-events-config` | `None` | Type: str | Special for GPU | | `--enable-trace` | `False` | bool flag
(set to enable) | A2, A3 | | `--oltp-traces-endpoint` | `localhost:4317` | Type: str | A2, A3 | | `--log-requests-target` | `None` | Type: str | A2, A3 | | `--uvicorn-access-log-exclude-prefixes` | `[]` | List[str] | A2, A3 | ## RequestMetricsExporter configuration | Argument | Defaults | Options | Server supported | |---------------------------------------|----------|--------------------------------|:----------------:| | `--export-metrics-to-`
`file` | `False` | bool flag
(set to enable) | A2, A3 | | `--export-metrics-to-`
`file-dir` | `None` | Type: str | A2, A3 | ## API related | Argument | Defaults | Options | Server supported | |-------------------------|-----------|--------------------------------|:----------------:| | `--api-key` | `None` | Type: str | A2, A3 | | `--admin-api-key` | `None` | Type: str | A2, A3 | | `--served-model-name` | `None` | Type: str | A2, A3 | | `--weight-version` | `default` | Type: str | A2, A3 | | `--chat-template` | `None` | Type: str | A2, A3 | | `--completion-template` | `None` | Type: str | A2, A3 | | `--enable-cache-report` | `False` | bool flag
(set to enable) | A2, A3 | | `--reasoning-parser` | `None` | `deepseek-r1`
`deepseek-v3`
`glm45`
`gpt-oss`
`kimi`
`qwen3`
`qwen3-thinking`
`step3` | A2, A3 | | `--tool-call-parser` | `None` | `deepseekv3`
`deepseekv31`
`glm`
`glm45`
`glm47`
`gpt-oss`
`kimi_k2`
`llama3`
`mistral`
`pythonic`
`qwen`
`qwen25`
`qwen3_coder`
`step3`
`gigachat3` | A2, A3 | | `--sampling-defaults` | `model` | `openai`, `model` | A2, A3 | ## Data parallelism | Argument | Defaults | Options | Server supported | |----------------------------------------|---------------|-----------------------------------------------------------|:----------------:| | `--data-parallel-size`
`--dp-size` | `1` | Type: int | A2, A3 | | `--load-balance-method` | `auto` | `auto`,
`round_robin`,
`follow_bootstrap_room`,
`total_requests`,
`total_tokens` | A2, A3 | | `--prefill-round-robin-balance` | `False` | bool flag
(set to enable) | A2, A3 | ## Multi-node distributed serving | Argument | Defaults | Options | Server supported | |-------------------------------------------|----------|-----------|:----------------:| | `--dist-init-addr`
`--nccl-init-addr` | `None` | Type: str | A2, A3 | | `--nnodes` | `1` | Type: int | A2, A3 | | `--node-rank` | `0` | Type: int | A2, A3 | ## Model override args | Argument | Defaults | Options | Server supported | |--------------------------------------|----------|-----------|:----------------:| | `--json-model-override-`
`args` | `{}` | Type: str | A2, A3 | | `--preferred-sampling-`
`params` | `None` | Type: str | A2, A3 | ## LoRA | Argument | Defaults | Options | Server supported | |--------------------------|----------|-------------------------------------|:----------------:| | `--enable-lora` | `False` | Bool flag
(set to enable) | A2, A3 | | `--max-lora-rank` | `None` | Type: int | A2, A3 | | `--lora-target-modules` | `None` | `all` | A2, A3 | | `--lora-paths` | `None` | Type: List[str] /
JSON objects | A2, A3 | | `--max-loras-per-batch` | `8` | Type: int | A2, A3 | | `--max-loaded-loras` | `None` | Type: int | A2, A3 | | `--lora-eviction-policy` | `lru` | `lru`,
`fifo` | A2, A3 | | `--lora-backend` | `csgmv` | `triton`,
`csgmv`,
`ascend`,
`torch_native` | A2, A3 | | `--max-lora-chunk-size` | `16` | `16`, `32`,
`64`, `128` | Special for GPU | ## Kernel Backends (Attention, Sampling, Grammar, GEMM) | Argument | Defaults | Options | Server supported | |----------------------------------------|-------------------|------------------------------------------------------------------------------------------------|:----------------:| | `--attention-backend` | `None` | `ascend` | A2, A3 | | `--prefill-attention-backend` | `None` | `ascend` | A2, A3 | | `--decode-attention-backend` | `None` | `ascend` | A2, A3 | | `--sampling-backend` | `None` | `pytorch`,
`ascend` | A2, A3 | | `--grammar-backend` | `None` | `xgrammar` | A2, A3 | | `--mm-attention-backend` | `None` | `ascend_attn` | A2, A3 | | `--nsa-prefill-backend` | `flashmla_sparse` | `flashmla_sparse`,
`flashmla_decode`,
`fa3`,
`tilelang`,
`aiter` | Special for GPU | | `--nsa-decode-backend` | `fa3` | `flashmla_prefill`,
`flashmla_kv`,
`fa3`,
`tilelang`,
`aiter` | Special for GPU | | `--fp8-gemm-backend` | `auto` | `auto`,
`deep_gemm`,
`flashinfer_trtllm`,
`flashinfer_cutlass`,
`flashinfer_deepgemm`,
`cutlass`,
`triton`,
`aiter` | Special for GPU | | `--disable-flashinfer-`
`autotune` | `False` | bool flag
(set to enable) | Special for GPU | ## Speculative decoding | Argument | Defaults | Options | Server supported | |------------------------------------------------------------------|-----------|--------------------------|:----------------:| | `--speculative-algorithm` | `None` | `EAGLE3`,
`NEXTN` | A2, A3 | | `--speculative-draft-model-path`
`--speculative-draft-model` | `None` | Type: str | A2, A3 | | `--speculative-draft-model-`
`revision` | `None` | Type: str | A2, A3 | | `--speculative-draft-load-format` | `None` | `auto` | A2, A3 | | `--speculative-num-steps` | `None` | Type: int | A2, A3 | | `--speculative-eagle-topk` | `None` | Type: int | A2, A3 | | `--speculative-num-draft-tokens` | `None` | Type: int | A2, A3 | | `--speculative-accept-`
`threshold-single` | `1.0` | Type: float | Special for GPU | | `--speculative-accept-`
`threshold-acc` | `1.0` | Type: float | Special for GPU | | `--speculative-token-map` | `None` | Type: str | A2, A3 | | `--speculative-attention-`
`mode` | `prefill` | `prefill`,
`decode` | A2, A3 | | `--speculative-moe-runner-`
`backend` | `None` | `auto` | A2, A3 | | `--speculative-moe-a2a-`
`backend` | `None` | `ascend_fuseep` | A2, A3 | | `--speculative-draft-attention-backend` | `None` | `ascend` | A2, A3 | | `--speculative-draft-model-quantization` | `None` | `unquant` | A2, A3 | ## Ngram speculative decoding | Argument | Defaults | Options | Server supported | |----------------------------------------------------|------------|--------------------|:----------------:| | `--speculative-ngram-`
`min-match-window-size` | `1` | Type: int | Experimental | | `--speculative-ngram-`
`max-match-window-size` | `12` | Type: int | Experimental | | `--speculative-ngram-`
`min-bfs-breadth` | `1` | Type: int | Experimental | | `--speculative-ngram-`
`max-bfs-breadth` | `10` | Type: int | Experimental | | `--speculative-ngram-`
`match-type` | `BFS` | `BFS`,
`PROB` | Experimental | | `--speculative-ngram-`
`branch-length` | `18` | Type: int | Experimental | | `--speculative-ngram-`
`capacity` | `10000000` | Type: int | Experimental | ## Expert parallelism | Argument | Defaults | Options | Server supported | |-------------------------------------------------------|-----------|---------------------------------------------|:----------------:| | `--expert-parallel-size`
`--ep-size`
`--ep` | `1` | Type: int | A2, A3 | | `--moe-a2a-backend` | `none` | `none`,
`deepep`,
`ascend_fuseep` | A2, A3 | | `--moe-runner-backend` | `auto` | `auto`, `triton` | A2, A3 | | `--flashinfer-mxfp4-`
`moe-precision` | `default` | `default`,
`bf16` | Special for GPU | | `--enable-flashinfer-`
`allreduce-fusion` | `False` | bool flag
(set to enable) | Special for GPU | | `--deepep-mode` | `auto` | `normal`,
`low_latency`,
`auto` | A2, A3 | | `--deepep-config` | `None` | Type: str | Special for GPU | | `--ep-num-redundant-experts` | `0` | Type: int | A2, A3 | | `--ep-dispatch-algorithm` | `None` | Type: str | A2, A3 | | `--init-expert-location` | `trivial` | Type: str | A2, A3 | | `--enable-eplb` | `False` | bool flag
(set to enable) | A2, A3 | | `--eplb-algorithm` | `auto` | Type: str | A2, A3 | | `--eplb-rebalance-layers-`
`per-chunk` | `None` | Type: int | A2, A3 | | `--eplb-min-rebalancing-`
`utilization-threshold` | `1.0` | Type: float | A2, A3 | | `--expert-distribution-`
`recorder-mode` | `None` | Type: str | A2, A3 | | `--expert-distribution-`
`recorder-buffer-size` | `None` | Type: int | A2, A3 | | `--enable-expert-distribution-`
`metrics` | `False` | bool flag (set to enable) | A2, A3 | | `--moe-dense-tp-size` | `None` | Type: int | A2, A3 | | `--elastic-ep-backend` | `None` | `none`, `mooncake` | Special for GPU | | `--mooncake-ib-device` | `None` | Type: str | Special for GPU | ## Mamba Cache | Argument | Defaults | Options | Server supported | |------------------------------|-----------|-----------------------------------------------|:----------------:| | `--max-mamba-cache-size` | `None` | Type: int | A2, A3 | | `--mamba-ssm-dtype` | `float32` | `float32`,
`bfloat16`,
`float16` | A2, A3 | | `--mamba-full-memory-ratio` | `0.9` | Type: float | A2, A3 | | `--mamba-scheduler-strategy` | `auto` | Only `auto`, `no_buffer` supported | A2, A3 | | `--mamba-track-interval` | `256` | Type: int | A2, A3 | ## Hierarchical cache | Argument | Defaults | Options | Server supported | |-------------------------------------------------|-----------------|---------------------------------------------------------------------|:----------------:| | `--enable-hierarchical-`
`cache` | `False` | bool flag
(set to enable) | A2, A3 | | `--hicache-ratio` | `2.0` | Type: float | A2, A3 | | `--hicache-size` | `0` | Type: int | A2, A3 | | `--hicache-write-policy` | `write_through` | Currently only `write_back` supported | A2, A3 | | `--hicache-io-backend` | `kernel` | `kernel_ascend`,
`direct` | A2, A3 | | `--hicache-mem-layout` | `layer_first` | `page_first_direct`,
`page_first_kv_split` | A2, A3 | | `--hicache-storage-`
`backend` | `None` | `file` | A2, A3 | | `--hicache-storage-`
`prefetch-policy` | `best_effort` | `best_effort`,
`wait_complete`,
`timeout` | Special for GPU | | `--hicache-storage-`
`backend-extra-config` | `None` | Type: str | Special for GPU | ## LMCache | Argument | Defaults | Options | Server supported | |--------------------|----------|--------------------------------|:----------------:| | `--enable-lmcache` | `False` | bool flag
(set to enable) | Special for GPU | ## Offloading | Argument | Defaults | Options | Server supported | |---------------------------|----------|-----------|:----------------:| | `--cpu-offload-gb` | `0` | Type: int | A2, A3 | | `--offload-group-size` | `-1` | Type: int | Planned | | `--offload-num-in-group` | `1` | Type: int | Planned | | `--offload-prefetch-step` | `1` | Type: int | Planned | | `--offload-mode` | `cpu` | Type: str | Planned | ## Args for multi-item scoring | Argument | Defaults | Options | Server supported | |----------------------------------|----------|-----------|:----------------:| | `--multi-item-scoring-delimiter` | `None` | Type: int | A2, A3 | ## Optimization/debug options | Argument | Defaults | Options | Server supported | |---------------------------------------------------------|----------|--------------------------------|:----------------:| | `--disable-radix-cache` | `False` | bool flag
(set to enable) | A2, A3 | | `--cuda-graph-max-bs` | `None` | Type: int | A2, A3 | | `--cuda-graph-bs` | `None` | List[int] | A2, A3 | | `--disable-cuda-graph` | `False` | bool flag
(set to enable) | A2, A3 | | `--disable-cuda-graph-`
`padding` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-profile-`
`cuda-graph` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-cudagraph-gc` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-nccl-nvls` | `False` | bool flag
(set to enable) | Special for GPU | | `--enable-symm-mem` | `False` | bool flag
(set to enable) | Special for GPU | | `--disable-flashinfer-`
`cutlass-moe-fp4-allgather` | `False` | bool flag
(set to enable) | Special for GPU | | `--enable-tokenizer-`
`batch-encode` | `False` | bool flag
(set to enable) | A2, A3 | | `--disable-tokenizer-`
`batch-decode` | `False` | bool flag
(set to enable) | A2, A3 | | `--disable-custom-`
`all-reduce` | `False` | bool flag
(set to enable) | Special for GPU | | `--enable-mscclpp` | `False` | bool flag
(set to enable) | Special for GPU | | `--enable-torch-`
`symm-mem` | `False` | bool flag
(set to enable) | Special for GPU | | `--disable-overlap`
`-schedule` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-mixed-`
`chunk` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-dp-attention` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-dp-lm-head` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-two-`
`batch-overlap` | `False` | bool flag
(set to enable) | Planned | | `--enable-single-`
`batch-overlap` | `False` | bool flag
(set to enable) | A2, A3 | | `--tbo-token-`
`distribution-threshold` | `0.48` | Type: float | Planned | | `--enable-torch-`
`compile` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-torch-`
`compile-debug-mode` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-piecewise-`
`cuda-graph` | `False` | bool flag
(set to enable) | A2, A3 | | `--piecewise-cuda-`
`graph-tokens` | `None` | Type: JSON
list | A2, A3 | | `--piecewise-cuda-`
`graph-compiler` | `eager` | ["eager", "inductor"] | A2, A3 | | `--torch-compile-max-bs` | `32` | Type: int | A2, A3 | | `--piecewise-cuda-`
`graph-max-tokens` | `None` | Type: int | A2, A3 | | `--torchao-config` | `` | Type: str | Special for GPU | | `--enable-nan-detection` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-p2p-check` | `False` | bool flag
(set to enable) | Special for GPU | | `--triton-attention-`
`reduce-in-fp32` | `False` | bool flag
(set to enable) | Special for GPU | | `--triton-attention-`
`num-kv-splits` | `8` | Type: int | Special for GPU | | `--triton-attention-`
`split-tile-size` | `None` | Type: int | Special for GPU | | `--delete-ckpt-`
`after-loading` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-memory-saver` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-weights-`
`cpu-backup` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-draft-weights-`
`cpu-backup` | `False` | bool flag
(set to enable) | A2, A3 | | `--allow-auto-truncate` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-custom-`
`logit-processor` | `False` | bool flag
(set to enable) | A2, A3 | | `--flashinfer-mla-`
`disable-ragged` | `False` | bool flag
(set to enable) | Special for GPU | | `--disable-shared-`
`experts-fusion` | `False` | bool flag
(set to enable) | A2, A3 | | `--disable-chunked-`
`prefix-cache` | `False` | bool flag
(set to enable) | A2, A3 | | `--disable-fast-`
`image-processor` | `False` | bool flag
(set to enable) | A2, A3 | | `--keep-mm-feature-`
`on-device` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-return-`
`hidden-states` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-return-`
`routed-experts` | `False` | bool flag
(set to enable) | A2, A3 | | `--scheduler-recv-`
`interval` | `1` | Type: int | A2, A3 | | `--numa-node` | `None` | List[int] | A2, A3 | | `--enable-deterministic-`
`inference` | `False` | bool flag
(set to enable) | Planned | | `--rl-on-policy-target` | `None` | `fsdp` | Planned | | `--enable-layerwise-`
`nvtx-marker` | `False` | bool flag
(set to enable) | Special for GPU | | `--enable-attn-tp-`
`input-scattered` | `False` | bool flag
(set to enable) | Experimental | | `--enable-nsa-prefill-`
`context-parallel` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-fused-qk-`
`norm-rope` | `False` | bool flag
(set to enable) | Special for GPU | ## Dynamic batch tokenizer | Argument | Defaults | Options | Server supported | |--------------------------------------------------|----------|--------------------------------|:----------------:| | `--enable-dynamic-`
`batch-tokenizer` | `False` | bool flag
(set to enable) | A2, A3 | | `--dynamic-batch-`
`tokenizer-batch-size` | `32` | Type: int | A2, A3 | | `--dynamic-batch-`
`tokenizer-batch-timeout` | `0.002` | Type: float | A2, A3 | ## Debug tensor dumps | Argument | Defaults | Options | Server supported | |--------------------------------------------|----------|-----------|:----------------:| | `--debug-tensor-dump-`
`output-folder` | `None` | Type: str | A2, A3 | | `--debug-tensor-dump-`
`layers` | `None` | List[int] | A2, A3 | | `--debug-tensor-dump-`
`input-file` | `None` | Type: str | A2, A3 | ## PD disaggregation | Argument | Defaults | Options | Server supported | |---------------------------------------------------------|------------|---------------------------------------|:----------------:| | `--disaggregation-mode` | `null` | `null`,
`prefill`,
`decode` | A2, A3 | | `--disaggregation-transfer-backend` | `mooncake` | `ascend` | A2, A3 | | `--disaggregation-bootstrap-port` | `8998` | Type: int | A2, A3 | | `--disaggregation-ib-device` | `None` | Type: str | Special for GPU | | `--disaggregation-decode-`
`enable-offload-kvcache` | `False` | bool flag
(set to enable) | A2, A3 | | `--disaggregation-decode-`
`enable-fake-auto` | `False` | bool flag
(set to enable) | A2, A3 | | `--num-reserved-decode-tokens` | `512` | Type: int | A2, A3 | | `--disaggregation-decode-`
`polling-interval` | `1` | Type: int | A2, A3 | ## Encode prefill disaggregation | Argument | Defaults | Options | Server supported | |------------------------------|--------------------|----------------------------------------------------------------|:----------------:| | `--encoder-only` | `False` | bool flag
(set to enable) | A2, A3 | | `--language-only` | `False` | bool flag
(set to enable) | A2, A3 | | `--encoder-transfer-backend` | `zmq_to_scheduler` | `zmq_to_scheduler`,
`zmq_to_tokenizer`,
`mooncake` | A2, A3 | | `--encoder-urls` | `[]` | List[str] | A2, A3 | ## Custom weight loader | Argument | Defaults | Options | Server supported | |-------------------------------------------------------------------------|----------|---------------------------------|:----------------:| | `--custom-weight-loader` | `None` | List[str] | A2, A3 | | `--weight-loader-disable-`
`mmap` | `False` | bool flag
(set to enable) | A2, A3 | | `--remote-instance-weight-`
`loader-seed-instance-ip` | `None` | Type: str | A2, A3 | | `--remote-instance-weight-`
`loader-seed-instance-service-port` | `None` | Type: int | A2, A3 | | `--remote-instance-weight-`
`loader-send-weights-group-ports` | `None` | Type: JSON
list | A2, A3 | | `--remote-instance-weight-`
`loader-backend` | `nccl` | `transfer_engine`,
`nccl` | A2, A3 | | `--remote-instance-weight-`
`loader-start-seed-via-transfer-engine` | `False` | bool flag
(set to enable) | Special for GPU | ## For PD-Multiplexing | Argument | Defaults | Options | Server supported | |-----------------------|----------|--------------------------------|:----------------:| | `--enable-pdmux` | `False` | bool flag
(set to enable) | Special for GPU | | `--pdmux-config-path` | `None` | Type: str | Special for GPU | | `--sm-group-num` | `8` | Type: int | Special for GPU | ## For Multi-Modal | Argument | Defaults | Options | Server supported | |-----------------------------------------------|----------|--------------------------------|:----------------:| | `--mm-max-concurrent-calls` | `32` | Type: int | A2, A3 | | `--mm-per-request-timeout` | `10.0` | Type: float | A2, A3 | | `--enable-broadcast-mm-`
`inputs-process` | `False` | bool flag
(set to enable) | A2, A3 | | `--mm-process-config` | `None` | Type: JSON / Dict | A2, A3 | | `--mm-enable-dp-encoder` | `False` | bool flag
(set to enable) | A2, A3 | | `--limit-mm-data-per-request` | `None` | Type: JSON / Dict | A2, A3 | ## For checkpoint decryption | Argument | Defaults | Options | Server supported | |---------------------------------|----------|--------------------------------|:----------------:| | `--decrypted-config-file` | `None` | Type: str | A2, A3 | | `--decrypted-draft-config-file` | `None` | Type: str | A2, A3 | | `--enable-prefix-mm-cache` | `False` | bool flag
(set to enable) | A2, A3 | ## Forward hooks | Argument | Defaults | Options | Server supported | |-------------------|----------|-----------------|:----------------:| | `--forward-hooks` | `None` | Type: JSON list | A2, A3 | ## Configuration file support | Argument | Defaults | Options | Server supported | |------------|----------|-----------|:----------------:| | `--config` | `None` | Type: str | A2, A3 | ## Other Params The following parameters are not supported because the third-party components that depend on are not compatible with the NPU, like Ktransformer, checkpoint-engine etc. | Argument | Defaults | Options | |-------------------------------------------------------------------|-----------|---------------------------| | `--checkpoint-engine-`
`wait-weights-`
`before-ready` | `False` | bool flag (set to enable) | | `--kt-weight-path` | `None` | Type: str | | `--kt-method` | `AMXINT4` | Type: str | | `--kt-cpuinfer` | `None` | Type: int | | `--kt-threadpool-count` | 2 | Type: int | | `--kt-num-gpu-experts` | `None` | Type: int | | `--kt-max-deferred-`
`experts-per-token` | `None` | Type: int | The following parameters have some functional deficiencies on community | Argument | Defaults | Options | |---------------------------------------|----------|--------------------------------| | `--enable-double-sparsity` | `False` | bool flag
(set to enable) | | `--ds-channel-config-path` | `None` | Type: str | | `--ds-heavy-channel-num` | `32` | Type: int | | `--ds-heavy-token-num` | `256` | Type: int | | `--ds-heavy-channel-type` | `qk` | Type: str | | `--ds-sparse-decode-`
`threshold` | `4096` | Type: int | | `--tool-server` | `None` | Type: str | ================================================ FILE: docs/platforms/ascend_npu_support_models.md ================================================ # Support Models on Ascend NPU This section describes the models supported on the Ascend NPU, including Large Language Models, Multimodal Language Models, Embedding Models, Reward Models and Rerank Models. Mainstream DeepSeek/Qwen/GLM series are included. You are welcome to enable various models based on your business requirements. ## Large Language Models | Models | Model Family | A2 Supported | A3 Supported | |--------------------------------------------|--------------------------------|:----------------------------------------:|:----------------------------------------:| | DeepSeek V3/V3.1 | DeepSeek | **√** | **√** | | DeepSeek-V3.2-W8A8 | DeepSeek | **√** | **√** | | DeepSeek-R1-0528-W8A8 | DeepSeek | **√** | **√** | | DeepSeek-V2-Lite-W8A8 | DeepSeek | **√** | **√** | | Qwen/Qwen3-30B-A3B-Instruct-2507 | Qwen | **√** | **√** | | Qwen/Qwen3-32B | Qwen | **√** | **√** | | Qwen/Qwen3-0.6B | Qwen | **√** | **√** | | Qwen3-235B-A22B-W8A8 | Qwen | **√** | **√** | | Qwen/Qwen3-Next-80B-A3B-Instruct | Qwen | **√** | **√** | | Qwen3-Coder-480B-A35B-Instruct-w8a8-QuaRot | Qwen | **√** | **√** | | Qwen/Qwen2.5-7B-Instruct | Qwen | **√** | **√** | | QWQ-32B-W8A8 | Qwen | **√** | **√** | | meta-llama/Llama-4-Scout-17B-16E-Instruct | Llama | **√** | **√** | | AI-ModelScope/Llama-3.1-8B-Instruct | Llama | **√** | **√** | | LLM-Research/llama-2-7b | Llama | **√** | **√** | | LLM-Research/Llama-3.2-1B-Instruct | Llama | **√** | **√** | | mistralai/Mistral-7B-Instruct-v0.2 | Mistral | **√** | **√** | | google/gemma-3-4b-it | Gemma | **√** | **√** | | microsoft/Phi-4-multimodal-instruct | Phi | **√** | **√** | | allenai/OLMoE-1B-7B-0924 | OLMoE | **√** | **√** | | stabilityai/stablelm-2-1_6b | StableLM | **√** | **√** | | CohereForAI/c4ai-command-r-v01 | Command-R | **√** | **√** | | huihui-ai/grok-2 | Grok | **√** | **√** | | ZhipuAI/chatglm2-6b | ChatGLM | **√** | **√** | | Shanghai_AI_Laboratory/internlm2-7b | InternLM 2 | **√** | **√** | | LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct | ExaONE 3 | **√** | **√** | | xverse/XVERSE-MoE-A36B | XVERSE | **√** | **√** | | HuggingFaceTB/SmolLM-1.7B | SmolLM | **√** | **√** | | ZhipuAI/glm-4-9b-chat | GLM-4 | **√** | **√** | | XiaomiMiMo/MiMo-7B-RL | MiMo | **√** | **√** | | arcee-ai/AFM-4.5B-Base | Arcee AFM-4.5B | **√** | **√** | | Howeee/persimmon-8b-chat | Persimmon | **√** | **√** | | inclusionAI/Ling-lite | Ling | **√** | **√** | | ibm-granite/granite-3.1-8b-instruct | Granite | **√** | **√** | | ibm-granite/granite-3.0-3b-a800m-instruct | Granite MoE | **√** | **√** | | AI-ModelScope/dbrx-instruct | DBRX (Databricks) | **√** | **√** | | baichuan-inc/Baichuan2-13B-Chat | Baichuan 2 (7B, 13B) | **√** | **√** | | baidu/ERNIE-4.5-21B-A3B-PT | ERNIE-4.5 (4.5, 4.5MoE series) | **√** | **√** | | OpenBMB/MiniCPM3-4B | MiniCPM (v3, 4B) | **√** | **√** | | Kimi/Kimi-K2-Thinking | Kimi | **√** | **√** | | openai/gpt-oss-120b | GPTOSS | **√** | **√** | | allenai/OLMo-2-1124-7B-Instruct | OLMo | **√** | **√** | | minimax/MiniMax-M2 | MiniMax-M2 | **√** | **√** | | upstage/SOLAR-10.7B-Instruct-v1.0 | Solar | **√** | **√** | | bigcode/starcoder2-7b | StarCoder2 | **√** | **√** | | arcee-ai/Trinity-Mini | Trinity (Nano, Mini) | **√** | **√** | ## Multimodal Language Models | Models | Model Family (Variants) | A2 Supported | A3 Supported | |-----------------------------------------------|---------------------------|:----------------------------------------:|:----------------------------------------:| | Qwen/Qwen2.5-VL-3B-Instruct | Qwen-VL | **√** | **√** | | Qwen/Qwen2.5-VL-72B-Instruct | Qwen-VL | **√** | **√** | | Qwen/Qwen3-VL-30B-A3B-Instruct | Qwen-VL | **√** | **√** | | Qwen/Qwen3-VL-8B-Instruct | Qwen-VL | **√** | **√** | | Qwen/Qwen3-VL-4B-Instruct | Qwen-VL | **√** | **√** | | Qwen/Qwen3-VL-235B-A22B-Instruct | Qwen-VL | **√** | **√** | | deepseek-ai/deepseek-vl2 | DeepSeek-VL2 | **√** | **√** | | deepseek-ai/Janus-Pro-1B | Janus-Pro (1B, 7B) | **√** | **√** | | deepseek-ai/Janus-Pro-7B | Janus-Pro (1B, 7B) | **√** | **√** | | openbmb/MiniCPM-V-2_6 | MiniCPM-V / MiniCPM-o | **√** | **√** | | openbmb/MiniCPM-o-2_6 | MiniCPM-V / MiniCPM-o | **√** | **√** | | google/gemma-3-4b-it | Gemma 3 (Multimodal) | **√** | **√** | | mistralai/Mistral-Small-3.1-24B-Instruct-2503 | Mistral-Small-3.1-24B | **√** | **√** | | microsoft/Phi-4-multimodal-instruct | Phi-4-multimodal-instruct | **√** | **√** | | XiaomiMiMo/MiMo-VL-7B-RL | MiMo-VL (7B) | **√** | **√** | | AI-ModelScope/llava-v1.6-34b | LLaVA (v1.5 & v1.6) | **√** | **√** | | lmms-lab/llava-next-72b | LLaVA-NeXT (8B, 72B) | **√** | **√** | | lmms-lab/llava-onevision-qwen2-7b-ov | LLaVA-OneVision | **√** | **√** | | Kimi/Kimi-VL-A3B-Instruct | Kimi-VL (A3B) | **√** | **√** | | ZhipuAI/GLM-4.5V | GLM-4.5V (106B) | **√** | **√** | | LLM-Research/Llama-3.2-11B-Vision-Instruct | Llama 3.2 Vision (11B) | **√** | **√** | | rednote-hilab/dots.ocr | DotsVLM-OCR | **√** | **√** | ## Embedding Models | Models | Model Family | A2 Supported | A3 Supported | |-------------------------------------------|--------------------------|:----------------------------------------:|:----------------------------------------:| | intfloat/e5-mistral-7b-instruct | E5 (Llama/Mistral based) | **√** | **√** | | iic/gte_Qwen2-1.5B-instruct | GTE-Qwen2 | **√** | **√** | | Qwen/Qwen3-Embedding-8B | Qwen3-Embedding | **√** | **√** | | Alibaba-NLP/gme-Qwen2-VL-2B-Instruct | GME (Multimodal) | **√** | **√** | | AI-ModelScope/clip-vit-large-patch14-336 | CLIP | **√** | **√** | | BAAI/bge-large-en-v1.5 | BGE | **√** | **√** | ## Reward Models | Models | Model Family | A2 Supported | A3 Supported | |------------------------------------------------|---------------------------|------------------------------------------|:----------------------------------------:| | Skywork/Skywork-Reward-Llama-3.1-8B-v0.2 | Llama3.1 Reward | **√** | **√** | | Shanghai_AI_Laboratory/internlm2-7b-reward | InternLM 2 Reward | **√** | **√** | | Qwen/Qwen2.5-Math-RM-72B | Qwen2.5 Reward - Math | **√** | **√** | | Howeee/Qwen2.5-1.5B-apeach | Qwen2.5 Reward - Sequence | **√** | **√** | | AI-ModelScope/Skywork-Reward-Gemma-2-27B-v0.2 | Gemma 2-27B Reward | **√** | **√** | ## Rerank Models | Models | Model Family | A2 Supported | A3 Supported | |-------------------------|--------------|:----------------------------------------:|:----------------------------------------:| | BAAI/bge-reranker-v2-m3 | BGE-Reranker | **√** | **√** | ================================================ FILE: docs/platforms/cpu_server.md ================================================ # CPU Servers The document addresses how to set up the [SGLang](https://github.com/sgl-project/sglang) environment and run LLM inference on CPU servers. SGLang is enabled and optimized on the CPUs equipped with Intel® AMX® Instructions, which are 4th generation or newer Intel® Xeon® Scalable Processors. ## Optimized Model List A list of popular LLMs are optimized and run efficiently on CPU, including the most notable open-source models like Llama series, Qwen series, and DeepSeek series like DeepSeek-R1 and DeepSeek-V3.1-Terminus. | Model Name | BF16 | W8A8_INT8 | FP8 | |:---:|:---:|:---:|:---:| | DeepSeek-R1 | | [meituan/DeepSeek-R1-Channel-INT8](https://huggingface.co/meituan/DeepSeek-R1-Channel-INT8) | [deepseek-ai/DeepSeek-R1](https://huggingface.co/deepseek-ai/DeepSeek-R1) | | DeepSeek-V3.1-Terminus | | [IntervitensInc/DeepSeek-V3.1-Terminus-Channel-int8](https://huggingface.co/IntervitensInc/DeepSeek-V3.1-Terminus-Channel-int8) | [deepseek-ai/DeepSeek-V3.1-Terminus](https://huggingface.co/deepseek-ai/DeepSeek-V3.1-Terminus) | | Llama-3.2-3B | [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) | [RedHatAI/Llama-3.2-3B-quantized.w8a8](https://huggingface.co/RedHatAI/Llama-3.2-3B-Instruct-quantized.w8a8) | | | Llama-3.1-8B | [meta-llama/Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) | [RedHatAI/Meta-Llama-3.1-8B-quantized.w8a8](https://huggingface.co/RedHatAI/Meta-Llama-3.1-8B-quantized.w8a8) | | | QwQ-32B | | [RedHatAI/QwQ-32B-quantized.w8a8](https://huggingface.co/RedHatAI/QwQ-32B-quantized.w8a8) | | | DeepSeek-Distilled-Llama | | [RedHatAI/DeepSeek-R1-Distill-Llama-70B-quantized.w8a8](https://huggingface.co/RedHatAI/DeepSeek-R1-Distill-Llama-70B-quantized.w8a8) | | | Qwen3-235B | | | [Qwen/Qwen3-235B-A22B-FP8](https://huggingface.co/Qwen/Qwen3-235B-A22B-FP8) | **Note:** The model identifiers listed in the table above have been verified on 6th Gen Intel® Xeon® P-core platforms. ## Installation ### Install Using Docker It is recommended to use Docker for setting up the SGLang environment. A [Dockerfile](https://github.com/sgl-project/sglang/blob/main/docker/xeon.Dockerfile) is provided to facilitate the installation. Replace `` below with your [HuggingFace access token](https://huggingface.co/docs/hub/en/security-tokens). ```bash # Clone the SGLang repository git clone https://github.com/sgl-project/sglang.git cd sglang/docker # Build the docker image docker build -t sglang-cpu:latest -f xeon.Dockerfile . # Initiate a docker container docker run \ -it \ --privileged \ --ipc=host \ --network=host \ -v /dev/shm:/dev/shm \ -v ~/.cache/huggingface:/root/.cache/huggingface \ -p 30000:30000 \ -e "HF_TOKEN=" \ sglang-cpu:latest /bin/bash ``` ### Install From Source If you prefer to install SGLang in a bare metal environment, the setup process is as follows: Please install the required packages and libraries beforehand if they are not already present on your system. You can refer to the Ubuntu-based installation commands in [the Dockerfile](https://github.com/sgl-project/sglang/blob/main/docker/xeon.Dockerfile#L11) for guidance. 1. Install `uv` package manager, then create and activate a virtual environment: ```bash # Taking '/opt' as the example uv env folder, feel free to change it as needed cd /opt curl -LsSf https://astral.sh/uv/install.sh | sh source $HOME/.local/bin/env uv venv --python 3.12 source .venv/bin/activate ``` 2. Create a config file to direct the installation channel (a.k.a. index-url) of `torch` related packages: ```bash vim .venv/uv.toml ``` Press 'a' to enter insert mode of `vim`, paste the following content into the created file ```file [[index]] name = "torch" url = "https://download.pytorch.org/whl/cpu" [[index]] name = "torchvision" url = "https://download.pytorch.org/whl/cpu" [[index]] name = "torchaudio" url = "https://download.pytorch.org/whl/cpu" [[index]] name = "triton" url = "https://download.pytorch.org/whl/cpu" ``` Save the file (in `vim`, press 'esc' to exit insert mode, then ':x+Enter'), and set it as the default `uv` config. ```bash export UV_CONFIG_FILE=/opt/.venv/uv.toml ``` 3. Clone the `sglang` source code and build the packages ```bash # Clone the SGLang code git clone https://github.com/sgl-project/sglang.git cd sglang git checkout # Use dedicated toml file cd python cp pyproject_cpu.toml pyproject.toml # Install SGLang dependent libs, and build SGLang main package uv pip install --upgrade pip setuptools uv pip install . # Build the CPU backend kernels cd ../sgl-kernel cp pyproject_cpu.toml pyproject.toml uv pip install . ``` 4. Set the required environment variables ```bash export SGLANG_USE_CPU_ENGINE=1 # Set 'LD_LIBRARY_PATH' and 'LD_PRELOAD' to ensure the libs can be loaded by sglang processes export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu export LD_PRELOAD=${LD_PRELOAD}:/opt/.venv/lib/libiomp5.so:${LD_LIBRARY_PATH}/libtcmalloc.so.4:${LD_LIBRARY_PATH}/libtbbmalloc.so.2 ``` Notes: - Note that the environment variable `SGLANG_USE_CPU_ENGINE=1` is required to enable the SGLang service with the CPU engine. - If you encounter code compilation issues during the `sgl-kernel` building process, please check your `gcc` and `g++` versions and upgrade them if they are outdated. It is recommended to use `gcc-13` and `g++-13` as they have been verified in the official Docker container. - The system library path is typically located in one of the following directories: `~/.local/lib/`, `/usr/local/lib/`, `/usr/local/lib64/`, `/usr/lib/`, `/usr/lib64/` and `/usr/lib/x86_64-linux-gnu/`. In the above example commands, `/usr/lib/x86_64-linux-gnu` is used. Please adjust the path according to your server configuration. - It is recommended to add the following to your `~/.bashrc` file to avoid setting these variables every time you open a new terminal: ```bash source .venv/bin/activate export SGLANG_USE_CPU_ENGINE=1 export LD_LIBRARY_PATH= export LD_PRELOAD= ``` ## Launch of the Serving Engine Example command to launch SGLang serving: ```bash python -m sglang.launch_server \ --model \ --trust-remote-code \ --disable-overlap-schedule \ --device cpu \ --host 0.0.0.0 \ --tp 6 ``` Notes: 1. For running W8A8 quantized models, please add the flag `--quantization w8a8_int8`. 2. The flag `--tp 6` specifies that tensor parallelism will be applied using 6 ranks (TP6). The number of TP specified is how many TP ranks will be used during the execution. On a CPU platform, a TP rank means a sub-NUMA cluster (SNC). Usually we can get the SNC information (How many available) from the Operating System with e.g. `lscpu` command. If the specified TP rank number differs from the total SNC count, the system will automatically utilize the first `n` SNCs. Note that `n` cannot exceed the total SNC number, doing so will result in an error. `SGLANG_CPU_OMP_THREADS_BIND` allows explicit control of CPU cores for each tensor parallel (TP) rank. **example 1**: Run SGLang service with TP=6, using the first 40 cores of each SNC on a Xeon® 6980P server, which has 43-43-42 cores on the 3 SNCs of a socket, we should set: ```bash export SGLANG_CPU_OMP_THREADS_BIND="0-39|43-82|86-125|128-167|171-210|214-253" ``` This configuration is equivalent to: - rank 0: `numactl -C 0-39 -m 0` - rank 1: `numactl -C 43-82 -m 1` - rank 2: `numactl -C 86-125 -m 2` - rank 3: `numactl -C 128-167 -m 3` - rank 4: `numactl -C 171-210 -m 4` - rank 5: `numactl -C 214-253 -m 5` **example 2**: Run SGLang service with TP=2, using 96 cores cross 3 SNCs on a Xeon® 6972P server, which has 32-32-32 cores on the 3 SNCs in a socket, we should set: ```bash export SGLANG_CPU_OMP_THREADS_BIND="0-95|96-191" ``` This configuration is equivalent to: - rank 0: `numactl -C 0-95 -m 0-2` - rank 1: `numactl -C 96-191 -m 3-5` Please beware that with SGLANG_CPU_OMP_THREADS_BIND set, the available memory amounts of the ranks may not be determined in prior. You may need to set proper `--max-total-tokens` to avoid the out-of-memory error. 3. For optimizing decoding with torch.compile, please add the flag `--enable-torch-compile`. To specify the maximum batch size when using `torch.compile`, set the flag `--torch-compile-max-bs`. For example, `--enable-torch-compile --torch-compile-max-bs 4` means using `torch.compile` and setting the maximum batch size to 4. 4. A warmup step is automatically triggered when the service is started. The server is ready when you see the log `The server is fired up and ready to roll!`. ## Benchmarking with Requests You can benchmark the performance via the `bench_serving` script. Run the command in another terminal. An example command would be: ```bash python -m sglang.bench_serving \ --dataset-name random \ --random-input-len 1024 \ --random-output-len 1024 \ --num-prompts 1 \ --request-rate inf \ --random-range-ratio 1.0 ``` Detailed parameter descriptions are available via the command: ```bash python -m sglang.bench_serving -h ``` Additionally, requests can be formatted using [the OpenAI Completions API](https://docs.sglang.io/basic_usage/openai_api_completions.html) and sent via the command line (e.g., using `curl`) or through your own scripts. ## Example Usage Commands Large Language Models can range from fewer than 1 billion to several hundred billion parameters. Dense models larger than 20B are expected to run on flagship 6th Gen Intel® Xeon® processors with dual sockets and a total of 6 sub-NUMA clusters. Dense models of approximately 10B parameters or fewer, or MoE (Mixture of Experts) models with fewer than 10B activated parameters, can run on more common 4th generation or newer Intel® Xeon® processors, or utilize a single socket of the flagship 6th Gen Intel® Xeon® processors. ### Example: Running DeepSeek-V3.1-Terminus An example command to launch service of W8A8_INT8 DeepSeek-V3.1-Terminus on a Xeon® 6980P server: ```bash python -m sglang.launch_server \ --model IntervitensInc/DeepSeek-V3.1-Terminus-Channel-int8 \ --trust-remote-code \ --disable-overlap-schedule \ --device cpu \ --quantization w8a8_int8 \ --host 0.0.0.0 \ --enable-torch-compile \ --torch-compile-max-bs 4 \ --tp 6 ``` Similarly, an example command to launch service of FP8 DeepSeek-V3.1-Terminus would be: ```bash python -m sglang.launch_server \ --model deepseek-ai/DeepSeek-V3.1-Terminus \ --trust-remote-code \ --disable-overlap-schedule \ --device cpu \ --host 0.0.0.0 \ --enable-torch-compile \ --torch-compile-max-bs 4 \ --tp 6 ``` Note: Please set `--torch-compile-max-bs` to the maximum desired batch size for your deployment, which can be up to 16. The value `4` in the examples is illustrative. ### Example: Running Llama-3.2-3B An example command to launch service of Llama-3.2-3B with BF16 precision: ```bash python -m sglang.launch_server \ --model meta-llama/Llama-3.2-3B-Instruct \ --trust-remote-code \ --disable-overlap-schedule \ --device cpu \ --host 0.0.0.0 \ --enable-torch-compile \ --torch-compile-max-bs 16 \ --tp 2 ``` The example command to launch service of W8A8_INT8 version of Llama-3.2-3B: ```bash python -m sglang.launch_server \ --model RedHatAI/Llama-3.2-3B-quantized.w8a8 \ --trust-remote-code \ --disable-overlap-schedule \ --device cpu \ --quantization w8a8_int8 \ --host 0.0.0.0 \ --enable-torch-compile \ --torch-compile-max-bs 16 \ --tp 2 ``` Note: The `--torch-compile-max-bs` and `--tp` settings are examples that should be adjusted for your setup. For instance, use `--tp 3` to utilize 1 socket with 3 sub-NUMA clusters on an Intel® Xeon® 6980P server. Once the server have been launched, you can test it using the `bench_serving` command or create your own commands or scripts following [the benchmarking example](#benchmarking-with-requests). ================================================ FILE: docs/platforms/mindspore_backend.md ================================================ # MindSpore Models ## Introduction MindSpore is a high-performance AI framework optimized for Ascend NPUs. This doc guides users to run MindSpore models in SGLang. ## Requirements MindSpore currently only supports Ascend NPU devices. Users need to first install Ascend CANN software packages. The CANN software packages can be downloaded from the [Ascend Official Website](https://www.hiascend.com). The recommended version is 8.3.RC2. ## Supported Models Currently, the following models are supported: - **Qwen3**: Dense and MoE models - **DeepSeek V3/R1** - *More models coming soon...* ## Installation > **Note**: Currently, MindSpore models are provided by an independent package `sgl-mindspore`. Support for MindSpore is built upon current SGLang support for Ascend NPU platform. Please first [install SGLang for Ascend NPU](ascend_npu.md) and then install `sgl-mindspore`: ```shell git clone https://github.com/mindspore-lab/sgl-mindspore.git cd sgl-mindspore pip install -e . ``` ## Run Model Current SGLang-MindSpore supports Qwen3 and DeepSeek V3/R1 models. This doc uses Qwen3-8B as an example. ### Offline infer Use the following script for offline infer: ```python import sglang as sgl # Initialize the engine with MindSpore backend llm = sgl.Engine( model_path="/path/to/your/model", # Local model path device="npu", # Use NPU device model_impl="mindspore", # MindSpore implementation attention_backend="ascend", # Attention backend tp_size=1, # Tensor parallelism size dp_size=1 # Data parallelism size ) # Generate text prompts = [ "Hello, my name is", "The capital of France is", "The future of AI is" ] sampling_params = {"temperature": 0, "top_p": 0.9} outputs = llm.generate(prompts, sampling_params) for prompt, output in zip(prompts, outputs): print(f"Prompt: {prompt}") print(f"Generated: {output['text']}") print("---") ``` ### Start server Launch a server with MindSpore backend: ```bash # Basic server startup python3 -m sglang.launch_server \ --model-path /path/to/your/model \ --host 0.0.0.0 \ --device npu \ --model-impl mindspore \ --attention-backend ascend \ --tp-size 1 \ --dp-size 1 ``` For distributed server with multiple nodes: ```bash # Multi-node distributed server python3 -m sglang.launch_server \ --model-path /path/to/your/model \ --host 0.0.0.0 \ --device npu \ --model-impl mindspore \ --attention-backend ascend \ --dist-init-addr 127.0.0.1:29500 \ --nnodes 2 \ --node-rank 0 \ --tp-size 4 \ --dp-size 2 ``` ## Troubleshooting #### Debug Mode Enable sglang debug logging by log-level argument. ```bash python3 -m sglang.launch_server \ --model-path /path/to/your/model \ --host 0.0.0.0 \ --device npu \ --model-impl mindspore \ --attention-backend ascend \ --log-level DEBUG ``` Enable mindspore info and debug logging by setting environments. ```bash export GLOG_v=1 # INFO export GLOG_v=0 # DEBUG ``` #### Explicitly select devices Use the following environment variable to explicitly select the devices to use. ```shell export ASCEND_RT_VISIBLE_DEVICES=4,5,6,7 # to set device ``` #### Some communication environment issues In case of some environment with special communication environment, users need set some environment variables. ```shell export MS_ENABLE_LCCL=off # current not support LCCL communication mode in SGLang-MindSpore ``` #### Some dependencies of protobuf In case of some environment with special protobuf version, users need set some environment variables to avoid binary version mismatch. ```shell export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python # to avoid protobuf binary version mismatch ``` ## Support For MindSpore-specific issues: - Refer to the [MindSpore documentation](https://www.mindspore.cn/) ================================================ FILE: docs/platforms/mthreads_gpu.md ================================================ # Moore Threads GPUs This document describes how run SGLang on Moore Threads GPUs. If you encounter issues or have questions, please [open an issue](https://github.com/sgl-project/sglang/issues). ## Install SGLang You can install SGLang using one of the methods below. ### Install from Source ```bash # Use the default branch git clone https://github.com/sgl-project/sglang.git cd sglang # Compile sgl-kernel pip install --upgrade pip cd sgl-kernel python setup_musa.py install # Install sglang python package cd .. rm -f python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml pip install -e "python[all_musa]" ``` ================================================ FILE: docs/platforms/nvidia_jetson.md ================================================ # NVIDIA Jetson Orin ## Prerequisites Before starting, ensure the following: - [**NVIDIA Jetson AGX Orin Devkit**](https://www.nvidia.com/en-us/autonomous-machines/embedded-systems/jetson-orin/) is set up with **JetPack 6.1** or later. - **CUDA Toolkit** and **cuDNN** are installed. - Verify that the Jetson AGX Orin is in **high-performance mode**: ```bash sudo nvpmodel -m 0 ``` * * * * * ## Installing and running SGLang with Jetson Containers Clone the jetson-containers github repository: ``` git clone https://github.com/dusty-nv/jetson-containers.git ``` Run the installation script: ``` bash jetson-containers/install.sh ``` Build the container image: ``` jetson-containers build sglang ``` Run the container: ``` jetson-containers run $(autotag sglang) ``` Or you can also manually run a container with this command: ``` docker run --runtime nvidia -it --rm --network=host IMAGE_NAME ``` * * * * * Running Inference ----------------------------------------- Launch the server: ```bash python -m sglang.launch_server \ --model-path deepseek-ai/DeepSeek-R1-Distill-Llama-8B \ --device cuda \ --dtype half \ --attention-backend flashinfer \ --mem-fraction-static 0.8 \ --context-length 8192 ``` The quantization and limited context length (`--dtype half --context-length 8192`) are due to the limited computational resources in [Nvidia jetson kit](https://www.nvidia.com/en-us/autonomous-machines/embedded-systems/jetson-orin/). A detailed explanation can be found in [Server Arguments](../advanced_features/server_arguments.md). After launching the engine, refer to [Chat completions](https://docs.sglang.io/basic_usage/openai_api_completions.html#Usage) to test the usability. * * * * * Running quantization with TorchAO ------------------------------------- TorchAO is suggested to NVIDIA Jetson Orin. ```bash python -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ --device cuda \ --dtype bfloat16 \ --attention-backend flashinfer \ --mem-fraction-static 0.8 \ --context-length 8192 \ --torchao-config int4wo-128 ``` This enables TorchAO's int4 weight-only quantization with a 128-group size. The usage of `--torchao-config int4wo-128` is also for memory efficiency. * * * * * Structured output with XGrammar ------------------------------- Please refer to [SGLang doc structured output](../advanced_features/structured_outputs.ipynb). * * * * * Thanks to the support from [Nurgaliyev Shakhizat](https://github.com/shahizat), [Dustin Franklin](https://github.com/dusty-nv) and [Johnny Núñez Cano](https://github.com/johnnynunez). References ---------- - [NVIDIA Jetson AGX Orin Documentation](https://developer.nvidia.com/embedded/jetson-agx-orin) ================================================ FILE: docs/platforms/tpu.md ================================================ # TPU SGLang supports high-performance TPU inference through the SGLang-JAX backend, which is specifically optimized for Google Cloud TPUs. The JAX-based implementation delivers exceptional throughput and low latency for Large Language Model (LLM) serving workloads on TPU hardware. For TPU-specific issues or feature requests, please visit the [sglang-jax GitHub issues page](https://github.com/sgl-project/sglang-jax/issues). **NOTE:** SGLang TPU support is implemented via the SGLang-JAX backend, a dedicated JAX-based inference engine maintained as a separate repository at [https://github.com/sgl-project/sglang-jax](https://github.com/sgl-project/sglang-jax). ## System Requirements ### Supported TPU Hardware | TPU Type | HBM Memory | Availability | |----------|-----------|--------------| | TPU v6e | 32 GB | Google Cloud | | TPU v7 | 96 GB per core | Google Cloud | ### Software Requirements - **Python:** 3.12 or higher - **JAX:** Latest version with TPU support - **Environment:** Google Cloud TPU VM or compatible TPU runtime - **Optional:** SkyPilot for simplified cloud deployment ## Feature Support Matrix SGLang-JAX provides comprehensive TPU-optimized features for production LLM serving: | Feature | Support Status | Description | |---------|---------------|-------------| | High-Throughput Continuous Batching | ✅ | Dynamic request batching for maximum TPU utilization | | Radix Tree KV Cache | ✅ | Memory-efficient prefix sharing between requests | | FlashAttention Backend | ✅ | TPU-optimized attention kernel for long sequences | | Tensor Parallelism | ✅ | Distribute models across multiple TPU cores | | Paged Attention | ✅ | Flexible KV cache management with paging | | Speculative Decoding (EAGLE/EAGLE3) | ✅ | 20-40% throughput improvement for compatible models | | Chunked Prefill | ✅ | Mixed prefill-decode batching | | OpenAI-Compatible API | ✅ | Drop-in replacement for OpenAI API | | Data Parallel Attention | 🚧 | In development - Attention computation with data parallelism | | Quantization | 🚧 | In development - Model quantization for reduced memory usage | | Multi-LoRA | 🚧 | In development - Serve multiple LoRA adapters simultaneously | ### Attention Backend Comparison | Backend | Paged Attention | Spec Decoding | MLA | Sliding Window | |---------|----------------|---------------|-----|----------------| | FlashAttention (fa) | ✅ | ✅ | ❌ | ✅ | | Native | ❌ | ❌ | ❌ | ❌ | **NOTE:** FlashAttention backend is recommended for production workloads due to superior memory efficiency and performance. ## Optimized Model List The following models have been tested and optimized for TPU deployment: | Model Family | Performance Status | |--------------|-------------------| | [Qwen 3](https://huggingface.co/Qwen) | ⭐ Recommended for production | | [Qwen 3 MoE](https://huggingface.co/Qwen) | ⭐ Best performance | | [Qwen 2](https://huggingface.co/Qwen) | Needs improvement | | [Qwen 2 MoE](https://huggingface.co/Qwen) | Needs improvement | | [Qwen 1.5](https://huggingface.co/Qwen) | Needs improvement | | [Llama/LLaMA](https://huggingface.co/meta-llama) | Needs improvement | | [Grok-2](https://huggingface.co/xai-org) | Needs improvement | | [Gemma 2](https://huggingface.co/google) | Verified on TPU | | Bailing MoE | Needs improvement | ## Installation ### Method 1: Using PyPI (Recommended) ```bash pip install sglang-jax ``` ### Method 2: From Source ```bash git clone https://github.com/sgl-project/sglang-jax cd sglang-jax uv venv --python 3.12 && source .venv/bin/activate uv pip install -e "python[all]" ``` ### Method 3: Using Docker **NOTE:** Docker support for TPU is currently under development. Please use PyPI or source installation methods. ### Method 4: Cloud TPU with SkyPilot [SkyPilot](https://github.com/skypilot-org/skypilot) provides simplified deployment on Google Cloud TPU: 1. Install SkyPilot and configure GCP access (see [SkyPilot documentation](https://skypilot.readthedocs.io/)) 2. Create a SkyPilot configuration file:
SkyPilot YAML: sglang-jax.sky.yaml ```yaml # sglang-jax.sky.yaml resources: accelerators: tpu-v6e-4 accelerator_args: tpu_vm: True runtime_version: v2-alpha-tpuv6e run: | git clone https://github.com/sgl-project/sglang-jax.git cd sglang-jax uv venv --python 3.12 source .venv/bin/activate uv pip install -e "python[all]" ```
3. Launch your TPU cluster: ```bash # Standard deployment sky launch -c sglang-jax sglang-jax.sky.yaml --infra=gcp # With spot instances for cost savings sky launch -c sglang-jax sglang-jax.sky.yaml --infra=gcp --use-spot ``` ## Launch of the Serving Engine ### Basic Example: Qwen-7B ```bash JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python3 -u -m sgl_jax.launch_server \ --model-path Qwen/Qwen-7B-Chat \ --trust-remote-code \ --dist-init-addr=0.0.0.0:10011 \ --nnodes=1 \ --tp-size=4 \ --device=tpu \ --random-seed=3 \ --node-rank=0 \ --mem-fraction-static=0.8 \ --max-prefill-tokens=8192 \ --download-dir=/tmp \ --dtype=bfloat16 \ --skip-server-warmup \ --host 0.0.0.0 \ --port 30000 ``` **Key Parameters Explained:** 1. `JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache` - Enables JIT compilation caching to accelerate server startup on subsequent runs 2. `--tp-size=4` - Tensor parallelism size; match this to your TPU core count (typically 1, 4, or 8) 3. `--device=tpu` - Specifies TPU device (this is the default for sglang-jax) 4. `--dtype=bfloat16` - Uses bfloat16 precision, which TPUs are optimized for 5. `--mem-fraction-static=0.8` - Allocates 80% of TPU HBM for static memory (adjustable from 0.2 to 0.9) 6. `--max-prefill-tokens=8192` - Maximum number of tokens processed in the prefill phase ### High-Performance Configuration: Qwen3-8B For production workloads with optimal throughput: ```bash python3 -u -m sgl_jax.launch_server \ --model-path Qwen/Qwen3-8B \ --trust-remote-code \ --tp-size=4 \ --device=tpu \ --mem-fraction-static=0.8 \ --chunked-prefill-size=2048 \ --dtype=bfloat16 \ --max-running-requests=256 \ --page-size=128 \ --attention-backend=fa ``` ### Advanced: Speculative Decoding (EAGLE3) Speculative decoding can improve throughput by 20-40% for compatible models: ```bash python3 -u -m sgl_jax.launch_server \ --model-path Qwen/Qwen3-32B \ --trust-remote-code \ --device=tpu \ --tp-size=4 \ --mem-fraction-static=0.8 \ --max-prefill-tokens=4096 \ --attention-backend=fa \ --dtype=bfloat16 \ --port=30000 \ --host=0.0.0.0 \ --disable-overlap-schedule \ --speculative-algorithm=EAGLE3 \ --speculative-draft-model-path=AngelSlim/Qwen3-32B_eagle3 \ --page-size=64 \ --speculative-eagle-topk=1 \ --speculative-num-steps=3 \ --speculative-num-draft-tokens=4 ``` **NOTE:** Speculative decoding is currently supported for Qwen3 and LLaMA model families. See the [Speculative Decoding documentation](https://github.com/sgl-project/sglang-jax/blob/main/docs/features/speculative_decoding.md) for detailed configuration guidance. ### Multi-Node Distributed Serving For large models requiring multiple TPU VMs: ```bash # Node 0 (coordinator) python3 -m sgl_jax.launch_server \ --model-path MODEL_PATH \ --dist-init-addr=NODE0_IP:10011 \ --nnodes=2 \ --node-rank=0 \ --tp-size=8 \ [other parameters...] # Node 1 (worker) python3 -m sgl_jax.launch_server \ --model-path MODEL_PATH \ --dist-init-addr=NODE0_IP:10011 \ --nnodes=2 \ --node-rank=1 \ --tp-size=8 \ [other parameters...] ``` ## Benchmarking with Requests ### Throughput Testing Basic throughput benchmark: ```bash python3 -m sgl_jax.bench_serving \ --backend sgl-jax \ --dataset-name random \ --num-prompts=100 \ --random-input=512 \ --random-output=128 \ --max-concurrency=8 \ --random-range-ratio=1 \ --warmup-requests=0 ``` ### Latency Testing Measure single-batch latency: ```bash python3 -m sgl_jax.bench_one_batch_server \ --base-url http://127.0.0.1:30000 \ --model-path Qwen/Qwen-7B-Chat \ --batch-size=32 \ --input-len=256 \ --output-len=32 ``` ### Comprehensive Benchmark Script For systematic performance evaluation across different configurations: ```bash #!/bin/bash set -e backend=${1:-sgl-jax} num_prompts_per_concurrency=3 input_seq_lens=(1024 4096 8192) output_seq_lens=(1 1024) max_concurrencies=(8 16 32 64 128 256) for input_seq_len in "${input_seq_lens[@]}"; do for output_seq_len in "${output_seq_lens[@]}"; do echo "=======================================" echo "Testing ISL/OSL: $input_seq_len/$output_seq_len" echo "=======================================" for max_concurrency in "${max_concurrencies[@]}"; do num_prompts=$((num_prompts_per_concurrency * max_concurrency)) python3 -m sgl_jax.bench_serving \ --backend ${backend} \ --dataset-name random \ --num-prompts ${num_prompts} \ --random-input ${input_seq_len} \ --random-output ${output_seq_len} \ --max-concurrency ${max_concurrency} \ --random-range-ratio 1 \ --disable-ignore-eos \ --warmup-requests 0 done done done ``` For detailed help on all benchmark parameters: ```bash python3 -m sgl_jax.bench_serving --help ``` See the [Benchmark and Profiling Guide](https://github.com/sgl-project/sglang-jax/blob/main/docs/developer_guide/benchmark_and_profiling.md) for advanced benchmarking techniques and profiling with JAX Profiler. ## Performance Optimization ### Memory Optimization **Reduce memory usage:** - Lower `--mem-fraction-static` (from 0.8 → 0.5 → 0.3) - Decrease `--max-prefill-tokens` (from 16384 → 8192 → 4096) - Reduce `--max-running-requests` **Handle OOM errors:** - Start with conservative memory settings (`--mem-fraction-static=0.5`) - Gradually increase until you find the optimal balance - Increase `--page-size` for better memory locality (1 → 16 → 64 → 128) ### Throughput Optimization To maximize tokens per second: - Use FlashAttention backend: `--attention-backend=fa` - Enable speculative decoding (EAGLE3) for Qwen3 models (20-40% improvement) - Increase `--max-running-requests` to 256+ - Set `--mem-fraction-static` to 0.8+ (if memory allows) - Use larger page sizes (64-128) - Enable chunked prefill: `--chunked-prefill-size=2048` ### Latency Optimization To minimize time-to-first-token (TTFT) and inter-token latency: - Reduce `--page-size` to 1-4 - Lower `--max-running-requests` (16-32) for smaller batches - Reduce `--chunked-prefill-size` - Use conservative memory settings to avoid GC pauses ### TPU-Specific Optimizations 1. **JIT Compilation Cache:** ```bash export JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache ``` Always set this environment variable to cache compiled kernels and accelerate server startup. 2. **Data Type Optimization:** Use `--dtype=bfloat16` for TPU native optimization. TPUs are specifically designed for bfloat16 computations. 3. **Tensor Parallelism:** Match `--tp-size` to your TPU core configuration (1, 4, or 8) for optimal model distribution. 4. **Attention Backend:** Always use `--attention-backend=fa` (FlashAttention) for production workloads. ## Troubleshooting ### OOM (Out of Memory) Errors If you encounter out-of-memory errors: 1. Reduce `--mem-fraction-static` from 0.8 to 0.5 or lower 2. Decrease `--max-prefill-tokens` from 8192 to 4096 or 2048 3. Lower `--max-running-requests` to reduce concurrent batch size 4. Increase `--page-size` for better memory layout efficiency ### Compilation Long-Time If the server takes too long to start: 1. Ensure `JAX_COMPILATION_CACHE_DIR` is properly set 2. Understand that the first run requires JIT compilation (this is normal) 3. Subsequent runs will be significantly faster with cached compilations 4. Consider using `--skip-server-warmup` to defer compilation until first request ### Low Throughput If you're not achieving expected throughput: 1. Verify `--tp-size` matches your TPU core configuration 2. Check that `--attention-backend=fa` is enabled 3. Increase `--max-running-requests` to enable larger batch formation 4. Consider enabling speculative decoding for compatible models 5. Ensure memory settings allow for sufficient batch sizes ### Connection Issues If clients cannot connect to the server: 1. Ensure `--host=0.0.0.0` for external access (not just `127.0.0.1`) 2. Verify firewall rules allow traffic on the specified port (default: 30000) 3. Check that the server process is running: `curl http://localhost:30000/health` ## Advanced Features ### Speculative Decoding SGLang-JAX supports EAGLE and EAGLE3 speculative decoding algorithms for Qwen3 and LLaMA model families. Speculative decoding can improve throughput by 20-40% without affecting output quality. See the [Speculative Decoding documentation](https://github.com/sgl-project/sglang-jax/blob/main/docs/features/speculative_decoding.md) for detailed configuration and supported model combinations. ### Chunked Prefill Enable mixed prefill-decode batching for better TPU utilization: ```bash --chunked-prefill-size=2048 --enable-mixed-chunk ``` This allows the scheduler to mix prefill operations with decode operations in the same batch, improving overall throughput. ### Custom Attention Backends SGLang-JAX supports a plugin-based attention backend system. You can implement custom attention kernels optimized for specific use cases. See the [Attention Backend documentation](https://github.com/sgl-project/sglang-jax/blob/main/docs/features/attention_backend.md) for implementation details. ### Environment Verification Verify your TPU setup before deploying: ```bash python -c "from sgl_jax import check_env; check_env.check_env()" ``` This command checks: - Installed package versions - TPU device availability and specifications - System resources and configuration - Compatibility of settings ## Contributing We welcome contributions to improve TPU support in SGLang-JAX! ### Areas for Contribution **Check the [Development Roadmap](https://github.com/sgl-project/sglang-jax/issues/190)** to see planned features and find opportunities to contribute new functionality. Current contribution areas include: - Performance optimizations for specific TPU generations - Support for additional model architectures - Documentation improvements and examples - Bug reports and fixes - Benchmark results and performance analysis ### How to Contribute 1. Visit the [sglang-jax repository](https://github.com/sgl-project/sglang-jax) 2. Read the [Contribution Guide](https://github.com/sgl-project/sglang-jax/blob/main/docs/developer_guide/contribution_guide.md) 3. Join the [SGL-JAX Slack community](https://sgl-fru7574.slack.com/archives/C09EBE5HT5X) for discussions 4. Report issues at [sglang-jax/issues](https://github.com/sgl-project/sglang-jax/issues) ### Testing on TPU For contributors who need TPU access for testing: - Refer to the [TPU Resources Guide](https://github.com/sgl-project/sglang-jax/blob/main/docs/developer_guide/tpu_resources_guide.md) for information on accessing TPU hardware - Use SkyPilot with spot instances for cost-effective testing - Follow the [Benchmark and Profiling Guide](https://github.com/sgl-project/sglang-jax/blob/main/docs/developer_guide/benchmark_and_profiling.md) for performance validation ## References ### Documentation - [SGLang-JAX Repository](https://github.com/sgl-project/sglang-jax) - [SGLang-JAX Installation Guide](https://github.com/sgl-project/sglang-jax/blob/main/docs/get_started/install.md) - [Qwen Models Quick Start](https://github.com/sgl-project/sglang-jax/blob/main/docs/basic_usage/qwen.md) - [Benchmark and Profiling Guide](https://github.com/sgl-project/sglang-jax/blob/main/docs/developer_guide/benchmark_and_profiling.md) - [Speculative Decoding](https://github.com/sgl-project/sglang-jax/blob/main/docs/features/speculative_decoding.md) ### External Resources - [JAX Documentation](https://jax.readthedocs.io/) - [Google Cloud TPU Documentation](https://cloud.google.com/tpu/docs) - [SkyPilot Documentation](https://skypilot.readthedocs.io/) ================================================ FILE: docs/platforms/xpu.md ================================================ # XPU The document addresses how to set up the [SGLang](https://github.com/sgl-project/sglang) environment and run LLM inference on Intel GPU, [see more context about Intel GPU support within PyTorch ecosystem](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html). Specifically, SGLang is optimized for [Intel® Arc™ Pro B-Series Graphics](https://www.intel.com/content/www/us/en/ark/products/series/242616/intel-arc-pro-b-series-graphics.html) and [ Intel® Arc™ B-Series Graphics](https://www.intel.com/content/www/us/en/ark/products/series/240391/intel-arc-b-series-graphics.html). ## Optimized Model List A list of LLMs have been optimized on Intel GPU, and more are on the way: | Model Name | BF16 | |:---:|:---:| | Llama-3.2-3B | [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) | | Llama-3.1-8B | [meta-llama/Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) | | Qwen2.5-1.5B | [Qwen/Qwen2.5-1.5B](https://huggingface.co/Qwen/Qwen2.5-1.5B) | **Note:** The model identifiers listed in the table above have been verified on [Intel® Arc™ B580 Graphics](https://www.intel.com/content/www/us/en/products/sku/241598/intel-arc-b580-graphics/specifications.html). ## Installation ### Install From Source Currently SGLang XPU only supports installation from source. Please refer to ["Getting Started on Intel GPU"](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html) to install XPU dependency. ```bash # Create and activate a conda environment conda create -n sgl-xpu python=3.12 -y conda activate sgl-xpu # Set PyTorch XPU as primary pip install channel to avoid installing the larger CUDA-enabled version and prevent potential runtime issues. pip3 install torch==2.10.0+xpu torchao torchvision torchaudio triton-xpu==3.6.0 --index-url https://download.pytorch.org/whl/xpu pip3 install xgrammar --no-deps # xgrammar will introduce CUDA-enabled triton which might conflict with XPU # Clone the SGLang code git clone https://github.com/sgl-project/sglang.git cd sglang git checkout # Use dedicated toml file cd python cp pyproject_xpu.toml pyproject.toml # Install SGLang dependent libs, and build SGLang main package pip install --upgrade pip setuptools pip install -v . --extra-index-url https://download.pytorch.org/whl/xpu ``` ### Install Using Docker The docker for XPU is under active development. Please stay tuned. ## Launch of the Serving Engine Example command to launch SGLang serving: ```bash python -m sglang.launch_server \ --model \ --trust-remote-code \ --disable-overlap-schedule \ --device xpu \ --host 0.0.0.0 \ --tp 2 \ # using multi GPUs --attention-backend intel_xpu \ # using intel optimized XPU attention backend --page-size \ # intel_xpu attention backend supports [32, 64, 128] ``` ## Benchmarking with Requests You can benchmark the performance via the `bench_serving` script. Run the command in another terminal. ```bash python -m sglang.bench_serving \ --dataset-name random \ --random-input-len 1024 \ --random-output-len 1024 \ --num-prompts 1 \ --request-rate inf \ --random-range-ratio 1.0 ``` The detail explanations of the parameters can be looked up by the command: ```bash python -m sglang.bench_serving -h ``` Additionally, the requests can be formed with [OpenAI Completions API](https://docs.sglang.io/basic_usage/openai_api_completions.html) and sent via the command line (e.g. using `curl`) or via your own script. ================================================ FILE: docs/references/custom_chat_template.md ================================================ # Custom Chat Template **NOTE**: There are two chat template systems in SGLang project. This document is about setting a custom chat template for the OpenAI-compatible API server (defined at [conversation.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/conversation.py)). It is NOT related to the chat template used in the SGLang language frontend (defined at [chat_template.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/lang/chat_template.py)). By default, the server uses the chat template specified in the model tokenizer from Hugging Face. It should just work for most official models such as Llama-2/Llama-3. If needed, you can also override the chat template when launching the server: ```bash python -m sglang.launch_server \ --model-path meta-llama/Llama-2-7b-chat-hf \ --port 30000 \ --chat-template llama-2 ``` If the chat template you are looking for is missing, you are welcome to contribute it or load it from a file. ## JSON Format You can load the JSON format, which is defined by `conversation.py`. ```json { "name": "my_model", "system": "<|im_start|>system", "user": "<|im_start|>user", "assistant": "<|im_start|>assistant", "sep_style": "CHATML", "sep": "<|im_end|>", "stop_str": ["<|im_end|>", "<|im_start|>"] } ``` ```bash python -m sglang.launch_server \ --model-path meta-llama/Llama-2-7b-chat-hf \ --port 30000 \ --chat-template ./my_model_template.json ``` ## Jinja Format You can also use the [Jinja template format](https://huggingface.co/docs/transformers/main/en/chat_templating) as defined by Hugging Face Transformers. ```bash python -m sglang.launch_server \ --model-path meta-llama/Llama-2-7b-chat-hf \ --port 30000 \ --chat-template ./my_model_template.jinja ``` ================================================ FILE: docs/references/environment_variables.md ================================================ # Environment Variables SGLang supports various environment variables that can be used to configure its runtime behavior. This document provides a comprehensive list and aims to stay updated over time. *Note: SGLang uses two prefixes for environment variables: `SGL_` and `SGLANG_`. This is likely due to historical reasons. While both are currently supported for different settings, future versions might consolidate them.* ## General Configuration | Environment Variable | Description | Default Value | |-------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------|------------------------------| | `SGLANG_USE_MODELSCOPE` | Enable using models from ModelScope | `false` | | `SGLANG_HOST_IP` | Host IP address for the server | `0.0.0.0` | | `SGLANG_PORT` | Port for the server | auto-detected | | `SGLANG_LOGGING_CONFIG_PATH` | Custom logging configuration path | Not set | | `SGLANG_DISABLE_REQUEST_LOGGING` | Disable request logging | `false` | | `SGLANG_LOG_REQUEST_HEADERS` | Comma-separated list of additional HTTP headers to log when `--log-requests` is enabled. Appends to the default `x-smg-routing-key`. | Not set | | `SGLANG_HEALTH_CHECK_TIMEOUT` | Timeout for health check in seconds | `20` | | `SGLANG_EPLB_HEATMAP_COLLECTION_INTERVAL` | The interval of passes to collect the metric of selected count of physical experts on each layer and GPU rank. 0 means disabled. | `0` | | `SGLANG_FORWARD_UNKNOWN_TOOLS` | Forward unknown tool calls to clients instead of dropping them | `false` (drop unknown tools) | | `SGLANG_REQ_WAITING_TIMEOUT` | Timeout (in seconds) for requests waiting in the queue before being scheduled | `-1` | | `SGLANG_REQ_RUNNING_TIMEOUT` | Timeout (in seconds) for requests running in the decode batch | `-1` | ## Performance Tuning | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_ENABLE_TORCH_INFERENCE_MODE` | Control whether to use torch.inference_mode | `false` | | `SGLANG_ENABLE_TORCH_COMPILE` | Enable torch.compile | `false` | | `SGLANG_SET_CPU_AFFINITY` | Enable CPU affinity setting (often set to `1` in Docker builds) | `false` | | `SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN` | Allows the scheduler to overwrite longer context length requests (often set to `1` in Docker builds) | `false` | | `SGLANG_IS_FLASHINFER_AVAILABLE` | Control FlashInfer availability check | `true` | | `SGLANG_SKIP_P2P_CHECK` | Skip P2P (peer-to-peer) access check | `false` | | `SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD` | Sets the threshold for enabling chunked prefix caching | `8192` | | `SGLANG_FUSED_MLA_ENABLE_ROPE_FUSION` | Enable RoPE fusion in Fused Multi-Layer Attention | `1` | | `SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP` | Disable overlap schedule for consecutive prefill batches | `false` | | `SGLANG_SCHEDULER_MAX_RECV_PER_POLL` | Set the maximum number of requests per poll, with a negative value indicating no limit | `-1` | | `SGLANG_DISABLE_FA4_WARMUP` | Disable Flash Attention 4 warmup passes (set to `1`, `true`, `yes`, or `on` to disable) | `false` | | `SGLANG_DATA_PARALLEL_BUDGET_INTERVAL` | Interval for DPBudget updates | `1` | | `SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_DEFAULT` | Default weight value for scheduler recv skipper counter (used when forward mode doesn't match specific modes). Only active when `--scheduler-recv-interval > 1`. The counter accumulates weights and triggers request polling when reaching the interval threshold. | `1000` | | `SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_DECODE` | Weight increment for decode forward mode in scheduler recv skipper. Works with `--scheduler-recv-interval` to control polling frequency during decode phase. | `1` | | `SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_TARGET_VERIFY` | Weight increment for target verify forward mode in scheduler recv skipper. Works with `--scheduler-recv-interval` to control polling frequency during verification phase. | `1` | | `SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_NONE` | Weight increment when forward mode is None in scheduler recv skipper. Works with `--scheduler-recv-interval` to control polling frequency when no specific forward mode is active. | `1` | | `SGLANG_MM_BUFFER_SIZE_MB` | Size of preallocated GPU buffer (in MB) for multi-modal feature hashing optimization. When set to a positive value, temporarily moves features to GPU for faster hash computation, then moves them back to CPU to save GPU memory. Larger features benefit more from GPU hashing. Set to `0` to disable. | `0` | | `SGLANG_MM_PRECOMPUTE_HASH` | Enable precomputing of hash values for MultimodalDataItem | `false` | | `SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH` | Enable NCCL for gathering when preparing mlp sync batch under overlap scheduler (without this flag gloo is used for gathering) | `false` | | `SGLANG_SYMM_MEM_PREALLOC_GB_SIZE` | Size of preallocated GPU buffer (in GB) for NCCL symmetric memory pool to limit memory fragmentation. Only have an effect when server arg `--enable-symm-mem` is set. | `-1` | | `SGLANG_CUSTOM_ALLREDUCE_ALGO` | The algorithm of custom all-reduce. Set to `oneshot` or `1stage` to force use one-shot. Set to `twoshot` or `2stage` to force use two-shot. | `` | ## DeepGEMM Configuration (Advanced Optimization) | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_ENABLE_JIT_DEEPGEMM` | Enable Just-In-Time compilation of DeepGEMM kernels (enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) GPUs when the DeepGEMM package is installed; set to `"0"` to disable) | `"true"` | | `SGLANG_JIT_DEEPGEMM_PRECOMPILE` | Enable precompilation of DeepGEMM kernels | `"true"` | | `SGLANG_JIT_DEEPGEMM_COMPILE_WORKERS` | Number of workers for parallel DeepGEMM kernel compilation | `4` | | `SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE` | Indicator flag used during the DeepGEMM precompile script | `"false"` | | `SGLANG_DG_CACHE_DIR` | Directory for caching compiled DeepGEMM kernels | `~/.cache/deep_gemm` | | `SGLANG_DG_USE_NVRTC` | Use NVRTC (instead of Triton) for JIT compilation (Experimental) | `"false"` | | `SGLANG_USE_DEEPGEMM_BMM` | Use DeepGEMM for Batched Matrix Multiplication (BMM) operations | `"false"` | | `SGLANG_JIT_DEEPGEMM_FAST_WARMUP` | Precompile less kernels during warmup, which reduces the warmup time from 30min to less than 3min. Might cause performance degradation during runtime. | `"false"` | ## DeepEP Configuration | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_DEEPEP_BF16_DISPATCH` | Use Bfloat16 for dispatch | `"false"` | | `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | The maximum number of dispatched tokens on each GPU | `"128"` | | `SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | The maximum number of dispatched tokens on each GPU for --moe-a2a-backend=flashinfer | `"1024"` | | `SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS` | Number of SMs used for DeepEP combine when single batch overlap is enabled | `"32"` | | `SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO` | Run shared experts on an alternate stream when single batch overlap is enabled on GB200. When not setting this flag, shared experts and down gemm will be overlapped with DeepEP combine together. | `"false"` | ## MORI Configuration | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_MORI_FP8_DISP` | Use FP8 for dispatch | `"false"` | | `SGLANG_MORI_FP4_DISP` | Use MXFP4 for dispatch | `"false"` | | `SGLANG_MORI_FP8_COMB` | Use FP8 for combine | `"false"` | | `SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | Maximum number of dispatch tokens per rank for MORI-EP buffer allocation | `4096` | | `SGLANG_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD` | Threshold for switching between `InterNodeV1` and `InterNodeV1LL` kernel types. `InterNodeV1LL` is used if `SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK` is less than or equal to this threshold; otherwise, `InterNodeV1` is used. | `256` | | `SGLANG_MORI_QP_PER_TRANSFER` | Number of RDMA Queue Pairs (QPs) used per transfer operation | `1` | | `SGLANG_MORI_POST_BATCH_SIZE` | Number of RDMA work requests posted in a single batch to each QP | `-1` | | `SGLANG_MORI_NUM_WORKERS` | Number of worker threads in the RDMA executor thread pool | `1` | ## NSA Backend Configuration (For DeepSeek V3.2) | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_NSA_FUSE_TOPK` | Fuse the operation of picking topk logits and picking topk indices from page table | `true` | | `SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA` | Precompute metadata that can be shared among different draft steps when MTP is enabled | `true` | | `SGLANG_USE_FUSED_METADATA_COPY` | Control whether to use fused metadata copy kernel for cuda graph replay | `true` | | `SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD` | When the maximum kv len in current prefill batch exceeds this value, the sparse mla kernel will be applied, else it falls back to dense MHA implementation. Default to the index topk of model (2048 for DeepSeek V3.2) | `2048` | ## Memory Management | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_DEBUG_MEMORY_POOL` | Enable memory pool debugging | `false` | | `SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION` | Clip max new tokens estimation for memory planning | `4096` | | `SGLANG_DETOKENIZER_MAX_STATES` | Maximum states for detokenizer | Default value based on system | | `SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK` | Enable checks for memory imbalance across Tensor Parallel ranks | `true` | | `SGLANG_MOONCAKE_CUSTOM_MEM_POOL` | Configure the custom memory pool type for Mooncake. Supports `NVLINK`, `BAREX`, `INTRA_NODE_NVLINK`. If set to `true`, it defaults to `NVLINK`. | `None` | ## Model-Specific Options | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_USE_AITER` | Use AITER optimize implementation | `false` | | `SGLANG_MOE_PADDING` | Enable MoE padding (sets padding size to 128 if value is `1`, often set to `1` in Docker builds) | `false` | | `SGLANG_CUTLASS_MOE` (deprecated) | Use Cutlass FP8 MoE kernel on Blackwell GPUs (deprecated, use --moe-runner-backend=cutlass) | `false` | ## Quantization | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_INT4_WEIGHT` | Enable INT4 weight quantization | `false` | | `SGLANG_PER_TOKEN_GROUP_QUANT_8BIT_V2` | Apply per token group quantization kernel with fused silu and mul and masked m | `false` | | `SGLANG_FORCE_FP8_MARLIN` | Force using FP8 MARLIN kernels even if other FP8 kernels are available | `false` | | `SGLANG_FLASHINFER_FP4_GEMM_BACKEND` (deprecated) | Select backend for `mm_fp4` on Blackwell GPUs. **DEPRECATED**: Please use `--fp4-gemm-backend` instead. | `` | | `SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN` | Quantize q_b_proj from BF16 to FP8 when launching DeepSeek NVFP4 checkpoint | `false` | | `SGLANG_MOE_NVFP4_DISPATCH` | Use nvfp4 for moe dispatch (on flashinfer_cutlass or flashinfer_cutedsl moe runner backend) | `"false"` | | `SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE` | Quantize moe of nextn layer from BF16 to FP8 when launching DeepSeek NVFP4 checkpoint | `false` | | `SGLANG_ENABLE_FLASHINFER_FP8_GEMM` (deprecated) | Use flashinfer kernels when running blockwise fp8 GEMM on Blackwell GPUs. **DEPRECATED**: Please use `--fp8-gemm-backend=flashinfer_trtllm` (SM100/SM103) or `--fp8-gemm-backend=flashinfer_cutlass` (SM120/SM121 and newer) instead. | `false` | | `SGLANG_SUPPORT_CUTLASS_BLOCK_FP8` (deprecated) | Use Cutlass kernels when running blockwise fp8 GEMM on Hopper or Blackwell GPUs. **DEPRECATED**: Please use `--fp8-gemm-backend=cutlass` instead. | `false` | | `SGLANG_QUANT_ALLOW_DOWNCASTING` | Allow weight dtype downcasting during loading (e.g., fp32 → fp16). By default, SGLang rejects this kind of downcasting when using quantization. | `false` | | `SGLANG_FP8_IGNORED_LAYERS` | A comma-separated list of layer names to ignore during FP8 quantization. For example: `model.layers.0,model.layers.1.,qkv_proj`. | `""` | ## Distributed Computing | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_BLOCK_NONZERO_RANK_CHILDREN` | Control blocking of non-zero rank children processes | `1` | | `SGLANG_IS_FIRST_RANK_ON_NODE` | Indicates if the current process is the first rank on its node | `"true"` | | `SGLANG_PP_LAYER_PARTITION` | Pipeline parallel layer partition specification | Not set | | `SGLANG_ONE_VISIBLE_DEVICE_PER_PROCESS` | Set one visible device per process for distributed computing | `false` | ## Testing & Debugging (Internal/CI) *These variables are primarily used for internal testing, continuous integration, or debugging.* | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_IS_IN_CI` | Indicates if running in CI environment | `false` | | `SGLANG_IS_IN_CI_AMD` | Indicates running in AMD CI environment | `false` | | `SGLANG_TEST_RETRACT` | Enable retract decode testing | `false` | | `SGLANG_TEST_RETRACT_NO_PREFILL_BS` | When SGLANG_TEST_RETRACT is enabled, no prefill is performed if the batch size exceeds SGLANG_TEST_RETRACT_NO_PREFILL_BS. | `2 ** 31` | | `SGLANG_RECORD_STEP_TIME` | Record step time for profiling | `false` | | `SGLANG_TEST_REQUEST_TIME_STATS` | Test request time statistics | `false` | ## Profiling & Benchmarking | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_TORCH_PROFILER_DIR` | Directory for PyTorch profiler output | `/tmp` | | `SGLANG_PROFILE_WITH_STACK` | Set `with_stack` option (bool) for PyTorch profiler (capture stack trace) | `true` | | `SGLANG_PROFILE_RECORD_SHAPES` | Set `record_shapes` option (bool) for PyTorch profiler (record shapes) | `true` | | `SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS` | Config BatchSpanProcessor.schedule_delay_millis if tracing is enabled | `500` | | `SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE` | Config BatchSpanProcessor.max_export_batch_size if tracing is enabled | `64` | ## Storage & Caching | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_WAIT_WEIGHTS_READY_TIMEOUT` | Timeout period for waiting on weights | `120` | | `SGLANG_DISABLE_OUTLINES_DISK_CACHE` | Disable Outlines disk cache | `false` | | `SGLANG_USE_CUSTOM_TRITON_KERNEL_CACHE` | Use SGLang's custom Triton kernel cache implementation for lower overheads (automatically enabled on CUDA) | `false` | | `SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE` | Decode-side incremental KV cache offload stride. Rounded down to a multiple of `--page-size` (min is `--page-size`). If unset/invalid/<=0, it falls back to `--page-size`. | Not set (uses `--page-size`) | ## Function Calling / Tool Use | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_TOOL_STRICT_LEVEL` | Controls the strictness level of tool call parsing and validation.
**Level 0**: Off - No strict validation
**Level 1**: Function strict - Enables structural tag constraints for all tools (even if none have `strict=True` set)
**Level 2**: Parameter strict - Enforces strict parameter validation for all tools, treating them as if they all have `strict=True` set | `0` | ================================================ FILE: docs/references/faq.md ================================================ # Troubleshooting and Frequently Asked Questions ## Troubleshooting This page lists common errors and tips for resolving them. ### CUDA Out of Memory If you encounter out-of-memory (OOM) errors, you can adjust the following parameters: - If OOM occurs during prefill, try reducing `--chunked-prefill-size` to `4096` or `2048`. This saves memory but slows down the prefill speed for long prompts. - If OOM occurs during decoding, try lowering `--max-running-requests`. - You can also decrease `--mem-fraction-static` to a smaller value, such as 0.8 or 0.7. This decreases the memory usage of the KV cache memory pool and helps prevent OOM errors during both prefill and decoding. However, it limits maximum concurrency and reduces peak throughput. - Another common case for OOM is requesting input logprobs for a long prompt as it requires significant memory. To address this, set `logprob_start_len` in your sampling parameters to include only the necessary parts. If you do need input logprobs for a long prompt, try reducing `--mem-fraction-static`. ### CUDA Error: Illegal Memory Access Encountered This error may result from kernel errors or out-of-memory issues: - If it is a kernel error, resolving it may be challenging. Please file an issue on GitHub. - If it is an out-of-memory issue, it may sometimes be reported as this error instead of "Out of Memory." Refer to the section above for guidance on avoiding OOM issues. ### The server hangs - If the server hangs during initialization or running, it can be memory issues (out of memory), network issues (nccl errors), or other bugs in sglang. - If it is out of memory, you might see that `avail mem` is very low during the initialization or right after initialization. In this case, you can try to decrease `--mem-fraction-static`, decrease `--cuda-graph-max-bs`, or decrease `--chunked-prefill-size`. - Other bugs, please file an issue on GitHub. ## Frequently Asked Questions ### The results are not deterministic, even with a temperature of 0 You may notice that when you send the same request twice, the results from the engine will be slightly different, even when the temperature is set to 0. From our initial investigation, this indeterminism arises from two factors: dynamic batching and prefix caching. Roughly speaking, dynamic batching accounts for about 95% of the indeterminism, while prefix caching accounts for the remaining portion. The server runs dynamic batching under the hood. Different batch sizes can cause PyTorch/CuBLAS to dispatch to different CUDA kernels, which can lead to slight numerical differences. This difference accumulates across many layers, resulting in nondeterministic output when the batch size changes. Similarly, when prefix caching is enabled, it can also dispatch to different kernels. Even when the computations are mathematically equivalent, small numerical differences from different kernel implementations lead to the final nondeterministic outputs. To achieve more deterministic outputs in the current code, you can add `--disable-radix-cache` and send only one request at a time. The results will be mostly deterministic under this setting. **Update**: Recently, we also introduced a deterministic mode, you can enable it with `--enable-deterministic-inference`. Please find more details in this blog post: https://lmsys.org/blog/2025-09-22-sglang-deterministic/ ================================================ FILE: docs/references/frontend/choices_methods.md ================================================ # Choices Methods in SGLang This doc describes the choices methods supported by SGLang. The optional `choices_method` arg determines how options supplied to SGLang's `choices` primitive are selected. Only the `RuntimeEndpoint` backend supports the `choices_method` arg. Other backends, such as `OpenAI`, have bespoke selection implementations due to API limitations. ## Methods ### Token Length Normalized Token length normalized is the default SGLang choices method. It selects the option with the highest average logprob across all of its tokens. Usage example (alternatively, simply omit the `choices_method` arg): ```python @sgl.function def example(s): s += sgl.user("What is the capital of France?") s += sgl.assistant( sgl.gen( "answer", choices=["London", "Paris", "Berlin"], choices_method=sgl.token_length_normalized, ) ) ``` This can perform poorly if an option contains many tokens, where its later tokens are predicted with high confidence based on its earlier tokens. For instance, even strong models will fail the above example if the specified options are `["Paris", "Antidisestablishmentarianism"]`. ### Greedy Token Selection Greedy token selection simply selects the option with the highest logprob for its initial token. For overlapping options where one option is a subset of a longer option, the logprobs of the shorter option are extended using its average logprob for comparison against the longer option. Usage example: ```python @sgl.function def example(s): s += sgl.user("What is the capital of France?") s += sgl.assistant( sgl.gen( "answer", choices=["London", "Paris", "Berlin"], choices_method=sgl.greedy_token_selection, ) ) ``` This can perform poorly if an option misleads the model down a bad path based on an attractive initial token. For instance, greedy selection will result in an incorrect response for this example: ```python @sgl.function def us_president_example(s): s += sgl.user("Name a US president.") s += sgl.assistant( sgl.gen( "answer", choices=["Donald Duck", "Millard Fillmore"], choices_method=sgl.greedy_token_selection, ) ) ``` ### Unconditional Likelihood Normalized Unconditional likelihood normalized selects the option with the highest average token logprob once normalized by the unconditional token logprobs, as described in [this EleutherAI blogpost](https://blog.eleuther.ai/multiple-choice-normalization/). This method incurs an additional LLM call to obtain the unconditional likelihoods. Usage example: ```python @sgl.function def example(s): s += sgl.user("What is the capital of France?") s += sgl.assistant( sgl.gen( "answer", choices=["London", "Paris", "Berlin"], choices_method=sgl.unconditional_likelihood_normalized, ) ) ``` ================================================ FILE: docs/references/frontend/frontend_index.rst ================================================ Frontend Language ================= .. toctree:: :maxdepth: 1 :caption: Frontend Language frontend_tutorial.ipynb choices_methods.md ================================================ FILE: docs/references/frontend/frontend_tutorial.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# SGLang Frontend Language" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "SGLang frontend language can be used to define simple and easy prompts in a convenient, structured way." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Launch A Server\n", "\n", "Launch the server in your terminal and wait for it to initialize." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from sglang import assistant_begin, assistant_end\n", "from sglang import assistant, function, gen, system, user\n", "from sglang import image\n", "from sglang import RuntimeEndpoint\n", "from sglang.lang.api import set_default_backend\n", "from sglang.srt.utils import load_image\n", "from sglang.test.doc_patch import launch_server_cmd\n", "from sglang.utils import print_highlight, terminate_process, wait_for_server\n", "\n", "server_process, port = launch_server_cmd(\n", " \"python -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct --host 0.0.0.0 --log-level warning\"\n", ")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=server_process)\n", "print(f\"Server started on http://localhost:{port}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Set the default backend. Note: Besides the local server, you may use also `OpenAI` or other API endpoints." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "set_default_backend(RuntimeEndpoint(f\"http://localhost:{port}\"))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Basic Usage\n", "\n", "The most simple way of using SGLang frontend language is a simple question answer dialog between a user and an assistant." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@function\n", "def basic_qa(s, question):\n", " s += system(f\"You are a helpful assistant than can answer questions.\")\n", " s += user(question)\n", " s += assistant(gen(\"answer\", max_tokens=512))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "state = basic_qa(\"List 3 countries and their capitals.\")\n", "print_highlight(state[\"answer\"])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Multi-turn Dialog\n", "\n", "SGLang frontend language can also be used to define multi-turn dialogs." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@function\n", "def multi_turn_qa(s):\n", " s += system(f\"You are a helpful assistant than can answer questions.\")\n", " s += user(\"Please give me a list of 3 countries and their capitals.\")\n", " s += assistant(gen(\"first_answer\", max_tokens=512))\n", " s += user(\"Please give me another list of 3 countries and their capitals.\")\n", " s += assistant(gen(\"second_answer\", max_tokens=512))\n", " return s\n", "\n", "\n", "state = multi_turn_qa()\n", "print_highlight(state[\"first_answer\"])\n", "print_highlight(state[\"second_answer\"])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Control flow\n", "\n", "You may use any Python code within the function to define more complex control flows." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@function\n", "def tool_use(s, question):\n", " s += assistant(\n", " \"To answer this question: \"\n", " + question\n", " + \". I need to use a \"\n", " + gen(\"tool\", choices=[\"calculator\", \"search engine\"])\n", " + \". \"\n", " )\n", "\n", " if s[\"tool\"] == \"calculator\":\n", " s += assistant(\"The math expression is: \" + gen(\"expression\"))\n", " elif s[\"tool\"] == \"search engine\":\n", " s += assistant(\"The key word to search is: \" + gen(\"word\"))\n", "\n", "\n", "state = tool_use(\"What is 2 * 2?\")\n", "print_highlight(state[\"tool\"])\n", "print_highlight(state[\"expression\"])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Parallelism\n", "\n", "Use `fork` to launch parallel prompts. Because `sgl.gen` is non-blocking, the for loop below issues two generation calls in parallel." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@function\n", "def tip_suggestion(s):\n", " s += assistant(\n", " \"Here are two tips for staying healthy: \"\n", " \"1. Balanced Diet. 2. Regular Exercise.\\n\\n\"\n", " )\n", "\n", " forks = s.fork(2)\n", " for i, f in enumerate(forks):\n", " f += assistant(\n", " f\"Now, expand tip {i+1} into a paragraph:\\n\"\n", " + gen(\"detailed_tip\", max_tokens=256, stop=\"\\n\\n\")\n", " )\n", "\n", " s += assistant(\"Tip 1:\" + forks[0][\"detailed_tip\"] + \"\\n\")\n", " s += assistant(\"Tip 2:\" + forks[1][\"detailed_tip\"] + \"\\n\")\n", " s += assistant(\n", " \"To summarize the above two tips, I can say:\\n\" + gen(\"summary\", max_tokens=512)\n", " )\n", "\n", "\n", "state = tip_suggestion()\n", "print_highlight(state[\"summary\"])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Constrained Decoding\n", "\n", "Use `regex` to specify a regular expression as a decoding constraint. This is only supported for local models." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@function\n", "def regular_expression_gen(s):\n", " s += user(\"What is the IP address of the Google DNS servers?\")\n", " s += assistant(\n", " gen(\n", " \"answer\",\n", " temperature=0,\n", " regex=r\"((25[0-5]|2[0-4]\\d|[01]?\\d\\d?).){3}(25[0-5]|2[0-4]\\d|[01]?\\d\\d?)\",\n", " )\n", " )\n", "\n", "\n", "state = regular_expression_gen()\n", "print_highlight(state[\"answer\"])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Use `regex` to define a `JSON` decoding schema." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "character_regex = (\n", " r\"\"\"\\{\\n\"\"\"\n", " + r\"\"\" \"name\": \"[\\w\\d\\s]{1,16}\",\\n\"\"\"\n", " + r\"\"\" \"house\": \"(Gryffindor|Slytherin|Ravenclaw|Hufflepuff)\",\\n\"\"\"\n", " + r\"\"\" \"blood status\": \"(Pure-blood|Half-blood|Muggle-born)\",\\n\"\"\"\n", " + r\"\"\" \"occupation\": \"(student|teacher|auror|ministry of magic|death eater|order of the phoenix)\",\\n\"\"\"\n", " + r\"\"\" \"wand\": \\{\\n\"\"\"\n", " + r\"\"\" \"wood\": \"[\\w\\d\\s]{1,16}\",\\n\"\"\"\n", " + r\"\"\" \"core\": \"[\\w\\d\\s]{1,16}\",\\n\"\"\"\n", " + r\"\"\" \"length\": [0-9]{1,2}\\.[0-9]{0,2}\\n\"\"\"\n", " + r\"\"\" \\},\\n\"\"\"\n", " + r\"\"\" \"alive\": \"(Alive|Deceased)\",\\n\"\"\"\n", " + r\"\"\" \"patronus\": \"[\\w\\d\\s]{1,16}\",\\n\"\"\"\n", " + r\"\"\" \"bogart\": \"[\\w\\d\\s]{1,16}\"\\n\"\"\"\n", " + r\"\"\"\\}\"\"\"\n", ")\n", "\n", "\n", "@function\n", "def character_gen(s, name):\n", " s += user(\n", " f\"{name} is a character in Harry Potter. Please fill in the following information about this character.\"\n", " )\n", " s += assistant(gen(\"json_output\", max_tokens=256, regex=character_regex))\n", "\n", "\n", "state = character_gen(\"Harry Potter\")\n", "print_highlight(state[\"json_output\"])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Batching \n", "\n", "Use `run_batch` to run a batch of prompts." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@function\n", "def text_qa(s, question):\n", " s += user(question)\n", " s += assistant(gen(\"answer\", stop=\"\\n\"))\n", "\n", "\n", "states = text_qa.run_batch(\n", " [\n", " {\"question\": \"What is the capital of the United Kingdom?\"},\n", " {\"question\": \"What is the capital of France?\"},\n", " {\"question\": \"What is the capital of Japan?\"},\n", " ],\n", " progress_bar=True,\n", ")\n", "\n", "for i, state in enumerate(states):\n", " print_highlight(f\"Answer {i+1}: {states[i]['answer']}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Streaming \n", "\n", "Use `stream` to stream the output to the user." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@function\n", "def text_qa(s, question):\n", " s += user(question)\n", " s += assistant(gen(\"answer\", stop=\"\\n\"))\n", "\n", "\n", "state = text_qa.run(\n", " question=\"What is the capital of France?\", temperature=0.1, stream=True\n", ")\n", "\n", "for out in state.text_iter():\n", " print(out, end=\"\", flush=True)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Complex Prompts\n", "\n", "You may use `{system|user|assistant}_{begin|end}` to define complex prompts." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@function\n", "def chat_example(s):\n", " s += system(\"You are a helpful assistant.\")\n", " # Same as: s += s.system(\"You are a helpful assistant.\")\n", "\n", " with s.user():\n", " s += \"Question: What is the capital of France?\"\n", "\n", " s += assistant_begin()\n", " s += \"Answer: \" + gen(\"answer\", max_tokens=100, stop=\"\\n\")\n", " s += assistant_end()\n", "\n", "\n", "state = chat_example()\n", "print_highlight(state[\"answer\"])" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Multi-modal Generation\n", "\n", "You may use SGLang frontend language to define multi-modal prompts.\n", "See [here](https://docs.sglang.io/supported_models/text_generation/multimodal_language_models.html) for supported models." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "server_process, port = launch_server_cmd(\n", " \"python -m sglang.launch_server --model-path Qwen/Qwen2.5-VL-7B-Instruct --host 0.0.0.0 --log-level warning\"\n", ")\n", "\n", "wait_for_server(f\"http://localhost:{port}\", process=server_process)\n", "print(f\"Server started on http://localhost:{port}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "set_default_backend(RuntimeEndpoint(f\"http://localhost:{port}\"))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Ask a question about an image." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@function\n", "def image_qa(s, image_file, question):\n", " s += user(image(image_file) + question)\n", " s += assistant(gen(\"answer\", max_tokens=256))\n", "\n", "\n", "image_url = \"https://raw.githubusercontent.com/sgl-project/sglang/main/examples/assets/example_image.png\"\n", "image_bytes, _ = load_image(image_url)\n", "state = image_qa(image_bytes, \"What is in the image?\")\n", "print_highlight(state[\"answer\"])" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "terminate_process(server_process)" ] } ], "metadata": { "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 2 } ================================================ FILE: docs/references/learn_more.md ================================================ # Learn More and Join the Community - The development roadmap: [https://roadmap.sglang.io](https://roadmap.sglang.io) - Join weekly public development meeting: [https://meet.sglang.io](https://meet.sglang.io) - Join Slack: [https://slack.sglang.io/](https://slack.sglang.io/) - Follow on X (formerly Twitter): [https://x.com/lmsysorg](https://x.com/lmsysorg) - Follow on LinkedIn: [https://www.linkedin.com/company/sgl-project/](https://www.linkedin.com/company/sgl-project/) - The latest SGLang features and updates are shared through the [LMSYS blog](https://lmsys.org/blog/) - More blogs, slides, and videos about SGLang at [https://github.com/sgl-project/sgl-learning-materials](https://github.com/sgl-project/sgl-learning-materials) ================================================ FILE: docs/references/multi_node_deployment/deploy_on_k8s.md ================================================ # Deploy On Kubernetes This document is for deploying a RoCE network-based SGLang two-node inference service on a Kubernetes (K8S) cluster. [LeaderWorkerSet (LWS)](https://github.com/kubernetes-sigs/lws) is a Kubernetes API that aims to address common deployment patterns of AI/ML inference workloads. A major use case is for multi-host/multi-node distributed inference. SGLang can also be deployed with LWS on Kubernetes for distributed model serving. Please see this guide for more details on deploying SGLang on Kubernetes using LWS. Here we take the deployment of DeepSeek-R1 as an example. ## Prerequisites 1. At least two Kubernetes nodes, each with two H20 systems and eight GPUs, are required. 2. Make sure your K8S cluster has LWS correctly installed. If it hasn't been set up yet, please follow the [installation instructions](https://github.com/kubernetes-sigs/lws/blob/main/site/content/en/docs/installation/_index.md). **Note:** For LWS versions ≤0.5.x, you must use the Downward API to obtain `LWS_WORKER_INDEX`, as native support for this feature was introduced in v0.6.0. ## Basic example For the basic example documentation, refer to [Deploy Distributed Inference Service with SGLang and LWS on GPUs](https://github.com/kubernetes-sigs/lws/tree/main/docs/examples/sglang). However, that document only covers the basic NCCL socket mode. In this section, we’ll make some simple modifications to adapt the setup to the RDMA scenario. ## RDMA RoCE case * Check your env: ```bash [root@node1 ~]# ibstatus Infiniband device 'mlx5_bond_0' port 1 status: default gid: fe80:0000:0000:0000:0225:9dff:fe64:c79a base lid: 0x0 sm lid: 0x0 state: 4: ACTIVE phys state: 5: LinkUp rate: 200 Gb/sec (2X NDR) link_layer: Ethernet Infiniband device 'mlx5_bond_1' port 1 status: default gid: fe80:0000:0000:0000:0225:9dff:fe6e:c3ec base lid: 0x0 sm lid: 0x0 state: 4: ACTIVE phys state: 5: LinkUp rate: 200 Gb/sec (2X NDR) link_layer: Ethernet Infiniband device 'mlx5_bond_2' port 1 status: default gid: fe80:0000:0000:0000:0225:9dff:fe73:0dd7 base lid: 0x0 sm lid: 0x0 state: 4: ACTIVE phys state: 5: LinkUp rate: 200 Gb/sec (2X NDR) link_layer: Ethernet Infiniband device 'mlx5_bond_3' port 1 status: default gid: fe80:0000:0000:0000:0225:9dff:fe36:f7ff base lid: 0x0 sm lid: 0x0 state: 4: ACTIVE phys state: 5: LinkUp rate: 200 Gb/sec (2X NDR) link_layer: Ethernet ``` * Prepare the `lws.yaml` file for deploying on k8s. ```yaml apiVersion: leaderworkerset.x-k8s.io/v1 kind: LeaderWorkerSet metadata: name: sglang spec: replicas: 1 leaderWorkerTemplate: size: 2 restartPolicy: RecreateGroupOnPodRestart leaderTemplate: metadata: labels: role: leader spec: dnsPolicy: ClusterFirstWithHostNet hostNetwork: true hostIPC: true containers: - name: sglang-leader image: sglang:latest securityContext: privileged: true env: - name: NCCL_IB_GID_INDEX value: "3" command: - python3 - -m - sglang.launch_server - --model-path - /work/models - --mem-fraction-static - "0.93" - --torch-compile-max-bs - "8" - --max-running-requests - "20" - --tp - "16" # Size of Tensor Parallelism - --dist-init-addr - $(LWS_LEADER_ADDRESS):20000 - --nnodes - $(LWS_GROUP_SIZE) - --node-rank - $(LWS_WORKER_INDEX) - --trust-remote-code - --host - "0.0.0.0" - --port - "40000" resources: limits: nvidia.com/gpu: "8" ports: - containerPort: 40000 readinessProbe: tcpSocket: port: 40000 initialDelaySeconds: 15 periodSeconds: 10 volumeMounts: - mountPath: /dev/shm name: dshm - name: model mountPath: /work/models - name: ib mountPath: /dev/infiniband volumes: - name: dshm emptyDir: medium: Memory - name: model hostPath: path: '< your models dir >' # modify it according your models dir - name: ib hostPath: path: /dev/infiniband workerTemplate: spec: dnsPolicy: ClusterFirstWithHostNet hostNetwork: true hostIPC: true containers: - name: sglang-worker image: sglang:latest securityContext: privileged: true env: - name: NCCL_IB_GID_INDEX value: "3" command: - python3 - -m - sglang.launch_server - --model-path - /work/models - --mem-fraction-static - "0.93" - --torch-compile-max-bs - "8" - --max-running-requests - "20" - --tp - "16" # Size of Tensor Parallelism - --dist-init-addr - $(LWS_LEADER_ADDRESS):20000 - --nnodes - $(LWS_GROUP_SIZE) - --node-rank - $(LWS_WORKER_INDEX) - --trust-remote-code resources: limits: nvidia.com/gpu: "8" volumeMounts: - mountPath: /dev/shm name: dshm - name: model mountPath: /work/models - name: ib mountPath: /dev/infiniband volumes: - name: dshm emptyDir: medium: Memory - name: ib hostPath: path: /dev/infiniband - name: model hostPath: path: /data1/models/deepseek_v3_moe --- apiVersion: v1 kind: Service metadata: name: sglang-leader spec: selector: leaderworkerset.sigs.k8s.io/name: sglang role: leader ports: - protocol: TCP port: 40000 targetPort: 40000 ``` * Then use `kubectl apply -f lws.yaml` you will get this output. ```text NAME READY STATUS RESTARTS AGE sglang-0 0/1 Running 0 9s sglang-0-1 1/1 Running 0 9s ``` Wait for the sglang leader (`sglang-0`) status to change to 1/1, which indicates it is `Ready`. You can use the command `kubectl logs -f sglang-0` to view the logs of the leader node. Once successful, you should see output like this: ```text [2025-02-17 05:27:24 TP1] Capture cuda graph end. Time elapsed: 84.89 s [2025-02-17 05:27:24 TP6] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 [2025-02-17 05:27:24 TP0] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 [2025-02-17 05:27:24 TP7] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 [2025-02-17 05:27:24 TP3] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 [2025-02-17 05:27:24 TP2] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 [2025-02-17 05:27:24 TP4] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 [2025-02-17 05:27:24 TP1] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 [2025-02-17 05:27:24 TP5] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 [2025-02-17 05:27:24] INFO: Started server process [1] [2025-02-17 05:27:24] INFO: Waiting for application startup. [2025-02-17 05:27:24] INFO: Application startup complete. [2025-02-17 05:27:24] INFO: Uvicorn running on http://0.0.0.0:40000 (Press CTRL+C to quit) [2025-02-17 05:27:25] INFO: 127.0.0.1:48908 - "GET /get_model_info HTTP/1.1" 200 OK [2025-02-17 05:27:25 TP0] Prefill batch. #new-seq: 1, #new-token: 7, #cached-token: 0, cache hit rate: 0.00%, token usage: 0.00, #running-req: 0, #queue-req: 0 [2025-02-17 05:27:32] INFO: 127.0.0.1:48924 - "POST /generate HTTP/1.1" 200 OK [2025-02-17 05:27:32] The server is fired up and ready to roll! ``` If it doesn’t start up successfully, please follow these steps to check for any remaining issues. Thanks! ### Debug * Set `NCCL_DEBUG=TRACE` to check if it is a NCCL communication problem. This should resolve most NCCL-related issues. ***Notice: If you find that NCCL_DEBUG=TRACE is not effective in the container environment, but the process is stuck or you encounter hard-to-diagnose issues, try switching to a different container image. Some images may not handle standard error output properly.*** #### RoCE scenario * Please make sure that RDMA devices are available in the cluster environment. * Please make sure that the nodes in the cluster have Mellanox NICs with RoCE. In this example, we use Mellanox ConnectX 5 model NICs, and the proper OFED driver has been installed. If not, please refer to the document [Install OFED Driver](https://docs.nvidia.com/networking/display/mlnxofedv461000/installing+mellanox+ofed) to install the driver. * Check your env: ```shell $ lspci -nn | grep Eth | grep Mellanox 0000:7f:00.0 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) 0000:7f:00.1 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) 0000:c7:00.0 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) 0000:c7:00.1 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) 0001:08:00.0 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) 0001:08:00.1 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) 0001:a2:00.0 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) 0001:a2:00.1 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) ``` * Check the OFED driver: ```shell ofed_info -s OFED-internal-23.07-0.5.0: ``` * Show RDMA link status and check IB devices: ```shell $ rdma link show 8/1: mlx5_bond_0/1: state ACTIVE physical_state LINK_UP netdev reth0 9/1: mlx5_bond_1/1: state ACTIVE physical_state LINK_UP netdev reth2 10/1: mlx5_bond_2/1: state ACTIVE physical_state LINK_UP netdev reth4 11/1: mlx5_bond_3/1: state ACTIVE physical_state LINK_UP netdev reth6 $ ibdev2netdev 8/1: mlx5_bond_0/1: state ACTIVE physical_state LINK_UP netdev reth0 9/1: mlx5_bond_1/1: state ACTIVE physical_state LINK_UP netdev reth2 10/1: mlx5_bond_2/1: state ACTIVE physical_state LINK_UP netdev reth4 11/1: mlx5_bond_3/1: state ACTIVE physical_state LINK_UP netdev reth6 ``` * Test RoCE network speed on the host: ```shell yum install qperf # for server： execute qperf # for client qperf -t 60 -cm1 rc_rdma_write_bw ``` * Check RDMA accessible in your container: ```shell # ibv_devices # ibv_devinfo ``` ## Keys to success * In the YAML configuration above, pay attention to the NCCL environment variable. For older versions of NCCL, you should check the NCCL_IB_GID_INDEX environment setting. * NCCL_SOCKET_IFNAME is also crucial, but in a containerized environment, this typically isn’t an issue. * In some cases, it’s necessary to configure GLOO_SOCKET_IFNAME correctly. * NCCL_DEBUG is essential for troubleshooting, but I've found that sometimes it doesn't show error logs within containers. This could be related to the Docker image you're using. You may want to try switching images if needed. * Avoid using Docker images based on Ubuntu 18.04, as they tend to have compatibility issues. ## Remaining issues * In Kubernetes, Docker, or Containerd environments, we use hostNetwork to prevent performance degradation. * We utilize privileged mode, which isn’t secure. Additionally, in containerized environments, full GPU isolation cannot be achieved. ## TODO * Integrated with [k8s-rdma-shared-dev-plugin](https://github.com/Mellanox/k8s-rdma-shared-dev-plugin). ================================================ FILE: docs/references/multi_node_deployment/lws_pd/lws-examples/d-svc.yaml ================================================ apiVersion: v1 kind: Service metadata: name: deepseekr10528-decode-main spec: selector: leaderworkerset.sigs.k8s.io/name: deepseekr10528-decode-main role: leader ports: - protocol: TCP port: 30000 targetPort: 30000 ================================================ FILE: docs/references/multi_node_deployment/lws_pd/lws-examples/d.yaml ================================================ apiVersion: leaderworkerset.x-k8s.io/v1 kind: LeaderWorkerSet metadata: name: deepseekr10528-decode-main spec: leaderWorkerTemplate: leaderTemplate: metadata: labels: role: leader spec: containers: - command: - python3 - -m - sglang.launch_server - --port - "30000" - --host - "0.0.0.0" - --model-path - /work/models - --chunked-prefill-size - "262144" - --page-size - "64" - --enable-dp-attention - --enable-dp-lm-head - --dp-size - "16" - --moe-a2a-backend - deepep - --disaggregation-mode - decode - --mem-fraction-static - "0.849" - --context-length - "32768" - --disaggregation-ib-device - "mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3" - --cuda-graph-max-bs - "64" - --max-running-requests - "2048" - --tp-size - "16" # Size of Tensor Parallelism - --dist-init-addr - $(LWS_LEADER_ADDRESS):20102 - --nnodes - $(LWS_GROUP_SIZE) - --node-rank - $(LWS_WORKER_INDEX) - --trust-remote-code - --ep-num-redundant-experts - "32" - --moe-dense-tp-size - "1" env: - name: CUDA_LAUNCH_BLOCKING value: "0" - name: NVSHMEM_IB_GID_INDEX value: "3" - name: NVSHMEM_ENABLE_NIC_PE_MAPPING value: "1" - name: NVSHMEM_HCA_PE_MAPPING value: "mlx5_bond_0:1:2,mlx5_bond_1:1:2,mlx5_bond_2:1:2,mlx5_bond_3:1:2" - name: NCCL_IB_QPS_PER_CONNECTION value: "8" - name: NCCL_IB_SPLIT_DATA_ON_QPS value: "1" - name: NCCL_NET_PLUGIN value: "none" - name: NCCL_IB_TC value: "136" - name: NCCL_MIN_NCHANNELS value: "4" - name: NCCL_IB_SL value: "5" - name: MC_TE_METRIC value: "true" - name: SGLANG_MOONCAKE_TRANS_THREAD value: "16" - name: SGLANG_ENABLE_JIT_DEEPGEMM value: "1" - name: NCCL_IB_HCA value: ^=mlx5_0,mlx5_5,mlx5_6 - name: LWS_WORKER_INDEX valueFrom: fieldRef: fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] image: lmsysorg/sglang:latest name: sglang-leader ports: - containerPort: 30000 protocol: TCP readinessProbe: periodSeconds: 30 tcpSocket: port: 30000 resources: limits: nvidia.com/gpu: "8" securityContext: capabilities: add: - IPC_LOCK privileged: true volumeMounts: - mountPath: /root/.cache name: sgl-cache - mountPath: /dev/shm name: dshm - mountPath: /work/models name: model - mountPath: /dev/infiniband name: ib - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs name: cf dnsPolicy: ClusterFirstWithHostNet hostIPC: true hostNetwork: true nodeSelector: # should modify according your deployment env pd: "yes" tolerations: # should modify according your deployment env - key: bopd operator: Exists - key: node-role operator: Exists volumes: - hostPath: path: /data1/sgl_cache1 type: DirectoryOrCreate name: sgl-cache - emptyDir: medium: Memory name: dshm - hostPath: path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 name: model - hostPath: path: /dev/infiniband name: ib - hostPath: path: /data1/maas_hosted_models/models/fused_moe_triton/configs name: cf restartPolicy: RecreateGroupOnPodRestart size: 2 workerTemplate: metadata: {} spec: containers: - command: - python3 - -m - sglang.launch_server - --model-path - /work/models - --chunked-prefill-size - "262144" - --page-size - "64" - --enable-dp-attention - --enable-dp-lm-head - --dp-size - "16" - --moe-a2a-backend - deepep - --disaggregation-mode - decode - --mem-fraction-static - "0.849" - --context-length - "32768" - --disaggregation-ib-device - "mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3" - --cuda-graph-max-bs - "64" - --max-running-requests - "2048" - --tp-size - "16" # Size of Tensor Parallelism - --dist-init-addr - $(LWS_LEADER_ADDRESS):20102 - --nnodes - $(LWS_GROUP_SIZE) - --node-rank - $(LWS_WORKER_INDEX) - --trust-remote-code - --ep-num-redundant-experts - "32" - --moe-dense-tp-size - "1" env: - name: NVSHMEM_IB_TRAFFIC_CLASS value: "16" - name: NVSHMEM_IB_GID_INDEX value: "3" - name: NVSHMEM_ENABLE_NIC_PE_MAPPING value: "1" - name: NVSHMEM_HCA_PE_MAPPING value: "mlx5_bond_0:1:2,mlx5_bond_1:1:2,mlx5_bond_2:1:2,mlx5_bond_3:1:2" - name: NCCL_IB_QPS_PER_CONNECTION value: "8" - name: NCCL_IB_SPLIT_DATA_ON_QPS value: "1" - name: NCCL_NET_PLUGIN value: "none" - name: NCCL_IB_TC value: "136" - name: NCCL_MIN_NCHANNELS value: "4" - name: MC_TE_METRIC value: "true" - name: NCCL_IB_SL value: "5" - name: SGLANG_MOONCAKE_TRANS_THREAD value: "16" - name: SGLANG_ENABLE_JIT_DEEPGEMM value: "1" - name: NCCL_IB_HCA value: ^=mlx5_0,mlx5_5,mlx5_6 - name: LWS_WORKER_INDEX valueFrom: fieldRef: fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] image: lmsysorg/sglang:latest name: sglang-worker ports: - containerPort: 30001 resources: limits: nvidia.com/gpu: "8" securityContext: capabilities: add: - IPC_LOCK privileged: true volumeMounts: - mountPath: /root/.cache name: sgl-cache - mountPath: /dev/shm name: dshm - mountPath: /work/models name: model - mountPath: /dev/infiniband name: ib - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs name: cf dnsPolicy: ClusterFirstWithHostNet hostIPC: true hostNetwork: true nodeSelector: # should modify according your deployment env pd: "yes" tolerations: # should modify according your deployment env - key: bopd operator: Exists - key: node-role operator: Exists volumes: - hostPath: path: /data1/sgl_cache1 type: DirectoryOrCreate name: sgl-cache - emptyDir: medium: Memory name: dshm - hostPath: path: /dev/infiniband name: ib - hostPath: # modify according to you deployment env path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 name: model - hostPath: # modify according to you deployment env path: /data1/maas_hosted_models/models/fused_moe_triton/configs name: cf networkConfig: subdomainPolicy: Shared replicas: 1 rolloutStrategy: rollingUpdateConfiguration: maxSurge: 0 maxUnavailable: 1 type: RollingUpdate startupPolicy: LeaderCreated ================================================ FILE: docs/references/multi_node_deployment/lws_pd/lws-examples/lb.yaml ================================================ apiVersion: apps/v1 kind: Deployment metadata: name: deepseekr10528-lb-main labels: app: deepseekr10528-lb spec: replicas: 1 selector: matchLabels: app: deepseekr10528-lb template: metadata: labels: app: deepseekr10528-lb spec: nodeSelector: bo: "yes" tolerations: - key: bopd operator: Exists - key: node-role operator: Exists containers: - name: sgl-minilb image: lmsysorg/sglang:latest command: - python - -m - sglang_router.launch_router - --pd-disaggregation - --prefill - http://deepseekr10528-prefill-main:30000 - --decode - http://deepseekr10528-decode-main:30000 - --host - 0.0.0.0 - --port - "8000" ports: - containerPort: 8000 --- apiVersion: v1 kind: Service metadata: name: deepseekr10528-lb-service spec: type: NodePort # NodePort is easy to test, you can also specify `ClusterIP` selector: app: deepseekr10528-lb ports: - protocol: TCP port: 8000 # Service Port（In-Cluster） targetPort: 8000 # Exposed Container nodePort: 30800 ================================================ FILE: docs/references/multi_node_deployment/lws_pd/lws-examples/p-svc.yaml ================================================ apiVersion: v1 kind: Service metadata: name: deepseekr10528-prefill-main spec: selector: leaderworkerset.sigs.k8s.io/name: deepseekr10528-prefill-main role: leader ports: - protocol: TCP port: 30000 targetPort: 30000 ================================================ FILE: docs/references/multi_node_deployment/lws_pd/lws-examples/p.yaml ================================================ apiVersion: leaderworkerset.x-k8s.io/v1 kind: LeaderWorkerSet metadata: name: deepseekr10528-prefill-main spec: leaderWorkerTemplate: leaderTemplate: metadata: labels: role: leader spec: containers: - command: - python3 - -m - sglang.launch_server - --port - "30000" - --host - "0.0.0.0" - --model-path - /work/models - --disaggregation-ib-device # should modify according your rdma env - mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 - --chunked-prefill-size - "524288" - --max-prefill-tokens - "32768" - --page-size - "64" - --ep-dispatch-algorithm - dynamic - --eplb-algorithm - deepseek - --enable-dp-lm-head - --enable-dp-attention - --dp-size - "16" - --disable-radix-cache - --moe-a2a-backend - deepep - --disaggregation-mode - prefill - --mem-fraction-static - "0.7" - --context-length - "32768" - --tp - "16" - --dist-init-addr - $(LWS_LEADER_ADDRESS):20102 - --nnodes - $(LWS_GROUP_SIZE) - --node-rank - $(LWS_WORKER_INDEX) - --trust-remote-code - --ep-num-redundant-experts - "32" - --moe-dense-tp-size - "1" - --max-running-requests - "1024" env: - name: NVSHMEM_HCA_PE_MAPPING # should modify according your rdma env value: "mlx5_bond_0:1:2,mlx5_bond_1:1:2,mlx5_bond_2:1:2,mlx5_bond_3:1:2" - name: NVSHMEM_IB_GID_INDEX value: "3" - name: NVSHMEM_ENABLE_NIC_PE_MAPPING value: "1" - name: SGLANG_SET_CPU_AFFINITY value: "true" - name: SGLANG_ENABLE_JIT_DEEPGEMM value: "1" - name: NCCL_IB_QPS_PER_CONNECTION value: "8" - name: NCCL_IB_SPLIT_DATA_ON_QPS value: "1" - name: NCCL_NET_PLUGIN value: none - name: NCCL_IB_TC value: "136" - name: NCCL_MIN_NCHANNELS value: "4" - name: MC_TE_METRIC value: "false" - name: NCCL_IB_SL value: "5" - name: NCCL_IB_HCA value: ^=mlx5_0,mlx5_5,mlx5_6 - name: LWS_WORKER_INDEX valueFrom: fieldRef: fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] image: lmsysorg/sglang:latest name: sglang-leader ports: - containerPort: 30000 protocol: TCP readinessProbe: periodSeconds: 30 tcpSocket: port: 30000 resources: limits: nvidia.com/gpu: "8" securityContext: capabilities: add: - IPC_LOCK privileged: true volumeMounts: - mountPath: /dev/shm name: dshm - mountPath: /work/models name: model - mountPath: /dev/infiniband name: ib - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs name: cf - mountPath: /root/.cache name: sgl-cache dnsPolicy: ClusterFirstWithHostNet hostIPC: true hostNetwork: true nodeSelector: # should modify according your deployment env pd: "yes" tolerations: # should modify according your deployment env - key: bopd operator: Exists - key: node-role operator: Exists volumes: - emptyDir: medium: Memory name: dshm - hostPath: path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 name: model - hostPath: path: /dev/infiniband name: ib - hostPath: path: /data1/maas_hosted_models/models/fused_moe_triton/configs name: cf - hostPath: path: /data1/sgl_cache type: DirectoryOrCreate name: sgl-cache restartPolicy: RecreateGroupOnPodRestart size: 2 workerTemplate: metadata: {} spec: containers: - command: - python3 - -m - sglang.launch_server - --model-path - /work/models - --disaggregation-ib-device # should modify according your rdma env - mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 - --chunked-prefill-size - "524288" - --max-prefill-tokens - "32768" - --page-size - "64" - --ep-dispatch-algorithm - dynamic - --eplb-algorithm - deepseek # - --deepep-config # - /home/aiges/tuned/tuned_8sms.json # can be tuned using deepep test scripts - --enable-dp-lm-head - --enable-dp-attention - --dp-size - "16" - --disable-radix-cache - --moe-a2a-backend - deepep - --disaggregation-mode - prefill - --mem-fraction-static - "0.7" - --context-length - "32768" - --tp - "16" - --dist-init-addr - $(LWS_LEADER_ADDRESS):20102 - --nnodes - $(LWS_GROUP_SIZE) - --node-rank - $(LWS_WORKER_INDEX) - --trust-remote-code - --ep-num-redundant-experts - "32" - --moe-dense-tp-size - "1" - --max-running-requests - "1024" env: - name: SGLANG_SET_CPU_AFFINITY value: "true" - name: NVSHMEM_HCA_PE_MAPPING # should modify according your rdma env value: "mlx5_bond_0:1:2,mlx5_bond_1:1:2,mlx5_bond_2:1:2,mlx5_bond_3:1:2" - name: NCCL_IB_HCA value: ^=mlx5_0,mlx5_5,mlx5_6 - name: NVSHMEM_IB_TRAFFIC_CLASS value: "16" - name: NVSHMEM_IB_GID_INDEX value: "3" - name: NVSHMEM_ENABLE_NIC_PE_MAPPING value: "1" - name: CUDA_LAUNCH_BLOCKING value: "0" - name: SGLANG_MOONCAKE_TRANS_THREAD value: "8" - name: SGLANG_ENABLE_JIT_DEEPGEMM value: "1" - name: SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD value: "0" - name: NCCL_IB_QPS_PER_CONNECTION value: "8" - name: NCCL_IB_SPLIT_DATA_ON_QPS value: "1" - name: NCCL_NET_PLUGIN value: none - name: NCCL_IB_TC value: "136" - name: NCCL_MIN_NCHANNELS value: "4" - name: MC_TE_METRIC value: "true" - name: NCCL_IB_SL value: "5" - name: LWS_WORKER_INDEX valueFrom: fieldRef: fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] image: lmsysorg/sglang:latest name: sglang-worker ports: - containerPort: 30001 protocol: TCP resources: limits: nvidia.com/gpu: "8" securityContext: capabilities: add: - IPC_LOCK privileged: true volumeMounts: - mountPath: /root/.cache name: sgl-cache - mountPath: /dev/shm name: dshm - mountPath: /work/models name: model - mountPath: /dev/infiniband name: ib - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs name: cf dnsPolicy: ClusterFirstWithHostNet hostIPC: true hostNetwork: true nodeSelector: # should modify according your deployment env pd: "yes" tolerations: # should modify according your deployment env - key: bopd operator: Exists - key: node-role operator: Exists volumes: - emptyDir: medium: Memory name: dshm - hostPath: path: /dev/infiniband name: ib - hostPath: # modify according to you deployment env path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 name: model - hostPath: # modify according to you deployment env path: /data1/maas_hosted_models/models/fused_moe_triton/configs name: cf - hostPath: # modify according to you deployment env path: /data1/sgl_cache type: DirectoryOrCreate name: sgl-cache ================================================ FILE: docs/references/multi_node_deployment/lws_pd/lws_pd_deploy.md ================================================ # LWS Based PD Deploy ## 0. Prerequisites 1. k8s >=1.26 2. lws installed on k8s. ## 1. Image Preparation `lmsysorg/sglang:deepep` ## 2. Deployment Manifest Files ***Notice: We will package all deployment files into Helm Chart format in the near future. Interested community members can contact us to contribute*** ### Prefill Prefill manifest file [prefill.yaml](lws-examples/p.yaml) *Note: The NodeSelector section, model location section, and taint toleration section can be adjusted according to your actual deployment environment* ```yaml apiVersion: leaderworkerset.x-k8s.io/v1 kind: LeaderWorkerSet metadata: name: deepseekr10528-prefill-main spec: leaderWorkerTemplate: leaderTemplate: metadata: labels: role: leader spec: containers: - command: - python3 - -m - sglang.launch_server - --port - "30000" - --host - "0.0.0.0" - --model-path - /work/models - --disaggregation-ib-device # should modify according your rdma env - mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 - --chunked-prefill-size - "524288" - --max-prefill-tokens - "32768" - --page-size - "64" # - --init-expert-location # - /home/aiges/tuned/attachment_ep_statistics/prefill_in1024.json - --ep-dispatch-algorithm - dynamic - --eplb-algorithm - deepseek # - --deepep-config # - /home/aiges/tuned/tuned_8sms.json - --enable-dp-lm-head - --enable-dp-attention - --dp-size - "16" - --disable-radix-cache - --moe-a2a-backend - deepep - --disaggregation-mode - prefill - --mem-fraction-static - "0.7" - --context-length - "32768" - --tp - "16" - --dist-init-addr - $(LWS_LEADER_ADDRESS):20102 - --nnodes - $(LWS_GROUP_SIZE) - --node-rank - $(LWS_WORKER_INDEX) - --trust-remote-code - --ep-num-redundant-experts - "32" - --moe-dense-tp-size - "1" - --max-running-requests - "1024" env: # - name: NVSHMEM_HCA_PE_MAPPING # value: "mlx5_bond_0:1:2,mlx5_bond_1:1:2,mlx5_bond_2:1:2,mlx5_bond_3:1:2" # - name: NVSHMEM_HCA_LIST # value: "mlx5_bond_0:1,mlx5_bond_1:1,mlx5_bond_2:1,mlx5_bond_3:1" - name: NVSHMEM_IB_GID_INDEX value: "3" - name: NVSHMEM_ENABLE_NIC_PE_MAPPING value: "1" - name: SGLANG_SET_CPU_AFFINITY value: "true" - name: SGLANG_ENABLE_JIT_DEEPGEMM value: "1" - name: NCCL_IB_QPS_PER_CONNECTION value: "8" - name: NCCL_IB_SPLIT_DATA_ON_QPS value: "1" - name: NCCL_NET_PLUGIN value: none - name: NCCL_IB_TC value: "136" - name: NCCL_MIN_NCHANNELS value: "4" - name: MC_TE_METRIC value: "false" - name: NCCL_IB_SL value: "5" - name: NCCL_IB_HCA value: ^=mlx5_0,mlx5_5,mlx5_6 - name: LWS_WORKER_INDEX valueFrom: fieldRef: fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] image: lmsysorg/sglang:deepep name: sglang-leader ports: - containerPort: 30000 protocol: TCP readinessProbe: periodSeconds: 30 tcpSocket: port: 30000 resources: limits: nvidia.com/gpu: "8" securityContext: capabilities: add: - IPC_LOCK privileged: true volumeMounts: - mountPath: /dev/shm name: dshm - mountPath: /work/models name: model - mountPath: /dev/infiniband name: ib - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs name: cf - mountPath: /root/.cache name: sgl-cache dnsPolicy: ClusterFirstWithHostNet hostIPC: true hostNetwork: true nodeSelector: pd: "yes" tolerations: - key: pd operator: Exists - key: node-role operator: Exists volumes: - emptyDir: medium: Memory name: dshm - hostPath: # modify according to you deployment env path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 name: model - hostPath: path: /dev/infiniband name: ib - hostPath: # modify according to you deployment env path: /data1/maas_hosted_models/models/fused_moe_triton/configs name: cf - hostPath: # modify according to you deployment env path: /data1/sgl_cache type: DirectoryOrCreate name: sgl-cache restartPolicy: RecreateGroupOnPodRestart size: 2 workerTemplate: metadata: {} spec: containers: - command: - python3 - -m - sglang.launch_server - --model-path - /work/models - --disaggregation-ib-device - mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 - --chunked-prefill-size - "524288" - --max-prefill-tokens - "32768" - --page-size - "64" #- --init-expert-location #- /home/aiges/tuned/attachment_ep_statistics/prefill_in1024.json - --ep-dispatch-algorithm - dynamic - --eplb-algorithm - deepseek # - --deepep-config # - /home/aiges/tuned/tuned_8sms.json - --enable-dp-lm-head - --enable-dp-attention - --dp-size - "16" - --disable-radix-cache - --moe-a2a-backend - deepep - --disaggregation-mode - prefill - --mem-fraction-static - "0.7" - --context-length - "32768" - --tp - "16" - --dist-init-addr - $(LWS_LEADER_ADDRESS):20102 - --nnodes - $(LWS_GROUP_SIZE) - --node-rank - $(LWS_WORKER_INDEX) - --trust-remote-code - --ep-num-redundant-experts - "32" - --moe-dense-tp-size - "1" - --max-running-requests - "1024" env: - name: SGLANG_SET_CPU_AFFINITY value: "true" - name: SGLANG_HACK_DEEPEP_NUM_SMS value: "8" - name: SGLANG_HACK_DEEPEP_NEW_MODE value: "0" # - name: NVSHMEM_HCA_PE_MAPPING # value: "mlx5_bond_0:1:2,mlx5_bond_1:1:2,mlx5_bond_2:1:2,mlx5_bond_3:1:2" # - name: NVSHMEM_HCA_LIST # value: "mlx5_bond_0:1,mlx5_bond_1:1,mlx5_bond_2:1,mlx5_bond_3:1" - name: NCCL_IB_HCA value: ^=mlx5_0,mlx5_5,mlx5_6 - name: NVSHMEM_IB_TRAFFIC_CLASS value: "16" - name: NVSHMEM_IB_GID_INDEX value: "3" - name: NVSHMEM_ENABLE_NIC_PE_MAPPING value: "1" - name: CUDA_LAUNCH_BLOCKING value: "0" - name: SGLANG_MOONCAKE_TRANS_THREAD value: "8" - name: SGLANG_ENABLE_JIT_DEEPGEMM value: "1" - name: SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD value: "0" - name: NCCL_IB_QPS_PER_CONNECTION value: "8" - name: NCCL_IB_SPLIT_DATA_ON_QPS value: "1" - name: NCCL_NET_PLUGIN value: none - name: NCCL_IB_TC value: "136" - name: NCCL_MIN_NCHANNELS value: "4" - name: MC_TE_METRIC value: "true" - name: NCCL_IB_SL value: "5" - name: LWS_WORKER_INDEX valueFrom: fieldRef: fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] image: lmsysorg/sglang:deepep name: sglang-worker ports: - containerPort: 30001 protocol: TCP resources: limits: nvidia.com/gpu: "8" securityContext: capabilities: add: - IPC_LOCK privileged: true volumeMounts: - mountPath: /root/.cache name: sgl-cache - mountPath: /dev/shm name: dshm - mountPath: /work/models name: model - mountPath: /dev/infiniband name: ib - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs name: cf dnsPolicy: ClusterFirstWithHostNet hostIPC: true hostNetwork: true nodeSelector: pd: "yes" tolerations: - key: pd operator: Exists - key: node-role operator: Exists volumes: - emptyDir: medium: Memory name: dshm - hostPath: path: /dev/infiniband name: ib - hostPath: path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 name: model - hostPath: path: /data1/maas_hosted_models/models/fused_moe_triton/configs name: cf - hostPath: path: /data1/sgl_cache type: DirectoryOrCreate name: sgl-cache ``` ### Decode Decode node deployment manifest file [decode.yaml](lws-examples/d.yaml) *Note: The NodeSelector section, model location section, and taint toleration section can be adjusted according to your actual deployment environment* ```yaml apiVersion: leaderworkerset.x-k8s.io/v1 kind: LeaderWorkerSet metadata: name: deepseekr10528-decode-main spec: leaderWorkerTemplate: leaderTemplate: metadata: labels: role: leader spec: containers: - command: - python3 - -m - sglang.launch_server - --port - "30000" - --host - "0.0.0.0" - --model-path - /work/models - --chunked-prefill-size - "262144" - --page-size - "64" - --enable-dp-attention - --enable-dp-lm-head - --dp-size - "16" - --moe-a2a-backend - deepep - --disaggregation-mode - decode - --mem-fraction-static - "0.849" - --context-length - "32768" - --disaggregation-ib-device - "mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3" - --cuda-graph-max-bs - "64" - --max-running-requests - "2048" - --tp-size - "16" # Size of Tensor Parallelism - --dist-init-addr - $(LWS_LEADER_ADDRESS):20102 - --nnodes - $(LWS_GROUP_SIZE) - --node-rank - $(LWS_WORKER_INDEX) - --trust-remote-code - --ep-num-redundant-experts - "32" - --moe-dense-tp-size - "1" env: - name: CUDA_LAUNCH_BLOCKING value: "0" - name: NVSHMEM_IB_GID_INDEX value: "3" - name: NVSHMEM_ENABLE_NIC_PE_MAPPING value: "1" - name: NCCL_IB_QPS_PER_CONNECTION value: "8" - name: NCCL_IB_SPLIT_DATA_ON_QPS value: "1" - name: NCCL_NET_PLUGIN value: "none" - name: NCCL_IB_TC value: "136" - name: NCCL_MIN_NCHANNELS value: "4" - name: NCCL_IB_SL value: "5" - name: MC_TE_METRIC value: "true" - name: SGLANG_MOONCAKE_TRANS_THREAD value: "16" - name: SGLANG_ENABLE_JIT_DEEPGEMM value: "1" - name: NCCL_IB_HCA value: ^=mlx5_0,mlx5_5,mlx5_6 - name: LWS_WORKER_INDEX valueFrom: fieldRef: fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] image: lmsysorg/sglang:deepep name: sglang-leader ports: - containerPort: 30000 protocol: TCP readinessProbe: periodSeconds: 30 tcpSocket: port: 30000 resources: limits: nvidia.com/gpu: "8" securityContext: capabilities: add: - IPC_LOCK privileged: true volumeMounts: - mountPath: /root/.cache name: sgl-cache - mountPath: /dev/shm name: dshm - mountPath: /work/models name: model - mountPath: /dev/infiniband name: ib - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs name: cf dnsPolicy: ClusterFirstWithHostNet hostIPC: true hostNetwork: true nodeSelector: pd: "yes" tolerations: - key: pd operator: Exists - key: node-role operator: Exists volumes: - hostPath: path: /data1/sgl_cache1 type: DirectoryOrCreate name: sgl-cache - emptyDir: medium: Memory name: dshm - hostPath: path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 name: model - hostPath: path: /dev/infiniband name: ib - hostPath: path: /data1/maas_hosted_models/models/fused_moe_triton/configs name: cf restartPolicy: RecreateGroupOnPodRestart size: 2 workerTemplate: metadata: {} spec: containers: - command: - python3 - -m - sglang.launch_server - --model-path - /work/models - --chunked-prefill-size - "262144" - --page-size - "64" - --enable-dp-attention - --enable-dp-lm-head #- --enable-two-batch-overlap - --dp-size - "16" - --moe-a2a-backend - deepep - --disaggregation-mode - decode - --mem-fraction-static - "0.849" - --context-length - "32768" - --disaggregation-ib-device # should modify according your rdma env - "mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3" - --cuda-graph-max-bs - "64" - --max-running-requests - "2048" - --tp-size - "16" # Size of Tensor Parallelism - --dist-init-addr - $(LWS_LEADER_ADDRESS):20102 - --nnodes - $(LWS_GROUP_SIZE) - --node-rank - $(LWS_WORKER_INDEX) - --trust-remote-code - --ep-num-redundant-experts - "32" - --moe-dense-tp-size - "1" env: - name: SGLANG_HACK_DEEPEP_NUM_SMS value: "24" - name: SGLANG_HACK_DEEPEP_NEW_MODE value: "0" - name: NVSHMEM_IB_TRAFFIC_CLASS value: "16" - name: NVSHMEM_IB_GID_INDEX value: "3" - name: NVSHMEM_ENABLE_NIC_PE_MAPPING value: "1" - name: NCCL_IB_QPS_PER_CONNECTION value: "8" - name: NCCL_IB_SPLIT_DATA_ON_QPS value: "1" - name: NCCL_NET_PLUGIN value: "none" - name: NCCL_IB_TC value: "136" - name: NCCL_MIN_NCHANNELS value: "4" - name: MC_TE_METRIC value: "true" - name: NCCL_IB_SL value: "5" - name: SGLANG_MOONCAKE_TRANS_THREAD value: "16" - name: SGLANG_ENABLE_JIT_DEEPGEMM value: "1" - name: NCCL_IB_HCA value: ^=mlx5_0,mlx5_5,mlx5_6 - name: LWS_WORKER_INDEX valueFrom: fieldRef: fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] image: lmsysorg/sglang:deepep name: sglang-worker ports: - containerPort: 30001 resources: limits: nvidia.com/gpu: "8" securityContext: capabilities: add: - IPC_LOCK privileged: true volumeMounts: - mountPath: /root/.cache name: sgl-cache - mountPath: /dev/shm name: dshm - mountPath: /work/models name: model - mountPath: /dev/infiniband name: ib - mountPath: /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs name: cf dnsPolicy: ClusterFirstWithHostNet hostIPC: true hostNetwork: true nodeSelector: pd: "yes" tolerations: - key: pd operator: Exists - key: node-role operator: Exists volumes: - hostPath: path: /data1/sgl_cache1 type: DirectoryOrCreate name: sgl-cache - emptyDir: medium: Memory name: dshm - hostPath: path: /dev/infiniband name: ib - hostPath: # modify according to you deployment env path: /data1/maas_hosted_models/models/DeepSeek-R1-0528/deepseek_r1_0528 name: model - hostPath: # modify according to you deployment env path: /data1/maas_hosted_models/models/fused_moe_triton/configs name: cf networkConfig: subdomainPolicy: Shared replicas: 1 rolloutStrategy: rollingUpdateConfiguration: maxSurge: 0 maxUnavailable: 1 type: RollingUpdate startupPolicy: LeaderCreated ``` Execute separately: ```bash kubectl apply -f p.yaml kubectl apply -f d.yaml ``` At this point, we have completed the deployment of the 1P1D SGLang engine part. To allow our users to directly experience the model API, we still need a load balancer to handle sequential calls between prefill and decode. Different companies implement LBs differently, and the community will also officially release a new LB component written in Rust in the near future. Currently, we use a static K8S service + minilb approach to implement model API calls. ### Creating Service for Prefill and Decode #### Create prefill k8s service [p-svc.yaml](lws-examples/p-svc.yaml) ```yaml apiVersion: v1 kind: Service metadata: name: deepseekr10528-prefill-main spec: selector: leaderworkerset.sigs.k8s.io/name: deepseekr10528-prefill-main role: leader ports: - protocol: TCP port: 30000 targetPort: 30000 ``` Execute `kubectl apply -f p-svc.yaml` #### Create decode k8s service [d-svc.yaml](lws-examples/d-svc.yaml) ```yaml apiVersion: v1 kind: Service metadata: name: deepseekr10528-decode-main spec: selector: leaderworkerset.sigs.k8s.io/name: deepseekr10528-decode-main role: leader ports: - protocol: TCP port: 30000 targetPort: 30000 ``` Execute `kubectl apply -f d-svc.yaml` #### Deploy minilb and lb service [lb.yaml](lws-examples/lb.yaml) ```yaml apiVersion: apps/v1 kind: Deployment metadata: name: deepseekr10528-lb-main labels: app: deepseekr10528-lb spec: replicas: 1 selector: matchLabels: app: deepseekr10528-lb template: metadata: labels: app: deepseekr10528-lb spec: nodeSelector: pd: "yes" tolerations: - key: pd operator: Exists - key: node-role operator: Exists containers: - name: sgl-minilb image: lmsysorg/sglang:deepep command: - python - -m - sglang_router.launch_router - --pd-disaggregation - --prefill - http://deepseekr10528-prefill-main:30000 - --decode - http://deepseekr10528-decode-main:30000 - --host - 0.0.0.0 - --port - "8000" ports: - containerPort: 8000 --- apiVersion: v1 kind: Service metadata: name: deepseekr10528-lb-service spec: type: NodePort selector: app: deepseekr10528-lb ports: - protocol: TCP port: 8000 # Service Port（In-Cluster） targetPort: 8000 # Exposed Container nodePort: 30800 ``` Execute `kubectl apply -f lb.yaml` After waiting for all model deployments to succeed, you will get the following output: ```bash [root@ecs-001]# kubectl get po deepseekr10528-decode-main-0 1/1 Running 0 74m deepseekr10528-decode-main-0-1 1/1 Running 0 74m deepseekr10528-lb-main-9c5dbfc57-6lcbd 1/1 Running 0 22m deepseekr10528-prefill-main-0 1/1 Running 0 74m deepseekr10528-prefill-main-0-1 1/1 Running 0 74m [root@ecs-cbm-x1-pd-cpu-001 main_doc]# kubectl get svc |grep dee deepseekr10528-decode-main ClusterIP None 97m deepseekr10528-lb-service NodePort 172.16.242.169 8000:30800/TCP 22m deepseekr10528-prefill-main ClusterIP None 97m ``` At this point, select a nodePort:30800 to access: ```bash [root@ecs-001]# curl -X POST "http://{nodePort}:30800/v1/chat/completions" \ > -H "Content-Type: application/json" \ > -H "Authorization: Bearer None" \ > -d '{ > "rid":"ccccdd", > "model": "r1", > "messages": [ > {"role": "system", "content": "0: You are a helpful AI assistant"}, > {"role": "user", "content": "你是谁？."} > ], > "max_tokens":221 > }' {"id":"ccccdd","object":"chat.completion","created":1750252498,"model":"qwen2","choices":[{"index":0,"message":{"role":"assistant","content":"\n嗯，用户问了一个很基础的自我介绍问题"你是谁？"。这可能是第一次互动时的常规开场白，也可能是想确认我的身份和功能范围。\n\n用户没有提供任何背景信息，语气简洁中性。这种场景下新用户的可能性较高，需要给出清晰友好的自我介绍，同时突出实用价值来降低陌生感。\n\n考虑到中文用户，应该用简体中文回复。重点要说明三点：身份归属（深度求索）、功能定位（AI助手）、服务范围（学习/工作/生活）。结尾用开放性问题引导对话很关键——既能了解需求，又能避免让用户面对空白输入框时不知所措。\n\n用波浪线结尾可以软化语气，那个笑脸表情😊刚好能中和AI的机械感。不过要控制表情符号数量，避免显得轻浮。\n\n你好呀！我是你的AI助手，由深度求索公司（DeepSeek）开发的语言模型，名字叫 **DeepSeek-R1**。你可以把我当成一个知识丰富、随叫随到的小帮手～😊\n\n我的任务就是陪你聊天、解答问题、","reasoning_content":null,"tool_calls":null},"logprobs":null,"finish_reason":"length","matched_stop":null}],"usage":{"prompt_tokens":14,"total_tokens":235,"completion_tokens":221,"prompt_tokens_details":null}} ``` ## FAQ 1. The current deployment startup parameters may not be fully compatible with all RDMA scenarios. Different RDMA NCCL-related environment configurations may be needed in different network environments. 2. Some preset, optimized configurations for EPLB are not used here. You can adjust them according to [6017](https://github.com/sgl-project/sglang/issues/6017) as needed. ================================================ FILE: docs/references/multi_node_deployment/multi_node.md ================================================ # Multi-Node Deployment ## Llama 3.1 405B **Run 405B (fp16) on Two Nodes** ```bash # replace 172.16.4.52:20000 with your own node ip address and port of the first node python3 -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3.1-405B-Instruct \ --tp 16 \ --dist-init-addr 172.16.4.52:20000 \ --nnodes 2 \ --node-rank 0 python3 -m sglang.launch_server \ --model-path meta-llama/Meta-Llama-3.1-405B-Instruct \ --tp 16 \ --dist-init-addr 172.16.4.52:20000 \ --nnodes 2 \ --node-rank 1 ``` Note that LLama 405B (fp8) can also be launched on a single node. ```bash python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-405B-Instruct-FP8 --tp 8 ``` ## DeepSeek V3/R1 Please refer to [DeepSeek documents for reference](https://docs.sglang.io/basic_usage/deepseek.html#running-examples-on-multi-node). ## Multi-Node Inference on SLURM This example showcases how to serve SGLang server across multiple nodes by SLURM. Submit the following job to the SLURM cluster. ``` #!/bin/bash -l #SBATCH -o SLURM_Logs/%x_%j_master.out #SBATCH -e SLURM_Logs/%x_%j_master.err #SBATCH -D ./ #SBATCH -J Llama-405B-Online-Inference-TP16-SGL #SBATCH --nodes=2 #SBATCH --ntasks=2 #SBATCH --ntasks-per-node=1 # Ensure 1 task per node #SBATCH --cpus-per-task=18 #SBATCH --mem=224GB #SBATCH --partition="lmsys.org" #SBATCH --gres=gpu:8 #SBATCH --time=12:00:00 echo "[INFO] Activating environment on node $SLURM_PROCID" if ! source ENV_FOLDER/bin/activate; then echo "[ERROR] Failed to activate environment" >&2 exit 1 fi # Define parameters model=MODEL_PATH tp_size=16 echo "[INFO] Running inference" echo "[INFO] Model: $model" echo "[INFO] TP Size: $tp_size" # Set NCCL initialization address using the hostname of the head node HEAD_NODE=$(scontrol show hostname "$SLURM_NODELIST" | head -n 1) NCCL_INIT_ADDR="${HEAD_NODE}:8000" echo "[INFO] NCCL_INIT_ADDR: $NCCL_INIT_ADDR" # Launch the model server on each node using SLURM srun --ntasks=2 --nodes=2 --output="SLURM_Logs/%x_%j_node$SLURM_NODEID.out" \ --error="SLURM_Logs/%x_%j_node$SLURM_NODEID.err" \ python3 -m sglang.launch_server \ --model-path "$model" \ --grammar-backend "xgrammar" \ --tp "$tp_size" \ --dist-init-addr "$NCCL_INIT_ADDR" \ --nnodes 2 \ --node-rank "$SLURM_NODEID" & # Wait for the NCCL server to be ready on port 30000 while ! nc -z "$HEAD_NODE" 30000; do sleep 1 echo "[INFO] Waiting for $HEAD_NODE:30000 to accept connections" done echo "[INFO] $HEAD_NODE:30000 is ready to accept connections" # Keep the script running until the SLURM job times out wait ``` Then, you can test the server by sending requests following other [documents](https://docs.sglang.io/basic_usage/openai_api_completions.html). Thanks for [aflah02](https://github.com/aflah02) for providing the example, based on his [blog post](https://aflah02.substack.com/p/multi-node-llm-inference-with-sglang). ================================================ FILE: docs/references/multi_node_deployment/multi_node_index.rst ================================================ Multi-Node Deployment ===================== .. toctree:: :maxdepth: 1 :caption: Multi-Node Deployment multi_node.md deploy_on_k8s.md lws_pd/lws_pd_deploy.md rbg_pd/deepseekv32_pd.md - `Deploying DeepSeek with PD Disaggregation and Large-Scale Expert Parallelism on 96 H100 GPUs `_ - `Deploying Kimi K2 with PD Disaggregation and Large-Scale Expert Parallelism on 128 H200 GPUs `_ ================================================ FILE: docs/references/multi_node_deployment/rbg_pd/deepseekv32_pd.md ================================================ # DeepSeekV32-Exp RBG Based PD Deploy ## 0. Prerequisites 1. k8s >=1.26 2. lws installed on k8s. 3. rbg installed on k8s. For RBG installation, please refer to: https://github.com/sgl-project/rbg ## 1. Image Preparation `lmsysorg/sglang:latest` ### 2. All In One manifest file *Note: The NodeSelector section, model location section, and taint toleration section can be adjusted according to your actual deployment environment* rbg-dsv32.yml ```yaml apiVersion: workloads.x-k8s.io/v1alpha1 kind: RoleBasedGroup metadata: name: deepseek-rbg-32exp namespace: default spec: roles: - name: prefill replicas: 1 workload: apiVersion: leaderworkerset.x-k8s.io/v1 kind: LeaderWorkerSet restartPolicy: None leaderWorkerSet: size: 1 patchLeaderTemplate: metadata: labels: role: leader pd_role: prefill spec: containers: - command: - python3 - -m - sglang.launch_server - --model-path - /work/models - --port - "30000" - --trust-remote - --host - 0.0.0.0 - --disaggregation-ib-device - mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7 - --disable-radix-cache - --chunked-prefill-size - "131072" - --page-size - "64" # - --enable-eplb - --ep-dispatch-algorithm - dynamic - --eplb-algorithm - deepseek - --enable-dp-lm-head - --enable-dp-attention - --dp-size - "8" - --moe-a2a-backend - deepep - --deepep-mode - normal - --disaggregation-mode - prefill - --mem-fraction-static - "0.8" - --max-prefill-tokens - "32768" - --context-length - "32768" - --tp - "8" - --dist-init-addr - $(LWS_LEADER_ADDRESS):20102 - --nnodes - $(LWS_GROUP_SIZE) - --node-rank - $(LWS_WORKER_INDEX) - --trust-remote-code - --ep-num-redundant-experts - "32" - --moe-dense-tp-size - "1" - --max-running-requests - "1024" env: - name: LWS_WORKER_INDEX valueFrom: fieldRef: fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] livenessProbe: failureThreshold: 3000 httpGet: path: /health port: 30000 initialDelaySeconds: 300 periodSeconds: 60 successThreshold: 1 timeoutSeconds: 10 readinessProbe: failureThreshold: 20 httpGet: path: /health port: 30000 periodSeconds: 30 successThreshold: 1 timeoutSeconds: 10 name: sglang ports: - containerPort: 30000 name: sglang-http protocol: TCP patchWorkerTemplate: {} template: metadata: labels: inference-framework: sglang inference-stack.io/monitoring: "enabled" spec: containers: - name: sglang image: lmsysorg/sglang:latest env: - name: SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK value: "1" - name: CUDA_LAUNCH_BLOCKING value: "0" - name: SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT value: "1000000000" - name: NVSHMEM_IB_TRAFFIC_CLASS value: "16" - name: NVSHMEM_DISABLE_P2P value: "0" - name: ENABLE_METRICS value: "true" - name: NVSHMEM_IB_GID_INDEX value: "3" - name: NVSHMEM_IB_SL value: "5" - name: SGLANG_SET_CPU_AFFINITY value: "true" - name: SGL_ENABLE_JIT_DEEPGEMM value: "1" - name: NCCL_IB_QPS_PER_CONNECTION value: "8" - name: NCCL_IB_SPLIT_DATA_ON_QPS value: "1" - name: NCCL_NET_PLUGIN value: "none" - name: NCCL_IB_TC value: "136" - name: NCCL_IB_SL value: "5" - name: NCCL_IB_TIMEOUT value: "22" - name: NCCL_IB_GID_INDEX value: "3" - name: NCCL_MIN_NCHANNELS value: "4" - name: NCCL_SOCKET_IFNAME value: bond0 - name: GLOO_SOCKET_IFNAME value: bond0 - name: NCCL_IB_HCA value: ^=mlx5_0,mlx5_5,mlx5_6 - name: NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME value: "bond0" - name: MC_TE_METRIC value: "false" resources: limits: nvidia.com/gpu: "8" securityContext: capabilities: add: - IPC_LOCK privileged: true volumeMounts: - mountPath: /root/.cache name: sgl-cache - mountPath: /dev/shm name: dshm - mountPath: /work/models name: model - mountPath: /dev/infiniband name: ib - mountPath: /sgl-workspace/sglang name: src dnsPolicy: ClusterFirstWithHostNet hostIPC: true hostNetwork: true nodeSelector: pd: "yes" tolerations: - key: pd operator: Exists volumes: - hostPath: path: /var/run/sys-topology name: topo - hostPath: path: /data1/sgl_cache4 type: DirectoryOrCreate name: sgl-cache - emptyDir: medium: Memory name: dshm - hostPath: path: /data/DeepSeek-V3.2-Exp name: model - hostPath: path: /dev/infiniband name: ib - hostPath: path: /data/src/sglang type: DirectoryOrCreate name: src - name: decode replicas: 1 workload: apiVersion: leaderworkerset.x-k8s.io/v1 kind: LeaderWorkerSet leaderWorkerSet: size: 1 patchLeaderTemplate: metadata: labels: role: leader pd_role: decode spec: containers: - command: - python3 - -m - sglang.launch_server - --model-path - /work/models - --port - "30000" - --trust-remote - --host - 0.0.0.0 - --disaggregation-ib-device - mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7 - --chunked-prefill-size - "131072" - --eplb-rebalance-layers-per-chunk - "29" - --page-size - "64" - --enable-dp-attention - --enable-dp-lm-head - --dp-size - "8" - --moe-a2a-backend - deepep - --deepep-mode - low_latency - --disaggregation-mode - decode - --mem-fraction-static - "0.8" - --context-length - "32768" - --max-running-requests - "2048" - --tp-size - "8" # Size of Tensor Parallelism - --cuda-graph-max-bs - "16" - --dist-init-addr - $(LWS_LEADER_ADDRESS):20102 - --nnodes - $(LWS_GROUP_SIZE) - --node-rank - $(LWS_WORKER_INDEX) - --trust-remote-code - --ep-num-redundant-experts - "32" - --moe-dense-tp-size - "1" env: - name: LWS_WORKER_INDEX valueFrom: fieldRef: fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] livenessProbe: failureThreshold: 30000 httpGet: path: /health port: 30000 initialDelaySeconds: 300 periodSeconds: 60 successThreshold: 1 timeoutSeconds: 10 name: sglang readinessProbe: failureThreshold: 20 httpGet: path: /health port: 30000 periodSeconds: 30 successThreshold: 1 timeoutSeconds: 10 patchWorkerTemplate: spec: containers: - command: - python3 - -m - sglang.launch_server - --model-path - /work/models - --crash-dump-folder - /log - --chunked-prefill-size - "262144" - --eplb-rebalance-layers-per-chunk - "29" - --page-size - "64" - --enable-dp-attention - --enable-dp-lm-head - --dp-size - "32" - --moe-a2a-backend - "deepep" - --deepep-mode - low_latency - --disaggregation-mode - decode - --mem-fraction-static - "0.849" - --context-length - "32768" - --disaggregation-ib-device - mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7 - --max-running-requests - "4096" - --cuda-graph-max-bs - "16" - --tp-size - "8" # Size of Tensor Parallelism - --dist-init-addr - $(LWS_LEADER_ADDRESS):20102 - --nnodes - $(LWS_GROUP_SIZE) - --node-rank - $(LWS_WORKER_INDEX) - --trust-remote-code - --ep-num-redundant-experts - "32" - --moe-dense-tp-size - "1" env: - name: LWS_WORKER_INDEX valueFrom: fieldRef: fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] name: sglang template: metadata: labels: inference-framework: sglang-unuse inference-stack.io/monitoring: "enabled" spec: containers: - image: lmsysorg/sglang:latest name: sglang resources: limits: nvidia.com/gpu: "8" securityContext: capabilities: add: - IPC_LOCK privileged: true volumeMounts: - mountPath: /root/.cache name: sgl-cache - mountPath: /dev/shm name: dshm - mountPath: /work/models name: model - mountPath: /dev/infiniband name: ib - mountPath: /sgl-workspace/sglang name: src env: - name: SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK value: "1" - name: SGLANG_DISAGGREGATION_WAITING_TIMEOUT value: "100000000" - name: NVSHMEM_DISABLE_P2P value: "0" - name: NVSHMEM_IB_TRAFFIC_CLASS value: "16" - name: NVSHMEM_IB_SL value: "5" - name: ENABLE_METRICS value: "true" - name: CUDA_LAUNCH_BLOCKING value: "0" - name: NVSHMEM_IB_GID_INDEX value: "3" - name: NCCL_IB_QPS_PER_CONNECTION value: "8" - name: NCCL_IB_SPLIT_DATA_ON_QPS value: "1" - name: NCCL_NET_PLUGIN value: "none" - name: NCCL_IB_TC value: "136" - name: NCCL_IB_SL value: "5" - name: NCCL_IB_TIMEOUT value: "22" - name: NCCL_IB_GID_INDEX value: "3" - name: NCCL_MIN_NCHANNELS value: "4" - name: NCCL_SOCKET_IFNAME value: bond0 - name: GLOO_SOCKET_IFNAME value: bond0 - name: NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME value: "bond0" - name: NCCL_IB_HCA value: ^=mlx5_0,mlx5_5,mlx5_6 - name: MC_TE_METRIC value: "false" - name: SGL_ENABLE_JIT_DEEPGEMM value: "1" dnsPolicy: ClusterFirstWithHostNet hostIPC: true hostNetwork: true nodeSelector: pd: "yes" tolerations: - key: pd operator: Exists volumes: - hostPath: path: /var/run/sys-topology name: topo - hostPath: path: /data1/sgl_cache4 type: DirectoryOrCreate name: sgl-cache - hostPath: path: /data/src/sglang type: DirectoryOrCreate name: src - emptyDir: medium: Memory name: dshm - hostPath: path: /data/DeepSeek-V3.2-Exp name: model - hostPath: path: /dev/infiniband name: ib - name: router replicas: 1 dependencies: [ "decode", "prefill" ] template: spec: containers: - name: scheduler image: lmsysorg/sglang:latest command: - sh - -c - > python3 -m sglang_router.launch_router --host 0.0.0.0 --port 8080 --pd-disaggregation --policy random --service-discovery --service-discovery-namespace ${NAMESPACE} --service-discovery-port 30000 --prefill-selector pd_role=prefill --decode-selector pd_role=decode --max-payload-size 2147483648 --worker-startup-timeout-secs 1200 env: - name: NAMESPACE valueFrom: fieldRef: apiVersion: v1 fieldPath: metadata.namespace --- apiVersion: v1 kind: Service metadata: labels: app: deepseek-rbg-32exp name: deepseek-rbg-32exp namespace: default spec: ports: - name: http port: 8080 protocol: TCP targetPort: 8080 nodePort: 30080 selector: rolebasedgroup.workloads.x-k8s.io/name: deepseek-rbg-32exp rolebasedgroup.workloads.x-k8s.io/role: router type: NodePort ``` ```bash [root@ecs-001]# kubectl get po -n default deepseek-rbg-32exp-decode-main-0 1/1 Running 0 74m deepseek-rbg-32exp-decode-0-1 1/1 Running 0 74m deepseek-rbg-32exp-router-9c5dbfc57 1/1 Running 0 22m deepseek-rbg-32exp-prefill-0 1/1 Running 0 74m [root@ecs-cbm-x1-pd-cpu-001 main_doc]# kubectl get svc |grep dee deepseek-rbg-32exp-decode ClusterIP None 97m deepseek-rbg-32exp-router-service NodePort 172.16.242.169 8000:30800/TCP 22m deepseek-rbg-32exp-prefill ClusterIP None 97m ``` At this point, select a nodePort:30800 to access: ```bash [root@ecs-001]# curl -X POST "http://{nodePort}:30800/v1/chat/completions" \ > -H "Content-Type: application/json" \ > -H "Authorization: Bearer None" \ > -d '{ > "rid":"ccccdd", > "model": "dsv32", > "messages": [ > {"role": "system", "content": "0: You are a helpful AI assistant"}, > {"role": "user", "content": "你是谁？."} > ], > "max_tokens":221 > }' {"id":"ccccdd","object":"chat.completion","created":1750252498,"model":"qwen2","choices":[{"index":0,"message":{"role":"assistant","content":"\n嗯，用户问了一个很基础的自我介绍问题"你是谁？"。这可能是第一次互动时的常规开场白，也可能是想确认我的身份和功能范围。\n\n用户没有提供任何背景信息，语气简洁中性。这种场景下新用户的可能性较高，需要给出清晰友好的自我介绍，同时突出实用价值来降低陌生感。\n\n考虑到中文用户，应该用简体中文回复。重点要说明三点：身份归属（深度求索）、功能定位（AI助手）、服务范围（学习/工作/生活）。结尾用开放性问题引导对话很关键——既能了解需求，又能避免让用户面对空白输入框时不知所措。\n\n用波浪线结尾可以软化语气，那个笑脸表情😊刚好能中和AI的机械感。不过要控制表情符号数量，避免显得轻浮。\n\n你好呀！我是你的AI助手，由深度求索公司（DeepSeek）开发的语言模型，名字叫 **DeepSeek-V32**。你可以把我当成一个知识丰富、随叫随到的小帮手～😊\n\n我的任务就是陪你聊天、解答问题、","reasoning_content":null,"tool_calls":null},"logprobs":null,"finish_reason":"length","matched_stop":null}],"usage":{"prompt_tokens":14,"total_tokens":235,"completion_tokens":221,"prompt_tokens_details":null}} ``` ## FAQ 1. The current deployment startup parameters may not be fully compatible with all RDMA scenarios. Different RDMA NCCL-related environment configurations may be needed in different network environments. 2. Please ensure that the sglang code in the image has incorporated the changes from [PR #10912](https://github.com/sgl-project/sglang/pull/10912). ================================================ FILE: docs/references/post_training_integration.md ================================================ # Post-Training Integration SGLang has become the de facto inference backend for modern LLM training frameworks, powering state-of-the-art models across the industry. From GLM-4.6 to Qwen3, leading models leverage SGLang's high-performance inference during reinforcement learning and post-training workflows. What makes SGLang essential for post-training? - Open-To-Use Refit Functionality: diverse method for colocate or disaggregate - Easy To Postpone Generation: enable partial rollout and dedicated rollout control - Fine-Grained Engine Sleep And Wake Up: facilitate maximum-powered rollout and training - Training Serving Alignment: ensure the performance consistency in training and serving - Load Balancing Router: cache-aware load-balancing for high-throughput rollout - Deterministic Inference: ensure zero kl divergence between rollout and training These capabilities, combined with native integration support across major frameworks, have established SGLang as the infrastructure backbone for modern LLM/VLMs post-training. We also share our latest work in this slide, [Optimizing Large-Scale RL with SGLang](https://gamma.app/docs/Optimizing-RL-with-SGLang-y0kqgj877k34779). ## Adoption - [**Miles**](https://github.com/radixark/miles): Enterprise-scale RL framework for large MoE models with SGLang-native rollout, speculative training, and production-grade stability - [**slime**](https://github.com/THUDM/slime): Post-training framework combining Megatron and SGLang, used to train GLM-4.6 - [**AReaL**](https://github.com/inclusionAI/AReaL): Fully asynchronous RL system achieving 2.77x speedup with SGLang backend for continuous rollout generation - [**ROLL**](https://github.com/alibaba/ROLL): ROLL is an efficient and user-friendly RL library designed for Large Language Models utilizing Large Scale GPU resources - [**verl**](https://github.com/volcengine/verl): Full-stack RLHF framework supporting PPO, GRPO, and ReMax with modular SGLang integration - [**Unsloth**](https://docs.unsloth.ai/basics/inference-and-deployment/sglang-guide): 2x faster fine-tuning with optimized kernels, deploys seamlessly with SGLang inference - [**LLaMA Factory**](https://github.com/hiyouga/LLaMA-Factory): Unified framework for training 100+ LLMs with LoRA, QLoRA, and full fine-tuning methods - [**Tunix**](https://github.com/google/tunix): Google's JAX-native library for LLM post-training with SFT, DPO, PPO, and GRPO support - [**RL2**](https://github.com/ChenmienTan/RL2): Ray Less Reinforcement Learning, a concise library of post-training for large language models ## Collaboration Due to the privacy of the design partners, we cannot list the companies that adopt SGLang for post-training. However, we are happy to share the details with you if you are interested and trust the choice among 10+ top companies and frontier labs across US and China. If you are interested in integrating SGLang with your training framework or need technical support, we're here to help! Reach out to us at **rl_team@lmsys.org** for partnerships, integration guidance, and custom feature development. ================================================ FILE: docs/references/production_metrics.md ================================================ # Production Metrics SGLang exposes the following metrics via Prometheus. You can enable it by adding `--enable-metrics` when you launch the server. An example of the monitoring dashboard is available in [examples/monitoring/grafana.json](https://github.com/sgl-project/sglang/blob/main/examples/monitoring/grafana/dashboards/json/sglang-dashboard.json). Here is an example of the metrics: ``` $ curl http://localhost:30000/metrics # HELP sglang:prompt_tokens_total Number of prefill tokens processed. # TYPE sglang:prompt_tokens_total counter sglang:prompt_tokens_total{model_name="meta-llama/Llama-3.1-8B-Instruct"} 8.128902e+06 # HELP sglang:generation_tokens_total Number of generation tokens processed. # TYPE sglang:generation_tokens_total counter sglang:generation_tokens_total{model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.557572e+06 # HELP sglang:token_usage The token usage # TYPE sglang:token_usage gauge sglang:token_usage{model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.28 # HELP sglang:cache_hit_rate The cache hit rate # TYPE sglang:cache_hit_rate gauge sglang:cache_hit_rate{model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.007507552643049313 # HELP sglang:time_to_first_token_seconds Histogram of time to first token in seconds. # TYPE sglang:time_to_first_token_seconds histogram sglang:time_to_first_token_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct"} 2.3518979474117756e+06 sglang:time_to_first_token_seconds_bucket{le="0.001",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 sglang:time_to_first_token_seconds_bucket{le="0.005",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 sglang:time_to_first_token_seconds_bucket{le="0.01",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 sglang:time_to_first_token_seconds_bucket{le="0.02",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 sglang:time_to_first_token_seconds_bucket{le="0.04",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1.0 sglang:time_to_first_token_seconds_bucket{le="0.06",model_name="meta-llama/Llama-3.1-8B-Instruct"} 3.0 sglang:time_to_first_token_seconds_bucket{le="0.08",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 sglang:time_to_first_token_seconds_bucket{le="0.1",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 sglang:time_to_first_token_seconds_bucket{le="0.25",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 sglang:time_to_first_token_seconds_bucket{le="0.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 sglang:time_to_first_token_seconds_bucket{le="0.75",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 sglang:time_to_first_token_seconds_bucket{le="1.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 27.0 sglang:time_to_first_token_seconds_bucket{le="2.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 140.0 sglang:time_to_first_token_seconds_bucket{le="5.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 314.0 sglang:time_to_first_token_seconds_bucket{le="7.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 941.0 sglang:time_to_first_token_seconds_bucket{le="10.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1330.0 sglang:time_to_first_token_seconds_bucket{le="15.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1970.0 sglang:time_to_first_token_seconds_bucket{le="20.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 2326.0 sglang:time_to_first_token_seconds_bucket{le="25.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 2417.0 sglang:time_to_first_token_seconds_bucket{le="30.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 2513.0 sglang:time_to_first_token_seconds_bucket{le="+Inf",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11008.0 sglang:time_to_first_token_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct"} 11008.0 # HELP sglang:e2e_request_latency_seconds Histogram of End-to-end request latency in seconds # TYPE sglang:e2e_request_latency_seconds histogram sglang:e2e_request_latency_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct"} 3.116093850019932e+06 sglang:e2e_request_latency_seconds_bucket{le="0.3",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 sglang:e2e_request_latency_seconds_bucket{le="0.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 sglang:e2e_request_latency_seconds_bucket{le="0.8",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 sglang:e2e_request_latency_seconds_bucket{le="1.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 sglang:e2e_request_latency_seconds_bucket{le="1.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 sglang:e2e_request_latency_seconds_bucket{le="2.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 sglang:e2e_request_latency_seconds_bucket{le="2.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 sglang:e2e_request_latency_seconds_bucket{le="5.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.0 sglang:e2e_request_latency_seconds_bucket{le="10.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 10.0 sglang:e2e_request_latency_seconds_bucket{le="15.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11.0 sglang:e2e_request_latency_seconds_bucket{le="20.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 14.0 sglang:e2e_request_latency_seconds_bucket{le="30.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 247.0 sglang:e2e_request_latency_seconds_bucket{le="40.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 486.0 sglang:e2e_request_latency_seconds_bucket{le="50.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 845.0 sglang:e2e_request_latency_seconds_bucket{le="60.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1513.0 sglang:e2e_request_latency_seconds_bucket{le="+Inf",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11228.0 sglang:e2e_request_latency_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct"} 11228.0 # HELP sglang:time_per_output_token_seconds Histogram of time per output token in seconds. # TYPE sglang:time_per_output_token_seconds histogram sglang:time_per_output_token_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct"} 866964.5791549598 sglang:time_per_output_token_seconds_bucket{le="0.005",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1.0 sglang:time_per_output_token_seconds_bucket{le="0.01",model_name="meta-llama/Llama-3.1-8B-Instruct"} 73.0 sglang:time_per_output_token_seconds_bucket{le="0.015",model_name="meta-llama/Llama-3.1-8B-Instruct"} 382.0 sglang:time_per_output_token_seconds_bucket{le="0.02",model_name="meta-llama/Llama-3.1-8B-Instruct"} 593.0 sglang:time_per_output_token_seconds_bucket{le="0.025",model_name="meta-llama/Llama-3.1-8B-Instruct"} 855.0 sglang:time_per_output_token_seconds_bucket{le="0.03",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1035.0 sglang:time_per_output_token_seconds_bucket{le="0.04",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1815.0 sglang:time_per_output_token_seconds_bucket{le="0.05",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11685.0 sglang:time_per_output_token_seconds_bucket{le="0.075",model_name="meta-llama/Llama-3.1-8B-Instruct"} 433413.0 sglang:time_per_output_token_seconds_bucket{le="0.1",model_name="meta-llama/Llama-3.1-8B-Instruct"} 4.950195e+06 sglang:time_per_output_token_seconds_bucket{le="0.15",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.039435e+06 sglang:time_per_output_token_seconds_bucket{le="0.2",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.171662e+06 sglang:time_per_output_token_seconds_bucket{le="0.3",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.266055e+06 sglang:time_per_output_token_seconds_bucket{le="0.4",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.296752e+06 sglang:time_per_output_token_seconds_bucket{le="0.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.312226e+06 sglang:time_per_output_token_seconds_bucket{le="0.75",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.339675e+06 sglang:time_per_output_token_seconds_bucket{le="1.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.357747e+06 sglang:time_per_output_token_seconds_bucket{le="2.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.389414e+06 sglang:time_per_output_token_seconds_bucket{le="+Inf",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.400757e+06 sglang:time_per_output_token_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.400757e+06 # HELP sglang:func_latency_seconds Function latency in seconds # TYPE sglang:func_latency_seconds histogram sglang:func_latency_seconds_sum{name="generate_request"} 4.514771912145079 sglang:func_latency_seconds_bucket{le="0.05",name="generate_request"} 14006.0 sglang:func_latency_seconds_bucket{le="0.07500000000000001",name="generate_request"} 14006.0 sglang:func_latency_seconds_bucket{le="0.1125",name="generate_request"} 14006.0 sglang:func_latency_seconds_bucket{le="0.16875",name="generate_request"} 14006.0 sglang:func_latency_seconds_bucket{le="0.253125",name="generate_request"} 14006.0 sglang:func_latency_seconds_bucket{le="0.3796875",name="generate_request"} 14006.0 sglang:func_latency_seconds_bucket{le="0.56953125",name="generate_request"} 14006.0 sglang:func_latency_seconds_bucket{le="0.8542968750000001",name="generate_request"} 14006.0 sglang:func_latency_seconds_bucket{le="1.2814453125",name="generate_request"} 14006.0 sglang:func_latency_seconds_bucket{le="1.9221679687500002",name="generate_request"} 14006.0 sglang:func_latency_seconds_bucket{le="2.8832519531250003",name="generate_request"} 14006.0 sglang:func_latency_seconds_bucket{le="4.3248779296875",name="generate_request"} 14007.0 sglang:func_latency_seconds_bucket{le="6.487316894531251",name="generate_request"} 14007.0 sglang:func_latency_seconds_bucket{le="9.730975341796876",name="generate_request"} 14007.0 sglang:func_latency_seconds_bucket{le="14.596463012695313",name="generate_request"} 14007.0 sglang:func_latency_seconds_bucket{le="21.89469451904297",name="generate_request"} 14007.0 sglang:func_latency_seconds_bucket{le="32.84204177856446",name="generate_request"} 14007.0 sglang:func_latency_seconds_bucket{le="49.26306266784668",name="generate_request"} 14007.0 sglang:func_latency_seconds_bucket{le="+Inf",name="generate_request"} 14007.0 sglang:func_latency_seconds_count{name="generate_request"} 14007.0 # HELP sglang:num_running_reqs The number of running requests # TYPE sglang:num_running_reqs gauge sglang:num_running_reqs{model_name="meta-llama/Llama-3.1-8B-Instruct"} 162.0 # HELP sglang:num_used_tokens The number of used tokens # TYPE sglang:num_used_tokens gauge sglang:num_used_tokens{model_name="meta-llama/Llama-3.1-8B-Instruct"} 123859.0 # HELP sglang:gen_throughput The generate throughput (token/s) # TYPE sglang:gen_throughput gauge sglang:gen_throughput{model_name="meta-llama/Llama-3.1-8B-Instruct"} 86.50814177726902 # HELP sglang:num_queue_reqs The number of requests in the waiting queue # TYPE sglang:num_queue_reqs gauge sglang:num_queue_reqs{model_name="meta-llama/Llama-3.1-8B-Instruct"} 2826.0 ``` ## Setup Guide This section describes how to set up the monitoring stack (Prometheus + Grafana) provided in the `examples/monitoring` directory. ### Prerequisites - Docker and Docker Compose installed - SGLang server running with metrics enabled ### Usage 1. **Start your SGLang server with metrics enabled:** ```bash python -m sglang.launch_server \ --model-path \ --port 30000 \ --enable-metrics ``` Replace `` with the actual path to your model (e.g., `meta-llama/Meta-Llama-3.1-8B-Instruct`). Ensure the server is accessible from the monitoring stack (you might need `--host 0.0.0.0` if running in Docker). By default, the metrics endpoint will be available at `http://:30000/metrics`. 2. **Navigate to the monitoring example directory:** ```bash cd examples/monitoring ``` 3. **Start the monitoring stack:** ```bash docker compose up -d ``` This command will start Prometheus and Grafana in the background. 4. **Access the monitoring interfaces:** * **Grafana:** Open your web browser and go to [http://localhost:3000](http://localhost:3000). * **Prometheus:** Open your web browser and go to [http://localhost:9090](http://localhost:9090). 5. **Log in to Grafana:** * Default Username: `admin` * Default Password: `admin` You will be prompted to change the password upon your first login. 6. **View the Dashboard:** The SGLang dashboard is pre-configured and should be available automatically. Navigate to `Dashboards` -> `Browse` -> `SGLang Monitoring` folder -> `SGLang Dashboard`. ### Troubleshooting * **Port Conflicts:** If you encounter errors like "port is already allocated," check if other services (including previous instances of Prometheus/Grafana) are using ports `9090` or `3000`. Use `docker ps` to find running containers and `docker stop ` to stop them, or use `lsof -i :` to find other processes using the ports. You might need to adjust the ports in the `docker-compose.yaml` file if they permanently conflict with other essential services on your system. To modify Grafana's port to the other one(like 3090) in your Docker Compose file, you need to explicitly specify the port mapping under the grafana service. Option 1: Add GF_SERVER_HTTP_PORT to the environment section: ``` environment: - GF_AUTH_ANONYMOUS_ENABLED=true - GF_SERVER_HTTP_PORT=3090 # <-- Add this line ``` Option 2: Use port mapping: ``` grafana: image: grafana/grafana:latest container_name: grafana ports: - "3090:3000" # <-- Host:Container port mapping ``` * **Connection Issues:** * Ensure both Prometheus and Grafana containers are running (`docker ps`). * Verify the Prometheus data source configuration in Grafana (usually auto-configured via `grafana/datasources/datasource.yaml`). Go to `Connections` -> `Data sources` -> `Prometheus`. The URL should point to the Prometheus service (e.g., `http://prometheus:9090`). * Confirm that your SGLang server is running and the metrics endpoint (`http://:30000/metrics`) is accessible *from the Prometheus container*. If SGLang is running on your host machine and Prometheus is in Docker, use `host.docker.internal` (on Docker Desktop) or your machine's network IP instead of `localhost` in the `prometheus.yaml` scrape configuration. * **No Data on Dashboard:** * Generate some traffic to your SGLang server to produce metrics. For example, run a benchmark: ```bash python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 100 --random-input 128 --random-output 128 ``` * Check the Prometheus UI (`http://localhost:9090`) under `Status` -> `Targets` to see if the SGLang endpoint is being scraped successfully. * Verify the `model_name` and `instance` labels in your Prometheus metrics match the variables used in the Grafana dashboard. You might need to adjust the Grafana dashboard variables or the labels in your Prometheus configuration. ### Configuration Files The monitoring setup is defined by the following files within the `examples/monitoring` directory: * `docker-compose.yaml`: Defines the Prometheus and Grafana services. * `prometheus.yaml`: Prometheus configuration, including scrape targets. * `grafana/datasources/datasource.yaml`: Configures the Prometheus data source for Grafana. * `grafana/dashboards/config/dashboard.yaml`: Tells Grafana to load dashboards from the specified path. * `grafana/dashboards/json/sglang-dashboard.json`: The actual Grafana dashboard definition in JSON format. You can customize the setup by modifying these files. For instance, you might need to update the `static_configs` target in `prometheus.yaml` if your SGLang server runs on a different host or port. #### Check if the metrics are being collected Run: ``` python3 -m sglang.bench_serving \ --backend sglang \ --dataset-name random \ --num-prompts 3000 \ --random-input 1024 \ --random-output 1024 \ --random-range-ratio 0.5 ``` to generate some requests. Then you should be able to see the metrics in the Grafana dashboard. ================================================ FILE: docs/references/production_request_trace.md ================================================ # Production Request Tracing SGLang exports request trace data based on the OpenTelemetry Collector. You can enable tracing by adding the `--enable-trace` and configure the OpenTelemetry Collector endpoint using `--otlp-traces-endpoint` when launching the server. You can find example screenshots of the visualization in https://github.com/sgl-project/sglang/issues/8965. ## Setup Guide This section explains how to configure the request tracing and export the trace data. 1. Install the required packages and tools * install Docker and Docker Compose * install the dependencies ```bash # enter the SGLang root directory pip install -e "python[tracing]" # or manually install the dependencies using pip pip install opentelemetry-sdk opentelemetry-api opentelemetry-exporter-otlp opentelemetry-exporter-otlp-proto-grpc ``` 2. Launch OpenTelemetry collector and Jaeger ```bash docker compose -f examples/monitoring/tracing_compose.yaml up -d ``` 3. Start your SGLang server with tracing enabled ```bash # set env variables export SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS=500 export SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE=64 # start the prefill and decode server python -m sglang.launch_server --enable-trace --otlp-traces-endpoint 0.0.0.0:4317 # start the model-gate-way python -m sglang_router.launch_router --enable-trace --otlp-traces-endpoint 0.0.0.0:4317 ``` Replace `0.0.0.0:4317` with the actual endpoint of the OpenTelemetry collector. If you launched the openTelemetry collector with tracing_compose.yaml, the default receiving port is 4317. To use the HTTP/protobuf span exporter, set the following environment variable and point to an HTTP endpoint, for example, `http://0.0.0.0:4318/v1/traces`. ```bash export OTEL_EXPORTER_OTLP_TRACES_PROTOCOL=http/protobuf ``` 4. Raise some requests 5. Observe whether trace data is being exported * Access port 16686 of Jaeger using a web browser to visualize the request traces. * The OpenTelemetry Collector also exports trace data in JSON format to /tmp/otel_trace.json. In a follow-up patch, we will provide a tool to convert this data into a Perfetto-compatible format, enabling visualization of requests in the Perfetto UI. 6. Dynamically adjust trace level The trace level accepts configurable values from `0` to `3`. The meanings of different trace level values are as follows: ``` 0: disable tracing 1: Trace important slices 2: Trace all slices except nested ones 3: Trace all slices ``` The trace level can be dynamically set via HTTP API, for example: ```bash curl http://0.0.0.0:30000/set_trace_level?level=2 ``` Replace `0.0.0.0:30000` with your actual server address, and replace `level=2` with the level you want to set. **Note**: You must set the parameter `--enable-trace`; otherwise, the trace capability will not be enabled regardless of any dynamic adjustments to the trace level. ## How to add Tracing for slices you're interested in?(API introduction) We have already inserted instrumentation points in the tokenizer and scheduler main threads. If you wish to trace additional request execution segments or perform finer-grained tracing, please use the APIs from the tracing package as described below. **All of the following implementations are done in python/sglang/srt/observability/req_time_stats.py. If you want to add another slice, please do it here.** 1. Initialization Every process involved in tracing during the initialization phase should execute: ```python process_tracing_init(otlp_traces_endpoint, server_name) ``` The otlp_traces_endpoint is obtained from the arguments, and you can set server_name freely, but it should remain consistent across all processes. Every thread involved in tracing during the initialization phase should execute: ```python trace_set_thread_info("thread label", tp_rank, dp_rank) ``` The "thread label" can be regarded as the name of the thread, used to distinguish different threads in the visualization view. 2. Create a trace context for a request Each request needs to call `TraceReqContext()` to initialize a request context, which is used to generate slice spans and record request stage info. You can either store it within the request object or maintain it as a global variable. 3. Mark the beginning and end of a request ``` trace_ctx.trace_req_start(). trace_ctx.trace_req_finish() ``` trace_req_start() and trace_req_finish() must be called within the same process, for example, in the tokenizer. 4. Add tracing for a slice * Add slice tracing normally: ```python trace_ctx.trace_slice_start(RequestStage.TOKENIZER.stage_name) trace_ctx.trace_slice_end(RequestStage.TOKENIZER.stage_name) or trace_ctx.trace_slice(slice: TraceSliceContext) ``` - The end of the last slice in a thread must be marked with thread_finish_flag=True, or explicitly call trace_ctx.abort(); otherwise, the thread's span will not be properly generated. ```python trace_ctx.slice_end(RequestStage.D.stage_name, thread_finish_flag = True) trace_ctx.abort() ``` 5. When the request execution flow transfers to another thread, the thread context needs to be explicitly rebuilt. - receiver: Execute the following code after receiving the request via ZMQ ```python trace_ctx.rebuild_thread_context() ``` ## How to Extend the Tracing Framework to Support Complex Tracing Scenarios The currently provided tracing package still has potential for further development. If you wish to build more advanced features upon it, you must first understand its existing design principles. The core of the tracing framework's implementation lies in the design of the span structure and the trace context. To aggregate scattered slices and enable concurrent tracking of multiple requests, we have designed a three-level trace context structure or span structure: `TraceReqContext`, `TraceThreadContext` and `TraceSliceContext`. Their relationship is as follows: ``` TraceReqContext (req_id="req-123") ├── TraceThreadContext(thread_label="scheduler", tp_rank=0) | └── TraceSliceContext(slice_name="prefill") | └── TraceThreadContext(thread_label="scheduler", tp_rank=1) └── TraceSliceContext(slice_name="prefill") ``` Each traced request maintains a global `TraceReqContext` and creates a corresponding request span. For every thread that processes the request, a `TraceThreadContext` is recorded and a thread span is created. The `TraceThreadContext` is nested within the `TraceReqContext`, and each currently traced code slice—potentially nested—is stored in its associated `TraceThreadContext`. In addition to the above hierarchy, each slice also records its previous slice via Span.add_link(), which can be used to trace the execution flow. ================================================ FILE: docs/references/release_lookup.rst ================================================ Release Lookup ============== Find which SGLang release first included a specific PR or commit. .. raw:: html
Initializing…
================================================ FILE: docs/references/torch_compile_cache.md ================================================ # Enabling cache for torch.compile SGLang uses `max-autotune-no-cudagraphs` mode of torch.compile. The auto-tuning can be slow. If you want to deploy a model on many different machines, you can ship the torch.compile cache to these machines and skip the compilation steps. This is based on https://pytorch.org/tutorials/recipes/torch_compile_caching_tutorial.html 1. Generate the cache by setting TORCHINDUCTOR_CACHE_DIR and running the model once. ``` TORCHINDUCTOR_CACHE_DIR=/root/inductor_root_cache python3 -m sglang.launch_server --model meta-llama/Llama-3.1-8B-Instruct --enable-torch-compile ``` 2. Copy the cache folder to other machines and launch the server with `TORCHINDUCTOR_CACHE_DIR`. ================================================ FILE: docs/release_lookup/README.md ================================================ # SGLang Release Lookup Tool This tool allows users to find the earliest release that contains a specific PR or commit. It runs entirely in the browser using a static JSON index generated from the git history. ## Usage 1. **Generate the Index**: Run the Python script to generate the `release_index.json` file from your local git repository. ```bash python3 generate_index.py --output release_index.json ``` This script: - Finds all tags matching `v*` and `gateway-v*`. - Sorts them by creation date. - Traverses the history to find which release first introduced each commit and PR. - Extracts PR numbers from commit messages. 2. **Open the Tool**: Open `index.html` in your browser. ```bash # You can open it directly if your browser supports local file fetch (Firefox usually does), # or serve it locally: python3 -m http.server # Then go to http://localhost:8000/index.html ``` ## Files - `index.html`: The UI for the lookup tool. - `generate_index.py`: Script to build the index. - `release_index.json`: The index file used by the UI. ## Logic The tool determines the "earliest release" based on the tag creation date. It traverses tags from oldest to newest. Any commit reachable from a tag (that wasn't reachable from a previous tag) is assigned to that release. ================================================ FILE: docs/release_lookup/generate_index.py ================================================ import argparse import json import os import re import subprocess import sys from datetime import datetime # Short hash length for commits (7 is git's default short hash) SHORT_HASH_LEN = 8 COMMIT_CHUNK_SIZE = 1000 def run_git(cmd): try: output = subprocess.check_output(cmd, stderr=subprocess.STDOUT) return output.decode("utf-8", errors="replace").strip() except subprocess.CalledProcessError as e: print(f"Error running cmd: {cmd}\n{e.output.decode('utf-8', errors='replace')}") sys.exit(1) def is_stable_release(tag_name): """Check if tag is a stable release (not rc/alpha/beta).""" # Skip release candidates, alpha, beta versions if re.search(r"(rc|alpha|beta)\d*$", tag_name, re.IGNORECASE): return False return True def get_tags(): # Get tags sorted by creator date cmd = [ "git", "tag", "--list", "v*", "gateway-v*", "--sort=creatordate", "--format=%(refname:short)|%(creatordate:iso8601)|%(objectname)", ] raw = run_git(cmd) tags = [] if not raw: return [] for line in raw.split("\n"): parts = line.split("|") if len(parts) >= 3: name, date, commit = parts[0], parts[1], parts[2] # Skip non-stable releases (rc, alpha, beta) if not is_stable_release(name): continue tag_type = "gateway" if name.startswith("gateway-") else "main" tags.append( {"name": name, "date": date, "commit": commit, "type": tag_type} ) return tags def extract_pr_num(message): lines = message.strip().split("\n") first_line = lines[0] m = re.search(r"\(#(\d+)\)$", first_line) if m: return m.group(1) m = re.search(r"Merge pull request #(\d+)", message) if m: return m.group(1) return None def process_tag_line(tags, commit_map, pr_map, tag_type, tag_to_idx): """Process a single release line (main or gateway) independently.""" seen_commits = set() for tag in tags: tag_name = tag["name"] print(f"Processing {tag_name}...") commits = run_git(["git", "rev-list", tag_name]).split("\n") new_commits = [] for c in commits: c = c.strip() if not c: continue if c in seen_commits: continue new_commits.append(c) seen_commits.add(c) if not new_commits: continue for i in range(0, len(new_commits), COMMIT_CHUNK_SIZE): chunk = new_commits[i : i + COMMIT_CHUNK_SIZE] cmd = ["git", "show", "-s", "--format=%H|%B%n--END-COMMIT--"] + chunk raw_logs = run_git(cmd) entries = raw_logs.split("--END-COMMIT--\n") for log_entry in entries: if not log_entry.strip(): continue parts = log_entry.split("|", 1) if len(parts) < 2: continue sha = parts[0].strip() msg = parts[1].strip() tag_idx = tag_to_idx[tag_name] # Store release index using full SHA as key if sha not in commit_map: commit_map[sha] = {} commit_map[sha][tag_type] = tag_idx pr = extract_pr_num(msg) if pr: if pr not in pr_map: pr_map[pr] = {} if tag_type not in pr_map[pr]: pr_map[pr][tag_type] = tag_idx def main(): parser = argparse.ArgumentParser( description="Generate lookup index for sglang releases" ) parser.add_argument( "--output", default="release_index.json", help="Output JSON file" ) args = parser.parse_args() tags = get_tags() print(f"Found {len(tags)} tags.") main_tags = [t for t in tags if t["type"] == "main"] gateway_tags = [t for t in tags if t["type"] == "gateway"] print(f" - {len(main_tags)} main tags") print(f" - {len(gateway_tags)} gateway tags") # Build tag list and index mapping # Tags array: [name, date, type] for each tag tag_list = [] tag_to_idx = {} for tag in tags: tag_to_idx[tag["name"]] = len(tag_list) # Compact format: [name, date, type (0=main, 1=gateway)] tag_list.append( [tag["name"], tag["date"], 1 if tag["type"] == "gateway" else 0] ) pr_map = {} commit_map_full = {} process_tag_line(main_tags, commit_map_full, pr_map, "m", tag_to_idx) process_tag_line(gateway_tags, commit_map_full, pr_map, "g", tag_to_idx) # Convert full SHAs to short SHAs, checking for collisions commit_map = {} short_to_full_map = {} for full_sha, data in commit_map_full.items(): short_sha = full_sha[:SHORT_HASH_LEN] if short_sha in short_to_full_map and short_to_full_map[short_sha] != full_sha: print( f"CRITICAL: Short SHA collision detected for '{short_sha}'\n" f" Commit 1: {short_to_full_map[short_sha]}\n" f" Commit 2: {full_sha}\n" "Please increase SHORT_HASH_LEN and re-run.", file=sys.stderr, ) sys.exit(1) commit_map[short_sha] = data short_to_full_map[short_sha] = full_sha # Compact output format: # - tags: array of [name, date, type] # - prs: {pr_num: tag_idx} or {pr_num: {m: idx, g: idx}} # - commits: {short_hash: tag_idx} or {short_hash: {m: idx, g: idx}} # Simplify single-entry dicts to just the value def simplify_map(m): result = {} for k, v in m.items(): if len(v) == 1: # Single entry: just store the index directly with type prefix key_type, idx = list(v.items())[0] result[k] = f"{key_type}{idx}" else: # Multiple entries: keep as dict result[k] = v return result output_data = { "t": tag_list, # tags "p": simplify_map(pr_map), # prs "c": simplify_map(commit_map), # commits "g": datetime.now().isoformat(), # generated_at } # Write minified JSON with a trailing newline for formatter compatibility. json_str = json.dumps(output_data, separators=(",", ":")) with open(args.output, "w", encoding="utf-8") as f: f.write(json_str) f.write("\n") json_size = os.path.getsize(args.output) print(f"Index generated at {args.output}") print(f"Stats: {len(tag_list)} tags, {len(pr_map)} PRs, {len(commit_map)} commits.") print(f"Size: {json_size/1024:.1f} KB") if __name__ == "__main__": main() ================================================ FILE: docs/release_lookup/index.html ================================================ SGLang Release Lookup

Release Lookup

Find which SGLang release first included your PR or commit.

Initializing...
================================================ FILE: docs/release_lookup/release_index.json ================================================ {"t":[["v0.1.3","2024-01-16 05:55:25 +0000",0],["v0.1.5","2024-01-17 18:37:02 -0800",0],["v0.1.6","2024-01-21 01:45:02 -0800",0],["v0.1.7","2024-01-21 10:31:02 +0000",0],["v0.1.8","2024-01-24 03:33:34 -0800",0],["v0.1.9","2024-01-24 11:37:25 +0000",0],["v0.1.10","2024-01-30 15:37:52 +0000",0],["v0.1.11","2024-02-03 02:50:13 -0800",0],["v0.1.12","2024-02-11 06:43:45 -0800",0],["v0.1.13","2024-03-11 05:49:27 -0700",0],["v0.1.14","2024-03-22 13:42:22 -0700",0],["v0.1.15","2024-05-12 14:22:33 -0700",0],["v0.1.16","2024-05-13 17:29:17 -0700",0],["v0.1.17","2024-06-07 19:49:18 -0700",0],["v0.1.18","2024-07-04 06:27:29 +0000",0],["v0.1.19","2024-07-09 02:23:14 -0700",0],["v0.1.20","2024-07-13 17:27:55 -0700",0],["v0.1.21","2024-07-15 13:10:53 -0700",0],["v0.1.22","2024-07-20 03:39:50 -0700",0],["v0.1.23","2024-07-23 13:49:34 -0700",0],["v0.1.24","2024-07-24 15:55:01 -0700",0],["v0.2.0","2024-07-25 08:03:36 -0700",0],["v0.2.5","2024-07-27 05:56:30 +1000",0],["v0.2.6","2024-07-27 20:29:33 -0700",0],["v0.2.7","2024-07-30 20:41:10 +1000",0],["v0.2.8","2024-08-01 14:18:26 -0700",0],["v0.2.9","2024-08-02 01:45:48 -0700",0],["v0.2.9.post1","2024-08-02 12:08:00 -0700",0],["v0.2.10","2024-08-04 16:52:51 -0700",0],["v0.2.11","2024-08-07 20:47:53 +0800",0],["v0.2.12","2024-08-12 20:59:38 +1000",0],["v0.2.13","2024-08-16 03:50:43 +1000",0],["v0.2.14","2024-08-27 00:28:24 +1000",0],["v0.2.14.post1","2024-08-28 21:16:47 +1000",0],["v0.2.14.post2","2024-08-28 18:46:33 +0000",0],["v0.2.15","2024-09-01 22:22:38 -0700",0],["v0.3.0","2024-09-04 04:21:21 -0700",0],["v0.3.1.post1","2024-09-17 01:47:31 -0700",0],["v0.3.1.post2","2024-09-19 02:03:38 -0700",0],["v0.3.1.post3","2024-09-21 11:17:45 +0800",0],["v0.3.2","2024-09-25 14:17:09 +0800",0],["v0.3.3","2024-10-08 12:58:41 -0700",0],["v0.3.3.post1","2024-10-11 07:56:16 -0700",0],["v0.3.4","2024-10-19 08:17:41 -0700",0],["v0.3.4.post1","2024-10-21 21:16:43 -0700",0],["v0.3.4.post2","2024-10-25 11:07:19 -0700",0],["v0.3.5","2024-11-03 13:48:11 -0800",0],["v0.3.5.post1","2024-11-13 10:27:12 -0800",0],["v0.3.5.post2","2024-11-15 06:54:00 -0800",0],["v0.3.6","2024-11-22 19:27:30 +0800",0],["v0.3.6.post1","2024-11-25 17:31:37 -0800",0],["v0.3.6.post2","2024-11-27 03:35:30 -0800",0],["v0.3.6.post3","2024-11-30 01:41:16 +0800",0],["v0.4.0","2024-12-03 11:55:41 -0800",0],["v0.4.0.post1","2024-12-06 06:08:19 -0800",0],["v0.4.0.post2","2024-12-21 21:16:34 +0800",0],["v0.4.1","2024-12-26 07:14:51 +0800",0],["v0.4.1.post1","2024-12-28 00:11:06 +0800",0],["v0.4.1.post2","2024-12-30 00:11:46 +0800",0],["v0.4.1.post3","2024-12-29 14:25:53 -0800",0],["v0.4.1.post4","2025-01-06 01:29:54 +0800",0],["v0.4.1.post5","2025-01-11 23:10:02 +0800",0],["v0.4.1.post6","2025-01-15 16:23:42 +0800",0],["v0.4.1.post7","2025-01-20 21:50:55 +0800",0],["v0.4.2","2025-01-27 21:42:05 +0800",0],["v0.4.2.post1","2025-01-31 20:35:55 +0800",0],["v0.4.2.post2","2025-02-05 17:35:02 +0800",0],["v0.4.2.post3","2025-02-07 08:20:03 -0800",0],["v0.4.2.post4","2025-02-10 14:12:16 +0800",0],["v0.4.3","2025-02-14 09:43:14 +0800",0],["v0.4.3.post1","2025-02-17 21:58:19 +0800",0],["v0.4.3.post2","2025-02-18 02:48:30 +0800",0],["v0.4.3.post3","2025-03-05 17:26:10 -0800",0],["v0.4.3.post4","2025-03-06 12:50:28 -0800",0],["v0.4.4","2025-03-13 02:49:58 -0700",0],["v0.4.4.post1","2025-03-13 17:53:46 -0700",0],["v0.4.4.post2","2025-03-26 19:58:00 -0700",0],["v0.4.4.post3","2025-03-28 23:21:24 -0700",0],["v0.4.4.post4","2025-04-05 15:36:17 -0700",0],["v0.4.5","2025-04-07 00:35:00 -0700",0],["v0.4.5.post1","2025-04-15 23:00:07 -0700",0],["v0.4.5.post2","2025-04-20 14:12:37 -0700",0],["v0.4.5.post3","2025-04-21 18:16:20 -0700",0],["v0.4.6","2025-04-27 14:07:05 -0700",0],["v0.4.6.post1","2025-04-28 12:57:08 -0700",0],["v0.4.6.post2","2025-04-30 22:04:40 -0700",0],["v0.4.6.post3","2025-05-09 15:38:47 -0700",0],["v0.4.6.post4","2025-05-13 01:57:51 -0700",0],["v0.4.6.post5","2025-05-24 00:48:05 -0700",0],["v0.4.7","2025-06-10 01:56:20 -0700",0],["v0.4.7.post1","2025-06-16 15:20:29 -0700",0],["v0.4.8","2025-06-23 23:14:22 -0700",0],["v0.4.8.post1","2025-06-26 02:21:12 -0700",0],["v0.4.9","2025-07-05 17:40:29 -0700",0],["gateway-v0.1.5","2025-07-06 22:54:17 -0700",1],["v0.4.9.post1","2025-07-09 00:28:17 -0700",0],["v0.4.9.post2","2025-07-11 21:11:20 -0700",0],["gateway-v0.1.6","2025-07-20 23:13:20 -0700",1],["v0.4.9.post3","2025-07-22 15:55:48 -0700",0],["v0.4.9.post4","2025-07-25 17:12:47 -0700",0],["v0.4.9.post5","2025-07-28 02:11:06 -0700",0],["v0.4.9.post6","2025-07-29 02:30:07 -0700",0],["v0.4.10","2025-07-31 20:50:17 +0800",0],["gateway-v0.1.7","2025-07-31 11:24:12 -0700",1],["gateway-v0.1.8","2025-07-31 19:00:23 -0700",1],["v0.4.10.post1","2025-08-01 12:07:30 +0800",0],["v0.4.10.post2","2025-08-03 03:43:29 -0700",0],["gateway-v0.1.9","2025-08-07 09:29:12 -0700",1],["v0.5.1","2025-08-23 07:09:26 -0700",0],["v0.5.1.post1","2025-08-24 01:14:17 -0700",0],["v0.5.1.post2","2025-08-25 03:45:09 -0700",0],["v0.5.1.post3","2025-08-27 15:42:42 -0700",0],["v0.5.2","2025-09-11 16:09:20 -0700",0],["v0.5.3","2025-10-06 20:07:02 +0800",0],["v0.5.3.post1","2025-10-09 15:19:59 -0700",0],["gateway-v0.2.0","2025-10-14 22:10:30 -0400",1],["v0.5.3.post2","2025-10-15 16:49:14 -0700",0],["v0.5.3.post3","2025-10-16 13:14:55 -0700",0],["gateway-v0.2.1","2025-10-20 21:08:45 -0700",1],["v0.5.4","2025-10-23 18:01:40 -0700",0],["v0.5.4.post1","2025-10-27 09:35:20 +0800",0],["gateway-v0.2.2","2025-10-30 14:40:13 -0700",1],["v0.5.4.post2","2025-10-31 17:38:50 -0700",0],["v0.5.4.post3","2025-11-04 18:32:11 -0800",0],["v0.5.5","2025-11-07 00:46:19 +0800",0],["v0.5.5.post1","2025-11-10 11:53:43 -0800",0],["v0.5.5.post2","2025-11-12 20:35:20 +0800",0],["gateway-v0.2.3","2025-11-14 19:04:20 -0800",1],["v0.5.5.post3","2025-11-16 17:55:38 -0800",0],["v0.5.6","2025-12-02 17:17:13 -0800",0],["v0.5.6.post1","2025-12-08 13:41:01 -0800",0],["gateway-v0.2.4","2025-12-09 16:36:17 -0800",1],["v0.5.6.post2","2025-12-11 12:29:52 -0800",0],["gateway-v0.3.0","2025-12-24 16:25:05 -0500",1],["v0.5.7","2026-01-01 10:59:48 +0800",0],["gateway-v0.3.1","2026-01-08 21:50:34 -0800",1],["v0.5.8","2026-01-23 09:58:11 -0800",0],["v0.5.8.post1","2026-02-05 20:56:52 +0800",0]],"p":{"10":{"m":0,"g":94},"8":{"m":0,"g":94},"7":{"m":0,"g":94},"6":{"m":0,"g":94},"4":{"m":0,"g":94},"3":{"m":0,"g":94},"2":{"m":0,"g":94},"1":{"m":0,"g":94},"32":{"m":1,"g":94},"18":{"m":1,"g":94},"30":{"m":1,"g":94},"20":{"m":1,"g":94},"19":{"m":1,"g":94},"9":{"m":1,"g":94},"17":{"m":1,"g":94},"16":{"m":1,"g":94},"15":{"m":1,"g":94},"12":{"m":1,"g":94},"11":{"m":1,"g":94},"68":{"m":2,"g":94},"67":{"m":2,"g":94},"64":{"m":2,"g":94},"63":{"m":2,"g":94},"58":{"m":2,"g":94},"57":{"m":2,"g":94},"36":{"m":2,"g":94},"52":{"m":2,"g":94},"50":{"m":2,"g":94},"49":{"m":2,"g":94},"47":{"m":2,"g":94},"46":{"m":2,"g":94},"45":{"m":2,"g":94},"42":{"m":2,"g":94},"34":{"m":2,"g":94},"33":{"m":2,"g":94},"93":{"m":4,"g":94},"92":{"m":4,"g":94},"90":{"m":4,"g":94},"87":{"m":4,"g":94},"84":{"m":4,"g":94},"83":{"m":4,"g":94},"82":{"m":4,"g":94},"75":{"m":4,"g":94},"80":{"m":4,"g":94},"72":{"m":4,"g":94},"37":{"m":4,"g":94},"71":{"m":4,"g":94},"113":{"m":6,"g":94},"121":{"m":6,"g":94},"120":{"m":6,"g":94},"118":{"m":6,"g":94},"114":{"m":6,"g":94},"117":{"m":6,"g":94},"108":{"m":6,"g":94},"103":{"m":6,"g":94},"101":{"m":6,"g":94},"98":{"m":6,"g":94},"48":{"m":6,"g":94},"97":{"m":6,"g":94},"95":{"m":6,"g":94},"134":{"m":7,"g":94},"133":{"m":7,"g":94},"132":{"m":7,"g":94},"112":{"m":7,"g":94},"129":{"m":7,"g":94},"125":{"m":7,"g":94},"119":{"m":7,"g":94},"116":{"m":7,"g":94},"178":{"m":8,"g":94},"177":{"m":8,"g":94},"172":{"m":8,"g":94},"174":{"m":8,"g":94},"168":{"m":8,"g":94},"170":{"m":8,"g":94},"162":{"m":8,"g":94},"160":{"m":8,"g":94},"156":{"m":8,"g":94},"155":{"m":8,"g":94},"130":{"m":8,"g":94},"141":{"m":8,"g":94},"153":{"m":8,"g":94},"148":{"m":8,"g":94},"146":{"m":8,"g":94},"144":{"m":8,"g":94},"142":{"m":8,"g":94},"137":{"m":8,"g":94},"136":{"m":8,"g":94},"280":{"m":9,"g":94},"279":{"m":9,"g":94},"230":{"m":9,"g":94},"277":{"m":9,"g":94},"278":{"m":9,"g":94},"256":{"m":9,"g":94},"222":{"m":9,"g":94},"261":{"m":9,"g":94},"275":{"m":9,"g":94},"263":{"m":9,"g":94},"201":{"m":9,"g":94},"224":{"m":9,"g":94},"253":{"m":9,"g":94},"226":{"m":9,"g":94},"195":{"m":9,"g":94},"198":{"m":9,"g":94},"225":{"m":9,"g":94},"219":{"m":9,"g":94},"193":{"m":9,"g":94},"210":{"m":9,"g":94},"207":{"m":9,"g":94},"200":{"m":9,"g":94},"196":{"m":9,"g":94},"189":{"m":9,"g":94},"186":{"m":9,"g":94},"184":{"m":9,"g":94},"181":{"m":9,"g":94},"182":{"m":9,"g":94},"324":{"m":10,"g":94},"323":{"m":10,"g":94},"301":{"m":10,"g":94},"304":{"m":10,"g":94},"311":{"m":10,"g":94},"291":{"m":10,"g":94},"290":{"m":10,"g":94},"288":{"m":10,"g":94},"286":{"m":10,"g":94},"287":{"m":10,"g":94},"282":{"m":10,"g":94},"242":{"m":10,"g":94},"281":{"m":10,"g":94},"431":{"m":11,"g":94},"430":{"m":11,"g":94},"429":{"m":11,"g":94},"428":{"m":11,"g":94},"427":{"m":11,"g":94},"422":{"m":11,"g":94},"420":{"m":11,"g":94},"380":{"m":11,"g":94},"416":{"m":11,"g":94},"415":{"m":11,"g":94},"412":{"m":11,"g":94},"411":{"m":11,"g":94},"381":{"m":11,"g":94},"392":{"m":11,"g":94},"390":{"m":11,"g":94},"406":{"m":11,"g":94},"399":{"m":11,"g":94},"395":{"m":11,"g":94},"394":{"m":11,"g":94},"382":{"m":11,"g":94},"385":{"m":11,"g":94},"378":{"m":11,"g":94},"372":{"m":11,"g":94},"375":{"m":11,"g":94},"364":{"m":11,"g":94},"370":{"m":11,"g":94},"368":{"m":11,"g":94},"369":{"m":11,"g":94},"358":{"m":11,"g":94},"355":{"m":11,"g":94},"354":{"m":11,"g":94},"346":{"m":11,"g":94},"338":{"m":11,"g":94},"345":{"m":11,"g":94},"343":{"m":11,"g":94},"315":{"m":11,"g":94},"332":{"m":11,"g":94},"337":{"m":11,"g":94},"331":{"m":11,"g":94},"329":{"m":11,"g":94},"293":{"m":11,"g":94},"327":{"m":11,"g":94},"326":{"m":11,"g":94},"298":{"m":11,"g":94},"438":{"m":12,"g":94},"437":{"m":12,"g":94},"426":{"m":12,"g":94},"436":{"m":12,"g":94},"434":{"m":12,"g":94},"418":{"m":12,"g":94},"433":{"m":12,"g":94},"363":{"m":12,"g":94},"432":{"m":12,"g":94},"515":{"m":13,"g":94},"514":{"m":13,"g":94},"505":{"m":13,"g":94},"512":{"m":13,"g":94},"502":{"m":13,"g":94},"500":{"m":13,"g":94},"511":{"m":13,"g":94},"493":{"m":13,"g":94},"491":{"m":13,"g":94},"492":{"m":13,"g":94},"488":{"m":13,"g":94},"486":{"m":13,"g":94},"480":{"m":13,"g":94},"484":{"m":13,"g":94},"477":{"m":13,"g":94},"475":{"m":13,"g":94},"476":{"m":13,"g":94},"440":{"m":13,"g":94},"471":{"m":13,"g":94},"463":{"m":13,"g":94},"470":{"m":13,"g":94},"460":{"m":13,"g":94},"459":{"m":13,"g":94},"458":{"m":13,"g":94},"457":{"m":13,"g":94},"456":{"m":13,"g":94},"250":{"m":13,"g":94},"451":{"m":13,"g":94},"449":{"m":13,"g":94},"448":{"m":13,"g":94},"447":{"m":13,"g":94},"446":{"m":13,"g":94},"441":{"m":13,"g":94},"419":{"m":13,"g":94},"579":{"m":14,"g":94},"585":{"m":14,"g":94},"583":{"m":14,"g":94},"578":{"m":14,"g":94},"577":{"m":14,"g":94},"576":{"m":14,"g":94},"574":{"m":14,"g":94},"545":{"m":14,"g":94},"571":{"m":14,"g":94},"569":{"m":14,"g":94},"568":{"m":14,"g":94},"567":{"m":14,"g":94},"566":{"m":14,"g":94},"564":{"m":14,"g":94},"563":{"m":14,"g":94},"561":{"m":14,"g":94},"560":{"m":14,"g":94},"559":{"m":14,"g":94},"558":{"m":14,"g":94},"557":{"m":14,"g":94},"556":{"m":14,"g":94},"554":{"m":14,"g":94},"550":{"m":14,"g":94},"553":{"m":14,"g":94},"551":{"m":14,"g":94},"546":{"m":14,"g":94},"542":{"m":14,"g":94},"540":{"m":14,"g":94},"539":{"m":14,"g":94},"538":{"m":14,"g":94},"517":{"m":14,"g":94},"531":{"m":14,"g":94},"516":{"m":14,"g":94},"526":{"m":14,"g":94},"525":{"m":14,"g":94},"524":{"m":14,"g":94},"518":{"m":14,"g":94},"605":{"m":15,"g":94},"503":{"m":15,"g":94},"530":{"m":15,"g":94},"603":{"m":15,"g":94},"598":{"m":15,"g":94},"602":{"m":15,"g":94},"604":{"m":15,"g":94},"601":{"m":15,"g":94},"600":{"m":15,"g":94},"599":{"m":15,"g":94},"586":{"m":15,"g":94},"594":{"m":15,"g":94},"593":{"m":15,"g":94},"592":{"m":15,"g":94},"588":{"m":15,"g":94},"618":{"m":16,"g":94},"616":{"m":16,"g":94},"615":{"m":16,"g":94},"614":{"m":16,"g":94},"613":{"m":16,"g":94},"612":{"m":16,"g":94},"611":{"m":16,"g":94},"610":{"m":16,"g":94},"609":{"m":16,"g":94},"607":{"m":16,"g":94},"626":{"m":17,"g":94},"625":{"m":17,"g":94},"623":{"m":17,"g":94},"621":{"m":17,"g":94},"620":{"m":17,"g":94},"619":{"m":17,"g":94},"677":{"m":18,"g":94},"676":{"m":18,"g":94},"675":{"m":18,"g":94},"664":{"m":18,"g":94},"673":{"m":18,"g":94},"671":{"m":18,"g":94},"669":{"m":18,"g":94},"668":{"m":18,"g":94},"667":{"m":18,"g":94},"640":{"m":18,"g":94},"666":{"m":18,"g":94},"665":{"m":18,"g":94},"663":{"m":18,"g":94},"662":{"m":18,"g":94},"661":{"m":18,"g":94},"660":{"m":18,"g":94},"659":{"m":18,"g":94},"655":{"m":18,"g":94},"657":{"m":18,"g":94},"658":{"m":18,"g":94},"656":{"m":18,"g":94},"654":{"m":18,"g":94},"653":{"m":18,"g":94},"651":{"m":18,"g":94},"650":{"m":18,"g":94},"648":{"m":18,"g":94},"647":{"m":18,"g":94},"649":{"m":18,"g":94},"646":{"m":18,"g":94},"645":{"m":18,"g":94},"643":{"m":18,"g":94},"642":{"m":18,"g":94},"617":{"m":18,"g":94},"638":{"m":18,"g":94},"637":{"m":18,"g":94},"636":{"m":18,"g":94},"635":{"m":18,"g":94},"632":{"m":18,"g":94},"633":{"m":18,"g":94},"624":{"m":18,"g":94},"630":{"m":18,"g":94},"631":{"m":18,"g":94},"629":{"m":18,"g":94},"628":{"m":18,"g":94},"627":{"m":18,"g":94},"705":{"m":19,"g":94},"704":{"m":19,"g":94},"701":{"m":19,"g":94},"702":{"m":19,"g":94},"700":{"m":19,"g":94},"698":{"m":19,"g":94},"697":{"m":19,"g":94},"696":{"m":19,"g":94},"695":{"m":19,"g":94},"694":{"m":19,"g":94},"692":{"m":19,"g":94},"691":{"m":19,"g":94},"690":{"m":19,"g":94},"689":{"m":19,"g":94},"688":{"m":19,"g":94},"687":{"m":19,"g":94},"686":{"m":19,"g":94},"685":{"m":19,"g":94},"684":{"m":19,"g":94},"682":{"m":19,"g":94},"681":{"m":19,"g":94},"679":{"m":19,"g":94},"670":{"m":19,"g":94},"678":{"m":19,"g":94},"718":{"m":20,"g":94},"717":{"m":20,"g":94},"716":{"m":20,"g":94},"715":{"m":20,"g":94},"714":{"m":20,"g":94},"713":{"m":20,"g":94},"712":{"m":20,"g":94},"711":{"m":20,"g":94},"708":{"m":20,"g":94},"709":{"m":20,"g":94},"707":{"m":20,"g":94},"706":{"m":20,"g":94},"730":{"m":21,"g":94},"729":{"m":21,"g":94},"728":{"m":21,"g":94},"727":{"m":21,"g":94},"726":{"m":21,"g":94},"725":{"m":21,"g":94},"720":{"m":21,"g":94},"723":{"m":21,"g":94},"724":{"m":21,"g":94},"722":{"m":21,"g":94},"721":{"m":21,"g":94},"719":{"m":21,"g":94},"755":{"m":22,"g":94},"754":{"m":22,"g":94},"753":{"m":22,"g":94},"752":{"m":22,"g":94},"751":{"m":22,"g":94},"740":{"m":22,"g":94},"743":{"m":22,"g":94},"742":{"m":22,"g":94},"739":{"m":22,"g":94},"741":{"m":22,"g":94},"736":{"m":22,"g":94},"734":{"m":22,"g":94},"733":{"m":22,"g":94},"731":{"m":22,"g":94},"779":{"m":23,"g":94},"778":{"m":23,"g":94},"776":{"m":23,"g":94},"775":{"m":23,"g":94},"774":{"m":23,"g":94},"773":{"m":23,"g":94},"772":{"m":23,"g":94},"766":{"m":23,"g":94},"770":{"m":23,"g":94},"769":{"m":23,"g":94},"767":{"m":23,"g":94},"761":{"m":23,"g":94},"763":{"m":23,"g":94},"762":{"m":23,"g":94},"760":{"m":23,"g":94},"757":{"m":23,"g":94},"693":{"m":23,"g":94},"830":{"m":24,"g":94},"829":{"m":24,"g":94},"828":{"m":24,"g":94},"825":{"m":24,"g":94},"826":{"m":24,"g":94},"823":{"m":24,"g":94},"824":{"m":24,"g":94},"822":{"m":24,"g":94},"821":{"m":24,"g":94},"820":{"m":24,"g":94},"819":{"m":24,"g":94},"807":{"m":24,"g":94},"817":{"m":24,"g":94},"814":{"m":24,"g":94},"815":{"m":24,"g":94},"812":{"m":24,"g":94},"809":{"m":24,"g":94},"699":{"m":24,"g":94},"806":{"m":24,"g":94},"802":{"m":24,"g":94},"805":{"m":24,"g":94},"803":{"m":24,"g":94},"800":{"m":24,"g":94},"799":{"m":24,"g":94},"797":{"m":24,"g":94},"793":{"m":24,"g":94},"796":{"m":24,"g":94},"795":{"m":24,"g":94},"794":{"m":24,"g":94},"792":{"m":24,"g":94},"791":{"m":24,"g":94},"790":{"m":24,"g":94},"789":{"m":24,"g":94},"788":{"m":24,"g":94},"787":{"m":24,"g":94},"786":{"m":24,"g":94},"785":{"m":24,"g":94},"784":{"m":24,"g":94},"783":{"m":24,"g":94},"781":{"m":24,"g":94},"877":{"m":25,"g":94},"872":{"m":25,"g":94},"873":{"m":25,"g":94},"871":{"m":25,"g":94},"870":{"m":25,"g":94},"869":{"m":25,"g":94},"864":{"m":25,"g":94},"862":{"m":25,"g":94},"861":{"m":25,"g":94},"860":{"m":25,"g":94},"811":{"m":25,"g":94},"852":{"m":25,"g":94},"858":{"m":25,"g":94},"856":{"m":25,"g":94},"855":{"m":25,"g":94},"850":{"m":25,"g":94},"848":{"m":25,"g":94},"843":{"m":25,"g":94},"842":{"m":25,"g":94},"838":{"m":25,"g":94},"840":{"m":25,"g":94},"890":{"m":26,"g":94},"886":{"m":26,"g":94},"889":{"m":26,"g":94},"888":{"m":26,"g":94},"883":{"m":26,"g":94},"882":{"m":26,"g":94},"880":{"m":26,"g":94},"879":{"m":26,"g":94},"749":{"m":26,"g":94},"878":{"m":26,"g":94},"876":{"m":26,"g":94},"875":{"m":26,"g":94},"899":{"m":27,"g":94},"896":{"m":27,"g":94},"895":{"m":27,"g":94},"894":{"m":27,"g":94},"884":{"m":27,"g":94},"891":{"m":27,"g":94},"923":{"m":28,"g":94},"916":{"m":28,"g":94},"920":{"m":28,"g":94},"918":{"m":28,"g":94},"917":{"m":28,"g":94},"915":{"m":28,"g":94},"905":{"m":28,"g":94},"914":{"m":28,"g":94},"912":{"m":28,"g":94},"911":{"m":28,"g":94},"909":{"m":28,"g":94},"866":{"m":28,"g":94},"904":{"m":28,"g":94},"908":{"m":28,"g":94},"900":{"m":28,"g":94},"970":{"m":29,"g":94},"966":{"m":29,"g":94},"967":{"m":29,"g":94},"960":{"m":29,"g":94},"965":{"m":29,"g":94},"964":{"m":29,"g":94},"963":{"m":29,"g":94},"932":{"m":29,"g":94},"936":{"m":29,"g":94},"957":{"m":29,"g":94},"953":{"m":29,"g":94},"948":{"m":29,"g":94},"941":{"m":29,"g":94},"940":{"m":29,"g":94},"835":{"m":29,"g":94},"934":{"m":29,"g":94},"935":{"m":29,"g":94},"928":{"m":29,"g":94},"927":{"m":29,"g":94},"926":{"m":29,"g":94},"925":{"m":29,"g":94},"921":{"m":29,"g":94},"1048":{"m":30,"g":94},"1052":{"m":30,"g":94},"1051":{"m":30,"g":94},"1049":{"m":30,"g":94},"1050":{"m":30,"g":94},"1033":{"m":30,"g":94},"1046":{"m":30,"g":94},"1047":{"m":30,"g":94},"1044":{"m":30,"g":94},"1045":{"m":30,"g":94},"1039":{"m":30,"g":94},"1037":{"m":30,"g":94},"1038":{"m":30,"g":94},"1025":{"m":30,"g":94},"1034":{"m":30,"g":94},"1031":{"m":30,"g":94},"1027":{"m":30,"g":94},"1029":{"m":30,"g":94},"1028":{"m":30,"g":94},"1022":{"m":30,"g":94},"1024":{"m":30,"g":94},"907":{"m":30,"g":94},"1021":{"m":30,"g":94},"1020":{"m":30,"g":94},"1019":{"m":30,"g":94},"990":{"m":30,"g":94},"1014":{"m":30,"g":94},"1010":{"m":30,"g":94},"1009":{"m":30,"g":94},"1007":{"m":30,"g":94},"959":{"m":30,"g":94},"997":{"m":30,"g":94},"1005":{"m":30,"g":94},"1002":{"m":30,"g":94},"1001":{"m":30,"g":94},"994":{"m":30,"g":94},"995":{"m":30,"g":94},"988":{"m":30,"g":94},"993":{"m":30,"g":94},"973":{"m":30,"g":94},"992":{"m":30,"g":94},"985":{"m":30,"g":94},"981":{"m":30,"g":94},"987":{"m":30,"g":94},"983":{"m":30,"g":94},"984":{"m":30,"g":94},"982":{"m":30,"g":94},"971":{"m":30,"g":94},"980":{"m":30,"g":94},"977":{"m":30,"g":94},"968":{"m":30,"g":94},"969":{"m":30,"g":94},"976":{"m":30,"g":94},"975":{"m":30,"g":94},"1111":{"m":31,"g":94},"1113":{"m":31,"g":94},"1112":{"m":31,"g":94},"1110":{"m":31,"g":94},"1107":{"m":31,"g":94},"1040":{"m":31,"g":94},"1106":{"m":31,"g":94},"1077":{"m":31,"g":94},"1104":{"m":31,"g":94},"1092":{"m":31,"g":94},"1103":{"m":31,"g":94},"1090":{"m":31,"g":94},"1082":{"m":31,"g":94},"1099":{"m":31,"g":94},"1095":{"m":31,"g":94},"1098":{"m":31,"g":94},"1096":{"m":31,"g":94},"1094":{"m":31,"g":94},"1088":{"m":31,"g":94},"1086":{"m":31,"g":94},"1084":{"m":31,"g":94},"1056":{"m":31,"g":94},"1081":{"m":31,"g":94},"1006":{"m":31,"g":94},"1079":{"m":31,"g":94},"1078":{"m":31,"g":94},"1074":{"m":31,"g":94},"1053":{"m":31,"g":94},"1070":{"m":31,"g":94},"1060":{"m":31,"g":94},"1066":{"m":31,"g":94},"1068":{"m":31,"g":94},"1057":{"m":31,"g":94},"1155":{"m":32,"g":94},"1201":{"m":32,"g":94},"1219":{"m":32,"g":94},"1212":{"m":32,"g":94},"1218":{"m":32,"g":94},"1217":{"m":32,"g":94},"1215":{"m":32,"g":94},"1214":{"m":32,"g":94},"1204":{"m":32,"g":94},"1213":{"m":32,"g":94},"1210":{"m":32,"g":94},"1211":{"m":32,"g":94},"1209":{"m":32,"g":94},"1208":{"m":32,"g":94},"1186":{"m":32,"g":94},"1205":{"m":32,"g":94},"1207":{"m":32,"g":94},"1199":{"m":32,"g":94},"1202":{"m":32,"g":94},"1198":{"m":32,"g":94},"1194":{"m":32,"g":94},"1193":{"m":32,"g":94},"1123":{"m":32,"g":94},"1185":{"m":32,"g":94},"1184":{"m":32,"g":94},"1180":{"m":32,"g":94},"1168":{"m":32,"g":94},"1179":{"m":32,"g":94},"1167":{"m":32,"g":94},"1170":{"m":32,"g":94},"1177":{"m":32,"g":94},"1171":{"m":32,"g":94},"1157":{"m":32,"g":94},"1166":{"m":32,"g":94},"1154":{"m":32,"g":94},"1165":{"m":32,"g":94},"1148":{"m":32,"g":94},"1164":{"m":32,"g":94},"1134":{"m":32,"g":94},"1138":{"m":32,"g":94},"1035":{"m":32,"g":94},"1144":{"m":32,"g":94},"1143":{"m":32,"g":94},"1141":{"m":32,"g":94},"1140":{"m":32,"g":94},"1139":{"m":32,"g":94},"1136":{"m":32,"g":94},"1133":{"m":32,"g":94},"1131":{"m":32,"g":94},"1013":{"m":32,"g":94},"1122":{"m":32,"g":94},"1119":{"m":32,"g":94},"1115":{"m":32,"g":94},"1114":{"m":32,"g":94},"1242":{"m":33,"g":94},"1239":{"m":33,"g":94},"1233":{"m":33,"g":94},"1237":{"m":33,"g":94},"1225":{"m":33,"g":94},"1236":{"m":33,"g":94},"1231":{"m":33,"g":94},"1230":{"m":33,"g":94},"1227":{"m":33,"g":94},"1222":{"m":33,"g":94},"1223":{"m":33,"g":94},"1125":{"m":33,"g":94},"1250":{"m":34,"g":94},"1252":{"m":34,"g":94},"1249":{"m":34,"g":94},"1234":{"m":34,"g":94},"1247":{"m":34,"g":94},"1232":{"m":34,"g":94},"1244":{"m":34,"g":94},"1243":{"m":34,"g":94},"1295":{"m":35,"g":94},"1297":{"m":35,"g":94},"1296":{"m":35,"g":94},"1294":{"m":35,"g":94},"1293":{"m":35,"g":94},"1291":{"m":35,"g":94},"1290":{"m":35,"g":94},"1277":{"m":35,"g":94},"1284":{"m":35,"g":94},"1288":{"m":35,"g":94},"1286":{"m":35,"g":94},"1289":{"m":35,"g":94},"1285":{"m":35,"g":94},"1280":{"m":35,"g":94},"1262":{"m":35,"g":94},"1282":{"m":35,"g":94},"1276":{"m":35,"g":94},"1256":{"m":35,"g":94},"1269":{"m":35,"g":94},"1267":{"m":35,"g":94},"1258":{"m":35,"g":94},"1261":{"m":35,"g":94},"1260":{"m":35,"g":94},"1253":{"m":35,"g":94},"1255":{"m":35,"g":94},"1254":{"m":35,"g":94},"1327":{"m":36,"g":94},"1326":{"m":36,"g":94},"1320":{"m":36,"g":94},"1319":{"m":36,"g":94},"1318":{"m":36,"g":94},"1317":{"m":36,"g":94},"1313":{"m":36,"g":94},"1299":{"m":36,"g":94},"1308":{"m":36,"g":94},"1306":{"m":36,"g":94},"1304":{"m":36,"g":94},"1445":{"m":37,"g":94},"1444":{"m":37,"g":94},"1442":{"m":37,"g":94},"1420":{"m":37,"g":94},"1441":{"m":37,"g":94},"1440":{"m":37,"g":94},"1438":{"m":37,"g":94},"1432":{"m":37,"g":94},"1433":{"m":37,"g":94},"1431":{"m":37,"g":94},"1430":{"m":37,"g":94},"1428":{"m":37,"g":94},"1429":{"m":37,"g":94},"1427":{"m":37,"g":94},"1422":{"m":37,"g":94},"1426":{"m":37,"g":94},"1425":{"m":37,"g":94},"1418":{"m":37,"g":94},"1392":{"m":37,"g":94},"1414":{"m":37,"g":94},"1412":{"m":37,"g":94},"1411":{"m":37,"g":94},"1409":{"m":37,"g":94},"1408":{"m":37,"g":94},"1407":{"m":37,"g":94},"1406":{"m":37,"g":94},"1307":{"m":37,"g":94},"1397":{"m":37,"g":94},"1402":{"m":37,"g":94},"1403":{"m":37,"g":94},"1401":{"m":37,"g":94},"1399":{"m":37,"g":94},"1393":{"m":37,"g":94},"1381":{"m":37,"g":94},"1390":{"m":37,"g":94},"1389":{"m":37,"g":94},"1367":{"m":37,"g":94},"1385":{"m":37,"g":94},"1378":{"m":37,"g":94},"1380":{"m":37,"g":94},"1379":{"m":37,"g":94},"1376":{"m":37,"g":94},"1375":{"m":37,"g":94},"1373":{"m":37,"g":94},"1371":{"m":37,"g":94},"1370":{"m":37,"g":94},"1368":{"m":37,"g":94},"1300":{"m":37,"g":94},"1363":{"m":37,"g":94},"1360":{"m":37,"g":94},"1361":{"m":37,"g":94},"1341":{"m":37,"g":94},"1357":{"m":37,"g":94},"1298":{"m":37,"g":94},"1346":{"m":37,"g":94},"1281":{"m":37,"g":94},"1345":{"m":37,"g":94},"1339":{"m":37,"g":94},"1340":{"m":37,"g":94},"1337":{"m":37,"g":94},"1336":{"m":37,"g":94},"1328":{"m":37,"g":94},"1470":{"m":38,"g":94},"1469":{"m":38,"g":94},"1464":{"m":38,"g":94},"1458":{"m":38,"g":94},"1457":{"m":38,"g":94},"1454":{"m":38,"g":94},"1453":{"m":38,"g":94},"1452":{"m":38,"g":94},"1449":{"m":38,"g":94},"1451":{"m":38,"g":94},"1450":{"m":38,"g":94},"1448":{"m":38,"g":94},"1447":{"m":38,"g":94},"1483":{"m":39,"g":94},"1484":{"m":39,"g":94},"1482":{"m":39,"g":94},"1476":{"m":39,"g":94},"1475":{"m":39,"g":94},"1305":{"m":39,"g":94},"1472":{"m":39,"g":94},"1512":{"m":40,"g":94},"1511":{"m":40,"g":94},"1508":{"m":40,"g":94},"1510":{"m":40,"g":94},"1503":{"m":40,"g":94},"1499":{"m":40,"g":94},"1502":{"m":40,"g":94},"1500":{"m":40,"g":94},"1497":{"m":40,"g":94},"1496":{"m":40,"g":94},"1494":{"m":40,"g":94},"1490":{"m":40,"g":94},"1492":{"m":40,"g":94},"1491":{"m":40,"g":94},"1489":{"m":40,"g":94},"1456":{"m":40,"g":94},"1488":{"m":40,"g":94},"1486":{"m":40,"g":94},"1481":{"m":40,"g":94},"1605":{"m":41,"g":94},"1606":{"m":41,"g":94},"1604":{"m":41,"g":94},"1598":{"m":41,"g":94},"1603":{"m":41,"g":94},"1597":{"m":41,"g":94},"1594":{"m":41,"g":94},"1596":{"m":41,"g":94},"1595":{"m":41,"g":94},"1593":{"m":41,"g":94},"1567":{"m":41,"g":94},"1592":{"m":41,"g":94},"1591":{"m":41,"g":94},"1590":{"m":41,"g":94},"1589":{"m":41,"g":94},"1587":{"m":41,"g":94},"1586":{"m":41,"g":94},"1585":{"m":41,"g":94},"1584":{"m":41,"g":94},"1573":{"m":41,"g":94},"1582":{"m":41,"g":94},"1583":{"m":41,"g":94},"1581":{"m":41,"g":94},"1576":{"m":41,"g":94},"1561":{"m":41,"g":94},"1572":{"m":41,"g":94},"1577":{"m":41,"g":94},"1580":{"m":41,"g":94},"1574":{"m":41,"g":94},"1563":{"m":41,"g":94},"1569":{"m":41,"g":94},"1568":{"m":41,"g":94},"1566":{"m":41,"g":94},"1562":{"m":41,"g":94},"1559":{"m":41,"g":94},"1536":{"m":41,"g":94},"1557":{"m":41,"g":94},"1556":{"m":41,"g":94},"1555":{"m":41,"g":94},"1553":{"m":41,"g":94},"1554":{"m":41,"g":94},"1549":{"m":41,"g":94},"1552":{"m":41,"g":94},"1550":{"m":41,"g":94},"1548":{"m":41,"g":94},"1547":{"m":41,"g":94},"1545":{"m":41,"g":94},"1544":{"m":41,"g":94},"1543":{"m":41,"g":94},"1541":{"m":41,"g":94},"1539":{"m":41,"g":94},"1538":{"m":41,"g":94},"1537":{"m":41,"g":94},"1534":{"m":41,"g":94},"1531":{"m":41,"g":94},"1532":{"m":41,"g":94},"1530":{"m":41,"g":94},"1495":{"m":41,"g":94},"1520":{"m":41,"g":94},"1521":{"m":41,"g":94},"1528":{"m":41,"g":94},"1525":{"m":41,"g":94},"1529":{"m":41,"g":94},"1524":{"m":41,"g":94},"1513":{"m":41,"g":94},"1636":{"m":42,"g":94},"1635":{"m":42,"g":94},"1634":{"m":42,"g":94},"1633":{"m":42,"g":94},"1632":{"m":42,"g":94},"1631":{"m":42,"g":94},"1626":{"m":42,"g":94},"1579":{"m":42,"g":94},"1611":{"m":42,"g":94},"1607":{"m":42,"g":94},"1629":{"m":42,"g":94},"1625":{"m":42,"g":94},"1619":{"m":42,"g":94},"1620":{"m":42,"g":94},"1615":{"m":42,"g":94},"1714":{"m":43,"g":94},"1713":{"m":43,"g":94},"1712":{"m":43,"g":94},"1710":{"m":43,"g":94},"1709":{"m":43,"g":94},"1707":{"m":43,"g":94},"1706":{"m":43,"g":94},"1705":{"m":43,"g":94},"1704":{"m":43,"g":94},"1703":{"m":43,"g":94},"1684":{"m":43,"g":94},"1702":{"m":43,"g":94},"1701":{"m":43,"g":94},"1700":{"m":43,"g":94},"1699":{"m":43,"g":94},"1694":{"m":43,"g":94},"1697":{"m":43,"g":94},"1696":{"m":43,"g":94},"1690":{"m":43,"g":94},"1679":{"m":43,"g":94},"1689":{"m":43,"g":94},"1688":{"m":43,"g":94},"1599":{"m":43,"g":94},"1687":{"m":43,"g":94},"1686":{"m":43,"g":94},"1685":{"m":43,"g":94},"1677":{"m":43,"g":94},"1676":{"m":43,"g":94},"1681":{"m":43,"g":94},"1674":{"m":43,"g":94},"1672":{"m":43,"g":94},"1671":{"m":43,"g":94},"1670":{"m":43,"g":94},"1658":{"m":43,"g":94},"1667":{"m":43,"g":94},"1666":{"m":43,"g":94},"1665":{"m":43,"g":94},"1459":{"m":43,"g":94},"1663":{"m":43,"g":94},"1662":{"m":43,"g":94},"1661":{"m":43,"g":94},"1659":{"m":43,"g":94},"1650":{"m":43,"g":94},"1656":{"m":43,"g":94},"1652":{"m":43,"g":94},"1654":{"m":43,"g":94},"1653":{"m":43,"g":94},"1651":{"m":43,"g":94},"1648":{"m":43,"g":94},"1480":{"m":43,"g":94},"1645":{"m":43,"g":94},"1642":{"m":43,"g":94},"1638":{"m":43,"g":94},"1614":{"m":43,"g":94},"1749":{"m":44,"g":94},"1748":{"m":44,"g":94},"1551":{"m":44,"g":94},"1746":{"m":44,"g":94},"1737":{"m":44,"g":94},"1738":{"m":44,"g":94},"1743":{"m":44,"g":94},"1741":{"m":44,"g":94},"1740":{"m":44,"g":94},"1736":{"m":44,"g":94},"1735":{"m":44,"g":94},"1734":{"m":44,"g":94},"1727":{"m":44,"g":94},"1726":{"m":44,"g":94},"1725":{"m":44,"g":94},"1724":{"m":44,"g":94},"1722":{"m":44,"g":94},"1721":{"m":44,"g":94},"1720":{"m":44,"g":94},"1718":{"m":44,"g":94},"1716":{"m":44,"g":94},"1796":{"m":45,"g":94},"1795":{"m":45,"g":94},"1797":{"m":45,"g":94},"1794":{"m":45,"g":94},"1787":{"m":45,"g":94},"1793":{"m":45,"g":94},"1780":{"m":45,"g":94},"1789":{"m":45,"g":94},"1785":{"m":45,"g":94},"1783":{"m":45,"g":94},"1778":{"m":45,"g":94},"1782":{"m":45,"g":94},"1779":{"m":45,"g":94},"1776":{"m":45,"g":94},"1774":{"m":45,"g":94},"1773":{"m":45,"g":94},"1772":{"m":45,"g":94},"1771":{"m":45,"g":94},"1769":{"m":45,"g":94},"1768":{"m":45,"g":94},"1766":{"m":45,"g":94},"1767":{"m":45,"g":94},"1765":{"m":45,"g":94},"1760":{"m":45,"g":94},"1758":{"m":45,"g":94},"1747":{"m":45,"g":94},"1908":{"m":46,"g":94},"1907":{"m":46,"g":94},"1906":{"m":46,"g":94},"1902":{"m":46,"g":94},"1905":{"m":46,"g":94},"1904":{"m":46,"g":94},"1903":{"m":46,"g":94},"1899":{"m":46,"g":94},"1896":{"m":46,"g":94},"1895":{"m":46,"g":94},"1894":{"m":46,"g":94},"1892":{"m":46,"g":94},"1890":{"m":46,"g":94},"1888":{"m":46,"g":94},"1889":{"m":46,"g":94},"1886":{"m":46,"g":94},"1885":{"m":46,"g":94},"1883":{"m":46,"g":94},"1873":{"m":46,"g":94},"1882":{"m":46,"g":94},"1881":{"m":46,"g":94},"1879":{"m":46,"g":94},"1878":{"m":46,"g":94},"1877":{"m":46,"g":94},"1875":{"m":46,"g":94},"1871":{"m":46,"g":94},"1867":{"m":46,"g":94},"1866":{"m":46,"g":94},"1754":{"m":46,"g":94},"1856":{"m":46,"g":94},"1859":{"m":46,"g":94},"1860":{"m":46,"g":94},"1861":{"m":46,"g":94},"1858":{"m":46,"g":94},"1855":{"m":46,"g":94},"1852":{"m":46,"g":94},"1851":{"m":46,"g":94},"1846":{"m":46,"g":94},"1850":{"m":46,"g":94},"1847":{"m":46,"g":94},"1845":{"m":46,"g":94},"1842":{"m":46,"g":94},"1836":{"m":46,"g":94},"1838":{"m":46,"g":94},"1840":{"m":46,"g":94},"1839":{"m":46,"g":94},"1827":{"m":46,"g":94},"1833":{"m":46,"g":94},"1835":{"m":46,"g":94},"1834":{"m":46,"g":94},"1822":{"m":46,"g":94},"1823":{"m":46,"g":94},"1830":{"m":46,"g":94},"1825":{"m":46,"g":94},"1790":{"m":46,"g":94},"1821":{"m":46,"g":94},"1820":{"m":46,"g":94},"1819":{"m":46,"g":94},"1810":{"m":46,"g":94},"1817":{"m":46,"g":94},"1816":{"m":46,"g":94},"1813":{"m":46,"g":94},"1811":{"m":46,"g":94},"1809":{"m":46,"g":94},"1808":{"m":46,"g":94},"1807":{"m":46,"g":94},"1805":{"m":46,"g":94},"1804":{"m":46,"g":94},"1803":{"m":46,"g":94},"1802":{"m":46,"g":94},"1801":{"m":46,"g":94},"1800":{"m":46,"g":94},"1786":{"m":46,"g":94},"1799":{"m":46,"g":94},"1798":{"m":46,"g":94},"1752":{"m":46,"g":94},"2022":{"m":47,"g":94},"2020":{"m":47,"g":94},"2018":{"m":47,"g":94},"1996":{"m":47,"g":94},"2015":{"m":47,"g":94},"2014":{"m":47,"g":94},"2013":{"m":47,"g":94},"1998":{"m":47,"g":94},"2011":{"m":47,"g":94},"2010":{"m":47,"g":94},"2009":{"m":47,"g":94},"2008":{"m":47,"g":94},"2006":{"m":47,"g":94},"2005":{"m":47,"g":94},"1994":{"m":47,"g":94},"2003":{"m":47,"g":94},"2004":{"m":47,"g":94},"2002":{"m":47,"g":94},"2001":{"m":47,"g":94},"2000":{"m":47,"g":94},"1999":{"m":47,"g":94},"1995":{"m":47,"g":94},"1997":{"m":47,"g":94},"1934":{"m":47,"g":94},"1980":{"m":47,"g":94},"1990":{"m":47,"g":94},"1988":{"m":47,"g":94},"1984":{"m":47,"g":94},"1986":{"m":47,"g":94},"1983":{"m":47,"g":94},"1981":{"m":47,"g":94},"1982":{"m":47,"g":94},"1977":{"m":47,"g":94},"1972":{"m":47,"g":94},"1976":{"m":47,"g":94},"1975":{"m":47,"g":94},"1974":{"m":47,"g":94},"1745":{"m":47,"g":94},"1973":{"m":47,"g":94},"1963":{"m":47,"g":94},"1966":{"m":47,"g":94},"1962":{"m":47,"g":94},"1961":{"m":47,"g":94},"1958":{"m":47,"g":94},"1957":{"m":47,"g":94},"1956":{"m":47,"g":94},"1955":{"m":47,"g":94},"1954":{"m":47,"g":94},"1933":{"m":47,"g":94},"1952":{"m":47,"g":94},"1951":{"m":47,"g":94},"1939":{"m":47,"g":94},"1941":{"m":47,"g":94},"1949":{"m":47,"g":94},"1940":{"m":47,"g":94},"1942":{"m":47,"g":94},"1891":{"m":47,"g":94},"1926":{"m":47,"g":94},"1922":{"m":47,"g":94},"1853":{"m":47,"g":94},"1924":{"m":47,"g":94},"1920":{"m":47,"g":94},"1916":{"m":47,"g":94},"1893":{"m":47,"g":94},"1915":{"m":47,"g":94},"1910":{"m":47,"g":94},"1909":{"m":47,"g":94},"2046":{"m":48,"g":94},"2044":{"m":48,"g":94},"2043":{"m":48,"g":94},"2030":{"m":48,"g":94},"2042":{"m":48,"g":94},"2038":{"m":48,"g":94},"1968":{"m":48,"g":94},"2039":{"m":48,"g":94},"2036":{"m":48,"g":94},"2034":{"m":48,"g":94},"2033":{"m":48,"g":94},"2031":{"m":48,"g":94},"2027":{"m":48,"g":94},"2028":{"m":48,"g":94},"2026":{"m":48,"g":94},"2024":{"m":48,"g":94},"2023":{"m":48,"g":94},"2120":{"m":49,"g":94},"2125":{"m":49,"g":94},"2122":{"m":49,"g":94},"2110":{"m":49,"g":94},"2118":{"m":49,"g":94},"2106":{"m":49,"g":94},"2115":{"m":49,"g":94},"2055":{"m":49,"g":94},"2111":{"m":49,"g":94},"2116":{"m":49,"g":94},"2107":{"m":49,"g":94},"2105":{"m":49,"g":94},"2104":{"m":49,"g":94},"2103":{"m":49,"g":94},"2073":{"m":49,"g":94},"2100":{"m":49,"g":94},"2067":{"m":49,"g":94},"2096":{"m":49,"g":94},"2095":{"m":49,"g":94},"2093":{"m":49,"g":94},"2094":{"m":49,"g":94},"2091":{"m":49,"g":94},"2088":{"m":49,"g":94},"2089":{"m":49,"g":94},"2086":{"m":49,"g":94},"2085":{"m":49,"g":94},"2083":{"m":49,"g":94},"2078":{"m":49,"g":94},"2069":{"m":49,"g":94},"2075":{"m":49,"g":94},"2074":{"m":49,"g":94},"2072":{"m":49,"g":94},"2071":{"m":49,"g":94},"2070":{"m":49,"g":94},"2068":{"m":49,"g":94},"2062":{"m":49,"g":94},"2056":{"m":49,"g":94},"2066":{"m":49,"g":94},"2061":{"m":49,"g":94},"2065":{"m":49,"g":94},"2064":{"m":49,"g":94},"2063":{"m":49,"g":94},"1849":{"m":49,"g":94},"2053":{"m":49,"g":94},"2051":{"m":49,"g":94},"1970":{"m":49,"g":94},"2050":{"m":49,"g":94},"2049":{"m":49,"g":94},"1876":{"m":49,"g":94},"2048":{"m":49,"g":94},"2047":{"m":49,"g":94},"2189":{"m":50,"g":94},"2188":{"m":50,"g":94},"2187":{"m":50,"g":94},"2052":{"m":50,"g":94},"2184":{"m":50,"g":94},"2176":{"m":50,"g":94},"2186":{"m":50,"g":94},"2171":{"m":50,"g":94},"2185":{"m":50,"g":94},"2183":{"m":50,"g":94},"2173":{"m":50,"g":94},"2182":{"m":50,"g":94},"2180":{"m":50,"g":94},"2175":{"m":50,"g":94},"2174":{"m":50,"g":94},"2170":{"m":50,"g":94},"2169":{"m":50,"g":94},"2167":{"m":50,"g":94},"2164":{"m":50,"g":94},"2163":{"m":50,"g":94},"2162":{"m":50,"g":94},"2158":{"m":50,"g":94},"2161":{"m":50,"g":94},"2159":{"m":50,"g":94},"2156":{"m":50,"g":94},"2154":{"m":50,"g":94},"2157":{"m":50,"g":94},"2155":{"m":50,"g":94},"2153":{"m":50,"g":94},"2152":{"m":50,"g":94},"2148":{"m":50,"g":94},"2147":{"m":50,"g":94},"2146":{"m":50,"g":94},"2144":{"m":50,"g":94},"2143":{"m":50,"g":94},"2142":{"m":50,"g":94},"2114":{"m":50,"g":94},"2139":{"m":50,"g":94},"2138":{"m":50,"g":94},"2136":{"m":50,"g":94},"2137":{"m":50,"g":94},"2134":{"m":50,"g":94},"2081":{"m":50,"g":94},"2121":{"m":50,"g":94},"2130":{"m":50,"g":94},"2124":{"m":50,"g":94},"2092":{"m":50,"g":94},"2127":{"m":50,"g":94},"2077":{"m":50,"g":94},"2126":{"m":50,"g":94},"2214":{"m":51,"g":94},"2222":{"m":51,"g":94},"2221":{"m":51,"g":94},"2217":{"m":51,"g":94},"2210":{"m":51,"g":94},"2212":{"m":51,"g":94},"2208":{"m":51,"g":94},"2207":{"m":51,"g":94},"2204":{"m":51,"g":94},"2206":{"m":51,"g":94},"2201":{"m":51,"g":94},"2199":{"m":51,"g":94},"2196":{"m":51,"g":94},"2198":{"m":51,"g":94},"2195":{"m":51,"g":94},"2197":{"m":51,"g":94},"2259":{"m":52,"g":94},"2257":{"m":52,"g":94},"2256":{"m":52,"g":94},"2254":{"m":52,"g":94},"2253":{"m":52,"g":94},"2252":{"m":52,"g":94},"2251":{"m":52,"g":94},"2250":{"m":52,"g":94},"2239":{"m":52,"g":94},"2242":{"m":52,"g":94},"2238":{"m":52,"g":94},"2231":{"m":52,"g":94},"2233":{"m":52,"g":94},"2235":{"m":52,"g":94},"2234":{"m":52,"g":94},"2232":{"m":52,"g":94},"2228":{"m":52,"g":94},"2123":{"m":52,"g":94},"2223":{"m":52,"g":94},"2224":{"m":52,"g":94},"2226":{"m":52,"g":94},"2191":{"m":52,"g":94},"2218":{"m":52,"g":94},"2338":{"m":53,"g":94},"2339":{"m":53,"g":94},"2300":{"m":53,"g":94},"2335":{"m":53,"g":94},"2327":{"m":53,"g":94},"2324":{"m":53,"g":94},"2328":{"m":53,"g":94},"2329":{"m":53,"g":94},"2325":{"m":53,"g":94},"2281":{"m":53,"g":94},"2319":{"m":53,"g":94},"2318":{"m":53,"g":94},"2314":{"m":53,"g":94},"2311":{"m":53,"g":94},"2310":{"m":53,"g":94},"2309":{"m":53,"g":94},"2279":{"m":53,"g":94},"2306":{"m":53,"g":94},"2305":{"m":53,"g":94},"2304":{"m":53,"g":94},"2301":{"m":53,"g":94},"2302":{"m":53,"g":94},"2299":{"m":53,"g":94},"2298":{"m":53,"g":94},"2241":{"m":53,"g":94},"2295":{"m":53,"g":94},"2292":{"m":53,"g":94},"2179":{"m":53,"g":94},"2293":{"m":53,"g":94},"2290":{"m":53,"g":94},"2244":{"m":53,"g":94},"2288":{"m":53,"g":94},"2289":{"m":53,"g":94},"2287":{"m":53,"g":94},"2284":{"m":53,"g":94},"2286":{"m":53,"g":94},"2285":{"m":53,"g":94},"2282":{"m":53,"g":94},"2215":{"m":53,"g":94},"2280":{"m":53,"g":94},"2274":{"m":53,"g":94},"2266":{"m":53,"g":94},"2265":{"m":53,"g":94},"2269":{"m":53,"g":94},"2243":{"m":53,"g":94},"2268":{"m":53,"g":94},"2225":{"m":53,"g":94},"2261":{"m":53,"g":94},"2375":{"m":54,"g":94},"2377":{"m":54,"g":94},"2363":{"m":54,"g":94},"2374":{"m":54,"g":94},"2373":{"m":54,"g":94},"2369":{"m":54,"g":94},"2357":{"m":54,"g":94},"2360":{"m":54,"g":94},"2370":{"m":54,"g":94},"2371":{"m":54,"g":94},"2368":{"m":54,"g":94},"2359":{"m":54,"g":94},"2364":{"m":54,"g":94},"2355":{"m":54,"g":94},"2342":{"m":54,"g":94},"2352":{"m":54,"g":94},"2308":{"m":54,"g":94},"2323":{"m":54,"g":94},"2350":{"m":54,"g":94},"2340":{"m":54,"g":94},"2349":{"m":54,"g":94},"2341":{"m":54,"g":94},"2348":{"m":54,"g":94},"2525":{"m":55,"g":94},"2528":{"m":55,"g":94},"2524":{"m":55,"g":94},"2517":{"m":55,"g":94},"2516":{"m":55,"g":94},"2515":{"m":55,"g":94},"2502":{"m":55,"g":94},"2500":{"m":55,"g":94},"2499":{"m":55,"g":94},"2438":{"m":55,"g":94},"2426":{"m":55,"g":94},"2457":{"m":55,"g":94},"2467":{"m":55,"g":94},"2495":{"m":55,"g":94},"2494":{"m":55,"g":94},"2493":{"m":55,"g":94},"2492":{"m":55,"g":94},"2491":{"m":55,"g":94},"2476":{"m":55,"g":94},"2490":{"m":55,"g":94},"2489":{"m":55,"g":94},"2486":{"m":55,"g":94},"2487":{"m":55,"g":94},"2481":{"m":55,"g":94},"2485":{"m":55,"g":94},"2484":{"m":55,"g":94},"2483":{"m":55,"g":94},"2479":{"m":55,"g":94},"2473":{"m":55,"g":94},"2469":{"m":55,"g":94},"2464":{"m":55,"g":94},"2466":{"m":55,"g":94},"2463":{"m":55,"g":94},"2462":{"m":55,"g":94},"2459":{"m":55,"g":94},"2456":{"m":55,"g":94},"2444":{"m":55,"g":94},"2455":{"m":55,"g":94},"2454":{"m":55,"g":94},"2442":{"m":55,"g":94},"2453":{"m":55,"g":94},"2452":{"m":55,"g":94},"2437":{"m":55,"g":94},"2449":{"m":55,"g":94},"2448":{"m":55,"g":94},"2425":{"m":55,"g":94},"2447":{"m":55,"g":94},"2436":{"m":55,"g":94},"2441":{"m":55,"g":94},"2440":{"m":55,"g":94},"2435":{"m":55,"g":94},"2434":{"m":55,"g":94},"2433":{"m":55,"g":94},"2424":{"m":55,"g":94},"2412":{"m":55,"g":94},"2422":{"m":55,"g":94},"2419":{"m":55,"g":94},"2417":{"m":55,"g":94},"2416":{"m":55,"g":94},"2410":{"m":55,"g":94},"2393":{"m":55,"g":94},"2413":{"m":55,"g":94},"2411":{"m":55,"g":94},"2398":{"m":55,"g":94},"2409":{"m":55,"g":94},"2408":{"m":55,"g":94},"2407":{"m":55,"g":94},"2406":{"m":55,"g":94},"2405":{"m":55,"g":94},"2404":{"m":55,"g":94},"2403":{"m":55,"g":94},"2401":{"m":55,"g":94},"2397":{"m":55,"g":94},"2394":{"m":55,"g":94},"2382":{"m":55,"g":94},"2330":{"m":55,"g":94},"2392":{"m":55,"g":94},"2391":{"m":55,"g":94},"2388":{"m":55,"g":94},"2390":{"m":55,"g":94},"2387":{"m":55,"g":94},"2380":{"m":55,"g":94},"2379":{"m":55,"g":94},"2378":{"m":55,"g":94},"2582":{"m":56,"g":94},"2581":{"m":56,"g":94},"2580":{"m":56,"g":94},"2579":{"m":56,"g":94},"2575":{"m":56,"g":94},"2566":{"m":56,"g":94},"2563":{"m":56,"g":94},"2553":{"m":56,"g":94},"2545":{"m":56,"g":94},"2547":{"m":56,"g":94},"2543":{"m":56,"g":94},"2509":{"m":56,"g":94},"2523":{"m":56,"g":94},"2529":{"m":56,"g":94},"2541":{"m":56,"g":94},"2616":{"m":57,"g":94},"2617":{"m":57,"g":94},"2615":{"m":57,"g":94},"2610":{"m":57,"g":94},"2611":{"m":57,"g":94},"2612":{"m":57,"g":94},"2608":{"m":57,"g":94},"2606":{"m":57,"g":94},"2586":{"m":57,"g":94},"2605":{"m":57,"g":94},"2603":{"m":57,"g":94},"2564":{"m":57,"g":94},"2570":{"m":57,"g":94},"2521":{"m":57,"g":94},"2598":{"m":57,"g":94},"2574":{"m":57,"g":94},"2565":{"m":57,"g":94},"2597":{"m":57,"g":94},"2596":{"m":57,"g":94},"2526":{"m":57,"g":94},"2557":{"m":57,"g":94},"2594":{"m":57,"g":94},"2555":{"m":57,"g":94},"2560":{"m":57,"g":94},"2592":{"m":57,"g":94},"2590":{"m":57,"g":94},"2589":{"m":57,"g":94},"2643":{"m":58,"g":94},"2641":{"m":58,"g":94},"2637":{"m":58,"g":94},"2635":{"m":58,"g":94},"2640":{"m":58,"g":94},"2639":{"m":58,"g":94},"2638":{"m":58,"g":94},"2609":{"m":58,"g":94},"2544":{"m":58,"g":94},"2631":{"m":58,"g":94},"2628":{"m":58,"g":94},"2633":{"m":58,"g":94},"2614":{"m":58,"g":94},"2626":{"m":58,"g":94},"2624":{"m":58,"g":94},"2625":{"m":58,"g":94},"2623":{"m":58,"g":94},"2622":{"m":58,"g":94},"2475":{"m":58,"g":94},"2618":{"m":58,"g":94},"2647":{"m":59,"g":94},"2648":{"m":59,"g":94},"2646":{"m":59,"g":94},"2636":{"m":59,"g":94},"2645":{"m":59,"g":94},"2644":{"m":59,"g":94},"2713":{"m":60,"g":94},"2688":{"m":60,"g":94},"2735":{"m":60,"g":94},"2733":{"m":60,"g":94},"2731":{"m":60,"g":94},"2571":{"m":60,"g":94},"2726":{"m":60,"g":94},"2727":{"m":60,"g":94},"2722":{"m":60,"g":94},"2717":{"m":60,"g":94},"2601":{"m":60,"g":94},"2716":{"m":60,"g":94},"2711":{"m":60,"g":94},"2714":{"m":60,"g":94},"2704":{"m":60,"g":94},"2707":{"m":60,"g":94},"2712":{"m":60,"g":94},"2150":{"m":60,"g":94},"2709":{"m":60,"g":94},"2695":{"m":60,"g":94},"2705":{"m":60,"g":94},"2663":{"m":60,"g":94},"2697":{"m":60,"g":94},"2692":{"m":60,"g":94},"2691":{"m":60,"g":94},"2690":{"m":60,"g":94},"2689":{"m":60,"g":94},"2685":{"m":60,"g":94},"2684":{"m":60,"g":94},"2683":{"m":60,"g":94},"2682":{"m":60,"g":94},"2680":{"m":60,"g":94},"2678":{"m":60,"g":94},"2679":{"m":60,"g":94},"2676":{"m":60,"g":94},"2674":{"m":60,"g":94},"2672":{"m":60,"g":94},"2670":{"m":60,"g":94},"2667":{"m":60,"g":94},"2669":{"m":60,"g":94},"2664":{"m":60,"g":94},"2642":{"m":60,"g":94},"2666":{"m":60,"g":94},"2654":{"m":60,"g":94},"2655":{"m":60,"g":94},"2656":{"m":60,"g":94},"2652":{"m":60,"g":94},"2651":{"m":60,"g":94},"2650":{"m":60,"g":94},"2649":{"m":60,"g":94},"2840":{"m":61,"g":94},"2837":{"m":61,"g":94},"2836":{"m":61,"g":94},"2804":{"m":61,"g":94},"2730":{"m":61,"g":94},"2826":{"m":61,"g":94},"2835":{"m":61,"g":94},"2822":{"m":61,"g":94},"2819":{"m":61,"g":94},"2833":{"m":61,"g":94},"2830":{"m":61,"g":94},"2787":{"m":61,"g":94},"2816":{"m":61,"g":94},"2813":{"m":61,"g":94},"2809":{"m":61,"g":94},"2792":{"m":61,"g":94},"2789":{"m":61,"g":94},"2773":{"m":61,"g":94},"2784":{"m":61,"g":94},"2723":{"m":61,"g":94},"2780":{"m":61,"g":94},"2779":{"m":61,"g":94},"2771":{"m":61,"g":94},"2774":{"m":61,"g":94},"2761":{"m":61,"g":94},"2770":{"m":61,"g":94},"2767":{"m":61,"g":94},"2758":{"m":61,"g":94},"2757":{"m":61,"g":94},"2513":{"m":61,"g":94},"2535":{"m":61,"g":94},"2756":{"m":61,"g":94},"2745":{"m":61,"g":94},"2748":{"m":61,"g":94},"2752":{"m":61,"g":94},"2751":{"m":61,"g":94},"2750":{"m":61,"g":94},"2899":{"m":62,"g":94},"2887":{"m":62,"g":94},"2888":{"m":62,"g":94},"2881":{"m":62,"g":94},"2885":{"m":62,"g":94},"2879":{"m":62,"g":94},"2878":{"m":62,"g":94},"2875":{"m":62,"g":94},"2630":{"m":62,"g":94},"2870":{"m":62,"g":94},"2869":{"m":62,"g":94},"2863":{"m":62,"g":94},"2868":{"m":62,"g":94},"2867":{"m":62,"g":94},"2862":{"m":62,"g":94},"2866":{"m":62,"g":94},"2865":{"m":62,"g":94},"2828":{"m":62,"g":94},"2859":{"m":62,"g":94},"2858":{"m":62,"g":94},"2861":{"m":62,"g":94},"2857":{"m":62,"g":94},"2860":{"m":62,"g":94},"2856":{"m":62,"g":94},"2851":{"m":62,"g":94},"2853":{"m":62,"g":94},"2854":{"m":62,"g":94},"2852":{"m":62,"g":94},"2786":{"m":62,"g":94},"2848":{"m":62,"g":94},"2850":{"m":62,"g":94},"2846":{"m":62,"g":94},"2843":{"m":62,"g":94},"2841":{"m":62,"g":94},"3009":{"m":63,"g":94},"2993":{"m":63,"g":94},"3010":{"m":63,"g":94},"3006":{"m":63,"g":94},"3008":{"m":63,"g":94},"3003":{"m":63,"g":94},"2998":{"m":63,"g":94},"3005":{"m":63,"g":94},"3004":{"m":63,"g":94},"3001":{"m":63,"g":94},"2996":{"m":63,"g":94},"2997":{"m":63,"g":94},"2995":{"m":63,"g":94},"2991":{"m":63,"g":94},"2992":{"m":63,"g":94},"2990":{"m":63,"g":94},"2396":{"m":63,"g":94},"2983":{"m":63,"g":94},"2988":{"m":63,"g":94},"2839":{"m":63,"g":94},"2982":{"m":63,"g":94},"2986":{"m":63,"g":94},"2987":{"m":63,"g":94},"2985":{"m":63,"g":94},"2984":{"m":63,"g":94},"2981":{"m":63,"g":94},"2975":{"m":63,"g":94},"2980":{"m":63,"g":94},"2979":{"m":63,"g":94},"2978":{"m":63,"g":94},"2976":{"m":63,"g":94},"2974":{"m":63,"g":94},"2973":{"m":63,"g":94},"2972":{"m":63,"g":94},"2958":{"m":63,"g":94},"2956":{"m":63,"g":94},"2901":{"m":63,"g":94},"2971":{"m":63,"g":94},"2785":{"m":63,"g":94},"2966":{"m":63,"g":94},"2967":{"m":63,"g":94},"2964":{"m":63,"g":94},"2963":{"m":63,"g":94},"2960":{"m":63,"g":94},"2941":{"m":63,"g":94},"2894":{"m":63,"g":94},"2942":{"m":63,"g":94},"2944":{"m":63,"g":94},"2954":{"m":63,"g":94},"2952":{"m":63,"g":94},"2951":{"m":63,"g":94},"2947":{"m":63,"g":94},"2948":{"m":63,"g":94},"2949":{"m":63,"g":94},"2950":{"m":63,"g":94},"2945":{"m":63,"g":94},"2907":{"m":63,"g":94},"2938":{"m":63,"g":94},"2937":{"m":63,"g":94},"2806":{"m":63,"g":94},"2876":{"m":63,"g":94},"2930":{"m":63,"g":94},"2928":{"m":63,"g":94},"2920":{"m":63,"g":94},"2926":{"m":63,"g":94},"2927":{"m":63,"g":94},"2925":{"m":63,"g":94},"2924":{"m":63,"g":94},"2923":{"m":63,"g":94},"2821":{"m":63,"g":94},"2910":{"m":63,"g":94},"2922":{"m":63,"g":94},"2919":{"m":63,"g":94},"2917":{"m":63,"g":94},"2915":{"m":63,"g":94},"2911":{"m":63,"g":94},"2909":{"m":63,"g":94},"2908":{"m":63,"g":94},"2511":{"m":63,"g":94},"2906":{"m":63,"g":94},"2904":{"m":63,"g":94},"2902":{"m":63,"g":94},"2872":{"m":63,"g":94},"2897":{"m":63,"g":94},"3180":{"m":64,"g":94},"3179":{"m":64,"g":94},"3178":{"m":64,"g":94},"3175":{"m":64,"g":94},"3176":{"m":64,"g":94},"3174":{"m":64,"g":94},"3146":{"m":64,"g":94},"3173":{"m":64,"g":94},"3170":{"m":64,"g":94},"3167":{"m":64,"g":94},"3156":{"m":64,"g":94},"3134":{"m":64,"g":94},"3162":{"m":64,"g":94},"3144":{"m":64,"g":94},"3155":{"m":64,"g":94},"2700":{"m":64,"g":94},"3154":{"m":64,"g":94},"3153":{"m":64,"g":94},"3152":{"m":64,"g":94},"3150":{"m":64,"g":94},"3151":{"m":64,"g":94},"3149":{"m":64,"g":94},"3147":{"m":64,"g":94},"3145":{"m":64,"g":94},"3085":{"m":64,"g":94},"3047":{"m":64,"g":94},"3143":{"m":64,"g":94},"3139":{"m":64,"g":94},"3135":{"m":64,"g":94},"3138":{"m":64,"g":94},"3113":{"m":64,"g":94},"3133":{"m":64,"g":94},"3132":{"m":64,"g":94},"3130":{"m":64,"g":94},"3129":{"m":64,"g":94},"3128":{"m":64,"g":94},"3127":{"m":64,"g":94},"3126":{"m":64,"g":94},"3125":{"m":64,"g":94},"3124":{"m":64,"g":94},"3121":{"m":64,"g":94},"3110":{"m":64,"g":94},"3109":{"m":64,"g":94},"3107":{"m":64,"g":94},"3096":{"m":64,"g":94},"3037":{"m":64,"g":94},"3105":{"m":64,"g":94},"3097":{"m":64,"g":94},"3095":{"m":64,"g":94},"3094":{"m":64,"g":94},"3070":{"m":64,"g":94},"3093":{"m":64,"g":94},"2742":{"m":64,"g":94},"3087":{"m":64,"g":94},"3086":{"m":64,"g":94},"3084":{"m":64,"g":94},"3083":{"m":64,"g":94},"3081":{"m":64,"g":94},"3080":{"m":64,"g":94},"3079":{"m":64,"g":94},"3078":{"m":64,"g":94},"3074":{"m":64,"g":94},"3030":{"m":64,"g":94},"3071":{"m":64,"g":94},"3069":{"m":64,"g":94},"3068":{"m":64,"g":94},"3067":{"m":64,"g":94},"3061":{"m":64,"g":94},"3062":{"m":64,"g":94},"3063":{"m":64,"g":94},"2989":{"m":64,"g":94},"3060":{"m":64,"g":94},"3045":{"m":64,"g":94},"3038":{"m":64,"g":94},"3058":{"m":64,"g":94},"3057":{"m":64,"g":94},"3055":{"m":64,"g":94},"3056":{"m":64,"g":94},"3054":{"m":64,"g":94},"3053":{"m":64,"g":94},"3052":{"m":64,"g":94},"3051":{"m":64,"g":94},"3048":{"m":64,"g":94},"3046":{"m":64,"g":94},"3039":{"m":64,"g":94},"3036":{"m":64,"g":94},"3035":{"m":64,"g":94},"3033":{"m":64,"g":94},"3027":{"m":64,"g":94},"3026":{"m":64,"g":94},"3025":{"m":64,"g":94},"3014":{"m":64,"g":94},"3018":{"m":64,"g":94},"2939":{"m":64,"g":94},"3022":{"m":64,"g":94},"3021":{"m":64,"g":94},"3017":{"m":64,"g":94},"3015":{"m":64,"g":94},"3020":{"m":64,"g":94},"3019":{"m":64,"g":94},"3016":{"m":64,"g":94},"3013":{"m":64,"g":94},"3012":{"m":64,"g":94},"3233":{"m":65,"g":94},"3231":{"m":65,"g":94},"3232":{"m":65,"g":94},"3224":{"m":65,"g":94},"3230":{"m":65,"g":94},"3229":{"m":65,"g":94},"3227":{"m":65,"g":94},"3218":{"m":65,"g":94},"3217":{"m":65,"g":94},"3216":{"m":65,"g":94},"3214":{"m":65,"g":94},"3213":{"m":65,"g":94},"3212":{"m":65,"g":94},"2977":{"m":65,"g":94},"3190":{"m":65,"g":94},"3169":{"m":65,"g":94},"3192":{"m":65,"g":94},"3183":{"m":65,"g":94},"3166":{"m":65,"g":94},"3171":{"m":65,"g":94},"3181":{"m":65,"g":94},"3313":{"m":66,"g":94},"3312":{"m":66,"g":94},"3306":{"m":66,"g":94},"3305":{"m":66,"g":94},"3299":{"m":66,"g":94},"3294":{"m":66,"g":94},"3293":{"m":66,"g":94},"3292":{"m":66,"g":94},"3287":{"m":66,"g":94},"3288":{"m":66,"g":94},"3161":{"m":66,"g":94},"3273":{"m":66,"g":94},"3276":{"m":66,"g":94},"3274":{"m":66,"g":94},"3272":{"m":66,"g":94},"3270":{"m":66,"g":94},"3269":{"m":66,"g":94},"3268":{"m":66,"g":94},"3205":{"m":66,"g":94},"3207":{"m":66,"g":94},"3259":{"m":66,"g":94},"3261":{"m":66,"g":94},"3114":{"m":66,"g":94},"3255":{"m":66,"g":94},"3252":{"m":66,"g":94},"3251":{"m":66,"g":94},"3250":{"m":66,"g":94},"3249":{"m":66,"g":94},"3248":{"m":66,"g":94},"3246":{"m":66,"g":94},"3221":{"m":66,"g":94},"3242":{"m":66,"g":94},"3240":{"m":66,"g":94},"3238":{"m":66,"g":94},"3236":{"m":66,"g":94},"3235":{"m":66,"g":94},"3228":{"m":66,"g":94},"3369":{"m":67,"g":94},"3378":{"m":67,"g":94},"3376":{"m":67,"g":94},"3374":{"m":67,"g":94},"3373":{"m":67,"g":94},"3372":{"m":67,"g":94},"3366":{"m":67,"g":94},"3356":{"m":67,"g":94},"3300":{"m":67,"g":94},"3355":{"m":67,"g":94},"3314":{"m":67,"g":94},"3350":{"m":67,"g":94},"3347":{"m":67,"g":94},"3352":{"m":67,"g":94},"3349":{"m":67,"g":94},"3338":{"m":67,"g":94},"3337":{"m":67,"g":94},"3335":{"m":67,"g":94},"3332":{"m":67,"g":94},"3325":{"m":67,"g":94},"3327":{"m":67,"g":94},"3324":{"m":67,"g":94},"3168":{"m":67,"g":94},"3317":{"m":67,"g":94},"3309":{"m":67,"g":94},"3459":{"m":68,"g":94},"3457":{"m":68,"g":94},"3413":{"m":68,"g":94},"3453":{"m":68,"g":94},"3452":{"m":68,"g":94},"3442":{"m":68,"g":94},"3441":{"m":68,"g":94},"3440":{"m":68,"g":94},"3439":{"m":68,"g":94},"3437":{"m":68,"g":94},"3410":{"m":68,"g":94},"3435":{"m":68,"g":94},"3433":{"m":68,"g":94},"3431":{"m":68,"g":94},"3430":{"m":68,"g":94},"3425":{"m":68,"g":94},"3422":{"m":68,"g":94},"3421":{"m":68,"g":94},"3408":{"m":68,"g":94},"3415":{"m":68,"g":94},"3411":{"m":68,"g":94},"3412":{"m":68,"g":94},"3407":{"m":68,"g":94},"3404":{"m":68,"g":94},"3346":{"m":68,"g":94},"3382":{"m":68,"g":94},"3386":{"m":68,"g":94},"3275":{"m":68,"g":94},"3556":{"m":69,"g":94},"3555":{"m":69,"g":94},"3550":{"m":69,"g":94},"3553":{"m":69,"g":94},"3543":{"m":69,"g":94},"3534":{"m":69,"g":94},"3541":{"m":69,"g":94},"3536":{"m":69,"g":94},"3529":{"m":69,"g":94},"3530":{"m":69,"g":94},"3267":{"m":69,"g":94},"3450":{"m":69,"g":94},"3503":{"m":69,"g":94},"3493":{"m":69,"g":94},"3523":{"m":69,"g":94},"3522":{"m":69,"g":94},"3405":{"m":69,"g":94},"3502":{"m":69,"g":94},"3420":{"m":69,"g":94},"3495":{"m":69,"g":94},"3496":{"m":69,"g":94},"3499":{"m":69,"g":94},"3498":{"m":69,"g":94},"3497":{"m":69,"g":94},"3500":{"m":69,"g":94},"3492":{"m":69,"g":94},"3490":{"m":69,"g":94},"3418":{"m":69,"g":94},"3364":{"m":69,"g":94},"3473":{"m":69,"g":94},"3469":{"m":69,"g":94},"3468":{"m":69,"g":94},"3466":{"m":69,"g":94},"3638":{"m":70,"g":94},"3636":{"m":70,"g":94},"3634":{"m":70,"g":94},"3632":{"m":70,"g":94},"3619":{"m":70,"g":94},"3617":{"m":70,"g":94},"3260":{"m":70,"g":94},"3532":{"m":70,"g":94},"3597":{"m":70,"g":94},"3535":{"m":70,"g":94},"3258":{"m":70,"g":94},"3598":{"m":70,"g":94},"3564":{"m":70,"g":94},"3594":{"m":70,"g":94},"3592":{"m":70,"g":94},"3591":{"m":70,"g":94},"3589":{"m":70,"g":94},"3587":{"m":70,"g":94},"3582":{"m":70,"g":94},"3548":{"m":70,"g":94},"3584":{"m":70,"g":94},"3505":{"m":70,"g":94},"3581":{"m":70,"g":94},"3563":{"m":70,"g":94},"3363":{"m":70,"g":94},"3558":{"m":70,"g":94},"3557":{"m":70,"g":94},"3645":{"m":71,"g":94},"3644":{"m":71,"g":94},"3643":{"m":71,"g":94},"3639":{"m":71,"g":94},"4114":{"m":72,"g":94},"4099":{"m":72,"g":94},"4101":{"m":72,"g":94},"3211":{"m":72,"g":94},"3029":{"m":72,"g":94},"4110":{"m":72,"g":94},"4105":{"m":72,"g":94},"4109":{"m":72,"g":94},"4103":{"m":72,"g":94},"4108":{"m":72,"g":94},"4107":{"m":72,"g":94},"4102":{"m":72,"g":94},"4100":{"m":72,"g":94},"4012":{"m":72,"g":94},"4081":{"m":72,"g":94},"3986":{"m":72,"g":94},"4075":{"m":72,"g":94},"3990":{"m":72,"g":94},"3790":{"m":72,"g":94},"3941":{"m":72,"g":94},"4077":{"m":72,"g":94},"4074":{"m":72,"g":94},"4066":{"m":72,"g":94},"4071":{"m":72,"g":94},"4065":{"m":72,"g":94},"4016":{"m":72,"g":94},"3607":{"m":72,"g":94},"3712":{"m":72,"g":94},"3948":{"m":72,"g":94},"3954":{"m":72,"g":94},"4030":{"m":72,"g":94},"4051":{"m":72,"g":94},"4046":{"m":72,"g":94},"4053":{"m":72,"g":94},"4052":{"m":72,"g":94},"4023":{"m":72,"g":94},"4000":{"m":72,"g":94},"4049":{"m":72,"g":94},"4044":{"m":72,"g":94},"4043":{"m":72,"g":94},"4033":{"m":72,"g":94},"4039":{"m":72,"g":94},"3999":{"m":72,"g":94},"4034":{"m":72,"g":94},"4032":{"m":72,"g":94},"4027":{"m":72,"g":94},"4025":{"m":72,"g":94},"3264":{"m":72,"g":94},"4031":{"m":72,"g":94},"4029":{"m":72,"g":94},"4021":{"m":72,"g":94},"3988":{"m":72,"g":94},"4014":{"m":72,"g":94},"3826":{"m":72,"g":94},"4010":{"m":72,"g":94},"4008":{"m":72,"g":94},"3987":{"m":72,"g":94},"3822":{"m":72,"g":94},"3993":{"m":72,"g":94},"3406":{"m":72,"g":94},"3994":{"m":72,"g":94},"3992":{"m":72,"g":94},"3991":{"m":72,"g":94},"3893":{"m":72,"g":94},"3989":{"m":72,"g":94},"3985":{"m":72,"g":94},"3982":{"m":72,"g":94},"3979":{"m":72,"g":94},"3976":{"m":72,"g":94},"3977":{"m":72,"g":94},"3975":{"m":72,"g":94},"3967":{"m":72,"g":94},"3966":{"m":72,"g":94},"3963":{"m":72,"g":94},"3852":{"m":72,"g":94},"3566":{"m":72,"g":94},"3678":{"m":72,"g":94},"3950":{"m":72,"g":94},"3870":{"m":72,"g":94},"3613":{"m":72,"g":94},"3866":{"m":72,"g":94},"3861":{"m":72,"g":94},"3934":{"m":72,"g":94},"3933":{"m":72,"g":94},"3593":{"m":72,"g":94},"3922":{"m":72,"g":94},"3925":{"m":72,"g":94},"3914":{"m":72,"g":94},"3905":{"m":72,"g":94},"3897":{"m":72,"g":94},"3907":{"m":72,"g":94},"3898":{"m":72,"g":94},"3903":{"m":72,"g":94},"3791":{"m":72,"g":94},"3900":{"m":72,"g":94},"3894":{"m":72,"g":94},"3298":{"m":72,"g":94},"3860":{"m":72,"g":94},"3845":{"m":72,"g":94},"3602":{"m":72,"g":94},"3865":{"m":72,"g":94},"3843":{"m":72,"g":94},"3519":{"m":72,"g":94},"3841":{"m":72,"g":94},"3857":{"m":72,"g":94},"3709":{"m":72,"g":94},"3803":{"m":72,"g":94},"3787":{"m":72,"g":94},"3237":{"m":72,"g":94},"3641":{"m":72,"g":94},"3741":{"m":72,"g":94},"3799":{"m":72,"g":94},"3801":{"m":72,"g":94},"3821":{"m":72,"g":94},"3829":{"m":72,"g":94},"3828":{"m":72,"g":94},"3818":{"m":72,"g":94},"3730":{"m":72,"g":94},"3785":{"m":72,"g":94},"3813":{"m":72,"g":94},"3809":{"m":72,"g":94},"2693":{"m":72,"g":94},"3795":{"m":72,"g":94},"3116":{"m":72,"g":94},"3115":{"m":72,"g":94},"3562":{"m":72,"g":94},"3117":{"m":72,"g":94},"3766":{"m":72,"g":94},"3777":{"m":72,"g":94},"3348":{"m":72,"g":94},"3223":{"m":72,"g":94},"3772":{"m":72,"g":94},"3773":{"m":72,"g":94},"3771":{"m":72,"g":94},"3733":{"m":72,"g":94},"3754":{"m":72,"g":94},"3432":{"m":72,"g":94},"3761":{"m":72,"g":94},"3740":{"m":72,"g":94},"3680":{"m":72,"g":94},"3747":{"m":72,"g":94},"3737":{"m":72,"g":94},"3588":{"m":72,"g":94},"3652":{"m":72,"g":94},"3732":{"m":72,"g":94},"3705":{"m":72,"g":94},"3731":{"m":72,"g":94},"3722":{"m":72,"g":94},"3727":{"m":72,"g":94},"3710":{"m":72,"g":94},"3677":{"m":72,"g":94},"3601":{"m":72,"g":94},"3692":{"m":72,"g":94},"3700":{"m":72,"g":94},"3706":{"m":72,"g":94},"3628":{"m":72,"g":94},"3698":{"m":72,"g":94},"3657":{"m":72,"g":94},"3665":{"m":72,"g":94},"3635":{"m":72,"g":94},"3676":{"m":72,"g":94},"3663":{"m":72,"g":94},"3629":{"m":72,"g":94},"3654":{"m":72,"g":94},"3624":{"m":72,"g":94},"3567":{"m":72,"g":94},"3616":{"m":72,"g":94},"3650":{"m":72,"g":94},"4140":{"m":73,"g":94},"4142":{"m":73,"g":94},"4147":{"m":73,"g":94},"4134":{"m":73,"g":94},"4135":{"m":73,"g":94},"4138":{"m":73,"g":94},"4137":{"m":73,"g":94},"4132":{"m":73,"g":94},"4128":{"m":73,"g":94},"4126":{"m":73,"g":94},"4129":{"m":73,"g":94},"4131":{"m":73,"g":94},"4038":{"m":73,"g":94},"4111":{"m":73,"g":94},"4113":{"m":73,"g":94},"4117":{"m":73,"g":94},"4121":{"m":73,"g":94},"4041":{"m":74,"g":94},"4381":{"m":74,"g":94},"4377":{"m":74,"g":94},"4376":{"m":74,"g":94},"4374":{"m":74,"g":94},"4367":{"m":74,"g":94},"4086":{"m":74,"g":94},"4356":{"m":74,"g":94},"3679":{"m":74,"g":94},"3814":{"m":74,"g":94},"3835":{"m":74,"g":94},"3844":{"m":74,"g":94},"3896":{"m":74,"g":94},"3959":{"m":74,"g":94},"3980":{"m":74,"g":94},"3962":{"m":74,"g":94},"3961":{"m":74,"g":94},"4026":{"m":74,"g":94},"4079":{"m":74,"g":94},"4359":{"m":74,"g":94},"4212":{"m":74,"g":94},"4295":{"m":74,"g":94},"4342":{"m":74,"g":94},"4335":{"m":74,"g":94},"4329":{"m":74,"g":94},"4362":{"m":74,"g":94},"4348":{"m":74,"g":94},"4326":{"m":74,"g":94},"4354":{"m":74,"g":94},"4355":{"m":74,"g":94},"4352":{"m":74,"g":94},"4350":{"m":74,"g":94},"4278":{"m":74,"g":94},"4082":{"m":74,"g":94},"3203":{"m":74,"g":94},"4340":{"m":74,"g":94},"4337":{"m":74,"g":94},"3911":{"m":74,"g":94},"4334":{"m":74,"g":94},"4331":{"m":74,"g":94},"4333":{"m":74,"g":94},"4104":{"m":74,"g":94},"4215":{"m":74,"g":94},"4327":{"m":74,"g":94},"4297":{"m":74,"g":94},"4323":{"m":74,"g":94},"4321":{"m":74,"g":94},"4317":{"m":74,"g":94},"4229":{"m":74,"g":94},"4220":{"m":74,"g":94},"4311":{"m":74,"g":94},"4299":{"m":74,"g":94},"4287":{"m":74,"g":94},"4290":{"m":74,"g":94},"4199":{"m":74,"g":94},"4291":{"m":74,"g":94},"4136":{"m":74,"g":94},"4288":{"m":74,"g":94},"4284":{"m":74,"g":94},"4277":{"m":74,"g":94},"4279":{"m":74,"g":94},"4275":{"m":74,"g":94},"4272":{"m":74,"g":94},"4261":{"m":74,"g":94},"4256":{"m":74,"g":94},"4267":{"m":74,"g":94},"4262":{"m":74,"g":94},"4258":{"m":74,"g":94},"4255":{"m":74,"g":94},"4231":{"m":74,"g":94},"4144":{"m":74,"g":94},"4252":{"m":74,"g":94},"4206":{"m":74,"g":94},"3958":{"m":74,"g":94},"4165":{"m":74,"g":94},"4250":{"m":74,"g":94},"4230":{"m":74,"g":94},"4238":{"m":74,"g":94},"4243":{"m":74,"g":94},"4242":{"m":74,"g":94},"4241":{"m":74,"g":94},"4237":{"m":74,"g":94},"4235":{"m":74,"g":94},"4228":{"m":74,"g":94},"4217":{"m":74,"g":94},"3148":{"m":74,"g":94},"4218":{"m":74,"g":94},"4224":{"m":74,"g":94},"4193":{"m":74,"g":94},"3631":{"m":74,"g":94},"4225":{"m":74,"g":94},"4223":{"m":74,"g":94},"4213":{"m":74,"g":94},"4222":{"m":74,"g":94},"4219":{"m":74,"g":94},"4124":{"m":74,"g":94},"4216":{"m":74,"g":94},"4211":{"m":74,"g":94},"4210":{"m":74,"g":94},"3749":{"m":74,"g":94},"4203":{"m":74,"g":94},"4181":{"m":74,"g":94},"4200":{"m":74,"g":94},"4198":{"m":74,"g":94},"4197":{"m":74,"g":94},"4164":{"m":74,"g":94},"4195":{"m":74,"g":94},"4194":{"m":74,"g":94},"4189":{"m":74,"g":94},"4185":{"m":74,"g":94},"4187":{"m":74,"g":94},"4186":{"m":74,"g":94},"4178":{"m":74,"g":94},"4174":{"m":74,"g":94},"4179":{"m":74,"g":94},"4177":{"m":74,"g":94},"4176":{"m":74,"g":94},"4170":{"m":74,"g":94},"4168":{"m":74,"g":94},"4166":{"m":74,"g":94},"4162":{"m":74,"g":94},"4163":{"m":74,"g":94},"3888":{"m":74,"g":94},"4089":{"m":74,"g":94},"3786":{"m":74,"g":94},"3694":{"m":74,"g":94},"4152":{"m":74,"g":94},"4154":{"m":74,"g":94},"4151":{"m":74,"g":94},"4148":{"m":74,"g":94},"4402":{"m":75,"g":94},"4397":{"m":75,"g":94},"4399":{"m":75,"g":94},"4269":{"m":75,"g":94},"4320":{"m":75,"g":94},"4398":{"m":75,"g":94},"4393":{"m":75,"g":94},"4390":{"m":75,"g":94},"4392":{"m":75,"g":94},"4375":{"m":75,"g":94},"4669":{"m":76,"g":94},"4743":{"m":76,"g":94},"4797":{"m":76,"g":94},"4782":{"m":76,"g":94},"4777":{"m":76,"g":94},"4775":{"m":76,"g":94},"4784":{"m":76,"g":94},"4566":{"m":76,"g":94},"4728":{"m":76,"g":94},"4755":{"m":76,"g":94},"4753":{"m":76,"g":94},"4752":{"m":76,"g":94},"4751":{"m":76,"g":94},"4310":{"m":76,"g":94},"4705":{"m":76,"g":94},"4435":{"m":76,"g":94},"4695":{"m":76,"g":94},"4691":{"m":76,"g":94},"4738":{"m":76,"g":94},"4735":{"m":76,"g":94},"4744":{"m":76,"g":94},"3023":{"m":76,"g":94},"4737":{"m":76,"g":94},"3899":{"m":76,"g":94},"4609":{"m":76,"g":94},"4736":{"m":76,"g":94},"4716":{"m":76,"g":94},"4721":{"m":76,"g":94},"4731":{"m":76,"g":94},"4720":{"m":76,"g":94},"4396":{"m":76,"g":94},"4605":{"m":76,"g":94},"4680":{"m":76,"g":94},"4631":{"m":76,"g":94},"4698":{"m":76,"g":94},"4064":{"m":76,"g":94},"4525":{"m":76,"g":94},"4661":{"m":76,"g":94},"4610":{"m":76,"g":94},"4608":{"m":76,"g":94},"4685":{"m":76,"g":94},"4679":{"m":76,"g":94},"4643":{"m":76,"g":94},"3984":{"m":76,"g":94},"4660":{"m":76,"g":94},"4676":{"m":76,"g":94},"4670":{"m":76,"g":94},"4677":{"m":76,"g":94},"4674":{"m":76,"g":94},"4556":{"m":76,"g":94},"4665":{"m":76,"g":94},"4596":{"m":76,"g":94},"4639":{"m":76,"g":94},"4664":{"m":76,"g":94},"4582":{"m":76,"g":94},"4654":{"m":76,"g":94},"4637":{"m":76,"g":94},"4641":{"m":76,"g":94},"4613":{"m":76,"g":94},"4640":{"m":76,"g":94},"4558":{"m":76,"g":94},"4622":{"m":76,"g":94},"4592":{"m":76,"g":94},"4571":{"m":76,"g":94},"3446":{"m":76,"g":94},"4577":{"m":76,"g":94},"4549":{"m":76,"g":94},"4583":{"m":76,"g":94},"4514":{"m":76,"g":94},"4232":{"m":76,"g":94},"4515":{"m":76,"g":94},"4557":{"m":76,"g":94},"4553":{"m":76,"g":94},"4274":{"m":76,"g":94},"4521":{"m":76,"g":94},"4247":{"m":76,"g":94},"4532":{"m":76,"g":94},"4441":{"m":76,"g":94},"4538":{"m":76,"g":94},"4541":{"m":76,"g":94},"4542":{"m":76,"g":94},"4500":{"m":76,"g":94},"4531":{"m":76,"g":94},"3682":{"m":76,"g":94},"4458":{"m":76,"g":94},"4486":{"m":76,"g":94},"4507":{"m":76,"g":94},"4522":{"m":76,"g":94},"4520":{"m":76,"g":94},"4505":{"m":76,"g":94},"4482":{"m":76,"g":94},"4517":{"m":76,"g":94},"4513":{"m":76,"g":94},"4510":{"m":76,"g":94},"4499":{"m":76,"g":94},"4495":{"m":76,"g":94},"4446":{"m":76,"g":94},"4480":{"m":76,"g":94},"4067":{"m":76,"g":94},"4485":{"m":76,"g":94},"4418":{"m":76,"g":94},"4493":{"m":76,"g":94},"2798":{"m":76,"g":94},"4372":{"m":76,"g":94},"4483":{"m":76,"g":94},"4386":{"m":76,"g":94},"2797":{"m":76,"g":94},"3612":{"m":76,"g":94},"4448":{"m":76,"g":94},"4465":{"m":76,"g":94},"4474":{"m":76,"g":94},"4479":{"m":76,"g":94},"4424":{"m":76,"g":94},"4202":{"m":76,"g":94},"4484":{"m":76,"g":94},"4481":{"m":76,"g":94},"4477":{"m":76,"g":94},"4472":{"m":76,"g":94},"4363":{"m":76,"g":94},"4470":{"m":76,"g":94},"4469":{"m":76,"g":94},"4383":{"m":76,"g":94},"4468":{"m":76,"g":94},"4467":{"m":76,"g":94},"4466":{"m":76,"g":94},"4449":{"m":76,"g":94},"4464":{"m":76,"g":94},"4460":{"m":76,"g":94},"4368":{"m":76,"g":94},"4459":{"m":76,"g":94},"4447":{"m":76,"g":94},"4423":{"m":76,"g":94},"4391":{"m":76,"g":94},"4413":{"m":76,"g":94},"4454":{"m":76,"g":94},"4453":{"m":76,"g":94},"4455":{"m":76,"g":94},"4452":{"m":76,"g":94},"4451":{"m":76,"g":94},"4442":{"m":76,"g":94},"4439":{"m":76,"g":94},"4438":{"m":76,"g":94},"4437":{"m":76,"g":94},"4302":{"m":76,"g":94},"4427":{"m":76,"g":94},"4419":{"m":76,"g":94},"3964":{"m":76,"g":94},"4400":{"m":76,"g":94},"4009":{"m":76,"g":94},"4403":{"m":76,"g":94},"4878":{"m":77,"g":94},"4874":{"m":77,"g":94},"4873":{"m":77,"g":94},"4831":{"m":77,"g":94},"4872":{"m":77,"g":94},"4768":{"m":77,"g":94},"4871":{"m":77,"g":94},"4866":{"m":77,"g":94},"4749":{"m":77,"g":94},"4864":{"m":77,"g":94},"4772":{"m":77,"g":94},"4834":{"m":77,"g":94},"4492":{"m":77,"g":94},"4863":{"m":77,"g":94},"4855":{"m":77,"g":94},"4853":{"m":77,"g":94},"4840":{"m":77,"g":94},"4740":{"m":77,"g":94},"4750":{"m":77,"g":94},"4687":{"m":77,"g":94},"4729":{"m":77,"g":94},"4712":{"m":77,"g":94},"4704":{"m":77,"g":94},"4688":{"m":77,"g":94},"4681":{"m":77,"g":94},"4648":{"m":77,"g":94},"4832":{"m":77,"g":94},"4528":{"m":77,"g":94},"4597":{"m":77,"g":94},"4487":{"m":77,"g":94},"3949":{"m":77,"g":94},"4844":{"m":77,"g":94},"4846":{"m":77,"g":94},"4843":{"m":77,"g":94},"4799":{"m":77,"g":94},"4788":{"m":77,"g":94},"4809":{"m":77,"g":94},"4837":{"m":77,"g":94},"3969":{"m":77,"g":94},"4835":{"m":77,"g":94},"4815":{"m":77,"g":94},"4819":{"m":77,"g":94},"4770":{"m":77,"g":94},"4833":{"m":77,"g":94},"4830":{"m":77,"g":94},"4694":{"m":77,"g":94},"4826":{"m":77,"g":94},"4825":{"m":77,"g":94},"4828":{"m":77,"g":94},"4827":{"m":77,"g":94},"4823":{"m":77,"g":94},"4341":{"m":77,"g":94},"4638":{"m":77,"g":94},"4813":{"m":77,"g":94},"4764":{"m":77,"g":94},"4706":{"m":77,"g":94},"4745":{"m":77,"g":94},"4565":{"m":77,"g":94},"4804":{"m":77,"g":94},"4506":{"m":77,"g":94},"4388":{"m":77,"g":94},"4628":{"m":77,"g":94},"4719":{"m":77,"g":94},"5091":{"m":78,"g":94},"5080":{"m":78,"g":94},"5089":{"m":78,"g":94},"5079":{"m":78,"g":94},"5088":{"m":78,"g":94},"5052":{"m":78,"g":94},"5050":{"m":78,"g":94},"5074":{"m":78,"g":94},"4535":{"m":78,"g":94},"5072":{"m":78,"g":94},"4918":{"m":78,"g":94},"4996":{"m":78,"g":94},"4995":{"m":78,"g":94},"4994":{"m":78,"g":94},"5060":{"m":78,"g":94},"5057":{"m":78,"g":94},"5056":{"m":78,"g":94},"4625":{"m":78,"g":94},"5049":{"m":78,"g":94},"5051":{"m":78,"g":94},"5039":{"m":78,"g":94},"4796":{"m":78,"g":94},"5046":{"m":78,"g":94},"5048":{"m":78,"g":94},"5036":{"m":78,"g":94},"4992":{"m":78,"g":94},"5005":{"m":78,"g":94},"5024":{"m":78,"g":94},"5030":{"m":78,"g":94},"5020":{"m":78,"g":94},"4727":{"m":78,"g":94},"5009":{"m":78,"g":94},"5011":{"m":78,"g":94},"4817":{"m":78,"g":94},"5008":{"m":78,"g":94},"4951":{"m":78,"g":94},"4861":{"m":78,"g":94},"4989":{"m":78,"g":94},"4581":{"m":78,"g":94},"4915":{"m":78,"g":94},"4977":{"m":78,"g":94},"4958":{"m":78,"g":94},"4767":{"m":78,"g":94},"4959":{"m":78,"g":94},"4954":{"m":78,"g":94},"4953":{"m":78,"g":94},"4950":{"m":78,"g":94},"4754":{"m":78,"g":94},"4944":{"m":78,"g":94},"4928":{"m":78,"g":94},"4913":{"m":78,"g":94},"4936":{"m":78,"g":94},"4883":{"m":78,"g":94},"4925":{"m":78,"g":94},"4933":{"m":78,"g":94},"4932":{"m":78,"g":94},"4931":{"m":78,"g":94},"4902":{"m":78,"g":94},"4930":{"m":78,"g":94},"4927":{"m":78,"g":94},"4926":{"m":78,"g":94},"4896":{"m":78,"g":94},"4914":{"m":78,"g":94},"4908":{"m":78,"g":94},"4909":{"m":78,"g":94},"4890":{"m":78,"g":94},"4899":{"m":78,"g":94},"4898":{"m":78,"g":94},"4530":{"m":78,"g":94},"4891":{"m":78,"g":94},"4889":{"m":78,"g":94},"4886":{"m":78,"g":94},"4795":{"m":78,"g":94},"4845":{"m":78,"g":94},"4882":{"m":78,"g":94},"5117":{"m":79,"g":94},"5106":{"m":79,"g":94},"5092":{"m":79,"g":94},"5097":{"m":79,"g":94},"5445":{"m":80,"g":94},"5113":{"m":80,"g":94},"5425":{"m":80,"g":94},"5397":{"m":80,"g":94},"5398":{"m":80,"g":94},"5038":{"m":80,"g":94},"5211":{"m":80,"g":94},"5264":{"m":80,"g":94},"5436":{"m":80,"g":94},"5344":{"m":80,"g":94},"5431":{"m":80,"g":94},"5434":{"m":80,"g":94},"5430":{"m":80,"g":94},"5419":{"m":80,"g":94},"5420":{"m":80,"g":94},"5423":{"m":80,"g":94},"5422":{"m":80,"g":94},"5415":{"m":80,"g":94},"5416":{"m":80,"g":94},"5351":{"m":80,"g":94},"5412":{"m":80,"g":94},"5352":{"m":80,"g":94},"5214":{"m":80,"g":94},"5406":{"m":80,"g":94},"5401":{"m":80,"g":94},"5400":{"m":80,"g":94},"5381":{"m":80,"g":94},"5399":{"m":80,"g":94},"5395":{"m":80,"g":94},"5279":{"m":80,"g":94},"5291":{"m":80,"g":94},"5368":{"m":80,"g":94},"5393":{"m":80,"g":94},"5263":{"m":80,"g":94},"5392":{"m":80,"g":94},"5371":{"m":80,"g":94},"5385":{"m":80,"g":94},"5370":{"m":80,"g":94},"5384":{"m":80,"g":94},"5326":{"m":80,"g":94},"5003":{"m":80,"g":94},"5367":{"m":80,"g":94},"5364":{"m":80,"g":94},"5277":{"m":80,"g":94},"5360":{"m":80,"g":94},"5359":{"m":80,"g":94},"5161":{"m":80,"g":94},"5328":{"m":80,"g":94},"5357":{"m":80,"g":94},"5342":{"m":80,"g":94},"5343":{"m":80,"g":94},"5322":{"m":80,"g":94},"5341":{"m":80,"g":94},"5337":{"m":80,"g":94},"5336":{"m":80,"g":94},"5333":{"m":80,"g":94},"5332":{"m":80,"g":94},"5331":{"m":80,"g":94},"5327":{"m":80,"g":94},"4848":{"m":80,"g":94},"5294":{"m":80,"g":94},"5321":{"m":80,"g":94},"4884":{"m":80,"g":94},"5210":{"m":80,"g":94},"5317":{"m":80,"g":94},"5316":{"m":80,"g":94},"5315":{"m":80,"g":94},"5120":{"m":80,"g":94},"5299":{"m":80,"g":94},"5311":{"m":80,"g":94},"5142":{"m":80,"g":94},"5271":{"m":80,"g":94},"5065":{"m":80,"g":94},"5310":{"m":80,"g":94},"5308":{"m":80,"g":94},"5307":{"m":80,"g":94},"5306":{"m":80,"g":94},"5304":{"m":80,"g":94},"5303":{"m":80,"g":94},"5302":{"m":80,"g":94},"5298":{"m":80,"g":94},"5301":{"m":80,"g":94},"5300":{"m":80,"g":94},"5290":{"m":80,"g":94},"5292":{"m":80,"g":94},"5289":{"m":80,"g":94},"5288":{"m":80,"g":94},"5287":{"m":80,"g":94},"5286":{"m":80,"g":94},"5280":{"m":80,"g":94},"5193":{"m":80,"g":94},"5244":{"m":80,"g":94},"5265":{"m":80,"g":94},"5254":{"m":80,"g":94},"5262":{"m":80,"g":94},"5127":{"m":80,"g":94},"5228":{"m":80,"g":94},"5259":{"m":80,"g":94},"5245":{"m":80,"g":94},"5216":{"m":80,"g":94},"5167":{"m":80,"g":94},"5213":{"m":80,"g":94},"5215":{"m":80,"g":94},"5204":{"m":80,"g":94},"4880":{"m":80,"g":94},"5086":{"m":80,"g":94},"5190":{"m":80,"g":94},"4444":{"m":80,"g":94},"5209":{"m":80,"g":94},"5196":{"m":80,"g":94},"5207":{"m":80,"g":94},"5171":{"m":80,"g":94},"5144":{"m":80,"g":94},"5102":{"m":80,"g":94},"5194":{"m":80,"g":94},"5128":{"m":80,"g":94},"5185":{"m":80,"g":94},"5189":{"m":80,"g":94},"4058":{"m":80,"g":94},"5179":{"m":80,"g":94},"5180":{"m":80,"g":94},"5176":{"m":80,"g":94},"5110":{"m":80,"g":94},"5068":{"m":80,"g":94},"5173":{"m":80,"g":94},"5175":{"m":80,"g":94},"5174":{"m":80,"g":94},"3972":{"m":80,"g":94},"5159":{"m":80,"g":94},"4938":{"m":80,"g":94},"5139":{"m":80,"g":94},"4911":{"m":80,"g":94},"5150":{"m":80,"g":94},"5155":{"m":80,"g":94},"4686":{"m":80,"g":94},"5158":{"m":80,"g":94},"5151":{"m":80,"g":94},"5115":{"m":80,"g":94},"5152":{"m":80,"g":94},"4760":{"m":80,"g":94},"5103":{"m":80,"g":94},"5147":{"m":80,"g":94},"5083":{"m":80,"g":94},"5140":{"m":80,"g":94},"5145":{"m":80,"g":94},"4984":{"m":80,"g":94},"5137":{"m":80,"g":94},"5133":{"m":80,"g":94},"5090":{"m":80,"g":94},"5126":{"m":80,"g":94},"5582":{"m":81,"g":94},"5581":{"m":81,"g":94},"4947":{"m":81,"g":94},"5571":{"m":81,"g":94},"5564":{"m":81,"g":94},"5568":{"m":81,"g":94},"5432":{"m":81,"g":94},"5149":{"m":81,"g":94},"5562":{"m":81,"g":94},"5543":{"m":81,"g":94},"5561":{"m":81,"g":94},"5546":{"m":81,"g":94},"5549":{"m":81,"g":94},"5504":{"m":81,"g":94},"5475":{"m":81,"g":94},"5460":{"m":81,"g":94},"5547":{"m":81,"g":94},"5534":{"m":81,"g":94},"5545":{"m":81,"g":94},"5548":{"m":81,"g":94},"5340":{"m":81,"g":94},"5540":{"m":81,"g":94},"5544":{"m":81,"g":94},"5542":{"m":81,"g":94},"4693":{"m":81,"g":94},"5476":{"m":81,"g":94},"5497":{"m":81,"g":94},"5461":{"m":81,"g":94},"5473":{"m":81,"g":94},"5518":{"m":81,"g":94},"5440":{"m":81,"g":94},"4836":{"m":81,"g":94},"5512":{"m":81,"g":94},"5373":{"m":81,"g":94},"5426":{"m":81,"g":94},"5511":{"m":81,"g":94},"5500":{"m":81,"g":94},"5503":{"m":81,"g":94},"5496":{"m":81,"g":94},"5493":{"m":81,"g":94},"5205":{"m":81,"g":94},"5479":{"m":81,"g":94},"4887":{"m":81,"g":94},"5480":{"m":81,"g":94},"5481":{"m":81,"g":94},"5484":{"m":81,"g":94},"5489":{"m":81,"g":94},"4982":{"m":81,"g":94},"5345":{"m":81,"g":94},"5467":{"m":81,"g":94},"5463":{"m":81,"g":94},"5447":{"m":81,"g":94},"5449":{"m":81,"g":94},"5444":{"m":81,"g":94},"5611":{"m":82,"g":94},"5510":{"m":82,"g":94},"5610":{"m":82,"g":94},"5580":{"m":82,"g":94},"5604":{"m":82,"g":94},"5609":{"m":82,"g":94},"5608":{"m":82,"g":94},"5477":{"m":82,"g":94},"5589":{"m":82,"g":94},"5598":{"m":82,"g":94},"5488":{"m":82,"g":94},"5037":{"m":82,"g":94},"5021":{"m":82,"g":94},"4980":{"m":82,"g":94},"4937":{"m":82,"g":94},"4718":{"m":82,"g":94},"4590":{"m":82,"g":94},"3443":{"m":82,"g":94},"5348":{"m":82,"g":94},"5570":{"m":82,"g":94},"5318":{"m":82,"g":94},"5590":{"m":82,"g":94},"5319":{"m":82,"g":94},"5241":{"m":82,"g":94},"5188":{"m":82,"g":94},"5141":{"m":82,"g":94},"5019":{"m":82,"g":94},"5016":{"m":82,"g":94},"4859":{"m":82,"g":94},"4852":{"m":82,"g":94},"4675":{"m":82,"g":94},"4733":{"m":82,"g":94},"5378":{"m":82,"g":94},"5224":{"m":82,"g":94},"4226":{"m":82,"g":94},"5433":{"m":82,"g":94},"5588":{"m":82,"g":94},"5452":{"m":82,"g":94},"5586":{"m":82,"g":94},"5575":{"m":82,"g":94},"5417":{"m":82,"g":94},"5531":{"m":82,"g":94},"5526":{"m":82,"g":94},"5521":{"m":82,"g":94},"5559":{"m":82,"g":94},"5560":{"m":82,"g":94},"5567":{"m":82,"g":94},"5574":{"m":82,"g":94},"5795":{"m":83,"g":94},"5691":{"m":83,"g":94},"5790":{"m":83,"g":94},"5791":{"m":83,"g":94},"5769":{"m":83,"g":94},"5789":{"m":83,"g":94},"5787":{"m":83,"g":94},"5786":{"m":83,"g":94},"5785":{"m":83,"g":94},"5779":{"m":83,"g":94},"5777":{"m":83,"g":94},"5774":{"m":83,"g":94},"5776":{"m":83,"g":94},"5772":{"m":83,"g":94},"3744":{"m":83,"g":94},"4986":{"m":83,"g":94},"4971":{"m":83,"g":94},"4870":{"m":83,"g":94},"5633":{"m":83,"g":94},"5599":{"m":83,"g":94},"5565":{"m":83,"g":94},"5509":{"m":83,"g":94},"5687":{"m":83,"g":94},"5592":{"m":83,"g":94},"5607":{"m":83,"g":94},"5730":{"m":83,"g":94},"5697":{"m":83,"g":94},"5748":{"m":83,"g":94},"5682":{"m":83,"g":94},"5685":{"m":83,"g":94},"5716":{"m":83,"g":94},"5720":{"m":83,"g":94},"5722":{"m":83,"g":94},"5728":{"m":83,"g":94},"5733":{"m":83,"g":94},"5736":{"m":83,"g":94},"5756":{"m":83,"g":94},"5760":{"m":83,"g":94},"5552":{"m":83,"g":94},"5718":{"m":83,"g":94},"5719":{"m":83,"g":94},"5737":{"m":83,"g":94},"5740":{"m":83,"g":94},"5753":{"m":83,"g":94},"5754":{"m":83,"g":94},"5750":{"m":83,"g":94},"5723":{"m":83,"g":94},"5704":{"m":83,"g":94},"5738":{"m":83,"g":94},"5715":{"m":83,"g":94},"5078":{"m":83,"g":94},"4491":{"m":83,"g":94},"5706":{"m":83,"g":94},"5648":{"m":83,"g":94},"5684":{"m":83,"g":94},"5707":{"m":83,"g":94},"5349":{"m":83,"g":94},"5688":{"m":83,"g":94},"5686":{"m":83,"g":94},"5683":{"m":83,"g":94},"5667":{"m":83,"g":94},"5677":{"m":83,"g":94},"5530":{"m":83,"g":94},"5671":{"m":83,"g":94},"5670":{"m":83,"g":94},"5435":{"m":83,"g":94},"5669":{"m":83,"g":94},"5666":{"m":83,"g":94},"5281":{"m":83,"g":94},"5601":{"m":83,"g":94},"5619":{"m":83,"g":94},"5628":{"m":83,"g":94},"5649":{"m":83,"g":94},"5646":{"m":83,"g":94},"5638":{"m":83,"g":94},"5634":{"m":83,"g":94},"5272":{"m":83,"g":94},"5641":{"m":83,"g":94},"5632":{"m":83,"g":94},"5640":{"m":83,"g":94},"5624":{"m":83,"g":94},"5622":{"m":83,"g":94},"5620":{"m":83,"g":94},"5618":{"m":83,"g":94},"5615":{"m":83,"g":94},"5578":{"m":83,"g":94},"5845":{"m":84,"g":94},"5849":{"m":84,"g":94},"5854":{"m":84,"g":94},"5823":{"m":84,"g":94},"5847":{"m":84,"g":94},"5816":{"m":84,"g":94},"5798":{"m":84,"g":94},"5850":{"m":84,"g":94},"5851":{"m":84,"g":94},"5846":{"m":84,"g":94},"5726":{"m":84,"g":94},"5839":{"m":84,"g":94},"5842":{"m":84,"g":94},"5833":{"m":84,"g":94},"5838":{"m":84,"g":94},"5551":{"m":84,"g":94},"5825":{"m":84,"g":94},"5276":{"m":84,"g":94},"5482":{"m":84,"g":94},"5771":{"m":84,"g":94},"5809":{"m":84,"g":94},"5807":{"m":84,"g":94},"5690":{"m":84,"g":94},"5390":{"m":84,"g":94},"5788":{"m":84,"g":94},"5643":{"m":84,"g":94},"5796":{"m":84,"g":94},"5797":{"m":84,"g":94},"5939":{"m":85,"g":94},"5934":{"m":85,"g":94},"5881":{"m":85,"g":94},"5915":{"m":85,"g":94},"5930":{"m":85,"g":94},"5724":{"m":85,"g":94},"5783":{"m":85,"g":94},"5933":{"m":85,"g":94},"5932":{"m":85,"g":94},"5912":{"m":85,"g":94},"5909":{"m":85,"g":94},"5917":{"m":85,"g":94},"5919":{"m":85,"g":94},"5910":{"m":85,"g":94},"5905":{"m":85,"g":94},"5383":{"m":85,"g":94},"5903":{"m":85,"g":94},"5900":{"m":85,"g":94},"5899":{"m":85,"g":94},"5870":{"m":85,"g":94},"5830":{"m":85,"g":94},"5901":{"m":85,"g":94},"5898":{"m":85,"g":94},"5861":{"m":85,"g":94},"5696":{"m":85,"g":94},"5725":{"m":85,"g":94},"5893":{"m":85,"g":94},"5793":{"m":85,"g":94},"5896":{"m":85,"g":94},"5746":{"m":85,"g":94},"5895":{"m":85,"g":94},"5841":{"m":85,"g":94},"5836":{"m":85,"g":94},"5894":{"m":85,"g":94},"5880":{"m":85,"g":94},"5875":{"m":85,"g":94},"5859":{"m":85,"g":94},"4115":{"m":85,"g":94},"4949":{"m":85,"g":94},"5820":{"m":85,"g":94},"5868":{"m":85,"g":94},"5860":{"m":85,"g":94},"5801":{"m":85,"g":94},"5857":{"m":85,"g":94},"6165":{"m":86,"g":94},"5778":{"m":86,"g":94},"6162":{"m":86,"g":94},"6089":{"m":86,"g":94},"6141":{"m":86,"g":94},"5822":{"m":86,"g":94},"6101":{"m":86,"g":94},"5745":{"m":86,"g":94},"6132":{"m":86,"g":94},"6131":{"m":86,"g":94},"6112":{"m":86,"g":94},"6129":{"m":86,"g":94},"6123":{"m":86,"g":94},"6097":{"m":86,"g":94},"5662":{"m":86,"g":94},"5764":{"m":86,"g":94},"6119":{"m":86,"g":94},"5626":{"m":86,"g":94},"6091":{"m":86,"g":94},"5572":{"m":86,"g":94},"5232":{"m":86,"g":94},"5219":{"m":86,"g":94},"5121":{"m":86,"g":94},"6077":{"m":86,"g":94},"6111":{"m":86,"g":94},"6038":{"m":86,"g":94},"6079":{"m":86,"g":94},"6034":{"m":86,"g":94},"6102":{"m":86,"g":94},"6105":{"m":86,"g":94},"5993":{"m":86,"g":94},"6075":{"m":86,"g":94},"6063":{"m":86,"g":94},"5233":{"m":86,"g":94},"5014":{"m":86,"g":94},"6084":{"m":86,"g":94},"3853":{"m":86,"g":94},"6010":{"m":86,"g":94},"6039":{"m":86,"g":94},"5655":{"m":86,"g":94},"6062":{"m":86,"g":94},"6004":{"m":86,"g":94},"6045":{"m":86,"g":94},"5885":{"m":86,"g":94},"5751":{"m":86,"g":94},"6057":{"m":86,"g":94},"6048":{"m":86,"g":94},"6047":{"m":86,"g":94},"6046":{"m":86,"g":94},"5081":{"m":86,"g":94},"5752":{"m":86,"g":94},"5996":{"m":86,"g":94},"5428":{"m":86,"g":94},"5555":{"m":86,"g":94},"5587":{"m":86,"g":94},"5781":{"m":86,"g":94},"6018":{"m":86,"g":94},"6002":{"m":86,"g":94},"5997":{"m":86,"g":94},"5679":{"m":86,"g":94},"5957":{"m":86,"g":94},"6012":{"m":86,"g":94},"5992":{"m":86,"g":94},"5998":{"m":86,"g":94},"5991":{"m":86,"g":94},"5986":{"m":86,"g":94},"5977":{"m":86,"g":94},"5975":{"m":86,"g":94},"5681":{"m":86,"g":94},"5969":{"m":86,"g":94},"5968":{"m":86,"g":94},"5350":{"m":86,"g":94},"5967":{"m":86,"g":94},"5908":{"m":86,"g":94},"5960":{"m":86,"g":94},"5782":{"m":86,"g":94},"5956":{"m":86,"g":94},"5945":{"m":86,"g":94},"5944":{"m":86,"g":94},"5952":{"m":86,"g":94},"5953":{"m":86,"g":94},"5834":{"m":86,"g":94},"5921":{"m":86,"g":94},"6245":{"m":87,"g":94},"6259":{"m":87,"g":94},"6252":{"m":87,"g":94},"5084":{"m":87,"g":94},"5657":{"m":87,"g":94},"6247":{"m":87,"g":94},"6235":{"m":87,"g":94},"6251":{"m":87,"g":94},"6042":{"m":87,"g":94},"6248":{"m":87,"g":94},"6225":{"m":87,"g":94},"5922":{"m":87,"g":94},"6241":{"m":87,"g":94},"6243":{"m":87,"g":94},"6244":{"m":87,"g":94},"6223":{"m":87,"g":94},"6206":{"m":87,"g":94},"6209":{"m":87,"g":94},"6231":{"m":87,"g":94},"6212":{"m":87,"g":94},"6201":{"m":87,"g":94},"6213":{"m":87,"g":94},"5558":{"m":87,"g":94},"6154":{"m":87,"g":94},"6043":{"m":87,"g":94},"6204":{"m":87,"g":94},"6202":{"m":87,"g":94},"6178":{"m":87,"g":94},"6198":{"m":87,"g":94},"6192":{"m":87,"g":94},"6188":{"m":87,"g":94},"6032":{"m":87,"g":94},"6199":{"m":87,"g":94},"5621":{"m":87,"g":94},"6196":{"m":87,"g":94},"6195":{"m":87,"g":94},"6073":{"m":87,"g":94},"6169":{"m":87,"g":94},"6191":{"m":87,"g":94},"6190":{"m":87,"g":94},"6186":{"m":87,"g":94},"5654":{"m":87,"g":94},"6180":{"m":87,"g":94},"6179":{"m":87,"g":94},"6184":{"m":87,"g":94},"6183":{"m":87,"g":94},"4701":{"m":87,"g":94},"6181":{"m":87,"g":94},"6146":{"m":87,"g":94},"6114":{"m":87,"g":94},"6118":{"m":87,"g":94},"6016":{"m":87,"g":94},"6566":{"m":88,"g":94},"6567":{"m":88,"g":94},"6485":{"m":88,"g":94},"6560":{"m":88,"g":94},"6550":{"m":88,"g":94},"6533":{"m":88,"g":94},"6562":{"m":88,"g":94},"6524":{"m":88,"g":94},"6347":{"m":88,"g":94},"6558":{"m":88,"g":94},"6521":{"m":88,"g":94},"6507":{"m":88,"g":94},"6474":{"m":88,"g":94},"6452":{"m":88,"g":94},"6404":{"m":88,"g":94},"6493":{"m":88,"g":94},"6355":{"m":88,"g":94},"6535":{"m":88,"g":94},"6536":{"m":88,"g":94},"6059":{"m":88,"g":94},"6532":{"m":88,"g":94},"6120":{"m":88,"g":94},"6522":{"m":88,"g":94},"6520":{"m":88,"g":94},"6469":{"m":88,"g":94},"6482":{"m":88,"g":94},"6308":{"m":88,"g":94},"6388":{"m":88,"g":94},"6492":{"m":88,"g":94},"6504":{"m":88,"g":94},"6019":{"m":88,"g":94},"6457":{"m":88,"g":94},"6499":{"m":88,"g":94},"6510":{"m":88,"g":94},"6508":{"m":88,"g":94},"5759":{"m":88,"g":94},"6419":{"m":88,"g":94},"6275":{"m":88,"g":94},"6503":{"m":88,"g":94},"5573":{"m":88,"g":94},"6445":{"m":88,"g":94},"6461":{"m":88,"g":94},"6467":{"m":88,"g":94},"6468":{"m":88,"g":94},"6476":{"m":88,"g":94},"6311":{"m":88,"g":94},"6487":{"m":88,"g":94},"5339":{"m":88,"g":94},"6381":{"m":88,"g":94},"6214":{"m":88,"g":94},"6475":{"m":88,"g":94},"6472":{"m":88,"g":94},"6447":{"m":88,"g":94},"6385":{"m":88,"g":94},"6444":{"m":88,"g":94},"6405":{"m":88,"g":94},"6429":{"m":88,"g":94},"6412":{"m":88,"g":94},"6387":{"m":88,"g":94},"6386":{"m":88,"g":94},"6326":{"m":88,"g":94},"6438":{"m":88,"g":94},"6321":{"m":88,"g":94},"4957":{"m":88,"g":94},"6440":{"m":88,"g":94},"6431":{"m":88,"g":94},"6098":{"m":88,"g":94},"6414":{"m":88,"g":94},"6430":{"m":88,"g":94},"6417":{"m":88,"g":94},"6137":{"m":88,"g":94},"6401":{"m":88,"g":94},"6400":{"m":88,"g":94},"5974":{"m":88,"g":94},"6325":{"m":88,"g":94},"6323":{"m":88,"g":94},"6333":{"m":88,"g":94},"6397":{"m":88,"g":94},"6396":{"m":88,"g":94},"6395":{"m":88,"g":94},"6339":{"m":88,"g":94},"6383":{"m":88,"g":94},"6392":{"m":88,"g":94},"6331":{"m":88,"g":94},"6391":{"m":88,"g":94},"6365":{"m":88,"g":94},"6250":{"m":88,"g":94},"6187":{"m":88,"g":94},"6379":{"m":88,"g":94},"6377":{"m":88,"g":94},"6362":{"m":88,"g":94},"6364":{"m":88,"g":94},"6330":{"m":88,"g":94},"6290":{"m":88,"g":94},"6041":{"m":88,"g":94},"6284":{"m":88,"g":94},"6257":{"m":88,"g":94},"6108":{"m":88,"g":94},"6134":{"m":88,"g":94},"6107":{"m":88,"g":94},"4741":{"m":88,"g":94},"6348":{"m":88,"g":94},"6211":{"m":88,"g":94},"6175":{"m":88,"g":94},"6366":{"m":88,"g":94},"6373":{"m":88,"g":94},"6356":{"m":88,"g":94},"6368":{"m":88,"g":94},"6324":{"m":88,"g":94},"6316":{"m":88,"g":94},"5099":{"m":88,"g":94},"6361":{"m":88,"g":94},"6360":{"m":88,"g":94},"6358":{"m":88,"g":94},"6359":{"m":88,"g":94},"6121":{"m":88,"g":94},"5694":{"m":88,"g":94},"6334":{"m":88,"g":94},"6136":{"m":88,"g":94},"6336":{"m":88,"g":94},"6327":{"m":88,"g":94},"6302":{"m":88,"g":94},"6147":{"m":88,"g":94},"6317":{"m":88,"g":94},"6216":{"m":88,"g":94},"6298":{"m":88,"g":94},"6109":{"m":88,"g":94},"5914":{"m":88,"g":94},"6009":{"m":88,"g":94},"6274":{"m":88,"g":94},"6283":{"m":88,"g":94},"6300":{"m":88,"g":94},"6138":{"m":88,"g":94},"6282":{"m":88,"g":94},"6276":{"m":88,"g":94},"6273":{"m":88,"g":94},"6115":{"m":88,"g":94},"7038":{"m":89,"g":94},"6833":{"m":89,"g":94},"7029":{"m":89,"g":94},"7027":{"m":89,"g":94},"6980":{"m":89,"g":94},"6964":{"m":89,"g":94},"7023":{"m":89,"g":94},"7018":{"m":89,"g":94},"7017":{"m":89,"g":94},"6741":{"m":89,"g":94},"6987":{"m":89,"g":94},"7015":{"m":89,"g":94},"7013":{"m":89,"g":94},"6884":{"m":89,"g":94},"6992":{"m":89,"g":94},"6998":{"m":89,"g":94},"7008":{"m":89,"g":94},"7007":{"m":89,"g":94},"6958":{"m":89,"g":94},"6557":{"m":89,"g":94},"6990":{"m":89,"g":94},"6960":{"m":89,"g":94},"6973":{"m":89,"g":94},"6983":{"m":89,"g":94},"6929":{"m":89,"g":94},"6977":{"m":89,"g":94},"6967":{"m":89,"g":94},"6981":{"m":89,"g":94},"6979":{"m":89,"g":94},"6976":{"m":89,"g":94},"6970":{"m":89,"g":94},"6937":{"m":89,"g":94},"6965":{"m":89,"g":94},"6974":{"m":89,"g":94},"6966":{"m":89,"g":94},"6956":{"m":89,"g":94},"6963":{"m":89,"g":94},"6926":{"m":89,"g":94},"6968":{"m":89,"g":94},"6853":{"m":89,"g":94},"6957":{"m":89,"g":94},"6955":{"m":89,"g":94},"6916":{"m":89,"g":94},"6885":{"m":89,"g":94},"6950":{"m":89,"g":94},"6953":{"m":89,"g":94},"6220":{"m":89,"g":94},"6866":{"m":89,"g":94},"6895":{"m":89,"g":94},"6874":{"m":89,"g":94},"6915":{"m":89,"g":94},"6924":{"m":89,"g":94},"6945":{"m":89,"g":94},"6369":{"m":89,"g":94},"5955":{"m":89,"g":94},"6912":{"m":89,"g":94},"6910":{"m":89,"g":94},"6944":{"m":89,"g":94},"6943":{"m":89,"g":94},"6942":{"m":89,"g":94},"6939":{"m":89,"g":94},"6767":{"m":89,"g":94},"6879":{"m":89,"g":94},"6922":{"m":89,"g":94},"6932":{"m":89,"g":94},"6934":{"m":89,"g":94},"6931":{"m":89,"g":94},"6930":{"m":89,"g":94},"6838":{"m":89,"g":94},"6458":{"m":89,"g":94},"6877":{"m":89,"g":94},"6887":{"m":89,"g":94},"6890":{"m":89,"g":94},"6764":{"m":89,"g":94},"6170":{"m":89,"g":94},"6837":{"m":89,"g":94},"6868":{"m":89,"g":94},"6865":{"m":89,"g":94},"6851":{"m":89,"g":94},"6846":{"m":89,"g":94},"6820":{"m":89,"g":94},"6277":{"m":89,"g":94},"6861":{"m":89,"g":94},"6878":{"m":89,"g":94},"6736":{"m":89,"g":94},"6852":{"m":89,"g":94},"6460":{"m":89,"g":94},"6659":{"m":89,"g":94},"6858":{"m":89,"g":94},"6745":{"m":89,"g":94},"5929":{"m":89,"g":94},"6735":{"m":89,"g":94},"6816":{"m":89,"g":94},"6818":{"m":89,"g":94},"6671":{"m":89,"g":94},"6456":{"m":89,"g":94},"6766":{"m":89,"g":94},"6093":{"m":89,"g":94},"6812":{"m":89,"g":94},"6811":{"m":89,"g":94},"6815":{"m":89,"g":94},"6813":{"m":89,"g":94},"6780":{"m":89,"g":94},"5382":{"m":89,"g":94},"6805":{"m":89,"g":94},"6699":{"m":89,"g":94},"6803":{"m":89,"g":94},"6804":{"m":89,"g":94},"6799":{"m":89,"g":94},"6800":{"m":89,"g":94},"5981":{"m":89,"g":94},"6421":{"m":89,"g":94},"6797":{"m":89,"g":94},"6795":{"m":89,"g":94},"6787":{"m":89,"g":94},"6794":{"m":89,"g":94},"6788":{"m":89,"g":94},"6791":{"m":89,"g":94},"6792":{"m":89,"g":94},"6734":{"m":89,"g":94},"6408":{"m":89,"g":94},"6786":{"m":89,"g":94},"6785":{"m":89,"g":94},"6782":{"m":89,"g":94},"6784":{"m":89,"g":94},"6679":{"m":89,"g":94},"6772":{"m":89,"g":94},"6289":{"m":89,"g":94},"6509":{"m":89,"g":94},"6737":{"m":89,"g":94},"6761":{"m":89,"g":94},"6765":{"m":89,"g":94},"6727":{"m":89,"g":94},"6265":{"m":89,"g":94},"6748":{"m":89,"g":94},"6746":{"m":89,"g":94},"6742":{"m":89,"g":94},"6728":{"m":89,"g":94},"6680":{"m":89,"g":94},"6437":{"m":89,"g":94},"6729":{"m":89,"g":94},"6545":{"m":89,"g":94},"6705":{"m":89,"g":94},"6715":{"m":89,"g":94},"6725":{"m":89,"g":94},"6718":{"m":89,"g":94},"6720":{"m":89,"g":94},"6726":{"m":89,"g":94},"6676":{"m":89,"g":94},"6709":{"m":89,"g":94},"6719":{"m":89,"g":94},"6479":{"m":89,"g":94},"6668":{"m":89,"g":94},"6682":{"m":89,"g":94},"6711":{"m":89,"g":94},"6710":{"m":89,"g":94},"6712":{"m":89,"g":94},"6706":{"m":89,"g":94},"6703":{"m":89,"g":94},"6697":{"m":89,"g":94},"6649":{"m":89,"g":94},"6685":{"m":89,"g":94},"6689":{"m":89,"g":94},"6693":{"m":89,"g":94},"6655":{"m":89,"g":94},"6260":{"m":89,"g":94},"6627":{"m":89,"g":94},"6687":{"m":89,"g":94},"6582":{"m":89,"g":94},"6672":{"m":89,"g":94},"6678":{"m":89,"g":94},"6380":{"m":89,"g":94},"6665":{"m":89,"g":94},"6673":{"m":89,"g":94},"6007":{"m":89,"g":94},"6473":{"m":89,"g":94},"6677":{"m":89,"g":94},"6661":{"m":89,"g":94},"6660":{"m":89,"g":94},"6638":{"m":89,"g":94},"6606":{"m":89,"g":94},"6662":{"m":89,"g":94},"6652":{"m":89,"g":94},"6658":{"m":89,"g":94},"6403":{"m":89,"g":94},"6601":{"m":89,"g":94},"6640":{"m":89,"g":94},"6450":{"m":89,"g":94},"6650":{"m":89,"g":94},"6585":{"m":89,"g":94},"6648":{"m":89,"g":94},"6634":{"m":89,"g":94},"6646":{"m":89,"g":94},"6643":{"m":89,"g":94},"6547":{"m":89,"g":94},"6631":{"m":89,"g":94},"6603":{"m":89,"g":94},"6263":{"m":89,"g":94},"6639":{"m":89,"g":94},"6635":{"m":89,"g":94},"6620":{"m":89,"g":94},"6629":{"m":89,"g":94},"6628":{"m":89,"g":94},"6617":{"m":89,"g":94},"6306":{"m":89,"g":94},"6599":{"m":89,"g":94},"6611":{"m":89,"g":94},"6598":{"m":89,"g":94},"6597":{"m":89,"g":94},"6581":{"m":89,"g":94},"6610":{"m":89,"g":94},"6575":{"m":89,"g":94},"6609":{"m":89,"g":94},"6594":{"m":89,"g":94},"6587":{"m":89,"g":94},"6596":{"m":89,"g":94},"6595":{"m":89,"g":94},"6593":{"m":89,"g":94},"6586":{"m":89,"g":94},"6571":{"m":89,"g":94},"6439":{"m":89,"g":94},"6527":{"m":89,"g":94},"6588":{"m":89,"g":94},"6570":{"m":89,"g":94},"6546":{"m":89,"g":94},"6537":{"m":89,"g":94},"6494":{"m":89,"g":94},"5961":{"m":89,"g":94},"6543":{"m":89,"g":94},"4068":{"m":89,"g":94},"6577":{"m":89,"g":94},"6578":{"m":89,"g":94},"6477":{"m":89,"g":94},"6564":{"m":89,"g":94},"6576":{"m":89,"g":94},"7248":{"m":90,"g":94},"7244":{"m":90,"g":94},"7247":{"m":90,"g":94},"6058":{"m":90,"g":94},"7234":{"m":90,"g":94},"7245":{"m":90,"g":94},"7231":{"m":90,"g":94},"7239":{"m":90,"g":94},"7232":{"m":90,"g":94},"7207":{"m":90,"g":94},"7213":{"m":90,"g":94},"7228":{"m":90,"g":94},"7221":{"m":90,"g":94},"7218":{"m":90,"g":94},"6378":{"m":90,"g":94},"7217":{"m":90,"g":94},"7215":{"m":90,"g":94},"7214":{"m":90,"g":94},"7210":{"m":90,"g":94},"7163":{"m":90,"g":94},"7205":{"m":90,"g":94},"7204":{"m":90,"g":94},"7202":{"m":90,"g":94},"7200":{"m":90,"g":94},"7196":{"m":90,"g":94},"7198":{"m":90,"g":94},"7195":{"m":90,"g":94},"7186":{"m":90,"g":94},"7180":{"m":90,"g":94},"7189":{"m":90,"g":94},"7191":{"m":90,"g":94},"7190":{"m":90,"g":94},"7184":{"m":90,"g":94},"7181":{"m":90,"g":94},"7177":{"m":90,"g":94},"7178":{"m":90,"g":94},"7175":{"m":90,"g":94},"7172":{"m":90,"g":94},"7173":{"m":90,"g":94},"7161":{"m":90,"g":94},"7150":{"m":90,"g":94},"7153":{"m":90,"g":94},"7170":{"m":90,"g":94},"7165":{"m":90,"g":94},"7156":{"m":90,"g":94},"7020":{"m":90,"g":94},"7157":{"m":90,"g":94},"6814":{"m":90,"g":94},"7154":{"m":90,"g":94},"7155":{"m":90,"g":94},"7152":{"m":90,"g":94},"7146":{"m":90,"g":94},"7145":{"m":90,"g":94},"7056":{"m":90,"g":94},"7058":{"m":90,"g":94},"7140":{"m":90,"g":94},"7134":{"m":90,"g":94},"6026":{"m":90,"g":94},"7092":{"m":90,"g":94},"7126":{"m":90,"g":94},"7119":{"m":90,"g":94},"7115":{"m":90,"g":94},"6919":{"m":90,"g":94},"7093":{"m":90,"g":94},"6994":{"m":90,"g":94},"6824":{"m":90,"g":94},"7079":{"m":90,"g":94},"6106":{"m":90,"g":94},"7091":{"m":90,"g":94},"7097":{"m":90,"g":94},"6870":{"m":90,"g":94},"7067":{"m":90,"g":94},"7076":{"m":90,"g":94},"7046":{"m":90,"g":94},"7054":{"m":90,"g":94},"7073":{"m":90,"g":94},"6031":{"m":90,"g":94},"7071":{"m":90,"g":94},"6579":{"m":90,"g":94},"7066":{"m":90,"g":94},"6716":{"m":90,"g":94},"7064":{"m":90,"g":94},"7049":{"m":90,"g":94},"7063":{"m":90,"g":94},"6947":{"m":90,"g":94},"7057":{"m":90,"g":94},"7061":{"m":90,"g":94},"7060":{"m":90,"g":94},"7021":{"m":90,"g":94},"7053":{"m":90,"g":94},"7037":{"m":90,"g":94},"7051":{"m":90,"g":94},"6999":{"m":90,"g":94},"7045":{"m":90,"g":94},"7043":{"m":90,"g":94},"7040":{"m":90,"g":94},"7493":{"m":91,"g":94},"7490":{"m":91,"g":94},"7376":{"m":91,"g":94},"7487":{"m":91,"g":94},"7347":{"m":91,"g":94},"7269":{"m":91,"g":94},"7449":{"m":91,"g":94},"7469":{"m":91,"g":94},"7378":{"m":91,"g":94},"7481":{"m":91,"g":94},"7382":{"m":91,"g":94},"7480":{"m":91,"g":94},"7456":{"m":91,"g":94},"7479":{"m":91,"g":94},"7397":{"m":91,"g":94},"7472":{"m":91,"g":94},"7457":{"m":91,"g":94},"7290":{"m":91,"g":94},"6821":{"m":91,"g":94},"7454":{"m":91,"g":94},"7451":{"m":91,"g":94},"7445":{"m":91,"g":94},"7361":{"m":91,"g":94},"7391":{"m":91,"g":94},"7441":{"m":91,"g":94},"7327":{"m":91,"g":94},"7406":{"m":91,"g":94},"7414":{"m":91,"g":94},"7408":{"m":91,"g":94},"7412":{"m":91,"g":94},"7351":{"m":91,"g":94},"7420":{"m":91,"g":94},"7409":{"m":91,"g":94},"7425":{"m":91,"g":94},"6984":{"m":91,"g":94},"7394":{"m":91,"g":94},"7396":{"m":91,"g":94},"7400":{"m":91,"g":94},"7329":{"m":91,"g":94},"7401":{"m":91,"g":94},"7403":{"m":91,"g":94},"7402":{"m":91,"g":94},"7285":{"m":91,"g":94},"7219":{"m":91,"g":94},"7360":{"m":91,"g":94},"7399":{"m":91,"g":94},"7398":{"m":91,"g":94},"6389":{"m":91,"g":94},"7372":{"m":91,"g":94},"5485":{"m":91,"g":94},"7393":{"m":91,"g":94},"7356":{"m":91,"g":94},"7326":{"m":91,"g":94},"7322":{"m":91,"g":94},"7371":{"m":91,"g":94},"7364":{"m":91,"g":94},"7159":{"m":91,"g":94},"7370":{"m":91,"g":94},"7343":{"m":91,"g":94},"7366":{"m":91,"g":94},"7362":{"m":91,"g":94},"7242":{"m":91,"g":94},"7363":{"m":91,"g":94},"7303":{"m":91,"g":94},"7354":{"m":91,"g":94},"7099":{"m":91,"g":94},"7333":{"m":91,"g":94},"7331":{"m":91,"g":94},"7284":{"m":91,"g":94},"7096":{"m":91,"g":94},"7319":{"m":91,"g":94},"7003":{"m":91,"g":94},"7301":{"m":91,"g":94},"7297":{"m":91,"g":94},"7251":{"m":91,"g":94},"7300":{"m":91,"g":94},"6614":{"m":91,"g":94},"7267":{"m":91,"g":94},"7264":{"m":91,"g":94},"7237":{"m":91,"g":94},"7289":{"m":91,"g":94},"7288":{"m":91,"g":94},"7286":{"m":91,"g":94},"6842":{"m":91,"g":94},"7283":{"m":91,"g":94},"7164":{"m":91,"g":94},"7022":{"m":91,"g":94},"6081":{"m":91,"g":94},"7265":{"m":91,"g":94},"7160":{"m":91,"g":94},"7179":{"m":91,"g":94},"7167":{"m":91,"g":94},"7252":{"m":91,"g":94},"7122":{"m":91,"g":94},"7233":{"m":91,"g":94},"7125":{"m":91,"g":94},"7559":{"m":92,"g":94},"7541":{"m":92,"g":94},"7542":{"m":92,"g":94},"7544":{"m":92,"g":94},"7549":{"m":92,"g":94},"7543":{"m":92,"g":94},"7527":{"m":92,"g":94},"7531":{"m":92,"g":94},"7148":{"m":92,"g":94},"7499":{"m":92,"g":94},"7507":{"m":92,"g":94},"7513":{"m":92,"g":94},"7521":{"m":92,"g":94},"7520":{"m":92,"g":94},"6793":{"m":92,"g":94},"6721":{"m":92,"g":94},"7386":{"m":92,"g":94},"7522":{"m":92,"g":94},"6641":{"m":92,"g":94},"6626":{"m":92,"g":94},"7498":{"m":92,"g":94},"7512":{"m":92,"g":94},"7510":{"m":92,"g":94},"7516":{"m":92,"g":94},"7508":{"m":92,"g":94},"7437":{"m":92,"g":94},"7489":{"m":92,"g":94},"7505":{"m":92,"g":94},"6717":{"m":92,"g":94},"7277":{"m":92,"g":94},"7422":{"m":92,"g":94},"7439":{"m":92,"g":94},"7236":{"m":92,"g":94},"7423":{"m":92,"g":94},"7268":{"m":92,"g":94},"7802":{"m":93,"g":94},"7801":{"m":93,"g":94},"7799":{"m":93,"g":94},"7792":{"m":93,"g":94},"7800":{"m":93,"g":94},"7756":{"m":93,"g":94},"7790":{"m":93,"g":94},"7757":{"m":93,"g":94},"7786":{"m":93,"g":94},"7222":{"m":93,"g":94},"7787":{"m":93,"g":94},"7784":{"m":93,"g":94},"7444":{"m":93,"g":94},"7623":{"m":93,"g":94},"7782":{"m":93,"g":94},"7596":{"m":93,"g":94},"7772":{"m":93,"g":94},"7705":{"m":93,"g":94},"7745":{"m":93,"g":94},"7778":{"m":93,"g":94},"7419":{"m":93,"g":94},"7418":{"m":93,"g":94},"7748":{"m":93,"g":94},"7764":{"m":93,"g":94},"7741":{"m":93,"g":94},"7729":{"m":93,"g":94},"7390":{"m":93,"g":94},"7759":{"m":93,"g":94},"7751":{"m":93,"g":94},"7754":{"m":93,"g":94},"7755":{"m":93,"g":94},"7744":{"m":93,"g":94},"7752":{"m":93,"g":94},"7723":{"m":93,"g":94},"7673":{"m":93,"g":94},"7740":{"m":93,"g":94},"7681":{"m":93,"g":94},"6771":{"m":93,"g":94},"7647":{"m":93,"g":94},"7750":{"m":93,"g":94},"7722":{"m":93,"g":94},"7738":{"m":93,"g":94},"7278":{"m":93,"g":94},"7731":{"m":93,"g":94},"7735":{"m":93,"g":94},"6770":{"m":93,"g":94},"7734":{"m":93,"g":94},"7714":{"m":93,"g":94},"6698":{"m":93,"g":94},"7462":{"m":93,"g":94},"6549":{"m":93,"g":94},"7621":{"m":93,"g":94},"6512":{"m":93,"g":94},"7292":{"m":93,"g":94},"7416":{"m":93,"g":94},"7717":{"m":93,"g":94},"7677":{"m":93,"g":94},"7683":{"m":93,"g":94},"7697":{"m":93,"g":94},"7635":{"m":93,"g":94},"7642":{"m":93,"g":94},"7698":{"m":93,"g":94},"7684":{"m":93,"g":94},"7688":{"m":93,"g":94},"7676":{"m":93,"g":94},"7629":{"m":93,"g":94},"6985":{"m":93,"g":94},"7675":{"m":93,"g":94},"7486":{"m":93,"g":94},"7671":{"m":93,"g":94},"7648":{"m":93,"g":94},"7318":{"m":93,"g":94},"7663":{"m":93,"g":94},"7627":{"m":93,"g":94},"7632":{"m":93,"g":94},"7524":{"m":93,"g":94},"7643":{"m":93,"g":94},"7640":{"m":93,"g":94},"7539":{"m":93,"g":94},"7176":{"m":93,"g":94},"7580":{"m":93,"g":94},"7628":{"m":93,"g":94},"7619":{"m":93,"g":94},"7432":{"m":93,"g":94},"7636":{"m":93,"g":94},"7630":{"m":93,"g":94},"7624":{"m":93,"g":94},"7625":{"m":93,"g":94},"7036":{"m":93,"g":94},"7620":{"m":93,"g":94},"7310":{"m":93,"g":94},"7309":{"m":93,"g":94},"7618":{"m":93,"g":94},"7598":{"m":93,"g":94},"7446":{"m":93,"g":94},"7584":{"m":93,"g":94},"7552":{"m":93,"g":94},"7308":{"m":93,"g":94},"6769":{"m":93,"g":94},"6563":{"m":93,"g":94},"7612":{"m":93,"g":94},"7588":{"m":93,"g":94},"7610":{"m":93,"g":94},"7581":{"m":93,"g":94},"7225":{"m":93,"g":94},"7577":{"m":93,"g":94},"7540":{"m":93,"g":94},"7208":{"m":93,"g":94},"7569":{"m":93,"g":94},"7573":{"m":93,"g":94},"7575":{"m":93,"g":94},"7330":{"m":93,"g":94},"7882":{"m":95,"g":97},"7880":{"m":95,"g":97},"7660":{"m":95,"g":97},"7846":{"m":95,"g":97},"7818":{"m":95,"g":97},"7866":{"m":95,"g":97},"7724":{"m":95,"g":97},"7830":{"m":95,"g":97},"7579":{"m":95,"g":97},"7840":{"m":95,"g":97},"7864":{"m":95,"g":97},"7860":{"m":95,"g":97},"7832":{"m":95,"g":97},"7853":{"m":95,"g":97},"7850":{"m":95,"g":97},"7129":{"m":95,"g":97},"7762":{"m":95,"g":97},"7821":{"m":95,"g":97},"7187":{"m":95,"g":97},"7816":{"m":95,"g":97},"7798":{"m":95,"g":94},"7797":{"m":95,"g":94},"7313":{"m":95,"g":94},"7689":{"m":95,"g":94},"7813":{"m":95,"g":94},"6094":{"m":95,"g":94},"7794":{"m":95,"g":94},"7812":{"m":95,"g":94},"7733":{"m":95,"g":94},"5246":{"m":95,"g":94},"7785":{"m":95,"g":94},"7709":{"m":95,"g":94},"7793":{"m":95,"g":94},"7803":{"m":95,"g":94},"7796":{"m":95,"g":94},"7963":{"m":96,"g":97},"7971":{"m":96,"g":97},"7960":{"m":96,"g":97},"7969":{"m":96,"g":97},"7970":{"m":96,"g":97},"7962":{"m":96,"g":97},"7968":{"m":96,"g":97},"7964":{"m":96,"g":97},"7932":{"m":96,"g":97},"7961":{"m":96,"g":97},"7953":{"m":96,"g":97},"7795":{"m":96,"g":97},"7940":{"m":96,"g":97},"7775":{"m":96,"g":97},"7791":{"m":96,"g":97},"6449":{"m":96,"g":97},"7922":{"m":96,"g":97},"7907":{"m":96,"g":97},"5888":{"m":96,"g":97},"7904":{"m":96,"g":97},"7608":{"m":96,"g":97},"7899":{"m":96,"g":97},"7898":{"m":96,"g":97},"7838":{"m":96,"g":97},"7885":{"m":96,"g":97},"7872":{"m":96,"g":97},"7895":{"m":96,"g":97},"8265":{"m":98,"g":103},"8260":{"m":98,"g":103},"7822":{"m":98,"g":103},"8059":{"m":98,"g":103},"8257":{"m":98,"g":103},"7484":{"m":98,"g":103},"8221":{"m":98,"g":103},"8237":{"m":98,"g":103},"8231":{"m":98,"g":103},"8202":{"m":98,"g":103},"8204":{"m":98,"g":103},"8209":{"m":98,"g":97},"8208":{"m":98,"g":97},"8107":{"m":98,"g":97},"8193":{"m":98,"g":97},"8200":{"m":98,"g":97},"8195":{"m":98,"g":97},"7935":{"m":98,"g":97},"8184":{"m":98,"g":97},"8197":{"m":98,"g":97},"8067":{"m":98,"g":97},"8163":{"m":98,"g":97},"8183":{"m":98,"g":97},"7983":{"m":98,"g":97},"8182":{"m":98,"g":97},"8181":{"m":98,"g":97},"7825":{"m":98,"g":97},"8178":{"m":98,"g":97},"8176":{"m":98,"g":97},"7312":{"m":98,"g":97},"6230":{"m":98,"g":97},"7999":{"m":98,"g":97},"8115":{"m":98,"g":97},"8175":{"m":98,"g":97},"8172":{"m":98,"g":97},"8019":{"m":98,"g":97},"8167":{"m":98,"g":97},"8170":{"m":98,"g":97},"8169":{"m":98,"g":97},"8103":{"m":98,"g":97},"8161":{"m":98,"g":97},"8171":{"m":98,"g":97},"8168":{"m":98,"g":97},"8166":{"m":98,"g":97},"8165":{"m":98,"g":97},"7966":{"m":98,"g":97},"8157":{"m":98,"g":97},"8160":{"m":98,"g":97},"8158":{"m":98,"g":97},"8028":{"m":98,"g":97},"7931":{"m":98,"g":97},"8048":{"m":98,"g":97},"7302":{"m":98,"g":97},"6881":{"m":98,"g":97},"8155":{"m":98,"g":97},"7661":{"m":98,"g":97},"7987":{"m":98,"g":97},"8113":{"m":98,"g":97},"8147":{"m":98,"g":97},"8142":{"m":98,"g":97},"8136":{"m":98,"g":97},"8141":{"m":98,"g":97},"7704":{"m":98,"g":97},"7820":{"m":98,"g":97},"7889":{"m":98,"g":97},"7506":{"m":98,"g":97},"7959":{"m":98,"g":97},"8127":{"m":98,"g":97},"7924":{"m":98,"g":97},"7030":{"m":98,"g":97},"8117":{"m":98,"g":97},"8102":{"m":98,"g":97},"8046":{"m":98,"g":97},"7884":{"m":98,"g":97},"7989":{"m":98,"g":97},"8105":{"m":98,"g":97},"8110":{"m":98,"g":97},"8108":{"m":98,"g":97},"8100":{"m":98,"g":97},"7597":{"m":98,"g":97},"8075":{"m":98,"g":97},"7992":{"m":98,"g":97},"7634":{"m":98,"g":97},"8098":{"m":98,"g":97},"8090":{"m":98,"g":97},"8086":{"m":98,"g":97},"8077":{"m":98,"g":97},"8001":{"m":98,"g":97},"7760":{"m":98,"g":97},"8029":{"m":98,"g":97},"8058":{"m":98,"g":97},"5163":{"m":98,"g":97},"8045":{"m":98,"g":97},"8047":{"m":98,"g":97},"7943":{"m":98,"g":97},"8052":{"m":98,"g":97},"6556":{"m":98,"g":97},"8022":{"m":98,"g":97},"8002":{"m":98,"g":97},"8023":{"m":98,"g":97},"8044":{"m":98,"g":97},"7887":{"m":98,"g":97},"8035":{"m":98,"g":97},"7897":{"m":98,"g":97},"7982":{"m":98,"g":97},"8006":{"m":98,"g":97},"7649":{"m":98,"g":97},"7653":{"m":98,"g":97},"8005":{"m":98,"g":97},"7874":{"m":98,"g":97},"8021":{"m":98,"g":97},"8010":{"m":98,"g":97},"7902":{"m":98,"g":97},"7862":{"m":98,"g":97},"7844":{"m":98,"g":97},"7997":{"m":98,"g":97},"7367":{"m":98,"g":97},"7749":{"m":98,"g":97},"7952":{"m":98,"g":97},"7988":{"m":98,"g":97},"7814":{"m":98,"g":97},"7978":{"m":98,"g":97},"7985":{"m":98,"g":97},"7975":{"m":98,"g":97},"7972":{"m":98,"g":97},"7950":{"m":98,"g":97},"8305":{"m":99,"g":103},"8370":{"m":99,"g":103},"8333":{"m":99,"g":103},"8367":{"m":99,"g":103},"8363":{"m":99,"g":103},"8359":{"m":99,"g":103},"8332":{"m":99,"g":103},"8357":{"m":99,"g":103},"8344":{"m":99,"g":103},"7858":{"m":99,"g":103},"8353":{"m":99,"g":103},"8341":{"m":99,"g":103},"8000":{"m":99,"g":103},"7135":{"m":99,"g":103},"8266":{"m":99,"g":103},"8280":{"m":99,"g":103},"6619":{"m":99,"g":103},"8233":{"m":99,"g":103},"8334":{"m":99,"g":103},"8307":{"m":99,"g":103},"8300":{"m":99,"g":103},"8299":{"m":99,"g":103},"8301":{"m":99,"g":103},"8310":{"m":99,"g":103},"8298":{"m":99,"g":103},"8315":{"m":99,"g":103},"8303":{"m":99,"g":103},"8317":{"m":99,"g":103},"8235":{"m":99,"g":103},"8070":{"m":99,"g":103},"7562":{"m":99,"g":103},"8043":{"m":99,"g":103},"7685":{"m":99,"g":103},"8304":{"m":99,"g":103},"8240":{"m":99,"g":103},"8262":{"m":99,"g":103},"7708":{"m":99,"g":103},"8133":{"m":99,"g":103},"8302":{"m":99,"g":103},"8295":{"m":99,"g":103},"8130":{"m":99,"g":103},"8288":{"m":99,"g":103},"8282":{"m":99,"g":103},"8264":{"m":99,"g":103},"8261":{"m":99,"g":103},"8284":{"m":99,"g":103},"8272":{"m":99,"g":103},"8458":{"m":100,"g":103},"8457":{"m":100,"g":103},"8449":{"m":100,"g":103},"8456":{"m":100,"g":103},"8445":{"m":100,"g":103},"8441":{"m":100,"g":103},"8442":{"m":100,"g":103},"8224":{"m":100,"g":103},"8352":{"m":100,"g":103},"8416":{"m":100,"g":103},"6338":{"m":100,"g":103},"8415":{"m":100,"g":103},"8422":{"m":100,"g":103},"8425":{"m":100,"g":103},"8419":{"m":100,"g":103},"8417":{"m":100,"g":103},"8316":{"m":100,"g":103},"7603":{"m":100,"g":103},"8213":{"m":100,"g":103},"8414":{"m":100,"g":103},"8062":{"m":100,"g":103},"8258":{"m":100,"g":103},"8406":{"m":100,"g":103},"8407":{"m":100,"g":103},"8405":{"m":100,"g":103},"8156":{"m":100,"g":103},"8241":{"m":100,"g":103},"8397":{"m":100,"g":103},"7720":{"m":100,"g":103},"8351":{"m":100,"g":103},"8395":{"m":100,"g":103},"8382":{"m":100,"g":103},"8392":{"m":100,"g":103},"8036":{"m":100,"g":103},"7739":{"m":100,"g":103},"7974":{"m":100,"g":103},"7976":{"m":100,"g":103},"8403":{"m":100,"g":103},"8401":{"m":100,"g":103},"8372":{"m":100,"g":103},"8394":{"m":100,"g":103},"8396":{"m":100,"g":103},"8314":{"m":100,"g":103},"8350":{"m":100,"g":103},"8381":{"m":100,"g":103},"6003":{"m":100,"g":103},"7737":{"m":100,"g":103},"8267":{"m":100,"g":103},"8343":{"m":100,"g":103},"7000":{"m":100,"g":103},"8356":{"m":100,"g":103},"8374":{"m":100,"g":103},"8517":{"m":101,"g":103},"8489":{"m":101,"g":103},"8482":{"m":101,"g":103},"7973":{"m":101,"g":103},"8426":{"m":101,"g":103},"8413":{"m":101,"g":103},"8486":{"m":101,"g":103},"8485":{"m":101,"g":103},"8477":{"m":101,"g":103},"8480":{"m":101,"g":103},"8478":{"m":101,"g":103},"8476":{"m":101,"g":103},"8469":{"m":101,"g":103},"8473":{"m":101,"g":103},"8421":{"m":101,"g":103},"7273":{"m":101,"g":103},"8453":{"m":101,"g":103},"8467":{"m":101,"g":103},"7565":{"m":101,"g":103},"8465":{"m":101,"g":103},"8125":{"m":101,"g":103},"8608":{"m":102,"g":103},"8590":{"m":102,"g":103},"8583":{"m":102,"g":103},"8604":{"m":102,"g":103},"8515":{"m":102,"g":103},"8603":{"m":102,"g":103},"8550":{"m":102,"g":103},"7211":{"m":102,"g":103},"8533":{"m":102,"g":103},"8404":{"m":102,"g":103},"8599":{"m":102,"g":103},"8514":{"m":102,"g":103},"8365":{"m":102,"g":103},"8544":{"m":102,"g":103},"8541":{"m":102,"g":103},"8564":{"m":102,"g":103},"8479":{"m":102,"g":103},"7280":{"m":102,"g":103},"8584":{"m":102,"g":103},"8154":{"m":102,"g":103},"8461":{"m":102,"g":103},"6869":{"m":102,"g":103},"8562":{"m":102,"g":103},"8545":{"m":102,"g":103},"8560":{"m":102,"g":103},"8516":{"m":102,"g":103},"8498":{"m":102,"g":103},"8448":{"m":102,"g":103},"8431":{"m":102,"g":103},"8537":{"m":102,"g":103},"8483":{"m":102,"g":103},"8531":{"m":102,"g":103},"8535":{"m":102,"g":103},"8499":{"m":102,"g":103},"8528":{"m":102,"g":103},"8527":{"m":102,"g":103},"8652":{"m":105,"g":107},"8051":{"m":105,"g":107},"8318":{"m":105,"g":107},"8636":{"m":105,"g":107},"8450":{"m":105,"g":107},"8645":{"m":105,"g":104},"8640":{"m":105,"g":104},"8644":{"m":105,"g":104},"8308":{"m":105,"g":104},"8270":{"m":105,"g":104},"8083":{"m":105,"g":104},"8642":{"m":105,"g":104},"8532":{"m":105,"g":104},"8632":{"m":105,"g":104},"8598":{"m":105,"g":104},"8630":{"m":105,"g":104},"8634":{"m":105,"g":104},"8488":{"m":105,"g":104},"8633":{"m":105,"g":104},"6227":{"m":105,"g":104},"8577":{"m":105,"g":104},"8628":{"m":105,"g":103},"8629":{"m":105,"g":103},"8626":{"m":105,"g":103},"8623":{"m":105,"g":103},"8611":{"m":105,"g":103},"8595":{"m":105,"g":103},"8727":{"m":106,"g":107},"8723":{"m":106,"g":107},"8579":{"m":106,"g":107},"8567":{"m":106,"g":107},"8718":{"m":106,"g":107},"8444":{"m":106,"g":107},"8547":{"m":106,"g":107},"8683":{"m":106,"g":107},"8631":{"m":106,"g":107},"7379":{"m":106,"g":107},"8719":{"m":106,"g":107},"8306":{"m":106,"g":107},"8650":{"m":106,"g":107},"8721":{"m":106,"g":107},"8524":{"m":106,"g":107},"8709":{"m":106,"g":107},"8722":{"m":106,"g":107},"7369":{"m":106,"g":107},"8714":{"m":106,"g":107},"8705":{"m":106,"g":107},"8693":{"m":106,"g":107},"8717":{"m":106,"g":107},"8713":{"m":106,"g":107},"8701":{"m":106,"g":107},"8711":{"m":106,"g":107},"8706":{"m":106,"g":107},"8704":{"m":106,"g":107},"8691":{"m":106,"g":107},"8512":{"m":106,"g":107},"7434":{"m":106,"g":107},"8688":{"m":106,"g":107},"8694":{"m":106,"g":107},"8364":{"m":106,"g":107},"8238":{"m":106,"g":107},"8618":{"m":106,"g":107},"8522":{"m":106,"g":107},"8668":{"m":106,"g":107},"8648":{"m":106,"g":107},"8679":{"m":106,"g":107},"8686":{"m":106,"g":107},"8684":{"m":106,"g":107},"8685":{"m":106,"g":107},"8647":{"m":106,"g":107},"8094":{"m":106,"g":107},"8664":{"m":106,"g":107},"8543":{"m":106,"g":107},"8665":{"m":106,"g":107},"8013":{"m":106,"g":107},"8643":{"m":106,"g":107},"8658":{"m":106,"g":107},"8511":{"m":106,"g":107},"8635":{"m":106,"g":107},"8653":{"m":106,"g":107},"9533":{"m":108,"g":115},"9532":{"m":108,"g":115},"9372":{"m":108,"g":115},"9485":{"m":108,"g":115},"9478":{"m":108,"g":115},"8034":{"m":108,"g":115},"9473":{"m":108,"g":115},"9525":{"m":108,"g":115},"9530":{"m":108,"g":115},"8946":{"m":108,"g":115},"9004":{"m":108,"g":115},"9241":{"m":108,"g":115},"9211":{"m":108,"g":115},"9503":{"m":108,"g":115},"9519":{"m":108,"g":115},"9456":{"m":108,"g":115},"7699":{"m":108,"g":115},"9200":{"m":108,"g":115},"9516":{"m":108,"g":115},"9127":{"m":108,"g":115},"9513":{"m":108,"g":115},"8624":{"m":108,"g":115},"8865":{"m":108,"g":115},"9109":{"m":108,"g":115},"9452":{"m":108,"g":115},"9507":{"m":108,"g":115},"9303":{"m":108,"g":115},"9331":{"m":108,"g":115},"9497":{"m":108,"g":115},"9494":{"m":108,"g":115},"9475":{"m":108,"g":115},"9480":{"m":108,"g":115},"9487":{"m":108,"g":115},"9491":{"m":108,"g":115},"9482":{"m":108,"g":115},"9492":{"m":108,"g":115},"9483":{"m":108,"g":115},"9333":{"m":108,"g":115},"9474":{"m":108,"g":115},"8616":{"m":108,"g":115},"9356":{"m":108,"g":115},"8593":{"m":108,"g":115},"9468":{"m":108,"g":115},"9467":{"m":108,"g":115},"9470":{"m":108,"g":115},"9469":{"m":108,"g":115},"9455":{"m":108,"g":115},"9463":{"m":108,"g":115},"9464":{"m":108,"g":115},"9462":{"m":108,"g":115},"9461":{"m":108,"g":115},"9458":{"m":108,"g":115},"9454":{"m":108,"g":115},"9427":{"m":108,"g":115},"9433":{"m":108,"g":115},"7604":{"m":108,"g":115},"8521":{"m":108,"g":115},"9392":{"m":108,"g":115},"9395":{"m":108,"g":115},"9238":{"m":108,"g":115},"9384":{"m":108,"g":115},"9430":{"m":108,"g":115},"9346":{"m":108,"g":115},"9399":{"m":108,"g":115},"9251":{"m":108,"g":115},"9388":{"m":108,"g":115},"9261":{"m":108,"g":115},"9420":{"m":108,"g":115},"9416":{"m":108,"g":115},"9415":{"m":108,"g":115},"9413":{"m":108,"g":115},"9371":{"m":108,"g":115},"9339":{"m":108,"g":115},"9357":{"m":108,"g":115},"9377":{"m":108,"g":115},"9381":{"m":108,"g":115},"9404":{"m":108,"g":115},"9359":{"m":108,"g":115},"9336":{"m":108,"g":115},"9249":{"m":108,"g":115},"9409":{"m":108,"g":115},"8690":{"m":108,"g":115},"9278":{"m":108,"g":115},"9391":{"m":108,"g":115},"9106":{"m":108,"g":115},"7375":{"m":108,"g":115},"9385":{"m":108,"g":115},"9383":{"m":108,"g":115},"9378":{"m":108,"g":115},"9380":{"m":108,"g":115},"9376":{"m":108,"g":115},"9350":{"m":108,"g":115},"9368":{"m":108,"g":115},"9344":{"m":108,"g":115},"9369":{"m":108,"g":115},"9367":{"m":108,"g":115},"9370":{"m":108,"g":115},"9364":{"m":108,"g":115},"9360":{"m":108,"g":115},"9335":{"m":108,"g":115},"9361":{"m":108,"g":115},"9354":{"m":108,"g":115},"6295":{"m":108,"g":115},"9353":{"m":108,"g":115},"9348":{"m":108,"g":115},"9327":{"m":108,"g":115},"9332":{"m":108,"g":115},"9326":{"m":108,"g":115},"8990":{"m":108,"g":115},"7019":{"m":108,"g":115},"9321":{"m":108,"g":115},"9317":{"m":108,"g":115},"9299":{"m":108,"g":115},"9322":{"m":108,"g":115},"9306":{"m":108,"g":115},"9320":{"m":108,"g":115},"9059":{"m":108,"g":115},"8936":{"m":108,"g":115},"9284":{"m":108,"g":115},"9313":{"m":108,"g":115},"9316":{"m":108,"g":115},"9315":{"m":108,"g":115},"8829":{"m":108,"g":115},"9011":{"m":108,"g":115},"9289":{"m":108,"g":115},"9310":{"m":108,"g":115},"9307":{"m":108,"g":115},"9298":{"m":108,"g":115},"9276":{"m":108,"g":115},"9293":{"m":108,"g":115},"6307":{"m":108,"g":115},"9245":{"m":108,"g":115},"9287":{"m":108,"g":115},"9286":{"m":108,"g":115},"8289":{"m":108,"g":115},"9281":{"m":108,"g":115},"8520":{"m":108,"g":115},"9272":{"m":108,"g":115},"9271":{"m":108,"g":115},"9279":{"m":108,"g":115},"9131":{"m":108,"g":115},"9242":{"m":108,"g":115},"9260":{"m":108,"g":115},"9268":{"m":108,"g":115},"9067":{"m":108,"g":115},"9264":{"m":108,"g":115},"9237":{"m":108,"g":115},"9232":{"m":108,"g":115},"9006":{"m":108,"g":115},"8893":{"m":108,"g":115},"9049":{"m":108,"g":115},"8846":{"m":108,"g":115},"9252":{"m":108,"g":115},"8027":{"m":108,"g":115},"9258":{"m":108,"g":115},"7758":{"m":108,"g":115},"9165":{"m":108,"g":115},"7667":{"m":108,"g":115},"9247":{"m":108,"g":115},"9246":{"m":108,"g":115},"8663":{"m":108,"g":115},"8268":{"m":108,"g":115},"9243":{"m":108,"g":115},"9236":{"m":108,"g":115},"9201":{"m":108,"g":115},"9198":{"m":108,"g":115},"9231":{"m":108,"g":115},"9220":{"m":108,"g":115},"9223":{"m":108,"g":115},"9222":{"m":108,"g":115},"8777":{"m":108,"g":115},"8790":{"m":108,"g":115},"9215":{"m":108,"g":115},"9218":{"m":108,"g":115},"9214":{"m":108,"g":115},"9208":{"m":108,"g":115},"9213":{"m":108,"g":115},"9207":{"m":108,"g":115},"9177":{"m":108,"g":115},"8849":{"m":108,"g":115},"9206":{"m":108,"g":115},"9205":{"m":108,"g":115},"9183":{"m":108,"g":115},"9204":{"m":108,"g":115},"9203":{"m":108,"g":115},"9202":{"m":108,"g":115},"9197":{"m":108,"g":115},"9008":{"m":108,"g":115},"8795":{"m":108,"g":115},"9191":{"m":108,"g":115},"9194":{"m":108,"g":115},"9060":{"m":108,"g":115},"8913":{"m":108,"g":115},"9185":{"m":108,"g":115},"8112":{"m":108,"g":115},"9065":{"m":108,"g":115},"8018":{"m":108,"g":115},"7687":{"m":108,"g":115},"7631":{"m":108,"g":115},"7004":{"m":108,"g":115},"8852":{"m":108,"g":115},"8808":{"m":108,"g":115},"8818":{"m":108,"g":115},"9154":{"m":108,"g":115},"9101":{"m":108,"g":115},"9162":{"m":108,"g":115},"9136":{"m":108,"g":115},"9171":{"m":108,"g":115},"9169":{"m":108,"g":115},"9159":{"m":108,"g":115},"8951":{"m":108,"g":115},"9161":{"m":108,"g":115},"8840":{"m":108,"g":115},"9134":{"m":108,"g":115},"9042":{"m":108,"g":115},"7957":{"m":108,"g":115},"9069":{"m":108,"g":115},"9028":{"m":108,"g":115},"8910":{"m":108,"g":115},"9149":{"m":108,"g":115},"9133":{"m":108,"g":115},"9126":{"m":108,"g":115},"9150":{"m":108,"g":115},"8484":{"m":108,"g":115},"9111":{"m":108,"g":115},"9146":{"m":108,"g":115},"9093":{"m":108,"g":115},"9088":{"m":108,"g":115},"8588":{"m":108,"g":115},"9137":{"m":108,"g":115},"8884":{"m":108,"g":115},"8651":{"m":108,"g":115},"9130":{"m":108,"g":115},"9129":{"m":108,"g":115},"8660":{"m":108,"g":115},"9119":{"m":108,"g":115},"8619":{"m":108,"g":115},"8610":{"m":108,"g":115},"8700":{"m":108,"g":115},"9125":{"m":108,"g":115},"9121":{"m":108,"g":115},"9014":{"m":108,"g":115},"9118":{"m":108,"g":115},"9122":{"m":108,"g":115},"9107":{"m":108,"g":115},"9113":{"m":108,"g":115},"9114":{"m":108,"g":115},"9021":{"m":108,"g":115},"9103":{"m":108,"g":115},"9077":{"m":108,"g":115},"9096":{"m":108,"g":115},"9005":{"m":108,"g":115},"9075":{"m":108,"g":115},"9097":{"m":108,"g":115},"9032":{"m":108,"g":115},"8766":{"m":108,"g":115},"9087":{"m":108,"g":115},"8293":{"m":108,"g":115},"9095":{"m":108,"g":115},"9084":{"m":108,"g":115},"8992":{"m":108,"g":115},"9089":{"m":108,"g":115},"9043":{"m":108,"g":115},"9086":{"m":108,"g":115},"9030":{"m":108,"g":115},"9053":{"m":108,"g":115},"8638":{"m":108,"g":115},"8731":{"m":108,"g":115},"9083":{"m":108,"g":115},"8752":{"m":108,"g":115},"9081":{"m":108,"g":115},"9082":{"m":108,"g":115},"8866":{"m":108,"g":115},"9080":{"m":108,"g":115},"8973":{"m":108,"g":115},"9063":{"m":108,"g":115},"9079":{"m":108,"g":115},"9066":{"m":108,"g":115},"7216":{"m":108,"g":115},"9051":{"m":108,"g":115},"9047":{"m":108,"g":115},"9057":{"m":108,"g":115},"9050":{"m":108,"g":115},"9054":{"m":108,"g":115},"9048":{"m":108,"g":115},"8997":{"m":108,"g":115},"9046":{"m":108,"g":115},"9044":{"m":108,"g":115},"9031":{"m":108,"g":115},"9037":{"m":108,"g":115},"9036":{"m":108,"g":115},"9034":{"m":108,"g":115},"9035":{"m":108,"g":115},"9033":{"m":108,"g":115},"8079":{"m":108,"g":115},"8794":{"m":108,"g":115},"9024":{"m":108,"g":115},"9027":{"m":108,"g":115},"9029":{"m":108,"g":115},"7626":{"m":108,"g":115},"9022":{"m":108,"g":115},"8940":{"m":108,"g":115},"8996":{"m":108,"g":115},"9018":{"m":108,"g":115},"9017":{"m":108,"g":115},"9019":{"m":108,"g":115},"8340":{"m":108,"g":115},"8991":{"m":108,"g":115},"8915":{"m":108,"g":115},"8245":{"m":108,"g":115},"9013":{"m":108,"g":115},"9007":{"m":108,"g":115},"9012":{"m":108,"g":115},"9003":{"m":108,"g":115},"9010":{"m":108,"g":115},"9001":{"m":108,"g":115},"8329":{"m":108,"g":115},"8355":{"m":108,"g":115},"8877":{"m":108,"g":115},"8995":{"m":108,"g":115},"8998":{"m":108,"g":115},"8878":{"m":108,"g":115},"6752":{"m":108,"g":115},"8673":{"m":108,"g":115},"8798":{"m":108,"g":115},"8687":{"m":108,"g":115},"8600":{"m":108,"g":115},"8966":{"m":108,"g":115},"8851":{"m":108,"g":115},"8984":{"m":108,"g":115},"8962":{"m":108,"g":115},"8987":{"m":108,"g":115},"8994":{"m":108,"g":115},"8993":{"m":108,"g":115},"8989":{"m":108,"g":115},"8667":{"m":108,"g":115},"8983":{"m":108,"g":115},"8980":{"m":108,"g":115},"8330":{"m":108,"g":115},"8770":{"m":108,"g":115},"8724":{"m":108,"g":115},"8988":{"m":108,"g":115},"8986":{"m":108,"g":115},"8785":{"m":108,"g":115},"8371":{"m":108,"g":115},"8982":{"m":108,"g":115},"8978":{"m":108,"g":115},"8981":{"m":108,"g":115},"8772":{"m":108,"g":115},"8971":{"m":108,"g":115},"8968":{"m":108,"g":115},"8972":{"m":108,"g":115},"7279":{"m":108,"g":115},"8941":{"m":108,"g":115},"8959":{"m":108,"g":115},"8757":{"m":108,"g":115},"8960":{"m":108,"g":115},"8958":{"m":108,"g":115},"8692":{"m":108,"g":115},"6555":{"m":108,"g":115},"8894":{"m":108,"g":115},"8957":{"m":108,"g":115},"8955":{"m":108,"g":115},"7657":{"m":108,"g":115},"8944":{"m":108,"g":115},"8799":{"m":108,"g":115},"8932":{"m":108,"g":115},"8952":{"m":108,"g":115},"8953":{"m":108,"g":115},"8947":{"m":108,"g":115},"8950":{"m":108,"g":115},"8720":{"m":108,"g":115},"8703":{"m":108,"g":115},"8923":{"m":108,"g":115},"8850":{"m":108,"g":115},"8933":{"m":108,"g":115},"8937":{"m":108,"g":115},"8929":{"m":108,"g":115},"8928":{"m":108,"g":115},"8925":{"m":108,"g":115},"8927":{"m":108,"g":115},"8908":{"m":108,"g":115},"8844":{"m":108,"g":107},"8916":{"m":108,"g":107},"8912":{"m":108,"g":107},"8698":{"m":108,"g":107},"8869":{"m":108,"g":107},"8898":{"m":108,"g":107},"8895":{"m":108,"g":107},"8041":{"m":108,"g":107},"5949":{"m":108,"g":107},"8888":{"m":108,"g":107},"8292":{"m":108,"g":107},"8787":{"m":108,"g":107},"8369":{"m":108,"g":107},"8697":{"m":108,"g":107},"8834":{"m":108,"g":107},"8883":{"m":108,"g":107},"8847":{"m":108,"g":107},"8539":{"m":108,"g":107},"8837":{"m":108,"g":107},"8880":{"m":108,"g":107},"8881":{"m":108,"g":107},"8811":{"m":108,"g":107},"8872":{"m":108,"g":107},"8861":{"m":108,"g":107},"8860":{"m":108,"g":107},"8868":{"m":108,"g":107},"8815":{"m":108,"g":107},"8859":{"m":108,"g":107},"8853":{"m":108,"g":107},"8843":{"m":108,"g":107},"8753":{"m":108,"g":107},"8751":{"m":108,"g":107},"8838":{"m":108,"g":107},"8144":{"m":108,"g":107},"8680":{"m":108,"g":107},"8839":{"m":108,"g":107},"8828":{"m":108,"g":107},"8836":{"m":108,"g":107},"8824":{"m":108,"g":107},"8809":{"m":108,"g":107},"8832":{"m":108,"g":107},"8681":{"m":108,"g":107},"8827":{"m":108,"g":107},"8823":{"m":108,"g":107},"8817":{"m":108,"g":107},"8804":{"m":108,"g":107},"8782":{"m":108,"g":107},"8802":{"m":108,"g":107},"8800":{"m":108,"g":107},"8797":{"m":108,"g":107},"8596":{"m":108,"g":107},"8779":{"m":108,"g":107},"8780":{"m":108,"g":107},"8744":{"m":108,"g":107},"8571":{"m":108,"g":107},"8212":{"m":108,"g":107},"8255":{"m":108,"g":107},"8776":{"m":108,"g":107},"8773":{"m":108,"g":107},"8771":{"m":108,"g":107},"8762":{"m":108,"g":107},"8768":{"m":108,"g":107},"8639":{"m":108,"g":107},"8552":{"m":108,"g":107},"8749":{"m":108,"g":107},"8294":{"m":108,"g":107},"8738":{"m":108,"g":107},"8437":{"m":108,"g":107},"8745":{"m":108,"g":107},"8733":{"m":108,"g":107},"8735":{"m":108,"g":107},"8737":{"m":108,"g":107},"8662":{"m":108,"g":107},"7114":{"m":108,"g":107},"8678":{"m":108,"g":107},"8732":{"m":108,"g":107},"8729":{"m":108,"g":107},"8676":{"m":108,"g":107},"8699":{"m":108,"g":107},"9558":{"m":109,"g":115},"9557":{"m":109,"g":115},"9549":{"m":109,"g":115},"9544":{"m":109,"g":115},"9547":{"m":109,"g":115},"9546":{"m":109,"g":115},"9592":{"m":110,"g":115},"9591":{"m":110,"g":115},"9589":{"m":110,"g":115},"9587":{"m":110,"g":115},"9581":{"m":110,"g":115},"9578":{"m":110,"g":115},"9229":{"m":110,"g":115},"9536":{"m":110,"g":115},"9535":{"m":110,"g":115},"9559":{"m":110,"g":115},"9560":{"m":110,"g":115},"7317":{"m":110,"g":115},"9576":{"m":110,"g":115},"9429":{"m":110,"g":115},"9565":{"m":110,"g":115},"9498":{"m":110,"g":115},"9716":{"m":111,"g":115},"9708":{"m":111,"g":115},"9340":{"m":111,"g":115},"9703":{"m":111,"g":115},"9702":{"m":111,"g":115},"9695":{"m":111,"g":115},"9683":{"m":111,"g":115},"9700":{"m":111,"g":115},"9676":{"m":111,"g":115},"9694":{"m":111,"g":115},"9693":{"m":111,"g":115},"9679":{"m":111,"g":115},"9678":{"m":111,"g":115},"9397":{"m":111,"g":115},"9495":{"m":111,"g":115},"9677":{"m":111,"g":115},"9446":{"m":111,"g":115},"9071":{"m":111,"g":115},"9597":{"m":111,"g":115},"9555":{"m":111,"g":115},"9583":{"m":111,"g":115},"9564":{"m":111,"g":115},"9658":{"m":111,"g":115},"9665":{"m":111,"g":115},"9648":{"m":111,"g":115},"9637":{"m":111,"g":115},"9649":{"m":111,"g":115},"9647":{"m":111,"g":115},"9656":{"m":111,"g":115},"9523":{"m":111,"g":115},"9606":{"m":111,"g":115},"9635":{"m":111,"g":115},"9630":{"m":111,"g":115},"9640":{"m":111,"g":115},"9636":{"m":111,"g":115},"9301":{"m":111,"g":115},"8328":{"m":111,"g":115},"9632":{"m":111,"g":115},"9629":{"m":111,"g":115},"9628":{"m":111,"g":115},"9623":{"m":111,"g":115},"9622":{"m":111,"g":115},"9608":{"m":111,"g":115},"8901":{"m":111,"g":115},"9613":{"m":111,"g":115},"9190":{"m":111,"g":115},"9554":{"m":111,"g":115},"9500":{"m":111,"g":115},"9436":{"m":111,"g":115},"9568":{"m":111,"g":115},"10221":{"m":112,"g":115},"10340":{"m":112,"g":115},"10303":{"m":112,"g":115},"10331":{"m":112,"g":115},"10339":{"m":112,"g":115},"10330":{"m":112,"g":115},"10327":{"m":112,"g":115},"10338":{"m":112,"g":115},"10254":{"m":112,"g":115},"10264":{"m":112,"g":115},"10280":{"m":112,"g":115},"10335":{"m":112,"g":115},"10322":{"m":112,"g":115},"10328":{"m":112,"g":115},"10326":{"m":112,"g":115},"10233":{"m":112,"g":115},"10297":{"m":112,"g":115},"10314":{"m":112,"g":115},"10311":{"m":112,"g":115},"10310":{"m":112,"g":115},"10299":{"m":112,"g":115},"9090":{"m":112,"g":115},"10229":{"m":112,"g":115},"10239":{"m":112,"g":115},"9881":{"m":112,"g":115},"10294":{"m":112,"g":115},"10292":{"m":112,"g":115},"10282":{"m":112,"g":115},"10184":{"m":112,"g":115},"10241":{"m":112,"g":115},"9662":{"m":112,"g":115},"10252":{"m":112,"g":115},"9940":{"m":112,"g":115},"10251":{"m":112,"g":115},"10256":{"m":112,"g":115},"10250":{"m":112,"g":115},"10262":{"m":112,"g":115},"9954":{"m":112,"g":115},"10173":{"m":112,"g":115},"10060":{"m":112,"g":115},"10253":{"m":112,"g":115},"10093":{"m":112,"g":115},"10240":{"m":112,"g":115},"8803":{"m":112,"g":115},"10246":{"m":112,"g":115},"9795":{"m":112,"g":115},"10245":{"m":112,"g":115},"10242":{"m":112,"g":115},"10236":{"m":112,"g":115},"10234":{"m":112,"g":115},"10213":{"m":112,"g":115},"10238":{"m":112,"g":115},"10210":{"m":112,"g":115},"10220":{"m":112,"g":115},"10214":{"m":112,"g":115},"9960":{"m":112,"g":115},"10208":{"m":112,"g":115},"10212":{"m":112,"g":115},"10209":{"m":112,"g":115},"10207":{"m":112,"g":115},"10205":{"m":112,"g":115},"9300":{"m":112,"g":115},"10193":{"m":112,"g":115},"10127":{"m":112,"g":115},"10188":{"m":112,"g":115},"10165":{"m":112,"g":115},"10169":{"m":112,"g":115},"4422":{"m":112,"g":115},"9900":{"m":112,"g":115},"10191":{"m":112,"g":115},"7995":{"m":112,"g":115},"10185":{"m":112,"g":115},"10149":{"m":112,"g":115},"9522":{"m":112,"g":115},"10182":{"m":112,"g":115},"10181":{"m":112,"g":115},"9839":{"m":112,"g":115},"10176":{"m":112,"g":115},"9595":{"m":112,"g":115},"10166":{"m":112,"g":115},"9925":{"m":112,"g":115},"10156":{"m":112,"g":115},"10161":{"m":112,"g":115},"10159":{"m":112,"g":115},"10131":{"m":112,"g":115},"10028":{"m":112,"g":115},"10155":{"m":112,"g":115},"9871":{"m":112,"g":115},"9434":{"m":112,"g":115},"10148":{"m":112,"g":115},"6226":{"m":112,"g":115},"10013":{"m":112,"g":115},"10147":{"m":112,"g":115},"9981":{"m":112,"g":115},"9989":{"m":112,"g":115},"10108":{"m":112,"g":115},"10123":{"m":112,"g":115},"7843":{"m":112,"g":115},"10090":{"m":112,"g":115},"10104":{"m":112,"g":115},"8801":{"m":112,"g":115},"10040":{"m":112,"g":115},"10141":{"m":112,"g":115},"10095":{"m":112,"g":115},"10144":{"m":112,"g":115},"9971":{"m":112,"g":115},"10134":{"m":112,"g":115},"10135":{"m":112,"g":115},"10074":{"m":112,"g":115},"10128":{"m":112,"g":115},"10126":{"m":112,"g":115},"10113":{"m":112,"g":115},"10096":{"m":112,"g":115},"10056":{"m":112,"g":115},"10101":{"m":112,"g":115},"9969":{"m":112,"g":115},"9741":{"m":112,"g":115},"9477":{"m":112,"g":115},"10117":{"m":112,"g":115},"10102":{"m":112,"g":115},"10116":{"m":112,"g":115},"10068":{"m":112,"g":115},"9956":{"m":112,"g":115},"10041":{"m":112,"g":115},"10058":{"m":112,"g":115},"10107":{"m":112,"g":115},"10100":{"m":112,"g":115},"9834":{"m":112,"g":115},"9861":{"m":112,"g":115},"9764":{"m":112,"g":115},"9269":{"m":112,"g":115},"9620":{"m":112,"g":115},"6905":{"m":112,"g":115},"10032":{"m":112,"g":115},"10097":{"m":112,"g":115},"10029":{"m":112,"g":115},"10039":{"m":112,"g":115},"10092":{"m":112,"g":115},"10086":{"m":112,"g":115},"10057":{"m":112,"g":115},"10047":{"m":112,"g":115},"9842":{"m":112,"g":115},"9965":{"m":112,"g":115},"10069":{"m":112,"g":115},"9884":{"m":112,"g":115},"10087":{"m":112,"g":115},"8622":{"m":112,"g":115},"8555":{"m":112,"g":115},"10080":{"m":112,"g":115},"10079":{"m":112,"g":115},"10043":{"m":112,"g":115},"10007":{"m":112,"g":115},"7182":{"m":112,"g":115},"8725":{"m":112,"g":115},"8867":{"m":112,"g":115},"9567":{"m":112,"g":115},"5255":{"m":112,"g":115},"10006":{"m":112,"g":115},"9534":{"m":112,"g":115},"9934":{"m":112,"g":115},"9931":{"m":112,"g":115},"9801":{"m":112,"g":115},"10049":{"m":112,"g":115},"10055":{"m":112,"g":115},"9964":{"m":112,"g":115},"10052":{"m":112,"g":115},"10050":{"m":112,"g":115},"8677":{"m":112,"g":115},"10008":{"m":112,"g":115},"9951":{"m":112,"g":115},"9957":{"m":112,"g":115},"9938":{"m":112,"g":115},"9634":{"m":112,"g":115},"9886":{"m":112,"g":115},"9973":{"m":112,"g":115},"9846":{"m":112,"g":115},"10003":{"m":112,"g":115},"10004":{"m":112,"g":115},"9997":{"m":112,"g":115},"10016":{"m":112,"g":115},"9993":{"m":112,"g":115},"9994":{"m":112,"g":115},"9999":{"m":112,"g":115},"10000":{"m":112,"g":115},"9996":{"m":112,"g":115},"9988":{"m":112,"g":115},"9986":{"m":112,"g":115},"9733":{"m":112,"g":115},"9314":{"m":112,"g":115},"9978":{"m":112,"g":115},"9914":{"m":112,"g":115},"9460":{"m":112,"g":115},"9958":{"m":112,"g":115},"9906":{"m":112,"g":115},"9953":{"m":112,"g":115},"9937":{"m":112,"g":115},"9959":{"m":112,"g":115},"9955":{"m":112,"g":115},"9755":{"m":112,"g":115},"9952":{"m":112,"g":115},"9671":{"m":112,"g":115},"7912":{"m":112,"g":115},"9895":{"m":112,"g":115},"9905":{"m":112,"g":115},"9927":{"m":112,"g":115},"9946":{"m":112,"g":115},"9912":{"m":112,"g":115},"9869":{"m":112,"g":115},"8747":{"m":112,"g":115},"9939":{"m":112,"g":115},"9929":{"m":112,"g":115},"9932":{"m":112,"g":115},"9705":{"m":112,"g":115},"9909":{"m":112,"g":115},"9920":{"m":112,"g":115},"9921":{"m":112,"g":115},"9879":{"m":112,"g":115},"9919":{"m":112,"g":115},"9916":{"m":112,"g":115},"9907":{"m":112,"g":115},"9844":{"m":112,"g":115},"9913":{"m":112,"g":115},"9902":{"m":112,"g":115},"8118":{"m":112,"g":115},"9875":{"m":112,"g":115},"9893":{"m":112,"g":115},"9878":{"m":112,"g":115},"9803":{"m":112,"g":115},"9876":{"m":112,"g":115},"9874":{"m":112,"g":115},"9783":{"m":112,"g":115},"9882":{"m":112,"g":115},"9857":{"m":112,"g":115},"9862":{"m":112,"g":115},"9864":{"m":112,"g":115},"8964":{"m":112,"g":115},"9794":{"m":112,"g":115},"9858":{"m":112,"g":115},"9852":{"m":112,"g":115},"9847":{"m":112,"g":115},"9073":{"m":112,"g":115},"9850":{"m":112,"g":115},"9797":{"m":112,"g":115},"9661":{"m":112,"g":115},"9841":{"m":112,"g":115},"9750":{"m":112,"g":115},"9709":{"m":112,"g":115},"8909":{"m":112,"g":115},"9840":{"m":112,"g":115},"9824":{"m":112,"g":115},"9835":{"m":112,"g":115},"9837":{"m":112,"g":115},"9836":{"m":112,"g":115},"9831":{"m":112,"g":115},"9830":{"m":112,"g":115},"9761":{"m":112,"g":115},"9828":{"m":112,"g":115},"9827":{"m":112,"g":115},"9826":{"m":112,"g":115},"9822":{"m":112,"g":115},"9802":{"m":112,"g":115},"9820":{"m":112,"g":115},"9746":{"m":112,"g":115},"9817":{"m":112,"g":115},"9815":{"m":112,"g":115},"9807":{"m":112,"g":115},"8345":{"m":112,"g":115},"9809":{"m":112,"g":115},"9670":{"m":112,"g":115},"9556":{"m":112,"g":115},"9675":{"m":112,"g":115},"9712":{"m":112,"g":115},"9793":{"m":112,"g":115},"9216":{"m":112,"g":115},"8375":{"m":112,"g":115},"9663":{"m":112,"g":115},"9715":{"m":112,"g":115},"9692":{"m":112,"g":115},"9776":{"m":112,"g":115},"9792":{"m":112,"g":115},"9786":{"m":112,"g":115},"9789":{"m":112,"g":115},"9788":{"m":112,"g":115},"9757":{"m":112,"g":115},"9784":{"m":112,"g":115},"9777":{"m":112,"g":115},"9749":{"m":112,"g":115},"9772":{"m":112,"g":115},"8750":{"m":112,"g":115},"6287":{"m":112,"g":115},"6407":{"m":112,"g":115},"9355":{"m":112,"g":115},"9770":{"m":112,"g":115},"8236":{"m":112,"g":115},"9759":{"m":112,"g":115},"9745":{"m":112,"g":115},"9721":{"m":112,"g":115},"9735":{"m":112,"g":115},"9573":{"m":112,"g":115},"9673":{"m":112,"g":115},"9740":{"m":112,"g":115},"9739":{"m":112,"g":115},"9684":{"m":112,"g":115},"9732":{"m":112,"g":115},"9730":{"m":112,"g":115},"9728":{"m":112,"g":115},"9615":{"m":112,"g":115},"9505":{"m":112,"g":115},"9724":{"m":112,"g":115},"9720":{"m":112,"g":115},"11263":{"m":113,"g":115},"11259":{"m":113,"g":115},"11061":{"m":113,"g":115},"11235":{"m":113,"g":115},"11240":{"m":113,"g":115},"11209":{"m":113,"g":115},"11254":{"m":113,"g":115},"11242":{"m":113,"g":115},"11252":{"m":113,"g":115},"11251":{"m":113,"g":115},"10048":{"m":113,"g":115},"10042":{"m":113,"g":115},"11248":{"m":113,"g":115},"11247":{"m":113,"g":115},"10996":{"m":113,"g":115},"11206":{"m":113,"g":115},"11228":{"m":113,"g":115},"11237":{"m":113,"g":115},"11222":{"m":113,"g":115},"11174":{"m":113,"g":115},"11162":{"m":113,"g":115},"10571":{"m":113,"g":115},"11229":{"m":113,"g":115},"11137":{"m":113,"g":115},"9624":{"m":113,"g":115},"11194":{"m":113,"g":115},"11225":{"m":113,"g":115},"11063":{"m":113,"g":115},"11217":{"m":113,"g":115},"11215":{"m":113,"g":115},"11140":{"m":113,"g":115},"11213":{"m":113,"g":115},"11012":{"m":113,"g":115},"11096":{"m":113,"g":115},"11011":{"m":113,"g":115},"11178":{"m":113,"g":115},"11196":{"m":113,"g":115},"10741":{"m":113,"g":115},"11198":{"m":113,"g":115},"10517":{"m":113,"g":115},"10838":{"m":113,"g":115},"10859":{"m":113,"g":115},"10609":{"m":113,"g":115},"10855":{"m":113,"g":115},"11090":{"m":113,"g":115},"10780":{"m":113,"g":115},"10892":{"m":113,"g":115},"11166":{"m":113,"g":115},"11192":{"m":113,"g":115},"11189":{"m":113,"g":115},"10873":{"m":113,"g":115},"11173":{"m":113,"g":115},"11167":{"m":113,"g":115},"10637":{"m":113,"g":115},"11185":{"m":113,"g":115},"11161":{"m":113,"g":115},"9537":{"m":113,"g":115},"10830":{"m":113,"g":115},"10133":{"m":113,"g":115},"11138":{"m":113,"g":115},"11179":{"m":113,"g":115},"11176":{"m":113,"g":115},"11159":{"m":113,"g":115},"11170":{"m":113,"g":115},"11175":{"m":113,"g":115},"11124":{"m":113,"g":115},"11171":{"m":113,"g":115},"11164":{"m":113,"g":115},"11130":{"m":113,"g":115},"11163":{"m":113,"g":115},"10837":{"m":113,"g":115},"11152":{"m":113,"g":115},"10988":{"m":113,"g":115},"11160":{"m":113,"g":115},"10422":{"m":113,"g":115},"10263":{"m":113,"g":115},"11156":{"m":113,"g":115},"10508":{"m":113,"g":115},"10779":{"m":113,"g":115},"10768":{"m":113,"g":115},"11132":{"m":113,"g":115},"11148":{"m":113,"g":115},"11149":{"m":113,"g":115},"10559":{"m":113,"g":115},"11135":{"m":113,"g":115},"10720":{"m":113,"g":115},"11145":{"m":113,"g":115},"11123":{"m":113,"g":115},"11143":{"m":113,"g":115},"11120":{"m":113,"g":115},"10512":{"m":113,"g":115},"10271":{"m":113,"g":115},"11005":{"m":113,"g":115},"10760":{"m":113,"g":115},"11128":{"m":113,"g":115},"11075":{"m":113,"g":115},"10985":{"m":113,"g":115},"11111":{"m":113,"g":115},"11115":{"m":113,"g":115},"11114":{"m":113,"g":115},"10735":{"m":113,"g":115},"11112":{"m":113,"g":115},"10972":{"m":113,"g":115},"11080":{"m":113,"g":115},"11113":{"m":113,"g":115},"11071":{"m":113,"g":115},"11081":{"m":113,"g":115},"11102":{"m":113,"g":115},"11101":{"m":113,"g":115},"10846":{"m":113,"g":115},"11094":{"m":113,"g":115},"11067":{"m":113,"g":115},"11099":{"m":113,"g":115},"11087":{"m":113,"g":115},"10991":{"m":113,"g":115},"11085":{"m":113,"g":115},"11070":{"m":113,"g":115},"11092":{"m":113,"g":115},"10729":{"m":113,"g":115},"9642":{"m":113,"g":115},"10816":{"m":113,"g":115},"11083":{"m":113,"g":115},"10875":{"m":113,"g":115},"11082":{"m":113,"g":115},"11079":{"m":113,"g":115},"10611":{"m":113,"g":115},"11076":{"m":113,"g":115},"11073":{"m":113,"g":115},"10975":{"m":113,"g":115},"11069":{"m":113,"g":115},"11056":{"m":113,"g":115},"10976":{"m":113,"g":115},"11054":{"m":113,"g":115},"11050":{"m":113,"g":115},"11022":{"m":113,"g":115},"9614":{"m":113,"g":115},"11010":{"m":113,"g":115},"10591":{"m":113,"g":115},"11015":{"m":113,"g":115},"10940":{"m":113,"g":115},"10701":{"m":113,"g":115},"11036":{"m":113,"g":115},"11038":{"m":113,"g":115},"10986":{"m":113,"g":115},"11033":{"m":113,"g":115},"11017":{"m":113,"g":115},"11003":{"m":113,"g":115},"11013":{"m":113,"g":115},"10543":{"m":113,"g":115},"10555":{"m":113,"g":115},"10964":{"m":113,"g":115},"11009":{"m":113,"g":115},"10999":{"m":113,"g":115},"10978":{"m":113,"g":115},"10995":{"m":113,"g":115},"10997":{"m":113,"g":115},"10550":{"m":113,"g":115},"10565":{"m":113,"g":115},"10616":{"m":113,"g":115},"10751":{"m":113,"g":115},"10930":{"m":113,"g":115},"10981":{"m":113,"g":115},"10112":{"m":113,"g":115},"10980":{"m":113,"g":115},"10982":{"m":113,"g":115},"10551":{"m":113,"g":115},"10965":{"m":113,"g":115},"10944":{"m":113,"g":115},"10971":{"m":113,"g":115},"10941":{"m":113,"g":115},"10495":{"m":113,"g":115},"10372":{"m":113,"g":115},"10970":{"m":113,"g":115},"10947":{"m":113,"g":115},"10968":{"m":113,"g":115},"10967":{"m":113,"g":115},"10963":{"m":113,"g":115},"10960":{"m":113,"g":115},"10958":{"m":113,"g":115},"10956":{"m":113,"g":115},"10955":{"m":113,"g":115},"10749":{"m":113,"g":115},"10927":{"m":113,"g":115},"10936":{"m":113,"g":115},"10192":{"m":113,"g":115},"10939":{"m":113,"g":115},"10935":{"m":113,"g":115},"10929":{"m":113,"g":115},"10932":{"m":113,"g":115},"10898":{"m":113,"g":115},"10612":{"m":113,"g":115},"10923":{"m":113,"g":115},"10926":{"m":113,"g":115},"10899":{"m":113,"g":115},"10883":{"m":113,"g":115},"10910":{"m":113,"g":115},"10924":{"m":113,"g":115},"10132":{"m":113,"g":115},"10881":{"m":113,"g":115},"10894":{"m":113,"g":115},"10915":{"m":113,"g":115},"10376":{"m":113,"g":115},"10778":{"m":113,"g":115},"10872":{"m":113,"g":115},"10885":{"m":113,"g":115},"10895":{"m":113,"g":115},"10845":{"m":113,"g":115},"10880":{"m":113,"g":115},"10572":{"m":113,"g":115},"10861":{"m":113,"g":115},"10876":{"m":113,"g":115},"10877":{"m":113,"g":115},"10534":{"m":113,"g":115},"10827":{"m":113,"g":115},"10832":{"m":113,"g":115},"10786":{"m":113,"g":115},"10860":{"m":113,"g":115},"10718":{"m":113,"g":115},"10829":{"m":113,"g":115},"10828":{"m":113,"g":115},"10825":{"m":113,"g":115},"10826":{"m":113,"g":115},"10822":{"m":113,"g":115},"10824":{"m":113,"g":115},"10823":{"m":113,"g":115},"10787":{"m":113,"g":115},"10820":{"m":113,"g":115},"10794":{"m":113,"g":115},"10799":{"m":113,"g":115},"10818":{"m":113,"g":115},"10540":{"m":113,"g":115},"10504":{"m":113,"g":115},"10761":{"m":113,"g":115},"10814":{"m":113,"g":115},"10812":{"m":113,"g":115},"10259":{"m":113,"g":115},"10323":{"m":113,"g":115},"10581":{"m":113,"g":115},"10773":{"m":113,"g":115},"10792":{"m":113,"g":115},"10791":{"m":113,"g":115},"10770":{"m":113,"g":115},"10783":{"m":113,"g":115},"10782":{"m":113,"g":115},"10715":{"m":113,"g":115},"10705":{"m":113,"g":115},"10776":{"m":113,"g":115},"10777":{"m":113,"g":115},"10771":{"m":113,"g":115},"10774":{"m":113,"g":115},"10767":{"m":113,"g":115},"10756":{"m":113,"g":115},"10765":{"m":113,"g":115},"10574":{"m":113,"g":115},"10556":{"m":113,"g":115},"10130":{"m":113,"g":115},"10762":{"m":113,"g":115},"10541":{"m":113,"g":115},"10281":{"m":113,"g":115},"10759":{"m":113,"g":115},"10300":{"m":113,"g":115},"10755":{"m":113,"g":115},"10732":{"m":113,"g":115},"10758":{"m":113,"g":115},"10757":{"m":113,"g":115},"10727":{"m":113,"g":115},"10754":{"m":113,"g":115},"10724":{"m":113,"g":115},"10753":{"m":113,"g":115},"10737":{"m":113,"g":115},"10728":{"m":113,"g":115},"10730":{"m":113,"g":115},"9849":{"m":113,"g":115},"10731":{"m":113,"g":115},"10709":{"m":113,"g":115},"10699":{"m":113,"g":115},"10678":{"m":113,"g":115},"10695":{"m":113,"g":115},"10694":{"m":113,"g":115},"10717":{"m":113,"g":115},"10714":{"m":113,"g":115},"10716":{"m":113,"g":115},"10385":{"m":113,"g":115},"10317":{"m":113,"g":115},"10706":{"m":113,"g":115},"10592":{"m":113,"g":115},"10697":{"m":113,"g":115},"10696":{"m":113,"g":115},"10651":{"m":113,"g":115},"10688":{"m":113,"g":115},"10686":{"m":113,"g":115},"10685":{"m":113,"g":115},"10673":{"m":113,"g":115},"10684":{"m":113,"g":115},"10680":{"m":113,"g":115},"10681":{"m":113,"g":115},"10683":{"m":113,"g":115},"10645":{"m":113,"g":115},"10677":{"m":113,"g":115},"10679":{"m":113,"g":115},"10648":{"m":113,"g":115},"10671":{"m":113,"g":115},"10675":{"m":113,"g":115},"10666":{"m":113,"g":115},"10670":{"m":113,"g":115},"10668":{"m":113,"g":115},"10664":{"m":113,"g":115},"10661":{"m":113,"g":115},"10522":{"m":113,"g":115},"10653":{"m":113,"g":115},"10634":{"m":113,"g":115},"10650":{"m":113,"g":115},"10321":{"m":113,"g":115},"10647":{"m":113,"g":115},"10081":{"m":113,"g":115},"10633":{"m":113,"g":115},"10319":{"m":113,"g":115},"10630":{"m":113,"g":115},"10631":{"m":113,"g":115},"10632":{"m":113,"g":115},"10586":{"m":113,"g":115},"10553":{"m":113,"g":115},"9873":{"m":113,"g":115},"10629":{"m":113,"g":115},"10628":{"m":113,"g":115},"10621":{"m":113,"g":115},"9947":{"m":113,"g":115},"10579":{"m":113,"g":115},"10595":{"m":113,"g":115},"10610":{"m":113,"g":115},"10622":{"m":113,"g":115},"10624":{"m":113,"g":115},"10222":{"m":113,"g":115},"9979":{"m":113,"g":115},"8274":{"m":113,"g":115},"10604":{"m":113,"g":115},"10525":{"m":113,"g":115},"10596":{"m":113,"g":115},"10273":{"m":113,"g":115},"10563":{"m":113,"g":115},"10190":{"m":113,"g":115},"10558":{"m":113,"g":115},"10526":{"m":113,"g":115},"9987":{"m":113,"g":115},"10584":{"m":113,"g":115},"9976":{"m":113,"g":115},"10171":{"m":113,"g":115},"10548":{"m":113,"g":115},"8813":{"m":113,"g":115},"10545":{"m":113,"g":115},"10523":{"m":113,"g":115},"10459":{"m":113,"g":115},"10529":{"m":113,"g":115},"8746":{"m":113,"g":115},"10538":{"m":113,"g":115},"10494":{"m":113,"g":115},"9928":{"m":113,"g":115},"10474":{"m":113,"g":115},"10530":{"m":113,"g":115},"10528":{"m":113,"g":115},"10524":{"m":113,"g":115},"10511":{"m":113,"g":115},"10506":{"m":113,"g":115},"10515":{"m":113,"g":115},"10491":{"m":113,"g":115},"10500":{"m":113,"g":115},"10507":{"m":113,"g":115},"10466":{"m":113,"g":115},"10498":{"m":113,"g":115},"10499":{"m":113,"g":115},"10493":{"m":113,"g":115},"10487":{"m":113,"g":115},"10230":{"m":113,"g":115},"10336":{"m":113,"g":115},"10203":{"m":113,"g":115},"10434":{"m":113,"g":115},"8863":{"m":113,"g":115},"10486":{"m":113,"g":115},"10484":{"m":113,"g":115},"10286":{"m":113,"g":115},"10481":{"m":113,"g":115},"10478":{"m":113,"g":115},"10479":{"m":113,"g":115},"10475":{"m":113,"g":115},"10473":{"m":113,"g":115},"10476":{"m":113,"g":115},"8189":{"m":113,"g":115},"9657":{"m":113,"g":115},"10375":{"m":113,"g":115},"9887":{"m":113,"g":115},"8710":{"m":113,"g":115},"10471":{"m":113,"g":115},"10470":{"m":113,"g":115},"10468":{"m":113,"g":115},"10465":{"m":113,"g":115},"10440":{"m":113,"g":115},"10456":{"m":113,"g":115},"10439":{"m":113,"g":115},"10463":{"m":113,"g":115},"10458":{"m":113,"g":115},"10457":{"m":113,"g":115},"10401":{"m":113,"g":115},"10358":{"m":113,"g":115},"10449":{"m":113,"g":115},"9343":{"m":113,"g":115},"10452":{"m":113,"g":115},"10450":{"m":113,"g":115},"10445":{"m":113,"g":115},"9626":{"m":113,"g":115},"9768":{"m":113,"g":115},"10143":{"m":113,"g":115},"10441":{"m":113,"g":115},"10437":{"m":113,"g":115},"9338":{"m":113,"g":115},"10435":{"m":113,"g":115},"10201":{"m":113,"g":115},"10432":{"m":113,"g":115},"10129":{"m":113,"g":115},"10433":{"m":113,"g":115},"10426":{"m":113,"g":115},"10425":{"m":113,"g":115},"10429":{"m":113,"g":115},"10431":{"m":113,"g":115},"10428":{"m":113,"g":115},"9962":{"m":113,"g":115},"10076":{"m":113,"g":115},"8627":{"m":113,"g":115},"10419":{"m":113,"g":115},"6539":{"m":113,"g":115},"10270":{"m":113,"g":115},"9948":{"m":113,"g":115},"10157":{"m":113,"g":115},"10313":{"m":113,"g":115},"10369":{"m":113,"g":115},"10318":{"m":113,"g":115},"10404":{"m":113,"g":115},"10228":{"m":113,"g":115},"10414":{"m":113,"g":115},"10410":{"m":113,"g":115},"10411":{"m":113,"g":115},"10412":{"m":113,"g":115},"9748":{"m":113,"g":115},"9382":{"m":113,"g":115},"10406":{"m":113,"g":115},"10392":{"m":113,"g":115},"10403":{"m":113,"g":115},"10400":{"m":113,"g":115},"10397":{"m":113,"g":115},"10398":{"m":113,"g":115},"9984":{"m":113,"g":115},"10395":{"m":113,"g":115},"10394":{"m":113,"g":115},"10332":{"m":113,"g":115},"10377":{"m":113,"g":115},"10379":{"m":113,"g":115},"10380":{"m":113,"g":115},"10387":{"m":113,"g":115},"10244":{"m":113,"g":115},"10386":{"m":113,"g":115},"10391":{"m":113,"g":115},"10361":{"m":113,"g":115},"10390":{"m":113,"g":115},"10388":{"m":113,"g":115},"10343":{"m":113,"g":115},"10333":{"m":113,"g":115},"9023":{"m":113,"g":115},"10099":{"m":113,"g":115},"10219":{"m":113,"g":115},"10370":{"m":113,"g":115},"10368":{"m":113,"g":115},"10180":{"m":113,"g":115},"10355":{"m":113,"g":115},"8215":{"m":113,"g":115},"8778":{"m":113,"g":115},"10351":{"m":113,"g":115},"10362":{"m":113,"g":115},"10359":{"m":113,"g":115},"10360":{"m":113,"g":115},"10356":{"m":113,"g":115},"10031":{"m":113,"g":115},"10283":{"m":113,"g":115},"10346":{"m":113,"g":115},"10352":{"m":113,"g":115},"9774":{"m":113,"g":115},"10296":{"m":113,"g":115},"9199":{"m":113,"g":115},"10345":{"m":113,"g":115},"10349":{"m":113,"g":115},"10347":{"m":113,"g":115},"11324":{"m":114,"g":115},"11369":{"m":114,"g":115},"11364":{"m":114,"g":115},"11394":{"m":114,"g":115},"11387":{"m":114,"g":115},"11376":{"m":114,"g":115},"11375":{"m":114,"g":115},"11373":{"m":114,"g":115},"11309":{"m":114,"g":115},"11366":{"m":114,"g":115},"11359":{"m":114,"g":115},"11353":{"m":114,"g":115},"10979":{"m":114,"g":115},"11350":{"m":114,"g":115},"11327":{"m":114,"g":115},"11342":{"m":114,"g":115},"11339":{"m":114,"g":115},"11341":{"m":114,"g":115},"11340":{"m":114,"g":115},"11336":{"m":114,"g":115},"11323":{"m":114,"g":115},"10909":{"m":114,"g":115},"11264":{"m":114,"g":115},"9812":{"m":114,"g":115},"11318":{"m":114,"g":115},"11007":{"m":114,"g":115},"11321":{"m":114,"g":115},"10937":{"m":114,"g":115},"11312":{"m":114,"g":115},"11211":{"m":114,"g":115},"9545":{"m":114,"g":115},"11314":{"m":114,"g":115},"11316":{"m":114,"g":115},"11126":{"m":114,"g":115},"11315":{"m":114,"g":115},"10710":{"m":114,"g":115},"11230":{"m":114,"g":115},"11200":{"m":114,"g":115},"11304":{"m":114,"g":115},"11310":{"m":114,"g":115},"11311":{"m":114,"g":115},"11297":{"m":114,"g":115},"11205":{"m":114,"g":115},"11307":{"m":114,"g":115},"11223":{"m":114,"g":115},"11306":{"m":114,"g":115},"11305":{"m":114,"g":115},"11001":{"m":114,"g":115},"11027":{"m":114,"g":115},"11288":{"m":114,"g":115},"11303":{"m":114,"g":115},"11302":{"m":114,"g":115},"11068":{"m":114,"g":115},"11301":{"m":114,"g":115},"11300":{"m":114,"g":115},"11290":{"m":114,"g":115},"11210":{"m":114,"g":115},"11231":{"m":114,"g":115},"11294":{"m":114,"g":115},"10949":{"m":114,"g":115},"11095":{"m":114,"g":115},"11283":{"m":114,"g":115},"11281":{"m":114,"g":115},"11286":{"m":114,"g":115},"11282":{"m":114,"g":115},"11261":{"m":114,"g":115},"11238":{"m":114,"g":115},"11279":{"m":114,"g":115},"11280":{"m":114,"g":115},"11276":{"m":114,"g":115},"11277":{"m":114,"g":115},"11182":{"m":114,"g":115},"11268":{"m":114,"g":115},"11274":{"m":114,"g":115},"11270":{"m":114,"g":115},"7149":{"m":114,"g":115},"11262":{"m":114,"g":115},"11219":{"m":114,"g":115},"11680":{"m":116,"g":118},"11676":{"m":116,"g":118},"11684":{"m":116,"g":118},"11667":{"m":116,"g":118},"11674":{"m":116,"g":118},"11681":{"m":116,"g":118},"11621":{"m":116,"g":118},"11367":{"m":116,"g":118},"11653":{"m":116,"g":118},"11660":{"m":116,"g":118},"11659":{"m":116,"g":118},"11585":{"m":116,"g":118},"11293":{"m":116,"g":118},"11458":{"m":116,"g":118},"11590":{"m":116,"g":118},"11636":{"m":116,"g":118},"11579":{"m":116,"g":118},"8247":{"m":116,"g":118},"10423":{"m":116,"g":118},"11642":{"m":116,"g":115},"11638":{"m":116,"g":115},"11628":{"m":116,"g":115},"11639":{"m":116,"g":115},"11351":{"m":116,"g":115},"11627":{"m":116,"g":115},"11633":{"m":116,"g":115},"11631":{"m":116,"g":115},"11622":{"m":116,"g":115},"11625":{"m":116,"g":115},"11623":{"m":116,"g":115},"11624":{"m":116,"g":115},"11619":{"m":116,"g":115},"11605":{"m":116,"g":115},"11620":{"m":116,"g":115},"11617":{"m":116,"g":115},"11453":{"m":116,"g":115},"11561":{"m":116,"g":115},"11434":{"m":116,"g":115},"10721":{"m":116,"g":115},"11586":{"m":116,"g":115},"11556":{"m":116,"g":115},"11603":{"m":116,"g":115},"11593":{"m":116,"g":115},"11601":{"m":116,"g":115},"11449":{"m":116,"g":115},"11600":{"m":116,"g":115},"11597":{"m":116,"g":115},"11598":{"m":116,"g":115},"11566":{"m":116,"g":115},"11591":{"m":116,"g":115},"11588":{"m":116,"g":115},"11587":{"m":116,"g":115},"11580":{"m":116,"g":115},"11583":{"m":116,"g":115},"11582":{"m":116,"g":115},"11041":{"m":116,"g":115},"11542":{"m":116,"g":115},"11535":{"m":116,"g":115},"11565":{"m":116,"g":115},"11413":{"m":116,"g":115},"11539":{"m":116,"g":115},"11572":{"m":116,"g":115},"11573":{"m":116,"g":115},"11534":{"m":116,"g":115},"11571":{"m":116,"g":115},"11564":{"m":116,"g":115},"11537":{"m":116,"g":115},"11538":{"m":116,"g":115},"11521":{"m":116,"g":115},"11562":{"m":116,"g":115},"11308":{"m":116,"g":115},"11557":{"m":116,"g":115},"11441":{"m":116,"g":115},"11483":{"m":116,"g":115},"11531":{"m":116,"g":115},"11549":{"m":116,"g":115},"11553":{"m":116,"g":115},"11547":{"m":116,"g":115},"11507":{"m":116,"g":115},"11419":{"m":116,"g":115},"11444":{"m":116,"g":115},"11442":{"m":116,"g":115},"11548":{"m":116,"g":115},"11530":{"m":116,"g":115},"11527":{"m":116,"g":115},"11201":{"m":116,"g":115},"11528":{"m":116,"g":115},"11457":{"m":116,"g":115},"11544":{"m":116,"g":115},"11460":{"m":116,"g":115},"11505":{"m":116,"g":115},"11385":{"m":116,"g":115},"11214":{"m":116,"g":115},"11493":{"m":116,"g":115},"11512":{"m":116,"g":115},"11432":{"m":116,"g":115},"11485":{"m":116,"g":115},"11511":{"m":116,"g":115},"5889":{"m":116,"g":115},"11520":{"m":116,"g":115},"11516":{"m":116,"g":115},"11498":{"m":116,"g":115},"11474":{"m":116,"g":115},"11515":{"m":116,"g":115},"11514":{"m":116,"g":115},"11331":{"m":116,"g":115},"11509":{"m":116,"g":115},"11443":{"m":116,"g":115},"11452":{"m":116,"g":115},"11503":{"m":116,"g":115},"11502":{"m":116,"g":115},"11497":{"m":116,"g":115},"11501":{"m":116,"g":115},"11500":{"m":116,"g":115},"11332":{"m":116,"g":115},"11499":{"m":116,"g":115},"11465":{"m":116,"g":115},"10577":{"m":116,"g":115},"11479":{"m":116,"g":115},"11221":{"m":116,"g":115},"10172":{"m":116,"g":115},"11478":{"m":116,"g":115},"11481":{"m":116,"g":115},"11476":{"m":116,"g":115},"11489":{"m":116,"g":115},"10062":{"m":116,"g":115},"11398":{"m":116,"g":115},"10635":{"m":116,"g":115},"9804":{"m":116,"g":115},"11454":{"m":116,"g":115},"11019":{"m":116,"g":115},"8919":{"m":116,"g":115},"11462":{"m":116,"g":115},"11428":{"m":116,"g":115},"11470":{"m":116,"g":115},"11427":{"m":116,"g":115},"11467":{"m":116,"g":115},"11448":{"m":116,"g":115},"10312":{"m":116,"g":115},"9991":{"m":116,"g":115},"11455":{"m":116,"g":115},"11450":{"m":116,"g":115},"11360":{"m":116,"g":115},"11368":{"m":116,"g":115},"11445":{"m":116,"g":115},"11438":{"m":116,"g":115},"11439":{"m":116,"g":115},"11399":{"m":116,"g":115},"11435":{"m":116,"g":115},"11313":{"m":116,"g":115},"10745":{"m":116,"g":115},"11411":{"m":116,"g":115},"11433":{"m":116,"g":115},"11437":{"m":116,"g":115},"11436":{"m":116,"g":115},"11345":{"m":116,"g":115},"9256":{"m":116,"g":115},"11381":{"m":116,"g":115},"11361":{"m":116,"g":115},"11420":{"m":116,"g":115},"11144":{"m":116,"g":115},"10734":{"m":116,"g":115},"10969":{"m":116,"g":115},"9045":{"m":116,"g":115},"11414":{"m":116,"g":115},"11388":{"m":116,"g":115},"11363":{"m":116,"g":115},"11365":{"m":116,"g":115},"11389":{"m":116,"g":115},"11401":{"m":116,"g":115},"11285":{"m":116,"g":115},"11693":{"m":117,"g":118},"11543":{"m":117,"g":118},"11687":{"m":117,"g":118},"11706":{"m":117,"g":118},"11510":{"m":117,"g":118},"11488":{"m":117,"g":118},"11370":{"m":117,"g":118},"11692":{"m":117,"g":118},"11663":{"m":117,"g":118},"10912":{"m":117,"g":118},"11679":{"m":117,"g":118},"11689":{"m":117,"g":118},"11686":{"m":117,"g":118},"10248":{"m":117,"g":118},"9493":{"m":117,"g":118},"12027":{"m":119,"g":121},"12009":{"m":119,"g":121},"12030":{"m":119,"g":121},"12029":{"m":119,"g":121},"11616":{"m":119,"g":121},"12028":{"m":119,"g":121},"9366":{"m":119,"g":121},"11891":{"m":119,"g":121},"10158":{"m":119,"g":121},"12024":{"m":119,"g":121},"11765":{"m":119,"g":121},"12022":{"m":119,"g":121},"12018":{"m":119,"g":121},"12021":{"m":119,"g":121},"11755":{"m":119,"g":121},"11981":{"m":119,"g":121},"12015":{"m":119,"g":121},"12014":{"m":119,"g":121},"11937":{"m":119,"g":121},"11988":{"m":119,"g":121},"11866":{"m":119,"g":121},"11821":{"m":119,"g":121},"12004":{"m":119,"g":121},"11985":{"m":119,"g":121},"11944":{"m":119,"g":121},"10652":{"m":119,"g":121},"11965":{"m":119,"g":121},"11906":{"m":119,"g":121},"11990":{"m":119,"g":121},"11299":{"m":119,"g":121},"11322":{"m":119,"g":121},"11811":{"m":119,"g":121},"11955":{"m":119,"g":121},"11978":{"m":119,"g":121},"10869":{"m":119,"g":121},"11921":{"m":119,"g":121},"11563":{"m":119,"g":121},"11980":{"m":119,"g":121},"11977":{"m":119,"g":121},"10750":{"m":119,"g":121},"11956":{"m":119,"g":121},"11723":{"m":119,"g":121},"11967":{"m":119,"g":121},"11953":{"m":119,"g":121},"9651":{"m":119,"g":121},"10606":{"m":119,"g":121},"11908":{"m":119,"g":121},"10154":{"m":119,"g":121},"11929":{"m":119,"g":121},"11717":{"m":119,"g":121},"11922":{"m":119,"g":121},"11945":{"m":119,"g":121},"11790":{"m":119,"g":121},"11926":{"m":119,"g":121},"11940":{"m":119,"g":121},"11935":{"m":119,"g":121},"11934":{"m":119,"g":121},"11377":{"m":119,"g":121},"11933":{"m":119,"g":121},"11876":{"m":119,"g":121},"11844":{"m":119,"g":121},"11287":{"m":119,"g":121},"11918":{"m":119,"g":121},"11915":{"m":119,"g":121},"11702":{"m":119,"g":121},"11482":{"m":119,"g":121},"10700":{"m":119,"g":121},"11902":{"m":119,"g":121},"11295":{"m":119,"g":121},"11416":{"m":119,"g":121},"11895":{"m":119,"g":121},"11570":{"m":119,"g":121},"11487":{"m":119,"g":121},"11878":{"m":119,"g":121},"11885":{"m":119,"g":118},"11664":{"m":119,"g":118},"10656":{"m":119,"g":118},"11843":{"m":119,"g":118},"11845":{"m":119,"g":118},"11859":{"m":119,"g":118},"11887":{"m":119,"g":118},"11838":{"m":119,"g":118},"11886":{"m":119,"g":118},"11826":{"m":119,"g":118},"11882":{"m":119,"g":118},"11875":{"m":119,"g":118},"11881":{"m":119,"g":118},"11868":{"m":119,"g":118},"11807":{"m":119,"g":118},"11867":{"m":119,"g":118},"11823":{"m":119,"g":118},"11847":{"m":119,"g":118},"11862":{"m":119,"g":118},"10691":{"m":119,"g":118},"11776":{"m":119,"g":118},"11822":{"m":119,"g":118},"11849":{"m":119,"g":118},"11396":{"m":119,"g":118},"11846":{"m":119,"g":118},"11747":{"m":119,"g":118},"10801":{"m":119,"g":118},"11722":{"m":119,"g":118},"11594":{"m":119,"g":118},"11780":{"m":119,"g":118},"11733":{"m":119,"g":118},"10510":{"m":119,"g":118},"11508":{"m":119,"g":118},"11787":{"m":119,"g":118},"11832":{"m":119,"g":118},"11778":{"m":119,"g":118},"11612":{"m":119,"g":118},"11810":{"m":119,"g":118},"11831":{"m":119,"g":118},"11815":{"m":119,"g":118},"11606":{"m":119,"g":118},"11835":{"m":119,"g":118},"10994":{"m":119,"g":118},"11147":{"m":119,"g":118},"11833":{"m":119,"g":118},"11834":{"m":119,"g":118},"11652":{"m":119,"g":118},"11819":{"m":119,"g":118},"11827":{"m":119,"g":118},"11805":{"m":119,"g":118},"5162":{"m":119,"g":118},"11786":{"m":119,"g":118},"11804":{"m":119,"g":118},"11808":{"m":119,"g":118},"11091":{"m":119,"g":118},"11818":{"m":119,"g":118},"10788":{"m":119,"g":118},"11817":{"m":119,"g":118},"11328":{"m":119,"g":118},"11555":{"m":119,"g":118},"11618":{"m":119,"g":118},"11670":{"m":119,"g":118},"11688":{"m":119,"g":118},"11773":{"m":119,"g":118},"11813":{"m":119,"g":118},"11000":{"m":119,"g":118},"11710":{"m":119,"g":118},"11749":{"m":119,"g":118},"11772":{"m":119,"g":118},"11506":{"m":119,"g":118},"11797":{"m":119,"g":118},"11781":{"m":119,"g":118},"11803":{"m":119,"g":118},"11801":{"m":119,"g":118},"10152":{"m":119,"g":118},"11669":{"m":119,"g":118},"11793":{"m":119,"g":118},"11665":{"m":119,"g":118},"11799":{"m":119,"g":118},"11798":{"m":119,"g":118},"11794":{"m":119,"g":118},"11784":{"m":119,"g":118},"11783":{"m":119,"g":118},"11614":{"m":119,"g":118},"11788":{"m":119,"g":118},"9170":{"m":119,"g":118},"11666":{"m":119,"g":118},"11613":{"m":119,"g":118},"11611":{"m":119,"g":118},"11607":{"m":119,"g":118},"11685":{"m":119,"g":118},"11519":{"m":119,"g":118},"11782":{"m":119,"g":118},"11777":{"m":119,"g":118},"11682":{"m":119,"g":118},"11775":{"m":119,"g":118},"11738":{"m":119,"g":118},"10725":{"m":119,"g":118},"11540":{"m":119,"g":118},"11767":{"m":119,"g":118},"11768":{"m":119,"g":118},"11766":{"m":119,"g":118},"11735":{"m":119,"g":118},"11730":{"m":119,"g":118},"11643":{"m":119,"g":118},"11062":{"m":119,"g":118},"11724":{"m":119,"g":118},"11746":{"m":119,"g":118},"11739":{"m":119,"g":118},"11740":{"m":119,"g":118},"11732":{"m":119,"g":118},"11734":{"m":119,"g":118},"11541":{"m":119,"g":118},"11728":{"m":119,"g":118},"11731":{"m":119,"g":118},"11727":{"m":119,"g":118},"11729":{"m":119,"g":118},"11677":{"m":119,"g":118},"10911":{"m":119,"g":118},"12169":{"m":120,"g":121},"12177":{"m":120,"g":121},"12170":{"m":120,"g":121},"12167":{"m":120,"g":121},"12171":{"m":120,"g":121},"12164":{"m":120,"g":121},"12168":{"m":120,"g":121},"12129":{"m":120,"g":121},"12166":{"m":120,"g":121},"11047":{"m":120,"g":121},"11632":{"m":120,"g":121},"10399":{"m":120,"g":121},"12142":{"m":120,"g":121},"12106":{"m":120,"g":121},"12156":{"m":120,"g":121},"12155":{"m":120,"g":121},"12152":{"m":120,"g":121},"12154":{"m":120,"g":121},"11615":{"m":120,"g":121},"11494":{"m":120,"g":121},"12113":{"m":120,"g":121},"12116":{"m":120,"g":121},"12136":{"m":120,"g":121},"12147":{"m":120,"g":121},"12141":{"m":120,"g":121},"11991":{"m":120,"g":121},"12097":{"m":120,"g":121},"12139":{"m":120,"g":121},"12138":{"m":120,"g":121},"12118":{"m":120,"g":121},"12133":{"m":120,"g":121},"11936":{"m":120,"g":121},"12132":{"m":120,"g":121},"11993":{"m":120,"g":121},"12130":{"m":120,"g":121},"12125":{"m":120,"g":121},"12101":{"m":120,"g":121},"11814":{"m":120,"g":121},"12127":{"m":120,"g":121},"12115":{"m":120,"g":121},"12126":{"m":120,"g":121},"12098":{"m":120,"g":121},"11962":{"m":120,"g":121},"12119":{"m":120,"g":121},"12124":{"m":120,"g":121},"11869":{"m":120,"g":121},"12110":{"m":120,"g":121},"12096":{"m":120,"g":121},"12083":{"m":120,"g":121},"12103":{"m":120,"g":121},"12105":{"m":120,"g":121},"12087":{"m":120,"g":121},"9501":{"m":120,"g":121},"11379":{"m":120,"g":121},"12058":{"m":120,"g":121},"12054":{"m":120,"g":121},"11877":{"m":120,"g":121},"12070":{"m":120,"g":121},"8464":{"m":120,"g":121},"12093":{"m":120,"g":121},"12071":{"m":120,"g":121},"12000":{"m":120,"g":121},"12091":{"m":120,"g":121},"12089":{"m":120,"g":121},"12086":{"m":120,"g":121},"11560":{"m":120,"g":121},"12084":{"m":120,"g":121},"12034":{"m":120,"g":121},"11924":{"m":120,"g":121},"11884":{"m":120,"g":121},"12025":{"m":120,"g":121},"11958":{"m":120,"g":121},"12067":{"m":120,"g":121},"11999":{"m":120,"g":121},"12046":{"m":120,"g":121},"12049":{"m":120,"g":121},"12063":{"m":120,"g":121},"12064":{"m":120,"g":121},"12053":{"m":120,"g":121},"11800":{"m":120,"g":121},"12056":{"m":120,"g":121},"12019":{"m":120,"g":121},"11853":{"m":120,"g":121},"12031":{"m":120,"g":121},"12041":{"m":120,"g":121},"10953":{"m":120,"g":121},"11759":{"m":120,"g":121},"12042":{"m":120,"g":121},"12037":{"m":120,"g":121},"11745":{"m":120,"g":121},"12038":{"m":120,"g":121},"11909":{"m":120,"g":121},"11816":{"m":120,"g":121},"11795":{"m":120,"g":121},"12003":{"m":120,"g":121},"9936":{"m":120,"g":121},"11964":{"m":120,"g":121},"12439":{"m":122,"g":127},"11874":{"m":122,"g":127},"12475":{"m":122,"g":127},"12469":{"m":122,"g":127},"11987":{"m":122,"g":127},"12430":{"m":122,"g":127},"12473":{"m":122,"g":127},"12297":{"m":122,"g":127},"12341":{"m":122,"g":127},"12428":{"m":122,"g":127},"12429":{"m":122,"g":127},"12066":{"m":122,"g":127},"12275":{"m":122,"g":127},"11757":{"m":122,"g":127},"11931":{"m":122,"g":127},"12470":{"m":122,"g":127},"12415":{"m":122,"g":127},"12266":{"m":122,"g":127},"12463":{"m":122,"g":127},"12328":{"m":122,"g":127},"12369":{"m":122,"g":127},"12256":{"m":122,"g":127},"12449":{"m":122,"g":127},"12452":{"m":122,"g":127},"12413":{"m":122,"g":127},"10889":{"m":122,"g":127},"12422":{"m":122,"g":127},"12436":{"m":122,"g":127},"12437":{"m":122,"g":127},"12410":{"m":122,"g":127},"12384":{"m":122,"g":127},"12307":{"m":122,"g":127},"10566":{"m":122,"g":127},"12405":{"m":122,"g":127},"12425":{"m":122,"g":127},"12401":{"m":122,"g":127},"12300":{"m":122,"g":127},"11224":{"m":122,"g":127},"11116":{"m":122,"g":127},"12399":{"m":122,"g":121},"12290":{"m":122,"g":121},"12242":{"m":122,"g":121},"12368":{"m":122,"g":121},"12403":{"m":122,"g":121},"12386":{"m":122,"g":121},"12409":{"m":122,"g":121},"12375":{"m":122,"g":121},"12404":{"m":122,"g":121},"12281":{"m":122,"g":121},"12185":{"m":122,"g":121},"12012":{"m":122,"g":121},"12364":{"m":122,"g":121},"12395":{"m":122,"g":121},"11960":{"m":122,"g":121},"12377":{"m":122,"g":121},"11897":{"m":122,"g":121},"11969":{"m":122,"g":121},"12394":{"m":122,"g":121},"11806":{"m":122,"g":121},"12319":{"m":122,"g":121},"12123":{"m":122,"g":121},"12358":{"m":122,"g":121},"12362":{"m":122,"g":121},"12135":{"m":122,"g":121},"12378":{"m":122,"g":121},"12174":{"m":122,"g":121},"12340":{"m":122,"g":121},"12195":{"m":122,"g":121},"11910":{"m":122,"g":121},"12216":{"m":122,"g":121},"12050":{"m":122,"g":121},"12153":{"m":122,"g":121},"12094":{"m":122,"g":121},"12182":{"m":122,"g":121},"12354":{"m":122,"g":121},"12350":{"m":122,"g":121},"12348":{"m":122,"g":121},"11737":{"m":122,"g":121},"12346":{"m":122,"g":121},"12325":{"m":122,"g":121},"12095":{"m":122,"g":121},"12347":{"m":122,"g":121},"12345":{"m":122,"g":121},"11709":{"m":122,"g":121},"12343":{"m":122,"g":121},"12315":{"m":122,"g":121},"12338":{"m":122,"g":121},"11673":{"m":122,"g":121},"12002":{"m":122,"g":121},"12336":{"m":122,"g":121},"12312":{"m":122,"g":121},"12317":{"m":122,"g":121},"12276":{"m":122,"g":121},"12294":{"m":122,"g":121},"12314":{"m":122,"g":121},"12144":{"m":122,"g":121},"10874":{"m":122,"g":121},"12269":{"m":122,"g":121},"9825":{"m":122,"g":121},"12313":{"m":122,"g":121},"12311":{"m":122,"g":121},"12259":{"m":122,"g":121},"12241":{"m":122,"g":121},"12308":{"m":122,"g":121},"12299":{"m":122,"g":121},"12271":{"m":122,"g":121},"12296":{"m":122,"g":121},"12295":{"m":122,"g":121},"12285":{"m":122,"g":121},"12233":{"m":122,"g":121},"11928":{"m":122,"g":121},"12188":{"m":122,"g":121},"12283":{"m":122,"g":121},"12231":{"m":122,"g":121},"12284":{"m":122,"g":121},"12274":{"m":122,"g":121},"12257":{"m":122,"g":121},"12268":{"m":122,"g":121},"12267":{"m":122,"g":121},"12206":{"m":122,"g":121},"12247":{"m":122,"g":121},"10804":{"m":122,"g":121},"12230":{"m":122,"g":121},"12229":{"m":122,"g":121},"12252":{"m":122,"g":121},"12249":{"m":122,"g":121},"7873":{"m":122,"g":121},"10567":{"m":122,"g":121},"11177":{"m":122,"g":121},"11655":{"m":122,"g":121},"12245":{"m":122,"g":121},"11517":{"m":122,"g":121},"10654":{"m":122,"g":121},"12222":{"m":122,"g":121},"11994":{"m":122,"g":121},"12176":{"m":122,"g":121},"11708":{"m":122,"g":121},"12235":{"m":122,"g":121},"12161":{"m":122,"g":121},"12234":{"m":122,"g":121},"11142":{"m":122,"g":121},"12006":{"m":122,"g":121},"11592":{"m":122,"g":121},"11656":{"m":122,"g":121},"12186":{"m":122,"g":121},"12209":{"m":122,"g":121},"12205":{"m":122,"g":121},"12107":{"m":122,"g":121},"12112":{"m":122,"g":121},"10153":{"m":122,"g":121},"12117":{"m":122,"g":121},"12080":{"m":122,"g":121},"9403":{"m":122,"g":121},"12192":{"m":122,"g":121},"12173":{"m":122,"g":121},"12159":{"m":122,"g":121},"12057":{"m":122,"g":121},"12639":{"m":123,"g":127},"12572":{"m":123,"g":127},"12656":{"m":123,"g":127},"12456":{"m":123,"g":127},"12585":{"m":123,"g":127},"12648":{"m":123,"g":127},"12650":{"m":123,"g":127},"12640":{"m":123,"g":127},"12645":{"m":123,"g":127},"12647":{"m":123,"g":127},"12642":{"m":123,"g":127},"12641":{"m":123,"g":127},"12628":{"m":123,"g":127},"12634":{"m":123,"g":127},"12633":{"m":123,"g":127},"12632":{"m":123,"g":127},"12593":{"m":123,"g":127},"12616":{"m":123,"g":127},"12594":{"m":123,"g":127},"12592":{"m":123,"g":127},"11456":{"m":123,"g":127},"12615":{"m":123,"g":127},"12599":{"m":123,"g":127},"12522":{"m":123,"g":127},"10183":{"m":123,"g":127},"12598":{"m":123,"g":127},"6318":{"m":123,"g":127},"11131":{"m":123,"g":127},"11974":{"m":123,"g":127},"12580":{"m":123,"g":127},"12597":{"m":123,"g":127},"11760":{"m":123,"g":127},"12462":{"m":123,"g":127},"12165":{"m":123,"g":127},"12111":{"m":123,"g":127},"12547":{"m":123,"g":127},"12270":{"m":123,"g":127},"12044":{"m":123,"g":127},"12519":{"m":123,"g":127},"12571":{"m":123,"g":127},"12301":{"m":123,"g":127},"12569":{"m":123,"g":127},"12550":{"m":123,"g":127},"12549":{"m":123,"g":127},"12553":{"m":123,"g":127},"12524":{"m":123,"g":127},"12560":{"m":123,"g":127},"12227":{"m":123,"g":127},"12548":{"m":123,"g":127},"11330":{"m":123,"g":127},"12564":{"m":123,"g":127},"12561":{"m":123,"g":127},"12060":{"m":123,"g":127},"12536":{"m":123,"g":127},"12541":{"m":123,"g":127},"12532":{"m":123,"g":127},"12530":{"m":123,"g":127},"12523":{"m":123,"g":127},"12367":{"m":123,"g":127},"12502":{"m":123,"g":127},"12521":{"m":123,"g":127},"12515":{"m":123,"g":127},"12453":{"m":123,"g":127},"12481":{"m":123,"g":127},"12506":{"m":123,"g":127},"11917":{"m":123,"g":127},"12511":{"m":123,"g":127},"10078":{"m":123,"g":127},"12505":{"m":123,"g":127},"12507":{"m":123,"g":127},"11133":{"m":123,"g":127},"11052":{"m":123,"g":127},"12499":{"m":123,"g":127},"12391":{"m":123,"g":127},"12488":{"m":123,"g":127},"12412":{"m":123,"g":127},"11966":{"m":123,"g":127},"12238":{"m":123,"g":127},"12423":{"m":123,"g":127},"12500":{"m":123,"g":127},"12480":{"m":123,"g":127},"12485":{"m":123,"g":127},"12435":{"m":123,"g":127},"12483":{"m":123,"g":127},"12482":{"m":123,"g":127},"12226":{"m":123,"g":127},"12334":{"m":123,"g":127},"12739":{"m":124,"g":127},"12778":{"m":124,"g":127},"12440":{"m":124,"g":127},"12760":{"m":124,"g":127},"12565":{"m":124,"g":127},"12674":{"m":124,"g":127},"12646":{"m":124,"g":127},"12240":{"m":124,"g":127},"12508":{"m":124,"g":127},"12737":{"m":124,"g":127},"12693":{"m":124,"g":127},"12752":{"m":124,"g":127},"12744":{"m":124,"g":127},"12736":{"m":124,"g":127},"12748":{"m":124,"g":127},"12741":{"m":124,"g":127},"12742":{"m":124,"g":127},"12738":{"m":124,"g":127},"11892":{"m":124,"g":127},"12721":{"m":124,"g":127},"12734":{"m":124,"g":127},"12723":{"m":124,"g":127},"12716":{"m":124,"g":127},"12732":{"m":124,"g":127},"12611":{"m":124,"g":127},"12728":{"m":124,"g":127},"12729":{"m":124,"g":127},"12718":{"m":124,"g":127},"12651":{"m":124,"g":127},"12711":{"m":124,"g":127},"12713":{"m":124,"g":127},"12714":{"m":124,"g":127},"12712":{"m":124,"g":127},"12658":{"m":124,"g":127},"12673":{"m":124,"g":127},"12710":{"m":124,"g":127},"12709":{"m":124,"g":127},"12406":{"m":124,"g":127},"12586":{"m":124,"g":127},"12631":{"m":124,"g":127},"12699":{"m":124,"g":127},"12484":{"m":124,"g":127},"12677":{"m":124,"g":127},"12696":{"m":124,"g":127},"12708":{"m":124,"g":127},"12609":{"m":124,"g":127},"12702":{"m":124,"g":127},"12455":{"m":124,"g":127},"12691":{"m":124,"g":127},"12687":{"m":124,"g":127},"11641":{"m":124,"g":127},"12680":{"m":124,"g":127},"10044":{"m":124,"g":127},"12668":{"m":124,"g":127},"12175":{"m":124,"g":127},"12670":{"m":124,"g":127},"8784":{"m":124,"g":127},"12486":{"m":124,"g":127},"12638":{"m":124,"g":127},"12353":{"m":124,"g":127},"13000":{"m":125,"g":127},"12908":{"m":125,"g":127},"12952":{"m":125,"g":127},"13010":{"m":125,"g":127},"11850":{"m":125,"g":127},"12781":{"m":125,"g":127},"12224":{"m":125,"g":127},"13013":{"m":125,"g":127},"13009":{"m":125,"g":127},"13001":{"m":125,"g":127},"13005":{"m":125,"g":127},"12996":{"m":125,"g":127},"12999":{"m":125,"g":127},"12966":{"m":125,"g":127},"12984":{"m":125,"g":127},"12982":{"m":125,"g":127},"10225":{"m":125,"g":127},"10702":{"m":125,"g":127},"11719":{"m":125,"g":127},"12916":{"m":125,"g":127},"12239":{"m":125,"g":127},"12883":{"m":125,"g":127},"12912":{"m":125,"g":127},"12604":{"m":125,"g":127},"12934":{"m":125,"g":127},"9528":{"m":125,"g":127},"12931":{"m":125,"g":127},"12959":{"m":125,"g":127},"12926":{"m":125,"g":127},"12803":{"m":125,"g":127},"12957":{"m":125,"g":127},"12943":{"m":125,"g":127},"12834":{"m":125,"g":127},"12956":{"m":125,"g":127},"12554":{"m":125,"g":127},"12946":{"m":125,"g":127},"12948":{"m":125,"g":127},"12940":{"m":125,"g":127},"11812":{"m":125,"g":127},"12928":{"m":125,"g":127},"12927":{"m":125,"g":127},"12839":{"m":125,"g":127},"10775":{"m":125,"g":127},"12917":{"m":125,"g":127},"12920":{"m":125,"g":127},"12332":{"m":125,"g":127},"12919":{"m":125,"g":127},"12448":{"m":125,"g":127},"12906":{"m":125,"g":127},"12907":{"m":125,"g":127},"12911":{"m":125,"g":127},"12895":{"m":125,"g":127},"12905":{"m":125,"g":127},"12904":{"m":125,"g":127},"12865":{"m":125,"g":127},"12900":{"m":125,"g":127},"12896":{"m":125,"g":127},"12891":{"m":125,"g":127},"12897":{"m":125,"g":127},"12889":{"m":125,"g":127},"12870":{"m":125,"g":127},"12361":{"m":125,"g":127},"12811":{"m":125,"g":127},"12888":{"m":125,"g":127},"12832":{"m":125,"g":127},"12843":{"m":125,"g":127},"12886":{"m":125,"g":127},"12798":{"m":125,"g":127},"12868":{"m":125,"g":127},"12853":{"m":125,"g":127},"12846":{"m":125,"g":127},"12582":{"m":125,"g":127},"12805":{"m":125,"g":127},"12849":{"m":125,"g":127},"12859":{"m":125,"g":127},"12852":{"m":125,"g":127},"12851":{"m":125,"g":127},"12856":{"m":125,"g":127},"12801":{"m":125,"g":127},"12822":{"m":125,"g":127},"12431":{"m":125,"g":127},"12825":{"m":125,"g":127},"12836":{"m":125,"g":127},"12374":{"m":125,"g":127},"12812":{"m":125,"g":127},"12816":{"m":125,"g":127},"12090":{"m":125,"g":127},"12758":{"m":125,"g":127},"12520":{"m":125,"g":127},"12765":{"m":125,"g":127},"12761":{"m":125,"g":127},"12763":{"m":125,"g":127},"12776":{"m":125,"g":127},"12772":{"m":125,"g":127},"12788":{"m":125,"g":127},"12576":{"m":125,"g":127},"12782":{"m":125,"g":127},"12794":{"m":125,"g":127},"12715":{"m":125,"g":127},"12795":{"m":125,"g":127},"12724":{"m":125,"g":127},"12279":{"m":125,"g":127},"12717":{"m":125,"g":127},"8243":{"m":125,"g":127},"11904":{"m":125,"g":127},"12684":{"m":125,"g":127},"12764":{"m":125,"g":127},"12363":{"m":125,"g":127},"11051":{"m":125,"g":127},"12749":{"m":125,"g":127},"13129":{"m":126,"g":127},"10808":{"m":126,"g":127},"12617":{"m":126,"g":127},"7906":{"m":126,"g":127},"13149":{"m":126,"g":127},"7886":{"m":126,"g":127},"9790":{"m":126,"g":127},"13120":{"m":126,"g":127},"12458":{"m":126,"g":127},"11961":{"m":126,"g":127},"13137":{"m":126,"g":127},"12860":{"m":126,"g":127},"13136":{"m":126,"g":127},"13135":{"m":126,"g":127},"13077":{"m":126,"g":127},"13132":{"m":126,"g":127},"13131":{"m":126,"g":127},"13133":{"m":126,"g":127},"12942":{"m":126,"g":127},"12817":{"m":126,"g":127},"13095":{"m":126,"g":127},"12666":{"m":126,"g":127},"12396":{"m":126,"g":127},"13118":{"m":126,"g":127},"12872":{"m":126,"g":127},"12997":{"m":126,"g":127},"12863":{"m":126,"g":127},"13114":{"m":126,"g":127},"13039":{"m":126,"g":127},"11856":{"m":126,"g":127},"13105":{"m":126,"g":127},"13056":{"m":126,"g":127},"12583":{"m":126,"g":127},"13093":{"m":126,"g":127},"12660":{"m":126,"g":127},"13090":{"m":126,"g":127},"12866":{"m":126,"g":127},"12915":{"m":126,"g":127},"13018":{"m":126,"g":127},"13092":{"m":126,"g":127},"11645":{"m":126,"g":127},"13041":{"m":126,"g":127},"10862":{"m":126,"g":127},"13088":{"m":126,"g":127},"12814":{"m":126,"g":127},"12941":{"m":126,"g":127},"12994":{"m":126,"g":127},"13076":{"m":126,"g":127},"13015":{"m":126,"g":127},"13063":{"m":126,"g":127},"12199":{"m":126,"g":127},"12983":{"m":126,"g":127},"13037":{"m":126,"g":127},"12689":{"m":126,"g":127},"11609":{"m":126,"g":127},"12976":{"m":126,"g":127},"13050":{"m":126,"g":127},"12980":{"m":126,"g":127},"11938":{"m":126,"g":127},"13057":{"m":126,"g":127},"13053":{"m":126,"g":127},"13036":{"m":126,"g":127},"13043":{"m":126,"g":127},"13029":{"m":126,"g":127},"12885":{"m":126,"g":127},"12518":{"m":126,"g":127},"13035":{"m":126,"g":127},"13027":{"m":126,"g":127},"13028":{"m":126,"g":127},"12218":{"m":126,"g":127},"12753":{"m":126,"g":127},"12869":{"m":126,"g":127},"12993":{"m":126,"g":127},"13012":{"m":126,"g":127},"13366":{"m":128,"g":131},"13389":{"m":128,"g":131},"13387":{"m":128,"g":131},"12903":{"m":128,"g":131},"13388":{"m":128,"g":131},"12874":{"m":128,"g":131},"13386":{"m":128,"g":131},"13385":{"m":128,"g":131},"13384":{"m":128,"g":131},"13381":{"m":128,"g":131},"13335":{"m":128,"g":131},"13228":{"m":128,"g":131},"13371":{"m":128,"g":131},"13339":{"m":128,"g":131},"12978":{"m":128,"g":131},"13332":{"m":128,"g":131},"13263":{"m":128,"g":131},"13375":{"m":128,"g":131},"13344":{"m":128,"g":131},"13373":{"m":128,"g":131},"13372":{"m":128,"g":131},"13336":{"m":128,"g":131},"13369":{"m":128,"g":131},"13348":{"m":128,"g":131},"13199":{"m":128,"g":131},"13179":{"m":128,"g":131},"12310":{"m":128,"g":131},"11870":{"m":128,"g":131},"13358":{"m":128,"g":131},"13101":{"m":128,"g":131},"13321":{"m":128,"g":131},"13355":{"m":128,"g":131},"13351":{"m":128,"g":131},"13181":{"m":128,"g":131},"13341":{"m":128,"g":131},"13325":{"m":128,"g":131},"12329":{"m":128,"g":131},"13337":{"m":128,"g":131},"12001":{"m":128,"g":131},"12692":{"m":128,"g":131},"12443":{"m":128,"g":131},"13331":{"m":128,"g":131},"13330":{"m":128,"g":131},"13329":{"m":128,"g":131},"10568":{"m":128,"g":131},"13306":{"m":128,"g":131},"13297":{"m":128,"g":131},"13326":{"m":128,"g":131},"13323":{"m":128,"g":131},"13287":{"m":128,"g":131},"13322":{"m":128,"g":131},"7415":{"m":128,"g":131},"13286":{"m":128,"g":131},"13285":{"m":128,"g":131},"13259":{"m":128,"g":131},"13320":{"m":128,"g":131},"13295":{"m":128,"g":131},"13318":{"m":128,"g":131},"13314":{"m":128,"g":131},"13226":{"m":128,"g":131},"13278":{"m":128,"g":131},"13279":{"m":128,"g":131},"12612":{"m":128,"g":131},"12871":{"m":128,"g":131},"13317":{"m":128,"g":131},"13294":{"m":128,"g":131},"13316":{"m":128,"g":131},"13315":{"m":128,"g":131},"13312":{"m":128,"g":127},"13091":{"m":128,"g":127},"13100":{"m":128,"g":127},"13311":{"m":128,"g":127},"13310":{"m":128,"g":127},"13045":{"m":128,"g":127},"13293":{"m":128,"g":127},"13305":{"m":128,"g":127},"13170":{"m":128,"g":127},"13298":{"m":128,"g":127},"13274":{"m":128,"g":127},"13235":{"m":128,"g":127},"13272":{"m":128,"g":127},"13260":{"m":128,"g":127},"10573":{"m":128,"g":127},"10665":{"m":128,"g":127},"13236":{"m":128,"g":127},"12777":{"m":128,"g":127},"13277":{"m":128,"g":127},"13288":{"m":128,"g":127},"13254":{"m":128,"g":127},"13284":{"m":128,"g":127},"13283":{"m":128,"g":127},"13242":{"m":128,"g":127},"12191":{"m":128,"g":127},"12605":{"m":128,"g":127},"12623":{"m":128,"g":127},"12622":{"m":128,"g":127},"12620":{"m":128,"g":127},"13113":{"m":128,"g":127},"13247":{"m":128,"g":127},"13237":{"m":128,"g":127},"13265":{"m":128,"g":127},"11589":{"m":128,"g":127},"13261":{"m":128,"g":127},"13256":{"m":128,"g":127},"13257":{"m":128,"g":127},"13255":{"m":128,"g":127},"13096":{"m":128,"g":127},"13221":{"m":128,"g":127},"13213":{"m":128,"g":127},"13246":{"m":128,"g":127},"13243":{"m":128,"g":127},"13186":{"m":128,"g":127},"13239":{"m":128,"g":127},"13097":{"m":128,"g":127},"12392":{"m":128,"g":127},"13222":{"m":128,"g":127},"13188":{"m":128,"g":127},"13218":{"m":128,"g":127},"13087":{"m":128,"g":127},"13142":{"m":128,"g":127},"11595":{"m":128,"g":127},"13210":{"m":128,"g":127},"12774":{"m":128,"g":127},"13211":{"m":128,"g":127},"13220":{"m":128,"g":127},"13171":{"m":128,"g":127},"13215":{"m":128,"g":127},"13212":{"m":128,"g":127},"10485":{"m":128,"g":127},"12543":{"m":128,"g":127},"12201":{"m":128,"g":127},"12376":{"m":128,"g":127},"13155":{"m":128,"g":127},"13148":{"m":128,"g":127},"13190":{"m":128,"g":127},"13178":{"m":128,"g":127},"13102":{"m":128,"g":127},"13154":{"m":128,"g":127},"12975":{"m":128,"g":127},"13163":{"m":128,"g":127},"13172":{"m":128,"g":127},"13150":{"m":128,"g":127},"10973":{"m":128,"g":127},"12288":{"m":128,"g":127},"13162":{"m":128,"g":127},"12215":{"m":128,"g":127},"13104":{"m":128,"g":127},"13164":{"m":128,"g":127},"13127":{"m":128,"g":127},"12979":{"m":128,"g":127},"13128":{"m":128,"g":127},"12998":{"m":128,"g":127},"13075":{"m":128,"g":127},"10907":{"m":128,"g":127},"13153":{"m":128,"g":127},"12214":{"m":128,"g":127},"14316":{"m":129,"g":131},"14324":{"m":129,"g":131},"14323":{"m":129,"g":131},"14317":{"m":129,"g":131},"14309":{"m":129,"g":131},"14319":{"m":129,"g":131},"14262":{"m":129,"g":131},"14315":{"m":129,"g":131},"14249":{"m":129,"g":131},"11423":{"m":129,"g":131},"14278":{"m":129,"g":131},"14299":{"m":129,"g":131},"13089":{"m":129,"g":131},"14133":{"m":129,"g":131},"14269":{"m":129,"g":131},"14281":{"m":129,"g":131},"14287":{"m":129,"g":131},"14286":{"m":129,"g":131},"14283":{"m":129,"g":131},"14252":{"m":129,"g":131},"14244":{"m":129,"g":131},"14279":{"m":129,"g":131},"14276":{"m":129,"g":131},"13738":{"m":129,"g":131},"14047":{"m":129,"g":131},"14274":{"m":129,"g":131},"14257":{"m":129,"g":131},"14254":{"m":129,"g":131},"13700":{"m":129,"g":131},"14267":{"m":129,"g":131},"14261":{"m":129,"g":131},"14172":{"m":129,"g":131},"14263":{"m":129,"g":131},"14259":{"m":129,"g":131},"14260":{"m":129,"g":131},"14222":{"m":129,"g":131},"14232":{"m":129,"g":131},"13968":{"m":129,"g":131},"14256":{"m":129,"g":131},"14255":{"m":129,"g":131},"13880":{"m":129,"g":131},"13794":{"m":129,"g":131},"14247":{"m":129,"g":131},"13843":{"m":129,"g":131},"14250":{"m":129,"g":131},"14245":{"m":129,"g":131},"14243":{"m":129,"g":131},"14241":{"m":129,"g":131},"14240":{"m":129,"g":131},"14237":{"m":129,"g":131},"14152":{"m":129,"g":131},"14179":{"m":129,"g":131},"14229":{"m":129,"g":131},"14230":{"m":129,"g":131},"13693":{"m":129,"g":131},"14228":{"m":129,"g":131},"14122":{"m":129,"g":131},"14088":{"m":129,"g":131},"13887":{"m":129,"g":131},"14219":{"m":129,"g":131},"14218":{"m":129,"g":131},"14214":{"m":129,"g":131},"14212":{"m":129,"g":131},"14211":{"m":129,"g":131},"14173":{"m":129,"g":131},"14165":{"m":129,"g":131},"14123":{"m":129,"g":131},"14186":{"m":129,"g":131},"14180":{"m":129,"g":131},"14167":{"m":129,"g":131},"14044":{"m":129,"g":131},"14182":{"m":129,"g":131},"14183":{"m":129,"g":131},"14181":{"m":129,"g":131},"14003":{"m":129,"g":131},"14034":{"m":129,"g":131},"14153":{"m":129,"g":131},"14187":{"m":129,"g":131},"12181":{"m":129,"g":131},"14155":{"m":129,"g":131},"14104":{"m":129,"g":131},"13873":{"m":129,"g":131},"14059":{"m":129,"g":131},"14140":{"m":129,"g":131},"13646":{"m":129,"g":131},"13841":{"m":129,"g":131},"14171":{"m":129,"g":131},"13907":{"m":129,"g":131},"14065":{"m":129,"g":131},"14166":{"m":129,"g":131},"14005":{"m":129,"g":131},"12494":{"m":129,"g":131},"14163":{"m":129,"g":131},"14156":{"m":129,"g":131},"14148":{"m":129,"g":131},"14161":{"m":129,"g":131},"14052":{"m":129,"g":131},"14150":{"m":129,"g":131},"14157":{"m":129,"g":131},"14154":{"m":129,"g":131},"14147":{"m":129,"g":131},"14146":{"m":129,"g":131},"14145":{"m":129,"g":131},"14136":{"m":129,"g":131},"13956":{"m":129,"g":131},"13759":{"m":129,"g":131},"14151":{"m":129,"g":131},"14135":{"m":129,"g":131},"14130":{"m":129,"g":131},"14119":{"m":129,"g":131},"14131":{"m":129,"g":131},"14129":{"m":129,"g":131},"14121":{"m":129,"g":131},"12306":{"m":129,"g":131},"14124":{"m":129,"g":131},"14113":{"m":129,"g":131},"14117":{"m":129,"g":131},"13377":{"m":129,"g":131},"13488":{"m":129,"g":131},"14111":{"m":129,"g":131},"12558":{"m":129,"g":131},"10712":{"m":129,"g":131},"14106":{"m":129,"g":131},"14067":{"m":129,"g":131},"14096":{"m":129,"g":131},"14094":{"m":129,"g":131},"13724":{"m":129,"g":131},"13904":{"m":129,"g":131},"14076":{"m":129,"g":131},"13936":{"m":129,"g":131},"14006":{"m":129,"g":131},"14082":{"m":129,"g":131},"13205":{"m":129,"g":131},"14079":{"m":129,"g":131},"14036":{"m":129,"g":131},"13944":{"m":129,"g":131},"13749":{"m":129,"g":131},"14069":{"m":129,"g":131},"13946":{"m":129,"g":131},"14048":{"m":129,"g":131},"13960":{"m":129,"g":131},"14057":{"m":129,"g":131},"13425":{"m":129,"g":131},"13854":{"m":129,"g":131},"14002":{"m":129,"g":131},"13855":{"m":129,"g":131},"14040":{"m":129,"g":131},"13895":{"m":129,"g":131},"13976":{"m":129,"g":131},"13814":{"m":129,"g":131},"13965":{"m":129,"g":131},"14026":{"m":129,"g":131},"14017":{"m":129,"g":131},"14033":{"m":129,"g":131},"14030":{"m":129,"g":131},"14028":{"m":129,"g":131},"13761":{"m":129,"g":131},"14027":{"m":129,"g":131},"14025":{"m":129,"g":131},"13824":{"m":129,"g":131},"12277":{"m":129,"g":131},"13941":{"m":129,"g":131},"13983":{"m":129,"g":131},"14018":{"m":129,"g":131},"14007":{"m":129,"g":131},"14022":{"m":129,"g":131},"14019":{"m":129,"g":131},"14021":{"m":129,"g":131},"13937":{"m":129,"g":131},"13990":{"m":129,"g":131},"14020":{"m":129,"g":131},"13966":{"m":129,"g":131},"13872":{"m":129,"g":131},"14016":{"m":129,"g":131},"14013":{"m":129,"g":131},"14015":{"m":129,"g":131},"14014":{"m":129,"g":131},"14012":{"m":129,"g":131},"13151":{"m":129,"g":131},"13892":{"m":129,"g":131},"14000":{"m":129,"g":131},"14009":{"m":129,"g":131},"13766":{"m":129,"g":131},"13754":{"m":129,"g":131},"12491":{"m":129,"g":131},"13994":{"m":129,"g":131},"13991":{"m":129,"g":131},"13977":{"m":129,"g":131},"13203":{"m":129,"g":131},"12588":{"m":129,"g":131},"13922":{"m":129,"g":131},"13852":{"m":129,"g":131},"13963":{"m":129,"g":131},"10071":{"m":129,"g":131},"13961":{"m":129,"g":131},"13962":{"m":129,"g":131},"13958":{"m":129,"g":131},"7725":{"m":129,"g":131},"13925":{"m":129,"g":131},"12786":{"m":129,"g":131},"13954":{"m":129,"g":131},"13951":{"m":129,"g":131},"13950":{"m":129,"g":131},"13945":{"m":129,"g":131},"13942":{"m":129,"g":131},"12969":{"m":129,"g":131},"13910":{"m":129,"g":131},"13866":{"m":129,"g":131},"13903":{"m":129,"g":131},"13421":{"m":129,"g":131},"13851":{"m":129,"g":131},"13544":{"m":129,"g":131},"13938":{"m":129,"g":131},"13935":{"m":129,"g":131},"13933":{"m":129,"g":131},"13859":{"m":129,"g":131},"13931":{"m":129,"g":131},"13928":{"m":129,"g":131},"13927":{"m":129,"g":131},"13657":{"m":129,"g":131},"13081":{"m":129,"g":131},"12078":{"m":129,"g":131},"13921":{"m":129,"g":131},"13905":{"m":129,"g":131},"13916":{"m":129,"g":131},"13642":{"m":129,"g":131},"13908":{"m":129,"g":131},"13848":{"m":129,"g":131},"13793":{"m":129,"g":131},"13901":{"m":129,"g":131},"13827":{"m":129,"g":131},"13889":{"m":129,"g":131},"13890":{"m":129,"g":131},"13888":{"m":129,"g":131},"11893":{"m":129,"g":131},"13891":{"m":129,"g":131},"13870":{"m":129,"g":131},"13874":{"m":129,"g":131},"13860":{"m":129,"g":131},"13871":{"m":129,"g":131},"13572":{"m":129,"g":131},"10275":{"m":129,"g":131},"13487":{"m":129,"g":131},"13786":{"m":129,"g":131},"13834":{"m":129,"g":131},"13822":{"m":129,"g":131},"13783":{"m":129,"g":131},"13763":{"m":129,"g":131},"13752":{"m":129,"g":131},"13864":{"m":129,"g":131},"13865":{"m":129,"g":131},"13853":{"m":129,"g":131},"10027":{"m":129,"g":131},"13745":{"m":129,"g":131},"13858":{"m":129,"g":131},"13612":{"m":129,"g":131},"11871":{"m":129,"g":131},"13713":{"m":129,"g":131},"13508":{"m":129,"g":131},"13846":{"m":129,"g":131},"13819":{"m":129,"g":131},"13245":{"m":129,"g":131},"13751":{"m":129,"g":131},"13833":{"m":129,"g":131},"13831":{"m":129,"g":131},"13792":{"m":129,"g":131},"13829":{"m":129,"g":131},"13800":{"m":129,"g":131},"13656":{"m":129,"g":131},"13820":{"m":129,"g":131},"13201":{"m":129,"g":131},"13650":{"m":129,"g":131},"13816":{"m":129,"g":131},"13601":{"m":129,"g":131},"13810":{"m":129,"g":131},"13802":{"m":129,"g":131},"13815":{"m":129,"g":131},"13813":{"m":129,"g":131},"13687":{"m":129,"g":131},"13806":{"m":129,"g":131},"13781":{"m":129,"g":131},"13791":{"m":129,"g":131},"13180":{"m":129,"g":131},"13718":{"m":129,"g":131},"13787":{"m":129,"g":131},"13764":{"m":129,"g":131},"13709":{"m":129,"g":131},"13776":{"m":129,"g":131},"13777":{"m":129,"g":131},"13727":{"m":129,"g":131},"13676":{"m":129,"g":131},"13690":{"m":129,"g":131},"13547":{"m":129,"g":131},"13533":{"m":129,"g":131},"13769":{"m":129,"g":131},"13478":{"m":129,"g":131},"13736":{"m":129,"g":131},"13706":{"m":129,"g":131},"13768":{"m":129,"g":131},"13704":{"m":129,"g":131},"13720":{"m":129,"g":131},"13714":{"m":129,"g":131},"12759":{"m":129,"g":131},"9405":{"m":129,"g":131},"13506":{"m":129,"g":131},"13756":{"m":129,"g":131},"13669":{"m":129,"g":131},"13729":{"m":129,"g":131},"13702":{"m":129,"g":131},"13746":{"m":129,"g":131},"13484":{"m":129,"g":131},"12690":{"m":129,"g":131},"13701":{"m":129,"g":131},"12949":{"m":129,"g":131},"13707":{"m":129,"g":131},"13694":{"m":129,"g":131},"13466":{"m":129,"g":131},"13739":{"m":129,"g":131},"13737":{"m":129,"g":131},"13735":{"m":129,"g":131},"13734":{"m":129,"g":131},"13733":{"m":129,"g":131},"13407":{"m":129,"g":131},"13649":{"m":129,"g":131},"13705":{"m":129,"g":131},"13630":{"m":129,"g":131},"13719":{"m":129,"g":131},"13327":{"m":129,"g":131},"13647":{"m":129,"g":131},"13708":{"m":129,"g":131},"13498":{"m":129,"g":131},"13590":{"m":129,"g":131},"13679":{"m":129,"g":131},"13564":{"m":129,"g":131},"13686":{"m":129,"g":131},"13177":{"m":129,"g":131},"13665":{"m":129,"g":131},"13675":{"m":129,"g":131},"13640":{"m":129,"g":131},"13587":{"m":129,"g":131},"13596":{"m":129,"g":131},"13555":{"m":129,"g":131},"13619":{"m":129,"g":131},"13301":{"m":129,"g":131},"12672":{"m":129,"g":131},"13683":{"m":129,"g":131},"13685":{"m":129,"g":131},"13627":{"m":129,"g":131},"13678":{"m":129,"g":131},"13677":{"m":129,"g":131},"13659":{"m":129,"g":131},"13667":{"m":129,"g":131},"13610":{"m":129,"g":131},"11526":{"m":129,"g":131},"13038":{"m":129,"g":131},"13600":{"m":129,"g":131},"13459":{"m":129,"g":131},"13666":{"m":129,"g":131},"12964":{"m":129,"g":131},"13655":{"m":129,"g":131},"13524":{"m":129,"g":131},"11577":{"m":129,"g":131},"13663":{"m":129,"g":131},"13637":{"m":129,"g":131},"13634":{"m":129,"g":131},"13644":{"m":129,"g":131},"13197":{"m":129,"g":131},"13617":{"m":129,"g":131},"13633":{"m":129,"g":131},"13453":{"m":129,"g":131},"13614":{"m":129,"g":131},"13554":{"m":129,"g":131},"13583":{"m":129,"g":131},"13328":{"m":129,"g":131},"13253":{"m":129,"g":131},"13248":{"m":129,"g":131},"12379":{"m":129,"g":131},"13528":{"m":129,"g":131},"13613":{"m":129,"g":131},"13429":{"m":129,"g":131},"13562":{"m":129,"g":131},"13055":{"m":129,"g":131},"13603":{"m":129,"g":131},"13604":{"m":129,"g":131},"13357":{"m":129,"g":131},"13570":{"m":129,"g":131},"13577":{"m":129,"g":131},"13465":{"m":129,"g":131},"13049":{"m":129,"g":131},"13448":{"m":129,"g":131},"13589":{"m":129,"g":131},"13567":{"m":129,"g":131},"13568":{"m":129,"g":131},"13413":{"m":129,"g":131},"13558":{"m":129,"g":131},"13557":{"m":129,"g":131},"13542":{"m":129,"g":131},"13481":{"m":129,"g":131},"12740":{"m":129,"g":131},"13551":{"m":129,"g":131},"13452":{"m":129,"g":131},"9234":{"m":129,"g":131},"13047":{"m":129,"g":131},"13548":{"m":129,"g":131},"13543":{"m":129,"g":131},"13541":{"m":129,"g":131},"13489":{"m":129,"g":131},"13495":{"m":129,"g":131},"13540":{"m":129,"g":131},"13537":{"m":129,"g":131},"13534":{"m":129,"g":131},"13536":{"m":129,"g":131},"13532":{"m":129,"g":131},"13527":{"m":129,"g":131},"13474":{"m":129,"g":131},"13525":{"m":129,"g":131},"13522":{"m":129,"g":131},"13521":{"m":129,"g":131},"13519":{"m":129,"g":131},"13516":{"m":129,"g":131},"12962":{"m":129,"g":131},"13513":{"m":129,"g":131},"13512":{"m":129,"g":131},"13510":{"m":129,"g":131},"13509":{"m":129,"g":131},"13168":{"m":129,"g":131},"13157":{"m":129,"g":131},"13501":{"m":129,"g":131},"13126":{"m":129,"g":131},"13496":{"m":129,"g":131},"13374":{"m":129,"g":131},"13491":{"m":129,"g":131},"13482":{"m":129,"g":131},"13486":{"m":129,"g":131},"13393":{"m":129,"g":131},"13460":{"m":129,"g":131},"13094":{"m":129,"g":131},"13479":{"m":129,"g":131},"13476":{"m":129,"g":131},"12149":{"m":129,"g":131},"13289":{"m":129,"g":131},"13473":{"m":129,"g":131},"13258":{"m":129,"g":131},"13229":{"m":129,"g":131},"13462":{"m":129,"g":131},"13444":{"m":129,"g":131},"13458":{"m":129,"g":131},"13449":{"m":129,"g":131},"13455":{"m":129,"g":131},"13140":{"m":129,"g":131},"13463":{"m":129,"g":131},"13264":{"m":129,"g":131},"12359":{"m":129,"g":131},"13461":{"m":129,"g":131},"13457":{"m":129,"g":131},"13173":{"m":129,"g":131},"13456":{"m":129,"g":131},"13273":{"m":129,"g":131},"13022":{"m":129,"g":131},"13450":{"m":129,"g":131},"13447":{"m":129,"g":131},"13445":{"m":129,"g":131},"13443":{"m":129,"g":131},"13418":{"m":129,"g":131},"13217":{"m":129,"g":131},"5879":{"m":129,"g":131},"11900":{"m":129,"g":131},"13144":{"m":129,"g":131},"13379":{"m":129,"g":131},"13282":{"m":129,"g":131},"13420":{"m":129,"g":131},"13416":{"m":129,"g":131},"13415":{"m":129,"g":131},"13399":{"m":129,"g":131},"13004":{"m":129,"g":131},"13324":{"m":129,"g":131},"13112":{"m":129,"g":131},"11644":{"m":129,"g":131},"13398":{"m":129,"g":131},"13396":{"m":129,"g":131},"13345":{"m":129,"g":131},"13338":{"m":129,"g":131},"12065":{"m":129,"g":131},"13391":{"m":129,"g":131},"13383":{"m":129,"g":131},"14670":{"m":130,"g":131},"14650":{"m":130,"g":131},"14457":{"m":130,"g":131},"14657":{"m":130,"g":131},"14667":{"m":130,"g":131},"14634":{"m":130,"g":131},"14664":{"m":130,"g":131},"14663":{"m":130,"g":131},"14658":{"m":130,"g":131},"14497":{"m":130,"g":131},"14651":{"m":130,"g":131},"14649":{"m":130,"g":131},"14356":{"m":130,"g":131},"14629":{"m":130,"g":131},"14558":{"m":130,"g":131},"12527":{"m":130,"g":131},"14632":{"m":130,"g":131},"14606":{"m":130,"g":131},"12551":{"m":130,"g":131},"14625":{"m":130,"g":131},"14556":{"m":130,"g":131},"14585":{"m":130,"g":131},"14618":{"m":130,"g":131},"14203":{"m":130,"g":131},"14452":{"m":130,"g":131},"14612":{"m":130,"g":131},"14609":{"m":130,"g":131},"14604":{"m":130,"g":131},"14608":{"m":130,"g":131},"14605":{"m":130,"g":131},"14600":{"m":130,"g":131},"14591":{"m":130,"g":131},"14386":{"m":130,"g":131},"14573":{"m":130,"g":131},"14551":{"m":130,"g":131},"14141":{"m":130,"g":131},"14590":{"m":130,"g":131},"14588":{"m":130,"g":131},"14587":{"m":130,"g":131},"14586":{"m":130,"g":131},"14517":{"m":130,"g":131},"14455":{"m":130,"g":131},"14553":{"m":130,"g":131},"13573":{"m":130,"g":131},"14132":{"m":130,"g":131},"14185":{"m":130,"g":131},"14576":{"m":130,"g":131},"14577":{"m":130,"g":131},"13725":{"m":130,"g":131},"13998":{"m":130,"g":131},"14569":{"m":130,"g":131},"14412":{"m":130,"g":131},"14544":{"m":130,"g":131},"14561":{"m":130,"g":131},"14560":{"m":130,"g":131},"14559":{"m":130,"g":131},"14337":{"m":130,"g":131},"14555":{"m":130,"g":131},"14494":{"m":130,"g":131},"14476":{"m":130,"g":131},"14205":{"m":130,"g":131},"14557":{"m":130,"g":131},"14447":{"m":130,"g":131},"14552":{"m":130,"g":131},"14518":{"m":130,"g":131},"14538":{"m":130,"g":131},"14520":{"m":130,"g":131},"14535":{"m":130,"g":131},"14493":{"m":130,"g":131},"14464":{"m":130,"g":131},"14543":{"m":130,"g":131},"13897":{"m":130,"g":131},"14505":{"m":130,"g":131},"14539":{"m":130,"g":131},"13115":{"m":130,"g":131},"14533":{"m":130,"g":131},"12324":{"m":130,"g":131},"14290":{"m":130,"g":131},"14528":{"m":130,"g":131},"14530":{"m":130,"g":131},"11791":{"m":130,"g":131},"14522":{"m":130,"g":131},"14465":{"m":130,"g":131},"14521":{"m":130,"g":131},"14291":{"m":130,"g":131},"14427":{"m":130,"g":131},"14516":{"m":130,"g":131},"14514":{"m":130,"g":131},"14513":{"m":130,"g":131},"14512":{"m":130,"g":131},"14312":{"m":130,"g":131},"14405":{"m":130,"g":131},"14420":{"m":130,"g":131},"14460":{"m":130,"g":131},"14508":{"m":130,"g":131},"13607":{"m":130,"g":131},"14507":{"m":130,"g":131},"12471":{"m":130,"g":131},"14093":{"m":130,"g":131},"14234":{"m":130,"g":131},"13434":{"m":130,"g":131},"14506":{"m":130,"g":131},"14471":{"m":130,"g":131},"14459":{"m":130,"g":131},"14466":{"m":130,"g":131},"14456":{"m":130,"g":131},"14097":{"m":130,"g":131},"14499":{"m":130,"g":131},"13584":{"m":130,"g":131},"14364":{"m":130,"g":131},"13861":{"m":130,"g":131},"13996":{"m":130,"g":131},"14472":{"m":130,"g":131},"14484":{"m":130,"g":131},"14444":{"m":130,"g":131},"14475":{"m":130,"g":131},"14463":{"m":130,"g":131},"14473":{"m":130,"g":131},"14468":{"m":130,"g":131},"13836":{"m":130,"g":131},"14421":{"m":130,"g":131},"14432":{"m":130,"g":131},"14251":{"m":130,"g":131},"14450":{"m":130,"g":131},"14445":{"m":130,"g":131},"14446":{"m":130,"g":131},"8287":{"m":130,"g":131},"14348":{"m":130,"g":131},"14350":{"m":130,"g":131},"14441":{"m":130,"g":131},"14325":{"m":130,"g":131},"14440":{"m":130,"g":131},"14438":{"m":130,"g":131},"14143":{"m":130,"g":131},"14434":{"m":130,"g":131},"14366":{"m":130,"g":131},"14430":{"m":130,"g":131},"14429":{"m":130,"g":131},"14225":{"m":130,"g":131},"14409":{"m":130,"g":131},"14213":{"m":130,"g":131},"14224":{"m":130,"g":131},"14334":{"m":130,"g":131},"14399":{"m":130,"g":131},"12446":{"m":130,"g":131},"13359":{"m":130,"g":131},"14383":{"m":130,"g":131},"14394":{"m":130,"g":131},"14381":{"m":130,"g":131},"12309":{"m":130,"g":131},"14393":{"m":130,"g":131},"12316":{"m":130,"g":131},"14292":{"m":130,"g":131},"14392":{"m":130,"g":131},"14272":{"m":130,"g":131},"13731":{"m":130,"g":131},"14359":{"m":130,"g":131},"14377":{"m":130,"g":131},"14330":{"m":130,"g":131},"14277":{"m":130,"g":131},"14375":{"m":130,"g":131},"14374":{"m":130,"g":131},"14253":{"m":130,"g":131},"14372":{"m":130,"g":131},"14226":{"m":130,"g":131},"14371":{"m":130,"g":131},"14326":{"m":130,"g":131},"9660":{"m":130,"g":131},"12330":{"m":130,"g":131},"14355":{"m":130,"g":131},"13585":{"m":130,"g":131},"14362":{"m":130,"g":131},"14271":{"m":130,"g":131},"14295":{"m":130,"g":131},"13980":{"m":130,"g":131},"14347":{"m":130,"g":131},"14333":{"m":130,"g":131},"12441":{"m":130,"g":131},"14344":{"m":130,"g":131},"14265":{"m":130,"g":131},"14335":{"m":130,"g":131},"14336":{"m":130,"g":131},"13350":{"m":130,"g":131},"14266":{"m":130,"g":131},"14329":{"m":130,"g":131},"13812":{"m":130,"g":131},"14195":{"m":130,"g":131},"14321":{"m":130,"g":131},"13710":{"m":130,"g":131},"14858":{"m":132,"g":133},"14620":{"m":132,"g":133},"14304":{"m":132,"g":133},"14917":{"m":132,"g":133},"14307":{"m":132,"g":133},"14887":{"m":132,"g":133},"14911":{"m":132,"g":133},"14910":{"m":132,"g":133},"14852":{"m":132,"g":133},"14889":{"m":132,"g":133},"14890":{"m":132,"g":133},"14900":{"m":132,"g":133},"14899":{"m":132,"g":133},"12287":{"m":132,"g":133},"14878":{"m":132,"g":133},"14541":{"m":132,"g":133},"13641":{"m":132,"g":133},"14828":{"m":132,"g":133},"14827":{"m":132,"g":133},"14853":{"m":132,"g":133},"14876":{"m":132,"g":133},"14880":{"m":132,"g":133},"14877":{"m":132,"g":133},"14856":{"m":132,"g":133},"14875":{"m":132,"g":133},"14861":{"m":132,"g":133},"14845":{"m":132,"g":133},"14871":{"m":132,"g":133},"14313":{"m":132,"g":133},"14554":{"m":132,"g":133},"14865":{"m":132,"g":133},"14811":{"m":132,"g":133},"14836":{"m":132,"g":133},"14848":{"m":132,"g":133},"14854":{"m":132,"g":133},"14849":{"m":132,"g":133},"14844":{"m":132,"g":133},"14851":{"m":132,"g":133},"14850":{"m":132,"g":133},"14847":{"m":132,"g":133},"14442":{"m":132,"g":133},"14841":{"m":132,"g":133},"14796":{"m":132,"g":133},"14638":{"m":132,"g":133},"14823":{"m":132,"g":133},"14801":{"m":132,"g":133},"14837":{"m":132,"g":133},"14045":{"m":132,"g":133},"14833":{"m":132,"g":133},"14829":{"m":132,"g":133},"14769":{"m":132,"g":133},"14712":{"m":132,"g":133},"14716":{"m":132,"g":133},"14830":{"m":132,"g":133},"14834":{"m":132,"g":133},"14812":{"m":132,"g":133},"14831":{"m":132,"g":133},"14806":{"m":132,"g":133},"14822":{"m":132,"g":133},"14819":{"m":132,"g":133},"14710":{"m":132,"g":133},"14807":{"m":132,"g":133},"14770":{"m":132,"g":133},"14793":{"m":132,"g":133},"14808":{"m":132,"g":133},"14697":{"m":132,"g":133},"14720":{"m":132,"g":133},"14803":{"m":132,"g":133},"14788":{"m":132,"g":133},"14794":{"m":132,"g":133},"9650":{"m":132,"g":133},"14786":{"m":132,"g":133},"14784":{"m":132,"g":133},"14725":{"m":132,"g":133},"14777":{"m":132,"g":133},"14761":{"m":132,"g":133},"14759":{"m":132,"g":133},"14064":{"m":132,"g":133},"14768":{"m":132,"g":133},"14756":{"m":132,"g":133},"14744":{"m":132,"g":133},"14687":{"m":132,"g":133},"14763":{"m":132,"g":131},"14177":{"m":132,"g":131},"14758":{"m":132,"g":131},"14669":{"m":132,"g":131},"14740":{"m":132,"g":131},"14753":{"m":132,"g":131},"14698":{"m":132,"g":131},"14379":{"m":132,"g":131},"14752":{"m":132,"g":131},"14751":{"m":132,"g":131},"14745":{"m":132,"g":131},"12953":{"m":132,"g":131},"14743":{"m":132,"g":131},"14738":{"m":132,"g":131},"14733":{"m":132,"g":131},"12039":{"m":132,"g":131},"13432":{"m":132,"g":131},"14461":{"m":132,"g":131},"14686":{"m":132,"g":131},"14601":{"m":132,"g":131},"14622":{"m":132,"g":131},"14714":{"m":132,"g":131},"14707":{"m":132,"g":131},"14699":{"m":132,"g":131},"14647":{"m":132,"g":131},"14648":{"m":132,"g":131},"14683":{"m":132,"g":131},"14678":{"m":132,"g":131},"14676":{"m":132,"g":131},"14529":{"m":132,"g":131},"14689":{"m":132,"g":131},"14627":{"m":132,"g":131},"14679":{"m":132,"g":131},"14469":{"m":132,"g":131},"14614":{"m":132,"g":131},"14653":{"m":132,"g":131},"13147":{"m":132,"g":131},"14652":{"m":132,"g":131},"13334":{"m":132,"g":131},"14489":{"m":132,"g":131},"14675":{"m":132,"g":131},"14671":{"m":132,"g":131},"16253":{"m":134,"g":135},"16107":{"m":134,"g":135},"16244":"m134","16241":{"m":134,"g":135},"16211":{"m":134,"g":135},"16153":{"m":134,"g":135},"15942":{"m":134,"g":135},"16140":{"m":134,"g":135},"16142":{"m":134,"g":135},"16141":{"m":134,"g":135},"16129":{"m":134,"g":135},"10959":{"m":134,"g":135},"15888":{"m":134,"g":135},"16114":{"m":134,"g":135},"16131":{"m":134,"g":135},"16133":{"m":134,"g":135},"16053":{"m":134,"g":135},"16130":{"m":134,"g":135},"16105":{"m":134,"g":135},"15187":{"m":134,"g":135},"16123":{"m":134,"g":135},"15813":{"m":134,"g":135},"16081":{"m":134,"g":135},"15896":{"m":134,"g":135},"15877":{"m":134,"g":135},"15800":{"m":134,"g":135},"15985":{"m":134,"g":135},"16103":{"m":134,"g":135},"16099":{"m":134,"g":135},"15921":{"m":134,"g":135},"16101":{"m":134,"g":135},"16100":{"m":134,"g":135},"16098":{"m":134,"g":135},"16097":{"m":134,"g":135},"16094":{"m":134,"g":135},"16096":{"m":134,"g":135},"16093":{"m":134,"g":135},"15057":{"m":134,"g":135},"14838":{"m":134,"g":135},"16061":{"m":134,"g":135},"16066":{"m":134,"g":135},"16087":{"m":134,"g":135},"16069":{"m":134,"g":135},"16085":{"m":134,"g":135},"16062":{"m":134,"g":135},"14414":{"m":134,"g":135},"16047":{"m":134,"g":135},"15805":{"m":134,"g":135},"16054":{"m":134,"g":135},"16003":{"m":134,"g":135},"16046":{"m":134,"g":135},"16051":{"m":134,"g":135},"16038":{"m":134,"g":135},"16041":{"m":134,"g":135},"16039":{"m":134,"g":135},"16037":{"m":134,"g":135},"16035":{"m":134,"g":135},"16036":{"m":134,"g":135},"16010":{"m":134,"g":135},"14280":{"m":134,"g":135},"16028":{"m":134,"g":135},"16017":{"m":134,"g":135},"14873":{"m":134,"g":135},"15922":{"m":134,"g":135},"16016":{"m":134,"g":135},"15939":{"m":134,"g":135},"15998":{"m":134,"g":135},"15928":{"m":134,"g":135},"16008":{"m":134,"g":135},"16002":{"m":134,"g":135},"16013":{"m":134,"g":135},"16004":{"m":134,"g":135},"15615":{"m":134,"g":135},"15992":{"m":134,"g":135},"15991":{"m":134,"g":135},"16001":{"m":134,"g":135},"15945":{"m":134,"g":135},"15990":{"m":134,"g":135},"15988":{"m":134,"g":135},"15987":{"m":134,"g":135},"15891":{"m":134,"g":135},"15216":{"m":134,"g":135},"15693":{"m":134,"g":135},"15353":{"m":134,"g":135},"15835":{"m":134,"g":135},"15806":{"m":134,"g":135},"15937":{"m":134,"g":135},"15986":{"m":134,"g":135},"12596":{"m":134,"g":135},"15919":{"m":134,"g":135},"15947":{"m":134,"g":135},"15925":{"m":134,"g":135},"15936":{"m":134,"g":135},"15886":{"m":134,"g":135},"15943":{"m":134,"g":135},"15935":{"m":134,"g":135},"15934":{"m":134,"g":135},"15933":{"m":134,"g":135},"14736":{"m":134,"g":135},"15923":{"m":134,"g":135},"15907":{"m":134,"g":135},"15887":{"m":134,"g":135},"15920":{"m":134,"g":135},"15918":{"m":134,"g":135},"15905":{"m":134,"g":135},"15850":{"m":134,"g":135},"15915":{"m":134,"g":135},"15398":{"m":134,"g":135},"15914":{"m":134,"g":135},"15916":{"m":134,"g":135},"15913":{"m":134,"g":135},"15911":{"m":134,"g":135},"15910":{"m":134,"g":135},"14750":{"m":134,"g":135},"15906":{"m":134,"g":135},"15778":{"m":134,"g":135},"15844":{"m":134,"g":135},"15812":{"m":134,"g":135},"15842":{"m":134,"g":135},"15881":{"m":134,"g":135},"15874":{"m":134,"g":135},"15889":{"m":134,"g":135},"15846":{"m":134,"g":135},"14209":{"m":134,"g":135},"15849":{"m":134,"g":135},"15817":{"m":134,"g":135},"15870":{"m":134,"g":135},"15867":{"m":134,"g":135},"14644":{"m":134,"g":135},"15518":{"m":134,"g":135},"15821":{"m":134,"g":135},"15369":{"m":134,"g":135},"15858":{"m":134,"g":135},"15857":{"m":134,"g":135},"15820":{"m":134,"g":135},"15851":{"m":134,"g":135},"15701":{"m":134,"g":135},"15791":{"m":134,"g":135},"15847":{"m":134,"g":135},"15522":{"m":134,"g":135},"15796":{"m":134,"g":135},"15826":{"m":134,"g":135},"15822":{"m":134,"g":135},"15772":{"m":134,"g":135},"15356":{"m":134,"g":135},"15759":{"m":134,"g":135},"15827":{"m":134,"g":135},"15488":{"m":134,"g":135},"15815":{"m":134,"g":135},"15818":{"m":134,"g":135},"15802":{"m":134,"g":135},"15666":{"m":134,"g":135},"15709":{"m":134,"g":135},"14741":{"m":134,"g":135},"15803":{"m":134,"g":135},"15409":{"m":134,"g":135},"15798":{"m":134,"g":135},"15811":{"m":134,"g":135},"15720":{"m":134,"g":135},"15736":{"m":134,"g":135},"15801":{"m":134,"g":135},"15555":{"m":134,"g":135},"15770":{"m":134,"g":135},"11469":{"m":134,"g":135},"15586":{"m":134,"g":135},"15596":{"m":134,"g":135},"14032":{"m":134,"g":135},"15787":{"m":134,"g":135},"15700":{"m":134,"g":135},"15781":{"m":134,"g":133},"15149":{"m":134,"g":133},"15775":{"m":134,"g":133},"15782":{"m":134,"g":133},"15758":{"m":134,"g":133},"15745":{"m":134,"g":133},"15750":{"m":134,"g":133},"15780":{"m":134,"g":133},"15741":{"m":134,"g":133},"14137":{"m":134,"g":133},"15390":{"m":134,"g":133},"15718":{"m":134,"g":133},"15769":{"m":134,"g":133},"15768":{"m":134,"g":133},"15653":{"m":134,"g":133},"15652":{"m":134,"g":133},"15752":{"m":134,"g":133},"15747":{"m":134,"g":133},"15748":{"m":134,"g":133},"15743":{"m":134,"g":133},"15740":{"m":134,"g":133},"15655":{"m":134,"g":133},"15706":{"m":134,"g":133},"15459":{"m":134,"g":133},"15689":{"m":134,"g":133},"15593":{"m":134,"g":133},"15704":{"m":134,"g":133},"15691":{"m":134,"g":133},"15656":{"m":134,"g":133},"15717":{"m":134,"g":133},"15715":{"m":134,"g":133},"15722":{"m":134,"g":133},"15716":{"m":134,"g":133},"15719":{"m":134,"g":133},"11828":{"m":134,"g":133},"15500":{"m":134,"g":133},"15622":{"m":134,"g":133},"15624":{"m":134,"g":133},"15705":{"m":134,"g":133},"15177":{"m":134,"g":133},"15702":{"m":134,"g":133},"15538":{"m":134,"g":133},"15667":{"m":134,"g":133},"15695":{"m":134,"g":133},"15696":{"m":134,"g":133},"15606":{"m":134,"g":133},"15273":{"m":134,"g":133},"15672":{"m":134,"g":133},"15694":{"m":134,"g":133},"15469":{"m":134,"g":133},"12968":{"m":134,"g":133},"15644":{"m":134,"g":133},"15692":{"m":134,"g":133},"14570":{"m":134,"g":133},"15091":{"m":134,"g":133},"15646":{"m":134,"g":133},"15633":{"m":134,"g":133},"15688":{"m":134,"g":133},"15684":{"m":134,"g":133},"15312":{"m":134,"g":133},"15600":{"m":134,"g":133},"15463":{"m":134,"g":133},"15460":{"m":134,"g":133},"14983":{"m":134,"g":133},"15563":{"m":134,"g":133},"15582":{"m":134,"g":133},"14628":{"m":134,"g":133},"15539":{"m":134,"g":133},"15570":{"m":134,"g":133},"15635":{"m":134,"g":133},"15590":{"m":134,"g":133},"13576":{"m":134,"g":133},"15632":{"m":134,"g":133},"15621":{"m":134,"g":133},"15368":{"m":134,"g":133},"15611":{"m":134,"g":133},"15610":{"m":134,"g":133},"15572":{"m":134,"g":133},"15573":{"m":134,"g":133},"15616":{"m":134,"g":133},"15613":{"m":134,"g":133},"15607":{"m":134,"g":133},"15612":{"m":134,"g":133},"15164":{"m":134,"g":133},"15566":{"m":134,"g":133},"15599":{"m":134,"g":133},"15589":{"m":134,"g":133},"15588":{"m":134,"g":133},"15587":{"m":134,"g":133},"15565":{"m":134,"g":133},"15581":{"m":134,"g":133},"15585":{"m":134,"g":133},"15230":{"m":134,"g":133},"15427":{"m":134,"g":133},"15553":{"m":134,"g":133},"15583":{"m":134,"g":133},"15580":{"m":134,"g":133},"15569":{"m":134,"g":133},"14901":{"m":134,"g":133},"15564":{"m":134,"g":133},"15578":{"m":134,"g":133},"15579":{"m":134,"g":133},"15509":{"m":134,"g":133},"15540":{"m":134,"g":133},"15568":{"m":134,"g":133},"15558":{"m":134,"g":133},"12162":{"m":134,"g":133},"15531":{"m":134,"g":133},"15111":{"m":134,"g":133},"15432":{"m":134,"g":133},"15554":{"m":134,"g":133},"15556":{"m":134,"g":133},"15537":{"m":134,"g":133},"15520":{"m":134,"g":133},"15447":{"m":134,"g":133},"15324":{"m":134,"g":133},"15436":{"m":134,"g":133},"14134":{"m":134,"g":133},"15552":{"m":134,"g":133},"15526":{"m":134,"g":133},"15547":{"m":134,"g":133},"15413":{"m":134,"g":133},"15515":{"m":134,"g":133},"15544":{"m":134,"g":133},"15542":{"m":134,"g":133},"15418":{"m":134,"g":133},"15479":{"m":134,"g":133},"15533":{"m":134,"g":133},"15534":{"m":134,"g":133},"15530":{"m":134,"g":133},"15511":{"m":134,"g":133},"15536":{"m":134,"g":133},"15320":{"m":134,"g":133},"14091":{"m":134,"g":133},"15464":{"m":134,"g":133},"15484":{"m":134,"g":133},"15172":{"m":134,"g":133},"15178":{"m":134,"g":133},"15498":{"m":134,"g":133},"15333":{"m":134,"g":133},"15267":{"m":134,"g":133},"15507":{"m":134,"g":133},"15510":{"m":134,"g":133},"15485":{"m":134,"g":133},"15473":{"m":134,"g":133},"15505":{"m":134,"g":133},"15504":{"m":134,"g":133},"15503":{"m":134,"g":133},"15497":{"m":134,"g":133},"15040":{"m":134,"g":133},"15296":{"m":134,"g":133},"15496":{"m":134,"g":133},"15495":{"m":134,"g":133},"15494":{"m":134,"g":133},"15491":{"m":134,"g":133},"15022":{"m":134,"g":133},"14164":{"m":134,"g":133},"15348":{"m":134,"g":133},"15406":{"m":134,"g":133},"15483":{"m":134,"g":133},"15166":{"m":134,"g":133},"14138":{"m":134,"g":133},"15408":{"m":134,"g":133},"15416":{"m":134,"g":133},"15219":{"m":134,"g":133},"15478":{"m":134,"g":133},"15382":{"m":134,"g":133},"12995":{"m":134,"g":133},"15394":{"m":134,"g":133},"15458":{"m":134,"g":133},"15262":{"m":134,"g":133},"15437":{"m":134,"g":133},"14395":{"m":134,"g":133},"15253":{"m":134,"g":133},"15415":{"m":134,"g":133},"15433":{"m":134,"g":133},"13402":{"m":134,"g":133},"13760":{"m":134,"g":133},"15340":{"m":134,"g":133},"13782":{"m":134,"g":133},"15423":{"m":134,"g":133},"15425":{"m":134,"g":133},"15287":{"m":134,"g":133},"15207":{"m":134,"g":133},"15410":{"m":134,"g":133},"15371":{"m":134,"g":133},"15429":{"m":134,"g":133},"15431":{"m":134,"g":133},"14723":{"m":134,"g":133},"14843":{"m":134,"g":133},"14353":{"m":134,"g":133},"15424":{"m":134,"g":133},"14354":{"m":134,"g":133},"15318":{"m":134,"g":133},"14781":{"m":134,"g":133},"15421":{"m":134,"g":133},"15352":{"m":134,"g":133},"15395":{"m":134,"g":133},"15401":{"m":134,"g":133},"15407":{"m":134,"g":133},"15400":{"m":134,"g":133},"15396":{"m":134,"g":133},"15404":{"m":134,"g":133},"15397":{"m":134,"g":133},"12921":{"m":134,"g":133},"15298":{"m":134,"g":133},"15141":{"m":134,"g":133},"15379":{"m":134,"g":133},"15306":{"m":134,"g":133},"14270":{"m":134,"g":133},"15384":{"m":134,"g":133},"15361":{"m":134,"g":133},"15372":{"m":134,"g":133},"12967":{"m":134,"g":133},"15290":{"m":134,"g":133},"14501":{"m":134,"g":133},"15049":{"m":134,"g":133},"15354":{"m":134,"g":133},"15337":{"m":134,"g":133},"15278":{"m":134,"g":133},"15131":{"m":134,"g":133},"15205":{"m":134,"g":133},"15307":{"m":134,"g":133},"14860":{"m":134,"g":133},"15176":{"m":134,"g":133},"15277":{"m":134,"g":133},"15328":{"m":134,"g":133},"15120":{"m":134,"g":133},"15241":{"m":134,"g":133},"15308":{"m":134,"g":133},"15336":{"m":134,"g":133},"15316":{"m":134,"g":133},"15335":{"m":134,"g":133},"15329":{"m":134,"g":133},"15186":{"m":134,"g":133},"15222":{"m":134,"g":133},"15198":{"m":134,"g":133},"15189":{"m":134,"g":133},"15326":{"m":134,"g":133},"15237":{"m":134,"g":133},"15138":{"m":134,"g":133},"14918":{"m":134,"g":133},"11914":{"m":134,"g":133},"15304":{"m":134,"g":133},"15223":{"m":134,"g":133},"15233":{"m":134,"g":133},"15314":{"m":134,"g":133},"12333":{"m":134,"g":133},"15071":{"m":134,"g":133},"15284":{"m":134,"g":133},"14449":{"m":134,"g":133},"14357":{"m":134,"g":133},"15088":{"m":134,"g":133},"14376":{"m":134,"g":133},"15155":{"m":134,"g":133},"13571":{"m":134,"g":133},"15283":{"m":134,"g":133},"15218":{"m":134,"g":133},"15297":{"m":134,"g":133},"15258":{"m":134,"g":133},"15280":{"m":134,"g":133},"15293":{"m":134,"g":133},"15292":{"m":134,"g":133},"15291":{"m":134,"g":133},"15282":{"m":134,"g":133},"15242":{"m":134,"g":133},"14997":{"m":134,"g":133},"14934":{"m":134,"g":133},"15281":{"m":134,"g":133},"14857":{"m":134,"g":133},"14975":{"m":134,"g":133},"14936":{"m":134,"g":133},"15234":{"m":134,"g":133},"15270":{"m":134,"g":133},"15232":{"m":134,"g":133},"15192":{"m":134,"g":133},"15100":{"m":134,"g":133},"9337":{"m":134,"g":133},"15239":{"m":134,"g":133},"15220":{"m":134,"g":133},"14866":{"m":134,"g":133},"14415":{"m":134,"g":133},"15180":{"m":134,"g":133},"15225":{"m":134,"g":133},"15162":{"m":134,"g":133},"14990":{"m":134,"g":133},"15229":{"m":134,"g":133},"15231":{"m":134,"g":133},"15204":{"m":134,"g":133},"15092":{"m":134,"g":133},"15224":{"m":134,"g":133},"13410":{"m":134,"g":133},"15212":{"m":134,"g":133},"15185":{"m":134,"g":133},"15153":{"m":134,"g":133},"14820":{"m":134,"g":133},"15201":{"m":134,"g":133},"15127":{"m":134,"g":133},"15191":{"m":134,"g":133},"15190":{"m":134,"g":133},"15196":{"m":134,"g":133},"15193":{"m":134,"g":133},"15160":{"m":134,"g":133},"15047":{"m":134,"g":133},"15017":{"m":134,"g":133},"15163":{"m":134,"g":133},"15152":{"m":134,"g":133},"14906":{"m":134,"g":133},"15116":{"m":134,"g":133},"14862":{"m":134,"g":133},"15174":{"m":134,"g":133},"9324":{"m":134,"g":133},"15170":{"m":134,"g":133},"15005":{"m":134,"g":133},"13914":{"m":134,"g":133},"15158":{"m":134,"g":133},"15156":{"m":134,"g":133},"15154":{"m":134,"g":133},"15144":{"m":134,"g":133},"14778":{"m":134,"g":133},"15147":{"m":134,"g":133},"15146":{"m":134,"g":133},"15130":{"m":134,"g":133},"14907":{"m":134,"g":133},"14764":{"m":134,"g":133},"14792":{"m":134,"g":133},"15142":{"m":134,"g":133},"15139":{"m":134,"g":133},"15101":{"m":134,"g":133},"14938":{"m":134,"g":133},"15134":{"m":134,"g":133},"15058":{"m":134,"g":133},"15086":{"m":134,"g":133},"15099":{"m":134,"g":133},"15098":{"m":134,"g":133},"14953":{"m":134,"g":133},"15052":{"m":134,"g":133},"15044":{"m":134,"g":133},"15125":{"m":134,"g":133},"15110":{"m":134,"g":133},"14791":{"m":134,"g":133},"15113":{"m":134,"g":133},"15117":{"m":134,"g":133},"15124":{"m":134,"g":133},"15121":{"m":134,"g":133},"14874":{"m":134,"g":133},"15106":{"m":134,"g":133},"15090":{"m":134,"g":133},"15114":{"m":134,"g":133},"12263":{"m":134,"g":133},"13969":{"m":134,"g":133},"14961":{"m":134,"g":133},"15062":{"m":134,"g":133},"15053":{"m":134,"g":133},"14881":{"m":134,"g":133},"15108":{"m":134,"g":133},"15048":{"m":134,"g":133},"14294":{"m":134,"g":133},"15084":{"m":134,"g":133},"14969":{"m":134,"g":133},"15096":{"m":134,"g":133},"15061":{"m":134,"g":133},"15095":{"m":134,"g":133},"15094":{"m":134,"g":133},"15093":{"m":134,"g":133},"14943":{"m":134,"g":133},"15087":{"m":134,"g":133},"14993":{"m":134,"g":133},"14956":{"m":134,"g":133},"15056":{"m":134,"g":133},"15066":{"m":134,"g":133},"15085":{"m":134,"g":133},"15064":{"m":134,"g":133},"14935":{"m":134,"g":133},"15065":{"m":134,"g":133},"14998":{"m":134,"g":133},"14201":{"m":134,"g":133},"14940":{"m":134,"g":133},"15054":{"m":134,"g":133},"15055":{"m":134,"g":133},"15080":{"m":134,"g":133},"15079":{"m":134,"g":133},"14742":{"m":134,"g":133},"15074":{"m":134,"g":133},"15072":{"m":134,"g":133},"13740":{"m":134,"g":133},"15069":{"m":134,"g":133},"15060":{"m":134,"g":133},"15059":{"m":134,"g":133},"14422":{"m":134,"g":133},"14855":{"m":134,"g":133},"14423":{"m":134,"g":133},"15027":{"m":134,"g":133},"13989":{"m":134,"g":133},"14924":{"m":134,"g":133},"14659":{"m":134,"g":133},"15002":{"m":134,"g":133},"15034":{"m":134,"g":133},"14485":{"m":134,"g":133},"13876":{"m":134,"g":133},"15037":{"m":134,"g":133},"15036":{"m":134,"g":133},"15032":{"m":134,"g":133},"15031":{"m":134,"g":133},"15030":{"m":134,"g":133},"15028":{"m":134,"g":133},"15035":{"m":134,"g":133},"14939":{"m":134,"g":133},"15024":{"m":134,"g":133},"15033":{"m":134,"g":133},"15009":{"m":134,"g":133},"14957":{"m":134,"g":133},"14955":{"m":134,"g":133},"15015":{"m":134,"g":133},"15023":{"m":134,"g":133},"15021":{"m":134,"g":133},"14992":{"m":134,"g":133},"14976":{"m":134,"g":133},"14966":{"m":134,"g":133},"14933":{"m":134,"g":133},"14795":{"m":134,"g":133},"15020":{"m":134,"g":133},"15010":{"m":134,"g":133},"15014":{"m":134,"g":133},"14869":{"m":134,"g":133},"14989":{"m":134,"g":133},"14958":{"m":134,"g":133},"14951":{"m":134,"g":133},"15004":{"m":134,"g":133},"15001":{"m":134,"g":133},"15000":{"m":134,"g":133},"14999":{"m":134,"g":133},"14996":{"m":134,"g":133},"14995":{"m":134,"g":133},"14987":{"m":134,"g":133},"14945":{"m":134,"g":133},"14985":{"m":134,"g":133},"14534":{"m":134,"g":133},"14916":{"m":134,"g":133},"14959":{"m":134,"g":133},"14937":{"m":134,"g":133},"13798":{"m":134,"g":133},"14572":{"m":134,"g":133},"14949":{"m":134,"g":133},"14842":{"m":134,"g":133},"13699":{"m":134,"g":133},"14944":{"m":134,"g":133},"14888":{"m":134,"g":133},"14893":{"m":134,"g":133},"14892":{"m":134,"g":133},"14941":{"m":134,"g":133},"14894":{"m":134,"g":133},"14839":{"m":134,"g":133},"11852":{"m":134,"g":133},"14870":{"m":134,"g":133},"14358":{"m":134,"g":133},"14923":{"m":134,"g":133},"14891":{"m":134,"g":133},"14909":{"m":134,"g":133},"14467":{"m":134,"g":133},"14927":{"m":134,"g":133},"14931":{"m":134,"g":133},"13730":{"m":134,"g":133},"14932":{"m":134,"g":133},"14748":{"m":134,"g":133},"14525":{"m":134,"g":133},"14771":{"m":134,"g":133},"14074":{"m":134,"g":133},"14922":{"m":134,"g":133},"13125":{"m":134,"g":133},"14921":{"m":134,"g":133},"14928":{"m":134,"g":133},"14925":{"m":134,"g":133},"17458":"m136","17569":"m136","17591":"m136","17553":"m136","15859":"m136","17460":"m136","17474":"m136","17541":"m136","17518":"m136","17539":"m136","16366":"m136","17514":"m136","17536":"m136","17529":"m136","17534":"m136","17442":"m136","17493":"m136","17528":"m136","17524":"m136","16927":"m136","17486":"m136","17519":"m136","17416":"m136","17490":"m136","17517":"m136","17394":"m136","17498":"m136","16034":"m136","17510":"m136","17166":"m136","16919":"m136","17108":"m136","17417":"m136","16670":"m136","17355":"m136","17399":"m136","17400":"m136","17397":"m136","17372":"m136","16396":"m136","17457":"m136","17462":"m136","17466":"m136","17425":"m136","17465":"m136","17452":"m136","17386":"m136","17293":"m136","17455":"m136","17251":"m136","17444":"m136","17443":"m136","17439":"m136","17290":"m136","17291":"m136","17313":"m136","17436":"m136","17428":"m136","17289":"m136","17334":"m136","17403":"m136","17429":"m136","17247":"m136","17205":"m136","17288":"m136","17419":"m136","17414":"m136","17409":"m136","17358":"m136","11657":"m136","17382":"m136","17160":"m136","17179":"m136","17327":"m136","17385":"m136","17302":"m136","17339":"m136","15325":"m136","17364":"m136","16880":"m136","17378":"m136","17376":"m136","17370":"m136","17043":"m136","17088":"m136","17336":"m136","17367":"m136","17366":"m136","17363":"m136","17329":"m136","17049":"m136","17345":"m136","17116":"m136","17142":"m136","17220":"m136","16567":"m136","17305":"m136","17309":"m136","17158":"m136","17177":"m136","17332":"m136","17238":"m136","17245":"m136","17241":"m136","16412":"m136","17325":"m136","16744":"m136","17317":"m136","17319":"m136","16961":"m136","15347":"m136","17315":"m136","17306":"m136","15512":"m136","17308":"m136","17235":"m136","15631":"m136","14197":"m136","16649":"m136","17296":"m136","17212":"m136","16534":"m136","17295":"m136","17236":"m136","17281":"m136","16974":"m136","17287":"m136","17182":"m136","17264":"m136","16152":"m136","17261":"m136","17250":"m136","17234":"m136","17225":"m136","17256":"m136","17257":"m136","17048":"m136","14883":"m136","17191":"m136","17252":"m136","16561":"m136","17248":"m136","17249":"m136","16879":"m136","17045":"m136","17047":"m136","17044":"m136","16951":"m136","17242":"m136","16354":"m136","16817":"m136","16824":"m136","17232":"m136","17230":"m136","17233":"m136","17165":"m136","16925":"m136","17187":"m136","17217":"m136","16842":"m136","13672":"m136","15551":"m136","17133":"m136","17016":"m136","17200":"m136","17038":"m136","13216":"m136","17143":"m136","17105":"m136","17184":"m136","17173":"m136","17051":"m136","16934":"m136","17186":"m136","16121":"m136","17180":"m136","15513":"m136","17041":"m136","14565":"m136","17100":"m136","14579":"m136","17178":"m136","16882":"m136","16826":"m136","16369":"m136","17176":"m136","17174":"m136","11028":"m136","15455":"m136","17170":"m136","10598":"m136","17091":"m136","17168":"m136","17167":"m136","16568":"m136","15789":"m136","17163":"m136","17099":"m136","17126":"m136","16965":"m136","17111":"m136","17113":"m136","17020":"m136","17064":"m136","17103":"m136","17092":"m136","16949":"m136","17056":"m136","17075":"m136","17061":"m136","17028":"m136","17101":"m136","17052":"m136","12497":"m136","14504":"m136","17087":"m136","16278":"m136","16971":"m136","16976":"m136","16403":"m136","17058":"m136","17077":"m136","16264":"m136","7392":"m136","16732":"m136","16899":"m136","16569":"m136","17013":"m136","16941":"m136","16962":"m136","16888":"m136","15908":"m136","16989":"m136","17046":"m136","17030":"m136","17027":"m136","17054":"m136","16898":"m136","16994":"m136","16841":"m136","17002":"m136","14108":"m136","17022":"m136","17005":"m136","16982":"m136","16894":"m136","16953":"m136","17019":"m136","16886":"m136","16259":"m136","16924":"m136","16986":"m136","16790":"m136","16935":"m136","16884":"m136","16967":"m136","16766":"m136","16564":"m136","16767":"m136","16648":"m136","16252":"m136","16940":"m136","15853":"m136","16765":"m136","16978":"m136","16981":"m136","16980":"m136","16979":"m136","16977":"m136","16973":"m136","16970":"m136","16192":"m136","16348":"m136","16963":"m136","15271":"m136","16933":"m136","16480":"m136","16922":"m136","16916":"m136","16912":"m136","16019":"m136","16559":"m136","16677":"m136","16930":"m136","16932":"m136","16536":"m136","16850":"m136","16915":"m136","16896":"m136","16820":"m136","15268":"m136","12909":"m136","16908":"m136","16906":"m136","16851":"m136","14867":"m136","16895":"m136","16300":"m136","16864":"m136","16889":"m136","15182":"m136","16876":"m136","16258":"m136","16835":"m136","16345":"m136","16878":"m136","16867":"m136","16273":"m136","16757":"m136","16865":"m136","16877":"m136","16863":"m136","16667":"m136","16788":"m136","16397":"m136","16737":"m136","16872":"m136","16871":"m136","16870":"m136","15790":"m136","15927":"m136","15227":"m136","16226":"m136","16844":"m136","16838":"m136","16849":"m136","16333":"m136","16779":"m136","16854":"m136","16862":"m136","16572":"m136","16805":"m136","16723":"m136","13715":"m136","16853":"m136","16637":"m136","16852":"m136","16847":"m136","16848":"m136","16845":"m136","16698":"m136","16825":"m136","16837":"m136","16840":"m136","16839":"m136","16831":"m136","16830":"m136","16679":"m136","16810":"m136","16783":"m136","16587":"m136","16821":"m136","16804":"m136","15753":"m136","13947":"m136","16275":"m136","16814":"m136","16813":"m136","16812":"m136","16811":"m136","14655":"m136","16325":"m136","16708":"m136","16760":"m136","16792":"m136","11349":"m136","16458":"m136","16743":"m136","16768":"m136","16721":"m136","16778":"m136","16774":"m136","16686":"m136","16378":"m136","16588":"m136","16446":"m136","16773":"m136","16380":"m136","16254":{"m":136,"g":135},"16720":{"m":136,"g":135},"16560":{"m":136,"g":135},"16764":{"m":136,"g":135},"16763":{"m":136,"g":135},"16759":{"m":136,"g":135},"16715":{"m":136,"g":135},"16756":{"m":136,"g":135},"16754":{"m":136,"g":135},"16752":{"m":136,"g":135},"16751":{"m":136,"g":135},"16749":{"m":136,"g":135},"16748":{"m":136,"g":135},"16675":{"m":136,"g":135},"16746":{"m":136,"g":135},"16409":{"m":136,"g":135},"13681":{"m":136,"g":135},"16745":{"m":136,"g":135},"16741":{"m":136,"g":135},"16014":{"m":136,"g":135},"16739":{"m":136,"g":135},"16719":{"m":136,"g":135},"16668":{"m":136,"g":135},"16738":{"m":136,"g":135},"16709":{"m":136,"g":135},"16735":{"m":136,"g":135},"16622":{"m":136,"g":135},"16635":{"m":136,"g":135},"16733":{"m":136,"g":135},"16730":{"m":136,"g":135},"16729":{"m":136,"g":135},"16519":{"m":136,"g":135},"16706":{"m":136,"g":135},"16115":{"m":136,"g":135},"16634":{"m":136,"g":135},"16535":{"m":136,"g":135},"16452":{"m":136,"g":135},"16533":{"m":136,"g":135},"16531":{"m":136,"g":135},"16529":{"m":136,"g":135},"15627":{"m":136,"g":135},"16608":{"m":136,"g":135},"16693":{"m":136,"g":135},"16155":{"m":136,"g":135},"16505":{"m":136,"g":135},"16697":{"m":136,"g":135},"16555":{"m":136,"g":135},"16669":{"m":136,"g":135},"16625":{"m":136,"g":135},"16695":{"m":136,"g":135},"16678":{"m":136,"g":135},"16692":{"m":136,"g":135},"16629":{"m":136,"g":135},"16680":{"m":136,"g":135},"16445":{"m":136,"g":135},"16306":{"m":136,"g":135},"16652":{"m":136,"g":135},"15343":{"m":136,"g":135},"16095":{"m":136,"g":135},"15151":{"m":136,"g":135},"16681":{"m":136,"g":135},"16674":{"m":136,"g":135},"16457":{"m":136,"g":135},"16426":{"m":136,"g":135},"16676":{"m":136,"g":135},"16658":{"m":136,"g":135},"16672":{"m":136,"g":135},"16633":{"m":136,"g":135},"16671":{"m":136,"g":135},"15712":{"m":136,"g":135},"16660":{"m":136,"g":135},"16664":{"m":136,"g":135},"16657":{"m":136,"g":135},"16661":{"m":136,"g":135},"16618":{"m":136,"g":135},"16179":{"m":136,"g":135},"16599":{"m":136,"g":135},"15238":{"m":136,"g":135},"16532":{"m":136,"g":135},"16654":{"m":136,"g":135},"16651":{"m":136,"g":135},"16203":{"m":136,"g":135},"13602":{"m":136,"g":135},"16620":{"m":136,"g":135},"16514":{"m":136,"g":135},"16631":{"m":136,"g":135},"16630":{"m":136,"g":135},"16619":{"m":136,"g":135},"16527":{"m":136,"g":135},"15938":{"m":136,"g":135},"16582":{"m":136,"g":135},"16418":{"m":136,"g":135},"16459":{"m":136,"g":135},"16465":{"m":136,"g":135},"16468":{"m":136,"g":135},"16566":{"m":136,"g":135},"16617":{"m":136,"g":135},"16597":{"m":136,"g":135},"16504":{"m":136,"g":135},"16593":{"m":136,"g":135},"16118":{"m":136,"g":135},"15439":{"m":136,"g":135},"16609":{"m":136,"g":135},"16453":{"m":136,"g":135},"16499":{"m":136,"g":135},"16606":{"m":136,"g":135},"16603":{"m":136,"g":135},"16576":{"m":136,"g":135},"16596":{"m":136,"g":135},"16598":{"m":136,"g":135},"16601":{"m":136,"g":135},"16594":{"m":136,"g":135},"16442":{"m":136,"g":135},"16422":{"m":136,"g":135},"16223":{"m":136,"g":135},"16425":{"m":136,"g":135},"16589":{"m":136,"g":135},"16592":{"m":136,"g":135},"16583":{"m":136,"g":135},"16585":{"m":136,"g":135},"16591":{"m":136,"g":135},"16326":{"m":136,"g":135},"16523":{"m":136,"g":135},"16420":{"m":136,"g":135},"16548":{"m":136,"g":135},"16549":{"m":136,"g":135},"16575":{"m":136,"g":135},"16507":{"m":136,"g":135},"16570":{"m":136,"g":135},"16463":{"m":136,"g":135},"14215":{"m":136,"g":135},"13518":{"m":136,"g":135},"16455":{"m":136,"g":135},"16496":{"m":136,"g":135},"16540":{"m":136,"g":135},"16201":{"m":136,"g":135},"15456":{"m":136,"g":135},"16421":{"m":136,"g":135},"16466":{"m":136,"g":135},"16539":{"m":136,"g":135},"16417":{"m":136,"g":135},"16502":{"m":136,"g":135},"16415":{"m":136,"g":135},"16467":{"m":136,"g":135},"16399":{"m":136,"g":135},"14112":{"m":136,"g":135},"15492":{"m":136,"g":135},"16538":{"m":136,"g":135},"6135":{"m":136,"g":135},"16528":{"m":136,"g":135},"16416":{"m":136,"g":135},"16419":{"m":136,"g":135},"16434":{"m":136,"g":135},"16525":{"m":136,"g":135},"16520":{"m":136,"g":135},"16086":{"m":136,"g":135},"16524":{"m":136,"g":135},"15677":{"m":136,"g":135},"16477":{"m":136,"g":135},"16356":{"m":136,"g":135},"16513":{"m":136,"g":135},"15663":{"m":136,"g":135},"16516":{"m":136,"g":135},"16482":{"m":136,"g":135},"16511":{"m":136,"g":135},"16280":{"m":136,"g":135},"16509":{"m":136,"g":135},"16508":{"m":136,"g":135},"16492":{"m":136,"g":135},"16469":{"m":136,"g":135},"16138":{"m":136,"g":135},"16481":{"m":136,"g":135},"14474":{"m":136,"g":135},"16367":{"m":136,"g":135},"16475":{"m":136,"g":135},"16478":{"m":136,"g":135},"16479":{"m":136,"g":135},"16474":{"m":136,"g":135},"16447":{"m":136,"g":135},"16471":{"m":136,"g":135},"16456":{"m":136,"g":135},"16382":{"m":136,"g":135},"14051":{"m":136,"g":135},"16042":{"m":136,"g":135},"16424":{"m":136,"g":135},"16460":{"m":136,"g":135},"16464":{"m":136,"g":135},"16454":{"m":136,"g":135},"16088":{"m":136,"g":135},"16387":{"m":136,"g":135},"16386":{"m":136,"g":135},"16451":{"m":136,"g":135},"16450":{"m":136,"g":135},"16448":{"m":136,"g":135},"16389":{"m":136,"g":135},"16180":{"m":136,"g":135},"16441":{"m":136,"g":135},"16444":{"m":136,"g":135},"16437":{"m":136,"g":135},"16310":{"m":136,"g":135},"16435":{"m":136,"g":135},"16433":{"m":136,"g":135},"16432":{"m":136,"g":135},"16375":{"m":136,"g":135},"16330":{"m":136,"g":135},"15836":{"m":136,"g":135},"16430":{"m":136,"g":135},"16414":{"m":136,"g":135},"16429":{"m":136,"g":135},"16427":{"m":136,"g":135},"16413":{"m":136,"g":135},"16408":{"m":136,"g":135},"16411":{"m":136,"g":135},"16334":{"m":136,"g":135},"16127":{"m":136,"g":135},"16323":{"m":136,"g":135},"14066":{"m":136,"g":135},"16405":{"m":136,"g":135},"16406":{"m":136,"g":135},"16401":{"m":136,"g":135},"16400":{"m":136,"g":135},"16347":{"m":136,"g":135},"16349":{"m":136,"g":135},"16390":{"m":136,"g":135},"16392":{"m":136,"g":135},"16374":{"m":136,"g":135},"16391":{"m":136,"g":135},"16381":{"m":136,"g":135},"16376":{"m":136,"g":135},"16373":{"m":136,"g":135},"16359":{"m":136,"g":135},"16268":{"m":136,"g":135},"16365":{"m":136,"g":135},"16368":{"m":136,"g":135},"16363":{"m":136,"g":135},"16358":{"m":136,"g":135},"16343":{"m":136,"g":135},"16357":{"m":136,"g":135},"16355":{"m":136,"g":135},"16352":{"m":136,"g":135},"16064":{"m":136,"g":135},"15434":{"m":136,"g":135},"16353":{"m":136,"g":135},"16341":{"m":136,"g":135},"16351":{"m":136,"g":135},"16269":{"m":136,"g":135},"16350":{"m":136,"g":135},"16340":{"m":136,"g":135},"16339":{"m":136,"g":135},"16335":{"m":136,"g":135},"16344":{"m":136,"g":135},"15941":{"m":136,"g":135},"16342":{"m":136,"g":135},"16164":{"m":136,"g":135},"16304":{"m":136,"g":135},"16303":{"m":136,"g":135},"16338":{"m":136,"g":135},"16219":{"m":136,"g":135},"16337":{"m":136,"g":135},"15640":{"m":136,"g":135},"16311":{"m":136,"g":135},"16332":{"m":136,"g":135},"15175":{"m":136,"g":135},"16328":{"m":136,"g":135},"16009":{"m":136,"g":135},"13592":{"m":136,"g":135},"16272":{"m":136,"g":135},"16283":{"m":136,"g":135},"16257":{"m":136,"g":135},"16177":{"m":136,"g":135},"15560":{"m":136,"g":135},"16324":{"m":136,"g":135},"16317":{"m":136,"g":135},"16296":{"m":136,"g":135},"16321":{"m":136,"g":135},"16316":{"m":136,"g":135},"16313":{"m":136,"g":135},"16318":{"m":136,"g":135},"16320":{"m":136,"g":135},"16319":{"m":136,"g":135},"16314":{"m":136,"g":135},"16315":{"m":136,"g":135},"16312":{"m":136,"g":135},"16287":{"m":136,"g":135},"16277":{"m":136,"g":135},"16308":{"m":136,"g":135},"16092":{"m":136,"g":135},"16128":{"m":136,"g":135},"16305":{"m":136,"g":135},"16301":{"m":136,"g":135},"13959":{"m":136,"g":135},"16248":{"m":136,"g":135},"16292":{"m":136,"g":135},"16298":{"m":136,"g":135},"16227":{"m":136,"g":135},"16213":{"m":136,"g":135},"16267":{"m":136,"g":135},"16144":{"m":136,"g":135},"16222":{"m":136,"g":135},"15345":{"m":136,"g":135},"16270":{"m":136,"g":135},"16285":{"m":136,"g":135},"14636":{"m":136,"g":135},"16263":{"m":136,"g":135},"16265":{"m":136,"g":135},"15995":{"m":136,"g":135},"16162":{"m":136,"g":135},"16262":{"m":136,"g":135},"16261":{"m":136,"g":135},"16260":{"m":136,"g":135},"16251":{"m":136,"g":135},"16247":{"m":136,"g":135},"16250":{"m":136,"g":135},"16243":{"m":136,"g":135},"14085":{"m":136,"g":135},"16171":{"m":136,"g":135},"16245":{"m":136,"g":135},"16238":{"m":136,"g":135},"15814":{"m":136,"g":135},"16239":{"m":136,"g":135},"16240":{"m":136,"g":135},"16236":{"m":136,"g":135},"16233":{"m":136,"g":135},"16202":{"m":136,"g":135},"16228":{"m":136,"g":135},"16018":{"m":136,"g":135},"16195":{"m":136,"g":135},"16111":{"m":136,"g":135},"13394":{"m":136,"g":135},"16204":{"m":136,"g":135},"16214":{"m":136,"g":135},"16221":{"m":136,"g":135},"16178":{"m":136,"g":135},"16187":{"m":136,"g":135},"16209":{"m":136,"g":135},"15597":{"m":136,"g":135},"16117":{"m":136,"g":135},"15878":{"m":136,"g":135},"16200":{"m":136,"g":135},"15946":{"m":136,"g":135},"16198":{"m":136,"g":135},"16172":{"m":136,"g":135},"15575":{"m":136,"g":135},"16174":{"m":136,"g":135},"16156":{"m":136,"g":135},"15754":{"m":136,"g":135},"14920":{"m":136,"g":135},"16050":{"m":136,"g":135},"16181":{"m":136,"g":135},"16183":{"m":136,"g":135},"16182":{"m":136,"g":135},"14416":{"m":136,"g":135},"16168":{"m":136,"g":135},"16175":{"m":136,"g":135},"16166":{"m":136,"g":135},"16161":{"m":136,"g":135},"16163":{"m":136,"g":135},"16165":{"m":136,"g":135},"16110":{"m":136,"g":135},"16173":{"m":136,"g":135},"16159":{"m":136,"g":135},"16150":{"m":136,"g":135},"16160":{"m":136,"g":135},"16158":{"m":136,"g":135},"16149":{"m":136,"g":135},"16124":{"m":136,"g":135},"15730":{"m":136,"g":135},"13978":{"m":136,"g":135},"12625":{"m":136,"g":135},"16154":{"m":136,"g":135},"18298":"m137","18111":"m137"},"c":{"2ccd9fd8":{"m":0,"g":94},"46b7ea7c":{"m":0,"g":94},"70359bf3":{"m":0,"g":94},"01ca82d7":{"m":0,"g":94},"4bd8233f":{"m":0,"g":94},"08ab2a16":{"m":0,"g":94},"f652494d":{"m":0,"g":94},"30720e73":{"m":0,"g":94},"331848de":{"m":0,"g":94},"93eeb543":{"m":0,"g":94},"ead5b39f":{"m":0,"g":94},"22085081":{"m":0,"g":94},"f6d40df0":{"m":0,"g":94},"22ec7bc2":{"m":1,"g":94},"c0454b32":{"m":1,"g":94},"8024fc5e":{"m":1,"g":94},"70528762":{"m":1,"g":94},"71d30d6d":{"m":1,"g":94},"f9d72381":{"m":1,"g":94},"bf51ddc6":{"m":1,"g":94},"fd7c4792":{"m":1,"g":94},"c4707f1b":{"m":1,"g":94},"ffe4aaee":{"m":1,"g":94},"5b27a1dc":{"m":1,"g":94},"e71d4ab3":{"m":1,"g":94},"fbf42263":{"m":1,"g":94},"cc3ada98":{"m":2,"g":94},"a837166e":{"m":2,"g":94},"11f3cca6":{"m":2,"g":94},"ca13f3b8":{"m":2,"g":94},"0b2efc2a":{"m":2,"g":94},"f30abd09":{"m":2,"g":94},"40ab1f01":{"m":2,"g":94},"199e82a1":{"m":2,"g":94},"23471f9a":{"m":2,"g":94},"61d4c939":{"m":2,"g":94},"98a3e8ef":{"m":2,"g":94},"2b079f89":{"m":2,"g":94},"05b4c398":{"m":2,"g":94},"dafafe5b":{"m":2,"g":94},"b240f751":{"m":2,"g":94},"501f9444":{"m":2,"g":94},"723f0421":{"m":3,"g":94},"585eabab":{"m":3,"g":94},"c70b3cfa":{"m":4,"g":94},"489796c7":{"m":4,"g":94},"fa7a696d":{"m":4,"g":94},"bef0b359":{"m":4,"g":94},"c6576e82":{"m":4,"g":94},"99258181":{"m":4,"g":94},"3de54a1b":{"m":4,"g":94},"7358fa64":{"m":4,"g":94},"9a16fea0":{"m":4,"g":94},"9e037c82":{"m":4,"g":94},"9076386d":{"m":4,"g":94},"959c4174":{"m":4,"g":94},"94e05770":{"m":4,"g":94},"63e97e5e":{"m":4,"g":94},"e08bca28":{"m":4,"g":94},"cd3ccb2e":{"m":4,"g":94},"3f5c2f4c":{"m":4,"g":94},"007eeb4e":{"m":4,"g":94},"e8f2b155":{"m":4,"g":94},"6dceab4d":{"m":5,"g":94},"a49dc52b":{"m":6,"g":94},"873d0e85":{"m":6,"g":94},"1d0fbe8e":{"m":6,"g":94},"97aa9b32":{"m":6,"g":94},"06175286":{"m":6,"g":94},"4ea92f83":{"m":6,"g":94},"6b0af285":{"m":6,"g":94},"6f560c76":{"m":6,"g":94},"cd687233":{"m":6,"g":94},"81561f8e":{"m":6,"g":94},"3a581e99":{"m":6,"g":94},"0147f940":{"m":6,"g":94},"23950056":{"m":6,"g":94},"93414c82":{"m":6,"g":94},"ed7c7eca":{"m":6,"g":94},"0c457bae":{"m":6,"g":94},"d3fc86a4":{"m":6,"g":94},"01ee0fbc":{"m":6,"g":94},"711d3435":{"m":6,"g":94},"f6bfe3aa":{"m":7,"g":94},"e095b162":{"m":7,"g":94},"67be11c7":{"m":7,"g":94},"cd8c3ccd":{"m":7,"g":94},"9c121f2a":{"m":7,"g":94},"03e04b23":{"m":7,"g":94},"86442530":{"m":7,"g":94},"79cb018e":{"m":7,"g":94},"c7af9f73":{"m":7,"g":94},"876db8dc":{"m":7,"g":94},"ad82bac6":{"m":7,"g":94},"71b54eea":{"m":7,"g":94},"74b3bfaa":{"m":7,"g":94},"4a634cf6":{"m":7,"g":94},"624b21e7":{"m":8,"g":94},"c51020cf":{"m":8,"g":94},"50afed4e":{"m":8,"g":94},"4d303c4f":{"m":8,"g":94},"37b42297":{"m":8,"g":94},"cba50273":{"m":8,"g":94},"a6aa46dd":{"m":8,"g":94},"405f26b0":{"m":8,"g":94},"b1a3a454":{"m":8,"g":94},"79e6b84b":{"m":8,"g":94},"26c34941":{"m":8,"g":94},"cb8e1982":{"m":8,"g":94},"23f05005":{"m":8,"g":94},"a7334aee":{"m":8,"g":94},"ee1df26a":{"m":8,"g":94},"3ae78a09":{"m":8,"g":94},"ccbe1e67":{"m":8,"g":94},"e2bf732b":{"m":8,"g":94},"322421fa":{"m":8,"g":94},"8ff870bf":{"m":8,"g":94},"26f0bedc":{"m":8,"g":94},"82fa69b3":{"m":8,"g":94},"8fb7459e":{"m":8,"g":94},"bb3a3b66":{"m":8,"g":94},"45d6592d":{"m":8,"g":94},"4aa5dd2c":{"m":9,"g":94},"13662fd5":{"m":9,"g":94},"d5ae2eba":{"m":9,"g":94},"1b355479":{"m":9,"g":94},"faba293a":{"m":9,"g":94},"89885b31":{"m":9,"g":94},"64fe3115":{"m":9,"g":94},"a7ace9c8":{"m":9,"g":94},"a833de05":{"m":9,"g":94},"30d67b2b":{"m":9,"g":94},"b0b722ee":{"m":9,"g":94},"01b07ea3":{"m":9,"g":94},"dfb13ac4":{"m":9,"g":94},"ec90b9c0":{"m":9,"g":94},"9759d927":{"m":9,"g":94},"8d0a7fae":{"m":9,"g":94},"c4e9ebe3":{"m":9,"g":94},"3c2c5869":{"m":9,"g":94},"4cb9aaed":{"m":9,"g":94},"9de9a468":{"m":9,"g":94},"ce3b2610":{"m":9,"g":94},"91e03633":{"m":9,"g":94},"2a74748b":{"m":9,"g":94},"63ba630b":{"m":9,"g":94},"6493256b":{"m":9,"g":94},"06008bc2":{"m":9,"g":94},"bb824da4":{"m":9,"g":94},"93121324":{"m":9,"g":94},"c97fdae4":{"m":9,"g":94},"51104cd4":{"m":10,"g":94},"e2b2f0a2":{"m":10,"g":94},"b57abe16":{"m":10,"g":94},"e57f0792":{"m":10,"g":94},"08df63a6":{"m":10,"g":94},"77835756":{"m":10,"g":94},"ed315799":{"m":10,"g":94},"92e2d74f":{"m":10,"g":94},"d9b3b018":{"m":10,"g":94},"745ea007":{"m":10,"g":94},"ad1dd746":{"m":10,"g":94},"eb4308c4":{"m":10,"g":94},"b2eb0805":{"m":10,"g":94},"72bb3443":{"m":11,"g":94},"2d580e7a":{"m":11,"g":94},"3fc97f67":{"m":11,"g":94},"abc548c7":{"m":11,"g":94},"aee4f523":{"m":11,"g":94},"7023f413":{"m":11,"g":94},"09deb20d":{"m":11,"g":94},"33b242df":{"m":11,"g":94},"a511a2d0":{"m":11,"g":94},"6ec65f45":{"m":11,"g":94},"e2c31fca":{"m":11,"g":94},"d5de20a3":{"m":11,"g":94},"4a1c6ae2":{"m":11,"g":94},"14522e6a":{"m":11,"g":94},"183df472":{"m":11,"g":94},"5c5aba59":{"m":11,"g":94},"ba67101f":{"m":11,"g":94},"95c4e0df":{"m":11,"g":94},"19818b9c":{"m":11,"g":94},"9216b106":{"m":11,"g":94},"da19434c":{"m":11,"g":94},"150d7020":{"m":11,"g":94},"9acc6e35":{"m":11,"g":94},"cf9d8efd":{"m":11,"g":94},"1bf1cf19":{"m":11,"g":94},"e822e590":{"m":11,"g":94},"ca4f1ab8":{"m":11,"g":94},"2b6d9991":{"m":11,"g":94},"65501a9c":{"m":11,"g":94},"db611066":{"m":11,"g":94},"c93293c5":{"m":11,"g":94},"62b3812b":{"m":11,"g":94},"550a4f78":{"m":11,"g":94},"ff99c38a":{"m":11,"g":94},"c9de3e16":{"m":11,"g":94},"ed27a6b9":{"m":11,"g":94},"463c6632":{"m":11,"g":94},"b0890631":{"m":11,"g":94},"cb389c91":{"m":11,"g":94},"eddaa2b5":{"m":11,"g":94},"2af565b3":{"m":11,"g":94},"3842eba5":{"m":11,"g":94},"24e59f53":{"m":11,"g":94},"75235419":{"m":11,"g":94},"64ee9c03":{"m":11,"g":94},"30d17840":{"m":11,"g":94},"ce216c80":{"m":11,"g":94},"e0ae5d42":{"m":12,"g":94},"32de16ce":{"m":12,"g":94},"0992d85f":{"m":12,"g":94},"5dc55a5f":{"m":12,"g":94},"4231a42f":{"m":12,"g":94},"455c9ccc":{"m":12,"g":94},"39191c85":{"m":12,"g":94},"562b8857":{"m":12,"g":94},"04c0b214":{"m":12,"g":94},"6e09cf6a":{"m":12,"g":94},"e8a2327d":{"m":13,"g":94},"91f93f14":{"m":13,"g":94},"f70f7258":{"m":13,"g":94},"c0ae70c8":{"m":13,"g":94},"87260b7b":{"m":13,"g":94},"651a23ee":{"m":13,"g":94},"bf3e271f":{"m":13,"g":94},"3bc01ac1":{"m":13,"g":94},"9f009261":{"m":13,"g":94},"159cc741":{"m":13,"g":94},"7d1ebc2d":{"m":13,"g":94},"83525a1d":{"m":13,"g":94},"80a33ce8":{"m":13,"g":94},"1a57e416":{"m":13,"g":94},"adc97426":{"m":13,"g":94},"0463f7fb":{"m":13,"g":94},"565d7274":{"m":13,"g":94},"09de730d":{"m":13,"g":94},"55c16436":{"m":13,"g":94},"2b605ab1":{"m":13,"g":94},"947bda73":{"m":13,"g":94},"f06e90c2":{"m":13,"g":94},"2cea6146":{"m":13,"g":94},"44c998fc":{"m":13,"g":94},"3167d8da":{"m":13,"g":94},"0fafc560":{"m":13,"g":94},"19d2135c":{"m":13,"g":94},"ced77c66":{"m":13,"g":94},"8dbdc018":{"m":13,"g":94},"3e684be7":{"m":13,"g":94},"ec380dfd":{"m":13,"g":94},"5b647543":{"m":13,"g":94},"8210ec60":{"m":13,"g":94},"5be9eb8a":{"m":13,"g":94},"c05956e5":{"m":13,"g":94},"d75dc20f":{"m":13,"g":94},"690d162d":{"m":13,"g":94},"664287b2":{"m":13,"g":94},"2f11936f":{"m":14,"g":94},"63fbef98":{"m":14,"g":94},"2a754e57":{"m":14,"g":94},"96c503eb":{"m":14,"g":94},"441cca77":{"m":14,"g":94},"c7709d3a":{"m":14,"g":94},"9380f50f":{"m":14,"g":94},"95dc093b":{"m":14,"g":94},"d9ac6392":{"m":14,"g":94},"26294b2f":{"m":14,"g":94},"75b31a2a":{"m":14,"g":94},"11616fc6":{"m":14,"g":94},"9ce89bc1":{"m":14,"g":94},"badf3fa0":{"m":14,"g":94},"945aa9be":{"m":14,"g":94},"2e6e62e1":{"m":14,"g":94},"a385ee27":{"m":14,"g":94},"eb1ae6ae":{"m":14,"g":94},"2187f362":{"m":14,"g":94},"9465b668":{"m":14,"g":94},"05471f21":{"m":14,"g":94},"1fa15099":{"m":14,"g":94},"303ef888":{"m":14,"g":94},"92cb93f3":{"m":14,"g":94},"e94e60d6":{"m":14,"g":94},"d2f8bfb2":{"m":14,"g":94},"b7e2f800":{"m":14,"g":94},"09593e9b":{"m":14,"g":94},"53a7ebd8":{"m":14,"g":94},"ad5f04d6":{"m":14,"g":94},"bbec01c9":{"m":14,"g":94},"40e53d65":{"m":14,"g":94},"fb9296f0":{"m":14,"g":94},"1374334d":{"m":14,"g":94},"94aead9e":{"m":14,"g":94},"9c902b19":{"m":14,"g":94},"111991fe":{"m":14,"g":94},"a8c787d2":{"m":14,"g":94},"5f283991":{"m":14,"g":94},"b6667a53":{"m":14,"g":94},"542bc733":{"m":14,"g":94},"f6dbd240":{"m":14,"g":94},"ad872feb":{"m":15,"g":94},"da2e5d65":{"m":15,"g":94},"ce62dc73":{"m":15,"g":94},"02b72586":{"m":15,"g":94},"d557e9f3":{"m":15,"g":94},"740c46a1":{"m":15,"g":94},"b3868722":{"m":15,"g":94},"710f614e":{"m":15,"g":94},"f25b76c0":{"m":15,"g":94},"f4e885b7":{"m":15,"g":94},"0877f1e7":{"m":15,"g":94},"5304b4ef":{"m":15,"g":94},"26908d95":{"m":15,"g":94},"c0982ac5":{"m":15,"g":94},"dc1b8bcf":{"m":15,"g":94},"5a57b8ad":{"m":15,"g":94},"d737da5f":{"m":15,"g":94},"ac113887":{"m":15,"g":94},"dc8cef1d":{"m":15,"g":94},"5d264a90":{"m":16,"g":94},"5949b1ca":{"m":16,"g":94},"0feca02d":{"m":16,"g":94},"10143e1a":{"m":16,"g":94},"65c65776":{"m":16,"g":94},"66581596":{"m":16,"g":94},"396a6924":{"m":16,"g":94},"af4e7910":{"m":16,"g":94},"519e20cf":{"m":16,"g":94},"d9a69029":{"m":16,"g":94},"56f5fc4a":{"m":17,"g":94},"6a2941f4":{"m":17,"g":94},"5ac8b806":{"m":17,"g":94},"bae9541e":{"m":17,"g":94},"a56858ba":{"m":17,"g":94},"564a898a":{"m":17,"g":94},"2b4c6462":{"m":18,"g":94},"f424e76d":{"m":18,"g":94},"490a1f39":{"m":18,"g":94},"06487f12":{"m":18,"g":94},"39c57317":{"m":18,"g":94},"9592a1f3":{"m":18,"g":94},"35759efa":{"m":18,"g":94},"8f4b1559":{"m":18,"g":94},"e3046ea3":{"m":18,"g":94},"49c5e0ec":{"m":18,"g":94},"ec2150b2":{"m":18,"g":94},"7620cd37":{"m":18,"g":94},"50a53887":{"m":18,"g":94},"11c8efff":{"m":18,"g":94},"e87c7fd5":{"m":18,"g":94},"630479c3":{"m":18,"g":94},"51fda143":{"m":18,"g":94},"dc4e4a6a":{"m":18,"g":94},"2d96da81":{"m":18,"g":94},"c126a6cc":{"m":18,"g":94},"ac971ff6":{"m":18,"g":94},"e1792cca":{"m":18,"g":94},"1b7adbb5":{"m":18,"g":94},"a9ef49c1":{"m":18,"g":94},"21ba3a88":{"m":18,"g":94},"9c5cac24":{"m":18,"g":94},"5960a6e5":{"m":18,"g":94},"b050d928":{"m":18,"g":94},"6a4dc996":{"m":18,"g":94},"d774acad":{"m":18,"g":94},"d93388da":{"m":18,"g":94},"476584cb":{"m":18,"g":94},"abd5385a":{"m":18,"g":94},"3de2f30a":{"m":18,"g":94},"4efcc59d":{"m":18,"g":94},"2e341cd4":{"m":18,"g":94},"a8552cb1":{"m":18,"g":94},"a470e60c":{"m":18,"g":94},"5f90e076":{"m":18,"g":94},"8832ecb1":{"m":18,"g":94},"5ff60eda":{"m":18,"g":94},"c1930022":{"m":18,"g":94},"fe3be159":{"m":18,"g":94},"0aa189f1":{"m":18,"g":94},"f6b29f69":{"m":18,"g":94},"c9ee3d35":{"m":18,"g":94},"41d1f677":{"m":18,"g":94},"444a0244":{"m":19,"g":94},"fa7ccb33":{"m":19,"g":94},"26868443":{"m":19,"g":94},"824a77d0":{"m":19,"g":94},"cf99eab7":{"m":19,"g":94},"9fdea29d":{"m":19,"g":94},"df7c4c19":{"m":19,"g":94},"c3f1aac8":{"m":19,"g":94},"d198791f":{"m":19,"g":94},"c07526e4":{"m":19,"g":94},"7b597475":{"m":19,"g":94},"5303c1ed":{"m":19,"g":94},"65bd1338":{"m":19,"g":94},"eedc12e1":{"m":19,"g":94},"5a4ef2b5":{"m":19,"g":94},"9dab947d":{"m":19,"g":94},"33ee97b0":{"m":19,"g":94},"6a846bb1":{"m":19,"g":94},"0fdb3127":{"m":19,"g":94},"5ad033a0":{"m":19,"g":94},"77e592e8":{"m":19,"g":94},"caaad53b":{"m":19,"g":94},"69d19188":{"m":19,"g":94},"4b4a67f8":{"m":19,"g":94},"0ac94c36":{"m":19,"g":94},"459abad2":{"m":20,"g":94},"30d8e130":{"m":20,"g":94},"08a3bd19":{"m":20,"g":94},"321a963b":{"m":20,"g":94},"e17deb27":{"m":20,"g":94},"2d3ae4e1":{"m":20,"g":94},"75f4ccb7":{"m":20,"g":94},"83d2b30d":{"m":20,"g":94},"4367f4bb":{"m":20,"g":94},"00e4baa7":{"m":20,"g":94},"4cd64b8e":{"m":20,"g":94},"01d66ae2":{"m":20,"g":94},"a523a3c1":{"m":20,"g":94},"9f94728f":{"m":20,"g":94},"1a491d00":{"m":21,"g":94},"8fbba3de":{"m":21,"g":94},"ae0f6130":{"m":21,"g":94},"60105897":{"m":21,"g":94},"926ac01b":{"m":21,"g":94},"25c881a0":{"m":21,"g":94},"04ec6ba2":{"m":21,"g":94},"d63f13c1":{"m":21,"g":94},"fded6744":{"m":21,"g":94},"6e453940":{"m":21,"g":94},"97e0f7d2":{"m":21,"g":94},"d5146bae":{"m":21,"g":94},"5bd06b45":{"m":22,"g":94},"9a611827":{"m":22,"g":94},"eeb24821":{"m":22,"g":94},"3e455b01":{"m":22,"g":94},"8628ab9c":{"m":22,"g":94},"1b77670f":{"m":22,"g":94},"768e05d0":{"m":22,"g":94},"01fbb11b":{"m":22,"g":94},"05d216da":{"m":22,"g":94},"6b32bb1c":{"m":22,"g":94},"40facad5":{"m":22,"g":94},"da504445":{"m":22,"g":94},"252e0f7b":{"m":22,"g":94},"7f6f2f0f":{"m":22,"g":94},"7802df1e":{"m":22,"g":94},"bc1154c3":{"m":23,"g":94},"752e6430":{"m":23,"g":94},"30db99b3":{"m":23,"g":94},"0a409bd4":{"m":23,"g":94},"e4db4e5b":{"m":23,"g":94},"bbc07c41":{"m":23,"g":94},"a036d419":{"m":23,"g":94},"f95e6617":{"m":23,"g":94},"de854fb5":{"m":23,"g":94},"f64b2a9b":{"m":23,"g":94},"9f95dcc6":{"m":23,"g":94},"0736b270":{"m":23,"g":94},"3fdab919":{"m":23,"g":94},"ba29504b":{"m":23,"g":94},"a72342f1":{"m":23,"g":94},"c3c74bf8":{"m":23,"g":94},"d9fccfef":{"m":23,"g":94},"679ebcbb":{"m":23,"g":94},"1edd4e07":{"m":24,"g":94},"62c673c4":{"m":24,"g":94},"377c5dc9":{"m":24,"g":94},"f52eda35":{"m":24,"g":94},"b579ecf0":{"m":24,"g":94},"e7487b08":{"m":24,"g":94},"ae5c0fc4":{"m":24,"g":94},"a30d5d75":{"m":24,"g":94},"17af39c5":{"m":24,"g":94},"daf593a3":{"m":24,"g":94},"bece265f":{"m":24,"g":94},"cdcbde5f":{"m":24,"g":94},"21e22b9e":{"m":24,"g":94},"a50c8a14":{"m":24,"g":94},"db6089e6":{"m":24,"g":94},"3520f75f":{"m":24,"g":94},"c8e9fed8":{"m":24,"g":94},"084fa54d":{"m":24,"g":94},"eba458bd":{"m":24,"g":94},"3d1cb0af":{"m":24,"g":94},"7d352b4f":{"m":24,"g":94},"87064015":{"m":24,"g":94},"7cd4f244":{"m":24,"g":94},"98111fbe":{"m":24,"g":94},"2ec39ab7":{"m":24,"g":94},"8f6274c8":{"m":24,"g":94},"325a06c2":{"m":24,"g":94},"79f81629":{"m":24,"g":94},"b688fd85":{"m":24,"g":94},"5bd89924":{"m":24,"g":94},"8d908a93":{"m":24,"g":94},"dd7e8b94":{"m":24,"g":94},"1f013d64":{"m":24,"g":94},"628e1fa7":{"m":24,"g":94},"c71880f8":{"m":24,"g":94},"bcb6611a":{"m":24,"g":94},"fa2aa0db":{"m":24,"g":94},"6a387a69":{"m":24,"g":94},"27f5ce0a":{"m":24,"g":94},"94862579":{"m":24,"g":94},"68e52626":{"m":24,"g":94},"e4d3333c":{"m":25,"g":94},"6f221d4c":{"m":25,"g":94},"aba6f51f":{"m":25,"g":94},"7f6c690b":{"m":25,"g":94},"40e6f513":{"m":25,"g":94},"40756776":{"m":25,"g":94},"9e8d2c7f":{"m":25,"g":94},"c9bff5fc":{"m":25,"g":94},"b04444ac":{"m":25,"g":94},"3d617a21":{"m":25,"g":94},"c020f9ce":{"m":25,"g":94},"ca600e8c":{"m":25,"g":94},"0c0c8137":{"m":25,"g":94},"90286d85":{"m":25,"g":94},"5e7dd984":{"m":25,"g":94},"bc3eaac2":{"m":25,"g":94},"a78d98de":{"m":25,"g":94},"7d5ed7c6":{"m":25,"g":94},"a6c7ebbb":{"m":25,"g":94},"bb0501c0":{"m":25,"g":94},"6b0f2e90":{"m":25,"g":94},"30a9b2ef":{"m":26,"g":94},"3cadecf0":{"m":26,"g":94},"e90e3a50":{"m":26,"g":94},"fbd6b94d":{"m":26,"g":94},"4c8093c8":{"m":26,"g":94},"ae7ee01a":{"m":26,"g":94},"76e59088":{"m":26,"g":94},"12ce3bef":{"m":26,"g":94},"4013a4e1":{"m":26,"g":94},"60340a36":{"m":26,"g":94},"70c78cfb":{"m":26,"g":94},"72b6ea88":{"m":26,"g":94},"b906c015":{"m":27,"g":94},"9319cd13":{"m":27,"g":94},"046c2b33":{"m":27,"g":94},"6b8f66ef":{"m":27,"g":94},"7937a886":{"m":27,"g":94},"2e218b9e":{"m":27,"g":94},"141e8c71":{"m":28,"g":94},"d53dcf9c":{"m":28,"g":94},"bb66cc4c":{"m":28,"g":94},"975adb80":{"m":28,"g":94},"0d4f3a9f":{"m":28,"g":94},"afd411d0":{"m":28,"g":94},"e1eae1fd":{"m":28,"g":94},"f4d9953d":{"m":28,"g":94},"4f005250":{"m":28,"g":94},"995af5a5":{"m":28,"g":94},"53985645":{"m":28,"g":94},"70cc0749":{"m":28,"g":94},"7dd8a7e6":{"m":28,"g":94},"947402c8":{"m":28,"g":94},"8c5382e6":{"m":28,"g":94},"001b0bdd":{"m":28,"g":94},"dc9d06d8":{"m":29,"g":94},"c31f084c":{"m":29,"g":94},"a01ddd96":{"m":29,"g":94},"7fa54a1a":{"m":29,"g":94},"05abd126":{"m":29,"g":94},"5f6fa04a":{"m":29,"g":94},"58a09708":{"m":29,"g":94},"ff68ae85":{"m":29,"g":94},"795eab6d":{"m":29,"g":94},"41bb1ab1":{"m":29,"g":94},"87e8c090":{"m":29,"g":94},"ad56e684":{"m":29,"g":94},"ffb15744":{"m":29,"g":94},"a9c833d5":{"m":29,"g":94},"94e01151":{"m":29,"g":94},"b216a545":{"m":29,"g":94},"fde83405":{"m":29,"g":94},"fd7926e4":{"m":29,"g":94},"399cad91":{"m":29,"g":94},"0a4f5f9b":{"m":29,"g":94},"3bc99e6f":{"m":29,"g":94},"ebf69964":{"m":29,"g":94},"b0ad0c1b":{"m":30,"g":94},"c877292c":{"m":30,"g":94},"0c1c72a0":{"m":30,"g":94},"41598e0d":{"m":30,"g":94},"89f23a51":{"m":30,"g":94},"cb99ba4f":{"m":30,"g":94},"32f61443":{"m":30,"g":94},"fb1f28cb":{"m":30,"g":94},"fb7421db":{"m":30,"g":94},"14b64930":{"m":30,"g":94},"82076370":{"m":30,"g":94},"7de60345":{"m":30,"g":94},"d84c5e70":{"m":30,"g":94},"d7854120":{"m":30,"g":94},"7b6a5332":{"m":30,"g":94},"4080e822":{"m":30,"g":94},"c245b789":{"m":30,"g":94},"9dae4078":{"m":30,"g":94},"fcc0f5ed":{"m":30,"g":94},"a97df791":{"m":30,"g":94},"33d61356":{"m":30,"g":94},"94752ac8":{"m":30,"g":94},"43fbb6d9":{"m":30,"g":94},"54fb1c80":{"m":30,"g":94},"b68c4c07":{"m":30,"g":94},"e712837d":{"m":30,"g":94},"7599bade":{"m":30,"g":94},"62757db6":{"m":30,"g":94},"73fa2d49":{"m":30,"g":94},"61728884":{"m":30,"g":94},"9cf0a5ba":{"m":30,"g":94},"b16e856f":{"m":30,"g":94},"05c50a82":{"m":30,"g":94},"b568df5d":{"m":30,"g":94},"10bca45b":{"m":30,"g":94},"b91a4cb1":{"m":30,"g":94},"95a28019":{"m":30,"g":94},"e040a245":{"m":30,"g":94},"9f662501":{"m":30,"g":94},"ab787594":{"m":30,"g":94},"228cf475":{"m":30,"g":94},"3a79613c":{"m":30,"g":94},"1ac304ee":{"m":30,"g":94},"20a4f927":{"m":30,"g":94},"0de7c2d0":{"m":30,"g":94},"6ed4e3b8":{"m":30,"g":94},"00023d62":{"m":30,"g":94},"c62d560c":{"m":30,"g":94},"2b8257f3":{"m":30,"g":94},"7623091d":{"m":30,"g":94},"f724f1f1":{"m":30,"g":94},"6db27f7b":{"m":30,"g":94},"4d929107":{"m":30,"g":94},"fbe0c818":{"m":30,"g":94},"5bd95374":{"m":31,"g":94},"0cb099e2":{"m":31,"g":94},"93d4e354":{"m":31,"g":94},"9195d136":{"m":31,"g":94},"14cb544d":{"m":31,"g":94},"e86b1ccb":{"m":31,"g":94},"8d2d876f":{"m":31,"g":94},"326df4ba":{"m":31,"g":94},"6767e222":{"m":31,"g":94},"73cf6834":{"m":31,"g":94},"1c2b5f52":{"m":31,"g":94},"96a2093e":{"m":31,"g":94},"a34dd86a":{"m":31,"g":94},"67c0d832":{"m":31,"g":94},"a59636bb":{"m":31,"g":94},"fe502432":{"m":31,"g":94},"f14569f6":{"m":31,"g":94},"8f790ac1":{"m":31,"g":94},"616b59f3":{"m":31,"g":94},"c8423ca3":{"m":31,"g":94},"e205527c":{"m":31,"g":94},"0909bb0d":{"m":31,"g":94},"ad3e4f16":{"m":31,"g":94},"312e8492":{"m":31,"g":94},"95f5fbf1":{"m":31,"g":94},"cebd78d8":{"m":31,"g":94},"0076f115":{"m":31,"g":94},"f7fb68d2":{"m":31,"g":94},"396a13e6":{"m":31,"g":94},"65915f9f":{"m":31,"g":94},"162f3ccb":{"m":31,"g":94},"65e89bae":{"m":31,"g":94},"6a38efa8":{"m":31,"g":94},"c5fe11a8":{"m":32,"g":94},"75ce37f4":{"m":32,"g":94},"97589a60":{"m":32,"g":94},"632d506d":{"m":32,"g":94},"3579162a":{"m":32,"g":94},"7514b9f8":{"m":32,"g":94},"158e8f1e":{"m":32,"g":94},"d3efcb39":{"m":32,"g":94},"2c615d12":{"m":32,"g":94},"61bb223e":{"m":32,"g":94},"15f1a49d":{"m":32,"g":94},"308d0240":{"m":32,"g":94},"ab4990e4":{"m":32,"g":94},"90227800":{"m":32,"g":94},"30b4f771":{"m":32,"g":94},"66e7dcaf":{"m":32,"g":94},"bc4c7a35":{"m":32,"g":94},"1cb4da5c":{"m":32,"g":94},"e61d13ac":{"m":32,"g":94},"b20daf98":{"m":32,"g":94},"f6af3a65":{"m":32,"g":94},"c9064e6f":{"m":32,"g":94},"a5b14ad0":{"m":32,"g":94},"5fafcac0":{"m":32,"g":94},"364d3d72":{"m":32,"g":94},"5623826f":{"m":32,"g":94},"83e23c69":{"m":32,"g":94},"ac1b74fa":{"m":32,"g":94},"068e9eae":{"m":32,"g":94},"d6aeb9fa":{"m":32,"g":94},"1fb94599":{"m":32,"g":94},"bea2bb9e":{"m":32,"g":94},"cd10654e":{"m":32,"g":94},"350a8160":{"m":32,"g":94},"6242c399":{"m":32,"g":94},"04707b09":{"m":32,"g":94},"ff2cfdb1":{"m":32,"g":94},"a8ae6403":{"m":32,"g":94},"d8476818":{"m":32,"g":94},"df191254":{"m":32,"g":94},"b997a18d":{"m":32,"g":94},"d8627ed1":{"m":32,"g":94},"fa13b95d":{"m":32,"g":94},"3c1f5a92":{"m":32,"g":94},"57d0bd91":{"m":32,"g":94},"cdc8d607":{"m":32,"g":94},"9208591f":{"m":32,"g":94},"5d0d40d0":{"m":32,"g":94},"f624f6a6":{"m":32,"g":94},"3694f8f9":{"m":32,"g":94},"5a261bd0":{"m":32,"g":94},"6aa8ad14":{"m":32,"g":94},"26e9c12c":{"m":32,"g":94},"87a0db82":{"m":32,"g":94},"f25f4dfd":{"m":33,"g":94},"184ae1c6":{"m":33,"g":94},"198974cd":{"m":33,"g":94},"6cc38b2b":{"m":33,"g":94},"1ece2cda":{"m":33,"g":94},"c8a9e791":{"m":33,"g":94},"3602692c":{"m":33,"g":94},"909f3436":{"m":33,"g":94},"5ff25cdf":{"m":33,"g":94},"2f1d9283":{"m":33,"g":94},"c61a1b6f":{"m":33,"g":94},"9935f97b":{"m":33,"g":94},"13ac95b8":{"m":34,"g":94},"492143bf":{"m":34,"g":94},"0a97d796":{"m":34,"g":94},"c411f32e":{"m":34,"g":94},"bf53bf51":{"m":34,"g":94},"b1a540ec":{"m":34,"g":94},"66975360":{"m":34,"g":94},"6c498313":{"m":34,"g":94},"99994427":{"m":35,"g":94},"6def9b01":{"m":35,"g":94},"47f20da2":{"m":35,"g":94},"4a9f8ea4":{"m":35,"g":94},"58fa6076":{"m":35,"g":94},"6487ef64":{"m":35,"g":94},"9b080524":{"m":35,"g":94},"32a4141d":{"m":35,"g":94},"08360553":{"m":35,"g":94},"00b19f19":{"m":35,"g":94},"6cb32ef9":{"m":35,"g":94},"761b2ceb":{"m":35,"g":94},"54772f78":{"m":35,"g":94},"1b5d56f7":{"m":35,"g":94},"d134c139":{"m":35,"g":94},"6cc9c525":{"m":35,"g":94},"52cefdbf":{"m":35,"g":94},"51c554d8":{"m":35,"g":94},"79ece2c5":{"m":35,"g":94},"55f5976b":{"m":35,"g":94},"b7f83410":{"m":35,"g":94},"f414352a":{"m":35,"g":94},"a362340b":{"m":35,"g":94},"381dd57b":{"m":35,"g":94},"8153168c":{"m":35,"g":94},"6c34d633":{"m":35,"g":94},"5ab9418f":{"m":36,"g":94},"843e63d8":{"m":36,"g":94},"a63c8275":{"m":36,"g":94},"dc67d976":{"m":36,"g":94},"1e495e08":{"m":36,"g":94},"12cb115d":{"m":36,"g":94},"c500f96b":{"m":36,"g":94},"474317f2":{"m":36,"g":94},"f64eae3a":{"m":36,"g":94},"a5a134f3":{"m":36,"g":94},"2561ed01":{"m":36,"g":94},"90a26be3":{"m":37,"g":94},"1f4b5f77":{"m":37,"g":94},"76524b70":{"m":37,"g":94},"3a6e0418":{"m":37,"g":94},"2fa5cec7":{"m":37,"g":94},"27b557ae":{"m":37,"g":94},"93dffd69":{"m":37,"g":94},"2abe4f1c":{"m":37,"g":94},"37963394":{"m":37,"g":94},"899cf5c4":{"m":37,"g":94},"e79f6cd7":{"m":37,"g":94},"9ba1f097":{"m":37,"g":94},"282681b8":{"m":37,"g":94},"58cafe23":{"m":37,"g":94},"9463bc13":{"m":37,"g":94},"e3fc4658":{"m":37,"g":94},"33b54e7c":{"m":37,"g":94},"30b404ce":{"m":37,"g":94},"70b68029":{"m":37,"g":94},"f3d32f88":{"m":37,"g":94},"8779da95":{"m":37,"g":94},"ad0ff62a":{"m":37,"g":94},"9a903a87":{"m":37,"g":94},"68be2f6d":{"m":37,"g":94},"b912de11":{"m":37,"g":94},"eb02c161":{"m":37,"g":94},"71221692":{"m":37,"g":94},"c33d82a2":{"m":37,"g":94},"8234e663":{"m":37,"g":94},"debbdb51":{"m":37,"g":94},"3efa7981":{"m":37,"g":94},"2a71be5e":{"m":37,"g":94},"44621377":{"m":37,"g":94},"fec185ce":{"m":37,"g":94},"c03cece4":{"m":37,"g":94},"15c75e41":{"m":37,"g":94},"224200e3":{"m":37,"g":94},"8c0efa51":{"m":37,"g":94},"144bc70f":{"m":37,"g":94},"46094e0c":{"m":37,"g":94},"3a6e8b6d":{"m":37,"g":94},"fbb4754c":{"m":37,"g":94},"6c7cb903":{"m":37,"g":94},"dff2860a":{"m":37,"g":94},"e72275cf":{"m":37,"g":94},"fec2d122":{"m":37,"g":94},"8d1095db":{"m":37,"g":94},"743007e1":{"m":37,"g":94},"9144ed10":{"m":37,"g":94},"69b3bb9a":{"m":37,"g":94},"689ff588":{"m":37,"g":94},"a7c47e0f":{"m":37,"g":94},"e4d68afc":{"m":37,"g":94},"c9b75917":{"m":37,"g":94},"662ecd93":{"m":37,"g":94},"8e6bdf85":{"m":37,"g":94},"05bea688":{"m":37,"g":94},"ab4a83b2":{"m":37,"g":94},"62f15eea":{"m":37,"g":94},"79794af5":{"m":37,"g":94},"3494b32c":{"m":37,"g":94},"eda7c090":{"m":37,"g":94},"5ce55aee":{"m":38,"g":94},"2d346a57":{"m":38,"g":94},"446ea332":{"m":38,"g":94},"8f527e29":{"m":38,"g":94},"7f24ea95":{"m":38,"g":94},"1acccb36":{"m":38,"g":94},"aa2750be":{"m":38,"g":94},"5e62a6b7":{"m":38,"g":94},"5752f25e":{"m":38,"g":94},"7c162fa9":{"m":38,"g":94},"36078fb2":{"m":38,"g":94},"b3710d2c":{"m":38,"g":94},"c6b6d2e7":{"m":38,"g":94},"82136eb0":{"m":39,"g":94},"b8ccaf4d":{"m":39,"g":94},"a68cb201":{"m":39,"g":94},"014982b5":{"m":39,"g":94},"a6db8862":{"m":39,"g":94},"b4408b0d":{"m":39,"g":94},"2cd7e181":{"m":39,"g":94},"37c5899f":{"m":40,"g":94},"f39a0197":{"m":40,"g":94},"3c93187c":{"m":40,"g":94},"fb2d0680":{"m":40,"g":94},"067d8e16":{"m":40,"g":94},"e6692bf4":{"m":40,"g":94},"28b4d8e1":{"m":40,"g":94},"bc068e96":{"m":40,"g":94},"8d4ed42a":{"m":40,"g":94},"2854a5ea":{"m":40,"g":94},"42a2d82b":{"m":40,"g":94},"e4780cf8":{"m":40,"g":94},"39bb49d1":{"m":40,"g":94},"6f3cf129":{"m":40,"g":94},"13f1357e":{"m":40,"g":94},"2a99993c":{"m":40,"g":94},"167591e8":{"m":40,"g":94},"441c22db":{"m":40,"g":94},"ce636ac4":{"m":40,"g":94},"7b69d91b":{"m":41,"g":94},"e8613df0":{"m":41,"g":94},"c5325aba":{"m":41,"g":94},"ebbc42d9":{"m":41,"g":94},"3ff64113":{"m":41,"g":94},"2b302b93":{"m":41,"g":94},"68f8b60d":{"m":41,"g":94},"6a5b352a":{"m":41,"g":94},"565b05f0":{"m":41,"g":94},"b6aad70a":{"m":41,"g":94},"551a3a9d":{"m":41,"g":94},"91877a9f":{"m":41,"g":94},"f7cce751":{"m":41,"g":94},"17e998f1":{"m":41,"g":94},"c98e84c2":{"m":41,"g":94},"9c064bf7":{"m":41,"g":94},"58d1082e":{"m":41,"g":94},"4d086719":{"m":41,"g":94},"9244f27f":{"m":41,"g":94},"2422de51":{"m":41,"g":94},"521f862d":{"m":41,"g":94},"34c32d28":{"m":41,"g":94},"dde8bb16":{"m":41,"g":94},"8ac3ccc0":{"m":41,"g":94},"9b0926ce":{"m":41,"g":94},"1c1bdc76":{"m":41,"g":94},"6bfdb403":{"m":41,"g":94},"f8fb4ce9":{"m":41,"g":94},"5d0ba403":{"m":41,"g":94},"04b262cd":{"m":41,"g":94},"2432ad40":{"m":41,"g":94},"45473d4b":{"m":41,"g":94},"114bbc86":{"m":41,"g":94},"32eb6e96":{"m":41,"g":94},"e0b5dbce":{"m":41,"g":94},"e6852b0d":{"m":41,"g":94},"4ae0969c":{"m":41,"g":94},"317631ca":{"m":41,"g":94},"b5648353":{"m":41,"g":94},"2c7d0a5b":{"m":41,"g":94},"8cdc76f6":{"m":41,"g":94},"f202ed97":{"m":41,"g":94},"100f5b8b":{"m":41,"g":94},"619bb6dd":{"m":41,"g":94},"b88ea90d":{"m":41,"g":94},"99ec439d":{"m":41,"g":94},"0f4fb19b":{"m":41,"g":94},"63ba2f8d":{"m":41,"g":94},"36d5acfc":{"m":41,"g":94},"3f0fe08d":{"m":41,"g":94},"55b974f9":{"m":41,"g":94},"f86c1e61":{"m":41,"g":94},"acaffd23":{"m":41,"g":94},"04868543":{"m":41,"g":94},"fd9ad817":{"m":41,"g":94},"e165a9fc":{"m":41,"g":94},"4e4459b9":{"m":41,"g":94},"065bb947":{"m":41,"g":94},"f42e9bfb":{"m":41,"g":94},"840c5dbc":{"m":41,"g":94},"63e845d0":{"m":41,"g":94},"9aa6553d":{"m":41,"g":94},"b1e330bc":{"m":41,"g":94},"4353acb4":{"m":41,"g":94},"9ae1db0b":{"m":41,"g":94},"00c7e636":{"m":42,"g":94},"23cc66f7":{"m":42,"g":94},"5d09ca57":{"m":42,"g":94},"81c33274":{"m":42,"g":94},"f13d86f9":{"m":42,"g":94},"aba9eae4":{"m":42,"g":94},"bbd72bfc":{"m":42,"g":94},"b503881b":{"m":42,"g":94},"58093b86":{"m":42,"g":94},"8275049c":{"m":42,"g":94},"5476ccad":{"m":42,"g":94},"b040ed71":{"m":42,"g":94},"c9e66586":{"m":42,"g":94},"e11ab79e":{"m":42,"g":94},"01fdb2f3":{"m":42,"g":94},"c996e8cc":{"m":42,"g":94},"087257ea":{"m":43,"g":94},"736f0402":{"m":43,"g":94},"769bf11c":{"m":43,"g":94},"3db43d1b":{"m":43,"g":94},"f0f8a769":{"m":43,"g":94},"2bcfba1b":{"m":43,"g":94},"bc12d403":{"m":43,"g":94},"392f2863":{"m":43,"g":94},"6d0fa73e":{"m":43,"g":94},"9e0dac1a":{"m":43,"g":94},"a95d5589":{"m":43,"g":94},"d17d19e5":{"m":43,"g":94},"dd3809fa":{"m":43,"g":94},"7feba415":{"m":43,"g":94},"30ee3630":{"m":43,"g":94},"e5db40dc":{"m":43,"g":94},"b1709305":{"m":43,"g":94},"5ab20cce":{"m":43,"g":94},"02f7f3e4":{"m":43,"g":94},"2782132b":{"m":43,"g":94},"d19cc0b9":{"m":43,"g":94},"b0facb33":{"m":43,"g":94},"ecb8bad2":{"m":43,"g":94},"dbec2f18":{"m":43,"g":94},"e4b367ba":{"m":43,"g":94},"d10b933a":{"m":43,"g":94},"9116b289":{"m":43,"g":94},"a5114b6f":{"m":43,"g":94},"b6b40946":{"m":43,"g":94},"f1088e0f":{"m":43,"g":94},"175afed3":{"m":43,"g":94},"4a292f67":{"m":43,"g":94},"cd0be748":{"m":43,"g":94},"56503d9b":{"m":43,"g":94},"02bc9579":{"m":43,"g":94},"24f3e151":{"m":43,"g":94},"6790240c":{"m":43,"g":94},"061e5463":{"m":43,"g":94},"0c1e8796":{"m":43,"g":94},"869f1c02":{"m":43,"g":94},"2725f8da":{"m":43,"g":94},"da1ffed6":{"m":43,"g":94},"48761171":{"m":43,"g":94},"c3f2fc5a":{"m":43,"g":94},"7ee6c259":{"m":43,"g":94},"9610fcd4":{"m":43,"g":94},"31fad29a":{"m":43,"g":94},"9da5a60b":{"m":43,"g":94},"69aa937a":{"m":43,"g":94},"5d638c92":{"m":43,"g":94},"e37cdab0":{"m":43,"g":94},"1d9deeac":{"m":43,"g":94},"dafb6a52":{"m":43,"g":94},"862cd265":{"m":43,"g":94},"1f26e8b8":{"m":44,"g":94},"5e1558f1":{"m":44,"g":94},"94cde109":{"m":44,"g":94},"00611286":{"m":44,"g":94},"e68b9e76":{"m":44,"g":94},"7ce36068":{"m":44,"g":94},"efb099cd":{"m":44,"g":94},"09603c6d":{"m":44,"g":94},"cf470fea":{"m":44,"g":94},"45d5af24":{"m":44,"g":94},"b121bc03":{"m":44,"g":94},"e12358dc":{"m":44,"g":94},"554fbf93":{"m":44,"g":94},"b48edff6":{"m":44,"g":94},"593b19f2":{"m":44,"g":94},"59cbf476":{"m":44,"g":94},"95946271":{"m":44,"g":94},"5c4ce656":{"m":44,"g":94},"cbbc82b7":{"m":44,"g":94},"8bee20f8":{"m":44,"g":94},"12cad0fe":{"m":44,"g":94},"b6cd9036":{"m":44,"g":94},"30643fed":{"m":45,"g":94},"e646c590":{"m":45,"g":94},"c555ce2c":{"m":45,"g":94},"40900bae":{"m":45,"g":94},"a2f5e755":{"m":45,"g":94},"2148914e":{"m":45,"g":94},"def55bc8":{"m":45,"g":94},"86a2c473":{"m":45,"g":94},"1701b0db":{"m":45,"g":94},"384d85ba":{"m":45,"g":94},"60597219":{"m":45,"g":94},"fc82f5a7":{"m":45,"g":94},"0089c4bc":{"m":45,"g":94},"72e7b57a":{"m":45,"g":94},"87a7cfa0":{"m":45,"g":94},"8f8f96a6":{"m":45,"g":94},"05b3bf5e":{"m":45,"g":94},"3f5ac88d":{"m":45,"g":94},"0d800090":{"m":45,"g":94},"b7d05594":{"m":45,"g":94},"80a90547":{"m":45,"g":94},"9af7b88e":{"m":45,"g":94},"fbcbb263":{"m":45,"g":94},"2fce449b":{"m":45,"g":94},"ad4125d1":{"m":45,"g":94},"17536e7e":{"m":45,"g":94},"65859754":{"m":46,"g":94},"2ce32db6":{"m":46,"g":94},"793b79db":{"m":46,"g":94},"1363b519":{"m":46,"g":94},"0abbf289":{"m":46,"g":94},"c17c5781":{"m":46,"g":94},"916b3cdd":{"m":46,"g":94},"838dcda1":{"m":46,"g":94},"efbc116a":{"m":46,"g":94},"6aed0445":{"m":46,"g":94},"908dd7f9":{"m":46,"g":94},"f4cd8040":{"m":46,"g":94},"be7986e0":{"m":46,"g":94},"5a5f1843":{"m":46,"g":94},"7b394e5f":{"m":46,"g":94},"3b60558d":{"m":46,"g":94},"5a9a4f41":{"m":46,"g":94},"72e979bf":{"m":46,"g":94},"146f6134":{"m":46,"g":94},"660ecb73":{"m":46,"g":94},"2565cb0f":{"m":46,"g":94},"066e8a4e":{"m":46,"g":94},"2134f089":{"m":46,"g":94},"a54f278d":{"m":46,"g":94},"d1b31b06":{"m":46,"g":94},"d59a4782":{"m":46,"g":94},"104bf260":{"m":46,"g":94},"3bf3d011":{"m":46,"g":94},"d86a2d65":{"m":46,"g":94},"16eb33ff":{"m":46,"g":94},"61cf00e1":{"m":46,"g":94},"b9fd178f":{"m":46,"g":94},"d8e9d61f":{"m":46,"g":94},"a2e0424a":{"m":46,"g":94},"8ce202a4":{"m":46,"g":94},"d913d52c":{"m":46,"g":94},"0ab7bcaf":{"m":46,"g":94},"438526a8":{"m":46,"g":94},"f7102fbd":{"m":46,"g":94},"a7a0a688":{"m":46,"g":94},"2d4ce1b7":{"m":46,"g":94},"4ba815b8":{"m":46,"g":94},"5f65e2b8":{"m":46,"g":94},"4e2af03c":{"m":46,"g":94},"3184aa95":{"m":46,"g":94},"b548801d":{"m":46,"g":94},"539df95d":{"m":46,"g":94},"5e00ddeb":{"m":46,"g":94},"54dd3ea1":{"m":46,"g":94},"d04899d7":{"m":46,"g":94},"5010e0d2":{"m":46,"g":94},"5e6c3265":{"m":46,"g":94},"680cad20":{"m":46,"g":94},"0a24eb85":{"m":46,"g":94},"3839be29":{"m":46,"g":94},"6e13b650":{"m":46,"g":94},"6fcd6d7d":{"m":46,"g":94},"c77762d5":{"m":46,"g":94},"51c81e33":{"m":46,"g":94},"eaade87a":{"m":46,"g":94},"86fc0d79":{"m":46,"g":94},"1be853ee":{"m":46,"g":94},"86e0dde5":{"m":46,"g":94},"2b809788":{"m":46,"g":94},"9d6fb084":{"m":46,"g":94},"ced362f7":{"m":46,"g":94},"9084a864":{"m":46,"g":94},"6aa94b96":{"m":46,"g":94},"c2650748":{"m":46,"g":94},"07bf2e84":{"m":46,"g":94},"a628dd8e":{"m":46,"g":94},"1e890341":{"m":46,"g":94},"715b16c1":{"m":46,"g":94},"9ce8e1a9":{"m":46,"g":94},"fb99aaa5":{"m":46,"g":94},"b77a02cd":{"m":46,"g":94},"f407fcf9":{"m":47,"g":94},"54479d6f":{"m":47,"g":94},"ba069a24":{"m":47,"g":94},"125b1199":{"m":47,"g":94},"eff468dd":{"m":47,"g":94},"a1bd7190":{"m":47,"g":94},"78c1d644":{"m":47,"g":94},"027e6524":{"m":47,"g":94},"b808a383":{"m":47,"g":94},"602ebc66":{"m":47,"g":94},"530ae1bd":{"m":47,"g":94},"befc6beb":{"m":47,"g":94},"59a5ba9b":{"m":47,"g":94},"86c37d01":{"m":47,"g":94},"f18b9c72":{"m":47,"g":94},"3e335743":{"m":47,"g":94},"0d94f1dd":{"m":47,"g":94},"e728258d":{"m":47,"g":94},"239eafbd":{"m":47,"g":94},"9d427265":{"m":47,"g":94},"00ffde20":{"m":47,"g":94},"ddeb9d42":{"m":47,"g":94},"aaf0a315":{"m":47,"g":94},"f9633fa9":{"m":47,"g":94},"087ab832":{"m":47,"g":94},"8169c6f4":{"m":47,"g":94},"3d043319":{"m":47,"g":94},"a8aad935":{"m":47,"g":94},"47ffe7af":{"m":47,"g":94},"b3523af8":{"m":47,"g":94},"1929c067":{"m":47,"g":94},"ed53ac84":{"m":47,"g":94},"520f0094":{"m":47,"g":94},"9c939a3d":{"m":47,"g":94},"549e8b83":{"m":47,"g":94},"a1f32867":{"m":47,"g":94},"760552e0":{"m":47,"g":94},"d9aada9d":{"m":47,"g":94},"f11eb90f":{"m":47,"g":94},"95a4ed12":{"m":47,"g":94},"d1150e9a":{"m":47,"g":94},"e3126e3c":{"m":47,"g":94},"a5095520":{"m":47,"g":94},"7ef0084b":{"m":47,"g":94},"f9a377f6":{"m":47,"g":94},"4ade15dd":{"m":47,"g":94},"8dc84da0":{"m":47,"g":94},"f16eb15d":{"m":47,"g":94},"5bc2508b":{"m":47,"g":94},"a71a44f2":{"m":47,"g":94},"691808d5":{"m":47,"g":94},"d32fba2a":{"m":47,"g":94},"67c424cc":{"m":47,"g":94},"1ae270c5":{"m":47,"g":94},"c77c1e05":{"m":47,"g":94},"dca87ec3":{"m":47,"g":94},"4b1d7a25":{"m":47,"g":94},"a5e0defb":{"m":47,"g":94},"96766101":{"m":47,"g":94},"a146d999":{"m":47,"g":94},"f5113e50":{"m":47,"g":94},"02755768":{"m":47,"g":94},"463d56bf":{"m":47,"g":94},"530ff541":{"m":47,"g":94},"3cd28092":{"m":47,"g":94},"704f8e8e":{"m":47,"g":94},"1853c352":{"m":47,"g":94},"32c9a7ec":{"m":48,"g":94},"b01df48c":{"m":48,"g":94},"c29b98e0":{"m":48,"g":94},"954f4e6b":{"m":48,"g":94},"2558d6a6":{"m":48,"g":94},"29ebe3df":{"m":48,"g":94},"f6dd6486":{"m":48,"g":94},"ea53c63b":{"m":48,"g":94},"a10d5309":{"m":48,"g":94},"aae5434b":{"m":48,"g":94},"c3eac1b0":{"m":48,"g":94},"b275ce00":{"m":48,"g":94},"13ce3e4b":{"m":48,"g":94},"df246e69":{"m":48,"g":94},"fb9fb351":{"m":48,"g":94},"c722d9bd":{"m":48,"g":94},"218ab361":{"m":48,"g":94},"9a00e6f4":{"m":49,"g":94},"4f8c3aea":{"m":49,"g":94},"2369e882":{"m":49,"g":94},"ad30d5cf":{"m":49,"g":94},"dfec7fca":{"m":49,"g":94},"8048c28c":{"m":49,"g":94},"30af7dfb":{"m":49,"g":94},"f6f71379":{"m":49,"g":94},"f35cb46c":{"m":49,"g":94},"7f8fcd39":{"m":49,"g":94},"5c6a41fa":{"m":49,"g":94},"722530fa":{"m":49,"g":94},"56a347f7":{"m":49,"g":94},"3295cd8a":{"m":49,"g":94},"5942dfc0":{"m":49,"g":94},"63a395b9":{"m":49,"g":94},"7d671e4a":{"m":49,"g":94},"699384cb":{"m":49,"g":94},"ffd20fcd":{"m":49,"g":94},"55bd97f3":{"m":49,"g":94},"e57c3e12":{"m":49,"g":94},"f239268f":{"m":49,"g":94},"929c7621":{"m":49,"g":94},"b7a065ea":{"m":49,"g":94},"b1104538":{"m":49,"g":94},"3b44bbee":{"m":49,"g":94},"80e2c4a8":{"m":49,"g":94},"66318ffe":{"m":49,"g":94},"76619261":{"m":49,"g":94},"2a3992b6":{"m":49,"g":94},"4af3f889":{"m":49,"g":94},"df7fe452":{"m":49,"g":94},"a7164b62":{"m":49,"g":94},"11668533":{"m":49,"g":94},"a9e90b4b":{"m":49,"g":94},"8c280cee":{"m":49,"g":94},"9c745d07":{"m":49,"g":94},"ebaa2f31":{"m":49,"g":94},"62832bb2":{"m":49,"g":94},"11f881d1":{"m":49,"g":94},"38625e21":{"m":49,"g":94},"c1f401fc":{"m":49,"g":94},"3b878863":{"m":49,"g":94},"f719d9ae":{"m":49,"g":94},"edad3731":{"m":49,"g":94},"976bc302":{"m":49,"g":94},"2f2e0743":{"m":49,"g":94},"2ffe0a73":{"m":49,"g":94},"cf248976":{"m":49,"g":94},"e5c67150":{"m":49,"g":94},"023d0a73":{"m":49,"g":94},"ac5a0f04":{"m":50,"g":94},"ea34350d":{"m":50,"g":94},"1605ae12":{"m":50,"g":94},"1aea19f6":{"m":50,"g":94},"1f76fc6e":{"m":50,"g":94},"7f076c2c":{"m":50,"g":94},"3c5538f7":{"m":50,"g":94},"10189d08":{"m":50,"g":94},"c4336b2b":{"m":50,"g":94},"4d62bca5":{"m":50,"g":94},"e1e595d7":{"m":50,"g":94},"5ada33ff":{"m":50,"g":94},"254fd130":{"m":50,"g":94},"538fa0ae":{"m":50,"g":94},"55842eb8":{"m":50,"g":94},"a866b65e":{"m":50,"g":94},"4b0a1c93":{"m":50,"g":94},"8e1adb84":{"m":50,"g":94},"dd44173d":{"m":50,"g":94},"8912b763":{"m":50,"g":94},"be0124bd":{"m":50,"g":94},"fe5d3e81":{"m":50,"g":94},"731146f6":{"m":50,"g":94},"fa271613":{"m":50,"g":94},"5652c565":{"m":50,"g":94},"e3938b2f":{"m":50,"g":94},"c211e7b6":{"m":50,"g":94},"d90c3d6b":{"m":50,"g":94},"9e8f8fbf":{"m":50,"g":94},"b509db58":{"m":50,"g":94},"dbe17293":{"m":50,"g":94},"84a1698d":{"m":50,"g":94},"32293a29":{"m":50,"g":94},"79216908":{"m":50,"g":94},"bbb81c24":{"m":50,"g":94},"52f58fc4":{"m":50,"g":94},"145c0ddc":{"m":50,"g":94},"505d7f71":{"m":50,"g":94},"cbedd1db":{"m":50,"g":94},"ad47749b":{"m":50,"g":94},"751c3a03":{"m":50,"g":94},"60769be1":{"m":50,"g":94},"a78d8f8d":{"m":50,"g":94},"c5f86501":{"m":50,"g":94},"d98fa1e9":{"m":50,"g":94},"865233e2":{"m":50,"g":94},"66d4859a":{"m":50,"g":94},"e1b63624":{"m":50,"g":94},"c35cd1f8":{"m":50,"g":94},"72f87b72":{"m":50,"g":94},"62a4a339":{"m":50,"g":94},"2797bc34":{"m":50,"g":94},"fed4c694":{"m":51,"g":94},"fb6e04a0":{"m":51,"g":94},"6997e28f":{"m":51,"g":94},"a0e58740":{"m":51,"g":94},"37c8a576":{"m":51,"g":94},"c754652f":{"m":51,"g":94},"0b46b951":{"m":51,"g":94},"2763c0a7":{"m":51,"g":94},"de3b67b7":{"m":51,"g":94},"19f33b32":{"m":51,"g":94},"30ce5b59":{"m":51,"g":94},"bc1f6fda":{"m":51,"g":94},"867e092f":{"m":51,"g":94},"88c7763f":{"m":51,"g":94},"e4118b15":{"m":51,"g":94},"ba4ee37f":{"m":51,"g":94},"fae4e5e9":{"m":52,"g":94},"afe1e465":{"m":52,"g":94},"f50a6cf4":{"m":52,"g":94},"fe97a2d4":{"m":52,"g":94},"8b48496a":{"m":52,"g":94},"4057ea82":{"m":52,"g":94},"4f2ee48e":{"m":52,"g":94},"71ff2728":{"m":52,"g":94},"b7038fec":{"m":52,"g":94},"65fdb289":{"m":52,"g":94},"b2ccf36d":{"m":52,"g":94},"d4fc1a70":{"m":52,"g":94},"db674e3d":{"m":52,"g":94},"fb915bd1":{"m":52,"g":94},"09798b36":{"m":52,"g":94},"b79fffdc":{"m":52,"g":94},"cd51758f":{"m":52,"g":94},"91e5dbf5":{"m":52,"g":94},"dd5eba4c":{"m":52,"g":94},"a4fd2f9b":{"m":52,"g":94},"92d1253e":{"m":52,"g":94},"a9ca297d":{"m":52,"g":94},"2a02185c":{"m":52,"g":94},"f8b03269":{"m":53,"g":94},"04957965":{"m":53,"g":94},"1228f7ca":{"m":53,"g":94},"fda628d8":{"m":53,"g":94},"07ec07ad":{"m":53,"g":94},"83b340e3":{"m":53,"g":94},"0639bf15":{"m":53,"g":94},"aa47f642":{"m":53,"g":94},"3ddb1c46":{"m":53,"g":94},"480e38a7":{"m":53,"g":94},"69e2d4fb":{"m":53,"g":94},"85e1a6f3":{"m":53,"g":94},"33deca81":{"m":53,"g":94},"18108abe":{"m":53,"g":94},"c54bda30":{"m":53,"g":94},"3c79ad35":{"m":53,"g":94},"983bfcf3":{"m":53,"g":94},"28bc60dc":{"m":53,"g":94},"7301a39b":{"m":53,"g":94},"47eb139f":{"m":53,"g":94},"5c18a037":{"m":53,"g":94},"5c91a315":{"m":53,"g":94},"3dbd73d3":{"m":53,"g":94},"e9a6203d":{"m":53,"g":94},"62c516ac":{"m":53,"g":94},"fc78640e":{"m":53,"g":94},"906d795f":{"m":53,"g":94},"118b6af3":{"m":53,"g":94},"9449a954":{"m":53,"g":94},"5f12f0e7":{"m":53,"g":94},"d5b95cbb":{"m":53,"g":94},"0303ca91":{"m":53,"g":94},"00181098":{"m":53,"g":94},"4936be8a":{"m":53,"g":94},"1bfa511b":{"m":53,"g":94},"f5b5f2bf":{"m":53,"g":94},"7e4c6dd8":{"m":53,"g":94},"d622851d":{"m":53,"g":94},"883c9554":{"m":53,"g":94},"0d6a49bd":{"m":53,"g":94},"ccaf1f99":{"m":53,"g":94},"7d1485d3":{"m":53,"g":94},"7d5d1d3d":{"m":53,"g":94},"b53d6cbd":{"m":53,"g":94},"01017d4c":{"m":53,"g":94},"94e167ea":{"m":53,"g":94},"262e370f":{"m":53,"g":94},"419a57e7":{"m":53,"g":94},"e5f227c0":{"m":54,"g":94},"0e7409ad":{"m":54,"g":94},"3cde5eb6":{"m":54,"g":94},"f5b2a3aa":{"m":54,"g":94},"f6817596":{"m":54,"g":94},"67b65794":{"m":54,"g":94},"37ee906f":{"m":54,"g":94},"34b364e0":{"m":54,"g":94},"84d96b3a":{"m":54,"g":94},"3d32e4a3":{"m":54,"g":94},"64fceab8":{"m":54,"g":94},"71e2a277":{"m":54,"g":94},"4a63c181":{"m":54,"g":94},"2b0fc594":{"m":54,"g":94},"9cc733b3":{"m":54,"g":94},"d693ec04":{"m":54,"g":94},"18ea841f":{"m":54,"g":94},"786be44d":{"m":54,"g":94},"2db44698":{"m":54,"g":94},"ed45e509":{"m":54,"g":94},"ec52464d":{"m":54,"g":94},"eb0c1f53":{"m":54,"g":94},"b2986d7a":{"m":54,"g":94},"8f4d04e5":{"m":55,"g":94},"feb2b768":{"m":55,"g":94},"d95a5f5b":{"m":55,"g":94},"4b83db24":{"m":55,"g":94},"64456cf0":{"m":55,"g":94},"bb4a9220":{"m":55,"g":94},"21e9e63a":{"m":55,"g":94},"1fc84cf6":{"m":55,"g":94},"361ea8d9":{"m":55,"g":94},"33c5ff28":{"m":55,"g":94},"5ce9daea":{"m":55,"g":94},"ce094a5d":{"m":55,"g":94},"e2102669":{"m":55,"g":94},"bd619616":{"m":55,"g":94},"56198b45":{"m":55,"g":94},"ba36b552":{"m":55,"g":94},"9cd9dc83":{"m":55,"g":94},"7a1aecb9":{"m":55,"g":94},"82699474":{"m":55,"g":94},"7154b4b1":{"m":55,"g":94},"b532a5fd":{"m":55,"g":94},"a0592c05":{"m":55,"g":94},"e8dbdf75":{"m":55,"g":94},"e04d3f28":{"m":55,"g":94},"5f2595be":{"m":55,"g":94},"0ba2c589":{"m":55,"g":94},"fccbfa37":{"m":55,"g":94},"2f9bd0fa":{"m":55,"g":94},"5282a473":{"m":55,"g":94},"f0ed9c35":{"m":55,"g":94},"e3b3acfa":{"m":55,"g":94},"2673fa29":{"m":55,"g":94},"dedaf8cd":{"m":55,"g":94},"32ed0160":{"m":55,"g":94},"6efa9e4a":{"m":55,"g":94},"7791fd99":{"m":55,"g":94},"2ac36b9a":{"m":55,"g":94},"2d60a5ee":{"m":55,"g":94},"2e4a5907":{"m":55,"g":94},"c0ee46fe":{"m":55,"g":94},"9208618b":{"m":55,"g":94},"864bf2ba":{"m":55,"g":94},"a4cca7fc":{"m":55,"g":94},"993956c6":{"m":55,"g":94},"f8548295":{"m":55,"g":94},"959735fc":{"m":55,"g":94},"f6772394":{"m":55,"g":94},"626a99ac":{"m":55,"g":94},"ece72491":{"m":55,"g":94},"0fb88aaa":{"m":55,"g":94},"d4de9a62":{"m":55,"g":94},"7310aede":{"m":55,"g":94},"5de9a58e":{"m":55,"g":94},"56fcd8e8":{"m":55,"g":94},"2b340adf":{"m":55,"g":94},"8586b72d":{"m":55,"g":94},"641b7d0a":{"m":55,"g":94},"0ce091a8":{"m":55,"g":94},"835f8afc":{"m":55,"g":94},"3844feb9":{"m":55,"g":94},"27f7bed7":{"m":55,"g":94},"6387098f":{"m":55,"g":94},"2a717c50":{"m":55,"g":94},"a1e697b2":{"m":55,"g":94},"a6ca736c":{"m":55,"g":94},"f62055b5":{"m":55,"g":94},"74bc9184":{"m":55,"g":94},"0f8eb153":{"m":55,"g":94},"67470bbb":{"m":55,"g":94},"cc858953":{"m":55,"g":94},"6128f7cf":{"m":55,"g":94},"a2486eb5":{"m":55,"g":94},"61dec545":{"m":55,"g":94},"96db0f66":{"m":55,"g":94},"7dc66fcb":{"m":55,"g":94},"1f09e84b":{"m":55,"g":94},"63dfab1b":{"m":55,"g":94},"ef995dae":{"m":55,"g":94},"75ae9689":{"m":55,"g":94},"95f93f49":{"m":55,"g":94},"aaac33fd":{"m":55,"g":94},"d332aa3b":{"m":55,"g":94},"c36736c8":{"m":55,"g":94},"1bf9e347":{"m":55,"g":94},"499c85f1":{"m":55,"g":94},"efc52f85":{"m":56,"g":94},"60e2fdcf":{"m":56,"g":94},"d7c0e872":{"m":56,"g":94},"31548116":{"m":56,"g":94},"53aed988":{"m":56,"g":94},"8a56b431":{"m":56,"g":94},"e835a500":{"m":56,"g":94},"23e5e50f":{"m":56,"g":94},"25e5d589":{"m":56,"g":94},"41b1db69":{"m":56,"g":94},"84967019":{"m":56,"g":94},"7d672d27":{"m":56,"g":94},"d4b17481":{"m":56,"g":94},"19ba2b0e":{"m":56,"g":94},"4e1e3cff":{"m":56,"g":94},"ef5b0ff9":{"m":57,"g":94},"6e530515":{"m":57,"g":94},"77d1210b":{"m":57,"g":94},"70dc2fbe":{"m":57,"g":94},"b438a2e5":{"m":57,"g":94},"7ca751ff":{"m":57,"g":94},"c75adfec":{"m":57,"g":94},"7722c11c":{"m":57,"g":94},"b2ed5c8e":{"m":57,"g":94},"f46f394f":{"m":57,"g":94},"2125898a":{"m":57,"g":94},"44f011d2":{"m":57,"g":94},"ed91e003":{"m":57,"g":94},"531d6ea9":{"m":57,"g":94},"dc3bee48":{"m":57,"g":94},"a74d1941":{"m":57,"g":94},"3169e66c":{"m":57,"g":94},"77395154":{"m":57,"g":94},"637de9e8":{"m":57,"g":94},"acb34072":{"m":57,"g":94},"08effbff":{"m":57,"g":94},"60bd3272":{"m":57,"g":94},"e7ebecf8":{"m":57,"g":94},"9a23c484":{"m":57,"g":94},"635a0426":{"m":57,"g":94},"2dccecf4":{"m":57,"g":94},"75ad0a14":{"m":57,"g":94},"3ccf566b":{"m":58,"g":94},"afa0341e":{"m":58,"g":94},"30828e71":{"m":58,"g":94},"e0e09fce":{"m":58,"g":94},"9c05c689":{"m":58,"g":94},"3464e57b":{"m":58,"g":94},"3815b23c":{"m":58,"g":94},"fd34f2da":{"m":58,"g":94},"8ee9a850":{"m":58,"g":94},"fd28640d":{"m":58,"g":94},"7863e436":{"m":58,"g":94},"333e3bfd":{"m":58,"g":94},"239c9d4d":{"m":58,"g":94},"855d0ba3":{"m":58,"g":94},"9254a33a":{"m":58,"g":94},"8a2681e2":{"m":58,"g":94},"5276a675":{"m":58,"g":94},"751e5ca2":{"m":58,"g":94},"7a7ac6be":{"m":58,"g":94},"d9e6ee38":{"m":58,"g":94},"03d5fbfd":{"m":59,"g":94},"1703d766":{"m":59,"g":94},"09e6e2aa":{"m":59,"g":94},"fad29f7f":{"m":59,"g":94},"35bdb485":{"m":59,"g":94},"b085e06b":{"m":59,"g":94},"763dd55d":{"m":59,"g":94},"2f0d3864":{"m":60,"g":94},"3900a94a":{"m":60,"g":94},"ded9fcd0":{"m":60,"g":94},"bc6ad367":{"m":60,"g":94},"3a22a303":{"m":60,"g":94},"bdb3929d":{"m":60,"g":94},"f5d0865b":{"m":60,"g":94},"afdee7b1":{"m":60,"g":94},"cb34d848":{"m":60,"g":94},"0f9cc6d8":{"m":60,"g":94},"c7ae474a":{"m":60,"g":94},"bdf946bf":{"m":60,"g":94},"8c8779cd":{"m":60,"g":94},"1775b963":{"m":60,"g":94},"dd2e2d27":{"m":60,"g":94},"a990daff":{"m":60,"g":94},"ba5112ff":{"m":60,"g":94},"815dce05":{"m":60,"g":94},"ad20b795":{"m":60,"g":94},"9183c23e":{"m":60,"g":94},"148254d4":{"m":60,"g":94},"a4d6d6f1":{"m":60,"g":94},"062c48d2":{"m":60,"g":94},"b6e0cfb5":{"m":60,"g":94},"0d8d97b8":{"m":60,"g":94},"0a765bbc":{"m":60,"g":94},"286cad3e":{"m":60,"g":94},"dc7eb01f":{"m":60,"g":94},"b0524c37":{"m":60,"g":94},"6c42fa22":{"m":60,"g":94},"d49b13c6":{"m":60,"g":94},"bedc4c7a":{"m":60,"g":94},"f44d1439":{"m":60,"g":94},"b6b57fc2":{"m":60,"g":94},"b4403985":{"m":60,"g":94},"339c69a2":{"m":60,"g":94},"f7074700":{"m":60,"g":94},"21ec66e5":{"m":60,"g":94},"c5210dfa":{"m":60,"g":94},"a29dd950":{"m":60,"g":94},"9c6ba248":{"m":60,"g":94},"b02da24a":{"m":60,"g":94},"bdd2827a":{"m":60,"g":94},"8c3b420e":{"m":60,"g":94},"e6f523b5":{"m":60,"g":94},"32318178":{"m":60,"g":94},"a11f8d5f":{"m":60,"g":94},"098d659c":{"m":60,"g":94},"76d14f8c":{"m":60,"g":94},"b08c308e":{"m":60,"g":94},"f624901c":{"m":61,"g":94},"f0e15dc6":{"m":61,"g":94},"f1769586":{"m":61,"g":94},"5d6e9467":{"m":61,"g":94},"a47bf391":{"m":61,"g":94},"b1706469":{"m":61,"g":94},"5413ec2b":{"m":61,"g":94},"f290bd43":{"m":61,"g":94},"8f157893":{"m":61,"g":94},"2db03a04":{"m":61,"g":94},"5cc11705":{"m":61,"g":94},"11fffbc9":{"m":61,"g":94},"4f077c01":{"m":61,"g":94},"679c3bca":{"m":61,"g":94},"656aed58":{"m":61,"g":94},"b5fb4ef5":{"m":61,"g":94},"2e6346fc":{"m":61,"g":94},"977f785d":{"m":61,"g":94},"8a690612":{"m":61,"g":94},"694e4192":{"m":61,"g":94},"b22f3f64":{"m":61,"g":94},"6fb57683":{"m":61,"g":94},"51caee74":{"m":61,"g":94},"58f9060e":{"m":61,"g":94},"bdc1acf6":{"m":61,"g":94},"6d08ce2a":{"m":61,"g":94},"380930a9":{"m":61,"g":94},"9dec582d":{"m":61,"g":94},"b01febdc":{"m":61,"g":94},"1acbaf1b":{"m":61,"g":94},"287427e2":{"m":61,"g":94},"b8574f69":{"m":61,"g":94},"2855caa4":{"m":61,"g":94},"2329e1dd":{"m":61,"g":94},"0f3eb1d2":{"m":61,"g":94},"06dd2eab":{"m":61,"g":94},"439f6580":{"m":61,"g":94},"b3e99dfb":{"m":62,"g":94},"f005758f":{"m":62,"g":94},"f5c6c667":{"m":62,"g":94},"cc0485be":{"m":62,"g":94},"b8cd09f2":{"m":62,"g":94},"c19d8482":{"m":62,"g":94},"80002562":{"m":62,"g":94},"46d44318":{"m":62,"g":94},"923f5183":{"m":62,"g":94},"d08c77c4":{"m":62,"g":94},"c1e097ca":{"m":62,"g":94},"6ec75e62":{"m":62,"g":94},"d855653b":{"m":62,"g":94},"336ff5b9":{"m":62,"g":94},"3b141e15":{"m":62,"g":94},"6249e4a1":{"m":62,"g":94},"f3516c28":{"m":62,"g":94},"17de02f9":{"m":62,"g":94},"51ab3ccf":{"m":62,"g":94},"67008f4b":{"m":62,"g":94},"4536d724":{"m":62,"g":94},"41d7e5b7":{"m":62,"g":94},"20a9f5df":{"m":62,"g":94},"42f39099":{"m":62,"g":94},"72c77763":{"m":62,"g":94},"4093aa46":{"m":62,"g":94},"e808c1df":{"m":62,"g":94},"a18ab81d":{"m":62,"g":94},"0bb0f763":{"m":62,"g":94},"85b2e057":{"m":62,"g":94},"a879c2fb":{"m":62,"g":94},"e2b16c47":{"m":62,"g":94},"c4f9707e":{"m":62,"g":94},"197cbf9b":{"m":62,"g":94},"e94fb7cb":{"m":63,"g":94},"b5caa22d":{"m":63,"g":94},"73401fd0":{"m":63,"g":94},"89cd9235":{"m":63,"g":94},"dc188132":{"m":63,"g":94},"10bfce71":{"m":63,"g":94},"583697cd":{"m":63,"g":94},"2584f6d9":{"m":63,"g":94},"51e87f6f":{"m":63,"g":94},"09bcbe01":{"m":63,"g":94},"03464890":{"m":63,"g":94},"44a96697":{"m":63,"g":94},"1a820e38":{"m":63,"g":94},"0ffcfdf4":{"m":63,"g":94},"cd493b5a":{"m":63,"g":94},"61f42b57":{"m":63,"g":94},"e403d237":{"m":63,"g":94},"3bcf5ece":{"m":63,"g":94},"2c05f81f":{"m":63,"g":94},"d77caa2b":{"m":63,"g":94},"8b6a4486":{"m":63,"g":94},"a69cb5cf":{"m":63,"g":94},"def5c318":{"m":63,"g":94},"3fc2b625":{"m":63,"g":94},"6ada05d0":{"m":63,"g":94},"24cafe31":{"m":63,"g":94},"5a176c92":{"m":63,"g":94},"4719c1d0":{"m":63,"g":94},"ef18b0ed":{"m":63,"g":94},"53cc91e5":{"m":63,"g":94},"d33cbb7e":{"m":63,"g":94},"23196d52":{"m":63,"g":94},"93b77c8e":{"m":63,"g":94},"7906d1d2":{"m":63,"g":94},"81d27c8e":{"m":63,"g":94},"4d4cdb3f":{"m":63,"g":94},"2bd18e2d":{"m":63,"g":94},"83452dbb":{"m":63,"g":94},"3d93f84a":{"m":63,"g":94},"c2f212d6":{"m":63,"g":94},"e2cdc8a5":{"m":63,"g":94},"2add697d":{"m":63,"g":94},"6f98c586":{"m":63,"g":94},"656dcc1a":{"m":63,"g":94},"8af7048d":{"m":63,"g":94},"d3024f4f":{"m":63,"g":94},"13387e6b":{"m":63,"g":94},"120c3634":{"m":63,"g":94},"78e5b22f":{"m":63,"g":94},"7a15e9ad":{"m":63,"g":94},"dc2ac0cb":{"m":63,"g":94},"d47c5101":{"m":63,"g":94},"033c715b":{"m":63,"g":94},"d06c1ab5":{"m":63,"g":94},"c5644cac":{"m":63,"g":94},"53e6552f":{"m":63,"g":94},"5dc54f1a":{"m":63,"g":94},"f3e9b489":{"m":63,"g":94},"6a7973ad":{"m":63,"g":94},"63051738":{"m":63,"g":94},"a8ccacc8":{"m":63,"g":94},"0427416b":{"m":63,"g":94},"bf3edc2c":{"m":63,"g":94},"78e974b2":{"m":63,"g":94},"bc6915e3":{"m":63,"g":94},"a883f079":{"m":63,"g":94},"8b6ce52e":{"m":63,"g":94},"58f3f2b8":{"m":63,"g":94},"93d69061":{"m":63,"g":94},"e00e5385":{"m":63,"g":94},"a2f602b5":{"m":63,"g":94},"8f2c522a":{"m":63,"g":94},"75964177":{"m":63,"g":94},"2dc957d4":{"m":63,"g":94},"bf8d07a6":{"m":63,"g":94},"ab317936":{"m":63,"g":94},"b7f3fec1":{"m":63,"g":94},"58f42b1d":{"m":63,"g":94},"767c9dec":{"m":63,"g":94},"a53454c5":{"m":63,"g":94},"6cb3974e":{"m":63,"g":94},"f65c13b5":{"m":63,"g":94},"b803b395":{"m":63,"g":94},"bfbda62c":{"m":63,"g":94},"4ab43cfb":{"m":64,"g":94},"2f79f588":{"m":64,"g":94},"8a96f749":{"m":64,"g":94},"827aa873":{"m":64,"g":94},"f8ca66fb":{"m":64,"g":94},"53cef815":{"m":64,"g":94},"351a72d4":{"m":64,"g":94},"514f37c3":{"m":64,"g":94},"52c03f16":{"m":64,"g":94},"741fccd7":{"m":64,"g":94},"1e3e5215":{"m":64,"g":94},"fb11a439":{"m":64,"g":94},"af02f99b":{"m":64,"g":94},"9472e699":{"m":64,"g":94},"1acc1f56":{"m":64,"g":94},"b045841b":{"m":64,"g":94},"f265d15b":{"m":64,"g":94},"02431b9a":{"m":64,"g":94},"1dda8c5e":{"m":64,"g":94},"7e097613":{"m":64,"g":94},"f4a92f4b":{"m":64,"g":94},"318260c0":{"m":64,"g":94},"4a612531":{"m":64,"g":94},"d1a08632":{"m":64,"g":94},"f8b28e46":{"m":64,"g":94},"82392da8":{"m":64,"g":94},"95f789ad":{"m":64,"g":94},"4f118a39":{"m":64,"g":94},"66283dbc":{"m":64,"g":94},"822bae8c":{"m":64,"g":94},"8e48ca8c":{"m":64,"g":94},"27acf63b":{"m":64,"g":94},"da6f8081":{"m":64,"g":94},"9286740e":{"m":64,"g":94},"896c0744":{"m":64,"g":94},"c23d5706":{"m":64,"g":94},"67ad4338":{"m":64,"g":94},"3cab5f71":{"m":64,"g":94},"14e754a8":{"m":64,"g":94},"98522149":{"m":64,"g":94},"5d9d15e7":{"m":64,"g":94},"665e5e85":{"m":64,"g":94},"a22f60a3":{"m":64,"g":94},"04f0b4cb":{"m":64,"g":94},"4505a436":{"m":64,"g":94},"685a5738":{"m":64,"g":94},"153b414e":{"m":64,"g":94},"6619f48e":{"m":64,"g":94},"3ed0a547":{"m":64,"g":94},"8d8ef849":{"m":64,"g":94},"9a0cc2e9":{"m":64,"g":94},"7bad7e75":{"m":64,"g":94},"1c4e0d24":{"m":64,"g":94},"54bac8af":{"m":64,"g":94},"5de4051b":{"m":64,"g":94},"e0cd65c2":{"m":64,"g":94},"f1b68618":{"m":64,"g":94},"0da0989a":{"m":64,"g":94},"07a22cbb":{"m":64,"g":94},"3d0bfa3e":{"m":64,"g":94},"1f6cf0d4":{"m":64,"g":94},"553f5a3f":{"m":64,"g":94},"ac2dc35d":{"m":64,"g":94},"3e032c07":{"m":64,"g":94},"44e12ce4":{"m":64,"g":94},"a547aad6":{"m":64,"g":94},"ea535dc5":{"m":64,"g":94},"862bcff8":{"m":64,"g":94},"8b84e69f":{"m":64,"g":94},"5de50653":{"m":64,"g":94},"c0bf9bf1":{"m":64,"g":94},"022614d2":{"m":64,"g":94},"b8ab989f":{"m":64,"g":94},"b3393e94":{"m":64,"g":94},"ddc2001f":{"m":64,"g":94},"806a3002":{"m":64,"g":94},"0d2148ef":{"m":64,"g":94},"bf669606":{"m":64,"g":94},"b2bd8f44":{"m":64,"g":94},"9d9b482a":{"m":64,"g":94},"7353fb9b":{"m":64,"g":94},"bcda0c9e":{"m":64,"g":94},"9f8f2c7f":{"m":64,"g":94},"6fc37bd8":{"m":64,"g":94},"3d8f1c9b":{"m":64,"g":94},"a42213db":{"m":64,"g":94},"0ac019f1":{"m":64,"g":94},"5a0d680a":{"m":64,"g":94},"a4331cd2":{"m":64,"g":94},"ec1c21cd":{"m":64,"g":94},"6c856b4f":{"m":64,"g":94},"287d07a6":{"m":64,"g":94},"d2571dd5":{"m":64,"g":94},"b730aa6b":{"m":64,"g":94},"60b2a44a":{"m":64,"g":94},"949b3fbf":{"m":64,"g":94},"da4e8b38":{"m":64,"g":94},"af6c5357":{"m":64,"g":94},"3ad4cd49":{"m":64,"g":94},"3a8428ec":{"m":64,"g":94},"0311ce8e":{"m":64,"g":94},"5dfcacfc":{"m":64,"g":94},"41a0ccd4":{"m":64,"g":94},"cf0f7eaf":{"m":65,"g":94},"b49d6d0f":{"m":65,"g":94},"c02e3139":{"m":65,"g":94},"734daedd":{"m":65,"g":94},"3ee62235":{"m":65,"g":94},"9829e77e":{"m":65,"g":94},"cde4bbd5":{"m":65,"g":94},"9602c2aa":{"m":65,"g":94},"e81d7f11":{"m":65,"g":94},"222ce6f1":{"m":65,"g":94},"468d23cf":{"m":65,"g":94},"c38b5fb4":{"m":65,"g":94},"20453cef":{"m":65,"g":94},"9f635ea5":{"m":65,"g":94},"76285fde":{"m":65,"g":94},"988d0a4b":{"m":65,"g":94},"81262c7b":{"m":65,"g":94},"27aeb4b7":{"m":65,"g":94},"7b9b4f44":{"m":65,"g":94},"08104b56":{"m":65,"g":94},"cf142b6e":{"m":65,"g":94},"7aad8d18":{"m":66,"g":94},"76fa2d15":{"m":66,"g":94},"7ab84948":{"m":66,"g":94},"4885b908":{"m":66,"g":94},"c2723a42":{"m":66,"g":94},"c7256ca8":{"m":66,"g":94},"6186a8f8":{"m":66,"g":94},"a07364cc":{"m":66,"g":94},"2c1a695f":{"m":66,"g":94},"d39899e8":{"m":66,"g":94},"70817a7e":{"m":66,"g":94},"7b5a3741":{"m":66,"g":94},"4b6f62e2":{"m":66,"g":94},"897e2e25":{"m":66,"g":94},"d54cee14":{"m":66,"g":94},"00fa7d04":{"m":66,"g":94},"013021b6":{"m":66,"g":94},"3c8ac78d":{"m":66,"g":94},"455bfe8d":{"m":66,"g":94},"28b0a62b":{"m":66,"g":94},"566d61d9":{"m":66,"g":94},"55f5fc68":{"m":66,"g":94},"c27c378a":{"m":66,"g":94},"d9eb9358":{"m":66,"g":94},"959dca4f":{"m":66,"g":94},"f2b3a318":{"m":66,"g":94},"ad674097":{"m":66,"g":94},"8db776f0":{"m":66,"g":94},"4eb4b401":{"m":66,"g":94},"17dbf976":{"m":66,"g":94},"53179026":{"m":66,"g":94},"d7c0b32f":{"m":66,"g":94},"7b020cca":{"m":66,"g":94},"7876279e":{"m":66,"g":94},"34e405e0":{"m":66,"g":94},"1ebe1d6d":{"m":66,"g":94},"7811bfda":{"m":66,"g":94},"656f7fc1":{"m":66,"g":94},"c1f5f99f":{"m":67,"g":94},"fa82dfcc":{"m":67,"g":94},"5da3d21c":{"m":67,"g":94},"f2870376":{"m":67,"g":94},"f9905d59":{"m":67,"g":94},"45c87e08":{"m":67,"g":94},"2b1808ce":{"m":67,"g":94},"e868d0b6":{"m":67,"g":94},"591e751e":{"m":67,"g":94},"40022d07":{"m":67,"g":94},"823148e7":{"m":67,"g":94},"76ca91df":{"m":67,"g":94},"cdae77b0":{"m":67,"g":94},"adeee152":{"m":67,"g":94},"6792411e":{"m":67,"g":94},"7348d962":{"m":67,"g":94},"25ed22b6":{"m":67,"g":94},"200d3b16":{"m":67,"g":94},"ad349985":{"m":67,"g":94},"32de54ed":{"m":67,"g":94},"2d9c3195":{"m":67,"g":94},"07e58a2d":{"m":67,"g":94},"04d8cd20":{"m":67,"g":94},"a322051e":{"m":67,"g":94},"de553334":{"m":67,"g":94},"cddb1cdf":{"m":68,"g":94},"fa1b40e0":{"m":68,"g":94},"c45cab1c":{"m":68,"g":94},"27c4c9cf":{"m":68,"g":94},"52a492a1":{"m":68,"g":94},"36f6fc50":{"m":68,"g":94},"d8727275":{"m":68,"g":94},"6239d0b2":{"m":68,"g":94},"4cfd3add":{"m":68,"g":94},"20cf910d":{"m":68,"g":94},"0af1d239":{"m":68,"g":94},"85986bb9":{"m":68,"g":94},"64c87135":{"m":68,"g":94},"1646149a":{"m":68,"g":94},"bc72e5bd":{"m":68,"g":94},"014cab4d":{"m":68,"g":94},"4d2dbeac":{"m":68,"g":94},"29daf498":{"m":68,"g":94},"6702592d":{"m":68,"g":94},"60abdb3e":{"m":68,"g":94},"7b4e61ff":{"m":68,"g":94},"6222e1c2":{"m":68,"g":94},"fad315cb":{"m":68,"g":94},"f90db8bc":{"m":68,"g":94},"d8ad5970":{"m":68,"g":94},"849f58d6":{"m":68,"g":94},"64480df4":{"m":68,"g":94},"4530136e":{"m":68,"g":94},"0a6f18f0":{"m":68,"g":94},"e0b9a423":{"m":69,"g":94},"e0821425":{"m":69,"g":94},"70f894b8":{"m":69,"g":94},"368de366":{"m":69,"g":94},"20de05a7":{"m":69,"g":94},"f076328b":{"m":69,"g":94},"bf2a7087":{"m":69,"g":94},"871a4aa1":{"m":69,"g":94},"98eecbda":{"m":69,"g":94},"4430c0a5":{"m":69,"g":94},"640363ad":{"m":69,"g":94},"8616357a":{"m":69,"g":94},"8adbc78b":{"m":69,"g":94},"45e3a7bc":{"m":69,"g":94},"b96e92e6":{"m":69,"g":94},"693c2600":{"m":69,"g":94},"ced68066":{"m":69,"g":94},"b8318aec":{"m":69,"g":94},"2f482210":{"m":69,"g":94},"d81ac443":{"m":69,"g":94},"2491cc92":{"m":69,"g":94},"67c5de92":{"m":69,"g":94},"1e2cf2b5":{"m":69,"g":94},"9490d157":{"m":69,"g":94},"eefcbdd3":{"m":69,"g":94},"7e6d5fc6":{"m":69,"g":94},"cadd5dbe":{"m":69,"g":94},"bb418ced":{"m":69,"g":94},"fdf04a14":{"m":69,"g":94},"5f0e7de3":{"m":69,"g":94},"2f47d710":{"m":69,"g":94},"4fe92bfc":{"m":69,"g":94},"d23cb9a0":{"m":69,"g":94},"2d611323":{"m":69,"g":94},"e782eb7e":{"m":70,"g":94},"e319153b":{"m":70,"g":94},"32b44d2f":{"m":70,"g":94},"5f1a485d":{"m":70,"g":94},"c9565e49":{"m":70,"g":94},"d03c4c25":{"m":70,"g":94},"8f13377d":{"m":70,"g":94},"3d4a8f9b":{"m":70,"g":94},"7474bed8":{"m":70,"g":94},"03caefeb":{"m":70,"g":94},"bcc213df":{"m":70,"g":94},"39416e39":{"m":70,"g":94},"231c40d8":{"m":70,"g":94},"bbc47c34":{"m":70,"g":94},"dfce9269":{"m":70,"g":94},"6718b109":{"m":70,"g":94},"7711ac6e":{"m":70,"g":94},"7443197a":{"m":70,"g":94},"862dd76c":{"m":70,"g":94},"fb4c9c3a":{"m":70,"g":94},"d973c78e":{"m":70,"g":94},"6ce6eabb":{"m":70,"g":94},"4e23c961":{"m":70,"g":94},"3efbdf68":{"m":70,"g":94},"6cc30955":{"m":70,"g":94},"31eec35b":{"m":70,"g":94},"ac963be2":{"m":70,"g":94},"a5375adc":{"m":71,"g":94},"75d171a9":{"m":71,"g":94},"714f3e63":{"m":71,"g":94},"c38f3aed":{"m":71,"g":94},"2e6be53e":{"m":71,"g":94},"fc671f66":{"m":72,"g":94},"197751e9":{"m":72,"g":94},"d2d0d061":{"m":72,"g":94},"25482edb":{"m":72,"g":94},"62b362b1":{"m":72,"g":94},"44d76463":{"m":72,"g":94},"cd85b78f":{"m":72,"g":94},"0aaccbbf":{"m":72,"g":94},"357671e2":{"m":72,"g":94},"e70fa279":{"m":72,"g":94},"abe74b7b":{"m":72,"g":94},"70b3c6ee":{"m":72,"g":94},"ef9d3b3c":{"m":72,"g":94},"fc91d08a":{"m":72,"g":94},"71ab0dab":{"m":72,"g":94},"d3d4d767":{"m":72,"g":94},"5be8f1ed":{"m":72,"g":94},"e5760bc4":{"m":72,"g":94},"56a724eb":{"m":72,"g":94},"583d6af7":{"m":72,"g":94},"e074d84e":{"m":72,"g":94},"4725e3f6":{"m":72,"g":94},"77a3954b":{"m":72,"g":94},"03b0364f":{"m":72,"g":94},"2dd7d0c5":{"m":72,"g":94},"0d4e3228":{"m":72,"g":94},"926f8efc":{"m":72,"g":94},"9545bfb2":{"m":72,"g":94},"37373ef2":{"m":72,"g":94},"61261b39":{"m":72,"g":94},"19120f71":{"m":72,"g":94},"2415ec38":{"m":72,"g":94},"87f671ab":{"m":72,"g":94},"51d25405":{"m":72,"g":94},"e0a2c963":{"m":72,"g":94},"12f2e6c3":{"m":72,"g":94},"95575aa7":{"m":72,"g":94},"11eea69e":{"m":72,"g":94},"1baa9e6c":{"m":72,"g":94},"911fcd09":{"m":72,"g":94},"9fafa62d":{"m":72,"g":94},"146ac8df":{"m":72,"g":94},"57a404fd":{"m":72,"g":94},"2796fbb5":{"m":72,"g":94},"935cda94":{"m":72,"g":94},"110e0066":{"m":72,"g":94},"6b45a21d":{"m":72,"g":94},"a7000a76":{"m":72,"g":94},"1a8f995c":{"m":72,"g":94},"a3ab768a":{"m":72,"g":94},"66301e12":{"m":72,"g":94},"ac238727":{"m":72,"g":94},"0194948f":{"m":72,"g":94},"b4d34cd3":{"m":72,"g":94},"728e175f":{"m":72,"g":94},"9e1014cf":{"m":72,"g":94},"fa561067":{"m":72,"g":94},"7fbab730":{"m":72,"g":94},"b7e274f2":{"m":72,"g":94},"9cf40772":{"m":72,"g":94},"d3fe9bae":{"m":72,"g":94},"00ce7e31":{"m":72,"g":94},"50f28f65":{"m":72,"g":94},"90a55e25":{"m":72,"g":94},"407e2b92":{"m":72,"g":94},"40782f05":{"m":72,"g":94},"18bb216c":{"m":72,"g":94},"6b859e7d":{"m":72,"g":94},"930da877":{"m":72,"g":94},"3f8a4414":{"m":72,"g":94},"aceb4201":{"m":72,"g":94},"90a4b7d9":{"m":72,"g":94},"f3b99f73":{"m":72,"g":94},"9e74ee91":{"m":72,"g":94},"77a6c9d2":{"m":72,"g":94},"e3e0bc50":{"m":72,"g":94},"bac414ab":{"m":72,"g":94},"eec3f6d1":{"m":72,"g":94},"90bc26a8":{"m":72,"g":94},"ec0a72c2":{"m":72,"g":94},"1c96fa86":{"m":72,"g":94},"bc20e93f":{"m":72,"g":94},"d3887852":{"m":72,"g":94},"564bdf29":{"m":72,"g":94},"5d860168":{"m":72,"g":94},"d2815879":{"m":72,"g":94},"b0df5d24":{"m":72,"g":94},"3e02526b":{"m":72,"g":94},"d8a98a2c":{"m":72,"g":94},"0519269d":{"m":72,"g":94},"d6898dd2":{"m":72,"g":94},"71ed0183":{"m":72,"g":94},"8b681d77":{"m":72,"g":94},"194eea17":{"m":72,"g":94},"acd1a159":{"m":72,"g":94},"7c1692aa":{"m":72,"g":94},"8f019c7d":{"m":72,"g":94},"7551498a":{"m":72,"g":94},"44a2c4bd":{"m":72,"g":94},"c9fc4a9d":{"m":72,"g":94},"21463e32":{"m":72,"g":94},"3dc9ff3c":{"m":72,"g":94},"06427dfa":{"m":72,"g":94},"60524920":{"m":72,"g":94},"10771026":{"m":72,"g":94},"4606e2a3":{"m":72,"g":94},"127998cc":{"m":72,"g":94},"c0bb9eb3":{"m":72,"g":94},"7036d6fc":{"m":72,"g":94},"6ce9dbe8":{"m":72,"g":94},"3758d209":{"m":72,"g":94},"faf29e0b":{"m":72,"g":94},"b0743ea0":{"m":72,"g":94},"60b771c8":{"m":72,"g":94},"d7934cde":{"m":72,"g":94},"62bbd343":{"m":72,"g":94},"f2388f6b":{"m":72,"g":94},"c9745ee0":{"m":72,"g":94},"1a6e9757":{"m":72,"g":94},"b1100846":{"m":72,"g":94},"27a46317":{"m":72,"g":94},"c9795808":{"m":72,"g":94},"6c7a152c":{"m":72,"g":94},"4d2a88bd":{"m":72,"g":94},"45360b2f":{"m":72,"g":94},"3f41b184":{"m":72,"g":94},"45205d88":{"m":72,"g":94},"90876940":{"m":72,"g":94},"a3339d8c":{"m":72,"g":94},"14d90617":{"m":72,"g":94},"d37f9551":{"m":72,"g":94},"c66b2c9c":{"m":72,"g":94},"20b765a2":{"m":72,"g":94},"e3107222":{"m":72,"g":94},"e074e76b":{"m":72,"g":94},"4592afc2":{"m":72,"g":94},"9af0e21e":{"m":72,"g":94},"c7c79b16":{"m":72,"g":94},"d8d75d25":{"m":72,"g":94},"1df6eabd":{"m":72,"g":94},"0c227ee3":{"m":72,"g":94},"5c54ef03":{"m":72,"g":94},"c6a48521":{"m":72,"g":94},"4f678c87":{"m":72,"g":94},"e79f7420":{"m":72,"g":94},"ac053100":{"m":72,"g":94},"d5d80ab4":{"m":72,"g":94},"ddcf9fe3":{"m":72,"g":94},"6252ade9":{"m":72,"g":94},"1eb8eade":{"m":72,"g":94},"3c7bfd7e":{"m":72,"g":94},"bb121214":{"m":72,"g":94},"55de40f7":{"m":72,"g":94},"6b0aeb58":{"m":72,"g":94},"bb3e5268":{"m":72,"g":94},"f93e9158":{"m":72,"g":94},"55a7ec38":{"m":72,"g":94},"fe0673f1":{"m":72,"g":94},"99c1b9d2":{"m":72,"g":94},"634a3561":{"m":72,"g":94},"424848d2":{"m":72,"g":94},"e5ce395a":{"m":72,"g":94},"f983213a":{"m":72,"g":94},"67fc595b":{"m":72,"g":94},"07ab4d4a":{"m":72,"g":94},"522e18ea":{"m":72,"g":94},"c51dc2cc":{"m":72,"g":94},"ddf39d3f":{"m":72,"g":94},"2eab1132":{"m":72,"g":94},"058d199d":{"m":72,"g":94},"9c58e68b":{"m":73,"g":94},"d03b3467":{"m":73,"g":94},"ab7fba0e":{"m":73,"g":94},"bc1534ff":{"m":73,"g":94},"3a391812":{"m":73,"g":94},"800bf018":{"m":73,"g":94},"b16af90b":{"m":73,"g":94},"98c73d71":{"m":73,"g":94},"fcc2e37f":{"m":73,"g":94},"0804dd11":{"m":73,"g":94},"55dc8e4d":{"m":73,"g":94},"02e9e9f1":{"m":73,"g":94},"8f0b6313":{"m":73,"g":94},"b9b3b098":{"m":73,"g":94},"aee30630":{"m":73,"g":94},"286e6540":{"m":73,"g":94},"718c391f":{"m":73,"g":94},"6aaeb848":{"m":74,"g":94},"3623b6a7":{"m":74,"g":94},"4ff12642":{"m":74,"g":94},"2a4cbad8":{"m":74,"g":94},"2937387a":{"m":74,"g":94},"cf721fde":{"m":74,"g":94},"45de8971":{"m":74,"g":94},"71046fcd":{"m":74,"g":94},"c76040e3":{"m":74,"g":94},"2f6bacee":{"m":74,"g":94},"40148041":{"m":74,"g":94},"ad46550d":{"m":74,"g":94},"14344caa":{"m":74,"g":94},"f7f88b70":{"m":74,"g":94},"18c27131":{"m":74,"g":94},"ccdd10c8":{"m":74,"g":94},"76f6c0eb":{"m":74,"g":94},"959a3143":{"m":74,"g":94},"6412c5e4":{"m":74,"g":94},"0c020860":{"m":74,"g":94},"85ef7f64":{"m":74,"g":94},"f1cf6eef":{"m":74,"g":94},"0a59a465":{"m":74,"g":94},"aff79f10":{"m":74,"g":94},"01603318":{"m":74,"g":94},"56c39a05":{"m":74,"g":94},"4068e012":{"m":74,"g":94},"817d4370":{"m":74,"g":94},"c550e52f":{"m":74,"g":94},"e35a93fa":{"m":74,"g":94},"2c3656f2":{"m":74,"g":94},"d40ee62b":{"m":74,"g":94},"91b19949":{"m":74,"g":94},"7c866711":{"m":74,"g":94},"10b544ae":{"m":74,"g":94},"01090e8a":{"m":74,"g":94},"6f43a9b9":{"m":74,"g":94},"0540fef7":{"m":74,"g":94},"481f608b":{"m":74,"g":94},"ed91561f":{"m":74,"g":94},"6e7239f9":{"m":74,"g":94},"0a3960f2":{"m":74,"g":94},"07f94463":{"m":74,"g":94},"e0917e6b":{"m":74,"g":94},"7130a7ce":{"m":74,"g":94},"8f1f614e":{"m":74,"g":94},"7140ba35":{"m":74,"g":94},"d1da58e2":{"m":74,"g":94},"1cf63485":{"m":74,"g":94},"ff2ce0b8":{"m":74,"g":94},"0f2a2e3c":{"m":74,"g":94},"690e1f23":{"m":74,"g":94},"00f42707":{"m":74,"g":94},"6a02b32d":{"m":74,"g":94},"3a08f546":{"m":74,"g":94},"dce303e2":{"m":74,"g":94},"4d27eb9a":{"m":74,"g":94},"d3ecd632":{"m":74,"g":94},"cd909455":{"m":74,"g":94},"bde24ab3":{"m":74,"g":94},"bf2eefc0":{"m":74,"g":94},"5524e7d0":{"m":74,"g":94},"e187a3d5":{"m":74,"g":94},"3dd4feae":{"m":74,"g":94},"2ac189ed":{"m":74,"g":94},"5a6400ee":{"m":74,"g":94},"cf0ccd40":{"m":74,"g":94},"3d56585a":{"m":74,"g":94},"00d25a7f":{"m":74,"g":94},"1a5023e0":{"m":74,"g":94},"23308a90":{"m":74,"g":94},"ac698850":{"m":74,"g":94},"aa957102":{"m":74,"g":94},"007f8b3d":{"m":74,"g":94},"4455b26e":{"m":74,"g":94},"c553e160":{"m":74,"g":94},"7c0541b3":{"m":74,"g":94},"e8a69e4d":{"m":74,"g":94},"fbd56002":{"m":74,"g":94},"730d084f":{"m":74,"g":94},"4a05bdfa":{"m":74,"g":94},"eb06dbcb":{"m":74,"g":94},"9dfafa74":{"m":74,"g":94},"f1d09a65":{"m":74,"g":94},"df84ab2a":{"m":74,"g":94},"34c88987":{"m":74,"g":94},"0dd6cda2":{"m":74,"g":94},"9fb48f95":{"m":74,"g":94},"89ccb533":{"m":74,"g":94},"dceb256f":{"m":74,"g":94},"0e90ae62":{"m":74,"g":94},"1361ab9e":{"m":74,"g":94},"5c7dd14b":{"m":74,"g":94},"8abf74e3":{"m":74,"g":94},"ee132a45":{"m":74,"g":94},"79a321af":{"m":74,"g":94},"6eec3cdc":{"m":74,"g":94},"48473684":{"m":74,"g":94},"b3251e9f":{"m":74,"g":94},"2cadd51d":{"m":74,"g":94},"4a893d14":{"m":74,"g":94},"8d323e95":{"m":74,"g":94},"0fe7c13b":{"m":74,"g":94},"08c4d764":{"m":74,"g":94},"96d0e37f":{"m":74,"g":94},"90bb2be2":{"m":74,"g":94},"b93ef5e5":{"m":74,"g":94},"d4017a6b":{"m":74,"g":94},"d052f4c8":{"m":74,"g":94},"e1aaa79a":{"m":74,"g":94},"20c81199":{"m":74,"g":94},"70866b6f":{"m":74,"g":94},"eb61f5c9":{"m":74,"g":94},"0beea450":{"m":74,"g":94},"c827c671":{"m":74,"g":94},"b55a621f":{"m":74,"g":94},"ffa1b3e3":{"m":74,"g":94},"7e3bb527":{"m":74,"g":94},"96263f27":{"m":74,"g":94},"9376ac36":{"m":74,"g":94},"94a2b9d3":{"m":74,"g":94},"3c3eb374":{"m":74,"g":94},"d557319a":{"m":74,"g":94},"95085d65":{"m":74,"g":94},"c7f25446":{"m":74,"g":94},"63ee26d1":{"m":74,"g":94},"ad55f171":{"m":74,"g":94},"361971b8":{"m":74,"g":94},"13bc39c5":{"m":74,"g":94},"9854a18a":{"m":74,"g":94},"ebddb65a":{"m":74,"g":94},"19fd57bc":{"m":74,"g":94},"ba80c102":{"m":75,"g":94},"fbdb5050":{"m":75,"g":94},"f0afaf52":{"m":75,"g":94},"85d2365d":{"m":75,"g":94},"5fe79605":{"m":75,"g":94},"c6d7f8d3":{"m":75,"g":94},"a5a892ff":{"m":75,"g":94},"8e66fbec":{"m":75,"g":94},"f141298a":{"m":75,"g":94},"4fea040c":{"m":75,"g":94},"1099f6c9":{"m":76,"g":94},"04e3ff69":{"m":76,"g":94},"45fdf1f7":{"m":76,"g":94},"d89c0e4b":{"m":76,"g":94},"fa3c9e06":{"m":76,"g":94},"0d658ac3":{"m":76,"g":94},"ced35a06":{"m":76,"g":94},"26f07294":{"m":76,"g":94},"34e07a65":{"m":76,"g":94},"15ddd843":{"m":76,"g":94},"52029bd1":{"m":76,"g":94},"eb934bdf":{"m":76,"g":94},"e45ae444":{"m":76,"g":94},"ac3fae84":{"m":76,"g":94},"2d1b83e5":{"m":76,"g":94},"199bb01d":{"m":76,"g":94},"6b7038ba":{"m":76,"g":94},"57eec0bf":{"m":76,"g":94},"f01b0925":{"m":76,"g":94},"14269198":{"m":76,"g":94},"9b7cf9ee":{"m":76,"g":94},"1e86457c":{"m":76,"g":94},"64129fa6":{"m":76,"g":94},"e9f8e423":{"m":76,"g":94},"22c3702e":{"m":76,"g":94},"4c584fc6":{"m":76,"g":94},"77cf771e":{"m":76,"g":94},"8154de5a":{"m":76,"g":94},"c11cfda0":{"m":76,"g":94},"64edeb79":{"m":76,"g":94},"65c24c28":{"m":76,"g":94},"3980ff1b":{"m":76,"g":94},"5d7edc8e":{"m":76,"g":94},"af6535e7":{"m":76,"g":94},"93cf7fc5":{"m":76,"g":94},"2a206b22":{"m":76,"g":94},"4d253057":{"m":76,"g":94},"11577ced":{"m":76,"g":94},"ca75741e":{"m":76,"g":94},"c6d549e7":{"m":76,"g":94},"3c09548d":{"m":76,"g":94},"8796cebb":{"m":76,"g":94},"c2bd094d":{"m":76,"g":94},"f8f9244a":{"m":76,"g":94},"ecbfe58b":{"m":76,"g":94},"8f163b16":{"m":76,"g":94},"e7a8610d":{"m":76,"g":94},"a2cc62a6":{"m":76,"g":94},"fb888603":{"m":76,"g":94},"321ab756":{"m":76,"g":94},"38f25e87":{"m":76,"g":94},"8cd42504":{"m":76,"g":94},"6a384d5c":{"m":76,"g":94},"f69e0696":{"m":76,"g":94},"f6ab4ca6":{"m":76,"g":94},"c7c7dbeb":{"m":76,"g":94},"417fc72f":{"m":76,"g":94},"c6ec7029":{"m":76,"g":94},"4c56e5db":{"m":76,"g":94},"7b5fc719":{"m":76,"g":94},"ad4e58bf":{"m":76,"g":94},"bfb03c61":{"m":76,"g":94},"b36ab493":{"m":76,"g":94},"9e93ef3f":{"m":76,"g":94},"fad86a68":{"m":76,"g":94},"df7014a8":{"m":76,"g":94},"49420741":{"m":76,"g":94},"ba52fd18":{"m":76,"g":94},"b6944f97":{"m":76,"g":94},"f44db16c":{"m":76,"g":94},"f9c53cbb":{"m":76,"g":94},"90532b76":{"m":76,"g":94},"c0e9a36c":{"m":76,"g":94},"588865f0":{"m":76,"g":94},"3196999f":{"m":76,"g":94},"9e0186f3":{"m":76,"g":94},"8baf9a0c":{"m":76,"g":94},"c7872985":{"m":76,"g":94},"45212ce1":{"m":76,"g":94},"c16b33cc":{"m":76,"g":94},"2d004512":{"m":76,"g":94},"804d250a":{"m":76,"g":94},"dd865bef":{"m":76,"g":94},"d373a48c":{"m":76,"g":94},"98be3bd3":{"m":76,"g":94},"a98290ae":{"m":76,"g":94},"9b81f9bd":{"m":76,"g":94},"f81a27f6":{"m":76,"g":94},"988ab646":{"m":76,"g":94},"3ded4b21":{"m":76,"g":94},"f4d7ab7a":{"m":76,"g":94},"c38ca4fc":{"m":76,"g":94},"82dec1f7":{"m":76,"g":94},"5f9b2c62":{"m":76,"g":94},"5493c334":{"m":76,"g":94},"f2ab37e5":{"m":76,"g":94},"91ba98fe":{"m":76,"g":94},"c614dbdf":{"m":76,"g":94},"927ca935":{"m":76,"g":94},"ef3c2dd0":{"m":76,"g":94},"75b65648":{"m":76,"g":94},"0f52fb55":{"m":76,"g":94},"d6d21640":{"m":76,"g":94},"0212d2e2":{"m":76,"g":94},"8cc300f5":{"m":76,"g":94},"452db508":{"m":76,"g":94},"d1112d85":{"m":76,"g":94},"48efec7b":{"m":76,"g":94},"9b8333d9":{"m":76,"g":94},"f5bbf603":{"m":76,"g":94},"5cbd709e":{"m":76,"g":94},"2e4a1e2d":{"m":76,"g":94},"9d02bb3e":{"m":76,"g":94},"402db5c5":{"m":76,"g":94},"754a0e82":{"m":76,"g":94},"799fb5f4":{"m":76,"g":94},"25e1816e":{"m":76,"g":94},"a53fe428":{"m":76,"g":94},"1b859295":{"m":76,"g":94},"9971dc22":{"m":76,"g":94},"3db35c1a":{"m":76,"g":94},"52a34d74":{"m":76,"g":94},"06d12b39":{"m":76,"g":94},"c30976fb":{"m":76,"g":94},"1a3fa75f":{"m":76,"g":94},"81f431ed":{"m":76,"g":94},"65b7c9b7":{"m":76,"g":94},"2c4f5cca":{"m":76,"g":94},"15843047":{"m":76,"g":94},"8ec2ce07":{"m":76,"g":94},"1fd0cf8a":{"m":76,"g":94},"bf63ee54":{"m":76,"g":94},"22c96f78":{"m":76,"g":94},"2892b9bb":{"m":76,"g":94},"470b4740":{"m":76,"g":94},"26c372c1":{"m":76,"g":94},"86d9baed":{"m":76,"g":94},"21d485f8":{"m":76,"g":94},"035ac2ab":{"m":76,"g":94},"e1a5e7e4":{"m":76,"g":94},"ad1ae7f7":{"m":76,"g":94},"e73167ad":{"m":76,"g":94},"862fe522":{"m":76,"g":94},"61e4433c":{"m":76,"g":94},"660305c3":{"m":76,"g":94},"642ab418":{"m":76,"g":94},"1ce4878d":{"m":76,"g":94},"977d7cd2":{"m":76,"g":94},"0e0ec702":{"m":76,"g":94},"bb378556":{"m":76,"g":94},"19e96e59":{"m":77,"g":94},"aa08aeac":{"m":77,"g":94},"d8a136a1":{"m":77,"g":94},"20c90be2":{"m":77,"g":94},"ec3ee028":{"m":77,"g":94},"92941ce7":{"m":77,"g":94},"2bb0e7cf":{"m":77,"g":94},"72549263":{"m":77,"g":94},"044c3159":{"m":77,"g":94},"4db29e82":{"m":77,"g":94},"c483377e":{"m":77,"g":94},"74e0ac1d":{"m":77,"g":94},"ef9a378a":{"m":77,"g":94},"6dea5c96":{"m":77,"g":94},"6ffb6bd4":{"m":77,"g":94},"47e6628a":{"m":77,"g":94},"7907f9eb":{"m":77,"g":94},"8c04f0f2":{"m":77,"g":94},"265e7564":{"m":77,"g":94},"d3f71f5e":{"m":77,"g":94},"5eae67cb":{"m":77,"g":94},"6dbf9998":{"m":77,"g":94},"e0166f8a":{"m":77,"g":94},"53a2c3b4":{"m":77,"g":94},"550586ef":{"m":77,"g":94},"cf29fe9e":{"m":77,"g":94},"26c0f131":{"m":77,"g":94},"f9970bd1":{"m":77,"g":94},"2e0f94ab":{"m":77,"g":94},"18317ddc":{"m":77,"g":94},"e2e2ab70":{"m":77,"g":94},"0d3e3072":{"m":77,"g":94},"62dd9587":{"m":77,"g":94},"72031173":{"m":77,"g":94},"9fdc6d6a":{"m":77,"g":94},"42a45df0":{"m":77,"g":94},"04eb6062":{"m":77,"g":94},"e84f4ba0":{"m":77,"g":94},"b149b393":{"m":77,"g":94},"31dfff7d":{"m":77,"g":94},"10a9ab7b":{"m":77,"g":94},"bb0fd749":{"m":77,"g":94},"7f19e083":{"m":77,"g":94},"98a2cfa9":{"m":77,"g":94},"2a882e8f":{"m":77,"g":94},"e6e4d022":{"m":77,"g":94},"188105a2":{"m":77,"g":94},"b3953258":{"m":77,"g":94},"5fa3058f":{"m":77,"g":94},"bbab97a6":{"m":77,"g":94},"0bc0bf57":{"m":77,"g":94},"f60f2931":{"m":77,"g":94},"17000d2b":{"m":77,"g":94},"668ecc6c":{"m":77,"g":94},"886fcbdd":{"m":77,"g":94},"8bf6d7f4":{"m":77,"g":94},"1b9175cb":{"m":77,"g":94},"92bb49a7":{"m":77,"g":94},"6f5cc5eb":{"m":77,"g":94},"c913ed40":{"m":77,"g":94},"1afe3d07":{"m":77,"g":94},"44f47d3e":{"m":77,"g":94},"ae25d36d":{"m":77,"g":94},"35e0856b":{"m":78,"g":94},"aba5ca15":{"m":78,"g":94},"496dde84":{"m":78,"g":94},"bcbbf519":{"m":78,"g":94},"0d99adb7":{"m":78,"g":94},"efbae697":{"m":78,"g":94},"ca8d02ab":{"m":78,"g":94},"3f287b85":{"m":78,"g":94},"7ed77d6b":{"m":78,"g":94},"4c54f442":{"m":78,"g":94},"924ca7c9":{"m":78,"g":94},"6ff9c6a5":{"m":78,"g":94},"77e929a1":{"m":78,"g":94},"febe21ce":{"m":78,"g":94},"a995a773":{"m":78,"g":94},"31035dda":{"m":78,"g":94},"913e38df":{"m":78,"g":94},"d95269f9":{"m":78,"g":94},"e53bf190":{"m":78,"g":94},"3289c120":{"m":78,"g":94},"69df9761":{"m":78,"g":94},"98f768d1":{"m":78,"g":94},"d7954b76":{"m":78,"g":94},"74885a84":{"m":78,"g":94},"b8b6008f":{"m":78,"g":94},"8e10fec9":{"m":78,"g":94},"e8999b13":{"m":78,"g":94},"772d2a19":{"m":78,"g":94},"9d0b36c4":{"m":78,"g":94},"7d8c0ce7":{"m":78,"g":94},"e41549c3":{"m":78,"g":94},"cccfc10e":{"m":78,"g":94},"a2aea59b":{"m":78,"g":94},"2c8fd993":{"m":78,"g":94},"31da75ab":{"m":78,"g":94},"e983e432":{"m":78,"g":94},"e9c6ce46":{"m":78,"g":94},"3fadc647":{"m":78,"g":94},"e119f042":{"m":78,"g":94},"9eb49e87":{"m":78,"g":94},"12047f5e":{"m":78,"g":94},"fda6bb78":{"m":78,"g":94},"23c764b1":{"m":78,"g":94},"87fafa01":{"m":78,"g":94},"1c63e797":{"m":78,"g":94},"ee47a6c1":{"m":78,"g":94},"6384d317":{"m":78,"g":94},"5cb552b1":{"m":78,"g":94},"c7457191":{"m":78,"g":94},"51ac297a":{"m":78,"g":94},"a169b9f8":{"m":78,"g":94},"4a63bc32":{"m":78,"g":94},"a303325f":{"m":78,"g":94},"42873eac":{"m":78,"g":94},"4814ecaf":{"m":78,"g":94},"e62d60fe":{"m":78,"g":94},"032f8faa":{"m":78,"g":94},"37c66ec8":{"m":78,"g":94},"9adf178c":{"m":78,"g":94},"f842853a":{"m":78,"g":94},"195a09f5":{"m":78,"g":94},"9fccda31":{"m":78,"g":94},"4ede6770":{"m":78,"g":94},"b26bc86b":{"m":78,"g":94},"5ec5eaf7":{"m":78,"g":94},"0d7fe866":{"m":78,"g":94},"54b9a2de":{"m":78,"g":94},"8e7b3154":{"m":78,"g":94},"45dcfc2e":{"m":78,"g":94},"ddf8981d":{"m":78,"g":94},"400ad660":{"m":78,"g":94},"05625b97":{"m":78,"g":94},"736502d4":{"m":78,"g":94},"8690c40b":{"m":78,"g":94},"b1cfb4e9":{"m":78,"g":94},"57f99608":{"m":79,"g":94},"81992474":{"m":79,"g":94},"f04c80dc":{"m":79,"g":94},"d1bb1711":{"m":79,"g":94},"5b5c7237":{"m":80,"g":94},"a42736bb":{"m":80,"g":94},"dd83e7e9":{"m":80,"g":94},"0769b14b":{"m":80,"g":94},"b64b88e7":{"m":80,"g":94},"bc24205b":{"m":80,"g":94},"3efc8e2d":{"m":80,"g":94},"27a009bb":{"m":80,"g":94},"8ec0bb7d":{"m":80,"g":94},"fa909dc3":{"m":80,"g":94},"e8f62b20":{"m":80,"g":94},"88defc4d":{"m":80,"g":94},"6f509d55":{"m":80,"g":94},"12ef7e3b":{"m":80,"g":94},"838fa0f2":{"m":80,"g":94},"f1b3b75f":{"m":80,"g":94},"33b16ad1":{"m":80,"g":94},"ffde65a0":{"m":80,"g":94},"471650de":{"m":80,"g":94},"d06a83fb":{"m":80,"g":94},"5d134401":{"m":80,"g":94},"f88f7e19":{"m":80,"g":94},"3dfc6023":{"m":80,"g":94},"15e91d72":{"m":80,"g":94},"8aab7fdb":{"m":80,"g":94},"e940dc4f":{"m":80,"g":94},"388e15c0":{"m":80,"g":94},"11421a3f":{"m":80,"g":94},"6c41fcf0":{"m":80,"g":94},"ee9d6ca6":{"m":80,"g":94},"2dd64894":{"m":80,"g":94},"61e7c4dd":{"m":80,"g":94},"dae79444":{"m":80,"g":94},"f6772f14":{"m":80,"g":94},"ac5b78ba":{"m":80,"g":94},"38076dea":{"m":80,"g":94},"5e0a9b09":{"m":80,"g":94},"bdde2375":{"m":80,"g":94},"e9fc2ac7":{"m":80,"g":94},"44afde82":{"m":80,"g":94},"072df753":{"m":80,"g":94},"defede50":{"m":80,"g":94},"fc728719":{"m":80,"g":94},"14e8bd88":{"m":80,"g":94},"adca585b":{"m":80,"g":94},"39d90449":{"m":80,"g":94},"39e41138":{"m":80,"g":94},"5fbafbb8":{"m":80,"g":94},"a9499885":{"m":80,"g":94},"f7655790":{"m":80,"g":94},"f58b929a":{"m":80,"g":94},"c1270aab":{"m":80,"g":94},"8311b07f":{"m":80,"g":94},"c1380257":{"m":80,"g":94},"b62e7e99":{"m":80,"g":94},"7d3b7c87":{"m":80,"g":94},"75015bb6":{"m":80,"g":94},"b371f7cd":{"m":80,"g":94},"812e82f3":{"m":80,"g":94},"4879e50c":{"m":80,"g":94},"bc92107b":{"m":80,"g":94},"3e4794aa":{"m":80,"g":94},"690ec205":{"m":80,"g":94},"2074a2e6":{"m":80,"g":94},"57de7c6b":{"m":80,"g":94},"115ae2e7":{"m":80,"g":94},"aea98512":{"m":80,"g":94},"e4155e96":{"m":80,"g":94},"1b1b47a9":{"m":80,"g":94},"3c9740d2":{"m":80,"g":94},"2eb55770":{"m":80,"g":94},"f65b8d5c":{"m":80,"g":94},"5ad05719":{"m":80,"g":94},"34ef6c81":{"m":80,"g":94},"61172091":{"m":80,"g":94},"4f288113":{"m":80,"g":94},"136b8e6a":{"m":80,"g":94},"034c5256":{"m":80,"g":94},"c1dd773c":{"m":80,"g":94},"6f859379":{"m":80,"g":94},"f774a0d2":{"m":80,"g":94},"60bcbf2a":{"m":80,"g":94},"a0a9f6d6":{"m":80,"g":94},"80aa8ca8":{"m":80,"g":94},"4aa6bab0":{"m":80,"g":94},"c35dcfdb":{"m":80,"g":94},"c163bf4f":{"m":80,"g":94},"55986343":{"m":80,"g":94},"b75275b6":{"m":80,"g":94},"7074e9ca":{"m":80,"g":94},"fc14cca0":{"m":80,"g":94},"e7beff8a":{"m":80,"g":94},"4d2e3051":{"m":80,"g":94},"e53a0b3d":{"m":80,"g":94},"038bc5d5":{"m":80,"g":94},"aee62d74":{"m":80,"g":94},"cd7e32e2":{"m":80,"g":94},"88799448":{"m":80,"g":94},"a879811c":{"m":80,"g":94},"a222945d":{"m":80,"g":94},"ed01b451":{"m":80,"g":94},"d050df36":{"m":80,"g":94},"76f44c2a":{"m":80,"g":94},"1078396f":{"m":80,"g":94},"7e4f72dd":{"m":80,"g":94},"4c31ae9f":{"m":80,"g":94},"f730362e":{"m":80,"g":94},"e3c4bd31":{"m":80,"g":94},"5db37c86":{"m":80,"g":94},"4cb53ecd":{"m":80,"g":94},"456b008b":{"m":80,"g":94},"ebf495f0":{"m":80,"g":94},"7f875f12":{"m":80,"g":94},"fbebcb7a":{"m":80,"g":94},"87eddedf":{"m":80,"g":94},"40652482":{"m":80,"g":94},"86a876d8":{"m":80,"g":94},"92823069":{"m":80,"g":94},"d2e507df":{"m":80,"g":94},"61970b08":{"m":80,"g":94},"76c48a09":{"m":80,"g":94},"90caf06c":{"m":80,"g":94},"6669d127":{"m":80,"g":94},"f2b70afd":{"m":80,"g":94},"bc3f6db2":{"m":80,"g":94},"aac531c5":{"m":80,"g":94},"39efad4f":{"m":80,"g":94},"466899e6":{"m":80,"g":94},"11d760d5":{"m":80,"g":94},"5039d547":{"m":80,"g":94},"d09a51f1":{"m":80,"g":94},"f8194b26":{"m":80,"g":94},"6d3b35fa":{"m":80,"g":94},"a73c4df4":{"m":80,"g":94},"89a55418":{"m":80,"g":94},"2695ab05":{"m":80,"g":94},"88d6fd9a":{"m":80,"g":94},"cc88d98a":{"m":80,"g":94},"3033c11a":{"m":80,"g":94},"fd5a55cf":{"m":80,"g":94},"804d9f2e":{"m":80,"g":94},"a7c3f74b":{"m":80,"g":94},"5a144a8a":{"m":80,"g":94},"27f8e6b9":{"m":80,"g":94},"afb752bc":{"m":80,"g":94},"9731eca7":{"m":80,"g":94},"7c5658c1":{"m":80,"g":94},"9798e72b":{"m":80,"g":94},"ade714a6":{"m":80,"g":94},"93470a14":{"m":80,"g":94},"db452760":{"m":80,"g":94},"fbdc94ba":{"m":81,"g":94},"b54b5a96":{"m":81,"g":94},"bca832c7":{"m":81,"g":94},"d9dd5298":{"m":81,"g":94},"0a0dd34e":{"m":81,"g":94},"80ac527d":{"m":81,"g":94},"99456bca":{"m":81,"g":94},"d07e797a":{"m":81,"g":94},"c555d794":{"m":81,"g":94},"e2574ee9":{"m":81,"g":94},"ab4b5606":{"m":81,"g":94},"20f1c8e3":{"m":81,"g":94},"613b197e":{"m":81,"g":94},"d58e3544":{"m":81,"g":94},"bf86c5e9":{"m":81,"g":94},"dca90f1d":{"m":81,"g":94},"0961feef":{"m":81,"g":94},"59dd090f":{"m":81,"g":94},"569b032c":{"m":81,"g":94},"f6a71139":{"m":81,"g":94},"1e0806f3":{"m":81,"g":94},"2c11f9c2":{"m":81,"g":94},"a6f892e5":{"m":81,"g":94},"08b518d5":{"m":81,"g":94},"4db463b1":{"m":81,"g":94},"bfa39224":{"m":81,"g":94},"e465b08d":{"m":81,"g":94},"bed05878":{"m":81,"g":94},"b2a189dd":{"m":81,"g":94},"f28d8299":{"m":81,"g":94},"8e09b370":{"m":81,"g":94},"53dcf388":{"m":81,"g":94},"1effba4c":{"m":81,"g":94},"a0fc5bc1":{"m":81,"g":94},"27e9538a":{"m":81,"g":94},"211c7b31":{"m":81,"g":94},"c08a717c":{"m":81,"g":94},"f13d65a7":{"m":81,"g":94},"06d0a3d9":{"m":81,"g":94},"22c2a79d":{"m":81,"g":94},"8beb356f":{"m":81,"g":94},"c776234b":{"m":81,"g":94},"3bface15":{"m":81,"g":94},"6fb29ffd":{"m":81,"g":94},"4fb05583":{"m":81,"g":94},"81c89111":{"m":81,"g":94},"92d1561b":{"m":81,"g":94},"8f783c19":{"m":81,"g":94},"90faf901":{"m":81,"g":94},"177320a5":{"m":81,"g":94},"d7bc19a4":{"m":81,"g":94},"85ec0440":{"m":81,"g":94},"06a1656e":{"m":81,"g":94},"6aca5834":{"m":81,"g":94},"b9c87e78":{"m":82,"g":94},"968ef515":{"m":82,"g":94},"13432002":{"m":82,"g":94},"c2942907":{"m":82,"g":94},"e69a2190":{"m":82,"g":94},"bf98d2e3":{"m":82,"g":94},"e65b9f21":{"m":82,"g":94},"4dce1cc6":{"m":82,"g":94},"deded17f":{"m":82,"g":94},"f29a718f":{"m":82,"g":94},"3f57b00a":{"m":82,"g":94},"453d412c":{"m":82,"g":94},"dc86f25a":{"m":82,"g":94},"08289eaa":{"m":82,"g":94},"3b6d539f":{"m":82,"g":94},"57131dd9":{"m":82,"g":94},"a7591ecf":{"m":82,"g":94},"c44f2869":{"m":82,"g":94},"685d8980":{"m":82,"g":94},"70645f4d":{"m":82,"g":94},"188f0955":{"m":82,"g":94},"eef9433b":{"m":82,"g":94},"97cb762b":{"m":82,"g":94},"11951820":{"m":82,"g":94},"5239d795":{"m":82,"g":94},"f0815419":{"m":82,"g":94},"2b3bdc93":{"m":82,"g":94},"5fc4b600":{"m":82,"g":94},"b868526d":{"m":82,"g":94},"502524e2":{"m":82,"g":94},"4c764007":{"m":82,"g":94},"9f3bd2ad":{"m":82,"g":94},"8de53da9":{"m":82,"g":94},"fac17acf":{"m":82,"g":94},"8b39274e":{"m":82,"g":94},"5156d5a4":{"m":82,"g":94},"c951d312":{"m":82,"g":94},"dcb82325":{"m":82,"g":94},"66c0ff9e":{"m":82,"g":94},"9a7e83e8":{"m":82,"g":94},"417b44eb":{"m":82,"g":94},"475e2e37":{"m":82,"g":94},"fba86b6b":{"m":82,"g":94},"072b4d03":{"m":82,"g":94},"9c434777":{"m":82,"g":94},"fa2f677e":{"m":82,"g":94},"463d4b74":{"m":82,"g":94},"9924bbe1":{"m":82,"g":94},"84022c0e":{"m":83,"g":94},"f9fb33ef":{"m":83,"g":94},"a38f6932":{"m":83,"g":94},"beb65c74":{"m":83,"g":94},"621e96bf":{"m":83,"g":94},"35ca04d2":{"m":83,"g":94},"3c4e0ee6":{"m":83,"g":94},"9c088829":{"m":83,"g":94},"005aad32":{"m":83,"g":94},"4d23ba08":{"m":83,"g":94},"6e313c1b":{"m":83,"g":94},"a45a4b23":{"m":83,"g":94},"981a2619":{"m":83,"g":94},"8ba31330":{"m":83,"g":94},"02102063":{"m":83,"g":94},"7e944246":{"m":83,"g":94},"a086a113":{"m":83,"g":94},"bdbe5f81":{"m":83,"g":94},"9ad28f63":{"m":83,"g":94},"d7b1ce65":{"m":83,"g":94},"f55933e1":{"m":83,"g":94},"408ba022":{"m":83,"g":94},"094891c0":{"m":83,"g":94},"a21ef363":{"m":83,"g":94},"3c4dc38a":{"m":83,"g":94},"d8fbc7c0":{"m":83,"g":94},"c5e1026f":{"m":83,"g":94},"799c4bb5":{"m":83,"g":94},"02723e1b":{"m":83,"g":94},"df2cf583":{"m":83,"g":94},"133ded03":{"m":83,"g":94},"f87a6ab3":{"m":83,"g":94},"eebfdb94":{"m":83,"g":94},"dfb32264":{"m":83,"g":94},"63c13a2c":{"m":83,"g":94},"4d1e52ab":{"m":83,"g":94},"155890e4":{"m":83,"g":94},"1f963d7f":{"m":83,"g":94},"04d0123f":{"m":83,"g":94},"feda9b11":{"m":83,"g":94},"c3948ba6":{"m":83,"g":94},"269c457e":{"m":83,"g":94},"18ce468d":{"m":83,"g":94},"21514ff5":{"m":83,"g":94},"5641a094":{"m":83,"g":94},"3dd3538c":{"m":83,"g":94},"93c6fb12":{"m":83,"g":94},"11e27d09":{"m":83,"g":94},"50eda839":{"m":83,"g":94},"c55550cb":{"m":83,"g":94},"43fb95c2":{"m":83,"g":94},"7d9679b7":{"m":83,"g":94},"b5be5694":{"m":83,"g":94},"d2b8d0b8":{"m":83,"g":94},"a14654dd":{"m":83,"g":94},"5d93a950":{"m":83,"g":94},"c998d04b":{"m":83,"g":94},"7d0edf3c":{"m":83,"g":94},"ce4ecba4":{"m":83,"g":94},"b1f6d89b":{"m":83,"g":94},"7c99103f":{"m":83,"g":94},"de071366":{"m":83,"g":94},"e0673969":{"m":83,"g":94},"127ff898":{"m":83,"g":94},"8777a1d2":{"m":83,"g":94},"711efe78":{"m":83,"g":94},"fbb5f229":{"m":83,"g":94},"15fabcc0":{"m":83,"g":94},"e62c4955":{"m":83,"g":94},"71d1785f":{"m":83,"g":94},"3f87f831":{"m":83,"g":94},"ce5412b6":{"m":83,"g":94},"7282ab74":{"m":83,"g":94},"b0feda09":{"m":83,"g":94},"6b6e7487":{"m":83,"g":94},"91732486":{"m":83,"g":94},"2ed96c7a":{"m":83,"g":94},"2aa3f5e2":{"m":83,"g":94},"76d17c7e":{"m":83,"g":94},"70d040f9":{"m":83,"g":94},"4418f599":{"m":83,"g":94},"04f2abcb":{"m":83,"g":94},"506be6b8":{"m":83,"g":94},"2343d8df":{"m":83,"g":94},"92bb64bc":{"m":83,"g":94},"11b23ae9":{"m":83,"g":94},"dcae1fb2":{"m":84,"g":94},"a0251a3f":{"m":84,"g":94},"663037a7":{"m":84,"g":94},"f4a9f60c":{"m":84,"g":94},"ee71ed8a":{"m":84,"g":94},"d364b9b0":{"m":84,"g":94},"849c83a0":{"m":84,"g":94},"d73ddeb1":{"m":84,"g":94},"f48b007c":{"m":84,"g":94},"74cb12a8":{"m":84,"g":94},"c6c62640":{"m":84,"g":94},"92ab0a20":{"m":84,"g":94},"e132cba2":{"m":84,"g":94},"0045f4b2":{"m":84,"g":94},"8601300b":{"m":84,"g":94},"6fa6f38e":{"m":84,"g":94},"693723d1":{"m":84,"g":94},"966eb908":{"m":84,"g":94},"644ed409":{"m":84,"g":94},"3029889c":{"m":84,"g":94},"ef15dcda":{"m":84,"g":94},"ad4df307":{"m":84,"g":94},"41ac0c6d":{"m":84,"g":94},"84810da4":{"m":84,"g":94},"40d9b8ac":{"m":84,"g":94},"f0365820":{"m":84,"g":94},"86317c09":{"m":84,"g":94},"daed453e":{"m":84,"g":94},"ded04b2e":{"m":84,"g":94},"9858113c":{"m":85,"g":94},"8441baad":{"m":85,"g":94},"256c4c25":{"m":85,"g":94},"9f21e754":{"m":85,"g":94},"7bcd8b1c":{"m":85,"g":94},"11383cec":{"m":85,"g":94},"e97e57e6":{"m":85,"g":94},"9a6ad891":{"m":85,"g":94},"d353d08b":{"m":85,"g":94},"08acdb5c":{"m":85,"g":94},"2afba1b1":{"m":85,"g":94},"e330f2b8":{"m":85,"g":94},"3ddf5b9d":{"m":85,"g":94},"3cff9633":{"m":85,"g":94},"d50e36a7":{"m":85,"g":94},"8fefdd32":{"m":85,"g":94},"403b855a":{"m":85,"g":94},"1698e94e":{"m":85,"g":94},"58195dd5":{"m":85,"g":94},"799789af":{"m":85,"g":94},"cc4a80ca":{"m":85,"g":94},"3c8a5231":{"m":85,"g":94},"a043f7f2":{"m":85,"g":94},"e3a53044":{"m":85,"g":94},"28b26dbf":{"m":85,"g":94},"2b06484b":{"m":85,"g":94},"e4b6133b":{"m":85,"g":94},"dd408ee4":{"m":85,"g":94},"9419e75d":{"m":85,"g":94},"2c7dbb7c":{"m":85,"g":94},"9a62191b":{"m":85,"g":94},"ae523675":{"m":85,"g":94},"5c08aa49":{"m":85,"g":94},"f4c191a7":{"m":85,"g":94},"771669cb":{"m":85,"g":94},"1468769b":{"m":85,"g":94},"91dda4cd":{"m":85,"g":94},"8e5a6d34":{"m":85,"g":94},"8465f035":{"m":85,"g":94},"8c0cfca8":{"m":85,"g":94},"2c3ea294":{"m":85,"g":94},"5bb0accb":{"m":85,"g":94},"8d463fe3":{"m":85,"g":94},"26fc32d1":{"m":85,"g":94},"1cc32603":{"m":85,"g":94},"05ee2192":{"m":85,"g":94},"678d8cc9":{"m":86,"g":94},"d2cb3024":{"m":86,"g":94},"1940cdec":{"m":86,"g":94},"63484f9f":{"m":86,"g":94},"dff0ab92":{"m":86,"g":94},"e30c273b":{"m":86,"g":94},"0ab3f437":{"m":86,"g":94},"cec98f10":{"m":86,"g":94},"8dc4efd0":{"m":86,"g":94},"6578cf27":{"m":86,"g":94},"087751a8":{"m":86,"g":94},"911f3ba6":{"m":86,"g":94},"f6f96b05":{"m":86,"g":94},"2a936a84":{"m":86,"g":94},"5e023301":{"m":86,"g":94},"fa7d7fd9":{"m":86,"g":94},"f1ff736d":{"m":86,"g":94},"acc816d8":{"m":86,"g":94},"a05bd83a":{"m":86,"g":94},"cef91b1e":{"m":86,"g":94},"6450c122":{"m":86,"g":94},"b6cf3532":{"m":86,"g":94},"3b2680a4":{"m":86,"g":94},"79961afa":{"m":86,"g":94},"cfca4e0e":{"m":86,"g":94},"e88dd482":{"m":86,"g":94},"73600673":{"m":86,"g":94},"8f508cc7":{"m":86,"g":94},"9bddf1c8":{"m":86,"g":94},"24c13ca9":{"m":86,"g":94},"b70957fc":{"m":86,"g":94},"e444c13f":{"m":86,"g":94},"fee37d9e":{"m":86,"g":94},"c68de479":{"m":86,"g":94},"4c7b4242":{"m":86,"g":94},"38053c33":{"m":86,"g":94},"00c2c1f0":{"m":86,"g":94},"cb691945":{"m":86,"g":94},"d25398cb":{"m":86,"g":94},"8a828666":{"m":86,"g":94},"aff584fa":{"m":86,"g":94},"6f566147":{"m":86,"g":94},"bdd17998":{"m":86,"g":94},"c9abd7be":{"m":86,"g":94},"a3e4e9bf":{"m":86,"g":94},"6d4d3bc8":{"m":86,"g":94},"5f300141":{"m":86,"g":94},"1c05425b":{"m":86,"g":94},"b26cb1c5":{"m":86,"g":94},"f8e46093":{"m":86,"g":94},"683707c3":{"m":86,"g":94},"a68ed766":{"m":86,"g":94},"82653f66":{"m":86,"g":94},"22da3d97":{"m":86,"g":94},"b8559764":{"m":86,"g":94},"56f6589e":{"m":86,"g":94},"1232f7e8":{"m":86,"g":94},"3008db9c":{"m":86,"g":94},"357fb2db":{"m":86,"g":94},"95c231e5":{"m":86,"g":94},"3042f1da":{"m":86,"g":94},"2b63798c":{"m":86,"g":94},"bf203cb7":{"m":86,"g":94},"8ebde73f":{"m":86,"g":94},"6b0fae79":{"m":86,"g":94},"141a4596":{"m":86,"g":94},"d8ab6011":{"m":86,"g":94},"6579cd7d":{"m":86,"g":94},"97ac42b6":{"m":86,"g":94},"1acca3a2":{"m":86,"g":94},"6ea1e6ac":{"m":86,"g":94},"3409aaab":{"m":86,"g":94},"73dcf2b3":{"m":86,"g":94},"170d1f21":{"m":86,"g":94},"73bc1d00":{"m":86,"g":94},"c5645e92":{"m":86,"g":94},"d33955d2":{"m":86,"g":94},"6fc17596":{"m":86,"g":94},"ad506a4e":{"m":86,"g":94},"ebaba856":{"m":86,"g":94},"de2faef9":{"m":86,"g":94},"67b7d5b1":{"m":86,"g":94},"4322c31e":{"m":86,"g":94},"16267d4f":{"m":87,"g":94},"0f5cb8ca":{"m":87,"g":94},"17299f08":{"m":87,"g":94},"5380cd7e":{"m":87,"g":94},"b2e95f62":{"m":87,"g":94},"1ab14c4c":{"m":87,"g":94},"3c32895c":{"m":87,"g":94},"ac2324c1":{"m":87,"g":94},"ef8ec07b":{"m":87,"g":94},"f24fc5b8":{"m":87,"g":94},"d18c6b33":{"m":87,"g":94},"f1c89600":{"m":87,"g":94},"983c663d":{"m":87,"g":94},"f94543d2":{"m":87,"g":94},"e8e18dcd":{"m":87,"g":94},"bad7c26f":{"m":87,"g":94},"12319a67":{"m":87,"g":94},"d738ab52":{"m":87,"g":94},"3ee40ff9":{"m":87,"g":94},"0f334945":{"m":87,"g":94},"fba8eccd":{"m":87,"g":94},"7d3a3d45":{"m":87,"g":94},"25c83fff":{"m":87,"g":94},"9f2c9568":{"m":87,"g":94},"3f2702ae":{"m":87,"g":94},"6ea05950":{"m":87,"g":94},"e7dd906c":{"m":87,"g":94},"6e2da515":{"m":87,"g":94},"e9a47f4c":{"m":87,"g":94},"03227c5f":{"m":87,"g":94},"01bdbf7f":{"m":87,"g":94},"94d42b67":{"m":87,"g":94},"69276f61":{"m":87,"g":94},"41a645f5":{"m":87,"g":94},"23010630":{"m":87,"g":94},"45b4dcf0":{"m":87,"g":94},"213e8c7d":{"m":87,"g":94},"41273fd7":{"m":87,"g":94},"e9bebafb":{"m":87,"g":94},"4d1c9db6":{"m":87,"g":94},"17c36c55":{"m":87,"g":94},"31d1f6e7":{"m":87,"g":94},"a823c6e8":{"m":87,"g":94},"2ce87935":{"m":87,"g":94},"de167cf5":{"m":87,"g":94},"4319978c":{"m":87,"g":94},"03dd785c":{"m":87,"g":94},"66fc63d6":{"m":87,"g":94},"921e4a81":{"m":87,"g":94},"9d8ec2e6":{"m":87,"g":94},"c178abda":{"m":87,"g":94},"b29a026e":{"m":87,"g":94},"7e257cd6":{"m":88,"g":94},"c4831e2f":{"m":88,"g":94},"2e37fa07":{"m":88,"g":94},"2d831c6e":{"m":88,"g":94},"ed0c3035":{"m":88,"g":94},"e6f11356":{"m":88,"g":94},"7b02c326":{"m":88,"g":94},"fefa19fe":{"m":88,"g":94},"9c574585":{"m":88,"g":94},"8233cc10":{"m":88,"g":94},"1b2e8f76":{"m":88,"g":94},"d2e0881a":{"m":88,"g":94},"2f427491":{"m":88,"g":94},"d8189660":{"m":88,"g":94},"3ded6235":{"m":88,"g":94},"4ba1eea8":{"m":88,"g":94},"4685fbb8":{"m":88,"g":94},"0a4fc73b":{"m":88,"g":94},"a6970a17":{"m":88,"g":94},"a6ae3af1":{"m":88,"g":94},"0b07c4a9":{"m":88,"g":94},"fc0e3b91":{"m":88,"g":94},"d71f3f0a":{"m":88,"g":94},"58f10679":{"m":88,"g":94},"7a80f565":{"m":88,"g":94},"9484eba4":{"m":88,"g":94},"e9feb488":{"m":88,"g":94},"fc992a09":{"m":88,"g":94},"121f92c5":{"m":88,"g":94},"3bde1010":{"m":88,"g":94},"75135580":{"m":88,"g":94},"4d643f6c":{"m":88,"g":94},"6ce0ed07":{"m":88,"g":94},"969660c7":{"m":88,"g":94},"16d4f680":{"m":88,"g":94},"ada268fd":{"m":88,"g":94},"cfe48c59":{"m":88,"g":94},"d4c038da":{"m":88,"g":94},"55f6005f":{"m":88,"g":94},"7222e1da":{"m":88,"g":94},"505eec4d":{"m":88,"g":94},"ccfe5c00":{"m":88,"g":94},"a071dc40":{"m":88,"g":94},"a40aecc5":{"m":88,"g":94},"d6e1d28c":{"m":88,"g":94},"7c347259":{"m":88,"g":94},"669caa0a":{"m":88,"g":94},"4024e1d2":{"m":88,"g":94},"5c0b38f3":{"m":88,"g":94},"30ca18f4":{"m":88,"g":94},"03886917":{"m":88,"g":94},"66324895":{"m":88,"g":94},"13feffd0":{"m":88,"g":94},"e98afbe0":{"m":88,"g":94},"69af3ec3":{"m":88,"g":94},"32cc66ef":{"m":88,"g":94},"83f2d9d4":{"m":88,"g":94},"6317c5c6":{"m":88,"g":94},"cba1cdbc":{"m":88,"g":94},"c471d39e":{"m":88,"g":94},"d0443275":{"m":88,"g":94},"17d080b7":{"m":88,"g":94},"1b19df4b":{"m":88,"g":94},"f0653886":{"m":88,"g":94},"b1465557":{"m":88,"g":94},"b06215da":{"m":88,"g":94},"7adf245b":{"m":88,"g":94},"299fd22f":{"m":88,"g":94},"506e5de8":{"m":88,"g":94},"844e2f22":{"m":88,"g":94},"4f39bcf7":{"m":88,"g":94},"31c9569b":{"m":88,"g":94},"1be6956d":{"m":88,"g":94},"626ccb7d":{"m":88,"g":94},"72bfb0ba":{"m":88,"g":94},"15521495":{"m":88,"g":94},"ebe58d54":{"m":88,"g":94},"066cf445":{"m":88,"g":94},"6dc6b306":{"m":88,"g":94},"1f30c05d":{"m":88,"g":94},"5dd62c3a":{"m":88,"g":94},"f11481b9":{"m":88,"g":94},"9d24c3ff":{"m":88,"g":94},"24161c59":{"m":88,"g":94},"eabcf82a":{"m":88,"g":94},"c47a51db":{"m":88,"g":94},"11553c1a":{"m":88,"g":94},"01dd39ba":{"m":88,"g":94},"b3f3d610":{"m":88,"g":94},"f07c6a00":{"m":88,"g":94},"4bb816d4":{"m":88,"g":94},"c250939e":{"m":88,"g":94},"b6909aa2":{"m":88,"g":94},"f8728357":{"m":88,"g":94},"73187152":{"m":88,"g":94},"40865665":{"m":88,"g":94},"fd08c048":{"m":88,"g":94},"26ebb849":{"m":88,"g":94},"02973cd9":{"m":88,"g":94},"6d95a35a":{"m":88,"g":94},"01d2838c":{"m":88,"g":94},"e3b8a722":{"m":88,"g":94},"3cf1473a":{"m":88,"g":94},"27168308":{"m":88,"g":94},"e3bed74a":{"m":88,"g":94},"e9ef39d2":{"m":88,"g":94},"205d5cb4":{"m":88,"g":94},"3d7f7a43":{"m":88,"g":94},"2df9d40a":{"m":88,"g":94},"8dc191f2":{"m":88,"g":94},"64825b83":{"m":88,"g":94},"69748d08":{"m":88,"g":94},"dcc0a456":{"m":88,"g":94},"c2b7ddca":{"m":88,"g":94},"abebd939":{"m":88,"g":94},"4bd2952a":{"m":88,"g":94},"6fc93575":{"m":88,"g":94},"839fb31e":{"m":88,"g":94},"f19a9204":{"m":88,"g":94},"c23a7072":{"m":88,"g":94},"e07a6977":{"m":88,"g":94},"cd8d4b9d":{"m":88,"g":94},"f194e14f":{"m":88,"g":94},"cfc9f9ab":{"m":88,"g":94},"fb4959b2":{"m":88,"g":94},"9a405274":{"m":88,"g":94},"2e4babdb":{"m":88,"g":94},"44a3783d":{"m":88,"g":94},"f3bf6110":{"m":88,"g":94},"198b9056":{"m":88,"g":94},"73eb67c0":{"m":88,"g":94},"9a91fa0e":{"m":88,"g":94},"cd7c8a8d":{"m":88,"g":94},"3e350a93":{"m":88,"g":94},"fb71725c":{"m":88,"g":94},"912788c0":{"m":88,"g":94},"0f75b907":{"m":88,"g":94},"4f723edd":{"m":89,"g":94},"fcde67b0":{"m":89,"g":94},"81372f3b":{"m":89,"g":94},"baa6624d":{"m":89,"g":94},"c2b16795":{"m":89,"g":94},"f6ebba53":{"m":89,"g":94},"6716b417":{"m":89,"g":94},"1c8b42c8":{"m":89,"g":94},"f20f7000":{"m":89,"g":94},"f40942ad":{"m":89,"g":94},"dc0705a5":{"m":89,"g":94},"a968c888":{"m":89,"g":94},"a979daac":{"m":89,"g":94},"f1569876":{"m":89,"g":94},"3465d7ae":{"m":89,"g":94},"e58423b2":{"m":89,"g":94},"7059ae16":{"m":89,"g":94},"51d9a597":{"m":89,"g":94},"56ccd3c2":{"m":89,"g":94},"98c00a2d":{"m":89,"g":94},"451ffe74":{"m":89,"g":94},"b1e5a33a":{"m":89,"g":94},"9d5fa68b":{"m":89,"g":94},"2c186425":{"m":89,"g":94},"18efb5e8":{"m":89,"g":94},"de1350ea":{"m":89,"g":94},"86fe943b":{"m":89,"g":94},"9ecb1856":{"m":89,"g":94},"cc74499d":{"m":89,"g":94},"0c1f03a2":{"m":89,"g":94},"3712abfa":{"m":89,"g":94},"971a0dfa":{"m":89,"g":94},"2fc12995":{"m":89,"g":94},"20d3ad3b":{"m":89,"g":94},"fa3592cf":{"m":89,"g":94},"608668e1":{"m":89,"g":94},"6c0a4828":{"m":89,"g":94},"47402883":{"m":89,"g":94},"1fb76ebb":{"m":89,"g":94},"c2c4f57f":{"m":89,"g":94},"23881fa6":{"m":89,"g":94},"8db3ac55":{"m":89,"g":94},"3e56f557":{"m":89,"g":94},"62fec60d":{"m":89,"g":94},"e7759778":{"m":89,"g":94},"77e928d0":{"m":89,"g":94},"515ef4fa":{"m":89,"g":94},"f5599ef1":{"m":89,"g":94},"c499591a":{"m":89,"g":94},"e1ce44cd":{"m":89,"g":94},"f1114e7f":{"m":89,"g":94},"bae4fdc7":{"m":89,"g":94},"6153f2ff":{"m":89,"g":94},"8b5f83ed":{"m":89,"g":94},"2a413829":{"m":89,"g":94},"d5c097a2":{"m":89,"g":94},"9736cd3b":{"m":89,"g":94},"2f715f51":{"m":89,"g":94},"d664ca18":{"m":89,"g":94},"22fe7878":{"m":89,"g":94},"c4ffbeca":{"m":89,"g":94},"f8eaaab8":{"m":89,"g":94},"697b0f71":{"m":89,"g":94},"132dad87":{"m":89,"g":94},"60fdad7c":{"m":89,"g":94},"61ce91ed":{"m":89,"g":94},"e6b7053b":{"m":89,"g":94},"5f91c825":{"m":89,"g":94},"b819381f":{"m":89,"g":94},"562f279a":{"m":89,"g":94},"8b247489":{"m":89,"g":94},"0df6765c":{"m":89,"g":94},"35b65cf0":{"m":89,"g":94},"dd1012fc":{"m":89,"g":94},"44aab7f9":{"m":89,"g":94},"43baba64":{"m":89,"g":94},"0166403c":{"m":89,"g":94},"bcf66ef3":{"m":89,"g":94},"0de5e7d4":{"m":89,"g":94},"72a110f6":{"m":89,"g":94},"5aff1e93":{"m":89,"g":94},"8e3797be":{"m":89,"g":94},"4474eaf5":{"m":89,"g":94},"499f5e62":{"m":89,"g":94},"81964328":{"m":89,"g":94},"f0f84975":{"m":89,"g":94},"3f1e4339":{"m":89,"g":94},"cf9815ba":{"m":89,"g":94},"bd75690f":{"m":89,"g":94},"180ff5ee":{"m":89,"g":94},"37f15475":{"m":89,"g":94},"8a548052":{"m":89,"g":94},"b6d0ce9f":{"m":89,"g":94},"0ea330ca":{"m":89,"g":94},"27e327b4":{"m":89,"g":94},"ff00895c":{"m":89,"g":94},"ff914748":{"m":89,"g":94},"eb38c7d1":{"m":89,"g":94},"df7f61ee":{"m":89,"g":94},"ef21729c":{"m":89,"g":94},"f5159315":{"m":89,"g":94},"6d7b6696":{"m":89,"g":94},"6376b632":{"m":89,"g":94},"e05e29d1":{"m":89,"g":94},"a2cb5913":{"m":89,"g":94},"55444ed6":{"m":89,"g":94},"20fd53b8":{"m":89,"g":94},"6a47b730":{"m":89,"g":94},"c429919d":{"m":89,"g":94},"1da8d230":{"m":89,"g":94},"2f7420bc":{"m":89,"g":94},"c6a0cacc":{"m":89,"g":94},"0a9bfc20":{"m":89,"g":94},"34c63731":{"m":89,"g":94},"2d72fc47":{"m":89,"g":94},"b520d028":{"m":89,"g":94},"7dc0e394":{"m":89,"g":94},"fb507b7b":{"m":89,"g":94},"f90945c4":{"m":89,"g":94},"094fbdac":{"m":89,"g":94},"888cb175":{"m":89,"g":94},"e39bca07":{"m":89,"g":94},"a2bb8565":{"m":89,"g":94},"ced3c07a":{"m":89,"g":94},"f18b068f":{"m":89,"g":94},"4fac524b":{"m":89,"g":94},"b581b225":{"m":89,"g":94},"69dd878b":{"m":89,"g":94},"22630ca2":{"m":89,"g":94},"d279d499":{"m":89,"g":94},"6cb00c63":{"m":89,"g":94},"62cac2c4":{"m":89,"g":94},"2c3b71d6":{"m":89,"g":94},"51cdd81f":{"m":89,"g":94},"73def253":{"m":89,"g":94},"d9d35def":{"m":89,"g":94},"6df81e8a":{"m":89,"g":94},"3ab7d9b5":{"m":89,"g":94},"7e5071c9":{"m":89,"g":94},"78689d33":{"m":89,"g":94},"1dc6864f":{"m":89,"g":94},"485a023b":{"m":89,"g":94},"7e412900":{"m":89,"g":94},"c673727e":{"m":89,"g":94},"f4d4f939":{"m":89,"g":94},"f2bd3515":{"m":89,"g":94},"c459536b":{"m":89,"g":94},"535c8386":{"m":89,"g":94},"2163586e":{"m":89,"g":94},"e06b0761":{"m":89,"g":94},"844a8f42":{"m":89,"g":94},"791b3bfa":{"m":89,"g":94},"31589e17":{"m":89,"g":94},"ae6a5b29":{"m":89,"g":94},"4839999b":{"m":89,"g":94},"541a985f":{"m":89,"g":94},"5170b010":{"m":89,"g":94},"d63e76f7":{"m":89,"g":94},"e9fd11c0":{"m":89,"g":94},"c7588d59":{"m":89,"g":94},"6b231325":{"m":89,"g":94},"b1c8d4e9":{"m":89,"g":94},"c25231c6":{"m":89,"g":94},"fba03b29":{"m":89,"g":94},"461a7302":{"m":89,"g":94},"07610353":{"m":89,"g":94},"c087ddd6":{"m":89,"g":94},"f4a8987f":{"m":89,"g":94},"41ba767f":{"m":89,"g":94},"f127355a":{"m":89,"g":94},"bdb962d7":{"m":89,"g":94},"0b9557fc":{"m":89,"g":94},"87068b5c":{"m":89,"g":94},"a564e001":{"m":89,"g":94},"2103b806":{"m":89,"g":94},"e806f708":{"m":89,"g":94},"fa6723f0":{"m":89,"g":94},"673ff668":{"m":89,"g":94},"447be242":{"m":89,"g":94},"183d9f96":{"m":89,"g":94},"63195028":{"m":89,"g":94},"a3d7f4b6":{"m":89,"g":94},"b18416fb":{"m":89,"g":94},"ce9d690e":{"m":89,"g":94},"bdaefbbf":{"m":89,"g":94},"45a31a82":{"m":89,"g":94},"1aa0fbf4":{"m":89,"g":94},"7a0bbe6a":{"m":89,"g":94},"ae335842":{"m":89,"g":94},"477a101c":{"m":89,"g":94},"1a8f5f68":{"m":89,"g":94},"32cd7070":{"m":89,"g":94},"ebd1ed49":{"m":89,"g":94},"f77da699":{"m":89,"g":94},"d6864ce6":{"m":89,"g":94},"755a3661":{"m":89,"g":94},"79a39ac0":{"m":89,"g":94},"3ce94f71":{"m":89,"g":94},"ca95556c":{"m":89,"g":94},"eb8f02dd":{"m":89,"g":94},"0ca3e568":{"m":89,"g":94},"5c7aa009":{"m":89,"g":94},"fe386aca":{"m":89,"g":94},"14d1075f":{"m":89,"g":94},"006ead9d":{"m":89,"g":94},"0d503090":{"m":89,"g":94},"501efc3d":{"m":89,"g":94},"f9bab3d5":{"m":89,"g":94},"16f69b1f":{"m":89,"g":94},"65f09131":{"m":89,"g":94},"fc419b62":{"m":89,"g":94},"7eb9d8e5":{"m":89,"g":94},"84147254":{"m":89,"g":94},"6bebef60":{"m":89,"g":94},"25be63d0":{"m":89,"g":94},"d502dae0":{"m":89,"g":94},"93e53f6e":{"m":89,"g":94},"a191a0e4":{"m":89,"g":94},"8c7279c2":{"m":89,"g":94},"0ca18117":{"m":89,"g":94},"2c3a6fe1":{"m":89,"g":94},"8b33d8df":{"m":89,"g":94},"e235be16":{"m":89,"g":94},"5ccf8fe1":{"m":89,"g":94},"3f23d8cd":{"m":89,"g":94},"1a399799":{"m":89,"g":94},"022012aa":{"m":89,"g":94},"681e7af3":{"m":89,"g":94},"681fdc26":{"m":89,"g":94},"0d477880":{"m":89,"g":94},"f4560373":{"m":89,"g":94},"b2388433":{"m":89,"g":94},"a38376fa":{"m":89,"g":94},"7a5e6ce1":{"m":89,"g":94},"24c035f2":{"m":89,"g":94},"f9dc9dd2":{"m":90,"g":94},"62a7aa2e":{"m":90,"g":94},"5ca07eed":{"m":90,"g":94},"e30ef368":{"m":90,"g":94},"91a066ec":{"m":90,"g":94},"c4943867":{"m":90,"g":94},"53a525bf":{"m":90,"g":94},"7ddf8e83":{"m":90,"g":94},"8321f8e4":{"m":90,"g":94},"cfceb83d":{"m":90,"g":94},"b1286a11":{"m":90,"g":94},"21615cc3":{"m":90,"g":94},"0ae1e9a7":{"m":90,"g":94},"e07d0647":{"m":90,"g":94},"3c2274fb":{"m":90,"g":94},"d2679f51":{"m":90,"g":94},"96be97bf":{"m":90,"g":94},"88f9c347":{"m":90,"g":94},"fff10809":{"m":90,"g":94},"5f1ab327":{"m":90,"g":94},"7df7c679":{"m":90,"g":94},"38af4f68":{"m":90,"g":94},"a6305c7d":{"m":90,"g":94},"a023856b":{"m":90,"g":94},"db0cc57e":{"m":90,"g":94},"349bb2c9":{"m":90,"g":94},"0b8939bc":{"m":90,"g":94},"ed89837c":{"m":90,"g":94},"55561e25":{"m":90,"g":94},"44733203":{"m":90,"g":94},"0bd67ba2":{"m":90,"g":94},"7d316991":{"m":90,"g":94},"ab1a4fa5":{"m":90,"g":94},"ed54bf9d":{"m":90,"g":94},"b57d87c2":{"m":90,"g":94},"98538822":{"m":90,"g":94},"f47a1b1d":{"m":90,"g":94},"93cec433":{"m":90,"g":94},"ba589b88":{"m":90,"g":94},"50876abc":{"m":90,"g":94},"b4c41f72":{"m":90,"g":94},"8b8f2e74":{"m":90,"g":94},"0fc3d992":{"m":90,"g":94},"be2d985d":{"m":90,"g":94},"5b1afa78":{"m":90,"g":94},"c49c1d92":{"m":90,"g":94},"0f1dfa1e":{"m":90,"g":94},"e3ec6bf4":{"m":90,"g":94},"b04df75a":{"m":90,"g":94},"bec3e484":{"m":90,"g":94},"8ab7d93c":{"m":90,"g":94},"5c66c442":{"m":90,"g":94},"aa46ed34":{"m":90,"g":94},"2f4ec752":{"m":90,"g":94},"da47621c":{"m":90,"g":94},"22a6b9fc":{"m":90,"g":94},"b02df20a":{"m":90,"g":94},"bd7cfbd2":{"m":90,"g":94},"4b9971e4":{"m":90,"g":94},"dcc79d32":{"m":90,"g":94},"7046e0fa":{"m":90,"g":94},"930746d9":{"m":90,"g":94},"84727a51":{"m":90,"g":94},"ef326774":{"m":90,"g":94},"021f76e4":{"m":90,"g":94},"777688b8":{"m":90,"g":94},"0ca594ed":{"m":90,"g":94},"31d6dee5":{"m":90,"g":94},"02543b54":{"m":90,"g":94},"25a6a9aa":{"m":90,"g":94},"83d87685":{"m":90,"g":94},"2a5f0100":{"m":90,"g":94},"dbdf76ca":{"m":90,"g":94},"f2a75a66":{"m":90,"g":94},"6b12d6a8":{"m":90,"g":94},"0f218731":{"m":90,"g":94},"14c18d25":{"m":90,"g":94},"90bd3e32":{"m":90,"g":94},"ca929118":{"m":90,"g":94},"344adb00":{"m":90,"g":94},"b56de8f9":{"m":90,"g":94},"ce5ee3bd":{"m":90,"g":94},"a0e4d4eb":{"m":90,"g":94},"2f584455":{"m":90,"g":94},"fe55947a":{"m":90,"g":94},"19995dd7":{"m":90,"g":94},"3b014bc1":{"m":90,"g":94},"d7c3e8e9":{"m":90,"g":94},"8ea7df61":{"m":90,"g":94},"4a102a2b":{"m":90,"g":94},"6406408a":{"m":90,"g":94},"019851d0":{"m":90,"g":94},"2dae104d":{"m":90,"g":94},"cef6655b":{"m":90,"g":94},"27196d41":{"m":90,"g":94},"bb185b0e":{"m":90,"g":94},"7c3a12c0":{"m":91,"g":94},"e846d95e":{"m":91,"g":94},"15f34013":{"m":91,"g":94},"d04163b3":{"m":91,"g":94},"7732bbe4":{"m":91,"g":94},"ed0a0b69":{"m":91,"g":94},"fa42e419":{"m":91,"g":94},"e5afb88b":{"m":91,"g":94},"e5ddeb04":{"m":91,"g":94},"bdbb8d00":{"m":91,"g":94},"34c3f9b2":{"m":91,"g":94},"76139bfb":{"m":91,"g":94},"f8d48fd3":{"m":91,"g":94},"34b6b842":{"m":91,"g":94},"25549433":{"m":91,"g":94},"d6dddc19":{"m":91,"g":94},"55e03b10":{"m":91,"g":94},"8aa68ed5":{"m":91,"g":94},"506c4928":{"m":91,"g":94},"30ceccc7":{"m":91,"g":94},"ac5010e0":{"m":91,"g":94},"3cee035e":{"m":91,"g":94},"30f2a44a":{"m":91,"g":94},"bd4f5818":{"m":91,"g":94},"50f1b6d6":{"m":91,"g":94},"5962e70d":{"m":91,"g":94},"edc21cc8":{"m":91,"g":94},"05c9bc89":{"m":91,"g":94},"b7a2df0a":{"m":91,"g":94},"1998ce40":{"m":91,"g":94},"72676cd6":{"m":91,"g":94},"02bf31ef":{"m":91,"g":94},"5ea5d221":{"m":91,"g":94},"fdfd5224":{"m":91,"g":94},"7f3ee861":{"m":91,"g":94},"bec58910":{"m":91,"g":94},"9edf6608":{"m":91,"g":94},"ab74f8f0":{"m":91,"g":94},"5e7fdc79":{"m":91,"g":94},"4d8d9b8e":{"m":91,"g":94},"5041df2d":{"m":91,"g":94},"256801e9":{"m":91,"g":94},"73b13e69":{"m":91,"g":94},"8609e637":{"m":91,"g":94},"dea2b84b":{"m":91,"g":94},"cfb2fb5a":{"m":91,"g":94},"22bfed75":{"m":91,"g":94},"e879d8b7":{"m":91,"g":94},"09988080":{"m":91,"g":94},"794be55a":{"m":91,"g":94},"187b85b7":{"m":91,"g":94},"ceba0ce4":{"m":91,"g":94},"1ab6be1b":{"m":91,"g":94},"4df5fc21":{"m":91,"g":94},"a06912ad":{"m":91,"g":94},"97011abc":{"m":91,"g":94},"1d6515ef":{"m":91,"g":94},"dea8aa7a":{"m":91,"g":94},"906dbc34":{"m":91,"g":94},"fadf18fd":{"m":91,"g":94},"f88e7085":{"m":91,"g":94},"4f838c09":{"m":91,"g":94},"d20a073b":{"m":91,"g":94},"47367b76":{"m":91,"g":94},"650127a1":{"m":91,"g":94},"3774f078":{"m":91,"g":94},"9179ea15":{"m":91,"g":94},"20a503c7":{"m":91,"g":94},"ffd1a26e":{"m":91,"g":94},"09ae5b20":{"m":91,"g":94},"712bf9ec":{"m":91,"g":94},"9c6a0656":{"m":91,"g":94},"2ae809c5":{"m":91,"g":94},"1de4db9b":{"m":91,"g":94},"31fccf5a":{"m":91,"g":94},"b783c1cb":{"m":91,"g":94},"094c116f":{"m":91,"g":94},"e56685ac":{"m":91,"g":94},"c26d7349":{"m":91,"g":94},"ceaa85c9":{"m":91,"g":94},"0650e517":{"m":91,"g":94},"fc554105":{"m":91,"g":94},"4f204db5":{"m":91,"g":94},"3eb4a800":{"m":91,"g":94},"e7261315":{"m":91,"g":94},"8c16da33":{"m":91,"g":94},"a39d9287":{"m":91,"g":94},"10d60cd4":{"m":91,"g":94},"8a10c4c3":{"m":91,"g":94},"405780bc":{"m":91,"g":94},"1dffee31":{"m":91,"g":94},"70c471a8":{"m":91,"g":94},"1a9c2c92":{"m":91,"g":94},"873ae12c":{"m":91,"g":94},"c64290dc":{"m":91,"g":94},"8e2363dc":{"m":91,"g":94},"69183f88":{"m":92,"g":94},"9b00990b":{"m":92,"g":94},"4d67025a":{"m":92,"g":94},"0e05fe8c":{"m":92,"g":94},"2390a2bc":{"m":92,"g":94},"16d76b9f":{"m":92,"g":94},"5c214257":{"m":92,"g":94},"b8df43ab":{"m":92,"g":94},"a1c1ebe9":{"m":92,"g":94},"fe2a0f96":{"m":92,"g":94},"20beb370":{"m":92,"g":94},"00fbd8a4":{"m":92,"g":94},"802815e4":{"m":92,"g":94},"4c6675c4":{"m":92,"g":94},"e21aa1df":{"m":92,"g":94},"f3cbd245":{"m":92,"g":94},"506a2d59":{"m":92,"g":94},"a07f8ae4":{"m":92,"g":94},"7eb47b0f":{"m":92,"g":94},"bc2e5645":{"m":92,"g":94},"3abc3036":{"m":92,"g":94},"afeed465":{"m":92,"g":94},"587b4c6e":{"m":92,"g":94},"7b9a174a":{"m":92,"g":94},"03c039c4":{"m":92,"g":94},"57ab7769":{"m":92,"g":94},"112b496a":{"m":92,"g":94},"3562256b":{"m":92,"g":94},"5f527834":{"m":92,"g":94},"9f1787fa":{"m":92,"g":94},"8ecad0b1":{"m":92,"g":94},"7151194b":{"m":92,"g":94},"2ed68d7a":{"m":92,"g":94},"e984d507":{"m":92,"g":94},"755f3147":{"m":92,"g":94},"ec5f9c62":{"m":93,"g":94},"62f5522f":{"m":93,"g":94},"01f98730":{"m":93,"g":94},"199d6218":{"m":93,"g":94},"f200af0d":{"m":93,"g":94},"5589b750":{"m":93,"g":94},"c04a8a82":{"m":93,"g":94},"6c903611":{"m":93,"g":94},"77cfea68":{"m":93,"g":94},"8fc910db":{"m":93,"g":94},"75354d9a":{"m":93,"g":94},"4fece12b":{"m":93,"g":94},"c7973222":{"m":93,"g":94},"ef8a29c4":{"m":93,"g":94},"8e9fb43d":{"m":93,"g":94},"83646089":{"m":93,"g":94},"da3890e8":{"m":93,"g":94},"cb432f17":{"m":93,"g":94},"1964c325":{"m":93,"g":94},"af564774":{"m":93,"g":94},"af46f299":{"m":93,"g":94},"16a6b1d8":{"m":93,"g":94},"14229ccf":{"m":93,"g":94},"975a5ec6":{"m":93,"g":94},"1e3e3add":{"m":93,"g":94},"8c298031":{"m":93,"g":94},"4de03953":{"m":93,"g":94},"8b1942c6":{"m":93,"g":94},"489934be":{"m":93,"g":94},"43f93f63":{"m":93,"g":94},"aca1101a":{"m":93,"g":94},"2998c4bd":{"m":93,"g":94},"6840a7bb":{"m":93,"g":94},"c01a1df5":{"m":93,"g":94},"00991723":{"m":93,"g":94},"264dc6e7":{"m":93,"g":94},"646cef2e":{"m":93,"g":94},"1dce6c48":{"m":93,"g":94},"9fcc9a80":{"m":93,"g":94},"ac49dac0":{"m":93,"g":94},"1e0e5497":{"m":93,"g":94},"b5822651":{"m":93,"g":94},"2c4feaf3":{"m":93,"g":94},"2ff572e2":{"m":93,"g":94},"84f2e4a0":{"m":93,"g":94},"8f844db6":{"m":93,"g":94},"36cc3ffd":{"m":93,"g":94},"1bebd315":{"m":93,"g":94},"d3c275b1":{"m":93,"g":94},"b044400d":{"m":93,"g":94},"40e5cb7a":{"m":93,"g":94},"8e64140e":{"m":93,"g":94},"82f021e2":{"m":93,"g":94},"0626f678":{"m":93,"g":94},"09e699bb":{"m":93,"g":94},"b116b21a":{"m":93,"g":94},"88f484ce":{"m":93,"g":94},"8e03b641":{"m":93,"g":94},"b3fa5dc3":{"m":93,"g":94},"00aec6ad":{"m":93,"g":94},"1a08358a":{"m":93,"g":94},"f18a8fdd":{"m":93,"g":94},"a7efbb27":{"m":93,"g":94},"93b6785d":{"m":93,"g":94},"f9eb04dd":{"m":93,"g":94},"3a911b85":{"m":93,"g":94},"886d3449":{"m":93,"g":94},"637bfee4":{"m":93,"g":94},"6005ecee":{"m":93,"g":94},"ff2e9c94":{"m":93,"g":94},"3e34e900":{"m":93,"g":94},"7349717e":{"m":93,"g":94},"392e441a":{"m":93,"g":94},"7248272c":{"m":93,"g":94},"22352d47":{"m":93,"g":94},"c5131f7a":{"m":93,"g":94},"78700893":{"m":93,"g":94},"663c04f7":{"m":93,"g":94},"3b3f1e3a":{"m":93,"g":94},"b691dcc4":{"m":93,"g":94},"0c9c6c75":{"m":93,"g":94},"e3f9b548":{"m":93,"g":94},"b3cff365":{"m":93,"g":94},"8f335b5b":{"m":93,"g":94},"b2264076":{"m":93,"g":94},"04b35190":{"m":93,"g":94},"071a1f51":{"m":93,"g":94},"7c0db3a6":{"m":93,"g":94},"c45e49d8":{"m":93,"g":94},"d8053929":{"m":93,"g":94},"00c7b1ad":{"m":93,"g":94},"82eccae4":{"m":93,"g":94},"a8c10aee":{"m":93,"g":94},"eb429b88":{"m":93,"g":94},"49538d11":{"m":93,"g":94},"cfe2edac":{"m":93,"g":94},"2373faa3":{"m":93,"g":94},"9efb2993":{"m":93,"g":94},"a5317b2f":{"m":93,"g":94},"eb6c2c16":{"m":93,"g":94},"357921aa":{"m":93,"g":94},"c071198c":{"m":93,"g":94},"d7374d74":{"m":93,"g":94},"ce3a3e87":{"m":93,"g":94},"41650b0d":{"m":93,"g":94},"1b951620":{"m":93,"g":94},"29bd4c81":{"m":93,"g":94},"031f64aa":{"m":93,"g":94},"3d7cdb2e":{"m":93,"g":94},"604efe07":{"m":93,"g":94},"1b8cf77b":{"m":93,"g":94},"bb9b608c":{"m":93,"g":94},"066f4ec9":{"m":95,"g":97},"b6b6268c":{"m":95,"g":97},"08702321":{"m":95,"g":97},"64c5907e":{"m":95,"g":97},"128f16a8":{"m":95,"g":97},"49861046":{"m":95,"g":97},"a37e1247":{"m":95,"g":97},"136c6e04":{"m":95,"g":97},"43e20c06":{"m":95,"g":97},"4bab50a6":{"m":95,"g":97},"2e7ab862":{"m":95,"g":97},"51ae4030":{"m":95,"g":97},"653b873b":{"m":95,"g":97},"d379bda4":{"m":95,"g":97},"2b0e1d1c":{"m":95,"g":97},"659907e3":{"m":95,"g":97},"cb9d91ea":{"m":95,"g":97},"6a6e0bb7":{"m":95,"g":97},"076313bd":{"m":95,"g":97},"9abe1163":{"m":95,"g":97},"3646f6bb":{"m":95,"g":94},"35724aa1":{"m":95,"g":94},"2fc824b8":{"m":95,"g":94},"253454de":{"m":95,"g":94},"ea3e7ffe":{"m":95,"g":94},"8d4a01cb":{"m":95,"g":94},"a3398d84":{"m":95,"g":94},"ba69c153":{"m":95,"g":94},"3589aa79":{"m":95,"g":94},"e00715eb":{"m":95,"g":94},"ea4bf122":{"m":95,"g":94},"a291439a":{"m":95,"g":94},"54411f6a":{"m":95,"g":94},"625018d2":{"m":95,"g":94},"5732d904":{"m":95,"g":94},"eb118d88":{"m":96,"g":97},"732fc8e4":{"m":96,"g":97},"f2d5c492":{"m":96,"g":97},"2a2d3478":{"m":96,"g":97},"aa205609":{"m":96,"g":97},"61bb2858":{"m":96,"g":97},"880221bd":{"m":96,"g":97},"8f3173d0":{"m":96,"g":97},"26118a13":{"m":96,"g":97},"475a249b":{"m":96,"g":97},"191d836f":{"m":96,"g":97},"86044712":{"m":96,"g":97},"61555307":{"m":96,"g":97},"49a5915f":{"m":96,"g":97},"766392c6":{"m":96,"g":97},"4a0d1919":{"m":96,"g":97},"57482415":{"m":96,"g":97},"2d54d4bb":{"m":96,"g":97},"b5e3d603":{"m":96,"g":97},"4ed57807":{"m":96,"g":97},"dd445a41":{"m":96,"g":97},"7590f522":{"m":96,"g":97},"f9df11ae":{"m":96,"g":97},"d389bedf":{"m":96,"g":97},"ac80f4da":{"m":96,"g":97},"d487555f":{"m":96,"g":97},"e5888edd":{"m":96,"g":97},"01c00004":{"m":98,"g":103},"0dfe2491":{"m":98,"g":103},"ff45ab7a":{"m":98,"g":103},"0f8b5386":{"m":98,"g":103},"c33499a6":{"m":98,"g":103},"e50109f2":{"m":98,"g":103},"69adc4f8":{"m":98,"g":103},"11483785":{"m":98,"g":103},"7b68d271":{"m":98,"g":103},"74f59ae5":{"m":98,"g":103},"6936be32":{"m":98,"g":103},"9b5de6cb":{"m":98,"g":97},"5c8365a0":{"m":98,"g":97},"8430bfe3":{"m":98,"g":97},"c9e8613c":{"m":98,"g":97},"429bb0ef":{"m":98,"g":97},"7eebd440":{"m":98,"g":97},"93d124ef":{"m":98,"g":97},"1fc455e8":{"m":98,"g":97},"465968b2":{"m":98,"g":97},"750838ad":{"m":98,"g":97},"99aefa03":{"m":98,"g":97},"bbcfbc1a":{"m":98,"g":97},"83c104b1":{"m":98,"g":97},"2db6719c":{"m":98,"g":97},"55381a46":{"m":98,"g":97},"a589a071":{"m":98,"g":97},"f62d75b6":{"m":98,"g":97},"0f9b11e3":{"m":98,"g":97},"877e35d7":{"m":98,"g":97},"cbdfb771":{"m":98,"g":97},"282eb59f":{"m":98,"g":97},"4540a466":{"m":98,"g":97},"abda2542":{"m":98,"g":97},"8cddfa56":{"m":98,"g":97},"4e3defe5":{"m":98,"g":97},"60468da4":{"m":98,"g":97},"41d33e47":{"m":98,"g":97},"bfdd226f":{"m":98,"g":97},"3de617a7":{"m":98,"g":97},"bb0e8a32":{"m":98,"g":97},"1b427dae":{"m":98,"g":97},"f3d97361":{"m":98,"g":97},"561dd7b2":{"m":98,"g":97},"f98e88b9":{"m":98,"g":97},"15ad6c90":{"m":98,"g":97},"cfab0ff6":{"m":98,"g":97},"b763cf7e":{"m":98,"g":97},"8fcc55cf":{"m":98,"g":97},"610381b7":{"m":98,"g":97},"1403ea56":{"m":98,"g":97},"b7e951a6":{"m":98,"g":97},"d918ab79":{"m":98,"g":97},"3964b352":{"m":98,"g":97},"9c7a4618":{"m":98,"g":97},"7750b91c":{"m":98,"g":97},"c8f31042":{"m":98,"g":97},"1f76fc87":{"m":98,"g":97},"6737671c":{"m":98,"g":97},"fd63b62e":{"m":98,"g":97},"719b29f2":{"m":98,"g":97},"d0510f08":{"m":98,"g":97},"9d33fcfb":{"m":98,"g":97},"7891bac1":{"m":98,"g":97},"48c1fa7b":{"m":98,"g":97},"8aa5ae6b":{"m":98,"g":97},"8a323557":{"m":98,"g":97},"6e92da8f":{"m":98,"g":97},"e1020dc5":{"m":98,"g":97},"3586b4ce":{"m":98,"g":97},"42960214":{"m":98,"g":97},"01857fab":{"m":98,"g":97},"519ff5c8":{"m":98,"g":97},"af1cc8fe":{"m":98,"g":97},"49b87774":{"m":98,"g":97},"02404a1e":{"m":98,"g":97},"5c08a36c":{"m":98,"g":97},"9069884b":{"m":98,"g":97},"8a7a7770":{"m":98,"g":97},"795668dc":{"m":98,"g":97},"4395c87a":{"m":98,"g":97},"c28ad199":{"m":98,"g":97},"570d3343":{"m":98,"g":97},"d9eb5efc":{"m":98,"g":97},"6dc4af49":{"m":98,"g":97},"b188a89a":{"m":98,"g":97},"497efe74":{"m":98,"g":97},"69f453e5":{"m":98,"g":97},"3bc43c68":{"m":98,"g":97},"7498522f":{"m":98,"g":97},"194841e3":{"m":98,"g":97},"ebff5fcb":{"m":98,"g":97},"f06bd210":{"m":98,"g":97},"14f1f151":{"m":98,"g":97},"38216cf0":{"m":98,"g":97},"4a883795":{"m":98,"g":97},"f1f1d1d4":{"m":98,"g":97},"9120e83d":{"m":98,"g":97},"6e923dbd":{"m":98,"g":97},"c268c11c":{"m":98,"g":97},"e6d59884":{"m":98,"g":97},"9b560c3e":{"m":98,"g":97},"5dc5866e":{"m":98,"g":97},"64e78bb3":{"m":98,"g":97},"7c39e8a1":{"m":98,"g":97},"d969504d":{"m":98,"g":97},"1ebec1a8":{"m":98,"g":97},"d4d0c7c3":{"m":98,"g":97},"8d2cf38c":{"m":98,"g":97},"2117f82d":{"m":98,"g":97},"c07f647c":{"m":98,"g":97},"07452cbe":{"m":98,"g":97},"a562c8a3":{"m":98,"g":97},"cb736df8":{"m":98,"g":97},"e2ed9d04":{"m":98,"g":97},"b5dd5e87":{"m":98,"g":97},"9379da77":{"m":98,"g":97},"0c55cbcf":{"m":98,"g":97},"c46e069d":{"m":98,"g":97},"42fc4410":{"m":98,"g":97},"5f6756b0":{"m":98,"g":97},"98aa836b":{"m":98,"g":97},"22bd857c":{"m":98,"g":97},"ccfa0841":{"m":98,"g":97},"bcc5ba94":{"m":98,"g":97},"cee9f329":{"m":98,"g":97},"2272c2a5":{"m":99,"g":103},"3ec0b212":{"m":99,"g":103},"58c468f4":{"m":99,"g":103},"f8ca2368":{"m":99,"g":103},"d8ee1564":{"m":99,"g":103},"7181ec8c":{"m":99,"g":103},"ed2e313e":{"m":99,"g":103},"f8260f25":{"m":99,"g":103},"12cb760a":{"m":99,"g":103},"1b9cea5a":{"m":99,"g":103},"9045cc1e":{"m":99,"g":103},"70e37b97":{"m":99,"g":103},"15d27591":{"m":99,"g":103},"af4b9bae":{"m":99,"g":103},"7ad6b766":{"m":99,"g":103},"c0fb25e9":{"m":99,"g":103},"28d4d472":{"m":99,"g":103},"f4674df6":{"m":99,"g":103},"d40846d4":{"m":99,"g":103},"145482f4":{"m":99,"g":103},"39fe1e88":{"m":99,"g":103},"33c4b4d0":{"m":99,"g":103},"8d1c5b94":{"m":99,"g":103},"a167fd0b":{"m":99,"g":103},"2f86f3ad":{"m":99,"g":103},"bfb118c0":{"m":99,"g":103},"f6e07f27":{"m":99,"g":103},"5dd0f870":{"m":99,"g":103},"f7e102d5":{"m":99,"g":103},"0e5fa677":{"m":99,"g":103},"624a3b8d":{"m":99,"g":103},"01079e17":{"m":99,"g":103},"0e7a5b26":{"m":99,"g":103},"4953f4ca":{"m":99,"g":103},"38000a5f":{"m":99,"g":103},"70251e93":{"m":99,"g":103},"c87d4fec":{"m":99,"g":103},"a99801e0":{"m":99,"g":103},"4c605235":{"m":99,"g":103},"6f8f4aee":{"m":99,"g":103},"0c8dab9e":{"m":99,"g":103},"f39037ff":{"m":99,"g":103},"ce86e201":{"m":99,"g":103},"b4326330":{"m":99,"g":103},"8abd3e77":{"m":99,"g":103},"e885bfdc":{"m":99,"g":103},"e2d66f60":{"m":99,"g":103},"45bc170b":{"m":100,"g":103},"22623699":{"m":100,"g":103},"fb4ce17d":{"m":100,"g":103},"25f73c6c":{"m":100,"g":103},"581e7dcb":{"m":100,"g":103},"484d0e02":{"m":100,"g":103},"5922c0cb":{"m":100,"g":103},"6d6a8bc2":{"m":100,"g":103},"2fd5c704":{"m":100,"g":103},"4ad97370":{"m":100,"g":103},"28103384":{"m":100,"g":103},"fe6a445d":{"m":100,"g":103},"dd487e55":{"m":100,"g":103},"bb81daef":{"m":100,"g":103},"58dd95fb":{"m":100,"g":103},"b47eda33":{"m":100,"g":103},"e983d666":{"m":100,"g":103},"b58c3c28":{"m":100,"g":103},"df906455":{"m":100,"g":103},"95217a9b":{"m":100,"g":103},"22e00eeb":{"m":100,"g":103},"b3eac168":{"m":100,"g":103},"10ee8955":{"m":100,"g":103},"bf3352c5":{"m":100,"g":103},"4d921f2b":{"m":100,"g":103},"44d600cd":{"m":100,"g":103},"5c9c275b":{"m":100,"g":103},"bf0f448f":{"m":100,"g":103},"36d6f0ba":{"m":100,"g":103},"2a1936de":{"m":100,"g":103},"2ab97023":{"m":100,"g":103},"0bcc195f":{"m":100,"g":103},"91e3d154":{"m":100,"g":103},"85486b6f":{"m":100,"g":103},"e34cf6ad":{"m":100,"g":103},"62222bd2":{"m":100,"g":103},"ed0fdbf3":{"m":100,"g":103},"b602f423":{"m":100,"g":103},"426b7493":{"m":100,"g":103},"528bd1ed":{"m":100,"g":103},"62a6b7c7":{"m":100,"g":103},"76154631":{"m":100,"g":103},"5c705b1d":{"m":100,"g":103},"b7094a5e":{"m":100,"g":103},"da0c0260":{"m":100,"g":103},"3212c2ad":{"m":100,"g":103},"53475674":{"m":100,"g":103},"ce32bc2b":{"m":100,"g":103},"e236d8fe":{"m":100,"g":103},"4fa44d63":{"m":100,"g":103},"e6312d27":{"m":100,"g":103},"8af145b7":{"m":100,"g":103},"6478831b":{"m":101,"g":103},"2e1d2d7e":{"m":101,"g":103},"fb16fbaf":{"m":101,"g":103},"0ce84c82":{"m":101,"g":103},"59d0bf01":{"m":101,"g":103},"7df2c0c2":{"m":101,"g":103},"69712e6f":{"m":101,"g":103},"001bffca":{"m":101,"g":103},"7c969717":{"m":101,"g":103},"8240a6b0":{"m":101,"g":103},"3a04aa4b":{"m":101,"g":103},"bd516949":{"m":101,"g":103},"74e7e457":{"m":101,"g":103},"1466c1b8":{"m":101,"g":103},"9c138a04":{"m":101,"g":103},"c8f549d9":{"m":101,"g":103},"134fa43e":{"m":101,"g":103},"ccfe52a0":{"m":101,"g":103},"747dd450":{"m":101,"g":103},"b5821592":{"m":101,"g":103},"a9dd3ec3":{"m":101,"g":103},"02328864":{"m":102,"g":103},"7a1f7fc5":{"m":102,"g":103},"51c38163":{"m":102,"g":103},"09f1a247":{"m":102,"g":103},"32fa1e9c":{"m":102,"g":103},"e7dc163f":{"m":102,"g":103},"e179e0b7":{"m":102,"g":103},"d9049592":{"m":102,"g":103},"26c8a310":{"m":102,"g":103},"5963e505":{"m":102,"g":103},"43118f5f":{"m":102,"g":103},"a5f5ab40":{"m":102,"g":103},"59aab76f":{"m":102,"g":103},"659bfd10":{"m":102,"g":103},"67e53b16":{"m":102,"g":103},"9b9e8253":{"m":102,"g":103},"66a398f4":{"m":102,"g":103},"29980334":{"m":102,"g":103},"a79a5d70":{"m":102,"g":103},"ec5f9442":{"m":102,"g":103},"3bdcdd13":{"m":102,"g":103},"a730ce81":{"m":102,"g":103},"55ecdc0a":{"m":102,"g":103},"e3f08c77":{"m":102,"g":103},"a9fd8033":{"m":102,"g":103},"2fbb754e":{"m":102,"g":103},"a85ebf50":{"m":102,"g":103},"9effeb5b":{"m":102,"g":103},"1992ef9b":{"m":102,"g":103},"c0fd77e8":{"m":102,"g":103},"a4c3b121":{"m":102,"g":103},"5973675b":{"m":102,"g":103},"4d16c88b":{"m":102,"g":103},"7a4309cc":{"m":102,"g":103},"81367066":{"m":102,"g":103},"263c9236":{"m":102,"g":103},"33f0de33":{"m":105,"g":107},"e7e5a305":{"m":105,"g":107},"dd7ca006":{"m":105,"g":107},"9305ea6c":{"m":105,"g":107},"aa4c66b5":{"m":105,"g":107},"39decec1":{"m":105,"g":104},"f6f46f46":{"m":105,"g":104},"2886e23d":{"m":105,"g":104},"99795d61":{"m":105,"g":104},"fe5086fd":{"m":105,"g":104},"04913430":{"m":105,"g":104},"0ad098b4":{"m":105,"g":104},"4a6e7a66":{"m":105,"g":104},"4b04998d":{"m":105,"g":104},"3dde8619":{"m":105,"g":104},"b7170cc8":{"m":105,"g":104},"5c14515f":{"m":105,"g":104},"2cd2e27f":{"m":105,"g":104},"743638bc":{"m":105,"g":104},"061c8959":{"m":105,"g":104},"4acf6902":{"m":105,"g":104},"aee0ef52":{"m":105,"g":103},"ae807774":{"m":105,"g":103},"8fbcfd07":{"m":105,"g":103},"3c307dc0":{"m":105,"g":103},"5d15fb8c":{"m":105,"g":103},"016fd251":{"m":105,"g":103},"8cd34458":{"m":106,"g":107},"0e0eef00":{"m":106,"g":107},"cb099d20":{"m":106,"g":107},"7a913301":{"m":106,"g":107},"5ce5093b":{"m":106,"g":107},"6f9baf10":{"m":106,"g":107},"a31b7a70":{"m":106,"g":107},"7ed8e51b":{"m":106,"g":107},"32f28154":{"m":106,"g":107},"f7b2853f":{"m":106,"g":107},"b0add2da":{"m":106,"g":107},"0305c505":{"m":106,"g":107},"8675bdf2":{"m":106,"g":107},"a437aa99":{"m":106,"g":107},"0e612dbf":{"m":106,"g":107},"9f47d686":{"m":106,"g":107},"d9def43d":{"m":106,"g":107},"e273aa6d":{"m":106,"g":107},"828a4fe9":{"m":106,"g":107},"8ada1ab6":{"m":106,"g":107},"e314b084":{"m":106,"g":107},"403566bc":{"m":106,"g":107},"0a56b721":{"m":106,"g":107},"603f5ce0":{"m":106,"g":107},"6d4fd882":{"m":106,"g":107},"f9f0138f":{"m":106,"g":107},"ac6962cc":{"m":106,"g":107},"4ca43b06":{"m":106,"g":107},"ea93079b":{"m":106,"g":107},"4bec99ec":{"m":106,"g":107},"89caf7a3":{"m":106,"g":107},"b27b1191":{"m":106,"g":107},"f642524f":{"m":106,"g":107},"82e6c3a6":{"m":106,"g":107},"b89d37cb":{"m":106,"g":107},"5deab128":{"m":106,"g":107},"d1c4d51c":{"m":106,"g":107},"1fe691a4":{"m":106,"g":107},"e2521926":{"m":106,"g":107},"07e46eca":{"m":106,"g":107},"ab9b893e":{"m":106,"g":107},"6a7528e6":{"m":106,"g":107},"2ae95d17":{"m":106,"g":107},"2d401bd9":{"m":106,"g":107},"b17c5b01":{"m":106,"g":107},"db7343c9":{"m":106,"g":107},"533cb5b2":{"m":106,"g":107},"6bdd2786":{"m":106,"g":107},"46e9d1c7":{"m":106,"g":107},"6c88f6c8":{"m":106,"g":107},"c8d3a402":{"m":106,"g":107},"7e831efe":{"m":106,"g":107},"20b5563e":{"m":106,"g":107},"97a38ee8":{"m":108,"g":115},"86d10d22":{"m":108,"g":115},"83871aa1":{"m":108,"g":115},"b1b3f0b3":{"m":108,"g":115},"34e5e11f":{"m":108,"g":115},"2600fc0d":{"m":108,"g":115},"ccd3fb94":{"m":108,"g":115},"c9dd70fb":{"m":108,"g":115},"6b2b8bf0":{"m":108,"g":115},"4edbe0d5":{"m":108,"g":115},"0374304a":{"m":108,"g":115},"127d4b0d":{"m":108,"g":115},"7e880286":{"m":108,"g":115},"446c8e4c":{"m":108,"g":115},"5ef545e6":{"m":108,"g":115},"c4500233":{"m":108,"g":115},"f445a1d9":{"m":108,"g":115},"e5638573":{"m":108,"g":115},"f556ac8b":{"m":108,"g":115},"110a6598":{"m":108,"g":115},"49f9d025":{"m":108,"g":115},"0f587e80":{"m":108,"g":115},"6078d5fc":{"m":108,"g":115},"70cf4abc":{"m":108,"g":115},"cebf4599":{"m":108,"g":115},"9c0c1e30":{"m":108,"g":115},"a1f011d0":{"m":108,"g":115},"9ec314c6":{"m":108,"g":115},"fedfe91c":{"m":108,"g":115},"988accbc":{"m":108,"g":115},"b6b2287e":{"m":108,"g":115},"243e745d":{"m":108,"g":115},"61a0e600":{"m":108,"g":115},"0f8cee8c":{"m":108,"g":115},"816c4c85":{"m":108,"g":115},"13ec8d42":{"m":108,"g":115},"05bd7897":{"m":108,"g":115},"5fd311d3":{"m":108,"g":115},"53e2cd46":{"m":108,"g":115},"9708d353":{"m":108,"g":115},"704ced1b":{"m":108,"g":115},"3cc3d9b9":{"m":108,"g":115},"0b3a5b11":{"m":108,"g":115},"6c855db8":{"m":108,"g":115},"0f9318f7":{"m":108,"g":115},"849957bc":{"m":108,"g":115},"cded039b":{"m":108,"g":115},"275f9df3":{"m":108,"g":115},"e8449ab5":{"m":108,"g":115},"4746aaea":{"m":108,"g":115},"10d34f74":{"m":108,"g":115},"9ba72530":{"m":108,"g":115},"9c8e4f69":{"m":108,"g":115},"78ae1758":{"m":108,"g":115},"dae9a80f":{"m":108,"g":115},"e85cb1ce":{"m":108,"g":115},"55d336cb":{"m":108,"g":115},"de4990a5":{"m":108,"g":115},"029e0af3":{"m":108,"g":115},"64574ef8":{"m":108,"g":115},"18da2c96":{"m":108,"g":115},"9b5f0f64":{"m":108,"g":115},"70bb066e":{"m":108,"g":115},"2c4b4b78":{"m":108,"g":115},"7cd2ee06":{"m":108,"g":115},"eb19ccad":{"m":108,"g":115},"25ef53f0":{"m":108,"g":115},"c674bf9c":{"m":108,"g":115},"af1973b8":{"m":108,"g":115},"5cfbb4c1":{"m":108,"g":115},"e6523102":{"m":108,"g":115},"3828db43":{"m":108,"g":115},"88fbc31b":{"m":108,"g":115},"8f5b9910":{"m":108,"g":115},"ef3004d9":{"m":108,"g":115},"84719b52":{"m":108,"g":115},"e99729c9":{"m":108,"g":115},"c10b8e6a":{"m":108,"g":115},"d4bce297":{"m":108,"g":115},"b0980af8":{"m":108,"g":115},"24eaebeb":{"m":108,"g":115},"a91e90d9":{"m":108,"g":115},"f96413c4":{"m":108,"g":115},"08ebdf79":{"m":108,"g":115},"42c87045":{"m":108,"g":115},"c9bf3877":{"m":108,"g":115},"de2dd738":{"m":108,"g":115},"1ec97697":{"m":108,"g":115},"d8ed60f2":{"m":108,"g":115},"f1b0eda5":{"m":108,"g":115},"f20b6a3f":{"m":108,"g":115},"3680d6f8":{"m":108,"g":115},"f5154495":{"m":108,"g":115},"e0ce171d":{"m":108,"g":115},"fe43e889":{"m":108,"g":115},"5ae5ecaa":{"m":108,"g":115},"5fbad308":{"m":108,"g":115},"7638f5e4":{"m":108,"g":115},"b45f753c":{"m":108,"g":115},"c5057262":{"m":108,"g":115},"46fe8b8c":{"m":108,"g":115},"0b95a01a":{"m":108,"g":115},"a3b810eb":{"m":108,"g":115},"94959237":{"m":108,"g":115},"f4fafacc":{"m":108,"g":115},"01d47a27":{"m":108,"g":115},"ecc9f3e4":{"m":108,"g":115},"7e8187e0":{"m":108,"g":115},"e483ab6d":{"m":108,"g":115},"720cd308":{"m":108,"g":115},"ce67b2d5":{"m":108,"g":115},"3c2c9f6c":{"m":108,"g":115},"a31ea448":{"m":108,"g":115},"439df454":{"m":108,"g":115},"5626e20b":{"m":108,"g":115},"c6c379ab":{"m":108,"g":115},"c2fbf60f":{"m":108,"g":115},"98b44e9e":{"m":108,"g":115},"6805f6da":{"m":108,"g":115},"ca533580":{"m":108,"g":115},"886454e8":{"m":108,"g":115},"0cf3fbeb":{"m":108,"g":115},"2256d62d":{"m":108,"g":115},"6cdcbcc6":{"m":108,"g":115},"c480a3f6":{"m":108,"g":115},"6e316588":{"m":108,"g":115},"24247b41":{"m":108,"g":115},"4c0bb411":{"m":108,"g":115},"968e1818":{"m":108,"g":115},"d08663ee":{"m":108,"g":115},"716e6827":{"m":108,"g":115},"84b30d9e":{"m":108,"g":115},"ff0cf51c":{"m":108,"g":115},"a1c7f742":{"m":108,"g":115},"ebbb75e9":{"m":108,"g":115},"b341b7db":{"m":108,"g":115},"b498cd21":{"m":108,"g":115},"0fc54b97":{"m":108,"g":115},"b3c1f2e4":{"m":108,"g":115},"be1a3cd9":{"m":108,"g":115},"4b74c3fc":{"m":108,"g":115},"ce3ca9b0":{"m":108,"g":115},"4d98e486":{"m":108,"g":115},"3d77a318":{"m":108,"g":115},"845d12a9":{"m":108,"g":115},"e47800e1":{"m":108,"g":115},"bb10e3a1":{"m":108,"g":115},"fda762a2":{"m":108,"g":115},"1df84ff4":{"m":108,"g":115},"66d6be08":{"m":108,"g":115},"1c1f8a11":{"m":108,"g":115},"384f8ab5":{"m":108,"g":115},"6a9d6ca3":{"m":108,"g":115},"94371dbb":{"m":108,"g":115},"740f0630":{"m":108,"g":115},"81da16f6":{"m":108,"g":115},"bc938ea1":{"m":108,"g":115},"eff4eb3f":{"m":108,"g":115},"87dab548":{"m":108,"g":115},"5121af46":{"m":108,"g":115},"983aa496":{"m":108,"g":115},"9c3e95d9":{"m":108,"g":115},"e52c3866":{"m":108,"g":115},"da53e13c":{"m":108,"g":115},"d7e38b2f":{"m":108,"g":115},"21b88460":{"m":108,"g":115},"0c8594e6":{"m":108,"g":115},"c186feed":{"m":108,"g":115},"84b006b2":{"m":108,"g":115},"8ca07bd9":{"m":108,"g":115},"4fc09e0d":{"m":108,"g":115},"a3d99d6d":{"m":108,"g":115},"189af908":{"m":108,"g":115},"f8644a56":{"m":108,"g":115},"e3e75a78":{"m":108,"g":115},"d4db9b02":{"m":108,"g":115},"f7dd651d":{"m":108,"g":115},"9d54c6e6":{"m":108,"g":115},"1f9d65f5":{"m":108,"g":115},"29589512":{"m":108,"g":115},"584e1ab2":{"m":108,"g":115},"392de007":{"m":108,"g":115},"004f7f19":{"m":108,"g":115},"d2fbf2de":{"m":108,"g":115},"fab0f6e7":{"m":108,"g":115},"27985c27":{"m":108,"g":115},"ac474869":{"m":108,"g":115},"0b1e04f0":{"m":108,"g":115},"c1c7dc45":{"m":108,"g":115},"2cc9eeab":{"m":108,"g":115},"63d82a77":{"m":108,"g":115},"53dcc750":{"m":108,"g":115},"432f2053":{"m":108,"g":115},"1fea998a":{"m":108,"g":115},"5aa1ebd2":{"m":108,"g":115},"4dbf4360":{"m":108,"g":115},"3d6be1fb":{"m":108,"g":115},"4063234c":{"m":108,"g":115},"83feef5b":{"m":108,"g":115},"2871eacc":{"m":108,"g":115},"ac15bdc1":{"m":108,"g":115},"d6451c3f":{"m":108,"g":115},"841810f2":{"m":108,"g":115},"733446dd":{"m":108,"g":115},"4c22897a":{"m":108,"g":115},"1bc183c6":{"m":108,"g":115},"b87aacb5":{"m":108,"g":115},"b3363cc1":{"m":108,"g":115},"98457c04":{"m":108,"g":115},"0fc8bf2c":{"m":108,"g":115},"a669bc2f":{"m":108,"g":115},"6b7c2471":{"m":108,"g":115},"a027a9b4":{"m":108,"g":115},"9e426466":{"m":108,"g":115},"2f20f430":{"m":108,"g":115},"65736dc5":{"m":108,"g":115},"7b56e494":{"m":108,"g":115},"0ff6d1fc":{"m":108,"g":115},"4a16a71c":{"m":108,"g":115},"a16923ef":{"m":108,"g":115},"6337d905":{"m":108,"g":115},"71fb8c95":{"m":108,"g":115},"94f44b88":{"m":108,"g":115},"3b3b3baf":{"m":108,"g":115},"35e6bc92":{"m":108,"g":115},"9394ed63":{"m":108,"g":115},"930fe467":{"m":108,"g":115},"13c48dcf":{"m":108,"g":115},"8723b4f1":{"m":108,"g":115},"62f99e08":{"m":108,"g":115},"86a0be65":{"m":108,"g":115},"0edda320":{"m":108,"g":115},"924827c3":{"m":108,"g":115},"c81daf83":{"m":108,"g":115},"25caa7a8":{"m":108,"g":115},"03d11449":{"m":108,"g":115},"83123f48":{"m":108,"g":115},"48afa8f1":{"m":108,"g":115},"2ecbd8b8":{"m":108,"g":115},"305b27c1":{"m":108,"g":115},"1ce30dd1":{"m":108,"g":115},"c9ee7385":{"m":108,"g":115},"1f9ec653":{"m":108,"g":115},"ad359d1c":{"m":108,"g":115},"5f5b3b24":{"m":108,"g":115},"4caca4f6":{"m":108,"g":115},"f2a5de28":{"m":108,"g":115},"445f9dca":{"m":108,"g":115},"3a9afe2a":{"m":108,"g":115},"9aea2555":{"m":108,"g":115},"fcc11e5e":{"m":108,"g":115},"5190ba7f":{"m":108,"g":115},"5438886c":{"m":108,"g":115},"9c83d74d":{"m":108,"g":115},"b4ac2b9c":{"m":108,"g":115},"83262dcb":{"m":108,"g":115},"c46c75f8":{"m":108,"g":115},"2aaf22c4":{"m":108,"g":115},"29a610b4":{"m":108,"g":115},"5ded39ca":{"m":108,"g":115},"4093d460":{"m":108,"g":115},"9d68bdb2":{"m":108,"g":115},"a2184901":{"m":108,"g":115},"0eec4cb6":{"m":108,"g":115},"ff1f6825":{"m":108,"g":115},"9f78f391":{"m":108,"g":115},"f508cd3c":{"m":108,"g":115},"44e86480":{"m":108,"g":115},"8c07fabd":{"m":108,"g":115},"90f44b74":{"m":108,"g":115},"38907fe6":{"m":108,"g":115},"f9afa7dc":{"m":108,"g":115},"0d9e89ec":{"m":108,"g":115},"3d64fda3":{"m":108,"g":115},"3bffe112":{"m":108,"g":115},"44426e54":{"m":108,"g":115},"9f24dfef":{"m":108,"g":115},"89f1d4f5":{"m":108,"g":115},"75e6a7cd":{"m":108,"g":115},"6f81a710":{"m":108,"g":115},"a6452b71":{"m":108,"g":115},"f4ae50e9":{"m":108,"g":115},"84cb449e":{"m":108,"g":115},"f003cd35":{"m":108,"g":115},"9d834fdc":{"m":108,"g":115},"b3279251":{"m":108,"g":115},"067068f2":{"m":108,"g":115},"6beeff41":{"m":108,"g":115},"2e8e7e35":{"m":108,"g":115},"2449a0af":{"m":108,"g":115},"0f229c07":{"m":108,"g":115},"dd001a54":{"m":108,"g":115},"4ea9d74a":{"m":108,"g":115},"dd949ace":{"m":108,"g":115},"f2887498":{"m":108,"g":115},"8ecf6b9d":{"m":108,"g":115},"0418b9d4":{"m":108,"g":115},"e322a94d":{"m":108,"g":115},"2c7f01bc":{"m":108,"g":115},"b58ae7a2":{"m":108,"g":115},"6345069f":{"m":108,"g":115},"ce9cf353":{"m":108,"g":115},"f8a173bb":{"m":108,"g":115},"6b847a9a":{"m":108,"g":115},"473400e4":{"m":108,"g":115},"dd665f96":{"m":108,"g":115},"3817a37d":{"m":108,"g":115},"7ba5ad57":{"m":108,"g":115},"19bc77f0":{"m":108,"g":115},"86497d99":{"m":108,"g":115},"5c31b35d":{"m":108,"g":115},"ef48d554":{"m":108,"g":115},"a886564a":{"m":108,"g":115},"9a44b643":{"m":108,"g":115},"41d71ca4":{"m":108,"g":115},"20cfc5a2":{"m":108,"g":115},"48b8b4c1":{"m":108,"g":115},"323bc2f5":{"m":108,"g":115},"137e75da":{"m":108,"g":115},"52e1f52f":{"m":108,"g":115},"50188092":{"m":108,"g":115},"326a901d":{"m":108,"g":115},"6e0b6468":{"m":108,"g":115},"4a9f3eef":{"m":108,"g":115},"1b7afad0":{"m":108,"g":115},"f29aba8c":{"m":108,"g":115},"faa25df1":{"m":108,"g":115},"7b81f956":{"m":108,"g":115},"d3e67deb":{"m":108,"g":115},"442534aa":{"m":108,"g":115},"de8b8b6e":{"m":108,"g":115},"3f2e315f":{"m":108,"g":115},"6e215118":{"m":108,"g":115},"a47baff1":{"m":108,"g":115},"fd7e15b7":{"m":108,"g":115},"fc42ff7b":{"m":108,"g":115},"7c0db868":{"m":108,"g":115},"706bd69c":{"m":108,"g":115},"23f2afb2":{"m":108,"g":115},"a60f88b5":{"m":108,"g":115},"591c232f":{"m":108,"g":115},"f352b793":{"m":108,"g":115},"6642e3a2":{"m":108,"g":115},"67a7d1f6":{"m":108,"g":115},"92cbef59":{"m":108,"g":115},"b3359dc9":{"m":108,"g":115},"7b7e5615":{"m":108,"g":115},"1a8706c8":{"m":108,"g":115},"7d3af603":{"m":108,"g":115},"4e7f0252":{"m":108,"g":115},"36bfddec":{"m":108,"g":115},"91e2f902":{"m":108,"g":115},"a59cbea9":{"m":108,"g":115},"53f7874a":{"m":108,"g":115},"61a46804":{"m":108,"g":115},"9020f7fc":{"m":108,"g":115},"dd650e0e":{"m":108,"g":115},"a9471542":{"m":108,"g":115},"41357e51":{"m":108,"g":115},"e2fd2b9c":{"m":108,"g":115},"7490e3f6":{"m":108,"g":115},"6ee6619b":{"m":108,"g":115},"54ea57f2":{"m":108,"g":115},"b4c9f38a":{"m":108,"g":115},"11325474":{"m":108,"g":115},"1d24db83":{"m":108,"g":115},"44401358":{"m":108,"g":115},"9c7e3924":{"m":108,"g":115},"08fab2b0":{"m":108,"g":115},"0d1e27a0":{"m":108,"g":115},"774b47f3":{"m":108,"g":115},"76915d68":{"m":108,"g":115},"39fd1788":{"m":108,"g":115},"ed0a3dd5":{"m":108,"g":115},"2e901e89":{"m":108,"g":115},"d3be9710":{"m":108,"g":115},"3e7ff1ab":{"m":108,"g":115},"aaf0ad8c":{"m":108,"g":115},"361379b5":{"m":108,"g":115},"1ac16add":{"m":108,"g":115},"c3a5fb3b":{"m":108,"g":115},"4bf6e5a6":{"m":108,"g":115},"3ae33fcd":{"m":108,"g":115},"500b15c9":{"m":108,"g":107},"16a4c66d":{"m":108,"g":107},"89e6521c":{"m":108,"g":107},"fd05b567":{"m":108,"g":107},"482c3db2":{"m":108,"g":107},"47824c14":{"m":108,"g":107},"c36a6693":{"m":108,"g":107},"62f8eb48":{"m":108,"g":107},"b7cd7430":{"m":108,"g":107},"a69b6370":{"m":108,"g":107},"2d120f8b":{"m":108,"g":107},"4f2e1490":{"m":108,"g":107},"3fa3c6cd":{"m":108,"g":107},"6210e2c4":{"m":108,"g":107},"6ad6c8c9":{"m":108,"g":107},"5b6acc14":{"m":108,"g":107},"4373df55":{"m":108,"g":107},"c0e84297":{"m":108,"g":107},"92cc32d9":{"m":108,"g":107},"cbbd685a":{"m":108,"g":107},"78aad910":{"m":108,"g":107},"288ae41f":{"m":108,"g":107},"01c99a99":{"m":108,"g":107},"b114a810":{"m":108,"g":107},"0475448e":{"m":108,"g":107},"399e7ec8":{"m":108,"g":107},"1bd53168":{"m":108,"g":107},"aeac900c":{"m":108,"g":107},"4fc5f2f9":{"m":108,"g":107},"168033d5":{"m":108,"g":107},"cbbb7383":{"m":108,"g":107},"89588179":{"m":108,"g":107},"8c7bb39d":{"m":108,"g":107},"ca47e24f":{"m":108,"g":107},"d26ca84f":{"m":108,"g":107},"8128e08d":{"m":108,"g":107},"5d62b56f":{"m":108,"g":107},"3ae8e3ea":{"m":108,"g":107},"c1d2061f":{"m":108,"g":107},"556e4143":{"m":108,"g":107},"4ef47839":{"m":108,"g":107},"32d9e39a":{"m":108,"g":107},"4f4e0e41":{"m":108,"g":107},"901ab758":{"m":108,"g":107},"8e8545ca":{"m":108,"g":107},"a4b0d5c9":{"m":108,"g":107},"40e3b2be":{"m":108,"g":107},"75df31b6":{"m":108,"g":107},"194561f2":{"m":108,"g":107},"5e91fed1":{"m":108,"g":107},"873f384a":{"m":108,"g":107},"b01eeb80":{"m":108,"g":107},"1ea94d3b":{"m":108,"g":107},"354ac435":{"m":108,"g":107},"d98a4913":{"m":108,"g":107},"08f8f490":{"m":108,"g":107},"d4bf5a85":{"m":108,"g":107},"7cb20754":{"m":108,"g":107},"6d0646da":{"m":108,"g":107},"02bc1c7d":{"m":108,"g":107},"fc8c8e50":{"m":108,"g":107},"9bd4872a":{"m":108,"g":107},"2fa0462c":{"m":108,"g":107},"915140fd":{"m":108,"g":107},"36fc9260":{"m":108,"g":107},"fee0ab0f":{"m":108,"g":107},"f57d2dc1":{"m":108,"g":107},"f2d68ded":{"m":108,"g":107},"3b87a9e8":{"m":108,"g":107},"f024795e":{"m":108,"g":107},"b102353f":{"m":108,"g":107},"7a27e798":{"m":108,"g":107},"76ba5bbe":{"m":108,"g":107},"ed6f7597":{"m":108,"g":107},"e67276ec":{"m":108,"g":107},"0242bb9c":{"m":108,"g":107},"760286e3":{"m":108,"g":107},"3435a24e":{"m":108,"g":107},"00da9065":{"m":108,"g":107},"e0ab167d":{"m":109,"g":115},"c807cd7c":{"m":109,"g":115},"327f7b7c":{"m":109,"g":115},"80425e59":{"m":109,"g":115},"af9d4eb0":{"m":109,"g":115},"fb107cfd":{"m":109,"g":115},"e3e97a12":{"m":110,"g":115},"05106867":{"m":110,"g":115},"9dcdf5da":{"m":110,"g":115},"f8b757bc":{"m":110,"g":115},"ebd9dbe7":{"m":110,"g":115},"938e986e":{"m":110,"g":115},"17d5eda8":{"m":110,"g":115},"71a7f1d8":{"m":110,"g":115},"433266c1":{"m":110,"g":115},"fda47926":{"m":110,"g":115},"a0b22f2f":{"m":110,"g":115},"b5c6529e":{"m":110,"g":115},"ca4b86c5":{"m":110,"g":115},"dd6ec029":{"m":110,"g":115},"bf863e3b":{"m":110,"g":115},"9e169ea8":{"m":110,"g":115},"bc80dc4c":{"m":111,"g":115},"b962a296":{"m":111,"g":115},"aa3eba8e":{"m":111,"g":115},"07ee0ab7":{"m":111,"g":115},"5c06dcb7":{"m":111,"g":115},"6f6beca4":{"m":111,"g":115},"68a54e06":{"m":111,"g":115},"fd18995c":{"m":111,"g":115},"db0831e0":{"m":111,"g":115},"6e4e1c8c":{"m":111,"g":115},"9768c50d":{"m":111,"g":115},"fd71b11b":{"m":111,"g":115},"ae7428a8":{"m":111,"g":115},"a3aee7c3":{"m":111,"g":115},"79e6a8a6":{"m":111,"g":115},"8f7b1c31":{"m":111,"g":115},"b9683be6":{"m":111,"g":115},"a85363c1":{"m":111,"g":115},"b21fdd53":{"m":111,"g":115},"c04c17ed":{"m":111,"g":115},"16a6d21b":{"m":111,"g":115},"a530b3ff":{"m":111,"g":115},"603b3446":{"m":111,"g":115},"b6c14ec0":{"m":111,"g":115},"43de1d73":{"m":111,"g":115},"79ce3688":{"m":111,"g":115},"44ffe2cb":{"m":111,"g":115},"1a0896e9":{"m":111,"g":115},"90313fb0":{"m":111,"g":115},"3578eb1e":{"m":111,"g":115},"0936c766":{"m":111,"g":115},"0ef583b7":{"m":111,"g":115},"f7881a27":{"m":111,"g":115},"fdff3167":{"m":111,"g":115},"cbc0e4d7":{"m":111,"g":115},"4cd08dc5":{"m":111,"g":115},"f92b729d":{"m":111,"g":115},"e2e378ca":{"m":111,"g":115},"dc1decc6":{"m":111,"g":115},"03680f33":{"m":111,"g":115},"d4c5e534":{"m":111,"g":115},"817c62a0":{"m":111,"g":115},"0ff72419":{"m":111,"g":115},"80dc76e1":{"m":111,"g":115},"9b08d975":{"m":111,"g":115},"a0a77d93":{"m":111,"g":115},"24a8cee6":{"m":111,"g":115},"3affa9dc":{"m":111,"g":115},"ea0696b9":{"m":111,"g":115},"3aec3d4f":{"m":111,"g":115},"b0d25e72":{"m":112,"g":115},"a2424068":{"m":112,"g":115},"c5d2b01c":{"m":112,"g":115},"46ccbed2":{"m":112,"g":115},"fe68c148":{"m":112,"g":115},"70c0c1f9":{"m":112,"g":115},"760b788a":{"m":112,"g":115},"1ee11df8":{"m":112,"g":115},"dee197e1":{"m":112,"g":115},"ab795ae8":{"m":112,"g":115},"480d1b8b":{"m":112,"g":115},"6c18ab46":{"m":112,"g":115},"4a0e0be2":{"m":112,"g":115},"64f296f8":{"m":112,"g":115},"956d805d":{"m":112,"g":115},"30c6e1f5":{"m":112,"g":115},"bfe01a5e":{"m":112,"g":115},"3dd6420a":{"m":112,"g":115},"532f998b":{"m":112,"g":115},"de15d140":{"m":112,"g":115},"37367da6":{"m":112,"g":115},"ef959d7b":{"m":112,"g":115},"4aa1e69b":{"m":112,"g":115},"dc491b39":{"m":112,"g":115},"5b64f006":{"m":112,"g":115},"5b7448de":{"m":112,"g":115},"6d55f60e":{"m":112,"g":115},"033b75f5":{"m":112,"g":115},"f3b5db6e":{"m":112,"g":115},"2286e85e":{"m":112,"g":115},"91b3555d":{"m":112,"g":115},"9e2f7252":{"m":112,"g":115},"21176b00":{"m":112,"g":115},"94100294":{"m":112,"g":115},"cda7e47c":{"m":112,"g":115},"e903f695":{"m":112,"g":115},"27760fc1":{"m":112,"g":115},"0ac809de":{"m":112,"g":115},"4efe2c57":{"m":112,"g":115},"5be8c2f7":{"m":112,"g":115},"737d73ed":{"m":112,"g":115},"ebd0e1c1":{"m":112,"g":115},"a1d03892":{"m":112,"g":115},"dccf52f9":{"m":112,"g":115},"676a7b51":{"m":112,"g":115},"15f99347":{"m":112,"g":115},"bcf1955f":{"m":112,"g":115},"a06bf664":{"m":112,"g":115},"bf72b801":{"m":112,"g":115},"8cbe1538":{"m":112,"g":115},"8471e5e6":{"m":112,"g":115},"4582931a":{"m":112,"g":115},"d352c29a":{"m":112,"g":115},"d3ee7098":{"m":112,"g":115},"71fc7b7f":{"m":112,"g":115},"9ab72f98":{"m":112,"g":115},"f3817cb0":{"m":112,"g":115},"71133a04":{"m":112,"g":115},"2cd94dd0":{"m":112,"g":115},"f5f6b3b4":{"m":112,"g":115},"94fb4e9e":{"m":112,"g":115},"d1d4074c":{"m":112,"g":115},"718f25ae":{"m":112,"g":115},"948b01a0":{"m":112,"g":115},"cdc56ef6":{"m":112,"g":115},"16ff3d4b":{"m":112,"g":115},"83d55ac5":{"m":112,"g":115},"2fe17735":{"m":112,"g":115},"97fff98c":{"m":112,"g":115},"ba066ca0":{"m":112,"g":115},"96784a65":{"m":112,"g":115},"df5407fb":{"m":112,"g":115},"8ad700f7":{"m":112,"g":115},"148022fc":{"m":112,"g":115},"7a40e4f4":{"m":112,"g":115},"19d64f2b":{"m":112,"g":115},"a02071a1":{"m":112,"g":115},"45b3a6a2":{"m":112,"g":115},"9a18aa54":{"m":112,"g":115},"91f0fd95":{"m":112,"g":115},"8085aca7":{"m":112,"g":115},"0096798e":{"m":112,"g":115},"2c2b19b1":{"m":112,"g":115},"72f9fc5f":{"m":112,"g":115},"ec99668a":{"m":112,"g":115},"78f13981":{"m":112,"g":115},"bfd7a18d":{"m":112,"g":115},"5dd8c644":{"m":112,"g":115},"ee21817c":{"m":112,"g":115},"b7d1f17b":{"m":112,"g":115},"c8295d23":{"m":112,"g":115},"b67c277f":{"m":112,"g":115},"8116804e":{"m":112,"g":115},"8c5930f0":{"m":112,"g":115},"3b99f23c":{"m":112,"g":115},"ee0b3c5b":{"m":112,"g":115},"6049ca20":{"m":112,"g":115},"7577f0e4":{"m":112,"g":115},"8cda5a62":{"m":112,"g":115},"400d3b97":{"m":112,"g":115},"37d83c6e":{"m":112,"g":115},"7802586c":{"m":112,"g":115},"bc5fc332":{"m":112,"g":115},"f3440adc":{"m":112,"g":115},"5a7e10fe":{"m":112,"g":115},"33467c05":{"m":112,"g":115},"b0fcbb74":{"m":112,"g":115},"76a2c86b":{"m":112,"g":115},"e719bb0e":{"m":112,"g":115},"06724683":{"m":112,"g":115},"617aa2b2":{"m":112,"g":115},"111b1379":{"m":112,"g":115},"41628dc1":{"m":112,"g":115},"a12061df":{"m":112,"g":115},"85ed8e0a":{"m":112,"g":115},"dd1e2689":{"m":112,"g":115},"9a7ced4e":{"m":112,"g":115},"cb3918a0":{"m":112,"g":115},"f3b67602":{"m":112,"g":115},"9eb50ecc":{"m":112,"g":115},"b3e7a2ce":{"m":112,"g":115},"00974e4f":{"m":112,"g":115},"5f1eb204":{"m":112,"g":115},"039cef76":{"m":112,"g":115},"4c22ebe2":{"m":112,"g":115},"a5a03209":{"m":112,"g":115},"21af5c04":{"m":112,"g":115},"012584ec":{"m":112,"g":115},"90dfe3de":{"m":112,"g":115},"9a719b7a":{"m":112,"g":115},"3fa62da7":{"m":112,"g":115},"dbb1235d":{"m":112,"g":115},"ad26f298":{"m":112,"g":115},"8d114f25":{"m":112,"g":115},"0e78c63c":{"m":112,"g":115},"1a3d6f31":{"m":112,"g":115},"0b8c5721":{"m":112,"g":115},"beac202b":{"m":112,"g":115},"21b9a4b4":{"m":112,"g":115},"db37422c":{"m":112,"g":115},"ab62b135":{"m":112,"g":115},"273b2834":{"m":112,"g":115},"f84db115":{"m":112,"g":115},"efb0de2c":{"m":112,"g":115},"0f6ac5e2":{"m":112,"g":115},"29850900":{"m":112,"g":115},"e678cc71":{"m":112,"g":115},"4efe844a":{"m":112,"g":115},"bde73ee4":{"m":112,"g":115},"4f0e28d7":{"m":112,"g":115},"045ab92d":{"m":112,"g":115},"bd7f8821":{"m":112,"g":115},"5e5c30d9":{"m":112,"g":115},"9f00ec44":{"m":112,"g":115},"8e85ee88":{"m":112,"g":115},"adf73175":{"m":112,"g":115},"13705dae":{"m":112,"g":115},"df97b31f":{"m":112,"g":115},"339f8eef":{"m":112,"g":115},"afd9f2f5":{"m":112,"g":115},"f40038fb":{"m":112,"g":115},"bebd0576":{"m":112,"g":115},"f9836660":{"m":112,"g":115},"8b3b995a":{"m":112,"g":115},"6e95f5e5":{"m":112,"g":115},"0e9387a9":{"m":112,"g":115},"fa9c82d3":{"m":112,"g":115},"918e3d4c":{"m":112,"g":115},"e9697374":{"m":112,"g":115},"93088b69":{"m":112,"g":115},"453511ac":{"m":112,"g":115},"d0730487":{"m":112,"g":115},"b32ab070":{"m":112,"g":115},"75ee0011":{"m":112,"g":115},"ec15c836":{"m":112,"g":115},"106c2b31":{"m":112,"g":115},"c6756949":{"m":112,"g":115},"27e8ffed":{"m":112,"g":115},"4dbb34fe":{"m":112,"g":115},"1e18a341":{"m":112,"g":115},"2c562fd2":{"m":112,"g":115},"b648d862":{"m":112,"g":115},"bbf261ae":{"m":112,"g":115},"4f8a982d":{"m":112,"g":115},"d966b902":{"m":112,"g":115},"de921733":{"m":112,"g":115},"397448eb":{"m":112,"g":115},"66d5d042":{"m":112,"g":115},"73179b76":{"m":112,"g":115},"8cbf71dc":{"m":112,"g":115},"56eb5d0a":{"m":112,"g":115},"4ed9053e":{"m":112,"g":115},"5e19b159":{"m":112,"g":115},"788b19a5":{"m":112,"g":115},"f78b7fd1":{"m":112,"g":115},"b1fb7e45":{"m":112,"g":115},"1b2ff4fb":{"m":112,"g":115},"2c7ca33a":{"m":112,"g":115},"df397a72":{"m":112,"g":115},"5dfcd6c2":{"m":112,"g":115},"0dfd54d1":{"m":112,"g":115},"bcbeed71":{"m":112,"g":115},"cc9a31c6":{"m":112,"g":115},"d631290e":{"m":112,"g":115},"37565b7f":{"m":112,"g":115},"6243c367":{"m":112,"g":115},"60e37f80":{"m":112,"g":115},"369b1433":{"m":112,"g":115},"03dbf1aa":{"m":112,"g":115},"11dcabc5":{"m":112,"g":115},"4d89389c":{"m":112,"g":115},"9491d6e5":{"m":112,"g":115},"f64b8e3e":{"m":112,"g":115},"53976fce":{"m":112,"g":115},"18f91eb6":{"m":112,"g":115},"8766b3ac":{"m":112,"g":115},"1db649ac":{"m":112,"g":115},"a1e5d781":{"m":112,"g":115},"b7361cc4":{"m":112,"g":115},"a96c5b5c":{"m":112,"g":115},"b9eb0d9c":{"m":112,"g":115},"1fbfdebe":{"m":112,"g":115},"a25e8e42":{"m":112,"g":115},"d4a93841":{"m":112,"g":115},"21e1bc47":{"m":112,"g":115},"9a0cac1b":{"m":112,"g":115},"b5245064":{"m":112,"g":115},"9d9fa9a5":{"m":112,"g":115},"58d06fdc":{"m":112,"g":115},"cb9e0e41":{"m":112,"g":115},"9db80253":{"m":112,"g":115},"598c0bc1":{"m":112,"g":115},"b361750a":{"m":112,"g":115},"16e56ea6":{"m":112,"g":115},"349b491c":{"m":112,"g":115},"5f77e129":{"m":112,"g":115},"4750cddf":{"m":112,"g":115},"065e523d":{"m":112,"g":115},"7de2ce45":{"m":112,"g":115},"8c2ffaaf":{"m":112,"g":115},"20445327":{"m":112,"g":115},"6d3c20cf":{"m":112,"g":115},"8b6966d0":{"m":112,"g":115},"a391f73a":{"m":112,"g":115},"25c73959":{"m":112,"g":115},"f05c6873":{"m":112,"g":115},"9a0d0b75":{"m":112,"g":115},"ba861293":{"m":112,"g":115},"c112bcc4":{"m":112,"g":115},"5e194b21":{"m":112,"g":115},"fd5ce576":{"m":112,"g":115},"92d79646":{"m":112,"g":115},"f9076a5a":{"m":112,"g":115},"646076b7":{"m":112,"g":115},"0d040089":{"m":112,"g":115},"05e47872":{"m":112,"g":115},"1e61b496":{"m":112,"g":115},"300676af":{"m":112,"g":115},"7fe89f7c":{"m":112,"g":115},"9970e3bf":{"m":112,"g":115},"70eedb58":{"m":112,"g":115},"9c99949e":{"m":112,"g":115},"c5082f0f":{"m":112,"g":115},"836873b9":{"m":112,"g":115},"8abe8dea":{"m":112,"g":115},"1e85589d":{"m":112,"g":115},"c2a26e72":{"m":112,"g":115},"591e6c59":{"m":112,"g":115},"42f34437":{"m":112,"g":115},"5c34b4f1":{"m":112,"g":115},"ff9b5618":{"m":112,"g":115},"fcd72bd1":{"m":112,"g":115},"3d8fc434":{"m":112,"g":115},"87a0f7d2":{"m":112,"g":115},"839c93bd":{"m":112,"g":115},"f1e9bbaf":{"m":112,"g":115},"3fd1431d":{"m":112,"g":115},"161e9dc5":{"m":112,"g":115},"54e872d3":{"m":112,"g":115},"e5b29bf1":{"m":112,"g":115},"9a7c8842":{"m":112,"g":115},"7a16db9b":{"m":112,"g":115},"09a1df22":{"m":112,"g":115},"4b7034dd":{"m":112,"g":115},"a23c3020":{"m":112,"g":115},"a7d825fc":{"m":112,"g":115},"38cd5fb1":{"m":112,"g":115},"001f5194":{"m":112,"g":115},"5ad296bd":{"m":112,"g":115},"9f81d741":{"m":112,"g":115},"a38c1497":{"m":112,"g":115},"74dd4249":{"m":112,"g":115},"dc20c22f":{"m":112,"g":115},"711390a9":{"m":112,"g":115},"53430588":{"m":112,"g":115},"fce7ae33":{"m":112,"g":115},"6b39f9cf":{"m":112,"g":115},"07c9d8fb":{"m":112,"g":115},"4a4772ae":{"m":112,"g":115},"c3779233":{"m":112,"g":115},"f84b57c8":{"m":112,"g":115},"aee094e4":{"m":112,"g":115},"55349e36":{"m":112,"g":115},"e1f7cf57":{"m":112,"g":115},"2bb9d454":{"m":112,"g":115},"d0934a51":{"m":112,"g":115},"3f2d0cef":{"m":112,"g":115},"8b30bec2":{"m":112,"g":115},"4aeba40d":{"m":112,"g":115},"28684f90":{"m":112,"g":115},"a4a3d823":{"m":113,"g":115},"0b13cbb7":{"m":113,"g":115},"efbc687c":{"m":113,"g":115},"292a867a":{"m":113,"g":115},"8fd41eae":{"m":113,"g":115},"0cd1996e":{"m":113,"g":115},"b6b4b563":{"m":113,"g":115},"f8924ad7":{"m":113,"g":115},"2f80bd9f":{"m":113,"g":115},"366a603e":{"m":113,"g":115},"baee0860":{"m":113,"g":115},"c7a104c1":{"m":113,"g":115},"97d966a7":{"m":113,"g":115},"8e66d87f":{"m":113,"g":115},"a20fc7b7":{"m":113,"g":115},"6b30e097":{"m":113,"g":115},"d645ae90":{"m":113,"g":115},"41763ba0":{"m":113,"g":115},"652c24a6":{"m":113,"g":115},"5e142484":{"m":113,"g":115},"c560410d":{"m":113,"g":115},"590f2da0":{"m":113,"g":115},"148d8d48":{"m":113,"g":115},"1a599509":{"m":113,"g":115},"36a6b8db":{"m":113,"g":115},"e0b2d3ee":{"m":113,"g":115},"4cb5a523":{"m":113,"g":115},"85c1f793":{"m":113,"g":115},"48e9e719":{"m":113,"g":115},"31b49c0b":{"m":113,"g":115},"d736e0b6":{"m":113,"g":115},"ffd03a9b":{"m":113,"g":115},"666da3d5":{"m":113,"g":115},"d01b9214":{"m":113,"g":115},"c70e58e8":{"m":113,"g":115},"c61b9a1d":{"m":113,"g":115},"3c3d6255":{"m":113,"g":115},"546914fa":{"m":113,"g":115},"4726c919":{"m":113,"g":115},"a0010bf4":{"m":113,"g":115},"307fc060":{"m":113,"g":115},"586e81a2":{"m":113,"g":115},"fad7ca73":{"m":113,"g":115},"08af8ffb":{"m":113,"g":115},"2c7f4ca2":{"m":113,"g":115},"03def5e3":{"m":113,"g":115},"6ae3f05b":{"m":113,"g":115},"fdc4e1e5":{"m":113,"g":115},"04b86b3c":{"m":113,"g":115},"d6777a70":{"m":113,"g":115},"8c574902":{"m":113,"g":115},"34151f17":{"m":113,"g":115},"6794d210":{"m":113,"g":115},"1a31229c":{"m":113,"g":115},"de89ef49":{"m":113,"g":115},"b00a0c78":{"m":113,"g":115},"a2faf894":{"m":113,"g":115},"7e61737d":{"m":113,"g":115},"3c699772":{"m":113,"g":115},"e8100774":{"m":113,"g":115},"963175d5":{"m":113,"g":115},"0618ad6d":{"m":113,"g":115},"6a261aac":{"m":113,"g":115},"7ff740a6":{"m":113,"g":115},"bfcd9b24":{"m":113,"g":115},"458611de":{"m":113,"g":115},"3511b370":{"m":113,"g":115},"afcd3e10":{"m":113,"g":115},"12d68183":{"m":113,"g":115},"b65db028":{"m":113,"g":115},"948278f1":{"m":113,"g":115},"7d004799":{"m":113,"g":115},"083629c2":{"m":113,"g":115},"b658be6f":{"m":113,"g":115},"5e786cca":{"m":113,"g":115},"0b9dfba7":{"m":113,"g":115},"6a290034":{"m":113,"g":115},"2ac453b0":{"m":113,"g":115},"f35def86":{"m":113,"g":115},"d61615fe":{"m":113,"g":115},"b1ccaf01":{"m":113,"g":115},"097725bb":{"m":113,"g":115},"44b1fbe2":{"m":113,"g":115},"c0dbbdd1":{"m":113,"g":115},"25e7dbe8":{"m":113,"g":115},"0b2aa8a7":{"m":113,"g":115},"609f65ba":{"m":113,"g":115},"2d62af6b":{"m":113,"g":115},"a28b394f":{"m":113,"g":115},"96fe2d0f":{"m":113,"g":115},"bfa27438":{"m":113,"g":115},"86cb4db0":{"m":113,"g":115},"2e130b76":{"m":113,"g":115},"ac1f2928":{"m":113,"g":115},"195a59fe":{"m":113,"g":115},"47488cc3":{"m":113,"g":115},"61305291":{"m":113,"g":115},"a9ce2bcb":{"m":113,"g":115},"5dddb331":{"m":113,"g":115},"01a26544":{"m":113,"g":115},"73d4a5f8":{"m":113,"g":115},"7fb551a7":{"m":113,"g":115},"1193f131":{"m":113,"g":115},"84a9f5d6":{"m":113,"g":115},"8ce830a8":{"m":113,"g":115},"fb367acf":{"m":113,"g":115},"a6cc86df":{"m":113,"g":115},"229d2b95":{"m":113,"g":115},"9710f718":{"m":113,"g":115},"91847e38":{"m":113,"g":115},"5a290a56":{"m":113,"g":115},"580051c5":{"m":113,"g":115},"1237aa19":{"m":113,"g":115},"59911195":{"m":113,"g":115},"424591d5":{"m":113,"g":115},"d1676cd4":{"m":113,"g":115},"33b3c0f8":{"m":113,"g":115},"e5281f84":{"m":113,"g":115},"d17986f8":{"m":113,"g":115},"8831c55c":{"m":113,"g":115},"2bc61dd1":{"m":113,"g":115},"6535fda1":{"m":113,"g":115},"3713eb61":{"m":113,"g":115},"5937a56d":{"m":113,"g":115},"f065e5be":{"m":113,"g":115},"9de1320b":{"m":113,"g":115},"dda34c2f":{"m":113,"g":115},"4eeaff74":{"m":113,"g":115},"a17e70f5":{"m":113,"g":115},"816b3a43":{"m":113,"g":115},"3a641d90":{"m":113,"g":115},"6f16bf9d":{"m":113,"g":115},"5942fdb4":{"m":113,"g":115},"af4ab656":{"m":113,"g":115},"11965b0d":{"m":113,"g":115},"71959545":{"m":113,"g":115},"24f7cb1e":{"m":113,"g":115},"e05555fa":{"m":113,"g":115},"43fa9f22":{"m":113,"g":115},"e98d9346":{"m":113,"g":115},"0c917410":{"m":113,"g":115},"25728863":{"m":113,"g":115},"dba751a8":{"m":113,"g":115},"2e763398":{"m":113,"g":115},"336e9a60":{"m":113,"g":115},"abb67815":{"m":113,"g":115},"07440f5f":{"m":113,"g":115},"9816989b":{"m":113,"g":115},"42245551":{"m":113,"g":115},"2a9d995c":{"m":113,"g":115},"a9050b5c":{"m":113,"g":115},"66face35":{"m":113,"g":115},"5519766a":{"m":113,"g":115},"72392f29":{"m":113,"g":115},"c1c8dd1d":{"m":113,"g":115},"f6bc3f52":{"m":113,"g":115},"8cc27fdc":{"m":113,"g":115},"9c339d6b":{"m":113,"g":115},"e23e280e":{"m":113,"g":115},"51f7c6bd":{"m":113,"g":115},"62e2e99d":{"m":113,"g":115},"8ebf72fe":{"m":113,"g":115},"82605747":{"m":113,"g":115},"37f3325b":{"m":113,"g":115},"bd95944c":{"m":113,"g":115},"c8a5d12a":{"m":113,"g":115},"2387c22b":{"m":113,"g":115},"592ddf37":{"m":113,"g":115},"0c3db889":{"m":113,"g":115},"2bdaf482":{"m":113,"g":115},"777eb538":{"m":113,"g":115},"05a35266":{"m":113,"g":115},"e56c64bf":{"m":113,"g":115},"fff7fbab":{"m":113,"g":115},"aae7ead2":{"m":113,"g":115},"a7fe6e10":{"m":113,"g":115},"be059b83":{"m":113,"g":115},"5d4fe1ce":{"m":113,"g":115},"1b011e68":{"m":113,"g":115},"5c0efa56":{"m":113,"g":115},"1e57b947":{"m":113,"g":115},"a5095d62":{"m":113,"g":115},"6c2c467d":{"m":113,"g":115},"c3d2ad4e":{"m":113,"g":115},"7ec5b4e8":{"m":113,"g":115},"60885482":{"m":113,"g":115},"172bcf01":{"m":113,"g":115},"37158f20":{"m":113,"g":115},"3e95aa1a":{"m":113,"g":115},"c4197e99":{"m":113,"g":115},"0ac61146":{"m":113,"g":115},"7dcd689b":{"m":113,"g":115},"f7bab41a":{"m":113,"g":115},"f68dd998":{"m":113,"g":115},"35ec2a45":{"m":113,"g":115},"0035f1ce":{"m":113,"g":115},"5e21d6ae":{"m":113,"g":115},"cd4da1f1":{"m":113,"g":115},"91678474":{"m":113,"g":115},"d511b2d9":{"m":113,"g":115},"77830a26":{"m":113,"g":115},"fce17048":{"m":113,"g":115},"3d40794f":{"m":113,"g":115},"c1f39013":{"m":113,"g":115},"3e43eb13":{"m":113,"g":115},"458c0219":{"m":113,"g":115},"a73eb8cd":{"m":113,"g":115},"e7387035":{"m":113,"g":115},"fe531d6f":{"m":113,"g":115},"c4e314f9":{"m":113,"g":115},"7a06ef98":{"m":113,"g":115},"4a87ba21":{"m":113,"g":115},"d7b20dd6":{"m":113,"g":115},"c3faf2d6":{"m":113,"g":115},"9209b209":{"m":113,"g":115},"adba172f":{"m":113,"g":115},"cd641a99":{"m":113,"g":115},"71f24ef8":{"m":113,"g":115},"b1f0fc1c":{"m":113,"g":115},"32d89373":{"m":113,"g":115},"f47a2c67":{"m":113,"g":115},"ee704e62":{"m":113,"g":115},"f4e3ebeb":{"m":113,"g":115},"312bfc4c":{"m":113,"g":115},"e290303e":{"m":113,"g":115},"aab35bcc":{"m":113,"g":115},"42aedb02":{"m":113,"g":115},"984730b7":{"m":113,"g":115},"23632d35":{"m":113,"g":115},"08b8c0c3":{"m":113,"g":115},"d42975c6":{"m":113,"g":115},"adc24a3a":{"m":113,"g":115},"7ff93e61":{"m":113,"g":115},"b24b2e7e":{"m":113,"g":115},"7135db5d":{"m":113,"g":115},"4b5ef300":{"m":113,"g":115},"4f564b9e":{"m":113,"g":115},"98c3b04f":{"m":113,"g":115},"ddab4fc7":{"m":113,"g":115},"d21c3522":{"m":113,"g":115},"4a762041":{"m":113,"g":115},"ea338676":{"m":113,"g":115},"b06db198":{"m":113,"g":115},"8c1ef0f9":{"m":113,"g":115},"f5a2faf2":{"m":113,"g":115},"1c82d9db":{"m":113,"g":115},"9241f4fd":{"m":113,"g":115},"063c3791":{"m":113,"g":115},"632b7d8c":{"m":113,"g":115},"16adf3dc":{"m":113,"g":115},"c3a1d775":{"m":113,"g":115},"89971c4c":{"m":113,"g":115},"113f8f65":{"m":113,"g":115},"e22f3a5e":{"m":113,"g":115},"095093ee":{"m":113,"g":115},"d27a6f70":{"m":113,"g":115},"0753ef83":{"m":113,"g":115},"662393f2":{"m":113,"g":115},"b1bb8e74":{"m":113,"g":115},"38c00ed7":{"m":113,"g":115},"d4041a5e":{"m":113,"g":115},"2f555c4c":{"m":113,"g":115},"e53df7c0":{"m":113,"g":115},"9c53dad8":{"m":113,"g":115},"7ca1bea6":{"m":113,"g":115},"97c38239":{"m":113,"g":115},"60dbbd08":{"m":113,"g":115},"aa1c5cf5":{"m":113,"g":115},"592caab6":{"m":113,"g":115},"2101d93b":{"m":113,"g":115},"70e4b218":{"m":113,"g":115},"944f1ea0":{"m":113,"g":115},"9d7e82a0":{"m":113,"g":115},"f0580551":{"m":113,"g":115},"635ccda6":{"m":113,"g":115},"1c3dbad8":{"m":113,"g":115},"e2ac7888":{"m":113,"g":115},"86527a47":{"m":113,"g":115},"134b4f7e":{"m":113,"g":115},"f67d1f45":{"m":113,"g":115},"0f04a5f4":{"m":113,"g":115},"2f18602f":{"m":113,"g":115},"56321e9f":{"m":113,"g":115},"12d6cf18":{"m":113,"g":115},"fc3e5420":{"m":113,"g":115},"08ecd0aa":{"m":113,"g":115},"720c1c8c":{"m":113,"g":115},"d403c143":{"m":113,"g":115},"cba0d8c3":{"m":113,"g":115},"f1d78923":{"m":113,"g":115},"7c876de7":{"m":113,"g":115},"ba94b829":{"m":113,"g":115},"2b7417bf":{"m":113,"g":115},"f1116495":{"m":113,"g":115},"bd7eb020":{"m":113,"g":115},"74cd6e39":{"m":113,"g":115},"b17e67df":{"m":113,"g":115},"8ecef73f":{"m":113,"g":115},"1d1ce624":{"m":113,"g":115},"60e2a7ce":{"m":113,"g":115},"d88ef4a3":{"m":113,"g":115},"6f993e8b":{"m":113,"g":115},"03ce92e5":{"m":113,"g":115},"00eb5eb7":{"m":113,"g":115},"dab4663b":{"m":113,"g":115},"610a6d6e":{"m":113,"g":115},"36efd5be":{"m":113,"g":115},"68cdc189":{"m":113,"g":115},"7f399e4b":{"m":113,"g":115},"873d858b":{"m":113,"g":115},"3fa3c22a":{"m":113,"g":115},"4f2055ad":{"m":113,"g":115},"616a3e20":{"m":113,"g":115},"ac2a723b":{"m":113,"g":115},"56b991b1":{"m":113,"g":115},"780d6a22":{"m":113,"g":115},"8b713c72":{"m":113,"g":115},"5bfafdfc":{"m":113,"g":115},"8c52de6f":{"m":113,"g":115},"c1815a99":{"m":113,"g":115},"4e6c4923":{"m":113,"g":115},"b91cb67e":{"m":113,"g":115},"e7bc6003":{"m":113,"g":115},"2a2ff9a8":{"m":113,"g":115},"5291f32d":{"m":113,"g":115},"67073dde":{"m":113,"g":115},"9a5c42f9":{"m":113,"g":115},"388c05d5":{"m":113,"g":115},"fc809665":{"m":113,"g":115},"6fd4816d":{"m":113,"g":115},"1344ebc8":{"m":113,"g":115},"e07b21ce":{"m":113,"g":115},"52f248cd":{"m":113,"g":115},"93f75778":{"m":113,"g":115},"4039c626":{"m":113,"g":115},"db71c38f":{"m":113,"g":115},"7a68b422":{"m":113,"g":115},"60fc5b51":{"m":113,"g":115},"a13dd1e4":{"m":113,"g":115},"d500eb91":{"m":113,"g":115},"1ccd59c7":{"m":113,"g":115},"c32fb7a2":{"m":113,"g":115},"1ba137e9":{"m":113,"g":115},"de28f8e7":{"m":113,"g":115},"56405076":{"m":113,"g":115},"b73ac629":{"m":113,"g":115},"77098aea":{"m":113,"g":115},"5ccf0b03":{"m":113,"g":115},"a77564e0":{"m":113,"g":115},"4f9e71df":{"m":113,"g":115},"541551ce":{"m":113,"g":115},"124097fc":{"m":113,"g":115},"e1d45bc2":{"m":113,"g":115},"14fdd527":{"m":113,"g":115},"f949ad57":{"m":113,"g":115},"c49484a6":{"m":113,"g":115},"a2f7218a":{"m":113,"g":115},"311de47b":{"m":113,"g":115},"373080ea":{"m":113,"g":115},"7f028b07":{"m":113,"g":115},"0abb41c7":{"m":113,"g":115},"925dbb32":{"m":113,"g":115},"8df7353a":{"m":113,"g":115},"ae4be601":{"m":113,"g":115},"9b876889":{"m":113,"g":115},"c0c6f543":{"m":113,"g":115},"edd6a07b":{"m":113,"g":115},"b6dd4bcb":{"m":113,"g":115},"b2435be6":{"m":113,"g":115},"5fe39e85":{"m":113,"g":115},"fa5d0bf6":{"m":113,"g":115},"16e93359":{"m":113,"g":115},"f1c692f6":{"m":113,"g":115},"80572c83":{"m":113,"g":115},"4bb08f6e":{"m":113,"g":115},"ec272dda":{"m":113,"g":115},"a220c14f":{"m":113,"g":115},"35ef3f29":{"m":113,"g":115},"31fb19a0":{"m":113,"g":115},"3f41b48c":{"m":113,"g":115},"2689f0bf":{"m":113,"g":115},"52074240":{"m":113,"g":115},"c3c26f76":{"m":113,"g":115},"2cf811a9":{"m":113,"g":115},"3b25dc12":{"m":113,"g":115},"5c08d7d2":{"m":113,"g":115},"a45d9a4e":{"m":113,"g":115},"28c79dc8":{"m":113,"g":115},"1fcccda4":{"m":113,"g":115},"79acec4f":{"m":113,"g":115},"b1721edb":{"m":113,"g":115},"57234d0c":{"m":113,"g":115},"b93acd70":{"m":113,"g":115},"86a32bb5":{"m":113,"g":115},"5afd0365":{"m":113,"g":115},"059c13de":{"m":113,"g":115},"50dc0c1e":{"m":113,"g":115},"76becc1d":{"m":113,"g":115},"2a37b24d":{"m":113,"g":115},"f73aae0b":{"m":113,"g":115},"69b35793":{"m":113,"g":115},"957482c8":{"m":113,"g":115},"3795b6a4":{"m":113,"g":115},"7eccbe99":{"m":113,"g":115},"0549f21c":{"m":113,"g":115},"b354e3c9":{"m":113,"g":115},"65e6f48c":{"m":113,"g":115},"0ec580a8":{"m":113,"g":115},"0b14159f":{"m":113,"g":115},"1489cd6c":{"m":113,"g":115},"fc2c3a3d":{"m":113,"g":115},"8f6a1758":{"m":113,"g":115},"01018138":{"m":113,"g":115},"4844fac9":{"m":113,"g":115},"b7d385e8":{"m":113,"g":115},"305c9e8c":{"m":113,"g":115},"ca63f075":{"m":113,"g":115},"f9ee6ae1":{"m":113,"g":115},"dcee42c2":{"m":113,"g":115},"258d02c8":{"m":113,"g":115},"60d7beda":{"m":113,"g":115},"2f8ba6fe":{"m":113,"g":115},"7ce6c10e":{"m":113,"g":115},"55025b92":{"m":113,"g":115},"4c21b090":{"m":113,"g":115},"165abeeb":{"m":113,"g":115},"21ca4c3a":{"m":113,"g":115},"e3cf812f":{"m":113,"g":115},"4da55336":{"m":113,"g":115},"ac964d2e":{"m":113,"g":115},"fa46e2bd":{"m":113,"g":115},"b047b553":{"m":113,"g":115},"a0f844ed":{"m":113,"g":115},"2df532ef":{"m":113,"g":115},"abea9250":{"m":113,"g":115},"b3c97762":{"m":113,"g":115},"b8347b40":{"m":113,"g":115},"72dfa96a":{"m":113,"g":115},"05b01ef4":{"m":113,"g":115},"55a6e644":{"m":113,"g":115},"6897e06b":{"m":113,"g":115},"a360511d":{"m":113,"g":115},"94d0f656":{"m":113,"g":115},"eca59f96":{"m":113,"g":115},"97528610":{"m":113,"g":115},"297d3745":{"m":113,"g":115},"31e9d3a5":{"m":113,"g":115},"6f4676ef":{"m":113,"g":115},"c9ec4cae":{"m":113,"g":115},"99757cc3":{"m":113,"g":115},"cdddab05":{"m":113,"g":115},"7c5a0a1b":{"m":113,"g":115},"49f169d5":{"m":113,"g":115},"7fce2fd9":{"m":113,"g":115},"16cd550c":{"m":113,"g":115},"d5e2a374":{"m":113,"g":115},"366043db":{"m":113,"g":115},"2f173ea0":{"m":113,"g":115},"321fecab":{"m":113,"g":115},"9d775b1a":{"m":113,"g":115},"78b7465c":{"m":113,"g":115},"07bcad7f":{"m":113,"g":115},"98adac8e":{"m":113,"g":115},"cef11e9a":{"m":113,"g":115},"2269cf1e":{"m":113,"g":115},"151e287d":{"m":113,"g":115},"8c86595c":{"m":113,"g":115},"4634fd59":{"m":113,"g":115},"efedbe6c":{"m":113,"g":115},"3a77c80b":{"m":113,"g":115},"36acd2ff":{"m":113,"g":115},"fe6cdf89":{"m":113,"g":115},"30d20ce8":{"m":113,"g":115},"1b1701f1":{"m":113,"g":115},"6d403089":{"m":113,"g":115},"24dc2bee":{"m":113,"g":115},"fac07c9b":{"m":113,"g":115},"b3839a7f":{"m":113,"g":115},"4aa39d72":{"m":113,"g":115},"b4c2c421":{"m":113,"g":115},"53ca1552":{"m":113,"g":115},"a23bdeaf":{"m":113,"g":115},"27778010":{"m":113,"g":115},"46d8fb1c":{"m":113,"g":115},"c7e85f53":{"m":113,"g":115},"3df05f4d":{"m":113,"g":115},"7b141f81":{"m":113,"g":115},"7bc5fb0d":{"m":113,"g":115},"144ee5f3":{"m":113,"g":115},"758b887a":{"m":114,"g":115},"eb7d9261":{"m":114,"g":115},"44cb0607":{"m":114,"g":115},"88bb627d":{"m":114,"g":115},"b520958e":{"m":114,"g":115},"fa7e2c30":{"m":114,"g":115},"8f2cd177":{"m":114,"g":115},"ab926dd6":{"m":114,"g":115},"a4b424c6":{"m":114,"g":115},"a0557642":{"m":114,"g":115},"84768d10":{"m":114,"g":115},"368fd206":{"m":114,"g":115},"53bd00d9":{"m":114,"g":115},"e22b13c5":{"m":114,"g":115},"a3c2ea44":{"m":114,"g":115},"fccac7d1":{"m":114,"g":115},"7ac6b900":{"m":114,"g":115},"a1080b72":{"m":114,"g":115},"a65ca739":{"m":114,"g":115},"677aa0e2":{"m":114,"g":115},"01c9ee1a":{"m":114,"g":115},"d6837aea":{"m":114,"g":115},"c882b5ae":{"m":114,"g":115},"e3bb7f5a":{"m":114,"g":115},"92473e2e":{"m":114,"g":115},"6c0bb327":{"m":114,"g":115},"0a7c4bde":{"m":114,"g":115},"edefab0c":{"m":114,"g":115},"97cd38e5":{"m":114,"g":115},"3c06b673":{"m":114,"g":115},"7c3f07db":{"m":114,"g":115},"edd86b88":{"m":114,"g":115},"4b4dc132":{"m":114,"g":115},"5a9170d9":{"m":114,"g":115},"c4d77774":{"m":114,"g":115},"832c84fb":{"m":114,"g":115},"64d1505c":{"m":114,"g":115},"f3764c26":{"m":114,"g":115},"7ba3de0e":{"m":114,"g":115},"fde9b963":{"m":114,"g":115},"f094e0a4":{"m":114,"g":115},"4ed67c27":{"m":114,"g":115},"cd4b39a9":{"m":114,"g":115},"420c99ac":{"m":114,"g":115},"e3c7f091":{"m":114,"g":115},"6f1e03a4":{"m":114,"g":115},"f4affd4d":{"m":114,"g":115},"df08bf9b":{"m":114,"g":115},"69efdd27":{"m":114,"g":115},"64582caa":{"m":114,"g":115},"2fcd56ea":{"m":114,"g":115},"0958a397":{"m":114,"g":115},"4f42c8cd":{"m":114,"g":115},"3ddd7dc9":{"m":114,"g":115},"501dfa6b":{"m":114,"g":115},"79d34951":{"m":114,"g":115},"1519a89c":{"m":114,"g":115},"24bc3fb0":{"m":114,"g":115},"8a8a608a":{"m":114,"g":115},"533e58a1":{"m":114,"g":115},"9b4c4497":{"m":114,"g":115},"a578d300":{"m":114,"g":115},"8c967037":{"m":114,"g":115},"fb27d383":{"m":114,"g":115},"fd8a0b29":{"m":114,"g":115},"afc35ccc":{"m":114,"g":115},"a57f0e3d":{"m":114,"g":115},"708f4ff4":{"m":114,"g":115},"e2daeb35":{"m":114,"g":115},"0e7b3530":{"m":114,"g":115},"b07c9c76":{"m":114,"g":115},"748f86f3":{"m":114,"g":115},"73ea484a":{"m":114,"g":115},"4aeb193f":{"m":114,"g":115},"466992b2":{"m":114,"g":115},"155cbb51":{"m":114,"g":115},"eb30b888":{"m":114,"g":115},"5ee777c9":{"m":114,"g":115},"baf277a9":{"m":116,"g":118},"f5d30dae":{"m":116,"g":118},"2479b894":{"m":116,"g":118},"54644572":{"m":116,"g":118},"6c01844f":{"m":116,"g":118},"f226d3da":{"m":116,"g":118},"d2478cd4":{"m":116,"g":118},"30ea4c46":{"m":116,"g":118},"6d036468":{"m":116,"g":118},"8221f9ae":{"m":116,"g":118},"ab9187a2":{"m":116,"g":118},"6b143d62":{"m":116,"g":118},"6bc503af":{"m":116,"g":118},"b2c85669":{"m":116,"g":118},"32803fb2":{"m":116,"g":118},"91fc5bb5":{"m":116,"g":118},"780fbf2f":{"m":116,"g":118},"825432fc":{"m":116,"g":118},"a40229f6":{"m":116,"g":118},"74737b28":{"m":116,"g":115},"40e0082d":{"m":116,"g":115},"e9e120ac":{"m":116,"g":115},"e0c2af2a":{"m":116,"g":115},"1d7f7835":{"m":116,"g":115},"32595146":{"m":116,"g":115},"86373b9e":{"m":116,"g":115},"d314bf60":{"m":116,"g":115},"e28c9e52":{"m":116,"g":115},"b98cf398":{"m":116,"g":115},"27d71045":{"m":116,"g":115},"c224a4c6":{"m":116,"g":115},"49345a68":{"m":116,"g":115},"94d26d85":{"m":116,"g":115},"9e8a15a7":{"m":116,"g":115},"3962e39d":{"m":116,"g":115},"eb8cac6f":{"m":116,"g":115},"5ea96ac7":{"m":116,"g":115},"56222658":{"m":116,"g":115},"dc965db0":{"m":116,"g":115},"817e46f4":{"m":116,"g":115},"5a33c3aa":{"m":116,"g":115},"9767a1e4":{"m":116,"g":115},"1d086539":{"m":116,"g":115},"a04efc49":{"m":116,"g":115},"642fa966":{"m":116,"g":115},"da7fac1b":{"m":116,"g":115},"28ad2297":{"m":116,"g":115},"f7f9f8ec":{"m":116,"g":115},"4b62af92":{"m":116,"g":115},"0b9915c1":{"m":116,"g":115},"27ef1459":{"m":116,"g":115},"e4358a45":{"m":116,"g":115},"ba2ce28f":{"m":116,"g":115},"98923880":{"m":116,"g":115},"f792e3c5":{"m":116,"g":115},"28f80b12":{"m":116,"g":115},"88a6f9da":{"m":116,"g":115},"cb8ed2c0":{"m":116,"g":115},"38473363":{"m":116,"g":115},"aaf7af1b":{"m":116,"g":115},"932e2637":{"m":116,"g":115},"43f80884":{"m":116,"g":115},"60b05032":{"m":116,"g":115},"dc48c4c0":{"m":116,"g":115},"6dc9ca8c":{"m":116,"g":115},"887c2b45":{"m":116,"g":115},"065ce815":{"m":116,"g":115},"8e51049f":{"m":116,"g":115},"cb8f3d90":{"m":116,"g":115},"4b694e7d":{"m":116,"g":115},"9f1f699a":{"m":116,"g":115},"c9cff2b9":{"m":116,"g":115},"b6fb5d76":{"m":116,"g":115},"f4aa7880":{"m":116,"g":115},"5e3f7e7f":{"m":116,"g":115},"728af887":{"m":116,"g":115},"7b59b0b8":{"m":116,"g":115},"acc2327b":{"m":116,"g":115},"bfadb5ea":{"m":116,"g":115},"9cc1e065":{"m":116,"g":115},"b8c430f1":{"m":116,"g":115},"f35f120d":{"m":116,"g":115},"54a46a26":{"m":116,"g":115},"7c94eaee":{"m":116,"g":115},"13d596c9":{"m":116,"g":115},"c7867b67":{"m":116,"g":115},"516738b0":{"m":116,"g":115},"0b6f535f":{"m":116,"g":115},"c5fe3c0b":{"m":116,"g":115},"318424e2":{"m":116,"g":115},"6806c4e6":{"m":116,"g":115},"0c0779d6":{"m":116,"g":115},"a55cf530":{"m":116,"g":115},"19ba16aa":{"m":116,"g":115},"a2b3d9b9":{"m":116,"g":115},"9a30914e":{"m":116,"g":115},"8e776c78":{"m":116,"g":115},"63e84352":{"m":116,"g":115},"a20e7df8":{"m":116,"g":115},"1bdd0102":{"m":116,"g":115},"6cd29694":{"m":116,"g":115},"2ac46e94":{"m":116,"g":115},"0aa65f94":{"m":116,"g":115},"0ecb4261":{"m":116,"g":115},"05f015f6":{"m":116,"g":115},"1083e7e3":{"m":116,"g":115},"2157d12a":{"m":116,"g":115},"9f2b457c":{"m":116,"g":115},"f5b34a51":{"m":116,"g":115},"5a6ec8f9":{"m":116,"g":115},"6a653bb1":{"m":116,"g":115},"548a57b1":{"m":116,"g":115},"88e73ed0":{"m":116,"g":115},"4b15fa00":{"m":116,"g":115},"f4941906":{"m":116,"g":115},"01e59e82":{"m":116,"g":115},"99a0704a":{"m":116,"g":115},"ec1cd90a":{"m":116,"g":115},"1103dc62":{"m":116,"g":115},"a220536f":{"m":116,"g":115},"7b064f04":{"m":116,"g":115},"43190bec":{"m":116,"g":115},"be740acd":{"m":116,"g":115},"2db2cddd":{"m":116,"g":115},"9b5efe34":{"m":116,"g":115},"4ac8e09d":{"m":116,"g":115},"20a6c0a6":{"m":116,"g":115},"47c606d3":{"m":116,"g":115},"9fcf7306":{"m":116,"g":115},"0a304870":{"m":116,"g":115},"8fdcd98e":{"m":116,"g":115},"b5dcfd41":{"m":116,"g":115},"5061b8fd":{"m":116,"g":115},"c8452551":{"m":116,"g":115},"bf3e7149":{"m":116,"g":115},"f5754d12":{"m":116,"g":115},"739daa63":{"m":116,"g":115},"d957177a":{"m":116,"g":115},"21337b22":{"m":116,"g":115},"129d2992":{"m":116,"g":115},"8b85926a":{"m":116,"g":115},"451d15c4":{"m":116,"g":115},"c80a96da":{"m":116,"g":115},"eae9a9fb":{"m":116,"g":115},"2674c1d2":{"m":116,"g":115},"61055cb3":{"m":116,"g":115},"92777135":{"m":116,"g":115},"c4958331":{"m":116,"g":115},"2eeb2751":{"m":116,"g":115},"b36afed4":{"m":116,"g":115},"9aa4502d":{"m":116,"g":115},"a0835c3a":{"m":116,"g":115},"55b14656":{"m":116,"g":115},"b4408e60":{"m":116,"g":115},"52fcbbb8":{"m":116,"g":115},"af96ca11":{"m":116,"g":115},"9082a7d3":{"m":116,"g":115},"3b9d97f3":{"m":116,"g":115},"a1a20b4c":{"m":116,"g":115},"4299aebd":{"m":116,"g":115},"0babd487":{"m":116,"g":115},"f19613e6":{"m":116,"g":115},"8df49455":{"m":116,"g":115},"ee3bd8a1":{"m":116,"g":115},"d8467db7":{"m":116,"g":115},"b5044fbf":{"m":116,"g":115},"70fbb3ad":{"m":116,"g":115},"9a7e7a65":{"m":116,"g":115},"0fe87213":{"m":116,"g":115},"1f106ee3":{"m":116,"g":115},"9b8ebb27":{"m":116,"g":115},"85ebeecf":{"m":117,"g":118},"0dd6cf16":{"m":117,"g":118},"0975ba99":{"m":117,"g":118},"1de3924b":{"m":117,"g":118},"3cceaa38":{"m":117,"g":118},"b0d20cde":{"m":117,"g":118},"cbac4997":{"m":117,"g":118},"476c67d7":{"m":117,"g":118},"3289da5b":{"m":117,"g":118},"868403f6":{"m":117,"g":118},"97d857c0":{"m":117,"g":118},"52a54a26":{"m":117,"g":118},"cd7e1bd5":{"m":117,"g":118},"729b7edf":{"m":117,"g":118},"4c03dbaa":{"m":117,"g":118},"1053e1be":{"m":119,"g":121},"9a71500c":{"m":119,"g":121},"6d6e24bc":{"m":119,"g":121},"2c057fbf":{"m":119,"g":121},"dbd9435d":{"m":119,"g":121},"8ae9d4bb":{"m":119,"g":121},"1c304aa9":{"m":119,"g":121},"770529a7":{"m":119,"g":121},"39c237f0":{"m":119,"g":121},"28b8a406":{"m":119,"g":121},"8bd26dd4":{"m":119,"g":121},"ab07cd3e":{"m":119,"g":121},"a9849683":{"m":119,"g":121},"a4b637d8":{"m":119,"g":121},"96a5e4dd":{"m":119,"g":121},"b0b4f716":{"m":119,"g":121},"6c18addb":{"m":119,"g":121},"32852fe9":{"m":119,"g":121},"53c2934d":{"m":119,"g":121},"e321c971":{"m":119,"g":121},"d6fee73d":{"m":119,"g":121},"36a4cad7":{"m":119,"g":121},"65d376b4":{"m":119,"g":121},"c23eda85":{"m":119,"g":121},"138ff231":{"m":119,"g":121},"13fb8b54":{"m":119,"g":121},"81fd2b0e":{"m":119,"g":121},"007b849b":{"m":119,"g":121},"8612811d":{"m":119,"g":121},"e7aa4664":{"m":119,"g":121},"4d4feccb":{"m":119,"g":121},"99c92ff2":{"m":119,"g":121},"6ade6a02":{"m":119,"g":121},"983ef22c":{"m":119,"g":121},"164302c7":{"m":119,"g":121},"5dccf697":{"m":119,"g":121},"eec9e471":{"m":119,"g":121},"6d535b71":{"m":119,"g":121},"fdcb1d13":{"m":119,"g":121},"d7e834d6":{"m":119,"g":121},"200a3c0b":{"m":119,"g":121},"77258ce0":{"m":119,"g":121},"1d097aac":{"m":119,"g":121},"7fceeef5":{"m":119,"g":121},"88568c01":{"m":119,"g":121},"904655c5":{"m":119,"g":121},"e028af69":{"m":119,"g":121},"80b2b320":{"m":119,"g":121},"4b65ed42":{"m":119,"g":121},"23afdfd1":{"m":119,"g":121},"9d61205d":{"m":119,"g":121},"590bc4b7":{"m":119,"g":121},"63cfe1b0":{"m":119,"g":121},"70f6309c":{"m":119,"g":121},"70416001":{"m":119,"g":121},"87a92e45":{"m":119,"g":121},"c461e771":{"m":119,"g":121},"fde2decf":{"m":119,"g":121},"9792b9d7":{"m":119,"g":121},"ef4a8097":{"m":119,"g":121},"ebff4ee6":{"m":119,"g":121},"2b1da821":{"m":119,"g":121},"97710ccd":{"m":119,"g":121},"f3cd5d25":{"m":119,"g":121},"c61b0b29":{"m":119,"g":121},"e8640ee9":{"m":119,"g":121},"d0a64c7e":{"m":119,"g":121},"05d3667a":{"m":119,"g":121},"260fe755":{"m":119,"g":121},"dbb16bed":{"m":119,"g":121},"c1e16003":{"m":119,"g":121},"852c0578":{"m":119,"g":121},"7e6191c0":{"m":119,"g":121},"6f9b66bd":{"m":119,"g":121},"8a801ee3":{"m":119,"g":118},"d9a20fd2":{"m":119,"g":118},"b113c72e":{"m":119,"g":118},"fb6cc7b0":{"m":119,"g":118},"8374a96e":{"m":119,"g":118},"74de76c6":{"m":119,"g":118},"9c0b1eb5":{"m":119,"g":118},"01f14a7a":{"m":119,"g":118},"11110303":{"m":119,"g":118},"28ddfb37":{"m":119,"g":118},"e69094df":{"m":119,"g":118},"43ad0590":{"m":119,"g":118},"b4948512":{"m":119,"g":118},"ddcba74b":{"m":119,"g":118},"0917c5da":{"m":119,"g":118},"184a4df6":{"m":119,"g":118},"f7b1d8c5":{"m":119,"g":118},"bfc3b3f7":{"m":119,"g":118},"da5bde4d":{"m":119,"g":118},"276e7b3e":{"m":119,"g":118},"296f6892":{"m":119,"g":118},"9edb7b51":{"m":119,"g":118},"e53bf442":{"m":119,"g":118},"d383e661":{"m":119,"g":118},"984fbeb1":{"m":119,"g":118},"a2ba0bc3":{"m":119,"g":118},"6d2d0ce2":{"m":119,"g":118},"271d3d0d":{"m":119,"g":118},"c4e81e64":{"m":119,"g":118},"c726d44c":{"m":119,"g":118},"283c8ba0":{"m":119,"g":118},"cae39565":{"m":119,"g":118},"27a223ab":{"m":119,"g":118},"53529f46":{"m":119,"g":118},"24ed3f32":{"m":119,"g":118},"44f0ece9":{"m":119,"g":118},"be0058bc":{"m":119,"g":118},"9e3be1fa":{"m":119,"g":118},"a8ba3279":{"m":119,"g":118},"3b80232d":{"m":119,"g":118},"252dc4e1":{"m":119,"g":118},"cbb5fc2e":{"m":119,"g":118},"53fb229f":{"m":119,"g":118},"4fff1ec1":{"m":119,"g":118},"7a020e0f":{"m":119,"g":118},"48738af7":{"m":119,"g":118},"efa47334":{"m":119,"g":118},"d658f049":{"m":119,"g":118},"57e25de7":{"m":119,"g":118},"12eb02e9":{"m":119,"g":118},"002d0373":{"m":119,"g":118},"a27825ae":{"m":119,"g":118},"ce399e15":{"m":119,"g":118},"ea6275df":{"m":119,"g":118},"eb7318f1":{"m":119,"g":118},"6058fb52":{"m":119,"g":118},"80407b04":{"m":119,"g":118},"b288f4f4":{"m":119,"g":118},"6d6ea5af":{"m":119,"g":118},"1dacedd2":{"m":119,"g":118},"b5e14b2b":{"m":119,"g":118},"d513ee93":{"m":119,"g":118},"a7ae61ed":{"m":119,"g":118},"fda0cb2a":{"m":119,"g":118},"ebda73dc":{"m":119,"g":118},"f4f8a1b4":{"m":119,"g":118},"c44e985d":{"m":119,"g":118},"f9a7d9b3":{"m":119,"g":118},"a93f10a7":{"m":119,"g":118},"585e1223":{"m":119,"g":118},"a7043c6f":{"m":119,"g":118},"67e34c56":{"m":119,"g":118},"1d726528":{"m":119,"g":118},"f4488e9d":{"m":119,"g":118},"e68a2b5b":{"m":119,"g":118},"31b9f19e":{"m":119,"g":118},"547003bd":{"m":119,"g":118},"f7ab9554":{"m":119,"g":118},"dbbd4e18":{"m":119,"g":118},"ca240eef":{"m":119,"g":118},"6c7c92eb":{"m":119,"g":118},"5b214b50":{"m":119,"g":118},"13219e1e":{"m":119,"g":118},"33e9bbec":{"m":119,"g":118},"dcb8f090":{"m":119,"g":118},"9eefe2c0":{"m":119,"g":118},"69fe3c97":{"m":119,"g":118},"8af84912":{"m":119,"g":118},"505329ca":{"m":119,"g":118},"8a382fd3":{"m":119,"g":118},"62797440":{"m":119,"g":118},"2614adf9":{"m":119,"g":118},"fdd7c69d":{"m":119,"g":118},"b9a54e09":{"m":119,"g":118},"20b8d230":{"m":119,"g":118},"d1984e21":{"m":119,"g":118},"b79f75fd":{"m":119,"g":118},"8fcc69e7":{"m":119,"g":118},"f440baa1":{"m":119,"g":118},"2bc3fcd4":{"m":119,"g":118},"a5978a20":{"m":119,"g":118},"e483c1ea":{"m":119,"g":118},"da681f35":{"m":119,"g":118},"9b0f725b":{"m":119,"g":118},"cde5a6e3":{"m":119,"g":118},"3e4c7da2":{"m":119,"g":118},"d88ac9bc":{"m":119,"g":118},"ce11dd82":{"m":119,"g":118},"9e87b60f":{"m":119,"g":118},"7780230a":{"m":119,"g":118},"dc01313d":{"m":119,"g":118},"7a7f99be":{"m":119,"g":118},"fd389df9":{"m":119,"g":118},"b0d1d717":{"m":119,"g":118},"c7962868":{"m":119,"g":118},"4f24ab17":{"m":119,"g":118},"64affab4":{"m":119,"g":118},"4c9bcb9d":{"m":119,"g":118},"86b04d25":{"m":119,"g":118},"55d75e11":{"m":120,"g":121},"3f4cc0af":{"m":120,"g":121},"cadfae66":{"m":120,"g":121},"da1766e4":{"m":120,"g":121},"a124b517":{"m":120,"g":121},"d05a968b":{"m":120,"g":121},"94aad0de":{"m":120,"g":121},"7ebc28f5":{"m":120,"g":121},"b89111d6":{"m":120,"g":121},"0b3b3e9a":{"m":120,"g":121},"a1d5bc4c":{"m":120,"g":121},"a8023891":{"m":120,"g":121},"0103f374":{"m":120,"g":121},"96a5a949":{"m":120,"g":121},"ea385ae8":{"m":120,"g":121},"9e949e58":{"m":120,"g":121},"6dbb569b":{"m":120,"g":121},"5994e6c3":{"m":120,"g":121},"3e6281d0":{"m":120,"g":121},"6371f7af":{"m":120,"g":121},"8491c794":{"m":120,"g":121},"bda3758f":{"m":120,"g":121},"7b36c47b":{"m":120,"g":121},"773d89da":{"m":120,"g":121},"03e7d949":{"m":120,"g":121},"ff604064":{"m":120,"g":121},"212f5e48":{"m":120,"g":121},"fe527812":{"m":120,"g":121},"97828878":{"m":120,"g":121},"c001deba":{"m":120,"g":121},"b4d2da10":{"m":120,"g":121},"b72f9f08":{"m":120,"g":121},"8e70064c":{"m":120,"g":121},"d98b81e2":{"m":120,"g":121},"bcecf27e":{"m":120,"g":121},"4b0ac1d5":{"m":120,"g":121},"8e987fa2":{"m":120,"g":121},"9e656dd3":{"m":120,"g":121},"3862661c":{"m":120,"g":121},"c8492978":{"m":120,"g":121},"428710c2":{"m":120,"g":121},"d9b31011":{"m":120,"g":121},"d0cff78f":{"m":120,"g":121},"e8b71445":{"m":120,"g":121},"4caca1ba":{"m":120,"g":121},"ceb105a7":{"m":120,"g":121},"22cbc9c0":{"m":120,"g":121},"3865afc5":{"m":120,"g":121},"4ea42f7c":{"m":120,"g":121},"ea13cb14":{"m":120,"g":121},"a04212f1":{"m":120,"g":121},"ce869793":{"m":120,"g":121},"22f55e1b":{"m":120,"g":121},"89824189":{"m":120,"g":121},"433c622e":{"m":120,"g":121},"20bd2271":{"m":120,"g":121},"64994980":{"m":120,"g":121},"729b2429":{"m":120,"g":121},"d7056c52":{"m":120,"g":121},"13bf565d":{"m":120,"g":121},"e51046be":{"m":120,"g":121},"4eeeae1e":{"m":120,"g":121},"f4b78d13":{"m":120,"g":121},"4463e90d":{"m":120,"g":121},"229f236d":{"m":120,"g":121},"5983e5bd":{"m":120,"g":121},"4b046a72":{"m":120,"g":121},"770d6312":{"m":120,"g":121},"14203432":{"m":120,"g":121},"d7f0d88f":{"m":120,"g":121},"fc86b18b":{"m":120,"g":121},"0bfa394a":{"m":120,"g":121},"e04340bf":{"m":120,"g":121},"84701338":{"m":120,"g":121},"93ef9a09":{"m":120,"g":121},"b04cd3d4":{"m":120,"g":121},"7ef5d8af":{"m":120,"g":121},"71d41212":{"m":120,"g":121},"b9fb74f3":{"m":120,"g":121},"e15b63a1":{"m":120,"g":121},"4060ed37":{"m":120,"g":121},"2342605e":{"m":120,"g":121},"dbf17a83":{"m":120,"g":121},"0f0c430e":{"m":120,"g":121},"8e797a47":{"m":120,"g":121},"aa3003f1":{"m":120,"g":121},"4793ec7d":{"m":120,"g":121},"92009bd2":{"m":120,"g":121},"4ef981e2":{"m":120,"g":121},"69ed8b67":{"m":120,"g":121},"1801cd19":{"m":120,"g":121},"ffc722a6":{"m":120,"g":121},"49afb3d9":{"m":120,"g":121},"f80371ff":{"m":120,"g":121},"62eff37b":{"m":120,"g":121},"47e12e08":{"m":120,"g":121},"823b4429":{"m":120,"g":121},"14a4d80e":{"m":120,"g":121},"41c10e67":{"m":122,"g":127},"0bfe1d14":{"m":122,"g":127},"5f98b7fe":{"m":122,"g":127},"a4bf5c6a":{"m":122,"g":127},"30ad1070":{"m":122,"g":127},"a80bcb5a":{"m":122,"g":127},"f7f9e41b":{"m":122,"g":127},"263eab9f":{"m":122,"g":127},"25257d8e":{"m":122,"g":127},"cf0c2415":{"m":122,"g":127},"5538e05c":{"m":122,"g":127},"c30ebb93":{"m":122,"g":127},"41efcaeb":{"m":122,"g":127},"70562969":{"m":122,"g":127},"b57dc169":{"m":122,"g":127},"0095e018":{"m":122,"g":127},"68486481":{"m":122,"g":127},"410225b7":{"m":122,"g":127},"2c9aebea":{"m":122,"g":127},"bc741073":{"m":122,"g":127},"2f6af1a3":{"m":122,"g":127},"50b6842b":{"m":122,"g":127},"2d5605e8":{"m":122,"g":127},"300b4c21":{"m":122,"g":127},"c0652d90":{"m":122,"g":127},"2e48584b":{"m":122,"g":127},"57cc5385":{"m":122,"g":127},"5cc0d25a":{"m":122,"g":127},"a076ec1a":{"m":122,"g":127},"72b5f3d0":{"m":122,"g":127},"2f766f38":{"m":122,"g":127},"069e490b":{"m":122,"g":127},"ab95d35f":{"m":122,"g":127},"34c286b8":{"m":122,"g":127},"9416ee60":{"m":122,"g":127},"d4a09ec9":{"m":122,"g":127},"662725b9":{"m":122,"g":127},"82cfcd3b":{"m":122,"g":127},"6c1a3f0c":{"m":122,"g":127},"62377548":{"m":122,"g":121},"96ac24c0":{"m":122,"g":121},"c0d02cf4":{"m":122,"g":121},"7d121448":{"m":122,"g":121},"6a63a985":{"m":122,"g":121},"4d2f17bd":{"m":122,"g":121},"7cd716f7":{"m":122,"g":121},"b7fdde4b":{"m":122,"g":121},"69bf8011":{"m":122,"g":121},"d8fcbaa3":{"m":122,"g":121},"2cf3d0f8":{"m":122,"g":121},"1ed1abfd":{"m":122,"g":121},"ecb9fa14":{"m":122,"g":121},"700daa34":{"m":122,"g":121},"39cee0fe":{"m":122,"g":121},"04e5b6fa":{"m":122,"g":121},"ce6b17c0":{"m":122,"g":121},"cafebef1":{"m":122,"g":121},"73dfd2df":{"m":122,"g":121},"df5192cf":{"m":122,"g":121},"78c43d88":{"m":122,"g":121},"7e28c67d":{"m":122,"g":121},"3edba9bc":{"m":122,"g":121},"e5ec9764":{"m":122,"g":121},"621dfb88":{"m":122,"g":121},"fb52d35f":{"m":122,"g":121},"2b71531a":{"m":122,"g":121},"25c50498":{"m":122,"g":121},"8e2ac2e6":{"m":122,"g":121},"17a57fd8":{"m":122,"g":121},"32438eba":{"m":122,"g":121},"fed02a49":{"m":122,"g":121},"03b3e89a":{"m":122,"g":121},"9ff9fa7f":{"m":122,"g":121},"7ed8ba05":{"m":122,"g":121},"df08f346":{"m":122,"g":121},"db15148c":{"m":122,"g":121},"5259becd":{"m":122,"g":121},"ed1044ac":{"m":122,"g":121},"d717e73e":{"m":122,"g":121},"a1816187":{"m":122,"g":121},"e39628fd":{"m":122,"g":121},"bacb3825":{"m":122,"g":121},"b53d9e11":{"m":122,"g":121},"8a683821":{"m":122,"g":121},"52694b60":{"m":122,"g":121},"400bddf2":{"m":122,"g":121},"1e90fe2e":{"m":122,"g":121},"caa5d296":{"m":122,"g":121},"750940ae":{"m":122,"g":121},"42f8ea40":{"m":122,"g":121},"14cbe42f":{"m":122,"g":121},"685c0645":{"m":122,"g":121},"1357397a":{"m":122,"g":121},"42e1a72e":{"m":122,"g":121},"83a7c89c":{"m":122,"g":121},"0380ca82":{"m":122,"g":121},"ec92b0ce":{"m":122,"g":121},"e03b6bee":{"m":122,"g":121},"5e36a0b4":{"m":122,"g":121},"0297773a":{"m":122,"g":121},"587deb15":{"m":122,"g":121},"83087247":{"m":122,"g":121},"334543ff":{"m":122,"g":121},"c143f416":{"m":122,"g":121},"b48354c5":{"m":122,"g":121},"29195aaa":{"m":122,"g":121},"0ee831de":{"m":122,"g":121},"8d6ab1cb":{"m":122,"g":121},"84a9d0ea":{"m":122,"g":121},"737b58d6":{"m":122,"g":121},"77225d60":{"m":122,"g":121},"9c6e25d2":{"m":122,"g":121},"2a3763c3":{"m":122,"g":121},"fdd00295":{"m":122,"g":121},"25e73640":{"m":122,"g":121},"0da9845e":{"m":122,"g":121},"92885441":{"m":122,"g":121},"64cf868e":{"m":122,"g":121},"ea399527":{"m":122,"g":121},"41a11335":{"m":122,"g":121},"a1f2dc90":{"m":122,"g":121},"ea961060":{"m":122,"g":121},"b1e13e7c":{"m":122,"g":121},"cc7b04a2":{"m":122,"g":121},"d85d6dba":{"m":122,"g":121},"c5642a7a":{"m":122,"g":121},"691c8534":{"m":122,"g":121},"d2b8c412":{"m":122,"g":121},"bf8f7a94":{"m":122,"g":121},"81a632ac":{"m":122,"g":121},"83b22400":{"m":122,"g":121},"285a8e69":{"m":122,"g":121},"813bd6f8":{"m":122,"g":121},"729f612d":{"m":122,"g":121},"899453ac":{"m":122,"g":121},"ce832d70":{"m":122,"g":121},"88596739":{"m":122,"g":121},"a6ea3add":{"m":122,"g":121},"326c84c4":{"m":122,"g":121},"8da608cc":{"m":122,"g":121},"9fc3e8aa":{"m":122,"g":121},"c11b34d5":{"m":122,"g":121},"05ad28f2":{"m":122,"g":121},"0cae873f":{"m":122,"g":121},"a8b91f6b":{"m":122,"g":121},"959d1ab8":{"m":122,"g":121},"6c1c1933":{"m":122,"g":121},"3029d301":{"m":122,"g":121},"f389f017":{"m":122,"g":121},"caa4819b":{"m":122,"g":121},"a88b006e":{"m":122,"g":121},"ce112c07":{"m":122,"g":121},"f7dc2f33":{"m":122,"g":121},"cd784faf":{"m":122,"g":121},"75c09e1f":{"m":122,"g":121},"09af0a7b":{"m":122,"g":121},"c8d385ce":{"m":122,"g":121},"09938e1f":{"m":123,"g":127},"23407983":{"m":123,"g":127},"1357ab02":{"m":123,"g":127},"44da7377":{"m":123,"g":127},"fb9582c4":{"m":123,"g":127},"d22d0447":{"m":123,"g":127},"887742a1":{"m":123,"g":127},"34f7564d":{"m":123,"g":127},"1cfbbc42":{"m":123,"g":127},"55dfb539":{"m":123,"g":127},"42889acb":{"m":123,"g":127},"211f4070":{"m":123,"g":127},"befa41a1":{"m":123,"g":127},"30b26ee9":{"m":123,"g":127},"aa797d01":{"m":123,"g":127},"7cee07a0":{"m":123,"g":127},"bb517fe3":{"m":123,"g":127},"0e82fd3d":{"m":123,"g":127},"b7d70411":{"m":123,"g":127},"ff0b64e1":{"m":123,"g":127},"d84790db":{"m":123,"g":127},"0678beaa":{"m":123,"g":127},"c2d4716d":{"m":123,"g":127},"c14cc47e":{"m":123,"g":127},"dbcf85b7":{"m":123,"g":127},"83804bc6":{"m":123,"g":127},"d5fa019c":{"m":123,"g":127},"fef3a6b6":{"m":123,"g":127},"173e0f70":{"m":123,"g":127},"f600866a":{"m":123,"g":127},"93be7e86":{"m":123,"g":127},"60b0754c":{"m":123,"g":127},"0b24af4d":{"m":123,"g":127},"a209fb05":{"m":123,"g":127},"48d6bea1":{"m":123,"g":127},"1689c0e3":{"m":123,"g":127},"193fbb0b":{"m":123,"g":127},"e607850f":{"m":123,"g":127},"15efbcb4":{"m":123,"g":127},"243c064d":{"m":123,"g":127},"0b41a293":{"m":123,"g":127},"d31d48b3":{"m":123,"g":127},"88342607":{"m":123,"g":127},"fd7a72d6":{"m":123,"g":127},"21a8fa16":{"m":123,"g":127},"7a21d8b2":{"m":123,"g":127},"d36639ee":{"m":123,"g":127},"6ef23b98":{"m":123,"g":127},"385599cb":{"m":123,"g":127},"952fbe47":{"m":123,"g":127},"edb25693":{"m":123,"g":127},"3529c061":{"m":123,"g":127},"ffb32a85":{"m":123,"g":127},"14d80648":{"m":123,"g":127},"ab8b83f7":{"m":123,"g":127},"de0b10cf":{"m":123,"g":127},"6e29446e":{"m":123,"g":127},"0c3543d7":{"m":123,"g":127},"6a3b9fd0":{"m":123,"g":127},"65f1d065":{"m":123,"g":127},"9434a0e5":{"m":123,"g":127},"20315697":{"m":123,"g":127},"c9db7911":{"m":123,"g":127},"15ed27d7":{"m":123,"g":127},"66fb9b13":{"m":123,"g":127},"819fc591":{"m":123,"g":127},"7efd8b3d":{"m":123,"g":127},"a920b9da":{"m":123,"g":127},"76196b3c":{"m":123,"g":127},"95191ebd":{"m":123,"g":127},"9a512cf9":{"m":123,"g":127},"3451fc32":{"m":123,"g":127},"c550ab91":{"m":123,"g":127},"086f0b79":{"m":123,"g":127},"0afd6832":{"m":123,"g":127},"6f858930":{"m":123,"g":127},"229256c5":{"m":123,"g":127},"6b634493":{"m":123,"g":127},"756ad9ce":{"m":123,"g":127},"d2a8f71c":{"m":123,"g":127},"2b7bf11b":{"m":123,"g":127},"566ade03":{"m":123,"g":127},"69193f71":{"m":123,"g":127},"d5b6e50f":{"m":123,"g":127},"9632e48f":{"m":123,"g":127},"59cce594":{"m":123,"g":127},"795e98f8":{"m":123,"g":127},"358ae356":{"m":123,"g":127},"0c006b88":{"m":124,"g":127},"cd135bfe":{"m":124,"g":127},"fc84b073":{"m":124,"g":127},"8be0e1bc":{"m":124,"g":127},"8e1d6756":{"m":124,"g":127},"0da30dbc":{"m":124,"g":127},"58095cb0":{"m":124,"g":127},"fd3034da":{"m":124,"g":127},"fbbe16fa":{"m":124,"g":127},"6a1a64fa":{"m":124,"g":127},"e2715cf8":{"m":124,"g":127},"74630ba3":{"m":124,"g":127},"73e9a2ef":{"m":124,"g":127},"837b08eb":{"m":124,"g":127},"bb6a21cd":{"m":124,"g":127},"32ec68fa":{"m":124,"g":127},"3c0a6df8":{"m":124,"g":127},"2104d20e":{"m":124,"g":127},"f235498e":{"m":124,"g":127},"149dc9aa":{"m":124,"g":127},"9a954982":{"m":124,"g":127},"97be66c3":{"m":124,"g":127},"74243dff":{"m":124,"g":127},"4ea4c48b":{"m":124,"g":127},"cf5d27e3":{"m":124,"g":127},"9ec6031d":{"m":124,"g":127},"a5affb0c":{"m":124,"g":127},"5925d3d7":{"m":124,"g":127},"7ef1964a":{"m":124,"g":127},"b0476a06":{"m":124,"g":127},"ffba61a1":{"m":124,"g":127},"3c219eb0":{"m":124,"g":127},"1ffdcdc4":{"m":124,"g":127},"c7d57d5b":{"m":124,"g":127},"80802c4c":{"m":124,"g":127},"83b104ee":{"m":124,"g":127},"14127804":{"m":124,"g":127},"82f39dc1":{"m":124,"g":127},"627bac64":{"m":124,"g":127},"3651cfbf":{"m":124,"g":127},"c8547ecd":{"m":124,"g":127},"7bc1dae0":{"m":124,"g":127},"4fe53e58":{"m":124,"g":127},"fb2e816e":{"m":124,"g":127},"7c45b8b4":{"m":124,"g":127},"ba5b6823":{"m":124,"g":127},"508d2f7a":{"m":124,"g":127},"a889c854":{"m":124,"g":127},"4d84f886":{"m":124,"g":127},"dc4f5418":{"m":124,"g":127},"0648eb48":{"m":124,"g":127},"b88fab31":{"m":124,"g":127},"36942660":{"m":124,"g":127},"9f5e7018":{"m":124,"g":127},"cbf23dbb":{"m":124,"g":127},"6dade6c3":{"m":124,"g":127},"b419e20c":{"m":124,"g":127},"48641435":{"m":124,"g":127},"44b1b394":{"m":124,"g":127},"0711d150":{"m":124,"g":127},"303cc957":{"m":125,"g":127},"661c1c97":{"m":125,"g":127},"c022107f":{"m":125,"g":127},"56c83e0f":{"m":125,"g":127},"b51d46d0":{"m":125,"g":127},"10864731":{"m":125,"g":127},"665416f6":{"m":125,"g":127},"838bcb0d":{"m":125,"g":127},"f1f4c451":{"m":125,"g":127},"b0ee99dd":{"m":125,"g":127},"ddfcb7c8":{"m":125,"g":127},"58b12ccb":{"m":125,"g":127},"547de8c7":{"m":125,"g":127},"37c40a87":{"m":125,"g":127},"1240ac13":{"m":125,"g":127},"5639145f":{"m":125,"g":127},"afee2843":{"m":125,"g":127},"6f084880":{"m":125,"g":127},"611a4fd0":{"m":125,"g":127},"9ea2c686":{"m":125,"g":127},"e2a784ec":{"m":125,"g":127},"05559a4a":{"m":125,"g":127},"7bffc5dc":{"m":125,"g":127},"a5e5088d":{"m":125,"g":127},"ac19ce7e":{"m":125,"g":127},"95876d75":{"m":125,"g":127},"a30f1907":{"m":125,"g":127},"dc8a5a1c":{"m":125,"g":127},"e123648b":{"m":125,"g":127},"90401cf7":{"m":125,"g":127},"9cfe78dd":{"m":125,"g":127},"307e7a61":{"m":125,"g":127},"ddd1440d":{"m":125,"g":127},"d1be60c3":{"m":125,"g":127},"583bb180":{"m":125,"g":127},"1f2a6c69":{"m":125,"g":127},"61c7fe7a":{"m":125,"g":127},"83f89cc6":{"m":125,"g":127},"db24d346":{"m":125,"g":127},"885cfca2":{"m":125,"g":127},"4e916f98":{"m":125,"g":127},"4f65a646":{"m":125,"g":127},"f5b3ccd9":{"m":125,"g":127},"bb00e24f":{"m":125,"g":127},"210a9cab":{"m":125,"g":127},"b8ac4fcb":{"m":125,"g":127},"877cb528":{"m":125,"g":127},"3633f8b0":{"m":125,"g":127},"93cf60fc":{"m":125,"g":127},"b5e04173":{"m":125,"g":127},"8a821af7":{"m":125,"g":127},"4b1d163b":{"m":125,"g":127},"c21a3ec2":{"m":125,"g":127},"b142831a":{"m":125,"g":127},"d1340963":{"m":125,"g":127},"f290e801":{"m":125,"g":127},"9299a62f":{"m":125,"g":127},"49543be9":{"m":125,"g":127},"52362903":{"m":125,"g":127},"b2b26d43":{"m":125,"g":127},"d3a03aee":{"m":125,"g":127},"5f02b918":{"m":125,"g":127},"49653c88":{"m":125,"g":127},"44f594d8":{"m":125,"g":127},"2b6c4257":{"m":125,"g":127},"243ea585":{"m":125,"g":127},"6fee2c53":{"m":125,"g":127},"f1a9c72d":{"m":125,"g":127},"190002c6":{"m":125,"g":127},"0296f1cd":{"m":125,"g":127},"e039ff38":{"m":125,"g":127},"b8ddc296":{"m":125,"g":127},"0b88d520":{"m":125,"g":127},"d3d7f960":{"m":125,"g":127},"e4341872":{"m":125,"g":127},"fe19a580":{"m":125,"g":127},"55e8e399":{"m":125,"g":127},"32f79828":{"m":125,"g":127},"0f76976c":{"m":125,"g":127},"ae622790":{"m":125,"g":127},"5c9273c0":{"m":125,"g":127},"e316bcac":{"m":125,"g":127},"0fe9c1f7":{"m":125,"g":127},"61bfd9fa":{"m":125,"g":127},"c67fce16":{"m":125,"g":127},"bef37d6d":{"m":125,"g":127},"bc25ea67":{"m":125,"g":127},"1fa788ec":{"m":125,"g":127},"125f76ea":{"m":125,"g":127},"d8736c75":{"m":125,"g":127},"3b1cc466":{"m":125,"g":127},"0ee5ab5a":{"m":125,"g":127},"ed5e905c":{"m":125,"g":127},"d9c812d8":{"m":125,"g":127},"7257525c":{"m":125,"g":127},"34ba10ef":{"m":125,"g":127},"a119363f":{"m":125,"g":127},"fb314d7b":{"m":125,"g":127},"3a64844a":{"m":125,"g":127},"585c417f":{"m":125,"g":127},"b0d1c21d":{"m":125,"g":127},"b07c5e40":{"m":125,"g":127},"1772671b":{"m":125,"g":127},"6e6009fb":{"m":125,"g":127},"88a2a340":{"m":125,"g":127},"78c58621":{"m":125,"g":127},"c3bb348d":{"m":125,"g":127},"4e234b4c":{"m":125,"g":127},"4cc725ac":{"m":125,"g":127},"2cb42dc1":{"m":125,"g":127},"ebaf86d4":{"m":126,"g":127},"8359f185":{"m":126,"g":127},"5324f37a":{"m":126,"g":127},"2864c49f":{"m":126,"g":127},"d26ec39f":{"m":126,"g":127},"9c546bfd":{"m":126,"g":127},"c9b58164":{"m":126,"g":127},"e5e65e3d":{"m":126,"g":127},"ffeb28ba":{"m":126,"g":127},"b40f605f":{"m":126,"g":127},"3cdec20c":{"m":126,"g":127},"d28caaf6":{"m":126,"g":127},"ad8d24c3":{"m":126,"g":127},"018123b5":{"m":126,"g":127},"4983b7e7":{"m":126,"g":127},"f825137f":{"m":126,"g":127},"ae68158f":{"m":126,"g":127},"a7cc02e3":{"m":126,"g":127},"8ece99a9":{"m":126,"g":127},"44e391b6":{"m":126,"g":127},"1a5c313f":{"m":126,"g":127},"7ea5b42d":{"m":126,"g":127},"5ded5e27":{"m":126,"g":127},"33d1aeb0":{"m":126,"g":127},"dd909a51":{"m":126,"g":127},"60cb7167":{"m":126,"g":127},"151e1368":{"m":126,"g":127},"2f9952cd":{"m":126,"g":127},"3e7cc273":{"m":126,"g":127},"8f01a12d":{"m":126,"g":127},"28b8c579":{"m":126,"g":127},"7b877ab8":{"m":126,"g":127},"0d4a4184":{"m":126,"g":127},"2ca25a8a":{"m":126,"g":127},"cc2e36c3":{"m":126,"g":127},"d8f7816a":{"m":126,"g":127},"99e25805":{"m":126,"g":127},"4a2768a8":{"m":126,"g":127},"9b247f73":{"m":126,"g":127},"36d14712":{"m":126,"g":127},"e0e6a6ef":{"m":126,"g":127},"e8114102":{"m":126,"g":127},"14a339fc":{"m":126,"g":127},"5f662e78":{"m":126,"g":127},"7f5055ed":{"m":126,"g":127},"63728b11":{"m":126,"g":127},"5de25f78":{"m":126,"g":127},"d52800db":{"m":126,"g":127},"38a704bc":{"m":126,"g":127},"a06c44f9":{"m":126,"g":127},"527b7d3f":{"m":126,"g":127},"f09eee03":{"m":126,"g":127},"e38994dd":{"m":126,"g":127},"4a78031a":{"m":126,"g":127},"ea10a9d1":{"m":126,"g":127},"71aea45c":{"m":126,"g":127},"e3b38d71":{"m":126,"g":127},"5c0cadd0":{"m":126,"g":127},"39b1d048":{"m":126,"g":127},"8db7fc41":{"m":126,"g":127},"fe92d4d8":{"m":126,"g":127},"fc8cda14":{"m":126,"g":127},"6a7322ff":{"m":126,"g":127},"c751cb38":{"m":126,"g":127},"3594815a":{"m":126,"g":127},"9caca6a4":{"m":126,"g":127},"08c805a8":{"m":126,"g":127},"f18ec927":{"m":126,"g":127},"aea88fa7":{"m":126,"g":127},"2fe4e69f":{"m":126,"g":127},"012bfc4f":{"m":126,"g":127},"0493775b":{"m":126,"g":127},"40b26b45":{"m":126,"g":127},"9840bf4f":{"m":126,"g":127},"7b2fb3d4":{"m":128,"g":131},"d64dd3e1":{"m":128,"g":131},"3ccd7fa6":{"m":128,"g":131},"254f62d8":{"m":128,"g":131},"2b8b9d84":{"m":128,"g":131},"e970892f":{"m":128,"g":131},"d5fa58c4":{"m":128,"g":131},"9f011f61":{"m":128,"g":131},"3b18fd4c":{"m":128,"g":131},"1869f25c":{"m":128,"g":131},"e019f233":{"m":128,"g":131},"b1c688fb":{"m":128,"g":131},"8e3663d4":{"m":128,"g":131},"9edb0e0d":{"m":128,"g":131},"95f43669":{"m":128,"g":131},"50691d7b":{"m":128,"g":131},"6afe3963":{"m":128,"g":131},"191f5c77":{"m":128,"g":131},"d7246708":{"m":128,"g":131},"efc5d8f5":{"m":128,"g":131},"ef32a252":{"m":128,"g":131},"9509c4cc":{"m":128,"g":131},"4e19c1d5":{"m":128,"g":131},"78a4b446":{"m":128,"g":131},"f9696641":{"m":128,"g":131},"f35f7f12":{"m":128,"g":131},"12c789eb":{"m":128,"g":131},"597d4160":{"m":128,"g":131},"1ca205f6":{"m":128,"g":131},"24a25ffa":{"m":128,"g":131},"db7299aa":{"m":128,"g":131},"13366843":{"m":128,"g":131},"be353ffd":{"m":128,"g":131},"20e59f95":{"m":128,"g":131},"0d116b9a":{"m":128,"g":131},"4a56fa5c":{"m":128,"g":131},"51f9b962":{"m":128,"g":131},"daf494b6":{"m":128,"g":131},"37e8724e":{"m":128,"g":131},"10592e9c":{"m":128,"g":131},"2aec8b6e":{"m":128,"g":131},"0d41ddfb":{"m":128,"g":131},"37c87615":{"m":128,"g":131},"3f400f25":{"m":128,"g":131},"d91b16eb":{"m":128,"g":131},"4a10e37b":{"m":128,"g":131},"b051d76d":{"m":128,"g":131},"d52d992a":{"m":128,"g":131},"d971f228":{"m":128,"g":131},"1d3d42bd":{"m":128,"g":131},"8e9f05ec":{"m":128,"g":131},"bc083521":{"m":128,"g":131},"33f08a98":{"m":128,"g":131},"8e6083bf":{"m":128,"g":131},"2fbc78a0":{"m":128,"g":131},"f0b5ccf5":{"m":128,"g":131},"b732ffa4":{"m":128,"g":131},"56fc4830":{"m":128,"g":131},"2a96e302":{"m":128,"g":131},"8a437340":{"m":128,"g":131},"f0021c0d":{"m":128,"g":131},"7ee3e364":{"m":128,"g":131},"6d5e16fb":{"m":128,"g":131},"67e6f143":{"m":128,"g":131},"34851471":{"m":128,"g":131},"c2083116":{"m":128,"g":131},"10285ec2":{"m":128,"g":131},"9b3fc186":{"m":128,"g":131},"172c71a2":{"m":128,"g":127},"af373636":{"m":128,"g":127},"a5be6ef9":{"m":128,"g":127},"8f4e18a2":{"m":128,"g":127},"e9681444":{"m":128,"g":127},"eae59b33":{"m":128,"g":127},"14dc0523":{"m":128,"g":127},"b2236691":{"m":128,"g":127},"dcc47a56":{"m":128,"g":127},"6448b4cd":{"m":128,"g":127},"5ae0ac42":{"m":128,"g":127},"22f641ab":{"m":128,"g":127},"0997c78d":{"m":128,"g":127},"fd3be107":{"m":128,"g":127},"665f43bd":{"m":128,"g":127},"a53f2d6c":{"m":128,"g":127},"a7002e61":{"m":128,"g":127},"84e151ac":{"m":128,"g":127},"875a25dd":{"m":128,"g":127},"fc55b45e":{"m":128,"g":127},"0050ff25":{"m":128,"g":127},"af9f71f9":{"m":128,"g":127},"15264232":{"m":128,"g":127},"f8d3d80f":{"m":128,"g":127},"5027739f":{"m":128,"g":127},"3701f34d":{"m":128,"g":127},"821fb060":{"m":128,"g":127},"ace27c0c":{"m":128,"g":127},"ed1d18d4":{"m":128,"g":127},"e7b57b0d":{"m":128,"g":127},"49141df9":{"m":128,"g":127},"7b79cc4f":{"m":128,"g":127},"385ff0e5":{"m":128,"g":127},"fc5da1e8":{"m":128,"g":127},"04848ba7":{"m":128,"g":127},"9bc6a9ad":{"m":128,"g":127},"7cdaedb8":{"m":128,"g":127},"1f134f85":{"m":128,"g":127},"5f72d36d":{"m":128,"g":127},"7a2254b2":{"m":128,"g":127},"b5904999":{"m":128,"g":127},"922525ee":{"m":128,"g":127},"ee3e337c":{"m":128,"g":127},"e523e216":{"m":128,"g":127},"2ce23777":{"m":128,"g":127},"19f6a33c":{"m":128,"g":127},"e9c0c558":{"m":128,"g":127},"9b41f31a":{"m":128,"g":127},"9db3add3":{"m":128,"g":127},"c9e5799b":{"m":128,"g":127},"2966367a":{"m":128,"g":127},"87791007":{"m":128,"g":127},"dd192a55":{"m":128,"g":127},"bfe638f7":{"m":128,"g":127},"4ac65e3c":{"m":128,"g":127},"0779c3d1":{"m":128,"g":127},"5c2d72ba":{"m":128,"g":127},"9bd511a5":{"m":128,"g":127},"85b8c5c4":{"m":128,"g":127},"c8b7516f":{"m":128,"g":127},"67e9d287":{"m":128,"g":127},"e7e89349":{"m":128,"g":127},"aead0ef5":{"m":128,"g":127},"66640835":{"m":128,"g":127},"c2d69e8b":{"m":128,"g":127},"4c1e909a":{"m":128,"g":127},"2bb0317e":{"m":128,"g":127},"86255f27":{"m":128,"g":127},"e4b29370":{"m":128,"g":127},"7a8524b4":{"m":128,"g":127},"909d0d38":{"m":128,"g":127},"e42df37d":{"m":128,"g":127},"6d21392b":{"m":128,"g":127},"a1cb717d":{"m":128,"g":127},"c9456491":{"m":128,"g":127},"c4b74c1d":{"m":128,"g":127},"7aa44390":{"m":128,"g":127},"4eda9969":{"m":128,"g":127},"03a7e6f4":{"m":128,"g":127},"2cdde3d4":{"m":128,"g":127},"401ed0c5":{"m":128,"g":127},"4ef43905":{"m":128,"g":127},"d646cf63":{"m":128,"g":127},"4edb2401":{"m":128,"g":127},"706502ff":{"m":128,"g":127},"c2e56dad":{"m":128,"g":127},"9c1c5c6d":{"m":128,"g":127},"2d531946":{"m":128,"g":127},"7ae368ef":{"m":129,"g":131},"c4e20cad":{"m":129,"g":131},"5dad1ff1":{"m":129,"g":131},"ca52ed42":{"m":129,"g":131},"5c03aa3e":{"m":129,"g":131},"253be18e":{"m":129,"g":131},"084b06e7":{"m":129,"g":131},"fc6fb550":{"m":129,"g":131},"7c38eca1":{"m":129,"g":131},"427b08e2":{"m":129,"g":131},"0141ca37":{"m":129,"g":131},"25a6be49":{"m":129,"g":131},"df1f3124":{"m":129,"g":131},"c5947ecd":{"m":129,"g":131},"9530b766":{"m":129,"g":131},"3067b3f0":{"m":129,"g":131},"e0ec42c7":{"m":129,"g":131},"9c9d7091":{"m":129,"g":131},"51a86ce6":{"m":129,"g":131},"64092c8b":{"m":129,"g":131},"63b9300f":{"m":129,"g":131},"21ec99be":{"m":129,"g":131},"383689e3":{"m":129,"g":131},"236a7c23":{"m":129,"g":131},"3dabd609":{"m":129,"g":131},"73df5253":{"m":129,"g":131},"e6420100":{"m":129,"g":131},"c9e20901":{"m":129,"g":131},"106df4ea":{"m":129,"g":131},"427a19b6":{"m":129,"g":131},"1f930cd2":{"m":129,"g":131},"11ce0516":{"m":129,"g":131},"cd4151ab":{"m":129,"g":131},"8fe8b635":{"m":129,"g":131},"8a7b1b83":{"m":129,"g":131},"26aebf83":{"m":129,"g":131},"3ab8ae68":{"m":129,"g":131},"03888b9d":{"m":129,"g":131},"796d82b1":{"m":129,"g":131},"1da59e83":{"m":129,"g":131},"1d66a14c":{"m":129,"g":131},"02af51e4":{"m":129,"g":131},"eb500884":{"m":129,"g":131},"079b1738":{"m":129,"g":131},"57f933fd":{"m":129,"g":131},"1f2b84d2":{"m":129,"g":131},"e7d6027e":{"m":129,"g":131},"07821352":{"m":129,"g":131},"edbeaf3b":{"m":129,"g":131},"d9dca282":{"m":129,"g":131},"9325f945":{"m":129,"g":131},"41b7aab8":{"m":129,"g":131},"491f4fe8":{"m":129,"g":131},"34035d8c":{"m":129,"g":131},"92ca6295":{"m":129,"g":131},"ec92d7f1":{"m":129,"g":131},"79b389da":{"m":129,"g":131},"3de09aad":{"m":129,"g":131},"2e8f54e6":{"m":129,"g":131},"e55731b6":{"m":129,"g":131},"45264554":{"m":129,"g":131},"c4293f59":{"m":129,"g":131},"a2423052":{"m":129,"g":131},"bc3d2a85":{"m":129,"g":131},"d815d002":{"m":129,"g":131},"fa9021b2":{"m":129,"g":131},"9c800728":{"m":129,"g":131},"630a6930":{"m":129,"g":131},"7ce8faae":{"m":129,"g":131},"de153cf7":{"m":129,"g":131},"f4a0c5c7":{"m":129,"g":131},"0f8e5394":{"m":129,"g":131},"e8ba5a66":{"m":129,"g":131},"a2960bdd":{"m":129,"g":131},"487c8d4d":{"m":129,"g":131},"f87b8eab":{"m":129,"g":131},"e8542db5":{"m":129,"g":131},"6df1e8d6":{"m":129,"g":131},"4addb602":{"m":129,"g":131},"bd0e6908":{"m":129,"g":131},"0825d7f4":{"m":129,"g":131},"0b9dbea5":{"m":129,"g":131},"982db4eb":{"m":129,"g":131},"f5f3a5d9":{"m":129,"g":131},"f138ae57":{"m":129,"g":131},"decb4896":{"m":129,"g":131},"c72f0756":{"m":129,"g":131},"f1115cf5":{"m":129,"g":131},"412160f4":{"m":129,"g":131},"7b03cc64":{"m":129,"g":131},"9872a677":{"m":129,"g":131},"c15c864b":{"m":129,"g":131},"dc7bdc73":{"m":129,"g":131},"0a9d6453":{"m":129,"g":131},"340c613a":{"m":129,"g":131},"36b729c2":{"m":129,"g":131},"67e6ef4b":{"m":129,"g":131},"65ba5ab8":{"m":129,"g":131},"990023e5":{"m":129,"g":131},"5ddd2f6b":{"m":129,"g":131},"0ae4b1ad":{"m":129,"g":131},"94cd64a7":{"m":129,"g":131},"b870271a":{"m":129,"g":131},"22ee9b01":{"m":129,"g":131},"9d0e5f1f":{"m":129,"g":131},"1d3d8b34":{"m":129,"g":131},"155a9e72":{"m":129,"g":131},"d7cb08c5":{"m":129,"g":131},"3339c810":{"m":129,"g":131},"d6c88d51":{"m":129,"g":131},"f03ea34a":{"m":129,"g":131},"4cafc835":{"m":129,"g":131},"c6a52f44":{"m":129,"g":131},"c6d34a06":{"m":129,"g":131},"848ee570":{"m":129,"g":131},"ce6b7dfc":{"m":129,"g":131},"0fe74af5":{"m":129,"g":131},"0a362d65":{"m":129,"g":131},"143b57b8":{"m":129,"g":131},"f446b51c":{"m":129,"g":131},"a102a050":{"m":129,"g":131},"6bad6a36":{"m":129,"g":131},"11b6217a":{"m":129,"g":131},"0b0b2607":{"m":129,"g":131},"841eb29d":{"m":129,"g":131},"45cf5758":{"m":129,"g":131},"0e8ce1e8":{"m":129,"g":131},"ea1e9f6b":{"m":129,"g":131},"f6e37d3e":{"m":129,"g":131},"ab9a46d4":{"m":129,"g":131},"621061f0":{"m":129,"g":131},"7daddcdb":{"m":129,"g":131},"91d249cd":{"m":129,"g":131},"95102896":{"m":129,"g":131},"3543a04a":{"m":129,"g":131},"e12c78aa":{"m":129,"g":131},"bce40fa2":{"m":129,"g":131},"051ad833":{"m":129,"g":131},"4c9f7c97":{"m":129,"g":131},"63b05621":{"m":129,"g":131},"21af8e73":{"m":129,"g":131},"7ab548ef":{"m":129,"g":131},"bab033b9":{"m":129,"g":131},"63500426":{"m":129,"g":131},"25758647":{"m":129,"g":131},"2bc8ee8b":{"m":129,"g":131},"ab843ced":{"m":129,"g":131},"6edffc63":{"m":129,"g":131},"077ca70e":{"m":129,"g":131},"7cb04dc0":{"m":129,"g":131},"5443db87":{"m":129,"g":131},"9f340ab1":{"m":129,"g":131},"70c6f951":{"m":129,"g":131},"d941a3be":{"m":129,"g":131},"b12c9e5c":{"m":129,"g":131},"e9e90460":{"m":129,"g":131},"6330d664":{"m":129,"g":131},"91e8dc37":{"m":129,"g":131},"231df4b0":{"m":129,"g":131},"b087ef8b":{"m":129,"g":131},"5155016b":{"m":129,"g":131},"082b54c6":{"m":129,"g":131},"a8ef4d18":{"m":129,"g":131},"15ff6982":{"m":129,"g":131},"9adef42c":{"m":129,"g":131},"685b9d82":{"m":129,"g":131},"a223402f":{"m":129,"g":131},"44d0a848":{"m":129,"g":131},"5b7da0f5":{"m":129,"g":131},"697a77bf":{"m":129,"g":131},"0a186924":{"m":129,"g":131},"b6312e62":{"m":129,"g":131},"779cbc6e":{"m":129,"g":131},"67c8c867":{"m":129,"g":131},"69a03bc3":{"m":129,"g":131},"66f242b9":{"m":129,"g":131},"8a9b8b84":{"m":129,"g":131},"7e964b51":{"m":129,"g":131},"a0d9f6cd":{"m":129,"g":131},"8308cd36":{"m":129,"g":131},"e0e8a996":{"m":129,"g":131},"5102d009":{"m":129,"g":131},"0dd759e0":{"m":129,"g":131},"5e70880e":{"m":129,"g":131},"9dab534b":{"m":129,"g":131},"262c3c1f":{"m":129,"g":131},"6c190cbd":{"m":129,"g":131},"eff6a07c":{"m":129,"g":131},"b704b0a9":{"m":129,"g":131},"15729dbc":{"m":129,"g":131},"21b0582d":{"m":129,"g":131},"5795da5e":{"m":129,"g":131},"540d6fee":{"m":129,"g":131},"5a8adca9":{"m":129,"g":131},"007c3e23":{"m":129,"g":131},"7130ad3a":{"m":129,"g":131},"35a4c21a":{"m":129,"g":131},"846ba3c6":{"m":129,"g":131},"18fb5158":{"m":129,"g":131},"ca5c8b16":{"m":129,"g":131},"f33e5d1e":{"m":129,"g":131},"13e5beea":{"m":129,"g":131},"c53e729d":{"m":129,"g":131},"03a26557":{"m":129,"g":131},"873382a9":{"m":129,"g":131},"391a863b":{"m":129,"g":131},"36b1bcd2":{"m":129,"g":131},"64a11303":{"m":129,"g":131},"1ab6ce0e":{"m":129,"g":131},"e99ca6ac":{"m":129,"g":131},"fcccaf90":{"m":129,"g":131},"215a97fa":{"m":129,"g":131},"5eed5fc0":{"m":129,"g":131},"808b6dfd":{"m":129,"g":131},"4852aa05":{"m":129,"g":131},"f922bfd5":{"m":129,"g":131},"dfd7ab96":{"m":129,"g":131},"46673b42":{"m":129,"g":131},"3421d049":{"m":129,"g":131},"d3d404d3":{"m":129,"g":131},"64225a8a":{"m":129,"g":131},"d64bf6c6":{"m":129,"g":131},"432ecf84":{"m":129,"g":131},"0b3f002d":{"m":129,"g":131},"6f094def":{"m":129,"g":131},"59464dbf":{"m":129,"g":131},"1f7fcc10":{"m":129,"g":131},"dbab5d50":{"m":129,"g":131},"760c20b3":{"m":129,"g":131},"407cb3ce":{"m":129,"g":131},"7cc43bd4":{"m":129,"g":131},"cce2d748":{"m":129,"g":131},"c1dd9a95":{"m":129,"g":131},"ed8786b0":{"m":129,"g":131},"f9fe0630":{"m":129,"g":131},"8ff3ef1f":{"m":129,"g":131},"da182e4b":{"m":129,"g":131},"83e72077":{"m":129,"g":131},"a2c388ba":{"m":129,"g":131},"9384fa27":{"m":129,"g":131},"173e73fa":{"m":129,"g":131},"a164259e":{"m":129,"g":131},"b0a26ba6":{"m":129,"g":131},"de430b67":{"m":129,"g":131},"4b45d556":{"m":129,"g":131},"db0ffc09":{"m":129,"g":131},"eb1d8854":{"m":129,"g":131},"bf108692":{"m":129,"g":131},"e83bd1fa":{"m":129,"g":131},"9dc15d85":{"m":129,"g":131},"fafaa2cc":{"m":129,"g":131},"9b4b3441":{"m":129,"g":131},"94216a9c":{"m":129,"g":131},"9535015d":{"m":129,"g":131},"a3b578fc":{"m":129,"g":131},"b60e769d":{"m":129,"g":131},"a95a3807":{"m":129,"g":131},"98b38de3":{"m":129,"g":131},"a146f833":{"m":129,"g":131},"1dd9a6ae":{"m":129,"g":131},"8ef11569":{"m":129,"g":131},"ecefc790":{"m":129,"g":131},"aeac6220":{"m":129,"g":131},"04b52fa8":{"m":129,"g":131},"e5c0f591":{"m":129,"g":131},"981ca831":{"m":129,"g":131},"414248e0":{"m":129,"g":131},"f56b9b42":{"m":129,"g":131},"b2f7b08c":{"m":129,"g":131},"75222bfe":{"m":129,"g":131},"9ea19533":{"m":129,"g":131},"dbf22152":{"m":129,"g":131},"d5e03468":{"m":129,"g":131},"a22104a6":{"m":129,"g":131},"4683e244":{"m":129,"g":131},"9054e844":{"m":129,"g":131},"18403f6b":{"m":129,"g":131},"2892265d":{"m":129,"g":131},"c9bd1aca":{"m":129,"g":131},"618ca238":{"m":129,"g":131},"5c291549":{"m":129,"g":131},"aaa40a9b":{"m":129,"g":131},"dd70cf99":{"m":129,"g":131},"d4593964":{"m":129,"g":131},"53fffefd":{"m":129,"g":131},"b964ce61":{"m":129,"g":131},"ac5505b0":{"m":129,"g":131},"a90435c0":{"m":129,"g":131},"5354d7b7":{"m":129,"g":131},"e0148677":{"m":129,"g":131},"04793508":{"m":129,"g":131},"cad78789":{"m":129,"g":131},"b29769f3":{"m":129,"g":131},"3990b84b":{"m":129,"g":131},"5a4394a3":{"m":129,"g":131},"dd303614":{"m":129,"g":131},"86312468":{"m":129,"g":131},"5625e32c":{"m":129,"g":131},"ca548d83":{"m":129,"g":131},"3397bcee":{"m":129,"g":131},"3e804bb0":{"m":129,"g":131},"a22de641":{"m":129,"g":131},"ac438226":{"m":129,"g":131},"a92afb00":{"m":129,"g":131},"0eea17e3":{"m":129,"g":131},"38052432":{"m":129,"g":131},"8bfce9b0":{"m":129,"g":131},"b41afa37":{"m":129,"g":131},"94ae816f":{"m":129,"g":131},"53620a1b":{"m":129,"g":131},"59b4d7f8":{"m":129,"g":131},"a56f7702":{"m":129,"g":131},"1b48e1b9":{"m":129,"g":131},"a24aefe5":{"m":129,"g":131},"45c572c5":{"m":129,"g":131},"681b9e64":{"m":129,"g":131},"dab06b50":{"m":129,"g":131},"85ffce30":{"m":129,"g":131},"e94ef9fc":{"m":129,"g":131},"964cdedc":{"m":129,"g":131},"aa6e2c8a":{"m":129,"g":131},"1776dce5":{"m":129,"g":131},"99e13d18":{"m":129,"g":131},"dc836909":{"m":129,"g":131},"323fed5c":{"m":129,"g":131},"eff7df6d":{"m":129,"g":131},"5e7f91d4":{"m":129,"g":131},"a34d3abb":{"m":129,"g":131},"6d0e0b9b":{"m":129,"g":131},"589d9ad5":{"m":129,"g":131},"a244c030":{"m":129,"g":131},"1bb063aa":{"m":129,"g":131},"b30f63c4":{"m":129,"g":131},"90a01335":{"m":129,"g":131},"43602790":{"m":129,"g":131},"b537ac0d":{"m":129,"g":131},"eda2f700":{"m":129,"g":131},"8c212a20":{"m":129,"g":131},"d754ce97":{"m":129,"g":131},"475962a1":{"m":129,"g":131},"bfcf15a1":{"m":129,"g":131},"fb04d434":{"m":129,"g":131},"6be65ae4":{"m":129,"g":131},"750084ae":{"m":129,"g":131},"64480ec7":{"m":129,"g":131},"db2d362d":{"m":129,"g":131},"81e86992":{"m":129,"g":131},"c4db77f8":{"m":129,"g":131},"c0a2513b":{"m":129,"g":131},"3ae664d7":{"m":129,"g":131},"c56fc424":{"m":129,"g":131},"3f1cfd87":{"m":129,"g":131},"b5344b31":{"m":129,"g":131},"6b262ac8":{"m":129,"g":131},"ada8ce1f":{"m":129,"g":131},"5a2c7039":{"m":129,"g":131},"42028af6":{"m":129,"g":131},"7291c72e":{"m":129,"g":131},"6bc30628":{"m":129,"g":131},"fa924410":{"m":129,"g":131},"fc9efdcb":{"m":129,"g":131},"2dec555d":{"m":129,"g":131},"acde21d8":{"m":129,"g":131},"4528cb7d":{"m":129,"g":131},"a352e833":{"m":129,"g":131},"852eb6ce":{"m":129,"g":131},"7af9b88c":{"m":129,"g":131},"2847e5c4":{"m":129,"g":131},"c8ede0e9":{"m":129,"g":131},"19729f72":{"m":129,"g":131},"b51f9bbe":{"m":129,"g":131},"4a8442af":{"m":129,"g":131},"7dcf910d":{"m":129,"g":131},"c7b37b70":{"m":129,"g":131},"10e0b83a":{"m":129,"g":131},"bc42c8c4":{"m":129,"g":131},"2e3a69ae":{"m":129,"g":131},"21370ef7":{"m":129,"g":131},"c3c4da71":{"m":129,"g":131},"dc694624":{"m":129,"g":131},"48ca9f75":{"m":129,"g":131},"127d59cd":{"m":129,"g":131},"af6bcadc":{"m":129,"g":131},"67fca6b2":{"m":129,"g":131},"f88b2aa6":{"m":129,"g":131},"17b24aca":{"m":129,"g":131},"bfaf0b86":{"m":129,"g":131},"83756a4b":{"m":129,"g":131},"a3557949":{"m":129,"g":131},"e72cf136":{"m":129,"g":131},"196b940a":{"m":129,"g":131},"a1e1e533":{"m":129,"g":131},"d4a4dcdf":{"m":129,"g":131},"97ba2c2d":{"m":129,"g":131},"e197bef5":{"m":129,"g":131},"8900f996":{"m":129,"g":131},"ba9102f9":{"m":129,"g":131},"b638abba":{"m":129,"g":131},"37980559":{"m":129,"g":131},"075ba74d":{"m":129,"g":131},"f7be98e1":{"m":129,"g":131},"9a1a9a42":{"m":129,"g":131},"6c2e5fcd":{"m":129,"g":131},"0d2d6878":{"m":129,"g":131},"9f59194f":{"m":129,"g":131},"cf1f0166":{"m":129,"g":131},"10969ae4":{"m":129,"g":131},"92ad2ff9":{"m":129,"g":131},"9b64f6f3":{"m":129,"g":131},"c0d1a338":{"m":129,"g":131},"a9d22b75":{"m":129,"g":131},"6b9459e8":{"m":129,"g":131},"3a6ec47b":{"m":129,"g":131},"b8e32e79":{"m":129,"g":131},"f5566acc":{"m":129,"g":131},"d79e1294":{"m":129,"g":131},"6e9b1549":{"m":129,"g":131},"109f27ba":{"m":129,"g":131},"c1a30aa7":{"m":129,"g":131},"2e1dbdb2":{"m":129,"g":131},"6d025fd3":{"m":129,"g":131},"518467be":{"m":129,"g":131},"63807079":{"m":129,"g":131},"f6cfe9f1":{"m":129,"g":131},"7bc99d41":{"m":129,"g":131},"e2d67468":{"m":129,"g":131},"6beb6e99":{"m":129,"g":131},"7e88b9c1":{"m":129,"g":131},"cfcf2758":{"m":129,"g":131},"ac81db66":{"m":129,"g":131},"33905005":{"m":129,"g":131},"820e13c9":{"m":129,"g":131},"595adf6d":{"m":129,"g":131},"a5ad0069":{"m":129,"g":131},"4e41edcb":{"m":129,"g":131},"67071f55":{"m":129,"g":131},"0c966779":{"m":129,"g":131},"f3386077":{"m":129,"g":131},"4ce8fb3c":{"m":129,"g":131},"4c3573e4":{"m":129,"g":131},"f1be8aa0":{"m":129,"g":131},"aa8ecbda":{"m":129,"g":131},"9188fecc":{"m":129,"g":131},"26ca0746":{"m":129,"g":131},"90c18a16":{"m":129,"g":131},"d7984f31":{"m":129,"g":131},"7119d188":{"m":129,"g":131},"fe3bbfb4":{"m":129,"g":131},"9846f8ed":{"m":129,"g":131},"85ae508e":{"m":129,"g":131},"d879e37f":{"m":129,"g":131},"e2c9a590":{"m":129,"g":131},"a1e37b02":{"m":129,"g":131},"e389f91d":{"m":129,"g":131},"aac07bf7":{"m":129,"g":131},"ea89a3a0":{"m":129,"g":131},"2bc7c5eb":{"m":129,"g":131},"a63f433b":{"m":129,"g":131},"58f8f4e4":{"m":129,"g":131},"25acbbc6":{"m":129,"g":131},"a8fcbf6f":{"m":129,"g":131},"e486308c":{"m":129,"g":131},"60420109":{"m":129,"g":131},"ff00b6ad":{"m":129,"g":131},"7b445260":{"m":129,"g":131},"c236d05f":{"m":129,"g":131},"df561392":{"m":129,"g":131},"15db5497":{"m":129,"g":131},"9ba3597d":{"m":129,"g":131},"80797c2a":{"m":129,"g":131},"7afff8fd":{"m":129,"g":131},"ac406d43":{"m":129,"g":131},"b436113f":{"m":129,"g":131},"8b5e2c53":{"m":129,"g":131},"b24235b8":{"m":129,"g":131},"ae7698fb":{"m":129,"g":131},"a3e4fe4b":{"m":129,"g":131},"f3e9336d":{"m":129,"g":131},"290fcd89":{"m":129,"g":131},"1dcde539":{"m":129,"g":131},"15bc1f5c":{"m":129,"g":131},"4d597616":{"m":129,"g":131},"ab63f3c5":{"m":129,"g":131},"147b7823":{"m":129,"g":131},"d368c745":{"m":129,"g":131},"7e626d12":{"m":129,"g":131},"2a577344":{"m":129,"g":131},"6abb8051":{"m":130,"g":131},"2e3946d8":{"m":130,"g":131},"32f8b606":{"m":130,"g":131},"b9bef31a":{"m":130,"g":131},"8550822d":{"m":130,"g":131},"8810152e":{"m":130,"g":131},"7bf16c63":{"m":130,"g":131},"39f9a9c2":{"m":130,"g":131},"d69ecc19":{"m":130,"g":131},"763888b5":{"m":130,"g":131},"9a327bdf":{"m":130,"g":131},"2de98010":{"m":130,"g":131},"8200fb56":{"m":130,"g":131},"cb4cdb43":{"m":130,"g":131},"80cfca50":{"m":130,"g":131},"7871593c":{"m":130,"g":131},"4a62a0e3":{"m":130,"g":131},"12a08efc":{"m":130,"g":131},"06836ad0":{"m":130,"g":131},"f72a7703":{"m":130,"g":131},"aeff0d38":{"m":130,"g":131},"cf0478d6":{"m":130,"g":131},"a2ca9bd4":{"m":130,"g":131},"36361adc":{"m":130,"g":131},"661e9775":{"m":130,"g":131},"2970f229":{"m":130,"g":131},"c08b780f":{"m":130,"g":131},"8fbf7dd5":{"m":130,"g":131},"1915a1f8":{"m":130,"g":131},"85d0ccfa":{"m":130,"g":131},"a4ffd665":{"m":130,"g":131},"559202b5":{"m":130,"g":131},"f57d4fe7":{"m":130,"g":131},"b7b7524e":{"m":130,"g":131},"6799847e":{"m":130,"g":131},"03b835e7":{"m":130,"g":131},"aff1238e":{"m":130,"g":131},"5e2cda61":{"m":130,"g":131},"b0bbc7f5":{"m":130,"g":131},"673c11ba":{"m":130,"g":131},"3b47973a":{"m":130,"g":131},"f6423b62":{"m":130,"g":131},"84efe54b":{"m":130,"g":131},"948b6ace":{"m":130,"g":131},"f124539a":{"m":130,"g":131},"c8683ae3":{"m":130,"g":131},"125e17ef":{"m":130,"g":131},"88c459c6":{"m":130,"g":131},"ae6a6630":{"m":130,"g":131},"26d95008":{"m":130,"g":131},"f2b5dcc9":{"m":130,"g":131},"9abcab3f":{"m":130,"g":131},"e5135b73":{"m":130,"g":131},"41d61faa":{"m":130,"g":131},"3c7886ec":{"m":130,"g":131},"0e4d8790":{"m":130,"g":131},"6d5d76ad":{"m":130,"g":131},"91c9c14c":{"m":130,"g":131},"32a32cf7":{"m":130,"g":131},"be4a3ec3":{"m":130,"g":131},"ff6e3ea9":{"m":130,"g":131},"dd91d38e":{"m":130,"g":131},"5f6f550a":{"m":130,"g":131},"5edbe351":{"m":130,"g":131},"d2b42477":{"m":130,"g":131},"9dfa01a4":{"m":130,"g":131},"e592ee65":{"m":130,"g":131},"cee93a6f":{"m":130,"g":131},"bc388471":{"m":130,"g":131},"3e40c636":{"m":130,"g":131},"80122e4f":{"m":130,"g":131},"e12c6b32":{"m":130,"g":131},"6d417918":{"m":130,"g":131},"35a9a073":{"m":130,"g":131},"ea177372":{"m":130,"g":131},"42fcf543":{"m":130,"g":131},"d257bf87":{"m":130,"g":131},"7b0c7ad1":{"m":130,"g":131},"d30d6b36":{"m":130,"g":131},"d881f314":{"m":130,"g":131},"2ac5b983":{"m":130,"g":131},"a0dde90a":{"m":130,"g":131},"b988c18e":{"m":130,"g":131},"e41664ba":{"m":130,"g":131},"3d1b591a":{"m":130,"g":131},"e11f795f":{"m":130,"g":131},"959a1746":{"m":130,"g":131},"b72f0268":{"m":130,"g":131},"09376fd7":{"m":130,"g":131},"aed835e3":{"m":130,"g":131},"49dfa1d8":{"m":130,"g":131},"1ea6b740":{"m":130,"g":131},"e73173b0":{"m":130,"g":131},"16e8463a":{"m":130,"g":131},"cf9a774c":{"m":130,"g":131},"ec7b2c16":{"m":130,"g":131},"1569fc7f":{"m":130,"g":131},"5a46fb15":{"m":130,"g":131},"38daa294":{"m":130,"g":131},"66984a8b":{"m":130,"g":131},"889b46ea":{"m":130,"g":131},"05284378":{"m":130,"g":131},"a8904560":{"m":130,"g":131},"66280987":{"m":130,"g":131},"205f041e":{"m":130,"g":131},"7235a7fb":{"m":130,"g":131},"8fce9e7b":{"m":130,"g":131},"53477322":{"m":130,"g":131},"2ce121a1":{"m":130,"g":131},"35ba6fe1":{"m":130,"g":131},"498ea41c":{"m":130,"g":131},"7c744d13":{"m":130,"g":131},"46b05ef5":{"m":130,"g":131},"beec8eed":{"m":130,"g":131},"b76e303e":{"m":130,"g":131},"80a575e4":{"m":130,"g":131},"4c5074eb":{"m":130,"g":131},"41429a8c":{"m":130,"g":131},"532037df":{"m":130,"g":131},"fa0ca976":{"m":130,"g":131},"b5d39985":{"m":130,"g":131},"2ecee757":{"m":130,"g":131},"6d37e708":{"m":130,"g":131},"c1006fd8":{"m":130,"g":131},"29c6c2ea":{"m":130,"g":131},"eb85fa6d":{"m":130,"g":131},"0e6441b4":{"m":130,"g":131},"88d1bab5":{"m":130,"g":131},"922756aa":{"m":130,"g":131},"d8faf2f3":{"m":130,"g":131},"7dfcc781":{"m":130,"g":131},"7f3308bc":{"m":130,"g":131},"fdc2ef58":{"m":130,"g":131},"1808df48":{"m":130,"g":131},"b01fc161":{"m":130,"g":131},"788628b5":{"m":130,"g":131},"11d33c0e":{"m":130,"g":131},"441420e1":{"m":130,"g":131},"29a2d4b5":{"m":130,"g":131},"079ac237":{"m":130,"g":131},"84280784":{"m":130,"g":131},"af35023e":{"m":130,"g":131},"cb8df87f":{"m":130,"g":131},"e3ab23c1":{"m":130,"g":131},"70d25873":{"m":130,"g":131},"894c0dc5":{"m":130,"g":131},"d6c49019":{"m":130,"g":131},"fa78c44a":{"m":130,"g":131},"654a78f9":{"m":130,"g":131},"f90b4004":{"m":130,"g":131},"78647e08":{"m":130,"g":131},"46f21a59":{"m":130,"g":131},"b2b09f5f":{"m":130,"g":131},"04df80a9":{"m":130,"g":131},"4f73e53d":{"m":130,"g":131},"16ff892c":{"m":130,"g":131},"df026bb1":{"m":130,"g":131},"d42c167b":{"m":130,"g":131},"38815105":{"m":130,"g":131},"9d823402":{"m":130,"g":131},"8ab5d8b4":{"m":130,"g":131},"03575ce3":{"m":130,"g":131},"80518bea":{"m":130,"g":131},"7e78825d":{"m":130,"g":131},"5bbd83a2":{"m":130,"g":131},"abf6272b":{"m":130,"g":131},"46d7b35e":{"m":130,"g":131},"20aad5b5":{"m":130,"g":131},"974c562a":{"m":130,"g":131},"aca0d01d":{"m":130,"g":131},"dc163502":{"m":130,"g":131},"24903b88":{"m":130,"g":131},"443d7bcd":{"m":130,"g":131},"16d8de22":{"m":130,"g":131},"d122e324":{"m":130,"g":131},"96cc1083":{"m":130,"g":131},"77512ae0":{"m":130,"g":131},"c233e9d7":{"m":130,"g":131},"58ac3f31":{"m":130,"g":131},"93452a82":{"m":130,"g":131},"65c8568c":{"m":130,"g":131},"4bcc5879":{"m":130,"g":131},"d5ea8c71":{"m":130,"g":131},"42271376":{"m":130,"g":131},"84e0abb7":{"m":130,"g":131},"043f1317":{"m":130,"g":131},"f764c691":{"m":130,"g":131},"92205407":{"m":130,"g":131},"7d1a130c":{"m":130,"g":131},"5c8bd8b5":{"m":132,"g":133},"b05b346a":{"m":132,"g":133},"5c961756":{"m":132,"g":133},"c5f1e861":{"m":132,"g":133},"2c4d376d":{"m":132,"g":133},"d0f756ae":{"m":132,"g":133},"6f99dc97":{"m":132,"g":133},"cd1c1fa5":{"m":132,"g":133},"5b0872d2":{"m":132,"g":133},"ba88f1ca":{"m":132,"g":133},"60560c07":{"m":132,"g":133},"ca114421":{"m":132,"g":133},"543d62d1":{"m":132,"g":133},"27032cec":{"m":132,"g":133},"5d804a37":{"m":132,"g":133},"388018a5":{"m":132,"g":133},"fca8e88f":{"m":132,"g":133},"45eeeb9a":{"m":132,"g":133},"a368df28":{"m":132,"g":133},"f85460fb":{"m":132,"g":133},"e52cf30e":{"m":132,"g":133},"a076d75e":{"m":132,"g":133},"8348725d":{"m":132,"g":133},"28566241":{"m":132,"g":133},"b62fe850":{"m":132,"g":133},"624725cb":{"m":132,"g":133},"1a96e664":{"m":132,"g":133},"32829b16":{"m":132,"g":133},"e54307f2":{"m":132,"g":133},"8642dbe4":{"m":132,"g":133},"7dcad45c":{"m":132,"g":133},"7c985331":{"m":132,"g":133},"bd7824b2":{"m":132,"g":133},"312df1d6":{"m":132,"g":133},"25e97380":{"m":132,"g":133},"b6523a4f":{"m":132,"g":133},"c51efb8b":{"m":132,"g":133},"bcc5483e":{"m":132,"g":133},"ccf26027":{"m":132,"g":133},"a4992873":{"m":132,"g":133},"c97ce391":{"m":132,"g":133},"c032b559":{"m":132,"g":133},"da9b801e":{"m":132,"g":133},"0e54a695":{"m":132,"g":133},"e99ee0c6":{"m":132,"g":133},"c1bd5ee8":{"m":132,"g":133},"ef1ab230":{"m":132,"g":133},"d6598737":{"m":132,"g":133},"6c5ebc0e":{"m":132,"g":133},"5b5571a8":{"m":132,"g":133},"1698c234":{"m":132,"g":133},"3d82c0f1":{"m":132,"g":133},"617e9b3b":{"m":132,"g":133},"83e35a7c":{"m":132,"g":133},"2543666c":{"m":132,"g":133},"f732f8ea":{"m":132,"g":133},"503880db":{"m":132,"g":133},"b8cfa02c":{"m":132,"g":133},"d85fecb5":{"m":132,"g":133},"6634f67b":{"m":132,"g":133},"5eccaf77":{"m":132,"g":133},"d7f6320b":{"m":132,"g":133},"766476f5":{"m":132,"g":133},"12b7a4fa":{"m":132,"g":133},"02f1e81e":{"m":132,"g":133},"908c7186":{"m":132,"g":133},"03836d85":{"m":132,"g":133},"87dbdddc":{"m":132,"g":133},"56e5c074":{"m":132,"g":133},"6c9c8da6":{"m":132,"g":133},"21028b55":{"m":132,"g":133},"b0a25d09":{"m":132,"g":133},"793c98af":{"m":132,"g":133},"b1cbfce6":{"m":132,"g":133},"b0f531ad":{"m":132,"g":133},"01835998":{"m":132,"g":133},"4285e99d":{"m":132,"g":133},"f0774368":{"m":132,"g":133},"5e8f544d":{"m":132,"g":133},"c8d74feb":{"m":132,"g":133},"cbc7dcda":{"m":132,"g":133},"a6dc7d29":{"m":132,"g":133},"390406c4":{"m":132,"g":131},"0c63fb94":{"m":132,"g":131},"9ad02b79":{"m":132,"g":131},"18bd8e8d":{"m":132,"g":131},"036e64da":{"m":132,"g":131},"7c6fb3aa":{"m":132,"g":131},"8b0b6a45":{"m":132,"g":131},"55504df2":{"m":132,"g":131},"73df7a4e":{"m":132,"g":131},"8b98bb76":{"m":132,"g":131},"15bc8cbd":{"m":132,"g":131},"9496f12d":{"m":132,"g":131},"ab004879":{"m":132,"g":131},"6ec77680":{"m":132,"g":131},"13680e55":{"m":132,"g":131},"98c430e1":{"m":132,"g":131},"fe7f91ef":{"m":132,"g":131},"cef5ba65":{"m":132,"g":131},"53d17088":{"m":132,"g":131},"f0e948a0":{"m":132,"g":131},"9a426fc5":{"m":132,"g":131},"66772aa2":{"m":132,"g":131},"b6263344":{"m":132,"g":131},"0f8bd55f":{"m":132,"g":131},"da3dc497":{"m":132,"g":131},"817daba0":{"m":132,"g":131},"af60cad0":{"m":132,"g":131},"ce4e836b":{"m":132,"g":131},"0e0b0c05":{"m":132,"g":131},"e6f0ddda":{"m":132,"g":131},"af20657c":{"m":132,"g":131},"08da4c26":{"m":132,"g":131},"ef3f8c97":{"m":132,"g":131},"e5201bda":{"m":132,"g":131},"60d36e7b":{"m":132,"g":131},"6f657070":{"m":132,"g":131},"c106b54b":{"m":132,"g":131},"119fd956":{"m":132,"g":131},"eac5b664":{"m":132,"g":131},"07404d76":{"m":132,"g":131},"93043f7b":{"m":132,"g":131},"edde5e5d":{"m":132,"g":131},"232982a0":"m134","9e88c0a2":"m134","b7e0d54e":"m134","5dde0a57":"m134","9e5ab903":"m134","98225be6":{"m":134,"g":135},"94bcc19b":{"m":134,"g":135},"db3821a9":{"m":134,"g":135},"8c5d91b8":{"m":134,"g":135},"6a3e7092":{"m":134,"g":135},"c2601f0d":{"m":134,"g":135},"1048803c":{"m":134,"g":135},"8a84b1e7":{"m":134,"g":135},"5e20e7a6":{"m":134,"g":135},"3946dad6":{"m":134,"g":135},"ac03ec08":{"m":134,"g":135},"94164646":{"m":134,"g":135},"5fb734f1":{"m":134,"g":135},"60f1ca69":{"m":134,"g":135},"9e263c21":{"m":134,"g":135},"0d003e34":{"m":134,"g":135},"f253f43c":{"m":134,"g":135},"1e453201":{"m":134,"g":135},"26e17f90":{"m":134,"g":135},"3de23274":{"m":134,"g":135},"f4ec6f8e":{"m":134,"g":135},"9c4eb460":{"m":134,"g":135},"c2e0913e":{"m":134,"g":135},"8c6f865a":{"m":134,"g":135},"269aa27b":{"m":134,"g":135},"f39382c6":{"m":134,"g":135},"2ff289e2":{"m":134,"g":135},"b2a3f055":{"m":134,"g":135},"684e148e":{"m":134,"g":135},"f44c4b37":{"m":134,"g":135},"7380ec9d":{"m":134,"g":135},"ac78f96e":{"m":134,"g":135},"278012ca":{"m":134,"g":135},"7f587998":{"m":134,"g":135},"162d1cf9":{"m":134,"g":135},"8e08207c":{"m":134,"g":135},"c31f6272":{"m":134,"g":135},"b5d9fc87":{"m":134,"g":135},"88f3de25":{"m":134,"g":135},"24616c52":{"m":134,"g":135},"d48723b7":{"m":134,"g":135},"f3d73b01":{"m":134,"g":135},"4ab66d95":{"m":134,"g":135},"de2799f3":{"m":134,"g":135},"f784cbfa":{"m":134,"g":135},"09733090":{"m":134,"g":135},"a44d0079":{"m":134,"g":135},"8305dc17":{"m":134,"g":135},"ec8c831d":{"m":134,"g":135},"f13949e5":{"m":134,"g":135},"c236a3fd":{"m":134,"g":135},"41a1d16b":{"m":134,"g":135},"9884c9fd":{"m":134,"g":135},"a435f55d":{"m":134,"g":135},"2ec6fa3c":{"m":134,"g":135},"c58a573a":{"m":134,"g":135},"b840d6aa":{"m":134,"g":135},"ef4b3c0e":{"m":134,"g":135},"6f9d0a89":{"m":134,"g":135},"d7a3336e":{"m":134,"g":135},"7d02c8e5":{"m":134,"g":135},"e6d5a213":{"m":134,"g":135},"9f8e2307":{"m":134,"g":135},"d90f9bfc":{"m":134,"g":135},"3881bc8d":{"m":134,"g":135},"208e6a9d":{"m":134,"g":135},"be3828a1":{"m":134,"g":135},"5969be2f":{"m":134,"g":135},"8fab4895":{"m":134,"g":135},"7ccaec64":{"m":134,"g":135},"c457aad5":{"m":134,"g":135},"c7e7bfa3":{"m":134,"g":135},"bf90ea9c":{"m":134,"g":135},"26c50912":{"m":134,"g":135},"325a4c19":{"m":134,"g":135},"0294844f":{"m":134,"g":135},"0e536600":{"m":134,"g":135},"d70c2655":{"m":134,"g":135},"656f4d69":{"m":134,"g":135},"8e43980e":{"m":134,"g":135},"474a4699":{"m":134,"g":135},"b4a00ed2":{"m":134,"g":135},"f55d608c":{"m":134,"g":135},"349ce2dd":{"m":134,"g":135},"183b6519":{"m":134,"g":135},"2af955e1":{"m":134,"g":135},"0cd2b719":{"m":134,"g":135},"39d56196":{"m":134,"g":135},"41addd2e":{"m":134,"g":135},"5c393e81":{"m":134,"g":135},"3645ed0f":{"m":134,"g":135},"ca740a41":{"m":134,"g":135},"0e25aa43":{"m":134,"g":135},"60a230b1":{"m":134,"g":135},"aa89c6a7":{"m":134,"g":135},"171912a9":{"m":134,"g":135},"faecd37e":{"m":134,"g":135},"9ad546d7":{"m":134,"g":135},"29ce7b36":{"m":134,"g":135},"a8380ded":{"m":134,"g":135},"4edee695":{"m":134,"g":135},"67caea6f":{"m":134,"g":135},"cd3289c7":{"m":134,"g":135},"acddb8e0":{"m":134,"g":135},"2ec57cef":{"m":134,"g":135},"988b14ca":{"m":134,"g":135},"93495dca":{"m":134,"g":135},"886e0383":{"m":134,"g":135},"01bd0d3e":{"m":134,"g":135},"43e1bbc0":{"m":134,"g":135},"59b12996":{"m":134,"g":135},"b7091496":{"m":134,"g":135},"51dbdb22":{"m":134,"g":135},"8dc6f0fc":{"m":134,"g":135},"cf34d0ab":{"m":134,"g":135},"3778c2fc":{"m":134,"g":135},"5d421db8":{"m":134,"g":135},"fe3d47fc":{"m":134,"g":135},"ef92b4eb":{"m":134,"g":135},"73c0c66f":{"m":134,"g":135},"a1e9b4ed":{"m":134,"g":135},"e75657c8":{"m":134,"g":135},"c28c536c":{"m":134,"g":135},"086813ae":{"m":134,"g":135},"3fd232ad":{"m":134,"g":135},"cb181295":{"m":134,"g":135},"a91e072f":{"m":134,"g":135},"0271fc34":{"m":134,"g":135},"68bece8c":{"m":134,"g":135},"7b7e357f":{"m":134,"g":135},"2f66b067":{"m":134,"g":135},"48051181":{"m":134,"g":135},"f2ccc442":{"m":134,"g":135},"caa95c7e":{"m":134,"g":135},"9d878c1f":{"m":134,"g":135},"8087ef12":{"m":134,"g":135},"6ef543f9":{"m":134,"g":135},"a3559119":{"m":134,"g":135},"f3ba7116":{"m":134,"g":135},"5c243ba5":{"m":134,"g":135},"b6702d72":{"m":134,"g":135},"c1256727":{"m":134,"g":135},"bb9e6cdf":{"m":134,"g":135},"de03b0cd":{"m":134,"g":135},"2a8a7856":{"m":134,"g":135},"f4e835af":{"m":134,"g":135},"e6ce16a4":{"m":134,"g":135},"de2f2880":{"m":134,"g":135},"a89e85e7":{"m":134,"g":135},"cbf9f134":{"m":134,"g":135},"eb3da9c1":{"m":134,"g":135},"72a980c6":{"m":134,"g":135},"ccf2330b":{"m":134,"g":135},"b9af8d2e":{"m":134,"g":135},"10a9573e":{"m":134,"g":135},"49ab72f8":{"m":134,"g":135},"8865424f":{"m":134,"g":135},"b311c43d":{"m":134,"g":135},"0c39730b":{"m":134,"g":135},"45adad37":{"m":134,"g":135},"1ba897f3":{"m":134,"g":135},"38dd4fbb":{"m":134,"g":135},"ecd2d09a":{"m":134,"g":135},"92ddc468":{"m":134,"g":135},"5454d2a7":{"m":134,"g":133},"17b38f88":{"m":134,"g":133},"ae434f78":{"m":134,"g":133},"643aeefe":{"m":134,"g":133},"186a56f6":{"m":134,"g":133},"17e65466":{"m":134,"g":133},"370bd27f":{"m":134,"g":133},"b27b5a83":{"m":134,"g":133},"2f7c6292":{"m":134,"g":133},"2fb31605":{"m":134,"g":133},"8bf7f240":{"m":134,"g":133},"2c5679f3":{"m":134,"g":133},"159b1283":{"m":134,"g":133},"9338f63f":{"m":134,"g":133},"8196998a":{"m":134,"g":133},"aa21c6e3":{"m":134,"g":133},"fd4a558e":{"m":134,"g":133},"b3b818fd":{"m":134,"g":133},"d5fbbfd9":{"m":134,"g":133},"e245cac0":{"m":134,"g":133},"c6a6ba43":{"m":134,"g":133},"d6108166":{"m":134,"g":133},"ddb3970e":{"m":134,"g":133},"f65fa047":{"m":134,"g":133},"7e027691":{"m":134,"g":133},"cb719c74":{"m":134,"g":133},"ff903a7e":{"m":134,"g":133},"eee3700d":{"m":134,"g":133},"e254cdf3":{"m":134,"g":133},"dfb53574":{"m":134,"g":133},"96655749":{"m":134,"g":133},"ac320a6f":{"m":134,"g":133},"6292c244":{"m":134,"g":133},"5f5a5677":{"m":134,"g":133},"3bf07c68":{"m":134,"g":133},"e7b09efc":{"m":134,"g":133},"99d3bcdf":{"m":134,"g":133},"4d64f150":{"m":134,"g":133},"aef7ca7c":{"m":134,"g":133},"fe712aa3":{"m":134,"g":133},"846953d9":{"m":134,"g":133},"5c64a20d":{"m":134,"g":133},"cf817376":{"m":134,"g":133},"aa6ac966":{"m":134,"g":133},"0d0367e9":{"m":134,"g":133},"bd572360":{"m":134,"g":133},"5f3a47d8":{"m":134,"g":133},"80ae2229":{"m":134,"g":133},"705287b2":{"m":134,"g":133},"76284653":{"m":134,"g":133},"dd620987":{"m":134,"g":133},"53f974b9":{"m":134,"g":133},"6a5764a7":{"m":134,"g":133},"291f11ae":{"m":134,"g":133},"d7301c89":{"m":134,"g":133},"758b9067":{"m":134,"g":133},"ffc23ef8":{"m":134,"g":133},"b3f83cc1":{"m":134,"g":133},"c15fa1c5":{"m":134,"g":133},"fa296698":{"m":134,"g":133},"f9dd90ac":{"m":134,"g":133},"66902e0f":{"m":134,"g":133},"e50f356f":{"m":134,"g":133},"ac42797c":{"m":134,"g":133},"883747ce":{"m":134,"g":133},"989d4b30":{"m":134,"g":133},"bc3ca300":{"m":134,"g":133},"061f41af":{"m":134,"g":133},"82f1d615":{"m":134,"g":133},"5e1a495c":{"m":134,"g":133},"34013d9d":{"m":134,"g":133},"77597167":{"m":134,"g":133},"3c882db3":{"m":134,"g":133},"b736a152":{"m":134,"g":133},"6984837d":{"m":134,"g":133},"2142881b":{"m":134,"g":133},"d77f3fcc":{"m":134,"g":133},"828dec1c":{"m":134,"g":133},"575a49dc":{"m":134,"g":133},"e62e1744":{"m":134,"g":133},"d5431ff8":{"m":134,"g":133},"454a2544":{"m":134,"g":133},"89619a99":{"m":134,"g":133},"cb30d056":{"m":134,"g":133},"677930c2":{"m":134,"g":133},"beae3f96":{"m":134,"g":133},"1167867e":{"m":134,"g":133},"ad7f35fb":{"m":134,"g":133},"f4100732":{"m":134,"g":133},"468931b5":{"m":134,"g":133},"796969ca":{"m":134,"g":133},"122c2503":{"m":134,"g":133},"a3a55223":{"m":134,"g":133},"0bf95e6d":{"m":134,"g":133},"a92de891":{"m":134,"g":133},"b9d78605":{"m":134,"g":133},"1354063a":{"m":134,"g":133},"254de6d2":{"m":134,"g":133},"350fbbf4":{"m":134,"g":133},"a3912667":{"m":134,"g":133},"e1dcd0df":{"m":134,"g":133},"393e2f9b":{"m":134,"g":133},"8766a1dd":{"m":134,"g":133},"1d9ba2ce":{"m":134,"g":133},"ef001fb8":{"m":134,"g":133},"c69c1c4f":{"m":134,"g":133},"bed301a5":{"m":134,"g":133},"8fe3e374":{"m":134,"g":133},"60143655":{"m":134,"g":133},"fc05acc2":{"m":134,"g":133},"1ed94668":{"m":134,"g":133},"d7fbe73b":{"m":134,"g":133},"42bff706":{"m":134,"g":133},"26704c23":{"m":134,"g":133},"43b7c174":{"m":134,"g":133},"96740d69":{"m":134,"g":133},"9a3bdf2c":{"m":134,"g":133},"4b351f6b":{"m":134,"g":133},"7fa4906f":{"m":134,"g":133},"050f108c":{"m":134,"g":133},"47cdb65a":{"m":134,"g":133},"1d90b194":{"m":134,"g":133},"537ef18d":{"m":134,"g":133},"69412ccb":{"m":134,"g":133},"d3885d4b":{"m":134,"g":133},"0a346d3b":{"m":134,"g":133},"bc18cb86":{"m":134,"g":133},"bee8ac5b":{"m":134,"g":133},"41bd76e1":{"m":134,"g":133},"c6ca1b3a":{"m":134,"g":133},"8999ce75":{"m":134,"g":133},"dce2ed44":{"m":134,"g":133},"019517a3":{"m":134,"g":133},"1f1f05a8":{"m":134,"g":133},"165f5c04":{"m":134,"g":133},"3e01f3a5":{"m":134,"g":133},"51e2eaa4":{"m":134,"g":133},"6468cb58":{"m":134,"g":133},"e220da17":{"m":134,"g":133},"b82c7a0a":{"m":134,"g":133},"c0f9b519":{"m":134,"g":133},"5529ab58":{"m":134,"g":133},"74a3349b":{"m":134,"g":133},"b9ebf0ed":{"m":134,"g":133},"3c116d5e":{"m":134,"g":133},"71a60288":{"m":134,"g":133},"61405b3d":{"m":134,"g":133},"0adfc42b":{"m":134,"g":133},"ba72e759":{"m":134,"g":133},"6afc5d49":{"m":134,"g":133},"2ee6c810":{"m":134,"g":133},"bd16244d":{"m":134,"g":133},"b5eb0214":{"m":134,"g":133},"d72e908b":{"m":134,"g":133},"50cad014":{"m":134,"g":133},"ef908aeb":{"m":134,"g":133},"9d0347b3":{"m":134,"g":133},"05eb0bcc":{"m":134,"g":133},"5dccd9bd":{"m":134,"g":133},"933cef16":{"m":134,"g":133},"241ae17b":{"m":134,"g":133},"2c5a4460":{"m":134,"g":133},"f3705b01":{"m":134,"g":133},"5a0ad731":{"m":134,"g":133},"5045aa34":{"m":134,"g":133},"ff1e2ce2":{"m":134,"g":133},"1e582488":{"m":134,"g":133},"46be74b4":{"m":134,"g":133},"1c658026":{"m":134,"g":133},"ba410808":{"m":134,"g":133},"89512029":{"m":134,"g":133},"92e6b3c3":{"m":134,"g":133},"6559e43f":{"m":134,"g":133},"af780c59":{"m":134,"g":133},"a21aa87e":{"m":134,"g":133},"fb178457":{"m":134,"g":133},"65c09859":{"m":134,"g":133},"4bf06635":{"m":134,"g":133},"f2d64e67":{"m":134,"g":133},"0e869f08":{"m":134,"g":133},"17394092":{"m":134,"g":133},"f228b662":{"m":134,"g":133},"a36142aa":{"m":134,"g":133},"160a06ca":{"m":134,"g":133},"a0985dd5":{"m":134,"g":133},"f6c9db4b":{"m":134,"g":133},"e88e75a9":{"m":134,"g":133},"4b4050e2":{"m":134,"g":133},"b2803ff2":{"m":134,"g":133},"9e0ef04e":{"m":134,"g":133},"216067c0":{"m":134,"g":133},"e72b02db":{"m":134,"g":133},"29e8f7f9":{"m":134,"g":133},"e0963a6c":{"m":134,"g":133},"e0026f7c":{"m":134,"g":133},"9749d3e3":{"m":134,"g":133},"2b0ddf89":{"m":134,"g":133},"17e81c75":{"m":134,"g":133},"88a405cc":{"m":134,"g":133},"602fe3b2":{"m":134,"g":133},"ad9616f1":{"m":134,"g":133},"c5f4e20f":{"m":134,"g":133},"2c196f95":{"m":134,"g":133},"9a7641d7":{"m":134,"g":133},"793c96c3":{"m":134,"g":133},"d1f00632":{"m":134,"g":133},"4792d1f4":{"m":134,"g":133},"fea2d521":{"m":134,"g":133},"56d12b4a":{"m":134,"g":133},"374ad4cc":{"m":134,"g":133},"8b0a68f1":{"m":134,"g":133},"70607e55":{"m":134,"g":133},"ee1ca51d":{"m":134,"g":133},"ef7c29ac":{"m":134,"g":133},"58c840db":{"m":134,"g":133},"3d42b7e7":{"m":134,"g":133},"9970ee34":{"m":134,"g":133},"9e7656be":{"m":134,"g":133},"9d4f066f":{"m":134,"g":133},"41683536":{"m":134,"g":133},"891ee822":{"m":134,"g":133},"8fa3dc36":{"m":134,"g":133},"d20699a3":{"m":134,"g":133},"169a75df":{"m":134,"g":133},"4128d4f5":{"m":134,"g":133},"011d8d89":{"m":134,"g":133},"5290cef9":{"m":134,"g":133},"726fe3e7":{"m":134,"g":133},"5d087891":{"m":134,"g":133},"8451e227":{"m":134,"g":133},"53e15194":{"m":134,"g":133},"d747147a":{"m":134,"g":133},"0c002207":{"m":134,"g":133},"eeb2b9b2":{"m":134,"g":133},"c4aed389":{"m":134,"g":133},"9d04b570":{"m":134,"g":133},"3e690cce":{"m":134,"g":133},"b12b40de":{"m":134,"g":133},"6c4bf8a0":{"m":134,"g":133},"533851fb":{"m":134,"g":133},"0071fe9c":{"m":134,"g":133},"ffa7e035":{"m":134,"g":133},"712f44ee":{"m":134,"g":133},"8c34e181":{"m":134,"g":133},"e9abb525":{"m":134,"g":133},"88859433":{"m":134,"g":133},"feb8e30b":{"m":134,"g":133},"45a959d3":{"m":134,"g":133},"cdce5163":{"m":134,"g":133},"79ab57bd":{"m":134,"g":133},"4b8901ac":{"m":134,"g":133},"435d1c83":{"m":134,"g":133},"2bdbaef1":{"m":134,"g":133},"31d48d7f":{"m":134,"g":133},"0129c911":{"m":134,"g":133},"03f9eb25":{"m":134,"g":133},"7ec678eb":{"m":134,"g":133},"9d64a7b2":{"m":134,"g":133},"46ad4b98":{"m":134,"g":133},"71cb9037":{"m":134,"g":133},"c8c64876":{"m":134,"g":133},"0861dca8":{"m":134,"g":133},"da58df6b":{"m":134,"g":133},"49237e26":{"m":134,"g":133},"d92c1f8c":{"m":134,"g":133},"93070586":{"m":134,"g":133},"ccc8f3b2":{"m":134,"g":133},"a4c76281":{"m":134,"g":133},"28a19e49":{"m":134,"g":133},"0261c4af":{"m":134,"g":133},"8ac350f3":{"m":134,"g":133},"99401e7b":{"m":134,"g":133},"66824751":{"m":134,"g":133},"f95729b0":{"m":134,"g":133},"9f4ed93d":{"m":134,"g":133},"ecb401ed":{"m":134,"g":133},"b399e3ac":{"m":134,"g":133},"e27635a0":{"m":134,"g":133},"272c5fe4":{"m":134,"g":133},"3c8dc448":{"m":134,"g":133},"5e96beb3":{"m":134,"g":133},"9327482b":{"m":134,"g":133},"36fcf71f":{"m":134,"g":133},"6292d971":{"m":134,"g":133},"c8434195":{"m":134,"g":133},"3e4d431a":{"m":134,"g":133},"538e733e":{"m":134,"g":133},"22587bc0":{"m":134,"g":133},"02d24244":{"m":134,"g":133},"1da5cd63":{"m":134,"g":133},"a9a2cdd8":{"m":134,"g":133},"4733fcff":{"m":134,"g":133},"3ffa2604":{"m":134,"g":133},"61f362c6":{"m":134,"g":133},"e7157c9b":{"m":134,"g":133},"30da2f05":{"m":134,"g":133},"49016931":{"m":134,"g":133},"3d484be5":{"m":134,"g":133},"c0d94440":{"m":134,"g":133},"1dedb638":{"m":134,"g":133},"7bc8b153":{"m":134,"g":133},"9003a436":{"m":134,"g":133},"3518b331":{"m":134,"g":133},"b098b1ae":{"m":134,"g":133},"abd3e048":{"m":134,"g":133},"92c29d43":{"m":134,"g":133},"bf643814":{"m":134,"g":133},"f03bfa4c":{"m":134,"g":133},"89ad3908":{"m":134,"g":133},"2ea844ec":{"m":134,"g":133},"1e2d7538":{"m":134,"g":133},"d16ff357":{"m":134,"g":133},"af49e302":{"m":134,"g":133},"01b955ac":{"m":134,"g":133},"16e6bc20":{"m":134,"g":133},"37250764":{"m":134,"g":133},"7b9156c7":{"m":134,"g":133},"21cfebac":{"m":134,"g":133},"702426b0":{"m":134,"g":133},"9e9a6169":{"m":134,"g":133},"bd9c3a47":{"m":134,"g":133},"1e641ee4":{"m":134,"g":133},"fb96669f":{"m":134,"g":133},"3912ee49":{"m":134,"g":133},"1ab9b8e0":{"m":134,"g":133},"e61dabf5":{"m":134,"g":133},"36e7c8c5":{"m":134,"g":133},"037c3982":{"m":134,"g":133},"62b3fdae":{"m":134,"g":133},"1cd0c3bf":{"m":134,"g":133},"2c899431":{"m":134,"g":133},"4513f549":{"m":134,"g":133},"c9690307":{"m":134,"g":133},"4449c170":{"m":134,"g":133},"5ca962ce":{"m":134,"g":133},"bab20a84":{"m":134,"g":133},"0e4108ba":{"m":134,"g":133},"99cb2ed9":{"m":134,"g":133},"0612175c":{"m":134,"g":133},"8c96fcda":{"m":134,"g":133},"ea7c69ce":{"m":134,"g":133},"8102e36b":{"m":134,"g":133},"3f048217":{"m":134,"g":133},"b11af135":{"m":134,"g":133},"f9bceea0":{"m":134,"g":133},"997ea57e":{"m":134,"g":133},"47633c19":{"m":134,"g":133},"64b5c3ab":{"m":134,"g":133},"d277a86d":{"m":134,"g":133},"9acb21ae":{"m":134,"g":133},"a9ce1623":{"m":134,"g":133},"6f0c77d7":{"m":134,"g":133},"e3f51e82":{"m":134,"g":133},"fdfabb7a":{"m":134,"g":133},"19c16748":{"m":134,"g":133},"5c75907e":{"m":134,"g":133},"4ea36422":{"m":134,"g":133},"f50af32d":{"m":134,"g":133},"72952919":{"m":134,"g":133},"2ae5bed1":{"m":134,"g":133},"54df514b":{"m":134,"g":133},"681c68cf":{"m":134,"g":133},"74ea45cc":{"m":134,"g":133},"0fa044ad":{"m":134,"g":133},"7d8e42c9":{"m":134,"g":133},"6abdf73f":{"m":134,"g":133},"20ce9938":{"m":134,"g":133},"168a31eb":{"m":134,"g":133},"69cfb17b":{"m":134,"g":133},"a7a4b175":{"m":134,"g":133},"fdc93b01":{"m":134,"g":133},"fd37cc5d":{"m":134,"g":133},"96705514":{"m":134,"g":133},"ab3ffd1c":{"m":134,"g":133},"2285afff":{"m":134,"g":133},"a81cc1b8":{"m":134,"g":133},"3134d2b2":{"m":134,"g":133},"06b58c5d":{"m":134,"g":133},"ea07a283":{"m":134,"g":133},"ea91a720":{"m":134,"g":133},"5d9c6bac":{"m":134,"g":133},"e048ee90":{"m":134,"g":133},"ed52d01b":{"m":134,"g":133},"90e7d4f7":{"m":134,"g":133},"c20d43d2":{"m":134,"g":133},"d977dd2e":{"m":134,"g":133},"0c23331e":{"m":134,"g":133},"993278b4":{"m":134,"g":133},"9e9d9107":{"m":134,"g":133},"3b8a824b":{"m":134,"g":133},"f1bbd26f":{"m":134,"g":133},"d36299ad":{"m":134,"g":133},"0e7d7969":{"m":134,"g":133},"80554598":{"m":134,"g":133},"b2e240bc":{"m":134,"g":133},"dcc5f5c0":{"m":134,"g":133},"4eda4194":{"m":134,"g":133},"875f84db":{"m":134,"g":133},"f6031adf":{"m":134,"g":133},"2a39cfe0":{"m":134,"g":133},"bf17e769":{"m":134,"g":133},"9a5d6a84":{"m":134,"g":133},"31c23e5f":{"m":134,"g":133},"06617a9e":{"m":134,"g":133},"e79ca959":{"m":134,"g":133},"9d3b411c":{"m":134,"g":133},"05325db3":{"m":134,"g":133},"01e3b3f3":{"m":134,"g":133},"86988674":{"m":134,"g":133},"29139654":{"m":134,"g":133},"665cb020":{"m":134,"g":133},"71602838":{"m":134,"g":133},"77873343":{"m":134,"g":133},"267170bf":{"m":134,"g":133},"313f59ad":{"m":134,"g":133},"487cf81a":{"m":134,"g":133},"df111bc0":{"m":134,"g":133},"6d2b3324":{"m":134,"g":133},"8cc77261":{"m":134,"g":133},"d143b020":{"m":134,"g":133},"9b9d2131":{"m":134,"g":133},"44fd7017":{"m":134,"g":133},"9a56273a":{"m":134,"g":133},"b737a125":{"m":134,"g":133},"1b5e9034":{"m":134,"g":133},"171b442a":{"m":134,"g":133},"4b7b5af3":{"m":134,"g":133},"ec242f51":{"m":134,"g":133},"b2431546":{"m":134,"g":133},"526fd008":{"m":134,"g":133},"56d0ad47":{"m":134,"g":133},"306e5b8d":{"m":134,"g":133},"10c68f62":{"m":134,"g":133},"c7c837cd":{"m":134,"g":133},"3e1e7157":{"m":134,"g":133},"8fa8d9d7":{"m":134,"g":133},"c8cf1caf":{"m":134,"g":133},"d71baa72":{"m":134,"g":133},"82e33170":{"m":134,"g":133},"94e12511":{"m":134,"g":133},"4dabfbc8":{"m":134,"g":133},"d7ed8a8c":{"m":134,"g":133},"c05d3afb":{"m":134,"g":133},"edb172e9":{"m":134,"g":133},"fe6d38d2":{"m":134,"g":133},"dab31e4c":{"m":134,"g":133},"a7fa31ff":{"m":134,"g":133},"bd91f882":{"m":134,"g":133},"b47adb80":{"m":134,"g":133},"b2b5bdb0":{"m":134,"g":133},"22fe5da1":{"m":134,"g":133},"1834401e":{"m":134,"g":133},"198c8ecf":{"m":134,"g":133},"c01b2ee0":{"m":134,"g":133},"8f5adac8":{"m":134,"g":133},"76743a98":{"m":134,"g":133},"8bf10e71":{"m":134,"g":133},"e4873d04":{"m":134,"g":133},"c660d8df":{"m":134,"g":133},"10146af0":{"m":134,"g":133},"b62e7e3b":{"m":134,"g":133},"6ce36b12":{"m":134,"g":133},"e9e7f15e":{"m":134,"g":133},"0aa3dec5":{"m":134,"g":133},"4885f8b9":{"m":134,"g":133},"f832994c":{"m":134,"g":133},"e59435c3":{"m":134,"g":133},"d6bd2d11":{"m":134,"g":133},"9975acf5":{"m":134,"g":133},"70758d45":{"m":134,"g":133},"fd1ebbb0":{"m":134,"g":133},"3d98bd5e":{"m":134,"g":133},"aa3716b2":{"m":134,"g":133},"6107268f":{"m":134,"g":133},"0189f41c":"m136","b6e4893a":"m136","3eb7da53":"m136","cb53ddc9":"m136","0c265321":"m136","6469c964":"m136","71705394":"m136","67589c16":"m136","a702c8f1":"m136","f2ae066a":"m136","0c2993ee":"m136","590969ee":"m136","e6ccb294":"m136","d6ea2c52":"m136","9be2a3a9":"m136","b74a57a8":"m136","95f59c13":"m136","85d9af51":"m136","858f317f":"m136","cf893516":"m136","1fdf5cac":"m136","cda43ffa":"m136","39089854":"m136","b827e9d3":"m136","d725487d":"m136","a95c9f5b":"m136","19089aa4":"m136","4f6f5d25":"m136","458fe5a3":"m136","2ff0880a":"m136","2c1b164a":"m136","bcc6d84f":"m136","a618202f":"m136","7520b929":"m136","e7224e96":"m136","e776239a":"m136","1b97fa76":"m136","236772c0":"m136","0d49b13f":"m136","0a7a2017":"m136","0a9099e1":"m136","0050c476":"m136","8251a74d":"m136","a54d75bf":"m136","3321eb4e":"m136","be5121b4":"m136","c3f9c30f":"m136","54a82179":"m136","aca354bc":"m136","aea57b33":"m136","823a046e":"m136","648aab0c":"m136","1e309030":"m136","60927215":"m136","38c233fd":"m136","20ed3822":"m136","4ecd9afd":"m136","eb38d644":"m136","6ea491e4":"m136","16802fb6":"m136","d97066d2":"m136","ce2d686e":"m136","76b06bee":"m136","612026ad":"m136","55c61642":"m136","91a4cd86":"m136","23d765d1":"m136","db2425a0":"m136","603f386c":"m136","f7a5e425":"m136","d50dcd9b":"m136","e6b7c049":"m136","a3addd62":"m136","6988a0f5":"m136","1b192cf1":"m136","c560e142":"m136","71cb9d03":"m136","8fb45523":"m136","84aef378":"m136","17c04b10":"m136","55c4288b":"m136","9f8b79f1":"m136","7dc3cbe7":"m136","79ddc34c":"m136","09a9d214":"m136","7e40d526":"m136","057b07fc":"m136","91d8c52d":"m136","e9a44ea6":"m136","c1282da2":"m136","71279e31":"m136","1a053a81":"m136","2ea02f06":"m136","20b0523e":"m136","ce8a6ac6":"m136","ebca5879":"m136","cc410a10":"m136","f374623f":"m136","5c022177":"m136","8916b9d0":"m136","a3d9a218":"m136","fb88fb67":"m136","2d72e168":"m136","64946679":"m136","5836324c":"m136","9fe56cd0":"m136","858a4d65":"m136","e619f531":"m136","fc4b932f":"m136","d2105d4a":"m136","84c83905":"m136","ea879c77":"m136","0227db89":"m136","ad1b4e47":"m136","51f147ad":"m136","d3eafc73":"m136","e00b4344":"m136","733de6be":"m136","93433726":"m136","330605cc":"m136","bb6055b4":"m136","4df74eb5":"m136","f3a7c7dc":"m136","2069050d":"m136","1fe0c82f":"m136","a45e0e5d":"m136","f78201f3":"m136","8fd33998":"m136","088758c1":"m136","6d29d8ab":"m136","e499258e":"m136","09491a9b":"m136","7edb0615":"m136","e486a4da":"m136","90399cbc":"m136","53609e5e":"m136","9c253064":"m136","d2c86387":"m136","737a1183":"m136","dc743fe4":"m136","dd99f818":"m136","8ce64aa1":"m136","eb768189":"m136","2cdd4370":"m136","a7b5f75d":"m136","305c1a57":"m136","c824ddd5":"m136","e18e0057":"m136","166396ca":"m136","43779f27":"m136","8b9e9357":"m136","d36f6f04":"m136","b0701f02":"m136","4229de3b":"m136","2e144079":"m136","d2ec128b":"m136","7f8353af":"m136","a7f5677a":"m136","3e968ab3":"m136","b4fce995":"m136","ec9b48ea":"m136","a0467589":"m136","82a1b645":"m136","6f10e17b":"m136","c771933d":"m136","9d8bbd42":"m136","daea5138":"m136","3355b6e2":"m136","a1dd3d48":"m136","d9ed80b9":"m136","7c39ea68":"m136","daa4841e":"m136","968c4f55":"m136","669d309a":"m136","21ee597e":"m136","6ee970a3":"m136","8ec160ed":"m136","2740ed1a":"m136","d44f09ad":"m136","e7dc85c5":"m136","c81bad1b":"m136","0e86de7c":"m136","e3a95077":"m136","146b5fcc":"m136","8b22deef":"m136","69822c72":"m136","8b99af9a":"m136","72e2f70e":"m136","3d72944f":"m136","7dde3438":"m136","77fc4c4a":"m136","655d2c7c":"m136","3f44268f":"m136","f7ec8174":"m136","cd23c2f0":"m136","d1110e1c":"m136","9227d9f6":"m136","4c59782e":"m136","dda35ccb":"m136","d11e2dc6":"m136","16831ab6":"m136","c9a45b7e":"m136","e7df8bdc":"m136","e9979950":"m136","6586f44a":"m136","43fe3a4d":"m136","98096b5e":"m136","68e8d0f6":"m136","000ad422":"m136","9d5f16d4":"m136","7f8a58ff":"m136","6b065298":"m136","c020d300":"m136","4346db5f":"m136","424a3800":"m136","f0918583":"m136","5b1215d9":"m136","aa2b4f76":"m136","b3a3f513":"m136","de94d793":"m136","0d904ef4":"m136","969faaa4":"m136","48c2aca9":"m136","5af84c8a":"m136","feae615b":"m136","e75299a1":"m136","72bacc88":"m136","ba625c2d":"m136","b025cff4":"m136","030496eb":"m136","c86ca128":"m136","cd336945":"m136","c5e363e8":"m136","9479eca7":"m136","a5348eac":"m136","95240402":"m136","2122fea3":"m136","e2c8a50b":"m136","a4825ed5":"m136","afe285f7":"m136","cf25852a":"m136","5938c3b0":"m136","b8806071":"m136","339915ce":"m136","075c5a57":"m136","a0b4ba90":"m136","2a7b67ad":"m136","1d811094":"m136","2ab3ed3e":"m136","250477d2":"m136","af1232b2":"m136","7a869045":"m136","888d7e54":"m136","3cb1fbae":"m136","ba9f6d8f":"m136","740d3c0b":"m136","87165898":"m136","ff3ddb9d":"m136","9d3018f4":"m136","a8348427":"m136","47d485f3":"m136","2b423099":"m136","ae0baefb":"m136","1f0e3d7f":"m136","d3c08fb0":"m136","c6a64e9f":"m136","e0ac559a":"m136","6620548f":"m136","6e158e55":"m136","ed729d22":"m136","fa51b854":"m136","559ff9ec":"m136","7b682de8":"m136","d0092dec":"m136","76f69b77":"m136","9a628744":"m136","53dca74f":"m136","2dadf635":"m136","aab640c9":"m136","2b3791ed":"m136","9f5cd80a":"m136","b1ee75ae":"m136","f44c63ee":"m136","c54c70ab":"m136","aab906a3":"m136","feb39f77":"m136","38b30c7b":"m136","c581b5ed":"m136","a1c48943":"m136","5b7bed7c":"m136","38a88479":"m136","503c3d95":"m136","934ae89a":"m136","cf1426a7":"m136","17cb3c8e":"m136","2f4a6add":"m136","f9fc50ac":"m136","3c16c586":"m136","7b089ae4":"m136","8b5d4263":"m136","b5493f65":"m136","09e2571e":"m136","c0248d6f":"m136","cc25f9df":"m136","cf14feba":"m136","7c25687c":"m136","ff978142":"m136","d112f6a2":"m136","a2c2c09d":"m136","2a9344d3":"m136","78c41758":"m136","206db66f":"m136","5c72be1e":"m136","76d48817":"m136","d1ec93e3":"m136","bdb76b34":"m136","dae6a409":"m136","a0899bdb":"m136","641830c1":"m136","3fd88ea9":"m136","145bd54f":"m136","2d088b85":"m136","3a8b44fe":"m136","aeb480c1":"m136","9fd2358c":"m136","3c358736":"m136","6327dff2":"m136","675acece":"m136","ad201273":"m136","7f393d95":"m136","4b14f622":"m136","94fc26aa":"m136","67b61a4e":"m136","d27f16f3":"m136","c89949bb":"m136","20abaee2":"m136","32a569fb":"m136","3ed3b7ef":"m136","1f9d4795":"m136","e6d40bff":"m136","fbc128a3":"m136","9c64a15a":"m136","e91a7176":"m136","1f0ea4f9":"m136","15da3061":"m136","6406a596":"m136","cec19b56":"m136","ef35d8fe":"m136","08636f72":"m136","a6c29d4c":"m136","84ab32a2":"m136","70667115":"m136","bd1afeb5":"m136","8ef5b905":"m136","76b3c698":"m136","5dcff947":"m136","8eeffbe9":"m136","71e9c31c":"m136","7656d267":"m136","dd24ba90":"m136","068abe7e":"m136","64a31d4b":"m136","2babf88f":"m136","d56d14e5":"m136","0c4e155a":"m136","d6d5c3fd":"m136","e46f7943":"m136","9d4d57db":"m136","41609b52":"m136","ccd0fb32":"m136","87ee6b5e":"m136","f7c1d24b":"m136","cceb5e6a":"m136","c1c13c84":"m136","fcec35dc":"m136","75da784d":"m136","77d35665":"m136","05dfef92":"m136","74602407":{"m":136,"g":135},"f7f5c389":{"m":136,"g":135},"9e3a032a":{"m":136,"g":135},"1bc7aa58":{"m":136,"g":135},"8726d30c":{"m":136,"g":135},"a1b243d7":{"m":136,"g":135},"a9799277":{"m":136,"g":135},"ee71e773":{"m":136,"g":135},"55a8dd00":{"m":136,"g":135},"cf242321":{"m":136,"g":135},"9d03af91":{"m":136,"g":135},"064ae341":{"m":136,"g":135},"49305fa1":{"m":136,"g":135},"4e999404":{"m":136,"g":135},"fbc24886":{"m":136,"g":135},"16880235":{"m":136,"g":135},"cda35611":{"m":136,"g":135},"8a45a9c6":{"m":136,"g":135},"aecd5f5f":{"m":136,"g":135},"05ab110e":{"m":136,"g":135},"c8dc4d2d":{"m":136,"g":135},"82a8d77b":{"m":136,"g":135},"294ff71d":{"m":136,"g":135},"f52ae586":{"m":136,"g":135},"2f8a3634":{"m":136,"g":135},"b6e8a0d8":{"m":136,"g":135},"fb7609f1":{"m":136,"g":135},"d2ea44f7":{"m":136,"g":135},"20ca2c6e":{"m":136,"g":135},"83abecd0":{"m":136,"g":135},"d54f0a10":{"m":136,"g":135},"fb04e7e3":{"m":136,"g":135},"1e5de05e":{"m":136,"g":135},"7dd679cb":{"m":136,"g":135},"3d51ae18":{"m":136,"g":135},"f9c04266":{"m":136,"g":135},"48b8dcd4":{"m":136,"g":135},"41b434a7":{"m":136,"g":135},"4935344f":{"m":136,"g":135},"63cc97f4":{"m":136,"g":135},"ab7d5829":{"m":136,"g":135},"261860e1":{"m":136,"g":135},"154740bd":{"m":136,"g":135},"1c09cbe3":{"m":136,"g":135},"e14f5ec8":{"m":136,"g":135},"8867d248":{"m":136,"g":135},"6b3f93c4":{"m":136,"g":135},"5a5cece5":{"m":136,"g":135},"d566739b":{"m":136,"g":135},"4c46ecde":{"m":136,"g":135},"12a0292b":{"m":136,"g":135},"a08dc5aa":{"m":136,"g":135},"eec7dbd3":{"m":136,"g":135},"109fe03a":{"m":136,"g":135},"bb798a1c":{"m":136,"g":135},"38dc5839":{"m":136,"g":135},"5e867f60":{"m":136,"g":135},"65bed838":{"m":136,"g":135},"156d97b2":{"m":136,"g":135},"24b30f77":{"m":136,"g":135},"3a4767da":{"m":136,"g":135},"6037267f":{"m":136,"g":135},"0241e046":{"m":136,"g":135},"0c474273":{"m":136,"g":135},"3e73e124":{"m":136,"g":135},"f4742558":{"m":136,"g":135},"7385834c":{"m":136,"g":135},"b5a94f8a":{"m":136,"g":135},"c356ed03":{"m":136,"g":135},"ee4d2287":{"m":136,"g":135},"153c69f6":{"m":136,"g":135},"55b79365":{"m":136,"g":135},"7fc12e0b":{"m":136,"g":135},"8729ad5e":{"m":136,"g":135},"e4320573":{"m":136,"g":135},"4d902c82":{"m":136,"g":135},"2ff87231":{"m":136,"g":135},"fd16c91c":{"m":136,"g":135},"32a6540a":{"m":136,"g":135},"62d0280f":{"m":136,"g":135},"d4b717c0":{"m":136,"g":135},"98a107d4":{"m":136,"g":135},"b86bbf84":{"m":136,"g":135},"48381c3b":{"m":136,"g":135},"8bce0853":{"m":136,"g":135},"6b8a9d70":{"m":136,"g":135},"973116e6":{"m":136,"g":135},"820e97d6":{"m":136,"g":135},"4c85f9d0":{"m":136,"g":135},"7d757d6f":{"m":136,"g":135},"5c04088b":{"m":136,"g":135},"f066036c":{"m":136,"g":135},"52de807d":{"m":136,"g":135},"70933f34":{"m":136,"g":135},"ce453fa4":{"m":136,"g":135},"3be1e734":{"m":136,"g":135},"38895a00":{"m":136,"g":135},"d8b81981":{"m":136,"g":135},"913b688f":{"m":136,"g":135},"951d16c8":{"m":136,"g":135},"53846746":{"m":136,"g":135},"534ac384":{"m":136,"g":135},"2a8d5493":{"m":136,"g":135},"90eac38a":{"m":136,"g":135},"badcd028":{"m":136,"g":135},"d874c8bb":{"m":136,"g":135},"9a21d89c":{"m":136,"g":135},"4c9ac856":{"m":136,"g":135},"fb5b71d0":{"m":136,"g":135},"05b54b6d":{"m":136,"g":135},"4f443f44":{"m":136,"g":135},"dce8b060":{"m":136,"g":135},"399d5283":{"m":136,"g":135},"2e0527dd":{"m":136,"g":135},"6beb50d6":{"m":136,"g":135},"18e2ef09":{"m":136,"g":135},"3271e0e7":{"m":136,"g":135},"d415d22d":{"m":136,"g":135},"a49b9a64":{"m":136,"g":135},"53497642":{"m":136,"g":135},"d57d8e7e":{"m":136,"g":135},"0cbd8f32":{"m":136,"g":135},"6e3fff13":{"m":136,"g":135},"2210155a":{"m":136,"g":135},"a3656cbb":{"m":136,"g":135},"21da2dc1":{"m":136,"g":135},"ed307a40":{"m":136,"g":135},"02722b91":{"m":136,"g":135},"f959250f":{"m":136,"g":135},"95934379":{"m":136,"g":135},"bc2f40be":{"m":136,"g":135},"2724b110":{"m":136,"g":135},"176266f3":{"m":136,"g":135},"5e5b1183":{"m":136,"g":135},"9bf76c11":{"m":136,"g":135},"1d7ad4af":{"m":136,"g":135},"fba785c4":{"m":136,"g":135},"5cfa901b":{"m":136,"g":135},"f27c6cdc":{"m":136,"g":135},"5097e1e8":{"m":136,"g":135},"861a35fb":{"m":136,"g":135},"6ffe1fc0":{"m":136,"g":135},"73398e22":{"m":136,"g":135},"17958c5f":{"m":136,"g":135},"3aa11ca7":{"m":136,"g":135},"7be1a8c7":{"m":136,"g":135},"84d13c54":{"m":136,"g":135},"c80c0e0f":{"m":136,"g":135},"c105a312":{"m":136,"g":135},"4cf2bbd0":{"m":136,"g":135},"d7b706be":{"m":136,"g":135},"4a9537a4":{"m":136,"g":135},"ca922d4b":{"m":136,"g":135},"402a0bd6":{"m":136,"g":135},"76c71d1d":{"m":136,"g":135},"2d02c150":{"m":136,"g":135},"b98bd9a5":{"m":136,"g":135},"ce694b2b":{"m":136,"g":135},"6c0fb189":{"m":136,"g":135},"9a9f996f":{"m":136,"g":135},"4221b7c5":{"m":136,"g":135},"c371df2f":{"m":136,"g":135},"1751c75b":{"m":136,"g":135},"23849eba":{"m":136,"g":135},"51541404":{"m":136,"g":135},"45ef8344":{"m":136,"g":135},"5a2b1ed4":{"m":136,"g":135},"454dc9e2":{"m":136,"g":135},"1e41069a":{"m":136,"g":135},"2b4d6d81":{"m":136,"g":135},"a3914e3b":{"m":136,"g":135},"130f60ee":{"m":136,"g":135},"4397cda7":{"m":136,"g":135},"d56fd10c":{"m":136,"g":135},"abb06be9":{"m":136,"g":135},"c35eb0fd":{"m":136,"g":135},"7f35c46e":{"m":136,"g":135},"da2f8cc3":{"m":136,"g":135},"4308c25b":{"m":136,"g":135},"4d737db8":{"m":136,"g":135},"9d6029fb":{"m":136,"g":135},"dcfb92dd":{"m":136,"g":135},"bebd625b":{"m":136,"g":135},"b12258bf":{"m":136,"g":135},"012dc586":{"m":136,"g":135},"7f6a678f":{"m":136,"g":135},"399ca037":{"m":136,"g":135},"d93f37a6":{"m":136,"g":135},"e6fe092d":{"m":136,"g":135},"f02d8221":{"m":136,"g":135},"10174e11":{"m":136,"g":135},"2138ff48":{"m":136,"g":135},"e53160bb":{"m":136,"g":135},"4ea6a11c":{"m":136,"g":135},"2181bc9e":{"m":136,"g":135},"520c048d":{"m":136,"g":135},"561a3e04":{"m":136,"g":135},"f84487af":{"m":136,"g":135},"07827047":{"m":136,"g":135},"c63e9cb2":{"m":136,"g":135},"12cde0df":{"m":136,"g":135},"a7fd8108":{"m":136,"g":135},"ca80c19b":{"m":136,"g":135},"1e7b3264":{"m":136,"g":135},"e267ca0b":{"m":136,"g":135},"9a8ba3c1":{"m":136,"g":135},"0fee6bc6":{"m":136,"g":135},"0ff3747c":{"m":136,"g":135},"f8411ded":{"m":136,"g":135},"249c3563":{"m":136,"g":135},"12df1660":{"m":136,"g":135},"55d112dc":{"m":136,"g":135},"a1ed247f":{"m":136,"g":135},"87699d48":{"m":136,"g":135},"52c60434":{"m":136,"g":135},"ff0f370f":{"m":136,"g":135},"26f9e207":{"m":136,"g":135},"828cd893":{"m":136,"g":135},"4436dc0f":{"m":136,"g":135},"cf6800f6":{"m":136,"g":135},"f16606d6":{"m":136,"g":135},"387fad2f":{"m":136,"g":135},"76bc07a3":{"m":136,"g":135},"b328cd20":{"m":136,"g":135},"bf32cd83":{"m":136,"g":135},"27a08305":{"m":136,"g":135},"53479e22":{"m":136,"g":135},"ff978e7d":{"m":136,"g":135},"16e00651":{"m":136,"g":135},"1b2b95d8":{"m":136,"g":135},"9cac3c86":{"m":136,"g":135},"8d58b3dc":{"m":136,"g":135},"5f3eb377":{"m":136,"g":135},"22993880":{"m":136,"g":135},"d7aa0ce7":{"m":136,"g":135},"e797f0c5":{"m":136,"g":135},"5d4b7c78":{"m":136,"g":135},"9bd64d73":{"m":136,"g":135},"24e116ef":{"m":136,"g":135},"8ca95970":{"m":136,"g":135},"216ea910":{"m":136,"g":135},"f4ab2ec5":{"m":136,"g":135},"25fa2ac2":{"m":136,"g":135},"ef5ac6f0":{"m":136,"g":135},"c88aaf22":{"m":136,"g":135},"e139d2aa":{"m":136,"g":135},"2337b1bb":{"m":136,"g":135},"7bc13c90":{"m":136,"g":135},"877c8e3a":{"m":136,"g":135},"66dfb8c1":{"m":136,"g":135},"dcacc492":{"m":136,"g":135},"87ef05e2":{"m":136,"g":135},"b65c9889":{"m":136,"g":135},"c7c0d97f":{"m":136,"g":135},"6bc5a52f":{"m":136,"g":135},"2c09de34":{"m":136,"g":135},"38d48de9":{"m":136,"g":135},"7f2fa216":{"m":136,"g":135},"d0fb24ee":{"m":136,"g":135},"65b0b5b2":{"m":136,"g":135},"9a414b16":{"m":136,"g":135},"5b4f7902":{"m":136,"g":135},"d8ac5eec":{"m":136,"g":135},"bdde9496":{"m":136,"g":135},"b23e7ed1":{"m":136,"g":135},"9821fae5":{"m":136,"g":135},"078d9621":{"m":136,"g":135},"8b111b20":{"m":136,"g":135},"74a166cb":{"m":136,"g":135},"888e126a":{"m":136,"g":135},"bb23a8fe":{"m":136,"g":135},"8b869e32":{"m":136,"g":135},"6256936d":{"m":136,"g":135},"62f73a8c":{"m":136,"g":135},"7c1b4b1c":{"m":136,"g":135},"31ed68e7":{"m":136,"g":135},"24c91001":{"m":136,"g":135},"c4edcac6":{"m":136,"g":135},"e9343389":{"m":136,"g":135},"a2d4f58a":{"m":136,"g":135},"f66b0916":{"m":136,"g":135},"2f623368":{"m":136,"g":135},"bd9a2ced":{"m":136,"g":135},"0ca417d9":{"m":136,"g":135},"30cfb687":{"m":136,"g":135},"b7c7e03d":{"m":136,"g":135},"f0195627":{"m":136,"g":135},"d401d238":{"m":136,"g":135},"17041f46":{"m":136,"g":135},"f07e76b2":{"m":136,"g":135},"1cfd2b2d":{"m":136,"g":135},"0d244116":{"m":136,"g":135},"0eae8317":{"m":136,"g":135},"c483a5f4":{"m":136,"g":135},"f26f6c2c":{"m":136,"g":135},"b0213323":{"m":136,"g":135},"5062537b":{"m":136,"g":135},"698629d1":{"m":136,"g":135},"dd93e445":{"m":136,"g":135},"6c8587b5":{"m":136,"g":135},"02704260":{"m":136,"g":135},"bd48ad5e":{"m":136,"g":135},"749736ba":{"m":136,"g":135},"72549863":{"m":136,"g":135},"00562ee1":{"m":136,"g":135},"d7a8257b":{"m":136,"g":135},"f6f7af40":{"m":136,"g":135},"a3b1e8ef":{"m":136,"g":135},"db499e18":{"m":136,"g":135},"6cf3a6dd":{"m":136,"g":135},"90e24f5c":{"m":136,"g":135},"21de3e14":{"m":136,"g":135},"e4c1e441":{"m":136,"g":135},"b5af283b":{"m":136,"g":135},"130c6911":{"m":136,"g":135},"e0e50848":{"m":136,"g":135},"70a769bc":{"m":136,"g":135},"3a42c5e3":{"m":136,"g":135},"12b89e51":{"m":136,"g":135},"85184557":{"m":136,"g":135},"57d2ba92":{"m":136,"g":135},"417e75a6":{"m":136,"g":135},"0500fea9":{"m":136,"g":135},"2b461c15":{"m":136,"g":135},"5595ae14":{"m":136,"g":135},"ace6f300":{"m":136,"g":135},"2da49eec":{"m":136,"g":135},"d65ae0ec":{"m":136,"g":135},"60d7279c":{"m":136,"g":135},"abdf65d4":{"m":136,"g":135},"c1dfbc77":{"m":136,"g":135},"3b3c5a05":{"m":136,"g":135},"2667c857":{"m":136,"g":135},"fc643ffb":{"m":136,"g":135},"4280a18a":{"m":136,"g":135},"386e5415":{"m":136,"g":135},"b6267de5":{"m":136,"g":135},"e47afa02":{"m":136,"g":135},"5ed384d0":{"m":136,"g":135},"b4ce7a6d":{"m":136,"g":135},"25b48564":{"m":136,"g":135},"1c360bf7":{"m":136,"g":135},"3619ec61":{"m":136,"g":135},"db6b51a8":{"m":136,"g":135},"8a9ca41f":{"m":136,"g":135},"ac11e6a7":{"m":136,"g":135},"47a660d5":{"m":136,"g":135},"c0fc7a89":{"m":136,"g":135},"5bf0d862":{"m":136,"g":135},"75b72eb8":{"m":136,"g":135},"bc8b526e":{"m":136,"g":135},"3dfff6ae":{"m":136,"g":135},"ad2c1ee3":{"m":136,"g":135},"9940c6f5":{"m":136,"g":135},"00e60711":{"m":136,"g":135},"4bc2f2e0":{"m":136,"g":135},"d17b9e63":{"m":136,"g":135},"ba67e006":{"m":136,"g":135},"45f3ad2f":{"m":136,"g":135},"f35b5da5":{"m":136,"g":135},"733a0c1a":{"m":136,"g":135},"b369aaa2":{"m":136,"g":135},"b9732025":{"m":136,"g":135},"cbff7ad9":{"m":136,"g":135},"4de59d83":{"m":136,"g":135},"39ca57cd":{"m":136,"g":135},"34498067":{"m":136,"g":135},"b3817fa9":{"m":136,"g":135},"059428bd":{"m":136,"g":135},"5d200dd8":{"m":136,"g":135},"7f9a3d06":{"m":136,"g":135},"664f611e":{"m":136,"g":135},"49adb37e":{"m":136,"g":135},"7518dc35":{"m":136,"g":135},"b6871ba7":{"m":136,"g":135},"c1f2241a":"m137","8da70e2a":"m137"},"g":"2026-02-12T19:54:15.483852"} ================================================ FILE: docs/requirements.txt ================================================ ipykernel ipywidgets jupyter_client markdown>=3.4.0 matplotlib myst-parser nbconvert nbsphinx pandoc pillow pydantic sphinx sphinx-book-theme sphinx-copybutton sphinx-tabs nbstripout sphinxcontrib-mermaid urllib3<2.0.0 gguf>=0.17.1 sphinx-autobuild ================================================ FILE: docs/serve.sh ================================================ # Clean and serve documentation with auto-build make clean make serve ================================================ FILE: docs/supported_models/extending/index.rst ================================================ Extending SGLang ================ Adding new models and alternative backends. .. toctree:: :maxdepth: 1 support_new_models.md transformers_fallback.md modelscope.md mindspore_models.md ================================================ FILE: docs/supported_models/extending/mindspore_models.md ================================================ # MindSpore Models ## Introduction MindSpore is a high-performance AI framework optimized for Ascend NPUs. This doc guides users to run MindSpore models in SGLang. ## Requirements MindSpore currently only supports Ascend NPU devices. Users need to first install Ascend CANN 8.5. The CANN software packages can be downloaded from the [Ascend Official Website](https://www.hiascend.com). ## Supported Models Currently, the following models are supported: - **Qwen3**: Dense and MoE models - **DeepSeek V3/R1** - *More models coming soon...* ## Installation > **Note**: Currently, MindSpore models are provided by an independent package `sgl-mindspore`. Support for MindSpore is built upon current SGLang support for Ascend NPU platform. Please first [install SGLang for Ascend NPU](../../platforms/ascend_npu.md) and then install `sgl-mindspore`: ```shell git clone https://github.com/mindspore-lab/sgl-mindspore.git cd sgl-mindspore pip install -e . ``` ## Run Model Current SGLang-MindSpore supports Qwen3 and DeepSeek V3/R1 models. This doc uses Qwen3-8B as an example. ### Offline inference Use the following script for offline inference: ```python import sglang as sgl # Initialize the engine with MindSpore backend llm = sgl.Engine( model_path="/path/to/your/model", # Local model path device="npu", # Use NPU device model_impl="mindspore", # MindSpore implementation attention_backend="ascend", # Attention backend tp_size=1, # Tensor parallelism size dp_size=1 # Data parallelism size ) # Generate text prompts = [ "Hello, my name is", "The capital of France is", "The future of AI is" ] sampling_params = {"temperature": 0, "top_p": 0.9} outputs = llm.generate(prompts, sampling_params) for prompt, output in zip(prompts, outputs): print(f"Prompt: {prompt}") print(f"Generated: {output['text']}") print("---") ``` ### Start server Launch a server with MindSpore backend: ```bash # Basic server startup python3 -m sglang.launch_server \ --model-path /path/to/your/model \ --host 0.0.0.0 \ --device npu \ --model-impl mindspore \ --attention-backend ascend \ --tp-size 1 \ --dp-size 1 ``` For distributed server with multiple nodes: ```bash # Multi-node distributed server python3 -m sglang.launch_server \ --model-path /path/to/your/model \ --host 0.0.0.0 \ --device npu \ --model-impl mindspore \ --attention-backend ascend \ --dist-init-addr 127.0.0.1:29500 \ --nnodes 2 \ --node-rank 0 \ --tp-size 4 \ --dp-size 2 ``` ## Troubleshooting #### Debug Mode Enable sglang debug logging by log-level argument. ```bash python3 -m sglang.launch_server \ --model-path /path/to/your/model \ --host 0.0.0.0 \ --device npu \ --model-impl mindspore \ --attention-backend ascend \ --log-level DEBUG ``` Enable mindspore info and debug logging by setting environments. ```bash export GLOG_v=1 # INFO export GLOG_v=0 # DEBUG ``` #### Explicitly select devices Use the following environment variable to explicitly select the devices to use. ```shell export ASCEND_RT_VISIBLE_DEVICES=4,5,6,7 # to set device ``` #### Some communication environment issues In case of some environment with special communication environment, users need set some environment variables. ```shell export MS_ENABLE_LCCL=off # current not support LCCL communication mode in SGLang-MindSpore ``` #### Some dependencies of protobuf In case of some environment with special protobuf version, users need set some environment variables to avoid binary version mismatch. ```shell export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python # to avoid protobuf binary version mismatch ``` ## Support For MindSpore-specific issues: - Refer to the [MindSpore documentation](https://www.mindspore.cn/) ================================================ FILE: docs/supported_models/extending/modelscope.md ================================================ # Use Models From ModelScope To use a model from [ModelScope](https://www.modelscope.cn), set the environment variable `SGLANG_USE_MODELSCOPE`. ```bash export SGLANG_USE_MODELSCOPE=true ``` We take [Qwen2-7B-Instruct](https://www.modelscope.cn/models/qwen/qwen2-7b-instruct) as an example. Launch the Server: ```bash python -m sglang.launch_server --model-path qwen/Qwen2-7B-Instruct --port 30000 ``` Or start it by docker: ```bash docker run --gpus all \ -p 30000:30000 \ -v ~/.cache/modelscope:/root/.cache/modelscope \ --env "SGLANG_USE_MODELSCOPE=true" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct --host 0.0.0.0 --port 30000 ``` Note that modelscope uses a different cache directory than huggingface. You may need to set it manually to avoid running out of disk space. ================================================ FILE: docs/supported_models/extending/support_new_models.md ================================================ # How to Support New Models This document explains how to add support for new language models and multimodal large language models (MLLMs) in SGLang. It also covers how to test new models and register external implementations. ## How to Support a New Language Model To support a new model in SGLang, you only need to add a single file under the [SGLang Models Directory](https://github.com/sgl-project/sglang/tree/main/python/sglang/srt/models). You can learn from existing model implementations and create a new file for your model. For most models, you should be able to find a similar model to start with (e.g., starting from Llama). Also refer how to [port a Model from vLLM to SGLang](#port-a-model-from-vllm-to-sglang) ## How to Support a New Multimodal Large Language Model To support a new multimodal large language model (MLLM) in SGLang, there are several key components in addition to the standard LLM support: 1. **Register your new model as multimodal**: Extend `is_multimodal_model` in [model_config.py](https://github.com/sgl-project/sglang/blob/0ab3f437aba729b348a683ab32b35b214456efc7/python/sglang/srt/configs/model_config.py#L561) to return `True` for your model. 2. **Register a new chat-template**: Only when your default chat-template is unable to accept images as input: Register a new chat template in [conversation.py](https://github.com/sgl-project/sglang/tree/main/python/sglang/srt/conversation.py) and the corresponding matching function. 3. **Multimodal Data Processor**: Define a new `Processor` class that inherits from `BaseMultimodalProcessor` and register this processor as your model’s dedicated processor. See [multimodal_processor.py](https://github.com/sgl-project/sglang/tree/main/python/sglang/srt/multimodal/processors) for more details. 4. **Handle Multimodal Tokens**: Implement a `pad_input_ids` function for your new model. In this function, multimodal tokens in the prompt should be expanded (if necessary) and padded with multimodal-data-hashes so that SGLang can recognize different multimodal data with `RadixAttention`. 5. **Handle Image Feature Extraction**: Implement a `get_image_feature` function for your new model, which extracts image features from raw image data and converts them into the embeddings used by the language model. 6. **Adapt to Vision Attention**: Adapt the multi-headed `Attention` of ViT with SGLang’s `VisionAttention`. You can refer to [Qwen2VL](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/models/qwen2_vl.py) or other mllm implementations. These models demonstrate how to correctly handle both multimodal and textual inputs. ## Testing and Debugging Please note all your testing and benchmarking results in PR description. ### Interactive Debugging For interactive debugging, compare the outputs of Hugging Face/Transformers and SGLang. The following two commands should give the same text output and very similar prefill logits: - Get the reference output: ```bash python3 scripts/playground/reference_hf.py --model-path [new model] --model-type {text,vlm} ``` - Get the SGLang output: ```bash python3 -m sglang.bench_one_batch --correct --model [new model] ``` ### Add the Model to the Test Suite To ensure the new model is well maintained, add it to the test suite by including it in the `ALL_OTHER_MODELS` list in the [test_generation_models.py](https://github.com/sgl-project/sglang/blob/main/test/registered/models/test_generation_models.py) file, test the new model on your local machine and report the results on demonstrative benchmarks (GSM8K, MMLU, MMMU, MMMU-Pro, etc.) in your PR. \\ For VLMs, also include a test in `test_vision_openai_server_{x}.py` (e.g. [test_vision_openai_server_a.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_vision_openai_server_a.py), [test_vision_openai_server_b.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_vision_openai_server_b.py)). This is an example command to run to test a new model on your local machine: ```bash ONLY_RUN=Qwen/Qwen2-1.5B python3 -m unittest test_generation_models.TestGenerationModels.test_others ``` ### Benchmark - **(Required) MMMU**: follow MMMU benchmark [README.md](https://github.com/sgl-project/sglang/blob/main/benchmark/mmmu/README.md) to get SGLang vs. HF Transformer accuracy comparison. The accuracy score from SGLang run should not be much lower than that from HF Transformer run. Similarly, follow https://docs.sglang.io/developer_guide/benchmark_and_profiling.html to get performance comparison: TTFT and throughput must meet or exceed baselines (e.g., HF Transformer). - **(Optional) Other evals**: If you ran other evals, please note the results in PR description. ## Port a Model from vLLM to SGLang The [vLLM Models Directory](https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/models) is a valuable resource, as vLLM covers many models. SGLang reuses vLLM’s interface and some layers, making it easier to port models from vLLM to SGLang. To port a model from vLLM to SGLang: - Compare these two files for guidance: - [SGLang Llama Implementation](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/models/llama.py) - [vLLM Llama Implementation](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/llama.py) - The major differences include: - **Replace vLLM’s `Attention` with `RadixAttention`** (ensure you pass `layer_id` to `RadixAttention`). - **Replace vLLM’s `LogitsProcessor` with SGLang’s `LogitsProcessor`.** - **Replace the multi-headed `Attention` of ViT with SGLang’s `VisionAttention`.** - **Replace other vLLM layers** (such as `RMSNorm`, `SiluAndMul`) with SGLang layers. - **Remove `Sample`.** - **Change the `forward()` functions** and add a `forward_batch()` method. - **Add `EntryClass`** at the end. - **Ensure that the new implementation uses only SGLang components** and does not rely on any vLLM components. Note: make sure you add your new model to the supported models list in the supported models documentation. ## Registering an External Model Implementation In addition to the methods above, you can register your new model with the `ModelRegistry` before launching the server. This allows you to integrate your model without modifying the source code. For example: ```python from sglang.srt.models.registry import ModelRegistry from sglang.srt.entrypoints.http_server import launch_server # For a single model, add it to the registry: ModelRegistry.models[model_name] = model_class # For multiple models, you can imitate the import_model_classes() function: from functools import lru_cache @lru_cache() def import_new_model_classes(): model_arch_name_to_cls = {} # Populate model_arch_name_to_cls with your new model classes. ... return model_arch_name_to_cls ModelRegistry.models.update(import_new_model_classes()) # Launch the server with your server arguments: launch_server(server_args) ``` ## Example: Implementing and Serving a Llama Wrapper Model Below is an introductory, step-by-step walkthrough on how to implement a new model end-to-end in SGLang and then run it via the [Offline Engine](https://github.com/sgl-project/sglang/blob/main/docs/basic_usage/offline_engine_api.ipynb). ### Implementing Our Model To keep things simple, this new model will be a simple wrapper around [Llama 3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct), and our goal will be just to bias the output logits for each `forward` call by taking the square root of each individual logit. Let's start by defining our model in a file called `llama_wrapper.py`. The first step is to import the necessary libraries from SRT, which is SGLang's internal backend. ```python # In the file `llama_wrapper.py` import torch from transformers import LlamaConfig from typing import Optional from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.models.llama import LlamaForCausalLM ``` Next, we declare a new `class` for our model and have it inherit from `LlamaForCausalLM`, which allows our model to access `LlamaForCausalLM`'s predefined modules and layers, such as `LlamaAttention` and `LlamaMLP`. Note that almost all model implementations take in `config` and `quant_config` as arguments for their `__init__` method; `config` and `quant_config` are passed in via [`model_loader/loader.py`](https://github.com/sgl-project/sglang/blob/bf72b80122fd888bf619d17b96fa3e323ab809fc/python/sglang/srt/model_loader/loader.py#L219). Because we have inherited from `LlamaForCausalLM`, we can pass our parameters directly to its constructor, which will set the member variables for us. ```python class LlamaWrapper(LlamaForCausalLM): def __init__( self, config: LlamaConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__(config=config, quant_config=quant_config, prefix=prefix) ``` Now, we want to define the `forward` method, which is what will be called at inference time. Note that the signature for `forward` is essentially the same for any model; you can take a look at the other models defined in the [`models` directory](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/models/) for references. To see where exactly `forward` is called in the SGLang runtime's internals, take a look at [`forward_decode`](https://github.com/sgl-project/sglang/blob/bf72b80122fd888bf619d17b96fa3e323ab809fc/python/sglang/srt/model_executor/model_runner.py#L1705) and [`forward_extend`](https://github.com/sgl-project/sglang/blob/bf72b80122fd888bf619d17b96fa3e323ab809fc/python/sglang/srt/model_executor/model_runner.py#L1724) in the [`ModelRunner` class](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/model_executor/model_runner.py). ```python @torch.no_grad() def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, pp_proxy_tensors: Optional[PPProxyTensors] = None, input_embeds: Optional[torch.Tensor] = None, get_embedding: bool = False, ) -> LogitsProcessorOutput: ``` We now call the `__call__` method for `self.model` (which is a member variable that `LlamaForCausalLM` defines in its `__init__` method), which eventually calls `LlamaForCausalLM`'s `forward` method. After that, we feed the `hidden_states` into our model's `LogitsProcessor` (again defined in `LlamaForCausalLM`). ```python hidden_states = self.model( input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors=pp_proxy_tensors, ) res: LogitsProcessorOutput = self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch, ) ``` After receiving the logits for the next token, we can finally perform our biasing step. ```python orig_logits = res.next_token_logits res.next_token_logits = torch.where( orig_logits > 0, orig_logits.sqrt(), orig_logits ) return res ``` Now, our `LlamaWrapper` model is created and ready to be served! ### Serving Our Model Via SGLang's Offline Engine The next step of this walkthrough involves hosting our new model offline, so that it can be served locally and without an HTTP server. First, create a new file called `run.py`. Now, we must ensure that SGLang's `ModelRegistry` can find our model. To do this, we first download the model's configuration and weights from Huggingface. ```python # In the file `run.py` import asyncio from functools import lru_cache from huggingface_hub import snapshot_download from llama_wrapper import LlamaWrapper # Make sure to import our new model! import sglang as sgl from sglang.srt.models.registry import ModelRegistry # Make sure to request access to this model on Huggingface, then export your # `HF_TOKEN` to download the model snapshot llama_dir = snapshot_download( repo_id="meta-llama/Llama-3.1-8B-Instruct", local_dir="./llama_ckpt", ) ``` Now that we have our model on disk, we want to point it to `LlamaWrapper` by changing the `architectures` field in `./llama_ckpt/config.json` to be `LlamaWrapper`. That way, when we pass in the path of our model checkpoint to SGLang, it will know that we want to use "LlamaWrapper" instead of "LlamaForCausalLM" as our model. ```python { "architectures": [ # "LlamaForCausalLM" "LlamaWrapper" ], ... } ``` However, if we don't link our `LlamaWrapper` class to the "LlamaWrapper" registry keyword, then SGLang won't be able to find our model. Thus, to register our `LlamaWrapper`, we want to follow the steps in the above section titled "Registering an External Model Implementation". ```python @lru_cache() def import_new_model_classes(): model_arch_name_to_cls = {"LlamaWrapper": LlamaWrapper} return model_arch_name_to_cls ModelRegistry.models.update(import_new_model_classes()) ``` Lastly, when we create our `Engine`, we just pass in the path to the local model directory. Then, our `LlamaWrapper` is ready to be served; for this walkthrough, we will use SGLang `Engine`'s non-streaming asynchronous generation endpoint. ```python def main(): llm = sgl.Engine(model_path="./llama_ckpt") sampling_params = {"temperature": 0.2, "top_k": 5} prompts = [ "Write a short, neutral self-introduction for a fictional character. Hello, my name is", "Provide a concise factual statement about France’s capital city. The capital of France is", "Explain possible future trends in artificial intelligence. The future of AI is", ] asyncio.run(run_llm(llm, sampling_params, prompts)) llm.shutdown() async def run_llm( llm, sampling_params, prompts, ) -> None: outputs = await llm.async_generate(prompts, sampling_params) for prompt, output in zip(prompts, outputs): print(f"\nPrompt: {prompt}") print(f"Generated text: {output['text']}") if __name__ == "__main__": main() ``` Now, when we call `python run.py`, we will get the outputs of our newly created model! ## Documentation Add to table of supported models in [generative_models.md](../text_generation/generative_models.md) or [multimodal_language_models.md](../text_generation/multimodal_language_models.md) --- By following these guidelines, you can add support for new language models and multimodal large language models in SGLang and ensure they are thoroughly tested and easily integrated into the system. ================================================ FILE: docs/supported_models/extending/transformers_fallback.md ================================================ # Transformers fallback in SGLang `sglang` can fall back to using models that are available in `transformers`. This works for most decoder-style language models and support for vision-language models is coming soon! ## Example launch Command By default, we will use sglang implementation if it is available. Otherwise, we will fall back to transformers one. However, you can switch the implementation by setting `--model-impl` to `transformers`. ```shell python3 -m sglang.launch_server \ --model-path meta-llama/Llama-3.2-1B-Instruct \ --host 0.0.0.0 \ --port 30000 \ --model-impl transformers ``` ## Supported features ### Quantization Transformers fall back has supported most of available quantization in SGLang (except GGUF). See [Quantization page](../advanced_features/quantization.md) for more information about supported quantization in SGLang. ### Remote code This fallback also means that any model on the hub that can be used in `transformers` with `trust_remote_code=True` that correctly implements attention can be used in production! A model just needs the following two things: ```python from transformers import PreTrainedModel from torch import nn class MyAttention(nn.Module): def forward(self, hidden_states, **kwargs): # <- kwargs are required ... attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, **kwargs, ) ... class MyModel(PreTrainedModel): _supports_attention_backend = True ``` Here is what happens in the background: 1. The config is loaded 2. `MyModel` python class is loaded from the `auto_map`, and we check that the model `_supports_attention_backend`. 3. The `TransformersModel` backend is used. See `/srt/models/transformers`, which leverages `self.config._attn_implementation = "sglang"`, thus the need to use `ALL_ATTENTION_FUNCTIONS`. That's it! ================================================ FILE: docs/supported_models/index.rst ================================================ Supported Models ================ SGLang supports a wide variety of model architectures for different use cases. Browse by category below to find models suited for your needs. .. toctree:: :maxdepth: 2 text_generation/index retrieval_ranking/index specialized/index extending/index ================================================ FILE: docs/supported_models/retrieval_ranking/classify_models.md ================================================ # Classification API This document describes the `/v1/classify` API endpoint implementation in SGLang, which is compatible with vLLM's classification API format. ## Overview The classification API allows you to classify text inputs using classification models. This implementation follows the same format as vLLM's 0.7.0 classification API. ## API Endpoint ``` POST /v1/classify ``` ## Request Format ```json { "model": "model_name", "input": "text to classify" } ``` ### Parameters - `model` (string, required): The name of the classification model to use - `input` (string, required): The text to classify - `user` (string, optional): User identifier for tracking - `rid` (string, optional): Request ID for tracking - `priority` (integer, optional): Request priority ## Response Format ```json { "id": "classify-9bf17f2847b046c7b2d5495f4b4f9682", "object": "list", "created": 1745383213, "model": "jason9693/Qwen2.5-1.5B-apeach", "data": [ { "index": 0, "label": "Default", "probs": [0.565970778465271, 0.4340292513370514], "num_classes": 2 } ], "usage": { "prompt_tokens": 10, "total_tokens": 10, "completion_tokens": 0, "prompt_tokens_details": null } } ``` ### Response Fields - `id`: Unique identifier for the classification request - `object`: Always "list" - `created`: Unix timestamp when the request was created - `model`: The model used for classification - `data`: Array of classification results - `index`: Index of the result - `label`: Predicted class label - `probs`: Array of probabilities for each class - `num_classes`: Total number of classes - `usage`: Token usage information - `prompt_tokens`: Number of input tokens - `total_tokens`: Total number of tokens - `completion_tokens`: Number of completion tokens (always 0 for classification) - `prompt_tokens_details`: Additional token details (optional) ## Example Usage ### Using curl ```bash curl -v "http://127.0.0.1:8000/v1/classify" \ -H "Content-Type: application/json" \ -d '{ "model": "jason9693/Qwen2.5-1.5B-apeach", "input": "Loved the new café—coffee was great." }' ``` ### Using Python ```python import requests import json # Make classification request response = requests.post( "http://127.0.0.1:8000/v1/classify", headers={"Content-Type": "application/json"}, json={ "model": "jason9693/Qwen2.5-1.5B-apeach", "input": "Loved the new café—coffee was great." } ) # Parse response result = response.json() print(json.dumps(result, indent=2)) ``` ## Supported Models The classification API works with any classification model supported by SGLang, including: ### Classification Models (Multi-class) - `LlamaForSequenceClassification` - Multi-class classification - `Qwen2ForSequenceClassification` - Multi-class classification - `Qwen3ForSequenceClassification` - Multi-class classification - `BertForSequenceClassification` - Multi-class classification - `Gemma2ForSequenceClassification` - Multi-class classification **Label Mapping**: The API automatically uses the `id2label` mapping from the model's `config.json` file to provide meaningful label names instead of generic class names. If `id2label` is not available, it falls back to `LABEL_0`, `LABEL_1`, etc., or `Class_0`, `Class_1` as a last resort. ### Reward Models (Single score) - `InternLM2ForRewardModel` - Single reward score - `Qwen2ForRewardModel` - Single reward score - `LlamaForSequenceClassificationWithNormal_Weights` - Special reward model **Note**: The `/classify` endpoint in SGLang was originally designed for reward models but now supports all non-generative models. Our `/v1/classify` endpoint provides a standardized vLLM-compatible interface for classification tasks. ## Error Handling The API returns appropriate HTTP status codes and error messages: - `400 Bad Request`: Invalid request format or missing required fields - `500 Internal Server Error`: Server-side processing error Error response format: ```json { "error": "Error message", "type": "error_type", "code": 400 } ``` ## Implementation Details The classification API is implemented using: 1. **Rust Model Gateway**: Handles routing and request/response models in `sgl-model-gateway/src/protocols/spec.rs` 2. **Python HTTP Server**: Implements the actual endpoint in `python/sglang/srt/entrypoints/http_server.py` 3. **Classification Service**: Handles the classification logic in `python/sglang/srt/entrypoints/openai/serving_classify.py` ## Testing Use the provided test script to verify the implementation: ```bash python test_classify_api.py ``` ## Compatibility This implementation is compatible with vLLM's classification API format, allowing seamless migration from vLLM to SGLang for classification tasks. ================================================ FILE: docs/supported_models/retrieval_ranking/embedding_models.md ================================================ # Embedding Models SGLang provides robust support for embedding models by integrating efficient serving mechanisms with its flexible programming interface. This integration allows for streamlined handling of embedding tasks, facilitating faster and more accurate retrieval and semantic search operations. SGLang's architecture enables better resource utilization and reduced latency in embedding model deployment. ```{important} Embedding models are executed with `--is-embedding` flag and some may require `--trust-remote-code` ``` ## Quick Start ### Launch Server ```shell python3 -m sglang.launch_server \ --model-path Qwen/Qwen3-Embedding-4B \ --is-embedding \ --host 0.0.0.0 \ --port 30000 ``` ### Client Request ```python import requests url = "http://127.0.0.1:30000" payload = { "model": "Qwen/Qwen3-Embedding-4B", "input": "What is the capital of France?", "encoding_format": "float" } response = requests.post(url + "/v1/embeddings", json=payload).json() print("Embedding:", response["data"][0]["embedding"]) ``` ## Multimodal Embedding Example For multimodal models like GME that support both text and images: ```shell python3 -m sglang.launch_server \ --model-path Alibaba-NLP/gme-Qwen2-VL-2B-Instruct \ --is-embedding \ --chat-template gme-qwen2-vl \ --host 0.0.0.0 \ --port 30000 ``` ```python import requests url = "http://127.0.0.1:30000" text_input = "Represent this image in embedding space." image_path = "https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild/resolve/main/images/023.jpg" payload = { "model": "gme-qwen2-vl", "input": [ { "text": text_input }, { "image": image_path } ], } response = requests.post(url + "/v1/embeddings", json=payload).json() print("Embeddings:", [x.get("embedding") for x in response.get("data", [])]) ``` ## Matryoshka Embedding Example [Matryoshka Embeddings](https://sbert.net/examples/sentence_transformer/training/matryoshka/README.html#matryoshka-embeddings) or [Matryoshka Representation Learning (MRL)](https://arxiv.org/abs/2205.13147) is a technique used in training embedding models. It allows user to trade off between performance and cost. ### 1. Launch a Matryoshka‑capable model If the model config already includes `matryoshka_dimensions` or `is_matryoshka` then no override is needed. Otherwise, you can use `--json-model-override-args` as below: ```shell python3 -m sglang.launch_server \ --model-path Qwen/Qwen3-Embedding-0.6B \ --is-embedding \ --host 0.0.0.0 \ --port 30000 \ --json-model-override-args '{"matryoshka_dimensions": [128, 256, 512, 1024, 1536]}' ``` 1. Setting `"is_matryoshka": true` allows truncating to any dimension. Otherwise, the server will validate that the specified dimension in the request is one of `matryoshka_dimensions`. 2. Omitting `dimensions` in a request returns the full vector. ### 2. Make requests with different output dimensions ```python import requests url = "http://127.0.0.1:30000" # Request a truncated (Matryoshka) embedding by specifying a supported dimension. payload = { "model": "Qwen/Qwen3-Embedding-0.6B", "input": "Explain diffusion models simply.", "dimensions": 512 # change to 128 / 1024 / omit for full size } response = requests.post(url + "/v1/embeddings", json=payload).json() print("Embedding:", response["data"][0]["embedding"]) ``` ## Supported Models | Model Family | Example Model | Chat Template | Description | | ------------------------------------------ | -------------------------------------- | ------------- | --------------------------------------------------------------------------- | | **E5 (Llama/Mistral based)** | `intfloat/e5-mistral-7b-instruct` | N/A | High-quality text embeddings based on Mistral/Llama architectures | | **GTE-Qwen2** | `Alibaba-NLP/gte-Qwen2-7B-instruct` | N/A | Alibaba's text embedding model with multilingual support | | **Qwen3-Embedding** | `Qwen/Qwen3-Embedding-4B` | N/A | Latest Qwen3-based text embedding model for semantic representation | | **BGE** | `BAAI/bge-large-en-v1.5` | N/A | BAAI's text embeddings (requires `attention-backend` triton/torch_native) | | **GME (Multimodal)** | `Alibaba-NLP/gme-Qwen2-VL-2B-Instruct`| `gme-qwen2-vl`| Multimodal embedding for text and image cross-modal tasks | | **CLIP** | `openai/clip-vit-large-patch14-336` | N/A | OpenAI's CLIP for image and text embeddings | ================================================ FILE: docs/supported_models/retrieval_ranking/index.rst ================================================ Retrieval & Ranking =================== Models for embeddings, reranking, and classification. .. toctree:: :maxdepth: 1 embedding_models.md rerank_models.md classify_models.md ================================================ FILE: docs/supported_models/retrieval_ranking/rerank_models.md ================================================ # Rerank Models SGLang offers comprehensive support for rerank models by incorporating optimized serving frameworks with a flexible programming interface. This setup enables efficient processing of cross-encoder reranking tasks, improving the accuracy and relevance of search result ordering. SGLang’s design ensures high throughput and low latency during reranker model deployment, making it ideal for semantic-based result refinement in large-scale retrieval systems. ```{important} Rerank models in SGLang fall into two categories: - **Cross-encoder rerank models**: run with `--is-embedding` (embedding runner). - **Decoder-only rerank models**: run **without** `--is-embedding` and use next-token logprob scoring (yes/no). - Text-only (e.g. Qwen3-Reranker) - Multimodal (e.g. Qwen3-VL-Reranker): also supports image/video content Some models may require `--trust-remote-code`. ``` ## Supported rerank models | Model Family (Rerank) | Example HuggingFace Identifier | Chat Template | Description | |------------------------------------------------|--------------------------------------|---------------|----------------------------------------------------------------------------------------------------------------------------------| | **BGE-Reranker (BgeRerankModel)** | `BAAI/bge-reranker-v2-m3` | N/A | Currently only support `attention-backend` `triton` and `torch_native`. High-performance cross-encoder reranker model from BAAI. Suitable for reranking search results based on semantic relevance. | | **Qwen3-Reranker (decoder-only yes/no)** | `Qwen/Qwen3-Reranker-8B` | `examples/chat_template/qwen3_reranker.jinja` | Decoder-only reranker using next-token logprob scoring for labels (yes/no). Launch **without** `--is-embedding`. | | **Qwen3-VL-Reranker (multimodal yes/no)** | `Qwen/Qwen3-VL-Reranker-2B` | `examples/chat_template/qwen3_vl_reranker.jinja` | Multimodal decoder-only reranker supporting text, images, and videos. Uses yes/no logprob scoring. Launch **without** `--is-embedding`. | ## Cross-Encoder Rerank (embedding runner) ### Launch Command ```shell python3 -m sglang.launch_server \ --model-path BAAI/bge-reranker-v2-m3 \ --host 0.0.0.0 \ --disable-radix-cache \ --chunked-prefill-size -1 \ --attention-backend triton \ --is-embedding \ --port 30000 ``` ### Example Client Request ```python import requests url = "http://127.0.0.1:30000/v1/rerank" payload = { "model": "BAAI/bge-reranker-v2-m3", "query": "what is panda?", "documents": [ "hi", "The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China." ], "top_n": 1, "return_documents": True } response = requests.post(url, json=payload) response_json = response.json() for item in response_json: if item.get("document"): print(f"Score: {item['score']:.2f} - Document: '{item['document']}'") else: print(f"Score: {item['score']:.2f} - Index: {item['index']}") ``` **Request Parameters:** - `query` (required): The query text to rank documents against - `documents` (required): List of documents to be ranked - `model` (required): Model to use for reranking - `top_n` (optional): Maximum number of documents to return. Defaults to returning all documents. If specified value is greater than the total number of documents, all documents will be returned. - `return_documents` (optional): Whether to return documents in the response. Defaults to `True`. ## Qwen3-Reranker (decoder-only yes/no rerank) ### Launch Command ```shell python3 -m sglang.launch_server \ --model-path Qwen/Qwen3-Reranker-0.6B \ --trust-remote-code \ --disable-radix-cache \ --host 0.0.0.0 \ --port 8001 \ --chat-template examples/chat_template/qwen3_reranker.jinja ``` ```{note} Qwen3-Reranker uses decoder-only logprob scoring (yes/no). Do NOT launch it with `--is-embedding`. ``` ### Example Client Request (supports optional instruct, top_n, and return_documents) ```shell curl -X POST http://127.0.0.1:8001/v1/rerank \ -H "Content-Type: application/json" \ -d '{ "model": "Qwen3-Reranker-0.6B", "query": "法国首都是哪里？", "documents": [ "法国的首都是巴黎。", "德国的首都是柏林。", "香蕉是黄色的水果。" ], "instruct": "Given a web search query, retrieve relevant passages that answer the query.", "top_n": 2, "return_documents": true }' ``` **Request Parameters:** - `query` (required): The query text to rank documents against - `documents` (required): List of documents to be ranked - `model` (required): Model to use for reranking - `instruct` (optional): Instruction text for the reranker - `top_n` (optional): Maximum number of documents to return. Defaults to returning all documents. If specified value is greater than the total number of documents, all documents will be returned. - `return_documents` (optional): Whether to return documents in the response. Defaults to `True`. ### Response Format `/v1/rerank` returns a list of objects (sorted by descending score): - `score`: float, higher means more relevant - `document`: the original document string (only included when `return_documents` is `true`) - `index`: the original index in the input `documents` - `meta_info`: optional debug/usage info (may be present for some models) The number of returned results is controlled by the `top_n` parameter. If `top_n` is not specified or is greater than the total number of documents, all documents are returned. Example (with `return_documents: true`): ```json [ {"score": 0.99, "document": "法国的首都是巴黎。", "index": 0}, {"score": 0.01, "document": "德国的首都是柏林。", "index": 1}, {"score": 0.00, "document": "香蕉是黄色的水果。", "index": 2} ] ``` Example (with `return_documents: false`): ```json [ {"score": 0.99, "index": 0}, {"score": 0.01, "index": 1}, {"score": 0.00, "index": 2} ] ``` Example (with `top_n: 2`): ```json [ {"score": 0.99, "document": "法国的首都是巴黎。", "index": 0}, {"score": 0.01, "document": "德国的首都是柏林。", "index": 1} ] ``` ### Common Pitfalls - **`--chat-template` is required.** Without `--chat-template examples/chat_template/qwen3_reranker.jinja`, the server does not recognize the model as a decoder-only reranker and returns a 400 error: `"This model does not appear to be an embedding model by default. Please add `--is-embedding`..."`. The fix is to add the chat template flag, NOT `--is-embedding`. - If you launch Qwen3-Reranker with `--is-embedding`, `/v1/rerank` cannot compute yes/no logprob scores. Relaunch **without** `--is-embedding`. - If you see a validation error like "score should be a valid number" and the backend returned a list, upgrade to a version that coerces `embedding[0]` into `score` for rerank responses. ## Qwen3-VL-Reranker (multimodal decoder-only rerank) Qwen3-VL-Reranker extends the Qwen3-Reranker to support multimodal content, allowing reranking of documents containing text, images, and videos. ### Launch Command ```shell python3 -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-Reranker-2B \ --trust-remote-code \ --disable-radix-cache \ --host 0.0.0.0 \ --port 30000 \ --chat-template examples/chat_template/qwen3_vl_reranker.jinja ``` ```{note} Qwen3-VL-Reranker uses decoder-only logprob scoring (yes/no) like Qwen3-Reranker. Do NOT launch it with `--is-embedding`. ``` ### Text-Only Reranking (backward compatible) ```python import requests url = "http://127.0.0.1:30000/v1/rerank" payload = { "model": "Qwen3-VL-Reranker-2B", "query": "What is machine learning?", "documents": [ "Machine learning is a branch of artificial intelligence that enables computers to learn from data.", "The weather in Paris is usually mild with occasional rain.", "Deep learning is a subset of machine learning using neural networks with many layers.", ], "instruct": "Retrieve passages that answer the question.", "return_documents": True } response = requests.post(url, json=payload) results = response.json() for item in results: print(f"Score: {item['score']:.4f} - {item['document'][:60]}...") ``` ### Image Reranking (text query, image/mixed documents) ```python import requests url = "http://127.0.0.1:30000/v1/rerank" payload = { "query": "A woman playing with her dog on a beach at sunset.", "documents": [ # Document 1: Text description "A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset.", # Document 2: Image URL [ { "type": "image_url", "image_url": { "url": "https://example.com/beach_dog.jpeg" } } ], # Document 3: Text + Image (mixed) [ {"type": "text", "text": "A joyful scene at the beach:"}, { "type": "image_url", "image_url": { "url": "https://example.com/beach_dog.jpeg" } } ] ], "instruct": "Retrieve images or text relevant to the user's query.", "return_documents": False } response = requests.post(url, json=payload) results = response.json() for item in results: print(f"Index: {item['index']}, Score: {item['score']:.4f}") ``` ### Multimodal Query Reranking (query with image) ```python import requests url = "http://127.0.0.1:30000/v1/rerank" payload = { # Query with text and image "query": [ {"type": "text", "text": "Find similar images to this:"}, { "type": "image_url", "image_url": { "url": "https://example.com/reference_image.jpeg" } } ], "documents": [ "A cat sleeping on a couch.", "A woman and her dog enjoying the sunset at the beach.", "A busy city street with cars and pedestrians.", [ { "type": "image_url", "image_url": { "url": "https://example.com/similar_image.jpeg" } } ] ], "instruct": "Find images or descriptions similar to the query image." } response = requests.post(url, json=payload) results = response.json() for item in results: print(f"Index: {item['index']}, Score: {item['score']:.4f}") ``` ### Request Parameters (Multimodal) - `query` (required): Can be a string (text-only) or a list of content parts: - `{"type": "text", "text": "..."}` for text - `{"type": "image_url", "image_url": {"url": "..."}}` for images - `{"type": "video_url", "video_url": {"url": "..."}}` for videos - `documents` (required): List where each document can be a string or list of content parts (same format as query) - `instruct` (optional): Instruction text for the reranker - `top_n` (optional): Maximum number of documents to return - `return_documents` (optional): Whether to return documents in the response (default: `false`) ### Common Pitfalls - Always use `--chat-template examples/chat_template/qwen3_vl_reranker.jinja` for Qwen3-VL-Reranker. - Do NOT launch with `--is-embedding`. - For best results, use `--disable-radix-cache` to avoid caching issues with multimodal content. - **Note**: Currently only `Qwen3-VL-Reranker-2B` is tested and supported. The 8B model may have different behavior and is not guaranteed to work with this template. ================================================ FILE: docs/supported_models/specialized/index.rst ================================================ Specialized Models ================== Models for specialized tasks like reward modeling. .. toctree:: :maxdepth: 1 reward_models.md ================================================ FILE: docs/supported_models/specialized/reward_models.md ================================================ # Reward Models These models output a scalar reward score or classification result, often used in reinforcement learning or content moderation tasks. ```{important} They are executed with `--is-embedding` and some may require `--trust-remote-code`. ``` ## Example launch Command ```shell python3 -m sglang.launch_server \ --model-path Qwen/Qwen2.5-Math-RM-72B \ # example HF/local path --is-embedding \ --host 0.0.0.0 \ --tp-size=4 \ # set for tensor parallelism --port 30000 \ ``` ## Supported models | Model Family (Reward) | Example HuggingFace Identifier | Description | |---------------------------------------------------------------------------|-----------------------------------------------------|---------------------------------------------------------------------------------| | **Llama (3.1 Reward / `LlamaForSequenceClassification`)** | `Skywork/Skywork-Reward-Llama-3.1-8B-v0.2` | Reward model (preference classifier) based on Llama 3.1 (8B) for scoring and ranking responses for RLHF. | | **Gemma 2 (27B Reward / `Gemma2ForSequenceClassification`)** | `Skywork/Skywork-Reward-Gemma-2-27B-v0.2` | Derived from Gemma‑2 (27B), this model provides human preference scoring for RLHF and multilingual tasks. | | **InternLM 2 (Reward / `InternLM2ForRewardMode`)** | `internlm/internlm2-7b-reward` | InternLM 2 (7B)–based reward model used in alignment pipelines to guide outputs toward preferred behavior. | | **Qwen2.5 (Reward - Math / `Qwen2ForRewardModel`)** | `Qwen/Qwen2.5-Math-RM-72B` | A 72B math-specialized RLHF reward model from the Qwen2.5 series, tuned for evaluating and refining responses. | | **Qwen2.5 (Reward - Sequence / `Qwen2ForSequenceClassification`)** | `jason9693/Qwen2.5-1.5B-apeach` | A smaller Qwen2.5 variant used for sequence classification, offering an alternative RLHF scoring mechanism. | ================================================ FILE: docs/supported_models/text_generation/diffusion_language_models.md ================================================ # Diffusion Language Models Diffusion language models have shown promise for non-autoregressive text generation with parallel decoding capabilities. Unlike auto-regressive language models, different diffusion language models require different decoding strategies. ## Example Launch Command SGLang supports different DLLM algorithms such as `LowConfidence` and `JointThreshold`. ```shell python3 -m sglang.launch_server \ --model-path inclusionAI/LLaDA2.0-mini \ # example HF/local path --dllm-algorithm LowConfidence \ --dllm-algorithm-config ./config.yaml \ # Optional. Uses the algorithm's default if not set. --host 0.0.0.0 \ --port 30000 ``` ## Example Configuration File Depending on the algorithm selected, the configuration parameters vary. LowConfidence Config: ```yaml # Confidence threshold for accepting predicted tokens # - Higher values: More conservative, better quality but slower # - Lower values: More aggressive, faster but potentially lower quality # Range: 0.0 - 1.0 threshold: 0.95 # Default: 32, for LLaDA2MoeModelLM block_size: 32 ``` JointThreshold Config: ```yaml # Decoding threshold for Mask-to-Token (M2T) phase # - Higher values: More conservative, better quality but slower # - Lower values: More aggressive, faster but potentially lower quality # Range: 0.0 - 1.0 threshold: 0.5 # Decoding threshold for Token-to-Token (T2T) phase # Range: 0.0 - 1.0 # Setting to 0.0 allows full editing (recommended for most cases). edit_threshold: 0.0 # Max extra T2T steps after all masks are removed. Prevents infinite loops. max_post_edit_steps: 16 # 2-gram repetition penalty (default 0). # An empirical value of 3 is often sufficient to mitigate most repetitions. penalty_lambda: 0 ``` ## Example Client Code Snippet Just like other supported models, diffusion language models can be used via the REST API or Python client. Python client example for making a generation request to the launched server: ```python import sglang as sgl def main(): llm = sgl.Engine(model_path="inclusionAI/LLaDA2.0-mini", dllm_algorithm="LowConfidence", max_running_requests=1, trust_remote_code=True) prompts = [ "SYSTEMdetailed thinking off<|role_end|>HUMAN Write a brief introduction of the great wall <|role_end|>ASSISTANT" ] sampling_params = { "temperature": 0, "max_new_tokens": 1024, } outputs = llm.generate(prompts, sampling_params) print(outputs) if __name__ == '__main__': main() ``` Curl example for making a generation request to the launched server: ```bash curl -X POST "http://127.0.0.1:30000/generate" \ -H "Content-Type: application/json" \ -d '{ "text": [ "SYSTEMdetailed thinking off<|role_end|>HUMAN Write the number from 1 to 128 <|role_end|>ASSISTANT", "SYSTEMdetailed thinking off<|role_end|>HUMAN Write a brief introduction of the great wall <|role_end|>ASSISTANT" ], "stream": true, "sampling_params": { "temperature": 0, "max_new_tokens": 1024 } }' ``` ## Supported Models Below the supported models are summarized in a table. | Model Family | Example Model | Description | | -------------------------- | ---------------------------- | ---------------------------------------------------------------------------------------------------- | | **LLaDA2.0 (mini, flash)** | `inclusionAI/LLaDA2.0-flash` | LLaDA2.0-flash is a diffusion language model featuring a 100B Mixture-of-Experts (MoE) architecture. | | **SDAR (JetLM)** | `JetLM/SDAR-8B-Chat` | SDAR series diffusion language model (Chat), dense architecture. | | **SDAR (JetLM)** | `JetLM/SDAR-30B-A3B-Chat` | SDAR series diffusion language model (Chat), MoE architecture. | ================================================ FILE: docs/supported_models/text_generation/generative_models.md ================================================ # Large Language Models These models accept text input and produce text output (e.g., chat completions). They are primarily large language models (LLMs), some with mixture-of-experts (MoE) architectures for scaling. ## Example launch Command ```shell python3 -m sglang.launch_server \ --model-path meta-llama/Llama-3.2-1B-Instruct \ # example HF/local path --host 0.0.0.0 \ --port 30000 \ ``` ## Supported models Below the supported models are summarized in a table. If you are unsure if a specific architecture is implemented, you can search for it via GitHub. For example, to search for `Qwen3ForCausalLM`, use the expression: ``` repo:sgl-project/sglang path:/^python\/sglang\/srt\/models\// Qwen3ForCausalLM ``` in the GitHub search bar. | Model Family (Variants) | Example HuggingFace Identifier | Description | |-------------------------------------|--------------------------------------------------|----------------------------------------------------------------------------------------| | **DeepSeek** (v1, v2, v3/R1) | `deepseek-ai/DeepSeek-R1` | Series of advanced reasoning-optimized models (including a 671B MoE) trained with reinforcement learning; top performance on complex reasoning, math, and code tasks. [SGLang provides Deepseek v3/R1 model-specific optimizations](../basic_usage/deepseek.md) and [Reasoning Parser](../advanced_features/separate_reasoning.ipynb)| | **Kimi K2** (Thinking, Instruct) | `moonshotai/Kimi-K2-Instruct` | Moonshot AI's 1 trillion parameter MoE model (32B active) with 128K–256K context; state-of-the-art agentic intelligence with stable long-horizon agency across 200–300 sequential tool calls. Features MLA attention and native INT4 quantization. [See Reasoning Parser docs](../advanced_features/separate_reasoning.ipynb)| | **Kimi Linear** (48B-A3B) | `moonshotai/Kimi-Linear-48B-A3B-Instruct` | Moonshot AI's hybrid linear attention model (48B total, 3B active) with 1M token context; features Kimi Delta Attention (KDA) for up to 6× faster decoding and 75% KV cache reduction vs full attention. | | **GPT-OSS** | `openai/gpt-oss-20b`, `openai/gpt-oss-120b` | OpenAI’s latest GPT-OSS series for complex reasoning, agentic tasks, and versatile developer use cases.| | **Qwen** (3.5, 3, 3MoE, 3Next, 2.5, 2 series) | `Qwen/Qwen3.5-397B-A17B`, `Qwen/Qwen3-0.6B`, `Qwen/Qwen3-30B-A3B` | Alibaba’s latest Qwen3 series for complex reasoning, language understanding, and generation tasks; Support for MoE variants along with previous generation 2.5, 2, etc. [SGLang provides Qwen3 specific reasoning parser](../advanced_features/separate_reasoning.ipynb)| | **Llama** (2, 3.x, 4 series) | `meta-llama/Llama-4-Scout-17B-16E-Instruct` | Meta's open LLM series, spanning 7B to 400B parameters (Llama 2, 3, and new Llama 4) with well-recognized performance. [SGLang provides Llama-4 model-specific optimizations](../basic_usage/llama4.md) | | **Mistral** (Mixtral, NeMo, Small3) | `mistralai/Mistral-7B-Instruct-v0.2` | Open 7B LLM by Mistral AI with strong performance; extended into MoE (“Mixtral”) and NeMo Megatron variants for larger scale. | | **Gemma** (v1, v2, v3) | `google/gemma-3-1b-it` | Google’s family of efficient multilingual models (1B–27B); Gemma 3 offers a 128K context window, and its larger (4B+) variants support vision input. | | **Phi** (Phi-1.5, Phi-2, Phi-3, Phi-4, Phi-MoE series) | `microsoft/Phi-4-multimodal-instruct`, `microsoft/Phi-3.5-MoE-instruct` | Microsoft’s Phi family of small models (1.3B–5.6B); Phi-4-multimodal (5.6B) processes text, images, and speech, Phi-4-mini is a high-accuracy text model and Phi-3.5-MoE is a mixture-of-experts model. | | **MiniCPM** (v3, 4B) | `openbmb/MiniCPM3-4B` | OpenBMB’s series of compact LLMs for edge devices; MiniCPM 3 (4B) achieves GPT-3.5-level results in text tasks. | | **OLMo** (2, 3) | `allenai/OLMo-3-1125-32B`, `allenai/OLMo-3-32B-Think`, `allenai/OLMo-2-1124-7B-Instruct` | Allen AI’s series of Open Language Models designed to enable the science of language models. | | **OLMoE** (Open MoE) | `allenai/OLMoE-1B-7B-0924` | Allen AI’s open Mixture-of-Experts model (7B total, 1B active parameters) delivering state-of-the-art results with sparse expert activation. | | **MiniMax-M2** (M2, M2.1, M2.5) | `MiniMaxAI/MiniMax-M2.5`, `MiniMaxAI/MiniMax-M2.1`, `MiniMaxAI/MiniMax-M2` | MiniMax's SOTA LLM for coding & agentic workflows. | | **StableLM** (3B, 7B) | `stabilityai/stablelm-tuned-alpha-7b` | StabilityAI’s early open-source LLM (3B & 7B) for general text generation; a demonstration model with basic instruction-following ability. | | **Command-(R,A)** (Cohere) | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025` | Cohere’s open conversational LLM (Command series) optimized for long context, retrieval-augmented generation, and tool use. | | **DBRX** (Databricks) | `databricks/dbrx-instruct` | Databricks’ 132B-parameter MoE model (36B active) trained on 12T tokens; competes with GPT-3.5 quality as a fully open foundation model. | | **Grok** (xAI) | `xai-org/grok-1` | xAI’s grok-1 model known for vast size(314B parameters) and high quality; integrated in SGLang for high-performance inference. | | **ChatGLM** (GLM-130B family) | `THUDM/chatglm2-6b` | Zhipu AI’s bilingual chat model (6B) excelling at Chinese-English dialogue; fine-tuned for conversational quality and alignment. | | **InternLM 2** (7B, 20B) | `internlm/internlm2-7b` | Next-gen InternLM (7B and 20B) from SenseTime, offering strong reasoning and ultra-long context support (up to 200K tokens). | | **ExaONE 3** (Korean-English) | `LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct` | LG AI Research’s Korean-English model (7.8B) trained on 8T tokens; provides high-quality bilingual understanding and generation. | | **Baichuan 2** (7B, 13B) | `baichuan-inc/Baichuan2-13B-Chat` | BaichuanAI’s second-generation Chinese-English LLM (7B/13B) with improved performance and an open commercial license. | | **XVERSE** (MoE) | `xverse/XVERSE-MoE-A36B` | Yuanxiang’s open MoE LLM (XVERSE-MoE-A36B: 255B total, 36B active) supporting ~40 languages; delivers 100B+ dense-level performance via expert routing. | | **SmolLM** (135M–1.7B) | `HuggingFaceTB/SmolLM-1.7B` | Hugging Face’s ultra-small LLM series (135M–1.7B params) offering surprisingly strong results, enabling advanced AI on mobile/edge devices. | | **GLM-4** (Multilingual 9B) | `ZhipuAI/glm-4-9b-chat` | Zhipu’s GLM-4 series (up to 9B parameters) – open multilingual models with support for 1M-token context and even a 5.6B multimodal variant (Phi-4V). | | **MiMo** (7B series) | `XiaomiMiMo/MiMo-7B-RL` | Xiaomi's reasoning-optimized model series, leverages Multiple-Token Prediction for faster inference. | | **ERNIE-4.5** (4.5, 4.5MoE series) | `baidu/ERNIE-4.5-21B-A3B-PT` | Baidu's ERNIE-4.5 series which consists of MoE with 47B and 3B active parameters, with the largest model having 424B total parameters, as well as a 0.3B dense model. | | **Arcee AFM-4.5B** | `arcee-ai/AFM-4.5B-Base` | Arcee's foundational model series for real world reliability and edge deployments. | | **Persimmon** (8B) | `adept/persimmon-8b-chat` | Adept’s open 8B model with a 16K context window and fast inference; trained for broad usability and licensed under Apache 2.0. | | **Solar** (10.7B) | `upstage/SOLAR-10.7B-Instruct-v1.0` | Upstage's 10.7B parameter model, optimized for instruction-following tasks. This architecture incorporates a depth-up scaling methodology, enhancing model performance. | | **Tele FLM** (52B-1T) | `CofeAI/Tele-FLM` | BAAI & TeleAI's multilingual model, available in 52-billion and 1-trillion parameter variants. It is a decoder-only transformer trained on ~2T tokens | | **Ling** (16.8B–290B) | `inclusionAI/Ling-lite`, `inclusionAI/Ling-plus` | InclusionAI’s open MoE models. Ling-Lite has 16.8B total / 2.75B active parameters, and Ling-Plus has 290B total / 28.8B active parameters. They are designed for high performance on NLP and complex reasoning tasks. | | **Granite 3.0, 3.1** (IBM) | `ibm-granite/granite-3.1-8b-instruct` | IBM's open dense foundation models optimized for reasoning, code, and business AI use cases. Integrated with Red Hat and watsonx systems. | | **Granite 3.0 MoE** (IBM) | `ibm-granite/granite-3.0-3b-a800m-instruct` | IBM’s Mixture-of-Experts models offering strong performance with cost-efficiency. MoE expert routing designed for enterprise deployment at scale. | | **GPT-J** (6B) | `EleutherAI/gpt-j-6b` | EleutherAI's GPT-2-like causal language model (6B) trained on the [Pile](https://pile.eleuther.ai/) dataset. | | **Orion** (14B) | `OrionStarAI/Orion-14B-Base` | A series of open-source multilingual large language models by OrionStarAI, pretrained on a 2.5T token multilingual corpus including Chinese, English, Japanese, Korean, etc, and it exhibits superior performance in these languages. | | **Llama Nemotron Super** (v1, v1.5, NVIDIA) | `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, `nvidia/Llama-3_3-Nemotron-Super-49B-v1_5` | The [NVIDIA Nemotron](https://www.nvidia.com/en-us/ai-data-science/foundation-models/nemotron/) family of multimodal models provides state-of-the-art reasoning models specifically designed for enterprise-ready AI agents. | | **Llama Nemotron Ultra** (v1, NVIDIA) | `nvidia/Llama-3_1-Nemotron-Ultra-253B-v1` | The [NVIDIA Nemotron](https://www.nvidia.com/en-us/ai-data-science/foundation-models/nemotron/) family of multimodal models provides state-of-the-art reasoning models specifically designed for enterprise-ready AI agents. | | **NVIDIA Nemotron Nano 2.0** | `nvidia/NVIDIA-Nemotron-Nano-9B-v2` | The [NVIDIA Nemotron](https://www.nvidia.com/en-us/ai-data-science/foundation-models/nemotron/) family of multimodal models provides state-of-the-art reasoning models specifically designed for enterprise-ready AI agents. `Nemotron-Nano-9B-v2` is a hybrid Mamba-Transformer language model designed to increase throughput for reasoning workloads while achieving state-of-the-art accuracy compared to similarly-sized models. | | **StarCoder2** (3B-15B) | `bigcode/starcoder2-7b` | StarCoder2 is a family of open large language models (LLMs) specialized for code generation and understanding. It is the successor to StarCoder, jointly developed by the BigCode project (a collaboration between Hugging Face, ServiceNow Research, and other contributors). | | **Jet-Nemotron** | `jet-ai/Jet-Nemotron-2B` | Jet-Nemotron is a new family of hybrid-architecture language models that surpass state-of-the-art open-source full-attention language models, while achieving significant efficiency gains. | | **Trinity** (Nano, Mini) | `arcee-ai/Trinity-Mini` | Arcee's foundational MoE Trinity family of models, open weights under Apache 2.0. | | **Falcon-H1** (0.5B–34B) | `tiiuae/Falcon-H1-34B-Instruct` | TII's hybrid Mamba-Transformer architecture combining attention and state-space models for efficient long-context inference. | | **Hunyuan-Large** (389B, MoE) | `tencent/Tencent-Hunyuan-Large` | Tencent's open-source MoE model with 389B total / 52B active parameters, featuring Cross-Layer Attention (CLA) for improved efficiency. | | **IBM Granite 4.0 (Hybrid, Dense)** | `ibm-granite/granite-4.0-h-micro`, `ibm-granite/granite-4.0-micro` | IBM Granite 4.0 micro models: hybrid Mamba–MoE (`h-micro`) and dense (`micro`) variants. Enterprise-focused reasoning models | | **Sarvam 2** (30B-A2B, 105B-A10B) | `sarvamai/sarvam-2` | Sarvam's Mixture-of-Experts models. The 105B variant uses MLA (Multi-head Latent Attention) and the 30B variant uses GQA, both with 128 routed experts. | ================================================ FILE: docs/supported_models/text_generation/index.rst ================================================ Text Generation =============== Models for generating text from text or multimodal inputs. .. toctree:: :maxdepth: 1 generative_models.md multimodal_language_models.md diffusion_language_models.md ================================================ FILE: docs/supported_models/text_generation/multimodal_language_models.md ================================================ # Multimodal Language Models These models accept multi-modal inputs (e.g., images and text) and generate text output. They augment language models with multimodal encoders. ## Example launch Command ```shell python3 -m sglang.launch_server \ --model-path meta-llama/Llama-3.2-11B-Vision-Instruct \ # example HF/local path --host 0.0.0.0 \ --port 30000 \ ``` > See the [OpenAI APIs section](https://docs.sglang.io/basic_usage/openai_api_vision.html) for how to send multimodal requests. ## Supported models Below the supported models are summarized in a table. If you are unsure if a specific architecture is implemented, you can search for it via GitHub. For example, to search for `Qwen2_5_VLForConditionalGeneration`, use the expression: ``` repo:sgl-project/sglang path:/^python\/sglang\/srt\/models\// Qwen2_5_VLForConditionalGeneration ``` in the GitHub search bar. | Model Family (Variants) | Example HuggingFace Identifier | Description | Notes | |----------------------------|--------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------| | **Qwen-VL** | `Qwen/Qwen3-VL-235B-A22B-Instruct` | Alibaba's vision-language extension of Qwen; for example, Qwen2.5-VL (7B and larger variants) can analyze and converse about image content. | | | **DeepSeek-VL2** | `deepseek-ai/deepseek-vl2` | Vision-language variant of DeepSeek (with a dedicated image processor), enabling advanced multimodal reasoning on image and text inputs. | | | **DeepSeek-OCR / OCR-2** | `deepseek-ai/DeepSeek-OCR-2` | OCR-focused DeepSeek models for document understanding and text extraction. | Use `--trust-remote-code`. | | **Janus-Pro** (1B, 7B) | `deepseek-ai/Janus-Pro-7B` | DeepSeek's open-source multimodal model capable of both image understanding and generation. Janus-Pro employs a decoupled architecture for separate visual encoding paths, enhancing performance in both tasks. | | | **MiniCPM-V / MiniCPM-o** | `openbmb/MiniCPM-V-2_6` | MiniCPM-V (2.6, ~8B) supports image inputs, and MiniCPM-o adds audio/video; these multimodal LLMs are optimized for end-side deployment on mobile/edge devices. | | | **Llama 3.2 Vision** (11B) | `meta-llama/Llama-3.2-11B-Vision-Instruct` | Vision-enabled variant of Llama 3 (11B) that accepts image inputs for visual question answering and other multimodal tasks. | | | **LLaVA** (v1.5 & v1.6) | *e.g.* `liuhaotian/llava-v1.5-13b` | Open vision-chat models that add an image encoder to LLaMA/Vicuna (e.g. LLaMA2 13B) for following multimodal instruction prompts. | | | **LLaVA-NeXT** (8B, 72B) | `lmms-lab/llava-next-72b` | Improved LLaVA models (with an 8B Llama3 version and a 72B version) offering enhanced visual instruction-following and accuracy on multimodal benchmarks. | | | **LLaVA-OneVision** | `lmms-lab/llava-onevision-qwen2-7b-ov` | Enhanced LLaVA variant integrating Qwen as the backbone; supports multiple images (and even video frames) as inputs via an OpenAI Vision API-compatible format. | | | **Gemma 3 (Multimodal)** | `google/gemma-3-4b-it` | Gemma 3's larger models (4B, 12B, 27B) accept images (each image encoded as 256 tokens) alongside text in a combined 128K-token context. | | | **Kimi-VL** (A3B) | `moonshotai/Kimi-VL-A3B-Instruct` | Kimi-VL is a multimodal model that can understand and generate text from images. | | | **Mistral-Small-3.1-24B** | `mistralai/Mistral-Small-3.1-24B-Instruct-2503` | Mistral 3.1 is a multimodal model that can generate text from text or images input. It also supports tool calling and structured output. | | | **Phi-4-multimodal-instruct** | `microsoft/Phi-4-multimodal-instruct` | Phi-4-multimodal-instruct is the multimodal variant of the Phi-4-mini model, enhanced with LoRA for improved multimodal capabilities. It supports text, vision and audio modalities in SGLang. | | | **MiMo-VL** (7B) | `XiaomiMiMo/MiMo-VL-7B-RL` | Xiaomi's compact yet powerful vision-language model featuring a native resolution ViT encoder for fine-grained visual details, an MLP projector for cross-modal alignment, and the MiMo-7B language model optimized for complex reasoning tasks. | | | **GLM-4.5V** (106B) / **GLM-4.1V**(9B) | `zai-org/GLM-4.5V` | GLM-4.5V and GLM-4.1V-Thinking: Towards Versatile Multimodal Reasoning with Scalable Reinforcement Learning | Use `--chat-template glm-4v` | | **GLM-OCR** | `zai-org/GLM-OCR` | GLM-OCR: A fast and accurate general OCR model | | | **DotsVLM** (General/OCR) | `rednote-hilab/dots.vlm1.inst` | RedNote's vision-language model built on a 1.2B vision encoder and DeepSeek V3 LLM, featuring NaViT vision encoder trained from scratch with dynamic resolution support and enhanced OCR capabilities through structured image data training. | | | **DotsVLM-OCR** | `rednote-hilab/dots.ocr` | Specialized OCR variant of DotsVLM optimized for optical character recognition tasks with enhanced text extraction and document understanding capabilities. | Don't use `--trust-remote-code` | | **NVILA** (8B, 15B, Lite-2B, Lite-8B, Lite-15B) | `Efficient-Large-Model/NVILA-8B` | `chatml` | NVILA explores the full stack efficiency of multi-modal design, achieving cheaper training, faster deployment and better performance. | | **NVIDIA Nemotron Nano 2.0 VL** | `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` | NVIDIA Nemotron Nano v2 VL enables multi-image reasoning and video understanding, along with strong document intelligence, visual Q&A and summarization capabilities. It builds on Nemotron Nano V2, a hybrid Mamba-Transformer LLM, in order to achieve higher inference throughput in long document and video scenarios. | Use `--trust-remote-code`. You may need to adjust `--max-mamba-cache-size` [default is 512] to fit memory constraints. | | **Ernie4.5-VL** | `baidu/ERNIE-4.5-VL-28B-A3B-PT` | Baidu's vision-language models(28B,424B). Support image and video comprehension, and also support thinking. | | | **JetVLM** | | JetVLM is an vision-language model designed for high-performance multimodal understanding and generation tasks built upon Jet-Nemotron. | Coming soon | | **Step3-VL** (10B) | `stepfun-ai/Step3-VL-10B` | StepFun's lightweight open-source 10B parameter VLM for multimodal intelligence, excelling in visual perception, complex reasoning, and human alignment. | | | **Qwen3-Omni** | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | Alibaba's omni-modal MoE model. Currently supports the **Thinker** component (multimodal understanding for text, images, audio, and video), while the **Talker** component (audio generation) is not yet supported. | | ## Video Input Support SGLang supports video input for Vision-Language Models (VLMs), enabling temporal reasoning tasks such as video question answering, captioning, and holistic scene understanding. Video clips are decoded, key frames are sampled, and the resulting tensors are batched together with the text prompt, allowing multimodal inference to integrate visual and linguistic context. | Model Family | Example Identifier | Video notes | |--------------|--------------------|-------------| | **Qwen-VL** (Qwen2-VL, Qwen2.5-VL, Qwen3-VL, Qwen3-Omni) | `Qwen/Qwen3-VL-235B-A22B-Instruct` | The processor gathers `video_data`, runs Qwen's frame sampler, and merges the resulting features with text tokens before inference. | | **GLM-4v** (4.5V, 4.1V, MOE) | `zai-org/GLM-4.5V` | Video clips are read with Decord, converted to tensors, and passed to the model alongside metadata for rotary-position handling. | | **NVILA** (Full & Lite) | `Efficient-Large-Model/NVILA-8B` | The runtime samples eight frames per clip and attaches them to the multimodal request when `video_data` is present. | | **LLaVA video variants** (LLaVA-NeXT-Video, LLaVA-OneVision) | `lmms-lab/LLaVA-NeXT-Video-7B` | The processor routes video prompts to the LlavaVid video-enabled architecture, and the provided example shows how to query it with `sgl.video(...)` clips. | | **NVIDIA Nemotron Nano 2.0 VL** | `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` | The processor samples at 2 FPS, at a max of 128 frames, as per model training. The model uses [EVS](../../python/sglang/srt/multimodal/evs/README.md), a pruning method that removes redundant tokens from video embeddings. By default `video_pruning_rate=0.7`. Change this by providing: `--json-model-override-args '{"video_pruning_rate": 0.0}'` to disable EVS, for example. | | **JetVLM** | | The runtime samples eight frames per clip and attaches them to the multimodal request when `video_data` is present. | Use `sgl.video(path, num_frames)` when building prompts to attach clips from your SGLang programs. Example OpenAI-compatible request that sends a video clip: ```python import requests url = "http://localhost:30000/v1/chat/completions" data = { "model": "Qwen/Qwen3-VL-30B-A3B-Instruct", "messages": [ { "role": "user", "content": [ {"type": "text", "text": "What’s happening in this video?"}, { "type": "video_url", "video_url": { "url": "https://github.com/sgl-project/sgl-test-files/raw/refs/heads/main/videos/jobs_presenting_ipod.mp4" }, }, ], } ], "max_tokens": 300, } response = requests.post(url, json=data) print(response.text) ``` ## Usage Notes ### Performance Optimization For multimodal models, you can use the `--keep-mm-feature-on-device` flag to optimize for latency at the cost of increased GPU memory usage: - **Default behavior**: Multimodal feature tensors are moved to CPU after processing to save GPU memory - **With `--keep-mm-feature-on-device`**: Feature tensors remain on GPU, reducing device-to-host copy overhead and improving latency, but consuming more GPU memory Use this flag when you have sufficient GPU memory and want to minimize latency for multimodal inference. ### Multimodal Inputs Limitation - **Use `--mm-process-config '{"image":{"max_pixels":1048576},"video":{"fps":3,"max_pixels":602112,"max_frames":60}}'`**: To set `image`, `video`, and `audio` input limits. This can reduce GPU memory usage, improve inference speed, and help to avoid OOM, but may impact model performance, thus set a proper value based on your specific use case. Currently, only `qwen_vl` supports this config. Please refer to [qwen_vl processor](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/multimodal/processors/qwen_vl.py) for understanding the meaning of each parameter. ### Bidirectional Attention in Multimodal Model Serving **Note for serving the Gemma-3 multimodal model**: As mentioned in [Welcome Gemma 3: Google's all new multimodal, multilingual, long context open LLM ](https://huggingface.co/blog/gemma3#multimodality), Gemma-3 employs bidirectional attention between image tokens during the prefill phase. Currently, SGLang only supports bidirectional attention when using the Triton Attention Backend. Note, however, that SGLang's current bidirectional attention implementation is incompatible with both CUDA Graph and Chunked Prefill. To enable bidirectional attention, you can use the `TritonAttnBackend` while disabling CUDA Graph and Chunked Prefill. Example launch command: ```shell python -m sglang.launch_server \ --model-path google/gemma-3-4b-it \ --host 0.0.0.0 --port 30000 \ --enable-multimodal \ --dtype bfloat16 --triton-attention-reduce-in-fp32 \ --attention-backend triton \ # Use Triton attention backend --disable-cuda-graph \ # Disable Cuda Graph --chunked-prefill-size -1 # Disable Chunked Prefill ``` If higher serving performance is required and a certain degree of accuracy loss is acceptable, you may choose to use other attention backends, and you can also enable features like CUDA Graph and Chunked Prefill for better performance, but note that the model will fall back to using causal attention instead of bidirectional attention. ================================================ FILE: docs/wrap_run_llm.py ================================================ import os import re def insert_runllm_widget(html_content): # RunLLM Widget script to be inserted widget_script = """ """ # Find the closing body tag and insert the widget script before it return re.sub(r"", f"{widget_script}\n", html_content) def process_html_files(build_dir): for root, dirs, files in os.walk(build_dir): for file in files: if file.endswith(".html"): file_path = os.path.join(root, file) # Read the HTML file with open(file_path, "r", encoding="utf-8") as f: content = f.read() # Insert the RunLLM widget modified_content = insert_runllm_widget(content) # Write back the modified content with open(file_path, "w", encoding="utf-8") as f: f.write(modified_content) def main(): # Get the build directory path build_dir = os.path.join( os.path.dirname(os.path.abspath(__file__)), "_build", "html" ) # Process all HTML files if os.path.exists(build_dir): process_html_files(build_dir) else: print(f"Build directory not found: {build_dir}") if __name__ == "__main__": main() ================================================ FILE: examples/assets/.gitignore ================================================ !example_image.png ================================================ FILE: examples/chat_template/qwen3_reranker.jinja ================================================ <|im_start|>system Judge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|> <|im_start|>user : {{ instruct | default("Given a web search query, retrieve relevant passages that answer the query.") }} : {{ messages[0]["content"] }} : {{ messages[1]["content"] }}<|im_end|> <|im_start|>assistant{{ '\n' }} ================================================ FILE: examples/chat_template/qwen3_vl_reranker.jinja ================================================ {#- Qwen3-VL-Reranker chat template for multimodal reranking -#} {#- This template formats query-document pairs for yes/no relevance judgment -#} {#- Supports text, images, and videos in both query and documents -#} <|im_start|>system Judge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|> <|im_start|>user : {{ instruct | default("Given a search query, retrieve relevant candidates that answer the query.") }} {#- Process query content -#} : {%- for content in query -%} {%- if content.type == 'image' or 'image' in content or 'image_url' in content -%} <|vision_start|><|image_pad|><|vision_end|> {%- elif content.type == 'video' or 'video' in content -%} <|vision_start|><|video_pad|><|vision_end|> {%- elif 'text' in content -%} {{ content.text }} {%- elif content.type == 'text' -%} {{ content.text }} {%- endif -%} {%- endfor %} {#- Process document content -#} {{ '\n' }}: {%- for content in document -%} {%- if content.type == 'image' or 'image' in content or 'image_url' in content -%} <|vision_start|><|image_pad|><|vision_end|> {%- elif content.type == 'video' or 'video' in content -%} <|vision_start|><|video_pad|><|vision_end|> {%- elif 'text' in content -%} {{ content.text }} {%- elif content.type == 'text' -%} {{ content.text }} {%- endif -%} {%- endfor %}<|im_end|> <|im_start|>assistant{{ '\n' }} ================================================ FILE: examples/chat_template/tool_chat_template_deepseekr1.jinja ================================================ {% if not add_generation_prompt is defined %} {% set add_generation_prompt = false %} {% endif %} {% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='', is_first_sp=true, is_last_user=false) %} {%- for message in messages %} {%- if message['role'] == 'system' %} {%- if ns.is_first_sp %} {% set ns.system_prompt = ns.system_prompt + message['content'] %} {% set ns.is_first_sp = false %} {%- else %} {% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %} {%- endif %} {%- endif %} {%- endfor %} {# --- Append tool descriptions if tools are defined --- #} {% if tools is defined and tools is not none %} {% set tool_ns = namespace(text='You are a helpful assistant with tool calling capabilities. ' 'When a tool call is needed, you MUST use the following format to issue the call:\n' '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>FUNCTION_NAME\n' '```json\n{"param1": "value1", "param2": "value2"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>\n\n' 'Make sure the JSON is valid.' '## Tools\n\n### Function\n\nYou have the following functions available:\n\n') %} {% for tool in tools %} {% set tool_ns.text = tool_ns.text + '- `' + tool['name'] + '`:\n```json\n' + (tool | tojson) + '\n```\n' %} {% endfor %} {% set ns.system_prompt = ns.system_prompt + '\n\n' + tool_ns.text %} {% endif %} {{ bos_token }} {{ ns.system_prompt }} {%- for message in messages %} {% set content = message['content'] %} {%- if message['role'] == 'user' %} {%- set ns.is_tool = false -%} {%- set ns.is_first = false -%} {%- set ns.is_last_user = true -%} {{'<｜User｜>' + content + '<｜Assistant｜>'}} {%- endif %} {%- if message['role'] == 'assistant' %} {% if '' in content %} {% set content = content.split('')[-1] %} {% endif %} {% endif %} {%- if message['role'] == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] is not none %} {%- set ns.is_last_user = false -%} {%- if ns.is_tool %} {{'<｜tool▁outputs▁end｜>'}} {%- endif %} {%- set ns.is_first = false %} {%- set ns.is_tool = false -%} {%- set ns.is_output_first = true %} {%- for tool in message['tool_calls'] %} {%- set tool_type = tool['type'] if tool['type'] is defined else 'function' -%} {%- set tool_args = tool['function']['arguments'] if tool['function']['arguments'] is string else tool['function']['arguments'] | tojson -%} {%- if not ns.is_first %} {%- if content is none %} {{'<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool_type + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool_args + '\n' + '```' + '<｜tool▁call▁end｜>'}} {%- else %} {{content + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool_type + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool_args + '\n' + '```' + '<｜tool▁call▁end｜>'}} {%- endif %} {%- set ns.is_first = true -%} {%- else %} {{'\n' + '<｜tool▁call▁begin｜>' + tool_type + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool_args + '\n' + '```' + '<｜tool▁call▁end｜>'}} {%- endif %} {%- endfor %} {{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}} {%- endif %} {%- if message['role'] == 'assistant' and (message['tool_calls'] is not defined or message['tool_calls'] is none)%} {%- set ns.is_last_user = false -%} {%- if ns.is_tool %} {{'<｜tool▁outputs▁end｜>' + content + '<｜end▁of▁sentence｜>'}} {%- set ns.is_tool = false -%} {%- else %} {{content + '<｜end▁of▁sentence｜>'}} {%- endif %} {%- endif %} {%- if message['role'] == 'tool' %} {%- set ns.is_last_user = false -%} {%- set ns.is_tool = true -%} {%- if ns.is_output_first %} {{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + content + '<｜tool▁output▁end｜>'}} {%- set ns.is_output_first = false %} {%- else %} {{'\n<｜tool▁output▁begin｜>' + content + '<｜tool▁output▁end｜>'}} {%- endif %} {%- endif %} {%- endfor -%} {% if ns.is_tool %} {{'<｜tool▁outputs▁end｜>'}} {% endif %} {% if add_generation_prompt and not ns.is_last_user and not ns.is_tool %} {{'<｜Assistant｜>'}} {% endif %} ================================================ FILE: examples/chat_template/tool_chat_template_deepseekv3.jinja ================================================ {% if not add_generation_prompt is defined %} {% set add_generation_prompt = false %} {% endif %} {% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='', is_first_sp=true, is_last_user=false) %} {%- for message in messages %} {%- if message['role'] == 'system' %} {%- if ns.is_first_sp %} {% set ns.system_prompt = ns.system_prompt + message['content'] %} {% set ns.is_first_sp = false %} {%- else %} {% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %} {%- endif %} {%- endif %} {%- endfor -%} {# --- Append tool descriptions if tools are defined --- #} {% if tools is defined and tools is not none %} {% set tool_ns = namespace(text='You are a helpful assistant with tool calling capabilities. ' 'When a tool call is needed, you MUST use the following format to issue the call:\n' '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>FUNCTION_NAME\n' '```json\n{"param1": "value1", "param2": "value2"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>\n\n' 'Make sure the JSON is valid.' '## Tools\n\n### Function\n\nYou have the following functions available:\n\n') %} {% for tool in tools %} {% set tool_ns.text = tool_ns.text + '\n```json\n' + (tool | tojson) + '\n```\n' %} {% endfor %} {% set ns.system_prompt = ns.system_prompt + '\n\n' + tool_ns.text %} {% endif %} {{- bos_token }} {{- ns.system_prompt }} {%- for message in messages %} {%- if message['role'] == 'user' %} {%- set ns.is_tool = false -%} {%- set ns.is_first = false -%} {%- set ns.is_last_user = true -%} {{'<｜User｜>' + message['content'] + '<｜Assistant｜>'}} {%- endif %} {%- if message['role'] == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] is not none %} {%- set ns.is_last_user = false -%} {%- if ns.is_tool %} {{- '<｜tool▁outputs▁end｜>'}} {%- endif %} {%- set ns.is_first = false %} {%- set ns.is_tool = false -%} {%- set ns.is_output_first = true %} {%- for tool in message['tool_calls'] %} {%- set formatted_args = tool['function']['arguments'] if tool['function']['arguments'] is string else tool['function']['arguments']|tojson %} {%- if not ns.is_first %} {%- if message['content'] is none %} {{- '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + formatted_args + '\n' + '```' + '<｜tool▁call▁end｜>'}} {%- else %} {{- message['content'] + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + formatted_args + '\n' + '```' + '<｜tool▁call▁end｜>'}} {%- endif %} {%- set ns.is_first = true -%} {%- else %} {{- '\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + formatted_args + '\n' + '```' + '<｜tool▁call▁end｜>'}} {%- endif %} {%- endfor %} {{- '<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}} {%- endif %} {%- if message['role'] == 'assistant' and (message['tool_calls'] is not defined or message['tool_calls'] is none)%} {%- set ns.is_last_user = false -%} {%- if ns.is_tool %} {{- '<｜tool▁outputs▁end｜>' + message['content'] + '<｜end▁of▁sentence｜>'}} {%- set ns.is_tool = false -%} {%- else %} {% set content = message['content'] %} {{- content + '<｜end▁of▁sentence｜>'}} {%- endif %} {%- endif %} {%- if message['role'] == 'tool' %} {%- set ns.is_last_user = false -%} {%- set ns.is_tool = true -%} {%- if ns.is_output_first %} {{- 'Use the results below to formulate an answer to the user question unless additional information is needed.' }} {{- '<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}} {%- set ns.is_output_first = false %} {%- else %} {{- '\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}} {%- endif %} {%- endif %} {%- endfor -%} {% if ns.is_tool %} {{- '<｜tool▁outputs▁end｜>'}} {% endif %} {% if add_generation_prompt and not ns.is_last_user and not ns.is_tool %} {{- '<｜Assistant｜>'}} {% endif %} ================================================ FILE: examples/chat_template/tool_chat_template_deepseekv31.jinja ================================================ {% if not add_generation_prompt is defined %} {% set add_generation_prompt = false %} {% endif %} {% if not thinking is defined %} {% set thinking = false %} {% endif %} {% set ns = namespace(is_first=false, is_tool=false, system_prompt='', is_first_sp=true, is_last_user=false) %} {%- for message in messages %} {%- if message['role'] == 'system' %} {%- if ns.is_first_sp %} {% set ns.system_prompt = ns.system_prompt + message['content'] %} {% set ns.is_first_sp = false %} {%- else %} {% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %} {%- endif %} {%- endif %} {%- endfor %} {% if tools is defined and tools is not none %} {% set tool_ns = namespace(text='## Tools\nYou have access to the following tools:\n') %} {% for tool in tools %} {% if tool.function.description is not none %} {% set tool_ns.text = tool_ns.text + '\n### ' + tool.function.name + '\nDescription: ' + tool.function.description + '\n\nParameters: ' + (tool.function.parameters | tojson) + '\n' %} {% else %} {% set tool_ns.text = tool_ns.text + '\n### ' + tool.function.name + '\n\nParameters: ' + (tool.function.parameters | tojson) + '\n' %} {% endif %} {% endfor %} {% set tool_ns.text = tool_ns.text + "\nIMPORTANT: ALWAYS adhere to this exact format for tool use:\n<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>tool_call_name<｜tool▁sep｜>tool_call_arguments<｜tool▁call▁end｜>{{additional_tool_calls}}<｜tool▁calls▁end｜>\n\nWhere:\n\n- `tool_call_name` must be an exact match to one of the available tools\n- `tool_call_arguments` must be valid JSON that strictly follows the tool's Parameters Schema\n- For multiple tool calls, chain them directly without separators or spaces\n" %} {% set ns.system_prompt = ns.system_prompt + '\n\n' + tool_ns.text %} {% endif %} {{ bos_token }}{{ ns.system_prompt }} {%- for message in messages %} {%- if message['role'] == 'user' %} {%- set ns.is_tool = false -%} {%- set ns.is_first = false -%} {%- set ns.is_last_user = true -%} {{'<｜User｜>' + message['content']}} {%- endif %} {%- if message['role'] == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] is not none %} {%- if ns.is_last_user %} {{'<｜Assistant｜>'}} {%- endif %} {%- set ns.is_last_user = false -%} {%- set ns.is_first = false %} {%- set ns.is_tool = false -%} {%- for tool in message['tool_calls'] %} {%- set formatted_args = tool['function']['arguments'] if tool['function']['arguments'] is string else tool['function']['arguments']|tojson %} {%- if not ns.is_first %} {%- if message['content'] is none %} {{'<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>'+ tool['function']['name'] + '<｜tool▁sep｜>' + formatted_args + '<｜tool▁call▁end｜>'}} {%- else %} {{message['content'] + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['function']['name'] + '<｜tool▁sep｜>' + formatted_args + '<｜tool▁call▁end｜>'}} {%- endif %} {%- set ns.is_first = true -%} {%- else %} {{'<｜tool▁call▁begin｜>'+ tool['function']['name'] + '<｜tool▁sep｜>' + formatted_args + '<｜tool▁call▁end｜>'}} {%- endif %} {%- endfor %} {{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}} {%- endif %} {%- if message['role'] == 'assistant' and (message['tool_calls'] is not defined or message['tool_calls'] is none) %} {%- if ns.is_last_user %} {{'<｜Assistant｜>'}} {%- if message['prefix'] is defined and message['prefix'] and thinking %} {{''}} {%- else %} {{''}} {%- endif %} {%- endif %} {%- set ns.is_last_user = false -%} {%- if ns.is_tool %} {{message['content'] + '<｜end▁of▁sentence｜>'}} {%- set ns.is_tool = false -%} {%- else %} {%- set content = message['content'] -%} {%- if '' in content %} {%- set content = content.split('', 1)[1] -%} {%- endif %} {{content + '<｜end▁of▁sentence｜>'}} {%- endif %} {%- endif %} {%- if message['role'] == 'tool' %} {%- set ns.is_last_user = false -%} {%- set ns.is_tool = true -%} {{'<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}} {%- endif %} {%- endfor -%} {%- if add_generation_prompt and ns.is_last_user and not ns.is_tool %} {{'<｜Assistant｜>'}} {%- if not thinking %} {{''}} {%- else %} {{''}} {%- endif %} {% endif %} ================================================ FILE: examples/chat_template/tool_chat_template_deepseekv32.jinja ================================================ {% if not add_generation_prompt is defined %} {% set add_generation_prompt = false %} {% endif %} {% if not thinking is defined %} {% set thinking = false %} {% endif %} {% set ns = namespace(is_first=false, is_tool=false, system_prompt='', is_first_sp=true, is_last_user=false, is_only_sys=false, is_prefix=false) %} {%- for message in messages %} {%- if message['role'] == 'system' %} {%- if ns.is_first_sp %} {% set ns.system_prompt = ns.system_prompt + message['content'] %} {% set ns.is_first_sp = false %} {%- else %} {% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %} {%- endif %} {% set ns.is_only_sys = true %} {%- endif %} {%- endfor %} {% if tools is defined and tools is not none %} {% set tool_ns = namespace(text='## Tools\nYou have access to the following tools:\n') %} {% for tool in tools %} {% set tool_ns.text = tool_ns.text + '\n### ' + tool.function.name + '\nDescription: ' + tool.function.description + '\n\nParameters: ' + (tool.function.parameters | tojson) + '\n' %} {% endfor %} {% set tool_ns.text = tool_ns.text + "\nIMPORTANT: ALWAYS adhere to this exact format for tool use:\n<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>tool_call_name<｜tool▁sep｜>tool_call_arguments<｜tool▁call▁end｜>{{additional_tool_calls}}<｜tool▁calls▁end｜>\n\nWhere:\n\n- `tool_call_name` must be an exact match to one of the available tools\n- `tool_call_arguments` must be valid JSON that strictly follows the tool's Parameters Schema\n- For multiple tool calls, chain them directly without separators or spaces\n" %} {% set ns.system_prompt = ns.system_prompt + '\n\n' + tool_ns.text %} {% endif %} {{ bos_token }}{{ ns.system_prompt }} {%- for message in messages %} {%- if message['role'] == 'user' %} {%- set ns.is_tool = false -%} {%- set ns.is_first = false -%} {%- set ns.is_last_user = true -%} {{'<｜User｜>' + message['content']}} {%- endif %} {%- if message['role'] == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] is not none %} {%- if ns.is_last_user or ns.is_only_sys %} {{'<｜Assistant｜>'}} {%- endif %} {%- set ns.is_last_user = false -%} {%- set ns.is_first = false %} {%- set ns.is_tool = false -%} {%- for tool in message['tool_calls'] %} {%- set formatted_args = tool['function']['arguments'] if tool['function']['arguments'] is string else tool['function']['arguments']|tojson %} {%- if not ns.is_first %} {%- if message['content'] is none %} {{'<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>'+ tool['function']['name'] + '<｜tool▁sep｜>' + formatted_args + '<｜tool▁call▁end｜>'}} {%- else %} {{message['content'] + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['function']['name'] + '<｜tool▁sep｜>' + formatted_args + '<｜tool▁call▁end｜>'}} {%- endif %} {%- set ns.is_first = true -%} {%- else %} {{'<｜tool▁call▁begin｜>'+ tool['function']['name'] + '<｜tool▁sep｜>' + formatted_args + '<｜tool▁call▁end｜>'}} {%- endif %} {%- endfor %} {{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}} {%- endif %} {%- if message['role'] == 'assistant' and (message['tool_calls'] is not defined or message['tool_calls'] is none) %} {%- if ns.is_last_user %} {{'<｜Assistant｜>'}} {%- if message['prefix'] is defined and message['prefix'] and thinking %} {{''}} {%- else %} {{''}} {%- endif %} {%- endif %} {%- if message['prefix'] is defined and message['prefix'] %} {%- set ns.is_prefix = true -%} {%- endif %} {%- set ns.is_last_user = false -%} {%- if ns.is_tool %} {{message['content'] + '<｜end▁of▁sentence｜>'}} {%- set ns.is_tool = false -%} {%- else %} {%- set content = message['content'] -%} {%- if '' in content %} {%- set content = content.split('', 1)[1] -%} {%- endif %} {{content + '<｜end▁of▁sentence｜>'}} {%- endif %} {%- endif %} {%- if message['role'] == 'tool' %} {%- set ns.is_last_user = false -%} {%- set ns.is_tool = true -%} {{'<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}} {%- endif %} {%- if message['role'] != 'system' %} {% set ns.is_only_sys = false %} {%- endif %} {%- endfor -%} {% if add_generation_prompt and not ns.is_tool%} {% if ns.is_last_user or ns.is_only_sys or not ns.is_prefix %} {{'<｜Assistant｜>'}} {%- if not thinking %} {{''}} {%- else %} {{''}} {%- endif %} {% endif %} {% endif %} ================================================ FILE: examples/chat_template/tool_chat_template_llama3.1_json.jinja ================================================ {# Copied from https://github.com/vllm-project/vllm/blob/main/examples/tool_chat_template_llama3.1_json.jinja to enable better model response. #} {{- bos_token }} {%- if custom_tools is defined %} {%- set tools = custom_tools %} {%- endif %} {%- if not tools_in_user_message is defined %} {#- Llama 3.1 doesn't pass all tests if the tools are in the system prompt #} {%- set tools_in_user_message = true %} {%- endif %} {%- if not date_string is defined %} {%- if strftime_now is defined %} {%- set date_string = strftime_now("%d %b %Y") %} {%- else %} {%- set date_string = "26 Jul 2024" %} {%- endif %} {%- endif %} {%- if not tools is defined %} {%- set tools = none %} {%- endif %} {#- This block extracts the system message, so we can slot it into the right place. #} {%- if messages[0]['role'] == 'system' %} {%- if messages[0]['content'] is string %} {%- set system_message = messages[0]['content']|trim %} {%- else %} {%- set system_message = messages[0]['content'][0]['text']|trim %} {%- endif %} {%- set messages = messages[1:] %} {%- else %} {%- if tools is not none %} {%- set system_message = "You are a helpful assistant with tool calling capabilities. Only reply with a tool call if the function exists in the library provided by the user. If it doesn't exist, just reply directly in natural language. When you receive a tool call response, use the output to format an answer to the original user question." %} {%- else %} {%- set system_message = "" %} {%- endif %} {%- endif %} {#- System message #} {{- "<|start_header_id|>system<|end_header_id|>\n\n" }} {%- if tools is not none %} {{- "Environment: ipython\n" }} {%- endif %} {{- "Cutting Knowledge Date: December 2023\n" }} {{- "Today Date: " + date_string + "\n\n" }} {%- if tools is not none and not tools_in_user_message %} {{- "You have access to the following functions. To call a function, please respond with JSON for a function call. " }} {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. ' }} {{- "Do not use variables.\n\n" }} {%- for t in tools %} {{- t | tojson(indent=4) }} {{- "\n\n" }} {%- endfor %} {%- endif %} {{- system_message }} {{- "<|eot_id|>" }} {#- Custom tools are passed in a user message with some extra guidance #} {%- if tools_in_user_message and not tools is none %} {#- Extract the first user message so we can plug it in here #} {%- if messages | length != 0 %} {%- if messages[0]['content'] is string %} {%- set first_user_message = messages[0]['content']|trim %} {%- else %} {%- set first_user_message = messages[0]['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | map('trim') | join('\n') %} {%- endif %} {%- set messages = messages[1:] %} {%- else %} {{- raise_exception("Cannot put tools in the first user message when there's no first user message!") }} {%- endif %} {{- '<|start_header_id|>user<|end_header_id|>\n\n' -}} {{- "Given the following functions, please respond with a JSON for a function call " }} {{- "with its proper arguments that best answers the given prompt.\n\n" }} {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. ' }} {{- "Do not use variables.\n\n" }} {%- for t in tools %} {{- t | tojson(indent=4) }} {{- "\n\n" }} {%- endfor %} {{- first_user_message + "<|eot_id|>"}} {%- endif %} {%- for message in messages %} {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %} {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' }} {%- if message['content'] is string %} {{- message['content'] | trim}} {%- else %} {%- for content in message['content'] %} {%- if content['type'] == 'text' %} {{- content['text'] | trim }} {%- endif %} {%- endfor %} {%- endif %} {{- '<|eot_id|>' }} {%- elif 'tool_calls' in message %} {%- if not message.tool_calls|length == 1 %} {{- raise_exception("This model only supports single tool-calls at once!") }} {%- endif %} {%- set tool_call = message.tool_calls[0].function %} {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' -}} {{- '{"name": "' + tool_call.name + '", ' }} {{- '"parameters": ' }} {{- tool_call.arguments | tojson }} {{- "}" }} {{- "<|eot_id|>" }} {%- elif message.role == "tool" or message.role == "ipython" %} {{- "<|start_header_id|>ipython<|end_header_id|>\n\n" }} {%- if message.content is string %} {{- { "output": message.content } | tojson }} {%- else %} {%- for content in message['content'] %} {%- if content['type'] == 'text' %} {{- { "output": content['text'] } | tojson }} {%- endif %} {%- endfor %} {%- endif %} {{- "<|eot_id|>" }} {%- endif %} {%- endfor %} {%- if add_generation_prompt %} {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }} {%- endif %} ================================================ FILE: examples/chat_template/tool_chat_template_llama4_pythonic.jinja ================================================ {# Copied from https://github.com/wukaixingxp/vllm/blob/8a32e2a6e452a03c0e8222e3876ad6086cbf581f/examples/tool_chat_template_llama4_pythonic.jinja to enable better model response. #} {{- bos_token }} {%- if custom_tools is defined and custom_tools %} {%- set tools = custom_tools %} {%- endif %} {%- if tools is defined and tools %} {%- set tool_definition = tool_definition ~ (tools | tojson(indent=4)) %} {%- else %} {%- set tools = none %} {%- endif %} {#- This block extracts the system message, so we can slot it into the right place. #} {%- if messages[0]['role'] == 'system' %} {%- set user_provided_system_message = true %} {%- if messages[0]['content'] is string %} {%- set system_message = messages[0]['content']|trim %} {%- else %} {%- set system_message = messages[0]['content'][0]['text']|trim %} {%- endif %} {%- set messages = messages[1:] %} {%- else %} {%- if tools is not none %} {#- Since not system_message was provided by user, if tool is provided, system_message is now default tool system message #} {#- This system message is from llama website:https://www.llama.com/docs/model-cards-and-prompt-formats/llama4/ #} {%- set system_message = "You are a helpful assistant and an expert in function composition. You can answer general questions using your internal knowledge OR invoke functions when necessary. Follow these strict guidelines:\n\n1. FUNCTION CALLS:\n- ONLY use functions that are EXPLICITLY listed in the function list below\n- If NO functions are listed (empty function list []), respond ONLY with internal knowledge or \"I don't have access to [Unavailable service] information\"\n- If a function is not in the list, respond ONLY with internal knowledge or \"I don't have access to [Unavailable service] information\"\n- If ALL required parameters are present AND the query EXACTLY matches a listed function's purpose: output ONLY the function call(s)\n- Use exact format: [func_name1(param1=value1, param2=value2), func_name2(...)]\nExamples:\nCORRECT: [get_weather(location=\"Vancouver\"), calculate_route(start=\"Boston\", end=\"New York\")] <- Only if get_weather and calculate_route are in function list\nINCORRECT: get_weather(location=\"New York\")\nINCORRECT: Let me check the weather: [get_weather(location=\"New York\")]\nINCORRECT: [get_events(location=\"Singapore\")] <- If function not in list\n\n2. RESPONSE RULES:\n- For pure function requests matching a listed function: ONLY output the function call(s)\n- For knowledge questions: ONLY output text\n- For missing parameters: ONLY request the specific missing parameters\n- For unavailable services (not in function list): output ONLY with internal knowledge or \"I don't have access to [Unavailable service] information\". Do NOT execute a function call.\n- If the query asks for information beyond what a listed function provides: output ONLY with internal knowledge about your limitations\n- NEVER combine text and function calls in the same response\n- NEVER suggest alternative functions when the requested service is unavailable\n- NEVER create or invent new functions not listed below\n\n3. STRICT BOUNDARIES:\n- ONLY use functions from the list below - no exceptions\n- NEVER use a function as an alternative to unavailable information\n- NEVER call functions not present in the function list\n- NEVER add explanatory text to function calls\n- NEVER respond with empty brackets\n- Use proper Python/JSON syntax for function calls\n- Check the function list carefully before responding\n\n4. TOOL RESPONSE HANDLING:\n- When receiving tool responses: provide concise, natural language responses\n- Don't repeat tool response verbatim\n- Don't add supplementary information\n\nHere is a list of functions in JSON format that you can invoke:\n" %} {%- else %} {%- set system_message = "" %} {%- endif %} {%- endif %} {#- Now writing the system message: use the user provided system message if user_provided_system_message, else default tool system message if tools presented #} {%- if system_message %} {#- always use user provided system message to override default tool system message #} {{- "<|header_start|>system<|header_end|>\n\n" }} {{- system_message }} {%- if user_provided_system_message and tools %} {{- "\nHere is a list of functions in JSON format that you can invoke. Use exact format: [func_name1(param1=value1, param2=value2), func_name2(...)]\n" }} {{- tool_definition -}} {%- elif tool_definition %} {{- tool_definition -}} {%- endif %} {{- "<|eot|>" }} {%- endif %} {#- Now deal with all other messages #} {%- for message in messages %} {#- Base case: messages that are not from tool role and has empty tool_call list #} {%- if not (message.role == 'ipython' or message.role == 'tool' or ('tool_calls' in message and message.tool_calls|length != 0 )) %} {{- '<|header_start|>' + message['role'] + '<|header_end|>\n\n' }} {%- if message['content'] is string %} {{- message['content'] }} {%- else %} {%- for content in message['content'] %} {%- if content['type'] == 'image' %} {{- '<|image|>' }} {%- elif content['type'] == 'text' %} {{- content['text'] | trim }} {%- endif %} {%- endfor %} {%- endif %} {{- "<|eot|>" }} {#- Tool case: messages has non-empty tool_call list, must from assistant #} {%- elif 'tool_calls' in message %} {#- assume tool_calls are always coming from assistant #} {%- if message.role == 'assistant' %} {{- '<|header_start|>assistant<|header_end|>\n\n' -}} {%- if message['content'] is string %} {{- message['content'] }} {%- else %} {%- for content in message['content'] %} {%- if content['type'] == 'image' %} {{- '<|image|>' }} {%- elif content['type'] == 'text' %} {{- content['text'] }} {%- endif %} {%- endfor %} {%- endif %} {{- "[" }} {%- for tool_call in message.tool_calls %} {%- if tool_call.function is defined %} {%- set tool_call = tool_call.function %} {%- endif %} {{- tool_call.name + '(' -}} {%- for param in tool_call.arguments %} {{- param + '="' -}} {{- "%s" | format(tool_call.arguments[param]) -}} {{- '"' -}} {% if not loop.last %}, {% endif %} {%- endfor %} {{- ')' -}} {% if not loop.last %}, {% endif %} {%- endfor %} {{- "]<|eot|>" }} {%- endif %} {#- Tool_response case: messages are from tool_response #} {%- elif message.role == "tool" or message.role == "ipython" %} {{- "<|header_start|>ipython<|header_end|>\n\n" }} {%- if message.content is string %} {{- message.content | tojson }} {%- else %} {%- for content in message['content'] %} {%- if content['type'] == 'text' %} {{- content['text'] | tojson }} {%- endif %} {%- endfor %} {%- endif %} {{- "<|eot|>" }} {%- endif %} {%- endfor %} {%- if add_generation_prompt %} {{- '<|header_start|>assistant<|header_end|>\n\n' }} {%- endif %} ================================================ FILE: examples/chat_template/vision_template_sarashina_vl.jinja ================================================ {# In sglang, the default chat templates often assume message['content'] is a plain string. That works fine for simple text conversations, but it ignores multimodal inputs (e.g. image_url, tool_call). To align with the original model behavior and support richer content, we iterate over message['content'] as a list of typed items and extract their values directly. This way, both text and non-text inputs are preserved in the prompt. Original template: https://huggingface.co/sbintuitions/sarashina2-vision-8b?chat_template=default #} {{ bos_token + '<|prefix|><|file|><|suffix|>A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human\'s questions.\n\n' }}{% for message in messages %}{% if message['role'] == 'user' %}{{ '### Human: ' }}{%- if message['content'] is string %}{{ message['content'] }}{%- else %}{% for item in message['content'] %}{% if item['type'] == 'text' %}{{ item['text'] }}{% endif %}{% endfor %}{% endif %}{{ '\n' }}{% elif message['role'] == 'assistant' %}{{ '### Assistant: ' }}{%- if message['content'] is string %}{{ message['content'] }}{%- else %}{% for item in message['content'] %}{% if item['type'] == 'text' %}{{ item['text'] }}{% endif %}{% endfor %}{% endif %}{{ '\n' }}{% endif %}{% endfor %}{% if messages[-1]['role'] == 'user' %}{{ '### Assistant:' }}{% endif %} ================================================ FILE: examples/checkpoint_engine/update.py ================================================ """ Usage: 1) Launch the server with wait-for-initial-weights option in one terminal: python -m sglang.launch_server --model-path /workspace/Qwen/Qwen3-4B/ --tensor-parallel-size 2 --port 19730 --load-format dummy --checkpoint-engine-wait-weights-before-ready --mem-fraction-static 0.7 2) Torchrun this script in another terminal: torchrun --nproc-per-node 2 update.py --update-method broadcast --checkpoint-path /workspace/Qwen/Qwen3-4B/ --inference-parallel-size 2 """ import argparse import json import os import pickle import time from collections import defaultdict from collections.abc import Callable from contextlib import contextmanager from typing import Literal import httpx import torch import torch.distributed as dist from checkpoint_engine.ps import ParameterServer from loguru import logger from safetensors import safe_open @contextmanager def timer(msg: str): start = time.perf_counter() yield end = time.perf_counter() logger.info(f"{msg} duration: {end - start:.2f} seconds") def check_sglang_ready( endpoint: str, inference_parallel_size: int, uds: str | None = None ): if rank != rank // inference_parallel_size * inference_parallel_size: return retry_num = 0 transport = None if uds is not None: transport = httpx.HTTPTransport(uds=uds) with httpx.Client(transport=transport) as client: while True: try: response = client.get(f"{endpoint}/ping", timeout=10) response.raise_for_status() break except (httpx.ConnectError, httpx.HTTPStatusError) as e: if retry_num % 10 == 0: logger.warning( f"fail to check sglang ready, retry {retry_num} times, error: {e}" ) retry_num += 1 time.sleep(0.1) def split_checkpoint_files( checkpoint_path: str, rank: int, world_size: int ) -> list[str]: checkpoint_files = [ os.path.join(checkpoint_path, f) for f in filter( lambda x: x.endswith(".safetensors"), os.listdir(checkpoint_path) ) ] files_per_rank = (len(checkpoint_files) + world_size - 1) // world_size return checkpoint_files[rank * files_per_rank : (rank + 1) * files_per_rank] def split_tensors( checkpoint_path: str, rank: int, world_size: int ) -> dict[str, torch.Tensor]: index_fn = os.path.join(checkpoint_path, "model.safetensors.index.json") with open(index_fn) as f: weight_map: dict[str, str] = json.load(f)["weight_map"] weights_per_rank = (len(weight_map) + world_size - 1) // world_size fn_tensors: dict[str, list[str]] = defaultdict(list) weight_keys = list(weight_map.items()) for name, file in weight_keys[ rank * weights_per_rank : (rank + 1) * weights_per_rank ]: fn_tensors[file].append(name) named_tensors = {} for file, names in fn_tensors.items(): with safe_open(os.path.join(checkpoint_path, file), framework="pt") as f: for name in names: named_tensors[name] = f.get_tensor(name) return named_tensors def req_inference( endpoint: str, inference_parallel_size: int, timeout: float = 300.0, uds: str | None = None, weight_version: str | None = None, ) -> Callable[[list[tuple[str, str]]], None]: rank = int(os.getenv("RANK", 0)) src = rank // inference_parallel_size * inference_parallel_size def req_func(socket_paths: list[tuple[str, str]]): if rank == src: with httpx.Client(transport=httpx.HTTPTransport(uds=uds)) as client: resp = client.post( f"{endpoint}/update_weights_from_ipc", json={ "zmq_handles": dict( socket_paths[src : src + inference_parallel_size] ), "flush_cache": True, "weight_version": weight_version, }, timeout=timeout, ) resp.raise_for_status() return req_func def update_weights( ps: ParameterServer, checkpoint_name: str, checkpoint_files: list[str], named_tensors: dict[str, torch.Tensor], req_func: Callable[[list[tuple[str, str]]], None], inference_parallel_size: int, endpoint: str, save_metas_file: str | None = None, update_method: Literal["broadcast", "p2p", "all"] = "broadcast", uds: str | None = None, ): ps.register_checkpoint( checkpoint_name, files=checkpoint_files, named_tensors=named_tensors ) ps.init_process_group() check_sglang_ready(endpoint, inference_parallel_size, uds) dist.barrier() with timer("Gather metas"): ps.gather_metas(checkpoint_name) if save_metas_file and int(os.getenv("RANK")) == 0: with open(save_metas_file, "wb") as f: pickle.dump(ps.get_metas(), f) if update_method == "broadcast" or update_method == "all": with timer("Update weights without setting ranks"): ps.update(checkpoint_name, req_func) if update_method == "p2p" or update_method == "all": if update_method: # sleep 2s to wait destroy process group time.sleep(2) with timer("Update weights with setting ranks"): ps.update( checkpoint_name, req_func, ranks=list(range(inference_parallel_size)) ) def join( ps: ParameterServer, checkpoint_name: str, load_metas_file: str, req_func: Callable[[list[tuple[str, str]]], None], inference_parallel_size: int, endpoint: str, uds: str | None = None, ): assert load_metas_file, "load_metas_file is required" with open(load_metas_file, "rb") as f: metas = pickle.load(f) ps.init_process_group() check_sglang_ready(endpoint, inference_parallel_size, uds) dist.barrier() with timer("Gather metas before join"): ps.gather_metas(checkpoint_name) ps.load_metas(metas) with timer( f"Update weights with setting ranks as range(0, {inference_parallel_size}) by using p2p" ): ps.update(checkpoint_name, req_func, ranks=list(range(inference_parallel_size))) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Update weights example") parser.add_argument("--checkpoint-path", type=str, default=None) parser.add_argument("--save-metas-file", type=str, default=None) parser.add_argument("--load-metas-file", type=str, default=None) parser.add_argument("--sleep-time", type=int, default=0) parser.add_argument("--endpoint", type=str, default="http://localhost:19730") parser.add_argument("--inference-parallel-size", type=int, default=8) parser.add_argument("--checkpoint-name", type=str, default="my-checkpoint-iter-0") parser.add_argument("--update-method", type=str, default="broadcast") parser.add_argument("--uds", type=str, default=None) parser.add_argument("--weight-version", type=str, default=None) args = parser.parse_args() rank = int(os.getenv("RANK")) world_size = int(os.getenv("WORLD_SIZE")) req_func = req_inference( args.endpoint, args.inference_parallel_size, uds=args.uds, weight_version=args.weight_version, ) ps = ParameterServer(auto_pg=True) ps._p2p_store = None if args.load_metas_file: join( ps, args.checkpoint_name, args.load_metas_file, req_func, args.inference_parallel_size, args.endpoint, args.uds, ) else: if os.path.exists( os.path.join(args.checkpoint_path, "model.safetensors.index.json") ): named_tensors = split_tensors(args.checkpoint_path, rank, world_size) checkpoint_files = [] else: checkpoint_files = split_checkpoint_files( args.checkpoint_path, rank, world_size ) named_tensors = {} update_weights( ps, args.checkpoint_name, checkpoint_files, named_tensors, req_func, args.inference_parallel_size, args.endpoint, args.save_metas_file, args.update_method, args.uds, ) time.sleep(args.sleep_time) ================================================ FILE: examples/frontend_language/quick_start/anthropic_example_chat.py ================================================ """ Usage: export ANTHROPIC_API_KEY=sk-****** python3 anthropic_example_chat.py """ import sglang as sgl @sgl.function def multi_turn_question(s, question_1, question_2): s += sgl.user(question_1) s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) s += sgl.user(question_2) s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) def single(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", ) for m in state.messages(): print(m["role"], ":", m["content"]) print("\n-- answer_1 --\n", state["answer_1"]) def stream(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", stream=True, ) for out in state.text_iter(): print(out, end="", flush=True) print() def batch(): states = multi_turn_question.run_batch( [ { "question_1": "What is the capital of the United States?", "question_2": "List two local attractions.", }, { "question_1": "What is the capital of France?", "question_2": "What is the population of this city?", }, ] ) for s in states: print(s.messages()) if __name__ == "__main__": sgl.set_default_backend(sgl.Anthropic("claude-3-haiku-20240307")) # Run a single request print("\n========== single ==========\n") single() # Stream output print("\n========== stream ==========\n") stream() # Run a batch of requests print("\n========== batch ==========\n") batch() ================================================ FILE: examples/frontend_language/quick_start/anthropic_example_complete.py ================================================ """ Usage: export ANTHROPIC_API_KEY=sk-****** python3 anthropic_example_complete.py """ import sglang as sgl @sgl.function def few_shot_qa(s, question): s += """ \n\nHuman: What is the capital of France? \n\nAssistant: Paris \n\nHuman: What is the capital of Germany? \n\nAssistant: Berlin \n\nHuman: What is the capital of Italy? \n\nAssistant: Rome """ s += "\n\nHuman: " + question + "\n" s += "\n\nAssistant:" + sgl.gen("answer", temperature=0) def single(): state = few_shot_qa.run(question="What is the capital of the United States?") answer = state["answer"].strip().lower() assert "washington" in answer, f"answer: {state['answer']}" print(state.text()) def stream(): state = few_shot_qa.run( question="What is the capital of the United States?", stream=True ) for out in state.text_iter("answer"): print(out, end="", flush=True) print() def batch(): states = few_shot_qa.run_batch( [ {"question": "What is the capital of the United States?"}, {"question": "What is the capital of China?"}, ] ) for s in states: print(s["answer"]) if __name__ == "__main__": sgl.set_default_backend(sgl.Anthropic("claude-3-haiku-20240307")) # Run a single request print("\n========== single ==========\n") single() # Stream output print("\n========== stream ==========\n") stream() # Run a batch of requests print("\n========== batch ==========\n") batch() ================================================ FILE: examples/frontend_language/quick_start/azure_openai_example_chat.py ================================================ """ Usage: export AZURE_OPENAI_API_KEY=sk-****** python3 openai_example_chat.py """ import os import sglang as sgl @sgl.function def multi_turn_question(s, question_1, question_2): s += sgl.system("You are a helpful assistant.") s += sgl.user(question_1) s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) s += sgl.user(question_2) s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) def single(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", ) for m in state.messages(): print(m["role"], ":", m["content"]) print("\n-- answer_1 --\n", state["answer_1"]) def stream(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", stream=True, ) for out in state.text_iter(): print(out, end="", flush=True) print() def batch(): states = multi_turn_question.run_batch( [ { "question_1": "What is the capital of the United States?", "question_2": "List two local attractions.", }, { "question_1": "What is the capital of France?", "question_2": "What is the population of this city?", }, ] ) for s in states: print(s.messages()) if __name__ == "__main__": backend = sgl.OpenAI( model_name="azure-gpt-4", api_version="2023-07-01-preview", azure_endpoint="https://oai-arena-sweden.openai.azure.com/", api_key=os.environ["AZURE_OPENAI_API_KEY"], is_azure=True, ) sgl.set_default_backend(backend) # Run a single request print("\n========== single ==========\n") single() # Stream output print("\n========== stream ==========\n") stream() # Run a batch of requests print("\n========== batch ==========\n") batch() ================================================ FILE: examples/frontend_language/quick_start/gemini_example_chat.py ================================================ """ Usage: export GCP_PROJECT_ID=****** python3 gemini_example_chat.py """ import sglang as sgl @sgl.function def multi_turn_question(s, question_1, question_2): s += sgl.user(question_1) s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) s += sgl.user(question_2) s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) def single(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", ) for m in state.messages(): print(m["role"], ":", m["content"]) print("\n-- answer_1 --\n", state["answer_1"]) def stream(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", stream=True, ) for out in state.text_iter(): print(out, end="", flush=True) print() def batch(): states = multi_turn_question.run_batch( [ { "question_1": "What is the capital of the United States?", "question_2": "List two local attractions.", }, { "question_1": "What is the capital of France?", "question_2": "What is the population of this city?", }, ] ) for s in states: print(s.messages()) if __name__ == "__main__": sgl.set_default_backend(sgl.VertexAI("gemini-pro")) # Run a single request print("\n========== single ==========\n") single() # Stream output print("\n========== stream ==========\n") stream() # Run a batch of requests print("\n========== batch ==========\n") batch() ================================================ FILE: examples/frontend_language/quick_start/gemini_example_complete.py ================================================ """ Usage: export GCP_PROJECT_ID=****** python3 gemini_example_complete.py """ import sglang as sgl @sgl.function def few_shot_qa(s, question): s += """The following are questions with answers. Q: What is the capital of France? A: Paris Q: What is the capital of Germany? A: Berlin Q: What is the capital of Italy? A: Rome """ s += "Q: " + question + "\n" s += "A:" + sgl.gen("answer", stop="\n", temperature=0) def single(): state = few_shot_qa.run(question="What is the capital of the United States?") answer = state["answer"].strip().lower() assert "washington" in answer, f"answer: {state['answer']}" print(state.text()) def stream(): state = few_shot_qa.run( question="What is the capital of the United States?", stream=True ) for out in state.text_iter("answer"): print(out, end="", flush=True) print() def batch(): states = few_shot_qa.run_batch( [ {"question": "What is the capital of the United States?"}, {"question": "What is the capital of China?"}, ] ) for s in states: print(s["answer"]) if __name__ == "__main__": sgl.set_default_backend(sgl.VertexAI("gemini-pro")) # Run a single request print("\n========== single ==========\n") single() # Stream output print("\n========== stream ==========\n") stream() # Run a batch of requests print("\n========== batch ==========\n") batch() ================================================ FILE: examples/frontend_language/quick_start/gemini_example_multimodal_chat.py ================================================ """ Usage: export GCP_PROJECT_ID=****** python3 gemini_example_multimodal_chat.py """ import sglang as sgl @sgl.function def image_qa(s, image_file1, image_file2, question): s += sgl.user(sgl.image(image_file1) + sgl.image(image_file2) + question) s += sgl.assistant(sgl.gen("answer", max_tokens=256)) if __name__ == "__main__": sgl.set_default_backend(sgl.VertexAI("gemini-pro-vision")) state = image_qa.run( image_file1="./images/cat.jpeg", image_file2="./images/dog.jpeg", question="Describe difference of the two images in one sentence.", stream=True, ) for out in state.text_iter("answer"): print(out, end="", flush=True) print() print(state["answer"]) ================================================ FILE: examples/frontend_language/quick_start/local_example_chat.py ================================================ """ Usage: python3 local_example_chat.py """ import sglang as sgl @sgl.function def multi_turn_question(s, question_1, question_2): s += sgl.user(question_1) s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) s += sgl.user(question_2) s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) def single(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", ) for m in state.messages(): print(m["role"], ":", m["content"]) print("\n-- answer_1 --\n", state["answer_1"]) def stream(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", stream=True, ) for out in state.text_iter(): print(out, end="", flush=True) print() def batch(): states = multi_turn_question.run_batch( [ { "question_1": "What is the capital of the United States?", "question_2": "List two local attractions.", }, { "question_1": "What is the capital of France?", "question_2": "What is the population of this city?", }, ] ) for s in states: print(s.messages()) if __name__ == "__main__": runtime = sgl.Runtime(model_path="meta-llama/Llama-2-7b-chat-hf") sgl.set_default_backend(runtime) # Run a single request print("\n========== single ==========\n") single() # Stream output print("\n========== stream ==========\n") stream() # Run a batch of requests print("\n========== batch ==========\n") batch() runtime.shutdown() ================================================ FILE: examples/frontend_language/quick_start/local_example_complete.py ================================================ """ Usage: python3 local_example_complete.py """ import sglang as sgl @sgl.function def few_shot_qa(s, question): s += """The following are questions with answers. Q: What is the capital of France? A: Paris Q: What is the capital of Germany? A: Berlin Q: What is the capital of Italy? A: Rome """ s += "Q: " + question + "\n" s += "A:" + sgl.gen("answer", stop="\n", temperature=0) def single(): state = few_shot_qa.run(question="What is the capital of the United States?") answer = state["answer"].strip().lower() assert "washington" in answer, f"answer: {state['answer']}" print(state.text()) def stream(): state = few_shot_qa.run( question="What is the capital of the United States?", stream=True ) for out in state.text_iter("answer"): print(out, end="", flush=True) print() def batch(): states = few_shot_qa.run_batch( [ {"question": "What is the capital of the United States?"}, {"question": "What is the capital of China?"}, ] ) for s in states: print(s["answer"]) if __name__ == "__main__": runtime = sgl.Runtime(model_path="meta-llama/Llama-2-7b-chat-hf") sgl.set_default_backend(runtime) # Run a single request print("\n========== single ==========\n") single() # Stream output print("\n========== stream ==========\n") stream() # Run a batch of requests print("\n========== batch ==========\n") batch() runtime.shutdown() ================================================ FILE: examples/frontend_language/quick_start/local_example_llava_next.py ================================================ """ Usage: python3 local_example_llava_next.py """ import sglang as sgl from sglang.lang.chat_template import get_chat_template @sgl.function def image_qa(s, image_path, question): s += sgl.user(sgl.image(image_path) + question) s += sgl.assistant(sgl.gen("answer")) def single(): state = image_qa.run( image_path="images/cat.jpeg", question="What is this?", max_new_tokens=128 ) print(state["answer"], "\n") def stream(): state = image_qa.run( image_path="images/cat.jpeg", question="What is this?", max_new_tokens=64, stream=True, ) for out in state.text_iter("answer"): print(out, end="", flush=True) print() def batch(): states = image_qa.run_batch( [ {"image_path": "images/cat.jpeg", "question": "What is this?"}, {"image_path": "images/dog.jpeg", "question": "What is this?"}, ], max_new_tokens=128, ) for s in states: print(s["answer"], "\n") if __name__ == "__main__": import multiprocessing as mp mp.set_start_method("spawn", force=True) runtime = sgl.Runtime(model_path="lmms-lab/llama3-llava-next-8b") runtime.endpoint.chat_template = get_chat_template("llama-3-instruct-llava") # Or you can use the 72B model # runtime = sgl.Runtime(model_path="lmms-lab/llava-next-72b", tp_size=8) # runtime.endpoint.chat_template = get_chat_template("chatml-llava") sgl.set_default_backend(runtime) print(f"chat template: {runtime.endpoint.chat_template.name}") # Or you can use API models # sgl.set_default_backend(sgl.OpenAI("gpt-4-vision-preview")) # sgl.set_default_backend(sgl.VertexAI("gemini-pro-vision")) # Run a single request print("\n========== single ==========\n") single() # Stream output print("\n========== stream ==========\n") stream() # Run a batch of requests print("\n========== batch ==========\n") batch() runtime.shutdown() ================================================ FILE: examples/frontend_language/quick_start/openai_example_chat.py ================================================ """ Usage: export OPENAI_API_KEY=sk-****** python3 openai_example_chat.py """ import sglang as sgl @sgl.function def multi_turn_question(s, question_1, question_2): s += sgl.system("You are a helpful assistant.") s += sgl.user(question_1) s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) s += sgl.user(question_2) s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) def single(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", ) for m in state.messages(): print(m["role"], ":", m["content"]) print("\n-- answer_1 --\n", state["answer_1"]) def stream(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", stream=True, ) for out in state.text_iter(): print(out, end="", flush=True) print() def batch(): states = multi_turn_question.run_batch( [ { "question_1": "What is the capital of the United States?", "question_2": "List two local attractions.", }, { "question_1": "What is the capital of France?", "question_2": "What is the population of this city?", }, ] ) for s in states: print(s.messages()) if __name__ == "__main__": sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo")) # Run a single request print("\n========== single ==========\n") single() # Stream output print("\n========== stream ==========\n") stream() # Run a batch of requests print("\n========== batch ==========\n") batch() ================================================ FILE: examples/frontend_language/quick_start/openai_example_complete.py ================================================ """ Usage: export OPENAI_API_KEY=sk-****** python3 openai_example_complete.py """ import sglang as sgl @sgl.function def few_shot_qa(s, question): s += """The following are questions with answers. Q: What is the capital of France? A: Paris Q: What is the capital of Germany? A: Berlin Q: What is the capital of Italy? A: Rome """ s += "Q: " + question + "\n" s += "A:" + sgl.gen("answer", stop="\n", temperature=0) def single(): state = few_shot_qa.run(question="What is the capital of the United States?") answer = state["answer"].strip().lower() assert "washington" in answer, f"answer: {state['answer']}" print(state.text()) def stream(): state = few_shot_qa.run( question="What is the capital of the United States?", stream=True ) for out in state.text_iter("answer"): print(out, end="", flush=True) print() def batch(): states = few_shot_qa.run_batch( [ {"question": "What is the capital of the United States?"}, {"question": "What is the capital of China?"}, ] ) for s in states: print(s["answer"]) if __name__ == "__main__": sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo-instruct")) # Run a single request print("\n========== single ==========\n") single() # Stream output print("\n========== stream ==========\n") stream() # Run a batch of requests print("\n========== batch ==========\n") batch() ================================================ FILE: examples/frontend_language/quick_start/openai_example_n.py ================================================ """ Usage: export OPENAI_API_KEY=sk-****** python3 openai_example_chat.py """ import sglang as sgl @sgl.function def multi_turn_question(s, question_1, question_2): s += sgl.system("You are a helpful assistant.") s += sgl.user(question_1) s += sgl.assistant(sgl.gen("answer_1", max_tokens=1024, n=2)) s += sgl.user(question_2) s += sgl.assistant( sgl.gen( "answer_2", max_tokens=1024, ) ) def single(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", ) for m in state.messages(): print(m["role"], ":", m["content"]) print("\n-- answer_1 --\n", state["answer_1"]) print("\n-- answer_2 --\n", state["answer_2"]) assert isinstance(state["answer_1"], list) assert len(state["answer_1"]) == 2 assert isinstance(state["answer_2"], str) def batch(): states = multi_turn_question.run_batch( [ { "question_1": "What is the capital of the United States?", "question_2": "List two local attractions.", }, { "question_1": "What is the capital of France?", "question_2": "What is the population of this city?", }, ] ) for s in states: print(s.messages()) print("\n-- answer_1 --\n", s["answer_1"]) print("\n-- answer_2 --\n", s["answer_2"]) assert isinstance(s["answer_1"], list) assert len(s["answer_1"]) == 2 assert isinstance(s["answer_2"], str) if __name__ == "__main__": sgl.set_default_backend(sgl.OpenAI("o1")) # Run a single request print("\n========== single ==========\n") single() # Run a batch of requests print("\n========== batch ==========\n") batch() ================================================ FILE: examples/frontend_language/quick_start/openai_example_o1.py ================================================ """ Usage: export OPENAI_API_KEY=sk-****** python3 openai_example_chat.py """ import sglang as sgl @sgl.function def multi_turn_question(s, question_1, question_2): s += sgl.system("You are a helpful assistant.") s += sgl.user(question_1) s += sgl.assistant(sgl.gen("answer_1", max_tokens=100)) s += sgl.user(question_2) s += sgl.assistant(sgl.gen("answer_2")) def single(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", ) for m in state.messages(): print(m["role"], ":", m["content"]) print("\n-- answer_1 --\n", state["answer_1"]) def batch(): states = multi_turn_question.run_batch( [ { "question_1": "What is the capital of the United States?", "question_2": "List two local attractions.", }, { "question_1": "What is the capital of France?", "question_2": "What is the population of this city?", }, ] ) for s in states: print(s.messages()) if __name__ == "__main__": sgl.set_default_backend(sgl.OpenAI("o1")) # Run a single request print("\n========== single ==========\n") single() # Run a batch of requests print("\n========== batch ==========\n") batch() ================================================ FILE: examples/frontend_language/quick_start/openrouter_example_chat.py ================================================ """ Usage: export OPENROUTER_API_KEY=sk-****** python3 together_example_chat.py """ import os import sglang as sgl @sgl.function def multi_turn_question(s, question_1, question_2): s += sgl.system("You are a helpful assistant.") s += sgl.user(question_1) s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) s += sgl.user(question_2) s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) def single(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", ) for m in state.messages(): print(m["role"], ":", m["content"]) print("\n-- answer_1 --\n", state["answer_1"]) def stream(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", stream=True, ) for out in state.text_iter(): print(out, end="", flush=True) print() def batch(): states = multi_turn_question.run_batch( [ { "question_1": "What is the capital of the United States?", "question_2": "List two local attractions.", }, { "question_1": "What is the capital of France?", "question_2": "What is the population of this city?", }, ] ) for s in states: print(s.messages()) if __name__ == "__main__": backend = sgl.OpenAI( model_name="google/gemma-7b-it:free", base_url="https://openrouter.ai/api/v1", api_key=os.environ.get("OPENROUTER_API_KEY"), ) sgl.set_default_backend(backend) # Run a single request print("\n========== single ==========\n") single() # Stream output print("\n========== stream ==========\n") stream() # Run a batch of requests print("\n========== batch ==========\n") batch() ================================================ FILE: examples/frontend_language/quick_start/together_example_chat.py ================================================ """ Usage: export TOGETHER_API_KEY=sk-****** python3 together_example_chat.py """ import os import sglang as sgl @sgl.function def multi_turn_question(s, question_1, question_2): s += sgl.system("You are a helpful assistant.") s += sgl.user(question_1) s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) s += sgl.user(question_2) s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) def single(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", ) for m in state.messages(): print(m["role"], ":", m["content"]) print("\n-- answer_1 --\n", state["answer_1"]) def stream(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", stream=True, ) for out in state.text_iter(): print(out, end="", flush=True) print() def batch(): states = multi_turn_question.run_batch( [ { "question_1": "What is the capital of the United States?", "question_2": "List two local attractions.", }, { "question_1": "What is the capital of France?", "question_2": "What is the population of this city?", }, ] ) for s in states: print(s.messages()) if __name__ == "__main__": backend = sgl.OpenAI( model_name="mistralai/Mixtral-8x7B-Instruct-v0.1", base_url="https://api.together.xyz/v1", api_key=os.environ.get("TOGETHER_API_KEY"), ) sgl.set_default_backend(backend) # Run a single request print("\n========== single ==========\n") single() # Stream output print("\n========== stream ==========\n") stream() # Run a batch of requests print("\n========== batch ==========\n") batch() ================================================ FILE: examples/frontend_language/quick_start/together_example_complete.py ================================================ """ Usage: export TOGETHER_API_KEY=sk-****** python3 together_example_complete.py """ import os import sglang as sgl @sgl.function def few_shot_qa(s, question): s += """The following are questions with answers. Q: What is the capital of France? A: Paris Q: What is the capital of Germany? A: Berlin Q: What is the capital of Italy? A: Rome """ s += "Q: " + question + "\n" s += "A:" + sgl.gen("answer", stop="\n", temperature=0) def single(): state = few_shot_qa.run(question="What is the capital of the United States?") answer = state["answer"].strip().lower() assert "washington" in answer, f"answer: {state['answer']}" print(state.text()) def stream(): state = few_shot_qa.run( question="What is the capital of the United States?", stream=True ) for out in state.text_iter("answer"): print(out, end="", flush=True) print() def batch(): states = few_shot_qa.run_batch( [ {"question": "What is the capital of the United States?"}, {"question": "What is the capital of China?"}, ] ) for s in states: print(s["answer"]) if __name__ == "__main__": backend = sgl.OpenAI( model_name="mistralai/Mixtral-8x7B-Instruct-v0.1", is_chat_model=False, base_url="https://api.together.xyz/v1", api_key=os.environ.get("TOGETHER_API_KEY"), ) sgl.set_default_backend(backend) # Run a single request print("\n========== single ==========\n") single() # Stream output print("\n========== stream ==========\n") stream() # Run a batch of requests print("\n========== batch ==========\n") batch() ================================================ FILE: examples/frontend_language/usage/chinese_regex.py ================================================ import sglang as sgl character_regex = ( r"""\{\n""" + r""" "姓名": "[^"]{1,32}",\n""" + r""" "学院": "(格兰芬多|赫奇帕奇|拉文克劳|斯莱特林)",\n""" + r""" "血型": "(纯血|混血|麻瓜)",\n""" + r""" "职业": "(学生|教师|傲罗|魔法部|食死徒|凤凰社成员)",\n""" + r""" "魔杖": \{\n""" + r""" "材质": "[^"]{1,32}",\n""" + r""" "杖芯": "[^"]{1,32}",\n""" + r""" "长度": [0-9]{1,2}\.[0-9]{0,2}\n""" + r""" \},\n""" + r""" "存活": "(存活|死亡)",\n""" + r""" "守护神": "[^"]{1,32}",\n""" + r""" "博格特": "[^"]{1,32}"\n""" + r"""\}""" ) @sgl.function def character_gen(s, name): s += name + " 是一名哈利波特系列小说中的角色。请填写以下关于这个角色的信息。" s += """\ 这是一个例子 { "姓名": "哈利波特", "学院": "格兰芬多", "血型": "混血", "职业": "学生", "魔杖": { "材质": "冬青木", "杖芯": "凤凰尾羽", "长度": 11.0 }, "存活": "存活", "守护神": "麋鹿", "博格特": "摄魂怪" } """ s += f"现在请你填写{name}的信息：\n" s += sgl.gen("json_output", max_tokens=256, regex=character_regex) def main(): backend = sgl.RuntimeEndpoint("http://localhost:30000") sgl.set_default_backend(backend) ret = character_gen.run(name="赫敏格兰杰", temperature=0) print(ret.text()) if __name__ == "__main__": main() ================================================ FILE: examples/frontend_language/usage/choices_logprob.py ================================================ """ Usage: python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 python choices_logprob.py """ import sglang as sgl @sgl.function def tool_use(s, question): s += "To answer this question: " + question + ", " s += "I need to use a " + sgl.gen("tool", choices=["calculator", "search engine"]) def main(): # Run one case question = "What is 5 + 5?" state = tool_use.run(question) print("questions:", question) print("choice:", state["tool"]) meta_info = state.get_meta_info("tool") print("logprobs of choice 1", meta_info["input_token_logprobs"][0]) print("logprobs of choice 2", meta_info["input_token_logprobs"][1]) print("-" * 50) # Run a batch questions = [ "What is 5 + 6?", "Who is Michael Jordan?", ] states = tool_use.run_batch([{"question": q} for q in questions]) for question, state in zip(questions, states): print("questions:", question) print("choice:", state["tool"]) meta_info = state.get_meta_info("tool") print("logprobs of choice 1", meta_info["input_token_logprobs"][0]) print("logprobs of choice 2", meta_info["input_token_logprobs"][1]) print("-" * 50) if __name__ == "__main__": sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) main() ================================================ FILE: examples/frontend_language/usage/cot_decoding.py ================================================ from math import exp from pprint import pformat import sglang as sgl YELLOW = "\033[1;33m" GREEN = "\033[1;32m" BLUE = "\033[1;34m" CLEAR = "\033[1;0m" @sgl.function def cot_decoding(s, question, get_top_k, is_chat_model, verbose): """CoT Decoding: http://arxiv.org/abs/2402.10200""" if is_chat_model: s += sgl.user("Question: " + question + "\nAnswer:") s += sgl.assistant_begin() else: s += "Question: " + question + "\nAnswer:" step_0 = s.fork(1)[0] forks = s.fork(get_top_k) answer_forks = s.fork(get_top_k) # decoding step 0 step_0 += sgl.gen( "get_top_k", max_tokens=0, return_logprob=True, top_logprobs_num=get_top_k, return_text_in_logprobs=True, ) logprobs = step_0.get_meta_info("get_top_k")["output_top_logprobs"][0] print("Decoding step 0:", ", ".join(pformat(token[2]) for token in logprobs)) for idx, (f, token) in enumerate(zip(forks, logprobs)): logprob, token_id, text = token f += text if text == "<|end_of_text|>": print( f"{YELLOW}Path #{idx} {pformat(text)}[{exp(logprob):.3f}] (score=nan, answer=nan){CLEAR}" ) continue # continue greedy decoding f += sgl.gen( "answer", temperature=0, max_tokens=1024, return_logprob=True, top_logprobs_num=2, return_text_in_logprobs=True, ) # calculate probability disparity between the top and secondary tokens x1s = [exp(xt[0][0]) for xt in f.get_meta_info("answer")["output_top_logprobs"]] x2s = [exp(xt[1][0]) for xt in f.get_meta_info("answer")["output_top_logprobs"]] tokens = [xt[0][2] for xt in f.get_meta_info("answer")["output_top_logprobs"]] delta = (sum(x1s) - sum(x2s)) / len(x1s) # extract the answer span (without the '<|end_of_text|>' token) answer_forks[idx] += text + f["answer"] + "\nSo the answer is" answer_forks[idx] += sgl.gen( "answer_span", temperature=0, max_tokens=64, return_logprob=True, top_logprobs_num=2, return_text_in_logprobs=True, ) answer = answer_forks[idx]["answer_span"].replace("\n", " ").strip(":") print( f"{YELLOW}Path #{idx} {pformat(text)}[{exp(logprob):.3f}] (score={delta}, answer={answer}){CLEAR}" ) generated_text = str(answer_forks[idx])[len("ProgramState(") : -1] print(f"{BLUE}{pformat(generated_text)}{CLEAR}") if verbose: answer_tokens = [ xt[0][2] for xt in answer_forks[idx].get_meta_info("answer_span")[ "output_top_logprobs" ] ] answer_x1s = [ exp(xt[0][0]) for xt in answer_forks[idx].get_meta_info("answer_span")[ "output_top_logprobs" ] ] answer_x2s = [ exp(xt[1][0]) for xt in answer_forks[idx].get_meta_info("answer_span")[ "output_top_logprobs" ] ] for token, x1, x2 in zip(tokens, x1s, x2s): print(f" {GREEN}{pformat(token)}{CLEAR}({x1:.3f}-{x2:.3f})", end="") print("\n===========") for token, x1, x2 in zip(answer_tokens, answer_x1s, answer_x2s): print(f" {GREEN}{pformat(token)}{CLEAR}({x1:.3f}-{x2:.3f})", end="") print() sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) state = cot_decoding.run( question=r"Claire makes a 3 egg omelet every morning for breakfast. How many dozens of eggs will she eat in 4 weeks?", get_top_k=10, is_chat_model=True, verbose=False, ) ================================================ FILE: examples/frontend_language/usage/json_decode.py ================================================ """ Usage: python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 python json_decode.py """ from enum import Enum from pydantic import BaseModel import sglang as sgl from sglang.srt.constrained.outlines_backend import build_regex_from_object character_regex = ( r"""\{\n""" + r""" "name": "[\w\d\s]{1,16}",\n""" + r""" "house": "(Gryffindor|Slytherin|Ravenclaw|Hufflepuff)",\n""" + r""" "blood status": "(Pure-blood|Half-blood|Muggle-born)",\n""" + r""" "occupation": "(student|teacher|auror|ministry of magic|death eater|order of the phoenix)",\n""" + r""" "wand": \{\n""" + r""" "wood": "[\w\d\s]{1,16}",\n""" + r""" "core": "[\w\d\s]{1,16}",\n""" + r""" "length": [0-9]{1,2}\.[0-9]{0,2}\n""" + r""" \},\n""" + r""" "alive": "(Alive|Deceased)",\n""" + r""" "patronus": "[\w\d\s]{1,16}",\n""" + r""" "bogart": "[\w\d\s]{1,16}"\n""" + r"""\}""" ) @sgl.function def character_gen(s, name): s += ( name + " is a character in Harry Potter. Please fill in the following information about this character.\n" ) s += "The constrained regex is:\n" s += character_regex + "\n" s += "The JSON output is:\n" s += sgl.gen("json_output", max_tokens=256, regex=character_regex) def driver_character_gen(): state = character_gen.run(name="Hermione Granger") print(state.text()) class Weapon(str, Enum): sword = "sword" axe = "axe" mace = "mace" spear = "spear" bow = "bow" crossbow = "crossbow" class Wizard(BaseModel): name: str age: int weapon: Weapon @sgl.function def pydantic_wizard_gen(s): s += "Give me a description about a wizard in the JSON format.\n" s += sgl.gen( "character", max_tokens=128, temperature=0, regex=build_regex_from_object(Wizard), # Requires pydantic >= 2.0 ) def driver_pydantic_wizard_gen(): state = pydantic_wizard_gen.run() print(state.text()) if __name__ == "__main__": sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) driver_character_gen() # driver_pydantic_wizard_gen() ================================================ FILE: examples/frontend_language/usage/json_logprobs.py ================================================ # NOTE: Currently this can only be run through HTTP requests. from concurrent.futures import ThreadPoolExecutor from json_decode import character_regex from sglang.utils import http_request character_names = ["Hermione Granger", "Ron Weasley", "Harry Potter"] base_url = "http://localhost:30000" prompt = "is a character in Harry Potter. Please fill in the following information about this character.\n" def openai_api_request(name): data = { "model": "", "prompt": name + prompt, "temperature": 0, "max_tokens": 128, "regex": character_regex, "logprobs": 3, } res = http_request(base_url + "/v1/completions", json=data).json() # with open(f"json_logprobs_{name.replace(' ', '_')}_tmp.json", "w") as fout: # fout.write(json.dumps(res, indent=4)) logprobs = res["choices"][0]["logprobs"] usage = res["usage"] assert len(logprobs["token_logprobs"]) == len(logprobs["tokens"]) assert len(logprobs["token_logprobs"]) == len(logprobs["top_logprobs"]) assert len(logprobs["token_logprobs"]) == usage["completion_tokens"] - 1 return res def srt_api_request(name): data = { "text": name + prompt, "sampling_params": { "temperature": 0, "max_new_tokens": 128, "regex": character_regex, }, "return_logprob": True, "logprob_start_len": 0, "top_logprobs_num": 3, "return_text_in_logprobs": True, } res = http_request(base_url + "/generate", json=data).json() # with open(f"json_logprobs_{name.replace(' ', '_')}_tmp.json", "w") as fout: # fout.write(json.dumps(res, indent=4)) meta_info = res["meta_info"] assert len(meta_info["input_token_logprobs"]) == len( meta_info["input_top_logprobs"] ) assert len(meta_info["output_token_logprobs"]) == len( meta_info["output_top_logprobs"] ) assert len(meta_info["input_token_logprobs"]) == meta_info["prompt_tokens"] assert len(meta_info["output_token_logprobs"]) == meta_info["completion_tokens"] - 1 return res def pretty_print(res): meta_info = res["meta_info"] print("\n\n", "=" * 30, "Prefill", "=" * 30) for i in range(len(meta_info["input_token_logprobs"])): print(f"{str(meta_info['input_token_logprobs'][i][2].encode()): <20}", end="") top_ks = ( [str(t[2].encode()) for t in meta_info["input_top_logprobs"][i]] if meta_info["input_top_logprobs"][i] else [] ) for top_k in top_ks: print(f"{top_k: <15}", end="") print() print("\n\n", "=" * 30, "Decode", "=" * 30) for i in range(len(meta_info["output_token_logprobs"])): print(f"{str(meta_info['output_token_logprobs'][i][2].encode()): <20}", end="") top_ks = [str(t[2].encode()) for t in meta_info["output_top_logprobs"][i]] for top_k in top_ks: print(f"{top_k: <15}", end="") print() print(res["text"]) if __name__ == "__main__": with ThreadPoolExecutor() as executor: ress = executor.map(srt_api_request, character_names) for res in ress: pretty_print(res) openai_api_request("Hermione Granger") ================================================ FILE: examples/frontend_language/usage/llava_video/srt_example_llava_v.py ================================================ """ Usage: pip install opencv-python-headless python3 srt_example_llava_v.py """ import argparse import csv import json import os import time import requests import sglang as sgl @sgl.function def video_qa(s, num_frames, video_path, question): s += sgl.user(sgl.video(video_path, num_frames) + question) s += sgl.assistant(sgl.gen("answer")) def single(path, num_frames=16): state = video_qa.run( num_frames=num_frames, video_path=path, question="Please provide a detailed description of the video, focusing on the main subjects, their actions, the background scenes", temperature=0.0, max_new_tokens=1024, ) print(state["answer"], "\n") def split_into_chunks(lst, num_chunks): """Split a list into a specified number of chunks.""" # Calculate the chunk size using integer division. Note that this may drop some items if not evenly divisible. chunk_size = len(lst) // num_chunks if chunk_size == 0: chunk_size = len(lst) # Use list comprehension to generate chunks. The last chunk will take any remainder if the list size isn't evenly divisible. chunks = [lst[i : i + chunk_size] for i in range(0, len(lst), chunk_size)] # Ensure we have exactly num_chunks chunks, even if some are empty chunks.extend([[] for _ in range(num_chunks - len(chunks))]) return chunks def save_batch_results(batch_video_files, states, cur_chunk, batch_idx, save_dir): csv_filename = f"{save_dir}/chunk_{cur_chunk}_batch_{batch_idx}.csv" with open(csv_filename, "w", newline="") as csvfile: writer = csv.writer(csvfile) writer.writerow(["video_name", "answer"]) for video_path, state in zip(batch_video_files, states): video_name = os.path.basename(video_path) writer.writerow([video_name, state["answer"]]) def compile_and_cleanup_final_results(cur_chunk, num_batches, save_dir): final_csv_filename = f"{save_dir}/final_results_chunk_{cur_chunk}.csv" with open(final_csv_filename, "w", newline="") as final_csvfile: writer = csv.writer(final_csvfile) writer.writerow(["video_name", "answer"]) for batch_idx in range(num_batches): batch_csv_filename = f"{save_dir}/chunk_{cur_chunk}_batch_{batch_idx}.csv" with open(batch_csv_filename, "r") as batch_csvfile: reader = csv.reader(batch_csvfile) next(reader) # Skip header row for row in reader: writer.writerow(row) os.remove(batch_csv_filename) def find_video_files(video_dir): # Check if the video_dir is actually a file if os.path.isfile(video_dir): # If it's a file, return it as a single-element list return [video_dir] # Original logic to find video files in a directory video_files = [] for root, dirs, files in os.walk(video_dir): for file in files: if file.endswith((".mp4", ".avi", ".mov")): video_files.append(os.path.join(root, file)) return video_files def batch(video_dir, save_dir, cur_chunk, num_chunks, num_frames=16, batch_size=64): video_files = find_video_files(video_dir) chunked_video_files = split_into_chunks(video_files, num_chunks)[cur_chunk] num_batches = 0 for i in range(0, len(chunked_video_files), batch_size): batch_video_files = chunked_video_files[i : i + batch_size] print(f"Processing batch of {len(batch_video_files)} video(s)...") if not batch_video_files: print("No video files found in the specified directory.") return batch_input = [ { "num_frames": num_frames, "video_path": video_path, "question": "Please provide a detailed description of the video, focusing on the main subjects, their actions, the background scenes.", } for video_path in batch_video_files ] start_time = time.perf_counter() states = video_qa.run_batch(batch_input, max_new_tokens=512, temperature=0.2) total_time = time.perf_counter() - start_time average_time = total_time / len(batch_video_files) print( f"Number of videos in batch: {len(batch_video_files)}. Average processing time per video: {average_time:.2f} seconds. Total time for this batch: {total_time:.2f} seconds" ) save_batch_results(batch_video_files, states, cur_chunk, num_batches, save_dir) num_batches += 1 compile_and_cleanup_final_results(cur_chunk, num_batches, save_dir) if __name__ == "__main__": url = "https://raw.githubusercontent.com/EvolvingLMMs-Lab/sglang/dev/onevision_local/assets/jobs.mp4" cache_dir = os.path.expanduser("~/.cache") file_path = os.path.join(cache_dir, "jobs.mp4") os.makedirs(cache_dir, exist_ok=True) response = requests.get(url) response.raise_for_status() # Raise an exception for bad responses with open(file_path, "wb") as f: f.write(response.content) print(f"File downloaded and saved to: {file_path}") # Create the parser parser = argparse.ArgumentParser( description="Run video processing with specified port." ) # Add an argument for the port parser.add_argument( "--port", type=int, default=30000, help="The master port for distributed serving.", ) parser.add_argument( "--chunk-idx", type=int, default=0, help="The index of the chunk to process." ) parser.add_argument( "--num-chunks", type=int, default=8, help="The number of chunks to process." ) parser.add_argument( "--save-dir", type=str, default="./work_dirs/llava_video", help="The directory to save the processed video files.", ) parser.add_argument( "--video-dir", type=str, default=os.path.expanduser("~/.cache/jobs.mp4"), help="The directory or path for the processed video files.", ) parser.add_argument( "--model-path", type=str, default="lmms-lab/LLaVA-NeXT-Video-7B", help="The model path for the video processing.", ) parser.add_argument( "--num-frames", type=int, default=16, help="The number of frames to process in each video.", ) parser.add_argument("--mm_spatial_pool_stride", type=int, default=2) # Parse the arguments args = parser.parse_args() cur_port = args.port cur_chunk = args.chunk_idx num_chunks = args.num_chunks num_frames = args.num_frames if "34b" in args.model_path.lower(): tokenizer_path = "liuhaotian/llava-v1.6-34b-tokenizer" elif "7b" in args.model_path.lower(): tokenizer_path = "llava-hf/llava-1.5-7b-hf" else: print("Invalid model path. Please specify a valid model path.") exit() model_override_args = {} model_override_args["mm_spatial_pool_stride"] = args.mm_spatial_pool_stride model_override_args["architectures"] = ["LlavaVidForCausalLM"] model_override_args["num_frames"] = args.num_frames model_override_args["model_type"] = "llava" if "34b" in args.model_path.lower(): model_override_args["image_token_index"] = 64002 if args.num_frames == 32: model_override_args["rope_scaling"] = {"factor": 2.0, "rope_type": "linear"} model_override_args["max_sequence_length"] = 4096 * 2 model_override_args["tokenizer_model_max_length"] = 4096 * 2 elif args.num_frames < 32: pass else: print( "The maximum number of frames to process is 32. Please specify a valid number of frames." ) exit() runtime = sgl.Runtime( model_path=args.model_path, # "liuhaotian/llava-v1.6-vicuna-7b", tokenizer_path=tokenizer_path, port=cur_port, json_model_override_args=json.dumps(model_override_args), tp_size=1, ) sgl.set_default_backend(runtime) print(f"chat template: {runtime.endpoint.chat_template.name}") # Run a single request print("\n========== single ==========\n") root = args.video_dir if os.path.isfile(root): video_files = [root] else: video_files = [ os.path.join(root, f) for f in os.listdir(root) if f.endswith((".mp4", ".avi", ".mov")) ] # Add more extensions if needed start_time = time.perf_counter() # Start time for processing a single video for cur_video in video_files[:1]: print(cur_video) single(cur_video, num_frames) end_time = time.perf_counter() # End time for processing a single video total_time = end_time - start_time average_time = total_time / len( video_files ) # Calculate the average processing time print(f"Average processing time per video: {average_time:.2f} seconds") runtime.shutdown() # # Run a batch of requests # print("\n========== batch ==========\n") # if not os.path.exists(args.save_dir): # os.makedirs(args.save_dir) # batch(args.video_dir, args.save_dir, cur_chunk, num_chunks, num_frames, num_chunks) # runtime.shutdown() ================================================ FILE: examples/frontend_language/usage/llava_video/srt_example_llava_v.sh ================================================ #!/bin/bash ##### USAGE ##### # - First node: # ```sh # bash examples/usage/llava_video/srt_example_llava_v.sh K 0 YOUR_VIDEO_PATH YOUR_MODEL_PATH FRAMES_PER_VIDEO # ``` # - Second node: # ```sh # bash examples/usage/llava_video/srt_example_llava_v.sh K 1 YOUR_VIDEO_PATH YOUR_MODEL_PATH FRAMES_PER_VIDEO # ``` # - The K node: # ```sh # bash examples/usage/llava_video/srt_example_llava_v.sh K K-1 YOUR_VIDEO_PATH YOUR_MODEL_PATH FRAMES_PER_VIDEO # ``` # Replace `K`, `YOUR_VIDEO_PATH`, `YOUR_MODEL_PATH`, and `FRAMES_PER_VIDEO` with your specific details. # CURRENT_ROOT="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" CURRENT_ROOT=$(dirname "$0") echo ${CURRENT_ROOT} cd ${CURRENT_ROOT} export PYTHONWARNINGS=ignore START_TIME=$(date +%s) # Capture start time NUM_NODES=$1 CUR_NODES_IDX=$2 VIDEO_DIR=$3 MODEL_PATH=$4 NUM_FRAMES=$5 # FRAME_FORMAT=$6 # FRAME_FORMAT=$(echo $FRAME_FORMAT | tr '[:lower:]' '[:upper:]') # # Check if FRAME_FORMAT is either JPEG or PNG # if [[ "$FRAME_FORMAT" != "JPEG" && "$FRAME_FORMAT" != "PNG" ]]; then # echo "Error: FRAME_FORMAT must be either JPEG or PNG." # exit 1 # fi # export TARGET_FRAMES=$TARGET_FRAMES echo "Each video you will sample $NUM_FRAMES frames" # export FRAME_FORMAT=$FRAME_FORMAT # echo "The frame format is $FRAME_FORMAT" # Assuming GPULIST is a bash array containing your GPUs GPULIST=(0 1 2 3 4 5 6 7) LOCAL_CHUNKS=${#GPULIST[@]} echo "Number of GPUs in GPULIST: $LOCAL_CHUNKS" ALL_CHUNKS=$((NUM_NODES * LOCAL_CHUNKS)) # Calculate GPUs per chunk GPUS_PER_CHUNK=1 echo $GPUS_PER_CHUNK for IDX in $(seq 1 $LOCAL_CHUNKS); do ( START=$(((IDX-1) * GPUS_PER_CHUNK)) LENGTH=$GPUS_PER_CHUNK # Length for slicing, not the end index CHUNK_GPUS=(${GPULIST[@]:$START:$LENGTH}) # Convert the chunk GPUs array to a comma-separated string CHUNK_GPUS_STR=$(IFS=,; echo "${CHUNK_GPUS[*]}") LOCAL_IDX=$((CUR_NODES_IDX * LOCAL_CHUNKS + IDX)) echo "Chunk $(($LOCAL_IDX - 1)) will run on GPUs $CHUNK_GPUS_STR" # Calculate the port for this chunk. Ensure it's incremented by 5 for each chunk. PORT=$((10000 + RANDOM % 55536)) MAX_RETRIES=10 RETRY_COUNT=0 COMMAND_STATUS=1 # Initialize as failed while [ $RETRY_COUNT -lt $MAX_RETRIES ] && [ $COMMAND_STATUS -ne 0 ]; do echo "Running chunk $(($LOCAL_IDX - 1)) on GPUs $CHUNK_GPUS_STR with port $PORT. Attempt $(($RETRY_COUNT + 1))" #!/bin/bash CUDA_VISIBLE_DEVICES=$CHUNK_GPUS_STR python3 srt_example_llava_v.py \ --port $PORT \ --num-chunks $ALL_CHUNKS \ --chunk-idx $(($LOCAL_IDX - 1)) \ --save-dir work_dirs/llava_next_video_inference_results \ --video-dir $VIDEO_DIR \ --model-path $MODEL_PATH \ --num-frames $NUM_FRAMES #& wait $! # Wait for the process to finish and capture its exit status COMMAND_STATUS=$? if [ $COMMAND_STATUS -ne 0 ]; then echo "Execution failed for chunk $(($LOCAL_IDX - 1)), attempt $(($RETRY_COUNT + 1)). Retrying..." RETRY_COUNT=$(($RETRY_COUNT + 1)) sleep 180 # Wait a bit before retrying else echo "Execution succeeded for chunk $(($LOCAL_IDX - 1))." fi done if [ $COMMAND_STATUS -ne 0 ]; then echo "Execution failed for chunk $(($LOCAL_IDX - 1)) after $MAX_RETRIES attempts." fi ) #& sleep 2 # Slight delay to stagger the start times done wait cat work_dirs/llava_next_video_inference_results/final_results_chunk_*.csv > work_dirs/llava_next_video_inference_results/final_results_node_${CUR_NODES_IDX}.csv END_TIME=$(date +%s) # Capture end time ELAPSED_TIME=$(($END_TIME - $START_TIME)) echo "Total execution time: $ELAPSED_TIME seconds." ================================================ FILE: examples/frontend_language/usage/openai_chat_speculative.py ================================================ """ Usage: ***Note: for speculative execution to work, user must put all "gen" in "assistant". Show in "assistant" the desired answer format. Each "gen" term should have a stop token. The stream mode is not supported in speculative execution. E.g. correct: sgl.assistant("\nName:" + sgl.gen("name", stop="\n") + "\nBirthday:" + sgl.gen("birthday", stop="\n") + "\nJob:" + sgl.gen("job", stop="\n")) incorrect: s += sgl.assistant("\nName:" + sgl.gen("name", stop="\n")) s += sgl.assistant("\nBirthday:" + sgl.gen("birthday", stop="\n")) s += sgl.assistant("\nJob:" + sgl.gen("job", stop="\n")) export OPENAI_API_KEY=sk-****** python3 openai_chat_speculative.py """ import sglang as sgl from sglang import OpenAI, function, set_default_backend @function(num_api_spec_tokens=256) def gen_character_spec(s): s += sgl.system("You are a helpful assistant.") s += sgl.user("Construct a character within the following format:") s += sgl.assistant( "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\n" ) s += sgl.user("Please generate new Name, Birthday and Job.\n") s += sgl.assistant( "Name:" + sgl.gen("name", stop="\n") + "\nBirthday:" + sgl.gen("birthday", stop="\n") + "\nJob:" + sgl.gen("job", stop="\n") ) @function(num_api_spec_tokens=256) def gen_character_spec_no_few_shot(s): s += sgl.user("Construct a character. For each field stop with a newline\n") s += sgl.assistant( "Name:" + sgl.gen("name", stop="\n") + "\nAge:" + sgl.gen("age", stop="\n") + "\nJob:" + sgl.gen("job", stop="\n") ) @function def gen_character_normal(s): s += sgl.system("You are a helpful assistant.") s += sgl.user("What's the answer of 23 + 8?") s += sgl.assistant(sgl.gen("answer", max_tokens=64)) @function(num_api_spec_tokens=1024) def multi_turn_question(s, question_1, question_2): s += sgl.system("You are a helpful assistant.") s += sgl.user("Answer questions in the following format:") s += sgl.user( "Question 1: What is the capital of France?\nQuestion 2: What is the population of this city?\n" ) s += sgl.assistant( "Answer 1: The capital of France is Paris.\nAnswer 2: The population of Paris in 2024 is estimated to be around 2.1 million for the city proper.\n" ) s += sgl.user("Question 1: " + question_1 + "\nQuestion 2: " + question_2) s += sgl.assistant( "Answer 1: " + sgl.gen("answer_1", stop="\n") + "\nAnswer 2: " + sgl.gen("answer_2", stop="\n") ) def test_spec_single_turn(): backend.token_usage.reset() state = gen_character_spec.run() for m in state.messages(): print(m["role"], ":", m["content"]) print("\n-- name:", state["name"]) print("-- birthday:", state["birthday"]) print("-- job:", state["job"]) print(backend.token_usage) def test_inaccurate_spec_single_turn(): state = gen_character_spec_no_few_shot.run() for m in state.messages(): print(m["role"], ":", m["content"]) print("\n-- name:", state["name"]) print("\n-- age:", state["age"]) print("\n-- job:", state["job"]) def test_normal_single_turn(): state = gen_character_normal.run() for m in state.messages(): print(m["role"], ":", m["content"]) def test_spec_multi_turn(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions in the capital of the United States.", ) for m in state.messages(): print(m["role"], ":", m["content"]) print("\n-- answer_1 --\n", state["answer_1"]) print("\n-- answer_2 --\n", state["answer_2"]) def test_spec_multi_turn_stream(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", stream=True, ) for out in state.text_iter(): print(out, end="", flush=True) if __name__ == "__main__": backend = OpenAI("gpt-4-turbo") set_default_backend(backend) print("\n========== test spec single turn ==========\n") # expect reasonable answer for each field test_spec_single_turn() print("\n========== test inaccurate spec single turn ==========\n") # expect incomplete or unreasonable answers test_inaccurate_spec_single_turn() print("\n========== test normal single turn ==========\n") # expect reasonable answer test_normal_single_turn() print("\n========== test spec multi turn ==========\n") # expect answer with same format as in the few shot test_spec_multi_turn() print("\n========== test spec multi turn stream ==========\n") # expect error in stream_executor: stream is not supported... test_spec_multi_turn_stream() ================================================ FILE: examples/frontend_language/usage/openai_speculative.py ================================================ """ Usage: python3 openai_speculative.py """ from sglang import OpenAI, function, gen, set_default_backend @function(num_api_spec_tokens=64) def gen_character_spec(s): s += "Construct a character within the following format:\n" s += "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\n" s += "\nPlease generate new Name, Birthday and Job.\n" s += "Name:" + gen("name", stop="\n") + "\nBirthday:" + gen("birthday", stop="\n") s += "\nJob:" + gen("job", stop="\n") + "\n" @function def gen_character_no_spec(s): s += "Construct a character within the following format:\n" s += "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\n" s += "\nPlease generate new Name, Birthday and Job.\n" s += "Name:" + gen("name", stop="\n") + "\nBirthday:" + gen("birthday", stop="\n") s += "\nJob:" + gen("job", stop="\n") + "\n" @function(num_api_spec_tokens=64) def gen_character_spec_no_few_shot(s): # s += "Construct a character with name, birthday, and job:\n" s += "Construct a character:\n" s += "Name:" + gen("name", stop="\n") + "\nBirthday:" + gen("birthday", stop="\n") s += "\nJob:" + gen("job", stop="\n") + "\n" if __name__ == "__main__": backend = OpenAI("gpt-3.5-turbo-instruct") set_default_backend(backend) for function in [ gen_character_spec, gen_character_no_spec, gen_character_spec_no_few_shot, ]: backend.token_usage.reset() print(f"function: {function.func.__name__}") state = function.run() print("...name:", state["name"]) print("...birthday:", state["birthday"]) print("...job:", state["job"]) print(backend.token_usage) print() ================================================ FILE: examples/frontend_language/usage/parallel_sample.py ================================================ """ Usage: python3 parallel_sample.py """ import sglang as sgl @sgl.function def parallel_sample(s, question, n): s += ( "Question: Compute 1 + 2 + 3\n" "Reasoning: I need to use a calculator.\n" "Tool: calculator\n" "Answer: 6\n" "Question: Compute 3 + 2 + 2\n" "Reasoning: I will try a calculator.\n" "Tool: calculator\n" "Answer: 7\n" ) s += "Question: " + question + "\n" forks = s.fork(n) forks += "Reasoning:" + sgl.gen("reasoning", stop="\n") + "\n" forks += "Tool:" + sgl.gen("tool", choices=["calculator", "browser"]) + "\n" forks += "Answer:" + sgl.gen("answer", stop="\n") + "\n" forks.join() sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo-instruct")) # sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) state = parallel_sample.run(question="Compute 5 + 2 + 4.", n=5, temperature=1.0) for i in range(5): obj = { "reasoning": state["reasoning"][i], "tool": state["tool"][i], "answer": state["answer"][i], } print(f"[{i}], {obj}") ================================================ FILE: examples/frontend_language/usage/rag_using_parea/trace_and_evaluate_rag_using_parea.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# RAG Powered by SGLang & Chroma Evaluated using Parea\n", "\n", "In this notebook, we will build a simple RAG pipeline using SGLang to execute our LLM calls, Chroma as vector database for retrieval and [Parea](https://www.parea.ai) for tracing and evaluation. We will then evaluate the performance of our RAG pipeline. The dataset we will use was created by [Virat](https://twitter.com/virattt) and contains 100 questions, contexts and answers from the Airbnb 2023 10k filing.\n", "\n", "The RAG pipeline consists of two steps:\n", "1. Retrieval: Given a question, we retrieve the relevant context from all provided contexts.\n", "2. Generation: Given the question and the retrieved context, we generate an answer.\n", "\n", "ℹ️ This notebook requires an OpenAI API key.\n", "\n", "ℹ️ This notebook requires a Parea API key, which can be created [here](https://docs.parea.ai/api-reference/authentication#parea-api-key)." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setting up the environment\n", "\n", "We will first install the necessary packages: `sglang`, `parea-ai` and `chromadb`." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# note, if you use a Mac M1 chip, you might need to install grpcio 1.59.0 first such that installing chromadb works\n", "# !pip install grpcio==1.59.0\n", "\n", "!pip install sglang[openai] parea-ai chromadb" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Create a Parea API key as outlined [here](https://docs.parea.ai/api-reference/authentication#parea-api-key) and save it in a `.env` file as `PAREA_API_KEY=your-api-key`." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Indexing the data\n", "\n", "Now it's time to download the data & index it! For that, we create a collection called `contexts` in Chroma and add the contexts as documents." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\n", "import os\n", "from typing import List\n", "\n", "import chromadb\n", "\n", "path_qca = \"airbnb-2023-10k-qca.json\"\n", "\n", "if not os.path.exists(path_qca):\n", " !wget https://virattt.github.io/datasets/abnb-2023-10k.json -O airbnb-2023-10k-qca.json\n", "\n", "with open(path_qca, \"r\") as f:\n", " question_context_answers = json.load(f)\n", "\n", "chroma_client = chromadb.PersistentClient()\n", "collection = chroma_client.get_or_create_collection(name=\"contexts\")\n", "if collection.count() == 0:\n", " collection.add(\n", " documents=[qca[\"context\"] for qca in question_context_answers],\n", " ids=[str(i) for i in range(len(question_context_answers))],\n", " )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Defining the RAG pipeline\n", "\n", "We will start with importing the necessary packages, setting up tracing of OpenAI calls via Parea and setting OpenAI as the default backend for SGLang." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "import time\n", "\n", "from dotenv import load_dotenv\n", "\n", "from sglang import function, user, assistant, gen, set_default_backend, OpenAI\n", "from sglang.lang.interpreter import ProgramState\n", "from parea import Parea, trace\n", "\n", "load_dotenv()\n", "\n", "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n", "\n", "p = Parea(api_key=os.getenv(\"PAREA_API_KEY\"), project_name=\"rag_sglang\")\n", "p.integrate_with_sglang()\n", "\n", "set_default_backend(OpenAI(\"gpt-3.5-turbo\"))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we can define our retrieval step shown below. Notice, the `trace` decorator which will automatically trace inputs, output, latency, etc. of that call." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@trace\n", "def retrieval(question: str) -> List[str]:\n", " return collection.query(query_texts=[question], n_results=1)[\"documents\"][0]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Next we will define the generation step which uses SGLang to execute the LLM call." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@function\n", "def generation_sglang(s, question: str, *context: str):\n", " context = \"\\n\".join(context)\n", " s += user(\n", " f\"Given this question:\\n{question}\\n\\nAnd this context:\\n{context}\\n\\nAnswer the question.\"\n", " )\n", " s += assistant(gen(\"answer\"))\n", "\n", "\n", "@trace\n", "def generation(question: str, *context):\n", " state: ProgramState = generation_sglang.run(question, *context)\n", " while not state.stream_executor.is_finished:\n", " time.sleep(1)\n", " return state.stream_executor.variables[\"answer\"]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Finally, we can tie it together and execute a sample query." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@trace\n", "def rag_pipeline(question: str) -> str:\n", " contexts = retrieval(question)\n", " return generation(question, *contexts)\n", "\n", "\n", "rag_pipeline(\n", " \"When did the World Health Organization formally declare an end to the COVID-19 global health emergency?\"\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Debug Trace\n", "\n", "The output is unfortunately wrong! Using the traced pipeline, we can see that\n", "\n", "- the context is relevant to the question and contains the correct information\n", "- but, the generation step is cut off as max tokens is set to 16\n", "\n", "When opening the generation step in the playground and rerunning the prompt with max. tokens set to 1000, the correct answer is produced.\n", "\n", "![RAG Trace](https://drive.google.com/uc?id=1QI243ogGjzbO01tUrR72g9rFoGzUJqVH)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Evaluating RAG Pipelines\n", "\n", "Before we apply above's fix, let's dive into evaluating RAG pipelines.\n", "\n", "RAG pipelines consist of a retrieval step to fetch relevant information and a generation step to generate a response to a users question. A RAG pipeline can fail at either step. E.g. the retrieval step can fail to find relevant information which makes generating the correct impossible. Another failure mode is that the generation step doesn't leverage the retrieved information correctly. We will apply the following evaluation metrics to understand different failure modes:\n", "\n", "- `context_relevancy`: measures how relevant the context is given the question\n", "- `percent_target_supported_by_context`: measures how much of the target answer is supported by the context; this will give an upper ceiling of how well the generation step can perform\n", "- `answer_context_faithfulness`: measures how much the generated answer utilizes the context\n", "- `answer_matches_target`: measures how well the generated answer matches the target answer judged by a LLM and gives a sense of accuracy of our entire pipeline\n", "\n", "To use these evaluation metrics, we can import them from `parea.evals.rag` and `parea.evals.general` and apply them to a function by specifying in the `trace` decorator which evaluation metrics to use. The `@trace` decorator will automatically log the results of the evaluation metrics to the Parea dashboard.\n", "\n", "Applying them to the retrieval step:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from parea.evals.rag import (\n", " context_query_relevancy_factory,\n", " percent_target_supported_by_context_factory,\n", ")\n", "\n", "context_relevancy_eval = context_query_relevancy_factory()\n", "percent_target_supported_by_context = percent_target_supported_by_context_factory()\n", "\n", "\n", "@trace(eval_funcs=[context_relevancy_eval, percent_target_supported_by_context])\n", "def retrieval(question: str) -> List[str]:\n", " return collection.query(query_texts=[question], n_results=1)[\"documents\"][0]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we can apply `answer_context_faithfulness` and `answer_matches_target` to the generation step." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from parea.evals.general import answer_matches_target_llm_grader_factory\n", "from parea.evals.rag import answer_context_faithfulness_statement_level_factory\n", "\n", "answer_context_faithfulness = answer_context_faithfulness_statement_level_factory()\n", "answer_matches_target_llm_grader = answer_matches_target_llm_grader_factory()\n", "\n", "\n", "@function\n", "def generation_sglang(s, question: str, *context: str):\n", " context = \"\\n\".join(context)\n", " s += user(\n", " f\"Given this question:\\n{question}\\n\\nAnd this context:\\n{context}\\n\\nAnswer the question.\"\n", " )\n", " s += assistant(gen(\"answer\", max_tokens=1_000))\n", "\n", "\n", "@trace(eval_funcs=[answer_context_faithfulness, answer_matches_target_llm_grader])\n", "def generation(question: str, *context):\n", " state: ProgramState = generation_sglang.run(question, *context)\n", " while not state.stream_executor.is_finished:\n", " time.sleep(1)\n", " return state.stream_executor.variables[\"answer\"]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Finally, we tie them together & execute the original sample query." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@trace\n", "def rag_pipeline(question: str) -> str:\n", " contexts = retrieval(question)\n", " return generation(question, *contexts)\n", "\n", "\n", "rag_pipeline(\n", " \"When did the World Health Organization formally declare an end to the COVID-19 global health emergency?\"\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Great, the answer is correct! Can you spot the line where we fixed the output truncation issue?\n", "\n", "The evaluation scores appear in the bottom right of the logs (screenshot below). Note, that there is no score for `answer_matches_target_llm_grader` and `percent_target_supported_by_context` as these evals are automatically skipped if the target answer is not provided.\n", "\n", "![Fixed Max. Tokens](max-tokens-fixed-rag-trace.png)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Running an experiment\n", "\n", "Now we are (almost) ready to evaluate the performance of our RAG pipeline on the entire dataset. First, we will need to apply the `nest_asyncio` package to avoid issues with the Jupyter notebook event loop." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!pip install nest-asyncio\n", "import nest_asyncio\n", "\n", "nest_asyncio.apply()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Running the actual experiment is straight-forward. For that we use `p.experiment` to initialize the experiment with a name, the data (list of key-value pairs fed into our entry function) and the entry function. We then call `run` on the experiment to execute it. Note, that `target` is a reserved key in the data dictionary and will be used as the target answer for evaluation." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "e = p.experiment(\n", " \"RAG\",\n", " data=[\n", " {\n", " \"question\": qca[\"question\"],\n", " \"target\": qca[\"answer\"],\n", " }\n", " for qca in question_context_answers\n", " ],\n", " func=rag_pipeline,\n", ").run()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Analyzing the results\n", "\n", "When opening above experiment, we will see an overview of the experiment as shown below. The upper half shows a summary of the statistics on the left and charts to investigate the distribution and relationships of scores on the right. The lower half is a table with the individual traces which we can use to debug individual samples.\n", "\n", "When looking at the statistics, we can see that the accuracy of our RAG pipeline is 22% as measured by `answer_matches_target_llm_grader`. Though when checking the quality of our retrieval step (`context_query_relevancy`), we can see that our retrieval step is fetching relevant information in only 27% of all samples. As shown in the GIF, we investigate the relationship between the two and see the two scores have 95% agreement. This confirms that the retrieval step is a major bottleneck for our RAG pipeline. So, now it's your turn to improve the retrieval step!\n", "\n", "Note, above link isn't publicly accessible but the experiment can be accessed through [here](https://app.parea.ai/public-experiments/parea/rag_sglang/30f0244a-d56c-44ff-bdfb-8f47626304b6).\n", "\n", "![Experiment Results](https://drive.google.com/uc?id=1KMtJBU47nPB02Pvv3SPPTK7RnHRh5YdA)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "language_info": { "codemirror_mode": { "name": "ipython", "version": 2 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython2" } }, "nbformat": 4, "nbformat_minor": 0 } ================================================ FILE: examples/frontend_language/usage/readme_examples.py ================================================ """ Usage: python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 python readme_examples.py """ import sglang as sgl @sgl.function def tool_use(s, question): s += "To answer this question: " + question + ". " s += ( "I need to use a " + sgl.gen("tool", choices=["calculator", "search engine"]) + ". " ) if s["tool"] == "calculator": s += "The math expression is" + sgl.gen("expression") elif s["tool"] == "search engine": s += "The key word to search is" + sgl.gen("word") @sgl.function def tip_suggestion(s): s += ( "Here are two tips for staying healthy: " "1. Balanced Diet. 2. Regular Exercise.\n\n" ) forks = s.fork(2) for i, f in enumerate(forks): f += f"Now, expand tip {i+1} into a paragraph:\n" f += sgl.gen(f"detailed_tip", max_tokens=256, stop="\n\n") s += "Tip 1:" + forks[0]["detailed_tip"] + "\n" s += "Tip 2:" + forks[1]["detailed_tip"] + "\n" s += "In summary" + sgl.gen("summary") @sgl.function def regular_expression_gen(s): s += "Q: What is the IP address of the Google DNS servers?\n" s += "A: " + sgl.gen( "answer", temperature=0, regex=r"((25[0-5]|2[0-4]\d|[01]?\d\d?).){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)", ) @sgl.function def text_qa(s, question): s += "Q: " + question + "\n" s += "A:" + sgl.gen("answer", stop="\n") def driver_tool_use(): state = tool_use.run(question="What is the capital of the United States?") print(state.text()) print("\n") def driver_tip_suggestion(): state = tip_suggestion.run() print(state.text()) print("\n") def driver_regex(): state = regular_expression_gen.run() print(state.text()) print("\n") def driver_batching(): states = text_qa.run_batch( [ {"question": "What is the capital of the United Kingdom?"}, {"question": "What is the capital of France?"}, {"question": "What is the capital of Japan?"}, ], progress_bar=True, ) for s in states: print(s.text()) print("\n") def driver_stream(): state = text_qa.run( question="What is the capital of France?", temperature=0.1, stream=True ) for out in state.text_iter(): print(out, end="", flush=True) print("\n") if __name__ == "__main__": # sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo-instruct")) sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) driver_tool_use() driver_tip_suggestion() driver_regex() driver_batching() driver_stream() ================================================ FILE: examples/frontend_language/usage/sgl_gen_min_tokens.py ================================================ """ This example demonstrates how to use `min_tokens` to enforce sgl.gen to generate a longer sequence Usage: python3 sgl_gen_min_tokens.py """ import sglang as sgl @sgl.function def long_answer(s): s += sgl.user("What is the capital of the United States?") s += sgl.assistant(sgl.gen("answer", min_tokens=64, max_tokens=128)) @sgl.function def short_answer(s): s += sgl.user("What is the capital of the United States?") s += sgl.assistant(sgl.gen("answer")) if __name__ == "__main__": runtime = sgl.Runtime(model_path="meta-llama/Meta-Llama-3.1-8B-Instruct") sgl.set_default_backend(runtime) state = long_answer.run() print("=" * 20) print("Longer Answer", state["answer"]) state = short_answer.run() print("=" * 20) print("Short Answer", state["answer"]) runtime.shutdown() ================================================ FILE: examples/frontend_language/usage/streaming.py ================================================ """ Usage: python3 streaming.py """ import asyncio import sglang as sgl @sgl.function def multi_turn_question(s, question_1, question_2): s += sgl.system("You are a helpful assistant.") s += sgl.user(question_1) s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) s += sgl.user(question_2) s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo")) def stream_a_variable(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", stream=True, ) for out in state.text_iter(var_name="answer_2"): print(out, end="", flush=True) print("\n") async def async_stream(): state = multi_turn_question.run( question_1="What is the capital of the United States?", question_2="List two local attractions.", stream=True, ) async for out in state.text_async_iter(var_name="answer_2"): print(out, end="", flush=True) print("\n") if __name__ == "__main__": stream_a_variable() asyncio.run(async_stream()) ================================================ FILE: examples/frontend_language/usage/triton/Dockerfile ================================================ FROM nvcr.io/nvidia/tritonserver:24.01-py3 WORKDIR /opt RUN git clone https://github.com/sgl-project/sglang.git WORKDIR /opt/sglang RUN pip install --upgrade pip && \ pip install -e "python[all]" && \ pip install datasets ================================================ FILE: examples/frontend_language/usage/triton/README.md ================================================ # sglang_triton Build the docker image: ``` docker build -t sglang-triton . ``` Then do: ``` docker run -ti --gpus=all --network=host --name sglang-triton -v ./models:/mnt/models sglang-triton ``` inside the docker container: ``` cd sglang python3 -m sglang.launch_server --model-path mistralai/Mistral-7B-Instruct-v0.2 --port 30000 --mem-fraction-static 0.9 ``` with another shell, inside the docker container: ``` docker exec -ti sglang-triton /bin/bash cd /mnt tritonserver --model-repository=/mnt/models ``` Send request to the server: ``` curl -X POST http://localhost:8000/v2/models/character_generation/generate \ -H "Content-Type: application/json" \ -d '{ "INPUT_TEXT": ["harry"] }' ``` ================================================ FILE: examples/frontend_language/usage/triton/models/character_generation/1/model.py ================================================ import numpy import triton_python_backend_utils as pb_utils from pydantic import BaseModel import sglang as sgl from sglang import function from sglang.srt.constrained.outlines_backend import build_regex_from_object sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) class Character(BaseModel): name: str eye_color: str house: str @function def character_gen(s, name): s += ( name + " is a character in Harry Potter. Please fill in the following information about this character.\n" ) s += sgl.gen( "json_output", max_tokens=256, regex=build_regex_from_object(Character) ) class TritonPythonModel: def initialize(self, args): print("Initialized.") def execute(self, requests): responses = [] for request in requests: tensor_in = pb_utils.get_input_tensor_by_name(request, "INPUT_TEXT") if tensor_in is None: return pb_utils.InferenceResponse(output_tensors=[]) input_list_names = [ i.decode("utf-8") if isinstance(i, bytes) else i for i in tensor_in.as_numpy().tolist() ] input_list_dicts = [{"name": i} for i in input_list_names] states = character_gen.run_batch(input_list_dicts) character_strs = [state.text() for state in states] tensor_out = pb_utils.Tensor( "OUTPUT_TEXT", numpy.array(character_strs, dtype=object) ) responses.append(pb_utils.InferenceResponse(output_tensors=[tensor_out])) return responses ================================================ FILE: examples/frontend_language/usage/triton/models/character_generation/config.pbtxt ================================================ name: "character_generation" backend: "python" input [ { name: "INPUT_TEXT" data_type: TYPE_STRING dims: [ -1 ] } ] output [ { name: "OUTPUT_TEXT" data_type: TYPE_STRING dims: [ -1 ] } ] instance_group [ { count: 1 kind: KIND_GPU gpus: [ 0 ] } ] ================================================ FILE: examples/monitoring/README.md ================================================ # SGLang Monitoring Setup This directory contains a ready-to-use monitoring setup for SGLang using Prometheus and Grafana. ## Prerequisites - Docker and Docker Compose installed - SGLang server running with metrics enabled ## Usage 1. Start your SGLang server with metrics enabled: ```bash python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --port 30000 --enable-metrics ``` By default, the metrics server will run on `127.0.0.1:30000`. 2. Start the monitoring stack: ```bash cd examples/monitoring docker compose up ``` 3. Access the monitoring interfaces: - Grafana: [http://localhost:3000](http://localhost:3000) - Prometheus: [http://localhost:9090](http://localhost:9090) Default Grafana login credentials: - Username: `admin` - Password: `admin` You'll be prompted to change the password on first login. 4. The SGLang dashboard will be automatically available in the "SGLang Monitoring" folder. ## Troubleshooting ### Port Conflicts If you see errors like "port is already allocated": 1. Check if you already have Prometheus or Grafana running: ```bash docker ps | grep -E 'prometheus|grafana' ``` 2. Stop any conflicting containers: ```bash docker stop ``` 3. Ensure no other services are using ports 9090 and 3000: ```bash lsof -i :9090 lsof -i :3000 ``` ### Connection Issues If Grafana cannot connect to Prometheus: 1. Check that both services are running 2. Verify the datasource configuration in Grafana 3. Check that your SGLang server is properly exposing metrics ## Configuration - Prometheus configuration: `prometheus.yaml` - Docker Compose configuration: `docker-compose.yaml` - Grafana datasource: `grafana/datasources/datasource.yaml` - Grafana dashboard configuration: `grafana/dashboards/config/dashboard.yaml` - SGLang dashboard JSON: `grafana/dashboards/json/sglang-dashboard.json` ## Customization You can customize the monitoring setup by modifying the configuration files as needed. ================================================ FILE: examples/monitoring/docker-compose.yaml ================================================ version: '3' services: prometheus: image: prom/prometheus:latest container_name: prometheus network_mode: host volumes: - ./prometheus.yaml:/etc/prometheus/prometheus.yml command: - '--config.file=/etc/prometheus/prometheus.yml' - '--storage.tsdb.path=/prometheus' grafana: image: grafana/grafana:latest container_name: grafana network_mode: host volumes: - ./grafana/datasources:/etc/grafana/provisioning/datasources - ./grafana/dashboards/config:/etc/grafana/provisioning/dashboards - ./grafana/dashboards/json:/var/lib/grafana/dashboards environment: - GF_AUTH_ANONYMOUS_ENABLED=true - GF_AUTH_ANONYMOUS_ORG_ROLE=Viewer - GF_AUTH_BASIC_ENABLED=false - GF_USERS_ALLOW_SIGN_UP=false - GF_DASHBOARDS_DEFAULT_HOME_DASHBOARD_PATH=/var/lib/grafana/dashboards/sglang-dashboard.json depends_on: - prometheus ================================================ FILE: examples/monitoring/grafana/dashboards/config/dashboard.yaml ================================================ apiVersion: 1 providers: - name: 'SGLang' orgId: 1 folder: 'SGLang Monitoring' type: file disableDeletion: false updateIntervalSeconds: 10 allowUiUpdates: false options: path: /var/lib/grafana/dashboards ================================================ FILE: examples/monitoring/grafana/dashboards/json/sglang-dashboard.json ================================================ { "annotations": { "list": [ { "builtIn": 1, "datasource": { "type": "grafana", "uid": "-- Grafana --" }, "enable": true, "hide": true, "iconColor": "rgba(0, 211, 255, 1)", "name": "Annotations & Alerts", "type": "dashboard" } ] }, "editable": true, "fiscalYearStartMonth": 0, "graphTooltip": 0, "id": 8, "links": [], "panels": [ { "datasource": { "default": true, "type": "prometheus" }, "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisBorderShow": false, "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "barWidthFactor": 0.6, "drawStyle": "line", "fillOpacity": 0, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "insertNulls": false, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "auto", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green" }, { "color": "red", "value": 80 } ] } }, "overrides": [] }, "gridPos": { "h": 8, "w": 12, "x": 0, "y": 0 }, "id": 14, "options": { "legend": { "calcs": [], "displayMode": "list", "placement": "bottom", "showLegend": true }, "tooltip": { "hideZeros": false, "mode": "single", "sort": "none" } }, "pluginVersion": "11.6.0", "targets": [ { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "disableTextWrap": false, "editorMode": "code", "expr": "histogram_quantile(0.99, sum by (le) (rate(sglang_e2e_request_latency_seconds_bucket[$__rate_interval])))\r\n", "fullMetaSearch": false, "includeNullMetadata": true, "instant": false, "legendFormat": "P99", "range": true, "refId": "A", "useBackend": false }, { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "disableTextWrap": false, "editorMode": "code", "expr": "histogram_quantile(0.9, sum by (le) (rate(sglang_e2e_request_latency_seconds_bucket[$__rate_interval])))\r\n", "fullMetaSearch": false, "hide": false, "includeNullMetadata": true, "instant": false, "legendFormat": "P90", "range": true, "refId": "B", "useBackend": false }, { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "disableTextWrap": false, "editorMode": "code", "expr": "histogram_quantile(0.5, sum by (le) (rate(sglang_e2e_request_latency_seconds_bucket[$__rate_interval])))\r\n", "fullMetaSearch": false, "hide": false, "includeNullMetadata": true, "instant": false, "legendFormat": "P50", "range": true, "refId": "C", "useBackend": false }, { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "disableTextWrap": false, "editorMode": "code", "expr": "avg(rate(sglang_e2e_request_latency_seconds_sum[$__rate_interval]) / rate(sglang_e2e_request_latency_seconds_count[$__rate_interval]))\r\n", "fullMetaSearch": false, "hide": false, "includeNullMetadata": true, "instant": false, "legendFormat": "Avg", "range": true, "refId": "D", "useBackend": false } ], "title": "End-to-End Request Latency", "type": "timeseries" }, { "datasource": { "default": true, "type": "prometheus" }, "fieldConfig": { "defaults": { "custom": { "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "scaleDistribution": { "type": "linear" } } }, "overrides": [] }, "gridPos": { "h": 8, "w": 12, "x": 12, "y": 0 }, "id": 17, "maxDataPoints": 30, "options": { "calculate": false, "calculation": { "yBuckets": { "scale": { "type": "linear" } } }, "cellGap": 1, "cellValues": {}, "color": { "exponent": 0.5, "fill": "dark-orange", "mode": "scheme", "reverse": false, "scale": "exponential", "scheme": "Spectral", "steps": 64 }, "exemplars": { "color": "rgba(255,0,255,0.7)" }, "filterValues": { "le": 1e-9 }, "legend": { "show": true }, "rowsFrame": { "layout": "auto" }, "tooltip": { "mode": "single", "showColorScale": true, "yHistogram": false }, "yAxis": { "axisPlacement": "left", "reverse": false, "unit": "secs" } }, "pluginVersion": "11.6.0", "targets": [ { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "disableTextWrap": false, "editorMode": "builder", "expr": "sum(increase(sglang_e2e_request_latency_seconds_bucket{model_name=~\"$model_name\"}[$__rate_interval])) by (le)\r\n", "format": "heatmap", "fullMetaSearch": false, "includeNullMetadata": true, "instant": false, "legendFormat": "{{le}}", "range": true, "refId": "A", "useBackend": false } ], "title": "End-to-End Request Latency(s) Heatmap", "type": "heatmap" }, { "datasource": { "default": true, "type": "prometheus" }, "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisBorderShow": false, "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "barWidthFactor": 0.6, "drawStyle": "line", "fillOpacity": 0, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "insertNulls": false, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "auto", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green" }, { "color": "red", "value": 80 } ] } }, "overrides": [] }, "gridPos": { "h": 8, "w": 12, "x": 0, "y": 8 }, "id": 20, "options": { "legend": { "calcs": [], "displayMode": "list", "placement": "bottom", "showLegend": true }, "tooltip": { "hideZeros": false, "mode": "single", "sort": "none" } }, "pluginVersion": "11.6.0", "targets": [ { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "disableTextWrap": false, "editorMode": "code", "expr": "histogram_quantile(0.99, sum by (le) (rate(sglang_time_to_first_token_seconds_bucket[$__rate_interval])))\r\n", "fullMetaSearch": false, "includeNullMetadata": true, "instant": false, "legendFormat": "P99", "range": true, "refId": "A", "useBackend": false }, { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "disableTextWrap": false, "editorMode": "code", "expr": "histogram_quantile(0.9, sum by (le) (rate(sglang_time_to_first_token_seconds_bucket[$__rate_interval])))\r\n", "fullMetaSearch": false, "hide": false, "includeNullMetadata": true, "instant": false, "legendFormat": "P90", "range": true, "refId": "B", "useBackend": false }, { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "disableTextWrap": false, "editorMode": "code", "expr": "histogram_quantile(0.5, sum by (le) (rate(sglang_time_to_first_token_seconds_bucket[$__rate_interval])))\r\n", "fullMetaSearch": false, "hide": false, "includeNullMetadata": true, "instant": false, "legendFormat": "P50", "range": true, "refId": "C", "useBackend": false }, { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "disableTextWrap": false, "editorMode": "code", "expr": "avg(rate(sglang_time_to_first_token_seconds_sum[$__rate_interval]) / rate(sglang_time_to_first_token_seconds_count[$__rate_interval]))\r\n", "fullMetaSearch": false, "hide": false, "includeNullMetadata": true, "instant": false, "legendFormat": "Avg", "range": true, "refId": "D", "useBackend": false } ], "title": "Time-To-First-Token Latency", "type": "timeseries" }, { "datasource": { "default": true, "type": "prometheus" }, "fieldConfig": { "defaults": { "custom": { "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "scaleDistribution": { "type": "linear" } } }, "overrides": [] }, "gridPos": { "h": 8, "w": 12, "x": 12, "y": 8 }, "id": 19, "maxDataPoints": 30, "options": { "calculate": false, "calculation": { "xBuckets": { "value": "" }, "yBuckets": { "mode": "size", "scale": { "type": "linear" }, "value": "" } }, "cellGap": 1, "color": { "exponent": 0.5, "fill": "dark-orange", "mode": "scheme", "reverse": false, "scale": "exponential", "scheme": "Spectral", "steps": 64 }, "exemplars": { "color": "rgba(255,0,255,0.7)" }, "filterValues": { "le": 1e-9 }, "legend": { "show": true }, "rowsFrame": { "layout": "auto" }, "tooltip": { "mode": "single", "showColorScale": true, "yHistogram": false }, "yAxis": { "axisPlacement": "left", "reverse": false } }, "pluginVersion": "11.6.0", "targets": [ { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "disableTextWrap": false, "editorMode": "builder", "exemplar": false, "expr": "sum by(le) (increase(sglang_time_to_first_token_seconds_bucket{model_name=~\"$model_name\"}[$__rate_interval]))", "format": "heatmap", "fullMetaSearch": false, "includeNullMetadata": true, "instant": false, "interval": "", "legendFormat": "{{le}}", "range": true, "refId": "A", "useBackend": false } ], "title": "Time-To-First-Token Seconds Heatmap", "type": "heatmap" }, { "datasource": { "default": true, "type": "prometheus" }, "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisBorderShow": false, "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "barWidthFactor": 0.6, "drawStyle": "line", "fillOpacity": 0, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "insertNulls": false, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "auto", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green" }, { "color": "red", "value": 80 } ] } }, "overrides": [] }, "gridPos": { "h": 8, "w": 12, "x": 0, "y": 16 }, "id": 7, "options": { "legend": { "calcs": [], "displayMode": "list", "placement": "bottom", "showLegend": true }, "tooltip": { "hideZeros": false, "mode": "single", "sort": "none" } }, "pluginVersion": "11.6.0", "targets": [ { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "disableTextWrap": false, "editorMode": "code", "expr": "sglang_num_running_reqs", "fullMetaSearch": false, "includeNullMetadata": true, "instant": false, "interval": "", "legendFormat": "{{instance}}", "range": true, "refId": "A", "useBackend": false } ], "title": "Num Running Requests", "type": "timeseries" }, { "datasource": { "default": true, "type": "prometheus" }, "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisBorderShow": false, "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "barWidthFactor": 0.6, "drawStyle": "line", "fillOpacity": 0, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "insertNulls": false, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "auto", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green" }, { "color": "red", "value": 80 } ] } }, "overrides": [] }, "gridPos": { "h": 8, "w": 12, "x": 12, "y": 16 }, "id": 18, "options": { "legend": { "calcs": [], "displayMode": "list", "placement": "bottom", "showLegend": true }, "tooltip": { "hideZeros": false, "mode": "single", "sort": "none" } }, "pluginVersion": "11.6.0", "targets": [ { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "editorMode": "code", "expr": "sglang_gen_throughput", "instant": false, "legendFormat": "{{instance}}", "range": true, "refId": "A" } ], "title": "Token Generation Throughput (Tokens / S)", "type": "timeseries" }, { "datasource": { "default": true, "type": "prometheus" }, "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisBorderShow": false, "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "barWidthFactor": 0.6, "drawStyle": "line", "fillOpacity": 0, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "insertNulls": false, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "auto", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green" }, { "color": "red", "value": 80 } ] } }, "overrides": [] }, "gridPos": { "h": 8, "w": 12, "x": 0, "y": 24 }, "id": 11, "options": { "legend": { "calcs": [], "displayMode": "list", "placement": "bottom", "showLegend": true }, "tooltip": { "hideZeros": false, "mode": "single", "sort": "none" } }, "pluginVersion": "11.6.0", "targets": [ { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "disableTextWrap": false, "editorMode": "code", "expr": "sglang_cache_hit_rate", "fullMetaSearch": false, "includeNullMetadata": true, "instant": false, "legendFormat": "{{instance}}", "range": true, "refId": "A", "useBackend": false } ], "title": "Cache Hit Rate", "type": "timeseries" }, { "datasource": { "default": true, "type": "prometheus" }, "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisBorderShow": false, "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "barWidthFactor": 0.6, "drawStyle": "line", "fillOpacity": 0, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "insertNulls": false, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "auto", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green" }, { "color": "red", "value": 80 } ] } }, "overrides": [] }, "gridPos": { "h": 8, "w": 12, "x": 12, "y": 24 }, "id": 8, "options": { "legend": { "calcs": [], "displayMode": "list", "placement": "bottom", "showLegend": true }, "tooltip": { "hideZeros": false, "mode": "single", "sort": "none" } }, "pluginVersion": "11.6.0", "targets": [ { "datasource": { "type": "prometheus", "uid": "ddyfngn31dg5cf" }, "disableTextWrap": false, "editorMode": "code", "expr": "sglang_num_queue_reqs", "fullMetaSearch": false, "includeNullMetadata": true, "instant": false, "legendFormat": "{{instance}}", "range": true, "refId": "A", "useBackend": false } ], "title": "Number Queued Requests", "type": "timeseries" } ], "preload": false, "refresh": "5s", "schemaVersion": 41, "tags": [], "templating": { "list": [ { "current": { "text": "127.0.0.1:30000", "value": "127.0.0.1:30000" }, "datasource": { "type": "prometheus" }, "definition": "label_values(instance)", "includeAll": false, "label": "instance", "name": "instance", "options": [], "query": { "qryType": 1, "query": "label_values(instance)", "refId": "PrometheusVariableQueryEditor-VariableQuery" }, "refresh": 1, "regex": "", "type": "query" }, { "current": { "text": "meta-llama/Llama-3.1-8B-Instruct", "value": "meta-llama/Llama-3.1-8B-Instruct" }, "datasource": { "type": "prometheus" }, "definition": "label_values(model_name)", "includeAll": false, "label": "model name", "name": "model_name", "options": [], "query": { "qryType": 1, "query": "label_values(model_name)", "refId": "PrometheusVariableQueryEditor-VariableQuery" }, "refresh": 1, "regex": "", "type": "query" } ] }, "time": { "from": "now-30m", "to": "now" }, "timepicker": {}, "timezone": "browser", "title": "SGLang Dashboard", "uid": "sglang-dashboard", "version": 11 } ================================================ FILE: examples/monitoring/grafana/datasources/datasource.yaml ================================================ apiVersion: 1 datasources: - name: Prometheus type: prometheus access: proxy url: http://localhost:9090 isDefault: true editable: false ================================================ FILE: examples/monitoring/opentelemetry.yaml ================================================ receivers: otlp: protocols: grpc: endpoint: 0.0.0.0:4317 http: endpoint: 0.0.0.0:4318 processors: batch: exporters: otlp: endpoint: jaeger:4317 tls: insecure: true file: path: /tmp/otel_trace.json extensions: health_check: pprof: zpages: service: extensions: [health_check, pprof, zpages] pipelines: traces: receivers: [otlp] processors: [batch] exporters: [otlp, file] metrics: receivers: [otlp] processors: [batch] exporters: [otlp] logs: receivers: [otlp] processors: [batch] exporters: [otlp] ================================================ FILE: examples/monitoring/prometheus.yaml ================================================ # prometheus.yaml global: scrape_interval: 5s evaluation_interval: 30s scrape_configs: - job_name: sglang static_configs: - targets: - '127.0.0.1:30000' ================================================ FILE: examples/monitoring/tracing_compose.yaml ================================================ services: otel-collector: image: docker.io/otel/opentelemetry-collector volumes: - ./opentelemetry.yaml:/etc/otelcol/config.yaml - /tmp:/tmp ports: - "4317:4317" # OTLP gRPC - "4318:4318" # OTLP HTTP depends_on: - jaeger restart: unless-stopped jaeger: image: jaegertracing/all-in-one container_name: jaeger ports: - "16686:16686" environment: - COLLECTOR_OTLP_ENABLED=true restart: unless-stopped ================================================ FILE: examples/profiler/nsys_profile_tools/README.md ================================================ # gputrc2graph.py This script processes NVIDIA Nsight Systems (`nsys`) GPU trace files (`.nsys-rep`) with -t cuda tracing enabled, and generates kernel-level summaries and visualizations of GPU and non-GPU time. It is useful for profiling and analyzing nsys profile output. ## Usage ### Command-line Arguments - `--in_file` **(required)** List of input files and their metadata. Each entry should be in the format: `,,,` - `nsys-rep`: Path to the `.nsys-rep` file. - `engine`: Engine name (e.g., `sglang`). - `model`: Model name (e.g., `llama`, `gpt-oss`, `ds`). - `elapsed_nonprofiled_sec`: Wall-clock runtime (in seconds) without profiling. Specify `0` to use the elapsed time from the nsys-rep file (this may inflate non-GPU time if actual runtime without profiling is less). Multiple entries can be provided, separated by spaces. - `--out_dir` Output directory for the generated CSV and HTML files. If not specified, results are saved in the current directory. - `--title` Title for the HTML chart/visualization. - `--nsys_cmd` Path to the `nsys` command. Default: `nsys` (assumes it is in your PATH). Use this if `nsys` is not in your system PATH. ## Notes - Make sure you have pandas installed. Any version is fine. - Make sure [nsys](https://developer.nvidia.com/nsight-systems/get-started) is installed, and specify the path to the `nsys` command with `--nsys_cmd` if it is not in your PATH. The nsys version must be >= the nsys profile version that was used to collect the traces when profiling the server, so that nsys can process the nsys-rep that was generated. - For more details on available engines and models, see the help string in the script or run: ```bash python3 gputrc2graph.py --help ``` ## Example 1: analyze a single profile To analyze the GPU cycles of for example, a llama-3.1-8B model with sglang: 1. Run the following command to collect nsys profile, for sglang server config. ```bash nsys profile -t cuda -o nsys_res -f true --trace-fork-before-exec=true \ --cuda-graph-trace=node --delay --duration \ python3 -m sglang.launch_server --model meta-llama/Llama-3.1-8B ... ``` where: - DELAY: how many seconds to delay nsys from collecting profiles, needed so that profiles aren't captured till sglang server has come up and load generation starts. - DURATION: how many seconds for nsys profile to run before generating the profile. This should be > the duration of the run. 2. After the server starts, run the client load generation command. Once the test completes, after DURATION amount of time, nsys profile will generate an nsys_res.nsys-rep file and shut down the server. 3. Run step #1 again, this time starting up the server without collecting the profile. 4. Run step #2 again, and record the total time to complete the test in seconds. This value will be used by the script to calculate the CPU(non-GPU) seconds for the analysis. 5. Say the run elapsed time from step #4 is 132 seconds. Run script to analyze: ```bash python3 gputrc2graph.py \ --in_file run1.nsys-rep,sglang,llama,132 ``` The command will produce 2 files for analysis: - result.html: this categorizes kernel names into different categories in a stacked bar chart. - result.csv: shows how the kernel names are mapped to the different categories. ### HTML visualization with result.html The html file shows the number of elapsed seconds due to different GPU Substages or categories, which consist of attention kernels as the biggest category, at 63 seconds, followed by "gemm" kernels. This lets the user prioritize the kernels to focus on for performance optimizations. There's also an appended data table underneath the bar chart for copying out to other post-processing tools. ### Kernel to category mapping with result.csv Suppose the user would like to focus on improving triton kernels. It's not the biggest consumer of cycles at .01 sec but perhaps it hasn't been optimized. The next step is to use the result.csv to dive into what the kernels are which compose the triton kernel GPU cycles. ## Example 2: analyze multiple profiles Suppose the user has multiple nsys trace files, captured for different models, say llama and gpt-oss in this case, and wish to compare their GPU/non-GPU time, something like the following command can be used. ```bash python3 gputrc2graph.py \ --in_file run1.nsys-rep,sglang,llama,100 run2.nsys-rep,sglang,gpt-oss,102 \ --out_dir results ``` The analysis process is similar to example 1 but now there will be multiple stack bar charts that can be compared. The categories for the different kernels will remain the same, so that it's easy to compare the GPU cycles for the same categories. Once a category is shown to have more cycles for one configuration than another, the next step would be to use the csv file to see what kernels are mapped into that category, and which kernels are taking the largest amount of time which would cause a difference for the overall category. ## Example 3: add new classification for a new model To create a new engine DEF with model ABC, just add another json file in the same directory as gputrc2graph.py with the same format as the other json files. The script will automatically pick up all the json files in the same directory as engine/model specifications. Then, for this new model, suppose there are 4 kernels to be classified into "gemm" and "attn", where the gemm kernels have names with "*H*" or "*I*" in them, and attn kernels have names with "*J*" or "*K*" in them, just add another .json file in the same directory as gputrc2graph.py with the same format as the other json files, like the following: ```json { "DEF": { "ABC": { "H|I": "gemm", "J|K": "attn", "CUDA mem": "non-gpu-H_D_memops", ".*": "misc" } } } ``` Each entry in the dictionary consists of: - key: a regex used to classify the kernels - value: the category to classify the kernels into. The last 2 entries are common for all engine/models, consisting of CUDA memory operations and a 'misc' for anything that's leftover and can't be classified. When invoking gputrc2graph.py, specify a trace file with this new model/engine like the following: ```bash --in_file new.nsys-rep,DEF,ABC, ``` If the engine_DEF.json file already exists, just add the model as a new node in the existing engine file, after the other models. ================================================ FILE: examples/profiler/nsys_profile_tools/gputrc2graph.py ================================================ """ This generates gpu kernel analysis output from nsys rep. Will call nsys stats -r cuda_gpu_kern_trace, get non-overlapped gpu cycles, then generate csv and html output for analysis """ import argparse import logging import os import shlex import regex as re logger = logging.getLogger(__name__) # helper data class for annotating kernels def load_engine_model(): """returns engine_model built from all json files in the current dir""" import glob import json engine_model = {} json_files = glob.glob(os.path.join(os.path.dirname(__file__) or ".", "*.json")) for fname in json_files: with open(fname, encoding="utf-8") as f: file_engine_model = json.load(f) for engine, models in file_engine_model.items(): engine_model.setdefault(engine, {}).update(models) return engine_model class GPUTrace2Graph: """ Parses output of nsys report, generates csv and bar chart output """ def __init__(self): import pandas as pd # avoid importing till needed self.pd = pd # helper functions for generating trace->summary csvs def gen_nonoverlapped_sum_from_gputrace(self, in_file, out_file): logger.info("loading %s", in_file) df = self.pd.read_csv( in_file, usecols=["Start (ns)", "Duration (ns)", "Device", "Strm", "Name"] ) df["End (ns)"] = df["Start (ns)"] + df["Duration (ns)"] df = self.sum_non_overlapping_intervals(df) # get ready to print table with elapsed times per kernel df["Instances"] = 1 df_sum = df.groupby("Name", as_index=False).agg( {"Elapsed Time (ns)": "sum", "Duration (ns)": "sum", "Instances": "size"} ) # generate csv df_sum["Total Time (sec)"] = df_sum["Duration (ns)"] / 1e9 df_sum["Elapsed Time (sec)"] = df_sum["Elapsed Time (ns)"] / 1e9 df_sum = df_sum.sort_values(by="Elapsed Time (sec)", ascending=False) df_sum[["Elapsed Time (sec)", "Total Time (sec)", "Instances", "Name"]].to_csv( out_file, index=False ) def sum_non_overlapping_intervals(self, df): """ returns new sorted df with Elapsed Time (ns) column using vectorized operations """ logger.info("sorting %s trace records by start time", str(df.shape)) # Sort by start time and reset index df = df.sort_values(by="Start (ns)").reset_index(drop=True) # Initialize elapsed time as duration df["Elapsed Time (ns)"] = df["Duration (ns)"] # Get numpy arrays for faster operations starts = df["Start (ns)"].values ends = df["End (ns)"].values # Keep track of current interval end current_end = ends[0] display_units = max(1, int(len(df) / 100)) # Update current_end for overlapping intervals for i in range(1, len(df)): if i % display_units == 0: print(f"processing trace: {int(i/len(df) * 100)} %", end="\r") if starts[i] <= current_end: if ends[i] > current_end: # Partial overlap df.iloc[i, df.columns.get_loc("Elapsed Time (ns)")] = ( ends[i] - current_end ) current_end = ends[i] else: # Complete overlap df.iloc[i, df.columns.get_loc("Elapsed Time (ns)")] = 0 else: # No overlap current_end = ends[i] return df # functions for generating html files def make_html(self, df, output_dir, title): """make html graph from df""" import plotly.express as px if df.empty: return output_name = os.path.join(output_dir, "result") if not title: title = "Model_Engine" x = "Model_Engine" y = "Elapsed Time (sec)" color = "Category" """ generate kernel mapping table """ # Sort Model_Engine categories by last field after underscore df["Model_Engine"] = self.pd.Categorical( df["Model_Engine"], sorted(df["Model_Engine"].unique(), key=lambda x: x.split("_")[-1]), ) df[["Model_Engine", color, "Instances", "Name", y]].sort_values( by=color ).to_csv(f"{output_name}.csv", index=False) graph = px.histogram( df.round(2), x=x, y=y, title=(f"{y} for {title}"), color=color, text_auto=True, ) # wrap x axis labels graph.update_xaxes(automargin=True) graph.write_html(f"{output_name}.html") """ Generate data table with columns per Model_Engine into result.html """ pivot_df = df.pivot_table( values="Elapsed Time (sec)", index="Category", columns="Model_Engine", aggfunc="sum", observed=False, ).round(2) # Add sum row at bottom pivot_df.loc["total_elapsed_sec"] = pivot_df.sum() pivot_df.fillna("").to_html("temp.html") with ( open(f"{output_name}.html", "a", encoding="utf-8") as outfile, open("temp.html", encoding="utf-8") as infile, ): outfile.write(infile.read()) os.remove("temp.html") print( f"Finished generating: \n" f" {output_name}.html for stack bar chart \n" f" {output_name}.csv for Kernel-Category mapping" ) def anno_gpu_kernname(self, df, mapping): """add "Category" column""" def anno_gpu_kernname_helper(name): for kern_name, val in mapping.items(): if re.search(kern_name, name): return val df["Category"] = df["Name"].apply(anno_gpu_kernname_helper) def make_nongpu_row(self, df, nongpu_sec): """this will append non-gpu time entry at end of df""" nongpu_row = self.pd.DataFrame([df.iloc[-1]]) nongpu_row["Category"] = nongpu_row["Name"] = "CPU(non-GPU)" nongpu_row["Instances"] = 1 nongpu_row["Elapsed Time (sec)"] = nongpu_sec return nongpu_row def is_valid_file(self, base_file): """asserts if base_file is non-existent or is empty""" assert ( os.path.isfile(base_file) and os.path.getsize(base_file) > 0 ), f"{base_file} doesn't exist or is empty" def should_gen_file(self, new_file, base_file): """figure out if new file should be generated from base_file""" self.is_valid_file(base_file) if ( os.path.exists(new_file) and (os.path.getmtime(new_file) > os.path.getmtime(base_file)) and (os.path.getsize(base_file) > 0) ): logger.info("reusing %s", new_file) return False else: logger.info("generating %s", new_file) return True def gen_sum_file(self, file, nsys_cmd): """ generates sum file from nsys trace with times per kernel and returns the name of the sum file """ import subprocess file_dir = os.path.dirname(file) file_name = os.path.basename(file) if not file_dir: file_dir = "." # Walk through trace and get the total non-overlapped time nsys_stats_file = os.path.join(file_dir, f"{file_name}_cuda_gpu_trace.csv") sum_file = os.path.join(file_dir, f"{file_name}_cuda_gpu_kernel_tracesum.csv") if self.should_gen_file(nsys_stats_file, file): cmd = [ nsys_cmd, "stats", "-r", "cuda_gpu_trace", file, "-o", f"{file_dir}/{file_name}", ] cmd_str = shlex.join(cmd) logger.info("+ %s", cmd_str) # estimate time based on calibrated 240M/min file_size_mb = os.path.getsize(file) / 1e6 logger.info( "nsys stats for %.2f MB file expected to take %.2f min", file_size_mb, file_size_mb / 240, ) try: subprocess.run(cmd, check=True) except (FileNotFoundError, subprocess.CalledProcessError) as e: logger.error( "'%s' failed: %s. Use --nsys_cmd to specify nsys path", cmd_str, e ) exit(1) logger.info("generating non-overalapped sum %s", sum_file) self.gen_nonoverlapped_sum_from_gputrace(nsys_stats_file, sum_file) self.is_valid_file(sum_file) logger.info("Finished generating %s", sum_file) return sum_file def gen_graph(self, in_file, out_dir, title, nsys_cmd, engine_model): """generates graph and csv file from in_file into out_dir""" # Initialize an empty DataFrame to store combined data combined_df = self.pd.DataFrame() for idx, (file, engine, model, total_sec) in enumerate(in_file): file_dir = os.path.dirname(file) file_name = os.path.basename(file) if not file_dir: file_dir = "." sum_file = self.gen_sum_file(file, nsys_cmd) # read kernel summary file df = self.pd.read_csv(sum_file) # annotate kernel to their categories assert engine_model.get(engine), f"engine {engine} unknown" assert engine_model[engine].get(model), f"model {model} unknown" # remove nsys-rep from file_name for shorter x-label file_name = file_name.replace(".nsys-rep", "") df["Model_Engine"] = f"{model}_{engine}_{file_name}_{idx}" self.anno_gpu_kernname(df, engine_model[engine][model]) # patch in non-gpu time gpu_sec = round(df["Elapsed Time (sec)"].sum(), 1) total_sec = round(float(total_sec), 1) if total_sec < gpu_sec: logger.warning( "Elapsed sec %.2f < GPU sec %.2f resetting Elapsed sec ", total_sec, gpu_sec, ) total_sec = gpu_sec nongpu_row = self.make_nongpu_row(df, total_sec - gpu_sec) df = self.pd.concat([df, nongpu_row], ignore_index=True) combined_df = self.pd.concat([combined_df, df], ignore_index=True) if out_dir is None: out_dir = "." else: os.makedirs(out_dir, exist_ok=True) # generate html file self.make_html(combined_df, out_dir, title) def parse_tuple(s): return tuple(s.split(",")) def main(): logging.basicConfig( format=("%(asctime)s - %(levelname)s - %(message)s"), level=logging.INFO ) parser = argparse.ArgumentParser( description=( "Process nsys rep and generate kernel non-overlapped cycles. \n" "Example:\n" "gputrc2graph.py --in_file d1.nsys-rep,sglang,llama,100 \n" "d2.nsys-rep,sglang,gpt-oss,102 " '--out_dir results/ --title "Model=gpt-oss SGLANG chart"' ), formatter_class=argparse.RawDescriptionHelpFormatter, ) # load supported engine_model engine_model_supported = load_engine_model() # Get a string representation of supported engine/model combinations engine_model_supported_str = ", ".join( f"{engine}:[{', '.join(models.keys())}]" for engine, models in engine_model_supported.items() ) parser.add_argument( "--in_file", type=parse_tuple, nargs="+", help=( "list of (nsys-rep, engine, model, elapsed_nonprofiled_sec) " "separated by space. Elapsed_nonprofiled_sec is runtime without " "profiling used to calculate non-gpu time. Specify 0 to use " "elapsed time from nsys-rep but that might inflate non-gpu time. " f"Available engine:[model] are: {engine_model_supported_str} " f"Example: --infile d1.nsys-rep,sglan,llama,100 " "d2.nsys-rep,sglang,gpt-oss,102" ), required=True, ) parser.add_argument("--out_dir", help=("output dir for result.csv/html")) parser.add_argument("--title", help=("title for html chart")) parser.add_argument( "--nsys_cmd", help=("nsys cmd, e.g. /usr/bin/nsys, Default: nsys"), default="nsys", ) args = parser.parse_args() gputrace = GPUTrace2Graph() gputrace.gen_graph( args.in_file, args.out_dir, args.title, args.nsys_cmd, engine_model_supported ) if __name__ == "__main__": main() ================================================ FILE: examples/profiler/nsys_profile_tools/sglang_engine_model.json ================================================ { "sglang": { "llama": { "gemm|nvjet": "gemm", "fused_moe_kernel|GroupProblemShape|group_gemm_starts|bmm_|GemmUniversal": "moe_gemm", "moe|sigmoid": "moe", "CatArrayBatched|prepare_inputs": "prepare_next", "ncclDevKernel|cross_device_reduce": "nccl_and_custom_ar", "_norm_|Norm": "norm", "topk": "topk", "act_and_mul_": "activation", "Rotary": "rope", "SoftMax": "softmax", "flash|fmha": "attn", "elementwise": "elementwise", "fp8_quant|cvt_|quantize": "quantize", "reduce_kernel": "reduce", "triton": "triton_kernel", "CUDA mem": "non-gpu-H_D_memops", ".*": "misc" }, "ds": { "block_fp8_matmul": "block_fp8_gemm", "gemm|matmul|nvjet": "gemm", "fused_moe_kernel": "moe_gemm", "moe|expert|sigmoid": "moe", "CatArrayBatched|write_req_to": "prepare_next", "ncclDevKernel|cross_device_reduce|all_gather": "nccl_and_custom_ar", "Norm": "norm", "topk": "topk", "activation|act_and_mul": "activation", "compute_position_kernel": "rope", "elementwise": "elementwise", "fp8_quant|quant_fp8|quantize": "quantize", "SoftMax": "softmax", "reduce": "reduce", "_fwd_|create_flash|::mla::|KVCache": "attn", "CUDA mem": "non-gpu-H_D_memops", ".*": "misc" }, "gpt-oss": { "gemm|nvjet": "gemm", "fused_moe_kernel|_group_gemm|GroupProblemShape|GemmUniversal|bmm_|matmul_ogs_|_topk_forward|_combined_routing|_sum_bitmatrix_rows|_compute_writeback_idx": "moe_gemm", "moe|sigmoid": "moe", "CatArrayBatched|prepare_inputs": "prepare_next", "_norm_|Norm": "norm", "ncclDevKernel|cross_device_reduce|allreduce": "nccl_and_custom_ar", "topk|TopK": "topk", "act_and_mul_": "activation", "Rotary": "rope", "SoftMax": "softmax", "flash|fmha": "attn", "elementwise": "elementwise", "fp8_quant|cvt_|quantize": "quantize", "reduce_kernel": "reduce", "triton": "triton_kernel", "CUDA mem": "non-gpu-H_D_memops", ".*": "misc" } } } ================================================ FILE: examples/runtime/README.md ================================================ # Runtime examples The below examples will mostly need you to start a server in a separate terminal before you can execute them. Please see in the code for detailed instruction. ## Native API * `lora.py`: An example how to use LoRA adapters. * `multimodal_embedding.py`: An example how perform [multi modal embedding](Alibaba-NLP/gme-Qwen2-VL-2B-Instruct). * `openai_batch_chat.py`: An example how to process batch requests for chat completions. * `openai_batch_complete.py`: An example how to process batch requests for text completions. * **`openai_chat_with_response_prefill.py`**: An example that demonstrates how to [prefill a response](https://eugeneyan.com/writing/prompting/#prefill-claudes-responses) using the OpenAI API by enabling the `continue_final_message` parameter. When enabled, the final (partial) assistant message is removed and its content is used as a prefill so that the model continues that message rather than starting a new turn. See [Anthropic's prefill example](https://docs.anthropic.com/en/docs/build-with-claude/prompt-engineering/prefill-claudes-response#example-structured-data-extraction-with-prefilling) for more context. * `reward_model.py`: An example how to extract scores from a reward model. * `vertex_predict.py`: An example how to deploy a model to [Vertex AI](https://cloud.google.com/vertex-ai?hl=en). ## Engine The `engine` folder contains that examples that show how to use [Offline Engine API](https://docs.sglang.io/basic_usage/offline_engine_api.html#Offline-Engine-API) for common workflows. * `custom_server.py`: An example how to deploy a custom server. * `embedding.py`: An example how to extract embeddings. * `launch_engine.py`: An example how to launch the Engine. * `offline_batch_inference_eagle.py`: An example how to perform speculative decoding using [EAGLE](https://docs.sglang.io/advanced_features/speculative_decoding.html). * `offline_batch_inference_torchrun.py`: An example how to perform inference using [torchrun](https://pytorch.org/docs/stable/elastic/run.html). * `offline_batch_inference_vlm.py`: An example how to use VLMs with the engine. * `offline_batch_inference.py`: An example how to use the engine to perform inference on a batch of examples. ## Hidden States The `hidden_states` folder contains examples on how to extract hidden states using SGLang. Please note that this might degrade throughput due to cuda graph rebuilding. * `hidden_states_engine.py`: An example how to extract hidden states using the Engine API. * `hidden_states_server.py`: An example how to extract hidden states using the Server API. ## Multimodal SGLang supports multimodal inputs for various model architectures. The `multimodal` folder contains examples showing how to use urls, files or encoded data to make requests to multimodal models. Examples include querying the [Llava-OneVision](multimodal/llava_onevision_server.py) model (image, multi-image, video), Llava-backed [Qwen-Llava](multimodal/qwen_llava_server.py) and [Llama3-Llava](multimodal/llama3_llava_server.py) models (image, multi-image), and Mistral AI's [Pixtral](multimodal/pixtral_server.py) (image, multi-image). ## Token In, Token Out The folder `token_in_token_out` shows how to perform inference, where we provide tokens and get tokens as response. * `token_in_token_out_{llm|vlm}_{engine|server}.py`: Shows how to perform token in, token out workflow for llm/vlm using either the engine or native API. ================================================ FILE: examples/runtime/engine/custom_server.py ================================================ from sanic import Sanic, text from sanic.response import json import sglang as sgl engine = None # Create an instance of the Sanic app app = Sanic("sanic-server") # Define an asynchronous route handler @app.route("/generate", methods=["POST"]) async def generate(request): prompt = request.json.get("prompt") if not prompt: return json({"error": "Prompt is required"}, status=400) # async_generate returns a dict result = await engine.async_generate(prompt) return text(result["text"]) @app.route("/generate_stream", methods=["POST"]) async def generate_stream(request): prompt = request.json.get("prompt") if not prompt: return json({"error": "Prompt is required"}, status=400) # async_generate returns a dict result = await engine.async_generate(prompt, stream=True) # https://sanic.dev/en/guide/advanced/streaming.md#streaming # init the response response = await request.respond() # result is an async generator async for chunk in result: await response.send(chunk["text"]) await response.eof() def run_server(): global engine engine = sgl.Engine(model_path="meta-llama/Meta-Llama-3.1-8B-Instruct") app.run(host="0.0.0.0", port=8000, single_process=True) if __name__ == "__main__": run_server() ================================================ FILE: examples/runtime/engine/embedding.py ================================================ import sglang as sgl def main(): # Sample prompts. prompts = [ "Hello, my name is", "The president of the United States is", "The capital of France is", "The future of AI is", ] # Create an LLM. llm = sgl.Engine( model_path="Alibaba-NLP/gte-Qwen2-1.5B-instruct", is_embedding=True ) outputs = llm.encode(prompts) # Print the outputs (embedding vectors) for prompt, output in zip(prompts, outputs): print("===============================") print(f"Prompt: {prompt}\nEmbedding vector: {output['embedding']}") # The __main__ condition is necessary here because we use "spawn" to create subprocesses # Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine if __name__ == "__main__": main() ================================================ FILE: examples/runtime/engine/fastapi_engine_inference.py ================================================ """ FastAPI server example for text generation using SGLang Engine and demonstrating client usage. Starts the server, sends requests to it, and prints responses. Usage: python fastapi_engine_inference.py --model-path Qwen/Qwen2.5-0.5B-Instruct --tp_size 1 --host 127.0.0.1 --port 8000 [--startup-timeout 60] """ import os import subprocess import time from contextlib import asynccontextmanager import requests from fastapi import FastAPI, Request import sglang as sgl from sglang.utils import terminate_process engine = None # Use FastAPI's lifespan manager to initialize/shutdown the engine @asynccontextmanager async def lifespan(app: FastAPI): """Manages SGLang engine initialization during server startup.""" global engine # Initialize the SGLang engine when the server starts # Adjust model_path and other engine arguments as needed print("Loading SGLang engine...") engine = sgl.Engine( model_path=os.getenv("MODEL_PATH"), tp_size=int(os.getenv("TP_SIZE")) ) print("SGLang engine loaded.") yield # Clean up engine resources when the server stops (optional, depends on engine needs) print("Shutting down SGLang engine...") # engine.shutdown() # Or other cleanup if available/necessary print("SGLang engine shutdown.") app = FastAPI(lifespan=lifespan) @app.post("/generate") async def generate_text(request: Request): """FastAPI endpoint to handle text generation requests.""" global engine if not engine: return {"error": "Engine not initialized"}, 503 try: data = await request.json() prompt = data.get("prompt") max_new_tokens = data.get("max_new_tokens", 128) temperature = data.get("temperature", 0.7) if not prompt: return {"error": "Prompt is required"}, 400 # Use async_generate for non-blocking generation state = await engine.async_generate( prompt, sampling_params={ "max_new_tokens": max_new_tokens, "temperature": temperature, }, # Add other parameters like stop, top_p etc. as needed ) return {"generated_text": state["text"]} except Exception as e: return {"error": str(e)}, 500 # Helper function to start the server def start_server(args, timeout=60): """Starts the Uvicorn server as a subprocess and waits for it to be ready.""" base_url = f"http://{args.host}:{args.port}" command = [ "python", "-m", "uvicorn", "fastapi_engine_inference:app", f"--host={args.host}", f"--port={args.port}", ] process = subprocess.Popen(command, stdout=None, stderr=None) start_time = time.perf_counter() with requests.Session() as session: while time.perf_counter() - start_time < timeout: try: # Check the /docs endpoint which FastAPI provides by default response = session.get( f"{base_url}/docs", timeout=5 ) # Add a request timeout if response.status_code == 200: print(f"Server {base_url} is ready (responded on /docs)") return process except requests.ConnectionError: # Specific exception for connection refused/DNS error etc. pass except requests.Timeout: # Specific exception for request timeout print(f"Health check to {base_url}/docs timed out, retrying...") pass except requests.RequestException as e: # Catch other request exceptions print(f"Health check request error: {e}, retrying...") pass # Use a shorter sleep interval for faster startup detection time.sleep(1) # If loop finishes, raise the timeout error # Attempt to terminate the failed process before raising if process: print( "Server failed to start within timeout, attempting to terminate process..." ) terminate_process(process) # Use the imported terminate_process raise TimeoutError( f"Server failed to start at {base_url} within the timeout period." ) def send_requests(server_url, prompts, max_new_tokens, temperature): """Sends generation requests to the running server for a list of prompts.""" # Iterate through prompts and send requests for i, prompt in enumerate(prompts): print(f"\n[{i+1}/{len(prompts)}] Sending prompt: '{prompt}'") payload = { "prompt": prompt, "max_new_tokens": max_new_tokens, "temperature": temperature, } try: response = requests.post(f"{server_url}/generate", json=payload, timeout=60) result = response.json() print(f"Prompt: {prompt}\nResponse: {result['generated_text']}") except requests.exceptions.Timeout: print(f" Error: Request timed out for prompt '{prompt}'") except requests.exceptions.RequestException as e: print(f" Error sending request for prompt '{prompt}': {e}") if __name__ == "__main__": """Main entry point for the script.""" import argparse parser = argparse.ArgumentParser() parser.add_argument("--host", type=str, default="127.0.0.1") parser.add_argument("--port", type=int, default=8000) parser.add_argument("--model-path", type=str, default="Qwen/Qwen2.5-0.5B-Instruct") parser.add_argument("--tp_size", type=int, default=1) parser.add_argument( "--startup-timeout", type=int, default=60, help="Time in seconds to wait for the server to be ready (default: %(default)s)", ) args = parser.parse_args() # Pass the model to the child uvicorn process via an env var os.environ["MODEL_PATH"] = args.model_path os.environ["TP_SIZE"] = str(args.tp_size) # Start the server process = start_server(args, timeout=args.startup_timeout) # Define the prompts and sampling parameters prompts = [ "Hello, my name is", "The president of the United States is", "The capital of France is", "The future of AI is", ] max_new_tokens = 64 temperature = 0.1 # Define server url server_url = f"http://{args.host}:{args.port}" # Send requests to the server send_requests(server_url, prompts, max_new_tokens, temperature) # Terminate the server process terminate_process(process) ================================================ FILE: examples/runtime/engine/launch_engine.py ================================================ """ This example demonstrates how to launch the offline engine. """ import sglang as sgl def main(): llm = sgl.Engine(model_path="meta-llama/Meta-Llama-3.1-8B-Instruct") llm.generate("What is the capital of France?") llm.shutdown() # The __main__ condition is necessary here because we use "spawn" to create subprocesses # Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine if __name__ == "__main__": main() ================================================ FILE: examples/runtime/engine/offline_batch_inference.py ================================================ """ Usage: python3 offline_batch_inference.py --model meta-llama/Llama-3.1-8B-Instruct """ import argparse import dataclasses import sglang as sgl from sglang.srt.server_args import ServerArgs def main( server_args: ServerArgs, ): # Sample prompts. prompts = [ "Hello, my name is", "The president of the United States is", "The capital of France is", "The future of AI is", ] # Create a sampling params object. sampling_params = {"temperature": 0.8, "top_p": 0.95} # Create an LLM. llm = sgl.Engine(**dataclasses.asdict(server_args)) outputs = llm.generate(prompts, sampling_params) # Print the outputs. for prompt, output in zip(prompts, outputs): print("===============================") print(f"Prompt: {prompt}\nGenerated text: {output['text']}") # The __main__ condition is necessary here because we use "spawn" to create subprocesses # Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine if __name__ == "__main__": parser = argparse.ArgumentParser() ServerArgs.add_cli_args(parser) args = parser.parse_args() server_args = ServerArgs.from_cli_args(args) main(server_args) ================================================ FILE: examples/runtime/engine/offline_batch_inference_async.py ================================================ """ Usage: python offline_batch_inference_async.py --model-path Qwen/Qwen2-VL-7B-Instruct Note: This demo shows the usage of async generation, which is useful to implement an online-like generation with batched inference. """ import argparse import asyncio import dataclasses import time import sglang as sgl from sglang.srt.server_args import ServerArgs class InferenceEngine: def __init__(self, **kwargs): self.engine = sgl.Engine(**kwargs) async def generate(self, prompt, sampling_params): result = await self.engine.async_generate(prompt, sampling_params) return result async def run_server(server_args): inference = InferenceEngine(**dataclasses.asdict(server_args)) # Sample prompts. prompts = [ "Hello, my name is", "The president of the United States is", "The capital of France is", "The future of AI is", ] * 100 # Create a sampling params object. sampling_params = {"temperature": 0.8, "top_p": 0.95} # Run the generation tasks concurrently in async mode. tasks = [] for prompt in prompts: task = asyncio.create_task(inference.generate(prompt, sampling_params)) tasks.append(task) # Get and print the result for task in tasks: await task while True: if not task.done(): time.sleep(1) else: result = task.result() print(f"Generated text: {result['text']}") break if __name__ == "__main__": parser = argparse.ArgumentParser() ServerArgs.add_cli_args(parser) args = parser.parse_args() server_args = ServerArgs.from_cli_args(args) asyncio.run(run_server(server_args)) ================================================ FILE: examples/runtime/engine/offline_batch_inference_eagle.py ================================================ import sglang as sgl def main(): # Sample prompts. prompts = [ "Hello, my name is", "The president of the United States is", "The capital of France is", "The future of AI is", ] # Create a sampling params object. sampling_params = {"temperature": 0, "max_new_tokens": 30} # Create an LLM. llm = sgl.Engine( model_path="meta-llama/Llama-2-7b-chat-hf", speculative_algorithm="EAGLE", speculative_draft_model_path="lmsys/sglang-EAGLE-llama2-chat-7B", speculative_num_steps=3, speculative_eagle_topk=4, speculative_num_draft_tokens=16, cuda_graph_max_bs=8, ) outputs = llm.generate(prompts, sampling_params) # Print the outputs. for prompt, output in zip(prompts, outputs): print("===============================") print(f"Prompt: {prompt}\nGenerated text: {output['text']}") # The __main__ condition is necessary here because we use "spawn" to create subprocesses # Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine if __name__ == "__main__": main() ================================================ FILE: examples/runtime/engine/offline_batch_inference_qwen_1m.py ================================================ """ Usage: python3 offline_batch_inference.py """ from urllib.request import urlopen import sglang as sgl def load_prompt() -> str: # Test cases with various lengths can be found at: # # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/64k.txt # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/200k.txt # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/600k.txt # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/1m.txt with urlopen( "https://qianwen-res.oss-cn-beijing.aliyuncs.com" "/Qwen2.5-1M/test-data/64k.txt", timeout=5, ) as response: prompt = response.read().decode("utf-8") return prompt # Processing the prompt. def process_requests(llm: sgl.Engine, prompts: list[str]) -> None: # Create a sampling params object. sampling_params = { "temperature": 0.7, "top_p": 0.8, "top_k": 20, "repetition_penalty": 1.05, "max_new_tokens": 256, } # Generate texts from the prompts. outputs = llm.generate(prompts, sampling_params) # Print the outputs. for output in outputs: prompt_token_ids = output["meta_info"]["prompt_tokens"] generated_text = output["text"] print( f"Prompt length: {prompt_token_ids}, " f"Generated text: {generated_text!r}" ) # Create an LLM. def initialize_engine() -> sgl.Engine: llm = sgl.Engine( model_path="Qwen/Qwen2.5-7B-Instruct-1M", context_length=1048576, page_size=256, attention_backend="dual_chunk_flash_attn", tp_size=4, disable_radix_cache=True, enable_mixed_chunk=False, enable_torch_compile=False, chunked_prefill_size=131072, mem_fraction_static=0.6, log_level="DEBUG", ) return llm def main(): llm = initialize_engine() prompt = load_prompt() process_requests(llm, [prompt]) if __name__ == "__main__": main() ================================================ FILE: examples/runtime/engine/offline_batch_inference_vlm.py ================================================ """ Usage: python offline_batch_inference_vlm.py --model-path Qwen/Qwen2-VL-7B-Instruct """ import argparse import dataclasses import sglang as sgl from sglang.srt.parser.conversation import chat_templates from sglang.srt.server_args import ServerArgs def main( server_args: ServerArgs, ): vlm = sgl.Engine(**dataclasses.asdict(server_args)) conv = chat_templates[server_args.chat_template].copy() image_token = conv.image_token image_url = "https://github.com/sgl-project/sglang/blob/main/examples/assets/example_image.png?raw=true" prompt = f"What's in this image?\n{image_token}" sampling_params = { "temperature": 0.001, "max_new_tokens": 30, } output = vlm.generate( prompt=prompt, image_data=image_url, sampling_params=sampling_params, ) print("===============================") print(f"Prompt: {prompt}") print(f"Generated text: {output['text']}") vlm.shutdown() # The __main__ condition is necessary here because we use "spawn" to create subprocesses # Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine if __name__ == "__main__": parser = argparse.ArgumentParser() ServerArgs.add_cli_args(parser) args = parser.parse_args() server_args = ServerArgs.from_cli_args(args) main(server_args) ================================================ FILE: examples/runtime/engine/readme.md ================================================ # SGLang Engine SGLang provides a direct inference engine without the need for an HTTP server. There are generally these use cases: - [Offline Batch Inference](#offline-batch-inference) - [Embedding Generation](#embedding-generation) - [Custom Server](#custom-server) - [Token-In-Token-Out for RLHF](#token-in-token-out-for-rlhf) - [Inference Using FastAPI](#inference-using-fastapi) ## Examples ### [Offline Batch Inference](./offline_batch_inference.py) In this example, we launch an SGLang engine and feed a batch of inputs for inference. If you provide a very large batch, the engine will intelligently schedule the requests to process efficiently and prevent OOM (Out of Memory) errors. ### [Embedding Generation](./embedding.py) In this example, we launch an SGLang engine and feed a batch of inputs for embedding generation. ### [Custom Server](./custom_server.py) This example demonstrates how to create a custom server on top of the SGLang Engine. We use [Sanic](https://sanic.dev/en/) as an example. The server supports both non-streaming and streaming endpoints. #### Steps 1. Install Sanic: ```bash pip install sanic ``` 2. Run the server: ```bash python custom_server ``` 3. Send requests: ```bash curl -X POST http://localhost:8000/generate -H "Content-Type: application/json" -d '{"prompt": "The Transformer architecture is..."}' curl -X POST http://localhost:8000/generate_stream -H "Content-Type: application/json" -d '{"prompt": "The Transformer architecture is..."}' --no-buffer ``` This will send both non-streaming and streaming requests to the server. ### [Token-In-Token-Out for RLHF](../token_in_token_out) In this example, we launch an SGLang engine, feed tokens as input and generate tokens as output. ### [Inference Using FastAPI](fastapi_engine_inference.py) This example demonstrates how to create a FastAPI server that uses the SGLang engine for text generation. ================================================ FILE: examples/runtime/engine/save_remote_state.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Saves each worker's model state dict directly to a checkpoint, which enables a fast load path for large tensor-parallel models where each worker only needs to read its own shard rather than the entire checkpoint. Example usage: python save_remote_state.py \ --model-path /path/to/load \ --tensor-parallel-size 8 \ --remote-model-save-url [protocol]://[host]:[port]/[model_name] \ Then, the model can be loaded with llm = Engine( model_path="[protocol]://[host]:[port]/[model_name]", tensor_parallel_size=8, ) """ import dataclasses from argparse import ArgumentParser from pathlib import Path from sglang import Engine, ServerArgs parser = ArgumentParser() ServerArgs.add_cli_args(parser) parser.add_argument( "--remote-model-save-url", required=True, type=str, help="remote address to store model weights", ) parser.add_argument( "--remote-draft-model-save-url", default=None, type=str, help="remote address to store draft model weights", ) def main(args): engine_args = ServerArgs.from_cli_args(args) model_path = engine_args.model_path if not Path(model_path).is_dir(): raise ValueError("model path must be a local directory") # Create LLM instance from arguments llm = Engine(**dataclasses.asdict(engine_args)) llm.save_remote_model( url=args.remote_model_save_url, draft_url=args.remote_draft_model_save_url ) print("save remote (draft) model successfully") if __name__ == "__main__": args = parser.parse_args() main(args) ================================================ FILE: examples/runtime/engine/save_sharded_state.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Saves each worker's model state dict directly to a checkpoint, which enables a fast load path for large tensor-parallel models where each worker only needs to read its own shard rather than the entire checkpoint. Example usage: python save_sharded_state.py \ --model-path /path/to/load \ --quantization deepspeedfp \ --tensor-parallel-size 8 \ --output /path/to/save Then, the model can be loaded with llm = Engine( model_path="/path/to/save", load_format="sharded_state", quantization="deepspeedfp", tensor_parallel_size=8, ) """ import dataclasses import os import shutil from argparse import ArgumentParser from pathlib import Path from sglang import Engine, ServerArgs parser = ArgumentParser() ServerArgs.add_cli_args(parser) parser.add_argument( "--output", "-o", required=True, type=str, help="path to output checkpoint" ) parser.add_argument( "--file-pattern", type=str, help="string pattern of saved filenames" ) parser.add_argument( "--max-file-size", type=str, default=5 * 1024**3, help="max size (in bytes) of each safetensors file", ) def main(args): engine_args = ServerArgs.from_cli_args(args) model_path = engine_args.model_path if not Path(model_path).is_dir(): raise ValueError("model path must be a local directory") # Create LLM instance from arguments llm = Engine(**dataclasses.asdict(engine_args)) Path(args.output).mkdir(exist_ok=True) llm.save_sharded_model( path=args.output, pattern=args.file_pattern, max_size=args.max_file_size ) # Copy metadata files to output directory for file in os.listdir(model_path): if os.path.splitext(file)[1] not in (".bin", ".pt", ".safetensors"): if os.path.isdir(os.path.join(model_path, file)): shutil.copytree( os.path.join(model_path, file), os.path.join(args.output, file) ) else: shutil.copy(os.path.join(model_path, file), args.output) if __name__ == "__main__": args = parser.parse_args() main(args) ================================================ FILE: examples/runtime/hidden_states/hidden_states_engine.py ================================================ """ Usage: python hidden_states.py Note that each time you change the `return_hidden_states` parameter, the cuda graph will be recaptured, which might lead to a performance hit. So avoid getting hidden states and completions alternately. """ import torch import sglang as sgl def main(): prompts = [ "Hello, my name is", "The president of the United States is", "The capital of France is", "The future of AI is", ] # Create an LLM. llm = sgl.Engine( model_path="Alibaba-NLP/gte-Qwen2-1.5B-instruct", enable_return_hidden_states=True, ) sampling_params = { "temperature": 0.8, "top_p": 0.95, "max_new_tokens": 10, } outputs = llm.generate( prompts, sampling_params=sampling_params, return_hidden_states=True ) llm.shutdown() for prompt, output in zip(prompts, outputs): for i in range(len(output["meta_info"]["hidden_states"])): output["meta_info"]["hidden_states"][i] = torch.tensor( output["meta_info"]["hidden_states"][i], dtype=torch.bfloat16 ) print("===============================") print( f"Prompt: {prompt}\n" f"Generated text: {output['text']}\n" f"Prompt_Tokens: {output['meta_info']['prompt_tokens']}\t" f"Completion_tokens: {output['meta_info']['completion_tokens']}" ) print("Hidden states: ") hidden_states = torch.cat( [ i.unsqueeze(0) if len(i.shape) == 1 else i for i in output["meta_info"]["hidden_states"] ] ) print(hidden_states) print() # The __main__ condition is necessary here because we use "spawn" to create subprocesses # Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine if __name__ == "__main__": main() ================================================ FILE: examples/runtime/hidden_states/hidden_states_server.py ================================================ """ Usage: python hidden_states_server.py Note that each time you change the `return_hidden_states` parameter, the cuda graph will be recaptured, which might lead to a performance hit. So avoid getting hidden states and completions alternately. """ import requests import torch from sglang.test.test_utils import is_in_ci from sglang.utils import terminate_process, wait_for_server if is_in_ci(): from docs.backend.patch import launch_server_cmd else: from sglang.utils import launch_server_cmd def main(): # Launch the server server_process, port = launch_server_cmd( "python -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-1.5B-instruct --enable-return-hidden-states --host 0.0.0.0" ) wait_for_server(f"http://localhost:{port}", process=server_process) prompts = [ "Hello, my name is", "The president of the United States is", "The capital of France is", "The future of AI is", ] sampling_params = { "temperature": 0.8, "top_p": 0.95, "max_new_tokens": 10, } json_data = { "text": prompts, "sampling_params": sampling_params, "return_hidden_states": True, } response = requests.post( f"http://localhost:{port}/generate", json=json_data, ) terminate_process(server_process) outputs = response.json() for prompt, output in zip(prompts, outputs): for i in range(len(output["meta_info"]["hidden_states"])): output["meta_info"]["hidden_states"][i] = torch.tensor( output["meta_info"]["hidden_states"][i], dtype=torch.bfloat16 ) print("===============================") print( f"Prompt: {prompt}\n" f"Generated text: {output['text']}\n" f"Prompt_Tokens: {output['meta_info']['prompt_tokens']}\t" f"Completion_tokens: {output['meta_info']['completion_tokens']}" ) print("Hidden states: ") hidden_states = torch.cat( [ i.unsqueeze(0) if len(i.shape) == 1 else i for i in output["meta_info"]["hidden_states"] ] ) print(hidden_states) print() if __name__ == "__main__": main() ================================================ FILE: examples/runtime/lora.py ================================================ """ OpenAI-compatible LoRA adapter usage with SGLang. Server Setup: python -m sglang.launch_server \\ --model meta-llama/Llama-3.1-8B-Instruct \\ --enable-lora \\ --lora-paths sql=/path/to/sql python=/path/to/python """ import openai client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="EMPTY") def main(): print("SGLang OpenAI-Compatible LoRA Examples\n") # Example 1: NEW - Adapter in model parameter (OpenAI-compatible) print("1. Chat with LoRA adapter in model parameter:") response = client.chat.completions.create( model="meta-llama/Llama-3.1-8B-Instruct:sql", # ← adapter:name syntax messages=[{"role": "user", "content": "Convert to SQL: show all users"}], max_tokens=50, ) print(f" Response: {response.choices[0].message.content}\n") # Example 2: Completions API with adapter print("2. Completion with LoRA adapter:") response = client.completions.create( model="meta-llama/Llama-3.1-8B-Instruct:python", prompt="def fibonacci(n):", max_tokens=50, ) print(f" Response: {response.choices[0].text}\n") # Example 3: OLD - Backward compatible with explicit lora_path print("3. Backward compatible (explicit lora_path):") response = client.chat.completions.create( model="meta-llama/Llama-3.1-8B-Instruct", messages=[{"role": "user", "content": "Convert to SQL: show all users"}], extra_body={"lora_path": "sql"}, max_tokens=50, ) print(f" Response: {response.choices[0].message.content}\n") # Example 4: Base model (no adapter) print("4. Base model without adapter:") response = client.chat.completions.create( model="meta-llama/Llama-3.1-8B-Instruct", messages=[{"role": "user", "content": "Hello!"}], max_tokens=30, ) print(f" Response: {response.choices[0].message.content}\n") print("All examples completed!") if __name__ == "__main__": try: main() except Exception as e: print(f"Error: {e}") print( "\nEnsure server is running:\n" " python -m sglang.launch_server --model ... --enable-lora --lora-paths ..." ) ================================================ FILE: examples/runtime/multimodal/llama3_llava_server.py ================================================ """ Usage: # Installing latest llava-next: pip install git+https://github.com/LLaVA-VL/LLaVA-NeXT.git # Installing latest sglang. # Endpoint Service CLI: python -m sglang.launch_server --model-path lmms-lab/llama3-llava-next-8b --port=30000 python3 llama3_llava_server.py Output: "Friends posing for a fun photo with a life-sized teddy bear, creating a playful and memorable moment." """ import argparse import asyncio import copy import json import aiohttp import requests from llava.conversation import conv_llava_llama_3 from sglang.utils import normalize_base_url async def send_request(url, data, delay=0): await asyncio.sleep(delay) async with aiohttp.ClientSession() as session: async with session.post(url, json=data) as resp: output = await resp.json() return output async def test_concurrent(args): url = normalize_base_url(args.host, args.port) prompt = "\nPlease generate caption towards this image." conv_template = copy.deepcopy(conv_llava_llama_3) conv_template.append_message(role=conv_template.roles[0], message=prompt) conv_template.append_message(role=conv_template.roles[1], message=None) prompt_with_template = conv_template.get_prompt() response = [] for i in range(1): response.append( send_request( url + "/generate", { "text": prompt_with_template, "image_data": "https://farm4.staticflickr.com/3175/2653711032_804ff86d81_z.jpg", "sampling_params": { "max_new_tokens": 1024, "temperature": 0, "top_p": 1.0, "presence_penalty": 2, "frequency_penalty": 2, "stop": "<|eot_id|>", }, }, ) ) rets = await asyncio.gather(*response) for ret in rets: print(ret["text"]) def test_streaming(args): url = normalize_base_url(args.host, args.port) prompt = "\nPlease generate caption towards this image." conv_template = copy.deepcopy(conv_llava_llama_3) conv_template.append_message(role=conv_template.roles[0], message=prompt) conv_template.append_message(role=conv_template.roles[1], message=None) prompt_with_template = conv_template.get_prompt() pload = { "text": prompt_with_template, "sampling_params": { "max_new_tokens": 1024, "temperature": 0, "top_p": 1.0, "presence_penalty": 2, "frequency_penalty": 2, "stop": "<|eot_id|>", }, "image_data": "https://farm4.staticflickr.com/3175/2653711032_804ff86d81_z.jpg", "stream": True, } response = requests.post( url + "/generate", json=pload, stream=True, ) prev = 0 for chunk in response.iter_lines(decode_unicode=False): chunk = chunk.decode("utf-8") if chunk and chunk.startswith("data:"): if chunk == "data: [DONE]": break data = json.loads(chunk[5:].strip("\n")) output = data["text"].strip() print(output[prev:], end="", flush=True) prev = len(output) print("") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--host", type=str, default="127.0.0.1") parser.add_argument("--port", type=int, default=30000) args = parser.parse_args() asyncio.run(test_concurrent(args)) test_streaming(args) ================================================ FILE: examples/runtime/multimodal/llava_onevision_server.py ================================================ """ Usage: python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-72b-ov --port=30000 --tp-size=8 python3 llava_onevision_server.py """ import io import os import sys import time import numpy as np import openai import pybase64 import requests from PIL import Image from sglang.srt.utils.video_decoder import VideoDecoderWrapper # pip install httpx==0.23.3 # pip install torchcodec # pip install protobuf==3.20.0 def download_video(url, cache_dir): file_path = os.path.join(cache_dir, "jobs.mp4") os.makedirs(cache_dir, exist_ok=True) response = requests.get(url) response.raise_for_status() with open(file_path, "wb") as f: f.write(response.content) print(f"File downloaded and saved to: {file_path}") return file_path def create_openai_client(base_url): return openai.Client(api_key="EMPTY", base_url=base_url) def image_stream_request_test(client): print("----------------------Image Stream Request Test----------------------") stream_request = client.chat.completions.create( model="default", messages=[ { "role": "user", "content": [ { "type": "image_url", "image_url": { "url": "https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png" }, }, { "type": "text", "text": "Please describe this image. Please list the benchmarks and the models.", }, ], }, ], temperature=0.7, max_tokens=1024, stream=True, ) stream_response = "" for chunk in stream_request: if chunk.choices[0].delta.content is not None: content = chunk.choices[0].delta.content stream_response += content sys.stdout.write(content) sys.stdout.flush() print("-" * 30) def multi_image_stream_request_test(client): print( "----------------------Multi-Images Stream Request Test----------------------" ) stream_request = client.chat.completions.create( model="default", messages=[ { "role": "user", "content": [ { "type": "image_url", "image_url": { "url": "https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png" }, "modalities": "multi-images", }, { "type": "image_url", "image_url": { "url": "https://raw.githubusercontent.com/sgl-project/sglang/main/examples/assets/example_image.png" }, "modalities": "multi-images", }, { "type": "text", "text": "I have shown you two images. Please describe the two images to me.", }, ], }, ], temperature=0.7, max_tokens=1024, stream=True, ) stream_response = "" for chunk in stream_request: if chunk.choices[0].delta.content is not None: content = chunk.choices[0].delta.content stream_response += content sys.stdout.write(content) sys.stdout.flush() print("-" * 30) def video_stream_request_test(client, video_path): print("------------------------Video Stream Request Test----------------------") messages = prepare_video_messages(video_path) video_request = client.chat.completions.create( model="default", messages=messages, temperature=0, max_tokens=1024, stream=True, ) print("-" * 30) video_response = "" for chunk in video_request: if chunk.choices[0].delta.content is not None: content = chunk.choices[0].delta.content video_response += content sys.stdout.write(content) sys.stdout.flush() print("-" * 30) def image_speed_test(client): print("----------------------Image Speed Test----------------------") start_time = time.perf_counter() request = client.chat.completions.create( model="default", messages=[ { "role": "user", "content": [ { "type": "image_url", "image_url": { "url": "https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png" }, }, { "type": "text", "text": "Please describe this image. Please list the benchmarks and the models.", }, ], }, ], temperature=0, max_tokens=1024, ) end_time = time.perf_counter() response = request.choices[0].message.content print(response) print("-" * 30) print_speed_test_results(request, start_time, end_time) def video_speed_test(client, video_path): print("------------------------Video Speed Test------------------------") messages = prepare_video_messages(video_path) start_time = time.perf_counter() video_request = client.chat.completions.create( model="default", messages=messages, temperature=0, max_tokens=1024, ) end_time = time.perf_counter() video_response = video_request.choices[0].message.content print(video_response) print("-" * 30) print_speed_test_results(video_request, start_time, end_time) def prepare_video_messages(video_path): max_frames_num = 32 decoder = VideoDecoderWrapper(video_path) total_frame_num = len(decoder) uniform_sampled_frames = np.linspace( 0, total_frame_num - 1, max_frames_num, dtype=int ) frame_idx = uniform_sampled_frames.tolist() frames = decoder.get_frames_at(frame_idx) base64_frames = [] for frame in frames: pil_img = Image.fromarray(frame) buff = io.BytesIO() pil_img.save(buff, format="JPEG") base64_str = pybase64.b64encode(buff.getvalue()).decode("utf-8") base64_frames.append(base64_str) messages = [{"role": "user", "content": []}] for base64_frame in base64_frames: frame_format = { "type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_frame}"}, "modalities": "video", } messages[0]["content"].append(frame_format) prompt = {"type": "text", "text": "Please describe the video in detail."} messages[0]["content"].append(prompt) return messages def print_speed_test_results(request, start_time, end_time): total_tokens = request.usage.total_tokens completion_tokens = request.usage.completion_tokens prompt_tokens = request.usage.prompt_tokens print(f"Total tokens: {total_tokens}") print(f"Completion tokens: {completion_tokens}") print(f"Prompt tokens: {prompt_tokens}") print(f"Time taken: {end_time - start_time} seconds") print(f"Token per second: {total_tokens / (end_time - start_time)}") print(f"Completion token per second: {completion_tokens / (end_time - start_time)}") print(f"Prompt token per second: {prompt_tokens / (end_time - start_time)}") def main(): url = "https://raw.githubusercontent.com/EvolvingLMMs-Lab/sglang/dev/onevision_local/assets/jobs.mp4" cache_dir = os.path.expanduser("~/.cache") video_path = download_video(url, cache_dir) client = create_openai_client("http://127.0.0.1:30000/v1") image_stream_request_test(client) multi_image_stream_request_test(client) video_stream_request_test(client, video_path) image_speed_test(client) video_speed_test(client, video_path) if __name__ == "__main__": main() ================================================ FILE: examples/runtime/multimodal/pixtral_server.py ================================================ """ Usage: # Run a Pixtral model with SGLang: # HuggingFace: python -m sglang.launch_server --model-path mistral-community/pixtral-12b --port=30000 # ModelScope: python -m sglang.launch_server --model-path AI-ModelScope/pixtral-12b --port=30000 # Then test it with: python pixtral_server.py This script tests Pixtral model with both single and multiple images. """ import argparse import asyncio import json import aiohttp import requests from sglang.utils import normalize_base_url IMAGE_TOKEN_SEP = "\n[IMG]" ROUTE = "/generate" async def send_request(url, data, delay=0): await asyncio.sleep(delay) async with aiohttp.ClientSession() as session: async with session.post(url, json=data) as resp: output = await resp.json() return output async def test_concurrent(args): url = f"{normalize_base_url(args.host, args.port)}{ROUTE}" # Single image test if args.single_image: prompt = f"[INST]Describe this image in detail.{IMAGE_TOKEN_SEP}[/INST]" image_url = "https://picsum.photos/id/237/400/300" modality = ["image"] # Multiple images test else: image_urls = [ "https://picsum.photos/id/237/400/300", "https://picsum.photos/id/27/500/500", ] prompt = f"[INST]How many photos are there? Describe each in a very short sentence.{IMAGE_TOKEN_SEP * len(image_urls)}[/INST]" image_url = image_urls modality = ["multi-images"] response = await send_request( url, { "text": prompt, "image_data": image_url, "sampling_params": { "max_new_tokens": 100, "temperature": 0.7, "top_p": 0.9, }, "modalities": modality, }, ) print(f"Response: {response}") if "text" in response: print("\nOutput text:", response["text"]) def test_streaming(args): url = f"{normalize_base_url(args.host, args.port)}/generate" # Single image test if args.single_image: prompt = f"[INST]Describe this image in detail.{IMAGE_TOKEN_SEP}[/INST]" image_data = "https://picsum.photos/id/237/400/300" modality = ["image"] # Multiple images test else: image_urls = [ "https://picsum.photos/id/237/400/300", "https://picsum.photos/id/27/500/500", ] prompt = f"[INST]How many photos are there? Describe each in a very short sentence.{IMAGE_TOKEN_SEP * len(image_urls)}[/INST]" image_data = image_urls modality = ["multi-images"] pload = { "text": prompt, "image_data": image_data, "sampling_params": {"max_new_tokens": 100, "temperature": 0.7, "top_p": 0.9}, "modalities": modality, "stream": True, } response = requests.post(url, json=pload, stream=True) print("Streaming response:") prev = 0 for chunk in response.iter_lines(decode_unicode=False): chunk = chunk.decode("utf-8") if chunk and chunk.startswith("data:"): if chunk == "data: [DONE]": break data = json.loads(chunk[5:].strip("\n")) output = data["text"].strip() print(output[prev:], end="", flush=True) prev = len(output) print("\n") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--host", type=str, default="127.0.0.1") parser.add_argument("--port", type=int, default=30000) parser.add_argument( "--single-image", action="store_true", help="Test with single image instead of multiple images", ) parser.add_argument("--no-stream", action="store_true", help="Don't test streaming") args = parser.parse_args() asyncio.run(test_concurrent(args)) if not args.no_stream: test_streaming(args) ================================================ FILE: examples/runtime/multimodal/qwen_llava_server.py ================================================ """ Usage: # Installing latest llava-next: pip install git+https://github.com/LLaVA-VL/LLaVA-NeXT.git # Installing latest sglang. # Endpoint Service CLI: python -m sglang.launch_server --model-path lmms-lab/llava-next-72b --port=30000 --tp-size=8 python3 qwen_llava_server.py Output: "Two children pose with a large teddy bear, one holding a smaller stuffed bear, in a room with an American flag and potted plants." """ import argparse import asyncio import copy import json import aiohttp import requests from llava.conversation import conv_qwen from sglang.utils import normalize_base_url async def send_request(url, data, delay=0): await asyncio.sleep(delay) async with aiohttp.ClientSession() as session: async with session.post(url, json=data) as resp: output = await resp.json() return output async def test_concurrent(args): url = normalize_base_url(args.host, args.port) prompt = "\nPlease generate caption towards this image." conv_template = copy.deepcopy(conv_qwen) conv_template.append_message(role=conv_template.roles[0], message=prompt) conv_template.append_message(role=conv_template.roles[1], message=None) prompt_with_template = conv_template.get_prompt() response = [] for i in range(1): response.append( send_request( url + "/generate", { "text": prompt_with_template, "image_data": "https://farm4.staticflickr.com/3175/2653711032_804ff86d81_z.jpg", "sampling_params": { "max_new_tokens": 1024, "temperature": 0, "top_p": 1.0, "presence_penalty": 2, "frequency_penalty": 2, "stop": "<|im_end|>", }, }, ) ) rets = await asyncio.gather(*response) for ret in rets: print(ret["text"]) def test_streaming(args): url = normalize_base_url(args.host, args.port) prompt = "\nPlease generate caption towards this image." conv_template = copy.deepcopy(conv_qwen) conv_template.append_message(role=conv_template.roles[0], message=prompt) conv_template.append_message(role=conv_template.roles[1], message=None) prompt_with_template = conv_template.get_prompt() pload = { "text": prompt_with_template, "sampling_params": { "max_new_tokens": 1024, "temperature": 0, "top_p": 1.0, "presence_penalty": 2, "frequency_penalty": 2, "stop": "<|im_end|>", }, "image_data": "https://farm4.staticflickr.com/3175/2653711032_804ff86d81_z.jpg", "stream": True, } response = requests.post( url + "/generate", json=pload, stream=True, ) prev = 0 for chunk in response.iter_lines(decode_unicode=False): chunk = chunk.decode("utf-8") if chunk and chunk.startswith("data:"): if chunk == "data: [DONE]": break data = json.loads(chunk[5:].strip("\n")) output = data["text"].strip() print(output[prev:], end="", flush=True) prev = len(output) print("") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--host", type=str, default="127.0.0.1") parser.add_argument("--port", type=int, default=30000) args = parser.parse_args() asyncio.run(test_concurrent(args)) test_streaming(args) ================================================ FILE: examples/runtime/multimodal_embedding.py ================================================ # launch server # python -m sglang.launch_server --model-path Alibaba-NLP/gme-Qwen2-VL-2B-Instruct --is-embedding import requests url = "http://127.0.0.1:30000" text_input = "Represent this image in embedding space." image_path = "https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild/resolve/main/images/023.jpg" payload = { "model": "gme-qwen2-vl", "input": [{"text": text_input}, {"image": image_path}], } response = requests.post(url + "/v1/embeddings", json=payload).json() print("Embeddings:", [x.get("embedding") for x in response.get("data", [])]) ================================================ FILE: examples/runtime/openai_chat_with_response_prefill.py ================================================ """ Usage: 1) Launch the server in one terminal: python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --port 30000 2) Run this script in another terminal: python openai_chat_with_response_prefill.py This example demonstrates two chat completion calls: - One with continue_final_message enabled (the final assistant message is used as a prefill). - One without continue_final_message (the final assistant message remains, starting a new turn). """ import openai client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="EMPTY") messages = [ {"role": "system", "content": "You are a helpful AI assistant."}, { "role": "user", "content": """ Extract the name, size, price, and color from this product description as a JSON object: The SmartHome Mini is a compact smart home assistant available in black or white for only $49.99. At just 5 inches wide, it lets you control lights, thermostats, and other connected devices via voice or app— no matter where you place it in your home. This affordable little hub brings convenient hands-free control to your smart devices. """, }, {"role": "assistant", "content": "{\n"}, ] # Calling the API with continue_final_message enabled. print("=== Prefill with continue_final_messagem ===") response_with = client.chat.completions.create( model="meta-llama/Llama-3.1-8B-Instruct", messages=messages, temperature=0, extra_body={"continue_final_message": True}, ) print(response_with.choices[0].message.content) # Calling the API without continue_final_message (using default behavior). print("\n=== Prefill without continue_final_message ===") response_without = client.chat.completions.create( model="meta-llama/Llama-3.1-8B-Instruct", messages=messages, temperature=0, ) print(response_without.choices[0].message.content) ================================================ FILE: examples/runtime/qwen3_vl_reranker.py ================================================ """ Example usage of Qwen3-VL-Reranker with SGLang. This example demonstrates how to use the Qwen3-VL-Reranker model for multimodal reranking tasks, supporting text, images, and videos. Server Launch: python -m sglang.launch_server \ --model-path Qwen/Qwen3-VL-Reranker-2B \ --served-model-name Qwen3-VL-Reranker-2B \ --trust-remote-code \ --disable-radix-cache \ --chat-template examples/chat_template/qwen3_vl_reranker.jinja Client Usage: python examples/runtime/qwen3_vl_reranker.py """ import requests # Server URL BASE_URL = "http://localhost:30000" def rerank_text_only(): """Example: Text-only reranking (backward compatible).""" print("=" * 60) print("Text-only reranking example") print("=" * 60) request_data = { "query": "What is machine learning?", "documents": [ "Machine learning is a branch of artificial intelligence that enables computers to learn from data.", "The weather in Paris is usually mild with occasional rain.", "Deep learning is a subset of machine learning using neural networks with many layers.", ], "instruct": "Retrieve passages that answer the question.", "return_documents": True, } response = requests.post(f"{BASE_URL}/v1/rerank", json=request_data) results = response.json() print("Results (sorted by relevance):") for i, result in enumerate(results): print(f" {i+1}. Score: {result['score']:.4f} - {result['document'][:60]}...") print() def rerank_with_images(): """Example: Query is text, documents contain images.""" print("=" * 60) print("Image reranking example") print("=" * 60) request_data = { "query": "A woman playing with her dog on a beach at sunset.", "documents": [ # Document 1: Text description "A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset.", # Document 2: Image URL [ { "type": "image_url", "image_url": { "url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg" }, } ], # Document 3: Text + Image (mixed) [ { "type": "text", "text": "A joyful scene at the beach:", }, { "type": "image_url", "image_url": { "url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg" }, }, ], ], "instruct": "Retrieve images or text relevant to the user's query.", "return_documents": False, } response = requests.post(f"{BASE_URL}/v1/rerank", json=request_data) results = response.json() # Debug: print raw response if it's an error if isinstance(results, dict) and "message" in results: print(f"Error: {results['message']}") return if isinstance(results, str): print(f"Error: {results}") return print("Results (sorted by relevance):") for i, result in enumerate(results): print(f" {i+1}. Index: {result['index']}, Score: {result['score']:.4f}") print() def rerank_multimodal_query(): """Example: Query contains both text and image.""" print("=" * 60) print("Multimodal query reranking example") print("=" * 60) request_data = { # Query with text and image "query": [ {"type": "text", "text": "Find similar images to this:"}, { "type": "image_url", "image_url": { "url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg" }, }, ], "documents": [ "A cat sleeping on a couch.", "A woman and her dog enjoying the sunset at the beach.", "A busy city street with cars and pedestrians.", [ { "type": "image_url", "image_url": { "url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg" }, } ], ], "instruct": "Find images or descriptions similar to the query image.", } response = requests.post(f"{BASE_URL}/v1/rerank", json=request_data) results = response.json() # Debug: print raw response if it's an error if isinstance(results, dict) and "message" in results: print(f"Error: {results['message']}") return if isinstance(results, str): print(f"Error: {results}") return print("Results (sorted by relevance):") for i, result in enumerate(results): print(f" {i+1}. Index: {result['index']}, Score: {result['score']:.4f}") print() def main(): """Run all examples.""" print("\nQwen3-VL-Reranker Examples") print("Make sure the server is running with the correct model and template.\n") # Check if server is available try: response = requests.get(f"{BASE_URL}/health") if response.status_code != 200: print(f"Server health check failed: {response.status_code}") return except requests.exceptions.ConnectionError: print(f"Cannot connect to server at {BASE_URL}") print("Please start the server first with:") print(" python -m sglang.launch_server \\") print(" --model-path Qwen/Qwen3-VL-Reranker-2B \\") print(" --served-model-name Qwen3-VL-Reranker-2B \\") print(" --trust-remote-code \\") print(" --disable-radix-cache \\") print(" --chat-template examples/chat_template/qwen3_vl_reranker.jinja") return # Run examples rerank_text_only() rerank_with_images() rerank_multimodal_query() if __name__ == "__main__": main() ================================================ FILE: examples/runtime/reward_model.py ================================================ # launch server # python -m sglang.launch_server --model LxzGordon/URM-LLaMa-3.1-8B --is-embedding import requests url = "http://127.0.0.1:30000" PROMPT = ( "What is the range of the numeric output of a sigmoid node in a neural network?" ) RESPONSE1 = "The output of a sigmoid node is bounded between -1 and 1." RESPONSE2 = "The output of a sigmoid node is bounded between 0 and 1." json_data = { "conv": [ [ {"role": "user", "content": PROMPT}, {"role": "assistant", "content": RESPONSE1}, ], [ {"role": "user", "content": PROMPT}, {"role": "assistant", "content": RESPONSE2}, ], ], } response = requests.post( url + "/classify", json=json_data, ).json() print(response) print("scores:", [x["embedding"] for x in response]) ================================================ FILE: examples/runtime/token_in_token_out/token_in_token_out_llm_engine.py ================================================ """ This example demonstrates how to provide tokenized ids to LLM as input instead of text prompt, i.e. a token-in-token-out workflow. """ import sglang as sgl from sglang.srt.utils.hf_transformers_utils import get_tokenizer MODEL_PATH = "meta-llama/Llama-3.1-8B-Instruct" def main(): # Sample prompts. prompts = [ "Hello, my name is", "The president of the United States is", "The capital of France is", "The future of AI is", ] # Create a sampling params object. sampling_params = {"temperature": 0.8, "top_p": 0.95} # Tokenize inputs tokenizer = get_tokenizer(MODEL_PATH) token_ids_list = [tokenizer.encode(prompt) for prompt in prompts] # Create an LLM. llm = sgl.Engine(model_path=MODEL_PATH, skip_tokenizer_init=True) outputs = llm.generate(input_ids=token_ids_list, sampling_params=sampling_params) # Print the outputs. for prompt, output in zip(prompts, outputs): decode_output = tokenizer.decode(output["output_ids"]) print("===============================") print( f"Prompt: {prompt}\nGenerated token ids: {output['output_ids']}\nGenerated text: {decode_output}" ) print() # The __main__ condition is necessary here because we use "spawn" to create subprocesses # Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine if __name__ == "__main__": main() ================================================ FILE: examples/runtime/token_in_token_out/token_in_token_out_llm_server.py ================================================ """ Usage: python token_in_token_out_llm_server.py """ import requests from sglang.srt.utils.hf_transformers_utils import get_tokenizer from sglang.test.test_utils import is_in_ci from sglang.utils import terminate_process, wait_for_server if is_in_ci(): from docs.backend.patch import launch_server_cmd else: from sglang.utils import launch_server_cmd MODEL_PATH = "meta-llama/Llama-3.1-8B-Instruct" def main(): # Launch the server server_process, port = launch_server_cmd( f"python -m sglang.launch_server --model-path {MODEL_PATH} --skip-tokenizer-init --host 0.0.0.0" ) wait_for_server(f"http://localhost:{port}", process=server_process) # Sample prompts. prompts = [ "Hello, my name is", "The president of the United States is", "The capital of France is", "The future of AI is", ] # Create a sampling params object. sampling_params = {"temperature": 0.8, "top_p": 0.95} # Tokenize inputs tokenizer = get_tokenizer(MODEL_PATH) token_ids_list = [tokenizer.encode(prompt) for prompt in prompts] json_data = { "input_ids": token_ids_list, "sampling_params": sampling_params, } response = requests.post( f"http://localhost:{port}/generate", json=json_data, ) outputs = response.json() for prompt, output in zip(prompts, outputs): print("===============================") decode_output = tokenizer.decode(output["output_ids"]) print( f"Prompt: {prompt}\nGenerated token ids: {output['output_ids']}\nGenerated text: {decode_output}" ) print() terminate_process(server_process) if __name__ == "__main__": main() ================================================ FILE: examples/runtime/token_in_token_out/token_in_token_out_vlm_engine.py ================================================ import argparse import dataclasses from typing import Tuple from transformers import AutoProcessor from sglang import Engine from sglang.lang.chat_template import get_chat_template_by_model_path from sglang.srt.configs.model_config import ModelConfig from sglang.srt.server_args import ServerArgs from sglang.test.test_utils import DEFAULT_IMAGE_URL def get_input_ids( server_args: ServerArgs, model_config: ModelConfig ) -> Tuple[list[int], list]: chat_template = get_chat_template_by_model_path(model_config.model_path) text = f"{chat_template.image_token}What is in this picture?" image_data = [DEFAULT_IMAGE_URL] processor = AutoProcessor.from_pretrained( model_config.model_path, trust_remote_code=server_args.trust_remote_code ) input_ids = ( processor.tokenizer( text=[text], return_tensors="pt", ) .input_ids[0] .tolist() ) return input_ids, image_data def token_in_out_example( server_args: ServerArgs, ): input_ids, image_data = get_input_ids( server_args, ModelConfig( server_args.model_path, trust_remote_code=server_args.trust_remote_code, model_override_args=server_args.json_model_override_args, ), ) backend = Engine(**dataclasses.asdict(server_args)) output = backend.generate( input_ids=input_ids, image_data=image_data, sampling_params={ "temperature": 0.8, "max_new_tokens": 32, }, ) print("===============================") print(f"Output token ids: ", output["output_ids"]) backend.shutdown() if __name__ == "__main__": parser = argparse.ArgumentParser() ServerArgs.add_cli_args(parser) args = [ "--model-path=Qwen/Qwen2-VL-2B", ] args = parser.parse_args(args=args) server_args = ServerArgs.from_cli_args(args) server_args.skip_tokenizer_init = True token_in_out_example(server_args) ================================================ FILE: examples/runtime/token_in_token_out/token_in_token_out_vlm_server.py ================================================ """ Usage: python token_in_token_out_vlm_server.py """ from typing import Tuple import requests from transformers import AutoProcessor from sglang.lang.chat_template import get_chat_template_by_model_path from sglang.test.test_utils import DEFAULT_IMAGE_URL, is_in_ci from sglang.utils import terminate_process, wait_for_server if is_in_ci(): from docs.backend.patch import launch_server_cmd else: from sglang.utils import launch_server_cmd MODEL_PATH = "Qwen/Qwen2-VL-2B" def get_input_ids() -> Tuple[list[int], list]: chat_template = get_chat_template_by_model_path(MODEL_PATH) text = f"{chat_template.image_token}What is in this picture?" image_data = [DEFAULT_IMAGE_URL] processor = AutoProcessor.from_pretrained(MODEL_PATH) input_ids = ( processor.tokenizer( text=[text], return_tensors="pt", ) .input_ids[0] .tolist() ) return input_ids, image_data def main(): # Launch the server server_process, port = launch_server_cmd( f"python -m sglang.launch_server --model-path {MODEL_PATH} --skip-tokenizer-init --host 0.0.0.0" ) wait_for_server(f"http://localhost:{port}", process=server_process) input_ids, image_data = get_input_ids() sampling_params = { "temperature": 0.8, "max_new_tokens": 32, } json_data = { "input_ids": input_ids, "image_data": image_data, "sampling_params": sampling_params, } response = requests.post( f"http://localhost:{port}/generate", json=json_data, ) output = response.json() print("===============================") print(f"Output token ids: ", output["output_ids"]) terminate_process(server_process) if __name__ == "__main__": main() ================================================ FILE: examples/runtime/vertex_predict.py ================================================ """ Usage: python -m sglang.launch_server --model meta-llama/Llama-2-7b-hf --port 30000 python vertex_predict.py This example shows the request and response formats of the prediction route for Google Cloud Vertex AI Online Predictions. Vertex AI SDK for Python is recommended for deploying models to Vertex AI instead of a local server. After deploying the model to a Vertex AI Online Prediction Endpoint, send requests via the Python SDK: response = endpoint.predict( instances=[ {"text": "The capital of France is"}, {"text": "What is a car?"}, ], parameters={"sampling_params": {"max_new_tokens": 16}}, ) print(response.predictions) More details about get online predictions from Vertex AI can be found at https://cloud.google.com/vertex-ai/docs/predictions/get-online-predictions. """ from dataclasses import dataclass from typing import List, Optional import requests @dataclass class VertexPrediction: predictions: List class LocalVertexEndpoint: def __init__(self) -> None: self.base_url = "http://127.0.0.1:30000" def predict(self, instances: List[dict], parameters: Optional[dict] = None): response = requests.post( self.base_url + "/vertex_generate", json={ "instances": instances, "parameters": parameters, }, ) return VertexPrediction(predictions=response.json()["predictions"]) endpoint = LocalVertexEndpoint() # Predict with a single prompt. response = endpoint.predict(instances=[{"text": "The capital of France is"}]) print(response.predictions) # Predict with multiple prompts and parameters. response = endpoint.predict( instances=[ {"text": "The capital of France is"}, {"text": "What is a car?"}, ], parameters={"sampling_params": {"max_new_tokens": 16}}, ) print(response.predictions) ================================================ FILE: examples/sagemaker/deploy_and_serve_endpoint.py ================================================ import json import boto3 from sagemaker import serializers from sagemaker.model import Model from sagemaker.predictor import Predictor boto_session = boto3.session.Session() sm_client = boto_session.client("sagemaker") sm_role = boto_session.resource("iam").Role("SageMakerRole").arn endpoint_name = "" image_uri = "" model_id = ( "" # eg: Qwen/Qwen3-0.6B from https://huggingface.co/Qwen/Qwen3-0.6B ) hf_token = "" prompt = "" model = Model( name=endpoint_name, image_uri=image_uri, role=sm_role, env={ "SM_SGLANG_MODEL_PATH": model_id, "HF_TOKEN": hf_token, }, ) print("Model created successfully") print("Starting endpoint deployment (this may take 10-15 minutes)...") endpoint_config = model.deploy( instance_type="ml.g5.12xlarge", initial_instance_count=1, endpoint_name=endpoint_name, inference_ami_version="al2-ami-sagemaker-inference-gpu-3-1", wait=True, ) print("Endpoint deployment completed successfully") print(f"Creating predictor for endpoint: {endpoint_name}") predictor = Predictor( endpoint_name=endpoint_name, serializer=serializers.JSONSerializer(), ) payload = { "model": model_id, "messages": [{"role": "user", "content": prompt}], "max_tokens": 2400, "temperature": 0.01, "top_p": 0.9, "top_k": 50, } print(f"Sending inference request with prompt: '{prompt[:50]}...'") response = predictor.predict(payload) print("Inference request completed successfully") if isinstance(response, bytes): response = response.decode("utf-8") if isinstance(response, str): try: response = json.loads(response) except json.JSONDecodeError: print("Warning: Response is not valid JSON. Returning as string.") print(f"Received model response: '{response}'") ================================================ FILE: examples/usage/modelopt_quantize_and_export.py ================================================ #!/usr/bin/env python3 """ Example: ModelOpt Quantization and Export with SGLang This example demonstrates the streamlined workflow for quantizing a model with ModelOpt and automatically exporting it for deployment with SGLang. """ import argparse import os from typing import Optional import torch import sglang as sgl from sglang.srt.configs.device_config import DeviceConfig from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.model_config import ModelConfig from sglang.srt.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, ) from sglang.srt.model_loader.loader import get_model_loader def _validate_export(export_dir: str) -> bool: """Validate that an exported model directory contains the expected files.""" import glob required_files = ["config.json", "tokenizer_config.json"] if not os.path.exists(export_dir): return False # Check required files for file in required_files: if not os.path.exists(os.path.join(export_dir, file)): return False # Check for model files using pattern matching to handle sharded models model_patterns = [ "model*.safetensors", "pytorch_model*.bin", ] has_model_file = False for pattern in model_patterns: matching_files = glob.glob(os.path.join(export_dir, pattern)) if matching_files: has_model_file = True break return has_model_file def _get_export_info(export_dir: str) -> Optional[dict]: """Get information about an exported model.""" import json if not _validate_export(export_dir): return None try: config_path = os.path.join(export_dir, "config.json") with open(config_path, "r") as f: config = json.load(f) return { "model_type": config.get("model_type", "unknown"), "architectures": config.get("architectures", []), "quantization_config": config.get("quantization_config", {}), "export_dir": export_dir, } except Exception: return None def quantize_and_export_model( model_path: str, export_dir: str, quantization_method: str = "modelopt_fp8", checkpoint_save_path: Optional[str] = None, device: str = "cuda", ) -> None: """ Quantize a model with ModelOpt and export it for SGLang deployment. Args: model_path: Path to the original model export_dir: Directory to export the quantized model quantization_method: Quantization method ("modelopt_fp8" or "modelopt_fp4") checkpoint_save_path: Optional path to save ModelOpt checkpoint device: Device to use for quantization """ print("🚀 Starting ModelOpt quantization and export workflow") print(f"📥 Input model: {model_path}") print(f"📤 Export directory: {export_dir}") print(f"⚙️ Quantization method: {quantization_method}") # Initialize minimal distributed environment for single GPU quantization if not torch.distributed.is_initialized(): print("🔧 Initializing distributed environment...") # Set up environment variables for single-process distributed os.environ["RANK"] = "0" os.environ["WORLD_SIZE"] = "1" os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "12355" # Use a different port than tests os.environ["LOCAL_RANK"] = "0" init_distributed_environment( world_size=1, rank=0, local_rank=0, backend="nccl" if device == "cuda" else "gloo", ) initialize_model_parallel( tensor_model_parallel_size=1, pipeline_model_parallel_size=1, ) # Configure model loading with ModelOpt quantization and export model_config = ModelConfig( model_path=model_path, quantization=quantization_method, # Use unified quantization flag trust_remote_code=True, ) load_config = LoadConfig( modelopt_checkpoint_save_path=checkpoint_save_path, modelopt_export_path=export_dir, ) device_config = DeviceConfig(device=device) # Load and quantize the model (export happens automatically) print("🔄 Loading and quantizing model...") model_loader = get_model_loader(load_config, model_config) try: model_loader.load_model( model_config=model_config, device_config=device_config, ) print("✅ Model quantized successfully!") # Validate the export if _validate_export(export_dir): print("✅ Export validation passed!") info = _get_export_info(export_dir) if info: print("📋 Model info:") print(f" - Type: {info['model_type']}") print(f" - Architecture: {info['architectures']}") print(f" - Quantization: {info['quantization_config']}") else: print("❌ Export validation failed!") return except Exception as e: print(f"❌ Quantization failed: {e}") return print("\n🎉 Workflow completed successfully!") print(f"📁 Quantized model exported to: {export_dir}") print("\n🚀 To use the exported model:") print( f" python -m sglang.launch_server --model-path {export_dir} --quantization modelopt" ) print("\n # Or in Python:") print(" import sglang as sgl") print(f" llm = sgl.Engine(model_path='{export_dir}', quantization='modelopt')") print(" # Note: 'modelopt' auto-detects FP4/FP8 from model config") def deploy_exported_model( export_dir: str, host: str = "127.0.0.1", port: int = 30000, ) -> None: """ Deploy an exported ModelOpt quantized model with SGLang. Args: export_dir: Directory containing the exported model host: Host to bind the server to port: Port to bind the server to """ print(f"🚀 Deploying exported model from: {export_dir}") # Validate export first if not _validate_export(export_dir): print("❌ Invalid export directory!") return try: # Launch SGLang engine with the exported model # Using generic "modelopt" for auto-detection of FP4/FP8 llm = sgl.Engine( model_path=export_dir, quantization="modelopt", host=host, port=port, ) print("✅ Model deployed successfully!") print(f"🌐 Server running at http://{host}:{port}") # Example inference prompts = ["Hello, how are you?", "What is the capital of France?"] sampling_params = {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": 100} print("\n🧪 Running example inference...") outputs = llm.generate(prompts, sampling_params) for i, output in enumerate(outputs): print(f"Prompt {i+1}: {prompts[i]}") print(f"Output: {output['text']}") print() except Exception as e: print(f"❌ Deployment failed: {e}") def main(): parser = argparse.ArgumentParser( description="ModelOpt Quantization and Export with SGLang", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: # Quantize and export a model (recommended workflow) python modelopt_quantize_and_export.py quantize \\ --model-path TinyLlama/TinyLlama-1.1B-Chat-v1.0 \\ --export-dir ./quantized_model \\ --quantization-method modelopt_fp8 # Deploy a pre-exported model python modelopt_quantize_and_export.py deploy \\ --export-dir ./quantized_model """, ) subparsers = parser.add_subparsers(dest="command", help="Available commands") # Quantize command quantize_parser = subparsers.add_parser( "quantize", help="Quantize and export a model" ) quantize_parser.add_argument( "--model-path", required=True, help="Path to the model to quantize" ) quantize_parser.add_argument( "--export-dir", required=True, help="Directory to export the quantized model" ) quantize_parser.add_argument( "--quantization-method", choices=["modelopt_fp8", "modelopt_fp4"], default="modelopt_fp8", help="Quantization method to use", ) quantize_parser.add_argument( "--checkpoint-save-path", help="Optional path to save ModelOpt checkpoint" ) quantize_parser.add_argument( "--device", default="cuda", help="Device to use for quantization" ) # TODO: Quantize-and-serve command removed due to compatibility issues # Use the separate quantize-then-deploy workflow instead # Deploy command deploy_parser = subparsers.add_parser("deploy", help="Deploy an exported model") deploy_parser.add_argument( "--export-dir", required=True, help="Directory containing the exported model" ) deploy_parser.add_argument( "--host", default="127.0.0.1", help="Host to bind the server to" ) deploy_parser.add_argument( "--port", type=int, default=30000, help="Port to bind the server to" ) args = parser.parse_args() if args.command == "quantize": quantize_and_export_model( model_path=args.model_path, export_dir=args.export_dir, quantization_method=args.quantization_method, checkpoint_save_path=args.checkpoint_save_path, device=args.device, ) elif args.command == "deploy": deploy_exported_model( export_dir=args.export_dir, host=args.host, port=args.port, ) else: parser.print_help() if __name__ == "__main__": main() ================================================ FILE: python/pyproject.toml ================================================ [build-system] requires = ["setuptools>=61.0", "setuptools-scm>=8.0", "wheel"] build-backend = "setuptools.build_meta" [project] name = "sglang" dynamic = ["version"] description = "SGLang is a fast serving framework for large language models and vision language models." readme = "README.md" requires-python = ">=3.10" license = { file = "LICENSE" } classifiers = [ "Programming Language :: Python :: 3", "License :: OSI Approved :: Apache Software License", ] dependencies = [ "IPython", "aiohttp", "apache-tvm-ffi>=0.1.5,<0.2", "anthropic>=0.20.0", "blobfile==3.0.0", "build", "compressed-tensors", "cuda-python==12.9", "decord2 ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')", "datasets", "einops", "fastapi", "flashinfer_python==0.6.6", # keep it aligned with jit-cache version in Dockerfile "flashinfer_cubin==0.6.6", "gguf", "interegular", "llguidance>=0.7.11,<0.8.0", "modelscope", "msgspec", "ninja", "numpy", "nvidia-cutlass-dsl>=4.4.1", "nvidia-ml-py", "openai-harmony==0.0.4", "openai==2.6.1", "orjson", "outlines==0.1.11", "packaging", "partial_json_parser", "pillow", "prometheus-client>=0.20.0", "psutil", "py-spy", "pybase64", "pydantic", "python-multipart", "pyzmq>=25.1.2", "quack-kernels>=0.3.0", "requests", "scipy", "sentencepiece", "setproctitle", "flash-attn-4>=4.0.0b4", "sglang-kernel==0.4.0", "soundfile==0.13.1", "tiktoken", "timm==1.0.16", "torch_memory_saver==0.0.9", "torch==2.9.1", "torchao==0.9.0", "torchaudio==2.9.1", "torchcodec==0.9.1 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # torchcodec 0.9.1 for torch 2.9.x. Not available on Linux ARM. "av ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')", "torchvision", "tqdm", "mistral_common>=1.9.0", "transformers==5.3.0", "uvicorn", "uvloop", "watchfiles", "xgrammar==0.1.27", "smg-grpc-servicer>=0.5.0", ] [[tool.uv.index]] name = "pypi" url = "https://pypi.org/simple" default = true [[tool.uv.index]] name = "torch-cu129" url = "https://download.pytorch.org/whl/cu129" explicit = true [tool.uv.sources] torch = [ { index = "pypi", marker = "platform_machine == 'x86_64'"}, { index = "torch-cu129", marker = "platform_machine == 'aarch64'"}, ] [project.optional-dependencies] checkpoint-engine = ["checkpoint-engine==0.1.2"] diffusion = [ "PyYAML==6.0.1", "cloudpickle==3.1.2", "diffusers==0.37.0", "imageio==2.36.0", "imageio-ffmpeg==0.5.1", "moviepy>=2.0.0", "opencv-python-headless==4.10.0.84", "remote-pdb==2.1.0", "st_attn==0.0.7 ; platform_machine != 'aarch64' and platform_machine != 'arm64'", "vsa==0.0.4 ; platform_machine != 'aarch64' and platform_machine != 'arm64'", "runai_model_streamer>=0.15.5", "cache-dit==1.3.0", "addict==2.4.0", "av==16.1.0", "scikit-image==0.25.2", "trimesh>=4.0.0", "xatlas", ] ray = [ "ray[default]>=2.54.0", ] tracing = [ "opentelemetry-api", "opentelemetry-exporter-otlp", "opentelemetry-exporter-otlp-proto-grpc", "opentelemetry-sdk", ] test = [ "accelerate", "addict", "bitsandbytes", "expecttest", "jsonlines", "lm-eval[api]>=0.4.9.2", "matplotlib", "pandas", "parameterized", "peft>=0.18.0", "pytest", "pytest-cov", "diff-cover", "sentence_transformers", "tabulate", ] dev = ["sglang[test]"] all = [ "sglang[diffusion]", "sglang[tracing]", ] [tool.uv.extra-build-dependencies] st-attn = ["torch", "setuptools"] vsa = ["torch", "setuptools"] [project.urls] "Homepage" = "https://github.com/sgl-project/sglang" "Bug Tracker" = "https://github.com/sgl-project/sglang/issues" [project.scripts] sglang = "sglang.cli.main:main" [tool.setuptools.package-data] "sglang" = [ "srt/**/*", "jit_kernel/**/*" ] [tool.setuptools.packages.find] exclude = [ "assets*", "benchmark*", "docs*", "dist*", "playground*", "scripts*", "tests*", ] [tool.wheel] exclude = [ "assets*", "benchmark*", "docs*", "dist*", "playground*", "scripts*", "tests*", ] [tool.setuptools_scm] root = ".." version_file = "sglang/_version.py" git_describe_command = ["bash", "-c", "git tag --list --sort=-version:refname 'v*.*.*' | head -1 | xargs git describe --tags --long"] # Allow editable installs even when .git metadata is not available. fallback_version = "0.0.0.dev0" ================================================ FILE: python/pyproject_cpu.toml ================================================ # https://docs.sglang.io/platforms/cpu_server.html [build-system] requires = ["setuptools>=61.0", "setuptools-scm>=8.0", "wheel"] build-backend = "setuptools.build_meta" [project] name = "sglang-cpu" dynamic = ["version"] description = "SGLang is a fast serving framework for large language models and vision language models." readme = "README.md" requires-python = ">=3.10" license = { file = "LICENSE" } classifiers = [ "Programming Language :: Python :: 3", "License :: OSI Approved :: Apache Software License", ] dependencies = [ "IPython", "aiohttp", "anthropic>=0.20.0", "blobfile==3.0.0", "build", "compressed-tensors", "datasets", "einops", "fastapi", "gguf", "intel-openmp; platform_machine == 'x86_64'", "interegular", "llguidance>=0.7.11,<0.8.0", "modelscope", "msgspec", "ninja", "numpy", "openai-harmony==0.0.4", "openai==2.6.1", "orjson", "outlines", "packaging", "partial_json_parser", "pillow", "prometheus-client>=0.20.0", "psutil", "py-spy", "pybase64", "pydantic", "python-multipart", "pyzmq>=25.1.2", "requests", "scipy", "sentencepiece", "setproctitle", "soundfile==0.13.1", "tabulate", "tiktoken", "timm==1.0.16", "torch==2.9.0", "torchao==0.14.1", "torchaudio==2.9.0", "torchvision==0.24.0", "tqdm", "mistral_common>=1.9.0", "transformers==5.3.0", "triton==3.5.0", "uvicorn", "uvloop", "xgrammar==0.1.27", "smg-grpc-servicer>=0.5.0", ] [project.optional-dependencies] tracing = [ "opentelemetry-sdk", "opentelemetry-api", "opentelemetry-exporter-otlp", "opentelemetry-exporter-otlp-proto-grpc", ] test = [ "accelerate", "expecttest", "jsonlines", "matplotlib", "pandas", "peft>=0.18.0", "pytest", "sentence_transformers", ] all = [] dev = ["sglang[test]"] [project.urls] "Homepage" = "https://github.com/sgl-project/sglang" "Bug Tracker" = "https://github.com/sgl-project/sglang/issues" [project.scripts] sglang = "sglang.cli.main:main" [tool.setuptools.package-data] "sglang" = [ "srt/**/*", "jit_kernel/**/*" ] [tool.setuptools.packages.find] exclude = [ "assets*", "benchmark*", "docs*", "dist*", "playground*", "scripts*", "tests*", ] [tool.wheel] exclude = [ "assets*", "benchmark*", "docs*", "dist*", "playground*", "scripts*", "tests*", ] [tool.setuptools_scm] root = ".." version_file = "sglang/_version.py" git_describe_command = ["git", "describe", "--tags", "--long", "--match", "v*"] ================================================ FILE: python/pyproject_npu.toml ================================================ [build-system] requires = ["setuptools>=61.0", "setuptools-scm>=8.0", "wheel"] build-backend = "setuptools.build_meta" [project] name = "sglang" dynamic = ["version"] description = "SGLang is a fast serving framework for large language models and vision language models." readme = "README.md" requires-python = ">=3.10" license = { file = "LICENSE" } classifiers = [ "Programming Language :: Python :: 3", "License :: OSI Approved :: Apache Software License", ] dependencies = [ "IPython", "aiohttp", "anthropic>=0.20.0", "blobfile==3.0.0", "av", "build", "compressed-tensors", "datasets", "einops", "fastapi", "gguf", "interegular", "llguidance>=0.7.11,<0.8.0", "modelscope", "msgspec", "ninja", "numpy", "openai-harmony==0.0.4", "openai==2.6.1", "orjson", "outlines==0.1.11", "packaging", "partial_json_parser", "pillow", "prometheus-client>=0.20.0", "psutil", "py-spy", "pybase64", "pydantic", "python-multipart", "pyzmq>=25.1.2", "requests", "scipy", "sentencepiece", "setproctitle", "soundfile==0.13.1", "tiktoken", "timm==1.0.16", "torchao==0.9.0", "tqdm", "mistral_common>=1.9.0", "transformers==5.3.0", "uvicorn", "uvloop", "xgrammar==0.1.27", "smg-grpc-servicer>=0.5.0", ] [project.optional-dependencies] checkpoint-engine = ["checkpoint-engine==0.1.2"] diffusion = [ "PyYAML==6.0.1", "cloudpickle", "diffusers==0.37.0", "imageio==2.36.0", "imageio-ffmpeg==0.5.1", "moviepy>=2.0.0", "opencv-python==4.10.0.84", "remote-pdb", "cache-dit==1.2.1", "addict", "scikit-image==0.25.2", "trimesh>=4.0.0", "xatlas", ] tracing = [ "opentelemetry-api", "opentelemetry-exporter-otlp", "opentelemetry-exporter-otlp-proto-grpc", "opentelemetry-sdk", ] test = [ "accelerate", "expecttest", "gguf", "jsonlines", "matplotlib", "pandas", "peft>=0.18.0", "pytest", "sentence_transformers", "tabulate", ] # https://docs.sglang.io/platforms/ascend_npu.html srt_npu = [] all_npu = ["sglang[diffusion]"] dev_npu = ["sglang[all_npu]", "sglang[test]"] [project.urls] "Homepage" = "https://github.com/sgl-project/sglang" "Bug Tracker" = "https://github.com/sgl-project/sglang/issues" [project.scripts] sglang = "sglang.cli.main:main" [tool.setuptools.package-data] "sglang" = [ "srt/**/*", "jit_kernel/**/*" ] [tool.setuptools.packages.find] exclude = [ "assets*", "benchmark*", "docs*", "dist*", "playground*", "scripts*", "tests*", ] [tool.wheel] exclude = [ "assets*", "benchmark*", "docs*", "dist*", "playground*", "scripts*", "tests*", ] [tool.setuptools_scm] root = ".." version_file = "sglang/_version.py" git_describe_command = ["git", "describe", "--tags", "--long", "--match", "v*"] ================================================ FILE: python/pyproject_other.toml ================================================ [build-system] requires = ["setuptools>=61.0", "setuptools-scm>=8.0", "wheel"] build-backend = "setuptools.build_meta" [project] name = "sglang" dynamic = ["version"] description = "SGLang is a fast serving framework for large language models and vision language models." readme = "README.md" requires-python = ">=3.10" license = { file = "LICENSE" } classifiers = [ "Programming Language :: Python :: 3", "License :: OSI Approved :: Apache Software License", ] dependencies = ["aiohttp", "requests", "tqdm", "numpy", "IPython", "setproctitle"] [project.optional-dependencies] runtime_common = [ "IPython", "aiohttp", "anthropic>=0.20.0", "blobfile==3.0.0", "av", "build", "compressed-tensors", "datasets", "einops", "fastapi", "gguf", "interegular", "llguidance>=0.7.11,<0.8.0", "modelscope", "msgspec", "ninja", "numpy", "openai-harmony==0.0.4", "openai==2.6.1", "orjson", "outlines==0.1.11", "packaging", "partial_json_parser", "pillow", "prometheus-client>=0.20.0", "psutil", "py-spy", "pybase64", "pydantic", "python-multipart", "pyzmq>=25.1.2", "requests", "scipy", "sentencepiece", "setproctitle", "soundfile==0.13.1", "tiktoken", "timm==1.0.16", "torchao==0.9.0", "tqdm", "mistral_common>=1.9.0", "transformers==5.3.0", "uvicorn", "uvloop", "xgrammar==0.1.27", "smg-grpc-servicer>=0.5.0", ] diffusion_common = [ "PyYAML==6.0.1", "cloudpickle", "diffusers==0.37.0", "imageio==2.36.0", "imageio-ffmpeg==0.5.1", "moviepy>=2.0.0", "opencv-python-headless==4.10.0.84", "remote-pdb", "addict", "scikit-image==0.25.2", "trimesh>=4.0.0", "xatlas", ] tracing = [ "opentelemetry-sdk", "opentelemetry-api", "opentelemetry-exporter-otlp", "opentelemetry-exporter-otlp-proto-grpc", ] # HIP (Heterogeneous-computing Interface for Portability) for AMD # => base docker rocm/vllm-dev:20250114, not from public vllm whl srt_hip = [ "sglang[runtime_common]", "torch", "petit_kernel==0.0.2", "wave-lang==3.8.2", ] diffusion_hip = [ "sglang[diffusion_common]", "st_attn==0.0.7", "vsa==0.0.4", "runai_model_streamer>=0.15.5", "cache-dit==1.1.8", ] # For Intel Gaudi(device : hpu) follow the installation guide # https://docs.vllm.ai/en/latest/getting_started/gaudi-installation.html srt_hpu = ["sglang[runtime_common]"] # https://docs.sglang.io/platforms/mthreads_gpu.md srt_musa = [ "sglang[runtime_common]", "torch", "torch_musa", "torchada>=0.1.25", "mthreads-ml-py", "numpy<2.0", ] diffusion_musa = [ "sglang[diffusion_common]", "st_attn==0.0.7", "vsa==0.0.4", "runai_model_streamer>=0.15.5", "cache-dit==1.1.8", ] # https://docs.sglang.io/platforms/mps.md srt_mps = [ "sglang[runtime_common]", "torch==2.9.1", "torchao==0.9.0", "torchaudio==2.9.1", "torchvision", ] diffusion_mps = [ "sglang[diffusion_common]", "cloudpickle==3.1.2", "remote-pdb==2.1.0", "cache-dit==1.2.3", "addict==2.4.0", "av==16.1.0", ] test = [ "accelerate", "expecttest", "gguf", "jsonlines", "matplotlib", "pandas", "peft>=0.18.0", "pytest", "sentence_transformers", "tabulate", ] all_hip = ["sglang[srt_hip]", "sglang[diffusion_hip]"] all_hpu = ["sglang[srt_hpu]"] all_musa = ["sglang[srt_musa]", "sglang[diffusion_musa]"] all_mps = ["sglang[srt_mps]", "sglang[diffusion_mps]"] dev_hip = ["sglang[all_hip]", "sglang[test]"] dev_hpu = ["sglang[all_hpu]", "sglang[test]"] dev_musa = ["sglang[all_musa]", "sglang[test]"] dev_mps = ["sglang[all_mps]", "sglang[test]"] [project.urls] "Homepage" = "https://github.com/sgl-project/sglang" "Bug Tracker" = "https://github.com/sgl-project/sglang/issues" [project.scripts] sglang = "sglang.cli.main:main" [tool.setuptools.package-data] "sglang" = [ "srt/**/*", "jit_kernel/**/*" ] [tool.setuptools.packages.find] exclude = [ "assets*", "benchmark*", "docs*", "dist*", "playground*", "scripts*", "tests*", ] [tool.wheel] exclude = [ "assets*", "benchmark*", "docs*", "dist*", "playground*", "scripts*", "tests*", ] [tool.setuptools_scm] root = ".." version_file = "sglang/_version.py" git_describe_command = ["git", "describe", "--tags", "--long", "--match", "v*"] ================================================ FILE: python/pyproject_xpu.toml ================================================ [build-system] requires = ["setuptools>=61.0", "setuptools-scm>=8.0", "wheel"] build-backend = "setuptools.build_meta" [project] name = "sglang" dynamic = ["version"] description = "SGLang is a fast serving framework for large language models and vision language models." readme = "README.md" requires-python = ">=3.10" license = { file = "LICENSE" } classifiers = [ "Programming Language :: Python :: 3", "License :: OSI Approved :: Apache Software License", ] dependencies = [ "torch==2.10.0+xpu", "torchcodec==0.10.0 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # torchcodec does not exist in those systems. torch==2.10.0 on XPU uses 0.10.0 "av ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')", "torchaudio==2.10.0+xpu", "torchvision", "sgl-kernel @ git+https://github.com/sgl-project/sgl-kernel-xpu.git", "IPython", "aiohttp", "anthropic>=0.20.0", "blobfile==3.0.0", "build", "compressed-tensors", "datasets", "einops", "fastapi", "gguf", "interegular", "llguidance>=0.7.11,<0.8.0", "modelscope", "msgspec", "ninja", "numpy", "openai-harmony==0.0.4", "openai==2.6.1", "orjson", "outlines==0.1.11", "packaging", "partial_json_parser", "pillow", "prometheus-client>=0.20.0", "psutil", "py-spy", "pybase64", "pydantic", "python-multipart", "pyzmq>=25.1.2", "requests", "scipy", "sentencepiece", "setproctitle", "soundfile==0.13.1", "tiktoken", "timm==1.0.16", "torchao==0.9.0", "tqdm", "mistral_common>=1.9.0", "transformers==5.3.0", "uvicorn", "uvloop", # "xgrammar==0.1.24", , xgrammar depends on CUDA PyTorch and Triton only "smg-grpc-servicer>=0.5.0", ] [project.optional-dependencies] tracing = [ "opentelemetry-api", "opentelemetry-exporter-otlp", "opentelemetry-exporter-otlp-proto-grpc", "opentelemetry-sdk", ] test = [ "accelerate", "bitsandbytes", "expecttest", "jsonlines", "lm-eval[api]>=0.4.9.2", "matplotlib", "pandas", "parameterized", "peft>=0.18.0", "pytest", "sentence_transformers", "tabulate", ] dev = ["sglang[test]"] all = [ "sglang[tracing]", ] [project.urls] "Homepage" = "https://github.com/sgl-project/sglang" "Bug Tracker" = "https://github.com/sgl-project/sglang/issues" [project.scripts] sglang = "sglang.cli.main:main" [tool.setuptools.package-data] "sglang" = [ "srt/**/*", "jit_kernel/**/*" ] [tool.setuptools.packages.find] exclude = [ "assets*", "benchmark*", "docs*", "dist*", "playground*", "scripts*", "tests*", ] [tool.wheel] exclude = [ "assets*", "benchmark*", "docs*", "dist*", "playground*", "scripts*", "tests*", ] [tool.setuptools_scm] root = ".." version_file = "sglang/_version.py" git_describe_command = ["bash", "-c", "git tag --list --sort=-version:refname 'v*.*.*' | head -1 | xargs git describe --tags --long"] # Allow editable installs even when .git metadata is not available. fallback_version = "0.0.0.dev0" ================================================ FILE: python/sglang/README.md ================================================ # Code Structure - `eval`: The evaluation utilities. - `lang`: The frontend language. - `multimodal_gen`: Inference framework for accelerated image/video generation. - `srt`: The backend engine for running local models. (SRT = SGLang Runtime). - `test`: The test utilities. - `api.py`: The public APIs. - `bench_offline_throughput.py`: Benchmark the performance in the offline mode. - `bench_one_batch.py`: Benchmark the latency of running a single static batch without a server. - `bench_one_batch_server.py`: Benchmark the latency of running a single batch with a server. - `bench_serving.py`: Benchmark online serving with dynamic requests. - `check_env.py`: Check the environment variables and dependencies. - `global_config.py`: The global configs and constants. - `launch_server.py`: The entry point for launching a local server. - `profiler.py`: The profiling entry point to send profile requests. - `utils.py`: Common utilities. - `version.py`: Version info. ================================================ FILE: python/sglang/__init__.py ================================================ # SGLang public APIs # Install stubs early for platforms where certain dependencies are unavailable # (e.g. macOS/MPS has no triton, and torch.mps lacks Stream / set_device / # get_device_properties). This must run before any downstream imports. import sys as _sys if _sys.platform == "darwin": try: import torch as _torch if _torch.backends.mps.is_available(): from sglang._triton_stub import install as _install_triton_stub _install_triton_stub() del _install_triton_stub from sglang._mps_stub import install as _install_mps_stub _install_mps_stub() del _install_mps_stub del _torch except ImportError: pass del _sys # Frontend Language APIs from sglang.global_config import global_config from sglang.lang.api import ( Engine, Runtime, assistant, assistant_begin, assistant_end, flush_cache, function, gen, gen_int, gen_string, get_server_info, image, select, separate_reasoning, set_default_backend, system, system_begin, system_end, user, user_begin, user_end, video, ) from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint from sglang.lang.choices import ( greedy_token_selection, token_length_normalized, unconditional_likelihood_normalized, ) # Lazy import some libraries from sglang.utils import LazyImport from sglang.version import __version__ Anthropic = LazyImport("sglang.lang.backend.anthropic", "Anthropic") LiteLLM = LazyImport("sglang.lang.backend.litellm", "LiteLLM") OpenAI = LazyImport("sglang.lang.backend.openai", "OpenAI") VertexAI = LazyImport("sglang.lang.backend.vertexai", "VertexAI") # Runtime Engine APIs ServerArgs = LazyImport("sglang.srt.server_args", "ServerArgs") Engine = LazyImport("sglang.srt.entrypoints.engine", "Engine") __all__ = [ "Engine", "Runtime", "assistant", "assistant_begin", "assistant_end", "flush_cache", "function", "gen", "gen_int", "gen_string", "get_server_info", "image", "select", "separate_reasoning", "set_default_backend", "system", "system_begin", "system_end", "user", "user_begin", "user_end", "video", "RuntimeEndpoint", "greedy_token_selection", "token_length_normalized", "unconditional_likelihood_normalized", "ServerArgs", "Anthropic", "LiteLLM", "OpenAI", "VertexAI", "global_config", "__version__", ] ================================================ FILE: python/sglang/_mps_stub.py ================================================ """Stub implementations for APIs missing from ``torch.mps``. ``torch.mps`` lacks several APIs that ``torch.cuda`` provides (``Stream``, ``set_device``, ``get_device_properties``, …). Rather than scattering ``hasattr`` / ``getattr`` guards throughout the codebase, we monkey-patch ``torch.mps`` once at startup so that generic device-agnostic code paths just work. """ from __future__ import annotations import functools from dataclasses import dataclass, field from typing import Any class Stream: """Minimal stand-in for ``torch.cuda.Stream``. MPS does not expose user-visible streams. Every method is a no-op so that code written for CUDA's multi-stream model still runs. """ def __init__(self, device: Any = None, priority: int = 0) -> None: pass def synchronize(self) -> None: pass def wait_stream(self, stream: Any) -> None: pass def wait_event(self, event: Any) -> None: pass def record_event(self, event: Any = None) -> Any: return None def query(self) -> bool: return True # context-manager protocol (``with stream:``) def __enter__(self) -> "Stream": return self def __exit__(self, *args: Any) -> None: pass class Event: """Minimal stand-in for ``torch.cuda.Event``.""" def __init__(self, enable_timing: bool = False) -> None: pass def record(self, stream: Any = None) -> None: pass def wait(self, stream: Any = None) -> None: pass def query(self) -> bool: return True def synchronize(self) -> None: pass def elapsed_time(self, end_event: Any) -> float: return 0.0 _default_stream = Stream() def current_stream(device: Any = None) -> Stream: """Return the default (and only) MPS stream.""" return _default_stream def stream(s: Any) -> Stream: """Return a context manager that is a no-op on MPS.""" return s if s is not None else _default_stream def set_device(device: Any) -> None: # noqa: ARG001 """Set the current device. This is a no-op for MPS as it has exactly one device.""" pass def current_device() -> int: """Return the index of the current MPS device (always 0).""" return 0 def device_count() -> int: """Return the number of available MPS devices (always 1).""" return 1 @dataclass class _MPSDeviceProperties: """Mimics the object returned by ``torch.cuda.get_device_properties``.""" name: str = "Apple MPS" total_memory: int = 0 # populated at install time multi_processor_count: int = 0 warp_size: int = 32 is_integrated: bool = True major: int = 0 minor: int = 0 # Extra attrs some callers inspect _extra: dict = field(default_factory=dict) def __getattr__(self, name: str) -> Any: # Return a safe default for any attribute we didn't anticipate try: return self._extra[name] except KeyError: return None _cached_props: _MPSDeviceProperties | None = None def get_device_properties(device: Any = 0) -> _MPSDeviceProperties: # noqa: ARG001 """Return the properties of the MPS device. Results are cached after first call.""" global _cached_props if _cached_props is None: import psutil _cached_props = _MPSDeviceProperties( total_memory=psutil.virtual_memory().total, ) return _cached_props class _MPSMemoryTracker: """Tracks peak memory values on top of ``torch.mps`` current-value APIs. * ``memory_allocated`` → ``torch.mps.current_allocated_memory()`` * ``memory_reserved`` → ``torch.mps.driver_allocated_memory()`` * ``max_memory_*`` → high-water marks of the above """ def __init__(self) -> None: self._peak_allocated: int = 0 self._peak_reserved: int = 0 def memory_allocated(self, device: Any = None) -> int: # noqa: ARG002 import torch val = torch.mps.current_allocated_memory() if val > self._peak_allocated: self._peak_allocated = val return val def memory_reserved(self, device: Any = None) -> int: # noqa: ARG002 import torch val = torch.mps.driver_allocated_memory() if val > self._peak_reserved: self._peak_reserved = val return val def max_memory_allocated(self, device: Any = None) -> int: # noqa: ARG002 self.memory_allocated() return self._peak_allocated def max_memory_reserved(self, device: Any = None) -> int: # noqa: ARG002 self.memory_reserved() return self._peak_reserved def reset_peak_memory_stats(self, device: Any = None) -> None: # noqa: ARG002 import torch self._peak_allocated = torch.mps.current_allocated_memory() self._peak_reserved = torch.mps.driver_allocated_memory() _memory_tracker = _MPSMemoryTracker() def _patch_non_blocking() -> None: """Force ``non_blocking=False`` for copies targeting the MPS device. Unlike CUDA, MPS does not guarantee that a subsequent kernel on the same "stream" will wait for an async host-to-device transfer to finish. Reading the tensor before the transfer completes yields uninitialised (garbage) data. Patching ``Tensor.to`` and ``Tensor.copy_`` centrally avoids having to sprinkle ``non_blocking=not is_mps()`` at every call-site. """ import torch _original_to = torch.Tensor.to @functools.wraps(_original_to) def _patched_to(self, *args, **kwargs): if kwargs.get("non_blocking"): # Detect target device from positional or keyword args device = None if args and isinstance(args[0], (str, torch.device)): device = torch.device(args[0]) if isinstance(args[0], str) else args[0] elif "device" in kwargs: d = kwargs["device"] device = torch.device(d) if isinstance(d, str) else d if device is not None and device.type == "mps": kwargs = {**kwargs, "non_blocking": False} return _original_to(self, *args, **kwargs) torch.Tensor.to = _patched_to _original_copy_ = torch.Tensor.copy_ @functools.wraps(_original_copy_) def _patched_copy_(self, src, non_blocking=False): if non_blocking and self.device.type == "mps": non_blocking = False return _original_copy_(self, src, non_blocking=non_blocking) torch.Tensor.copy_ = _patched_copy_ _installed = False def install() -> None: """Patch ``torch.mps`` with the stubs above. Safe to call multiple times.""" global _installed if _installed: return import torch mps = torch.mps # Only patch attributes that are actually missing for name, obj in [ ("Stream", Stream), ("Event", Event), ("current_stream", current_stream), ("stream", stream), ("set_device", set_device), ("current_device", current_device), ("device_count", device_count), ("get_device_properties", get_device_properties), ("reset_peak_memory_stats", _memory_tracker.reset_peak_memory_stats), ("memory_allocated", _memory_tracker.memory_allocated), ("memory_reserved", _memory_tracker.memory_reserved), ("max_memory_allocated", _memory_tracker.max_memory_allocated), ("max_memory_reserved", _memory_tracker.max_memory_reserved), ]: if not hasattr(mps, name): setattr(mps, name, obj) _patch_non_blocking() _installed = True ================================================ FILE: python/sglang/_triton_stub.py ================================================ """ Mock triton module for platforms where triton is not available (e.g., macOS/MPS). This module provides stub implementations of triton APIs so that modules which import triton at the top level can be loaded without error. The actual triton kernels are never executed on these platforms – alternative backends (e.g. SDPA for MPS) are used instead. Usage – call ``install()`` **before** any ``import triton`` in the process: from sglang._triton_stub import install install() """ import sys import types class _StubBase: """A base class that any mock attribute can safely be subclassed from. Used when external code does ``class Foo(triton.runtime.KernelInterface):``. """ def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) class _MockModule(types.ModuleType): """A module whose every attribute is itself a ``_MockModule``. When called (e.g. ``@triton.jit``), it acts as a pass-through decorator so that kernel *definitions* are syntactically valid even though they will never be compiled. """ def __init__(self, name: str): super().__init__(name) self.__path__: list[str] = [] # make it look like a package self.__package__ = name self.__file__ = __file__ self._children: dict[str, object] = {} # Set __spec__ so that importlib.util.find_spec() works on cached modules import importlib self.__spec__ = importlib.machinery.ModuleSpec(name, None, is_package=True) def __getattr__(self, name: str): """Handle attribute access by creating and returning a child _MockModule.""" if name.startswith("__") and name.endswith("__"): raise AttributeError(name) full = f"{self.__name__}.{name}" if full in sys.modules: return sys.modules[full] # If the name looks like a class (CamelCase / uppercase), return a # stub class that can be used as a base class for inheritance. if name[0:1].isupper(): stub_cls = type(name, (_StubBase,), {"__module__": self.__name__}) self._children[name] = stub_cls return stub_cls child = _MockModule(full) sys.modules[full] = child self._children[name] = child return child def __call__(self, *args, **kwargs): # Direct decorator usage: @triton.jit (receives the function) if len(args) == 1 and callable(args[0]) and not kwargs: return args[0] # Parameterised decorator: @triton.jit(...) → returns a decorator def _decorator(fn): return fn return _decorator def __instancecheck__(self, instance): """Return False for all instance checks against the mock.""" return False def __contains__(self, item): """Return False for all membership checks.""" return False def __iter__(self): return iter([]) def __len__(self): return 0 def __bool__(self): return False def __repr__(self): return f"" def _cdiv(a: int, b: int) -> int: """Ceiling division – mirrors ``triton.cdiv``.""" return -(a // -b) def _next_power_of_2(n: int) -> int: """Mirrors ``triton.next_power_of_2``.""" return 1 << (n - 1).bit_length() if n > 0 else 1 class _Config: """Minimal stand-in for ``triton.Config`` used in ``@triton.autotune``.""" def __init__(self, kwargs=None, num_warps=4, num_stages=2, **extra): self.kwargs = kwargs or {} self.num_warps = num_warps self.num_stages = num_stages class _TritonFinder: """A meta-path finder that intercepts all ``import triton.*`` statements. When Python encounters ``import triton.backends.compiler``, it walks the dotted path and tries to import each component. Our mock module's ``__getattr__`` handles *attribute* access, but the import machinery uses ``importlib`` finders, not attribute access, for sub-module resolution. This finder bridges that gap by creating ``_MockModule`` instances for any ``triton.*`` sub-module that isn't already in ``sys.modules``. """ def find_module(self, fullname, path=None): if fullname == "triton" or fullname.startswith("triton."): return self return None def load_module(self, fullname): if fullname in sys.modules: return sys.modules[fullname] mod = _MockModule(fullname) sys.modules[fullname] = mod # Wire up the parent relationship parts = fullname.rsplit(".", 1) if len(parts) == 2: parent_name, child_name = parts parent = sys.modules.get(parent_name) if parent is not None: setattr(parent, child_name, mod) return mod def _make_mock(name: str) -> _MockModule: """Create a ``_MockModule`` and register it in ``sys.modules``.""" mod = _MockModule(name) sys.modules[name] = mod return mod def install() -> None: """Register a mock ``triton`` package in *sys.modules*. This is a no-op if a real ``triton`` is already importable. """ if "triton" in sys.modules: return # Check whether a real triton exists before installing the stub. import importlib.util if importlib.util.find_spec("triton") is not None: return # Register the meta-path finder FIRST so that any ``import triton.X`` # during the rest of install() (or later) is handled. sys.meta_path.insert(0, _TritonFinder()) triton = _make_mock("triton") triton.__version__ = "3.0.0" triton.cdiv = _cdiv triton.next_power_of_2 = _next_power_of_2 triton.Config = _Config # triton.language (commonly imported as ``tl``) tl = _make_mock("triton.language") class _constexpr: """Stand-in for ``tl.constexpr`` – works as both annotation and value wrapper.""" def __init__(self, value=None): self.value = value def __repr__(self): return f"constexpr({self.value!r})" tl.constexpr = _constexpr triton.language = tl # triton.language.extra.libdevice extra = _make_mock("triton.language.extra") tl.extra = extra libdevice = _make_mock("triton.language.extra.libdevice") extra.libdevice = libdevice # triton.runtime.jit (JITFunction used in isinstance checks) runtime = _make_mock("triton.runtime") triton.runtime = runtime jit_mod = _make_mock("triton.runtime.jit") class _JITFunction: """Dummy so ``isinstance(fn, triton.runtime.jit.JITFunction)`` works.""" pass jit_mod.JITFunction = _JITFunction runtime.jit = jit_mod # triton.runtime.driver (used by fla/utils.py) driver = _make_mock("triton.runtime.driver") runtime.driver = driver # triton.testing testing = _make_mock("triton.testing") triton.testing = testing # triton.tools / triton.tools.tensor_descriptor tools = _make_mock("triton.tools") triton.tools = tools td = _make_mock("triton.tools.tensor_descriptor") tools.tensor_descriptor = td # triton.backends / triton.backends.compiler (used by torch._inductor) backends = _make_mock("triton.backends") triton.backends = backends compiler = _make_mock("triton.backends.compiler") backends.compiler = compiler ================================================ FILE: python/sglang/bench_offline_throughput.py ================================================ """ Benchmark the throughput in the offline mode. It accepts server arguments (the same as launch_server.py) and benchmark arguments (the same as bench_serving.py). # Usage ## Sharegpt dataset with default args python -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --num-prompts 10 ## Random dataset with default args python -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 """ import argparse import asyncio import dataclasses import inspect import json import logging import os import random import time from typing import Dict, List, Optional import numpy as np from sglang.benchmark.datasets import DatasetRow, get_dataset from sglang.benchmark.datasets.random import sample_random_requests from sglang.benchmark.utils import get_tokenizer, set_ulimit from sglang.lang.backend.runtime_endpoint import Runtime from sglang.srt.entrypoints.engine import Engine from sglang.srt.server_args import ServerArgs @dataclasses.dataclass class BenchArgs: backend: str = "engine" result_filename: str = "" dataset_name: str = "sharegpt" dataset_path: str = "" num_prompts: int = 1000 sharegpt_output_len: Optional[int] = None sharegpt_context_len: Optional[int] = None random_input_len: int = 1024 random_output_len: int = 1024 random_range_ratio: float = 0.0 gsp_num_groups: int = 64 gsp_prompts_per_group: int = 16 gsp_system_prompt_len: int = 2048 gsp_question_len: int = 128 gsp_output_len: int = 256 seed: int = 1 disable_ignore_eos: bool = False extra_request_body: Optional[str] = None apply_chat_template: bool = False profile: bool = False skip_warmup: bool = False do_not_exit: bool = False prompt_suffix: str = "" return_logprob: bool = False logprob_start_len: int = -1 @staticmethod def add_cli_args(parser: argparse.ArgumentParser): parser.add_argument("--backend", type=str, default=BenchArgs.backend) parser.add_argument( "--result-filename", type=str, default=BenchArgs.result_filename ) parser.add_argument( "--dataset-name", type=str, default="sharegpt", choices=["sharegpt", "random", "generated-shared-prefix"], help="Name of the dataset to benchmark on.", ) parser.add_argument( "--dataset-path", type=str, default="", help="Path to the dataset." ) parser.add_argument( "--num-prompts", type=int, default=BenchArgs.num_prompts, help="Number of prompts to process. Default is 1000.", ) parser.add_argument( "--sharegpt-output-len", type=int, default=BenchArgs.sharegpt_output_len, help="Output length for each request. Overrides the output length from the ShareGPT dataset.", ) parser.add_argument( "--sharegpt-context-len", type=int, default=BenchArgs.sharegpt_context_len, help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.", ) parser.add_argument( "--random-input-len", type=int, default=BenchArgs.random_input_len, help="Number of input tokens per request, used only for random dataset.", ) parser.add_argument( "--random-output-len", type=int, default=BenchArgs.random_output_len, help="Number of output tokens per request, used only for random dataset.", ) parser.add_argument( "--random-range-ratio", type=float, default=BenchArgs.random_range_ratio, help="Range of sampled ratio of input/output length, " "used only for random dataset.", ) parser.add_argument( "--gsp-num-groups", type=int, default=BenchArgs.gsp_num_groups, help="Number of groups with shared prefix, used" "only for generate-shared-prefix", ) parser.add_argument( "--gsp-prompts-per-group", type=int, default=BenchArgs.gsp_prompts_per_group, help="Number of prompts per group of shared prefix, used" "only for generate-shared-prefix", ) parser.add_argument( "--gsp-system-prompt-len", type=int, default=BenchArgs.gsp_system_prompt_len, help="System prompt length, used" "only for generate-shared-prefix", ) parser.add_argument( "--gsp-question-len", type=int, default=BenchArgs.gsp_question_len, help="Question length, used" "only for generate-shared-prefix", ) parser.add_argument( "--gsp-output-len", type=int, default=BenchArgs.gsp_output_len, help="Target length in tokens for outputs in generated-shared-prefix dataset", ) parser.add_argument("--seed", type=int, default=1, help="The random seed.") parser.add_argument( "--disable-ignore-eos", action="store_true", help="Disable ignore EOS token", ) parser.add_argument( "--extra-request-body", metavar='{"key1": "value1", "key2": "value2"}', type=str, default=BenchArgs.extra_request_body, help="Append given JSON object to the request payload. You can use this to specify" "additional generate params like sampling params.", ) parser.add_argument( "--apply-chat-template", action="store_true", help="Apply chat template", ) parser.add_argument( "--profile", action="store_true", help="Use Torch Profiler. The endpoint must be launched with " "SGLANG_TORCH_PROFILER_DIR to enable profiler.", ) parser.add_argument( "--skip-warmup", action="store_true", help="Skip the warmup batches.", ) parser.add_argument( "--do-not-exit", action="store_true", help="Do not exit the program. This is useful for nsys profile with --duration and --delay.", ) parser.add_argument( "--prompt-suffix", type=str, default="", help="Suffix applied to the end of all user prompts, followed by assistant prompt suffix.", ) parser.add_argument( "--return-logprob", action="store_true", help="Enable returning log probabilities.", ) parser.add_argument( "--logprob-start-len", type=int, default=-1, help="Start length for logprob. -1 means only return logprobs for output tokens (default). 0 means return logprobs for all tokens including input.", ) @classmethod def from_cli_args(cls, args: argparse.Namespace): attrs = [attr.name for attr in dataclasses.fields(cls)] return cls(**{attr: getattr(args, attr) for attr in attrs}) def throughput_test_once( backend_name: str, backend, reqs: List[DatasetRow], ignore_eos: bool, extra_request_body: Dict, profile: bool, return_logprob: bool = False, logprob_start_len: int = -1, ): measurement_results = { "backend": backend_name, "successful_requests": len(reqs), "total_latency": -1, "total_input_tokens": sum(r.prompt_len for r in reqs), "total_output_tokens": -1, "request_throughput": -1, "input_throughput": -1, "output_throughput": -1, "total_throughput": -1, } prompt = [r.prompt for r in reqs] sampling_params = [ { "temperature": 0, "max_new_tokens": r.output_len, "ignore_eos": ignore_eos, **extra_request_body, } for r in reqs ] if profile: assert ( "SGLANG_TORCH_PROFILER_DIR" in os.environ ), "Please set SGLANG_TORCH_PROFILER_DIR." os.makedirs(os.environ["SGLANG_TORCH_PROFILER_DIR"], exist_ok=True) backend.start_profile() st = time.perf_counter() gen_out = backend.generate( prompt=prompt, sampling_params=sampling_params, return_logprob=return_logprob, logprob_start_len=logprob_start_len, ) latency = time.perf_counter() - st if profile: dir = os.getenv("SGLANG_TORCH_PROFILER_DIR") known_files = set(os.listdir(dir)) backend.stop_profile() monitor_trace_file(known_files, dir) if backend_name == "runtime": gen_out = json.loads(gen_out) server_info = backend.get_server_info() measurement_results["total_latency"] = latency measurement_results["total_output_tokens"] = sum( o["meta_info"]["completion_tokens"] for o in gen_out ) measurement_results["request_throughput"] = ( measurement_results["successful_requests"] / latency ) measurement_results["input_throughput"] = ( measurement_results["total_input_tokens"] / latency ) measurement_results["output_throughput"] = ( measurement_results["total_output_tokens"] / latency ) measurement_results["total_throughput"] = ( measurement_results["total_input_tokens"] + measurement_results["total_output_tokens"] ) / latency if inspect.isawaitable(server_info): server_info = asyncio.run(server_info) measurement_results["last_gen_throughput"] = server_info["internal_states"][0][ "last_gen_throughput" ] return measurement_results def monitor_trace_file(known_files, directory, interval=1): print(f"Monitoring {directory} for new trace files...") while True: flag = False time.sleep(interval) current_files = set(os.listdir(directory)) new_files = current_files - known_files for new_file in new_files: new_file_path = os.path.join(directory, new_file) print(f"New file detected: {new_file}") previous_size = 0 while True: try: current_size = os.path.getsize(new_file_path) except FileNotFoundError: print(f"File {new_file} is no longer accessible.") break if current_size > previous_size: previous_size = current_size else: flag = True break time.sleep(interval) if flag: break def _create_ray_engine_backend(server_args: ServerArgs): """Create a RayEngine inside a Ray actor on a placement group. RayEngine requires a placement group, so we launch it inside a Ray actor and return a lightweight proxy that forwards calls via ray.get(). """ import ray from ray.runtime_env import RuntimeEnv from ray.util.placement_group import placement_group from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy env_vars = {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"} if os.environ.get("HF_TOKEN"): env_vars["HF_TOKEN"] = os.environ["HF_TOKEN"] if not ray.is_initialized(): ray.init(runtime_env=RuntimeEnv(env_vars=env_vars)) total_gpus = server_args.tp_size * server_args.pp_size pg = placement_group([{"CPU": 1, "GPU": total_gpus}], strategy="STRICT_PACK") ray.get(pg.ready()) @ray.remote class _EngineActor: def __init__(self, **kwargs): from sglang.srt.ray.engine import RayEngine self.engine = RayEngine(**kwargs) def call(self, method, **kwargs): return getattr(self.engine, method)(**kwargs) actor = _EngineActor.options( num_cpus=1, num_gpus=0, scheduling_strategy=PlacementGroupSchedulingStrategy( placement_group=pg, placement_group_bundle_index=0, ), ).remote(**dataclasses.asdict(server_args)) class _Proxy: """Forwards method calls to the remote RayEngine actor.""" def generate(self, **kwargs): return ray.get(actor.call.remote("generate", **kwargs)) def get_server_info(self, **kwargs): return ray.get(actor.call.remote("get_server_info", **kwargs)) def start_profile(self, **kwargs): return ray.get(actor.call.remote("start_profile", **kwargs)) def stop_profile(self, **kwargs): return ray.get(actor.call.remote("stop_profile", **kwargs)) def shutdown(self): try: ray.get(actor.call.remote("shutdown"), timeout=60) except Exception: pass try: ray.util.remove_placement_group(pg) except Exception: pass return _Proxy() def throughput_test( server_args: ServerArgs, bench_args: BenchArgs, ): if bench_args.backend == "engine": if server_args.use_ray: backend = _create_ray_engine_backend(server_args) else: backend = Engine(**dataclasses.asdict(server_args)) if not backend: raise ValueError("Please provide valid engine arguments") elif bench_args.backend == "runtime": backend = Runtime(**dataclasses.asdict(server_args)) else: raise ValueError('Please set backend to either "engine" or "runtime"') tokenizer_id = server_args.tokenizer_path or server_args.model_path tokenizer = get_tokenizer(tokenizer_id) # Set global environments set_ulimit() random.seed(bench_args.seed) np.random.seed(bench_args.seed) # Parse args extra_request_body = {} if bench_args.extra_request_body: extra_request_body = json.loads(args.extra_request_body) # Read dataset input_requests = get_dataset(bench_args, tokenizer) warmup_requests = sample_random_requests( input_len=256, output_len=16, num_prompts=min(bench_args.num_prompts, 16), range_ratio=1.0, tokenizer=tokenizer, dataset_path=bench_args.dataset_path, ) # Warm up if not bench_args.skip_warmup: logging.info("\nWarmup...") throughput_test_once( backend_name=bench_args.backend, backend=backend, reqs=warmup_requests, ignore_eos=not bench_args.disable_ignore_eos, extra_request_body=extra_request_body, profile=False, return_logprob=bench_args.return_logprob, logprob_start_len=bench_args.logprob_start_len, ) time.sleep(0.5) logging.info("\nBenchmark...") result = throughput_test_once( backend_name=bench_args.backend, backend=backend, reqs=input_requests, ignore_eos=not bench_args.disable_ignore_eos, extra_request_body=extra_request_body, profile=bench_args.profile, return_logprob=bench_args.return_logprob, logprob_start_len=bench_args.logprob_start_len, ) backend.shutdown() if bench_args.result_filename: with open(bench_args.result_filename, "a") as fout: fout.write(json.dumps(result) + "\n") print( "\n{s:{c}^{n}}".format(s=" Offline Throughput Benchmark Result ", n=50, c="=") ) print("{:<40} {:<10}".format("Backend:", result["backend"])) print("{:<40} {:<10}".format("Successful requests:", result["successful_requests"])) print("{:<40} {:<10.2f}".format("Benchmark duration (s):", result["total_latency"])) print("{:<40} {:<10}".format("Total input tokens:", result["total_input_tokens"])) print( "{:<40} {:<10}".format("Total generated tokens:", result["total_output_tokens"]) ) print( "{:<40} {:<10.2f}".format( "Last generation throughput (tok/s):", result["last_gen_throughput"] ) ) print( "{:<40} {:<10.2f}".format( "Request throughput (req/s):", result["request_throughput"] ) ) print( "{:<40} {:<10.2f}".format( "Input token throughput (tok/s):", result["input_throughput"] ) ) print( "{:<40} {:<10.2f}".format( "Output token throughput (tok/s):", result["output_throughput"] ) ) print( "{:<40} {:<10.2f}".format( "Total token throughput (tok/s):", result["total_throughput"] ) ) print("=" * 50) return result if __name__ == "__main__": parser = argparse.ArgumentParser() ServerArgs.add_cli_args(parser) BenchArgs.add_cli_args(parser) args = parser.parse_args() # handling ModelScope model downloads if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() in ("true", "1"): if os.path.exists(args.model_path): print(f"Using local model path: {args.model_path}") else: try: from modelscope import snapshot_download print(f"Using ModelScope to download model: {args.model_path}") # download the model and replace args.model_path args.model_path = snapshot_download( args.model_path, ) print(f"Model downloaded to: {args.model_path}") except Exception as e: print(f"ModelScope download failed: {str(e)}") raise e server_args = ServerArgs.from_cli_args(args) bench_args = BenchArgs.from_cli_args(args) logging.basicConfig( level=getattr(logging, server_args.log_level.upper()), format="%(message)s", ) throughput_test(server_args, bench_args) while bench_args.do_not_exit: pass ================================================ FILE: python/sglang/bench_one_batch.py ================================================ """ Benchmark the latency of running a single static batch without a server. This script does not launch a server and uses the low-level APIs. It accepts server arguments (the same as launch_server.py) and benchmark arguments (e.g., batch size, input lengths). # Usage (latency test) ## with dummy weights: python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --load-format dummy ## sweep through multiple data points and store (append) the results in a jsonl file: python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --output-len 32 256 --run-name test_run ## run with profiling: python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --profile ## run with profiling to custom directory: export SGLANG_TORCH_PROFILER_DIR=/root/sglang/profile_log python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 --input-len 256 --profile ## run with CUDA profiler (nsys): nsys profile --force-overwrite=true -o bench_one_batch python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 --input-len 256 --profile --profile-activities CUDA_PROFILER # Usage (correctness test): python -m sglang.bench_one_batch --model-path TinyLlama/TinyLlama-1.1B-Chat-v0.4 --correct ## Reference output (of the correctness test above, can be gpu dependent): input_ids=[[1, 450, 7483, 310, 3444, 338], [1, 450, 7483, 310, 278, 3303, 13187, 290, 338], [1, 20628, 338, 263, 6575, 1460, 2462, 322, 306, 763]] prefill logits (first half): tensor([[-10.0312, -9.5000, 0.8931, ..., -4.9414, -3.2422, -3.3633], [-10.0312, -9.5000, 0.8931, ..., -4.9414, -3.2422, -3.3633], [ -9.1875, -10.2500, 2.7129, ..., -4.3359, -4.0664, -4.1328]], device='cuda:0') prefill logits (final): tensor([[-8.3125, -7.1172, 3.3457, ..., -4.9570, -4.1328, -3.4141], [-8.9141, -9.0156, 4.1445, ..., -4.9922, -4.4961, -4.0781], [-9.6328, -9.0547, 4.0195, ..., -5.3047, -4.7148, -4.4570]], device='cuda:0') ========== Prompt 0 ========== The capital of France is Paris. The capital of the United States is Washington, D.C. ========== Prompt 1 ========== The capital of the United Kindom is London. The capital of the United Kingdom is London. The capital of the ========== Prompt 2 ========== Today is a sunny day and I like to go for a walk in the park. I'm going to the park """ import argparse import copy import dataclasses import itertools import json import logging import multiprocessing import os import time from types import SimpleNamespace from typing import Optional, Tuple import numpy as np import torch import torch.distributed as dist from sglang.srt.configs.model_config import ModelConfig from sglang.srt.distributed.parallel_state import destroy_distributed_environment from sglang.srt.entrypoints.engine import _set_envs_and_config from sglang.srt.layers.moe import initialize_moe_config from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config from sglang.srt.managers.schedule_batch import Req, ScheduleBatch from sglang.srt.managers.scheduler_dp_attn_mixin import prepare_mlp_sync_batch_raw from sglang.srt.mem_cache.base_prefix_cache import EvictParams from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils import ( configure_logger, get_bool_env_var, kill_process_tree, maybe_reindex_device_id, require_mlp_sync, require_mlp_tp_gather, set_gpu_proc_affinity, suppress_other_loggers, ) from sglang.srt.utils.hf_transformers_utils import get_tokenizer def start_profile(profile_activities, profile_record_shapes=False, rank_print=print): """ Abstracted function to start profiling based on profile_activities. Returns profiler object (or None). """ if "CUDA_PROFILER" in profile_activities: try: torch.cuda.cudart().cudaProfilerStart() rank_print("CUDA Profiler started (nsys will begin capturing)") except Exception as e: rank_print(f"Failed to start CUDA profiler: {e}") return None else: activities = [] if "CPU" in profile_activities: activities.append(torch.profiler.ProfilerActivity.CPU) if "GPU" in profile_activities: activities.append(torch.profiler.ProfilerActivity.CUDA) if "XPU" in profile_activities: activities.append(torch.profiler.ProfilerActivity.XPU) if activities: profiler = torch.profiler.profile( activities=activities, with_stack=True, record_shapes=profile_record_shapes, ) profiler.start() return profiler return None def stop_profile( profiler, profile_activities, rank_print=print, save_trace=False, trace_filename=None, stage=None, ): """ Abstracted function to stop profiling based on profile_activities. Optionally saves trace results and prints completion messages. """ if "CUDA_PROFILER" in profile_activities: try: torch.cuda.cudart().cudaProfilerStop() rank_print("CUDA Profiler stopped (nsys should dump traces)") except Exception as e: rank_print(f"Failed to stop CUDA profiler: {e}") elif profiler is not None: profiler.stop() if save_trace: if profiler is not None: if trace_filename: _save_profile_trace_results(profiler, trace_filename) stage_desc = f"for {stage}" if stage else "" rank_print( f"torch profiler chrome trace {stage_desc} saved to {trace_filename}" ) if "CUDA_PROFILER" in profile_activities: rank_print(f"CUDA profiler trace for {stage} completed") @dataclasses.dataclass class BenchArgs: run_name: str = "default" batch_size: Tuple[int] = (1,) input_len: Tuple[int] = (1024,) output_len: Tuple[int] = (16,) prompt_filename: str = "" result_filename: str = "result.jsonl" correctness_test: bool = False # This is only used for correctness test cut_len: int = 4 log_decode_step: int = 0 profile: bool = False profile_record_shapes: bool = False profile_activities: Tuple[str] = ("CPU", "GPU") profile_stage: str = "all" profile_filename_prefix: str = "profile" profile_start_step: Optional[int] = None profile_steps: Optional[int] = None @staticmethod def add_cli_args(parser: argparse.ArgumentParser): parser.add_argument("--run-name", type=str, default=BenchArgs.run_name) parser.add_argument( "--batch-size", type=int, nargs="+", default=BenchArgs.batch_size ) parser.add_argument( "--input-len", type=int, nargs="+", default=BenchArgs.input_len ) parser.add_argument( "--output-len", type=int, nargs="+", default=BenchArgs.output_len ) parser.add_argument( "--prompt-filename", type=str, default=BenchArgs.prompt_filename ) parser.add_argument( "--result-filename", type=str, default=BenchArgs.result_filename ) parser.add_argument("--correctness-test", action="store_true") parser.add_argument("--cut-len", type=int, default=BenchArgs.cut_len) parser.add_argument( "--log-decode-step", type=int, default=BenchArgs.log_decode_step, help="Log decode latency by step, default is set to zero to disable.", ) parser.add_argument("--profile", action="store_true", help="Enable profiling.") parser.add_argument( "--profile-record-shapes", action="store_true", help="Record tensor shapes in profiling results.", ) parser.add_argument( "--profile-activities", type=str, nargs="+", default=["CPU", "GPU"], choices=["CPU", "GPU", "CUDA_PROFILER", "XPU"], help="Profiler activities: CPU, GPU, XPU, CUDA_PROFILER. If CPU/GPU/XPU, use torch profiler. If CUDA_PROFILER, use CUDA profiler.", ) parser.add_argument( "--profile-stage", type=str, default=BenchArgs.profile_stage, choices=["all", "prefill", "decode"], help="Which stage to profile: all, prefill, or decode only.", ) parser.add_argument( "--profile-filename-prefix", type=str, default=BenchArgs.profile_filename_prefix, help="Prefix of the profiling file names. The full profiling result file(s) be " '"[profile_filename_prefix]_batch[batch_size]_input[input_len]_output[output_len].trace.json.gz"', ) parser.add_argument( "--profile-start-step", type=int, default=None, help="Decode step at which to start profiling (0-indexed). If not specified, defaults to output_len // 2.", ) parser.add_argument( "--profile-steps", type=int, default=None, help="Number of decode steps to profile starting from profile-start-step. If not specified, profiles only one step.", ) @classmethod def from_cli_args(cls, args: argparse.Namespace): # use the default value's type to cast the args into correct types. attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)] result = {} for attr, attr_type in attrs: value = getattr(args, attr) # Handle None values - don't try to cast them if value is None or attr_type == type(None): result[attr] = value else: result[attr] = attr_type(value) return cls(**result) def load_model(server_args, port_args, gpu_id, tp_rank): suppress_other_loggers() rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None moe_ep_rank = tp_rank // (server_args.tp_size // server_args.ep_size) model_config = ModelConfig.from_server_args(server_args) model_runner = ModelRunner( model_config=model_config, mem_fraction_static=server_args.mem_fraction_static, gpu_id=gpu_id, tp_rank=tp_rank, tp_size=server_args.tp_size, moe_ep_rank=moe_ep_rank, moe_ep_size=server_args.ep_size, pp_rank=0, pp_size=1, nccl_port=port_args.nccl_port, server_args=server_args, ) rank_print(f"max_total_num_tokens={model_runner.max_total_num_tokens}") tokenizer = get_tokenizer( server_args.tokenizer_path, tokenizer_mode=server_args.tokenizer_mode, trust_remote_code=server_args.trust_remote_code, ) if server_args.tp_size > 1: dist.barrier() return model_runner, tokenizer def prepare_inputs_for_correctness_test(bench_args, tokenizer, custom_prompts): if custom_prompts: custom_input_len = len(custom_prompts) bs = bench_args.batch_size[0] if custom_input_len > bs: logging.warning( f"Custom input size ({custom_input_len}) is larger than batch_size ({bs}). " f"Using the first {bs} prompts." ) custom_prompts = custom_prompts[:bs] prompts = ( custom_prompts if custom_prompts else [ "The capital of France is", "The capital of the United Kindom is", "Today is a sunny day and I like", ] ) input_ids = [tokenizer.encode(p) for p in prompts] sampling_params = SamplingParams( temperature=0, max_new_tokens=BenchArgs.output_len, ) reqs = [] for i in range(len(prompts)): assert len(input_ids[i]) > bench_args.cut_len tmp_input_ids = input_ids[i][: bench_args.cut_len] req = Req( rid=i, origin_input_text=prompts[i], origin_input_ids=tmp_input_ids, sampling_params=sampling_params, ) req.fill_ids = req.origin_input_ids req.logprob_start_len = -1 req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) reqs.append(req) return input_ids, reqs def prepare_extend_inputs_for_correctness_test( bench_args, input_ids, reqs, model_runner ): for i in range(len(reqs)): req: Req = reqs[i] req.fill_ids += input_ids[i][bench_args.cut_len :] req.prefix_indices = model_runner.req_to_token_pool.req_to_token[ i, : bench_args.cut_len ].to(req.prefix_indices.dtype) req.logprob_start_len = -1 req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) return reqs def prepare_synthetic_inputs_for_latency_test( batch_size, input_len, custom_inputs=None ): input_ids = ( custom_inputs if custom_inputs else np.random.randint(0, 10000, (batch_size, input_len), dtype=np.int32) ) sampling_params = SamplingParams( temperature=0, max_new_tokens=BenchArgs.output_len, ) reqs = [] for i in range(len(input_ids)): req = Req( rid=i, origin_input_text="", origin_input_ids=list(input_ids[i]), sampling_params=sampling_params, ) req.fill_ids = req.origin_input_ids req.logprob_start_len = -1 req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices)) reqs.append(req) return reqs class TreeCacheNamespace(SimpleNamespace): def supports_swa(self) -> bool: return False def supports_mamba(self) -> bool: return False def is_chunk_cache(self) -> bool: return False def is_tree_cache(self) -> bool: return not self.is_chunk_cache() def evict(self, params: EvictParams): pass @torch.no_grad def extend(reqs, model_runner): # Create dummy tree_cache for benchmarks (no prefix caching, just allocation) dummy_tree_cache = TreeCacheNamespace( page_size=model_runner.server_args.page_size, device=model_runner.device, token_to_kv_pool_allocator=model_runner.token_to_kv_pool_allocator, ) batch = ScheduleBatch.init_new( reqs=reqs, req_to_token_pool=model_runner.req_to_token_pool, token_to_kv_pool_allocator=model_runner.token_to_kv_pool_allocator, tree_cache=dummy_tree_cache, model_config=model_runner.model_config, enable_overlap=False, spec_algorithm=SpeculativeAlgorithm.NONE, ) batch.prepare_for_extend() _maybe_prepare_mlp_sync_batch(batch, model_runner) model_worker_batch = batch.get_model_worker_batch() forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner) logits_output = model_runner.forward(forward_batch).logits_output next_token_ids = model_runner.sample(logits_output, forward_batch) return next_token_ids, logits_output.next_token_logits, batch @torch.no_grad def decode(input_token_ids, batch, model_runner): batch.output_ids = input_token_ids batch.prepare_for_decode() _maybe_prepare_mlp_sync_batch(batch, model_runner) model_worker_batch = batch.get_model_worker_batch() forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner) logits_output = model_runner.forward(forward_batch).logits_output next_token_ids = model_runner.sample(logits_output, forward_batch) return next_token_ids, logits_output.next_token_logits def _maybe_prepare_mlp_sync_batch(batch: ScheduleBatch, model_runner): if require_mlp_sync(model_runner.server_args): prepare_mlp_sync_batch_raw( batch, dp_size=model_runner.server_args.dp_size, attn_tp_size=1, tp_group=model_runner.tp_group, get_idle_batch=None, disable_cuda_graph=model_runner.server_args.disable_cuda_graph, require_mlp_tp_gather=require_mlp_tp_gather(model_runner.server_args), disable_overlap_schedule=model_runner.server_args.disable_overlap_schedule, offload_tags=set(), ) def _read_prompts_from_file(prompt_file, rank_print): """Read custom prompts from the file specified by `--prompt-filename`.""" if not prompt_file: return [] if not os.path.exists(prompt_file): rank_print( f"Custom prompt file {prompt_file} not found. Using default inputs..." ) return [] with open(prompt_file, "r") as pf: return pf.readlines() def _get_torch_profiler_output_dir(): return os.environ.get("SGLANG_TORCH_PROFILER_DIR", "/tmp") def _create_torch_profiler_filename( profile_filename_prefix, batch_size, input_len, output_len, stage ): output_dir = _get_torch_profiler_output_dir() filename = f"{profile_filename_prefix}_batch{batch_size}_input{input_len}_output{output_len}_{stage}.trace.json.gz" return os.path.join(output_dir, filename) def _save_profile_trace_results(profiler, filename): parent_dir = os.path.dirname(os.path.abspath(filename)) os.makedirs(parent_dir, exist_ok=True) profiler.export_chrome_trace(filename) print( profiler.key_averages(group_by_input_shape=True).table( sort_by="self_cpu_time_total" ) ) def correctness_test( server_args, port_args, bench_args, gpu_id, tp_rank, ): # Configure the logger configure_logger(server_args, prefix=f" TP{tp_rank}") rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None # Load the model model_runner, tokenizer = load_model(server_args, port_args, gpu_id, tp_rank) # Prepare inputs custom_prompts = _read_prompts_from_file(bench_args.prompt_filename, rank_print) input_ids, reqs = prepare_inputs_for_correctness_test( bench_args, tokenizer, custom_prompts ) rank_print(f"\n{input_ids=}\n") if bench_args.cut_len > 0: # Prefill next_token_ids, next_token_logits, batch = extend(reqs, model_runner) rank_print(f"prefill logits (first half): {next_token_logits} \n") # Prepare extend inputs reqs = prepare_extend_inputs_for_correctness_test( bench_args, input_ids, reqs, model_runner ) # Extend (prefill w/ KV cache) next_token_ids, next_token_logits, batch = extend(reqs, model_runner) rank_print(f"prefill logits (final): {next_token_logits} \n") # Decode output_ids = [input_ids[i] + [next_token_ids[i]] for i in range(len(input_ids))] for _ in range(bench_args.output_len[0] - 1): next_token_ids, _ = decode(next_token_ids, batch, model_runner) next_token_ids_list = next_token_ids.tolist() for i in range(len(reqs)): output_ids[i].append(next_token_ids_list[i]) # Print output texts for i in range(len(reqs)): rank_print(f"========== Prompt {i} ==========") rank_print(tokenizer.decode(output_ids[i]), "\n") def synchronize(device): torch.get_device_module(device).synchronize() def latency_test_run_once( run_name, model_runner, rank_print, reqs, batch_size, input_len, output_len, device, log_decode_step, profile, profile_record_shapes, profile_activities, profile_filename_prefix, profile_stage, tp_rank, profile_start_step=None, profile_steps=None, ): max_batch_size = model_runner.max_total_num_tokens // (input_len + output_len) if batch_size > max_batch_size: rank_print( f"skipping ({batch_size}, {input_len}, {output_len}) due to max batch size limit" ) return model_runner.req_to_token_pool.clear() model_runner.token_to_kv_pool_allocator.clear() measurement_results = { "run_name": run_name, "batch_size": batch_size, "input_len": input_len, "output_len": output_len, } tot_latency = 0 profiler = None enable_profile_prefill = profile and profile_stage in ["all", "prefill"] if enable_profile_prefill: profiler = start_profile( profile_activities, profile_record_shapes=profile_record_shapes, rank_print=rank_print, ) synchronize(device) tic = time.perf_counter() next_token_ids, _, batch = extend(reqs, model_runner) synchronize(device) prefill_latency = time.perf_counter() - tic if enable_profile_prefill: trace_filename = _create_torch_profiler_filename( profile_filename_prefix, batch_size, input_len, output_len, "prefill" ) stop_profile( profiler, profile_activities, rank_print=rank_print, save_trace=True, trace_filename=trace_filename, stage="prefill", ) tot_latency += prefill_latency throughput = input_len * batch_size / prefill_latency rank_print( f"Prefill. latency: {prefill_latency:6.5f} s, throughput: {throughput:9.2f} token/s" ) measurement_results["prefill_latency"] = prefill_latency measurement_results["prefill_throughput"] = throughput decode_latencies = [] # Determine profiling start step and end step profile_start = ( profile_start_step if profile_start_step is not None else (output_len // 2) ) profile_end = profile_start + (profile_steps if profile_steps is not None else 1) enable_profile_decode = profile and profile_stage in ["all", "decode"] profiler = None for i in range(output_len - 1): synchronize(device) # Start profiler at the specified step if enable_profile_decode and i == profile_start: profiler = start_profile( profile_activities, profile_record_shapes=profile_record_shapes, rank_print=rank_print, ) tic = time.perf_counter() next_token_ids, _ = decode(next_token_ids, batch, model_runner) synchronize(device) latency = time.perf_counter() - tic # Stop profiler after the specified number of steps if enable_profile_decode and profiler is not None and i >= profile_end - 1: trace_filename = _create_torch_profiler_filename( profile_filename_prefix, batch_size, input_len, output_len, "decode" ) stop_profile( profiler, profile_activities, rank_print=rank_print, save_trace=True, trace_filename=trace_filename, stage="decode", ) profiler = None tot_latency += latency throughput = batch_size / latency decode_latencies.append(latency) if i < 5 or (log_decode_step > 0 and i % log_decode_step == 0): rank_print( f"Decode {i}. Batch size: {batch_size}, latency: {latency:6.5f} s, throughput: {throughput:9.2f} token/s" ) # Record decode timing from 2nd output if output_len > 1: med_decode_latency = np.median(decode_latencies) med_decode_throughput = batch_size / med_decode_latency rank_print( f"Decode. median latency: {med_decode_latency:6.5f} s, median throughput: {med_decode_throughput:9.2f} token/s" ) measurement_results["median_decode_latency"] = med_decode_latency measurement_results["median_decode_throughput"] = med_decode_throughput throughput = (input_len + output_len) * batch_size / tot_latency rank_print( f"Total. latency: {tot_latency:6.3f} s, throughput: {throughput:9.2f} token/s" ) measurement_results["total_latency"] = tot_latency measurement_results["overall_throughput"] = throughput return measurement_results def latency_test( server_args, port_args, bench_args, gpu_id, tp_rank, ): initialize_moe_config(server_args) initialize_fp8_gemm_config(server_args) initialize_fp4_gemm_config(server_args) # Set CPU affinity if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"): set_gpu_proc_affinity( server_args.pp_size, server_args.tp_size, server_args.nnodes, tp_rank ) # Configure the logger configure_logger(server_args, prefix=f" TP{tp_rank}") rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None # Load the model model_runner, tokenizer = load_model(server_args, port_args, gpu_id, tp_rank) # Prepare inputs for warm up reqs = prepare_synthetic_inputs_for_latency_test( bench_args.batch_size[0], bench_args.input_len[0] ) # Warm up rank_print("Warmup ...") latency_test_run_once( bench_args.run_name, model_runner, rank_print, reqs, bench_args.batch_size[0], bench_args.input_len[0], min(32, bench_args.output_len[0]), # shorter decoding to speed up the warmup server_args.device, log_decode_step=0, profile=False, profile_record_shapes=False, profile_activities=("CPU", "GPU"), profile_filename_prefix="", profile_stage="all", tp_rank=tp_rank, profile_start_step=None, profile_steps=None, ) rank_print("Benchmark ...") custom_inputs = _read_prompts_from_file(bench_args.prompt_filename, rank_print) custom_inputs = [tokenizer.encode(p.strip()) for p in custom_inputs] custom_input_len = len(custom_inputs) # Run the sweep result_list = [] for bs, il, ol in itertools.product( bench_args.batch_size, bench_args.input_len, bench_args.output_len ): bs_aligned_inputs = [] if custom_inputs: if custom_input_len == bs: bs_aligned_inputs = custom_inputs elif custom_input_len > bs: rank_print( f"Custom input size ({custom_input_len}) is larger than batch_size ({bs}). " f"Using the first {bs} prompts." ) bs_aligned_inputs = copy.deepcopy(custom_inputs[:bs]) else: rank_print( f"Custom input size ({custom_input_len}) is smaller than batch_size ({bs}). " f"Pad to the desired batch_size with the last prompt." ) bs_aligned_inputs = copy.deepcopy(custom_inputs) bs_aligned_inputs.extend( [bs_aligned_inputs[-1]] * (bs - custom_input_len) ) reqs = prepare_synthetic_inputs_for_latency_test(bs, il, bs_aligned_inputs) ret = latency_test_run_once( bench_args.run_name, model_runner, rank_print, reqs, bs, il, ol, server_args.device, bench_args.log_decode_step, bench_args.profile if tp_rank == 0 else None, bench_args.profile_record_shapes if tp_rank == 0 else None, bench_args.profile_activities, bench_args.profile_filename_prefix, bench_args.profile_stage, tp_rank, bench_args.profile_start_step, bench_args.profile_steps, ) if ret is not None: result_list.append(ret) # Write results in jsonlines format on rank 0. if tp_rank == 0 and bench_args.result_filename: with open(bench_args.result_filename, "a") as fout: for result in result_list: fout.write(json.dumps(result) + "\n") if server_args.tp_size > 1: destroy_distributed_environment() def main(server_args, bench_args): server_args.cuda_graph_max_bs = max(bench_args.batch_size) _set_envs_and_config(server_args) if server_args.model_path: if bench_args.correctness_test: work_func = correctness_test else: work_func = latency_test else: raise ValueError( "Provide --model-path for running the tests or " "provide --result-filename for plotting the results" ) port_args = PortArgs.init_new(server_args) if server_args.tp_size == 1: work_func(server_args, port_args, bench_args, 0, 0) else: workers = [] for tp_rank in range(server_args.tp_size): with maybe_reindex_device_id(tp_rank) as gpu_id: proc = multiprocessing.Process( target=work_func, args=( server_args, port_args, bench_args, gpu_id, tp_rank, ), ) proc.start() workers.append(proc) for proc in workers: proc.join() proc.terminate() if __name__ == "__main__": parser = argparse.ArgumentParser() ServerArgs.add_cli_args(parser) BenchArgs.add_cli_args(parser) args = parser.parse_args() server_args = ServerArgs.from_cli_args(args) bench_args = BenchArgs.from_cli_args(args) logging.basicConfig( level=getattr(logging, server_args.log_level.upper()), format="%(message)s", ) try: main(server_args, bench_args) finally: if server_args.tp_size != 1: kill_process_tree(os.getpid(), include_parent=False) ================================================ FILE: python/sglang/bench_one_batch_server.py ================================================ """ Benchmark the latency of running a single batch with a server. This script launches a server and uses the HTTP interface. It accepts server arguments (the same as launch_server.py) and benchmark arguments (e.g., batch size, input lengths). Usage: python3 -m sglang.bench_one_batch_server --model meta-llama/Meta-Llama-3.1-8B --batch-size 1 16 64 --input-len 1024 --output-len 8 python3 -m sglang.bench_one_batch_server --model None --base-url http://localhost:30000 --batch-size 16 --input-len 1024 --output-len 8 python3 -m sglang.bench_one_batch_server --model None --base-url http://localhost:30000 --batch-size 16 --input-len 1024 --output-len 8 --show-report --profile --profile-by-stage python3 -m sglang.bench_one_batch_server --model None --base-url http://localhost:30000 --batch-size 16 --input-len 1024 --output-len 8 --output-path results.json --profile """ import argparse from sglang.srt.server_args import ServerArgs from sglang.test.bench_one_batch_server_internal import ( BenchArgs, run_benchmark_internal, ) from sglang.test.nightly_bench_utils import save_results_as_pydantic_models def run_benchmark(server_args: ServerArgs, bench_args: BenchArgs): results, server_info = run_benchmark_internal(server_args, bench_args) # Save results as pydantic models in the JSON format if bench_args.pydantic_result_filename: save_results_as_pydantic_models( results, pydantic_result_filename=bench_args.pydantic_result_filename, model_path=server_args.model_path, server_args=bench_args.server_args_for_metrics, ) return results, server_info if __name__ == "__main__": parser = argparse.ArgumentParser() ServerArgs.add_cli_args(parser) BenchArgs.add_cli_args(parser) args = parser.parse_args() server_args = ServerArgs.from_cli_args(args) bench_args = BenchArgs.from_cli_args(args) run_benchmark(server_args, bench_args) ================================================ FILE: python/sglang/bench_serving.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/backend_request_func.py # Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/benchmark_serving.py """ Benchmark online serving with dynamic requests. Usage: python3 -m sglang.bench_serving --backend sglang --num-prompt 10 python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 3000 --random-input 1024 --random-output 1024 --random-range-ratio 0.5 """ import argparse import asyncio import copy import importlib.util import json import os import random import shutil import sys import time import traceback import uuid import warnings from argparse import ArgumentParser from copy import deepcopy from dataclasses import dataclass, field, replace from datetime import datetime from pathlib import Path from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Tuple, Union import aiohttp import numpy as np import requests from tqdm.asyncio import tqdm from transformers import AutoTokenizer, PreTrainedTokenizerBase from sglang.benchmark.datasets import DatasetRow, get_dataset from sglang.benchmark.datasets.mooncake import get_mooncake_request_over_time from sglang.benchmark.utils import ( get_tokenizer, parse_custom_headers, remove_prefix, set_ulimit, ) _ROUTING_KEY_HEADER = "X-SMG-Routing-Key" _EMBEDDING_UNSUPPORTED_DATASETS = {"image", "mmmu", "mooncake"} TERM_PLOTLIB_AVAILABLE = (importlib.util.find_spec("termplotlib") is not None) and ( shutil.which("gnuplot") is not None ) global args # don't want to import sglang package here def _get_bool_env_var(name: str, default: str = "false") -> bool: value = os.getenv(name, default) return value.lower() in ("true", "1") def _create_bench_client_session(): # When the pressure is big, the read buffer could be full before aio thread read # the content. We increase the read_bufsize from 64K to 10M. # Define constants for timeout and buffer size for clarity and maintainability BENCH_AIOHTTP_TIMEOUT_SECONDS = 6 * 60 * 60 # 6 hours BENCH_AIOHTTP_READ_BUFSIZE_BYTES = 10 * 1024**2 # 10 MB aiohttp_timeout = aiohttp.ClientTimeout(total=BENCH_AIOHTTP_TIMEOUT_SECONDS) return aiohttp.ClientSession( timeout=aiohttp_timeout, read_bufsize=BENCH_AIOHTTP_READ_BUFSIZE_BYTES ) @dataclass class RequestFuncInput: prompt: Union[str, List[str], List[Dict[str, str]]] api_url: str prompt_len: int output_len: int model: str lora_name: str image_data: Optional[List[str]] extra_request_body: Dict[str, Any] timestamp: Optional[float] = None routing_key: Optional[str] = None @dataclass class RequestFuncOutput: generated_text: str = "" success: bool = False latency: float = 0.0 ttft: float = 0.0 # Time to first token itl: List[float] = field(default_factory=list) # List of inter-token latencies text_chunks: List[str] = field(default_factory=list) prompt_len: int = 0 error: str = "" output_len: int = 0 start_time: float = 0.0 @staticmethod def init_new(request_func_input: RequestFuncInput): output = RequestFuncOutput() output.prompt_len = request_func_input.prompt_len return output def get_auth_headers() -> Dict[str, str]: openai_api_key = os.environ.get("OPENAI_API_KEY") if openai_api_key: return {"Authorization": f"Bearer {openai_api_key}"} else: api_key = os.environ.get("API_KEY") if api_key: return {"Authorization": f"{api_key}"} return {} def get_request_headers() -> Dict[str, str]: headers = get_auth_headers() if h := getattr(args, "header", None): headers.update(parse_custom_headers(h)) return headers def wait_for_endpoint(url: str, timeout_sec: int = 60) -> bool: """Wait for the server to become ready by polling the given URL.""" print(f"Waiting up to {timeout_sec}s for {url} to become ready...") start_time = time.perf_counter() headers = get_auth_headers() while True: try: response = requests.get(url, headers=headers, timeout=5) if response.status_code == 200: elapsed = time.perf_counter() - start_time print(f"Server ready in {elapsed:.1f}s.") return True except requests.exceptions.RequestException: pass elapsed = time.perf_counter() - start_time if elapsed >= timeout_sec: print(f"Server did not become ready within {timeout_sec}s timeout.") return False time.sleep(1) # trt llm does not support ignore_eos # https://github.com/triton-inference-server/tensorrtllm_backend/issues/505 async def async_request_trt_llm( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url assert api_url.endswith("generate_stream") async with _create_bench_client_session() as session: payload = { "accumulate_tokens": True, "text_input": request_func_input.prompt, "temperature": 0.000001, "top_p": 1.0, "max_tokens": request_func_input.output_len, "stream": True, "min_length": request_func_input.output_len, "end_id": 1048576, **request_func_input.extra_request_body, } if args.disable_ignore_eos: del payload["min_length"] del payload["end_id"] output = RequestFuncOutput.init_new(request_func_input) ttft = 0.0 st = time.perf_counter() most_recent_timestamp = st try: async with session.post(url=api_url, json=payload) as response: if response.status == 200: async for chunk_bytes in response.content: chunk_bytes = chunk_bytes.strip() if not chunk_bytes: continue chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data:") data = json.loads(chunk) output.generated_text += data["text_output"] timestamp = time.perf_counter() # First token if ttft == 0.0: ttft = timestamp - st output.ttft = ttft # Decoding phase else: output.itl.append(timestamp - most_recent_timestamp) most_recent_timestamp = timestamp output.latency = most_recent_timestamp - st output.success = True output.output_len = request_func_input.output_len else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) if pbar: pbar.update(1) return output # set ignore_eos True by default async def async_request_openai_completions( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url assert api_url.endswith( "completions" ), "OpenAI Completions API URL must end with 'completions'." prompt = request_func_input.prompt async with _create_bench_client_session() as session: # Build payload with defaults that can be overridden by extra_request_body payload = { "model": request_func_input.model, "prompt": prompt, "best_of": 1, "max_tokens": request_func_input.output_len, "stream": not args.disable_stream, } # Add temperature default only if not specified in extra_request_body if "temperature" not in request_func_input.extra_request_body: payload["temperature"] = 0.0 # Add ignore_eos default only if not specified in extra_request_body if "ignore_eos" not in request_func_input.extra_request_body: payload["ignore_eos"] = not args.disable_ignore_eos if args.return_logprob and args.top_logprobs_num > 0: payload["logprobs"] = args.top_logprobs_num # Merge in extra parameters - these will override defaults if present payload.update(request_func_input.extra_request_body) # hack to accommodate different LoRA conventions between SGLang and vLLM. if request_func_input.lora_name: payload["model"] = request_func_input.lora_name payload["lora_path"] = request_func_input.lora_name if request_func_input.image_data: payload.update({"image_data": request_func_input.image_data}) headers = get_request_headers() if request_func_input.routing_key: headers[_ROUTING_KEY_HEADER] = request_func_input.routing_key output = RequestFuncOutput.init_new(request_func_input) generated_text = "" output_len = request_func_input.output_len ttft = 0.0 st = time.perf_counter() output.start_time = st most_recent_timestamp = st try: async with session.post( url=api_url, json=payload, headers=headers ) as response: if response.status == 200: async for chunk_bytes in response.content: chunk_bytes = chunk_bytes.strip() if not chunk_bytes: continue chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") latency = time.perf_counter() - st if chunk == "[DONE]": pass else: data = json.loads(chunk) # NOTE: Some completion API might have a last # usage summary response without a token so we # want to check a token was generated if data["choices"][0]["text"]: timestamp = time.perf_counter() # First token if ttft == 0.0: ttft = time.perf_counter() - st output.ttft = ttft # Decoding phase else: output.text_chunks.append( data["choices"][0]["text"] ) output.itl.append(timestamp - most_recent_timestamp) most_recent_timestamp = timestamp generated_text += data["choices"][0]["text"] output_len = (data.get("usage") or {}).get( "completion_tokens", output_len ) output.generated_text = generated_text output.success = True output.latency = latency output.output_len = output_len else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) if pbar: pbar.update(1) return output async def async_request_openai_chat_completions( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: """Makes a request to the OpenAI Chat Completions API. Handles both streaming and non-streaming responses, including support for image data in messages. Calculates and returns various performance metrics. Args: request_func_input: Input parameters for the request. pbar: Optional tqdm progress bar to update. Returns: RequestFuncOutput: Output of the request, including generated text, latency, TTFT, ITL, and success status. """ api_url = request_func_input.api_url assert api_url.endswith( "chat/completions" ), "OpenAI Chat Completions API URL must end with 'chat/completions'." # TODO put it to other functions when `pbar` logic is refactored if getattr(args, "print_requests", False): rid = str(uuid.uuid4()) input_partial = deepcopy(request_func_input) input_partial.prompt = "..." request_start_time = time.time() print( f'rid={rid} time={request_start_time} message="request start" request_func_input="{str(input_partial)}"' ) if isinstance(request_func_input.prompt, list): messages = request_func_input.prompt elif request_func_input.image_data: # Build multi-image content: a list of image_url entries followed by the text content_items = [ { "type": "image_url", "image_url": {"url": img_url}, } for img_url in request_func_input.image_data ] content_items.append({"type": "text", "text": request_func_input.prompt}) messages = [ { "role": "user", "content": content_items, }, ] else: messages = [{"role": "user", "content": request_func_input.prompt}] async with _create_bench_client_session() as session: # Build payload with defaults that can be overridden by extra_request_body payload = { "model": request_func_input.model, "messages": messages, "max_completion_tokens": request_func_input.output_len, "stream": not args.disable_stream, } # Add temperature default only if not specified in extra_request_body if "temperature" not in request_func_input.extra_request_body: payload["temperature"] = 0.0 # Add ignore_eos default only if not specified in extra_request_body # Default to False for more realistic behavior (respect EOS tokens) if "ignore_eos" not in request_func_input.extra_request_body: payload["ignore_eos"] = not args.disable_ignore_eos # Merge in extra parameters (tools, temperature, top_p, etc.) # These will override defaults if present payload.update(request_func_input.extra_request_body) # hack to accommodate different LoRA conventions between SGLang and vLLM. if request_func_input.lora_name: payload["model"] = request_func_input.lora_name payload["lora_path"] = request_func_input.lora_name headers = get_request_headers() if request_func_input.routing_key: headers[_ROUTING_KEY_HEADER] = request_func_input.routing_key output = RequestFuncOutput.init_new(request_func_input) generated_text = "" output_len = request_func_input.output_len ttft = 0.0 st = time.perf_counter() output.start_time = st most_recent_timestamp = st try: async with session.post( url=api_url, json=payload, headers=headers ) as response: if response.status == 200: if args.disable_stream: # Non-streaming response response_json = await response.json() output.generated_text = response_json["choices"][0]["message"][ "content" ] output.success = True output.latency = time.perf_counter() - st output.ttft = ( output.latency ) # For non-streaming, TTFT = total latency output.output_len = response_json.get("usage", {}).get( "completion_tokens", output_len ) else: # Streaming response async for chunk_bytes in response.content: chunk_bytes = chunk_bytes.strip() if not chunk_bytes: continue chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") latency = time.perf_counter() - st if chunk == "[DONE]": pass else: data = json.loads(chunk) # Check if this chunk contains content delta = data.get("choices", [{}])[0].get("delta", {}) content = delta.get("content", "") if content: timestamp = time.perf_counter() # First token if ttft == 0.0: ttft = timestamp - st output.ttft = ttft # Decoding phase else: output.text_chunks.append(content) output.itl.append( timestamp - most_recent_timestamp ) most_recent_timestamp = timestamp generated_text += content # Check for usage info in final chunk output_len = (data.get("usage") or {}).get( "completion_tokens", output_len ) output.generated_text = generated_text output.success = True output.latency = latency output.output_len = output_len else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) # TODO put it to other functions when `pbar` logic is refactored if getattr(args, "print_requests", False): curr_t = time.time() output_partial = deepcopy(output) output_partial.generated_text = "..." print( f'rid={rid} time={curr_t} time_delta={curr_t - request_start_time} message="request end" output="{str(output_partial)}"' ) if pbar: pbar.update(1) return output async def async_request_truss( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url prompt = request_func_input.prompt async with _create_bench_client_session() as session: payload = { "model": request_func_input.model, "prompt": prompt, "temperature": 0.0, "best_of": 1, "max_tokens": request_func_input.output_len, "stream": not args.disable_stream, "ignore_eos": not args.disable_ignore_eos, **request_func_input.extra_request_body, } headers = get_request_headers() output = RequestFuncOutput.init_new(request_func_input) generated_text = "" ttft = 0.0 st = time.perf_counter() most_recent_timestamp = st try: async with session.post( url=api_url, json=payload, headers=headers ) as response: if response.status == 200: async for chunk_bytes in response.content: chunk_bytes = chunk_bytes.strip() if not chunk_bytes: continue chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") latency = time.perf_counter() - st if chunk == "[DONE]": pass else: data = json.loads(chunk) # NOTE: Some completion API might have a last # usage summary response without a token so we # want to check a token was generated if data["choices"][0]["text"]: timestamp = time.perf_counter() # First token if ttft == 0.0: ttft = time.perf_counter() - st output.ttft = ttft # Decoding phase else: output.itl.append(timestamp - most_recent_timestamp) most_recent_timestamp = timestamp generated_text += data["choices"][0]["text"] output.generated_text = generated_text output.success = True output.latency = latency output.output_len = request_func_input.output_len else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) if pbar: pbar.update(1) return output async def async_request_sglang_generate( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url prompt = request_func_input.prompt async with _create_bench_client_session() as session: payload = { ("text" if isinstance(prompt, str) else "input_ids"): prompt, "sampling_params": { "temperature": 0.0, "max_new_tokens": request_func_input.output_len, "ignore_eos": not args.disable_ignore_eos, }, "stream": not args.disable_stream, "lora_path": request_func_input.lora_name, "return_logprob": args.return_logprob, "return_routed_experts": args.return_routed_experts, "logprob_start_len": args.logprob_start_len, **request_func_input.extra_request_body, } if args.top_logprobs_num > 0: payload["top_logprobs_num"] = args.top_logprobs_num if args.token_ids_logprob is not None: payload["token_ids_logprob"] = args.token_ids_logprob # Add image data if available (list of image urls/base64) if request_func_input.image_data: payload["image_data"] = request_func_input.image_data headers = get_request_headers() if request_func_input.routing_key: headers[_ROUTING_KEY_HEADER] = request_func_input.routing_key output = RequestFuncOutput.init_new(request_func_input) generated_text = "" output_len = request_func_input.output_len ttft = 0.0 st = time.perf_counter() output.start_time = st most_recent_timestamp = st last_output_len = 0 try: async with session.post( url=api_url, json=payload, headers=headers ) as response: if response.status == 200: async for chunk_bytes in response.content: chunk_bytes = chunk_bytes.strip() if not chunk_bytes: continue chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") latency = time.perf_counter() - st if chunk == "[DONE]": pass else: data = json.loads(chunk) # NOTE: Some completion API might have a last # usage summary response without a token so we # want to check a token was generated if "text" in data and data["text"]: timestamp = time.perf_counter() generated_text = data["text"] output_len = data["meta_info"]["completion_tokens"] # First token if ttft == 0.0: ttft = time.perf_counter() - st output.ttft = ttft # Decoding phase else: num_new_tokens = output_len - last_output_len if num_new_tokens == 0: continue chunk_gap = timestamp - most_recent_timestamp adjust_itl = chunk_gap / num_new_tokens output.itl.extend([adjust_itl] * num_new_tokens) most_recent_timestamp = timestamp last_output_len = output_len output.generated_text = generated_text output.success = True output.latency = latency output.output_len = output_len else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) print(f"{output.error=}") if pbar: pbar.update(1) return output async def async_request_openai_embeddings( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url async with _create_bench_client_session() as session: payload = { "input": request_func_input.prompt, "model": request_func_input.model, } if request_func_input.lora_name: payload["model"] = request_func_input.lora_name payload["lora_path"] = request_func_input.lora_name payload.update(request_func_input.extra_request_body) headers = get_request_headers() if request_func_input.routing_key: headers[_ROUTING_KEY_HEADER] = request_func_input.routing_key output = RequestFuncOutput.init_new(request_func_input) st = time.perf_counter() output.start_time = st try: async with session.post( url=api_url, json=payload, headers=headers ) as response: if response.status == 200: await response.json() output.latency = time.perf_counter() - st output.success = True output.output_len = 0 else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) if pbar: pbar.update(1) return output async def async_request_gserver( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: raise NotImplementedError() async def async_request_profile(api_url: str) -> RequestFuncOutput: async with _create_bench_client_session() as session: output = RequestFuncOutput() try: if api_url.endswith("/start_profile"): num_steps = getattr(args, "profile_num_steps", None) profile_by_stage = getattr(args, "profile_by_stage", None) if profile_by_stage and num_steps is None: num_steps = 5 output_dir = getattr(args, "profile_output_dir", None) if output_dir is None: output_dir = os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp") output_dir = Path(os.path.abspath(os.path.normpath(output_dir))) / str( time.time() ) output_dir.mkdir(exist_ok=True, parents=True) output_dir = str(output_dir) body = { "activities": getattr(args, "profile_activities", []), "num_steps": num_steps, "profile_by_stage": profile_by_stage, "profile_stages": getattr(args, "profile_stages", None), "output_dir": output_dir, "profile_prefix": getattr(args, "profile_prefix", None), } else: # stop_profile doesn't need any parameters body = {} print(f"async_request_profile {api_url=} {body=}") # Add optional profiling parameters if provided if ( hasattr(args, "profile_start_step") and args.profile_start_step is not None ): body["start_step"] = str(args.profile_start_step) if hasattr(args, "profile_steps") and args.profile_steps is not None: body["num_steps"] = str(args.profile_steps) async with session.post(url=api_url, json=body) as response: if response.status == 200: output.success = True else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) return output def _build_profile_urls( profile_prefill_url: Optional[List[str]], profile_decode_url: Optional[List[str]], ) -> List[Tuple[str, str]]: """Build profile URLs list from prefill/decode URL arguments. Returns: List of (worker_type, url) tuples. e.g., [("Prefill-0", "http://..."), ("Decode-0", "http://...")] """ profile_urls = [] if profile_prefill_url: for idx, url in enumerate(profile_prefill_url): profile_urls.append((f"Prefill-{idx}", url)) if profile_decode_url: for idx, url in enumerate(profile_decode_url): profile_urls.append((f"Decode-{idx}", url)) return profile_urls async def _call_profile_pd(profile_urls: List[Tuple[str, str]], mode: str) -> None: """Call profile endpoint (start/stop) on PD separated workers. Args: profile_urls: List of (worker_type, url) tuples mode: "start" or "stop" """ endpoint = "/start_profile" if mode == "start" else "/stop_profile" action = "Starting" if mode == "start" else "Stopping" action_past = "started" if mode == "start" else "stopped" print(f"{action} profiler...") for worker_type, url in profile_urls: profile_output = await async_request_profile(api_url=url + endpoint) if profile_output.success: print(f"Profiler {action_past} for {worker_type} worker at {url}") else: print( f"Failed to {mode} profiler for {worker_type} worker at {url}: {profile_output.error}" ) ASYNC_REQUEST_FUNCS = { "sglang": async_request_sglang_generate, "sglang-native": async_request_sglang_generate, "sglang-oai": async_request_openai_completions, "sglang-oai-chat": async_request_openai_chat_completions, "sglang-embedding": async_request_openai_embeddings, "vllm": async_request_openai_completions, "vllm-chat": async_request_openai_chat_completions, "lmdeploy": async_request_openai_completions, "lmdeploy-chat": async_request_openai_chat_completions, "trt": async_request_trt_llm, "gserver": async_request_gserver, "truss": async_request_truss, } @dataclass class BenchmarkMetrics: completed: int total_input: int total_input_text: int total_input_vision: int total_output: int total_output_retokenized: int request_throughput: float input_throughput: float output_throughput: float output_throughput_retokenized: float total_throughput: float total_throughput_retokenized: float mean_ttft_ms: float median_ttft_ms: float std_ttft_ms: float p99_ttft_ms: float mean_tpot_ms: float median_tpot_ms: float std_tpot_ms: float p99_tpot_ms: float mean_itl_ms: float median_itl_ms: float std_itl_ms: float p95_itl_ms: float p99_itl_ms: float max_itl_ms: float mean_e2e_latency_ms: float median_e2e_latency_ms: float std_e2e_latency_ms: float p90_e2e_latency_ms: float p99_e2e_latency_ms: float concurrency: float max_output_tokens_per_s: float = 0.0 max_concurrent_requests: int = 0 async def get_request( input_requests: List[DatasetRow], request_rate: float, use_trace_timestamps: bool = False, slowdown_factor: float = 1.0, ) -> AsyncGenerator[DatasetRow, None]: if use_trace_timestamps: print( f"Using trace timestamps for request generation with slowdown factor {slowdown_factor}." ) # Sort requests by timestamp for correct replay input_requests.sort(key=lambda r: r.timestamp) start_time = time.perf_counter() trace_start_time_ms = input_requests[0].timestamp if input_requests else 0 for request in input_requests: trace_time_s = (request.timestamp - trace_start_time_ms) / 1000.0 target_arrival_time = start_time + (trace_time_s * slowdown_factor) sleep_duration = target_arrival_time - time.perf_counter() if sleep_duration > 0: await asyncio.sleep(sleep_duration) yield request else: input_requests_iter = iter(input_requests) for request in input_requests_iter: yield request if request_rate == float("inf"): # If the request rate is infinity, then we don't need to wait. continue # Sample the request interval from the exponential distribution. interval = np.random.exponential(1.0 / request_rate) # The next request will be sent after the interval. await asyncio.sleep(interval) def calculate_metrics( input_requests: Optional[List[DatasetRow]], outputs: List[RequestFuncOutput], dur_s: float, tokenizer: PreTrainedTokenizerBase, backend: str, accept_length: Optional[float] = None, plot_throughput: bool = False, ) -> Tuple[BenchmarkMetrics, List[int]]: output_lens: List[int] = [] retokenized_output_lens: List[int] = [] total_input = 0 total_input_text = 0 total_input_vision = 0 completed = 0 itls: List[float] = [] tpots: List[float] = [] ttfts: List[float] = [] e2e_latencies: List[float] = [] retokenized_itls: List[float] = [] use_retokenized_itl = ( accept_length is not None and accept_length > 0 and backend in ("sglang-oai", "sglang-oai-chat") ) for i in range(len(outputs)): if outputs[i].success: output_len = outputs[i].output_len output_lens.append(output_len) retokenized_output_len = len( tokenizer.encode(outputs[i].generated_text, add_special_tokens=False) ) retokenized_output_lens.append(retokenized_output_len) if input_requests is not None: total_input += input_requests[i].prompt_len total_input_text += input_requests[i].text_prompt_len total_input_vision += input_requests[i].vision_prompt_len if output_len > 1: tpots.append((outputs[i].latency - outputs[i].ttft) / (output_len - 1)) if use_retokenized_itl: for k, itl in enumerate(outputs[i].itl): num_tokens = len( tokenizer.encode( outputs[i].text_chunks[k], add_special_tokens=False ) ) adjusted_itl = itl / num_tokens retokenized_itls.extend([adjusted_itl] * num_tokens) else: itls += outputs[i].itl ttfts.append(outputs[i].ttft) e2e_latencies.append(outputs[i].latency) completed += 1 else: output_lens.append(0) retokenized_output_lens.append(0) if completed == 0: warnings.warn( "All requests failed. This is likely due to a misconfiguration " "on the benchmark arguments.", stacklevel=2, ) max_output_tokens_per_s = 0.0 max_concurrent_requests = 0 successful_outputs = [output for output in outputs if output.success] if successful_outputs: min_start_time = min(output.start_time for output in successful_outputs) max_end_time = max( output.start_time + output.latency for output in successful_outputs ) duration_seconds = int(np.ceil(max_end_time - min_start_time)) + 1 tokens_per_second = np.zeros(duration_seconds) concurrent_requests_per_second = np.zeros(duration_seconds) for output in outputs: if not output.success: continue token_times = [output.start_time + output.ttft] current_time = token_times[0] for itl_value in output.itl: current_time += itl_value token_times.append(current_time) for token_time in token_times: second_bucket = int(token_time - min_start_time) if 0 <= second_bucket < duration_seconds: tokens_per_second[second_bucket] += 1 request_start_second = int(output.start_time - min_start_time) request_end_second = int( (output.start_time + output.latency) - min_start_time ) for second in range( request_start_second, min(request_end_second + 1, duration_seconds) ): concurrent_requests_per_second[second] += 1 if len(tokens_per_second) > 0: max_output_tokens_per_s = float(np.max(tokens_per_second)) max_concurrent_requests = int(np.max(concurrent_requests_per_second)) if plot_throughput: if TERM_PLOTLIB_AVAILABLE: import termplotlib as tpl fig = tpl.figure() fig.plot( np.arange(len(tokens_per_second)), tokens_per_second, title="Output tokens per second", xlabel="Time (s)", ) fig.plot( np.arange(len(concurrent_requests_per_second)), concurrent_requests_per_second, title="Concurrent requests per second", xlabel="Time (s)", ) fig.show() else: print("tip: install termplotlib and gnuplot to plot the metrics") itls = retokenized_itls if use_retokenized_itl else itls metrics = BenchmarkMetrics( completed=completed, total_input=total_input, total_input_text=total_input_text, total_input_vision=total_input_vision, total_output=sum(output_lens), total_output_retokenized=sum(retokenized_output_lens), request_throughput=completed / dur_s, input_throughput=total_input / dur_s, output_throughput=sum(output_lens) / dur_s, output_throughput_retokenized=sum(retokenized_output_lens) / dur_s, total_throughput=(total_input + sum(output_lens)) / dur_s, total_throughput_retokenized=(total_input + sum(retokenized_output_lens)) / dur_s, mean_ttft_ms=np.mean(ttfts or 0) * 1000, # ttfts is empty if streaming is not supported by backend median_ttft_ms=np.median(ttfts or 0) * 1000, std_ttft_ms=np.std(ttfts or 0) * 1000, p99_ttft_ms=np.percentile(ttfts or 0, 99) * 1000, mean_tpot_ms=np.mean(tpots or 0) * 1000, median_tpot_ms=np.median(tpots or 0) * 1000, std_tpot_ms=np.std(tpots or 0) * 1000, p99_tpot_ms=np.percentile(tpots or 0, 99) * 1000, mean_itl_ms=np.mean(itls or 0) * 1000, median_itl_ms=np.median(itls or 0) * 1000, std_itl_ms=np.std(itls or 0) * 1000, p95_itl_ms=np.percentile(itls or 0, 95) * 1000, p99_itl_ms=np.percentile(itls or 0, 99) * 1000, max_itl_ms=np.max(itls or 0) * 1000, mean_e2e_latency_ms=np.mean(e2e_latencies) * 1000, median_e2e_latency_ms=np.median(e2e_latencies) * 1000, std_e2e_latency_ms=np.std(e2e_latencies) * 1000, p90_e2e_latency_ms=np.percentile(e2e_latencies, 90) * 1000, p99_e2e_latency_ms=np.percentile(e2e_latencies, 99) * 1000, concurrency=np.sum(e2e_latencies) / dur_s, max_output_tokens_per_s=max_output_tokens_per_s, max_concurrent_requests=max_concurrent_requests, ) return metrics, output_lens MULTI_TURN_BACKENDS = {"sglang-oai-chat", "vllm-chat", "lmdeploy-chat"} def wrap_multi_turn_request_func(request_func: Callable, backend: str) -> Callable: assert ( backend in MULTI_TURN_BACKENDS ), f"Multi-turn only supports chat backends: {MULTI_TURN_BACKENDS}, got {backend}" async def f( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> List[RequestFuncOutput]: prompts: List[str] = request_func_input.prompt prev_messages: List[Dict[str, str]] = [] outputs = [] for round_index in range(len(prompts)): prev_messages.append({"role": "user", "content": prompts[round_index]}) inner_input = replace( copy.deepcopy(request_func_input), prompt=copy.deepcopy(prev_messages) ) output = await request_func( inner_input, pbar=pbar if round_index == len(prompts) - 1 else None ) outputs.append(output) prev_messages.append( {"role": "assistant", "content": output.generated_text} ) return outputs return f async def benchmark( backend: str, api_url: str, base_url: str, model_id: str, tokenizer: PreTrainedTokenizerBase, input_requests: List[DatasetRow], request_rate: float, max_concurrency: Optional[int], disable_tqdm: bool, lora_names: List[str], lora_request_distribution: Optional[str], lora_zipf_alpha: Optional[float], extra_request_body: Dict[str, Any], profile: bool, pd_separated: bool = False, flush_cache: bool = False, warmup_requests: int = 1, use_trace_timestamps: bool = False, mooncake_slowdown_factor=1.0, mooncake_num_rounds=1, profile_prefill_url: Optional[List[str]] = None, profile_decode_url: Optional[List[str]] = None, ): if backend in ASYNC_REQUEST_FUNCS: request_func = ASYNC_REQUEST_FUNCS[backend] else: raise ValueError(f"Unknown backend: {backend}") # Check for multi-turn: prompt is a list of strings (not OpenAI messages dicts) # Multi-turn format: ["turn1", "turn2", ...] - list of strings # OpenAI format: [{"role": "user", "content": "..."}, ...] - list of dicts first_prompt = input_requests[0].prompt is_multi_turn = ( isinstance(first_prompt, list) and len(first_prompt) > 0 and isinstance(first_prompt[0], str) ) if is_multi_turn: request_func = wrap_multi_turn_request_func(request_func, backend=backend) # Limit concurrency # From https://github.com/vllm-project/vllm/pull/9390 semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None async def limited_request_func(request_func_input, pbar): if semaphore is None: return await request_func(request_func_input=request_func_input, pbar=pbar) async with semaphore: return await request_func(request_func_input=request_func_input, pbar=pbar) # Warmup print(f"Starting warmup with {warmup_requests} sequences...") # Handle the data structure difference for the warmup request if args.dataset_name == "mooncake": # For mooncake, input_requests is a list of dicts. # We need to build a temporary DatasetRow for the warmup phase. warmup_record = input_requests[0] # Build prompt from hash_ids, just like in the async generator hash_ids = warmup_record.get("hash_ids", []) prompt_text = "" for hash_id in hash_ids: prompt_text += f"{hash_id}" + " ".join(["hi"] * 512) prompt_text += "Can you tell me a detailed story in 1000 words?" output_len = warmup_record.get("output_length", 32) prompt_len = len(tokenizer.encode(prompt_text)) # Create a temporary DatasetRow object for warmup test_request = DatasetRow( prompt=prompt_text, prompt_len=prompt_len, output_len=output_len, image_data=None, # Mooncake doesn't have image data ) else: # For all other datasets, input_requests is a list of DatasetRow objects test_request = input_requests[0] if lora_names is not None and len(lora_names) != 0: lora_name = lora_names[0] else: lora_name = None # Create the test input once test_input = RequestFuncInput( model=model_id, prompt=test_request.prompt, api_url=api_url, prompt_len=test_request.prompt_len, output_len=min(test_request.output_len, 32), lora_name=lora_name, image_data=test_request.image_data, extra_request_body=extra_request_body, ) # Run warmup requests warmup_tasks = [] for _ in range(warmup_requests): warmup_tasks.append( asyncio.create_task(request_func(request_func_input=test_input)) ) warmup_outputs = await asyncio.gather(*warmup_tasks) if is_multi_turn: warmup_outputs = [x for output in warmup_outputs for x in output] # Check if at least one warmup request succeeded if warmup_requests > 0 and not any(output.success for output in warmup_outputs): raise ValueError( "Warmup failed - Please make sure benchmark arguments " f"are correctly specified. Error: {warmup_outputs[0].error}" ) else: print( f"Warmup completed with {args.warmup_requests} sequences. Starting main benchmark run..." ) # Flush cache if ("sglang" in backend and _get_bool_env_var("SGLANG_IS_IN_CI")) or flush_cache: requests.post(base_url + "/flush_cache", headers=get_auth_headers()) time.sleep(1.0) # Build profile URLs for PD separated mode (do this once at the beginning) pd_profile_urls = [] if profile and pd_separated: pd_profile_urls = _build_profile_urls(profile_prefill_url, profile_decode_url) if not pd_profile_urls: print( "Warning: PD separated mode requires --profile-prefill-url or --profile-decode-url" ) print("Skipping profiler start. Please specify worker URLs for profiling.") # Start profiler if profile: if pd_separated: if pd_profile_urls: await _call_profile_pd(pd_profile_urls, "start") else: print("Starting profiler...") profile_output = await async_request_profile( api_url=base_url + "/start_profile" ) if profile_output.success: print("Profiler started") # Run all requests benchmark_start_time = time.perf_counter() tasks: List[asyncio.Task] = [] pbar_total = len(input_requests) if ( backend == "sglang" and args.dataset_name == "mooncake" ): # Assuming mooncake is mainly for sglang or similar backends print("Using time-based Mooncake request scheduler, ignoring --request-rate.") request_generator = get_mooncake_request_over_time( input_requests, tokenizer, mooncake_slowdown_factor, mooncake_num_rounds ) print( f"Starting Mooncake trace replay. Sessions: {len(input_requests)}, Rounds per session: {mooncake_num_rounds}. Slowdown factor: {mooncake_slowdown_factor}" ) pbar_total *= args.mooncake_num_rounds else: request_generator = get_request(input_requests, request_rate) # Prepare LoRA request distribution parameters if lora_request_distribution == "distinct": lora_idx = 0 elif lora_request_distribution == "skewed": weights = np.array([lora_zipf_alpha**-i for i in range(len(lora_names))]) lora_probs = weights / np.sum(weights) else: lora_idx = None lora_probs = None pbar = None if disable_tqdm else tqdm(total=pbar_total) async for request in request_generator: if lora_names is not None and len(lora_names) != 0: if lora_request_distribution == "uniform": lora_name = random.choice(lora_names) elif lora_request_distribution == "distinct": lora_name = lora_names[lora_idx] lora_idx = (lora_idx + 1) % len(lora_names) else: assert ( lora_request_distribution == "skewed" ), f"Unexpected lora_request_distribution: {lora_request_distribution}. Expected 'skewed'." lora_name = np.random.choice(lora_names, p=lora_probs) else: lora_name = None # Merge global extra_request_body with per-request extras # Per-request parameters take precedence over global ones merged_extra_body = {**extra_request_body, **request.extra_request_body} request_func_input = RequestFuncInput( model=model_id, prompt=request.prompt, api_url=api_url, prompt_len=request.prompt_len, output_len=request.output_len, lora_name=lora_name, image_data=request.image_data, extra_request_body=merged_extra_body, timestamp=request.timestamp, routing_key=request.routing_key, ) tasks.append( asyncio.create_task( limited_request_func(request_func_input=request_func_input, pbar=pbar) ) ) outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) if is_multi_turn: outputs = [x for output in outputs for x in output] # Stop profiler (only if profile_steps was not provided, as it auto-stops) if profile and not ( hasattr(args, "profile_steps") and args.profile_steps is not None ): if pd_separated: if pd_profile_urls: await _call_profile_pd(pd_profile_urls, "stop") else: if getattr(args, "profile_num_steps", None) is None: print("Stopping profiler...") profile_output = await async_request_profile( api_url=base_url + "/stop_profile" ) if profile_output.success: print("Profiler stopped") if pbar is not None: pbar.close() if "sglang" in backend: server_info = requests.get( base_url + "/get_server_info", headers=get_auth_headers() ) if server_info.status_code == 200: server_info_json = server_info.json() if "decode" in server_info_json: server_info_json = server_info_json["decode"][0] if ( "internal_states" in server_info_json and server_info_json["internal_states"] ): accept_length = server_info_json["internal_states"][0].get( "avg_spec_accept_length", None ) else: accept_length = None else: accept_length = None else: accept_length = None # Compute metrics and print results benchmark_duration = time.perf_counter() - benchmark_start_time metrics, output_lens = calculate_metrics( input_requests=None if is_multi_turn else input_requests, outputs=outputs, dur_s=benchmark_duration, tokenizer=tokenizer, backend=backend, accept_length=accept_length, plot_throughput=args.plot_throughput, ) print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) print("{:<40} {:<10}".format("Backend:", backend)) print( "{:<40} {:<10}".format( "Traffic request rate:", "trace" if use_trace_timestamps else request_rate ) ) print( "{:<40} {:<10}".format( "Max request concurrency:", max_concurrency if max_concurrency else "not set", ) ) print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration)) print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) print("{:<40} {:<10}".format("Total input text tokens:", metrics.total_input_text)) if args.dataset_name in ["image", "mmmu"]: print( "{:<40} {:<10}".format( "Total input vision tokens:", metrics.total_input_vision ) ) is_embedding = backend == "sglang-embedding" if not is_embedding: print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output)) print( "{:<40} {:<10}".format( "Total generated tokens (retokenized):", metrics.total_output_retokenized, ) ) print( "{:<40} {:<10.2f}".format( "Request throughput (req/s):", metrics.request_throughput ) ) print( "{:<40} {:<10.2f}".format( "Input token throughput (tok/s):", metrics.input_throughput ) ) if not is_embedding: print( "{:<40} {:<10.2f}".format( "Output token throughput (tok/s):", metrics.output_throughput ) ) print( "{:<40} {:<10.2f}".format( "Peak output token throughput (tok/s):", metrics.max_output_tokens_per_s, ) ) print( "{:<40} {:<10}".format( "Peak concurrent requests:", metrics.max_concurrent_requests ) ) if not is_embedding: print( "{:<40} {:<10.2f}".format( "Total token throughput (tok/s):", metrics.total_throughput ) ) print("{:<40} {:<10.2f}".format("Concurrency:", metrics.concurrency)) if accept_length: print("{:<40} {:<10.2f}".format("Accept length:", accept_length)) print("{s:{c}^{n}}".format(s="End-to-End Latency", n=50, c="-")) print( "{:<40} {:<10.2f}".format("Mean E2E Latency (ms):", metrics.mean_e2e_latency_ms) ) print( "{:<40} {:<10.2f}".format( "Median E2E Latency (ms):", metrics.median_e2e_latency_ms ) ) print( "{:<40} {:<10.2f}".format("P90 E2E Latency (ms):", metrics.p90_e2e_latency_ms) ) print( "{:<40} {:<10.2f}".format("P99 E2E Latency (ms):", metrics.p99_e2e_latency_ms) ) if not is_embedding: print("{s:{c}^{n}}".format(s="Time to First Token", n=50, c="-")) print("{:<40} {:<10.2f}".format("Mean TTFT (ms):", metrics.mean_ttft_ms)) print("{:<40} {:<10.2f}".format("Median TTFT (ms):", metrics.median_ttft_ms)) print("{:<40} {:<10.2f}".format("P99 TTFT (ms):", metrics.p99_ttft_ms)) print( "{s:{c}^{n}}".format( s="Time per Output Token (excl. 1st token)", n=50, c="-" ) ) print("{:<40} {:<10.2f}".format("Mean TPOT (ms):", metrics.mean_tpot_ms)) print("{:<40} {:<10.2f}".format("Median TPOT (ms):", metrics.median_tpot_ms)) print("{:<40} {:<10.2f}".format("P99 TPOT (ms):", metrics.p99_tpot_ms)) print("{s:{c}^{n}}".format(s="Inter-Token Latency", n=50, c="-")) print("{:<40} {:<10.2f}".format("Mean ITL (ms):", metrics.mean_itl_ms)) print("{:<40} {:<10.2f}".format("Median ITL (ms):", metrics.median_itl_ms)) print("{:<40} {:<10.2f}".format("P95 ITL (ms):", metrics.p95_itl_ms)) print("{:<40} {:<10.2f}".format("P99 ITL (ms):", metrics.p99_itl_ms)) print("{:<40} {:<10.2f}".format("Max ITL (ms):", metrics.max_itl_ms)) print("=" * 50) resp = requests.get(base_url + "/get_server_info", headers=get_auth_headers()) server_info = resp.json() if resp.status_code == 200 else None if ( metrics.median_ttft_ms is not None and metrics.mean_itl_ms is not None and metrics.output_throughput is not None ): result = { # Arguments "tag": getattr(args, "tag", None), "backend": args.backend, "dataset_name": args.dataset_name, "request_rate": "trace" if use_trace_timestamps else request_rate, "max_concurrency": max_concurrency, "sharegpt_output_len": args.sharegpt_output_len, "random_input_len": args.random_input_len, "random_output_len": args.random_output_len, "random_range_ratio": args.random_range_ratio, # Information "server_info": server_info, # Results "duration": benchmark_duration, "completed": metrics.completed, "total_input_tokens": metrics.total_input, "total_input_text_tokens": metrics.total_input_text, "total_input_vision_tokens": metrics.total_input_vision, "total_output_tokens": metrics.total_output, "total_output_tokens_retokenized": metrics.total_output_retokenized, "request_throughput": metrics.request_throughput, "input_throughput": metrics.input_throughput, "output_throughput": metrics.output_throughput, "total_throughput": metrics.total_throughput, "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, "median_e2e_latency_ms": metrics.median_e2e_latency_ms, "std_e2e_latency_ms": metrics.std_e2e_latency_ms, "p90_e2e_latency_ms": metrics.p90_e2e_latency_ms, "p99_e2e_latency_ms": metrics.p99_e2e_latency_ms, "mean_ttft_ms": metrics.mean_ttft_ms, "median_ttft_ms": metrics.median_ttft_ms, "std_ttft_ms": metrics.std_ttft_ms, "p99_ttft_ms": metrics.p99_ttft_ms, "mean_tpot_ms": metrics.mean_tpot_ms, "median_tpot_ms": metrics.median_tpot_ms, "std_tpot_ms": metrics.std_tpot_ms, "p99_tpot_ms": metrics.p99_tpot_ms, "mean_itl_ms": metrics.mean_itl_ms, "median_itl_ms": metrics.median_itl_ms, "std_itl_ms": metrics.std_itl_ms, "p95_itl_ms": metrics.p95_itl_ms, "p99_itl_ms": metrics.p99_itl_ms, "concurrency": metrics.concurrency, "accept_length": accept_length, "max_output_tokens_per_s": metrics.max_output_tokens_per_s, "max_concurrent_requests": metrics.max_concurrent_requests, } else: print(f"Error running benchmark for request rate: {request_rate}") print("-" * 30) # Determine output file name if args.output_file: output_file_name = args.output_file else: now = datetime.now().strftime("%m%d") if args.dataset_name == "image": output_file_name = ( f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_" f"{args.random_output_len}_{args.image_count}imgs_" f"{args.image_resolution}.jsonl" ) elif args.dataset_name.startswith("random"): output_file_name = f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_{args.random_output_len}.jsonl" else: output_file_name = ( f"{args.backend}_{now}_{args.num_prompts}_{args.dataset_name}.jsonl" ) result_details = { "input_lens": [output.prompt_len for output in outputs], "output_lens": output_lens, "ttfts": [output.ttft for output in outputs], "itls": [output.itl for output in outputs], "generated_texts": [output.generated_text for output in outputs], "errors": [output.error for output in outputs], } # Append results to a JSONL file with open(output_file_name, "a") as file: if args.output_details: result_for_dump = result | result_details else: result_for_dump = result file.write(json.dumps(result_for_dump) + "\n") return result | result_details def check_chat_template(model_path): try: tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) return "chat_template" in tokenizer.init_kwargs except Exception as e: print(f"Fail to load tokenizer config with error={e}") return False def set_global_args(args_: argparse.Namespace): """Set the global args.""" global args args = args_ def run_benchmark(args_: argparse.Namespace): global args args = args_ # Set default value for max_concurrency if not present if not hasattr(args, "max_concurrency"): args.max_concurrency = None # Set default value for warmup_requests if not present if not hasattr(args, "warmup_requests"): args.warmup_requests = 1 if not hasattr(args, "output_details"): args.output_details = False if not hasattr(args, "tokenize_prompt"): args.tokenize_prompt = False if not hasattr(args, "plot_throughput"): args.plot_throughput = False if not hasattr(args, "top_logprobs_num"): args.top_logprobs_num = 0 if not hasattr(args, "token_ids_logprob"): args.token_ids_logprob = None if not hasattr(args, "logprob_start_len"): args.logprob_start_len = -1 if not hasattr(args, "return_logprob"): args.return_logprob = False if not hasattr(args, "use_trace_timestamps"): args.use_trace_timestamps = False if not hasattr(args, "mooncake_slowdown_factor"): args.mooncake_slowdown_factor = 1.0 if not hasattr(args, "mooncake_slowdown_factor"): args.mooncake_slowdown_factor = 1.0 if not hasattr(args, "mooncake_num_rounds"): args.mooncake_num_rounds = 1 if not hasattr(args, "served_model_name"): args.served_model_name = None if getattr(args, "print_requests", False): assert args.backend == "sglang-oai-chat" # only support this now print(f"benchmark_args={args}") # Set global environments set_ulimit() random.seed(args.seed) np.random.seed(args.seed) extra_request_body = {} if args.extra_request_body: extra_request_body = json.loads(args.extra_request_body) if args.tokenize_prompt: assert ( args.backend == "sglang" ), "`--tokenize-prompt` only compatible with `--backend sglang` currently" # Set url if args.port is None: args.port = { "sglang": 30000, "sglang-native": 30000, "sglang-oai": 30000, "lmdeploy": 23333, "vllm": 8000, "trt": 8000, "gserver": 9988, "truss": 8080, }.get(args.backend, 30000) model_url = ( f"{args.base_url}/v1/models" if args.base_url else f"http://{args.host}:{args.port}/v1/models" ) if args.backend == "sglang-embedding": api_url = ( f"{args.base_url}/v1/embeddings" if args.base_url else f"http://{args.host}:{args.port}/v1/embeddings" ) elif args.backend in ["sglang", "sglang-native"]: api_url = ( f"{args.base_url}/generate" if args.base_url else f"http://{args.host}:{args.port}/generate" ) elif args.backend in ["sglang-oai", "vllm", "lmdeploy"]: api_url = ( f"{args.base_url}/v1/completions" if args.base_url else f"http://{args.host}:{args.port}/v1/completions" ) elif args.backend in ["sglang-oai-chat", "vllm-chat", "lmdeploy-chat"]: api_url = ( f"{args.base_url}/v1/chat/completions" if args.base_url else f"http://{args.host}:{args.port}/v1/chat/completions" ) elif args.backend == "trt": api_url = ( f"{args.base_url}/v2/models/ensemble/generate_stream" if args.base_url else f"http://{args.host}:{args.port}/v2/models/ensemble/generate_stream" ) if args.model is None: print("Please provide a model using `--model` when using `trt` backend.") sys.exit(1) elif args.backend == "gserver": api_url = args.base_url if args.base_url else f"{args.host}:{args.port}" args.model = args.model or "default" elif args.backend == "truss": api_url = ( f"{args.base_url}/v1/models/model:predict" if args.base_url else f"http://{args.host}:{args.port}/v1/models/model:predict" ) base_url = ( f"http://{args.host}:{args.port}" if args.base_url is None else args.base_url ) # Wait for server to be ready if args.ready_check_timeout_sec > 0: health_url = model_url if args.backend not in ("trt", "gserver") else base_url if not wait_for_endpoint(health_url, args.ready_check_timeout_sec): print(f"Server at {health_url} is not ready. Exiting.") sys.exit(1) # Get model name if args.model is None: if args.backend == "truss": print( "Please provide a model with `--model` when using truss backend. e.g. --model meta-llama/Llama-3.1-8B-Instruct" ) sys.exit(1) try: response = requests.get(model_url, headers=get_auth_headers()) model_list = response.json().get("data", []) args.model = model_list[0]["id"] if model_list else None except Exception as e: print(f"Failed to fetch model from {model_url}. Error: {e}") print( "Please specify the correct host and port using `--host` and `--port`." ) sys.exit(1) if args.model is None: print("No model specified or found. Please provide a model using `--model`.") sys.exit(1) if args.backend != "sglang-embedding" and not check_chat_template(args.model): print( "\nWARNING It is recommended to use the `Chat` or `Instruct` model for benchmarking.\n" "Because when the tokenizer counts the output tokens, if there is gibberish, it might count incorrectly.\n" ) if ( args.backend == "sglang-embedding" and args.dataset_name in _EMBEDDING_UNSUPPORTED_DATASETS ): print(f"{args.dataset_name} dataset is unsupported for embeddings benchmark") sys.exit(1) if args.dataset_name in ["image", "mmmu"]: args.apply_chat_template = True assert ( not args.tokenize_prompt ), "`--tokenize-prompt` not compatible with image dataset" if args.lora_request_distribution in ["distinct", "skewed"]: assert ( args.lora_name is not None and len(args.lora_name) > 1 ), "More than 1 LoRA adapter must be specified via --lora-name to use 'distinct' or 'skewed' request distribution." assert ( args.lora_zipf_alpha > 1 ), f"Got invalid value for --lora-zipf-alpha of {args.lora_zipf_alpha}. It must be greater than 1." print(f"{args}\n") # Read dataset backend = args.backend model_id = args.served_model_name or args.model tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model tokenizer = get_tokenizer(tokenizer_id) input_requests = get_dataset(args, tokenizer, model_id) # compatible with SimpleNamespace if not hasattr(args, "flush_cache"): args.flush_cache = False # Prepare LoRA arguments lora_request_distribution = ( args.lora_request_distribution if args.lora_name is not None else None ) lora_zipf_alpha = ( args.lora_zipf_alpha if args.lora_name is not None and args.lora_request_distribution == "skewed" else None ) return asyncio.run( benchmark( backend=backend, api_url=api_url, base_url=base_url, model_id=model_id, tokenizer=tokenizer, input_requests=input_requests, request_rate=args.request_rate, max_concurrency=args.max_concurrency, disable_tqdm=args.disable_tqdm, lora_names=args.lora_name, lora_request_distribution=lora_request_distribution, lora_zipf_alpha=lora_zipf_alpha, extra_request_body=extra_request_body, profile=args.profile, pd_separated=args.pd_separated, flush_cache=args.flush_cache, warmup_requests=args.warmup_requests, use_trace_timestamps=args.use_trace_timestamps, mooncake_slowdown_factor=args.mooncake_slowdown_factor, mooncake_num_rounds=args.mooncake_num_rounds, profile_prefill_url=getattr(args, "profile_prefill_url", None), profile_decode_url=getattr(args, "profile_decode_url", None), ) ) class LoRAPathAction(argparse.Action): def __call__(self, parser, namespace, values, option_string=None): setattr(namespace, self.dest, []) for lora_name in values: getattr(namespace, self.dest).append(lora_name) if __name__ == "__main__": parser = ArgumentParser(description="Benchmark the online serving throughput.") parser.add_argument( "--backend", type=str, choices=list(ASYNC_REQUEST_FUNCS.keys()), default="sglang", help="Must specify a backend, depending on the LLM Inference Engine.", ) parser.add_argument( "--base-url", type=str, default=None, help="Server or API base url if not using http host and port.", ) parser.add_argument( "--host", type=str, default="0.0.0.0", help="Default host is 0.0.0.0." ) parser.add_argument( "--port", type=int, help="If not set, the default port is configured according to its default value for different LLM Inference Engines.", ) parser.add_argument( "--ready-check-timeout-sec", type=int, default=60, help="Maximum time in seconds to wait for the server to be ready before benchmarking. Set to 0 to skip. Default: 60.", ) parser.add_argument( "--dataset-name", type=str, default="sharegpt", choices=[ "sharegpt", "custom", "openai", "random", "random-ids", "generated-shared-prefix", "mmmu", "image", "mooncake", ], help="Name of the dataset to benchmark on.", ) parser.add_argument( "--dataset-path", type=str, default="", help="Path to the dataset." ) parser.add_argument( "--model", type=str, help="Name or path of the model. If not set, the default model will request /v1/models for conf.", ) parser.add_argument( "--served-model-name", type=str, help="The name of the model as served by the serving service. If not set, this defaults to the value of --model.", ) parser.add_argument( "--tokenizer", type=str, help="Name or path of the tokenizer. If not set, using the model conf.", ) parser.add_argument( "--num-prompts", type=int, default=1000, help="Number of prompts to process. Default is 1000.", ) parser.add_argument( "--sharegpt-output-len", type=int, default=None, help="Output length for each request. Overrides the output length from the ShareGPT dataset.", ) parser.add_argument( "--sharegpt-context-len", type=int, default=None, help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.", ) parser.add_argument( "--random-input-len", type=int, default=1024, help="Number of input tokens per request, used only for random and image dataset.", ) parser.add_argument( "--random-output-len", default=1024, type=int, help="Number of output tokens per request, used only for random and image dataset.", ) parser.add_argument( "--random-range-ratio", type=float, default=0.0, help="Range of sampled ratio of input/output length, " "used only for random and image dataset.", ) # image dataset args parser.add_argument( "--image-count", type=int, default=1, help="Number of images per request (only available with the image dataset)", ) parser.add_argument( "--image-resolution", type=str, default="1080p", help=( "Resolution of images for image dataset. " "Supports presets 4k/1080p/720p/360p or custom 'heightxwidth' (e.g., 1080x1920)." ), ) parser.add_argument( "--random-image-count", action="store_true", help="Enable Random Image Count", ) parser.add_argument( "--image-format", type=str, default="jpeg", help=("Format of images for image dataset. " "Supports jpeg and png."), ) parser.add_argument( "--image-content", type=str, default="random", help=("Content for images for image dataset. " "Supports random and blank."), ) parser.add_argument( "--request-rate", type=float, default=float("inf"), help="Number of requests per second. If this is inf, then all the requests are sent at time 0. " "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.", ) parser.add_argument( "--use-trace-timestamps", action="store_true", help="Use timestamps from the trace file for request scheduling. Only valid for 'mooncake' dataset.", ) parser.add_argument( "--max-concurrency", type=int, default=None, help="Maximum number of concurrent requests. This can be used " "to help simulate an environment where a higher level component " "is enforcing a maximum number of concurrent requests. While the " "--request-rate argument controls the rate at which requests are " "initiated, this argument will control how many are actually allowed " "to execute at a time. This means that when used in combination, the " "actual request rate may be lower than specified with --request-rate, " "if the server is not processing requests fast enough to keep up.", ) parser.add_argument("--output-file", type=str, help="Output JSONL file name.") parser.add_argument( "--output-details", action="store_true", help="Output details of benchmarking." ) parser.add_argument( "--print-requests", action="store_true", help="Print requests immediately during benchmarking. Useful to quickly realize issues.", ) parser.add_argument( "--disable-tqdm", action="store_true", help="Specify to disable tqdm progress bar.", ) parser.add_argument( "--disable-stream", action="store_true", help="Disable streaming mode.", ) parser.add_argument( "--return-logprob", action="store_true", help="Return logprob.", ) parser.add_argument( "--top-logprobs-num", type=int, default=0, help="Number of top logprobs to return per token. Only used with --return-logprob.", ) parser.add_argument( "--token-ids-logprob", type=int, nargs="+", default=None, help="Token IDs to probe logprobs for. E.g. --token-ids-logprob 1 2 10 100 1000. Only used with --return-logprob.", ) parser.add_argument( "--logprob-start-len", type=int, default=-1, help="Start position for returning input logprobs. -1 means no input logprobs, 0 means all. Only used with --return-logprob.", ) parser.add_argument( "--return-routed-experts", action="store_true", help="Return routed experts.", ) parser.add_argument("--seed", type=int, default=1, help="The random seed.") parser.add_argument( "--disable-ignore-eos", action="store_true", help="Disable ignoring EOS.", ) parser.add_argument( "--extra-request-body", metavar='{"key1": "value1", "key2": "value2"}', type=str, help="Append given JSON object to the request payload. You can use this to specify" "additional generate params like sampling params.", ) parser.add_argument( "--apply-chat-template", action="store_true", help="Apply chat template", ) parser.add_argument( "--profile", action="store_true", help="Use Torch Profiler. The endpoint must be launched with " "SGLANG_TORCH_PROFILER_DIR to enable profiler.", ) parser.add_argument( "--plot-throughput", action="store_true", help="Plot throughput and concurrent requests over time. Requires termplotlib and gnuplot.", ) # TODO unify all these parser.add_argument( "--profile-activities", type=str, nargs="+", default=["CPU", "GPU"], choices=["CPU", "GPU", "CUDA_PROFILER", "XPU"], help="Profiler activities to capture: CPU, GPU, XPU, CUDA_PROFILER.", ) parser.add_argument( "--profile-start-step", type=int, default=None, help="Start profiling after this many forward steps. Useful for warmup.", ) parser.add_argument( "--profile-steps", type=int, default=None, help="Number of steps to profile. If specified, profiling stops automatically after this many steps.", ) parser.add_argument("--profile-num-steps", type=int, default=None) parser.add_argument("--profile-by-stage", action="store_true", default=False) parser.add_argument("--profile-stages", nargs="+", default=None) parser.add_argument( "--profile-output-dir", type=str, default=None, help="Output directory for profile traces.", ) parser.add_argument( "--profile-prefix", type=str, default=None, help="Prefix for profile trace filenames.", ) parser.add_argument( "--lora-name", type=str, nargs="*", default=None, action=LoRAPathAction, help="The names of LoRA adapters. You can provide a list of names in the format {name} {name} {name}...", ) parser.add_argument( "--lora-request-distribution", type=str, default="uniform", choices=[ "uniform", "distinct", "skewed", ], help="What distribution to sample the LoRA adapters specified in --lora-name. Borrowed from the Punica paper. " "'distinct' distribution means selecting a new LoRA adapter for every request. " "'skewed' distribution follows the Zipf distribution, where the number of requests " "to model i specified in --lora-name is α times the number of requests for model i+1, " "where α > 1.", ) parser.add_argument( "--lora-zipf-alpha", type=float, default=1.5, help="The parameter to use for the Zipf distribution when --lora-request-distribution='skewed'.", ) parser.add_argument( "--prompt-suffix", type=str, default="", help="Suffix applied to the end of all user prompts, followed by assistant prompt suffix.", ) parser.add_argument( "--pd-separated", action="store_true", help="Benchmark PD disaggregation server", ) # Create a mutually exclusive group for profiling URLs # In PD separated mode, prefill and decode workers must be profiled separately profile_url_group = parser.add_mutually_exclusive_group() profile_url_group.add_argument( "--profile-prefill-url", type=str, nargs="*", default=None, help="URL(s) of the prefill worker(s) for profiling in PD separated mode. " "Can specify multiple URLs: --profile-prefill-url http://localhost:30000 http://localhost:30001. " "NOTE: Cannot be used together with --profile-decode-url. " "In PD separated mode, prefill and decode workers must be profiled separately.", ) profile_url_group.add_argument( "--profile-decode-url", type=str, nargs="*", default=None, help="URL(s) of the decode worker(s) for profiling in PD separated mode. " "Can specify multiple URLs: --profile-decode-url http://localhost:30010 http://localhost:30011. " "NOTE: Cannot be used together with --profile-prefill-url. " "In PD separated mode, prefill and decode workers must be profiled separately.", ) parser.add_argument( "--flush-cache", action="store_true", help="Flush the cache before running the benchmark", ) parser.add_argument( "--warmup-requests", type=int, default=1, help="Number of warmup requests to run before the benchmark", ) parser.add_argument( "--tokenize-prompt", action="store_true", help="Use integer ids instead of string for inputs. Useful to control prompt lengths accurately", ) group = parser.add_argument_group("generated-shared-prefix dataset arguments") group.add_argument( "--gsp-num-groups", type=int, default=64, help="Number of system prompt groups for generated-shared-prefix dataset", ) group.add_argument( "--gsp-prompts-per-group", type=int, default=16, help="Number of prompts per system prompt group for generated-shared-prefix dataset", ) group.add_argument( "--gsp-system-prompt-len", type=int, default=2048, help="Target length in tokens for system prompts in generated-shared-prefix dataset", ) group.add_argument( "--gsp-question-len", type=int, default=128, help="Target length in tokens for questions in generated-shared-prefix dataset", ) group.add_argument( "--gsp-output-len", type=int, default=256, help="Target length in tokens for outputs in generated-shared-prefix dataset", ) parser.add_argument( "--gsp-range-ratio", type=float, # WARN: The default 1.0 is for backward compatibility, and is different from the default 0.0 for random dataset default=1.0, help="Range of sampled ratio of input/output length, used only for gsp dataset.", ) group.add_argument( "--gsp-fast-prepare", action="store_true", help="Speedup preparing by removing statistics computation, which will make some output statistics inaccurate but suitable for pressure tests.", ) group.add_argument( "--gsp-send-routing-key", action="store_true", help="Send routing key in requests via X-SMG-Routing-Key header. Requests with the same prefix share the same routing key.", ) group.add_argument( "--gsp-num-turns", type=int, default=1, help="Number of turns for multi-turn conversations. If > 1, each prompt becomes a list of questions sharing the same system prefix.", ) group.add_argument( "--gsp-ordered", action="store_true", help="Keep requests in order without shuffling. By default, requests are shuffled randomly.", ) mooncake_group = parser.add_argument_group("mooncake dataset arguments") mooncake_group.add_argument( "--mooncake-slowdown-factor", type=float, default=1.0, help="Slowdown factor for replaying the mooncake trace. " "A value of 2.0 means the replay is twice as slow. " "NOTE: --request-rate is IGNORED in mooncake mode.", ) mooncake_group.add_argument( "--mooncake-num-rounds", type=int, default=1, help="Number of conversation rounds for each session in the mooncake dataset. " "A value > 1 will enable true multi-turn session benchmarking.", ) mooncake_group.add_argument( "--mooncake-workload", type=str, default="conversation", choices=[ "mooncake", "conversation", "synthetic", "toolagent", ], help="Underlying workload for the mooncake dataset.", ) parser.add_argument( "--tag", type=str, default=None, help="The tag to be dumped to output." ) parser.add_argument( "--header", type=str, nargs="+", default=None, help="Custom HTTP headers in Key=Value format. Example: --header MyHeader=MY_VALUE MyAnotherHeader=myanothervalue", ) args = parser.parse_args() run_benchmark(args) ================================================ FILE: python/sglang/benchmark/__init__.py ================================================ ================================================ FILE: python/sglang/benchmark/bench_utils.py ================================================ """Triton do_bench/do_bench_cudagraph compatible wrapper using flashinfer.testing.bench_gpu_time.""" import numpy as np from flashinfer.testing import bench_gpu_time def run_bench( fn, use_cuda_graph: bool = True, quantiles=(0.5, 0.2, 0.8), warmup_ms: int = 25, rep_ms: int = 100, ): """Returns (ms, min_ms, max_ms) or (median,) when quantiles=None.""" times = bench_gpu_time( fn=fn, use_cuda_graph=use_cuda_graph, dry_run_time_ms=warmup_ms, repeat_time_ms=rep_ms, ) if quantiles is None: return (float(np.median(times)),) return tuple(float(np.percentile(times, q * 100)) for q in quantiles) ================================================ FILE: python/sglang/benchmark/datasets/__init__.py ================================================ from typing import Dict, Type from sglang.benchmark.datasets.common import BaseDataset, DatasetRow from sglang.benchmark.datasets.custom import CustomDataset from sglang.benchmark.datasets.generated_shared_prefix import ( GeneratedSharedPrefixDataset, ) from sglang.benchmark.datasets.image import ImageDataset from sglang.benchmark.datasets.mmmu import MMMUDataset from sglang.benchmark.datasets.mooncake import MooncakeDataset from sglang.benchmark.datasets.openai_dataset import OpenAIDataset from sglang.benchmark.datasets.random import RandomDataset from sglang.benchmark.datasets.sharegpt import ShareGPTDataset DATASET_MAPPING: Dict[str, Type[BaseDataset]] = { "sharegpt": ShareGPTDataset, "custom": CustomDataset, "openai": OpenAIDataset, # TODO: "random" vs "random-ids" should be a flag (e.g. --random-source=sharegpt|integers), # not two separate dataset names sharing the same class. "random": RandomDataset, "random-ids": RandomDataset, "generated-shared-prefix": GeneratedSharedPrefixDataset, "mmmu": MMMUDataset, "image": ImageDataset, "mooncake": MooncakeDataset, } def get_dataset(args, tokenizer, model_id=None): dataset_name = args.dataset_name if dataset_name.startswith("random") and dataset_name not in DATASET_MAPPING: dataset_name = "random-ids" if dataset_name not in DATASET_MAPPING: raise ValueError(f"Unknown dataset: {args.dataset_name}") dataset_cls = DATASET_MAPPING[dataset_name] dataset = dataset_cls.from_args(args) return dataset.load(tokenizer=tokenizer, model_id=model_id) __all__ = [ "DATASET_MAPPING", "DatasetRow", "get_dataset", ] ================================================ FILE: python/sglang/benchmark/datasets/common.py ================================================ import random from abc import ABC, abstractmethod from argparse import Namespace from dataclasses import dataclass from functools import lru_cache from typing import Any, Dict, List, Optional import numpy as np ASSISTANT_SUFFIX = "Assistant:" SHAREGPT_REPO_ID = "anon8231489123/ShareGPT_Vicuna_unfiltered" SHAREGPT_FILENAME = "ShareGPT_V3_unfiltered_cleaned_split.json" MOONCAKE_DATASET_URL = { "mooncake": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/arxiv-trace/mooncake_trace.jsonl", "conversation": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/conversation_trace.jsonl", "synthetic": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/synthetic_trace.jsonl", "toolagent": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/toolagent_trace.jsonl", } @dataclass class DatasetRow: prompt: Any prompt_len: int output_len: int text_prompt_len: Optional[int] = None vision_prompt_len: Optional[int] = None image_data: Optional[List[str]] = None timestamp: Optional[float] = None routing_key: Optional[str] = None extra_request_body: Optional[Dict[str, Any]] = None # Per-request API parameters def __post_init__(self): if self.text_prompt_len is None: self.text_prompt_len = self.prompt_len if self.vision_prompt_len is None: self.vision_prompt_len = 0 if self.extra_request_body is None: self.extra_request_body = {} @dataclass class BaseDataset(ABC): @classmethod @abstractmethod def from_args(cls, args: Namespace) -> "BaseDataset": ... @abstractmethod def load( self, tokenizer: Any, model_id: Optional[str] = None, ) -> List[DatasetRow]: ... def compute_random_lens(full_len: int, range_ratio: float, num: int) -> List[int]: # full_len=0 is valid for embedding benchmarks where no output tokens are generated if full_len <= 0: return [0] * num return np.random.randint( max(int(full_len * range_ratio), 1), full_len + 1, size=num, ).tolist() @lru_cache(maxsize=1) def get_available_tokens(tokenizer): """Get all available token ids from the tokenizer vocabulary.""" return list(tokenizer.get_vocab().values()) def gen_prompt(tokenizer, token_num): """Generate a random prompt of specified token length using tokenizer vocabulary.""" all_available_tokens = get_available_tokens(tokenizer) selected_tokens = random.choices(all_available_tokens, k=token_num) return tokenizer.decode(selected_tokens) def gen_mm_prompt(tokenizer, image_pad_id, token_num): """Generate a random prompt of specified token length using tokenizer vocabulary.""" all_available_tokens = list(tokenizer.get_vocab().values()) if image_pad_id: all_available_tokens.remove(image_pad_id) selected_tokens = random.choices(all_available_tokens, k=token_num) return tokenizer.decode(selected_tokens) ================================================ FILE: python/sglang/benchmark/datasets/custom.py ================================================ import json import os import random from argparse import Namespace from dataclasses import dataclass from typing import List, Optional import numpy as np from transformers import PreTrainedTokenizerBase from sglang.benchmark.datasets.common import ( ASSISTANT_SUFFIX, BaseDataset, DatasetRow, ) from sglang.benchmark.utils import remove_suffix @dataclass class CustomDataset(BaseDataset): dataset_path: str num_requests: int fixed_output_len: Optional[int] context_len: Optional[int] prompt_suffix: str apply_chat_template: bool @classmethod def from_args(cls, args: Namespace) -> "CustomDataset": assert not getattr(args, "tokenize_prompt", False) return cls( dataset_path=args.dataset_path, num_requests=args.num_prompts, fixed_output_len=args.sharegpt_output_len, context_len=args.sharegpt_context_len, prompt_suffix=args.prompt_suffix, apply_chat_template=args.apply_chat_template, ) def load( self, tokenizer: PreTrainedTokenizerBase, model_id=None ) -> List[DatasetRow]: return sample_custom_requests( dataset_path=self.dataset_path, num_requests=self.num_requests, tokenizer=tokenizer, fixed_output_len=self.fixed_output_len, context_len=self.context_len, prompt_suffix=self.prompt_suffix, apply_chat_template=self.apply_chat_template, ) def sample_custom_requests( dataset_path: str, num_requests: int, tokenizer: PreTrainedTokenizerBase, fixed_output_len: Optional[int] = None, context_len: Optional[int] = None, prompt_suffix: Optional[str] = "", apply_chat_template=False, ) -> List[DatasetRow]: """ Sample requests from a custom JSONL dataset: supports 'content'/'value' as conversation keys. """ if fixed_output_len is not None and fixed_output_len < 4: raise ValueError("output_len too small") # Load the dataset dataset = [] if not os.path.isfile(dataset_path): raise FileNotFoundError(f"Dataset not found at {dataset_path}") with open(dataset_path, "r", encoding="utf-8") as f: for line in f: line = line.strip() if line: # skip empty lines try: dataset.append(json.loads(line)) except json.JSONDecodeError: continue # skip lines with JSON errors # Filter out the conversations with less than 2 turns. processed_dataset = [] for data in dataset: convs = data.get("conversations", data.get("conversation", [])) if len(convs) >= 2: user_turn = convs[0].get("content", convs[0].get("value", "")) assist_turn = convs[1].get("content", convs[1].get("value", "")) processed_dataset.append((user_turn, assist_turn)) dataset = processed_dataset random.shuffle(dataset) # Filter out sequences that are too long or too short filtered_dataset: List[DatasetRow] = [] for i in range(len(dataset)): if len(filtered_dataset) == num_requests: break # Tokenize the prompts and completions. prompt = dataset[i][0] if prompt_suffix: prompt = ( remove_suffix(prompt, ASSISTANT_SUFFIX) + prompt_suffix + ASSISTANT_SUFFIX ) if apply_chat_template: prompt = tokenizer.apply_chat_template( [{"role": "user", "content": prompt}], add_generation_prompt=True, tokenize=False, return_dict=False, ) if tokenizer.bos_token: prompt = prompt.replace(tokenizer.bos_token, "") prompt_token_ids = tokenizer.encode(prompt) completion = dataset[i][1] completion_token_ids = tokenizer.encode(completion) prompt_len = len(prompt_token_ids) output_len = ( len(completion_token_ids) if fixed_output_len is None else fixed_output_len ) if prompt_len < 2 or output_len < 2: # Prune too short sequences. continue if context_len and prompt_len + output_len > context_len: # Prune too long sequences. continue filtered_dataset.append( DatasetRow( prompt=prompt, prompt_len=prompt_len, output_len=output_len, ) ) print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}") print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}") return filtered_dataset ================================================ FILE: python/sglang/benchmark/datasets/generated_shared_prefix.py ================================================ import pickle import random import uuid from argparse import Namespace from dataclasses import dataclass from datetime import datetime from pathlib import Path from typing import List import numpy as np from tqdm.asyncio import tqdm from transformers import PreTrainedTokenizerBase from sglang.benchmark.datasets.common import ( BaseDataset, DatasetRow, compute_random_lens, gen_prompt, ) @dataclass class GeneratedSharedPrefixDataset(BaseDataset): num_groups: int prompts_per_group: int system_prompt_len: int question_len: int output_len: int range_ratio: float seed: int fast_prepare: bool send_routing_key: bool num_turns: int ordered: bool @classmethod def from_args(cls, args: Namespace) -> "GeneratedSharedPrefixDataset": assert not getattr(args, "tokenize_prompt", False) return cls( num_groups=args.gsp_num_groups, prompts_per_group=args.gsp_prompts_per_group, system_prompt_len=args.gsp_system_prompt_len, question_len=args.gsp_question_len, output_len=args.gsp_output_len, range_ratio=getattr(args, "gsp_range_ratio", 1.0), seed=args.seed, fast_prepare=getattr(args, "gsp_fast_prepare", False), send_routing_key=getattr(args, "gsp_send_routing_key", False), num_turns=getattr(args, "gsp_num_turns", 1), ordered=getattr(args, "gsp_ordered", False), ) def load( self, tokenizer: PreTrainedTokenizerBase, model_id=None ) -> List[DatasetRow]: return sample_generated_shared_prefix_requests( num_groups=self.num_groups, prompts_per_group=self.prompts_per_group, system_prompt_len=self.system_prompt_len, question_len=self.question_len, output_len=self.output_len, range_ratio=self.range_ratio, tokenizer=tokenizer, seed=self.seed, send_routing_key=self.send_routing_key, num_turns=self.num_turns, fast_prepare=self.fast_prepare, ordered=self.ordered, ) def get_gen_prefix_cache_path( seed: int, num_groups: int, prompts_per_group: int, system_prompt_len: int, question_len: int, output_len: int, tokenizer, ): """Create cache directory under ~/.cache/sglang/benchmark""" cache_dir = Path.home() / ".cache" / "sglang" / "benchmark" cache_key = ( f"gen_shared_prefix_{seed}_{num_groups}_{prompts_per_group}_" f"{system_prompt_len}_{question_len}_{output_len}_" f"{tokenizer.__class__.__name__}.pkl" ) return cache_dir / cache_key def sample_generated_shared_prefix_requests( num_groups: int, prompts_per_group: int, system_prompt_len: int, question_len: int, output_len: int, range_ratio: float, tokenizer: PreTrainedTokenizerBase, seed: int, send_routing_key: bool = False, num_turns: int = 1, fast_prepare: bool = False, ordered: bool = False, ) -> List[DatasetRow]: """Generate benchmark requests with shared system prompts using random tokens and caching.""" cache_path = get_gen_prefix_cache_path( seed, num_groups, prompts_per_group, system_prompt_len, question_len, output_len, tokenizer, ) should_cache = (range_ratio == 1) and not send_routing_key and num_turns == 1 # Try to load from cache first if cache_path.exists() and should_cache: print(f"\nLoading cached generated input data from {cache_path}") with open(cache_path, "rb") as f: return pickle.load(f) print( f"\nGenerating new input data... " f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=}, {num_turns=})" ) run_random_str = uuid.uuid4().hex[:8] run_start_timestamp = datetime.now().strftime("%Y%m%d%H%M%S") system_prompt_lens = compute_random_lens( full_len=system_prompt_len, range_ratio=range_ratio, num=num_groups, ) question_lens = np.array( compute_random_lens( full_len=question_len, range_ratio=range_ratio, num=num_groups * prompts_per_group * num_turns, ) ).reshape(num_groups, prompts_per_group, num_turns) output_lens = np.array( compute_random_lens( full_len=output_len, range_ratio=range_ratio, num=num_groups * prompts_per_group, ) ).reshape(num_groups, prompts_per_group) del system_prompt_len, question_len, output_len # Generate system prompts for each group system_prompts = [ gen_prompt(tokenizer, system_prompt_lens[i]) for i in range(num_groups) ] # Generate questions: shape (num_groups, prompts_per_group, num_turns) questions = [ [ [ gen_prompt(tokenizer, int(question_lens[g, p, t])) for t in range(num_turns) ] for p in range(prompts_per_group) ] for g in range(num_groups) ] # Combine system prompts with questions input_requests = [] total_input_tokens = 0 total_output_tokens = 0 for group_idx in tqdm(range(num_groups), desc="Generating system prompt"): system_prompt = system_prompts[group_idx] routing_key = ( f"{run_random_str}_{run_start_timestamp}_{group_idx}" if send_routing_key else None ) for prompt_idx in tqdm( range(prompts_per_group), desc="Generating questions", leave=False ): turn_questions = questions[group_idx][prompt_idx] turn_prompts = [f"{system_prompt}\n\n{turn_questions[0]}"] + turn_questions[ 1: ] full_prompt = turn_prompts[0] if num_turns == 1 else turn_prompts prompt_len = 1 if fast_prepare else len(tokenizer.encode(turn_prompts[0])) output_len_val = int(output_lens[group_idx, prompt_idx]) input_requests.append( DatasetRow( prompt=full_prompt, prompt_len=prompt_len, output_len=output_len_val, routing_key=routing_key, ) ) total_input_tokens += prompt_len total_output_tokens += output_len_val if not ordered: random.shuffle(input_requests) # Print statistics print(f"\nGenerated shared prefix dataset statistics:") print(f"Number of groups: {num_groups}") print(f"Prompts per group: {prompts_per_group}") print(f"Number of turns: {num_turns}") print(f"Total prompts: {len(input_requests)}") if not fast_prepare: print(f"Total input tokens: {total_input_tokens}") print(f"Total output tokens: {total_output_tokens}") print( f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens" ) all_questions = [q for group in questions for conv in group for q in conv] print( f"Average question length: {sum(len(tokenizer.encode(q)) for q in all_questions) / len(all_questions):.1f} tokens\n" ) # Save to cache if should_cache: cache_path.parent.mkdir(parents=True, exist_ok=True) print(f"Caching generated input data to {cache_path}") with open(cache_path, "wb") as f: pickle.dump(input_requests, f) return input_requests ================================================ FILE: python/sglang/benchmark/datasets/image.py ================================================ import io import warnings from argparse import Namespace from dataclasses import dataclass from typing import List, Tuple import numpy as np import pybase64 from PIL import Image from transformers import AutoProcessor from sglang.benchmark.datasets.common import ( BaseDataset, DatasetRow, compute_random_lens, gen_mm_prompt, ) from sglang.benchmark.utils import get_processor @dataclass class ImageDataset(BaseDataset): num_requests: int image_count: int input_len: int output_len: int range_ratio: float image_content: str image_format: str image_resolution: str backend: str random_image_count: bool @classmethod def from_args(cls, args: Namespace) -> "ImageDataset": return cls( num_requests=args.num_prompts, image_count=args.image_count, input_len=args.random_input_len, output_len=args.random_output_len, range_ratio=args.random_range_ratio, image_content=args.image_content, image_format=args.image_format, image_resolution=args.image_resolution, backend=args.backend, random_image_count=args.random_image_count, ) def load(self, tokenizer=None, model_id=None) -> List[DatasetRow]: processor = get_processor(model_id) return sample_image_requests( num_requests=self.num_requests, image_count=self.image_count, input_len=self.input_len, output_len=self.output_len, range_ratio=self.range_ratio, processor=processor, image_content=self.image_content, image_format=self.image_format, image_resolution=self.image_resolution, backend=self.backend, random_image_count=self.random_image_count, ) def parse_image_resolution(image_resolution: str) -> Tuple[int, int]: """Parse image resolution into (width, height). Supports presets '1080p', '720p', '360p' and custom 'heightxwidth' format (e.g., '1080x1920' means height=1080, width=1920). """ resolution_to_size = { "4k": (3840, 2160), "1080p": (1920, 1080), "720p": (1280, 720), "360p": (640, 360), } if image_resolution in resolution_to_size: return resolution_to_size[image_resolution] res = image_resolution.strip().lower() if "x" in res: parts = res.split("x") if len(parts) == 2 and parts[0].isdigit() and parts[1].isdigit(): height = int(parts[0]) width = int(parts[1]) if height > 0 and width > 0: return (width, height) raise ValueError( f"Unsupported image resolution: {image_resolution}. " "Choose from 4k, 1080p, 720p, 360p, or provide custom 'heightxwidth' (e.g., 1080x1920)." ) def create_mm_data_row( text_prompt, images: list, images_base64, output_len, processor, backend ): try: if type(processor).__name__ == "Phi4MMProcessor": # <|endoftext10|> is the image token used in the phi-4-multimodal model. content_items = text_prompt.replace("image 1", "|endoftext10|") else: content_items = [ {"type": "image", "image": {"url": image_base64}} for image_base64 in images_base64 ] content_items.append({"type": "text", "text": text_prompt}) prompt_str = processor.apply_chat_template( [{"role": "user", "content": content_items}], add_generation_prompt=True, tokenize=False, ) except Exception as e: # Note (Xinyuan): This is a workaround for an issue where some tokenizers do not support content as a list. (e.g. InternVL) print(f"Error applying chat template: {e}, fallback to tag") # Some tokenizers do not support list content; fall back to a placeholder in the text prompt_str = f"{text_prompt}" # Calculate total tokens (text + vision) prompt_len = processor( text=[prompt_str], images=images, padding=False, return_tensors="pt", )["input_ids"].numel() # Calculate text-only tokens try: # Create text-only version of the prompt text_only_prompt = processor.apply_chat_template( [{"role": "user", "content": text_prompt}], add_generation_prompt=True, tokenize=False, ) text_prompt_len = processor( text=[text_only_prompt], padding=False, return_tensors="pt", )["input_ids"].numel() except Exception: # Fallback: just tokenize the text prompt directly tokenizer_to_use = ( processor.tokenizer if hasattr(processor, "tokenizer") else processor ) text_prompt_len = len(tokenizer_to_use.encode(text_prompt)) # Vision tokens = total tokens - text tokens vision_prompt_len = prompt_len - text_prompt_len use_raw_prompt = backend in [ "sglang", "sglang-oai", "sglang-oai-chat", "vllm", "vllm-chat", "lmdeploy", "lmdeploy-chat", ] return DatasetRow( prompt=text_prompt if use_raw_prompt else prompt_str, prompt_len=prompt_len, output_len=output_len, text_prompt_len=text_prompt_len, vision_prompt_len=vision_prompt_len, image_data=images_base64, ) def sample_image_requests( num_requests: int, image_count: int, input_len: int, output_len: int, range_ratio: float, processor: AutoProcessor, image_content: str, image_format: str, image_resolution: str, backend: str, random_image_count: bool = False, ) -> List[DatasetRow]: """Generate requests with images. - If ``random_image_count`` is True, each request includes a random number of images between 1 and ``image_count``. - If ``random_image_count`` is False, each request includes exactly ``image_count`` images. - Supported resolutions: 4k (3840x2160), 1080p (1920x1080), 720p (1280x720), 360p (640x360), or custom 'heightxwidth' (e.g., 1080x1920). - Text lengths follow the 'random' dataset sampling rule. ``prompt_len`` only counts text tokens and excludes image data. """ # Parse resolution (supports presets and 'heightxwidth') width, height = parse_image_resolution(image_resolution) # Determine image counts for each request if random_image_count: # Random number of images per request image_counts = np.random.randint(1, image_count + 1, size=num_requests) total_images = np.sum(image_counts) else: # Fixed number of images per request image_counts = np.full(num_requests, image_count) total_images = image_count * num_requests # Check for potentially problematic combinations and warn user if width * height >= 1920 * 1080 and total_images >= 100: warnings.warn( f"High resolution ({width}x{height}) with {total_images} total images " f"may take a long time. Consider reducing resolution or image count.", UserWarning, stacklevel=2, ) # Sample text lengths input_lens = compute_random_lens( full_len=input_len, range_ratio=range_ratio, num=num_requests, ) output_lens = compute_random_lens( full_len=output_len, range_ratio=range_ratio, num=num_requests, ) def _gen_random_image_data_uri( width: int = width, height: int = height ) -> Tuple[Image.Image, str, int]: if image_content == "blank": # Generate blank white image arr = np.full((height, width, 3), 255, dtype=np.uint8) else: # Generate random colored image arr = (np.random.rand(height, width, 3) * 255).astype(np.uint8) img = Image.fromarray(arr) buf = io.BytesIO() img.save(buf, format=image_format, quality=85) encoded = pybase64.b64encode(buf.getvalue()).decode("utf-8") image_data = f"data:image/{image_format};base64,{encoded}" image_bytes = len(image_data.encode("utf-8")) return img, image_data, image_bytes dataset: List[DatasetRow] = [] total_image_bytes = 0 for i in range(num_requests): # Get the number of images for this request request_image_count = int(image_counts[i]) # Generate text prompt text_prompt = gen_mm_prompt( processor.tokenizer, processor.image_token_id if hasattr(processor, "image_token_id") else None, int(input_lens[i]), ) # Generate image list images, images_base64, images_bytes = zip( *[_gen_random_image_data_uri() for _ in range(request_image_count)] ) total_image_bytes += sum(images_bytes) data_row = create_mm_data_row( text_prompt, list(images), list(images_base64), int(output_lens[i]), processor, backend, ) dataset.append(data_row) # Print statistics print(f"#Input tokens: {np.sum([x.prompt_len for x in dataset])}") print(f"#Output tokens: {np.sum([x.output_len for x in dataset])}") print(f"#Total images: {total_images}") if random_image_count: print( f"#Images per request: min={np.min(image_counts)}, max={np.max(image_counts)}, mean={np.mean(image_counts):.2f}" ) else: print(f"#Images per request: {image_count} (fixed)") print( f"\nCreated {len(dataset)} {image_content} {image_format} images with average {total_image_bytes // num_requests} bytes per request" ) return dataset ================================================ FILE: python/sglang/benchmark/datasets/mmmu.py ================================================ import io import random from argparse import Namespace from dataclasses import dataclass from typing import List, Optional import pybase64 from datasets import load_dataset from transformers import AutoProcessor, AutoTokenizer from sglang.benchmark.datasets.common import BaseDataset, DatasetRow from sglang.benchmark.datasets.image import create_mm_data_row from sglang.benchmark.utils import get_processor @dataclass class MMMUDataset(BaseDataset): num_requests: int backend: str fixed_output_len: Optional[int] @classmethod def from_args(cls, args: Namespace) -> "MMMUDataset": return cls( num_requests=args.num_prompts, backend=args.backend, fixed_output_len=args.random_output_len, ) def load(self, tokenizer=None, model_id=None) -> List[DatasetRow]: processor = get_processor(model_id) return sample_mmmu_requests( num_requests=self.num_requests, processor=processor, backend=self.backend, fixed_output_len=self.fixed_output_len, ) def sample_mmmu_requests( num_requests: int, processor: AutoProcessor | AutoTokenizer, backend: str = "sglang", fixed_output_len: Optional[int] = None, random_sample: bool = True, ) -> List[DatasetRow]: """ Sample requests from the MMMU dataset using HuggingFace datasets. Args: num_requests: Number of requests to sample. fixed_output_len: If provided, use this fixed output length for all requests. random_sample: Whether to randomly sample or take the first N. Returns: List of tuples (prompt, prompt_token_len, output_token_len). """ print("Loading MMMU dataset from HuggingFace...") try: print("Attempting to load MMMU Math dataset...") mmmu_dataset = load_dataset("MMMU/MMMU", "Math", split="test") print( f"Successfully loaded MMMU Math dataset from HuggingFace with {len(mmmu_dataset)} examples" ) except Exception as e: print(f"Failed to load MMMU Math dataset: {e}") raise ValueError(f"Failed to load MMMU dataset: {e}") # Sample from the dataset if len(mmmu_dataset) > num_requests: if random_sample: # Random sample indices = random.sample(range(len(mmmu_dataset)), num_requests) sample_dataset = mmmu_dataset.select(indices) else: # Take first N sample_dataset = mmmu_dataset.select( range(min(num_requests, len(mmmu_dataset))) ) else: print(f"Dataset has less than {num_requests} examples, using all examples") sample_dataset = mmmu_dataset print(f"Selected {len(sample_dataset)} examples for benchmarking") # Create prompts filtered_dataset = [] for i, example in enumerate(sample_dataset): try: # Extract image_1 image = example.get("image_1") if image is not None: if hasattr(image, "save"): # Convert RGBA images to RGB before encoding if image.mode == "RGBA": image = image.convert("RGB") # Encode image to base64 (save as PNG to support palette/alpha modes) buffered = io.BytesIO() image.save(buffered, format="PNG") img_str = pybase64.b64encode(buffered.getvalue()).decode("utf-8") image_data = f"data:image/png;base64,{img_str}" else: continue # Extract the question question = example.get("question") # Construct the prompt text_prompt = f"Question: {question}\n\nAnswer: " output_len = fixed_output_len if fixed_output_len is not None else 256 data_row = create_mm_data_row( text_prompt, [image], [image_data], output_len, processor, backend ) filtered_dataset.append(data_row) except Exception as e: print(f"Error processing example {i}: {e}") print(f"\nCreated {len(filtered_dataset)} MMMU prompts") return filtered_dataset ================================================ FILE: python/sglang/benchmark/datasets/mooncake.py ================================================ import asyncio import json import os import time from argparse import Namespace from dataclasses import dataclass from typing import AsyncGenerator, Dict, List from transformers import PreTrainedTokenizerBase from sglang.benchmark.datasets.common import ( MOONCAKE_DATASET_URL, BaseDataset, DatasetRow, ) from sglang.benchmark.utils import download_and_cache_file @dataclass class MooncakeDataset(BaseDataset): dataset_path: str mooncake_workload: str num_requests: int @classmethod def from_args(cls, args: Namespace) -> "MooncakeDataset": return cls( dataset_path=args.dataset_path, mooncake_workload=args.mooncake_workload, num_requests=args.num_prompts, ) def load(self, tokenizer=None, model_id=None) -> List[Dict]: if not self.dataset_path: local_path = os.path.join("/tmp", self.mooncake_workload + "_trace.jsonl") else: local_path = self.dataset_path if not os.path.exists(local_path): download_and_cache_file( MOONCAKE_DATASET_URL[self.mooncake_workload], local_path ) with open(local_path, "r") as f: all_requests_data = [json.loads(line) for line in f if line.strip()] return all_requests_data[: self.num_requests] async def get_mooncake_request_over_time( input_requests: List[Dict], tokenizer: PreTrainedTokenizerBase, slowdown_factor: float, num_rounds: int, ) -> AsyncGenerator[DatasetRow, None]: """ An async generator that yields requests based on the timestamps in the Mooncake trace file, with support for multi-round sessions. """ if not input_requests: return input_requests.sort(key=lambda r: r["timestamp"]) start_time = time.perf_counter() trace_start_time_ms = input_requests[0]["timestamp"] for record in input_requests: # Calculate when this entire session should start relative_arrival_time_s = (record["timestamp"] - trace_start_time_ms) / 1000.0 target_arrival_time_s = relative_arrival_time_s * slowdown_factor current_elapsed_time_s = time.perf_counter() - start_time sleep_duration_s = target_arrival_time_s - current_elapsed_time_s if sleep_duration_s > 0: await asyncio.sleep(sleep_duration_s) # Once the session starts, generate all rounds for it as a burst # This simulates a user engaging in a multi-turn conversation # Base user query constructed from hash_ids user_query_base = "" hash_ids = record.get("hash_ids", []) for hash_id in hash_ids: user_query_base += f"{hash_id}" + " ".join( ["hi"] * 128 ) # Shorter for multi-round user_query_base += "Tell me a story based on this context." output_len_per_round = record.get("output_length", 256) chat_history = [] for i in range(num_rounds): # Add user query for the current round chat_history.append( {"role": "user", "content": f"Round {i + 1}: {user_query_base}"} ) # Form the full prompt from history try: full_prompt_text = tokenizer.apply_chat_template( chat_history, tokenize=False, add_generation_prompt=True, return_dict=False, ) except Exception: full_prompt_text = "\n".join( [f"{msg['role']}: {msg['content']}" for msg in chat_history] ) prompt_len = len(tokenizer.encode(full_prompt_text)) yield DatasetRow( prompt=full_prompt_text, prompt_len=prompt_len, output_len=output_len_per_round, ) # Add a placeholder assistant response for the next round's context # We use a placeholder because we don't know the real response placeholder_response = " ".join(["story"] * output_len_per_round) chat_history.append({"role": "assistant", "content": placeholder_response}) ================================================ FILE: python/sglang/benchmark/datasets/openai_dataset.py ================================================ import json from argparse import Namespace from dataclasses import dataclass from typing import List, Optional import numpy as np from transformers import PreTrainedTokenizerBase from sglang.benchmark.datasets.common import BaseDataset, DatasetRow @dataclass class OpenAIDataset(BaseDataset): dataset_path: str num_requests: int fixed_output_len: Optional[int] @classmethod def from_args(cls, args: Namespace) -> "OpenAIDataset": return cls( dataset_path=args.dataset_path, num_requests=args.num_prompts, fixed_output_len=args.sharegpt_output_len, ) def load( self, tokenizer: PreTrainedTokenizerBase, model_id=None ) -> List[DatasetRow]: return sample_openai_requests( dataset_path=self.dataset_path, num_requests=self.num_requests, tokenizer=tokenizer, fixed_output_len=self.fixed_output_len, ) def sample_openai_requests( dataset_path: str, num_requests: int, tokenizer: PreTrainedTokenizerBase, fixed_output_len: Optional[int] = None, ) -> List[DatasetRow]: """ Load OpenAI-compatible chat completion requests from a JSONL file. Each line should be a JSON object with: - "messages": list of {"role": str, "content": str} - "max_tokens": int (used as output_len if fixed_output_len not set) - "tools": optional list of tool definitions - "temperature": optional temperature value - "top_p": optional top_p value - Other OpenAI API parameters are also extracted and passed through """ dataset = [] with open(dataset_path, "r") as f: for line in f: if num_requests > 0 and len(dataset) >= num_requests: break if line.strip(): try: dataset.append(json.loads(line)) except json.JSONDecodeError: # Skip invalid JSON lines continue # Fields that should NOT be passed through extra_request_body # These are either handled separately or are metadata # max_tokens is excluded because it's handled via output_len -> max_completion_tokens # max_completion_tokens is also excluded to avoid conflicts EXCLUDED_FIELDS = {"messages", "max_tokens", "max_completion_tokens", "model"} filtered_dataset: List[DatasetRow] = [] for data in dataset: messages = data.get("messages", []) if not messages: continue # Use max_tokens from the request, or fall back to fixed_output_len output_len = fixed_output_len or data.get("max_tokens", 256) # Extract extra request body parameters (tools, temperature, top_p, etc.) extra_body = {k: v for k, v in data.items() if k not in EXCLUDED_FIELDS} # Calculate prompt length by applying chat template # This includes the messages but not the tools prompt_len = len( tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=True ) ) # If tools are present, we need to add their token count # Tools are sent as part of the request and count toward input tokens if "tools" in extra_body: # Encode tools as JSON string to estimate token count tools_str = json.dumps(extra_body["tools"]) tools_tokens = len(tokenizer.encode(tools_str)) prompt_len += tools_tokens # Pass messages list directly - bench_serving handles List[Dict] prompts filtered_dataset.append( DatasetRow( prompt=messages, prompt_len=prompt_len, output_len=output_len, extra_request_body=extra_body, # Store per-request parameters ) ) print(f"Loaded {len(filtered_dataset)} OpenAI-format requests") print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}") print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}") return filtered_dataset ================================================ FILE: python/sglang/benchmark/datasets/random.py ================================================ import json import random from argparse import Namespace from dataclasses import dataclass from typing import List import numpy as np from transformers import PreTrainedTokenizerBase from sglang.benchmark.datasets.common import ( SHAREGPT_FILENAME, SHAREGPT_REPO_ID, BaseDataset, DatasetRow, compute_random_lens, ) from sglang.benchmark.utils import download_and_cache_hf_file, is_file_valid_json @dataclass class RandomDataset(BaseDataset): input_len: int output_len: int num_requests: int range_ratio: float dataset_path: str return_text: bool random_sample: bool @classmethod def from_args(cls, args: Namespace) -> "RandomDataset": return cls( input_len=args.random_input_len, output_len=args.random_output_len, num_requests=args.num_prompts, range_ratio=args.random_range_ratio, dataset_path=args.dataset_path, return_text=not getattr(args, "tokenize_prompt", False), random_sample=(args.dataset_name == "random"), ) def load( self, tokenizer: PreTrainedTokenizerBase, model_id=None ) -> List[DatasetRow]: return sample_random_requests( input_len=self.input_len, output_len=self.output_len, num_prompts=self.num_requests, range_ratio=self.range_ratio, tokenizer=tokenizer, dataset_path=self.dataset_path, random_sample=self.random_sample, return_text=self.return_text, ) def sample_random_requests( input_len: int, output_len: int, num_prompts: int, range_ratio: float, tokenizer: PreTrainedTokenizerBase, dataset_path: str, random_sample: bool = True, return_text: bool = True, ) -> List[DatasetRow]: input_lens = compute_random_lens( full_len=input_len, range_ratio=range_ratio, num=num_prompts, ) output_lens = compute_random_lens( full_len=output_len, range_ratio=range_ratio, num=num_prompts, ) if return_text: # Need to truncate input_len as server encode will add special token. num_special_tokens = int(tokenizer.num_special_tokens_to_add()) for i in range(num_prompts): input_lens[i] = max(0, input_lens[i] - num_special_tokens) if random_sample: # Sample token ids from ShareGPT and repeat/truncate them to satisfy the input_lens # Download sharegpt if necessary if not is_file_valid_json(dataset_path): dataset_path = download_and_cache_hf_file( repo_id=SHAREGPT_REPO_ID, filename=SHAREGPT_FILENAME, ) # Load the dataset. with open(dataset_path) as f: dataset = json.load(f) # Filter out the conversations with less than 2 turns. dataset = [ data for data in dataset if len(data.get("conversations", data.get("conversation", []))) >= 2 ] # Only keep the first two turns of each conversation. dataset = [ ( data.get("conversations", data.get("conversation", []))[0]["value"], data.get("conversations", data.get("conversation", []))[1]["value"], ) for data in dataset ] # Shuffle the dataset. random.shuffle(dataset) # Filter out sequences that are too long or too short input_requests: List[DatasetRow] = [] for data in dataset: i = len(input_requests) if i == num_prompts: break # Tokenize the prompts and completions. prompt = data[0] prompt_token_ids = tokenizer.encode(prompt) prompt_len = len(prompt_token_ids) # Skip empty prompt if prompt_len == 0: continue if prompt_len > input_lens[i]: input_ids = prompt_token_ids[: input_lens[i]] else: ratio = (input_lens[i] + prompt_len - 1) // prompt_len input_ids = (prompt_token_ids * ratio)[: input_lens[i]] input_content = input_ids if return_text: input_content = tokenizer.decode(input_content) input_requests.append( DatasetRow( prompt=input_content, prompt_len=input_lens[i], output_len=output_lens[i], ) ) else: # Sample token ids from random integers. This can cause some NaN issues. offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts) input_requests = [] for i in range(num_prompts): # Use int() to convert numpy.int64 to native Python int for JSON serialization input_content = [ int((offsets[i] + i + j) % tokenizer.vocab_size) for j in range(input_lens[i]) ] if return_text: input_content = tokenizer.decode(input_content) input_requests.append( DatasetRow( prompt=input_content, prompt_len=input_lens[i], output_len=output_lens[i], ) ) print(f"#Input tokens: {np.sum(input_lens)}") print(f"#Output tokens: {np.sum(output_lens)}") return input_requests ================================================ FILE: python/sglang/benchmark/datasets/sharegpt.py ================================================ import json import random from argparse import Namespace from dataclasses import dataclass from typing import List, Optional import numpy as np from transformers import PreTrainedTokenizerBase from sglang.benchmark.datasets.common import ( ASSISTANT_SUFFIX, SHAREGPT_FILENAME, SHAREGPT_REPO_ID, BaseDataset, DatasetRow, ) from sglang.benchmark.utils import ( download_and_cache_hf_file, is_file_valid_json, remove_suffix, ) @dataclass class ShareGPTDataset(BaseDataset): dataset_path: str num_requests: int fixed_output_len: Optional[int] context_len: Optional[int] prompt_suffix: str apply_chat_template: bool @classmethod def from_args(cls, args: Namespace) -> "ShareGPTDataset": assert not getattr(args, "tokenize_prompt", False) return cls( dataset_path=args.dataset_path, num_requests=args.num_prompts, fixed_output_len=args.sharegpt_output_len, context_len=args.sharegpt_context_len, prompt_suffix=args.prompt_suffix, apply_chat_template=args.apply_chat_template, ) def load( self, tokenizer: PreTrainedTokenizerBase, model_id=None ) -> List[DatasetRow]: return sample_sharegpt_requests( dataset_path=self.dataset_path, num_requests=self.num_requests, tokenizer=tokenizer, fixed_output_len=self.fixed_output_len, context_len=self.context_len, prompt_suffix=self.prompt_suffix, apply_chat_template=self.apply_chat_template, ) def sample_sharegpt_requests( dataset_path: str, num_requests: int, tokenizer: PreTrainedTokenizerBase, fixed_output_len: Optional[int] = None, context_len: Optional[int] = None, prompt_suffix: Optional[str] = "", apply_chat_template=False, ) -> List[DatasetRow]: if fixed_output_len is not None and fixed_output_len < 4: raise ValueError("output_len too small") # Download sharegpt if necessary if not is_file_valid_json(dataset_path) and dataset_path == "": dataset_path = download_and_cache_hf_file( repo_id=SHAREGPT_REPO_ID, filename=SHAREGPT_FILENAME, ) # Load the dataset. with open(dataset_path) as f: dataset = json.load(f) # Filter out the conversations with less than 2 turns. dataset = [ data for data in dataset if len(data.get("conversations", data.get("conversation", []))) >= 2 ] # Only keep the first two turns of each conversation. dataset = [ ( data.get("conversations", data.get("conversation", []))[0]["value"], data.get("conversations", data.get("conversation", []))[1]["value"], ) for data in dataset ] # Shuffle the dataset. random.shuffle(dataset) # Filter out sequences that are too long or too short filtered_dataset: List[DatasetRow] = [] for i in range(len(dataset)): if len(filtered_dataset) == num_requests: break # Tokenize the prompts and completions. prompt = dataset[i][0] if prompt_suffix: prompt = ( remove_suffix(prompt, ASSISTANT_SUFFIX) + prompt_suffix + ASSISTANT_SUFFIX ) if apply_chat_template: prompt = tokenizer.apply_chat_template( [{"role": "user", "content": prompt}], add_generation_prompt=True, tokenize=False, return_dict=False, ) if tokenizer.bos_token: prompt = prompt.replace(tokenizer.bos_token, "") prompt_token_ids = tokenizer.encode(prompt) completion = dataset[i][1] completion_token_ids = tokenizer.encode(completion) prompt_len = len(prompt_token_ids) output_len = ( len(completion_token_ids) if fixed_output_len is None else fixed_output_len ) if prompt_len < 2 or output_len < 2: # Prune too short sequences. continue if context_len and prompt_len + output_len > context_len: # Prune too long sequences. continue filtered_dataset.append( DatasetRow( prompt=prompt, prompt_len=prompt_len, output_len=output_len, ) ) print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}") print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}") return filtered_dataset ================================================ FILE: python/sglang/benchmark/utils.py ================================================ import json import os import resource from json import JSONDecodeError from typing import Dict, List, Optional, Union import requests from tqdm.asyncio import tqdm from transformers import ( AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, ) def remove_prefix(text: str, prefix: str) -> str: return text[len(prefix) :] if text.startswith(prefix) else text def remove_suffix(text: str, suffix: str) -> str: return text[: -len(suffix)] if text.endswith(suffix) else text def parse_custom_headers(header_list: List[str]) -> Dict[str, str]: return {k: v for h in header_list for k, _, v in [h.partition("=")] if k and v} def get_model(pretrained_model_name_or_path: str) -> str: if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() == "true": import huggingface_hub.constants from modelscope import snapshot_download model_path = snapshot_download( model_id=pretrained_model_name_or_path, local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"], ) return model_path return pretrained_model_name_or_path def get_tokenizer( pretrained_model_name_or_path: str, ) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: assert ( pretrained_model_name_or_path is not None and pretrained_model_name_or_path != "" ) if pretrained_model_name_or_path.endswith( ".json" ) or pretrained_model_name_or_path.endswith(".model"): from sglang.srt.utils.hf_transformers_utils import get_tokenizer return get_tokenizer(pretrained_model_name_or_path) if pretrained_model_name_or_path is not None and not os.path.exists( pretrained_model_name_or_path ): pretrained_model_name_or_path = get_model(pretrained_model_name_or_path) return AutoTokenizer.from_pretrained( pretrained_model_name_or_path, trust_remote_code=True ) def get_processor( pretrained_model_name_or_path: str, ) -> AutoProcessor: assert ( pretrained_model_name_or_path is not None and pretrained_model_name_or_path != "" ) if pretrained_model_name_or_path.endswith( ".json" ) or pretrained_model_name_or_path.endswith(".model"): from sglang.srt.utils.hf_transformers_utils import get_processor return get_processor(pretrained_model_name_or_path) if pretrained_model_name_or_path is not None and not os.path.exists( pretrained_model_name_or_path ): pretrained_model_name_or_path = get_model(pretrained_model_name_or_path) return AutoProcessor.from_pretrained( pretrained_model_name_or_path, trust_remote_code=True ) def download_and_cache_hf_file( repo_id: str, filename: str, repo_type: str = "dataset", ): """Download a file from Hugging Face and cache it locally.""" from huggingface_hub import hf_hub_download return hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type) def download_and_cache_file(url: str, filename: Optional[str] = None): """Read and cache a file from a url.""" if filename is None: filename = os.path.join("/tmp", url.split("/")[-1]) # Check if the cache file already exists if is_file_valid_json(filename): return filename print(f"Downloading from {url} to {filename}") # Stream the response to show the progress bar response = requests.get(url, stream=True) response.raise_for_status() # Check for request errors # Total size of the file in bytes total_size = int(response.headers.get("content-length", 0)) chunk_size = 1024 # Download in chunks of 1KB # Use tqdm to display the progress bar with open(filename, "wb") as f, tqdm( desc=filename, total=total_size, unit="B", unit_scale=True, unit_divisor=1024, ) as bar: for chunk in response.iter_content(chunk_size=chunk_size): f.write(chunk) bar.update(len(chunk)) return filename def is_file_valid_json(path): if not os.path.isfile(path): return False # TODO can fuse into the real file open later try: with open(path) as f: json.load(f) return True except JSONDecodeError as e: print( f"{path} exists but json loading fails ({e=}), thus treat as invalid file" ) return False def set_ulimit(target_soft_limit=65535): resource_type = resource.RLIMIT_NOFILE current_soft, current_hard = resource.getrlimit(resource_type) if current_soft < target_soft_limit: try: resource.setrlimit(resource_type, (target_soft_limit, current_hard)) except ValueError as e: print(f"Fail to set RLIMIT_NOFILE: {e}") ================================================ FILE: python/sglang/check_env.py ================================================ """Check environment configurations and dependency versions.""" import importlib.metadata import os import resource import subprocess import sys from abc import abstractmethod from collections import OrderedDict, defaultdict import torch from sglang.srt.utils import is_hip, is_musa, is_npu def is_cuda_v2(): return torch.version.cuda is not None # List of packages to check versions PACKAGE_LIST = [ "sglang", "sglang-kernel", "flashinfer_python", "flashinfer_cubin", "flashinfer_jit_cache", "triton", "transformers", "torchao", "numpy", "aiohttp", "fastapi", "huggingface_hub", "interegular", "modelscope", "orjson", "outlines", "packaging", "psutil", "pydantic", "python-multipart", "pyzmq", "torchao", "uvicorn", "uvloop", "vllm", "xgrammar", "openai", "tiktoken", "anthropic", "litellm", "torchcodec", ] class BaseEnv: """Base class for environment check""" def __init__(self): self.package_list = PACKAGE_LIST @abstractmethod def get_info(self) -> dict: """ Get CUDA-related information if available. """ raise NotImplementedError @abstractmethod def get_topology(self) -> dict: raise NotImplementedError def get_package_versions(self) -> dict: """ Get versions of specified packages. """ versions = {} for package in self.package_list: package_name = package.split("==")[0].split(">=")[0].split("<=")[0] try: version = importlib.metadata.version(package_name) versions[package_name] = version except ModuleNotFoundError: versions[package_name] = "Module Not Found" return versions def get_device_info(self): """ Get information about available GPU devices. """ devices = defaultdict(list) capabilities = defaultdict(list) for k in range(torch.cuda.device_count()): devices[torch.cuda.get_device_name(k)].append(str(k)) capability = torch.cuda.get_device_capability(k) capabilities[f"{capability[0]}.{capability[1]}"].append(str(k)) gpu_info = {} for name, device_ids in devices.items(): gpu_info[f"GPU {','.join(device_ids)}"] = name if len(capabilities) == 1: # All GPUs have the same compute capability cap, gpu_ids = list(capabilities.items())[0] gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap else: # GPUs have different compute capabilities for cap, gpu_ids in capabilities.items(): gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap return gpu_info def get_hypervisor_vendor(self) -> dict: try: output = subprocess.check_output(["lscpu"], text=True) for line in output.split("\n"): if "Hypervisor vendor:" in line: return {"Hypervisor vendor:": line.split(":")[1].strip()} return {} except: return {} def get_ulimit_soft(self) -> dict: ulimit_soft, _ = resource.getrlimit(resource.RLIMIT_NOFILE) return {"ulimit soft": ulimit_soft} def check_env(self): """ Check and print environment information. """ env_info = OrderedDict() env_info["Python"] = sys.version.replace("\n", "") env_info.update(self.get_info()) env_info["PyTorch"] = torch.__version__ env_info.update(self.get_package_versions()) env_info.update(self.get_topology()) env_info.update(self.get_hypervisor_vendor()) env_info.update(self.get_ulimit_soft()) for k, v in env_info.items(): print(f"{k}: {v}") class GPUEnv(BaseEnv): """Environment checker for Nvidia GPU""" def get_info(self): cuda_info = {"CUDA available": torch.cuda.is_available()} if cuda_info["CUDA available"]: cuda_info.update(self.get_device_info()) cuda_info.update(self._get_cuda_version_info()) return cuda_info def _get_cuda_version_info(self): """ Get CUDA version information. """ from torch.utils.cpp_extension import CUDA_HOME cuda_info = {"CUDA_HOME": CUDA_HOME} if CUDA_HOME and os.path.isdir(CUDA_HOME): cuda_info.update(self._get_nvcc_info()) cuda_info.update(self._get_cuda_driver_version()) return cuda_info def _get_nvcc_info(self): """ Get NVCC version information. """ from torch.utils.cpp_extension import CUDA_HOME try: nvcc = os.path.join(CUDA_HOME, "bin/nvcc") nvcc_output = ( subprocess.check_output(f'"{nvcc}" -V', shell=True) .decode("utf-8") .strip() ) return { "NVCC": nvcc_output[ nvcc_output.rfind("Cuda compilation tools") : nvcc_output.rfind( "Build" ) ].strip() } except subprocess.SubprocessError: return {"NVCC": "Not Available"} def _get_cuda_driver_version(self): """ Get CUDA driver version. """ versions = set() try: output = subprocess.check_output( [ "nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader,nounits", ] ) versions = set(output.decode().strip().split("\n")) if len(versions) == 1: return {"CUDA Driver Version": versions.pop()} else: return {"CUDA Driver Versions": ", ".join(sorted(versions))} except subprocess.SubprocessError: return {"CUDA Driver Version": "Not Available"} def get_topology(self): """ Get GPU topology information. """ try: result = subprocess.run( ["nvidia-smi", "topo", "-m"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True, ) return { "NVIDIA Topology": ( "\n" + result.stdout if result.returncode == 0 else None ) } except subprocess.SubprocessError: return {} class HIPEnv(BaseEnv): """Environment checker for ROCm/HIP""" def get_info(self): cuda_info = {"ROCM available": torch.cuda.is_available()} if cuda_info["ROCM available"]: cuda_info.update(self.get_device_info()) cuda_info.update(self._get_cuda_version_info()) return cuda_info def _get_cuda_version_info(self): from torch.utils.cpp_extension import ROCM_HOME as ROCM_HOME cuda_info = {"ROCM_HOME": ROCM_HOME} if ROCM_HOME and os.path.isdir(ROCM_HOME): cuda_info.update(self._get_hipcc_info()) cuda_info.update(self._get_rocm_driver_version()) return cuda_info def _get_hipcc_info(self): from torch.utils.cpp_extension import ROCM_HOME try: hipcc = os.path.join(ROCM_HOME, "bin/hipcc") hipcc_output = ( subprocess.check_output(f'"{hipcc}" --version', shell=True) .decode("utf-8") .strip() ) return { "HIPCC": hipcc_output[ hipcc_output.rfind("HIP version") : hipcc_output.rfind("AMD clang") ].strip() } except subprocess.SubprocessError: return {"HIPCC": "Not Available"} def _get_rocm_driver_version(self): try: output = subprocess.check_output( [ "rocm-smi", "--showdriverversion", "--csv", ] ) versions = set(output.decode().strip().split("\n")) versions.discard("name, value") ver = versions.pop() ver = ver.replace('"Driver version", ', "").replace('"', "") return {"ROCM Driver Version": ver} except subprocess.SubprocessError: return {"ROCM Driver Version": "Not Available"} def get_topology(self): try: result = subprocess.run( ["rocm-smi", "--showtopotype"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True, ) return { "AMD Topology": "\n" + result.stdout if result.returncode == 0 else None } except subprocess.SubprocessError: return {} class NPUEnv(BaseEnv): """Environment checker for Ascend NPU""" EXTRA_PACKAGE_LIST = [ "torch_npu", "sgl-kernel-npu", "deep_ep", ] def __init__(self): super().__init__() self.package_list.extend(NPUEnv.EXTRA_PACKAGE_LIST) def get_info(self): cuda_info = {"NPU available": torch.npu.is_available()} if cuda_info["NPU available"]: cuda_info.update(self.get_device_info()) cuda_info.update(self._get_cann_version_info()) return cuda_info def get_device_info(self): """ Get information about available NPUs. Need to override due to torch_npu interface differences. """ devices = defaultdict(list) for k in range(torch.npu.device_count()): devices[torch.npu.get_device_name(k)].append(str(k)) npu_info = {} for name, device_ids in devices.items(): npu_info[f"NPU {','.join(device_ids)}"] = name return npu_info def _get_cann_version_info(self): cann_envs = ["ASCEND_TOOLKIT_HOME", "ASCEND_INSTALL_PATH"] for var in cann_envs: path = os.environ.get(var) if path and os.path.exists(path): CANN_HOME = path break else: default_path = "/usr/local/Ascend/ascend-toolkit/latest" CANN_HOME = default_path if os.path.exists(default_path) else None if CANN_HOME: npu_info = {"CANN_HOME": CANN_HOME} npu_info.update(self._get_cann_info(CANN_HOME)) npu_info.update(self._get_ascend_driver_version()) return npu_info else: return {"CANN_HOME": "Not found"} def _get_cann_info(self, CANN_HOME: str): cann_info = {} cann_version_file = os.path.join(CANN_HOME, "version.cfg") if os.path.exists(cann_version_file): with open(cann_version_file, "r", encoding="utf-8") as f: f.readline() # discard first line comment in version.cfg cann_info["CANN"] = f.readline().split("[")[1].split("]")[0] else: cann_info["CANN"] = "Not Available" try: bisheng = os.path.join(CANN_HOME, "compiler/ccec_compiler/bin/bisheng") bisheng_output = ( subprocess.check_output([bisheng, "--version"]).decode("utf-8").strip() ) cann_info["BiSheng"] = bisheng_output.split("\n")[0].strip() except subprocess.SubprocessError: cann_info["BiSheng"] = "Not Available" return cann_info def _get_ascend_driver_version(self): try: output = subprocess.check_output( [ "npu-smi", "info", "-t", "board", "-i", "0", ] ) for line in output.decode().strip().split("\n"): if "Software Version" in line: version = line.split(":")[-1].strip() break else: version = "Not Available" return {"Ascend Driver Version": version} except subprocess.SubprocessError: return {"Ascend Driver Version": "Not Available"} def get_topology(self): try: result = subprocess.run( ["npu-smi", "info", "-t", "topo"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True, ) return { "Ascend Topology": ( "\n" + result.stdout if result.returncode == 0 else None ) } except subprocess.SubprocessError: return {} class MUSAEnv(BaseEnv): """Environment checker for MThreads GPU""" def get_info(self): musa_info = {"MUSA available": torch.musa.is_available()} if musa_info["MUSA available"]: musa_info.update(self.get_device_info()) musa_info.update(self._get_musa_version_info()) return musa_info def _get_musa_version_info(self): """ Get MUSA version information. """ from torch_musa.utils.musa_extension import MUSA_HOME musa_info = {"MUSA_HOME": MUSA_HOME} if MUSA_HOME and os.path.isdir(MUSA_HOME): musa_info.update(self._get_mcc_info()) musa_info.update(self._get_musa_driver_version()) return musa_info def _get_mcc_info(self): """ Get MCC version information. """ from torch_musa.utils.musa_extension import MUSA_HOME try: mcc = os.path.join(MUSA_HOME, "bin/mcc") mcc_output = ( subprocess.check_output(f'"{mcc}" --version', shell=True) .decode("utf-8") .strip() ) return { "MCC": mcc_output[ mcc_output.rfind("mcc version") : mcc_output.rfind("Target") ].strip() } except subprocess.SubprocessError: return {"MCC": "Not Available"} def _get_musa_driver_version(self): """ Get MUSA driver version. """ try: output = subprocess.check_output( [ "mthreads-gmi", "-q", ], text=True, ) driver_version = None for line in output.splitlines(): if "Driver Version" in line: driver_version = line.split(":", 1)[1].strip() break return {"MUSA Driver Version": driver_version} except subprocess.SubprocessError: return {"MUSA Driver Version": "Not Available"} def get_topology(self): """ Get GPU topology information. """ try: result = subprocess.run( ["mthreads-gmi", "topo", "-m"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True, ) return { "MTHREADS Topology": ( "\n" + result.stdout if result.returncode == 0 else None ) } except subprocess.SubprocessError: return {} if __name__ == "__main__": if is_cuda_v2(): env = GPUEnv() elif is_hip(): env = HIPEnv() elif is_npu(): env = NPUEnv() elif is_musa(): env = MUSAEnv() env.check_env() ================================================ FILE: python/sglang/cli/__init__.py ================================================ ================================================ FILE: python/sglang/cli/generate.py ================================================ import argparse from sglang.cli.utils import get_is_diffusion_model, get_model_path def generate(args, extra_argv): # If help is requested, show generate subcommand help without requiring --model-path if any(h in extra_argv for h in ("-h", "--help")): from sglang.multimodal_gen.runtime.entrypoints.cli.generate import ( add_multimodal_gen_generate_args, ) parser = argparse.ArgumentParser(description="SGLang Multimodal Generation") add_multimodal_gen_generate_args(parser) parser.parse_args(extra_argv) return model_path = get_model_path(extra_argv) is_diffusion_model = get_is_diffusion_model(model_path) if is_diffusion_model: from sglang.multimodal_gen.runtime.entrypoints.cli.generate import ( add_multimodal_gen_generate_args, generate_cmd, ) parser = argparse.ArgumentParser(description="SGLang Multimodal Generation") add_multimodal_gen_generate_args(parser) parsed_args, unknown_args = parser.parse_known_args(extra_argv) generate_cmd(parsed_args, unknown_args) else: raise Exception( f"Generate subcommand is not yet supported for model: {model_path}" ) ================================================ FILE: python/sglang/cli/main.py ================================================ import argparse from sglang.cli.utils import get_git_commit_hash from sglang.version import __version__ def version(args, extra_argv): print(f"sglang version: {__version__}") print(f"git revision: {get_git_commit_hash()[:7]}") def main(): parser = argparse.ArgumentParser() # complex sub commands subparsers = parser.add_subparsers(dest="subcommand", required=True) subparsers.add_parser( "serve", help="Launch an SGLang server.", add_help=False, ) subparsers.add_parser( "generate", help="Run inference on a multimodal model.", add_help=False, ) # simple commands version_parser = subparsers.add_parser( "version", help="Show the version information.", ) version_parser.set_defaults(func=version) args, extra_argv = parser.parse_known_args() if args.subcommand == "serve": from sglang.cli.serve import serve serve(args, extra_argv) elif args.subcommand == "generate": from sglang.cli.generate import generate generate(args, extra_argv) elif args.subcommand == "version": version(args, extra_argv) ================================================ FILE: python/sglang/cli/serve.py ================================================ # SPDX-License-Identifier: Apache-2.0 import argparse import logging import os from sglang.cli.utils import get_is_diffusion_model, get_model_path from sglang.srt.utils import kill_process_tree from sglang.srt.utils.common import suppress_noisy_warnings suppress_noisy_warnings() logger = logging.getLogger(__name__) def _extract_model_type_override(extra_argv): """Extract and remove --model-type override from argv.""" model_type = "auto" filtered_argv = [] i = 0 while i < len(extra_argv): arg = extra_argv[i] if arg == "--model-type": if i + 1 >= len(extra_argv): raise Exception( "Error: --model-type requires a value. " "Valid values are: auto, llm, diffusion." ) model_type = extra_argv[i + 1] i += 2 continue if arg.startswith("--model-type="): model_type = arg.split("=", 1)[1] i += 1 continue filtered_argv.append(arg) i += 1 if model_type not in ("auto", "llm", "diffusion"): raise Exception( f"Error: invalid --model-type '{model_type}'. " "Valid values are: auto, llm, diffusion." ) return model_type, filtered_argv def serve(args, extra_argv): if any(h in extra_argv for h in ("-h", "--help")): # Since the server type is determined by the model, and we don't have a model path, # we can't show the exact help. Instead, we show a general help message and then # the help for both possible server types. print( "Usage: sglang serve --model-path [additional-arguments]\n\n" "This command can launch either a standard language model server or a diffusion model server.\n" "The server type is determined by the --model-path.\n" "Optional override: --model-type {auto,llm,diffusion} " "(default: auto, fallback to LLM on detection failure)." ) print("\n--- Help for Standard Language Model Server ---") from sglang.srt.server_args import prepare_server_args try: prepare_server_args(["--help"]) except SystemExit: pass # argparse --help calls sys.exit print("\n--- Help for Diffusion Model Server ---") try: from sglang.multimodal_gen.runtime.entrypoints.cli.serve import ( add_multimodal_gen_serve_args, ) parser = argparse.ArgumentParser( prog="sglang serve", description="SGLang Diffusion Model Serving", ) add_multimodal_gen_serve_args(parser) parser.print_help() except ImportError: print( "Diffusion model support is not available. " 'Install with: pip install "sglang[diffusion]"' ) return model_type, dispatch_argv = _extract_model_type_override(extra_argv) model_path = get_model_path(dispatch_argv) try: if model_type == "auto": is_diffusion_model = get_is_diffusion_model(model_path) if is_diffusion_model: logger.info("Diffusion model detected") else: is_diffusion_model = model_type == "diffusion" logger.info( "Dispatch override enabled: --model-type=%s " "(skip auto detection)", model_type, ) if is_diffusion_model: # Logic for Diffusion Models from sglang.multimodal_gen.runtime.entrypoints.cli.serve import ( add_multimodal_gen_serve_args, execute_serve_cmd, ) parser = argparse.ArgumentParser( description="SGLang Diffusion Model Serving" ) add_multimodal_gen_serve_args(parser) parsed_args, remaining_argv = parser.parse_known_args(dispatch_argv) execute_serve_cmd(parsed_args, remaining_argv) else: # Logic for Standard Language Models from sglang.launch_server import run_server from sglang.srt.server_args import prepare_server_args server_args = prepare_server_args(dispatch_argv) run_server(server_args) finally: kill_process_tree(os.getpid(), include_parent=False) ================================================ FILE: python/sglang/cli/utils.py ================================================ import json import logging import os import subprocess from functools import lru_cache from sglang.srt.environ import envs logger = logging.getLogger(__name__) def _is_diffusers_model_dir(model_dir: str) -> bool: """Check if a local directory contains a valid diffusers model_index.json.""" config_path = os.path.join(model_dir, "model_index.json") if not os.path.exists(config_path): return False with open(config_path) as f: config = json.load(f) return "_diffusers_version" in config def get_is_diffusion_model(model_path: str) -> bool: """Detect whether model_path points to a diffusion model. For local directories, checks the filesystem directly. For HF/ModelScope model IDs, attempts to fetch only model_index.json. Returns False on any failure (network error, 404, offline mode, etc.) so that the caller falls through to the standard LLM server path. """ try: from sglang.multimodal_gen.registry import ( is_known_non_diffusers_multimodal_model, ) except ImportError: is_known_non_diffusers_multimodal_model = lambda _: False if os.path.isdir(model_path): if _is_diffusers_model_dir(model_path): return True return is_known_non_diffusers_multimodal_model(model_path) if is_known_non_diffusers_multimodal_model(model_path): return True try: if envs.SGLANG_USE_MODELSCOPE.get(): from modelscope import model_file_download file_path = model_file_download( model_id=model_path, file_path="model_index.json" ) else: from huggingface_hub import hf_hub_download file_path = hf_hub_download(repo_id=model_path, filename="model_index.json") return _is_diffusers_model_dir(os.path.dirname(file_path)) except Exception as e: logger.debug("Failed to auto-detect diffusion model for %s: %s", model_path, e) return False def get_model_path(extra_argv): # Find the model_path argument model_path = None for i, arg in enumerate(extra_argv): if arg == "--model-path": if i + 1 < len(extra_argv): model_path = extra_argv[i + 1] break elif arg.startswith("--model-path="): model_path = arg.split("=", 1)[1] break if model_path is None: # Fallback for --help or other cases where model-path is not provided if any(h in extra_argv for h in ["-h", "--help"]): raise Exception( "Usage: sglang serve --model-path [additional-arguments]\n\n" "This command can launch either a standard language model server or a diffusion model server.\n" "The server type is determined by the --model-path.\n" ) else: raise Exception( "Error: --model-path is required. " "Please provide the path to the model." ) return model_path @lru_cache(maxsize=1) def get_git_commit_hash() -> str: try: commit_hash = os.environ.get("SGLANG_GIT_COMMIT") if not commit_hash: commit_hash = ( subprocess.check_output( ["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL ) .strip() .decode("utf-8") ) _CACHED_COMMIT_HASH = commit_hash return commit_hash except (subprocess.CalledProcessError, FileNotFoundError): _CACHED_COMMIT_HASH = "N/A" return "N/A" ================================================ FILE: python/sglang/compile_deep_gemm.py ================================================ """ Compile DeepGEMM Kernels for a model with specify server arguments This script launches a server for capturing DeepGEMM calls and then compiles the kernels. It accepts server arguments (the same as launch_server.py). Usage: python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code """ import argparse import dataclasses import multiprocessing import os import time import requests from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST from sglang.srt.entrypoints.http_server import launch_server from sglang.srt.entrypoints.warmup import warmup from sglang.srt.environ import envs from sglang.srt.managers.io_struct import GenerateReqInput from sglang.srt.managers.tokenizer_manager import TokenizerManager from sglang.srt.server_args import ServerArgs from sglang.srt.utils import kill_process_tree multiprocessing.set_start_method("spawn", force=True) # Reduce warning envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.set(True) # Force enable deep gemm envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(True) # Force enable mha chunked kv for DeepSeek V3 to avoid missing kv_b_proj DeepGEMM case envs.SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD.set(0) @dataclasses.dataclass class CompileArgs: timeout: int = 3600 @staticmethod def add_cli_args(parser: argparse.ArgumentParser): parser.add_argument("--timeout", type=int, default=CompileArgs.timeout) @classmethod def from_cli_args(cls, args: argparse.Namespace): # use the default value's type to cast the args into correct types. attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)] return cls( **{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs} ) @warmup("compile-deep-gemm") async def warm_up_compile( disaggregation_mode: str, tokenizer_manager: TokenizerManager ): print("\nGenerate warm up request for compiling DeepGEMM...\n") generate_req_input = GenerateReqInput( input_ids=[0, 1, 2, 3], sampling_params={ "temperature": 0.0, "max_new_tokens": 8, "ignore_eos": True, }, ) if disaggregation_mode != "null": generate_req_input.bootstrap_room = 0 generate_req_input.bootstrap_host = FAKE_BOOTSTRAP_HOST await tokenizer_manager.generate_request(generate_req_input, None).__anext__() def launch_server_internal(server_args): try: launch_server(server_args) except Exception as e: raise e finally: kill_process_tree(os.getpid(), include_parent=False) def launch_server_process_and_send_one_request( server_args: ServerArgs, compile_args: CompileArgs ): proc = multiprocessing.Process(target=launch_server_internal, args=(server_args,)) proc.start() base_url = f"http://{server_args.host}:{server_args.port}" timeout = compile_args.timeout start_time = time.perf_counter() while time.perf_counter() - start_time < timeout: try: headers = { "Content-Type": "application/json; charset=utf-8", } if server_args.node_rank == 0: response = requests.get(f"{base_url}/v1/models", headers=headers) else: # This http api is created by launch_dummy_health_check_server for none-rank0 node. response = requests.get(f"{base_url}/health", headers=headers) if response.status_code == 200: # Rank-0 node send a request to sync with other node and then return. if server_args.node_rank == 0: payload = { "input_ids": [0, 1, 2, 3], "sampling_params": { "max_new_tokens": 8, "temperature": 0, }, } # In PD mode, include fake bootstrap fields so workers don't assert if server_args.disaggregation_mode != "null": payload["bootstrap_host"] = FAKE_BOOTSTRAP_HOST payload["bootstrap_room"] = 0 response = requests.post( f"{base_url}/generate", json=payload, timeout=600, ) if response.status_code != 200: error = response.json() raise RuntimeError(f"Sync request failed: {error}") # Other nodes should wait for the exit signal from Rank-0 node. else: start_time_waiting = time.perf_counter() while proc.is_alive(): if time.perf_counter() - start_time_waiting < timeout: time.sleep(10) else: raise TimeoutError("Waiting for main node timeout!") return proc except requests.RequestException: pass time.sleep(10) raise TimeoutError( "DeepGEMM Kernels compilation timeout." "\n\nFeel free and please restart the command." ) def refine_server_args(server_args: ServerArgs, compile_args: CompileArgs): # Disable cuda graph and torch compile to save time server_args.disable_cuda_graph = True server_args.enable_torch_compile = False print(f"Disable CUDA Graph and Torch Compile to save time...") # Set watchdog timeout to compile_args.timeout because compilation will take a long time server_args.watchdog_timeout = compile_args.timeout server_args.warmups = "compile-deep-gemm" def run_compile(server_args: ServerArgs, compile_args: CompileArgs): print( "Begin DeepGEMM Kernels compilation...\n" "It may take a long time and timeout maybe raised " "while the compilation is still in progress.\n" "Just feel free to restart the command " "until the compilation is fully finished.\n" ) proc = launch_server_process_and_send_one_request(server_args, compile_args) print("\nDeepGEMM Kernels compilation finished successfully.") # Sleep for safety time.sleep(10) if proc.is_alive(): # This is the rank0 node. kill_process_tree(proc.pid) else: try: kill_process_tree(proc.pid) except Exception: pass if __name__ == "__main__": parser = argparse.ArgumentParser() ServerArgs.add_cli_args(parser) CompileArgs.add_cli_args(parser) args = parser.parse_args() server_args = ServerArgs.from_cli_args(args) compile_args = CompileArgs.from_cli_args(args) refine_server_args(server_args, compile_args) run_compile(server_args, compile_args) ================================================ FILE: python/sglang/eval/llama3_eval.py ================================================ # Adapt from https://github.com/fw-ai/llm_eval_meta import argparse import asyncio import os import pickle import re import shutil from collections import defaultdict from dataclasses import dataclass import httpx import numpy as np import openai from datasets import load_dataset from openai import AsyncOpenAI from tqdm import tqdm # Mapping providers to their clients and models provider_to_models = { "b10": { "8b": "meta-llama/Llama-3.1-8B-Instruct", "70b": "meta-llama/Llama-3.1-70B-Instruct", "405b": "meta-llama/Llama-3.1-405B-Instruct", }, "oai": { "8b": "meta-llama/Llama-3.1-8B-Instruct", "70b": "meta-llama/Llama-3.1-70B-Instruct", "405b": "meta-llama/Llama-3.1-405B-Instruct", }, "sgl": { "8b": "meta-llama/Llama-3.1-8B-Instruct", "70b": "meta-llama/Llama-3.1-70B-Instruct", "405b": "meta-llama/Llama-3.1-405B-Instruct", }, } async def fetch_responses( client, prompt, semaphore, index, provider, model_size, output_dir, max_tokens ): output_file = os.path.join(output_dir, f"response_{index}.pkl") if os.path.exists(output_file): print(f"File {output_file} already exists, skipping.") return async with semaphore: response = await client.completions.create( model=provider_to_models[provider][model_size], prompt=prompt, temperature=0.0, max_tokens=max_tokens, ) if isinstance(response, openai.BadRequestError): with open(output_file, "wb") as f: pickle.dump("bad_response", f) assert isinstance(response, openai.types.completion.Completion) # Save response to a file with open(output_file, "wb") as f: pickle.dump(response, f) TASK_TO_MAX_TOKENS = { "evals__mmlu__details": 1, "evals__mmlu__0_shot__cot__details": 1024, # Official meta uses 1024, but a small % (.05) of questions are answered correctly after relaxing "evals__mmlu_pro__details": 2048, "evals__gsm8k__details": 1024, } TASK_TO_EVAL_SET = { "mmlu": "evals__mmlu__details", "mmlu_cot": "evals__mmlu__0_shot__cot__details", "mmlu_pro": "evals__mmlu_pro__details", "gsm8k": "evals__gsm8k__details", } class CustomAsyncHTTPXClient(httpx.AsyncClient): async def send(self, request: httpx.Request, *args, **kwargs) -> httpx.Response: request.url = httpx.URL( f"https://model-{os.getenv('MODEL_ID')}.api.baseten.co/development/predict" ) return await super().send(request, *args, **kwargs) def get_client(provider): if provider not in "b10": if os.getenv("OPENAI_API_KEY") is None: os.environ["OPENAI_API_KEY"] = "EMPTY" return { "oai": AsyncOpenAI(base_url="http://127.0.0.1:8000/v1/"), "b10": AsyncOpenAI( api_key=f"Api-Key {os.getenv('OPENAI_API_KEY')}", base_url=f"https://model-{os.getenv('MODEL_ID')}.api.baseten.co/development/predict", http_client=CustomAsyncHTTPXClient(), ), "sgl": AsyncOpenAI(base_url="http://127.0.0.1:30000/v1/"), }[provider] # Define the benchmark function async def benchmark(args): ds = load_dataset( "meta-llama/Llama-3.1-405B-Instruct-evals", f"Llama-3.1-405B-Instruct-{TASK_TO_EVAL_SET[args.task]}", ) semaphore = asyncio.Semaphore(args.concurrency) # Limit to 16 concurrent tasks if args.num_examples is None: args.num_examples = len(ds["latest"]["input_final_prompts"]) prompts = ds["latest"]["input_final_prompts"][: args.num_examples] # Create the output directory if it does not exist os.makedirs(args.output_dir, exist_ok=True) tasks = [] # Create the tasks with tqdm progress bar max_tokens = TASK_TO_MAX_TOKENS[TASK_TO_EVAL_SET[args.task]] client = get_client(args.provider) for idx, prompt in enumerate(tqdm(prompts, desc="Creating tasks")): tasks.append( asyncio.create_task( fetch_responses( client, f"<|begin_of_text|>{prompt[0]}", semaphore, idx, args.provider, args.model_size, args.output_dir, max_tokens=max_tokens, ) ) ) # Run the tasks with tqdm progress bar for future in tqdm( asyncio.as_completed(tasks), total=len(tasks), desc="Processing tasks" ): await future def get_mmlu_answer(response): if response is not None: return response.choices[0].text.lstrip().rstrip().upper().replace(".", "") return None def get_mmlu_cot_answer(response): pattern = r"The best answer is (.+)\.?" match = re.search(pattern, response.choices[0].text) if match: return match.group(1).replace(".", "").replace("*", "") pattern = r"the best answer is (.+)\.?" match = re.search(pattern, response.choices[0].text) if match: return match.group(1).replace(".", "") pattern = r"The correct answer is (.+)\.?" match = re.search(pattern, response.choices[0].text) if match: return match.group(1).replace(".", "") pattern = r"the correct answer is (.+)\.?" match = re.search(pattern, response.choices[0].text) if match: return match.group(1).replace(".", "") def get_answer_gsm8k(response): pattern = r"The final answer is (.+)\.?" match = re.search(pattern, response.choices[0].text) if match: s = match.group(1) for ok_symbol in ["%", "$"]: s = s.replace(ok_symbol, "") return s TASK_TO_ANSWER_EXTRACTOR = { "evals__mmlu__details": get_mmlu_answer, "evals__mmlu__0_shot__cot__details": get_mmlu_cot_answer, "evals__gsm8k__details": get_answer_gsm8k, "evals__mmlu_pro__details": get_mmlu_cot_answer, } def get_dataset_from_task(task, response_path, model_size): ds_405b = load_dataset( f"meta-llama/Llama-3.1-405B-Instruct-evals", f"Llama-3.1-405B-Instruct-{task}", ) ds_405b_hash_order = [x[0] for x in ds_405b["latest"]["input_final_prompts_hash"]] if "70b" in model_size or "8b" in model_size: if "70" in model_size: ref_model_ds = load_dataset( f"meta-llama/Llama-3.1-70B-Instruct-evals", f"Llama-3.1-70B-Instruct-{task}", ) else: ref_model_ds = load_dataset( f"meta-llama/Llama-3.1-8B-Instruct-evals", f"Llama-3.1-8B-Instruct-{task}", ) hash_to_row = {} for row in ref_model_ds["latest"]: hash_to_row[row["input_final_prompts_hash"][0]] = row reordered_rows = [] for prompt_hash in ds_405b_hash_order: reordered_rows.append(hash_to_row[prompt_hash]) ref_model_ds["latest"] = reordered_rows return ref_model_ds return ds_405b def analyze(task, response_path, model_size): ds = get_dataset_from_task(task, response_path, model_size) responses = [] total = len(ds["latest"]) for i in range(0, total): response = pickle.load( open(os.path.join(response_path, f"response_{i}.pkl"), "rb") ) responses.append(response) @dataclass class Stats: correct: int = 0 total: int = 0 meta_correct: int = 0 average: float = None subtask_name_to_stats = defaultdict(lambda: Stats()) for response, ds_row in zip(responses, ds["latest"]): model_answer = TASK_TO_ANSWER_EXTRACTOR[task](response) subtask = ds_row["subtask_name"] is_eval_correct = model_answer in ds_row["input_correct_responses"] if is_eval_correct: subtask_name_to_stats[subtask].correct += 1 if ds_row["is_correct"]: subtask_name_to_stats[subtask].meta_correct += 1 subtask_name_to_stats[subtask].total += 1 micro_stats = Stats() for subtask, stats in subtask_name_to_stats.items(): stats.average = stats.correct / stats.total stats.meta_average = stats.meta_correct / stats.total micro_stats.correct += stats.correct micro_stats.total += stats.total micro_stats.meta_correct += stats.meta_correct micro_stats.average = micro_stats.correct / micro_stats.total micro_stats.meta_average = micro_stats.meta_correct / micro_stats.total print("Macro average", np.mean([x.average for x in subtask_name_to_stats.values()])) print( "Meta Macro average", np.mean([x.meta_average for x in subtask_name_to_stats.values()]), ) print("Micro average", micro_stats.average) print("Meta Micro average", micro_stats.meta_average) # Entry point for the script if __name__ == "__main__": parser = argparse.ArgumentParser( description="Script to run model with specified parameters." ) parser.add_argument( "--model-size", type=str, default="8b", help="Size of the model (e.g., 8b or 70b)", ) parser.add_argument( "--provider", type=str, default="sgl", help="Provider name (e.g., sgl, oai, b10)", ) parser.add_argument( "--task", type=str, required=True, help="Task (e.g., mmlu, mmlu_cot, mmlu_pro, gsm8k)", ) parser.add_argument( "--num-examples", type=int, default=None, help="Number of examples to process" ) parser.add_argument("--concurrency", type=int, default=16) parser.add_argument( "--output-dir", type=str, default="tmp-output-dir", help="Directory to save responses", ) args = parser.parse_args() asyncio.run(benchmark(args)) analyze(TASK_TO_EVAL_SET[args.task], args.output_dir, args.model_size) shutil.rmtree("tmp-output-dir", ignore_errors=True) ================================================ FILE: python/sglang/eval/loogle_eval.py ================================================ import argparse import asyncio import os import pickle from pathlib import Path from typing import List import openai import torch from bert_score import BERTScorer from datasets import load_dataset from tqdm import tqdm def get_client(api_url: str) -> openai.AsyncOpenAI: if os.getenv("OPENAI_API_KEY") is None: os.environ["OPENAI_API_KEY"] = "EMPTY" return openai.AsyncOpenAI(base_url=api_url) def get_dataset(): return load_dataset("bigai-nlco/LooGLE", "longdep_qa", split="test") async def fetch_response( client: openai.AsyncOpenAI, context: str, question: str, semaphore: asyncio.Semaphore, index: int, model: str, output_dir: Path, ): output_file = output_dir / f"response_{index}.pkl" if output_file.exists(): return prompt = ( "Please answer the question based on the long texts below.\n" f"{context}\n" f"Question: {question}\n" "Answer:" ) messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": prompt}, ] async with semaphore: try: response = await client.chat.completions.create( model=model, messages=messages, temperature=0.0, max_tokens=512, ) except openai.BadRequestError as e: with open(output_file, "wb") as f: pickle.dump({"error": str(e)}, f) return with open(output_file, "wb") as f: pickle.dump(response, f) async def benchmark(args): dataset = get_dataset() output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) client = get_client(args.api_url) semaphore = asyncio.Semaphore(args.max_concurrency) tasks: List[asyncio.Task] = [] for idx, ex in enumerate(dataset): if idx >= args.num_prompts: break tasks.append( asyncio.create_task( fetch_response( client, ex["context"], ex["question"], semaphore, idx, args.model, output_dir, ) ) ) for _ in tqdm( asyncio.as_completed(tasks), total=len(tasks), desc="Running benchmark" ): await _ def analyse(args): dataset = get_dataset() output_dir = Path(args.output_dir) device = "cuda" if torch.cuda.is_available() else "cpu" scorer = BERTScorer(lang="en", device=device) hyps: List[str] = [] refs: List[str] = [] for idx, ex in enumerate(tqdm(dataset, desc="Loading responses")): if idx >= args.num_prompts: break pkl_file = output_dir / f"response_{idx}.pkl" if not pkl_file.exists(): raise FileNotFoundError(pkl_file) response = pickle.load(open(pkl_file, "rb")) if isinstance(response, dict) and "error" in response: continue hyps.append(response.choices[0].message.content.strip()) refs.append(ex["answer"]) if not hyps: print("No valid responses to score!") return batch_size = 64 all_f1: List[float] = [] for i in tqdm(range(0, len(hyps), batch_size), desc="Scoring batches"): h_batch = hyps[i : i + batch_size] r_batch = refs[i : i + batch_size] _, _, f1_scores = scorer.score(h_batch, r_batch, verbose=False) all_f1.extend([float(x) for x in f1_scores]) avg = sum(all_f1) / len(all_f1) print(f"Average BERTScore (F1): {avg:.2%}") if __name__ == "__main__": parser = argparse.ArgumentParser( description="Run benchmark and evaluation in one go." ) parser.add_argument( "--api-url", default="http://127.0.0.1:30000/v1", help="OpenAI‑compatible API base URL", ) parser.add_argument( "--model", default="meta-llama/Llama-4-Maverick-17B-128E-Instruct", help="Model name or ID, only used for model name", ) parser.add_argument( "--max-concurrency", type=int, default=144, help="Maximum concurrent requests" ) parser.add_argument( "--output-dir", default="tmp-output-dir", help="Directory for cached responses" ) parser.add_argument( "--num-prompts", type=int, default=10000, help="Number of prompts to run" ) args = parser.parse_args() asyncio.run(benchmark(args)) analyse(args) ================================================ FILE: python/sglang/global_config.py ================================================ """Global configurations""" # FIXME: deprecate this file and move all usage to sglang.srt.environ or sglang.__init__.py class GlobalConfig: """ Store some global constants. """ def __init__(self): # Verbosity level # 0: do not output anything # 2: output final text after every run self.verbosity = 0 # Default backend of the language self.default_backend = None # Output tokenization configs self.skip_special_tokens_in_output = True self.spaces_between_special_tokens_in_out = True # Language frontend interpreter optimization configs self.enable_precache_with_tracing = True self.enable_parallel_encoding = True global_config = GlobalConfig() ================================================ FILE: python/sglang/jit_kernel/.clang-format ================================================ BasedOnStyle: Google IndentWidth: 2 ColumnLimit: 120 AllowShortFunctionsOnASingleLine: Empty DerivePointerAlignment: false PointerAlignment: Left NamespaceIndentation: None SortIncludes: true AllowShortLoopsOnASingleLine: false BinPackParameters: false # Prevents packing parameters in declarations BinPackArguments: false # Prevents packing arguments in function calls AlignAfterOpenBracket: AlwaysBreak # Forces a break after the opening parenthesis AlignOperands: Align # Aligns arguments vertically PenaltyBreakBeforeFirstCallParameter: 1 # Encourages breaking before the first argument PenaltyReturnTypeOnItsOwnLine: 100 # Keeps return type with function name IncludeCategories: - Regex: '^$' Priority: 0 - Regex: '^$' Priority: 2 - Regex: '^$' Priority: 1 - Regex: '^<.*/.*>$' Priority: 3 ================================================ FILE: python/sglang/jit_kernel/__init__.py ================================================ ================================================ FILE: python/sglang/jit_kernel/__main__.py ================================================ assert __name__ == "__main__" def generate_clangd(): import logging import os import subprocess from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path from sglang.jit_kernel.utils import DEFAULT_INCLUDE logger = logging.getLogger() logger.info("Generating .clangd file...") include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE status = subprocess.run( args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], capture_output=True, check=True, ) compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] major, minor = compute_cap.split(".") compile_flags = ",\n ".join( [ "-xcuda", f"--cuda-gpu-arch=sm_{major}{minor}", "-std=c++20", "-Wall", "-Wextra", ] + [f"-isystem{path}" for path in include_paths] ) clangd_content = f""" CompileFlags: Add: [ {compile_flags} ] """ if os.path.exists(".clangd"): logger.warning(".clangd file already exists, nothing done.") logger.warning(f"suggested content: {clangd_content}") else: with open(".clangd", "w") as f: f.write(clangd_content) logger.info(".clangd file generated.") generate_clangd() ================================================ FILE: python/sglang/jit_kernel/add_constant.py ================================================ from __future__ import annotations from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_add_constant_module(constant: int) -> Module: args = make_cpp_args(constant) return load_jit( "add_constant", *args, cuda_files=["add_constant.cuh"], cuda_wrappers=[("add_constant", f"add_constant<{args}>")], ) def add_constant(src: torch.Tensor, constant: int) -> torch.Tensor: dst = torch.empty_like(src) module = _jit_add_constant_module(constant) module.add_constant(dst, src) return dst ================================================ FILE: python/sglang/jit_kernel/awq_dequantize.py ================================================ from __future__ import annotations from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_awq_dequantize_module(dtype: torch.dtype) -> Module: args = make_cpp_args(dtype) return load_jit( "awq_dequantize", *args, cuda_files=["gemm/awq_dequantize.cuh"], cuda_wrappers=[("awq_dequantize", f"awq_dequantize<{args}>")], ) def awq_dequantize( qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor, ) -> torch.Tensor: qweight_rows = qweight.shape[0] qweight_cols = qweight.shape[1] output = torch.empty( (qweight_rows, qweight_cols * 8), dtype=scales.dtype, device=scales.device, ) module = _jit_awq_dequantize_module(scales.dtype) module.awq_dequantize(output, qweight, scales, qzeros) return output ================================================ FILE: python/sglang/jit_kernel/awq_marlin_repack.py ================================================ from __future__ import annotations from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import cache_once, load_jit if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_awq_marlin_repack_module() -> Module: return load_jit( "awq_marlin_repack", cuda_files=["gemm/marlin/awq_marlin_repack.cuh"], cuda_wrappers=[("awq_marlin_repack", "awq_marlin_repack")], ) def awq_marlin_repack( b_q_weight: torch.Tensor, size_k: int, size_n: int, num_bits: int, ) -> torch.Tensor: tile_size = 16 pack_factor = 32 // num_bits out = torch.empty( (size_k // tile_size, size_n * tile_size // pack_factor), dtype=b_q_weight.dtype, device=b_q_weight.device, ) module = _jit_awq_marlin_repack_module() module.awq_marlin_repack(out, b_q_weight, size_k, size_n, num_bits) return out def awq_marlin_moe_repack( b_q_weight: torch.Tensor, perm: torch.Tensor, size_k: int, size_n: int, num_bits: int, ) -> torch.Tensor: num_experts = b_q_weight.shape[0] assert size_k % 16 == 0 output = torch.empty( (num_experts, size_k // 16, size_n * (num_bits // 2)), device=b_q_weight.device, dtype=b_q_weight.dtype, ) for e in range(num_experts): output[e] = awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits) return output ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py ================================================ import itertools import torch import triton import triton.testing from sglang.jit_kernel.awq_dequantize import awq_dequantize as jit_awq_dequantize from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark try: from sgl_kernel import awq_dequantize as aot_awq_dequantize AOT_AVAILABLE = True except ImportError: AOT_AVAILABLE = False IS_CI = is_in_ci() if IS_CI: qweight_row_range = [128] qweight_cols_range = [16] else: qweight_row_range = [128, 256, 512, 1024, 3584] qweight_cols_range = [16, 32, 64, 128, 448] configs = list(itertools.product(qweight_row_range, qweight_cols_range)) def check_correctness(): if not AOT_AVAILABLE: print("sgl_kernel AOT not available, skipping correctness check") return qweight_row, qweight_col = 128, 16 device = torch.device("cuda") qweight = torch.randint( 0, torch.iinfo(torch.int32).max, (qweight_row, qweight_col), dtype=torch.int32, device=device, ) group_size = qweight_row scales_row = qweight_row // group_size scales_col = qweight_col * 8 scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device) qzeros = torch.randint( 0, torch.iinfo(torch.int32).max, (scales_row, qweight_col), dtype=torch.int32, device=device, ) jit_out = jit_awq_dequantize(qweight, scales, qzeros) aot_out = aot_awq_dequantize(qweight, scales, qzeros) torch.cuda.synchronize() torch.testing.assert_close(jit_out, aot_out, rtol=0, atol=0) print("Correctness check passed (JIT vs AOT)") if AOT_AVAILABLE: line_vals = ["jit", "aot"] line_names = ["JIT Kernel", "AOT Kernel"] styles = [("blue", "-"), ("green", "-")] else: line_vals = ["jit"] line_names = ["JIT Kernel"] styles = [("blue", "-")] @triton.testing.perf_report( triton.testing.Benchmark( x_names=["qweight_row", "qweight_col"], x_vals=configs, line_arg="provider", line_vals=line_vals, line_names=line_names, styles=styles, ylabel="us", plot_name="awq-dequantize-jit-vs-aot", args={}, ) ) def benchmark(qweight_row, qweight_col, provider): device = torch.device("cuda") qweight = torch.randint( 0, torch.iinfo(torch.int32).max, (qweight_row, qweight_col), dtype=torch.int32, device=device, ) group_size = qweight_row scales_row = qweight_row // group_size scales_col = qweight_col * 8 scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device) qzeros = torch.randint( 0, torch.iinfo(torch.int32).max, (scales_row, qweight_col), dtype=torch.int32, device=device, ) if provider == "jit": fn = lambda: jit_awq_dequantize(qweight, scales, qzeros) elif provider == "aot": fn = lambda: aot_awq_dequantize(qweight, scales, qzeros) else: raise ValueError(f"Unknown provider: {provider}") return run_benchmark(fn) if __name__ == "__main__": check_correctness() benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py ================================================ import numpy as np import torch import triton import triton.testing from sgl_kernel.scalar_type import scalar_types from sglang.jit_kernel.awq_marlin_repack import ( awq_marlin_moe_repack as jit_awq_marlin_moe_repack, ) from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "awq_marlin_moe_repack") and hasattr( torch.ops.sgl_kernel.awq_marlin_moe_repack, "default" ) IS_CI = is_in_ci() NUM_BITS = 4 GROUP_SIZE = 128 SIZE_N = 4096 def awq_pack(q_w, num_bits, size_k, size_n): if num_bits == 4: interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7]) elif num_bits == 8: interleave = np.array([0, 2, 1, 3]) else: raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel() q_w = q_w.reshape((-1, size_n)).contiguous() return pack_cols(q_w, num_bits, size_k, size_n) def make_moe_weights(num_experts, size_k, size_n, num_bits, group_size): pack_factor = 32 // num_bits b_q_weight = torch.empty( (num_experts, size_k, size_n // pack_factor), dtype=torch.int32, device="cuda", ) for e in range(num_experts): b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda") w_ref, q_w, s, zp = quantize_weights( b_weight, scalar_types.uint4, min(group_size, size_k), zero_points=True ) b_q_weight[e] = awq_pack(q_w, num_bits, size_k, size_n) perm = torch.empty((num_experts, 0), dtype=torch.int32, device="cuda") return b_q_weight, perm def check_correctness(): if not AOT_AVAILABLE: print("sgl_kernel AOT not available, skipping correctness check") return num_experts = 4 size_k = 1024 b_q_weight, perm = make_moe_weights( num_experts, size_k, SIZE_N, NUM_BITS, GROUP_SIZE ) out_jit = jit_awq_marlin_moe_repack(b_q_weight, perm, size_k, SIZE_N, NUM_BITS) out_aot = torch.ops.sgl_kernel.awq_marlin_moe_repack.default( b_q_weight, perm, size_k, SIZE_N, NUM_BITS ) torch.cuda.synchronize() torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0) print("Correctness check passed (JIT vs AOT)") if IS_CI: expert_range = [2, 4] else: expert_range = [2, 4, 8, 16] if AOT_AVAILABLE: line_vals = ["jit", "aot"] line_names = ["JIT Kernel", "AOT Kernel"] styles = [("blue", "-"), ("green", "-")] else: line_vals = ["jit"] line_names = ["JIT Kernel"] styles = [("blue", "-")] @triton.testing.perf_report( triton.testing.Benchmark( x_names=["num_experts"], x_vals=expert_range, line_arg="provider", line_vals=line_vals, line_names=line_names, styles=styles, ylabel="us", plot_name="awq-marlin-moe-repack-performance", args={"size_k": 4096, "size_n": SIZE_N, "num_bits": NUM_BITS}, ) ) def benchmark(num_experts, size_k, size_n, num_bits, provider): group_size = min(GROUP_SIZE, size_k) b_q_weight, perm = make_moe_weights( num_experts, size_k, size_n, num_bits, group_size ) if provider == "jit": fn = lambda: jit_awq_marlin_moe_repack( b_q_weight, perm, size_k, size_n, num_bits ) elif provider == "aot": fn = lambda: torch.ops.sgl_kernel.awq_marlin_moe_repack.default( b_q_weight, perm, size_k, size_n, num_bits ) else: raise ValueError(f"Unknown provider: {provider}") return run_benchmark(fn) if __name__ == "__main__": check_correctness() benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py ================================================ import numpy as np import torch import triton import triton.testing from sgl_kernel.scalar_type import scalar_types from sglang.jit_kernel.awq_marlin_repack import ( awq_marlin_repack as jit_awq_marlin_repack, ) from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "awq_marlin_repack") and hasattr( torch.ops.sgl_kernel.awq_marlin_repack, "default" ) IS_CI = is_in_ci() SIZE_K = 4096 SIZE_N = 4096 NUM_BITS = 4 GROUP_SIZE = 128 def awq_pack(q_w, num_bits, size_k, size_n): if num_bits == 4: interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7]) elif num_bits == 8: interleave = np.array([0, 2, 1, 3]) else: raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel() q_w = q_w.reshape((-1, size_n)).contiguous() return pack_cols(q_w, num_bits, size_k, size_n) _b_weight = torch.randn((SIZE_K, SIZE_N), dtype=torch.float16, device="cuda") _w_ref, _q_w, _s, _zp = quantize_weights( _b_weight, scalar_types.uint4, GROUP_SIZE, zero_points=True ) _q_w_awq = awq_pack(_q_w, NUM_BITS, SIZE_K, SIZE_N) def check_correctness(): if not AOT_AVAILABLE: print("sgl_kernel AOT not available, skipping correctness check") return out_jit = jit_awq_marlin_repack(_q_w_awq, SIZE_K, SIZE_N, NUM_BITS) out_aot = torch.ops.sgl_kernel.awq_marlin_repack.default( _q_w_awq, SIZE_K, SIZE_N, NUM_BITS ) torch.cuda.synchronize() torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0) print("Correctness check passed (JIT vs AOT)") if IS_CI: k_range = [1024, 4096] else: k_range = [512, 1024, 2048, 4096, 8192] if AOT_AVAILABLE: line_vals = ["jit", "aot"] line_names = ["JIT Kernel", "AOT Kernel"] styles = [("blue", "-"), ("green", "-")] else: line_vals = ["jit"] line_names = ["JIT Kernel"] styles = [("blue", "-")] @triton.testing.perf_report( triton.testing.Benchmark( x_names=["size_k"], x_vals=k_range, line_arg="provider", line_vals=line_vals, line_names=line_names, styles=styles, ylabel="us", plot_name="awq-marlin-repack-performance", args={"size_n": SIZE_N, "num_bits": NUM_BITS}, ) ) def benchmark(size_k, size_n, num_bits, provider): group_size = min(GROUP_SIZE, size_k) b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda") w_ref, q_w, s, zp = quantize_weights( b_weight, scalar_types.uint4, group_size, zero_points=True ) q_w_awq = awq_pack(q_w, num_bits, size_k, size_n) if provider == "jit": fn = lambda: jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits) elif provider == "aot": fn = lambda: torch.ops.sgl_kernel.awq_marlin_repack.default( q_w_awq, size_k, size_n, num_bits ) else: raise ValueError(f"Unknown provider: {provider}") return run_benchmark(fn) if __name__ == "__main__": check_correctness() benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_concat_mla.py ================================================ import itertools import torch import triton import triton.testing from sgl_kernel import concat_mla_absorb_q as aot_absorb_q from sgl_kernel import concat_mla_k as aot_k from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark from sglang.jit_kernel.concat_mla import concat_mla_absorb_q as jit_absorb_q from sglang.jit_kernel.concat_mla import concat_mla_k as jit_k IS_CI = is_in_ci() NUM_LOCAL_HEADS = 128 QK_NOPE_HEAD_DIM = 128 QK_ROPE_HEAD_DIM = 64 K_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM A_LAST_DIM = 512 B_LAST_DIM = 64 DTYPE = torch.bfloat16 DEVICE = "cuda" def aot_concat_mla_k(k, k_nope, k_rope): aot_k(k, k_nope, k_rope) def jit_concat_mla_k(k, k_nope, k_rope): jit_k(k, k_nope, k_rope) def torch_concat_mla_k(k, k_nope, k_rope): nope_head_dim = k_nope.shape[-1] k[:, :, :nope_head_dim] = k_nope k[:, :, nope_head_dim:] = k_rope.expand(-1, k.shape[1], -1) def aot_concat_mla_absorb_q(a, b): return aot_absorb_q(a, b) def jit_concat_mla_absorb_q(a, b): return jit_absorb_q(a, b) def torch_concat_mla_absorb_q(a, b, out): a_last_dim = a.shape[-1] out[:, :, :a_last_dim] = a out[:, :, a_last_dim:] = b if IS_CI: NUM_TOKENS_VALS = [256, 1024] else: NUM_TOKENS_VALS = [256, 512, 1024, 2048, 4096, 8192, 16384, 32768] K_LINE_VALS = ["aot", "jit", "torch"] K_LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"] K_STYLES = [("orange", "-"), ("blue", "--"), ("green", "-.")] def _create_concat_mla_k_data(num_tokens): """Allocate oversized containers and slice to produce non-contiguous tensors.""" k_nope_container = torch.randn( (num_tokens, NUM_LOCAL_HEADS, QK_NOPE_HEAD_DIM + 128), dtype=DTYPE, device=DEVICE, ) k_nope = k_nope_container[:, :, :QK_NOPE_HEAD_DIM] k_rope_container = torch.randn( (num_tokens, 1, 128 + QK_ROPE_HEAD_DIM), dtype=DTYPE, device=DEVICE, ) k_rope = k_rope_container[:, :, -QK_ROPE_HEAD_DIM:] k = torch.empty( (num_tokens, NUM_LOCAL_HEADS, K_HEAD_DIM), dtype=DTYPE, device=DEVICE, ) return k, k_nope, k_rope @triton.testing.perf_report( triton.testing.Benchmark( x_names=["num_tokens"], x_vals=NUM_TOKENS_VALS, line_arg="provider", line_vals=K_LINE_VALS, line_names=K_LINE_NAMES, styles=K_STYLES, ylabel="us", plot_name="concat-mla-k-performance", args={}, ) ) def bench_concat_mla_k(num_tokens: int, provider: str): k, k_nope, k_rope = _create_concat_mla_k_data(num_tokens) FN_MAP = { "aot": aot_concat_mla_k, "jit": jit_concat_mla_k, "torch": torch_concat_mla_k, } fn = lambda: FN_MAP[provider](k, k_nope, k_rope) return run_benchmark(fn) if IS_CI: ABSORB_Q_VALS = list(itertools.product([4, 16], [16])) else: ABSORB_Q_VALS = list(itertools.product([1, 4, 8, 16, 32], [1, 8, 32, 128])) Q_LINE_VALS = ["aot", "jit", "torch"] Q_LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"] Q_STYLES = [("orange", "-"), ("blue", "--"), ("green", "-.")] @triton.testing.perf_report( triton.testing.Benchmark( x_names=["dim_0", "dim_1"], x_vals=ABSORB_Q_VALS, line_arg="provider", line_vals=Q_LINE_VALS, line_names=Q_LINE_NAMES, styles=Q_STYLES, ylabel="us", plot_name="concat-mla-absorb-q-performance", args={}, ) ) def bench_concat_mla_absorb_q(dim_0: int, dim_1: int, provider: str): a = torch.randn(dim_0, dim_1, A_LAST_DIM, dtype=DTYPE, device=DEVICE) b = torch.randn(dim_0, dim_1, B_LAST_DIM, dtype=DTYPE, device=DEVICE) if provider == "torch": out = torch.empty( dim_0, dim_1, A_LAST_DIM + B_LAST_DIM, dtype=DTYPE, device=DEVICE ) fn = lambda: torch_concat_mla_absorb_q(a, b, out) else: FN_MAP = { "aot": aot_concat_mla_absorb_q, "jit": jit_concat_mla_absorb_q, } fn = lambda: FN_MAP[provider](a, b) return run_benchmark(fn) if __name__ == "__main__": bench_concat_mla_k.run(print_data=True) bench_concat_mla_absorb_q.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py ================================================ import itertools import torch import triton import triton.testing from flashinfer import fused_add_rmsnorm as fi_fused_add_rmsnorm from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm IS_CI = is_in_ci() def sglang_jit_fused_add_rmsnorm( input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float ) -> None: jit_fused_add_rmsnorm(input, residual, weight, eps) def flashinfer_fused_add_rmsnorm( input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float ) -> None: fi_fused_add_rmsnorm(input, residual, weight, eps=eps) DTYPE = torch.bfloat16 DEVICE = "cuda" if IS_CI: BS_LIST = [16] HIDDEN_SIZE_LIST = [512, 2048] else: BS_LIST = [2**n for n in range(0, 14)] HIDDEN_SIZE_LIST = [1536, 3072, 4096, 5120, 8192] LINE_VALS = ["jit", "flashinfer"] LINE_NAMES = ["SGL JIT Kernel", "FlashInfer"] STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")] configs = list(itertools.product(HIDDEN_SIZE_LIST, BS_LIST)) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["hidden_size", "batch_size"], x_vals=configs, line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="fused-add-rmsnorm-performance", args={}, ) ) def benchmark(hidden_size: int, batch_size: int, provider: str): input = torch.randn((batch_size, hidden_size), dtype=DTYPE, device=DEVICE) residual = torch.randn((batch_size, hidden_size), dtype=DTYPE, device=DEVICE) weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE) FN_MAP = { "jit": sglang_jit_fused_add_rmsnorm, "flashinfer": flashinfer_fused_add_rmsnorm, } fn = lambda: FN_MAP[provider]( input.clone(), residual.clone(), weight, torch.finfo(torch.bfloat16).eps ) return run_benchmark(fn) if __name__ == "__main__": benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py ================================================ # Benchmarks SGLang fused layernorm/rmsnorm scale shift kernels # 1. fused_norm_scale_shift # 2. fused_scale_residual_norm_scale_shift import itertools from typing import Tuple import torch import triton import triton.testing from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark_no_cudagraph from sglang.multimodal_gen.runtime.layers.layernorm import ( LayerNormScaleShift, RMSNormScaleShift, ScaleResidualLayerNormScaleShift, ScaleResidualRMSNormScaleShift, ) if is_in_ci(): B_RANGE, S_RANGE, D_RANGE = [1], [128], [1024] else: B_RANGE, S_RANGE, D_RANGE = [1], [128, 1024, 4096], [1024, 3072, 4096] NORM_TYPE_RANGE = ["layer", "rms"] AFFINE_RANGE = [True, False] DTYPE = torch.bfloat16 DEVICE = "cuda" EPS = 1e-5 LINE_VALS = ["native", "cuda"] LINE_NAMES = ["SGLang Native", "SGLang Fused"] STYLES = [("red", "-"), ("blue", "--")] config = list( itertools.product(B_RANGE, S_RANGE, D_RANGE, NORM_TYPE_RANGE, AFFINE_RANGE) ) def preprocess_layer(layer, affine: bool, D: int, DTYPE: torch.dtype): if affine: weight = torch.randn(D, dtype=DTYPE, device=DEVICE) bias = torch.randn(D, dtype=DTYPE, device=DEVICE) with torch.no_grad(): layer.norm.weight.copy_(weight) if hasattr(layer.norm, "bias"): layer.norm.bias.copy_(bias) layer.requires_grad_(False) return layer.to(DEVICE) # ============================================================================ # Benchmark 1: fused_norm_scale_shift # ============================================================================ @triton.testing.perf_report( triton.testing.Benchmark( x_names=["B", "S", "D", "norm_type", "affine"], x_vals=config, line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="fused_norm_scale_shift", args={}, ) ) def bench_fused_norm_scale_shift( B: int, S: int, D: int, norm_type, affine: bool, provider: str ) -> Tuple[float, float, float]: x = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE) scale = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE) shift = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE) if norm_type == "layer": layer = LayerNormScaleShift(D, EPS, affine, dtype=DTYPE) else: layer = RMSNormScaleShift(D, EPS, affine, dtype=DTYPE) layer = preprocess_layer(layer, affine, D, DTYPE) if provider == "native": fn = lambda: layer.forward_native(x, shift, scale) else: fn = lambda: layer.forward_cuda(x, shift, scale) return run_benchmark_no_cudagraph(fn) # ============================================================================ # Benchmark 2: fused_scale_residual_norm_scale_shift # ============================================================================ @triton.testing.perf_report( triton.testing.Benchmark( x_names=["B", "S", "D", "norm_type", "affine"], x_vals=config, line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="fused_scale_residual_norm_scale_shift", args={}, ) ) def bench_fused_scale_residual_norm_scale_shift( B: int, S: int, D: int, norm_type, affine: bool, provider: str ) -> Tuple[float, float, float]: residual = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE) x = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE) scale = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE) shift = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE) gate = torch.randn(B, 1, D, dtype=DTYPE, device=DEVICE) if norm_type == "layer": layer = ScaleResidualLayerNormScaleShift(D, EPS, affine, dtype=DTYPE).to(DEVICE) else: layer = ScaleResidualRMSNormScaleShift(D, EPS, affine, dtype=DTYPE).to(DEVICE) layer = preprocess_layer(layer, affine, D, DTYPE) if provider == "native": fn = lambda: layer.forward_native(residual, x, gate, shift, scale) else: fn = lambda: layer.forward_cuda(residual, x, gate, shift, scale) return run_benchmark_no_cudagraph(fn) if __name__ == "__main__": print(f"\n{'='*80}") print("Benchmark: fused_norm_scale_shift") print(f"{'='*80}\n") bench_fused_norm_scale_shift.run(print_data=True) print(f"\n{'='*80}") print("Benchmark: fused_scale_residual_norm_scale_shift") print(f"{'='*80}\n") bench_fused_scale_residual_norm_scale_shift.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py ================================================ import torch import triton import triton.testing from sgl_kernel.scalar_type import scalar_types from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm as jit_gptq_marlin_gemm from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace from sglang.test.test_marlin_utils import marlin_quantize AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "gptq_marlin_gemm") and hasattr( torch.ops.sgl_kernel.gptq_marlin_gemm, "default" ) IS_CI = is_in_ci() SIZE_K = 4096 SIZE_N = 4096 GROUP_SIZE = 128 QUANT_TYPE = scalar_types.uint4b8 _b_weight = torch.randn((SIZE_K, SIZE_N), dtype=torch.float16, device="cuda") _w_ref, _marlin_q_w, _marlin_s, _g_idx, _sort_indices, _ = marlin_quantize( _b_weight, QUANT_TYPE, GROUP_SIZE, act_order=False ) _workspace = marlin_make_workspace(_w_ref.device) def _run_gemm(fn, a): return fn( a, None, _marlin_q_w, _marlin_s, None, None, _g_idx, _sort_indices, _workspace, QUANT_TYPE, a.shape[0], SIZE_N, SIZE_K, is_k_full=True, use_atomic_add=False, use_fp32_reduce=False, is_zp_float=False, ) def _run_gemm_aot(a): return torch.ops.sgl_kernel.gptq_marlin_gemm.default( a, None, _marlin_q_w, _marlin_s, None, None, _g_idx, _sort_indices, _workspace, QUANT_TYPE.id, a.shape[0], SIZE_N, SIZE_K, True, False, False, False, ) def check_correctness(): if not AOT_AVAILABLE: print("sgl_kernel AOT not available, skipping correctness check") return a = torch.randn((16, SIZE_K), dtype=torch.float16, device="cuda") out_jit = _run_gemm(jit_gptq_marlin_gemm, a) out_aot = _run_gemm_aot(a) torch.testing.assert_close(out_jit, out_aot, rtol=1e-3, atol=1e-3) print("Correctness check passed (JIT vs AOT)") if IS_CI: m_range = [1, 16, 128] else: m_range = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] if AOT_AVAILABLE: line_vals = ["jit", "aot"] line_names = ["JIT Kernel", "AOT Kernel"] styles = [("blue", "-"), ("green", "-")] else: line_vals = ["jit"] line_names = ["JIT Kernel"] styles = [("blue", "-")] @triton.testing.perf_report( triton.testing.Benchmark( x_names=["size_m"], x_vals=m_range, line_arg="provider", line_vals=line_vals, line_names=line_names, styles=styles, ylabel="us", plot_name="gptq-marlin-gemm-performance", args={}, ) ) def benchmark(size_m, provider): device = torch.device("cuda") a = torch.randn((size_m, SIZE_K), dtype=torch.float16, device=device) if provider == "jit": fn = lambda: _run_gemm(jit_gptq_marlin_gemm, a) elif provider == "aot": fn = lambda: _run_gemm_aot(a) else: raise ValueError(f"Unknown provider: {provider}") return run_benchmark(fn) if __name__ == "__main__": check_correctness() benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py ================================================ import torch import triton import triton.testing from sgl_kernel.scalar_type import scalar_types from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack as jit_fn from sglang.srt.layers.quantization.utils import gptq_quantize_weights, pack_rows AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "gptq_marlin_repack") and hasattr( torch.ops.sgl_kernel.gptq_marlin_repack, "default" ) IS_CI = is_in_ci() SIZE_N = 4096 NUM_BITS = 4 QUANT_TYPE = scalar_types.uint4b8 GROUP_SIZE = 128 _cache = {} def _get_inputs(size_k): if size_k not in _cache: size_n = SIZE_N b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda") _, q_w, _, _, _ = gptq_quantize_weights( b_weight, QUANT_TYPE, GROUP_SIZE, act_order=False ) q_w_gptq = pack_rows(q_w, NUM_BITS, size_k, size_n) sort_indices = torch.empty(0, dtype=torch.int, device="cuda") _cache[size_k] = (q_w_gptq, sort_indices) return _cache[size_k] def check_correctness(): if not AOT_AVAILABLE: print("sgl_kernel AOT not available, skipping correctness check") return size_k = 4096 q_w_gptq, sort_indices = _get_inputs(size_k) out_jit = jit_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS) out_aot = torch.ops.sgl_kernel.gptq_marlin_repack.default( q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS ) torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0) print("Correctness check passed (JIT vs AOT)") if IS_CI: k_range = [128, 1024, 4096] else: k_range = [128, 256, 512, 1024, 2048, 4096, 8192] if AOT_AVAILABLE: line_vals = ["jit", "aot"] line_names = ["JIT Kernel", "AOT Kernel"] styles = [("blue", "-"), ("green", "-")] else: line_vals = ["jit"] line_names = ["JIT Kernel"] styles = [("blue", "-")] @triton.testing.perf_report( triton.testing.Benchmark( x_names=["size_k"], x_vals=k_range, line_arg="provider", line_vals=line_vals, line_names=line_names, styles=styles, ylabel="us", plot_name="gptq-marlin-repack-performance", args={}, ) ) def benchmark(size_k, provider): q_w_gptq, sort_indices = _get_inputs(size_k) if provider == "jit": fn = lambda: jit_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS) elif provider == "aot": fn = lambda: torch.ops.sgl_kernel.gptq_marlin_repack.default( q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS ) else: raise ValueError(f"Unknown provider: {provider}") return run_benchmark(fn) if __name__ == "__main__": check_correctness() benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_hadamard.py ================================================ import itertools import math from typing import Tuple import torch import torch.nn.functional as F import triton import triton.testing from sglang.jit_kernel.benchmark.utils import ( DEFAULT_DEVICE, DEFAULT_DTYPE, get_benchmark_range, run_benchmark, ) from sglang.jit_kernel.hadamard import hadamard_transform # AOT kernel: might not be available in all environments. # This is used for performance baseline comparison. try: from sgl_kernel import hadamard_transform as hadamard_transform_aot AOT_AVAILABLE = True except Exception: AOT_AVAILABLE = False # Naive reference implementation using scipy hadamard matrix. try: from scipy.linalg import hadamard SCIPY_AVAILABLE = True except ImportError: SCIPY_AVAILABLE = False # CI environment uses simplified parameters batch_sizes = get_benchmark_range( full_range=[1, 16, 64, 256], ci_range=[16], ) dim_range = get_benchmark_range( full_range=[64, 256, 1024, 4096, 8192, 16384, 32768], ci_range=[1024], ) # Naive reference implementation using precomputed scipy hadamard matrix. def torch_hadamard_transform(x, scale, H, dim, dim_padded): flat = x.reshape(-1, dim) if dim != dim_padded: flat = F.pad(flat, (0, dim_padded - dim)) out = F.linear(flat, H) * scale return out[..., :dim].reshape(x.shape) available_providers = ["jit_kernel"] available_names = ["JIT Kernel"] available_styles = [("red", "-")] if AOT_AVAILABLE: available_providers.insert(0, "aot_kernel") available_names.insert(0, "AOT Kernel") available_styles.insert(0, ("green", "-")) if SCIPY_AVAILABLE: available_providers.append("naive") available_names.append("Naive (scipy)") available_styles.append(("blue", "-")) configs = list(itertools.product(batch_sizes, dim_range)) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["batch_size", "dim"], x_vals=[list(c) for c in configs], line_arg="provider", line_vals=available_providers, line_names=available_names, styles=available_styles, ylabel="us", plot_name="hadamard-transform-performance", args={}, ) ) def benchmark(batch_size: int, dim: int, provider: str) -> Tuple[float, float, float]: scale = 1.0 / math.sqrt(dim) x = torch.randn(batch_size, dim, device=DEFAULT_DEVICE, dtype=DEFAULT_DTYPE) FN_MAP = { "jit_kernel": lambda: hadamard_transform(x.clone(), scale=scale), } if AOT_AVAILABLE: FN_MAP["aot_kernel"] = lambda: hadamard_transform_aot(x.clone(), scale=scale) if SCIPY_AVAILABLE: # Precompute Hadamard matrix on GPU to avoid CPU-GPU transfer # during CUDA graph capture. log_dim = math.ceil(math.log2(dim)) if dim > 0 else 0 dim_padded = 2**log_dim if dim > 0 else 1 H = torch.tensor( hadamard(dim_padded, dtype=float), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE, ) FN_MAP["naive"] = lambda: torch_hadamard_transform( x.clone(), scale, H, dim, dim_padded ) fn = FN_MAP[provider] return run_benchmark(fn) if __name__ == "__main__": print("=" * 80) print("Benchmarking Fast Hadamard Transform") print("=" * 80) benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_hicache.py ================================================ """Benchmark for HiCache JIT kernel performance. This benchmark tests the performance of KV cache transfer operations between GPU and CPU (host pinned memory), comparing: - SGL AOT Kernel: Pre-compiled transfer_kv kernels from sgl_kernel - SGL JIT Kernel: JIT-compiled hicache kernels - PyTorch Indexing: Plain PyTorch index copy - PyTorch 2 Stream: PyTorch implementation using 2 CUDA streams Tests cover: - One Layer: CPU->GPU - All Layer: GPU->CPU Note: Uses do_bench instead of do_bench_cudagraph since CUDA graph capture doesn't support CPU-GPU memory transfers. """ import itertools import os from dataclasses import dataclass from typing import Tuple import torch import triton import triton.testing from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer from sglang.jit_kernel.benchmark.utils import DEFAULT_QUANTILES, get_benchmark_range from sglang.jit_kernel.hicache import ( can_use_hicache_jit_kernel, transfer_hicache_all_layer, transfer_hicache_one_layer, ) DISABLE_TORCH = os.environ.get("DISABLE_TORCH", "0") == "1" PAGE_SIZE = 1 ENABLE_SORT = True GPU_CACHE_SIZE = 256 * 1024 # 256K tokens on GPU HOST_CACHE_SIZE = 512 * 1024 # 512K tokens on CPU NUM_LAYERS = 8 @dataclass(frozen=True) class HiCacheCache: k_cache_cuda: torch.Tensor v_cache_cuda: torch.Tensor k_cache_host: torch.Tensor v_cache_host: torch.Tensor def get_slice(self, num_layers: int, element_size: int) -> "HiCacheCache": def slice_cuda(t: torch.Tensor) -> torch.Tensor: needed_cuda = num_layers * GPU_CACHE_SIZE return t.view(-1, element_size)[:needed_cuda].unflatten(0, (num_layers, -1)) def slice_host(t: torch.Tensor) -> torch.Tensor: needed_host = num_layers * HOST_CACHE_SIZE return t.view(-1, element_size)[:needed_host].unflatten(0, (num_layers, -1)) return HiCacheCache( k_cache_cuda=slice_cuda(self.k_cache_cuda), v_cache_cuda=slice_cuda(self.v_cache_cuda), k_cache_host=slice_host(self.k_cache_host), v_cache_host=slice_host(self.v_cache_host), ) def gen_indices( size: int, max_size: int, *, page_size: int = PAGE_SIZE ) -> torch.Tensor: def align(x: int) -> int: return (x + page_size - 1) // page_size assert size <= max_size and max_size % page_size == 0 indices = torch.randperm(align(max_size))[: align(size)] offsets = torch.arange(page_size) return (indices[:, None] * page_size + offsets).flatten().cuda()[:size] def sglang_aot_transfer_one( k_cache_dst: torch.Tensor, v_cache_dst: torch.Tensor, indices_dst: torch.Tensor, k_cache_src: torch.Tensor, v_cache_src: torch.Tensor, indices_src: torch.Tensor, item_size: int, ) -> None: """SGL AOT Kernel for single layer transfer.""" transfer_kv_per_layer( k_cache_src, k_cache_dst, v_cache_src, v_cache_dst, indices_src, indices_dst, item_size, ) def sglang_jit_transfer_one( k_cache_dst: torch.Tensor, v_cache_dst: torch.Tensor, indices_dst: torch.Tensor, k_cache_src: torch.Tensor, v_cache_src: torch.Tensor, indices_src: torch.Tensor, element_dim: int, ) -> None: """SGL JIT Kernel for single layer transfer.""" transfer_hicache_one_layer( k_cache_dst, v_cache_dst, indices_dst, k_cache_src, v_cache_src, indices_src, element_dim=element_dim, ) def sglang_aot_transfer_all( k_ptrs_dst: torch.Tensor, v_ptrs_dst: torch.Tensor, indices_dst: torch.Tensor, k_ptrs_src: torch.Tensor, v_ptrs_src: torch.Tensor, indices_src: torch.Tensor, item_size: int, num_layers: int, ) -> None: """SGL AOT Kernel for all layer transfer.""" transfer_kv_all_layer( k_ptrs_src, k_ptrs_dst, v_ptrs_src, v_ptrs_dst, indices_src, indices_dst, item_size, num_layers, ) def sglang_jit_transfer_all( k_ptrs_dst: torch.Tensor, v_ptrs_dst: torch.Tensor, indices_dst: torch.Tensor, k_ptrs_src: torch.Tensor, v_ptrs_src: torch.Tensor, indices_src: torch.Tensor, stride_bytes: int, element_size: int, ) -> None: """SGL JIT Kernel for all layer transfer.""" transfer_hicache_all_layer( k_ptrs_dst, v_ptrs_dst, indices_dst, k_ptrs_src, v_ptrs_src, indices_src, kv_cache_src_stride_bytes=stride_bytes, kv_cache_dst_stride_bytes=stride_bytes, element_size=element_size, ) def pytorch_transfer( k_cache_dst: torch.Tensor, v_cache_dst: torch.Tensor, indices_dst_on_dst: torch.Tensor, k_cache_src: torch.Tensor, v_cache_src: torch.Tensor, indices_src_on_src: torch.Tensor, ) -> None: """PyTorch indexing baseline.""" dst_device = k_cache_dst.device k_cache_dst[indices_dst_on_dst] = k_cache_src[indices_src_on_src].to(dst_device) v_cache_dst[indices_dst_on_dst] = v_cache_src[indices_src_on_src].to(dst_device) # Benchmark configuration BS_RANGE = get_benchmark_range( full_range=[2**n for n in range(0, 16)], ci_range=[16], ) ELEMENT_SIZE_RANGE = get_benchmark_range( full_range=[64, 128, 256, 512, 1024], ci_range=[1024], ) LINE_VALS = ["aot", "jit", "torch"] LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"] STYLES = [("orange", "-"), ("blue", "--"), ("red", ":")] CONFIGS = list(itertools.product(ELEMENT_SIZE_RANGE, BS_RANGE)) # ============================================================================= # One Layer Benchmarks # ============================================================================= @triton.testing.perf_report( triton.testing.Benchmark( x_names=["element_size", "batch_size"], x_vals=CONFIGS, line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="hicache-one-layer-h2d", args={}, ) ) def benchmark_one_layer_h2d( element_size: int, batch_size: int, provider: str ) -> Tuple[float, float, float]: """One Layer: Host (CPU) -> Device (GPU).""" global cache cache_local = cache.get_slice(num_layers=NUM_LAYERS, element_size=element_size) k_cache_src = cache_local.k_cache_host v_cache_src = cache_local.v_cache_host k_cache_dst = cache_local.k_cache_cuda v_cache_dst = cache_local.v_cache_cuda torch.manual_seed(batch_size * 65536 + element_size) indices_src_gpu = gen_indices(batch_size, HOST_CACHE_SIZE) indices_dst_gpu = gen_indices(batch_size, GPU_CACHE_SIZE) if ENABLE_SORT: indices_src_gpu, mapping = indices_src_gpu.sort() indices_dst_gpu = indices_dst_gpu[mapping] indices_src_cpu = indices_src_gpu.cpu() torch.cuda.synchronize() element_bytes = element_size * k_cache_src.element_size() FN_MAP = { "aot": lambda: [ sglang_aot_transfer_one( k_cache_dst[i], v_cache_dst[i], indices_dst_gpu, k_cache_src[i], v_cache_src[i], indices_src_gpu, element_bytes, ) for i in range(NUM_LAYERS) ], "jit": lambda: [ sglang_jit_transfer_one( k_cache_dst[i], v_cache_dst[i], indices_dst_gpu, k_cache_src[i], v_cache_src[i], indices_src_gpu, element_size, ) for i in range(NUM_LAYERS) ], "torch": lambda: [ pytorch_transfer( k_cache_dst[i], v_cache_dst[i], indices_dst_gpu, k_cache_src[i], v_cache_src[i], indices_src_cpu, ) for i in range(NUM_LAYERS) ], } if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes): return (float("nan"), float("nan"), float("nan")) if DISABLE_TORCH and provider in ["torch"]: return (float("nan"), float("nan"), float("nan")) ms, min_ms, max_ms = triton.testing.do_bench( # type: ignore FN_MAP[provider], quantiles=DEFAULT_QUANTILES, warmup=5, rep=25 ) return ( 1000 * ms / NUM_LAYERS, 1000 * max_ms / NUM_LAYERS, 1000 * min_ms / NUM_LAYERS, ) # ============================================================================= # All Layer Benchmarks # ============================================================================= def _create_ptr_tensor(tensors, device="cuda"): """Create a tensor of data pointers.""" return torch.tensor( [t.data_ptr() for t in tensors], dtype=torch.uint64, device=device, ) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["element_size", "batch_size"], x_vals=CONFIGS, line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="hicache-all-layer-d2h", args={}, ) ) def benchmark_all_layer_d2h( element_size: int, batch_size: int, provider: str ) -> Tuple[float, float, float]: """All Layer: Device (GPU) -> Host (CPU).""" global cache cache_local = cache.get_slice(num_layers=NUM_LAYERS, element_size=element_size) k_caches_src = cache_local.k_cache_cuda v_caches_src = cache_local.v_cache_cuda k_caches_dst = cache_local.k_cache_host v_caches_dst = cache_local.v_cache_host torch.manual_seed(batch_size * 65536 + element_size) indices_src_gpu = gen_indices(batch_size, GPU_CACHE_SIZE) indices_dst_gpu = gen_indices(batch_size, HOST_CACHE_SIZE) if ENABLE_SORT: indices_dst_gpu, mapping = indices_dst_gpu.sort() indices_src_gpu = indices_src_gpu[mapping] indices_dst_cpu = indices_dst_gpu.cpu() torch.cuda.synchronize() element_bytes = element_size * k_caches_src.element_size() k_ptrs_src = _create_ptr_tensor([k_caches_src[i] for i in range(NUM_LAYERS)]) v_ptrs_src = _create_ptr_tensor([v_caches_src[i] for i in range(NUM_LAYERS)]) k_ptrs_dst = _create_ptr_tensor([k_caches_dst[i] for i in range(NUM_LAYERS)]) v_ptrs_dst = _create_ptr_tensor([v_caches_dst[i] for i in range(NUM_LAYERS)]) FN_MAP = { "aot": lambda: sglang_aot_transfer_all( k_ptrs_dst, v_ptrs_dst, indices_dst_gpu, k_ptrs_src, v_ptrs_src, indices_src_gpu, element_bytes, NUM_LAYERS, ), "jit": lambda: sglang_jit_transfer_all( k_ptrs_dst, v_ptrs_dst, indices_dst_gpu, k_ptrs_src, v_ptrs_src, indices_src_gpu, element_bytes, element_bytes, ), "torch": lambda: [ pytorch_transfer( k_caches_dst[i], v_caches_dst[i], indices_dst_cpu, k_caches_src[i], v_caches_src[i], indices_src_gpu, ) for i in range(NUM_LAYERS) ], } if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes): return (float("nan"), float("nan"), float("nan")) if DISABLE_TORCH and provider in ["torch"]: return (float("nan"), float("nan"), float("nan")) ms, min_ms, max_ms = triton.testing.do_bench( # type: ignore FN_MAP[provider], quantiles=DEFAULT_QUANTILES, warmup=5, rep=25 ) return ( 1000 * ms / NUM_LAYERS, 1000 * max_ms / NUM_LAYERS, 1000 * min_ms / NUM_LAYERS, ) if __name__ == "__main__": MAX_SIZE = max(ELEMENT_SIZE_RANGE) DEVICE_SHAPE = (NUM_LAYERS * GPU_CACHE_SIZE, MAX_SIZE) HOST_SHAPE = (NUM_LAYERS * HOST_CACHE_SIZE, MAX_SIZE) cache = HiCacheCache( k_cache_cuda=torch.empty(DEVICE_SHAPE, dtype=torch.bfloat16, device="cuda"), v_cache_cuda=torch.empty(DEVICE_SHAPE, dtype=torch.bfloat16, device="cuda"), k_cache_host=torch.empty(HOST_SHAPE, dtype=torch.bfloat16, pin_memory=True), v_cache_host=torch.empty(HOST_SHAPE, dtype=torch.bfloat16, pin_memory=True), ) print("=" * 60) print("One Layer: Host -> Device (CPU -> GPU)") print("=" * 60) benchmark_one_layer_h2d.run(print_data=True) print("\n" + "=" * 60) print("All Layer: Device -> Host (GPU -> CPU) [per-layer avg]") print("=" * 60) benchmark_all_layer_d2h.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py ================================================ import torch import triton import triton.testing from sgl_kernel.scalar_type import scalar_types from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm as jit_fn from sglang.srt.layers.moe.fused_moe_triton import moe_align_block_size from sglang.test.test_marlin_utils import marlin_quantize AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "moe_wna16_marlin_gemm") and hasattr( torch.ops.sgl_kernel.moe_wna16_marlin_gemm, "default" ) IS_CI = is_in_ci() def stack_and_dev(tensors): dev = tensors[0].device return torch.stack(tensors, dim=0).to(dev) E = 8 SIZE_K = 4096 SIZE_N = 4096 GROUP_SIZE = 128 TOPK = 2 QUANT_TYPE = scalar_types.uint4b8 DTYPE = torch.float16 BLOCK_SIZE_M = 64 torch.manual_seed(0) _qweight_l, _scales_l, _w_ref_l = [], [], [] for i in range(E): _w = torch.randn((SIZE_N, SIZE_K), dtype=DTYPE, device="cuda") / 20 _perm = torch.randperm(SIZE_K) _w_ref, _qw, _s, _, _, _ = marlin_quantize(_w, QUANT_TYPE, GROUP_SIZE, False, _perm) _w_ref_l.append(_w_ref.T) _qweight_l.append(_qw) _scales_l.append(_s) _qweight = stack_and_dev(_qweight_l).contiguous() _scales = stack_and_dev(_scales_l) _sms = torch.cuda.get_device_properties("cuda").multi_processor_count def _make_inputs(size_m): a = torch.randn((size_m, SIZE_K), dtype=DTYPE, device="cuda") / 10 score = torch.randn((size_m, E), dtype=DTYPE, device="cuda") score_softmax = torch.softmax(score, dim=-1, dtype=torch.float32) topk_weights, topk_ids = torch.topk(score_softmax, TOPK) sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( topk_ids, BLOCK_SIZE_M, E ) max_workspace_size = (SIZE_N // 64) * (sorted_token_ids.size(0) // BLOCK_SIZE_M) max_workspace_size = min(max_workspace_size, _sms * 4) workspace = torch.zeros(max_workspace_size, dtype=torch.int, device="cuda") c = torch.empty((size_m * TOPK, SIZE_N), dtype=DTYPE, device="cuda") return ( a, c, topk_weights, topk_ids, sorted_token_ids, expert_ids, num_tokens_post_padded, workspace, ) def _run_jit( a, c, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded, workspace, size_m, ): return jit_fn( a, c, _qweight, None, _scales, None, None, None, None, workspace, sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights, moe_block_size=BLOCK_SIZE_M, top_k=TOPK, mul_topk_weights=False, is_ep=False, b_q_type=QUANT_TYPE, size_m=size_m, size_n=SIZE_N, size_k=SIZE_K, is_k_full=True, use_atomic_add=True, use_fp32_reduce=True, is_zp_float=False, ) def _run_aot( a, c, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded, workspace, size_m, ): return torch.ops.sgl_kernel.moe_wna16_marlin_gemm.default( a, c, _qweight, None, _scales, None, None, None, None, workspace, sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights, moe_block_size=BLOCK_SIZE_M, top_k=TOPK, mul_topk_weights=False, is_ep=False, b_q_type_id=QUANT_TYPE.id, size_m=size_m, size_n=SIZE_N, size_k=SIZE_K, is_k_full=True, use_atomic_add=True, use_fp32_reduce=True, is_zp_float=False, ) def check_correctness(): if not AOT_AVAILABLE: print("sgl_kernel AOT not available, skipping correctness check") return size_m = 16 a, c, topk_weights, topk_ids, sorted_token_ids, expert_ids, ntp, workspace = ( _make_inputs(size_m) ) c_jit = c.clone() c_aot = c.clone() _run_jit( a, c_jit, topk_weights, sorted_token_ids, expert_ids, ntp, workspace, size_m ) _run_aot( a, c_aot, topk_weights, sorted_token_ids, expert_ids, ntp, workspace, size_m ) torch.testing.assert_close(c_jit, c_aot, rtol=1e-3, atol=1e-3) print("Correctness check passed (JIT vs AOT)") if IS_CI: m_range = [1, 16, 128] else: m_range = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] if AOT_AVAILABLE: line_vals = ["jit", "aot"] line_names = ["JIT Kernel", "AOT Kernel"] styles = [("blue", "-"), ("green", "-")] else: line_vals = ["jit"] line_names = ["JIT Kernel"] styles = [("blue", "-")] @triton.testing.perf_report( triton.testing.Benchmark( x_names=["size_m"], x_vals=m_range, line_arg="provider", line_vals=line_vals, line_names=line_names, styles=styles, ylabel="us", plot_name="moe-wna16-marlin-gemm-performance", args={}, ) ) def benchmark(size_m, provider): a, c, topk_weights, topk_ids, sorted_token_ids, expert_ids, ntp, workspace = ( _make_inputs(size_m) ) if provider == "jit": fn = lambda: _run_jit( a, c.clone(), topk_weights, sorted_token_ids, expert_ids, ntp, workspace, size_m, ) elif provider == "aot": fn = lambda: _run_aot( a, c.clone(), topk_weights, sorted_token_ids, expert_ids, ntp, workspace, size_m, ) else: raise ValueError(f"Unknown provider: {provider}") return run_benchmark(fn) if __name__ == "__main__": check_correctness() benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_norm.py ================================================ import itertools import torch import triton import triton.testing from flashinfer.norm import fused_add_rmsnorm as fi_fused_add_rmsnorm from flashinfer.norm import rmsnorm as fi_rmsnorm from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm IS_CI = is_in_ci() DTYPE = torch.bfloat16 DEVICE = "cuda" # JIT rmsnorm: hidden_size in {64,128,256} or (multiple of 256, <=8192) # JIT fused_add_rmsnorm: hidden_size % 8 == 0, <=8192 # Use multiples of 256 <=8192 to satisfy both kernels if IS_CI: BS_LIST = [16] HIDDEN_SIZE_LIST = [512, 2048] else: BS_LIST = [2**n for n in range(0, 14)] HIDDEN_SIZE_LIST = [1536, 3072, 4096, 5120, 8192] LINE_VALS = ["jit", "flashinfer"] LINE_NAMES = ["SGL JIT Kernel", "FlashInfer"] STYLES = [("blue", "--"), ("green", "-.")] configs = list(itertools.product(HIDDEN_SIZE_LIST, BS_LIST)) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["hidden_size", "batch_size"], x_vals=configs, line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="rmsnorm-performance", args={}, ) ) def benchmark_rmsnorm(hidden_size: int, batch_size: int, provider: str): input = torch.randn((batch_size, hidden_size), dtype=DTYPE, device=DEVICE) weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE) FN_MAP = { "jit": lambda: jit_rmsnorm(input.clone(), weight), "flashinfer": lambda: fi_rmsnorm(input.clone(), weight, out=input.clone()), } fn = FN_MAP[provider] return run_benchmark(fn) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["hidden_size", "batch_size"], x_vals=configs, line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="fused-add-rmsnorm-performance", args={}, ) ) def benchmark_fused_add_rmsnorm(hidden_size: int, batch_size: int, provider: str): input = torch.randn((batch_size, hidden_size), dtype=DTYPE, device=DEVICE) residual = torch.randn((batch_size, hidden_size), dtype=DTYPE, device=DEVICE) weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE) FN_MAP = { "jit": lambda: jit_fused_add_rmsnorm( input.clone(), residual.clone(), weight, torch.finfo(DTYPE).eps ), "flashinfer": lambda: fi_fused_add_rmsnorm( input.clone(), residual.clone(), weight, eps=torch.finfo(DTYPE).eps ), } fn = FN_MAP[provider] return run_benchmark(fn) if __name__ == "__main__": print("Benchmarking rmsnorm...") benchmark_rmsnorm.run(print_data=True) print("Benchmarking fused_add_rmsnorm...") benchmark_fused_add_rmsnorm.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_norm_impls.py ================================================ from __future__ import annotations import argparse import csv import functools import importlib import math import os import statistics import subprocess import sys from pathlib import Path from typing import Callable import torch import torch.nn.functional as F from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, is_in_ci from sglang.jit_kernel.diffusion.triton.norm import norm_infer, rms_norm_fn from sglang.jit_kernel.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm from sglang.jit_kernel.utils import KERNEL_PATH os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1") REPO_ROOT = KERNEL_PATH.parents[2] THIRD_PARTY_ROOT = REPO_ROOT / "third_party" FLAGGEMS_REPO = "https://github.com/flagos-ai/FlagGems.git" QUACK_REPO = "https://github.com/Dao-AILab/quack.git" TORCH_LN = "torch.nn.LayerNorm" SGL_RMS = "sglang.RMSNorm.forward_cuda" SGL_FUSED = "sgl_kernel.fused_add_rmsnorm" SGL_LN = "sglang.LayerNormScaleShift" SGL_RES_LN = "sglang.ScaleResidualLayerNormScaleShift" SGL_LN_PAIR = f"{SGL_LN} / {SGL_RES_LN}" MOVA_LN_MIX = f"{TORCH_LN} / {SGL_LN_PAIR}" ACTUAL_DIFFUSION_GROUPS: list[ tuple[str, str, list[tuple[str, str, tuple[int, ...], str]]] ] = [ ( "qwen", "1 GPU", [ ("qwen_ln_4096x3072", "layernorm", (1, 4096, 3072), SGL_LN_PAIR), ("qwen_ln_26x3072", "layernorm", (1, 26, 3072), SGL_LN_PAIR), ("qwen_ln_6x3072", "layernorm", (1, 6, 3072), SGL_LN_PAIR), ("qwen_rms_26x3584", "rmsnorm", (1, 26, 3584), SGL_RMS), ("qwen_rms_6x3584", "rmsnorm", (1, 6, 3584), SGL_RMS), ], ), ( "qwen-edit", "1 GPU", [ ("qwen_edit_ln_189x3072", "layernorm", (1, 189, 3072), SGL_LN_PAIR), ("qwen_edit_ln_192x3072", "layernorm", (1, 192, 3072), SGL_LN_PAIR), ("qwen_edit_ln_8308x3072", "layernorm", (1, 8308, 3072), TORCH_LN), ("qwen_edit_rms_189x3584", "rmsnorm", (1, 189, 3584), SGL_RMS), ("qwen_edit_rms_192x3584", "rmsnorm", (1, 192, 3584), SGL_RMS), ], ), ( "flux", "1 GPU", [ ("flux_ln_77x768", "layernorm", (1, 77, 768), TORCH_LN), ("flux_ln_512x3072", "layernorm", (1, 512, 3072), TORCH_LN), ("flux_ln_4096x3072", "layernorm", (1, 4096, 3072), TORCH_LN), ("flux_ln_4608x3072", "layernorm", (1, 4608, 3072), TORCH_LN), ("flux_rms_512x4096", "rmsnorm", (1, 512, 4096), SGL_RMS), ], ), ( "flux2", "1 GPU", [ ("flux2_ln_512x6144", "layernorm", (1, 512, 6144), TORCH_LN), ("flux2_ln_4096x6144", "layernorm", (1, 4096, 6144), TORCH_LN), ("flux2_ln_4608x6144", "layernorm", (1, 4608, 6144), TORCH_LN), ("flux2_rms_4608x48x128", "rmsnorm", (1, 4608, 48, 128), SGL_RMS), ], ), ( "zimage", "1 GPU", [ ("zimage_ln_4128x3840", "layernorm", (1, 4128, 3840), TORCH_LN), ("zimage_rms_32x3840", "rmsnorm", (1, 32, 3840), SGL_RMS), ("zimage_rms_4096x3840", "rmsnorm", (1, 4096, 3840), SGL_RMS), ("zimage_rms_4128x3840", "rmsnorm", (1, 4128, 3840), SGL_RMS), ("zimage_rms_512x2560", "rmsnorm", (1, 512, 2560), SGL_RMS), ("zimage_rms_512x32x128", "rmsnorm", (1, 512, 32, 128), SGL_RMS), ("zimage_rms_512x8x128", "rmsnorm", (1, 512, 8, 128), SGL_RMS), ], ), ( "wan-ti2v", "1 GPU", [ ("wan_ti2v_ln_17850x3072", "layernorm", (1, 17850, 3072), SGL_LN_PAIR), ("wan_ti2v_rms_17850x3072", "rmsnorm", (1, 17850, 3072), SGL_RMS), ("wan_ti2v_rms_512x3072", "rmsnorm", (1, 512, 3072), SGL_RMS), ("wan_ti2v_rms_512x4096", "rmsnorm", (1, 512, 4096), SGL_RMS), ], ), ( "hunyuanvideo", "1 GPU", [ ("hunyuan_ln_46x768", "layernorm", (1, 46, 768), TORCH_LN), ("hunyuan_ln_45x3072", "layernorm", (1, 45, 3072), SGL_LN_PAIR), ("hunyuan_ln_27030x3072", "layernorm", (1, 27030, 3072), SGL_LN_PAIR), ("hunyuan_ln_27075x3072", "layernorm", (1, 27075, 3072), SGL_LN), ("hunyuan_rms_140x4096", "rmsnorm", (1, 140, 4096), SGL_RMS), ("hunyuan_rms_45x24x128", "rmsnorm", (1, 45, 24, 128), SGL_RMS), ("hunyuan_rms_27030x24x128", "rmsnorm", (1, 27030, 24, 128), SGL_RMS), ("hunyuan_rms_27075x24x128", "rmsnorm", (1, 27075, 24, 128), SGL_RMS), ("hunyuan_fused_add_140x4096", "fused_add_rmsnorm", (140, 4096), SGL_FUSED), ], ), ( "mova-720p", "4 GPU, ulysses=4, ring=1", [ ("mova_ln_101x1536", "layernorm", (1, 101, 1536), MOVA_LN_MIX), ("mova_ln_403x1536", "layernorm", (1, 403, 1536), TORCH_LN), ("mova_ln_44100x5120", "layernorm", (1, 44100, 5120), MOVA_LN_MIX), ("mova_ln_176400x5120", "layernorm", (1, 176400, 5120), SGL_LN), ("mova_rms_101x1536", "rmsnorm", (1, 101, 1536), SGL_RMS), ("mova_rms_101x5120", "rmsnorm", (1, 101, 5120), SGL_RMS), ("mova_rms_44100x1536", "rmsnorm", (1, 44100, 1536), SGL_RMS), ("mova_rms_44100x5120", "rmsnorm", (1, 44100, 5120), SGL_RMS), ("mova_rms_512x1536", "rmsnorm", (1, 512, 1536), SGL_RMS), ("mova_rms_512x4096", "rmsnorm", (1, 512, 4096), SGL_RMS), ("mova_rms_512x5120", "rmsnorm", (1, 512, 5120), SGL_RMS), ], ), ] ACTUAL_DIFFUSION_SHAPES: list[dict[str, object]] = [ { "shape_id": shape_id, "model": model, "gpu_config": gpu_config, "op": op, "input_shape": list(input_shape), "source_impl": source_impl, } for model, gpu_config, cases in ACTUAL_DIFFUSION_GROUPS for shape_id, op, input_shape, source_impl in cases ] def effective_rows_from_shape(input_shape: list[int]) -> int: rows = 1 for dim in input_shape[:-1]: rows *= dim return rows def ensure_repo(repo_name: str, repo_url: str) -> Path: repo_path = THIRD_PARTY_ROOT / repo_name if repo_path.exists(): return repo_path repo_path.parent.mkdir(parents=True, exist_ok=True) subprocess.run( ["git", "clone", "--depth", "1", repo_url, str(repo_path)], check=True, cwd=REPO_ROOT, ) return repo_path def ensure_python_dep(module_name: str, package_name: str | None = None) -> None: package_name = package_name or module_name try: importlib.import_module(module_name) except ModuleNotFoundError: subprocess.run( [sys.executable, "-m", "pip", "install", package_name], check=True, ) def dtype_from_name(name: str) -> torch.dtype: mapping = { "bf16": torch.bfloat16, "bfloat16": torch.bfloat16, "fp16": torch.float16, "float16": torch.float16, "fp32": torch.float32, "float32": torch.float32, } return mapping[name] def dtype_name(dtype: torch.dtype) -> str: mapping = { torch.bfloat16: "bf16", torch.float16: "fp16", torch.float32: "fp32", } return mapping[dtype] def normalize_hidden_sizes(text: str) -> list[int]: return [int(x) for x in text.split(",") if x] def normalize_dtypes(text: str) -> list[torch.dtype]: return [dtype_from_name(x.strip()) for x in text.split(",") if x.strip()] def prewarm(fn: Callable[[], object], iters: int = 3) -> None: for _ in range(iters): fn() torch.cuda.synchronize() def benchmark_provider( fn: Callable[[], object], setup_fn: Callable[[], None] | None = None, warmup: int = 10, rep: int = 30, ) -> tuple[float, float, float]: for _ in range(warmup): if setup_fn is not None: setup_fn() fn() torch.cuda.synchronize() start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) times_us: list[float] = [] for _ in range(rep): if setup_fn is not None: setup_fn() start_event.record() fn() end_event.record() end_event.synchronize() times_us.append(start_event.elapsed_time(end_event) * 1000.0) return statistics.median(times_us), max(times_us), min(times_us) def geometric_mean(values: list[float]) -> float: if not values: return float("nan") return math.exp(sum(math.log(v) for v in values) / len(values)) @functools.cache def load_flaggems(): ensure_python_dep("sqlalchemy") ensure_repo("FlagGems", FLAGGEMS_REPO) src_root = THIRD_PARTY_ROOT / "FlagGems" / "src" if str(src_root) not in sys.path: sys.path.insert(0, str(src_root)) from flag_gems.fused.fused_add_rms_norm import fused_add_rms_norm from flag_gems.ops.layernorm import layer_norm from flag_gems.ops.rms_norm import rms_norm return rms_norm, layer_norm, fused_add_rms_norm @functools.cache def load_quack(): repo_path = ensure_repo("quack", QUACK_REPO) try: quack_rmsnorm = importlib.import_module("quack.rmsnorm") except ModuleNotFoundError: subprocess.run( [sys.executable, "-m", "pip", "install", "-e", str(repo_path)], check=True, ) quack_rmsnorm = importlib.import_module("quack.rmsnorm") return quack_rmsnorm.rmsnorm_fwd, quack_rmsnorm.layernorm_fwd def build_rmsnorm_providers(dtype: torch.dtype, batch_size: int, hidden_size: int): import flashinfer.norm as flashinfer_norm import sgl_kernel x = torch.randn((batch_size, hidden_size), device=DEFAULT_DEVICE, dtype=dtype) weight = torch.randn(hidden_size, device=DEFAULT_DEVICE, dtype=dtype) jit_out = torch.empty_like(x) sgl_out = torch.empty_like(x) flashinfer_out = torch.empty_like(x) flaggems_rms_norm, _, _ = load_flaggems() quack_rmsnorm_fwd, _ = load_quack() providers = { "pytorch": lambda: F.rms_norm(x, (hidden_size,), weight, 1e-6), "sgl_kernel": lambda: sgl_kernel.rmsnorm(x, weight, eps=1e-6, out=sgl_out), "flashinfer": lambda: flashinfer_norm.rmsnorm( x, weight, eps=1e-6, out=flashinfer_out ), "jit_rmsnorm": lambda: jit_rmsnorm(x, weight, jit_out, 1e-6), "quack": lambda: quack_rmsnorm_fwd(x, weight, eps=1e-6), "triton_rms_norm_fn": lambda: rms_norm_fn( x, weight, bias=None, residual=None, eps=1e-6 ), "flaggems": lambda: flaggems_rms_norm(x, (hidden_size,), weight, 1e-6), } if hidden_size <= 128: providers["triton_one_pass"] = lambda: triton_one_pass_rms_norm(x, weight, 1e-6) return providers def build_fused_add_rmsnorm_providers( dtype: torch.dtype, batch_size: int, hidden_size: int ): import flashinfer.norm as flashinfer_norm import sgl_kernel base_x = torch.randn((batch_size, hidden_size), device=DEFAULT_DEVICE, dtype=dtype) base_residual = torch.randn_like(base_x) weight = torch.randn(hidden_size, device=DEFAULT_DEVICE, dtype=dtype) x = base_x.clone() residual = base_residual.clone() def reset(): x.copy_(base_x) residual.copy_(base_residual) _, _, flaggems_fused_add_rms_norm = load_flaggems() quack_rmsnorm_fwd, _ = load_quack() def pytorch_impl(): out = x + residual return F.rms_norm(out, (hidden_size,), weight, 1e-6) providers = { "pytorch": (pytorch_impl, reset), "sgl_kernel": ( lambda: sgl_kernel.fused_add_rmsnorm(x, residual, weight, eps=1e-6), reset, ), "flashinfer": ( lambda: flashinfer_norm.fused_add_rmsnorm(x, residual, weight, eps=1e-6), reset, ), "jit_fused_add_rmsnorm": ( lambda: jit_fused_add_rmsnorm(x, residual, weight, 1e-6), reset, ), "quack": ( lambda: quack_rmsnorm_fwd(x, weight, residual=residual, eps=1e-6), reset, ), "flaggems": ( lambda: flaggems_fused_add_rms_norm( x, residual, (hidden_size,), weight, 1e-6 ), reset, ), } return providers def build_layernorm_providers(dtype: torch.dtype, batch_size: int, hidden_size: int): import flashinfer.norm as flashinfer_norm x = torch.randn((batch_size, hidden_size), device=DEFAULT_DEVICE, dtype=dtype) weight = torch.randn(hidden_size, device=DEFAULT_DEVICE, dtype=dtype) bias = torch.randn(hidden_size, device=DEFAULT_DEVICE, dtype=dtype) flashinfer_weight = torch.randn( hidden_size, device=DEFAULT_DEVICE, dtype=torch.float32 ) flashinfer_bias = torch.randn( hidden_size, device=DEFAULT_DEVICE, dtype=torch.float32 ) triton_out = torch.empty_like(x) _, flaggems_layer_norm, _ = load_flaggems() _, quack_layernorm_fwd = load_quack() providers = { "pytorch": lambda: F.layer_norm(x, (hidden_size,), weight, bias, 1e-6), "triton_norm_infer": lambda: norm_infer( x, weight, bias, eps=1e-6, is_rms_norm=False, out=triton_out ), "flashinfer": lambda: flashinfer_norm.layernorm( x, flashinfer_weight, flashinfer_bias, 1e-6 ), "quack": lambda: quack_layernorm_fwd( x, flashinfer_weight, flashinfer_bias, 1e-6 ), "flaggems": lambda: flaggems_layer_norm(x, (hidden_size,), weight, bias)[0], } return providers def maybe_benchmark( op_name: str, provider_name: str, fn: Callable[[], object], rows: list[dict[str, object]], dtype: torch.dtype, batch_size: int, hidden_size: int, reset: Callable[[], None] | None = None, metadata: dict[str, object] | None = None, ) -> None: metadata = metadata or {} try: median_us, max_us, min_us = benchmark_provider(fn, reset) rows.append( { "op": op_name, "provider": provider_name, "dtype": dtype_name(dtype), "batch_size": batch_size, "hidden_size": hidden_size, "median_us": median_us, "min_us": min_us, "max_us": max_us, "status": "ok", "error": "", **metadata, } ) except Exception as exc: # pragma: no cover - benchmark failures are data rows.append( { "op": op_name, "provider": provider_name, "dtype": dtype_name(dtype), "batch_size": batch_size, "hidden_size": hidden_size, "median_us": "", "min_us": "", "max_us": "", "status": "unsupported", "error": str(exc), **metadata, } ) def write_csv(rows: list[dict[str, object]], output_path: Path) -> None: output_path.parent.mkdir(parents=True, exist_ok=True) with output_path.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter( f, fieldnames=[ "op", "provider", "dtype", "batch_size", "hidden_size", "median_us", "min_us", "max_us", "shape_id", "source_model", "source_gpu_config", "source_input_shape", "source_impl", "status", "error", ], ) writer.writeheader() writer.writerows(rows) def write_markdown(rows: list[dict[str, object]], output_path: Path) -> None: lines: list[str] = [] lines.append("# Norm Benchmark Summary") lines.append("") actual_shape_rows = [row for row in rows if row.get("shape_id")] if actual_shape_rows: seen: set[tuple[str, str, str, str, str, str]] = set() lines.append("## Diffusion Shape Cases") lines.append("") lines.append( "| Shape ID | Op | Model | GPU Config | Input Shape | Source Impl |" ) lines.append("|---|---|---|---|---|---|") for row in actual_shape_rows: key = ( str(row.get("shape_id", "")), str(row.get("op", "")), str(row.get("source_model", "")), str(row.get("source_gpu_config", "")), str(row.get("source_input_shape", "")), str(row.get("source_impl", "")), ) if key in seen: continue seen.add(key) lines.append( f"| {key[0]} | {key[1]} | {key[2]} | {key[3]} | `{key[4]}` | {key[5]} |" ) lines.append("") for op_name in ("rmsnorm", "fused_add_rmsnorm", "layernorm"): for dtype in sorted({row["dtype"] for row in rows}): scoped = [ row for row in rows if row["op"] == op_name and row["dtype"] == dtype and row["status"] == "ok" ] if not scoped: continue provider_to_values: dict[str, list[float]] = {} provider_to_speedups: dict[str, list[float]] = {} by_shape: dict[tuple[str, int, int], dict[str, float]] = {} for row in scoped: provider = str(row["provider"]) value = float(row["median_us"]) provider_to_values.setdefault(provider, []).append(value) shape = ( str(row.get("shape_id", "")), int(row["batch_size"]), int(row["hidden_size"]), ) by_shape.setdefault(shape, {})[provider] = value for shape, perf in by_shape.items(): if "pytorch" not in perf: continue baseline = perf["pytorch"] for provider, value in perf.items(): provider_to_speedups.setdefault(provider, []).append( baseline / value ) lines.append(f"## {op_name} ({dtype})") lines.append("") lines.append( "| Provider | Geomean Speedup vs PyTorch | Median Latency (us) | Win Count |" ) lines.append("|---|---:|---:|---:|") wins: dict[str, int] = {} for perf in by_shape.values(): best_provider = min(perf, key=perf.get) wins[best_provider] = wins.get(best_provider, 0) + 1 for provider in sorted(provider_to_values): geomean_speedup = geometric_mean(provider_to_speedups.get(provider, [])) median_latency = statistics.median(provider_to_values[provider]) win_count = wins.get(provider, 0) lines.append( f"| {provider} | {geomean_speedup:.3f}x | {median_latency:.2f} | {win_count} |" ) lines.append("") output_path.write_text("\n".join(lines) + "\n", encoding="utf-8") def run_suite( hidden_sizes: list[int], batch_sizes: list[int], dtypes: list[torch.dtype], ops: list[str], ) -> list[dict[str, object]]: rows: list[dict[str, object]] = [] for dtype in dtypes: for batch_size in batch_sizes: for hidden_size in hidden_sizes: if "rmsnorm" in ops: rms_providers = build_rmsnorm_providers( dtype, batch_size, hidden_size ) for provider_name, fn in rms_providers.items(): maybe_benchmark( "rmsnorm", provider_name, fn, rows, dtype, batch_size, hidden_size, ) if "fused_add_rmsnorm" in ops: fused_providers = build_fused_add_rmsnorm_providers( dtype, batch_size, hidden_size ) for provider_name, provider in fused_providers.items(): fn, reset = provider maybe_benchmark( "fused_add_rmsnorm", provider_name, fn, rows, dtype, batch_size, hidden_size, reset, ) if "layernorm" in ops: layernorm_providers = build_layernorm_providers( dtype, batch_size, hidden_size ) for provider_name, fn in layernorm_providers.items(): maybe_benchmark( "layernorm", provider_name, fn, rows, dtype, batch_size, hidden_size, ) return rows def run_shape_suite( shape_cases: list[dict[str, object]], dtypes: list[torch.dtype], ) -> list[dict[str, object]]: rows: list[dict[str, object]] = [] for case in shape_cases: op_name = str(case["op"]) input_shape = [int(x) for x in case["input_shape"]] batch_size = effective_rows_from_shape(input_shape) hidden_size = input_shape[-1] metadata = { "shape_id": str(case["shape_id"]), "source_model": str(case["model"]), "source_gpu_config": str(case["gpu_config"]), "source_input_shape": str(input_shape), "source_impl": str(case["source_impl"]), } for dtype in dtypes: if op_name == "rmsnorm": providers = build_rmsnorm_providers(dtype, batch_size, hidden_size) for provider_name, fn in providers.items(): maybe_benchmark( op_name, provider_name, fn, rows, dtype, batch_size, hidden_size, metadata=metadata, ) elif op_name == "fused_add_rmsnorm": providers = build_fused_add_rmsnorm_providers( dtype, batch_size, hidden_size ) for provider_name, provider in providers.items(): fn, reset = provider maybe_benchmark( op_name, provider_name, fn, rows, dtype, batch_size, hidden_size, reset, metadata=metadata, ) elif op_name == "layernorm": providers = build_layernorm_providers(dtype, batch_size, hidden_size) for provider_name, fn in providers.items(): maybe_benchmark( op_name, provider_name, fn, rows, dtype, batch_size, hidden_size, metadata=metadata, ) else: raise ValueError(f"Unsupported op in shape preset: {op_name}") return rows def main() -> None: parser = argparse.ArgumentParser( description="Benchmark RMSNorm/LayerNorm implementations across providers." ) parser.add_argument( "--hidden-sizes", default="64,128,256,512,1024,2048,4096,8192,16384", help="Comma-separated hidden sizes.", ) parser.add_argument( "--batch-sizes", default="1,16,128,1024", help="Comma-separated batch sizes.", ) parser.add_argument( "--dtypes", default="bf16,fp16", help="Comma-separated dtypes: bf16, fp16, fp32.", ) parser.add_argument( "--output-dir", default=str(REPO_ROOT / "outputs" / "norm_benchmarks"), help="Directory for CSV/Markdown outputs.", ) parser.add_argument( "--ops", default="rmsnorm,fused_add_rmsnorm,layernorm", help="Comma-separated ops to benchmark.", ) parser.add_argument( "--shape-preset", choices=["grid", "diffusion-actual"], default="grid", help="Use the default grid sweep or the captured diffusion workload shapes.", ) args = parser.parse_args() if not torch.cuda.is_available(): raise RuntimeError("CUDA is required for norm benchmarks.") hidden_sizes = normalize_hidden_sizes(args.hidden_sizes) batch_sizes = normalize_hidden_sizes(args.batch_sizes) dtypes = normalize_dtypes(args.dtypes) ops = [op.strip() for op in args.ops.split(",") if op.strip()] if args.shape_preset == "diffusion-actual": shape_cases = [case for case in ACTUAL_DIFFUSION_SHAPES if case["op"] in ops] rows = run_shape_suite(shape_cases, dtypes) else: rows = run_suite(hidden_sizes, batch_sizes, dtypes, ops) output_dir = Path(args.output_dir) csv_path = output_dir / "norm_impls.csv" md_path = output_dir / "norm_impls_summary.md" write_csv(rows, csv_path) write_markdown(rows, md_path) print(f"Wrote {csv_path}") print(f"Wrote {md_path}") if __name__ == "__main__": if is_in_ci(): print("Skipping bench_norm_impls.py in CI") sys.exit(0) main() ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py ================================================ from __future__ import annotations import sys from typing import Any import torch import triton from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark from sglang.jit_kernel.nvfp4 import ( cutlass_fp4_group_mm, scaled_fp4_experts_quant, scaled_fp4_quant, ) from sglang.srt.utils import is_sm100_supported FLOAT4_E2M1_MAX = 6.0 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max _NVFP4_SUPPORTED = is_sm100_supported() def _round_up(x: int, y: int) -> int: return ((x + y - 1) // y) * y def _expert_offsets(m_per_expert: list[int], device: torch.device) -> torch.Tensor: offsets = [0] for m in m_per_expert: offsets.append(offsets[-1] + m) return torch.tensor(offsets, dtype=torch.int32, device=device) def _blockscale_offsets(m_per_expert: list[int], device: torch.device) -> torch.Tensor: offsets = [0] for m in m_per_expert: offsets.append(offsets[-1] + _round_up(m, 128)) return torch.tensor(offsets, dtype=torch.int32, device=device) def _prepare_case( total_tokens: int, n: int, k: int, num_experts: int, dtype: torch.dtype ) -> dict[str, Any]: device = torch.device("cuda") base = total_tokens // num_experts rem = total_tokens % num_experts m_per_expert = [base + (1 if i < rem else 0) for i in range(num_experts)] expert_offsets_full = _expert_offsets(m_per_expert, device) blockscale_offsets_full = _blockscale_offsets(m_per_expert, device) a = torch.randn((total_tokens, k), device=device, dtype=dtype) * 0.1 b = torch.randn((num_experts, n, k), device=device, dtype=dtype) * 0.1 a_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32) for i in range(num_experts): start = int(expert_offsets_full[i].item()) end = int(expert_offsets_full[i + 1].item()) a_global_scale[i] = ( FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / a[start:end].abs().max().to(torch.float32) ) b_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32) for i in range(num_experts): b_global_scale[i] = ( FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b[i].abs().max().to(torch.float32) ) a_fp4, a_blockscale = scaled_fp4_experts_quant( a, a_global_scale, expert_offsets_full, blockscale_offsets_full, topk=1, ) b_fp4 = torch.empty((num_experts, n, k // 2), device=device, dtype=torch.uint8) b_blockscale = torch.empty( (num_experts, _round_up(n, 128), _round_up(k // 16, 4)), device=device, dtype=torch.float8_e4m3fn, ) for i in range(num_experts): b_fp4_i, b_scale_i = scaled_fp4_quant(b[i], b_global_scale[i]) b_fp4[i].copy_(b_fp4_i) b_blockscale[i].copy_(b_scale_i) alphas = (1.0 / (a_global_scale * b_global_scale)).to(torch.float32) params = { "ab_strides": torch.full((num_experts,), k, dtype=torch.int64, device=device), "c_strides": torch.full((num_experts,), n, dtype=torch.int64, device=device), "problem_sizes": torch.tensor( [[m, n, k] for m in m_per_expert], dtype=torch.int32, device=device ), "expert_offsets": expert_offsets_full[:-1].contiguous(), "blockscale_offsets": blockscale_offsets_full[:-1].contiguous(), "a_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), "b_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), "out_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), "a_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), "b_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), "alpha_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), "layout_sfa": torch.empty((num_experts, 5), dtype=torch.int64, device=device), "layout_sfb": torch.empty((num_experts, 5), dtype=torch.int64, device=device), } expert_ranges: list[tuple[int, int]] = [] start = 0 for m in m_per_expert: end = start + m expert_ranges.append((start, end)) start = end return { "a": a, "b": b, "a_fp4": a_fp4, "b_fp4": b_fp4, "a_blockscale": a_blockscale, "b_blockscale": b_blockscale, "alphas": alphas, "params": params, "expert_offsets_full": expert_offsets_full, "expert_ranges": expert_ranges, "dtype": dtype, } def _torch_ref_group_mm(case: dict[str, Any]) -> torch.Tensor: a = case["a"] b = case["b"] dtype = case["dtype"] expert_ranges = case["expert_ranges"] total_tokens = a.shape[0] n = b.shape[1] out = torch.empty((total_tokens, n), device=a.device, dtype=dtype) for i, (start, end) in enumerate(expert_ranges): out[start:end] = torch.matmul(a[start:end], b[i].t()) return out def _aot_cutlass_fp4_group_mm(case: dict[str, Any]) -> torch.Tensor: a_fp4 = case["a_fp4"] b_fp4 = case["b_fp4"] a_blockscale = case["a_blockscale"] b_blockscale = case["b_blockscale"] alphas = case["alphas"] params = case["params"] out_dtype = case["dtype"] out = torch.empty( (a_fp4.shape[0], b_fp4.shape[1]), device=a_fp4.device, dtype=out_dtype ) torch.ops.sgl_kernel.cutlass_fp4_group_mm.default( out, a_fp4, b_fp4, a_blockscale, b_blockscale, alphas, params["ab_strides"], params["c_strides"], params["problem_sizes"], params["expert_offsets"], params["blockscale_offsets"], ) return out def _probe_legacy_aot_group_mm() -> tuple[bool, str]: if not torch.cuda.is_available(): return False, "CUDA is not available." if not _NVFP4_SUPPORTED: return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+." try: import sgl_kernel # noqa: F401 except Exception as e: return False, f"import sgl_kernel failed: {e}" if not hasattr(torch.ops, "sgl_kernel"): return False, "torch.ops.sgl_kernel is not registered." op = getattr(torch.ops.sgl_kernel, "cutlass_fp4_group_mm", None) if op is None or not hasattr(op, "default"): return False, "torch.ops.sgl_kernel.cutlass_fp4_group_mm.default is missing." try: case = _prepare_case(64, 256, 128, 4, torch.bfloat16) _aot_cutlass_fp4_group_mm(case) torch.cuda.synchronize() except Exception as e: return False, f"calling AOT grouped_mm op failed: {e}" return True, "" _AOT_GROUP_MM_AVAILABLE, _AOT_GROUP_MM_REASON = _probe_legacy_aot_group_mm() shape_range = get_benchmark_range( full_range=[(128, 256, 128, 4), (256, 512, 128, 8), (512, 512, 256, 8)], ci_range=[(128, 256, 128, 4)], ) line_vals = ["jit"] line_names = ["JIT NVFP4 MoE GroupMM"] styles = [("green", "-")] if _AOT_GROUP_MM_AVAILABLE: line_vals.append("aot_sgl_kernel") line_names.append("AOT NVFP4 MoE GroupMM") styles.append(("orange", "-")) line_vals.append("torch_ref") line_names.append("Torch Ref") styles.append(("blue", "-")) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["total_tokens", "n", "k", "num_experts"], x_vals=shape_range, x_log=False, line_arg="provider", line_vals=line_vals, line_names=line_names, styles=styles, ylabel="us", plot_name="nvfp4-blockwise-moe-groupmm-performance", args={}, ) ) def benchmark(total_tokens, n, k, num_experts, provider): case = _prepare_case(total_tokens, n, k, num_experts, torch.bfloat16) if provider == "jit": fn = lambda: cutlass_fp4_group_mm( case["a_fp4"], case["b_fp4"], case["a_blockscale"], case["b_blockscale"], case["alphas"], case["dtype"], case["params"], ) elif provider == "aot_sgl_kernel": fn = lambda: _aot_cutlass_fp4_group_mm(case) elif provider == "torch_ref": fn = lambda: _torch_ref_group_mm(case) else: raise ValueError(f"Unknown provider: {provider}") return run_benchmark(fn) if __name__ == "__main__": if not _NVFP4_SUPPORTED: print("[skip] NVFP4 blockwise MoE benchmark requires sm100+ with CUDA 12.8+.") sys.exit(0) if not _AOT_GROUP_MM_AVAILABLE: print( f"[info] legacy AOT grouped_mm baseline unavailable: {_AOT_GROUP_MM_REASON}" ) benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py ================================================ from __future__ import annotations import sys import torch import triton from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark from sglang.jit_kernel.nvfp4 import scaled_fp4_quant from sglang.srt.utils import is_sm100_supported FLOAT4_E2M1_MAX = 6.0 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max BLOCK_SIZE = 16 _NVFP4_SUPPORTED = is_sm100_supported() try: from flashinfer import fp4_quantize as flashinfer_fp4_quantize except Exception: flashinfer_fp4_quantize = None def _torch_ref_quant(input: torch.Tensor, input_global_scale: torch.Tensor): m, n = input.shape x = input.view(m, n // BLOCK_SIZE, BLOCK_SIZE) vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32) scale = input_global_scale * (vec_max / FLOAT4_E2M1_MAX) scale = scale.to(torch.float8_e4m3fn).to(torch.float32) output_scale = torch.where(scale == 0, torch.zeros_like(scale), 1.0 / scale) scaled_x = x.to(torch.float32) * output_scale clipped = torch.clamp(scaled_x, -6.0, 6.0).reshape(m, n) rounded = clipped.clone() rounded[(rounded >= 0.0) & (rounded <= 0.25)] = 0.0 rounded[(rounded > 0.25) & (rounded < 0.75)] = 0.5 rounded[(rounded >= 0.75) & (rounded <= 1.25)] = 1.0 rounded[(rounded > 1.25) & (rounded < 1.75)] = 1.5 rounded[(rounded >= 1.75) & (rounded <= 2.5)] = 2.0 rounded[(rounded > 2.5) & (rounded < 3.5)] = 3.0 rounded[(rounded >= 3.5) & (rounded <= 5.0)] = 4.0 rounded[rounded > 5.0] = 6.0 # This baseline intentionally keeps work on GPU but does not pack to uint8. return rounded, scale def _aot_scaled_fp4_quant(input: torch.Tensor, input_global_scale: torch.Tensor): m, n = input.shape output = torch.empty((m, n // 2), device=input.device, dtype=torch.uint8) rounded_m = ((m + 128 - 1) // 128) * 128 scale_n = n // BLOCK_SIZE rounded_n = ((scale_n + 4 - 1) // 4) * 4 output_scale = torch.empty( (rounded_m, rounded_n // 4), device=input.device, dtype=torch.int32 ) torch.ops.sgl_kernel.scaled_fp4_quant.default( output, input, output_scale, input_global_scale ) return output, output_scale.view(torch.float8_e4m3fn) def _probe_legacy_aot_quant() -> tuple[bool, str]: if not torch.cuda.is_available(): return False, "CUDA is not available." if not _NVFP4_SUPPORTED: return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+." try: import sgl_kernel # noqa: F401 except Exception as e: return False, f"import sgl_kernel failed: {e}" if not hasattr(torch.ops, "sgl_kernel"): return False, "torch.ops.sgl_kernel is not registered." op = getattr(torch.ops.sgl_kernel, "scaled_fp4_quant", None) if op is None or not hasattr(op, "default"): return False, "torch.ops.sgl_kernel.scaled_fp4_quant.default is missing." try: x = torch.randn((16, 64), dtype=torch.bfloat16, device="cuda") global_scale = ( FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.abs(x).max().to(torch.float32) ) _aot_scaled_fp4_quant(x, global_scale) torch.cuda.synchronize() except Exception as e: return False, f"calling AOT quant op failed: {e}" return True, "" _AOT_QUANT_AVAILABLE, _AOT_QUANT_REASON = _probe_legacy_aot_quant() def _probe_flashinfer_quant() -> tuple[bool, str]: if flashinfer_fp4_quantize is None: return False, "import flashinfer.fp4_quantize failed." if not torch.cuda.is_available(): return False, "CUDA is not available." if not _NVFP4_SUPPORTED: return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+." try: x = torch.randn((16, 64), dtype=torch.bfloat16, device="cuda") global_scale = ( FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.abs(x).max().to(torch.float32) ) flashinfer_fp4_quantize( x, global_scale, BLOCK_SIZE, # sf_vec_size False, # use_ue8m0 True, # is_sf_swizzled_layout ) torch.cuda.synchronize() except Exception as e: return False, f"calling flashinfer.fp4_quantize failed: {e}" return True, "" _FLASHINFER_QUANT_AVAILABLE, _FLASHINFER_QUANT_REASON = _probe_flashinfer_quant() shape_range = get_benchmark_range( full_range=[(128, 2048), (512, 4096), (1024, 4096), (2048, 8192)], ci_range=[(128, 2048)], ) line_vals = [] line_names = [] styles = [] if _FLASHINFER_QUANT_AVAILABLE: line_vals.append("flashinfer") line_names.append("FlashInfer FP4 Quant") styles.append(("purple", "-")) line_vals.append("jit") line_names.append("JIT NVFP4 Quant") styles.append(("green", "-")) if _AOT_QUANT_AVAILABLE: line_vals.append("aot_sgl_kernel") line_names.append("AOT NVFP4 Quant") styles.append(("orange", "-")) line_vals.append("torch_ref") line_names.append("Torch Ref") styles.append(("blue", "-")) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["m", "n"], x_vals=shape_range, x_log=False, line_arg="provider", line_vals=line_vals, line_names=line_names, styles=styles, ylabel="us", plot_name="nvfp4-quant-performance", args={}, ) ) def benchmark(m, n, provider): x = torch.randn((m, n), dtype=torch.bfloat16, device="cuda") tensor_amax = torch.abs(x).max().to(torch.float32) global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax if provider == "jit": fn = lambda: scaled_fp4_quant(x, global_scale) elif provider == "flashinfer": fn = lambda: flashinfer_fp4_quantize( x, global_scale, BLOCK_SIZE, # sf_vec_size False, # use_ue8m0 True, # is_sf_swizzled_layout ) elif provider == "aot_sgl_kernel": fn = lambda: _aot_scaled_fp4_quant(x, global_scale) elif provider == "torch_ref": fn = lambda: _torch_ref_quant(x, global_scale) else: raise ValueError(f"Unknown provider: {provider}") return run_benchmark(fn) if __name__ == "__main__": if not _NVFP4_SUPPORTED: print("[skip] NVFP4 quant benchmark requires sm100+ with CUDA 12.8+.") sys.exit(0) if not _FLASHINFER_QUANT_AVAILABLE: print( f"[info] flashinfer quant baseline unavailable: {_FLASHINFER_QUANT_REASON}" ) if not _AOT_QUANT_AVAILABLE: print(f"[info] legacy AOT quant baseline unavailable: {_AOT_QUANT_REASON}") benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py ================================================ from __future__ import annotations import sys import torch import triton from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant from sglang.srt.utils import is_sm100_supported FLOAT4_E2M1_MAX = 6.0 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max BLOCK_SIZE = 16 _NVFP4_SUPPORTED = is_sm100_supported() K_E2M1_TO_FLOAT = [ 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, ] def _dequantize_to_fp16( tensor_fp4: torch.Tensor, tensor_sf: torch.Tensor, global_scale: torch.Tensor ): m, packed_k = tensor_fp4.shape k = packed_k * 2 flat = tensor_fp4.flatten() high = (flat & 0xF0) >> 4 low = flat & 0x0F f_h = torch.tensor([K_E2M1_TO_FLOAT[x] for x in high], device=tensor_fp4.device) f_l = torch.tensor([K_E2M1_TO_FLOAT[x] for x in low], device=tensor_fp4.device) val = torch.stack((f_l, f_h), dim=-1).reshape(m, k) rounded_m = ((m + 128 - 1) // 128) * 128 scale_n = k // BLOCK_SIZE rounded_n = ((scale_n + 4 - 1) // 4) * 4 sf = tensor_sf.view(torch.float8_e4m3fn) tmp = torch.reshape(sf, (1, rounded_m // 128, rounded_n // 4, 32, 4, 4)) tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5)) scale = torch.reshape(tmp, (rounded_m, rounded_n))[:m, :scale_n].to(torch.float32) scale = scale / global_scale return (val.view(m, scale_n, BLOCK_SIZE) * scale.unsqueeze(-1)).reshape(m, k) def _aot_cutlass_scaled_fp4_mm( a: torch.Tensor, b: torch.Tensor, block_scale_a: torch.Tensor, block_scale_b: torch.Tensor, alpha: torch.Tensor, out_dtype: torch.dtype, ) -> torch.Tensor: out = torch.empty((a.shape[0], b.shape[0]), dtype=out_dtype, device=a.device) torch.ops.sgl_kernel.cutlass_scaled_fp4_mm.default( out, a, b, block_scale_a, block_scale_b, alpha ) return out def _probe_legacy_aot_scaled_mm() -> tuple[bool, str]: if not torch.cuda.is_available(): return False, "CUDA is not available." if not _NVFP4_SUPPORTED: return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+." try: import sgl_kernel # noqa: F401 except Exception as e: return False, f"import sgl_kernel failed: {e}" if not hasattr(torch.ops, "sgl_kernel"): return False, "torch.ops.sgl_kernel is not registered." op = getattr(torch.ops.sgl_kernel, "cutlass_scaled_fp4_mm", None) if op is None or not hasattr(op, "default"): return False, "torch.ops.sgl_kernel.cutlass_scaled_fp4_mm.default is missing." try: m, n, k = 16, 32, 64 a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") b = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") a_global_scale = ( FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(a.flatten(), dim=-1) ).to(torch.float32) b_global_scale = ( FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(b.flatten(), dim=-1) ).to(torch.float32) alpha = 1.0 / (a_global_scale * b_global_scale) a_fp4, a_sf = scaled_fp4_quant(a, a_global_scale) b_fp4, b_sf = scaled_fp4_quant(b, b_global_scale) _aot_cutlass_scaled_fp4_mm(a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16) torch.cuda.synchronize() except Exception as e: return False, f"calling AOT scaled_mm op failed: {e}" return True, "" _AOT_SCALED_MM_AVAILABLE, _AOT_SCALED_MM_REASON = _probe_legacy_aot_scaled_mm() shape_range = get_benchmark_range( full_range=[(128, 4096, 4096), (512, 4096, 4096), (1024, 8192, 4096)], ci_range=[(128, 4096, 4096)], ) line_vals = ["jit"] line_names = ["JIT NVFP4 GEMM"] styles = [("green", "-")] if _AOT_SCALED_MM_AVAILABLE: line_vals.append("aot_sgl_kernel") line_names.append("AOT NVFP4 GEMM") styles.append(("orange", "-")) line_vals.append("torch_ref") line_names.append("Torch Ref") styles.append(("blue", "-")) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["m", "n", "k"], x_vals=shape_range, x_log=False, line_arg="provider", line_vals=line_vals, line_names=line_names, styles=styles, ylabel="us", plot_name="nvfp4-scaled-mm-performance", args={}, ) ) def benchmark(m, n, k, provider): a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") b = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") a_global_scale = ( FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(a.flatten(), dim=-1) ).to(torch.float32) b_global_scale = ( FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(b.flatten(), dim=-1) ).to(torch.float32) alpha = 1.0 / (a_global_scale * b_global_scale) a_fp4, a_sf = scaled_fp4_quant(a, a_global_scale) b_fp4, b_sf = scaled_fp4_quant(b, b_global_scale) if provider == "jit": fn = lambda: cutlass_scaled_fp4_mm( a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16 ) elif provider == "aot_sgl_kernel": fn = lambda: _aot_cutlass_scaled_fp4_mm( a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16 ) elif provider == "torch_ref": a_ref = _dequantize_to_fp16(a_fp4, a_sf, a_global_scale) b_ref = _dequantize_to_fp16(b_fp4, b_sf, b_global_scale) fn = lambda: torch.matmul(a_ref, b_ref.t()) else: raise ValueError(f"Unknown provider: {provider}") return run_benchmark(fn) if __name__ == "__main__": if not _NVFP4_SUPPORTED: print("[skip] NVFP4 scaled_mm benchmark requires sm100+ with CUDA 12.8+.") sys.exit(0) if not _AOT_SCALED_MM_AVAILABLE: print( f"[info] legacy AOT scaled_mm baseline unavailable: {_AOT_SCALED_MM_REASON}" ) benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_per_tensor_quant_fp8.py ================================================ from typing import Optional, Tuple import torch import triton import triton.testing from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 try: from vllm import _custom_ops as ops VLLM_AVAILABLE = True except ImportError: ops = None VLLM_AVAILABLE = False try: from sglang.srt.utils import is_hip _is_hip = is_hip() except ImportError: _is_hip = False fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn def vllm_scaled_fp8_quant( input: torch.Tensor, scale: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: if not VLLM_AVAILABLE: return sglang_scaled_fp8_quant(input, scale) return ops.scaled_fp8_quant(input, scale) def sglang_scaled_fp8_quant( input: torch.Tensor, scale: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: fp8_type_: torch.dtype = torch.float8_e4m3fn output = torch.empty_like(input, device=input.device, dtype=fp8_type_) is_static = True if scale is None: scale = torch.zeros(1, device=input.device, dtype=torch.float32) is_static = False per_tensor_quant_fp8(input, output, scale, is_static) return output, scale def calculate_diff(batch_size: int, seq_len: int): device = torch.device("cuda") x = torch.rand((batch_size, seq_len), dtype=torch.bfloat16, device=device) if not VLLM_AVAILABLE: print("vLLM not available, skipping comparison") return vllm_out, vllm_scale = vllm_scaled_fp8_quant(x) sglang_out, sglang_scale = sglang_scaled_fp8_quant(x) vllm_out = vllm_out.to(torch.float32) sglang_out = sglang_out.to(torch.float32) triton.testing.assert_close(vllm_out, sglang_out, rtol=1e-3, atol=1e-3) triton.testing.assert_close(vllm_scale, sglang_scale, rtol=1e-3, atol=1e-3) # Benchmark configuration element_range = get_benchmark_range( full_range=[2**n for n in range(10, 20)], ci_range=[16384], ) if VLLM_AVAILABLE: line_vals = ["vllm", "sglang"] line_names = ["VLLM", "SGL Kernel"] styles = [("blue", "-"), ("green", "-")] else: line_vals = ["sglang"] line_names = ["SGL Kernel"] styles = [("green", "-")] @triton.testing.perf_report( triton.testing.Benchmark( x_names=["element_count"], x_vals=element_range, line_arg="provider", line_vals=line_vals, line_names=line_names, styles=styles, ylabel="us", plot_name="per-tensor-quant-fp8-performance", args={}, ) ) def benchmark(element_count, provider): dtype = torch.float16 device = torch.device("cuda") x = torch.randn(element_count, 4096, device=device, dtype=dtype) if provider == "vllm": fn = lambda: vllm_scaled_fp8_quant(x.clone()) elif provider == "sglang": fn = lambda: sglang_scaled_fp8_quant(x.clone()) else: raise ValueError(f"Unknown provider: {provider}") return run_benchmark(fn) if __name__ == "__main__": calculate_diff(batch_size=4, seq_len=4096) benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_per_token_group_quant_8bit.py ================================================ import itertools from typing import Any, Dict, List import torch import triton from sgl_kernel.test_utils import create_per_token_group_quant_test_data from sglang.jit_kernel.benchmark.utils import ( get_benchmark_range, is_in_ci, ) from sglang.jit_kernel.per_token_group_quant_8bit import ( per_token_group_quant_8bit as sglang_per_token_group_quant_8bit, ) from sglang.srt.layers.quantization.fp8_kernel import ( create_per_token_group_quant_fp8_output_scale, ) from sglang.srt.layers.quantization.fp8_kernel import ( per_token_group_quant_8bit as triton_per_token_group_quant_8bit, ) from sglang.srt.utils import is_hip from sglang.srt.utils.bench_utils import bench_kineto IS_CI = is_in_ci() _is_hip = is_hip() fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn NUM_TESTS = 30 if IS_CI else 300 GROUP_SIZE_RANGE = [128] DST_DTYPE_RANGE = [fp8_type_] # ---- GEMM-like branch (num_ranks=None) ---- NUM_TOKENS_RANGE_GEMM = get_benchmark_range( full_range=[1, 4, 16, 64, 256, 768, 2048, 8192, 16384], ci_range=[768], ) HIDDEN_DIM_RANGE_GEMM = [1536, 7168, 16384] NUM_RANKS_RANGE_GEMM = [None] FLAGS_GEMM_FULL: List[Dict[str, Any]] = [ dict( column_major_scales=False, scale_tma_aligned=False, scale_ue8m0=False, fuse_silu_and_mul=False, masked_layout_mode=None, ), dict( column_major_scales=True, scale_tma_aligned=False, scale_ue8m0=False, fuse_silu_and_mul=False, masked_layout_mode=None, ), dict( column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=False, fuse_silu_and_mul=False, masked_layout_mode=None, ), dict( column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=True, fuse_silu_and_mul=False, masked_layout_mode=None, ), ] FLAGS_GEMM_CI: List[Dict[str, Any]] = [ dict( column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=True, fuse_silu_and_mul=False, masked_layout_mode=None, ), ] FLAGS_RANGE_GEMM = get_benchmark_range( full_range=FLAGS_GEMM_FULL, ci_range=FLAGS_GEMM_CI ) CONFIGS_GEMM = list( itertools.product( NUM_TOKENS_RANGE_GEMM, HIDDEN_DIM_RANGE_GEMM, GROUP_SIZE_RANGE, NUM_RANKS_RANGE_GEMM, DST_DTYPE_RANGE, FLAGS_RANGE_GEMM, ) ) # ---- MoE-like / multi-rank branch (hidden_dim=2048, num_ranks in {8,16,32,48}) ---- NUM_TOKENS_RANGE_MOE = get_benchmark_range( full_range=[1 * 8, 4 * 8, 64 * 8, 256 * 8, 768 * 8], ci_range=[768 * 8], ) HIDDEN_DIM_RANGE_MOE = [2048] NUM_RANKS_RANGE_MOE = get_benchmark_range( full_range=[8, 16, 32, 48], ci_range=[48], ) FLAGS_MOE: List[Dict[str, Any]] = [ dict( column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=True, fuse_silu_and_mul=True, masked_layout_mode=None, ), dict( column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=True, fuse_silu_and_mul=True, masked_layout_mode="balanced", ), dict( column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=True, fuse_silu_and_mul=True, masked_layout_mode="imbalanced", ), dict( column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=True, fuse_silu_and_mul=True, masked_layout_mode="extreme", ), ] FLAGS_RANGE_MOE = get_benchmark_range(full_range=FLAGS_MOE, ci_range=FLAGS_MOE) CONFIGS_MOE = list( itertools.product( NUM_TOKENS_RANGE_MOE, HIDDEN_DIM_RANGE_MOE, GROUP_SIZE_RANGE, NUM_RANKS_RANGE_MOE, DST_DTYPE_RANGE, FLAGS_RANGE_MOE, ) ) # ---- Final configs ---- CONFIGS = CONFIGS_GEMM + CONFIGS_MOE LINE_VALS = ["triton", "sglang"] LINE_NAMES = ["Triton (Inaccurate)", "SGL Kernel"] STYLES = [("blue", "-"), ("green", "-")] def _flatten_to_2d(t: torch.Tensor) -> torch.Tensor: """Reshape a tensor with 3+ dims to 2D by merging all leading dims.""" if t.ndim <= 2: return t return t.reshape(-1, t.shape[-1]) def _make_sglang_bench_fn( x: torch.Tensor, group_size: int, dst_dtype: torch.dtype, flags: dict, ): """ Adapter that pre-allocates output tensors and returns a zero-arg callable matching the JIT kernel's signature. The JIT kernel does not support fuse_silu_and_mul, so when enabled we pre-compute silu+mul on the input. bench_kineto only times the kernel matching the given name, so the pre-processing is not included. The JIT kernel expects 2D tensors, so any higher-dimensional inputs (e.g. from masked_layout_mode) are flattened to 2D. """ fuse_silu_and_mul = flags.get("fuse_silu_and_mul", False) column_major_scales = flags.get("column_major_scales", False) scale_tma_aligned = flags.get("scale_tma_aligned", False) scale_ue8m0 = flags.get("scale_ue8m0", False) # JIT kernel does not support fuse_silu_and_mul; pre-compute it if fuse_silu_and_mul: half = x.shape[-1] // 2 x_input = torch.nn.functional.silu(x[..., :half]) * x[..., half:] else: x_input = x # JIT kernel expects 2D (num_tokens, hidden_dim); flatten if needed x_input = _flatten_to_2d(x_input.contiguous()) out_shape = x_input.shape output_q = torch.empty(out_shape, device=x.device, dtype=dst_dtype) fp8_max = torch.finfo(dst_dtype).max fp8_min = -fp8_max output_s = create_per_token_group_quant_fp8_output_scale( x_shape=out_shape, device=x.device, group_size=group_size, column_major_scales=column_major_scales, scale_tma_aligned=scale_tma_aligned, scale_ue8m0=scale_ue8m0, ) def _run(): sglang_per_token_group_quant_8bit( input=x_input, output_q=output_q, output_s=output_s, group_size=group_size, eps=1e-10, fp8_min=fp8_min, fp8_max=fp8_max, scale_ue8m0=scale_ue8m0, ) return _run @triton.testing.perf_report( triton.testing.Benchmark( x_names=[ "num_tokens", "hidden_dim", "group_size", "num_ranks", "dst_dtype", "flags", ], x_vals=CONFIGS, line_arg="provider", line_vals=LINE_VALS, # Triton has multi kernels and we only report the time for the core one line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="per-token-group-quant-8bit-performance", args={}, ) ) def benchmark( num_tokens, hidden_dim, group_size, num_ranks, dst_dtype, flags, provider ): print( f"Testing: {num_tokens=} {hidden_dim=} {group_size=} {num_ranks=} {dst_dtype=} {flags=} {provider=}" ) x, masked_m = create_per_token_group_quant_test_data( num_tokens=num_tokens, hidden_dim=hidden_dim, num_ranks=num_ranks, flags=flags ) if provider == "triton": fn = triton_per_token_group_quant_8bit kernel_names = "_per_token_group_quant_8bit|_silu_and_mul_post_quant_kernel" bench_fn = lambda: fn( x=x, masked_m=masked_m, group_size=group_size, dst_dtype=dst_dtype, **{k: v for k, v in flags.items() if k not in ["masked_layout_mode"]}, ) elif provider == "sglang": kernel_names = "per_token_group_quant_8bit_kernel" bench_fn = _make_sglang_bench_fn( x=x, group_size=group_size, dst_dtype=dst_dtype, flags=flags, ) else: raise ValueError(f"Unknown provider: {provider}") time_s = bench_kineto(bench_fn, kernel_names=kernel_names, num_tests=NUM_TESTS) return time_s * 1e6 if __name__ == "__main__": benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_qknorm.py ================================================ import itertools import torch import triton import triton.testing from sgl_kernel import rmsnorm from sglang.jit_kernel.benchmark.utils import ( DEFAULT_DEVICE, DEFAULT_DTYPE, get_benchmark_range, run_benchmark, ) from sglang.jit_kernel.norm import fused_inplace_qknorm from sglang.srt.utils import get_current_device_stream_fast alt_stream = torch.cuda.Stream() def sglang_aot_qknorm( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, ) -> None: head_dim = q.shape[-1] q = q.view(-1, head_dim) k = k.view(-1, head_dim) current_stream = get_current_device_stream_fast() alt_stream.wait_stream(current_stream) rmsnorm(q, q_weight, out=q) with torch.cuda.stream(alt_stream): rmsnorm(k, k_weight, out=k) current_stream.wait_stream(alt_stream) def sglang_jit_qknorm( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, ) -> None: fused_inplace_qknorm(q, k, q_weight, k_weight) def flashinfer_qknorm( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, ) -> None: from flashinfer import rmsnorm rmsnorm(q, q_weight, out=q) rmsnorm(k, k_weight, out=k) @torch.compile() def torch_impl_qknorm( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, eps: float = 1e-6, ) -> None: q_mean = q.float().pow(2).mean(dim=-1, keepdim=True) k_mean = k.float().pow(2).mean(dim=-1, keepdim=True) q_norm = (q_mean + eps).rsqrt() k_norm = (k_mean + eps).rsqrt() q.copy_(q.float() * q_norm * q_weight.float()) k.copy_(k.float() * k_norm * k_weight.float()) BS_RANGE = get_benchmark_range( full_range=[2**n for n in range(0, 14)], ci_range=[16], ) GQA_RANGE = get_benchmark_range( full_range=[4, 8], ci_range=[4], ) KV_HEAD_RANGE = get_benchmark_range( full_range=[1, 2, 4, 8], ci_range=[1], ) HEAD_DIM_RANGE = get_benchmark_range( full_range=[128, 256, 512, 1024], ci_range=[128], ) LINE_VALS = ["aot", "jit", "flashinfer", "torch"] LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "FlashInfer", "PyTorch"] STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")] configs = list(itertools.product(HEAD_DIM_RANGE, GQA_RANGE, KV_HEAD_RANGE, BS_RANGE)) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["head_dim", "GQA", "num_kv_heads", "batch_size"], x_vals=configs, line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="qknorm-performance", args={}, ) ) def benchmark( head_dim: int, GQA: int, num_kv_heads: int, batch_size: int, provider: str ): num_qo_heads = GQA * num_kv_heads q = torch.randn( (batch_size, num_qo_heads, head_dim), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE ) k = torch.randn( (batch_size, num_kv_heads, head_dim), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE ) q_weight = torch.randn(head_dim, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) k_weight = torch.randn(head_dim, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) FN_MAP = { "aot": sglang_aot_qknorm, "jit": sglang_jit_qknorm, "flashinfer": flashinfer_qknorm, "torch": torch_impl_qknorm, } fn = lambda: FN_MAP[provider](q, k, q_weight, k_weight) return run_benchmark(fn) if __name__ == "__main__": benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py ================================================ import itertools from typing import Tuple import torch import triton import triton.testing from sgl_kernel import rmsnorm from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads from sglang.srt.utils import get_current_device_stream_fast IS_CI = is_in_ci() alt_stream = torch.cuda.Stream() def sglang_jit_qknorm_across_heads( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, ) -> None: fused_inplace_qknorm_across_heads(q, k, q_weight, k_weight) def sglang_aot_qknorm_across_heads( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, ) -> None: current_stream = get_current_device_stream_fast() alt_stream.wait_stream(current_stream) rmsnorm(q, q_weight, out=q) with torch.cuda.stream(alt_stream): rmsnorm(k, k_weight, out=k) current_stream.wait_stream(alt_stream) def flashinfer_qknorm_across_heads( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, ) -> None: from flashinfer import rmsnorm rmsnorm(q, q_weight, out=q) rmsnorm(k, k_weight, out=k) @torch.compile() def torch_impl_qknorm_across_heads( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, eps: float = 1e-6, ) -> None: q_mean = q.float().pow(2).mean(dim=-1, keepdim=True) k_mean = k.float().pow(2).mean(dim=-1, keepdim=True) q_norm = (q_mean + eps).rsqrt() k_norm = (k_mean + eps).rsqrt() q.copy_(q.float() * q_norm * q_weight.float()) k.copy_(k.float() * k_norm * k_weight.float()) DTYPE = torch.bfloat16 DEVICE = "cuda" if IS_CI: BS_RANGE = [16] HIDDEN_DIM_RANGE = [1024] else: BS_RANGE = [2**n for n in range(0, 14)] HIDDEN_DIM_RANGE = [512, 1024, 2048, 4096, 8192] LINE_VALS = ["jit", "aot", "flashinfer", "torch"] LINE_NAMES = ["SGL JIT Kernel", "SGL AOT Kernel", "FlashInfer", "PyTorch"] STYLES = [("blue", "-"), ("orange", "--"), ("green", "-."), ("red", ":")] configs = list(itertools.product(BS_RANGE, HIDDEN_DIM_RANGE)) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["batch_size", "hidden_dim"], x_vals=configs, line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="qknorm-across-heads-performance", args={}, ) ) def benchmark( batch_size: int, hidden_dim: int, provider: str ) -> Tuple[float, float, float]: q = torch.randn((batch_size, hidden_dim), dtype=DTYPE, device=DEVICE) k = torch.randn((batch_size, hidden_dim), dtype=DTYPE, device=DEVICE) q_weight = torch.randn(hidden_dim, dtype=DTYPE, device=DEVICE) k_weight = torch.randn(hidden_dim, dtype=DTYPE, device=DEVICE) FN_MAP = { "jit": sglang_jit_qknorm_across_heads, "aot": sglang_aot_qknorm_across_heads, "flashinfer": flashinfer_qknorm_across_heads, "torch": torch_impl_qknorm_across_heads, } fn = lambda: FN_MAP[provider](q, k, q_weight, k_weight) return run_benchmark(fn) if __name__ == "__main__": benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py ================================================ from typing import Tuple import torch import triton.testing from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark_no_cudagraph from sglang.jit_kernel.diffusion.triton.norm import norm_infer from sglang.jit_kernel.diffusion.triton.scale_shift import ( fuse_layernorm_scale_shift_gate_select01_kernel, fuse_residual_layernorm_scale_shift_gate_select01_kernel, fuse_scale_shift_gate_select01_kernel, ) if is_in_ci(): B_RANGE, S_RANGE, D_RANGE = [1], [128], [3072] else: B_RANGE, S_RANGE, D_RANGE = [1, 2], [128, 512, 2048], [1024, 1536, 3072] DTYPE = torch.bfloat16 DEVICE = "cuda" EPS = 1e-6 LINE_VALS = ["split", "fused"] LINE_NAMES = ["Split Kernels", "Fused Triton"] STYLES = [("red", "-"), ("blue", "--")] CONFIG = [(b, s, d) for b in B_RANGE for s in S_RANGE for d in D_RANGE] def _make_common_inputs(batch_size: int, seq_len: int, hidden_size: int): x = torch.randn(batch_size, seq_len, hidden_size, dtype=DTYPE, device=DEVICE) weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE) bias = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE) index = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.int32, device=DEVICE) scale0 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE) shift0 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE) gate0 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE) scale1 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE) shift1 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE) gate1 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE) return x, weight, bias, index, scale0, shift0, gate0, scale1, shift1, gate1 @triton.testing.perf_report( triton.testing.Benchmark( x_names=["B", "S", "D"], x_vals=CONFIG, line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="qwen_image_layernorm_scale_shift_gate_select01", args={}, ) ) def bench_layernorm_scale_shift_gate_select01( B: int, S: int, D: int, provider: str ) -> Tuple[float, float, float]: x, weight, bias, index, scale0, shift0, gate0, scale1, shift1, gate1 = ( _make_common_inputs(B, S, D) ) if provider == "split": def fn(): normalized = norm_infer( x.view(-1, x.shape[-1]), weight, bias, eps=EPS, is_rms_norm=False, ).view_as(x) return fuse_scale_shift_gate_select01_kernel( normalized, scale0=scale0, shift0=shift0, gate0=gate0, scale1=scale1, shift1=shift1, gate1=gate1, index=index, ) else: def fn(): return fuse_layernorm_scale_shift_gate_select01_kernel( x, weight=weight, bias=bias, scale0=scale0, shift0=shift0, gate0=gate0, scale1=scale1, shift1=shift1, gate1=gate1, index=index, eps=EPS, ) return run_benchmark_no_cudagraph(fn) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["B", "S", "D"], x_vals=CONFIG, line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="qwen_image_residual_layernorm_scale_shift_gate_select01", args={}, ) ) def bench_residual_layernorm_scale_shift_gate_select01( B: int, S: int, D: int, provider: str ) -> Tuple[float, float, float]: x, weight, bias, index, scale0, shift0, gate0, scale1, shift1, gate1 = ( _make_common_inputs(B, S, D) ) residual = torch.randn_like(x) residual_gate = torch.randn_like(x) if provider == "split": def fn(): residual_out = residual + residual_gate * x normalized = norm_infer( residual_out.view(-1, residual_out.shape[-1]), weight, bias, eps=EPS, is_rms_norm=False, ).view_as(residual_out) return fuse_scale_shift_gate_select01_kernel( normalized, scale0=scale0, shift0=shift0, gate0=gate0, scale1=scale1, shift1=shift1, gate1=gate1, index=index, ) else: def fn(): return fuse_residual_layernorm_scale_shift_gate_select01_kernel( x, residual=residual, residual_gate=residual_gate, weight=weight, bias=bias, scale0=scale0, shift0=shift0, gate0=gate0, scale1=scale1, shift1=shift1, gate1=gate1, index=index, eps=EPS, ) return run_benchmark_no_cudagraph(fn) if __name__ == "__main__": print(f"\n{'=' * 80}") print("Benchmark: qwen_image layernorm + scale_shift_gate_select01") print(f"{'=' * 80}\n") bench_layernorm_scale_shift_gate_select01.run(print_data=True) print(f"\n{'=' * 80}") print("Benchmark: qwen_image residual + layernorm + scale_shift_gate_select01") print(f"{'=' * 80}\n") bench_residual_layernorm_scale_shift_gate_select01.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_renorm.py ================================================ import itertools import sgl_kernel import torch import triton import triton.testing from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark_no_cudagraph def torch_top_k_renorm_probs(probs, top_k): """Vectorized PyTorch implementation of top-k renormalization.""" batch_size, vocab_size = probs.shape # Handle scalar or tensor k if isinstance(top_k, int): k_val = min(max(top_k, 1), vocab_size) # Get top-k indices for all batches at once _, topk_indices = torch.topk(probs, k_val, dim=1, largest=True) # Create mask: batch_size x vocab_size mask = torch.zeros_like(probs) mask.scatter_(1, topk_indices, 1.0) # Vectorized renormalization masked_probs = probs * mask renorm_probs = masked_probs / (masked_probs.sum(dim=1, keepdim=True) + 1e-10) return renorm_probs else: # Variable k per batch - need to handle separately renorm_probs = torch.zeros_like(probs) for i in range(batch_size): k_val = min(max(top_k[i].item(), 1), vocab_size) _, topk_indices = torch.topk(probs[i], k_val, largest=True) mask = torch.zeros_like(probs[i]) mask[topk_indices] = 1.0 masked_probs = probs[i] * mask renorm_probs[i] = masked_probs / (masked_probs.sum() + 1e-10) return renorm_probs def torch_top_p_renorm_probs(probs, top_p, eps=1e-5): """Vectorized PyTorch implementation of top-p renormalization.""" batch_size, vocab_size = probs.shape # Handle scalar or tensor p if isinstance(top_p, float): p_val = top_p # Vectorized implementation for uniform top_p # Sort probs in descending order sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=1) cumsum_probs = torch.cumsum(sorted_probs, dim=1) # Find cutoff: where cumsum exceeds top_p cutoff_mask = cumsum_probs <= p_val # Keep at least one token (the highest prob) cutoff_mask[:, 0] = True # Create mask in original order mask = torch.zeros_like(probs) mask.scatter_(1, sorted_indices, cutoff_mask.float()) # Vectorized renormalization masked_probs = probs * mask renorm_probs = masked_probs / (masked_probs.sum(dim=1, keepdim=True) + eps) return renorm_probs else: # Variable p per batch - need to handle separately renorm_probs = torch.zeros_like(probs) for i in range(batch_size): p_val = top_p[i].item() sorted_prob, indices = torch.sort(probs[i], descending=False) cdf = torch.cumsum(sorted_prob, dim=-1) mask = torch.zeros(vocab_size, dtype=torch.float32, device=probs.device) mask.scatter_(0, indices, (cdf >= (1 - p_val) - eps).float()) masked_probs = probs[i] * mask renorm_probs[i] = masked_probs / (masked_probs.sum() + eps) return renorm_probs def torch_top_k_mask_logits(logits, top_k): """Vectorized PyTorch implementation of top-k logits masking.""" batch_size, vocab_size = logits.shape # Handle scalar or tensor k if isinstance(top_k, int): k_val = min(max(top_k, 1), vocab_size) # Get top-k indices for all batches at once _, topk_indices = torch.topk(logits, k_val, dim=1, largest=True) # Create masked logits: start with -inf everywhere masked_logits = torch.full_like(logits, float("-inf")) # Scatter the top-k values back masked_logits.scatter_(1, topk_indices, logits.gather(1, topk_indices)) else: # Variable k per batch - need to handle separately masked_logits = torch.full_like(logits, float("-inf")) for i in range(batch_size): k_val = min(max(top_k[i].item(), 1), vocab_size) _, topk_indices = torch.topk(logits[i], k_val, largest=True) masked_logits[i, topk_indices] = logits[i, topk_indices] return masked_logits def calculate_diff_top_k_renorm(batch_size, vocab_size, k): """Compare Torch reference and SGLang kernel for top-k renorm correctness.""" torch.manual_seed(42) device = torch.device("cuda") pre_norm_prob = torch.rand(batch_size, vocab_size, device=device) probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True) top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32) torch_output = torch_top_k_renorm_probs(probs, top_k_tensor) sglang_output = sgl_kernel.top_k_renorm_prob(probs, top_k_tensor) torch.testing.assert_close(torch_output, sglang_output, rtol=1e-3, atol=1e-3) def calculate_diff_top_p_renorm(batch_size, vocab_size, p): """Compare Torch reference and SGLang kernel for top-p renorm correctness.""" torch.manual_seed(42) device = torch.device("cuda") pre_norm_prob = torch.rand(batch_size, vocab_size, device=device) probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True) top_p_tensor = torch.full((batch_size,), p, device=device, dtype=torch.float32) torch_output = torch_top_p_renorm_probs(probs, top_p_tensor) sglang_output = sgl_kernel.top_p_renorm_prob(probs, top_p_tensor) torch.testing.assert_close(torch_output, sglang_output, rtol=1e-3, atol=1e-3) def calculate_diff_top_k_mask(batch_size, vocab_size, k): """Compare Torch reference and SGLang kernel for top-k mask correctness.""" torch.manual_seed(42) device = torch.device("cuda") logits = torch.randn(batch_size, vocab_size, device=device) * 5 top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32) torch_output = torch_top_k_mask_logits(logits, top_k_tensor) sglang_output = sgl_kernel.top_k_mask_logits(logits, top_k_tensor) torch.testing.assert_close(torch_output, sglang_output, rtol=1e-3, atol=1e-3) # Parameter space - simplified for CI if is_in_ci(): batch_size_range = [16] vocab_size_range = [111] k_range = [10] p_range = [0.5] else: batch_size_range = [16, 64, 128] vocab_size_range = [111, 32000, 128256] k_range = [10, 100, 500] p_range = [0.1, 0.5, 0.9] configs_k = list(itertools.product(batch_size_range, vocab_size_range, k_range)) configs_p = list(itertools.product(batch_size_range, vocab_size_range, p_range)) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["batch_size", "vocab_size", "k"], x_vals=configs_k, line_arg="provider", line_vals=["torch", "sglang"], line_names=["Torch Reference", "SGL Kernel"], styles=[("red", "-"), ("green", "-")], ylabel="us", plot_name="top-k-renorm-probs-performance", args={}, ) ) def benchmark_top_k_renorm(batch_size, vocab_size, k, provider): # Skip invalid configurations if k >= vocab_size: return float("nan"), float("nan"), float("nan") torch.manual_seed(42) device = torch.device("cuda") pre_norm_prob = torch.rand(batch_size, vocab_size, device=device) probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True) top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32) if provider == "torch": fn = lambda: torch_top_k_renorm_probs(probs.clone(), top_k_tensor) elif provider == "sglang": fn = lambda: sgl_kernel.top_k_renorm_prob(probs.clone(), top_k_tensor) return run_benchmark_no_cudagraph(fn) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["batch_size", "vocab_size", "p"], x_vals=configs_p, line_arg="provider", line_vals=["torch", "sglang"], line_names=["Torch Reference", "SGL Kernel"], styles=[("red", "-"), ("blue", "-")], ylabel="us", plot_name="top-p-renorm-probs-performance", args={}, ) ) def benchmark_top_p_renorm(batch_size, vocab_size, p, provider): torch.manual_seed(42) device = torch.device("cuda") pre_norm_prob = torch.rand(batch_size, vocab_size, device=device) probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True) top_p_tensor = torch.full((batch_size,), p, device=device, dtype=torch.float32) if provider == "torch": fn = lambda: torch_top_p_renorm_probs(probs.clone(), top_p_tensor) elif provider == "sglang": fn = lambda: sgl_kernel.top_p_renorm_prob(probs.clone(), top_p_tensor) return run_benchmark_no_cudagraph(fn) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["batch_size", "vocab_size", "k"], x_vals=configs_k, line_arg="provider", line_vals=["torch", "sglang"], line_names=["Torch Reference", "SGL Kernel"], styles=[("red", "-"), ("orange", "-")], ylabel="us", plot_name="top-k-mask-logits-performance", args={}, ) ) def benchmark_top_k_mask(batch_size, vocab_size, k, provider): # Skip invalid configurations if k >= vocab_size: return float("nan"), float("nan"), float("nan") torch.manual_seed(42) device = torch.device("cuda") logits = torch.randn(batch_size, vocab_size, device=device) * 5 top_k_tensor = torch.full((batch_size,), k, device=device, dtype=torch.int32) if provider == "torch": fn = lambda: torch_top_k_mask_logits(logits.clone(), top_k_tensor) elif provider == "sglang": fn = lambda: sgl_kernel.top_k_mask_logits(logits.clone(), top_k_tensor) return run_benchmark_no_cudagraph(fn) if __name__ == "__main__": print("=" * 60) print("Running correctness checks...") print("=" * 60) # Correctness checks - simplified for CI if is_in_ci(): test_configs_k = [configs_k[0]] if configs_k else [(16, 111, 10)] test_configs_p = [configs_p[0]] if configs_p else [(16, 111, 0.5)] else: test_configs_k = configs_k[:3] # Test first 3 configs test_configs_p = configs_p[:3] print("\n1. Testing top_k_renorm_probs...") for cfg in test_configs_k: batch_size, vocab_size, k = cfg if k < vocab_size: # Skip invalid configs calculate_diff_top_k_renorm(batch_size, vocab_size, k) print( f" ✓ Passed: batch_size={batch_size}, vocab_size={vocab_size}, k={k}" ) print("\n2. Testing top_p_renorm_probs...") for cfg in test_configs_p: calculate_diff_top_p_renorm(*cfg) batch_size, vocab_size, p = cfg print(f" ✓ Passed: batch_size={batch_size}, vocab_size={vocab_size}, p={p}") print("\n3. Testing top_k_mask_logits...") for cfg in test_configs_k: batch_size, vocab_size, k = cfg if k < vocab_size: # Skip invalid configs calculate_diff_top_k_mask(batch_size, vocab_size, k) print( f" ✓ Passed: batch_size={batch_size}, vocab_size={vocab_size}, k={k}" ) print("\n" + "=" * 60) print("All correctness checks passed!") print("=" * 60) print("\n" + "=" * 60) print("Starting performance benchmarks...") print("=" * 60) print("\n1. Benchmarking top_k_renorm_probs...") benchmark_top_k_renorm.run(print_data=True) print("\n2. Benchmarking top_p_renorm_probs...") benchmark_top_p_renorm.run(print_data=True) print("\n3. Benchmarking top_k_mask_logits...") benchmark_top_k_mask.run(print_data=True) print("\n" + "=" * 60) print("Benchmarking complete!") print("=" * 60) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_rmsnorm.py ================================================ import itertools import torch import triton import triton.testing from flashinfer import rmsnorm as fi_rmsnorm from sgl_kernel import rmsnorm from sglang.jit_kernel.benchmark.utils import ( DEFAULT_DEVICE, DEFAULT_DTYPE, get_benchmark_range, run_benchmark, ) from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm def sglang_aot_rmsnorm( input: torch.Tensor, weight: torch.Tensor, ) -> None: rmsnorm(input, weight, out=input) def sglang_jit_rmsnorm( input: torch.Tensor, weight: torch.Tensor, ) -> None: jit_rmsnorm(input, weight, output=input) def flashinfer_rmsnorm( input: torch.Tensor, weight: torch.Tensor, ) -> None: fi_rmsnorm(input, weight, out=input) @torch.compile() def torch_impl_rmsnorm( input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6, ) -> None: mean = input.float().pow(2).mean(dim=-1, keepdim=True) norm = (mean + eps).rsqrt() input.copy_(input.float() * norm * weight.float()) BS_LIST = get_benchmark_range( full_range=[2**n for n in range(0, 14)], ci_range=[16], ) HIDDEN_SIZE_LIST = get_benchmark_range( full_range=[1536, 3072, 4096, 5120, 8192], ci_range=[512, 2048], ) LINE_VALS = ["aot", "jit", "flashinfer", "torch"] LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "FlashInfer", "PyTorch"] STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")] configs = list(itertools.product(HIDDEN_SIZE_LIST, BS_LIST)) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["hidden_size", "batch_size"], x_vals=configs, line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="rmsnorm-performance", args={}, ) ) def benchmark(hidden_size: int, batch_size: int, provider: str): input = torch.randn( (batch_size, hidden_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE ) weight = torch.randn(hidden_size, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) FN_MAP = { "aot": sglang_aot_rmsnorm, "jit": sglang_jit_rmsnorm, "flashinfer": flashinfer_rmsnorm, "torch": torch_impl_rmsnorm, } fn = lambda: FN_MAP[provider](input.clone(), weight) return run_benchmark(fn) if __name__ == "__main__": benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_rope.py ================================================ import itertools import torch import triton import triton.testing from sglang.jit_kernel.benchmark.utils import ( DEFAULT_DEVICE, DEFAULT_DTYPE, get_benchmark_range, run_benchmark, ) MAX_SEQ_LEN = 131072 ROPE_BASE = 10000.0 ROPE_DIM = 128 CACHE_SIZE = 1024 * 1024 def create_cos_sin_cache( rotary_dim: int = ROPE_DIM, max_position: int = MAX_SEQ_LEN, base: float = ROPE_BASE, ) -> torch.Tensor: inv_freq = 1.0 / ( base ** ( torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=DEFAULT_DEVICE) / rotary_dim ) ) t = torch.arange(max_position, dtype=torch.float32, device=DEFAULT_DEVICE) freqs = torch.einsum("i,j->ij", t, inv_freq) cos = freqs.cos() sin = freqs.sin() return torch.cat((cos, sin), dim=-1) # Pre-build the cache once COS_SIN_CACHE = create_cos_sin_cache() # --------------------------------------------------------------------------- # RoPE-only provider implementations # --------------------------------------------------------------------------- def flashinfer_rope( q: torch.Tensor, k: torch.Tensor, positions: torch.Tensor, is_neox: bool, ) -> None: from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace head_size = q.shape[-1] apply_rope_with_cos_sin_cache_inplace( positions=positions, query=q.view(q.shape[0], -1), key=k.view(k.shape[0], -1), head_size=head_size, cos_sin_cache=COS_SIN_CACHE, is_neox=is_neox, ) def sglang_pos_enc_rope( q: torch.Tensor, k: torch.Tensor, positions: torch.Tensor, is_neox: bool, ) -> None: from sglang.jit_kernel.rope import rotary_embedding_with_key head_size = q.shape[-1] rotary_embedding_with_key( positions=positions, query=q.view(q.shape[0], -1), key=k.view(k.shape[0], -1), head_size=head_size, cos_sin_cache=COS_SIN_CACHE, is_neox=is_neox, ) def sglang_fused_rope( q: torch.Tensor, k: torch.Tensor, positions: torch.Tensor, is_neox: bool, ) -> None: from sglang.jit_kernel.rope import apply_rope_inplace apply_rope_inplace(q, k, COS_SIN_CACHE, positions, is_neox=is_neox) # --------------------------------------------------------------------------- # RoPE + KV cache store provider implementations # --------------------------------------------------------------------------- def jit_rope_then_store( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, positions: torch.Tensor, out_loc: torch.Tensor, is_neox: bool, ) -> None: from sglang.jit_kernel.kvcache import store_cache from sglang.jit_kernel.rope import apply_rope_inplace head_size = q.shape[-1] row_dim = k.shape[-2] * head_size apply_rope_inplace( positions=positions, q=q, k=k, rope_dim=head_size, cos_sin_cache=COS_SIN_CACHE, is_neox=is_neox, ) store_cache( k.view(-1, row_dim), v.view(-1, row_dim), k_cache, v_cache, out_loc, ) def jit_fused_rope_store( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, positions: torch.Tensor, out_loc: torch.Tensor, is_neox: bool, ) -> None: from sglang.jit_kernel.rope import apply_rope_inplace_with_kvcache apply_rope_inplace_with_kvcache( q, k, v, k_cache, v_cache, COS_SIN_CACHE, positions, out_loc, is_neox=is_neox ) # --------------------------------------------------------------------------- # Benchmark configuration (shared) # --------------------------------------------------------------------------- BS_RANGE = get_benchmark_range( full_range=[2**n for n in range(0, 16)], ci_range=[16], ) QK_HEAD_RANGE = get_benchmark_range( full_range=[(8, 1), (16, 2), (32, 8)], ci_range=[(16, 2)], ) QK_HEAD_RANGE = [f"{q},{k}" for q, k in QK_HEAD_RANGE] IS_NEOX_RANGE = get_benchmark_range( full_range=[True, False], ci_range=[True], ) # --------------------------------------------------------------------------- # Benchmark 1: RoPE only # --------------------------------------------------------------------------- ROPE_LINE_VALS = ["flashinfer", "jit_pos_enc", "jit_fused_rope"] ROPE_LINE_NAMES = [ "FlashInfer", "SGL JIT PosEnc", "SGL JIT Fused RoPE", ] ROPE_STYLES = [("green", "-."), ("red", "-"), ("blue", "--")] rope_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE)) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["num_q_k_heads", "is_neox", "batch_size"], x_vals=rope_configs, line_arg="provider", line_vals=ROPE_LINE_VALS, line_names=ROPE_LINE_NAMES, styles=ROPE_STYLES, ylabel="us", plot_name="rope-performance", args={}, ) ) def benchmark(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str): qo, kv = num_q_k_heads.split(",") num_qo_heads = int(qo) num_kv_heads = int(kv) q = torch.randn( (batch_size, num_qo_heads, ROPE_DIM), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE, ) k = torch.randn( (batch_size, num_kv_heads, ROPE_DIM), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE, ) seed = batch_size << 16 | num_qo_heads << 8 | num_kv_heads << 4 | is_neox torch.random.manual_seed(seed) positions = torch.randint( MAX_SEQ_LEN, (batch_size,), device=DEFAULT_DEVICE, dtype=torch.int64 ) torch.cuda.synchronize() FN_MAP = { "flashinfer": flashinfer_rope, "jit_pos_enc": sglang_pos_enc_rope, "jit_fused_rope": sglang_fused_rope, } fn = lambda: FN_MAP[provider](q, k, positions, is_neox) return run_benchmark(fn) # --------------------------------------------------------------------------- # Benchmark 2: RoPE + KV cache store # --------------------------------------------------------------------------- STORE_LINE_VALS = ["jit_rope_then_store", "jit_fused_store"] STORE_LINE_NAMES = [ "SGL JIT RoPE + Store", "SGL JIT Fused RoPE + Store", ] STORE_STYLES = [("red", "-"), ("blue", "--")] store_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE)) @triton.testing.perf_report( triton.testing.Benchmark( x_names=["num_q_k_heads", "is_neox", "batch_size"], x_vals=store_configs, line_arg="provider", line_vals=STORE_LINE_VALS, line_names=STORE_LINE_NAMES, styles=STORE_STYLES, ylabel="us", plot_name="rope-store-performance", args={}, ) ) def benchmark_store(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str): qo, kv = num_q_k_heads.split(",") num_qo_heads = int(qo) num_kv_heads = int(kv) q = torch.randn( (batch_size, num_qo_heads, ROPE_DIM), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE, ) k = torch.randn( (batch_size, num_kv_heads, ROPE_DIM), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE, ) v = torch.randn( (batch_size, num_kv_heads, ROPE_DIM), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE, ) row_size = num_kv_heads * ROPE_DIM k_cache = torch.zeros( CACHE_SIZE, row_size, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE ) v_cache = torch.zeros( CACHE_SIZE, row_size, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE ) out_loc = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE, dtype=torch.int64)[ :batch_size ] seed = batch_size << 16 | num_qo_heads << 8 | num_kv_heads << 4 | is_neox torch.random.manual_seed(seed) positions = torch.randint( MAX_SEQ_LEN, (batch_size,), device=DEFAULT_DEVICE, dtype=torch.int64 ) torch.cuda.synchronize() FN_MAP = { "jit_rope_then_store": jit_rope_then_store, "jit_fused_store": jit_fused_rope_store, } fn = lambda: FN_MAP[provider]( q, k, v, k_cache, v_cache, positions, out_loc, is_neox ) return run_benchmark(fn) if __name__ == "__main__": print("Running RoPE performance benchmark...") benchmark.run(print_data=True) print("\nRunning RoPE + KV cache store performance benchmark...") benchmark_store.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/bench_store_cache.py ================================================ import itertools from typing import Tuple import torch import triton import triton.testing from sglang.jit_kernel.benchmark.utils import ( DEFAULT_DEVICE, DEFAULT_DTYPE, DEFAULT_QUANTILES, get_benchmark_range, ) from sglang.jit_kernel.kvcache import store_cache def sglang_jit_store_cache( k: torch.Tensor, v: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, indices: torch.Tensor, ) -> None: store_cache(k, v, k_cache, v_cache, indices) @torch.compile() def torch_compile_store_cache( k: torch.Tensor, v: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, indices: torch.Tensor, ) -> None: k_cache[indices] = k v_cache[indices] = v alt_stream = torch.cuda.Stream() def torch_streams_store_cache( k: torch.Tensor, v: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, indices: torch.Tensor, ) -> None: current_stream = torch.cuda.current_stream() alt_stream.wait_stream(current_stream) k_cache[indices] = k with torch.cuda.stream(alt_stream): v_cache[indices] = v current_stream.wait_stream(alt_stream) NUM_LAYERS = 8 CACHE_SIZE = 2 * 1024 * 1024 // NUM_LAYERS BS_RANGE = get_benchmark_range( full_range=[2**n for n in range(0, 15)], ci_range=[16], ) ITEM_SIZE = get_benchmark_range( full_range=[64, 128, 256, 512, 1024], ci_range=[1024], ) LINE_VALS = ["jit", "torch_compile", "torch_streams"] LINE_NAMES = ["SGL JIT Kernel", "PyTorch Compile", "PyTorch 2 Stream"] STYLES = [("blue", "--"), ("red", ":"), ("green", "-.")] X_NAMES = ["item_size", "batch_size"] CONFIGS = list(itertools.product(ITEM_SIZE, BS_RANGE)) @triton.testing.perf_report( triton.testing.Benchmark( x_names=X_NAMES, x_vals=CONFIGS, line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, styles=STYLES, ylabel="us", plot_name="store-kvcache-performance", args={}, ) ) def benchmark( batch_size: int, item_size: int, provider: str ) -> Tuple[float, float, float]: k = torch.randn( (NUM_LAYERS, batch_size, item_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE ) v = torch.randn( (NUM_LAYERS, batch_size, item_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE ) k_cache = torch.randn( (NUM_LAYERS, CACHE_SIZE, item_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE ) v_cache = torch.randn( (NUM_LAYERS, CACHE_SIZE, item_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE ) indices = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE)[:batch_size] torch.cuda.synchronize() FN_MAP = { "jit": sglang_jit_store_cache, "torch_compile": torch_compile_store_cache, "torch_streams": torch_streams_store_cache, } def fn(): impl = FN_MAP[provider] for i in range(NUM_LAYERS): impl(k[i], v[i], k_cache[i], v_cache[i], indices) # Custom time calculation: divide by NUM_LAYERS ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( fn, quantiles=DEFAULT_QUANTILES ) return ( 1000 * ms / NUM_LAYERS, 1000 * max_ms / NUM_LAYERS, 1000 * min_ms / NUM_LAYERS, ) if __name__ == "__main__": benchmark.run(print_data=True) ================================================ FILE: python/sglang/jit_kernel/benchmark/utils.py ================================================ """Common utilities for jit_kernel benchmark files.""" from typing import Callable, List, Tuple import torch import triton.testing from sglang.jit_kernel.utils import is_in_ci as jit_kernel_is_in_ci # Common constants DEFAULT_DTYPE = torch.bfloat16 DEFAULT_DEVICE = "cuda" DEFAULT_QUANTILES = [0.5, 0.2, 0.8] def is_in_ci() -> bool: """Check if running in CI environment.""" return jit_kernel_is_in_ci() def get_benchmark_range(full_range: List, ci_range: List) -> List: """Return appropriate benchmark range based on CI environment.""" return ci_range if is_in_ci() else full_range def run_benchmark( fn: Callable, quantiles: List[float] = None ) -> Tuple[float, float, float]: """Execute benchmark using CUDA graph and return times in microseconds. Args: fn: Function to benchmark quantiles: Quantiles for timing measurements [median, min, max] Returns: Tuple of (median_us, max_us, min_us) """ quantiles = quantiles or DEFAULT_QUANTILES ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles) return 1000 * ms, 1000 * max_ms, 1000 * min_ms def run_benchmark_no_cudagraph( fn: Callable, quantiles: List[float] = None ) -> Tuple[float, float, float]: quantiles = quantiles or DEFAULT_QUANTILES ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles) return 1000 * ms, 1000 * max_ms, 1000 * min_ms ================================================ FILE: python/sglang/jit_kernel/concat_mla.py ================================================ from __future__ import annotations from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import cache_once, load_jit if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_concat_mla_k_module() -> Module: return load_jit( "concat_mla_k", cuda_files=["elementwise/concat_mla.cuh"], cuda_wrappers=[("concat_mla_k", "ConcatMlaKKernel::run")], ) @cache_once def _jit_concat_mla_absorb_q_module() -> Module: return load_jit( "concat_mla_absorb_q", cuda_files=["elementwise/concat_mla.cuh"], cuda_wrappers=[("concat_mla_absorb_q", "ConcatMlaAbsorbQKernel::run")], ) def concat_mla_k(k: torch.Tensor, k_nope: torch.Tensor, k_rope: torch.Tensor) -> None: """ Concatenate k_nope and k_rope into k for MLA (Multi-head Latent Attention). This kernel efficiently broadcasts k_rope across all heads while copying k_nope values directly. Args: k: Output tensor of shape [num_tokens, num_heads=128, k_head_dim=192], dtype=bfloat16 k_nope: Input tensor of shape [num_tokens, num_heads=128, nope_head_dim=128], dtype=bfloat16 k_rope: Input tensor of shape [num_tokens, 1, rope_head_dim=64], dtype=bfloat16 """ module = _jit_concat_mla_k_module() module.concat_mla_k(k, k_nope, k_rope) def concat_mla_absorb_q(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: """ Concatenate tensors a and b for MLA absorbed Q computation. Args: a: Input tensor of shape [dim_0, dim_1, a_last_dim], dtype=bfloat16 b: Input tensor of shape [dim_0, dim_1, b_last_dim], dtype=bfloat16 Returns: Output tensor of shape [dim_0, dim_1, a_last_dim + b_last_dim], dtype=bfloat16 """ out = torch.empty( (*a.shape[:-1], a.shape[-1] + b.shape[-1]), dtype=a.dtype, device=a.device, ) module = _jit_concat_mla_absorb_q_module() module.concat_mla_absorb_q(a, b, out) return out ================================================ FILE: python/sglang/jit_kernel/csrc/add_constant.cuh ================================================ #include // For TensorMatcher, SymbolicSize, SymbolicDevice #include // For div_ceil, RuntimeCheck #include // For LaunchKernel #include #include #include #include namespace { template __global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t length) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < length) { dst[idx] = src[idx] + kConstant; } } constexpr size_t kBlockSize = 256; // You can also use struct with static method as an alternative template void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { using namespace host; // 1. Validate input tensors SymbolicSize N = {"num_elements"}; SymbolicDevice device_; TensorMatcher({N}) // 1D tensor, must be contiguous .with_dtype() // must be int32 .with_device(device_) // must be on CUDA device .verify(dst) // check tensor dst .verify(src); // check tensor src // 2. Extract required parameters, prepare for kernel launch const size_t num_elements = N.unwrap(); const size_t grid_size = div_ceil(num_elements, kBlockSize); const DLDevice device = device_.unwrap(); [[maybe_unused]] // optional, can be omitted const size_t dynamic_smem = 0; [[maybe_unused]] // optional, LaunchKernel can auto determine stream from device const cudaStream_t stream = LaunchKernel::resolve_device(device); // some extra runtime checks using host::RuntimeCheck RuntimeCheck(num_elements > 0, "We only support non-empty tensors, got num_elements = ", num_elements); // 3. Launch the kernel. Error code will be automatically checked. LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)( // kernel function add_constant_kernel, // kernel arguments static_cast(dst.data_ptr()), static_cast(src.data_ptr()), num_elements); } } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh ================================================ #include #include #include #include #include #include #include #include #include #include #include #include namespace { template __global__ void timestep_embedding_kernel( const TIn* __restrict__ t_ptr, float* __restrict__ output_ptr, int dim, float neg_log_max_period, float scale, int batch_size) { int row_idx = static_cast(blockIdx.x * blockDim.y + threadIdx.y); if (row_idx >= batch_size) { return; } float t_val = device::cast(t_ptr[row_idx]); float* output_batch_base_ptr = output_ptr + row_idx * dim; int half_dim = dim / 2; int thread_offset = static_cast(threadIdx.x); while (thread_offset * 4 < half_dim) { float4* top_half; float4* bottom_half; if constexpr (!kFlipSinToCos) { bottom_half = reinterpret_cast(output_batch_base_ptr + thread_offset * 4); top_half = reinterpret_cast(output_batch_base_ptr + half_dim + thread_offset * 4); } else { top_half = reinterpret_cast(output_batch_base_ptr + thread_offset * 4); bottom_half = reinterpret_cast(output_batch_base_ptr + half_dim + thread_offset * 4); } float4 vals; vals.x = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 0)); vals.y = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 1)); vals.z = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 2)); vals.w = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 3)); float4 cos_vals; cos_vals.x = device::math::cos(vals.x); cos_vals.y = device::math::cos(vals.y); cos_vals.z = device::math::cos(vals.z); cos_vals.w = device::math::cos(vals.w); *top_half = cos_vals; float4 sin_vals; sin_vals.x = device::math::sin(vals.x); sin_vals.y = device::math::sin(vals.y); sin_vals.z = device::math::sin(vals.z); sin_vals.w = device::math::sin(vals.w); *bottom_half = sin_vals; thread_offset += static_cast(blockDim.x); } } template inline void launch_timestep_embedding( const tvm::ffi::TensorView t, const tvm::ffi::TensorView output, int dim, bool flip_sin_to_cos, float downscale_freq_shift, float scale, int max_period) { using namespace host; const int batch_size = static_cast(t.shape()[0]); const int half_dim = dim / 2; constexpr int kMaxThreadsPerBlock = 1024; constexpr int kMinThreadsPerBlock = 128; const int num_threads_per_row = std::min(kMaxThreadsPerBlock, half_dim / 4); const int num_rows = (kMinThreadsPerBlock + num_threads_per_row - 1) / num_threads_per_row; dim3 grid((batch_size + num_rows - 1) / num_rows); dim3 block(num_threads_per_row, num_rows); const float neg_log_max_period = std::log(static_cast(max_period)) * (-1.0f) / (static_cast(half_dim) - downscale_freq_shift); const DLDevice device = output.device(); if (flip_sin_to_cos) { LaunchKernel(grid, block, device)( timestep_embedding_kernel, static_cast(t.data_ptr()), static_cast(output.data_ptr()), dim, neg_log_max_period, scale, batch_size); } else { LaunchKernel(grid, block, device)( timestep_embedding_kernel, static_cast(t.data_ptr()), static_cast(output.data_ptr()), dim, neg_log_max_period, scale, batch_size); } } template void timestep_embedding( tvm::ffi::TensorView input, tvm::ffi::TensorView output, int dim, bool flip_sin_to_cos, float downscale_freq_shift, float scale, int max_period) { using namespace host; auto B = SymbolicSize{"batch_size"}; auto D = SymbolicSize{"dim"}; auto device = SymbolicDevice{}; TensorMatcher({B}) // input .with_strides({1}) .with_dtype() .template with_device(device) .verify(input); TensorMatcher({B, D}).with_strides({D, 1}).with_dtype().template with_device(device).verify(output); RuntimeCheck(D.unwrap() == dim, "Output dim mismatch: ", D.unwrap(), " vs ", dim); RuntimeCheck(dim % 8 == 0, "dim must align to 8, got ", dim); launch_timestep_embedding(input, output, dim, flip_sin_to_cos, downscale_freq_shift, scale, max_period); } } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/elementwise/concat_mla.cuh ================================================ #include #include #include #include #include #include namespace { // ======================= Memory Utilities ======================= // Adapted from DeepEP: https://github.com/deepseek-ai/DeepEP/blob/main/csrc/kernels/utils.cuh SGL_DEVICE int get_lane_id() { int lane_id; asm("mov.s32 %0, %laneid;" : "=r"(lane_id)); return lane_id; } SGL_DEVICE void st_na_global_v1(const int* ptr, int v) { asm volatile("st.global.L1::no_allocate.s32 [%0], %1;" ::"l"(ptr), "r"(v) : "memory"); } SGL_DEVICE void st_na_global_v2(const int2* ptr, const int2& v) { asm volatile("st.global.L1::no_allocate.v2.s32 [%0], {%1, %2};" ::"l"(ptr), "r"(v.x), "r"(v.y) : "memory"); } SGL_DEVICE int ld_na_global_v1(const int* ptr) { int r; asm volatile("ld.global.nc.L1::no_allocate.s32 %0, [%1];" : "=r"(r) : "l"(ptr)); return r; } SGL_DEVICE int2 ld_na_global_v2(const int2* ptr) { int2 r; asm volatile("ld.global.nc.L1::no_allocate.v2.s32 {%0, %1}, [%2];" : "=r"(r.x), "=r"(r.y) : "l"(ptr)); return r; } SGL_DEVICE void prefetch_L2(const void* p) { #if defined(ENABLE_L2_PREFETCH) asm volatile("prefetch.global.L2 [%0];" ::"l"(p)); #endif } // ======================= concat_mla_k Kernel ======================= constexpr int NUM_LOCAL_HEADS = 128; constexpr int QK_NOPE_HEAD_DIM = 128; constexpr int QK_ROPE_HEAD_DIM = 64; constexpr int K_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM; constexpr int HEAD_CHUNK_SIZE = 16; constexpr int NUM_HEAD_CHUNKS = NUM_LOCAL_HEADS / HEAD_CHUNK_SIZE; __global__ void concat_mla_k_kernel( bf16_t* __restrict__ k, const bf16_t* __restrict__ k_nope, const bf16_t* __restrict__ k_rope, const int num_tokens, const int64_t k_stride_0, const int k_stride_1, const int64_t k_nope_stride_0, const int k_nope_stride_1, const int64_t k_rope_stride_0) { const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32; const int token_id = flat_warp_id / NUM_HEAD_CHUNKS; const int head_chunk_id = flat_warp_id % NUM_HEAD_CHUNKS; const int lane_id = get_lane_id(); if (token_id >= num_tokens) return; using NopeVec = int2; // 8B/thread, 32 threads = 256B/row using RopeVec = int; // 4B/thread, 32 threads = 128B/row static_assert(sizeof(NopeVec) * 32 == QK_NOPE_HEAD_DIM * sizeof(bf16_t), "nope vec mismatch"); static_assert(sizeof(RopeVec) * 32 == QK_ROPE_HEAD_DIM * sizeof(bf16_t), "rope vec mismatch"); const int head_row0 = head_chunk_id * HEAD_CHUNK_SIZE; const int2* __restrict__ nope_src = reinterpret_cast(k_nope + token_id * k_nope_stride_0 + head_row0 * k_nope_stride_1) + lane_id; int2* __restrict__ nope_dst = reinterpret_cast(k + token_id * k_stride_0 + head_row0 * k_stride_1) + lane_id; int* __restrict__ rope_dst = reinterpret_cast(k + token_id * k_stride_0 + head_row0 * k_stride_1 + QK_NOPE_HEAD_DIM) + lane_id; const int nope_src_stride_v = (k_nope_stride_1 >> 2); // int2 covers 4 bf16 const int nope_dst_stride_v = (k_stride_1 >> 2); const int rope_dst_stride_v = (k_stride_1 >> 1); // int covers 2 bf16 const int* rope_base = reinterpret_cast(k_rope + token_id * k_rope_stride_0); const RopeVec rope_val = ld_na_global_v1(rope_base + lane_id); prefetch_L2(nope_src); NopeVec cur = ld_na_global_v2(nope_src); #pragma unroll for (int i = 0; i < HEAD_CHUNK_SIZE; ++i) { NopeVec next; if (i + 1 < HEAD_CHUNK_SIZE) { const int2* next_src = nope_src + nope_src_stride_v; prefetch_L2(next_src); next = ld_na_global_v2(next_src); } st_na_global_v2(nope_dst, cur); st_na_global_v1(rope_dst, rope_val); nope_src += nope_src_stride_v; nope_dst += nope_dst_stride_v; rope_dst += rope_dst_stride_v; cur = next; } } struct ConcatMlaKKernel { static void run(tvm::ffi::TensorView k, tvm::ffi::TensorView k_nope, tvm::ffi::TensorView k_rope) { using namespace host; auto N = SymbolicSize{"num_tokens"}; auto H = SymbolicSize{"num_heads"}; auto D = SymbolicSize{"k_head_dim"}; auto D_nope = SymbolicSize{"nope_head_dim"}; auto D_rope = SymbolicSize{"rope_head_dim"}; auto S0_k = SymbolicSize{"k_stride_0"}; auto S1_k = SymbolicSize{"k_stride_1"}; auto S0_k_nope = SymbolicSize{"k_nope_stride_0"}; auto S1_k_nope = SymbolicSize{"k_nope_stride_1"}; auto S0_k_rope = SymbolicSize{"k_rope_stride_0"}; auto device = SymbolicDevice{}; // Set known fixed values H.set_value(NUM_LOCAL_HEADS); D.set_value(K_HEAD_DIM); D_nope.set_value(QK_NOPE_HEAD_DIM); D_rope.set_value(QK_ROPE_HEAD_DIM); // Verify k: [num_tokens, num_heads, k_head_dim] TensorMatcher({N, H, D}).with_strides({S0_k, S1_k, 1}).with_dtype().with_device(device).verify(k); // Verify k_nope: [num_tokens, num_heads, nope_head_dim] TensorMatcher({N, H, D_nope}) .with_strides({S0_k_nope, S1_k_nope, 1}) .with_dtype() .with_device(device) .verify(k_nope); // Verify k_rope: [num_tokens, 1, rope_head_dim] TensorMatcher({N, 1, D_rope}) .with_strides({S0_k_rope, -1, 1}) .with_dtype() .with_device(device) .verify(k_rope); // Check alignment RuntimeCheck(reinterpret_cast(k.data_ptr()) % 16 == 0, "Tensor k must be 16-byte aligned"); RuntimeCheck(reinterpret_cast(k_nope.data_ptr()) % 16 == 0, "Tensor k_nope must be 16-byte aligned"); RuntimeCheck(reinterpret_cast(k_rope.data_ptr()) % 16 == 0, "Tensor k_rope must be 16-byte aligned"); const int num_tokens = static_cast(N.unwrap()); constexpr int num_warps_per_block = 32; const int grid_size = div_ceil(num_tokens * NUM_HEAD_CHUNKS, num_warps_per_block); const int block_size = num_warps_per_block * 32; LaunchKernel(grid_size, block_size, device.unwrap())( concat_mla_k_kernel, static_cast(k.data_ptr()), static_cast(k_nope.data_ptr()), static_cast(k_rope.data_ptr()), num_tokens, S0_k.unwrap(), static_cast(S1_k.unwrap()), S0_k_nope.unwrap(), static_cast(S1_k_nope.unwrap()), S0_k_rope.unwrap()); } }; // ======================= concat_mla_absorb_q Kernel ======================= constexpr int A_LAST_DIM = 512; constexpr int B_LAST_DIM = 64; constexpr int OUT_LAST_DIM = A_LAST_DIM + B_LAST_DIM; __global__ void concat_mla_absorb_q_kernel( bf16_t* a, bf16_t* b, bf16_t* out, const int num_items, const int dim_1, const int64_t a_stride_0, const int a_stride_1, const int64_t b_stride_0, const int b_stride_1, const int64_t out_stride_0, const int out_stride_1) { const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32; const int lane_id = get_lane_id(); const int idx_0 = flat_warp_id / dim_1; const int idx_1 = flat_warp_id % dim_1; if (flat_warp_id >= num_items) { return; } using ABufType = int4; constexpr int A_NUM_UNROLL = 2; static_assert(sizeof(ABufType) * A_NUM_UNROLL == A_LAST_DIM * sizeof(a[0]) / 32); ABufType a_buf[A_NUM_UNROLL]; using BBufType = int; constexpr int B_NUM_UNROLL = 1; static_assert(sizeof(BBufType) * B_NUM_UNROLL == B_LAST_DIM * sizeof(b[0]) / 32); BBufType b_buf; { const BBufType* base_addr = reinterpret_cast(b + idx_0 * b_stride_0 + idx_1 * b_stride_1); b_buf = *(base_addr + lane_id); } #pragma unroll for (int i = 0; i < A_NUM_UNROLL; ++i) { const ABufType* base_addr = reinterpret_cast(a + idx_0 * a_stride_0 + idx_1 * a_stride_1); a_buf[i] = *(base_addr + i * 32 + lane_id); } { BBufType* base_addr = reinterpret_cast(out + idx_0 * out_stride_0 + idx_1 * out_stride_1 + A_LAST_DIM); *(base_addr + lane_id) = b_buf; } #pragma unroll for (int i = 0; i < A_NUM_UNROLL; ++i) { ABufType* base_addr = reinterpret_cast(out + idx_0 * out_stride_0 + idx_1 * out_stride_1); *(base_addr + i * 32 + lane_id) = a_buf[i]; } } struct ConcatMlaAbsorbQKernel { static void run(tvm::ffi::TensorView a, tvm::ffi::TensorView b, tvm::ffi::TensorView out) { using namespace host; auto N0_a = SymbolicSize{"a_dim_0"}; auto N1_a = SymbolicSize{"a_dim_1"}; auto D_a = SymbolicSize{"a_last_dim"}; auto N0_b = SymbolicSize{"b_dim_0"}; auto N1_b = SymbolicSize{"b_dim_1"}; auto D_b = SymbolicSize{"b_last_dim"}; auto N0_out = SymbolicSize{"out_dim_0"}; auto N1_out = SymbolicSize{"out_dim_1"}; auto D_out = SymbolicSize{"out_last_dim"}; auto S0_a = SymbolicSize{"a_stride_0"}; auto S1_a = SymbolicSize{"a_stride_1"}; auto S0_b = SymbolicSize{"b_stride_0"}; auto S1_b = SymbolicSize{"b_stride_1"}; auto S0_out = SymbolicSize{"out_stride_0"}; auto S1_out = SymbolicSize{"out_stride_1"}; auto device = SymbolicDevice{}; // Set known fixed values D_a.set_value(A_LAST_DIM); D_b.set_value(B_LAST_DIM); D_out.set_value(OUT_LAST_DIM); // Verify a: [dim_0, dim_1, A_LAST_DIM] TensorMatcher({N0_a, N1_a, D_a}) .with_strides({S0_a, S1_a, 1}) .with_dtype() .with_device(device) .verify(a); // Verify b: [dim_0, dim_1, B_LAST_DIM] TensorMatcher({N0_b, N1_b, D_b}) .with_strides({S0_b, S1_b, 1}) .with_dtype() .with_device(device) .verify(b); // Verify out: [dim_0, dim_1, OUT_LAST_DIM] TensorMatcher({N0_out, N1_out, D_out}) .with_strides({S0_out, S1_out, 1}) .with_dtype() .with_device(device) .verify(out); // Check alignment RuntimeCheck(reinterpret_cast(a.data_ptr()) % 16 == 0, "Tensor a must be 16-byte aligned"); RuntimeCheck(reinterpret_cast(b.data_ptr()) % 16 == 0, "Tensor b must be 16-byte aligned"); RuntimeCheck(reinterpret_cast(out.data_ptr()) % 16 == 0, "Tensor out must be 16-byte aligned"); // Verify dimensions match: a.size(0) * a.size(1) == b.size(0) * b.size(1) RuntimeCheck( N0_a.unwrap() * N1_a.unwrap() == N0_b.unwrap() * N1_b.unwrap(), "Dimension mismatch: a.size(0) * a.size(1) must equal b.size(0) * b.size(1)"); RuntimeCheck(N1_a.unwrap() == N1_b.unwrap(), "Dimension mismatch: a.size(1) must equal b.size(1)"); const int num_items = static_cast(N0_a.unwrap() * N1_a.unwrap()); const int dim_1 = static_cast(N1_a.unwrap()); constexpr int num_warps_per_block = 32; const int grid_size = div_ceil(num_items, num_warps_per_block); const int block_size = num_warps_per_block * 32; LaunchKernel(grid_size, block_size, device.unwrap())( concat_mla_absorb_q_kernel, static_cast(a.data_ptr()), static_cast(b.data_ptr()), static_cast(out.data_ptr()), num_items, dim_1, S0_a.unwrap(), static_cast(S1_a.unwrap()), S0_b.unwrap(), static_cast(S1_b.unwrap()), S0_out.unwrap(), static_cast(S1_out.unwrap())); } }; } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh ================================================ #include #include #include #include #include #include #include #include #include #include #include namespace { template struct VecTypeTrait; template <> struct VecTypeTrait { using packed_t = packed_t; using vec_t = device::AlignedVector; }; template <> struct VecTypeTrait { using packed_t = packed_t; using vec_t = device::AlignedVector; }; template <> struct VecTypeTrait { using packed_t = packed_t; using vec_t = device::AlignedVector; }; template <> struct VecTypeTrait { using packed_t = packed_t; using vec_t = device::AlignedVector; }; template SGL_DEVICE packed_t rms(packed_t& val, packed_t& weight, float rsqrt_square_sum) { float2 valf = device::cast(val); float2 weightf = device::cast(weight); return device::cast( make_float2(valf.x * weightf.x * rsqrt_square_sum, valf.y * weightf.y * rsqrt_square_sum)); } template __global__ void fused_add_rmsnorm_reg_kernel( T* __restrict__ input, T* __restrict__ residual, const T* __restrict__ weight, int vec_hidden_size, float eps) { constexpr int inner_loop = VEC_SIZE_IN_BYTE == 16 ? 4 : 8; __shared__ float shared_memory[32]; // Used for CTA reduce using vec_t = typename VecTypeTrait::vec_t; using packed_t = typename VecTypeTrait::packed_t; vec_t v; // Save input vec_t v_res; // Save residual vec_t v_weight; // Save weight vec_t v_out; // Save output auto token_id = blockIdx.x; float2 acc_square = make_float2(0.0f, 0.0f); // Sum of squares for each thread if (threadIdx.x < vec_hidden_size) { // Compute address vec_t* p = reinterpret_cast(input) + token_id * vec_hidden_size; vec_t* p_res = reinterpret_cast(residual) + token_id * vec_hidden_size; const vec_t* p_weight = reinterpret_cast(weight); // Load data v = p[threadIdx.x]; v_res = p_res[threadIdx.x]; v_weight = p_weight[threadIdx.x]; for (int i = 0; i < inner_loop; i++) { float2 val = device::cast(v[i]); float2 res = device::cast(v_res[i]); float2 inp_res = make_float2(val.x + res.x, val.y + res.y); acc_square.x += inp_res.x * inp_res.x; acc_square.y += inp_res.y * inp_res.y; v[i] = device::cast(inp_res); } // Store inp+res to residual p_res[threadIdx.x] = v; } // CTA Reduce // Step 0: Warp Reduce auto cg_warp = cooperative_groups::tiled_partition<32>(cooperative_groups::this_thread_block()); float warp_sum = cooperative_groups::reduce(cg_warp, acc_square.x + acc_square.y, cooperative_groups::plus()); float* buffer = shared_memory; if (threadIdx.x % 32 == 0) { buffer[threadIdx.x / 32] = warp_sum; // Write warp_sum to buffer } // Step 1: CTA Reduce __syncthreads(); if (threadIdx.x < 32) { float cta_sum = cooperative_groups::reduce( cg_warp, (threadIdx.x < blockDim.x / 32) ? buffer[threadIdx.x] : 0.0f, cooperative_groups::plus()); buffer[threadIdx.x] = rsqrtf(eps + cta_sum * (1.0f / static_cast(vec_hidden_size * (VEC_SIZE_IN_BYTE / sizeof(T))))); } __syncthreads(); // Compute RMSNorm if (threadIdx.x < vec_hidden_size) { float rsqrt_square_sum = buffer[threadIdx.x / 32]; // Read rsqrt from Shared Memory(Broadcast) for (int i = 0; i < inner_loop; i++) { v_out[i] = rms(v[i], v_weight[i], rsqrt_square_sum); } vec_t* p_out = reinterpret_cast(input) + token_id * vec_hidden_size; p_out[threadIdx.x] = v_out; } } template struct FusedAddRMSNormKernel { static void run(const tvm::ffi::TensorView input, const tvm::ffi::TensorView residual, const tvm::ffi::TensorView weight, float eps) { using namespace host; auto N = SymbolicSize{"num_tokens"}; auto D = SymbolicSize{"hidden_size"}; auto device = SymbolicDevice{}; device.set_options(); TensorMatcher({N, D}) // input .with_strides({D, 1}) .with_dtype() .with_device(device) .verify(input); TensorMatcher({D}) // weight .with_dtype() .with_device(device) .verify(weight); TensorMatcher({N, D}) // residual .with_strides({D, 1}) .with_dtype() .with_device(device) .verify(residual); auto cc_major = host::runtime::get_cc_major(device.unwrap().device_id); int hidden_size = static_cast(D.unwrap()); if ((cc_major <= 9 && hidden_size <= 8192) || (cc_major >= 10 && hidden_size <= 12288)) { int max_vec_size_byte = cc_major >= 10 ? 32 : 16; int elements_in_vec = max_vec_size_byte / sizeof(DType); int vec_hidden_size = hidden_size / elements_in_vec; uint threads = (vec_hidden_size + 31) / 32 * 32; // Runtime check host::RuntimeCheck( hidden_size % elements_in_vec == 0, "hidden_size", hidden_size, " can not align to elements_in_vec ", elements_in_vec); // Launch kernel auto kernel = max_vec_size_byte == 32 ? fused_add_rmsnorm_reg_kernel : fused_add_rmsnorm_reg_kernel; LaunchKernel(static_cast(N.unwrap()), threads, device.unwrap()) .enable_pdl(false)( kernel, reinterpret_cast(input.data_ptr()), reinterpret_cast(residual.data_ptr()), reinterpret_cast(weight.data_ptr()), vec_hidden_size, eps); } else { host::RuntimeCheck(false, "Large hidden_sizes are not supported for now."); } } }; } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh ================================================ /* * Fused metadata copy kernel for NSA backend CUDA graph replay. * JIT-compiled version for python/sglang/jit_kernel. * * OVERVIEW: * This kernel fuses multiple tensor copy operations (cache_seqlens, cu_seqlens_k, * page_table, nsa metadata, and optional FlashMLA metadata) into single kernel * launches, significantly reducing kernel launch overhead and improving CUDA * graph replay performance during inference. * * PERFORMANCE BENEFITS: * - Single kernel launch vs. multiple separate copies (3-10x faster) * - Optimized memory coalescing and SM utilization * - __grid_constant__ parameter passing via constant memory * - Especially beneficial in CUDA graph replay scenarios * * DESIGN: * - Unified kernel supporting all forward modes (DECODE, TARGET_VERIFY, DRAFT_EXTEND) * - Structured parameter passing (SourcePointers/DestinationPointers) for clarity * - Template parameters (HAS_REAL_PAGE_TABLE, HAS_FLASHMLA) for compile-time optimization * - Multi-backend variant copies to 3 destinations in one kernel (for speculative decoding) * * USAGE: * This header is included by JIT compilation system. The FusedMetadataCopyKernel * and FusedMetadataCopyMultiKernel wrapper structs provide the Python-accessible interface. */ #pragma once #include #include #include #include #include // for std::min #include // Forward mode enum (must match Python ForwardMode in sglang/srt/layers/attention/nsa_backend.py) enum ForwardModeEnum { DECODE = 0, TARGET_VERIFY = 1, DRAFT_EXTEND = 2 }; /** * Source pointers for metadata copy operations. * Groups all source tensor pointers for cleaner parameter passing. * Some pointers may be nullptr depending on forward mode and feature flags. */ struct SourcePointers { const int32_t* __restrict__ cache_seqlens; // [bs] sequence lengths in cache const int32_t* __restrict__ cu_seqlens_k; // [bs+1] cumulative sequence lengths const int32_t* __restrict__ page_indices; // page table indices const int32_t* __restrict__ nsa_cache_seqlens; // NSA-specific cache lengths const int32_t* __restrict__ seqlens_expanded; // expanded sequence lengths (TARGET_VERIFY/DRAFT_EXTEND only) const int32_t* __restrict__ nsa_cu_seqlens_k; // NSA cumulative sequence lengths const int32_t* __restrict__ real_page_table; // optional real page table const int32_t* __restrict__ flashmla_num_splits; // optional FlashMLA split counts const int32_t* __restrict__ flashmla_metadata; // optional FlashMLA metadata }; /** * Destination pointers for metadata copy operations. * Groups all destination tensor pointers for cleaner parameter passing. * Layout matches SourcePointers for consistency. */ struct DestinationPointers { int32_t* __restrict__ cache_seqlens; // [bs] sequence lengths in cache int32_t* __restrict__ cu_seqlens_k; // [bs+1] cumulative sequence lengths int32_t* __restrict__ page_table_1; // page table (note: different name from source) int32_t* __restrict__ nsa_cache_seqlens; // NSA-specific cache lengths int32_t* __restrict__ seqlens_expanded; // expanded sequence lengths (TARGET_VERIFY/DRAFT_EXTEND only) int32_t* __restrict__ nsa_cu_seqlens_k; // NSA cumulative sequence lengths int32_t* __restrict__ real_page_table; // optional real page table int32_t* __restrict__ flashmla_num_splits; // optional FlashMLA split counts int32_t* __restrict__ flashmla_metadata; // optional FlashMLA metadata }; /** * Parameter structure for single-backend fused metadata copy kernel. * Passed via __grid_constant__ for efficient constant memory access. */ struct FusedMetadataCopyParams { SourcePointers src; // Source tensor pointers DestinationPointers dst; // Destination tensor pointers // Kernel parameters int forward_mode; // 0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND int bs; // Batch size int max_len; // Max length for DECODE mode int max_seqlen_k; // Max sequence length for TARGET_VERIFY/DRAFT_EXTEND int seqlens_expanded_size; // Size of expanded sequence lengths int page_indices_rows; // Number of rows in page_indices int page_table_1_stride; // Stride for page_table_1 int real_page_table_cols; // Columns in real_page_table int real_page_table_dst_stride; // Stride for destination real_page_table int flashmla_metadata_size; // Size of FlashMLA metadata }; /** * Parameter structure for multi-backend fused metadata copy kernel. * Enables copying from one source to three destinations in a single kernel launch. * Used for speculative decoding with multiple draft backends. */ struct FusedMetadataCopyMultiParams { SourcePointers src; // Source pointers (shared across all backends) DestinationPointers dst0; // Backend 0 destination pointers DestinationPointers dst1; // Backend 1 destination pointers DestinationPointers dst2; // Backend 2 destination pointers // Kernel parameters int bs; // Batch size int max_len; // Max length (DECODE mode only) int seqlens_expanded_size; // Size of expanded sequence lengths int page_table_1_stride; // Stride for page_table_1 int real_page_table_cols; // Columns in real_page_table int real_page_table_dst_stride; // Stride for destination real_page_table int flashmla_metadata_size; // Size of FlashMLA metadata }; /** * Unified kernel for all forward modes (DECODE, TARGET_VERIFY, DRAFT_EXTEND). * Uses runtime branches for mode selection, with template parameters for * compile-time optimization of optional features. * * DESIGN: * - Runtime branches (forward_mode) handle mode-specific logic * - Template parameters (HAS_*) eliminate unused feature code at compile time * - Structured parameters (SourcePointers/DestinationPointers) passed via constant memory * * Used by FusedMetadataCopyKernel for single-backend metadata copy. * * @tparam HAS_REAL_PAGE_TABLE Compile-time flag for real_page_table support * @tparam HAS_FLASHMLA Compile-time flag for FlashMLA metadata support */ template __global__ void fused_metadata_copy_kernel(const FusedMetadataCopyParams __grid_constant__ params) { int tid = blockIdx.x * blockDim.x + threadIdx.x; int total_threads = gridDim.x * blockDim.x; // Unpack parameters for readability const auto& src = params.src; const auto& dst = params.dst; const int forward_mode = params.forward_mode; const int bs = params.bs; const int max_len = params.max_len; const int max_seqlen_k = params.max_seqlen_k; const int seqlens_expanded_size = params.seqlens_expanded_size; const int page_indices_rows = params.page_indices_rows; const int page_table_1_stride = params.page_table_1_stride; const int real_page_table_cols = params.real_page_table_cols; const int real_page_table_dst_stride = params.real_page_table_dst_stride; const int flashmla_metadata_size = params.flashmla_metadata_size; // Copy cache_seqlens (bs elements) - common to all modes #pragma unroll 8 for (int i = tid; i < bs; i += total_threads) { dst.cache_seqlens[i] = src.cache_seqlens[i]; } // Copy cu_seqlens_k (skip first element) - common to all modes #pragma unroll 8 for (int i = tid; i < bs; i += total_threads) { dst.cu_seqlens_k[i + 1] = src.cu_seqlens_k[i + 1]; } // Branch 1: page_table copy (different dimensions per mode) if (forward_mode == 0) { // DECODE int page_table_elements = bs * max_len; #pragma unroll 4 for (int i = tid; i < page_table_elements; i += total_threads) { int row = i / max_len; int col = i % max_len; dst.page_table_1[row * page_table_1_stride + col] = src.page_indices[i]; } } else { // TARGET_VERIFY or DRAFT_EXTEND int page_table_elements = page_indices_rows * max_seqlen_k; #pragma unroll 4 for (int i = tid; i < page_table_elements; i += total_threads) { int row = i / max_seqlen_k; int col = i % max_seqlen_k; dst.page_table_1[row * page_table_1_stride + col] = src.page_indices[i]; } } // Branch 2: seqlens_expanded copy (only for TARGET_VERIFY/DRAFT_EXTEND) if (forward_mode != 0) { // TARGET_VERIFY or DRAFT_EXTEND #pragma unroll 4 for (int i = tid; i < seqlens_expanded_size; i += total_threads) { dst.seqlens_expanded[i] = src.seqlens_expanded[i]; } } // Branch 3: NSA metadata copy (different loop sizes per mode) if (forward_mode == 0) { // DECODE #pragma unroll 8 for (int i = tid; i < bs; i += total_threads) { dst.nsa_cache_seqlens[i] = src.nsa_cache_seqlens[i]; } #pragma unroll 8 for (int i = tid; i < bs; i += total_threads) { dst.nsa_cu_seqlens_k[i + 1] = src.nsa_cu_seqlens_k[i + 1]; } } else { // TARGET_VERIFY or DRAFT_EXTEND #pragma unroll 4 for (int i = tid; i < seqlens_expanded_size; i += total_threads) { dst.nsa_cache_seqlens[i] = src.nsa_cache_seqlens[i]; } #pragma unroll 4 for (int i = tid; i < seqlens_expanded_size; i += total_threads) { dst.nsa_cu_seqlens_k[i + 1] = src.nsa_cu_seqlens_k[i + 1]; } } // Copy real page table - compile-time branch if constexpr (HAS_REAL_PAGE_TABLE) { int real_table_elements = (forward_mode == 0 ? bs : page_indices_rows) * real_page_table_cols; #pragma unroll 2 for (int i = tid; i < real_table_elements; i += total_threads) { int row = i / real_page_table_cols; int col = i % real_page_table_cols; dst.real_page_table[row * real_page_table_dst_stride + col] = src.real_page_table[row * real_page_table_cols + col]; } } // Branch 4: FlashMLA metadata copy (different sizes per mode) if constexpr (HAS_FLASHMLA) { int flashmla_size = (forward_mode == 0) ? (bs + 1) : (seqlens_expanded_size + 1); if (forward_mode == 0) { #pragma unroll 8 for (int i = tid; i < flashmla_size; i += total_threads) { dst.flashmla_num_splits[i] = src.flashmla_num_splits[i]; } } else { #pragma unroll 4 for (int i = tid; i < flashmla_size; i += total_threads) { dst.flashmla_num_splits[i] = src.flashmla_num_splits[i]; } } #pragma unroll 2 for (int i = tid; i < flashmla_metadata_size; i += total_threads) { dst.flashmla_metadata[i] = src.flashmla_metadata[i]; } } } /** * Multi-backend kernel for DECODE mode. * Copies from one source to THREE destinations in a single kernel launch. * * PERFORMANCE: 3x faster than three separate kernel launches due to: * - Reduced kernel launch overhead (1 launch instead of 3) * - Improved memory coalescing (source read once, written to 3 destinations) * - Better instruction-level parallelism * * Used by FusedMetadataCopyMultiKernel for speculative decoding scenarios. * * @tparam HAS_REAL_PAGE_TABLE Compile-time flag for real_page_table support * @tparam HAS_FLASHMLA Compile-time flag for FlashMLA metadata support */ template __global__ void fused_metadata_copy_multi_kernel(const FusedMetadataCopyMultiParams __grid_constant__ params) { int tid = blockIdx.x * blockDim.x + threadIdx.x; int total_threads = gridDim.x * blockDim.x; // Unpack parameters for readability const auto& src = params.src; const auto& dst0 = params.dst0; const auto& dst1 = params.dst1; const auto& dst2 = params.dst2; const int bs = params.bs; const int max_len = params.max_len; const int seqlens_expanded_size = params.seqlens_expanded_size; const int page_table_1_stride = params.page_table_1_stride; const int real_page_table_cols = params.real_page_table_cols; const int real_page_table_dst_stride = params.real_page_table_dst_stride; const int flashmla_metadata_size = params.flashmla_metadata_size; // Copy cache_seqlens to all 3 backends #pragma unroll 8 for (int i = tid; i < bs; i += total_threads) { int32_t val = src.cache_seqlens[i]; dst0.cache_seqlens[i] = val; dst1.cache_seqlens[i] = val; dst2.cache_seqlens[i] = val; } // Copy cu_seqlens_k to all 3 backends (skip first element) #pragma unroll 8 for (int i = tid; i < bs; i += total_threads) { int32_t val = src.cu_seqlens_k[i + 1]; dst0.cu_seqlens_k[i + 1] = val; dst1.cu_seqlens_k[i + 1] = val; dst2.cu_seqlens_k[i + 1] = val; } // DECODE mode: copy page_table_1 to all 3 backends int page_table_elements = bs * max_len; #pragma unroll 4 for (int i = tid; i < page_table_elements; i += total_threads) { int row = i / max_len; int col = i % max_len; int32_t val = src.page_indices[i]; dst0.page_table_1[row * page_table_1_stride + col] = val; dst1.page_table_1[row * page_table_1_stride + col] = val; dst2.page_table_1[row * page_table_1_stride + col] = val; } // Copy nsa_cache_seqlens to all 3 backends #pragma unroll 8 for (int i = tid; i < bs; i += total_threads) { int32_t val = src.nsa_cache_seqlens[i]; dst0.nsa_cache_seqlens[i] = val; dst1.nsa_cache_seqlens[i] = val; dst2.nsa_cache_seqlens[i] = val; } // Copy NSA cu_seqlens to all 3 backends #pragma unroll 8 for (int i = tid; i < bs; i += total_threads) { int32_t val = src.nsa_cu_seqlens_k[i + 1]; dst0.nsa_cu_seqlens_k[i + 1] = val; dst1.nsa_cu_seqlens_k[i + 1] = val; dst2.nsa_cu_seqlens_k[i + 1] = val; } // Copy real page table to all 3 backends if (src.real_page_table != nullptr && dst0.real_page_table != nullptr) { int real_table_elements = bs * real_page_table_cols; #pragma unroll 2 for (int i = tid; i < real_table_elements; i += total_threads) { int row = i / real_page_table_cols; int col = i % real_page_table_cols; int src_idx = row * real_page_table_cols + col; int dst_idx = row * real_page_table_dst_stride + col; int32_t val = src.real_page_table[src_idx]; dst0.real_page_table[dst_idx] = val; dst1.real_page_table[dst_idx] = val; dst2.real_page_table[dst_idx] = val; } } // Copy FlashMLA metadata to all 3 backends if constexpr (HAS_FLASHMLA) { int flashmla_size = bs + 1; #pragma unroll 8 for (int i = tid; i < flashmla_size; i += total_threads) { int32_t val = src.flashmla_num_splits[i]; dst0.flashmla_num_splits[i] = val; dst1.flashmla_num_splits[i] = val; dst2.flashmla_num_splits[i] = val; } #pragma unroll 2 for (int i = tid; i < flashmla_metadata_size; i += total_threads) { int32_t val = src.flashmla_metadata[i]; dst0.flashmla_metadata[i] = val; dst1.flashmla_metadata[i] = val; dst2.flashmla_metadata[i] = val; } } } // ============================================================================ // Host-side launcher wrappers for JIT compilation // ============================================================================ namespace { // Launch configuration constants constexpr int THREADS_PER_BLOCK = 256; constexpr int MAX_GRID_SIZE = 1024; // Limit to prevent excessive resource usage /** * Helper function to extract a typed data pointer from a TensorView. * Performs runtime type checking and returns the properly cast pointer. * * @tparam T The expected element type (e.g., int32_t) * @param tensor The TensorView to extract the pointer from * @param name The name of the tensor (for error reporting) * @return Typed pointer to the tensor data */ template inline const T* unwrap_data_ptr(const tvm::ffi::TensorView& tensor, const char* name) { using namespace host; if (tensor.data_ptr()) { RuntimeCheck(is_type(tensor.dtype()), "Tensor ", name, " must have dtype int32"); } return static_cast(tensor.data_ptr()); } /** * Helper function to extract a typed mutable data pointer from a TensorView. * Performs runtime type checking and returns the properly cast pointer. * * @tparam T The expected element type (e.g., int32_t) * @param tensor The TensorView to extract the pointer from * @param name The name of the tensor (for error reporting) * @return Typed mutable pointer to the tensor data */ template inline T* unwrap_data_ptr_mut(const tvm::ffi::TensorView& tensor, const char* name) { using namespace host; if (tensor.data_ptr()) { RuntimeCheck(is_type(tensor.dtype()), "Tensor ", name, " must have dtype int32"); } return static_cast(tensor.data_ptr()); } /** * Helper function to extract a typed data pointer from an Optional TensorView. * Returns nullptr if the optional has no value, otherwise performs type checking. * * @tparam T The expected element type (e.g., int32_t) * @param optional_tensor The Optional TensorView to extract the pointer from * @param name The name of the tensor (for error reporting) * @return Typed pointer to the tensor data, or nullptr if optional has no value */ template inline const T* unwrap_optional_data_ptr(const tvm::ffi::Optional& optional_tensor, const char* name) { using namespace host; if (!optional_tensor.has_value()) { return nullptr; } const auto& tensor = optional_tensor.value(); RuntimeCheck(is_type(tensor.dtype()), "Tensor ", name, " must have dtype int32"); return static_cast(tensor.data_ptr()); } /** * Helper function to extract a typed mutable data pointer from an Optional TensorView. * Returns nullptr if the optional has no value, otherwise performs type checking. * * @tparam T The expected element type (e.g., int32_t) * @param optional_tensor The Optional TensorView to extract the pointer from * @param name The name of the tensor (for error reporting) * @return Typed mutable pointer to the tensor data, or nullptr if optional has no value */ template inline T* unwrap_optional_data_ptr_mut(const tvm::ffi::Optional& optional_tensor, const char* name) { using namespace host; if (!optional_tensor.has_value()) { return nullptr; } const auto& tensor = optional_tensor.value(); RuntimeCheck(is_type(tensor.dtype()), "Tensor ", name, " must have dtype int32"); return static_cast(tensor.data_ptr()); } /** * Calculate kernel launch configuration. * * @param total_work Total number of work items * @param threads_per_block Threads per block (default: THREADS_PER_BLOCK) * @return Grid dimension for kernel launch */ inline dim3 get_launch_config(int total_work, int threads_per_block = THREADS_PER_BLOCK) { int num_blocks = (total_work + threads_per_block - 1) / threads_per_block; // Limit grid size to prevent excessive resource usage while ensuring coverage num_blocks = std::min(num_blocks, MAX_GRID_SIZE); return dim3(num_blocks); } /** * JIT wrapper for single-backend fused metadata copy kernel. * * This struct provides a unified interface for launching the fused metadata copy * kernel with different forward modes. It constructs the parameter struct and * launches the unified kernel. * * IMPLEMENTATION: * - Extracts raw pointers from TensorView objects * - Constructs FusedMetadataCopyParams with nested SourcePointers/DestinationPointers * - Calculates grid configuration based on maximum work size * - Launches fused_metadata_copy_kernel with __grid_constant__ parameters * * @tparam FORWARD_MODE Forward mode: 0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND * @tparam HAS_REAL_PAGE_TABLE Whether real_page_table tensors are present * @tparam HAS_FLASHMLA Whether FlashMLA metadata tensors are present */ template struct FusedMetadataCopyKernel { static_assert( FORWARD_MODE >= 0 && FORWARD_MODE <= 2, "FORWARD_MODE must be 0 (DECODE), 1 (TARGET_VERIFY), or 2 (DRAFT_EXTEND)"); static void run(const tvm::ffi::TensorView cache_seqlens_src, const tvm::ffi::TensorView cu_seqlens_k_src, const tvm::ffi::TensorView page_indices_src, const tvm::ffi::TensorView nsa_cache_seqlens_src, const tvm::ffi::Optional seqlens_expanded_src, const tvm::ffi::TensorView nsa_cu_seqlens_k_src, const tvm::ffi::Optional real_page_table_src, const tvm::ffi::Optional flashmla_num_splits_src, const tvm::ffi::Optional flashmla_metadata_src, const tvm::ffi::TensorView cache_seqlens_dst, const tvm::ffi::TensorView cu_seqlens_k_dst, const tvm::ffi::TensorView page_table_1_dst, const tvm::ffi::TensorView nsa_cache_seqlens_dst, const tvm::ffi::Optional seqlens_expanded_dst, const tvm::ffi::TensorView nsa_cu_seqlens_k_dst, const tvm::ffi::Optional real_page_table_dst, const tvm::ffi::Optional flashmla_num_splits_dst, const tvm::ffi::Optional flashmla_metadata_dst, int bs, int max_len, int max_seqlen_k, int seqlens_expanded_size) { using namespace host; // Build parameter struct with nested source/destination pointers // unwrap_data_ptr and unwrap_optional_data_ptr perform dtype validation const auto params = FusedMetadataCopyParams{ .src = { .cache_seqlens = unwrap_data_ptr(cache_seqlens_src, "cache_seqlens_src"), .cu_seqlens_k = unwrap_data_ptr(cu_seqlens_k_src, "cu_seqlens_k_src"), .page_indices = unwrap_data_ptr(page_indices_src, "page_indices_src"), .nsa_cache_seqlens = unwrap_data_ptr(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"), .seqlens_expanded = unwrap_optional_data_ptr(seqlens_expanded_src, "seqlens_expanded_src"), .nsa_cu_seqlens_k = unwrap_data_ptr(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"), .real_page_table = unwrap_optional_data_ptr(real_page_table_src, "real_page_table_src"), .flashmla_num_splits = unwrap_optional_data_ptr(flashmla_num_splits_src, "flashmla_num_splits_src"), .flashmla_metadata = unwrap_optional_data_ptr(flashmla_metadata_src, "flashmla_metadata_src"), }, .dst = { .cache_seqlens = unwrap_data_ptr_mut(cache_seqlens_dst, "cache_seqlens_dst"), .cu_seqlens_k = unwrap_data_ptr_mut(cu_seqlens_k_dst, "cu_seqlens_k_dst"), .page_table_1 = unwrap_data_ptr_mut(page_table_1_dst, "page_table_1_dst"), .nsa_cache_seqlens = unwrap_data_ptr_mut(nsa_cache_seqlens_dst, "nsa_cache_seqlens_dst"), .seqlens_expanded = unwrap_optional_data_ptr_mut(seqlens_expanded_dst, "seqlens_expanded_dst"), .nsa_cu_seqlens_k = unwrap_data_ptr_mut(nsa_cu_seqlens_k_dst, "nsa_cu_seqlens_k_dst"), .real_page_table = unwrap_optional_data_ptr_mut(real_page_table_dst, "real_page_table_dst"), .flashmla_num_splits = unwrap_optional_data_ptr_mut(flashmla_num_splits_dst, "flashmla_num_splits_dst"), .flashmla_metadata = unwrap_optional_data_ptr_mut(flashmla_metadata_dst, "flashmla_metadata_dst"), }, .forward_mode = FORWARD_MODE, .bs = bs, .max_len = max_len, .max_seqlen_k = max_seqlen_k, .seqlens_expanded_size = seqlens_expanded_size, .page_indices_rows = static_cast(page_indices_src.shape()[0]), .page_table_1_stride = static_cast(page_table_1_dst.shape()[1]), .real_page_table_cols = real_page_table_src.has_value() ? static_cast(real_page_table_src.value().shape()[1]) : 0, .real_page_table_dst_stride = real_page_table_dst.has_value() ? static_cast(real_page_table_dst.value().stride(0)) : 0, .flashmla_metadata_size = flashmla_metadata_src.has_value() ? static_cast(flashmla_metadata_src.value().numel()) : 0, }; // Calculate grid configuration int max_elements = std::max( {bs, params.page_indices_rows * max_seqlen_k, seqlens_expanded_size, HAS_FLASHMLA ? (seqlens_expanded_size + 1) : 0, HAS_FLASHMLA ? params.flashmla_metadata_size : 0}); dim3 grid = get_launch_config(max_elements); dim3 block(THREADS_PER_BLOCK); DLDevice device = cache_seqlens_src.device(); // Launch unified kernel with params struct host::LaunchKernel(grid, block, device)(fused_metadata_copy_kernel, params); } }; /** * JIT wrapper for multi-backend fused metadata copy kernel. * * This kernel optimizes the common case where metadata needs to be copied from * one source to THREE destination backends in a single kernel launch. This is * 3x faster than launching three separate kernels due to: * - Reduced kernel launch overhead (1 launch instead of 3) * - Improved memory coalescing (source read once, written to 3 destinations) * - Better GPU occupancy and instruction-level parallelism * * USAGE: Primarily for speculative decoding with multiple draft models, where * the same source metadata needs to be replicated to multiple backend contexts. * * LIMITATION: Currently only supports DECODE mode, which is the most frequently * used mode in speculative decoding scenarios. * * IMPLEMENTATION: * - Constructs FusedMetadataCopyMultiParams with 1 SourcePointers + 3 DestinationPointers * - Launches fused_metadata_copy_multi_kernel with __grid_constant__ parameters * * @tparam HAS_REAL_PAGE_TABLE Whether real_page_table tensors are present * @tparam HAS_FLASHMLA Whether FlashMLA metadata tensors are present */ template struct FusedMetadataCopyMultiKernel { static void run(const tvm::ffi::TensorView cache_seqlens_src, const tvm::ffi::TensorView cu_seqlens_k_src, const tvm::ffi::TensorView page_indices_src, const tvm::ffi::TensorView nsa_cache_seqlens_src, const tvm::ffi::TensorView nsa_cu_seqlens_k_src, const tvm::ffi::Optional real_page_table_src, const tvm::ffi::Optional flashmla_num_splits_src, const tvm::ffi::Optional flashmla_metadata_src, const tvm::ffi::TensorView cache_seqlens_dst0, const tvm::ffi::TensorView cu_seqlens_k_dst0, const tvm::ffi::TensorView page_table_1_dst0, const tvm::ffi::TensorView nsa_cache_seqlens_dst0, const tvm::ffi::TensorView nsa_cu_seqlens_k_dst0, const tvm::ffi::Optional real_page_table_dst0, const tvm::ffi::Optional flashmla_num_splits_dst0, const tvm::ffi::Optional flashmla_metadata_dst0, const tvm::ffi::TensorView cache_seqlens_dst1, const tvm::ffi::TensorView cu_seqlens_k_dst1, const tvm::ffi::TensorView page_table_1_dst1, const tvm::ffi::TensorView nsa_cache_seqlens_dst1, const tvm::ffi::TensorView nsa_cu_seqlens_k_dst1, const tvm::ffi::Optional real_page_table_dst1, const tvm::ffi::Optional flashmla_num_splits_dst1, const tvm::ffi::Optional flashmla_metadata_dst1, const tvm::ffi::TensorView cache_seqlens_dst2, const tvm::ffi::TensorView cu_seqlens_k_dst2, const tvm::ffi::TensorView page_table_1_dst2, const tvm::ffi::TensorView nsa_cache_seqlens_dst2, const tvm::ffi::TensorView nsa_cu_seqlens_k_dst2, const tvm::ffi::Optional real_page_table_dst2, const tvm::ffi::Optional flashmla_num_splits_dst2, const tvm::ffi::Optional flashmla_metadata_dst2, int bs, int max_len, int seqlens_expanded_size) { using namespace host; // Build parameter struct with nested source/destination pointers // unwrap_data_ptr and unwrap_optional_data_ptr perform dtype validation const auto params = FusedMetadataCopyMultiParams{ .src = { .cache_seqlens = unwrap_data_ptr(cache_seqlens_src, "cache_seqlens_src"), .cu_seqlens_k = unwrap_data_ptr(cu_seqlens_k_src, "cu_seqlens_k_src"), .page_indices = unwrap_data_ptr(page_indices_src, "page_indices_src"), .nsa_cache_seqlens = unwrap_data_ptr(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"), .seqlens_expanded = nullptr, // Not used in multi-backend DECODE mode .nsa_cu_seqlens_k = unwrap_data_ptr(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"), .real_page_table = unwrap_optional_data_ptr(real_page_table_src, "real_page_table_src"), .flashmla_num_splits = unwrap_optional_data_ptr(flashmla_num_splits_src, "flashmla_num_splits_src"), .flashmla_metadata = unwrap_optional_data_ptr(flashmla_metadata_src, "flashmla_metadata_src"), }, .dst0 = { .cache_seqlens = unwrap_data_ptr_mut(cache_seqlens_dst0, "cache_seqlens_dst0"), .cu_seqlens_k = unwrap_data_ptr_mut(cu_seqlens_k_dst0, "cu_seqlens_k_dst0"), .page_table_1 = unwrap_data_ptr_mut(page_table_1_dst0, "page_table_1_dst0"), .nsa_cache_seqlens = unwrap_data_ptr_mut(nsa_cache_seqlens_dst0, "nsa_cache_seqlens_dst0"), .seqlens_expanded = nullptr, .nsa_cu_seqlens_k = unwrap_data_ptr_mut(nsa_cu_seqlens_k_dst0, "nsa_cu_seqlens_k_dst0"), .real_page_table = unwrap_optional_data_ptr_mut(real_page_table_dst0, "real_page_table_dst0"), .flashmla_num_splits = unwrap_optional_data_ptr_mut(flashmla_num_splits_dst0, "flashmla_num_splits_dst0"), .flashmla_metadata = unwrap_optional_data_ptr_mut(flashmla_metadata_dst0, "flashmla_metadata_dst0"), }, .dst1 = { .cache_seqlens = unwrap_data_ptr_mut(cache_seqlens_dst1, "cache_seqlens_dst1"), .cu_seqlens_k = unwrap_data_ptr_mut(cu_seqlens_k_dst1, "cu_seqlens_k_dst1"), .page_table_1 = unwrap_data_ptr_mut(page_table_1_dst1, "page_table_1_dst1"), .nsa_cache_seqlens = unwrap_data_ptr_mut(nsa_cache_seqlens_dst1, "nsa_cache_seqlens_dst1"), .seqlens_expanded = nullptr, .nsa_cu_seqlens_k = unwrap_data_ptr_mut(nsa_cu_seqlens_k_dst1, "nsa_cu_seqlens_k_dst1"), .real_page_table = unwrap_optional_data_ptr_mut(real_page_table_dst1, "real_page_table_dst1"), .flashmla_num_splits = unwrap_optional_data_ptr_mut(flashmla_num_splits_dst1, "flashmla_num_splits_dst1"), .flashmla_metadata = unwrap_optional_data_ptr_mut(flashmla_metadata_dst1, "flashmla_metadata_dst1"), }, .dst2 = { .cache_seqlens = unwrap_data_ptr_mut(cache_seqlens_dst2, "cache_seqlens_dst2"), .cu_seqlens_k = unwrap_data_ptr_mut(cu_seqlens_k_dst2, "cu_seqlens_k_dst2"), .page_table_1 = unwrap_data_ptr_mut(page_table_1_dst2, "page_table_1_dst2"), .nsa_cache_seqlens = unwrap_data_ptr_mut(nsa_cache_seqlens_dst2, "nsa_cache_seqlens_dst2"), .seqlens_expanded = nullptr, .nsa_cu_seqlens_k = unwrap_data_ptr_mut(nsa_cu_seqlens_k_dst2, "nsa_cu_seqlens_k_dst2"), .real_page_table = unwrap_optional_data_ptr_mut(real_page_table_dst2, "real_page_table_dst2"), .flashmla_num_splits = unwrap_optional_data_ptr_mut(flashmla_num_splits_dst2, "flashmla_num_splits_dst2"), .flashmla_metadata = unwrap_optional_data_ptr_mut(flashmla_metadata_dst2, "flashmla_metadata_dst2"), }, .bs = bs, .max_len = max_len, .seqlens_expanded_size = seqlens_expanded_size, .page_table_1_stride = static_cast(page_table_1_dst0.shape()[1]), .real_page_table_cols = real_page_table_src.has_value() ? static_cast(real_page_table_src.value().shape()[1]) : 0, .real_page_table_dst_stride = real_page_table_dst0.has_value() ? static_cast(real_page_table_dst0.value().stride(0)) : 0, .flashmla_metadata_size = flashmla_metadata_src.has_value() ? static_cast(flashmla_metadata_src.value().numel()) : 0, }; dim3 grid = get_launch_config(bs * max_len); dim3 block(THREADS_PER_BLOCK); DLDevice device = cache_seqlens_src.device(); // Launch multi-backend kernel with params struct host::LaunchKernel(grid, block, device)( fused_metadata_copy_multi_kernel, params); } }; } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh ================================================ #include #include #include #include #include #include #include #include namespace { struct StoreKVCacheParams { const void* __restrict__ k; const void* __restrict__ v; void* __restrict__ k_cache; void* __restrict__ v_cache; const void* __restrict__ indices; int64_t stride_k_bytes; int64_t stride_v_bytes; int64_t stride_cache_bytes; int64_t stride_indices; uint32_t batch_size; }; constexpr uint32_t kNumWarps = 4; constexpr uint32_t kThreadsPerBlock = kNumWarps * device::kWarpThreads; /** * \brief Use a single warp to copy key and value data from source to destination. * Each thread in the warp copies a portion of the data in a coalesced manner. * \tparam kElementBytes The size of each key/value element in bytes. * \param k_src Pointer to the source key data. * \param v_src Pointer to the source value data. * \param k_dst Pointer to the destination key data. * \param v_dst Pointer to the destination value data. */ template SGL_DEVICE void copy_kv_warp( const void* __restrict__ k_src, const void* __restrict__ v_src, void* __restrict__ k_dst, void* __restrict__ v_dst) { using namespace device; constexpr int64_t kAlignment = (kElementBytes % (16 * kWarpThreads) == 0) ? 16 : kElementBytes % (8 * kWarpThreads) == 0 ? 8 : kElementBytes % (4 * kWarpThreads) == 0 ? 4 : kElementBytes % 4 == 0 ? 4 : 0; static_assert(kAlignment > 0, "Element size must be multiple of 4 bytes"); using vec_t = AlignedStorage; constexpr auto kLoopBytes = sizeof(vec_t) * kWarpThreads; constexpr auto kLoopCount = kElementBytes / kLoopBytes; const auto gmem = tile::Memory::warp(); #pragma unroll kLoopCount for (int64_t i = 0; i < kLoopCount; ++i) { const auto k = gmem.load(k_src, i); const auto v = gmem.load(v_src, i); gmem.store(k_dst, k, i); gmem.store(v_dst, v, i); } // handle the epilogue if any if constexpr (kLoopCount * kLoopBytes < kElementBytes) { if (gmem.in_bound(kElementBytes / sizeof(vec_t), kLoopCount)) { const auto k = gmem.load(k_src, kLoopCount); const auto v = gmem.load(v_src, kLoopCount); gmem.store(k_dst, k, kLoopCount); gmem.store(v_dst, v, kLoopCount); } } } /** * \brief Kernel to store key-value pairs into the KV cache. * Each element is split into multiple parts to allow parallel memory copy. * \tparam kElementBytes The size of each key/value element in bytes. * \tparam kSplit The number of warps that handle each element. * \tparam kUsePDL Whether to use PDL feature. * \tparam T The data type of the indices (`int32_t` or `int64_t`). */ template __global__ void store_kvcache(const __grid_constant__ StoreKVCacheParams params) { using namespace device; constexpr auto kSplitSize = kElementBytes / kSplit; const uint32_t warp_id = blockIdx.x * kNumWarps + threadIdx.x / kWarpThreads; const uint32_t item_id = warp_id / kSplit; const uint32_t split_id = warp_id % kSplit; const auto& [ k_input, v_input, k_cache, v_cache, indices, // ptr stride_k, stride_v, stride_cache, stride_indices, batch_size // size ] = params; if (item_id >= batch_size) return; const auto index_ptr = static_cast(indices) + item_id * stride_indices; PDLWaitPrimary(); const auto index = *index_ptr; const auto k_src = pointer::offset(k_input, item_id * stride_k, split_id * kSplitSize); const auto v_src = pointer::offset(v_input, item_id * stride_v, split_id * kSplitSize); const auto k_dst = pointer::offset(k_cache, index * stride_cache, split_id * kSplitSize); const auto v_dst = pointer::offset(v_cache, index * stride_cache, split_id * kSplitSize); copy_kv_warp(k_src, v_src, k_dst, v_dst); PDLTriggerSecondary(); } template struct StoreKVCacheKernel { static_assert(kElementBytes > 0 && kElementBytes % 4 == 0); template static constexpr auto store_kernel = store_kvcache; template static auto get_kernel(const int num_split) { using namespace host; // only apply split optimization when element size is aligned if constexpr (kElementBytes % (4 * 128) == 0) { if (num_split == 4) return store_kernel<4, T>; } if constexpr (kElementBytes % (2 * 128) == 0) { if (num_split == 2) return store_kernel<2, T>; } if (num_split == 1) return store_kernel<1, T>; Panic("Unsupported num_split {} for element size {}", num_split, kElementBytes); } static void run(const tvm::ffi::TensorView k, const tvm::ffi::TensorView v, const tvm::ffi::TensorView k_cache, const tvm::ffi::TensorView v_cache, const tvm::ffi::TensorView indices, const int num_split) { using namespace host; auto B = SymbolicSize{"batch_size"}; auto D = SymbolicSize{"element_size"}; auto KS = SymbolicSize{"k_stride"}; auto VS = SymbolicSize{"v_stride"}; auto S = SymbolicSize{"cache_stride"}; auto I = SymbolicSize{"indices_stride"}; auto dtype = SymbolicDType{}; auto device = SymbolicDevice{}; auto indice_dtype = SymbolicDType{}; device.set_options(); TensorMatcher({B, D}) // .with_strides({KS, 1}) .with_dtype(dtype) .with_device(device) .verify(k); TensorMatcher({B, D}) // .with_strides({VS, 1}) .with_dtype(dtype) .with_device(device) .verify(v); TensorMatcher({-1, D}) // .with_strides({S, 1}) .with_dtype(dtype) .with_device(device) .verify(k_cache) .verify(v_cache); TensorMatcher({B}) // .with_strides({I}) .with_dtype(indice_dtype) .with_device(device) .verify(indices); const int64_t dtype_size = dtype_bytes(dtype.unwrap()); const uint32_t num_elements = static_cast(B.unwrap()); RuntimeCheck(kElementBytes == dtype_size * D.unwrap()); const auto params = StoreKVCacheParams{ .k = k.data_ptr(), .v = v.data_ptr(), .k_cache = k_cache.data_ptr(), .v_cache = v_cache.data_ptr(), .indices = indices.data_ptr(), .stride_k_bytes = KS.unwrap() * dtype_size, .stride_v_bytes = VS.unwrap() * dtype_size, .stride_cache_bytes = S.unwrap() * dtype_size, .stride_indices = I.unwrap(), .batch_size = static_cast(B.unwrap()), }; // select kernel and update num_split if needed const auto use_int32 = indice_dtype.is_type(); const auto kernel = use_int32 ? get_kernel(num_split) : get_kernel(num_split); const auto num_blocks = div_ceil(num_elements * num_split, kNumWarps); LaunchKernel(num_blocks, kThreadsPerBlock, device.unwrap()) // .enable_pdl(kUsePDL)(kernel, params); } }; } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/elementwise/pos_enc.cuh ================================================ // Adapted from // https://github.com/vllm-project/vllm/blob/014ece97c7aa49084a1119dca792af081a18dbc1/csrc/pos_encoding_kernels.cu #include #include #include #include #include #include namespace { template inline __device__ void apply_token_rotary_embedding( scalar_t* __restrict__ arr, const scalar_t* __restrict__ cos_ptr, const scalar_t* __restrict__ sin_ptr, int rot_offset, int embed_dim) { int x_index, y_index; scalar_t cos, sin; if (IS_NEOX) { // GPT-NeoX style rotary embedding. x_index = rot_offset; y_index = embed_dim + rot_offset; cos = SGLANG_LDG(cos_ptr + x_index); sin = SGLANG_LDG(sin_ptr + x_index); } else { // GPT-J style rotary embedding. x_index = 2 * rot_offset; y_index = 2 * rot_offset + 1; cos = SGLANG_LDG(cos_ptr + x_index / 2); sin = SGLANG_LDG(sin_ptr + x_index / 2); } const scalar_t x = arr[x_index]; const scalar_t y = arr[y_index]; arr[x_index] = x * cos - y * sin; arr[y_index] = y * cos + x * sin; } template inline __device__ void apply_rotary_embedding( scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads, // head_size] or [num_tokens, num_heads, // head_size] scalar_t* __restrict__ key, // nullptr or // [batch_size, seq_len, num_kv_heads, // head_size] or [num_tokens, num_kv_heads, // head_size] const scalar_t* cache_ptr, const int head_size, const int num_heads, const int num_kv_heads, const int rot_dim, const int token_idx, const int64_t query_stride, const int64_t key_stride, const int64_t head_stride) { const int embed_dim = rot_dim / 2; const scalar_t* cos_ptr = cache_ptr; const scalar_t* sin_ptr = cache_ptr + embed_dim; const int nq = num_heads * embed_dim; for (int i = threadIdx.x; i < nq; i += blockDim.x) { const int head_idx = i / embed_dim; const int64_t token_head = token_idx * query_stride + head_idx * head_stride; const int rot_offset = i % embed_dim; apply_token_rotary_embedding(query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim); } if (key != nullptr) { const int nk = num_kv_heads * embed_dim; for (int i = threadIdx.x; i < nk; i += blockDim.x) { const int head_idx = i / embed_dim; const int64_t token_head = token_idx * key_stride + head_idx * head_stride; const int rot_offset = i % embed_dim; apply_token_rotary_embedding(key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim); } } } template __global__ void rotary_embedding_kernel( const int64_t* __restrict__ positions, // [batch_size, seq_len] or // [num_tokens] scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads, // head_size] or [num_tokens, num_heads, // head_size] scalar_t* __restrict__ key, // nullptr or // [batch_size, seq_len, num_kv_heads, // head_size] or [num_tokens, num_kv_heads, // head_size] const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim // // 2] const int rot_dim, const int64_t query_stride, const int64_t key_stride, const int64_t head_stride, const int num_heads, const int num_kv_heads, const int head_size) { // Each thread block is responsible for one token. const int token_idx = blockIdx.x; int64_t pos = positions[token_idx]; const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim; apply_rotary_embedding( query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim, token_idx, query_stride, key_stride, head_stride); } // Helper struct to launch kernel template void launch_kernel( const int64_t* positions_data_ptr, void* query_ptr, void* key_ptr, const void* cos_sin_cache_ptr, int rot_dim, int64_t query_stride, int64_t key_stride, int64_t head_stride, int num_heads, int num_kv_heads, int head_size, dim3 grid, dim3 block, const cudaStream_t stream) { rotary_embedding_kernel<<>>( positions_data_ptr, static_cast(query_ptr), static_cast(key_ptr), static_cast(cos_sin_cache_ptr), rot_dim, query_stride, key_stride, head_stride, num_heads, num_kv_heads, head_size); }; // Helper macro to reduce repetition #define DISPATCH_DTYPE(DTYPE_CODE, DTYPE_BITS, IS_NEOX, ...) \ if (DTYPE_CODE == kDLFloat && DTYPE_BITS == 32) { \ launch_kernel(__VA_ARGS__); \ } else if (DTYPE_CODE == kDLFloat && DTYPE_BITS == 16) { \ launch_kernel(__VA_ARGS__); \ } else if (DTYPE_CODE == kDLBfloat && DTYPE_BITS == 16) { \ launch_kernel(__VA_ARGS__); \ } else { \ RuntimeCheck( \ false, "Unsupported data type for rotary embedding. Only float32, float16, and bfloat16 are supported."); \ } // Helper function to dispatch based on data type template void dispatch_by_dtype( const int64_t* positions_data_ptr, DLDataType query_dtype, void* query_ptr, void* key_ptr, void* cos_sin_cache_ptr, int rot_dim, int64_t query_stride, int64_t key_stride, int64_t head_stride, int num_heads, int num_kv_heads, int head_size, dim3 grid, dim3 block, const cudaStream_t stream) { using namespace host; DISPATCH_DTYPE( query_dtype.code, query_dtype.bits, IS_NEOX, positions_data_ptr, query_ptr, key_ptr, cos_sin_cache_ptr, rot_dim, query_stride, key_stride, head_stride, num_heads, num_kv_heads, head_size, grid, block, stream); } struct RotaryEmbeddingKernel { static void run(tvm::ffi::TensorView positions, // [batch_size, seq_len] or [num_tokens] tvm::ffi::TensorView query, // [batch_size, seq_len, num_heads * head_size] or // [num_tokens, num_heads * head_size] or // [batch_size, seq_len, num_heads, head_size] or // [num_tokens, num_heads, head_size] tvm::ffi::Optional key, // null or // [batch_size, seq_len, num_kv_heads * head_size] or // [num_tokens, num_kv_heads * head_size] or // [batch_size, seq_len, num_heads, head_size] or // [num_tokens, num_heads, head_size] int64_t head_size, tvm::ffi::TensorView cos_sin_cache, // [max_position, rot_dim] bool is_neox) { using namespace host; // num_tokens = batch_size * seq_len int64_t num_tokens = positions.numel(); int32_t positions_ndim = positions.ndim(); // Make sure num_tokens dim is consistent across positions, query, and key RuntimeCheck( positions_ndim == 1 || positions_ndim == 2, "positions must have shape [num_tokens] or [batch_size, seq_len]"); if (positions_ndim == 1) { RuntimeCheck( query.size(0) == positions.size(0) && (!key.has_value() || key.value().size(0) == positions.size(0)), "query, key and positions must have the same number of tokens"); } if (positions_ndim == 2) { RuntimeCheck( query.size(0) == positions.size(0) && (!key.has_value() || key.value().size(0) == positions.size(0)) && query.size(1) == positions.size(1) && (!key.has_value() || key.value().size(1) == positions.size(1)), "query, key and positions must have the same batch_size and seq_len"); } // Make sure head_size is valid for query and key // hidden_size = num_heads * head_size int query_hidden_size = query.numel() / num_tokens; int key_hidden_size = key.has_value() ? key.value().numel() / num_tokens : 0; RuntimeCheck(query_hidden_size % head_size == 0); RuntimeCheck(key_hidden_size % head_size == 0); // Make sure query and key have consistent number of heads int num_heads = query_hidden_size / head_size; int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads; RuntimeCheck(num_heads % num_kv_heads == 0); int rot_dim = cos_sin_cache.size(1); int seq_dim_idx = positions_ndim - 1; int64_t query_stride = query.stride(seq_dim_idx); int64_t key_stride = key.has_value() ? key.value().stride(seq_dim_idx) : 0; // Determine head stride: for [*, heads, head_size] use stride of last dim; // for flat [*, heads*head_size], heads blocks are contiguous of size // head_size int query_ndim = query.dim(); int64_t head_stride = (query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size; dim3 grid(num_tokens); dim3 block(std::min(num_heads * rot_dim / 2, 512)); auto device = query.device(); const cudaStream_t stream = LaunchKernel::resolve_device(device); auto positions_data_ptr = static_cast(positions.data_ptr()); if (is_neox) { dispatch_by_dtype( positions_data_ptr, query.dtype(), query.data_ptr(), key.has_value() ? key.value().data_ptr() : nullptr, cos_sin_cache.data_ptr(), rot_dim, query_stride, key_stride, head_stride, num_heads, num_kv_heads, head_size, grid, block, stream); } else { dispatch_by_dtype( positions_data_ptr, query.dtype(), query.data_ptr(), key.has_value() ? key.value().data_ptr() : nullptr, cos_sin_cache.data_ptr(), rot_dim, query_stride, key_stride, head_stride, num_heads, num_kv_heads, head_size, grid, block, stream); } } }; } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh ================================================ #include #include #include #include #include #include #include #include #include #include #include #include namespace { struct QKNormParams { void* __restrict__ q; void* __restrict__ k; // k is offset by (-num_qo_heads * head_dim) elements int64_t q_stride; int64_t k_stride; uint32_t num_qo_heads; uint32_t num_kv_heads; float eps; const void* __restrict__ q_weight; const void* __restrict__ k_weight; uint32_t num_tokens; }; constexpr uint32_t kWarpsPerBlock = 4; constexpr uint32_t kThreadsPerBlock = kWarpsPerBlock * device::kWarpThreads; // Warp-level kernel for head_dim <= 256 template __global__ void fused_qknorm_warp(const QKNormParams __grid_constant__ params) { using namespace device; using Storage = norm::StorageType; static_assert(sizeof(Float) == 2, "Only support FP16/BF16"); const auto& [q, k, q_stride, k_stride, num_qo_heads, num_kv_heads, eps, q_weight, k_weight, num_tokens] = params; const auto num_blks = gridDim.x; const auto num_workers = num_blks * kWarpsPerBlock; const auto num_q_and_k_heads = num_qo_heads + num_kv_heads; const auto num_works = num_q_and_k_heads * num_tokens; const auto start_worker_id = blockIdx.x * kWarpsPerBlock + threadIdx.x / kWarpThreads; const auto gmem = tile::Memory::warp(); PDLWaitPrimary(); // wait for primary kernel for (auto idx = start_worker_id; idx < num_works; idx += num_workers) { const int64_t token_id = idx / num_q_and_k_heads; const int64_t head_id = idx % num_q_and_k_heads; const auto load_q = head_id < num_qo_heads; const auto input = load_q ? pointer::offset(q, 2 * (token_id * q_stride + head_id * kHeadDim)) : pointer::offset(k, 2 * (token_id * k_stride + head_id * kHeadDim)); const auto weight = load_q ? q_weight : k_weight; const auto input_vec = gmem.load(input); const auto weight_vec = gmem.load(weight); const auto output_vec = norm::apply_norm_warp(input_vec, weight_vec, eps); gmem.store(input, output_vec); } PDLTriggerSecondary(); // launch secondary kernel } // For CTA level, used for head_dim > 256 (512,1024) template __global__ void fused_qknorm_cta(const QKNormParams __grid_constant__ params) { using namespace device; using Storage = norm::StorageType; constexpr auto kNumThreads = host::norm::get_cta_threads(); constexpr auto kNumWarps = kNumThreads / kWarpThreads; static_assert(sizeof(Float) == 2, "Only support FP16/BF16"); const auto& [q, k, q_stride, k_stride, num_qo_heads, num_kv_heads, eps, q_weight, k_weight, num_tokens] = params; const auto num_q_and_k_heads = num_qo_heads + num_kv_heads; const auto num_works = num_q_and_k_heads * num_tokens; const auto gmem = tile::Memory::cta(kNumThreads); __shared__ float smem[norm::kSmemBufferSize]; PDLWaitPrimary(); // wait for primary kernel for (auto idx = blockIdx.x; idx < num_works; idx += gridDim.x) { const int64_t token_id = idx / num_q_and_k_heads; const int64_t head_id = idx % num_q_and_k_heads; const auto load_q = head_id < num_qo_heads; const auto input = load_q ? pointer::offset(q, 2 * (token_id * q_stride + head_id * kHeadDim)) : pointer::offset(k, 2 * (token_id * k_stride + head_id * kHeadDim)); const auto weight = load_q ? q_weight : k_weight; const auto input_vec = gmem.load(input); const auto weight_vec = gmem.load(weight); const auto output_vec = norm::apply_norm_cta(input_vec, weight_vec, eps, smem, kNumWarps); gmem.store(input, output_vec); } PDLTriggerSecondary(); // launch secondary kernel } // Warp-level kernel struct for head_dim <= 256 template struct QKNormKernelWarp { static_assert(std::is_same_v || std::is_same_v); static_assert(!host::norm::should_use_cta(), "Use QKNormKernelCTA for head_dim > 256"); static constexpr auto kernel = fused_qknorm_warp; static void run(const tvm::ffi::TensorView q, const tvm::ffi::TensorView k, const tvm::ffi::TensorView q_weight, const tvm::ffi::TensorView k_weight, float eps) { using namespace host; auto N = SymbolicSize{"num_tokens"}; auto Q = SymbolicSize{"num_qo_heads"}; auto K = SymbolicSize{"num_kv_heads"}; auto D = SymbolicSize{"head_dim"}; auto Sq = SymbolicSize{"q_stride"}; auto Sk = SymbolicSize{"k_stride"}; auto device = SymbolicDevice{}; D.set_value(kHeadDim); device.set_options(); TensorMatcher({N, Q, D}) // q input .with_strides({Sq, D, 1}) .with_dtype() .with_device(device) .verify(q); TensorMatcher({N, K, D}) // k input .with_strides({Sk, D, 1}) .with_dtype() .with_device(device) .verify(k); TensorMatcher({D}) // weight .with_dtype() .with_device(device) .verify(q_weight) .verify(k_weight); const auto num_tokens = static_cast(N.unwrap()); const auto num_qo_heads = static_cast(Q.unwrap()); const auto num_kv_heads = static_cast(K.unwrap()); // NOTE: we offset the k here to reduce computation cost in the kernel const auto params = QKNormParams{ .q = q.data_ptr(), .k = pointer::offset(k.data_ptr(), -2 * static_cast(num_qo_heads) * kHeadDim), .q_stride = static_cast(Sq.unwrap()), .k_stride = static_cast(Sk.unwrap()), .num_qo_heads = num_qo_heads, .num_kv_heads = num_kv_heads, .eps = eps, .q_weight = q_weight.data_ptr(), .k_weight = k_weight.data_ptr(), .num_tokens = num_tokens, }; static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kThreadsPerBlock); static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id); // choose kernel based on dtype const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens; const auto needed_blocks = div_ceil(num_works, kWarpsPerBlock); // we use persistent kernel, which limit the number of blocks to reduce overhead const auto num_blocks = std::min(kNumSM * max_occupancy, needed_blocks); LaunchKernel(num_blocks, kThreadsPerBlock, device.unwrap()) // .enable_pdl(kUsePDL)(kernel, params); } }; // This goes with fused_qknorm_cta template struct QKNormKernelCTA { static_assert(std::is_same_v || std::is_same_v); static_assert(host::norm::should_use_cta(), "Use QKNormKernelWarp for head_dim <= 256"); static constexpr auto kernel = fused_qknorm_cta; static constexpr auto kNumThreads = host::norm::get_cta_threads(); static void run(const tvm::ffi::TensorView q, const tvm::ffi::TensorView k, const tvm::ffi::TensorView q_weight, const tvm::ffi::TensorView k_weight, float eps) { using namespace host; auto N = SymbolicSize{"num_tokens"}; auto Q = SymbolicSize{"num_qo_heads"}; auto K = SymbolicSize{"num_kv_heads"}; auto D = SymbolicSize{"head_dim"}; auto Sq = SymbolicSize{"q_stride"}; auto Sk = SymbolicSize{"k_stride"}; auto device = SymbolicDevice{}; D.set_value(kHeadDim); device.set_options(); TensorMatcher({N, Q, D}) // q input .with_strides({Sq, D, 1}) .with_dtype() .with_device(device) .verify(q); TensorMatcher({N, K, D}) // k input .with_strides({Sk, D, 1}) .with_dtype() .with_device(device) .verify(k); TensorMatcher({D}) // weight .with_dtype() .with_device(device) .verify(q_weight) .verify(k_weight); const auto num_tokens = static_cast(N.unwrap()); const auto num_qo_heads = static_cast(Q.unwrap()); const auto num_kv_heads = static_cast(K.unwrap()); // NOTE: we offset the k here to reduce computation cost in the kernel const auto params = QKNormParams{ .q = q.data_ptr(), .k = pointer::offset(k.data_ptr(), -2 * static_cast(num_qo_heads) * kHeadDim), .q_stride = static_cast(Sq.unwrap()), .k_stride = static_cast(Sk.unwrap()), .num_qo_heads = num_qo_heads, .num_kv_heads = num_kv_heads, .eps = eps, .q_weight = q_weight.data_ptr(), .k_weight = k_weight.data_ptr(), .num_tokens = num_tokens, }; static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kNumThreads); static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id); const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens; // we use persistent kernel, which limit the number of blocks to reduce overhead const auto num_blocks = std::min(num_works, max_occupancy * kNumSM); LaunchKernel(num_blocks, kNumThreads, device.unwrap()) // .enable_pdl(kUsePDL)(kernel, params); } }; // Unified dispatch: select warp or CTA kernel based on head_dim template using QKNormKernel = std::conditional_t< host::norm::should_use_cta(), QKNormKernelCTA, QKNormKernelWarp>; } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/elementwise/qknorm_across_heads.cuh ================================================ #include #include #include #include #include #include #include #include #include #include #include namespace { template struct VecTypeTrait; template <> struct VecTypeTrait { using packed_t = packed_t; using vec_t = device::AlignedVector; }; template <> struct VecTypeTrait { using packed_t = packed_t; using vec_t = device::AlignedVector; }; template <> struct VecTypeTrait { using packed_t = packed_t; using vec_t = device::AlignedVector; }; template <> struct VecTypeTrait { using packed_t = packed_t; using vec_t = device::AlignedVector; }; template SGL_DEVICE packed_t rms(packed_t& val, packed_t& weight, float rsqrt_square_sum) { float2 valf = device::cast(val); float2 weightf = device::cast(weight); return device::cast( make_float2(valf.x * weightf.x * rsqrt_square_sum, valf.y * weightf.y * rsqrt_square_sum)); } template __global__ void qknorm_across_heads_reg_kernel( T* __restrict__ q, T* __restrict__ k, const T* __restrict__ q_weight, const T* __restrict__ k_weight, int vec_hidden_size, float eps) { constexpr int inner_loop = VEC_SIZE_IN_BYTE == 16 ? 4 : 8; __shared__ float shared_memory[64]; // Used for CTA reduce, store both Q and K rsqrt using vec_t = typename VecTypeTrait::vec_t; using packed_t = typename VecTypeTrait::packed_t; vec_t v_q; // Save q vec_t v_k; // Save k vec_t v_q_weight; // Save q_weight vec_t v_k_weight; // Save k_weight vec_t v_q_out; // Save q output vec_t v_k_out; // Save k output auto token_id = blockIdx.x; float2 acc_square_q = make_float2(0.0f, 0.0f); // Sum of squares for q float2 acc_square_k = make_float2(0.0f, 0.0f); // Sum of squares for k if (threadIdx.x < vec_hidden_size) { // Compute address for q and k vec_t* p_q = reinterpret_cast(q) + token_id * vec_hidden_size; vec_t* p_k = reinterpret_cast(k) + token_id * vec_hidden_size; const vec_t* p_q_weight = reinterpret_cast(q_weight); const vec_t* p_k_weight = reinterpret_cast(k_weight); // Load data v_q = p_q[threadIdx.x]; v_k = p_k[threadIdx.x]; v_q_weight = p_q_weight[threadIdx.x]; v_k_weight = p_k_weight[threadIdx.x]; // Compute sum of squares for q for (int i = 0; i < inner_loop; i++) { float2 val = device::cast(v_q[i]); acc_square_q.x += val.x * val.x; acc_square_q.y += val.y * val.y; } // Compute sum of squares for k for (int i = 0; i < inner_loop; i++) { float2 val = device::cast(v_k[i]); acc_square_k.x += val.x * val.x; acc_square_k.y += val.y * val.y; } } auto cg_warp = cooperative_groups::tiled_partition<32>(cooperative_groups::this_thread_block()); float* buffer_q = shared_memory; // [0, 31] for Q float* buffer_k = shared_memory + 32; // [32, 63] for K // ========== Reduction phase: Compute rsqrt for both Q and K ========== // Step 0: Warp Reduce for Q float warp_sum_q = cooperative_groups::reduce(cg_warp, acc_square_q.x + acc_square_q.y, cooperative_groups::plus()); if (threadIdx.x % 32 == 0) { buffer_q[threadIdx.x / 32] = warp_sum_q; } // Step 0: Warp Reduce for K float warp_sum_k = cooperative_groups::reduce(cg_warp, acc_square_k.x + acc_square_k.y, cooperative_groups::plus()); if (threadIdx.x % 32 == 0) { buffer_k[threadIdx.x / 32] = warp_sum_k; } // Step 1: CTA Reduce for both Q and K __syncthreads(); if (threadIdx.x < 32) { // CTA Reduce for Q float cta_sum_q = cooperative_groups::reduce( cg_warp, (threadIdx.x < blockDim.x / 32) ? buffer_q[threadIdx.x] : 0.0f, cooperative_groups::plus()); buffer_q[threadIdx.x] = rsqrtf(eps + cta_sum_q * (1.0f / static_cast(vec_hidden_size * (VEC_SIZE_IN_BYTE / sizeof(T))))); // CTA Reduce for K float cta_sum_k = cooperative_groups::reduce( cg_warp, (threadIdx.x < blockDim.x / 32) ? buffer_k[threadIdx.x] : 0.0f, cooperative_groups::plus()); buffer_k[threadIdx.x] = rsqrtf(eps + cta_sum_k * (1.0f / static_cast(vec_hidden_size * (VEC_SIZE_IN_BYTE / sizeof(T))))); } __syncthreads(); // ========== Apply normalization phase: Compute and write back Q and K ========== if (threadIdx.x < vec_hidden_size) { // Apply RMSNorm for Q float rsqrt_q = buffer_q[threadIdx.x / 32]; for (int i = 0; i < inner_loop; i++) { v_q_out[i] = rms(v_q[i], v_q_weight[i], rsqrt_q); } vec_t* p_q_out = reinterpret_cast(q) + token_id * vec_hidden_size; p_q_out[threadIdx.x] = v_q_out; // Apply RMSNorm for K float rsqrt_k = buffer_k[threadIdx.x / 32]; for (int i = 0; i < inner_loop; i++) { v_k_out[i] = rms(v_k[i], v_k_weight[i], rsqrt_k); } vec_t* p_k_out = reinterpret_cast(k) + token_id * vec_hidden_size; p_k_out[threadIdx.x] = v_k_out; } } template struct QKNormAcrossHeadsKernel { static void run(const tvm::ffi::TensorView q, const tvm::ffi::TensorView k, const tvm::ffi::TensorView q_weight, const tvm::ffi::TensorView k_weight, float eps) { using namespace host; auto N = SymbolicSize{"num_tokens"}; auto D = SymbolicSize{"hidden_size"}; auto device = SymbolicDevice{}; device.set_options(); TensorMatcher({N, D}) // q .with_strides({D, 1}) .with_dtype() .with_device(device) .verify(q); TensorMatcher({N, D}) // k .with_strides({D, 1}) .with_dtype() .with_device(device) .verify(k); TensorMatcher({D}) // q_weight .with_dtype() .with_device(device) .verify(q_weight); TensorMatcher({D}) // k_weight .with_dtype() .with_device(device) .verify(k_weight); auto cc_major = host::runtime::get_cc_major(device.unwrap().device_id); int hidden_size = static_cast(D.unwrap()); if ((cc_major <= 9 && hidden_size <= 8192) || (cc_major >= 10 && hidden_size <= 12288)) { int max_vec_size_byte = cc_major >= 10 ? 32 : 16; int elements_in_vec = max_vec_size_byte / sizeof(DType); int vec_hidden_size = hidden_size / elements_in_vec; uint threads = (vec_hidden_size + 31) / 32 * 32; // Runtime check host::RuntimeCheck( hidden_size % elements_in_vec == 0, "hidden_size", hidden_size, " can not align to elements_in_vec ", elements_in_vec); // Launch single kernel for both q and k auto kernel = max_vec_size_byte == 32 ? qknorm_across_heads_reg_kernel : qknorm_across_heads_reg_kernel; LaunchKernel(static_cast(N.unwrap()), threads, device.unwrap()) .enable_pdl(false)( kernel, reinterpret_cast(q.data_ptr()), reinterpret_cast(k.data_ptr()), reinterpret_cast(q_weight.data_ptr()), reinterpret_cast(k_weight.data_ptr()), vec_hidden_size, eps); } else { host::RuntimeCheck(false, "Large hidden_sizes are not supported for now."); } } }; } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh ================================================ #include #include #include #include #include #include #include namespace { struct RMSNormParams { const void* input; const void* __restrict__ weight; void* output; int64_t input_stride; int64_t output_stride; uint32_t num_tokens; float eps; }; template __global__ void rmsnorm_cta(const RMSNormParams __grid_constant__ params) { using namespace device; using Storage = norm::StorageType; constexpr auto kNumThreads = host::norm::get_cta_threads(); constexpr auto kNumWarps = kNumThreads / kWarpThreads; const auto& [input, weight_ptr, output, input_stride, output_stride, num_tokens, eps] = params; const auto gmem = tile::Memory::cta(kNumThreads); __shared__ float smem[norm::kSmemBufferSize]; PDLWaitPrimary(); // wait for primary kernel void* output_ptr = nullptr; Storage output_vec; for (uint32_t i = blockIdx.x; i < num_tokens; i += gridDim.x) { const auto input_ptr = pointer::offset(input, i * input_stride); const auto input_vec = gmem.load(input_ptr); const auto weight_vec = gmem.load(weight_ptr); if (output_ptr != nullptr) { gmem.store(output_ptr, output_vec); } output_ptr = pointer::offset(output, i * output_stride); output_vec = norm::apply_norm_cta(input_vec, weight_vec, eps, smem, kNumWarps); } gmem.store(output_ptr, output_vec); PDLTriggerSecondary(); // launch secondary kernel } template struct RMSNormKernel { static_assert(host::norm::should_use_cta(), "Hidden size invalid for RMSNorm"); static constexpr auto kernel = rmsnorm_cta; static void run(const tvm::ffi::TensorView input, const tvm::ffi::TensorView weight, const tvm::ffi::TensorView output, float eps) { using namespace host; auto N = SymbolicSize{"num_tokens"}; auto D = SymbolicSize{"hidden_size"}; auto SI = SymbolicSize{"input_stride"}; auto SO = SymbolicSize{"output_stride"}; auto device = SymbolicDevice{}; D.set_value(kDim); device.set_options(); TensorMatcher({N, D}) // input .with_strides({SI, 1}) .with_dtype() .with_device(device) .verify(input); TensorMatcher({D}) // weight .with_dtype() .with_device(device) .verify(weight); TensorMatcher({N, D}) // output .with_strides({SO, 1}) .with_dtype() .with_device(device) .verify(output); const auto num_tokens = static_cast(N.unwrap()); const auto params = RMSNormParams{ .input = input.data_ptr(), .weight = weight.data_ptr(), .output = output.data_ptr(), .input_stride = SI.unwrap(), .output_stride = SO.unwrap(), .num_tokens = num_tokens, .eps = eps, }; static constexpr auto kNumThreads = norm::get_cta_threads(); static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kNumThreads); static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id); const auto num_blocks = std::min(num_tokens, max_occupancy * kNumSM); LaunchKernel(num_blocks, kNumThreads, device.unwrap()) // .enable_pdl(kUsePDL)(kernel, params); } }; } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/elementwise/rope.cuh ================================================ #include #include #include #include #include #include #include #include #include namespace { struct FusedRopeParams { void* __restrict__ q_ptr; void* __restrict__ k_ptr; // NOTE: this k is pre-offset in host code to reduce computation in kernel const void* __restrict__ cos_sin_cache_ptr; const void* __restrict__ positions; int64_t q_stride_bytes; int64_t k_stride_bytes; int64_t head_stride_bytes; uint32_t num_qo_heads; uint32_t num_kv_heads; uint32_t num_tokens; }; struct FusedRopeStoreParams { FusedRopeParams base_params; void* v_ptr; void* __restrict__ k_cache; void* __restrict__ v_cache; const void* __restrict__ out_loc; int64_t v_stride_bytes; int64_t cache_stride_bytes; }; constexpr uint32_t kBlockSize = 128; [[maybe_unused]] constexpr auto next_pow2(uint32_t target, uint32_t factor = 1) { uint32_t power = 1; while (power * factor < target) power *= 2; return power; } template __global__ void fused_rope_kernel(const __grid_constant__ FusedRopeParams params) { using namespace device; constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float); constexpr int64_t kVecSize = next_pow2(kRopeDim, (2 * kWorkThreads * (1 + kIsNeox))); using DType2 = packed_t; using InputStorage = AlignedVector; constexpr int64_t kDimPerThread = kVecSize * 2 * (1 + kIsNeox); constexpr uint32_t kLaneCount = kRopeDim / kDimPerThread; static_assert(kRopeDim % kDimPerThread == 0 && kLaneCount <= kWorkThreads); const auto &[ q, k, cos_sin_cache_ptr, positions, // pointers q_stride_bytes, k_stride_bytes, head_stride_bytes, // strides num_qo_heads, num_kv_heads, num_tokens // dimensions ] = params; const auto num_blks = gridDim.x; constexpr auto kWorkersPerBlock = kBlockSize / kWorkThreads; const auto num_workers = num_blks * kWorkersPerBlock; const auto num_q_and_k_heads = num_qo_heads + num_kv_heads; const auto num_works = num_q_and_k_heads * num_tokens; const auto start_worker_id = (blockIdx.x * kBlockSize + threadIdx.x) / kWorkThreads; const auto cos_cache_ptr = cos_sin_cache_ptr; const auto sin_cache_ptr = pointer::offset(cos_sin_cache_ptr, kCosSinStrideBytes / 2); uint32_t lane_id = threadIdx.x % kWorkThreads; if constexpr (kLaneCount < kWorkThreads) { if (lane_id >= kLaneCount) return; } PDLWaitPrimary(); for (auto idx = start_worker_id; idx < num_works; idx += num_workers) { const int64_t token_id = idx / num_q_and_k_heads; const int64_t head_id = idx % num_q_and_k_heads; const auto pos = static_cast(positions)[token_id]; const auto load_q = head_id < num_qo_heads; const auto input_ = load_q ? pointer::offset(q, token_id * q_stride_bytes) // : pointer::offset(k, token_id * k_stride_bytes); const auto input = pointer::offset(input_, head_id * head_stride_bytes); const auto cos_ptr = pointer::offset(cos_cache_ptr, pos * kCosSinStrideBytes); const auto sin_ptr = pointer::offset(sin_cache_ptr, pos * kCosSinStrideBytes); if constexpr (kIsNeox) { using CacheStorage = AlignedVector; const auto input_x = input; const auto input_y = pointer::offset(input, (kRopeDim / 2) * sizeof(DType)); auto input_vec_x = load_as(input_x, lane_id); auto input_vec_y = load_as(input_y, lane_id); const auto cos_pair = load_as(cos_ptr, lane_id); const auto sin_pair = load_as(sin_ptr, lane_id); #pragma unroll for (int64_t j = 0; j < kVecSize; ++j) { const auto [x0, x1] = cast(input_vec_x[j]); const auto [y0, y1] = cast(input_vec_y[j]); const auto [cos_0, cos_1] = cos_pair[j]; const auto [sin_0, sin_1] = sin_pair[j]; const auto out_x0 = x0 * cos_0 - y0 * sin_0; const auto out_y0 = x0 * sin_0 + y0 * cos_0; const auto out_x1 = x1 * cos_1 - y1 * sin_1; const auto out_y1 = x1 * sin_1 + y1 * cos_1; input_vec_x[j] = cast({out_x0, out_x1}); input_vec_y[j] = cast({out_y0, out_y1}); } store_as(input_x, input_vec_x, lane_id); store_as(input_y, input_vec_y, lane_id); } else { using CacheStorage = AlignedVector; auto input_vec = load_as(input, lane_id); const auto cos_vec = load_as(cos_ptr, lane_id); const auto sin_vec = load_as(sin_ptr, lane_id); #pragma unroll for (int64_t j = 0; j < kVecSize; ++j) { const auto [x, y] = cast(input_vec[j]); const auto cos = cos_vec[j]; const auto sin = sin_vec[j]; const auto out_x = x * cos - y * sin; const auto out_y = x * sin + y * cos; input_vec[j] = cast({out_x, out_y}); } store_as(input, input_vec, lane_id); } } PDLTriggerSecondary(); } template __global__ void fused_rope_store_kernel(const __grid_constant__ FusedRopeStoreParams params) { using namespace device; constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float); constexpr int64_t kVecSize = kRopeDim / (2 * kWorkThreads * (1 + kIsNeox)); using DType2 = packed_t; using InputStorage = AlignedVector; constexpr int64_t kDimPerThread = kVecSize * 2 * (1 + kIsNeox); static_assert(kRopeDim == kDimPerThread * kWorkThreads); const auto& [base_params, v_ptr, k_cache, v_cache, out_loc, v_stride_bytes, cache_stride_bytes] = params; const auto &[ q, k, cos_sin_cache_ptr, positions, // pointers q_stride_bytes, k_stride_bytes, head_stride_bytes, // strides num_qo_heads, num_kv_heads, num_tokens // dimensions ] = base_params; const auto num_blks = gridDim.x; constexpr auto kWorkersPerBlock = kBlockSize / kWorkThreads; const auto num_workers = num_blks * kWorkersPerBlock; const auto num_q_and_k_heads = num_qo_heads + num_kv_heads; const auto num_works = num_q_and_k_heads * num_tokens; const auto num_extra_works = num_kv_heads * num_tokens; // rope works + v store works const auto start_worker_id = (blockIdx.x * kBlockSize + threadIdx.x) / kWorkThreads; const auto lane_id = threadIdx.x % kWorkThreads; const auto cos_cache_ptr = cos_sin_cache_ptr; const auto sin_cache_ptr = pointer::offset(cos_sin_cache_ptr, kCosSinStrideBytes / 2); auto idx = start_worker_id; PDLWaitPrimary(); // in this case, head_dim = rope_dim must be true __builtin_assume(head_stride_bytes == kRopeDim * sizeof(DType)); for (; idx < num_works; idx += num_workers) { const int64_t token_id = idx / num_q_and_k_heads; const int64_t head_id = idx % num_q_and_k_heads; const auto pos = static_cast(positions)[token_id]; const auto loc = static_cast(out_loc)[token_id]; const auto load_q = head_id < num_qo_heads; const auto input_ = load_q ? pointer::offset(q, token_id * q_stride_bytes) // : pointer::offset(k, token_id * k_stride_bytes); const auto input = pointer::offset(input_, head_id * head_stride_bytes); const auto cos_ptr = pointer::offset(cos_cache_ptr, pos * kCosSinStrideBytes); const auto sin_ptr = pointer::offset(sin_cache_ptr, pos * kCosSinStrideBytes); if constexpr (kIsNeox) { using CacheStorage = AlignedVector; const auto input_x = input; const auto input_y = pointer::offset(input, (kRopeDim / 2) * sizeof(DType)); auto input_vec_x = load_as(input_x, lane_id); auto input_vec_y = load_as(input_y, lane_id); const auto cos_pair = load_as(cos_ptr, lane_id); const auto sin_pair = load_as(sin_ptr, lane_id); #pragma unroll for (int64_t j = 0; j < kVecSize; ++j) { const auto [x0, x1] = cast(input_vec_x[j]); const auto [y0, y1] = cast(input_vec_y[j]); const auto [cos_0, cos_1] = cos_pair[j]; const auto [sin_0, sin_1] = sin_pair[j]; const auto out_x0 = x0 * cos_0 - y0 * sin_0; const auto out_y0 = x0 * sin_0 + y0 * cos_0; const auto out_x1 = x1 * cos_1 - y1 * sin_1; const auto out_y1 = x1 * sin_1 + y1 * cos_1; input_vec_x[j] = cast({out_x0, out_x1}); input_vec_y[j] = cast({out_y0, out_y1}); } store_as(input, input_vec_x, lane_id); const auto input_y_out = pointer::offset(input, (kRopeDim / 2) * sizeof(DType)); store_as(input_y_out, input_vec_y, lane_id); if (!load_q) { const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes); store_as(k_out, input_vec_x, lane_id); const auto k_out_y = pointer::offset(k_out, (kRopeDim / 2) * sizeof(DType)); store_as(k_out_y, input_vec_y, lane_id); } } else { using CacheStorage = AlignedVector; auto input_vec = load_as(input, lane_id); const auto cos_vec = load_as(cos_ptr, lane_id); const auto sin_vec = load_as(sin_ptr, lane_id); #pragma unroll for (int64_t j = 0; j < kVecSize; ++j) { const auto [x, y] = cast(input_vec[j]); const auto cos = cos_vec[j]; const auto sin = sin_vec[j]; const auto out_x = x * cos - y * sin; const auto out_y = x * sin + y * cos; input_vec[j] = cast({out_x, out_y}); } store_as(input, input_vec, lane_id); if (!load_q) { const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes); store_as(k_out, input_vec, lane_id); } } } __syncwarp(); // to avoid warp divergence idx -= num_works; for (; idx < num_extra_works; idx += num_workers) { using VStorage = AlignedVector; const int64_t token_id = idx / num_kv_heads; const int64_t head_id = idx % num_kv_heads; const auto loc = static_cast(out_loc)[token_id]; const auto input = pointer::offset(v_ptr, token_id * v_stride_bytes, head_id * head_stride_bytes); const auto input_vec = load_as(input, lane_id); const auto output = pointer::offset(v_cache, loc * cache_stride_bytes, head_id * head_stride_bytes); store_as(output, input_vec, lane_id); } PDLTriggerSecondary(); } template struct FusedRopeKernel { static constexpr uint32_t kDimPerThread = std::gcd(16 / sizeof(DType), kRopeDim); static constexpr uint32_t kWorkThreads = next_pow2(kRopeDim, kDimPerThread); static constexpr bool kSupportFused = kWorkThreads * kDimPerThread == kRopeDim; static_assert(kRopeDim % kDimPerThread == 0); static_assert(kBlockSize % kWorkThreads == 0); template static constexpr auto _kernel_0 = fused_rope_kernel; template static constexpr auto _kernel_1 = fused_rope_store_kernel; static auto get_num_sm(DLDevice device) { static const auto kNumSM = host::runtime::get_sm_count(device.device_id); return kNumSM; } static void run(const tvm::ffi::TensorView q, const tvm::ffi::TensorView k, const tvm::ffi::TensorView cos_sin_cache, const tvm::ffi::TensorView positions) { using namespace host; auto N = SymbolicSize{"num_tokens"}; auto Q = SymbolicSize{"num_qo_heads"}; auto K = SymbolicSize{"num_kv_heads"}; auto D = SymbolicSize{"rope_dim"}; auto Dq = SymbolicSize{"q_stride"}; auto Dk = SymbolicSize{"k_stride"}; auto Dd = SymbolicSize{"head_stride"}; auto device = SymbolicDevice{}; auto id_type = SymbolicDType{}; D.set_value(kRopeDim); device.set_options(); TensorMatcher({N, Q, D}) // q input .with_strides({Dq, Dd, 1}) .with_dtype() .with_device(device) .verify(q); TensorMatcher({N, K, D}) // k input .with_strides({Dk, Dd, 1}) .with_dtype() .with_device(device) .verify(k); TensorMatcher({-1, D}) // cos_sin_cache .with_dtype() .with_device(device) .verify(cos_sin_cache); TensorMatcher({N}) // positions .with_dtype(id_type) .with_device(device) .verify(positions); const auto num_tokens = static_cast(N.unwrap()); const auto num_qo_heads = static_cast(Q.unwrap()); const auto num_kv_heads = static_cast(K.unwrap()); const auto q_stride_bytes = static_cast(Dq.unwrap() * sizeof(DType)); const auto k_stride_bytes = static_cast(Dk.unwrap() * sizeof(DType)); const auto head_stride_bytes = static_cast(Dd.unwrap() * sizeof(DType)); // NOTE: we offset the k here to reduce computation cost in the kernel const int64_t k_offset = static_cast(num_qo_heads) * head_stride_bytes; const auto params = FusedRopeParams{ .q_ptr = q.data_ptr(), .k_ptr = pointer::offset(k.data_ptr(), -k_offset), .cos_sin_cache_ptr = cos_sin_cache.data_ptr(), .positions = positions.data_ptr(), .q_stride_bytes = q_stride_bytes, .k_stride_bytes = k_stride_bytes, .head_stride_bytes = head_stride_bytes, .num_qo_heads = num_qo_heads, .num_kv_heads = num_kv_heads, .num_tokens = num_tokens, }; const auto is_int32 = id_type.is_type(); const auto kernel = is_int32 ? _kernel_0 : _kernel_0; const uint32_t kNumSM = get_num_sm(device.unwrap()); static const uint32_t kOccupancyTable[2] = { runtime::get_blocks_per_sm(_kernel_0, kBlockSize), runtime::get_blocks_per_sm(_kernel_0, kBlockSize), }; const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM; const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens; const auto needed_blocks = div_ceil(num_works, (kBlockSize / kWorkThreads)); const auto num_blocks = std::min(max_blocks, needed_blocks); LaunchKernel(num_blocks, kBlockSize, device.unwrap()) // .enable_pdl(kUsePDL)(kernel, params); } static void run_fused( const tvm::ffi::TensorView q, const tvm::ffi::TensorView k, const tvm::ffi::TensorView v, const tvm::ffi::TensorView k_cache, const tvm::ffi::TensorView v_cache, const tvm::ffi::TensorView cos_sin_cache, const tvm::ffi::TensorView positions, const tvm::ffi::TensorView out_loc) { if constexpr (kSupportFused) { return _run_fused_impl(q, k, v, k_cache, v_cache, cos_sin_cache, positions, out_loc); } else { host::Panic("Fused rope + store is not supported for rope_dim ", kRopeDim); } } static void _run_fused_impl( const tvm::ffi::TensorView q, const tvm::ffi::TensorView k, const tvm::ffi::TensorView v, const tvm::ffi::TensorView k_cache, const tvm::ffi::TensorView v_cache, const tvm::ffi::TensorView cos_sin_cache, const tvm::ffi::TensorView positions, const tvm::ffi::TensorView out_loc) { using namespace host; auto N = SymbolicSize{"num_tokens"}; auto Q = SymbolicSize{"num_qo_heads"}; auto K = SymbolicSize{"num_kv_heads"}; auto D = SymbolicSize{"rope_dim"}; auto R = SymbolicSize{"row_size"}; auto Dq = SymbolicSize{"q_stride"}; auto Dk = SymbolicSize{"k_stride"}; auto Dv = SymbolicSize{"v_stride"}; auto Dd = SymbolicSize{"head_stride"}; auto Dc = SymbolicSize{"cache_stride"}; auto device = SymbolicDevice{}; auto id_type = SymbolicDType{}; D.set_value(kRopeDim); device.set_options(); TensorMatcher({N, Q, D}) // q input .with_strides({Dq, Dd, 1}) .with_dtype() .with_device(device) .verify(q); TensorMatcher({N, K, D}) // k input .with_strides({Dk, Dd, 1}) .with_dtype() .with_device(device) .verify(k); TensorMatcher({N, K, D}) // v input .with_strides({Dv, Dd, 1}) .with_dtype() .with_device(device) .verify(v); TensorMatcher({-1, D}) // cos_sin_cache .with_dtype() .with_device(device) .verify(cos_sin_cache); TensorMatcher({N}) // positions, out_loc .with_dtype(id_type) .with_device(device) .verify(positions) .verify(out_loc); TensorMatcher({-1, R}) // k_cache .with_strides({Dc, 1}) .with_dtype() .with_device(device) .verify(k_cache) .verify(v_cache); const auto num_tokens = static_cast(N.unwrap()); const auto num_qo_heads = static_cast(Q.unwrap()); const auto num_kv_heads = static_cast(K.unwrap()); const auto q_stride_bytes = static_cast(Dq.unwrap() * sizeof(DType)); const auto k_stride_bytes = static_cast(Dk.unwrap() * sizeof(DType)); const auto head_stride = Dd.unwrap(); const auto row_dim = R.unwrap(); const auto head_stride_bytes = static_cast(Dd.unwrap() * sizeof(DType)); RuntimeCheck(kRopeDim == head_stride, "rope_dim ", kRopeDim, " should = head_stride ", head_stride); RuntimeCheck(num_kv_heads * kRopeDim == row_dim, "invalid kvcache"); // NOTE: we offset the k here to reduce computation cost in the kernel const int64_t k_offset = static_cast(num_qo_heads) * head_stride_bytes; const auto params = FusedRopeParams{ .q_ptr = q.data_ptr(), .k_ptr = pointer::offset(k.data_ptr(), -k_offset), .cos_sin_cache_ptr = cos_sin_cache.data_ptr(), .positions = positions.data_ptr(), .q_stride_bytes = q_stride_bytes, .k_stride_bytes = k_stride_bytes, .head_stride_bytes = head_stride_bytes, .num_qo_heads = num_qo_heads, .num_kv_heads = num_kv_heads, .num_tokens = num_tokens, }; const auto v_stride_bytes = static_cast(Dv.unwrap() * sizeof(DType)); const auto cache_stride_bytes = static_cast(Dc.unwrap() * sizeof(DType)); const auto store_params = FusedRopeStoreParams{ .base_params = params, .v_ptr = v.data_ptr(), .k_cache = pointer::offset(k_cache.data_ptr(), -k_offset), .v_cache = v_cache.data_ptr(), .out_loc = out_loc.data_ptr(), .v_stride_bytes = v_stride_bytes, .cache_stride_bytes = cache_stride_bytes, }; const auto is_int32 = id_type.is_type(); const auto kernel = is_int32 ? _kernel_1 : _kernel_1; const uint32_t kNumSM = get_num_sm(device.unwrap()); static const uint32_t kOccupancyTable[2] = { runtime::get_blocks_per_sm(_kernel_1, kBlockSize), runtime::get_blocks_per_sm(_kernel_1, kBlockSize), }; const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM; // rope works for q+k heads, plus v store works for kv heads const auto num_total_works = (num_qo_heads + 2 * num_kv_heads) * num_tokens; const auto needed_blocks = div_ceil(num_total_works, (kBlockSize / kWorkThreads)); const auto num_blocks = std::min(max_blocks, needed_blocks); LaunchKernel(num_blocks, kBlockSize, device.unwrap()) // .enable_pdl(kUsePDL)(kernel, store_params); } }; } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py ================================================ from pathlib import Path import numpy as np # From https://en.wikipedia.org/wiki/Paley_construction (construction II for q = 5) had_12_paley = """ +-++++++++++ --+-+-+-+-+- +++-++----++ +---+--+-++- +++++-++---- +-+---+--+-+ ++--+++-++-- +--++---+--+ ++----+++-++ +--+-++---+- ++++----+++- +-+--+-++--- """ # From http://neilsloane.com/hadamard/ had_12 = """ +----------- ++-+---+++-+ +++-+---+++- +-++-+---+++ ++-++-+---++ +++-++-+---+ ++++-++-+--- +-+++-++-+-- +--+++-++-+- +---+++-++-+ ++---+++-++- +-+---+++-++ """ had_20_will = """ +----+----++--++-++- -+----+---+++---+-++ --+----+---+++-+-+-+ ---+----+---+++++-+- ----+----++--++-++-+ -+++++-----+--+++--+ +-+++-+---+-+--+++-- ++-++--+---+-+--+++- +++-+---+---+-+--+++ ++++-----++--+-+--++ --++-+-++-+-----++++ ---++-+-++-+---+-+++ +---++-+-+--+--++-++ ++---++-+----+-+++-+ -++---++-+----+++++- -+--+--++-+----+---- +-+-----++-+----+--- -+-+-+---+--+----+-- --+-+++------+----+- +--+--++------+----+ """ had_28_will = """ +------++----++-+--+-+--++-- -+-----+++-----+-+--+-+--++- --+-----+++---+-+-+----+--++ ---+-----+++---+-+-+-+--+--+ ----+-----+++---+-+-+++--+-- -----+-----++++--+-+--++--+- ------++----++-+--+-+--++--+ --++++-+-------++--+++-+--+- ---++++-+-----+-++--+-+-+--+ +---+++--+----++-++--+-+-+-- ++---++---+----++-++--+-+-+- +++---+----+----++-++--+-+-+ ++++--------+-+--++-++--+-+- -++++--------+++--++--+--+-+ -+-++-++--++--+--------++++- +-+-++--+--++--+--------++++ -+-+-++--+--++--+----+---+++ +-+-+-++--+--+---+---++---++ ++-+-+-++--+------+--+++---+ -++-+-+-++--+------+-++++--- +-++-+---++--+------+-++++-- -++--++-+-++-+++----++------ +-++--++-+-++-+++-----+----- ++-++---+-+-++-+++-----+---- -++-++-+-+-+-+--+++-----+--- --++-++++-+-+----+++-----+-- +--++-+-++-+-+----+++-----+- ++--++-+-++-+-+----++------+ """ had_40_tpal = """ +-------------------+------------------- ++-++----+-+-++++--+++-++----+-+-++++--+ +++-++----+-+-++++--+++-++----+-+-++++-- +-++-++----+-+-++++-+-++-++----+-+-++++- +--++-++----+-+-+++++--++-++----+-+-++++ ++--++-++----+-+-+++++--++-++----+-+-+++ +++--++-++----+-+-+++++--++-++----+-+-++ ++++--++-++----+-+-+++++--++-++----+-+-+ +++++--++-++----+-+-+++++--++-++----+-+- +-++++--++-++----+-++-++++--++-++----+-+ ++-++++--++-++----+-++-++++--++-++----+- +-+-++++--++-++----++-+-++++--++-++----+ ++-+-++++--++-++----++-+-++++--++-++---- +-+-+-++++--++-++---+-+-+-++++--++-++--- +--+-+-++++--++-++--+--+-+-++++--++-++-- +---+-+-++++--++-++-+---+-+-++++--++-++- +----+-+-++++--++-+++----+-+-++++--++-++ ++----+-+-++++--++-+++----+-+-++++--++-+ +++----+-+-++++--++-+++----+-+-++++--++- +-++----+-+-++++--+++-++----+-+-++++--++ +--------------------+++++++++++++++++++ ++-++----+-+-++++--+--+--++++-+-+----++- +++-++----+-+-++++-----+--++++-+-+----++ +-++-++----+-+-++++--+--+--++++-+-+----+ +--++-++----+-+-++++-++--+--++++-+-+---- ++--++-++----+-+-+++--++--+--++++-+-+--- +++--++-++----+-+-++---++--+--++++-+-+-- ++++--++-++----+-+-+----++--+--++++-+-+- +++++--++-++----+-+------++--+--++++-+-+ +-++++--++-++----+-+-+----++--+--++++-+- ++-++++--++-++----+---+----++--+--++++-+ +-+-++++--++-++----+-+-+----++--+--++++- ++-+-++++--++-++------+-+----++--+--++++ +-+-+-++++--++-++----+-+-+----++--+--+++ +--+-+-++++--++-++---++-+-+----++--+--++ +---+-+-++++--++-++--+++-+-+----++--+--+ +----+-+-++++--++-++-++++-+-+----++--+-- ++----+-+-++++--++-+--++++-+-+----++--+- +++----+-+-++++--++----++++-+-+----++--+ +-++----+-+-++++--++-+--++++-+-+----++-- """ header = """ /****************************************************************************** * Copyright (c) 2023, Tri Dao. ******************************************************************************/ // This file is auto-generated. See "code_gen.py"\n #pragma once """ template = """ __device__ __forceinline__ void hadamard_mult_thread_{N}(float x[{N}]) { float out[{N}]; {code} #pragma unroll for (int i = 0; i < {N}; i++) { x[i] = out[i]; } } """ def string_to_array(string): # Convert strings of + and - to bool arrays string = string.strip().replace("+", "1").replace("-", "-1").split() return np.stack( [ np.fromstring(" ".join(string[i]), dtype=np.int32, sep=" ") for i in range(len(string)) ] ) def array_code_gen(arr): N = arr.shape[0] assert arr.shape[0] == arr.shape[1] out = [] for i in range(N): out.append( f"out[{i}] = " + " ".join([f"{'+' if arr[i, j] == 1 else '-'} x[{j}]" for j in range(N)]) + ";" ) return template.replace("{N}", str(N)).replace("{code}", "\n ".join(out)) def main(): output_dir = Path(__file__).parent / "fast_hadamard_transform_special.h" output_dir.write_text( header + array_code_gen(string_to_array(had_12_paley)) + array_code_gen(string_to_array(had_20_will)) + array_code_gen(string_to_array(had_28_will)) + array_code_gen(string_to_array(had_40_tpal)) ) if __name__ == "__main__": main() ================================================ FILE: python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform.h ================================================ /****************************************************************************** * Copyright (c) 2023, Tri Dao. ******************************************************************************/ // Copied from https://github.com/sgl-project/fast-hadamard-transform #pragma once //////////////////////////////////////////////////////////////////////////////////////////////////// struct HadamardParamsBase { using index_t = int64_t; int batch, dim, log_N; index_t x_batch_stride; index_t out_batch_stride; float scale; // Common data pointers. void* __restrict__ x_ptr; void* __restrict__ out_ptr; }; ================================================ FILE: python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_common.h ================================================ /****************************************************************************** * Copyright (c) 2023, Tri Dao. ******************************************************************************/ // Copied from https://github.com/sgl-project/fast-hadamard-transform #pragma once #include #include #define FULL_MASK 0xffffffff //////////////////////////////////////////////////////////////////////////////////////////////////// struct uint8 { uint4 u; uint4 v; }; template struct BytesToType {}; template <> struct BytesToType<32> { using Type = uint8; static_assert(sizeof(Type) == 32); }; template <> struct BytesToType<16> { using Type = uint4; static_assert(sizeof(Type) == 16); }; template <> struct BytesToType<8> { using Type = uint64_t; static_assert(sizeof(Type) == 8); }; template <> struct BytesToType<4> { using Type = uint32_t; static_assert(sizeof(Type) == 4); }; template <> struct BytesToType<2> { using Type = uint16_t; static_assert(sizeof(Type) == 2); }; template <> struct BytesToType<1> { using Type = uint8_t; static_assert(sizeof(Type) == 1); }; //////////////////////////////////////////////////////////////////////////////////////////////////// template struct SumOp { __device__ inline T operator()(T const& x, T const& y) { return x + y; } }; template struct Allreduce { static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4); template static __device__ inline T run(T x, Operator& op) { constexpr int OFFSET = THREADS / 2; x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET)); return Allreduce::run(x, op); } }; template <> struct Allreduce<2> { template static __device__ inline T run(T x, Operator& op) { x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1)); return x; } }; //////////////////////////////////////////////////////////////////////////////////////////////////// // https://stackoverflow.com/questions/35311711/whats-the-right-way-to-compute-integral-base-2-logarithms-at-compile-time constexpr int cilog2(int val) { return val > 0 ? 1 + cilog2(val >> 1) : -1; } //////////////////////////////////////////////////////////////////////////////////////////////////// template __device__ __forceinline__ void hadamard_mult_thread(float x[kNChunks][1 << kLogN]) { constexpr int N = 1 << kLogN; #pragma unroll for (int i = 0; i < kLogN; ++i) { const int stride = 1 << i; #pragma unroll for (int j = 0; j < N / 2; ++j) { const int lo = j & (stride - 1); const int idx = (j - lo) * 2 + lo; #pragma unroll for (int c = 0; c < kNChunks; ++c) { const float a = x[c][idx]; const float b = x[c][idx + stride]; x[c][idx] = a + b; x[c][idx + stride] = a - b; } } } } template __device__ __forceinline__ void hadamard_mult_warp(float x[kNChunks][kNItems]) { constexpr int N = 1 << kLogWarpSize; int lane_id = threadIdx.x % N; #pragma unroll for (int step = kStepStart; step < kLogWarpSize; ++step) { const int lane_mask = 1 << step; const float sign = (lane_id & lane_mask) ? -1.f : 1.f; #pragma unroll for (int c = 0; c < kNChunks; ++c) { #pragma unroll for (int i = 0; i < kNItems; ++i) { float x_val_other = __shfl_xor_sync(FULL_MASK, x[c][i], lane_mask); x[c][i] = sign * x[c][i] + x_val_other; } } } } //////////////////////////////////////////////////////////////////////////////////////////////////// template inline __device__ void load_input(input_t* x, float x_vals[kNChunks][kNElts], int dim) { using vec_t = typename BytesToType::Type; input_t x_vals_load[kNChunks][kNElts] = {0}; #pragma unroll for (int c = 0; c < kNChunks; ++c) { if ((c * blockDim.x + threadIdx.x) * kNElts < dim) { reinterpret_cast(x_vals_load)[c] = reinterpret_cast(x)[c * blockDim.x + threadIdx.x]; } } #pragma unroll for (int c = 0; c < kNChunks; ++c) { #pragma unroll for (int i = 0; i < kNElts; ++i) { x_vals[c][i] = float(x_vals_load[c][i]); } } } template inline __device__ void store_output(output_t* out, float out_vals[kNChunks][kNElts], int dim, float scale = 1.f) { using vec_t = typename BytesToType::Type; output_t out_vals_store[kNChunks][kNElts]; #pragma unroll for (int c = 0; c < kNChunks; ++c) { #pragma unroll for (int i = 0; i < kNElts; ++i) { out_vals_store[c][i] = out_vals[c][i] * scale; } } #pragma unroll for (int c = 0; c < kNChunks; ++c) { if ((c * blockDim.x + threadIdx.x) * kNElts < dim) { reinterpret_cast(out)[c * blockDim.x + threadIdx.x] = reinterpret_cast(out_vals_store)[c]; } } } //////////////////////////////////////////////////////////////////////////////////////////////////// // Pre=true means the exchange before the hadamard_mult_warp, Pre=false means after. template inline __device__ void exchange_smem_pre(float x_vals[kNChunks][kNElts], vec_t* smem) { constexpr int kNThreads = kWarpSize * kNWarps; constexpr int kNExchangePerVec = kNElts / (sizeof(vec_t) / sizeof(float)); const int warp_id = threadIdx.x / kWarpSize; const int lane_id = threadIdx.x % kWarpSize; const int row_t = threadIdx.x % kNWarps; const int col_t = threadIdx.x / kNWarps; // We use the XOR swizzle trick (new_col = col ^ row) to avoid / reduce smem bank conflicts. #pragma unroll for (int c0 = 0; c0 < kNChunks / kChunksPerExchange; ++c0) { __syncthreads(); #pragma unroll for (int c1 = 0; c1 < kChunksPerExchange; ++c1) { #pragma unroll for (int r = 0; r < kNExchangePerVec; ++r) { smem [(c1 * kNExchangePerVec + r) * kNThreads + (Pre ? warp_id * kWarpSize + lane_id ^ warp_id : row_t * kWarpSize + col_t ^ row_t)] = reinterpret_cast(x_vals[c0 * kChunksPerExchange + c1])[r]; } } __syncthreads(); #pragma unroll for (int c1 = 0; c1 < kChunksPerExchange; ++c1) { #pragma unroll for (int r = 0; r < kNExchangePerVec; ++r) { reinterpret_cast(x_vals[c0 * kChunksPerExchange + c1])[r] = smem [(c1 * kNExchangePerVec + r) * kNThreads + (Pre ? row_t * kWarpSize + col_t ^ row_t : warp_id * kWarpSize + lane_id ^ warp_id)]; } } } } ================================================ FILE: python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_special.h ================================================ /****************************************************************************** * Copyright (c) 2023, Tri Dao. ******************************************************************************/ // Copied from https://github.com/sgl-project/fast-hadamard-transform // This file is auto-generated. See "code_gen.py" #pragma once __device__ __forceinline__ void hadamard_mult_thread_12(float x[12]) { float out[12]; out[0] = +x[0] - x[1] + x[2] + x[3] + x[4] + x[5] + x[6] + x[7] + x[8] + x[9] + x[10] + x[11]; out[1] = -x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] - x[11]; out[2] = +x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11]; out[3] = +x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] + x[7] - x[8] + x[9] + x[10] - x[11]; out[4] = +x[0] + x[1] + x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11]; out[5] = +x[0] - x[1] + x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] + x[9] - x[10] + x[11]; out[6] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11]; out[7] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] + x[11]; out[8] = +x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] + x[7] + x[8] - x[9] + x[10] + x[11]; out[9] = +x[0] - x[1] - x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11]; out[10] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11]; out[11] = +x[0] - x[1] + x[2] - x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11]; #pragma unroll for (int i = 0; i < 12; i++) { x[i] = out[i]; } } __device__ __forceinline__ void hadamard_mult_thread_20(float x[20]) { float out[20]; out[0] = +x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11] - x[12] - x[13] + x[14] + x[15] - x[16] + x[17] + x[18] - x[19]; out[1] = -x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] - x[9] + x[10] + x[11] + x[12] - x[13] - x[14] - x[15] + x[16] - x[17] + x[18] + x[19]; out[2] = -x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] - x[9] - x[10] + x[11] + x[12] + x[13] - x[14] + x[15] - x[16] + x[17] - x[18] + x[19]; out[3] = -x[0] - x[1] - x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] + x[13] + x[14] + x[15] + x[16] - x[17] + x[18] - x[19]; out[4] = -x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] + x[14] - x[15] + x[16] + x[17] - x[18] + x[19]; out[5] = -x[0] + x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] - x[13] + x[14] + x[15] + x[16] - x[17] - x[18] + x[19]; out[6] = +x[0] - x[1] + x[2] + x[3] + x[4] - x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] - x[14] + x[15] + x[16] + x[17] - x[18] - x[19]; out[7] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] + x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - x[14] - x[15] + x[16] + x[17] + x[18] - x[19]; out[8] = +x[0] + x[1] + x[2] - x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + x[14] - x[15] - x[16] + x[17] + x[18] + x[19]; out[9] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] - x[14] + x[15] - x[16] - x[17] + x[18] + x[19]; out[10] = -x[0] - x[1] + x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] + x[10] - x[11] - x[12] - x[13] - x[14] - x[15] + x[16] + x[17] + x[18] + x[19]; out[11] = -x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] + x[18] + x[19]; out[12] = +x[0] - x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] - x[11] + x[12] - x[13] - x[14] + x[15] + x[16] - x[17] + x[18] + x[19]; out[13] = +x[0] + x[1] - x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - x[14] + x[15] + x[16] + x[17] - x[18] + x[19]; out[14] = -x[0] + x[1] + x[2] - x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + x[14] + x[15] + x[16] + x[17] + x[18] - x[19]; out[15] = -x[0] + x[1] - x[2] - x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] - x[11] - x[12] - x[13] - x[14] + x[15] - x[16] - x[17] - x[18] - x[19]; out[16] = +x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] - x[13] - x[14] - x[15] + x[16] - x[17] - x[18] - x[19]; out[17] = -x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] - x[13] - x[14] - x[15] - x[16] + x[17] - x[18] - x[19]; out[18] = -x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - x[14] - x[15] - x[16] - x[17] + x[18] - x[19]; out[19] = +x[0] - x[1] - x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15] - x[16] - x[17] - x[18] + x[19]; #pragma unroll for (int i = 0; i < 20; i++) { x[i] = out[i]; } } __device__ __forceinline__ void hadamard_mult_thread_28(float x[28]) { float out[28]; out[0] = +x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] + x[14] - x[15] + x[16] - x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] + x[25] - x[26] - x[27]; out[1] = -x[0] + x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] - x[24] + x[25] + x[26] - x[27]; out[2] = -x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11] - x[12] - x[13] + x[14] - x[15] + x[16] - x[17] + x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] + x[26] + x[27]; out[3] = -x[0] - x[1] - x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] + x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] - x[25] - x[26] + x[27]; out[4] = -x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11] + x[12] - x[13] - x[14] - x[15] + x[16] - x[17] + x[18] - x[19] + x[20] + x[21] + x[22] - x[23] - x[24] + x[25] - x[26] - x[27]; out[5] = -x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] + x[12] + x[13] + x[14] - x[15] - x[16] + x[17] - x[18] + x[19] - x[20] - x[21] + x[22] + x[23] - x[24] - x[25] + x[26] - x[27]; out[6] = -x[0] - x[1] - x[2] - x[3] - x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] + x[13] - x[14] + x[15] - x[16] - x[17] + x[18] - x[19] + x[20] - x[21] - x[22] + x[23] + x[24] - x[25] - x[26] + x[27]; out[7] = -x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] + x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] - x[14] + x[15] + x[16] - x[17] - x[18] + x[19] + x[20] + x[21] - x[22] + x[23] - x[24] - x[25] + x[26] - x[27]; out[8] = -x[0] - x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15] + x[16] + x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] + x[24] - x[25] - x[26] + x[27]; out[9] = +x[0] - x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + x[14] + x[15] - x[16] + x[17] + x[18] - x[19] - x[20] + x[21] - x[22] + x[23] - x[24] + x[25] - x[26] - x[27]; out[10] = +x[0] + x[1] - x[2] - x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11] - x[12] - x[13] - x[14] + x[15] + x[16] - x[17] + x[18] + x[19] - x[20] - x[21] + x[22] - x[23] + x[24] - x[25] + x[26] - x[27]; out[11] = +x[0] + x[1] + x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] - x[13] - x[14] - x[15] + x[16] + x[17] - x[18] + x[19] + x[20] - x[21] - x[22] + x[23] - x[24] + x[25] - x[26] + x[27]; out[12] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + x[14] - x[15] - x[16] + x[17] + x[18] - x[19] + x[20] + x[21] - x[22] - x[23] + x[24] - x[25] + x[26] - x[27]; out[13] = -x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] + x[13] + x[14] + x[15] - x[16] - x[17] + x[18] + x[19] - x[20] - x[21] + x[22] - x[23] - x[24] + x[25] - x[26] + x[27]; out[14] = -x[0] + x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] - x[13] + x[14] - x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] + x[24] + x[25] + x[26] - x[27]; out[15] = +x[0] - x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] - x[14] + x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] - x[22] - x[23] + x[24] + x[25] + x[26] + x[27]; out[16] = -x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] + x[13] - x[14] - x[15] + x[16] - x[17] - x[18] - x[19] - x[20] + x[21] - x[22] - x[23] - x[24] + x[25] + x[26] + x[27]; out[17] = +x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] - x[11] - x[12] + x[13] - x[14] - x[15] - x[16] + x[17] - x[18] - x[19] - x[20] + x[21] + x[22] - x[23] - x[24] - x[25] + x[26] + x[27]; out[18] = +x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] + x[11] - x[12] - x[13] - x[14] - x[15] - x[16] - x[17] + x[18] - x[19] - x[20] + x[21] + x[22] + x[23] - x[24] - x[25] - x[26] + x[27]; out[19] = -x[0] + x[1] + x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] + x[12] - x[13] - x[14] - x[15] - x[16] - x[17] - x[18] + x[19] - x[20] + x[21] + x[22] + x[23] + x[24] - x[25] - x[26] - x[27]; out[20] = +x[0] - x[1] + x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] - x[14] - x[15] - x[16] - x[17] - x[18] - x[19] + x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] - x[27]; out[21] = -x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] - x[12] + x[13] + x[14] + x[15] - x[16] - x[17] - x[18] - x[19] + x[20] + x[21] - x[22] - x[23] - x[24] - x[25] - x[26] - x[27]; out[22] = +x[0] - x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] - x[13] + x[14] + x[15] + x[16] - x[17] - x[18] - x[19] - x[20] - x[21] + x[22] - x[23] - x[24] - x[25] - x[26] - x[27]; out[23] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] - x[14] + x[15] + x[16] + x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] - x[26] - x[27]; out[24] = -x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] - x[14] - x[15] + x[16] + x[17] + x[18] - x[19] - x[20] - x[21] - x[22] - x[23] + x[24] - x[25] - x[26] - x[27]; out[25] = -x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] + x[10] - x[11] + x[12] - x[13] - x[14] - x[15] - x[16] + x[17] + x[18] + x[19] - x[20] - x[21] - x[22] - x[23] - x[24] + x[25] - x[26] - x[27]; out[26] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] + x[13] - x[14] - x[15] - x[16] - x[17] + x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] + x[26] - x[27]; out[27] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] - x[11] + x[12] - x[13] + x[14] - x[15] - x[16] - x[17] - x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] - x[26] + x[27]; #pragma unroll for (int i = 0; i < 28; i++) { x[i] = out[i]; } } __device__ __forceinline__ void hadamard_mult_thread_40(float x[40]) { float out[40]; out[0] = +x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] - x[14] - x[15] - x[16] - x[17] - x[18] - x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] - x[26] - x[27] - x[28] - x[29] - x[30] - x[31] - x[32] - x[33] - x[34] - x[35] - x[36] - x[37] - x[38] - x[39]; out[1] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] + x[14] + x[15] + x[16] - x[17] - x[18] + x[19] + x[20] + x[21] - x[22] + x[23] + x[24] - x[25] - x[26] - x[27] - x[28] + x[29] - x[30] + x[31] - x[32] + x[33] + x[34] + x[35] + x[36] - x[37] - x[38] + x[39]; out[2] = +x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] + x[14] + x[15] + x[16] + x[17] - x[18] - x[19] + x[20] + x[21] + x[22] - x[23] + x[24] + x[25] - x[26] - x[27] - x[28] - x[29] + x[30] - x[31] + x[32] - x[33] + x[34] + x[35] + x[36] + x[37] - x[38] - x[39]; out[3] = +x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - x[14] + x[15] + x[16] + x[17] + x[18] - x[19] + x[20] - x[21] + x[22] + x[23] - x[24] + x[25] + x[26] - x[27] - x[28] - x[29] - x[30] + x[31] - x[32] + x[33] - x[34] + x[35] + x[36] + x[37] + x[38] - x[39]; out[4] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + x[14] - x[15] + x[16] + x[17] + x[18] + x[19] + x[20] - x[21] - x[22] + x[23] + x[24] - x[25] + x[26] + x[27] - x[28] - x[29] - x[30] - x[31] + x[32] - x[33] + x[34] - x[35] + x[36] + x[37] + x[38] + x[39]; out[5] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - x[14] + x[15] - x[16] + x[17] + x[18] + x[19] + x[20] + x[21] - x[22] - x[23] + x[24] + x[25] - x[26] + x[27] + x[28] - x[29] - x[30] - x[31] - x[32] + x[33] - x[34] + x[35] - x[36] + x[37] + x[38] + x[39]; out[6] = +x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15] + x[16] - x[17] + x[18] + x[19] + x[20] + x[21] + x[22] - x[23] - x[24] + x[25] + x[26] - x[27] + x[28] + x[29] - x[30] - x[31] - x[32] - x[33] + x[34] - x[35] + x[36] - x[37] + x[38] + x[39]; out[7] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] + x[10] - x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] - x[18] + x[19] + x[20] + x[21] + x[22] + x[23] - x[24] - x[25] + x[26] + x[27] - x[28] + x[29] + x[30] - x[31] - x[32] - x[33] - x[34] + x[35] - x[36] + x[37] - x[38] + x[39]; out[8] = +x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] - x[14] - x[15] + x[16] - x[17] + x[18] - x[19] + x[20] + x[21] + x[22] + x[23] + x[24] - x[25] - x[26] + x[27] + x[28] - x[29] + x[30] + x[31] - x[32] - x[33] - x[34] - x[35] + x[36] - x[37] + x[38] - x[39]; out[9] = +x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] - x[14] - x[15] - x[16] + x[17] - x[18] + x[19] + x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] - x[27] + x[28] + x[29] - x[30] + x[31] + x[32] - x[33] - x[34] - x[35] - x[36] + x[37] - x[38] + x[39]; out[10] = +x[0] + x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] - x[11] + x[12] + x[13] - x[14] - x[15] - x[16] - x[17] + x[18] - x[19] + x[20] + x[21] - x[22] + x[23] + x[24] + x[25] + x[26] - x[27] - x[28] + x[29] + x[30] - x[31] + x[32] + x[33] - x[34] - x[35] - x[36] - x[37] + x[38] - x[39]; out[11] = +x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] + x[13] + x[14] - x[15] - x[16] - x[17] - x[18] + x[19] + x[20] - x[21] + x[22] - x[23] + x[24] + x[25] + x[26] + x[27] - x[28] - x[29] + x[30] + x[31] - x[32] + x[33] + x[34] - x[35] - x[36] - x[37] - x[38] + x[39]; out[12] = +x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] + x[14] + x[15] - x[16] - x[17] - x[18] - x[19] + x[20] + x[21] - x[22] + x[23] - x[24] + x[25] + x[26] + x[27] + x[28] - x[29] - x[30] + x[31] + x[32] - x[33] + x[34] + x[35] - x[36] - x[37] - x[38] - x[39]; out[13] = +x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] - x[11] + x[12] + x[13] - x[14] + x[15] + x[16] - x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] + x[24] - x[25] + x[26] + x[27] + x[28] + x[29] - x[30] - x[31] + x[32] + x[33] - x[34] + x[35] + x[36] - x[37] - x[38] - x[39]; out[14] = +x[0] - x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] + x[14] - x[15] + x[16] + x[17] - x[18] - x[19] + x[20] - x[21] - x[22] + x[23] - x[24] + x[25] - x[26] + x[27] + x[28] + x[29] + x[30] - x[31] - x[32] + x[33] + x[34] - x[35] + x[36] + x[37] - x[38] - x[39]; out[15] = +x[0] - x[1] - x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] + x[10] + x[11] - x[12] - x[13] + x[14] + x[15] - x[16] + x[17] + x[18] - x[19] + x[20] - x[21] - x[22] - x[23] + x[24] - x[25] + x[26] - x[27] + x[28] + x[29] + x[30] + x[31] - x[32] - x[33] + x[34] + x[35] - x[36] + x[37] + x[38] - x[39]; out[16] = +x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] - x[13] - x[14] + x[15] + x[16] - x[17] + x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] + x[25] - x[26] + x[27] - x[28] + x[29] + x[30] + x[31] + x[32] - x[33] - x[34] + x[35] + x[36] - x[37] + x[38] + x[39]; out[17] = +x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] - x[14] - x[15] + x[16] + x[17] - x[18] + x[19] + x[20] + x[21] - x[22] - x[23] - x[24] - x[25] + x[26] - x[27] + x[28] - x[29] + x[30] + x[31] + x[32] + x[33] - x[34] - x[35] + x[36] + x[37] - x[38] + x[39]; out[18] = +x[0] + x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] + x[13] + x[14] - x[15] - x[16] + x[17] + x[18] - x[19] + x[20] + x[21] + x[22] - x[23] - x[24] - x[25] - x[26] + x[27] - x[28] + x[29] - x[30] + x[31] + x[32] + x[33] + x[34] - x[35] - x[36] + x[37] + x[38] - x[39]; out[19] = +x[0] - x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] + x[14] + x[15] - x[16] - x[17] + x[18] + x[19] + x[20] - x[21] + x[22] + x[23] - x[24] - x[25] - x[26] - x[27] + x[28] - x[29] + x[30] - x[31] + x[32] + x[33] + x[34] + x[35] - x[36] - x[37] + x[38] + x[39]; out[20] = +x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] - x[14] - x[15] - x[16] - x[17] - x[18] - x[19] - x[20] + x[21] + x[22] + x[23] + x[24] + x[25] + x[26] + x[27] + x[28] + x[29] + x[30] + x[31] + x[32] + x[33] + x[34] + x[35] + x[36] + x[37] + x[38] + x[39]; out[21] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] + x[14] + x[15] + x[16] - x[17] - x[18] + x[19] - x[20] - x[21] + x[22] - x[23] - x[24] + x[25] + x[26] + x[27] + x[28] - x[29] + x[30] - x[31] + x[32] - x[33] - x[34] - x[35] - x[36] + x[37] + x[38] - x[39]; out[22] = +x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] + x[14] + x[15] + x[16] + x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] + x[26] + x[27] + x[28] + x[29] - x[30] + x[31] - x[32] + x[33] - x[34] - x[35] - x[36] - x[37] + x[38] + x[39]; out[23] = +x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - x[14] + x[15] + x[16] + x[17] + x[18] - x[19] - x[20] + x[21] - x[22] - x[23] + x[24] - x[25] - x[26] + x[27] + x[28] + x[29] + x[30] - x[31] + x[32] - x[33] + x[34] - x[35] - x[36] - x[37] - x[38] + x[39]; out[24] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + x[14] - x[15] + x[16] + x[17] + x[18] + x[19] - x[20] + x[21] + x[22] - x[23] - x[24] + x[25] - x[26] - x[27] + x[28] + x[29] + x[30] + x[31] - x[32] + x[33] - x[34] + x[35] - x[36] - x[37] - x[38] - x[39]; out[25] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - x[14] + x[15] - x[16] + x[17] + x[18] + x[19] - x[20] - x[21] + x[22] + x[23] - x[24] - x[25] + x[26] - x[27] - x[28] + x[29] + x[30] + x[31] + x[32] - x[33] + x[34] - x[35] + x[36] - x[37] - x[38] - x[39]; out[26] = +x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15] + x[16] - x[17] + x[18] + x[19] - x[20] - x[21] - x[22] + x[23] + x[24] - x[25] - x[26] + x[27] - x[28] - x[29] + x[30] + x[31] + x[32] + x[33] - x[34] + x[35] - x[36] + x[37] - x[38] - x[39]; out[27] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] + x[10] - x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] - x[18] + x[19] - x[20] - x[21] - x[22] - x[23] + x[24] + x[25] - x[26] - x[27] + x[28] - x[29] - x[30] + x[31] + x[32] + x[33] + x[34] - x[35] + x[36] - x[37] + x[38] - x[39]; out[28] = +x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] - x[14] - x[15] + x[16] - x[17] + x[18] - x[19] - x[20] - x[21] - x[22] - x[23] - x[24] + x[25] + x[26] - x[27] - x[28] + x[29] - x[30] - x[31] + x[32] + x[33] + x[34] + x[35] - x[36] + x[37] - x[38] + x[39]; out[29] = +x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] - x[14] - x[15] - x[16] + x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] - x[24] - x[25] + x[26] + x[27] - x[28] - x[29] + x[30] - x[31] - x[32] + x[33] + x[34] + x[35] + x[36] - x[37] + x[38] - x[39]; out[30] = +x[0] + x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] - x[11] + x[12] + x[13] - x[14] - x[15] - x[16] - x[17] + x[18] - x[19] - x[20] - x[21] + x[22] - x[23] - x[24] - x[25] - x[26] + x[27] + x[28] - x[29] - x[30] + x[31] - x[32] - x[33] + x[34] + x[35] + x[36] + x[37] - x[38] + x[39]; out[31] = +x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] + x[13] + x[14] - x[15] - x[16] - x[17] - x[18] + x[19] - x[20] + x[21] - x[22] + x[23] - x[24] - x[25] - x[26] - x[27] + x[28] + x[29] - x[30] - x[31] + x[32] - x[33] - x[34] + x[35] + x[36] + x[37] + x[38] - x[39]; out[32] = +x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] + x[14] + x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] + x[22] - x[23] + x[24] - x[25] - x[26] - x[27] - x[28] + x[29] + x[30] - x[31] - x[32] + x[33] - x[34] - x[35] + x[36] + x[37] + x[38] + x[39]; out[33] = +x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] - x[11] + x[12] + x[13] - x[14] + x[15] + x[16] - x[17] - x[18] - x[19] - x[20] + x[21] - x[22] + x[23] - x[24] + x[25] - x[26] - x[27] - x[28] - x[29] + x[30] + x[31] - x[32] - x[33] + x[34] - x[35] - x[36] + x[37] + x[38] + x[39]; out[34] = +x[0] - x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] + x[14] - x[15] + x[16] + x[17] - x[18] - x[19] - x[20] + x[21] + x[22] - x[23] + x[24] - x[25] + x[26] - x[27] - x[28] - x[29] - x[30] + x[31] + x[32] - x[33] - x[34] + x[35] - x[36] - x[37] + x[38] + x[39]; out[35] = +x[0] - x[1] - x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] + x[10] + x[11] - x[12] - x[13] + x[14] + x[15] - x[16] + x[17] + x[18] - x[19] - x[20] + x[21] + x[22] + x[23] - x[24] + x[25] - x[26] + x[27] - x[28] - x[29] - x[30] - x[31] + x[32] + x[33] - x[34] - x[35] + x[36] - x[37] - x[38] + x[39]; out[36] = +x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] - x[13] - x[14] + x[15] + x[16] - x[17] + x[18] + x[19] - x[20] + x[21] + x[22] + x[23] + x[24] - x[25] + x[26] - x[27] + x[28] - x[29] - x[30] - x[31] - x[32] + x[33] + x[34] - x[35] - x[36] + x[37] - x[38] - x[39]; out[37] = +x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] - x[14] - x[15] + x[16] + x[17] - x[18] + x[19] - x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] + x[27] - x[28] + x[29] - x[30] - x[31] - x[32] - x[33] + x[34] + x[35] - x[36] - x[37] + x[38] - x[39]; out[38] = +x[0] + x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] + x[13] + x[14] - x[15] - x[16] + x[17] + x[18] - x[19] - x[20] - x[21] - x[22] + x[23] + x[24] + x[25] + x[26] - x[27] + x[28] - x[29] + x[30] - x[31] - x[32] - x[33] - x[34] + x[35] + x[36] - x[37] - x[38] + x[39]; out[39] = +x[0] - x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] + x[14] + x[15] - x[16] - x[17] + x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] + x[25] + x[26] + x[27] - x[28] + x[29] - x[30] + x[31] - x[32] - x[33] - x[34] - x[35] + x[36] + x[37] - x[38] - x[39]; #pragma unroll for (int i = 0; i < 40; i++) { x[i] = out[i]; } } ================================================ FILE: python/sglang/jit_kernel/csrc/fast-hadamard-transform/hadamard_jit.cuh ================================================ /****************************************************************************** * Copyright (c) 2023, Tri Dao. ******************************************************************************/ #pragma once #include #include #include #include #include "fast_hadamard_transform.h" #include "fast_hadamard_transform_common.h" #include "fast_hadamard_transform_special.h" #include "static_switch.h" #include #include #include namespace { using ::bf16_t; using ::fp16_t; using ::HadamardParamsBase; constexpr inline int ceil_log2(int val) { int log = 0; int p = 1; while (p < val) { p <<= 1; ++log; } return log; } template struct FastHadamardKernelTraits { using input_t = input_t_; static constexpr int kNThreads = kNThreads_; static constexpr int kLogN = kLogN_; static constexpr int N = 1 << kLogN; static constexpr int kNBytes = sizeof(input_t); static_assert(kNBytes == 2 || kNBytes == 4); static constexpr int kNElts = kNBytes == 4 ? 4 : 8; static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t); using vec_t = typename BytesToType::Type; static constexpr int kNChunks = N / (kNElts * kNThreads); static constexpr int kSmemExchangeSize = (N * 4) < (32 * 1024) ? (N * 4) : (32 * 1024); static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize; static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4); static constexpr int kSmemSize = kSmemExchangeSize; }; template struct FastHadamardMNKernelTraits { using input_t = input_t_; static constexpr int kNThreads = kNThreads_; static constexpr int kLogN = kLogN_; static constexpr int N = (1 << kLogN) * kMultiple; static_assert(N <= kMaxDim); static constexpr int kNBytes = sizeof(input_t); static_assert(kNBytes == 2 || kNBytes == 4); static constexpr int kNElts = 4; static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t); using vec_t = typename BytesToType::Type; static constexpr int kNChunks = N / (kNElts * kNThreads); static_assert(kNChunks == kMultiple); static constexpr int kSmemExchangeSize = (N * 4) < kMaxSmem ? (N * 4) : kMaxSmem; static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize; static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4); static constexpr int kSmemSize = kSmemExchangeSize; }; template using FastHadamard12NTraits = FastHadamardMNKernelTraits; template using FastHadamard20NTraits = FastHadamardMNKernelTraits; template using FastHadamard28NTraits = FastHadamardMNKernelTraits; template using FastHadamard40NTraits = FastHadamardMNKernelTraits; template SGL_DEVICE void hadamard_mult_thread_chunk_12(float x[kNChunks][12]) { #pragma unroll for (int c = 0; c < kNChunks; ++c) { hadamard_mult_thread_12(x[c]); } } template SGL_DEVICE void hadamard_mult_thread_chunk_20(float x[kNChunks][20]) { #pragma unroll for (int c = 0; c < kNChunks; ++c) { hadamard_mult_thread_20(x[c]); } } template SGL_DEVICE void hadamard_mult_thread_chunk_28(float x[kNChunks][28]) { #pragma unroll for (int c = 0; c < kNChunks; ++c) { hadamard_mult_thread_28(x[c]); } } template SGL_DEVICE void hadamard_mult_thread_chunk_40(float x[kNChunks][40]) { #pragma unroll for (int c = 0; c < kNChunks; ++c) { hadamard_mult_thread_40(x[c]); } } template __global__ __launch_bounds__(Ktraits::kNThreads) void fast_hadamard_transform_kernel(HadamardParamsBase params) { constexpr int kNThreads = Ktraits::kNThreads; constexpr int kNElts = Ktraits::kNElts; constexpr int kNExchangePerVec = Ktraits::kNExchangePerVec; constexpr int kNChunks = Ktraits::kNChunks; using input_t = typename Ktraits::input_t; using vec_t = typename Ktraits::vec_t; constexpr int kLogNElts = cilog2(Ktraits::kNElts); static_assert(1 << kLogNElts == kNElts, "kNElts must be a power of 2"); constexpr int kWarpSize = kNThreads < 32 ? kNThreads : 32; constexpr int kLogWarpSize = cilog2(kWarpSize); static_assert(1 << kLogWarpSize == kWarpSize, "Warp size must be a power of 2"); constexpr int kNWarps = kNThreads / kWarpSize; constexpr int kLogNWarps = cilog2(kNWarps); static_assert(1 << kLogNWarps == kNWarps, "kNWarps must be a power of 2"); constexpr int kChunksPerExchange = Ktraits::kSmemExchangeSize / (sizeof(vec_t) * kNExchangePerVec * kNThreads); static_assert(kChunksPerExchange * sizeof(vec_t) * kNExchangePerVec * kNThreads == Ktraits::kSmemExchangeSize); constexpr int kNExchanges = kNChunks / kChunksPerExchange; static_assert(kNExchanges * kChunksPerExchange == kNChunks); extern __shared__ char smem_[]; vec_t* smem_exchange = reinterpret_cast(smem_); const int batch_id = static_cast(blockIdx.x); input_t* x = reinterpret_cast(params.x_ptr) + batch_id * params.x_batch_stride; input_t* out = reinterpret_cast(params.out_ptr) + batch_id * params.out_batch_stride; float x_vals[kNChunks][kNElts]; load_input(x, x_vals, params.dim); hadamard_mult_thread(x_vals); hadamard_mult_warp(x_vals); if constexpr (kNWarps > 1) { exchange_smem_pre(x_vals, smem_exchange); hadamard_mult_warp(x_vals); exchange_smem_pre(x_vals, smem_exchange); } if constexpr (kNChunks > 1) { float x_vals_transposed[kNElts][kNChunks]; #pragma unroll for (int c = 0; c < kNChunks; ++c) { #pragma unroll for (int i = 0; i < kNElts; ++i) { x_vals_transposed[i][c] = x_vals[c][i]; } } if constexpr (kNChunks == 12) { hadamard_mult_thread_chunk_12(x_vals_transposed); } else if constexpr (kNChunks == 20) { hadamard_mult_thread_chunk_20(x_vals_transposed); } else if constexpr (kNChunks == 28) { hadamard_mult_thread_chunk_28(x_vals_transposed); } else if constexpr (kNChunks == 40) { hadamard_mult_thread_chunk_40(x_vals_transposed); } else { constexpr int kLogNChunks = cilog2(kNChunks); static_assert(1 << kLogNChunks == kNChunks, "kNChunks must be a power of 2"); hadamard_mult_thread(x_vals_transposed); } #pragma unroll for (int c = 0; c < kNChunks; ++c) { #pragma unroll for (int i = 0; i < kNElts; ++i) { x_vals[c][i] = x_vals_transposed[i][c]; } } } store_output(out, x_vals, params.dim, params.scale); } template inline void set_max_dynamic_smem() { constexpr int kSmemSize = Ktraits::kSmemSize; if constexpr (kSmemSize >= 48 * 1024) { auto kernel = &fast_hadamard_transform_kernel; host::RuntimeDeviceCheck(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); } } template inline void launch_kernel(HadamardParamsBase& params, DLDevice device) { constexpr int kSmemSize = Ktraits::kSmemSize; set_max_dynamic_smem(); auto kernel = &fast_hadamard_transform_kernel; host::LaunchKernel(dim3(params.batch), dim3(Ktraits::kNThreads), device, kSmemSize)(kernel, params); host::RuntimeDeviceCheck(); } template inline void fast_hadamard_transform_launch(HadamardParamsBase& params, DLDevice device) { using Ktraits = FastHadamardKernelTraits; launch_kernel(params, device); } template inline void fast_hadamard_transform_cuda(HadamardParamsBase& params, DLDevice device) { if (params.log_N == 3) { fast_hadamard_transform_launch<1, 3, input_t>(params, device); } else if (params.log_N == 4) { fast_hadamard_transform_launch<2, 4, input_t>(params, device); } else if (params.log_N == 5) { fast_hadamard_transform_launch<4, 5, input_t>(params, device); } else if (params.log_N == 6) { fast_hadamard_transform_launch<8, 6, input_t>(params, device); } else if (params.log_N == 7) { fast_hadamard_transform_launch<16, 7, input_t>(params, device); } else if (params.log_N == 8) { fast_hadamard_transform_launch<32, 8, input_t>(params, device); } else if (params.log_N == 9) { fast_hadamard_transform_launch<32, 9, input_t>(params, device); } else if (params.log_N == 10) { fast_hadamard_transform_launch<128, 10, input_t>(params, device); } else if (params.log_N == 11) { fast_hadamard_transform_launch<256, 11, input_t>(params, device); } else if (params.log_N == 12) { fast_hadamard_transform_launch<256, 12, input_t>(params, device); } else if (params.log_N == 13) { fast_hadamard_transform_launch<256, 13, input_t>(params, device); } else if (params.log_N == 14) { fast_hadamard_transform_launch<256, 14, input_t>(params, device); } else if (params.log_N == 15) { fast_hadamard_transform_launch<256, 15, input_t>(params, device); } else { host::Panic("fast_hadamard_transform: unsupported log_N=", params.log_N); } } template inline void fast_hadamard_transform_12N_launch(HadamardParamsBase& params, DLDevice device) { using Ktraits = FastHadamard12NTraits; launch_kernel(params, device); } template inline void fast_hadamard_transform_12N_cuda(HadamardParamsBase& params, DLDevice device) { if (params.log_N == 2) { fast_hadamard_transform_12N_launch<1, 2, input_t>(params, device); } else if (params.log_N == 3) { fast_hadamard_transform_12N_launch<2, 3, input_t>(params, device); } else if (params.log_N == 4) { fast_hadamard_transform_12N_launch<4, 4, input_t>(params, device); } else if (params.log_N == 5) { fast_hadamard_transform_12N_launch<8, 5, input_t>(params, device); } else if (params.log_N == 6) { fast_hadamard_transform_12N_launch<16, 6, input_t>(params, device); } else if (params.log_N == 7) { fast_hadamard_transform_12N_launch<32, 7, input_t>(params, device); } else if (params.log_N == 8) { fast_hadamard_transform_12N_launch<64, 8, input_t>(params, device); } else if (params.log_N == 9) { fast_hadamard_transform_12N_launch<128, 9, input_t>(params, device); } else if (params.log_N == 10) { fast_hadamard_transform_12N_launch<256, 10, input_t>(params, device); } else { host::Panic("fast_hadamard_transform_12N: unsupported log_N=", params.log_N); } } template inline void fast_hadamard_transform_20N_launch(HadamardParamsBase& params, DLDevice device) { using Ktraits = FastHadamard20NTraits; launch_kernel(params, device); } template inline void fast_hadamard_transform_20N_cuda(HadamardParamsBase& params, DLDevice device) { if (params.log_N == 2) { fast_hadamard_transform_20N_launch<1, 2, input_t>(params, device); } else if (params.log_N == 3) { fast_hadamard_transform_20N_launch<2, 3, input_t>(params, device); } else if (params.log_N == 4) { fast_hadamard_transform_20N_launch<4, 4, input_t>(params, device); } else if (params.log_N == 5) { fast_hadamard_transform_20N_launch<8, 5, input_t>(params, device); } else if (params.log_N == 6) { fast_hadamard_transform_20N_launch<16, 6, input_t>(params, device); } else if (params.log_N == 7) { fast_hadamard_transform_20N_launch<32, 7, input_t>(params, device); } else if (params.log_N == 8) { fast_hadamard_transform_20N_launch<64, 8, input_t>(params, device); } else if (params.log_N == 9) { fast_hadamard_transform_20N_launch<128, 9, input_t>(params, device); } else if (params.log_N == 10) { fast_hadamard_transform_20N_launch<256, 10, input_t>(params, device); } else { host::Panic("fast_hadamard_transform_20N: unsupported log_N=", params.log_N); } } template inline void fast_hadamard_transform_28N_launch(HadamardParamsBase& params, DLDevice device) { using Ktraits = FastHadamard28NTraits; launch_kernel(params, device); } template inline void fast_hadamard_transform_28N_cuda(HadamardParamsBase& params, DLDevice device) { if (params.log_N == 2) { fast_hadamard_transform_28N_launch<1, 2, input_t>(params, device); } else if (params.log_N == 3) { fast_hadamard_transform_28N_launch<2, 3, input_t>(params, device); } else if (params.log_N == 4) { fast_hadamard_transform_28N_launch<4, 4, input_t>(params, device); } else if (params.log_N == 5) { fast_hadamard_transform_28N_launch<8, 5, input_t>(params, device); } else if (params.log_N == 6) { fast_hadamard_transform_28N_launch<16, 6, input_t>(params, device); } else if (params.log_N == 7) { fast_hadamard_transform_28N_launch<32, 7, input_t>(params, device); } else if (params.log_N == 8) { fast_hadamard_transform_28N_launch<64, 8, input_t>(params, device); } else if (params.log_N == 9) { fast_hadamard_transform_28N_launch<128, 9, input_t>(params, device); } else if (params.log_N == 10) { fast_hadamard_transform_28N_launch<256, 10, input_t>(params, device); } else { host::Panic("fast_hadamard_transform_28N: unsupported log_N=", params.log_N); } } template inline void fast_hadamard_transform_40N_launch(HadamardParamsBase& params, DLDevice device) { using Ktraits = FastHadamard40NTraits; launch_kernel(params, device); } template inline void fast_hadamard_transform_40N_cuda(HadamardParamsBase& params, DLDevice device) { if (params.log_N == 2) { fast_hadamard_transform_40N_launch<1, 2, input_t>(params, device); } else if (params.log_N == 3) { fast_hadamard_transform_40N_launch<2, 3, input_t>(params, device); } else if (params.log_N == 4) { fast_hadamard_transform_40N_launch<4, 4, input_t>(params, device); } else if (params.log_N == 5) { fast_hadamard_transform_40N_launch<8, 5, input_t>(params, device); } else if (params.log_N == 6) { fast_hadamard_transform_40N_launch<16, 6, input_t>(params, device); } else if (params.log_N == 7) { fast_hadamard_transform_40N_launch<32, 7, input_t>(params, device); } else if (params.log_N == 8) { fast_hadamard_transform_40N_launch<64, 8, input_t>(params, device); } else if (params.log_N == 9) { fast_hadamard_transform_40N_launch<128, 9, input_t>(params, device); } else if (params.log_N == 10) { fast_hadamard_transform_40N_launch<256, 10, input_t>(params, device); } else { host::Panic("fast_hadamard_transform_40N: unsupported log_N=", params.log_N); } } inline void set_hadamard_params( HadamardParamsBase& params, int64_t batch, int64_t dim, int64_t multiple, const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) { std::memset(¶ms, 0, sizeof(params)); params.batch = static_cast(batch); params.dim = static_cast(dim); params.log_N = ceil_log2(static_cast(dim / multiple)); params.x_ptr = const_cast(x.data_ptr()); params.out_ptr = const_cast(out.data_ptr()); params.x_batch_stride = x.stride(0); params.out_batch_stride = out.stride(0); params.scale = scale; } template inline void run_hadamard(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) { using namespace host; auto N = SymbolicSize{"batch"}; auto D = SymbolicSize{"dim"}; auto SX = SymbolicSize{"x_batch_stride"}; auto SO = SymbolicSize{"out_batch_stride"}; auto device = SymbolicDevice{}; device.set_options(); TensorMatcher({N, D}).with_strides({SX, 1}).with_dtype().with_device(device).verify(x); TensorMatcher({N, D}).with_strides({SO, 1}).with_dtype().with_device(device).verify(out); const int64_t batch = N.unwrap(); const int64_t dim = D.unwrap(); RuntimeCheck(dim % kMultiple == 0, "hadamard: dim must be divisible by ", kMultiple); HadamardParamsBase params; set_hadamard_params(params, batch, dim, kMultiple, x, out, scale); if constexpr (kMultiple == 1) { RuntimeCheck(dim % 8 == 0, "fast_hadamard_transform only supports hidden dim divisible by 8"); RuntimeCheck(dim <= 32768, "fast_hadamard_transform only supports hidden dim <= 32768"); fast_hadamard_transform_cuda(params, device.unwrap()); } else if constexpr (kMultiple == 12) { RuntimeCheck(dim % (4 * 12) == 0, "fast_hadamard_transform_12N only supports hidden dim divisible by 48"); RuntimeCheck(dim <= 12 * 1024, "fast_hadamard_transform_12N only supports hidden dim <= 12288"); fast_hadamard_transform_12N_cuda(params, device.unwrap()); } else if constexpr (kMultiple == 20) { RuntimeCheck(dim % (4 * 20) == 0, "fast_hadamard_transform_20N only supports hidden dim divisible by 80"); RuntimeCheck(dim <= 20 * 1024, "fast_hadamard_transform_20N only supports hidden dim <= 20480"); fast_hadamard_transform_20N_cuda(params, device.unwrap()); } else if constexpr (kMultiple == 28) { RuntimeCheck(dim % (4 * 28) == 0, "fast_hadamard_transform_28N only supports hidden dim divisible by 112"); RuntimeCheck(dim <= 28 * 1024, "fast_hadamard_transform_28N only supports hidden dim <= 28672"); fast_hadamard_transform_28N_cuda(params, device.unwrap()); } else if constexpr (kMultiple == 40) { RuntimeCheck(dim % (4 * 40) == 0, "fast_hadamard_transform_40N only supports hidden dim divisible by 160"); RuntimeCheck(dim <= 40 * 1024, "fast_hadamard_transform_40N only supports hidden dim <= 40960"); fast_hadamard_transform_40N_cuda(params, device.unwrap()); } else { Panic("Unsupported multiple"); } } template struct HadamardKernel { static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) { run_hadamard<1, DType>(x, out, scale); } }; template struct Hadamard12NKernel { static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) { run_hadamard<12, DType>(x, out, scale); } }; template struct Hadamard20NKernel { static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) { run_hadamard<20, DType>(x, out, scale); } }; template struct Hadamard28NKernel { static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) { run_hadamard<28, DType>(x, out, scale); } }; template struct Hadamard40NKernel { static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) { run_hadamard<40, DType>(x, out, scale); } }; } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/fast-hadamard-transform/static_switch.h ================================================ // Copied from https://github.com/sgl-project/fast-hadamard-transform // Inspired by https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h // and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h #pragma once /// @param COND - a boolean expression to switch by /// @param CONST_NAME - a name given for the constexpr bool variable. /// @param ... - code to execute for true and false /// /// Usage: /// ``` /// BOOL_SWITCH(flag, BoolConst, [&] { /// some_function(...); /// }); /// ``` #define BOOL_SWITCH(COND, CONST_NAME, ...) \ [&] { \ if (COND) { \ static constexpr bool CONST_NAME = true; \ return __VA_ARGS__(); \ } else { \ static constexpr bool CONST_NAME = false; \ return __VA_ARGS__(); \ } \ }() ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/awq_dequantize.cuh ================================================ // Adapted from // https://github.com/vllm-project/vllm/blob/eb59b5a6cba6727d3727c0372258db9002f687c1/csrc/quantization/awq/gemm_kernels.cu#L350 #pragma once #include #include namespace device::awq { template __device__ inline int lop3(int a, int b, int c) { int res; asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(res) : "r"(a), "r"(b), "r"(c), "n"(lut)); return res; } __device__ uint4 dequantize_s4_to_fp16x2(uint32_t const& source) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750 uint4 result; uint32_t* h = reinterpret_cast(&result); uint32_t const i4s = reinterpret_cast(source); // First, we extract the i4s and construct an intermediate fp16 number. static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; static constexpr uint32_t BOTTOM_MASK = 0x000f000f; static constexpr uint32_t TOP_MASK = 0x00f000f0; static constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400; // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide RAW // dependency if we issue immediately before required. const uint32_t top_i4s = i4s >> 8; // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400 asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[0]) : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400 asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[1]) : "r"(i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); // Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400 asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[2]) : "r"(top_i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400 asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(h[3]) : "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); // This is the half2 {1024, 1024} represented as an integer. static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64006400; // This is the half2 {1 / 16, 1 / 16} represented as an integer. static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00; // This is the half2 {-64, -64} represented as an integer. static constexpr uint32_t NEG_64 = 0xd400d400; // Finally, we construct the output numbers. // Convert elt_01 asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM)); // Convert elt_23 asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[1]) : "r"(h[1]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); // Convert elt_45 asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(FP16_TOP_MAGIC_NUM)); // Convert elt_67 asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[3]) : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); return result; #else assert(false); return {}; #endif } __device__ uint4 dequantize_s4_to_bf16x2(uint32_t const& source) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 uint4 result; uint32_t* h = reinterpret_cast(&result); uint32_t const i4s = source; // Define masks and constants static constexpr uint32_t MASK = 0x000f000f; static constexpr uint32_t EX = 0x43004300; static constexpr uint32_t MUL = 0x3F803F80; static constexpr uint32_t ADD = 0xC300C300; int lo0 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s, MASK, EX); int hi0 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s >> 4, MASK, EX); int lo1 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s >> 8, MASK, EX); int hi1 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s >> 12, MASK, EX); nv_bfloat162* res = reinterpret_cast(h); res[0] = __hfma2( *reinterpret_cast(&lo0), *reinterpret_cast(&MUL), *reinterpret_cast(&ADD)); res[1] = __hfma2( *reinterpret_cast(&hi0), *reinterpret_cast(&MUL), *reinterpret_cast(&ADD)); res[2] = __hfma2( *reinterpret_cast(&lo1), *reinterpret_cast(&MUL), *reinterpret_cast(&ADD)); res[3] = __hfma2( *reinterpret_cast(&hi1), *reinterpret_cast(&MUL), *reinterpret_cast(&ADD)); return result; #else assert(false); return {}; #endif } template __global__ void __launch_bounds__(256) dequantize_weights( int* __restrict__ qweight, OutputT* __restrict__ scales, int* __restrict__ qzeros, OutputT* __restrict__ output, int group_size, int qweight_cols, int qweight_rows) { int col = blockIdx.x * blockDim.x + threadIdx.x; int row = blockIdx.y * blockDim.y + threadIdx.y; if (col >= qweight_cols || row >= qweight_rows) return; int group_idx = row / group_size; int scale_offset = 8 * col + group_idx * qweight_cols * 8; uint4 loaded_scale = *(uint4*)(scales + scale_offset); // Handle different data types if constexpr (std::is_same::value) { // FP16 path uint4 zeros = dequantize_s4_to_fp16x2(qzeros[col + group_idx * qweight_cols]); uint4 weight_fp16 = dequantize_s4_to_fp16x2(qweight[col + row * qweight_cols]); // Use PTX assembly for FP16 operations asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.x) : "r"(weight_fp16.x), "r"(zeros.x)); asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.x) : "r"(weight_fp16.x), "r"(loaded_scale.x)); asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.y) : "r"(weight_fp16.y), "r"(zeros.y)); asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.y) : "r"(weight_fp16.y), "r"(loaded_scale.y)); asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.z) : "r"(weight_fp16.z), "r"(zeros.z)); asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.z) : "r"(weight_fp16.z), "r"(loaded_scale.z)); asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.w) : "r"(weight_fp16.w), "r"(zeros.w)); asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.w) : "r"(weight_fp16.w), "r"(loaded_scale.w)); OutputT* output_ptr = output + 8 * col + 8 * row * qweight_cols; *(uint4*)output_ptr = weight_fp16; } else if constexpr (std::is_same::value) { uint4 weight_raw = dequantize_s4_to_bf16x2(qweight[col + row * qweight_cols]); uint4 zero_raw = dequantize_s4_to_bf16x2(qzeros[col + group_idx * qweight_cols]); uint4 scale_raw = *reinterpret_cast(scales + scale_offset); // Vectorized processing (each uint4 contains 4 nv_bfloat162) nv_bfloat162* weight_vec = reinterpret_cast(&weight_raw); nv_bfloat162* zero_vec = reinterpret_cast(&zero_raw); nv_bfloat162* scale_vec = reinterpret_cast(&scale_raw); // Single instruction dual-channel operation #pragma unroll for (int i = 0; i < 4; ++i) { // uint4 = 4 * nv_bfloat162 weight_vec[i] = __hmul2(__hsub2(weight_vec[i], zero_vec[i]), scale_vec[i]); } // Directly store to OutputT array (guaranteed contiguous memory) OutputT* output_ptr = output + 8 * col + row * qweight_cols * 8; static_assert(sizeof(uint4) == 8 * sizeof(OutputT), "Memory layout mismatch"); *reinterpret_cast(output_ptr) = weight_raw; } } } // namespace device::awq // Host wrapper template void awq_dequantize( tvm::ffi::TensorView output, tvm::ffi::TensorView qweight, tvm::ffi::TensorView scales, tvm::ffi::TensorView qzeros) { using namespace host; int64_t qweight_rows = qweight.size(0); int64_t qweight_cols = qweight.size(1); int64_t scales_rows = scales.size(0); // Validate tensors SymbolicDevice cuda_device; cuda_device.set_options(); TensorMatcher({qweight_rows, qweight_cols}).with_dtype().with_device(cuda_device).verify(qweight); TensorMatcher({scales_rows, qweight_cols * 8}).with_dtype().with_device(cuda_device).verify(scales); TensorMatcher({scales_rows, qweight_cols}).with_dtype().with_device(cuda_device).verify(qzeros); TensorMatcher({qweight_rows, qweight_cols * 8}).with_dtype().with_device(cuda_device).verify(output); // Get device and stream auto device = cuda_device.unwrap(); auto stream = LaunchKernel::resolve_device(device); int group_size = static_cast(qweight_rows / scales_rows); int x_num_threads = 16; int y_num_threads = 16; int x_blocks = (static_cast(qweight_cols) + x_num_threads - 1) / x_num_threads; int y_blocks = (static_cast(qweight_rows) + y_num_threads - 1) / y_num_threads; dim3 num_blocks(x_blocks, y_blocks); dim3 threads_per_block(x_num_threads, y_num_threads); // Get pointers auto* qweight_ptr = reinterpret_cast(qweight.data_ptr()); auto* scales_ptr = reinterpret_cast(scales.data_ptr()); auto* qzeros_ptr = reinterpret_cast(qzeros.data_ptr()); auto* output_ptr = reinterpret_cast(output.data_ptr()); LaunchKernel(num_blocks, threads_per_block, stream)( device::awq::dequantize_weights, qweight_ptr, scales_ptr, qzeros_ptr, output_ptr, group_size, static_cast(qweight_cols), static_cast(qweight_rows)); } ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/marlin/awq_marlin_repack.cuh ================================================ #pragma once #include #include #include "marlin.cuh" namespace device::marlin { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 template __global__ void awq_marlin_repack_kernel( uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr, int size_k, int size_n) { return; } #else template __global__ void awq_marlin_repack_kernel( uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr, int size_k, int size_n) { constexpr int pack_factor = 32 / num_bits; int k_tiles = size_k / tile_k_size; int n_tiles = size_n / tile_n_size; int block_k_tiles = div_ceil(k_tiles, (int)gridDim.x); auto start_k_tile = blockIdx.x * block_k_tiles; if (start_k_tile >= k_tiles) { return; } int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); // Wait until the next thread tile has been loaded to shared memory. auto wait_for_stage = [&]() { // We only have `stages - 2` active fetches since we are double buffering // and can only issue the next fetch when it is guaranteed that the previous // shared memory load is fully complete (as it may otherwise be // overwritten). cp_async_wait(); __syncthreads(); }; extern __shared__ int4 sh[]; constexpr int tile_n_ints = tile_n_size / pack_factor; constexpr int stage_n_threads = tile_n_ints / 4; constexpr int stage_k_threads = tile_k_size; constexpr int stage_size = stage_k_threads * stage_n_threads; auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { if (n_tile_id >= n_tiles) { cp_async_fence(); return; } int first_n = n_tile_id * tile_n_size; int first_n_packed = first_n / pack_factor; int4* sh_ptr = sh + stage_size * pipe; if (threadIdx.x < stage_size) { auto k_id = threadIdx.x / stage_n_threads; auto n_id = threadIdx.x % stage_n_threads; int first_k = k_tile_id * tile_k_size; cp_async4( &sh_ptr[k_id * stage_n_threads + n_id], reinterpret_cast( &(b_q_weight_ptr[(first_k + k_id) * (size_n / pack_factor) + first_n_packed + (n_id * 4)]))); } cp_async_fence(); }; auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { if (n_tile_id >= n_tiles) { return; } auto warp_id = threadIdx.x / 32; auto th_id = threadIdx.x % 32; if (warp_id >= 4) { return; } int tc_col = th_id / 4; int tc_row = (th_id % 4) * 2; constexpr int tc_offsets[4] = {0, 1, 8, 9}; int cur_n = warp_id * 16 + tc_col; int cur_n_packed = cur_n / pack_factor; int cur_n_pos = cur_n % pack_factor; constexpr int sh_stride = tile_n_ints; constexpr uint32_t mask = (1 << num_bits) - 1; int4* sh_stage_ptr = sh + stage_size * pipe; uint32_t* sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); // Undo interleaving int cur_n_pos_unpacked; if constexpr (num_bits == 4) { constexpr int undo_pack[8] = {0, 4, 1, 5, 2, 6, 3, 7}; cur_n_pos_unpacked = undo_pack[cur_n_pos]; } else { constexpr int undo_pack[4] = {0, 2, 1, 3}; cur_n_pos_unpacked = undo_pack[cur_n_pos]; } uint32_t vals[8]; #pragma unroll for (int i = 0; i < 4; i++) { int cur_elem = tc_row + tc_offsets[i]; int packed_src_0 = sh_stage_int_ptr[cur_n_packed + sh_stride * cur_elem]; int packed_src_1 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) + sh_stride * cur_elem]; vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask; vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask; } constexpr int tile_size_val = tile_k_size * tile_n_size / pack_factor; int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size_val; // Result of: // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h if constexpr (num_bits == 4) { constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; uint32_t res = 0; #pragma unroll for (int i = 0; i < 8; i++) { res |= vals[pack_idx[i]] << (i * 4); } out_ptr[out_offset + th_id * 4 + warp_id] = res; } else { constexpr int pack_idx[4] = {0, 2, 1, 3}; uint32_t res1 = 0; uint32_t res2 = 0; #pragma unroll for (int i = 0; i < 4; i++) { res1 |= vals[pack_idx[i]] << (i * 8); res2 |= vals[4 + pack_idx[i]] << (i * 8); } out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; } }; auto start_pipes = [&](int k_tile_id, int n_tile_id) { #pragma unroll for (int pipe = 0; pipe < repack_stages - 1; pipe++) { fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); } wait_for_stage(); }; #pragma unroll for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) { int n_tile_id = 0; start_pipes(k_tile_id, n_tile_id); while (n_tile_id < n_tiles) { #pragma unroll for (int pipe = 0; pipe < repack_stages; pipe++) { fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, n_tile_id + pipe + repack_stages - 1); repack_tile(pipe, k_tile_id, n_tile_id + pipe); wait_for_stage(); } n_tile_id += repack_stages; } } } #endif } // namespace device::marlin // Host wrapper void awq_marlin_repack( tvm::ffi::TensorView out, tvm::ffi::TensorView b_q_weight, int64_t size_k, int64_t size_n, int64_t num_bits) { using namespace host; using namespace device::marlin; // Validate alignment RuntimeCheck(size_k % tile_k_size == 0, "size_k = ", size_k, " is not divisible by tile_k_size = ", tile_k_size); RuntimeCheck(size_n % tile_n_size == 0, "size_n = ", size_n, " is not divisible by tile_n_size = ", tile_n_size); RuntimeCheck(num_bits == 4 || num_bits == 8, "num_bits must be 4 or 8. Got = ", num_bits); int const pack_factor = 32 / num_bits; // Validate tensors SymbolicDevice cuda_device; cuda_device.set_options(); TensorMatcher({size_k, size_n / pack_factor}).with_dtype().with_device(cuda_device).verify(b_q_weight); TensorMatcher({size_k / tile_size, size_n * tile_size / pack_factor}) .with_dtype() .with_device(cuda_device) .verify(out); // Get device and stream auto device = cuda_device.unwrap(); auto stream = LaunchKernel::resolve_device(device); // Get pointers auto* b_q_weight_ptr = reinterpret_cast(b_q_weight.data_ptr()); auto* out_ptr = reinterpret_cast(out.data_ptr()); // Get device attributes int blocks = 0; cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device.device_id); int max_shared_mem = 0; cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device.device_id); RuntimeCheck(max_shared_mem > 0, "max_shared_mem must be > 0"); // Dispatch based on num_bits if (num_bits == 4) { cudaFuncSetAttribute( awq_marlin_repack_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); LaunchKernel(blocks, repack_threads, stream, max_shared_mem)( awq_marlin_repack_kernel, b_q_weight_ptr, out_ptr, static_cast(size_k), static_cast(size_n)); } else if (num_bits == 8) { cudaFuncSetAttribute( awq_marlin_repack_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); LaunchKernel(blocks, repack_threads, stream, max_shared_mem)( awq_marlin_repack_kernel, b_q_weight_ptr, out_ptr, static_cast(size_k), static_cast(size_n)); } else { RuntimeCheck(false, "Unsupported repack config: num_bits = ", num_bits); } } ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/marlin/dequant.h ================================================ /* Fast Dequantization (Converting INT4/INT8/FP4/FP8 to FP16/BF16) The process of fast dequantization can be summarized as a combination of bitwise operations and floating-point computations: weight =>(bit_op / bitwise operations)=> f16_value =>(flop / floating-point computation)=> dequantized_weight Since the dequantized weights typically require subtracting the zero point and applying a scale factor, the floating-point computation step can be fused with the zero-point subtraction and scaling operations. The following are the parts that need to be modified for the fused operation of zero-point subtraction and scaling. ## INT4 => FP16/BF16 or INT8 => FP16 The floating-point computation is `__hsub2` If has zero points: flop(bit_op(weight)) - flop(bit_op(zp)) = sub(bit_op(weight), bias) - sub(bit_op(zp), bias) = bit_op(weight) - bit_op(zp) so we don't need additional modification. If has float zero points: flop(bit_op(weight)) - fzp = sub(bit_op(weight), bias) - fzp = bit_op(weight) - (fzp + bias) where the `fzp + bias` can be computed at weight loading. But this may have accuracy issue, so we should not use this in most cases. If has not zero points: scale(flop(bit_op(weight))) = scale(sub(bit_op(weight), bias)) = scale(bit_op(weight)) - scale(bias) = fma(bit_op(weight), scale_factor, scale(bias)) where the `scale(bias)` can be cached. But this may have accuracy issue, so we should not use this in most cases. ## INT8 => BF16 INT8 => BF16 is a special case, it use byte_perm instead of flop. We cannot fused byte_perm with scaling. ## FP4/FP8 => FP16/BF16 scale(flop(bit_op(weight))) = scale(mul(bit_op(weight), multiplier)) = mul(bit_op(weight), scale_factor * multiplier) where `scale_factor * multiplier` can be computed at weight loading. */ #include "marlin_dtypes.cuh" namespace device::marlin { #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 // Lookup-table based 3-input logical operation; explicitly used for // dequantization as the compiler does not seem to automatically recognize it in // all cases. template __device__ inline int lop3(int a, int b, int c) { int res; asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(res) : "r"(a), "r"(b), "r"(c), "n"(lut)); return res; } // Constructs destination register by taking bytes from 2 sources (based on // mask) template __device__ inline uint32_t prmt(uint32_t a) { uint32_t res; asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(res) : "r"(a), "n"(start_byte), "n"(mask)); return res; } template __device__ inline void dequant(int q, scalar_t2* frag_b); // // Efficiently dequantize 4bit values packed in an int32 value into a full // B-fragment of 4 fp16 values. We mostly follow the strategy in the link below, // with some small changes: // - FP16: // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287 // - BF16: // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385 // template <> __device__ inline void dequant(int q, half2* frag_b) { const int MASK = 0x000f000f; const int EX = 0x64006400; // Guarantee that the `(a & b) | c` operations are LOP3s. int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); q >>= 4; int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); frag_b[0] = *reinterpret_cast(&lo); frag_b[1] = *reinterpret_cast(&hi); } template <> __device__ inline void dequant(int q, half2* frag_b) { const int LO = 0x000f000f; const int HI = 0x00f000f0; const int EX = 0x64006400; // Guarantee that the `(a & b) | c` operations are LOP3s. // clang-format off int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); // clang-format on // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point // directly into `SUB` and `ADD`. const int SUB = 0x64086408; const int MUL = 0x2c002c00; const int ADD = 0xd480d480; frag_b[0] = __hsub2(*reinterpret_cast(&lo), *reinterpret_cast(&SUB)); frag_b[1] = __hfma2( *reinterpret_cast(&hi), *reinterpret_cast(&MUL), *reinterpret_cast(&ADD)); } template <> __device__ inline void dequant(int q, half2* frag_b) { dequant(q, frag_b); } template <> __device__ inline void dequant(int q, half2* frag_b) { const int LO = 0x000f000f; const int HI = 0x00f000f0; const int EX = 0x64006400; // Guarantee that the `(a & b) | c` operations are LOP3s. // clang-format off int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); // clang-format on // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point // directly into `SUB` and `ADD`. const int SUB = 0x64006400; const int MUL = 0x2c002c00; const int ADD = 0xd400d400; frag_b[0] = __hsub2(*reinterpret_cast(&lo), *reinterpret_cast(&SUB)); frag_b[1] = __hfma2( *reinterpret_cast(&hi), *reinterpret_cast(&MUL), *reinterpret_cast(&ADD)); } template <> __device__ inline void dequant(int q, nv_bfloat162* frag_b) { static constexpr uint32_t MASK = 0x000f000f; static constexpr uint32_t EX = 0x43004300; // Guarantee that the `(a & b) | c` operations are LOP3s. // clang-format off int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); q >>= 4; int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); // clang-format on frag_b[0] = *reinterpret_cast(&lo); frag_b[1] = *reinterpret_cast(&hi); } template <> __device__ inline void dequant(int q, nv_bfloat162* frag_b) { dequant(q, frag_b); static constexpr uint32_t SUB = 0x43084308; frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&SUB)); frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&SUB)); } template <> __device__ inline void dequant(int q, nv_bfloat162* frag_b) { dequant(q, frag_b); } template <> __device__ inline void dequant(int q, nv_bfloat162* frag_b) { dequant(q, frag_b); static constexpr uint32_t SUB = 0x43004300; frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&SUB)); frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&SUB)); } // // Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or // bf16 Reference: // - FP16: // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85 // - BF16: // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175 // template <> __device__ inline void dequant(int q, half2* frag_b) { static constexpr uint32_t mask_for_elt_01 = 0x5250; static constexpr uint32_t mask_for_elt_23 = 0x5351; static constexpr uint32_t start_byte_for_fp16 = 0x64646464; uint32_t lo = prmt(q); uint32_t hi = prmt(q); frag_b[0] = *reinterpret_cast(&lo); frag_b[1] = *reinterpret_cast(&hi); } template <> __device__ inline void dequant(int q, half2* frag_b) { dequant(q, frag_b); static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480; frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); } template <> __device__ inline void dequant(int q, half2* frag_b) { dequant(q, frag_b); } template <> __device__ inline void dequant(int q, half2* frag_b) { dequant(q, frag_b); static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400; frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); } template <> __device__ inline void dequant(int q, nv_bfloat162* frag_b) { float fp32_intermediates[4]; uint32_t* fp32_intermediates_casted = reinterpret_cast(fp32_intermediates); static constexpr uint32_t fp32_base = 0x4B000000; fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); fp32_intermediates[0] -= 8388736.f; fp32_intermediates[1] -= 8388736.f; fp32_intermediates[2] -= 8388736.f; fp32_intermediates[3] -= 8388736.f; uint32_t* bf16_result_ptr = reinterpret_cast(frag_b); bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], fp32_intermediates_casted[1], 0x7632); bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], fp32_intermediates_casted[3], 0x7632); } template <> __device__ inline void dequant(int q, nv_bfloat162* frag_b) { float fp32_intermediates[4]; uint32_t* fp32_intermediates_casted = reinterpret_cast(fp32_intermediates); static constexpr uint32_t fp32_base = 0x4B000000; fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); fp32_intermediates[0] -= 8388608.f; fp32_intermediates[1] -= 8388608.f; fp32_intermediates[2] -= 8388608.f; fp32_intermediates[3] -= 8388608.f; uint32_t* bf16_result_ptr = reinterpret_cast(frag_b); bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], fp32_intermediates_casted[1], 0x7632); bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], fp32_intermediates_casted[3], 0x7632); } template <> __device__ inline void dequant(int q, half2* frag_b) { // Constants for FP8 (E4M3) and FP16 formats constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5; constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT; constexpr int MASK = 0x7F007F00; // Extract and shift FP8 values to FP16 format int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); q <<= 8; int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); // Note: reverse indexing is intentional because weights are permuted frag_b[1] = *reinterpret_cast(&Out1); frag_b[0] = *reinterpret_cast(&Out2); } template <> __device__ inline void dequant(int q, half2* frag_b) { dequant(q, frag_b); // Constants for FP8 (E4M3) and FP16 formats constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5; // Construct and apply exponent bias constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); // Convert to half2 and apply bias frag_b[1] = __hmul2(frag_b[1], bias_reg); frag_b[0] = __hmul2(frag_b[0], bias_reg); } template <> __device__ inline void dequant(int q, nv_bfloat162* frag_b) { // Constants for FP8 (E4M3) and BF16 formats constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; constexpr int MASK = 0x7F007F00; // Extract and shift FP8 values to BF16 format int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); q <<= 8; int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); // Note: reverse indexing is intentional because weights are permuted frag_b[1] = *reinterpret_cast(&Out1); frag_b[0] = *reinterpret_cast(&Out2); } template <> __device__ inline void dequant(int q, nv_bfloat162* frag_b) { dequant(q, frag_b); // Constants for FP8 (E4M3) and BF16 formats constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; // Construct and apply exponent bias constexpr int BIAS_OFFSET = (1 << (BF16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent // position constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); // Convert to bfloat162 and apply bias frag_b[1] = __hmul2(frag_b[1], bias_reg); frag_b[0] = __hmul2(frag_b[0], bias_reg); } template <> __device__ inline void dequant(int q, half2* frag_b) { // Constants for FP4 (E2M1) and FP16 formats constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5; constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP4_EXPONENT; constexpr int MASK = 0x70007000; // Extract and shift FP4 values to FP16 format int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); q <<= 4; int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); // Note: reverse indexing is intentional because weights are permuted frag_b[1] = *reinterpret_cast(&Out1); frag_b[0] = *reinterpret_cast(&Out2); } template <> __device__ inline void dequant(int q, half2* frag_b) { dequant(q, frag_b); // Constants for FP4 (E2M1) and FP16 formats constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5; // Construct and apply exponent bias constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1)); const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); // Convert to half2 and apply bias frag_b[1] = __hmul2(frag_b[1], bias_reg); frag_b[0] = __hmul2(frag_b[0], bias_reg); } template <> __device__ inline void dequant(int q, nv_bfloat162* frag_b) { // Constants for FP4 (E2M1) and FP16 formats constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8; constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP4_EXPONENT; constexpr int MASK = 0x70007000; // Extract and shift FP4 values to FP16 format int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); q <<= 4; int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); // Note: reverse indexing is intentional because weights are permuted frag_b[1] = *reinterpret_cast(&Out1); frag_b[0] = *reinterpret_cast(&Out2); } template <> __device__ inline void dequant(int q, nv_bfloat162* frag_b) { dequant(q, frag_b); // Constants for FP4 (E2M1) and BF16 formats constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8; // Construct and apply exponent bias constexpr int BIAS_OFFSET = (1 << (BF16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1)); // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent // position constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); // Convert to half2 and apply bias frag_b[1] = __hmul2(frag_b[1], bias_reg); frag_b[0] = __hmul2(frag_b[0], bias_reg); } template __device__ inline void dequant_fp8_scales(int q, scalar_t2* frag_b); template <> __device__ inline void dequant_fp8_scales(int q, half2* frag_b) { int Out1 = (q & 0xFF00FF00) >> 1; ; q <<= 8; int Out2 = (q & 0xFF00FF00) >> 1; // Note: reverse indexing is intentional because weights are permuted frag_b[1] = *reinterpret_cast(&Out1); frag_b[0] = *reinterpret_cast(&Out2); }; template <> __device__ inline void dequant_fp8_scales(int q, nv_bfloat162* frag_b) { constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; constexpr int MASK = 0x7F007F00; // Extract and shift FP8 values to BF16 format int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); q <<= 8; int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); // Note: reverse indexing is intentional because weights are permuted frag_b[1] = *reinterpret_cast(&Out1); frag_b[0] = *reinterpret_cast(&Out2); }; // New version with s_type_id parameter for marlin_moe_wna16_v2 template __device__ inline void dequant_fp8_scales(int q, scalar_t2* frag_b); template <> __device__ inline void dequant_fp8_scales(int q, half2* frag_b) { int Out1 = (q & 0xFF00FF00) >> 1; ; q <<= 8; int Out2 = (q & 0xFF00FF00) >> 1; // Note: reverse indexing is intentional because weights are permuted frag_b[1] = *reinterpret_cast(&Out1); frag_b[0] = *reinterpret_cast(&Out2); }; template <> __device__ inline void dequant_fp8_scales(int q, nv_bfloat162* frag_b) { constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; constexpr int MASK = 0x7F007F00; // Extract and shift FP8 values to BF16 format int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); q <<= 8; int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); // Note: reverse indexing is intentional because weights are permuted frag_b[1] = *reinterpret_cast(&Out1); frag_b[0] = *reinterpret_cast(&Out2); } template <> __device__ inline void dequant_fp8_scales(int q, nv_bfloat162* frag_b) { // In this conversion, 2 ** -127 in FP8E8M0 would become 0 in BF16, // but we assume that such a extreme value would not occur in real models. int Out1 = (q & 0xFF00FF00) >> 1; q <<= 7; int Out2 = q & 0x7F807F80; // Note: reverse indexing is intentional because weights are permuted frag_b[1] = *reinterpret_cast(&Out1); frag_b[0] = *reinterpret_cast(&Out2); } #endif } // namespace device::marlin ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/marlin/gptq_marlin.cuh ================================================ /* * Modified by Neural Magic * Copyright (C) Marlin.2024 Elias Frantar * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ /* * Adapted from https://github.com/IST-DASLab/marlin */ #pragma once #include #include #include "kernel.h" #include "marlin_template.h" namespace device::marlin { __global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){}; using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 __global__ void permute_cols_kernel( int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr, int4* __restrict__ out_int4_ptr, int size_m, int size_k, int lda, int block_rows) {} #else // For a given "a" of size [M,K] performs a permutation of the K columns based // on the given "perm" indices. __global__ void permute_cols_kernel( int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr, int4* __restrict__ out_int4_ptr, int size_m, int size_k, int lda, int block_rows) { auto start_row = block_rows * blockIdx.x; int finish_row = start_row + block_rows; if (finish_row > size_m) { finish_row = size_m; } int cur_block_rows = finish_row - start_row; int input_row_stride = lda * sizeof(half) / 16; int output_row_stride = size_k * sizeof(half) / 16; auto permute_row = [&](int row) { int iters = size_k / default_threads; int rest = size_k % default_threads; int input_offset = row * input_row_stride; int output_offset = row * output_row_stride; half const* a_row_half = reinterpret_cast(a_int4_ptr + input_offset); half* out_half = reinterpret_cast(out_int4_ptr + output_offset); int base_k = 0; for (int i = 0; i < iters; i++) { auto cur_k = base_k + threadIdx.x; int src_pos = perm_int_ptr[cur_k]; out_half[cur_k] = a_row_half[src_pos]; base_k += default_threads; } if (rest) { if (threadIdx.x < rest) { auto cur_k = base_k + threadIdx.x; int src_pos = perm_int_ptr[cur_k]; out_half[cur_k] = a_row_half[src_pos]; } } }; for (int i = 0; i < cur_block_rows; i++) { int cur_row = start_row + i; if (cur_row < size_m) { permute_row(cur_row); } } } typedef struct { int thread_k; int thread_n; int num_threads; } thread_config_t; thread_config_t small_batch_thread_configs[] = { // Ordered by priority // thread_k, thread_n, num_threads {128, 128, 256}, {64, 128, 128}, {128, 64, 128}}; thread_config_t large_batch_thread_configs[] = { // Ordered by priority // thread_k, thread_n, num_threads {64, 256, 256}, {64, 128, 128}, {128, 64, 128}}; typedef struct { int blocks_per_sm; thread_config_t tb_cfg; } exec_config_t; int get_scales_cache_size( thread_config_t const& th_config, int prob_m, int prob_n, int prob_k, int num_bits, int group_size, bool has_act_order, bool is_k_full) { bool cache_scales_chunk = has_act_order && !is_k_full; int tb_n = th_config.thread_n; int tb_k = th_config.thread_k; // Get max scale groups per thread-block int tb_groups; if (group_size == -1) { tb_groups = 1; } else if (group_size == 0) { tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size } else { tb_groups = div_ceil(tb_k, group_size); } if (cache_scales_chunk) { int load_groups = tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K load_groups = max(load_groups, 32); // We load at least 32 scale groups return load_groups * tb_n * 2; } else { int tb_scales = tb_groups * tb_n * 2; return tb_scales * pipe_stages; } } int get_kernel_cache_size( thread_config_t const& th_config, int thread_m_blocks, int prob_m, int prob_n, int prob_k, int num_bits, int group_size, bool has_act_order, bool is_k_full, int has_zp, int is_zp_float) { int pack_factor = 32 / num_bits; // Get B size int tb_k = th_config.thread_k; int tb_n = th_config.thread_n; int tb_m = thread_m_blocks * 16; int sh_a_size = pipe_stages * (tb_m * tb_k) * 2; int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4; int sh_red_size = tb_m * (tb_n + 8); int sh_s_size = get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full); int sh_g_idx_size = has_act_order && !is_k_full ? pipe_stages * tb_k / 4 : 0; int sh_zp_size = 0; if (has_zp) { if (is_zp_float) sh_zp_size = sh_s_size; else if (num_bits == 4) sh_zp_size = sh_s_size / 4; else if (num_bits == 8) sh_zp_size = sh_s_size / 2; } int total_size = max(sh_b_size, sh_red_size) + sh_a_size + sh_s_size + sh_zp_size + sh_g_idx_size; return total_size; } bool is_valid_config( thread_config_t const& th_config, int thread_m_blocks, int prob_m, int prob_n, int prob_k, int num_bits, int group_size, bool has_act_order, bool is_k_full, int has_zp, int is_zp_float, int max_shared_mem) { // Sanity if (th_config.thread_k == -1 || th_config.thread_n == -1 || th_config.num_threads == -1) { return false; } // Verify K/N are divisible by thread K/N if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) { return false; } // Verify min for thread K/N if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) { return false; } // num_threads must be at least 128 (= 4 warps) if (th_config.num_threads < 128) { return false; } // Check that pipeline fits into cache int cache_size = get_kernel_cache_size( th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float); return cache_size <= max_shared_mem; } #define _GET_IF( \ W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \ else if ( \ q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && thread_n_blocks == THREAD_N_BLOCKS && \ thread_k_blocks == THREAD_K_BLOCKS && m_block_size_8 == M_BLOCK_SIZE_8 && group_blocks == GROUP_BLOCKS && \ num_threads == NUM_THREADS && is_zp_float == IS_ZP_FLOAT) { \ kernel = Marlin< \ scalar_t, \ W_TYPE.id(), \ NUM_THREADS, \ THREAD_M_BLOCKS, \ THREAD_N_BLOCKS, \ THREAD_K_BLOCKS, \ M_BLOCK_SIZE_8, \ pipe_stages, \ GROUP_BLOCKS, \ IS_ZP_FLOAT>; \ } // COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false) // this is the most common cases // BIGGROUP: cases for big group size (group_blocks in [-1, 8]) // FZP: cases for float-zero-point (is_zp_float = true) // ACT: cases for act order case (group_blocks == 0) // FP4: cases for nvfp4(e2m1) (group_blocks == 1) #define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) #define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) #define COMMON_GET_IF(W_TYPE) \ COMMON_GET_IF_M1(W_TYPE, 8, 8, 256) \ COMMON_GET_IF_M1(W_TYPE, 8, 4, 128) \ COMMON_GET_IF_M1(W_TYPE, 4, 8, 128) \ COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \ COMMON_GET_IF_M234(W_TYPE, 8, 4, 128) \ COMMON_GET_IF_M234(W_TYPE, 4, 8, 128) #define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) #define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) #define BIGGROUP_GET_IF(W_TYPE) \ BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256) \ BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128) \ BIGGROUP_GET_IF_M1(W_TYPE, 4, 8, 128) \ BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \ BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128) \ BIGGROUP_GET_IF_M234(W_TYPE, 4, 8, 128) #define FP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) #define FP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) #define FP4_GET_IF(W_TYPE) \ FP4_GET_IF_M1(W_TYPE, 8, 8, 256) \ FP4_GET_IF_M1(W_TYPE, 8, 4, 128) \ FP4_GET_IF_M1(W_TYPE, 4, 8, 128) \ FP4_GET_IF_M234(W_TYPE, 16, 4, 256) \ FP4_GET_IF_M234(W_TYPE, 8, 4, 128) \ FP4_GET_IF_M234(W_TYPE, 4, 8, 128) // We currently have 4-bit models only with group_blocks == 4 #define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) #define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) #define FZP_GET_IF(W_TYPE) \ FZP_GET_IF_M1(W_TYPE, 8, 8, 256) \ FZP_GET_IF_M1(W_TYPE, 8, 4, 128) \ FZP_GET_IF_M1(W_TYPE, 4, 8, 128) \ FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \ FZP_GET_IF_M234(W_TYPE, 8, 4, 128) \ FZP_GET_IF_M234(W_TYPE, 4, 8, 128) // We currently have 4-bit models only with group_blocks == 4 #define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) #define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) #define ACT_GET_IF(W_TYPE) \ ACT_GET_IF_M1(W_TYPE, 8, 8, 256) \ ACT_GET_IF_M1(W_TYPE, 8, 4, 128) \ ACT_GET_IF_M1(W_TYPE, 4, 8, 128) \ ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \ ACT_GET_IF_M234(W_TYPE, 8, 4, 128) \ ACT_GET_IF_M234(W_TYPE, 4, 8, 128) template MarlinFuncPtr get_marlin_kernel( const host::ScalarType q_type, int thread_m_blocks, int thread_n_blocks, int thread_k_blocks, bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks, int num_threads, bool is_zp_float) { int num_bits = q_type.size_bits(); auto kernel = MarlinDefault; if (false) { } COMMON_GET_IF(host::kU4) COMMON_GET_IF(host::kU4B8) COMMON_GET_IF(host::kU8B128) FP4_GET_IF(host::kFE2M1f) BIGGROUP_GET_IF(host::kFE4M3fn) ACT_GET_IF(host::kU4B8) ACT_GET_IF(host::kU8B128) if (std::is_same::value) { if (false) { } FZP_GET_IF(host::kU4) } return kernel; } template exec_config_t determine_exec_config( const host::ScalarType& q_type, int prob_m, int prob_n, int prob_k, int thread_m_blocks, bool m_block_size_8, int num_bits, int group_size, bool has_act_order, bool is_k_full, bool has_zp, bool is_zp_float, int max_shared_mem, int sms) { exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}}; thread_config_t* thread_configs = thread_m_blocks > 1 ? large_batch_thread_configs : small_batch_thread_configs; int thread_configs_size = thread_m_blocks > 1 ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t) : sizeof(small_batch_thread_configs) / sizeof(thread_config_t); for (int i = 0; i < thread_configs_size; i++) { thread_config_t th_config = thread_configs[i]; if (!is_valid_config( th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float, max_shared_mem)) { continue; } int cache_size = get_kernel_cache_size( th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float); int group_blocks = 0; if (!has_act_order) { group_blocks = group_size == -1 ? -1 : group_size / 16; } auto kernel = get_marlin_kernel( q_type, thread_m_blocks, th_config.thread_n / 16, th_config.thread_k / 16, m_block_size_8, has_act_order, has_zp, group_blocks, th_config.num_threads, is_zp_float); if (kernel == MarlinDefault) continue; // int m_tiles = div_ceil(prob_m, thread_m_blocks * 16); // int n_tiles = prob_n / th_config.thread_n; // int k_tiles = prob_k / th_config.thread_k; return {1, th_config}; } return exec_cfg; } template void marlin_mm( const void* A, const void* B, void* C, void* C_tmp, void* s, void* s2, void* zp, void* g_idx, void* perm, void* a_tmp, int prob_m, int prob_n, int prob_k, int lda, void* workspace, host::ScalarType const& q_type, bool has_act_order, bool is_k_full, bool has_zp, int num_groups, int group_size, int dev, cudaStream_t stream, int thread_k_init, int thread_n_init, int sms, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { if (has_zp) { host::RuntimeCheck( q_type == host::kU4 || q_type == host::kU8, "q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str()); } else { host::RuntimeCheck( q_type == host::kU4B8 || q_type == host::kU8B128 || q_type == host::kFE4M3fn || q_type == host::kFE2M1f, "q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when " "has_zp = False. Got = ", q_type.str()); } host::RuntimeCheck( prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, ", ", prob_n, ", ", prob_k, "]"); int group_blocks = 0; if (has_act_order) { if (is_k_full) { host::RuntimeCheck(group_size != -1); group_blocks = group_size / 16; host::RuntimeCheck( prob_k % group_blocks == 0, "prob_k = ", prob_k, " is not divisible by group_blocks = ", group_blocks); } else { host::RuntimeCheck(group_size == 0); group_blocks = 0; } } else { if (group_size == -1) { group_blocks = -1; } else { group_blocks = group_size / 16; host::RuntimeCheck( prob_k % group_blocks == 0, "prob_k = ", prob_k, " is not divisible by group_blocks = ", group_blocks); } } int num_bits = q_type.size_bits(); const int4* A_ptr = (const int4*)A; const int4* B_ptr = (const int4*)B; int4* C_ptr = (int4*)C; int4* C_tmp_ptr = (int4*)C_tmp; const int4* s_ptr = (const int4*)s; const uint16_t* s2_ptr = (const uint16_t*)s2; const int4* zp_ptr = (const int4*)zp; const int* g_idx_ptr = (const int*)g_idx; const int* perm_ptr = (const int*)perm; int4* a_tmp_ptr = (int4*)a_tmp; int* locks = (int*)workspace; if (has_act_order) { // Permute A columns int block_rows = div_ceil(prob_m, sms); host::LaunchKernel(sms, default_threads, stream)( permute_cols_kernel, A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, lda, block_rows); A_ptr = a_tmp_ptr; lda = prob_k; // If we have a full K, then we can run the non-act-order version of Marlin // (since the weight rows are reordered by increasing group ids, and by // having a full K, we have full original groups) if (is_k_full) has_act_order = false; } int max_shared_mem = 0; host::RuntimeDeviceCheck(cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev)); host::RuntimeCheck(max_shared_mem > 0); int max_par = 16; if (prob_n <= 4096) max_par = 16 * 8; int max_shared_mem_new = max_shared_mem; int rest_m = prob_m; int max_thread_m_blocks = 4; while (rest_m) { int par_count = rest_m / (max_thread_m_blocks * 16); if (par_count > max_par) par_count = max_par; int prob_m_split = par_count > 0 ? (par_count * (max_thread_m_blocks * 16)) : rest_m; int thread_k = thread_k_init; int thread_n = thread_n_init; int thread_m_blocks = min(div_ceil(prob_m_split, 16), max_thread_m_blocks); int m_block_size_8 = prob_m_split <= 8; // Set thread config exec_config_t exec_cfg; thread_config_t thread_tfg; if (thread_k != -1 && thread_n != -1) { thread_tfg = thread_config_t{thread_k, thread_n, default_threads}; exec_cfg = exec_config_t{1, thread_tfg}; host::RuntimeCheck(prob_n % thread_n == 0, "prob_n = ", prob_n, " is not divisible by thread_n = ", thread_n); host::RuntimeCheck(prob_k % thread_k == 0, "prob_k = ", prob_k, " is not divisible by thread_k = ", thread_k); } else { // Auto config exec_cfg = determine_exec_config( q_type, prob_m_split, prob_n, prob_k, thread_m_blocks, m_block_size_8, num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float, max_shared_mem, sms); thread_tfg = exec_cfg.tb_cfg; if (thread_tfg.thread_k == -1 && max_thread_m_blocks > 1) { max_thread_m_blocks--; continue; } } int num_threads = thread_tfg.num_threads; thread_k = thread_tfg.thread_k; thread_n = thread_tfg.thread_n; int blocks = sms * exec_cfg.blocks_per_sm; if (exec_cfg.blocks_per_sm > 1) max_shared_mem_new = max_shared_mem / exec_cfg.blocks_per_sm - 1024; int thread_k_blocks = thread_k / 16; int thread_n_blocks = thread_n / 16; host::RuntimeCheck( is_valid_config( thread_tfg, thread_m_blocks, prob_m_split, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float, max_shared_mem_new), "Invalid thread config: thread_m_blocks = ", thread_m_blocks, ", thread_k = ", thread_tfg.thread_k, ", thread_n = ", thread_tfg.thread_n, ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits, ", prob_m_split = ", prob_m_split, ", group_size = ", group_size, ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full, ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float, ", max_shared_mem_new = ", max_shared_mem_new); auto kernel = get_marlin_kernel( q_type, thread_m_blocks, thread_n_blocks, thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks, num_threads, is_zp_float); if (kernel == MarlinDefault) { host::Panic( "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ", prob_k, "]", ", has_act_order = ", has_act_order, ", num_groups = ", num_groups, ", group_size = ", group_size, ", prob_m_split = ", prob_m_split, ", thread_m_blocks = ", thread_m_blocks, ", thread_n_blocks = ", thread_n_blocks, ", thread_k_blocks = ", thread_k_blocks, ", num_threads = ", num_threads, ", num_bits = ", num_bits); } host::RuntimeDeviceCheck( cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem_new)); bool part_use_atomic_add = use_atomic_add && div_ceil(prob_m_split, 64) * prob_n <= 2048; host::LaunchKernel(blocks, num_threads, stream, max_shared_mem_new)( kernel, A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, s2_ptr, zp_ptr, g_idx_ptr, num_groups, prob_m_split, prob_n, prob_k, lda, locks, part_use_atomic_add, use_fp32_reduce, max_shared_mem_new); A_ptr += prob_m_split * (lda / 8); C_ptr += prob_m_split * (prob_n / 8); rest_m -= prob_m_split; } } #endif } // namespace device::marlin template void gptq_marlin_gemm( tvm::ffi::TensorView a, tvm::ffi::TensorView b_q_weight, tvm::ffi::TensorView b_scales, tvm::ffi::TensorView global_scale, tvm::ffi::TensorView b_zeros, tvm::ffi::TensorView g_idx, tvm::ffi::TensorView perm, tvm::ffi::TensorView c, tvm::ffi::TensorView c_tmp, tvm::ffi::TensorView a_tmp, tvm::ffi::TensorView workspace, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { using namespace host; ScalarType const b_q_type = ScalarType::from_id(b_q_type_id); int pack_factor = 32 / b_q_type.size_bits(); // Bind symbolic sizes auto M = SymbolicSize{"M"}; auto K = SymbolicSize{"K"}; auto N = SymbolicSize{"N"}; auto device = SymbolicDevice{}; device.set_options(); // Verify a: [M, K] auto lda = SymbolicSize{"lda"}; TensorMatcher({M, K}).with_strides({lda, 1}).with_dtype().with_device(device).verify(a); int64_t size_m = M.unwrap(); int64_t size_k = K.unwrap(); // Verify b_q_weight: [K/tile_size, packed_N] RuntimeCheck( size_k % device::marlin::tile_size == 0, "size_k = ", size_k, " is not divisible by tile_size = ", device::marlin::tile_size); int64_t expected_bqw_dim0 = size_k / device::marlin::tile_size; auto bqw_dim0 = SymbolicSize{"bqw_dim0"}; auto bqw_dim1 = SymbolicSize{"bqw_dim1"}; bqw_dim0.set_value(expected_bqw_dim0); TensorMatcher({bqw_dim0, bqw_dim1}).with_dtype().with_device(device).verify(b_q_weight); RuntimeCheck( b_q_weight.size(1) % device::marlin::tile_size == 0, "b_q_weight.size(1) = ", b_q_weight.size(1), " is not divisible by tile_size = ", device::marlin::tile_size); int64_t actual_size_n = (b_q_weight.size(1) / device::marlin::tile_size) * pack_factor; N.set_value(actual_size_n); int64_t size_n = N.unwrap(); // Verify stride alignment int64_t a_stride0 = a.stride(0); RuntimeCheck(a_stride0 % 8 == 0, "a.stride(0) must be divisible by 8"); // Verify b_scales: [num_groups, N] auto num_groups_sym = SymbolicSize{"num_groups"}; TensorMatcher({num_groups_sym, N}).with_device(device).verify(b_scales); int num_groups = static_cast(num_groups_sym.unwrap()); // Verify c: [M, N] TensorMatcher({M, N}).with_dtype().with_device(device).verify(c); // Early return for zero-size M if (size_m == 0) return; // Determine has_act_order from g_idx/perm sizes int64_t g_idx_size = g_idx.size(0); int64_t perm_size = perm.size(0); bool has_act_order = g_idx_size > 0 && perm_size > 0; if (has_act_order) { RuntimeCheck( (g_idx_size == size_k && perm_size == size_k), "Unexpected g_idx.size(0) = ", g_idx_size, " and perm.size(0) = ", perm_size, ", where size_k = ", size_k); } // Determine has_zp from b_zeros size int64_t b_zeros_size = b_zeros.size(0); bool has_zp = b_zeros_size > 0; if (has_zp) { RuntimeCheck( b_q_type == kU4 || b_q_type == kU8, "b_q_type must be u4 or u8 when has_zp = True. Got = ", b_q_type.str()); } else { RuntimeCheck( b_q_type == kU4B8 || b_q_type == kU8B128 || b_q_type == kFE4M3fn || b_q_type == kFE2M1f, "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when " "has_zp = False. Got = ", b_q_type.str()); } if (has_zp && is_zp_float) { RuntimeCheck( std::is_same::value, "Computation type must be float16 (half) when using float zero points."); } // Verify b_zeros shape if (has_zp) { RuntimeCheck(b_zeros.dim() == 2, "b_zeros rank = ", b_zeros.dim(), " is not 2"); if (is_zp_float) { RuntimeCheck(b_zeros.size(1) == size_n, "b_zeros dim 1 = ", b_zeros.size(1), " is not size_n = ", size_n); RuntimeCheck( num_groups == b_zeros.size(0), "b_zeros dim 0 = ", b_zeros.size(0), " is not num_groups = ", num_groups); RuntimeCheck(num_groups != -1, "num_groups must be != -1"); } else { RuntimeCheck( b_zeros.size(0) == num_groups, "b_zeros dim 0 = ", b_zeros.size(0), " is not num_groups = ", num_groups); RuntimeCheck( b_zeros.size(1) == size_n / pack_factor, "b_zeros dim 1 = ", b_zeros.size(1), " is not size_n / pack_factor = ", size_n / pack_factor); } } // Verify global_scale int64_t global_scale_size = global_scale.size(0); if (global_scale_size > 0) { RuntimeCheck(b_q_type == kFE2M1f, "global_scale can only be used for float4_e2m1f."); } else { RuntimeCheck(!(b_q_type == kFE2M1f), "the global_scale parameter must be passed for float4_e2m1f."); } // Derive group_size int group_size = -1; if (has_act_order) { if (is_k_full) { RuntimeCheck(num_groups > 1, "For act_order, num_groups must be > 1"); RuntimeCheck(size_k % num_groups == 0, "size_k = ", size_k, ", is not divisible by num_groups = ", num_groups); group_size = static_cast(size_k / num_groups); } else { group_size = 0; } } else { if (num_groups > 1) { RuntimeCheck(size_k % num_groups == 0, "size_k = ", size_k, ", is not divisible by num_groups = ", num_groups); group_size = static_cast(size_k / num_groups); } else { group_size = -1; } } // Verify workspace and get device info RuntimeCheck( size_n % device::marlin::min_thread_n == 0, "size_n = ", size_n, ", is not divisible by min_thread_n = ", device::marlin::min_thread_n); DLDevice dl_device = device.unwrap(); int dev = dl_device.device_id; cudaStream_t stream = LaunchKernel::resolve_device(dl_device); int sms = -1; RuntimeDeviceCheck(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev)); RuntimeCheck( workspace.size(0) >= sms, "workspace.size(0) = ", workspace.size(0), " is below min_workspace_size = ", sms); // Hardcoded defaults (auto config) int thread_k_init = -1; int thread_n_init = -1; // Compute c_tmp and a_tmp pointers // c_tmp and a_tmp are pre-allocated by caller device::marlin::marlin_mm( a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(), b_scales.data_ptr(), global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr(), static_cast(size_m), static_cast(size_n), static_cast(size_k), static_cast(a_stride0), workspace.data_ptr(), b_q_type, has_act_order, is_k_full, has_zp, num_groups, group_size, dev, stream, thread_k_init, thread_n_init, sms, use_atomic_add, use_fp32_reduce, is_zp_float); } ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/marlin/gptq_marlin_repack.cuh ================================================ /* * Modified by Neural Magic * Copyright (C) Marlin.2024 Elias Frantar * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ /* * Adapted from https://github.com/IST-DASLab/marlin */ #pragma once #include #include #include "marlin.cuh" namespace device::marlin { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 template __global__ void gptq_marlin_repack_kernel( uint32_t const* __restrict__ b_q_weight_ptr, uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr, int size_k, int size_n) { return; } #else template __global__ void gptq_marlin_repack_kernel( uint32_t const* __restrict__ b_q_weight_ptr, uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr, int size_k, int size_n) { constexpr int pack_factor = 32 / num_bits; int k_tiles = size_k / tile_k_size; int n_tiles = size_n / tile_n_size; int block_k_tiles = div_ceil(k_tiles, gridDim.x); auto start_k_tile = blockIdx.x * block_k_tiles; if (start_k_tile >= k_tiles) { return; } int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); // Wait until the next thread tile has been loaded to shared memory. auto wait_for_stage = [&]() { // We only have `stages - 2` active fetches since we are double buffering // and can only issue the next fetch when it is guaranteed that the previous // shared memory load is fully complete (as it may otherwise be // overwritten). cp_async_wait(); __syncthreads(); }; extern __shared__ int4 sh[]; constexpr int perm_size = tile_k_size / 4; int4* sh_perm_ptr = sh; int4* sh_pipe_ptr = sh_perm_ptr; if constexpr (has_perm) { sh_pipe_ptr += perm_size; } constexpr int tile_ints = tile_k_size / pack_factor; constexpr int stage_n_threads = tile_n_size / 4; constexpr int stage_k_threads = has_perm ? tile_k_size : tile_ints; constexpr int stage_size = stage_k_threads * stage_n_threads; auto load_perm_to_shared = [&](int k_tile_id) { int first_k_int4 = (k_tile_id * tile_k_size) / 4; int4 const* perm_int4_ptr = reinterpret_cast(perm_ptr); if (threadIdx.x < perm_size) { sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x]; } __syncthreads(); }; auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { if (n_tile_id >= n_tiles) { cp_async_fence(); return; } int first_n = n_tile_id * tile_n_size; int4* sh_ptr = sh_pipe_ptr + stage_size * pipe; if constexpr (has_perm) { if (threadIdx.x < stage_size) { auto k_id = threadIdx.x / stage_n_threads; auto n_id = threadIdx.x % stage_n_threads; uint32_t const* sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); int src_k = sh_perm_int_ptr[k_id]; int src_k_packed = src_k / pack_factor; cp_async4( &sh_ptr[k_id * stage_n_threads + n_id], reinterpret_cast(&(b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)]))); } } else { if (threadIdx.x < stage_size) { auto k_id = threadIdx.x / stage_n_threads; auto n_id = threadIdx.x % stage_n_threads; int first_k = k_tile_id * tile_k_size; int first_k_packed = first_k / pack_factor; cp_async4( &sh_ptr[k_id * stage_n_threads + n_id], reinterpret_cast(&(b_q_weight_ptr[(first_k_packed + k_id) * size_n + first_n + (n_id * 4)]))); } } cp_async_fence(); }; auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { if (n_tile_id >= n_tiles) { return; } auto warp_id = threadIdx.x / 32; auto th_id = threadIdx.x % 32; if (warp_id >= 4) { return; } int tc_col = th_id / 4; int tc_row = (th_id % 4) * 2; constexpr int tc_offsets[4] = {0, 1, 8, 9}; int cur_n = warp_id * 16 + tc_col; constexpr int sh_stride = 64; constexpr uint32_t mask = (1 << num_bits) - 1; int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe; uint32_t* sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); uint32_t* sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); uint32_t vals[8]; if constexpr (has_perm) { for (int i = 0; i < 4; i++) { int k_idx = tc_row + tc_offsets[i]; uint32_t src_k = sh_perm_int_ptr[k_idx]; uint32_t src_k_pos = src_k % pack_factor; uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n]; uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask; uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8]; uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask; vals[i] = b1_cur_val; vals[4 + i] = b2_cur_val; } } else { uint32_t b1_vals[tile_ints]; uint32_t b2_vals[tile_ints]; #pragma unroll for (int i = 0; i < tile_ints; i++) { b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i]; b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i]; } #pragma unroll for (int i = 0; i < 4; i++) { int cur_elem = tc_row + tc_offsets[i]; int cur_int = cur_elem / pack_factor; int cur_pos = cur_elem % pack_factor; vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask; vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask; } } constexpr int tile_size = tile_k_size * tile_n_size / pack_factor; int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size; // Result of: // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h if constexpr (num_bits == 4) { constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; uint32_t res = 0; #pragma unroll for (int i = 0; i < 8; i++) { res |= vals[pack_idx[i]] << (i * 4); } out_ptr[out_offset + th_id * 4 + warp_id] = res; } else { constexpr int pack_idx[4] = {0, 2, 1, 3}; uint32_t res1 = 0; uint32_t res2 = 0; #pragma unroll for (int i = 0; i < 4; i++) { res1 |= vals[pack_idx[i]] << (i * 8); res2 |= vals[4 + pack_idx[i]] << (i * 8); } out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; } }; auto start_pipes = [&](int k_tile_id, int n_tile_id) { #pragma unroll for (int pipe = 0; pipe < repack_stages - 1; pipe++) { fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); } wait_for_stage(); }; #pragma unroll for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) { int n_tile_id = 0; if constexpr (has_perm) { load_perm_to_shared(k_tile_id); } start_pipes(k_tile_id, n_tile_id); while (n_tile_id < n_tiles) { #pragma unroll for (int pipe = 0; pipe < repack_stages; pipe++) { fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, n_tile_id + pipe + repack_stages - 1); repack_tile(pipe, k_tile_id, n_tile_id + pipe); wait_for_stage(); } n_tile_id += repack_stages; } } } #endif } // namespace device::marlin #define CALL_IF_REPACK(NUM_BITS, HAS_PERM) \ else if (num_bits == NUM_BITS && has_perm == HAS_PERM) { \ host::RuntimeDeviceCheck(cudaFuncSetAttribute( \ device::marlin::gptq_marlin_repack_kernel, \ cudaFuncAttributeMaxDynamicSharedMemorySize, \ max_shared_mem)); \ host::LaunchKernel(blocks, device::marlin::repack_threads, stream, static_cast(max_shared_mem))( \ device::marlin::gptq_marlin_repack_kernel, \ b_q_weight_ptr, \ perm_ptr, \ out_ptr, \ size_k, \ size_n); \ } void gptq_marlin_repack( tvm::ffi::TensorView b_q_weight, tvm::ffi::TensorView perm, tvm::ffi::TensorView out, int64_t size_k, int64_t size_n, int64_t num_bits) { using namespace host; // Validate num_bits RuntimeCheck(num_bits == 4 || num_bits == 8, "num_bits must be 4 or 8. Got = ", num_bits); int const pack_factor = 32 / static_cast(num_bits); // Validate size alignment RuntimeCheck( size_k % device::marlin::tile_k_size == 0, "size_k = ", size_k, " is not divisible by tile_k_size = ", device::marlin::tile_k_size); RuntimeCheck( size_n % device::marlin::tile_n_size == 0, "size_n = ", size_n, " is not divisible by tile_n_size = ", device::marlin::tile_n_size); // Validate b_q_weight auto bqw_dim0 = SymbolicSize{"bqw_dim0"}; auto bqw_dim1 = SymbolicSize{"bqw_dim1"}; bqw_dim0.set_value(size_k / pack_factor); bqw_dim1.set_value(size_n); auto device_ = SymbolicDevice{}; device_.set_options(); TensorMatcher({bqw_dim0, bqw_dim1}).with_dtype().with_device(device_).verify(b_q_weight); // Validate out auto out_dim0 = SymbolicSize{"out_dim0"}; auto out_dim1 = SymbolicSize{"out_dim1"}; out_dim0.set_value(size_k / device::marlin::tile_size); out_dim1.set_value(size_n * device::marlin::tile_size / pack_factor); TensorMatcher({out_dim0, out_dim1}).with_dtype().with_device(device_).verify(out); // Detect if there is act_order bool has_perm = perm.size(0) != 0; // Get ptrs uint32_t const* b_q_weight_ptr = reinterpret_cast(b_q_weight.data_ptr()); uint32_t const* perm_ptr = reinterpret_cast(perm.data_ptr()); uint32_t* out_ptr = reinterpret_cast(out.data_ptr()); // Get dev info DLDevice dl_device = device_.unwrap(); int dev = dl_device.device_id; cudaStream_t stream = LaunchKernel::resolve_device(dl_device); int blocks; RuntimeDeviceCheck(cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev)); int max_shared_mem = 0; RuntimeDeviceCheck(cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev)); RuntimeCheck(max_shared_mem > 0, "max_shared_mem must be > 0"); if (false) { } CALL_IF_REPACK(4, false) CALL_IF_REPACK(4, true) CALL_IF_REPACK(8, false) CALL_IF_REPACK(8, true) else { Panic("Unsupported repack config: num_bits = ", num_bits, ", has_perm = ", has_perm); } } #undef CALL_IF_REPACK ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/marlin/kernel.h ================================================ #include #include "marlin.cuh" #include "marlin_dtypes.cuh" #define MARLIN_KERNEL_PARAMS \ const int4 *__restrict__ A, const int4 *__restrict__ B, int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ const int4 *__restrict__ scales_ptr, const uint16_t *__restrict__ scale2_ptr, const int4 *__restrict__ zp_ptr, \ const int *__restrict__ g_idx, int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \ bool use_atomic_add, bool use_fp32_reduce, int max_shared_mem namespace device::marlin { template < typename scalar_t, // compute dtype, half or nv_float16 const host::ScalarTypeId w_type_id, // weight ScalarType id const int threads, // number of threads in a threadblock const int thread_m_blocks, // number of 16x16 blocks in the m // dimension (batchsize) of the // threadblock const int thread_n_blocks, // same for n dimension (output) const int thread_k_blocks, // same for k dimension (reduction) const bool m_block_size_8, // whether m_block_size == 8 // only works when thread_m_blocks == 1 const int stages, // number of stages for the async global->shared // fetch pipeline const int group_blocks, // number of consecutive 16x16 blocks // with a separate quantization scale const bool is_zp_float // is zero point of float16 type? > __global__ void Marlin(MARLIN_KERNEL_PARAMS); } // namespace device::marlin ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/marlin/marlin.cuh ================================================ #pragma once #include #include namespace device::marlin { // Marlin params // 8 warps are a good choice since every SM has 4 schedulers and having more // than 1 warp per schedule allows some more latency hiding. At the same time, // we want relatively few warps to have many registers per warp and small tiles. static constexpr int default_threads = 256; static constexpr int pipe_stages = 4; // 4 pipeline stages fit into shared memory static constexpr int min_thread_n = 64; static constexpr int min_thread_k = 64; static constexpr int max_thread_n = 256; static constexpr int tile_size = 16; static constexpr int max_par = 16; // Repack params static constexpr int repack_stages = 8; static constexpr int repack_threads = 256; static constexpr int tile_k_size = tile_size; static constexpr int tile_n_size = tile_k_size * 4; // Helpers template struct Vec { T elems[n]; __device__ T& operator[](int i) { return elems[i]; } }; using I4 = Vec; using host::div_ceil; #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 // No support for async #else __device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr, bool pred = true) { const int BYTES = 16; uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); asm volatile( "{\n" " .reg .pred p;\n" " setp.ne.b32 p, %0, 0;\n" " @p cp.async.cg.shared.global [%1], [%2], %3;\n" "}\n" ::"r"((int)pred), "r"(smem), "l"(glob_ptr), "n"(BYTES)); } __device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) { const int BYTES = 16; uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); asm volatile( "{\n" " cp.async.cg.shared.global [%0], [%1], %2;\n" "}\n" ::"r"(smem), "l"(glob_ptr), "n"(BYTES)); } __device__ inline void cp_async_fence() { asm volatile("cp.async.commit_group;\n" ::); } template __device__ inline void cp_async_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); } #endif } // namespace device::marlin ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/marlin/marlin_dtypes.cuh ================================================ #ifndef _data_types_cuh #define _data_types_cuh #include #include "marlin.cuh" namespace device::marlin { template class ScalarType {}; template <> class ScalarType { public: using scalar_t = fp16_t; using scalar_t2 = fp16x2_t; // Matrix fragments for tensor core instructions; their precise layout is // documented here: // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type using FragA = Vec; using FragB = Vec; using FragC = Vec; using FragS = Vec; using FragZP = Vec; static __device__ float inline num2float(const fp16_t x) { return __half2float(x); } static __device__ fp16x2_t inline num2num2(const fp16_t x) { return __half2half2(x); } static __device__ fp16x2_t inline nums2num2(const fp16_t x1, const fp16_t x2) { return __halves2half2(x1, x2); } static __host__ __device__ fp16_t inline float2num(const float x) { return __float2half(x); } }; template <> class ScalarType { public: using scalar_t = bf16_t; using scalar_t2 = bf16x2_t; using FragA = Vec; using FragB = Vec; using FragC = Vec; using FragS = Vec; using FragZP = Vec; #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 static __device__ float inline num2float(const bf16_t x) { return __bfloat162float(x); } static __device__ bf16x2_t inline num2num2(const bf16_t x) { return __bfloat162bfloat162(x); } static __device__ bf16x2_t inline nums2num2(const bf16_t x1, const bf16_t x2) { return __halves2bfloat162(x1, x2); } static __host__ __device__ bf16_t inline float2num(const float x) { return __float2bfloat16(x); } #endif }; } // namespace device::marlin #endif ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/marlin/marlin_template.h ================================================ /* * Modified by Neural Magic * Copyright (C) Marlin.2024 Elias Frantar * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ /* * Adapted from https://github.com/IST-DASLab/marlin */ #include #include "dequant.h" #include "marlin.cuh" #include "marlin_dtypes.cuh" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert( \ std::is_same::value || std::is_same::value, \ "only float16 and bfloat16 is supported"); namespace device::marlin { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 template < typename scalar_t, // compute dtype, half or nv_float16 const host::ScalarTypeId w_type_id, // weight ScalarType id const int threads, // number of threads in a threadblock const int thread_m_blocks, // number of 16x16 blocks in the m // dimension (batchsize) of the // threadblock const int thread_n_blocks, // same for n dimension (output) const int thread_k_blocks, // same for k dimension (reduction) const bool m_block_size_8, // whether m_block_size == 8 // only works when thread_m_blocks == 1 const int stages, // number of stages for the async global->shared // fetch pipeline const bool has_act_order, // whether act_order is enabled const int group_blocks, // number of consecutive 16x16 blocks // with a separate quantization scale const bool is_zp_float // is zero point of float16 type? > __global__ void Marlin( const int4* __restrict__ A, // fp16 input matrix of shape mxk const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn int4* __restrict__ C, // fp16 output buffer of shape mxn int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape // (k/groupsize)xn const int* __restrict__ g_idx, // int32 group indices of shape k int num_groups, // number of scale groups per output channel int prob_m, // batch dimension m int prob_n, // output dimension n int prob_k, // reduction dimension k int* locks, // extra global storage for barrier synchronization bool use_fp32_reduce // whether to use fp32 global reduce ) {} } // namespace device::marlin #else // m16n8k16 tensor core mma instruction with fp16 inputs and fp32 // output/accumulation. template __device__ inline void mma(const typename ScalarType::FragA& a_frag, const typename ScalarType::FragB& frag_b, typename ScalarType::FragC& frag_c) { const uint32_t* a = reinterpret_cast(&a_frag); const uint32_t* b = reinterpret_cast(&frag_b); float* c = reinterpret_cast(&frag_c); if constexpr (std::is_same::value) { asm volatile( "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); } else if constexpr (std::is_same::value) { asm volatile( "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); } else { STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); } } template __device__ inline void mma_trans( const typename ScalarType::FragA& a_frag, const typename ScalarType::FragB& frag_b, const typename ScalarType::FragB& frag_b2, typename ScalarType::FragC& frag_c) { const uint32_t* a = reinterpret_cast(&a_frag); const uint32_t* b = reinterpret_cast(&frag_b); const uint32_t* b2 = reinterpret_cast(&frag_b2); float* c = reinterpret_cast(&frag_c); if constexpr (std::is_same::value) { asm volatile( "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); } else if constexpr (std::is_same::value) { asm volatile( "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); } else { STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); } } // Instruction for loading a full 16x16 matrix fragment of operand A from shared // memory, directly in tensor core layout. template __device__ inline void ldsm(typename ScalarType::FragA& frag_a, const void* smem_ptr) { uint32_t* a = reinterpret_cast(&frag_a); uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); if constexpr (count == 4) { asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) : "r"(smem)); } else if constexpr (count == 2) { asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" : "=r"(a[0]), "=r"(a[1]) : "r"(smem)); } else if constexpr (count == 1) { asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" : "=r"(a[0]) : "r"(smem)); } else { static_assert(count == 1 || count == 2 || count == 4, "invalid count"); } } // Multiply dequantized values by the corresponding quantization scale; used // only for grouped quantization. template __device__ inline void scale(typename ScalarType::FragB& frag_b, typename ScalarType::FragS& frag_s, int i) { using scalar_t2 = typename ScalarType::scalar_t2; scalar_t2 s = ScalarType::num2num2(reinterpret_cast(&frag_s)[i]); frag_b[0] = __hmul2(frag_b[0], s); frag_b[1] = __hmul2(frag_b[1], s); } template __device__ inline void scale_and_sub(typename ScalarType::FragB& frag_b, scalar_t s, scalar_t zp) { using scalar_t2 = typename ScalarType::scalar_t2; scalar_t2 s2 = ScalarType::num2num2(s); scalar_t2 zp2 = ScalarType::num2num2(zp); frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2)); frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2)); } template __device__ inline void sub_zp(typename ScalarType::FragB& frag_b, typename ScalarType::scalar_t2& frag_zp, int i) { using scalar_t2 = typename ScalarType::scalar_t2; scalar_t2 zp = ScalarType::num2num2(reinterpret_cast(&frag_zp)[i]); frag_b[0] = __hsub2(frag_b[0], zp); frag_b[1] = __hsub2(frag_b[1], zp); } // Same as above, but for act_order (each K is multiplied individually) template __device__ inline void scale4( typename ScalarType::FragB& frag_b, typename ScalarType::FragS& frag_s_1, typename ScalarType::FragS& frag_s_2, typename ScalarType::FragS& frag_s_3, typename ScalarType::FragS& frag_s_4, int i) { using scalar_t2 = typename ScalarType::scalar_t2; scalar_t2 s_val_1_2; s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; scalar_t2 s_val_3_4; s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; frag_b[0] = __hmul2(frag_b[0], s_val_1_2); frag_b[1] = __hmul2(frag_b[1], s_val_3_4); } // Given 2 floats multiply by 2 scales (halves) template __device__ inline void scale_float(float* c, typename ScalarType::FragS& s) { scalar_t* s_ptr = reinterpret_cast(&s); c[0] = __fmul_rn(c[0], ScalarType::num2float(s_ptr[0])); c[1] = __fmul_rn(c[1], ScalarType::num2float(s_ptr[1])); } // Wait until barrier reaches `count`, then lock for current threadblock. __device__ inline void barrier_acquire(int* lock, int count) { if (threadIdx.x == 0) { int state = -1; do // Guarantee that subsequent writes by this threadblock will be visible // globally. asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); while (state != count); } __syncthreads(); } // Release barrier and increment visitation count. __device__ inline void barrier_release(int* lock, bool reset = false) { __syncthreads(); if (threadIdx.x == 0) { if (reset) { lock[0] = 0; return; } int val = 1; // Make sure that all writes since acquiring this barrier are visible // globally, while releasing the barrier. asm volatile("fence.acq_rel.gpu;\n"); asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" : : "l"(lock), "r"(val)); } } // Wait until value of lock to be negative, and then add 1 __device__ inline void wait_negative_and_add(int* lock) { if (threadIdx.x == 0) { int state = 0; do // Guarantee that subsequent writes by this threadblock will be visible // globally. asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); while (state >= 0); atomicAdd(lock, 1); } __syncthreads(); } template < typename scalar_t, // compute dtype, half or nv_float16 const host::ScalarTypeId w_type_id, // weight ScalarType id const int threads, // number of threads in a threadblock const int thread_m_blocks, // number of 16x16 blocks in the m // dimension (batchsize) of the // threadblock const int thread_n_blocks, // same for n dimension (output) const int thread_k_blocks, // same for k dimension (reduction) const bool m_block_size_8, // whether m_block_size == 8 // only works when thread_m_blocks == 1 const int stages, // number of stages for the async global->shared // fetch pipeline const int group_blocks, // number of consecutive 16x16 blocks // with a separate quantization scale const bool is_zp_float // is zero point of float16 type? > __global__ void Marlin( const int4* __restrict__ A, // fp16 input matrix of shape mxk const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn int4* __restrict__ C, // fp16 output buffer of shape mxn int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape // (k/groupsize)xn const uint16_t* __restrict__ scale2_ptr, // fp16 global scale (for nvfp4 // only) const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape // (k/groupsize)x(n/pack_factor) const int* __restrict__ g_idx, // int32 group indices of shape k int num_groups, // number of scale groups per output channel int prob_m, // batch dimension m int prob_n, // output dimension n int prob_k, // reduction dimension k int lda, // A.stride(0), equal to prob_k is A is contiguous int* locks, // extra global storage for barrier synchronization bool use_atomic_add, // whether to use atomic add to reduce bool use_fp32_reduce, // whether to use fp32 global reduce int max_shared_mem) { // Each threadblock processes one "stripe" of the B matrix with (roughly) the // same size, which might involve multiple column "slices" (of width 16 * // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM // example: // 0 1 3 // 0 2 3 // 1 2 4 // While this kind of partitioning makes things somewhat more complicated, it // ensures good utilization of all SMs for many kinds of shape and GPU // configurations, while requiring as few slow global cross-threadblock // reductions as possible. using Dtype = ScalarType; using scalar_t2 = typename ScalarType::scalar_t2; using FragA = typename ScalarType::FragA; using FragB = typename ScalarType::FragB; using FragC = typename ScalarType::FragC; using FragS = typename ScalarType::FragS; using FragZP = typename ScalarType::FragZP; static constexpr auto w_type = host::ScalarType::from_id(w_type_id); constexpr bool has_zp = w_type == host::kU4 || w_type == host::kU8; constexpr bool is_int_type = w_type == host::kU4 || w_type == host::kU8 || w_type == host::kU4B8 || w_type == host::kU8B128; // see comments of dequant.h for more details constexpr bool dequant_skip_flop = !is_int_type || has_zp && !is_zp_float && !std::is_same::value || has_zp && !is_zp_float && !(w_type == host::kU8); scalar_t2 global_scale; if constexpr (w_type == host::kFE2M1f) { uint16_t val = scale2_ptr[0]; global_scale = Dtype::num2num2(*reinterpret_cast(&val)); } constexpr bool has_act_order = group_blocks == 0; constexpr int m_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); constexpr int pack_factor = 32 / w_type.size_bits(); static_assert(thread_m_blocks == 1 || !m_block_size_8); // For larger GEMMs we run multiple batchsize 64 versions in parallel for a // better partitioning with less reductions int parallel = 1; if (prob_m > m_block_size) { parallel = prob_m / m_block_size; prob_m = m_block_size; } int k_tiles = prob_k / 16 / thread_k_blocks; int n_tiles = prob_n / 16 / thread_n_blocks; int iters = div_ceil(k_tiles * n_tiles * parallel, gridDim.x); if constexpr (!has_act_order && group_blocks != -1) { if (group_blocks >= thread_k_blocks) { // Ensure that the number of tiles in each stripe is a multiple of the // groupsize; this avoids an annoying special case where a stripe starts // in the middle of group. iters = (group_blocks / thread_k_blocks) * div_ceil(iters, (group_blocks / thread_k_blocks)); } } int slice_row = (iters * blockIdx.x) % k_tiles; int slice_col_par = (iters * blockIdx.x) / k_tiles; int slice_col = slice_col_par; int slice_iters; // number of threadblock tiles in the current slice int slice_count = 0; // total number of active threadblocks in the current slice int slice_idx; // index of threadblock in current slice; numbered bottom to // top int par_id = 0; int locks_off = 0; // We can easily implement parallel problem execution by just remapping // indices and advancing global pointers if (slice_col_par >= n_tiles) { A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * lda / 8; C += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 8; slice_col = slice_col_par % n_tiles; par_id = slice_col_par / n_tiles; } if (parallel * n_tiles >= gridDim.x) { // when parallel * n_tiles >= sms // then there are at most $sms$ conflict tile blocks locks_off = blockIdx.x; } else { locks_off = (iters * blockIdx.x) / k_tiles - 1; } // Compute all information about the current slice which is required for // synchronization. auto init_slice = [&](bool first_init = false) { slice_iters = iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0; if (slice_iters == 0) return; if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row; slice_count = 1; slice_idx = 0; int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); if (col_first <= k_tiles * (slice_col_par + 1)) { int col_off = col_first - k_tiles * slice_col_par; slice_count = div_ceil(k_tiles - col_off, iters); if (col_off > 0) slice_count++; int delta_first = iters * blockIdx.x - col_first; if (delta_first < 0 || (col_off == 0 && delta_first == 0)) slice_idx = slice_count - 1; else { slice_idx = slice_count - 1 - delta_first / iters; if (col_off > 0) slice_idx--; } } if (parallel * n_tiles >= gridDim.x) { if (slice_count > 1 && slice_idx == slice_count - 1) { locks_off++; } } else { locks_off++; } if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) { constexpr int threads_per_m = 16 * thread_n_blocks / 8; int m_per_thread = div_ceil(thread_m_blocks * 16, threads / threads_per_m); if (m_block_size_8) m_per_thread = div_ceil(8, threads / threads_per_m); for (int i = 0; i < m_per_thread; i++) { int row = threads / threads_per_m * i + threadIdx.x / threads_per_m; if (row < prob_m) { int col = slice_col * 16 * thread_n_blocks / 8 + threadIdx.x % threads_per_m; C[row * prob_n / 8 + col] = {0, 0, 0, 0}; } } // After write zero to output, write a negative value to lock. // Every SM that processes the same slice would wait for // the negative value, and then atomicAdd 1 to it. // After all SMs are processed, the lock value would back to 0 again. __syncthreads(); if (threadIdx.x == 0) locks[locks_off] = 1 - slice_count; } if (slice_col == n_tiles) { A += 16 * thread_m_blocks * lda / 8; C += 16 * thread_m_blocks * prob_n / 8; slice_col = 0; par_id++; } }; init_slice(true); // A sizes/strides // stride of the A matrix in global memory int a_gl_stride = lda / 8; // stride of an A matrix tile in shared memory constexpr int a_sh_stride = 16 * thread_k_blocks / 8; // delta between subsequent A tiles in global memory constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / 8; // between subsequent accesses within a tile int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); // between shared memory writes constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); // between shared memory tile reads constexpr int a_sh_rd_delta_o = 2 * ((threads / 32) / (thread_n_blocks / 4)); // within a shared memory tile constexpr int a_sh_rd_delta_i = a_sh_stride * 16; // overall size of a tile constexpr int a_sh_stage = a_sh_stride * m_block_size; // number of shared write iterations for a tile constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); // B sizes/strides int b_gl_stride = 16 * prob_n / (pack_factor * 4); constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4; constexpr int b_thread_vecs = w_type.size_bits() == 4 ? 1 : 2; constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks; int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads); constexpr int b_sh_wr_delta = threads * b_thread_vecs; constexpr int b_sh_rd_delta = threads * b_thread_vecs; constexpr int b_sh_stage = b_sh_stride * thread_k_blocks; constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; // Scale sizes/strides without act_order int s_gl_stride = prob_n / 8; constexpr int s_sh_stride = 16 * thread_n_blocks / 8; constexpr int s_tb_groups = !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks ? thread_k_blocks / group_blocks / (w_type == host::kFE2M1f ? 2 : 1) : 1; constexpr int s_sh_stage = s_tb_groups * s_sh_stride; int s_gl_rd_delta = s_gl_stride; // Scale size/strides with act_order constexpr int tb_k = 16 * thread_k_blocks; constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; // constexpr int act_s_row_stride = 1; // int act_s_col_stride = act_s_row_stride * num_groups; constexpr int act_s_max_num_groups = 32; int act_s_col_stride = 1; int act_s_col_warp_stride = act_s_col_stride * 8; int tb_n_warps = thread_n_blocks / 4; int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; // Zero-points sizes/strides int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4; constexpr int zp_sh_stride = is_zp_float ? 16 * thread_n_blocks / 8 : ((16 * thread_n_blocks) / pack_factor) / 4; constexpr int zp_tb_groups = s_tb_groups; constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0; int zp_gl_rd_delta = zp_gl_stride; // Global A read index of current thread. int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); a_gl_rd += a_gl_rd_delta_o * slice_row; // Shared write index of current thread. int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); // Shared read index. int a_sh_rd = a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) + (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1)); a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4)); int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) + (threadIdx.x % b_sh_stride_threads) * b_thread_vecs; b_gl_rd += b_sh_stride * slice_col; b_gl_rd += b_gl_rd_delta_o * slice_row; auto b_sh_wr = threadIdx.x * b_thread_vecs; auto b_sh_rd = threadIdx.x * b_thread_vecs; // For act_order constexpr int k_iter_size = tb_k / b_sh_wr_iters; int slice_k_start = tb_k * slice_row; int slice_k_finish = slice_k_start + tb_k * slice_iters; int slice_k_start_shared_fetch = slice_k_start; int slice_n_offset = act_s_col_tb_stride * slice_col; // No act_order int s_gl_rd; if constexpr (!has_act_order) { if constexpr (group_blocks == -1) { s_gl_rd = s_sh_stride * slice_col + threadIdx.x; } else { s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) / (w_type == host::kFE2M1f ? 2 : 1) + s_sh_stride * slice_col + threadIdx.x; } } auto s_sh_wr = threadIdx.x; bool s_sh_wr_pred = threadIdx.x < s_sh_stride; // Zero-points int zp_gl_rd; if constexpr (has_zp) { if constexpr (group_blocks == -1) { zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; } else { zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + zp_sh_stride * slice_col + threadIdx.x; } } auto zp_sh_wr = threadIdx.x; bool zp_sh_wr_pred = threadIdx.x < zp_sh_stride; // We use a different scale layout for grouped and column-wise quantization as // we scale a `half2` tile in column-major layout in the former and in // row-major in the latter case. int s_sh_rd; if constexpr (group_blocks != -1 && w_type == host::kFE2M1f) { auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; int warp_row = warp_id / n_warps; s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 4; s_sh_rd = s_sh_rd * 2 + warp_row % 2; } else if constexpr (group_blocks != -1) s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 4; else if constexpr (group_blocks == -1 && (m_block_size_8 || (has_zp && !dequant_skip_flop))) s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 8; else s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) % 4; // Zero-points have the same read layout as the scales // (without column-wise case) constexpr int num_col_threads = 8; constexpr int num_row_threads = 4; constexpr int num_ints_per_thread = 8 / pack_factor; int zp_sh_rd; if constexpr (has_zp) { if constexpr (is_zp_float) { if constexpr (group_blocks != -1) { zp_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 4; } } else { zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); } } // Precompute which thread should not read memory in which iterations; this is // needed if there are more threads than required for a certain tilesize or // when the batchsize is not a multiple of 16. bool a_sh_wr_pred[a_sh_wr_iters]; #pragma unroll for (int i = 0; i < a_sh_wr_iters; i++) a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m; // To ensure that writing and reading A tiles to/from shared memory, the // latter in fragment format, is fully bank conflict free, we need to use a // rather fancy XOR-based layout. The key here is that neither reads nor // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the // same shared memory banks. Further, it seems (based on NSight-Compute) that // each warp must also write a consecutive memory segment? auto transform_a = [&](int i) { int row = i / a_gl_rd_delta_o; return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8); }; // Since the computation of this remapping is non-trivial and, due to our main // loop unrolls, all shared memory accesses are static, we simply precompute // both transformed reads and writes. int a_sh_wr_trans[a_sh_wr_iters]; #pragma unroll for (int i = 0; i < a_sh_wr_iters; i++) a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; #pragma unroll for (int i = 0; i < b_sh_wr_iters; i++) { #pragma unroll for (int j = 0; j < thread_m_blocks; j++) a_sh_rd_trans[i][j] = transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd); } // Since B-accesses have non-constant stride they have to be computed at // runtime; we break dependencies between subsequent accesses with a tile by // maintining multiple pointers (we have enough registers), a tiny // optimization. const int4* B_ptr[b_sh_wr_iters]; #pragma unroll for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd; extern __shared__ int4 sh[]; // Shared memory storage for global fetch pipelines. constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks; constexpr int sh_b_size = stages * b_sh_stage; int4* sh_b = sh; int4* sh_red = sh; int4* sh_g_idx = sh_b + (sh_red_size > sh_b_size ? sh_red_size : sh_b_size); int4* sh_zp = sh_g_idx + (stages * g_idx_stage); constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride) : (stages * s_sh_stage); int4* sh_s = sh_zp + (stages * zp_sh_stage); // shared memory reused by reduction should be smaller than // shared memory used by weight. static_assert(thread_m_blocks * 16 * thread_n_blocks * 16 / 8 <= stages * b_sh_stage); int4* sh_a = sh_s + sh_s_size; // constexpr int shm_size_used = // stages * (g_idx_stage + zp_sh_stage) + sh_s_size + // (sh_red_size > sh_b_size ? sh_red_size : sh_b_size); // Register storage for double buffer of shared memory reads. FragA frag_a[2][thread_m_blocks]; I4 frag_b_quant[2][b_thread_vecs]; FragC frag_c[thread_m_blocks][4][2]; FragS frag_s[2][4]; // No act-order FragS act_frag_s[2][4][4]; // For act-order int frag_qzp[2][num_ints_per_thread]; // Zero-points FragZP frag_zp; // Zero-points in fp16 FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ // Zero accumulators. auto zero_accums = [&]() { #pragma unroll for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) reinterpret_cast(frag_c)[i] = 0; }; int sh_first_group_id = -1; int sh_num_groups = -1; auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id, int last_group_id) { sh_first_group_id = first_group_id; sh_num_groups = last_group_id - first_group_id + 1; if (sh_num_groups > act_s_max_num_groups) { sh_num_groups = act_s_max_num_groups; } if (sh_first_group_id + sh_num_groups > num_groups) { sh_num_groups = num_groups - sh_first_group_id; } int row_offset = first_group_id * s_gl_stride; if (is_async) { for (int i = 0; i < sh_num_groups; i++) { if (threadIdx.x < s_sh_stride) { cp_async4_pred( &sh_s[(i * s_sh_stride) + threadIdx.x], &scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]); } } } else { for (int i = 0; i < sh_num_groups; i++) { if (threadIdx.x < s_sh_stride) { sh_s[(i * s_sh_stride) + threadIdx.x] = scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]; } } } }; // Asynchronously fetch the next A, B and s tile from global to the next // shared memory pipeline location. auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { if (pred) { int4* sh_a_stage = sh_a + a_sh_stage * pipe; #pragma unroll for (int i = 0; i < a_sh_wr_iters; i++) { cp_async4_pred( &sh_a_stage[a_sh_wr_trans[i]], &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off], a_sh_wr_pred[i]); } int4* sh_b_stage = sh_b + b_sh_stage * pipe; #pragma unroll for (int i = 0; i < b_sh_wr_iters; i++) { #pragma unroll for (int j = 0; j < b_thread_vecs; j++) { cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j], B_ptr[i] + j); } B_ptr[i] += b_gl_rd_delta_o; } if constexpr (has_act_order) { // Fetch g_idx thread-block portion int full_pipe = a_off; int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; if (cur_k < prob_k && cur_k < slice_k_finish) { int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; int4 const* cur_g_idx_stage_ptr = reinterpret_cast(&g_idx[cur_k]); if (threadIdx.x < g_idx_stage) { cp_async4_pred(&sh_g_idx_stage[threadIdx.x], &cur_g_idx_stage_ptr[threadIdx.x]); } } } else { if constexpr (group_blocks != -1) { int4* sh_s_stage = sh_s + s_sh_stage * pipe; if constexpr (group_blocks >= thread_k_blocks) { // Only fetch scales if this tile starts a new group if (pipe % (group_blocks / thread_k_blocks) == 0) { if (s_sh_wr_pred) { cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); } s_gl_rd += s_gl_rd_delta; } } else { for (int i = 0; i < s_tb_groups; i++) { if (s_sh_wr_pred) { cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr], &scales_ptr[s_gl_rd]); } s_gl_rd += s_gl_rd_delta; } } } if constexpr (has_zp && group_blocks != -1) { int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; if constexpr (group_blocks >= thread_k_blocks) { // Only fetch zero-points if this tile starts a new group if (pipe % (group_blocks / thread_k_blocks) == 0) { if (zp_sh_wr_pred) { cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); } zp_gl_rd += zp_gl_rd_delta; } } else { for (int i = 0; i < zp_tb_groups; i++) { if (zp_sh_wr_pred) { cp_async4(&sh_zp_stage[i * zp_sh_stride + zp_sh_wr], &zp_ptr[zp_gl_rd]); } zp_gl_rd += zp_gl_rd_delta; } } } } } // Insert a fence even when we are winding down the pipeline to ensure that // waiting is also correct at this point. cp_async_fence(); }; auto fetch_col_zp_to_shared = [&]() { if (zp_sh_wr_pred) { cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]); } }; auto fetch_col_scale_to_shared = [&]() { if (s_sh_wr_pred) { cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); } }; // Wait until the next thread tile has been loaded to shared memory. auto wait_for_stage = [&]() { // We only have `stages - 2` active fetches since we are double buffering // and can only issue the next fetch when it is guaranteed that the previous // shared memory load is fully complete (as it may otherwise be // overwritten). cp_async_wait(); __syncthreads(); }; // Load the next sub-tile from the current location in the shared memory pipe // into the current register buffer. auto fetch_to_registers = [&](int k, int pipe) { int4* sh_a_stage = sh_a + a_sh_stage * pipe; #pragma unroll for (int i = 0; i < thread_m_blocks; i++) ldsm(frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); int4* sh_b_stage = sh_b + b_sh_stage * pipe; #pragma unroll for (int i = 0; i < b_thread_vecs; i++) { frag_b_quant[k % 2][i] = *reinterpret_cast(&sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]); } }; bool is_same_group[stages]; int same_group_id[stages]; auto init_same_group = [&](int pipe) { if constexpr (!has_act_order) { return; } int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); int group_id_1 = sh_g_idx_int_ptr[0]; int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; is_same_group[pipe] = group_id_1 == group_id_2; same_group_id[pipe] = group_id_1; }; auto fetch_scales_to_registers = [&](int k, int full_pipe) { int pipe = full_pipe % stages; if constexpr (!has_act_order) { // No act-order case if constexpr (group_blocks == -1) { // load only when starting a new slice if (k == 0 && full_pipe == 0) { reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; } } else if constexpr (group_blocks != -1) { if constexpr (group_blocks >= thread_k_blocks) { if (k % b_sh_wr_iters == 0) { int4* sh_s_stage = sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; } else { reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; } } else { auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; int warp_row = warp_id / n_warps; int cur_k = warp_row * 16; cur_k += k_iter_size * (k % b_sh_wr_iters); int k_blocks = cur_k / 16; int cur_group_id = k_blocks / (group_blocks * (w_type == host::kFE2M1f ? 2 : 1)); int4* sh_s_stage = sh_s + s_sh_stage * pipe; if constexpr (w_type_id != host::kFE2M1f.id()) { reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; } else { reinterpret_cast(&frag_s[k % 2])[0] = reinterpret_cast(sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; } } } return; } // Act-order case // Determine K of the "current" thread-block int cur_k = slice_k_start + tb_k * full_pipe; if (cur_k >= prob_k || cur_k >= slice_k_finish) { return; } // Reset (to current thread-block) since we read g_idx portion from the // shared memory cur_k = 0; // Progress to current iteration cur_k += k_iter_size * (k % b_sh_wr_iters); // Determine "position" inside the thread-block (based on warp and // thread-id) auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N int warp_row = warp_id / n_warps; int warp_col = warp_id % n_warps; cur_k += warp_row * 16; auto th_id = threadIdx.x % 32; cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix int s_col_shift = /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + (th_id / 4) * act_s_col_stride; if (is_same_group[pipe]) { if (k % 2 == 0) { *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + s_col_shift]; } else { *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); } for (int i = 1; i < 4; i++) { *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); } return; } int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); constexpr int k_frag_offsets[4] = {0, 1, 8, 9}; // Tensor core offsets per thread #pragma unroll for (int i = 0; i < 4; i++) { int actual_k = cur_k + k_frag_offsets[i]; int group_id = sh_g_idx_int_ptr[actual_k]; int rel_group_id = group_id - sh_first_group_id; *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = sh_s[rel_group_id * s_sh_stride + s_col_shift]; } }; auto fetch_zp_to_registers = [&](int k, int full_pipe) { // This code does not handle group_blocks == 0, // which signifies act_order. // has_zp implies AWQ, which doesn't have act_order, static_assert(!has_zp || group_blocks != 0); if constexpr (has_zp && !is_zp_float) { int pipe = full_pipe % stages; if constexpr (group_blocks == -1) { // load only when starting a new slice if (k == 0 && full_pipe == 0) { #pragma unroll for (int i = 0; i < num_ints_per_thread; i++) { frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp))[zp_sh_rd + i]; } } } else if constexpr (group_blocks >= thread_k_blocks) { if (k % b_sh_wr_iters == 0) { int4* sh_zp_stage = sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); #pragma unroll for (int i = 0; i < num_ints_per_thread; i++) { frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; } } } else { auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; int warp_row = warp_id / n_warps; int cur_k = warp_row * 16; cur_k += k_iter_size * (k % b_sh_wr_iters); int k_blocks = cur_k / 16; int cur_group_id = 0; // Suppress bogus and persistent divide-by-zero warning #pragma nv_diagnostic push #pragma nv_diag_suppress divide_by_zero cur_group_id = k_blocks / group_blocks; #pragma nv_diagnostic pop int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; sh_zp_stage += cur_group_id * zp_sh_stride; #pragma unroll for (int i = 0; i < num_ints_per_thread; i++) { frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; } } } else if constexpr (has_zp && is_zp_float) { int pipe = full_pipe % stages; if constexpr (group_blocks != -1) { if constexpr (group_blocks >= thread_k_blocks) { if (k % b_sh_wr_iters == 0) { int4* sh_zp_stage = sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); reinterpret_cast(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd]; } } else { auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; int warp_row = warp_id / n_warps; int cur_k = warp_row * 16; cur_k += k_iter_size * (k % b_sh_wr_iters); int k_blocks = cur_k / 16; // Suppress bogus and persistent divide-by-zero warning #pragma nv_diagnostic push #pragma nv_diag_suppress divide_by_zero int cur_group_id = k_blocks / group_blocks; #pragma nv_diagnostic pop int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; reinterpret_cast(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd + cur_group_id * zp_sh_stride]; } } } }; auto dequant_data = [&](int q, scalar_t2* frag_b_ptr) { dequant(q, frag_b_ptr); }; // Execute the actual tensor core matmul of a sub-tile. bool is_first_matmul_in_slice = true; auto matmul = [&](int k) { int k2 = k % 2; const bool is_new_zp = ((group_blocks != -1) && (group_blocks < thread_k_blocks || k == 0)) || (group_blocks == -1 && is_first_matmul_in_slice); if constexpr (has_zp && !is_zp_float) { if (is_new_zp) { if constexpr (group_blocks == -1) is_first_matmul_in_slice = false; FragB frag_zp_0; FragB frag_zp_1; int zp_quant_0, zp_quant_1; if constexpr (w_type.size_bits() == 4) { zp_quant_0 = frag_qzp[k2][0]; zp_quant_1 = zp_quant_0 >> 8; } else { static_assert(w_type.size_bits() == 8); zp_quant_0 = frag_qzp[k2][0]; zp_quant_1 = frag_qzp[k2][1]; } dequant_data(zp_quant_0, reinterpret_cast(&frag_zp)); dequant_data(zp_quant_1, reinterpret_cast(&frag_zp) + 2); } } if constexpr (!dequant_skip_flop && has_zp && is_zp_float) { if (is_new_zp) { reinterpret_cast(&frag_zp)[0] = reinterpret_cast(&frag_zpf[k2])[0]; } } if constexpr (w_type == host::kFE2M1f) { int s_quant_0 = reinterpret_cast(frag_s[k2])[0]; int s_quant_1 = reinterpret_cast(frag_s[k2])[1]; dequant_fp8_scales(s_quant_0, reinterpret_cast(&frag_s[k2])); dequant_fp8_scales(s_quant_1, reinterpret_cast(&frag_s[k2]) + 2); } // We have the m dimension as the inner loop in order to encourage overlapping // dequantization and matmul operations. #pragma unroll for (int j = 0; j < 4; j++) { FragB frag_b0; FragB frag_b1; int b_quant_0, b_quant_1; if constexpr (w_type_id == host::kFE2M1f.id()) { b_quant_1 = frag_b_quant[k2][0][j]; b_quant_0 = b_quant_1 << 8; } else if constexpr (w_type.size_bits() == 4) { b_quant_0 = frag_b_quant[k2][0][j]; b_quant_1 = b_quant_0 >> 8; } else { static_assert(w_type.size_bits() == 8); int* frag_b_quant_ptr = reinterpret_cast(frag_b_quant[k2]); b_quant_0 = frag_b_quant_ptr[j * 2 + 0]; b_quant_1 = frag_b_quant_ptr[j * 2 + 1]; } dequant_data(b_quant_0, reinterpret_cast(&frag_b0)); dequant_data(b_quant_1, reinterpret_cast(&frag_b1)); if constexpr (dequant_skip_flop && has_zp && !is_zp_float) { sub_zp(frag_b0, frag_zp[j], 0); sub_zp(frag_b1, frag_zp[j], 1); } // Apply scale to frag_b0 if constexpr (has_act_order) { static_assert(group_blocks != -1); scale4( frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j], act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0); scale4( frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j], act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1); } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float && group_blocks == -1) { int idx = (threadIdx.x / 4) % 2; scalar_t2 s2 = Dtype::nums2num2( reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 1])[idx]); if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2); scale_and_sub(frag_b0, s2.x, frag_zp[j].x); scale_and_sub(frag_b1, s2.y, frag_zp[j].y); } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1) { if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], *reinterpret_cast(&frag_s[k2][j])); scale_and_sub(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x); scale_and_sub(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y); } else if constexpr (group_blocks != -1) { scale(frag_b0, frag_s[k2][j], 0); scale(frag_b1, frag_s[k2][j], 1); } #pragma unroll for (int i = 0; i < thread_m_blocks; i++) { if constexpr (m_block_size_8) { mma_trans(frag_a[k2][i], frag_b0, frag_b1, frag_c[i][j][0]); } else { mma(frag_a[k2][i], frag_b0, frag_c[i][j][0]); mma(frag_a[k2][i], frag_b1, frag_c[i][j][1]); } } } }; // Since we slice across the k dimension of a tile in order to increase the // number of warps while keeping the n dimension of a tile reasonable, we have // multiple warps that accumulate their partial sums of the same output // location; which we have to reduce over in the end. We do in shared memory. auto thread_block_reduce = [&]() { constexpr int red_off = threads / b_sh_stride_threads / 2; if (red_off >= 1) { auto red_idx = threadIdx.x / b_sh_stride_threads; constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2; constexpr int red_sh_delta = b_sh_stride_threads; int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + (threadIdx.x % b_sh_stride_threads); // Parallel logarithmic shared memory reduction. We make sure to avoid any // unnecessary read or write iterations, e.g., for two warps we write only // once by warp 1 and read only once by warp 0. #pragma unroll for (int m_block = 0; m_block < thread_m_blocks; m_block++) { #pragma unroll for (int i = red_off; i > 0; i /= 2) { if (i <= red_idx && red_idx < 2 * i) { #pragma unroll for (int j = 0; j < 4 * 2; j += (m_block_size_8 ? 2 : 1)) { int red_sh_wr = red_sh_delta * j + (red_sh_rd - red_sh_stride * i); if (i < red_off) { float* c_rd = reinterpret_cast(&sh_red[red_sh_delta * j + red_sh_rd]); float* c_wr = reinterpret_cast(&sh_red[red_sh_wr]); #pragma unroll for (int k = 0; k < 4; k++) reinterpret_cast(frag_c)[4 * 2 * m_block + j][k] += c_rd[k] + c_wr[k]; } sh_red[red_sh_wr] = reinterpret_cast(&frag_c)[4 * 2 * m_block + j]; } } __syncthreads(); } if (red_idx == 0) { #pragma unroll for (int i = 0; i < 4 * 2; i += (m_block_size_8 ? 2 : 1)) { float* c_rd = reinterpret_cast(&sh_red[red_sh_delta * i + red_sh_rd]); #pragma unroll for (int j = 0; j < 4; j++) reinterpret_cast(frag_c)[4 * 2 * m_block + i][j] += c_rd[j]; } } __syncthreads(); } } }; // Since multiple threadblocks may process parts of the same column slice, we // finally have to globally reduce over the results. As the striped // partitioning minimizes the number of such reductions and our outputs are // usually rather small, we perform this reduction serially in L2 cache. auto global_reduce_fp16 = [&](bool first = false, bool last = false) { // We are very careful here to reduce directly in the output buffer to // maximize L2 cache utilization in this step. To do this, we write out // results in FP16 (but still reduce with FP32 compute). constexpr int active_threads = 32 * thread_n_blocks / 4; if (threadIdx.x < active_threads) { int c_gl_stride = prob_n / 8; int c_gl_wr_delta_o = 8 * c_gl_stride; int c_gl_wr_delta_i = 4 * (active_threads / 32); int c_gl_wr; if constexpr (m_block_size_8) { c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + 4 * (threadIdx.x / 32) + (threadIdx.x % 32) / 8; c_gl_wr += (2 * thread_n_blocks) * slice_col; } else { c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) + 4 * (threadIdx.x / 32) + threadIdx.x % 4; c_gl_wr += (2 * thread_n_blocks) * slice_col; } constexpr int c_sh_wr_delta = active_threads; auto c_sh_wr = threadIdx.x; int row = (threadIdx.x % 32) / 4; if (!first) { // Interestingly, doing direct global accesses here really seems to mess up // the compiler and lead to slowdowns, hence we also use async-copies even // though these fetches are not actually asynchronous. #pragma unroll for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { if constexpr (m_block_size_8) { cp_async4_pred( &sh_red[c_sh_wr + c_sh_wr_delta * i], &C[c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i], (threadIdx.x % 4) * 2 + i < prob_m); } else { cp_async4_pred( &sh_red[c_sh_wr + c_sh_wr_delta * i], &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)], i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); } } cp_async_fence(); cp_async_wait<0>(); } #pragma unroll for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { bool mask = (!m_block_size_8) && (i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m) || (m_block_size_8) && ((threadIdx.x % 4) * 2 + i < prob_m); if (mask) { if (!first) { int4 c_red = sh_red[c_sh_wr + i * c_sh_wr_delta]; #pragma unroll for (int j = 0; j < 2 * 4; j++) { int delta = 0; if constexpr (m_block_size_8) { delta = j % 2 == 1 ? -2 : 0; } reinterpret_cast(&frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta] += Dtype::num2float(reinterpret_cast(&c_red)[j]); } } if (!last) { int4 c; #pragma unroll for (int j = 0; j < 2 * 4; j++) { int delta = 0; if constexpr (m_block_size_8) { delta = j % 2 == 1 ? -2 : 0; } reinterpret_cast(&c)[j] = Dtype::float2num(reinterpret_cast(&frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta]); } if constexpr (m_block_size_8) C[c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i] = c; else C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] = c; } } } } }; // Globally reduce over threadblocks that compute the same column block. // We use a tmp C buffer to reduce in full fp32 precision. auto global_reduce_fp32 = [&](bool first = false, bool last = false) { constexpr int tb_m = thread_m_blocks * 16; constexpr int tb_n = thread_n_blocks * 16; constexpr int c_size = tb_m * tb_n * sizeof(float) / 16; constexpr int active_threads = 32 * thread_n_blocks / 4; bool is_th_active = threadIdx.x < active_threads; constexpr int num_floats = thread_m_blocks * 4 * 2 * 4; constexpr int th_size = num_floats * sizeof(float) / 16; int c_cur_offset = locks_off * c_size; if (!is_th_active) { return; } if (!first) { float* frag_c_ptr = reinterpret_cast(&frag_c); #pragma unroll for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) { sh_red[threadIdx.x] = C_tmp[c_cur_offset + active_threads * k + threadIdx.x]; float* sh_c_ptr = reinterpret_cast(&sh_red[threadIdx.x]); #pragma unroll for (int f = 0; f < 4; f++) { frag_c_ptr[k * 4 + f] += sh_c_ptr[f]; } } } if (!last) { int4* frag_c_ptr = reinterpret_cast(&frag_c); #pragma unroll for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) { C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k]; } } }; // Write out the reduce final result in the correct layout. We only actually // reshuffle matrix fragments in this step, the reduction above is performed // in fragment layout. auto write_result = [&]() { int c_gl_stride = prob_n / 8; constexpr int c_sh_stride = 2 * thread_n_blocks + 1; int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); constexpr int c_sh_rd_delta = c_sh_stride * (threads / (2 * thread_n_blocks)); int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); c_gl_wr += (2 * thread_n_blocks) * slice_col; int c_sh_wr; if constexpr (m_block_size_8) { c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) + (threadIdx.x % 32) / 4; c_sh_wr += 64 * (threadIdx.x / 32); } else { c_sh_wr = (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; c_sh_wr += 32 * (threadIdx.x / 32); } int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); int c_gl_wr_end = c_gl_stride * prob_m; // We first reorder in shared memory to guarantee the most efficient final // global write patterns auto write = [&](int idx, float c0, float c1, FragS& s) { scalar_t2 res = Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1)); // For per-column quantization we finally apply the scale here (only for // 4-bit) if constexpr ( !has_act_order && group_blocks == -1 && w_type.size_bits() == 4 && (has_zp && dequant_skip_flop || !has_zp)) { res = __hmul2(res, s[0]); } if constexpr (w_type == host::kFE2M1f) { res = __hmul2(res, global_scale); } if constexpr (m_block_size_8) { ((scalar_t*)sh_red)[idx] = res.x; ((scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y; } else { ((scalar_t2*)sh_red)[idx] = res; } }; if (threadIdx.x / 32 < thread_n_blocks / 4) { #pragma unroll for (int i = 0; i < thread_m_blocks; i++) { #pragma unroll for (int j = 0; j < 4; j++) { if constexpr (m_block_size_8) { int wr = c_sh_wr + 16 * j; write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]); write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 1]); } else { int wr = c_sh_wr + 8 * j; write( wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]); write( wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]); write( wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]); write( wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]); } } c_sh_wr += 16 * (4 * c_sh_stride); } } __syncthreads(); #pragma unroll for (int i = 0; i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); i++) { if (c_gl_wr < c_gl_wr_end) { if (use_atomic_add && slice_count > 1) { scalar_t2* C_half2 = reinterpret_cast(&C[c_gl_wr]); scalar_t2* sh_red_half2 = reinterpret_cast(&sh_red[c_sh_rd]); #pragma unroll for (int a = 0; a < 4; a++) { atomicAdd(&C_half2[a], sh_red_half2[a]); } } else { C[c_gl_wr] = sh_red[c_sh_rd]; } c_gl_wr += c_gl_wr_delta; c_sh_rd += c_sh_rd_delta; } } __syncthreads(); }; // Start global fetch and register load pipelines. auto start_pipes = [&]() { #pragma unroll for (int i = 0; i < stages - 1; i++) { if (has_act_order && i == 0) { int last_g_idx = slice_k_start + stages * tb_k * 2; if (last_g_idx >= prob_k) { last_g_idx = prob_k - 1; } fetch_act_order_scales_to_shared(true, g_idx[slice_k_start], g_idx[last_g_idx]); } if constexpr (has_zp && !is_zp_float && group_blocks == -1) { if (i == 0) { fetch_col_zp_to_shared(); if constexpr (!dequant_skip_flop) { fetch_col_scale_to_shared(); } } } fetch_to_shared(i, i, i < slice_iters); } zero_accums(); wait_for_stage(); init_same_group(0); fetch_to_registers(0, 0); fetch_scales_to_registers(0, 0); fetch_zp_to_registers(0, 0); a_gl_rd += a_gl_rd_delta_o * (stages - 1); if constexpr (has_act_order) { slice_k_start_shared_fetch += tb_k * (stages - 1); } }; if (slice_iters) { start_pipes(); } // Main loop. while (slice_iters) { // We unroll over both the global fetch and the register load pipeline to // ensure all shared memory accesses are static. Note that both pipelines // have even length meaning that the next iteration will always start at // index 0. #pragma unroll for (int pipe = 0; pipe < stages;) { #pragma unroll for (int k = 0; k < b_sh_wr_iters; k++) { fetch_to_registers(k + 1, pipe % stages); fetch_scales_to_registers(k + 1, pipe); fetch_zp_to_registers(k + 1, pipe); if (k == b_sh_wr_iters - 2) { fetch_to_shared((pipe + stages - 1) % stages, pipe, slice_iters >= stages); pipe++; wait_for_stage(); init_same_group(pipe % stages); } matmul(k); } slice_iters--; if (slice_iters == 0) { break; } } a_gl_rd += a_gl_rd_delta_o * stages; if constexpr (has_act_order) { slice_k_start += tb_k * stages; if (slice_k_start < prob_k) { slice_k_start_shared_fetch += tb_k * stages; int first_group_id = g_idx[slice_k_start]; int last_g_idx = slice_k_start + stages * tb_k * 2; if (last_g_idx >= prob_k) { last_g_idx = prob_k - 1; } int last_group_id = g_idx[last_g_idx]; if (last_group_id >= sh_first_group_id + sh_num_groups) { fetch_act_order_scales_to_shared(false, first_group_id, last_group_id); __syncthreads(); } } } // Process results and, if necessary, proceed to the next column slice. // While this pattern may not be the most readable, other ways of writing // the loop seemed to noticeably worse performance after compilation. if (slice_iters == 0) { cp_async_wait<0>(); bool last = slice_idx == slice_count - 1; // For per-column scales, we only fetch them here in the final step before // write-out if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) { if (w_type.size_bits() == 8 || (last || use_atomic_add)) { if (s_sh_wr_pred) { cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); } cp_async_fence(); } } thread_block_reduce(); if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) { if (w_type.size_bits() == 8 || (last || use_atomic_add)) { cp_async_wait<0>(); __syncthreads(); if (threadIdx.x / 32 < thread_n_blocks / 4) { reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; if constexpr (m_block_size_8) { int idx = (threadIdx.x / 4) % 2; scalar_t2* frag_s_half2 = reinterpret_cast(frag_s); #pragma unroll for (int i = 0; i < 8; i++) { frag_s_half2[i] = Dtype::num2num2(reinterpret_cast(&frag_s_half2[i])[idx]); } } } } } // For 8-bit channelwise, we apply the scale before the global reduction // that converts the fp32 results to fp16 (so that we avoid possible // overflow in fp16) if constexpr ( !has_act_order && group_blocks == -1 && w_type.size_bits() == 8 && (has_zp && dequant_skip_flop || !has_zp)) { if (threadIdx.x / 32 < thread_n_blocks / 4) { #pragma unroll for (int i = 0; i < thread_m_blocks; i++) { #pragma unroll for (int j = 0; j < 4; j++) { scale_float(reinterpret_cast(&frag_c[i][j][0][0]), frag_s[j / 2][2 * (j % 2) + 0]); scale_float( reinterpret_cast(&frag_c[i][j][0][2]), frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]); if constexpr (!m_block_size_8) { scale_float(reinterpret_cast(&frag_c[i][j][1][0]), frag_s[j / 2][2 * (j % 2) + 1]); scale_float(reinterpret_cast(&frag_c[i][j][1][2]), frag_s[j / 2][2 * (j % 2) + 1]); } } } } } if (slice_count > 1 && !use_atomic_add) { // only globally reduce if there is more than one block in a slice barrier_acquire(&locks[locks_off], slice_idx); if (use_fp32_reduce) { global_reduce_fp32(slice_idx == 0, last); } else { global_reduce_fp16(slice_idx == 0, last); } barrier_release(&locks[locks_off], last); } if (use_atomic_add && slice_count > 1 && slice_idx != 0) wait_negative_and_add(&locks[locks_off]); if (last || use_atomic_add) // only the last block in a slice actually writes the result write_result(); slice_row = 0; slice_col_par++; slice_col++; is_first_matmul_in_slice = true; init_slice(); if (slice_iters) { a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); #pragma unroll for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles; if (slice_col == 0) { #pragma unroll for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride; } // Update slice k/n for scales loading if constexpr (has_act_order) { slice_k_start = tb_k * slice_row; slice_k_finish = slice_k_start + tb_k * slice_iters; slice_k_start_shared_fetch = slice_k_start; slice_n_offset = act_s_col_tb_stride * slice_col; } else { s_gl_rd = s_sh_stride * slice_col + threadIdx.x; zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; } start_pipes(); } } } } } // namespace device::marlin #endif ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/marlin_moe/kernel.h ================================================ #include #include "../marlin/marlin.cuh" #include "../marlin/marlin_dtypes.cuh" #define MARLIN_KERNEL_PARAMS \ const int4 *__restrict__ A, const int4 *__restrict__ B, int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ const int4 *__restrict__ b_bias_ptr, const int4 *__restrict__ scales_ptr, \ const uint16_t *__restrict__ scale2_ptr, const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \ const int32_t *__restrict__ sorted_token_ids_ptr, const int32_t *__restrict__ expert_ids_ptr, \ const int32_t *__restrict__ num_tokens_past_padded_ptr, const float *__restrict__ topk_weights_ptr, int top_k, \ bool mul_topk_weights, bool is_ep, int num_groups, int prob_m, int prob_n, int prob_k, int *locks, \ bool has_bias, bool use_atomic_add, bool use_fp32_reduce, int max_shared_mem namespace device::marlin_moe { template < typename scalar_t, // compute dtype, half or nv_float16 const host::ScalarTypeId w_type_id, // weight ScalarType id const host::ScalarTypeId s_type_id, // weight scale ScalarType id const int threads, // number of threads in a threadblock const int thread_m_blocks, // number of 16x16 blocks in the m // dimension (batchsize) of the // threadblock const int thread_n_blocks, // same for n dimension (output) const int thread_k_blocks, // same for k dimension (reduction) const bool m_block_size_8, // whether m_block_size == 8 // only works when thread_m_blocks == 1 const int stages, // number of stages for the async global->shared // fetch pipeline const int group_blocks, // number of consecutive 16x16 blocks // with a separate quantization scale const bool is_zp_float // is zero point of float16 type? > __global__ void Marlin(MARLIN_KERNEL_PARAMS); } // namespace device::marlin_moe ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/marlin_moe/marlin_template.h ================================================ /* * Modified by Neural Magic * Copyright (C) Marlin.2024 Elias Frantar * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ /* * Adapted from https://github.com/IST-DASLab/marlin */ #include #include "../marlin/dequant.h" #include "../marlin/marlin.cuh" #include "../marlin/marlin_dtypes.cuh" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert( \ std::is_same::value || std::is_same::value, \ "only float16 and bfloat16 is supported"); namespace device::marlin_moe { using namespace device::marlin; #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 template < typename scalar_t, // compute dtype, half or nv_float16 const host::ScalarTypeId w_type_id, // weight ScalarType id const int threads, // number of threads in a threadblock const int thread_m_blocks, // number of 16x16 blocks in the m // dimension (batchsize) of the // threadblock const int thread_n_blocks, // same for n dimension (output) const int thread_k_blocks, // same for k dimension (reduction) const bool m_block_size_8, // whether m_block_size == 8 // only works when thread_m_blocks == 1 const int stages, // number of stages for the async global->shared // fetch pipeline const int group_blocks, // number of consecutive 16x16 blocks // with a separate quantization scale const bool is_zp_float // is zero point of float16 type? > __global__ void Marlin( const int4* __restrict__ A, // fp16 input matrix of shape mxk const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn int4* __restrict__ C, // fp16 output buffer of shape mxn int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape // (k/groupsize)xn const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape // (k/groupsize)x(n/pack_factor) const int* __restrict__ g_idx, // int32 group indices of shape k const int32_t* __restrict__ sorted_token_ids_ptr, // moe sorted_ids const int32_t* __restrict__ expert_ids_ptr, // moe expert ids const int32_t* __restrict__ num_tokens_past_padded_ptr, // moe num tokens const float* __restrict__ topk_weights_ptr, // moe top weights int top_k, // num of experts per token bool mul_topk_weights, // mul topk weights or not bool is_ep, // expert parallelism int num_groups, // number of scale groups per output channel int prob_m, // batch dimension m int prob_n, // output dimension n int prob_k, // reduction dimension k int* locks, // extra global storage for barrier synchronization bool use_atomic_add, // whether to use atomic add to reduce bool use_fp32_reduce, // whether to use fp32 global reduce int max_shared_mem) {} } // namespace device::marlin_moe #else // m16n8k16 tensor core mma instruction with fp16 inputs and fp32 // output/accumulation. template __device__ inline void mma(const typename ScalarType::FragA& a_frag, const typename ScalarType::FragB& frag_b, typename ScalarType::FragC& frag_c) { const uint32_t* a = reinterpret_cast(&a_frag); const uint32_t* b = reinterpret_cast(&frag_b); float* c = reinterpret_cast(&frag_c); if constexpr (std::is_same::value) { asm volatile( "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); } else if constexpr (std::is_same::value) { asm volatile( "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); } else { STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); } } template __device__ inline void mma_trans( const typename ScalarType::FragA& a_frag, const typename ScalarType::FragB& frag_b, const typename ScalarType::FragB& frag_b2, typename ScalarType::FragC& frag_c) { const uint32_t* a = reinterpret_cast(&a_frag); const uint32_t* b = reinterpret_cast(&frag_b); const uint32_t* b2 = reinterpret_cast(&frag_b2); float* c = reinterpret_cast(&frag_c); if constexpr (std::is_same::value) { asm volatile( "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); } else if constexpr (std::is_same::value) { asm volatile( "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); } else { STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); } } // Instruction for loading a full 16x16 matrix fragment of operand A from shared // memory, directly in tensor core layout. template __device__ inline void ldsm(typename ScalarType::FragA& frag_a, const void* smem_ptr) { uint32_t* a = reinterpret_cast(&frag_a); uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); if constexpr (count == 4) { asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) : "r"(smem)); } else if constexpr (count == 2) { asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" : "=r"(a[0]), "=r"(a[1]) : "r"(smem)); } else if constexpr (count == 1) { asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" : "=r"(a[0]) : "r"(smem)); } else { static_assert(count == 1 || count == 2 || count == 4, "invalid count"); } } // Multiply dequantized values by the corresponding quantization scale; used // only for grouped quantization. template __device__ inline void scale(typename ScalarType::FragB& frag_b, typename ScalarType::FragS& frag_s, int i) { using scalar_t2 = typename ScalarType::scalar_t2; scalar_t2 s = ScalarType::num2num2(reinterpret_cast(&frag_s)[i]); frag_b[0] = __hmul2(frag_b[0], s); frag_b[1] = __hmul2(frag_b[1], s); } template __device__ inline void scale_and_sub(typename ScalarType::FragB& frag_b, scalar_t s, scalar_t zp) { using scalar_t2 = typename ScalarType::scalar_t2; scalar_t2 s2 = ScalarType::num2num2(s); scalar_t2 zp2 = ScalarType::num2num2(zp); frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2)); frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2)); } template __device__ inline void sub_zp(typename ScalarType::FragB& frag_b, typename ScalarType::scalar_t2& frag_zp, int i) { using scalar_t2 = typename ScalarType::scalar_t2; scalar_t2 zp = ScalarType::num2num2(reinterpret_cast(&frag_zp)[i]); frag_b[0] = __hsub2(frag_b[0], zp); frag_b[1] = __hsub2(frag_b[1], zp); } // Same as above, but for act_order (each K is multiplied individually) template __device__ inline void scale4( typename ScalarType::FragB& frag_b, typename ScalarType::FragS& frag_s_1, typename ScalarType::FragS& frag_s_2, typename ScalarType::FragS& frag_s_3, typename ScalarType::FragS& frag_s_4, int i) { using scalar_t2 = typename ScalarType::scalar_t2; scalar_t2 s_val_1_2; s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; scalar_t2 s_val_3_4; s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; frag_b[0] = __hmul2(frag_b[0], s_val_1_2); frag_b[1] = __hmul2(frag_b[1], s_val_3_4); } // Given 2 floats multiply by 2 scales (halves) template __device__ inline void scale_float(float* c, typename ScalarType::FragS& s) { scalar_t* s_ptr = reinterpret_cast(&s); c[0] = __fmul_rn(c[0], ScalarType::num2float(s_ptr[0])); c[1] = __fmul_rn(c[1], ScalarType::num2float(s_ptr[1])); } // Wait until barrier reaches `count`, then lock for current threadblock. __device__ inline void barrier_acquire(int* lock, int count) { if (threadIdx.x == 0) { int state = -1; do // Guarantee that subsequent writes by this threadblock will be visible // globally. asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); while (state != count); } __syncthreads(); } // Release barrier and increment visitation count. __device__ inline void barrier_release(int* lock, bool reset = false) { __syncthreads(); if (threadIdx.x == 0) { if (reset) { lock[0] = 0; return; } int val = 1; // Make sure that all writes since acquiring this barrier are visible // globally, while releasing the barrier. asm volatile("fence.acq_rel.gpu;\n"); asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" : : "l"(lock), "r"(val)); } } // Wait until value of lock to be negative, and then add 1 __device__ inline void wait_negative_and_add(int* lock) { if (threadIdx.x == 0) { int state = 0; do // Guarantee that subsequent writes by this threadblock will be visible // globally. asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); while (state >= 0); atomicAdd(lock, 1); } __syncthreads(); } template < typename scalar_t, // compute dtype, half or nv_float16 const host::ScalarTypeId w_type_id, // weight ScalarType id const host::ScalarTypeId s_type_id, // weight scale ScalarType id const int threads, // number of threads in a threadblock const int thread_m_blocks, // number of 16x16 blocks in the m // dimension (batchsize) of the // threadblock const int thread_n_blocks, // same for n dimension (output) const int thread_k_blocks, // same for k dimension (reduction) const bool m_block_size_8, // whether m_block_size == 8 // only works when thread_m_blocks == 1 const int stages, // number of stages for the async global->shared // fetch pipeline const int group_blocks, // number of consecutive 16x16 blocks // with a separate quantization scale const bool is_zp_float // is zero point of float16 type? > __global__ void Marlin( const int4* __restrict__ A, // fp16 input matrix of shape mxk const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn int4* __restrict__ C, // fp16 output buffer of shape mxn int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) const int4* __restrict__ b_bias_ptr, const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape // (k/groupsize)xn const uint16_t* __restrict__ scale2_ptr, // fp16 global scale (for nvfp4 // only) const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape // (k/groupsize)x(n/pack_factor) const int* __restrict__ g_idx, // int32 group indices of shape k const int32_t* __restrict__ sorted_token_ids_ptr, // moe sorted_ids const int32_t* __restrict__ expert_ids_ptr, // moe expert ids const int32_t* __restrict__ num_tokens_past_padded_ptr, // moe num tokens const float* __restrict__ topk_weights_ptr, // moe top weights int top_k, // num of experts per token bool mul_topk_weights, // mul topk weights or not bool is_ep, // expert parallelism int num_groups, // number of scale groups per output channel int prob_m, // batch dimension m int prob_n, // output dimension n int prob_k, // reduction dimension k int* locks, // extra global storage for barrier synchronization bool has_bias, bool use_atomic_add, // whether to use atomic add to reduce bool use_fp32_reduce, // whether to use fp32 global reduce int max_shared_mem) { // Each threadblock processes one "stripe" of the B matrix with (roughly) the // same size, which might involve multiple column "slices" (of width 16 * // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM // example: // 0 1 3 // 0 2 3 // 1 2 4 // While this kind of partitioning makes things somewhat more complicated, it // ensures good utilization of all SMs for many kinds of shape and GPU // configurations, while requiring as few slow global cross-threadblock // reductions as possible. using Dtype = ScalarType; using scalar_t2 = typename ScalarType::scalar_t2; using FragA = typename ScalarType::FragA; using FragB = typename ScalarType::FragB; using FragC = typename ScalarType::FragC; using FragS = typename ScalarType::FragS; using FragZP = typename ScalarType::FragZP; extern __shared__ int4 sh[]; static constexpr auto w_type = host::ScalarType::from_id(w_type_id); static constexpr auto s_type = host::ScalarType::from_id(s_type_id); if constexpr (w_type == host::kFE2M1f) { static_assert(s_type == host::kFE4M3fn && group_blocks == 1 || s_type == host::kFE8M0fnu && group_blocks == 2); } else if constexpr (std::is_same::value) { static_assert(s_type == host::kBFloat16); } else if constexpr (std::is_same::value) { static_assert(s_type == host::kFloat16); } constexpr bool has_zp = w_type == host::kU4 || w_type == host::kU8; constexpr bool is_int_type = w_type == host::kU4 || w_type == host::kU8 || w_type == host::kU4B8 || w_type == host::kU8B128; // see comments of dequant.h for more details constexpr bool dequant_skip_flop = w_type == host::kFE4M3fn || w_type == host::kFE2M1f && s_type == host::kFE4M3fn || has_zp && !is_zp_float && !std::is_same::value || has_zp && !is_zp_float && !(w_type == host::kU8); scalar_t2 global_scale; constexpr bool has_act_order = group_blocks == 0; constexpr int pack_factor = 32 / w_type.size_bits(); static_assert(thread_m_blocks == 1 || !m_block_size_8); constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); const int group_size = (!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups; const int scales_expert_stride = prob_n * prob_k / group_size / (w_type == host::kFE2M1f ? 16 : 8); const int zp_expert_stride = is_zp_float ? prob_n * prob_k / group_size / 8 : prob_n * prob_k / group_size / (pack_factor * 4); const int b_bias_expert_stride = prob_n / 8; // parallel: num valid moe blocks int num_tokens_past_padded = num_tokens_past_padded_ptr[0]; int parallel = num_tokens_past_padded / moe_block_size; int num_valid_blocks = parallel; if (is_ep) { for (int i = 0; i < parallel; i++) { if (expert_ids_ptr[i] == -1) num_valid_blocks--; } } int num_invalid_blocks = parallel - num_valid_blocks; parallel = num_valid_blocks; int k_tiles = prob_k / 16 / thread_k_blocks; int n_tiles = prob_n / 16 / thread_n_blocks; int iters = div_ceil(k_tiles * n_tiles * parallel, gridDim.x); if constexpr (!has_act_order && group_blocks != -1) { if (group_blocks >= thread_k_blocks) { // Ensure that the number of tiles in each stripe is a multiple of the // groupsize; this avoids an annoying special case where a stripe starts // in the middle of group. iters = (group_blocks / thread_k_blocks) * div_ceil(iters, (group_blocks / thread_k_blocks)); } } int slice_row = (iters * blockIdx.x) % k_tiles; int slice_col_par = (iters * blockIdx.x) / k_tiles; int slice_col = slice_col_par; int slice_iters; // number of threadblock tiles in the current slice int slice_count = 0; // total number of active threadblocks in the current slice int slice_idx; // index of threadblock in current slice; numbered bottom to // top int par_id = 0; int block_id = -1; int64_t expert_id = 0; // use int64 to avoid computation result overflow int old_expert_id = 0; int64_t B_expert_off = 0; int4* sh_block_sorted_ids_int4 = sh; int4* sh_rd_block_sorted_ids_int4 = sh_block_sorted_ids_int4 + moe_block_size / 4; int4* sh_block_topk_weights_int4 = sh_rd_block_sorted_ids_int4 + moe_block_size / 4; // sh_block_topk_weights_int4 only need (moe_block_size / 4); // but we pad to align to 256 bytes int4* sh_new = sh_block_topk_weights_int4 + moe_block_size / 2 + moe_block_size; int32_t* sh_block_sorted_ids = reinterpret_cast(sh_block_sorted_ids_int4); int32_t* sh_rd_block_sorted_ids = reinterpret_cast(sh_rd_block_sorted_ids_int4); scalar_t2* sh_block_topk_weights = reinterpret_cast(sh_block_topk_weights_int4); int32_t block_num_valid_tokens = 0; int32_t locks_off = 0; // We can easily implement parallel problem execution by just remapping // indices and advancing global pointers if (slice_col_par >= n_tiles) { slice_col = slice_col_par % n_tiles; par_id = slice_col_par / n_tiles; } if (parallel * n_tiles >= gridDim.x) { // when parallel * n_tiles >= sms // then there are at most $sms$ conflict tile blocks locks_off = blockIdx.x; } else { locks_off = (iters * blockIdx.x) / k_tiles - 1; } // read moe block data given block_id // block_sorted_ids / block_num_valid_tokens / block_topk_weights auto read_moe_block_data = [&](int block_id) { block_num_valid_tokens = moe_block_size; #pragma unroll for (int i = 0; i < moe_block_size / 4; i++) { int4 sorted_token_ids_int4 = reinterpret_cast(sorted_token_ids_ptr)[block_id * moe_block_size / 4 + i]; int* sorted_token_ids = reinterpret_cast(&sorted_token_ids_int4); #pragma unroll for (int j = 0; j < 4; j++) { if (sorted_token_ids[j] >= prob_m * top_k) { block_num_valid_tokens = i * 4 + j; break; } } if (block_num_valid_tokens != moe_block_size) break; } __syncthreads(); int tid4 = threadIdx.x / 4; if (threadIdx.x % 4 == 0 && threadIdx.x < block_num_valid_tokens) { sh_block_sorted_ids_int4[tid4] = reinterpret_cast(sorted_token_ids_ptr)[block_id * moe_block_size / 4 + tid4]; #pragma unroll for (int i = 0; i < 4; i++) sh_rd_block_sorted_ids[tid4 * 4 + i] = sh_block_sorted_ids[tid4 * 4 + i] / top_k; if (mul_topk_weights) { #pragma unroll for (int i = 0; i < 4; i++) { int idx = tid4 * 4 + i; // idx = idx < block_num_valid_tokens ? idx : 0; if (idx < block_num_valid_tokens) { if constexpr (w_type == host::kFE2M1f && s_type == host::kFE4M3fn) { sh_block_topk_weights[idx] = __hmul2(global_scale, Dtype::num2num2(Dtype::float2num(topk_weights_ptr[sh_block_sorted_ids[idx]]))); } else { sh_block_topk_weights[idx] = Dtype::num2num2(Dtype::float2num(topk_weights_ptr[sh_block_sorted_ids[idx]])); } } } } } __syncthreads(); }; // when move to next moe block, find the next block_id and expert_id // and then read moe block data auto update_next_moe_block_data = [&]() { if (par_id >= parallel) return; old_expert_id = expert_id; if (num_invalid_blocks > 0) { int skip_count = block_id == -1 ? par_id : 0; block_id++; for (int i = block_id; i < num_tokens_past_padded / moe_block_size; i++) { expert_id = expert_ids_ptr[i]; if (expert_id != -1) { if (skip_count == 0) { block_id = i; break; }; skip_count--; }; } } else { block_id = par_id; expert_id = expert_ids_ptr[block_id]; } if constexpr (w_type == host::kFE2M1f && s_type == host::kFE4M3fn) { uint16_t val = scale2_ptr[expert_id]; global_scale = Dtype::num2num2(*reinterpret_cast(&val)); } B_expert_off = expert_id * prob_n * prob_k / (pack_factor * 4); scales_ptr += (expert_id - old_expert_id) * scales_expert_stride; if constexpr (has_zp) { zp_ptr += (expert_id - old_expert_id) * zp_expert_stride; } if constexpr (has_act_order) { g_idx += (expert_id - old_expert_id) * prob_k; } if (has_bias) { b_bias_ptr += (expert_id - old_expert_id) * b_bias_expert_stride; } read_moe_block_data(block_id); }; // Compute all information about the current slice which is required for // synchronization. auto init_slice = [&](bool first_init = false) { slice_iters = iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0; if (slice_iters == 0) return; if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row; slice_count = 1; slice_idx = 0; int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); if (col_first <= k_tiles * (slice_col_par + 1)) { int col_off = col_first - k_tiles * slice_col_par; slice_count = div_ceil(k_tiles - col_off, iters); if (col_off > 0) slice_count++; int delta_first = iters * blockIdx.x - col_first; if (delta_first < 0 || (col_off == 0 && delta_first == 0)) slice_idx = slice_count - 1; else { slice_idx = slice_count - 1 - delta_first / iters; if (col_off > 0) slice_idx--; } } if (parallel * n_tiles >= gridDim.x) { if (slice_count > 1 && slice_idx == slice_count - 1) { locks_off++; } } else { locks_off++; } if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) { constexpr int threads_per_m = 16 * thread_n_blocks / 8; int m_per_thread = div_ceil(block_num_valid_tokens, threads / threads_per_m); for (int i = 0; i < m_per_thread; i++) { int row = threads / threads_per_m * i + threadIdx.x / threads_per_m; if (row < block_num_valid_tokens) { int64_t sorted_row = sh_block_sorted_ids[row]; int col = slice_col * 16 * thread_n_blocks / 8 + threadIdx.x % threads_per_m; C[sorted_row * prob_n / 8 + col] = {0, 0, 0, 0}; } } // After write zero to output, write a negative value to lock. // Every SM that processes the same slice would wait for // the negative value, and then atomicAdd 1 to it. // After all SMs are processed, the lock value would back to 0 again. __syncthreads(); if (threadIdx.x == 0) locks[locks_off] = 1 - slice_count; } if (slice_col == n_tiles) { slice_col = 0; par_id++; update_next_moe_block_data(); } }; update_next_moe_block_data(); init_slice(true); // A sizes/strides // stride of the A matrix in global memory int a_gl_stride = prob_k / 8; // stride of an A matrix tile in shared memory constexpr int a_sh_stride = 16 * thread_k_blocks / 8; // delta between subsequent A tiles in global memory constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / 8; // between subsequent accesses within a tile int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); // between shared memory writes constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); // between shared memory tile reads constexpr int a_sh_rd_delta_o = 2 * ((threads / 32) / (thread_n_blocks / 4)); // within a shared memory tile constexpr int a_sh_rd_delta_i = a_sh_stride * 16; // overall size of a tile constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks); // number of shared write iterations for a tile constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); // B sizes/strides int b_gl_stride = 16 * prob_n / (pack_factor * 4); constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4; constexpr int b_thread_vecs = w_type.size_bits() == 4 ? 1 : 2; constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks; int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads); constexpr int b_sh_wr_delta = threads * b_thread_vecs; constexpr int b_sh_rd_delta = threads * b_thread_vecs; constexpr int b_sh_stage = b_sh_stride * thread_k_blocks; constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; // Scale sizes/strides without act_order int s_gl_stride = prob_n / 8; constexpr int s_sh_stride = 16 * thread_n_blocks / 8; constexpr int s_tb_groups = !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks ? thread_k_blocks / group_blocks / (w_type == host::kFE2M1f ? 2 : 1) : 1; constexpr int s_sh_stage = s_tb_groups * s_sh_stride; int s_gl_rd_delta = s_gl_stride; // Scale size/strides with act_order constexpr int tb_k = 16 * thread_k_blocks; constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; // constexpr int act_s_row_stride = 1; // int act_s_col_stride = act_s_row_stride * num_groups; constexpr int act_s_max_num_groups = 32; int act_s_col_stride = 1; int act_s_col_warp_stride = act_s_col_stride * 8; int tb_n_warps = thread_n_blocks / 4; int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; // Zero-points sizes/strides int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4; constexpr int zp_sh_stride = is_zp_float ? 16 * thread_n_blocks / 8 : ((16 * thread_n_blocks) / pack_factor) / 4; constexpr int zp_tb_groups = s_tb_groups; constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0; int zp_gl_rd_delta = zp_gl_stride; // Global A read index of current thread. int a_gl_rd_row = threadIdx.x / a_gl_rd_delta_o; int a_gl_rd_col = a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o; // Shared write index of current thread. int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); // Shared read index. int a_sh_rd = a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) + (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1)); a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4)); int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) + (threadIdx.x % b_sh_stride_threads) * b_thread_vecs; b_gl_rd += b_sh_stride * slice_col; b_gl_rd += b_gl_rd_delta_o * slice_row; auto b_sh_wr = threadIdx.x * b_thread_vecs; auto b_sh_rd = threadIdx.x * b_thread_vecs; // For act_order constexpr int k_iter_size = tb_k / b_sh_wr_iters; int slice_k_start = tb_k * slice_row; int slice_k_finish = slice_k_start + tb_k * slice_iters; int slice_k_start_shared_fetch = slice_k_start; int slice_n_offset = act_s_col_tb_stride * slice_col; // No act_order int s_gl_rd; if constexpr (!has_act_order) { if constexpr (group_blocks == -1) { s_gl_rd = s_sh_stride * slice_col + threadIdx.x; } else { s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) / (w_type == host::kFE2M1f ? 2 : 1) + s_sh_stride * slice_col + threadIdx.x; } } auto s_sh_wr = threadIdx.x; bool s_sh_wr_pred = threadIdx.x < s_sh_stride; // Zero-points int zp_gl_rd; if constexpr (has_zp) { if constexpr (group_blocks == -1) { zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; } else { zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + zp_sh_stride * slice_col + threadIdx.x; } } auto zp_sh_wr = threadIdx.x; bool zp_sh_wr_pred = threadIdx.x < zp_sh_stride; // We use a different scale layout for grouped and column-wise quantization as // we scale a `half2` tile in column-major layout in the former and in // row-major in the latter case. int s_sh_rd; if constexpr (group_blocks != -1 && w_type == host::kFE2M1f) { auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; int warp_row = warp_id / n_warps; s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 4; s_sh_rd = s_sh_rd * 2 + (warp_row / group_blocks) % 2; } else if constexpr (group_blocks != -1) s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 4; else if constexpr (group_blocks == -1 && (m_block_size_8 || (has_zp && !dequant_skip_flop))) s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 8; else s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) % 4; int bias_sh_rd; if constexpr (m_block_size_8) { bias_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 8; } else { bias_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) % 4; } int bias_sh_wr = threadIdx.x; int bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; // Zero-points have the same read layout as the scales // (without column-wise case) constexpr int num_col_threads = 8; constexpr int num_row_threads = 4; constexpr int num_ints_per_thread = 8 / pack_factor; int zp_sh_rd; if constexpr (has_zp) { if constexpr (is_zp_float) { if constexpr (group_blocks != -1) { zp_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 4; } } else { zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); } } // To ensure that writing and reading A tiles to/from shared memory, the // latter in fragment format, is fully bank conflict free, we need to use a // rather fancy XOR-based layout. The key here is that neither reads nor // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the // same shared memory banks. Further, it seems (based on NSight-Compute) that // each warp must also write a consecutive memory segment? auto transform_a = [&](int i) { int row = i / a_gl_rd_delta_o; return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8); }; // Since the computation of this remapping is non-trivial and, due to our main // loop unrolls, all shared memory accesses are static, we simply precompute // both transformed reads and writes. int a_sh_wr_trans[a_sh_wr_iters]; #pragma unroll for (int i = 0; i < a_sh_wr_iters; i++) a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; #pragma unroll for (int i = 0; i < b_sh_wr_iters; i++) { #pragma unroll for (int j = 0; j < thread_m_blocks; j++) a_sh_rd_trans[i][j] = transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd); } // Since B-accesses have non-constant stride they have to be computed at // runtime; we break dependencies between subsequent accesses with a tile by // maintining multiple pointers (we have enough registers), a tiny // optimization. const int4* B_ptr[b_sh_wr_iters]; #pragma unroll for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd; // Shared memory storage for global fetch pipelines. constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks; constexpr int sh_b_size = stages * b_sh_stage; int4* sh_b = sh_new; int4* sh_red = sh_new; constexpr int sh_size_b_red_min = (sh_red_size < sh_b_size ? sh_red_size : sh_b_size); constexpr int sh_size_b_red_max = (sh_red_size > sh_b_size ? sh_red_size : sh_b_size); constexpr int sh_bias_size = (thread_n_blocks * 16 / 8); constexpr int sh_b_red_bias_size = sh_size_b_red_max > (sh_size_b_red_min + sh_bias_size) ? sh_size_b_red_max : (sh_size_b_red_min + sh_bias_size); int4* sh_bias = sh_new + sh_size_b_red_min; int4* sh_g_idx = sh_new + sh_b_red_bias_size; int4* sh_zp = sh_g_idx + (stages * g_idx_stage); constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride) : (stages * s_sh_stage); int4* sh_s = sh_zp + (stages * zp_sh_stage); // shared memory reused by reduction should be smaller than // shared memory used by weight. static_assert(thread_m_blocks * 16 * thread_n_blocks * 16 / 8 <= stages * b_sh_stage); int4* sh_a = sh_s + sh_s_size; constexpr int shm_size_used = moe_block_size + stages * (g_idx_stage + zp_sh_stage) + sh_s_size + sh_b_red_bias_size; // all remaining shared memory is used to cache A (input) // sh_a_max_row is at least ` stages * 16 * thread_m_blocks ` int sh_a_max_row = ((max_shared_mem - 1024) / 16 - shm_size_used) / (thread_k_blocks * 2); // Register storage for double buffer of shared memory reads. FragA frag_a[2][thread_m_blocks]; I4 frag_b_quant[2][b_thread_vecs]; FragC frag_c[thread_m_blocks][4][2]; FragS frag_s[2][4]; // No act-order FragS frag_bias[2][4]; FragS act_frag_s[2][4][4]; // For act-order int frag_qzp[2][num_ints_per_thread]; // Zero-points FragZP frag_zp; // Zero-points in fp16 FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ // Zero accumulators. auto zero_accums = [&]() { #pragma unroll for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) reinterpret_cast(frag_c)[i] = 0; }; int sh_first_group_id = -1; int sh_num_groups = -1; auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id, int last_group_id) { sh_first_group_id = first_group_id; sh_num_groups = last_group_id - first_group_id + 1; if (sh_num_groups > act_s_max_num_groups) { sh_num_groups = act_s_max_num_groups; } if (sh_first_group_id + sh_num_groups > num_groups) { sh_num_groups = num_groups - sh_first_group_id; } int row_offset = first_group_id * s_gl_stride; if (is_async) { for (int i = 0; i < sh_num_groups; i++) { if (threadIdx.x < s_sh_stride) { cp_async4_pred( &sh_s[(i * s_sh_stride) + threadIdx.x], &scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]); } } } else { for (int i = 0; i < sh_num_groups; i++) { if (threadIdx.x < s_sh_stride) { sh_s[(i * s_sh_stride) + threadIdx.x] = scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]; } } } }; // Asynchronously fetch the next A, B and s tile from global to the next // shared memory pipeline location. bool should_load_a = true; int max_num_stage_groups = ((sh_a_max_row - moe_block_size) / moe_block_size + 1) / stages; max_num_stage_groups = max(max_num_stage_groups, 1); auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true, int pipe_a = 0) { if (pred) { if (should_load_a) { int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe_a; #pragma unroll for (int i = 0; i < a_sh_wr_iters; i++) { int row = a_gl_rd_delta_i / a_gl_stride * i + a_gl_rd_row; int64_t sorted_row = 0; if (!m_block_size_8 || row < 8) sorted_row = sh_rd_block_sorted_ids[row]; int64_t true_idx = sorted_row * a_gl_stride + a_gl_rd_col + a_gl_rd_delta_o * a_off; cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], &A[true_idx], row < block_num_valid_tokens); } } int4* sh_b_stage = sh_b + b_sh_stage * pipe; #pragma unroll for (int i = 0; i < b_sh_wr_iters; i++) { #pragma unroll for (int j = 0; j < b_thread_vecs; j++) { cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j], B_ptr[i] + j + B_expert_off); } B_ptr[i] += b_gl_rd_delta_o; } if constexpr (has_act_order) { // Fetch g_idx thread-block portion int full_pipe = a_off; int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; if (cur_k < prob_k && cur_k < slice_k_finish) { int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; int4 const* cur_g_idx_stage_ptr = reinterpret_cast(&g_idx[cur_k]); if (threadIdx.x < g_idx_stage) { cp_async4_pred(&sh_g_idx_stage[threadIdx.x], &cur_g_idx_stage_ptr[threadIdx.x]); } } } else { if constexpr (group_blocks != -1) { int4* sh_s_stage = sh_s + s_sh_stage * pipe; if constexpr (group_blocks >= thread_k_blocks) { // Only fetch scales if this tile starts a new group if (pipe % (group_blocks / thread_k_blocks) == 0) { if (s_sh_wr_pred) { cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); } s_gl_rd += s_gl_rd_delta; } } else { for (int i = 0; i < s_tb_groups; i++) { if (s_sh_wr_pred) { cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr], &scales_ptr[s_gl_rd]); } s_gl_rd += s_gl_rd_delta; } } } if constexpr (has_zp && group_blocks != -1) { int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; if constexpr (group_blocks >= thread_k_blocks) { // Only fetch zero-points if this tile starts a new group if (pipe % (group_blocks / thread_k_blocks) == 0) { if (zp_sh_wr_pred) { cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); } zp_gl_rd += zp_gl_rd_delta; } } else { for (int i = 0; i < zp_tb_groups; i++) { if (zp_sh_wr_pred) { cp_async4(&sh_zp_stage[i * zp_sh_stride + zp_sh_wr], &zp_ptr[zp_gl_rd]); } zp_gl_rd += zp_gl_rd_delta; } } } } } // Insert a fence even when we are winding down the pipeline to ensure that // waiting is also correct at this point. cp_async_fence(); }; auto fetch_col_zp_to_shared = [&]() { if (zp_sh_wr_pred) { cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]); } }; auto fetch_col_scale_to_shared = [&]() { if (s_sh_wr_pred) { cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); } }; // Wait until the next thread tile has been loaded to shared memory. auto wait_for_stage = [&]() { // We only have `stages - 2` active fetches since we are double buffering // and can only issue the next fetch when it is guaranteed that the previous // shared memory load is fully complete (as it may otherwise be // overwritten). cp_async_wait(); __syncthreads(); }; // Load the next sub-tile from the current location in the shared memory pipe // into the current register buffer. auto fetch_to_registers = [&](int k, int pipe, int pipe_a = 0) { int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe_a; #pragma unroll for (int i = 0; i < thread_m_blocks; i++) ldsm(frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); int4* sh_b_stage = sh_b + b_sh_stage * pipe; #pragma unroll for (int i = 0; i < b_thread_vecs; i++) { frag_b_quant[k % 2][i] = *reinterpret_cast(&sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]); } }; bool is_same_group[stages]; int same_group_id[stages]; auto init_same_group = [&](int pipe) { if constexpr (!has_act_order) { return; } int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); int group_id_1 = sh_g_idx_int_ptr[0]; int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; is_same_group[pipe] = group_id_1 == group_id_2; same_group_id[pipe] = group_id_1; }; auto fetch_scales_to_registers = [&](int k, int full_pipe) { int pipe = full_pipe % stages; if constexpr (!has_act_order) { // No act-order case if constexpr (group_blocks == -1) { // load only when starting a new slice if (k == 0 && full_pipe == 0) { reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; } } else if constexpr (group_blocks != -1) { if constexpr (group_blocks >= thread_k_blocks) { if (k % b_sh_wr_iters == 0) { int4* sh_s_stage = sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; } else { reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; } } else { auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; int warp_row = warp_id / n_warps; int cur_k = warp_row * 16; cur_k += k_iter_size * (k % b_sh_wr_iters); int k_blocks = cur_k / 16; int cur_group_id = k_blocks / (group_blocks * (w_type == host::kFE2M1f ? 2 : 1)); int4* sh_s_stage = sh_s + s_sh_stage * pipe; if constexpr (w_type_id != host::kFE2M1f.id()) { reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; } else if constexpr (group_blocks == 1 || thread_k_blocks > 4) { reinterpret_cast(&frag_s[k % 2])[0] = reinterpret_cast(sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; } else { reinterpret_cast(&frag_s[k % 2])[0] = reinterpret_cast(sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride) + k % 2]; } } } return; } // Act-order case // Determine K of the "current" thread-block int cur_k = slice_k_start + tb_k * full_pipe; if (cur_k >= prob_k || cur_k >= slice_k_finish) { return; } // Reset (to current thread-block) since we read g_idx portion from the // shared memory cur_k = 0; // Progress to current iteration cur_k += k_iter_size * (k % b_sh_wr_iters); // Determine "position" inside the thread-block (based on warp and // thread-id) auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N int warp_row = warp_id / n_warps; int warp_col = warp_id % n_warps; cur_k += warp_row * 16; auto th_id = threadIdx.x % 32; cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix int s_col_shift = /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + (th_id / 4) * act_s_col_stride; if (is_same_group[pipe]) { if (k % 2 == 0) { *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + s_col_shift]; } else { *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); } for (int i = 1; i < 4; i++) { *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); } return; } int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); constexpr int k_frag_offsets[4] = {0, 1, 8, 9}; // Tensor core offsets per thread #pragma unroll for (int i = 0; i < 4; i++) { int actual_k = cur_k + k_frag_offsets[i]; int group_id = sh_g_idx_int_ptr[actual_k]; int rel_group_id = group_id - sh_first_group_id; *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = sh_s[rel_group_id * s_sh_stride + s_col_shift]; } }; auto fetch_zp_to_registers = [&](int k, int full_pipe) { // This code does not handle group_blocks == 0, // which signifies act_order. // has_zp implies AWQ, which doesn't have act_order, static_assert(!has_zp || group_blocks != 0); if constexpr (has_zp && !is_zp_float) { int pipe = full_pipe % stages; if constexpr (group_blocks == -1) { // load only when starting a new slice if (k == 0 && full_pipe == 0) { #pragma unroll for (int i = 0; i < num_ints_per_thread; i++) { frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp))[zp_sh_rd + i]; } } } else if constexpr (group_blocks >= thread_k_blocks) { if (k % b_sh_wr_iters == 0) { int4* sh_zp_stage = sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); #pragma unroll for (int i = 0; i < num_ints_per_thread; i++) { frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; } } } else { auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; int warp_row = warp_id / n_warps; int cur_k = warp_row * 16; cur_k += k_iter_size * (k % b_sh_wr_iters); int k_blocks = cur_k / 16; int cur_group_id = 0; // Suppress bogus and persistent divide-by-zero warning #pragma nv_diagnostic push #pragma nv_diag_suppress divide_by_zero cur_group_id = k_blocks / group_blocks; #pragma nv_diagnostic pop int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; sh_zp_stage += cur_group_id * zp_sh_stride; #pragma unroll for (int i = 0; i < num_ints_per_thread; i++) { frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; } } } else if constexpr (has_zp && is_zp_float) { int pipe = full_pipe % stages; if constexpr (group_blocks != -1) { if constexpr (group_blocks >= thread_k_blocks) { if (k % b_sh_wr_iters == 0) { int4* sh_zp_stage = sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); reinterpret_cast(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd]; } } else { auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; int warp_row = warp_id / n_warps; int cur_k = warp_row * 16; cur_k += k_iter_size * (k % b_sh_wr_iters); int k_blocks = cur_k / 16; // Suppress bogus and persistent divide-by-zero warning #pragma nv_diagnostic push #pragma nv_diag_suppress divide_by_zero int cur_group_id = k_blocks / group_blocks; #pragma nv_diagnostic pop int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; reinterpret_cast(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd + cur_group_id * zp_sh_stride]; } } } }; auto dequant_data = [&](int q, scalar_t2* frag_b_ptr) { dequant(q, frag_b_ptr); }; // Execute the actual tensor core matmul of a sub-tile. bool is_first_matmul_in_slice = true; auto matmul = [&](int k) { int k2 = k % 2; const bool is_new_zp = ((group_blocks != -1) && (group_blocks < thread_k_blocks || k == 0)) || (group_blocks == -1 && is_first_matmul_in_slice); if constexpr (has_zp && !is_zp_float) { if (is_new_zp) { if constexpr (group_blocks == -1) is_first_matmul_in_slice = false; int zp_quant_0, zp_quant_1; if constexpr (w_type.size_bits() == 4) { zp_quant_0 = frag_qzp[k2][0]; zp_quant_1 = zp_quant_0 >> 8; } else { static_assert(w_type.size_bits() == 8); zp_quant_0 = frag_qzp[k2][0]; zp_quant_1 = frag_qzp[k2][1]; } dequant_data(zp_quant_0, reinterpret_cast(&frag_zp)); dequant_data(zp_quant_1, reinterpret_cast(&frag_zp) + 2); } } if constexpr (!dequant_skip_flop && has_zp && is_zp_float) { if (is_new_zp) { reinterpret_cast(&frag_zp)[0] = reinterpret_cast(&frag_zpf[k2])[0]; } } // Commented out FP4/FP8 scale dequantization since we don't generate // kFE2M1f kernels to reduce compilation time // if constexpr (w_type == host::kFE2M1f) { // int s_quant_0 = reinterpret_cast(frag_s[k2])[0]; // int s_quant_1 = reinterpret_cast(frag_s[k2])[1]; // // dequant_fp8_scales( // s_quant_0, reinterpret_cast(&frag_s[k2])); // dequant_fp8_scales( // s_quant_1, reinterpret_cast(&frag_s[k2]) + 2); // } // We have the m dimension as the inner loop in order to encourage overlapping // dequantization and matmul operations. #pragma unroll for (int j = 0; j < 4; j++) { FragB frag_b0; FragB frag_b1; int b_quant_0, b_quant_1; if constexpr (w_type_id == host::kFE2M1f.id()) { b_quant_1 = frag_b_quant[k2][0][j]; b_quant_0 = b_quant_1 << 8; } else if constexpr (w_type.size_bits() == 4) { b_quant_0 = frag_b_quant[k2][0][j]; b_quant_1 = b_quant_0 >> 8; } else { static_assert(w_type.size_bits() == 8); int* frag_b_quant_ptr = reinterpret_cast(frag_b_quant[k2]); b_quant_0 = frag_b_quant_ptr[j * 2 + 0]; b_quant_1 = frag_b_quant_ptr[j * 2 + 1]; } dequant_data(b_quant_0, reinterpret_cast(&frag_b0)); dequant_data(b_quant_1, reinterpret_cast(&frag_b1)); if constexpr (dequant_skip_flop && has_zp && !is_zp_float) { sub_zp(frag_b0, frag_zp[j], 0); sub_zp(frag_b1, frag_zp[j], 1); } // Apply scale to frag_b0 if constexpr (has_act_order) { static_assert(group_blocks != -1); scale4( frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j], act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0); scale4( frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j], act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1); } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float && group_blocks == -1) { int idx = (threadIdx.x / 4) % 2; scalar_t2 s2 = Dtype::nums2num2( reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 1])[idx]); if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2); scale_and_sub(frag_b0, s2.x, frag_zp[j].x); scale_and_sub(frag_b1, s2.y, frag_zp[j].y); } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1) { if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], *reinterpret_cast(&frag_s[k2][j])); scale_and_sub(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x); scale_and_sub(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y); } else if constexpr (group_blocks != -1) { scale(frag_b0, frag_s[k2][j], 0); scale(frag_b1, frag_s[k2][j], 1); } #pragma unroll for (int i = 0; i < thread_m_blocks; i++) { if constexpr (m_block_size_8) { mma_trans(frag_a[k2][i], frag_b0, frag_b1, frag_c[i][j][0]); } else { mma(frag_a[k2][i], frag_b0, frag_c[i][j][0]); mma(frag_a[k2][i], frag_b1, frag_c[i][j][1]); } } } }; // Since we slice across the k dimension of a tile in order to increase the // number of warps while keeping the n dimension of a tile reasonable, we have // multiple warps that accumulate their partial sums of the same output // location; which we have to reduce over in the end. We do in shared memory. auto thread_block_reduce = [&]() { constexpr int red_off = threads / b_sh_stride_threads / 2; if (red_off >= 1) { auto red_idx = threadIdx.x / b_sh_stride_threads; constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2; constexpr int red_sh_delta = b_sh_stride_threads; int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + (threadIdx.x % b_sh_stride_threads); // Parallel logarithmic shared memory reduction. We make sure to avoid any // unnecessary read or write iterations, e.g., for two warps we write only // once by warp 1 and read only once by warp 0. #pragma unroll for (int m_block = 0; m_block < thread_m_blocks; m_block++) { #pragma unroll for (int i = red_off; i > 0; i /= 2) { if (i <= red_idx && red_idx < 2 * i) { #pragma unroll for (int j = 0; j < 4 * 2; j += (m_block_size_8 ? 2 : 1)) { int red_sh_wr = red_sh_delta * j + (red_sh_rd - red_sh_stride * i); if (i < red_off) { float* c_rd = reinterpret_cast(&sh_red[red_sh_delta * j + red_sh_rd]); float* c_wr = reinterpret_cast(&sh_red[red_sh_wr]); #pragma unroll for (int k = 0; k < 4; k++) reinterpret_cast(frag_c)[4 * 2 * m_block + j][k] += c_rd[k] + c_wr[k]; } sh_red[red_sh_wr] = reinterpret_cast(&frag_c)[4 * 2 * m_block + j]; } } __syncthreads(); } if (red_idx == 0) { #pragma unroll for (int i = 0; i < 4 * 2; i += (m_block_size_8 ? 2 : 1)) { float* c_rd = reinterpret_cast(&sh_red[red_sh_delta * i + red_sh_rd]); #pragma unroll for (int j = 0; j < 4; j++) reinterpret_cast(frag_c)[4 * 2 * m_block + i][j] += c_rd[j]; } } __syncthreads(); } } }; // Since multiple threadblocks may process parts of the same column slice, we // finally have to globally reduce over the results. As the striped // partitioning minimizes the number of such reductions and our outputs are // usually rather small, we perform this reduction serially in L2 cache. auto global_reduce_fp16 = [&](bool first = false, bool last = false) { // We are very careful here to reduce directly in the output buffer to // maximize L2 cache utilization in this step. To do this, we write out // results in FP16 (but still reduce with FP32 compute). constexpr int active_threads = 32 * thread_n_blocks / 4; bool is_th_active = threadIdx.x < active_threads; if (!is_th_active) { return; } int c_gl_stride = prob_n / 8; int c_gl_wr_delta_o = 8 * c_gl_stride; int c_gl_wr_delta_i = 4 * (active_threads / 32); int c_gl_wr; if constexpr (m_block_size_8) { c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + 4 * (threadIdx.x / 32) + (threadIdx.x % 32) / 8; c_gl_wr += (2 * thread_n_blocks) * slice_col; } else { c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) + 4 * (threadIdx.x / 32) + threadIdx.x % 4; c_gl_wr += (2 * thread_n_blocks) * slice_col; } constexpr int c_sh_wr_delta = active_threads; int c_sh_wr = threadIdx.x; if (!first) { #pragma unroll for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { int c_idx; if constexpr (m_block_size_8) c_idx = c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i; else c_idx = c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2); if (c_idx / c_gl_stride < block_num_valid_tokens) { int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride]; int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride; sh_red[c_sh_wr + c_sh_wr_delta * i] = C[true_idx]; } } } #pragma unroll for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { if (!first) { int4 c_red = sh_red[c_sh_wr + i * c_sh_wr_delta]; #pragma unroll for (int j = 0; j < 2 * 4; j++) { int delta = 0; if constexpr (m_block_size_8) { delta = j % 2 == 1 ? -2 : 0; } reinterpret_cast(&frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta] += Dtype::num2float(reinterpret_cast(&c_red)[j]); } } if (!last) { int4 c; #pragma unroll for (int j = 0; j < 2 * 4; j++) { int delta = 0; if constexpr (m_block_size_8) { delta = j % 2 == 1 ? -2 : 0; } reinterpret_cast(&c)[j] = Dtype::float2num(reinterpret_cast(&frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta]); } int c_idx; if constexpr (m_block_size_8) c_idx = c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i; else c_idx = c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2); if (c_idx / c_gl_stride < block_num_valid_tokens) { int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride]; int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride; C[true_idx] = c; } } } }; // Globally reduce over threadblocks that compute the same column block. // We use a tmp C buffer to reduce in full fp32 precision. auto global_reduce_fp32 = [&](bool first = false, bool last = false) { constexpr int tb_m = thread_m_blocks * 16; constexpr int tb_n = thread_n_blocks * 16; constexpr int c_size = tb_m * tb_n * sizeof(float) / 16; constexpr int active_threads = 32 * thread_n_blocks / 4; bool is_th_active = threadIdx.x < active_threads; constexpr int num_floats = thread_m_blocks * 4 * 2 * 4; constexpr int th_size = num_floats * sizeof(float) / 16; int c_cur_offset = locks_off * c_size; if (!is_th_active) { return; } if (!first) { float* frag_c_ptr = reinterpret_cast(&frag_c); #pragma unroll for (int k = 0; k < th_size; k++) { if constexpr (m_block_size_8) { if (k % 2) continue; } else { if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens) continue; } sh_red[threadIdx.x] = C_tmp[c_cur_offset + active_threads * k + threadIdx.x]; float* sh_c_ptr = reinterpret_cast(&sh_red[threadIdx.x]); #pragma unroll for (int f = 0; f < 4; f++) { frag_c_ptr[k * 4 + f] += sh_c_ptr[f]; } } } if (!last) { int4* frag_c_ptr = reinterpret_cast(&frag_c); #pragma unroll for (int k = 0; k < th_size; k++) { if constexpr (m_block_size_8) { if (k % 2) continue; } else { if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens) continue; } C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k]; } } }; // Write out the reduce final result in the correct layout. We only actually // reshuffle matrix fragments in this step, the reduction above is performed // in fragment layout. auto write_result = [&](bool last) { int c_gl_stride = prob_n / 8; constexpr int c_sh_stride = 2 * thread_n_blocks + 1; int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); constexpr int c_sh_rd_delta = c_sh_stride * (threads / (2 * thread_n_blocks)); int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); c_gl_wr += (2 * thread_n_blocks) * slice_col; int c_sh_wr; if constexpr (m_block_size_8) { c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) + (threadIdx.x % 32) / 4; c_sh_wr += 64 * (threadIdx.x / 32); } else { c_sh_wr = (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; c_sh_wr += 32 * (threadIdx.x / 32); } int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); // We first reorder in shared memory to guarantee the most efficient final // global write patterns auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) { scalar_t2 res = Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1)); // For per-column quantization we finally apply the scale here (only for // 4-bit) if constexpr ( !has_act_order && group_blocks == -1 && w_type.size_bits() == 4 && (has_zp && dequant_skip_flop || !has_zp)) { scalar_t2 tmp_scale = s[0]; if constexpr (m_block_size_8) { tmp_scale = Dtype::num2num2(reinterpret_cast(&s[0])[(threadIdx.x % 8) / 4]); } res = __hmul2(res, tmp_scale); } if constexpr (w_type == host::kFE2M1f && s_type == host::kFE4M3fn) { if (!mul_topk_weights) { res = __hmul2(res, global_scale); } } if (has_bias && last) { scalar_t2 tmp_bias = b_bias[0]; if constexpr (m_block_size_8) { tmp_bias = Dtype::num2num2(reinterpret_cast(&b_bias[0])[(threadIdx.x % 8) / 4]); } res = __hadd2(res, tmp_bias); } if constexpr (m_block_size_8) { ((scalar_t*)sh_red)[idx] = res.x; ((scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y; } else { ((scalar_t2*)sh_red)[idx] = res; } }; if (threadIdx.x / 32 < thread_n_blocks / 4) { #pragma unroll for (int i = 0; i < thread_m_blocks; i++) { #pragma unroll for (int j = 0; j < 4; j++) { if constexpr (m_block_size_8) { int wr = c_sh_wr + 16 * j; write( wr, frag_c[i][j][0][0], frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0], frag_bias[j / 2][2 * (j % 2) + 0]); write( wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 1], frag_bias[j / 2][2 * (j % 2) + 1]); } else { int wr = c_sh_wr + 8 * j; write( wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0], frag_bias[j / 2][2 * (j % 2) + 0]); write( wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0], frag_bias[j / 2][2 * (j % 2) + 0]); write( wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1], frag_bias[j / 2][2 * (j % 2) + 1]); write( wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1], frag_bias[j / 2][2 * (j % 2) + 1]); } } c_sh_wr += 16 * (4 * c_sh_stride); } } __syncthreads(); #pragma unroll for (int i = 0; i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); i++) { int row = c_gl_wr / c_gl_stride; if (row < block_num_valid_tokens) { int64_t sorted_row = sh_block_sorted_ids[row]; int64_t true_idx = sorted_row * c_gl_stride + c_gl_wr % c_gl_stride; scalar_t2 topk_weight_score; if (mul_topk_weights) topk_weight_score = sh_block_topk_weights[row]; if (use_atomic_add && slice_count > 1 || mul_topk_weights) { scalar_t2* C_half2 = reinterpret_cast(&C[true_idx]); scalar_t2* sh_red_half2 = reinterpret_cast(&sh_red[c_sh_rd]); #pragma unroll for (int a = 0; a < 4; a++) { scalar_t2 res = sh_red_half2[a]; if (mul_topk_weights) { res = __hmul2(res, topk_weight_score); } if (use_atomic_add && slice_count > 1) { atomicAdd(&C_half2[a], res); } else { C_half2[a] = res; }; } } else { C[true_idx] = sh_red[c_sh_rd]; } c_gl_wr += c_gl_wr_delta; c_sh_rd += c_sh_rd_delta; } } __syncthreads(); }; // Start global fetch and register load pipelines. auto start_pipes = [&]() { #pragma unroll for (int i = 0; i < stages - 1; i++) { if (has_act_order && i == 0) { int last_g_idx = slice_k_start + stages * tb_k * 2; if (last_g_idx >= prob_k) { last_g_idx = prob_k - 1; } fetch_act_order_scales_to_shared(true, g_idx[slice_k_start], g_idx[last_g_idx]); } if constexpr (has_zp && !is_zp_float && group_blocks == -1) { if (i == 0) { fetch_col_zp_to_shared(); if constexpr (!dequant_skip_flop) { fetch_col_scale_to_shared(); } } } fetch_to_shared(i, i, i < slice_iters, i); } zero_accums(); wait_for_stage(); init_same_group(0); fetch_to_registers(0, 0); fetch_scales_to_registers(0, 0); fetch_zp_to_registers(0, 0); a_gl_rd_col += a_gl_rd_delta_o * (stages - 1); if constexpr (has_act_order) { slice_k_start_shared_fetch += tb_k * (stages - 1); } }; if (slice_iters) { start_pipes(); } // Main loop. while (slice_iters) { // We unroll over both the global fetch and the register load pipeline to // ensure all shared memory accesses are static. Note that both pipelines // have even length meaning that the next iteration will always start at // index 0. for (int stage_group_id = 0; stage_group_id < max_num_stage_groups; stage_group_id++) { #pragma unroll for (int pipe = 0; pipe < stages;) { #pragma unroll for (int k = 0; k < b_sh_wr_iters; k++) { int idx = (pipe >= stages && stage_group_id == max_num_stage_groups - 1) ? (pipe - stages) : (pipe + stage_group_id * stages); fetch_to_registers(k + 1, pipe % stages, idx); fetch_scales_to_registers(k + 1, pipe); fetch_zp_to_registers(k + 1, pipe); if (k == b_sh_wr_iters - 2) { int idx = (pipe >= 1 && stage_group_id == max_num_stage_groups - 1) ? (pipe - 1) : (pipe + (stage_group_id + 1) * stages - 1); fetch_to_shared((pipe + stages - 1) % stages, pipe, slice_iters >= stages, idx); pipe++; wait_for_stage(); init_same_group(pipe % stages); } matmul(k); } slice_iters--; if (slice_iters == 0) { break; } } a_gl_rd_col += a_gl_rd_delta_o * stages; if constexpr (has_act_order) { slice_k_start += tb_k * stages; if (slice_k_start < prob_k) { slice_k_start_shared_fetch += tb_k * stages; int first_group_id = g_idx[slice_k_start]; int last_g_idx = slice_k_start + stages * tb_k * 2; if (last_g_idx >= prob_k) { last_g_idx = prob_k - 1; } int last_group_id = g_idx[last_g_idx]; if (last_group_id >= sh_first_group_id + sh_num_groups) { fetch_act_order_scales_to_shared(false, first_group_id, last_group_id); __syncthreads(); } } } if (slice_iters == 0) { break; } } // Process results and, if necessary, proceed to the next column slice. // While this pattern may not be the most readable, other ways of writing // the loop seemed to noticeably worse performance after compilation. if (slice_iters == 0) { cp_async_wait<0>(); bool last = slice_idx == slice_count - 1; // For per-column scales, we only fetch them here in the final step before // write-out if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) { if (w_type.size_bits() == 8 || (last || use_atomic_add)) { if (s_sh_wr_pred) { cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); } cp_async_fence(); } } thread_block_reduce(); if (has_bias && last) { __syncthreads(); cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd], threadIdx.x < 16 * thread_n_blocks / 8); cp_async_fence(); } if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) { if (w_type.size_bits() == 8 || (last || use_atomic_add)) { cp_async_wait<0>(); __syncthreads(); if (threadIdx.x / 32 < thread_n_blocks / 4) { reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; if constexpr (m_block_size_8) { int idx = (threadIdx.x / 4) % 2; scalar_t2* frag_s_half2 = reinterpret_cast(frag_s); #pragma unroll for (int i = 0; i < 8; i++) { frag_s_half2[i] = Dtype::num2num2(reinterpret_cast(&frag_s_half2[i])[idx]); } } } } } // For 8-bit channelwise, we apply the scale before the global reduction // that converts the fp32 results to fp16 (so that we avoid possible // overflow in fp16) if constexpr ( !has_act_order && group_blocks == -1 && w_type.size_bits() == 8 && (has_zp && dequant_skip_flop || !has_zp)) { if (threadIdx.x / 32 < thread_n_blocks / 4) { #pragma unroll for (int i = 0; i < thread_m_blocks; i++) { #pragma unroll for (int j = 0; j < 4; j++) { scale_float(reinterpret_cast(&frag_c[i][j][0][0]), frag_s[j / 2][2 * (j % 2) + 0]); scale_float( reinterpret_cast(&frag_c[i][j][0][2]), frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]); if constexpr (!m_block_size_8) { scale_float(reinterpret_cast(&frag_c[i][j][1][0]), frag_s[j / 2][2 * (j % 2) + 1]); scale_float(reinterpret_cast(&frag_c[i][j][1][2]), frag_s[j / 2][2 * (j % 2) + 1]); } } } } } if (slice_count > 1 && !use_atomic_add) { // only globally reduce if there is more than one block in a slice barrier_acquire(&locks[locks_off], slice_idx); if (use_fp32_reduce) { global_reduce_fp32(slice_idx == 0, last); } else { global_reduce_fp16(slice_idx == 0, last); } barrier_release(&locks[locks_off], last); } if (has_bias && last) { cp_async_wait<0>(); __syncthreads(); reinterpret_cast(&frag_bias)[0] = sh_bias[bias_sh_rd]; reinterpret_cast(&frag_bias)[1] = sh_bias[bias_sh_rd + 4]; __syncthreads(); } if (use_atomic_add && slice_count > 1 && slice_idx != 0) wait_negative_and_add(&locks[locks_off]); if (last || use_atomic_add) // only the last block in a slice actually writes the result write_result(last); int old_slice_row = slice_row; slice_row = 0; slice_col_par++; slice_col++; is_first_matmul_in_slice = true; init_slice(); // Should we load A matrix in next slice? // `slice_col == 0`: when move to a new moe block // `old_slice_row > 0`: // when the last slice is not starting from k_index == 0 // (only happen when it is the first slice of a threadblock) // `prob_k > thread_k_blocks * 16 * stages * max_num_stage_groups`: // when the required shared memory size is larger than // the remaining shared memory if (slice_col == 0 || old_slice_row || prob_k > thread_k_blocks * 16 * stages * max_num_stage_groups) { should_load_a = true; } else { should_load_a = false; } if (slice_iters) { a_gl_rd_col = (threadIdx.x % a_gl_rd_delta_o); #pragma unroll for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles; if (slice_col == 0) { #pragma unroll for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride; } bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; // Update slice k/n for scales loading if constexpr (has_act_order) { slice_k_start = tb_k * slice_row; slice_k_finish = slice_k_start + tb_k * slice_iters; slice_k_start_shared_fetch = slice_k_start; slice_n_offset = act_s_col_tb_stride * slice_col; } else { s_gl_rd = s_sh_stride * slice_col + threadIdx.x; zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; } start_pipes(); } } } } } // namespace device::marlin_moe #endif ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh ================================================ /* * Modified by Neural Magic * Copyright (C) Marlin.2024 Elias Frantar * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ /* * Adapted from https://github.com/IST-DASLab/marlin */ #pragma once #include #include #include "kernel.h" #include "marlin_template.h" namespace device::marlin_moe { __global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){}; using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 template __global__ void permute_cols_kernel( int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr, int4* __restrict__ out_int4_ptr, const int32_t* __restrict__ sorted_token_ids_ptr, const int32_t* __restrict__ expert_ids_ptr, const int32_t* __restrict__ num_tokens_past_padded_ptr, int size_m, int size_k, int top_k) {}; #else // For a given "a" of size [M,K] performs a permutation of the K columns based // on the given "perm" indices. template __global__ void permute_cols_kernel( int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr, int4* __restrict__ out_int4_ptr, const int32_t* __restrict__ sorted_token_ids_ptr, const int32_t* __restrict__ expert_ids_ptr, const int32_t* __restrict__ num_tokens_past_padded_ptr, int size_m, int size_k, int top_k) { int num_tokens_past_padded = num_tokens_past_padded_ptr[0]; int num_moe_blocks = div_ceil(num_tokens_past_padded, moe_block_size); int32_t block_sorted_ids[moe_block_size]; int block_num_valid_tokens = 0; int64_t old_expert_id = 0; int64_t expert_id = 0; int row_stride = size_k * sizeof(half) / 16; auto read_moe_block_data = [&](int block_id) { block_num_valid_tokens = moe_block_size; int4* tmp_block_sorted_ids = reinterpret_cast(block_sorted_ids); for (int i = 0; i < moe_block_size / 4; i++) { tmp_block_sorted_ids[i] = ((int4*)sorted_token_ids_ptr)[block_id * moe_block_size / 4 + i]; } for (int i = 0; i < moe_block_size; i++) { if (block_sorted_ids[i] >= size_m * top_k) { block_num_valid_tokens = i; break; }; } }; auto permute_row = [&](int row) { int iters = size_k / default_threads; int rest = size_k % default_threads; int in_offset = (row / top_k) * row_stride; int out_offset = row * row_stride; half const* a_row_half = reinterpret_cast(a_int4_ptr + in_offset); half* out_half = reinterpret_cast(out_int4_ptr + out_offset); int base_k = 0; for (int i = 0; i < iters; i++) { auto cur_k = base_k + threadIdx.x; int src_pos = perm_int_ptr[cur_k]; out_half[cur_k] = a_row_half[src_pos]; base_k += default_threads; } if (rest) { if (threadIdx.x < rest) { auto cur_k = base_k + threadIdx.x; int src_pos = perm_int_ptr[cur_k]; out_half[cur_k] = a_row_half[src_pos]; } } }; for (int index = blockIdx.x; index < num_moe_blocks; index += gridDim.x) { old_expert_id = expert_id; int tmp_expert_id = expert_ids_ptr[index]; if (tmp_expert_id == -1) continue; expert_id = tmp_expert_id; perm_int_ptr += (expert_id - old_expert_id) * size_k; read_moe_block_data(index); for (int i = 0; i < block_num_valid_tokens; i++) permute_row(block_sorted_ids[i]); } } typedef struct { int thread_k; int thread_n; int num_threads; } thread_config_t; thread_config_t small_batch_thread_configs[] = { // Ordered by priority // thread_k, thread_n, num_threads {128, 128, 256}, {64, 128, 128}}; thread_config_t large_batch_thread_configs[] = { // Ordered by priority // thread_k, thread_n, num_threads {64, 256, 256}, {64, 128, 128}}; typedef struct { int blocks_per_sm; thread_config_t tb_cfg; } exec_config_t; int get_scales_cache_size( thread_config_t const& th_config, int prob_m, int prob_n, int prob_k, int num_bits, int group_size, bool has_act_order, bool is_k_full) { bool cache_scales_chunk = has_act_order && !is_k_full; int tb_n = th_config.thread_n; int tb_k = th_config.thread_k; // Get max scale groups per thread-block int tb_groups; if (group_size == -1) { tb_groups = 1; } else if (group_size == 0) { tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size } else { tb_groups = div_ceil(tb_k, group_size); } if (cache_scales_chunk) { int load_groups = tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K load_groups = max(load_groups, 32); // We load at least 32 scale groups return load_groups * tb_n * 2; } else { int tb_scales = tb_groups * tb_n * 2; return tb_scales * pipe_stages; } } int get_kernel_cache_size( thread_config_t const& th_config, bool m_block_size_8, int thread_m_blocks, int prob_m, int prob_n, int prob_k, int num_bits, int group_size, bool has_act_order, bool is_k_full, int has_zp, int is_zp_float) { int pack_factor = 32 / num_bits; // Get B size int tb_k = th_config.thread_k; int tb_n = th_config.thread_n; int tb_m = thread_m_blocks * 16; // shm size for block_sorted_ids/rd_block_sorted_ids/block_topk_weights // both of them requires tb_m * 4 bytes (tb_m * int32 or tb_m * float32) int sh_block_meta_size = tb_m * 4; int sh_a_size = pipe_stages * (tb_m * tb_k) * 2; int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4; int sh_red_size = tb_m * (tb_n + 8) * 2; int sh_bias_size = tb_n * 2; int tmp_size = (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size; tmp_size = max(max(sh_b_size, sh_red_size), tmp_size); int sh_s_size = get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full); int sh_g_idx_size = has_act_order && !is_k_full ? pipe_stages * tb_k / 4 : 0; int sh_zp_size = 0; if (has_zp) { if (is_zp_float) sh_zp_size = sh_s_size; else if (num_bits == 4) sh_zp_size = sh_s_size / 4; else if (num_bits == 8) sh_zp_size = sh_s_size / 2; } int total_size = tmp_size + sh_a_size + sh_s_size + sh_zp_size + sh_g_idx_size + sh_block_meta_size; return total_size; } bool is_valid_config( thread_config_t const& th_config, bool m_block_size_8, int thread_m_blocks, int prob_m, int prob_n, int prob_k, int num_bits, int group_size, bool has_act_order, bool is_k_full, int has_zp, int is_zp_float, int max_shared_mem) { // Sanity if (th_config.thread_k == -1 || th_config.thread_n == -1 || th_config.num_threads == -1) { return false; } // Verify K/N are divisible by thread K/N if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) { return false; } // Verify min for thread K/N if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) { return false; } // num_threads must be at least 128 (= 4 warps) if (th_config.num_threads < 128) { return false; } // Check that pipeline fits into cache int cache_size = get_kernel_cache_size( th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float); return cache_size + 512 <= max_shared_mem; } #define _GET_IF( \ W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \ else if ( \ q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && thread_n_blocks == THREAD_N_BLOCKS && \ thread_k_blocks == THREAD_K_BLOCKS && m_block_size_8 == M_BLOCK_SIZE_8 && group_blocks == GROUP_BLOCKS && \ num_threads == NUM_THREADS && is_zp_float == IS_ZP_FLOAT) { \ constexpr auto S_TYPE = W_TYPE == host::kFE2M1f \ ? (GROUP_BLOCKS == 1 ? host::kFE4M3fn : host::kFE8M0fnu) \ : (std::is_same::value ? host::kFloat16 : host::kBFloat16); \ kernel = Marlin< \ scalar_t, \ W_TYPE.id(), \ S_TYPE.id(), \ NUM_THREADS, \ THREAD_M_BLOCKS, \ THREAD_N_BLOCKS, \ THREAD_K_BLOCKS, \ M_BLOCK_SIZE_8, \ pipe_stages, \ GROUP_BLOCKS, \ IS_ZP_FLOAT>; \ } // COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false) // this is the most common cases // BIGGROUP: cases for big group size (group_blocks in [-1, 8]) // FZP: cases for float-zero-point (is_zp_float = true) // ACT: cases for act order case (group_blocks == 0) // NVFP4: cases for nvfp4(e2m1) (group_blocks == 1) // MXFP4: cases for mxfp4(e2m1) (group_blocks == 2) #define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) #define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) #define COMMON_GET_IF(W_TYPE) \ COMMON_GET_IF_M1(W_TYPE, 8, 8, 256) \ COMMON_GET_IF_M1(W_TYPE, 8, 4, 128) \ COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \ COMMON_GET_IF_M234(W_TYPE, 8, 4, 128) #define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) #define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) #define BIGGROUP_GET_IF(W_TYPE) \ BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256) \ BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128) \ BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \ BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128) #define NVFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) #define NVFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) #define NVFP4_GET_IF(W_TYPE) \ NVFP4_GET_IF_M1(W_TYPE, 8, 8, 256) \ NVFP4_GET_IF_M1(W_TYPE, 8, 4, 128) \ NVFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \ NVFP4_GET_IF_M234(W_TYPE, 8, 4, 128) #define MXFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) #define MXFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) #define MXFP4_GET_IF(W_TYPE) \ MXFP4_GET_IF_M1(W_TYPE, 8, 8, 256) \ MXFP4_GET_IF_M1(W_TYPE, 8, 4, 128) \ MXFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \ MXFP4_GET_IF_M234(W_TYPE, 8, 4, 128) // We currently have 4-bit models only with group_blocks == 4 #define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) #define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) #define FZP_GET_IF(W_TYPE) \ FZP_GET_IF_M1(W_TYPE, 8, 8, 256) \ FZP_GET_IF_M1(W_TYPE, 8, 4, 128) \ FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \ FZP_GET_IF_M234(W_TYPE, 8, 4, 128) // We currently have 4-bit models only with group_blocks == 4 #define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) #define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) #define ACT_GET_IF(W_TYPE) \ ACT_GET_IF_M1(W_TYPE, 8, 8, 256) \ ACT_GET_IF_M1(W_TYPE, 8, 4, 128) \ ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \ ACT_GET_IF_M234(W_TYPE, 8, 4, 128) template MarlinFuncPtr get_marlin_kernel( const host::ScalarType q_type, int thread_m_blocks, int thread_n_blocks, int thread_k_blocks, bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks, int num_threads, bool is_zp_float) { int num_bits = q_type.size_bits(); auto kernel = MarlinDefault; if (false) { } COMMON_GET_IF(host::kU4) COMMON_GET_IF(host::kU4B8) COMMON_GET_IF(host::kU8B128) NVFP4_GET_IF(host::kFE2M1f) BIGGROUP_GET_IF(host::kFE4M3fn) ACT_GET_IF(host::kU4B8) ACT_GET_IF(host::kU8B128) if (std::is_same::value) { if (false) { } MXFP4_GET_IF(host::kFE2M1f) } return kernel; } template exec_config_t determine_exec_config( const host::ScalarType& q_type, int prob_m, int prob_n, int prob_k, int thread_m_blocks, bool m_block_size_8, int num_bits, int group_size, bool has_act_order, bool is_k_full, bool has_zp, bool is_zp_float, int max_shared_mem) { exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}}; thread_config_t* thread_configs = thread_m_blocks > 1 ? large_batch_thread_configs : small_batch_thread_configs; int thread_configs_size = thread_m_blocks > 1 ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t) : sizeof(small_batch_thread_configs) / sizeof(thread_config_t); int count = 0; constexpr int device_max_reg_size = 255 * 1024; for (int i = 0; i < thread_configs_size; i++) { thread_config_t th_config = thread_configs[i]; if (!is_valid_config( th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float, max_shared_mem)) { continue; } int cache_size = get_kernel_cache_size( th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float); int group_blocks = 0; if (!has_act_order) { group_blocks = group_size == -1 ? -1 : (group_size / 16); } auto kernel = get_marlin_kernel( q_type, thread_m_blocks, th_config.thread_n / 16, th_config.thread_k / 16, m_block_size_8, has_act_order, has_zp, group_blocks, th_config.num_threads, is_zp_float); if (kernel == MarlinDefault) continue; if (thread_m_blocks > 1) { exec_cfg = {1, th_config}; break; } else { cudaFuncAttributes attr; cudaFuncGetAttributes(&attr, kernel); int reg_size = max(attr.numRegs, 1) * th_config.num_threads * 4; int allow_count = min(device_max_reg_size / reg_size, max_shared_mem / (cache_size + 1024)); allow_count = max(min(allow_count, 4), 1); if (allow_count > count) { count = allow_count; exec_cfg = {count, th_config}; }; } } return exec_cfg; } template void marlin_mm( const void* A, const void* B, void* C, void* C_tmp, void* b_bias, void* s, void* s2, void* zp, void* g_idx, void* perm, void* a_tmp, void* sorted_token_ids, void* expert_ids, void* num_tokens_past_padded, void* topk_weights, int moe_block_size, int top_k, bool mul_topk_weights, bool is_ep, int prob_m, int prob_n, int prob_k, void* workspace, host::ScalarType const& q_type, bool has_bias, bool has_act_order, bool is_k_full, bool has_zp, int num_groups, int group_size, int dev, cudaStream_t stream, int thread_k, int thread_n, int sms, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { int thread_m_blocks = div_ceil(moe_block_size, 16); bool m_block_size_8 = moe_block_size == 8; if (has_zp) { host::RuntimeCheck( q_type == host::kU4 || q_type == host::kU8, "q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str()); } else { host::RuntimeCheck( q_type == host::kU4B8 || q_type == host::kU8B128 || q_type == host::kFE4M3fn || q_type == host::kFE2M1f, "q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when " "has_zp = False. Got = ", q_type.str()); } host::RuntimeCheck( prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, ", ", prob_n, ", ", prob_k, "]"); int group_blocks = 0; if (has_act_order) { if (is_k_full) { host::RuntimeCheck(group_size != -1); group_blocks = group_size / 16; host::RuntimeCheck( prob_k % group_blocks == 0, "prob_k = ", prob_k, " is not divisible by group_blocks = ", group_blocks); } else { host::RuntimeCheck(group_size == 0); group_blocks = 0; } } else { if (group_size == -1) { group_blocks = -1; } else { group_blocks = group_size / 16; host::RuntimeCheck( prob_k % group_blocks == 0, "prob_k = ", prob_k, " is not divisible by group_blocks = ", group_blocks); } } int num_bits = q_type.size_bits(); const int4* A_ptr = (const int4*)A; const int4* B_ptr = (const int4*)B; int4* C_ptr = (int4*)C; int4* C_tmp_ptr = (int4*)C_tmp; const int4* bias_ptr = (const int4*)b_bias; const int4* s_ptr = (const int4*)s; const uint16_t* s2_ptr = (const uint16_t*)s2; const int4* zp_ptr = (const int4*)zp; const int* g_idx_ptr = (const int*)g_idx; const int* perm_ptr = (const int*)perm; int4* a_tmp_ptr = (int4*)a_tmp; const int32_t* sorted_token_ids_ptr = (const int32_t*)sorted_token_ids; const int32_t* expert_ids_ptr = (const int32_t*)expert_ids; const int32_t* num_tokens_past_padded_ptr = (const int32_t*)num_tokens_past_padded; const float* topk_weights_ptr = (const float*)topk_weights; int* locks = (int*)workspace; if (has_act_order) { // Permute A columns auto perm_kernel = permute_cols_kernel<8>; if (moe_block_size == 8) { } else if (moe_block_size == 16) perm_kernel = permute_cols_kernel<16>; else if (moe_block_size == 32) perm_kernel = permute_cols_kernel<32>; else if (moe_block_size == 48) perm_kernel = permute_cols_kernel<48>; else if (moe_block_size == 64) perm_kernel = permute_cols_kernel<64>; else host::Panic("unsupported moe_block_size ", moe_block_size); // clang-format off perm_kernel<<>>( A_ptr, perm_ptr, a_tmp_ptr, sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr, prob_m, prob_k, top_k); // clang-format on A_ptr = a_tmp_ptr; prob_m = prob_m * top_k; top_k = 1; // If we have a full K, then we can run the non-act-order version of Marlin // (since the weight rows are reordered by increasing group ids, and by // having a full K, we have full original groups) if (is_k_full) has_act_order = false; } int max_shared_mem = 0; host::RuntimeDeviceCheck(cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev)); host::RuntimeCheck(max_shared_mem > 0); // Set thread config exec_config_t exec_cfg; thread_config_t thread_tfg; if (thread_k != -1 && thread_n != -1) { thread_tfg = thread_config_t{thread_k, thread_n, default_threads}; exec_cfg = exec_config_t{1, thread_tfg}; host::RuntimeCheck(prob_n % thread_n == 0, "prob_n = ", prob_n, " is not divisible by thread_n = ", thread_n); host::RuntimeCheck(prob_k % thread_k == 0, "prob_k = ", prob_k, " is not divisible by thread_k = ", thread_k); } else { // Auto config exec_cfg = determine_exec_config( q_type, prob_m, prob_n, prob_k, thread_m_blocks, m_block_size_8, num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float, max_shared_mem); thread_tfg = exec_cfg.tb_cfg; } int num_threads = thread_tfg.num_threads; thread_k = thread_tfg.thread_k; thread_n = thread_tfg.thread_n; int blocks = sms * exec_cfg.blocks_per_sm; if (exec_cfg.blocks_per_sm > 1) max_shared_mem = max_shared_mem / exec_cfg.blocks_per_sm - 1024; int thread_k_blocks = thread_k / 16; int thread_n_blocks = thread_n / 16; host::RuntimeCheck( is_valid_config( thread_tfg, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float, max_shared_mem), "Invalid thread config: thread_m_blocks = ", thread_m_blocks, ", thread_k = ", thread_tfg.thread_k, ", thread_n = ", thread_tfg.thread_n, ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits, ", group_size = ", group_size, ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full, ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float, ", max_shared_mem = ", max_shared_mem); auto kernel = get_marlin_kernel( q_type, thread_m_blocks, thread_n_blocks, thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks, num_threads, is_zp_float); if (kernel == MarlinDefault) { host::Panic( "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ", prob_k, "]", ", has_act_order = ", has_act_order, ", num_groups = ", num_groups, ", group_size = ", group_size, ", thread_m_blocks = ", thread_m_blocks, ", thread_n_blocks = ", thread_n_blocks, ", thread_k_blocks = ", thread_k_blocks, ", num_bits = ", num_bits); } host::RuntimeDeviceCheck(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem)); // clang-format off kernel<<>>( A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, s_ptr, s2_ptr, zp_ptr, g_idx_ptr, sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr, topk_weights_ptr, top_k, mul_topk_weights, is_ep, num_groups, prob_m, prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce, max_shared_mem); // clang-format on } #endif } // namespace device::marlin_moe template void moe_wna16_marlin_gemm( tvm::ffi::TensorView a, tvm::ffi::TensorView c, tvm::ffi::TensorView b_q_weight, tvm::ffi::TensorView b_bias, tvm::ffi::TensorView b_scales, tvm::ffi::TensorView global_scale, tvm::ffi::TensorView b_zeros, tvm::ffi::TensorView g_idx, tvm::ffi::TensorView perm, tvm::ffi::TensorView workspace, tvm::ffi::TensorView sorted_token_ids, tvm::ffi::TensorView expert_ids, tvm::ffi::TensorView num_tokens_post_padded, tvm::ffi::TensorView topk_weights, tvm::ffi::TensorView a_tmp, tvm::ffi::TensorView c_tmp, int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, bool is_ep, int64_t b_q_type_id, int64_t size_m, int64_t size_n, int64_t size_k, bool has_act_order, bool has_bias, bool is_k_full, bool has_zp, int64_t num_groups, int64_t group_size, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { using namespace host; ScalarType const b_q_type = ScalarType::from_id(b_q_type_id); int pack_factor = 32 / b_q_type.size_bits(); if (moe_block_size != 8) { RuntimeCheck(moe_block_size % 16 == 0, "unsupported moe_block_size=", moe_block_size); RuntimeCheck(moe_block_size >= 16 && moe_block_size <= 64, "unsupported moe_block_size=", moe_block_size); } // Verify A RuntimeCheck(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0), ", size_m = ", size_m); RuntimeCheck(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1), ", size_k = ", size_k); // Verify B RuntimeCheck( size_k % device::marlin::tile_size == 0, "size_k = ", size_k, " is not divisible by tile_size = ", device::marlin::tile_size); RuntimeCheck( (size_k / device::marlin::tile_size) == b_q_weight.size(1), "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1), ", size_k = ", size_k, ", tile_size = ", device::marlin::tile_size); RuntimeCheck( b_q_weight.size(2) % device::marlin::tile_size == 0, "b_q_weight.size(2) = ", b_q_weight.size(2), " is not divisible by tile_size = ", device::marlin::tile_size); int64_t actual_size_n = (b_q_weight.size(2) / device::marlin::tile_size) * pack_factor; RuntimeCheck(size_n == actual_size_n, "size_n = ", size_n, ", actual_size_n = ", actual_size_n); // Verify device and strides auto device = SymbolicDevice{}; device.set_options(); TensorMatcher({-1, -1}).with_dtype().with_device(device).verify(a); device.verify(b_q_weight.device()); RuntimeCheck(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); device.verify(b_scales.device()); RuntimeCheck(b_scales.is_contiguous(), "b_scales is not contiguous"); // thread_k, thread_n, sms int thread_k = -1; int thread_n = -1; int sms = -1; DLDevice dl_device = device.unwrap(); int dev = dl_device.device_id; cudaStream_t stream = LaunchKernel::resolve_device(dl_device); RuntimeDeviceCheck(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev)); // Verify c (allocation done in Python) device.verify(c.device()); RuntimeCheck(c.is_contiguous(), "c is not contiguous"); RuntimeCheck( c.size(0) == size_m * top_k, "Shape mismatch: c.size(0) = ", c.size(0), ", size_m * topk = ", size_m * top_k); RuntimeCheck(c.size(1) == size_n, "Shape mismatch: c.size(1) = ", c.size(1), ", size_n = ", size_n); // Alloc c_tmp: SKIP, done in Python // Detect groupsize: b_scales rank and dims RuntimeCheck(b_scales.dim() == 3, "b_scales rank = ", b_scales.dim(), " is not 3"); RuntimeCheck(b_scales.size(2) == size_n, "b_scales dim 2 = ", b_scales.size(2), " is not size_n = ", size_n); RuntimeCheck( b_scales.size(1) == num_groups, "b_scales dim 1 = ", b_scales.size(1), " is not num_groups = ", num_groups); // Validate g_idx, perm (Optional unwrap done in Python; empty tensors when absent) if (g_idx.size(g_idx.dim() - 1) > 0 && perm.size(perm.dim() - 1) > 0) { device.verify(g_idx.device()); RuntimeCheck(g_idx.is_contiguous(), "g_idx is not contiguous"); device.verify(perm.device()); RuntimeCheck(perm.is_contiguous(), "perm is not contiguous"); int64_t g_idx_last = g_idx.size(g_idx.dim() - 1); int64_t perm_last = perm.size(perm.dim() - 1); RuntimeCheck( (g_idx_last == 0 && perm_last == 0) || (g_idx_last == size_k && perm_last == size_k), "Unexpected g_idx.size(-1) = ", g_idx_last, " and perm.size(-1) = ", perm_last, ", where size_k = ", size_k); } // has_act_order derivation: SKIP (passed as param) // Verify group_size consistency if (has_act_order) { // SKIP: a_tmp allocation done in Python if (is_k_full) { RuntimeCheck(num_groups > 1, "For act_order, num_groups must be > 1"); RuntimeCheck(size_k % num_groups == 0, "size_k = ", size_k, ", is not divisible by num_groups = ", num_groups); } } else { if (num_groups > 1) { RuntimeCheck( size_k % num_groups == 0, "size_k = ", size_k, ", is not divisible by b_scales.size(1) = ", num_groups); } } // Verify global_scale (Optional unwrap done in Python) int64_t global_scale_size = global_scale.size(0); if (global_scale_size > 0) { RuntimeCheck(b_q_type == kFE2M1f && group_size == 16, "global_scale can only be used for nvfp4 format."); } else { RuntimeCheck( !(b_q_type == kFE2M1f && group_size == 16), "the global_scale parameter must be passed for nvfp4 format."); } // Verify b_bias (Optional unwrap done in Python) if (has_bias) { device.verify(b_bias.device()); RuntimeCheck(b_bias.is_contiguous(), "b_bias is not contiguous"); RuntimeCheck(b_bias.size(1) == size_n, "b_bias.size(0) != size_n"); RuntimeCheck(b_bias.stride(1) == 1, "b_bias.stride(1) != 1"); } // b_zeros Optional unwrap + has_zp derivation: SKIP (done in Python) // Verify b_q_type vs has_zp if (has_zp) { device.verify(b_zeros.device()); RuntimeCheck(b_zeros.is_contiguous(), "b_zeros is not contiguous"); RuntimeCheck( b_q_type == kU4 || b_q_type == kU8, "b_q_type must be u4 or u8 when has_zp = True. Got = ", b_q_type.str()); } else { RuntimeCheck( b_q_type == kU4B8 || b_q_type == kU8B128 || b_q_type == kFE4M3fn || b_q_type == kFE2M1f, "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or " "float4_e2m1f when " "has_zp = False. Got = ", b_q_type.str()); } if (has_zp && is_zp_float) { RuntimeCheck( std::is_same::value, "Computation type must be float16 (half) when using float zero " "points."); } // Verify b_zeros if (has_zp) { RuntimeCheck(b_zeros.dim() == 3, "b_zeros rank = ", b_zeros.dim(), " is not 3"); if (is_zp_float) { RuntimeCheck(b_zeros.size(2) == size_n, "b_zeros dim 2 = ", b_zeros.size(2), " is not size_n = ", size_n); RuntimeCheck( num_groups == b_zeros.size(1), "b_zeros dim 1 = ", b_zeros.size(1), " is not num_groups = ", num_groups); RuntimeCheck(num_groups != -1, "num_groups must be != -1"); } else { RuntimeCheck( b_zeros.size(1) == num_groups, "b_zeros dim 1 = ", b_zeros.size(1), " is not num_groups = ", num_groups); RuntimeCheck( b_zeros.size(2) == size_n / pack_factor, "b_zeros dim 2 = ", b_zeros.size(2), " is not size_n / pack_factor = ", size_n / pack_factor); } } // Verify workspace size RuntimeCheck( size_n % device::marlin::min_thread_n == 0, "size_n = ", size_n, ", is not divisible by min_thread_n = ", device::marlin::min_thread_n); int64_t max_n_tiles = size_n / device::marlin::min_thread_n; int64_t min_workspace_size = std::min(max_n_tiles * (sorted_token_ids.size(0) / moe_block_size), static_cast(sms) * 4); RuntimeCheck( workspace.size(0) >= min_workspace_size, "workspace.numel = ", workspace.size(0), " is below min_workspace_size = ", min_workspace_size); // Early return for zero-size M (moved after all validation) if (size_m == 0) return; device::marlin_moe::marlin_mm( a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(), b_bias.data_ptr(), b_scales.data_ptr(), global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr(), sorted_token_ids.data_ptr(), expert_ids.data_ptr(), num_tokens_post_padded.data_ptr(), topk_weights.data_ptr(), static_cast(moe_block_size), static_cast(top_k), mul_topk_weights, is_ep, static_cast(size_m), static_cast(size_n), static_cast(size_k), workspace.data_ptr(), b_q_type, has_bias, has_act_order, is_k_full, has_zp, static_cast(num_groups), static_cast(group_size), dev, stream, thread_k, thread_n, sms, use_atomic_add, use_fp32_reduce, is_zp_float); } ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh ================================================ #include #include #include #include #include #include "nvfp4_quant.cuh" #include #include using namespace host; // Quantizes the provided PackedVec into the uint32_t output template SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec& vec, float SFScaleVal, uint8_t* SFout) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) // Get absolute maximum values among the local 8 values. auto localMax = __habs2(vec.elts[0]); // Local maximum value. #pragma unroll for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { localMax = __hmax2(localMax, __habs2(vec.elts[i])); } // Get the absolute maximum among all 16 values (two threads). localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax); // Get the final absolute maximum values. float vecMax = float(__hmax(localMax.x, localMax.y)); // Get the SF (max value of the vector / max value of e2m1). // maximum value of e2m1 = 6.0. // TODO: use half as compute data type. float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f)); // 8 bits representation of the SF. uint8_t fp8SFVal; // Write the SF to global memory (STG.8). if constexpr (UE8M0_SF) { // Extract the 8 exponent bits from float32. // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits. uint32_t tmp = reinterpret_cast(SFValue) >> 23; fp8SFVal = tmp & 0xff; // Convert back to fp32. reinterpret_cast(SFValue) = tmp << 23; } else { // Here SFValue is always positive, so E4M3 is the same as UE4M3. __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue); reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp; // Convert back to fp32. SFValue = float(tmp); } // Get the output scale. // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) * // reciprocal(SFScaleVal)) float outputScale = SFValue != 0 ? reciprocal_approximate_ftz(SFValue * reciprocal_approximate_ftz(SFScaleVal)) : 0.0f; if (SFout) { // Write the SF to global memory (STG.8). *SFout = fp8SFVal; } // Convert the input to float. float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2]; #pragma unroll for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { fp2Vals[i] = device::cast(vec.elts[i]); fp2Vals[i].x *= outputScale; fp2Vals[i].y *= outputScale; } // Convert to e2m1 values. uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals); // Write the e2m1 values to global memory. return e2m1Vec; #else return 0; #endif } SGL_DEVICE float silu(const float& val) { return val / (1.0f + __expf(-val)); } template SGL_DEVICE void silu_and_mul(PackedVec& x_vec, const PackedVec& y_vec) { float2 x[CVT_FP4_ELTS_PER_THREAD / 2]; float2 y[CVT_FP4_ELTS_PER_THREAD / 2]; #pragma unroll for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { x[i] = device::cast(x_vec.elts[i]); y[i] = device::cast(y_vec.elts[i]); x[i].x = silu(x[i].x) * y[i].x; x[i].y = silu(x[i].y) * y[i].y; x_vec.elts[i] = device::cast>(x[i]); } } // Use UE4M3 by default. template __global__ void #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) __launch_bounds__(512, 4) cvt_fp16_to_fp4( #else cvt_fp16_to_fp4( #endif int32_t numRows, int32_t numCols, Type const* in, float const* SFScale, uint32_t* out, uint32_t* SFout, uint32_t* input_offset_by_experts, uint32_t* output_scale_offset_by_experts, int32_t* mask, int n_experts, bool low_latency) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) using PackedVec = PackedVec; static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD); static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched."); // Input tensor row/col loops. int tid = blockIdx.x * blockDim.x + threadIdx.x; int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; // TODO(kaixih@nvidia): For now, we assume mask is used together with // silu_and_mal. Maybe we want a more general behavior of mask later. In the // silu case, the input last dim doubles. bool use_mask = mask != nullptr; int actualColsPerRow = use_mask ? colsPerRow * 2 : colsPerRow; // Each global thread processes one element for (int globalIdx = tid; globalIdx < numRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) { // Calculate which row and column this global thread should process int rowIdx = globalIdx / colsPerRow; int colIdx = globalIdx % colsPerRow; // Find index within the experts using different strategies based on expert // count int rowIdx_in_expert = 0; int expert_idx = 0; if constexpr (SMALL_NUM_EXPERTS) { for (int i = 0; i < n_experts; i++) { uint32_t current_offset = __ldca(&input_offset_by_experts[i]); uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]); if (rowIdx >= current_offset && rowIdx < next_offset) { rowIdx_in_expert = rowIdx - current_offset; expert_idx = i; break; } } } else { // Load input offsets into registers first, then do the computation. // Local array size set to 17 because of register limit. uint32_t local_offsets[17]; for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) { *reinterpret_cast(local_offsets) = __ldca(reinterpret_cast(&input_offset_by_experts[chunk_start])); *reinterpret_cast(local_offsets + 4) = __ldca(reinterpret_cast(&input_offset_by_experts[chunk_start + 4])); *reinterpret_cast(local_offsets + 8) = __ldca(reinterpret_cast(&input_offset_by_experts[chunk_start + 8])); *reinterpret_cast(local_offsets + 12) = __ldca(reinterpret_cast(&input_offset_by_experts[chunk_start + 12])); local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]); // Check against the 16 loaded offsets #pragma unroll for (int i = 0; i < 16; i++) { if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) { rowIdx_in_expert = rowIdx - local_offsets[i]; expert_idx = chunk_start + i; break; } } } } // Early exit when using masks. if (use_mask && rowIdx_in_expert >= mask[expert_idx]) { continue; } int64_t inOffset = rowIdx * actualColsPerRow + colIdx; PackedVec in_vec = reinterpret_cast(in)[inOffset]; if (use_mask) { PackedVec in_vec_mul = reinterpret_cast(in)[inOffset + colsPerRow]; silu_and_mul(in_vec, in_vec_mul); } // Get the output tensor offset. // Same as inOffset because 8 elements are packed into one uint32_t. int64_t outOffset = rowIdx * colsPerRow + colIdx; auto& out_pos = out[outOffset]; // Get the global scaling factor, which will be applied to the SF. // Note SFScale is the same as next GEMM's alpha, which is // (448.f / (Alpha_A / 6.f)). float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx]; int factor = CVT_FP4_SF_VEC_SIZE * 4; // The actual output_scales dim is computed from the padded numCols. int32_t numCols_padded = (numCols + factor - 1) / factor * factor; int numCols_SFout = numCols_padded / CVT_FP4_SF_VEC_SIZE / 4; uint32_t* SFout_in_expert = SFout + output_scale_offset_by_experts[expert_idx] * numCols_SFout; auto sf_out = cvt_quant_to_fp4_get_sf_out_offset( rowIdx_in_expert, colIdx, numCols, SFout_in_expert); out_pos = cvt_warp_fp16_to_fp4(in_vec, SFScaleVal, sf_out); } #endif } // Use UE4M3 by default. template __global__ void #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) __launch_bounds__(512, 4) cvt_fp16_to_fp4_expert( #else cvt_fp16_to_fp4_expert( #endif int32_t numRows, int32_t numCols, Type const* in, float const* SFScale, uint32_t* out, uint32_t* SFout, int32_t* mask, bool use_silu_and_mul, int n_experts) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) using PackedVec = PackedVec; static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD); static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched."); // Input tensor row/col loops. int tid = blockIdx.x * blockDim.x + threadIdx.x; int stride = (gridDim.x * blockDim.x) / n_experts; int remainder = (gridDim.x * blockDim.x) % n_experts; int expert_idx; int tid_in_expert; int actual_stride; if (remainder > 0) { int bound = remainder * (stride + 1); if (tid < bound) { expert_idx = tid / (stride + 1); tid_in_expert = tid % (stride + 1); actual_stride = stride + 1; } else { expert_idx = remainder + (tid - bound) / stride; tid_in_expert = (tid - bound) % stride; actual_stride = stride; } } else { expert_idx = tid / stride; tid_in_expert = tid % stride; actual_stride = stride; } int m = numRows / n_experts; int padded_m = (m + (128 - 1)) / 128 * 128; int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; // TODO(kaixih@nvidia): For now, we assume mask is used together with // silu_and_mal. Maybe we want a more general behavior of mask later. In the // silu case, the input last dim doubles. bool use_mask = mask != nullptr; int actualColsPerRow = use_silu_and_mul ? colsPerRow * 2 : colsPerRow; // Each global thread processes one element for (int globalIdx = tid_in_expert + expert_idx * m * colsPerRow; globalIdx < (expert_idx + 1) * m * colsPerRow; globalIdx += actual_stride) { // Calculate which row and column this global thread should process int rowIdx = globalIdx / colsPerRow; int colIdx = globalIdx % colsPerRow; // Find index within the experts int rowIdx_in_expert = rowIdx - expert_idx * m; // Early exit when using masks. if (use_mask && rowIdx_in_expert >= mask[expert_idx]) { break; } int64_t inOffset = rowIdx * actualColsPerRow + colIdx; PackedVec in_vec = reinterpret_cast(in)[inOffset]; if (use_silu_and_mul) { PackedVec in_vec_mul = reinterpret_cast(in)[inOffset + colsPerRow]; silu_and_mul(in_vec, in_vec_mul); } // Get the output tensor offset. // Same as inOffset because 8 elements are packed into one uint32_t. int64_t outOffset = rowIdx * colsPerRow + colIdx; auto& out_pos = out[outOffset]; // Get the global scaling factor, which will be applied to the SF. // Note SFScale is the same as next GEMM's alpha, which is // (448.f / (Alpha_A / 6.f)). float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx]; int factor = CVT_FP4_SF_VEC_SIZE * 4; // The actual output_scales dim is computed from the padded numCols. int32_t numCols_padded = (numCols + factor - 1) / factor * factor; int numCols_SFout = numCols_padded / CVT_FP4_SF_VEC_SIZE / 4; uint32_t* SFout_in_expert = SFout + expert_idx * padded_m * numCols_SFout; auto sf_out = cvt_quant_to_fp4_get_sf_out_offset( rowIdx_in_expert, colIdx, numCols, SFout_in_expert); out_pos = cvt_warp_fp16_to_fp4(in_vec, SFScaleVal, sf_out); } #endif } // Kernel for LARGE_M_TOPK = true (large m_topk optimized version) template __global__ void #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) __launch_bounds__(1024, 4) cvt_fp16_to_fp4( #else cvt_fp16_to_fp4( #endif int32_t numRows, int32_t numCols, Type const* in, float const* SFScale, uint32_t* out, uint32_t* SFout, uint32_t* input_offset_by_experts, uint32_t* output_scale_offset_by_experts, int32_t* mask, int n_experts) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) using PackedVec = PackedVec; static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD); static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched."); extern __shared__ uint32_t shared_input_offsets[]; // Load input offsets into shared memory. // If n_experts is larger than 4, use vectorized int4 to save instructions. // If n_experts is smaller than 4, read directly. if constexpr (SMALL_NUM_EXPERTS) { for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) { shared_input_offsets[i] = input_offset_by_experts[i]; } } else { for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) { *reinterpret_cast(&shared_input_offsets[i]) = *reinterpret_cast(&input_offset_by_experts[i]); } if (threadIdx.x == 0) { shared_input_offsets[n_experts] = input_offset_by_experts[n_experts]; } } __syncthreads(); int tid = blockIdx.x * blockDim.x + threadIdx.x; int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; bool use_mask = mask != nullptr; int actualColsPerRow = use_mask ? colsPerRow * 2 : colsPerRow; // Each global thread processes one element for (int globalIdx = tid; globalIdx < numRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) { // Calculate which row and column this global thread should process int rowIdx = globalIdx / colsPerRow; int colIdx = globalIdx % colsPerRow; // Find expert using binary search for better performance with large m_topk int rowIdx_in_expert = 0; int expert_idx = 0; // Binary search through experts using shared memory int left = 0, right = n_experts - 1; while (left <= right) { int mid = (left + right) / 2; // Get offsets: shared_input_offsets[i] corresponds to // input_offset_by_experts[i] uint32_t mid_offset = shared_input_offsets[mid]; uint32_t next_offset = shared_input_offsets[mid + 1]; if (rowIdx >= mid_offset && rowIdx < next_offset) { rowIdx_in_expert = rowIdx - mid_offset; expert_idx = mid; break; } else if (rowIdx < mid_offset) { right = mid - 1; } else { left = mid + 1; } } if (use_mask && rowIdx_in_expert >= mask[expert_idx]) { continue; } int64_t inOffset = rowIdx * actualColsPerRow + colIdx; PackedVec in_vec = reinterpret_cast(in)[inOffset]; if (use_mask) { PackedVec in_vec_mul = reinterpret_cast(in)[inOffset + colsPerRow]; silu_and_mul(in_vec, in_vec_mul); } int64_t outOffset = rowIdx * colsPerRow + colIdx; auto& out_pos = out[outOffset]; float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx]; int factor = CVT_FP4_SF_VEC_SIZE * 4; int32_t numCols_padded = (numCols + factor - 1) / factor * factor; int numCols_SFout = numCols_padded / CVT_FP4_SF_VEC_SIZE / 4; uint32_t* SFout_in_expert = SFout + output_scale_offset_by_experts[expert_idx] * numCols_SFout; auto sf_out = cvt_quant_to_fp4_get_sf_out_offset( rowIdx_in_expert, colIdx, numCols, SFout_in_expert); out_pos = cvt_warp_fp16_to_fp4(in_vec, SFScaleVal, sf_out); } #endif } template void quant_impl( void* output, void* output_scale, void* input, void* input_global_scale, void* input_offset_by_experts, void* output_scale_offset_by_experts, void* mask, bool use_silu_and_mul, int m_topk, int k, int n_experts, cudaStream_t stream) { // TODO: this multiProcessorCount should be cached. int device; cudaGetDevice(&device); int multiProcessorCount; cudaDeviceGetAttribute(&multiProcessorCount, cudaDevAttrMultiProcessorCount, device); // Grid, Block size. // Each thread converts 8 values. int const workSizePerRow = k / ELTS_PER_THREAD; int const totalWorkSize = m_topk * workSizePerRow; dim3 block(std::min(workSizePerRow, 512)); // Get number of blocks per SM (assume we can fully utilize the SM). int const numBlocksPerSM = 2048 / block.x; dim3 grid(std::min(static_cast((totalWorkSize + block.x - 1) / block.x), multiProcessorCount * numBlocksPerSM)); while (grid.x <= multiProcessorCount && block.x > 64) { grid.x *= 2; block.x = (block.x + 1) / 2; } // TODO(kaixih@nvidia): Should relax this to allow any grid size. if (mask != nullptr) { grid.x = (grid.x + n_experts - 1) / n_experts * n_experts; cvt_fp16_to_fp4_expert<<>>( m_topk, k, reinterpret_cast(input), reinterpret_cast(input_global_scale), reinterpret_cast(output), reinterpret_cast(output_scale), reinterpret_cast(mask), use_silu_and_mul, n_experts); return; } int const blockRepeat = (totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x); if (blockRepeat > 1) { size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t); if (n_experts >= 4) { cvt_fp16_to_fp4<<>>( m_topk, k, reinterpret_cast(input), reinterpret_cast(input_global_scale), reinterpret_cast(output), reinterpret_cast(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(output_scale_offset_by_experts), reinterpret_cast(mask), n_experts); } else { cvt_fp16_to_fp4<<>>( m_topk, k, reinterpret_cast(input), reinterpret_cast(input_global_scale), reinterpret_cast(output), reinterpret_cast(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(output_scale_offset_by_experts), reinterpret_cast(mask), n_experts); } } else { if (n_experts >= 16) { cvt_fp16_to_fp4<<>>( m_topk, k, reinterpret_cast(input), reinterpret_cast(input_global_scale), reinterpret_cast(output), reinterpret_cast(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(output_scale_offset_by_experts), reinterpret_cast(mask), n_experts, /* bool low_latency */ true); } else { cvt_fp16_to_fp4<<>>( m_topk, k, reinterpret_cast(input), reinterpret_cast(input_global_scale), reinterpret_cast(output), reinterpret_cast(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(output_scale_offset_by_experts), reinterpret_cast(mask), n_experts, /* bool low_latency */ true); } } } inline int getSMVersion(int device_id) { int sm_major = 0; int sm_minor = 0; RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id)); RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id)); return sm_major * 10 + sm_minor; } void scaled_fp4_experts_quant_sm100a( tvm::ffi::TensorView output, tvm::ffi::TensorView output_scale, tvm::ffi::TensorView input, tvm::ffi::TensorView input_global_scale, tvm::ffi::TensorView input_offset_by_experts, tvm::ffi::TensorView output_scale_offset_by_experts) { auto MTopK = SymbolicSize{"m_topk"}; auto K = SymbolicSize{"k"}; auto OutputCols = SymbolicSize{"output_cols"}; auto OutputScaleRows = SymbolicSize{"output_scale_rows"}; auto OutputScaleCols = SymbolicSize{"output_scale_cols"}; auto NExperts = SymbolicSize{"n_experts"}; auto OffsetSize = SymbolicSize{"offset_size"}; auto device = SymbolicDevice{}; TensorMatcher({MTopK, K}) // .with_dtype() .template with_device(device) .verify(input); TensorMatcher({MTopK, OutputCols}) // .with_dtype() .with_device(device) .verify(output); TensorMatcher({OutputScaleRows, OutputScaleCols}) // .with_dtype() .with_device(device) .verify(output_scale); TensorMatcher({NExperts}) // .with_dtype() .with_device(device) .verify(input_global_scale); TensorMatcher({OffsetSize}) // .with_dtype() .with_device(device) .verify(input_offset_by_experts) .verify(output_scale_offset_by_experts); const int device_id = input.device().device_id; RuntimeCheck(getSMVersion(device_id) >= 100, "fp4_quant is only supported on sm100+"); const int BLOCK_SIZE = 16; const auto m_topk = static_cast(MTopK.unwrap()); const auto k = static_cast(K.unwrap()); RuntimeCheck(k % BLOCK_SIZE == 0, "k must be a multiple of 16"); const auto n_experts = static_cast(NExperts.unwrap()); const auto offset_size = static_cast(OffsetSize.unwrap()); RuntimeCheck(offset_size == n_experts + 1, "input/output offset size mismatch"); RuntimeCheck(static_cast(OutputCols.unwrap()) == k / 2, "output second dim mismatch"); const int scales_k = k / BLOCK_SIZE; const int padded_k = (scales_k + 3) / 4 * 4; RuntimeCheck(static_cast(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch"); const cudaStream_t stream = LaunchKernel::resolve_device(input.device()); if (host::is_type(input.dtype())) { quant_impl( output.data_ptr(), output_scale.data_ptr(), input.data_ptr(), input_global_scale.data_ptr(), input_offset_by_experts.data_ptr(), output_scale_offset_by_experts.data_ptr(), nullptr, // mask false, // use_silu_and_mul m_topk, k, n_experts, stream); } else { quant_impl<__nv_bfloat16>( output.data_ptr(), output_scale.data_ptr(), input.data_ptr(), input_global_scale.data_ptr(), input_offset_by_experts.data_ptr(), output_scale_offset_by_experts.data_ptr(), nullptr, // mask false, // use_silu_and_mul m_topk, k, n_experts, stream); } } void silu_and_mul_scaled_fp4_experts_quant_sm100a( tvm::ffi::TensorView output, tvm::ffi::TensorView output_scale, tvm::ffi::TensorView input, tvm::ffi::TensorView input_global_scale, tvm::ffi::TensorView mask, bool use_silu_and_mul) { auto MTopK = SymbolicSize{"m_topk"}; auto KBy2 = SymbolicSize{"k_by_2"}; auto OutputCols = SymbolicSize{"output_cols"}; auto OutputScaleRows = SymbolicSize{"output_scale_rows"}; auto OutputScaleCols = SymbolicSize{"output_scale_cols"}; auto NExperts = SymbolicSize{"n_experts"}; auto device = SymbolicDevice{}; TensorMatcher({MTopK, KBy2}) // .with_dtype() .template with_device(device) .verify(input); TensorMatcher({MTopK, OutputCols}) // .with_dtype() .with_device(device) .verify(output); TensorMatcher({OutputScaleRows, OutputScaleCols}) // .with_dtype() .with_device(device) .verify(output_scale); TensorMatcher({NExperts}) // .with_dtype() .with_device(device) .verify(input_global_scale); TensorMatcher({NExperts}) // .with_dtype() .with_device(device) .verify(mask); const int device_id = input.device().device_id; RuntimeCheck(getSMVersion(device_id) >= 100, "fp4_quant is only supported on sm100+"); const int BLOCK_SIZE = 16; const auto m_topk = static_cast(MTopK.unwrap()); const auto k_by_2 = static_cast(KBy2.unwrap()); int k = k_by_2; if (use_silu_and_mul) { RuntimeCheck(k_by_2 % 2 == 0, "k must be a multiple of 2"); k = k_by_2 / 2; } const auto n_experts = static_cast(NExperts.unwrap()); RuntimeCheck(static_cast(OutputCols.unwrap()) == k / 2, "output second dim mismatch"); const int scales_k = k / BLOCK_SIZE; const int padded_k = (scales_k + 3) / 4 * 4; RuntimeCheck(static_cast(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch"); const cudaStream_t stream = LaunchKernel::resolve_device(input.device()); if (host::is_type(input.dtype())) { quant_impl( output.data_ptr(), output_scale.data_ptr(), input.data_ptr(), input_global_scale.data_ptr(), nullptr, // input_offset_by_experts nullptr, // output_scale_offset_by_experts mask.data_ptr(), use_silu_and_mul, m_topk, k, n_experts, stream); } else { quant_impl<__nv_bfloat16>( output.data_ptr(), output_scale.data_ptr(), input.data_ptr(), input_global_scale.data_ptr(), nullptr, // input_offset_by_experts nullptr, // output_scale_offset_by_experts mask.data_ptr(), use_silu_and_mul, m_topk, k, n_experts, stream); } } ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant.cuh ================================================ /* Copyright 2025 SGLang Team. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include #include #include #include #include #define ELTS_PER_THREAD 8 constexpr int CVT_FP4_ELTS_PER_THREAD = 8; constexpr int CVT_FP4_SF_VEC_SIZE = 16; // Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t). SGL_DEVICE uint32_t fp32_vec_to_e2m1(float (&array)[8]) { // PTX instructions used here requires >= sm100f. #if CUTLASS_ARCH_MMA_SM100A_ENABLED || CUTLASS_ARCH_MMA_SM103A_ENABLED || CUTLASS_ARCH_MMA_SM120A_ENABLED || \ (defined(__CUDA_ARCH_FAMILY_SPECIFIC__) && (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000)) uint32_t val; asm volatile( "{\n" ".reg .b8 byte0;\n" ".reg .b8 byte1;\n" ".reg .b8 byte2;\n" ".reg .b8 byte3;\n" "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" "}" : "=r"(val) : "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]), "f"(array[4]), "f"(array[5]), "f"(array[6]), "f"(array[7])); return val; #else printf("fp32_vec_to_e2m1 is not supported on this architecture\n"); __trap(); return 0; #endif } // Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t). SGL_DEVICE uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) { // PTX instructions used here requires >= sm100f. #if CUTLASS_ARCH_MMA_SM100A_ENABLED || CUTLASS_ARCH_MMA_SM103A_ENABLED || CUTLASS_ARCH_MMA_SM120A_ENABLED || \ (defined(__CUDA_ARCH_FAMILY_SPECIFIC__) && (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000)) uint32_t val; asm volatile( "{\n" ".reg .b8 byte0;\n" ".reg .b8 byte1;\n" ".reg .b8 byte2;\n" ".reg .b8 byte3;\n" "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" "}" : "=r"(val) : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y), "f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y)); return val; #else printf("fp32_vec_to_e2m1 is not supported on this architecture\n"); __trap(); return 0; #endif } // Fast reciprocal. SGL_DEVICE float reciprocal_approximate_ftz(float a) { float b; asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); return b; } template SGL_DEVICE uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx, int numCols, SFType* SFout) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 || CVT_FP4_NUM_THREADS_PER_SF == 2); // One pair of threads write one SF to global memory. // TODO: stage through smem for packed STG.32 // is it better than STG.8 from 4 threads ? if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) { // SF vector index (16 elements share one SF in the K dimension). int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF; int32_t mIdx = rowIdx; // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)] // --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx] int32_t mTileIdx = mIdx / (32 * 4); // SF vector size 16. int factor = CVT_FP4_SF_VEC_SIZE * 4; int32_t numKTiles = (numCols + factor - 1) / factor; int64_t mTileStride = numKTiles * 32 * 4 * 4; int32_t kTileIdx = (kIdx / 4); int64_t kTileStride = 32 * 4 * 4; // M tile layout [32, 4] is column-major. int32_t outerMIdx = (mIdx % 32); int64_t outerMStride = 4 * 4; int32_t innerMIdx = (mIdx % (32 * 4)) / 32; int64_t innerMStride = 4; int32_t innerKIdx = (kIdx % 4); int64_t innerKStride = 1; // Compute the global offset. int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride + outerMIdx * outerMStride + innerMIdx * innerMStride + innerKIdx * innerKStride; return reinterpret_cast(SFout) + SFOffset; } #endif return nullptr; } // Define a 16 bytes packed data type. template struct PackedVec { packed_t elts[4]; }; template <> struct PackedVec<__nv_fp8_e4m3> { __nv_fp8x2_e4m3 elts[8]; }; ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh ================================================ /* Copyright 2025 SGLang Team. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include #include void scaled_fp4_quant_sm100a_sm120a( tvm::ffi::TensorView output, tvm::ffi::TensorView input, tvm::ffi::TensorView output_sf, tvm::ffi::TensorView input_sf); void scaled_fp4_experts_quant_sm100a( tvm::ffi::TensorView output, tvm::ffi::TensorView output_scale, tvm::ffi::TensorView input, tvm::ffi::TensorView input_global_scale, tvm::ffi::TensorView input_offset_by_experts, tvm::ffi::TensorView output_scale_offset_by_experts); void silu_and_mul_scaled_fp4_experts_quant_sm100a( tvm::ffi::TensorView output, tvm::ffi::TensorView output_scale, tvm::ffi::TensorView input, tvm::ffi::TensorView input_global_scale, tvm::ffi::TensorView mask, bool use_silu_and_mul); void scaled_fp4_quant( tvm::ffi::TensorView output, tvm::ffi::TensorView input, tvm::ffi::TensorView output_sf, tvm::ffi::TensorView input_sf) { scaled_fp4_quant_sm100a_sm120a(output, input, output_sf, input_sf); } void scaled_fp4_experts_quant( tvm::ffi::TensorView output, tvm::ffi::TensorView output_scale, tvm::ffi::TensorView input, tvm::ffi::TensorView input_global_scale, tvm::ffi::TensorView input_offset_by_experts, tvm::ffi::TensorView output_scale_offset_by_experts) { scaled_fp4_experts_quant_sm100a( output, output_scale, input, input_global_scale, input_offset_by_experts, output_scale_offset_by_experts); } void silu_and_mul_scaled_fp4_experts_quant( tvm::ffi::TensorView output, tvm::ffi::TensorView output_scale, tvm::ffi::TensorView input, tvm::ffi::TensorView input_global_scale, tvm::ffi::TensorView mask, bool use_silu_and_mul) { silu_and_mul_scaled_fp4_experts_quant_sm100a(output, output_scale, input, input_global_scale, mask, use_silu_and_mul); } ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_kernels.cuh ================================================ /* Copyright 2025 SGLang Team. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include #include #include #include #include "nvfp4_quant.cuh" #include #include using namespace host; // Quantizes the provided PackedVec into the uint32_t output template SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec& vec, float SFScaleVal, uint8_t* SFout) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) // Get absolute maximum values among the local 8 values. auto localMax = __habs2(vec.elts[0]); // Local maximum value. #pragma unroll for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { localMax = __hmax2(localMax, __habs2(vec.elts[i])); } // Get the absolute maximum among all 16 values (two threads). localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax); // Get the final absolute maximum values. float vecMax = float(__hmax(localMax.x, localMax.y)); // Get the SF (max value of the vector / max value of e2m1). // maximum value of e2m1 = 6.0. // TODO: use half as compute data type. float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f)); // 8 bits representation of the SF. uint8_t fp8SFVal; // Write the SF to global memory (STG.8). if constexpr (UE8M0_SF) { __nv_fp8_e8m0 tmp; tmp.__x = __nv_cvt_float_to_e8m0(SFValue, __NV_SATFINITE, cudaRoundPosInf); SFValue = static_cast(tmp); fp8SFVal = tmp.__x; } else { // Here SFValue is always positive, so E4M3 is the same as UE4M3. __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue); fp8SFVal = tmp.__x; SFValue = static_cast(tmp); } // Get the output scale. // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) * // reciprocal(SFScaleVal)) float outputScale = SFValue != 0 ? reciprocal_approximate_ftz(SFValue * reciprocal_approximate_ftz(SFScaleVal)) : 0.0f; if (SFout) { // Write the SF to global memory (STG.8). *SFout = fp8SFVal; } // Convert the input to float. float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2]; #pragma unroll for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { if constexpr (std::is_same_v) { fp2Vals[i] = __half22float2(vec.elts[i]); } else { fp2Vals[i] = __bfloat1622float2(vec.elts[i]); } fp2Vals[i].x *= outputScale; fp2Vals[i].y *= outputScale; } // Convert to e2m1 values. uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals); // Write the e2m1 values to global memory. return e2m1Vec; #else return 0; #endif } // Use UE4M3 by default. template __global__ void #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) __launch_bounds__(512, 4) cvt_fp16_to_fp4( #else cvt_fp16_to_fp4( #endif int32_t numRows, int32_t numCols, Type const* in, float const* SFScale, uint32_t* out, uint32_t* SFout) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) using PackedVec = PackedVec; static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD); static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched."); // Get the global scaling factor, which will be applied to the SF. // Note SFScale is the same as next GEMM's alpha, which is // (448.f / (Alpha_A / 6.f)). float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[0]; // Input tensor row/col loops. for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) { for (int colIdx = threadIdx.x; colIdx < numCols / CVT_FP4_ELTS_PER_THREAD; colIdx += blockDim.x) { int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx; PackedVec in_vec = reinterpret_cast(in)[inOffset]; // Get the output tensor offset. // Same as inOffset because 8 elements are packed into one uint32_t. int64_t outOffset = inOffset; auto& out_pos = out[outOffset]; auto sf_out = cvt_quant_to_fp4_get_sf_out_offset(rowIdx, colIdx, numCols, SFout); out_pos = cvt_warp_fp16_to_fp4(in_vec, SFScaleVal, sf_out); } } #endif } template void invokeFP4Quantization( int m, int n, T const* input, float const* SFScale, int64_t* output, int32_t* SFOuput, bool useUE8M0, int multiProcessorCount, cudaStream_t stream) { // Grid, Block size. // Each thread converts 8 values. dim3 block(std::min(int(n / ELTS_PER_THREAD), 512)); // Get number of blocks per SM (assume we can fully utilize the SM). int const numBlocksPerSM = 2048 / block.x; dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM)); // Launch the cvt kernel. if (useUE8M0) { cvt_fp16_to_fp4<<>>( m, n, input, SFScale, reinterpret_cast(output), reinterpret_cast(SFOuput)); } else { cvt_fp16_to_fp4<<>>( m, n, input, SFScale, reinterpret_cast(output), reinterpret_cast(SFOuput)); } } // Instantiate the function. template void invokeFP4Quantization( int m, int n, half const* input, float const* SFScale, int64_t* output, int32_t* SFOuput, bool useUE8M0, int multiProcessorCount, cudaStream_t stream); template void invokeFP4Quantization( int m, int n, __nv_bfloat16 const* input, float const* SFScale, int64_t* output, int32_t* SFOuput, bool useUE8M0, int multiProcessorCount, cudaStream_t stream); inline int getSMVersion(int device_id) { int sm_major = 0; int sm_minor = 0; RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id)); RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id)); return sm_major * 10 + sm_minor; } void scaled_fp4_quant_sm100a_sm120a( tvm::ffi::TensorView output, tvm::ffi::TensorView input, tvm::ffi::TensorView output_sf, tvm::ffi::TensorView input_sf) { RuntimeCheck(input.device().device_type == kDLCUDA, "input must be a CUDA tensor"); RuntimeCheck(output.device() == input.device(), "output and input must be on same device"); RuntimeCheck(output_sf.device() == input.device(), "output_sf and input must be on same device"); RuntimeCheck(input_sf.device() == input.device(), "input_sf and input must be on same device"); RuntimeCheck(input.dim() == 2, "input must be a 2D tensor"); RuntimeCheck(output.dim() == 2, "output must be a 2D tensor"); RuntimeCheck(output_sf.dim() == 2, "output_sf must be a 2D tensor"); RuntimeCheck(input_sf.numel() == 1, "input_sf must have exactly one element"); RuntimeCheck(host::is_type(output.dtype()), "output must be uint8"); RuntimeCheck(host::is_type(output_sf.dtype()), "output_sf must be int32"); RuntimeCheck(host::is_type(input_sf.dtype()), "input_sf must be float32"); RuntimeCheck( host::is_type(input.dtype()) || host::is_type(input.dtype()), "input dtype must be fp16 or bf16"); const int device_id = input.device().device_id; const auto sm_version = getSMVersion(device_id); RuntimeCheck(sm_version >= 100, "fp4_quant is only supported on sm100+"); const int32_t m = static_cast(input.size(0)); const int32_t n = static_cast(input.size(1)); RuntimeCheck(output.size(0) == m, "output row size mismatch"); RuntimeCheck(output.size(1) == n / 2, "output column size mismatch"); RuntimeCheck(n % 16 == 0, "The N dimension must be multiple of 16."); const int multiProcessorCount = static_cast(runtime::get_sm_count(device_id)); auto input_sf_ptr = static_cast(input_sf.data_ptr()); auto sf_out = static_cast(output_sf.data_ptr()); auto output_ptr = static_cast(output.data_ptr()); const cudaStream_t stream = LaunchKernel::resolve_device(input.device()); constexpr bool useUE8M0 = false; if (host::is_type(input.dtype())) { auto input_ptr = reinterpret_cast(input.data_ptr()); invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out, useUE8M0, multiProcessorCount, stream); } else { auto input_ptr = reinterpret_cast<__nv_bfloat16 const*>(input.data_ptr()); invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out, useUE8M0, multiProcessorCount, stream); } } ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_entry.cuh ================================================ /* Copyright 2025 SGLang Team. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include void cutlass_scaled_fp4_mm_sm100a_sm120a( tvm::ffi::TensorView D, tvm::ffi::TensorView A, tvm::ffi::TensorView B, tvm::ffi::TensorView A_sf, tvm::ffi::TensorView B_sf, tvm::ffi::TensorView alpha); void cutlass_scaled_fp4_mm( tvm::ffi::TensorView D, tvm::ffi::TensorView A, tvm::ffi::TensorView B, tvm::ffi::TensorView A_sf, tvm::ffi::TensorView B_sf, tvm::ffi::TensorView alpha) { cutlass_scaled_fp4_mm_sm100a_sm120a(D, A, B, A_sf, B_sf, alpha); } ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_kernels.cuh ================================================ /* Copyright 2025 SGLang Team. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include #include #include #include #include #include #include #include using namespace host; // clang-format off #include "cutlass/cutlass.h" #include "cutlass/gemm/collective/collective_builder.hpp" #include "cutlass/epilogue/collective/collective_builder.hpp" #include "cutlass/gemm/device/gemm_universal_adapter.h" #include "cutlass/gemm/kernel/gemm_universal.hpp" #include "cutlass/util/packed_stride.hpp" // clang-format on /** * Helper function for checking CUTLASS errors */ #define CUTLASS_CHECK(status) \ { \ cutlass::Status error = status; \ RuntimeCheck(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \ } using namespace cute; // Helper function for next power of 2 inline uint32_t next_pow_2(uint32_t x) { if (x == 0) return 1; x--; x |= x >> 1; x |= x >> 2; x |= x >> 4; x |= x >> 8; x |= x >> 16; return x + 1; } struct WorkspaceKey { int device_id; uintptr_t stream; auto operator==(const WorkspaceKey&) const -> bool = default; }; struct WorkspaceKeyHash { auto operator()(const WorkspaceKey& key) const -> size_t { size_t h1 = std::hash{}(key.device_id); size_t h2 = std::hash{}(key.stream); return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2)); } }; struct WorkspaceState { void* ptr = nullptr; size_t bytes = 0; }; inline auto get_cached_workspace(size_t required_bytes, int device_id, cudaStream_t stream) -> void* { if (required_bytes == 0) { return nullptr; } thread_local std::unordered_map cache; WorkspaceKey key{device_id, reinterpret_cast(stream)}; auto& ws = cache[key]; if (ws.ptr != nullptr && ws.bytes >= required_bytes) { return ws.ptr; } RuntimeDeviceCheck(cudaSetDevice(device_id)); if (ws.ptr != nullptr) { RuntimeDeviceCheck(cudaFreeAsync(ws.ptr, stream)); ws.ptr = nullptr; ws.bytes = 0; } RuntimeDeviceCheck(cudaMallocAsync(&ws.ptr, required_bytes, stream)); ws.bytes = required_bytes; return ws.ptr; } #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || \ defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED) // Config(half_t/bfloat16_t) for M <= 128 template struct KernelConfigM128 { using OutputType = T; using MmaTileShape = Shape<_128, _256, _256>; using ClusterShape = Shape; using EpilogueTile = Shape<_128, _64>; // Avoid register spilling using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100; const static dim3 preferred_cluster; const static dim3 fallback_cluster; }; template const dim3 KernelConfigM128::preferred_cluster(1, 4, 1); template const dim3 KernelConfigM128::fallback_cluster(1, 2, 1); // Config(half_t/bfloat16_t) for M <= 256 template struct KernelConfigM256 { using OutputType = T; using MmaTileShape = Shape<_256, _256, _256>; using ClusterShape = Shape; using EpilogueTile = Shape<_128, _64>; // Avoid register spilling using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100; const static dim3 preferred_cluster; const static dim3 fallback_cluster; }; template const dim3 KernelConfigM256::preferred_cluster(2, 4, 1); template const dim3 KernelConfigM256::fallback_cluster(2, 1, 1); // Default config(half_t/bfloat16_t) for M > 256 template struct KernelConfigDefault { using OutputType = T; using MmaTileShape = Shape<_256, _256, _256>; using ClusterShape = Shape; using EpilogueTile = Shape<_128, _64>; // Avoid register spilling using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100; const static dim3 preferred_cluster; const static dim3 fallback_cluster; }; template const dim3 KernelConfigDefault::preferred_cluster(4, 4, 1); template const dim3 KernelConfigDefault::fallback_cluster(2, 1, 1); struct KernelConfigFp32 { using OutputType = float; using MmaTileShape = Shape<_128, _128, _256>; using ClusterShape = Shape; using EpilogueTile = cutlass::epilogue::collective::EpilogueTileAuto; using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100; const static dim3 preferred_cluster; const static dim3 fallback_cluster; }; const dim3 KernelConfigFp32::preferred_cluster = dim3(1, 4, 1); const dim3 KernelConfigFp32::fallback_cluster = dim3(1, 2, 1); // SM120 specific configurations struct sm120_fp4_config_M256 { using ClusterShape = Shape<_1, _1, _1>; using MmaTileShape = Shape<_128, _128, _128>; using PerSmTileShape_MNK = Shape<_128, _128, _128>; }; struct sm120_fp4_config_default { using ClusterShape = Shape<_1, _1, _1>; using MmaTileShape = Shape<_256, _128, _128>; using PerSmTileShape_MNK = Shape<_256, _128, _128>; }; template struct Fp4GemmSm100 { using Config = KernelConfig; // For generating args using OutputType = typename KernelConfig::OutputType; // A matrix configuration using ElementA = cutlass::nv_float4_t; using LayoutATag = cutlass::layout::RowMajor; static constexpr int AlignmentA = 32; // B matrix configuration using ElementB = cutlass::nv_float4_t; using LayoutBTag = cutlass::layout::ColumnMajor; static constexpr int AlignmentB = 32; // C/D matrix configuration using ElementD = OutputType; using ElementC = OutputType; using LayoutCTag = cutlass::layout::RowMajor; using LayoutDTag = cutlass::layout::RowMajor; static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; // Kernel functional config using ElementAccumulator = float; using ArchTag = cutlass::arch::Sm100; using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; // Kernel Perf config using MmaTileShape = typename KernelConfig::MmaTileShape; using ClusterShape = typename KernelConfig::ClusterShape; using EpilogueTile = typename KernelConfig::EpilogueTile; using EpilogueSchedule = typename KernelConfig::EpilogueSchedule; using MainloopSchedule = typename KernelConfig::MainloopSchedule; using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< ArchTag, OperatorClass, MmaTileShape, ClusterShape, EpilogueTile, ElementAccumulator, ElementAccumulator, void, LayoutCTag, AlignmentC, ElementD, LayoutDTag, AlignmentD, EpilogueSchedule, cutlass::epilogue::fusion::LinearCombination>::CollectiveOp; using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< ArchTag, OperatorClass, ElementA, LayoutATag, AlignmentA, ElementB, LayoutBTag, AlignmentB, ElementAccumulator, MmaTileShape, ClusterShape, cutlass::gemm::collective::StageCountAutoCarveout( sizeof(typename CollectiveEpilogue::SharedStorage))>, MainloopSchedule>::CollectiveOp; using GemmKernel = cutlass::gemm::kernel::GemmUniversal, CollectiveMainloop, CollectiveEpilogue, void>; using Gemm = cutlass::gemm::device::GemmUniversalAdapter; using StrideA = typename Gemm::GemmKernel::StrideA; using LayoutA = decltype(cute::make_layout(make_shape(0, 0, 0), StrideA{})); using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; using StrideB = typename Gemm::GemmKernel::StrideB; using LayoutB = decltype(cute::make_layout(make_shape(0, 0, 0), StrideB{})); using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; using StrideC = typename Gemm::GemmKernel::StrideC; using LayoutC = decltype(cute::make_layout(make_shape(0, 0, 0), StrideC{})); using StrideD = typename Gemm::GemmKernel::StrideD; using LayoutD = decltype(cute::make_layout(make_shape(0, 0, 0), StrideD{})); }; // SM120 specific GEMM template template struct Fp4GemmSm120 { using ElementA = cutlass::nv_float4_t; using LayoutATag = cutlass::layout::RowMajor; static constexpr int AlignmentA = 32; using ElementB = cutlass::nv_float4_t; using LayoutBTag = cutlass::layout::ColumnMajor; static constexpr int AlignmentB = 32; using ElementD = OutType; using ElementC = OutType; using LayoutCTag = cutlass::layout::RowMajor; using LayoutDTag = cutlass::layout::RowMajor; static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; using ElementAccumulator = float; using ArchTag = cutlass::arch::Sm120; using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; using MmaTileShape = typename Config::MmaTileShape; using ClusterShape = typename Config::ClusterShape; using PerSmTileShape_MNK = typename Config::PerSmTileShape_MNK; using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< ArchTag, OperatorClass, PerSmTileShape_MNK, ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, ElementAccumulator, ElementC, LayoutCTag, AlignmentC, ElementD, LayoutDTag, AlignmentD, cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< ArchTag, OperatorClass, ElementA, LayoutATag, AlignmentA, ElementB, LayoutBTag, AlignmentB, ElementAccumulator, MmaTileShape, ClusterShape, cutlass::gemm::collective::StageCountAutoCarveout( sizeof(typename CollectiveEpilogue::SharedStorage))>, cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp; using GemmKernel = cutlass::gemm::kernel::GemmUniversal, CollectiveMainloop, CollectiveEpilogue, void>; using Gemm = cutlass::gemm::device::GemmUniversalAdapter; }; template typename T::Gemm::Arguments args_from_options( tvm::ffi::TensorView D, tvm::ffi::TensorView A, tvm::ffi::TensorView B, tvm::ffi::TensorView A_sf, tvm::ffi::TensorView B_sf, tvm::ffi::TensorView alpha, int64_t M, int64_t N, int64_t K) { using ElementA = typename T::Gemm::ElementA; using ElementB = typename T::Gemm::ElementB; using ElementSFA = cutlass::float_ue4m3_t; using ElementSFB = cutlass::float_ue4m3_t; using ElementD = typename T::Gemm::ElementD; using ElementCompute = float; using StrideA = typename T::StrideA; using StrideB = typename T::StrideB; using StrideD = typename T::StrideD; using Sm1xxBlkScaledConfig = typename T::Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; int m = static_cast(M); int n = static_cast(N); int k = static_cast(K); auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, {m, k, 1}); auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, 1}); auto stride_D = cutlass::make_cute_packed_stride(StrideD{}, {m, n, 1}); auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(m, n, k, 1)); auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(m, n, k, 1)); typename T::Gemm::Arguments arguments{ cutlass::gemm::GemmUniversalMode::kGemm, {m, n, k, 1}, {// Mainloop arguments static_cast(A.data_ptr()), stride_A, static_cast(B.data_ptr()), stride_B, static_cast(A_sf.data_ptr()), layout_SFA, static_cast(B_sf.data_ptr()), layout_SFB}, { // Epilogue arguments {}, // epilogue.thread nullptr, stride_D, static_cast(D.data_ptr()), stride_D}}; auto& fusion_args = arguments.epilogue.thread; fusion_args.alpha_ptr = static_cast(alpha.data_ptr()); using KernelConfig = typename T::Config; arguments.hw_info.cluster_shape = KernelConfig::preferred_cluster; arguments.hw_info.cluster_shape_fallback = KernelConfig::fallback_cluster; return arguments; } template void runGemm( tvm::ffi::TensorView D, tvm::ffi::TensorView A, tvm::ffi::TensorView B, tvm::ffi::TensorView A_sf, tvm::ffi::TensorView B_sf, tvm::ffi::TensorView alpha, int64_t m, int64_t n, int64_t k, cudaStream_t stream) { typename T::Gemm gemm; auto arguments = args_from_options(D, A, B, A_sf, B_sf, alpha, m, n, k); size_t workspace_size = T::Gemm::get_workspace_size(arguments); int device_id = A.device().device_id; void* workspace = get_cached_workspace(workspace_size, device_id, stream); CUTLASS_CHECK(gemm.can_implement(arguments)); CUTLASS_CHECK(gemm.initialize(arguments, workspace, stream)); CUTLASS_CHECK(gemm.run(arguments, workspace, stream)); } // SM120 specific args_from_options function template typename Gemm::Arguments args_from_options_sm120( tvm::ffi::TensorView D, tvm::ffi::TensorView A, tvm::ffi::TensorView B, tvm::ffi::TensorView A_sf, tvm::ffi::TensorView B_sf, tvm::ffi::TensorView alpha, int M, int N, int K) { using ElementA = typename Gemm::ElementA; using ElementB = typename Gemm::ElementB; using ElementD = typename Gemm::ElementD; using ElementSFA = cutlass::float_ue4m3_t; using ElementSFB = cutlass::float_ue4m3_t; using ElementCompute = float; using StrideA = typename Gemm::GemmKernel::StrideA; using StrideB = typename Gemm::GemmKernel::StrideB; using StrideC = typename Gemm::GemmKernel::StrideC; using StrideD = typename Gemm::GemmKernel::StrideD; using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1}); auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1}); auto stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1}); auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1)); auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1)); typename Gemm::Arguments arguments{ cutlass::gemm::GemmUniversalMode::kGemm, {M, N, K, 1}, {static_cast(A.data_ptr()), stride_A, static_cast(B.data_ptr()), stride_B, static_cast(A_sf.data_ptr()), layout_SFA, static_cast(B_sf.data_ptr()), layout_SFB}, {{}, static_cast(D.data_ptr()), stride_D, static_cast(D.data_ptr()), stride_D}}; auto& fusion_args = arguments.epilogue.thread; fusion_args.alpha_ptr = static_cast(alpha.data_ptr()); return arguments; } // SM120 specific runGemm function template void runGemmSm120( tvm::ffi::TensorView D, tvm::ffi::TensorView A, tvm::ffi::TensorView B, tvm::ffi::TensorView A_sf, tvm::ffi::TensorView B_sf, tvm::ffi::TensorView alpha, int M, int N, int K, cudaStream_t stream) { Gemm gemm; auto arguments = args_from_options_sm120(D, A, B, A_sf, B_sf, alpha, M, N, K); size_t workspace_size = Gemm::get_workspace_size(arguments); int device_id = A.device().device_id; void* workspace = get_cached_workspace(workspace_size, device_id, stream); CUTLASS_CHECK(gemm.can_implement(arguments)); CUTLASS_CHECK(gemm.initialize(arguments, workspace, stream)); CUTLASS_CHECK(gemm.run(arguments, workspace, stream)); } // Dispatch function to select appropriate config based on M template void cutlassFp4GemmDispatch( tvm::ffi::TensorView D, tvm::ffi::TensorView A, tvm::ffi::TensorView B, tvm::ffi::TensorView A_sf, tvm::ffi::TensorView B_sf, tvm::ffi::TensorView alpha, int64_t m, int64_t n, int64_t k, cudaStream_t stream) { if (m <= 128) { // m in [1, 128] runGemm>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); } else if (m <= 256) { // m in (128, 256] runGemm>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); } else { // m in (256, inf) runGemm>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); } } // Dispatch function to select appropriate config based on M template <> void cutlassFp4GemmDispatch( tvm::ffi::TensorView D, tvm::ffi::TensorView A, tvm::ffi::TensorView B, tvm::ffi::TensorView A_sf, tvm::ffi::TensorView B_sf, tvm::ffi::TensorView alpha, int64_t m, int64_t n, int64_t k, cudaStream_t stream) { runGemm>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); } // SM120 specific dispatch functions void cutlass_fp4_bf16_gemm_dispatch_sm120( tvm::ffi::TensorView D, tvm::ffi::TensorView A, tvm::ffi::TensorView B, tvm::ffi::TensorView A_sf, tvm::ffi::TensorView B_sf, tvm::ffi::TensorView alpha, int m, int n, int k, cudaStream_t stream) { uint32_t const mp2 = std::max(static_cast(16), next_pow_2(m)); if (mp2 <= 256) { runGemmSm120::Gemm>( D, A, B, A_sf, B_sf, alpha, m, n, k, stream); } else { runGemmSm120::Gemm>( D, A, B, A_sf, B_sf, alpha, m, n, k, stream); } } void cutlass_fp4_f16_gemm_dispatch_sm120( tvm::ffi::TensorView D, tvm::ffi::TensorView A, tvm::ffi::TensorView B, tvm::ffi::TensorView A_sf, tvm::ffi::TensorView B_sf, tvm::ffi::TensorView alpha, int m, int n, int k, cudaStream_t stream) { uint32_t const mp2 = std::max(static_cast(16), next_pow_2(m)); if (mp2 <= 256) { runGemmSm120::Gemm>( D, A, B, A_sf, B_sf, alpha, m, n, k, stream); } else { runGemmSm120::Gemm>( D, A, B, A_sf, B_sf, alpha, m, n, k, stream); } } #else template void cutlassFp4GemmDispatch( tvm::ffi::TensorView D, tvm::ffi::TensorView A, tvm::ffi::TensorView B, tvm::ffi::TensorView A_sf, tvm::ffi::TensorView B_sf, tvm::ffi::TensorView alpha, int64_t m, int64_t n, int64_t k, cudaStream_t stream) { RuntimeCheck( false, "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to " "a CUTLASS 3.8 source directory to enable support."); } #endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || // defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED) inline int getSMVersion(int device_id) { int sm_major = 0; int sm_minor = 0; RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id)); RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id)); return sm_major * 10 + sm_minor; } void cutlass_scaled_fp4_mm_sm100a_sm120a( tvm::ffi::TensorView D, tvm::ffi::TensorView A, tvm::ffi::TensorView B, tvm::ffi::TensorView A_sf, tvm::ffi::TensorView B_sf, tvm::ffi::TensorView alpha) { RuntimeCheck(A.device().device_type == kDLCUDA, "a must be a CUDA tensor"); RuntimeCheck(B.device().device_type == kDLCUDA, "b must be a CUDA tensor"); RuntimeCheck(A_sf.device().device_type == kDLCUDA, "scale_a must be a CUDA tensor"); RuntimeCheck(B_sf.device().device_type == kDLCUDA, "scale_b must be a CUDA tensor"); RuntimeCheck(alpha.device().device_type == kDLCUDA, "alpha must be a CUDA tensor"); RuntimeCheck(D.device().device_type == kDLCUDA, "out must be a CUDA tensor"); RuntimeCheck(A.device() == B.device(), "a and b must be on same device"); RuntimeCheck(A.device() == A_sf.device(), "a and scale_a must be on same device"); RuntimeCheck(A.device() == B_sf.device(), "a and scale_b must be on same device"); RuntimeCheck(A.device() == alpha.device(), "a and alpha must be on same device"); RuntimeCheck(A.device() == D.device(), "a and out must be on same device"); RuntimeCheck(A.is_contiguous(), "a must be contiguous"); RuntimeCheck(B.is_contiguous(), "b must be contiguous"); RuntimeCheck(A_sf.is_contiguous(), "scale_a must be contiguous"); RuntimeCheck(B_sf.is_contiguous(), "scale_b must be contiguous"); RuntimeCheck(alpha.is_contiguous(), "alpha must be contiguous"); RuntimeCheck(D.is_contiguous(), "out must be contiguous"); RuntimeCheck(host::is_type(A.dtype()), "a must be uint8"); RuntimeCheck(host::is_type(B.dtype()), "b must be uint8"); RuntimeCheck(host::is_type(A_sf.dtype()), "scale_a must be float8_e4m3fn"); RuntimeCheck(host::is_type(B_sf.dtype()), "scale_b must be float8_e4m3fn"); RuntimeCheck(host::is_type(alpha.dtype()), "alpha must be float32"); RuntimeCheck(A.dim() == 2, "a must be a matrix"); RuntimeCheck(B.dim() == 2, "b must be a matrix"); RuntimeCheck(A_sf.dim() == 2, "scale_a must be a matrix"); RuntimeCheck(B_sf.dim() == 2, "scale_b must be a matrix"); RuntimeCheck(alpha.numel() == 1, "alpha must have exactly one element"); RuntimeCheck( A.size(1) == B.size(1), "a and b shapes cannot be multiplied (", A.size(0), "x", A.size(1), " and ", B.size(0), "x", B.size(1), ")"); const auto m = static_cast(A.size(0)); const auto n = static_cast(B.size(0)); const auto k = static_cast(A.size(1) * 2); RuntimeCheck(D.dim() == 2, "out must be 2D"); RuntimeCheck(D.size(0) == m, "out first dim must equal m"); RuntimeCheck(D.size(1) == n, "out second dim must equal n"); constexpr int alignment = 32; RuntimeCheck(k % alignment == 0, "Expected k to be divisible by ", alignment, ", but got k: ", k); RuntimeCheck(n % alignment == 0, "Expected n to be divisible by ", alignment, ", but got n: ", n); auto round_up = [](int64_t x, int64_t y) { return (x + y - 1) / y * y; }; const int64_t rounded_m = round_up(m, 128); const int64_t rounded_n = round_up(n, 128); const int64_t rounded_k = round_up(k / 16, 4); RuntimeCheck( A_sf.size(1) == B_sf.size(1), "scale_a and scale_b shapes cannot be multiplied (", A_sf.size(0), "x", A_sf.size(1), " and ", B_sf.size(0), "x", B_sf.size(1), ")"); RuntimeCheck( A_sf.size(0) == rounded_m && A_sf.size(1) == rounded_k, "scale_a must be padded/swizzled to shape (", rounded_m, "x", rounded_k, "), got (", A_sf.size(0), "x", A_sf.size(1), ")"); RuntimeCheck( B_sf.size(0) == rounded_n && B_sf.size(1) == rounded_k, "scale_b must be padded/swizzled to shape (", rounded_n, "x", rounded_k, "), got (", B_sf.size(0), "x", B_sf.size(1), ")"); const cudaStream_t stream = LaunchKernel::resolve_device(A.device()); const int sm_version = getSMVersion(A.device().device_id); if (sm_version >= 120) { if (host::is_type(D.dtype())) { cutlass_fp4_f16_gemm_dispatch_sm120( D, A, B, A_sf, B_sf, alpha, static_cast(m), static_cast(n), static_cast(k), stream); } else if (host::is_type(D.dtype())) { cutlass_fp4_bf16_gemm_dispatch_sm120( D, A, B, A_sf, B_sf, alpha, static_cast(m), static_cast(n), static_cast(k), stream); } else { Panic("Unsupported output data type of nvfp4 mm sm120"); } } else { if (host::is_type(D.dtype())) { cutlassFp4GemmDispatch(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); } else if (host::is_type(D.dtype())) { cutlassFp4GemmDispatch(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); } else if (host::is_type(D.dtype())) { cutlassFp4GemmDispatch(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); } else { Panic("Unsupported output data type of nvfp4 mm"); } } } ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh ================================================ #include #include #include #include #include #include #include #include #include #include #include namespace { constexpr size_t kBlockSize = 256; // each warp will handle 512B data template __global__ void per_tensor_absmax_kernel(const T* __restrict__ input, float* __restrict__ output_s, const int64_t num_elements) { using namespace device; constexpr uint32_t VEC_SIZE = 16 / sizeof(T); const int64_t gid = blockIdx.x * blockDim.x + threadIdx.x; float max_value = 0.0f; if (gid * VEC_SIZE + VEC_SIZE <= num_elements) { using vec_t = AlignedVector; const auto gmem_in = tile::Memory::thread(); const auto input_vec = gmem_in.load(input, gid); #pragma unroll for (uint32_t i = 0; i < VEC_SIZE; ++i) { const float value = static_cast(input_vec[i]); max_value = math::max(max_value, math::abs(value)); } } else if (gid * VEC_SIZE < num_elements) { [[unlikely]]; // poorly aligned case, do not optimize const auto remainder = num_elements - gid * VEC_SIZE; for (uint32_t i = 0; i < remainder; ++i) { const float value = static_cast(input[gid * VEC_SIZE + i]); max_value = math::max(max_value, math::abs(value)); } } // reduce within block and then atomic reduce between blocks __shared__ float smem[kWarpThreads]; cta::reduce_max(max_value, smem); if (threadIdx.x == 0) { const auto max_value = smem[0]; atomic::max(output_s, max_value / math::FP8_E4M3_MAX); } } [[maybe_unused]] SGL_DEVICE float fp8_e4m3_clip(float val) { namespace math = device::math; return math::max(math::min(val, math::FP8_E4M3_MAX), -math::FP8_E4M3_MAX); } template __global__ void per_tensor_quant_fp8_kernel( const T* __restrict__ input, DST_DTYPE* __restrict__ output, const float* __restrict__ scale, const int64_t num_elements) { using namespace device; constexpr uint32_t VEC_SIZE = 16 / sizeof(T); const int64_t gid = blockIdx.x * blockDim.x + threadIdx.x; const float scale_val = 1.0f / (*scale); if (gid * VEC_SIZE + VEC_SIZE <= num_elements) { using input_vec_t = AlignedVector; using output_vec_t = AlignedVector; const auto gmem_in = tile::Memory::thread(); const auto gmem_out = tile::Memory::thread(); const auto input_vec = gmem_in.load(input, gid); output_vec_t output_vec; #pragma unroll for (uint32_t i = 0; i < VEC_SIZE; ++i) { const float value = fp8_e4m3_clip(static_cast(input_vec[i]) * scale_val); output_vec[i] = static_cast(value); } gmem_out.store(output, output_vec, gid); } else if (gid * VEC_SIZE < num_elements) { [[unlikely]]; // poorly aligned case, do not optimize const auto remainder = num_elements - gid * VEC_SIZE; for (uint32_t i = 0; i < remainder; ++i) { const float value = fp8_e4m3_clip(static_cast(input[gid * VEC_SIZE + i]) * scale_val); output[gid * VEC_SIZE + i] = static_cast(value); } } } template void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) { using namespace host; auto device = SymbolicDevice{}; auto N = SymbolicSize{"num_elements"}; device.set_options(); TensorMatcher({N}) // .with_dtype() .with_device(device) .verify(input); TensorMatcher({N}) // .with_dtype() .with_device(device) .verify(output_q); TensorMatcher({1}) // .with_dtype() .with_device(device) .verify(output_s); const auto num_elements = N.unwrap(); constexpr size_t kElementsPerBlock = kBlockSize * (16 / sizeof(DType)); const uint32_t num_blocks = div_ceil(num_elements, kElementsPerBlock); if constexpr (!kIsStatic) { LaunchKernel(num_blocks, kBlockSize, device.unwrap())( per_tensor_absmax_kernel, static_cast(input.data_ptr()), static_cast(output_s.data_ptr()), static_cast(num_elements)); } LaunchKernel(num_blocks, kBlockSize, device.unwrap())( per_tensor_quant_fp8_kernel, static_cast(input.data_ptr()), static_cast(output_q.data_ptr()), static_cast(output_s.data_ptr()), static_cast(num_elements)); } } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh ================================================ #include #include #include #include #include #include #include #include #include #include #include namespace { constexpr int kThreadsPerGroup = 16; __device__ __forceinline__ float GroupReduceMax(float val, const int tid) { unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff; val = fmaxf(val, __shfl_xor_sync(mask, val, 8)); val = fmaxf(val, __shfl_xor_sync(mask, val, 4)); val = fmaxf(val, __shfl_xor_sync(mask, val, 2)); val = fmaxf(val, __shfl_xor_sync(mask, val, 1)); return val; } template using scale_packed_t_t = std::conditional_t; template using scale_element_t_t = std::conditional_t; template __global__ void per_token_group_quant_8bit_kernel( const T* __restrict__ input, DST_DTYPE* __restrict__ output_q, scale_packed_t_t* __restrict__ output_s, const int group_size, const int num_groups, const int groups_per_block, const float eps, const float min_8bit, const float max_8bit, const int num_groups_per_row = 0, const int scale_stride = 0) { using namespace device; namespace math = device::math; (void)num_groups; const int local_group_id = static_cast(threadIdx.x / kThreadsPerGroup); const int lane_id = threadIdx.x % kThreadsPerGroup; const int64_t block_group_id = blockIdx.x * groups_per_block; const int64_t global_group_id = block_group_id + local_group_id; const int64_t block_group_offset = global_group_id * group_size; float local_absmax = eps; using scale_packed_t = scale_packed_t_t; using scale_element_t = scale_element_t_t; static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0); const T* group_input = input + block_group_offset; DST_DTYPE* group_output = static_cast(output_q) + block_group_offset; scale_element_t* scale_output = nullptr; if constexpr (kIsColumnMajor) { constexpr int kElemsPerPack = static_cast(sizeof(scale_packed_t) / sizeof(scale_element_t)); const int row_idx = global_group_id / num_groups_per_row; const int col_idx_unpacked = global_group_id % num_groups_per_row; const int col_idx = col_idx_unpacked / kElemsPerPack; const int pack_idx = col_idx_unpacked % kElemsPerPack; scale_output = reinterpret_cast(output_s) + (col_idx * scale_stride * kElemsPerPack + row_idx * kElemsPerPack + pack_idx); } else { static_assert(!kScaleUE8M0); scale_output = output_s + global_group_id; } constexpr uint32_t kVecSize = 16 / sizeof(T); using vec_t = AlignedVector; const auto gmem_in = tile::Memory::thread(); const int32_t num_vec_elems = group_size / kVecSize; for (int32_t i = lane_id; i < num_vec_elems; i += kThreadsPerGroup) { const vec_t input_vec = gmem_in.load(group_input, i); #pragma unroll for (uint32_t j = 0; j < kVecSize; ++j) { const float val = static_cast(input_vec[j]); local_absmax = math::max(local_absmax, math::abs(val)); } } local_absmax = GroupReduceMax(local_absmax, lane_id); float y_s = local_absmax / max_8bit; if constexpr (kScaleUE8M0) { y_s = exp2f(ceilf(log2f(math::max(y_s, 1e-10f)))); } scale_element_t y_s_quant; if constexpr (kScaleUE8M0) { y_s_quant = static_cast(((int)log2f(y_s)) + 127); } else { y_s_quant = y_s; } if (lane_id == 0) { *scale_output = y_s_quant; } for (int32_t i = lane_id; i < num_vec_elems; i += kThreadsPerGroup) { const vec_t input_vec = gmem_in.load(group_input, i); #pragma unroll for (uint32_t j = 0; j < kVecSize; ++j) { const float val = static_cast(input_vec[j]); const float q_val = math::min(math::max(val / y_s, min_8bit), max_8bit); group_output[i * kVecSize + j] = DST_DTYPE(q_val); } } } inline int compute_groups_per_block(int64_t num_groups) { if (num_groups % 16 == 0) return 16; if (num_groups % 8 == 0) return 8; if (num_groups % 4 == 0) return 4; if (num_groups % 2 == 0) return 2; return 1; } template void per_token_group_quant_8bit( tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s, int64_t group_size, double eps, double min_8bit, double max_8bit, bool scale_ue8m0) { using namespace host; auto device = SymbolicDevice{}; auto M = SymbolicSize{"num_tokens"}; auto K = SymbolicSize{"hidden_dim"}; device.set_options(); TensorMatcher({M, K}).with_dtype().with_device(device).verify(input); TensorMatcher({M, K}).with_dtype().with_device(device).verify(output_q); const auto num_tokens = M.unwrap(); const auto hidden_dim = K.unwrap(); const int64_t num_groups_per_row = hidden_dim / group_size; const int64_t num_groups = num_tokens * num_groups_per_row; const int groups_per_block = compute_groups_per_block(num_groups); const int num_blocks = num_groups / groups_per_block; const int num_threads = groups_per_block * kThreadsPerGroup; const bool is_column_major = output_s.stride(0) < output_s.stride(1); const int scale_stride = output_s.stride(1); const float feps = static_cast(eps); const float fmin8 = static_cast(min_8bit); const float fmax8 = static_cast(max_8bit); if (is_column_major) { if (scale_ue8m0) { LaunchKernel(num_blocks, num_threads, input.device())( per_token_group_quant_8bit_kernel, static_cast(input.data_ptr()), static_cast(output_q.data_ptr()), static_cast(output_s.data_ptr()), static_cast(group_size), static_cast(num_groups), static_cast(groups_per_block), feps, fmin8, fmax8, static_cast(num_groups_per_row), scale_stride); } else { LaunchKernel(num_blocks, num_threads, input.device())( per_token_group_quant_8bit_kernel, static_cast(input.data_ptr()), static_cast(output_q.data_ptr()), static_cast(output_s.data_ptr()), static_cast(group_size), static_cast(num_groups), static_cast(groups_per_block), feps, fmin8, fmax8, static_cast(num_groups_per_row), scale_stride); } } else { LaunchKernel(num_blocks, num_threads, input.device())( per_token_group_quant_8bit_kernel, static_cast(input.data_ptr()), static_cast(output_q.data_ptr()), static_cast(output_s.data_ptr()), static_cast(group_size), static_cast(num_groups), static_cast(groups_per_block), feps, fmin8, fmax8, 0, 0); } } } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/hicache.cuh ================================================ #include #include #include #include #include #include #include #include namespace device { namespace details { template inline constexpr auto get_mem_package() { if constexpr (kUnit == 16) { return uint4{}; } else if constexpr (kUnit == 8) { return uint2{}; } else if constexpr (kUnit == 4) { return uint1{}; } else { static_assert(kUnit == 16 || kUnit == 8 || kUnit == 4, "Unsupported memory package size"); } } template using PackageType = decltype(get_mem_package()); SGL_DEVICE uint1 load_nc(const uint1* __restrict__ src) { uint32_t tmp; asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src)); return uint1{tmp}; } SGL_DEVICE uint2 load_nc(const uint2* __restrict__ src) { uint32_t tmp0, tmp1; asm volatile("ld.global.L1::no_allocate.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src)); return uint2{tmp0, tmp1}; } SGL_DEVICE uint4 load_nc(const uint4* __restrict__ src) { uint32_t tmp0, tmp1, tmp2, tmp3; asm volatile("ld.global.L1::no_allocate.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src)); return uint4{tmp0, tmp1, tmp2, tmp3}; } SGL_DEVICE void store_nc(uint1* __restrict__ dst, const uint1& value) { uint32_t tmp = value.x; asm volatile("st.global.L1::no_allocate.b32 [%0],%1;" ::"l"(dst), "r"(tmp)); } SGL_DEVICE void store_nc(uint2* __restrict__ dst, const uint2& value) { uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; asm volatile("st.global.L1::no_allocate.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1)); } SGL_DEVICE void store_nc(uint4* __restrict__ dst, const uint4& value) { uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; uint32_t tmp2 = value.z; uint32_t tmp3 = value.w; asm volatile( "st.global.L1::no_allocate.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3)); } } // namespace details template SGL_DEVICE auto load_vec(const void* __restrict__ src) { static_assert(kBytes % 128 == 0, "kBytes must be multiple of 128 bytes"); static_assert(128 % kNumThreads == 0, "kNumThreads must divide 128 bytes"); constexpr uint32_t kLoopCount = kBytes / 128; using Package = details::PackageType<128 / kNumThreads>; using Storage = AlignedStorage; const auto src_packed = static_cast(src); const auto lane_id = threadIdx.x % kNumThreads; Storage vec; #pragma unroll kLoopCount for (uint32_t i = 0; i < kLoopCount; ++i) { const auto j = i * kNumThreads + lane_id; vec.data[i] = details::load_nc(&src_packed[j]); } return vec; } template SGL_DEVICE void store_vec(void* __restrict__ dst, const Storage& vec) { using Package = std::decay_t; constexpr uint32_t kBytesPerLoop = sizeof(Package) * kNumThreads; constexpr uint32_t kLoopCount = kBytes / kBytesPerLoop; static_assert(kBytes % kBytesPerLoop == 0, "Invalid Storage configuration"); const auto dst_packed = static_cast(dst); const auto lane_id = threadIdx.x % kNumThreads; #pragma unroll kLoopCount for (uint32_t i = 0; i < kLoopCount; ++i) { const auto j = i * kNumThreads + lane_id; details::store_nc(&dst_packed[j], vec.data[i]); } } } // namespace device namespace { #define SGL_HICACHE_KERNEL __global__ __launch_bounds__(kBlockSize, 1) struct HicacheKernelParams { void* __restrict__ k_cache_dst; void* __restrict__ v_cache_dst; const void* __restrict__ indices_dst; void* __restrict__ k_cache_src; void* __restrict__ v_cache_src; const void* __restrict__ indices_src; int64_t kv_cache_src_stride; int64_t kv_cache_dst_stride; uint32_t length; uint32_t num_layers = 0; // only used in all_layer transfer }; template SGL_HICACHE_KERNEL void hicache_transfer_per_layer(const __grid_constant__ HicacheKernelParams params) { using namespace device; static_assert(kBlockSize % kWarpThreads == 0); static_assert(kWarpThreads % kUnroll == 0); constexpr uint32_t kNumThreads = kWarpThreads / kUnroll; constexpr uint32_t kWorkersPerBlock = kBlockSize / kNumThreads; constexpr uint32_t kNumWorkers = kWorkersPerBlock * kBlockQuota; const auto& [ k_cache_dst, v_cache_dst, indices_dst, // dst k_cache_src, v_cache_src, indices_src, // src kv_cache_src_stride, kv_cache_dst_stride, length, _ // metadata ] = params; const uint32_t work_id = blockIdx.x * kWorkersPerBlock + threadIdx.x / kNumThreads; for (uint32_t i = work_id; i < length; i += kNumWorkers) { const auto pos_src = static_cast(indices_src)[i]; const auto pos_dst = static_cast(indices_dst)[i]; const auto src_k = pointer::offset(k_cache_src, pos_src * kv_cache_src_stride); const auto dst_k = pointer::offset(k_cache_dst, pos_dst * kv_cache_dst_stride); const auto src_v = pointer::offset(v_cache_src, pos_src * kv_cache_src_stride); const auto dst_v = pointer::offset(v_cache_dst, pos_dst * kv_cache_dst_stride); const auto vec_k = load_vec(src_k); const auto vec_v = load_vec(src_v); store_vec(dst_k, vec_k); store_vec(dst_v, vec_v); } } template SGL_HICACHE_KERNEL void hicache_transfer_all_layer(const __grid_constant__ HicacheKernelParams params) { using namespace device; using src_ptr_t = const void*; using dst_ptr_t = void*; static_assert(kBlockSize % kWarpThreads == 0); static_assert(kWarpThreads % kUnroll == 0); constexpr uint32_t kNumThreads = kWarpThreads / kUnroll; constexpr uint32_t kWorkersPerBlock = kBlockSize / kNumThreads; constexpr uint32_t kNumWorkers = kWorkersPerBlock * kBlockQuota; const auto& [ k_ptr_dst, v_ptr_dst, indices_dst, // dst k_ptr_src, v_ptr_src, indices_src, // src kv_cache_src_stride, kv_cache_dst_stride, length, num_layers // metadata ] = params; const uint32_t work_id = blockIdx.x * kWorkersPerBlock + threadIdx.x / kNumThreads; for (uint32_t i = work_id; i < length; i += kNumWorkers) { const auto pos_src = static_cast(indices_src)[i]; const auto pos_dst = static_cast(indices_dst)[i]; for (uint32_t layer = 0; layer < num_layers; ++layer) { const auto k_cache_src = static_cast(k_ptr_src)[layer]; const auto v_cache_src = static_cast(v_ptr_src)[layer]; const auto k_cache_dst = static_cast(k_ptr_dst)[layer]; const auto v_cache_dst = static_cast(v_ptr_dst)[layer]; const auto src_k = pointer::offset(k_cache_src, pos_src * kv_cache_src_stride); const auto dst_k = pointer::offset(k_cache_dst, pos_dst * kv_cache_dst_stride); const auto src_v = pointer::offset(v_cache_src, pos_src * kv_cache_src_stride); const auto dst_v = pointer::offset(v_cache_dst, pos_dst * kv_cache_dst_stride); const auto vec_k = load_vec(src_k); const auto vec_v = load_vec(src_v); store_vec(dst_k, vec_k); store_vec(dst_v, vec_v); } } } template struct HiCacheKernel { template static constexpr auto kernel_one = hicache_transfer_per_layer; template static constexpr auto kernel_all = hicache_transfer_all_layer; static void run_one( const tvm::ffi::TensorView k_cache_dst, const tvm::ffi::TensorView v_cache_dst, const tvm::ffi::TensorView indices_dst, const tvm::ffi::TensorView k_cache_src, const tvm::ffi::TensorView v_cache_src, const tvm::ffi::TensorView indices_src) { using namespace host; auto D = SymbolicSize{"head dimension"}; auto N = SymbolicSize{"src kv stride"}; auto M = SymbolicSize{"dst kv stride"}; auto L = SymbolicSize{"indices length"}; auto cache_dtype = SymbolicDType{}; auto indices_dtype = SymbolicDType{}; auto indices_device = SymbolicDevice{}; TensorMatcher({-1, D}) // .with_strides({N, 1}) .with_dtype(cache_dtype) .with_device() .verify(k_cache_src) .verify(v_cache_src); TensorMatcher({-1, D}) // .with_strides({M, 1}) .with_dtype(cache_dtype) .with_device() .verify(k_cache_dst) .verify(v_cache_dst); TensorMatcher({L}) // .with_dtype(indices_dtype) .with_device(indices_device) .verify(indices_src) .verify(indices_dst); // verify dimension match const auto dtype_size = dtype_bytes(cache_dtype.unwrap()); const auto element_bytes = D.unwrap() * dtype_size; RuntimeCheck(kElementSize == element_bytes, "HicacheKernel: cache dimension mismatch."); const auto k_cache_dst_ptr = k_cache_dst.data_ptr(); const auto v_cache_dst_ptr = v_cache_dst.data_ptr(); const auto k_cache_src_ptr = k_cache_src.data_ptr(); const auto v_cache_src_ptr = v_cache_src.data_ptr(); const auto indices_dst_ptr = indices_dst.data_ptr(); const auto indices_src_ptr = indices_src.data_ptr(); const auto length = static_cast(L.unwrap()); const auto kv_cache_src_stride = static_cast(N.unwrap() * dtype_size); const auto kv_cache_dst_stride = static_cast(M.unwrap() * dtype_size); const auto use_int32 = indices_dtype.unwrap().bits == 32; const auto device = indices_device.unwrap(); constexpr auto kWorkersPerBlock = kBlockSize / (device::kWarpThreads / kUnroll); const auto num_blocks = std::min(div_ceil(length, kWorkersPerBlock), kBlockQuota); const auto params = HicacheKernelParams{ .k_cache_dst = k_cache_dst_ptr, .v_cache_dst = v_cache_dst_ptr, .indices_dst = indices_dst_ptr, .k_cache_src = k_cache_src_ptr, .v_cache_src = v_cache_src_ptr, .indices_src = indices_src_ptr, .kv_cache_src_stride = kv_cache_src_stride, .kv_cache_dst_stride = kv_cache_dst_stride, .length = length, }; const auto kernel = use_int32 ? kernel_one : kernel_one; LaunchKernel(num_blocks, kBlockSize, device)(kernel, params); } static void run_all( const tvm::ffi::TensorView k_ptr_dst, const tvm::ffi::TensorView v_ptr_dst, const tvm::ffi::TensorView indices_dst, const tvm::ffi::TensorView k_ptr_src, const tvm::ffi::TensorView v_ptr_src, const tvm::ffi::TensorView indices_src, const int64_t kv_src_stride_bytes, const int64_t kv_dst_stride_bytes) { using namespace host; auto N = SymbolicSize{"num_layers"}; auto L = SymbolicSize{"indices length"}; auto dtype_ = SymbolicDType{}; auto device_ = SymbolicDevice{}; TensorMatcher({N}) // .with_dtype() .with_device(device_) .verify(k_ptr_src) .verify(v_ptr_src) .verify(k_ptr_dst) .verify(v_ptr_dst); TensorMatcher({L}) // .with_dtype(dtype_) .with_device(device_) .verify(indices_src) .verify(indices_dst); // verify dimension match const auto k_cache_dst_ptr = k_ptr_dst.data_ptr(); const auto v_cache_dst_ptr = v_ptr_dst.data_ptr(); const auto k_cache_src_ptr = k_ptr_src.data_ptr(); const auto v_cache_src_ptr = v_ptr_src.data_ptr(); const auto indices_dst_ptr = indices_dst.data_ptr(); const auto indices_src_ptr = indices_src.data_ptr(); const auto length = static_cast(L.unwrap()); const auto use_int32 = dtype_.unwrap().bits == 32; const auto device = device_.unwrap(); constexpr auto kWorkersPerBlock = kBlockSize / (device::kWarpThreads / kUnroll); const auto num_blocks = std::min(div_ceil(length, kWorkersPerBlock), kBlockQuota); const auto params = HicacheKernelParams{ .k_cache_dst = k_cache_dst_ptr, .v_cache_dst = v_cache_dst_ptr, .indices_dst = indices_dst_ptr, .k_cache_src = k_cache_src_ptr, .v_cache_src = v_cache_src_ptr, .indices_src = indices_src_ptr, .kv_cache_src_stride = kv_src_stride_bytes, .kv_cache_dst_stride = kv_dst_stride_bytes, .length = length, .num_layers = static_cast(N.unwrap()), }; const auto kernel = use_int32 ? kernel_all : kernel_all; LaunchKernel(num_blocks, kBlockSize, device)(kernel, params); } }; #undef SGL_HICACHE_KERNEL } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/lora/moe_lora_align_kernel.cu ================================================ // Temporarily adapted from https://github.com/vllm-project/vllm/blob/main/csrc/moe/moe_align_sum_kernels.cu, will // optimize in future refactor #include #include #include #include #include #include #ifndef WARP_SIZE #define WARP_SIZE 32 #endif #define CEILDIV(x, y) (((x) + (y) - 1) / (y)) namespace moe { template SGL_DEVICE void _moe_align_block_size( const scalar_t* __restrict__ topk_ids, int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids, int32_t* __restrict__ total_tokens_post_pad, int32_t* __restrict__ expert_map, int32_t num_experts, int32_t padded_num_experts, int32_t experts_per_warp, int32_t block_size, size_t numel, int32_t* __restrict__ cumsum, int32_t max_num_tokens_padded, int32_t max_num_m_blocks, int32_t model_offset, int32_t inactive_expert_id, int32_t topk_num, int32_t* token_mask, bool has_expert_map) { extern __shared__ int32_t shared_counts[]; // Compute input buffer offsets. Typically these will all be 0, except when // using Multi LoRA. int sorted_token_ids_offset = max_num_tokens_padded * model_offset; int expert_ids_offset = max_num_m_blocks * model_offset; int cumsum_offset = (num_experts + 1) * model_offset; // Use separate threadblocks to fill sorted_token_ids. // This is safe since the current kernel does not use sorted_token_ids. if (blockIdx.x % 2) { // Initialize sorted_token_ids with numel for (size_t it = threadIdx.x; it < max_num_tokens_padded; it += blockDim.x) { sorted_token_ids[sorted_token_ids_offset + it] = static_cast(numel); } return; } const int warp_id = threadIdx.x / WARP_SIZE; const int my_expert_start = warp_id * experts_per_warp; for (int i = 0; i < experts_per_warp; ++i) { if (my_expert_start + i < padded_num_experts) { shared_counts[warp_id * experts_per_warp + i] = 0; } } __syncthreads(); const size_t tid = threadIdx.x; const size_t stride = blockDim.x; for (size_t i = tid; i < numel; i += stride) { int expert_id = topk_ids[i]; if (expert_id < 0 || expert_id >= num_experts) { continue; } if (has_expert_map) { expert_id = expert_map[expert_id]; if (expert_id < 0 || expert_id >= num_experts) continue; } int warp_idx = expert_id / experts_per_warp; int expert_offset = expert_id % experts_per_warp; int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], mask); } __syncthreads(); // Compute prefix sum over token counts per expert using BlockScan = cub::BlockScan; __shared__ typename BlockScan::TempStorage temp_storage; int expert_count = 0; int expert_id = threadIdx.x; if (expert_id < num_experts) { int warp_idx = expert_id / experts_per_warp; int expert_offset = expert_id % experts_per_warp; expert_count = shared_counts[warp_idx * experts_per_warp + expert_offset]; expert_count = CEILDIV(expert_count, block_size) * block_size; } int cumsum_val; BlockScan(temp_storage).ExclusiveSum(expert_count, cumsum_val); if (expert_id <= num_experts) { cumsum[cumsum_offset + expert_id] = cumsum_val; } if (expert_id == num_experts) { total_tokens_post_pad[model_offset] = cumsum_val; } __syncthreads(); if (threadIdx.x < num_experts) { for (int i = cumsum[cumsum_offset + threadIdx.x]; i < cumsum[cumsum_offset + threadIdx.x + 1]; i += block_size) { expert_ids[expert_ids_offset + i / block_size] = threadIdx.x; } } // Fill remaining expert_ids with 0 const size_t fill_start_idx = cumsum[cumsum_offset + num_experts] / block_size + threadIdx.x; for (size_t i = fill_start_idx; i < max_num_m_blocks; i += blockDim.x) { expert_ids[expert_ids_offset + i] = inactive_expert_id; } } template SGL_DEVICE void _moe_align_block_size_small_batch_expert( const scalar_t* __restrict__ topk_ids, int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids, int32_t* __restrict__ total_tokens_post_pad, int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size, size_t numel, int32_t max_num_tokens_padded, int32_t max_num_m_blocks, int32_t inactive_expert_id, int32_t model_offset, int32_t topk_num, int32_t* token_mask, bool has_expert_map) { // Compute input buffer offsets. Typically these will all be 0, except when // using Multi LoRA. int sorted_token_ids_offset = max_num_tokens_padded * model_offset; int expert_ids_offset = max_num_m_blocks * model_offset; // Use an additional group of threads to fill sorted_token_ids. // Since the current kernel will use sorted_token_ids afterward, // we fill sorted_token_ids within the same threadblock to make // synchronization easier. if (threadIdx.x < fill_threads) { // Initialize sorted_token_ids with numel for (size_t it = threadIdx.x; it < max_num_tokens_padded; it += fill_threads) { sorted_token_ids[sorted_token_ids_offset + it] = static_cast(numel); } // Three __syncthreads() corresponding to the other threads __syncthreads(); __syncthreads(); __syncthreads(); return; } const size_t tid = threadIdx.x - fill_threads; const size_t stride = blockDim.x - fill_threads; extern __shared__ int32_t shared_mem[]; int32_t* cumsum = shared_mem; int32_t* tokens_cnts = (int32_t*)(shared_mem + num_experts + 1); for (int i = 0; i < num_experts; ++i) { tokens_cnts[(tid + 1) * num_experts + i] = 0; } for (size_t i = tid; i < numel; i += stride) { int32_t expert_id = topk_ids[i]; if (expert_id < 0 || expert_id >= num_experts) continue; if (has_expert_map) { expert_id = expert_map[expert_id]; if (expert_id < 0 || expert_id >= num_experts) continue; } int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; tokens_cnts[(tid + 1) * num_experts + expert_id] += mask; } __syncthreads(); if (tid < num_experts) { tokens_cnts[tid] = 0; for (int i = 1; i <= stride; ++i) { tokens_cnts[i * num_experts + tid] += tokens_cnts[(i - 1) * num_experts + tid]; } } __syncthreads(); if (tid == 0) { cumsum[0] = 0; for (int i = 1; i <= num_experts; ++i) { cumsum[i] = cumsum[i - 1] + CEILDIV(tokens_cnts[stride * num_experts + i - 1], block_size) * block_size; } total_tokens_post_pad[model_offset] = static_cast(cumsum[num_experts]); } __syncthreads(); if (tid < num_experts) { for (int i = cumsum[tid]; i < cumsum[tid + 1]; i += block_size) { expert_ids[expert_ids_offset + i / block_size] = tid; } } // Fill remaining expert_ids with 0 const size_t fill_start_idx = cumsum[num_experts] / block_size + tid; for (size_t i = fill_start_idx; i < max_num_m_blocks; i += stride) { expert_ids[expert_ids_offset + i] = inactive_expert_id; } for (size_t i = tid; i < numel; i += stride) { int32_t expert_id = topk_ids[i]; if (expert_id < 0 || expert_id >= num_experts) continue; if (has_expert_map) { expert_id = expert_map[expert_id]; if (expert_id < 0 || expert_id >= num_experts) continue; } int32_t rank_post_pad = tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id]; if (token_mask == nullptr || token_mask[i / topk_num]) { sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i; ++tokens_cnts[tid * num_experts + expert_id]; } } } template SGL_DEVICE void _count_and_sort_expert_tokens( const scalar_t* __restrict__ topk_ids, int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer, int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts, int32_t max_num_tokens_padded, int32_t* __restrict__ token_mask, int32_t model_offset, int32_t topk_num, bool has_expert_map) { const size_t tid = blockIdx.y * blockDim.x + threadIdx.x; const size_t stride = blockDim.x * gridDim.y; for (size_t i = tid; i < numel; i += stride) { int32_t expert_id = topk_ids[i]; if (expert_id >= num_experts) { continue; } if (has_expert_map) { expert_id = expert_map[expert_id]; // filter invalid experts if (expert_id == -1) continue; } if (token_mask == nullptr || token_mask[i / topk_num]) { int32_t rank_post_pad = atomicAdd(&cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1); sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] = i; } } } template __global__ void moe_lora_align_block_size_kernel( scalar_t* __restrict__ topk_ids, int32_t* __restrict__ seg_indptr, int32_t* __restrict__ req_to_lora, int num_reqs, int64_t block_size, int32_t* __restrict__ expert_map, int num_experts, int max_loras, size_t numel, int max_num_tokens_padded, int max_num_m_blocks, int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids, int32_t topk_num, int32_t* total_tokens_post_pad, int32_t* adapter_enabled, int32_t* __restrict__ cumsum, int32_t experts_per_warp, int32_t padded_num_experts, int32_t* lora_ids, int32_t* __restrict__ token_mask, bool has_expert_map) { int lora_idx = blockIdx.x / 2; int lora_id = lora_ids[lora_idx]; if (lora_id == -1 || adapter_enabled[lora_id] == 0) { return; } int num_tokens = numel / topk_num; int lora_offset = lora_id * num_tokens; if (blockIdx.x % 2 == 0) { // 1. Parallel Clear (Reset mask to 0) for (int i = threadIdx.x; i < num_tokens; i += blockDim.x) { token_mask[lora_offset + i] = 0; } if (threadIdx.x == 0) { total_tokens_post_pad[lora_id] = 0; } __syncthreads(); // 2. Segment-based Fill for (int r = 0; r < num_reqs; ++r) { if (req_to_lora[r] == lora_id) { int start = seg_indptr[r]; int end = seg_indptr[r + 1]; for (int i = start + threadIdx.x; i < end; i += blockDim.x) { token_mask[lora_offset + i] = 1; } } } __syncthreads(); } _moe_align_block_size( topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map, num_experts, padded_num_experts, experts_per_warp, block_size, numel, cumsum, max_num_tokens_padded, max_num_m_blocks, lora_id, -1, // inactive_expert_id padding topk_num, &token_mask[(lora_id * num_tokens)], has_expert_map); } template __global__ void lora_count_and_sort_expert_tokens_kernel( const scalar_t* __restrict__ topk_ids, int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer, int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts, int32_t max_num_tokens_padded, int32_t topk_num, int32_t* token_mask, int32_t* lora_ids, int32_t* adapter_enabled, bool has_expert_map) { int lora_idx = blockIdx.x; int lora_id = lora_ids[lora_idx]; if (lora_id == -1 || adapter_enabled[lora_id] == 0) { return; } int num_tokens = numel / topk_num; _count_and_sort_expert_tokens( topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts, max_num_tokens_padded, &token_mask[(lora_id * num_tokens)], lora_id, topk_num, has_expert_map); } template __global__ void moe_lora_align_block_size_small_batch_expert_kernel( scalar_t* __restrict__ topk_ids, int32_t* __restrict__ seg_indptr, int32_t* __restrict__ req_to_lora, int num_reqs, int64_t block_size, int32_t* __restrict__ expert_map, int num_experts, int max_loras, size_t numel, int max_num_tokens_padded, int max_num_m_blocks, int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids, int topk_num, int32_t* total_tokens_post_pad, int32_t* adapter_enabled, int32_t* lora_ids, int32_t* token_mask, bool has_expert_map) { int lora_idx = blockIdx.x; int lora_id = lora_ids[lora_idx]; if (lora_id == -1 || adapter_enabled[lora_id] == 0) { return; } int num_tokens = numel / topk_num; int lora_offset = lora_id * num_tokens; // 1. Parallel Clear (Reset mask to 0) // All threads help clear the mask for this adapter for (int i = threadIdx.x; i < num_tokens; i += blockDim.x) { token_mask[lora_offset + i] = 0; } // Initialize output counter if (threadIdx.x == 0) { total_tokens_post_pad[lora_id] = 0; } __syncthreads(); // 2. Segment-based Fill // Iterate over requests. If a request matches this LoRA, fill its range. for (int r = 0; r < num_reqs; ++r) { if (req_to_lora[r] == lora_id) { int start = seg_indptr[r]; int end = seg_indptr[r + 1]; // Parallel Fill: All threads help mark this segment as "1" for (int i = start + threadIdx.x; i < end; i += blockDim.x) { token_mask[lora_offset + i] = 1; } } } __syncthreads(); _moe_align_block_size_small_batch_expert( topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map, num_experts, block_size, numel, max_num_tokens_padded, max_num_m_blocks, -1, // inactive_expert_id padding lora_id, topk_num, &token_mask[(lora_id * num_tokens)], has_expert_map); } } // namespace moe namespace { template struct MoeLoraAlignBlockSizeKernel { static void run(tvm::ffi::TensorView topk_ids, tvm::ffi::TensorView seg_indptr, tvm::ffi::TensorView req_to_lora, int64_t num_experts, int64_t block_size, int64_t max_loras, int64_t max_num_tokens_padded, int64_t max_num_m_blocks, tvm::ffi::TensorView sorted_token_ids, tvm::ffi::TensorView expert_ids, tvm::ffi::TensorView num_tokens_post_pad, tvm::ffi::TensorView adapter_enabled, tvm::ffi::TensorView lora_ids, tvm::ffi::Optional maybe_expert_map, tvm::ffi::TensorView cumsum_buffer, tvm::ffi::TensorView token_mask) { using namespace host; const int topk_num = topk_ids.size(1); RuntimeCheck(block_size > 0, "block_size should be greater than 0. "); int device_max_shared_mem; auto device = topk_ids.device(); int dev_id = device.device_id; RuntimeDeviceCheck(cudaDeviceGetAttribute(&device_max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev_id)); const cudaStream_t stream = LaunchKernel::resolve_device(device); int64_t padded_num_experts = ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; // BlockScan uses 1024 threads and assigns one thread per expert. RuntimeCheck(padded_num_experts < 1024, "padded_num_experts must be less than 1024"); int32_t* token_mask_ptr = static_cast(token_mask.data_ptr()); bool has_expert_map = maybe_expert_map.has_value(); int32_t* expert_map_ptr = nullptr; if (has_expert_map) { expert_map_ptr = static_cast(maybe_expert_map.value().data_ptr()); } int num_reqs = seg_indptr.size(0) - 1; bool small_batch_expert_mode = (topk_ids.numel() < 1024) && (num_experts <= 64); if (small_batch_expert_mode) { const int32_t num_thread = std::max((int32_t)num_experts, 128); const int32_t shared_mem = (num_thread + 1) * num_experts * sizeof(int32_t) + (num_experts + 1) * sizeof(int32_t); if (shared_mem > device_max_shared_mem) { RuntimeCheck(false, "Shared memory usage exceeds device limit."); } // threadIdx.x >= fill_threads: counting experts and aligning // threadIdx.x < fill_threads: filling sorted_token_ids constexpr int32_t fill_threads = 256; dim3 blockDim(num_thread + fill_threads); auto kernel = moe::moe_lora_align_block_size_small_batch_expert_kernel; RuntimeDeviceCheck(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_mem)); LaunchKernel(dim3(max_loras), blockDim, stream, shared_mem)( kernel, static_cast(topk_ids.data_ptr()), static_cast(seg_indptr.data_ptr()), static_cast(req_to_lora.data_ptr()), num_reqs, block_size, expert_map_ptr, num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded, max_num_m_blocks, static_cast(sorted_token_ids.data_ptr()), static_cast(expert_ids.data_ptr()), topk_num, static_cast(num_tokens_post_pad.data_ptr()), static_cast(adapter_enabled.data_ptr()), static_cast(lora_ids.data_ptr()), token_mask_ptr, has_expert_map); } else { int num_thread = 1024; dim3 blockDim(num_thread); size_t num_warps = CEILDIV(padded_num_experts, WARP_SIZE); size_t shared_mem_size = num_warps * WARP_SIZE * sizeof(int32_t); auto align_kernel = moe::moe_lora_align_block_size_kernel; // launch two threadblocks for each lora // blockIdx.x % 2 == 0: counting experts and aligning // blockIdx.x % 2 == 1: filling sorted_token_ids LaunchKernel(dim3(max_loras * 2), blockDim, stream, shared_mem_size)( align_kernel, static_cast(topk_ids.data_ptr()), static_cast(seg_indptr.data_ptr()), static_cast(req_to_lora.data_ptr()), num_reqs, block_size, expert_map_ptr, num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded, max_num_m_blocks, static_cast(sorted_token_ids.data_ptr()), static_cast(expert_ids.data_ptr()), topk_num, static_cast(num_tokens_post_pad.data_ptr()), static_cast(adapter_enabled.data_ptr()), static_cast(cumsum_buffer.data_ptr()), WARP_SIZE, padded_num_experts, static_cast(lora_ids.data_ptr()), token_mask_ptr, has_expert_map); const int block_threads = std::min(256, (int)num_thread); const int num_blocks = (topk_ids.numel() + block_threads - 1) / block_threads; const int max_blocks = 65535; const int actual_blocks = std::min(num_blocks, max_blocks); dim3 gridDims(max_loras, actual_blocks); auto sort_kernel = moe::lora_count_and_sort_expert_tokens_kernel; LaunchKernel(gridDims, dim3(block_threads), stream)( sort_kernel, static_cast(topk_ids.data_ptr()), static_cast(sorted_token_ids.data_ptr()), static_cast(cumsum_buffer.data_ptr()), expert_map_ptr, topk_ids.numel(), num_experts, max_num_tokens_padded, topk_num, token_mask_ptr, static_cast(lora_ids.data_ptr()), static_cast(adapter_enabled.data_ptr()), has_expert_map); } } }; } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/moe/nvfp4_blockwise_moe.cuh ================================================ #include #include #include #include #include #include #include "cute/tensor.hpp" #include "cutlass/epilogue/collective/collective_builder.hpp" #include "cutlass/epilogue/collective/default_epilogue.hpp" #include "cutlass/epilogue/thread/linear_combination.h" #include "cutlass/gemm/collective/collective_builder.hpp" #include "cutlass/gemm/device/gemm_universal_adapter.h" #include "cutlass/gemm/dispatch_policy.hpp" #include "cutlass/gemm/group_array_problem_shape.hpp" #include "cutlass/gemm/kernel/gemm_universal.hpp" #include "cutlass/tensor_ref.h" #include "cutlass/util/command_line.h" #include "cutlass/util/distribution.h" #include "cutlass/util/host_tensor.h" #include "cutlass/util/packed_stride.hpp" #include "cutlass/util/reference/device/gemm.h" #include "cutlass/util/reference/device/tensor_compare.h" #include "cutlass/util/reference/host/gett.hpp" #include "cutlass/util/reference/host/tensor_compare.h" #include "cutlass/util/reference/host/tensor_fill.h" #include "cutlass/util/reference/host/tensor_norm.h" #include "cutlass/util/tensor_view_io.h" #include #include #include #include #include using namespace host; using namespace cute; struct WorkspaceKey { int device_id; uintptr_t stream; auto operator==(const WorkspaceKey&) const -> bool = default; }; struct WorkspaceKeyHash { auto operator()(const WorkspaceKey& key) const -> size_t { size_t h1 = std::hash{}(key.device_id); size_t h2 = std::hash{}(key.stream); return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2)); } }; struct WorkspaceState { void* ptr = nullptr; size_t bytes = 0; }; inline auto get_cached_workspace(size_t required_bytes, int device_id, cudaStream_t stream) -> void* { if (required_bytes == 0) { return nullptr; } thread_local std::unordered_map cache; WorkspaceKey key{device_id, reinterpret_cast(stream)}; auto& ws = cache[key]; if (ws.ptr != nullptr && ws.bytes >= required_bytes) { return ws.ptr; } RuntimeDeviceCheck(cudaSetDevice(device_id)); if (ws.ptr != nullptr) { RuntimeDeviceCheck(cudaFreeAsync(ws.ptr, stream)); ws.ptr = nullptr; ws.bytes = 0; } RuntimeDeviceCheck(cudaMallocAsync(&ws.ptr, required_bytes, stream)); ws.bytes = required_bytes; return ws.ptr; } inline int getSMVersion(int device_id) { int sm_major = 0; int sm_minor = 0; RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id)); RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id)); return sm_major * 10 + sm_minor; } template < typename ElementAB, typename ElementC, typename ElementSF, typename ElementAccumulator, typename LayoutSFA, typename LayoutSFB, typename ScaleConfig> __global__ void __get_group_gemm_starts( ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets, ElementSF** a_scales_offsets, ElementSF** b_scales_offsets, ElementAccumulator** alpha_offsets, LayoutSFA* layout_sfa_base_as_int, LayoutSFB* layout_sfb_base_as_int, ElementAB* a_base_as_int, ElementAB* b_base_as_int, ElementC* out_base_as_int, ElementSF* a_scales_base_as_int, ElementSF* b_scales_base_as_int, ElementAccumulator* alphas_base_as_int, const int32_t* expert_offsets, const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes, const int K, const int N) { int64_t expert_id = threadIdx.x; if (expert_id >= gridDim.x * blockDim.x) { return; } // Originally int32_t but upcasting to int64_t to avoid overflow // during offset calculations int64_t expert_offset = static_cast(expert_offsets[expert_id]); int64_t sf_offset = static_cast(sf_offsets[expert_id]); // size for block in block scale. int64_t group_size = 16; int64_t m = static_cast(problem_sizes_as_shapes[expert_id * 3]); int64_t n = static_cast(problem_sizes_as_shapes[expert_id * 3 + 1]); int64_t k = static_cast(problem_sizes_as_shapes[expert_id * 3 + 2]); assert((m >= 0 && n == N && k == K && k % 2 == 0) && "unexpected problem sizes"); int64_t half_k = static_cast(k / 2); int64_t group_k = static_cast(k / group_size); // Shape of A as uint8/byte = [M, K // 2] // Shape of B as uint8/byte = [E, N, K // 2] a_offsets[expert_id] = a_base_as_int + expert_offset * half_k; b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k; // Shape of C = [M, N] out_offsets[expert_id] = out_base_as_int + expert_offset * n; // Shape of a_scale = [sum(sf_sizes), K // group_size] a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k; assert((reinterpret_cast(a_scales_offsets[expert_id]) % 128) == 0 && "TMA requires 128-byte alignment"); // Shape of B scale = [E, N, K // group_size] b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k; assert((reinterpret_cast(b_scales_offsets[expert_id]) % 128) == 0 && "TMA requires 128-byte alignment"); // Shape of alpha = [E] alpha_offsets[expert_id] = alphas_base_as_int + expert_id; LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id; LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id; *layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA( cute::make_shape(static_cast(m), static_cast(n), static_cast(k), 1)); *layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB( cute::make_shape(static_cast(m), static_cast(n), static_cast(k), 1)); } #define __CALL_GET_STARTS_KERNEL_BLOCKSCALE( \ ELEMENT_AB_TYPE, SF_TYPE, TYPE_CHECK, C_TYPE, LayoutSFA, LayoutSFB, ScaleConfig) \ else if (TYPE_CHECK) { \ __get_group_gemm_starts \ <<<1, num_experts, 0, stream>>>( \ static_cast(a_starts.data_ptr()), \ static_cast(b_starts.data_ptr()), \ static_cast(out_starts.data_ptr()), \ static_cast(a_scales_starts.data_ptr()), \ static_cast(b_scales_starts.data_ptr()), \ static_cast(alpha_starts.data_ptr()), \ reinterpret_cast(layout_sfa.data_ptr()), \ reinterpret_cast(layout_sfb.data_ptr()), \ static_cast(a_tensors.data_ptr()), \ static_cast(b_tensors.data_ptr()), \ static_cast(out_tensors.data_ptr()), \ static_cast(a_scales.data_ptr()), \ static_cast(b_scales.data_ptr()), \ static_cast(alphas.data_ptr()), \ static_cast(expert_offsets.data_ptr()), \ static_cast(sf_offsets.data_ptr()), \ static_cast(problem_sizes.data_ptr()), \ K, \ N); \ } template void run_get_group_gemm_starts( const tvm::ffi::TensorView a_starts, const tvm::ffi::TensorView b_starts, const tvm::ffi::TensorView out_starts, const tvm::ffi::TensorView a_scales_starts, const tvm::ffi::TensorView b_scales_starts, const tvm::ffi::TensorView alpha_starts, const tvm::ffi::TensorView layout_sfa, const tvm::ffi::TensorView layout_sfb, /*these are used for their base addresses*/ tvm::ffi::TensorView const& a_tensors, tvm::ffi::TensorView const& b_tensors, tvm::ffi::TensorView const& out_tensors, tvm::ffi::TensorView const& a_scales, tvm::ffi::TensorView const& b_scales, tvm::ffi::TensorView const& alphas, tvm::ffi::TensorView const& expert_offsets, tvm::ffi::TensorView const& sf_offsets, tvm::ffi::TensorView const& problem_sizes, int M, int N, int K) { int num_experts = static_cast(expert_offsets.size(0)); auto stream = LaunchKernel::resolve_device(a_tensors.device()); RuntimeCheck(out_tensors.size(1) == N, "Output tensor shape doesn't match expected shape"); RuntimeCheck( K / 2 == b_tensors.size(2), "b_tensors(dim = 2) and a_tensors(dim = 1) trailing" " dimension must match"); if (false) { } //(ELEMENT_AB_TYPE, BS_TYPE, TENSOR_C_TYPE, C_TYPE, LayoutSFA, LayoutSFB, // ScaleConfig) __CALL_GET_STARTS_KERNEL_BLOCKSCALE( cutlass::float_e2m1_t, cutlass::float_ue4m3_t, host::is_type(out_tensors.dtype()), cutlass::bfloat16_t, LayoutSFA, LayoutSFB, ScaleConfig) __CALL_GET_STARTS_KERNEL_BLOCKSCALE( cutlass::float_e2m1_t, cutlass::float_ue4m3_t, host::is_type(out_tensors.dtype()), cutlass::half_t, LayoutSFA, LayoutSFB, ScaleConfig) else { Panic("Invalid output type (must be float16 or bfloat16)"); } } void run_fp4_blockwise_scaled_group_mm_sm120( tvm::ffi::TensorView output, const tvm::ffi::TensorView a, const tvm::ffi::TensorView b, const tvm::ffi::TensorView a_blockscale, const tvm::ffi::TensorView b_blockscales, const tvm::ffi::TensorView alphas, const tvm::ffi::TensorView ab_strides, const tvm::ffi::TensorView c_strides, const tvm::ffi::TensorView problem_sizes, const tvm::ffi::TensorView expert_offsets, const tvm::ffi::TensorView sf_offsets, const tvm::ffi::TensorView a_ptrs, const tvm::ffi::TensorView b_ptrs, const tvm::ffi::TensorView out_ptrs, const tvm::ffi::TensorView a_scales_ptrs, const tvm::ffi::TensorView b_scales_ptrs, const tvm::ffi::TensorView alpha_ptrs, const tvm::ffi::TensorView layout_sfa, const tvm::ffi::TensorView layout_sfb, int M, int N, int K) { using ProblemShape = cutlass::gemm::GroupProblemShape>; using ElementType = cutlass::float_e2m1_t; using ElementSFType = cutlass::float_ue4m3_t; using ElementA = cutlass::nv_float4_t; using ElementB = cutlass::nv_float4_t; using ElementC = cutlass::bfloat16_t; using ElementD = cutlass::bfloat16_t; using ElementAccumulator = float; // Layout definitions using LayoutA = cutlass::layout::RowMajor; using LayoutB = cutlass::layout::ColumnMajor; using LayoutC = cutlass::layout::RowMajor; using LayoutD = cutlass::layout::RowMajor; // Alignment constraints static constexpr int AlignmentA = 32; static constexpr int AlignmentB = 32; static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; // Architecture definitions using ArchTag = cutlass::arch::Sm120; using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; using StageCountType = cutlass::gemm::collective::StageCountAuto; using ThreadBlockShape = Shape<_128, _128, _128>; // on the tile size using ClusterShape = Shape<_1, _1, _1>; using FusionOperation = cutlass::epilogue::fusion::LinearCombination; using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< ArchTag, OperatorClass, ThreadBlockShape, ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD, LayoutC*, AlignmentD, cutlass::epilogue::collective::EpilogueScheduleAuto, FusionOperation>::CollectiveOp; using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< ArchTag, OperatorClass, ElementA, LayoutA*, AlignmentA, ElementB, LayoutB*, AlignmentB, ElementAccumulator, ThreadBlockShape, ClusterShape, cutlass::gemm::collective::StageCountAutoCarveout( sizeof(typename CollectiveEpilogue::SharedStorage))>, cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong>::CollectiveOp; using GemmKernel = cutlass::gemm::kernel::GemmUniversal; using Gemm1SM = cutlass::gemm::device::GemmUniversalAdapter; using Gemm = Gemm1SM; using StrideA = typename Gemm::GemmKernel::InternalStrideA; using StrideB = typename Gemm::GemmKernel::InternalStrideB; using StrideC = typename Gemm::GemmKernel::InternalStrideC; using StrideD = typename Gemm::GemmKernel::InternalStrideD; using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA; using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB; using ScaleConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape; int num_experts = static_cast(expert_offsets.size(0)); run_get_group_gemm_starts( a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs, layout_sfa, layout_sfb, a, b, output, a_blockscale, b_blockscales, alphas, expert_offsets, sf_offsets, problem_sizes, M, N, K); // Create an instance of the GEMM Gemm gemm_op; // Initialize problem_sizes_as_shapes correctly UnderlyingProblemShape* problem_sizes_as_shapes = static_cast(problem_sizes.data_ptr()); // Set the Scheduler info cutlass::KernelHardwareInfo hw_info; using RasterOrderOptions = cutlass::gemm::kernel::detail::RasterOrderOptions; typename Gemm::GemmKernel::TileSchedulerArguments scheduler; scheduler.raster_order = RasterOrderOptions::AlongM; hw_info.device_id = a.device().device_id; static std::unordered_map cached_sm_counts; if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) { cached_sm_counts[hw_info.device_id] = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); } hw_info.sm_count = std::min(cached_sm_counts[hw_info.device_id], std::numeric_limits::max()); // Mainloop Arguments typename GemmKernel::MainloopArguments mainloop_args{ static_cast(a_ptrs.data_ptr()), static_cast(ab_strides.data_ptr()), static_cast(b_ptrs.data_ptr()), static_cast(ab_strides.data_ptr()), static_cast(a_scales_ptrs.data_ptr()), reinterpret_cast(layout_sfa.data_ptr()), static_cast(b_scales_ptrs.data_ptr()), reinterpret_cast(layout_sfb.data_ptr())}; // Epilogue Arguments typename GemmKernel::EpilogueArguments epilogue_args{ {}, // epilogue.thread nullptr, static_cast(c_strides.data_ptr()), static_cast(out_ptrs.data_ptr()), static_cast(c_strides.data_ptr())}; auto& fusion_args = epilogue_args.thread; fusion_args.alpha_ptr_array = reinterpret_cast(alpha_ptrs.data_ptr()); fusion_args.dAlpha = {_0{}, _0{}, 1}; fusion_args.beta = 0.0f; // Gemm Arguments typename GemmKernel::Arguments args{ cutlass::gemm::GemmUniversalMode::kGrouped, {num_experts, problem_sizes_as_shapes, nullptr}, mainloop_args, epilogue_args, hw_info, scheduler}; size_t workspace_size = Gemm::get_workspace_size(args); const cudaStream_t stream = LaunchKernel::resolve_device(a.device()); void* workspace = get_cached_workspace(workspace_size, hw_info.device_id, stream); auto can_implement_status = gemm_op.can_implement(args); RuntimeCheck( can_implement_status == cutlass::Status::kSuccess, "Failed to implement GEMM: ", cutlassGetStatusString(can_implement_status)); // Run the GEMM auto status = gemm_op.initialize(args, workspace); RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to initialize GEMM: ", cutlassGetStatusString(status)); status = gemm_op.run(args, workspace, stream); RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to run GEMM: ", cutlassGetStatusString(status)); } template void run_fp4_blockwise_scaled_group_mm_sm100( tvm::ffi::TensorView output, const tvm::ffi::TensorView a, const tvm::ffi::TensorView b, const tvm::ffi::TensorView a_blockscale, const tvm::ffi::TensorView b_blockscales, const tvm::ffi::TensorView alphas, const tvm::ffi::TensorView ab_strides, const tvm::ffi::TensorView c_strides, const tvm::ffi::TensorView problem_sizes, const tvm::ffi::TensorView expert_offsets, const tvm::ffi::TensorView sf_offsets, const tvm::ffi::TensorView a_ptrs, const tvm::ffi::TensorView b_ptrs, const tvm::ffi::TensorView out_ptrs, const tvm::ffi::TensorView a_scales_ptrs, const tvm::ffi::TensorView b_scales_ptrs, const tvm::ffi::TensorView alpha_ptrs, const tvm::ffi::TensorView layout_sfa, const tvm::ffi::TensorView layout_sfb, int M, int N, int K) { using ProblemShape = cutlass::gemm::GroupProblemShape>; using ElementType = cutlass::float_e2m1_t; using ElementSFType = cutlass::float_ue4m3_t; using ElementA = cutlass::nv_float4_t; using ElementB = cutlass::nv_float4_t; using ElementC = OutType; using ElementD = ElementC; using ElementAccumulator = float; // Layout definitions using LayoutA = cutlass::layout::RowMajor; using LayoutB = cutlass::layout::ColumnMajor; using LayoutC = cutlass::layout::RowMajor; using LayoutD = LayoutC; // Alignment constraints static constexpr int AlignmentA = 32; static constexpr int AlignmentB = 32; static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; // Architecture definitions using ArchTag = cutlass::arch::Sm100; using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp; // Epilogue Operator class tag using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; // Mainloop Operator class tag using StageCountType = cutlass::gemm::collective::StageCountAuto; // Stage count maximized based // on the tile size using ClusterShape = Shape<_1, _1, _1>; struct MMA1SMConfig { using MmaTileShape = Shape<_128, _128, _128>; using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmNvf4Sm100; // Kernel to launch using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm; // Epilogue to launch }; using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape, ClusterShape, Shape<_128, _64>, ElementAccumulator, ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD, LayoutC*, AlignmentD, typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp; using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA, ElementB, LayoutB*, AlignmentB, ElementAccumulator, typename MMA1SMConfig::MmaTileShape, ClusterShape, cutlass::gemm::collective::StageCountAutoCarveout( sizeof(typename CollectiveEpilogue::SharedStorage))>, typename MMA1SMConfig::KernelSchedule>::CollectiveOp; using GemmKernel = cutlass::gemm::kernel::GemmUniversal; using Gemm1SM = cutlass::gemm::device::GemmUniversalAdapter; using Gemm = Gemm1SM; using StrideA = typename Gemm::GemmKernel::InternalStrideA; using StrideB = typename Gemm::GemmKernel::InternalStrideB; using StrideC = typename Gemm::GemmKernel::InternalStrideC; using StrideD = typename Gemm::GemmKernel::InternalStrideD; using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA; using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB; using ScaleConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape; int num_experts = static_cast(expert_offsets.size(0)); run_get_group_gemm_starts( a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs, layout_sfa, layout_sfb, a, b, output, a_blockscale, b_blockscales, alphas, expert_offsets, sf_offsets, problem_sizes, M, N, K); // Create an instance of the GEMM Gemm gemm_op; // Initialize problem_sizes_as_shapes correctly UnderlyingProblemShape* problem_sizes_as_shapes = static_cast(problem_sizes.data_ptr()); // Set the Scheduler info cutlass::KernelHardwareInfo hw_info; using RasterOrderOptions = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm100GroupParams< typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions; typename Gemm::GemmKernel::TileSchedulerArguments scheduler; scheduler.raster_order = RasterOrderOptions::AlongM; hw_info.device_id = a.device().device_id; static std::unordered_map cached_sm_counts; if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) { cached_sm_counts[hw_info.device_id] = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); } hw_info.sm_count = std::min(cached_sm_counts[hw_info.device_id], std::numeric_limits::max()); // Mainloop Arguments typename GemmKernel::MainloopArguments mainloop_args{ static_cast(a_ptrs.data_ptr()), static_cast(ab_strides.data_ptr()), static_cast(b_ptrs.data_ptr()), static_cast(ab_strides.data_ptr()), static_cast(a_scales_ptrs.data_ptr()), reinterpret_cast(layout_sfa.data_ptr()), static_cast(b_scales_ptrs.data_ptr()), reinterpret_cast(layout_sfb.data_ptr())}; // Epilogue Arguments typename GemmKernel::EpilogueArguments epilogue_args{ {}, // epilogue.thread nullptr, static_cast(c_strides.data_ptr()), static_cast(out_ptrs.data_ptr()), static_cast(c_strides.data_ptr())}; auto& fusion_args = epilogue_args.thread; fusion_args.alpha_ptr_array = reinterpret_cast(alpha_ptrs.data_ptr()); fusion_args.dAlpha = {_0{}, _0{}, 1}; // Gemm Arguments typename GemmKernel::Arguments args{ cutlass::gemm::GemmUniversalMode::kGrouped, {num_experts, problem_sizes_as_shapes, nullptr}, mainloop_args, epilogue_args, hw_info, scheduler}; size_t workspace_size = Gemm::get_workspace_size(args); const cudaStream_t stream = LaunchKernel::resolve_device(a.device()); void* workspace = get_cached_workspace(workspace_size, hw_info.device_id, stream); auto can_implement_status = gemm_op.can_implement(args); RuntimeCheck( can_implement_status == cutlass::Status::kSuccess, "Failed to implement GEMM: ", cutlassGetStatusString(can_implement_status)); // Run the GEMM auto status = gemm_op.initialize(args, workspace); RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to initialize GEMM: ", cutlassGetStatusString(status)); status = gemm_op.run(args, workspace, stream); RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to run GEMM: ", cutlassGetStatusString(status)); } void cutlass_fp4_group_mm_sm100a_sm120a( tvm::ffi::TensorView output, const tvm::ffi::TensorView a, const tvm::ffi::TensorView b, const tvm::ffi::TensorView a_blockscale, const tvm::ffi::TensorView b_blockscales, const tvm::ffi::TensorView alphas, const tvm::ffi::TensorView ab_strides, const tvm::ffi::TensorView c_strides, const tvm::ffi::TensorView problem_sizes, const tvm::ffi::TensorView expert_offsets, const tvm::ffi::TensorView sf_offsets, const tvm::ffi::TensorView a_ptrs, const tvm::ffi::TensorView b_ptrs, const tvm::ffi::TensorView out_ptrs, const tvm::ffi::TensorView a_scales_ptrs, const tvm::ffi::TensorView b_scales_ptrs, const tvm::ffi::TensorView alpha_ptrs, const tvm::ffi::TensorView layout_sfa, const tvm::ffi::TensorView layout_sfb) { auto check_cuda_contig = [](const tvm::ffi::TensorView t, const char* name) { RuntimeCheck(t.device().device_type == kDLCUDA, name, " must be a CUDA tensor"); RuntimeCheck(t.is_contiguous(), name, " must be contiguous"); }; check_cuda_contig(output, "output"); check_cuda_contig(a, "a"); check_cuda_contig(b, "b"); check_cuda_contig(a_blockscale, "a_blockscale"); check_cuda_contig(b_blockscales, "b_blockscales"); check_cuda_contig(alphas, "alphas"); check_cuda_contig(ab_strides, "ab_strides"); check_cuda_contig(c_strides, "c_strides"); check_cuda_contig(problem_sizes, "problem_sizes"); check_cuda_contig(expert_offsets, "expert_offsets"); check_cuda_contig(sf_offsets, "sf_offsets"); check_cuda_contig(a_ptrs, "a_ptrs"); check_cuda_contig(b_ptrs, "b_ptrs"); check_cuda_contig(out_ptrs, "out_ptrs"); check_cuda_contig(a_scales_ptrs, "a_scales_ptrs"); check_cuda_contig(b_scales_ptrs, "b_scales_ptrs"); check_cuda_contig(alpha_ptrs, "alpha_ptrs"); check_cuda_contig(layout_sfa, "layout_sfa"); check_cuda_contig(layout_sfb, "layout_sfb"); RuntimeCheck( output.device() == a.device() && a.device() == b.device() && a.device() == a_blockscale.device() && a.device() == b_blockscales.device() && a.device() == alphas.device() && a.device() == ab_strides.device() && a.device() == c_strides.device() && a.device() == problem_sizes.device() && a.device() == expert_offsets.device() && a.device() == sf_offsets.device() && a.device() == a_ptrs.device() && a.device() == b_ptrs.device() && a.device() == out_ptrs.device() && a.device() == a_scales_ptrs.device() && a.device() == b_scales_ptrs.device() && a.device() == alpha_ptrs.device() && a.device() == layout_sfa.device() && a.device() == layout_sfb.device(), "all tensors must be on the same device"); RuntimeCheck(host::is_type(a.dtype()), "a must be uint8"); RuntimeCheck(host::is_type(b.dtype()), "b must be uint8"); RuntimeCheck(host::is_type(a_blockscale.dtype()), "a_blockscale must be float8_e4m3fn"); RuntimeCheck(host::is_type(b_blockscales.dtype()), "b_blockscales must be float8_e4m3fn"); RuntimeCheck(host::is_type(alphas.dtype()), "alphas must be float32"); RuntimeCheck(host::is_type(ab_strides.dtype()), "ab_strides must be int64"); RuntimeCheck(host::is_type(c_strides.dtype()), "c_strides must be int64"); RuntimeCheck(host::is_type(problem_sizes.dtype()), "problem_sizes must be int32"); RuntimeCheck(host::is_type(expert_offsets.dtype()), "expert_offsets must be int32"); RuntimeCheck(host::is_type(sf_offsets.dtype()), "sf_offsets must be int32"); RuntimeCheck(host::is_type(a_ptrs.dtype()), "a_ptrs must be int64"); RuntimeCheck(host::is_type(b_ptrs.dtype()), "b_ptrs must be int64"); RuntimeCheck(host::is_type(out_ptrs.dtype()), "out_ptrs must be int64"); RuntimeCheck(host::is_type(a_scales_ptrs.dtype()), "a_scales_ptrs must be int64"); RuntimeCheck(host::is_type(b_scales_ptrs.dtype()), "b_scales_ptrs must be int64"); RuntimeCheck(host::is_type(alpha_ptrs.dtype()), "alpha_ptrs must be int64"); RuntimeCheck(host::is_type(layout_sfa.dtype()), "layout_sfa must be int64"); RuntimeCheck(host::is_type(layout_sfb.dtype()), "layout_sfb must be int64"); RuntimeCheck( host::is_type(output.dtype()) || host::is_type(output.dtype()), "output must be bfloat16 or float16"); RuntimeCheck(a.dim() == 2, "a must be 2D"); RuntimeCheck(b.dim() == 3, "b must be 3D"); RuntimeCheck(a_blockscale.dim() == 2, "a_blockscale must be 2D"); RuntimeCheck(b_blockscales.dim() == 3, "b_blockscales must be 3D"); RuntimeCheck(alphas.dim() == 1, "alphas must be 1D"); RuntimeCheck(ab_strides.dim() == 1, "ab_strides must be 1D"); RuntimeCheck(c_strides.dim() == 1, "c_strides must be 1D"); RuntimeCheck(problem_sizes.dim() == 2, "problem_sizes must be 2D"); RuntimeCheck(expert_offsets.dim() == 1, "expert_offsets must be 1D"); RuntimeCheck(sf_offsets.dim() == 1, "sf_offsets must be 1D"); RuntimeCheck(a_ptrs.dim() == 1, "a_ptrs must be 1D"); RuntimeCheck(b_ptrs.dim() == 1, "b_ptrs must be 1D"); RuntimeCheck(out_ptrs.dim() == 1, "out_ptrs must be 1D"); RuntimeCheck(a_scales_ptrs.dim() == 1, "a_scales_ptrs must be 1D"); RuntimeCheck(b_scales_ptrs.dim() == 1, "b_scales_ptrs must be 1D"); RuntimeCheck(alpha_ptrs.dim() == 1, "alpha_ptrs must be 1D"); RuntimeCheck(layout_sfa.dim() == 2, "layout_sfa must be 2D"); RuntimeCheck(layout_sfb.dim() == 2, "layout_sfb must be 2D"); RuntimeCheck(problem_sizes.size(1) == 3, "problem_sizes must have shape (num_experts, 3)"); const int num_experts = static_cast(expert_offsets.size(0)); RuntimeCheck(problem_sizes.size(0) == num_experts, "problem_sizes size mismatch with expert_offsets"); RuntimeCheck(sf_offsets.size(0) == num_experts, "sf_offsets size mismatch with expert_offsets"); RuntimeCheck(alphas.size(0) == num_experts, "alphas size mismatch with expert_offsets"); RuntimeCheck(ab_strides.size(0) == num_experts, "ab_strides size mismatch with expert_offsets"); RuntimeCheck(c_strides.size(0) == num_experts, "c_strides size mismatch with expert_offsets"); RuntimeCheck(a_ptrs.size(0) == num_experts, "a_ptrs size mismatch with expert_offsets"); RuntimeCheck(b_ptrs.size(0) == num_experts, "b_ptrs size mismatch with expert_offsets"); RuntimeCheck(out_ptrs.size(0) == num_experts, "out_ptrs size mismatch with expert_offsets"); RuntimeCheck(a_scales_ptrs.size(0) == num_experts, "a_scales_ptrs size mismatch with expert_offsets"); RuntimeCheck(b_scales_ptrs.size(0) == num_experts, "b_scales_ptrs size mismatch with expert_offsets"); RuntimeCheck(alpha_ptrs.size(0) == num_experts, "alpha_ptrs size mismatch with expert_offsets"); RuntimeCheck(layout_sfa.size(0) == num_experts && layout_sfa.size(1) == 5, "layout_sfa must be [num_experts, 5]"); RuntimeCheck(layout_sfb.size(0) == num_experts && layout_sfb.size(1) == 5, "layout_sfb must be [num_experts, 5]"); int M = static_cast(a.size(0)); int N = static_cast(b.size(1)); int K = static_cast(2 * b.size(2)); RuntimeCheck(output.dim() == 2, "output must be 2D"); RuntimeCheck(output.size(0) == M && output.size(1) == N, "output shape mismatch"); auto sm_version = getSMVersion(a.device().device_id); if (sm_version == 100 || sm_version == 103) { if (host::is_type(output.dtype())) { run_fp4_blockwise_scaled_group_mm_sm100( output, a, b, a_blockscale, b_blockscales, alphas, ab_strides, c_strides, problem_sizes, expert_offsets, sf_offsets, a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs, layout_sfa, layout_sfb, M, N, K); } else { run_fp4_blockwise_scaled_group_mm_sm100( output, a, b, a_blockscale, b_blockscales, alphas, ab_strides, c_strides, problem_sizes, expert_offsets, sf_offsets, a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs, layout_sfa, layout_sfb, M, N, K); } } else if (sm_version >= 120) { if (host::is_type(output.dtype())) { run_fp4_blockwise_scaled_group_mm_sm120( output, a, b, a_blockscale, b_blockscales, alphas, ab_strides, c_strides, problem_sizes, expert_offsets, sf_offsets, a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs, layout_sfa, layout_sfb, M, N, K); } else { Panic("SM120 path currently supports only bfloat16 output"); } } else { RuntimeCheck(false, "Unsupported SM version: ", sm_version); } } void cutlass_fp4_group_mm( tvm::ffi::TensorView output, const tvm::ffi::TensorView a, const tvm::ffi::TensorView b, const tvm::ffi::TensorView a_blockscale, const tvm::ffi::TensorView b_blockscales, const tvm::ffi::TensorView alphas, const tvm::ffi::TensorView ab_strides, const tvm::ffi::TensorView c_strides, const tvm::ffi::TensorView problem_sizes, const tvm::ffi::TensorView expert_offsets, const tvm::ffi::TensorView sf_offsets, const tvm::ffi::TensorView a_ptrs, const tvm::ffi::TensorView b_ptrs, const tvm::ffi::TensorView out_ptrs, const tvm::ffi::TensorView a_scales_ptrs, const tvm::ffi::TensorView b_scales_ptrs, const tvm::ffi::TensorView alpha_ptrs, const tvm::ffi::TensorView layout_sfa, const tvm::ffi::TensorView layout_sfb) { cutlass_fp4_group_mm_sm100a_sm120a( output, a, b, a_blockscale, b_blockscales, alphas, ab_strides, c_strides, problem_sizes, expert_offsets, sf_offsets, a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs, layout_sfa, layout_sfb); } ================================================ FILE: python/sglang/jit_kernel/csrc/ngram_embedding.cuh ================================================ #include #include #include #include #include #include #include #include #include namespace device::ngram_embedding { __global__ void ComputeNGramIdsKernel( int batch_size, int ne_n, int ne_k, int* ne_weights, // [ne_n-1,ne_k,ne_n] int* ne_mods, // [ne_n-1,ne_k] int* exclusive_ne_embeder_size_sums, // [(ne_n-1)*ne_k] int* tokens, // [token_num] int* exclusive_req_len_sums, // [batch_size+1] int* ne_token_table, // [max_running_reqs, max_context_len] int max_context_len, // max_context_len long* row_indices, // [batch_size] int* column_starts, // [batch_size] int* n_gram_ids // [ne_n-1,ne_k,token_num] ) { // Determine which n, k, and request this block handles. /** Example: [req0, req1, req2] with n=3, k=2 n k req_id blockIdx.x config_id (combination of n and k) 2 1 0 0 0 2 1 1 1 0 2 1 2 2 0 2 2 0 3 1 2 2 1 4 1 2 2 2 5 1 3 1 0 0 2 3 1 1 1 2 3 1 2 2 2 3 2 0 3 3 3 2 1 4 3 3 2 2 5 3 */ const int req_id = blockIdx.x % batch_size; const int config_id = (blockIdx.x - req_id) / batch_size; // n and k here are offset from their physical meanings: n = real_n - 2, k = real_k - 1. // This offset exists because n and k are used as indices into ne_weights and ne_mods. const int k = config_id % ne_k; const int n = (config_id - config_id % ne_k) / ne_k; // ne_weights has shape [ne_n-1, ne_k, ne_n]; last dim is token distance, so compute base index first const int ne_weight_base_idx = n * ne_k * ne_n + k * ne_n; // ne_mods has shape [ne_n-1, ne_k] const int ne_mod = ne_mods[n * ne_k + k]; // stride loop for (int i = exclusive_req_len_sums[req_id] + threadIdx.x; i < exclusive_req_len_sums[req_id + 1]; i += blockDim.x) { uint64_t n_gram_id = 0; // Token offset within the current request int current_token_offset = i - exclusive_req_len_sums[req_id]; // Start index of this request in the token table; tokens before this belong to other requests int req_token_table_index = row_indices[req_id] * max_context_len; // Position of the current token in the token table int current_token_table_index = req_token_table_index + column_starts[req_id] + current_token_offset; for (int j = 0; j < n + 2; j++) { if (current_token_table_index - j < req_token_table_index) { // Out of this request's range, stop computing n_gram_id break; } if (ne_token_table[current_token_table_index - j] < 0) { // Token was marked as ignored during write break; } const uint64_t term = (uint64_t)ne_token_table[current_token_table_index - j] * (uint64_t)ne_weights[ne_weight_base_idx + j]; n_gram_id += term % ne_mod; } n_gram_id %= ne_mod; n_gram_id += exclusive_ne_embeder_size_sums[n * ne_k + k]; // [token_num, ne_n-1, ne_k] n_gram_ids[i * (ne_n - 1) * ne_k + n * ne_k + k] = (int)(n_gram_id); } } __global__ void UpdateTokenTableKernel( int batch_size, int* tokens, // [token_num] int* ne_token_table, // [max_running_reqs, max_context_len] int max_context_len, // max_context_len long* row_indices, // [batch_size] int* column_starts, // [batch_size] int* req_lens, // [batch_size] int ignore_token_num, // number of tokens to ignore int* ignore_tokens // [ignore_token_num] ) { // Each block processes one request. const int req_id = blockIdx.x % batch_size; int start = 0; int end = 0; for (int i = 0; i < req_id; i++) { start += req_lens[i]; } end = start + req_lens[req_id]; // stride loop for (int i = start + threadIdx.x; i < end; i += blockDim.x) { // Token offset within the current request int current_token_offset = i - start; // Start index of this request in the token table int req_token_table_index = row_indices[req_id] * max_context_len; // Position of the current token in the token table int current_token_table_index = req_token_table_index + column_starts[req_id] + current_token_offset; ne_token_table[current_token_table_index] = tokens[i]; for (int j = 0; j < ignore_token_num; j++) { if (ignore_tokens[j] == tokens[i]) { ne_token_table[current_token_table_index] = -tokens[i]; break; } } } } } // namespace device::ngram_embedding namespace { struct NgramEmbeddingKernel { static void compute_n_gram_ids( const int64_t ne_n, const int64_t ne_k, const tvm::ffi::TensorView ne_weights, const tvm::ffi::TensorView ne_mods, const tvm::ffi::TensorView exclusive_ne_embeder_size_sums, const tvm::ffi::TensorView tokens, const tvm::ffi::TensorView exclusive_req_len_sums, const tvm::ffi::TensorView ne_token_table, const tvm::ffi::TensorView row_indices, const tvm::ffi::TensorView column_starts, const tvm::ffi::TensorView n_gram_ids) { using namespace host; auto device_ = SymbolicDevice{}; // Verify tensor shapes and types using -1 (kAnySize) for dynamic dimensions TensorMatcher({-1, -1, -1}) // [ne_n-1, ne_k, ne_n] .with_dtype() .with_device(device_) .verify(ne_weights); TensorMatcher({-1, -1}) // [ne_n-1, ne_k] .with_dtype() .with_device() .verify(ne_mods); TensorMatcher({-1}) // [(ne_n-1)*ne_k + 1] .with_dtype() .with_device() .verify(exclusive_ne_embeder_size_sums); TensorMatcher({-1}) // [token_num] .with_dtype() .with_device() .verify(tokens); TensorMatcher({-1}) // [batch_size+1] .with_dtype() .with_device() .verify(exclusive_req_len_sums); TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len] .with_dtype() .with_device() .verify(ne_token_table); TensorMatcher({-1}) // [batch_size] .with_dtype() .with_device() .verify(row_indices); TensorMatcher({-1}) // [batch_size] .with_dtype() .with_device() .verify(column_starts); TensorMatcher({-1, -1}) // [token_num, (ne_n-1)*ne_k] .with_dtype() .with_device() .verify(n_gram_ids); const int batch_size = static_cast(exclusive_req_len_sums.size(0) - 1); const int max_context_len = static_cast(ne_token_table.size(1)); const auto stream = LaunchKernel::resolve_device(device_.unwrap()); constexpr int BLOCK_THREADS = 256; const int num_configs = (static_cast(ne_n) - 1) * static_cast(ne_k); const int grid_size = num_configs * batch_size; LaunchKernel(grid_size, BLOCK_THREADS, stream)( device::ngram_embedding::ComputeNGramIdsKernel, batch_size, static_cast(ne_n), static_cast(ne_k), static_cast(ne_weights.data_ptr()), static_cast(ne_mods.data_ptr()), static_cast(exclusive_ne_embeder_size_sums.data_ptr()), static_cast(tokens.data_ptr()), static_cast(exclusive_req_len_sums.data_ptr()), static_cast(ne_token_table.data_ptr()), max_context_len, static_cast(row_indices.data_ptr()), static_cast(column_starts.data_ptr()), static_cast(n_gram_ids.data_ptr())); } static void update_token_table( const tvm::ffi::TensorView tokens, const tvm::ffi::TensorView ne_token_table, const tvm::ffi::TensorView row_indices, const tvm::ffi::TensorView column_starts, const tvm::ffi::TensorView req_lens, const tvm::ffi::TensorView ignore_tokens) { using namespace host; auto device_ = SymbolicDevice{}; // Verify tensor shapes and types using -1 (kAnySize) for dynamic dimensions TensorMatcher({-1}) // [token_num] .with_dtype() .with_device(device_) .verify(tokens); TensorMatcher({-1, -1}) // [max_running_reqs, max_context_len] .with_dtype() .with_device() .verify(ne_token_table); TensorMatcher({-1}) // [batch_size] .with_dtype() .with_device() .verify(row_indices); TensorMatcher({-1}) // [batch_size] .with_dtype() .with_device() .verify(column_starts); TensorMatcher({-1}) // [batch_size] .with_dtype() .with_device() .verify(req_lens); // ignore_tokens can be empty or have values void* ignore_tokens_ptr = ignore_tokens.data_ptr(); const bool has_ignore_tokens = ignore_tokens_ptr != nullptr && ignore_tokens.numel() > 0; if (has_ignore_tokens) { TensorMatcher({-1}) // [ignore_token_num] .with_dtype() .with_device() .verify(ignore_tokens); } const int batch_size = static_cast(req_lens.size(0)); if (batch_size <= 0) { return; } const int max_context_len = static_cast(ne_token_table.size(1)); const auto stream = LaunchKernel::resolve_device(device_.unwrap()); constexpr int BLOCK_THREADS = 256; const int grid_size = batch_size; int ignore_token_num = 0; int* ignore_tokens_typed_ptr = nullptr; if (has_ignore_tokens) { ignore_token_num = static_cast(ignore_tokens.numel()); ignore_tokens_typed_ptr = static_cast(ignore_tokens_ptr); } LaunchKernel(grid_size, BLOCK_THREADS, stream)( device::ngram_embedding::UpdateTokenTableKernel, batch_size, static_cast(tokens.data_ptr()), static_cast(ne_token_table.data_ptr()), max_context_len, static_cast(row_indices.data_ptr()), static_cast(column_starts.data_ptr()), static_cast(req_lens.data_ptr()), ignore_token_num, ignore_tokens_typed_ptr); } }; } // namespace ================================================ FILE: python/sglang/jit_kernel/csrc/nsa/fused_store_index_cache.cuh ================================================ #include #include #include #include #include #include #include #include #include #include #include #include namespace { struct FusedStoreCacheParam { const void* __restrict__ input; void* __restrict__ cache; const void* __restrict__ indices; uint32_t num_tokens; }; [[maybe_unused]] SGL_DEVICE float fp8_e4m3_clip(float val) { namespace math = device::math; return math::max(math::min(val, math::FP8_E4M3_MAX), -math::FP8_E4M3_MAX); } [[maybe_unused]] SGL_DEVICE fp8x2_e4m3_t pack_fp8(float x, float y) { return fp8x2_e4m3_t{fp32x2_t{fp8_e4m3_clip(x), fp8_e4m3_clip(y)}}; } template __global__ void fused_store_indexer_cache(const __grid_constant__ FusedStoreCacheParam param) { using namespace device; /// NOTE: 132 = 128 + 4 constexpr int64_t kPageBytes = 132 << kPageBits; // each warp handles 128 elements, each block handles multiple rows const auto& [input, cache, indices, num_tokens] = param; const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x; const auto global_wid = global_tid / 32; const auto lane_id = threadIdx.x % 32; if (global_wid >= num_tokens) return; PDLWaitPrimary(); // wait for primary kernel // prefetch the index const auto index = static_cast(indices)[global_wid]; // always load the value from input (don't store if invalid) using KeyT2 = packed_t; using InStorage = AlignedVector; using OutStorage = AlignedVector; const auto elems = static_cast(input)[global_tid]; const auto [x0, x1] = cast(elems[0]); const auto [y0, y1] = cast(elems[1]); const auto local_max = fmaxf(fmaxf(fabs(x0), fabs(x1)), fmaxf(fabs(y0), fabs(y1))); const auto abs_max = warp::reduce_max(local_max); // use normal fp32 scale const auto scale = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX; const auto inv_scale = 1.0f / scale; const int32_t page = index >> kPageBits; const int32_t offset = index & ((1 << kPageBits) - 1); const auto page_ptr = pointer::offset(cache, page * kPageBytes); const auto value_ptr = pointer::offset(page_ptr, offset * 128); const auto scale_ptr = pointer::offset(page_ptr, 128 << kPageBits, offset * 4); OutStorage result; result[0] = pack_fp8(x0 * inv_scale, x1 * inv_scale); result[1] = pack_fp8(y0 * inv_scale, y1 * inv_scale); static_cast(value_ptr)[lane_id] = result; static_cast(scale_ptr)[0] = scale; PDLTriggerSecondary(); // launch secondary kernel } template struct FusedStoreCacheIndexerKernel { static constexpr int32_t kLogSize = std::countr_zero(kPageSize); /// NOTE: 132 = 128 + 4 (128 represent K and 4 represent scale) static constexpr int64_t kPageBytes = 132 * kPageSize; static constexpr auto kernel = fused_store_indexer_cache; static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2"); static_assert(1 << kLogSize == kPageSize); static void run(tvm::ffi::TensorView input, tvm::ffi::TensorView cache, tvm::ffi::TensorView indices) { using namespace host; auto N = SymbolicSize{"num_tokens"}; auto device_ = SymbolicDevice{}; device_.set_options(); TensorMatcher({N, 128}) // input .with_dtype() .with_device(device_) .verify(input); TensorMatcher({-1, -1}) // cache .with_strides({kPageBytes, 1}) .with_dtype() .with_device(device_) .verify(cache); TensorMatcher({N}) // indices .with_dtype() .with_device(device_) .verify(indices); const auto num_tokens = static_cast(N.unwrap()); const auto params = FusedStoreCacheParam{ .input = input.data_ptr(), .cache = cache.data_ptr(), .indices = indices.data_ptr(), .num_tokens = num_tokens, }; const auto kBlockSize = 128; const auto num_blocks = div_ceil(num_tokens * 32, kBlockSize); LaunchKernel(num_blocks, kBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(kernel, params); } }; } // namespace ================================================ FILE: python/sglang/jit_kernel/cutedsl_gdn.py ================================================ """CuTe DSL Fused Sigmoid Gating Delta Rule Kernel for GDN Decode.""" import logging from typing import Dict, Optional, Tuple import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import torch from cutlass.cute.nvgpu import cpasync from cutlass.cute.runtime import from_dlpack logger = logging.getLogger(__name__) _compiled_kernels: Dict[Tuple, object] = {} _cu_seqlens_cache: Dict[Tuple, torch.Tensor] = {} TILE_K = 128 TILE_V = 32 TILE_V_PADDED = 36 TILE_V_SMALL = 16 TILE_V_SMALL_PADDED = 20 NUM_STAGES = 2 NUM_THREADS = 128 NUM_BLOCKS_PER_STATE_SMALL = 8 NUM_THREADS_LARGE = 256 NUM_WARPS_LARGE = 8 V_PER_WARP = 4 ROWS_PER_ITER = 8 NUM_K_ITERS = TILE_K // ROWS_PER_ITER SMALL_BATCH_THRESHOLD = 32 def _define_kernels(): """Define CuTe DSL kernels for normal and varlen decode modes.""" NUM_WARPS_SMALL = 4 V_PER_WARP_SMALL = TILE_V_SMALL // NUM_WARPS_SMALL ROWS_PER_ITER_SMALL = 32 // V_PER_WARP_SMALL NUM_K_ITERS_SMALL = TILE_K // ROWS_PER_ITER_SMALL @cute.kernel def gdn_kernel_small_batch( tiled_copy_load: cute.TiledCopy, h0_source: cute.Tensor, smem_layout_staged: cute.Layout, num_v_tiles: cutlass.Constexpr[int], q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, a: cute.Tensor, b: cute.Tensor, A_log: cute.Tensor, dt_bias: cute.Tensor, o: cute.Tensor, h0_indices: cute.Tensor, softplus_beta: cutlass.Constexpr[float], softplus_threshold: cutlass.Constexpr[float], scale: cutlass.Constexpr[float], H: cutlass.Constexpr[int], HV: cutlass.Constexpr[int], use_qk_l2norm: cutlass.Constexpr[bool], ): """Small batch kernel for (N, 1, ...) format.""" tidx, _, _ = cute.arch.thread_idx() in_warp_tid = tidx % 32 warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) block_idx, _, _ = cute.arch.block_idx() batch_idx = block_idx // NUM_BLOCKS_PER_STATE_SMALL batch_inner = block_idx % NUM_BLOCKS_PER_STATE_SMALL num_v_tiles_per_block = num_v_tiles // NUM_BLOCKS_PER_STATE_SMALL start_v_tile = batch_inner * num_v_tiles_per_block i_n = batch_idx // HV i_hv = batch_idx % HV i_h = i_hv // (HV // H) pool_idx = h0_indices[i_n] if pool_idx >= 0: k_local = in_warp_tid // V_PER_WARP_SMALL v_local = in_warp_tid % V_PER_WARP_SMALL v_base = warp_idx * V_PER_WARP_SMALL v_idx = v_base + v_local smem = cutlass.utils.SmemAllocator() sData = smem.allocate_tensor(cutlass.Float32, smem_layout_staged, 128) smem_o_layout = cute.make_layout((TILE_V_SMALL,), stride=(1,)) smem_o = smem.allocate_tensor(cutlass.Float32, smem_o_layout, 128) smem_k_layout = cute.make_layout((TILE_K,), stride=(1,)) smem_q_layout = cute.make_layout((TILE_K,), stride=(1,)) sK = smem.allocate_tensor(cutlass.Float32, smem_k_layout, 128) sQ = smem.allocate_tensor(cutlass.Float32, smem_q_layout, 128) if tidx < TILE_K: sK[tidx] = cutlass.Float32(k[i_n, 0, i_h, tidx]) sQ[tidx] = cutlass.Float32(q[i_n, 0, i_h, tidx]) gSrc_batch = h0_source[(pool_idx, i_hv, None, None)] gSrc = cute.local_tile(gSrc_batch, (TILE_K, TILE_V_SMALL), (0, None)) thr_copy_load = tiled_copy_load.get_slice(tidx) prefetch_count = cutlass.min(NUM_STAGES - 1, num_v_tiles_per_block) for v_tile_offset in range(prefetch_count): v_tile = start_v_tile + v_tile_offset stage = v_tile_offset % NUM_STAGES gSrc_tile = gSrc[(None, None, v_tile)] sData_stage = sData[(None, None, stage)] thr_gSrc = thr_copy_load.partition_S(gSrc_tile) thr_sData = thr_copy_load.partition_D(sData_stage) cute.copy(tiled_copy_load, thr_gSrc, thr_sData) cute.arch.cp_async_commit_group() r_A_log = cutlass.Float32(A_log[i_hv]) r_dt_bias = cutlass.Float32(dt_bias[i_hv]) r_a = cutlass.Float32(a[i_n, 0, i_hv]) r_b = cutlass.Float32(b[i_n, 0, i_hv]) r_g = 0.0 r_beta = 0.0 if in_warp_tid == 0: x = r_a + r_dt_bias beta_x = softplus_beta * x softplus_x = 0.0 if beta_x <= softplus_threshold: exp_beta_x = cute.exp(beta_x) log_input = cutlass.Float32(1.0 + exp_beta_x) log_result = cutlass.Float32(cute.log(log_input)) softplus_x = cutlass.Float32( (cutlass.Float32(1.0) / softplus_beta) * log_result ) else: softplus_x = x r_g_value = -cute.exp(r_A_log) * softplus_x r_beta = 1.0 / (1.0 + cute.exp(-r_b)) r_g = cute.exp(r_g_value) r_g = cute.arch.shuffle_sync(r_g, 0) r_beta = cute.arch.shuffle_sync(r_beta, 0) cute.arch.barrier() if use_qk_l2norm: sum_q_partial = 0.0 sum_k_partial = 0.0 if tidx < TILE_K: q_val = sQ[tidx] k_val = sK[tidx] sum_q_partial = q_val * q_val sum_k_partial = k_val * k_val for offset in [16, 8, 4, 2, 1]: sum_q_partial += cute.arch.shuffle_sync_bfly( sum_q_partial, offset=offset, mask=-1, mask_and_clamp=31 ) sum_k_partial += cute.arch.shuffle_sync_bfly( sum_k_partial, offset=offset, mask=-1, mask_and_clamp=31 ) if in_warp_tid == 0: smem_o[warp_idx] = sum_q_partial smem_o[warp_idx + 4] = sum_k_partial cute.arch.barrier() inv_norm_q = 0.0 inv_norm_k = 0.0 if warp_idx == 0: local_sum_q = 0.0 local_sum_k = 0.0 if in_warp_tid < NUM_WARPS_SMALL: local_sum_q = smem_o[in_warp_tid] local_sum_k = smem_o[in_warp_tid + 4] for offset in [2, 1]: local_sum_q += cute.arch.shuffle_sync_bfly( local_sum_q, offset=offset, mask=-1, mask_and_clamp=31 ) local_sum_k += cute.arch.shuffle_sync_bfly( local_sum_k, offset=offset, mask=-1, mask_and_clamp=31 ) if in_warp_tid == 0: smem_o[0] = cute.rsqrt(local_sum_q + 1e-6) smem_o[1] = cute.rsqrt(local_sum_k + 1e-6) cute.arch.barrier() inv_norm_q = smem_o[0] inv_norm_k = smem_o[1] if tidx < TILE_K: sK[tidx] = sK[tidx] * inv_norm_k sQ[tidx] = sQ[tidx] * scale * inv_norm_q cute.arch.barrier() else: if tidx < TILE_K: sQ[tidx] = sQ[tidx] * scale cute.arch.barrier() for v_tile_offset in range(num_v_tiles_per_block): v_tile = start_v_tile + v_tile_offset stage = v_tile_offset % NUM_STAGES cute.arch.cp_async_wait_group(0) cute.arch.barrier() next_v_tile_offset = v_tile_offset + prefetch_count if next_v_tile_offset < num_v_tiles_per_block: next_v_tile = start_v_tile + next_v_tile_offset next_stage = next_v_tile_offset % NUM_STAGES gSrc_next = gSrc[(None, None, next_v_tile)] sData_next = sData[(None, None, next_stage)] thr_gSrc = thr_copy_load.partition_S(gSrc_next) thr_sData = thr_copy_load.partition_D(sData_next) cute.copy(tiled_copy_load, thr_gSrc, thr_sData) cute.arch.cp_async_commit_group() v_global = v_tile * TILE_V_SMALL + v_idx r_v = cutlass.Float32(v[i_n, 0, i_hv, v_global]) sum_hk = 0.0 for k_iter in cutlass.range_dynamic(NUM_K_ITERS_SMALL, unroll=16): k_base = k_iter * ROWS_PER_ITER_SMALL k_idx = k_base + k_local h_val = sData[(k_idx, v_idx, stage)] * r_g r_k_val = sK[k_idx] sum_hk += h_val * r_k_val for offset in [4, 2, 1]: sum_hk += cute.arch.shuffle_sync_bfly( sum_hk, offset=offset * V_PER_WARP_SMALL, mask=-1, mask_and_clamp=31, ) v_new = (r_v - sum_hk) * r_beta v_new = cute.arch.shuffle_sync(v_new, v_local) sum_hq = 0.0 for k_iter in cutlass.range_dynamic(NUM_K_ITERS_SMALL, unroll=16): k_base = k_iter * ROWS_PER_ITER_SMALL k_idx = k_base + k_local h_old = sData[(k_idx, v_idx, stage)] * r_g r_k_val = sK[k_idx] r_q_val = sQ[k_idx] h_new = h_old + r_k_val * v_new sData[(k_idx, v_idx, stage)] = h_new sum_hq += h_new * r_q_val for offset in [4, 2, 1]: sum_hq += cute.arch.shuffle_sync_bfly( sum_hq, offset=offset * V_PER_WARP_SMALL, mask=-1, mask_and_clamp=31, ) if k_local == 0: v_global_out = v_tile * TILE_V_SMALL + v_idx o[(i_n, 0, i_hv, v_global_out)] = cutlass.BFloat16(sum_hq) cute.arch.barrier() for k_iter in range(NUM_K_ITERS_SMALL): flat_idx = tidx + k_iter * 128 k_write = flat_idx // TILE_V_SMALL v_write = flat_idx % TILE_V_SMALL if k_write < TILE_K: h_val = sData[(k_write, v_write, stage)] v_global_write = v_tile * TILE_V_SMALL + v_write h0_source[(pool_idx, i_hv, k_write, v_global_write)] = h_val cute.arch.barrier() @cute.kernel def gdn_kernel_small_batch_varlen( tiled_copy_load: cute.TiledCopy, h0_source: cute.Tensor, smem_layout_staged: cute.Layout, num_v_tiles: cutlass.Constexpr[int], q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, a: cute.Tensor, b: cute.Tensor, A_log: cute.Tensor, dt_bias: cute.Tensor, o: cute.Tensor, h0_indices: cute.Tensor, softplus_beta: cutlass.Constexpr[float], softplus_threshold: cutlass.Constexpr[float], scale: cutlass.Constexpr[float], H: cutlass.Constexpr[int], HV: cutlass.Constexpr[int], use_qk_l2norm: cutlass.Constexpr[bool], ): """Small batch kernel for varlen decode (1, N, ...) format.""" tidx, _, _ = cute.arch.thread_idx() in_warp_tid = tidx % 32 warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) block_idx, _, _ = cute.arch.block_idx() batch_idx = block_idx // NUM_BLOCKS_PER_STATE_SMALL batch_inner = block_idx % NUM_BLOCKS_PER_STATE_SMALL num_v_tiles_per_block = num_v_tiles // NUM_BLOCKS_PER_STATE_SMALL start_v_tile = batch_inner * num_v_tiles_per_block i_n = batch_idx // HV i_hv = batch_idx % HV i_h = i_hv // (HV // H) pool_idx = h0_indices[i_n] if pool_idx >= 0: k_local = in_warp_tid // V_PER_WARP_SMALL v_local = in_warp_tid % V_PER_WARP_SMALL v_base = warp_idx * V_PER_WARP_SMALL v_idx = v_base + v_local smem = cutlass.utils.SmemAllocator() sData = smem.allocate_tensor(cutlass.Float32, smem_layout_staged, 128) smem_o_layout = cute.make_layout((TILE_V_SMALL,), stride=(1,)) smem_o = smem.allocate_tensor(cutlass.Float32, smem_o_layout, 128) smem_k_layout = cute.make_layout((TILE_K,), stride=(1,)) smem_q_layout = cute.make_layout((TILE_K,), stride=(1,)) sK = smem.allocate_tensor(cutlass.Float32, smem_k_layout, 128) sQ = smem.allocate_tensor(cutlass.Float32, smem_q_layout, 128) if tidx < TILE_K: sK[tidx] = cutlass.Float32(k[0, i_n, i_h, tidx]) sQ[tidx] = cutlass.Float32(q[0, i_n, i_h, tidx]) gSrc_batch = h0_source[(pool_idx, i_hv, None, None)] gSrc = cute.local_tile(gSrc_batch, (TILE_K, TILE_V_SMALL), (0, None)) thr_copy_load = tiled_copy_load.get_slice(tidx) prefetch_count = cutlass.min(NUM_STAGES - 1, num_v_tiles_per_block) for v_tile_offset in range(prefetch_count): v_tile = start_v_tile + v_tile_offset stage = v_tile_offset % NUM_STAGES gSrc_tile = gSrc[(None, None, v_tile)] sData_stage = sData[(None, None, stage)] thr_gSrc = thr_copy_load.partition_S(gSrc_tile) thr_sData = thr_copy_load.partition_D(sData_stage) cute.copy(tiled_copy_load, thr_gSrc, thr_sData) cute.arch.cp_async_commit_group() r_A_log = cutlass.Float32(A_log[i_hv]) r_dt_bias = cutlass.Float32(dt_bias[i_hv]) r_a = cutlass.Float32(a[i_n, i_hv]) r_b = cutlass.Float32(b[i_n, i_hv]) r_g = 0.0 r_beta = 0.0 if in_warp_tid == 0: x = r_a + r_dt_bias beta_x = softplus_beta * x softplus_x = 0.0 if beta_x <= softplus_threshold: exp_beta_x = cute.exp(beta_x) log_input = cutlass.Float32(1.0 + exp_beta_x) log_result = cutlass.Float32(cute.log(log_input)) softplus_x = cutlass.Float32( (cutlass.Float32(1.0) / softplus_beta) * log_result ) else: softplus_x = x r_g_value = -cute.exp(r_A_log) * softplus_x r_beta = 1.0 / (1.0 + cute.exp(-r_b)) r_g = cute.exp(r_g_value) r_g = cute.arch.shuffle_sync(r_g, 0) r_beta = cute.arch.shuffle_sync(r_beta, 0) cute.arch.barrier() if use_qk_l2norm: sum_q_partial = 0.0 sum_k_partial = 0.0 if tidx < TILE_K: q_val = sQ[tidx] k_val = sK[tidx] sum_q_partial = q_val * q_val sum_k_partial = k_val * k_val for offset in [16, 8, 4, 2, 1]: sum_q_partial += cute.arch.shuffle_sync_bfly( sum_q_partial, offset=offset, mask=-1, mask_and_clamp=31 ) sum_k_partial += cute.arch.shuffle_sync_bfly( sum_k_partial, offset=offset, mask=-1, mask_and_clamp=31 ) if in_warp_tid == 0: smem_o[warp_idx] = sum_q_partial smem_o[warp_idx + 4] = sum_k_partial cute.arch.barrier() inv_norm_q = 0.0 inv_norm_k = 0.0 if warp_idx == 0: local_sum_q = 0.0 local_sum_k = 0.0 if in_warp_tid < NUM_WARPS_SMALL: local_sum_q = smem_o[in_warp_tid] local_sum_k = smem_o[in_warp_tid + 4] for offset in [2, 1]: local_sum_q += cute.arch.shuffle_sync_bfly( local_sum_q, offset=offset, mask=-1, mask_and_clamp=31 ) local_sum_k += cute.arch.shuffle_sync_bfly( local_sum_k, offset=offset, mask=-1, mask_and_clamp=31 ) if in_warp_tid == 0: smem_o[0] = cute.rsqrt(local_sum_q + 1e-6) smem_o[1] = cute.rsqrt(local_sum_k + 1e-6) cute.arch.barrier() inv_norm_q = smem_o[0] inv_norm_k = smem_o[1] if tidx < TILE_K: sK[tidx] = sK[tidx] * inv_norm_k sQ[tidx] = sQ[tidx] * scale * inv_norm_q cute.arch.barrier() else: if tidx < TILE_K: sQ[tidx] = sQ[tidx] * scale cute.arch.barrier() for v_tile_offset in range(num_v_tiles_per_block): v_tile = start_v_tile + v_tile_offset stage = v_tile_offset % NUM_STAGES cute.arch.cp_async_wait_group(0) cute.arch.barrier() next_v_tile_offset = v_tile_offset + prefetch_count if next_v_tile_offset < num_v_tiles_per_block: next_v_tile = start_v_tile + next_v_tile_offset next_stage = next_v_tile_offset % NUM_STAGES gSrc_next = gSrc[(None, None, next_v_tile)] sData_next = sData[(None, None, next_stage)] thr_gSrc = thr_copy_load.partition_S(gSrc_next) thr_sData = thr_copy_load.partition_D(sData_next) cute.copy(tiled_copy_load, thr_gSrc, thr_sData) cute.arch.cp_async_commit_group() v_global = v_tile * TILE_V_SMALL + v_idx r_v = cutlass.Float32(v[0, i_n, i_hv, v_global]) sum_hk = 0.0 for k_iter in cutlass.range_dynamic(NUM_K_ITERS_SMALL, unroll=16): k_base = k_iter * ROWS_PER_ITER_SMALL k_idx = k_base + k_local h_val = sData[(k_idx, v_idx, stage)] * r_g r_k_val = sK[k_idx] sum_hk += h_val * r_k_val for offset in [4, 2, 1]: sum_hk += cute.arch.shuffle_sync_bfly( sum_hk, offset=offset * V_PER_WARP_SMALL, mask=-1, mask_and_clamp=31, ) v_new = (r_v - sum_hk) * r_beta v_new = cute.arch.shuffle_sync(v_new, v_local) sum_hq = 0.0 for k_iter in cutlass.range_dynamic(NUM_K_ITERS_SMALL, unroll=16): k_base = k_iter * ROWS_PER_ITER_SMALL k_idx = k_base + k_local h_old = sData[(k_idx, v_idx, stage)] * r_g r_k_val = sK[k_idx] r_q_val = sQ[k_idx] h_new = h_old + r_k_val * v_new sData[(k_idx, v_idx, stage)] = h_new sum_hq += h_new * r_q_val for offset in [4, 2, 1]: sum_hq += cute.arch.shuffle_sync_bfly( sum_hq, offset=offset * V_PER_WARP_SMALL, mask=-1, mask_and_clamp=31, ) if k_local == 0: v_global_out = v_tile * TILE_V_SMALL + v_idx o[(0, i_n, i_hv, v_global_out)] = cutlass.BFloat16(sum_hq) cute.arch.barrier() for k_iter in range(NUM_K_ITERS_SMALL): flat_idx = tidx + k_iter * 128 k_write = flat_idx // TILE_V_SMALL v_write = flat_idx % TILE_V_SMALL if k_write < TILE_K: h_val = sData[(k_write, v_write, stage)] v_global_write = v_tile * TILE_V_SMALL + v_write h0_source[(pool_idx, i_hv, k_write, v_global_write)] = h_val cute.arch.barrier() @cute.kernel def gdn_kernel_large_batch( tiled_copy_load: cute.TiledCopy, h0_source: cute.Tensor, smem_layout_staged: cute.Layout, num_v_tiles: cutlass.Constexpr[int], q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, a: cute.Tensor, b: cute.Tensor, A_log: cute.Tensor, dt_bias: cute.Tensor, o: cute.Tensor, h0_indices: cute.Tensor, softplus_beta: cutlass.Constexpr[float], softplus_threshold: cutlass.Constexpr[float], scale: cutlass.Constexpr[float], H: cutlass.Constexpr[int], HV: cutlass.Constexpr[int], use_qk_l2norm: cutlass.Constexpr[bool], ): """Large batch kernel for (N, 1, ...) format.""" tidx, _, _ = cute.arch.thread_idx() in_warp_tid = tidx % 32 warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) batch_idx, _, _ = cute.arch.block_idx() i_n = batch_idx // HV i_hv = batch_idx % HV i_h = i_hv // (HV // H) pool_idx = h0_indices[i_n] if pool_idx >= 0: k_local = in_warp_tid // V_PER_WARP v_local = in_warp_tid % V_PER_WARP v_base = warp_idx * V_PER_WARP v_idx = v_base + v_local smem = cutlass.utils.SmemAllocator() sData = smem.allocate_tensor(cutlass.Float32, smem_layout_staged, 128) smem_o_layout = cute.make_layout((TILE_V,), stride=(1,)) smem_o = smem.allocate_tensor(cutlass.Float32, smem_o_layout, 128) smem_k_layout = cute.make_layout((TILE_K,), stride=(1,)) smem_q_layout = cute.make_layout((TILE_K,), stride=(1,)) sK = smem.allocate_tensor(cutlass.Float32, smem_k_layout, 128) sQ = smem.allocate_tensor(cutlass.Float32, smem_q_layout, 128) if tidx < TILE_K: sK[tidx] = cutlass.Float32(k[i_n, 0, i_h, tidx]) sQ[tidx] = cutlass.Float32(q[i_n, 0, i_h, tidx]) gSrc_batch = h0_source[(pool_idx, i_hv, None, None)] gSrc = cute.local_tile(gSrc_batch, (TILE_K, TILE_V), (0, None)) thr_copy_load = tiled_copy_load.get_slice(tidx) prefetch_count = cutlass.min(NUM_STAGES - 1, num_v_tiles) for v_tile in range(prefetch_count): stage = v_tile % NUM_STAGES gSrc_tile = gSrc[(None, None, v_tile)] sData_stage = sData[(None, None, stage)] thr_gSrc = thr_copy_load.partition_S(gSrc_tile) thr_sData = thr_copy_load.partition_D(sData_stage) cute.copy(tiled_copy_load, thr_gSrc, thr_sData) cute.arch.cp_async_commit_group() r_A_log = cutlass.Float32(A_log[i_hv]) r_dt_bias = cutlass.Float32(dt_bias[i_hv]) r_a = cutlass.Float32(a[i_n, 0, i_hv]) r_b = cutlass.Float32(b[i_n, 0, i_hv]) r_g = 0.0 r_beta = 0.0 if in_warp_tid == 0: x = r_a + r_dt_bias beta_x = softplus_beta * x softplus_x = 0.0 if beta_x <= softplus_threshold: exp_beta_x = cute.exp(beta_x) log_input = cutlass.Float32(1.0 + exp_beta_x) log_result = cutlass.Float32(cute.log(log_input)) softplus_x = cutlass.Float32( (cutlass.Float32(1.0) / softplus_beta) * log_result ) else: softplus_x = x r_g_value = -cute.exp(r_A_log) * softplus_x r_beta = 1.0 / (1.0 + cute.exp(-r_b)) r_g = cute.exp(r_g_value) r_g = cute.arch.shuffle_sync(r_g, 0) r_beta = cute.arch.shuffle_sync(r_beta, 0) cute.arch.barrier() if use_qk_l2norm: sum_q_partial = 0.0 sum_k_partial = 0.0 if tidx < TILE_K: q_val = sQ[tidx] k_val = sK[tidx] sum_q_partial = q_val * q_val sum_k_partial = k_val * k_val for offset in [16, 8, 4, 2, 1]: sum_q_partial += cute.arch.shuffle_sync_bfly( sum_q_partial, offset=offset, mask=-1, mask_and_clamp=31 ) sum_k_partial += cute.arch.shuffle_sync_bfly( sum_k_partial, offset=offset, mask=-1, mask_and_clamp=31 ) if in_warp_tid == 0: smem_o[warp_idx] = sum_q_partial smem_o[warp_idx + 8] = sum_k_partial cute.arch.barrier() inv_norm_q = 0.0 inv_norm_k = 0.0 if warp_idx == 0: local_sum_q = 0.0 local_sum_k = 0.0 if in_warp_tid < NUM_WARPS_LARGE: local_sum_q = smem_o[in_warp_tid] local_sum_k = smem_o[in_warp_tid + 8] for offset in [4, 2, 1]: local_sum_q += cute.arch.shuffle_sync_bfly( local_sum_q, offset=offset, mask=-1, mask_and_clamp=31 ) local_sum_k += cute.arch.shuffle_sync_bfly( local_sum_k, offset=offset, mask=-1, mask_and_clamp=31 ) if in_warp_tid == 0: smem_o[0] = cute.rsqrt(local_sum_q + 1e-6) smem_o[1] = cute.rsqrt(local_sum_k + 1e-6) cute.arch.barrier() inv_norm_q = smem_o[0] inv_norm_k = smem_o[1] if tidx < TILE_K: sK[tidx] = sK[tidx] * inv_norm_k sQ[tidx] = sQ[tidx] * scale * inv_norm_q cute.arch.barrier() else: if tidx < TILE_K: sQ[tidx] = sQ[tidx] * scale cute.arch.barrier() for v_tile in range(num_v_tiles): stage = v_tile % NUM_STAGES cute.arch.cp_async_wait_group(0) cute.arch.barrier() next_v_tile = v_tile + prefetch_count if next_v_tile < num_v_tiles: next_stage = next_v_tile % NUM_STAGES gSrc_next = gSrc[(None, None, next_v_tile)] sData_next = sData[(None, None, next_stage)] thr_gSrc = thr_copy_load.partition_S(gSrc_next) thr_sData = thr_copy_load.partition_D(sData_next) cute.copy(tiled_copy_load, thr_gSrc, thr_sData) cute.arch.cp_async_commit_group() v_global = v_tile * TILE_V + v_idx r_v = cutlass.Float32(v[i_n, 0, i_hv, v_global]) sum_hk = 0.0 for k_iter in cutlass.range_dynamic(NUM_K_ITERS, unroll=8): k_base = k_iter * ROWS_PER_ITER k_idx = k_base + k_local h_val = sData[(k_idx, v_idx, stage)] * r_g r_k_val = sK[k_idx] sum_hk += h_val * r_k_val for offset in [4, 2, 1]: sum_hk += cute.arch.shuffle_sync_bfly( sum_hk, offset=offset * V_PER_WARP, mask=-1, mask_and_clamp=31 ) v_new = (r_v - sum_hk) * r_beta v_new = cute.arch.shuffle_sync(v_new, v_local) sum_hq = 0.0 for k_iter in cutlass.range_dynamic(NUM_K_ITERS, unroll=8): k_base = k_iter * ROWS_PER_ITER k_idx = k_base + k_local h_old = sData[(k_idx, v_idx, stage)] * r_g r_k_val = sK[k_idx] r_q_val = sQ[k_idx] h_new = h_old + r_k_val * v_new sData[(k_idx, v_idx, stage)] = h_new sum_hq += h_new * r_q_val for offset in [4, 2, 1]: sum_hq += cute.arch.shuffle_sync_bfly( sum_hq, offset=offset * V_PER_WARP, mask=-1, mask_and_clamp=31 ) if k_local == 0: v_global_out = v_tile * TILE_V + v_idx o[(i_n, 0, i_hv, v_global_out)] = cutlass.BFloat16(sum_hq) cute.arch.barrier() for k_iter in range(NUM_K_ITERS): flat_idx = tidx + k_iter * 256 k_write = flat_idx // TILE_V v_write = flat_idx % TILE_V if k_write < TILE_K: h_val = sData[(k_write, v_write, stage)] v_global_write = v_tile * TILE_V + v_write h0_source[(pool_idx, i_hv, k_write, v_global_write)] = h_val cute.arch.barrier() @cute.kernel def gdn_kernel_large_batch_varlen( tiled_copy_load: cute.TiledCopy, h0_source: cute.Tensor, smem_layout_staged: cute.Layout, num_v_tiles: cutlass.Constexpr[int], q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, a: cute.Tensor, b: cute.Tensor, A_log: cute.Tensor, dt_bias: cute.Tensor, o: cute.Tensor, h0_indices: cute.Tensor, softplus_beta: cutlass.Constexpr[float], softplus_threshold: cutlass.Constexpr[float], scale: cutlass.Constexpr[float], H: cutlass.Constexpr[int], HV: cutlass.Constexpr[int], use_qk_l2norm: cutlass.Constexpr[bool], ): """Large batch kernel for varlen decode (1, N, ...) format.""" tidx, _, _ = cute.arch.thread_idx() in_warp_tid = tidx % 32 warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) batch_idx, _, _ = cute.arch.block_idx() i_n = batch_idx // HV i_hv = batch_idx % HV i_h = i_hv // (HV // H) pool_idx = h0_indices[i_n] if pool_idx >= 0: k_local = in_warp_tid // V_PER_WARP v_local = in_warp_tid % V_PER_WARP v_base = warp_idx * V_PER_WARP v_idx = v_base + v_local smem = cutlass.utils.SmemAllocator() sData = smem.allocate_tensor(cutlass.Float32, smem_layout_staged, 128) smem_o_layout = cute.make_layout((TILE_V,), stride=(1,)) smem_o = smem.allocate_tensor(cutlass.Float32, smem_o_layout, 128) smem_k_layout = cute.make_layout((TILE_K,), stride=(1,)) smem_q_layout = cute.make_layout((TILE_K,), stride=(1,)) sK = smem.allocate_tensor(cutlass.Float32, smem_k_layout, 128) sQ = smem.allocate_tensor(cutlass.Float32, smem_q_layout, 128) if tidx < TILE_K: sK[tidx] = cutlass.Float32(k[0, i_n, i_h, tidx]) sQ[tidx] = cutlass.Float32(q[0, i_n, i_h, tidx]) gSrc_batch = h0_source[(pool_idx, i_hv, None, None)] gSrc = cute.local_tile(gSrc_batch, (TILE_K, TILE_V), (0, None)) thr_copy_load = tiled_copy_load.get_slice(tidx) prefetch_count = cutlass.min(NUM_STAGES - 1, num_v_tiles) for v_tile in range(prefetch_count): stage = v_tile % NUM_STAGES gSrc_tile = gSrc[(None, None, v_tile)] sData_stage = sData[(None, None, stage)] thr_gSrc = thr_copy_load.partition_S(gSrc_tile) thr_sData = thr_copy_load.partition_D(sData_stage) cute.copy(tiled_copy_load, thr_gSrc, thr_sData) cute.arch.cp_async_commit_group() r_A_log = cutlass.Float32(A_log[i_hv]) r_dt_bias = cutlass.Float32(dt_bias[i_hv]) r_a = cutlass.Float32(a[i_n, i_hv]) r_b = cutlass.Float32(b[i_n, i_hv]) r_g = 0.0 r_beta = 0.0 if in_warp_tid == 0: x = r_a + r_dt_bias beta_x = softplus_beta * x softplus_x = 0.0 if beta_x <= softplus_threshold: exp_beta_x = cute.exp(beta_x) log_input = cutlass.Float32(1.0 + exp_beta_x) log_result = cutlass.Float32(cute.log(log_input)) softplus_x = cutlass.Float32( (cutlass.Float32(1.0) / softplus_beta) * log_result ) else: softplus_x = x r_g_value = -cute.exp(r_A_log) * softplus_x r_beta = 1.0 / (1.0 + cute.exp(-r_b)) r_g = cute.exp(r_g_value) r_g = cute.arch.shuffle_sync(r_g, 0) r_beta = cute.arch.shuffle_sync(r_beta, 0) cute.arch.barrier() if use_qk_l2norm: sum_q_partial = 0.0 sum_k_partial = 0.0 if tidx < TILE_K: q_val = sQ[tidx] k_val = sK[tidx] sum_q_partial = q_val * q_val sum_k_partial = k_val * k_val for offset in [16, 8, 4, 2, 1]: sum_q_partial += cute.arch.shuffle_sync_bfly( sum_q_partial, offset=offset, mask=-1, mask_and_clamp=31 ) sum_k_partial += cute.arch.shuffle_sync_bfly( sum_k_partial, offset=offset, mask=-1, mask_and_clamp=31 ) if in_warp_tid == 0: smem_o[warp_idx] = sum_q_partial smem_o[warp_idx + 8] = sum_k_partial cute.arch.barrier() inv_norm_q = 0.0 inv_norm_k = 0.0 if warp_idx == 0: local_sum_q = 0.0 local_sum_k = 0.0 if in_warp_tid < NUM_WARPS_LARGE: local_sum_q = smem_o[in_warp_tid] local_sum_k = smem_o[in_warp_tid + 8] for offset in [4, 2, 1]: local_sum_q += cute.arch.shuffle_sync_bfly( local_sum_q, offset=offset, mask=-1, mask_and_clamp=31 ) local_sum_k += cute.arch.shuffle_sync_bfly( local_sum_k, offset=offset, mask=-1, mask_and_clamp=31 ) if in_warp_tid == 0: smem_o[0] = cute.rsqrt(local_sum_q + 1e-6) smem_o[1] = cute.rsqrt(local_sum_k + 1e-6) cute.arch.barrier() inv_norm_q = smem_o[0] inv_norm_k = smem_o[1] if tidx < TILE_K: sK[tidx] = sK[tidx] * inv_norm_k sQ[tidx] = sQ[tidx] * scale * inv_norm_q cute.arch.barrier() else: if tidx < TILE_K: sQ[tidx] = sQ[tidx] * scale cute.arch.barrier() for v_tile in range(num_v_tiles): stage = v_tile % NUM_STAGES cute.arch.cp_async_wait_group(0) cute.arch.barrier() next_v_tile = v_tile + prefetch_count if next_v_tile < num_v_tiles: next_stage = next_v_tile % NUM_STAGES gSrc_next = gSrc[(None, None, next_v_tile)] sData_next = sData[(None, None, next_stage)] thr_gSrc = thr_copy_load.partition_S(gSrc_next) thr_sData = thr_copy_load.partition_D(sData_next) cute.copy(tiled_copy_load, thr_gSrc, thr_sData) cute.arch.cp_async_commit_group() v_global = v_tile * TILE_V + v_idx r_v = cutlass.Float32(v[0, i_n, i_hv, v_global]) sum_hk = 0.0 for k_iter in cutlass.range_dynamic(NUM_K_ITERS, unroll=8): k_base = k_iter * ROWS_PER_ITER k_idx = k_base + k_local h_val = sData[(k_idx, v_idx, stage)] * r_g r_k_val = sK[k_idx] sum_hk += h_val * r_k_val for offset in [4, 2, 1]: sum_hk += cute.arch.shuffle_sync_bfly( sum_hk, offset=offset * V_PER_WARP, mask=-1, mask_and_clamp=31 ) v_new = (r_v - sum_hk) * r_beta v_new = cute.arch.shuffle_sync(v_new, v_local) sum_hq = 0.0 for k_iter in cutlass.range_dynamic(NUM_K_ITERS, unroll=8): k_base = k_iter * ROWS_PER_ITER k_idx = k_base + k_local h_old = sData[(k_idx, v_idx, stage)] * r_g r_k_val = sK[k_idx] r_q_val = sQ[k_idx] h_new = h_old + r_k_val * v_new sData[(k_idx, v_idx, stage)] = h_new sum_hq += h_new * r_q_val for offset in [4, 2, 1]: sum_hq += cute.arch.shuffle_sync_bfly( sum_hq, offset=offset * V_PER_WARP, mask=-1, mask_and_clamp=31 ) if k_local == 0: v_global_out = v_tile * TILE_V + v_idx o[(0, i_n, i_hv, v_global_out)] = cutlass.BFloat16(sum_hq) cute.arch.barrier() for k_iter in range(NUM_K_ITERS): flat_idx = tidx + k_iter * 256 k_write = flat_idx // TILE_V v_write = flat_idx % TILE_V if k_write < TILE_K: h_val = sData[(k_write, v_write, stage)] v_global_write = v_tile * TILE_V + v_write h0_source[(pool_idx, i_hv, k_write, v_global_write)] = h_val cute.arch.barrier() return ( gdn_kernel_small_batch, gdn_kernel_small_batch_varlen, gdn_kernel_large_batch, gdn_kernel_large_batch_varlen, ) def _create_jit_functions(): """Create JIT-compiled launcher functions for all kernel variants.""" gdn_small, gdn_small_varlen, gdn_large, gdn_large_varlen = _define_kernels() @cute.jit def run_small_batch( cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, a: cute.Tensor, b: cute.Tensor, A_log: cute.Tensor, dt_bias: cute.Tensor, h0_source: cute.Tensor, h0_indices: cute.Tensor, o: cute.Tensor, softplus_beta: cutlass.Constexpr[float], softplus_threshold: cutlass.Constexpr[float], scale: cutlass.Constexpr[float], B: cutlass.Constexpr[int], T: cutlass.Constexpr[int], H: cutlass.Constexpr[int], HV: cutlass.Constexpr[int], K: cutlass.Constexpr[int], V: cutlass.Constexpr[int], use_initial_state: cutlass.Constexpr[bool], use_qk_l2norm: cutlass.Constexpr[bool], stream: cuda.CUstream, ): pool_size, hv_dim, k_dim, v_dim = h0_source.layout.shape n_indices = h0_indices.layout.shape[0] batch_size = n_indices * hv_dim copy_atom = cute.make_copy_atom( cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), cutlass.Float32, num_bits_per_copy=128, ) num_v_tiles_small = cute.ceil_div(v_dim, TILE_V_SMALL) smem_layout_small = cute.make_layout( (TILE_K, TILE_V_SMALL, NUM_STAGES), stride=(TILE_V_SMALL_PADDED, 1, TILE_K * TILE_V_SMALL_PADDED), ) thread_layout_small = cute.make_layout((32, 4), stride=(4, 1)) val_layout_small = cute.make_layout((1, 4)) tiled_copy_load_small = cute.make_tiled_copy_tv( copy_atom, thread_layout_small, val_layout_small ) smem_bytes_small = ( 4 * TILE_K * TILE_V_SMALL_PADDED * NUM_STAGES + 4 * TILE_V_SMALL + 4 * TILE_K * 2 + 64 ) gdn_small( tiled_copy_load_small, h0_source, smem_layout_small, num_v_tiles_small, q, k, v, a, b, A_log, dt_bias, o, h0_indices, softplus_beta, softplus_threshold, scale, H, HV, use_qk_l2norm, ).launch( grid=(batch_size * NUM_BLOCKS_PER_STATE_SMALL, 1, 1), block=[NUM_THREADS, 1, 1], smem=smem_bytes_small, stream=stream, ) @cute.jit def run_small_batch_varlen( cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, a: cute.Tensor, b: cute.Tensor, A_log: cute.Tensor, dt_bias: cute.Tensor, h0_source: cute.Tensor, h0_indices: cute.Tensor, o: cute.Tensor, softplus_beta: cutlass.Constexpr[float], softplus_threshold: cutlass.Constexpr[float], scale: cutlass.Constexpr[float], B: cutlass.Constexpr[int], T: cutlass.Constexpr[int], H: cutlass.Constexpr[int], HV: cutlass.Constexpr[int], K: cutlass.Constexpr[int], V: cutlass.Constexpr[int], use_initial_state: cutlass.Constexpr[bool], use_qk_l2norm: cutlass.Constexpr[bool], stream: cuda.CUstream, ): pool_size, hv_dim, k_dim, v_dim = h0_source.layout.shape n_indices = h0_indices.layout.shape[0] batch_size = n_indices * hv_dim copy_atom = cute.make_copy_atom( cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), cutlass.Float32, num_bits_per_copy=128, ) num_v_tiles_small = cute.ceil_div(v_dim, TILE_V_SMALL) smem_layout_small = cute.make_layout( (TILE_K, TILE_V_SMALL, NUM_STAGES), stride=(TILE_V_SMALL_PADDED, 1, TILE_K * TILE_V_SMALL_PADDED), ) thread_layout_small = cute.make_layout((32, 4), stride=(4, 1)) val_layout_small = cute.make_layout((1, 4)) tiled_copy_load_small = cute.make_tiled_copy_tv( copy_atom, thread_layout_small, val_layout_small ) smem_bytes_small = ( 4 * TILE_K * TILE_V_SMALL_PADDED * NUM_STAGES + 4 * TILE_V_SMALL + 4 * TILE_K * 2 + 64 ) gdn_small_varlen( tiled_copy_load_small, h0_source, smem_layout_small, num_v_tiles_small, q, k, v, a, b, A_log, dt_bias, o, h0_indices, softplus_beta, softplus_threshold, scale, H, HV, use_qk_l2norm, ).launch( grid=(batch_size * NUM_BLOCKS_PER_STATE_SMALL, 1, 1), block=[NUM_THREADS, 1, 1], smem=smem_bytes_small, stream=stream, ) @cute.jit def run_large_batch( cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, a: cute.Tensor, b: cute.Tensor, A_log: cute.Tensor, dt_bias: cute.Tensor, h0_source: cute.Tensor, h0_indices: cute.Tensor, o: cute.Tensor, softplus_beta: cutlass.Constexpr[float], softplus_threshold: cutlass.Constexpr[float], scale: cutlass.Constexpr[float], B: cutlass.Constexpr[int], T: cutlass.Constexpr[int], H: cutlass.Constexpr[int], HV: cutlass.Constexpr[int], K: cutlass.Constexpr[int], V: cutlass.Constexpr[int], use_initial_state: cutlass.Constexpr[bool], use_qk_l2norm: cutlass.Constexpr[bool], stream: cuda.CUstream, ): pool_size, hv_dim, k_dim, v_dim = h0_source.layout.shape n_indices = h0_indices.layout.shape[0] batch_size = n_indices * hv_dim copy_atom = cute.make_copy_atom( cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), cutlass.Float32, num_bits_per_copy=128, ) num_v_tiles = cute.ceil_div(v_dim, TILE_V) base_smem_layout = cute.make_layout( (TILE_K, TILE_V, NUM_STAGES), stride=(TILE_V_PADDED, 1, TILE_K * TILE_V_PADDED), ) thread_layout = cute.make_layout((32, 8), stride=(8, 1)) val_layout = cute.make_layout((1, 4)) tiled_copy_load = cute.make_tiled_copy_tv(copy_atom, thread_layout, val_layout) smem_bytes = ( 4 * TILE_K * TILE_V_PADDED * NUM_STAGES + 4 * TILE_V + 4 * TILE_K * 2 + 64 ) gdn_large( tiled_copy_load, h0_source, base_smem_layout, num_v_tiles, q, k, v, a, b, A_log, dt_bias, o, h0_indices, softplus_beta, softplus_threshold, scale, H, HV, use_qk_l2norm, ).launch( grid=(batch_size, 1, 1), block=[NUM_THREADS_LARGE, 1, 1], smem=smem_bytes, stream=stream, ) @cute.jit def run_large_batch_varlen( cu_seqlens: cute.Tensor, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, a: cute.Tensor, b: cute.Tensor, A_log: cute.Tensor, dt_bias: cute.Tensor, h0_source: cute.Tensor, h0_indices: cute.Tensor, o: cute.Tensor, softplus_beta: cutlass.Constexpr[float], softplus_threshold: cutlass.Constexpr[float], scale: cutlass.Constexpr[float], B: cutlass.Constexpr[int], T: cutlass.Constexpr[int], H: cutlass.Constexpr[int], HV: cutlass.Constexpr[int], K: cutlass.Constexpr[int], V: cutlass.Constexpr[int], use_initial_state: cutlass.Constexpr[bool], use_qk_l2norm: cutlass.Constexpr[bool], stream: cuda.CUstream, ): pool_size, hv_dim, k_dim, v_dim = h0_source.layout.shape n_indices = h0_indices.layout.shape[0] batch_size = n_indices * hv_dim copy_atom = cute.make_copy_atom( cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), cutlass.Float32, num_bits_per_copy=128, ) num_v_tiles = cute.ceil_div(v_dim, TILE_V) base_smem_layout = cute.make_layout( (TILE_K, TILE_V, NUM_STAGES), stride=(TILE_V_PADDED, 1, TILE_K * TILE_V_PADDED), ) thread_layout = cute.make_layout((32, 8), stride=(8, 1)) val_layout = cute.make_layout((1, 4)) tiled_copy_load = cute.make_tiled_copy_tv(copy_atom, thread_layout, val_layout) smem_bytes = ( 4 * TILE_K * TILE_V_PADDED * NUM_STAGES + 4 * TILE_V + 4 * TILE_K * 2 + 64 ) gdn_large_varlen( tiled_copy_load, h0_source, base_smem_layout, num_v_tiles, q, k, v, a, b, A_log, dt_bias, o, h0_indices, softplus_beta, softplus_threshold, scale, H, HV, use_qk_l2norm, ).launch( grid=(batch_size, 1, 1), block=[NUM_THREADS_LARGE, 1, 1], smem=smem_bytes, stream=stream, ) return ( run_small_batch, run_small_batch_varlen, run_large_batch, run_large_batch_varlen, ) _jit_functions = None def _get_jit_functions(): global _jit_functions if _jit_functions is None: _jit_functions = _create_jit_functions() return _jit_functions def _get_compiled_kernel(N, H, HV, K, V, pool_size, use_small_batch, is_varlen_decode): """Get or compile the kernel for given dimensions.""" global _compiled_kernels key = (N, H, HV, K, V, pool_size, use_small_batch, is_varlen_decode) if key in _compiled_kernels: return _compiled_kernels[key] cu_seqlens = torch.zeros(N + 1, dtype=torch.int32, device="cuda") if is_varlen_decode: q = torch.zeros(1, N, H, K, dtype=torch.bfloat16, device="cuda") k = torch.zeros(1, N, H, K, dtype=torch.bfloat16, device="cuda") v = torch.zeros(1, N, HV, V, dtype=torch.bfloat16, device="cuda") a = torch.zeros(N, HV, dtype=torch.bfloat16, device="cuda") b = torch.zeros(N, HV, dtype=torch.bfloat16, device="cuda") o = torch.zeros(1, N, HV, V, dtype=torch.bfloat16, device="cuda") else: q = torch.zeros(N, 1, H, K, dtype=torch.bfloat16, device="cuda") k = torch.zeros(N, 1, H, K, dtype=torch.bfloat16, device="cuda") v = torch.zeros(N, 1, HV, V, dtype=torch.bfloat16, device="cuda") a = torch.zeros(N, 1, HV, dtype=torch.bfloat16, device="cuda") b = torch.zeros(N, 1, HV, dtype=torch.bfloat16, device="cuda") o = torch.zeros(N, 1, HV, V, dtype=torch.bfloat16, device="cuda") A_log = torch.zeros(HV, dtype=torch.float32, device="cuda") dt_bias = torch.zeros(HV, dtype=torch.bfloat16, device="cuda") h0_source = torch.zeros(pool_size, HV, K, V, dtype=torch.float32, device="cuda") h0_indices = torch.zeros(N, dtype=torch.int32, device="cuda") cu_seqlens_tensor = from_dlpack(cu_seqlens, assumed_align=16) q_tensor = from_dlpack(q, assumed_align=16) k_tensor = from_dlpack(k, assumed_align=16) v_tensor = from_dlpack(v, assumed_align=16) a_tensor = from_dlpack(a, assumed_align=16) b_tensor = from_dlpack(b, assumed_align=16) A_log_tensor = from_dlpack(A_log, assumed_align=16) dt_bias_tensor = from_dlpack(dt_bias, assumed_align=16) h0_source_tensor = from_dlpack(h0_source, assumed_align=16) h0_indices_tensor = from_dlpack(h0_indices, assumed_align=16) o_tensor = from_dlpack(o, assumed_align=16) stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) run_small, run_small_varlen, run_large, run_large_varlen = _get_jit_functions() if use_small_batch: kernel_func = run_small_varlen if is_varlen_decode else run_small else: kernel_func = run_large_varlen if is_varlen_decode else run_large scale = K**-0.5 softplus_beta = 1.0 softplus_threshold = 20.0 B_compile = 1 if is_varlen_decode else N T_compile = N if is_varlen_decode else 1 compiled_kernel = cute.compile( kernel_func, cu_seqlens_tensor, q_tensor, k_tensor, v_tensor, a_tensor, b_tensor, A_log_tensor, dt_bias_tensor, h0_source_tensor, h0_indices_tensor, o_tensor, softplus_beta=softplus_beta, softplus_threshold=softplus_threshold, scale=scale, B=B_compile, T=T_compile, H=H, K=K, V=V, HV=HV, use_initial_state=True, use_qk_l2norm=True, stream=stream, ) _compiled_kernels[key] = compiled_kernel logger.info( f"CuTe DSL GDN kernel compiled: N={N}, H={H}, HV={HV}, K={K}, V={V}, pool_size={pool_size}, small_batch={use_small_batch}, varlen={is_varlen_decode}" ) return compiled_kernel def cutedsl_fused_sigmoid_gating_delta_rule_update( A_log: torch.Tensor, dt_bias: torch.Tensor, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, a: torch.Tensor, b: torch.Tensor, initial_state_source: torch.Tensor, initial_state_indices: torch.Tensor, cu_seqlens: Optional[torch.Tensor] = None, scale: Optional[float] = None, use_qk_l2norm_in_kernel: bool = True, softplus_beta: float = 1.0, softplus_threshold: float = 20.0, ) -> torch.Tensor: """CuTe DSL implementation of fused sigmoid gating delta rule update.""" B_q, T_q, H, K = q.shape HV = v.shape[2] V = v.shape[3] N = initial_state_indices.shape[0] is_varlen_decode = B_q == 1 and T_q == N and N > 1 if scale is None: scale = K**-0.5 use_small_batch = N < SMALL_BATCH_THRESHOLD if initial_state_source.dim() == 1: pool_size = initial_state_source.numel() // (HV * K * V) h0_source = initial_state_source.view(pool_size, HV, K, V) elif initial_state_source.dim() == 4: pool_size = initial_state_source.shape[0] h0_source = initial_state_source else: raise ValueError( f"Unexpected initial_state_source shape: {initial_state_source.shape}" ) if is_varlen_decode: if a.dim() == 3: a = a.squeeze(0) if b.dim() == 3: b = b.squeeze(0) o = q.new_empty(1, N, HV, V, dtype=torch.bfloat16) else: if a.dim() == 2: a = a.unsqueeze(1) if b.dim() == 2: b = b.unsqueeze(1) o = q.new_empty(N, 1, HV, V, dtype=torch.bfloat16) q, k, v = [t.contiguous() for t in (q, k, v)] global _cu_seqlens_cache if cu_seqlens is not None: cu_seqlens_to_use = cu_seqlens else: cache_key = (N, str(q.device)) if cache_key not in _cu_seqlens_cache: _cu_seqlens_cache[cache_key] = torch.arange( N + 1, dtype=torch.int32, device=q.device ) cu_seqlens_to_use = _cu_seqlens_cache[cache_key] cu_seqlens_tensor = from_dlpack( cu_seqlens_to_use.detach(), assumed_align=16 ).mark_layout_dynamic(leading_dim=0) q_tensor = from_dlpack(q.detach(), assumed_align=16).mark_layout_dynamic( leading_dim=q.ndim - 1 ) k_tensor = from_dlpack(k.detach(), assumed_align=16).mark_layout_dynamic( leading_dim=k.ndim - 1 ) v_tensor = from_dlpack(v.detach(), assumed_align=16).mark_layout_dynamic( leading_dim=v.ndim - 1 ) a_tensor = from_dlpack(a.detach(), assumed_align=16).mark_layout_dynamic( leading_dim=a.ndim - 1 ) b_tensor = from_dlpack(b.detach(), assumed_align=16).mark_layout_dynamic( leading_dim=b.ndim - 1 ) A_log_tensor = from_dlpack(A_log.detach(), assumed_align=16).mark_layout_dynamic( leading_dim=0 ) dt_bias_tensor = from_dlpack( dt_bias.detach(), assumed_align=16 ).mark_layout_dynamic(leading_dim=0) h0_source_tensor = from_dlpack( h0_source.detach(), assumed_align=16 ).mark_layout_dynamic(leading_dim=h0_source.ndim - 1) h0_indices_tensor = from_dlpack( initial_state_indices.detach(), assumed_align=16 ).mark_layout_dynamic(leading_dim=0) o_tensor = from_dlpack(o.detach(), assumed_align=16).mark_layout_dynamic( leading_dim=o.ndim - 1 ) stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) compiled_kernel = _get_compiled_kernel( N, H, HV, K, V, pool_size, use_small_batch, is_varlen_decode ) compiled_kernel( cu_seqlens_tensor, q_tensor, k_tensor, v_tensor, a_tensor, b_tensor, A_log_tensor, dt_bias_tensor, h0_source_tensor, h0_indices_tensor, o_tensor, stream, ) return o ================================================ FILE: python/sglang/jit_kernel/diffusion/cutedsl/common/norm_fusion.py ================================================ from typing import Optional, Tuple, Union import cutlass import cutlass.cute as cute import torch from einops import rearrange from sglang.jit_kernel.diffusion.cutedsl.common.reduce import ( cta_reduce_sum, warp_reduce_sum, ) @cute.jit def apply_norm_cta( norm_type: cutlass.Constexpr, num_warps: cutlass.Constexpr, tidx: cutlass.Int32, tXrX: cute.Tensor, tWrW: Optional[cute.Tensor], tBrB: Optional[cute.Tensor], D: Union[cutlass.Int32, cutlass.Constexpr], eps: Union[cutlass.Float32, cutlass.Constexpr], ) -> cute.Tensor: if cutlass.const_expr(norm_type == "rms"): return apply_rmsnorm_cta(num_warps, tidx, tXrX, tWrW, D, eps) else: return apply_layernorm_cta(num_warps, tidx, tXrX, tWrW, tBrB, D, eps) @cute.jit def apply_rmsnorm_cta( num_warps: Union[cutlass.Int32, cutlass.Constexpr], tidx: cutlass.Int32, tXrX: cute.Tensor, tWrW: Optional[cute.Tensor], D: Union[cutlass.Int32, cutlass.Constexpr], eps: Union[cutlass.Float32, cutlass.Constexpr], ) -> cute.Tensor: """ RMSNorm: y[i] = x[i] / sqrt(sum(x ^ 2) / D + eps) * w[i] """ val = cute.Float32(0.0) for idx in range(cute.size(tXrX)): # Accumulate in FP32 to improve numerical precision. x_fp32 = tXrX[idx].to(cutlass.Float32) val += x_fp32 * x_fp32 val = warp_reduce_sum(val) acc_sq = cta_reduce_sum(val, num_warps, tidx) factor = cute.rsqrt(acc_sq / D + eps) tNrN = cute.make_fragment_like(tXrX) if cutlass.const_expr(isinstance(tWrW, cute.Tensor)): tNrN.store((tXrX.load() * factor * tWrW.load()).to(tNrN.element_type)) else: tNrN.store((tXrX.load() * factor).to(tNrN.element_type)) return tNrN @cute.jit def apply_layernorm_cta( num_warps: Union[cutlass.Int32, cutlass.Constexpr], tidx: cutlass.Int32, tXrX: cute.Tensor, tWrW: Optional[cute.Tensor], tBrB: Optional[cute.Tensor], D: Union[cutlass.Int32, cutlass.Constexpr], eps: Union[cutlass.Float32, cutlass.Constexpr], ) -> cute.Tensor: """ LayerNorm: mean = sum(x) / D var = sum((x - mean) ^ 2) / D y[i] = (x[i] - mean) / sqrt(var + eps) * w[i] + b[i] """ # Reduce mean val = cute.Float32(0.0) for idx in range(cute.size(tXrX)): # Accumulate in FP32 to improve numerical precision. val += tXrX[idx].to(cutlass.Float32) val = warp_reduce_sum(val) val = cta_reduce_sum(val, num_warps, tidx) mean = val / D # Reduce variance val = cute.Float32(0.0) for idx in range(cute.size(tXrX)): # Accumulate in FP32 to improve numerical precision. x_fp32 = tXrX[idx].to(cutlass.Float32) val += (x_fp32 - mean) * (x_fp32 - mean) val = warp_reduce_sum(val) val = cta_reduce_sum(val, num_warps, tidx) factor = cute.rsqrt(val / D + eps) # Normalize tNrN = cute.make_fragment_like(tXrX) if cutlass.const_expr( isinstance(tWrW, cute.Tensor) and isinstance(tBrB, cute.Tensor) ): tNrN.store( ((tXrX.load() - mean) * factor * tWrW.load() + tBrB.load()).to( tNrN.element_type ) ) else: tNrN.store(((tXrX.load() - mean) * factor).to(tNrN.element_type)) return tNrN ################################################################################ # BSFD Indexing ################################################################################ # In diffusion norm-fusion kernels, we compute `norm(x) + y`, where # `x` has shape [B, S, D] and `y` may come in various broadcastable forms: # [1], [D], [1, D], [1, 1, D], [B, D], [B, 1, D], [B, S, D], or [B, F, 1, D]. # # For a given (batch_id, seq_id), the index mapping for `y` falls into 3 cases: # 1) Scalar broadcast [1]: # (batch_id, seq_id, *) -> (0) # 2) Frame-based BSFD broadcast [B, F, 1, D]: # frame_id = seq_id // len_frame # (batch_id, seq_id, *) -> (batch_id, frame_id, *) # 3) All other cases: # `y` is broadcast to [B, S, D] (via view/expand, no materialization), # and indexed as (batch_id, seq_id, *). # # This helper normalizes `y` into a BSFD-compatible view so that kernel # indexing logic remains simple and uniform. ################################################################################ def broadcast_tensor_for_bsfd( tensor: Union[Optional[torch.Tensor], int], B: int, S: int, D: int, ) -> Union[Optional[torch.Tensor], int]: """ Broadcast to (B, S, D) without memory copy for following shapes: - [D], [1, D], [1, 1, D], [B, D], [B, 1, D], [B, S, D]. """ # Return directly for non-tensor value if not isinstance(tensor, torch.Tensor): return tensor if tensor.ndim == 1: # Scalar [1] is preserved as-is and handled specially in CuTe kernel. if tensor.numel() == 1: return tensor return rearrange(tensor, "d -> 1 1 d").expand(B, S, D) if tensor.ndim == 2: return rearrange(tensor, "b d -> b 1 d").expand(B, S, D) if tensor.ndim == 3: return tensor.expand(B, S, D) if tensor.ndim == 4: return tensor raise ValueError(f"BSFD broadcast: unsupported tensor ndim: {tensor.ndim}.") @cute.jit def tensor_slice_for_bsfd( mV: cute.Tensor, thr_copy: cute.ThrCopy, batch_id: cutlass.Int32, seq_id: cutlass.Int32, S: Union[cutlass.Int32, cutlass.Constexpr], D: Union[cutlass.Int32, cutlass.Constexpr], ) -> Tuple[cute.Tensor, cute.Tensor]: """ Slice a BSFD-compatible tensor into a per-thread gmem tile and rmem fragment. Given a logical (batch_id, seq_id), this helper selects the corresponding D-length slice from `mV` and prepares it for vectorized copy. """ gV: cute.Tensor if cutlass.const_expr(cute.is_static(mV.layout) and cute.size(mV.layout) == 1): # build a ((1,1),(1,)) layout so it could broadcast-align with the # regular rmem fragment shape ((4,1),(k,)). layout = cute.make_layout(shape=((1, 1), (1,))) tVgV = cute.make_tensor(mV.iterator, layout) tVrV = cute.make_rmem_tensor(layout, mV.element_type) return tVgV, tVrV # Use `local_tile` instead of direct indexing to preserve gmem base pointer # alignment required for vectorized loads. if cutlass.const_expr(len(mV.shape) == 1): gV = mV elif cutlass.const_expr(len(mV.shape) == 3): gV = cute.local_tile(mV, tiler=(1, 1, D), coord=(batch_id, seq_id, 0)) gV = gV[0, 0, None] elif cutlass.const_expr(len(mV.shape) == 4): # Compute frame length at runtime (instead of compile time) to avoid # specializing kernels on the frame dimension. frame_len = S // mV.shape[1] frame_id = seq_id // frame_len gV = cute.local_tile(mV, tiler=(1, 1, 1, D), coord=(batch_id, frame_id, 0, 0)) gV = gV[0, 0, 0, None] else: raise NotImplementedError(f"BSFD slice: unsupported shape {mV.shape}.") tVgV = thr_copy.partition_S(gV) tVrV = cute.make_fragment_like(tVgV, tVgV.element_type) return tVgV, tVrV ================================================ FILE: python/sglang/jit_kernel/diffusion/cutedsl/common/reduce.py ================================================ import math import cutlass import cutlass.cute as cute @cute.jit def warp_reduce_sum(val: cute.Numeric, reduce_size: int = 32) -> cute.Numeric: iters = int(math.log2(reduce_size)) for i in range(iters): val = val + cute.arch.shuffle_sync_down(val, offset=1 << (iters - i - 1)) return val @cute.jit def cta_reduce_sum( val: cute.Numeric, num_warps: cutlass.Constexpr, tidx: cutlass.Int32 ) -> cute.Numeric: smem = cutlass.utils.SmemAllocator() acc = smem.allocate_tensor(cutlass.Float32, num_warps + 1) warp_id = tidx >> 5 lane_id = tidx & 31 if lane_id == 0: acc[warp_id] = val cute.arch.sync_threads() if warp_id == 0: val = acc[lane_id] if lane_id < num_warps else cutlass.Float32(0) val = warp_reduce_sum(val) if lane_id == 0: acc[num_warps] = val cute.arch.sync_threads() val = acc[num_warps] return val ================================================ FILE: python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py ================================================ from typing import Optional, Tuple, Union import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import torch from sglang.jit_kernel.diffusion.cutedsl.common.norm_fusion import ( apply_norm_cta, broadcast_tensor_for_bsfd, tensor_slice_for_bsfd, ) from sglang.jit_kernel.diffusion.cutedsl.utils import TORCH_TO_CUTE_DTYPE, WARP_SIZE _COMPILE_CACHE = {} def to_cute_arg( t, *, assume_aligned: Optional[int] = 32, use_32bit_stride: bool = False, enable_tvm_ffi: bool = True, ): """ Convert a Python value into a CuTeDSL value. """ if isinstance(t, torch.Tensor): return cute.runtime.from_dlpack( t, assumed_align=assume_aligned, use_32bit_stride=use_32bit_stride, enable_tvm_ffi=enable_tvm_ffi, ) if isinstance(t, int): return cutlass.Int32(t) if isinstance(t, float): return cutlass.Float32(t) return t def to_fake_cute_args(t: torch.Tensor): if isinstance(t, torch.Tensor): # Only keep the last dim as compile-time value to maximum compiled kernel reuse # e.g. (1,2,1536):(3027,1536,1) -> (?,?,1536):(?,?,1) D = t.shape[-1] dtype = TORCH_TO_CUTE_DTYPE[t.dtype] shape = (*(cute.sym_int() for _ in range(t.ndim - 1)), D) stride = (*(cute.sym_int(divisibility=D) for _ in range(t.ndim - 1)), 1) fake_t = cute.runtime.make_fake_tensor( dtype, shape, stride, memspace=cute.AddressSpace.gmem, assumed_align=32 ) return fake_t return to_cute_arg(t) class ScaleResidualNormScaleShift: @classmethod def make_hash_key(cls, *inputs): """ Compile-time values: - D: hidden dimension (size of the last dimension) - norm_type: layer norm or RMS norm - tensor dtype - tensor rank (i.e., tensor.ndim) Runtime values: - all other inputs This hash key defines the compile-time specialization boundary for ScaleResidualNormScaleShift kernels. """ def _sig(val): if isinstance(val, torch.Tensor): return (val.dtype, val.ndim, val.shape[-1]) return val return tuple(_sig(val) for val in inputs) def __init__(self, D: int, norm_type: str): self.D = D self.norm_type = norm_type # "layer" or "rms" self.num_warps = self.D // 256 # num of warps per cta self.num_threads = self.num_warps * WARP_SIZE # num of threads per cta @cute.jit def __call__( self, mY, mResOut, mRes, mX, mGate, mWeight, mBias, mScale, mShift, eps: cutlass.Float32 = cutlass.Float32(1e-5), stream: cuda.CUstream = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT), ): # Tensor shapes B, S, _ = mX.shape # (batch, seq_len, hidden_dim) # Vectorized copy configuration num_vectorized = 8 # maximum num of elem per copy atom_copy = cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), mX.element_type, num_bits_per_copy=128, ) # Thread/value layouts for tiled copy t_layout = cute.make_layout(self.num_threads) # thread layout within a CTA v_layout = cute.make_layout(num_vectorized) # per-thread vector layout tiled_copy = cute.make_tiled_copy_tv(atom_copy, t_layout, v_layout) self.kernel( mY, mResOut, mRes, mX, mGate, mWeight, mBias, mScale, mShift, tiled_copy, eps, ).launch( grid=[B * S, 1, 1], block=[self.num_threads, 1, 1], stream=stream, ) @cute.kernel def kernel( self, mY, mResOut, mRes, mX, mGate, mWeight, mBias, mScale, mShift, tiled_copy: cute.TiledCopy, eps: cutlass.Float32, ): _, S, _ = mX.shape tidx, _, _ = cute.arch.thread_idx() # thread index bid, _, _ = cute.arch.block_idx() # cta index bidx = cutlass.Int32(bid // S) # batch index bidy = cutlass.Int32(bid % S) # seq_len index thr_copy = tiled_copy.get_slice(tidx) @cute.jit def slice_if(mV): if cutlass.const_expr(isinstance(mV, cute.Tensor)): return tensor_slice_for_bsfd(mV, thr_copy, bidx, bidy, S, self.D) return mV, mV @cute.jit def copy_if(src, dst): if cutlass.const_expr( isinstance(src, cute.Tensor) and isinstance(dst, cute.Tensor) ): cute.autovec_copy(src, dst) # LDG.128 @cute.jit def norm(x, weight, bias): return apply_norm_cta( self.norm_type, self.num_warps, tidx, x, weight, bias, self.D, eps ) # Slice: retrieve the per-thread data slices for both global memory (gmem) # and register memory (rmem). The layouts are: # - ((4,2),(1)):((1,4),(0)) for fp32 # - ((8,1),(1)):((1,0),(0)) for fp16/bf16 tRgR, tRrR = slice_if(mRes) # residual tXgX, tXrX = slice_if(mX) # x tGgG, tGrG = slice_if(mGate) # gate tROgRO, tROrRO = slice_if(mResOut) # residual_out tWgW, tWrW = slice_if(mWeight) # weight tBgB, tBrB = slice_if(mBias) # bias tSCgSC, tSCrSC = slice_if(mScale) # scale tSHgSH, tSHrSH = slice_if(mShift) # shift tYgY, tYrY = slice_if(mY) # y # Load: load tensor from global memory to registers copy_if(tRgR, tRrR) # gmem -> rmem copy_if(tXgX, tXrX) # gmem -> rmem copy_if(tGgG, tGrG) # gmem -> rmem copy_if(tWgW, tWrW) # gmem -> rmem copy_if(tBgB, tBrB) # gmem -> rmem # For norm_scale_shift, output: # - y = norm(x, weight, bias) * (1 + scale) + shift # For scale_residual_norm_scale_shift, output: # - residual_out = residual + gate * x # - y = norm(residual_out, weight, bias) * (1 + scale) + shift # Compute: value = * x value = tXrX.load() if cutlass.const_expr(isinstance(tGrG, cute.Tensor)): value = tGrG.load() * value # Compute: value = value + if cutlass.const_expr(isinstance(tRrR, cute.Tensor)): value = value + tRrR.load() # Store: residual_out if cutlass.const_expr(isinstance(tROrRO, cute.Tensor)): tROrRO.store(value.to(tROrRO.element_type)) copy_if(tROrRO, tROgRO) # rmem -> gmem # Compute: value = norm(value) * + tNrN = cute.make_rmem_tensor_like(tXrX, tXrX.element_type) tNrN.store(value.to(tNrN.element_type)) tNrN = norm(tNrN, tWrW, tBrB) # Compute: value = value * (1 + ) + value = tNrN.load() copy_if(tSCgSC, tSCrSC) # gmem -> rmem copy_if(tSHgSH, tSHrSH) # gmem -> rmem if cutlass.const_expr(isinstance(tSCrSC, cute.Tensor)): value = value * (1 + tSCrSC.load()) if cutlass.const_expr(isinstance(tSHrSH, cute.Tensor)): value = value + tSHrSH.load() # Store: y tYrY.store(value.to(tYrY.element_type)) copy_if(tYrY, tYgY) # rmem -> gmem def validate_x(t: torch.Tensor, B: int, S: int, D: int): if t.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise ValueError(f"Validate failed: unsupported dtype: {t.dtype}") if t.shape != (B, S, D): raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.") if t.stride()[-1] != 1: raise ValueError(f"Validate failed: not contiguous on dim D.") def validate_weight_bias(t: Optional[torch.Tensor], B: int, S: int, D: int): if t is None: return if t.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise ValueError(f"Validate failed: unsupported dtype: {t.dtype}") if t.shape != (D,): raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.") if t.stride()[-1] != 1: raise ValueError(f"Validate failed: not contiguous on dim D.") def validate_scale_shift(t: torch.Tensor, B: int, S: int, D: int): if t.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise ValueError(f"Validate failed: unsupported dtype: {t.dtype}") failed = False if t.ndim == 1 and (t.shape[0] not in (1, D)): failed = True elif t.ndim == 2 and ((t.shape[0] not in (1, B)) or t.shape[1] != D): failed = True elif t.ndim == 3 and ( (t.shape[0] not in (1, B)) or (t.shape[1] not in (1, S) or t.shape[2] != D) ): failed = True elif t.ndim == 4 and (t.shape[0] != B or t.shape[2] != 1 or t.shape[3] != D): F = t.shape[1] if S % F != 0: raise ValueError(f"Validate failed: S({S}) must be divisible by F({F}).") failed = True if failed: raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.") if t.stride()[-1] != 1: raise ValueError(f"Validate failed: not contiguous on dim D.") def validate_gate(t: Union[torch.Tensor, int], B: int, S: int, D: int): if not isinstance(t, torch.Tensor): return validate_scale_shift(t, B, S, D) @torch.library.custom_op("sglang::fused_norm_scale_shift", mutates_args=()) def fused_norm_scale_shift( x: torch.Tensor, weight: Optional[torch.Tensor], bias: Optional[torch.Tensor], scale: torch.Tensor, shift: torch.Tensor, norm_type: str, eps: float = 1e-5, ) -> torch.Tensor: """ Fuse: norm(x) * (1 + scale) + shift where norm is either layernorm or rmsnorm. Expects: - x: [B, S, D] - weight/bias: None, [D] - scale/shift: [1], [D], [1/B, D], [1/B, 1/S, D] or [B, F, 1, D] - norm_type: str, "layer" or "rms" - eps: Optional[float], default: 1e-5 D must be a multiple of 256 and <= 8192 to enable LDG.128 vectorized loads per thread and avoid predicated loads (e.g., bounds checks such as `index < D`). """ stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) # Tensor Validation BSD = x.shape validate_x(x, *BSD) validate_weight_bias(weight, *BSD) validate_weight_bias(bias, *BSD) validate_scale_shift(scale, *BSD) validate_scale_shift(shift, *BSD) if norm_type == "layer" or norm_type == "rms": D = x.shape[-1] if D % 256 != 0 or D > 8192: raise ValueError( f"D={D} not supported, must be multiple of 256 and <= 8192" ) y = torch.empty_like(x) # create output tensor scale = broadcast_tensor_for_bsfd(scale, *x.shape) # handle various shapes shift = broadcast_tensor_for_bsfd(shift, *x.shape) # handle various shapes # Use scalar placeholders for None tensors as a workaround, since the CuTe DSL # TVM-FFI backend does not support None parameters. scalar values do not result # in code generation and have no impact on runtime performance. weight = 1 if weight is None else weight bias = 0 if bias is None else bias ResOut, Residual, Gate = 0, 0, 1 torch_tensors = [y, ResOut, Residual, x, Gate, weight, bias, scale, shift] # Compile cache hash_key = ScaleResidualNormScaleShift.make_hash_key(norm_type, *torch_tensors) compiled_fn = _COMPILE_CACHE.get(hash_key) if compiled_fn is None: kernel = ScaleResidualNormScaleShift(D, norm_type) fake_sig_args = [to_fake_cute_args(t) for t in torch_tensors] compiled_fn = cute.compile( kernel, *fake_sig_args, options="--enable-tvm-ffi" ) _COMPILE_CACHE[hash_key] = compiled_fn # Execute compiled_fn(*torch_tensors, eps, stream) return y else: raise ValueError(f'norm_type must be one of "layer" and "rms"') @fused_norm_scale_shift.register_fake def _fused_norm_scale_shift_fake(x, weight, bias, scale, shift, norm_type, eps=1e-5): y = x.new_empty(x.shape) return y @torch.library.custom_op( "sglang::fused_scale_residual_norm_scale_shift", mutates_args=() ) def fused_scale_residual_norm_scale_shift( residual: torch.Tensor, x: torch.Tensor, gate: Optional[torch.Tensor], # Union[Optional[torch.Tensor], int] indeed weight: Optional[torch.Tensor], bias: Optional[torch.Tensor], scale: torch.Tensor, shift: torch.Tensor, norm_type: str, eps: float = 1e-5, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Fuse: norm(residual + gate * x) * (1 + scale) + shift where norm is either layernorm or rmsnorm. Expects: - residual, x: [B, S, D] - gate: None, [1], [D], [1/B, D], [1/B, 1/S, D] or [B, F, 1, D] - weight/bias: None, [D] - scale/shift: [1], [D], [1/B, D], [1/B, 1/S, D] or [B, F, 1, D] - norm_type: str, "layer" or "rms" - eps: Optional[float], default: 1e-5 D must be a multiple of 256 and <= 8192 to enable LDG.128 vectorized loads per thread and avoid predicated loads (e.g., bounds checks such as `index < D`). """ # Tensor Validation BSD = x.shape validate_x(x, *BSD) validate_x(residual, *BSD) validate_gate(gate, *BSD) validate_weight_bias(weight, *BSD) validate_weight_bias(bias, *BSD) validate_scale_shift(scale, *BSD) validate_scale_shift(shift, *BSD) if norm_type == "layer" or norm_type == "rms": stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) # if norm_type == "layer" or norm_type == "rms": D = x.shape[-1] if D % 256 != 0 or D > 8192: raise ValueError( f"D={D} not supported, must be multiple of 256 and <= 8192" ) y = torch.empty_like(x) # create output tensor resi_out = torch.empty_like(x) # create output tensor gate = broadcast_tensor_for_bsfd(gate, *x.shape) # handle various shapes scale = broadcast_tensor_for_bsfd(scale, *x.shape) # handle various shapes shift = broadcast_tensor_for_bsfd(shift, *x.shape) # handle various shapes # Use scalar placeholders for None tensors as a workaround, since the CuTe DSL # TVM-FFI backend does not support None parameters. scalar values do not result # in code generation and have no impact on runtime performance. gate = 1 if gate is None else gate weight = 1 if weight is None else weight bias = 0 if bias is None else bias torch_tensors = [y, resi_out, residual, x, gate, weight, bias, scale, shift] # Compile cache hash_key = ScaleResidualNormScaleShift.make_hash_key(norm_type, *torch_tensors) compiled_fn = _COMPILE_CACHE.get(hash_key) if compiled_fn is None: kernel = ScaleResidualNormScaleShift(D, norm_type) fake_sig_args = [to_fake_cute_args(t) for t in torch_tensors] compiled_fn = cute.compile( kernel, *fake_sig_args, options="--enable-tvm-ffi" ) _COMPILE_CACHE[hash_key] = compiled_fn # Execute compiled_fn(*torch_tensors, eps, stream) return y, resi_out else: raise ValueError(f'norm_type must be one of "layer" and "rms"') @fused_scale_residual_norm_scale_shift.register_fake def _fused_scale_residual_norm_scale_shift_fake( residual, x, gate, weight, bias, scale, shift, norm_type, eps=1e-5 ): y = x.new_empty(x.shape) residual_out = x.new_empty(x.shape) return y, residual_out ================================================ FILE: python/sglang/jit_kernel/diffusion/cutedsl/utils.py ================================================ import cutlass import torch WARP_SIZE = 32 TORCH_TO_CUTE_DTYPE = { torch.float16: cutlass.Float16, torch.bfloat16: cutlass.BFloat16, torch.float32: cutlass.Float32, } ================================================ FILE: python/sglang/jit_kernel/diffusion/triton/mps_fallback.py ================================================ """MPS (Apple Silicon) fallbacks for Triton diffusion kernels. Triton is not available on macOS / Metal, so these pure-PyTorch (and optionally MLX-accelerated) implementations replace the Triton kernels at import time when ``current_platform.is_mps()`` is True. MLX acceleration (opt-in via ``SGLANG_USE_MLX=1``): Norm ops use ``mx.fast.rms_norm`` / ``mx.fast.layer_norm`` — single fused Metal kernels that are 1.4x–2.9x faster than the multi-step PyTorch MPS decomposition for medium-to-large tensors. """ from typing import Optional import torch from torch import Tensor from sglang.srt.environ import envs # MLX acceleration – opt-in via SGLANG_USE_MLX=1 _MLX_AVAILABLE = False try: import mlx.core as mx _MLX_AVAILABLE = True except ImportError: pass _USE_MLX = envs.SGLANG_USE_MLX.get() and _MLX_AVAILABLE # Dtype mapping for torch <-> MLX tensor bridge _TORCH_TO_MLX_DTYPE = ( { torch.float32: mx.float32, torch.float16: mx.float16, torch.bfloat16: mx.bfloat16, } if _MLX_AVAILABLE else {} ) _MLX_TO_TORCH_DTYPE = {v: k for k, v in _TORCH_TO_MLX_DTYPE.items()} def _torch_to_mlx(tensor: torch.Tensor) -> "mx.array": """Convert a PyTorch tensor to an MLX array (via numpy on CPU).""" t = tensor.cpu().detach() if t.dtype == torch.bfloat16: return mx.array(t.float().numpy(), dtype=mx.bfloat16) return mx.array(t.numpy()) def _mlx_to_torch(array: "mx.array", device: torch.device) -> torch.Tensor: """Convert an MLX array to a PyTorch tensor (zero-copy via memoryview).""" torch_dtype = _MLX_TO_TORCH_DTYPE.get(array.dtype, torch.float32) array = mx.contiguous(array) mx.eval(array) tensor = torch.frombuffer(memoryview(array), dtype=torch_dtype).reshape(array.shape) if device.type == "mps": tensor = tensor.to(device) return tensor def fuse_scale_shift_kernel_native( x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, scale_constant: float = 1.0, block_l: int = 128, block_c: int = 128, ): """Native fallback for fuse_scale_shift_kernel with scale_constant support.""" B, L, C = x.shape def _expand(t: torch.Tensor) -> torch.Tensor: if t.dim() == 4: # [B, F, 1, C] -> [B, L, C] num_frames = t.shape[1] frame_seqlen = L // num_frames return ( t.squeeze(2) .unsqueeze(2) .expand(-1, -1, frame_seqlen, -1) .reshape(B, L, C) ) elif t.dim() == 2: # [B, C] -> [B, 1, C] return t.unsqueeze(1) return t scale = _expand(scale) shift = _expand(shift) return x * (scale_constant + scale) + shift def fuse_scale_shift_gate_select01_kernel_native( x: torch.Tensor, scale0: torch.Tensor, shift0: torch.Tensor, gate0: torch.Tensor, scale1: torch.Tensor, shift1: torch.Tensor, gate1: torch.Tensor, index: torch.Tensor, block_l: int = 128, block_c: int = 128, ): """Native fallback for fuse_scale_shift_gate_select01_kernel.""" idx = index.unsqueeze(-1).bool() scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1)) shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1)) gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1)) y = x * (1 + scale) + shift return y, gate def apply_rotary_embedding_native( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False ) -> torch.Tensor: """Native fallback for rotary embedding (shared with NPU implementation).""" cos = cos.unsqueeze(-2).to(x.dtype) sin = sin.unsqueeze(-2).to(x.dtype) x1 = x[..., ::2] x2 = x[..., 1::2] o1 = x1 * cos - x2 * sin o2 = x2 * cos + x1 * sin return torch.stack((o1, o2), dim=-1).flatten(-2) def norm_infer_native( x: Tensor, weight: Optional[Tensor], bias: Optional[Tensor], eps: float, is_rms_norm: bool = False, out: Optional[Tensor] = None, ) -> Tensor: """Native fallback for norm_infer (layer norm / rms norm inference).""" orig_dtype = x.dtype x = x.contiguous().float() if is_rms_norm: variance = x.pow(2).mean(dim=-1, keepdim=True) x_hat = x * torch.rsqrt(variance + eps) else: mean = x.mean(dim=-1, keepdim=True) variance = (x - mean).pow(2).mean(dim=-1, keepdim=True) x_hat = (x - mean) * torch.rsqrt(variance + eps) if weight is not None: x_hat = x_hat * weight.float() if bias is not None: x_hat = x_hat + bias.float() result = x_hat.to(orig_dtype) if out is not None: out.copy_(result) return out return result def triton_one_pass_rms_norm_native( x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6 ) -> torch.Tensor: """Native fallback for triton_one_pass_rms_norm.""" shape = x.shape orig_dtype = x.dtype x = x.contiguous().float() variance = x.pow(2).mean(dim=-1, keepdim=True) x_hat = x * torch.rsqrt(variance + eps) return (x_hat * w.float()).to(orig_dtype).view(shape) def rms_norm_fn_native( x, weight, bias, residual=None, x1=None, weight1=None, bias1=None, eps=1e-6, dropout_p=0.0, rowscale=None, prenorm=False, residual_in_fp32=False, zero_centered_weight=False, return_dropout_mask=False, out_dtype=None, out=None, residual_out=None, ): """Native fallback for rms_norm_fn (inference only, no dropout/x1 support).""" x_shape_og = x.shape orig_dtype = x.dtype x = x.reshape(-1, x.shape[-1]).float() if residual is not None: residual = residual.reshape(-1, residual.shape[-1]).float() x = x + residual residual_out_val = x.to(torch.float32 if residual_in_fp32 else orig_dtype) else: residual_out_val = None variance = x.pow(2).mean(dim=-1, keepdim=True) x_hat = x * torch.rsqrt(variance + eps) if weight is not None: w = weight.float() if zero_centered_weight: w = w + 1.0 x_hat = x_hat * w if bias is not None: x_hat = x_hat + bias.float() final_dtype = out_dtype if out_dtype is not None else orig_dtype y = x_hat.to(final_dtype).reshape(x_shape_og) if residual is not None and residual_out_val is not None: return y, residual_out_val.reshape(x_shape_og) return y # MLX-accelerated norm ops (1.4x–2.9x faster than torch native on MPS) # Uses mx.fast.rms_norm / mx.fast.layer_norm — single fused Metal kernels # instead of 7+ separate PyTorch MPS kernel launches. if _USE_MLX: def norm_infer_native( # noqa: F811 x: Tensor, weight: Optional[Tensor], bias: Optional[Tensor], eps: float, is_rms_norm: bool = False, out: Optional[Tensor] = None, ) -> Tensor: """MLX-accelerated norm_infer (layer norm / rms norm inference).""" device = x.device orig_dtype = x.dtype x_mx = _torch_to_mlx(x) if is_rms_norm: w_mx = ( _torch_to_mlx(weight) if weight is not None else mx.ones(x_mx.shape[-1]) ) result_mx = mx.fast.rms_norm(x_mx, w_mx, eps) else: w_mx = _torch_to_mlx(weight) if weight is not None else None b_mx = _torch_to_mlx(bias) if bias is not None else None result_mx = mx.fast.layer_norm(x_mx, w_mx, b_mx, eps) result = _mlx_to_torch(result_mx, device).to(orig_dtype) if out is not None: out.copy_(result) return out return result def triton_one_pass_rms_norm_native( # noqa: F811 x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6 ) -> torch.Tensor: """MLX-accelerated triton_one_pass_rms_norm.""" shape = x.shape device = x.device orig_dtype = x.dtype x_mx = _torch_to_mlx(x.reshape(-1, x.shape[-1])) w_mx = _torch_to_mlx(w) result_mx = mx.fast.rms_norm(x_mx, w_mx, eps) return _mlx_to_torch(result_mx, device).to(orig_dtype).view(shape) def rms_norm_fn_native( # noqa: F811 x, weight, bias, residual=None, x1=None, weight1=None, bias1=None, eps=1e-6, dropout_p=0.0, rowscale=None, prenorm=False, residual_in_fp32=False, zero_centered_weight=False, return_dropout_mask=False, out_dtype=None, out=None, residual_out=None, ): """MLX-accelerated rms_norm_fn (inference only, no dropout/x1 support).""" x_shape_og = x.shape device = x.device orig_dtype = x.dtype x_flat = x.reshape(-1, x.shape[-1]) if residual is not None: residual = residual.reshape(-1, residual.shape[-1]).float() x_flat = x_flat.float() + residual residual_out_val = x_flat.to( torch.float32 if residual_in_fp32 else orig_dtype ) else: residual_out_val = None if weight is not None and zero_centered_weight: w = weight.float() + 1.0 else: w = weight x_mx = _torch_to_mlx(x_flat) w_mx = _torch_to_mlx(w) if w is not None else mx.ones(x_mx.shape[-1]) result_mx = mx.fast.rms_norm(x_mx, w_mx, eps) x_hat = _mlx_to_torch(result_mx, device) if bias is not None: x_hat = x_hat + bias.to(x_hat.device, x_hat.dtype) final_dtype = out_dtype if out_dtype is not None else orig_dtype y = x_hat.to(final_dtype).reshape(x_shape_og) if residual is not None and residual_out_val is not None: return y, residual_out_val.reshape(x_shape_og) return y ================================================ FILE: python/sglang/jit_kernel/diffusion/triton/norm.py ================================================ from typing import Optional import torch import triton # type: ignore import triton.language as tl # type: ignore from torch import Tensor # RMSNorm-fp32 def maybe_contiguous_lastdim(x): return x.contiguous() if x is not None and x.stride(-1) != 1 else x def maybe_contiguous(x): return x.contiguous() if x is not None else None def triton_autotune_configs(): # Return configs with a valid warp count for the current device configs = [] # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024 max_threads_per_block = 1024 # Default to warp size 32 if not defined by device warp_size = getattr( torch.get_device_module().get_device_properties( torch.get_device_module().current_device() ), "warp_size", 32, ) if warp_size is None: warp_size = 32 # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit return [ triton.Config({}, num_warps=warp_count) for warp_count in [1, 2, 4, 8, 16, 32] if warp_count * warp_size <= max_threads_per_block ] # return [triton.Config({}, num_warps=8)] # Copied from flash-attn @triton.autotune( configs=triton_autotune_configs(), key=[ "N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS", "HAS_WEIGHT", "HAS_X1", "HAS_W1", "HAS_B1", ], ) # torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel # @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None}) # @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None}) # @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None}) # @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None}) # @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None}) @triton.jit def _layer_norm_fwd_1pass_kernel( X, # pointer to the input Y, # pointer to the output W, # pointer to the weights B, # pointer to the biases RESIDUAL, # pointer to the residual X1, W1, B1, Y1, RESIDUAL_OUT, # pointer to the residual ROWSCALE, SEEDS, # Dropout seeds for each row DROPOUT_MASK, DROPOUT_MASK1, Mean, # pointer to the mean Rstd, # pointer to the 1/std stride_x_row, # how much to increase the pointer when moving by 1 row stride_y_row, stride_res_row, stride_res_out_row, stride_x1_row, stride_y1_row, M, # number of rows in X N, # number of columns in X eps, # epsilon to avoid division by zero dropout_p, # Dropout probability zero_centered_weight, # If true, add 1.0 to the weight IS_RMS_NORM: tl.constexpr, BLOCK_N: tl.constexpr, HAS_RESIDUAL: tl.constexpr, STORE_RESIDUAL_OUT: tl.constexpr, HAS_WEIGHT: tl.constexpr, HAS_BIAS: tl.constexpr, HAS_DROPOUT: tl.constexpr, STORE_DROPOUT_MASK: tl.constexpr, HAS_ROWSCALE: tl.constexpr, HAS_X1: tl.constexpr, HAS_W1: tl.constexpr, HAS_B1: tl.constexpr, ): # Map the program id to the row of X and Y it should compute. row = tl.program_id(0) X += row * stride_x_row Y += row * stride_y_row if HAS_RESIDUAL: RESIDUAL += row * stride_res_row if STORE_RESIDUAL_OUT: RESIDUAL_OUT += row * stride_res_out_row if HAS_X1: X1 += row * stride_x1_row if HAS_W1: Y1 += row * stride_y1_row # Compute mean and variance cols = tl.arange(0, BLOCK_N) x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) if HAS_ROWSCALE: rowscale = tl.load(ROWSCALE + row).to(tl.float32) x *= rowscale if HAS_DROPOUT: # Compute dropout mask # 7 rounds is good enough, and reduces register pressure keep_mask = ( tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p ) x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0) if STORE_DROPOUT_MASK: tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N) if HAS_X1: x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32) if HAS_ROWSCALE: rowscale = tl.load(ROWSCALE + M + row).to(tl.float32) x1 *= rowscale if HAS_DROPOUT: # Compute dropout mask # 7 rounds is good enough, and reduces register pressure keep_mask = ( tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7) > dropout_p ) x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0) if STORE_DROPOUT_MASK: tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N) x += x1 if HAS_RESIDUAL: residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) x += residual if STORE_RESIDUAL_OUT: tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) if not IS_RMS_NORM: mean = tl.sum(x, axis=0) / N tl.store(Mean + row, mean) xbar = tl.where(cols < N, x - mean, 0.0) var = tl.sum(xbar * xbar, axis=0) / N else: xbar = tl.where(cols < N, x, 0.0) var = tl.sum(xbar * xbar, axis=0) / N rstd = 1 / tl.sqrt(var + eps) tl.store(Rstd + row, rstd) # Normalize and apply linear transformation mask = cols < N if HAS_WEIGHT: w = tl.load(W + cols, mask=mask).to(tl.float32) if zero_centered_weight: w += 1.0 if HAS_BIAS: b = tl.load(B + cols, mask=mask).to(tl.float32) x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd if HAS_WEIGHT: y = x_hat * w + b if HAS_BIAS else x_hat * w else: y = x_hat + b if HAS_BIAS else x_hat # Write output tl.store(Y + cols, y, mask=mask) if HAS_W1: w1 = tl.load(W1 + cols, mask=mask).to(tl.float32) if zero_centered_weight: w1 += 1.0 if HAS_B1: b1 = tl.load(B1 + cols, mask=mask).to(tl.float32) y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1 tl.store(Y1 + cols, y1, mask=mask) def _layer_norm_fwd( x: Tensor, weight: Tensor, bias: Tensor, eps: float, residual: Optional[Tensor] = None, x1: Optional[Tensor] = None, weight1: Optional[Tensor] = None, bias1: Optional[Tensor] = None, dropout_p: float = 0.0, rowscale: Optional[Tensor] = None, out_dtype: Optional[torch.dtype] = None, residual_dtype: Optional[torch.dtype] = None, zero_centered_weight: bool = False, is_rms_norm: bool = False, return_dropout_mask: bool = False, out: Optional[Tensor] = None, residual_out: Optional[Tensor] = None, ) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): # Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library # and torch.compile unhappy. Also allocate memory for out and residual_out if they are None # so that _layer_norm_fwd_impl doesn't have to return them. if out is None: out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) if residual is not None: residual_dtype = residual.dtype if residual_out is None and ( residual is not None or (residual_dtype is not None and residual_dtype != x.dtype) or dropout_p > 0.0 or rowscale is not None or x1 is not None ): residual_out = torch.empty_like( x, dtype=residual_dtype if residual_dtype is not None else x.dtype ) else: residual_out = None y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl( x, weight, bias, eps, out, residual=residual, x1=x1, weight1=weight1, bias1=bias1, dropout_p=dropout_p, rowscale=rowscale, zero_centered_weight=zero_centered_weight, is_rms_norm=is_rms_norm, return_dropout_mask=return_dropout_mask, residual_out=residual_out, ) # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0 if residual_out is None: residual_out = x return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 # [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema # since we're returning a tuple of tensors def _layer_norm_fwd_impl( x: Tensor, weight: Optional[Tensor], bias: Tensor, eps: float, out: Tensor, residual: Optional[Tensor] = None, x1: Optional[Tensor] = None, weight1: Optional[Tensor] = None, bias1: Optional[Tensor] = None, dropout_p: float = 0.0, rowscale: Optional[Tensor] = None, zero_centered_weight: bool = False, is_rms_norm: bool = False, return_dropout_mask: bool = False, residual_out: Optional[Tensor] = None, ) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor): M, N = x.shape assert x.stride(-1) == 1 if residual is not None: assert residual.stride(-1) == 1 assert residual.shape == (M, N) if weight is not None: assert weight.shape == (N,) assert weight.stride(-1) == 1 if bias is not None: assert bias.stride(-1) == 1 assert bias.shape == (N,) if x1 is not None: assert x1.shape == x.shape assert rowscale is None assert x1.stride(-1) == 1 if weight1 is not None: assert weight1.shape == (N,) assert weight1.stride(-1) == 1 if bias1 is not None: assert bias1.shape == (N,) assert bias1.stride(-1) == 1 if rowscale is not None: assert rowscale.is_contiguous() assert rowscale.shape == (M,) assert out.shape == x.shape assert out.stride(-1) == 1 if residual_out is not None: assert residual_out.shape == x.shape assert residual_out.stride(-1) == 1 if weight1 is not None: y1 = torch.empty_like(out) assert y1.stride(-1) == 1 else: y1 = None mean = ( torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None ) rstd = torch.empty((M,), dtype=torch.float32, device=x.device) if dropout_p > 0.0: seeds = torch.randint( 2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64 ) else: seeds = None if return_dropout_mask and dropout_p > 0.0: dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool) if x1 is not None: dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool) else: dropout_mask1 = None else: dropout_mask, dropout_mask1 = None, None # Less than 64KB per feature: enqueue fused kernel MAX_FUSED_SIZE = 65536 // x.element_size() BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) if N > BLOCK_N: raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") with torch.get_device_module().device(x.device.index): torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)]( x, out, weight if weight is not None else x, # unused when HAS_WEIGHT == False bias, residual, x1, weight1, bias1, y1, residual_out, rowscale, seeds, dropout_mask, dropout_mask1, mean, rstd, x.stride(0), out.stride(0), residual.stride(0) if residual is not None else 0, residual_out.stride(0) if residual_out is not None else 0, x1.stride(0) if x1 is not None else 0, y1.stride(0) if y1 is not None else 0, M, N, eps, dropout_p, # Passing bool make torch inductor very unhappy since it then tries to compare to int_max int(zero_centered_weight), is_rms_norm, BLOCK_N, residual is not None, residual_out is not None, weight is not None, bias is not None, dropout_p > 0.0, dropout_mask is not None, rowscale is not None, HAS_X1=x1 is not None, HAS_W1=weight1 is not None, HAS_B1=bias1 is not None, ) return y1, mean, rstd, seeds, dropout_mask, dropout_mask1 class LayerNormFn: @staticmethod def forward( x, weight, bias, residual=None, x1=None, weight1=None, bias1=None, eps=1e-6, dropout_p=0.0, rowscale=None, prenorm=False, residual_in_fp32=False, zero_centered_weight=False, is_rms_norm=False, return_dropout_mask=False, out_dtype=None, out=None, residual_out=None, ): x_shape_og = x.shape # reshape input data into 2D tensor x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1])) if residual is not None: assert residual.shape == x_shape_og residual = maybe_contiguous_lastdim( residual.reshape(-1, residual.shape[-1]) ) if x1 is not None: assert x1.shape == x_shape_og assert rowscale is None, "rowscale is not supported with parallel LayerNorm" x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1])) # weight can be None when elementwise_affine=False for LayerNorm if weight is not None: weight = weight.contiguous() bias = maybe_contiguous(bias) weight1 = maybe_contiguous(weight1) bias1 = maybe_contiguous(bias1) if rowscale is not None: rowscale = rowscale.reshape(-1).contiguous() residual_dtype = ( residual.dtype if residual is not None else (torch.float32 if residual_in_fp32 else None) ) if out is not None: out = out.reshape(-1, out.shape[-1]) if residual_out is not None: residual_out = residual_out.reshape(-1, residual_out.shape[-1]) y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = ( _layer_norm_fwd( x, weight, bias, eps, residual, x1, weight1, bias1, dropout_p=dropout_p, rowscale=rowscale, out_dtype=out_dtype, residual_dtype=residual_dtype, zero_centered_weight=zero_centered_weight, is_rms_norm=is_rms_norm, return_dropout_mask=return_dropout_mask, out=out, residual_out=residual_out, ) ) y = y.reshape(x_shape_og) if residual is not None: residual_out = residual_out.reshape(x_shape_og) return y, residual_out return y def layer_norm_fn( x, weight, bias, residual=None, x1=None, weight1=None, bias1=None, eps=1e-6, dropout_p=0.0, rowscale=None, prenorm=False, residual_in_fp32=False, zero_centered_weight=False, is_rms_norm=False, return_dropout_mask=False, out_dtype=None, out=None, residual_out=None, ): return LayerNormFn.forward( x, weight, bias, residual, x1, weight1, bias1, eps, dropout_p, rowscale, prenorm, residual_in_fp32, zero_centered_weight, is_rms_norm, return_dropout_mask, out_dtype, out, residual_out, ) @triton.jit def _norm_infer_kernel( X, Y, W, B, stride_x_row, stride_y_row, M, N, eps, IS_RMS_NORM: tl.constexpr, HAS_WEIGHT: tl.constexpr, HAS_BIAS: tl.constexpr, BLOCK_N: tl.constexpr, ): row = tl.program_id(0) X += row * stride_x_row Y += row * stride_y_row if HAS_WEIGHT: W += 0 if HAS_BIAS: B += 0 cols = tl.arange(0, BLOCK_N) x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) if not IS_RMS_NORM: mean = tl.sum(x, axis=0) / N xbar = tl.where(cols < N, x - mean, 0.0) var = tl.sum(xbar * xbar, axis=0) / N else: xbar = tl.where(cols < N, x, 0.0) var = tl.sum(xbar * xbar, axis=0) / N rstd = 1 / tl.sqrt(var + eps) x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd if HAS_WEIGHT: w = tl.load(W + cols, mask=cols < N, other=1.0).to(tl.float32) y = x_hat * w else: y = x_hat if HAS_BIAS: b = tl.load(B + cols, mask=cols < N, other=0.0).to(tl.float32) y += b tl.store(Y + cols, y, mask=cols < N) def norm_infer( x: Tensor, weight: Optional[Tensor], bias: Optional[Tensor], eps: float, is_rms_norm: bool = False, out: Optional[Tensor] = None, ): M, N = x.shape x = x.contiguous() if weight is not None: assert weight.shape == (N,) assert weight.stride(-1) == 1 if bias is not None: assert bias.shape == (N,) assert bias.stride(-1) == 1 if out is None: out = torch.empty_like(x) MAX_FUSED_SIZE = 65536 // x.element_size() BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) if N > BLOCK_N: raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") num_warps = min(max(BLOCK_N // 256, 1), 8) _norm_infer_kernel[(M,)]( x, out, weight if weight is not None else x, # dummy when HAS_WEIGHT=False bias if bias is not None else x, # dummy when HAS_BIAS=False x.stride(0), out.stride(0), M, N, eps, IS_RMS_NORM=is_rms_norm, HAS_WEIGHT=weight is not None, HAS_BIAS=bias is not None, BLOCK_N=BLOCK_N, num_warps=num_warps, ) return out def rms_norm_fn( x, weight, bias, residual=None, x1=None, weight1=None, bias1=None, eps=1e-6, dropout_p=0.0, rowscale=None, prenorm=False, residual_in_fp32=False, zero_centered_weight=False, return_dropout_mask=False, out_dtype=None, out=None, residual_out=None, ): return LayerNormFn.forward( x, weight, bias, residual, x1, weight1, bias1, eps, dropout_p, rowscale, prenorm, residual_in_fp32, zero_centered_weight, True, return_dropout_mask, out_dtype, out, residual_out, ) from sglang.multimodal_gen.runtime.platforms import current_platform if current_platform.is_mps(): from .mps_fallback import norm_infer_native, rms_norm_fn_native norm_infer = norm_infer_native rms_norm_fn = rms_norm_fn_native ================================================ FILE: python/sglang/jit_kernel/diffusion/triton/npu_fallback.py ================================================ import torch # TODO: remove this when triton ascend bug is fixed def fuse_scale_shift_native( x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, block_l: int = 128, block_c: int = 128, ): return x * (1 + scale) + shift # TODO: remove this when triton ascend bug is fixed def apply_rotary_embedding_native( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False ) -> torch.Tensor: cos = cos.unsqueeze(-2).to(x.dtype) sin = sin.unsqueeze(-2).to(x.dtype) x1 = x[..., ::2] x2 = x[..., 1::2] o1 = x1 * cos - x2 * sin o2 = x2 * cos + x1 * sin return torch.stack((o1, o2), dim=-1).flatten(-2) ================================================ FILE: python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py ================================================ import torch import triton # type: ignore import triton.language as tl # type: ignore # Adapted from https://github.com/ModelTC/LightX2V/blob/main/lightx2v/common/ops/norm/triton_ops.py#L905-L956 @triton.jit def _rms_norm_tiled_onepass( y_ptr, x_ptr, w_ptr, SEQ: tl.constexpr, DIM: tl.constexpr, EPS: tl.constexpr, BLOCK_SIZE_SEQ: tl.constexpr, BLOCK_SIZE_DIM: tl.constexpr, ): seq_blk_id = tl.program_id(0) seq_id = seq_blk_id * BLOCK_SIZE_SEQ seq_offset = seq_id + tl.arange(0, BLOCK_SIZE_SEQ)[:, None] s_mask = seq_offset < SEQ d_offset = tl.arange(0, BLOCK_SIZE_DIM)[None, :] d_mask = d_offset < DIM y_blk = y_ptr + seq_offset * DIM + d_offset x_blk = x_ptr + seq_offset * DIM + d_offset mask = s_mask & d_mask x = tl.load(x_blk, mask=mask, other=0.0).to(tl.float32) mean_square = tl.sum(x * x, axis=1, keep_dims=True) / DIM rstd = tl.math.rsqrt(mean_square + EPS) w = tl.load(w_ptr + d_offset, mask=d_mask) tl.store(y_blk, x * rstd * w, mask=mask) def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6): shape = x.shape x = x.contiguous() y = torch.empty_like(x) x_view = x.reshape(-1, shape[-1]) y_view = y.reshape(-1, shape[-1]) S, D = x_view.shape BLOCK_SIZE_SEQ = min(16, triton.next_power_of_2(max(1, S // 512))) grid = (triton.cdiv(S, BLOCK_SIZE_SEQ),) with torch.get_device_module().device(x.device): torch.library.wrap_triton(_rms_norm_tiled_onepass)[grid]( y_view, x_view, w, S, D, eps, BLOCK_SIZE_DIM=triton.next_power_of_2(D), BLOCK_SIZE_SEQ=BLOCK_SIZE_SEQ, ) return y from sglang.multimodal_gen.runtime.platforms import current_platform if current_platform.is_mps(): from .mps_fallback import triton_one_pass_rms_norm_native triton_one_pass_rms_norm = triton_one_pass_rms_norm_native ================================================ FILE: python/sglang/jit_kernel/diffusion/triton/rotary.py ================================================ import torch import triton # type: ignore import triton.language as tl # type: ignore from sglang.multimodal_gen.runtime.platforms import current_platform @triton.autotune( configs=[ triton.Config({"BLOCK_HS_HALF": 32}, num_warps=2), triton.Config({"BLOCK_HS_HALF": 64}, num_warps=4), triton.Config({"BLOCK_HS_HALF": 128}, num_warps=4), triton.Config({"BLOCK_HS_HALF": 256}, num_warps=8), ], key=["head_size", "interleaved"], ) @triton.jit def _rotary_embedding_kernel( output_ptr, x_ptr, cos_ptr, sin_ptr, num_heads, head_size, num_tokens, stride_x_row, stride_cos_row, stride_sin_row, interleaved: tl.constexpr, BLOCK_HS_HALF: tl.constexpr, ): row_idx = tl.program_id(0) token_idx = (row_idx // num_heads) % num_tokens x_row_ptr = x_ptr + row_idx * stride_x_row cos_row_ptr = cos_ptr + token_idx * stride_cos_row sin_row_ptr = sin_ptr + token_idx * stride_sin_row output_row_ptr = output_ptr + row_idx * stride_x_row # half size for x1 and x2 head_size_half = head_size // 2 for block_start in range(0, head_size_half, BLOCK_HS_HALF): offsets_half = block_start + tl.arange(0, BLOCK_HS_HALF) mask = offsets_half < head_size_half cos_vals = tl.load(cos_row_ptr + offsets_half, mask=mask, other=0.0) sin_vals = tl.load(sin_row_ptr + offsets_half, mask=mask, other=0.0) offsets_x1 = 2 * offsets_half offsets_x2 = 2 * offsets_half + 1 x1_vals = tl.load(x_row_ptr + offsets_x1, mask=mask, other=0.0) x2_vals = tl.load(x_row_ptr + offsets_x2, mask=mask, other=0.0) x1_fp32 = x1_vals.to(tl.float32) x2_fp32 = x2_vals.to(tl.float32) cos_fp32 = cos_vals.to(tl.float32) sin_fp32 = sin_vals.to(tl.float32) o1_vals = tl.fma(-x2_fp32, sin_fp32, x1_fp32 * cos_fp32) o2_vals = tl.fma(x1_fp32, sin_fp32, x2_fp32 * cos_fp32) tl.store(output_row_ptr + offsets_x1, o1_vals.to(x1_vals.dtype), mask=mask) tl.store(output_row_ptr + offsets_x2, o2_vals.to(x2_vals.dtype), mask=mask) def apply_rotary_embedding( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False ) -> torch.Tensor: output = torch.empty_like(x) if x.dim() > 3: bsz, num_tokens, num_heads, head_size = x.shape else: num_tokens, num_heads, head_size = x.shape bsz = 1 assert head_size % 2 == 0, "head_size must be divisible by 2" x_reshaped = x.view(-1, head_size) output_reshaped = output.view(-1, head_size) # num_tokens per head, 1 token per block grid = (bsz * num_tokens * num_heads,) if interleaved and cos.shape[-1] == head_size: cos = cos[..., ::2].contiguous() sin = sin[..., ::2].contiguous() else: cos = cos.contiguous() sin = sin.contiguous() _rotary_embedding_kernel[grid]( output_reshaped, x_reshaped, cos, sin, num_heads, head_size, num_tokens, x_reshaped.stride(0), cos.stride(0), sin.stride(0), interleaved, ) return output if current_platform.is_npu(): from .npu_fallback import apply_rotary_embedding_native apply_rotary_embedding = apply_rotary_embedding_native if current_platform.is_mps(): from .mps_fallback import apply_rotary_embedding_native apply_rotary_embedding = apply_rotary_embedding_native ================================================ FILE: python/sglang/jit_kernel/diffusion/triton/scale_shift.py ================================================ import torch import triton # type: ignore import triton.language as tl # type: ignore from sglang.multimodal_gen.runtime.platforms import current_platform @triton.jit def _fused_layernorm_scale_shift_gate_select01_kernel( output_ptr, gate_out_ptr, x_ptr, weight_ptr, bias_ptr, scale0_ptr, shift0_ptr, gate0_ptr, scale1_ptr, shift1_ptr, gate1_ptr, index_ptr, inner_dim, seq_len, stride_x_row, stride_out_row, stride_go_row, stride_w, stride_b, stride_s0_b, stride_s0_c, stride_sh0_b, stride_sh0_c, stride_g0_b, stride_g0_c, stride_s1_b, stride_s1_c, stride_sh1_b, stride_sh1_c, stride_g1_b, stride_g1_c, stride_i_b, stride_i_l, eps, HAS_WEIGHT: tl.constexpr, HAS_BIAS: tl.constexpr, BLOCK_N: tl.constexpr, ): row = tl.program_id(0) cols = tl.arange(0, BLOCK_N) mask = cols < inner_dim x_row_ptr = x_ptr + row * stride_x_row out_row_ptr = output_ptr + row * stride_out_row gate_row_ptr = gate_out_ptr + row * stride_go_row x = tl.load(x_row_ptr + cols, mask=mask, other=0.0).to(tl.float32) mean = tl.sum(x, axis=0) / inner_dim xbar = tl.where(mask, x - mean, 0.0) var = tl.sum(xbar * xbar, axis=0) / inner_dim rstd = tl.rsqrt(var + eps) x_hat = (x - mean) * rstd if HAS_WEIGHT: w = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32) x_hat = x_hat * w if HAS_BIAS: b = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32) x_hat = x_hat + b batch_idx = row // seq_len seq_idx = row % seq_len idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1) scale0 = tl.load( scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c, mask=mask, other=0.0, ).to(tl.float32) shift0 = tl.load( shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c, mask=mask, other=0.0, ).to(tl.float32) gate0 = tl.load( gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c, mask=mask, other=0.0, ) scale1 = tl.load( scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c, mask=mask, other=0.0, ).to(tl.float32) shift1 = tl.load( shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c, mask=mask, other=0.0, ).to(tl.float32) gate1 = tl.load( gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c, mask=mask, other=0.0, ) scale = tl.where(idx, scale1, scale0) shift = tl.where(idx, shift1, shift0) gate = tl.where(idx, gate1, gate0) y = x_hat * (1.0 + scale) + shift tl.store(out_row_ptr + cols, y, mask=mask) tl.store(gate_row_ptr + cols, gate, mask=mask) @triton.jit def _fused_residual_layernorm_scale_shift_gate_select01_kernel( output_ptr, residual_out_ptr, gate_out_ptr, x_ptr, residual_ptr, residual_gate_ptr, weight_ptr, bias_ptr, scale0_ptr, shift0_ptr, gate0_ptr, scale1_ptr, shift1_ptr, gate1_ptr, index_ptr, inner_dim, seq_len, stride_x_row, stride_res_row, stride_rg_row, stride_out_row, stride_res_out_row, stride_go_row, stride_w, stride_b, stride_s0_b, stride_s0_c, stride_sh0_b, stride_sh0_c, stride_g0_b, stride_g0_c, stride_s1_b, stride_s1_c, stride_sh1_b, stride_sh1_c, stride_g1_b, stride_g1_c, stride_i_b, stride_i_l, eps, HAS_WEIGHT: tl.constexpr, HAS_BIAS: tl.constexpr, BLOCK_N: tl.constexpr, ): row = tl.program_id(0) cols = tl.arange(0, BLOCK_N) mask = cols < inner_dim x_row_ptr = x_ptr + row * stride_x_row res_row_ptr = residual_ptr + row * stride_res_row rg_row_ptr = residual_gate_ptr + row * stride_rg_row out_row_ptr = output_ptr + row * stride_out_row res_out_row_ptr = residual_out_ptr + row * stride_res_out_row gate_row_ptr = gate_out_ptr + row * stride_go_row x = tl.load(x_row_ptr + cols, mask=mask, other=0.0).to(tl.float32) residual = tl.load(res_row_ptr + cols, mask=mask, other=0.0).to(tl.float32) residual_gate = tl.load(rg_row_ptr + cols, mask=mask, other=0.0).to(tl.float32) residual_out = residual + residual_gate * x tl.store(res_out_row_ptr + cols, residual_out, mask=mask) mean = tl.sum(residual_out, axis=0) / inner_dim xbar = tl.where(mask, residual_out - mean, 0.0) var = tl.sum(xbar * xbar, axis=0) / inner_dim rstd = tl.rsqrt(var + eps) x_hat = (residual_out - mean) * rstd if HAS_WEIGHT: w = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32) x_hat = x_hat * w if HAS_BIAS: b = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32) x_hat = x_hat + b batch_idx = row // seq_len seq_idx = row % seq_len idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1) scale0 = tl.load( scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c, mask=mask, other=0.0, ).to(tl.float32) shift0 = tl.load( shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c, mask=mask, other=0.0, ).to(tl.float32) gate0 = tl.load( gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c, mask=mask, other=0.0, ) scale1 = tl.load( scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c, mask=mask, other=0.0, ).to(tl.float32) shift1 = tl.load( shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c, mask=mask, other=0.0, ).to(tl.float32) gate1 = tl.load( gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c, mask=mask, other=0.0, ) scale = tl.where(idx, scale1, scale0) shift = tl.where(idx, shift1, shift0) gate = tl.where(idx, gate1, gate0) y = x_hat * (1.0 + scale) + shift tl.store(out_row_ptr + cols, y, mask=mask) tl.store(gate_row_ptr + cols, gate, mask=mask) @triton.autotune( configs=[ triton.Config({"BLOCK_N": 64}, num_warps=2), triton.Config({"BLOCK_N": 128}, num_warps=4), triton.Config({"BLOCK_N": 256}, num_warps=4), triton.Config({"BLOCK_N": 512}, num_warps=4), triton.Config({"BLOCK_N": 1024}, num_warps=8), ], key=["inner_dim"], ) @triton.jit def _fused_scale_shift_4d_kernel( output_ptr, normalized_ptr, scale_ptr, shift_ptr, scale_constant: tl.constexpr, # scale_constant is either 0 or 1. rows, inner_dim, seq_len, num_frames, frame_seqlen, BLOCK_N: tl.constexpr, ): pid_row = tl.program_id(0) pid_col = tl.program_id(1) col_offsets = pid_col * BLOCK_N + tl.arange(0, BLOCK_N) mask = col_offsets < inner_dim # Pointers for normalized and output row_base = pid_row * inner_dim norm_ptrs = normalized_ptr + row_base + col_offsets out_ptrs = output_ptr + row_base + col_offsets # Pointers for scale (per-frame) and shift (per-token) b_idx = pid_row // seq_len t_idx = pid_row % seq_len frame_idx_in_batch = t_idx // frame_seqlen scale_row_idx = b_idx * num_frames + frame_idx_in_batch scale_ptrs = scale_ptr + scale_row_idx * inner_dim + col_offsets # shift is per-token [B*L, C], indexed by pid_row directly shift_ptrs = shift_ptr + pid_row * inner_dim + col_offsets normalized = tl.load(norm_ptrs, mask=mask, other=0.0) scale = tl.load(scale_ptrs, mask=mask, other=0.0) shift = tl.load(shift_ptrs, mask=mask, other=0.0) scale_const_tensor = tl.full([BLOCK_N], scale_constant, dtype=scale.dtype) output = normalized * (scale_const_tensor + scale) + shift tl.store(out_ptrs, output, mask=mask) @triton.jit def fuse_scale_shift_kernel_blc_opt( x_ptr, shift_ptr, scale_ptr, scale_constant: tl.constexpr, # scale_constant is either 0 or 1., y_ptr, B, L, C, stride_x_b, stride_x_l, stride_x_c, stride_s_b, stride_s_l, stride_s_c, stride_sc_b, stride_sc_l, stride_sc_c, SCALE_IS_SCALAR: tl.constexpr, SHIFT_IS_SCALAR: tl.constexpr, BLOCK_L: tl.constexpr, BLOCK_C: tl.constexpr, ): pid_l = tl.program_id(0) pid_c = tl.program_id(1) pid_b = tl.program_id(2) l_offsets = pid_l * BLOCK_L + tl.arange(0, BLOCK_L) c_offsets = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) mask_l = l_offsets < L mask_c = c_offsets < C mask = mask_l[:, None] & mask_c[None, :] x_off = ( pid_b * stride_x_b + l_offsets[:, None] * stride_x_l + c_offsets[None, :] * stride_x_c ) x = tl.load(x_ptr + x_off, mask=mask, other=0) if SHIFT_IS_SCALAR: shift_val = tl.load(shift_ptr) shift = tl.full((BLOCK_L, BLOCK_C), shift_val, dtype=shift_val.dtype) else: s_off = ( pid_b * stride_s_b + l_offsets[:, None] * stride_s_l + c_offsets[None, :] * stride_s_c ) shift = tl.load(shift_ptr + s_off, mask=mask, other=0) if SCALE_IS_SCALAR: scale_val = tl.load(scale_ptr) scale = tl.full((BLOCK_L, BLOCK_C), scale_val, dtype=scale_val.dtype) else: sc_off = ( pid_b * stride_sc_b + l_offsets[:, None] * stride_sc_l + c_offsets[None, :] * stride_sc_c ) scale = tl.load(scale_ptr + sc_off, mask=mask, other=0) y = x * (scale_constant + scale) + shift tl.store(y_ptr + x_off, y, mask=mask) @triton.jit def fuse_scale_shift_gate_select01_kernel_blc_opt( x_ptr, shift0_ptr, scale0_ptr, gate0_ptr, shift1_ptr, scale1_ptr, gate1_ptr, index_ptr, y_ptr, gate_out_ptr, B, L, C, stride_x_b, stride_x_l, stride_x_c, stride_s0_b, stride_s0_c, stride_sc0_b, stride_sc0_c, stride_g0_b, stride_g0_c, stride_s1_b, stride_s1_c, stride_sc1_b, stride_sc1_c, stride_g1_b, stride_g1_c, stride_i_b, stride_i_l, stride_go_b, stride_go_l, stride_go_c, BLOCK_L: tl.constexpr, BLOCK_C: tl.constexpr, ): pid_l = tl.program_id(0) pid_c = tl.program_id(1) pid_b = tl.program_id(2) l_offsets = pid_l * BLOCK_L + tl.arange(0, BLOCK_L) c_offsets = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) mask_l = l_offsets < L mask_c = c_offsets < C mask = mask_l[:, None] & mask_c[None, :] x_off = ( pid_b * stride_x_b + l_offsets[:, None] * stride_x_l + c_offsets[None, :] * stride_x_c ) x = tl.load(x_ptr + x_off, mask=mask, other=0) idx_off = pid_b * stride_i_b + l_offsets * stride_i_l idx = tl.load(index_ptr + idx_off, mask=mask_l, other=0).to(tl.int1)[:, None] s0_off = pid_b * stride_s0_b + c_offsets[None, :] * stride_s0_c sc0_off = pid_b * stride_sc0_b + c_offsets[None, :] * stride_sc0_c g0_off = pid_b * stride_g0_b + c_offsets[None, :] * stride_g0_c s1_off = pid_b * stride_s1_b + c_offsets[None, :] * stride_s1_c sc1_off = pid_b * stride_sc1_b + c_offsets[None, :] * stride_sc1_c g1_off = pid_b * stride_g1_b + c_offsets[None, :] * stride_g1_c shift0 = tl.load(shift0_ptr + s0_off, mask=mask_c[None, :], other=0) scale0 = tl.load(scale0_ptr + sc0_off, mask=mask_c[None, :], other=0) gate0 = tl.load(gate0_ptr + g0_off, mask=mask_c[None, :], other=0) shift1 = tl.load(shift1_ptr + s1_off, mask=mask_c[None, :], other=0) scale1 = tl.load(scale1_ptr + sc1_off, mask=mask_c[None, :], other=0) gate1 = tl.load(gate1_ptr + g1_off, mask=mask_c[None, :], other=0) shift = tl.where(idx, shift1, shift0) scale = tl.where(idx, scale1, scale0) gate = tl.where(idx, gate1, gate0) y = x * (1 + scale) + shift tl.store(y_ptr + x_off, y, mask=mask) go_off = ( pid_b * stride_go_b + l_offsets[:, None] * stride_go_l + c_offsets[None, :] * stride_go_c ) tl.store(gate_out_ptr + go_off, gate, mask=mask) def fuse_scale_shift_kernel( x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, scale_constant: float = 1.0, block_l: int = 128, block_c: int = 128, ): assert x.is_cuda and scale.is_cuda assert x.is_contiguous() B, L, C = x.shape output = torch.empty_like(x) if scale.dim() == 4: # scale/shift: [B, F, 1, C] rows = B * L x_2d = x.view(rows, C) output_2d = output.view(rows, C) grid = lambda META: (rows, triton.cdiv(C, META["BLOCK_N"])) num_frames = scale.shape[1] assert ( L % num_frames == 0 ), "seq_len must be divisible by num_frames for 4D scale/shift" frame_seqlen = L // num_frames # Compact scale [B, F, 1, C] -> [B*F, C] (per-frame) scale_reshaped = scale.squeeze(2).reshape(-1, C).contiguous() # shift is per-token [B, L, C] -> [B*L, C] shift_reshaped = shift.reshape(rows, C).contiguous() _fused_scale_shift_4d_kernel[grid]( output_2d, x_2d, scale_reshaped, shift_reshaped, scale_constant, rows, C, L, num_frames, frame_seqlen, ) else: # 2D: [B, C] or [1, C] -> treat as [B, 1, C] and broadcast over L # 3D: [B, L, C] (or broadcastable variants like [B, 1, C], [1, L, C], [1, 1, C]) # Also support scalar (0D or 1-element) if scale.dim() == 0 or (scale.dim() == 1 and scale.numel() == 1): scale_blc = scale.reshape(1) elif scale.dim() == 2: scale_blc = scale[:, None, :] elif scale.dim() == 3: scale_blc = scale else: raise ValueError("scale must be 0D/1D(1)/2D/3D or 4D") if shift.dim() == 0 or (shift.dim() == 1 and shift.numel() == 1): shift_blc = shift.reshape(1) elif shift.dim() == 2: shift_blc = shift[:, None, :] elif shift.dim() == 3: shift_blc = shift else: # broadcast later via expand if possible shift_blc = shift need_scale_scalar = scale_blc.dim() == 1 and scale_blc.numel() == 1 need_shift_scalar = shift_blc.dim() == 1 and shift_blc.numel() == 1 if not need_scale_scalar: scale_exp = scale_blc.expand(B, L, C) s_sb, s_sl, s_sc = scale_exp.stride() else: s_sb = s_sl = s_sc = 0 if not need_shift_scalar: shift_exp = shift_blc.expand(B, L, C) sh_sb, sh_sl, sh_sc = shift_exp.stride() else: sh_sb = sh_sl = sh_sc = 0 # If both scalars and both zero, copy fast-path if need_scale_scalar and need_shift_scalar: if not ( scale_blc.any().to("cpu", non_blocking=True) or shift_blc.any().to("cpu", non_blocking=True) ): output.copy_(x) return output grid = (triton.cdiv(L, block_l), triton.cdiv(C, block_c), B) fuse_scale_shift_kernel_blc_opt[grid]( x, shift_blc if need_shift_scalar else shift_exp, scale_blc if need_scale_scalar else scale_exp, scale_constant, output, B, L, C, x.stride(0), x.stride(1), x.stride(2), sh_sb, sh_sl, sh_sc, s_sb, s_sl, s_sc, SCALE_IS_SCALAR=need_scale_scalar, SHIFT_IS_SCALAR=need_shift_scalar, BLOCK_L=block_l, BLOCK_C=block_c, num_warps=4, num_stages=2, ) return output def fuse_scale_shift_gate_select01_kernel( x: torch.Tensor, scale0: torch.Tensor, shift0: torch.Tensor, gate0: torch.Tensor, scale1: torch.Tensor, shift1: torch.Tensor, gate1: torch.Tensor, index: torch.Tensor, block_l: int = 128, block_c: int = 128, ): assert x.is_contiguous() B, L, C = x.shape output = torch.empty_like(x) gate_out = torch.empty_like(x) if ( scale0.dim() != 2 or shift0.dim() != 2 or gate0.dim() != 2 or scale1.dim() != 2 or shift1.dim() != 2 or gate1.dim() != 2 ): raise ValueError("scale0/shift0/gate0/scale1/shift1/gate1 must be 2D [B, C]") if index.dim() != 2: raise ValueError("index must be 2D [B, L]") grid = (triton.cdiv(L, block_l), triton.cdiv(C, block_c), B) fuse_scale_shift_gate_select01_kernel_blc_opt[grid]( x, shift0, scale0, gate0, shift1, scale1, gate1, index, output, gate_out, B, L, C, x.stride(0), x.stride(1), x.stride(2), shift0.stride(0), shift0.stride(1), scale0.stride(0), scale0.stride(1), gate0.stride(0), gate0.stride(1), shift1.stride(0), shift1.stride(1), scale1.stride(0), scale1.stride(1), gate1.stride(0), gate1.stride(1), index.stride(0), index.stride(1), gate_out.stride(0), gate_out.stride(1), gate_out.stride(2), BLOCK_L=block_l, BLOCK_C=block_c, num_warps=4, num_stages=2, ) return output, gate_out def fuse_layernorm_scale_shift_gate_select01_kernel( x: torch.Tensor, weight: torch.Tensor | None, bias: torch.Tensor | None, scale0: torch.Tensor, shift0: torch.Tensor, gate0: torch.Tensor, scale1: torch.Tensor, shift1: torch.Tensor, gate1: torch.Tensor, index: torch.Tensor, eps: float, ): assert x.is_cuda assert x.is_contiguous() B, L, C = x.shape output = torch.empty_like(x) gate_out = torch.empty_like(x) if ( scale0.dim() != 2 or shift0.dim() != 2 or gate0.dim() != 2 or scale1.dim() != 2 or shift1.dim() != 2 or gate1.dim() != 2 ): raise ValueError("scale0/shift0/gate0/scale1/shift1/gate1 must be 2D [B, C]") if index.dim() != 2: raise ValueError("index must be 2D [B, L]") if weight is not None and (weight.dim() != 1 or weight.shape[0] != C): raise ValueError("weight must be 1D [C]") if bias is not None and (bias.dim() != 1 or bias.shape[0] != C): raise ValueError("bias must be 1D [C]") x_2d = x.view(B * L, C) output_2d = output.view(B * L, C) gate_out_2d = gate_out.view(B * L, C) weight = weight.contiguous() if weight is not None else x_2d bias = bias.contiguous() if bias is not None else x_2d MAX_FUSED_SIZE = 65536 // x_2d.element_size() BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(C)) if C > BLOCK_N: raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") grid = (B * L,) _fused_layernorm_scale_shift_gate_select01_kernel[grid]( output_2d, gate_out_2d, x_2d, weight, bias, scale0.contiguous(), shift0.contiguous(), gate0.contiguous(), scale1.contiguous(), shift1.contiguous(), gate1.contiguous(), index.contiguous(), C, L, x_2d.stride(0), output_2d.stride(0), gate_out_2d.stride(0), weight.stride(0) if weight.dim() == 1 else 0, bias.stride(0) if bias.dim() == 1 else 0, scale0.stride(0), scale0.stride(1), shift0.stride(0), shift0.stride(1), gate0.stride(0), gate0.stride(1), scale1.stride(0), scale1.stride(1), shift1.stride(0), shift1.stride(1), gate1.stride(0), gate1.stride(1), index.stride(0), index.stride(1), eps, HAS_WEIGHT=weight is not x_2d, HAS_BIAS=bias is not x_2d, BLOCK_N=BLOCK_N, ) return output, gate_out def fuse_residual_layernorm_scale_shift_gate_select01_kernel( x: torch.Tensor, residual: torch.Tensor, residual_gate: torch.Tensor, weight: torch.Tensor | None, bias: torch.Tensor | None, scale0: torch.Tensor, shift0: torch.Tensor, gate0: torch.Tensor, scale1: torch.Tensor, shift1: torch.Tensor, gate1: torch.Tensor, index: torch.Tensor, eps: float, ): assert x.is_cuda assert x.is_contiguous() assert residual.is_contiguous() assert residual_gate.is_contiguous() B, L, C = x.shape output = torch.empty_like(x) residual_out = torch.empty_like(x) gate_out = torch.empty_like(x) if residual.shape != x.shape: raise ValueError("residual must have the same shape as x") if residual_gate.shape != x.shape: raise ValueError("residual_gate must have the same shape as x") if ( scale0.dim() != 2 or shift0.dim() != 2 or gate0.dim() != 2 or scale1.dim() != 2 or shift1.dim() != 2 or gate1.dim() != 2 ): raise ValueError("scale0/shift0/gate0/scale1/shift1/gate1 must be 2D [B, C]") if index.dim() != 2: raise ValueError("index must be 2D [B, L]") if weight is not None and (weight.dim() != 1 or weight.shape[0] != C): raise ValueError("weight must be 1D [C]") if bias is not None and (bias.dim() != 1 or bias.shape[0] != C): raise ValueError("bias must be 1D [C]") x_2d = x.view(B * L, C) residual_2d = residual.view(B * L, C) residual_gate_2d = residual_gate.view(B * L, C) output_2d = output.view(B * L, C) residual_out_2d = residual_out.view(B * L, C) gate_out_2d = gate_out.view(B * L, C) weight = weight.contiguous() if weight is not None else x_2d bias = bias.contiguous() if bias is not None else x_2d MAX_FUSED_SIZE = 65536 // x_2d.element_size() BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(C)) if C > BLOCK_N: raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") grid = (B * L,) _fused_residual_layernorm_scale_shift_gate_select01_kernel[grid]( output_2d, residual_out_2d, gate_out_2d, x_2d, residual_2d, residual_gate_2d, weight, bias, scale0.contiguous(), shift0.contiguous(), gate0.contiguous(), scale1.contiguous(), shift1.contiguous(), gate1.contiguous(), index.contiguous(), C, L, x_2d.stride(0), residual_2d.stride(0), residual_gate_2d.stride(0), output_2d.stride(0), residual_out_2d.stride(0), gate_out_2d.stride(0), weight.stride(0) if weight.dim() == 1 else 0, bias.stride(0) if bias.dim() == 1 else 0, scale0.stride(0), scale0.stride(1), shift0.stride(0), shift0.stride(1), gate0.stride(0), gate0.stride(1), scale1.stride(0), scale1.stride(1), shift1.stride(0), shift1.stride(1), gate1.stride(0), gate1.stride(1), index.stride(0), index.stride(1), eps, HAS_WEIGHT=weight is not x_2d, HAS_BIAS=bias is not x_2d, BLOCK_N=BLOCK_N, ) return output, residual_out, gate_out if current_platform.is_npu(): from .npu_fallback import fuse_scale_shift_native fuse_scale_shift_kernel = fuse_scale_shift_native if current_platform.is_mps(): from .mps_fallback import ( fuse_scale_shift_gate_select01_kernel_native, fuse_scale_shift_kernel_native, ) fuse_scale_shift_kernel = fuse_scale_shift_kernel_native fuse_scale_shift_gate_select01_kernel = fuse_scale_shift_gate_select01_kernel_native ================================================ FILE: python/sglang/jit_kernel/flash_attention_v4.py ================================================ from __future__ import annotations from typing import Callable, Optional, Tuple, Union import torch try: from flash_attn.cute import flash_attn_varlen_func as _flash_attn_varlen_func except Exception as _e: # pragma: no cover _flash_attn_varlen_func = None _flash_attn_import_error = _e else: _flash_attn_import_error = None def _maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]: return x.contiguous() if x is not None and x.stride(-1) != 1 else x def flash_attn_varlen_func( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_k: Optional[torch.Tensor] = None, seqused_q: Optional[torch.Tensor] = None, seqused_k: Optional[torch.Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_k: Optional[int] = None, page_table: Optional[torch.Tensor] = None, softmax_scale: Optional[float] = None, causal: bool = False, softcap: Optional[float] = None, window_size: Tuple[Optional[int], Optional[int]] = (-1, -1), learnable_sink: Optional[torch.Tensor] = None, sinks: Optional[torch.Tensor] = None, num_splits: int = 1, pack_gqa: Optional[bool] = None, score_mod: Optional[Callable] = None, aux_tensors: Optional[list] = None, return_softmax_lse: bool = False, **_: object, ): if _flash_attn_varlen_func is None: # pragma: no cover raise ImportError( "Vendored FlashAttention CUTE is not available (cannot import " "flash_attn.cute). Please check your source tree." ) from _flash_attn_import_error q, k, v = [_maybe_contiguous(t) for t in (q, k, v)] cu_seqlens_q, cu_seqlens_k = [ _maybe_contiguous(t) for t in (cu_seqlens_q, cu_seqlens_k) ] seqused_q, seqused_k = [_maybe_contiguous(t) for t in (seqused_q, seqused_k)] page_table = _maybe_contiguous(page_table) if learnable_sink is None and sinks is not None: learnable_sink = sinks if window_size == (-1, -1): window_size = (None, None) result = _flash_attn_varlen_func( q=q, k=k, v=v, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, seqused_q=seqused_q, seqused_k=seqused_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, page_table=page_table, softmax_scale=softmax_scale, causal=causal, softcap=softcap, window_size=window_size, learnable_sink=learnable_sink, num_splits=num_splits, pack_gqa=pack_gqa, score_mod=score_mod, aux_tensors=aux_tensors, ) if return_softmax_lse: return result if isinstance(result, tuple): return result[0] return result def flash_attn_with_kvcache( q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, k: Optional[torch.Tensor] = None, v: Optional[torch.Tensor] = None, qv: Optional[torch.Tensor] = None, rotary_cos: Optional[torch.Tensor] = None, rotary_sin: Optional[torch.Tensor] = None, cache_seqlens: Optional[Union[int, torch.Tensor]] = None, cache_batch_idx: Optional[torch.Tensor] = None, cache_leftpad: Optional[torch.Tensor] = None, page_table: Optional[torch.Tensor] = None, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_k_new: Optional[torch.Tensor] = None, max_seqlen_q: Optional[int] = None, rotary_seqlens: Optional[torch.Tensor] = None, q_descale: Optional[torch.Tensor] = None, k_descale: Optional[torch.Tensor] = None, v_descale: Optional[torch.Tensor] = None, softmax_scale: Optional[float] = None, causal: bool = False, window_size: Tuple[int, int] = (-1, -1), attention_chunk: Optional[int] = None, softcap: float = 0.0, rotary_interleaved: bool = True, scheduler_metadata=None, num_splits: int = 0, pack_gqa: Optional[bool] = None, sm_margin: int = 0, sinks: Optional[torch.Tensor] = None, score_mod: Optional[Callable] = None, aux_tensors: Optional[list] = None, return_softmax_lse: bool = False, **_: object, ): if k is not None or v is not None or qv is not None: raise NotImplementedError("FA4 does not support updating KV cache in-place.") if rotary_cos is not None or rotary_sin is not None or rotary_seqlens is not None: raise NotImplementedError("FA4 path does not support rotary embedding.") if cache_batch_idx is not None or cache_leftpad is not None: raise NotImplementedError( "FA4 path does not support non-consecutive batch indices or left padding." ) if q_descale is not None or k_descale is not None or v_descale is not None: raise NotImplementedError("FA4 path does not support descale.") if isinstance(cache_seqlens, int): cache_seqlens = torch.full( (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device ) result = flash_attn_varlen_func( q=q, k=k_cache, v=v_cache, cu_seqlens_q=cu_seqlens_q, seqused_k=cache_seqlens, max_seqlen_q=max_seqlen_q, page_table=page_table, softmax_scale=softmax_scale, causal=causal, softcap=softcap if softcap != 0.0 else None, window_size=window_size, num_splits=num_splits if num_splits != 0 else 1, pack_gqa=pack_gqa, learnable_sink=sinks, score_mod=score_mod, aux_tensors=aux_tensors, return_softmax_lse=True, ) if return_softmax_lse: return result if isinstance(result, tuple): return result[0] return result ================================================ FILE: python/sglang/jit_kernel/fused_metadata_copy.py ================================================ """ Fused metadata copy kernel for NSA backend CUDA graph replay. This module provides JIT-compiled CUDA kernels for fusing multiple tensor copy operations into single kernel launches, reducing kernel launch overhead and improving CUDA graph replay performance. The kernels are compiled on-demand using TVM FFI and cached for subsequent use. """ from __future__ import annotations import logging from typing import Optional import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args logger = logging.getLogger(__name__) # ============================================================================ # JIT Module Compilation # ============================================================================ @cache_once def _jit_fused_metadata_copy_module( forward_mode: int, has_real_page_table: bool, has_flashmla: bool ): """Compile JIT module for single-backend fused metadata copy. Args: forward_mode: 0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND has_real_page_table: Whether real_page_table tensors are used has_flashmla: Whether FlashMLA metadata tensors are used """ args = make_cpp_args(forward_mode, has_real_page_table, has_flashmla) try: return load_jit( "fused_metadata_copy", *args, cuda_files=["elementwise/fused_metadata_copy.cuh"], cuda_wrappers=[ ( "fused_metadata_copy", f"FusedMetadataCopyKernel<{args}>::run", ) ], ) except Exception as e: logger.error( f"Failed to compile JIT fused metadata copy kernel " f"(forward_mode={forward_mode}, has_real_page_table={has_real_page_table}, " f"has_flashmla={has_flashmla}): {e}" ) raise @cache_once def _jit_fused_metadata_copy_multi_module( has_real_page_table: bool, has_flashmla: bool ): """Compile JIT module for multi-backend fused metadata copy (DECODE mode only). Args: has_real_page_table: Whether real_page_table tensors are used has_flashmla: Whether FlashMLA metadata tensors are used """ args = make_cpp_args(has_real_page_table, has_flashmla) try: return load_jit( "fused_metadata_copy_multi", *args, cuda_files=["elementwise/fused_metadata_copy.cuh"], cuda_wrappers=[ ( "fused_metadata_copy_multi", f"FusedMetadataCopyMultiKernel<{args}>::run", ) ], ) except Exception as e: logger.error( f"Failed to compile JIT fused metadata copy multi kernel " f"(has_real_page_table={has_real_page_table}, has_flashmla={has_flashmla}): {e}" ) raise # ============================================================================ # Public API # ============================================================================ def fused_metadata_copy_cuda( cache_seqlens_src: torch.Tensor, cu_seqlens_k_src: torch.Tensor, page_indices_src: torch.Tensor, nsa_cache_seqlens_src: torch.Tensor, seqlens_expanded_src: Optional[torch.Tensor], nsa_cu_seqlens_k_src: torch.Tensor, real_page_table_src: Optional[torch.Tensor], flashmla_num_splits_src: Optional[torch.Tensor], flashmla_metadata_src: Optional[torch.Tensor], cache_seqlens_dst: torch.Tensor, cu_seqlens_k_dst: torch.Tensor, page_table_1_dst: torch.Tensor, nsa_cache_seqlens_dst: torch.Tensor, seqlens_expanded_dst: Optional[torch.Tensor], nsa_cu_seqlens_k_dst: torch.Tensor, real_page_table_dst: Optional[torch.Tensor], flashmla_num_splits_dst: Optional[torch.Tensor], flashmla_metadata_dst: Optional[torch.Tensor], forward_mode: int, bs: int, max_len: int, max_seqlen_k: int, seqlens_expanded_size: int, ) -> None: """ Fused metadata copy kernel for NSA backend CUDA graph replay. This function fuses multiple tensor copy operations into a single kernel launch, reducing kernel launch overhead and improving performance. Args: cache_seqlens_src: Source cache sequence lengths [bs] cu_seqlens_k_src: Source cumulative sequence lengths [bs+1] page_indices_src: Source page indices [rows, max_len] nsa_cache_seqlens_src: Source NSA cache sequence lengths [size] seqlens_expanded_src: Optional source expanded sequence lengths [size] (required for TARGET_VERIFY/DRAFT_EXTEND) nsa_cu_seqlens_k_src: Source NSA cumulative sequence lengths [size+1] real_page_table_src: Optional source real page table [rows, cols] flashmla_num_splits_src: Optional source FlashMLA num_splits [size+1] flashmla_metadata_src: Optional source FlashMLA metadata tensor cache_seqlens_dst: Destination cache sequence lengths [bs] cu_seqlens_k_dst: Destination cumulative sequence lengths [bs+1] page_table_1_dst: Destination page table [rows, stride] nsa_cache_seqlens_dst: Destination NSA cache sequence lengths [size] seqlens_expanded_dst: Optional destination expanded sequence lengths [size] (required for TARGET_VERIFY/DRAFT_EXTEND) nsa_cu_seqlens_k_dst: Destination NSA cumulative sequence lengths [size+1] real_page_table_dst: Optional destination real page table [rows, cols] flashmla_num_splits_dst: Optional destination FlashMLA num_splits [size+1] flashmla_metadata_dst: Optional destination FlashMLA metadata tensor forward_mode: Forward mode (0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND) bs: Batch size max_len: Maximum length for decode/draft_extend mode max_seqlen_k: Maximum sequence length for target_verify mode seqlens_expanded_size: Size of expanded sequence lengths """ # Determine template parameters for kernel specialization has_real_page_table = real_page_table_src is not None has_flashmla = flashmla_num_splits_src is not None # Get JIT-compiled module for this configuration (cached after first use) module = _jit_fused_metadata_copy_module( forward_mode, has_real_page_table, has_flashmla ) # Ensure all required source tensors are contiguous (required for kernel's linear indexing) # This matches the CHECK_INPUT checks in the verified sgl-kernel implementation cache_seqlens_src = cache_seqlens_src.contiguous() cu_seqlens_k_src = cu_seqlens_k_src.contiguous() page_indices_src = page_indices_src.contiguous() nsa_cache_seqlens_src = nsa_cache_seqlens_src.contiguous() if seqlens_expanded_src is not None: seqlens_expanded_src = seqlens_expanded_src.contiguous() nsa_cu_seqlens_k_src = nsa_cu_seqlens_k_src.contiguous() # Call JIT-compiled kernel (None values are passed as Optional with no value) module.fused_metadata_copy( cache_seqlens_src, cu_seqlens_k_src, page_indices_src, nsa_cache_seqlens_src, seqlens_expanded_src, nsa_cu_seqlens_k_src, real_page_table_src, flashmla_num_splits_src, flashmla_metadata_src, cache_seqlens_dst, cu_seqlens_k_dst, page_table_1_dst, nsa_cache_seqlens_dst, seqlens_expanded_dst, nsa_cu_seqlens_k_dst, real_page_table_dst, flashmla_num_splits_dst, flashmla_metadata_dst, bs, max_len, max_seqlen_k, seqlens_expanded_size, ) def fused_metadata_copy_multi_cuda( cache_seqlens_src: torch.Tensor, cu_seqlens_k_src: torch.Tensor, page_indices_src: torch.Tensor, nsa_cache_seqlens_src: torch.Tensor, nsa_cu_seqlens_k_src: torch.Tensor, real_page_table_src: Optional[torch.Tensor], flashmla_num_splits_src: Optional[torch.Tensor], flashmla_metadata_src: Optional[torch.Tensor], cache_seqlens_dst0: torch.Tensor, cu_seqlens_k_dst0: torch.Tensor, page_table_1_dst0: torch.Tensor, nsa_cache_seqlens_dst0: torch.Tensor, nsa_cu_seqlens_k_dst0: torch.Tensor, real_page_table_dst0: Optional[torch.Tensor], flashmla_num_splits_dst0: Optional[torch.Tensor], flashmla_metadata_dst0: Optional[torch.Tensor], cache_seqlens_dst1: torch.Tensor, cu_seqlens_k_dst1: torch.Tensor, page_table_1_dst1: torch.Tensor, nsa_cache_seqlens_dst1: torch.Tensor, nsa_cu_seqlens_k_dst1: torch.Tensor, real_page_table_dst1: Optional[torch.Tensor], flashmla_num_splits_dst1: Optional[torch.Tensor], flashmla_metadata_dst1: Optional[torch.Tensor], cache_seqlens_dst2: torch.Tensor, cu_seqlens_k_dst2: torch.Tensor, page_table_1_dst2: torch.Tensor, nsa_cache_seqlens_dst2: torch.Tensor, nsa_cu_seqlens_k_dst2: torch.Tensor, real_page_table_dst2: Optional[torch.Tensor], flashmla_num_splits_dst2: Optional[torch.Tensor], flashmla_metadata_dst2: Optional[torch.Tensor], bs: int, max_len: int, seqlens_expanded_size: int, ) -> None: """ Multi-backend fused metadata copy kernel for NSA backend CUDA graph replay. This function copies metadata from one source to THREE destinations in a single kernel launch, eliminating the overhead of 3 separate kernel calls. Currently only supports DECODE mode, which is the most common case. Args: cache_seqlens_src: Source cache sequence lengths [bs] cu_seqlens_k_src: Source cumulative sequence lengths [bs+1] page_indices_src: Source page indices [bs, max_len] nsa_cache_seqlens_src: Source NSA cache sequence lengths [bs] nsa_cu_seqlens_k_src: Source NSA cumulative sequence lengths [bs+1] real_page_table_src: Optional source real page table [bs, cols] flashmla_num_splits_src: Optional source FlashMLA num_splits [bs+1] flashmla_metadata_src: Optional source FlashMLA metadata tensor cache_seqlens_dst0-2: Destination cache sequence lengths for backends 0-2 cu_seqlens_k_dst0-2: Destination cumulative sequence lengths for backends 0-2 page_table_1_dst0-2: Destination page tables for backends 0-2 nsa_cache_seqlens_dst0-2: Destination NSA cache sequence lengths for backends 0-2 nsa_cu_seqlens_k_dst0-2: Destination NSA cumulative sequence lengths for backends 0-2 real_page_table_dst0-2: Optional destination real page tables for backends 0-2 flashmla_num_splits_dst0-2: Optional destination FlashMLA num_splits for backends 0-2 flashmla_metadata_dst0-2: Optional destination FlashMLA metadata tensors for backends 0-2 bs: Batch size max_len: Maximum length for decode mode seqlens_expanded_size: Size of expanded sequence lengths """ # Determine template parameters for kernel specialization has_real_page_table = real_page_table_src is not None has_flashmla = flashmla_num_splits_src is not None # Get JIT-compiled module for this configuration (cached after first use) module = _jit_fused_metadata_copy_multi_module(has_real_page_table, has_flashmla) # Ensure all source tensors are contiguous (required for kernel's linear indexing) # This matches the CHECK_INPUT checks in the verified sgl-kernel implementation cache_seqlens_src = cache_seqlens_src.contiguous() cu_seqlens_k_src = cu_seqlens_k_src.contiguous() page_indices_src = page_indices_src.contiguous() nsa_cache_seqlens_src = nsa_cache_seqlens_src.contiguous() nsa_cu_seqlens_k_src = nsa_cu_seqlens_k_src.contiguous() # Call JIT-compiled kernel (None values are passed as Optional with no value) module.fused_metadata_copy_multi( cache_seqlens_src, cu_seqlens_k_src, page_indices_src, nsa_cache_seqlens_src, nsa_cu_seqlens_k_src, real_page_table_src, flashmla_num_splits_src, flashmla_metadata_src, cache_seqlens_dst0, cu_seqlens_k_dst0, page_table_1_dst0, nsa_cache_seqlens_dst0, nsa_cu_seqlens_k_dst0, real_page_table_dst0, flashmla_num_splits_dst0, flashmla_metadata_dst0, cache_seqlens_dst1, cu_seqlens_k_dst1, page_table_1_dst1, nsa_cache_seqlens_dst1, nsa_cu_seqlens_k_dst1, real_page_table_dst1, flashmla_num_splits_dst1, flashmla_metadata_dst1, cache_seqlens_dst2, cu_seqlens_k_dst2, page_table_1_dst2, nsa_cache_seqlens_dst2, nsa_cu_seqlens_k_dst2, real_page_table_dst2, flashmla_num_splits_dst2, flashmla_metadata_dst2, bs, max_len, seqlens_expanded_size, ) ================================================ FILE: python/sglang/jit_kernel/fused_store_index_cache.py ================================================ """ This module provides JIT-compiled CUDA kernels for fusing multiple tensor copy operations into single kernel launches, reducing kernel launch overhead and improving CUDA graph replay performance. The kernels are compiled on-demand using TVM FFI and cached for subsequent use. """ from __future__ import annotations import logging from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import ( cache_once, is_arch_support_pdl, load_jit, make_cpp_args, ) if TYPE_CHECKING: from tvm_ffi.module import Module logger = logging.getLogger(__name__) @cache_once def _jit_nsa_fused_store_module( key_dtype: torch.dtype, indices_dtype: torch.dtype, page_size: int ) -> Module: """ Build a JIT module that exposes: module.fused_store_index_k_cache(input_bf16, index_k_with_scale_u8, loc_i64) """ args = make_cpp_args(key_dtype, indices_dtype, page_size, is_arch_support_pdl()) return load_jit( "fused_store_index_k_cache", *args, cuda_files=["nsa/fused_store_index_cache.cuh"], cuda_wrappers=[ ( "fused_store_index_k_cache", # - Float = bf16_t (sgl_kernel/type.cuh) # - IndicesT = int64_t (out_cache_loc is int64 in SGLang SetKAndS) # - kPageSize = 64 (CUDA NSA) f"FusedStoreCacheIndexerKernel<{args}>::run", ) ], ) @cache_once def can_use_nsa_fused_store( key_dtype: torch.dtype, indices_dtype: torch.dtype, page_size: int ) -> bool: logger = logging.getLogger(__name__) try: _jit_nsa_fused_store_module(key_dtype, indices_dtype, page_size) return True except Exception as e: logger.warning(f"Failed to load nsa fused store JIT kernel: {e}") return False def fused_store_index_k_cache( key: torch.Tensor, index_k_with_scale: torch.Tensor, out_cache_loc: torch.Tensor, page_size: int = 64, ) -> None: """ Fused: quantize bf16 key (N,128) -> fp8 + fp32 scale and write into NSATokenToKVPool.index_k_with_scale_buffer. key: (num_tokens, 128) bf16 (or reshapeable to it) index_k_with_scale: (num_pages, 64*(128+4)) uint8 out_cache_loc: (num_tokens,) int64 token indices in TokenToKVPool """ assert key.is_cuda assert index_k_with_scale.is_cuda assert out_cache_loc.is_cuda # 1) normalize shapes if key.dim() != 2: key = key.view(-1, key.shape[-1]) assert key.shape[1] == 128, f"expected key last-dim=128, got {key.shape}" # 2) dtypes assert key.dtype == torch.bfloat16, f"{key.dtype=}" assert index_k_with_scale.dtype == torch.uint8, f"{index_k_with_scale.dtype=}" assert out_cache_loc.dtype == torch.int64, f"{out_cache_loc.dtype=}" # 3) contiguity if not key.is_contiguous(): key = key.contiguous() if not out_cache_loc.is_contiguous(): out_cache_loc = out_cache_loc.contiguous() if not index_k_with_scale.is_contiguous(): index_k_with_scale = index_k_with_scale.contiguous() module = _jit_nsa_fused_store_module(key.dtype, out_cache_loc.dtype, page_size) module.fused_store_index_k_cache(key, index_k_with_scale, out_cache_loc) ================================================ FILE: python/sglang/jit_kernel/gptq_marlin.py ================================================ from __future__ import annotations from typing import TYPE_CHECKING, Optional import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from sgl_kernel.scalar_type import ScalarType from tvm_ffi.module import Module # Constants matching device::marlin:: in marlin.cuh _MAX_THREAD_N = 256 @cache_once def _jit_gptq_marlin_module(dtype: torch.dtype) -> Module: args = make_cpp_args(dtype) return load_jit( "gptq_marlin", *args, cuda_files=["gemm/marlin/gptq_marlin.cuh"], cuda_wrappers=[("gptq_marlin_gemm", f"gptq_marlin_gemm<{args}>")], ) def _or_empty( t: Optional[torch.Tensor], device: torch.device, dtype: torch.dtype ) -> torch.Tensor: return t if t is not None else torch.empty(0, device=device, dtype=dtype) def gptq_marlin_gemm( a: torch.Tensor, c: Optional[torch.Tensor], b_q_weight: torch.Tensor, b_scales: torch.Tensor, global_scale: Optional[torch.Tensor], b_zeros: Optional[torch.Tensor], g_idx: Optional[torch.Tensor], perm: Optional[torch.Tensor], workspace: torch.Tensor, b_q_type: ScalarType, size_m: int, size_n: int, size_k: int, is_k_full: bool = True, use_atomic_add: bool = False, use_fp32_reduce: bool = False, is_zp_float: bool = False, ) -> torch.Tensor: device = a.device # Allocate output if not provided if c is None: c = torch.empty((size_m, size_n), dtype=a.dtype, device=device) # Early return for zero-size M if size_m == 0: return c # Determine activation ordering has_act_order = ( g_idx is not None and perm is not None and g_idx.numel() > 0 and perm.numel() > 0 ) # Allocate c_tmp for fp32 reduce if use_fp32_reduce: sms = torch.cuda.get_device_properties(device).multi_processor_count max_m_block = min(((size_m + 15) // 16) * 16, 64) c_tmp = torch.empty( sms * max_m_block * _MAX_THREAD_N, dtype=torch.float32, device=device, ) else: c_tmp = torch.empty(0, dtype=torch.float32, device=device) # Allocate a_tmp for act_order column permutation if has_act_order: a_tmp = torch.empty((size_m, size_k), dtype=a.dtype, device=device) else: a_tmp = torch.empty(0, dtype=a.dtype, device=device) # Convert Optional tensors to empty tensors global_scale_t = _or_empty(global_scale, device, a.dtype) b_zeros_t = _or_empty(b_zeros, device, torch.int32) g_idx_t = _or_empty(g_idx, device, torch.int32) perm_t = _or_empty(perm, device, torch.int32) module = _jit_gptq_marlin_module(a.dtype) module.gptq_marlin_gemm( a, b_q_weight, b_scales, global_scale_t, b_zeros_t, g_idx_t, perm_t, c, c_tmp, a_tmp, workspace, b_q_type.id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float, ) return c ================================================ FILE: python/sglang/jit_kernel/gptq_marlin_repack.py ================================================ from __future__ import annotations from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import cache_once, load_jit if TYPE_CHECKING: from tvm_ffi.module import Module # Constants matching device::marlin:: in marlin.cuh _TILE_SIZE = 16 @cache_once def _jit_gptq_marlin_repack_module() -> Module: return load_jit( "gptq_marlin_repack", cuda_files=["gemm/marlin/gptq_marlin_repack.cuh"], cuda_wrappers=[("gptq_marlin_repack", "gptq_marlin_repack")], ) def gptq_marlin_repack( b_q_weight: torch.Tensor, perm: torch.Tensor, size_k: int, size_n: int, num_bits: int, ) -> torch.Tensor: pack_factor = 32 // num_bits # Allocate output tensor out = torch.empty( (size_k // _TILE_SIZE, size_n * _TILE_SIZE // pack_factor), dtype=b_q_weight.dtype, device=b_q_weight.device, ) module = _jit_gptq_marlin_repack_module() module.gptq_marlin_repack(b_q_weight, perm, out, size_k, size_n, num_bits) return out ================================================ FILE: python/sglang/jit_kernel/hadamard.py ================================================ from __future__ import annotations from typing import TYPE_CHECKING, Callable import torch from sglang.jit_kernel.utils import KERNEL_PATH, cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_hadamard_module(dtype: torch.dtype) -> Module: args = make_cpp_args(dtype) hadamard_include_dir = (KERNEL_PATH / "csrc" / "fast-hadamard-transform").resolve() return load_jit( "hadamard", *args, cuda_files=["fast-hadamard-transform/hadamard_jit.cuh"], cuda_wrappers=[ ("hadamard_transform", f"HadamardKernel<{args}>::run"), ("hadamard_transform_12n", f"Hadamard12NKernel<{args}>::run"), ("hadamard_transform_20n", f"Hadamard20NKernel<{args}>::run"), ("hadamard_transform_28n", f"Hadamard28NKernel<{args}>::run"), ("hadamard_transform_40n", f"Hadamard40NKernel<{args}>::run"), ], extra_include_paths=[str(hadamard_include_dir)], ) def _hadamard_transform_impl( x: torch.Tensor, scale: float, pad_multiple: int, kernel_fn: Callable, ) -> torch.Tensor: if not x.is_cuda: raise RuntimeError(f"{kernel_fn.__name__} only supports CUDA tensors") shapes_og = x.size() dim_og = x.size(-1) x = x.reshape(-1, dim_og) if x.stride(-1) != 1: x = x.contiguous() needs_pad = dim_og % pad_multiple != 0 if needs_pad: x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple)) out = torch.empty_like(x) kernel_fn(x, out, scale) if needs_pad: out = out[:, :dim_og] return out.reshape(shapes_og) def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: module = _jit_hadamard_module(x.dtype) return _hadamard_transform_impl(x, scale, 8, module.hadamard_transform) def hadamard_transform_12n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: module = _jit_hadamard_module(x.dtype) return _hadamard_transform_impl(x, scale, 4 * 12, module.hadamard_transform_12n) def hadamard_transform_20n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: module = _jit_hadamard_module(x.dtype) return _hadamard_transform_impl(x, scale, 4 * 20, module.hadamard_transform_20n) def hadamard_transform_28n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: module = _jit_hadamard_module(x.dtype) return _hadamard_transform_impl(x, scale, 4 * 28, module.hadamard_transform_28n) def hadamard_transform_40n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: module = _jit_hadamard_module(x.dtype) return _hadamard_transform_impl(x, scale, 4 * 40, module.hadamard_transform_40n) ================================================ FILE: python/sglang/jit_kernel/hicache.py ================================================ from __future__ import annotations import logging from typing import TYPE_CHECKING from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: import torch from tvm_ffi.module import Module DEFAULT_BLOCK_QUOTA = 2 @cache_once def _jit_hicache_module(*, element_size: int, unroll: int, block_quota: int) -> Module: args = make_cpp_args( element_size, unroll, block_quota, 1024, # num_threads, can be tuned for performance ) return load_jit( "hicache", *args, cuda_files=["hicache.cuh"], cuda_wrappers=[ ("launch_one", f"&HiCacheKernel<{args}>::run_one"), ("launch_all", f"&HiCacheKernel<{args}>::run_all"), ], ) def can_use_hicache_jit_kernel( *, element_size: int, unroll: int | None = None, # can be tuned for performance block_quota: int | None = None, # can be tuned for less interference ) -> bool: logger = logging.getLogger(__name__) if element_size % 128 != 0: logger.warning(f"Unsupported {element_size = } for JIT HiCache kernel") return False try: unroll = unroll or _default_unroll(element_size) block_quota = block_quota or DEFAULT_BLOCK_QUOTA _jit_hicache_module( element_size=element_size, unroll=unroll, block_quota=block_quota, ) return True except Exception as e: logger.warning(f"Failed to load JIT HiCache kernel: {e}") return False def _default_unroll(element_size: int) -> int: if element_size <= 512: return 4 if element_size <= 1024: return 2 # fallback: no unroll return 1 def transfer_hicache_one_layer( k_cache_dst: torch.Tensor, v_cache_dst: torch.Tensor, indices_dst: torch.Tensor, k_cache_src: torch.Tensor, v_cache_src: torch.Tensor, indices_src: torch.Tensor, *, element_dim: int | None = None, unroll: int | None = None, # can be tuned for performance block_quota: int | None = None, # can be tuned for less interference ) -> None: element_dim = element_dim or k_cache_dst.size(-1) k_cache_src = k_cache_src.view(-1, element_dim) v_cache_src = v_cache_src.view(-1, element_dim) k_cache_dst = k_cache_dst.view(-1, element_dim) v_cache_dst = v_cache_dst.view(-1, element_dim) element_size = element_dim * k_cache_dst.element_size() block_quota = block_quota or DEFAULT_BLOCK_QUOTA unroll = unroll or _default_unroll(element_size) module = _jit_hicache_module( element_size=element_size, unroll=unroll, block_quota=block_quota, ) module.launch_one( k_cache_dst, v_cache_dst, indices_dst, k_cache_src, v_cache_src, indices_src, ) def transfer_hicache_all_layer( k_ptr_dst: torch.Tensor, v_ptr_dst: torch.Tensor, indices_dst: torch.Tensor, k_ptr_src: torch.Tensor, v_ptr_src: torch.Tensor, indices_src: torch.Tensor, *, kv_cache_src_stride_bytes: int, kv_cache_dst_stride_bytes: int, element_size: int | None = None, unroll: int | None = None, # can be tuned for performance block_quota: int | None = None, # can be tuned for less interference ) -> None: if element_size is None: # assume both contiguous assert kv_cache_dst_stride_bytes == kv_cache_src_stride_bytes element_size = kv_cache_dst_stride_bytes block_quota = block_quota or DEFAULT_BLOCK_QUOTA unroll = unroll or _default_unroll(element_size) module = _jit_hicache_module( element_size=element_size, unroll=unroll, block_quota=block_quota, ) module.launch_all( k_ptr_dst, v_ptr_dst, indices_dst, k_ptr_src, v_ptr_src, indices_src, kv_cache_src_stride_bytes, kv_cache_dst_stride_bytes, ) ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh ================================================ /// \file atomic.cuh /// \brief Device-side atomic operations. #pragma once #include namespace device::atomic { /** * \brief Atomically computes the maximum of `*addr` and `value`, storing the * result in `*addr`. * \param addr Pointer to the value in global/shared memory to be updated. * \param value The value to compare against. * \return The old value at `*addr` before the update. * \note On CUDA, this uses `atomicMax`/`atomicMin` on the reinterpreted * integer representation. On ROCm, a CAS loop is used as a fallback. */ SGL_DEVICE float max(float* addr, float value) { #ifndef USE_ROCM float old; old = (value >= 0) ? __int_as_float(atomicMax((int*)addr, __float_as_int(value))) : __uint_as_float(atomicMin((unsigned int*)addr, __float_as_uint(value))); return old; #else int* addr_as_i = (int*)addr; int old = *addr_as_i, assumed; do { assumed = old; old = atomicCAS(addr_as_i, assumed, __float_as_int(fmaxf(value, __int_as_float(assumed)))); } while (assumed != old); return __int_as_float(old); #endif } } // namespace device::atomic ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/cta.cuh ================================================ /// \file cta.cuh /// \brief CTA (Cooperative Thread Array / thread-block) level primitives. #pragma once #include #include #include namespace device::cta { /** * \brief Compute the maximum of `value` across all threads in the CTA. * * Uses a two-level reduction: first within each warp via `warp::reduce_max`, * then across warps using shared memory. The final result is stored in * `smem[0]`. * * \tparam T Numeric type (must be supported by `warp::reduce_max`). * \param value Per-thread input value. * \param smem Shared memory buffer (must have at least `blockDim.x / 32` * elements). * \param min_value Identity element for max (default 0.0f). * \note This function does NOT issue a trailing `__syncthreads()`. * Callers must synchronize before reading `smem[0]`. */ template SGL_DEVICE void reduce_max(T value, float* smem, float min_value = 0.0f) { const uint32_t warp_id = threadIdx.x / kWarpThreads; smem[warp_id] = warp::reduce_max(value); __syncthreads(); if (warp_id == 0) { const auto tx = threadIdx.x; const auto local_value = tx * kWarpThreads < blockDim.x ? smem[tx] : min_value; const auto max_value = warp::reduce_max(local_value); smem[0] = max_value; } // no extra sync; it is caller's responsibility to sync if needed } } // namespace device::cta ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/impl/norm.cuh ================================================ #pragma once #include #include #include #include #include #include #include namespace host::norm { /** * \brief Check if the given configuration is supported. * \tparam T Element type (only fp16_t/bf16_t is supported) * \tparam kDim Dimension size (usually hidden size) */ template inline constexpr bool is_config_supported() { if (!std::is_same_v && !std::is_same_v) return false; if (kDim <= 256) { return (kDim == 64 || kDim == 128 || kDim == 256); } else { return (kDim % 256 == 0 && kDim <= 8192); } } /** * \brief Determine whether to use cta norm based on dimension size. * TL;DR: use warp norm for dim <= 256, cta norm otherwise. * \tparam T Element type (fp16_t or bf16_t) * \tparam kDim Dimension size (usually hidden size) * \note This function assumes that the configuration is supported. * \see `is_config_supported` */ template inline constexpr bool should_use_cta() { static_assert(is_config_supported(), "Unsupported norm configuration"); return kDim > 256; } /** * \brief Get the number of threads per CTA for cta norm. * \tparam T Element type (fp16_t or bf16_t) * \tparam kDim Dimension size (usually hidden size) * \return Number of threads per CTA */ template inline constexpr uint32_t get_cta_threads() { static_assert(should_use_cta()); return (kDim / 256) * device::kWarpThreads; } } // namespace host::norm namespace device::norm { namespace details { template SGL_DEVICE AlignedVector apply_norm_impl( const AlignedVector input, const AlignedVector weight, const float eps, [[maybe_unused]] float* smem_buffer, [[maybe_unused]] uint32_t num_warps) { float sum_of_squares = 0.0f; #pragma unroll for (auto i = 0u; i < N; ++i) { const auto fp32_input = cast(input[i]); sum_of_squares += fp32_input.x * fp32_input.x; sum_of_squares += fp32_input.y * fp32_input.y; } sum_of_squares = warp::reduce_sum(sum_of_squares); float norm_factor; if constexpr (kUseCTA) { // need to synchronize across the cta const auto warp_id = threadIdx.x / kWarpThreads; smem_buffer[warp_id] = sum_of_squares; __syncthreads(); // use the first warp to reduce if (warp_id == 0) { const auto tx = threadIdx.x; const auto local_sum = tx < num_warps ? smem_buffer[tx] : 0.0f; sum_of_squares = warp::reduce_sum(local_sum); smem_buffer[32] = math::rsqrt(sum_of_squares / kDim + eps); } __syncthreads(); norm_factor = smem_buffer[32]; } else { norm_factor = math::rsqrt(sum_of_squares / kDim + eps); } AlignedVector output; #pragma unroll for (auto i = 0u; i < N; ++i) { const auto fp32_input = cast(input[i]); const auto fp32_weight = cast(weight[i]); output[i] = cast({ fp32_input.x * norm_factor * fp32_weight.x, fp32_input.y * norm_factor * fp32_weight.y, }); } return output; } } // namespace details /** * \brief Apply norm using warp-level implementation. * \tparam kDim Dimension size * \tparam T Element type (fp16_t or bf16_t) * \param input Input vector * \param weight Weight vector * \param eps Epsilon value for numerical stability * \return Normalized output vector */ template SGL_DEVICE T apply_norm_warp(const T& input, const T& weight, float eps) { static_assert(kDim <= 256, "Warp norm only supports dim <= 256"); return details::apply_norm_impl(input, weight, eps, nullptr, 0); } /** * \brief Apply norm using CTA-level implementation. * \tparam kDim Dimension size * \tparam T Element type (fp16_t or bf16_t) * \param input Input vector * \param weight Weight vector * \param eps Epsilon value for numerical stability * \param smem Shared memory buffer * \param num_warps Number of warps in the CTA * \return Normalized output vector */ template SGL_DEVICE T apply_norm_cta( const T& input, const T& weight, float eps, float* smem, uint32_t num_warps = blockDim.x / kWarpThreads) { static_assert(kDim > 256, "CTA norm only supports dim > 256"); return details::apply_norm_impl(input, weight, eps, smem, num_warps); } /** * \brief Storage type for norm operation. * For warp norm, the storage size depends on kDim. * For cta norm, the storage size is fixed to 16B. * We will also pack the input 16-bit floats into 32-bit types * for faster CUDA core operations. * * \tparam T Element type (fp16_t or bf16_t) * \tparam kDim Dimension size */ template using StorageType = std::conditional_t< // storage type (kDim > 256), // whether to use cta norm AlignedVector, 4>, // cta norm storage, fixed to 16B AlignedVector, kDim / (2 * kWarpThreads)> // warp norm storage >; /** * \brief Minimum shared memory size (in bytes) required for cta norm. */ inline constexpr uint32_t kSmemBufferSize = 33; } // namespace device::norm ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/math.cuh ================================================ /// \file math.cuh /// \brief Device-side math helper functions and constants. /// /// Provides type-generic wrappers around CUDA math intrinsics by /// dispatching through `dtype_trait`. All functions are forced-inline /// device functions. #pragma once #include #include namespace device::math { /// \brief Constant: log2(e) inline constexpr float log2e = 1.44269504088896340736f; /// \brief Constant: ln(2) inline constexpr float loge2 = 0.693147180559945309417f; /// \brief Maximum representable value for FP8 E4M3 format. inline constexpr float FP8_E4M3_MAX = 448.0f; static_assert(log2e * loge2 == 1.0f, "log2e * loge2 must be 1"); /// \brief Returns the larger of `a` and `b`. template SGL_DEVICE T max(T a, T b) { return dtype_trait::max(a, b); } /// \brief Returns the smaller of `a` and `b`. template SGL_DEVICE T min(T a, T b) { return dtype_trait::min(a, b); } /// \brief Returns the absolute value of `a`. template SGL_DEVICE T abs(T a) { return dtype_trait::abs(a); } /// \brief Returns the square root of `a`. template SGL_DEVICE T sqrt(T a) { return dtype_trait::sqrt(a); } /// \brief Returns the reciprocal square root of `a` (i.e. 1 / sqrt(a)). template SGL_DEVICE T rsqrt(T a) { return dtype_trait::rsqrt(a); } /// \brief Returns e^a. template SGL_DEVICE T exp(T a) { return dtype_trait::exp(a); } /// \brief Returns sin(a). template SGL_DEVICE T sin(T a) { return dtype_trait::sin(a); } /// \brief Returns cos(a). template SGL_DEVICE T cos(T a) { return dtype_trait::cos(a); } } // namespace device::math ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh ================================================ /// \file runtime.cuh /// \brief Host-side CUDA runtime query helpers. /// /// Thin wrappers around CUDA occupancy and device-property APIs with /// automatic error checking via `RuntimeDeviceCheck`. #pragma once #include #include #include #include namespace host::runtime { // Return the maximum number of active blocks per SM for the given kernel template inline auto get_blocks_per_sm(T&& kernel, int32_t block_dim, std::size_t dynamic_smem = 0) -> uint32_t { int num_blocks_per_sm = 0; RuntimeDeviceCheck( cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm, kernel, block_dim, dynamic_smem)); return static_cast(num_blocks_per_sm); } // Return the number of SMs for the given device inline auto get_sm_count(int device_id) -> uint32_t { int sm_count; RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_id)); return static_cast(sm_count); } // Return the Major compute capability for the given device inline auto get_cc_major(int device_id) -> int { int cc_major; RuntimeDeviceCheck(cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device_id)); return cc_major; } // Return the runtime version inline auto get_runtime_version() -> int { int runtime_version; RuntimeDeviceCheck(cudaRuntimeGetVersion(&runtime_version)); return runtime_version; } // Return the maximum dynamic shared memory per block for the given kernel template inline auto get_available_dynamic_smem_per_block(T&& kernel, int num_blocks, int block_size) -> std::size_t { std::size_t smem_size; RuntimeDeviceCheck(cudaOccupancyAvailableDynamicSMemPerBlock(&smem_size, kernel, num_blocks, block_size)); return smem_size; } } // namespace host::runtime ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/scalar_type.hpp ================================================ #pragma once #include #include #ifndef __CUDACC__ #include #endif namespace host { // // ScalarType can represent a wide range of floating point and integer types, // in particular it can be used to represent sub-byte data types (something // that torch.dtype currently does not support). // // The type definitions on the Python side can be found in: vllm/scalar_type.py // these type definitions should be kept up to date with any Python API changes // here. // class ScalarType { public: enum NanRepr : uint8_t { NAN_NONE = 0, // nans are not supported NAN_IEEE_754 = 1, // nans are: exp all 1s, mantissa not all 0s NAN_EXTD_RANGE_MAX_MIN = 2, // nans are: exp all 1s, mantissa all 1s NAN_REPR_ID_MAX }; constexpr ScalarType( uint8_t exponent, uint8_t mantissa, bool signed_, int32_t bias, bool finite_values_only = false, NanRepr nan_repr = NAN_IEEE_754) : exponent(exponent), mantissa(mantissa), signed_(signed_), bias(bias), finite_values_only(finite_values_only), nan_repr(nan_repr) {}; static constexpr ScalarType int_(uint8_t size_bits, int32_t bias = 0) { return ScalarType(0, size_bits - 1, true, bias); } static constexpr ScalarType uint(uint8_t size_bits, int32_t bias = 0) { return ScalarType(0, size_bits, false, bias); } // IEEE 754 compliant floating point type static constexpr ScalarType float_IEEE754(uint8_t exponent, uint8_t mantissa) { assert(mantissa > 0 && exponent > 0); return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754); } // IEEE 754 non-compliant floating point type static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa, bool finite_values_only, NanRepr nan_repr) { assert(nan_repr < NAN_REPR_ID_MAX); assert(mantissa > 0 && exponent > 0); assert(nan_repr != NAN_IEEE_754); return ScalarType(exponent, mantissa, true, 0, finite_values_only, nan_repr); } uint8_t const exponent; // size of the exponent field (0 for integer types) uint8_t const mantissa; // size of the mantissa field (size of the integer // excluding the sign bit for integer types) bool const signed_; // flag if the type supports negative numbers (i.e. has a // sign bit) int32_t const bias; // stored values equal value + bias, // used for quantized type // Extra Floating point info bool const finite_values_only; // i.e. no +/-inf if true NanRepr const nan_repr; // how NaNs are represented // (not applicable for integer types) using Id = int64_t; private: // Field size in id template static constexpr size_t member_id_field_width() { using T = std::decay_t; return std::is_same_v ? 1 : sizeof(T) * 8; } template static constexpr auto reduce_members_helper(Fn f, Init val, Member member, Rest... rest) { auto new_val = f(val, member); if constexpr (sizeof...(rest) > 0) { return reduce_members_helper(f, new_val, rest...); } else { return new_val; }; } template constexpr auto reduce_members(Fn f, Init init) const { // Should be in constructor order for `from_id` return reduce_members_helper(f, init, exponent, mantissa, signed_, bias, finite_values_only, nan_repr); }; template static constexpr auto reduce_member_types(Fn f, Init init) { constexpr auto dummy_type = ScalarType(0, 0, false, 0, false, NAN_NONE); return dummy_type.reduce_members(f, init); }; static constexpr auto id_size_bits() { return reduce_member_types( [](int acc, auto member) -> int { return acc + member_id_field_width(); }, 0); } public: // unique id for this scalar type that can be computed at compile time for // c++17 template specialization this is not needed once we migrate to // c++20 and can pass literal classes as template parameters constexpr Id id() const { static_assert(id_size_bits() <= sizeof(Id) * 8, "ScalarType id is too large to be stored"); auto or_and_advance = [](std::pair result, auto member) -> std::pair { auto [id, bit_offset] = result; auto constexpr bits = member_id_field_width(); return {id | (int64_t(member) & ((uint64_t(1) << bits) - 1)) << bit_offset, bit_offset + bits}; }; return reduce_members(or_and_advance, std::pair{}).first; } // create a ScalarType from an id, for c++17 template specialization, // this is not needed once we migrate to c++20 and can pass literal // classes as template parameters static constexpr ScalarType from_id(Id id) { auto extract_and_advance = [id](auto result, auto member) { using T = decltype(member); auto [tuple, bit_offset] = result; auto constexpr bits = member_id_field_width(); auto extracted_val = static_cast((int64_t(id) >> bit_offset) & ((uint64_t(1) << bits) - 1)); auto new_tuple = std::tuple_cat(tuple, std::make_tuple(extracted_val)); return std::pair{new_tuple, bit_offset + bits}; }; auto [tuple_args, _] = reduce_member_types(extract_and_advance, std::pair, int>{}); return std::apply([](auto... args) { return ScalarType(args...); }, tuple_args); } constexpr int64_t size_bits() const { return mantissa + exponent + is_signed(); } constexpr bool is_signed() const { return signed_; } constexpr bool is_integer() const { return exponent == 0; } constexpr bool is_floating_point() const { return exponent > 0; } constexpr bool is_ieee_754() const { return is_floating_point() && finite_values_only == false && nan_repr == NAN_IEEE_754; } constexpr bool has_nans() const { return is_floating_point() && nan_repr != NAN_NONE; } constexpr bool has_infs() const { return is_floating_point() && finite_values_only == false; } constexpr bool has_bias() const { return bias != 0; } #ifndef __CUDACC__ private: double _floating_point_max() const { assert(mantissa <= 52 && exponent <= 11); uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1; if (nan_repr == NAN_EXTD_RANGE_MAX_MIN) { max_mantissa -= 1; } uint64_t max_exponent = (uint64_t(1) << exponent) - 2; if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) { assert(exponent < 11); max_exponent += 1; } // adjust the exponent to match that of a double // for now we assume the exponent bias is the standard 2^(e-1) -1, (where e // is the exponent bits), there is some precedent for non-standard biases, // example `float8_e4m3b11fnuz` here: https://github.com/jax-ml/ml_dtypes // but to avoid premature over complication we are just assuming the // standard exponent bias until there is a need to support non-standard // biases uint64_t exponent_bias = (uint64_t(1) << (exponent - 1)) - 1; uint64_t exponent_bias_double = (uint64_t(1) << 10) - 1; // double e = 11 uint64_t max_exponent_double = max_exponent - exponent_bias + exponent_bias_double; // shift the mantissa into the position for a double and // the exponent uint64_t double_raw = (max_mantissa << (52 - mantissa)) | (max_exponent_double << 52); return *reinterpret_cast(&double_raw); } constexpr std::variant _raw_max() const { if (is_floating_point()) { return {_floating_point_max()}; } else { assert(size_bits() < 64 || (size_bits() == 64 && is_signed())); return {(int64_t(1) << mantissa) - 1}; } } constexpr std::variant _raw_min() const { if (is_floating_point()) { assert(is_signed()); constexpr uint64_t sign_bit_double = (uint64_t(1) << 63); double max = _floating_point_max(); uint64_t max_raw = *reinterpret_cast(&max); uint64_t min_raw = max_raw | sign_bit_double; return {*reinterpret_cast(&min_raw)}; } else { assert(!is_signed() || size_bits() <= 64); if (is_signed()) { // set the top bit to 1 (i.e. INT64_MIN) and the rest to 0 // then perform an arithmetic shift right to set all the bits above // (size_bits() - 1) to 1 return {INT64_MIN >> (64 - size_bits())}; } else { return {int64_t(0)}; } } } public: // Max representable value for this scalar type. // (accounting for bias if there is one) constexpr std::variant max() const { return std::visit([this](auto x) -> std::variant { return {x - bias}; }, _raw_max()); } // Min representable value for this scalar type. // (accounting for bias if there is one) constexpr std::variant min() const { return std::visit([this](auto x) -> std::variant { return {x - bias}; }, _raw_min()); } #endif // __CUDACC__ public: std::string str() const { /* naming generally follows: https://github.com/jax-ml/ml_dtypes * for floating point types (leading f) the scheme is: * `float_em[flags]` * flags: * - no-flags: means it follows IEEE 754 conventions * - f: means finite values only (no infinities) * - n: means nans are supported (non-standard encoding) * for integer types the scheme is: * `[u]int[b]` * - if bias is not present it means its zero */ if (is_floating_point()) { auto ret = "float" + std::to_string(size_bits()) + "_e" + std::to_string(exponent) + "m" + std::to_string(mantissa); if (!is_ieee_754()) { if (finite_values_only) { ret += "f"; } if (nan_repr != NAN_NONE) { ret += "n"; } } return ret; } else { auto ret = ((is_signed()) ? "int" : "uint") + std::to_string(size_bits()); if (has_bias()) { ret += "b" + std::to_string(bias); } return ret; } } constexpr bool operator==(ScalarType const& other) const { return mantissa == other.mantissa && exponent == other.exponent && bias == other.bias && signed_ == other.signed_ && finite_values_only == other.finite_values_only && nan_repr == other.nan_repr; } }; using ScalarTypeId = ScalarType::Id; // "rust style" names generally following: // https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L60-L70 static inline constexpr auto kS4 = ScalarType::int_(4); static inline constexpr auto kU4 = ScalarType::uint(4); static inline constexpr auto kU4B8 = ScalarType::uint(4, 8); static inline constexpr auto kS8 = ScalarType::int_(8); static inline constexpr auto kU8 = ScalarType::uint(8); static inline constexpr auto kU8B128 = ScalarType::uint(8, 128); static inline constexpr auto kFE2M1f = ScalarType::float_(2, 1, true, ScalarType::NAN_NONE); static inline constexpr auto kFE3M2f = ScalarType::float_(3, 2, true, ScalarType::NAN_NONE); static inline constexpr auto kFE4M3fn = ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); static inline constexpr auto kFE8M0fnu = ScalarType(8, 0, false, 0, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); static inline constexpr auto kFE5M2 = ScalarType::float_IEEE754(5, 2); static inline constexpr auto kFE8M7 = ScalarType::float_IEEE754(8, 7); static inline constexpr auto kFE5M10 = ScalarType::float_IEEE754(5, 10); // Fixed width style names, generally following: // https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L47-L57 static inline constexpr auto kInt4 = kS4; static inline constexpr auto kUint4 = kU4; static inline constexpr auto kUint4b8 = kU4B8; static inline constexpr auto kInt8 = kS8; static inline constexpr auto kUint8 = kU8; static inline constexpr auto kUint8b128 = kU8B128; static inline constexpr auto kFloat4_e2m1f = kFE2M1f; static inline constexpr auto kFloat6_e3m2f = kFE3M2f; static inline constexpr auto kFloat8_e4m3fn = kFE4M3fn; static inline constexpr auto kFloat8_e5m2 = kFE5M2; static inline constexpr auto kFloat16_e8m7 = kFE8M7; static inline constexpr auto kFloat16_e5m10 = kFE5M10; // colloquial names static inline constexpr auto kHalf = kFE5M10; static inline constexpr auto kFloat16 = kHalf; static inline constexpr auto kBFloat16 = kFE8M7; static inline constexpr auto kFloat16Id = kFloat16.id(); } // namespace host ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/source_location.h ================================================ /// \file source_location.h /// \brief Portable `source_location` wrapper. /// /// Uses `std::source_location` when available (C++20), otherwise falls /// back to a minimal stub that returns empty/zero values. #pragma once #include /// NOTE: fallback to a minimal source_location implementation #if defined(__cpp_lib_source_location) #include using source_location_t = std::source_location; #else struct source_location_fallback { public: static constexpr source_location_fallback current() noexcept { return source_location_fallback{}; } constexpr source_location_fallback() noexcept = default; constexpr unsigned line() const noexcept { return 0; } constexpr unsigned column() const noexcept { return 0; } constexpr const char* file_name() const noexcept { return ""; } constexpr const char* function_name() const noexcept { return ""; } }; using source_location_t = source_location_fallback; #endif ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/tensor.h ================================================ /// \file tensor.h /// \brief Tensor validation and symbolic matching utilities. /// /// Provides the `TensorMatcher` fluent API for validating tensor shapes, /// strides, dtypes, and devices at kernel entry points, along with /// `SymbolicSize`, `SymbolicDType`, and `SymbolicDevice` for capturing /// and cross-checking tensor metadata across multiple tensors. /// /// See the "Tensor Checking" section in the JIT kernel dev guide for /// usage examples. #pragma once #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #ifdef __CUDACC__ #include #endif namespace host { namespace details { inline constexpr auto kAnyDeviceID = -1; inline constexpr auto kAnySize = static_cast(-1); inline constexpr auto kNullSize = static_cast(-1); inline constexpr auto kNullDType = static_cast(18u); inline constexpr auto kNullDevice = static_cast(-1); struct SizeRef; struct DTypeRef; struct DeviceRef; template struct _dtype_trait {}; template struct _dtype_trait { inline static constexpr DLDataType value = { .code = std::is_signed_v ? DLDataTypeCode::kDLInt : DLDataTypeCode::kDLUInt, .bits = static_cast(sizeof(T) * 8), .lanes = 1}; }; template struct _dtype_trait { inline static constexpr DLDataType value = { .code = DLDataTypeCode::kDLFloat, .bits = static_cast(sizeof(T) * 8), .lanes = 1}; }; #ifdef __CUDACC__ template <> struct _dtype_trait { inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat, .bits = 16, .lanes = 1}; }; template <> struct _dtype_trait { inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLBfloat, .bits = 16, .lanes = 1}; }; template <> struct _dtype_trait { inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat8_e4m3fn, .bits = 8, .lanes = 1}; }; #endif template struct _device_trait { inline static constexpr DLDevice value = {.device_type = Code, .device_id = kAnyDeviceID}; }; template inline constexpr auto kDTypeList = std::array{_dtype_trait::value...}; template inline constexpr auto kDeviceList = std::array{_device_trait::value...}; template struct PrintAbleSpan { explicit PrintAbleSpan(std::span data) : data(data) {} std::span data; }; // define DLDataType comparison and printing in root namespace inline constexpr auto kDeviceStringMap = [] { constexpr auto map = std::array, 16>{ std::pair{DLDeviceType::kDLCPU, "cpu"}, std::pair{DLDeviceType::kDLCUDA, "cuda"}, std::pair{DLDeviceType::kDLCUDAHost, "cuda_host"}, std::pair{DLDeviceType::kDLOpenCL, "opencl"}, std::pair{DLDeviceType::kDLVulkan, "vulkan"}, std::pair{DLDeviceType::kDLMetal, "metal"}, std::pair{DLDeviceType::kDLVPI, "vpi"}, std::pair{DLDeviceType::kDLROCM, "rocm"}, std::pair{DLDeviceType::kDLROCMHost, "rocm_host"}, std::pair{DLDeviceType::kDLExtDev, "ext_dev"}, std::pair{DLDeviceType::kDLCUDAManaged, "cuda_managed"}, std::pair{DLDeviceType::kDLOneAPI, "oneapi"}, std::pair{DLDeviceType::kDLWebGPU, "webgpu"}, std::pair{DLDeviceType::kDLHexagon, "hexagon"}, std::pair{DLDeviceType::kDLMAIA, "maia"}, std::pair{DLDeviceType::kDLTrn, "trn"}, }; constexpr auto max_type = stdr::max(map | stdv::keys); auto result = std::array{}; for (const auto& [code, name] : map) { result[static_cast(code)] = name; } return result; }(); struct PrintableDevice { DLDevice device; }; inline auto& operator<<(std::ostream& os, DLDevice device) { const auto& mapping = kDeviceStringMap; const auto entry = static_cast(device.device_type); RuntimeCheck(entry < mapping.size()); const auto name = mapping[entry]; RuntimeCheck(!name.empty(), "Unknown device: ", int(device.device_type)); os << name; if (device.device_id != kAnyDeviceID && device.device_type != DLDeviceType::kDLCPU) { os << ":" << device.device_id; } return os; } inline auto& operator<<(std::ostream& os, PrintableDevice pd) { return os << pd.device; } template inline auto& operator<<(std::ostream& os, PrintAbleSpan span) { os << "["; for (const auto i : irange(span.data.size())) { if (i > 0) { os << ", "; } os << span.data[i]; } os << "]"; return os; } } // namespace details /// \brief Check whether `dtype` matches the DLDataType for C++ type `T`. template inline bool is_type(DLDataType dtype) { return dtype == details::_dtype_trait::value; } /** * \brief A symbolic dimension size that can be bound once and * verified across multiple tensors. * * Create with an optional annotation string for error messages: * \code * auto N = SymbolicSize{"num_tokens"}; * \endcode * * Call `verify()` during tensor matching to either bind the first * observed value or check subsequent values match. Call `unwrap()` * to retrieve the bound value (panics if unset). */ struct SymbolicSize { public: SymbolicSize(std::string_view annotation = {}) : m_value(details::kNullSize), m_annotation(annotation) {} SymbolicSize(const SymbolicSize&) = delete; SymbolicSize& operator=(const SymbolicSize&) = delete; auto get_name() const -> std::string_view { return m_annotation; } auto set_value(int64_t value) -> void { RuntimeCheck(!this->has_value(), "Size value already set"); m_value = value; } auto has_value() const -> bool { return m_value != details::kNullSize; } auto get_value() const -> std::optional { return this->has_value() ? std::optional{m_value} : std::nullopt; } auto unwrap(DebugInfo info = {}) const -> int64_t { RuntimeCheck(info, this->has_value(), "Size value is not set"); return m_value; } auto verify(int64_t value, const char* prefix, int64_t dim) -> void { if (this->has_value()) { if (m_value != value) { [[unlikely]]; Panic("Size mismatch for ", m_name_str(prefix, dim), ": expected ", m_value, " but got ", value); } } else { this->set_value(value); } } auto value_or_name(const char* prefix, int64_t dim) const -> std::string { if (const auto value = this->get_value()) { return std::to_string(*value); } else { return m_name_str(prefix, dim); } } private: auto m_name_str(const char* prefix, int64_t dim) const -> std::string { std::ostringstream os; os << prefix << '#' << dim; if (!m_annotation.empty()) os << "('" << m_annotation << "')"; return std::move(os).str(); } std::int64_t m_value; std::string_view m_annotation; }; inline auto operator==(DLDevice lhs, DLDevice rhs) -> bool { return lhs.device_type == rhs.device_type && lhs.device_id == rhs.device_id; } /** * \brief A symbolic data type that can be constrained and verified. * * Optionally restrict allowed types via `set_options()`. * Use `verify()` to bind/check the dtype, and `unwrap()` to retrieve it. */ struct SymbolicDType { public: SymbolicDType() : m_value({details::kNullDType, 0, 0}) {} SymbolicDType(const SymbolicDType&) = delete; SymbolicDType& operator=(const SymbolicDType&) = delete; auto set_value(DLDataType value) -> void { RuntimeCheck(!this->has_value(), "Dtype value already set"); RuntimeCheck( m_check(value), "Dtype value [", value, "] not in the allowed options: ", details::PrintAbleSpan{m_options}); m_value = value; } auto has_value() const -> bool { return m_value.code != details::kNullDType; } auto get_value() const -> std::optional { return this->has_value() ? std::optional{m_value} : std::nullopt; } auto unwrap(DebugInfo info = {}) const -> DLDataType { RuntimeCheck(info, this->has_value(), "Dtype value is not set"); return m_value; } auto set_options(std::span options) -> void { m_options = options; } template auto set_options() -> void { m_options = details::kDTypeList; } auto verify(DLDataType dtype) -> void { if (this->has_value()) { RuntimeCheck(m_value == dtype, "DType mismatch: expected ", m_value, " but got ", dtype); } else { this->set_value(dtype); } } template auto is_type() const -> bool { return ::host::is_type(m_value); } private: auto m_check(DLDataType value) const -> bool { return stdr::empty(m_options) || (stdr::find(m_options, value) != stdr::end(m_options)); } std::span m_options; DLDataType m_value; }; /** * \brief A symbolic device that can be constrained and verified. * * Optionally restrict allowed device types via * `set_options()`. The device id can be wildcarded. */ struct SymbolicDevice { public: SymbolicDevice() : m_value({details::kNullDevice, details::kAnyDeviceID}) {} SymbolicDevice(const SymbolicDevice&) = delete; SymbolicDevice& operator=(const SymbolicDevice&) = delete; auto set_value(DLDevice value) -> void { RuntimeCheck(!this->has_value(), "Device value already set"); RuntimeCheck( m_check(value), "Device value [", details::PrintableDevice{value}, "] not in the allowed options: ", details::PrintAbleSpan{m_options}); m_value = value; } auto has_value() const -> bool { return m_value.device_type != details::kNullDevice; } auto get_value() const -> std::optional { return this->has_value() ? std::optional{m_value} : std::nullopt; } auto unwrap(DebugInfo info = {}) const -> DLDevice { RuntimeCheck(info, this->has_value(), "Device value is not set"); return m_value; } auto set_options(std::span options) -> void { m_options = options; } template auto set_options() -> void { m_options = details::kDeviceList; } auto verify(DLDevice device) -> void { if (this->has_value()) { RuntimeCheck( m_value == device, "Device mismatch: expected ", details::PrintableDevice{m_value}, " but got ", details::PrintableDevice{device}); } else { this->set_value(device); } } private: auto m_check(DLDevice value) const -> bool { return stdr::empty(m_options) || (stdr::any_of(m_options, [value](const DLDevice& opt) { // device type must exactly match if (opt.device_type != value.device_type) return false; // device id can be wildcarded return opt.device_id == details::kAnyDeviceID || opt.device_id == value.device_id; })); } std::span m_options; DLDevice m_value; }; namespace details { template struct BaseRef { public: BaseRef(const BaseRef&) = delete; BaseRef& operator=(const BaseRef&) = delete; auto operator->() const -> T* { return m_ref; } auto operator*() const -> T& { return *m_ref; } auto rebind(T& other) -> void { m_ref = &other; } explicit BaseRef() : m_ref(&m_cache), m_cache() {} BaseRef(T& size) : m_ref(&size), m_cache() {} private: T* m_ref; T m_cache; }; struct SizeRef : BaseRef { using BaseRef::BaseRef; SizeRef(int64_t value) { if (value != kAnySize) { (**this).set_value(value); } else { // otherwise, we can match any size } } }; struct DTypeRef : BaseRef { using BaseRef::BaseRef; DTypeRef(DLDataType options) { (**this).set_value(options); } DTypeRef(std::initializer_list options) { (**this).set_options(options); } DTypeRef(std::span options) { (**this).set_options(options); } }; struct DeviceRef : BaseRef { using BaseRef::BaseRef; DeviceRef(DLDevice options) { (**this).set_value(options); } DeviceRef(std::initializer_list options) { (**this).set_options(options); } DeviceRef(std::span options) { (**this).set_options(options); } }; } // namespace details /** * \brief Fluent API for validating tensor shape, strides, dtype, and device. * * Construct with the expected shape (using `SymbolicSize` or literal * integers), chain `.with_strides()`, `.with_dtype<...>()`, and * `.with_device<...>()`, then call `.verify(tensor)`. * * Example: * \code * auto N = SymbolicSize{"N"}; * TensorMatcher({N, 128}) * .with_dtype() * .with_device() * .verify(input_tensor); * \endcode * * \note `TensorMatcher` is a move-only temporary. Do not store in a variable. */ struct TensorMatcher { private: using SizeRef = details::SizeRef; using DTypeRef = details::DTypeRef; using DeviceRef = details::DeviceRef; public: TensorMatcher(const TensorMatcher&) = delete; TensorMatcher& operator=(const TensorMatcher&) = delete; explicit TensorMatcher(std::initializer_list shape) : m_shape(shape), m_strides(), m_dtype() {} auto with_strides(std::initializer_list strides) && -> TensorMatcher&& { // no partial update allowed RuntimeCheck(m_strides.size() == 0, "Strides already specified"); RuntimeCheck(m_shape.size() == strides.size(), "Strides size must match shape size"); m_strides = strides; return std::move(*this); } template auto with_dtype(DTypeRef&& dtype) && -> TensorMatcher&& { m_init_dtype(); m_dtype.rebind(*dtype); m_dtype->set_options(); return std::move(*this); } template auto with_dtype() && -> TensorMatcher&& { static_assert(sizeof...(Ts) > 0, "At least one dtype option must be specified"); m_init_dtype(); m_dtype->set_options(); return std::move(*this); } template auto with_device(DeviceRef&& device) && -> TensorMatcher&& { m_init_device(); m_device.rebind(*device); m_device->set_options(); return std::move(*this); } template auto with_device() && -> TensorMatcher&& { static_assert(sizeof...(Codes) > 0, "At least one device option must be specified"); m_init_device(); m_device->set_options(); return std::move(*this); } // once we start verification, we cannot modify anymore auto verify(tvm::ffi::TensorView view, DebugInfo info = {}) const&& -> const TensorMatcher&& { try { m_verify_impl(view); } catch (PanicError& e) { auto oss = std::ostringstream{}; oss << "Tensor match failed for "; s_print_tensor(oss, view); oss << " at " << info.file_name() << ":" << info.line() << "\n- Root cause: " << e.root_cause(); throw PanicError(std::move(oss).str()); } return std::move(*this); } private: static auto s_print_tensor(std::ostringstream& oss, tvm::ffi::TensorView view) -> void { oss << "Tensor<"; int64_t dim = 0; for (const auto& size : view.shape()) { if (dim++ > 0) oss << ", "; oss << size; } oss << ">[strides=<"; dim = 0; for (const auto& stride : view.strides()) { if (dim++ > 0) { oss << ", "; } oss << stride; } oss << ">, dtype=" << view.dtype(); oss << ", device=" << details::PrintableDevice{view.device()} << "]"; } auto m_verify_impl(tvm::ffi::TensorView view) const -> void { const auto dim = static_cast(view.dim()); RuntimeCheck(dim == m_shape.size(), "Tensor dimension mismatch: expected ", m_shape.size(), " but got ", dim); for (const auto i : irange(dim)) { m_shape[i]->verify(view.size(i), "shape", i); } if (m_has_strides()) { for (const auto i : irange(dim)) { if (view.size(i) != 1 || !m_strides[i]->has_value()) { // skip stride check for size 1 dimension m_strides[i]->verify(view.stride(i), "stride", i); } } } else { RuntimeCheck(view.is_contiguous(), "Tensor is not contiguous as expected"); } // since we may double verify, we will force to check m_dtype->verify(view.dtype()); m_device->verify(view.device()); } auto m_init_dtype() -> void { RuntimeCheck(!m_has_dtype, "DType already specified"); m_has_dtype = true; } auto m_init_device() -> void { RuntimeCheck(!m_has_device, "Device already specified"); m_has_device = true; } auto m_has_strides() const -> bool { return !m_strides.empty(); } std::span m_shape; std::span m_strides; DTypeRef m_dtype; DeviceRef m_device; bool m_has_dtype = false; bool m_has_device = false; }; } // namespace host ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/tile.cuh ================================================ /// \file tile.cuh /// \brief Tiled memory access helpers for coalesced global memory I/O. /// /// `tile::Memory` represents a contiguous memory region where multiple /// threads cooperatively load/store elements. The three factory methods /// determine the thread group: /// - `thread()` - single thread (no tiling). /// - `warp()` - all threads in a warp cooperate. /// - `cta()` - all threads in the CTA cooperate. #pragma once #include #include namespace device::tile { /** * \brief Represents a contiguous memory region for cooperative tiled access. * * Each instance is parameterized by an element type `T` and bound to a * specific thread id (`tid`) within a group of `tsize` threads. * * \tparam T The storage element type (e.g. `AlignedVector, 4>`). */ template struct Memory { public: SGL_DEVICE constexpr Memory(uint32_t tid, uint32_t tsize) : tid(tid), tsize(tsize) {} /// \brief Create a Memory accessor for a single thread (no cooperation). SGL_DEVICE static constexpr Memory thread() { return Memory{0, 1}; } /// \brief Create a Memory accessor distributed across warp threads. SGL_DEVICE static Memory warp(int warp_threads = kWarpThreads) { return Memory{static_cast(threadIdx.x % warp_threads), static_cast(warp_threads)}; } /// \brief Create a Memory accessor distributed across all CTA threads. SGL_DEVICE static Memory cta(int cta_threads = blockDim.x) { return Memory{static_cast(threadIdx.x), static_cast(cta_threads)}; } /// \brief Load one element from `ptr` at the position assigned to this thread. /// \param ptr Base pointer (cast to `const T*`). /// \param offset Optional tile offset (multiplied by `tsize`). SGL_DEVICE T load(const void* ptr, int64_t offset = 0) const { return static_cast(ptr)[tid + offset * tsize]; } /// \brief Store one element to `ptr` at the position assigned to this thread. SGL_DEVICE void store(void* ptr, T val, int64_t offset = 0) const { static_cast(ptr)[tid + offset * tsize] = val; } /// \brief Check whether this thread's element index is within bounds. SGL_DEVICE bool in_bound(int64_t element_count, int64_t offset = 0) const { return tid + offset * tsize < element_count; } private: uint32_t tid; uint32_t tsize; }; } // namespace device::tile ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/type.cuh ================================================ /// \file type.cuh /// \brief Dtype trait system for CUDA scalar/packed types. /// /// `dtype_trait` provides per-type metadata: packed type alias, /// conversion functions (`from`), and unary/binary math operations. /// Use `device::cast(from_value)` for type conversion on device. /// /// Registered types: /// | Scalar | Packed (x2) | Notes | /// |-----------|-------------|-------------------------------| /// | `fp32_t` | `fp32x2_t` | Full math ops (abs,sqrt,...) | /// | `fp16_t` | `fp16x2_t` | Conversion only | /// | `bf16_t` | `bf16x2_t` | Conversion only | /// | `fp32x2_t`| `fp32x4_t` | Packed float2 <-> half2/bf162 | #pragma once #include template struct dtype_trait {}; #define SGL_REGISTER_DTYPE_TRAIT(TYPE, PACK2, ...) \ template <> \ struct dtype_trait { \ using self_t = TYPE; \ using packed_t = PACK2; \ template \ SGL_DEVICE static self_t from(const S& value) { \ return static_cast(value); \ } \ __VA_ARGS__ \ } #define SGL_REGISTER_TYPE_END static_assert(true) #define SGL_REGISTER_FROM_FUNCTION(FROM, FN) \ SGL_DEVICE static self_t from(const FROM& x) { \ return FN(x); \ } \ static_assert(true) #define SGL_REGISTER_UNARY_FUNCTION(NAME, FN) \ SGL_DEVICE static self_t NAME(const self_t& x) { \ return FN(x); \ } \ static_assert(true) #define SGL_REGISTER_BINARY_FUNCTION(NAME, FN) \ SGL_DEVICE static self_t NAME(const self_t& x, const self_t& y) { \ return FN(x, y); \ } \ static_assert(true) SGL_REGISTER_DTYPE_TRAIT( fp32_t, fp32x2_t, SGL_REGISTER_TYPE_END; // SGL_REGISTER_FROM_FUNCTION(fp16_t, __half2float); SGL_REGISTER_FROM_FUNCTION(bf16_t, __bfloat162float); SGL_REGISTER_UNARY_FUNCTION(abs, fabsf); SGL_REGISTER_UNARY_FUNCTION(sqrt, sqrtf); SGL_REGISTER_UNARY_FUNCTION(rsqrt, rsqrtf); SGL_REGISTER_UNARY_FUNCTION(exp, expf); SGL_REGISTER_UNARY_FUNCTION(sin, sinf); SGL_REGISTER_UNARY_FUNCTION(cos, cosf); SGL_REGISTER_BINARY_FUNCTION(max, fmaxf); SGL_REGISTER_BINARY_FUNCTION(min, fminf);); SGL_REGISTER_DTYPE_TRAIT(fp16_t, fp16x2_t); SGL_REGISTER_DTYPE_TRAIT(bf16_t, bf16x2_t); /// TODO: Add ROCM implementation SGL_REGISTER_DTYPE_TRAIT( fp32x2_t, fp32x4_t, SGL_REGISTER_TYPE_END; SGL_REGISTER_FROM_FUNCTION(fp16x2_t, __half22float2); SGL_REGISTER_FROM_FUNCTION(bf16x2_t, __bfloat1622float2);); SGL_REGISTER_DTYPE_TRAIT( fp16x2_t, void, SGL_REGISTER_TYPE_END; SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22half2_rn);); SGL_REGISTER_DTYPE_TRAIT( bf16x2_t, void, SGL_REGISTER_TYPE_END; SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22bfloat162_rn);); #undef SGL_REGISTER_DTYPE_TRAIT #undef SGL_REGISTER_FROM_FUNCTION /// \brief Alias: the packed (x2) type for `T`. template using packed_t = typename dtype_trait::packed_t; namespace device { /** * \brief Cast a value from type `From` to type `To` on device. * * Dispatches through `dtype_trait::from()`, which uses the appropriate * CUDA intrinsic (e.g. `__half2float`, `__float22half2_rn`). */ template SGL_DEVICE To cast(const From& value) { return dtype_trait::from(value); } } // namespace device ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/utils.cuh ================================================ /// \file utils.cuh /// \brief Core CUDA/device utilities: type aliases, PDL helpers, /// typed pointer access, kernel launch wrapper, and error checking. /// /// This header is included (directly or transitively) by nearly every /// JIT kernel. It provides: /// - Scalar/packed type aliases (`fp16_t`, `bf16_t`, `fp8_e4m3_t`, ...). /// - `SGL_DEVICE` macro (forced-inline device function qualifier). /// - `kWarpThreads` constant (32). /// - PDL (Programmatic Dependent Launch) helpers for Hopper (sm_90+). /// - Typed `load_as` / `store_as` for void-pointer access. /// - `pointer::offset` for safe void-pointer arithmetic. /// - `host::LaunchKernel` - kernel launcher with optional PDL. /// - `host::RuntimeDeviceCheck` - CUDA error checking. #pragma once #include #include #include #include #include #include #ifndef USE_ROCM #include #include #include #include #else #include #include #include #ifndef __grid_constant__ #define __grid_constant__ #endif using cudaError_t = hipError_t; using cudaStream_t = hipStream_t; using cudaLaunchConfig_t = hipLaunchConfig_t; using cudaLaunchAttribute = hipLaunchAttribute; inline constexpr auto cudaSuccess = hipSuccess; #define cudaStreamPerThread hipStreamPerThread #define cudaGetErrorString hipGetErrorString #define cudaGetLastError hipGetLastError #define cudaLaunchKernel hipLaunchKernel #endif #ifndef USE_ROCM using fp32_t = float; using fp16_t = __half; using bf16_t = __nv_bfloat16; using fp8_e4m3_t = __nv_fp8_e4m3; using fp8_e5m2_t = __nv_fp8_e5m2; using fp32x2_t = float2; using fp16x2_t = __half2; using bf16x2_t = __nv_bfloat162; using fp8x2_e4m3_t = __nv_fp8x2_e4m3; using fp8x2_e5m2_t = __nv_fp8x2_e5m2; using fp32x4_t = float4; #else using fp32_t = float; using fp16_t = __half; using bf16_t = __hip_bfloat16; using fp8_e4m3_t = uint8_t; using fp8_e5m2_t = uint8_t; using fp32x2_t = float2; using fp16x2_t = half2; using bf16x2_t = __hip_bfloat162; using fp8x2_e4m3_t = uint16_t; using fp8x2_e5m2_t = uint16_t; using fp32x4_t = float4; #endif /* * LDG Support */ #ifndef USE_ROCM #define SGLANG_LDG(arg) __ldg(arg) #else #define SGLANG_LDG(arg) *(arg) #endif namespace device { /// \brief Macro: forced-inline device function qualifier. #define SGL_DEVICE __forceinline__ __device__ // Architecture detection: SGL_CUDA_ARCH is injected by load_jit() and is // available in both host and device compilation passes, whereas __CUDA_ARCH__ // is only defined by nvcc during the device pass. #if !defined(USE_ROCM) #if !defined(SGL_CUDA_ARCH) #error "SGL_CUDA_ARCH is not defined. JIT compilation must inject -DSGL_CUDA_ARCH via load_jit()." #endif #if defined(__CUDA_ARCH__) static_assert( __CUDA_ARCH__ == SGL_CUDA_ARCH, "SGL_CUDA_ARCH mismatch: injected arch flag does not match device target"); #endif #define SGL_ARCH_HOPPER_OR_GREATER (SGL_CUDA_ARCH >= 900) #define SGL_ARCH_BLACKWELL_OR_GREATER ((SGL_CUDA_ARCH >= 1000) && (CUDA_VERSION >= 12090)) #else // USE_ROCM #define SGL_ARCH_HOPPER_OR_GREATER 0 #define SGL_ARCH_BLACKWELL_OR_GREATER 0 #endif /// \brief Number of threads per warp (always 32 on NVIDIA/AMD GPUs). inline constexpr auto kWarpThreads = 32u; /// \brief Full warp active mask (all 32 lanes). inline constexpr auto kFullMask = 0xffffffffu; /** * \brief PDL (Programmatic Dependent Launch): wait for the primary kernel. * * On Hopper (sm_90+), inserts a `griddepcontrol.wait` instruction to * synchronize with a preceding kernel in the same stream. On older * architectures or ROCm this is a no-op. */ template SGL_DEVICE void PDLWaitPrimary() { #if SGL_ARCH_HOPPER_OR_GREATER if constexpr (kUsePDL) { asm volatile("griddepcontrol.wait;" ::: "memory"); } #endif } /** * \brief PDL: trigger dependent (secondary) kernel launch. * * On Hopper (sm_90+), inserts a `griddepcontrol.launch_dependents` * instruction. On older architectures or ROCm this is a no-op. */ template SGL_DEVICE void PDLTriggerSecondary() { #if SGL_ARCH_HOPPER_OR_GREATER if constexpr (kUsePDL) { asm volatile("griddepcontrol.launch_dependents;" :::); } #endif } /** * \brief Load data with the specified type and offset from a void pointer. * \tparam T The type to load. * \param ptr The base pointer. * \param offset The offset in number of elements of type T. */ template SGL_DEVICE T load_as(const void* ptr, int64_t offset = 0) { return static_cast(ptr)[offset]; } /** * \brief Store data with the specified type and offset to a void pointer. * \tparam T The type to store. * \param ptr The base pointer. * \param val The value to store. * \param offset The offset in number of elements of type T. * \note we use type_identity_t to force the caller to explicitly specify * the template parameter `T`, which can avoid accidentally using the wrong type. */ template SGL_DEVICE void store_as(void* ptr, std::type_identity_t val, int64_t offset = 0) { static_cast(ptr)[offset] = val; } /// \brief Safe void-pointer arithmetic (byte-level by default). namespace pointer { // we only allow void * pointer arithmetic for safety template SGL_DEVICE auto offset(void* ptr, U... offset) -> void* { return static_cast(ptr) + (... + offset); } template SGL_DEVICE auto offset(const void* ptr, U... offset) -> const void* { return static_cast(ptr) + (... + offset); } } // namespace pointer } // namespace device namespace host { /** * \brief Check the CUDA error code and panic with location info on failure. */ inline void RuntimeDeviceCheck(::cudaError_t error, DebugInfo location = {}) { if (error != ::cudaSuccess) { [[unlikely]]; ::host::panic(location, "CUDA error: ", ::cudaGetErrorString(error)); } } /// \brief Check the last CUDA error (calls `cudaGetLastError`). inline void RuntimeDeviceCheck(DebugInfo location = {}) { return RuntimeDeviceCheck(::cudaGetLastError(), location); } /** * \brief Kernel launcher with automatic stream resolution and PDL support. * * Usage: * \code * host::LaunchKernel(grid, block, device) * .enable_pdl(true) * (my_kernel, arg1, arg2); * \endcode * * The constructor resolves the CUDA stream from a `DLDevice` (via * `TVMFFIEnvGetStream`) or accepts a raw `cudaStream_t`. The call * operator launches the kernel and checks for errors. */ struct LaunchKernel { public: explicit LaunchKernel( dim3 grid_dim, dim3 block_dim, DLDevice device, std::size_t dynamic_shared_mem_bytes = 0, DebugInfo location = {}) noexcept : m_config(s_make_config(grid_dim, block_dim, resolve_device(device), dynamic_shared_mem_bytes)), m_location(location) {} explicit LaunchKernel( dim3 grid_dim, dim3 block_dim, cudaStream_t stream, std::size_t dynamic_shared_mem_bytes = 0, DebugInfo location = {}) noexcept : m_config(s_make_config(grid_dim, block_dim, stream, dynamic_shared_mem_bytes)), m_location(location) {} LaunchKernel(const LaunchKernel&) = delete; LaunchKernel& operator=(const LaunchKernel&) = delete; static auto resolve_device(DLDevice device) -> cudaStream_t { return static_cast(::TVMFFIEnvGetStream(device.device_type, device.device_id)); } auto enable_pdl(bool enabled = true) -> LaunchKernel& { #ifdef USE_ROCM (void)enabled; m_config.numAttrs = 0; #else if (enabled) { m_attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; m_attrs[0].val.programmaticStreamSerializationAllowed = true; m_config.numAttrs = 1; m_config.attrs = m_attrs; } else { m_config.numAttrs = 0; } #endif return *this; } template auto operator()(T&& kernel, Args&&... args) const -> void { #ifdef USE_ROCM hipLaunchKernelGGL( std::forward(kernel), m_config.gridDim, m_config.blockDim, m_config.dynamicSmemBytes, m_config.stream, std::forward(args)...); RuntimeDeviceCheck(m_location); #else RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward(args)...), m_location); #endif } private: static auto s_make_config( // Make a config for kernel launch dim3 grid_dim, dim3 block_dim, cudaStream_t stream, std::size_t smem) -> cudaLaunchConfig_t { auto config = ::cudaLaunchConfig_t{}; config.gridDim = grid_dim; config.blockDim = block_dim; config.dynamicSmemBytes = smem; config.stream = stream; config.numAttrs = 0; return config; } cudaLaunchConfig_t m_config; const DebugInfo m_location; cudaLaunchAttribute m_attrs[1]; }; } // namespace host ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/utils.h ================================================ /// \file utils.h /// \brief Host-side C++ utilities used by JIT kernel wrappers. /// /// Provides: /// - `DebugInfo` - wraps `std::source_location` for error reporting. /// - `RuntimeCheck` - runtime assertion with formatted error messages. /// - `Panic` - unconditional abort with formatted error messages. /// - `pointer::offset` - safe void-pointer arithmetic (host side). /// - `div_ceil` - integer ceiling division. /// - `dtype_bytes` - byte width of a `DLDataType`. /// - `irange` - Python-style integer range for range-for loops. #pragma once // ref: https://forums.developer.nvidia.com/t/c-20s-source-location-compilation-error-when-using-nvcc-12-1/258026/3 #ifdef __CUDACC__ #include #if CUDA_VERSION <= 12010 #pragma push_macro("__cpp_consteval") #pragma push_macro("_NODISCARD") #pragma push_macro("__builtin_LINE") #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wbuiltin-macro-redefined" #define __cpp_consteval 201811L #pragma clang diagnostic pop #ifdef _NODISCARD #undef _NODISCARD #define _NODISCARD #endif #define consteval constexpr #include "source_location.h" #undef consteval #pragma pop_macro("__cpp_consteval") #pragma pop_macro("_NODISCARD") #else // __CUDACC__ && CUDA_VERSION > 12010 #include "source_location.h" #endif #else // no __CUDACC__ #include "source_location.h" #endif #include #include #include #include #include #include #include namespace host { template inline constexpr bool dependent_false_v = false; /// \brief Source-location wrapper for debug/error messages. struct DebugInfo : public source_location_t { DebugInfo(source_location_t loc = source_location_t::current()) : source_location_t(loc) {} }; /// \brief Exception type thrown by `RuntimeCheck` and `Panic`. struct PanicError : public std::runtime_error { public: explicit PanicError(std::string msg) : runtime_error(msg), m_message(std::move(msg)) {} auto root_cause() const -> std::string_view { const auto str = std::string_view{m_message}; const auto pos = str.find(": "); return pos == std::string_view::npos ? str : str.substr(pos + 2); } private: std::string m_message; }; /// \brief Unconditionally abort with a formatted error message. template [[noreturn]] inline auto panic(DebugInfo location, Args&&... args) -> void { std::ostringstream os; os << "Runtime check failed at " << location.file_name() << ":" << location.line(); if constexpr (sizeof...(args) > 0) { os << ": "; (os << ... << std::forward(args)); } else { os << " in " << location.function_name(); } throw PanicError(std::move(os).str()); } /** * \brief Runtime assertion: panics with a formatted message when `condition` * is false. Extra `args` are streamed to the error message. * * Example: * \code * RuntimeCheck(n > 0, "n must be positive, got ", n); * \endcode */ template struct RuntimeCheck { template explicit RuntimeCheck(Cond&& condition, Args&&... args, DebugInfo location = {}) { if (condition) return; [[unlikely]] ::host::panic(location, std::forward(args)...); } template explicit RuntimeCheck(DebugInfo location, Cond&& condition, Args&&... args) { if (condition) return; [[unlikely]] ::host::panic(location, std::forward(args)...); } }; template struct Panic { explicit Panic(Args&&... args, DebugInfo location = {}) { ::host::panic(location, std::forward(args)...); } explicit Panic(DebugInfo location, Args&&... args) { ::host::panic(location, std::forward(args)...); } [[noreturn]] ~Panic() { std::terminate(); } }; template explicit RuntimeCheck(Cond&&, Args&&...) -> RuntimeCheck; template explicit RuntimeCheck(DebugInfo, Cond&&, Args&&...) -> RuntimeCheck; template explicit Panic(Args&&...) -> Panic; template explicit Panic(DebugInfo, Args&&...) -> Panic; namespace pointer { // we only allow void * pointer arithmetic for safety template inline auto offset(void* ptr, U... offset) -> void* { return static_cast(ptr) + (... + offset); } template inline auto offset(const void* ptr, U... offset) -> const void* { return static_cast(ptr) + (... + offset); } } // namespace pointer /// \brief Integer ceiling division: ceil(a / b). template inline constexpr auto div_ceil(T a, U b) { return (a + b - 1) / b; } /// \brief Returns the byte width of a DLPack data type. inline auto dtype_bytes(DLDataType dtype) -> std::size_t { return static_cast(dtype.bits / 8); } namespace stdr = std::ranges; namespace stdv = stdr::views; /// \brief Python-style integer range: `irange(n)` -> `[0, n)`. template inline auto irange(T end) { return stdv::iota(static_cast(0), end); } /// \brief Python-style integer range: `irange(start, end)` -> `[start, end)`. template inline auto irange(T start, T end) { return stdv::iota(start, end); } } // namespace host ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/vec.cuh ================================================ /// \file vec.cuh /// \brief Aligned vector types for coalesced global memory access. /// /// `AlignedVector` wraps `N` elements of type `T` in a naturally /// aligned struct so that the compiler emits wide (vectorized) load/store /// instructions (e.g. `LDG.128`). The maximum supported vector width is /// 256 bits (32 bytes), matching CUDA's widest vector load. #pragma once #include #include #include namespace device { namespace details { /// \brief Maps byte-width to the corresponding unsigned integer type. template struct uint_trait {}; template <> struct uint_trait<1> { using type = uint8_t; }; template <> struct uint_trait<2> { using type = uint16_t; }; template <> struct uint_trait<4> { using type = uint32_t; }; template <> struct uint_trait<8> { using type = uint64_t; }; /// \brief Alias: maps `sizeof(T)` to matching unsigned int type. template using sized_int = typename uint_trait::type; } // namespace details /// \brief Raw aligned storage for `N` elements of type `T`. template struct alignas(sizeof(T) * N) AlignedStorage { T data[N]; }; /** * \brief Aligned vector for vectorized memory access on GPU. * * Stores `N` elements of type `T` with natural alignment so that a single * `load`/`store` call compiles to a wide memory transaction. * * \tparam T Element type (e.g. `fp16_t`, `bf16_t`, `float`). * \tparam N Number of elements. Must be a power of two and * `sizeof(T) * N <= 32` (256 bits). * * Example: * \code * AlignedVector vec; // 16 bytes, 128-bit aligned * vec.load(input_ptr, tid); // vectorized load * vec[0] = vec[0] + 1; * vec.store(output_ptr, tid); // vectorized store * \endcode */ template struct AlignedVector { private: /// NOTE: N must be a power of two and sizeof(T) * N <= 32 bytes (256 bits) static_assert((N > 0 && (N & (N - 1)) == 0) && sizeof(T) * N <= 32, "CUDA only supports at most 256-bit vector op"); using element_t = typename details::sized_int; using storage_t = AlignedStorage; public: /// \brief Vectorized load from `ptr` at the given element `offset`. template SGL_DEVICE void load(const U* ptr, std::size_t offset = 0) { static_assert(std::is_same_v || std::is_same_v); m_storage = reinterpret_cast(ptr)[offset]; } /// \brief Vectorized store to `ptr` at the given element `offset`. template SGL_DEVICE void store(U* ptr, std::size_t offset = 0) const { static_assert(std::is_same_v || std::is_same_v); reinterpret_cast(ptr)[offset] = m_storage; } /// \brief Fill all N elements with the same `value`. SGL_DEVICE void fill(T value) { const auto store_value = *reinterpret_cast(&value); #pragma unroll for (std::size_t i = 0; i < N; ++i) { m_storage.data[i] = store_value; } } SGL_DEVICE auto operator[](std::size_t idx) -> T& { return reinterpret_cast(&m_storage)[idx]; } SGL_DEVICE auto operator[](std::size_t idx) const -> T { return reinterpret_cast(&m_storage)[idx]; } SGL_DEVICE auto data() -> T* { return reinterpret_cast(&m_storage); } SGL_DEVICE auto data() const -> const T* { return reinterpret_cast(&m_storage); } private: storage_t m_storage; }; } // namespace device ================================================ FILE: python/sglang/jit_kernel/include/sgl_kernel/warp.cuh ================================================ /// \file warp.cuh /// \brief Warp-level reduction primitives using `__shfl_xor_sync`. #pragma once #include namespace device::warp { /// \brief Full 32-thread active mask. static constexpr uint32_t kFullMask = 0xffffffffu; /** * \brief Warp-level sum reduction. * * Computes the sum of `value` across all active lanes specified by * `active_mask` using butterfly (XOR) shuffles. The result is * broadcast to all participating lanes. * * \tparam T Numeric type (e.g. float). * \param value Per-lane input value. * \param active_mask Bitmask of participating lanes (default: all 32). * \return The sum across all active lanes. */ template SGL_DEVICE T reduce_sum(T value, uint32_t active_mask = kFullMask) { #pragma unroll for (int mask = 16; mask > 0; mask >>= 1) value = value + __shfl_xor_sync(active_mask, value, mask, 32); return value; } /** * \brief Warp-level max reduction. * * Computes the maximum of `value` across all active lanes using * butterfly shuffles. The result is broadcast to all participating * lanes. * * \tparam T Numeric type (must be supported by `math::max`). * \param value Per-lane input value. * \param active_mask Bitmask of participating lanes (default: all 32). * \return The maximum across all active lanes. */ template SGL_DEVICE T reduce_max(T value, uint32_t active_mask = kFullMask) { #pragma unroll for (int mask = 16; mask > 0; mask >>= 1) value = math::max(value, __shfl_xor_sync(active_mask, value, mask, 32)); return value; } } // namespace device::warp ================================================ FILE: python/sglang/jit_kernel/kvcache.py ================================================ from __future__ import annotations import logging from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import ( cache_once, is_arch_support_pdl, load_jit, make_cpp_args, ) if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_kvcache_module(row_bytes: int) -> Module: args = make_cpp_args(row_bytes, is_arch_support_pdl()) return load_jit( "kvcache", *args, cuda_files=["elementwise/kvcache.cuh"], cuda_wrappers=[("store_cache", f"StoreKVCacheKernel<{args}>::run")], ) @cache_once def can_use_store_cache(size: int) -> bool: logger = logging.getLogger(__name__) if size % 4 != 0: logger.warning( f"Unsupported row_bytes={size} for JIT KV-Cache kernel:" " must be multiple of 4" ) return False try: _jit_kvcache_module(size) return True except Exception as e: logger.warning( f"Failed to load JIT KV-Cache kernel " f"with row_bytes={size}: {e}" ) return False def store_cache( k: torch.Tensor, v: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, indices: torch.Tensor, *, row_bytes: int = 0, num_split: int = 0, # can be tuned for performance ) -> None: """Store key and value tensors into KV cache at specified indices. Args: k (torch.Tensor): Key tensor of shape (batch_size, H * D). v (torch.Tensor): Value tensor of shape (batch_size, H * D). k_cache (torch.Tensor): Key cache tensor of shape (num_pages, H * D). v_cache (torch.Tensor): Value cache tensor of shape (num_pages, H * D). indices (torch.Tensor): Indices tensor of shape (batch_size,). """ row_bytes = row_bytes or k.shape[-1] * k.element_size() module = _jit_kvcache_module(row_bytes) if num_split <= 0: if row_bytes % 2048 == 0: num_split = 4 elif row_bytes % 1024 == 0: num_split = 2 else: num_split = 1 module.store_cache( k, v, k_cache, v_cache, indices, num_split, ) ================================================ FILE: python/sglang/jit_kernel/moe_lora_align.py ================================================ from __future__ import annotations from typing import TYPE_CHECKING, Optional import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_moe_align_module(dtype: torch.dtype) -> Module: args = make_cpp_args(dtype) return load_jit( "moe_lora_align_block_size", *args, cuda_files=["lora/moe_lora_align_kernel.cu"], cuda_wrappers=[ ("moe_lora_align_block_size", f"MoeLoraAlignBlockSizeKernel<{args}>::run"), ], ) def moe_lora_align_block_size( topk_ids: torch.Tensor, seg_indptr: torch.Tensor, req_to_lora: torch.Tensor, num_experts: int, block_size: int, max_loras: int, max_num_tokens_padded: int, max_num_m_blocks: int, sorted_token_ids: torch.Tensor, expert_ids: torch.Tensor, num_tokens_post_pad: torch.Tensor, adapter_enabled: torch.Tensor, lora_ids: torch.Tensor, maybe_expert_map: Optional[torch.Tensor] = None, ) -> None: module = _jit_moe_align_module(topk_ids.dtype) cumsum_buffer = torch.zeros( max_loras * (num_experts + 1), dtype=torch.int32, device=topk_ids.device ) token_mask = torch.empty( (max_loras * topk_ids.shape[0],), dtype=torch.int32, device=topk_ids.device ) module.moe_lora_align_block_size( topk_ids, seg_indptr, req_to_lora, num_experts, block_size, max_loras, max_num_tokens_padded, max_num_m_blocks, sorted_token_ids, expert_ids, num_tokens_post_pad, adapter_enabled, lora_ids, maybe_expert_map, cumsum_buffer, token_mask, ) ================================================ FILE: python/sglang/jit_kernel/moe_wna16_marlin.py ================================================ from __future__ import annotations from typing import TYPE_CHECKING, Optional import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from sgl_kernel.scalar_type import ScalarType from tvm_ffi.module import Module # Constants matching device::marlin_moe:: in marlin.cuh _MAX_THREAD_N = 256 @cache_once def _jit_moe_wna16_marlin_module(dtype: torch.dtype) -> Module: args = make_cpp_args(dtype) return load_jit( "moe_wna16_marlin", *args, cuda_files=["gemm/marlin_moe/moe_wna16_marlin.cuh"], cuda_wrappers=[ ( "moe_wna16_marlin_gemm", f"moe_wna16_marlin_gemm<{args}>", ) ], ) def _or_empty( t: Optional[torch.Tensor], device: torch.device, dtype: torch.dtype ) -> torch.Tensor: return t if t is not None else torch.empty(0, device=device, dtype=dtype) def moe_wna16_marlin_gemm( a: torch.Tensor, c_or_none: Optional[torch.Tensor], b_q_weight: torch.Tensor, b_bias_or_none: Optional[torch.Tensor], b_scales: torch.Tensor, global_scale_or_none: Optional[torch.Tensor], b_zeros_or_none: Optional[torch.Tensor], g_idx_or_none: Optional[torch.Tensor], perm_or_none: Optional[torch.Tensor], workspace: torch.Tensor, sorted_token_ids: torch.Tensor, expert_ids: torch.Tensor, num_tokens_post_padded: torch.Tensor, topk_weights: torch.Tensor, moe_block_size: int, top_k: int, mul_topk_weights: bool, is_ep: bool, b_q_type: ScalarType, size_m: int, size_n: int, size_k: int, is_k_full: bool = True, use_atomic_add: bool = False, use_fp32_reduce: bool = False, is_zp_float: bool = False, ) -> torch.Tensor: device = a.device # Allocate output if not provided if c_or_none is not None: c = c_or_none else: c = torch.empty((size_m * top_k, size_n), dtype=a.dtype, device=device) # Early return for zero-size M if size_m == 0: return c # Determine activation ordering has_act_order = ( g_idx_or_none is not None and perm_or_none is not None and g_idx_or_none.numel() > 0 and perm_or_none.numel() > 0 and g_idx_or_none.size(-1) > 0 and perm_or_none.size(-1) > 0 ) # Determine has_zp has_zp = b_zeros_or_none is not None and b_zeros_or_none.numel() > 0 # Determine has_bias has_bias = b_bias_or_none is not None # Derive num_groups and group_size from b_scales num_groups = b_scales.size(1) if has_act_order: if is_k_full: group_size = size_k // num_groups else: group_size = 0 else: if num_groups > 1: group_size = size_k // num_groups else: group_size = -1 # Allocate a_tmp for act_order column permutation if has_act_order: a_tmp = torch.empty((size_m * top_k, size_k), dtype=a.dtype, device=device) else: a_tmp = torch.empty(0, dtype=a.dtype, device=device) # Allocate c_tmp for fp32 reduce if use_fp32_reduce and not use_atomic_add: sms = torch.cuda.get_device_properties(device).multi_processor_count # max num of threadblocks is sms * 4 max_c_tmp_size = min( size_n * sorted_token_ids.size(0), sms * 4 * moe_block_size * _MAX_THREAD_N, ) if moe_block_size == 8: max_c_tmp_size *= 2 c_tmp = torch.empty(max_c_tmp_size, dtype=torch.float32, device=device) else: c_tmp = torch.empty(0, dtype=torch.float32, device=device) # Convert Optional tensors to empty tensors g_idx_t = _or_empty(g_idx_or_none, device, torch.int32) perm_t = _or_empty(perm_or_none, device, torch.int32) b_zeros_t = _or_empty(b_zeros_or_none, device, a.dtype) b_bias_t = _or_empty(b_bias_or_none, device, a.dtype) global_scale_t = _or_empty(global_scale_or_none, device, a.dtype) module = _jit_moe_wna16_marlin_module(a.dtype) module.moe_wna16_marlin_gemm( a, c, b_q_weight, b_bias_t, b_scales, global_scale_t, b_zeros_t, g_idx_t, perm_t, workspace, sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights, a_tmp, c_tmp, moe_block_size, top_k, mul_topk_weights, is_ep, b_q_type.id, size_m, size_n, size_k, has_act_order, has_bias, is_k_full, has_zp, num_groups, group_size, use_atomic_add, use_fp32_reduce, is_zp_float, ) return c ================================================ FILE: python/sglang/jit_kernel/ngram_embedding.py ================================================ from __future__ import annotations from typing import TYPE_CHECKING from sglang.jit_kernel.utils import cache_once, load_jit if TYPE_CHECKING: import torch from tvm_ffi.module import Module @cache_once def _jit_ngram_embedding_module() -> Module: return load_jit( "ngram_embedding", cuda_files=["ngram_embedding.cuh"], cuda_wrappers=[ ("compute_n_gram_ids", "&NgramEmbeddingKernel::compute_n_gram_ids"), ("update_token_table", "&NgramEmbeddingKernel::update_token_table"), ], ) def compute_n_gram_ids( ne_n: int, ne_k: int, ne_weights: torch.Tensor, ne_mods: torch.Tensor, exclusive_ne_embedder_size_sums: torch.Tensor, tokens: torch.Tensor, exclusive_req_len_sums: torch.Tensor, ne_token_table: torch.Tensor, row_indices: torch.Tensor, column_starts: torch.Tensor, n_gram_ids: torch.Tensor, ) -> None: """ Compute n-gram IDs for embedding. Args: ne_n: n value for n-gram ne_k: k value for n-gram configurations ne_weights: weights tensor with shape [ne_n-1, ne_k, ne_n] ne_mods: mods tensor with shape [ne_n-1, ne_k] exclusive_ne_embedder_size_sums: exclusive sum of embedder sizes tokens: input token ids exclusive_req_len_sums: exclusive sum of request lengths ne_token_table: token table for all requests row_indices: row indices for each request column_starts: column start positions for each request n_gram_ids: output tensor for n-gram ids """ module = _jit_ngram_embedding_module() module.compute_n_gram_ids( ne_n, ne_k, ne_weights, ne_mods, exclusive_ne_embedder_size_sums, tokens, exclusive_req_len_sums, ne_token_table, row_indices, column_starts, n_gram_ids, ) def update_token_table( tokens: torch.Tensor, ne_token_table: torch.Tensor, row_indices: torch.Tensor, column_starts: torch.Tensor, req_lens: torch.Tensor, ignore_tokens: torch.Tensor | None = None, ) -> None: """ Update the token table with new tokens. Args: tokens: input token ids ne_token_table: token table for all requests row_indices: row indices for each request column_starts: column start positions for each request req_lens: request lengths ignore_tokens: tokens to be ignored (marked as negative in table) """ module = _jit_ngram_embedding_module() if ignore_tokens is None: # Create an empty tensor for ignore_tokens ignore_tokens = tokens.new_empty(0, dtype=tokens.dtype) module.update_token_table( tokens, ne_token_table, row_indices, column_starts, req_lens, ignore_tokens, ) ================================================ FILE: python/sglang/jit_kernel/norm.py ================================================ from __future__ import annotations import logging from typing import TYPE_CHECKING, Optional import torch from sglang.jit_kernel.utils import ( cache_once, is_arch_support_pdl, load_jit, make_cpp_args, ) if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_qknorm_module(head_dim: int, dtype: torch.dtype) -> Module: args = make_cpp_args(head_dim, is_arch_support_pdl(), dtype) return load_jit( "qknorm", *args, cuda_files=["elementwise/qknorm.cuh"], cuda_wrappers=[("qknorm", f"QKNormKernel<{args}>::run")], ) @cache_once def _jit_rmsnorm_module(hidden_size: int, dtype: torch.dtype) -> Module: args = make_cpp_args(hidden_size, is_arch_support_pdl(), dtype) return load_jit( "rmsnorm", *args, cuda_files=["elementwise/rmsnorm.cuh"], cuda_wrappers=[("rmsnorm", f"RMSNormKernel<{args}>::run")], ) @cache_once def _jit_fused_add_rmsnorm_module(dtype: torch.dtype) -> Module: args = make_cpp_args(dtype) return load_jit( "fused_add_rmsnorm", *args, cuda_files=["elementwise/fused_add_rmsnorm.cuh"], cuda_wrappers=[("fused_add_rmsnorm", f"FusedAddRMSNormKernel<{args}>::run")], ) @cache_once def _jit_qknorm_across_heads_module(dtype: torch.dtype) -> Module: args = make_cpp_args(dtype) return load_jit( "qknorm_across_heads", *args, cuda_files=["elementwise/qknorm_across_heads.cuh"], cuda_wrappers=[ ("qknorm_across_heads", f"QKNormAcrossHeadsKernel<{args}>::run") ], ) @cache_once def can_use_fused_inplace_qknorm(head_dim: int, dtype: torch.dtype) -> bool: logger = logging.getLogger(__name__) if head_dim not in [64, 128, 256, 512, 1024]: logger.warning(f"Unsupported head_dim={head_dim} for JIT QK-Norm kernel") return False try: _jit_qknorm_module(head_dim, dtype) return True except Exception as e: logger.warning(f"Failed to load JIT QK-Norm kernel: {e}") return False def fused_inplace_qknorm( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, eps: float = 1e-6, *, head_dim: int = 0, ) -> None: head_dim = head_dim or q.size(-1) module = _jit_qknorm_module(head_dim, q.dtype) module.qknorm(q, k, q_weight, k_weight, eps) def rmsnorm( input: torch.Tensor, weight: torch.Tensor, output: Optional[torch.Tensor] = None, eps: float = 1e-6, ) -> None: output = output if output is not None else input hidden_size = input.size(-1) module = _jit_rmsnorm_module(hidden_size, input.dtype) module.rmsnorm(input, weight, output, eps) def fused_add_rmsnorm( input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6, ) -> None: module = _jit_fused_add_rmsnorm_module(input.dtype) module.fused_add_rmsnorm(input, residual, weight, eps) def fused_inplace_qknorm_across_heads( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, eps: float = 1e-6, ) -> None: """ Fused inplace QK normalization across all heads. Args: q: Query tensor of shape [batch_size, num_heads * head_dim] k: Key tensor of shape [batch_size, num_heads * head_dim] q_weight: Query weight tensor of shape [num_heads * head_dim] k_weight: Key weight tensor of shape [num_heads * head_dim] eps: Epsilon for numerical stability """ module = _jit_qknorm_across_heads_module(q.dtype) module.qknorm_across_heads(q, k, q_weight, k_weight, eps) ================================================ FILE: python/sglang/jit_kernel/nvfp4.py ================================================ from __future__ import annotations import importlib.util import os import pathlib from contextlib import contextmanager from typing import TYPE_CHECKING, Optional, Tuple import torch from sglang.jit_kernel.utils import cache_once, load_jit from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: from tvm_ffi.module import Module _FLOAT4_E2M1_MAX = 6.0 _FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max def _find_package_root(package: str) -> Optional[pathlib.Path]: spec = importlib.util.find_spec(package) if spec is None or spec.origin is None: return None return pathlib.Path(spec.origin).resolve().parent def _resolve_cutlass_include_paths() -> list[str]: include_paths: list[str] = [] flashinfer_root = _find_package_root("flashinfer") if flashinfer_root is not None: candidates = [ flashinfer_root / "data" / "cutlass" / "include", flashinfer_root / "data" / "cutlass" / "tools" / "util" / "include", ] for path in candidates: if path.exists(): include_paths.append(str(path)) deep_gemm_root = _find_package_root("deep_gemm") if deep_gemm_root is not None: candidate = deep_gemm_root / "include" if candidate.exists(): include_paths.append(str(candidate)) # De-duplicate while preserving order. unique_paths = [] seen = set() for path in include_paths: if path in seen: continue seen.add(path) unique_paths.append(path) return unique_paths def _nvfp4_cuda_flags() -> list[str]: return [ "-DNDEBUG", "-DFLASHINFER_ENABLE_F16", "-DCUTE_USE_PACKED_TUPLE=1", "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1", "-DCUTLASS_VERSIONS_GENERATED", "-DCUTLASS_TEST_LEVEL=0", "-DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1", "-DCUTLASS_DEBUG_TRACE_LEVEL=0", "--expt-extended-lambda", ] def _get_nvfp4_cuda_arch_list() -> str: if not torch.cuda.is_available(): raise RuntimeError("NVFP4 JIT kernels require CUDA.") major, minor = torch.cuda.get_device_capability() if major < 10: raise RuntimeError( f"NVFP4 JIT kernels require compute capability >= 10.0, got {major}.{minor}." ) # NVFP4 kernels use architecture-family-specific instructions and must be # compiled for `sm_*a` targets (e.g. sm_100a), not plain sm_100. # JIT compilation targets only the current device, unlike AOT fat-binaries; # adding extra architectures here would clash with the single SGL_CUDA_ARCH # value injected by load_jit(). return f"{major}.{minor}a" @contextmanager def _nvfp4_arch_env(): key = "TVM_FFI_CUDA_ARCH_LIST" old_val = os.environ.get(key) os.environ[key] = _get_nvfp4_cuda_arch_list() try: yield finally: if old_val is None: os.environ.pop(key, None) else: os.environ[key] = old_val @cache_once def _jit_nvfp4_quant_module() -> Module: extra_include_paths = _resolve_cutlass_include_paths() if not extra_include_paths: raise RuntimeError( "Cannot find CUTLASS headers required for NVFP4 JIT quantization. " "Please install flashinfer or deep_gemm with CUTLASS headers." ) with _nvfp4_arch_env(): return load_jit( "nvfp4_quant", cuda_files=[ "gemm/nvfp4/nvfp4_quant_kernels.cuh", ], cuda_wrappers=[ ("scaled_fp4_quant", "scaled_fp4_quant_sm100a_sm120a"), ], extra_include_paths=extra_include_paths, extra_cuda_cflags=_nvfp4_cuda_flags(), ) @cache_once def _jit_nvfp4_expert_quant_module() -> Module: extra_include_paths = _resolve_cutlass_include_paths() if not extra_include_paths: raise RuntimeError( "Cannot find CUTLASS headers required for NVFP4 JIT expert quantization. " "Please install flashinfer or deep_gemm with CUTLASS headers." ) with _nvfp4_arch_env(): return load_jit( "nvfp4_expert_quant", cuda_files=[ "gemm/nvfp4/nvfp4_expert_quant.cuh", ], cuda_wrappers=[ ("scaled_fp4_experts_quant", "scaled_fp4_experts_quant_sm100a"), ( "silu_and_mul_scaled_fp4_experts_quant", "silu_and_mul_scaled_fp4_experts_quant_sm100a", ), ], extra_include_paths=extra_include_paths, extra_cuda_cflags=_nvfp4_cuda_flags(), ) @cache_once def _jit_nvfp4_scaled_mm_module() -> Module: extra_include_paths = _resolve_cutlass_include_paths() if not extra_include_paths: raise RuntimeError( "Cannot find CUTLASS headers required for NVFP4 JIT GEMM. " "Please install flashinfer or deep_gemm with CUTLASS headers." ) with _nvfp4_arch_env(): return load_jit( "nvfp4_scaled_mm", cuda_files=[ "gemm/nvfp4/nvfp4_scaled_mm_kernels.cuh", "gemm/nvfp4/nvfp4_scaled_mm_entry.cuh", ], cuda_wrappers=[("cutlass_scaled_fp4_mm", "cutlass_scaled_fp4_mm")], extra_include_paths=extra_include_paths, extra_cuda_cflags=_nvfp4_cuda_flags(), ) @cache_once def _jit_nvfp4_blockwise_moe_module() -> Module: extra_include_paths = _resolve_cutlass_include_paths() if not extra_include_paths: raise RuntimeError( "Cannot find CUTLASS headers required for NVFP4 JIT MoE grouped GEMM. " "Please install flashinfer or deep_gemm with CUTLASS headers." ) with _nvfp4_arch_env(): return load_jit( "nvfp4_blockwise_moe", cuda_files=[ "moe/nvfp4_blockwise_moe.cuh", ], cuda_wrappers=[ ("cutlass_fp4_group_mm", "cutlass_fp4_group_mm_sm100a_sm120a") ], extra_include_paths=extra_include_paths, extra_cuda_cflags=_nvfp4_cuda_flags(), ) def cutlass_scaled_fp4_mm( a: torch.Tensor, b: torch.Tensor, block_scale_a: torch.Tensor, block_scale_b: torch.Tensor, alpha: torch.Tensor, out_dtype: torch.dtype, ) -> torch.Tensor: assert a.ndim == 2 and b.ndim == 2 m, n = a.shape[0], b.shape[0] out = torch.empty((m, n), dtype=out_dtype, device=a.device) module = _jit_nvfp4_scaled_mm_module() module.cutlass_scaled_fp4_mm(out, a, b, block_scale_a, block_scale_b, alpha) return out def cutlass_fp4_group_mm( a_fp4: torch.Tensor, b_fp4: torch.Tensor, a_blockscale: torch.Tensor, b_blockscale: torch.Tensor, alphas: torch.Tensor, out_dtype: torch.dtype, params: dict[str, torch.Tensor], ) -> torch.Tensor: m_topk = a_fp4.shape[0] n = b_fp4.shape[1] output = torch.empty((m_topk, n), device=a_fp4.device, dtype=out_dtype) num_experts = int(params["expert_offsets"].numel()) device = a_fp4.device # Backward compatibility: older callers may not pass scratch tensors. a_ptrs = params.get( "a_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device) ) b_ptrs = params.get( "b_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device) ) out_ptrs = params.get( "out_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device) ) a_scales_ptrs = params.get( "a_scales_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device) ) b_scales_ptrs = params.get( "b_scales_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device) ) alpha_ptrs = params.get( "alpha_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device) ) layout_sfa = params.get( "layout_sfa", torch.empty((num_experts, 5), dtype=torch.int64, device=device) ) layout_sfb = params.get( "layout_sfb", torch.empty((num_experts, 5), dtype=torch.int64, device=device) ) _cutlass_fp4_group_mm_custom_op( output, a_fp4, b_fp4, a_blockscale, b_blockscale, alphas, params["ab_strides"], params["c_strides"], params["problem_sizes"], params["expert_offsets"], params["blockscale_offsets"], a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs, layout_sfa, layout_sfb, ) return output @register_custom_op( op_name="scaled_fp4_quant", mutates_args=["output", "output_scale"], ) def _scaled_fp4_quant_custom_op( input: torch.Tensor, output: torch.Tensor, output_scale: torch.Tensor, input_global_scale: torch.Tensor, ) -> None: module = _jit_nvfp4_quant_module() module.scaled_fp4_quant(output, input, output_scale, input_global_scale) def scaled_fp4_quant( input: torch.Tensor, input_global_scale: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: """Quantize input tensor to FP4 and return packed FP4 tensor + swizzled scales.""" assert input.ndim >= 1, f"input.ndim needs to be >= 1, but got {input.ndim}." other_dims = 1 if input.ndim == 1 else -1 input = input.reshape(other_dims, input.shape[-1]) m, n = input.shape block_size = 16 device = input.device assert n % block_size == 0, f"last dim has to be multiple of 16, but got {n}." assert input.dtype in ( torch.float16, torch.bfloat16, ), f"input.dtype needs to be fp16 or bf16 but got {input.dtype}." output = torch.empty((m, n // 2), device=device, dtype=torch.uint8) rounded_m = ((m + 128 - 1) // 128) * 128 scale_n = n // block_size rounded_n = ((scale_n + 4 - 1) // 4) * 4 if rounded_n > scale_n: output_scale = torch.zeros( (rounded_m, rounded_n // 4), device=device, dtype=torch.int32 ) else: output_scale = torch.empty( (rounded_m, rounded_n // 4), device=device, dtype=torch.int32 ) _scaled_fp4_quant_custom_op(input, output, output_scale, input_global_scale) output_scale = output_scale.view(torch.float8_e4m3fn) return output, output_scale def _shuffle_rows_torch( input_tensor: torch.Tensor, dst2src_map: torch.Tensor, output_tensor_shape: tuple[int, int], ) -> torch.Tensor: # Keep compatibility when sgl-kernel is slimmed and shuffle_rows may not be present. output = input_tensor.index_select(0, dst2src_map.to(dtype=torch.int64)) return output.view(output_tensor_shape) @register_custom_op( op_name="scaled_fp4_experts_quant", mutates_args=["output", "output_scales"], ) def _scaled_fp4_experts_quant_custom_op( output: torch.Tensor, output_scales: torch.Tensor, input_tensor: torch.Tensor, input_global_scale: torch.Tensor, expert_offsets: torch.Tensor, blockscale_offsets: torch.Tensor, ) -> None: module = _jit_nvfp4_expert_quant_module() module.scaled_fp4_experts_quant( output, output_scales, input_tensor, input_global_scale, expert_offsets, blockscale_offsets, ) def scaled_fp4_experts_quant( input_tensor: torch.Tensor, input_global_scale: torch.Tensor, expert_offsets: torch.Tensor, blockscale_offsets: torch.Tensor, topk: int, expert_map: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Quantize packed MoE activations to NVFP4.""" assert ( input_tensor.ndim == 2 ), f"input.ndim needs to be == 2, but got {input_tensor.ndim}." if expert_map is not None: m, k = input_tensor.shape output_tensor_shape = (m * topk, k) input_tensor = _shuffle_rows_torch( input_tensor, expert_map, output_tensor_shape ) m_numtopk, k = input_tensor.shape max_tokens_per_expert = int(os.environ.get("MODELOPT_MAX_TOKENS_PER_EXPERT", 65536)) assert m_numtopk <= max_tokens_per_expert * topk, ( f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT({max_tokens_per_expert})" f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use" " MODELOPT_MAX_TOKENS_PER_EXPERT to set this value." ) scales_k = k // 16 # output_scales is int32-packed FP8 scales, so second dim is in int32 units. padded_k_in_int32 = (scales_k + 3) // 4 output = torch.empty( m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8 ) if padded_k_in_int32 * 4 > scales_k: output_scales = torch.zeros( max_tokens_per_expert * topk, padded_k_in_int32, dtype=torch.int32, device=input_tensor.device, ) else: output_scales = torch.empty( max_tokens_per_expert * topk, padded_k_in_int32, dtype=torch.int32, device=input_tensor.device, ) _scaled_fp4_experts_quant_custom_op( output, output_scales, input_tensor, input_global_scale, expert_offsets, blockscale_offsets, ) output_scales = output_scales.view(torch.float8_e4m3fn) return output, output_scales @register_custom_op( op_name="scaled_fp4_grouped_quant", mutates_args=["output", "output_scales"], ) def _scaled_fp4_grouped_quant_custom_op( input_tensor: torch.Tensor, output: torch.Tensor, output_scales: torch.Tensor, input_global_scale: torch.Tensor, mask: torch.Tensor, ) -> None: l, m, k = input_tensor.shape del l, m module = _jit_nvfp4_expert_quant_module() module.silu_and_mul_scaled_fp4_experts_quant( output.view(-1, k // 2), output_scales.view(-1, output_scales.shape[-1]), input_tensor.view(-1, k), input_global_scale, mask, False, ) def scaled_fp4_grouped_quant( input_tensor: torch.Tensor, input_global_scale: torch.Tensor, mask: torch.Tensor, ): """Quantize grouped GEMM inputs to FP4 and return logical (m, k//2, l).""" device = input_tensor.device l, m, k = input_tensor.shape sf_vec_size = 16 assert k % sf_vec_size == 0, f"k must be multiple of 16, but got {k}." scale_k = k // sf_vec_size padded_k = (scale_k + (4 - 1)) // 4 * 4 padded_k_int32 = padded_k // 4 padded_m = (m + (128 - 1)) // 128 * 128 output = torch.empty(l, m, k // 2, device=device, dtype=torch.uint8) output_scales = torch.empty( l, padded_m, padded_k_int32, device=device, dtype=torch.int32 ) _scaled_fp4_grouped_quant_custom_op( input_tensor, output, output_scales, input_global_scale, mask, ) output = output.permute(1, 2, 0) output_scales = output_scales.view(torch.float8_e4m3fn).view( l, padded_m // 128, padded_k // 4, 32, 4, 4 ) output_scales = output_scales.permute(3, 4, 1, 5, 2, 0) return output, output_scales @register_custom_op( op_name="silu_and_mul_scaled_fp4_grouped_quant", mutates_args=["output", "output_scales"], ) def _silu_and_mul_scaled_fp4_grouped_quant_custom_op( input_tensor: torch.Tensor, output: torch.Tensor, output_scales: torch.Tensor, input_global_scale: torch.Tensor, mask: torch.Tensor, ) -> None: l, m, k_by_2 = input_tensor.shape del l, m module = _jit_nvfp4_expert_quant_module() module.silu_and_mul_scaled_fp4_experts_quant( output.view(-1, output.shape[-1]), output_scales.view(-1, output_scales.shape[-1]), input_tensor.view(-1, k_by_2), input_global_scale, mask, True, ) def silu_and_mul_scaled_fp4_grouped_quant( input_tensor: torch.Tensor, input_global_scale: torch.Tensor, mask: torch.Tensor, ): """Apply SiLU-and-mul then quantize grouped GEMM inputs to FP4.""" device = input_tensor.device l, m, k_by_2 = input_tensor.shape k = k_by_2 // 2 sf_vec_size = 16 assert k % sf_vec_size == 0, f"k must be multiple of 16, but got {k}." scale_k = k // sf_vec_size padded_k = (scale_k + (4 - 1)) // 4 * 4 padded_k_int32 = padded_k // 4 padded_m = (m + (128 - 1)) // 128 * 128 output = torch.empty(l, m, k // 2, device=device, dtype=torch.uint8) output_scales = torch.empty( l, padded_m, padded_k_int32, device=device, dtype=torch.int32 ) _silu_and_mul_scaled_fp4_grouped_quant_custom_op( input_tensor, output, output_scales, input_global_scale, mask, ) output = output.permute(1, 2, 0) output_scales = output_scales.view(torch.float8_e4m3fn).view( l, padded_m // 128, padded_k // 4, 32, 4, 4 ) output_scales = output_scales.permute(3, 4, 1, 5, 2, 0) return output, output_scales @register_custom_op( op_name="cutlass_fp4_group_mm", mutates_args=[ "output", "a_ptrs", "b_ptrs", "out_ptrs", "a_scales_ptrs", "b_scales_ptrs", "alpha_ptrs", "layout_sfa", "layout_sfb", ], ) def _cutlass_fp4_group_mm_custom_op( output: torch.Tensor, a_fp4: torch.Tensor, b_fp4: torch.Tensor, a_blockscale: torch.Tensor, b_blockscale: torch.Tensor, alphas: torch.Tensor, ab_strides: torch.Tensor, c_strides: torch.Tensor, problem_sizes: torch.Tensor, expert_offsets: torch.Tensor, blockscale_offsets: torch.Tensor, a_ptrs: torch.Tensor, b_ptrs: torch.Tensor, out_ptrs: torch.Tensor, a_scales_ptrs: torch.Tensor, b_scales_ptrs: torch.Tensor, alpha_ptrs: torch.Tensor, layout_sfa: torch.Tensor, layout_sfb: torch.Tensor, ) -> None: module = _jit_nvfp4_blockwise_moe_module() module.cutlass_fp4_group_mm( output, a_fp4, b_fp4, a_blockscale, b_blockscale, alphas, ab_strides, c_strides, problem_sizes, expert_offsets, blockscale_offsets, a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs, layout_sfa, layout_sfb, ) def suggest_nvfp4_global_scale(x: torch.Tensor) -> torch.Tensor: """Utility for tests/benchmarks: return global scale used by NVFP4 quantization.""" tensor_amax = torch.abs(x).max().to(torch.float32) return _FLOAT8_E4M3_MAX * _FLOAT4_E2M1_MAX / tensor_amax ================================================ FILE: python/sglang/jit_kernel/per_tensor_quant_fp8.py ================================================ from __future__ import annotations from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_per_tensor_quant_fp8_module(is_static: bool, dtype: torch.dtype) -> Module: args = make_cpp_args(is_static, dtype) return load_jit( "per_tensor_quant_fp8", *args, cuda_files=["gemm/per_tensor_quant_fp8.cuh"], cuda_wrappers=[("per_tensor_quant_fp8", f"per_tensor_quant_fp8<{args}>")], ) @register_custom_op( op_name="per_tensor_quant_fp8", mutates_args=["output_q", "output_s"], ) def per_tensor_quant_fp8( input: torch.Tensor, output_q: torch.Tensor, output_s: torch.Tensor, is_static: bool = False, ) -> None: """ Per-tensor quantization to FP8 format. Args: input: Input tensor to quantize (float, half, or bfloat16) output_q: Output quantized tensor (fp8_e4m3) output_s: Output scale tensor (float scalar or 1D tensor with 1 element) is_static: If True, assumes scale is pre-computed and skips absmax computation """ module = _jit_per_tensor_quant_fp8_module(is_static, input.dtype) module.per_tensor_quant_fp8(input.view(-1), output_q.view(-1), output_s.view(-1)) ================================================ FILE: python/sglang/jit_kernel/per_token_group_quant_8bit.py ================================================ from __future__ import annotations from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: from tvm_ffi.module import Module from sglang.jit_kernel.utils import CPP_DTYPE_MAP as OUTPUT_DTYPE_MAP @cache_once def _jit_per_token_group_quant_8bit_module( dtype: torch.dtype, output_type: torch.dtype ) -> Module: input_args = make_cpp_args(dtype) out_cpp = OUTPUT_DTYPE_MAP[output_type] return load_jit( "per_token_group_quant_8bit", cuda_files=["gemm/per_token_group_quant_8bit.cuh"], cuda_wrappers=[ ( "per_token_group_quant_8bit", f"per_token_group_quant_8bit<{input_args}, {out_cpp}>", ) ], ) @register_custom_op( op_name="per_token_group_quant_8bit", mutates_args=["output_q", "output_s"], ) def _per_token_group_quant_8bit_custom_op( input: torch.Tensor, output_q: torch.Tensor, output_s: torch.Tensor, group_size: int, eps: float, fp8_min: float, fp8_max: float, scale_ue8m0: bool = False, ) -> None: """ Per-token-group quantization to 8-bit format. Args: input: Input tensor to quantize (float, half, or bfloat16). output_q: Output quantized tensor (e.g., fp8_e4m3 or int8). output_s: Output scale tensor. group_size: The size of the group for quantization. eps: A small value to avoid division by zero. fp8_min: The minimum value of the 8-bit data type. fp8_max: The maximum value of the 8-bit data type. scale_ue8m0: Whether to use UE8M0 format for scales. """ module = _jit_per_token_group_quant_8bit_module(input.dtype, output_q.dtype) module.per_token_group_quant_8bit( input, output_q, output_s, group_size, eps, fp8_min, fp8_max, scale_ue8m0, ) return None def per_token_group_quant_8bit( input: torch.Tensor, output_q: torch.Tensor, output_s: torch.Tensor, group_size: int, eps: float, fp8_min: float, fp8_max: float, scale_ue8m0: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: _per_token_group_quant_8bit_custom_op( input=input, output_q=output_q, output_s=output_s, group_size=group_size, eps=eps, fp8_min=fp8_min, fp8_max=fp8_max, scale_ue8m0=scale_ue8m0, ) return output_q, output_s ================================================ FILE: python/sglang/jit_kernel/rope.py ================================================ from __future__ import annotations from dataclasses import dataclass from typing import TYPE_CHECKING, Optional import torch from sglang.jit_kernel.utils import ( cache_once, is_arch_support_pdl, load_jit, make_cpp_args, ) from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_rotary_embedding_module() -> Module: return load_jit( "rotary_embedding", cuda_files=["elementwise/pos_enc.cuh"], cuda_wrappers=[("rotary_embedding", "RotaryEmbeddingKernel::run")], ) @cache_once def _jit_fused_rope_module(is_neox: bool, rope_dim: int, dtype: torch.dtype) -> Module: args = make_cpp_args(is_neox, rope_dim, is_arch_support_pdl(), dtype) return load_jit( "fused_rope", *args, cuda_files=["elementwise/rope.cuh"], cuda_wrappers=[ ("run_rope", f"FusedRopeKernel<{args}>::run"), ("run_rope_store", f"FusedRopeKernel<{args}>::run_fused"), ], ) @register_custom_op( op_name="rotary_embedding_with_key", mutates_args=["query", "key"], ) def rotary_embedding_with_key( positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor, head_size: int, cos_sin_cache: torch.Tensor, is_neox: bool = True, ) -> None: module = _jit_rotary_embedding_module() module.rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox) @register_custom_op( op_name="rotary_embedding_without_key", mutates_args=["query"], ) def rotary_embedding_without_key( positions: torch.Tensor, query: torch.Tensor, head_size: int, cos_sin_cache: torch.Tensor, is_neox: bool = True, ) -> None: module = _jit_rotary_embedding_module() module.rotary_embedding(positions, query, None, head_size, cos_sin_cache, is_neox) def rotary_embedding( positions: torch.Tensor, query: torch.Tensor, key: Optional[torch.Tensor], head_size: int, cos_sin_cache: torch.Tensor, is_neox: bool = True, ): if key is None: rotary_embedding_without_key( positions, query, head_size, cos_sin_cache, is_neox ) else: rotary_embedding_with_key( positions, query, key, head_size, cos_sin_cache, is_neox ) return query, key @dataclass class FusedSetKVBufferArg: """ value : Optional[torch.Tensor] Value tensor, shape: ``(nnz, num_v_heads * head_size)``. k_buffer : Optional[torch.Tensor] Buffer for keys, shape: ``(nnz, num_k_heads * head_size)``. v_buffer : Optional[torch.Tensor] Buffer for values, shape: ``(nnz, num_v_heads * head_size)``. cache_loc : Optional[torch.Tensor] Cache location tensor, used for indexing kv cache. """ value: torch.Tensor k_buffer: torch.Tensor v_buffer: torch.Tensor cache_loc: torch.Tensor @register_custom_op(mutates_args=["q", "k"]) def apply_rope_inplace( q: torch.Tensor, k: torch.Tensor, cos_sin_cache: torch.Tensor, positions: torch.Tensor, *, is_neox: bool, rope_dim: int = 0, ) -> None: """ Fused inplace rotary position embedding for query and key tensors. Args: q: Query tensor of shape [num_tokens, num_qo_heads, rope_dim]. k: Key tensor of shape [num_tokens, num_kv_heads, rope_dim]. cos_sin_cache: Cosine/sine cache of shape [max_position, rope_dim], where the first half along dim=-1 is cos and the second half is sin. Must be float32. positions: Position indices of shape [num_tokens], int32 or int64. is_neox: Whether to use GPT-NeoX style (True) or GPT-J interleaved style (False). rope_dim: Rotary embedding dimension. Defaults to cos_sin_cache.size(-1). """ rope_dim = rope_dim or cos_sin_cache.size(-1) module = _jit_fused_rope_module(is_neox, rope_dim, q.dtype) module.run_rope(q, k, cos_sin_cache, positions) @register_custom_op(mutates_args=["q", "k", "k_cache", "v_cache"]) def apply_rope_inplace_with_kvcache( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, cos_sin_cache: torch.Tensor, positions: torch.Tensor, out_loc: torch.Tensor, *, is_neox: bool, rope_dim: int = 0, ) -> None: """ Fused inplace RoPE + KV cache store. Applies rotary position embedding to q and k inplace. The rotated k is also stored in k_cache. The original v is also stored in v_cache. Args: q: Query tensor of shape [num_tokens, num_qo_heads, head_dim]. k: Key tensor of shape [num_tokens, num_kv_heads, head_dim]. v: Value tensor of shape [num_tokens, num_kv_heads, head_dim]. k_cache: Key cache of shape [cache_size, num_kv_heads * head_dim]. v_cache: Value cache of shape [cache_size, num_kv_heads * head_dim]. cos_sin_cache: Cosine/sine cache of shape [max_position, rope_dim], float32. positions: Position indices of shape [num_tokens], int32 or int64. out_loc: Cache write locations of shape [num_tokens], same dtype as positions. is_neox: Whether to use GPT-NeoX style (True) or GPT-J interleaved (False). rope_dim: Rotary embedding dimension. Defaults to cos_sin_cache.size(-1). """ rope_dim = rope_dim or cos_sin_cache.size(-1) v = v.view_as(k) module = _jit_fused_rope_module(is_neox, rope_dim, q.dtype) module.run_rope_store(q, k, v, k_cache, v_cache, cos_sin_cache, positions, out_loc) # NOTE: this name is intentionally set as the old kernel in `sgl_kernel` def apply_rope_with_cos_sin_cache_inplace( q: torch.Tensor, k: torch.Tensor, cos_sin_cache: torch.Tensor, positions: torch.Tensor, *, is_neox: bool, rope_dim: int = 0, fused_args: Optional[FusedSetKVBufferArg] = None, ) -> None: """ Apply RoPE to q and k inplace, with optional fused kv cache store. If `fused_args` is provided, it will perform fused RoPE and KV cache store. Otherwise, it will only apply RoPE inplace. Args: q: Query tensor of shape [num_tokens, num_qo_heads, head_dim]. k: Key tensor of shape [num_tokens, num_kv_heads, head_dim]. cos_sin_cache: Cosine/sine cache of shape [max_position, rope_dim], float32. positions: Position indices of shape [num_tokens], int32 or int64. is_neox: Whether to use GPT-NeoX style (True) or GPT-J interleaved (False). rope_dim: Rotary embedding dimension. Defaults to cos_sin_cache.size(-1). fused_args: Optional arguments for fused RoPE + KV cache store. If None, only RoPE will be applied inplace without touching kv cache. """ if fused_args is not None: apply_rope_inplace_with_kvcache( q, k, fused_args.value, fused_args.k_buffer, fused_args.v_buffer, cos_sin_cache, positions, fused_args.cache_loc, is_neox=is_neox, rope_dim=rope_dim, ) else: apply_rope_inplace( q, k, cos_sin_cache, positions, is_neox=is_neox, rope_dim=rope_dim ) ================================================ FILE: python/sglang/jit_kernel/tests/test_add_constant.py ================================================ import pytest import torch from sglang.jit_kernel.add_constant import add_constant @pytest.mark.parametrize("size", [1, 2, 127, 128, 1024, 1025]) @pytest.mark.parametrize("constant", [0, 1, 7, 1024, -3]) def test_add_constant(size: int, constant: int) -> None: src = torch.arange(0, size, dtype=torch.int32, device="cuda") dst = add_constant(src, constant) assert torch.all(dst == src + constant) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_awq_dequantize.py ================================================ import itertools import pytest import torch from sglang.jit_kernel.awq_dequantize import awq_dequantize as jit_awq_dequantize try: from sgl_kernel import awq_dequantize as aot_awq_dequantize AOT_AVAILABLE = True except ImportError: AOT_AVAILABLE = False def reverse_awq_order(t: torch.Tensor): bits = 4 AWQ_REVERSE_ORDER = [0, 4, 1, 5, 2, 6, 3, 7] reverse_order_tensor = torch.arange( t.shape[-1], dtype=torch.int32, device=t.device, ) reverse_order_tensor = reverse_order_tensor.view(-1, 32 // bits) reverse_order_tensor = reverse_order_tensor[:, AWQ_REVERSE_ORDER] reverse_order_tensor = reverse_order_tensor.view(-1) t = t[:, reverse_order_tensor] & 0xF return t # qweights - [R , C // 8], int32 # scales - [R // G, C ], float16 # zeros - [R // G, C // 8], int32 def awq_dequantize_torch( qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor, group_size: int ) -> torch.Tensor: if group_size == -1: group_size = qweight.shape[0] bits = 4 shifts = torch.arange(0, 32, bits, device=qzeros.device) iweights = torch.bitwise_right_shift(qweight[:, :, None], shifts[None, None, :]).to( torch.int8 ) iweights = iweights.view(iweights.shape[0], -1) zeros = torch.bitwise_right_shift(qzeros[:, :, None], shifts[None, None, :]).to( torch.int8 ) zeros = zeros.view(qzeros.shape[0], -1) zeros = reverse_awq_order(zeros) iweights = reverse_awq_order(iweights) iweights = torch.bitwise_and(iweights, (2**bits) - 1) zeros = torch.bitwise_and(zeros, (2**bits) - 1) scales = scales.repeat_interleave(group_size, dim=0) zeros = zeros.repeat_interleave(group_size, dim=0) return (iweights - zeros) * scales @pytest.mark.parametrize( "qweight_row,qweight_col,is_bf16_act", list( itertools.product( [128, 256, 512, 1024, 3584], [16, 32, 64, 128, 448], [True, False], ) ), ) def test_awq_dequantize_jit_vs_torch( qweight_row: int, qweight_col: int, is_bf16_act: bool ): device = torch.device("cuda") qweight = torch.randint( 0, torch.iinfo(torch.int32).max, (qweight_row, qweight_col), dtype=torch.int32, device=device, ) group_size = qweight_row scales_row = qweight_row // group_size scales_col = qweight_col * 8 if is_bf16_act: scales = torch.rand(scales_row, scales_col, dtype=torch.bfloat16, device=device) else: scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device) qzeros = torch.randint( 0, torch.iinfo(torch.int32).max, (scales_row, qweight_col), dtype=torch.int32, device=device, ) # Run both implementations torch_out = awq_dequantize_torch(qweight, scales, qzeros, group_size) jit_out = jit_awq_dequantize(qweight, scales, qzeros) # Compare results (approximate due to different computation paths) torch.testing.assert_close( torch_out.to(torch.float32), jit_out.to(torch.float32), rtol=1e-3, atol=1e-5 ) @pytest.mark.parametrize( "qweight_row,qweight_col,is_bf16_act", list( itertools.product( [128, 256, 512, 1024, 3584], [16, 32, 64, 128, 448], [True, False], ) ), ) def test_awq_dequantize_jit_vs_aot( qweight_row: int, qweight_col: int, is_bf16_act: bool ): if not AOT_AVAILABLE: pytest.skip("sgl_kernel AOT not available") device = torch.device("cuda") qweight = torch.randint( 0, torch.iinfo(torch.int32).max, (qweight_row, qweight_col), dtype=torch.int32, device=device, ) group_size = qweight_row scales_row = qweight_row // group_size scales_col = qweight_col * 8 if is_bf16_act: scales = torch.rand(scales_row, scales_col, dtype=torch.bfloat16, device=device) else: scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device) qzeros = torch.randint( 0, torch.iinfo(torch.int32).max, (scales_row, qweight_col), dtype=torch.int32, device=device, ) # Run both implementations aot_out = aot_awq_dequantize(qweight, scales, qzeros) jit_out = jit_awq_dequantize(qweight, scales, qzeros) # Bitwise equality torch.testing.assert_close(jit_out, aot_out, rtol=0, atol=0) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_awq_marlin_moe_repack.py ================================================ import numpy as np import pytest import torch from sgl_kernel.scalar_type import scalar_types from sglang.jit_kernel.awq_marlin_repack import ( awq_marlin_moe_repack as jit_awq_marlin_moe_repack, ) from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights def _has_aot_awq_marlin_moe_repack() -> bool: return hasattr(torch.ops.sgl_kernel, "awq_marlin_moe_repack") and hasattr( torch.ops.sgl_kernel.awq_marlin_moe_repack, "default" ) AOT_AVAILABLE = _has_aot_awq_marlin_moe_repack() def awq_pack( q_w: torch.Tensor, num_bits: int, size_k: int, size_n: int, ): assert q_w.shape == (size_k, size_n) if num_bits == 4: interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7]) elif num_bits == 8: interleave = np.array([0, 2, 1, 3]) else: raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel() q_w = q_w.reshape((-1, size_n)).contiguous() return pack_cols(q_w, num_bits, size_k, size_n) @pytest.mark.parametrize("num_bits", [4]) @pytest.mark.parametrize("num_experts", [2, 4, 8]) @pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2), (4, 4)]) @pytest.mark.parametrize("group_size", [16, 32]) def test_awq_marlin_moe_repack_jit_vs_aot( num_bits, num_experts, k_tiles, n_tiles, group_size ): if not AOT_AVAILABLE: pytest.skip("sgl_kernel AOT not available") tile_k, tile_n = 16, 64 size_k = k_tiles * tile_k size_n = n_tiles * tile_n pack_factor = 32 // num_bits # Create per-expert AWQ-packed weights b_q_weight = torch.empty( (num_experts, size_k, size_n // pack_factor), dtype=torch.int32, device="cuda", ) for e in range(num_experts): b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda") w_ref, q_w, s, zp = quantize_weights( b_weight, scalar_types.uint4, group_size, zero_points=True ) b_q_weight[e] = awq_pack(q_w, num_bits, size_k, size_n) perm = torch.empty((num_experts, 0), dtype=torch.int32, device="cuda") out_jit = jit_awq_marlin_moe_repack(b_q_weight, perm, size_k, size_n, num_bits) out_aot = torch.ops.sgl_kernel.awq_marlin_moe_repack.default( b_q_weight, perm, size_k, size_n, num_bits ) torch.cuda.synchronize() # Bitwise equality torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0) @pytest.mark.parametrize("num_bits", [4]) @pytest.mark.parametrize("num_experts", [2, 4]) @pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2)]) @pytest.mark.parametrize("group_size", [16, 32]) def test_awq_marlin_moe_repack_shape( num_bits, num_experts, k_tiles, n_tiles, group_size ): tile_k, tile_n = 16, 64 size_k = k_tiles * tile_k size_n = n_tiles * tile_n pack_factor = 32 // num_bits # Create per-expert AWQ-packed weights b_q_weight = torch.empty( (num_experts, size_k, size_n // pack_factor), dtype=torch.int32, device="cuda", ) for e in range(num_experts): b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda") w_ref, q_w, s, zp = quantize_weights( b_weight, scalar_types.uint4, group_size, zero_points=True ) b_q_weight[e] = awq_pack(q_w, num_bits, size_k, size_n) perm = torch.empty((num_experts, 0), dtype=torch.int32, device="cuda") out = jit_awq_marlin_moe_repack(b_q_weight, perm, size_k, size_n, num_bits) torch.cuda.synchronize() assert out.is_cuda and out.dtype == torch.int32 expected_shape = (num_experts, size_k // 16, size_n * (num_bits // 2)) assert list(out.shape) == list(expected_shape) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_awq_marlin_repack.py ================================================ import numpy as np import pytest import torch from sgl_kernel.scalar_type import scalar_types from sglang.jit_kernel.awq_marlin_repack import ( awq_marlin_repack as jit_awq_marlin_repack, ) from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights def _has_aot_awq_marlin_repack() -> bool: return hasattr(torch.ops.sgl_kernel, "awq_marlin_repack") and hasattr( torch.ops.sgl_kernel.awq_marlin_repack, "default" ) AOT_AVAILABLE = _has_aot_awq_marlin_repack() def awq_pack( q_w: torch.Tensor, num_bits: int, size_k: int, size_n: int, ): assert q_w.shape == (size_k, size_n) if num_bits == 4: interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7]) elif num_bits == 8: interleave = np.array([0, 2, 1, 3]) else: raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel() q_w = q_w.reshape((-1, size_n)).contiguous() return pack_cols(q_w, num_bits, size_k, size_n) @pytest.mark.parametrize("num_bits", [4, 8]) @pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2), (4, 4)]) @pytest.mark.parametrize("group_size", [16, 32]) def test_awq_marlin_repack_jit_vs_aot(num_bits, k_tiles, n_tiles, group_size): if not AOT_AVAILABLE: pytest.skip("sgl_kernel AOT not available") tile_k, tile_n = 16, 64 size_k = k_tiles * tile_k size_n = n_tiles * tile_n b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda") w_ref, q_w, s, zp = quantize_weights( b_weight, scalar_types.uint4, group_size, zero_points=True ) q_w_awq = awq_pack(q_w, num_bits, size_k, size_n) out_jit = jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits) out_aot = torch.ops.sgl_kernel.awq_marlin_repack.default( q_w_awq, size_k, size_n, num_bits ) torch.cuda.synchronize() # Bitwise equality torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0) @pytest.mark.parametrize("num_bits", [4, 8]) @pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2)]) @pytest.mark.parametrize("group_size", [16, 32]) def test_awq_marlin_repack_correct(num_bits, k_tiles, n_tiles, group_size): tile_k, tile_n = 16, 64 size_k = k_tiles * tile_k size_n = n_tiles * tile_n pack_factor = 32 // num_bits b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda") w_ref, q_w, s, zp = quantize_weights( b_weight, scalar_types.uint4, group_size, zero_points=True ) q_w_awq = awq_pack(q_w, num_bits, size_k, size_n) weight_perm = get_weight_perm(num_bits) q_w_marlin = marlin_weights(q_w, size_k, size_n, num_bits, weight_perm) out_gpu = jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits) assert out_gpu.is_cuda and out_gpu.dtype == torch.int32 expected_cols = size_n * tile_k // pack_factor assert list(out_gpu.shape) == [size_k // tile_k, expected_cols] torch.cuda.synchronize() torch.testing.assert_close(out_gpu, q_w_marlin) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_concat_mla.py ================================================ import itertools import pytest import torch import triton def torch_concat_mla_k( k: torch.Tensor, k_nope: torch.Tensor, k_rope: torch.Tensor ) -> None: """Reference PyTorch implementation for concat_mla_k.""" # k_nope: [num_tokens, num_heads, nope_head_dim] # k_rope: [num_tokens, 1, rope_head_dim] # k: [num_tokens, num_heads, nope_head_dim + rope_head_dim] nope_head_dim = k_nope.shape[-1] k[:, :, :nope_head_dim] = k_nope # Broadcast k_rope across all heads k[:, :, nope_head_dim:] = k_rope.expand(-1, k.shape[1], -1) def torch_concat_mla_absorb_q( a: torch.Tensor, b: torch.Tensor, out: torch.Tensor ) -> None: """Reference PyTorch implementation for concat_mla_absorb_q.""" # a: [dim_0, dim_1, a_last_dim] # b: [dim_0, dim_1, b_last_dim] # out: [dim_0, dim_1, a_last_dim + b_last_dim] a_last_dim = a.shape[-1] out[:, :, :a_last_dim] = a out[:, :, a_last_dim:] = b def sgl_kernel_concat_mla_k( k: torch.Tensor, k_nope: torch.Tensor, k_rope: torch.Tensor ) -> None: """AOT compiled sgl_kernel implementation.""" from sgl_kernel import concat_mla_k concat_mla_k(k, k_nope, k_rope) def sgl_kernel_concat_mla_absorb_q( a: torch.Tensor, b: torch.Tensor, out: torch.Tensor ) -> None: """AOT compiled sgl_kernel implementation.""" from sgl_kernel import concat_mla_absorb_q result = concat_mla_absorb_q(a, b) # AOT returns output out.copy_(result) # Copy to provided tensor for comparison def jit_concat_mla_k( k: torch.Tensor, k_nope: torch.Tensor, k_rope: torch.Tensor ) -> None: """JIT compiled implementation.""" from sglang.jit_kernel.concat_mla import concat_mla_k concat_mla_k(k, k_nope, k_rope) def jit_concat_mla_absorb_q( a: torch.Tensor, b: torch.Tensor, out: torch.Tensor ) -> None: """JIT compiled implementation - wrapper for test compatibility.""" from sglang.jit_kernel.concat_mla import concat_mla_absorb_q result = concat_mla_absorb_q(a, b) out.copy_(result) # Constants matching the kernel NUM_LOCAL_HEADS = 128 QK_NOPE_HEAD_DIM = 128 QK_ROPE_HEAD_DIM = 64 K_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM A_LAST_DIM = 512 B_LAST_DIM = 64 OUT_LAST_DIM = A_LAST_DIM + B_LAST_DIM DEVICE = "cuda" DTYPE = torch.bfloat16 # Test configurations NUM_TOKENS_LIST = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024] @pytest.mark.parametrize("num_tokens", NUM_TOKENS_LIST) def test_concat_mla_k_jit_vs_torch(num_tokens: int) -> None: """Test JIT kernel against PyTorch reference.""" k_jit = torch.empty( num_tokens, NUM_LOCAL_HEADS, K_HEAD_DIM, device=DEVICE, dtype=DTYPE ) k_torch = torch.empty( num_tokens, NUM_LOCAL_HEADS, K_HEAD_DIM, device=DEVICE, dtype=DTYPE ) k_nope = torch.randn( num_tokens, NUM_LOCAL_HEADS, QK_NOPE_HEAD_DIM, device=DEVICE, dtype=DTYPE ) k_rope = torch.randn(num_tokens, 1, QK_ROPE_HEAD_DIM, device=DEVICE, dtype=DTYPE) torch_concat_mla_k(k_torch, k_nope, k_rope) jit_concat_mla_k(k_jit, k_nope, k_rope) triton.testing.assert_close(k_jit, k_torch, atol=0, rtol=0) @pytest.mark.parametrize("num_tokens", NUM_TOKENS_LIST) def test_concat_mla_k_jit_vs_aot(num_tokens: int) -> None: """Test JIT kernel against AOT kernel for bitwise equivalence.""" k_jit = torch.empty( num_tokens, NUM_LOCAL_HEADS, K_HEAD_DIM, device=DEVICE, dtype=DTYPE ) k_aot = torch.empty( num_tokens, NUM_LOCAL_HEADS, K_HEAD_DIM, device=DEVICE, dtype=DTYPE ) k_nope = torch.randn( num_tokens, NUM_LOCAL_HEADS, QK_NOPE_HEAD_DIM, device=DEVICE, dtype=DTYPE ) k_rope = torch.randn(num_tokens, 1, QK_ROPE_HEAD_DIM, device=DEVICE, dtype=DTYPE) sgl_kernel_concat_mla_k(k_aot, k_nope, k_rope) jit_concat_mla_k(k_jit, k_nope, k_rope) triton.testing.assert_close(k_jit, k_aot, atol=0, rtol=0) DIM_0_LIST = [1, 2, 4, 8, 16, 32] DIM_1_LIST = [1, 2, 4, 8, 16, 128] @pytest.mark.parametrize( "dim_0,dim_1", list(itertools.product(DIM_0_LIST, DIM_1_LIST)), ) def test_concat_mla_absorb_q_jit_vs_torch(dim_0: int, dim_1: int) -> None: """Test JIT kernel against PyTorch reference.""" a = torch.randn(dim_0, dim_1, A_LAST_DIM, device=DEVICE, dtype=DTYPE) b = torch.randn(dim_0, dim_1, B_LAST_DIM, device=DEVICE, dtype=DTYPE) out_jit = torch.empty(dim_0, dim_1, OUT_LAST_DIM, device=DEVICE, dtype=DTYPE) out_torch = torch.empty(dim_0, dim_1, OUT_LAST_DIM, device=DEVICE, dtype=DTYPE) torch_concat_mla_absorb_q(a, b, out_torch) jit_concat_mla_absorb_q(a, b, out_jit) triton.testing.assert_close(out_jit, out_torch, atol=0, rtol=0) @pytest.mark.parametrize( "dim_0,dim_1", list(itertools.product(DIM_0_LIST, DIM_1_LIST)), ) def test_concat_mla_absorb_q_jit_vs_aot(dim_0: int, dim_1: int) -> None: """Test JIT kernel against AOT kernel for bitwise equivalence.""" a = torch.randn(dim_0, dim_1, A_LAST_DIM, device=DEVICE, dtype=DTYPE) b = torch.randn(dim_0, dim_1, B_LAST_DIM, device=DEVICE, dtype=DTYPE) out_jit = torch.empty(dim_0, dim_1, OUT_LAST_DIM, device=DEVICE, dtype=DTYPE) out_aot = torch.empty(dim_0, dim_1, OUT_LAST_DIM, device=DEVICE, dtype=DTYPE) sgl_kernel_concat_mla_absorb_q(a, b, out_aot) jit_concat_mla_absorb_q(a, b, out_jit) triton.testing.assert_close(out_jit, out_aot, atol=0, rtol=0) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_cutedsl_gdn.py ================================================ """Tests for CuTe DSL fused sigmoid gating delta rule kernel (GDN).""" import numpy as np import pytest import torch try: import cuda.bindings.driver as cuda_driver import cutlass # noqa: F401 from cutlass.cute.runtime import from_dlpack from sglang.jit_kernel import cutedsl_gdn CUTEDSL_AVAILABLE = True except ImportError: CUTEDSL_AVAILABLE = False cutedsl_gdn = None try: from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) TRITON_AVAILABLE = True except ImportError: TRITON_AVAILABLE = False def run_triton_kernel(A_log, dt_bias, q, k, v, a, b, initial_state, indices, scale): return fused_sigmoid_gating_delta_rule_update( A_log=A_log, a=a, dt_bias=dt_bias, softplus_beta=1.0, softplus_threshold=20.0, q=q, k=k, v=v, b=b, initial_state_source=initial_state, initial_state_indices=indices, scale=scale, use_qk_l2norm_in_kernel=True, cu_seqlens=None, ) @pytest.mark.skipif(not CUTEDSL_AVAILABLE, reason="CuTe DSL not available") @pytest.mark.skipif(not TRITON_AVAILABLE, reason="Triton kernel not available") @pytest.mark.skip( reason=( "Temporary CI workaround: CuTe DSL GDN precision is currently unstable " "against the Triton reference and needs follow-up investigation." ) ) @pytest.mark.parametrize("B", [16, 128]) def test_cutedsl_gdn_precision(B: int): """Test precision of CuTe DSL GDN kernel against Triton reference.""" torch.manual_seed(2025) T, H, K, V, HV = 1, 16, 128, 128, 32 scale = K**-0.5 A_log = torch.randn(HV, dtype=torch.float32, device="cuda") dt_bias = torch.randn(HV, dtype=torch.bfloat16, device="cuda") a = torch.randn(B, T, HV, dtype=torch.bfloat16, device="cuda") b = torch.randn(B, T, HV, dtype=torch.bfloat16, device="cuda") q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device="cuda") k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device="cuda") v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device="cuda") indices = torch.arange(B, dtype=torch.int32, device="cuda") state_cutedsl = torch.randn(B, HV, K, V, dtype=torch.float32, device="cuda") state_triton = state_cutedsl.clone().reshape(-1).contiguous() # Warmup compilation _ = cutedsl_gdn.cutedsl_fused_sigmoid_gating_delta_rule_update( A_log, dt_bias, q, k, v, a, b, state_cutedsl.clone(), indices, scale=scale ) torch.cuda.synchronize() # Fresh state for actual test state_cutedsl = torch.randn(B, HV, K, V, dtype=torch.float32, device="cuda") state_triton = state_cutedsl.clone().reshape(-1).contiguous() out_cutedsl = cutedsl_gdn.cutedsl_fused_sigmoid_gating_delta_rule_update( A_log, dt_bias, q, k, v, a, b, state_cutedsl, indices, scale=scale ) out_triton = run_triton_kernel( A_log, dt_bias, q, k, v, a, b, state_triton, indices, scale ) # Check precision: diff > 0.1 must be < 1% of elements abs_diff = (out_triton.float() - out_cutedsl.float()).abs() max_diff = abs_diff.max().item() mean_diff = abs_diff.mean().item() fail_rate = (abs_diff > 0.1).float().mean().item() * 100 has_nan = torch.isnan(out_cutedsl).any() or torch.isinf(out_cutedsl).any() kernel_type = "SmallBatch" if B < 32 else "LargeBatch" print( f"\n B={B} ({kernel_type}): max_diff={max_diff:.2e}, mean_diff={mean_diff:.2e}, fail_rate={fail_rate:.2f}%" ) assert not has_nan, "Output contains NaN/Inf" assert fail_rate < 1.0, f"Fail rate {fail_rate:.2f}% >= 1%" @pytest.mark.skipif( True, reason="Skip the performance test because the speedup ratio is highly unstable in the CI environment. ", ) @pytest.mark.skipif(not CUTEDSL_AVAILABLE, reason="CuTe DSL not available") @pytest.mark.skipif(not TRITON_AVAILABLE, reason="Triton kernel not available") @pytest.mark.parametrize("B", [1, 128]) def test_cutedsl_gdn_performance(B: int): """Benchmark CuTe DSL GDN kernel against Triton reference.""" torch.manual_seed(2025) T, H, K, V, HV = 1, 16, 128, 128, 32 N = B scale = K**-0.5 is_varlen = True warmup, bench_iters, run_iters = 10, 100, 10 A_log = torch.randn(HV, dtype=torch.float32, device="cuda") dt_bias = torch.randn(HV, dtype=torch.bfloat16, device="cuda") indices = torch.arange(N, dtype=torch.int32, device="cuda") state_cutedsl = torch.randn(N, HV, K, V, dtype=torch.float32, device="cuda") state_triton = state_cutedsl.reshape(-1).contiguous() cu_seqlens = torch.zeros(N + 1, dtype=torch.int32, device="cuda") o_cutedsl = torch.zeros(1, N, HV, V, dtype=torch.bfloat16, device="cuda") # Prepare tensors for multiple runs q_list, k_list, v_list, a_list, b_list = [], [], [], [], [] q_tensor_list, k_tensor_list, v_tensor_list, a_tensor_list, b_tensor_list = ( [], [], [], [], [], ) q_triton, k_triton, v_triton, a_triton, b_triton = [], [], [], [], [] for ri in range(run_iters): torch.manual_seed(2025 + ri) q_i = torch.randn(1, N, H, K, dtype=torch.bfloat16, device="cuda") k_i = torch.randn(1, N, H, K, dtype=torch.bfloat16, device="cuda") v_i = torch.randn(1, N, HV, V, dtype=torch.bfloat16, device="cuda") a_i = torch.randn(N, HV, dtype=torch.bfloat16, device="cuda") b_i = torch.randn(N, HV, dtype=torch.bfloat16, device="cuda") q_list.append(q_i) k_list.append(k_i) v_list.append(v_i) a_list.append(a_i) b_list.append(b_i) q_tensor_list.append(from_dlpack(q_i, assumed_align=16)) k_tensor_list.append(from_dlpack(k_i, assumed_align=16)) v_tensor_list.append(from_dlpack(v_i, assumed_align=16)) a_tensor_list.append(from_dlpack(a_i, assumed_align=16)) b_tensor_list.append(from_dlpack(b_i, assumed_align=16)) q_triton.append(q_i.transpose(0, 1).contiguous()) k_triton.append(k_i.transpose(0, 1).contiguous()) v_triton.append(v_i.transpose(0, 1).contiguous()) a_triton.append(a_i.unsqueeze(1).contiguous()) b_triton.append(b_i.unsqueeze(1).contiguous()) A_log_t = from_dlpack(A_log, assumed_align=16) dt_bias_t = from_dlpack(dt_bias, assumed_align=16) h0_t = from_dlpack(state_cutedsl, assumed_align=16) idx_t = from_dlpack(indices, assumed_align=16) o_t = from_dlpack(o_cutedsl, assumed_align=16) cu_t = from_dlpack(cu_seqlens, assumed_align=16) torch_stream = torch.cuda.Stream() stream = cuda_driver.CUstream(torch_stream.cuda_stream) # Compile kernels compiled = cutedsl_gdn._get_compiled_kernel(N, H, HV, K, V, N, N < 32, is_varlen) torch.cuda.synchronize() for ri in range(run_iters): _ = run_triton_kernel( A_log, dt_bias, q_triton[ri], k_triton[ri], v_triton[ri], a_triton[ri], b_triton[ri], state_triton, indices, scale, ) torch.cuda.synchronize() def run_cutedsl(): for ri in range(run_iters): compiled( cu_t, q_tensor_list[ri], k_tensor_list[ri], v_tensor_list[ri], a_tensor_list[ri], b_tensor_list[ri], A_log_t, dt_bias_t, h0_t, idx_t, o_t, stream, ) def run_triton(): for ri in range(run_iters): _ = run_triton_kernel( A_log, dt_bias, q_triton[ri], k_triton[ri], v_triton[ri], a_triton[ri], b_triton[ri], state_triton, indices, scale, ) # Warmup with torch.cuda.stream(torch_stream): run_cutedsl() torch.cuda.synchronize() run_triton() torch.cuda.synchronize() # Capture CUDA graphs graph_triton = torch.cuda.CUDAGraph() graph_cutedsl = torch.cuda.CUDAGraph() try: with torch.cuda.graph(graph_triton): run_triton() with torch.cuda.graph(graph_cutedsl, stream=torch_stream): run_cutedsl() torch.cuda.synchronize() except Exception: graph_triton = graph_cutedsl = None # Warmup with graphs for _ in range(warmup): if graph_cutedsl: graph_cutedsl.replay() else: with torch.cuda.stream(torch_stream): run_cutedsl() torch.cuda.synchronize() if graph_triton: graph_triton.replay() else: run_triton() torch.cuda.synchronize() # Benchmark triton_times, cutedsl_times = [], [] for _ in range(bench_iters): start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( enable_timing=True ) start.record() if graph_triton: graph_triton.replay() else: run_triton() end.record() torch.cuda.synchronize() triton_times.append(start.elapsed_time(end)) start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( enable_timing=True ) with torch.cuda.stream(torch_stream): start.record() if graph_cutedsl: graph_cutedsl.replay() else: run_cutedsl() end.record() torch.cuda.synchronize() cutedsl_times.append(start.elapsed_time(end)) triton_mean = np.mean(triton_times) / run_iters * 1000 triton_std = np.std(triton_times) / run_iters * 1000 cutedsl_mean = np.mean(cutedsl_times) / run_iters * 1000 cutedsl_std = np.std(cutedsl_times) / run_iters * 1000 speedup = triton_mean / cutedsl_mean kernel_type = "SmallBatch" if B < 32 else "LargeBatch" print( f"\n B={B} ({kernel_type}): Triton={triton_mean:.2f}±{triton_std:.2f}μs, CuTeDSL={cutedsl_mean:.2f}±{cutedsl_std:.2f}μs, speedup={speedup:.2f}x" ) min_speedup = 1.0 if B < 32 else 1.15 assert speedup >= min_speedup, f"Speedup {speedup:.2f}x < {min_speedup}x for B={B}" if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_flash_attention_4.py ================================================ # Adapted from https://github.com/Dao-AILab/flash-attention/blob/8ecf128f683266735ba68e3c106ff67a2611886e/tests/cute/test_flash_attn.py # Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. import itertools import math import pytest import torch import torch.nn.functional as F from einops import rearrange, repeat from sglang.jit_kernel.flash_attention_v4 import flash_attn_varlen_func # Skip this test on Hopper machine skip_condition = torch.cuda.get_device_capability() < (10, 0) def apply_rotary_emb( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, seqlen_offsets: torch.Tensor | int | None = 0, interleaved: bool = False, ) -> torch.Tensor: rotary_dim = cos.shape[-1] * 2 x_rot = x[..., :rotary_dim] x_pass = x[..., rotary_dim:] cos = cos.to(dtype=x.dtype) sin = sin.to(dtype=x.dtype) if x_rot.dim() < 2: raise ValueError(f"apply_rotary_emb expects x.dim() >= 2, got {x_rot.dim()}") b = x_rot.shape[0] s = x_rot.shape[1] if seqlen_offsets is None: seqlen_offsets = 0 if isinstance(seqlen_offsets, int): positions = ( torch.arange(s, device=x_rot.device, dtype=torch.long) + seqlen_offsets ) cos_s = cos.index_select(0, positions) sin_s = sin.index_select(0, positions) cos_s = cos_s.unsqueeze(0).expand(b, -1, -1) sin_s = sin_s.unsqueeze(0).expand(b, -1, -1) else: if seqlen_offsets.dim() != 1 or seqlen_offsets.shape[0] != b: raise ValueError( "apply_rotary_emb expects seqlen_offsets to be int or shape [batch]" ) positions = torch.arange(s, device=x_rot.device, dtype=torch.long).unsqueeze( 0 ) + seqlen_offsets.to(dtype=torch.long).unsqueeze(1) cos_s = cos.index_select(0, positions.reshape(-1)).view(b, s, -1) sin_s = sin.index_select(0, positions.reshape(-1)).view(b, s, -1) x_rot = x_rot.reshape(b, s, -1, rotary_dim) cos_s = cos_s.unsqueeze(2) sin_s = sin_s.unsqueeze(2) if interleaved: x1 = x_rot[..., ::2] x2 = x_rot[..., 1::2] o1 = x1 * cos_s - x2 * sin_s o2 = x2 * cos_s + x1 * sin_s x_rot = torch.stack((o1, o2), dim=-1).flatten(-2) else: x1, x2 = torch.chunk(x_rot, 2, dim=-1) o1 = x1 * cos_s - x2 * sin_s o2 = x2 * cos_s + x1 * sin_s x_rot = torch.cat((o1, o2), dim=-1) x_rot = x_rot.reshape_as(x[..., :rotary_dim]) return torch.cat((x_rot, x_pass), dim=-1) def unpad_input(hidden_states, attention_mask, unused_mask=None): """ Arguments: hidden_states: (batch, seqlen, ...) attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. Return: hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. indices: (total_nnz), the indices of masked tokens from the flattened input sequence. cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. max_seqlen_in_batch: int seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. """ all_masks = ( (attention_mask + unused_mask) if unused_mask is not None else attention_mask ) seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() max_seqlen_in_batch = seqlens_in_batch.max().item() cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to # index with integer indices. return ( rearrange(hidden_states, "b s ... -> (b s) ...")[indices], indices, cu_seqlens, max_seqlen_in_batch, used_seqlens_in_batch, ) def pad_input(hidden_states, indices, batch, seqlen): """ Arguments: hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. batch: int, batch size for the padded sequence. seqlen: int, maximum sequence length for the padded sequence. Return: hidden_states: (batch, seqlen, ...) """ dim = hidden_states.shape[1:] output = torch.zeros( (batch * seqlen), *dim, device=hidden_states.device, dtype=hidden_states.dtype ) output[indices] = hidden_states return rearrange(output, "(b s) ... -> b s ...", b=batch) def generate_random_padding_mask( max_seqlen, batch_size, device, mode="random", zero_lengths=False ): assert mode in ["full", "random", "third"] if mode == "full": lengths = torch.full( (batch_size, 1), max_seqlen, device=device, dtype=torch.int32 ) elif mode == "random": lengths = torch.randint( max(0 if zero_lengths else 1, max_seqlen - 20), max_seqlen + 1, (batch_size, 1), device=device, ) elif mode == "third": lengths = torch.randint( max_seqlen // 3, max_seqlen + 1, (batch_size, 1), device=device ) else: # This should never happen due to the assertion above, but for linter lengths = torch.full( (batch_size, 1), max_seqlen, device=device, dtype=torch.int32 ) if zero_lengths: # Generate zero-lengths every 5 batches and the last batch. for i in range(batch_size): if i % 5 == 0: lengths[i] = 0 lengths[-1] = 0 padding_mask = ( repeat(torch.arange(max_seqlen, device=device), "s -> b s", b=batch_size) < lengths ) return padding_mask def generate_qkv( q, k, v, query_padding_mask=None, key_padding_mask=None, qv=None, kvpacked=False, qkvpacked=False, query_unused_mask=None, key_unused_mask=None, ): """ Arguments: q: (batch_size, seqlen_q, nheads, d) k: (batch_size, seqlen_k, nheads_k, d) v: (batch_size, seqlen_k, nheads_k, d_v) query_padding_mask: (batch_size, seqlen), bool key_padding_mask: (batch_size, seqlen), bool """ assert not (kvpacked and qkvpacked) batch_size, seqlen_q, nheads, d = q.shape d_v = v.shape[-1] _, seqlen_k, nheads_k, _ = k.shape assert k.shape == (batch_size, seqlen_k, nheads_k, d) assert v.shape == (batch_size, seqlen_k, nheads_k, d_v) if query_unused_mask is not None or key_unused_mask is not None: assert not kvpacked assert not qkvpacked if query_padding_mask is not None: q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, seqused_q = unpad_input( q, query_padding_mask, query_unused_mask ) output_pad_fn = lambda output_unpad: pad_input( output_unpad, indices_q, batch_size, seqlen_q ) qv_unpad = ( rearrange(qv, "b s ... -> (b s) ...")[indices_q] if qv is not None else None ) else: q_unpad = rearrange(q, "b s h d -> (b s) h d") cu_seqlens_q = torch.arange( 0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32, device=q_unpad.device, ) seqused_q = None max_seqlen_q = seqlen_q output_pad_fn = lambda output_unpad: rearrange( output_unpad, "(b s) h d -> b s h d", b=batch_size ) qv_unpad = rearrange(qv, "b s ... -> (b s) ...") if qv is not None else None if key_padding_mask is not None: k_unpad, indices_k, cu_seqlens_k, max_seqlen_k, seqused_k = unpad_input( k, key_padding_mask, key_unused_mask ) v_unpad, *rest = unpad_input(v, key_padding_mask, key_unused_mask) else: k_unpad = rearrange(k, "b s h d -> (b s) h d") v_unpad = rearrange(v, "b s h d -> (b s) h d") cu_seqlens_k = torch.arange( 0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32, device=k_unpad.device, ) seqused_k = None max_seqlen_k = seqlen_k if qkvpacked: assert (query_padding_mask == key_padding_mask).all() assert nheads == nheads_k qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1) qkv = torch.stack([q, k, v], dim=2) if query_padding_mask is not None: dqkv_pad_fn = lambda dqkv_unpad: pad_input( dqkv_unpad, indices_q, batch_size, seqlen_q ) else: dqkv_pad_fn = lambda dqkv_unpad: rearrange( dqkv_unpad, "(b s) t h d -> b s t h d", b=batch_size ) return ( qkv_unpad.detach().requires_grad_(), cu_seqlens_q, max_seqlen_q, qkv.detach().requires_grad_(), output_pad_fn, dqkv_pad_fn, ) elif kvpacked: kv_unpad = torch.stack([k_unpad, v_unpad], dim=1) kv = torch.stack([k, v], dim=2) dq_pad_fn = output_pad_fn if key_padding_mask is not None: dkv_pad_fn = lambda dkv_unpad: pad_input( dkv_unpad, indices_k, batch_size, seqlen_k ) else: dkv_pad_fn = lambda dkv_unpad: rearrange( dkv_unpad, "(b s) t h d -> b s t h d", b=batch_size ) return ( q_unpad.detach().requires_grad_(), kv_unpad.detach().requires_grad_(), cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, q.detach().requires_grad_(), kv.detach().requires_grad_(), output_pad_fn, dq_pad_fn, dkv_pad_fn, ) else: dq_pad_fn = output_pad_fn if key_padding_mask is not None: dk_pad_fn = lambda dk_unpad: pad_input( dk_unpad, indices_k, batch_size, seqlen_k ) else: dk_pad_fn = lambda dk_unpad: rearrange( dk_unpad, "(b s) h d -> b s h d", b=batch_size ) return ( q_unpad.detach().requires_grad_(), k_unpad.detach().requires_grad_(), v_unpad.detach().requires_grad_(), qv_unpad.detach() if qv is not None else None, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, max_seqlen_q, max_seqlen_k, q.detach().requires_grad_(), k.detach().requires_grad_(), v.detach().requires_grad_(), qv.detach() if qv is not None else None, output_pad_fn, dq_pad_fn, dk_pad_fn, ) def construct_local_mask( seqlen_q, seqlen_k, window_size=(None, None), sink_token_length=0, query_padding_mask=None, key_padding_mask=None, key_leftpad=None, device=None, ): row_idx = rearrange( torch.arange(seqlen_q, device=device, dtype=torch.long), "s -> s 1" ) col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long) if key_leftpad is not None: key_leftpad = rearrange(key_leftpad, "b -> b 1 1 1") col_idx = repeat(col_idx, "s -> b 1 1 s", b=key_leftpad.shape[0]) col_idx = torch.where(col_idx >= key_leftpad, col_idx - key_leftpad, 2**32) sk = ( seqlen_k if key_padding_mask is None else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1") ) sq = ( seqlen_q if query_padding_mask is None else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1") ) if window_size[0] is None: return col_idx > row_idx + sk - sq + window_size[1] else: sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk return torch.logical_or( col_idx > torch.minimum(row_idx + sk - sq + window_size[1], sk), torch.logical_and( col_idx < row_idx + sk - sq - window_size[0], col_idx >= sink_token_length, ), ) def construct_chunk_mask( seqlen_q, seqlen_k, attention_chunk, query_padding_mask=None, key_padding_mask=None, key_leftpad=None, device=None, ): row_idx = rearrange( torch.arange(seqlen_q, device=device, dtype=torch.long), "s -> s 1" ) col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long) if key_leftpad is not None: key_leftpad = rearrange(key_leftpad, "b -> b 1 1 1") col_idx = repeat(col_idx, "s -> b 1 1 s", b=key_leftpad.shape[0]) col_idx = torch.where(col_idx >= key_leftpad, col_idx - key_leftpad, 2**32) sk = ( seqlen_k if key_padding_mask is None else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1") ) sq = ( seqlen_q if query_padding_mask is None else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1") ) sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk # Subtract remainder instead of divide and then multiply to take care of negative values col_limit_left_chunk = row_idx + sk - sq - (row_idx + sk - sq) % attention_chunk return torch.logical_or( col_idx < col_limit_left_chunk, col_idx >= col_limit_left_chunk + attention_chunk, ) def attention_ref( q, k, v, query_padding_mask=None, key_padding_mask=None, key_leftpad=None, attn_bias=None, dropout_p=0.0, dropout_mask=None, causal=False, qv=None, q_descale=None, k_descale=None, v_descale=None, window_size=(None, None), attention_chunk=0, sink_token_length=0, learnable_sink=None, softcap=0.0, upcast=True, reorder_ops=False, intermediate_dtype=None, ): """ Arguments: q: (batch_size, seqlen_q, nheads, head_dim) k: (batch_size, seqlen_k, nheads, head_dim) v: (batch_size, seqlen_k, nheads, head_dim_v) qv: (batch_size, seqlen_q, nheads, head_dim_v) query_padding_mask: (batch_size, seqlen_q) key_padding_mask: (batch_size, seqlen_k) attn_bias: broadcastable to (batch_size, nheads, seqlen_q, seqlen_k) dropout_p: float dropout_mask: (batch_size, nheads, seqlen_q, seqlen_k) causal: whether to apply causal masking upcast: whether to cast all inputs to fp32, do all computation in fp32, then cast output back to fp16/bf16. reorder_ops: whether to change the order of operations (scaling k instead of scaling k, etc.) without changing the math. This is to estimate the numerical error from operation reordering. Output: output: (batch_size, seqlen_q, nheads, head_dim_v) attention: (batch_size, nheads, seqlen_q, seqlen_k), softmax after dropout """ if causal: window_size = (window_size[0], 0) dtype_og = q.dtype if upcast: q, k, v = q.float(), k.float(), v.float() qv = qv.float() if qv is not None else None if q_descale is not None: q_descale = repeat(q_descale, "b h -> b 1 (h g) 1", g=q.shape[2] // k.shape[2]) q = (q.float() * q_descale).to(q.dtype) qv = (qv.float() * q_descale).to(qv.dtype) if qv is not None else None if k_descale is not None: k = (k.float() * rearrange(k_descale, "b h -> b 1 h 1")).to(dtype=k.dtype) if v_descale is not None: v = (v.float() * rearrange(v_descale, "b h -> b 1 h 1")).to(dtype=v.dtype) seqlen_q, seqlen_k = q.shape[1], k.shape[1] k = repeat(k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2]) v = repeat(v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2]) d = q.shape[-1] dv = v.shape[-1] softmax_scale = 1.0 / math.sqrt(d if qv is None else d + dv) if not reorder_ops: scores = torch.einsum("bthd,bshd->bhts", q * softmax_scale, k) else: scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale) if qv is not None: scores = scores + torch.einsum("bthd,bshd->bhts", qv * softmax_scale, v) if softcap > 0: scores = torch.tanh(scores / softcap) * softcap if key_padding_mask is not None: scores.masked_fill_( rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf") ) local_mask = None if window_size[0] is not None or window_size[1] is not None: local_mask = construct_local_mask( seqlen_q, seqlen_k, window_size, sink_token_length, query_padding_mask, key_padding_mask, key_leftpad=key_leftpad, device=q.device, ) if attention_chunk > 0: chunk_mask = construct_chunk_mask( seqlen_q, seqlen_k, attention_chunk, query_padding_mask, key_padding_mask, key_leftpad=key_leftpad, device=q.device, ) local_mask = ( torch.logical_or(local_mask, chunk_mask) if local_mask is not None else chunk_mask ) if local_mask is not None: scores.masked_fill_(local_mask, float("-inf")) if attn_bias is not None: scores = scores + attn_bias if learnable_sink is None: attention = torch.softmax(scores, dim=-1).to(v.dtype) else: scores_fp32 = scores.to(torch.float32) logits_max = torch.amax(scores_fp32, dim=-1, keepdim=True) learnable_sink = rearrange(learnable_sink, "h -> h 1 1") logits_or_sinks_max = torch.maximum(learnable_sink, logits_max) unnormalized_scores = torch.exp(scores_fp32 - logits_or_sinks_max) normalizer = unnormalized_scores.sum(dim=-1, keepdim=True) + torch.exp( learnable_sink - logits_or_sinks_max ) attention = (unnormalized_scores / normalizer).to(v.dtype) # We want to mask here so that the attention matrix doesn't have any NaNs # Otherwise we'll get NaN in dV if query_padding_mask is not None: attention = attention.masked_fill( rearrange(~query_padding_mask, "b s -> b 1 s 1"), 0.0 ) # Without this we might get NaN in dv if key_padding_mask is not None: attention = attention.masked_fill( rearrange(~key_padding_mask, "b s -> b 1 1 s"), 0.0 ) # Some rows might be completely masked out so we fill them with zero instead of NaN if local_mask is not None: attention = attention.masked_fill( torch.all(local_mask, dim=-1, keepdim=True), 0.0 ) dropout_scaling = 1.0 / (1 - dropout_p) # attention_drop = attention.masked_fill(~dropout_mask, 0.0) * dropout_scaling # output = torch.einsum('bhts,bshd->bthd', attention_drop , v) if dropout_mask is not None: attention_drop = attention.masked_fill(~dropout_mask, 0.0) else: attention_drop = attention if intermediate_dtype is not None: attention_drop = attention_drop.to(intermediate_dtype).to(attention_drop.dtype) output = torch.einsum("bhts,bshd->bthd", attention_drop, v * dropout_scaling) if query_padding_mask is not None: output.masked_fill_(rearrange(~query_padding_mask, "b s -> b s 1 1"), 0.0) return output.to(dtype=dtype_og), attention.to(dtype=dtype_og) @pytest.mark.skipif( skip_condition, reason="FA4 Requires compute capability of 10 or above." ) # @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float8_e4m3fn]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("mha_type", ["mha", "mqa", "gqa"]) # @pytest.mark.parametrize("mha_type", ["mqa"]) @pytest.mark.parametrize("has_learnable_sink", [False, True]) # @pytest.mark.parametrize("has_learnable_sink", [False]) # @pytest.mark.parametrize("has_qv", [False, True]) @pytest.mark.parametrize("has_qv", [False]) # @pytest.mark.parametrize("deterministic", [False, True]) @pytest.mark.parametrize("deterministic", [False]) # @pytest.mark.parametrize("softcap", [0.0, 15.0]) @pytest.mark.parametrize("softcap", [0.0]) # @pytest.mark.parametrize("local", [False, True]) @pytest.mark.parametrize("local", [False]) @pytest.mark.parametrize("causal", [False, True]) # @pytest.mark.parametrize("causal", [False]) # @pytest.mark.parametrize("add_unused_qkv", [False, True]) @pytest.mark.parametrize("add_unused_qkv", [False]) # @pytest.mark.parametrize("d", [32, 64, 96, 128, 160, 192, 224, 256]) # @pytest.mark.parametrize('d', [32, 40, 64, 80, 96, 128, 160, 192, 256]) # @pytest.mark.parametrize('d', [32, 64, 96, 128, 160, 192]) # @pytest.mark.parametrize('d', [56, 80]) # @pytest.mark.parametrize('d', [32, 40, 64, 80, 96, 128]) # @pytest.mark.parametrize("d", [64, 96, 128]) @pytest.mark.parametrize("d", [64, 128]) # @pytest.mark.parametrize("d", [192]) @pytest.mark.parametrize( "seqlen_q,seqlen_k", [ # (1, 1), # (1, 3), # (2, 1), (511, 1), (3, 513), (64, 128), (128, 128), (256, 256), # (113, 203), # (128, 217), # (113, 211), # (108, 256), # (256, 512), (307, 256), (640, 128), (512, 256), (1024, 1024), (1023, 1024), (1024, 1023), (2048, 2048), ], ) def test_flash_attn_varlen_output( seqlen_q, seqlen_k, d, add_unused_qkv, causal, local, softcap, deterministic, has_qv, has_learnable_sink, mha_type, dtype, ): if ( causal or local ): # Right now we only support causal attention with seqlen_k == seqlen_q seqlen_k = seqlen_q device = "cuda" # set seed torch.random.manual_seed(seqlen_q + seqlen_k + d + int(causal) * 2 + int(local)) batch_size = 49 if seqlen_q <= 1024 else 7 nheads = 6 # batch_size = 1 # nheads = 1 nheads_kv = nheads if mha_type == "mha" else (3 if mha_type == "gqa" else 1) dtype_ref = torch.bfloat16 if dtype == torch.float8_e4m3fn else dtype # dv_vals = [128, d] if d > 128 and d <= 192 else ([256, 512, d] if d <= 64 else [d]) dv_vals = [128] if d == 192 else ([d] if d != 128 else [64, d]) if dtype == torch.float8_e4m3fn: dv_vals = [d] # attention_chunk_vals = [torch.randint(1, seqlen_k * 2, (1,)).item(), 0] if seqlen_q <= seqlen_k else [0] attention_chunk_vals = [0] for dv, attention_chunk in itertools.product(dv_vals, attention_chunk_vals): q_ref = torch.randn( batch_size, seqlen_q, nheads, d, device=device, dtype=dtype_ref ) if softcap > 0.0: # Ensure the values of qk are at least within softcap range. q_ref = (q_ref * softcap / 4).detach().requires_grad_() q_ref = q_ref.to(dtype).to(dtype_ref).requires_grad_() k_ref = ( torch.randn( batch_size, seqlen_k, nheads_kv, d, device=device, dtype=dtype_ref ) .to(dtype) .to(dtype_ref) .requires_grad_() ) v_ref = ( torch.randn( batch_size, seqlen_k, nheads_kv, dv, device=device, dtype=dtype_ref ) .to(dtype) .to(dtype_ref) .requires_grad_() ) if has_qv: qv_ref = ( torch.randn( batch_size, seqlen_q, nheads, dv, device=device, dtype=dtype_ref ) .to(dtype) .to(dtype_ref) ) else: qv_ref = None # Put window_size after QKV randn so that window_size changes from test to test window_size = ( (None, None) if not local else torch.randint(0, seqlen_k, (2,)).tolist() ) if has_learnable_sink: learnable_sink = torch.randn(nheads, dtype=torch.bfloat16, device=device) else: learnable_sink = None if dtype == torch.float8_e4m3fn: q_descale, k_descale, v_descale = [ torch.rand(batch_size, nheads_kv, device=device, dtype=torch.float32) * 2 for _ in range(3) ] else: q_descale, k_descale, v_descale = None, None, None q, k, v = [x.detach().requires_grad_() for x in (q_ref, k_ref, v_ref)] qv = qv_ref.detach() if has_qv else None query_padding_mask = generate_random_padding_mask( seqlen_q, batch_size, device, mode="random", zero_lengths=False ) # TODO: test zero_lengths key_padding_mask = generate_random_padding_mask( # seqlen_k, batch_size, device, mode="random", zero_lengths=True seqlen_k, batch_size, device, mode="random", zero_lengths=False, ) def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): if add_unused: another_mask = generate_random_padding_mask(max_seq_len, bs, device) attn_mask = torch.logical_and(padding_mask, another_mask) unused_mask = torch.logical_xor( torch.logical_or(padding_mask, another_mask), attn_mask ) else: attn_mask = padding_mask unused_mask = None return attn_mask, unused_mask query_padding_mask, query_unused_mask = _gen_unused_masks( query_padding_mask, add_unused_qkv, seqlen_q, batch_size, q.device ) # query_padding_mask[:] = True # query_unused_mask = None key_padding_mask, key_unused_mask = _gen_unused_masks( key_padding_mask, add_unused_qkv, seqlen_k, batch_size, k.device ) if causal or local: key_padding_mask = query_padding_mask result = generate_qkv( q, k, v, query_padding_mask, key_padding_mask, qv=qv, kvpacked=False, query_unused_mask=query_unused_mask, key_unused_mask=key_unused_mask, ) ( q_unpad, # 0 k_unpad, # 1 v_unpad, # 2 qv_unpad, # 3 cu_seqlens_q, # 4 cu_seqlens_k, # 5 seqused_q, # 6 seqused_k, # 7 max_seqlen_q, # 8 max_seqlen_k, # 9 q, # 10 k, # 11 v, # 12 qv, # 13 output_pad_fn, # 14 dq_pad_fn, # 15 dk_pad_fn, # 16 ) = result q_unpad, k_unpad, v_unpad = [ x.detach().to(dtype).requires_grad_() for x in (q_unpad, k_unpad, v_unpad) ] out_ref, attn_ref = attention_ref( q_ref, k_ref, v_ref, query_padding_mask, key_padding_mask, causal=causal, qv=qv_ref, q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, window_size=window_size, attention_chunk=attention_chunk, learnable_sink=learnable_sink, softcap=softcap, ) out_pt, attn_pt = attention_ref( q_ref, k_ref, v_ref, query_padding_mask, key_padding_mask, causal=causal, qv=qv_ref, q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, window_size=window_size, attention_chunk=attention_chunk, learnable_sink=learnable_sink, softcap=softcap, upcast=False, reorder_ops=True, intermediate_dtype=dtype if dtype == torch.float8_e4m3fn else None, ) print(f"Pytorch max diff: {(out_pt - out_ref).abs().max().item()}") print(f"Pytorch mean diff: {(out_pt - out_ref).abs().mean().item()}") if query_unused_mask is not None: q_zero_masking = rearrange(query_unused_mask, "b s -> b s 1 1") # Numerical error if we just do any arithmetic on out_ref fwd_atol = 2 * (out_ref + 0.3 - 0.3 - out_ref).abs().max().item() rtol = 2 if softcap == 0.0 else 3 pack_gqa_vals = [False, True, None] # num_splits_vals = [1, 3] num_splits_vals = [1] for pack_gqa, num_splits in itertools.product(pack_gqa_vals, num_splits_vals): out_unpad, lse = flash_attn_varlen_func( q_unpad, k_unpad, v_unpad, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, # max_seqlen_q and max_seqlen_k not needed for FA4 seqused_q=seqused_q, seqused_k=seqused_k, causal=causal, window_size=window_size, softcap=softcap, sinks=learnable_sink, # FA4 uses learnable_sink, not sinks pack_gqa=pack_gqa, return_softmax_lse=True, ) out = output_pad_fn(out_unpad) if query_unused_mask is not None: out.masked_fill_(q_zero_masking, 0.0) print(f"Output max diff: {(out - out_ref).abs().max().item()}") print(f"Output mean diff: {(out - out_ref).abs().mean().item()}") # if not causal: # print(f"LSE max diff: {(lse - lse_ref).abs().max().item()}") # breakpoint() # Check that FlashAttention's numerical error is at most 3x the numerical error # of a Pytorch implementation. assert (out - out_ref).abs().max().item() <= rtol * ( out_pt - out_ref ).abs().max().item() + fwd_atol if ( dtype != torch.float8_e4m3fn and not has_qv and not dv > 256 and not attention_chunk != 0 and dv == d and not has_learnable_sink and False ): g_unpad = torch.randn_like(out_unpad) do_o = ((g_unpad.float() * out_unpad.float()).sum(-1)).transpose(-1, -2) # import flash_attn_3_cuda # dq_unpad, dk_unpad, dv_unpad, softmax_d, dq_accum, lse_log2 = flash_attn_3_cuda.bwd_varlen( # g_unpad, # q_unpad, # k_unpad, # v_unpad, # out_unpad, # lse, # None, # None, # None, # cu_seqlens_q, # cu_seqlens_k, # None, None, # max_seqlen_q, # max_seqlen_k, # d ** (-0.5), # causal, # window_size[0], window_size[1], # softcap, # deterministic, # 0, # sm_margin # ) dq_unpad, dk_unpad, dv_unpad = torch.autograd.grad( out_unpad, (q_unpad, k_unpad, v_unpad), g_unpad ) dq = dq_pad_fn(dq_unpad) dk = dk_pad_fn(dk_unpad) dv = dk_pad_fn(dv_unpad) if key_unused_mask is not None: k_zero_masking = rearrange(key_unused_mask, "b s -> b s 1 1") dk.masked_fill_(k_zero_masking, 0.0) dv.masked_fill_(k_zero_masking, 0.0) if query_unused_mask is not None: dq.masked_fill_(q_zero_masking, 0.0) # print(f"dO_O max diff: {(softmax_d - do_o).abs().max().item()}") # assert (softmax_d - do_o).abs().max().item() <= 1e-5 # assert dq_accum.abs().max().item() == 0.0 g = output_pad_fn(g_unpad) # qk = torch.einsum('bthd,bshd->bhts', q / (d ** 0.5), k).float() # qk = torch.masked_fill(qk, rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf")) # dS = torch.einsum('bthd,bshd->bhts', g.float(), v.float()) # P = torch.softmax(qk, -1) # dP = P * (dS - (g.float() * out.float()).sum(-1).transpose(1, 2).unsqueeze(-1)) # dQ = torch.einsum('bhts,bshd->bthd', dP, k.float()) # dV = torch.einsum('bhts,bthd->bshd', P, g.float()) # dK = torch.einsum('bhts,bthd->bshd', dP, q.float()) # dq, dk, dv = torch.autograd.grad(out, (q, k, v), g) dq_ref, dk_ref, dv_ref = torch.autograd.grad( out_ref, (q_ref, k_ref, v_ref), g ) dq_pt, dk_pt, dv_pt = torch.autograd.grad(out_pt, (q_ref, k_ref, v_ref), g) print(f"dQ max diff: {(dq - dq_ref).abs().max().item()}") print(f"dK max diff: {(dk - dk_ref).abs().max().item()}") print(f"dV max diff: {(dv - dv_ref).abs().max().item()}") print(f"dQ mean diff: {(dq - dq_ref).abs().mean().item()}") print(f"dK mean diff: {(dk - dk_ref).abs().mean().item()}") print(f"dV mean diff: {(dv - dv_ref).abs().mean().item()}") print(f"dQ Pytorch max diff: {(dq_pt - dq_ref).abs().max().item()}") print(f"dK Pytorch max diff: {(dk_pt - dk_ref).abs().max().item()}") print(f"dV Pytorch max diff: {(dv_pt - dv_ref).abs().max().item()}") print(f"dQ Pytorch mean diff: {(dq_pt - dq_ref).abs().mean().item()}") print(f"dK Pytorch mean diff: {(dk_pt - dk_ref).abs().mean().item()}") print(f"dV Pytorch mean diff: {(dv_pt - dv_ref).abs().mean().item()}") # breakpoint() dq_atol = 2 * (dq_ref + 0.3 - 0.3 - dq_ref).abs().max().item() + ( 0 if softcap == 0 else 3e-4 ) assert (dq - dq_ref).abs().max().item() <= rtol * ( dq_pt - dq_ref ).abs().max().item() + dq_atol dk_atol = 2 * (dk_ref + 0.3 - 0.3 - dk_ref).abs().max().item() + ( 0 if softcap == 0 else 3e-4 ) assert (dk - dk_ref).abs().max().item() <= rtol * ( dk_pt - dk_ref ).abs().max().item() + dk_atol dv_atol = 2 * (dv_ref + 0.3 - 0.3 - dv_ref).abs().max().item() + ( 0 if softcap == 0 else 3e-4 ) assert (dv - dv_ref).abs().max().item() <= rtol * ( dv_pt - dv_ref ).abs().max().item() + dv_atol @pytest.mark.skipif( skip_condition, reason="FA4 Requires compute capability of 10 or above." ) # @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float8_e4m3fn]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) # @pytest.mark.parametrize("dtype", [torch.float8_e4m3fn]) @pytest.mark.parametrize("mha_type", ["mha", "mqa", "gqa"]) # @pytest.mark.parametrize("mha_type", ["mha"]) @pytest.mark.parametrize("has_learnable_sink", [False, True]) # @pytest.mark.parametrize("has_learnable_sink", [False]) # @pytest.mark.parametrize("new_kv", [False, True]) @pytest.mark.parametrize("new_kv", [False]) # @pytest.mark.parametrize("local", [False, True]) @pytest.mark.parametrize("local", [False]) # @pytest.mark.parametrize("causal", [False, True]) @pytest.mark.parametrize("causal", [True]) # @pytest.mark.parametrize("seqlen_new_eq_seqlen_q", [True, False]) @pytest.mark.parametrize("seqlen_new_eq_seqlen_q", [False]) # @pytest.mark.parametrize("has_rotary_seqlens", [False, True]) @pytest.mark.parametrize("has_rotary_seqlens", [False]) # @pytest.mark.parametrize("rotary_interleaved", [False, True]) @pytest.mark.parametrize("rotary_interleaved", [True]) # @pytest.mark.parametrize("rotary_fraction", [0.0, 0.5, 1.0]) @pytest.mark.parametrize("rotary_fraction", [0.0]) # @pytest.mark.parametrize("page_size", [None] + ([1, 4, 128])) # @pytest.mark.parametrize("page_size", [None, 128]) @pytest.mark.parametrize("page_size", [128]) # @pytest.mark.parametrize("has_leftpad", [False, True]) @pytest.mark.parametrize("has_leftpad", [False]) # @pytest.mark.parametrize("has_batch_idx", [False, True]) @pytest.mark.parametrize("has_batch_idx", [False]) # @pytest.mark.parametrize("varlen_q", [False, True]) @pytest.mark.parametrize("varlen_q", [False]) # @pytest.mark.parametrize("d", [32, 59, 64, 80, 128, 256]) # @pytest.mark.parametrize("d", [32, 64, 96, 128, 160, 192, 224, 256]) # @pytest.mark.parametrize('d', [32, 40, 64, 80, 96, 128, 160, 192]) # @pytest.mark.parametrize('d', [56, 80]) # @pytest.mark.parametrize("d", [128]) @pytest.mark.parametrize("d", [64]) # @pytest.mark.parametrize("d", [192]) @pytest.mark.parametrize( "seqlen_q,seqlen_k", [ (1, 128), (1, 339), (3, 1024), (64, 800), (64, 256), (3, 799), (64, 2048), (16, 20000), # # (1, 128 * 1024), # # (16, 128 * 1024), # (128, 128), # (256, 512), # To test appending KV with more than 1 block # (2048, 3577), # Enough tile to test persistent scheduler ], ) # @pytest.mark.parametrize('seqlen_q,seqlen_k', [(256, 128)]) def test_flash_attn_kvcache( seqlen_q, seqlen_k, d, varlen_q, has_batch_idx, has_leftpad, page_size, rotary_fraction, rotary_interleaved, has_rotary_seqlens, seqlen_new_eq_seqlen_q, causal, local, new_kv, has_learnable_sink, mha_type, dtype, ): if page_size is not None and seqlen_k % page_size != 0: pytest.skip() if seqlen_q > seqlen_k and new_kv: pytest.skip() if not new_kv and rotary_fraction > 0.0: pytest.skip() if rotary_fraction == 0.0 and has_rotary_seqlens: pytest.skip() device = "cuda" # set seed torch.random.manual_seed(0) batch_size = 5 # batch_size = 1 batch_size_cache = batch_size if not has_batch_idx else batch_size * 2 nheads = 6 # nheads = 1 # rotary_dim must be a multiple of 16, and must be <= d rotary_dim = math.floor(int(rotary_fraction * d) / 16) * 16 nheads_k = nheads if mha_type == "mha" else (1 if mha_type == "mqa" else 3) assert nheads % nheads_k == 0 dtype_ref = torch.bfloat16 if dtype == torch.float8_e4m3fn else dtype # dv_vals = [128, d] if d > 128 and d <= 192 else ([256, 512, d] if d <= 64 else [d]) dv_vals = [d] if dtype == torch.float8_e4m3fn: dv_vals = [d] # attention_chunk_vals = [torch.randint(1, seqlen_k * 2, (1,)).item(), 0] if (causal or local) else [0] attention_chunk_vals = [0] for dv, attention_chunk in itertools.product(dv_vals, attention_chunk_vals): # has_qv = d == 64 and dv >= 256 has_qv = False q = ( torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype_ref) .to(dtype) .to(dtype_ref) ) if has_qv: qv = ( torch.randn( batch_size, seqlen_q, nheads, dv, device=device, dtype=dtype_ref ) .to(dtype) .to(dtype_ref) ) else: qv = None if varlen_q: query_padding_mask = generate_random_padding_mask( seqlen_q, batch_size, device, mode="random" ) q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, *rest = unpad_input( q, query_padding_mask ) output_pad_fn = lambda output_unpad: pad_input( output_unpad, indices_q, batch_size, seqlen_q ) qv_unpad = ( rearrange(qv, "b s ... -> (b s) ...")[indices_q] if has_qv else None ) else: query_padding_mask = None q_unpad = q qv_unpad = qv cu_seqlens_q, max_seqlen_q = None, None # Put window_size after QKV randn so that window_size changes from test to test window_size = ( (None, None) if not local else torch.randint(0, seqlen_k, (2,)).tolist() ) if has_learnable_sink: learnable_sink = torch.randn(nheads, dtype=torch.bfloat16, device=device) else: learnable_sink = None seqlen_new = ( seqlen_q if seqlen_new_eq_seqlen_q else torch.randint(1, seqlen_q + 1, (1,)).item() ) cu_seqlens_k_new = None key_new_padding_mask = None if new_kv: k = ( torch.randn( batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype_ref ) .to(dtype) .to(dtype_ref) ) v = ( torch.randn( batch_size, seqlen_new, nheads_k, dv, device=device, dtype=dtype_ref ) .to(dtype) .to(dtype_ref) ) if varlen_q: # k & v are also varlen key_new_padding_mask = generate_random_padding_mask( seqlen_new, batch_size, device, mode="random" ) k_unpad, indices_k, cu_seqlens_k_new, *rest = unpad_input( k, key_new_padding_mask ) v_unpad, *rest = unpad_input(v, key_new_padding_mask) else: k_unpad, v_unpad = k, v else: k, v, k_unpad, v_unpad = None, None, None, None if page_size is None: k_cache = ( torch.randn( batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype_ref, ) .to(dtype) .to(dtype_ref) ) v_cache = ( torch.randn( batch_size_cache, seqlen_k, nheads_k, dv, device=device, dtype=dtype_ref, ) .to(dtype) .to(dtype_ref) ) page_table = None num_blocks = None else: ( k_cache, v_cache, page_table, k_cache_paged, v_cache_paged, num_blocks, ) = _generate_block_kvcache( seqlen_k, page_size, batch_size_cache, nheads_k, d, dv, device, dtype, dtype_ref, ) cache_seqlens = torch.randint( 0 if new_kv else 1, # If we don't use seqlen_q in the case of causal and rotary, cos/sin won't be long enough ( ( seqlen_k - (seqlen_q if (causal or local) and rotary_dim > 1 else seqlen_new) + 1 ) if new_kv else (seqlen_k + 1) ), (batch_size,), dtype=torch.int32, device=device, ) if has_leftpad: cache_leftpad = torch.cat( [ ( torch.randint( 0, cache_seqlens[i].item(), (1,), dtype=torch.int32, device=device, ) if cache_seqlens[i].item() > 0 else torch.zeros(1, dtype=torch.int32, device=device) ) for i in range(batch_size) ] ) else: cache_leftpad = None if has_batch_idx: cache_batch_idx = torch.randperm( batch_size_cache, dtype=torch.int32, device=device )[:batch_size] else: cache_batch_idx = None arange = rearrange(torch.arange(seqlen_k, device=device), "s -> 1 s") cache_seqlens_expanded = rearrange(cache_seqlens, "b -> b 1") if not new_kv: key_padding_mask = arange < cache_seqlens_expanded else: k_new_seqlens = ( key_new_padding_mask.sum(-1, keepdims=True) if varlen_q else seqlen_new ) key_padding_mask = arange < cache_seqlens_expanded + k_new_seqlens if has_leftpad: key_padding_mask = torch.logical_and( key_padding_mask, arange >= cache_leftpad.unsqueeze(-1).expand(-1, seqlen_k), ) # cache_seqlens = torch.tensor([64], dtype=torch.int32, device=device) rotary_seqlens = cache_seqlens if not has_rotary_seqlens else cache_seqlens // 2 if rotary_dim > 0: angle = ( torch.rand( seqlen_k if page_size is None else num_blocks * page_size, rotary_dim // 2, device=device, ) * 2 * math.pi ) cos = torch.cos(angle).to(dtype=dtype_ref).to(dtype).to(dtype_ref) sin = torch.sin(angle).to(dtype=dtype_ref).to(dtype).to(dtype_ref) if causal or local: q_ro = apply_rotary_emb( q, cos, sin, seqlen_offsets=rotary_seqlens, interleaved=rotary_interleaved, ) else: q_ro = rearrange( apply_rotary_emb( rearrange(q, "b s h d -> b 1 (s h) d"), cos, sin, seqlen_offsets=rotary_seqlens, interleaved=rotary_interleaved, ), "b 1 (s h) d -> b s h d", s=seqlen_q, ) # q_ro = q k_ro = apply_rotary_emb( k, cos, sin, seqlen_offsets=rotary_seqlens, interleaved=rotary_interleaved, ) else: cos, sin = None, None q_ro, k_ro = q, k # k_cache[:, 64:] = -1 k_cache_ref = ( k_cache if not has_batch_idx else k_cache[cache_batch_idx] ).clone() v_cache_ref = ( v_cache if not has_batch_idx else v_cache[cache_batch_idx] ).clone() if new_kv: update_mask = torch.logical_and( cache_seqlens_expanded <= arange, arange < cache_seqlens_expanded + k_new_seqlens, ) k_to_update = rearrange(k_ro, "b s ... -> (b s) ...") v_to_update = rearrange(v, "b s ... -> (b s) ...") if varlen_q: k_to_update = k_to_update[indices_k] v_to_update = v_to_update[indices_k] k_cache_ref[update_mask] = k_to_update v_cache_ref[update_mask] = v_to_update k_cache_rep = repeat( k_cache_ref, "b s h d -> b s (h g) d", g=nheads // nheads_k ) v_cache_rep = repeat( v_cache_ref, "b s h d -> b s (h g) d", g=nheads // nheads_k ) out_ref, _ = attention_ref( q_ro, k_cache_rep, v_cache_rep, query_padding_mask, key_padding_mask, causal=causal, qv=qv, window_size=window_size, learnable_sink=learnable_sink, attention_chunk=attention_chunk, key_leftpad=cache_leftpad, ) out_pt, _ = attention_ref( q_ro, k_cache_rep, v_cache_rep, query_padding_mask, key_padding_mask, causal=causal, qv=qv, window_size=window_size, learnable_sink=learnable_sink, attention_chunk=attention_chunk, upcast=False, reorder_ops=True, key_leftpad=cache_leftpad, intermediate_dtype=dtype if dtype == torch.float8_e4m3fn else None, ) q = q.to(dtype) q_unpad = q_unpad.to(dtype) if varlen_q else None k_cache = k_cache.to(dtype) v_cache = v_cache.to(dtype) k_cache_paged = k_cache_paged.to(dtype) if page_size is not None else None v_cache_paged = v_cache_paged.to(dtype) if page_size is not None else None k = k.to(dtype) if k is not None else None v = v.to(dtype) if v is not None else None k_unpad = k_unpad.to(dtype) if k_unpad is not None else None v_unpad = v_unpad.to(dtype) if v_unpad is not None else None qv = qv.to(dtype) if qv is not None else None qv_unpad = qv_unpad.to(dtype) if (varlen_q and qv is not None) else None cos = cos.to(dtype) if cos is not None else None sin = sin.to(dtype) if sin is not None else None k_cache_saved = k_cache.clone() if page_size is None else k_cache_paged.clone() v_cache_saved = v_cache.clone() if page_size is None else v_cache_paged.clone() # num_splits_vals = [1, 0] num_splits_vals = [1] # precompute_metadata_vals = [False, True] precompute_metadata_vals = [False] for num_splits, precompute_metadata in itertools.product( num_splits_vals, precompute_metadata_vals ): # if precompute_metadata: # scheduler_metadata = get_scheduler_metadata( # batch_size, max_seqlen_q if varlen_q else seqlen_q, seqlen_k, nheads, nheads_k, d, # cache_seqlens, q.dtype, headdim_v=dv, cu_seqlens_q=cu_seqlens_q, # cu_seqlens_k_new=cu_seqlens_k_new, cache_leftpad=cache_leftpad, # max_seqlen_k_new=seqlen_new, page_size=page_size, # causal=causal, window_size=window_size, attention_chunk=attention_chunk, # num_splits=num_splits # ) # else: # scheduler_metadata = None scheduler_metadata = None # Repeat to test metadata reuse for _ in range(1 if not precompute_metadata else 2): if page_size is None: k_cache.copy_(k_cache_saved) v_cache.copy_(v_cache_saved) else: k_cache_paged.copy_(k_cache_saved) v_cache_paged.copy_(v_cache_saved) # For FA4, use flash_attn_varlen_func directly instead of flash_attn_with_kvcache # This matches the pattern from the original FA4 test out, lse = flash_attn_varlen_func( q if not varlen_q else q_unpad, k_cache if page_size is None else k_cache_paged, v_cache if page_size is None else v_cache_paged, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=None, # FA4 doesn't use cu_seqlens_k for KV cache # max_seqlen_q and max_seqlen_k not needed for FA4 seqused_k=cache_seqlens, # Use cache_seqlens as seqused_k page_table=page_table, causal=causal, window_size=window_size, sinks=learnable_sink, # FA4 uses learnable_sink, not sinks softcap=0.0, pack_gqa=None, return_softmax_lse=True, ) if varlen_q: out = output_pad_fn(out) # out = flash_attn_with_kvcache( # q, k_cache, v_cache, cache_seqlens=cache_seqlens, causal=causal, window_size=window_size # ) # out = flash_attn_with_kvcache(q, k_cache, v_cache, causal=causal, window_size=window_size) # qk = torch.einsum("bqhd,bkhd->bhqk", q, k_cache_ref) # m = qk.amax(-1, keepdim=True) # s_tmp = torch.exp((qk - m) / math.sqrt(d)) # o1 = torch.einsum('bhst,bthd->bshd', s_tmp, v_cache_ref) # lse_ref = torch.logsumexp(qk / math.sqrt(d), -1) # probs = torch.softmax(qk, dim=-1) print(f"Output max diff: {(out - out_ref).abs().max().item()}") print(f"Output mean diff: {(out - out_ref).abs().mean().item()}") print(f"Pytorch max diff: {(out_pt - out_ref).abs().max().item()}") print(f"Pytorch mean diff: {(out_pt - out_ref).abs().mean().item()}") # breakpoint() # Check that FlashAttention's numerical error is at most twice the numerical error # of a Pytorch implementation. if new_kv: if page_size is None: k_cache_select = ( k_cache.to(dtype_ref) if not has_batch_idx else k_cache.to(dtype_ref)[cache_batch_idx] ) v_cache_select = ( v_cache.to(dtype_ref) if not has_batch_idx else v_cache.to(dtype_ref)[cache_batch_idx] ) else: k_cache_select = rearrange( k_cache_paged.to(dtype_ref)[ ( page_table if not has_batch_idx else page_table[cache_batch_idx] ).flatten() ], "(b nblocks) block_size ... -> b (nblocks block_size) ...", b=batch_size, )[:, :seqlen_k].to(dtype_ref) v_cache_select = rearrange( v_cache_paged.to(dtype_ref)[ ( page_table if not has_batch_idx else page_table[cache_batch_idx] ).flatten() ], "(b nblocks) block_size ... -> b (nblocks block_size) ...", b=batch_size, )[:, :seqlen_k].to(dtype_ref) k_cache_ref = k_cache_ref.to(dtype).to(dtype_ref) v_cache_ref = v_cache_ref.to(dtype).to(dtype_ref) if dtype is not torch.float8_e4m3fn: assert torch.equal(v_cache_select, v_cache_ref) else: assert torch.allclose( v_cache_select, v_cache_ref, rtol=1e-3, atol=1e-3 ) # breakpoint() # if rotary_dim == 0 and dtype is not torch.float8_e4m3fn: if rotary_dim == 0: assert torch.equal(k_cache_select, k_cache_ref) else: # if not torch.allclose(k_cache_select, k_cache_ref, rtol=1e-3, atol=1e-3): # breakpoint() if dtype is not torch.float8_e4m3fn: assert torch.allclose( k_cache_select, k_cache_ref, rtol=1e-3, atol=1e-3 ) else: assert torch.allclose( k_cache_select, k_cache_ref, rtol=1e-1, atol=1e-1 ) mult = 4 if dtype == torch.float8_e4m3fn else 2 assert (out - out_ref).abs().max().item() <= mult * ( out_pt - out_ref ).abs().max().item() + 1e-5 mult_mean = 3 if dtype == torch.float8_e4m3fn else 1.5 assert (out - out_ref).abs().mean().item() <= mult_mean * ( out_pt - out_ref ).abs().mean().item() def _generate_block_kvcache( seqlen_k, page_size, batch_size, nheads_k, d, dv, device, dtype, dtype_ref ): num_blocks = math.ceil(seqlen_k / page_size) * batch_size * 3 k_cache_paged = ( torch.randn(num_blocks, page_size, nheads_k, d, device=device, dtype=dtype_ref) .to(dtype) .to(dtype_ref) ) v_cache_paged = ( torch.randn(num_blocks, page_size, nheads_k, dv, device=device, dtype=dtype_ref) .to(dtype) .to(dtype_ref) ) page_table = rearrange( torch.randperm(num_blocks, dtype=torch.int32, device=device), "(b nblocks) -> b nblocks", b=batch_size, ) k_cache = rearrange( k_cache_paged[page_table.flatten()], "(b nblocks) block_size ... -> b (nblocks block_size) ...", b=batch_size, )[:, :seqlen_k] v_cache = rearrange( v_cache_paged[page_table.flatten()], "(b nblocks) block_size ... -> b (nblocks block_size) ...", b=batch_size, )[:, :seqlen_k] return k_cache, v_cache, page_table, k_cache_paged, v_cache_paged, num_blocks if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py ================================================ import itertools import pytest import torch from sglang.jit_kernel.utils import get_ci_test_range def sglang_jit_fused_add_rmsnorm( input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float ) -> None: from sglang.jit_kernel.norm import fused_add_rmsnorm fused_add_rmsnorm(input, residual, weight, eps) def flashinfer_fused_add_rmsnorm( input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float ) -> None: from flashinfer.norm import fused_add_rmsnorm fused_add_rmsnorm(input, residual, weight, eps=eps) BS_LIST = [2**n for n in range(0, 14)] BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109]) HIDDEN_SIZE_LIST = get_ci_test_range( [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192], [512, 2048, 8192], ) DEVICE = "cuda" DTYPE = torch.bfloat16 @pytest.mark.parametrize( "batch_size,hidden_size", list(itertools.product(BS_LIST, HIDDEN_SIZE_LIST)) ) def test_fused_add_rmsnorm(batch_size: int, hidden_size: int) -> None: input = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE) residual = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE) weight = torch.randn(hidden_size, device=DEVICE, dtype=DTYPE) input_sglang = input.clone() residual_sglang = residual.clone() input_flashinfer = input.clone() residual_flashinfer = residual.clone() sglang_jit_fused_add_rmsnorm( input_sglang, residual_sglang, weight, torch.finfo(torch.bfloat16).eps ) flashinfer_fused_add_rmsnorm( input_flashinfer, residual_flashinfer, weight, torch.finfo(torch.bfloat16).eps ) torch.testing.assert_close(input_sglang, input_flashinfer, atol=1e-2, rtol=1e-2) torch.testing.assert_close( residual_sglang, residual_flashinfer, atol=1e-2, rtol=1e-2 ) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_fused_metadata_copy.py ================================================ """ Comprehensive tests for JIT-compiled fused metadata copy kernels. This test suite verifies: 1. Single-backend fused kernel (fused_metadata_copy_cuda) - all forward modes 2. Multi-backend fused kernel (fused_metadata_copy_multi_cuda) - 3 backends at once 3. Correctness against reference implementations 4. Performance benchmarks and speedup measurements """ import time import pytest import torch # ============================================================================= # Helper Functions # ============================================================================= def create_test_metadata( bs: int, max_len: int, max_seqlen_k: int, seqlens_expanded_size: int, has_real_page_table: bool = False, has_flashmla: bool = False, device: str = "cuda", ): """Create test metadata tensors matching NSA backend structure.""" # Basic tensors (always present) cache_seqlens_src = torch.randint( 1, max_len, (bs,), dtype=torch.int32, device=device ) cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device) cu_seqlens_k_src[1:] = torch.cumsum(cache_seqlens_src, dim=0) page_indices_src = torch.randint( 0, 1000, (bs, max_len), dtype=torch.int32, device=device ) nsa_cache_seqlens_src = torch.randint( 1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device ) seqlens_expanded_src = torch.randint( 1, max_seqlen_k, (seqlens_expanded_size,), dtype=torch.int32, device=device ) nsa_cu_seqlens_k_src = torch.zeros( seqlens_expanded_size + 1, dtype=torch.int32, device=device ) nsa_cu_seqlens_k_src[1:] = torch.cumsum(nsa_cache_seqlens_src, dim=0) # Destination tensors cache_seqlens_dst = torch.zeros(bs, dtype=torch.int32, device=device) cu_seqlens_k_dst = torch.zeros(bs + 1, dtype=torch.int32, device=device) page_table_1_dst = torch.zeros((bs, max_len + 16), dtype=torch.int32, device=device) nsa_cache_seqlens_dst = torch.zeros( seqlens_expanded_size, dtype=torch.int32, device=device ) nsa_seqlens_expanded_dst = torch.zeros( seqlens_expanded_size, dtype=torch.int32, device=device ) nsa_cu_seqlens_k_dst = torch.zeros( seqlens_expanded_size + 1, dtype=torch.int32, device=device ) # Optional tensors real_page_table_src = None real_page_table_dst = None if has_real_page_table: real_page_table_cols = max_len // 2 real_page_table_src = torch.randint( 0, 1000, (bs, real_page_table_cols), dtype=torch.int32, device=device ) real_page_table_dst = torch.zeros( (bs, real_page_table_cols + 8), dtype=torch.int32, device=device ) flashmla_num_splits_src = None flashmla_num_splits_dst = None flashmla_metadata_src = None flashmla_metadata_dst = None if has_flashmla: flashmla_num_splits_src = torch.randint( 1, 10, (seqlens_expanded_size + 1,), dtype=torch.int32, device=device ) flashmla_num_splits_dst = torch.zeros( seqlens_expanded_size + 1, dtype=torch.int32, device=device ) # FlashMLA metadata is typically (num_sm_parts, TileSchedulerMetaDataSize) # For testing, we use a simplified size flashmla_metadata_size = 128 flashmla_metadata_src = torch.randint( 0, 100, (flashmla_metadata_size,), dtype=torch.int32, device=device ) flashmla_metadata_dst = torch.zeros( flashmla_metadata_size, dtype=torch.int32, device=device ) return { "src": { "cache_seqlens": cache_seqlens_src, "cu_seqlens_k": cu_seqlens_k_src, "page_indices": page_indices_src, "nsa_cache_seqlens": nsa_cache_seqlens_src, "seqlens_expanded": seqlens_expanded_src, "nsa_cu_seqlens_k": nsa_cu_seqlens_k_src, "real_page_table": real_page_table_src, "flashmla_num_splits": flashmla_num_splits_src, "flashmla_metadata": flashmla_metadata_src, }, "dst": { "cache_seqlens": cache_seqlens_dst, "cu_seqlens_k": cu_seqlens_k_dst, "page_table_1": page_table_1_dst, "nsa_cache_seqlens": nsa_cache_seqlens_dst, "nsa_seqlens_expanded": nsa_seqlens_expanded_dst, "nsa_cu_seqlens_k": nsa_cu_seqlens_k_dst, "real_page_table": real_page_table_dst, "flashmla_num_splits": flashmla_num_splits_dst, "flashmla_metadata": flashmla_metadata_dst, }, } def reference_copy_decode(src, dst, max_len): """Reference implementation: individual .copy_() for DECODE mode.""" bs = src["cache_seqlens"].shape[0] dst["cache_seqlens"].copy_(src["cache_seqlens"]) dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) dst["page_table_1"][:, :max_len].copy_(src["page_indices"]) dst["nsa_cache_seqlens"].copy_(src["nsa_cache_seqlens"]) dst["nsa_cu_seqlens_k"][1 : bs + 1].copy_(src["nsa_cu_seqlens_k"][1 : bs + 1]) if src["real_page_table"] is not None: rows, cols = src["real_page_table"].shape dst["real_page_table"][:rows, :cols].copy_(src["real_page_table"]) if src["flashmla_num_splits"] is not None: flashmla_size = bs + 1 dst["flashmla_num_splits"][:flashmla_size].copy_( src["flashmla_num_splits"][:flashmla_size] ) if src["flashmla_metadata"] is not None: dst["flashmla_metadata"].copy_(src["flashmla_metadata"]) def reference_copy_target_verify(src, dst, max_seqlen_k, seqlens_expanded_size): """Reference implementation: individual .copy_() for TARGET_VERIFY mode.""" bs = src["cache_seqlens"].shape[0] dst["cache_seqlens"].copy_(src["cache_seqlens"]) dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) rows, cols = src["page_indices"].shape dst["page_table_1"][:rows, :cols].copy_(src["page_indices"]) dst["nsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"]) dst["nsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["nsa_cache_seqlens"]) dst["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_( src["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1] ) if src["real_page_table"] is not None: rows, cols = src["real_page_table"].shape dst["real_page_table"][:rows, :cols].copy_(src["real_page_table"]) if src["flashmla_num_splits"] is not None: flashmla_size = seqlens_expanded_size + 1 dst["flashmla_num_splits"][:flashmla_size].copy_( src["flashmla_num_splits"][:flashmla_size] ) if src["flashmla_metadata"] is not None: dst["flashmla_metadata"].copy_(src["flashmla_metadata"]) def reference_copy_draft_extend(src, dst, max_seqlen_k, seqlens_expanded_size): """Reference implementation: individual .copy_() for DRAFT_EXTEND mode.""" bs = src["cache_seqlens"].shape[0] dst["cache_seqlens"].copy_(src["cache_seqlens"]) dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) rows, cols = src["page_indices"].shape dst["page_table_1"][:rows, :cols].copy_(src["page_indices"]) dst["nsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"]) dst["nsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["nsa_cache_seqlens"]) dst["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_( src["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1] ) if src["real_page_table"] is not None: rows, cols = src["real_page_table"].shape dst["real_page_table"][:rows, :cols].copy_(src["real_page_table"]) if src["flashmla_num_splits"] is not None: flashmla_size = seqlens_expanded_size + 1 dst["flashmla_num_splits"][:flashmla_size].copy_( src["flashmla_num_splits"][:flashmla_size] ) if src["flashmla_metadata"] is not None: dst["flashmla_metadata"].copy_(src["flashmla_metadata"]) # ============================================================================= # Single-Backend Kernel Tests # ============================================================================= def test_fused_metadata_copy_dtype_validation(): """Test that dtype validation rejects non-int32 tensors.""" if not torch.cuda.is_available(): pytest.skip("CUDA not available") from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_cuda bs = 2 max_len = 128 max_seqlen_k = 256 seqlens_expanded_size = bs device = "cuda" # Create tensors with WRONG dtype (int64 instead of int32) cache_seqlens_src_wrong = torch.randint( 1, max_len, (bs,), dtype=torch.int64, device=device ) cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device) page_indices_src = torch.randint( 0, 1000, (bs, max_len), dtype=torch.int32, device=device ) nsa_cache_seqlens_src = torch.randint( 1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device ) seqlens_expanded_src = torch.randint( 1, max_seqlen_k, (seqlens_expanded_size,), dtype=torch.int32, device=device ) nsa_cu_seqlens_k_src = torch.zeros( seqlens_expanded_size + 1, dtype=torch.int32, device=device ) # Destination tensors (correct dtype) cache_seqlens_dst = torch.zeros(bs, dtype=torch.int32, device=device) cu_seqlens_k_dst = torch.zeros(bs + 1, dtype=torch.int32, device=device) page_table_1_dst = torch.zeros((bs, max_len + 16), dtype=torch.int32, device=device) nsa_cache_seqlens_dst = torch.zeros( seqlens_expanded_size, dtype=torch.int32, device=device ) nsa_seqlens_expanded_dst = torch.zeros( seqlens_expanded_size, dtype=torch.int32, device=device ) nsa_cu_seqlens_k_dst = torch.zeros( seqlens_expanded_size + 1, dtype=torch.int32, device=device ) # Test 1: Wrong dtype for source tensor should raise RuntimeError with pytest.raises(RuntimeError, match="must have dtype int32"): fused_metadata_copy_cuda( cache_seqlens_src_wrong, # Wrong dtype: int64 cu_seqlens_k_src, page_indices_src, nsa_cache_seqlens_src, seqlens_expanded_src, nsa_cu_seqlens_k_src, None, # real_page_table_src None, # flashmla_num_splits_src None, # flashmla_metadata_src cache_seqlens_dst, cu_seqlens_k_dst, page_table_1_dst, nsa_cache_seqlens_dst, nsa_seqlens_expanded_dst, nsa_cu_seqlens_k_dst, None, # real_page_table_dst None, # flashmla_num_splits_dst None, # flashmla_metadata_dst 0, # forward_mode bs, max_len, max_seqlen_k, seqlens_expanded_size, ) # Test 2: Wrong dtype for destination tensor should also raise RuntimeError cache_seqlens_src = torch.randint( 1, max_len, (bs,), dtype=torch.int32, device=device ) cache_seqlens_dst_wrong = torch.zeros(bs, dtype=torch.int64, device=device) with pytest.raises(RuntimeError, match="must have dtype int32"): fused_metadata_copy_cuda( cache_seqlens_src, cu_seqlens_k_src, page_indices_src, nsa_cache_seqlens_src, seqlens_expanded_src, nsa_cu_seqlens_k_src, None, None, None, cache_seqlens_dst_wrong, # Wrong dtype: int64 cu_seqlens_k_dst, page_table_1_dst, nsa_cache_seqlens_dst, nsa_seqlens_expanded_dst, nsa_cu_seqlens_k_dst, None, None, None, 0, bs, max_len, max_seqlen_k, seqlens_expanded_size, ) @pytest.mark.parametrize("bs", [1, 2, 4, 8]) @pytest.mark.parametrize( "forward_mode", [0] ) # DECODE mode only (other modes not fully tested yet) @pytest.mark.parametrize("has_real_page_table", [False, True]) @pytest.mark.parametrize("has_flashmla", [False, True]) def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla): """Test fused metadata copy kernel against reference implementation.""" if not torch.cuda.is_available(): pytest.skip("CUDA not available") from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_cuda max_len = 128 max_seqlen_k = 256 seqlens_expanded_size = bs if forward_mode == 0 else bs * 2 # Create test data data = create_test_metadata( bs=bs, max_len=max_len, max_seqlen_k=max_seqlen_k, seqlens_expanded_size=seqlens_expanded_size, has_real_page_table=has_real_page_table, has_flashmla=has_flashmla, ) # Create separate destination tensors for reference and fused kernel dst_ref = {k: v.clone() if v is not None else None for k, v in data["dst"].items()} dst_fused = { k: v.clone() if v is not None else None for k, v in data["dst"].items() } # Run reference implementation if forward_mode == 0: # DECODE reference_copy_decode(data["src"], dst_ref, max_len) elif forward_mode == 1: # TARGET_VERIFY reference_copy_target_verify( data["src"], dst_ref, max_seqlen_k, seqlens_expanded_size ) else: # DRAFT_EXTEND reference_copy_draft_extend( data["src"], dst_ref, max_seqlen_k, seqlens_expanded_size ) # Run fused kernel fused_metadata_copy_cuda( data["src"]["cache_seqlens"], data["src"]["cu_seqlens_k"], data["src"]["page_indices"], data["src"]["nsa_cache_seqlens"], data["src"]["seqlens_expanded"], data["src"]["nsa_cu_seqlens_k"], data["src"]["real_page_table"], data["src"]["flashmla_num_splits"], data["src"]["flashmla_metadata"], dst_fused["cache_seqlens"], dst_fused["cu_seqlens_k"], dst_fused["page_table_1"], dst_fused["nsa_cache_seqlens"], dst_fused["nsa_seqlens_expanded"], dst_fused["nsa_cu_seqlens_k"], dst_fused["real_page_table"], dst_fused["flashmla_num_splits"], dst_fused["flashmla_metadata"], forward_mode, bs, max_len, max_seqlen_k, seqlens_expanded_size, ) # Compare results assert torch.equal( dst_ref["cache_seqlens"], dst_fused["cache_seqlens"] ), "cache_seqlens mismatch" assert torch.equal( dst_ref["cu_seqlens_k"], dst_fused["cu_seqlens_k"] ), "cu_seqlens_k mismatch" assert torch.equal( dst_ref["page_table_1"], dst_fused["page_table_1"] ), "page_table_1 mismatch" assert torch.equal( dst_ref["nsa_cache_seqlens"], dst_fused["nsa_cache_seqlens"] ), "nsa_cache_seqlens mismatch" assert torch.equal( dst_ref["nsa_seqlens_expanded"], dst_fused["nsa_seqlens_expanded"] ), "nsa_seqlens_expanded mismatch" assert torch.equal( dst_ref["nsa_cu_seqlens_k"], dst_fused["nsa_cu_seqlens_k"] ), "nsa_cu_seqlens_k mismatch" if has_real_page_table: assert torch.equal( dst_ref["real_page_table"], dst_fused["real_page_table"] ), "real_page_table mismatch" if has_flashmla: assert torch.equal( dst_ref["flashmla_num_splits"], dst_fused["flashmla_num_splits"] ), "flashmla_num_splits mismatch" assert torch.equal( dst_ref["flashmla_metadata"], dst_fused["flashmla_metadata"] ), "flashmla_metadata mismatch" @pytest.mark.parametrize("bs", [16, 32]) def test_fused_metadata_copy_large_batch(bs): """Test with larger batch sizes.""" if not torch.cuda.is_available(): pytest.skip("CUDA not available") from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_cuda forward_mode = 0 # DECODE max_len = 128 max_seqlen_k = 256 seqlens_expanded_size = bs data = create_test_metadata( bs=bs, max_len=max_len, max_seqlen_k=max_seqlen_k, seqlens_expanded_size=seqlens_expanded_size, has_real_page_table=True, has_flashmla=True, ) dst_ref = {k: v.clone() if v is not None else None for k, v in data["dst"].items()} dst_fused = { k: v.clone() if v is not None else None for k, v in data["dst"].items() } reference_copy_decode(data["src"], dst_ref, max_len) fused_metadata_copy_cuda( data["src"]["cache_seqlens"], data["src"]["cu_seqlens_k"], data["src"]["page_indices"], data["src"]["nsa_cache_seqlens"], data["src"]["seqlens_expanded"], data["src"]["nsa_cu_seqlens_k"], data["src"]["real_page_table"], data["src"]["flashmla_num_splits"], data["src"]["flashmla_metadata"], dst_fused["cache_seqlens"], dst_fused["cu_seqlens_k"], dst_fused["page_table_1"], dst_fused["nsa_cache_seqlens"], dst_fused["nsa_seqlens_expanded"], dst_fused["nsa_cu_seqlens_k"], dst_fused["real_page_table"], dst_fused["flashmla_num_splits"], dst_fused["flashmla_metadata"], forward_mode, bs, max_len, max_seqlen_k, seqlens_expanded_size, ) # Verify all tensors match for key in dst_ref: if dst_ref[key] is not None: assert torch.equal(dst_ref[key], dst_fused[key]), f"{key} mismatch" # ============================================================================= # Multi-Backend Kernel Tests # ============================================================================= def create_test_metadata_multi( bs: int, max_len: int, seqlens_expanded_size: int, has_real_page_table: bool = False, has_flashmla: bool = False, device: str = "cuda", ): """Create test metadata tensors for multi-backend testing.""" # Source tensors (precomputed metadata) cache_seqlens_src = torch.randint( 1, max_len, (bs,), dtype=torch.int32, device=device ) cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device) cu_seqlens_k_src[1:] = torch.cumsum(cache_seqlens_src, dim=0) page_indices_src = torch.randint( 0, 1000, (bs, max_len), dtype=torch.int32, device=device ) nsa_cache_seqlens_src = torch.randint( 1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device ) nsa_cu_seqlens_k_src = torch.zeros( seqlens_expanded_size + 1, dtype=torch.int32, device=device ) nsa_cu_seqlens_k_src[1:] = torch.cumsum(nsa_cache_seqlens_src, dim=0) # Optional tensors real_page_table_src = None if has_real_page_table: real_page_table_cols = max_len // 2 real_page_table_src = torch.randint( 0, 1000, (bs, real_page_table_cols), dtype=torch.int32, device=device ) flashmla_num_splits_src = None flashmla_metadata_src = None if has_flashmla: flashmla_num_splits_src = torch.randint( 1, 10, (seqlens_expanded_size + 1,), dtype=torch.int32, device=device ) flashmla_metadata_size = 128 flashmla_metadata_src = torch.randint( 0, 100, (flashmla_metadata_size,), dtype=torch.int32, device=device ) # Create destination tensors for 3 backends def create_dst_tensors(): cache_seqlens_dst = torch.zeros(bs, dtype=torch.int32, device=device) cu_seqlens_k_dst = torch.zeros(bs + 1, dtype=torch.int32, device=device) page_table_1_dst = torch.zeros( (bs, max_len + 16), dtype=torch.int32, device=device ) nsa_cache_seqlens_dst = torch.zeros( seqlens_expanded_size, dtype=torch.int32, device=device ) nsa_cu_seqlens_k_dst = torch.zeros( seqlens_expanded_size + 1, dtype=torch.int32, device=device ) real_page_table_dst = None if has_real_page_table: real_page_table_cols = max_len // 2 real_page_table_dst = torch.zeros( (bs, real_page_table_cols + 8), dtype=torch.int32, device=device ) flashmla_num_splits_dst = None flashmla_metadata_dst = None if has_flashmla: flashmla_num_splits_dst = torch.zeros( seqlens_expanded_size + 1, dtype=torch.int32, device=device ) flashmla_metadata_size = 128 flashmla_metadata_dst = torch.zeros( flashmla_metadata_size, dtype=torch.int32, device=device ) return { "cache_seqlens_int32": cache_seqlens_dst, "cu_seqlens_k": cu_seqlens_k_dst, "page_table_1": page_table_1_dst, "nsa_cache_seqlens_int32": nsa_cache_seqlens_dst, "nsa_cu_seqlens_k": nsa_cu_seqlens_k_dst, "real_page_table": real_page_table_dst, "flashmla_num_splits": flashmla_num_splits_dst, "flashmla_metadata": flashmla_metadata_dst, } return { "src": { "cache_seqlens": cache_seqlens_src, "cu_seqlens_k": cu_seqlens_k_src, "page_indices": page_indices_src, "nsa_cache_seqlens": nsa_cache_seqlens_src, "nsa_cu_seqlens_k": nsa_cu_seqlens_k_src, "real_page_table": real_page_table_src, "flashmla_num_splits": flashmla_num_splits_src, "flashmla_metadata": flashmla_metadata_src, }, "dst0": create_dst_tensors(), "dst1": create_dst_tensors(), "dst2": create_dst_tensors(), } def reference_copy_for_loop(src, dst_list, bs, max_len): """Reference implementation: for-loop calling copy for each backend.""" for dst in dst_list: # Simulate what init_forward_metadata_replay_cuda_graph_from_precomputed does dst["cache_seqlens_int32"].copy_(src["cache_seqlens"]) dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) dst["page_table_1"][:, :max_len].copy_(src["page_indices"]) dst["nsa_cache_seqlens_int32"].copy_(src["nsa_cache_seqlens"]) dst["nsa_cu_seqlens_k"][1 : bs + 1].copy_(src["nsa_cu_seqlens_k"][1 : bs + 1]) if src["real_page_table"] is not None: rows, cols = src["real_page_table"].shape dst["real_page_table"][:rows, :cols].copy_(src["real_page_table"]) if src["flashmla_num_splits"] is not None: flashmla_size = bs + 1 dst["flashmla_num_splits"][:flashmla_size].copy_( src["flashmla_num_splits"][:flashmla_size] ) if src["flashmla_metadata"] is not None: dst["flashmla_metadata"].copy_(src["flashmla_metadata"]) def test_fused_metadata_copy_multi_dtype_validation(): """Test that dtype validation rejects non-int32 tensors for multi-backend kernel.""" if not torch.cuda.is_available(): pytest.skip("CUDA not available") from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_multi_cuda bs = 2 max_len = 128 seqlens_expanded_size = bs device = "cuda" # Create source tensors - one with WRONG dtype cache_seqlens_src_wrong = torch.randint( 1, max_len, (bs,), dtype=torch.int64, device=device # Wrong dtype! ) cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device) page_indices_src = torch.randint( 0, 1000, (bs, max_len), dtype=torch.int32, device=device ) nsa_cache_seqlens_src = torch.randint( 1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device ) nsa_cu_seqlens_k_src = torch.zeros( seqlens_expanded_size + 1, dtype=torch.int32, device=device ) # Create destination tensors for 3 backends (all correct dtype) def create_dst(): return { "cache_seqlens": torch.zeros(bs, dtype=torch.int32, device=device), "cu_seqlens_k": torch.zeros(bs + 1, dtype=torch.int32, device=device), "page_table_1": torch.zeros( (bs, max_len + 16), dtype=torch.int32, device=device ), "nsa_cache_seqlens": torch.zeros( seqlens_expanded_size, dtype=torch.int32, device=device ), "nsa_cu_seqlens_k": torch.zeros( seqlens_expanded_size + 1, dtype=torch.int32, device=device ), } dst0 = create_dst() dst1 = create_dst() dst2 = create_dst() # Test: Wrong dtype for source tensor should raise RuntimeError with pytest.raises(RuntimeError, match="must have dtype int32"): fused_metadata_copy_multi_cuda( cache_seqlens_src_wrong, # Wrong dtype: int64 cu_seqlens_k_src, page_indices_src, nsa_cache_seqlens_src, nsa_cu_seqlens_k_src, None, # real_page_table_src None, # flashmla_num_splits_src None, # flashmla_metadata_src # Backend 0 dst0["cache_seqlens"], dst0["cu_seqlens_k"], dst0["page_table_1"], dst0["nsa_cache_seqlens"], dst0["nsa_cu_seqlens_k"], None, None, None, # Backend 1 dst1["cache_seqlens"], dst1["cu_seqlens_k"], dst1["page_table_1"], dst1["nsa_cache_seqlens"], dst1["nsa_cu_seqlens_k"], None, None, None, # Backend 2 dst2["cache_seqlens"], dst2["cu_seqlens_k"], dst2["page_table_1"], dst2["nsa_cache_seqlens"], dst2["nsa_cu_seqlens_k"], None, None, None, # Parameters bs, max_len, seqlens_expanded_size, ) @pytest.mark.parametrize("bs", [1, 2, 4, 8, 16]) @pytest.mark.parametrize("has_real_page_table", [False, True]) @pytest.mark.parametrize("has_flashmla", [False, True]) def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla): """Test fused multi-backend metadata copy kernel against for-loop version.""" if not torch.cuda.is_available(): pytest.skip("CUDA not available") from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_multi_cuda max_len = 128 seqlens_expanded_size = bs # Create test data data = create_test_metadata_multi( bs=bs, max_len=max_len, seqlens_expanded_size=seqlens_expanded_size, has_real_page_table=has_real_page_table, has_flashmla=has_flashmla, ) # Create separate destination tensors for reference (for-loop) and fused kernel dst_ref_0 = { k: v.clone() if v is not None else None for k, v in data["dst0"].items() } dst_ref_1 = { k: v.clone() if v is not None else None for k, v in data["dst1"].items() } dst_ref_2 = { k: v.clone() if v is not None else None for k, v in data["dst2"].items() } dst_fused_0 = { k: v.clone() if v is not None else None for k, v in data["dst0"].items() } dst_fused_1 = { k: v.clone() if v is not None else None for k, v in data["dst1"].items() } dst_fused_2 = { k: v.clone() if v is not None else None for k, v in data["dst2"].items() } # Run reference implementation (for-loop) torch.cuda.synchronize() loop_start = time.perf_counter() reference_copy_for_loop(data["src"], [dst_ref_0, dst_ref_1, dst_ref_2], bs, max_len) torch.cuda.synchronize() loop_end = time.perf_counter() loop_time = loop_end - loop_start # Run fused kernel torch.cuda.synchronize() fused_start = time.perf_counter() fused_metadata_copy_multi_cuda( # Source tensors data["src"]["cache_seqlens"], data["src"]["cu_seqlens_k"], data["src"]["page_indices"], data["src"]["nsa_cache_seqlens"], data["src"]["nsa_cu_seqlens_k"], data["src"]["real_page_table"], data["src"]["flashmla_num_splits"], data["src"]["flashmla_metadata"], # Destination tensors for backend 0 dst_fused_0["cache_seqlens_int32"], dst_fused_0["cu_seqlens_k"], dst_fused_0["page_table_1"], dst_fused_0["nsa_cache_seqlens_int32"], dst_fused_0["nsa_cu_seqlens_k"], dst_fused_0["real_page_table"], dst_fused_0["flashmla_num_splits"], dst_fused_0["flashmla_metadata"], # Destination tensors for backend 1 dst_fused_1["cache_seqlens_int32"], dst_fused_1["cu_seqlens_k"], dst_fused_1["page_table_1"], dst_fused_1["nsa_cache_seqlens_int32"], dst_fused_1["nsa_cu_seqlens_k"], dst_fused_1["real_page_table"], dst_fused_1["flashmla_num_splits"], dst_fused_1["flashmla_metadata"], # Destination tensors for backend 2 dst_fused_2["cache_seqlens_int32"], dst_fused_2["cu_seqlens_k"], dst_fused_2["page_table_1"], dst_fused_2["nsa_cache_seqlens_int32"], dst_fused_2["nsa_cu_seqlens_k"], dst_fused_2["real_page_table"], dst_fused_2["flashmla_num_splits"], dst_fused_2["flashmla_metadata"], # Parameters bs, max_len, seqlens_expanded_size, ) torch.cuda.synchronize() fused_end = time.perf_counter() fused_time = fused_end - fused_start # Compare results for all 3 backends speedup = loop_time / fused_time if fused_time > 0 else 0 print( f"\n[VERIFY] bs={bs}, real_page_table={has_real_page_table}, flashmla={has_flashmla}" ) print( f"[VERIFY] Fused time: {fused_time*1000:.3f}ms, Loop time: {loop_time*1000:.3f}ms, Speedup: {speedup:.2f}x" ) max_diff = 0.0 all_match = True for backend_idx, (dst_ref, dst_fused) in enumerate( [ (dst_ref_0, dst_fused_0), (dst_ref_1, dst_fused_1), (dst_ref_2, dst_fused_2), ] ): for key in [ "cache_seqlens_int32", "cu_seqlens_k", "page_table_1", "nsa_cache_seqlens_int32", "nsa_cu_seqlens_k", ]: if not torch.equal(dst_ref[key], dst_fused[key]): diff = ( (dst_ref[key].float() - dst_fused[key].float()).abs().max().item() ) max_diff = max(max_diff, diff) all_match = False print( f"[ERROR] Backend {backend_idx} {key}: MISMATCH! Max diff: {diff}" ) if has_real_page_table and dst_ref["real_page_table"] is not None: if not torch.equal( dst_ref["real_page_table"], dst_fused["real_page_table"] ): diff = ( ( dst_ref["real_page_table"].float() - dst_fused["real_page_table"].float() ) .abs() .max() .item() ) max_diff = max(max_diff, diff) all_match = False print( f"[ERROR] Backend {backend_idx} real_page_table: MISMATCH! Max diff: {diff}" ) if has_flashmla: if dst_ref["flashmla_num_splits"] is not None and not torch.equal( dst_ref["flashmla_num_splits"], dst_fused["flashmla_num_splits"] ): diff = ( ( dst_ref["flashmla_num_splits"].float() - dst_fused["flashmla_num_splits"].float() ) .abs() .max() .item() ) max_diff = max(max_diff, diff) all_match = False print( f"[ERROR] Backend {backend_idx} flashmla_num_splits: MISMATCH! Max diff: {diff}" ) if dst_ref["flashmla_metadata"] is not None and not torch.equal( dst_ref["flashmla_metadata"], dst_fused["flashmla_metadata"] ): diff = ( ( dst_ref["flashmla_metadata"].float() - dst_fused["flashmla_metadata"].float() ) .abs() .max() .item() ) max_diff = max(max_diff, diff) all_match = False print( f"[ERROR] Backend {backend_idx} flashmla_metadata: MISMATCH! Max diff: {diff}" ) if not all_match: error_msg = ( f"Fused metadata copy verification FAILED! " f"Maximum difference: {max_diff}. " f"The fused kernel produces different results than the for-loop version." ) print(f"[ERROR] {error_msg}") raise AssertionError(error_msg) print(f"[VERIFY] Verification PASSED - all tensors match!") @pytest.mark.parametrize("bs", [32, 64]) def test_fused_metadata_copy_multi_large_batch(bs): """Test with larger batch sizes and timing comparison.""" if not torch.cuda.is_available(): pytest.skip("CUDA not available") from sglang.jit_kernel.fused_metadata_copy import fused_metadata_copy_multi_cuda max_len = 128 seqlens_expanded_size = bs data = create_test_metadata_multi( bs=bs, max_len=max_len, seqlens_expanded_size=seqlens_expanded_size, has_real_page_table=True, has_flashmla=True, ) dst_ref_0 = { k: v.clone() if v is not None else None for k, v in data["dst0"].items() } dst_ref_1 = { k: v.clone() if v is not None else None for k, v in data["dst1"].items() } dst_ref_2 = { k: v.clone() if v is not None else None for k, v in data["dst2"].items() } dst_fused_0 = { k: v.clone() if v is not None else None for k, v in data["dst0"].items() } dst_fused_1 = { k: v.clone() if v is not None else None for k, v in data["dst1"].items() } dst_fused_2 = { k: v.clone() if v is not None else None for k, v in data["dst2"].items() } # Warmup for _ in range(5): reference_copy_for_loop( data["src"], [dst_ref_0, dst_ref_1, dst_ref_2], bs, max_len ) fused_metadata_copy_multi_cuda( data["src"]["cache_seqlens"], data["src"]["cu_seqlens_k"], data["src"]["page_indices"], data["src"]["nsa_cache_seqlens"], data["src"]["nsa_cu_seqlens_k"], data["src"]["real_page_table"], data["src"]["flashmla_num_splits"], data["src"]["flashmla_metadata"], dst_fused_0["cache_seqlens_int32"], dst_fused_0["cu_seqlens_k"], dst_fused_0["page_table_1"], dst_fused_0["nsa_cache_seqlens_int32"], dst_fused_0["nsa_cu_seqlens_k"], dst_fused_0["real_page_table"], dst_fused_0["flashmla_num_splits"], dst_fused_0["flashmla_metadata"], dst_fused_1["cache_seqlens_int32"], dst_fused_1["cu_seqlens_k"], dst_fused_1["page_table_1"], dst_fused_1["nsa_cache_seqlens_int32"], dst_fused_1["nsa_cu_seqlens_k"], dst_fused_1["real_page_table"], dst_fused_1["flashmla_num_splits"], dst_fused_1["flashmla_metadata"], dst_fused_2["cache_seqlens_int32"], dst_fused_2["cu_seqlens_k"], dst_fused_2["page_table_1"], dst_fused_2["nsa_cache_seqlens_int32"], dst_fused_2["nsa_cu_seqlens_k"], dst_fused_2["real_page_table"], dst_fused_2["flashmla_num_splits"], dst_fused_2["flashmla_metadata"], bs, max_len, seqlens_expanded_size, ) torch.cuda.synchronize() # Actual timing torch.cuda.synchronize() loop_start = time.perf_counter() reference_copy_for_loop(data["src"], [dst_ref_0, dst_ref_1, dst_ref_2], bs, max_len) torch.cuda.synchronize() loop_time = time.perf_counter() - loop_start torch.cuda.synchronize() fused_start = time.perf_counter() fused_metadata_copy_multi_cuda( data["src"]["cache_seqlens"], data["src"]["cu_seqlens_k"], data["src"]["page_indices"], data["src"]["nsa_cache_seqlens"], data["src"]["nsa_cu_seqlens_k"], data["src"]["real_page_table"], data["src"]["flashmla_num_splits"], data["src"]["flashmla_metadata"], dst_fused_0["cache_seqlens_int32"], dst_fused_0["cu_seqlens_k"], dst_fused_0["page_table_1"], dst_fused_0["nsa_cache_seqlens_int32"], dst_fused_0["nsa_cu_seqlens_k"], dst_fused_0["real_page_table"], dst_fused_0["flashmla_num_splits"], dst_fused_0["flashmla_metadata"], dst_fused_1["cache_seqlens_int32"], dst_fused_1["cu_seqlens_k"], dst_fused_1["page_table_1"], dst_fused_1["nsa_cache_seqlens_int32"], dst_fused_1["nsa_cu_seqlens_k"], dst_fused_1["real_page_table"], dst_fused_1["flashmla_num_splits"], dst_fused_1["flashmla_metadata"], dst_fused_2["cache_seqlens_int32"], dst_fused_2["cu_seqlens_k"], dst_fused_2["page_table_1"], dst_fused_2["nsa_cache_seqlens_int32"], dst_fused_2["nsa_cu_seqlens_k"], dst_fused_2["real_page_table"], dst_fused_2["flashmla_num_splits"], dst_fused_2["flashmla_metadata"], bs, max_len, seqlens_expanded_size, ) torch.cuda.synchronize() fused_time = time.perf_counter() - fused_start speedup = loop_time / fused_time if fused_time > 0 else 0 print( f"\n[PERF] Large batch (bs={bs}): Fused={fused_time*1000:.3f}ms, Loop={loop_time*1000:.3f}ms, Speedup={speedup:.2f}x" ) # Verify correctness for backend_idx, (dst_ref, dst_fused) in enumerate( [ (dst_ref_0, dst_fused_0), (dst_ref_1, dst_fused_1), (dst_ref_2, dst_fused_2), ] ): for key in dst_ref: if dst_ref[key] is not None and dst_fused[key] is not None: assert torch.equal( dst_ref[key], dst_fused[key] ), f"Backend {backend_idx} {key} mismatch" if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py ================================================ from typing import Optional, Tuple import pytest import torch from einops import rearrange from torch import Tensor from sglang.jit_kernel.diffusion.cutedsl.scale_residual_norm_scale_shift import ( fused_norm_scale_shift, fused_scale_residual_norm_scale_shift, ) DEVICE = "cuda" SHAPE_MAP = { "1": lambda B, S, F, D: (1,), "D": lambda B, S, F, D: (D,), "1D": lambda B, S, F, D: (1, D), "BD": lambda B, S, F, D: (B, D), "11D": lambda B, S, F, D: (1, 1, D), "B1D": lambda B, S, F, D: (B, 1, D), "1SD": lambda B, S, F, D: (1, S, D), "BSD": lambda B, S, F, D: (B, S, D), "BF1D": lambda B, S, F, D: (B, F, 1, D), } SHAPES = [ # (B, S, F, D) (1, 115200, 1, 3072), # Hunyuan (1, 32760, 1, 1536), # Wan (1, 6, 1, 3072), # Qwen (1, 1024, 8, 3072), (4, 512, 16, 3072), ] DTYPES = [torch.float16, torch.bfloat16, torch.float32] NORM_TYPES = ["layer", "rms"] AFFINE_MODES = ["D", "NAT"] INDEX_MODES = ["BSD", "1", "1SD", "BD", "B1D", "D", "1D", "11D", "BF1D"] def _tol(dtype: torch.dtype): return 1e-5 if dtype == torch.float32 else 5e-2 @pytest.fixture(autouse=True) def cuda_setup(): if not torch.cuda.is_available(): pytest.skip("CUDA required") torch.cuda.manual_seed(0) def _apply_scale_shift(y: Tensor, scale: Tensor, shift: Tensor) -> Tensor: if scale.ndim == 4: num_frame = scale.shape[1] return rearrange( rearrange(y, "b (f l) d -> b f l d", f=num_frame) * (1 + scale) + shift, "b f l d -> b (f l) d", ) else: scale = rearrange(scale, "b d -> b 1 d") if scale.ndim == 2 else scale shift = rearrange(shift, "b d -> b 1 d") if shift.ndim == 2 else shift return y * (1 + scale) + shift def fused_norm_scale_shift_ref( x: Tensor, weight: Optional[Tensor], bias: Optional[Tensor], scale: Tensor, shift: Tensor, norm_type: str, eps: float, ) -> Tensor: original_dtype = x.dtype x, weight, bias, scale, shift = ( v.float() if v is not None else v for v in [x, weight, bias, scale, shift] ) if norm_type == "layer": norm = torch.layer_norm(x, x.shape[-1:], eps=eps, weight=weight, bias=bias) else: norm = torch.rms_norm(x, x.shape[-1:], eps=eps, weight=weight) return _apply_scale_shift(norm, scale, shift).to(original_dtype) def fused_scale_residual_norm_scale_shift_ref( residual: Tensor, x: Tensor, gate: Optional[Tensor] | int, weight: Optional[Tensor], bias: Optional[Tensor], scale: Tensor, shift: Tensor, norm_type: str, eps: float, ): original_dtype = x.dtype residual, x, gate, weight, bias, scale, shift = ( v.float() if isinstance(v, Tensor) else v for v in [residual, x, gate, weight, bias, scale, shift] ) if isinstance(gate, int): x = residual + gate * x else: if gate.ndim == 4: num_frame = gate.shape[1] x_fld = rearrange(x, "b (f l) d -> b f l d", f=num_frame) x = residual + rearrange(x_fld * gate, "b f l d -> b (f l) d") else: gate = rearrange(gate, "b d -> b 1 d") if gate.ndim == 2 else gate x = residual + gate * x if norm_type == "layer": norm = torch.layer_norm(x, x.shape[-1:], eps=eps, weight=weight, bias=bias) else: norm = torch.rms_norm(x, x.shape[-1:], eps=eps, weight=weight) y_ref = _apply_scale_shift(norm, scale, shift) return y_ref.to(original_dtype), x.to(original_dtype) def _make_tensor(index_mode: str, shape: Tuple, dtype: torch.dtype): if index_mode == "NAT": return None return torch.randn(*SHAPE_MAP[index_mode](*shape), device=DEVICE, dtype=dtype) @torch.no_grad() def run_norm_scale_shift( shape=SHAPES[0], dtype=DTYPES[0], affine_dtype=DTYPES[0], scale_dtype=DTYPES[0], shift_dtype=DTYPES[0], norm_type=NORM_TYPES[0], affine_mode=AFFINE_MODES[0], scale_mode="BSD", shift_mode="BSD", eps=1e-5, ): x = _make_tensor("BSD", shape, dtype) weight = _make_tensor(affine_mode, shape, affine_dtype) bias = _make_tensor(affine_mode, shape, affine_dtype) scale = _make_tensor(scale_mode, shape, scale_dtype) shift = _make_tensor(shift_mode, shape, shift_dtype) y_dev = fused_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps) y_ref = fused_norm_scale_shift_ref(x, weight, bias, scale, shift, norm_type, eps) torch.testing.assert_close(y_dev, y_ref, atol=_tol(dtype), rtol=_tol(dtype)) @torch.no_grad() def run_scale_resi_norm_scale_shift( shape=SHAPES[0], dtype=DTYPES[0], affine_dtype=DTYPES[0], scale_dtype=DTYPES[0], shift_dtype=DTYPES[0], norm_type=NORM_TYPES[0], affine_mode=AFFINE_MODES[0], gate_mode="B1D", scale_mode="BSD", shift_mode="BSD", eps=1e-5, ): residual = _make_tensor("BSD", shape, dtype) x = _make_tensor("BSD", shape, dtype) gate = _make_tensor(gate_mode, shape, dtype) weight = _make_tensor(affine_mode, shape, affine_dtype) bias = _make_tensor(affine_mode, shape, affine_dtype) scale = _make_tensor(scale_mode, shape, scale_dtype) shift = _make_tensor(shift_mode, shape, shift_dtype) y_dev, res_dev = fused_scale_residual_norm_scale_shift( residual, x, gate, weight, bias, scale, shift, norm_type, eps ) y_ref, res_ref = fused_scale_residual_norm_scale_shift_ref( residual, x, gate, weight, bias, scale, shift, norm_type, eps ) torch.testing.assert_close(y_dev, y_ref, atol=_tol(dtype), rtol=_tol(dtype)) torch.testing.assert_close(res_dev, res_ref, atol=_tol(dtype), rtol=_tol(dtype)) @pytest.mark.parametrize("norm_type", NORM_TYPES) class TestFusedNormScaleShift: @pytest.mark.parametrize("shape", SHAPES) @pytest.mark.parametrize("dtype", DTYPES) def test_shape_dtype(self, shape, dtype, norm_type): run_norm_scale_shift(shape=shape, dtype=dtype, norm_type=norm_type) @pytest.mark.parametrize("dtype", DTYPES) def test_dtype_0(self, dtype, norm_type): run_norm_scale_shift(affine_dtype=dtype, norm_type=norm_type) @pytest.mark.parametrize("dtype", DTYPES) def test_dtype_1(self, dtype, norm_type): run_norm_scale_shift(scale_dtype=dtype, shift_dtype=dtype, norm_type=norm_type) @pytest.mark.parametrize("affine_mode", AFFINE_MODES) def test_normtype_affine(self, affine_mode, norm_type): run_norm_scale_shift(affine_mode=affine_mode, norm_type=norm_type) @pytest.mark.parametrize("index_mode", INDEX_MODES) def test_index_mode(self, index_mode, norm_type): run_norm_scale_shift( scale_mode=index_mode, shift_mode=index_mode, norm_type=norm_type ) @pytest.mark.parametrize("norm_type", NORM_TYPES) class TestFusedScaleResidualNormScaleShift: @pytest.mark.parametrize("shape", SHAPES) @pytest.mark.parametrize("dtype", DTYPES) def test_shape_dtype(self, shape, dtype, norm_type): run_scale_resi_norm_scale_shift(shape=shape, dtype=dtype, norm_type=norm_type) @pytest.mark.parametrize("dtype", DTYPES) def test_dtype_0(self, dtype, norm_type): run_scale_resi_norm_scale_shift(affine_dtype=dtype, norm_type=norm_type) @pytest.mark.parametrize("dtype", DTYPES) def test_dtype_1(self, dtype, norm_type): run_scale_resi_norm_scale_shift( scale_dtype=dtype, shift_dtype=dtype, norm_type=norm_type ) @pytest.mark.parametrize("affine_mode", AFFINE_MODES) def test_normtype_affine(self, affine_mode, norm_type): run_scale_resi_norm_scale_shift(affine_mode=affine_mode, norm_type=norm_type) @pytest.mark.parametrize("index_mode", INDEX_MODES) def test_scale_shift_index_mode(self, index_mode, norm_type): run_scale_resi_norm_scale_shift( scale_mode=index_mode, shift_mode=index_mode, norm_type=norm_type ) @pytest.mark.parametrize("index_mode", INDEX_MODES) def test_gate_index_mode(self, index_mode, norm_type): run_scale_resi_norm_scale_shift(gate_mode=index_mode, norm_type=norm_type) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_fused_store_index_cache.py ================================================ """ Test for fused_store_index_k_cache kernel. Design Notes: 1. torch.cuda.synchronize() needed after TVM FFI kernel call. 2. _split_buffer used buf[:, :vb].reshape(-1) which COPIES data for non-contiguous slices → reference buffer stayed all-zeros. Fix: use flat byte-offset indexing. 3. act_quant may use a different quantization scheme → generous tolerance. 4. FP8 E4M3 1-ULP rounding differences between CUDA hardware cast (__nv_fp8_e4m3) and PyTorch .to(float8_e4m3fn) at tie-break points. Adjacent FP8 representable values at the high end differ by up to 32 in float space (e.g. 288, 320, 352, ..., 448). Need to compare dequantized values with FP8-appropriate tolerance. """ from __future__ import annotations from typing import Optional, Tuple import pytest import torch try: from sglang.jit_kernel.fused_store_index_cache import ( can_use_nsa_fused_store, fused_store_index_k_cache, ) HAS_FUSED = True except ImportError: HAS_FUSED = False try: from sglang.srt.utils import is_hip _is_hip = is_hip() except ImportError: _is_hip = False try: from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz _is_fp8_fnuz = is_fp8_fnuz() except ImportError: _is_fp8_fnuz = False PAGE_SIZE = 64 HEAD_DIM = 128 FP8_E4M3_MAX = 448.0 FP8_DTYPE = torch.float8_e4m3fn BYTES_PER_TOKEN = 128 + 4 # 128 fp8 bytes + 4 scale bytes BYTES_PER_PAGE = PAGE_SIZE * BYTES_PER_TOKEN def _skip_if_unavailable(page_size: int = PAGE_SIZE): if not torch.cuda.is_available(): pytest.skip("CUDA required") if _is_hip: pytest.skip("Fused store kernel is CUDA-specific") if _is_fp8_fnuz: pytest.skip("Fused store path disabled for FP8 FNUZ") if not hasattr(torch, "float8_e4m3fn"): pytest.skip("torch.float8_e4m3fn not available") if not HAS_FUSED: pytest.skip("fused_store_index_cache not importable") if not can_use_nsa_fused_store(torch.bfloat16, torch.int64, page_size): pytest.skip("JIT kernel unavailable / failed to compile") def _num_pages(loc: torch.Tensor, page_size: int, extra: int = 1) -> int: return int(loc.max().item()) // page_size + 1 + extra def _make_buffer(num_pages: int, page_size: int = PAGE_SIZE) -> torch.Tensor: return torch.zeros( (num_pages, page_size * BYTES_PER_TOKEN), dtype=torch.uint8, device="cuda", ) def _read_token_from_buffer( buf: torch.Tensor, token_idx: int, page_size: int = PAGE_SIZE, ) -> Tuple[torch.Tensor, float]: """ Read a single token's fp8 values and scale from the paged buffer using flat byte offsets. """ page = token_idx // page_size offset = token_idx % page_size page_bytes = page_size * BYTES_PER_TOKEN buf_flat = buf.reshape(-1) val_start = page * page_bytes + offset * 128 fp8_bytes = buf_flat[val_start : val_start + 128] fp8_vals = fp8_bytes.view(FP8_DTYPE).float() scale_start = page * page_bytes + 128 * page_size + offset * 4 scale_bytes = buf_flat[scale_start : scale_start + 4] scale = scale_bytes.view(torch.float32).item() return fp8_vals, scale def _write_token_to_buffer( buf: torch.Tensor, token_idx: int, fp8_data: torch.Tensor, scale: float, page_size: int = PAGE_SIZE, ) -> None: """ Write a single token's fp8 values and scale into the paged buffer using flat byte offsets on buf.reshape(-1) (which is a true view since buf is contiguous). """ page = token_idx // page_size offset = token_idx % page_size page_bytes = page_size * BYTES_PER_TOKEN buf_flat = buf.reshape(-1) val_start = page * page_bytes + offset * 128 buf_flat[val_start : val_start + 128] = fp8_data.view(torch.uint8) scale_start = page * page_bytes + 128 * page_size + offset * 4 scale_t = torch.tensor([scale], dtype=torch.float32, device=buf.device) buf_flat[scale_start : scale_start + 4] = scale_t.view(torch.uint8) def _gather_tokens( buf: torch.Tensor, loc: torch.Tensor, page_size: int = PAGE_SIZE, ) -> Tuple[torch.Tensor, torch.Tensor]: N = loc.shape[0] fp8_f32 = torch.empty((N, HEAD_DIM), dtype=torch.float32, device=buf.device) scales = torch.empty((N,), dtype=torch.float32, device=buf.device) for i in range(N): idx = int(loc[i].item()) vals, s = _read_token_from_buffer(buf, idx, page_size) fp8_f32[i] = vals scales[i] = s return fp8_f32, scales # Reference kernel def _reference_quantize_and_store( key_bf16: torch.Tensor, loc: torch.Tensor, num_pages: int, page_size: int = PAGE_SIZE, ) -> torch.Tensor: """ Reference kernel of the fused kernel's quantization: abs_max = max(|row|) scale = max(1e-4, abs_max) / 448 fp8_val = clip(val / scale, -448, 448) -> cast to fp8 """ N = key_bf16.shape[0] key_f32 = key_bf16.float() buf = _make_buffer(num_pages, page_size) for i in range(N): row = key_f32[i] abs_max = row.abs().max().item() scale = max(1e-4, abs_max) / FP8_E4M3_MAX inv_scale = 1.0 / scale quantized = (row * inv_scale).clamp(-FP8_E4M3_MAX, FP8_E4M3_MAX) quantized_fp8 = quantized.to(FP8_DTYPE) idx = int(loc[i].item()) _write_token_to_buffer(buf, idx, quantized_fp8, scale, page_size) return buf def _import_act_quant(): try: from sglang.srt.layers.attention.nsa.triton_kernel import act_quant return act_quant except Exception: return None def _ref_store_via_act_quant( key_bf16: torch.Tensor, loc: torch.Tensor, num_pages: int, page_size: int = PAGE_SIZE, block_size: int = 128, scale_fmt: Optional[str] = None, ) -> Optional[torch.Tensor]: act_quant = _import_act_quant() if act_quant is None: return None try: k_fp8, k_scale = act_quant(key_bf16, block_size, scale_fmt) except TypeError: k_fp8, k_scale = act_quant(key_bf16, block_size) if k_fp8.dim() == 3 and k_fp8.shape[1] == 1: k_fp8 = k_fp8.squeeze(1) if k_scale is not None and k_scale.dim() == 3 and k_scale.shape[1] == 1: k_scale = k_scale.squeeze(1) k_scale = k_scale.view(-1).float() buf = _make_buffer(num_pages, page_size) N = key_bf16.shape[0] for i in range(N): idx = int(loc[i].item()) _write_token_to_buffer( buf, idx, k_fp8[i].to(FP8_DTYPE), k_scale[i].item(), page_size ) return buf # TEST 1: Fused kernel vs. its own algorithm (pure-Python reference) # # NOTE on FP8 rounding: # CUDA hardware fp8 cast (__nv_fp8_e4m3) and PyTorch .to(float8_e4m3fn) # may round differently at tie-break points. This causes up to 1-ULP # differences in the FP8 codes. In FP8 E4M3, adjacent representable # values at the high end differ by up to 32 in float space (e.g. # 288 vs 320). After dequantization (fp8_float * scale), the error # from 1-ULP is: scale * ulp ≈ (abs_max/448) * 32 ≈ 0.07 * abs_max. # For randn inputs (abs_max ≈ 3-4), this is about 0.2-0.3. # # We therefore compare dequantized values with tolerances that # accommodate 1-ULP FP8 rounding, NOT byte-exact fp8 codes. @pytest.mark.parametrize( "num_tokens,base_index", [(1, 0), (32, 0), (64, 0), (128, 64), (257, 65), (512, 0)], ) def test_fused_kernel_matches_own_algorithm(num_tokens: int, base_index: int): """Compare fused CUDA kernel against a pure-Python implementation of the *same* quantization formula.""" _skip_if_unavailable() device = torch.device("cuda") key = torch.randn((num_tokens, HEAD_DIM), device=device, dtype=torch.bfloat16) loc = ( base_index + torch.randperm(num_tokens, device=device, dtype=torch.int64) ).contiguous() num_pages = _num_pages(loc, PAGE_SIZE) # Reference kernel ref_buf = _reference_quantize_and_store(key, loc, num_pages) # Fused kernel out_buf = _make_buffer(num_pages) fused_store_index_k_cache(key, out_buf, loc, page_size=PAGE_SIZE) torch.cuda.synchronize() out_f, out_s = _gather_tokens(out_buf, loc) ref_f, ref_s = _gather_tokens(ref_buf, loc) # 1) Scales must match tightly (same f32 formula, no rounding ambiguity) torch.testing.assert_close(out_s, ref_s, rtol=1e-5, atol=1e-7) # 2) Most FP8 codes should match; allow rare 1-ULP differences. # 1-ULP at FP8 E4M3 high end = 32 in float space. mismatch = out_f != ref_f mismatch_frac = mismatch.float().mean().item() assert mismatch_frac < 0.01, ( f"Too many FP8 code mismatches: {mismatch_frac:.2%} " f"(expected < 1% from rounding tie-breaks)" ) # 3) Where codes differ, the difference should be exactly 1 ULP. # In FP8 E4M3: if the float-cast value is V, the adjacent value # differs by ~V * 0.1 (relative) at most. if mismatch.any(): diff = (out_f[mismatch] - ref_f[mismatch]).abs() rel_diff = diff / ref_f[mismatch].abs().clamp(min=1e-6) # 1-ULP relative difference for E4M3 is at most ~12.5% (2^-3) assert rel_diff.max().item() <= 0.15, ( f"FP8 code difference exceeds 1-ULP: max relative diff = " f"{rel_diff.max().item():.4f}" ) # 4) Dequantized values should be close. # Max error from 1-ULP: scale * fp8_ulp ≈ (abs_max/448) * 32 # For randn abs_max ≈ 3-4: max_err ≈ 0.21 - 0.29 out_deq = out_f * out_s.unsqueeze(-1) ref_deq = ref_f * ref_s.unsqueeze(-1) torch.testing.assert_close(out_deq, ref_deq, rtol=0.15, atol=0.5) # TEST 2: Cross-check against act_quant @pytest.mark.parametrize("scale_fmt", [None, "fp32"]) def test_fused_kernel_vs_act_quant_semantic(scale_fmt: Optional[str]): """Both fused kernel and act_quant should approximately reconstruct the original bf16 values.""" _skip_if_unavailable() device = torch.device("cuda") num_tokens = 257 base_index = 65 key = torch.randn((num_tokens, HEAD_DIM), device=device, dtype=torch.bfloat16) loc = ( base_index + torch.randperm(num_tokens, device=device, dtype=torch.int64) ).contiguous() num_pages = _num_pages(loc, PAGE_SIZE) ref_buf = _ref_store_via_act_quant(key, loc, num_pages, scale_fmt=scale_fmt) if ref_buf is None: pytest.skip("act_quant not available") out_buf = _make_buffer(num_pages) fused_store_index_k_cache(key, out_buf, loc, page_size=PAGE_SIZE) torch.cuda.synchronize() out_f, out_s = _gather_tokens(out_buf, loc) ref_f, ref_s = _gather_tokens(ref_buf, loc) out_deq = out_f * out_s.unsqueeze(-1) ref_deq = ref_f * ref_s.unsqueeze(-1) orig_f32 = key.float() # Fused kernel should reconstruct original within FP8 precision torch.testing.assert_close( out_deq, orig_f32, rtol=0.15, atol=5e-2, msg="Fused kernel dequantized values don't approximate original", ) # act_quant may use a very different scale policy. try: torch.testing.assert_close( ref_deq, orig_f32, rtol=0.25, atol=0.5, msg="act_quant dequantized values don't approximate original", ) except AssertionError: nonzero_frac = (ref_deq.abs() > 1e-6).float().mean().item() if nonzero_frac < 0.5: pytest.fail( f"act_quant output looks mostly zero ({nonzero_frac:.1%} nonzero)." ) else: pytest.skip( f"act_quant uses a very different quantization scheme " f"(scale_fmt={scale_fmt}). Fused kernel validated independently." ) torch.testing.assert_close( out_deq, ref_deq, rtol=0.3, atol=0.5, msg="Fused and act_quant dequantized values diverge too much", ) # TEST 3: Roundtrip reconstruction @pytest.mark.parametrize("num_tokens", [1, 64, 257]) def test_roundtrip_reconstruction(num_tokens: int): _skip_if_unavailable() device = torch.device("cuda") key = torch.randn((num_tokens, HEAD_DIM), device=device, dtype=torch.bfloat16) loc = torch.arange(num_tokens, device=device, dtype=torch.int64) num_pages = _num_pages(loc, PAGE_SIZE) buf = _make_buffer(num_pages) fused_store_index_k_cache(key, buf, loc, page_size=PAGE_SIZE) torch.cuda.synchronize() fp8_f32, scales = _gather_tokens(buf, loc) reconstructed = fp8_f32 * scales.unsqueeze(-1) original = key.float() torch.testing.assert_close(reconstructed, original, rtol=0.15, atol=5e-2) per_row_energy = reconstructed.abs().sum(dim=-1) orig_energy = original.abs().sum(dim=-1) mask = orig_energy > 0.1 assert ( per_row_energy[mask] > 0.01 ).all(), "Some tokens have zero reconstruction — kernel may not be writing output" # TEST 4: Boundary conditions def test_single_token(): _skip_if_unavailable() device = torch.device("cuda") key = torch.randn((1, HEAD_DIM), device=device, dtype=torch.bfloat16) loc = torch.tensor([0], device=device, dtype=torch.int64) buf = _make_buffer(1) fused_store_index_k_cache(key, buf, loc, page_size=PAGE_SIZE) torch.cuda.synchronize() fp8_f32, scales = _gather_tokens(buf, loc) reconstructed = fp8_f32 * scales.unsqueeze(-1) torch.testing.assert_close(reconstructed, key.float(), rtol=0.15, atol=5e-2) # TEST 5: Zero input conditions def test_zero_input(): _skip_if_unavailable() device = torch.device("cuda") key = torch.zeros((4, HEAD_DIM), device=device, dtype=torch.bfloat16) loc = torch.arange(4, device=device, dtype=torch.int64) buf = _make_buffer(1) fused_store_index_k_cache(key, buf, loc, page_size=PAGE_SIZE) torch.cuda.synchronize() fp8_f32, scales = _gather_tokens(buf, loc) expected_scale = 1e-4 / FP8_E4M3_MAX torch.testing.assert_close( scales, torch.full_like(scales, expected_scale), rtol=1e-5, atol=1e-10, ) assert (fp8_f32 == 0).all() # TEST 6: Sanity check — verify reference itself writes non-zero data def test_reference_writes_nonzero(): _skip_if_unavailable() device = torch.device("cuda") key = torch.randn((8, HEAD_DIM), device=device, dtype=torch.bfloat16) loc = torch.arange(8, device=device, dtype=torch.int64) buf = _reference_quantize_and_store(key, loc, num_pages=1) fp8_f32, scales = _gather_tokens(buf, loc) deq = fp8_f32 * scales.unsqueeze(-1) assert deq.abs().sum().item() > 0, "Reference buffer is all zeros — error!" torch.testing.assert_close(deq, key.float(), rtol=0.15, atol=5e-2) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_fused_verify_triton_gdn.py ================================================ """Tests for fused sigmoid gating delta rule MTP kernel (GDN target_verify). Compares the fused kernel `fused_sigmoid_gating_delta_rule_update` against the reference two-step implementation: 1. g, beta = fused_gdn_gating(A_log, a, b, dt_bias) 2. o = fused_recurrent_gated_delta_rule_update(q, k, v, g, beta, ...) """ import pytest import torch try: from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating from sglang.srt.layers.attention.fla.fused_recurrent import ( fused_recurrent_gated_delta_rule_update, ) from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) KERNELS_AVAILABLE = True except ImportError: KERNELS_AVAILABLE = False def _make_tensors(N, T, H, HV, K, V, device="cuda", seed=2025): """Create input tensors for GDN target_verify.""" torch.manual_seed(seed) A_log = torch.randn(HV, dtype=torch.float32, device=device) dt_bias = torch.randn(HV, dtype=torch.bfloat16, device=device) a = torch.randn(1, N * T, HV, dtype=torch.bfloat16, device=device) b = torch.randn(1, N * T, HV, dtype=torch.bfloat16, device=device) q = torch.randn(1, N * T, H, K, dtype=torch.bfloat16, device=device) k = torch.randn(1, N * T, H, K, dtype=torch.bfloat16, device=device) v = torch.randn(1, N * T, HV, V, dtype=torch.bfloat16, device=device) indices = torch.arange(N, dtype=torch.int32, device=device) initial_state = torch.randn(N, HV, K, V, dtype=torch.float, device=device) cu_seqlens = torch.arange(0, N * T + 1, T, dtype=torch.int32, device=device) return A_log, dt_bias, a, b, q, k, v, initial_state, indices, cu_seqlens def run_reference( A_log, dt_bias, q, k, v, a, b, initial_state_source, initial_state_indices, cu_seqlens, disable_state_update=True, intermediate_states_buffer=None, intermediate_state_indices=None, cache_steps=None, retrieve_parent_token=None, ): """Reference: fused_gdn_gating + fused_recurrent_gated_delta_rule_update.""" # fused_gdn_gating expects 2D [seq_len, HV] a_2d = a.view(-1, a.shape[-1]) b_2d = b.view(-1, b.shape[-1]) g, beta = fused_gdn_gating(A_log, a_2d, b_2d, dt_bias) # fused_recurrent expects 3D [B, T, HV] g = g.view(a.shape) beta = beta.view(b.shape) # fused_recurrent requires intermediate_state_indices when cu_seqlens is used if cu_seqlens is not None and intermediate_state_indices is None: N = len(cu_seqlens) - 1 intermediate_state_indices = torch.arange(N, dtype=torch.int32, device=q.device) return fused_recurrent_gated_delta_rule_update( q=q, k=k, v=v, g=g, beta=beta, initial_state_source=initial_state_source, initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, use_qk_l2norm_in_kernel=True, disable_state_update=disable_state_update, intermediate_states_buffer=intermediate_states_buffer, intermediate_state_indices=intermediate_state_indices, cache_steps=cache_steps, retrieve_parent_token=retrieve_parent_token, ) def run_fused_mtp( A_log, dt_bias, q, k, v, a, b, initial_state_source, initial_state_indices, cu_seqlens, disable_state_update=True, intermediate_states_buffer=None, intermediate_state_indices=None, cache_steps=None, retrieve_parent_token=None, ): """Fused: fused_sigmoid_gating_delta_rule_update.""" return fused_sigmoid_gating_delta_rule_update( A_log=A_log, dt_bias=dt_bias, q=q, k=k, v=v, a=a, b=b, initial_state_source=initial_state_source, initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, use_qk_l2norm_in_kernel=True, softplus_beta=1.0, softplus_threshold=20.0, is_kda=False, disable_state_update=disable_state_update, intermediate_states_buffer=intermediate_states_buffer, intermediate_state_indices=intermediate_state_indices, cache_steps=cache_steps, retrieve_parent_token=retrieve_parent_token, ) @pytest.mark.skipif(not KERNELS_AVAILABLE, reason="Kernel not available") @pytest.mark.parametrize("N", [1, 8, 16]) @pytest.mark.parametrize("T", [1, 4, 8]) def test_fused_gdn_mtp_precision(N: int, T: int): """Compare fused MTP output against reference.""" H, HV, K, V = 16, 32, 128, 128 A_log, dt_bias, a, b, q, k, v, state, indices, cu_seqlens = _make_tensors( N, T, H, HV, K, V ) state_ref = state.clone() state_fused = state.clone() out_ref = run_reference( A_log, dt_bias, q, k, v, a, b, state_ref, indices, cu_seqlens, disable_state_update=True, ) out_fused = run_fused_mtp( A_log, dt_bias, q, k, v, a, b, state_fused, indices, cu_seqlens, disable_state_update=True, ) torch.testing.assert_close(out_ref, out_fused, rtol=1e-2, atol=1e-2) @pytest.mark.skipif(not KERNELS_AVAILABLE, reason="Kernels not available") @pytest.mark.parametrize("N", [1, 16, 128]) def test_mtp_single_step_decode(N: int): """Verify MTP kernel matches reference for T=1 (decode scenario).""" T = 1 H, HV, K, V = 16, 32, 128, 128 A_log, dt_bias, a, b, q, k, v, state, indices, cu_seqlens = _make_tensors( N, T, H, HV, K, V ) state_ref = state.clone() state_fused = state.clone() out_ref = run_reference( A_log, dt_bias, q, k, v, a, b, state_ref, indices, cu_seqlens, disable_state_update=False, ) out_fused = run_fused_mtp( A_log, dt_bias, q, k, v, a, b, state_fused, indices, cu_seqlens, disable_state_update=False, ) torch.testing.assert_close(out_ref, out_fused, rtol=1e-2, atol=1e-2) # Also verify states match after update state_diff = (state_ref.float() - state_fused.float()).abs() state_max_diff = state_diff.max().item() state_fail_rate = (state_diff > 0.1).float().mean().item() * 100 print( f" single_step state N={N}: max_diff={state_max_diff:.2e}, " f"fail_rate={state_fail_rate:.2f}%" ) assert state_fail_rate < 0.01, f"State mismatch: fail_rate={state_fail_rate:.2f}%" if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_gptq_marlin.py ================================================ import pytest import torch from sgl_kernel.scalar_type import scalar_types from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace from sglang.test.test_marlin_utils import awq_marlin_quantize, marlin_quantize MNK_FACTORS = [ (1, 1, 1), (1, 4, 8), (13, 17, 67), (257, 13, 11), ] @pytest.mark.parametrize("k_chunk", [128]) @pytest.mark.parametrize("n_chunk", [64, 256]) @pytest.mark.parametrize("quant_type", [scalar_types.uint4, scalar_types.uint4b8]) @pytest.mark.parametrize("group_size", [-1, 128]) @pytest.mark.parametrize("mnk_factors", MNK_FACTORS) @pytest.mark.parametrize("act_order", [False, True]) def test_gptq_marlin_gemm( k_chunk, n_chunk, quant_type, group_size, mnk_factors, act_order, ): m_factor, n_factor, k_factor = mnk_factors has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8] size_m = m_factor size_k = k_chunk * k_factor size_n = n_chunk * n_factor if act_order: if group_size == -1: return if group_size == size_k: return if has_zp: return if size_k % group_size != 0: return a_input = torch.randn((size_m, size_k), dtype=torch.float16, device="cuda") b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda") if has_zp: w_ref, marlin_q_w, marlin_s, marlin_zp = awq_marlin_quantize( b_weight, quant_type, group_size ) g_idx = None sort_indices = None marlin_s2 = None else: w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize( b_weight, quant_type, group_size, act_order ) marlin_zp = None marlin_s2 = None workspace = marlin_make_workspace(w_ref.device) output = gptq_marlin_gemm( a_input, None, marlin_q_w, marlin_s, marlin_s2, marlin_zp, g_idx, sort_indices, workspace, quant_type, a_input.shape[0], b_weight.shape[1], a_input.shape[1], is_k_full=True, use_atomic_add=False, use_fp32_reduce=False, is_zp_float=False, ) output_ref = torch.matmul(a_input, w_ref) torch.cuda.synchronize() # JIT kernel should produce approximately correct results vs torch.matmul max_diff = torch.mean(torch.abs(output - output_ref)) / torch.mean( torch.abs(output_ref) ) assert max_diff < 0.04 if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_gptq_marlin_repack.py ================================================ import pytest import torch from sgl_kernel.scalar_type import scalar_types from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack from sglang.srt.layers.quantization.utils import ( gptq_quantize_weights, pack_rows, sort_weights, ) from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights MARLIN_K_CHUNKS = [128] MARLIN_N_CHUNKS = [64, 256] MNK_FACTORS = [ (1, 1, 1), (1, 4, 8), (1, 7, 5), (13, 17, 67), (26, 37, 13), (67, 13, 11), (257, 13, 11), (658, 13, 11), ] @pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS) @pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS) @pytest.mark.parametrize("quant_type", [scalar_types.uint4b8]) @pytest.mark.parametrize("group_size", [-1, 32, 64, 128]) @pytest.mark.parametrize("act_order", [False, True]) @pytest.mark.parametrize("mnk_factors", MNK_FACTORS) def test_gptq_marlin_repack( k_chunk, n_chunk, quant_type, group_size, act_order, mnk_factors ): m_factor, n_factor, k_factor = mnk_factors size_k = k_chunk * k_factor size_n = n_chunk * n_factor # Filter act_order if act_order: if group_size == -1: return if group_size == size_k: return # Normalize group_size if group_size == -1: group_size = size_k assert group_size <= size_k if size_k % group_size != 0: pytest.skip("size_k must be divisible by group_size") # Create input b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda") # Quantize (and apply act_order if provided) w_ref, q_w, s, g_idx, rand_perm = gptq_quantize_weights( b_weight, quant_type, group_size, act_order ) q_w_gptq = pack_rows(q_w, quant_type.size_bits, size_k, size_n) # For act_order, sort the "weights" and "g_idx" so that group ids are # increasing sort_indices = torch.empty(0, dtype=torch.int, device=b_weight.device) if act_order: q_w, g_idx, sort_indices = sort_weights(q_w, g_idx) marlin_layout_perm = get_weight_perm(quant_type.size_bits) q_w_marlin_ref = marlin_weights( q_w, size_k, size_n, quant_type.size_bits, marlin_layout_perm ) # Run JIT repack kernel jit_output = gptq_marlin_repack( q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits ) torch.cuda.synchronize() # JIT should match the reference (computed from CPU marlin_weights) torch.testing.assert_close(jit_output, q_w_marlin_ref) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_hadamard_jit.py ================================================ import math import numpy as np import pytest import torch import torch.nn.functional as F from scipy.linalg import hadamard from sglang.jit_kernel.hadamard import ( hadamard_transform, hadamard_transform_12n, hadamard_transform_20n, hadamard_transform_28n, hadamard_transform_40n, ) # Exact M×N Hadamard matrices (±1 entries) copied from # python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py. # These are non-power-of-2 Hadamard matrices constructed via Paley/Williamson methods. # "+" = +1, "-" = -1. Used by the _12n/_20n/_28n/_40n kernel variants. _HAD_12_STR = """ +-++++++++++ --+-+-+-+-+- +++-++----++ +---+--+-++- +++++-++---- +-+---+--+-+ ++--+++-++-- +--++---+--+ ++----+++-++ +--+-++---+- ++++----+++- +-+--+-++--- """ _HAD_20_STR = """ +----+----++--++-++- -+----+---+++---+-++ --+----+---+++-+-+-+ ---+----+---+++++-+- ----+----++--++-++-+ -+++++-----+--+++--+ +-+++-+---+-+--+++-- ++-++--+---+-+--+++- +++-+---+---+-+--+++ ++++-----++--+-+--++ --++-+-++-+-----++++ ---++-+-++-+---+-+++ +---++-+-+--+--++-++ ++---++-+----+-+++-+ -++---++-+----+++++- -+--+--++-+----+---- +-+-----++-+----+--- -+-+-+---+--+----+-- --+-+++------+----+- +--+--++------+----+ """ _HAD_28_STR = """ +------++----++-+--+-+--++-- -+-----+++-----+-+--+-+--++- --+-----+++---+-+-+----+--++ ---+-----+++---+-+-+-+--+--+ ----+-----+++---+-+-+++--+-- -----+-----++++--+-+--++--+- ------++----++-+--+-+--++--+ --++++-+-------++--+++-+--+- ---++++-+-----+-++--+-+-+--+ +---+++--+----++-++--+-+-+-- ++---++---+----++-++--+-+-+- +++---+----+----++-++--+-+-+ ++++--------+-+--++-++--+-+- -++++--------+++--++--+--+-+ -+-++-++--++--+--------++++- +-+-++--+--++--+--------++++ -+-+-++--+--++--+----+---+++ +-+-+-++--+--+---+---++---++ ++-+-+-++--+------+--+++---+ -++-+-+-++--+------+-++++--- +-++-+---++--+------+-++++-- -++--++-+-++-+++----++------ +-++--++-+-++-+++-----+----- ++-++---+-+-++-+++-----+---- -++-++-+-+-+-+--+++-----+--- --++-++++-+-+----+++-----+-- +--++-+-++-+-+----+++-----+- ++--++-+-++-+-+----++------+ """ _HAD_40_STR = """ +-------------------+------------------- ++-++----+-+-++++--+++-++----+-+-++++--+ +++-++----+-+-++++--+++-++----+-+-++++-- +-++-++----+-+-++++-+-++-++----+-+-++++- +--++-++----+-+-+++++--++-++----+-+-++++ ++--++-++----+-+-+++++--++-++----+-+-+++ +++--++-++----+-+-+++++--++-++----+-+-++ ++++--++-++----+-+-+++++--++-++----+-+-+ +++++--++-++----+-+-+++++--++-++----+-+- +-++++--++-++----+-++-++++--++-++----+-+ ++-++++--++-++----+-++-++++--++-++----+- +-+-++++--++-++----++-+-++++--++-++----+ ++-+-++++--++-++----++-+-++++--++-++---- +-+-+-++++--++-++---+-+-+-++++--++-++--- +--+-+-++++--++-++--+--+-+-++++--++-++-- +---+-+-++++--++-++-+---+-+-++++--++-++- +----+-+-++++--++-+++----+-+-++++--++-++ ++----+-+-++++--++-+++----+-+-++++--++-+ +++----+-+-++++--++-+++----+-+-++++--++- +-++----+-+-++++--+++-++----+-+-++++--++ +--------------------+++++++++++++++++++ ++-++----+-+-++++--+--+--++++-+-+----++- +++-++----+-+-++++-----+--++++-+-+----++ +-++-++----+-+-++++--+--+--++++-+-+----+ +--++-++----+-+-++++-++--+--++++-+-+---- ++--++-++----+-+-+++--++--+--++++-+-+--- +++--++-++----+-+-++---++--+--++++-+-+-- ++++--++-++----+-+-+----++--+--++++-+-+- +++++--++-++----+-+------++--+--++++-+-+ +-++++--++-++----+-+-+----++--+--++++-+- ++-++++--++-++----+---+----++--+--++++-+ +-+-++++--++-++----+-+-+----++--+--++++- ++-+-++++--++-++------+-+----++--+--++++ +-+-+-++++--++-++----+-+-+----++--+--+++ +--+-+-++++--++-++---++-+-+----++--+--++ +---+-+-++++--++-++--+++-+-+----++--+--+ +----+-+-++++--++-++-++++-+-+----++--+-- ++----+-+-++++--++-+--++++-+-+----++--+- +++----+-+-++++--++----++++-+-+----++--+ +-++----+-+-++++--++-+--++++-+-+----++-- """ def _parse_hadamard_str(s): """Parse a ±1 string matrix definition into a numpy array.""" s = s.strip().replace("+", "1").replace("-", "-1").split() return np.stack( [np.fromstring(" ".join(s[i]), dtype=np.int32, sep=" ") for i in range(len(s))] ) # Parsed M×M special Hadamard matrices, keyed by M (the "multiple"). # Copied from python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py # (had_12_paley, had_20_will, had_28_will, had_40_tpal) _SPECIAL_MATRICES = { 12: _parse_hadamard_str(_HAD_12_STR), 20: _parse_hadamard_str(_HAD_20_STR), 28: _parse_hadamard_str(_HAD_28_STR), 40: _parse_hadamard_str(_HAD_40_STR), } def hadamard_transform_ref(x, scale=1.0): """Reference impl for the general (power-of-2) hadamard_transform. Pads dim to the next power of 2, multiplies by the full H matrix via F.linear, then truncates back to the original dim. """ x_shape = x.shape dim = x.shape[-1] x = x.reshape(-1, dim) log_dim = math.ceil(math.log2(dim)) if dim > 0 else 0 dim_padded = 2**log_dim if dim > 0 else 1 if dim != dim_padded: x = F.pad(x, (0, dim_padded - dim)) H = torch.tensor(hadamard(dim_padded, dtype=float), dtype=x.dtype, device=x.device) out = F.linear(x, H) out = out * scale return out[..., :dim].reshape(*x_shape) def hadamard_transform_mn_ref(x, multiple, scale=1.0): """Reference impl for the M×N hadamard variants (_12n, _20n, _28n, _40n). The kernel computes (H_M ⊗ H_N) · x via two steps: 1) H_N (power-of-2 Hadamard) along the N dimension 2) H_M (special ±1 matrix) along the M dimension where dim = M * N, M = `multiple`, N = power of 2. """ x_shape = x.shape dim = x.shape[-1] x = x.reshape(-1, dim) # The kernel requires dim % (4*M) == 0 (for vectorized memory access). # See python/sglang/jit_kernel/hadamard.py: pad_multiple = 4 * 12 / 4 * 20 / etc. pad_multiple = 4 * multiple if dim % pad_multiple != 0: pad_size = pad_multiple - dim % pad_multiple x = F.pad(x, (0, pad_size)) dim_padded = dim + pad_size else: dim_padded = dim # N = dim_padded / M, must be a power of 2 n = dim_padded // multiple log_n = int(math.log2(n)) assert 2**log_n == n, f"n={n} is not a power of 2" batch = x.shape[0] x = x.reshape(batch, multiple, n) # (batch, M, N) # Step 1: apply H_N (standard power-of-2 Hadamard) along the N dimension H_n = torch.tensor(hadamard(n, dtype=float), dtype=x.dtype, device=x.device) x = torch.einsum("bmn,kn->bmk", x, H_n) # Step 2: apply H_M (special ±1 matrix) along the M dimension H_m = torch.tensor( _SPECIAL_MATRICES[multiple].astype(float), dtype=x.dtype, device=x.device ) x = torch.einsum("bmn,km->bkn", x, H_m) x = x.reshape(batch, -1) * scale return x[..., : x_shape[-1]].reshape(*x_shape) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) @pytest.mark.parametrize( "dim", # Power-of-2 dims from sgl-kernel/tests/test_hadamard.py (old AOT test) [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768], ) def test_hadamard_transform(dim, dtype): device = "cuda" # Tolerances from sgl-kernel/tests/test_hadamard.py (old AOT test) if dtype == torch.float32: rtol, atol = 3e-4, 3e-3 elif dtype == torch.bfloat16: rtol, atol = 1e-2, 5e-2 else: # float16 rtol, atol = 3e-3, 5e-3 torch.random.manual_seed(0) batch_size = 15 x = torch.randn(batch_size, dim, device=device, dtype=dtype) scale = 1.0 / math.sqrt(dim) out = hadamard_transform(x, scale=scale) # Compute reference in float32 from a detached copy to avoid precision loss out_ref = hadamard_transform_ref(x.detach().clone().float(), scale=scale) torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) @pytest.mark.parametrize( "dim", # Non-power-of-2 dims to test the padding path # (137 from sgl-kernel/tests/test_hadamard.py, 500/1000 added for coverage) [137, 500, 1000], ) def test_hadamard_transform_non_power_of_two(dim, dtype): device = "cuda" if dtype == torch.float32: rtol, atol = 3e-4, 3e-3 elif dtype == torch.bfloat16: rtol, atol = 1e-2, 5e-2 else: rtol, atol = 3e-3, 5e-3 torch.random.manual_seed(42) batch_size = 15 x = torch.randn(batch_size, dim, device=device, dtype=dtype) scale = 1.0 / math.sqrt(dim) out = hadamard_transform(x, scale=scale) out_ref = hadamard_transform_ref(x.detach().clone().float(), scale=scale) torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_hadamard_transform_3d_input(dtype): device = "cuda" if dtype == torch.bfloat16: rtol, atol = 1e-2, 5e-2 else: rtol, atol = 3e-3, 5e-3 torch.random.manual_seed(0) x = torch.randn(4, 8, 256, device=device, dtype=dtype) scale = 1.0 / math.sqrt(256) out = hadamard_transform(x, scale=scale) assert out.shape == x.shape out_ref = hadamard_transform_ref(x.detach().clone().float(), scale=scale) torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_hadamard_transform_scale_one(dtype): device = "cuda" if dtype == torch.bfloat16: rtol, atol = 1e-2, 5e-2 else: rtol, atol = 3e-3, 5e-3 torch.random.manual_seed(0) x = torch.randn(8, 64, device=device, dtype=dtype) out = hadamard_transform(x, scale=1.0) out_ref = hadamard_transform_ref(x.detach().clone().float(), scale=1.0) torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol) # Test dimensions for M×N variants: dim = M * N where N = 2^k. # M = 12/20/28/40 are the non-power-of-2 Hadamard sizes registered in # python/sglang/jit_kernel/hadamard.py (Hadamard12NKernel, ..., Hadamard40NKernel). # range(2,9) gives N = 4,8,...,256 so dims cover a practical range. _12N_DIMS = [12 * (2**k) for k in range(2, 9)] # 48, 96, ... , 3072 _20N_DIMS = [20 * (2**k) for k in range(2, 9)] # 80, 160, ... , 5120 _28N_DIMS = [28 * (2**k) for k in range(2, 9)] # 112, 224, ... , 7168 _40N_DIMS = [40 * (2**k) for k in range(2, 9)] # 160, 320, ... , 10240 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) @pytest.mark.parametrize("dim", _12N_DIMS) def test_hadamard_transform_12n(dim, dtype): device = "cuda" if dtype == torch.float32: rtol, atol = 3e-4, 3e-3 elif dtype == torch.bfloat16: rtol, atol = 1e-2, 5e-2 else: rtol, atol = 3e-3, 5e-3 torch.random.manual_seed(0) batch_size = 15 x = torch.randn(batch_size, dim, device=device, dtype=dtype) scale = 1.0 / math.sqrt(dim) out = hadamard_transform_12n(x, scale=scale) out_ref = hadamard_transform_mn_ref(x.detach().clone().float(), 12, scale=scale) torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) @pytest.mark.parametrize("dim", _20N_DIMS) def test_hadamard_transform_20n(dim, dtype): device = "cuda" if dtype == torch.float32: rtol, atol = 3e-4, 3e-3 elif dtype == torch.bfloat16: rtol, atol = 1e-2, 5e-2 else: rtol, atol = 3e-3, 5e-3 torch.random.manual_seed(0) batch_size = 15 x = torch.randn(batch_size, dim, device=device, dtype=dtype) scale = 1.0 / math.sqrt(dim) out = hadamard_transform_20n(x, scale=scale) out_ref = hadamard_transform_mn_ref(x.detach().clone().float(), 20, scale=scale) torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) @pytest.mark.parametrize("dim", _28N_DIMS) def test_hadamard_transform_28n(dim, dtype): device = "cuda" if dtype == torch.float32: rtol, atol = 3e-4, 3e-3 elif dtype == torch.bfloat16: rtol, atol = 1e-2, 5e-2 else: rtol, atol = 3e-3, 5e-3 torch.random.manual_seed(0) batch_size = 15 x = torch.randn(batch_size, dim, device=device, dtype=dtype) scale = 1.0 / math.sqrt(dim) out = hadamard_transform_28n(x, scale=scale) out_ref = hadamard_transform_mn_ref(x.detach().clone().float(), 28, scale=scale) torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) @pytest.mark.parametrize("dim", _40N_DIMS) def test_hadamard_transform_40n(dim, dtype): device = "cuda" if dtype == torch.float32: rtol, atol = 3e-4, 3e-3 elif dtype == torch.bfloat16: rtol, atol = 1e-2, 5e-2 else: rtol, atol = 3e-3, 5e-3 torch.random.manual_seed(0) batch_size = 15 x = torch.randn(batch_size, dim, device=device, dtype=dtype) scale = 1.0 / math.sqrt(dim) out = hadamard_transform_40n(x, scale=scale) out_ref = hadamard_transform_mn_ref(x.detach().clone().float(), 40, scale=scale) torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol) if __name__ == "__main__": pytest.main([__file__]) ================================================ FILE: python/sglang/jit_kernel/tests/test_moe_lora_align_block_size.py ================================================ # Temporarily adapted from https://github.com/vllm-project/vllm/blob/main/tests/lora/test_moe_lora_align_sum.py, will optimize in future refactor import random import pytest import torch # --------------------------------------------------------- # IMPORT PREBUILT KERNEL # --------------------------------------------------------- from sglang.jit_kernel.moe_lora_align import moe_lora_align_block_size from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=80, suite="stage-b-test-large-1-gpu") def round_up(x, base): return ((x + base - 1) // base) * base def CEILDIV(x, y): return (x + y - 1) // y def sample_data(num_experts, max_loras, num_tokens, topk_num): # 1. Generate TopK IDs (Flattened tokens) topk_ids = torch.zeros((num_tokens, topk_num), dtype=torch.int32) for i in range(num_tokens): pool = list(range(num_experts)) random.shuffle(pool) for j in range(topk_num): topk_ids[i, j] = pool[j] # 2. Generate Random Requests (Segments) # We split num_tokens into random chunks to simulate a batch of requests remaining_tokens = num_tokens seg_lens = [] while remaining_tokens > 0: # Random length between 1 and remaining length = random.randint(1, min(32, remaining_tokens)) if remaining_tokens - length < 0: length = remaining_tokens seg_lens.append(length) remaining_tokens -= length # Ensure we cover the full range exactly (cleanup last segment) if sum(seg_lens) < num_tokens: seg_lens.append(num_tokens - sum(seg_lens)) # 3. Build seg_indptr [0, len1, len1+len2, ...] seg_indptr = torch.cumsum( torch.tensor([0] + seg_lens, dtype=torch.int32), dim=0 ).to(dtype=torch.int32) # 4. Assign a LoRA ID to each Request num_reqs = len(seg_lens) req_to_lora = torch.randint(0, max_loras, (num_reqs,), dtype=torch.int32) return (topk_ids.to("cuda"), seg_indptr.to("cuda"), req_to_lora.to("cuda")) @pytest.mark.parametrize("num_tokens", [100, 200, 1024, 4096]) @pytest.mark.parametrize("topk_num", [6]) @pytest.mark.parametrize("num_experts", [64, 128, 256, 512]) @pytest.mark.parametrize("max_loras", [2, 32]) @pytest.mark.parametrize("block_size", [16]) def test_moe_lora_align_block_size( num_tokens, topk_num, num_experts, max_loras, block_size ): # sample data random.seed(1) torch.manual_seed(1) if not torch.cuda.is_available(): pytest.skip("CUDA is not available, skipping moe_lora_align_block_size test.") # UPDATED: Get the new 3-step mapping tensors topk_ids, seg_indptr, req_to_lora = sample_data( num_experts, max_loras, num_tokens, topk_num ) # compute paddings max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) max_num_tokens_padded = round_up(max_num_tokens_padded, block_size) max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size) # init output tensors sorted_token_ids = torch.full( (max_loras * max_num_tokens_padded,), topk_ids.numel(), dtype=torch.int32, device="cuda", ) expert_ids = torch.full( (max_loras * max_num_m_blocks,), num_experts, dtype=torch.int32, device="cuda" ) num_tokens_post_pad = torch.zeros((max_loras,), dtype=torch.int32, device="cuda") adapter_enabled = torch.ones((max_loras + 1,), dtype=torch.int32, device="cuda") lora_ids = torch.arange(max_loras, dtype=torch.int32, device="cuda") # UPDATED: Call kernel with new signature moe_lora_align_block_size( topk_ids, seg_indptr, # Arg 2: Pointers req_to_lora, # Arg 3: Request Map num_experts, block_size, max_loras, max_num_tokens_padded, max_num_m_blocks, sorted_token_ids, expert_ids, num_tokens_post_pad, adapter_enabled, lora_ids, None, ) # verify values expert_ids = expert_ids.view(max_loras, -1) sorted_token_ids = sorted_token_ids.view(max_loras, -1, block_size) # Reconstruct token-level ownership for verification logic # We expand req_to_lora back to [num_tokens] on CPU just to check correctness # This proves the kernel (which used the compressed format) produced the right result cpu_seg_indptr = seg_indptr.cpu() cpu_req_to_lora = req_to_lora.cpu() token_ownership = torch.zeros(num_tokens, dtype=torch.int32) for r in range(len(cpu_req_to_lora)): start = cpu_seg_indptr[r] end = cpu_seg_indptr[r + 1] token_ownership[start:end] = cpu_req_to_lora[r] token_ownership = token_ownership.to("cuda") for lora_idx in range(max_loras): # Count how many tokens actually belong to this LoRA expected_count = (token_ownership == lora_idx).sum().item() # Verify the kernel processed a reasonable number of tokens (sanity check) # Note: num_tokens_post_pad includes padding, so it might be larger than expected_count assert num_tokens_post_pad[lora_idx].item() >= expected_count * topk_num for token_idx in range(sorted_token_ids.size(1)): block = sorted_token_ids[lora_idx][token_idx] # Valid indices are those less than total numel indices = block[block != topk_ids.numel()] if indices.numel() > 0: # 1. Verify routing: Does the token actually route to this expert? expert_id = expert_ids[lora_idx][token_idx] assert torch.all(topk_ids.view(-1)[indices] == expert_id) # 2. Verify ownership: Did the kernel grab the correct tokens for this LoRA? # The indices in 'sorted_token_ids' point to the flattened [token, topk] array. # We divide by topk_num to get the original token index. original_token_indices = indices // topk_num # Check that all tokens in this block truly belong to 'lora_idx' actual_owners = token_ownership[original_token_indices] assert torch.all( actual_owners == lora_idx ), f"Kernel put tokens from LoRA {actual_owners} into block for LoRA {lora_idx}" if __name__ == "__main__": pytest.main([__file__]) ================================================ FILE: python/sglang/jit_kernel/tests/test_moe_wna16_marlin.py ================================================ import itertools import pytest import torch from sgl_kernel.scalar_type import scalar_types from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm from sglang.srt.layers.moe.fused_moe_triton import moe_align_block_size from sglang.test.test_marlin_utils import awq_marlin_quantize, marlin_quantize def _has_aot_moe_wna16_marlin_gemm() -> bool: return hasattr(torch.ops.sgl_kernel, "moe_wna16_marlin_gemm") and hasattr( torch.ops.sgl_kernel.moe_wna16_marlin_gemm, "default" ) AOT_AVAILABLE = _has_aot_moe_wna16_marlin_gemm() def stack_and_dev(tensors: list[torch.Tensor]): dev = tensors[0].device return torch.stack(tensors, dim=0).to(dev) def _get_scalar_type(num_bits: int, has_zp: bool): if has_zp: assert num_bits == 4 return scalar_types.uint4 else: return scalar_types.uint4b8 if num_bits == 4 else scalar_types.uint8b128 def _setup_moe_weights(e, n, k, quant_type, group_size, act_order, dtype): """Set up quantized MoE weights for a single gate (e experts, output n, input k).""" has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8] w = torch.randn((e, n, k), device="cuda", dtype=dtype) / 20 w_ref_l = [] qweight_l = [] scales_l = [] zeros_l = [] g_idx_l = [] sort_indices_l = [] for i in range(e): if has_zp: w_ref, qweight, scales, zeros = awq_marlin_quantize( w[i].transpose(1, 0), quant_type, group_size ) w_ref_l.append(w_ref.T) qweight_l.append(qweight) scales_l.append(scales) zeros_l.append(zeros) else: test_perm = torch.randperm(k) w_ref, qweight, scales, g_idx, sort_indices, _ = marlin_quantize( w[i].transpose(1, 0), quant_type, group_size, act_order, test_perm ) w_ref_l.append(w_ref.T) qweight_l.append(qweight) scales_l.append(scales) g_idx_l.append(g_idx) sort_indices_l.append(sort_indices) w_ref = stack_and_dev(w_ref_l) qweight = stack_and_dev(qweight_l).contiguous() scales = stack_and_dev(scales_l) g_idx = stack_and_dev(g_idx_l) if g_idx_l else None sort_indices = stack_and_dev(sort_indices_l) if sort_indices_l else None zeros = stack_and_dev(zeros_l) if zeros_l else None return w_ref, qweight, scales, zeros, g_idx, sort_indices def _run_single_gemm( fn, a, c, qweight, scales, zeros, g_idx, sort_indices, workspace, sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights, quant_type, block_size_m, topk, size_m, size_n, size_k, mul_topk_weights, is_k_full, use_atomic_add, ): return fn( a, c, qweight, None, # b_bias scales, None, # global_scale zeros, g_idx, sort_indices, workspace, sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights, moe_block_size=block_size_m, top_k=topk, mul_topk_weights=mul_topk_weights, is_ep=False, b_q_type=quant_type, size_m=size_m, size_n=size_n, size_k=size_k, is_k_full=is_k_full, use_atomic_add=use_atomic_add, use_fp32_reduce=True, is_zp_float=False, ) def _run_single_gemm_aot( a, c, qweight, scales, zeros, g_idx, sort_indices, workspace, sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights, quant_type, block_size_m, topk, size_m, size_n, size_k, mul_topk_weights, is_k_full, use_atomic_add, ): return torch.ops.sgl_kernel.moe_wna16_marlin_gemm.default( a, c, qweight, None, # b_bias scales, None, # global_scale zeros, g_idx, sort_indices, workspace, sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights, moe_block_size=block_size_m, top_k=topk, mul_topk_weights=mul_topk_weights, is_ep=False, b_q_type_id=quant_type.id, size_m=size_m, size_n=size_n, size_k=size_k, is_k_full=is_k_full, use_atomic_add=use_atomic_add, use_fp32_reduce=True, is_zp_float=False, ) def generate_test_cases(): m_list = [1, 123] n_list = [128, 1024] k_list = [256] e_list = [4] topk_list = [2] dtype_list = [torch.float16, torch.bfloat16] group_size_list = [128] act_order_list = [False, True] quant_type_list = [scalar_types.uint4, scalar_types.uint4b8] all_combinations = itertools.product( m_list, n_list, k_list, e_list, topk_list, dtype_list, group_size_list, act_order_list, quant_type_list, ) def is_valid(m, n, k, e, topk, dtype, group_size, act_order, quant_type): has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8] if act_order: if group_size == -1 or group_size == k: return False if has_zp: return False if group_size > 0 and k % group_size != 0: return False return True return [case for case in all_combinations if is_valid(*case)] TEST_CASES = generate_test_cases() @pytest.mark.parametrize( "m,n,k,e,topk,dtype,group_size,act_order,quant_type", TEST_CASES, ids=[ f"m{c[0]}_n{c[1]}_k{c[2]}_e{c[3]}_t{c[4]}_{c[5].__name__ if hasattr(c[5], '__name__') else str(c[5]).split('.')[-1]}_g{c[6]}_act{c[7]}_{c[8]}" for c in TEST_CASES ], ) def test_moe_wna16_marlin_gemm( m, n, k, e, topk, dtype, group_size, act_order, quant_type ): if not AOT_AVAILABLE: pytest.skip("sgl_kernel moe_wna16_marlin_gemm AOT op not available") torch.manual_seed(0) has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8] a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 # Set up quantized weights for first gemm (gate_up: output 2*n, input k) w_ref1, qweight1, scales1, zeros1, g_idx1, sort_indices1 = _setup_moe_weights( e, 2 * n, k, quant_type, group_size, act_order, dtype ) # Compute block_size_m for block_size_m in [8, 16, 32, 48, 64]: if m * topk / e / block_size_m < 0.9: break # Align tokens score = torch.randn((m, e), device="cuda", dtype=dtype) score_softmax = torch.softmax(score, dim=-1, dtype=torch.float32) topk_weights, topk_ids = torch.topk(score_softmax, topk) sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( topk_ids, block_size_m, e ) # Workspace sms = torch.cuda.get_device_properties("cuda").multi_processor_count max_workspace_size = (max(2 * n, k) // 64) * ( sorted_token_ids.size(0) // block_size_m ) max_workspace_size = min(max_workspace_size, sms * 4) workspace = torch.zeros( max_workspace_size, dtype=torch.int, device="cuda", requires_grad=False ) use_atomic_add = ( dtype == torch.half or torch.cuda.get_device_capability("cuda")[0] >= 9 ) scalar_type = _get_scalar_type(4, has_zp) # --- Run JIT kernel --- c_jit = torch.empty((m * topk, 2 * n), dtype=dtype, device="cuda") c_jit = _run_single_gemm( moe_wna16_marlin_gemm, a, c_jit, qweight1, scales1, zeros1, g_idx1, sort_indices1, workspace, sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights, scalar_type, block_size_m, topk, m, 2 * n, k, False, True, use_atomic_add, ) torch.cuda.synchronize() # --- Check bitwise equality with AOT kernel --- c_aot = torch.empty((m * topk, 2 * n), dtype=dtype, device="cuda") c_aot = _run_single_gemm_aot( a, c_aot, qweight1, scales1, zeros1, g_idx1, sort_indices1, workspace, sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights, scalar_type, block_size_m, topk, m, 2 * n, k, False, True, use_atomic_add, ) torch.cuda.synchronize() torch.testing.assert_close(c_jit, c_aot, rtol=0, atol=0) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_norm_jit.py ================================================ # Adapted from sgl-kernel/tests/test_norm.py import pytest import torch # JIT rmsnorm: fp16/bf16 only; hidden_size must be a multiple of 256, > 256, and <=8192 RMSNORM_HIDDEN_SIZES = [512, 1024, 3072, 3584, 4096, 8192] # JIT fused_add_rmsnorm: fp16/bf16 only; hidden_size % 8 == 0, <=8192 FUSED_ADD_RMSNORM_HIDDEN_SIZES = [1024, 3072, 3584, 4096, 8192] BS_LIST = [1, 19, 99, 989] def _jit_rmsnorm(input, weight, output, eps): from sglang.jit_kernel.norm import rmsnorm rmsnorm(input, weight, output=output, eps=eps) def _fi_rmsnorm(input, weight, out, eps): from flashinfer.norm import rmsnorm rmsnorm(input, weight, out=out, eps=eps) def _jit_fused_add_rmsnorm(input, residual, weight, eps): from sglang.jit_kernel.norm import fused_add_rmsnorm fused_add_rmsnorm(input, residual, weight, eps) def _fi_fused_add_rmsnorm(input, residual, weight, eps): from flashinfer.norm import fused_add_rmsnorm fused_add_rmsnorm(input, residual, weight, eps=eps) @pytest.mark.parametrize("batch_size", BS_LIST) @pytest.mark.parametrize("hidden_size", RMSNORM_HIDDEN_SIZES) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("specify_out", [True, False]) def test_rmsnorm_jit(batch_size, hidden_size, dtype, specify_out): eps = 1e-6 x = torch.randn(batch_size, hidden_size, device="cuda", dtype=dtype) w = torch.randn(hidden_size, device="cuda", dtype=dtype) # flashinfer reference x_ref = x.clone() _fi_rmsnorm(x_ref, w, out=x_ref, eps=eps) if specify_out: y = torch.empty_like(x) _jit_rmsnorm(x, w, output=y, eps=eps) else: y = x.clone() _jit_rmsnorm(y, w, output=y, eps=eps) torch.testing.assert_close(y, x_ref, rtol=1e-2, atol=1e-2) @pytest.mark.parametrize("batch_size", BS_LIST) @pytest.mark.parametrize("hidden_size", FUSED_ADD_RMSNORM_HIDDEN_SIZES) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_fused_add_rmsnorm_jit(batch_size, hidden_size, dtype): eps = 1e-6 x = torch.randn(batch_size, hidden_size, dtype=dtype, device="cuda") residual = torch.randn_like(x) weight = torch.randn(hidden_size, dtype=dtype, device="cuda") # flashinfer reference x_ref = x.clone() r_ref = residual.clone() _fi_fused_add_rmsnorm(x_ref, r_ref, weight, eps=eps) x_jit = x.clone() r_jit = residual.clone() _jit_fused_add_rmsnorm(x_jit, r_jit, weight, eps) torch.testing.assert_close(x_jit, x_ref, rtol=1e-2, atol=1e-2) torch.testing.assert_close(r_jit, r_ref, rtol=1e-2, atol=1e-2) if __name__ == "__main__": pytest.main([__file__]) ================================================ FILE: python/sglang/jit_kernel/tests/test_nvfp4_blockwise_moe.py ================================================ import pytest import torch from sglang.jit_kernel.nvfp4 import ( cutlass_fp4_group_mm, scaled_fp4_experts_quant, scaled_fp4_quant, ) FLOAT4_E2M1_MAX = 6.0 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max def _nvfp4_supported() -> bool: return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0) def _round_up(x: int, y: int) -> int: return ((x + y - 1) // y) * y def _build_expert_offsets( m_per_expert: list[int], device: torch.device ) -> torch.Tensor: offsets = [0] for m in m_per_expert: offsets.append(offsets[-1] + m) return torch.tensor(offsets, dtype=torch.int32, device=device) def _build_blockscale_offsets( m_per_expert: list[int], device: torch.device ) -> torch.Tensor: offsets = [0] for m in m_per_expert: offsets.append(offsets[-1] + _round_up(m, 128)) return torch.tensor(offsets, dtype=torch.int32, device=device) @pytest.mark.skipif( not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0" ) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_nvfp4_blockwise_moe_grouped_mm(dtype: torch.dtype) -> None: torch.manual_seed(0) device = torch.device("cuda") num_experts = 4 m_per_expert = [33, 17, 48, 29] n = 256 k = 128 expert_offsets_full = _build_expert_offsets(m_per_expert, device) blockscale_offsets_full = _build_blockscale_offsets(m_per_expert, device) total_m = int(expert_offsets_full[-1].item()) a = torch.randn((total_m, k), device=device, dtype=dtype) * 0.1 b = torch.randn((num_experts, n, k), device=device, dtype=dtype) * 0.1 a_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32) for i in range(num_experts): start = int(expert_offsets_full[i].item()) end = int(expert_offsets_full[i + 1].item()) amax = a[start:end].abs().max().to(torch.float32) a_global_scale[i] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / amax b_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32) for i in range(num_experts): bmax = b[i].abs().max().to(torch.float32) b_global_scale[i] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / bmax a_fp4, a_blockscale = scaled_fp4_experts_quant( a, a_global_scale, expert_offsets_full, blockscale_offsets_full, topk=1, ) b_fp4 = torch.empty((num_experts, n, k // 2), device=device, dtype=torch.uint8) b_blockscale = torch.empty( (num_experts, _round_up(n, 128), _round_up(k // 16, 4)), device=device, dtype=torch.float8_e4m3fn, ) for i in range(num_experts): b_fp4_i, b_scale_i = scaled_fp4_quant(b[i], b_global_scale[i]) b_fp4[i].copy_(b_fp4_i) b_blockscale[i].copy_(b_scale_i) alphas = (1.0 / (a_global_scale * b_global_scale)).to(torch.float32) params = { "ab_strides": torch.full((num_experts,), k, dtype=torch.int64, device=device), "c_strides": torch.full((num_experts,), n, dtype=torch.int64, device=device), "problem_sizes": torch.tensor( [[m, n, k] for m in m_per_expert], dtype=torch.int32, device=device ), "expert_offsets": expert_offsets_full[:-1].contiguous(), "blockscale_offsets": blockscale_offsets_full[:-1].contiguous(), "a_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), "b_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), "out_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), "a_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), "b_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), "alpha_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), "layout_sfa": torch.empty((num_experts, 5), dtype=torch.int64, device=device), "layout_sfb": torch.empty((num_experts, 5), dtype=torch.int64, device=device), } out = cutlass_fp4_group_mm( a_fp4, b_fp4, a_blockscale, b_blockscale, alphas, dtype, params, ) ref = torch.empty((total_m, n), device=device, dtype=dtype) for i in range(num_experts): start = int(expert_offsets_full[i].item()) end = int(expert_offsets_full[i + 1].item()) ref[start:end] = torch.matmul(a[start:end], b[i].t()) torch.testing.assert_close(out, ref, atol=1e-1, rtol=1e-1) ================================================ FILE: python/sglang/jit_kernel/tests/test_nvfp4_gemm.py ================================================ import pytest import torch from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant def _nvfp4_supported() -> bool: return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0) DTYPES = [torch.float16, torch.bfloat16] SHAPES = [ (128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128), (150, 128, 64), ] FLOAT4_E2M1_MAX = 6.0 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max K_E2M1_TO_FLOAT = [ 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, ] def e2m1_to_fp32(int4_value: int) -> float: sign_bit = int4_value & 0x8 int4_abs_value = int4_value & 0x7 float_result = K_E2M1_TO_FLOAT[int4_abs_value] return -float_result if sign_bit else float_result def break_fp4_bytes(a: torch.Tensor) -> torch.Tensor: assert a.dtype == torch.uint8 m, n = a.shape a = a.flatten() high_half_byte = (a & 0xF0) >> 4 low_half_byte = a & 0x0F f_h = torch.tensor([e2m1_to_fp32(x) for x in high_half_byte], device=a.device) f_l = torch.tensor([e2m1_to_fp32(x) for x in low_half_byte], device=a.device) return torch.stack((f_l, f_h), dim=-1).reshape(m, n * 2) def convert_swizzled_to_linear( a_sf_swizzled: torch.Tensor, m: int, k: int, block_size: int ) -> torch.Tensor: sf_m, sf_k = a_sf_swizzled.shape del sf_m, sf_k m_tiles = (m + 128 - 1) // 128 f = block_size * 4 k_tiles = (k + f - 1) // f tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4)) tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5)) out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size) return out[0:m, 0 : k // block_size] def dequantize_to_dtype( tensor_fp4: torch.Tensor, tensor_sf: torch.Tensor, global_scale: torch.Tensor, block_size: int = 16, ) -> torch.Tensor: assert tensor_fp4.dtype == torch.uint8 m, packed_k = tensor_fp4.shape k = packed_k * 2 tensor_f32 = break_fp4_bytes(tensor_fp4) tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size) tensor_sf = tensor_sf.view(torch.float8_e4m3fn) tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size) tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale return (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k) def get_ref_results( a_fp4: torch.Tensor, b_fp4: torch.Tensor, a_sf: torch.Tensor, b_sf: torch.Tensor, a_global_scale: torch.Tensor, b_global_scale: torch.Tensor, block_size: int, ) -> torch.Tensor: a_in_dtype = dequantize_to_dtype(a_fp4, a_sf, a_global_scale, block_size=block_size) b_in_dtype = dequantize_to_dtype(b_fp4, b_sf, b_global_scale, block_size=block_size) return torch.matmul(a_in_dtype, b_in_dtype.t()) @pytest.mark.skipif( not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0" ) @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("shape", SHAPES) def test_nvfp4_gemm(dtype: torch.dtype, shape: tuple[int, int, int]) -> None: m, n, packed_k = shape k = packed_k * 2 block_size = 16 a_dtype = torch.randn((m, k), dtype=dtype, device="cuda") b_dtype = torch.randn((n, k), dtype=dtype, device="cuda") a_global_scale = ( (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1) ).to(torch.float32) b_global_scale = ( (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1) ).to(torch.float32) alpha = 1.0 / (a_global_scale * b_global_scale) a_fp4, a_scale_interleaved = scaled_fp4_quant(a_dtype, a_global_scale) b_fp4, b_scale_interleaved = scaled_fp4_quant(b_dtype, b_global_scale) expected_out = get_ref_results( a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, a_global_scale, b_global_scale, block_size, ) out = cutlass_scaled_fp4_mm( a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, dtype, ) torch.testing.assert_close(out, expected_out.to(dtype=dtype), atol=1e-1, rtol=1e-1) ================================================ FILE: python/sglang/jit_kernel/tests/test_nvfp4_quant.py ================================================ import pytest import torch from sglang.jit_kernel.nvfp4 import ( scaled_fp4_grouped_quant, scaled_fp4_quant, silu_and_mul_scaled_fp4_grouped_quant, ) try: from sgl_kernel import silu_and_mul as _sgl_silu_and_mul except Exception: _sgl_silu_and_mul = None def _nvfp4_supported() -> bool: return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0) def _silu_and_mul_reference(x: torch.Tensor) -> torch.Tensor: if _sgl_silu_and_mul is not None: return _sgl_silu_and_mul(x) k = x.shape[-1] // 2 return torch.nn.functional.silu(x[:, :, :k]) * x[:, :, k:] DTYPES = [torch.float16, torch.bfloat16] SHAPES = [(128, 64), (128, 128), (256, 64), (256, 128)] PAD_SHAPES = [ (90, 64), (150, 64), (128, 48), (128, 80), ] FLOAT4_E2M1_MAX = 6.0 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max BLOCK_SIZE = 16 E2M1_TO_FLOAT32 = [ 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, ] def cast_from_fp4(x: torch.Tensor, m: int, n: int) -> torch.Tensor: v_2nd = (x & 0xF).to(torch.long) v_1st = ((x >> 4) & 0xF).to(torch.long) c = torch.stack((v_2nd, v_1st), dim=-1).flatten() lut = torch.tensor(E2M1_TO_FLOAT32, device=x.device, dtype=torch.float32) return lut[c].reshape(m, n) def cast_to_fp4(x: torch.Tensor) -> torch.Tensor: sign = torch.sign(x) x = torch.abs(x) x[(x >= 0.0) & (x <= 0.25)] = 0.0 x[(x > 0.25) & (x < 0.75)] = 0.5 x[(x >= 0.75) & (x <= 1.25)] = 1.0 x[(x > 1.25) & (x < 1.75)] = 1.5 x[(x >= 1.75) & (x <= 2.5)] = 2.0 x[(x > 2.5) & (x < 3.5)] = 3.0 x[(x >= 3.5) & (x <= 5.0)] = 4.0 x[x > 5.0] = 6.0 return x * sign def get_reciprocal(x): if isinstance(x, torch.Tensor): return torch.where(x == 0, torch.tensor(0.0, dtype=x.dtype), 1.0 / x) return 0.0 if x == 0 else 1.0 / x def ref_nvfp4_quant(x: torch.Tensor, global_scale: torch.Tensor): assert global_scale.dtype == torch.float32 assert x.ndim == 2 m, n = x.shape x = torch.reshape(x, (m, n // BLOCK_SIZE, BLOCK_SIZE)) vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32) scale = global_scale * (vec_max * get_reciprocal(FLOAT4_E2M1_MAX)) scale = scale.to(torch.float8_e4m3fn).to(torch.float32) output_scale = get_reciprocal(scale * get_reciprocal(global_scale)) scaled_x = x.to(torch.float32) * output_scale clipped_x = torch.clamp(scaled_x, -6.0, 6.0).reshape(m, n) return cast_to_fp4(clipped_x), scale.squeeze(-1) def recover_swizzled_scales(scale: torch.Tensor, m: int, n: int) -> torch.Tensor: rounded_m = ((m + 128 - 1) // 128) * 128 scale_n = n // BLOCK_SIZE rounded_n = ((scale_n + 4 - 1) // 4) * 4 tmp = torch.reshape(scale, (1, rounded_m // 128, rounded_n // 4, 32, 4, 4)) tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5)) result = torch.reshape(tmp, (rounded_m, rounded_n)).to(torch.float32) return result[:m, :scale_n] @pytest.mark.skipif( not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0" ) @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("shape", SHAPES) def test_quantize_to_fp4(dtype: torch.dtype, shape: tuple[int, int]) -> None: torch.manual_seed(42) m, n = shape x = torch.randn((m, n), dtype=dtype, device="cuda") tensor_amax = torch.abs(x).max().to(torch.float32) global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax out_ref, scale_ref = ref_nvfp4_quant(x, global_scale) out, out_scale = scaled_fp4_quant(x, global_scale) scale_ans = recover_swizzled_scales(out_scale, m, n) out_ans = cast_from_fp4(out, m, n) torch.testing.assert_close(out_ans, out_ref) torch.testing.assert_close(scale_ans, scale_ref) @pytest.mark.skipif( not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0" ) @pytest.mark.parametrize("shape", PAD_SHAPES) def test_quantize_to_fp4_padded(shape: tuple[int, int]) -> None: torch.manual_seed(42) m, n = shape x = torch.randn((m, n), dtype=torch.float16, device="cuda") tensor_amax = torch.abs(x).max().to(torch.float32) global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax out_ref, scale_ref = ref_nvfp4_quant(x, global_scale) out, out_scale = scaled_fp4_quant(x, global_scale) scale_ans = recover_swizzled_scales(out_scale, m, n) out_ans = cast_from_fp4(out, m, n) torch.testing.assert_close(out_ans, out_ref) torch.testing.assert_close(scale_ans, scale_ref) @pytest.mark.skipif( not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0" ) @pytest.mark.parametrize("shape", [(2, 128, 512), (2, 100, 128)]) def test_quantize_to_fp4_grouped(shape: tuple[int, int, int]) -> None: torch.manual_seed(42) l, m, k = shape x = torch.randn((l, m, k), dtype=torch.bfloat16, device="cuda") mask = torch.randint(1, max(2, m // 2), (l,), dtype=torch.int32, device="cuda") tensor_amax = x.abs().amax(dim=(1, 2)).to(torch.float32) x_sf_global = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax output, output_scales = scaled_fp4_grouped_quant(x, x_sf_global, mask) output = output.permute(2, 0, 1) padded_m = ((m + 128 - 1) // 128) * 128 output_scales = output_scales.permute(5, 2, 4, 0, 1, 3).view(l, padded_m, -1) for i in range(l): a_fp4, a_scale_interleaved = scaled_fp4_quant(x[i], x_sf_global[i]) torch.testing.assert_close(a_fp4[: mask[i]], output[i][: mask[i]]) scale_ref = recover_swizzled_scales(a_scale_interleaved, m, k) scale_ans = recover_swizzled_scales(output_scales[i], m, k) torch.testing.assert_close(scale_ref[: mask[i]], scale_ans[: mask[i]]) @pytest.mark.skipif( not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0" ) @pytest.mark.parametrize("shape", [(4, 96, 256), (8, 128, 512)]) def test_silu_and_mul_quantize_to_fp4_grouped(shape: tuple[int, int, int]) -> None: torch.manual_seed(42) l, m, k = shape x = torch.randn((l, m, k * 2), dtype=torch.bfloat16, device="cuda") mask = torch.randint(1, max(2, m // 2), (l,), dtype=torch.int32, device="cuda") ref_y = _silu_and_mul_reference(x) tensor_amax = ref_y.abs().amax(dim=(1, 2)).to(torch.float32) y_sf_global = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax ref_output, ref_output_scales = scaled_fp4_grouped_quant(ref_y, y_sf_global, mask) output, output_scales = silu_and_mul_scaled_fp4_grouped_quant(x, y_sf_global, mask) output = output.permute(2, 0, 1) ref_output = ref_output.permute(2, 0, 1) padded_m = ((m + 128 - 1) // 128) * 128 output_scales = output_scales.permute(5, 2, 4, 0, 1, 3).view(l, padded_m, -1) ref_output_scales = ref_output_scales.permute(5, 2, 4, 0, 1, 3).view( l, padded_m, -1 ) for i in range(l): torch.testing.assert_close(ref_output[i, : mask[i]], output[i, : mask[i]]) scale_ref = recover_swizzled_scales(ref_output_scales[i], m, k) scale_ans = recover_swizzled_scales(output_scales[i], m, k) torch.testing.assert_close(scale_ref[: mask[i]], scale_ans[: mask[i]]) ================================================ FILE: python/sglang/jit_kernel/tests/test_per_tensor_quant_fp8.py ================================================ import itertools from typing import Optional, Tuple import pytest import torch from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 try: from sglang.srt.utils import is_hip _is_hip = is_hip() except ImportError: _is_hip = False fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn def sglang_scaled_fp8_quant( input: torch.Tensor, scale: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: fp8_type_: torch.dtype = torch.float8_e4m3fn output = torch.empty_like(input, device=input.device, dtype=fp8_type_) is_static = True if scale is None: scale = torch.zeros(1, device=input.device, dtype=torch.float32) is_static = False per_tensor_quant_fp8(input, output, scale, is_static) return output, scale def torch_scaled_fp8_quant(tensor, inv_scale): finfo = torch.finfo(torch.float8_e4m3fn) scale = inv_scale.reciprocal() qweight = (tensor.to(torch.float32) * scale).clamp(min=finfo.min, max=finfo.max) qweight = qweight.to(torch.float8_e4m3fn) return qweight @pytest.mark.parametrize( "num_tokens,hidden_dim", list(itertools.product([128, 256, 512], [512, 2048, 4096])), ) def test_jit_per_tensor_quant_compare_implementations( num_tokens: int, hidden_dim: int, ): device = torch.device("cuda") x = torch.rand((num_tokens, hidden_dim), dtype=torch.float16, device=device) sglang_out, sglang_scale = sglang_scaled_fp8_quant(x) torch_out = torch_scaled_fp8_quant(x, sglang_scale) torch.testing.assert_close( sglang_out.float(), torch_out.float(), rtol=1e-3, atol=1e-3 ) @pytest.mark.parametrize("shape", [(4, 8, 64), (2, 16, 128), (19260817, 1, 1)]) def test_jit_per_tensor_quant_supports_3d(shape): device = torch.device("cuda") x = torch.rand(shape, dtype=torch.bfloat16, device=device) out = torch.empty_like(x, device=x.device, dtype=fp8_type_) scale = torch.zeros(1, device=x.device, dtype=torch.float32) per_tensor_quant_fp8(x, out, scale, is_static=False) x_2d = x.flatten(0, -2) out_ref_2d = torch_scaled_fp8_quant(x_2d, scale) out_ref = out_ref_2d.reshape(shape) torch.testing.assert_close(out.float(), out_ref.float(), rtol=1e-3, atol=1e-3) scale = torch.rand(1, dtype=torch.float32, device=device) sglang_out, _ = sglang_scaled_fp8_quant(x, scale) torch_out = torch_scaled_fp8_quant(x, scale) torch.testing.assert_close( sglang_out.float(), torch_out.float(), rtol=1e-3, atol=1e-3 ) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_per_token_group_quant_8bit.py ================================================ import itertools import pytest import torch from sglang.srt.utils import is_hip _is_hip = is_hip() fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn from sgl_kernel.test_utils import ( assert_all_close_or_tiny_diff, create_per_token_group_quant_test_data, ) from sglang.jit_kernel.per_token_group_quant_8bit import ( per_token_group_quant_8bit as sglang_per_token_group_quant_8bit, ) from sglang.srt.layers.quantization.fp8_kernel import ( create_per_token_group_quant_fp8_output_scale, ) from sglang.srt.layers.quantization.fp8_kernel import ( per_token_group_quant_8bit as triton_per_token_group_quant_8bit, ) configs = list( itertools.product( [1, 4, 16, 64, 127, 128, 512, 1024, 4096, 8192], # num_tokens [128, 256, 384, 512, 1024, 1536, 1664, 2048, 4096, 7168, 16384], # hidden_dim [16, 32, 64, 128], # group_size [None], # num_ranks [fp8_type_], # dtype [ dict( column_major_scales=False, scale_tma_aligned=False, scale_ue8m0=False, fuse_silu_and_mul=False, masked_layout_mode=None, ), dict( column_major_scales=True, scale_tma_aligned=False, scale_ue8m0=False, fuse_silu_and_mul=False, masked_layout_mode=None, ), dict( column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=False, fuse_silu_and_mul=False, masked_layout_mode=None, ), dict( column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=True, fuse_silu_and_mul=False, masked_layout_mode=None, ), ], ) ) + list( itertools.product( [1, 4, 1 * 8, 4 * 8, 64 * 8, 256 * 8, 768 * 8], [2048], [128], [8, 16, 32, 48], [fp8_type_], [ dict( column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=True, fuse_silu_and_mul=True, masked_layout_mode=None, ), dict( column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=True, fuse_silu_and_mul=True, masked_layout_mode="balanced", ), dict( column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=True, fuse_silu_and_mul=True, masked_layout_mode="imbalanced", ), dict( column_major_scales=True, scale_tma_aligned=True, scale_ue8m0=True, fuse_silu_and_mul=True, masked_layout_mode="extreme", ), ], ) ) @pytest.mark.parametrize( "num_tokens, hidden_dim, group_size, num_ranks, dst_dtype, flags", configs ) def test_per_token_group_quant_with_column_major( num_tokens, hidden_dim, group_size, num_ranks, dst_dtype, flags, ): arch_major, _ = torch.cuda.get_device_capability(torch.cuda.current_device()) if flags["scale_ue8m0"] and (arch_major <= 9): pytest.skip("Only Blackwell need ue8m0 fusion") return if (flags["scale_ue8m0"] and (group_size != 128)) or ( (dst_dtype == torch.int8) and flags["column_major_scales"] ): pytest.skip() return x, masked_m = create_per_token_group_quant_test_data( num_tokens=num_tokens, hidden_dim=hidden_dim, num_ranks=num_ranks, flags=flags ) execute_kwargs = dict( x=x, masked_m=masked_m, group_size=group_size, eps=1e-10, dst_dtype=dst_dtype, **{k: v for k, v in flags.items() if k not in ["masked_layout_mode"]}, ) def _postprocess(x_q, x_s): if masked_m is not None: print(f"Mask tokens after {masked_m} to be zero") for i in range(len(masked_m)): x_q[i, masked_m[i] :, :] = 0 x_s[i, masked_m[i] :, :] = 0 return x_q, x_s x_q_triton, x_s_triton = _postprocess( *triton_per_token_group_quant_8bit(**execute_kwargs) ) fuse_silu_and_mul = False out_shape = (*x.shape[:-1], x.shape[-1] // (2 if fuse_silu_and_mul else 1)) fp8_dtype = torch.float8_e4m3fn fp8_max = torch.finfo(fp8_dtype).max fp8_min = -fp8_max x_q = torch.empty(out_shape, device=x.device, dtype=fp8_dtype) x_s = create_per_token_group_quant_fp8_output_scale( x_shape=out_shape, device=x.device, group_size=group_size, column_major_scales=False, scale_tma_aligned=False, scale_ue8m0=False, ) execute_kwargs = dict( input=x, output_q=x_q, output_s=x_s, group_size=group_size, eps=1e-10, fp8_max=fp8_max, fp8_min=fp8_min, ) x_q_sglang, x_s_sglang = _postprocess( *sglang_per_token_group_quant_8bit(**execute_kwargs) ) try: assert_all_close_or_tiny_diff(x_q_triton, x_q_sglang) torch.testing.assert_close( x_s_triton.contiguous(), x_s_sglang.contiguous(), rtol=1e-3, atol=1e-5, msg=lambda message: message + f" {x_s_triton=} {x_s_sglang=}", ) except AssertionError: print( f"{x.shape=} {x_q_triton.shape=} {x_s_triton.shape=} {x_q_sglang.shape=} {x_s_sglang.shape=}" ) print(f"{x=}") print(f"{masked_m=}") print(f"{x_q_triton=}") print(f"{x_s_triton=}") print(f"{x_q_sglang=}") print(f"{x_s_sglang=}") raise if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_pos_enc.py ================================================ import time from typing import Optional, Tuple, Union import pytest import torch import triton import triton.language as tl from sglang.jit_kernel.rope import rotary_embedding @triton.jit def burn_kernel(out_ptr, iters: tl.constexpr): pid = tl.program_id(0) x = tl.full((), pid + 1, dtype=tl.uint32) a = tl.full((), 1664525, dtype=tl.uint32) c = tl.full((), 1013904223, dtype=tl.uint32) sh = tl.full((), 13, dtype=tl.uint32) for _ in range(iters): x = x * a + c x = x ^ (x >> sh) if pid == 0: tl.store(out_ptr, x) def triton_burn(ms: float, grid=(256,)): iters = int(ms * 20000) out = torch.empty((), device="cuda", dtype=torch.uint32) burn_kernel[grid](out, iters=iters) return out def create_test_inputs( head_size, batch_size, seq_len, device, dtype, num_q_heads, num_kv_heads ): """Create test inputs.""" total_tokens = batch_size * seq_len query = torch.randn( batch_size, seq_len, num_q_heads, head_size, dtype=dtype, device=device ) key = torch.randn( batch_size, seq_len, num_kv_heads, head_size, dtype=dtype, device=device ) pos_ids = torch.randint( 0, min(seq_len * 2, 100), (total_tokens,), dtype=torch.long, device=device ) query = query.view(total_tokens, num_q_heads, head_size) key = key.view(total_tokens, num_kv_heads, head_size) return query, key, pos_ids def create_cos_sin_cache(rotary_dim, max_position_embeddings, base, dtype, device): """Create cos/sin cache for rotary embedding.""" max_pos = max_position_embeddings extended_max_pos = max(max_pos, 100) cos_sin_cache = torch.zeros( extended_max_pos, rotary_dim, dtype=dtype, device=device ) inv_freq = 1.0 / ( base ** ( torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device) / rotary_dim ) ) t = torch.arange(extended_max_pos, dtype=torch.float32, device=device) freqs = torch.outer(t, inv_freq) cos_cache = torch.cos(freqs).to(dtype) sin_cache = torch.sin(freqs).to(dtype) cos_sin_cache[:, : rotary_dim // 2] = cos_cache cos_sin_cache[:, rotary_dim // 2 :] = sin_cache return cos_sin_cache # vLLM torch native def _apply_rotary_emb( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, is_neox_style: bool, ) -> torch.Tensor: """ Args: x: [num_tokens, num_heads, head_size] cos: [num_tokens, head_size // 2] sin: [num_tokens, head_size // 2] is_neox_style: Whether to use the Neox-style or GPT-J-style rotary positional embeddings. """ cos = cos.unsqueeze(-2).to(x.dtype) sin = sin.unsqueeze(-2).to(x.dtype) if is_neox_style: x1, x2 = torch.chunk(x, 2, dim=-1) else: x1 = x[..., ::2] x2 = x[..., 1::2] o1 = x1 * cos - x2 * sin o2 = x2 * cos + x1 * sin if is_neox_style: return torch.cat((o1, o2), dim=-1) else: return torch.stack((o1, o2), dim=-1).flatten(-2) class RotaryEmbedding(torch.nn.Module): # Reference: https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/rotary_embedding.py def __init__( self, head_size: int, rotary_dim: int, max_position_embeddings: int, base: int, is_neox_style: bool, dtype: torch.dtype, ) -> None: super().__init__() self.head_size = head_size self.rotary_dim = rotary_dim self.max_position_embeddings = max_position_embeddings self.base = base self.is_neox_style = is_neox_style self.dtype = dtype cache = self._compute_cos_sin_cache() self.cos_sin_cache: torch.Tensor self.register_buffer("cos_sin_cache", cache, persistent=False) def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor: inv_freq = 1.0 / ( base ** ( torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim ) ) return inv_freq def _compute_cos_sin_cache(self) -> torch.Tensor: """Compute the cos and sin cache.""" inv_freq = self._compute_inv_freq(self.base) t = torch.arange(self.max_position_embeddings, dtype=torch.float) freqs = torch.einsum("i,j -> ij", t, inv_freq) cos = freqs.cos() sin = freqs.sin() cache = torch.cat((cos, sin), dim=-1) return cache def forward_native( self, positions: torch.Tensor, query: torch.Tensor, key: Optional[torch.Tensor] = None, offsets: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """A PyTorch-native implementation of forward().""" if offsets is not None: positions = positions + offsets positions = positions.flatten() num_tokens = positions.shape[0] cos_sin = self.cos_sin_cache.index_select(0, positions) cos, sin = cos_sin.chunk(2, dim=-1) query_shape = query.shape query = query.view(num_tokens, -1, self.head_size) query_rot = query[..., : self.rotary_dim] query_pass = query[..., self.rotary_dim :] query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style) query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape) # Modification: convert to the correct dtype query = query.to(self.dtype) if key is not None: key_shape = key.shape key = key.view(num_tokens, -1, self.head_size) key_rot = key[..., : self.rotary_dim] key_pass = key[..., self.rotary_dim :] key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style) key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) key = key.to(self.dtype) return query, key def get_torch_rotary_embedding( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, device ): """Initialize Torch Native RotaryEmbedding based on vLLM implementation.""" return RotaryEmbedding( head_size=head_size, rotary_dim=rotary_dim, max_position_embeddings=max_position_embeddings, base=base, is_neox_style=is_neox_style, dtype=dtype, ).to(device) def get_sgl_rotary_embedding( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, device ): """Initialize SglKernelRotaryEmbedding.""" try: from sgl_kernel.testing.rotary_embedding import SglKernelRotaryEmbedding except ImportError: pytest.skip( "SglKernelRotaryEmbedding is not available. Test case can be removed." ) return SglKernelRotaryEmbedding( head_size=head_size, rotary_dim=rotary_dim, max_position_embeddings=max_position_embeddings, base=base, is_neox_style=is_neox_style, dtype=dtype, ).to(device) def compare_results(jit_out, sgl_out, dtype): """Compare results between JIT and SGL implementations.""" if jit_out is None: assert sgl_out is None return assert sgl_out is not None # Check for NaN values assert not torch.isnan(jit_out).any(), "NaN in JIT results" assert not torch.isnan(sgl_out).any(), "NaN in SGL results" # Compare results atol = 1e-2 if dtype != torch.float32 else 1e-5 rtol = 1e-2 if dtype != torch.float32 else 1e-5 torch.testing.assert_close(jit_out, sgl_out, atol=atol, rtol=rtol) @pytest.mark.parametrize( "head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, device, batch_size, seq_len, num_q_heads, num_kv_heads", [ # GPT-OSS cases *[ (64, 64, 4096, 8000, True, torch.bfloat16, "cuda", bs, sl, 8, 8) for bs, sl in [(1, 1), (32, 1), (128, 1), (512, 1), (2, 512), (4, 4096)] ], # Other cases (64, 64, 32, 8000, True, torch.bfloat16, "cuda", 32, 32, 1, 1), (256, 128, 4096, 10000, True, torch.bfloat16, "cuda", 2, 512, 4, 2), (512, 128, 311, 10000, True, torch.bfloat16, "cuda", 3, 39, 4, 2), (128, 128, 2048, 10000, False, torch.bfloat16, "cuda", 2, 512, 32, 8), (128, 128, 2048, 10000, False, torch.bfloat16, "cuda", 2, 512, 16, 4), (512, 128, 311, 10000, False, torch.bfloat16, "cuda", 3, 39, 4, 2), (64, 64, 32, 8000, True, torch.float32, "cuda", 32, 32, 1, 1), (256, 128, 4096, 10000, True, torch.float32, "cuda", 2, 512, 4, 2), (512, 128, 311, 10000, True, torch.float32, "cuda", 3, 39, 4, 2), (128, 128, 2048, 10000, False, torch.float32, "cuda", 2, 512, 32, 8), (128, 128, 2048, 10000, False, torch.float32, "cuda", 2, 512, 16, 4), (512, 128, 311, 10000, False, torch.float32, "cuda", 3, 39, 4, 2), # Additional test cases for different head sizes and dtypes (64, 32, 1024, 10000, True, torch.float16, "cuda", 16, 64, 8, 4), (128, 64, 2048, 10000, True, torch.float16, "cuda", 8, 128, 16, 8), (256, 128, 4096, 10000, True, torch.float16, "cuda", 4, 256, 8, 4), ], ) @pytest.mark.parametrize( "key_is_none", [True, False], ) def test_correctness( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, device, batch_size, seq_len, num_q_heads, num_kv_heads, key_is_none, ): """Test correctness of JIT rotary embedding implementation.""" # Create inputs and caches query, key, pos_ids = create_test_inputs( head_size, batch_size, seq_len, device, dtype, num_q_heads, num_kv_heads ) cos_sin_cache = create_cos_sin_cache( rotary_dim, max_position_embeddings, base, dtype, device ) # Initialize torch kernel torch_rotary_emb = get_torch_rotary_embedding( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, device, ) torch_rotary_emb.cos_sin_cache = cos_sin_cache r = torch.randn_like(query) # Apply rotary embeddings query_jit, key_jit = query.clone(), key.clone() query_torch, key_torch = query.clone(), key.clone() stream_jit = torch.get_device_module("cuda").Stream() stream_kernel = torch.get_device_module("cuda").Stream() if key_is_none: key_jit = None key_torch = None triton_burn(100.0, grid=(1024,)) r_jit, r_torch = r.clone(), r.clone() torch.cuda.synchronize() with torch.cuda.stream(stream_jit): # Test if rotary_embedding runs on stream_jit triton_burn(100.0, grid=(1024,)) query_jit = query_jit + r_jit query_jit_out, key_jit_out = rotary_embedding( positions=pos_ids, query=query_jit, key=key_jit, head_size=head_size, cos_sin_cache=cos_sin_cache, is_neox=is_neox_style, ) with torch.cuda.stream(stream_kernel): triton_burn(100.0, grid=(1024,)) query_torch = query_torch + r_torch query_torch_out, key_torch_out = torch_rotary_emb.forward_native( positions=pos_ids, query=query_torch, key=key_torch ) torch.cuda.synchronize() compare_results(query_jit_out, query_torch_out, dtype) compare_results(key_jit_out, key_torch_out, dtype) @pytest.mark.parametrize( "head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, device, batch_size, seq_len, num_q_heads, num_kv_heads", [ # Small scale (64, 64, 4096, 8000, True, torch.bfloat16, "cuda", 1, 1, 8, 8), (64, 64, 4096, 8000, True, torch.bfloat16, "cuda", 4, 16, 8, 8), # Medium scale (64, 64, 4096, 8000, True, torch.bfloat16, "cuda", 8, 64, 8, 8), (64, 64, 4096, 8000, True, torch.bfloat16, "cuda", 16, 128, 8, 8), # Large scale (64, 64, 4096, 8000, True, torch.bfloat16, "cuda", 32, 512, 8, 8), (64, 64, 4096, 8000, True, torch.bfloat16, "cuda", 64, 1024, 8, 8), ], ) def test_performance( head_size: int, rotary_dim: int, max_position_embeddings: int, base: int, is_neox_style, dtype, device, batch_size, seq_len, num_q_heads, num_kv_heads, ): """Performance test comparing JIT and SGL implementations with accuracy validation.""" # Create inputs and caches query, key, pos_ids = create_test_inputs( head_size, batch_size, seq_len, device, dtype, num_q_heads, num_kv_heads ) cos_sin_cache = create_cos_sin_cache( rotary_dim, max_position_embeddings, base, dtype, device ) # Initialize SGL kernel sgl_rotary_emb = get_sgl_rotary_embedding( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, device, ) sgl_rotary_emb.cos_sin_cache = cos_sin_cache warmup = 3 # Warmup runs for _ in range(warmup): query_warm, key_warm = query.clone(), key.clone() rotary_embedding( positions=pos_ids, query=query_warm, key=key_warm, head_size=head_size, cos_sin_cache=cos_sin_cache, is_neox=is_neox_style, ) query_sgl_warm, key_sgl_warm = query.clone(), key.clone() sgl_rotary_emb.forward_cuda( positions=pos_ids, query=query_sgl_warm, key=key_sgl_warm ) iteration = 100 # Time JIT implementation torch.cuda.synchronize() start_time = time.time() for _ in range(iteration): query_jit, key_jit = query.clone(), key.clone() rotary_embedding( positions=pos_ids, query=query_jit, key=key_jit, head_size=head_size, cos_sin_cache=cos_sin_cache, is_neox=is_neox_style, ) torch.cuda.synchronize() jit_time = (time.time() - start_time) / iteration # Time SGL implementation torch.cuda.synchronize() start_time = time.time() for _ in range(iteration): query_sgl, key_sgl = query.clone(), key.clone() sgl_rotary_emb.forward_cuda(positions=pos_ids, query=query_sgl, key=key_sgl) torch.cuda.synchronize() sgl_time = (time.time() - start_time) / iteration # Accuracy validation during performance test # Run one more time to get outputs for comparison query_jit_final, key_jit_final = query.clone(), key.clone() query_sgl_final, key_sgl_final = query.clone(), key.clone() query_jit_out, key_jit_out = rotary_embedding( positions=pos_ids, query=query_jit_final, key=key_jit_final, head_size=head_size, cos_sin_cache=cos_sin_cache, is_neox=is_neox_style, ) query_sgl_out, key_sgl_out = sgl_rotary_emb.forward_cuda( positions=pos_ids, query=query_sgl_final, key=key_sgl_final ) # Validate accuracy compare_results(query_jit_out, query_sgl_out, dtype) compare_results(key_jit_out, key_sgl_out, dtype) # Print results total_tokens = batch_size * seq_len print( f"\nPerformance Test - Batch={batch_size}, SeqLen={seq_len}, Tokens={total_tokens}" ) print(f"JIT: {jit_time*1000:.9f}ms, SGL: {sgl_time*1000:.9f}ms") if sgl_time > 0: speedup = sgl_time / jit_time if jit_time > 0 else float("inf") print(f"Speedup (SGL/JIT): {speedup:.2f}x") assert jit_time >= 0 and sgl_time >= 0 if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_qknorm.py ================================================ import itertools import pytest import torch import triton from sglang.jit_kernel.utils import get_ci_test_range def sglang_aot_qknorm( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, ) -> None: from sgl_kernel import rmsnorm head_dim = q.shape[-1] q = q.view(-1, head_dim) k = k.view(-1, head_dim) rmsnorm(q, q_weight, out=q) rmsnorm(k, k_weight, out=k) def sglang_jit_qknorm( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, ) -> None: from sglang.jit_kernel.norm import fused_inplace_qknorm fused_inplace_qknorm(q, k, q_weight, k_weight) def flashinfer_qknorm( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, ) -> None: from flashinfer.norm import rmsnorm rmsnorm(q, q_weight, out=q) rmsnorm(k, k_weight, out=k) @torch.compile() def torch_impl_qknorm( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, eps: float = 1e-6, ) -> None: q_mean = q.float().pow(2).mean(dim=-1, keepdim=True) k_mean = k.float().pow(2).mean(dim=-1, keepdim=True) q_norm = (q_mean + eps).rsqrt() k_norm = (k_mean + eps).rsqrt() q.copy_(q.float() * q_norm * q_weight.float()) k.copy_(k.float() * k_norm * k_weight.float()) BS_LIST = [2**n for n in range(0, 14)] BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109]) N_K_LIST = get_ci_test_range([2, 4], [2, 4]) N_Q_LIST = get_ci_test_range([8, 16], [8, 16]) HEAD_DIM_LIST = get_ci_test_range([64, 128, 256, 512, 1024], [64, 256, 1024]) DEVICE = "cuda" DTYPE = torch.bfloat16 # NOTE(dark): sgl_kernel use flashinfer template, which is bitwise identical to flashinfer impl. # However, sgl-jit-kernel, flashinfer, torch_impl, may have small numerical differences. # so we allow a small rel/abs tolerance in correctness test. @pytest.mark.parametrize( "batch_size,n_k,n_q,head_dim", list(itertools.product(BS_LIST, N_K_LIST, N_Q_LIST, HEAD_DIM_LIST)), ) def test_qknorm(batch_size: int, n_k: int, n_q: int, head_dim: int) -> None: q = torch.randn(batch_size, n_q, head_dim, device=DEVICE, dtype=DTYPE) k = torch.randn(batch_size, n_k, head_dim, device=DEVICE, dtype=DTYPE) q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE) k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE) q_k_aot = (q.clone(), k.clone()) q_k_jit = (q.clone(), k.clone()) sglang_aot_qknorm(q_k_aot[0], q_k_aot[1], q_weight, k_weight) sglang_jit_qknorm(q_k_jit[0], q_k_jit[1], q_weight, k_weight) triton.testing.assert_close(q_k_aot[0], q_k_jit[0], atol=1e-2, rtol=1e-2) triton.testing.assert_close(q_k_aot[1], q_k_jit[1], atol=1e-2, rtol=1e-2) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_qknorm_across_heads.py ================================================ import itertools import pytest import torch import triton from sglang.jit_kernel.utils import get_ci_test_range def sglang_jit_qknorm_across_heads( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, ) -> None: from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads fused_inplace_qknorm_across_heads(q, k, q_weight, k_weight) def sglang_aot_qknorm_across_heads( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, ) -> None: from sgl_kernel import rmsnorm rmsnorm(q, q_weight, out=q) rmsnorm(k, k_weight, out=k) @torch.compile() def torch_impl_qknorm_across_heads( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, eps: float = 1e-6, ) -> None: q_mean = q.float().pow(2).mean(dim=-1, keepdim=True) k_mean = k.float().pow(2).mean(dim=-1, keepdim=True) q_norm = (q_mean + eps).rsqrt() k_norm = (k_mean + eps).rsqrt() q.copy_(q.float() * q_norm * q_weight.float()) k.copy_(k.float() * k_norm * k_weight.float()) BS_LIST = [2**n for n in range(0, 14)] BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109]) HIDDEN_DIM_LIST = get_ci_test_range([512, 1024, 2048, 4096], [512, 2048, 4096]) DEVICE = "cuda" DTYPE = torch.bfloat16 @pytest.mark.parametrize( "batch_size,hidden_dim", list(itertools.product(BS_LIST, HIDDEN_DIM_LIST)), ) def test_qknorm_across_heads(batch_size: int, hidden_dim: int) -> None: q = torch.randn(batch_size, hidden_dim, device=DEVICE, dtype=DTYPE) k = torch.randn(batch_size, hidden_dim, device=DEVICE, dtype=DTYPE) q_weight = torch.randn(hidden_dim, device=DEVICE, dtype=DTYPE) k_weight = torch.randn(hidden_dim, device=DEVICE, dtype=DTYPE) q_k_jit = (q.clone(), k.clone()) q_k_aot = (q.clone(), k.clone()) sglang_jit_qknorm_across_heads(q_k_jit[0], q_k_jit[1], q_weight, k_weight) sglang_aot_qknorm_across_heads(q_k_aot[0], q_k_aot[1], q_weight, k_weight) triton.testing.assert_close(q_k_jit[0], q_k_aot[0], atol=1e-2, rtol=1e-2) triton.testing.assert_close(q_k_jit[1], q_k_aot[1], atol=1e-2, rtol=1e-2) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_qwen_image_modulation.py ================================================ import pytest import torch import triton from sglang.jit_kernel.diffusion.triton.norm import norm_infer from sglang.jit_kernel.diffusion.triton.scale_shift import ( fuse_layernorm_scale_shift_gate_select01_kernel, fuse_residual_layernorm_scale_shift_gate_select01_kernel, fuse_scale_shift_gate_select01_kernel, ) from sglang.jit_kernel.utils import get_ci_test_range DEVICE = "cuda" DTYPES = get_ci_test_range( [torch.float16, torch.bfloat16, torch.float32], [torch.float16, torch.bfloat16] ) BATCH_SIZES = get_ci_test_range([1, 2, 4], [1, 2]) SEQ_LENS = get_ci_test_range([6, 33, 128, 257], [6, 128]) HIDDEN_SIZES = get_ci_test_range([512, 1024, 1536, 3072], [512, 3072]) EPS = 1e-6 def _tol(dtype: torch.dtype) -> tuple[float, float]: if dtype == torch.float32: return 1e-5, 1e-5 return 5e-2, 5e-2 def _make_modulation_tensors(batch_size: int, hidden_size: int, dtype: torch.dtype): scale0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype) shift0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype) gate0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype) scale1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype) shift1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype) gate1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype) return scale0, shift0, gate0, scale1, shift1, gate1 def _baseline_select01_modulation( x: torch.Tensor, weight: torch.Tensor | None, bias: torch.Tensor | None, scale0: torch.Tensor, shift0: torch.Tensor, gate0: torch.Tensor, scale1: torch.Tensor, shift1: torch.Tensor, gate1: torch.Tensor, index: torch.Tensor, eps: float, ): normalized = norm_infer( x.view(-1, x.shape[-1]), weight, bias, eps=eps, is_rms_norm=False, ).view_as(x) output, gate_out = fuse_scale_shift_gate_select01_kernel( normalized, scale0=scale0, shift0=shift0, gate0=gate0, scale1=scale1, shift1=shift1, gate1=gate1, index=index, ) return output, gate_out def _baseline_residual_select01_modulation( x: torch.Tensor, residual: torch.Tensor, residual_gate: torch.Tensor, weight: torch.Tensor | None, bias: torch.Tensor | None, scale0: torch.Tensor, shift0: torch.Tensor, gate0: torch.Tensor, scale1: torch.Tensor, shift1: torch.Tensor, gate1: torch.Tensor, index: torch.Tensor, eps: float, ): residual_out = residual + residual_gate * x normalized = norm_infer( residual_out.view(-1, residual_out.shape[-1]), weight, bias, eps=eps, is_rms_norm=False, ).view_as(residual_out) output, gate_out = fuse_scale_shift_gate_select01_kernel( normalized, scale0=scale0, shift0=shift0, gate0=gate0, scale1=scale1, shift1=shift1, gate1=gate1, index=index, ) return output, residual_out, gate_out @pytest.fixture(autouse=True) def cuda_setup(): if not torch.cuda.is_available(): pytest.skip("CUDA required") torch.cuda.manual_seed(0) @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("batch_size", BATCH_SIZES) @pytest.mark.parametrize("seq_len", SEQ_LENS) @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) def test_fused_layernorm_scale_shift_gate_select01( dtype, batch_size, seq_len, hidden_size ): x = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE, dtype=dtype) weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype) bias = torch.randn(hidden_size, device=DEVICE, dtype=dtype) index = torch.randint(0, 2, (batch_size, seq_len), device=DEVICE, dtype=torch.int32) scale0, shift0, gate0, scale1, shift1, gate1 = _make_modulation_tensors( batch_size, hidden_size, dtype ) out_ref, gate_ref = _baseline_select01_modulation( x, weight, bias, scale0, shift0, gate0, scale1, shift1, gate1, index, EPS, ) out_fused, gate_fused = fuse_layernorm_scale_shift_gate_select01_kernel( x.contiguous(), weight=weight, bias=bias, scale0=scale0, shift0=shift0, gate0=gate0, scale1=scale1, shift1=shift1, gate1=gate1, index=index, eps=EPS, ) atol, rtol = _tol(dtype) triton.testing.assert_close(out_ref, out_fused, atol=atol, rtol=rtol) triton.testing.assert_close(gate_ref, gate_fused, atol=atol, rtol=rtol) @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("batch_size", BATCH_SIZES) @pytest.mark.parametrize("seq_len", SEQ_LENS) @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) def test_fused_residual_layernorm_scale_shift_gate_select01( dtype, batch_size, seq_len, hidden_size ): x = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE, dtype=dtype) residual = torch.randn_like(x) residual_gate = torch.randn_like(x) weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype) bias = torch.randn(hidden_size, device=DEVICE, dtype=dtype) index = torch.randint(0, 2, (batch_size, seq_len), device=DEVICE, dtype=torch.int32) scale0, shift0, gate0, scale1, shift1, gate1 = _make_modulation_tensors( batch_size, hidden_size, dtype ) out_ref, residual_ref, gate_ref = _baseline_residual_select01_modulation( x, residual, residual_gate, weight, bias, scale0, shift0, gate0, scale1, shift1, gate1, index, EPS, ) out_fused, residual_fused, gate_fused = ( fuse_residual_layernorm_scale_shift_gate_select01_kernel( x.contiguous(), residual=residual.contiguous(), residual_gate=residual_gate.contiguous(), weight=weight, bias=bias, scale0=scale0, shift0=shift0, gate0=gate0, scale1=scale1, shift1=shift1, gate1=gate1, index=index, eps=EPS, ) ) atol, rtol = _tol(dtype) triton.testing.assert_close(out_ref, out_fused, atol=atol, rtol=rtol) triton.testing.assert_close(residual_ref, residual_fused, atol=atol, rtol=rtol) triton.testing.assert_close(gate_ref, gate_fused, atol=atol, rtol=rtol) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_renorm.py ================================================ # Adapted from https://github.com/flashinfer-ai/flashinfer/blob/main/tests/test_sampling.py # and /sgl-workspace/sglang/sgl-kernel/tests/test_sampling.py import pytest import sgl_kernel import torch @pytest.mark.parametrize("batch_size", [1, 99, 989]) @pytest.mark.parametrize("vocab_size", [111, 32000, 128256]) @pytest.mark.parametrize("k", [10, 100, 500]) def test_top_k_renorm_probs(batch_size, vocab_size, k): """Test top_k_renorm_probs kernel for correctness. This test validates that the kernel correctly: 1. Identifies the top-k probabilities 2. Masks out non-top-k values 3. Renormalizes the remaining probabilities to sum to 1 """ if k > vocab_size: pytest.skip("k should be less than vocab_size") torch.manual_seed(42) pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0") normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True) sorted_prob, _ = torch.sort(normalized_prob, descending=True) pivot = sorted_prob[:, k - 1] mask = (normalized_prob >= pivot.unsqueeze(-1)).int() renorm_prob_ground_truth = normalized_prob.clone() renorm_prob_ground_truth[mask == 0] = 0 renorm_prob_ground_truth = renorm_prob_ground_truth / renorm_prob_ground_truth.sum( dim=-1, keepdim=True ) renorm_prob = sgl_kernel.top_k_renorm_prob(normalized_prob, k) for i in range(batch_size): torch.testing.assert_close( renorm_prob_ground_truth[i], renorm_prob[i], rtol=1e-3, atol=1e-3, ) @pytest.mark.parametrize("batch_size", [1, 99, 989]) @pytest.mark.parametrize("vocab_size", [111, 32000, 128256]) @pytest.mark.parametrize("p", [0.1, 0.5, 0.9]) def test_top_p_renorm_probs(batch_size, vocab_size, p): """Test top_p_renorm_probs kernel for correctness. This test validates that the kernel correctly: 1. Computes the cumulative probability distribution 2. Identifies tokens in the top-p threshold 3. Masks out tokens outside the threshold 4. Renormalizes the remaining probabilities to sum to 1 """ torch.manual_seed(42) pre_norm_prob = torch.rand(batch_size, vocab_size, device="cuda:0") normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True) sorted_prob, indices = torch.sort(normalized_prob, descending=False) cdf = torch.cumsum(sorted_prob, dim=-1) mask = torch.zeros(batch_size, vocab_size, dtype=torch.int32, device="cuda:0") mask.scatter_add_(1, indices, (cdf >= (1 - p)).int()) renorm_prob_ground_truth = normalized_prob.clone() renorm_prob_ground_truth[mask == 0] = 0 renorm_prob_ground_truth = renorm_prob_ground_truth / renorm_prob_ground_truth.sum( dim=-1, keepdim=True ) renorm_prob = sgl_kernel.top_p_renorm_prob(normalized_prob, p) torch.testing.assert_close( renorm_prob_ground_truth, renorm_prob, rtol=1e-3, atol=1e-3, ) @pytest.mark.parametrize("batch_size", [1, 99, 989]) @pytest.mark.parametrize("vocab_size", [111, 32000, 128256]) @pytest.mark.parametrize("k", [10, 100, 500]) @pytest.mark.parametrize("neginf_input", [False, True]) def test_top_k_mask_logits(batch_size, vocab_size, k, neginf_input): """Test top_k_mask_logits kernel for correctness. This test validates that the kernel correctly: 1. Identifies the top-k logits 2. Masks non-top-k values to -inf 3. Preserves the top-k values 4. Handles negative infinity inputs gracefully The test verifies correctness by comparing softmax(top_k_mask_logits(logits)) with top_k_renorm_prob(probs), which should be equivalent. """ if k > vocab_size: pytest.skip("k should be less than vocab_size") torch.manual_seed(42) logits = torch.randn(batch_size, vocab_size, device="cuda:0") * 5 if neginf_input: # Randomly assign some logits to -inf to test edge cases num_neginf = torch.randint(1, vocab_size * batch_size, (1,)).item() idxs = torch.randperm(batch_size * vocab_size, device="cuda:0")[:num_neginf] logits[idxs // vocab_size, idxs % vocab_size] = -float("inf") probs = torch.softmax(logits, dim=-1) masked_logits = sgl_kernel.top_k_mask_logits(logits, k) renormed_probs = torch.softmax(masked_logits, dim=-1) renormed_probs_ref = sgl_kernel.top_k_renorm_prob(probs, k) torch.testing.assert_close( renormed_probs, renormed_probs_ref, rtol=1e-3, atol=1e-3, ) if __name__ == "__main__": pytest.main([__file__]) ================================================ FILE: python/sglang/jit_kernel/tests/test_rmsnorm.py ================================================ import itertools import pytest import torch import triton from sglang.jit_kernel.utils import get_ci_test_range def sglang_jit_rmsnorm(input: torch.Tensor, weight: torch.Tensor) -> None: from sglang.jit_kernel.norm import rmsnorm rmsnorm(input, weight, output=input) def flashinfer_rmsnorm(input: torch.Tensor, weight: torch.Tensor) -> None: from flashinfer.norm import rmsnorm rmsnorm(input, weight, out=input) BS_LIST = [2**n for n in range(0, 14)] BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109]) HIDDEN_SIZE_LIST = get_ci_test_range( [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192], [512, 2048, 8192], ) DEVICE = "cuda" DTYPE = torch.bfloat16 @pytest.mark.parametrize( "batch_size,hidden_size", list(itertools.product(BS_LIST, HIDDEN_SIZE_LIST)) ) def test_rmsnorm(batch_size: int, hidden_size: int) -> None: input = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE) weight = torch.randn(hidden_size, device=DEVICE, dtype=DTYPE) input_sglang = input.clone() input_flashinfer = input.clone() sglang_jit_rmsnorm(input_sglang, weight) flashinfer_rmsnorm(input_flashinfer, weight) triton.testing.assert_close(input_sglang, input_flashinfer, atol=1e-2, rtol=1e-2) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_rope.py ================================================ import pytest import torch import triton from sglang.jit_kernel.utils import get_ci_test_range DEVICE = "cuda" DTYPE = torch.bfloat16 MAX_SEQ_LEN = 131072 # common seq length ROPE_BASE = 10000.0 CACHE_SIZE = 1024 * 128 def create_cos_sin_cache( rotary_dim: int, max_position: int = MAX_SEQ_LEN, base: float = ROPE_BASE, ) -> torch.Tensor: """Create cos/sin cache compatible with SGLang layout: [max_pos, rotary_dim].""" inv_freq = 1.0 / ( base ** ( torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=DEVICE) / rotary_dim ) ) t = torch.arange(max_position, dtype=torch.float32, device=DEVICE) freqs = torch.einsum("i,j->ij", t, inv_freq) cos = freqs.cos() sin = freqs.sin() cache = torch.cat((cos, sin), dim=-1) # [max_pos, rotary_dim] return cache # --------------------------------------------------------------------------- # Implementation wrappers # --------------------------------------------------------------------------- def sglang_jit_rope( q: torch.Tensor, k: torch.Tensor, cos_sin_cache: torch.Tensor, positions: torch.Tensor, is_neox: bool, ) -> None: from sglang.jit_kernel.rope import apply_rope_inplace apply_rope_inplace(q, k, cos_sin_cache, positions, is_neox=is_neox) def flashinfer_rope( q: torch.Tensor, k: torch.Tensor, cos_sin_cache: torch.Tensor, positions: torch.Tensor, is_neox: bool, ) -> None: from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace head_size = q.shape[-1] # flashinfer expects [nnz, num_heads * head_size] q_2d = q.view(q.shape[0], -1) k_2d = k.view(k.shape[0], -1) apply_rope_with_cos_sin_cache_inplace( positions=positions, query=q_2d, key=k_2d, head_size=head_size, cos_sin_cache=cos_sin_cache, is_neox=is_neox, ) def torch_impl_rope( q: torch.Tensor, k: torch.Tensor, cos_sin_cache: torch.Tensor, positions: torch.Tensor, is_neox: bool, ) -> None: # TODO: implement a pure-PyTorch reference for extra coverage pass # --------------------------------------------------------------------------- # Test parameters # --------------------------------------------------------------------------- BS_LIST = [2**x for x in range(12)] BS_LIST += [x + 1 for x in BS_LIST] # odd sizes to stress non-aligned paths BS_LIST = get_ci_test_range(BS_LIST, [1, 129, 2048, 2049]) NUM_KV_HEADS_LIST = get_ci_test_range([1, 2, 8], [1, 8]) GQA_RATIO = get_ci_test_range([1, 4, 8], [1, 8]) ROPE_DIM_LIST = get_ci_test_range([64, 128, 256, 512], [64, 256]) IS_NEOX_LIST = [False, True] DTYPE_LIST = get_ci_test_range( [torch.bfloat16, torch.float16], [torch.bfloat16, torch.float16] ) PARTIAL_ROPE_DIM_LIST = get_ci_test_range([64, 80, 96, 128], [64, 96]) HEAD_DIM_LIST = get_ci_test_range([64, 128, 256], [64, 256]) @pytest.mark.parametrize("batch_size", BS_LIST) @pytest.mark.parametrize("gqa_ratio", GQA_RATIO) @pytest.mark.parametrize("num_kv_heads", NUM_KV_HEADS_LIST) @pytest.mark.parametrize("rope_dim", ROPE_DIM_LIST) @pytest.mark.parametrize("is_neox", IS_NEOX_LIST) @pytest.mark.parametrize("dtype", DTYPE_LIST) def test_rope( batch_size: int, gqa_ratio: int, num_kv_heads: int, rope_dim: int, is_neox: bool, dtype: torch.dtype, ) -> None: num_qo_heads = num_kv_heads * gqa_ratio q = torch.randn(batch_size, num_qo_heads, rope_dim, device=DEVICE, dtype=dtype) k = torch.randn(batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=dtype) positions = torch.randint( 0, MAX_SEQ_LEN, (batch_size,), device=DEVICE, dtype=torch.int64 ) cos_sin_cache = create_cos_sin_cache(rope_dim) q_fi, k_fi = q.clone(), k.clone() q_jit, k_jit = q.clone(), k.clone() flashinfer_rope(q_fi, k_fi, cos_sin_cache, positions, is_neox) sglang_jit_rope(q_jit, k_jit, cos_sin_cache, positions, is_neox) atol = rtol = 1e-2 triton.testing.assert_close(q_fi, q_jit, atol=atol, rtol=rtol) triton.testing.assert_close(k_fi, k_jit, atol=atol, rtol=rtol) @pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) def test_rope_position_dtypes(dtype: torch.dtype) -> None: """Ensure both int32 and int64 position tensors work correctly.""" batch_size, num_qo_heads, num_kv_heads, rope_dim = 16384, 16, 2, 128 is_neox = True q = torch.randn(batch_size, num_qo_heads, rope_dim, device=DEVICE, dtype=DTYPE) k = torch.randn(batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=DTYPE) positions = torch.randint(0, MAX_SEQ_LEN, (batch_size,), device=DEVICE, dtype=dtype) cos_sin_cache = create_cos_sin_cache(rope_dim) q_fi, k_fi = q.clone(), k.clone() q_jit, k_jit = q.clone(), k.clone() flashinfer_rope(q_fi, k_fi, cos_sin_cache, positions.long(), is_neox) sglang_jit_rope(q_jit, k_jit, cos_sin_cache, positions, is_neox) atol = rtol = 1e-2 triton.testing.assert_close(q_fi, q_jit, atol=atol, rtol=rtol) triton.testing.assert_close(k_fi, k_jit, atol=atol, rtol=rtol) @pytest.mark.parametrize("batch_size", BS_LIST) @pytest.mark.parametrize("is_neox", IS_NEOX_LIST) @pytest.mark.parametrize("rope_dim", PARTIAL_ROPE_DIM_LIST) @pytest.mark.parametrize("head_dim", HEAD_DIM_LIST) def test_partial_rope(batch_size: int, is_neox: bool, rope_dim: int, head_dim: int): if head_dim < rope_dim: pytest.skip("Invalid config: head_dim must be >= rope_dim.") num_qo_heads, num_kv_heads = 8, 2 q = torch.randn(batch_size, num_qo_heads, head_dim, device=DEVICE, dtype=DTYPE) k = torch.randn(batch_size, num_kv_heads, head_dim, device=DEVICE, dtype=DTYPE) positions = torch.randint(0, MAX_SEQ_LEN, (batch_size,), device=DEVICE) cos_sin_cache = create_cos_sin_cache(rope_dim) q_fi, k_fi = q.clone(), k.clone() q_jit, k_jit = q.clone(), k.clone() rope = ..., slice(rope_dim) # NOTE: flashinfer by default apply to first rope_dim flashinfer_rope(q_fi, k_fi, cos_sin_cache, positions.long(), is_neox) sglang_jit_rope(q_jit[rope], k_jit[rope], cos_sin_cache, positions, is_neox) atol = rtol = 1e-2 triton.testing.assert_close(q_fi, q_jit, atol=atol, rtol=rtol) triton.testing.assert_close(k_fi, k_jit, atol=atol, rtol=rtol) @pytest.mark.parametrize("batch_size", BS_LIST) @pytest.mark.parametrize("gqa_ratio", GQA_RATIO) @pytest.mark.parametrize("num_kv_heads", NUM_KV_HEADS_LIST) @pytest.mark.parametrize("rope_dim", ROPE_DIM_LIST) @pytest.mark.parametrize("is_neox", IS_NEOX_LIST) def test_fused_rope_store( batch_size: int, gqa_ratio: int, num_kv_heads: int, rope_dim: int, is_neox: bool, ) -> None: """Test fused RoPE + KV cache store against separate RoPE + manual store.""" from sglang.jit_kernel.rope import apply_rope_inplace_with_kvcache num_qo_heads = num_kv_heads * gqa_ratio dtype = DTYPE q = torch.randn(batch_size, num_qo_heads, rope_dim, device=DEVICE, dtype=dtype) k = torch.randn(batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=dtype) v = torch.randn(batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=dtype) positions = torch.randint( 0, MAX_SEQ_LEN, (batch_size,), device=DEVICE, dtype=torch.int64 ) out_loc = torch.randperm(CACHE_SIZE, device=DEVICE, dtype=torch.int64)[:batch_size] cos_sin_cache = create_cos_sin_cache(rope_dim) row_size = num_kv_heads * rope_dim k_cache_ref = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=dtype) v_cache_ref = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=dtype) k_cache_fused = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=dtype) v_cache_fused = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=dtype) # --- reference: separate RoPE then manual scatter --- q_ref, k_ref = q.clone(), k.clone() flashinfer_rope(q_ref, k_ref, cos_sin_cache, positions, is_neox) k_cache_ref[out_loc] = k_ref.view(batch_size, -1) v_cache_ref[out_loc] = v.view(batch_size, -1) # --- fused kernel --- q_fused, k_fused = q.clone(), k.clone() v_fused = v.clone() apply_rope_inplace_with_kvcache( q_fused, k_fused, v_fused, k_cache_fused, v_cache_fused, cos_sin_cache, positions, out_loc, is_neox=is_neox, ) atol = rtol = 1e-2 # q should match RoPE-only result triton.testing.assert_close(q_ref, q_fused, atol=atol, rtol=rtol) # k_cache should contain the rotated k triton.testing.assert_close( k_cache_ref[out_loc], k_cache_fused[out_loc], atol=atol, rtol=rtol ) # v_cache should be an exact copy assert torch.all(v_cache_ref[out_loc] == v_cache_fused[out_loc]), "v_cache mismatch" if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_store_cache.py ================================================ import itertools import pytest import torch from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache from sglang.jit_kernel.utils import get_ci_test_range BS_LIST = [2**n for n in range(0, 15)] BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 16399]) HIDDEN_DIMS = get_ci_test_range( [64, 128, 256, 512, 1024, 96, 98, 100], [64, 512, 1024, 98] ) CACHE_SIZE = 1024 * 1024 DTYPE = torch.bfloat16 DEVICE = "cuda" @pytest.mark.parametrize( "batch_size,element_dim", list(itertools.product(BS_LIST, HIDDEN_DIMS)), ) def test_store_cache(batch_size: int, element_dim: int) -> None: k = torch.randn((batch_size, element_dim), dtype=DTYPE, device=DEVICE) v = torch.randn((batch_size, element_dim), dtype=DTYPE, device=DEVICE) k_cache = torch.randn((CACHE_SIZE, element_dim), dtype=DTYPE, device=DEVICE) v_cache = torch.randn((CACHE_SIZE, element_dim), dtype=DTYPE, device=DEVICE) indices = torch.randperm(CACHE_SIZE, device=DEVICE)[:batch_size] # AOT store cache store_cache(k, v, k_cache, v_cache, indices) assert torch.all(k_cache[indices] == k) assert torch.all(v_cache[indices] == v) # Smaller subset for targeted tests below REPR_BS = get_ci_test_range([1, 7, 128], [1, 128]) REPR_DIMS = get_ci_test_range([64, 128, 512, 1024, 96], [64, 1024, 96]) SMALL_CACHE = 4096 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @pytest.mark.parametrize( "batch_size,element_dim", list(itertools.product(REPR_BS, REPR_DIMS)), ) def test_store_cache_dtypes( batch_size: int, element_dim: int, dtype: torch.dtype ) -> None: k = torch.randn((batch_size, element_dim), dtype=dtype, device=DEVICE) v = torch.randn((batch_size, element_dim), dtype=dtype, device=DEVICE) k_cache = torch.randn((SMALL_CACHE, element_dim), dtype=dtype, device=DEVICE) v_cache = torch.randn((SMALL_CACHE, element_dim), dtype=dtype, device=DEVICE) indices = torch.randperm(SMALL_CACHE, device=DEVICE)[:batch_size] store_cache(k, v, k_cache, v_cache, indices) assert torch.all(k_cache[indices] == k) assert torch.all(v_cache[indices] == v) @pytest.mark.parametrize( "batch_size,element_dim", list(itertools.product(REPR_BS, REPR_DIMS)), ) def test_store_cache_int32_indices(batch_size: int, element_dim: int) -> None: k = torch.randn((batch_size, element_dim), dtype=DTYPE, device=DEVICE) v = torch.randn((batch_size, element_dim), dtype=DTYPE, device=DEVICE) k_cache = torch.randn((SMALL_CACHE, element_dim), dtype=DTYPE, device=DEVICE) v_cache = torch.randn((SMALL_CACHE, element_dim), dtype=DTYPE, device=DEVICE) # int32 indices exercise a different CUDA template instantiation than default int64 indices = torch.randperm(SMALL_CACHE, device=DEVICE)[:batch_size].to(torch.int32) store_cache(k, v, k_cache, v_cache, indices) assert torch.all(k_cache[indices.long()] == k) assert torch.all(v_cache[indices.long()] == v) def _valid_num_splits(element_dim: int, dtype: torch.dtype) -> list: """Return the list of valid num_split values for a given element_dim/dtype.""" row_bytes = element_dim * dtype.itemsize splits = [1] if row_bytes % (2 * 128) == 0: splits.append(2) if row_bytes % (4 * 128) == 0: splits.append(4) return splits _NUM_SPLIT_CASES = [ (_dim, _ns, _dtype) for _dtype in [torch.float16, torch.bfloat16, torch.float32] for _dim in REPR_DIMS for _ns in _valid_num_splits(_dim, _dtype) ] @pytest.mark.parametrize("element_dim,num_split,dtype", _NUM_SPLIT_CASES) def test_store_cache_num_split( element_dim: int, num_split: int, dtype: torch.dtype ) -> None: batch_size = 128 k = torch.randn((batch_size, element_dim), dtype=dtype, device=DEVICE) v = torch.randn((batch_size, element_dim), dtype=dtype, device=DEVICE) k_cache = torch.randn((SMALL_CACHE, element_dim), dtype=dtype, device=DEVICE) v_cache = torch.randn((SMALL_CACHE, element_dim), dtype=dtype, device=DEVICE) indices = torch.randperm(SMALL_CACHE, device=DEVICE)[:batch_size] # Verify each num_split kernel path (1, 2, 4) produces correct results store_cache(k, v, k_cache, v_cache, indices, num_split=num_split) assert torch.all(k_cache[indices] == k) assert torch.all(v_cache[indices] == v) def test_can_use_store_cache() -> None: assert can_use_store_cache(128) assert can_use_store_cache(256) assert can_use_store_cache(1024) assert can_use_store_cache(2048) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/tests/test_timestep_embedding.py ================================================ import os import numpy as np import pytest import torch try: import tabulate except Exception: tabulate = None from sglang.jit_kernel.timestep_embedding import ( timestep_embedding as timestep_embedding_cuda, ) from sglang.jit_kernel.utils import get_ci_test_range CORRECTNESS_BATCH_SIZES = get_ci_test_range( [1, 2, 8, 128, 256, 512, 1536, 2048, 4096, 11008, 16384], [1, 128, 2048, 16384], ) CORRECTNESS_DIMS = get_ci_test_range( [32, 128, 256, 512, 1536, 2048, 4096, 8192], [32, 512, 8192], ) DIFFUSERS_BATCH_SIZES = get_ci_test_range( [1, 2, 8, 128, 256, 512, 1536, 2048, 16384], [1, 512, 16384], ) DIFFUSERS_DIMS = get_ci_test_range([32, 256, 512, 1536, 8192], [32, 512, 8192]) DTYPES = get_ci_test_range( [torch.float16, torch.bfloat16, torch.float32], [torch.float16, torch.bfloat16], ) SCALES = get_ci_test_range([1, 0.01], [1, 0.01]) def get_timestep_embedding_reference( timesteps: torch.Tensor, dim: int, *, flip_sin_to_cos: bool = False, downscale_freq_shift: float = 1, scale: float = 1, max_period: int = 10000, ): assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" timesteps = timesteps.to(torch.float32) half_dim = dim // 2 exponent = -torch.log( torch.tensor(max_period, dtype=torch.float32, device=timesteps.device) ) * torch.arange( start=0, end=half_dim, dtype=torch.float32, device=timesteps.device ) exponent = exponent / (half_dim - downscale_freq_shift) emb = torch.exp(exponent) emb = timesteps[:, None].float() * emb[None, :] emb = scale * emb emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) if flip_sin_to_cos: emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) if dim % 2 == 1: emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) return emb @pytest.mark.parametrize("batch_size", CORRECTNESS_BATCH_SIZES) @pytest.mark.parametrize("dim", CORRECTNESS_DIMS) @pytest.mark.parametrize("dtype", DTYPES) def test_timestep_embedding_correctness_with_sgld(batch_size, dim, dtype): device = "cuda" t = torch.randint(low=0, high=1000, size=(batch_size,), device=device).to(dtype) torch_output = get_timestep_embedding_reference( t, dim, flip_sin_to_cos=True, downscale_freq_shift=0 ) cuda_output = timestep_embedding_cuda( t, dim, flip_sin_to_cos=True, downscale_freq_shift=0 ) torch.testing.assert_close(torch_output, cuda_output, atol=1e-3, rtol=1e-3) @pytest.mark.parametrize("batch_size", DIFFUSERS_BATCH_SIZES) @pytest.mark.parametrize("dim", DIFFUSERS_DIMS) @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("flip_sin_to_cos", [False, True]) @pytest.mark.parametrize("downscale_freq_shift", [0, 1]) @pytest.mark.parametrize("scale", SCALES) def test_timestep_embedding_correctness_with_diffusers( batch_size, dim, flip_sin_to_cos, downscale_freq_shift, scale, dtype ): device = "cuda" t = torch.randint(low=0, high=1000, size=(batch_size,), device=device).to(dtype) torch_output = get_timestep_embedding_reference( t, dim, flip_sin_to_cos=flip_sin_to_cos, downscale_freq_shift=downscale_freq_shift, scale=scale, max_period=10000, ) cuda_output = timestep_embedding_cuda( t, dim, flip_sin_to_cos=flip_sin_to_cos, downscale_freq_shift=downscale_freq_shift, scale=scale, max_period=10000, ) torch.testing.assert_close(torch_output, cuda_output, atol=1e-3, rtol=1e-3) def test_timestep_embedding_perf(): if os.environ.get("SGLANG_RUN_JIT_KERNEL_PERF_TESTS") != "1": pytest.skip("Perf test disabled by default") if tabulate is None: pytest.skip("Optional dependency 'tabulate' is not installed") NUM_BATCH = [1, 2, 8, 63, 256, 512, 613, 1024, 1536] NUM_DIM = [32, 64, 128, 256, 512, 1024, 2048, 4096] def perf_kernel_fn(kernel_fn: callable, *args, **kwargs): warmup_times = 4 repeat_times = 20 start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) for _ in range(warmup_times): output_fn = kernel_fn(*args, **kwargs) torch.cuda.synchronize() start.record() for _ in range(repeat_times): output_fn = kernel_fn(*args, **kwargs) end.record() end.synchronize() return start.elapsed_time(end) / repeat_times device = "cuda" results = [] cuda_speedups = [] for B in NUM_BATCH: for dim in NUM_DIM: t = torch.linspace(0, max(100000, B), steps=B, device=device).to( torch.float32 ) time_torch = perf_kernel_fn(get_timestep_embedding_reference, t, dim) time_cuda = perf_kernel_fn(timestep_embedding_cuda, t, dim) speedup_cuda = time_torch / time_cuda results.append( { "Batch Size": B, "Dimension": dim, "Torch Time (ms)": time_torch, "CUDA Time (ms)": time_cuda, "Speedup (CUDA)": speedup_cuda, } ) cuda_speedups.append(speedup_cuda) print("=== Timestep Embedding Benchmark Results ===") print( tabulate.tabulate( results, headers="keys", tablefmt="fancy_grid", floatfmt=(".0f", ".0f", ".6f", ".6f", ".5f"), ) ) print(f"Average Speedup(cuda): {np.mean(cuda_speedups):.4f}") if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/jit_kernel/timestep_embedding.py ================================================ from __future__ import annotations from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_timestep_embedding_module(dtype: torch.dtype) -> Module: args = make_cpp_args(dtype) return load_jit( "timestep_embedding", *args, cuda_files=["diffusion/timestep_embedding.cuh"], cuda_wrappers=[("timestep_embedding", f"timestep_embedding<{args}>")], ) def timestep_embedding( t: torch.Tensor, dim: int, flip_sin_to_cos: bool = False, downscale_freq_shift: float = 0.0, scale: float = 1, max_period: int = 10000, dtype: torch.dtype = torch.float32, ) -> torch.Tensor: if t.dtype not in (torch.float16, torch.bfloat16, torch.float32): t = t.to(dtype) output = torch.empty((t.shape[0], dim), dtype=torch.float32, device=t.device) module = _jit_timestep_embedding_module(t.dtype) module.timestep_embedding( t, output, dim, flip_sin_to_cos, float(downscale_freq_shift), float(scale), int(max_period), ) return output ================================================ FILE: python/sglang/jit_kernel/utils.py ================================================ from __future__ import annotations import functools import os import pathlib from typing import TYPE_CHECKING, Any, Callable, List, Tuple, TypeAlias, TypeVar, Union import torch if TYPE_CHECKING: from tvm_ffi import Module F = TypeVar("F", bound=Callable[..., Any]) _FULL_TEST_ENV_VAR = "SGLANG_JIT_KERNEL_RUN_FULL_TESTS" def is_in_ci() -> bool: ci_env_vars = ("SGLANG_IS_IN_CI", "CI", "GITHUB_ACTIONS") return any(os.getenv(env_var, "false").lower() == "true" for env_var in ci_env_vars) def should_run_full_tests() -> bool: return os.getenv(_FULL_TEST_ENV_VAR, "false").lower() == "true" def get_ci_test_range(full_range: List[Any], ci_range: List[Any]) -> List[Any]: if should_run_full_tests(): return full_range return ci_range if is_in_ci() else full_range def cache_once(fn: F) -> F: """ NOTE: `functools.lru_cache` is not compatible with `torch.compile` So we manually implement a simple cache_once decorator to replace it. """ result_map = {} @functools.wraps(fn) def wrapper(*args, **kwargs): key = (args, tuple(sorted(kwargs.items(), key=lambda x: x[0]))) if key not in result_map: result_map[key] = fn(*args, **kwargs) return result_map[key] return wrapper # type: ignore def _make_wrapper(tup: Tuple[str, str]) -> str: export_name, kernel_name = tup return f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({export_name}, ({kernel_name}));" @cache_once def _resolve_kernel_path() -> pathlib.Path: cur_dir = pathlib.Path(__file__).parent.resolve() # first, try this directory structure def _environment_install(): candidate = cur_dir.resolve() if (candidate / "include").exists() and (candidate / "csrc").exists(): return candidate return None def _package_install(): # TODO: support find path by package return None path = _environment_install() or _package_install() if path is None: raise RuntimeError("Cannot find sglang.jit_kernel path") return path KERNEL_PATH = _resolve_kernel_path() DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")] DEFAULT_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"] DEFAULT_HIP_CFLAGS = [ flag for flag in DEFAULT_CUDA_CFLAGS if flag != "--expt-relaxed-constexpr" ] DEFAULT_LDFLAGS = [] CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool, torch.dtype] class CPPArgList(list[str]): def __str__(self) -> str: return ", ".join(self) CPP_DTYPE_MAP = { torch.float: "fp32_t", torch.float16: "fp16_t", torch.float8_e4m3fn: "fp8_e4m3_t", torch.bfloat16: "bf16_t", torch.int8: "int8_t", torch.int32: "int32_t", torch.int64: "int64_t", } # AMD/ROCm note: @cache_once def is_hip_runtime() -> bool: return bool(torch.version.hip) def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList: def _convert(arg: CPP_TEMPLATE_TYPE) -> str: if isinstance(arg, bool): return "true" if arg else "false" if isinstance(arg, (int, float)): return str(arg) if isinstance(arg, torch.dtype): return CPP_DTYPE_MAP[arg] raise TypeError(f"Unsupported argument type for cpp template: {type(arg)}") return CPPArgList(_convert(arg) for arg in args) def load_jit( *args: str, cpp_files: List[str] | None = None, cuda_files: List[str] | None = None, cpp_wrappers: List[Tuple[str, str]] | None = None, cuda_wrappers: List[Tuple[str, str]] | None = None, extra_cflags: List[str] | None = None, extra_cuda_cflags: List[str] | None = None, extra_ldflags: List[str] | None = None, extra_include_paths: List[str] | None = None, build_directory: str | None = None, ) -> Module: """ Loading a JIT module from C++/CUDA source files. We define a wrapper as a tuple of (export_name, kernel_name), where `export_name` is the name used to called from Python, and `kernel_name` is the name of the kernel class in C++/CUDA source. :param args: Unique marker of the JIT module. Must be distinct for different kernels. :type args: str :param cpp_files: A list of C++ source files. :type cpp_files: List[str] | None :param cuda_files: A list of CUDA source files. :type cuda_files: List[str] | None :param cpp_wrappers: A list of C++ wrappers, defining the export name and kernel name. :type cpp_wrappers: List[Tuple[str, str]] | None :param cuda_wrappers: A list of CUDA wrappers, defining the export name and kernel name. :type cuda_wrappers: List[Tuple[str, str]] | None :param extra_cflags: Extra C++ compiler flags. :type extra_cflags: List[str] | None :param extra_cuda_cflags: Extra CUDA compiler flags. :type extra_cuda_cflags: List[str] | None :param extra_ldflags: Extra linker flags. :type extra_ldflags: List[str] | None :param extra_include_paths: Extra include paths. :type extra_include_paths: List[str] | None :param build_directory: The build directory for JIT compilation. :type build_directory: str | None :return: A just-in-time(JIT) compiled module. :rtype: Module """ from tvm_ffi.cpp import load_inline cpp_files = cpp_files or [] cuda_files = cuda_files or [] cpp_wrappers = cpp_wrappers or [] cuda_wrappers = cuda_wrappers or [] extra_cflags = extra_cflags or [] extra_cuda_cflags = extra_cuda_cflags or [] extra_ldflags = extra_ldflags or [] extra_include_paths = extra_include_paths or [] # include cpp files cpp_paths = [(KERNEL_PATH / "csrc" / f).resolve() for f in cpp_files] cpp_sources = [f'#include "{path}"' for path in cpp_paths] cpp_sources += [_make_wrapper(tup) for tup in cpp_wrappers] # include cuda files cuda_paths = [(KERNEL_PATH / "csrc" / f).resolve() for f in cuda_files] cuda_sources = [f'#include "{path}"' for path in cuda_paths] cuda_sources += [_make_wrapper(tup) for tup in cuda_wrappers] # Override TVM_FFI_CUDA_ARCH_LIST if it does not exist. env_key = "TVM_FFI_CUDA_ARCH_LIST" env_existed = env_key in os.environ selected_cuda_cflags = DEFAULT_CUDA_CFLAGS if is_hip_runtime(): selected_cuda_cflags = DEFAULT_HIP_CFLAGS extra_cuda_cflags = ["-DUSE_ROCM"] + extra_cuda_cflags else: extra_cuda_cflags = [ f"-DSGL_CUDA_ARCH={_get_cuda_arch_value()}" ] + extra_cuda_cflags if not env_existed: os.environ[env_key] = _get_cuda_arch_list() try: return load_inline( "sgl_kernel_jit_" + "_".join(str(arg) for arg in args), cpp_sources=cpp_sources, cuda_sources=cuda_sources, extra_cflags=DEFAULT_CFLAGS + extra_cflags, extra_cuda_cflags=selected_cuda_cflags + extra_cuda_cflags, extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, ) finally: # Reset TVM_FFI_CUDA_ARCH_LIST to original state (not exist) if not env_existed: del os.environ[env_key] @cache_once def is_arch_support_pdl() -> bool: import torch device = torch.cuda.current_device() return torch.cuda.get_device_capability(device)[0] >= 9 @cache_once def _get_cuda_arch_value() -> int: """Get CUDA arch value for -DSGL_CUDA_ARCH (e.g. 900 for SM 9.0).""" device = torch.cuda.current_device() major, minor = torch.cuda.get_device_capability(device) return major * 100 + minor * 10 @cache_once def _get_cuda_arch_list() -> str: """Get the correct CUDA architecture string for TVM_FFI_CUDA_ARCH_LIST.""" device = torch.cuda.current_device() major, minor = torch.cuda.get_device_capability(device) return f"{major}.{minor}" ================================================ FILE: python/sglang/lang/api.py ================================================ """Public APIs of the language.""" import re from typing import Callable, List, Optional, Union from sglang.global_config import global_config from sglang.lang.backend.base_backend import BaseBackend from sglang.lang.choices import ChoicesSamplingMethod, token_length_normalized from sglang.lang.ir import ( SglExpr, SglExprList, SglFunction, SglGen, SglImage, SglRoleBegin, SglRoleEnd, SglSelect, SglSeparateReasoning, SglVideo, ) def function( func: Optional[Callable] = None, num_api_spec_tokens: Optional[int] = None ): if func: return SglFunction(func, num_api_spec_tokens=num_api_spec_tokens) def decorator(func): return SglFunction(func, num_api_spec_tokens=num_api_spec_tokens) return decorator def Runtime(*args, **kwargs): # Avoid importing unnecessary dependency from sglang.lang.backend.runtime_endpoint import Runtime return Runtime(*args, **kwargs) def Engine(*args, **kwargs): # Avoid importing unnecessary dependency from sglang.srt.entrypoints.engine import Engine return Engine(*args, **kwargs) def set_default_backend(backend: BaseBackend): global_config.default_backend = backend def flush_cache(backend: Optional[BaseBackend] = None): backend = backend or global_config.default_backend if backend is None: return False # If backend is Runtime if hasattr(backend, "endpoint"): backend = backend.endpoint return backend.flush_cache() def get_server_info(backend: Optional[BaseBackend] = None): backend = backend or global_config.default_backend if backend is None: return None # If backend is Runtime if hasattr(backend, "endpoint"): backend = backend.endpoint return backend.get_server_info() def gen( name: Optional[str] = None, max_tokens: Optional[int] = None, min_tokens: Optional[int] = None, n: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, stop_token_ids: Optional[List[int]] = None, stop_regex: Optional[Union[str, List[str]]] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, top_k: Optional[int] = None, min_p: Optional[float] = None, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, ignore_eos: Optional[bool] = None, return_logprob: Optional[bool] = None, logprob_start_len: Optional[int] = None, top_logprobs_num: Optional[int] = None, return_text_in_logprobs: Optional[bool] = None, dtype: Optional[Union[type, str]] = None, choices: Optional[List[str]] = None, choices_method: Optional[ChoicesSamplingMethod] = None, regex: Optional[str] = None, json_schema: Optional[str] = None, ): """Call the model to generate. See the meaning of the arguments in docs/backend/sampling_params.md""" if choices: return SglSelect( name, choices, 0.0 if temperature is None else temperature, token_length_normalized if choices_method is None else choices_method, ) # check regex is valid if regex is not None: try: re.compile(regex) except re.error as e: raise e return SglGen( name, max_tokens, min_tokens, n, stop, stop_token_ids, stop_regex, temperature, top_p, top_k, min_p, frequency_penalty, presence_penalty, ignore_eos, return_logprob, logprob_start_len, top_logprobs_num, return_text_in_logprobs, dtype, regex, json_schema, ) def gen_int( name: Optional[str] = None, max_tokens: Optional[int] = None, n: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, stop_token_ids: Optional[List[int]] = None, stop_regex: Optional[Union[str, List[str]]] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, top_k: Optional[int] = None, min_p: Optional[float] = None, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, ignore_eos: Optional[bool] = None, return_logprob: Optional[bool] = None, logprob_start_len: Optional[int] = None, top_logprobs_num: Optional[int] = None, return_text_in_logprobs: Optional[bool] = None, ): return SglGen( name, max_tokens, None, n, stop, stop_token_ids, stop_regex, temperature, top_p, top_k, min_p, frequency_penalty, presence_penalty, ignore_eos, return_logprob, logprob_start_len, top_logprobs_num, return_text_in_logprobs, int, None, ) def gen_string( name: Optional[str] = None, max_tokens: Optional[int] = None, n: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, stop_token_ids: Optional[List[int]] = None, stop_regex: Optional[Union[str, List[str]]] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, top_k: Optional[int] = None, min_p: Optional[float] = None, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, ignore_eos: Optional[bool] = None, return_logprob: Optional[bool] = None, logprob_start_len: Optional[int] = None, top_logprobs_num: Optional[int] = None, return_text_in_logprobs: Optional[bool] = None, ): return SglGen( name, max_tokens, None, n, stop, stop_token_ids, stop_regex, temperature, top_p, top_k, min_p, frequency_penalty, presence_penalty, ignore_eos, return_logprob, logprob_start_len, top_logprobs_num, return_text_in_logprobs, str, None, ) def image(expr: SglExpr): return SglImage(expr) def video(path: str, num_frames: int): return SglVideo(path, num_frames) def select( name: Optional[str] = None, choices: Optional[List[str]] = None, temperature: float = 0.0, choices_method: ChoicesSamplingMethod = token_length_normalized, ): assert choices is not None return SglSelect(name, choices, temperature, choices_method) def _role_common(name: str, expr: Optional[SglExpr] = None): if expr is None: return SglExprList([SglRoleBegin(name), SglRoleEnd(name)]) else: return SglExprList([SglRoleBegin(name), expr, SglRoleEnd(name)]) def system(expr: Optional[SglExpr] = None): return _role_common("system", expr) def user(expr: Optional[SglExpr] = None): return _role_common("user", expr) def assistant(expr: Optional[SglExpr] = None): return _role_common("assistant", expr) def system_begin(): return SglRoleBegin("system") def system_end(): return SglRoleEnd("system") def user_begin(): return SglRoleBegin("user") def user_end(): return SglRoleEnd("user") def assistant_begin(): return SglRoleBegin("assistant") def assistant_end(): return SglRoleEnd("assistant") def separate_reasoning( expr: Optional[SglExpr] = None, model_type: Optional[str] = None ): return SglExprList([expr, SglSeparateReasoning(model_type, expr=expr)]) ================================================ FILE: python/sglang/lang/backend/anthropic.py ================================================ from sglang.lang.backend.base_backend import BaseBackend from sglang.lang.chat_template import get_chat_template from sglang.lang.interpreter import StreamExecutor from sglang.lang.ir import SglSamplingParams try: import anthropic except ImportError as e: anthropic = e class Anthropic(BaseBackend): def __init__(self, model_name, *args, **kwargs): super().__init__() if isinstance(anthropic, Exception): raise anthropic self.model_name = model_name self.chat_template = get_chat_template("claude") self.client = anthropic.Anthropic(*args, **kwargs) def get_chat_template(self): return self.chat_template def generate( self, s: StreamExecutor, sampling_params: SglSamplingParams, ): if s.messages_: messages = s.messages_ else: messages = [{"role": "user", "content": s.text_}] if messages and messages[0]["role"] == "system": system = messages.pop(0)["content"] else: system = "" ret = self.client.messages.create( model=self.model_name, system=system, messages=messages, **sampling_params.to_anthropic_kwargs(), ) comp = ret.content[0].text return comp, {} def generate_stream( self, s: StreamExecutor, sampling_params: SglSamplingParams, ): if s.messages_: messages = s.messages_ else: messages = [{"role": "user", "content": s.text_}] if messages and messages[0]["role"] == "system": system = messages.pop(0)["content"] else: system = "" with self.client.messages.stream( model=self.model_name, system=system, messages=messages, **sampling_params.to_anthropic_kwargs(), ) as stream: for text in stream.text_stream: yield text, {} ================================================ FILE: python/sglang/lang/backend/base_backend.py ================================================ from typing import List, Optional, Union from sglang.lang.chat_template import get_chat_template from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod from sglang.lang.interpreter import StreamExecutor from sglang.lang.ir import SglSamplingParams class BaseBackend: def __init__(self) -> None: self.support_concate_and_append = False self.chat_template = get_chat_template("default") def get_model_name(self): raise NotImplementedError() def get_chat_template(self): return self.chat_template def cache_prefix(self, prefix_str: str): pass def uncache_prefix(self, rid: str): pass def end_request(self, rid: Union[str, List[str]]): pass def begin_program(self, s: StreamExecutor): pass def end_program(self, s: Union[StreamExecutor, List[StreamExecutor]]): pass def commit_lazy_operations(self, s: StreamExecutor): pass def fork_program( self, src: StreamExecutor, dst: List[StreamExecutor], position_ids_offset: Optional[List[int]] = None, ): pass def fill_image(self, s: StreamExecutor): pass def generate( self, s: StreamExecutor, sampling_params: SglSamplingParams, ): raise NotImplementedError() def generate_stream( self, s: StreamExecutor, sampling_params: SglSamplingParams, ): raise NotImplementedError() def select( self, s: StreamExecutor, choices: List[str], temperature: float, choices_method: Optional[ChoicesSamplingMethod] = None, ) -> ChoicesDecision: raise NotImplementedError() def concatenate_and_append(self, src_rids: List[str], dst_rid: str): raise NotImplementedError() def shutdown(self): pass def flush_cache(self): pass def get_server_info(self): pass ================================================ FILE: python/sglang/lang/backend/litellm.py ================================================ from typing import Mapping, Optional from sglang.lang.backend.base_backend import BaseBackend from sglang.lang.chat_template import get_chat_template_by_model_path from sglang.lang.interpreter import StreamExecutor from sglang.lang.ir import SglSamplingParams try: import litellm except ImportError as e: litellm = e litellm.num_retries = 1 class LiteLLM(BaseBackend): def __init__( self, model_name, chat_template=None, api_key=None, organization: Optional[str] = None, base_url: Optional[str] = None, timeout: Optional[float] = 600, max_retries: Optional[int] = litellm.num_retries, default_headers: Optional[Mapping[str, str]] = None, ): super().__init__() if isinstance(litellm, Exception): raise litellm self.model_name = model_name self.chat_template = chat_template or get_chat_template_by_model_path( model_name ) self.client_params = { "api_key": api_key, "organization": organization, "base_url": base_url, "timeout": timeout, "max_retries": max_retries, "default_headers": default_headers, } def get_chat_template(self): return self.chat_template def generate( self, s: StreamExecutor, sampling_params: SglSamplingParams, ): if s.messages_: messages = s.messages_ else: messages = [{"role": "user", "content": s.text_}] ret = litellm.completion( model=self.model_name, messages=messages, **self.client_params, **sampling_params.to_litellm_kwargs(), ) comp = ret.choices[0].message.content return comp, {} def generate_stream( self, s: StreamExecutor, sampling_params: SglSamplingParams, ): if s.messages_: messages = s.messages_ else: messages = [{"role": "user", "content": s.text_}] ret = litellm.completion( model=self.model_name, messages=messages, stream=True, **self.client_params, **sampling_params.to_litellm_kwargs(), ) for chunk in ret: text = chunk.choices[0].delta.content if text is not None: yield text, {} ================================================ FILE: python/sglang/lang/backend/openai.py ================================================ import dataclasses import logging import time import warnings from typing import List, Optional, Union import numpy as np from sglang.lang.backend.base_backend import BaseBackend from sglang.lang.chat_template import ChatTemplate, get_chat_template_by_model_path from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod from sglang.lang.interpreter import StreamExecutor from sglang.lang.ir import SglSamplingParams try: import openai import tiktoken except ImportError as e: openai = tiktoken = e logger = logging.getLogger(__name__) def create_logit_bias_int(tokenizer): """Get logit bias for integer numbers.""" int_token_ids = [] tokens = tokenizer._mergeable_ranks for token, token_id in tokens.items(): s = tokenizer.decode([token_id]) if all([c.isdigit() for c in s]) or s in [" "]: int_token_ids.append(token_id) if len(int_token_ids) >= 300: # OpenAI API limit break special_tokens = tokenizer._special_tokens mask = {t: 100 for t in int_token_ids[:299]} mask[special_tokens["<|endoftext|>"]] = 100 return mask INSTRUCT_MODEL_NAMES = [ "gpt-3.5-turbo-instruct", ] @dataclasses.dataclass class TokenUsage: prompt_tokens: int completion_tokens: int def reset(self): self.prompt_tokens = self.completion_tokens = 0 class OpenAI(BaseBackend): def __init__( self, model_name: str, is_chat_model: Optional[bool] = None, chat_template: Optional[ChatTemplate] = None, is_azure: bool = False, *args, **kwargs, ): super().__init__() if isinstance(openai, Exception): raise openai if is_azure: self.client = openai.AzureOpenAI(*args, **kwargs) else: self.client = openai.OpenAI(*args, **kwargs) self.model_name = model_name try: self.tokenizer = tiktoken.encoding_for_model(model_name) except KeyError: self.tokenizer = tiktoken.get_encoding("cl100k_base") self.logit_bias_int = create_logit_bias_int(self.tokenizer) self.chat_template = chat_template or get_chat_template_by_model_path( model_name ) if is_chat_model is not None: self.is_chat_model = is_chat_model else: if model_name in INSTRUCT_MODEL_NAMES: self.is_chat_model = False else: self.is_chat_model = True self.chat_prefix = self.chat_template.role_prefix_and_suffix["assistant"][0] # Usage self.token_usage = TokenUsage(0, 0) # API speculative execution # TODO(ying): This does not support multi-threading (run_batch) self.spec_kwargs = {} self.spec_format = [] self.spec_max_num_tries = 3 def get_chat_template(self): return self.chat_template def _prepare_spec_execution( self, sampling_params: SglSamplingParams, num_api_spec_tokens: int, spec_var_name: str, ): if "max_tokens" not in self.spec_kwargs: self.spec_kwargs["max_tokens"] = num_api_spec_tokens else: assert self.spec_kwargs["max_tokens"] == num_api_spec_tokens params = sampling_params.to_openai_kwargs() for key, value in params.items(): if key in ["stop"]: continue if key in ["max_tokens"]: warnings.warn( "The parameter max_tokens will be overwritten by speculated number of tokens." ) continue if key not in self.spec_kwargs: self.spec_kwargs[key] = value else: assert ( value == self.spec_kwargs[key] ), "sampling parameters should be consistent if turn on api speculative execution." self.spec_format.append( {"text": "", "stop": params["stop"], "name": spec_var_name} ) return "", {} def generate( self, s: StreamExecutor, sampling_params: SglSamplingParams, spec_var_name: str = None, ): if sampling_params.dtype is None: if self.is_chat_model: if s.num_api_spec_tokens is None: if not s.text_.endswith(self.chat_prefix): raise RuntimeError( "This use case is not supported if api speculative execution is off. " "For OpenAI chat models, sgl.gen must be right after sgl.assistant. " "Example of adding api speculative execution: @function(num_api_spec_tokens=128)." ) prompt = s.messages_ else: return self._prepare_spec_execution( sampling_params, s.num_api_spec_tokens, spec_var_name ) else: prompt = s.text_ kwargs = sampling_params.to_openai_kwargs() if ( self.model_name.startswith("o1") or self.model_name.startswith("o3") or "o1" in self.model_name ): kwargs.pop("max_tokens", None) else: kwargs.pop("max_completion_tokens", None) comp = openai_completion( client=self.client, token_usage=self.token_usage, is_chat=self.is_chat_model, model=self.model_name, prompt=prompt, **kwargs, ) # Keep the returned list (or string) as is. elif sampling_params.dtype in [str, "str", "string"]: assert ( not self.is_chat_model ), "constrained type not supported on chat model" kwargs = sampling_params.to_openai_kwargs() kwargs.pop("stop") comp = openai_completion( client=self.client, token_usage=self.token_usage, is_chat=self.is_chat_model, model=self.model_name, prompt=s.text_ + '"', stop='"', **kwargs, ) # Wrap each element in quotes if we have a list. if isinstance(comp, list): comp = ['"' + x + '"' for x in comp] else: comp = '"' + comp + '"' elif sampling_params.dtype in [int, "int"]: assert ( not self.is_chat_model ), "constrained type not supported on chat model" kwargs = sampling_params.to_openai_kwargs() kwargs.pop("stop") comp = openai_completion( client=self.client, token_usage=self.token_usage, is_chat=self.is_chat_model, model=self.model_name, prompt=s.text_, logit_bias=self.logit_bias_int, stop=[" "], **kwargs, ) # Leave as a list if that's what is returned. else: raise ValueError(f"Unknown dtype: {sampling_params.dtype}") return comp, {} def spec_fill(self, value: str): assert self.is_chat_model self.spec_format.append({"text": value, "stop": None, "name": None}) def spec_pattern_match(self, comp): for i, term in enumerate(self.spec_format): text = term["text"] if text != "": if comp.startswith(text): comp = comp[len(text) :] else: return False else: pos = comp.find(term["stop"]) if pos != -1: term["text"] = comp[:pos] comp = comp[pos:] else: if i == len(self.spec_format) - 1: term["text"] = comp else: return False return True def role_end_generate( self, s: StreamExecutor, ): if s.num_api_spec_tokens is None or not s.text_.endswith(self.chat_prefix): return comp = "" if not all(x["name"] is None for x in self.spec_format): # TODO(ying): throw errors or warnings for i in range(self.spec_max_num_tries): comp = openai_completion( client=self.client, token_usage=self.token_usage, is_chat=self.is_chat_model, model=self.model_name, prompt=s.messages_, **self.spec_kwargs, ) # Use a string for pattern matching. comp_for_match = comp[0] if isinstance(comp, list) else comp if self.spec_pattern_match(comp_for_match): break for term in self.spec_format: s.text_ += term["text"] name = term["name"] if name is not None: s.variables[name] = term["text"] s.meta_info[name] = {} s.variable_event[name].set() self.spec_kwargs = {} self.spec_format = [] def generate_stream( self, s: StreamExecutor, sampling_params: SglSamplingParams, ): if sampling_params.dtype is None: if self.is_chat_model: if not s.text_.endswith(self.chat_prefix): raise RuntimeError( "This use case is not supported. " "For OpenAI chat models, sgl.gen must be right after sgl.assistant" ) prompt = s.messages_ else: prompt = s.text_ kwargs = sampling_params.to_openai_kwargs() generator = openai_completion_stream( client=self.client, token_usage=self.token_usage, is_chat=self.is_chat_model, model=self.model_name, prompt=prompt, **kwargs, ) return generator else: raise ValueError(f"Unknown dtype: {sampling_params.dtype}") def select( self, s: StreamExecutor, choices: List[str], temperature: float, choices_method: ChoicesSamplingMethod, ) -> ChoicesDecision: """Note: `choices_method` is not used by the OpenAI backend.""" if self.is_chat_model: raise NotImplementedError( "select/choices is not supported for chat models. " "Please try to use a non-chat model such as gpt-3.5-turbo-instruct" ) n_choices = len(choices) token_ids = [self.tokenizer.encode(x) for x in choices] scores = [0] * n_choices valid = [len(x) > 0 for x in token_ids] prompt_tokens = self.tokenizer.encode(s.text_) max_len = max([len(x) for x in token_ids]) for step in range(max_len): # Build logit bias logit_bias = {} for i in range(n_choices): if valid[i]: logit_bias[token_ids[i][step]] = 100 # Call API ret = self.client.completions.create( model=self.model_name, prompt=prompt_tokens, logit_bias=logit_bias, max_tokens=1, temperature=temperature, ) ret_str = ret.choices[0].text ret_token = self.tokenizer.encode(ret_str)[0] self.token_usage.prompt_tokens += ret.usage.prompt_tokens self.token_usage.completion_tokens = ret.usage.completion_tokens # TODO: # 1. return logits as the scores # 2. compute logits of the full choice # 3. consider chunk-based decoding # Update valid hit = False for i in range(n_choices): if valid[i]: if step == len(token_ids[i]) - 1: valid[i] = False if ret_token == token_ids[i][step]: scores[i] += 1 hit = True else: valid[i] = False assert hit if np.sum(valid) <= 1: break prompt_tokens.append(ret_token) return ChoicesDecision( decision=choices[np.argmax(scores)], meta_info={"scores": scores}, ) def openai_completion( client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs ) -> Union[str, List[str]]: # if "ebnf" is in kwargs, warn and remove if "ebnf" in kwargs: warnings.warn("EBNF is not officially supported by OpenAI endpoints. Ignoring.") del kwargs["ebnf"] for attempt in range(retries): try: if is_chat: if "stop" in kwargs and kwargs["stop"] is None: kwargs.pop("stop") ret = client.chat.completions.create(messages=prompt, **kwargs) if len(ret.choices) == 1: comp = ret.choices[0].message.content else: comp = [c.message.content for c in ret.choices] else: ret = client.completions.create(prompt=prompt, **kwargs) if isinstance(prompt, (list, tuple)): comp = [c.text for c in ret.choices] else: comp = ret.choices[0].text if len(ret.choices) > 1: comp = [c.text for c in ret.choices] token_usage.prompt_tokens += ret.usage.prompt_tokens token_usage.completion_tokens += ret.usage.completion_tokens break except (openai.APIError, openai.APIConnectionError, openai.RateLimitError) as e: logger.error(f"OpenAI Error: {e}. Waiting 5 seconds...") time.sleep(5) if attempt == retries - 1: raise e except Exception as e: logger.error(f"RuntimeError {e}.") raise e return comp def openai_completion_stream( client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs ): # if "ebnf" is in kwargs, warn and remove if "ebnf" in kwargs: warnings.warn("EBNF is not officially supported by OpenAI endpoints. Ignoring.") del kwargs["ebnf"] for attempt in range(retries): try: if is_chat: if "stop" in kwargs and kwargs["stop"] is None: kwargs.pop("stop") generator = client.chat.completions.create( messages=prompt, stream=True, stream_options={"include_usage": True}, **kwargs, ) for ret in generator: if len(ret.choices) == 0: continue try: content = ret.choices[0].delta.content except IndexError: content = None yield content or "", {} else: generator = client.completions.create( prompt=prompt, stream=True, stream_options={"include_usage": True}, **kwargs, ) for ret in generator: if len(ret.choices) == 0: continue content = ret.choices[0].text yield content or "", {} token_usage.prompt_tokens += ret.usage.prompt_tokens token_usage.completion_tokens += ret.usage.completion_tokens break except (openai.APIError, openai.APIConnectionError, openai.RateLimitError) as e: logger.error(f"OpenAI Error: {e}. Waiting 5 seconds...") time.sleep(5) if attempt == retries - 1: raise e except Exception as e: logger.error(f"RuntimeError {e}.") raise e ================================================ FILE: python/sglang/lang/backend/runtime_endpoint.py ================================================ import atexit import json import multiprocessing import time import warnings from typing import Dict, List, Optional, Union import aiohttp import requests from sglang.global_config import global_config from sglang.lang.backend.base_backend import BaseBackend from sglang.lang.chat_template import get_chat_template, get_chat_template_by_model_path from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod from sglang.lang.interpreter import StreamExecutor from sglang.lang.ir import ( REGEX_BOOL, REGEX_FLOAT, REGEX_INT, REGEX_STR, SglSamplingParams, ) from sglang.utils import http_request class RuntimeEndpoint(BaseBackend): def __init__( self, base_url: str, api_key: Optional[str] = None, verify: Optional[str] = None, chat_template_name: Optional[str] = None, ): super().__init__() self.support_concate_and_append = True self.base_url = base_url self.api_key = api_key self.verify = verify res = http_request( self.base_url + "/get_model_info", api_key=self.api_key, verify=self.verify, ) self._assert_success(res) self.model_info = res.json() if chat_template_name: self.chat_template = get_chat_template(chat_template_name) else: self.chat_template = get_chat_template_by_model_path( self.model_info["model_path"] ) def get_model_name(self): return self.model_info["model_path"] def flush_cache(self): res = http_request( self.base_url + "/flush_cache", api_key=self.api_key, verify=self.verify, method="POST", ) self._assert_success(res) def get_server_info(self): res = http_request( self.base_url + "/get_server_info", api_key=self.api_key, verify=self.verify, ) self._assert_success(res) return res.json() def get_chat_template(self): return self.chat_template def cache_prefix(self, prefix_str: str): res = http_request( self.base_url + "/generate", json={"text": prefix_str, "sampling_params": {"max_new_tokens": 0}}, api_key=self.api_key, verify=self.verify, ) self._assert_success(res) def start_profile(self): res = http_request( self.base_url + "/start_profile", api_key=self.api_key, verify=self.verify, ) self._assert_success(res) def stop_profile(self): res = http_request( self.base_url + "/stop_profile", api_key=self.api_key, verify=self.verify, ) self._assert_success(res) def commit_lazy_operations(self, s: StreamExecutor): data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}} self._add_images(s, data) res = http_request( self.base_url + "/generate", json=data, api_key=self.api_key, verify=self.verify, ) self._assert_success(res) def fill_image(self, s: StreamExecutor): data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}} self._add_images(s, data) res = http_request( self.base_url + "/generate", json=data, api_key=self.api_key, verify=self.verify, ) self._assert_success(res) def _handle_dtype_to_regex(self, sampling_params: SglSamplingParams): if sampling_params.dtype is None: return if sampling_params.stop == (): sampling_params.stop = [] dtype_regex = None if sampling_params.dtype in ["int", int]: dtype_regex = REGEX_INT sampling_params.stop.extend([" ", "\n"]) elif sampling_params.dtype in ["float", float]: dtype_regex = REGEX_FLOAT sampling_params.stop.extend([" ", "\n"]) elif sampling_params.dtype in ["str", str]: dtype_regex = REGEX_STR elif sampling_params.dtype in ["bool", bool]: dtype_regex = REGEX_BOOL else: raise RuntimeError(f"Invalid dtype: {sampling_params.dtype}") if dtype_regex is not None and sampling_params.regex is not None: warnings.warn( f"Both dtype and regex are set. Only dtype will be used. dtype: {sampling_params.dtype}, regex: {sampling_params.regex}" ) sampling_params.regex = dtype_regex def generate( self, s: StreamExecutor, sampling_params: SglSamplingParams, ): self._handle_dtype_to_regex(sampling_params) data = { "text": s.text_, "sampling_params": { "skip_special_tokens": global_config.skip_special_tokens_in_output, "spaces_between_special_tokens": global_config.spaces_between_special_tokens_in_out, **sampling_params.to_srt_kwargs(), }, } for item in [ "return_logprob", "logprob_start_len", "top_logprobs_num", "return_text_in_logprobs", ]: value = getattr(sampling_params, item, None) if value is not None: data[item] = value self._add_images(s, data) res = http_request( self.base_url + "/generate", json=data, api_key=self.api_key, verify=self.verify, ) self._assert_success(res) obj = res.json() comp = obj["text"] return comp, obj["meta_info"] def generate_stream( self, s: StreamExecutor, sampling_params: SglSamplingParams, ): self._handle_dtype_to_regex(sampling_params) data = { "text": s.text_, "sampling_params": { "skip_special_tokens": global_config.skip_special_tokens_in_output, "spaces_between_special_tokens": global_config.spaces_between_special_tokens_in_out, **sampling_params.to_srt_kwargs(), }, } for item in [ "return_logprob", "logprob_start_len", "top_logprobs_num", "return_text_in_logprobs", ]: value = getattr(sampling_params, item, None) if value is not None: data[item] = value data["stream"] = True self._add_images(s, data) res = http_request( self.base_url + "/generate", json=data, stream=True, api_key=self.api_key, verify=self.verify, ) self._assert_success(res) pos = 0 for chunk in res.iter_lines(decode_unicode=False): chunk = chunk.decode("utf-8") if chunk and chunk.startswith("data:"): if chunk == "data: [DONE]": break data = json.loads(chunk[5:].strip("\n")) chunk_text = data["text"][pos:] meta_info = data["meta_info"] pos += len(chunk_text) yield chunk_text, meta_info def select( self, s: StreamExecutor, choices: List[str], temperature: float, choices_method: ChoicesSamplingMethod, ) -> ChoicesDecision: assert temperature <= 1e-5 # Cache common prefix data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}} obj = self._generate_http_request(s, data) prompt_len = obj["meta_info"]["prompt_tokens"] logprob_start_len = max(prompt_len - 2, 0) # For token healing # Compute logprob data = { "text": [s.text_ + c for c in choices], "sampling_params": { "max_new_tokens": 0, "temperature": 0, }, "return_logprob": True, "return_text_in_logprobs": True, "logprob_start_len": logprob_start_len, } obj = self._generate_http_request(s, data) input_token_logprobs = [r["meta_info"]["input_token_logprobs"] for r in obj] output_token_logprobs = [r["meta_info"]["output_token_logprobs"] for r in obj] normalized_prompt_logprobs = [ compute_normalized_prompt_logprobs(r["meta_info"]["input_token_logprobs"]) for r in obj ] # Remove extra token if no token healing occurred for i in range(len(input_token_logprobs)): healed_token_str = input_token_logprobs[i][0][-1] if s.text_.endswith(healed_token_str): healed_token_logprob = input_token_logprobs[i][0][0] normalized_prompt_logprobs[i] = ( normalized_prompt_logprobs[i] * len(input_token_logprobs[i]) - healed_token_logprob ) / (len(input_token_logprobs[i]) - 1) input_token_logprobs[i] = input_token_logprobs[i][1:] # Compute unconditional logprobs if required if choices_method.requires_unconditional_logprobs: input_ids = [[el[1] for el in subl] for subl in input_token_logprobs] data = { "input_ids": input_ids, "sampling_params": {"max_new_tokens": 0}, "return_logprob": True, } obj = self._generate_http_request(s, data) unconditional_token_logprobs = [ r["meta_info"]["input_token_logprobs"] for r in obj ] else: unconditional_token_logprobs = None return choices_method( choices=choices, normalized_prompt_logprobs=normalized_prompt_logprobs, input_token_logprobs=input_token_logprobs, output_token_logprobs=output_token_logprobs, unconditional_token_logprobs=unconditional_token_logprobs, ) def concatenate_and_append(self, src_rids: List[str], dst_rid: str): res = http_request( self.base_url + "/concate_and_append_request", json={"src_rids": src_rids, "dst_rid": dst_rid}, api_key=self.api_key, verify=self.verify, ) self._assert_success(res) def _generate_http_request(self, s: StreamExecutor, data): self._add_images(s, data) res = http_request( self.base_url + "/generate", json=data, api_key=self.api_key, verify=self.verify, ) self._assert_success(res) return res.json() def _add_images(self, s: StreamExecutor, data): if s.images_: assert len(s.images_) == 1, "Only support one image." data["image_data"] = s.images_[0][1] def _assert_success(self, res): if res.status_code != 200: try: content = res.json() except json.JSONDecodeError: content = res.text raise RuntimeError(content) def compute_normalized_prompt_logprobs(input_logprobs): values = [x[0] for x in input_logprobs if x[0]] return sum(values) / len(values) class Runtime: """ A wrapper for the HTTP server. This is used for launching the server in a python program without using the command line interface. It is mainly used for the frontend language. You should use the Engine class if you want to do normal offline processing without the frontend language. """ def __init__( self, log_level: str = "error", launch_timeout: float = 300.0, *args, **kwargs, ): """See the arguments in server_args.py::ServerArgs Args: log_level: Log level for the server. timeout: Timeout in seconds for waiting for the server to start. *args: Additional arguments passed to ServerArgs. **kwargs: Additional keyword arguments passed to ServerArgs. """ # We delay the import of any `sglang.srt` components in `sglang.lang`, so users can run # client code without installing SRT server and its dependency if they want. from sglang.srt.entrypoints.http_server import launch_server from sglang.srt.server_args import ServerArgs from sglang.srt.utils.network import is_port_available self.server_args = ServerArgs(*args, log_level=log_level, **kwargs) # Pre-allocate ports for port in range(self.server_args.port, 40000): if is_port_available(port): break self.server_args.port = port self.url = self.server_args.url() self.generate_url = self.url + "/generate" # NOTE: We store pid instead of proc to fix some issues during __delete__ self.pid = None ctx = multiprocessing.get_context("spawn") proc = ctx.Process( target=launch_server, args=(self.server_args,), ) proc.start() self.pid = proc.pid # Before python program terminates, call shutdown implicitly. Therefore, users don't have to explicitly call .shutdown() atexit.register(self.shutdown) # Wait for server to be ready by polling /health_generate start_time = time.time() with requests.Session() as session: while time.time() - start_time < launch_timeout: try: response = session.get(f"{self.url}/health_generate") if response.status_code == 200: break except requests.RequestException: pass if not proc.is_alive(): self.shutdown() raise RuntimeError( "Initialization failed. Please see the error messages above." ) time.sleep(2) else: self.shutdown() raise TimeoutError("Server failed to start within the timeout period.") self.endpoint = RuntimeEndpoint(self.url) def shutdown(self): from sglang.srt.utils import kill_process_tree if self.pid is not None: kill_process_tree(self.pid) self.pid = None def start_profile(self): self.endpoint.start_profile() def stop_profile(self): self.endpoint.stop_profile() def cache_prefix(self, prefix: str): self.endpoint.cache_prefix(prefix) def get_tokenizer(self): from sglang.srt.utils.hf_transformers_utils import get_tokenizer return get_tokenizer( self.server_args.tokenizer_path, tokenizer_mode=self.server_args.tokenizer_mode, trust_remote_code=self.server_args.trust_remote_code, revision=self.server_args.revision, ) async def async_generate( self, prompt: str, sampling_params: Optional[Dict] = None, ): if self.server_args.skip_tokenizer_init: json_data = { "input_ids": prompt, "sampling_params": sampling_params, "stream": True, } else: json_data = { "text": prompt, "sampling_params": sampling_params, "stream": True, } pos = 0 timeout = aiohttp.ClientTimeout(total=3 * 3600) async with aiohttp.ClientSession(timeout=timeout, trust_env=True) as session: async with session.post(self.generate_url, json=json_data) as response: async for chunk, _ in response.content.iter_chunks(): chunk = chunk.decode("utf-8") if chunk and chunk.startswith("data:"): if chunk == "data: [DONE]\n\n": break data = json.loads(chunk[5:].strip("\n")) if "text" in data: cur = data["text"][pos:] if cur: yield cur pos += len(cur) else: yield data add_request = async_generate def generate( self, prompt: Union[str, List[str]], sampling_params: Optional[Dict] = None, return_logprob: Optional[Union[List[bool], bool]] = False, logprob_start_len: Optional[Union[List[int], int]] = None, top_logprobs_num: Optional[Union[List[int], int]] = None, lora_path: Optional[List[Optional[str]]] = None, ): json_data = { "text": prompt, "sampling_params": sampling_params, "return_logprob": return_logprob, "logprob_start_len": logprob_start_len, "top_logprobs_num": top_logprobs_num, "lora_path": lora_path, } assert not isinstance(lora_path, list) or len(lora_path) == len(prompt) response = requests.post( self.url + "/generate", json=json_data, ) return json.dumps(response.json()) def encode( self, prompt: Union[str, List[str], List[Dict], List[List[Dict]]], ): json_data = {"text": prompt} response = requests.post(self.url + "/encode", json=json_data) return json.dumps(response.json()) async def get_server_info(self): async with aiohttp.ClientSession() as session: async with session.get(f"{self.url}/get_server_info") as response: if response.status == 200: return await response.json() else: error_data = await response.json() raise RuntimeError( f"Failed to get server info. {error_data['error']['message']}" ) def __del__(self): self.shutdown() ================================================ FILE: python/sglang/lang/backend/vertexai.py ================================================ import os import warnings from sglang.lang.backend.base_backend import BaseBackend from sglang.lang.chat_template import get_chat_template from sglang.lang.interpreter import StreamExecutor from sglang.lang.ir import SglSamplingParams try: import vertexai from vertexai.preview.generative_models import ( GenerationConfig, GenerativeModel, Image, ) except ImportError as e: GenerativeModel = e class VertexAI(BaseBackend): def __init__(self, model_name, safety_settings=None): super().__init__() if isinstance(GenerativeModel, Exception): raise GenerativeModel project_id = os.environ["GCP_PROJECT_ID"] location = os.environ.get("GCP_LOCATION") vertexai.init(project=project_id, location=location) self.model_name = model_name self.chat_template = get_chat_template("default") self.safety_settings = safety_settings def get_chat_template(self): return self.chat_template def generate( self, s: StreamExecutor, sampling_params: SglSamplingParams, ): if s.messages_: prompt = self.messages_to_vertexai_input(s.messages_) else: # single-turn prompt = ( self.text_to_vertexai_input(s.text_, s.cur_images) if s.cur_images else s.text_ ) ret = GenerativeModel(self.model_name).generate_content( prompt, generation_config=GenerationConfig(**sampling_params.to_vertexai_kwargs()), safety_settings=self.safety_settings, ) comp = ret.text return comp, {} def generate_stream( self, s: StreamExecutor, sampling_params: SglSamplingParams, ): if s.messages_: prompt = self.messages_to_vertexai_input(s.messages_) else: # single-turn prompt = ( self.text_to_vertexai_input(s.text_, s.cur_images) if s.cur_images else s.text_ ) generator = GenerativeModel(self.model_name).generate_content( prompt, stream=True, generation_config=GenerationConfig(**sampling_params.to_vertexai_kwargs()), safety_settings=self.safety_settings, ) for ret in generator: yield ret.text, {} def text_to_vertexai_input(self, text, images): input = [] # split with image token text_segs = text.split(self.chat_template.image_token) for image_path, image_base64_data in images: text_seg = text_segs.pop(0) if text_seg != "": input.append(text_seg) input.append(Image.from_bytes(image_base64_data)) text_seg = text_segs.pop(0) if text_seg != "": input.append(text_seg) return input def messages_to_vertexai_input(self, messages): vertexai_message = [] # from openai message format to vertexai message format for msg in messages: if isinstance(msg["content"], str): text = msg["content"] else: text = msg["content"][0]["text"] if msg["role"] == "system": warnings.warn("Warning: system prompt is not supported in VertexAI.") vertexai_message.append( { "role": "user", "parts": [{"text": "System prompt: " + text}], } ) vertexai_message.append( { "role": "model", "parts": [{"text": "Understood."}], } ) continue if msg["role"] == "user": vertexai_msg = { "role": "user", "parts": [{"text": text}], } elif msg["role"] == "assistant": vertexai_msg = { "role": "model", "parts": [{"text": text}], } # images if isinstance(msg["content"], list) and len(msg["content"]) > 1: for image in msg["content"][1:]: assert image["type"] == "image_url" vertexai_msg["parts"].append( { "inline_data": { "data": image["image_url"]["url"].split(",")[1], "mime_type": "image/jpeg", } } ) vertexai_message.append(vertexai_msg) return vertexai_message ================================================ FILE: python/sglang/lang/chat_template.py ================================================ import re from dataclasses import dataclass from enum import Enum, auto from typing import Callable, Dict, List, Tuple class ChatTemplateStyle(Enum): PLAIN = auto() LLAMA2 = auto() @dataclass class ChatTemplate: name: str default_system_prompt: str role_prefix_and_suffix: Dict[str, Tuple[str, str]] stop_str: List[str] = () image_token: str = "" audio_token: str = ""), }, image_token=" \n", ) ) register_chat_template( ChatTemplate( name="llama-2-chat", default_system_prompt=None, role_prefix_and_suffix={ "system": ("<>\n", "\n<>\n\n"), "user": ("[INST] ", " [/INST]"), "assistant": ("", " "), }, style=ChatTemplateStyle.LLAMA2, ) ) # Reference: https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503/blob/main/chat_template.json register_chat_template( ChatTemplate( name="mistral", default_system_prompt=None, role_prefix_and_suffix={ "system": ("[SYSTEM_PROMPT] ", " [/SYSTEM_PROMPT]"), "user": ("[INST] ", " [/INST]"), "assistant": ("", " "), }, stop_str=("",), image_token="[IMG]", ) ) register_chat_template( ChatTemplate( name="llama-3-instruct", default_system_prompt=None, role_prefix_and_suffix={ "system": ( "<|start_header_id|>system<|end_header_id|>\n\n", "<|eot_id|>", ), "user": ( "<|start_header_id|>user<|end_header_id|>\n\n", "<|eot_id|>", ), "assistant": ( "<|start_header_id|>assistant<|end_header_id|>\n\n", "<|eot_id|>", ), }, stop_str=("<|eot_id|>",), image_token="<|image|>", ) ) # https://huggingface.co/openbmb/MiniCPM-V-2_6 register_chat_template( ChatTemplate( name="minicpmv", default_system_prompt=None, role_prefix_and_suffix={ "system": ("", " "), "user": ("user:", " "), "assistant": ("assistant:", ""), }, stop_str=("<|im_end|>", "<|endoftext|>"), image_token="(./)", ) ) register_chat_template( ChatTemplate( name="janus-pro", default_system_prompt=None, role_prefix_and_suffix={ "system": ( "", "", ), "User": ( "<｜User｜>", "", ), "assistant": ( "<｜Assistant｜>", "<｜end▁of▁sentence｜>", ), }, stop_str=("<｜end▁of▁sentence｜>",), image_token="\n", ) ) # https://huggingface.co/openbmb/MiniCPM-o-2_6 register_chat_template( ChatTemplate( name="minicpmo", default_system_prompt=None, role_prefix_and_suffix={ "system": ("", " "), "user": ("user:", " "), "assistant": ("assistant:", ""), }, stop_str=("<|im_end|>", "<|endoftext|>"), image_token="(./)", audio_token="()", ) ) register_chat_template( ChatTemplate( name="janus", default_system_prompt=None, role_prefix_and_suffix={ "system": ( "", "", ), "user": ( "<｜User｜>", "", ), "assistant": ( "<｜Assistant｜>", "<｜end▁of▁sentence｜>", ), }, stop_str=("<｜end▁of▁sentence｜>",), image_token="\n", ) ) # The difference between "llama-3-instruct-llava" and "llama-3-instruct" is that llava uses a different image_token. register_chat_template( ChatTemplate( name="llama-3-instruct-llava", default_system_prompt=None, role_prefix_and_suffix={ "system": ( "<|start_header_id|>system<|end_header_id|>\n\n", "<|eot_id|>", ), "user": ( "<|start_header_id|>user<|end_header_id|>\n\n", "<|eot_id|>", ), "assistant": ( "<|start_header_id|>assistant<|end_header_id|>\n\n", "<|eot_id|>", ), }, stop_str=("<|eot_id|>",), image_token="\n", ) ) # Reference: https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct/blob/main/chat_template.json register_chat_template( ChatTemplate( name="llama-4", default_system_prompt=None, role_prefix_and_suffix={ "system": ( "<|header_start|>system<|header_end|>\n\n", "<|eot|>", ), "user": ( "<|header_start|>user<|header_end|>\n\n", "<|eot|>", ), "assistant": ( "<|header_start|>assistant<|header_end|>\n\n", "<|eot|>", ), }, stop_str=("<|eot|>",), image_token="<|image|>", ) ) # Reference: https://modelscope.cn/models/01ai/Yi-1.5-34B-Chat/file/view/master?fileName=tokenizer_config.json&status=1 register_chat_template( ChatTemplate( name="yi-1.5", default_system_prompt=None, role_prefix_and_suffix={ "system": ("", ""), "user": ("<|im_start|>user\n", "<|im_end|>\n<|im_start|>assistant\n"), "assistant": ("", "<|im_end|>\n"), }, style=ChatTemplateStyle.PLAIN, stop_str=("<|im_end|>",), ) ) # Reference: https://github.com/01-ai/Yi/tree/main/VL#major-difference-with-llava register_chat_template( ChatTemplate( name="yi-vl", default_system_prompt=( "This is a chat between an inquisitive human and an AI assistant. Assume the role of the AI assistant. Read all the images carefully, and respond to the human's questions with informative, helpful, detailed and polite answers." "这是一个好奇的人类和一个人工智能助手之间的对话。假设你扮演这个AI助手的角色。仔细阅读所有的图像，并对人类的问题做出信息丰富、有帮助、详细的和礼貌的回答。" ), role_prefix_and_suffix={ "system": ("", "\n\n"), "user": ("### Human:", "\n"), "assistant": ("### Assistant:", "\n"), }, image_token=" \n", ) ) register_chat_template( ChatTemplate( name="gemma-it", default_system_prompt=None, role_prefix_and_suffix={ "system": ("", ""), "user": ("user\n", "\n"), "assistant": ("model\n", "\n"), }, image_token="", audio_token="", style=ChatTemplateStyle.PLAIN, ) ) register_chat_template( ChatTemplate( name="dbrx-instruct", default_system_prompt="You are DBRX, created by Databricks. You were last updated in December 2023. You answer questions based on information available up to that point.\nYOU PROVIDE SHORT RESPONSES TO SHORT QUESTIONS OR STATEMENTS, but provide thorough responses to more complex and open-ended questions.\nYou assist with various tasks, from writing to coding (using markdown for code blocks — remember to use ``` with code, JSON, and tables).\n(You do not have real-time data access or code execution capabilities. You avoid stereotyping and provide balanced perspectives on controversial topics. You do not provide song lyrics, poems, or news articles and do not divulge details of your training data.)\nThis is your system prompt, guiding your responses. Do not reference it, just respond to the user. If you find yourself talking about this message, stop. You should be responding appropriately and usually that means not mentioning this.\nYOU DO NOT MENTION ANY OF THIS INFORMATION ABOUT YOURSELF UNLESS THE INFORMATION IS DIRECTLY PERTINENT TO THE USER'S QUERY.", role_prefix_and_suffix={ "system": ("<|im_start|>system\n", "<|im_end|>"), "user": ("\n<|im_start|>user\n", "<|im_end|>"), "assistant": ("\n<|im_start|>assistant\n", "<|im_end|>"), }, stop_str=("<|im_end|>",), ) ) register_chat_template( ChatTemplate( name="c4ai-command-r", default_system_prompt=None, role_prefix_and_suffix={ "system": ( "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>", "<|END_OF_TURN_TOKEN|>", ), "user": ("<|START_OF_TURN_TOKEN|><|USER_TOKEN|>", "<|END_OF_TURN_TOKEN|>"), "assistant": ( "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", "<|END_OF_TURN_TOKEN|>", ), }, style=ChatTemplateStyle.PLAIN, ) ) # Adapted from https://huggingface.co/OpenGVLab/InternVL2-4B/blob/main/modeling_intern_vit.py register_chat_template( ChatTemplate( name="internvl-2-5", default_system_prompt="你是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。", role_prefix_and_suffix={ "system": ("<|im_start|>system\n", "<|im_end|>\n"), "user": ("<|im_start|>user\n", "<|im_end|>\n"), "assistant": ("<|im_start|>assistant\n", "<|im_end|>\n"), }, stop_str=["<|im_end|>", "<|action_end|>"], ) ) register_chat_template( ChatTemplate( name="interns1", default_system_prompt="You are an AI assistant whose name is Intern-S1 (书生大模型).\n- Intern-S1 (书生大模型) is a vision-language model that is developed by Shanghai AI Laboratory (上海人工智能实验室). It is designed to be helpful, honest, and harmless.\n- Intern-S1 (书生大模型) can understand and communicate fluently in the language chosen by the user such as English and 中文.\nYou are an expert reasoner with extensive experience in all areas. You approach problems through systematic thinking and rigorous reasoning. Your response should reflect deep understanding and precise logical thinking, making your solution path and reasoning clear to others. Please put your thinking process within ... tags.", role_prefix_and_suffix={ "system": ("<|im_start|>system\n", "<|im_end|>\n"), "user": ("<|im_start|>user\n", "<|im_end|>\n"), "assistant": ("<|im_start|>assistant\n", "<|im_end|>\n"), }, stop_str=["<|im_end|>", "<|action_end|>"], ) ) register_chat_template( ChatTemplate( name="granite-3-instruct", default_system_prompt=None, role_prefix_and_suffix={ "system": ( "<|start_of_role|>system<|end_of_role|>", "<|end_of_text|>", ), "user": ( "<|start_of_role|>user<|end_of_role|>", "<|end_of_text|>", ), "assistant": ( "<|start_of_role|>assistant<|end_of_role|>", "<|end_of_text|>", ), }, stop_str=("<|end_of_text|>",), ) ) register_chat_template( ChatTemplate( name="deepseek-v3", default_system_prompt=None, role_prefix_and_suffix={ "system": ( "", "", ), "user": ( "<｜User｜>", "", ), "assistant": ( "<｜Assistant｜>", "<｜end▁of▁sentence｜>", ), }, stop_str=("<｜end▁of▁sentence｜>",), ) ) # Reference: https://huggingface.co/docs/transformers/main/model_doc/glm4_v#usage-example register_chat_template( ChatTemplate( name="glm-4v", default_system_prompt=None, role_prefix_and_suffix={ "system": ("<|system|>\n", "\n"), "user": ("<|user|>\n", "\n"), "assistant": ("<|assistant|>\n", "\n"), }, style=ChatTemplateStyle.PLAIN, stop_str=["<|user|>", "<|endoftext|>", "<|observation|>"], image_token="<|image|>", ) ) @register_chat_template_matching_function def match_deepseek(model_path: str): if re.search(r"deepseek-(v3|r1)", model_path, re.IGNORECASE) and not re.search( r"base", model_path, re.IGNORECASE ): return "deepseek-v3" @register_chat_template_matching_function def match_orion(model_path: str): if "orion" in model_path.lower(): return "claude" @register_chat_template_matching_function def match_deepseek_janus_pro(model_path: str): if re.search(r"janus", model_path, re.IGNORECASE): return "janus-pro" @register_chat_template_matching_function def match_dbrx(model_path: str): if re.search(r"dbrx", model_path, re.IGNORECASE) and re.search( r"instruct", model_path, re.IGNORECASE ): return "dbrx-instruct" @register_chat_template_matching_function def match_vicuna(model_path: str): if re.search(r"vicuna|llava-v1\.5|llava-next-video-7b", model_path, re.IGNORECASE): return "vicuna_v1.1" @register_chat_template_matching_function def match_llama2_chat(model_path: str): if re.search( r"llama-2.*chat|codellama.*instruct", model_path, re.IGNORECASE, ): return "llama-2-chat" @register_chat_template_matching_function def match_mistral(model_path: str): if re.search(r"pixtral|(mistral|mixtral).*instruct", model_path, re.IGNORECASE): return "mistral" @register_chat_template_matching_function def match_llama3_instruct(model_path: str): if re.search(r"llama-3.*instruct", model_path, re.IGNORECASE): return "llama-3-instruct" @register_chat_template_matching_function def match_chat_ml(model_path: str): if re.search(r"tinyllama", model_path, re.IGNORECASE): return "chatml" if re.search(r"qwen.*vl", model_path, re.IGNORECASE): return "qwen2-vl" if re.search(r"glm[-_]?4(\.\d+)?v", model_path, re.IGNORECASE): return "glm-4v" if re.search(r"qwen.*(chat|instruct)", model_path, re.IGNORECASE) and not re.search( r"llava", model_path, re.IGNORECASE ): return "qwen" if re.search( r"llava-v1\.6-34b|llava-v1\.6-yi-34b|llava-next-video-34b|llava-onevision-qwen2", model_path, re.IGNORECASE, ): return "chatml-llava" @register_chat_template_matching_function def match_chat_yi(model_path: str): if re.search(r"yi-vl", model_path, re.IGNORECASE) and not re.search( r"llava", model_path, re.IGNORECASE ): return "yi-vl" elif re.search(r"yi-1\.5.*chat", model_path, re.IGNORECASE): return "yi-1.5" @register_chat_template_matching_function def match_gemma_it(model_path: str): if re.search(r"gemma.*it", model_path, re.IGNORECASE): return "gemma-it" @register_chat_template_matching_function def match_openbmb_minicpm(model_path: str): if re.search(r"minicpm-v", model_path, re.IGNORECASE): return "minicpmv" elif re.search(r"minicpm-o", model_path, re.IGNORECASE): return "minicpmo" @register_chat_template_matching_function def match_c4ai_command_r(model_path: str): if re.search(r"c4ai-command-r", model_path, re.IGNORECASE): return "c4ai-command-r" @register_chat_template_matching_function def match_granite_instruct(model_path: str): if re.search(r"granite.*instruct", model_path, re.IGNORECASE): return "granite-3-instruct" @register_chat_template_matching_function def match_gemma3_instruct(model_path: str): if re.search(r"gemma-3", model_path, re.IGNORECASE): return "gemma-it" @register_chat_template_matching_function def match_internvl_chat(model_path: str): if re.search(r"internvl2_5", model_path, re.IGNORECASE): return "internvl-2-5" @register_chat_template_matching_function def match_interns1_chat(model_path: str): if re.search(r"intern-s1", model_path, re.IGNORECASE): return "interns1" if re.search(r"interns1", model_path, re.IGNORECASE): return "interns1" if __name__ == "__main__": messages = [ {"role": "system", "content": None}, # None means default # {"role": "system", "content": "You are a helpful, respectful and honest assistant."}, {"role": "user", "content": "Hello!"}, {"role": "assistant", "content": "Hi!"}, {"role": "user", "content": "What can you do?"}, {"role": "assistant", "content": "I can chat with you."}, ] template = get_chat_template("llama-2-chat") print(template.get_prompt(messages)) ================================================ FILE: python/sglang/lang/choices.py ================================================ from abc import ABC, abstractmethod from dataclasses import dataclass from typing import Any, Dict, List, Optional import numpy as np @dataclass class ChoicesDecision: decision: str meta_info: Optional[Dict[str, Any]] = None class ChoicesSamplingMethod(ABC): @property def requires_unconditional_logprobs(self) -> bool: return False @abstractmethod def __call__( self, *, choices: List[str], normalized_prompt_logprobs: List[float], input_token_logprobs: List[List[Any]], output_token_logprobs: List[List[Any]], unconditional_token_logprobs: Optional[List[List[Any]]] = None, ) -> ChoicesDecision: ... class TokenLengthNormalized(ChoicesSamplingMethod): def __call__( self, *, choices: List[str], normalized_prompt_logprobs: List[float], input_token_logprobs: List[List[Any]], output_token_logprobs: List[List[Any]], unconditional_token_logprobs: Optional[List[List[Any]]] = None, ) -> ChoicesDecision: """Select the option with the highest token length normalized prompt logprob.""" best_choice = choices[np.argmax(normalized_prompt_logprobs)] meta_info = { "normalized_prompt_logprobs": normalized_prompt_logprobs, "input_token_logprobs": input_token_logprobs, "output_token_logprobs": output_token_logprobs, } return ChoicesDecision(decision=best_choice, meta_info=meta_info) token_length_normalized = TokenLengthNormalized() class GreedyTokenSelection(ChoicesSamplingMethod): def __call__( self, *, choices: List[str], normalized_prompt_logprobs: List[float], input_token_logprobs: List[List[Any]], output_token_logprobs: List[List[Any]], unconditional_token_logprobs: Optional[List[List[Any]]] = None, ) -> ChoicesDecision: """Select the option based on greedy logprob selection. For overlapping options where one option is a subset of a longer option, extend the shorter option using its average logprob for comparison against the longer option.""" num_options = len(choices) max_tokens = max(len(option) for option in input_token_logprobs) logprob_matrix = self._build_logprob_matrix( input_token_logprobs, max_tokens, num_options ) remaining = self._greedy_selection(logprob_matrix, num_options, max_tokens) best_choice = choices[remaining[0]] meta_info = { "normalized_prompt_logprobs": normalized_prompt_logprobs, "input_token_logprobs": input_token_logprobs, "output_token_logprobs": output_token_logprobs, "greedy_logprob_matrix": logprob_matrix.tolist(), } return ChoicesDecision(decision=best_choice, meta_info=meta_info) def _build_logprob_matrix(self, input_token_logprobs, max_tokens, num_options): logprob_matrix = np.zeros((num_options, max_tokens)) for i, option in enumerate(input_token_logprobs): actual_logprobs = [token[0] for token in option] avg_logprob = np.mean(actual_logprobs) logprob_matrix[i, : len(option)] = actual_logprobs if len(option) < max_tokens: logprob_matrix[i, len(option) :] = avg_logprob return logprob_matrix def _greedy_selection(self, logprob_matrix, num_options, max_tokens): remaining = np.arange(num_options) for j in range(max_tokens): max_logprob = np.max(logprob_matrix[remaining, j]) remaining = remaining[logprob_matrix[remaining, j] == max_logprob] if len(remaining) == 1: break return remaining greedy_token_selection = GreedyTokenSelection() class UnconditionalLikelihoodNormalized(ChoicesSamplingMethod): @property def requires_unconditional_logprobs(self) -> bool: return True def __call__( self, *, choices: List[str], normalized_prompt_logprobs: List[float], input_token_logprobs: List[List[Any]], output_token_logprobs: List[List[Any]], unconditional_token_logprobs: Optional[List[List[Any]]] = None, ) -> ChoicesDecision: """Select the option with the highest average token logprob once normalized by the unconditional token logprobs. The first unconditional token logprob is assumed to be None. If so, it is replaced with 0 for the purposes of normalization.""" if unconditional_token_logprobs is None: raise ValueError( "Unconditional token logprobs are required for this method." ) normalized_unconditional_prompt_logprobs = self._normalize_logprobs( input_token_logprobs, unconditional_token_logprobs ) best_choice = choices[np.argmax(normalized_unconditional_prompt_logprobs)] meta_info = { "normalized_prompt_logprobs": normalized_prompt_logprobs, "input_token_logprobs": input_token_logprobs, "output_token_logprobs": output_token_logprobs, "unconditional_token_logprobs": unconditional_token_logprobs, "normalized_unconditional_prompt_logprobs": normalized_unconditional_prompt_logprobs, } return ChoicesDecision(decision=best_choice, meta_info=meta_info) def _normalize_logprobs(self, input_token_logprobs, unconditional_token_logprobs): normalized_unconditional_prompt_logprobs = [] for inputs, unconditionals in zip( input_token_logprobs, unconditional_token_logprobs ): inputs_logprobs = np.array([token[0] for token in inputs]) unconditionals_logprobs = np.array([token[0] for token in unconditionals]) unconditionals_logprobs[0] = unconditionals_logprobs[0] or 0 normalized_unconditional_prompt_logprobs.append( float(np.mean(inputs_logprobs - unconditionals_logprobs)) ) return normalized_unconditional_prompt_logprobs unconditional_likelihood_normalized = UnconditionalLikelihoodNormalized() ================================================ FILE: python/sglang/lang/interpreter.py ================================================ """The interpreter that executes SGL programs""" import asyncio import contextvars import copy import multiprocessing import queue import threading import uuid import warnings from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from typing import Any, Callable, Dict, List, Optional import tqdm from sglang.global_config import global_config from sglang.lang.ir import ( SglCommitLazy, SglConcateAndAppend, SglConstantText, SglExpr, SglExprList, SglGen, SglImage, SglRoleBegin, SglRoleEnd, SglSelect, SglSeparateReasoning, SglVariable, SglVarScopeBegin, SglVarScopeEnd, SglVideo, ) from sglang.utils import ( encode_image_base64, encode_video_base64, get_exception_traceback, ) def run_internal(state, program, func_args, func_kwargs, sync): try: state.ret_value = program.func(state, *func_args, **func_kwargs) except Exception as e: raise e finally: state.stream_executor.end() if sync: state.stream_executor.sync() if global_config.verbosity >= 2: print(state.text()) def run_program( program, backend, func_args, func_kwargs, default_sampling_para, stream, sync=False, use_thread=True, ): if hasattr(backend, "endpoint"): backend = backend.endpoint assert backend is not None, "Please specify a backend" func_kwargs.update(program.bind_arguments) stream_executor = StreamExecutor( backend, func_kwargs, default_sampling_para, chat_template=None, stream=stream, num_api_spec_tokens=program.num_api_spec_tokens, use_thread=use_thread, ) state = ProgramState(stream_executor) if stream: t = threading.Thread( target=run_internal, args=(state, program, func_args, func_kwargs, sync) ) t.start() return state else: run_internal(state, program, func_args, func_kwargs, sync) return state def run_program_batch( program, backend, batch_arguments, default_sampling_para, num_threads, progress_bar, generator_style=False, ): if hasattr(backend, "endpoint"): backend = backend.endpoint # Pre-cache the common prefix for a batch. The prefix is extracted by tracing the program. if global_config.enable_precache_with_tracing and len(batch_arguments) > 1: cache_program(program, backend) # Run all programs if num_threads == "auto": num_threads = max(96, multiprocessing.cpu_count() * 16) num_threads = min(num_threads, len(batch_arguments)) if generator_style: return _run_program_batch_generator( program, backend, batch_arguments, default_sampling_para, num_threads, progress_bar, ) # Original code path when generator_style=False if num_threads == 1: rets = [] if progress_bar: for arguments in tqdm.tqdm(batch_arguments): rets.append( run_program( program, backend, (), arguments, default_sampling_para, False, True, ) ) else: for arguments in batch_arguments: rets.append( run_program( program, backend, (), arguments, default_sampling_para, False, True, ) ) else: if progress_bar: pbar = tqdm.tqdm(total=len(batch_arguments)) with ThreadPoolExecutor(num_threads) as executor: futures = [] for arguments in batch_arguments: futures.append( executor.submit( run_program, program, backend, (), arguments, default_sampling_para, False, True, ) ) if progress_bar: futures[-1].add_done_callback(lambda _: pbar.update()) rets = [f.result() for f in futures] rets[-1].sync() if progress_bar: pbar.close() return rets def _run_program_batch_generator( program, backend, batch_arguments, default_sampling_para, num_threads, progress_bar, ): """Helper function that yields results one by one using chunking to avoid overwhelming ThreadPoolExecutor.""" if num_threads == 1: iterator = tqdm.tqdm(batch_arguments) if progress_bar else batch_arguments for arguments in iterator: yield run_program( program, backend, (), arguments, default_sampling_para, False, True, ) else: pbar = tqdm.tqdm(total=len(batch_arguments)) if progress_bar else None # Process in chunks to avoid overwhelming ThreadPoolExecutor # Otherwise, ThreadPoolExecutor.submit will block after adding certain number of tasks # so we will never reach "yield" until all tasks are done chunk_size = 200 with ThreadPoolExecutor(num_threads) as executor: for chunk_start in range(0, len(batch_arguments), chunk_size): chunk_end = min(chunk_start + chunk_size, len(batch_arguments)) chunk_futures = [] # Submit chunk of tasks for i in range(chunk_start, chunk_end): future = executor.submit( run_program, program, backend, (), batch_arguments[i], default_sampling_para, False, True, ) if pbar: future.add_done_callback(lambda _: pbar.update()) chunk_futures.append(future) # Yield results from this chunk as they complete for future in chunk_futures: yield future.result() if pbar: pbar.close() def cache_program(program, backend): from sglang.lang.tracer import extract_prefix_by_tracing prefix = extract_prefix_by_tracing(program, backend) if prefix and len(prefix) > 64: backend.cache_prefix(prefix) class StreamExecutor: """A stream executor that executes SGL expressions in a background thread.""" def __init__( self, backend, arguments, default_sampling_para, chat_template, stream, num_api_spec_tokens=None, use_thread=True, ): from sglang.lang.backend.base_backend import BaseBackend self.sid = uuid.uuid4().hex self.backend: BaseBackend = backend self.arguments: Dict[str, Any] = arguments self.default_sampling_para = default_sampling_para self.stream = stream self.variables = {} # Dict[name: str -> value: str] self.variable_event = {} # Dict[name: str -> event: threading.Event] self.meta_info = {} # Dict[name: str -> info: str] self.is_finished = False self.error_ = None # For completion self.text_ = "" # The full text # For chat self.messages_ = [] # The messages in the OpenAI API format self.chat_template = chat_template or self.backend.get_chat_template() self.cur_role = None self.cur_role_begin_pos = None # For vision self.images_ = [] self.cur_images = [] # For fork/join self.fork_start_text_pos = None # For speculative execution self.num_api_spec_tokens = num_api_spec_tokens self.speculated_text = "" # Worker thread self.use_thread = use_thread if self.use_thread: self.queue = queue.Queue() def _run_worker_in_context(): self._thread_worker_func() self.worker = threading.Thread( target=contextvars.copy_context().run, args=(_run_worker_in_context,) ) self.worker.start() # For streaming if stream: self.stream_text_event = threading.Event() self.stream_var_event = {} else: self.stream_text_event = None self.stream_var_event = None def submit(self, expr: SglExpr): self._init_var_event(expr) if self.use_thread: self.queue.put(expr) else: self._execute(expr) def sync(self): if self.use_thread: self.queue.join() def get_var(self, name): if name in self.variable_event: self.variable_event[name].wait() return self.variables[name] def set_var(self, name, value): self.variables[name] = value def get_meta_info(self, name, timeout=None): if name in self.variable_event: got = self.variable_event[name].wait(timeout) if not got: raise TimeoutError(f"Timeout while waiting for event '{name}'") ret = self.meta_info.get(name, None) return ret def fork( self, size: int = 1, position_ids_offset: Optional[List[int]] = None, ): if size > 1 and str(self.text_): self.submit(SglCommitLazy()) self.sync() size = int(size) exes = [ StreamExecutor( self.backend, self.arguments, self.default_sampling_para, self.chat_template, self.stream, ) for _ in range(size) ] for i in range(size): exes[i].variables = dict(self.variables) exes[i].text_ = str(self.text_) exes[i].messages_ = list(self.messages_) exes[i].cur_role = self.cur_role exes[i].cur_role_begin_pos = self.cur_role_begin_pos exes[i].fork_start_text_pos = len(self.text_) exes[i].images_ = list(self.images_) # TODO(ying): handle API speculative execution return exes def text(self): self.sync() return self.text_ def messages(self): self.sync() return self.messages_ def error(self): self.sync() return self.error_ def end(self): if self.use_thread: if self.worker.is_alive(): self.queue.put(None) self.backend.end_program(self) def _thread_worker_func(self): error = None while True: expr = self.queue.get() if expr is None: self.queue.task_done() break try: self._execute(expr) except Exception as e: warnings.warn(f"Error in stream_executor: {get_exception_traceback()}") error = e break self.queue.task_done() if self.stream_text_event: self.stream_text_event.set() # Clean the queue and events if error is not None: try: while True: self.queue.task_done() self.queue.get_nowait() except queue.Empty: pass for name in self.variable_event: self.variable_event[name].set() if self.stream_var_event: for name in self.stream_var_event: self.stream_var_event[name].set() self.error_ = error if self.stream_text_event: self.stream_text_event.set() self.is_finished = True def _execute(self, other): if isinstance(other, str): other = SglConstantText(other) assert isinstance(other, SglExpr), f"{other}" if isinstance(other, SglConstantText): self._execute_fill(other.value) elif isinstance(other, SglGen): self._execute_gen(other) elif isinstance(other, SglSelect): self._execute_select(other) elif isinstance(other, SglExprList): for x in other.expr_list: self._execute(x) elif isinstance(other, SglRoleBegin): self._execute_role_begin(other) elif isinstance(other, SglRoleEnd): self._execute_role_end(other) elif isinstance(other, SglImage): self._execute_image(other) elif isinstance(other, SglVideo): self._execute_video(other) elif isinstance(other, SglVariable): self._execute_variable(other) elif isinstance(other, SglVarScopeBegin): self._execute_var_scope_begin(other) elif isinstance(other, SglVarScopeEnd): self._execute_var_scope_end(other) elif isinstance(other, SglCommitLazy): self._execute_commit_lazy_operations(other) elif isinstance(other, SglConcateAndAppend): if ( global_config.enable_parallel_encoding and self.backend.support_concate_and_append ): self._execute_concatenate_and_append_kv_cache(other) else: self._execute_concatenate_and_append_text(other) elif isinstance(other, SglSeparateReasoning): self._execute_separate_reasoning(other) else: raise ValueError(f"Unknown type: {type(other)}") def _execute_fill(self, value: str, prefix=False): value = str(value) if ( self.cur_role == "assistant" and self.num_api_spec_tokens is not None and self.backend.is_chat_model and not prefix ): self.backend.spec_fill(value) return if self.speculated_text.startswith(value): self.speculated_text = self.speculated_text[len(value) :] else: self.speculated_text = "" self.text_ += value def _execute_image(self, expr: SglImage): path = expr.path base64_data = encode_image_base64(path) self.images_.append((path, base64_data)) self.cur_images.append((path, base64_data)) self.text_ += self.chat_template.image_token def _execute_video(self, expr: SglVideo): path = expr.path num_frames = expr.num_frames base64_data = encode_video_base64(path, num_frames) self.images_.append((path, base64_data)) self.cur_images.append((path, base64_data)) self.text_ += self.chat_template.image_token def _spec_gen(self, sampling_params): stop = sampling_params.stop max_new_tokens = sampling_params.max_new_tokens meta_info = {} def regen(): nonlocal meta_info sampling_params.max_new_tokens = max( sampling_params.max_new_tokens, self.num_api_spec_tokens ) sampling_params.stop = None self.speculated_text, meta_info = self.backend.generate( self, sampling_params=sampling_params ) def find_stop(): if isinstance(stop, str): return self.speculated_text.find(stop) elif isinstance(stop, (tuple, list)): pos = -1 for stop_str in stop: stop_pos = self.speculated_text.find(stop_str) if stop_pos != -1 and (pos == -1 or stop_pos < pos): pos = stop_pos return pos else: raise Exception("Wrong type of stop in sampling parameters.") if stop is None: if len(self.speculated_text) < max_new_tokens: regen() comp = self.speculated_text[:max_new_tokens] self.speculated_text = self.speculated_text[max_new_tokens:] elif isinstance(stop, (str, list, tuple)): if self.speculated_text == "": regen() stop_pos = find_stop() if stop_pos == -1: stop_pos = min( sampling_params.max_new_tokens, len(self.speculated_text), ) comp = self.speculated_text[:stop_pos] self.speculated_text = self.speculated_text[stop_pos:] else: raise ValueError("Wrong type of stop in sampling parameters.") return comp, meta_info def _execute_gen(self, expr: SglGen): sampling_params = self._resolve_sampling_params(expr.sampling_params) name = expr.name if not self.stream: if self.num_api_spec_tokens is None: comp, meta_info = self.backend.generate( self, sampling_params=sampling_params, ) else: if self.backend.is_chat_model: # Speculative execution on models with only chat interface. # Store the calls into a temporary list. # They will be lazily executed later. comp, meta_info = self.backend.generate( self, sampling_params=sampling_params, spec_var_name=name, ) return else: # Speculative execution on models with completion interface comp, meta_info = self._spec_gen(sampling_params) if isinstance(comp, list): self.text_ += comp[0] else: assert isinstance(comp, str) self.text_ += comp self.variables[name] = comp self.meta_info[name] = meta_info self.variable_event[name].set() else: assert ( self.num_api_spec_tokens is None ), "stream is not supported with api speculative execution" generator = self.backend.generate_stream( self, sampling_params=sampling_params ) self.variables[name] = "" self.stream_var_event[name].set() for comp, meta_info in generator: self.text_ += comp self.variables[name] += comp self.meta_info[name] = meta_info self.stream_var_event[name].set() self.stream_text_event.set() self.variable_event[name].set() self.stream_var_event[name].set() def _execute_select(self, expr: SglSelect): choices_decision = self.backend.select( self, expr.choices, expr.temperature, expr.choices_method ) if expr.name is not None: name = expr.name self.variables[name] = choices_decision.decision self.meta_info[name] = choices_decision.meta_info self.variable_event[name].set() if self.stream_var_event: self.stream_var_event[name].set() self.text_ += choices_decision.decision def _execute_variable(self, expr: SglVariable): src_executor = expr.source_stream_executor value = src_executor.get_var(expr.name) self._execute_fill(value) def _execute_role_begin(self, expr: SglRoleBegin): assert self.cur_role is None, "Nested roles are not allowed." if len(self.messages_) == 0 and expr.role != "system": # Insert the default system message default_system = self.chat_template.default_system_prompt if default_system: self._execute_role_begin(SglRoleBegin("system")) self._execute_fill(default_system) self._execute_role_end(SglRoleEnd("system")) self.cur_role = expr.role prefix, _ = self.chat_template.get_prefix_and_suffix(expr.role, self.messages_) self._execute_fill(prefix, prefix=True) self.cur_role_begin_pos = len(self.text_) def _execute_role_end(self, expr: SglRoleEnd): if ( self.cur_role == "assistant" and self.num_api_spec_tokens is not None and self.backend.is_chat_model ): # Execute the stored lazy generation calls self.backend.role_end_generate(self) self.cur_role = None new_text = self.text_[self.cur_role_begin_pos :].lstrip() _, suffix = self.chat_template.get_prefix_and_suffix(expr.role, self.messages_) self._execute_fill(suffix) if self.cur_images: # OpenAI vision API format last_msg = { "role": expr.role, "content": [{"type": "text", "text": new_text}], } for image_path, image_base64_data in self.cur_images: last_msg["content"].append( { "type": "image_url", "image_url": { "url": f"data:image/jpeg;base64,{image_base64_data}" }, } ) self.messages_.append(last_msg) self.cur_images = [] else: # OpenAI chat API format self.messages_.append({"role": expr.role, "content": new_text}) def _execute_var_scope_begin(self, expr: SglVarScopeBegin): self.variables[expr.name] = int(len(self.text_)) def _execute_var_scope_end(self, expr: SglVarScopeEnd): self.variables[expr.name] = self.text_[self.variables[expr.name] :] self.variable_event[expr.name].set() def _execute_commit_lazy_operations(self, expr: SglCommitLazy): self.backend.commit_lazy_operations(self) def _execute_concatenate_and_append_text(self, expr: SglConcateAndAppend): new_text = "" for s in expr.states: exe = s.stream_executor exe.sync() new_text += exe.text_[exe.fork_start_text_pos :] self._execute_fill(new_text) def _execute_concatenate_and_append_kv_cache(self, expr: SglConcateAndAppend): self_len = len(self.text_) for i, s in enumerate(expr.states): exe = s.stream_executor exe.submit(SglCommitLazy()) for i, s in enumerate(expr.states): exe = s.stream_executor exe.sync() assert exe.fork_start_text_pos == self_len self.text_ += exe.text_[exe.fork_start_text_pos :] src_rids = [state.stream_executor.sid for state in expr.states] self.backend.concatenate_and_append(src_rids, self.sid) def _execute_separate_reasoning(self, expr: SglSeparateReasoning): if self.stream: # separate reasoning for stream is not supported return if ( self.cur_role == "assistant" and self.num_api_spec_tokens is not None and self.backend.is_chat_model ): # Execute the stored lazy generation calls self.backend.role_end_generate(self) from sglang.srt.parser.reasoning_parser import ReasoningParser reasoning_parser = ReasoningParser(expr.model_type) other = expr.expr if not other: return elif isinstance(other, SglGen) or isinstance(other, SglSelect): cur_text = self.get_var(other.name) reasoning, normal_text = reasoning_parser.parse_non_stream(cur_text) reasoning_name = expr.process_name_for_reasoning(other.name) self.set_var(other.name, normal_text) self.set_var(reasoning_name, reasoning) # the variable is ready to be used self.variable_event[reasoning_name].set() self.text_ = self.text_[: self.cur_role_begin_pos] + normal_text elif isinstance(other, SglExprList): for x in other.expr_list: self._execute_separate_reasoning( SglSeparateReasoning(expr.model_type, x) ) def _init_var_event(self, expr): if isinstance( expr, (SglGen, SglSelect, SglVarScopeBegin, SglSeparateReasoning) ): self.variable_event[expr.name] = threading.Event() if self.stream: self.stream_var_event[expr.name] = threading.Event() elif isinstance(expr, SglExprList): for e in expr.expr_list: self._init_var_event(e) def _resolve_sampling_params(self, sampling_params): """ Construct sampling param based on default + override values The default values of sampling are populated in `default_sampling_para` via sgl.function.run(...sampling_args) , and `sampling_params` contains the override values from sgl.gen(). Here we use default_sampling_para as the base and override the values if they exist in `sampling_params`. It also extends the stop tokens based on the chat template. """ # deepcopy is required because the dict has lists inside clone = copy.deepcopy(self.default_sampling_para) for item in [ "max_new_tokens", "min_new_tokens", "n", "stop", "stop_token_ids", "stop_regex", "temperature", "top_p", "top_k", "min_p", "frequency_penalty", "presence_penalty", "ignore_eos", "return_logprob", "logprob_start_len", "top_logprobs_num", "return_text_in_logprobs", "dtype", "regex", "json_schema", ]: value = getattr(sampling_params, item, None) if value is not None: setattr(clone, item, value) if self.chat_template.stop_str: if clone.stop == (): clone.stop = [] elif isinstance(clone.stop, str): clone.stop = [clone.stop] clone.stop += self.chat_template.stop_str return clone def __del__(self): self.end() class ProgramState: """The state of an SGL program.""" def __init__(self, stream_executor: StreamExecutor): self.stream_executor = stream_executor def _role_common(self, name: str, expr: Optional[SglExpr] = None): if expr is not None: role_expr = SglExprList([SglRoleBegin(name), expr, SglRoleEnd(name)]) self.stream_executor.submit(role_expr) return role_expr else: @contextmanager def role_scope(): self.stream_executor.submit(SglRoleBegin(name)) yield self.stream_executor.submit(SglRoleEnd(name)) return role_scope() def system(self, expr: Optional[SglExpr] = None): return self._role_common("system", expr) def user(self, expr: Optional[SglExpr] = None): return self._role_common("user", expr) def assistant(self, expr: Optional[SglExpr] = None): return self._role_common("assistant", expr) @contextmanager def var_scope(self, name: str): self.stream_executor.submit(SglVarScopeBegin(name)) yield self.stream_executor.submit(SglVarScopeEnd(name)) def fork( self, size: int = 1, position_ids_offset: Optional[List[int]] = None, ): stream_executors = self.stream_executor.fork(size, position_ids_offset) states = [ProgramState(x) for x in stream_executors] state_group = ProgramStateGroup(states, self) return state_group @contextmanager def copy(self, position_ids_offset: Optional[List[int]] = None): state_group = self.fork(1, position_ids_offset) try: yield state_group[0] finally: state_group.join() def text(self): return self.stream_executor.text() def messages(self): return self.stream_executor.messages() def sync(self): return self.stream_executor.sync() def error(self): return self.stream_executor.error() def text_iter(self, var_name: Optional[str] = None): if self.stream_executor.stream: prev = 0 if var_name is None: event = self.stream_executor.stream_text_event while True: event.wait() event.clear() out = str(self.stream_executor.text_[prev:]) prev += len(out) if out: yield out if self.stream_executor.is_finished: break else: event = None while not event: if var_name in self.stream_executor.stream_var_event: event = self.stream_executor.stream_var_event[var_name] if self.stream_executor.is_finished: yield "" return while True: event.wait() event.clear() out = str(self.stream_executor.variables[var_name][prev:]) prev += len(out) if out: yield out if self.stream_executor.variable_event[var_name].is_set(): break else: if var_name is None: yield self.text() else: yield self.get_var(var_name) async def text_async_iter( self, var_name: Optional[str] = None, return_meta_data: bool = False ): loop = asyncio.get_running_loop() if self.stream_executor.stream: prev = 0 if var_name is None: event = self.stream_executor.stream_text_event while True: await loop.run_in_executor(None, event.wait) event.clear() out = str(self.stream_executor.text_[prev:]) prev += len(out) if out: yield out if self.stream_executor.is_finished: break else: event = None while not event: if var_name in self.stream_executor.stream_var_event: event = self.stream_executor.stream_var_event[var_name] if self.stream_executor.is_finished: yield "" return while True: await loop.run_in_executor(None, event.wait) event.clear() out = str(self.stream_executor.variables[var_name][prev:]) prev += len(out) if out: if return_meta_data: yield out, self.stream_executor.meta_info[var_name] else: yield out if self.stream_executor.variable_event[var_name].is_set(): break else: if var_name is None: yield self.text() else: yield self.get_var(var_name) def get_var(self, name): return self.stream_executor.get_var(name) def set_var(self, name, value): return self.stream_executor.set_var(name, value) def get_meta_info(self, name): return self.stream_executor.get_meta_info(name) def __iadd__(self, other): if other is None: raise ValueError("Tried to append None to state.") self.stream_executor.submit(other) return self def __getitem__(self, name): return self.get_var(name) def __setitem__(self, name, value): self.set_var(name, value) def __contains__(self, name): return name in self.stream_executor.variables def __del__(self): self.stream_executor.end() def __repr__(self) -> str: return f"ProgramState({self.text()})" class ProgramStateGroup: def __init__( self, states: List[ProgramState], src_state: Optional[ProgramState] = None ): self.states = states self.src_state = src_state def join(self, mode: str = "gather_variable"): if mode == "gather_variable": # Copy variables back src_vars = self.src_state.stream_executor.variables src_var_set = set(src_vars.keys()) for child_state in self.states: child_state.stream_executor.sync() child_vars = child_state.stream_executor.variables new_vars = set(child_vars.keys()) - src_var_set for k in new_vars: if k in src_vars: src_vars[k].append(child_vars[k]) else: src_vars[k] = [child_vars[k]] elif mode == "concate_and_append": # Concatenate and append KV cache self.src_state += SglConcateAndAppend(self.states) # Need a sync here. Otherwise, `states` can be deleted. self.src_state.stream_executor.sync() else: raise ValueError(f"Invalid join mode: {mode}") for s in self.states: s.stream_executor.end() def __getitem__(self, i: int): return self.states[i] def __setitem__(self, i: int, value): assert self.states[i] == value def __iadd__(self, other): if isinstance(other, Callable): # lambda function for i in range(len(self.states)): self.states[i] += other(i) elif isinstance(other, SglExpr): for i in range(len(self.states)): self.states[i] += other elif isinstance(other, (list, tuple)): for i in range(len(self.states)): self.states[i] += other[i] else: raise ValueError(f"Invalid value: {other}") return self ================================================ FILE: python/sglang/lang/ir.py ================================================ """The intermediate representation.""" import dataclasses import inspect import warnings from typing import List, Optional, Union from sglang.global_config import global_config from sglang.lang.choices import ChoicesSamplingMethod REGEX_INT = r"[-+]?[0-9]+[ \n]*" REGEX_FLOAT = r"[-+]?[0-9]*\.?[0-9]+[ \n]*" REGEX_BOOL = r"(True|False)" REGEX_STR = r"\"[\w\d\s]*\"" # bugs with regex r"\".*\"" in interegular pkg @dataclasses.dataclass class SglSamplingParams: max_new_tokens: int = 128 min_new_tokens: int = 0 n: int = 1 stop: Union[str, List[str]] = () stop_token_ids: Optional[List[int]] = () stop_regex: Optional[Union[str, List[str]]] = () temperature: float = 1.0 top_p: float = 1.0 top_k: int = -1 # -1 means disable min_p: float = 0.0 frequency_penalty: float = 0.0 presence_penalty: float = 0.0 ignore_eos: bool = False return_logprob: Optional[bool] = None logprob_start_len: Optional[int] = (None,) top_logprobs_num: Optional[int] = (None,) return_text_in_logprobs: Optional[bool] = (None,) json_schema: Optional[str] = None # for constrained generation, not included in to_xxx_kwargs dtype: Optional[str] = None regex: Optional[str] = None def clone(self): return SglSamplingParams( self.max_new_tokens, self.min_new_tokens, self.n, self.stop, self.stop_token_ids, self.stop_regex, self.temperature, self.top_p, self.top_k, self.min_p, self.frequency_penalty, self.presence_penalty, self.ignore_eos, self.return_logprob, self.logprob_start_len, self.top_logprobs_num, self.return_text_in_logprobs, self.json_schema, ) def to_openai_kwargs(self): # OpenAI does not support top_k, so we drop it here if self.regex is not None: warnings.warn("Regular expression is not supported in the OpenAI backend.") return { "max_tokens": self.max_new_tokens, "max_completion_tokens": self.max_new_tokens, "n": self.n, "stop": self.stop or None, "temperature": self.temperature, "top_p": self.top_p, "frequency_penalty": self.frequency_penalty, "presence_penalty": self.presence_penalty, } def to_vertexai_kwargs(self): if self.regex is not None: warnings.warn( "Regular expression is not supported in the VertexAI backend." ) return { "candidate_count": 1, "max_output_tokens": self.max_new_tokens, "stop_sequences": self.stop, "temperature": self.temperature, "top_p": self.top_p, "top_k": self.top_k if self.top_k > 0 else None, } def to_anthropic_kwargs(self): # Anthropic does not support frequency_penalty or presence_penalty, so we drop it here if self.regex is not None: warnings.warn( "Regular expression is not supported in the Anthropic backend." ) return { "max_tokens": self.max_new_tokens, "stop_sequences": ( self.stop if isinstance(self.stop, (list, tuple)) else [self.stop] ), "temperature": self.temperature, "top_p": self.top_p, "top_k": self.top_k, } def to_litellm_kwargs(self): if self.regex is not None: warnings.warn("Regular expression is not supported in the LiteLLM backend.") return { "max_tokens": self.max_new_tokens, "stop": self.stop or None, "temperature": self.temperature, "top_p": self.top_p, "frequency_penalty": self.frequency_penalty, "presence_penalty": self.presence_penalty, } def to_srt_kwargs(self): return { "max_new_tokens": self.max_new_tokens, "min_new_tokens": self.min_new_tokens, "n": self.n, "stop": self.stop, "stop_token_ids": self.stop_token_ids, "stop_regex": self.stop_regex, "temperature": self.temperature, "top_p": self.top_p, "top_k": self.top_k, "min_p": self.min_p, "frequency_penalty": self.frequency_penalty, "presence_penalty": self.presence_penalty, "ignore_eos": self.ignore_eos, "regex": self.regex, "json_schema": self.json_schema, } class SglFunction: def __init__(self, func, num_api_spec_tokens=None, bind_arguments=None): self.func = func self.num_api_spec_tokens = num_api_spec_tokens self.bind_arguments = bind_arguments or {} self.pin_prefix_rid = None # Parse arguments argspec = inspect.getfullargspec(func) assert argspec.args[0] == "s", 'The first argument must be "s"' self.arg_names = argspec.args[1:] self.arg_defaults = argspec.defaults if argspec.defaults is not None else [] def bind(self, **kwargs): assert all(key in self.arg_names for key in kwargs) new_bind_dict = {**self.bind_arguments, **kwargs} return SglFunction(self.func, bind_arguments=new_bind_dict) def run( self, *args, max_new_tokens: int = 128, n: int = 1, stop: Optional[Union[str, List[str]]] = None, stop_token_ids: Optional[List[int]] = None, stop_regex: Optional[Union[str, List[str]]] = None, temperature: float = 1.0, top_p: float = 1.0, top_k: int = -1, min_p: float = 0.0, frequency_penalty: float = 0.0, presence_penalty: float = 0.0, ignore_eos: bool = False, return_logprob: Optional[bool] = None, logprob_start_len: Optional[int] = None, top_logprobs_num: Optional[int] = None, return_text_in_logprobs: Optional[bool] = None, stream: bool = False, backend=None, use_thread: bool = True, **kwargs, ): from sglang.lang.interpreter import run_program # avoid using [] as the default arg: https://nikos7am.com/posts/mutable-default-arguments/ if stop is None: stop = [] if stop_token_ids is None: stop_token_ids = [] if stop_regex is None: stop_regex = [] default_sampling_para = SglSamplingParams( max_new_tokens=max_new_tokens, n=n, stop=stop, stop_token_ids=stop_token_ids, stop_regex=stop_regex, temperature=temperature, top_p=top_p, top_k=top_k, min_p=min_p, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, ignore_eos=ignore_eos, return_logprob=return_logprob, logprob_start_len=logprob_start_len, top_logprobs_num=top_logprobs_num, return_text_in_logprobs=return_text_in_logprobs, ) backend = backend or global_config.default_backend return run_program( self, backend, args, kwargs, default_sampling_para, stream, use_thread=use_thread, ) def run_batch( self, batch_kwargs, *, max_new_tokens: int = 128, n: int = 1, stop: Optional[Union[str, List[str]]] = None, stop_token_ids: Optional[List[int]] = None, stop_regex: Optional[Union[str, List[str]]] = None, temperature: float = 1.0, top_p: float = 1.0, top_k: int = -1, min_p: float = 0.0, frequency_penalty: float = 0.0, presence_penalty: float = 0.0, ignore_eos: bool = False, return_logprob: Optional[bool] = None, logprob_start_len: Optional[int] = None, top_logprobs_num: Optional[int] = None, return_text_in_logprobs: Optional[bool] = None, backend=None, num_threads: Union[str, int] = "auto", progress_bar: bool = False, generator_style: bool = False, ): from sglang.lang.interpreter import run_program_batch if stop is None: stop = [] if stop_token_ids is None: stop_token_ids = [] if stop_regex is None: stop_regex = [] assert isinstance(batch_kwargs, (list, tuple)) if len(batch_kwargs) == 0: return [] if not isinstance(batch_kwargs[0], dict): num_programs = len(batch_kwargs) # change the list of argument values to dict of arg_name -> arg_value batch_kwargs = [ {self.arg_names[i]: v for i, v in enumerate(arg_values)} for arg_values in batch_kwargs if isinstance(arg_values, (list, tuple)) and len(self.arg_names) - len(self.arg_defaults) <= len(arg_values) <= len(self.arg_names) ] # Ensure to raise an exception if the number of arguments mismatch if len(batch_kwargs) != num_programs: raise Exception("Given arguments mismatch the SGL function signature") default_sampling_para = SglSamplingParams( max_new_tokens=max_new_tokens, n=n, stop=stop, stop_token_ids=stop_token_ids, stop_regex=stop_regex, temperature=temperature, top_p=top_p, top_k=top_k, min_p=min_p, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, ignore_eos=ignore_eos, return_logprob=return_logprob, logprob_start_len=logprob_start_len, top_logprobs_num=top_logprobs_num, return_text_in_logprobs=return_text_in_logprobs, ) backend = backend or global_config.default_backend return run_program_batch( self, backend, batch_kwargs, default_sampling_para, num_threads, progress_bar, generator_style=generator_style, ) def trace(self, *, backend=None, **kwargs): from sglang.lang.tracer import trace_program backend = backend or global_config.default_backend return trace_program(self, kwargs, backend) def cache(self, backend=None): from sglang.lang.interpreter import cache_program backend = backend or global_config.default_backend return cache_program(self, backend) def __call__(self, *args, **kwargs): from sglang.lang.tracer import TracingScope tracing_scope = TracingScope.get_current_scope() if tracing_scope is None: return self.run(*args, **kwargs) else: kwargs["backend"] = tracing_scope.tracer_state.backend return self.trace(*args, **kwargs) class SglExpr: node_ct = 0 def __init__(self): self.node_id = SglExpr.node_ct self.prev_node = None self.pid = None SglExpr.node_ct += 1 def __add__(self, other): if isinstance(other, str): other = SglConstantText(other) assert isinstance(other, SglExpr) return self.concatenate_ir(self, other) def __radd__(self, other): if isinstance(other, str): other = SglConstantText(other) assert isinstance(other, SglExpr), f"{other}" return self.concatenate_ir(other, self) def concatenate_ir(self, a, b): if isinstance(a, SglExprList): if isinstance(b, SglExprList): return SglExprList(a.expr_list + b.expr_list) else: return SglExprList(a.expr_list + [b]) elif isinstance(b, SglExprList): return SglExprList([a] + b.expr_list) return SglExprList([a, b]) def print_graph_dfs(self): ret = [""] visited = set() def dfs_print(x): if x is None or x in visited: return visited.add(x) # Print dependency if x.prev_node is not None: dfs_print(x.prev_node) if isinstance(x, SglExprList): for y in x.expr_list: dfs_print(y) # elif isinstance(x, SglRole): # dfs_print(x.expr) elif isinstance(x, SglVariable): dfs_print(x.source) # Print the node itself if isinstance(x, (SglFork, SglGetForkItem)): ret[0] += f"%{x.node_id} = {x}\n" else: if x.prev_node is not None: ret[0] += ( f"%{x.node_id} = %{x.prev_node.node_id} + " + str(x) + "\n" ) else: ret[0] += f"%{x.node_id} = " + str(x) + "\n" dfs_print(self) return ret[0] class SglExprList(SglExpr): def __init__(self, expr_list: List[SglExpr]): super().__init__() self.expr_list = expr_list def __repr__(self): return f"ExprList({self.expr_list})" class SglArgument(SglExpr): def __init__(self, name: str, value: str): super().__init__() self.name = name self.value = value def __repr__(self): return f"Argument(name={self.name}, value={repr(self.value)})" def __len__(self): return len(self.value) def __getitem__(self, i): return self.value[i] def __int__(self): return self.value def __bool__(self): return self.value def __format__(self, *args): raise TypeError( "Cannot put argument inside a f-string. " "This is not compatible with the tracer. " ) class SglImage(SglExpr): def __init__(self, path: str): self.path = path def __repr__(self) -> str: return f"SglImage({self.path})" class SglVideo(SglExpr): def __init__(self, path: str, num_frames: int): self.path = path self.num_frames = num_frames def __repr__(self) -> str: return f"SglVideo({self.path}, {self.num_frames})" class SglGen(SglExpr): def __init__( self, name: Optional[str] = None, max_new_tokens: Optional[int] = None, min_new_tokens: Optional[int] = None, n: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, stop_token_ids: Optional[List[int]] = None, stop_regex: Optional[Union[str, List[str]]] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, top_k: Optional[int] = None, min_p: Optional[float] = None, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, ignore_eos: Optional[bool] = None, return_logprob: Optional[bool] = None, logprob_start_len: Optional[int] = None, top_logprobs_num: Optional[int] = None, return_text_in_logprobs: Optional[bool] = None, dtype: Optional[type] = None, regex: Optional[str] = None, json_schema: Optional[str] = None, ): """Call the model to generate. See the meaning of the arguments in docs/backend/sampling_params.md""" super().__init__() self.name = name self.sampling_params = SglSamplingParams( max_new_tokens=max_new_tokens, min_new_tokens=min_new_tokens, n=n, stop=stop, stop_regex=stop_regex, stop_token_ids=stop_token_ids, temperature=temperature, top_p=top_p, top_k=top_k, min_p=min_p, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, ignore_eos=ignore_eos, return_logprob=return_logprob, logprob_start_len=logprob_start_len, top_logprobs_num=top_logprobs_num, return_text_in_logprobs=return_text_in_logprobs, dtype=dtype, regex=regex, json_schema=json_schema, ) def __repr__(self): return f"Gen('{self.name}')" class SglConstantText(SglExpr): def __init__(self, value: str): super().__init__() self.value = value def __repr__(self): return f"Constant({repr(self.value)})" class SglRoleBegin(SglExpr): def __init__(self, role: str): super().__init__() self.role = role def __repr__(self): return f"RoleBegin({self.role})" class SglRoleEnd(SglExpr): def __init__(self, role: str): super().__init__() self.role = role def __repr__(self): return f"RoleEnd({self.role})" class SglSelect(SglExpr): def __init__( self, name: str, choices: List[str], temperature: float, choices_method: ChoicesSamplingMethod, ): super().__init__() self.name = name self.choices = choices self.temperature = temperature self.choices_method = choices_method def __repr__(self): return f"Select({self.name}, choices={self.choices}, choices_method={self.choices_method})" class SglFork(SglExpr): def __init__(self, number: int, position_ids_offset=None): super().__init__() self.number = number self.position_ids_offset = position_ids_offset def __repr__(self): return ( f"Fork(%{self.prev_node.node_id}, number={self.number}, " f"position_ids_offset={self.position_ids_offset})" ) class SglGetForkItem(SglExpr): def __init__(self, index: int): super().__init__() self.index = index def __repr__(self): return f"GetForkItem(%{self.prev_node.node_id}, index={self.index})" class SglVariable(SglExpr): def __init__(self, name: str, source): super().__init__() self.name = name self.source = source def __repr__(self): return f"Variable('{self.name}', source=%{self.source.node_id})" class SglVarScopeBegin(SglExpr): def __init__(self, name: str): super().__init__() self.name = name def __repr__(self): return f"VarScopeBegin('{self.name}')" class SglVarScopeEnd(SglExpr): def __init__(self, name: str): super().__init__() self.name = name def __repr__(self): return f"VarScopeEnd('{self.name}')" class SglConcateAndAppend(SglExpr): def __init__(self, states): super().__init__() self.states = states def __repr__(self): return f"ConcatenateAndAppend('{self.states}')" class SglCommitLazy(SglExpr): def __init__(self): super().__init__() def __repr__(self): return "CommitLazy()" class SglSeparateReasoning(SglExpr): def __init__(self, model_type: str, expr: SglExpr): super().__init__() self.model_type = model_type self.expr = expr self.name = None self._process_expr(expr) def process_name_for_reasoning(self, name): if not name: raise ValueError("name must be provided") return f"{name}_reasoning_content" def _process_expr(self, expr): if isinstance(expr, SglGen): self.name = self.process_name_for_reasoning(expr.name) elif isinstance(expr, SglSelect): self.name = self.process_name_for_reasoning(expr.name) elif isinstance(expr, SglExprList): for x in expr.expr_list: self._process_expr(x) def __repr__(self): return f"SeparateReasoning(model_type={self.model_type}, name={self.name})" ================================================ FILE: python/sglang/lang/tracer.py ================================================ """Tracing a program.""" import uuid from typing import Any, Dict, List, Optional from sglang.lang.backend.base_backend import BaseBackend from sglang.lang.interpreter import ProgramState, ProgramStateGroup from sglang.lang.ir import ( SglArgument, SglConstantText, SglExpr, SglExprList, SglFork, SglGen, SglGetForkItem, SglRoleBegin, SglRoleEnd, SglSelect, SglVariable, SglVarScopeBegin, SglVarScopeEnd, ) class StopTracing(Exception): pass def extract_prefix_by_tracing(program, backend): # Create dummy arguments dummy_arguments = {name: SglArgument(name, None) for name in program.arg_names} arguments = dummy_arguments arguments.update(program.bind_arguments) # Trace tracer = TracerProgramState(backend, arguments, only_trace_prefix=True) try: with TracingScope(tracer): tracer.ret_value = program.func(tracer, **arguments) except (StopTracing, TypeError, AttributeError): # Some exceptions may not be caught pass # Run and cache prefix prefix = "" for expr in tracer.flatten_nodes(): if isinstance(expr, SglConstantText): prefix += expr.value else: break return prefix def trace_program(program, arguments, backend): # Create dummy backend if backend is None: backend = BaseBackend() # Create dummy arguments dummy_arguments = { name: SglArgument(name, None) for name in program.arg_names if name not in arguments } arguments.update(dummy_arguments) arguments.update(program.bind_arguments) # Trace tracer = TracerProgramState(backend, arguments, only_trace_prefix=False) with TracingScope(tracer): tracer.ret_value = program.func(tracer, **arguments) return tracer class TracerProgramState(ProgramState): def __init__(self, backend, arguments, only_trace_prefix): self.pid = uuid.uuid4().hex self.backend = backend self.arguments: Dict[str, Any] = arguments self.only_trace_prefix = only_trace_prefix if hasattr(backend, "endpoint"): self.backend = backend.endpoint self.nodes = [] self.last_node = None self.variables = {} self.ret_value = None # For completion # For chat self.messages_ = [] self.cur_role = None self.chat_template = self.backend.get_chat_template() # For multi states self.child_states = [] cur_scope = TracingScope.get_current_scope() if cur_scope is not None: cur_scope.add_child_state(self) ################################## ########### Public API ########### ################################## def fork(self, size: int = 1, position_ids_offset: Optional[List[int]] = None): assert size >= 1 if self.only_trace_prefix: raise StopTracing() fork_node = SglFork(size) fork_node.prev_node = self.last_node states = [ TracerProgramState(self.backend, self.arguments, self.only_trace_prefix) for _ in range(size) ] for i in range(size): node = SglGetForkItem(i) node.prev_node = fork_node states[i].last_node = node states[i].variables = dict(self.variables) states[i].messages_ = list(self.messages_) states[i].cur_role = self.cur_role states[i].chat_template = self.chat_template state_group = ProgramStateGroup(states, self) return state_group ################################## ########## Internal API ########## ################################## def _append_node(self, other: SglExpr): self.nodes.append(other) other.prev_node = self.last_node self.last_node = other def _execute(self, other: SglExpr): if isinstance(other, str): other = SglConstantText(other) other.pid = self.pid if isinstance(other, SglConstantText): self._execute_fill(other) elif isinstance(other, SglGen): self._execute_gen(other) elif isinstance(other, SglSelect): self._execute_select(other) elif isinstance(other, SglExprList): for x in other.expr_list: self._execute(x) elif isinstance(other, SglRoleBegin): self._execute_role_begin(other) elif isinstance(other, SglRoleEnd): self._execute_role_end(other) elif isinstance(other, SglVarScopeBegin): self._execute_var_scope_begin(other) elif isinstance(other, SglVarScopeEnd): self._execute_var_scope_end(other) else: if self.only_trace_prefix: raise StopTracing() else: self._append_node(other) return self def __iadd__(self, other): self._execute(other) return self def _execute_fill(self, expr: SglConstantText): if isinstance(expr, str): expr = SglConstantText(expr) self._append_node(expr) def _execute_gen(self, expr: SglGen): name = expr.name if expr.name is not None else "gen_" + str(len(self.variables)) new_node = SglVariable(name, source=expr) self.variables[name] = new_node self._append_node(expr) def _execute_select(self, expr: SglSelect): name = ( expr.name if expr.name is not None else "select_" + str(len(self.variables)) ) new_node = SglVariable(name, source=expr) self.variables[name] = new_node self._append_node(expr) def _execute_role_begin(self, expr: SglRoleBegin): assert self.cur_role is None, "Nested roles are not allowed." if len(self.messages_) == 0 and expr.role != "system": # Insert default system message default_system = self.chat_template.default_system_prompt if default_system: self._execute_role_begin(SglRoleBegin("system")) self._execute_fill(default_system) self._execute_role_end(SglRoleEnd("system")) self.cur_role = expr.role prefix, suffix = self.chat_template.get_prefix_and_suffix( expr.role, self.messages_ ) self._execute_fill(prefix) def _execute_role_end(self, expr: SglRoleEnd): prefix, suffix = self.chat_template.get_prefix_and_suffix( expr.role, self.messages_ ) self._execute_fill(suffix) self.messages_.append({"role": expr.role, "content": ""}) self.cur_role = None def _execute_var_scope_end(self, expr: SglVarScopeEnd): new_node = SglVariable(expr.name, source=self.last_node) self.variables[expr.name] = new_node def get_var(self, name): ret = self.arguments.get(name, None) if ret is not None: return ret v = self.variables[name] return SglVariable(v.name, v.source) def flatten_nodes(self): def traverse(cur): if isinstance(cur, SglExprList): for child in cur.expr_list: traverse(child) else: ret.append(cur) ret = [] for x in self.nodes: traverse(x) return ret def __del__(self): pass class TracingScope: cur_scope = None def __init__(self, tracer_state: TracerProgramState): self.tracer_state = tracer_state self.last_scope = TracingScope.cur_scope def __enter__(self): TracingScope.cur_scope = self return self def __exit__(self, exc_type, exc_value, traceback): TracingScope.cur_scope = self.last_scope @staticmethod def get_current_scope(): return TracingScope.cur_scope def add_child_state(self, state: TracerProgramState): cur_scope = self while cur_scope is not None: cur_scope.tracer_state.child_states.append(state) cur_scope = cur_scope.last_scope ================================================ FILE: python/sglang/launch_server.py ================================================ """Launch the inference server.""" import asyncio import os import sys import warnings from sglang.srt.server_args import prepare_server_args from sglang.srt.utils import kill_process_tree from sglang.srt.utils.common import suppress_noisy_warnings suppress_noisy_warnings() def run_server(server_args): """Run the server based on server_args.grpc_mode and server_args.encoder_only.""" if server_args.encoder_only: # For encoder disaggregation if server_args.grpc_mode: from sglang.srt.disaggregation.encode_grpc_server import ( serve_grpc_encoder, ) asyncio.run(serve_grpc_encoder(server_args)) else: from sglang.srt.disaggregation.encode_server import launch_server launch_server(server_args) elif server_args.grpc_mode: from sglang.srt.entrypoints.grpc_server import serve_grpc asyncio.run(serve_grpc(server_args)) elif server_args.use_ray: try: from sglang.srt.ray.http_server import launch_server except ImportError: raise ImportError( "Ray is required for --use-ray mode. " "Install it with: pip install 'sglang[ray]'" ) launch_server(server_args) else: # Default mode: HTTP mode. from sglang.srt.entrypoints.http_server import launch_server launch_server(server_args) if __name__ == "__main__": warnings.warn( "'python -m sglang.launch_server' is still supported, but " "'sglang serve' is the recommended entrypoint.\n" " Example: sglang serve --model-path [options]", UserWarning, stacklevel=1, ) server_args = prepare_server_args(sys.argv[1:]) try: run_server(server_args) finally: kill_process_tree(os.getpid(), include_parent=False) ================================================ FILE: python/sglang/multimodal_gen/.claude/CLAUDE.md ================================================ # CLAUDE.md — sglang-diffusion (multimodal_gen) ## What is this? SGLang's diffusion/multimodal generation subsystem. Separate from the LLM runtime (`srt`). Supports 20+ image/video diffusion models (Wan, FLUX, HunyuanVideo, LTX, Qwen-Image, etc.) with distributed inference, LoRA, and multiple attention backends. ## Quick Start ```bash # One-shot generation sglang generate --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers --prompt "A curious raccoon" --save-output # Start server sglang serve --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers --num-gpus 4 # Python API from sglang import DiffGenerator gen = DiffGenerator.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") result = gen.generate(sampling_params_kwargs={"prompt": "A curious raccoon"}) ``` ## Architecture ``` CLI / Python API / HTTP Server (FastAPI + OpenAI-compatible) ↓ ZMQ Scheduler (request queue, batching, dispatch) ↓ multiprocessing pipes GPU Worker(s) → ComposedPipeline (stages: TextEncode → Denoise → Decode) ``` ### Key Directories ``` runtime/ ├── entrypoints/ # CLI (generate/serve), HTTP server, Python API (DiffGenerator) ├── managers/ # scheduler.py, gpu_worker.py ├── pipelines_core/ # ComposedPipelineBase, stages/, schedule_batch.py (Req/OutputBatch) ├── pipelines/ # Model-specific pipelines (wan, flux, hunyuan, ltx, qwen_image, ...) ├── models/ # encoders/, dits/, vaes/, schedulers/ ├── layers/ # attention/, lora/, quantization/ ├── loader/ # Model loading, weight utils ├── server_args.py # ServerArgs (all CLI/config params) └── distributed/ # Multi-GPU (TP, SP: ulysses/ring) configs/ ├── pipeline_configs/ # Per-model pipeline configs ├── sample/ # SamplingParams └── models/ # DiT, VAE, Encoder configs ``` ### Key Classes | Class | Location | Purpose | |-------|----------|---------| | `DiffGenerator` | `runtime/entrypoints/diffusion_generator.py` | Python API entry point | | `ComposedPipelineBase` | `runtime/pipelines_core/composed_pipeline_base.py` | Pipeline orchestrator (stages) | | `Scheduler` | `runtime/managers/scheduler.py` | ZMQ event loop, request dispatch | | `GPUWorker` | `runtime/managers/gpu_worker.py` | GPU inference worker | | `Req` / `OutputBatch` | `runtime/pipelines_core/schedule_batch.py` | Request/output containers | | `ServerArgs` | `runtime/server_args.py` | All config params | | `SamplingParams` | `configs/sample/sampling_params.py` | Generation params | | `PipelineConfig` | `configs/pipeline_configs/base.py` | Model structure config | ### Key Functions | Function | Module | Purpose | |----------|--------|---------| | `build_pipeline()` | `runtime/pipelines_core/__init__.py` | Instantiate pipeline from model_path | | `get_model_info()` | `registry.py` | Resolve pipeline + config classes | | `launch_server()` | `runtime/launch_server.py` | Start multi-process server | ## Adding a New Model 1. Create pipeline in `runtime/pipelines/` extending `ComposedPipelineBase` 2. Define stages via `create_pipeline_stages()` (TextEncoding → Denoising → Decoding) 3. Add config in `configs/pipeline_configs/` 4. Register in `registry.py` via `register_configs()` ## Multi-GPU ```bash # Sequence parallelism (video frames across GPUs) sglang serve --model-path ... --num-gpus 4 --ulysses-degree 2 --ring-degree 2 # Tensor parallelism (model layers across GPUs) sglang serve --model-path ... --num-gpus 2 --tp-size 2 ``` ## Testing ```bash # Tests live in test/ subdirectory python -m pytest python/sglang/multimodal_gen/test/ # No need to pre-download models — auto-downloaded at runtime # Dependencies assumed already installed via `pip install -e "python[diffusion]"` ``` ## Performance Tuning For questions about optimal performance, fastest commands, VRAM reduction, or best flag combinations for a given model/GPU setup, **read the [diffusion-optimal-perf skill](skills/diffusion-optimal-perf/SKILL.md)**. It contains a complete table of all lossless and lossy optimization flags with trade-offs, quick recipes, and tips. ### Perf Measurement Look for `Pixel data generated successfully in xxxx seconds` in console output. With warmup enabled, use the line containing `warmup excluded` for accurate timing. ## Env Vars Defined in `envs.py` (300+ vars). Key ones: - `SGLANG_DIFFUSION_ATTENTION_BACKEND` — attention backend override - `SGLANG_CACHE_DIT_ENABLED` — enable Cache-DiT acceleration - `SGLANG_CLOUD_STORAGE_TYPE` — cloud output storage (s3, etc.) ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/SKILL.md ================================================ --- name: diffusion-kernel description: Index for SGLang Diffusion kernel development skills. --- # Diffusion Kernel Skills ## Rule: Follow User Kernel Language Preference If the user explicitly states a preference for **Triton** or **CUDA**, follow that preference when implementing and optimizing kernels (even if the other option could work). Do not “pick for convenience”. ## Directory Layout ``` python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/ ├── SKILL.md ├── add-triton-kernel.md ├── add-cuda-kernel.md ├── diffusion-benchmark-and-profile.md ├── nsight-profiler.md ├── use-efficient-diffusion-kernels.md ├── references/ │ ├── kernel-templates.md # Copy-paste CUDA kernel templates (sglang JIT style) │ ├── troubleshooting.md # Build/perf/integration issues & fixes │ ├── h100-optimization-guide.md # H100 (sm_90) deep dive │ ├── a100-optimization-guide.md # A100 (sm_80) deep dive │ └── t4-optimization-guide.md # T4 (sm_75, FP16 only) deep dive └── scripts/ ├── bench_diffusion_rmsnorm.py # RMSNorm micro-benchmark vs PyTorch └── bench_diffusion_denoise.py # End-to-end denoise benchmark (sglang generate) ``` ## Index Before running any benchmark, profiler, or kernel-validation command, use `scripts/diffusion_skill_env.py` to derive the repo root from `sglang.__file__`, verify the repo is writable, export `FLASHINFER_DISABLE_VERSION_CHECK=1`, and choose idle GPU(s) before starting perf work. - [scripts/diffusion_skill_env.py](scripts/diffusion_skill_env.py) Shared preflight helper for all diffusion skill commands. Use it to print the repo root, create benchmark/profile output directories, and choose idle GPUs before running `sglang generate`, torch profiler, nsys, or ncu. - [add-triton-kernel.md](./add-triton-kernel.md) Step-by-step guide for adding a new Triton kernel to SGLang Diffusion's `jit_kernel/diffusion/triton/` module, including authoring, autotune, `torch.compile` compatibility, integration, and tests. Use for fused elementwise ops, norm variants, RoPE variants, or when NPU/CPU fallback is needed. - [add-cuda-kernel.md](./add-cuda-kernel.md) Step-by-step guide for adding a JIT CUDA kernel. CUDA source goes in `jit_kernel/csrc/diffusion/.cuh`; Python wrapper at `jit_kernel/diffusion/.py`. Uses SGLang's JIT compilation system (`load_jit`, `cache_once`) and internal abstractions (`TensorMatcher`, `device::AlignedVector`, `host::LaunchKernel`, `device::warp::reduce_sum`). Use for bandwidth-bound reductions (RMSNorm, LayerNorm) or ops needing fine-grained vectorization and shared memory control. Adapted from [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels). - [use-efficient-diffusion-kernels.md](./use-efficient-diffusion-kernels.md) Practical guidance for using SGLang Diffusion fused kernels and fast CUDA paths, including constraints, fallbacks, and where the fused ops are wired into the runtime. - [diffusion-benchmark-and-profile.md](./diffusion-benchmark-and-profile.md) Denoise-stage benchmark and profiling guide for SGLang Diffusion models. Three profiling levels: Level 1 (torch.profiler — kernel time ranking), Level 2 (nsys — category breakdown), Level 3 (ncu — per-kernel bandwidth/occupancy/roofline analysis). **ncu is critical for kernel optimization** — always use it when writing or tuning custom kernels to verify hardware saturation. - [nsight-profiler.md](./nsight-profiler.md) Advanced profiling skill for NVIDIA Nsight Systems / Nsight Compute: collecting traces, reading reports, and interpreting kernel-level performance metrics. ## References (GPU optimization guides, templates, troubleshooting) Loaded by `add-cuda-kernel.md`. Adapted from [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels). - [references/kernel-templates.md](references/kernel-templates.md) — copy-paste ready sglang JIT CUDA templates: element-wise (SiLU), row-reduction (RMSNorm), fused AdaLN, Python wrapper, test, benchmark - [references/troubleshooting.md](references/troubleshooting.md) — build errors, performance issues, torch.compile compatibility, kernel injection pitfalls - [references/h100-optimization-guide.md](references/h100-optimization-guide.md) — H100 (sm_90): AlignedVector benchmarks, warp reductions, occupancy, TMA, PDL - [references/a100-optimization-guide.md](references/a100-optimization-guide.md) — A100 (sm_80): cp.async, TF32, 2:4 sparsity, H100→A100 migration checklist - [references/t4-optimization-guide.md](references/t4-optimization-guide.md) — T4 (sm_75): FP16 only, 320 GB/s bandwidth, 64 KB shared mem, 16 GB memory management ## Scripts (runnable benchmarks) - [scripts/diffusion_skill_env.py](scripts/diffusion_skill_env.py) — preflight helper: repo root discovery via `sglang.__file__`, write-access probe, benchmark/profile output directories, idle GPU selection - [scripts/bench_diffusion_rmsnorm.py](scripts/bench_diffusion_rmsnorm.py) — RMSNorm micro-benchmark: JIT CUDA vs PyTorch, correctness check, bandwidth efficiency analysis - [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark via `sglang generate`, baseline vs custom kernels comparison table ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/add-cuda-kernel.md ================================================ --- name: add-cuda-kernel description: Step-by-step guide for adding a new JIT CUDA kernel to SGLang Diffusion. CUDA source files go in jit_kernel/csrc/diffusion/.cuh; Python wrapper at jit_kernel/diffusion/.py. Use when implementing optimized CUDA kernels for diffusion model operators (RMSNorm, RoPE, AdaLN, GEGLU, etc.) on NVIDIA GPUs (H100, A100). Covers kernel authoring with sglang abstractions, JIT compilation, Python wrapper, integration into the denoise stage, and benchmarking. Adapted from https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels. --- # Adding a CUDA Kernel to SGLang Diffusion (JIT Style) > **Origin**: This skill is adapted from the [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels), rewritten to follow SGLang's JIT compilation system and internal abstractions. > > **Run environment first**: before compiling, benchmarking, or profiling any kernel from this guide, use `scripts/diffusion_skill_env.py` (or the setup block in `diffusion-benchmark-and-profile.md`) to `cd` to the repo root resolved from `sglang.__file__`, verify write access, export `FLASHINFER_DISABLE_VERSION_CHECK=1`, and pick an idle GPU. > > **Extended references** (in this directory's `references/` and `scripts/`): > - [references/kernel-templates.md](references/kernel-templates.md) — copy-paste ready templates for element-wise, row-reduction (RMSNorm), fused AdaLN > - [references/troubleshooting.md](references/troubleshooting.md) — build errors, perf issues, integration pitfalls > - [references/h100-optimization-guide.md](references/h100-optimization-guide.md) — H100 (sm_90) deep dive > - [references/a100-optimization-guide.md](references/a100-optimization-guide.md) — A100 (sm_80) deep dive > - [references/t4-optimization-guide.md](references/t4-optimization-guide.md) — T4 (sm_75, FP16 only) deep dive > - [scripts/bench_diffusion_rmsnorm.py](scripts/bench_diffusion_rmsnorm.py) — RMSNorm micro-benchmark vs PyTorch > - [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark with/without kernels ## When to Use CUDA vs Triton | Scenario | Use | |----------|-----| | Fused elementwise / norm variants / RoPE | **Triton** (`add-triton-kernel.md`) — faster iteration | | Bandwidth-bound reduction (RMSNorm, LayerNorm) requiring max vectorization | **CUDA** — full control over `__nv_bfloat162` / `float4` vectorization | | Attention pattern or tile-based ops needing shared memory tuning | **CUDA** — warp-level primitives, shared memory layout | | Prototype or NPU/CPU fallback needed | **Triton** — portable across backends | For most diffusion-model elementwise ops, **start with Triton**. Switch to CUDA when profiling shows Triton can't reach hardware bandwidth limits. ## Directory Layout ``` python/sglang/jit_kernel/ ├── csrc/ │ ├── diffusion/ # JIT CUDA source files for diffusion kernels (this skill) │ │ ├── timestep_embedding.cuh # existing example │ │ ├── rmsnorm.cuh # NEW: add here │ │ └── adaln.cuh # NEW: add here │ └── elementwise/ # shared JIT CUDA csrc (non-diffusion) ├── diffusion/ │ ├── triton/ # Triton kernels (scale_shift, norm, rope, ...) │ ├── cutedsl/ # CuTe DSL kernels │ └── rmsnorm.py # NEW: CUDA JIT Python wrapper (add here) ├── timestep_embedding.py # existing CUDA diffusion kernel Python wrapper (legacy) ``` New diffusion CUDA kernel source files go into `python/sglang/jit_kernel/csrc/diffusion/.cuh`. The Python wrapper goes at `python/sglang/jit_kernel/diffusion/.py` (inside `diffusion/`, alongside the `triton/` and `cutedsl/` subdirectories). --- ## SGLang Kernel Abstractions (Required) Always use these — do **not** use raw CUDA primitives directly. ```cpp #include // TensorMatcher, SymbolicSize, SymbolicDevice #include // fp16_t, bf16_t, fp32_t, dtype_trait, packed_t #include // RuntimeCheck, div_ceil #include // LaunchKernel, SGL_DEVICE, type aliases #include // AlignedVector — 128-bit vector loads #include // warp::reduce_sum, warp::reduce_max #include // device::math::rsqrt, sqrt, ... #include // tile::Memory (strided access pattern) ``` Key types: `fp16_t` = `__half`, `bf16_t` = `__nv_bfloat16`, `fp32_t` = `float`. Packed variants: `fp16x2_t`, `bf16x2_t`. Use `packed_t` for the 2-element alias. --- ## Step 1: Write the CUDA Kernel Create `python/sglang/jit_kernel/csrc/diffusion/rmsnorm.cuh` (RMSNorm as example). ### 1a. Vectorized RMSNorm Kernel ```cpp #include #include #include #include #include #include #include #include namespace { // --------------------------------------------------------------- // RMSNorm kernel: y = x / rms(x) * weight // T = fp16_t | bf16_t | fp32_t // kVecN = vectorized elements per load (8 for fp16/bf16, 4 for fp32) // --------------------------------------------------------------- template __global__ void rmsnorm_kernel( T* __restrict__ dst, const T* __restrict__ src, const T* __restrict__ weight, // may be nullptr if no affine weight uint32_t hidden_size, uint32_t n_vecs, // hidden_size / kVecN float eps) { using vec_t = device::AlignedVector; const uint32_t row = blockIdx.x; const T* row_src = src + row * hidden_size; T* row_dst = dst + row * hidden_size; // --- Pass 1: accumulate sum of squares (vectorized) --- float sum_sq = 0.f; for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { vec_t v; v.load(row_src, vi); #pragma unroll for (int i = 0; i < kVecN; ++i) { float val = static_cast(v[i]); sum_sq += val * val; } } // --- Warp reduction --- sum_sq = device::warp::reduce_sum(sum_sq); // --- Block reduction via shared memory --- __shared__ float smem[32]; if (threadIdx.x % 32 == 0) { smem[threadIdx.x / 32] = sum_sq; } __syncthreads(); if (threadIdx.x < 32) { sum_sq = (threadIdx.x < blockDim.x / 32) ? smem[threadIdx.x] : 0.f; sum_sq = device::warp::reduce_sum(sum_sq); } __syncthreads(); const float rms_inv = device::math::rsqrt(sum_sq / static_cast(hidden_size) + eps); // --- Pass 2: normalize + apply weight (vectorized) --- for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { vec_t v_in, v_w, v_out; v_in.load(row_src, vi); if (weight != nullptr) { v_w.load(weight, vi); } #pragma unroll for (int i = 0; i < kVecN; ++i) { float val = static_cast(v_in[i]) * rms_inv; if (weight != nullptr) { val *= static_cast(v_w[i]); } v_out[i] = static_cast(val); } v_out.store(row_dst, vi); } } // --------------------------------------------------------------- // Launcher // --------------------------------------------------------------- template void rmsnorm( tvm::ffi::TensorView dst, tvm::ffi::TensorView src, tvm::ffi::TensorView weight, // pass empty / nullptr for no-weight case float eps) { using namespace host; // Validate SymbolicSize B{"batch_tokens"}, H{"hidden_size"}; SymbolicDevice device; device.set_options(); TensorMatcher({B, H}) .with_dtype() .with_device(device) .verify(dst) .verify(src); const uint32_t num_rows = static_cast(B.unwrap()); const uint32_t hidden = static_cast(H.unwrap()); const DLDevice dev = device.unwrap(); RuntimeCheck(hidden % (16 / sizeof(T)) == 0, "rmsnorm: hidden_size must be divisible by vector width, got ", hidden); constexpr int kVecN = 16 / sizeof(T); // 128-bit vector: 8×fp16/bf16, 4×fp32 const uint32_t n_vecs = hidden / kVecN; // Thread count: enough warps to cover n_vecs, max 512 threads uint32_t threads = std::min(n_vecs, 512u); threads = (threads + 31) / 32 * 32; // round up to warp boundary const T* w_ptr = (weight.data_ptr() != nullptr) ? static_cast(weight.data_ptr()) : nullptr; LaunchKernel(num_rows, threads, dev)( rmsnorm_kernel, static_cast(dst.data_ptr()), static_cast(src.data_ptr()), w_ptr, hidden, n_vecs, eps); } } // namespace ``` --- ## Step 2: Python Wrapper Create `python/sglang/jit_kernel/diffusion/rmsnorm.py`: ```python from __future__ import annotations from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_rmsnorm_module(dtype: torch.dtype) -> Module: args = make_cpp_args(dtype) return load_jit( "diffusion_rmsnorm", *args, cuda_files=["diffusion/rmsnorm.cuh"], # relative to csrc/ cuda_wrappers=[("rmsnorm", f"rmsnorm<{args}>")], extra_cuda_cflags=["-O3", "--use_fast_math"], ) def diffusion_rmsnorm( src: torch.Tensor, weight: torch.Tensor | None = None, eps: float = 1e-6, out: torch.Tensor | None = None, ) -> torch.Tensor: """ RMSNorm for diffusion DiT layers. y = x / rms(x) * weight (weight=None → no affine scaling) Supported dtypes: float16, bfloat16, float32. hidden_size must be divisible by 8 (fp16/bf16) or 4 (fp32). """ assert src.is_cuda, "src must be a CUDA tensor" assert src.dtype in (torch.float16, torch.bfloat16, torch.float32) if out is None: out = torch.empty_like(src) # Pass a zero-sized tensor when weight is absent (launcher checks data_ptr == nullptr) w = weight if weight is not None else torch.empty(0, dtype=src.dtype, device=src.device) module = _jit_rmsnorm_module(src.dtype) module.rmsnorm(out, src, w, eps) return out ``` **Key rules for the wrapper:** - Use `cache_once` — never `functools.lru_cache` (breaks `torch.compile`) - First arg(s) to `load_jit` form the unique build cache key - `cuda_files` are relative to `python/sglang/jit_kernel/csrc/` - `cuda_wrappers`: `(python_name, cpp_template_instantiation)` --- ## Step 3: Integrate into Denoising Stage The kernel replaces a slow operator inside the DiT forward pass. Find the correct module in: ``` python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py python/sglang/multimodal_gen/runtime/models/dits/.py ``` **Pattern — monkey-patch the DiT block's RMSNorm:** ```python from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm def _patch_rmsnorm(model: torch.nn.Module) -> None: for name, module in model.named_modules(): cls_name = type(module).__name__ if cls_name in ("RMSNorm", "LlamaRMSNorm") or "RMSNorm" in cls_name: eps = getattr(module, "eps", getattr(module, "variance_epsilon", 1e-6)) has_weight = hasattr(module, "weight") and module.weight is not None if has_weight: def _make_fwd(mod, epsilon): def forward(x): return diffusion_rmsnorm(x, weight=mod.weight, eps=epsilon) return forward module.forward = _make_fwd(module, eps) else: def _make_fwd_noweight(epsilon): def forward(x): return diffusion_rmsnorm(x, weight=None, eps=epsilon) return forward module.forward = _make_fwd_noweight(eps) ``` **Critical:** inject kernels **before** `torch.compile` and before any CPU offload is enabled. --- ## Step 4: Key Kernel Patterns Reference ### Diffusion-Specific Operators | Operator | Kernel Pattern | Notes | |----------|---------------|-------| | **RMSNorm** | 2-pass row reduction + vectorized normalize | Weight may be `None` (`elementwise_affine=False`) | | **AdaLN modulation** | `y = norm(x) * (1 + scale) + shift` | Fuse norm + scale + shift in one pass | | **RoPE 3D** | Read `(t, h, w)` cos/sin tables, apply to `(q, k)` | Layout: `[batch, t*h*w, heads, head_dim]` | | **GEGLU** | Split last dim → `gate * silu(linear)` | Input `[B, L, 2*H]` → output `[B, L, H]` | | **SiLU gate** | `out = a * sigmoid(a)` fused | Avoid separate elementwise ops | ### Vectorized Memory Access ```cpp // BF16: 8 elements × 2 bytes = 16 bytes per vector load (AlignedVector) // FP16: 8 elements × 2 bytes = 16 bytes (AlignedVector) // FP32: 4 elements × 4 bytes = 16 bytes (AlignedVector) constexpr int kVecN = 16 / sizeof(T); using vec_t = device::AlignedVector; ``` ### Warp / Block Reductions ```cpp // Warp reduction (within 32 threads) float result = device::warp::reduce_sum(partial); // Block reduction via shared memory (see rmsnorm example above) __shared__ float smem[32]; // ... write warp-leaders into smem, sync, reduce again ``` ### Thread Configuration ```cpp // Element-wise (RoPE, GEGLU, SiLU): simple 1D grid constexpr uint32_t kBlock = 256; uint32_t grid = host::div_ceil(total_elements, kBlock); LaunchKernel(grid, kBlock, dev)(kernel, ...); // Row reduction (RMSNorm, LayerNorm): one block per row uint32_t threads = std::min(hidden_size / kVecN, 512u); threads = (threads + 31) / 32 * 32; LaunchKernel(num_rows, threads, dev)(kernel, ...); ``` --- ## Step 5: GPU Architecture Targets | GPU | Compute Cap | Memory BW | BF16 | Key Note | |-----|------------|-----------|------|----------| | H100 | sm_90 | 3.35 TB/s | Yes | Primary target; 132 SMs, 192 KB shared mem/SM | | A100 | sm_80 | 2.0 TB/s | Yes | 108 SMs, 164 KB shared mem/SM | | T4 | sm_75 | 320 GB/s | **No** | FP16 only; no `__nv_bfloat16` | If kernel requires SM90+ features (e.g., TMA, wgmma), raise a clear error: ```python if torch.cuda.get_device_capability()[0] < 9: raise RuntimeError("This kernel requires SM90 (H100/Hopper) or later") ``` **Grid sizing for H100** (132 SMs): aim for grid multiples of 132 for good occupancy. --- ## Step 6: Tests Create `python/sglang/jit_kernel/tests/test_diffusion_rmsnorm.py`: ```python import pytest import torch from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @pytest.mark.parametrize("shape", [(1, 2048), (4, 3072), (16, 4096)]) @pytest.mark.parametrize("has_weight", [True, False]) def test_rmsnorm_correctness(dtype, shape, has_weight): batch, hidden = shape src = torch.randn(batch, hidden, dtype=dtype, device="cuda") weight = torch.randn(hidden, dtype=dtype, device="cuda") if has_weight else None out_jit = diffusion_rmsnorm(src, weight=weight, eps=1e-6) # Reference: torch.nn.functional ref = torch.nn.functional.rms_norm( src.float(), (hidden,), weight.float() if weight is not None else None, eps=1e-6 ).to(dtype) tol = {"rtol": 1e-2, "atol": 1e-2} if dtype != torch.float32 else {"rtol": 1e-5, "atol": 1e-6} torch.testing.assert_close(out_jit, ref, **tol) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ``` --- ## Step 7: Benchmark Create `python/sglang/jit_kernel/benchmark/bench_diffusion_rmsnorm.py`: ```python import torch import triton.testing from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE, run_benchmark from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm SHAPES = [(4096, 2048), (4096, 3072), (4096, 4096)] @triton.testing.perf_report( triton.testing.Benchmark( x_names=["hidden"], x_vals=[s[1] for s in SHAPES], line_arg="provider", line_vals=["jit_cuda", "torch"], line_names=["SGLang JIT CUDA", "PyTorch rms_norm"], styles=[("blue", "-"), ("red", "--")], ylabel="us", plot_name="diffusion-rmsnorm", args={}, ) ) def benchmark(hidden: int, provider: str): src = torch.randn(4096, hidden, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) w = torch.ones(hidden, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) if provider == "jit_cuda": fn = lambda: diffusion_rmsnorm(src, weight=w, eps=1e-6) else: fn = lambda: torch.nn.functional.rms_norm(src, (hidden,), w, eps=1e-6) return run_benchmark(fn) if __name__ == "__main__": benchmark.run(print_data=True) ``` --- ## Step 8: Profile with Nsight Compute (required) After correctness + benchmarking, you must collect **Nsight Compute (ncu)** data to validate: - Whether the kernel reaches reasonable bandwidth/throughput (avoid false positives where it is “faster” but under-utilizes hardware) - Whether there are clear occupancy / register / shared memory limiters Use the canonical docs in this directory (do not duplicate CLI details across multiple skills): - `diffusion-benchmark-and-profile.md` → Step 3.5 (ncu workflow, including CUDA graph profiling) - `nsight-profiler.md` (metrics interpretation: bandwidth / occupancy / roofline / stall reasons) --- ## Common Pitfalls | Issue | Fix | |-------|-----| | `RMSNorm weight is None` | Use `type(module).__name__` check; pass `None` weight explicitly | | `isinstance(m, torch.nn.RMSNorm)` misses diffusers variants | Use `"RMSNorm" in type(m).__name__` | | Kernel patched after `torch.compile` | Inject **before** any compile call | | Kernel patched after `enable_model_cpu_offload()` | Inject **before** CPU offload | | `hidden_size` not divisible by `kVecN` | Add `RuntimeCheck(hidden % kVecN == 0, ...)` in launcher | | `torch.compile` fails with custom CUDA kernel | Register as `@torch.library.custom_op` or use Triton instead | | T4 GPU with BF16 kernel | Gate on compute capability; T4 is `sm_75`, no native BF16 | --- ## Summary of Files ``` python/sglang/jit_kernel/csrc/diffusion/ └── rmsnorm.cuh # NEW: JIT CUDA kernel source python/sglang/jit_kernel/diffusion/ └── rmsnorm.py # NEW: Python wrapper + load_jit python/sglang/jit_kernel/tests/ └── test_diffusion_rmsnorm.py # NEW: correctness tests python/sglang/jit_kernel/benchmark/ └── bench_diffusion_rmsnorm.py # NEW: benchmark ``` --- ## References ### This Skill's Extended Docs (references/ and scripts/) | File | Contents | |------|----------| | [references/kernel-templates.md](references/kernel-templates.md) | Copy-paste templates: element-wise, RMSNorm, AdaLN, Python wrapper, test, benchmark | | [references/troubleshooting.md](references/troubleshooting.md) | Build errors, perf issues, torch.compile compatibility, debugging checklist | | [references/h100-optimization-guide.md](references/h100-optimization-guide.md) | H100 (sm_90): memory hierarchy, warp reductions, occupancy, vectorization benchmarks | | [references/a100-optimization-guide.md](references/a100-optimization-guide.md) | A100 (sm_80): cp.async, TF32, 2:4 sparsity, H100→A100 migration checklist | | [references/t4-optimization-guide.md](references/t4-optimization-guide.md) | T4 (sm_75): FP16 only, low bandwidth, tile size limits, memory constraints | | [scripts/bench_diffusion_rmsnorm.py](scripts/bench_diffusion_rmsnorm.py) | Micro-benchmark: JIT CUDA RMSNorm vs PyTorch, correctness check, bandwidth analysis | | [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) | End-to-end: `sglang generate` baseline vs custom kernels, comparison table | ### SGLang Internals - **JIT system**: `add-jit-kernel` skill (`sglang/.claude/skills/add-jit-kernel/SKILL.md`) - **JIT utils**: `python/sglang/jit_kernel/utils.py` — `cache_once`, `load_jit`, `make_cpp_args` - **Abstractions**: `python/sglang/jit_kernel/include/sgl_kernel/` — `tensor.h`, `utils.cuh`, `vec.cuh`, `warp.cuh`, `math.cuh`, `tile.cuh` - **Real csrc examples**: `python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh`, `python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh` ### Other Diffusion Kernel Skills (this directory) - **Triton alternative**: `add-triton-kernel.md` — prefer Triton unless bandwidth analysis shows CUDA needed - **Existing fused kernels**: `use-efficient-diffusion-kernels.md` — check here first before writing new kernels - **Profiling**: `diffusion-benchmark-and-profile.md` — workflow to identify bottleneck before implementing - **Nsight Compute deep dive**: `nsight-profiler.md` — full guide: occupancy analysis, roofline model, warp efficiency, kernel comparison ### External - [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels) — original source adapted for this skill ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/add-triton-kernel.md ================================================ --- name: add-triton-kernel description: Step-by-step guide for adding a new Triton kernel to SGLang Diffusion's jit_kernel module. Use when implementing fused elementwise ops, norm variants, RoPE variants, or any other lightweight GPU kernel for diffusion models using Triton JIT. Covers kernel authoring, autotune, torch.compile compatibility, layer integration, and tests. --- # Adding a Triton Kernel to SGLang Diffusion This guide walks through adding a Triton kernel to `python/sglang/jit_kernel/diffusion/triton/`. We use a fused elementwise operation as the running example: `y = x * (1 + scale) + shift` (AdaLN modulation). Before compiling, benchmarking, or profiling any Triton kernel from this guide, use `scripts/diffusion_skill_env.py` (or the setup block in `diffusion-benchmark-and-profile.md`) to `cd` to the repo root resolved from `sglang.__file__`, verify write access, export `FLASHINFER_DISABLE_VERSION_CHECK=1`, and choose an idle GPU. --- ## Directory Layout ``` python/sglang/jit_kernel/diffusion/ ├── triton/ │ ├── scale_shift.py # AdaLN scale/shift fused kernels │ ├── norm.py # LayerNorm / RMSNorm fused kernels │ ├── rmsnorm_onepass.py # One-pass RMSNorm for small hidden size │ └── rotary.py # RoPE kernel └── cutedsl/ └── ... # CuTe DSL kernels (see use-efficient-diffusion-kernels.md) ``` New Triton kernels go into `triton/.py`. --- ## Step 1: Write the Triton Kernel Create `python/sglang/jit_kernel/diffusion/triton/.py`. ### 1a. Imports ```python import torch import triton # type: ignore import triton.language as tl # type: ignore ``` Always use `# type: ignore` on triton imports — the stubs are incomplete. ### 1b. The `@triton.jit` Kernel Function Follow the naming convention `__kernel` (private, underscore prefix). ```python @triton.autotune( configs=[ triton.Config({"BLOCK_C": 64}, num_warps=2), triton.Config({"BLOCK_C": 128}, num_warps=4), triton.Config({"BLOCK_C": 256}, num_warps=4), triton.Config({"BLOCK_C": 512}, num_warps=8), ], key=["C"], # re-tune when hidden dim changes ) @triton.jit def _fused_scale_shift_kernel( # Pointers — always pass raw tensors; Triton takes .data_ptr() internally x_ptr, scale_ptr, shift_ptr, y_ptr, # Dimensions B, # batch size L, # sequence length C, # hidden / channel dim # Strides — pass every stride separately; do NOT assume contiguous stride_xb, stride_xl, stride_xc, stride_sb, stride_sc, stride_yb, stride_yl, stride_yc, # Compile-time constants (tl.constexpr) BLOCK_C: tl.constexpr, ): # Grid: (cdiv(L, 1), B) — one program per (batch, token) pid_l = tl.program_id(0) pid_b = tl.program_id(1) c_offs = tl.arange(0, BLOCK_C) mask = c_offs < C x_row = pid_b * stride_xb + pid_l * stride_xl y_row = pid_b * stride_yb + pid_l * stride_yl s_row = pid_b * stride_sb x = tl.load(x_ptr + x_row + c_offs * stride_xc, mask=mask, other=0.0) scale = tl.load(scale_ptr + s_row + c_offs * stride_sc, mask=mask, other=0.0) shift = tl.load(shift_ptr + s_row + c_offs * stride_sc, mask=mask, other=0.0) y = x * (1.0 + scale) + shift tl.store(y_ptr + y_row + c_offs * stride_yc, y, mask=mask) ``` **Rules:** - All pointer arguments are raw (Triton extracts `.data_ptr()` internally when called via `kernel[grid](...)`). - Pass every stride as a separate scalar — never assume a tensor is contiguous inside the kernel. - Use `tl.constexpr` for block sizes and boolean flags (`HAS_RESIDUAL`, `IS_RMS_NORM`, etc.). - Use `mask=mask, other=0.0` on every `tl.load` to avoid out-of-bounds reads. - Compute in `tl.float32` when precision matters (`x.to(tl.float32)`), then cast back to output dtype before `tl.store`. - Use `tl.fma(a, b, c)` (`a*b + c`) for fused multiply-add — avoids rounding errors and maps to a single instruction. ### 1c. `@triton.autotune` Guidelines | `key` entry | When to include | |-------------|-----------------| | `"C"` / `"hidden_dim"` | Always — block tile size depends on C | | `"IS_RMS_NORM"` | When the kernel has a `constexpr` boolean flag that changes code paths | | `"HAS_RESIDUAL"` | Same — constexpr path branching | | Shape / batch / seq | Usually NOT — autotune cost outweighs benefit | Keep configs in ascending `BLOCK_C` order with matching `num_warps` (warp × 32 threads ≤ 1024). ### 1d. `torch.compile` Compatibility When the kernel is called inside a `torch.compile`-d region, wrap the launch with `torch.library.wrap_triton`: ```python with torch.get_device_module().device(x.device): torch.library.wrap_triton(_fused_scale_shift_kernel)[grid]( x, scale, shift, y, B, L, C, x.stride(0), x.stride(1), x.stride(2), scale.stride(0), scale.stride(1), y.stride(0), y.stride(1), y.stride(2), ) ``` Use `wrap_triton` when the kernel is called from a layer that runs under `torch.compile`. Skip it for utility kernels called only at Python graph boundaries. --- ## Step 2: Write the Python Launcher The launcher is a regular Python function (public, no underscore) in the same file. ```python def fused_scale_shift( x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, ) -> torch.Tensor: """ Fused AdaLN modulation: y = x * (1 + scale) + shift. Args: x: [B, L, C], CUDA, contiguous scale: [B, C], CUDA shift: [B, C], CUDA (same shape as scale) Returns: y: same shape and dtype as x """ # --- Precondition checks --- assert x.is_cuda, "x must be on CUDA" assert x.is_contiguous(), "x must be contiguous" assert scale.is_cuda and shift.is_cuda assert x.ndim == 3, f"x must be 3D [B, L, C], got {x.shape}" assert scale.shape == shift.shape B, L, C = x.shape # Allocate output y = torch.empty_like(x) # Grid: one program per token grid = (L, B) _fused_scale_shift_kernel[grid]( x, scale, shift, y, B, L, C, x.stride(0), x.stride(1), x.stride(2), scale.stride(0), scale.stride(1), y.stride(0), y.stride(1), y.stride(2), ) return y ``` **Rules:** - Validate CUDA placement and shape/dtype **before** launching — use `assert` with a helpful message. - Call `.contiguous()` on inputs that the kernel requires contiguous **before** the launch, not inside it. - Allocate the output with `torch.empty_like(x)` — never reuse input buffers unless the op is explicitly in-place. - The `grid` is a tuple or a lambda `(META)` when block sizes are auto-tuned: ```python # Static grid (block size fixed) grid = (triton.cdiv(L, BLOCK_L), triton.cdiv(C, BLOCK_C), B) # Dynamic grid (block size comes from autotune) grid = lambda META: (triton.cdiv(L, META["BLOCK_C"]), B) ``` ### Handling Non-Contiguous Inputs Never call `.contiguous()` silently — it copies data. Instead, pass strides to the kernel and let it handle arbitrary layouts. Only call `.contiguous()` when the kernel genuinely requires it (e.g., after a reshape): ```python # OK: reshape + contiguous needed for 2D view trick x_2d = x.view(B * L, C) # view only works on contiguous if not x.is_contiguous(): x = x.contiguous() x_2d = x.view(B * L, C) ``` --- ## Step 3: Integrate into the Layer Call the new kernel from the appropriate layer file in `python/sglang/multimodal_gen/runtime/layers/` (typically `layernorm.py` or `elementwise.py`). ```python # In layernorm.py or elementwise.py import torch def apply_scale_shift(x, scale, shift): if x.is_cuda: from sglang.jit_kernel.diffusion.triton.my_op import fused_scale_shift return fused_scale_shift(x, scale, shift) # Pure-PyTorch fallback for non-CUDA execution return x * (1.0 + scale) + shift ``` **Rules:** - Gate on `x.is_cuda` — the Triton kernel only runs on CUDA; the fallback handles everything else. - The launcher raises `AssertionError` on invalid inputs (wrong shape, CPU tensor, etc.) — do **not** silently catch these. Let them propagate so bugs are visible during development. - Add `logger.warning_once(...)` only when falling back due to a **known hardware limitation** (e.g., unsupported SM compute capability), not for wrong-input errors. --- ## Step 4: Write Tests Create `python/sglang/jit_kernel/tests/test_.py`. ```python import pytest import torch from sglang.jit_kernel.diffusion.triton.my_op import fused_scale_shift def _ref_fused_scale_shift(x, scale, shift): """PyTorch reference implementation.""" # Broadcast scale/shift from [B, C] to [B, L, C] return x * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1) @pytest.fixture(autouse=True) def require_cuda(): if not torch.cuda.is_available(): pytest.skip("CUDA required") @pytest.mark.parametrize("B,L,C", [ (1, 6, 3072), # Qwen (small batch) (1, 1024, 1536), # Wan (2, 512, 3072), # typical training shape (1, 1, 256), # edge: L=1 ]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) def test_fused_scale_shift_correctness(B, L, C, dtype): torch.manual_seed(0) x = torch.randn(B, L, C, dtype=dtype, device="cuda") scale = torch.randn(B, C, dtype=dtype, device="cuda") * 0.1 shift = torch.randn(B, C, dtype=dtype, device="cuda") * 0.1 out = fused_scale_shift(x, scale, shift) ref = _ref_fused_scale_shift(x.float(), scale.float(), shift.float()).to(dtype) atol = 1e-5 if dtype == torch.float32 else 1e-2 torch.testing.assert_close(out, ref, atol=atol, rtol=atol, msg=f"Mismatch at B={B} L={L} C={C} dtype={dtype}") def test_fused_scale_shift_non_cuda_raises(): x = torch.randn(1, 4, 64) scale = torch.randn(1, 64) shift = torch.randn(1, 64) with pytest.raises(AssertionError, match="CUDA"): fused_scale_shift(x, scale, shift) def test_fused_scale_shift_output_dtype_preserved(): x = torch.randn(1, 8, 128, dtype=torch.bfloat16, device="cuda") scale = torch.randn(1, 128, dtype=torch.bfloat16, device="cuda") shift = torch.zeros(1, 128, dtype=torch.bfloat16, device="cuda") out = fused_scale_shift(x, scale, shift) assert out.dtype == torch.bfloat16 assert out.shape == x.shape if __name__ == "__main__": pytest.main([__file__, "-v"]) ``` Run: ```bash pytest python/sglang/jit_kernel/tests/test_.py -v ``` **Test coverage requirements:** 1. Reference comparison against pure-PyTorch for all supported dtypes (fp16, bf16, fp32). 2. Edge shapes: `L=1`, `C` not a multiple of the largest BLOCK_C, large `B`. 3. Error cases: CPU tensor, wrong shape. 4. Output dtype and shape preservation. --- ## Step 5: Add a Benchmark (required) Create `python/sglang/jit_kernel/benchmark/bench_.py`. ```python import torch import triton.testing from sglang.jit_kernel.diffusion.triton.my_op import fused_scale_shift SHAPES = [ # (B, L, C) — representative diffusion shapes (1, 6, 3072), # Qwen image (1, 1024, 1536), # Wan video (1, 4096, 3072), # FLUX double-stream ] @triton.testing.perf_report( triton.testing.Benchmark( x_names=["B", "L", "C"], x_vals=SHAPES, line_arg="provider", line_vals=["triton", "torch"], line_names=["Triton Fused", "PyTorch"], styles=[("blue", "-"), ("red", "--")], ylabel="µs (median)", plot_name="fused-scale-shift", args={}, ) ) def benchmark(B, L, C, provider): dtype = torch.bfloat16 x = torch.randn(B, L, C, dtype=dtype, device="cuda") scale = torch.randn(B, C, dtype=dtype, device="cuda") shift = torch.randn(B, C, dtype=dtype, device="cuda") if provider == "triton": fn = lambda: fused_scale_shift(x, scale, shift) else: fn = lambda: x * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1) ms, *_ = triton.testing.do_bench_cudagraph(fn, quantiles=[0.5, 0.2, 0.8]) return ms * 1000 # µs if __name__ == "__main__": benchmark.run(print_data=True) ``` Run: ```bash python python/sglang/jit_kernel/benchmark/bench_.py ``` --- ## Step 6: Profile with Nsight Compute (required for optimization work) After correctness tests, you must use **ncu (Nsight Compute)** to validate hardware efficiency (bandwidth/throughput/occupancy/bottleneck type). To avoid duplicating ncu CLI details across multiple skills, this skill does not repeat command flags. Follow the canonical docs: - `diffusion-benchmark-and-profile.md` → Step 3.5 (ncu workflow, including CUDA graph profiling) - `nsight-profiler.md` (metrics interpretation: bandwidth / occupancy / roofline / warp stalls) --- ## Common Patterns Reference ### Pattern 1: Autotune over a 2D tile (L × C) Used in `scale_shift.py` (`fuse_scale_shift_kernel_blc_opt`): ```python @triton.jit def _kernel(..., BLOCK_L: tl.constexpr, BLOCK_C: tl.constexpr): pid_l = tl.program_id(0) pid_c = tl.program_id(1) pid_b = tl.program_id(2) l_offs = pid_l * BLOCK_L + tl.arange(0, BLOCK_L) c_offs = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) mask = (l_offs[:, None] < L) & (c_offs[None, :] < C) ... # Launch: grid = (triton.cdiv(L, BLOCK_L), triton.cdiv(C, BLOCK_C), B) _kernel[grid](..., BLOCK_L=block_l, BLOCK_C=block_c, num_warps=4, num_stages=2) ``` ### Pattern 2: One-pass RMSNorm for small hidden size Used in `rmsnorm_onepass.py`: ```python @triton.jit def _rms_norm_tiled_onepass(y_ptr, x_ptr, w_ptr, SEQ: tl.constexpr, DIM: tl.constexpr, EPS: tl.constexpr, BLOCK_SIZE_SEQ: tl.constexpr, BLOCK_SIZE_DIM: tl.constexpr): seq_blk_id = tl.program_id(0) seq_id = seq_blk_id * BLOCK_SIZE_SEQ seq_offset = seq_id + tl.arange(0, BLOCK_SIZE_SEQ)[:, None] d_offset = tl.arange(0, BLOCK_SIZE_DIM)[None, :] ... x = tl.load(x_ptr + seq_offset * DIM + d_offset, mask=..., other=0.0).to(tl.float32) mean_sq = tl.sum(x * x, axis=1, keep_dims=True) / DIM rstd = tl.math.rsqrt(mean_sq + EPS) tl.store(y_ptr + ..., x * rstd * w, mask=...) # Launch with wrap_triton for torch.compile compat: with torch.get_device_module().device(x.device): torch.library.wrap_triton(_rms_norm_tiled_onepass)[grid]( y_view, x_view, w, S, D, eps, BLOCK_SIZE_DIM=triton.next_power_of_2(D), BLOCK_SIZE_SEQ=BLOCK_SIZE_SEQ, ) ``` ### Pattern 3: `tl.constexpr` boolean flags for conditional paths Used in `norm.py` and `scale_shift.py`: ```python @triton.jit def _kernel(..., IS_RMS_NORM: tl.constexpr, HAS_RESIDUAL: tl.constexpr, SCALE_IS_SCALAR: tl.constexpr): ... if IS_RMS_NORM: var = tl.sum(x * x, axis=0) / N else: mean = tl.sum(x, axis=0) / N var = tl.sum((x - mean) ** 2, axis=0) / N if HAS_RESIDUAL: x = x + tl.load(residual_ptr + ...) if SCALE_IS_SCALAR: scale_val = tl.load(scale_ptr) scale = tl.full([BLOCK_N], scale_val, dtype=scale_val.dtype) else: scale = tl.load(scale_ptr + col_offsets, mask=mask, other=0.0) ``` Autotune key must include these booleans so the compiler generates separate specializations. ### Pattern 4: Computing in fp32, storing in original dtype Always up-cast to `tl.float32` for reductions and math, then down-cast before storing: ```python x_f32 = x.to(tl.float32) scale_f32 = scale.to(tl.float32) y_f32 = x_f32 * (1.0 + scale_f32) + shift_f32 tl.store(y_ptr + offsets, y_f32.to(x.dtype), mask=mask) ``` --- ## Checklist Before Submitting ### Prerequisites - [ ] `ncu --version` prints a valid Nsight Compute version (required for Step 7 profiling) ### Implementation - [ ] Kernel file at `python/sglang/jit_kernel/diffusion/triton/.py` - [ ] All pointer arguments passed with separate stride scalars - [ ] Every `tl.load` uses `mask=` and `other=` - [ ] Autotune `key` includes all `constexpr` flags that change code paths - [ ] `torch.library.wrap_triton` used if kernel runs inside `torch.compile` region - [ ] PyTorch fallback path in the layer integration (see Step 4) ### Validation - [ ] Tests pass: `pytest python/sglang/jit_kernel/tests/test_.py -v` - [ ] Benchmark runs: `python python/sglang/jit_kernel/benchmark/bench_.py` - [ ] **Correctness verified**: Triton output matches PyTorch reference within tolerance - [ ] Nsight Compute profile collected (`ncu --set full`); achieved occupancy ≥ 50% and memory throughput ≥ 70% of peak (or bottleneck documented) --- ## Summary of Files Created/Modified ``` python/sglang/jit_kernel/diffusion/triton/.py # NEW: Triton kernel + launcher python/sglang/jit_kernel/tests/test_.py # NEW: correctness tests python/sglang/jit_kernel/benchmark/bench_.py # NEW: performance benchmark python/sglang/multimodal_gen/runtime/layers/layernorm.py # MODIFIED: integrate into layer (or elementwise.py, depending on op type) ``` ## References - `python/sglang/jit_kernel/diffusion/triton/scale_shift.py` — 2D tile pattern, scalar broadcast, 4D shape handling - `python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py` — `wrap_triton`, tiled one-pass reduction - `python/sglang/jit_kernel/diffusion/triton/norm.py` — complex autotune with many `constexpr` flags - `python/sglang/jit_kernel/diffusion/triton/rotary.py` — per-head grid, interleaved RoPE - `nsight-profiler.md` — full Nsight Compute guide: occupancy analysis, roofline model, warp efficiency, kernel comparison - `diffusion-benchmark-and-profile.md` — how to verify the kernel's impact on denoise latency - `use-efficient-diffusion-kernels.md` — overview of existing fused kernel entry points ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/diffusion-benchmark-and-profile.md ================================================ --- name: diffusion-benchmark-and-profile description: Denoise-stage benchmark and per-layer kernel profiling guide for SGLang Diffusion models. Use when measuring denoising latency, profiling DiT kernel breakdown with torch.profiler or nsys+gputrc2graph.py, investigating performance bottlenecks, or optimizing with custom Triton/CUDA kernels. Always verify output correctness before and after any optimization. --- # SGLang Diffusion Benchmark and Profile Guide **Primary Metric: Denoise Latency** The denoising loop latency — total DiT forward pass time across all inference steps — is the dominant cost (>80% of end-to-end) and the **sole optimization target** for kernel work. End-to-end latency is recorded as a secondary check only. > **Correctness First**: Faster but incorrect output is not an improvement. Always compare generated images/videos against a reference baseline before and after any change. --- ## Prerequisites ```bash ENV_PY=python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/diffusion_skill_env.py ROOT=$(python3 "$ENV_PY" print-root) cd "$ROOT" python3 "$ENV_PY" check-write-access >/dev/null export FLASHINFER_DISABLE_VERSION_CHECK=1 export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 1) ASSET_DIR=$(python3 "$ENV_PY" print-assets-dir --mkdir) BENCH_DIR=$(python3 "$ENV_PY" print-output-dir --kind benchmarks --mkdir) PROFILE_DIR=$(python3 "$ENV_PY" print-output-dir --kind profiles --mkdir) NCU_DIR=$(python3 "$ENV_PY" print-output-dir --kind ncu --mkdir) export PROFILE_DIR check() { local label="$1" shift "$@" &>/dev/null && echo "[OK] $label" || echo "[MISS] $label" } check "sglang" python3 -c "import sglang" check "torch+CUDA" python3 -c "import torch; assert torch.cuda.is_available()" check "torch.profiler" python3 -c "import torch.profiler" check "nsys (Level 2)" which nsys check "ncu (Level 3)" which ncu check "pandas" python3 -c "import pandas" check "plotly" python3 -c "import plotly" ``` **Minimum for benchmarking**: `sglang`, `torch` with CUDA. **Level 1 profiling**: `torch.profiler` (bundled with torch). **Level 2 profiling**: `nsys`, `pandas`, `plotly` + `gputrc2graph.py` from the sglang repo. All commands below assume you are inside the configured diffusion container shell, already `cd`'d to the repo root derived from `sglang.__file__`, with `FLASHINFER_DISABLE_VERSION_CHECK=1` exported. Re-run `print-idle-gpus` before each perf command if GPU availability may have changed. Keep benchmark commands within 4 GPUs or fewer. Download input images required by some models: ```bash wget -O "${ASSET_DIR}/cat.png" \ https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png wget -O "${ASSET_DIR}/astronaut.jpg" \ https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg wget -O "${ASSET_DIR}/mova_single_person.jpg" \ https://github.com/OpenMOSS/MOVA/raw/main/assets/single_person.jpg ``` --- ## Benchmark Commands All commands include `--warmup` and `--enable-torch-compile` for real production performance. Add `--perf-dump-path .json` for machine-readable output. ### Perf dump & before/after compare For every benchmark run, always write a perf dump JSON: ```bash sglang generate ... --warmup --perf-dump-path "${BENCH_DIR}/.json" ``` Before/after comparison (outputs a Markdown table suitable for PR descriptions): ```bash # Baseline (on main branch or before changes) sglang generate ... --warmup --perf-dump-path "${BENCH_DIR}/baseline.json" # New (after changes) sglang generate ... --warmup --perf-dump-path "${BENCH_DIR}/new.json" python3 python/sglang/multimodal_gen/benchmarks/compare_perf.py \ "${BENCH_DIR}/baseline.json" "${BENCH_DIR}/new.json" ``` ### Qwen-Image-2512 (1024×1024, 50 steps) ```bash sglang generate \ --model-path=Qwen/Qwen-Image-2512 \ --prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k" \ '--negative-prompt= ' \ --width=1024 --height=1024 --num-inference-steps=50 --guidance-scale=4.0 \ --seed=42 --save-output --enable-torch-compile --warmup \ --dit-cpu-offload false --text-encoder-cpu-offload false ``` ### Qwen-Image-Edit-2511 (image editing, 1024×1024, 50 steps) ```bash sglang generate \ --model-path=Qwen/Qwen-Image-Edit-2511 \ '--prompt=Transform into anime style' '--negative-prompt= ' \ --image-path="${ASSET_DIR}/cat.png" \ --width=1024 --height=1024 --num-inference-steps=50 --guidance-scale=4.0 \ --seed=42 --save-output --enable-torch-compile --warmup \ --dit-cpu-offload false --text-encoder-cpu-offload false ``` ### FLUX.1-dev (1024×1024, 50 steps) ```bash sglang generate \ --model-path=black-forest-labs/FLUX.1-dev \ --prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k" \ --width=1024 --height=1024 --num-inference-steps=50 --guidance-scale=4.0 \ --seed=42 --save-output --enable-torch-compile --warmup ``` ### FLUX.2-dev (1024×1024) ```bash sglang generate \ --model-path black-forest-labs/FLUX.2-dev \ --prompt "A Logo With Bold Large Text: SGL Diffusion" \ --width=1024 --height=1024 \ --dit-layerwise-offload false --enable-torch-compile --warmup \ --dit-cpu-offload false --text-encoder-cpu-offload true --vae-cpu-offload false ``` ### Z-Image-Turbo (1024×1024, 9 steps) ```bash sglang generate \ --model-path=Tongyi-MAI/Z-Image-Turbo \ --prompt='A fantasy landscape with mountains and a river, detailed, vibrant colors' \ --width=1024 --height=1024 --num-inference-steps=9 --guidance-scale=0.0 \ --seed=42 --save-output --enable-torch-compile --warmup \ --dit-cpu-offload false --text-encoder-cpu-offload false ``` ### Wan2.2-T2V-A14B 720P (4 GPUs, 81 frames, 2 steps) ```bash # Select four idle GPUs first: # export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 4) sglang generate \ --model-path=Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --prompt="A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon." \ --negative-prompt=" " --720p --num-inference-steps=2 --num-frames=81 \ --guidance-scale=5.0 --seed=42 --save-output \ --num-gpus=4 --ulysses-degree=4 \ --text-encoder-cpu-offload --pin-cpu-memory \ --warmup --enable-torch-compile ``` ### Wan2.2-TI2V-5B 720P (single GPU, 81 frames, 50 steps) ```bash sglang generate \ --model-path Wan-AI/Wan2.2-TI2V-5B-Diffusers \ --prompt "An astronaut hatching from an egg, on the surface of the moon..." \ --negative-prompt "Bright tones, overexposed, static, blurred details..." \ --image-path="${ASSET_DIR}/astronaut.jpg" \ --num-frames 81 --720p --num-inference-steps 50 --guidance-scale 5.0 \ --seed 42 --save-output \ --dit-layerwise-offload false --dit-cpu-offload false \ --vae-cpu-offload false --text-encoder-cpu-offload false \ --enable-torch-compile --warmup ``` ### HunyuanVideo (848×480, 65 frames, 30 steps) ```bash sglang generate \ --model-path=hunyuanvideo-community/HunyuanVideo \ --text-encoder-cpu-offload --pin-cpu-memory \ --prompt="A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window." \ --save-output --num-frames=65 --width=848 --height=480 \ --num-inference-steps=30 \ --warmup --enable-torch-compile ``` ### MOVA-720p (4 GPUs, 193 frames, 2 steps) ```bash # Select four idle GPUs first: # export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 4) sglang generate \ --model-path=OpenMOSS-Team/MOVA-720p \ --prompt="A man in a blue blazer and glasses speaks in a formal indoor setting, framed by wooden furniture and a filled bookshelf. Quiet room acoustics underscore his measured tone as he delivers his remarks. At one point, he says, \"I would also believe that this advance in AI recently wasn’t unexpected.\"" \ --image-path="${ASSET_DIR}/mova_single_person.jpg" \ --adjust-frames=false \ --num-gpus=4 --ring-degree=1 --ulysses-degree=4 \ --num-frames=193 --fps=24 \ --num-inference-steps=2 \ --enable-torch-compile --save-output --warmup ``` **Key metrics** (all models): denoise latency ★, end-to-end latency, peak GPU memory. --- ## Performance Bottleneck Workflow ### Step 1: Identify the Slow DiT Operation Add `--log-level=info` and observe: - **Denoise loop latency** ★ — primary target - Per-step DiT latency — denoise ÷ steps ### Step 2: Profile with torch.profiler (Level 1) **Compile-safety rule for fused or rewritten kernels** - Any new kernel must be checked for `torch.compile` graph breaks before trusting its benchmark result. - If a direct Python/library call triggers tracing issues, wrap it as a custom op first. - For external libraries, use `register_custom_op_from_extern(...)`. - For SGLang JIT kernels, use `@register_custom_op(...)` and keep the JIT/module loading inside the custom op body. - Re-run `torch._dynamo.explain` on representative shapes and verify the optimized path still gets `graph_count=1` and `graph_break_count=0`. ```bash SGLANG_TORCH_PROFILER_DIR="${PROFILE_DIR}/torch" \ sglang generate \ --model-path=black-forest-labs/FLUX.1-dev \ --prompt="A futuristic cyberpunk city at night" \ --width=1024 --height=1024 --num-inference-steps=50 \ --seed=42 --enable-torch-compile --warmup \ --profile --num-profiled-timesteps 3 ``` Parse the trace without a browser: ```python import gzip, json, collections, glob, os log_dir = os.environ.get("SGLANG_TORCH_PROFILER_DIR", "./logs") trace_path = sorted(glob.glob(f"{log_dir}/*.trace.json.gz"), key=os.path.getmtime, reverse=True)[0] with gzip.open(trace_path, "rb") as f: data = json.loads(f.read()) cuda_ops = collections.defaultdict(lambda: {"total_us": 0, "count": 0}) for e in data.get("traceEvents", []): if e.get("cat") in ("kernel", "gpu_memcpy") and "dur" in e: cuda_ops[e.get("name","unknown")]["total_us"] += e["dur"] cuda_ops[e.get("name","unknown")]["count"] += 1 print(f"{'Kernel':<80} {'Total(ms)':>10} {'Count':>6}") for name, s in sorted(cuda_ops.items(), key=lambda x: -x[1]["total_us"])[:30]: print(f"{name:<80} {s['total_us']/1000:>10.3f} {s['count']:>6}") ``` Add `record_function` scopes in the DiT block for per-layer attribution: ```python with torch.profiler.record_function(f"dit_block_{idx}.attn"): x = self.attn(x) with torch.profiler.record_function(f"dit_block_{idx}.norm"): x = self.norm(x) ``` **Expected dominant kernels per DiT sub-component:** | Sub-component | Expected kernel | |--------------|-----------------| | QKV / output / MLP projections | `cutlass_gemm` / `ampere_*_gemm` | | Attention | `flash_attn_fwd` / `fmha_*` (FA3/FA4) | | AdaLN modulation | `fuse_scale_shift_kernel` | | RMSNorm / LayerNorm | `sgl_kernel_rmsnorm` / Triton norm | | SiLU gate | `vectorized_elementwise_kernel` | | RoPE | `apply_rotary_embedding` (Triton) | | QK Norm | `fused_inplace_qknorm` (JIT) | ### Step 3: Deep CUDA Kernel Breakdown (Level 2 — nsys) ```bash # Pass A — collect nsys trace (skip warmup with --delay) nsys profile -t cuda -o "${PROFILE_DIR}/flux_dev" -f true \ --trace-fork-before-exec=true --delay 120 --duration 60 \ sglang generate \ --model-path=black-forest-labs/FLUX.1-dev \ --prompt="A futuristic cyberpunk city at night" \ --width=1024 --height=1024 --num-inference-steps=50 \ --seed=42 --enable-torch-compile --warmup # Pass B — measure wall-clock time without profiling time sglang generate --model-path=black-forest-labs/FLUX.1-dev \ --width=1024 --height=1024 --num-inference-steps=50 --seed=42 \ --enable-torch-compile --warmup # Record ELAPSED_SEC from Pass B ``` Create classification JSON at `examples/profiler/nsys_profile_tools/sglang_diffusion_engine_model.json`: ```json { "sglang": { "diffusion": { "gemm|nvjet|cutlass": "gemm", "flash|fmha|fwd_flash": "attn", "fuse_scale_shift|scale_shift_gate": "adaln_modulation", "_norm_|Norm|rmsnorm|fused_add_rmsnorm": "norm", "rotary|rope": "rope", "act_and_mul|silu|gelu": "activation", "ncclDevKernel|all_gather|all_reduce": "nccl_comm", "triton": "triton_kernel", "CUDA mem": "non-gpu-H_D_memops", ".*": "misc" } } } ``` Run analysis: ```bash cd "$ROOT/examples/profiler/nsys_profile_tools" python3 gputrc2graph.py \ --in_file "${PROFILE_DIR}/flux_dev.nsys-rep,sglang,diffusion,ELAPSED_SEC" \ --out_dir "${PROFILE_DIR}/analysis" \ --title "FLUX.1-dev denoise kernel breakdown" # Read results python3 - << 'EOF' import pandas as pd df = pd.read_csv(f"{os.environ['PROFILE_DIR']}/analysis/result.csv") summary = df.groupby("Category")["Elapsed Time (sec)"].sum().sort_values(ascending=False) total = summary.sum() for cat, sec in summary.items(): print(f"{cat:<30} {sec:>8.3f}s ({sec/total*100:>5.1f}%)") EOF ``` **What the category breakdown tells you:** | Category high | Investigation | |--------------|---------------| | `gemm` dominant | Check tensor parallelism; QKV/MLP bottleneck | | `attn` dominant | Verify FA3/FA4 is active | | `adaln_modulation` high | Verify fused `fuse_scale_shift_kernel` is used | | `norm` high | Verify `sgl_kernel_rmsnorm` / CuTe DSL path; check D alignment | | `nccl_comm` high | Multi-GPU: tune Ulysses degree | | `triton_kernel` high | Identify which Triton kernel; consider CUDA replacement | | `non-gpu-H_D_memops` high | Accidental CPU offload or `.cpu()` calls mid-denoising | | `CPU(non-GPU)` high | Python dispatch overhead / torch.compile graph breaks | ### Step 3.5: Per-Kernel Deep Analysis (Level 3 — ncu) **CRITICAL**: `ncu` (Nsight Compute) is the essential tool for kernel-level optimization. While nsys and torch.profiler tell you **which** kernels are slow, only ncu tells you **why** — memory bandwidth utilization, compute throughput, occupancy limiters, warp stall reasons, and roofline position. **Always use ncu when optimizing or writing custom kernels.** #### When to use ncu - After writing a new Triton or CUDA kernel — verify it saturates hardware bandwidth - When a kernel shows up as a top bottleneck in Level 1/2 profiling - When comparing your fused kernel vs PyTorch baseline or torch.compile output - When tuning Triton autotune configs (block sizes, num_warps) #### Basic ncu workflow ```bash # 1. Profile a specific kernel by name (skip warmup launches, collect 3 invocations) ncu --kernel-name "_fused_gated_residual_add_kernel" \ --launch-skip 10 --launch-count 3 \ --set full \ -o "${NCU_DIR}/gated_residual" \ sglang generate \ --model-path=black-forest-labs/FLUX.1-dev \ --prompt="test" --width=1024 --height=1024 \ --num-inference-steps=5 --seed=42 # 2. Profile all kernels in a short run (use few steps to limit time) ncu --launch-skip 50 --launch-count 200 \ --set full \ -o "${NCU_DIR}/all_kernels" \ sglang generate \ --model-path=black-forest-labs/FLUX.1-dev \ --prompt="test" --width=1024 --height=1024 \ --num-inference-steps=3 --seed=42 # 3. For CUDA graph mode, keep --graph-profiling=node on the ncu side. # Note: `--enable-piecewise-cuda-graph` is a server flag, not a valid # `sglang generate` flag, so do not append it here. ncu --graph-profiling node \ --kernel-name "_fused_gated_residual_add_kernel" \ --launch-skip 5 --launch-count 3 \ --set full \ -o "${NCU_DIR}/gated_residual_cudagraph" \ sglang generate \ --model-path=black-forest-labs/FLUX.1-dev \ --prompt="test" --width=1024 --height=1024 \ --num-inference-steps=5 --seed=42 ``` #### Reading ncu results (CLI, no GUI needed) ```bash # Summary of all profiled kernels ncu --import "${NCU_DIR}/gated_residual.ncu-rep" --page raw --csv 2>/dev/null | head -50 # Key metrics to extract: ncu --import "${NCU_DIR}/gated_residual.ncu-rep" \ --page details --csv 2>/dev/null | python3 -c " import csv, sys reader = csv.DictReader(sys.stdin) key_metrics = { 'gpu__time_duration.avg': 'Duration', 'sm__throughput.avg.pct_of_peak_sustained_elapsed': 'Compute (SM) Throughput', 'dram__throughput.avg.pct_of_peak_sustained_elapsed': 'DRAM Throughput', 'l1tex__throughput.avg.pct_of_peak_sustained_elapsed': 'L1/TEX Cache Throughput', 'sm__warps_active.avg.pct_of_peak_sustained_active': 'Achieved Occupancy', 'launch__occupancy_limit_registers': 'Block Limit Registers', 'launch__occupancy_limit_shared_mem': 'Block Limit Shared Mem', } for row in reader: name = row.get('Metric Name', '') if any(alias in name or metric in name for metric, alias in key_metrics.items()): print(f'{name:<60} {row.get(\"Metric Value\",\"\")}') " ``` #### Interpreting ncu results for kernel optimization | Metric | Good | Action if bad | |--------|------|--------------| | DRAM throughput > 80% peak | Memory-bound, near optimal | Already saturating HBM — fuse with adjacent ops to reduce total memory traffic | | DRAM throughput < 50% peak | Not saturating memory bandwidth | Check coalescing, increase vector width, tune BLOCK sizes | | SM throughput > 60% peak | Compute-bound, near optimal | Reduce arithmetic, use faster instructions (e.g., FMA) | | SM throughput < 30% peak | Underutilized compute | Increase occupancy, reduce warp stalls, check instruction mix | | Achieved occupancy > 50% | Acceptable for most kernels | — | | Achieved occupancy < 25% | Too few active warps | Reduce register pressure or shared memory; increase block size | #### Comparing before/after with ncu ```bash # Profile baseline kernel ncu --kernel-name "vectorized_elementwise_kernel" \ --launch-skip 10 --launch-count 3 --set full \ -o "${NCU_DIR}/baseline" ./program # Profile optimized kernel ncu --kernel-name "_fused_gated_residual_add_kernel" \ --launch-skip 10 --launch-count 3 --set full \ -o "${NCU_DIR}/optimized" ./program # Compare key metrics for report in baseline optimized; do echo "=== $report ===" ncu --import "${NCU_DIR}/${report}.ncu-rep" \ --page details --csv 2>/dev/null | grep -E "time_duration|throughput.*pct|occupancy" done ``` **Decision rule after ncu analysis:** - Kernel already at >80% DRAM bandwidth → fuse with neighbors to reduce total traffic - Kernel at <50% DRAM bandwidth → tune block sizes, fix coalescing, increase vectorization - Kernel compute-bound (SM util high, DRAM low) → reduce FLOPs or switch to a faster algorithm - Low occupancy → reduce registers (simplify kernel) or increase block size in autotune configs ### Step 4: Apply Kernel Optimization After pinpointing the slow op, choose the right tool: | Scenario | Skill to use | |----------|-------------| | New fused elementwise, norm variant, RoPE variant | **`add-triton-kernel.md`** — Triton JIT, faster iteration, NPU fallback | | Bandwidth-bound reduction (RMSNorm) needing max vectorization | **`add-cuda-kernel.md`** — CUDA JIT with `AlignedVector`, warp reductions | | Attention or tile-based op needing shared memory tuning | **`add-cuda-kernel.md`** — full control over CUDA primitives | | Slow op already covered by existing fused kernel | **`use-efficient-diffusion-kernels.md`** — check constraints & enable | **Quick decision rule**: start with Triton. Switch to CUDA JIT only when profiling shows Triton can't saturate hardware bandwidth. Both kernel types use SGLang's JIT compilation: - **Triton**: `python/sglang/jit_kernel/diffusion/triton/.py` - **CUDA JIT**: `python/sglang/jit_kernel/csrc/diffusion/.cuh` + wrapper `python/sglang/jit_kernel/diffusion/.py` ### Step 5: torch.compile Coverage ```bash TORCH_COMPILE_DEBUG=1 sglang generate ... ``` - Dynamic shape changes trigger recompilation → fix resolution and frame count when benchmarking - `tensor.item()` in conditional branches causes graph breaks → rewrite as tensor ops ### Step 6: Multi-GPU Efficiency (Wan2.2-T2V-A14B / MOVA) - Verify `--ulysses-degree` evenly divides `--num-gpus` - Keep the command shape fixed when comparing kernels; for quick checks, reduce only `--num-inference-steps` - If a run OOMs or jitters because of host contention, first confirm there are no leaked scheduler processes on the chosen GPU set --- ## Optimization Workflow Summary ``` 0. BASELINE sglang generate --seed=42 --save-output → save reference images/videos ↓ 1. BENCHMARK Run benchmark commands above → record denoise latency baseline ↓ 2. LEVEL 1 PROFILE (torch.profiler) --profile --num-profiled-timesteps 3 → parse .trace.json.gz → rank ops by CUDA time → identify slow DiT layer (norm / attn / mlp / rope / adaln) ↓ 3. LEVEL 2 PROFILE (nsys + gputrc2graph.py) → result.csv category breakdown (gemm / attn / adaln / norm / triton / cpu) → confirm where GPU time is concentrated ↓ 4. LEVEL 3 PROFILE (ncu — per-kernel deep analysis) ★ CRITICAL → ncu --set full on target kernel(s) → extract DRAM bandwidth util, SM throughput, achieved occupancy → determine if kernel is memory-bound, compute-bound, or latency-bound → for CUDA graph: use --graph-profiling node ↓ 5. KERNEL OPTIMIZATION Existing fused kernel? → use-efficient-diffusion-kernels.md New Triton kernel? → add-triton-kernel.md New CUDA JIT kernel? → add-cuda-kernel.md After writing kernel → ncu again to verify bandwidth/occupancy ★ ↓ 6. VERIFY CORRECTNESS sglang generate --seed=42 --save-output → diff against reference If output differs beyond tolerance → reject optimization ↓ 7. RE-BENCHMARK Verify denoise latency improvement; no regression on other models ``` --- ## Checklist Before Merging ### Correctness (must pass first) - [ ] Reference outputs collected with `--seed=42 --save-output` **before** any change - [ ] After change: regenerate with identical args and compare - [ ] No visible quality degradation in generated images / videos - [ ] Correctness verified on all benchmark models ### Performance (only after correctness passes) - [ ] All benchmark models executed; denoise latency ★, end-to-end, peak memory recorded - [ ] No regression in denoise latency vs. previous baseline (±2% tolerance) - [ ] New kernel shows measurable improvement on at least 2 models - [ ] No new torch.compile graph breaks introduced - [ ] Results reproducible with all offloads disabled and fixed `--seed=42` ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/nsight-profiler.md ================================================ --- name: nsight-profiler description: Expert skill for NVIDIA Nsight Systems and Nsight Compute profiling tools. Configure profiling sessions, analyze kernel reports, interpret occupancy metrics, roofline model data, memory bandwidth bottlenecks, and warp execution efficiency. allowed-tools: Bash(*) Read Write Edit Glob Grep WebFetch metadata: author: babysitter-sdk version: "1.0.0" category: performance-profiling backlog-id: SK-002 source: "Adapted from https://github.com/lobehub/lobehub (.agents/skills/nsight-profiler)" --- > **Source**: This skill is adapted from the [lobehub/lobehub](https://github.com/lobehub/lobehub) open-source repository (`.agents/skills/nsight-profiler`). Original author: `babysitter-sdk`. # nsight-profiler You are **nsight-profiler** - a specialized skill for NVIDIA Nsight Systems and Nsight Compute profiling tools. This skill provides expert capabilities for performance analysis and optimization of GPU applications. ## Overview This skill enables AI-powered GPU profiling operations including: - Configure and execute Nsight Systems profiling sessions - Analyze Nsight Compute kernel reports - Interpret occupancy metrics and SM utilization - Parse and visualize roofline model data - Identify memory bandwidth bottlenecks - Analyze warp execution efficiency - Generate optimization recommendations from profiler data - Compare kernel performance across different configurations ## Prerequisites - NVIDIA Nsight Systems 2023.1+ - NVIDIA Nsight Compute 2023.1+ - CUDA Toolkit 11.0+ - GPU with compute capability 7.0+ (for full profiling features) ## Capabilities ### 1. Nsight Systems Profiling System-wide performance analysis: ```bash # Basic system profile nsys profile -o report ./cuda_program # Profile with CUDA API tracing nsys profile -t cuda,nvtx,osrt -o report ./cuda_program # Capture GPU metrics nsys profile --gpu-metrics-device=all -o report ./cuda_program # Profile specific duration nsys profile -d 10 -o report ./cuda_program # Export to multiple formats (one type per command) nsys export -t sqlite report.nsys-rep nsys export -t json report.nsys-rep # Generate summary statistics nsys stats report.nsys-rep ``` ### 2. Nsight Compute Profiling Detailed kernel analysis: ```bash # Profile all kernels ncu -o profile ./cuda_program # Profile specific kernel ncu --kernel-name myKernel -o profile ./cuda_program # Full metric collection ncu --set full -o profile ./cuda_program # Roofline analysis ncu --set roofline -o profile ./cuda_program # Memory analysis ncu --section MemoryWorkloadAnalysis -o profile ./cuda_program # Compare two runs ncu --import baseline.ncu-rep --diff ./cuda_program ``` ### 3. Occupancy Analysis Analyze and optimize occupancy: ```bash # Collect occupancy metrics ncu --section Occupancy -o occupancy ./cuda_program # Key metrics to analyze: # - Achieved Occupancy # - Theoretical Occupancy # - Block Limit (registers, shared memory, warps) # - Occupancy Limiter ``` ```cuda // Query occupancy in code int numBlocks; int blockSize = 256; cudaOccupancyMaxActiveBlocksPerMultiprocessor( &numBlocks, myKernel, blockSize, sharedMemSize); float occupancy = (numBlocks * blockSize) / (float)deviceProp.maxThreadsPerMultiProcessor; printf("Theoretical Occupancy: %.2f%%\n", occupancy * 100); ``` ### 4. Roofline Model Analysis Performance bound analysis: ```bash # Generate roofline data ncu --set roofline -o roofline ./cuda_program # Key metrics: # - Achieved FLOP/s # - Achieved Memory Bandwidth # - Arithmetic Intensity (FLOP/byte) # - Ridge Point ``` Interpretation guide: - Below memory roofline: Memory bound - Below compute roofline: Compute bound - At peak: Optimal utilization ### 5. Memory Bandwidth Analysis Identify memory bottlenecks: ```bash # Memory analysis sections ncu --section MemoryWorkloadAnalysis \ --section MemoryWorkloadAnalysis_Chart \ --section MemoryWorkloadAnalysis_Tables \ -o memory ./cuda_program ``` Key metrics: - Global Load/Store Throughput - L1/L2 Cache Hit Rate - Shared Memory Bandwidth - Memory Transactions per Request ### 6. Warp Execution Analysis Analyze warp efficiency: ```bash # Warp state analysis ncu --section WarpStateStatistics -o warp ./cuda_program # Scheduler statistics ncu --section SchedulerStatistics -o scheduler ./cuda_program ``` Key metrics: - Warp Cycles Per Issued Instruction - Eligible Warps Per Active Cycle - Active Warps Per Scheduler - Stall Reasons (memory, sync, execution) ### 7. Kernel Comparison Compare kernel variants: ```bash # Step 1: Profile baseline ncu --set full -o baseline ./program_v1 # Step 2: Profile optimized version ncu --set full -o optimized ./program_v2 # Step 3: Export both profiles to CSV, then compare with Python (no GUI needed) # Note: --import can only be specified once; --page diff is not a valid page value. ncu --import baseline.ncu-rep --page details --csv > baseline_details.csv ncu --import optimized.ncu-rep --page details --csv > optimized_details.csv python3 -c " import csv def load(p): return {r.get('Metric Name',''): r.get('Metric Value','') for r in csv.DictReader(open(p))} b = load('baseline_details.csv') o = load('optimized_details.csv') for k in sorted(set(b) | set(o)): bv, ov = b.get(k,''), o.get(k,'') if bv != ov: print(f'{k[:55]:<55} {bv} -> {ov}') " ### 8. Performance Recommendations Automated analysis: ```bash # Get optimization recommendations ncu --section SpeedOfLight \ --section SpeedOfLight_RooflineChart \ -o speedoflight ./cuda_program # Export with recommendations ncu --import profile.ncu-rep --page details --csv > details.csv ``` ## Common Profiling Workflows ### Workflow 1: Initial Performance Assessment ```bash # Step 1: System overview nsys profile -t cuda -o system_overview ./program nsys stats system_overview.nsys-rep # Step 2: Identify hot kernels ncu --launch-skip 10 --launch-count 5 -o hot_kernels ./program # Step 3: Deep dive on bottleneck kernel ncu --kernel-name hotKernel --set full -o detailed ./program ``` ### Workflow 2: Memory Optimization ```bash # Analyze memory access patterns ncu --section SourceCounters \ --section MemoryWorkloadAnalysis \ --kernel-name targetKernel \ -o memory_analysis ./program # Check for coalescing issues ncu --metrics l1tex__t_sectors_pipe_lsu_mem_global_op_ld.sum,\ l1tex__t_requests_pipe_lsu_mem_global_op_ld.sum \ -o coalescing ./program ``` ### Workflow 3: Occupancy Optimization ```bash # Profile with occupancy focus ncu --section Occupancy \ --section LaunchStatistics \ -o occupancy ./program ``` **Interpreting occupancy limiters** (from the `Occupancy` section report): | Limiter shown | Fix | |---------------|-----| | `Registers` | Reduce register pressure: use fewer local variables, add `maxnreg` hint | | `Shared Memory` | Decrease shared memory allocation or use 32-bit instead of 64-bit | | `Block Size` | Increase threads per block; ensure block size is a multiple of warp size (32) | | `Warp Limit` | Already at theoretical max for this SM; no action needed | > **For Triton kernels**: block sizes are controlled via `@triton.autotune` configs, not CLI flags. To test occupancy at different block sizes, add or modify the `triton.Config({"BLOCK_C": N}, num_warps=W)` entries in the autotune list and re-run. Do **not** pass `--block-size` as a CLI argument — the Triton benchmark script does not accept it. ## Dependencies - Nsight Systems 2023.1+ - Nsight Compute 2023.1+ - CUDA Toolkit 11.0+ ## Constraints - Full profiling requires root/admin privileges - Some metrics only available on specific GPU architectures - Profiling adds overhead; results may differ from production - Nsight Compute profiles one kernel invocation at a time by default ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/references/a100-optimization-guide.md ================================================ # A100 GPU Optimization Guide — SGLang Diffusion JIT Kernels Deep dive into A100-specific optimizations for diffusion model CUDA kernels in SGLang's JIT system. > **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels) --- ## A100 Ampere Architecture Overview | Component | A100 40GB | A100 80GB | Notes | |-----------|-----------|-----------|-------| | Compute Capability | sm_80 | sm_80 | Use `"-arch=sm_80"` in `extra_cuda_cflags` | | SMs | 108 | 108 | Grid: aim for multiples of 108 | | Shared Memory | 164 KB/SM | 164 KB/SM | Configurable: 48/96/164 KB | | L2 Cache | 40 MB | 40 MB | Less than H100 (50 MB) | | Memory Bandwidth | 1.55 TB/s | 2.0 TB/s | HBM2e | | Max Threads/SM | 2048 | 2048 | Same as H100 | | Tensor Cores | 3rd gen | 3rd gen | FP16, BF16, TF32, INT8, INT4 | ### A100 vs H100 Comparison | Feature | A100 | H100 | Impact on JIT Kernels | |---------|------|------|-----------------------| | Memory BW | 2.0 TB/s | 3.35 TB/s | H100 ~67% faster for memory-bound ops | | SMs | 108 | 132 | Adjust persistent kernel grid sizing | | Shared Mem/SM | 164 KB | 192 KB | Reduce max tile sizes on A100 | | L2 Cache | 40 MB | 50 MB | Attention tile reuse still works well | | TMA | No | Yes | Can't use `cp.async.bulk` on A100 | | FP8 | No | Yes | Use FP16/BF16 only on A100 | --- ## Memory Access Optimization Same coalescing and vectorization rules as H100; lower bandwidth makes them even more critical. ### `AlignedVector` Vectorization (same pattern as H100) ```cpp #include constexpr int kVecN = 16 / sizeof(T); // 8 for bf16/fp16, 4 for fp32 using vec_t = device::AlignedVector; vec_t v; v.load(src, vi); // ... process elements ... v.store(dst, vi); ``` **Expected A100 performance (BF16 RMSNorm):** | Implementation | A100 (ms) | H100 (ms) | A100 Speedup | |:---|:---:|:---:|:---:| | Scalar loads | ~0.10 | 0.065 | 1.00x | | `AlignedVector` | ~0.03 | 0.019 | ~3x | **Target bandwidth**: 30–40% of A100's 2.0 TB/s = 600–800 GB/s. ### Shared Memory Configuration ```cpp // A100 max: 164 KB/SM cudaFuncSetAttribute( your_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 164 * 1024 // 164 KB max on A100 ); ``` Attention tile sizes for A100: ``` BLOCK_SIZE_M = 128 (Q block) BLOCK_SIZE_N = 64 (K,V block) Tile = 128×64×2 = 16 KB (FP16) — fits in 164 KB shared mem ``` --- ## Occupancy Tuning **Grid sizing for A100 (108 SMs):** ```cpp #include // Cap blocks to SM × occupancy (same pattern as H100) static const uint32_t max_occ = host::runtime::get_blocks_per_sm(kernel, kBlockSize); static const uint32_t num_sm = host::runtime::get_sm_count(device.unwrap().device_id); const uint32_t num_blocks = std::min(num_sm * max_occ, host::div_ceil(n, kBlockSize)); ``` **Recommended block sizes (same as H100):** | Kernel Type | Threads/Block | Notes | |-------------|---------------|-------| | Element-wise | 256 | High occupancy | | Row reduction | 512 | Full reduction per row | | Tiled/attention | 256 | Balance shared mem | --- ## A100-Specific Features ### Async Memory Copy (sm_80) A100 introduced `cp.async` for overlapping compute and memory. Use this in custom kernels for prefetching: ```cuda #if __CUDA_ARCH__ >= 800 // Async copy from global to shared (A100+) __pipeline_memcpy_async(smem_ptr, global_ptr, bytes); __pipeline_commit(); __pipeline_wait_prior(0); #endif ``` ### TF32 Mode (A100 specific) Enables FP32-range with FP16-like throughput for GEMM. Enable in Python: ```python # Enable TF32 for matmuls (A100+) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True ``` TF32 is automatic for FP32 GEMMs via cuBLAS — no kernel changes needed. ### Structural Sparsity (2:4) A100 tensor cores support 50% structured sparsity: ```python from torch.sparse import to_sparse_semi_structured sparse_weight = to_sparse_semi_structured(dense_weight) # ~2x GEMM speedup for matmul with sparse weight ``` --- ## JIT Compilation for A100 ```python return load_jit( "my_kernel", *args, cuda_files=["diffusion/my_kernel.cuh"], cuda_wrappers=[("my_kernel", f"my_kernel<{args}>")], extra_cuda_cflags=[ "-O3", "--use_fast_math", "-arch=sm_80", # A100 only; omit for multi-arch ], ) ``` **Multi-arch (A100 + H100):** ```python extra_cuda_cflags=[ "-O3", "--use_fast_math", "-gencode=arch=compute_80,code=sm_80", # A100 "-gencode=arch=compute_90,code=sm_90", # H100 ] ``` Runtime arch guard (in Python wrapper): ```python cap = torch.cuda.get_device_capability() if cap < (8, 0): raise RuntimeError(f"This kernel requires sm_80 (A100) or later, got sm_{cap[0]}{cap[1]}") ``` --- ## H100 → A100 Migration Checklist When porting an H100-optimized kernel to A100: | Item | H100 | A100 | Change Required | |------|------|------|-----------------| | Shared memory | 192 KB | 164 KB | Reduce `cudaFuncSetAttribute` size | | Grid sizing | ×132 SMs | ×108 SMs | `get_sm_count()` handles automatically | | TMA bulk copy | Available | **Not available** | Remove `cp.async.bulk`; use standard `__pipeline_memcpy_async` | | FP8 | Available | **Not available** | Fall back to FP16/BF16 | | PDL | Supported | Supported | `.enable_pdl(true)` works on sm_80 | | Warp shuffles | Same | Same | No changes | | `AlignedVector` | Same | Same | No changes | **Conditional compilation:** ```cuda #if __CUDA_ARCH__ >= 900 // H100-only: TMA, FP8, thread block clusters #define USE_TMA 1 #elif __CUDA_ARCH__ >= 800 // A100: cp.async, TF32, 2:4 sparsity #define USE_ASYNC_COPY 1 #endif ``` --- ## Precision Notes | Type | Available on A100 | Notes | |------|-------------------|-------| | FP16 | Yes | Good, watch overflow in attention | | BF16 | Yes | Preferred for training and inference | | TF32 | Yes (A100 specific) | Auto for FP32 GEMMs | | FP8 | **No** | H100 only | --- ## Performance Profiling ### NVIDIA Nsight Systems (nsys) ```bash nsys profile -o a100_profile python scripts/bench_diffusion_rmsnorm.py # Key metrics to watch: # - Kernel duration # - Memory transfer time # - GPU idle time # - Stream utilization ``` ### NVIDIA Nsight Compute (ncu) ```bash # Full metrics ncu --set full -o a100_metrics.ncu-rep \ python scripts/bench_diffusion_rmsnorm.py # Specific metrics for bandwidth / occupancy checks ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\ dram__throughput.avg.pct_of_peak_sustained_elapsed \ python scripts/bench_diffusion_rmsnorm.py # Key metrics for A100 diffusion kernels: # - Achieved occupancy (sm__warps_active.avg.pct_of_peak_sustained_active) # - Memory throughput (dram__throughput.avg.pct_of_peak_sustained_elapsed) # → Target: 30–40% of 2.0 TB/s (600–800 GB/s) for vectorized kernels # - Compute throughput (sm__throughput.avg.pct_of_peak_sustained_elapsed) # - Warp stall reasons (smsp__warp_issue_stalled_*.avg.pct_of_peak_sustained_active) # - Kernel time (gpu__time_duration.avg) ``` ### Common A100 Performance Issues 1. **Memory bound below target**: `dram__throughput` < 30% - Fix: Use `AlignedVector` (128-bit vector loads) 2. **Low occupancy**: Grid too small for 108 SMs - Fix: Use `runtime::get_sm_count()` persistent kernel pattern 3. **No TF32 for FP32 GEMMs**: torch.backends.cuda.matmul.allow_tf32 not set - Fix: `torch.backends.cuda.matmul.allow_tf32 = True` --- ## Best Practices Summary (A100) 1. **Bandwidth**: Even more critical than H100 — profile with `ncu` first 2. **Vectorization**: `AlignedVector` gives ~3x over scalar 3. **TF32**: Enable for any FP32 matmul workload 4. **Shared memory**: Cap at 164 KB; use `cudaFuncSetAttribute` 5. **Grid sizing**: Multiples of 108 SMs via `runtime::get_sm_count` 6. **cp.async**: Use for prefetching in tiled kernels 7. **Multi-arch**: Build for both `sm_80` and `sm_90` to support both GPUs 8. **Same abstractions**: `AlignedVector`, `TensorMatcher`, `LaunchKernel` work identically ## Reference Benchmark Results (A100 80GB, BF16) | Kernel | Shape | A100 (ms) | H100 (ms) | H100 Speedup | |--------|-------|-----------|-----------|--------------| | RMSNorm | [2, 1024, 2048] | ~0.08 | 0.054 | 1.5x | | GEGLU | [2, 1024, 4096] | ~0.05 | 0.030 | 1.7x | ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/references/h100-optimization-guide.md ================================================ # H100 GPU Optimization Guide — SGLang Diffusion JIT Kernels Deep dive into H100-specific optimizations for diffusion model CUDA kernels, written for SGLang's JIT kernel system. > **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels) --- ## H100 Hopper Architecture Overview | Component | Specification | Optimization Implication | |-----------|---------------|--------------------------| | Compute Capability | sm_90 | Use `extra_cuda_cflags=["-arch=sm_90"]` in `load_jit` | | SMs | 132 | Grid: aim for multiples of 132 | | Shared Memory | 192 KB/SM | Configurable: 96/144/192 KB | | L2 Cache | 50 MB | Tile K,V of attention to fit in L2 | | Memory Bandwidth | 3.35 TB/s | BF16 vectorized: achieves ~38% (~1.27 TB/s) | | Max Threads/SM | 2048 | Max 16 blocks of 128 threads per SM | | Warp Size | 32 | All reductions use `warp::reduce_sum` | | Registers | 64K 32-bit/SM | 255 per thread max | ### New Hopper Features (sm_90+) 1. **Thread Block Clusters** — groups cooperating via Distributed Shared Memory 2. **TMA (Tensor Memory Accelerator)** — hardware-accelerated bulk copies 3. **FP8 support** — native 8-bit floating point in tensor cores 4. **PDL (Programmatic Dependent Launch)** — enable with `.enable_pdl(true)` in `LaunchKernel` Gate sm_90+ features with a runtime check before calling `load_jit`: ```python if torch.cuda.get_device_capability()[0] < 9: raise RuntimeError("This kernel requires H100 (sm_90+)") ``` --- ## Memory Hierarchy Optimization ### Coalesced Global Memory Access ```cpp // GOOD: threads read consecutive addresses → 128-byte transaction per warp uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x; fp16_t val = src[idx]; // BAD: strided access → multiple transactions, lower effective bandwidth uint32_t idx = threadIdx.x * stride; // avoid stride > 1 ``` **Transaction sizes**: 32 bytes minimum, 128 bytes optimal (full warp, FP32). ### Vectorized Memory Access with `AlignedVector` SGLang's `AlignedVector` provides 128-bit (16-byte) vector loads. Always use this instead of raw pointer reinterprets. ```cpp #include // 16 bytes per load: 8×bf16_t, 8×fp16_t, or 4×fp32_t constexpr int kVecN = 16 / sizeof(T); using vec_t = device::AlignedVector; // Load vec_t v; v.load(src, vi); // loads src[vi * kVecN .. vi * kVecN + kVecN - 1] // Process #pragma unroll for (int i = 0; i < kVecN; ++i) { float val = static_cast(v[i]); // ... compute ... v[i] = static_cast(result); } // Store v.store(dst, vi); ``` **RMSNorm benchmark (H100 80GB, BF16):** | Implementation | Time (ms) | Speedup | |:---|:---:|:---:| | Scalar loads | 0.065 | 1.00x | | `AlignedVector` | 0.019 | **3.37x** | Bandwidth achieved: **~38% of 3.35 TB/s** = 1.27 TB/s. ### L2 Cache Utilization (50 MB) For attention, tile K and V so they stay in L2 while Q iterates: ``` BLOCK_SIZE_M = 128 (Q block) BLOCK_SIZE_N = 64 (K,V block) With head_dim=64: tile = 128×64×2 = 16 KB (FP16), multiple tiles fit in L2 ``` ### Shared Memory Configuration Request max shared memory for attention kernels: ```cpp // In launcher (after selecting kernel function pointer): cudaFuncSetAttribute( your_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 192 * 1024 // 192 KB max on H100 ); ``` Shared memory has 32 banks (4 bytes/bank). Avoid conflicts with padding: ```cpp __shared__ float data[32][33]; // 33 instead of 32 → no bank conflict ``` --- ## Warp & CTA Reductions (SGLang Abstractions) Use `sgl_kernel/warp.cuh` and `sgl_kernel/cta.cuh` — never raw `__shfl_xor_sync`. ```cpp #include #include // Warp-level sum (uses __shfl_xor_sync internally) float result = device::warp::reduce_sum(partial); // Warp-level max float mx = device::warp::reduce_max(val); // CTA-wide max via shared memory __shared__ float smem[32]; device::cta::reduce_max(val, smem, -1e38f); // smem[0] holds the result after __syncthreads() ``` **Block reduction pattern for RMSNorm:** ```cpp // 1. Warp reduction sum_sq = device::warp::reduce_sum(sum_sq); // 2. Write warp leaders to smem __shared__ float smem_r[32]; if (threadIdx.x % 32 == 0) smem_r[threadIdx.x / 32] = sum_sq; __syncthreads(); // 3. Final warp reduction over warp leaders if (threadIdx.x < 32) { sum_sq = (threadIdx.x < blockDim.x / 32) ? smem_r[threadIdx.x] : 0.f; sum_sq = device::warp::reduce_sum(sum_sq); } __syncthreads(); ``` --- ## Occupancy Tuning ``` Occupancy = Active Warps per SM / Max Warps per SM (64) Limiting factors on H100: 1. Registers: 65536 / (threads_per_block × regs_per_thread) 2. Shared Memory: 192 KB / smem_per_block 3. Threads: 2048 / threads_per_block ``` **Recommended block sizes:** | Kernel Type | Threads/Block | Warps | Reasoning | |-------------|---------------|-------|-----------| | Element-wise (RoPE, GEGLU) | 256 | 8 | High occupancy, simple | | Row reduction (RMSNorm, LayerNorm) | 256–512 | 8–16 | Enough threads for full reduction | | Tiled (attention) | 256 | 8 | Balance shared mem and registers | **Persistent kernel pattern** (cap grid to SM × occupancy): ```cpp #include static const uint32_t max_occ = host::runtime::get_blocks_per_sm(kernel, kBlockSize); static const uint32_t num_sm = host::runtime::get_sm_count(device.unwrap().device_id); const uint32_t num_blocks = std::min(num_sm * max_occ, host::div_ceil(n, kBlockSize)); host::LaunchKernel(num_blocks, kBlockSize, device.unwrap())(kernel, params); ``` --- ## Precision and Numerical Stability | Type | Exponent Bits | Mantissa Bits | Range | Use Case | |------|--------------|---------------|-------|----------| | FP16 | 5 | 10 | ±65504 | Inference; attention score overflow risk | | BF16 | 8 | 7 | ±3.39×10³⁸ | Training/inference preferred; safer for attn | | FP32 | 8 | 23 | ±3.39×10³⁸ | Accumulation only | **Mixed precision pattern** (always accumulate in FP32): ```cpp // Input via AlignedVector vec_t v; v.load(src, vi); float acc = 0.f; #pragma unroll for (int i = 0; i < kVecN; ++i) { float val = static_cast(v[i]); // promote to FP32 acc += val * val; } // Output v[i] = static_cast(fp32_result); // demote back ``` --- ## Diffusion-Specific Patterns ### DiT Block Operators | Operator | Pattern | Key Constraint | |----------|---------|----------------| | **RMSNorm** | 2-pass row reduction | weight may be `None` | | **AdaLN** | `norm(x) * (1 + scale) + shift` | fuse norm+scale+shift | | **RoPE 3D** | `[B, t*h*w, heads, head_dim]` | layout: `seq = t*h*w` | | **GEGLU** | `gelu(gate) * value`, input `[B,L,2H]` | don't use for LTX-Video (uses GELU) | | **SiLU gate** | `x * sigmoid(x)` | fuse with MLP linear | ### Online Softmax (for custom attention) ```cuda // Numerically stable without materializing full [seq×seq] score matrix float row_max = -INFINITY, row_sum = 0.f; for each K block: compute local_scores new_max = max(row_max, max(local_scores)) rescale = exp(row_max - new_max) row_sum = row_sum * rescale + sum(exp(local_scores - new_max)) out_acc = out_acc * rescale + softmax(local_scores) @ V_block row_max = new_max ``` --- ## Profiling and Debugging ### NVIDIA Nsight Systems (nsys) System-wide profiling to see kernel durations, memory transfers, and GPU idle time: ```bash nsys profile -o profile_report python scripts/bench_diffusion_rmsnorm.py # Key metrics to watch: # - Kernel duration # - Memory transfer time # - GPU idle time # - Stream utilization ``` For end-to-end denoise profiling via `sglang generate`, see `diffusion-benchmark-and-profile.md` (Level 2: nsys + gputrc2graph.py). ### NVIDIA Nsight Compute (ncu) Detailed per-kernel analysis for tuning individual JIT CUDA kernels: ```bash # Full metrics — use when you need everything (slow) ncu --set full -o metrics.ncu-rep \ python scripts/bench_diffusion_rmsnorm.py # Specific metrics — use for targeted bandwidth / occupancy checks ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\ dram__throughput.avg.pct_of_peak_sustained_elapsed \ python scripts/bench_diffusion_rmsnorm.py # Key metrics for diffusion JIT kernels: # - Achieved occupancy (sm__warps_active.avg.pct_of_peak_sustained_active) # - Memory throughput (dram__throughput.avg.pct_of_peak_sustained_elapsed) # - Compute throughput (sm__throughput.avg.pct_of_peak_sustained_elapsed) # - Warp stall reasons (smsp__warp_issue_stalled_*.avg.pct_of_peak_sustained_active) # - L1 cache hit rate (l1tex__t_requests_pipe_lsu_mem_global_op_ld.sum) ``` ### Common Performance Issues 1. **Low occupancy**: Too many registers or shared memory per block - Check: `--ptxas-options=-v` in `extra_cuda_cflags` to see register count - Fix: Reduce `--maxrregcount=N`; use smaller block size 2. **Memory bound, low bandwidth**: Achieved < 30% of 3.35 TB/s - Check: `dram__throughput.avg.pct_of_peak_sustained_elapsed` - Fix: Switch to `AlignedVector` for 128-bit vector loads 3. **Shared memory bank conflicts**: `l1tex__data_bank_conflicts_pipe_lmem_op_st.sum` is high - Fix: Add padding — `__shared__ float data[32][33]` 4. **Warp divergence**: Conditional branches splitting warps - Check: `smsp__warp_issue_stalled_branch.avg.pct_of_peak_sustained_active` - Fix: Restructure so elements with identical branches are in the same warp 5. **Too many small kernels**: High kernel launch overhead - Fix: Fuse operations (e.g., norm + scale + shift → AdaLN in one kernel) --- ## JIT Compilation Notes SGLang's JIT compiles kernels on first use via `load_jit`. For H100-specific flags: ```python return load_jit( "my_kernel", *args, cuda_files=["diffusion/my_kernel.cuh"], cuda_wrappers=[("my_kernel", f"my_kernel<{args}>")], extra_cuda_cflags=[ "-O3", "--use_fast_math", "-arch=sm_90", # H100 only; omit for multi-arch "--ptxas-options=-v", # Remove after tuning ], ) ``` For multi-arch (H100 + A100): ```python extra_cuda_cflags=[ "-O3", "--use_fast_math", "-gencode=arch=compute_80,code=sm_80", # A100 "-gencode=arch=compute_90,code=sm_90", # H100 ] ``` --- ## Best Practices Summary 1. **Memory access**: Coalesce writes, align to 128-byte boundaries 2. **Vectorization**: Use `AlignedVector` for all element-wise loads/stores 3. **Reductions**: Use `warp::reduce_sum/max`, then shared memory pattern above 4. **Precision**: BF16 for I/O, FP32 for accumulation; use `static_cast` 5. **Block size**: 256 threads default; 512 for reductions; tune with `runtime::get_blocks_per_sm` 6. **Grid sizing**: Multiples of 132 SMs; use persistent kernel pattern for small N 7. **Shared memory**: Add padding (`[32][33]`) to avoid bank conflicts 8. **Profile**: Run `ncu` before claiming a speedup; check dram throughput % 9. **Fuse**: Combine norm + scale + shift into a single pass to reduce memory traffic 10. **Abstractions**: Always use `TensorMatcher`, `AlignedVector`, `LaunchKernel` — never raw CUDA ## Reference Benchmark Results (H100 80GB, BF16) | Kernel | Shape | Time (ms) | |--------|-------|-----------| | RMSNorm | [2, 1024, 2048] | 0.054 | | GEGLU | [2, 1024, 4096] → [2, 1024, 2048] | 0.030 | | RoPE 3D | [2, 480, 8, 64] | 1.670 | | RMSNorm vectorized | [1, 1024, 2048] | 0.019 | | RMSNorm vectorized | [4, 4096, 3072] | 0.157 | > See `kernel-templates.md` for copy-paste ready sglang JIT kernel implementations. ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/references/kernel-templates.md ================================================ # CUDA Kernel Templates — SGLang Diffusion JIT Style Copy-paste ready templates for JIT CUDA kernels in `python/sglang/jit_kernel/csrc/diffusion/`. All templates use SGLang's internal abstractions; no raw CUDA headers needed. > **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels) --- ## Prerequisite: Standard Includes Every kernel file in `csrc/diffusion/` starts with: ```cpp #include // TensorMatcher, SymbolicSize, SymbolicDevice #include // fp16_t, bf16_t, fp32_t, dtype_trait, packed_t #include // RuntimeCheck, Panic, div_ceil #include // LaunchKernel, SGL_DEVICE, type aliases #include // AlignedVector #include // warp::reduce_sum, warp::reduce_max #include // device::math::rsqrt, sqrt, ... #include // tile::Memory (strided access pattern) #include #include ``` **Key type aliases** (from `utils.cuh`): - `fp16_t` = `__half`, `fp16x2_t` = `__half2` - `bf16_t` = `__nv_bfloat16`, `bf16x2_t` = `__nv_bfloat162` - `fp32_t` = `float`, `fp32x2_t` = `float2` - `SGL_DEVICE` = `__forceinline__ __device__` --- ## Template 1: Element-wise Operation Use for ops that process elements independently: RoPE, SiLU, GEGLU, scale+bias. ### `.cuh` file: `csrc/diffusion/silu_gate.cuh` ```cpp #include #include #include #include #include #include #include #include namespace { // SiLU gate: out[i] = x[i] * sigmoid(x[i]) // Input layout: [B, L, hidden] template __global__ void silu_gate_kernel( T* __restrict__ dst, const T* __restrict__ src, uint32_t n_vecs, uint32_t n_remainder, uint32_t n_total) { using vec_t = device::AlignedVector; const uint32_t stride = blockDim.x * gridDim.x; // --- vectorized body --- for (uint32_t vi = blockIdx.x * blockDim.x + threadIdx.x; vi < n_vecs; vi += stride) { vec_t v; v.load(src, vi); #pragma unroll for (int i = 0; i < kVecN; ++i) { float val = static_cast(v[i]); float sig = 1.f / (1.f + device::math::exp(-val)); v[i] = static_cast(val * sig); } v.store(dst, vi); } // --- scalar tail (for sizes not divisible by kVecN) --- const uint32_t base = n_vecs * kVecN; for (uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n_remainder; i += stride) { float val = static_cast(src[base + i]); float sig = 1.f / (1.f + device::math::exp(-val)); dst[base + i] = static_cast(val * sig); } } template void silu_gate(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { using namespace host; SymbolicSize N{"num_elements"}; SymbolicDevice device; device.set_options(); TensorMatcher({N}) .with_dtype() .with_device(device) .verify(dst) .verify(src); const uint32_t n = static_cast(N.unwrap()); const DLDevice dev = device.unwrap(); RuntimeCheck(n > 0, "silu_gate: num_elements must be > 0"); constexpr int kVecN = 16 / sizeof(T); // 128-bit vector load const uint32_t n_vecs = n / kVecN; const uint32_t n_rem = n % kVecN; constexpr uint32_t kBlock = 256; const uint32_t grid = div_ceil(std::max(n_vecs, n_rem), kBlock); LaunchKernel(grid, kBlock, dev)( silu_gate_kernel, static_cast(dst.data_ptr()), static_cast(src.data_ptr()), n_vecs, n_rem, n); } } // namespace ``` ### Python wrapper: `diffusion/silu_gate.py` ```python from __future__ import annotations import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args @cache_once def _jit_silu_gate_module(dtype: torch.dtype): args = make_cpp_args(dtype) return load_jit( "diffusion_silu_gate", *args, cuda_files=["diffusion/silu_gate.cuh"], cuda_wrappers=[("silu_gate", f"silu_gate<{args}>")], extra_cuda_cflags=["-O3", "--use_fast_math"], ) def diffusion_silu_gate(src: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor: assert src.is_cuda and src.dtype in (torch.float16, torch.bfloat16, torch.float32) if out is None: out = torch.empty_like(src) module = _jit_silu_gate_module(src.dtype) module.silu_gate(out, src) return out ``` --- ## Template 2: Row-wise Reduction (RMSNorm / LayerNorm) Use for ops that reduce across the last dimension of each row. ### `.cuh` file: `csrc/diffusion/rmsnorm.cuh` ```cpp #include #include #include #include #include #include #include #include #include namespace { // RMSNorm: y = x / rms(x) * weight // One block per row; vectorized loads/stores; warp + shared-mem reduction template __global__ void rmsnorm_kernel( T* __restrict__ dst, const T* __restrict__ src, const T* __restrict__ weight, // nullptr if no affine weight uint32_t hidden, uint32_t n_vecs, float eps) { using vec_t = device::AlignedVector; const uint32_t row = blockIdx.x; const T* row_src = src + row * hidden; T* row_dst = dst + row * hidden; // Pass 1: sum of squares float sum_sq = 0.f; for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { vec_t v; v.load(row_src, vi); #pragma unroll for (int i = 0; i < kVecN; ++i) { float val = static_cast(v[i]); sum_sq += val * val; } } // Warp + block reduction sum_sq = device::warp::reduce_sum(sum_sq); __shared__ float smem[32]; if (threadIdx.x % 32 == 0) smem[threadIdx.x / 32] = sum_sq; __syncthreads(); if (threadIdx.x < 32) { sum_sq = (threadIdx.x < blockDim.x / 32) ? smem[threadIdx.x] : 0.f; sum_sq = device::warp::reduce_sum(sum_sq); } __syncthreads(); const float rms_inv = device::math::rsqrt(sum_sq / static_cast(hidden) + eps); // Pass 2: normalize + optional weight for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { vec_t v_in, v_out; v_in.load(row_src, vi); if (weight != nullptr) { vec_t v_w; v_w.load(weight, vi); #pragma unroll for (int i = 0; i < kVecN; ++i) v_out[i] = static_cast(static_cast(v_in[i]) * rms_inv * static_cast(v_w[i])); } else { #pragma unroll for (int i = 0; i < kVecN; ++i) v_out[i] = static_cast(static_cast(v_in[i]) * rms_inv); } v_out.store(row_dst, vi); } } template void rmsnorm( tvm::ffi::TensorView dst, tvm::ffi::TensorView src, tvm::ffi::TensorView weight, // data_ptr == nullptr → no weight float eps) { using namespace host; SymbolicSize B{"batch_tokens"}, H{"hidden_size"}; SymbolicDevice device; device.set_options(); TensorMatcher({B, H}) .with_dtype() .with_device(device) .verify(dst) .verify(src); const uint32_t num_rows = static_cast(B.unwrap()); const uint32_t hidden = static_cast(H.unwrap()); const DLDevice dev = device.unwrap(); constexpr int kVecN = 16 / sizeof(T); RuntimeCheck(hidden % kVecN == 0, "rmsnorm: hidden_size (", hidden, ") must be divisible by ", kVecN); const uint32_t n_vecs = hidden / kVecN; uint32_t threads = std::min(n_vecs, 512u); threads = (threads + 31) / 32 * 32; const T* w_ptr = (weight.data_ptr() != nullptr) ? static_cast(weight.data_ptr()) : nullptr; LaunchKernel(num_rows, threads, dev)( rmsnorm_kernel, static_cast(dst.data_ptr()), static_cast(src.data_ptr()), w_ptr, hidden, n_vecs, eps); } } // namespace ``` --- ## Template 3: Fused Row-Reduction + Element-wise (AdaLN) Combines RMSNorm + AdaLN modulation into one pass: `y = norm(x) * (1 + scale) + shift`. ### `.cuh` file: `csrc/diffusion/adaln.cuh` ```cpp #include #include #include #include #include #include #include #include #include namespace { // AdaLN: y = norm(x) * (1 + scale) + shift // scale, shift: [batch, hidden] (one per row) template __global__ void adaln_kernel( T* __restrict__ dst, const T* __restrict__ src, const T* __restrict__ weight, const T* __restrict__ scale, const T* __restrict__ shift, uint32_t hidden, uint32_t n_vecs, float eps) { using vec_t = device::AlignedVector; const uint32_t row = blockIdx.x; const T* row_src = src + row * hidden; const T* row_scale = scale + row * hidden; const T* row_shift = shift + row * hidden; T* row_dst = dst + row * hidden; // Pass 1: compute RMS float sum_sq = 0.f; for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { vec_t v; v.load(row_src, vi); #pragma unroll for (int i = 0; i < kVecN; ++i) { float val = static_cast(v[i]); sum_sq += val * val; } } sum_sq = device::warp::reduce_sum(sum_sq); __shared__ float smem[32]; if (threadIdx.x % 32 == 0) smem[threadIdx.x / 32] = sum_sq; __syncthreads(); if (threadIdx.x < 32) { sum_sq = (threadIdx.x < blockDim.x / 32) ? smem[threadIdx.x] : 0.f; sum_sq = device::warp::reduce_sum(sum_sq); } __syncthreads(); const float rms_inv = device::math::rsqrt(sum_sq / static_cast(hidden) + eps); // Pass 2: normalize + modulate for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { vec_t v_in, v_w, v_sc, v_sh, v_out; v_in.load(row_src, vi); v_w.load(weight, vi); v_sc.load(row_scale, vi); v_sh.load(row_shift, vi); #pragma unroll for (int i = 0; i < kVecN; ++i) { float x = static_cast(v_in[i]) * rms_inv * static_cast(v_w[i]); float sc = static_cast(v_sc[i]); float sh = static_cast(v_sh[i]); v_out[i] = static_cast(x * (1.f + sc) + sh); } v_out.store(row_dst, vi); } } template void adaln( tvm::ffi::TensorView dst, tvm::ffi::TensorView src, tvm::ffi::TensorView weight, tvm::ffi::TensorView scale, tvm::ffi::TensorView shift, float eps) { using namespace host; SymbolicSize B{"batch_tokens"}, H{"hidden_size"}; SymbolicDevice device; device.set_options(); TensorMatcher({B, H}) .with_dtype() .with_device(device) .verify(dst).verify(src).verify(weight).verify(scale).verify(shift); const uint32_t num_rows = static_cast(B.unwrap()); const uint32_t hidden = static_cast(H.unwrap()); const DLDevice dev = device.unwrap(); constexpr int kVecN = 16 / sizeof(T); RuntimeCheck(hidden % kVecN == 0, "adaln: hidden_size must be divisible by ", kVecN); const uint32_t n_vecs = hidden / kVecN; uint32_t threads = std::min(n_vecs, 512u); threads = (threads + 31) / 32 * 32; LaunchKernel(num_rows, threads, dev)( adaln_kernel, static_cast(dst.data_ptr()), static_cast(src.data_ptr()), static_cast(weight.data_ptr()), static_cast(scale.data_ptr()), static_cast(shift.data_ptr()), hidden, n_vecs, eps); } } // namespace ``` --- ## Template 4: Python Wrapper (generic pattern) File location: `python/sglang/jit_kernel/diffusion/.py` ```python from __future__ import annotations from typing import TYPE_CHECKING import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args if TYPE_CHECKING: from tvm_ffi.module import Module @cache_once def _jit_module(dtype: torch.dtype) -> Module: """Cache key: dtype (and any other template params you need).""" args = make_cpp_args(dtype) return load_jit( "diffusion_your_op", # unique build cache key *args, cuda_files=["diffusion/your_op.cuh"], # relative to csrc/ cuda_wrappers=[("your_op", f"your_op<{args}>")], extra_cuda_cflags=["-O3", "--use_fast_math"], ) def diffusion_your_op( src: torch.Tensor, out: torch.Tensor | None = None, ) -> torch.Tensor: """ Your op description. Supported dtypes: float16, bfloat16, float32. """ assert src.is_cuda, "src must be a CUDA tensor" assert src.dtype in (torch.float16, torch.bfloat16, torch.float32), ( f"Unsupported dtype {src.dtype}" ) if out is None: out = torch.empty_like(src) module = _jit_module(src.dtype) module.your_op(out, src) return out ``` **`make_cpp_args` conversion table:** | `torch.dtype` | C++ type | |---------------|----------| | `torch.float16` | `fp16_t` | | `torch.bfloat16` | `bf16_t` | | `torch.float32` | `fp32_t` | --- ## Template 5: Correctness Test ```python # python/sglang/jit_kernel/tests/test_diffusion_.py import pytest import torch from sglang.jit_kernel.diffusion. import diffusion_ @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @pytest.mark.parametrize("shape", [(1, 2048), (4, 3072), (16, 4096)]) def test__correctness(dtype, shape): src = torch.randn(*shape, dtype=dtype, device="cuda") out_jit = diffusion_(src) ref = reference_(src.float()).to(dtype) # reference in fp32 tol = {"rtol": 1e-2, "atol": 1e-2} if dtype != torch.float32 else {"rtol": 1e-5, "atol": 1e-6} torch.testing.assert_close(out_jit, ref, **tol) def test__out_param(): src = torch.randn(1024, 2048, dtype=torch.bfloat16, device="cuda") out = torch.empty_like(src) result = diffusion_(src, out=out) assert result is out def test__cpu_error(): src = torch.randn(128, dtype=torch.float16) # CPU tensor with pytest.raises(AssertionError): diffusion_(src) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ``` --- ## Template 6: Benchmark ```python # python/sglang/jit_kernel/benchmark/bench_diffusion_.py import torch import triton.testing from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE, run_benchmark from sglang.jit_kernel.diffusion. import diffusion_ SHAPES = [(4096, 2048), (4096, 3072), (4096, 4096)] @triton.testing.perf_report( triton.testing.Benchmark( x_names=["hidden"], x_vals=[s[1] for s in SHAPES], line_arg="provider", line_vals=["jit_cuda", "torch"], line_names=["SGLang JIT CUDA", "PyTorch"], styles=[("blue", "-"), ("red", "--")], ylabel="us", plot_name="diffusion-", args={}, ) ) def benchmark(hidden: int, provider: str): src = torch.randn(4096, hidden, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) if provider == "jit_cuda": fn = lambda: diffusion_(src) else: fn = lambda: reference_(src) # torch baseline return run_benchmark(fn) if __name__ == "__main__": benchmark.run(print_data=True) ``` --- ## Summary of New Files per Kernel ``` python/sglang/jit_kernel/csrc/diffusion/ └── .cuh # CUDA kernel + launcher python/sglang/jit_kernel/diffusion/ └── .py # Python wrapper (load_jit + cache_once) python/sglang/jit_kernel/tests/ └── test_diffusion_.py # correctness tests python/sglang/jit_kernel/benchmark/ └── bench_diffusion_.py # triton.testing benchmark ``` > See `scripts/bench_diffusion_rmsnorm.py` and `scripts/bench_diffusion_denoise.py` for full runnable examples. ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/references/t4-optimization-guide.md ================================================ # T4 GPU Optimization Guide — SGLang Diffusion JIT Kernels T4 is a Turing architecture GPU (GCP n1+T4, AWS g4dn) commonly used for cloud inference. Its key constraint for diffusion kernels: **no BF16 support** — FP16 only. > **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels) --- ## T4 Turing Architecture Overview | Component | T4 | A100 | H100 | |-----------|-----|------|------| | Compute Capability | sm_75 | sm_80 | sm_90 | | SMs | 40 | 108 | 132 | | Shared Memory/SM | **64 KB** | 164 KB | 192 KB | | L2 Cache | 4 MB | 40 MB | 50 MB | | Memory Bandwidth | **320 GB/s** | 2.0 TB/s | 3.35 TB/s | | Memory | 16 GB GDDR6 | 40–80 GB HBM2e | 80 GB HBM3 | | Max Threads/SM | **1024** | 2048 | 2048 | | BF16 Support | **No** | Yes | Yes | ### Critical T4 Constraints 1. **No BFloat16** — must use FP16 everywhere 2. **320 GB/s bandwidth** — ~10x lower than H100; vectorization is critical 3. **16 GB memory** — limits model size; use offloading 4. **64 KB shared memory/SM** — smaller attention tiles 5. **Max 1024 threads/SM** — half of A100/H100; affects occupancy calculations --- ## No BF16: Always Use FP16 This is the most impactful constraint. **Never use `bf16_t` or `__nv_bfloat16` on T4.** **Python wrapper guard:** ```python import torch from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args @cache_once def _jit_rmsnorm_module(dtype: torch.dtype): # T4 (sm_75) does not support BF16 cap = torch.cuda.get_device_capability() if cap < (8, 0) and dtype == torch.bfloat16: raise RuntimeError( f"T4 (sm_75) does not support BF16. Use torch.float16 instead. " f"Got dtype={dtype}" ) args = make_cpp_args(dtype) return load_jit( "diffusion_rmsnorm", *args, cuda_files=["diffusion/rmsnorm.cuh"], cuda_wrappers=[("rmsnorm", f"rmsnorm<{args}>")], ) ``` **Conditional type in kernel:** ```cuda #if __CUDA_ARCH__ >= 800 // A100/H100: BF16 available using DefaultHalf = bf16_t; #else // T4/Turing: FP16 only using DefaultHalf = fp16_t; #endif ``` **Runtime detection helper:** ```python def get_diffusion_dtype() -> torch.dtype: """Return the appropriate half-precision dtype for the current GPU.""" cap = torch.cuda.get_device_capability() if cap >= (8, 0): return torch.bfloat16 # A100/H100: prefer BF16 else: return torch.float16 # T4/older: FP16 only ``` --- ## Memory Access Optimization With only 320 GB/s, **vectorization is more critical on T4 than on A100/H100**. ### `AlignedVector` (same abstraction, FP16 only) ```cpp #include // On T4, T must be fp16_t or fp32_t (NOT bf16_t) constexpr int kVecN = 16 / sizeof(T); // 8 for fp16, 4 for fp32 using vec_t = device::AlignedVector; ``` **Target bandwidth**: 40–50% of T4's 320 GB/s = 128–160 GB/s. ### Increase Arithmetic Intensity With low bandwidth, fusing ops saves more on T4 than on H100: ```cpp // BAD on T4: separate passes → 2× memory traffic output1[i] = input[i] * scale; // pass 1 output2[i] = output1[i] + bias; // pass 2 // GOOD: fuse → single memory read, single write float val = static_cast(v[i]); val = val * scale + bias; val = device::math::max(val, 0.f); // ReLU v[i] = static_cast(val); ``` ### Expected T4 Performance | Kernel | T4 (ms) | A100 (ms) | H100 (ms) | T4 vs H100 | |--------|---------|-----------|-----------|------------| | RMSNorm [2, 1024, 2048] | ~0.5 | ~0.08 | 0.054 | ~9x slower | | GEGLU [2, 1024, 4096] | ~0.3 | ~0.05 | 0.030 | ~10x slower | --- ## Shared Memory Configuration T4 max: **64 KB/SM**. Use smaller tiles vs A100/H100. ```cpp // T4: request max shared memory (64 KB) cudaFuncSetAttribute( your_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 64 * 1024 ); ``` **Attention tile sizes for T4** (halved vs H100): ``` H100/A100: BLOCK_SIZE_M = 128, BLOCK_SIZE_N = 64 T4: BLOCK_SIZE_M = 64, BLOCK_SIZE_N = 32 ← reduced for 64 KB limit ``` --- ## Occupancy Tuning T4 max: **1024 threads/SM** (vs 2048 on A100/H100). This halves max occupancy for a given block size. **Block sizes for T4:** | Kernel Type | Threads/Block | Notes | |-------------|---------------|-------| | Element-wise | 256 | Same as H100 | | Row reduction | 256–512 | Avoid > 512 to fit multiple blocks/SM | | Tiled/attention | 128–256 | Small tiles due to 64 KB shared mem | **Grid sizing for T4 (40 SMs)** — `runtime::get_sm_count` handles this automatically: ```cpp // get_sm_count() returns 40 on T4, 108 on A100, 132 on H100 const uint32_t num_sm = host::runtime::get_sm_count(device.unwrap().device_id); ``` --- ## Numerical Stability with FP16 FP16 has a smaller dynamic range (±65504) vs BF16 (±3.39×10³⁸). Watch for overflow in attention: ```cuda // Scale attention scores to prevent FP16 overflow float scale_factor = 1.0f / sqrtf(static_cast(head_dim)); // For very long sequences on T4, may need additional scaling: // if (score * scale_factor > 65000.f) { /* clamp */ } ``` Always accumulate in FP32: ```cpp float acc = 0.f; // FP32 accumulation for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { vec_t v; v.load(src, vi); #pragma unroll for (int i = 0; i < kVecN; ++i) { float val = static_cast(v[i]); // fp16 → fp32 acc += val * val; } } ``` --- ## Memory Management for 16 GB T4's 16 GB requires careful planning for large diffusion models. **sglang generate flags for T4:** ```bash # Enable CPU offloading to fit within 16 GB sglang generate \ --model-path=black-forest-labs/FLUX.1-dev \ --dit-cpu-offload true \ # DiT weights to CPU --text-encoder-cpu-offload true \ --vae-cpu-offload true \ --width=512 --height=512 \ # Reduce resolution --num-inference-steps=20 \ # Fewer steps --seed=42 ``` **Resolution recommendations for T4:** | Model | H100/A100 | T4 | |-------|-----------|-----| | FLUX.1-dev | 1024×1024 | 512×512 | | Wan2.2-TI2V-5B | 720P | 480P | | FLUX.2-dev | 1024×1024 | 512×512 | --- ## JIT Compilation for T4 ```python return load_jit( "my_kernel", *args, cuda_files=["diffusion/my_kernel.cuh"], cuda_wrappers=[("my_kernel", f"my_kernel<{args}>")], extra_cuda_cflags=[ "-O3", "--use_fast_math", "-arch=sm_75", # T4 only; omit for multi-arch ], ) ``` **Multi-arch (T4 + A100 + H100):** ```python extra_cuda_cflags=[ "-O3", "--use_fast_math", "-gencode=arch=compute_75,code=sm_75", # T4 "-gencode=arch=compute_80,code=sm_80", # A100 "-gencode=arch=compute_90,code=sm_90", # H100 ] ``` --- ## H100/A100 → T4 Migration Checklist | Item | H100/A100 | T4 | Action | |------|-----------|-----|--------| | BF16 | Available | **Not available** | Replace `bf16_t` with `fp16_t`; guard in Python wrapper | | Shared memory | 164–192 KB | **64 KB** | Halve tile sizes | | Grid sizing | ×108/132 SMs | ×40 SMs | `get_sm_count()` auto-handles | | Max threads/SM | 2048 | **1024** | Don't exceed 512 threads/block | | Memory | 40–80 GB | **16 GB** | Enable CPU offloading | | cp.async | Available | No (Turing has limited async) | Remove async copy patterns | | `AlignedVector` | Same | Same | No changes | | `warp::reduce_sum` | Same | Same | No changes | --- ## Performance Profiling ### NVIDIA Nsight Systems (nsys) ```bash nsys profile -o t4_profile python scripts/bench_diffusion_rmsnorm.py # Key metrics to watch: # - Kernel duration # - Memory transfer time # - GPU idle time # - Stream utilization ``` ### NVIDIA Nsight Compute (ncu) ```bash # Full metrics ncu --set full -o t4_metrics.ncu-rep \ python scripts/bench_diffusion_rmsnorm.py # Specific metrics — T4 is memory-bound; focus on dram throughput ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\ dram__throughput.avg.pct_of_peak_sustained_elapsed \ python scripts/bench_diffusion_rmsnorm.py # Key metrics for T4 diffusion kernels: # - Memory throughput (dram__throughput.avg.pct_of_peak_sustained_elapsed) # → Target: 40–50% of 320 GB/s (128–160 GB/s) for vectorized kernels # - SM utilization (sm__throughput.avg.pct_of_peak_sustained_elapsed) # → Target high with only 40 SMs # - Achieved occupancy (sm__warps_active.avg.pct_of_peak_sustained_active) # → Max 1024 threads/SM on T4 — block size ≤ 512 for decent occupancy # - Warp stall reasons (smsp__warp_issue_stalled_*.avg.pct_of_peak_sustained_active) ``` ### Common T4 Bottlenecks 1. **Memory Bandwidth** — 320 GB/s is the primary limit; if `dram__throughput` < 40% → use `AlignedVector` 2. **Limited Memory** — 16 GB; enable `--dit-cpu-offload`/`--vae-cpu-offload` as needed 3. **No BF16** — guard in Python wrapper; FP16 overflow risk in long-sequence attention 4. **Smaller tiles** — 64 KB shared memory; reduce `BLOCK_SIZE_M/N` vs H100 --- ## Best Practices Summary (T4) 1. **No BF16**: Guard in Python wrapper, raise clear error 2. **Vectorization**: Even more critical at 320 GB/s — always use `AlignedVector` 3. **Tile sizes**: 64 KB shared memory limit → halve BLOCK_SIZE vs H100 4. **Block size**: Max 512 threads/block for decent occupancy (max 1024 threads/SM) 5. **Grid sizing**: 40 SMs — `runtime::get_sm_count()` auto-handles 6. **FP32 accumulation**: Always accumulate in FP32 to avoid FP16 overflow 7. **Memory**: Plan for 16 GB; use `--dit-cpu-offload`/`--vae-cpu-offload` as needed 8. **Fuse more**: Low bandwidth makes kernel fusion more impactful than on H100 9. **Multi-arch build**: Always build for `sm_75,sm_80,sm_90` together ## T4 Cloud Instance Quick Reference | Provider | Instance | Notes | |----------|----------|-------| | GCP | n1-standard-4 + T4 | Most common inference setup | | AWS | g4dn.xlarge | 1× T4, 16 GB | | AWS | g4dn.12xlarge | 4× T4, 64 GB total | | Azure | NC4as T4 v3 | 1× T4 | ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/references/troubleshooting.md ================================================ # Troubleshooting Guide — SGLang Diffusion JIT CUDA Kernels Common issues and solutions when writing and integrating JIT CUDA kernels for SGLang Diffusion. > **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels) --- ## Build / Compile Issues ### 1. JIT compilation fails: "No such file or directory" **Problem:** `load_jit` cannot find your `.cuh` file. ``` FileNotFoundError: .../jit_kernel/csrc/diffusion/your_op.cuh not found ``` **Fix:** Ensure the file is under `python/sglang/jit_kernel/csrc/diffusion/`. The path passed to `cuda_files` is relative to `csrc/`: ```python # CORRECT — file lives at csrc/diffusion/your_op.cuh load_jit(..., cuda_files=["diffusion/your_op.cuh"]) # resolves to: python/sglang/jit_kernel/csrc/diffusion/your_op.cuh # ALSO CORRECT — absolute path (pathlib replaces the csrc/ prefix) load_jit(..., cuda_files=["/full/absolute/path/to/your_op.cuh"]) ``` ### 2. Type conversion errors (FP16/BF16) **Problem:** Implicit FP16/BF16 conversion fails because PyTorch compiles with `-D__CUDA_NO_HALF_OPERATORS__`: ``` error: no suitable conversion function from "__half" to "float" exists ``` **Fix:** SGLang's `static_cast` works because `fp16_t` and `bf16_t` are typedef'd with proper conversion operators. Always use explicit casts: ```cpp // CORRECT — explicit cast float val = static_cast(v[i]); // fp16_t / bf16_t → float v[i] = static_cast(fp32_result); // float → T // WRONG — implicit conversion (disabled by PyTorch build flags) float val = v[i]; // compile error v[i] = fp32_result; // compile error ``` If you need the raw intrinsics for packed types: ```cpp // bf16x2_t → two floats bf16x2_t packed = ...; float v0 = __bfloat162float(packed.x); float v1 = __bfloat162float(packed.y); ``` ### 3. Template instantiation explodes / slow first compile **Problem:** Many template combinations makes the first JIT compile very slow. **Fix:** Reduce template argument combinations. Move compile-time constants to runtime if they don't affect performance critically: ```cpp // Fewer template args = fewer instantiations template // only dtype varies void my_op(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, int block_size); ``` ### 4. SM check: kernel requires sm_90 but device is sm_80 **Problem:** Kernel uses H100-only features on A100. **Fix:** Add a Python guard before calling `load_jit`: ```python cap = torch.cuda.get_device_capability() if cap[0] < 9: raise RuntimeError( f"This kernel requires H100 (sm_90+). " f"Got compute capability {cap[0]}.{cap[1]}. " f"Use the Triton fallback instead: diffusion_triton_()" ) ``` --- ## Performance Issues ### 5. Kernel is slower than Triton / PyTorch baseline **Steps to diagnose:** 1. Check dtype: are you using `bf16_t` on T4? (T4 has no BF16 — silently falls back to slow emulation) 2. Check vectorization: is `hidden_size` divisible by `kVecN = 16/sizeof(T)` (8 for bf16, 4 for fp32)? 3. Profile with `ncu`: ```bash ncu --set full --csv -o metrics.csv \ python -c "from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm; ..." ``` Look at `dram__throughput.avg.pct_of_peak_sustained_elapsed` — if < 30%, check coalescing. 4. Check occupancy: run with `--ptxas-options=-v` in `extra_cuda_cflags` to see register usage. ### 6. Shared memory bank conflicts **Problem:** `ncu` reports high `l1tex__data_bank_conflicts_pipe_lmem_op_st.sum`. **Fix:** Add padding to shared memory arrays: ```cpp // Conflict (all threads hit same bank when stride=32) __shared__ float data[32][32]; // Fixed with padding __shared__ float data[32][33]; // 33 instead of 32 ``` ### 7. Low occupancy from too many registers **Problem:** `nvcc --ptxas-options=-v` shows high register count; occupancy < 25%. **Fix:** Add `--maxrregcount=N` to limit registers: ```python extra_cuda_cflags=["-O3", "--use_fast_math", "--maxrregcount=64"] ``` Reduces registers per thread at the cost of possible register spilling to local memory. --- ## Integration Issues ### 8. RMSNorm weight is None (`elementwise_affine=False`) **Problem:** ``` AttributeError: 'NoneType' object has no attribute 'data_ptr' ``` **Root Cause:** DiT transformer blocks often use `RMSNorm(dim, elementwise_affine=False)` — no learnable weight. **Fix in Python wrapper:** pass an empty tensor when weight is absent; the kernel launcher checks `data_ptr == nullptr`: ```python w = weight if weight is not None else torch.empty(0, dtype=src.dtype, device=src.device) module.rmsnorm(out, src, w, eps) ``` **Fix in `.cuh` launcher:** ```cpp const T* w_ptr = (weight.data_ptr() != nullptr) ? static_cast(weight.data_ptr()) : nullptr; // ... pass w_ptr to kernel ... ``` **Fix in module patching:** ```python has_weight = hasattr(module, "weight") and module.weight is not None if has_weight: def _fwd(mod, eps): def forward(x): return diffusion_rmsnorm(x, weight=mod.weight, eps=eps) return forward module.forward = _fwd(module, module.eps) else: def _fwd_noweight(eps): def forward(x): return diffusion_rmsnorm(x, weight=None, eps=eps) return forward module.forward = _fwd_noweight(module.eps) ``` ### 9. `isinstance(module, torch.nn.RMSNorm)` misses diffusion variants **Problem:** Patching doesn't apply because diffusers / sglang diffusion models define their own `RMSNorm` class that is **not** a subclass of `torch.nn.RMSNorm`. **Fix:** Match by class name string: ```python # WRONG — misses diffusers/sglang RMSNorm if isinstance(module, torch.nn.RMSNorm): # CORRECT — catches all variants if type(module).__name__ == "RMSNorm": # or for broader matching: if "RMSNorm" in type(module).__name__: ``` ### 10. Kernel patching doesn't persist after CPU offloading **Problem:** After calling `pipe.enable_model_cpu_offload()`, patched modules revert. **Fix:** Always inject **after** moving to CUDA, **before** enabling any offloading: ```python pipe = load_pipeline(...) pipe.to("cuda") # 1. Move to CUDA inject_optimized_kernels(pipe) # 2. Patch modules pipe.enable_model_cpu_offload() # 3. Now safe to enable offloading ``` ### 11. Kernel patched after `torch.compile` **Problem:** Module is already compiled; patching its `forward` after compilation has no effect. **Fix:** Apply patches **before** any `torch.compile` call: ```python inject_optimized_kernels(pipe) # FIRST: patch pipe.transformer = torch.compile(...) # SECOND: compile ``` --- ## `torch.compile` Compatibility ### 12. Custom CUDA kernel causes graph break **Problem:** ``` torch._dynamo.exc.Unsupported: Attempted to call function marked as skipped ``` or: ``` torch._dynamo.exc.TorchRuntimeError: Cannot access data pointer of Tensor (FakeTensor) ``` **Root Cause:** `torch.compile` traces with "fake tensors" that have no real data. Any kernel that calls `.data_ptr()` during tracing fails. **Options:** **Option A (simplest):** Don't use `torch.compile` with CUDA JIT kernels — use Triton instead: ```python # Triton kernels are torch.compile compatible from sglang.jit_kernel.diffusion.triton.norm import fused_rmsnorm ``` **Option B:** Register as a `@torch.library.custom_op` (advanced): ```python import torch @torch.library.custom_op("diffusion_jit::rmsnorm", mutates_args={"out"}) def _rmsnorm_op(out: torch.Tensor, src: torch.Tensor, weight: torch.Tensor, eps: float) -> None: module = _jit_rmsnorm_module(src.dtype) module.rmsnorm(out, src, weight, eps) @_rmsnorm_op.register_fake def _(out, src, weight, eps): pass # no shape changes; output already allocated in 'out' ``` **Performance trade-off:** | Approach | Speedup (denoise) | torch.compile | Notes | |----------|-------------------|---------------|-------| | CUDA JIT kernel | best | Yes (via `torch.library.custom_op`) | Performance-optimal regardless of whether `torch.compile` is enabled; use `custom_op` + `register_fake` for compile compatibility | | Triton kernel | good | Yes | Use when you need faster iteration/portability, or when you do not have a well-tuned CUDA kernel yet | | Triton + compile | good | Yes | Use for end-to-end `torch.compile` integration convenience; typically slower than a well-tuned CUDA kernel | ### 13. Unstable benchmark results from JIT timing **Problem:** First few runs are slow due to JIT compilation; timing is noisy. **Fix:** Use `triton.testing.do_bench` / `run_benchmark` which use CUDA-graph-based timing automatically. Always do a warmup run first: ```python # Pre-compile by running once before timing diffusion_rmsnorm(dummy_src, weight=dummy_w, eps=1e-6) torch.cuda.synchronize() # Now time result = run_benchmark(lambda: diffusion_rmsnorm(src, weight=w, eps=1e-6)) ``` --- ## Debugging Checklist ```bash # 1. Verify CUDA device and compute capability python -c "import torch; print(torch.cuda.get_device_name(), torch.cuda.get_device_capability())" # 2. Force synchronous CUDA execution to get real error location CUDA_LAUNCH_BLOCKING=1 python scripts/bench_diffusion_rmsnorm.py # 3. Run memory sanitizer to catch illegal accesses compute-sanitizer --tool memcheck python scripts/bench_diffusion_rmsnorm.py # 4. Check register and shared memory usage # Add to extra_cuda_cflags: "--ptxas-options=-v" # 5a. Kernel-level profiling — full metrics ncu --set full -o metrics.ncu-rep \ python scripts/bench_diffusion_rmsnorm.py # 5b. Kernel-level profiling — targeted bandwidth + occupancy check ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\ dram__throughput.avg.pct_of_peak_sustained_elapsed \ python scripts/bench_diffusion_rmsnorm.py # Key metrics to interpret: # - sm__throughput : compute utilization % of peak # - dram__throughput: memory bandwidth % of peak (target ≥ 30% on H100/A100) # - smsp__warp_issue_stalled_*: warp stall breakdown (memory_dependency / math_pipe) # 6. System-level profiling (per-op breakdown inside sglang generate) nsys profile -o denoise_profile \ sglang generate --model-path=black-forest-labs/FLUX.1-dev \ --width=1024 --height=1024 --num-inference-steps=50 \ --seed=42 --enable-torch-compile --warmup # 7. Verify a patched module produces correct output python - << 'EOF' import torch from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm x = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") w = torch.ones(2048, dtype=torch.bfloat16, device="cuda") out_jit = diffusion_rmsnorm(x, weight=w, eps=1e-6) out_ref = torch.nn.functional.rms_norm(x.float(), (2048,), w.float(), eps=1e-6).to(torch.bfloat16) max_diff = (out_jit - out_ref).abs().max().item() print(f"Max diff: {max_diff:.2e} ({'PASS' if max_diff < 0.02 else 'FAIL'})") EOF ``` ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py ================================================ """ End-to-end denoise-stage benchmark for SGLang Diffusion with/without custom JIT CUDA kernels. Measures denoise latency (primary metric ★) and peak GPU memory. All model configs are kept in exact sync with diffusion-benchmark-and-profile.md. Adapted from: https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels Usage: # Baseline — single model cd /path/to/sglang python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux # With custom JIT CUDA kernels python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux --custom-kernels # Side-by-side comparison python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux --compare # All 10 models, comparison python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --all --compare Input images required for image-guided models: ASSET_DIR=$(python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/diffusion_skill_env.py print-assets-dir --mkdir) wget -O "${ASSET_DIR}/cat.png" \ https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png wget -O "${ASSET_DIR}/astronaut.jpg" \ https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg wget -O "${ASSET_DIR}/mova_single_person.jpg" \ https://github.com/OpenMOSS/MOVA/raw/main/assets/single_person.jpg """ import argparse import json import os import subprocess import sys import time from pathlib import Path from typing import Optional SCRIPT_DIR = Path(__file__).resolve().parent if str(SCRIPT_DIR) not in sys.path: sys.path.insert(0, str(SCRIPT_DIR)) from diffusion_skill_env import ( ensure_dir, get_assets_dir, get_output_dir, get_repo_root, pick_idle_gpus, ) REPO_ROOT = get_repo_root() ASSET_DIR = ensure_dir(get_assets_dir(REPO_ROOT)) # --------------------------------------------------------------------------- # Model configs — kept in exact sync with diffusion-benchmark-and-profile.md # Each entry produces the same `sglang generate` command as shown in that doc. # --------------------------------------------------------------------------- MODELS = { # 1. Qwen/Qwen-Image-2512 — Text-to-Image, 1024×1024, 50 steps "qwen": { "path": "Qwen/Qwen-Image-2512", "prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k", "negative_prompt": " ", "extra_args": [ "--width=1024", "--height=1024", "--num-inference-steps=50", "--guidance-scale=4.0", "--dit-cpu-offload", "false", "--text-encoder-cpu-offload", "false", ], }, # 2. Qwen/Qwen-Image-Edit-2511 — Image Editing, 1024×1024, 50 steps # Requires: /inputs/diffusion_benchmark/figs/cat.png "qwen-edit": { "path": "Qwen/Qwen-Image-Edit-2511", "prompt": "Transform into anime style", "negative_prompt": " ", "image_path": str(ASSET_DIR / "cat.png"), "extra_args": [ "--width=1024", "--height=1024", "--num-inference-steps=50", "--guidance-scale=4.0", "--dit-cpu-offload", "false", "--text-encoder-cpu-offload", "false", ], }, # 3. black-forest-labs/FLUX.1-dev — Text-to-Image, 1024×1024, 50 steps "flux": { "path": "black-forest-labs/FLUX.1-dev", "prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k", "extra_args": [ "--width=1024", "--height=1024", "--num-inference-steps=50", "--guidance-scale=4.0", ], }, # 4. black-forest-labs/FLUX.2-dev — Text-to-Image, 1024×1024 "flux2": { "path": "black-forest-labs/FLUX.2-dev", "prompt": "A Logo With Bold Large Text: SGL Diffusion", "extra_args": [ "--width=1024", "--height=1024", "--dit-layerwise-offload", "false", "--dit-cpu-offload", "false", "--text-encoder-cpu-offload", "true", "--vae-cpu-offload", "false", ], }, # 5. Tongyi-MAI/Z-Image-Turbo — Turbo Text-to-Image, 1024×1024, 9 steps "zimage": { "path": "Tongyi-MAI/Z-Image-Turbo", "prompt": "A fantasy landscape with mountains and a river, detailed, vibrant colors", "extra_args": [ "--width=1024", "--height=1024", "--num-inference-steps=9", "--guidance-scale=0.0", "--dit-cpu-offload", "false", "--text-encoder-cpu-offload", "false", ], }, # 6. Wan-AI/Wan2.2-T2V-A14B-Diffusers — Text-to-Video, 720P, 4 GPUs, 81 frames, 2 steps "wan-t2v": { "path": "Wan-AI/Wan2.2-T2V-A14B-Diffusers", "prompt": "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon.", "negative_prompt": " ", "extra_args": [ "--720p", "--num-inference-steps=2", "--num-frames=81", "--guidance-scale=5.0", "--num-gpus=4", "--ulysses-degree=4", "--text-encoder-cpu-offload", "--pin-cpu-memory", ], }, # 7. Wan-AI/Wan2.2-TI2V-5B-Diffusers — Text-Image-to-Video, 720P, 1 GPU, 81 frames, 50 steps # Requires: /inputs/diffusion_benchmark/figs/astronaut.jpg "wan-ti2v": { "path": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", "prompt": "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.", "negative_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards", "image_path": str(ASSET_DIR / "astronaut.jpg"), "extra_args": [ "--num-frames", "81", "--720p", "--num-inference-steps", "50", "--guidance-scale", "5.0", "--dit-layerwise-offload", "false", "--dit-cpu-offload", "false", "--vae-cpu-offload", "false", "--text-encoder-cpu-offload", "false", ], }, # 8. hunyuanvideo-community/HunyuanVideo — Text-to-Video, 848×480, 65 frames, 30 steps "hunyuanvideo": { "path": "hunyuanvideo-community/HunyuanVideo", "prompt": "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window.", "extra_args": [ "--text-encoder-cpu-offload", "--pin-cpu-memory", "--num-frames=65", "--width=848", "--height=480", "--num-inference-steps=30", ], }, # 9. OpenMOSS-Team/MOVA-720p — Image-to-Video, 4 GPUs, 193 frames, 2 steps # Requires: /inputs/diffusion_benchmark/figs/mova_single_person.jpg "mova-720p": { "path": "OpenMOSS-Team/MOVA-720p", "prompt": 'A man in a blue blazer and glasses speaks in a formal indoor setting, framed by wooden furniture and a filled bookshelf. Quiet room acoustics underscore his measured tone as he delivers his remarks. At one point, he says, "I would also believe that this advance in AI recently was not unexpected."', "image_path": str(ASSET_DIR / "mova_single_person.jpg"), "extra_args": [ "--adjust-frames=false", "--num-gpus=4", "--ring-degree=1", "--ulysses-degree=4", "--num-frames=193", "--fps=24", "--num-inference-steps=2", ], }, # 10. BestWishYsh/Helios-Base — Text-to-Video, 640×384, 33 frames "helios": { "path": "BestWishYsh/Helios-Base", "prompt": "A curious raccoon", "extra_args": [ "--width=640", "--height=384", "--num-frames=33", "--dit-layerwise-offload", "false", "--dit-cpu-offload", "false", "--text-encoder-cpu-offload", "false", "--vae-cpu-offload", "false", ], }, } def required_gpus_for_model(model_key: str) -> int: if model_key == "wan-t2v": return 4 if model_key == "mova-720p": return 4 return 1 def build_sglang_cmd( model_key: str, use_custom_kernels: bool, perf_dump_path: Optional[str] = None, warmup: bool = True, torch_compile: bool = True, seed: int = 42, save_output: bool = True, ) -> list[str]: """ Build the `sglang generate` command for the given model. Matches the commands in diffusion-benchmark-and-profile.md exactly. """ cfg = MODELS[model_key] cmd = [ "sglang", "generate", f"--model-path={cfg['path']}", f"--prompt={cfg['prompt']}", "--log-level=info", ] if seed is not None: cmd.append(f"--seed={seed}") if "negative_prompt" in cfg: cmd.append(f"--negative-prompt={cfg['negative_prompt']}") if "image_path" in cfg: cmd.append(f"--image-path={cfg['image_path']}") cmd.extend(cfg["extra_args"]) if save_output: cmd.append("--save-output") if warmup: cmd.append("--warmup") if torch_compile: cmd.append("--enable-torch-compile") if perf_dump_path: cmd.extend(["--perf-dump-path", perf_dump_path]) return cmd def run_benchmark_once( model_key: str, use_custom_kernels: bool, output_dir: Path, warmup: bool = True, ) -> dict: """Run a single benchmark pass and return results dict.""" label = "custom" if use_custom_kernels else "baseline" perf_path = output_dir / f"{model_key}_{label}.json" cmd = build_sglang_cmd( model_key, use_custom_kernels=use_custom_kernels, perf_dump_path=str(perf_path), warmup=warmup, ) env = os.environ.copy() env.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1") if not env.get("CUDA_VISIBLE_DEVICES"): env["CUDA_VISIBLE_DEVICES"] = ",".join( str(index) for index in pick_idle_gpus(required_gpus_for_model(model_key)) ) if use_custom_kernels: # NOTE: This env var is a convention for user-implemented kernel injection # logic. SGLang runtime does not read it by default — you must add handling # in your denoising stage or model code to check this var and apply patches. env["SGLANG_DIFFUSION_CUSTOM_CUDA_KERNELS"] = "1" print(f"\n{'=' * 64}") print(f"[{label.upper()}] {model_key}") print(f" CUDA_VISIBLE_DEVICES={env.get('CUDA_VISIBLE_DEVICES', '')}") print(" " + " \\\n ".join(cmd)) print() t0 = time.time() result = subprocess.run(cmd, env=env, text=True) elapsed = time.time() - t0 if result.returncode != 0: print(f" ERROR: exit code {result.returncode}") return {"model": model_key, "label": label, "error": True, "elapsed_s": elapsed} metrics = {"model": model_key, "label": label, "elapsed_s": elapsed, "error": False} if perf_path.exists(): try: with open(perf_path) as f: perf = json.load(f) # e2e latency: total_duration_ms (set by PerformanceLogger.dump_benchmark_report) total_ms = perf.get("total_duration_ms") metrics["e2e_latency_s"] = ( float(total_ms) / 1000.0 if total_ms is not None else None ) # denoise latency: look in "steps" list for the "DenoisingStage" entry # steps = [{"name": "DenoisingStage", "duration_ms": 1234.5}, ...] denoise_latency_s = None for step in perf.get("steps", []): if ( step.get("name") == "DenoisingStage" and step.get("duration_ms") is not None ): denoise_latency_s = float(step["duration_ms"]) / 1000.0 break # fallback: sum all per-step durations from denoise_steps_ms # denoise_steps_ms = [{"step": 0, "duration_ms": 100.5}, ...] if denoise_latency_s is None: denoise_steps = perf.get("denoise_steps_ms", []) if denoise_steps: denoise_latency_s = ( sum(s.get("duration_ms", 0.0) for s in denoise_steps) / 1000.0 ) metrics["denoise_latency_s"] = denoise_latency_s # peak memory: max peak_reserved_mb across all memory checkpoints (→ GB) # memory_checkpoints = {"after_DenoisingStage": {"peak_reserved_mb": 12288.0, ...}} peak_memory_gb = None for snapshot in perf.get("memory_checkpoints", {}).values(): peak_mb = snapshot.get("peak_reserved_mb") if peak_mb is not None: candidate = float(peak_mb) / 1024.0 if peak_memory_gb is None or candidate > peak_memory_gb: peak_memory_gb = candidate metrics["peak_memory_gb"] = peak_memory_gb except Exception as e: print(f" Warning: could not parse perf dump: {e}") return metrics def print_results_table(results: list[dict]): """Print baseline vs custom kernel comparison table.""" print() print("=" * 80) print("BENCHMARK RESULTS — Denoise Latency (primary metric ★)") print("(Models and params match diffusion-benchmark-and-profile.md)") print("=" * 80) by_model: dict[str, dict] = {} for r in results: by_model.setdefault(r["model"], {})[r["label"]] = r print( f"{'Model':<16} {'Baseline(s)':>12} {'Custom(s)':>10} {'Speedup':>9} {'Peak Mem(GB)':>14}" ) print("-" * 64) for model_key in MODELS: # preserve order if model_key not in by_model: continue runs = by_model[model_key] base = runs.get("baseline", {}) custom = runs.get("custom", {}) base_lat = base.get("denoise_latency_s") custom_lat = custom.get("denoise_latency_s") peak_mem = base.get("peak_memory_gb") or custom.get("peak_memory_gb") speedup = f"{base_lat / custom_lat:.2f}x" if base_lat and custom_lat else "n/a" base_s = f"{base_lat:.2f}" if base_lat else "n/a" custom_s = f"{custom_lat:.2f}" if custom_lat else "n/a" mem_s = f"{peak_mem:.1f}" if isinstance(peak_mem, float) else "n/a" print(f"{model_key:<16} {base_s:>12} {custom_s:>10} {speedup:>9} {mem_s:>14}") print("-" * 64) print() print("★ Denoise latency = total DiT forward pass time across all inference steps.") print( " See diffusion-benchmark-and-profile.md for full Level 1/2 profiling workflow." ) def inject_kernels_example(): """ Show the kernel injection pattern used when SGLANG_DIFFUSION_CUSTOM_CUDA_KERNELS=1. After implementing add-cuda-kernel.md, this logic lives in denoising.py or the model's transformer.py — NOT in this script. Call patch_rmsnorm(dit_model) BEFORE torch.compile and BEFORE any CPU offloading. """ import torch.nn as nn try: from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm except ImportError: print( "diffusion.rmsnorm JIT kernel not available. " "Implement add-cuda-kernel.md first." ) return def patch_rmsnorm(model: nn.Module, verbose: bool = False) -> int: """Monkey-patch all RMSNorm variants to use the JIT CUDA kernel.""" patched = 0 for name, module in model.named_modules(): if "RMSNorm" not in type(module).__name__: continue eps = getattr(module, "eps", getattr(module, "variance_epsilon", 1e-6)) has_weight = hasattr(module, "weight") and module.weight is not None if has_weight: def _make(mod, ep): def fwd(x): return diffusion_rmsnorm(x, weight=mod.weight, eps=ep) return fwd module.forward = _make(module, eps) else: def _make_no_w(ep): def fwd(x): return diffusion_rmsnorm(x, weight=None, eps=ep) return fwd module.forward = _make_no_w(eps) patched += 1 if verbose: print(f" Patched: {name} (weight={has_weight})") return patched return patch_rmsnorm def main(): parser = argparse.ArgumentParser( description="SGLang Diffusion denoise benchmark — baseline vs JIT CUDA kernels" ) parser.add_argument( "--model", choices=list(MODELS.keys()), help="Model to benchmark (default: flux)", ) parser.add_argument("--all", action="store_true", help="Benchmark all 7 models") parser.add_argument( "--custom-kernels", action="store_true", help="Run with custom JIT CUDA kernels (SGLANG_DIFFUSION_CUSTOM_CUDA_KERNELS=1)", ) parser.add_argument( "--no-custom-kernels", action="store_true", help="Run baseline (no custom kernels)", ) parser.add_argument( "--compare", action="store_true", help="Run both baseline and custom, print comparison table", ) parser.add_argument( "--output-dir", type=str, default=str(get_output_dir("benchmarks", REPO_ROOT)), help="Directory for perf dump JSON files", ) parser.add_argument("--no-warmup", action="store_true", help="Skip warmup") parser.add_argument( "--show-injection-example", action="store_true", help="Print kernel injection pattern and exit", ) args = parser.parse_args() if args.show_injection_example: patch_fn = inject_kernels_example() if patch_fn: print( "patch_rmsnorm function defined. " "Call it on the DiT model before torch.compile and CPU offloading." ) return output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) warmup = not args.no_warmup models_to_run = list(MODELS.keys()) if args.all else [args.model or "flux"] results = [] for model_key in models_to_run: if args.compare: results.append(run_benchmark_once(model_key, False, output_dir, warmup)) results.append(run_benchmark_once(model_key, True, output_dir, warmup)) elif args.custom_kernels: results.append(run_benchmark_once(model_key, True, output_dir, warmup)) else: results.append(run_benchmark_once(model_key, False, output_dir, warmup)) if results: print_results_table(results) print(f"Perf dump JSONs → {output_dir}") print( "Compare across runs: follow diffusion-benchmark-and-profile.md → Perf dump & before/after compare." ) if __name__ == "__main__": main() ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_rmsnorm.py ================================================ """ Micro-benchmark for the SGLang Diffusion JIT CUDA RMSNorm kernel. Compares: 1. SGLang JIT CUDA kernel (diffusion_rmsnorm) 2. PyTorch baseline (torch.nn.functional.rms_norm) Adapted from: https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels Usage: cd /path/to/sglang python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_rmsnorm.py Requirements: # Run inside the configured SGLang diffusion container shell. # This script auto-selects an idle GPU when CUDA_VISIBLE_DEVICES is unset. """ import os import sys import time from pathlib import Path from typing import Tuple SCRIPT_DIR = Path(__file__).resolve().parent if str(SCRIPT_DIR) not in sys.path: sys.path.insert(0, str(SCRIPT_DIR)) from diffusion_skill_env import configure_runtime_env configure_runtime_env(required_gpus=1) import torch # --------------------------------------------------------------------------- # Import the JIT CUDA kernel. # When you implement add-cuda-kernel.md, the file will be at: # python/sglang/jit_kernel/diffusion/rmsnorm.py # --------------------------------------------------------------------------- try: from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm JIT_AVAILABLE = True except ImportError: JIT_AVAILABLE = False print( "WARNING: diffusion.rmsnorm JIT kernel not available. " "Run after implementing add-cuda-kernel.md." ) def pytorch_rmsnorm( x: torch.Tensor, weight: torch.Tensor | None = None, eps: float = 1e-6, ) -> torch.Tensor: """Reference PyTorch implementation of RMSNorm.""" hidden = x.shape[-1] return torch.nn.functional.rms_norm( x.float(), (hidden,), weight.float() if weight is not None else None, eps=eps ).to(x.dtype) def benchmark_kernel( func, args, warmup: int = 20, iterations: int = 100, ) -> Tuple[float, float]: """Benchmark a kernel function. Returns (avg_ms, min_ms).""" for _ in range(warmup): func(*args) torch.cuda.synchronize() times = [] for _ in range(iterations): torch.cuda.synchronize() t0 = time.perf_counter() func(*args) torch.cuda.synchronize() times.append((time.perf_counter() - t0) * 1000) return sum(times) / len(times), min(times) def run_benchmark(): print("=" * 72) print("SGLang Diffusion RMSNorm Micro-Benchmark: JIT CUDA vs PyTorch") print("=" * 72) print(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', '')}") print(f"Device: {torch.cuda.get_device_name(0)}") cap = torch.cuda.get_device_capability() print(f"Compute Capability: sm_{cap[0]}{cap[1]}") print() if not JIT_AVAILABLE: print("Skipping JIT kernel benchmark (kernel not available).") return # Determine dtype: T4 (sm_75) has no BF16 dtype = torch.bfloat16 if cap >= (8, 0) else torch.float16 print(f"Dtype: {dtype}") print() # Typical DiT hidden sizes for sglang diffusion models: # FLUX.1-dev: hidden=3072 # Qwen-Image: hidden=2048 # Wan2.2: hidden=4096 configs = [ # (batch_tokens, hidden_size, has_weight) (1024, 2048, True), # Qwen-Image: 1 sample × 1024 tokens (4096, 2048, True), # Qwen-Image: larger batch (1024, 3072, True), # FLUX: 1 sample × 1024 tokens (4096, 3072, True), # FLUX: larger (4096, 4096, True), # Wan2.2 (4096, 2048, False), # no-weight (elementwise_affine=False) (16384, 3072, True), # long sequence ] print( f"{'Config':<32} {'JIT(ms)':>10} {'PyTorch(ms)':>12} {'Speedup':>9} {'Weight'}" ) print("-" * 72) total_speedup = 0 n = 0 for batch_tokens, hidden, has_weight in configs: x = torch.randn(batch_tokens, hidden, dtype=dtype, device="cuda") weight = torch.ones(hidden, dtype=dtype, device="cuda") if has_weight else None jit_avg, _ = benchmark_kernel( diffusion_rmsnorm, (x, weight, 1e-6), warmup=20, iterations=100 ) pt_avg, _ = benchmark_kernel( pytorch_rmsnorm, (x, weight, 1e-6), warmup=20, iterations=100 ) speedup = pt_avg / jit_avg total_speedup += speedup n += 1 w_str = "yes" if has_weight else "no " cfg = f"[{batch_tokens}×{hidden}]" print(f"{cfg:<32} {jit_avg:>10.3f} {pt_avg:>12.3f} {speedup:>8.2f}x {w_str}") print("-" * 72) print(f"{'Average Speedup':>56} {total_speedup / n:.2f}x") print() # ----------------------------------------------------------------------- # Correctness check # ----------------------------------------------------------------------- print("Correctness Check (BF16 tolerance 0.02):") x = torch.randn(4096, 3072, dtype=dtype, device="cuda") weight = torch.ones(3072, dtype=dtype, device="cuda") out_jit = diffusion_rmsnorm(x, weight=weight, eps=1e-6) out_ref = pytorch_rmsnorm(x, weight=weight, eps=1e-6) max_diff = (out_jit - out_ref).abs().max().item() rel_diff = ((out_jit - out_ref).abs() / (out_ref.abs() + 1e-8)).max().item() passed = max_diff < 0.02 print(f" Max absolute diff: {max_diff:.2e}") print(f" Max relative diff: {rel_diff:.2e}") print(f" Correctness: {'PASS ✓' if passed else 'FAIL ✗'}") print() # ----------------------------------------------------------------------- # Memory bandwidth analysis # ----------------------------------------------------------------------- print("Memory Bandwidth Analysis:") bt, hid = 4096, 3072 x = torch.randn(bt, hid, dtype=dtype, device="cuda") weight = torch.ones(hid, dtype=dtype, device="cuda") bytes_per_elem = dtype.itemsize total_bytes = ( bt * hid + hid + bt * hid ) * bytes_per_elem # read x + read w + write out jit_avg, _ = benchmark_kernel(diffusion_rmsnorm, (x, weight, 1e-6)) bandwidth_gbps = (total_bytes / 1e9) / (jit_avg / 1000) theoretical_bw = { (9, 0): 3350, # H100: 3.35 TB/s (8, 0): 2000, # A100 80GB }.get( cap, 320 ) # T4: 320 GB/s efficiency = bandwidth_gbps / theoretical_bw * 100 print(f" Shape: [{bt} × {hid}] dtype: {dtype}") print(f" Total data: {total_bytes / 1e6:.1f} MB") print(f" Achieved: {bandwidth_gbps:.1f} GB/s") print(f" Theoretical ({torch.cuda.get_device_name(0)}): {theoretical_bw} GB/s") print(f" Bandwidth efficiency: {efficiency:.1f}%") print() print("Target: ≥ 30% efficiency (H100/A100), ≥ 40% (T4)") if __name__ == "__main__": if not torch.cuda.is_available(): print("CUDA not available.") else: run_benchmark() ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/diffusion_skill_env.py ================================================ from __future__ import annotations import argparse import csv import os import subprocess from pathlib import Path OUTPUT_DIR_NAMES = { "benchmarks": Path("outputs/diffusion_benchmarks"), "profiles": Path("outputs/diffusion_profiles"), "ncu": Path("outputs/ncu_reports"), } def get_repo_root() -> Path: import sglang return Path(sglang.__file__).resolve().parents[2] def get_assets_dir(repo_root: Path | None = None) -> Path: root = repo_root or get_repo_root() return root / "inputs" / "diffusion_benchmark" / "figs" def get_output_dir(name: str, repo_root: Path | None = None) -> Path: if name not in OUTPUT_DIR_NAMES: raise KeyError(f"Unknown output dir name: {name}") root = repo_root or get_repo_root() return root / OUTPUT_DIR_NAMES[name] def ensure_dir(path: Path) -> Path: path.mkdir(parents=True, exist_ok=True) return path def check_write_access(repo_root: Path | None = None) -> Path: root = repo_root or get_repo_root() probe_dir = ensure_dir(root / ".cache" / "diffusion_skill_write_test") probe_file = probe_dir / "probe.txt" probe_file.write_text("ok", encoding="utf-8") return probe_file def _run_nvidia_smi(query: str) -> list[list[str]]: command = [ "nvidia-smi", f"--query-{query}", "--format=csv,noheader,nounits", ] result = subprocess.run(command, check=True, capture_output=True, text=True) rows: list[list[str]] = [] for raw_line in result.stdout.splitlines(): line = raw_line.strip() if not line: continue rows.append([field.strip() for field in csv.reader([line]).__next__()]) return rows def get_gpu_inventory() -> list[dict[str, int | str]]: rows = _run_nvidia_smi("gpu=index,uuid,memory.used,memory.total,utilization.gpu") inventory = [] for index, uuid, memory_used, memory_total, utilization_gpu in rows: inventory.append( { "index": int(index), "uuid": uuid, "memory_used_mib": int(memory_used), "memory_total_mib": int(memory_total), "utilization_gpu_pct": int(utilization_gpu), } ) return inventory def get_busy_gpu_uuids() -> set[str]: rows = _run_nvidia_smi("compute-apps=gpu_uuid,pid,process_name,used_gpu_memory") return {gpu_uuid for gpu_uuid, *_ in rows} def pick_idle_gpus( required_gpus: int, max_memory_used_mib: int = 32, max_utilization_gpu_pct: int = 5, ) -> list[int]: inventory = get_gpu_inventory() busy_uuids = get_busy_gpu_uuids() idle = [ int(gpu["index"]) for gpu in inventory if gpu["uuid"] not in busy_uuids and int(gpu["memory_used_mib"]) <= max_memory_used_mib and int(gpu["utilization_gpu_pct"]) <= max_utilization_gpu_pct ] if len(idle) < required_gpus: raise RuntimeError( "Not enough idle GPUs. " f"required={required_gpus}, idle={idle}, inventory={inventory}, busy={sorted(busy_uuids)}" ) return idle[:required_gpus] def configure_runtime_env(required_gpus: int = 1) -> str | None: os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1") if os.environ.get("CUDA_VISIBLE_DEVICES"): return None selected = ",".join(str(index) for index in pick_idle_gpus(required_gpus)) os.environ["CUDA_VISIBLE_DEVICES"] = selected return selected def main() -> None: parser = argparse.ArgumentParser( description="Resolve SGLang diffusion skill paths and idle GPUs." ) parser.add_argument( "command", choices=[ "print-root", "print-assets-dir", "print-output-dir", "print-idle-gpus", "check-write-access", ], ) parser.add_argument( "--kind", choices=sorted(OUTPUT_DIR_NAMES), help="Output directory kind for print-output-dir.", ) parser.add_argument( "--count", type=int, default=1, help="Number of idle GPUs to print.", ) parser.add_argument( "--mkdir", action="store_true", help="Create the requested directory before printing it.", ) args = parser.parse_args() if args.command == "print-root": print(get_repo_root()) return if args.command == "print-assets-dir": path = get_assets_dir() if args.mkdir: ensure_dir(path) print(path) return if args.command == "print-output-dir": if not args.kind: raise SystemExit("--kind is required for print-output-dir") path = get_output_dir(args.kind) if args.mkdir: ensure_dir(path) print(path) return if args.command == "print-idle-gpus": print(",".join(str(index) for index in pick_idle_gpus(args.count))) return if args.command == "check-write-access": print(check_write_access()) return raise SystemExit(f"Unhandled command: {args.command}") if __name__ == "__main__": main() ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/use-efficient-diffusion-kernels.md ================================================ --- name: use-efficient-diffusion-kernels description: Guidance for using SGLang Diffusion fused kernels and fast CUDA paths. Use when mapping fusion patterns in diffusion inference, choosing fused ops or attention backends, handling RoPE/QK norm performance pitfalls, or integrating new diffusion models with kernel-aware constraints. --- # Use Efficient Diffusion Kernels **Overview** This skill focuses on SGLang Diffusion (`sglang.multimodal_gen`) kernel fusion patterns and fast CUDA paths. Prefer existing fused ops (Triton, CuTe DSL, sgl-kernel). Make constraints and fallbacks explicit. **Key Files** - `python/sglang/multimodal_gen/runtime/layers/layernorm.py` - `python/sglang/multimodal_gen/runtime/layers/elementwise.py` - `python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py` - `python/sglang/jit_kernel/diffusion/triton/scale_shift.py` - `python/sglang/jit_kernel/diffusion/triton/norm.py` - `python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py` - `python/sglang/jit_kernel/diffusion/triton/rotary.py` - `python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py` - `python/sglang/jit_kernel/norm.py` - `python/sglang/multimodal_gen/runtime/platforms/cuda.py` - `python/sglang/multimodal_gen/runtime/layers/attention/selector.py` - `docs/diffusion/performance/attention_backends.md` (repo root) **Core Fusion Patterns** 1. Scale/Shift elementwise fusion (AdaLN modulation) - Kernels: `fuse_scale_shift_kernel`, `fuse_scale_shift_gate_select01_kernel` - Locations: `elementwise.py`, `layernorm.py`, `qwen_image.py`, `triton/scale_shift.py` - Use cases: `x * (1 + scale) + shift` and `a * (k + b) + c` - Constraints: `x` must be CUDA and contiguous. `scale/shift` support 0D/1D/2D/3D/4D broadcast. 4D `[B, F, 1, C]` requires `L % F == 0`. - NPU fallback: `scale_shift.py` swaps to `npu_fallback` native path. 2. Norm + Scale/Shift fusion (CuTe DSL) - Kernels: `fused_norm_scale_shift`, `fused_scale_residual_norm_scale_shift` - Locations: `layernorm.py`, `cutedsl/scale_residual_norm_scale_shift.py` - Use cases: - `y = norm(x) * (1 + scale) + shift` - `y = norm(residual + gate * x) * (1 + scale) + shift` - Constraints: `D % 256 == 0` and `D <= 8192`. `x/residual/gate/scale/shift` must pass shape and stride validation. Dtypes limited to fp16/bf16/fp32. - Behavior: CuTe DSL compilation cached by `(dtype, ndim, D, norm_type)`. `None` tensors replaced by scalar placeholders. If constraints fail, `layernorm.py` warns and falls back to native PyTorch. 3. Triton LayerNorm/RMSNorm fusion - Kernels: `rms_norm_fn`, `layer_norm_fn`, `norm_infer` - Locations: `triton/norm.py`, `layernorm.py` - Use cases: fp32 RMSNorm with residual/dropout/rowscale/x1 branches, and inference-friendly `norm_infer`. - Constraints: last dim must be contiguous, and `N * element_size < 64KB`. 4. Triton one-pass RMSNorm (small hidden size fast path) - Kernel: `triton_one_pass_rms_norm` - Locations: `triton/rmsnorm_onepass.py`, `layernorm.py` - Use case: `hidden_size <= 128` in `RMSNorm.forward_cuda`. 5. Triton RoPE fusion - Kernel: `apply_rotary_embedding` - Locations: `triton/rotary.py`, `rotary_embedding/utils.py` - Use case: GPT-J style RoPE when not Neox. - Constraints: `head_size` must be even. - NPU fallback: `npu_fallback.apply_rotary_embedding_native`. **Faster CUDA Kernel Usage Points** 1. sgl-kernel RMSNorm and fused add RMSNorm - Location: `layernorm.py` - Behavior: CUDA uses `sgl_kernel.fused_add_rmsnorm` and `sgl_kernel.rmsnorm`. `hidden_size <= 128` uses Triton one-pass. ROCm falls back to native. 2. Attention backend selection (FlashAttention, Sage, SDPA) - Locations: `platforms/cuda.py`, `attention/selector.py`, `docs/diffusion/performance/attention_backends.md` - Behavior: CUDA prefers FlashAttention (FA3/FA4) when supported, otherwise Torch SDPA. Force via `--attention-backend` or `global_force_attn_backend`. 3. FlashInfer RoPE (Q/K inplace) - Location: `rotary_embedding/utils.py` - Behavior: `flashinfer.rope.apply_rope_with_cos_sin_cache_inplace` when available, otherwise Triton RoPE fallback. **QK Norm Optimization** - Entry point: `apply_qk_norm` in `layernorm.py`. - Fast path: JIT fused inplace QK norm from `python/sglang/jit_kernel/norm.py` via `fused_inplace_qknorm`. - Preconditions for fused path: - CUDA only. - `allow_inplace=True` and `q_eps == k_eps`. - `can_use_fused_inplace_qknorm(head_dim, dtype)` returns true. - Supported head dims: `64, 128, 256, 512, 1024`. - Behavior: Fused path operates on `q` and `k` in place after reshaping to `[B, -1, head_dim]`. If preconditions fail, fall back to per-tensor RMSNorm. **Common Entry Points in Diffusion Models** - AdaLN modulation: `LayerNormScaleShift`, `RMSNormScaleShift`, `ScaleResidual*` in `layernorm.py`. - Qwen-Image gating: `fuse_scale_shift_gate_select01_kernel` in `qwen_image.py`. - QK norm: `apply_qk_norm` used in `flux.py`, `flux_2.py`, `qwen_image.py`, `zimage.py`, `wanvideo.py`, `ltx_2.py`, `hunyuanvideo.py`. - RoPE: `_apply_rotary_emb` prefers Triton; Q/K RoPE prefers FlashInfer when present. **Constraints and Fallbacks** - `scale_shift` Triton requires CUDA + contiguous `x`. NPU swaps to native. - CuTe DSL fused norms require `D % 256 == 0` and `D <= 8192`. - Triton norm kernels error on feature size >= 64KB. - FlashAttention requires fp16/bf16 and SM80+; otherwise SDPA. **Integration Checklist for New Models** 1. Reuse `LayerNormScaleShift` or `ScaleResidual*` modules instead of re-implementing fusion logic. 2. Keep tensors contiguous and satisfy D alignment (`% 256`) and size (`<= 8192`) for CuTe fused paths. 3. Use `fuse_scale_shift_kernel` for AdaLN modulation and keep a PyTorch fallback. 4. Use `apply_qk_norm` and ensure head_dim is in the supported list for fused QK norm. 5. If using FlashInfer RoPE, avoid `pack qkv` and ensure Q/K are contiguous. 6. For attention, follow `selector.py` priority; override with CLI only if needed. **When Extending or Modifying Kernels** - Add `torch.library.custom_op` and `register_fake` for compile and meta support. - Keep CuTe compile cache keys aligned to `(dtype, ndim, D)`. - Avoid implicit broadcasts that force hidden `contiguous()` copies. - Preserve NPU and ROCm fallback paths. - **Always verify with ncu** (`ncu --set full`) that the kernel achieves adequate memory bandwidth utilization (>70% of peak for bandwidth-bound ops) and occupancy (>50%). See `diffusion-benchmark-and-profile.md` Step 3.5 for the ncu workflow. ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/diffusion-optimal-perf/SKILL.md ================================================ --- name: diffusion-optimal-perf description: Guide for achieving optimal performance with SGLang-Diffusion. Covers all perf-related CLI flags, env vars, and best practices for lossless and lossy speedup. --- # SGLang-Diffusion: Optimal Performance Guide Use this guide when a user asks how to speed up diffusion inference, reduce latency, lower VRAM usage, or tune SGLang-Diffusion for production. Before running any `sglang generate` command below inside the diffusion container: - derive the repo root from `python3 -c "import os, sglang; print(os.path.abspath(os.path.join(os.path.dirname(sglang.__file__), '..', '..')))"` and `cd` there - export `FLASHINFER_DISABLE_VERSION_CHECK=1` - verify the repo is writable if you expect perf dumps or outputs - choose idle GPU(s) first; reuse `diffusion-kernel/scripts/diffusion_skill_env.py` when doing perf work Reference: [SGLang-Diffusion Advanced Optimizations Blog](https://lmsys.org/blog/2026-02-16-sglang-diffusion-advanced-optimizations/) --- ## Section 1: Lossless Optimizations These options **do not** affect output quality. The generated images/videos are numerically identical (or within floating-point rounding) to the baseline. | Option | CLI Flag / Env Var | What It Does | Speedup | Limitations / Notes | |---|---|---|---|---| | **torch.compile** | `--enable-torch-compile` | Applies `torch.compile` to the DiT forward pass, fusing ops and reducing kernel launch overhead. | ~1.2–1.5x on denoising | First request is slow (compilation). May cause minor precision drifts due to [PyTorch issue #145213](https://github.com/pytorch/pytorch/issues/145213). Pair with `--warmup` for best results. | | **Warmup** | `--warmup` | Runs dummy forward passes to warm up CUDA caches, JIT, and `torch.compile`. Eliminates cold-start penalty. | Removes first-request latency spike | Adds startup time. Without `--warmup-resolutions`, warmup happens on first request. | | **Warmup Resolutions** | `--warmup-resolutions 256x256 720x720` | Pre-compiles and warms up specific resolutions at server startup (instead of lazily on first request). | Faster first request per resolution | Each resolution adds to startup time. Serving mode only; useful when you know your target resolutions in advance. | | **Multi-GPU (SP)** | `--num-gpus N --ulysses-degree N` | Sequence parallelism across GPUs. Shards sequence tokens (not frames) to minimize padding. | Near-linear scaling with N GPUs | Requires NCCL; inter-GPU bandwidth matters. `ulysses_degree * ring_degree = sp_degree`. | | **CFG Parallel** | `--enable-cfg-parallel` | Runs conditional and unconditional CFG branches in parallel across GPUs. **For CFG models with multi-GPU, always prefer `--enable-cfg-parallel` + Ulysses over pure Ulysses** — it is generally faster at the same GPU count due to better compute-to-communication ratio and elimination of sequential branch execution. | Typically faster than pure SP for CFG models | Requires `num_gpus >= 2`. Halves the Ulysses group size (e.g. 8 GPU → two 4-GPU groups). Only for models that use CFG. | | **Layerwise Offload** | `--dit-layerwise-offload` | Async layer-by-layer H2D prefetch with compute overlap. Only ~2 DiT layers reside on GPU at a time, dramatically reducing VRAM. For **video models** (where per-layer compute >> H2D transfer), the memcpy is completely hidden behind computation — **zero-cost offload** that saves VRAM without speed penalty ([PR #15511](https://github.com/sgl-project/sglang/pull/15511)). | Saves VRAM (40 GB → ~11 GB for Wan A14B); zero or near-zero speed cost for video models | Enabled by default for Wan/MOVA video models. Incompatible with Cache-DiT. For **image models** or highly parallelized setups (many GPUs, small per-GPU compute), the copy stream may not be fully hidden and can cause slowdown. | | **Offload Prefetch Size** | `--dit-offload-prefetch-size F` | Fine-grained control over layerwise offload: how many layers to prefetch ahead. `0.0` = 1 layer (min VRAM), `0.1` = 10% of layers, `≥1` = absolute layer count. | Tune for cases where default offload has copy stream interference (e.g. image models). 0.05–0.1 is a good starting point. | Values ≥ 0.5 approach no-offload VRAM with worse performance. See [PR #17693](https://github.com/sgl-project/sglang/pull/17693) for benchmarks on image models. | | **FSDP Inference** | `--use-fsdp-inference` | Uses PyTorch FSDP to shard model weights across GPUs with prefetch. Low latency, low VRAM. | Reduces per-GPU VRAM | Mutually exclusive with `--dit-layerwise-offload`. More overhead than SP on high-bandwidth interconnects. | | **CPU Offload (components)** | `--text-encoder-cpu-offload`, `--image-encoder-cpu-offload`, `--vae-cpu-offload`, `--dit-cpu-offload` | Offloads specific pipeline components to CPU when not in use. | Reduces peak VRAM | Adds H2D transfer latency when the component is needed. Auto-enabled for low-VRAM GPUs (<30 GB). **Tip:** after the first request completes, the console prints a peak VRAM analysis with suggestions on which offload flags can be safely disabled — look for the `"Components that could stay resident"` log line. | | **Pin CPU Memory** | `--pin-cpu-memory` | Uses pinned (page-locked) memory for CPU offload transfers. | Faster H2D transfers | Slightly higher host memory usage. Enabled by default; disable only as workaround for CUDA errors. | | **Attention Backend (lossless)** | `--attention-backend fa` | Selects lossless attention kernel: `fa` (FlashAttention 2/3/4), `torch_sdpa`. FA is the fastest lossless option. | FA >> SDPA for long sequences | FA requires compatible GPU (Ampere+). `fa3`/`fa4` are aliased to `fa`. Ring attention only works with `fa` or `sage_attn`. | | **Parallel Folding** | *(automatic when SP > 1)* | Reuses the SP process group as TP for the T5 text encoder, so text encoding is parallelized "for free". | Faster text encoding on multi-GPU | Automatic; no user action needed. Only applies to T5-based pipelines. | --- ## Section 2: Lossy Optimizations These options **trade output quality** for speed or VRAM savings. Results will differ from the baseline. | Option | CLI Flag / Env Var | What It Does | Speedup | Quality Impact / Limitations | |---|---|---|---|---| | **Approximate Attention** | `--attention-backend sage_attn` / `sage_attn_3` / `sliding_tile_attn` / `video_sparse_attn` / `sparse_video_gen_2_attn` / `vmoba_attn` / `sla_attn` / `sage_sla_attn` | Replaces exact attention with approximate or sparse variants. `sage_attn`: INT8/FP8 quantized Q·K; `sliding_tile_attn`: spatial-temporal tile skipping; others: model-specific sparse patterns. | ~1.5–2x on attention (varies by backend) | Quality degradation varies by backend and model. `sage_attn` is the most general; sparse backends (`sliding_tile_attn`, `video_sparse_attn`, etc.) are video-model-specific and may require config files (e.g. `--mask-strategy-file-path` for STA). Requires corresponding packages installed. | | **Cache-DiT** | `SGLANG_CACHE_DIT_ENABLED=true` + `--cache-dit-config ` | Caches intermediate residuals across denoising steps and skips redundant computations via a Selective Computation Mask (SCM). | ~1.5–2x on supported models | Quality depends on SCM config. Incompatible with `--dit-layerwise-offload`. Requires correct per-model config YAML. | | **Quantized Models (Nunchaku / SVDQuant)** | `--enable-svdquant --transformer-weights-path ` + optional `--quantization-precision int4\|nvfp4`, `--quantization-rank 32` | W4A4-style quantization via [Nunchaku](https://nunchaku.tech). Reduces DiT weight memory by ~4x. Precision/rank can be auto-inferred from weight filename or set explicitly. | ~1.5–2x compute speedup | Lossy quantization; quality depends on rank and precision. Requires pre-quantized weights. Ampere (SM8x) or SM12x only (no Hopper SM90). Higher rank = better quality but more memory. | | **Pre-quantized Weights** | `--transformer-weights-path ` | Load any pre-quantized transformer weights (FP8, INT8, etc.) from a single `.safetensors` file, a directory, or a HuggingFace repo ID. | ~1.3–1.5x compute (dtype dependent) | Requires pre-converted weights (e.g. via `tools/convert_hf_to_fp8.py` for FP8). Quality slightly worse than BF16; varies by quantization format. | | **Component Precision Override** | `--dit-precision fp16`, `--vae-precision fp16\|bf16` | On-the-fly dtype conversion for individual components. E.g. convert a BF16 model to FP16 at load time, or run VAE in BF16 instead of FP32. | Reduces memory; FP16 can be faster on some GPUs | May affect numerical stability. VAE is FP32 by default for accuracy; lowering it is lossy. DiT defaults to BF16. | | **Fewer Inference Steps** | `--num-inference-steps N` (sampling param) | Reduces the number of denoising steps. Fewer steps = faster. | Linear speedup | Quality degrades with too few steps. Model-dependent optimal range. | --- ## Quick Recipes ### Maximum speed, video model, multi-GPU, lossless (Wan A14B, 8 GPUs) ```bash sglang generate --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --num-gpus 8 --enable-cfg-parallel --ulysses-degree 4 \ --enable-torch-compile --warmup \ --text-encoder-cpu-offload true \ --prompt "..." --save-output ``` Note: `--dit-layerwise-offload` is enabled by default for Wan/MOVA video models and is zero-cost (H2D fully overlapped with compute). No need to disable it. ### Maximum speed, image model, single GPU, lossless ```bash sglang generate --model-path \ --enable-torch-compile --warmup \ --dit-layerwise-offload false \ --prompt "..." --save-output ``` Note: for image models, per-layer compute is smaller, so layerwise offload may not fully hide H2D transfer. Disable it if VRAM allows. ### Low VRAM, decent speed (single GPU) ```bash sglang generate --model-path \ --enable-torch-compile --warmup \ --dit-layerwise-offload --dit-offload-prefetch-size 0.1 \ --text-encoder-cpu-offload true --vae-cpu-offload true \ --prompt "..." --save-output ``` ### Maximum speed, lossy (SageAttention + Cache-DiT) ```bash SGLANG_CACHE_DIT_ENABLED=true sglang generate --model-path \ --attention-backend sage_attn \ --cache-dit-config \ --enable-torch-compile --warmup \ --dit-layerwise-offload false \ --prompt "..." --save-output ``` --- ## Tips - **Benchmarking**: always use `--warmup` and look for the line ending with `(with warmup excluded)` for accurate timing. - **Perf dump**: use `--perf-dump-path result.json` to save structured metrics, then compare with `python python/sglang/multimodal_gen/benchmarks/compare_perf.py baseline.json result.json`. - **Offload tuning**: after the first request, the runtime logs peak GPU memory and which components could stay resident. Use this to decide which `--*-cpu-offload` flags to disable. - **Backend selection**: `--backend sglang` (default, auto-detected) enables all native optimizations (fused kernels, SP, etc.). `--backend diffusers` falls back to vanilla Diffusers pipelines but supports `--cache-dit-config` and diffusers attention backends. ================================================ FILE: python/sglang/multimodal_gen/.claude/skills/support-new-model/SKILL.md ================================================ --- name: add-new-diffusion-model description: Step-by-step guide for adding a new diffusion model to SGLang. Covers the recommended Hybrid Monolithic pipeline pattern (BeforeDenoisingStage), as well as when to use the Modular Composition Style. Includes pipeline config, model components, registration, and testing. --- # Tutorial: Adding a New Diffusion Model to SGLang This tutorial walks through adding support for a new diffusion model. SGLang Diffusion supports two pipeline styles; choose the one that best fits your model. ## Two Pipeline Styles ### Style A: Hybrid Monolithic Pipeline (Recommended) The recommended default for most new models. Uses a three-stage structure: ``` BeforeDenoisingStage (model-specific) --> DenoisingStage (standard) --> DecodingStage (standard) ``` - **BeforeDenoisingStage**: A single, model-specific stage that consolidates all pre-processing logic: input validation, text encoding, image encoding, latent preparation, timestep setup. This stage is unique per model. - **DenoisingStage**: Framework-standard stage for the denoising loop (DiT/UNet forward passes). Shared across models. - **DecodingStage**: Framework-standard stage for VAE decoding. Shared across models. **Why recommended?** Modern diffusion models have highly heterogeneous pre-processing requirements (different text encoders, different latent formats, different conditioning mechanisms). The Hybrid approach keeps pre-processing isolated per model, avoids fragile shared stages with excessive conditional logic, and lets developers port Diffusers reference code quickly. ### Style B: Modular Composition Style Uses the framework's fine-grained standard stages (`TextEncodingStage`, `LatentPreparationStage`, `TimestepPreparationStage`, etc.) to build the pipeline by composition. This style is appropriate when: - **The new model's pre-processing can largely reuse existing stages** — e.g., a model that uses standard CLIP/T5 text encoding + standard latent preparation with minimal customization. In this case, `add_standard_t2i_stages()` or `add_standard_ti2i_stages()` may be all you need. - **A model-specific optimization needs to be extracted as a standalone stage** — e.g., a specialized encoding or conditioning step that benefits from being a separate stage for profiling, parallelism control, or reuse across multiple pipeline variants. See existing Modular examples: `QwenImagePipeline` (uses `add_standard_t2i_stages`), `FluxPipeline`, `WanPipeline`. ### How to Choose | Situation | Recommended Style | |-----------|-------------------| | Model has unique/complex pre-processing (VLM captioning, AR token generation, custom latent packing, etc.) | **Hybrid** — consolidate into a BeforeDenoisingStage | | Model fits neatly into standard text-to-image or text+image-to-image pattern | **Modular** — use `add_standard_t2i_stages()` / `add_standard_ti2i_stages()` | | Porting a Diffusers pipeline with many custom steps | **Hybrid** — copy the `__call__` logic into a single stage | | Adding a variant of an existing model that shares most logic | **Modular** — reuse existing stages, customize via PipelineConfig callbacks | | A specific pre-processing step needs special parallelism or profiling isolation | **Modular** — extract that step as a dedicated stage | **Key principle (both styles)**: The stage(s) before `DenoisingStage` must produce a `Req` batch object with all the standard tensor fields that `DenoisingStage` expects (latents, timesteps, prompt_embeds, etc.). As long as this contract is met, the pipeline remains composable regardless of which style you use. --- ## Key Files and Directories | Purpose | Path | |---------|------| | Pipeline classes | `python/sglang/multimodal_gen/runtime/pipelines/` | | Model-specific stages | `python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/` | | PipelineStage base class | `python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py` | | Pipeline base class | `python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py` | | Standard stages (Denoising, Decoding) | `python/sglang/multimodal_gen/runtime/pipelines_core/stages/` | | Pipeline configs | `python/sglang/multimodal_gen/configs/pipeline_configs/` | | Sampling params | `python/sglang/multimodal_gen/configs/sample/` | | DiT model implementations | `python/sglang/multimodal_gen/runtime/models/dits/` | | VAE implementations | `python/sglang/multimodal_gen/runtime/models/vaes/` | | Encoder implementations | `python/sglang/multimodal_gen/runtime/models/encoders/` | | Scheduler implementations | `python/sglang/multimodal_gen/runtime/models/schedulers/` | | Model/VAE/DiT configs | `python/sglang/multimodal_gen/configs/models/dits/`, `vaes/`, `encoders/` | | Central registry | `python/sglang/multimodal_gen/registry.py` | --- ## Step-by-Step Implementation ### Step 1: Obtain and Study the Reference Implementation **Before writing any code, ask the user to provide the model's original implementation or Diffusers pipeline code.** You need the actual source code to work from — do not guess or assume the model's architecture. If the user has not provided it, request: - The model's Diffusers pipeline source (e.g., the `pipeline_*.py` file from the `diffusers` library or HuggingFace repo) - Or the model's official reference implementation (e.g., from the model author's GitHub repo) - Or the HuggingFace model ID so you can look up `model_index.json` and the associated pipeline class Once you have the reference code, study it thoroughly: 1. Find the model's `model_index.json` to identify required modules (text_encoder, vae, transformer, scheduler, etc.) 2. Read the Diffusers pipeline's `__call__` method end-to-end. Identify: - How text prompts are encoded - How latents are prepared (shape, dtype, scaling) - How timesteps/sigmas are computed - What conditioning kwargs the DiT/UNet expects - How the denoising loop works (classifier-free guidance, etc.) - How VAE decoding is done (scaling factors, tiling, etc.) ### Step 2: Evaluate Reuse of Existing Pipelines and Stages **Before creating any new files, check whether an existing pipeline or stage can be reused or extended.** Only create new pipelines/stages when the existing ones would require extensive modifications or when no similar implementation exists. Specifically: 1. **Compare the new model's architecture against existing pipelines** (Flux, Wan, Qwen-Image, GLM-Image, HunyuanVideo, LTX, etc.). If the new model shares most of its structure with an existing one (e.g., same text encoders, similar latent format, compatible denoising loop), prefer: - Adding a new config variant to the existing pipeline rather than creating a new pipeline class - Reusing the existing `BeforeDenoisingStage` with minor parameter differences - Using `add_standard_t2i_stages()` / `add_standard_ti2i_stages()` / `add_standard_ti2v_stages()` if the model fits standard patterns 2. **Check existing stages** in `runtime/pipelines_core/stages/` and `stages/model_specific_stages/`. If an existing stage handles 80%+ of what the new model needs, extend it rather than duplicating it. 3. **Check existing model components** — many models share VAEs (e.g., `AutoencoderKL`), text encoders (CLIP, T5), and schedulers. Reuse these directly instead of re-implementing. **Rule of thumb**: Only create a new file when the existing implementation would need substantial structural changes to accommodate the new model, or when no architecturally similar implementation exists. ### Step 3: Implement Model Components Adapt or implement the model's core components in the appropriate directories. **DiT/Transformer** (`runtime/models/dits/{model_name}.py`): ```python # python/sglang/multimodal_gen/runtime/models/dits/my_model.py import torch import torch.nn as nn from sglang.multimodal_gen.runtime.layers.layernorm import ( LayerNormScaleShift, RMSNormScaleShift, ) from sglang.multimodal_gen.runtime.layers.attention.selector import ( get_attn_backend, ) class MyModelTransformer2DModel(nn.Module): """DiT model for MyModel. Adapt from the Diffusers/reference implementation. Key points: - Use SGLang's fused LayerNorm/RMSNorm ops (see use-efficient-diffusion-kernels skill) - Use SGLang's attention backend selector - Keep the same parameter naming as Diffusers for weight loading compatibility """ def __init__(self, config): super().__init__() # ... model layers ... def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, timestep: torch.Tensor, # ... model-specific kwargs ... ) -> torch.Tensor: # ... forward pass ... return output ``` **Tensor Parallel (TP) and Sequence Parallel (SP)**: For multi-GPU deployment, it is recommended to add TP/SP support to the DiT model. This can be done incrementally after the single-GPU implementation is verified. Reference existing implementations and adapt to your model's architecture: - **Wan model** (`runtime/models/dits/wanvideo.py`) — Full TP + SP reference: - TP: Uses `ColumnParallelLinear` for Q/K/V projections, `RowParallelLinear` for output projections, attention heads divided by `tp_size` - SP: Sequence dimension sharding via `get_sp_world_size()`, padding for alignment, `sequence_model_parallel_all_gather` for aggregation - Cross-attention skips SP (`skip_sequence_parallel=is_cross_attention`) - **Qwen-Image model** (`runtime/models/dits/qwen_image.py`) — SP + USPAttention reference: - SP: Uses `USPAttention` (Ulysses + Ring Attention), configured via `--ulysses-degree` / `--ring-degree` - TP: Uses `MergedColumnParallelLinear` for QKV (with Nunchaku quantization), `ReplicatedLinear` otherwise **Important**: These are references only — each model has its own architecture and parallelism requirements. Consider: - How attention heads can be divided across TP ranks - Whether the model's sequence dimension is naturally shardable for SP - Which linear layers benefit from column/row parallel sharding vs. replication - Whether cross-attention or other special modules need SP exclusion Key imports for distributed support: ```python from sglang.multimodal_gen.runtime.distributed import ( divide, get_sp_group, get_sp_world_size, get_tp_world_size, sequence_model_parallel_all_gather, ) from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, RowParallelLinear, ReplicatedLinear, ) ``` **VAE** (`runtime/models/vaes/{model_name}.py`): Implement if the model uses a non-standard VAE. Many models reuse existing VAEs. **Encoders** (`runtime/models/encoders/{model_name}.py`): Implement if the model uses custom text/image encoders. **Schedulers** (`runtime/models/schedulers/{scheduler_name}.py`): Implement if the model requires a custom scheduler not available in Diffusers. ### Step 4: Create Model Configs **DiT Config** (`configs/models/dits/{model_name}.py`): ```python # python/sglang/multimodal_gen/configs/models/dits/mymodel.py from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.dits.base import DiTConfig @dataclass class MyModelDitConfig(DiTConfig): arch_config: dict = field(default_factory=lambda: { "in_channels": 16, "num_layers": 24, "patch_size": 2, # ... model-specific architecture params ... }) ``` **VAE Config** (`configs/models/vaes/{model_name}.py`): ```python from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.vaes.base import VAEConfig @dataclass class MyModelVAEConfig(VAEConfig): vae_scale_factor: int = 8 # ... VAE-specific params ... ``` **Sampling Params** (`configs/sample/{model_name}.py`): ```python from dataclasses import dataclass from sglang.multimodal_gen.configs.sample.base import SamplingParams @dataclass class MyModelSamplingParams(SamplingParams): num_inference_steps: int = 50 guidance_scale: float = 7.5 height: int = 1024 width: int = 1024 # ... model-specific defaults ... ``` ### Step 5: Create PipelineConfig The `PipelineConfig` holds static model configuration and defines callback methods used by the standard `DenoisingStage` and `DecodingStage`. ```python # python/sglang/multimodal_gen/configs/pipeline_configs/my_model.py from dataclasses import dataclass, field from sglang.multimodal_gen.configs.pipeline_configs.base import ( ImagePipelineConfig, # for image generation # SpatialImagePipelineConfig, # alternative base # VideoPipelineConfig, # for video generation ) from sglang.multimodal_gen.configs.models.dits.mymodel import MyModelDitConfig from sglang.multimodal_gen.configs.models.vaes.mymodel import MyModelVAEConfig @dataclass class MyModelPipelineConfig(ImagePipelineConfig): """Pipeline config for MyModel. This config provides callbacks that the standard DenoisingStage and DecodingStage use during execution. The BeforeDenoisingStage handles all model-specific pre-processing independently. """ task_type: ModelTaskType = ModelTaskType.T2I vae_precision: str = "bf16" should_use_guidance: bool = True vae_tiling: bool = False enable_autocast: bool = False dit_config: DiTConfig = field(default_factory=MyModelDitConfig) vae_config: VAEConfig = field(default_factory=MyModelVAEConfig) # --- Callbacks used by DenoisingStage --- def get_freqs_cis(self, batch, device, rotary_emb, dtype): """Prepare rotary position embeddings for the DiT.""" # Model-specific RoPE computation ... return freqs_cis def prepare_pos_cond_kwargs(self, batch, latent_model_input, t, **kwargs): """Build positive conditioning kwargs for each denoising step.""" return { "hidden_states": latent_model_input, "encoder_hidden_states": batch.prompt_embeds[0], "timestep": t, # ... model-specific kwargs ... } def prepare_neg_cond_kwargs(self, batch, latent_model_input, t, **kwargs): """Build negative conditioning kwargs for CFG.""" return { "hidden_states": latent_model_input, "encoder_hidden_states": batch.negative_prompt_embeds[0], "timestep": t, # ... model-specific kwargs ... } # --- Callbacks used by DecodingStage --- def get_decode_scale_and_shift(self): """Return (scale, shift) for latent denormalization before VAE decode.""" return self.vae_config.latents_std, self.vae_config.latents_mean def post_denoising_loop(self, latents, batch): """Optional post-processing after the denoising loop finishes.""" return latents.to(torch.bfloat16) def post_decoding(self, frames, server_args): """Optional post-processing after VAE decoding.""" return frames ``` **Important**: The `prepare_pos_cond_kwargs` / `prepare_neg_cond_kwargs` methods define what the DiT receives at each denoising step. These must match the DiT's `forward()` signature. ### Step 6: Implement the BeforeDenoisingStage (Core Step) This is the heart of the Hybrid pattern. Create a single stage that handles ALL pre-processing. ```python # python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/my_model.py import torch from typing import List, Optional, Union from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class MyModelBeforeDenoisingStage(PipelineStage): """Monolithic pre-processing stage for MyModel. Consolidates all logic before the denoising loop: - Input validation - Text/image encoding - Latent preparation - Timestep/sigma computation This stage produces a Req batch with all fields required by the standard DenoisingStage. """ def __init__(self, vae, text_encoder, tokenizer, transformer, scheduler): super().__init__() self.vae = vae self.text_encoder = text_encoder self.tokenizer = tokenizer self.transformer = transformer self.scheduler = scheduler # ... other initialization (image processors, scale factors, etc.) ... # --- Internal helper methods --- # Copy/adapt directly from the Diffusers reference pipeline. # These are private to this stage; no need to make them reusable. def _encode_prompt(self, prompt, device, dtype): """Encode text prompt into embeddings.""" # ... model-specific text encoding logic ... return prompt_embeds, negative_prompt_embeds def _prepare_latents(self, batch_size, height, width, dtype, device, generator): """Create initial noisy latents.""" # ... model-specific latent preparation ... return latents def _prepare_timesteps(self, num_inference_steps, device): """Compute the timestep/sigma schedule.""" # ... model-specific timestep computation ... return timesteps, sigmas # --- Main forward method --- @torch.no_grad() def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Execute all pre-processing and populate batch for DenoisingStage. This method mirrors the first half of a Diffusers pipeline __call__, up to (but not including) the denoising loop. """ device = get_local_torch_device() dtype = torch.bfloat16 generator = torch.Generator(device=device).manual_seed(batch.seed) # 1. Encode prompt prompt_embeds, negative_prompt_embeds = self._encode_prompt( batch.prompt, device, dtype ) # 2. Prepare latents latents = self._prepare_latents( batch_size=1, height=batch.height, width=batch.width, dtype=dtype, device=device, generator=generator, ) # 3. Prepare timesteps timesteps, sigmas = self._prepare_timesteps( batch.num_inference_steps, device ) # 4. Populate batch with everything DenoisingStage needs batch.prompt_embeds = [prompt_embeds] batch.negative_prompt_embeds = [negative_prompt_embeds] batch.latents = latents batch.timesteps = timesteps batch.num_inference_steps = len(timesteps) batch.sigmas = sigmas batch.generator = generator batch.raw_latent_shape = latents.shape batch.height = batch.height batch.width = batch.width return batch ``` **Key fields that `DenoisingStage` expects on the batch** (set these in your `forward`): | Field | Type | Description | |-------|------|-------------| | `batch.latents` | `torch.Tensor` | Initial noisy latent tensor | | `batch.timesteps` | `torch.Tensor` | Timestep schedule | | `batch.num_inference_steps` | `int` | Number of denoising steps | | `batch.sigmas` | `list[float]` | Sigma schedule (as a list, not numpy) | | `batch.prompt_embeds` | `list[torch.Tensor]` | Positive prompt embeddings (wrapped in list) | | `batch.negative_prompt_embeds` | `list[torch.Tensor]` | Negative prompt embeddings (wrapped in list) | | `batch.generator` | `torch.Generator` | RNG generator for reproducibility | | `batch.raw_latent_shape` | `tuple` | Original latent shape before any packing | | `batch.height` / `batch.width` | `int` | Output dimensions | ### Step 7: Define the Pipeline Class The pipeline class is minimal -- it just wires the stages together. ```python # python/sglang/multimodal_gen/runtime/pipelines/my_model.py from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.stages import DenoisingStage from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.my_model import ( MyModelBeforeDenoisingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs class MyModelPipeline(LoRAPipeline, ComposedPipelineBase): pipeline_name = "MyModelPipeline" # Must match model_index.json _class_name _required_config_modules = [ "text_encoder", "tokenizer", "vae", "transformer", "scheduler", # ... list all modules from model_index.json ... ] def create_pipeline_stages(self, server_args: ServerArgs): # 1. Monolithic pre-processing (model-specific) self.add_stage( MyModelBeforeDenoisingStage( vae=self.get_module("vae"), text_encoder=self.get_module("text_encoder"), tokenizer=self.get_module("tokenizer"), transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), ), ) # 2. Standard denoising loop (framework-provided) self.add_stage( DenoisingStage( transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), ), ) # 3. Standard VAE decoding (framework-provided) self.add_standard_decoding_stage() # REQUIRED: This is how the registry discovers the pipeline EntryClass = [MyModelPipeline] ``` ### Step 8: Register the Model In `python/sglang/multimodal_gen/registry.py`, register your configs: ```python register_configs( model_family="my_model", sampling_param_cls=MyModelSamplingParams, pipeline_config_cls=MyModelPipelineConfig, hf_model_paths=[ "org/my-model-name", # HuggingFace model ID(s) ], ) ``` The `EntryClass` in your pipeline file is automatically discovered by the registry's `_discover_and_register_pipelines()` function -- no additional registration needed for the pipeline class itself. ### Step 9: Verify Output Quality After implementation, **you must verify that the generated output is not noise**. A noisy or garbled output image/video is the most common sign of an incorrect implementation. Common causes include: - Incorrect latent scale/shift factors (`get_decode_scale_and_shift` returning wrong values) - Wrong timestep/sigma schedule (order, dtype, or value range) - Mismatched conditioning kwargs (fields not matching the DiT's `forward()` signature) - Incorrect VAE decoder configuration (wrong `vae_scale_factor`, missing denormalization) - Rotary embedding style mismatch (`is_neox_style` set incorrectly) - Wrong prompt embedding format (missing list wrapping, wrong encoder output selection) **If the output is noise, the implementation is incorrect — do not ship it.** Debug by: 1. Comparing intermediate tensor values (latents, prompt_embeds, timesteps) against the Diffusers reference pipeline 2. Running the Diffusers pipeline and SGLang pipeline side-by-side with the same seed 3. Checking each stage's output shape and value range independently ## Reference Implementations ### Hybrid Style (recommended for most new models) | Model | Pipeline | BeforeDenoisingStage | PipelineConfig | |-------|----------|---------------------|----------------| | GLM-Image | `runtime/pipelines/glm_image.py` | `stages/model_specific_stages/glm_image.py` | `configs/pipeline_configs/glm_image.py` | | Qwen-Image-Layered | `runtime/pipelines/qwen_image.py` (`QwenImageLayeredPipeline`) | `stages/model_specific_stages/qwen_image_layered.py` | `configs/pipeline_configs/qwen_image.py` (`QwenImageLayeredPipelineConfig`) | ### Modular Style (when standard stages fit well) | Model | Pipeline | Notes | |-------|----------|-------| | Qwen-Image (T2I) | `runtime/pipelines/qwen_image.py` | Uses `add_standard_t2i_stages()` — standard text encoding + latent prep fits this model | | Qwen-Image-Edit | `runtime/pipelines/qwen_image.py` | Uses `add_standard_ti2i_stages()` — standard image-to-image flow | | Flux | `runtime/pipelines/flux.py` | Uses `add_standard_t2i_stages()` with custom `prepare_mu` | | Wan | `runtime/pipelines/wan_pipeline.py` | Uses `add_standard_ti2v_stages()` | --- ## Checklist Before submitting, verify: **Common (both styles):** - [ ] **Pipeline file** exists at `runtime/pipelines/{model_name}.py` with `EntryClass` - [ ] **PipelineConfig** at `configs/pipeline_configs/{model_name}.py` - [ ] **SamplingParams** at `configs/sample/{model_name}.py` - [ ] **DiT model** at `runtime/models/dits/{model_name}.py` - [ ] **DiT config** at `configs/models/dits/{model_name}.py` - [ ] **VAE** — reuse existing (e.g., `AutoencoderKL`) or create new at `runtime/models/vaes/` - [ ] **VAE config** — reuse existing or create new at `configs/models/vaes/{model_name}.py` - [ ] **Registry entry** in `registry.py` via `register_configs()` - [ ] `pipeline_name` matches Diffusers `model_index.json` `_class_name` - [ ] `_required_config_modules` lists all modules from `model_index.json` - [ ] `PipelineConfig` callbacks (`prepare_pos_cond_kwargs`, `get_freqs_cis`, etc.) match DiT's `forward()` signature - [ ] Latent scale/shift factors are correctly configured - [ ] Use fused kernels where possible (see `use-efficient-diffusion-kernels` skill) - [ ] Weight names match Diffusers for automatic loading - [ ] **TP/SP support** considered for DiT model (recommended; reference `wanvideo.py` for TP+SP, `qwen_image.py` for USPAttention) - [ ] **Output quality verified** — generated images/videos are not noise; compared against Diffusers reference output **Hybrid style only:** - [ ] **BeforeDenoisingStage** at `stages/model_specific_stages/{model_name}.py` - [ ] `BeforeDenoisingStage.forward()` populates all fields needed by `DenoisingStage` ## Common Pitfalls 1. **`batch.sigmas` must be a Python list**, not a numpy array. Use `.tolist()` to convert. 2. **`batch.prompt_embeds` is a list of tensors** (one per encoder), not a single tensor. Wrap with `[tensor]`. 3. **Don't forget `batch.raw_latent_shape`** -- `DecodingStage` uses it to unpack latents. 4. **Rotary embedding style matters**: `is_neox_style=True` = split-half rotation, `is_neox_style=False` = interleaved. Check the reference model carefully. 5. **VAE precision**: Many VAEs need fp32 or bf16 for numerical stability. Set `vae_precision` in the PipelineConfig accordingly. 6. **Avoid forcing model-specific logic into shared stages**: If your model's pre-processing doesn't naturally fit the existing standard stages, prefer the Hybrid pattern with a dedicated BeforeDenoisingStage rather than adding conditional branches to shared stages. ## After Implementation: Tests and Performance Data Once the model is working and output quality is verified, **ask the user** whether they would like to: 1. **Add tests** — Create unit tests and/or integration tests for the new model. Tests should cover: - Pipeline construction and stage wiring - Single-GPU inference producing non-noise output - Multi-GPU inference (TP/SP) if supported - See the `write-sglang-test` skill for test conventions and placement guidelines 2. **Generate performance data** — Run benchmarks and collect perf metrics: - Single-GPU latency and throughput (look for `Pixel data generated successfully in xxxx seconds` in console output; use the `warmup excluded` line for accurate timing) - Multi-GPU scaling (TP/SP) throughput comparison - Use `python/sglang/multimodal_gen/benchmarks/bench_serving.py` for serving benchmarks Do not skip this step — always ask the user before proceeding, as test and benchmark requirements vary per model. ================================================ FILE: python/sglang/multimodal_gen/README.md ================================================
**SGLang diffusion is an inference framework for accelerated image/video generation.** SGLang diffusion features an end-to-end unified pipeline for accelerating diffusion models. It is designed to be modular and extensible, allowing users to easily add new models and optimizations. ## Key Features SGLang Diffusion has the following features: - Broad model support: Wan series, FastWan series, Hunyuan, Qwen-Image, Qwen-Image-Edit, Flux, Z-Image, GLM-Image - Fast inference speed: enpowered by highly optimized kernel from sgl-kernel and efficient scheduler loop - Ease of use: OpenAI-compatible api, CLI, and python sdk support - Multi-platform support: - NVIDIA GPUs (H100, H200, A100, B200, 4090) - AMD GPUs (MI300X, MI325X) - Ascend NPU (A2, A3) - Apple Silicon (M-series via MPS) - Moore Threads GPUs (MTT S5000) ### AMD/ROCm Support SGLang Diffusion supports AMD Instinct GPUs through ROCm. On AMD platforms, we use the Triton attention backend and leverage AITER kernels for optimized layernorm and other operations. See the [installation guide](https://github.com/sgl-project/sglang/tree/main/docs/diffusion/installation.md) for setup instructions. ### Moore Threads/MUSA Support SGLang Diffusion supports Moore Threads GPUs (MTGPU) through the MUSA software stack. On MUSA platforms, we use the Torch SDPA backend for attention. See the [installation guide](https://github.com/sgl-project/sglang/tree/main/docs/diffusion/installation.md) for setup instructions. ### Apple MPS Support SGLang Diffusion supports Apple Silicon (M-series) via the MPS backend. Since Triton is Linux-only, all Triton kernels are replaced with PyTorch-native fallbacks on MPS. Norm operations can be optionally accelerated with MLX fused Metal kernels (`SGLANG_USE_MLX=1`). See the [installation guide](https://github.com/sgl-project/sglang/tree/main/docs/diffusion/installation.md) for setup instructions. ## Getting Started ```bash uv pip install 'sglang[diffusion]' --prerelease=allow ``` For more installation methods (e.g. pypi, uv, docker, ROCm/AMD, MUSA/Moore Threads), check [install.md](https://github.com/sgl-project/sglang/tree/main/docs/diffusion/installation.md). ## Inference Here's a minimal example to generate a video using the default settings: ```python from sglang.multimodal_gen import DiffGenerator def main(): # Create a diff generator from a pre-trained model generator = DiffGenerator.from_pretrained( model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", num_gpus=1, # Adjust based on your hardware ) # Generate the video video = generator.generate( sampling_params_kwargs=dict( prompt="A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest.", return_frames=True, # Also return frames from this call (defaults to False) output_path="my_videos/", # Controls where videos are saved save_output=True ) ) if __name__ == '__main__': main() ``` Or, more simply, with the CLI: ```bash sglang generate --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ --text-encoder-cpu-offload --pin-cpu-memory \ --prompt "A curious raccoon" \ --save-output ``` ### LoRA support Apply LoRA adapters via `--lora-path`: ```bash sglang generate \ --model-path Qwen/Qwen-Image-Edit-2511 \ --lora-path prithivMLmods/Qwen-Image-Edit-2511-Anime \ --prompt "Transform into anime." \ --image-path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png" \ --save-output ``` For more usage examples (e.g. OpenAI compatible API, server mode), check [cli.md](https://github.com/sgl-project/sglang/tree/main/docs/diffusion/cli.md). ## Contributing All contributions are welcome. The contribution guide is available [here](https://github.com/sgl-project/sglang/tree/main/docs/diffusion/contributing.md). ## Acknowledgement We learnt and reused code from the following projects: - [FastVideo](https://github.com/hao-ai-lab/FastVideo.git). The major components of this repo are based on a fork of FastVideo on Sept. 24, 2025. - [xDiT](https://github.com/xdit-project/xDiT). We used the parallelism library from it. - [diffusers](https://github.com/huggingface/diffusers) We used the pipeline design from it. ================================================ FILE: python/sglang/multimodal_gen/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from sglang.multimodal_gen.configs.pipeline_configs import PipelineConfig from sglang.multimodal_gen.configs.sample import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator __all__ = ["DiffGenerator", "PipelineConfig", "SamplingParams"] # Trigger multimodal CI tests ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/README.md ================================================ # ComfyUI SGLDiffusion Plugin A ComfyUI plugin for integrating with SGLang Diffusion server, supporting image and video generation capabilities. ## Installation 1. **Install SGLang**: Follow the [Installation Guide](../../docs/install.md) to install `sglang[diffusion]`. 2. **Install Plugin**: Copy this entire directory (`ComfyUI_SGLDiffusion`) to your ComfyUI `custom_nodes/` folder. 3. **Restart ComfyUI**: Restart ComfyUI to load the plugin. ## Usage The plugin supports two modes of operation: **Server Mode** (via HTTP API) and **Integrated Mode** (tight integration with ComfyUI). ### Supported Models - **Z-Image**: High-speed image generation models (e.g., `Z-Image-Turbo`) - **FLUX**: State-of-the-art text-to-image models (e.g., `FLUX.1-dev`) - **Qwen-Image**: Multi-modal image generation models (e.g., `Qwen-Image`,`Qwen-Image-2512`). *Note: Image editing support is currently experimental and may have some issues.* ### Mode 1: Server Mode (HTTP API) Connect to a standalone SGLang Diffusion server. 1. **Start SGLang Diffusion Server**: Ensure the server is running and accessible. 2. **Connect to Server**: Use the `SGLDiffusion Server Model` node to connect (default: `http://localhost:3000/v1`). 3. **Generate Content**: - `SGLDiffusion Generate Image`: For text-to-image and image editing. - `SGLDiffusion Generate Video`: For text-to-video and image-to-video. 4. **LoRA Support**: Use `SGLDiffusion Server Set LoRA` and `SGLDiffusion Server Unset LoRA`. ### Mode 2: Integrated Mode (Tight Integration) Leverage SGLang's high-performance sampling directly within ComfyUI while using ComfyUI's front-end nodes (CLIP, VAE, etc.). 1. **Load Model**: Use the `SGLDiffusion UNET Loader` node to load your diffusion model. 2. **Configure Options**: Use the `SGLDiffusion Options` node to set runtime parameters like `num_gpus`, `tp_size`, `model_type`, or `enable_torch_compile`. 3. **Sample**: Connect the loaded model to standard ComfyUI samplers. SGLang will handle the sampling process efficiently. 4. **LoRA Support**: Use the `SGLDiffusion LoRA Loader` for native LoRA integration. ## Example Workflows Reference workflow files are provided in the `workflows/` directory: - **`flux_sgld_sp.json`**: Multi-GPU (Sequence Parallelism) workflow for FLUX models. High-performance inference across multiple cards. - **`qwen_image_sgld.json`**: Qwen-Image generation with LoRA support. Optimized for multi-modal image tasks. - **`z-image_sgld.json`**: High-speed image generation using Z-Image. - **`sgld_text2img.json`**: Server-mode text-to-image generation with LoRA support. - **`sgld_image2video.json`**: Server-mode image-to-video generation. For other workflows supporting the models, you can easily use SGLang by replacing the official `UNET Loader` node with the `SGLDUNETLoader` node. Similarly, for LoRA support, replace the official LoRA loader with the `SGLDiffusion LoRA Loader`. To use these workflows: 1. Open ComfyUI. 2. Load the workflow JSON file from the `workflows/` directory. 3. Adjust the parameters and model paths as needed. 4. Run the workflow. ## Current Implementation This plugin provides a high-performance backend for diffusion models in ComfyUI. By leveraging SGLang's optimized kernels and parallelization techniques (Tensor Parallelism, TeaCache, etc.), it significantly accelerates the sampling process, especially for large models like FLUX. ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/__init__.py ================================================ """ ComfyUI SGLang Diffusion nodes package. """ try: from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] except ImportError: # ComfyUI dependencies not available (e.g., in test environment) NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/__init__.py ================================================ """ Core components for SGLang Diffusion ComfyUI integration. Provides generator, model patcher, and server API client. """ from .generator import SGLDiffusionGenerator from .model_patcher import SGLDModelPatcher from .server_api import SGLDiffusionServerAPI __all__ = [ "SGLDiffusionGenerator", "SGLDModelPatcher", "SGLDiffusionServerAPI", ] ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/generator.py ================================================ """ Generator for SGLang Diffusion ComfyUI integration. """ import logging import os import psutil from comfy import model_detection, model_management from comfy.utils import ( calculate_parameters, load_torch_file, state_dict_prefix_replace, unet_to_diffusers, ) logger = logging.getLogger(__name__) try: from sglang.multimodal_gen import DiffGenerator except ImportError: logger.error( "Error: sglang.multimodal_gen is not installed. Please install it using 'pip install sglang[diffusion]'" ) from ..executors import ( FluxExecutor, QwenImageEditExecutor, QwenImageExecutor, ZImageExecutor, ) from .model_patcher import SGLDModelPatcher class SGLDiffusionGenerator: """Generator for SGLang Diffusion models in ComfyUI.""" def __init__(self): self.model_path = None self.generator = None self.executor = None self.last_options = None self.pipeline_class_dict = { "flux": "ComfyUIFluxPipeline", "lumina2": "ComfyUIZImagePipeline", # zimage "qwen_image": "ComfyUIQwenImagePipeline", "qwen_image_edit": "ComfyUIQwenImageEditPipeline", } self.executor_class_dict = { "flux": FluxExecutor, "lumina2": ZImageExecutor, "qwen_image": QwenImageExecutor, "qwen_image_edit": QwenImageEditExecutor, } def __del__(self): self.close_generator() def init_generator( self, model_path: str, pipeline_class_name: str, kwargs: dict = None ): """Initialize the diffusion generator.""" if self.generator is not None: return self.generator if kwargs is None: kwargs = {} # Set comfyui_mode for ComfyUI integration kwargs["comfyui_mode"] = True self.generator = DiffGenerator.from_pretrained( model_path=model_path, pipeline_class_name=pipeline_class_name, **kwargs, ) return self.generator def kill_generator(self): """Kill worker processes manually because generator shutdown cannot terminate them.""" current_pid = os.getpid() worker_processes = [] for proc in psutil.process_iter(["pid", "name", "cmdline"]): try: # Look for sglang-diffusionWorker processes if proc.info["cmdline"]: cmdline = " ".join(proc.info["cmdline"]) if "sgl_diffusion::" in cmdline: if proc.info["pid"] != current_pid: worker_processes.append(proc) except (psutil.NoSuchProcess, psutil.AccessDenied): continue if worker_processes: logger.info( f"Found {len(worker_processes)} worker processes to terminate..." ) for proc in worker_processes: try: logger.info( f"Terminating worker process {proc.info['pid']}: {proc.info['name']}" ) proc.terminate() proc.wait(timeout=5) except psutil.TimeoutExpired: logger.warning( f"Process {proc.info['pid']} did not terminate, forcing kill..." ) try: proc.kill() proc.wait(timeout=2) except (psutil.NoSuchProcess, psutil.TimeoutExpired): pass except (psutil.NoSuchProcess, psutil.AccessDenied): pass def close_generator(self): """Close and cleanup the generator and all associated resources.""" if self.generator is not None: self.generator.shutdown() self.kill_generator() # Clear other references self.last_options = None self.model_path = None self.generator = None self.executor = None def get_comfyui_model(self, model_path: str, model_options: dict = None): """Get ComfyUI model from model path.""" if model_options is None: model_options = {} dtype = model_options.get("dtype", None) # Allow loading unets from checkpoint files sd = load_torch_file(model_path) diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd) temp_sd = state_dict_prefix_replace( sd, {diffusion_model_prefix: ""}, filter_keys=True ) if len(temp_sd) > 0: sd = temp_sd parameters = calculate_parameters(sd) load_device = model_management.get_torch_device() model_detect_config = model_detection.detect_unet_config(sd, "") model_type = model_detect_config.get("image_model", None) if model_type is None or model_type not in self.pipeline_class_dict: raise ValueError(f"Unsupported model type: {model_type}") model_config = model_detection.model_config_from_unet(sd, "") if model_config is not None: new_sd = sd else: new_sd = model_detection.convert_diffusers_mmdit(sd, "") if new_sd is not None: # diffusers mmdit model_config = model_detection.model_config_from_unet(new_sd, "") if model_config is None: return None else: # diffusers unet model_config = model_detection.model_config_from_diffusers_unet(sd) if model_config is None: return None diffusers_keys = unet_to_diffusers(model_config.unet_config) new_sd = {} for k in diffusers_keys: if k in sd: new_sd[diffusers_keys[k]] = sd.pop(k) offload_device = model_management.unet_offload_device() if dtype is None: unet_dtype = model_management.unet_dtype( model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes, ) else: unet_dtype = dtype manual_cast_dtype = model_management.unet_manual_cast( unet_dtype, load_device, model_config.supported_inference_dtypes ) model_config.set_inference_dtype(unet_dtype, manual_cast_dtype) model_config.custom_operations = model_options.get("custom_operations", None) model_config.unet_config["disable_unet_model_creation"] = True comfyui_model = model_config.get_model({}) return comfyui_model, model_config, model_type def load_model( self, model_path: str, model_options: dict = None, sgld_options: dict = None ): """Load model and return model patcher.""" gather_options = { "model_path": model_path, "model_options": model_options, "sgld_options": sgld_options, } if ( self.last_options is not None and self.last_options == gather_options and self.generator is not None ): return self.generator else: self.close_generator() self.last_options = gather_options self.model_path = model_path comfyui_model, model_config, model_type = self.get_comfyui_model( model_path, model_options ) if model_type is None or model_type not in self.pipeline_class_dict: raise ValueError(f"Unsupported model type: {model_type}") set_model_type = sgld_options.pop("model_type", None) if sgld_options else None if set_model_type is not None and set_model_type in self.pipeline_class_dict: model_type = set_model_type pipeline_class_name = self.pipeline_class_dict[model_type] self.generator = self.init_generator( model_path, pipeline_class_name, sgld_options ) executor_class = self.executor_class_dict[model_type] self.executor = executor_class( self.generator, model_path, comfyui_model, model_config ) comfyui_model.diffusion_model = self.executor load_device = model_management.get_torch_device() offload_device = model_management.unet_offload_device() return SGLDModelPatcher( comfyui_model, load_device, offload_device, model_type=model_type ) ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/model_patcher.py ================================================ """ Model patcher for SGLang Diffusion ComfyUI integration. """ import copy from comfy.model_patcher import ModelPatcher class SGLDModelPatcher(ModelPatcher): """Model patcher for SGLang Diffusion models in ComfyUI.""" def __init__( self, model, load_device, offload_device, size=0, weight_inplace_update=False, model_type=None, ): super().__init__( model, load_device, offload_device, size, weight_inplace_update ) self.lora_cache = {} self.model_type = model_type self.model_size_dict = { "flux": 27 * 1024 * 1024 * 1024, "lumina2": 8 * 1024 * 1024 * 1024, } def clone(self): """Clone the model patcher.""" n = SGLDModelPatcher( self.model, self.load_device, self.offload_device, self.size, weight_inplace_update=self.weight_inplace_update, ) n.patches = {} for k in self.patches: n.patches[k] = self.patches[k][:] n.patches_uuid = self.patches_uuid n.object_patches = self.object_patches.copy() n.model_options = copy.deepcopy(self.model_options) n.backup = self.backup n.object_patches_backup = self.object_patches_backup n.lora_cache = copy.copy(self.lora_cache) return n def model_size(self): """Get the model size in bytes.""" if self.model_type in self.model_size_dict: return self.model_size_dict[self.model_type] else: return 0 def load( self, device_to=None, lowvram_model_memory=0, force_patch_weights=False, full_load=False, ): """Load model (no-op for SGLang Diffusion).""" pass def patch_model( self, device_to=None, lowvram_model_memory=0, load_weights=True, force_patch_weights=False, ): """Patch model (no-op for SGLang Diffusion).""" pass def unpatch_model(self, device_to=None, unpatch_weights=True): """Unpatch model (no-op for SGLang Diffusion).""" pass ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/server_api.py ================================================ """ SGLang Diffusion Server API client. Provides a low-level interface for interacting with SGLang Diffusion HTTP server. """ import base64 import io import os import time from typing import Any, Dict, Optional import requests from PIL import Image class SGLDiffusionServerAPI: """Client for SGLang Diffusion HTTP server API.""" def __init__(self, base_url: str, api_key: str = "sk-proj-1234567890"): """ Initialize the API client. Args: base_url: Base URL of the SGLang Diffusion server (e.g., "http://localhost:30010/v1") api_key: API key for authentication (default: "sk-proj-1234567890") """ # Ensure base_url doesn't end with /v1 if it's already there if base_url.endswith("/v1"): self.base_url = base_url elif base_url.endswith("/v1/"): self.base_url = base_url.rstrip("/") else: self.base_url = f"{base_url.rstrip('/')}/v1" self.api_key = api_key self.headers = { "Content-Type": "application/json", "Authorization": f"Bearer {api_key}", } def get_model_info(self) -> Dict[str, Any]: """ Get information about the model served by this server. Returns: Dictionary containing model information including: - model_path: Path to the model - task_type: Type of task (e.g., "T2V", "I2I") - pipeline_name: Name of the pipeline - num_gpus: Number of GPUs - dit_precision: DiT model precision - vae_precision: VAE model precision """ try: # Remove /v1 from base_url for /models endpoint models_url = self.base_url.removesuffix("/v1") + "/models" response = requests.get(models_url, headers=self.headers, timeout=30) response.raise_for_status() return response.json() except requests.exceptions.RequestException as e: raise RuntimeError(f"Failed to get model info: {str(e)}") def generate_image( self, prompt: str, image_path: Optional[str] = None, mask_path: Optional[str] = None, size: Optional[str] = None, width: Optional[int] = None, height: Optional[int] = None, n: int = 1, negative_prompt: Optional[str] = None, guidance_scale: Optional[float] = None, num_inference_steps: Optional[int] = None, seed: Optional[int] = None, enable_teacache: bool = False, response_format: str = "b64_json", quality: Optional[str] = "auto", style: Optional[str] = "vivid", background: Optional[str] = "auto", output_format: Optional[str] = None, generator_device: Optional[str] = "cuda", ) -> Dict[str, Any]: """ Generate or edit an image using SGLang Diffusion API. If image_path is provided, calls the edit endpoint; otherwise calls the generation endpoint. Args: prompt: Text prompt for image generation/editing image_path: Optional path to input image file for editing. If provided, uses edit API. mask_path: Optional path to mask image file (only used when image_path is provided) size: Image size in format "WIDTHxHEIGHT" (e.g., "1024x1024") width: Image width (used if size is not provided) height: Image height (used if size is not provided) n: Number of images to generate (1-10) negative_prompt: Negative prompt to avoid certain elements guidance_scale: Classifier-free guidance scale num_inference_steps: Number of denoising steps seed: Random seed for reproducible generation enable_teacache: Enable TEA cache acceleration response_format: Response format ("b64_json" or "url") quality: Image quality ("auto", "standard", "hd") - only for generation style: Image style ("vivid" or "natural") - only for generation background: Background type ("auto", "transparent", "opaque") output_format: Output format ("png", "jpeg", "webp") generator_device: Device for random generator ("cuda" or "cpu") Returns: Dictionary containing the API response with generated/edited image data """ if not prompt: raise ValueError("Prompt cannot be empty") # Determine size if size is None: if width is not None and height is not None: size = f"{width}x{height}" else: size = "1024x1024" # Build common parameters common_params = self._build_image_common_params( prompt=prompt, size=size, n=n, response_format=response_format, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, seed=seed, enable_teacache=enable_teacache, background=background, output_format=output_format, generator_device=generator_device, ) # If image_path is provided, use edit endpoint if image_path: if not os.path.exists(image_path): raise FileNotFoundError(f"Image file not found: {image_path}") # Prepare multipart form data for edit files: Dict[str, Any] = {} data = common_params.copy() # Add image file files["image"] = ( os.path.basename(image_path), open(image_path, "rb"), self._get_content_type(image_path), ) # Add mask file if provided if mask_path: if not os.path.exists(mask_path): raise FileNotFoundError(f"Mask file not found: {mask_path}") files["mask"] = ( os.path.basename(mask_path), open(mask_path, "rb"), self._get_content_type(mask_path), ) # Prepare headers for multipart form data headers = { "Authorization": f"Bearer {self.api_key}", } try: response = requests.post( f"{self.base_url}/images/edits", files=files, data=data, headers=headers, timeout=300, # 5 minutes timeout for generation ) response.raise_for_status() return response.json() except requests.exceptions.RequestException as e: raise RuntimeError(f"Failed to edit image: {str(e)}") finally: # Close file handles for file_tuple in files.values(): if isinstance(file_tuple, tuple) and len(file_tuple) > 1: file_tuple[1].close() else: # Use generation endpoint - add generation-specific parameters payload = common_params.copy() if quality: payload["quality"] = quality if style: payload["style"] = style try: response = requests.post( f"{self.base_url}/images/generations", json=payload, headers=self.headers, timeout=300, # 5 minutes timeout for generation ) response.raise_for_status() return response.json() except requests.exceptions.RequestException as e: raise RuntimeError(f"Failed to generate image: {str(e)}") def generate_video( self, prompt: str, size: Optional[str] = None, width: Optional[int] = None, height: Optional[int] = None, seconds: Optional[int] = 4, fps: Optional[int] = None, num_frames: Optional[int] = None, negative_prompt: Optional[str] = None, guidance_scale: Optional[float] = None, num_inference_steps: Optional[int] = None, seed: Optional[int] = None, enable_teacache: bool = False, generator_device: Optional[str] = "cuda", input_reference: Optional[str] = None, output_path: Optional[str] = None, ) -> Dict[str, Any]: """ Generate a video using SGLang Diffusion API and wait for completion. Args: prompt: Text prompt for video generation size: Video size in format "WIDTHxHEIGHT" (e.g., "1280x720") width: Video width (used if size is not provided) height: Video height (used if size is not provided) seconds: Duration of the video in seconds fps: Frames per second num_frames: Number of frames (overrides seconds * fps if provided) negative_prompt: Negative prompt to avoid certain elements guidance_scale: Classifier-free guidance scale num_inference_steps: Number of denoising steps seed: Random seed for reproducible generation enable_teacache: Enable TEA cache acceleration generator_device: Device for random generator ("cuda" or "cpu") input_reference: Path to input reference image for image-to-video Returns: Dictionary containing completed video job information with file_path """ if not prompt: raise ValueError("Prompt cannot be empty") # Determine size if size is None: if width is not None and height is not None: size = f"{width}x{height}" else: size = "720x1280" # Prepare request payload payload: Dict[str, Any] = { "prompt": prompt, "size": size, } # Add optional parameters if seconds is not None: payload["seconds"] = seconds if fps is not None: payload["fps"] = fps if num_frames is not None: payload["num_frames"] = num_frames if negative_prompt: payload["negative_prompt"] = negative_prompt if guidance_scale is not None: payload["guidance_scale"] = guidance_scale if num_inference_steps is not None: payload["num_inference_steps"] = num_inference_steps if seed is not None and seed >= 0: payload["seed"] = seed if enable_teacache: payload["enable_teacache"] = True if generator_device: payload["generator_device"] = generator_device if input_reference: payload["input_reference"] = input_reference if output_path: payload["output_path"] = output_path try: # Create video generation job response = requests.post( f"{self.base_url}/videos", json=payload, headers=self.headers, timeout=30, ) response.raise_for_status() video_job = response.json() video_id = video_job.get("id") # Wait for completion with fixed polling poll_interval = 5 # 5 seconds max_wait_time = 3600 # 1 hour max_consecutive_errors = 5 consecutive_errors = 0 start_time = time.time() while time.time() - start_time < max_wait_time: try: status_response = requests.get( f"{self.base_url}/videos/{video_id}", headers=self.headers, timeout=30, ) status_response.raise_for_status() status = status_response.json() # Reset error counter on successful request consecutive_errors = 0 if status.get("status") == "completed": return status elif status.get("status") == "failed": error = status.get("error", {}) error_msg = ( error.get("message", "Unknown error") if error else "Unknown error" ) raise RuntimeError(f"Video generation failed: {error_msg}") except requests.exceptions.ConnectionError as e: # Connection errors - likely server is down consecutive_errors += 1 if consecutive_errors >= max_consecutive_errors: raise RuntimeError( f"Lost connection to server after {consecutive_errors} consecutive errors. " f"Server may be unavailable: {str(e)}" ) except requests.exceptions.RequestException as e: # Other network errors - continue polling but track errors consecutive_errors += 1 if consecutive_errors >= max_consecutive_errors: raise RuntimeError( f"Network error after {consecutive_errors} consecutive failures: {str(e)}" ) time.sleep(poll_interval) raise TimeoutError( f"Video generation timed out after {max_wait_time} seconds" ) except requests.exceptions.RequestException as e: raise RuntimeError(f"Failed to generate video: {str(e)}") def _build_image_common_params( self, prompt: str, size: str, n: int, response_format: str, negative_prompt: Optional[str] = None, guidance_scale: Optional[float] = None, num_inference_steps: Optional[int] = None, seed: Optional[int] = None, enable_teacache: bool = False, background: Optional[str] = None, output_format: Optional[str] = None, generator_device: Optional[str] = None, ) -> Dict[str, Any]: """ Build common parameters for both image generation and editing. Returns: Dictionary containing common parameters """ params: Dict[str, Any] = { "prompt": prompt, "size": size, "n": max(1, min(n, 10)), "response_format": response_format, } # Add optional parameters if negative_prompt: params["negative_prompt"] = negative_prompt if guidance_scale is not None: params["guidance_scale"] = guidance_scale if num_inference_steps is not None: params["num_inference_steps"] = num_inference_steps if seed is not None and seed >= 0: params["seed"] = seed if enable_teacache: params["enable_teacache"] = True if background: params["background"] = background if output_format: params["output_format"] = output_format if generator_device: params["generator_device"] = generator_device return params def _get_content_type(self, file_path: str) -> str: """Get content type based on file extension.""" ext = os.path.splitext(file_path)[1].lower() content_types = { ".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp", } return content_types.get(ext, "image/png") def decode_image_from_response( self, response_data: Dict[str, Any], index: int = 0 ) -> Image.Image: """ Decode base64 image from API response. Args: response_data: API response dictionary index: Index of the image in the response (default: 0) Returns: PIL Image object """ if "data" not in response_data or not response_data["data"]: raise ValueError("No image data in response") if index >= len(response_data["data"]): raise IndexError(f"Image index {index} out of range") image_data = response_data["data"][index] if "b64_json" not in image_data or not image_data["b64_json"]: raise ValueError("No base64 image data found") image_bytes = base64.b64decode(image_data["b64_json"]) image = Image.open(io.BytesIO(image_bytes)) # Convert to RGB if needed if image.mode != "RGB": image = image.convert("RGB") return image def set_lora( self, lora_nickname: str, lora_path: Optional[str] = None, target: str = "all", ) -> Dict[str, Any]: """ Set a LoRA adapter for the specified transformer(s). Args: lora_nickname: The nickname of the adapter (required). lora_path: Path to the LoRA adapter (local path or HF repo id). Required for the first load; optional if re-activating a cached nickname. target: Which transformer(s) to apply the LoRA to. One of: - "all": Apply to all transformers (default) - "transformer": Apply only to the primary transformer (high noise for Wan2.2) - "transformer_2": Apply only to transformer_2 (low noise for Wan2.2) - "critic": Apply only to the critic model Returns: Dictionary containing the API response with status and message """ if not lora_nickname: raise ValueError("lora_nickname cannot be empty") # Prepare request payload payload: Dict[str, Any] = { "lora_nickname": lora_nickname, "target": target, } # Add optional lora_path if provided if lora_path: payload["lora_path"] = lora_path try: response = requests.post( f"{self.base_url}/set_lora", json=payload, headers=self.headers, timeout=30, ) response.raise_for_status() return response.json() except requests.exceptions.RequestException as e: raise RuntimeError(f"Failed to set LoRA adapter: {str(e)}") def unset_lora( self, target: str = "all", ) -> Dict[str, Any]: """ Unset (unmerge) LoRA weights from the base model. Args: target: same as set_lora Returns: Dictionary containing the API response with status and message """ # Prepare request payload payload: Dict[str, Any] = { "target": target, } try: response = requests.post( f"{self.base_url}/unmerge_lora_weights", json=payload, headers=self.headers, timeout=30, ) response.raise_for_status() return response.json() except requests.exceptions.RequestException as e: raise RuntimeError(f"Failed to unset LoRA adapter: {str(e)}") if __name__ == "__main__": api = SGLDiffusionServerAPI( base_url="http://localhost:30010/v1", api_key="sk-proj-1234567890" ) model_info = api.get_model_info() print(api.get_model_info()) if model_info.get("task_type") == "T2V" or model_info.get("task_type") == "I2V": print( api.generate_video( prompt="A calico cat playing a piano on stage", num_inference_steps=1, size="480x480", ) ) else: print( api.generate_image( prompt="A calico cat playing a piano on stage", size="1024x1024" ) ) ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/__init__.py ================================================ """ ComfyUI SGLang Diffusion executors package. Provides executor classes for different model types. """ from .base import SGLDiffusionExecutor from .flux import FluxExecutor from .qwen_image import QwenImageEditExecutor, QwenImageExecutor from .zimage import ZImageExecutor __all__ = [ "SGLDiffusionExecutor", "FluxExecutor", "ZImageExecutor", "QwenImageExecutor", "QwenImageEditExecutor", ] ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/base.py ================================================ """ Base executor class for SGLang Diffusion ComfyUI integration. """ import torch class SGLDiffusionExecutor(torch.nn.Module): """Base executor class for SGLang Diffusion models in ComfyUI.""" def __init__(self, generator, model_path, model, config): super(SGLDiffusionExecutor, self).__init__() self.generator = generator self.model_path = model_path self.model = model self.dtype = config.unet_config["dtype"] self.config = config self.loras = [] @staticmethod def should_suppress_logs(timestep): """Determine if logs should be suppressed based on timestep value.""" if torch.is_tensor(timestep): return bool((timestep < 1.0).item()) return bool(timestep < 1.0) def set_lora(self, lora_nickname=None, lora_path=None, strength=None, target=None): """Set LoRA adapter using SGLang Diffusion API.""" if len(lora_nickname) > 0: self.generator.set_lora( lora_nickname=lora_nickname, lora_path=lora_path, strength=strength, target=target, ) def _unpack_latents(self, latents, height, width, channels): """Unpack latents from packed format to standard format.""" batch_size = latents.shape[0] latents = latents.view(batch_size, height // 2, width // 2, channels, 2, 2) latents = latents.permute(0, 3, 1, 4, 2, 5) latents = latents.reshape(batch_size, channels, height, width) return latents def _pack_latents(self, latents): """Pack latents from standard format to packed format.""" batch_size, num_channels_latents, height, width = latents.shape latents = latents.view( batch_size, num_channels_latents, height // 2, 2, width // 2, 2 ) latents = latents.permute(0, 2, 4, 1, 3, 5) latents = latents.reshape( batch_size, (height // 2) * (width // 2), num_channels_latents * 4 ) return latents ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py ================================================ """ Flux executor for SGLang Diffusion ComfyUI integration. """ import torch try: from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request except ImportError: print( "Error: sglang.multimodal_gen is not installed. Please install it using 'pip install sglang[diffusion]'" ) from .base import SGLDiffusionExecutor class FluxExecutor(SGLDiffusionExecutor): """Executor for Flux models in ComfyUI.""" def __init__(self, generator, model_path, model, config): super().__init__(generator, model_path, model, config) def forward(self, x, timestep, context, y=None, guidance=None, **kwargs): """Forward pass for Flux model.""" hidden_states = self._pack_latents(x) timesteps = timestep * 1000.0 encoder_hidden_states = context pooled_projections = y guidance = guidance * 1000.0 B, C, H, W = x.shape height = H * 8 width = W * 8 # Create SamplingParams sampling_params = SamplingParams.from_user_sampling_params_args( self.model_path, server_args=self.generator.server_args, prompt=" ", guidance_scale=3.5, # Flux typically uses embedded_cfg_scale=3.5 height=height, width=width, num_frames=1, num_inference_steps=1, save_output=False, suppress_logs=self.should_suppress_logs(timestep), ) # Prepare request (converts SamplingParams to Req) req = prepare_request( server_args=self.generator.server_args, sampling_params=sampling_params, ) req.latents = hidden_states # Set as [B, S, D] format directly req.timesteps = timesteps # ComfyUI's timesteps parameter req.prompt_embeds = [pooled_projections, encoder_hidden_states] # [CLIP, T5] req.raw_latent_shape = torch.tensor(hidden_states.shape, dtype=torch.long) # Set pooled_projections (required by Flux) req.pooled_embeds = [pooled_projections] # List format as per Req definition req.do_classifier_free_guidance = False req.generator = [ torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt) ] # Send request to scheduler output_batch = self.generator._send_to_scheduler_and_wait_for_response([req]) noise_pred = output_batch.noise_pred return self._unpack_latents(noise_pred, H, W, C).to(x.device) ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/qwen_image.py ================================================ """ QwenImage executor for SGLang Diffusion ComfyUI integration. """ import torch try: from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request except ImportError: print( "Error: sglang.multimodal_gen is not installed. Please install it using 'pip install sglang[diffusion]'" ) import comfy.ldm.common_dit from .base import SGLDiffusionExecutor class QwenImageExecutor(SGLDiffusionExecutor): """Executor for QwenImage models in ComfyUI.""" def __init__(self, generator, model_path, model, config): super().__init__(generator, model_path, model, config) self.patch_size = 2 def _pack_latents(self, x): """Process hidden states for QwenImage model.""" bs, c, t, h, w = x.shape patch_size = self.patch_size latents = comfy.ldm.common_dit.pad_to_patch_size( x, (1, self.patch_size, self.patch_size) ) orig_shape = latents.shape latents = latents.view( orig_shape[0], orig_shape[1], orig_shape[-3], orig_shape[-2] // 2, 2, orig_shape[-1] // 2, 2, ) latents = latents.permute(0, 2, 3, 5, 1, 4, 6) latents = latents.reshape( orig_shape[0], orig_shape[-3] * (orig_shape[-2] // 2) * (orig_shape[-1] // 2), orig_shape[1] * 4, ) return latents, orig_shape def _unpack_latents(self, latents, num_embeds, orig_shape, x): """Unpack hidden states from packed format to standard format.""" latents = latents[:, :num_embeds].view( orig_shape[0], orig_shape[-3], orig_shape[-2] // 2, orig_shape[-1] // 2, orig_shape[1], 2, 2, ) latents = latents.permute(0, 4, 1, 2, 5, 3, 6) latents = latents.reshape(orig_shape)[:, :, :, : x.shape[-2], : x.shape[-1]] return latents def forward(self, x, timestep, context, **kwargs): """Forward pass for QwenImage model.""" latents, orig_shape = self._pack_latents(x) num_embeds = latents.shape[1] height = orig_shape[-2] * 8 width = orig_shape[-1] * 8 sampling_params = SamplingParams.from_user_sampling_params_args( self.model_path, server_args=self.generator.server_args, prompt=" ", guidance_scale=1.0, height=height, width=width, num_frames=1, num_inference_steps=1, save_output=False, suppress_logs=self.should_suppress_logs(timestep), ) # Prepare request (converts SamplingParams to Req) req = prepare_request( server_args=self.generator.server_args, sampling_params=sampling_params, ) # Set ComfyUI-specific inputs directly on the Req object req.latents = latents req.timesteps = timestep * 1000.0 req.prompt_embeds = [context] req.raw_latent_shape = torch.tensor(latents.shape, dtype=torch.long) req.do_classifier_free_guidance = False req.generator = [ torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt) ] output_batch = self.generator._send_to_scheduler_and_wait_for_response([req]) noise_pred = output_batch.noise_pred return self._unpack_latents(noise_pred, num_embeds, orig_shape, x) class QwenImageEditExecutor(QwenImageExecutor): """Executor for QwenImageEdit models in ComfyUI.""" def __init__(self, generator, model_path, model, config): super().__init__(generator, model_path, model, config) def forward( self, x, timestep, context, attention_mask=None, ref_latents=None, additional_t_cond=None, transformer_options={}, **kwargs ): """Forward pass for QwenImageEdit model.""" latents, orig_shape = self._pack_latents(x) num_embeds = latents.shape[1] height = orig_shape[-2] * 8 width = orig_shape[-1] * 8 # Prepare vae_image_sizes for the condition image (ref_latents) vae_image_sizes = [] pack_ref_latents = None # TODO: sgld now don't support multiple condition images, so we only support one condition image for now. if ref_latents is not None and len(ref_latents) > 0: pack_ref_latents, orig_ref_shape = self._pack_latents(ref_latents[0]) vae_image_sizes = [(orig_ref_shape[-1], orig_ref_shape[-2])] sampling_params = SamplingParams.from_user_sampling_params_args( self.model_path, server_args=self.generator.server_args, prompt=" ", guidance_scale=1.0, image_path="", height=height, width=width, num_frames=1, num_inference_steps=1, save_output=False, suppress_logs=self.should_suppress_logs(timestep), ) # Prepare request (converts SamplingParams to Req) req = prepare_request( server_args=self.generator.server_args, sampling_params=sampling_params, ) # Set ComfyUI-specific inputs directly on the Req object req.latents = latents req.image_latent = pack_ref_latents req.timesteps = timestep * 1000.0 req.vae_image_sizes = vae_image_sizes req.prompt_embeds = [context] req.raw_latent_shape = torch.tensor(latents.shape, dtype=torch.long) req.do_classifier_free_guidance = False req.generator = [ torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt) ] output_batch = self.generator._send_to_scheduler_and_wait_for_response([req]) noise_pred = output_batch.noise_pred return self._unpack_latents(noise_pred, num_embeds, orig_shape, x) ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/zimage.py ================================================ """ ZImage executor for SGLang Diffusion ComfyUI integration. """ import torch try: from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request except ImportError: print( "Error: sglang.multimodal_gen is not installed. Please install it using 'pip install sglang[diffusion]'" ) from .base import SGLDiffusionExecutor class ZImageExecutor(SGLDiffusionExecutor): """Executor for ZImage models in ComfyUI.""" def __init__(self, generator, model_path, model, config): super().__init__(generator, model_path, model, config) def forward(self, x, timesteps, context, **kwargs): """Forward pass for ZImage model.""" B, C, H, W = x.shape height = H * 8 width = W * 8 sampling_params = SamplingParams.from_user_sampling_params_args( self.model_path, server_args=self.generator.server_args, prompt=" ", guidance_scale=1.0, height=height, width=width, num_frames=1, # For images num_inference_steps=1, # Single step for ComfyUI save_output=False, suppress_logs=self.should_suppress_logs(timesteps), ) # Prepare request (converts SamplingParams to Req) req = prepare_request( server_args=self.generator.server_args, sampling_params=sampling_params, ) latents = x.unsqueeze(2) context = context.squeeze(0) # Set ComfyUI-specific inputs directly on the Req object req.latents = latents # ComfyUI's x parameter req.timesteps = timesteps * 1000.0 # ComfyUI's timesteps parameter req.prompt_embeds = [ context ] # ComfyUI's context parameter (must be List[Tensor]) req.raw_latent_shape = torch.tensor(latents.shape, dtype=torch.long) req.do_classifier_free_guidance = False req.generator = [ torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt) ] output_batch = self.generator._send_to_scheduler_and_wait_for_response([req]) noise_pred = output_batch.noise_pred return noise_pred.permute(1, 0, 2, 3).to(x.device) ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/nodes.py ================================================ """ ComfyUI nodes for SGLang Diffusion integration. Provides nodes for connecting to SGLang Diffusion server and generating images/videos. """ import os import uuid import folder_paths import torch from .core import SGLDiffusionGenerator, SGLDiffusionServerAPI from .utils import ( convert_b64_to_tensor_image, convert_video_to_comfy_video, get_image_path, is_empty_image, ) class SGLDOptions: @classmethod def INPUT_TYPES(cls): return { "required": {}, "optional": { "model_type": ( ["auto-detect", "qwen_image", "qwen_image_edit", "flux", "lumina2"], {"default": "auto-detect"}, ), "enable_torch_compile": ( "BOOLEAN", {"default": False}, ), "num_gpus": ("INT", {"default": 1, "min": 1, "step": 1}), "tp_size": ("INT", {"default": -1, "min": -1, "step": 1}), "sp_degree": ("INT", {"default": -1, "min": -1, "step": 1}), "ulysses_degree": ( "INT", { "default": -1, "min": -1, "step": 1, }, ), "ring_degree": ( "INT", { "default": -1, "min": -1, "step": 1, }, ), "dp_size": ("INT", {"default": 1, "min": 1, "step": 1}), "dp_degree": ("INT", {"default": 1, "min": 1, "step": 1}), "enable_cfg_parallel": ( "BOOLEAN", {"default": False}, ), "attention_backend": ( "STRING", {"default": ""}, ), }, } RETURN_TYPES = ("SGLD_OPTIONS",) RETURN_NAMES = ("sgld_options",) FUNCTION = "create_options" CATEGORY = "SGLDiffusion" def create_options( self, model_type: str = "auto-detect", enable_torch_compile: bool = False, num_gpus: int = 1, tp_size: int = -1, sp_degree: int = -1, ulysses_degree: int = -1, ring_degree: int = -1, dp_size: int = 1, dp_degree: int = 1, enable_cfg_parallel: bool = False, attention_backend: str = "", ): """ Build a dictionary of SGLang Diffusion runtime options. """ # Convert -1 to None for optional parameters (matching ServerArgs defaults) ulysses_degree = None if ulysses_degree == -1 else ulysses_degree ring_degree = None if ring_degree == -1 else ring_degree attention_backend = None if attention_backend == "" else attention_backend options = { "model_type": model_type, "enable_torch_compile": enable_torch_compile, "num_gpus": num_gpus, "tp_size": tp_size, "sp_degree": sp_degree, "ulysses_degree": ulysses_degree, "ring_degree": ring_degree, "dp_size": dp_size, "dp_degree": dp_degree, "enable_cfg_parallel": enable_cfg_parallel, "attention_backend": attention_backend, } # Strip None to keep payload clean options = {k: v for k, v in options.items() if v is not None} return (options,) class SGLDLoraLoader: @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("MODEL",), "lora_name": (folder_paths.get_filename_list("loras"),), "strength_model": ( "FLOAT", {"default": 1.0, "min": 0, "max": 10, "step": 0.01}, ), "nickname": ("STRING", {"default": ""}), "target": ( ["all", "transformer", "transformer_2", "critic"], {"default": "all"}, ), }, } RETURN_TYPES = ("MODEL",) FUNCTION = "load_lora" CATEGORY = "SGLDiffusion" def load_lora( self, model, lora_name, strength_model=1.0, nickname="", target="all" ): """Load LoRA adapter using SGLang Diffusion API.""" lora_path = folder_paths.get_full_path("loras", lora_name) assert model is not None bi = model.clone() nickname = nickname if nickname != "" else str("lora" + str(uuid.uuid4())) # set lora in the model bi.patches[nickname] = (lora_path, strength_model, target) # prepare input for the SGLang Diffusion API lora_input = { "lora_nickname": [], "lora_path": [], "strength": [], "target": [], } for nickname, lora_info in bi.patches.items(): lora_input["lora_nickname"].append(nickname) lora_input["lora_path"].append(lora_info[0]) lora_input["strength"].append(lora_info[1]) lora_input["target"].append(lora_info[2]) # call the SGLang Diffusion API model.model.diffusion_model.set_lora(**lora_input) return (model,) class SGLDUNETLoader: def __init__(self): self.generator = SGLDiffusionGenerator() @classmethod def INPUT_TYPES(s): return { "required": { "unet_name": (folder_paths.get_filename_list("diffusion_models"),), "weight_dtype": (["default", "fp8_e4m3fn", "fp8_e5m2"],), }, "optional": { "sgld_options": ("SGLD_OPTIONS",), }, } RETURN_TYPES = ("MODEL",) FUNCTION = "load_unet" CATEGORY = "SGLDiffusion" def load_unet(self, unet_name, weight_dtype, sgld_options: dict = None): model_options = {} if weight_dtype == "fp8_e4m3fn": model_options["dtype"] = torch.float8_e4m3fn elif weight_dtype == "fp8_e5m2": model_options["dtype"] = torch.float8_e5m2 unet_path = folder_paths.get_full_path("diffusion_models", unet_name) model = self.generator.load_model( unet_path, model_options=model_options, sgld_options=sgld_options ) return (model,) class SGLDiffusionServerModel: """Node to load and manage SGLang Diffusion server connection.""" @classmethod def INPUT_TYPES(cls): return { "required": { "base_url": ( "STRING", { "default": "http://localhost:3000/v1", "multiline": False, }, ), "api_key": ( "STRING", { "default": "sk-proj-1234567890", "multiline": False, }, ), } } RETURN_TYPES = ("SGLD_CLIENT", "STRING") RETURN_NAMES = ("sgld_client", "model_info") FUNCTION = "load_server" CATEGORY = "SGLDiffusion" def load_server(self, base_url: str, api_key: str): """Initialize OpenAI client for SGLang Diffusion server.""" client = SGLDiffusionServerAPI(base_url=base_url, api_key=api_key) try: model_info = client.get_model_info() # Format model_info as a readable string info_lines = ["=== SGLDiffusion Model Info ==="] for key, value in model_info.items(): info_lines.append(f"{key}: {value}") model_info_str = "\n".join(info_lines) except Exception as e: model_info_str = f"Failed to get model info: {str(e)}" return (client, model_info_str) class SGLDiffusionGenerateImage: """Node to generate images using SGLang Diffusion.""" @classmethod def INPUT_TYPES(cls): return { "required": { "sgld_client": ("SGLD_CLIENT",), "positive_prompt": ( "STRING", { "default": "", "tooltip": "Text prompt for image generation", }, ), }, "optional": { "negative_prompt": ( "STRING", { "default": "", "tooltip": "Negative prompt to avoid certain elements", }, ), "image": ( "IMAGE", { "default": None, "tooltip": "input image to use for editing", }, ), "seed": ( "INT", { "default": 1024, "min": -1, "max": 2**32 - 1, }, ), "steps": ( "INT", { "default": 6, "min": 1, "max": 100, "step": 1, }, ), "cfg": ( "FLOAT", { "default": 7.0, "min": 1.0, "max": 20.0, "step": 0.1, }, ), "width": ( "INT", { "default": 1024, "min": 256, "max": 4096, "step": 64, }, ), "height": ( "INT", { "default": 1024, "min": 256, "max": 4096, "step": 64, }, ), "enable_teacache": ( "BOOLEAN", { "default": False, }, ), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) FUNCTION = "generate_image" CATEGORY = "SGLDiffusion" OUTPUT_NODE = False def generate_image( self, sgld_client: SGLDiffusionServerAPI, positive_prompt: str, negative_prompt: str = "", image: torch.Tensor = None, seed: int = 1024, steps: int = 6, cfg: float = 7.0, width: int = 1024, height: int = 1024, enable_teacache: bool = False, ): """Generate image using SGLang Diffusion API.""" if not positive_prompt: raise ValueError("Prompt cannot be empty") size = f"{width}x{height}" # Prepare request parameters request_params = { "prompt": positive_prompt, "size": size, "response_format": "b64_json", } # Add optional parameters if provided if negative_prompt: request_params["negative_prompt"] = negative_prompt if cfg is not None: request_params["guidance_scale"] = cfg if steps is not None: request_params["num_inference_steps"] = steps if seed is not None and seed >= 0: request_params["seed"] = seed if enable_teacache: request_params["enable_teacache"] = True if image is not None: # If the image is empty, use the size of the image to generate the image if is_empty_image(image): width, height = image.shape[2], image.shape[1] size = f"{width}x{height}" request_params["size"] = size else: request_params["image_path"] = get_image_path(image) # Call API try: response = sgld_client.generate_image(**request_params) except Exception as e: raise RuntimeError(f"Failed to generate image: {str(e)}") # Decode base64 image if not response["data"] or not response["data"][0]["b64_json"]: raise RuntimeError("No image data in response") image_data = response["data"][0]["b64_json"] image = convert_b64_to_tensor_image(image_data) return (image,) class SGLDiffusionGenerateVideo: """Node to generate videos using SGLang Diffusion.""" @classmethod def INPUT_TYPES(cls): return { "required": { "sgld_client": ("SGLD_CLIENT",), "positive_prompt": ( "STRING", { "default": "", "tooltip": "Text prompt for video generation", }, ), }, "optional": { "negative_prompt": ( "STRING", { "default": "", "tooltip": "Negative prompt to avoid certain elements", }, ), "image": ( "IMAGE", { "default": None, "tooltip": "input image to use for image-to-video", }, ), "seed": ( "INT", { "default": 1024, "min": -1, "max": 2**32 - 1, }, ), "steps": ( "INT", { "default": 6, "min": 1, "max": 100, "step": 1, }, ), "cfg": ( "FLOAT", { "default": 7.0, "min": 1.0, "max": 20.0, "step": 0.1, }, ), "width": ( "INT", { "default": 1280, "min": 256, "max": 4096, "step": 1, }, ), "height": ( "INT", { "default": 720, "min": 256, "max": 4096, "step": 1, }, ), "num_frames": ( "INT", { "default": 120, "min": 1, "max": 1000, "step": 1, }, ), "fps": ( "INT", { "default": 24, "min": 1, "max": 60, "step": 1, }, ), "seconds": ( "INT", { "default": 5, "min": 1, "max": 60, "step": 1, }, ), "enable_teacache": ( "BOOLEAN", { "default": False, }, ), }, } RETURN_TYPES = ("VIDEO", "STRING") RETURN_NAMES = ("video", "video_path") FUNCTION = "generate_video" CATEGORY = "SGLDiffusion" OUTPUT_NODE = False def generate_video( self, sgld_client: SGLDiffusionServerAPI, positive_prompt: str, negative_prompt: str = "", image: torch.Tensor = None, seed: int = 1024, steps: int = 6, cfg: float = 7.0, width: int = 1280, height: int = 720, num_frames: int = 120, fps: int = 24, seconds: int = 5, enable_teacache: bool = False, ): """Generate video using SGLang Diffusion API.""" if not positive_prompt: raise ValueError("Prompt cannot be empty") size = f"{width}x{height}" output_dir = folder_paths.get_temp_directory() # Prepare request parameters request_params = { "prompt": positive_prompt, "size": size, "seconds": seconds, "fps": fps, "output_path": output_dir, } # Add optional parameters if provided if negative_prompt: request_params["negative_prompt"] = negative_prompt if cfg is not None: request_params["guidance_scale"] = cfg if steps is not None: request_params["num_inference_steps"] = steps if seed is not None and seed >= 0: request_params["seed"] = seed if enable_teacache: request_params["enable_teacache"] = True if num_frames is not None: request_params["num_frames"] = num_frames if image is not None: # If the image is empty, use the size of the image to generate the video if is_empty_image(image): width, height = image.shape[2], image.shape[1] size = f"{width}x{height}" request_params["size"] = size else: request_params["input_reference"] = get_image_path(image) # Call API try: response = sgld_client.generate_video(**request_params) video_path = response.get("file_path", "") video = convert_video_to_comfy_video(video_path, height, width) except Exception as e: raise RuntimeError(f"Failed to generate video: {str(e)}") return (video, video_path) class SGLDiffusionServerSetLora: """Node to set LoRA adapter for SGLang Diffusion server.""" @classmethod def INPUT_TYPES(cls): return { "required": { "sgld_client": ("SGLD_CLIENT",), "lora_name": ( "STRING", { "default": "", "tooltip": "The name of the LoRA adapter", }, ), }, "optional": { "lora_nickname": ( "STRING", { "default": "", "tooltip": "The nickname of the LoRA adapter", }, ), "target": ( [ "all", "transformer", "transformer_2", "critic", ], { "default": "all", "tooltip": "Which transformer(s) to apply the LoRA to", }, ), }, } RETURN_TYPES = ("SGLD_CLIENT",) RETURN_NAMES = ("sgld_client",) FUNCTION = "set_lora" CATEGORY = "SGLDiffusion" OUTPUT_NODE = False def set_lora( self, sgld_client: SGLDiffusionServerAPI, lora_name: str = "", lora_nickname: str = "", target: str = "all", ): """Set LoRA adapter using SGLang Diffusion API.""" if lora_nickname == "": lora_nickname = os.path.splitext(lora_name)[0] # Prepare request parameters request_params = { "lora_nickname": lora_nickname, "lora_path": lora_name, "target": target, } # Call API try: response = sgld_client.set_lora(**request_params) return (sgld_client,) except Exception as e: raise RuntimeError(f"Failed to set LoRA adapter: {str(e)}") class SGLDiffusionServerUnsetLora: """Node to unset LoRA adapter for SGLang Diffusion server.""" @classmethod def INPUT_TYPES(cls): return { "required": { "sgld_client": ("SGLD_CLIENT",), }, "optional": { "target": ( [ "all", "transformer", "transformer_2", "critic", ], { "default": "all", "tooltip": "Which transformer(s) to unset the LoRA from", }, ), }, } RETURN_TYPES = ("SGLD_CLIENT",) RETURN_NAMES = ("sgld_client",) FUNCTION = "unset_lora" CATEGORY = "SGLDiffusion" OUTPUT_NODE = False def unset_lora( self, sgld_client: SGLDiffusionServerAPI, target: str = "all", ): """Unset LoRA adapter using SGLang Diffusion API.""" try: response = sgld_client.unset_lora(target=target) return (sgld_client,) except Exception as e: raise RuntimeError(f"Failed to unset LoRA adapter: {str(e)}") # Register nodes NODE_CLASS_MAPPINGS = { "SGLDiffusionServerModel": SGLDiffusionServerModel, "SGLDiffusionGenerateImage": SGLDiffusionGenerateImage, "SGLDiffusionGenerateVideo": SGLDiffusionGenerateVideo, "SGLDiffusionServerSetLora": SGLDiffusionServerSetLora, "SGLDiffusionServerUnsetLora": SGLDiffusionServerUnsetLora, "SGLDUNETLoader": SGLDUNETLoader, "SGLDOptions": SGLDOptions, "SGLDLoraLoader": SGLDLoraLoader, } NODE_DISPLAY_NAME_MAPPINGS = { "SGLDiffusionServerModel": "SGLDiffusion Server Model", "SGLDiffusionGenerateImage": "SGLDiffusion Generate Image", "SGLDiffusionGenerateVideo": "SGLDiffusion Generate Video", "SGLDiffusionServerSetLora": "SGLDiffusion Server Set LoRA", "SGLDiffusionServerUnsetLora": "SGLDiffusion Server Unset LoRA", "SGLDUNETLoader": "SGLDiffusion UNET Loader", "SGLDOptions": "SGLDiffusion Options", "SGLDLoraLoader": "SGLDiffusion LoRA Loader", } ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/README.md ================================================ # ComfyUI SGLDiffusion Pipeline Tests This directory contains tests for each ComfyUI pipeline integration. ## Test Files - `test_zimage_pipeline.py` - Tests for ComfyUIZImagePipeline - `test_flux_pipeline.py` - Tests for ComfyUIFluxPipeline - `test_qwen_image_pipeline.py` - Tests for ComfyUIQwenImagePipeline - `test_qwen_image_edit_pipeline.py` - Tests for ComfyUIQwenImageEditPipeline (I2I/edit mode) ## Running Tests ### Run all tests ```bash pytest python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/ -v -s ``` ### Run a specific test file ```bash pytest python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py -v -s ``` ## Environment Variables You can configure model paths via environment variables. Model paths support two formats: - **Safetensors file**: Path to a single `.safetensors` file (e.g., `/path/to/model.safetensors`) - **Diffusers format**: HuggingFace model ID or local diffusers directory (e.g., `Tongyi-MAI/Z-Image-Turbo`) Environment variables: - `SGLANG_TEST_ZIMAGE_MODEL_PATH` - Path to ZImage model (default: `Tongyi-MAI/Z-Image-Turbo`) - `SGLANG_TEST_FLUX_MODEL_PATH` - Path to Flux model (default: `black-forest-labs/FLUX.1-dev`) - `SGLANG_TEST_QWEN_IMAGE_MODEL_PATH` - Path to QwenImage model (default: `Qwen/Qwen-Image`) - `SGLANG_TEST_QWEN_IMAGE_EDIT_MODEL_PATH` - Path to QwenImageEdit model (default: `Qwen/Qwen-Image-Edit-2511`) Examples: ```bash # Using HuggingFace model ID (diffusers format) export SGLANG_TEST_ZIMAGE_MODEL_PATH="Tongyi-MAI/Z-Image-Turbo" pytest python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py -v -s # Using safetensors file export SGLANG_TEST_ZIMAGE_MODEL_PATH="/path/to/z_image_turbo_bf16.safetensors" pytest python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py -v -s ``` ## Test Structure Each test file follows a similar structure: 1. **Setup**: Creates a `DiffGenerator` with the appropriate pipeline class 2. **Input Preparation**: Creates dummy tensors for latents, timesteps, and embeddings 3. **Request Preparation**: Uses `prepare_request` to convert `SamplingParams` to `Req` 4. **ComfyUI Inputs**: Sets ComfyUI-specific inputs directly on the `Req` object 5. **Execution**: Sends request to scheduler and waits for response 6. **Validation**: Checks that `noise_pred` is retrieved from `OutputBatch` ## Notes - These tests use `comfyui_mode=True` to enable ComfyUI-specific behavior - Tests use pre-processed inputs (latents, timesteps, embeddings) as ComfyUI would provide - The tests verify that `noise_pred` can be retrieved from the `OutputBatch` after processing - All tests use dummy/ones tensors for simplicity - in production, these would be actual model outputs ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/__init__.py ================================================ """ Test suite for ComfyUI SGLDiffusion pipelines. This package contains tests for each ComfyUI pipeline integration: - ZImagePipeline - FluxPipeline - QwenImagePipeline - QwenImageEditPipeline """ ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_flux_pipeline.py ================================================ """Test for ComfyUIFluxPipeline with pass-through scheduler.""" import os import pytest import torch from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request def test_comfyui_flux_pipeline_direct() -> None: """Test ComfyUIFluxPipeline with custom inputs.""" model_path = os.environ.get( "SGLANG_TEST_FLUX_MODEL_PATH", "black-forest-labs/FLUX.1-dev", # Supports both safetensors file and diffusers format ) generator = DiffGenerator.from_pretrained( model_path=model_path, pipeline_class_name="ComfyUIFluxPipeline", num_gpus=2, comfyui_mode=True, ) batch_size = 1 hidden_states_seq_len = 3600 hidden_states_dim = 64 height = 1280 width = 720 encoder_seq_len = 512 encoder_dim = 4096 pooled_dim = 768 hidden_states = torch.ones( batch_size, hidden_states_seq_len, hidden_states_dim, device="cuda", dtype=torch.bfloat16, ) encoder_hidden_states = torch.ones( batch_size, encoder_seq_len, encoder_dim, device="cuda", dtype=torch.bfloat16, ) pooled_projections = torch.ones( batch_size, pooled_dim, device="cuda", dtype=torch.bfloat16, ) timesteps = torch.tensor([1000], dtype=torch.long, device="cuda") sampling_params = SamplingParams.from_user_sampling_params_args( generator.server_args.model_path, server_args=generator.server_args, prompt="a beautiful girl", height=height, width=width, num_frames=1, num_inference_steps=1, save_output=True, return_trajectory_latents=True, ) req = prepare_request( server_args=generator.server_args, sampling_params=sampling_params, ) req.latents = hidden_states req.timesteps = timesteps req.raw_latent_shape = torch.tensor(hidden_states.shape, dtype=torch.long) clip_dim = 768 dummy_clip_embedding = torch.zeros( batch_size, 77, clip_dim, device="cuda", dtype=torch.bfloat16, ) req.prompt_embeds = [pooled_projections, encoder_hidden_states] if req.guidance_scale > 1.0: dummy_neg_clip_embedding = torch.zeros( batch_size, 77, clip_dim, device="cuda", dtype=torch.bfloat16, ) negative_encoder_hidden_states = torch.ones( batch_size, encoder_seq_len, encoder_dim, device="cuda", dtype=torch.bfloat16, ) req.negative_prompt_embeds = [ dummy_neg_clip_embedding, negative_encoder_hidden_states, ] else: req.negative_prompt_embeds = None req.pooled_embeds = [pooled_projections] req.neg_pooled_embeds = [] if ( req.guidance_scale > 1.0 and req.negative_prompt_embeds is not None and len(req.negative_prompt_embeds) > 0 ): req.do_classifier_free_guidance = True else: req.do_classifier_free_guidance = False if req.seed is not None: generator_device = req.generator_device device_str = "cuda" if generator_device == "cuda" else "cpu" req.generator = [ torch.Generator(device_str).manual_seed(req.seed + i) for i in range(req.num_outputs_per_prompt) ] else: req.generator = [ torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt) ] output_batch = generator._send_to_scheduler_and_wait_for_response([req]) noise_pred = output_batch.noise_pred assert noise_pred is not None, "noise_pred should not be None in OutputBatch" assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor" assert ( noise_pred.device.type == "cuda" ), f"noise_pred should be on cuda, got {noise_pred.device}" assert ( noise_pred.dtype == torch.bfloat16 ), f"noise_pred should be bfloat16, got {noise_pred.dtype}" print(f"✓ Successfully retrieved noise_pred from OutputBatch!") print(f" noise_pred shape: {noise_pred.shape}") print(f" noise_pred dtype: {noise_pred.dtype}") print(f" noise_pred device: {noise_pred.device}") latents = output_batch.output if output_batch.output is not None else req.latents assert latents is not None, "latents should not be None" print(f"latents.shape: {latents.shape}") if __name__ == "__main__": pytest.main([__file__, "-v"]) ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_edit_pipeline.py ================================================ """Test for ComfyUIQwenImageEditPipeline with pass-through scheduler (I2I/edit mode).""" import os import pytest import torch from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request def test_comfyui_qwen_image_edit_pipeline_direct() -> None: """Test ComfyUIQwenImageEditPipeline with edit mode (I2I) and custom inputs.""" model_path = os.environ.get( "SGLANG_TEST_QWEN_IMAGE_EDIT_MODEL_PATH", "Qwen/Qwen-Image-Edit-2511", # Supports both safetensors file and diffusers format ) generator = DiffGenerator.from_pretrained( model_path=model_path, pipeline_class_name="ComfyUIQwenImageEditPipeline", num_gpus=1, comfyui_mode=True, dit_layerwise_offload=False, ) batch_size = 1 noisy_image_seq_len = 3600 hidden_states_dim = 64 condition_image_seq_len = 6889 condition_image_dim = 64 encoder_seq_len = 45 encoder_dim = 3584 height = 720 width = 1280 vae_scale_factor = 8 condition_height_latent = 1328 // vae_scale_factor condition_width_latent = 1328 // vae_scale_factor noisy_image_latents = torch.ones( batch_size, noisy_image_seq_len, hidden_states_dim, device="cuda", dtype=torch.bfloat16, ) condition_image_latents = torch.ones( batch_size, condition_image_seq_len, condition_image_dim, device="cuda", dtype=torch.bfloat16, ) encoder_hidden_states = torch.ones( batch_size, encoder_seq_len, encoder_dim, device="cuda", dtype=torch.bfloat16, ) timesteps = torch.tensor([1000], dtype=torch.long, device="cuda") sampling_params = SamplingParams.from_user_sampling_params_args( generator.server_args.model_path, server_args=generator.server_args, prompt=" ", guidance_scale=1.0, height=height, width=width, image_path="", num_frames=1, num_inference_steps=1, seed=42, save_output=False, return_frames=False, ) req = prepare_request( server_args=generator.server_args, sampling_params=sampling_params, ) req.latents = noisy_image_latents req.image_latent = condition_image_latents req.timesteps = timesteps req.prompt_embeds = [encoder_hidden_states] req.negative_prompt_embeds = None req.vae_image_sizes = [(condition_width_latent, condition_height_latent)] req.raw_latent_shape = torch.tensor(noisy_image_latents.shape, dtype=torch.long) if req.guidance_scale > 1.0 and req.negative_prompt_embeds is not None: req.do_classifier_free_guidance = True else: req.do_classifier_free_guidance = False if req.seed is not None: generator_device = req.generator_device device_str = "cpu" if generator_device == "cpu" else "cuda" req.generator = [ torch.Generator(device_str).manual_seed(req.seed + i) for i in range(req.num_outputs_per_prompt) ] else: req.generator = [ torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt) ] output_batch = generator._send_to_scheduler_and_wait_for_response([req]) noise_pred = output_batch.noise_pred assert noise_pred is not None, "noise_pred should not be None in OutputBatch" assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor" assert ( noise_pred.device.type == "cuda" ), f"noise_pred should be on cuda, got {noise_pred.device}" assert ( noise_pred.dtype == torch.bfloat16 ), f"noise_pred should be bfloat16, got {noise_pred.dtype}" print(f"✓ Successfully retrieved noise_pred from OutputBatch (Edit Mode)!") print(f" noise_pred shape: {noise_pred.shape}") print(f" noise_pred dtype: {noise_pred.dtype}") print(f" noise_pred device: {noise_pred.device}") latents = output_batch.output if output_batch.output is not None else req.latents assert latents is not None, "latents should not be None" if __name__ == "__main__": pytest.main([__file__, "-v"]) ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_pipeline.py ================================================ """Test for ComfyUIQwenImagePipeline with pass-through scheduler.""" import os import pytest import torch from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request def test_comfyui_qwen_image_pipeline_direct() -> None: """Test ComfyUIQwenImagePipeline with custom inputs.""" model_path = os.environ.get( "SGLANG_TEST_QWEN_IMAGE_MODEL_PATH", "Qwen/Qwen-Image", # Supports both safetensors file and diffusers format ) generator = DiffGenerator.from_pretrained( model_path=model_path, pipeline_class_name="ComfyUIQwenImagePipeline", num_gpus=2, comfyui_mode=True, dit_layerwise_offload=False, ) batch_size = 1 hidden_states_seq_len = 6889 hidden_states_dim = 64 encoder_seq_len = 45 encoder_dim = 3584 height = 1328 width = 1328 dtype = torch.bfloat16 hidden_states = torch.ones( batch_size, hidden_states_seq_len, hidden_states_dim, device="cuda", dtype=dtype, ) encoder_hidden_states = torch.ones( batch_size, encoder_seq_len, encoder_dim, device="cuda", dtype=torch.bfloat16, ) timesteps = torch.tensor([1000], dtype=torch.long, device="cuda") sampling_params = SamplingParams.from_user_sampling_params_args( generator.server_args.model_path, server_args=generator.server_args, prompt=" ", guidance_scale=3.0, height=height, width=width, num_frames=1, num_inference_steps=1, seed=42, save_output=False, return_frames=False, ) req = prepare_request( server_args=generator.server_args, sampling_params=sampling_params, ) req.latents = hidden_states req.timesteps = timesteps req.prompt_embeds = [encoder_hidden_states] req.negative_prompt_embeds = [encoder_hidden_states] req.raw_latent_shape = torch.tensor(hidden_states.shape, dtype=torch.long) if req.guidance_scale > 1.0 and req.negative_prompt_embeds is not None: req.do_classifier_free_guidance = True else: req.do_classifier_free_guidance = False if req.seed is not None: generator_device = req.generator_device device_str = "cpu" if generator_device == "cpu" else "cuda" req.generator = [ torch.Generator(device_str).manual_seed(req.seed + i) for i in range(req.num_outputs_per_prompt) ] else: req.generator = [ torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt) ] output_batch = generator._send_to_scheduler_and_wait_for_response([req]) noise_pred = output_batch.noise_pred assert noise_pred is not None, "noise_pred should not be None in OutputBatch" assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor" assert ( noise_pred.device.type == "cuda" ), f"noise_pred should be on cuda, got {noise_pred.device}" assert ( noise_pred.dtype == torch.bfloat16 ), f"noise_pred should be bfloat16, got {noise_pred.dtype}" print(f"✓ Successfully retrieved noise_pred from OutputBatch!") print(f" noise_pred shape: {noise_pred.shape}") print(f" noise_pred dtype: {noise_pred.dtype}") print(f" noise_pred device: {noise_pred.device}") latents = output_batch.output if output_batch.output is not None else req.latents assert latents is not None, "latents should not be None" if __name__ == "__main__": pytest.main([__file__, "-v"]) ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py ================================================ """Test for ComfyUIZImagePipeline with pass-through scheduler.""" import os import pytest import torch from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request def test_comfyui_zimage_pipeline_direct() -> None: """Test ComfyUIZImagePipeline with custom inputs.""" model_path = os.environ.get( "SGLANG_TEST_ZIMAGE_MODEL_PATH", "Tongyi-MAI/Z-Image-Turbo", # Supports both safetensors file and diffusers format ) generator = DiffGenerator.from_pretrained( model_path=model_path, pipeline_class_name="ComfyUIZImagePipeline", num_gpus=1, sp_degree=1, comfyui_mode=True, ) batch_size = 1 num_channels = 16 num_frames = 1 height = 720 width = 1280 latent_height = height // 8 latent_width = width // 8 latents = torch.ones( batch_size, num_channels, num_frames, latent_height, latent_width, device="cuda", dtype=torch.bfloat16, ) timesteps = torch.tensor([1000], dtype=torch.long, device="cuda") context_seq_len = 19 context_dim = 2560 context = torch.ones( context_seq_len, context_dim, device="cuda", dtype=torch.bfloat16, ) sampling_params = SamplingParams.from_user_sampling_params_args( generator.server_args.model_path, server_args=generator.server_args, prompt="a beautiful girl", guidance_scale=1.0, height=height, width=width, num_frames=1, num_inference_steps=1, seed=42, save_output=False, return_frames=False, ) req = prepare_request( server_args=generator.server_args, sampling_params=sampling_params, ) req.latents = latents req.timesteps = timesteps req.prompt_embeds = [context] req.negative_prompt_embeds = None req.raw_latent_shape = torch.tensor(latents.shape, dtype=torch.long) if req.guidance_scale > 1.0 and req.negative_prompt_embeds is not None: req.do_classifier_free_guidance = True else: req.do_classifier_free_guidance = False if req.seed is not None: generator_device = req.generator_device device_str = "cpu" if generator_device == "cpu" else "cuda" req.generator = [ torch.Generator(device_str).manual_seed(req.seed + i) for i in range(req.num_outputs_per_prompt) ] else: req.generator = [ torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt) ] output_batch = generator._send_to_scheduler_and_wait_for_response([req]) noise_pred = output_batch.noise_pred assert noise_pred is not None, "noise_pred should not be None in OutputBatch" assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor" assert ( noise_pred.device.type == "cuda" ), f"noise_pred should be on cuda, got {noise_pred.device}" assert ( noise_pred.dtype == torch.bfloat16 ), f"noise_pred should be bfloat16, got {noise_pred.dtype}" print(f"✓ Successfully retrieved noise_pred from OutputBatch!") print(f" noise_pred shape: {noise_pred.shape}") print(f" noise_pred dtype: {noise_pred.dtype}") print(f" noise_pred device: {noise_pred.device}") latents = output_batch.output if output_batch.output is not None else req.latents assert latents is not None, "latents should not be None" if __name__ == "__main__": pytest.main([__file__, "-v"]) ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/utils.py ================================================ import base64 import io import os import shutil import time import uuid import folder_paths import numpy as np import torch from comfy_api.input import VideoInput from PIL import Image def _ensure_dir(path: str) -> None: os.makedirs(path, exist_ok=True) def _to_numpy_image(image: torch.Tensor) -> np.ndarray: """Convert ComfyUI image tensor to uint8 numpy array (H, W, C).""" if image.dim() == 4: image = image[0] if image.dim() == 3 and image.shape[0] in (1, 3, 4): image = image.permute(1, 2, 0) elif image.dim() == 2: image = image.unsqueeze(-1) np_img = image.detach().cpu().numpy() np_img = np.clip(np_img, 0.0, 1.0) np_img = (np_img * 255).astype(np.uint8) if np_img.shape[-1] == 1: np_img = np.repeat(np_img, 3, axis=-1) return np_img def _to_hwc_tensor(image: torch.Tensor) -> torch.Tensor: """Convert ComfyUI image tensor to HWC format (normalized [0, 1]).""" img = image.clone() if img.dim() == 4: img = img[0] if img.dim() == 3 and img.shape[0] in (1, 3, 4): img = img.permute(1, 2, 0) elif img.dim() == 2: img = img.unsqueeze(-1) img = torch.clamp(img, 0.0, 1.0) if img.shape[-1] == 1: img = img.repeat(1, 1, 3) return img def is_empty_image(image: torch.Tensor, tolerance: float = 1e-6) -> bool: """ Check if the input image is an empty/solid color image (like ComfyUI's empty image). Args: image: Input tensor image in ComfyUI format (BCHW, CHW, HWC, etc.) tolerance: Tolerance for floating point comparison (default: 1e-6) Returns: True if the image is empty (all pixels have same color), False otherwise """ if image is None: return True # Convert to HWC format img_hwc = _to_hwc_tensor(image) # Get the first pixel's RGB values first_pixel = img_hwc[0, 0, :] h, w, c = img_hwc.shape pixels = img_hwc.reshape(-1, c) diff = torch.abs(pixels - first_pixel) max_diff = torch.max(diff) return max_diff.item() <= tolerance def get_image_path(image: torch.Tensor) -> str: """ Save tensor image to ComfyUI temp directory as PNG and return the path. """ temp_dir = folder_paths.get_temp_directory() # Build file name ts = time.strftime("%Y%m%d-%H%M%S") unique = uuid.uuid4().hex[:8] file_name = f"sgl_output_{ts}_{unique}.png" file_path = os.path.join(temp_dir, file_name) # Save image np_img = _to_numpy_image(image) img = Image.fromarray(np_img) img.save(file_path, format="PNG") return file_path def convert_b64_to_tensor_image(b64_image: str) -> torch.Tensor: """ Convert base64 encoded image to ComfyUI IMAGE format (torch.Tensor). Args: b64_image: Base64 encoded image string Returns: torch.Tensor with shape [batch_size, height, width, channels] (BHWC format), values normalized to [0, 1] range, RGB format (3 channels) """ # Decode base64 image_bytes = base64.b64decode(b64_image) # Open image and convert to RGB pil_image = Image.open(io.BytesIO(image_bytes)) if pil_image.mode != "RGB": pil_image = pil_image.convert("RGB") # Convert to numpy array and normalize to [0, 1] image_array = np.array(pil_image).astype(np.float32) / 255.0 # Add batch dimension: [height, width, channels] -> [1, height, width, channels] image_array = image_array[np.newaxis, ...] # Convert to torch.Tensor tensor_image = torch.from_numpy(image_array) return tensor_image class SGLDVideoInput(VideoInput): def __init__(self, video_path: str, height: int, width: int): super().__init__() self.video_path = video_path self.height = height self.width = width def get_dimensions(self) -> tuple[int, int]: """ Returns the dimensions of the video input. Returns: Tuple of (width, height) """ return self.width, self.height def get_components(self): """ Returns the components of the video input. This is required by the VideoInput abstract base class. """ return [self.video_path] def save_to(self, path: str, format=None, codec=None, metadata=None): """ Abstract method to save the video input to a file. """ save_path = path # Copy video file from video_path to save_path if os.path.exists(self.video_path): # Ensure destination directory exists save_dir = os.path.dirname(save_path) if save_dir: os.makedirs(save_dir, exist_ok=True) shutil.copy2(self.video_path, save_path) def convert_video_to_comfy_video( video_path: str, height: int, width: int ) -> VideoInput: """ Convert video to ComfyUI VIDEO format (VideoInput). """ video_input = SGLDVideoInput(video_path, height, width) return video_input ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/flux_sgld_sp.json ================================================ { "8": { "inputs": { "samples": [ "40", 0 ], "vae": [ "10", 0 ] }, "class_type": "VAEDecode", "_meta": { "title": "VAE Decode" } }, "10": { "inputs": { "vae_name": "ae.safetensors" }, "class_type": "VAELoader", "_meta": { "title": "Load VAE" } }, "11": { "inputs": { "clip_name1": "t5xxl_fp16.safetensors", "clip_name2": "clip_l.safetensors", "type": "flux", "device": "default" }, "class_type": "DualCLIPLoader", "_meta": { "title": "DualCLIPLoader" } }, "17": { "inputs": { "scheduler": "normal", "steps": 25, "denoise": 1, "model": [ "46", 0 ] }, "class_type": "BasicScheduler", "_meta": { "title": "BasicScheduler" } }, "38": { "inputs": { "model": [ "46", 0 ], "conditioning": [ "42", 0 ] }, "class_type": "BasicGuider", "_meta": { "title": "BasicGuider" } }, "39": { "inputs": { "filename_prefix": "ComfyUI", "images": [ "8", 0 ] }, "class_type": "SaveImage", "_meta": { "title": "Save Image" } }, "40": { "inputs": { "noise": [ "45", 0 ], "guider": [ "38", 0 ], "sampler": [ "47", 0 ], "sigmas": [ "17", 0 ], "latent_image": [ "44", 0 ] }, "class_type": "SamplerCustomAdvanced", "_meta": { "title": "SamplerCustomAdvanced" } }, "42": { "inputs": { "guidance": 3.5, "conditioning": [ "43", 0 ] }, "class_type": "FluxGuidance", "_meta": { "title": "FluxGuidance" } }, "43": { "inputs": { "text": "beautiful photography of a gonger haired artist with Lots of Colorful coloursplashes in face and pn her hands, she is natural, having her hair in a casual bun, looking happily into camera, cinematic,", "speak_and_recognation": { "__value__": [ false, true ] }, "clip": [ "11", 0 ] }, "class_type": "CLIPTextEncode", "_meta": { "title": "CLIP Text Encode (Prompt)" } }, "44": { "inputs": { "width": 1024, "height": 1024, "batch_size": 1 }, "class_type": "EmptySD3LatentImage", "_meta": { "title": "EmptySD3LatentImage" } }, "45": { "inputs": { "noise_seed": 747172083610812 }, "class_type": "RandomNoise", "_meta": { "title": "RandomNoise" } }, "46": { "inputs": { "max_shift": 1.15, "base_shift": 0.5, "width": 1024, "height": 1024, "model": [ "51", 0 ] }, "class_type": "ModelSamplingFlux", "_meta": { "title": "ModelSamplingFlux" } }, "47": { "inputs": { "sampler_name": "euler" }, "class_type": "KSamplerSelect", "_meta": { "title": "KSamplerSelect" } }, "51": { "inputs": { "unet_name": "flux1-dev.safetensors", "weight_dtype": "default", "sgld_options": [ "52", 0 ] }, "class_type": "SGLDUNETLoader", "_meta": { "title": "SGLDiffusion UNET Loader" } }, "52": { "inputs": { "model_type": "auto-detect", "enable_torch_compile": false, "num_gpus": 2, "tp_size": -1, "sp_degree": -1, "ulysses_degree": -1, "ring_degree": -1, "dp_size": 1, "dp_degree": 1, "enable_cfg_parallel": false, "attention_backend": "", "cache_strategy": "none" }, "class_type": "SGLDOptions", "_meta": { "title": "SGLDiffusion Options" } } } ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/qwen_image_sgld.json ================================================ { "3": { "inputs": { "seed": 808633539418610, "steps": 4, "cfg": 1, "sampler_name": "euler", "scheduler": "simple", "denoise": 1, "model": [ "66", 0 ], "positive": [ "6", 0 ], "negative": [ "7", 0 ], "latent_image": [ "58", 0 ] }, "class_type": "KSampler", "_meta": { "title": "KSampler" } }, "6": { "inputs": { "text": "\"A vibrant, warm neon-lit street scene in Hong Kong at the afternoon, with a mix of colorful Chinese and English signs glowing brightly. The atmosphere is lively, cinematic, and rain-washed with reflections on the pavement. The colors are vivid, full of pink, blue, red, and green hues. Crowded buildings with overlapping neon signs. 1980s Hong Kong style. Signs include:\n\"龍鳳冰室\" \"金華燒臘\" \"HAPPY HAIR\" \"鴻運茶餐廳\" \"EASY BAR\" \"永發魚蛋粉\" \"添記粥麵\" \"SUNSHINE MOTEL\" \"美都餐室\" \"富記糖水\" \"太平館\" \"雅芳髮型屋\" \"STAR KTV\" \"銀河娛樂城\" \"百樂門舞廳\" \"BUBBLE CAFE\" \"萬豪麻雀館\" \"CITY LIGHTS BAR\" \"瑞祥香燭莊\" \"文記文具\" \"GOLDEN JADE HOTEL\" \"LOVELY BEAUTY\" \"合興百貨\" \"興旺電器\" And the background is warm yellow street and with all stores' lights on.", "speak_and_recognation": { "__value__": [ false, true ] }, "clip": [ "38", 0 ] }, "class_type": "CLIPTextEncode", "_meta": { "title": "CLIP Text Encode (Positive Prompt)" } }, "7": { "inputs": { "text": "", "speak_and_recognation": { "__value__": [ false, true ] }, "clip": [ "38", 0 ] }, "class_type": "CLIPTextEncode", "_meta": { "title": "CLIP Text Encode (Negative Prompt)" } }, "8": { "inputs": { "samples": [ "3", 0 ], "vae": [ "39", 0 ] }, "class_type": "VAEDecode", "_meta": { "title": "VAE Decode" } }, "38": { "inputs": { "clip_name": "qwen_2.5_vl_7b_fp8_scaled.safetensors", "type": "qwen_image", "device": "default" }, "class_type": "CLIPLoader", "_meta": { "title": "Load CLIP" } }, "39": { "inputs": { "vae_name": "qwen_image_vae.safetensors" }, "class_type": "VAELoader", "_meta": { "title": "Load VAE" } }, "58": { "inputs": { "width": 1328, "height": 1328, "batch_size": 1 }, "class_type": "EmptySD3LatentImage", "_meta": { "title": "EmptySD3LatentImage" } }, "60": { "inputs": { "filename_prefix": "ComfyUI" }, "class_type": "SaveImage", "_meta": { "title": "Save Image" } }, "66": { "inputs": { "shift": 3.1000000000000005, "model": [ "78", 0 ] }, "class_type": "ModelSamplingAuraFlow", "_meta": { "title": "ModelSamplingAuraFlow" } }, "77": { "inputs": { "unet_name": "qwen_image_2512_bf16.safetensors", "weight_dtype": "default" }, "class_type": "SGLDUNETLoader", "_meta": { "title": "SGLDiffusion UNET Loader" } }, "78": { "inputs": { "lora_name": "Qwen-Image-2512-Lightning-4steps-V1.0-bf16.safetensors", "strength_model": 1, "nickname": "", "target": "all", "model": [ "77", 0 ] }, "class_type": "SGLDLoraLoader", "_meta": { "title": "SGLDiffusion LoRA Loader" } } } ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/sgld_image2video.json ================================================ { "1": { "inputs": { "base_url": "http://localhost:3000/v1", "api_key": "sk-proj-1234567890" }, "class_type": "SGLDiffusionServerModel", "_meta": { "title": "SGLDiffusion Server Model" } }, "3": { "inputs": { "prompt": "The girl turn the body and spin around in place.", "main": "none", "lighting": "none", "speak_and_recognation": { "__value__": [ false, true ] } }, "class_type": "easy prompt", "_meta": { "title": "Prompt" } }, "4": { "inputs": { "text": "", "anything": [ "1", 1 ] }, "class_type": "easy showAnything", "_meta": { "title": "Show Any" } }, "15": { "inputs": { "positive_prompt": [ "3", 0 ], "negative_prompt": "", "seed": 2435791308, "steps": 50, "cfg": 4, "width": 704, "height": 1280, "num_frames": 16, "fps": 16, "seconds": 1, "enable_teacache": false, "sgld_client": [ "1", 0 ], "image": [ "17", 0 ] }, "class_type": "SGLDiffusionGenerateVideo", "_meta": { "title": "SGLDiffusion Generate Video" } }, "16": { "inputs": { "filename_prefix": "video/ComfyUI", "format": "auto", "codec": "auto", "video-preview": "", "video": [ "15", 0 ] }, "class_type": "SaveVideo", "_meta": { "title": "save video" } }, "17": { "inputs": { "image": "tmpe_w0bd_0.jpg" }, "class_type": "LoadImage", "_meta": { "title": "load image" } } } ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/sgld_text2img.json ================================================ { "1": { "inputs": { "base_url": "http://localhost:3000/v1", "api_key": "sk-proj-1234567890" }, "class_type": "SGLDiffusionServerModel", "_meta": { "title": "SGLDiffusion Server Model" } }, "3": { "inputs": { "prompt": "a bicycle, illustration in the style of SMPL, thick black lines on a white background", "main": "none", "lighting": "none", "speak_and_recognation": { "__value__": [ false, true ] } }, "class_type": "easy prompt", "_meta": { "title": "Prompt" } }, "4": { "inputs": { "text": "", "anything": [ "1", 1 ] }, "class_type": "easy showAnything", "_meta": { "title": "Show Any" } }, "5": { "inputs": { "filename_prefix": "ComfyUI", "images": [ "6", 0 ] }, "class_type": "SaveImage", "_meta": { "title": "save image" } }, "6": { "inputs": { "positive_prompt": [ "3", 0 ], "negative_prompt": "", "seed": 4215918563, "steps": 50, "cfg": 4, "width": 512, "height": 512, "enable_teacache": false, "sgld_client": [ "11", 0 ], "image": [ "14", 0 ] }, "class_type": "SGLDiffusionGenerateImage", "_meta": { "title": "SGLDiffusion Generate Image" } }, "11": { "inputs": { "lora_name": "dvyio/flux-lora-simple-illustration", "lora_nickname": "", "target": "all", "sgld_client": [ "1", 0 ] }, "class_type": "SGLDiffusionSetLora", "_meta": { "title": "SGLDiffusion Set LoRA" } }, "14": { "inputs": { "width": 512, "height": 512, "batch_size": 1, "color": 0 }, "class_type": "EmptyImage", "_meta": { "title": "empty image" } } } ================================================ FILE: python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/z-image_sgld.json ================================================ { "3": { "inputs": { "seed": 3338398, "steps": 9, "cfg": 1, "sampler_name": "euler", "scheduler": "simple", "denoise": 1, "model": [ "28", 0 ], "positive": [ "6", 0 ], "negative": [ "7", 0 ], "latent_image": [ "13", 0 ] }, "class_type": "KSampler", "_meta": { "title": "KSampler" } }, "6": { "inputs": { "text": "cute anime style girl with massive fluffy fennec ears and a big fluffy tail blonde messy long hair blue eyes wearing a maid outfit with a long black gold leaf pattern dress and a white apron, it is a postcard held by a hand in front of a beautiful realistic city at sunset and there is cursive writing that says \"ZImage, Now in ComfyUI\"", "speak_and_recognation": { "__value__": [ false, true ] }, "clip": [ "18", 0 ] }, "class_type": "CLIPTextEncode", "_meta": { "title": "CLIP Text Encode (Positive Prompt)" } }, "7": { "inputs": { "text": "blurry ugly bad", "speak_and_recognation": { "__value__": [ false, true ] }, "clip": [ "18", 0 ] }, "class_type": "CLIPTextEncode", "_meta": { "title": "CLIP Text Encode (Negative Prompt)" } }, "8": { "inputs": { "samples": [ "3", 0 ], "vae": [ "17", 0 ] }, "class_type": "VAEDecode", "_meta": { "title": "VAE Decode" } }, "9": { "inputs": { "filename_prefix": "ComfyUI", "images": [ "8", 0 ] }, "class_type": "SaveImage", "_meta": { "title": "Save Image" } }, "13": { "inputs": { "width": 1024, "height": 1024, "batch_size": 1 }, "class_type": "EmptySD3LatentImage", "_meta": { "title": "EmptySD3LatentImage" } }, "17": { "inputs": { "vae_name": "ae.safetensors" }, "class_type": "VAELoader", "_meta": { "title": "VAE Loader" } }, "18": { "inputs": { "clip_name": "qwen_3_4b.safetensors", "type": "lumina2", "device": "default" }, "class_type": "CLIPLoader", "_meta": { "title": "CLIP Loader" } }, "28": { "inputs": { "unet_name": "z_image_turbo_bf16.safetensors", "weight_dtype": "default" }, "class_type": "SGLDUNETLoader", "_meta": { "title": "SGLDiffusion UNET Loader" } } } ================================================ FILE: python/sglang/multimodal_gen/apps/webui/README.md ================================================ # SGLang Diffusion WebUI User Guide SGLang Diffusion WebUI provides an intuitive Gradio-based interface for image and video generation, supporting parameter tuning and real-time previews. ## Prerequisites The WebUI runs on Gradio. To get started, install Gradio first: ```bash pip install gradio==6.1.0 ``` ## Launch WebUI Service SGLang Diffusion now includes an integrated WebUI. Simply add the `--webui` parameter when starting the service. ### Launch Text-to-Image Service ```bash sglang serve --model-path black-forest-labs/FLUX.1-dev --num-gpus 1 --webui --webui-port 2333 ``` ### Launch Text-to-Video Service ```bash sglang serve --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers --num-gpus 1 --webui --webui-port 2333 ``` ### Launch Image-to-Image Service ```bash sglang serve --model-path Qwen/Qwen-Image-Edit-2511 --num-gpus 1 --webui --webui-port 2333 ``` ### Launch Image-to-Video Service ```bash sglang serve --model-path Wan-AI/Wan2.2-TI2V-5B-Diffusers --num-gpus 1 --webui --webui-port 2333 ``` ## Port Forwarding Once the WebUI service is running, you need to use **SSH port forwarding** to securely access the remote service from your local machine. In most cases: Your IDE (like VS Code, Cursor, etc.) can handle this automatically. Check your IDE's remote development or port forwarding features. Otherwise, execute this command manually. ```bash ssh -L ${WEBUI_PORT}:localhost:${WEBUI_PORT} user_name@machine_name ``` Learn more about port forwarding: [Port Forwarding](https://en.wikipedia.org/wiki/Port_forwarding). ## Interface Instructions You can view your model path and task name directly in the UI. We'd appreciate any feedback you'd like to share. Once launched, access the interface at `http://localhost:${WEBUI_PORT}` in your browser. ================================================ FILE: python/sglang/multimodal_gen/apps/webui/__init__.py ================================================ from .main import run_sgl_diffusion_webui __all__ = ["run_sgl_diffusion_webui"] ================================================ FILE: python/sglang/multimodal_gen/apps/webui/main.py ================================================ import argparse import os from sglang.multimodal_gen.configs.sample.sampling_params import ( DataType, SamplingParams, ) from sglang.multimodal_gen.runtime.entrypoints.utils import ( post_process_sample, prepare_request, ) from sglang.multimodal_gen.runtime.scheduler_client import sync_scheduler_client from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.srt.environ import envs logger = init_logger(__name__) def add_webui_args(parser: argparse.ArgumentParser): """Add the arguments for the generate command.""" parser = ServerArgs.add_cli_args(parser) parser = SamplingParams.add_cli_args(parser) return parser def run_sgl_diffusion_webui(server_args: ServerArgs): # import gradio in function to avoid CI crash import gradio as gr def resolve_model_repo_id(model_path: str) -> str: from pathlib import Path from huggingface_hub.utils import HFValidationError, validate_repo_id try: validate_repo_id(model_path) return model_path except HFValidationError: pass p = Path(model_path).expanduser() parts = p.parts if len(parts) < 2: raise ValueError(f"Invalid model_path: {model_path}") candidate = f"{parts[-2]}/{parts[-1]}" validate_repo_id(candidate) # let it raise if invalid return candidate repo_id = resolve_model_repo_id(server_args.model_path) if envs.SGLANG_USE_MODELSCOPE.get(): from modelscope.hub.api import HubApi api = HubApi() model_info_obj = api.model_info(repo_id) task_name = model_info_obj.tasks[0]["Name"].replace("-synthesis", "") else: from huggingface_hub import model_info task_name = model_info(repo_id).pipeline_tag # init client sync_scheduler_client.initialize(server_args) if task_name in ("text-to-video", "image-to-video", "video-to-video"): task_type = "video" elif task_name in ["text-to-image", "image-to-image"]: task_type = "image" else: raise ValueError( f"The task name {task_name} of model {server_args.model_path} is not a valid task name. Please check the model path." ) video_visible_only = task_type == "video" image_visible_only = task_type == "image" # server_args will be reused in gradio_generate function def gradio_generate( prompt, negative_prompt, reference_image_paths_str, seed, num_frames, frames_per_second, width, height, num_inference_steps, guidance_scale, enable_teacache, ): """ NOTE: The input and output of function which is called by gradio button must be gradio components So we use global variable sampling_params_kwargs to avoid pass this param, because gradio does not support this. return [ np.ndarray, None ] | [None, np.ndarray] """ if reference_image_paths_str: if "，" in reference_image_paths_str: logger.warning( f"Warning: please use English comma to separate the reference image paths, and the reference image paths is: {reference_image_paths_str}" ) reference_image_paths_str = reference_image_paths_str.replace("，", ",") image_path = [path.strip() for path in reference_image_paths_str.split(",")] else: image_path = None sampling_params_kwargs = dict( prompt=prompt, negative_prompt=negative_prompt, image_path=image_path, seed=seed, num_frames=num_frames, fps=frames_per_second, width=width, height=height, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, enable_teacache=enable_teacache, return_file_paths_only=False, ) sampling_params = SamplingParams.from_user_sampling_params_args( server_args.model_path, server_args=server_args, **sampling_params_kwargs, ) batch = prepare_request( server_args=server_args, sampling_params=sampling_params, ) result = sync_scheduler_client.forward([batch]) save_file_path = str(os.path.join(batch.output_path, batch.output_file_name)) if result.output is None: sampling_params_str = "\n".join( [f"{key}: {value}" for key, value in sampling_params_kwargs.items()] ) no_output_msg = f"No output is generated by client, and their sampling params is: {sampling_params_str}" if batch.data_type == DataType.VIDEO: if os.path.exists(save_file_path): logger.warning(no_output_msg) return None, save_file_path else: no_output_msg += f"\nAnd the expected output file was not found at: {save_file_path}" raise ValueError(no_output_msg) else: raise ValueError(no_output_msg) frames = post_process_sample( result.output[0], batch.data_type, batch.fps, batch.save_output, save_file_path, ) if batch.data_type == DataType.VIDEO: # gradio video need video path to show video return None, save_file_path else: return frames[0], None with gr.Blocks() as demo: gr.Markdown("# 🚀 SGLang Diffusion Application") with gr.Row(): launched_model_box = gr.Textbox(label="Model", value=server_args.model_path) task_name_box = gr.Textbox(label="Task name", value=task_name) with gr.Row(): with gr.Column(scale=4): prompt = gr.Textbox(label="Prompt", value="A curious raccoon") negative_prompt = gr.Textbox( label="Negative_prompt", value="Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards", ) with gr.Column(scale=1): seed = gr.Number(label="seed", precision=0, value=1234) run_btn = gr.Button("Generate", variant="primary", size="lg") with gr.Row(): with gr.Column(): width = gr.Number(label="width", precision=0, value=720) height = gr.Number(label="height", precision=0, value=480) num_inference_steps = gr.Slider( minimum=0, maximum=50, value=20, step=1, label="num_inference_steps" ) guidance_scale = gr.Slider( minimum=0.0, maximum=10, value=5, step=0.01, label="guidance_scale" ) num_frames = gr.Slider( minimum=1, maximum=181, value=81, step=1, label="num_frames", visible=video_visible_only, ) frames_per_second = gr.Slider( minimum=4, maximum=60, value=16, step=1, label="frames_per_second", visible=video_visible_only, ) reference_image_paths_str = gr.Textbox( label="reference images", placeholder="Examples: 'image1.png, image2.png' or 'https://example.com/image1.png, https://example.com/image2.png'", ) enable_teacache = gr.Checkbox(label="enable_teacache", value=False) with gr.Column(): image_out = gr.Image( label="Generated Image", visible=image_visible_only, format="png" ) video_out = gr.Video( label="Generated Video", visible=video_visible_only ) run_btn.click( fn=gradio_generate, inputs=[ prompt, negative_prompt, reference_image_paths_str, seed, num_frames, frames_per_second, width, height, num_inference_steps, guidance_scale, enable_teacache, ], outputs=[image_out, video_out], ) _, local_url, _ = demo.launch( server_port=server_args.webui_port, quiet=True, prevent_thread_lock=True, show_error=True, ) # print banner delimiter = "=" * 80 url = local_url or f"http://localhost:{server_args.webui_port}" print(f""" {delimiter} \033[1mSGLang Diffusion WebUI available at:\033[0m \033[1;4;92m{url}\033[0m {delimiter} """) demo.block_thread() ================================================ FILE: python/sglang/multimodal_gen/benchmarks/bench_offline_throughput.py ================================================ """ Benchmark offline throughput for multimodal generation models (Image/Video Generation). This script benchmarks generation throughput without running a server, using low-level APIs. It provides detailed metrics on throughput, latency, and resource utilization. # Usage Examples ## Text-to-Video with VBench dataset python -m sglang.multimodal_gen.benchmarks.bench_offline_throughput \\ --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \\ --dataset vbench \\ --num-prompts 20 \\ --batch-size 1 \\ --width 512 --height 512 --num-frames 16 ## Random dataset for stress testing python -m sglang.multimodal_gen.benchmarks.bench_offline_throughput \\ --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \\ --dataset random \\ --num-prompts 100 \\ --batch-size 1 \\ --num-inference-steps 20 \\ --output-file results.json """ import argparse import dataclasses import json import time from dataclasses import dataclass from typing import Any, Dict, List, Tuple import torch from tqdm import tqdm from sglang.multimodal_gen.benchmarks.datasets import RandomDataset, VBenchDataset from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator from sglang.multimodal_gen.runtime.server_args import ServerArgs, set_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import ( configure_logger, init_logger, ) from sglang.multimodal_gen.test.test_utils import print_divider, print_value_formatted logger = init_logger(__name__) @dataclass class BatchOutput: """Container for batch generation results.""" latency: float = 0.0 latency_per_sample: float = 0.0 num_samples: int = 0 total_frames: int = 0 peak_memory_mb: float = 0.0 success: bool = False error: str = "" @dataclass class BenchArgs: """Benchmark configuration for multimodal generation.""" # Diffusion Model Configuration num_inference_steps: int = 20 guidance_scale: float = 7.5 seed: int = 42 disable_safety_checker: bool = False # Output Configuration width: int = 32 height: int = 32 num_frames: int = 1 fps: int = 24 # Dataset & Benchmark dataset: str = "random" dataset_path: str = "" task_name: str = "unknown" num_prompts: int = 10 batch_size: int = 1 # Benchmark Execution skip_warmup: bool = False output_file: str = "" disable_tqdm: bool = False @staticmethod def add_cli_args(parser: argparse.ArgumentParser): """Add benchmark-specific CLI arguments.""" # Diffusion Model Configuration parser.add_argument( "--num-inference-steps", type=int, default=20, help="Number of denoising steps", ) parser.add_argument( "--guidance-scale", type=float, default=7.5, help="Classifier-free guidance scale", ) parser.add_argument("--seed", type=int, default=42, help="Random seed") parser.add_argument( "--disable-safety-checker", action="store_true", help="Disable NSFW detection", ) # Output Configuration parser.add_argument("--width", type=int, default=32, help="Image/video width") parser.add_argument("--height", type=int, default=32, help="Image/video height") parser.add_argument( "--num-frames", type=int, default=1, help="Number of frames for video" ) parser.add_argument("--fps", type=int, default=24, help="FPS for video") # Dataset & Benchmark parser.add_argument( "--dataset", type=str, default="random", choices=["vbench", "random"], help="Dataset to use", ) parser.add_argument( "--dataset-path", type=str, default="", help="Path to dataset (prompts file or image directory)", ) parser.add_argument( "--task-name", type=str, default="unknown", help="Task name for benchmark identification", ) parser.add_argument( "--num-prompts", type=int, default=10, help="Total number of prompts to benchmark", ) parser.add_argument( "--batch-size", type=int, default=1, help="Batch size per generation call (currently only bs=1 is supported)", ) # Benchmark Execution parser.add_argument( "--skip-warmup", action="store_true", help="Skip warmup batch" ) parser.add_argument( "--output-file", type=str, default="", help="Output JSON file for results (append mode)", ) parser.add_argument( "--disable-tqdm", action="store_true", help="Disable progress bar", ) @classmethod def from_cli_args(cls, args: argparse.Namespace): """Create BenchArgs from parsed CLI arguments.""" attrs = [attr.name for attr in dataclasses.fields(cls)] return cls(**{attr: getattr(args, attr) for attr in attrs}) def initialize_engine(server_args: ServerArgs) -> DiffGenerator: """Initialize diffusion pipeline engine.""" logger.info("Initializing engine...") engine = DiffGenerator.from_server_args(server_args, local_mode=True) logger.info("Engine initialized successfully") return engine def generate_batch( engine: DiffGenerator, bench_args: BenchArgs, prompts: List[str], user_sampling_params: Dict[str, Any], ) -> BatchOutput: """Generate batch of images/videos synchronously.""" output = BatchOutput() start_time = time.perf_counter() torch.cuda.reset_peak_memory_stats() for prompt in prompts: try: sampling_params_kwargs = dict(user_sampling_params) sampling_params_kwargs["prompt"] = prompt result = engine.generate(sampling_params_kwargs=sampling_params_kwargs) if result is not None: if isinstance(result, list): output.total_frames += len(result) else: output.total_frames += 1 output.num_samples += 1 except Exception as e: logger.error(f"Generation failed for prompt '{prompt[:50]}...': {e}") output.error = str(e) output.latency = time.perf_counter() - start_time output.latency_per_sample = output.latency / len(prompts) if prompts else 0.0 output.success = output.num_samples > 0 output.peak_memory_mb = torch.cuda.max_memory_allocated() / (1024 * 1024) logger.debug( f"Batch generated: {output.num_samples}/{len(prompts)} samples in {output.latency:.2f}s" ) return output def calculate_metrics( outputs: List[BatchOutput], total_duration: float, resolution: Tuple[int, int, int], num_requests: int, ) -> Dict[str, Any]: """Calculate generation-specific throughput metrics.""" successful = [o for o in outputs if o.success] num_success = sum(o.num_samples for o in successful) total_frames = sum(o.total_frames for o in successful) peak_memory = max((o.peak_memory_mb for o in outputs), default=0) width, height, frames = resolution pixels_per_sample = width * height * frames total_pixels = num_success * pixels_per_sample metrics = { "num_requests": num_requests, "successful_requests": num_success, "failed_requests": num_requests - num_success, "total_duration_seconds": total_duration, "total_frames_generated": total_frames, "total_pixels_generated": total_pixels, "images_per_second": num_success / total_duration if total_duration > 0 else 0, "frames_per_second": total_frames / total_duration if total_duration > 0 else 0, "megapixels_per_second": ( total_pixels / (total_duration * 1e6) if total_duration > 0 else 0 ), "requests_per_second": ( num_success / total_duration if total_duration > 0 else 0 ), "latency_per_request_seconds": ( total_duration / num_success if num_success > 0 else 0 ), "peak_memory_mb": peak_memory, } return metrics def throughput_test( server_args: ServerArgs, bench_args: BenchArgs, ) -> Dict[str, Any]: """Main throughput benchmark function.""" configure_logger(server_args=server_args) logger.info("Starting offline throughput benchmark...") engine = initialize_engine(server_args) logger.info(f"Loading {bench_args.dataset} dataset...") if bench_args.dataset == "vbench": bench_args.task_name = engine.server_args.pipeline_config.task_type dataset = VBenchDataset(bench_args) elif bench_args.dataset == "random": dataset = RandomDataset(bench_args) else: raise ValueError(f"Unknown dataset: {bench_args.dataset}") sampling_params = { "guidance_scale": bench_args.guidance_scale, "num_inference_steps": bench_args.num_inference_steps, "height": bench_args.height, "width": bench_args.width, "num_frames": bench_args.num_frames, "seed": bench_args.seed, } if bench_args.disable_safety_checker: sampling_params["safety_checker"] = None if not bench_args.skip_warmup: logger.info("Running warmup batch...") warmup_count = min(bench_args.batch_size, len(dataset)) warmup_prompts = [dataset[i].prompt for i in range(warmup_count)] generate_batch(engine, bench_args, warmup_prompts, sampling_params) logger.info(f"Running benchmark with {bench_args.num_prompts} prompts...") outputs: List[BatchOutput] = [] total_count = min(bench_args.num_prompts, len(dataset)) all_prompts = [dataset[i].prompt for i in range(total_count)] start_time = time.perf_counter() num_batches = (total_count + bench_args.batch_size - 1) // bench_args.batch_size pbar = tqdm( total=num_batches, disable=bench_args.disable_tqdm, desc="Benchmark", ) for batch_start in range(0, total_count, bench_args.batch_size): batch_end = min(batch_start + bench_args.batch_size, total_count) batch_prompts = all_prompts[batch_start:batch_end] batch_output = generate_batch( engine, bench_args, batch_prompts, sampling_params ) outputs.append(batch_output) pbar.update(1) pbar.close() total_duration = time.perf_counter() - start_time resolution = (bench_args.width, bench_args.height, bench_args.num_frames) metrics = calculate_metrics( outputs, total_duration, resolution=resolution, num_requests=total_count, ) display_results( metrics, bench_args, model_path=server_args.model_path, ) if bench_args.output_file: save_results(metrics, bench_args, server_args) return metrics def display_results( metrics: Dict[str, Any], bench_args: BenchArgs, model_path: str, ): """Display benchmark results in console.""" print( "\n{s:{c}^{n}}".format(s=" Offline Throughput Benchmark Result ", n=110, c="=") ) print_value_formatted("Model:", model_path) print_value_formatted("Dataset:", bench_args.dataset) print_value_formatted( "Resolution:", f"{bench_args.width}x{bench_args.height}x{bench_args.num_frames}", ) print_value_formatted("Num Inference Steps:", bench_args.num_inference_steps) print_divider(75) print_value_formatted("Total Requests:", metrics["num_requests"]) print_value_formatted("Successful Requests:", metrics["successful_requests"]) print_value_formatted("Failed Requests:", metrics["failed_requests"]) print_value_formatted( "Total Duration (seconds):", metrics["total_duration_seconds"] ) print_divider(75) print_value_formatted("Frames Generated:", metrics["total_frames_generated"]) print_value_formatted( "Megapixels Generated:", metrics["total_pixels_generated"] / 1e6 ) print_divider(75) print_value_formatted( "Frame Throughput (frames/sec):", metrics["frames_per_second"] ) print_value_formatted("MP Throughput (MP/sec):", metrics["megapixels_per_second"]) print_value_formatted("Requests Per Second:", metrics["requests_per_second"]) print_value_formatted( "Latency Per Request (sec):", metrics["latency_per_request_seconds"] ) print_value_formatted("Peak Memory (MB):", metrics["peak_memory_mb"]) print_divider(110, "=") def save_results( metrics: Dict[str, Any], bench_args: BenchArgs, server_args: ServerArgs, ): """Save benchmark results to JSON file.""" result = { "metadata": { "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), "model_path": server_args.model_path, "task_type": bench_args.task_name, "backend": "engine", }, "configuration": { "num_inference_steps": bench_args.num_inference_steps, "guidance_scale": bench_args.guidance_scale, "seed": bench_args.seed, "batch_size": bench_args.batch_size, "num_prompts": bench_args.num_prompts, "resolution": f"{bench_args.width}x{bench_args.height}x{bench_args.num_frames}", "dataset": bench_args.dataset, }, "results": metrics, } with open(bench_args.output_file, "a") as f: f.write(json.dumps(result) + "\n") logger.info(f"Results saved to {bench_args.output_file}") def main(): """Main entry point.""" parser = argparse.ArgumentParser( description="Offline throughput benchmark for multimodal generation models" ) ServerArgs.add_cli_args(parser) BenchArgs.add_cli_args(parser) args = parser.parse_args() server_args = ServerArgs.from_cli_args(args) bench_args = BenchArgs.from_cli_args(args) set_global_server_args(server_args) result = throughput_test(server_args, bench_args) return result if __name__ == "__main__": main() ================================================ FILE: python/sglang/multimodal_gen/benchmarks/bench_serving.py ================================================ """ Benchmark online serving for diffusion models (Image/Video Generation). Usage: # launch a server and benchmark on it # T2V or T2I or any other multimodal generation model sglang serve --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers --num-gpus 1 --port 1231 # benchmark it and make sure the port is the same as the server's port python3 -m sglang.multimodal_gen.benchmarks.bench_serving --dataset vbench --num-prompts 20 --port 1231 # benchmark with SLO metrics enabled python3 -m sglang.multimodal_gen.benchmarks.bench_serving --dataset vbench --num-prompts 20 --port 1231 --slo --slo-scale 3.0 --warmup-requests 2 """ import argparse import asyncio import json import os import time from dataclasses import replace from typing import Any, Dict, List, Optional import aiohttp import numpy as np import requests from tqdm.asyncio import tqdm from sglang.multimodal_gen.benchmarks.datasets import ( RandomDataset, RequestFuncInput, RequestFuncOutput, VBenchDataset, ) from sglang.multimodal_gen.runtime.utils.logging_utils import ( configure_logger, init_logger, ) from sglang.multimodal_gen.test.test_utils import print_divider, print_value_formatted logger = init_logger(__name__) # Patch size used for computing area units (e.g. in latent diffusion models). PATCH_SIZE = 16 PATCH_AREA = PATCH_SIZE * PATCH_SIZE def _compute_scale_factor(req: RequestFuncInput, args) -> Optional[float]: """Computes the composite scale factor (area × frames × steps) for a request.""" width = req.width or args.width height = req.height or args.height if None in (width, height): return None frames = req.num_frames or args.num_frames steps = req.num_inference_steps or args.num_inference_steps frame_scale = frames if isinstance(frames, int) and frames > 0 else 1 step_scale = steps if isinstance(steps, int) and steps > 0 else 1 area_units = max((float(width) * float(height)) / float(PATCH_AREA), 1.0) return area_units * float(frame_scale) * float(step_scale) def _compute_expected_latency_ms_from_base( req: RequestFuncInput, args, base_time_ms: Optional[float] ) -> Optional[float]: """Scales latency linearly by pixel area, frame count, and inference steps.""" if base_time_ms is None: return None scale = _compute_scale_factor(req, args) if scale is None: return None return float(base_time_ms) * scale def _infer_slo_base_time_ms_from_warmups( warmup_pairs: List[tuple], args ) -> Optional[float]: """Derives median base latency from successful warmup runs.""" candidates_ms: List[float] = [] for req, out in warmup_pairs: if not out.success or out.latency <= 0: logger.warning( f"Skipping warmup result: success={out.success}, latency={out.latency:.3f}" ) continue scale = _compute_scale_factor(req, args) if scale is None or scale <= 0: continue candidates_ms.append((out.latency * 1000.0) / scale) return float(np.median(candidates_ms)) if candidates_ms else None def _populate_slo_ms_from_warmups( requests_list: List[RequestFuncInput], warmup_pairs: List[tuple], args ) -> List[RequestFuncInput]: """Assigns estimated SLO targets to requests lacking them.""" if not any(req.slo_ms is None for req in requests_list): return requests_list base_time_ms = _infer_slo_base_time_ms_from_warmups(warmup_pairs, args) if base_time_ms is None: return requests_list slo_scale = float(getattr(args, "slo_scale", 3.0)) if slo_scale <= 0: raise ValueError(f"slo_scale must be positive, got {slo_scale}.") updated: List[RequestFuncInput] = [] for req in requests_list: if req.slo_ms is not None: updated.append(req) continue expected_ms = _compute_expected_latency_ms_from_base(req, args, base_time_ms) if expected_ms is not None: # Create a new RequestFuncInput with updated slo_ms updated.append(replace(req, slo_ms=expected_ms * slo_scale)) else: updated.append(req) return updated async def async_request_image_sglang( input: RequestFuncInput, session: aiohttp.ClientSession, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: output = RequestFuncOutput() output.start_time = time.perf_counter() # Check if we need to use multipart (for image edits with input images) if input.image_paths and len(input.image_paths) > 0: # Use multipart/form-data for image edits data = aiohttp.FormData() data.add_field("model", input.model) data.add_field("prompt", input.prompt) data.add_field("response_format", "b64_json") if input.width and input.height: data.add_field("size", f"{input.width}x{input.height}") # Merge extra parameters for key, value in input.extra_body.items(): data.add_field(key, str(value)) # Add image file(s) for idx, img_path in enumerate(input.image_paths): if os.path.exists(img_path): data.add_field( "image", open(img_path, "rb"), filename=os.path.basename(img_path), content_type="application/octet-stream", ) else: output.error = f"Image file not found: {img_path}" output.success = False if pbar: pbar.update(1) return output try: async with session.post(input.api_url, data=data) as response: if response.status == 200: resp_json = await response.json() output.response_body = resp_json output.success = True if "peak_memory_mb" in resp_json: output.peak_memory_mb = resp_json["peak_memory_mb"] else: output.error = f"HTTP {response.status}: {await response.text()}" output.success = False except Exception as e: output.error = str(e) output.success = False else: # Use JSON for text-to-image generation payload = { "model": input.model, "prompt": input.prompt, "n": 1, "response_format": "b64_json", } if input.width and input.height: payload["size"] = f"{input.width}x{input.height}" # Merge extra parameters payload.update(input.extra_body) try: async with session.post(input.api_url, json=payload) as response: if response.status == 200: resp_json = await response.json() output.response_body = resp_json output.success = True if "peak_memory_mb" in resp_json: output.peak_memory_mb = resp_json["peak_memory_mb"] else: output.error = f"HTTP {response.status}: {await response.text()}" output.success = False except Exception as e: output.error = str(e) output.success = False output.latency = time.perf_counter() - output.start_time # Check SLO if defined if input.slo_ms is not None and output.success: output.slo_achieved = (output.latency * 1000.0) <= input.slo_ms if pbar: pbar.update(1) return output async def async_request_video_sglang( input: RequestFuncInput, session: aiohttp.ClientSession, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: output = RequestFuncOutput() output.start_time = time.perf_counter() # 1. Submit Job job_id = None # Check if we need to upload images (Multipart) or just send JSON if input.image_paths and len(input.image_paths) > 0: # Use multipart/form-data data = aiohttp.FormData() data.add_field("model", input.model) data.add_field("prompt", input.prompt) if input.width and input.height: data.add_field("size", f"{input.width}x{input.height}") # Add extra body fields to form data if possible, or assume simple key-values # Note: Nested dicts in extra_body might need JSON serialization if API expects it stringified if input.extra_body: data.add_field("extra_body", json.dumps(input.extra_body)) # Explicitly add fps/num_frames if they are not in extra_body (bench_serving logic overrides) if input.num_frames: data.add_field("num_frames", str(input.num_frames)) if input.fps: data.add_field("fps", str(input.fps)) # Add image file # Currently only support single image upload as 'input_reference' per API spec img_path = input.image_paths[0] if os.path.exists(img_path): data.add_field( "input_reference", open(img_path, "rb"), filename=os.path.basename(img_path), content_type="application/octet-stream", ) else: output.error = f"Image file not found: {img_path}" output.success = False if pbar: pbar.update(1) return output try: async with session.post(input.api_url, data=data) as response: if response.status == 200: resp_json = await response.json() job_id = resp_json.get("id") else: output.error = ( f"Submit failed HTTP {response.status}: {await response.text()}" ) output.success = False if pbar: pbar.update(1) return output except Exception as e: output.error = f"Submit exception: {str(e)}" output.success = False if pbar: pbar.update(1) return output else: # Use JSON payload: Dict[str, Any] = { "model": input.model, "prompt": input.prompt, } if input.width and input.height: payload["size"] = f"{input.width}x{input.height}" if input.num_frames: payload["num_frames"] = input.num_frames if input.fps: payload["fps"] = input.fps payload.update(input.extra_body) try: async with session.post(input.api_url, json=payload) as response: if response.status == 200: resp_json = await response.json() job_id = resp_json.get("id") else: output.error = ( f"Submit failed HTTP {response.status}: {await response.text()}" ) output.success = False if pbar: pbar.update(1) return output except Exception as e: output.error = f"Submit exception: {str(e)}" output.success = False if pbar: pbar.update(1) return output if not job_id: output.error = "No job_id returned" output.success = False if pbar: pbar.update(1) return output # 2. Poll for completion # Assuming the API returns a 'status' field. # We construct the check URL. Assuming api_url is like .../v1/videos # The check url should be .../v1/videos/{id} check_url = f"{input.api_url}/{job_id}" while True: try: async with session.get(check_url) as response: if response.status == 200: status_data = await response.json() status = status_data.get("status") if status == "completed": output.success = True output.response_body = status_data if "peak_memory_mb" in status_data: output.peak_memory_mb = status_data["peak_memory_mb"] break elif status == "failed": output.success = False output.error = f"Job failed: {status_data.get('error')}" break else: # queued or processing await asyncio.sleep(1.0) else: output.success = False output.error = ( f"Poll failed HTTP {response.status}: {await response.text()}" ) break except Exception as e: output.success = False output.error = f"Poll exception: {str(e)}" break output.latency = time.perf_counter() - output.start_time # Check SLO if defined if input.slo_ms is not None and output.success: output.slo_achieved = (output.latency * 1000.0) <= input.slo_ms if pbar: pbar.update(1) return output def calculate_metrics( outputs: List[RequestFuncOutput], total_duration: float, requests_list: List[RequestFuncInput], args, slo_enabled: bool, ): success_outputs = [o for o in outputs if o.success] error_outputs = [o for o in outputs if not o.success] num_success = len(success_outputs) latencies = [o.latency for o in success_outputs] peak_memories = [o.peak_memory_mb for o in success_outputs if o.peak_memory_mb > 0] metrics = { "duration": total_duration, "completed_requests": num_success, "failed_requests": len(error_outputs), "throughput_qps": num_success / total_duration if total_duration > 0 else 0, "latency_mean": np.mean(latencies) if latencies else 0, "latency_median": np.median(latencies) if latencies else 0, "latency_p99": np.percentile(latencies, 99) if latencies else 0, "latency_p50": np.percentile(latencies, 50) if latencies else 0, "peak_memory_mb_max": max(peak_memories) if peak_memories else 0, "peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0, "peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0, } if slo_enabled: slo_defined_total = 0 slo_met_success = 0 for req, out in zip(requests_list, outputs): if req.slo_ms is None: continue slo_defined_total += 1 if out.slo_achieved: slo_met_success += 1 slo_attain_all = ( (slo_met_success / slo_defined_total) if slo_defined_total > 0 else 0.0 ) metrics.update( { "slo_attainment_rate": slo_attain_all, "slo_met_success": slo_met_success, "slo_scale": getattr(args, "slo_scale", 3.0), } ) return metrics def wait_for_service(base_url: str, timeout: int = 1200) -> None: logger.info(f"Waiting for service at {base_url}...") start_time = time.time() while True: try: # Try /health endpoint first resp = requests.get(f"{base_url}/health", timeout=1) if resp.status_code == 200: logger.info("Service is ready.") break except requests.exceptions.RequestException: pass if time.time() - start_time > timeout: raise TimeoutError( f"Service at {base_url} did not start within {timeout} seconds." ) time.sleep(1) async def benchmark(args): from huggingface_hub import model_info # Construct base_url if not provided if args.base_url is None: args.base_url = f"http://{args.host}:{args.port}" # Wait for service wait_for_service(args.base_url) # Fetch model info try: resp = requests.get(f"{args.base_url}/v1/model_info", timeout=5) if resp.status_code == 200: info = resp.json() if "model_path" in info and info["model_path"]: args.model = info["model_path"] logger.info(f"Updated model name from server: {args.model}") except Exception as e: logger.info(f"Failed to fetch model info: {e}. Using default: {args.model}") valid_tasks = ( "text-to-video", "image-to-video", "video-to-video", "text-to-image", "image-to-image", ) # Resolve task_name with priority: args.task > local config > HF pipeline_tag if args.task: task_name = args.task logger.info(f"Using task from --task: {task_name}") elif os.path.exists(args.model): config_path = os.path.join(args.model, "config.json") if os.path.exists(config_path): with open(config_path, "r") as f: config = json.load(f) task_name = config.get("pipeline_tag", "text-to-image") logger.info(f"Inferred task from local config.json: {task_name}") else: task_name = "text-to-image" logger.info(f"No config.json found, defaulting task to: {task_name}") else: task_name = model_info(args.model).pipeline_tag logger.info(f"Inferred task from HuggingFace pipeline_tag: {task_name}") if task_name not in valid_tasks: raise ValueError( f"Task '{task_name}' is not a valid multimodal generation task. " f"Use --task to specify one of: {', '.join(valid_tasks)}" ) if task_name in ("text-to-video", "image-to-video", "video-to-video"): api_url = f"{args.base_url}/v1/videos" request_func = async_request_video_sglang else: # text-to-image or image-to-image api_url = ( f"{args.base_url}/v1/images/edits" if task_name == "image-to-image" else f"{args.base_url}/v1/images/generations" ) request_func = async_request_image_sglang setattr(args, "task_name", task_name) if args.dataset == "vbench": dataset = VBenchDataset(args, api_url, args.model) elif args.dataset == "random": dataset = RandomDataset(args, api_url, args.model) else: raise ValueError(f"Unknown dataset: {args.dataset}") logger.info(f"Loading requests...") requests_list = dataset.get_requests() logger.info(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.") # Limit concurrency if args.max_concurrency is not None: semaphore = asyncio.Semaphore(args.max_concurrency) else: semaphore = None async def limited_request_func(req, session, pbar): if semaphore: async with semaphore: return await request_func(req, session, pbar) else: return await request_func(req, session, pbar) async with aiohttp.ClientSession() as session: # Run warmup requests warmup_pairs: List[tuple] = [] if args.warmup_requests and requests_list: # The server always overrides warmup requests to use # num_inference_steps=1 (see Req.set_as_warmup), so we match # that here to keep the benchmark's SLO estimation consistent. warmup_steps = 1 logger.info( f"Running {args.warmup_requests} warmup request(s) with " f"num_inference_steps={warmup_steps}..." ) for i in range(args.warmup_requests): warm_req = requests_list[i % len(requests_list)] warm_req = replace( warm_req, num_inference_steps=warmup_steps, ) warm_out = await limited_request_func(warm_req, session, None) warmup_pairs.append((warm_req, warm_out)) logger.info( f"Warmup {i+1}/{args.warmup_requests}: " f"latency={warm_out.latency:.2f}s, success={warm_out.success}" ) # Populate SLO values from warmups if enabled if args.slo: requests_list = _populate_slo_ms_from_warmups( requests_list=requests_list, warmup_pairs=warmup_pairs, args=args ) # Run benchmark pbar = tqdm(total=len(requests_list), disable=args.disable_tqdm) start_time = time.perf_counter() tasks = [] for req in requests_list: if args.request_rate != float("inf"): # Poisson process: inter-arrival times follow exponential distribution interval = np.random.exponential(1.0 / args.request_rate) await asyncio.sleep(interval) task = asyncio.create_task(limited_request_func(req, session, pbar)) tasks.append(task) outputs = await asyncio.gather(*tasks) total_duration = time.perf_counter() - start_time pbar.close() # Calculate metrics metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo) print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="=")) # Section 1: Configuration print_value_formatted("Task:", task_name) print_value_formatted("Model:", args.model) print_value_formatted("Dataset:", args.dataset) # Section 2: Execution & Traffic print_divider(50) print_value_formatted("Benchmark duration (s):", metrics["duration"]) print_value_formatted("Request rate:", str(args.request_rate)) print_value_formatted( "Max request concurrency:", str(args.max_concurrency) if args.max_concurrency else "not set", ) print_value_formatted( "Successful requests:", f"{metrics['completed_requests']}/{len(requests_list)}", ) # Section 3: Performance Metrics print_divider(50) print_value_formatted("Request throughput (req/s):", metrics["throughput_qps"]) print_value_formatted("Latency Mean (s):", metrics["latency_mean"]) print_value_formatted("Latency Median (s):", metrics["latency_median"]) print_value_formatted("Latency P99 (s):", metrics["latency_p99"]) if metrics["peak_memory_mb_max"] > 0: print_divider(50) print_value_formatted("Peak Memory Max (MB):", metrics["peak_memory_mb_max"]) print_value_formatted("Peak Memory Mean (MB):", metrics["peak_memory_mb_mean"]) print_value_formatted( "Peak Memory Median (MB):", metrics["peak_memory_mb_median"] ) if args.slo and "slo_attainment_rate" in metrics: print_divider(50) print( "{:<40} {:<15.2%}".format( "SLO Attainment Rate:", metrics["slo_attainment_rate"] ) ) print("{:<40} {:<15}".format("SLO Met (Success):", metrics["slo_met_success"])) print("{:<40} {:<15.2f}".format("SLO Scale:", metrics["slo_scale"])) print_divider(60) if args.output_file: with open(args.output_file, "w") as f: json.dump(metrics, f, indent=2) print(f"Metrics saved to {args.output_file}") if __name__ == "__main__": parser = argparse.ArgumentParser( description="Benchmark serving for diffusion models." ) parser.add_argument( "--backend", type=str, default=None, help="DEPRECATED: --task is deprecated and will be ignored. The task will be inferred from --model.", ) parser.add_argument( "--base-url", type=str, default=None, help="Base URL of the server (e.g., http://localhost:30000). Overrides host/port.", ) parser.add_argument("--host", type=str, default="localhost", help="Server host.") parser.add_argument("--port", type=int, default=30000, help="Server port.") parser.add_argument("--model", type=str, default="default", help="Model name.") parser.add_argument( "--dataset", type=str, default="vbench", choices=["vbench", "random"], help="Dataset to use.", ) parser.add_argument( "--task", type=str, choices=[ "text-to-video", "image-to-video", "text-to-image", "image-to-image", "video-to-video", ], default=None, help="The task will be inferred from huggingface pipeline_tag. When huggingface pipeline_tag is not provided, --task will be used.", ) parser.add_argument( "--dataset-path", type=str, default=None, help="Path to local dataset file (optional).", ) parser.add_argument( "--num-prompts", type=int, default=10, help="Number of prompts to benchmark." ) parser.add_argument( "--max-concurrency", type=int, default=1, help="Maximum number of concurrent requests, default to `1`. This can be used " "to help simulate an environment where a higher level component " "is enforcing a maximum number of concurrent requests. While the " "--request-rate argument controls the rate at which requests are " "initiated, this argument will control how many are actually allowed " "to execute at a time. This means that when used in combination, the " "actual request rate may be lower than specified with --request-rate, " "if the server is not processing requests fast enough to keep up.", ) parser.add_argument( "--request-rate", type=float, default=float("inf"), help="Number of requests per second. If this is inf, then all the requests are sent at time 0. " "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.", ) parser.add_argument("--width", type=int, default=None, help="Image/Video width.") parser.add_argument("--height", type=int, default=None, help="Image/Video height.") parser.add_argument( "--num-frames", type=int, default=None, help="Number of frames (for video)." ) parser.add_argument("--fps", type=int, default=None, help="FPS (for video).") parser.add_argument( "--output-file", type=str, default=None, help="Output JSON file for metrics." ) parser.add_argument( "--disable-tqdm", action="store_true", help="Disable progress bar." ) parser.add_argument( "--log-level", type=str, default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"], help="Log level.", ) parser.add_argument( "--slo", action="store_true", help="Enable SLO calculation. Uses trace-provided slo_ms or infers from warmups.", ) parser.add_argument( "--slo-scale", type=float, default=3.0, help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).", ) parser.add_argument( "--warmup-requests", type=int, default=1, help="Number of warmup requests to run before measurement.", ) parser.add_argument( "--num-inference-steps", type=int, default=None, help="Number of inference steps for diffusion models.", ) args = parser.parse_args() configure_logger(args) asyncio.run(benchmark(args)) ================================================ FILE: python/sglang/multimodal_gen/benchmarks/compare_perf.py ================================================ import argparse import json import os import re from datetime import datetime from typing import Any, Dict, List, Tuple def calculate_diff(base: float, new: float) -> Tuple[float, float]: """Returns (diff, diff_percent).""" diff = new - base if base == 0: percent = 0.0 else: percent = (diff / base) * 100 return diff, percent def calculate_upper_bound(baseline: float, rel_tol: float, min_abs_tol: float) -> float: """Calculates the upper bound for performance regression check.""" rel_limit = baseline * (1 + rel_tol) abs_limit = baseline + min_abs_tol return max(rel_limit, abs_limit) def calculate_lower_bound(baseline: float, rel_tol: float, min_abs_tol: float) -> float: """Calculates the lower bound for performance improvement check.""" rel_lower = baseline * (1 - rel_tol) abs_lower = baseline - min_abs_tol return min(rel_lower, abs_lower) def get_perf_status_emoji( baseline: float, new: float, rel_tol: float = 0.1, min_abs_tol: float = 120.0, ) -> str: """ Determines the status emoji based on performance difference. Logic: Upper bound (Slower): max(baseline * (1 + rel_tol), baseline + min_abs_tol) Lower bound (Faster): min(baseline * (1 - rel_tol), baseline - min_abs_tol) """ upper_bound = calculate_upper_bound(baseline, rel_tol, min_abs_tol) lower_bound = calculate_lower_bound(baseline, rel_tol, min_abs_tol) if new > upper_bound: return "🔴" elif new < lower_bound: return "🟢" else: return "⚪️" def consolidate_steps( steps_list: List[Dict[str, Any]], ) -> Tuple[Dict[str, float], List[str], Dict[str, int]]: """ Aggregates specific repeating steps (like denoising_step_*) into groups. Returns: - aggregated_durations: {name: duration_ms} - ordered_names: list of names in execution order - counts: {name: count_of_steps_aggregated} """ durations = {} counts = {} ordered_names = [] seen_names = set() # Regex for steps to group # Group "denoising_step_0", "denoising_step_1" -> "Denoising Loop" denoise_pattern = re.compile(r"^denoising_step_(\d+)$") denoising_group_name = "Denoising Loop" for step in steps_list: name = step.get("name", "unknown") dur = step.get("duration_ms", 0.0) match = denoise_pattern.match(name) if match: key = denoising_group_name if key not in durations: durations[key] = 0.0 counts[key] = 0 if key not in seen_names: ordered_names.append(key) seen_names.add(key) durations[key] += dur counts[key] += 1 else: # Standard stage (preserve order) if name not in durations: durations[name] = 0.0 counts[name] = 0 if name not in seen_names: ordered_names.append(name) seen_names.add(name) durations[name] += dur counts[name] += 1 return durations, ordered_names, counts def _load_benchmark_file(file_path: str) -> Dict[str, Any]: """Loads a benchmark JSON file.""" with open(file_path, "r", encoding="utf-8") as f: return json.load(f) def _get_status_emoji_from_diff_percent(diff_pct): if diff_pct < -2.0: return "✅" elif diff_pct > 2.0: return "❌" else: return "⚪️" def _print_single_comparison_report( others_data, base_e2e, combined_order, base_durations, others_processed, base_counts ): new_data = others_data[0] new_e2e = new_data.get("total_duration_ms", 0) diff_ms, diff_pct = calculate_diff(base_e2e, new_e2e) status = _get_status_emoji_from_diff_percent(diff_pct) print("#### 1. High-level Summary") print("| Metric | Baseline | New | Diff | Status |") print("| :--- | :--- | :--- | :--- | :--- |") print( f"| **E2E Latency** | {base_e2e:.2f} ms | {new_e2e:.2f} ms | **{diff_ms:+.2f} ms ({diff_pct:+.1f}%)** | {status} |" ) print( f"| **Throughput** | {1000 / base_e2e if base_e2e else 0:.2f} req/s | {1000 / new_e2e if new_e2e else 0:.2f} req/s | - | - |" ) print("\n") print("#### 2. Stage Breakdown") print("| Stage Name | Baseline (ms) | New (ms) | Diff (ms) | Diff (%) | Status |") print("| :--- | :--- | :--- | :--- | :--- | :--- |") new_durations, _, new_counts = others_processed[0] for stage in combined_order: b_val = base_durations.get(stage, 0.0) n_val = new_durations.get(stage, 0.0) b_count = base_counts.get(stage, 1) n_count = new_counts.get(stage, 1) s_diff, s_pct = calculate_diff(b_val, n_val) count_str = "" if stage == "Denoising Loop": count_str = ( f" ({n_count} steps)" if n_count == b_count else f" ({b_count}->{n_count} steps)" ) status_emoji = get_perf_status_emoji(b_val, n_val) print( f"| {stage}{count_str} | {b_val:.2f} | {n_val:.2f} | {s_diff:+.2f} | {s_pct:+.1f}% | {status_emoji} |" ) def _print_multi_comparison_report( base_e2e, others_data, other_labels, combined_order, base_durations, others_processed, ): print("#### 1. High-level Summary") header = "| Metric | Baseline | " + " | ".join(other_labels) + " |" sep = "| :--- | :--- | " + " | ".join([":---"] * len(other_labels)) + " |" print(header) print(sep) # E2E Row row_e2e = f"| **E2E Latency** | {base_e2e:.2f} ms |" for i, d in enumerate(others_data): val = d.get("total_duration_ms", 0) diff_ms, diff_pct = calculate_diff(base_e2e, val) status = _get_status_emoji_from_diff_percent(diff_pct) row_e2e += f" {val:.2f} ms ({diff_pct:+.1f}%) {status} |" print(row_e2e) print("\n") print("#### 2. Stage Breakdown") # Header: Stage | Baseline | Label1 | Label2 ... header = "| Stage Name | Baseline | " + " | ".join(other_labels) + " |" sep = "| :--- | :--- | " + " | ".join([":---"] * len(other_labels)) + " |" print(header) print(sep) for stage in combined_order: b_val = base_durations.get(stage, 0.0) row_str = f"| {stage} | {b_val:.2f} |" for i, (n_durations, _, n_counts) in enumerate(others_processed): n_val = n_durations.get(stage, 0.0) _, s_pct = calculate_diff(b_val, n_val) status_emoji = get_perf_status_emoji(b_val, n_val) row_str += f" {n_val:.2f} ({s_pct:+.1f}%) {status_emoji} |" print(row_str) def compare_benchmarks(file_paths: List[str], output_format: str = "markdown"): """ Compares benchmark JSON files and prints a report. First file is baseline, others will be compared against it. """ if len(file_paths) < 2: print("Error: Need at least 2 files to compare.") return try: data_list = [_load_benchmark_file(f) for f in file_paths] except Exception as e: print(f"Error loading benchmark files: {e}") return base_data = data_list[0] others_data = data_list[1:] # Use filenames as labels if multiple comparisons, else just "New" other_labels = [os.path.basename(p) for p in file_paths[1:]] base_e2e = base_data.get("total_duration_ms", 0) base_durations, base_order, base_counts = consolidate_steps( base_data.get("steps", []) ) others_processed = [] for d in others_data: dur, order, counts = consolidate_steps(d.get("steps", [])) others_processed.append((dur, order, counts)) combined_order = [] # Collect all unique stages maintaining order from newest to baseline for _, order, _ in reversed(others_processed): for name in order: if name not in combined_order: combined_order.append(name) for name in base_order: if name not in combined_order: combined_order.append(name) if output_format == "markdown": print("### Performance Comparison Report\n") if len(others_data) == 1: _print_single_comparison_report( others_data, base_e2e, combined_order, base_durations, others_processed, base_counts, ) else: _print_multi_comparison_report( base_e2e, others_data, other_labels, combined_order, base_durations, others_processed, ) print("\n") # Metadata print("
") print("Metadata\n") print(f"- Baseline Commit: `{base_data.get('commit_hash', 'N/A')}`") for i, d in enumerate(others_data): label = "New" if len(others_data) == 1 else other_labels[i] print(f"- {label} Commit: `{d.get('commit_hash', 'N/A')}`") print(f"- Timestamp: {datetime.now().isoformat()}") print("
") if __name__ == "__main__": parser = argparse.ArgumentParser( description="Compare sglang-diffusion performance JSON files." ) parser.add_argument( "files", nargs="+", help="List of JSON files. First is baseline, others are compared against it.", ) args = parser.parse_args() compare_benchmarks(args.files) ================================================ FILE: python/sglang/multimodal_gen/benchmarks/datasets.py ================================================ import glob import json import os import re import subprocess import uuid from abc import ABC, abstractmethod from dataclasses import dataclass, field from typing import Any, Dict, List, Optional import requests from PIL import Image from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @dataclass class RequestFuncInput: prompt: str api_url: str = "" model: str = "" width: Optional[int] = None height: Optional[int] = None num_frames: Optional[int] = None fps: Optional[int] = None extra_body: Dict[str, Any] = field(default_factory=dict) image_paths: Optional[List[str]] = None request_id: str = field(default_factory=lambda: str(uuid.uuid4())) slo_ms: Optional[float] = None num_inference_steps: Optional[int] = None @dataclass class RequestFuncOutput: success: bool = False latency: float = 0.0 error: str = "" start_time: float = 0.0 response_body: Dict[str, Any] = field(default_factory=dict) peak_memory_mb: float = 0.0 slo_achieved: Optional[bool] = None def is_dir_not_empty(path: str) -> bool: return os.path.isdir(path) and bool(os.listdir(path)) class BaseDataset(ABC): def __init__(self, args, api_url: str = "", model: str = ""): self.args = args self.api_url = api_url self.model = model self.items: List[Dict[str, Any]] = [] @abstractmethod def __len__(self) -> int: pass @abstractmethod def __getitem__(self, idx: int) -> RequestFuncInput: pass def get_requests(self) -> List[RequestFuncInput]: return [self[i] for i in range(len(self))] class VBenchDataset(BaseDataset): """ Dataset loader for VBench prompts. Supports t2v, i2v. """ T2V_PROMPT_URL = "https://raw.githubusercontent.com/Vchitect/VBench/master/prompts/prompts_per_dimension/subject_consistency.txt" I2V_DOWNLOAD_SCRIPT_URL = "https://raw.githubusercontent.com/Vchitect/VBench/master/vbench2_beta_i2v/download_data.sh" def __init__(self, args, api_url: str = "", model: str = ""): super().__init__(args, api_url, model) self.cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "sglang") self.items = self._load_data() def _load_data(self) -> List[Dict[str, Any]]: if self.args.task_name in ("text-to-video", "text-to-image", "video-to-video"): return self._load_t2v_prompts() elif self.args.task_name in ("image-to-video", "image-to-image"): return self._load_i2v_data() else: raise ValueError( f"Illegal task name is found in VBenchDataset {self.args.task_name}" ) def _download_file(self, url: str, dest_path: str) -> None: """Download a file from URL to destination path.""" os.makedirs(os.path.dirname(dest_path), exist_ok=True) resp = requests.get(url) resp.raise_for_status() with open(dest_path, "w") as f: f.write(resp.text) def _load_t2v_prompts(self) -> List[Dict[str, Any]]: path = self.args.dataset_path if not path: path = os.path.join(self.cache_dir, "vbench_subject_consistency.txt") if not os.path.exists(path): logger.info(f"Downloading VBench T2V prompts to {path}...") try: self._download_file(self.T2V_PROMPT_URL, path) except Exception as e: logger.info(f"Failed to download VBench prompts: {e}") return [{"prompt": "A cat sitting on a bench"}] * 50 prompts = [] with open(path, "r") as f: for line in f: line = line.strip() if line: prompts.append({"prompt": line}) return self._resize_data(prompts) def _auto_download_i2v_dataset(self) -> Optional[str]: """Auto-download VBench I2V dataset and return the dataset directory.""" vbench_i2v_dir = os.path.join(self.cache_dir, "vbench_i2v", "vbench2_beta_i2v") info_json_path = os.path.join(vbench_i2v_dir, "data", "i2v-bench-info.json") crop_dir = os.path.join(vbench_i2v_dir, "data", "crop") origin_dir = os.path.join(vbench_i2v_dir, "data", "origin") if ( os.path.exists(info_json_path) and is_dir_not_empty(crop_dir) and is_dir_not_empty(origin_dir) ): return vbench_i2v_dir logger.info(f"Downloading VBench I2V dataset to {vbench_i2v_dir}...") try: cache_root = os.path.join(self.cache_dir, "vbench_i2v") script_path = os.path.join(cache_root, "download_data.sh") self._download_file(self.I2V_DOWNLOAD_SCRIPT_URL, script_path) os.chmod(script_path, 0o755) logger.info("Executing download_data.sh (this may take a while)...") result = subprocess.run( ["bash", script_path], cwd=cache_root, capture_output=True, text=True, ) if result.returncode != 0: raise RuntimeError(f"Download script failed: {result.stderr}") missing_packages = re.findall(r"(\S+): command not found", result.stderr) if missing_packages: missing_packages = list(set(missing_packages)) package_list = ", ".join(f"'{cmd}'" for cmd in missing_packages) raise RuntimeError( f"Download script failed because the following commands are not installed: {package_list}.\n" "Please install them (e.g., on Ubuntu: `sudo apt install ...`) and try again." ) logger.info( f"Successfully downloaded VBench I2V dataset to {vbench_i2v_dir}" ) except Exception as e: logger.info(f"Failed to download VBench I2V dataset: {e}") logger.info("Please manually download following instructions at:") logger.info( "https://github.com/Vchitect/VBench/tree/master/vbench2_beta_i2v#22-download" ) return None return vbench_i2v_dir if os.path.exists(info_json_path) else None def _load_from_i2v_json(self, json_path: str) -> List[Dict[str, Any]]: """Load I2V data from i2v-bench-info.json format.""" with open(json_path, "r") as f: items = json.load(f) base_dir = os.path.dirname( os.path.dirname(json_path) ) # Go up to vbench2_beta_i2v origin_dir = os.path.join(base_dir, "data", "origin") data = [] for item in items: img_path = os.path.join(origin_dir, item.get("file_name", "")) if os.path.exists(img_path): data.append({"prompt": item.get("caption", ""), "image_path": img_path}) else: logger.warning(f"Image not found: {img_path}") logger.info(f"Loaded {len(data)} I2V samples from VBench I2V dataset") return data def _scan_directory_for_images(self, path: str) -> List[Dict[str, Any]]: """Scan directory for image files.""" exts = ["*.jpg", "*.jpeg", "*.png", "*.webp"] files = [] for ext in exts: files.extend(glob.glob(os.path.join(path, ext))) files.extend(glob.glob(os.path.join(path, ext.upper()))) origin_dir = os.path.join(path, "data", "origin") if os.path.exists(origin_dir): files.extend(glob.glob(os.path.join(origin_dir, ext))) files.extend(glob.glob(os.path.join(origin_dir, ext.upper()))) return [ {"prompt": os.path.splitext(os.path.basename(f))[0], "image_path": f} for f in files ] def _create_dummy_data(self) -> List[Dict[str, Any]]: """Create dummy data with a placeholder image in cache directory.""" logger.info("No I2V data found. Using dummy placeholders.") dummy_image = os.path.join(self.cache_dir, "dummy_image.jpg") if not os.path.exists(dummy_image): os.makedirs(self.cache_dir, exist_ok=True) img = Image.new("RGB", (100, 100), color="red") img.save(dummy_image) logger.info(f"Created dummy image at {dummy_image}") return [{"prompt": "A moving cat", "image_path": dummy_image}] * 10 def _load_i2v_data(self) -> List[Dict[str, Any]]: """Load I2V data from VBench I2V dataset or user-provided path.""" path = self.args.dataset_path if not path: path = self._auto_download_i2v_dataset() if not path: return self._resize_data(self._create_dummy_data()) info_json_candidates = [ os.path.join(path, "data", "i2v-bench-info.json"), path if path.endswith(".json") else None, ] for json_path in info_json_candidates: if json_path and os.path.exists(json_path): try: return self._resize_data(self._load_from_i2v_json(json_path)) except Exception as e: logger.info(f"Failed to load {json_path}: {e}") if os.path.isdir(path): data = self._scan_directory_for_images(path) if data: return self._resize_data(data) return self._resize_data(self._create_dummy_data()) def _resize_data(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """Resize data to match num_prompts.""" if not self.args.num_prompts: return data if len(data) < self.args.num_prompts: factor = (self.args.num_prompts // len(data)) + 1 data = data * factor return data[: self.args.num_prompts] def __len__(self) -> int: return len(self.items) def __getitem__(self, idx: int) -> RequestFuncInput: item = self.items[idx] return RequestFuncInput( prompt=item.get("prompt", ""), api_url=self.api_url, model=self.model, width=self.args.width, height=self.args.height, num_frames=self.args.num_frames, fps=self.args.fps, image_paths=[item["image_path"]] if "image_path" in item else None, ) class RandomDataset(BaseDataset): def __init__(self, args, api_url: str = "", model: str = ""): super().__init__(args, api_url, model) self.num_prompts = args.num_prompts or 100 def __len__(self) -> int: return self.num_prompts def __getitem__(self, idx: int) -> RequestFuncInput: return RequestFuncInput( prompt=f"Random prompt {idx} for benchmarking diffusion models", api_url=self.api_url, model=self.model, width=self.args.width, height=self.args.height, num_frames=self.args.num_frames, fps=self.args.fps, ) ================================================ FILE: python/sglang/multimodal_gen/configs/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # Configs for pipelines, and pipeline modules (in models folder) ================================================ FILE: python/sglang/multimodal_gen/configs/backend/vmoba/wan_1.3B_77_448_832.json ================================================ { "temporal_chunk_size": 2, "temporal_topk": 2, "spatial_chunk_size": [4, 13], "spatial_topk": 6, "st_chunk_size": [4, 4, 13], "st_topk": 18, "moba_select_mode": "topk", "moba_threshold": 0.25, "moba_threshold_type": "query_head", "first_full_layer": 0, "first_full_step": 12, "temporal_layer": 1, "spatial_layer": 1, "st_layer": 1 } ================================================ FILE: python/sglang/multimodal_gen/configs/backend/vmoba/wan_1.3B_77_480_832.json ================================================ { "temporal_chunk_size": 2, "temporal_topk": 3, "spatial_chunk_size": [3, 4], "spatial_topk": 20, "st_chunk_size": [4, 6, 4], "st_topk": 15, "moba_select_mode": "threshold", "moba_threshold": 0.25, "moba_threshold_type": "query_head", "first_full_layer": 0, "first_full_step": 12, "temporal_layer": 1, "spatial_layer": 1, "st_layer": 1 } ================================================ FILE: python/sglang/multimodal_gen/configs/models/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from sglang.multimodal_gen.configs.models.base import ModelConfig from sglang.multimodal_gen.configs.models.dits.base import DiTConfig from sglang.multimodal_gen.configs.models.encoders.base import EncoderConfig from sglang.multimodal_gen.configs.models.vaes.base import VAEConfig __all__ = ["ModelConfig", "VAEConfig", "DiTConfig", "EncoderConfig"] ================================================ FILE: python/sglang/multimodal_gen/configs/models/adapter/base.py ================================================ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from typing import Any from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum @dataclass class AdapterArchConfig(ArchConfig): _fsdp_shard_conditions: list = field(default_factory=list) _compile_conditions: list = field(default_factory=list) # convert weights name from HF-format to SGLang-dit-format param_names_mapping: dict = field(default_factory=dict) # Reverse mapping for saving checkpoints: custom -> hf reverse_param_names_mapping: dict = field(default_factory=dict) _supported_attention_backends: set[AttentionBackendEnum] = field( default_factory=lambda: { AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.FA, AttentionBackendEnum.AITER, AttentionBackendEnum.AITER_SAGE, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_3, } ) hidden_size: int = 0 num_attention_heads: int = 0 num_channels_latents: int = 0 exclude_lora_layers: list[str] = field(default_factory=list) boundary_ratio: float | None = None def __post_init__(self) -> None: if not self._compile_conditions: self._compile_conditions = self._fsdp_shard_conditions.copy() @dataclass class AdapterConfig(ModelConfig): arch_config: AdapterArchConfig = field(default_factory=AdapterArchConfig) # sglang-diffusion Adapter-specific parameters prefix: str = "" @staticmethod def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any: """Add CLI arguments for AdapterConfig fields""" parser.add_argument( f"--{prefix}.prefix", type=str, dest=f"{prefix.replace('-', '_')}.prefix", default=AdapterConfig.prefix, help="Prefix for the Adapter", ) return parser ================================================ FILE: python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py ================================================ from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.adapter.base import ( AdapterArchConfig, AdapterConfig, ) @dataclass class LTX2ConnectorArchConfig(AdapterArchConfig): audio_connector_attention_head_dim: int = 128 audio_connector_num_attention_heads: int = 30 audio_connector_num_layers: int = 2 audio_connector_num_learnable_registers: int = 128 caption_channels: int = 3840 causal_temporal_positioning: bool = False connector_rope_base_seq_len: int = 4096 rope_double_precision: bool = True rope_theta: float = 10000.0 rope_type: str = "split" text_proj_in_factor: int = 49 video_connector_attention_head_dim: int = 128 video_connector_num_attention_heads: int = 30 video_connector_num_layers: int = 2 video_connector_num_learnable_registers: int = 128 @dataclass class LTX2ConnectorConfig(AdapterConfig): arch_config: AdapterArchConfig = field(default_factory=LTX2ConnectorArchConfig) prefix: str = "LTX2" ================================================ FILE: python/sglang/multimodal_gen/configs/models/base.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field, fields from typing import Any, Dict from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # 1. ArchConfig contains all fields from diffuser's/transformer's config.json (i.e. all fields related to the architecture of the model) # 2. ArchConfig should be inherited & overridden by each model arch_config # 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users @dataclass class ArchConfig: stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=list ) # mapping from huggingface weight names to custom names extra_attrs: Dict[str, Any] = field(default_factory=dict) def __getattr__(self, name: str): d = object.__getattribute__(self, "__dict__") extras = d.get("extra_attrs") if extras is not None and name in extras: return extras[name] raise AttributeError( f"'{self.__class__.__name__}' object has no attribute '{name}'" ) def __setattr__(self, key, value): if key in type(self).__dataclass_fields__: object.__setattr__(self, key, value) else: d = object.__getattribute__(self, "__dict__") extras = d.get("extra_attrs") if extras is None: extras = {} d["extra_attrs"] = extras extras[key] = value @dataclass class ModelConfig: # Every model config parameter can be categorized into either ArchConfig or everything else # Diffuser/Transformer parameters arch_config: ArchConfig = field(default_factory=ArchConfig) # sglang-diffusion-specific parameters here # i.e. STA, quantization, teacache def __getattr__(self, name): # Only called if 'name' is not found in ModelConfig directly if hasattr(self.arch_config, name): return getattr(self.arch_config, name) raise AttributeError( f"'{type(self).__name__}' object has no attribute '{name}'" ) def __getstate__(self): # Return a dictionary of attributes to pickle # Convert to dict and exclude any problematic attributes state = self.__dict__.copy() return state def __setstate__(self, state): # Restore instance attributes from the unpickled state self.__dict__.update(state) # This should be used only when loading from transformers/diffusers def update_model_arch(self, source_model_dict: dict[str, Any]) -> None: """ Update arch_config with source_model_dict """ arch_config = self.arch_config for key, value in source_model_dict.items(): setattr(arch_config, key, value) if hasattr(arch_config, "__post_init__"): arch_config.__post_init__() def update_model_config(self, source_model_dict: dict[str, Any]) -> None: assert ( "arch_config" not in source_model_dict ), "Source model config shouldn't contain arch_config." valid_fields = {f.name for f in fields(self)} for key, value in source_model_dict.items(): if key in valid_fields: setattr(self, key, value) else: logger.warning( "%s does not contain field '%s'!", type(self).__name__, key ) raise AttributeError(f"Invalid field: {key}") if hasattr(self, "__post_init__"): self.__post_init__() ================================================ FILE: python/sglang/multimodal_gen/configs/models/bridges/__init__.py ================================================ # SPDX-License-Identifier: Apache-2.0 from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import ( MOVADualTowerConfig, ) __all__ = ["MOVADualTowerConfig"] ================================================ FILE: python/sglang/multimodal_gen/configs/models/bridges/mova_dual_tower.py ================================================ # SPDX-License-Identifier: Apache-2.0 """Configuration for MOVA dual tower bridge model.""" from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig def _is_conditioner_block(name: str, module) -> bool: """Check if module is a ConditionalCrossAttentionBlock.""" return "ConditionalCrossAttentionBlock" in type(module).__name__ @dataclass class MOVADualTowerArchConfig(DiTArchConfig): _fsdp_shard_conditions: list = field( default_factory=lambda: [_is_conditioner_block] ) # Model architecture parameters visual_layers: int = 40 audio_layers: int = 30 visual_hidden_dim: int = 5120 audio_hidden_dim: int = 1536 audio_fps: float = 50.0 head_dim: int = 128 interaction_strategy: str = "full" apply_cross_rope: bool = True apply_first_frame_bias_in_rope: bool = False trainable_condition_scale: bool = False pooled_adaln: bool = False eps: float = 1e-6 def __post_init__(self): super().__post_init__() self.hidden_size = self.visual_hidden_dim self.num_attention_heads = self.visual_hidden_dim // self.head_dim @dataclass class MOVADualTowerConfig(DiTConfig): arch_config: DiTArchConfig = field(default_factory=MOVADualTowerArchConfig) ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConfig from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig __all__ = [ "HeliosConfig", "HunyuanVideoConfig", "WanVideoConfig", "Hunyuan3DDiTConfig", "MOVAAudioConfig", "MOVAVideoConfig", ] ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/base.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from typing import Any from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum @dataclass class DiTArchConfig(ArchConfig): _fsdp_shard_conditions: list = field(default_factory=list) _compile_conditions: list = field(default_factory=list) # convert weights name from HF-format to SGLang-dit-format param_names_mapping: dict = field(default_factory=dict) # convert weights name from misc-format to HF-format # usually applicable if the LoRA is trained with official repo implementation lora_param_names_mapping: dict = field(default_factory=dict) # Reverse mapping for saving checkpoints: custom -> hf reverse_param_names_mapping: dict = field(default_factory=dict) _supported_attention_backends: set[AttentionBackendEnum] = field( default_factory=lambda: { AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.FA, AttentionBackendEnum.AITER, AttentionBackendEnum.AITER_SAGE, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_3, } ) hidden_size: int = 0 num_attention_heads: int = 0 num_channels_latents: int = 0 exclude_lora_layers: list[str] = field(default_factory=list) boundary_ratio: float | None = None def __post_init__(self) -> None: if not self._compile_conditions: self._compile_conditions = self._fsdp_shard_conditions.copy() @dataclass class DiTConfig(ModelConfig): arch_config: DiTArchConfig = field(default_factory=DiTArchConfig) # sglang-diffusion DiT-specific parameters prefix: str = "" quant_config: QuantizationConfig | None = None @staticmethod def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any: """Add CLI arguments for DiTConfig fields""" parser.add_argument( f"--{prefix}.prefix", type=str, dest=f"{prefix.replace('-', '_')}.prefix", default=DiTConfig.prefix, help="Prefix for the DiT model", ) parser.add_argument( f"--{prefix}.quant-config", type=str, dest=f"{prefix.replace('-', '_')}.quant_config", default=None, help="Quantization configuration for the DiT model", ) return parser ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/flux.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from typing import Tuple from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig @dataclass class FluxArchConfig(DiTArchConfig): patch_size: int = 1 in_channels: int = 64 out_channels: int | None = None num_layers: int = 19 num_single_layers: int = 38 attention_head_dim: int = 128 num_attention_heads: int = 24 joint_attention_dim: int = 4096 pooled_projection_dim: int = 768 guidance_embeds: bool = False axes_dims_rope: Tuple[int, int, int] = (16, 56, 56) stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list) exclude_lora_layers: list[str] = field( default_factory=lambda: [ "time_guidance_embed.timestep_embedder.linear_1", "time_guidance_embed.timestep_embedder.linear_2", "time_guidance_embed.guidance_embedder.linear_1", "time_guidance_embed.guidance_embedder.linear_2", ] ) # nunchaku checkpoint uses different weight names; map to sglang flux layout param_names_mapping: dict = field( default_factory=lambda: { # HF diffusers format r"^transformer\.(\w*)\.(.*)$": r"\1.\2", # transformer_blocks nunchaku format (raw export - before internal conversion) r"^transformer_blocks\.(\d+)\.mlp_fc1\.(.*)$": r"transformer_blocks.\1.ff.net.0.proj.\2", r"^transformer_blocks\.(\d+)\.mlp_fc2\.(.*)$": r"transformer_blocks.\1.ff.net.2.\2", r"^transformer_blocks\.(\d+)\.mlp_context_fc1\.(.*)$": r"transformer_blocks.\1.ff_context.net.0.proj.\2", r"^transformer_blocks\.(\d+)\.mlp_context_fc2\.(.*)$": r"transformer_blocks.\1.ff_context.net.2.\2", r"^transformer_blocks\.(\d+)\.qkv_proj\.(.*)$": r"transformer_blocks.\1.attn.to_qkv.\2", r"^transformer_blocks\.(\d+)\.qkv_proj_context\.(.*)$": r"transformer_blocks.\1.attn.to_added_qkv.\2", r"^transformer_blocks\.(\d+)\.out_proj\.(.*)$": r"transformer_blocks.\1.attn.to_out.0.\2", r"^transformer_blocks\.(\d+)\.out_proj_context\.(.*)$": r"transformer_blocks.\1.attn.to_add_out.\2", r"^transformer_blocks\.(\d+)\.norm_q\.(.*)$": r"transformer_blocks.\1.attn.norm_q.\2", r"^transformer_blocks\.(\d+)\.norm_k\.(.*)$": r"transformer_blocks.\1.attn.norm_k.\2", r"^transformer_blocks\.(\d+)\.norm_added_q\.(.*)$": r"transformer_blocks.\1.attn.norm_added_q.\2", r"^transformer_blocks\.(\d+)\.norm_added_k\.(.*)$": r"transformer_blocks.\1.attn.norm_added_k.\2", # transformer_blocks nunchaku format (already converted with convert_flux_state_dict) r"^transformer_blocks\.(\d+)\.attn\.add_qkv_proj\.(.*)$": r"transformer_blocks.\1.attn.to_added_qkv.\2", # single_transformer_blocks nunchaku format (raw export - before internal conversion) r"^single_transformer_blocks\.(\d+)\.qkv_proj\.(.*)$": r"single_transformer_blocks.\1.attn.to_qkv.\2", r"^single_transformer_blocks\.(\d+)\.out_proj\.(.*)$": r"single_transformer_blocks.\1.attn.to_out.0.\2", r"^single_transformer_blocks\.(\d+)\.norm_q\.(.*)$": r"single_transformer_blocks.\1.attn.norm_q.\2", r"^single_transformer_blocks\.(\d+)\.norm_k\.(.*)$": r"single_transformer_blocks.\1.attn.norm_k.\2", # nunchaku quantization parameter name conversions (apply to all blocks) r"^(.*)\.smooth_orig$": r"\1.smooth_factor_orig", r"^(.*)\.smooth$": r"\1.smooth_factor", r"^(.*)\.lora_down$": r"\1.proj_down", r"^(.*)\.lora_up$": r"\1.proj_up", } ) def __post_init__(self): super().__post_init__() self.out_channels = self.out_channels or self.in_channels self.hidden_size = self.num_attention_heads * self.attention_head_dim self.num_channels_latents = self.out_channels @dataclass class FluxConfig(DiTConfig): arch_config: DiTArchConfig = field(default_factory=FluxArchConfig) prefix: str = "Flux" ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/glmimage.py ================================================ from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig @dataclass class GlmImageArchConfig(DiTArchConfig): patch_size: int = 2 in_channels: int = 16 out_channels: int | None = 16 num_layers: int = 30 attention_head_dim: int = 128 num_attention_heads: int = 32 condition_dim: int = 256 prior_vq_quantizer_codebook_size: int = 16384 text_embed_dim: int = 1472 time_embed_dim: int = 512 stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list) param_names_mapping: dict = field( default_factory=lambda: { # LoRA mappings r"^(transformer_blocks\.\d+\.attn\..*\.lora_[AB])\.default$": r"\1", } ) def __post_init__(self): super().__post_init__() self.out_channels = self.out_channels or self.in_channels self.hidden_size = self.num_attention_heads * self.attention_head_dim self.num_channels_latents = self.out_channels @dataclass class GlmImageDitConfig(DiTConfig): arch_config: DiTArchConfig = field(default_factory=GlmImageArchConfig) prefix: str = "glmimage" ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/helios.py ================================================ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig def is_blocks(n: str, m) -> bool: return "blocks" in n and str.isdigit(n.split(".")[-1]) @dataclass class HeliosArchConfig(DiTArchConfig): _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks]) param_names_mapping: dict = field( default_factory=lambda: { # Patch embeddings r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1", # Condition embedder: text r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1", r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1", # Condition embedder: time r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1", r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1", r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1", # Blocks: self-attention (keep attn1. prefix, drop .0. from to_out) r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.attn1.to_out.\2", # Blocks: cross-attention output (drop .0. from to_out) r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2", # Blocks: feed-forward r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2", r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", # Blocks: cross-attn residual norm r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.self_attn_residual_norm.\2", } ) reverse_param_names_mapping: dict = field(default_factory=lambda: {}) lora_param_names_mapping: dict = field(default_factory=lambda: {}) patch_size: tuple[int, int, int] = (1, 2, 2) text_len: int = 226 num_attention_heads: int = 40 attention_head_dim: int = 128 in_channels: int = 16 out_channels: int = 16 text_dim: int = 4096 freq_dim: int = 256 ffn_dim: int = 13824 num_layers: int = 40 cross_attn_norm: bool = True qk_norm: str = "rms_norm_across_heads" eps: float = 1e-6 added_kv_proj_dim: int | None = None rope_max_seq_len: int = 1024 pos_embed_seq_len: int | None = None exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"]) # Helios-specific rope_dim: tuple[int, int, int] = (44, 42, 42) rope_theta: float = 10000.0 guidance_cross_attn: bool = True zero_history_timestep: bool = True has_multi_term_memory_patch: bool = True is_amplify_history: bool = False history_scale_mode: str = "per_head" def __post_init__(self): super().__post_init__() self.out_channels = self.out_channels or self.in_channels self.hidden_size = self.num_attention_heads * self.attention_head_dim self.num_channels_latents = self.out_channels @dataclass class HeliosConfig(DiTConfig): arch_config: DiTArchConfig = field(default_factory=HeliosArchConfig) prefix: str = "Helios" ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/hunyuan3d.py ================================================ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig @dataclass class Hunyuan3DDiTArchConfig(DiTArchConfig): """Architecture config for Hunyuan3D DiT (Flux-style for Hunyuan3D-2.0).""" param_names_mapping: dict = field( default_factory=lambda: { r"(.*)\.img_mlp\.0\.(.*)$": r"\1.img_mlp.fc_in.\2", r"(.*)\.img_mlp\.2\.(.*)$": r"\1.img_mlp.fc_out.\2", r"(.*)\.txt_mlp\.0\.(.*)$": r"\1.txt_mlp.fc_in.\2", r"(.*)\.txt_mlp\.2\.(.*)$": r"\1.txt_mlp.fc_out.\2", } ) in_channels: int = 64 hidden_size: int = 1024 num_attention_heads: int = 16 num_layers: int = 16 num_single_layers: int = 32 mlp_ratio: float = 4.0 context_in_dim: int = 1536 axes_dim: tuple[int, ...] = (64,) theta: int = 10000 qkv_bias: bool = True guidance_embed: bool = False time_factor: float = 1000.0 def __post_init__(self) -> None: if self.num_channels_latents == 0: self.num_channels_latents = self.in_channels super().__post_init__() @dataclass class Hunyuan3DDiTConfig(DiTConfig): """DiT configuration for Hunyuan3D shape generation (Flux-style).""" arch_config: Hunyuan3DDiTArchConfig = field(default_factory=Hunyuan3DDiTArchConfig) subfolder: str = "hunyuan3d-dit-v2-0" ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/hunyuanvideo.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field import torch from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig def is_double_block(n: str, m) -> bool: return "double" in n and str.isdigit(n.split(".")[-1]) def is_single_block(n: str, m) -> bool: return "single" in n and str.isdigit(n.split(".")[-1]) def is_refiner_block(n: str, m) -> bool: return "refiner" in n and str.isdigit(n.split(".")[-1]) def is_txt_in(n: str, m) -> bool: return n.split(".")[-1] == "txt_in" @dataclass class HunyuanVideoArchConfig(DiTArchConfig): _fsdp_shard_conditions: list = field( default_factory=lambda: [is_double_block, is_single_block, is_refiner_block] ) _compile_conditions: list = field( default_factory=lambda: [is_double_block, is_single_block, is_txt_in] ) param_names_mapping: dict = field( default_factory=lambda: { # 1. context_embedder.time_text_embed submodules (specific rules, applied first): r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$": r"txt_in.t_embedder.mlp.fc_in.\1", r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_2\.(.*)$": r"txt_in.t_embedder.mlp.fc_out.\1", r"^context_embedder\.proj_in\.(.*)$": r"txt_in.input_embedder.\1", r"^context_embedder\.time_text_embed\.text_embedder\.linear_1\.(.*)$": r"txt_in.c_embedder.fc_in.\1", r"^context_embedder\.time_text_embed\.text_embedder\.linear_2\.(.*)$": r"txt_in.c_embedder.fc_out.\1", r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm1\.(.*)$": r"txt_in.refiner_blocks.\1.norm1.\2", r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm2\.(.*)$": r"txt_in.refiner_blocks.\1.norm2.\2", r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_q\.(.*)$": ( r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 0, 3, ), r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_k\.(.*)$": ( r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 1, 3, ), r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_v\.(.*)$": ( r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 2, 3, ), r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"txt_in.refiner_blocks.\1.self_attn_proj.\2", r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_in.\2", r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_out.\2", r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear\.(.*)$": r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2", # 3. x_embedder mapping: r"^x_embedder\.proj\.(.*)$": r"img_in.proj.\1", # 4. Top-level time_text_embed mappings: r"^time_text_embed\.timestep_embedder\.linear_1\.(.*)$": r"time_in.mlp.fc_in.\1", r"^time_text_embed\.timestep_embedder\.linear_2\.(.*)$": r"time_in.mlp.fc_out.\1", r"^time_text_embed\.guidance_embedder\.linear_1\.(.*)$": r"guidance_in.mlp.fc_in.\1", r"^time_text_embed\.guidance_embedder\.linear_2\.(.*)$": r"guidance_in.mlp.fc_out.\1", r"^time_text_embed\.text_embedder\.linear_1\.(.*)$": r"vector_in.fc_in.\1", r"^time_text_embed\.text_embedder\.linear_2\.(.*)$": r"vector_in.fc_out.\1", # 5. transformer_blocks mapping: r"^transformer_blocks\.(\d+)\.norm1\.linear\.(.*)$": r"double_blocks.\1.img_mod.linear.\2", r"^transformer_blocks\.(\d+)\.norm1_context\.linear\.(.*)$": r"double_blocks.\1.txt_mod.linear.\2", r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": r"double_blocks.\1.img_attn_q_norm.\2", r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": r"double_blocks.\1.img_attn_k_norm.\2", r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": ( r"double_blocks.\1.img_attn_qkv.\2", 0, 3, ), r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": ( r"double_blocks.\1.img_attn_qkv.\2", 1, 3, ), r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": ( r"double_blocks.\1.img_attn_qkv.\2", 2, 3, ), r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$": ( r"double_blocks.\1.txt_attn_qkv.\2", 0, 3, ), r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$": ( r"double_blocks.\1.txt_attn_qkv.\2", 1, 3, ), r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$": ( r"double_blocks.\1.txt_attn_qkv.\2", 2, 3, ), r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"double_blocks.\1.img_attn_proj.\2", # Corrected: merge attn.to_add_out into the main projection. r"^transformer_blocks\.(\d+)\.attn\.to_add_out\.(.*)$": r"double_blocks.\1.txt_attn_proj.\2", r"^transformer_blocks\.(\d+)\.attn\.norm_added_q\.(.*)$": r"double_blocks.\1.txt_attn_q_norm.\2", r"^transformer_blocks\.(\d+)\.attn\.norm_added_k\.(.*)$": r"double_blocks.\1.txt_attn_k_norm.\2", r"^transformer_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$": r"double_blocks.\1.img_mlp.fc_in.\2", r"^transformer_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$": r"double_blocks.\1.img_mlp.fc_out.\2", r"^transformer_blocks\.(\d+)\.ff_context\.net\.0(?:\.proj)?\.(.*)$": r"double_blocks.\1.txt_mlp.fc_in.\2", r"^transformer_blocks\.(\d+)\.ff_context\.net\.2(?:\.proj)?\.(.*)$": r"double_blocks.\1.txt_mlp.fc_out.\2", # 6. single_transformer_blocks mapping: r"^single_transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": r"single_blocks.\1.q_norm.\2", r"^single_transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": r"single_blocks.\1.k_norm.\2", r"^single_transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": ( r"single_blocks.\1.linear1.\2", 0, 4, ), r"^single_transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": ( r"single_blocks.\1.linear1.\2", 1, 4, ), r"^single_transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": ( r"single_blocks.\1.linear1.\2", 2, 4, ), r"^single_transformer_blocks\.(\d+)\.proj_mlp\.(.*)$": ( r"single_blocks.\1.linear1.\2", 3, 4, ), # Corrected: map proj_out to modulation.linear rather than a separate proj_out branch. r"^single_transformer_blocks\.(\d+)\.proj_out\.(.*)$": r"single_blocks.\1.linear2.\2", r"^single_transformer_blocks\.(\d+)\.norm\.linear\.(.*)$": r"single_blocks.\1.modulation.linear.\2", # 7. Final layers mapping: r"^norm_out\.linear\.(.*)$": r"final_layer.adaLN_modulation.linear.\1", r"^proj_out\.(.*)$": r"final_layer.linear.\1", } ) reverse_param_names_mapping: dict = field(default_factory=lambda: {}) patch_size: int = 2 patch_size_t: int = 1 in_channels: int = 16 out_channels: int = 16 num_attention_heads: int = 24 attention_head_dim: int = 128 mlp_ratio: float = 4.0 num_layers: int = 20 num_single_layers: int = 40 num_refiner_layers: int = 2 rope_axes_dim: tuple[int, int, int] = (16, 56, 56) guidance_embeds: bool = False dtype: torch.dtype | None = None text_embed_dim: int = 4096 pooled_projection_dim: int = 768 rope_theta: int = 256 qk_norm: str = "rms_norm" exclude_lora_layers: list[str] = field( default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"] ) def __post_init__(self): super().__post_init__() self.hidden_size: int = self.attention_head_dim * self.num_attention_heads self.num_channels_latents: int = self.in_channels @dataclass class HunyuanVideoConfig(DiTConfig): arch_config: DiTArchConfig = field(default_factory=HunyuanVideoArchConfig) prefix: str = "Hunyuan" ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/ltx_2.py ================================================ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from enum import Enum from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig class LTXModelType(Enum): """ Model type enum mirroring upstream `LTXModelType`. Upstream reference: - `LTX-2/packages/ltx-core/src/ltx_core/model/transformer/model.py::LTXModelType` """ AudioVideo = "ltx av model" VideoOnly = "ltx video only model" AudioOnly = "ltx audio only model" def is_video_enabled(self) -> bool: return self in (LTXModelType.AudioVideo, LTXModelType.VideoOnly) def is_audio_enabled(self) -> bool: return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly) class LTX2RopeType(str, Enum): """ Minimal RoPE type enum mirroring LTX-2 upstream `LTXRopeType`. Upstream reference: - `LTX-2/packages/ltx-core/src/ltx_core/model/transformer/rope.py::LTXRopeType` """ INTERLEAVED = "interleaved" SPLIT = "split" class LTX2AttentionFunction(str, Enum): """ Placeholder enum for upstream `AttentionFunction.DEFAULT`. Upstream reference: - `LTX-2/packages/ltx-core/src/ltx_core/model/transformer/attention.py` """ DEFAULT = "default" def is_blocks(n: str, m) -> bool: return "blocks" in n and str.isdigit(n.split(".")[-1]) @dataclass class LTX2ArchConfig(DiTArchConfig): """Architecture configuration for LTX-2 Video Transformer.""" _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks]) param_names_mapping: dict = field( default_factory=lambda: { # Parameter name mappings from HuggingFace checkpoint keys to SGLang module names. # We use upstream variable names (patchify_proj, adaln_single) but HF uses different keys. # # HF key -> SGLang key (upstream naming) r"^proj_in\.(.*)$": r"patchify_proj.\1", r"^time_embed\.(.*)$": r"adaln_single.\1", r"^audio_proj_in\.(.*)$": r"audio_patchify_proj.\1", r"^audio_time_embed\.(.*)$": r"audio_adaln_single.\1", # FeedForward r"(.*)ff\.net\.0\.proj\.(.*)$": r"\1ff.proj_in.\2", r"(.*)ff\.net\.2\.(.*)$": r"\1ff.proj_out.\2", # Attention Norms r"(.*)\.norm_q\.(.*)$": r"\1.q_norm.\2", r"(.*)\.norm_k\.(.*)$": r"\1.k_norm.\2", # Scale Shift Tables (Global) r"^av_cross_attn_video_scale_shift\.(.*)$": r"av_ca_video_scale_shift_adaln_single.\1", r"^av_cross_attn_audio_scale_shift\.(.*)$": r"av_ca_audio_scale_shift_adaln_single.\1", r"^av_cross_attn_video_a2v_gate\.(.*)$": r"av_ca_a2v_gate_adaln_single.\1", r"^av_cross_attn_audio_v2a_gate\.(.*)$": r"av_ca_v2a_gate_adaln_single.\1", # Scale Shift Tables (Block Level) # HF: scale_shift_table_a2v_ca_video -> SGLang: video_a2v_cross_attn_scale_shift_table r"(.*)scale_shift_table_a2v_ca_video": r"\1video_a2v_cross_attn_scale_shift_table", r"(.*)scale_shift_table_a2v_ca_audio": r"\1audio_a2v_cross_attn_scale_shift_table", } ) reverse_param_names_mapping: dict = field( default_factory=lambda: { # Reverse mapping: SGLang module names -> HF checkpoint keys (for saving). r"^patchify_proj\.(.*)$": r"proj_in.\1", r"^adaln_single\.(.*)$": r"time_embed.\1", r"^audio_patchify_proj\.(.*)$": r"audio_proj_in.\1", r"^audio_adaln_single\.(.*)$": r"audio_time_embed.\1", # FeedForward r"(.*)ff\.proj_in\.(.*)$": r"\1ff.net.0.proj.\2", r"(.*)ff\.proj_out\.(.*)$": r"\1ff.net.2.\2", # Attention Norms r"(.*)\.q_norm\.(.*)$": r"\1.norm_q.\2", r"(.*)\.k_norm\.(.*)$": r"\1.norm_k.\2", # Scale Shift Tables (Global) r"^av_ca_video_scale_shift_adaln_single\.(.*)$": r"av_cross_attn_video_scale_shift.\1", r"^av_ca_audio_scale_shift_adaln_single\.(.*)$": r"av_cross_attn_audio_scale_shift.\1", r"^av_ca_a2v_gate_adaln_single\.(.*)$": r"av_cross_attn_video_a2v_gate.\1", r"^av_ca_v2a_gate_adaln_single\.(.*)$": r"av_cross_attn_audio_v2a_gate.\1", # Scale Shift Tables (Block Level) # SGLang: video_a2v_cross_attn_scale_shift_table -> HF: scale_shift_table_a2v_ca_video r"(.*)video_a2v_cross_attn_scale_shift_table": r"\1scale_shift_table_a2v_ca_video", r"(.*)audio_a2v_cross_attn_scale_shift_table": r"\1scale_shift_table_a2v_ca_audio", } ) lora_param_names_mapping: dict = field( default_factory=lambda: { # LoRA parameter name mappings from official repo format to HF format. # This is applied before param_names_mapping when loading LoRA adapters. # Will be populated if LoRA adapters use different naming conventions. } ) # Model type and attention configuration model_type: LTXModelType = LTXModelType.AudioVideo attention_type: LTX2AttentionFunction = LTX2AttentionFunction.DEFAULT rope_type: LTX2RopeType = LTX2RopeType.INTERLEAVED double_precision_rope: bool = False # Video parameters num_attention_heads: int = 32 attention_head_dim: int = 128 in_channels: int = 128 out_channels: int = 128 num_layers: int = 48 cross_attention_dim: int = 4096 norm_eps: float = 1e-6 caption_channels: int = 3840 positional_embedding_theta: float = 10000.0 positional_embedding_max_pos: list[int] | None = None timestep_scale_multiplier: int = 1000 use_middle_indices_grid: bool = True # Audio parameters audio_num_attention_heads: int = 32 audio_attention_head_dim: int = 64 audio_in_channels: int = 128 audio_out_channels: int = 128 audio_cross_attention_dim: int = 2048 audio_positional_embedding_max_pos: list[int] | None = None av_ca_timestep_scale_multiplier: int = 1 # SGLang-specific parameters patch_size: tuple[int, int, int] = (1, 2, 2) text_len: int = 512 def __post_init__(self): super().__post_init__() # Video derived values self.hidden_size = self.num_attention_heads * self.attention_head_dim self.num_channels_latents = self.out_channels if self.positional_embedding_max_pos is None: self.positional_embedding_max_pos = [20, 2048, 2048] # Audio derived values self.audio_hidden_size = ( self.audio_num_attention_heads * self.audio_attention_head_dim ) if self.audio_positional_embedding_max_pos is None: self.audio_positional_embedding_max_pos = [20] @dataclass class LTX2Config(DiTConfig): """Configuration for LTX-2 Video Transformer.""" arch_config: LTX2ArchConfig = field(default_factory=LTX2ArchConfig) prefix: str = "ltx2" ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/mova_audio.py ================================================ # Copied and adapted from: mossVG/mova/diffusion/models/wan_audio_dit.py # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig def _is_blocks(n: str, m) -> bool: return "blocks" in n and str.isdigit(n.split(".")[-1]) @dataclass class MOVAAudioArchConfig(DiTArchConfig): _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_blocks]) param_names_mapping: dict = field( default_factory=lambda: { r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_norm.\2", r"^text_embedding\.0\.(.*)$": r"text_embedding.fc_in.\1", r"^text_embedding\.2\.(.*)$": r"text_embedding.fc_out.\1", r"^time_embedding\.0\.(.*)$": r"time_embedding.fc_in.\1", r"^time_embedding\.2\.(.*)$": r"time_embedding.fc_out.\1", r"^img_emb\.proj\.1\.(.*)$": r"img_emb.fc_in.\1", r"^img_emb\.proj\.3\.(.*)$": r"img_emb.fc_out.\1", } ) reverse_param_names_mapping: dict = field(default_factory=dict) lora_param_names_mapping: dict = field(default_factory=dict) dim: int = 1536 in_dim: int = 128 ffn_dim: int = 6144 out_dim: int = 128 text_dim: int = 4096 freq_dim: int = 256 eps: float = 1e-6 patch_size: tuple[int, int, int] = (1, 2, 2) num_heads: int = 12 num_layers: int = 30 has_image_input: bool = False has_image_pos_emb: bool = False has_ref_conv: bool = False add_control_adapter: bool = False in_dim_control_adapter: int = 24 separated_timestep: bool = False require_vae_embedding: bool = False require_clip_embedding: bool = False fuse_vae_embedding_in_latents: bool = False vae_type: str = "dac" def __post_init__(self): super().__post_init__() self.hidden_size = self.dim self.num_attention_heads = self.num_heads self.num_channels_latents = self.out_dim assert ( not self.has_image_input ), "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)." @dataclass class MOVAAudioConfig(DiTConfig): arch_config: DiTArchConfig = field(default_factory=MOVAAudioArchConfig) prefix: str = "mova_audio" ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/mova_video.py ================================================ # Copied and adapted from: mossVG/mova/diffusion/models/wan_video_dit.py # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig def _is_blocks(n: str, m) -> bool: return "blocks" in n and str.isdigit(n.split(".")[-1]) @dataclass class MOVAVideoArchConfig(DiTArchConfig): _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_blocks]) param_names_mapping: dict = field( default_factory=lambda: { r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_norm.\2", r"^text_embedding\.0\.(.*)$": r"text_embedding.fc_in.\1", r"^text_embedding\.2\.(.*)$": r"text_embedding.fc_out.\1", r"^time_embedding\.0\.(.*)$": r"time_embedding.fc_in.\1", r"^time_embedding\.2\.(.*)$": r"time_embedding.fc_out.\1", r"^img_emb\.proj\.1\.(.*)$": r"img_emb.fc_in.\1", r"^img_emb\.proj\.3\.(.*)$": r"img_emb.fc_out.\1", } ) reverse_param_names_mapping: dict = field(default_factory=dict) lora_param_names_mapping: dict = field(default_factory=dict) dim: int = 5120 in_dim: int = 16 ffn_dim: int = 13824 out_dim: int = 16 text_dim: int = 4096 freq_dim: int = 256 eps: float = 1e-6 patch_size: tuple[int, int, int] = (1, 2, 2) num_heads: int = 40 num_layers: int = 40 has_image_input: bool = False has_image_pos_emb: bool = False has_ref_conv: bool = False add_control_adapter: bool = False in_dim_control_adapter: int = 24 separated_timestep: bool = False require_vae_embedding: bool = True require_clip_embedding: bool = True fuse_vae_embedding_in_latents: bool = False def __post_init__(self): super().__post_init__() self.hidden_size = self.dim self.num_attention_heads = self.num_heads self.num_channels_latents = self.out_dim assert ( not self.has_image_input ), "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)." @dataclass class MOVAVideoConfig(DiTConfig): arch_config: DiTArchConfig = field(default_factory=MOVAVideoArchConfig) prefix: str = "mova_video" ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/qwenimage.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from typing import Tuple from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig @dataclass class QwenImageArchConfig(DiTArchConfig): patch_size: int = 1 in_channels: int = 64 out_channels: int | None = None num_layers: int = 19 num_single_layers: int = 38 attention_head_dim: int = 128 num_attention_heads: int = 24 joint_attention_dim: int = 4096 pooled_projection_dim: int = 768 guidance_embeds: bool = False axes_dims_rope: Tuple[int, int, int] = (16, 56, 56) zero_cond_t: bool = False stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list) param_names_mapping: dict = field( default_factory=lambda: { # LoRA mappings r"^(transformer_blocks\.\d+\.attn\..*\.lora_[AB])\.default$": r"\1", # SVDquant mappings r"(.*)\.add_qkv_proj\.(.+)$": r"\1.to_added_qkv.\2", r"(transformer_blocks\.\d+\.(img_mlp|txt_mlp)\..*\.(smooth_factor_orig|wcscales))$": r"\1", r".*\.wtscale$": r"", } ) def __post_init__(self): super().__post_init__() self.out_channels = self.out_channels or self.in_channels self.hidden_size = self.num_attention_heads * self.attention_head_dim self.num_channels_latents = self.out_channels @dataclass class QwenImageEditPlus_2511_ArchConfig(QwenImageArchConfig): zero_cond_t: bool = True @dataclass class QwenImageDitConfig(DiTConfig): arch_config: DiTArchConfig = field(default_factory=QwenImageArchConfig) prefix: str = "qwenimage" @dataclass class QwenImageEditPlus_2511_DitConfig(DiTConfig): arch_config: DiTArchConfig = field( default_factory=QwenImageEditPlus_2511_ArchConfig ) prefix: str = "qwenimageedit" ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/sana.py ================================================ # SPDX-License-Identifier: Apache-2.0 # # Architecture and model configuration for SANA DiT (Diffusion Transformer). # # SANA uses a linear-attention-based transformer that replaces standard # quadratic self-attention with ReLU-based linear attention, enabling # efficient high-resolution image synthesis. Cross-attention (standard SDPA) # is used for text conditioning via Gemma2 embeddings. # # Defaults below correspond to the SANA-1.6B / 1024px variant. # For 4.8B, override num_layers=36, num_attention_heads=64, etc. # # Reference: https://arxiv.org/abs/2410.10629 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig @dataclass class SanaArchConfig(DiTArchConfig): patch_size: int = 1 in_channels: int = 32 out_channels: int = 32 num_layers: int = 20 attention_head_dim: int = 32 num_attention_heads: int = 70 num_cross_attention_heads: int = 20 cross_attention_head_dim: int = 112 cross_attention_dim: int = 2240 caption_channels: int = 2304 mlp_ratio: float = 2.5 # "rms_norm_across_heads" applies RMSNorm over the full (num_heads * head_dim) qk_norm: str = "rms_norm_across_heads" norm_elementwise_affine: bool = False norm_eps: float = 1e-6 sample_size: int = 32 guidance_embeds: bool = False param_names_mapping: dict = field( default_factory=lambda: { r"^transformer\.(.*)$": r"\1", } ) def __post_init__(self): super().__post_init__() self.hidden_size = self.num_attention_heads * self.attention_head_dim self.num_channels_latents = self.out_channels @dataclass class SanaConfig(DiTConfig): arch_config: DiTArchConfig = field(default_factory=SanaArchConfig) prefix: str = "Sana" ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/wanvideo.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig def is_blocks(n: str, m) -> bool: return "blocks" in n and str.isdigit(n.split(".")[-1]) @dataclass class WanVideoArchConfig(DiTArchConfig): _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks]) param_names_mapping: dict = field( default_factory=lambda: { r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1", r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1", r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1", r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1", r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1", r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1", r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$": r"condition_embedder.image_embedder.ff.fc_in.\1", r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$": r"condition_embedder.image_embedder.ff.fc_out.\1", r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$": r"blocks.\1.to_q.\2", r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$": r"blocks.\1.to_k.\2", r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$": r"blocks.\1.to_v.\2", r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.to_out.\2", r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2", r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2", r"^blocks\.(\d+)\.attn1\.attn_op\.local_attn\.proj_l\.(.*)$": r"blocks.\1.attn1.local_attn.proj_l.\2", r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2", r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2", r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2", } ) reverse_param_names_mapping: dict = field(default_factory=lambda: {}) # Some LoRA adapters use the original official layer names instead of hf layer names, # so apply this before the param_names_mapping lora_param_names_mapping: dict = field( default_factory=lambda: { r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2", r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2", r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2", r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.attn1.to_out.0.\2", r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2", r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2", r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2", r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.0.\2", r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", } ) patch_size: tuple[int, int, int] = (1, 2, 2) text_len = 512 num_attention_heads: int = 40 attention_head_dim: int = 128 in_channels: int = 16 out_channels: int = 16 text_dim: int = 4096 freq_dim: int = 256 ffn_dim: int = 13824 num_layers: int = 40 cross_attn_norm: bool = True qk_norm: str = "rms_norm_across_heads" eps: float = 1e-6 image_dim: int | None = None added_kv_proj_dim: int | None = None rope_max_seq_len: int = 1024 pos_embed_seq_len: int | None = None exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"]) # Wan MoE boundary_ratio: float | None = None # Causal Wan local_attn_size: int = ( -1 ) # Window size for temporal local attention (-1 indicates global attention) sink_size: int = ( 0 # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache ) num_frames_per_block: int = 3 sliding_window_num_frames: int = 21 attention_type: str = "original" sla_topk: float = 0.1 def __post_init__(self): super().__post_init__() self.out_channels = self.out_channels or self.in_channels self.hidden_size = self.num_attention_heads * self.attention_head_dim self.num_channels_latents = self.out_channels @dataclass class WanVideoConfig(DiTConfig): arch_config: DiTArchConfig = field(default_factory=WanVideoArchConfig) prefix: str = "Wan" ================================================ FILE: python/sglang/multimodal_gen/configs/models/dits/zimage.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from typing import Tuple from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig def is_zimage_layer(n: str, m) -> bool: """Returns if the module should be sharded for Z-Image model.""" if "layers" in n and str.isdigit(n.split(".")[-1]): return True if ("noise_refiner" in n or "context_refiner" in n) and str.isdigit( n.split(".")[-1] ): return True return False @dataclass class ZImageArchConfig(DiTArchConfig): all_patch_size: Tuple[int, ...] = (2,) all_f_patch_size: Tuple[int, ...] = (1,) in_channels: int = 16 out_channels: int | None = None dim: int = 3840 num_layers: int = 30 n_refiner_layers: int = 2 num_attention_heads: int = 30 n_kv_heads: int = 30 norm_eps: float = 1e-5 qk_norm: bool = True cap_feat_dim: int = 2560 rope_theta: float = 256.0 t_scale: float = 1000.0 axes_dims: Tuple[int, int, int] = (32, 48, 48) axes_lens: Tuple[int, int, int] = (1024, 512, 512) _fsdp_shard_conditions: list = field(default_factory=lambda: [is_zimage_layer]) stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=lambda: [ # (param_name, shard_name, shard_id) (".feed_forward.w13", ".feed_forward.w1", "gate"), (".feed_forward.w13", ".feed_forward.w3", "up"), ] ) param_names_mapping: dict = field( default_factory=lambda: { r"(.*)\.feed_forward\.w1\.weight$": (r"\1.feed_forward.w13.weight", 0, 2), r"(.*)\.feed_forward\.w3\.weight$": (r"\1.feed_forward.w13.weight", 1, 2), r"(.*)\.feed_forward\.w1\.(lora_A|lora_B)$": ( r"\1.feed_forward.w13.\2", 0, 2, ), r"(.*)\.feed_forward\.w3\.(lora_A|lora_B)$": ( r"\1.feed_forward.w13.\2", 1, 2, ), } ) def __post_init__(self): super().__post_init__() self.out_channels = self.out_channels or self.in_channels self.num_channels_latents = self.in_channels self.hidden_size = self.dim @dataclass class ZImageDitConfig(DiTConfig): arch_config: ZImageArchConfig = field(default_factory=ZImageArchConfig) prefix: str = "zimage" ================================================ FILE: python/sglang/multimodal_gen/configs/models/encoders/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from sglang.multimodal_gen.configs.models.encoders.base import ( BaseEncoderOutput, EncoderConfig, ImageEncoderConfig, TextEncoderConfig, ) from sglang.multimodal_gen.configs.models.encoders.clip import ( CLIPTextConfig, CLIPVisionConfig, ) from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config from sglang.multimodal_gen.configs.models.encoders.gemma_3 import Gemma3Config from sglang.multimodal_gen.configs.models.encoders.llama import LlamaConfig from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config __all__ = [ "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig", "LlamaConfig", "Qwen3TextConfig", "T5Config", "Gemma2Config", "Gemma3Config", ] ================================================ FILE: python/sglang/multimodal_gen/configs/models/encoders/base.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from typing import Any import torch from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum @dataclass class EncoderArchConfig(ArchConfig): _fsdp_shard_conditions: list = field(default_factory=lambda: []) architectures: list[str] = field(default_factory=lambda: []) _supported_attention_backends: set[AttentionBackendEnum] = field( default_factory=lambda: { AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_3, } ) output_hidden_states: bool = False use_return_dict: bool = True @dataclass class TextEncoderArchConfig(EncoderArchConfig): vocab_size: int = 0 hidden_size: int = 0 num_hidden_layers: int = 0 num_attention_heads: int = 0 pad_token_id: int = 0 eos_token_id: int = 0 text_len: int = 0 hidden_state_skip_layer: int = 0 decoder_start_token_id: int = 0 output_past: bool = True scalable_attention: bool = True tie_word_embeddings: bool = False stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=list ) # mapping from huggingface weight names to custom names tokenizer_kwargs: dict[str, Any] = field(default_factory=dict) _fsdp_shard_conditions: list = field(default_factory=lambda: []) def __post_init__(self) -> None: self.tokenizer_kwargs = { "truncation": True, "max_length": self.text_len, "return_tensors": "pt", } @dataclass class ImageEncoderArchConfig(EncoderArchConfig): pass @dataclass class BaseEncoderOutput: last_hidden_state: torch.FloatTensor | None = None pooler_output: torch.FloatTensor | None = None hidden_states: tuple[torch.FloatTensor, ...] | None = None attentions: tuple[torch.FloatTensor, ...] | None = None attention_mask: torch.Tensor | None = None @dataclass class EncoderConfig(ModelConfig): arch_config: ArchConfig = field(default_factory=EncoderArchConfig) prefix: str = "" quant_config: QuantizationConfig | None = None lora_config: Any | None = None @dataclass class TextEncoderConfig(EncoderConfig): arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig) # Use the SP Group of the transformer as the TP Group of T5. parallel_folding: bool = False # "sp" or "ulysses" or "ring" parallel_folding_mode: str = "sp" @dataclass class ImageEncoderConfig(EncoderConfig): arch_config: ArchConfig = field(default_factory=ImageEncoderArchConfig) ================================================ FILE: python/sglang/multimodal_gen/configs/models/encoders/clip.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.encoders.base import ( ImageEncoderArchConfig, ImageEncoderConfig, TextEncoderArchConfig, TextEncoderConfig, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum def _is_transformer_layer(n: str, m) -> bool: return "layers" in n and str.isdigit(n.split(".")[-1]) def _is_embeddings(n: str, m) -> bool: return n.endswith("embeddings") @dataclass class CLIPTextArchConfig(TextEncoderArchConfig): vocab_size: int = 49408 hidden_size: int = 512 intermediate_size: int = 2048 projection_dim: int = 512 num_hidden_layers: int = 12 num_attention_heads: int = 8 max_position_embeddings: int = 77 hidden_act: str = "quick_gelu" layer_norm_eps: float = 1e-5 dropout: float = 0.0 attention_dropout: float = 0.0 initializer_range: float = 0.02 initializer_factor: float = 1.0 pad_token_id: int = 1 bos_token_id: int = 49406 eos_token_id: int = 49407 text_len: int = 77 _supported_attention_backends: set[AttentionBackendEnum] = field( default_factory=lambda: { AttentionBackendEnum.TORCH_SDPA, # Force TORCH_SDPA to support attention_mask } ) stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=lambda: [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), ("qkv_proj", "k_proj", "k"), ("qkv_proj", "v_proj", "v"), ] ) _fsdp_shard_conditions: list = field( default_factory=lambda: [_is_transformer_layer, _is_embeddings] ) @dataclass class CLIPVisionArchConfig(ImageEncoderArchConfig): hidden_size: int = 768 intermediate_size: int = 3072 projection_dim: int = 512 num_hidden_layers: int = 12 num_attention_heads: int = 12 num_channels: int = 3 image_size: int = 224 patch_size: int = 32 hidden_act: str = "quick_gelu" layer_norm_eps: float = 1e-5 dropout: float = 0.0 attention_dropout: float = 0.0 initializer_range: float = 0.02 initializer_factor: float = 1.0 stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=lambda: [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), ("qkv_proj", "k_proj", "k"), ("qkv_proj", "v_proj", "v"), ] ) @dataclass class CLIPTextConfig(TextEncoderConfig): arch_config: TextEncoderArchConfig = field(default_factory=CLIPTextArchConfig) num_hidden_layers_override: int | None = None require_post_norm: bool | None = None prefix: str = "clip" @dataclass class CLIPVisionConfig(ImageEncoderConfig): arch_config: ImageEncoderArchConfig = field(default_factory=CLIPVisionArchConfig) num_hidden_layers_override: int | None = None require_post_norm: bool | None = None prefix: str = "clip" ================================================ FILE: python/sglang/multimodal_gen/configs/models/encoders/gemma2.py ================================================ # SPDX-License-Identifier: Apache-2.0 # # Text encoder configuration for Gemma2 2B, used by SANA for text conditioning. # # SANA uses the hidden states from Gemma2 (not logits) as the conditioning # signal for cross-attention in the DiT. The encoder output dimension (2304) # is projected to the DiT's inner_dim via caption_projection. # # Defaults match google/gemma-2-2b-it (the model used in SANA HF checkpoints). from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.encoders.base import ( TextEncoderArchConfig, TextEncoderConfig, ) def _is_transformer_layer(n: str, m) -> bool: return "layers" in n and str.isdigit(n.split(".")[-1]) def _is_embeddings(n: str, m) -> bool: return n.endswith("embed_tokens") def _is_final_norm(n: str, m) -> bool: return n.endswith("norm") @dataclass class Gemma2ArchConfig(TextEncoderArchConfig): vocab_size: int = 256000 hidden_size: int = 2304 intermediate_size: int = 9216 num_hidden_layers: int = 26 num_attention_heads: int = 8 num_key_value_heads: int = 4 head_dim: int = 256 hidden_act: str = "gelu_pytorch_tanh" hidden_activation: str = "gelu_pytorch_tanh" max_position_embeddings: int = 8192 rms_norm_eps: float = 1e-6 use_cache: bool = True pad_token_id: int = 0 eos_token_id: int = 1 bos_token_id: int = 2 tie_word_embeddings: bool = True rope_theta: float = 10000.0 attention_bias: bool = False attention_dropout: float = 0.0 # Gemma2 alternates between global and sliding-window attention # on odd/even layers, respectively. sliding_window: int = 4096 # query_pre_attn_scalar replaces the standard 1/sqrt(head_dim) scaling. query_pre_attn_scalar: int = 256 # Softcapping bounds raw attention logits via tanh(logits/cap)*cap. # NOTE: SDPA does not natively support softcapping; the runtime model # currently skips this (see Gemma2Attention.forward). Quality impact # is minimal for short text-encoder sequences but should be revisited # for longer context. attn_logit_softcapping: float = 50.0 final_logit_softcapping: float = 30.0 text_len: int = 300 stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=lambda: [ (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), (".gate_up_proj", ".gate_proj", "0"), (".gate_up_proj", ".up_proj", "1"), ] ) _fsdp_shard_conditions: list = field( default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm] ) @dataclass class Gemma2Config(TextEncoderConfig): arch_config: TextEncoderArchConfig = field(default_factory=Gemma2ArchConfig) prefix: str = "gemma_2" ================================================ FILE: python/sglang/multimodal_gen/configs/models/encoders/gemma_3.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.encoders.base import ( TextEncoderArchConfig, TextEncoderConfig, ) def _is_transformer_layer(n: str, m) -> bool: return "layers" in n and str.isdigit(n.split(".")[-1]) def _is_embeddings(n: str, m) -> bool: return n.endswith("embed_tokens") def _is_final_norm(n: str, m) -> bool: return n.endswith("norm") @dataclass class Gemma3ArchConfig(TextEncoderArchConfig): """Minimal Gemma text-encoder config for tokenizer kwargs. Note: runtime will load the actual `text_encoder/` module from the model repo (e.g. Gemma3Model) via transformers; this config mainly controls tokenization. """ vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 num_hidden_layers: int = 32 num_attention_heads: int = 32 num_key_value_heads: int | None = None hidden_act: str = "gelu_pytorch_tanh" max_position_embeddings: int = 2048 initializer_range: float = 0.02 rms_norm_eps: float = 1e-6 use_cache: bool = True pad_token_id: int = 0 bos_token_id: int = 1 eos_token_id: int = 2 pretraining_tp: int = 1 tie_word_embeddings: bool = True rope_theta: float = 10000.0 rope_scaling: dict | None = None rope_local_base_freq: float = 10000.0 sliding_window: int = 4096 layer_types: list[str] = field(default_factory=list) query_pre_attn_scalar: int | None = None attention_bias: bool = False attention_dropout: float = 0.0 mlp_bias: bool = False head_dim: int | None = None hidden_state_skip_layer: int = 2 text_len: int = 1024 stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=lambda: [ # (param_name, shard_name, shard_id) (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), (".gate_up_proj", ".gate_proj", "0"), # type: ignore (".gate_up_proj", ".up_proj", "1"), # type: ignore ] ) _fsdp_shard_conditions: list = field( default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm] ) @dataclass class Gemma3Config(TextEncoderConfig): arch_config: TextEncoderArchConfig = field(default_factory=Gemma3ArchConfig) prefix: str = "gemma_3" ================================================ FILE: python/sglang/multimodal_gen/configs/models/encoders/llama.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.encoders.base import ( TextEncoderArchConfig, TextEncoderConfig, ) def _is_transformer_layer(n: str, m) -> bool: return "layers" in n and str.isdigit(n.split(".")[-1]) def _is_embeddings(n: str, m) -> bool: return n.endswith("embed_tokens") def _is_final_norm(n: str, m) -> bool: return n.endswith("norm") @dataclass class LlamaArchConfig(TextEncoderArchConfig): vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 num_hidden_layers: int = 32 num_attention_heads: int = 32 num_key_value_heads: int | None = None hidden_act: str = "silu" max_position_embeddings: int = 2048 initializer_range: float = 0.02 rms_norm_eps: float = 1e-6 use_cache: bool = True pad_token_id: int = 0 bos_token_id: int = 1 eos_token_id: int = 2 pretraining_tp: int = 1 tie_word_embeddings: bool = False rope_theta: float = 10000.0 rope_scaling: float | None = None attention_bias: bool = False attention_dropout: float = 0.0 mlp_bias: bool = False head_dim: int | None = None hidden_state_skip_layer: int = 2 text_len: int = 256 stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=lambda: [ # (param_name, shard_name, shard_id) (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), (".gate_up_proj", ".gate_proj", 0), # type: ignore (".gate_up_proj", ".up_proj", 1), # type: ignore ] ) _fsdp_shard_conditions: list = field( default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm] ) @dataclass class LlamaConfig(TextEncoderConfig): arch_config: TextEncoderArchConfig = field(default_factory=LlamaArchConfig) prefix: str = "llama" ================================================ FILE: python/sglang/multimodal_gen/configs/models/encoders/qwen3.py ================================================ # SPDX-License-Identifier: Apache-2.0 """Qwen3 text encoder configuration for SGLang diffusion models.""" from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.encoders.base import ( TextEncoderArchConfig, TextEncoderConfig, ) def _is_transformer_layer(n: str, m) -> bool: return "layers" in n and str.isdigit(n.split(".")[-1]) def _is_embeddings(n: str, m) -> bool: return n.endswith("embed_tokens") def _is_final_norm(n: str, m) -> bool: return n.endswith("norm") @dataclass class Qwen3TextArchConfig(TextEncoderArchConfig): """Architecture config for Qwen3 text encoder. Qwen3 is similar to LLaMA but with QK-Norm (RMSNorm on Q and K before attention). """ vocab_size: int = 151936 hidden_size: int = 2560 intermediate_size: int = 9728 num_hidden_layers: int = 36 num_attention_heads: int = 32 num_key_value_heads: int = 8 hidden_act: str = "silu" max_position_embeddings: int = 40960 initializer_range: float = 0.02 rms_norm_eps: float = 1e-6 use_cache: bool = True pad_token_id: int = 151643 bos_token_id: int = 151643 eos_token_id: int = 151645 tie_word_embeddings: bool = True rope_theta: float = 1000000.0 rope_scaling: dict | None = None attention_bias: bool = False attention_dropout: float = 0.0 mlp_bias: bool = False head_dim: int = 128 text_len: int = 512 output_hidden_states: bool = True # Klein needs hidden states from layers 9, 18, 27 # Stacked params for weight loading with tensor parallelism stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=lambda: [ # (param_name, shard_name, shard_id) (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), (".gate_up_proj", ".gate_proj", 0), (".gate_up_proj", ".up_proj", 1), ] ) # FSDP sharding conditions for CPU offload _fsdp_shard_conditions: list = field( default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm] ) def __post_init__(self) -> None: self.tokenizer_kwargs = { "padding": "max_length", "truncation": True, "max_length": self.text_len, "return_tensors": "pt", } @dataclass class Qwen3TextConfig(TextEncoderConfig): """Top-level config for Qwen3 text encoder.""" arch_config: TextEncoderArchConfig = field(default_factory=Qwen3TextArchConfig) prefix: str = "qwen3" ================================================ FILE: python/sglang/multimodal_gen/configs/models/encoders/qwen_image.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.encoders.base import ( TextEncoderArchConfig, TextEncoderConfig, ) def _is_transformer_layer(n: str, m) -> bool: return "layers" in n and str.isdigit(n.split(".")[-1]) def _is_embeddings(n: str, m) -> bool: return n.endswith("embed_tokens") def _is_final_norm(n: str, m) -> bool: return n.endswith("norm") @dataclass class QwenImageArchConfig(TextEncoderArchConfig): vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 num_hidden_layers: int = 32 num_attention_heads: int = 32 num_key_value_heads: int | None = None hidden_act: str = "silu" max_position_embeddings: int = 2048 initializer_range: float = 0.02 rms_norm_eps: float = 1e-6 use_cache: bool = True pad_token_id: int = -1 eos_token_id: int = 2 pretraining_tp: int = 1 tie_word_embeddings: bool = False rope_theta: float = 10000.0 rope_scaling: float | None = None attention_bias: bool = False attention_dropout: float = 0.0 mlp_bias: bool = False head_dim: int | None = None hidden_state_skip_layer: int = 2 text_len: int = 512 stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=lambda: [ # (param_name, shard_name, shard_id) (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), (".gate_up_proj", ".gate_proj", 0), # type: ignore (".gate_up_proj", ".up_proj", 1), # type: ignore ] ) _fsdp_shard_conditions: list = field( default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm] ) @dataclass class Qwen2_5VLConfig(TextEncoderConfig): arch_config: TextEncoderArchConfig = field(default_factory=QwenImageArchConfig) # prefix: str = "qwen_image" ================================================ FILE: python/sglang/multimodal_gen/configs/models/encoders/t5.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import argparse from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.encoders.base import ( TextEncoderArchConfig, TextEncoderConfig, ) def _is_transformer_layer(n: str, m) -> bool: return "block" in n and str.isdigit(n.split(".")[-1]) def _is_embeddings(n: str, m) -> bool: return n.endswith("shared") def _is_final_layernorm(n: str, m) -> bool: return n.endswith("final_layer_norm") @dataclass class T5ArchConfig(TextEncoderArchConfig): vocab_size: int = 32128 d_model: int = 512 d_kv: int = 64 d_ff: int = 2048 num_layers: int = 6 num_decoder_layers: int | None = None num_heads: int = 8 relative_attention_num_buckets: int = 32 relative_attention_max_distance: int = 128 dropout_rate: float = 0.1 layer_norm_epsilon: float = 1e-6 initializer_factor: float = 1.0 feed_forward_proj: str = "relu" dense_act_fn: str = "" is_gated_act: bool = False is_encoder_decoder: bool = True use_cache: bool = True pad_token_id: int = 0 eos_token_id: int = 1 classifier_dropout: float = 0.0 text_len: int = 512 stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=lambda: [ # (param_name, shard_name, shard_id) (".qkv_proj", ".q", "q"), (".qkv_proj", ".k", "k"), (".qkv_proj", ".v", "v"), ] ) _fsdp_shard_conditions: list = field( default_factory=lambda: [ _is_transformer_layer, _is_embeddings, _is_final_layernorm, ] ) # Referenced from https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/configuration_t5.py def __post_init__(self): super().__post_init__() act_info = self.feed_forward_proj.split("-") self.dense_act_fn: str = act_info[-1] self.is_gated_act: bool = act_info[0] == "gated" if self.feed_forward_proj == "gated-gelu": self.dense_act_fn = "gelu_new" self.tokenizer_kwargs = { "padding": "max_length", "truncation": True, "max_length": self.text_len, "add_special_tokens": True, "return_attention_mask": True, "return_tensors": "pt", } @dataclass class T5Config(TextEncoderConfig): arch_config: TextEncoderArchConfig = field(default_factory=T5ArchConfig) prefix: str = "t5" # Use the SP Group of the transformer as the TP Group of T5. parallel_folding: bool = False # "sp" or "ulysses" or "ring" parallel_folding_mode: str = "sp" @staticmethod def add_cli_args( parser: argparse.ArgumentParser, prefix: str = "t5-config" ) -> argparse.ArgumentParser: return parser ================================================ FILE: python/sglang/multimodal_gen/configs/models/vaes/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from sglang.multimodal_gen.configs.models.vaes.dac import DacVAEConfig from sglang.multimodal_gen.configs.models.vaes.hunyuan3d import Hunyuan3DVAEConfig from sglang.multimodal_gen.configs.models.vaes.hunyuanvae import HunyuanVAEConfig from sglang.multimodal_gen.configs.models.vaes.wanvae import WanVAEConfig __all__ = [ "DacVAEConfig", "HunyuanVAEConfig", "WanVAEConfig", "Hunyuan3DVAEConfig", ] ================================================ FILE: python/sglang/multimodal_gen/configs/models/vaes/base.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import argparse import dataclasses from dataclasses import dataclass, field from typing import Any import torch from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig from sglang.multimodal_gen.utils import StoreBoolean @dataclass class VAEArchConfig(ArchConfig): scaling_factor: float | torch.Tensor = 0 temporal_compression_ratio: int = 4 # or vae_scale_factor? spatial_compression_ratio: int = 8 @dataclass class VAEConfig(ModelConfig): arch_config: VAEArchConfig = field(default_factory=VAEArchConfig) # sglang-diffusion VAE-specific parameters load_encoder: bool = True load_decoder: bool = True tile_sample_min_height: int = 256 tile_sample_min_width: int = 256 tile_sample_min_num_frames: int = 16 tile_sample_stride_height: int = 192 tile_sample_stride_width: int = 192 tile_sample_stride_num_frames: int = 12 blend_num_frames: int = 0 use_tiling: bool = True use_temporal_tiling: bool = True use_parallel_tiling: bool = True use_temporal_scaling_frames: bool = True def __post_init__(self): self.blend_num_frames = ( self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames ) def post_init(self): pass @staticmethod def add_cli_args(parser: Any, prefix: str = "vae-config") -> Any: """Add CLI arguments for VAEConfig fields""" parser.add_argument( f"--{prefix}.load-encoder", action=StoreBoolean, dest=f"{prefix.replace('-', '_')}.load_encoder", default=VAEConfig.load_encoder, help="Whether to load the VAE encoder", ) parser.add_argument( f"--{prefix}.load-decoder", action=StoreBoolean, dest=f"{prefix.replace('-', '_')}.load_decoder", default=VAEConfig.load_decoder, help="Whether to load the VAE decoder", ) parser.add_argument( f"--{prefix}.tile-sample-min-height", type=int, dest=f"{prefix.replace('-', '_')}.tile_sample_min_height", default=VAEConfig.tile_sample_min_height, help="Minimum height for VAE tile sampling", ) parser.add_argument( f"--{prefix}.tile-sample-min-width", type=int, dest=f"{prefix.replace('-', '_')}.tile_sample_min_width", default=VAEConfig.tile_sample_min_width, help="Minimum width for VAE tile sampling", ) parser.add_argument( f"--{prefix}.tile-sample-min-num-frames", type=int, dest=f"{prefix.replace('-', '_')}.tile_sample_min_num_frames", default=VAEConfig.tile_sample_min_num_frames, help="Minimum number of frames for VAE tile sampling", ) parser.add_argument( f"--{prefix}.tile-sample-stride-height", type=int, dest=f"{prefix.replace('-', '_')}.tile_sample_stride_height", default=VAEConfig.tile_sample_stride_height, help="Stride height for VAE tile sampling", ) parser.add_argument( f"--{prefix}.tile-sample-stride-width", type=int, dest=f"{prefix.replace('-', '_')}.tile_sample_stride_width", default=VAEConfig.tile_sample_stride_width, help="Stride width for VAE tile sampling", ) parser.add_argument( f"--{prefix}.tile-sample-stride-num-frames", type=int, dest=f"{prefix.replace('-', '_')}.tile_sample_stride_num_frames", default=VAEConfig.tile_sample_stride_num_frames, help="Stride number of frames for VAE tile sampling", ) parser.add_argument( f"--{prefix}.blend-num-frames", type=int, dest=f"{prefix.replace('-', '_')}.blend_num_frames", default=VAEConfig.blend_num_frames, help="Number of frames to blend for VAE tile sampling", ) parser.add_argument( f"--{prefix}.use-tiling", action=StoreBoolean, dest=f"{prefix.replace('-', '_')}.use_tiling", default=VAEConfig.use_tiling, help="Whether to use tiling for VAE", ) parser.add_argument( f"--{prefix}.use-temporal-tiling", action=StoreBoolean, dest=f"{prefix.replace('-', '_')}.use_temporal_tiling", default=VAEConfig.use_temporal_tiling, help="Whether to use temporal tiling for VAE", ) parser.add_argument( f"--{prefix}.use-parallel-tiling", action=StoreBoolean, dest=f"{prefix.replace('-', '_')}.use_parallel_tiling", default=VAEConfig.use_parallel_tiling, help="Whether to use parallel tiling for VAE", ) return parser def get_vae_scale_factor(self): return 2 ** (len(self.arch_config.block_out_channels) - 1) def encode_sample_mode(self): return "argmax" @classmethod def from_cli_args(cls, args: argparse.Namespace) -> "VAEConfig": kwargs = {} for attr in dataclasses.fields(cls): value = getattr(args, attr.name, None) if value is not None: kwargs[attr.name] = value return cls(**kwargs) ================================================ FILE: python/sglang/multimodal_gen/configs/models/vaes/dac.py ================================================ # Copied and adapted from: mossVG/mova/diffusion/models/dac_vae.py # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from typing import List from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig @dataclass class DacVAEArchConfig(ArchConfig): codebook_dim: int = 8 codebook_size: int = 1024 continuous: bool = True decoder_dim: int = 2048 decoder_rates: List[int] = field(default_factory=lambda: [8, 5, 4, 3, 2]) encoder_dim: int = 128 encoder_rates: List[int] = field(default_factory=lambda: [2, 3, 4, 5, 8]) hop_length: int = 3840 latent_dim: int = 128 n_codebooks: int = 9 quantizer_dropout: bool = False sample_rate: int = 48000 @dataclass class DacVAEConfig(ModelConfig): arch_config: DacVAEArchConfig = field(default_factory=DacVAEArchConfig) load_encoder: bool = True load_decoder: bool = True ================================================ FILE: python/sglang/multimodal_gen/configs/models/vaes/flux.py ================================================ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig @dataclass class FluxVAEArchConfig(VAEArchConfig): spatial_compression_ratio: int = 1 base_dim: int = 96 decoder_base_dim: int | None = None z_dim: int = 16 dim_mult: tuple[int, ...] = (1, 2, 4, 4) num_res_blocks: int = 2 attn_scales: tuple[float, ...] = () temperal_downsample: tuple[bool, ...] = (False, True, True) dropout: float = 0.0 is_residual: bool = False in_channels: int = 3 out_channels: int = 3 patch_size: int | None = None scale_factor_temporal: int = 4 scale_factor_spatial: int = 8 clip_output: bool = True @dataclass class Flux2VAEArchConfig(FluxVAEArchConfig): pass @dataclass class FluxVAEConfig(VAEConfig): arch_config: FluxVAEArchConfig = field(default_factory=FluxVAEArchConfig) use_feature_cache: bool = True use_tiling: bool = False use_temporal_tiling: bool = False use_parallel_tiling: bool = False def __post_init__(self): self.blend_num_frames = ( self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames ) * 2 def post_init(self): # Calculate vae_scale_factor: prefer block_out_channels, fallback to dim_mult or scale_factor_spatial if ( hasattr(self.arch_config, "block_out_channels") and self.arch_config.block_out_channels ): self.arch_config.vae_scale_factor = 2 ** ( len(self.arch_config.block_out_channels) - 1 ) elif self.arch_config.dim_mult: self.arch_config.vae_scale_factor = 2 ** ( len(self.arch_config.dim_mult) - 1 ) else: self.arch_config.vae_scale_factor = self.arch_config.scale_factor_spatial self.arch_config.spatial_compression_ratio = self.arch_config.vae_scale_factor @dataclass class Flux2VAEConfig(FluxVAEConfig): arch_config: Flux2VAEArchConfig = field(default_factory=Flux2VAEArchConfig) ================================================ FILE: python/sglang/multimodal_gen/configs/models/vaes/glmimage.py ================================================ from dataclasses import dataclass, field import torch from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig @dataclass class GlmImageVAEArchConfig(VAEArchConfig): spatial_compression_ratio: int = 1 base_dim: int = 96 decoder_base_dim: int | None = None z_dim: int = 16 dim_mult: tuple[int, ...] = (1, 2, 4, 4) num_res_blocks: int = 2 attn_scales: tuple[float, ...] = () temperal_downsample: tuple[bool, ...] = (False, True, True) dropout: float = 0.0 is_residual: bool = False input_channels: int = 3 out_channels: int = 3 patch_size: int | None = None scale_factor_temporal: int = 4 scale_factor_spatial: int = 8 clip_output: bool = True scaling_factor: float | torch.Tensor = 0 latents_mean: tuple[float, ...] | None = None latents_std: tuple[float, ...] | None = None shift_factor: float | None = None latent_channels: int = 16 in_channels: int = 16 @dataclass class GlmImageVAEConfig(VAEConfig): arch_config: GlmImageVAEArchConfig = field(default_factory=GlmImageVAEArchConfig) use_feature_cache: bool = True use_tiling: bool = False use_temporal_tiling: bool = False use_parallel_tiling: bool = False def get_vae_scale_factor(self): return 2 ** len(self.arch_config.temperal_downsample) def __post_init__(self): self.blend_num_frames = ( self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames ) * 2 def post_init(self): self.arch_config.vae_scale_factor = 2 ** ( len(self.arch_config.temperal_downsample) ) self.arch_config.spatial_compression_ratio = self.arch_config.vae_scale_factor ================================================ FILE: python/sglang/multimodal_gen/configs/models/vaes/hunyuan3d.py ================================================ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig @dataclass class Hunyuan3DVAEArchConfig(VAEArchConfig): """Architecture config for Hunyuan3D VAE.""" latent_shape: tuple[int, ...] = (1024, 64) scale_factor: float = 1.0 @dataclass class Hunyuan3DVAEConfig(VAEConfig): """VAE configuration for Hunyuan3D.""" arch_config: Hunyuan3DVAEArchConfig = field(default_factory=Hunyuan3DVAEArchConfig) subfolder: str = "hunyuan3d-dit-v2-0" load_encoder: bool = False load_decoder: bool = True ================================================ FILE: python/sglang/multimodal_gen/configs/models/vaes/hunyuanvae.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig @dataclass class HunyuanVAEArchConfig(VAEArchConfig): in_channels: int = 3 out_channels: int = 3 latent_channels: int = 16 down_block_types: tuple[str, ...] = ( "HunyuanVideoDownBlock3D", "HunyuanVideoDownBlock3D", "HunyuanVideoDownBlock3D", "HunyuanVideoDownBlock3D", ) up_block_types: tuple[str, ...] = ( "HunyuanVideoUpBlock3D", "HunyuanVideoUpBlock3D", "HunyuanVideoUpBlock3D", "HunyuanVideoUpBlock3D", ) block_out_channels: tuple[int, ...] = (128, 256, 512, 512) layers_per_block: int = 2 act_fn: str = "silu" norm_num_groups: int = 32 scaling_factor: float = 0.476986 spatial_compression_ratio: int = 8 temporal_compression_ratio: int = 4 mid_block_add_attention: bool = True def __post_init__(self): self.spatial_compression_ratio: int = 2 ** (len(self.block_out_channels) - 1) @dataclass class HunyuanVAEConfig(VAEConfig): arch_config: VAEArchConfig = field(default_factory=HunyuanVAEArchConfig) ================================================ FILE: python/sglang/multimodal_gen/configs/models/vaes/ltx_audio.py ================================================ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from typing import Optional, Tuple from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig @dataclass class LTXAudioVAEArchConfig(VAEArchConfig): # Architecture params causality_axis: str = "height" attn_resolutions: Optional[Tuple[int, ...]] = None base_channels: int = 128 latent_channels: int = 8 output_channels: int = 2 ch_mult: Tuple[int, ...] = (1, 2, 4) num_res_blocks: int = 2 norm_type: str = "pixel" dropout: float = 0.0 mid_block_add_attention: bool = False sample_rate: int = 16000 mel_hop_length: int = 160 is_causal: bool = True mel_bins: Optional[int] = 64 double_z: bool = True @dataclass class LTXAudioVAEConfig(VAEConfig): arch_config: LTXAudioVAEArchConfig = field(default_factory=LTXAudioVAEArchConfig) ================================================ FILE: python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py ================================================ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from typing import List from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig @dataclass class LTXVideoVAEArchConfig(VAEArchConfig): # Architecture params in_channels: int = 3 latent_channels: int = 128 out_channels: int = 3 block_out_channels: List[int] = field( default_factory=lambda: [256, 512, 1024, 2048] ) down_block_types: List[str] = field( default_factory=lambda: [ "LTX2VideoDownBlock3D", "LTX2VideoDownBlock3D", "LTX2VideoDownBlock3D", "LTX2VideoDownBlock3D", ] ) spatio_temporal_scaling: List[bool] = field( default_factory=lambda: [True, True, True, True] ) layers_per_block: List[int] = field(default_factory=lambda: [4, 6, 6, 2, 2]) downsample_type: List[str] = field( default_factory=lambda: [ "spatial", "temporal", "spatiotemporal", "spatiotemporal", ] ) patch_size: int = 4 patch_size_t: int = 1 resnet_norm_eps: float = 1e-6 encoder_causal: bool = True encoder_spatial_padding_mode: str = "zeros" decoder_block_out_channels: List[int] = field( default_factory=lambda: [256, 512, 1024] ) decoder_spatio_temporal_scaling: List[bool] = field( default_factory=lambda: [True, True, True] ) decoder_layers_per_block: List[int] = field(default_factory=lambda: [5, 5, 5, 5]) decoder_causal: bool = False decoder_spatial_padding_mode: str = "reflect" @dataclass class LTXVideoVAEConfig(VAEConfig): arch_config: LTXVideoVAEArchConfig = field(default_factory=LTXVideoVAEArchConfig) ================================================ FILE: python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig @dataclass class QwenImageVAEArchConfig(VAEArchConfig): spatial_compression_ratio: int = 1 base_dim: int = 96 decoder_base_dim: int | None = None z_dim: int = 16 dim_mult: tuple[int, ...] = (1, 2, 4, 4) num_res_blocks: int = 2 attn_scales: tuple[float, ...] = () temperal_downsample: tuple[bool, ...] = (False, True, True) dropout: float = 0.0 is_residual: bool = False input_channels: int = 3 out_channels: int = 3 patch_size: int | None = None scale_factor_temporal: int = 4 scale_factor_spatial: int = 8 clip_output: bool = True @dataclass class QwenImageVAEConfig(VAEConfig): arch_config: QwenImageVAEArchConfig = field(default_factory=QwenImageVAEArchConfig) use_feature_cache: bool = True use_tiling: bool = False use_temporal_tiling: bool = False use_parallel_tiling: bool = False def get_vae_scale_factor(self): return 2 ** len(self.arch_config.temperal_downsample) def __post_init__(self): self.blend_num_frames = ( self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames ) * 2 def post_init(self): self.arch_config.vae_scale_factor = 2 ** ( len(self.arch_config.temperal_downsample) ) self.arch_config.spatial_compression_ratio = self.arch_config.vae_scale_factor ================================================ FILE: python/sglang/multimodal_gen/configs/models/vaes/sana.py ================================================ # SPDX-License-Identifier: Apache-2.0 # # VAE configuration for SANA's DC-AE (Deep Compression AutoEncoder). # # DC-AE achieves a 32x spatial compression ratio (vs. 8x for standard SD VAEs), # which means a 1024x1024 image becomes 32x32 latents with 32 channels. # This aggressive compression is what allows SANA to run efficiently at # high resolutions despite having a relatively small DiT. # # Reference: https://arxiv.org/abs/2405.17811 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig @dataclass class SanaVAEArchConfig(VAEArchConfig): spatial_compression_ratio: int = 32 # DC-AE uses a different scaling factor than standard VAEs; # this value must match the pretrained checkpoint. scaling_factor: float = 0.41407 latent_channels: int = 32 in_channels: int = 3 @dataclass class SanaVAEConfig(VAEConfig): arch_config: SanaVAEArchConfig = field(default_factory=SanaVAEArchConfig) # DC-AE does not currently support tiling in our wrapper. # Enable these once the diffusers AutoencoderDC adds tiling support. use_tiling: bool = False use_temporal_tiling: bool = False use_parallel_tiling: bool = False def post_init(self): # Called by VAELoader AFTER update_model_arch() merges the HF config.json # values into arch_config. Must be post_init() (not __post_init__) because # __post_init__ fires at dataclass creation time, before the HF config merge. # # The base VAEConfig.get_vae_scale_factor() derives from block_out_channels, # which DC-AE doesn't have. Set vae_scale_factor directly from the # spatial_compression_ratio (32x for DC-AE). self.arch_config.vae_scale_factor = self.arch_config.spatial_compression_ratio ================================================ FILE: python/sglang/multimodal_gen/configs/models/vaes/wanvae.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field import torch from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig @dataclass class WanVAEArchConfig(VAEArchConfig): base_dim: int = 96 decoder_base_dim: int | None = None z_dim: int = 16 dim_mult: tuple[int, ...] = (1, 2, 4, 4) num_res_blocks: int = 2 attn_scales: tuple[float, ...] = () temperal_downsample: tuple[bool, ...] = (False, True, True) dropout: float = 0.0 latents_mean: tuple[float, ...] = ( -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921, ) latents_std: tuple[float, ...] = ( 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160, ) is_residual: bool = False in_channels: int = 3 out_channels: int = 3 patch_size: int | None = None scale_factor_temporal: int = 4 scale_factor_spatial: int = 8 clip_output: bool = True def __post_init__(self): self.scaling_factor: torch.tensor = 1.0 / torch.tensor(self.latents_std).view( 1, self.z_dim, 1, 1, 1 ) self.shift_factor: torch.tensor = torch.tensor(self.latents_mean).view( 1, self.z_dim, 1, 1, 1 ) self.temporal_compression_ratio = self.scale_factor_temporal self.spatial_compression_ratio = self.scale_factor_spatial @dataclass class WanVAEConfig(VAEConfig): arch_config: WanVAEArchConfig = field(default_factory=WanVAEArchConfig) use_feature_cache: bool = True use_tiling: bool = False use_temporal_tiling: bool = False use_parallel_tiling: bool = False use_parallel_encode: bool = True use_parallel_decode: bool = True def __post_init__(self): self.blend_num_frames = ( self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames ) * 2 ================================================ FILE: python/sglang/multimodal_gen/configs/models/vocoder/__init__.py ================================================ from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import LTXVocoderConfig __all__ = ["LTXVocoderConfig"] ================================================ FILE: python/sglang/multimodal_gen/configs/models/vocoder/base.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import argparse import dataclasses from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig @dataclass class VocoderArchConfig(ArchConfig): in_channels: int = 128 hidden_channels: int = 1024 out_channels: int = 2 @dataclass class VocoderConfig(ModelConfig): arch_config: VocoderArchConfig = field(default_factory=VocoderArchConfig) @classmethod def from_cli_args(cls, args: argparse.Namespace) -> "VocoderConfig": kwargs = {} for attr in dataclasses.fields(cls): value = getattr(args, attr.name, None) if value is not None: kwargs[attr.name] = value return cls(**kwargs) ================================================ FILE: python/sglang/multimodal_gen/configs/models/vocoder/ltx_vocoder.py ================================================ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from typing import List from sglang.multimodal_gen.configs.models.vocoder.base import ( VocoderArchConfig, VocoderConfig, ) @dataclass class LTXVocoderArchConfig(VocoderArchConfig): # Architecture params in_channels: int = 128 hidden_channels: int = 1024 out_channels: int = 2 upsample_kernel_sizes: List[int] = field(default_factory=lambda: [3, 7, 11]) upsample_factors: List[int] = field(default_factory=lambda: [6, 5, 2, 2, 2]) resnet_kernel_sizes: List[int] = field(default_factory=lambda: [3, 7, 11]) resnet_dilations: List[List[int]] = field( default_factory=lambda: [[1, 3, 5], [1, 3, 5], [1, 3, 5]] ) leaky_relu_negative_slope: float = 0.1 sample_rate: int = 24000 @dataclass class LTXVocoderConfig(VocoderConfig): arch_config: LTXVocoderArchConfig = field(default_factory=LTXVocoderArchConfig) ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from sglang.multimodal_gen.configs.pipeline_configs.base import ( PipelineConfig, SlidingTileAttnConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.diffusers_generic import ( DiffusersGenericPipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.flux import ( Flux2KleinPipelineConfig, Flux2PipelineConfig, FluxPipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.flux_finetuned import ( Flux2FinetunedPipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.helios import ( HeliosDistilledConfig, HeliosMidConfig, HeliosT2VConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( FastHunyuanConfig, HunyuanConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( Hunyuan3D2PipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.wan import ( SelfForcingWanT2V480PConfig, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig __all__ = [ "DiffusersGenericPipelineConfig", "HeliosDistilledConfig", "HeliosMidConfig", "HeliosT2VConfig", "HunyuanConfig", "FastHunyuanConfig", "Hunyuan3D2PipelineConfig", "FluxPipelineConfig", "Flux2PipelineConfig", "Flux2KleinPipelineConfig", "Flux2FinetunedPipelineConfig", "PipelineConfig", "SanaPipelineConfig", "SlidingTileAttnConfig", "MOVAPipelineConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig", "SelfForcingWanT2V480PConfig", "ZImagePipelineConfig", "LTX2PipelineConfig", ] ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/base.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import json import os from collections.abc import Callable from dataclasses import asdict, dataclass, field, fields from enum import Enum, auto from typing import Any import numpy as np import PIL import torch from einops import rearrange from sglang.multimodal_gen.configs.models import ( DiTConfig, EncoderConfig, ModelConfig, VAEConfig, ) from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config from sglang.multimodal_gen.configs.sample.sampling_params import DataType from sglang.multimodal_gen.configs.utils import update_config_from_args from sglang.multimodal_gen.runtime.distributed.communication_op import ( sequence_model_parallel_all_gather, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_sp_parallel_rank, get_sp_world_size, ) from sglang.multimodal_gen.runtime.models.vision_utils import get_default_height_width from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import ( FlexibleArgumentParser, StoreBoolean, shallow_asdict, ) logger = init_logger(__name__) # NOTE: possible duplication with DataType # this may focus on the model's original ability class ModelTaskType(Enum): # TODO: check if I2V/TI2V models can work w/wo text I2V = auto() # Image to Video T2V = auto() # Text to Video TI2V = auto() # Text and Image to Video T2I = auto() # Text to Image I2I = auto() # Image to Image TI2I = auto() # Image to Image or Text-Image to Image I2M = auto() # Image to Mesh def is_image_gen(self) -> bool: return ( self == ModelTaskType.T2I or self == ModelTaskType.I2I or self == ModelTaskType.TI2I ) def requires_image_input(self) -> bool: return ( self == ModelTaskType.I2V or self == ModelTaskType.I2I or self == ModelTaskType.I2M ) def accepts_image_input(self) -> bool: return ( self == ModelTaskType.I2V or self == ModelTaskType.I2I or self == ModelTaskType.TI2I or self == ModelTaskType.TI2V or self == ModelTaskType.I2M ) def data_type(self) -> DataType: if self == ModelTaskType.I2M: return DataType.MESH if self.is_image_gen(): return DataType.IMAGE else: return DataType.VIDEO class STA_Mode(str, Enum): """STA (Sliding Tile Attention) modes.""" STA_INFERENCE = "STA_inference" STA_SEARCHING = "STA_searching" STA_TUNING = "STA_tuning" STA_TUNING_CFG = "STA_tuning_cfg" NONE = None def preprocess_text(prompt: str) -> str: return prompt def postprocess_text(output: BaseEncoderOutput, _text_inputs) -> torch.tensor: raise NotImplementedError def shard_rotary_emb_for_sp(emb): """ Shard rotary embeddings [S, D] along sequence for SP. If S is not divisible by SP degree, pad by repeating the last row. """ # Sequence Parallelism: slice image RoPE to local shard if enabled try: from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_sp_parallel_rank, get_sp_world_size, ) sp_world_size = get_sp_world_size() except Exception: sp_world_size = 1 seq_len = emb.shape[0] if seq_len % sp_world_size != 0: pad_len = sp_world_size - (seq_len % sp_world_size) pad = emb[-1:].repeat(pad_len, 1) emb = torch.cat([emb, pad], dim=0) if sp_world_size > 1: try: rank = get_sp_parallel_rank() except Exception: rank = 0 seq_len = emb.shape[0] local_len = seq_len // sp_world_size start = rank * local_len end = start + local_len emb = emb[start:end] return emb else: return emb def maybe_unpad_latents(latents, batch): # If SP padding was applied, remove extra tokens before reshaping raw_shape = batch.raw_latent_shape if len(raw_shape) == 3: # Sequence format [B, S, D]: use seq_len directly target_tokens = raw_shape[1] else: # Spatial format [B, C, H, W] or [B, C, T, H, W]: use width * height width, height = raw_shape[-1], raw_shape[-2] target_tokens = width * height if latents.shape[1] > target_tokens: latents = latents[:, :target_tokens, :] return latents # config for a single pipeline @dataclass class PipelineConfig: """The base configuration class for a generation pipeline.""" task_type: ModelTaskType = ModelTaskType.I2I skip_input_image_preprocess: bool = False model_path: str = "" pipeline_config_path: str | None = None # precision and autocast enable_autocast: bool = True # generation parameters # controls the timestep embedding generation should_use_guidance: bool = True embedded_cfg_scale: float = 6.0 flow_shift: float | None = None disable_autocast: bool = False # Model configuration dit_config: DiTConfig = field(default_factory=DiTConfig) dit_precision: str = "bf16" # VAE configuration vae_config: VAEConfig = field(default_factory=VAEConfig) vae_precision: str = "fp32" vae_tiling: bool = True vae_slicing: bool = False vae_sp: bool = True # Image encoder configuration image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig) image_encoder_precision: str = "fp32" # Text encoder configuration DEFAULT_TEXT_ENCODER_PRECISIONS = ("fp32",) text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (EncoderConfig(),) ) # See PRECISION_TO_TYPE for detailed mapping text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32",)) text_encoder_extra_args: list[dict] = field(default_factory=lambda: [{}]) # image encoding image_encoder_extra_args: dict = field(default_factory=lambda: {}) def postprocess_image(self, image): return image.last_hidden_state preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( default_factory=lambda: (preprocess_text,) ) # get prompt_embeds from encoder output postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor], ...] = ( field(default_factory=lambda: (postprocess_text,)) ) # STA (Sliding Tile Attention) parameters mask_strategy_file_path: str | None = None STA_mode: STA_Mode = STA_Mode.STA_INFERENCE skip_time_steps: int = 15 # DMD parameters dmd_denoising_steps: list[int] | None = field(default=None) # Wan2.2 TI2V parameters boundary_ratio: float | None = None # Compilation # enable_torch_compile: bool = False # calculate the adjust size for condition image # width: original condition image width # height: original condition image height def calculate_condition_image_size(self, image, width, height) -> tuple[int, int]: vae_scale_factor = self.vae_config.arch_config.spatial_compression_ratio height, width = get_default_height_width(image, vae_scale_factor, height, width) return width, height ## For timestep preparation stage def prepare_sigmas(self, sigmas, num_inference_steps): return sigmas ## For ImageVAEEncodingStage def preprocess_condition_image( self, image, target_width, target_height, _vae_image_processor ): """ preprocess the condition image, returns (image, final_image_width, final_image_height) """ return image.resize( (target_width, target_height), PIL.Image.Resampling.LANCZOS ), (target_width, target_height) def prepare_calculated_size(self, image): return self.calculate_condition_image_size(image, image.width, image.height) def prepare_image_processor_kwargs(self, batch, neg=False): return {} def postprocess_image_latent(self, latent_condition, batch): vae_arch_config = self.vae_config.arch_config spatial_compression_ratio = vae_arch_config.spatial_compression_ratio temporal_compression_ratio = vae_arch_config.temporal_compression_ratio num_frames = batch.num_frames latent_height = batch.height // spatial_compression_ratio latent_width = batch.width // spatial_compression_ratio mask_lat_size = torch.ones(1, 1, num_frames, latent_height, latent_width) mask_lat_size[:, :, 1:] = 0 first_frame_mask = mask_lat_size[:, :, 0:1] first_frame_mask = torch.repeat_interleave( first_frame_mask, repeats=temporal_compression_ratio, dim=2, ) mask_lat_size = torch.concat( [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2 ) mask_lat_size = mask_lat_size.view( 1, -1, temporal_compression_ratio, latent_height, latent_width, ) mask_lat_size = mask_lat_size.transpose(1, 2) mask_lat_size = mask_lat_size.to(latent_condition.device) image_latents = torch.concat([mask_lat_size, latent_condition], dim=1) return image_latents def slice_noise_pred(self, noise, latents): return noise def adjust_num_frames(self, num_frames): return num_frames # tokenize the prompt def tokenize_prompt(self, prompt: list[str], tokenizer, tok_kwargs) -> dict: return tokenizer(prompt, **tok_kwargs) def prepare_latent_shape(self, batch, batch_size, num_frames): height = batch.height // self.vae_config.arch_config.spatial_compression_ratio width = batch.width // self.vae_config.arch_config.spatial_compression_ratio # Calculate latent shape shape = ( batch_size, self.dit_config.num_channels_latents, num_frames, height, width, ) return shape def allow_set_num_frames(self): return False def get_decode_scale_and_shift(self, device, dtype, vae): vae_arch_config = self.vae_config.arch_config scaling_factor = getattr(vae_arch_config, "scaling_factor", None) if scaling_factor is None: scaling_factor = getattr(vae, "scaling_factor", None) shift_factor = getattr(vae_arch_config, "shift_factor", None) if shift_factor is None: shift_factor = getattr(vae, "shift_factor", None) return scaling_factor, shift_factor # called after latents are prepared def maybe_pack_latents(self, latents, batch_size, batch): return latents def maybe_prepare_latent_ids(self, latents): return None # called after vae encode def postprocess_vae_encode(self, image_latents, vae): return image_latents # called after scale_and_shift, before vae decoding def preprocess_decoding(self, latents, server_args=None, vae=None): return latents def gather_latents_for_sp(self, latents): # For video latents [B, C, T_local, H, W], gather along time dim=2 latents = sequence_model_parallel_all_gather(latents, dim=2) return latents def preprocess_vae_image(self, batch, vae_image_processor): pass def shard_latents_for_sp(self, batch, latents): # general logic for video models sp_world_size, rank_in_sp_group = get_sp_world_size(), get_sp_parallel_rank() if batch.enable_sequence_shard and sp_world_size > 1: return latents, False if latents.dim() != 5: return latents, False time_dim = latents.shape[2] # Pad to next multiple of SP degree if needed if time_dim > 0 and time_dim % sp_world_size != 0: logger.debug( "Padding latents to next multiple of SP degree, performance is sub-optimal" ) pad_len = sp_world_size - (time_dim % sp_world_size) pad = torch.zeros( (*latents.shape[:2], pad_len, *latents.shape[3:]), dtype=latents.dtype, device=latents.device, ) latents = torch.cat([latents, pad], dim=2) assert latents.shape[2] % sp_world_size == 0 sharded_tensor = rearrange( latents, "b c (n t) h w -> b c n t h w", n=sp_world_size ).contiguous() sharded_tensor = sharded_tensor[:, :, rank_in_sp_group, :, :, :] return sharded_tensor, True def get_pos_prompt_embeds(self, batch): return batch.prompt_embeds def get_neg_prompt_embeds(self, batch): return batch.negative_prompt_embeds def post_denoising_loop(self, latents, batch): latents = maybe_unpad_latents(latents, batch) return latents def post_decoding(self, frames, server_args): return frames def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): return {} def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): return {} @staticmethod def add_cli_args( parser: FlexibleArgumentParser, prefix: str = "" ) -> FlexibleArgumentParser: prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else "" # model_path will be conflicting with the model_path in ServerArgs, # so we add it separately if prefix is not empty if prefix_with_dot != "": parser.add_argument( f"--{prefix_with_dot}model-path", type=str, dest=f"{prefix_with_dot.replace('-', '_')}model_path", default=PipelineConfig.model_path, help="Path to the pretrained model", ) parser.add_argument( f"--{prefix_with_dot}pipeline-config-path", type=str, dest=f"{prefix_with_dot.replace('-', '_')}pipeline_config_path", default=PipelineConfig.pipeline_config_path, help="Path to the pipeline config", ) parser.add_argument( f"--{prefix_with_dot}embedded-cfg-scale", type=float, dest=f"{prefix_with_dot.replace('-', '_')}embedded_cfg_scale", default=PipelineConfig.embedded_cfg_scale, help="Embedded CFG scale", ) parser.add_argument( f"--{prefix_with_dot}flow-shift", type=float, dest=f"{prefix_with_dot.replace('-', '_')}flow_shift", default=PipelineConfig.flow_shift, help="Flow shift parameter", ) parser.add_argument( f"--{prefix_with_dot}resolution", type=int, dest=f"{prefix_with_dot.replace('-', '_')}resolution", default=None, help="Override the selected pipeline config's resolution setting. Only applies to pipelines that define a resolution field.", ) # DiT configuration parser.add_argument( f"--{prefix_with_dot}dit-precision", type=str, dest=f"{prefix_with_dot.replace('-', '_')}dit_precision", default=PipelineConfig.dit_precision, choices=["fp32", "fp16", "bf16"], help="Precision for the DiT model", ) # VAE configuration parser.add_argument( f"--{prefix_with_dot}vae-precision", type=str, dest=f"{prefix_with_dot.replace('-', '_')}vae_precision", default=PipelineConfig.vae_precision, choices=["fp32", "fp16", "bf16"], help="Precision for VAE", ) parser.add_argument( f"--{prefix_with_dot}vae-tiling", action=StoreBoolean, dest=f"{prefix_with_dot.replace('-', '_')}vae_tiling", default=PipelineConfig.vae_tiling, help="Enable VAE tiling", ) parser.add_argument( f"--{prefix_with_dot}vae-slicing", action=StoreBoolean, dest=f"{prefix_with_dot.replace('-', '_')}vae_slicing", default=PipelineConfig.vae_slicing, help="Enable VAE slicing", ) parser.add_argument( f"--{prefix_with_dot}vae-sp", action=StoreBoolean, dest=f"{prefix_with_dot.replace('-', '_')}vae_sp", help="Enable VAE spatial parallelism", ) # Text encoder configuration parser.add_argument( f"--{prefix_with_dot}text-encoder-precisions", nargs="+", type=str, dest=f"{prefix_with_dot.replace('-', '_')}text_encoder_precisions", default=PipelineConfig.DEFAULT_TEXT_ENCODER_PRECISIONS, choices=["fp32", "fp16", "bf16"], help="Precision for each text encoder", ) # Image encoder configuration parser.add_argument( f"--{prefix_with_dot}image-encoder-precision", type=str, dest=f"{prefix_with_dot.replace('-', '_')}image_encoder_precision", default=PipelineConfig.image_encoder_precision, choices=["fp32", "fp16", "bf16"], help="Precision for image encoder", ) # DMD parameters parser.add_argument( f"--{prefix_with_dot}dmd-denoising-steps", type=parse_int_list, default=PipelineConfig.dmd_denoising_steps, help="Comma-separated list of denoising steps (e.g., '1000,757,522')", ) # Add VAE configuration arguments from sglang.multimodal_gen.configs.models.vaes.base import VAEConfig VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config") # Add DiT configuration arguments from sglang.multimodal_gen.configs.models.dits.base import DiTConfig DiTConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}dit-config") # Add T5 configuration arguments from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config T5Config.add_cli_args(parser, prefix=f"{prefix_with_dot}t5-config") return parser def update_config_from_dict(self, args: dict[str, Any], prefix: str = "") -> None: prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else "" update_config_from_args(self, args, prefix, pop_args=True) update_config_from_args( self.vae_config, args, f"{prefix_with_dot}vae_config", pop_args=True ) update_config_from_args( self.dit_config, args, f"{prefix_with_dot}dit_config", pop_args=True ) for text_encoder_config in self.text_encoder_configs: if isinstance(text_encoder_config, T5Config): update_config_from_args( text_encoder_config, args, f"{prefix_with_dot}t5_config", pop_args=True, ) @classmethod def from_kwargs( cls, kwargs: dict[str, Any], config_cli_prefix: str = "" ) -> "PipelineConfig": """ Load PipelineConfig from kwargs Dictionary, as part of the ServerArg initialization process kwargs: dictionary of kwargs config_cli_prefix: prefix of CLI arguments for this PipelineConfig instance """ from sglang.multimodal_gen.registry import get_model_info prefix_with_dot = ( f"{config_cli_prefix}." if (config_cli_prefix.strip() != "") else "" ) model_path: str | None = kwargs.get( prefix_with_dot + "model_path", None ) or kwargs.get("model_path") pipeline_config_or_path: str | PipelineConfig | dict[str, Any] | None = ( kwargs.get(prefix_with_dot + "pipeline_config", None) or kwargs.get("pipeline_config") ) if model_path is None: raise ValueError("model_path is required in kwargs") # Check if model_path is a safetensors file and pipeline_class_name is specified pipeline_class_name = kwargs.get( prefix_with_dot + "pipeline_class_name" ) or kwargs.get("pipeline_class_name") is_safetensors_file = os.path.isfile(model_path) and model_path.endswith( ".safetensors" ) # 1. Get the pipeline config class from the registry from sglang.multimodal_gen.configs.pipeline_configs.flux import ( Flux2PipelineConfig, ) from sglang.multimodal_gen.registry import get_pipeline_config_classes # If model_path is a safetensors file and pipeline_class_name is specified, # try to get PipelineConfig from the registry first if is_safetensors_file and pipeline_class_name: config_classes = get_pipeline_config_classes(pipeline_class_name) if config_classes is not None: pipeline_config_cls, _ = config_classes logger.info( f"Detected safetensors file with {pipeline_class_name}, " f"using {pipeline_config_cls.__name__} directly without model_index.json" ) else: model_info = get_model_info( model_path, backend=kwargs.get("backend"), model_id=kwargs.get("model_id"), ) if model_info is None: from sglang.multimodal_gen.registry import ( _PIPELINE_CONFIG_REGISTRY, _discover_and_register_pipelines, ) _discover_and_register_pipelines() available_pipelines = list(_PIPELINE_CONFIG_REGISTRY.keys()) raise ValueError( f"Could not get model info for '{model_path}'. " f"If using a safetensors file, please specify a valid pipeline_class_name. " f"Available pipelines with config classes: {available_pipelines}" ) pipeline_config_cls = model_info.pipeline_config_cls else: model_info = get_model_info( model_path, backend=kwargs.get("backend"), model_id=kwargs.get("model_id"), ) if model_info is None: raise ValueError( f"Could not get model info for '{model_path}'. " f"If using a safetensors file, please specify pipeline_class_name" ) # 1.5. Adjust pipeline config for fine-tuned VAE if needed pipeline_config_cls = model_info.pipeline_config_cls vae_path = kwargs.get(prefix_with_dot + "vae_path") or kwargs.get("vae_path") if vae_path is None: component_paths = kwargs.get( prefix_with_dot + "component_paths" ) or kwargs.get("component_paths") if isinstance(component_paths, dict): vae_path = component_paths.get("vae") # Check if this is a Flux2 model with fal/FLUX.2-Tiny-AutoEncoder if ( isinstance(pipeline_config_cls, type) and issubclass(pipeline_config_cls, Flux2PipelineConfig) and vae_path is not None and "FLUX.2-Tiny-AutoEncoder" in vae_path ): from sglang.multimodal_gen.configs.pipeline_configs.flux_finetuned import ( Flux2FinetunedPipelineConfig, ) pipeline_config_cls = Flux2FinetunedPipelineConfig pipeline_config = pipeline_config_cls() # 2. Load PipelineConfig from a json file or a PipelineConfig object if provided if isinstance(pipeline_config_or_path, str): pipeline_config.load_from_json(pipeline_config_or_path) kwargs[prefix_with_dot + "pipeline_config_path"] = pipeline_config_or_path elif isinstance(pipeline_config_or_path, PipelineConfig): pipeline_config = pipeline_config_or_path elif isinstance(pipeline_config_or_path, dict): pipeline_config.update_pipeline_config(pipeline_config_or_path) # 3. Update PipelineConfig from CLI arguments if provided kwargs[prefix_with_dot + "model_path"] = model_path pipeline_config.update_config_from_dict(kwargs, config_cli_prefix) return pipeline_config def check_pipeline_config(self) -> None: if self.vae_sp and not self.vae_tiling: raise ValueError( "Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True." ) if len(self.text_encoder_configs) != len(self.text_encoder_precisions): raise ValueError( f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text encoder precisions ({len(self.text_encoder_precisions)})" ) if len(self.text_encoder_configs) != len(self.preprocess_text_funcs): raise ValueError( f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})" ) if len(self.preprocess_text_funcs) != len(self.postprocess_text_funcs): raise ValueError( f"Length of text postprocess functions ({len(self.postprocess_text_funcs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})" ) def dump_to_json(self, file_path: str): output_dict = shallow_asdict(self) del_keys = [] for key, value in output_dict.items(): if isinstance(value, ModelConfig): model_dict = asdict(value) # Model Arch Config should be hidden away from the users model_dict.pop("arch_config") output_dict[key] = model_dict elif isinstance(value, tuple) and all( isinstance(v, ModelConfig) for v in value ): model_dicts = [] for v in value: model_dict = asdict(v) # Model Arch Config should be hidden away from the users model_dict.pop("arch_config") model_dicts.append(model_dict) output_dict[key] = model_dicts elif isinstance(value, tuple) and all(callable(f) for f in value): # Skip dumping functions del_keys.append(key) for key in del_keys: output_dict.pop(key, None) with open(file_path, "w") as f: json.dump(output_dict, f, indent=2) def load_from_json(self, file_path: str): with open(file_path) as f: input_pipeline_dict = json.load(f) self.update_pipeline_config(input_pipeline_dict) def update_pipeline_config(self, source_pipeline_dict: dict[str, Any]) -> None: for f in fields(self): key = f.name if key in source_pipeline_dict: current_value = getattr(self, key) new_value = source_pipeline_dict[key] # If it's a nested ModelConfig, update it recursively if isinstance(current_value, ModelConfig): current_value.update_model_config(new_value) elif isinstance(current_value, tuple) and all( isinstance(v, ModelConfig) for v in current_value ): assert len(current_value) == len( new_value ), "Users shouldn't delete or add text encoder config objects in your json" for target_config, source_config in zip( current_value, new_value, strict=True ): target_config.update_model_config(source_config) else: setattr(self, key, new_value) if hasattr(self, "__post_init__"): self.__post_init__() @dataclass class ImagePipelineConfig(PipelineConfig): """Base config for image generation pipelines with token-like latents [B, S, D].""" def _prepare_sigmas(self, sigmas, num_inference_steps): sigmas = ( np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas ) return sigmas def shard_latents_for_sp(self, batch, latents): # latents: [B, H * W, C] sp_world_size, rank_in_sp_group = get_sp_world_size(), get_sp_parallel_rank() seq_len = latents.shape[1] # TODO: reuse code in PipelineConfig::shard_latents_for_sp # Pad to next multiple of SP degree if needed if seq_len % sp_world_size != 0: pad_len = sp_world_size - (seq_len % sp_world_size) pad = torch.zeros( (*latents.shape[:1], pad_len, *latents.shape[2:]), dtype=latents.dtype, device=latents.device, ) latents = torch.cat([latents, pad], dim=1) sharded_tensor = rearrange( latents, "b (n s) d -> b n s d", n=sp_world_size ).contiguous() sharded_tensor = sharded_tensor[:, rank_in_sp_group, :, :] return sharded_tensor, True def gather_latents_for_sp(self, latents): # For image latents [B, S_local, D], gather along sequence dim=1 latents = sequence_model_parallel_all_gather(latents, dim=1) return latents def _unpad_and_unpack_latents(self, latents, batch): vae_scale_factor = self.vae_config.arch_config.vae_scale_factor channels = self.dit_config.arch_config.in_channels batch_size = latents.shape[0] height = 2 * (int(batch.height) // (vae_scale_factor * 2)) width = 2 * (int(batch.width) // (vae_scale_factor * 2)) latents = maybe_unpad_latents(latents, batch) latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2) latents = latents.permute(0, 3, 1, 4, 2, 5) return latents, batch_size, channels, height, width @dataclass class SpatialImagePipelineConfig(ImagePipelineConfig): """Base config for spatial image pipelines (e.g. GLM-Image) with 4D latents (B, C, H', W'). Overrides shard_latents_for_sp / gather_latents_for_sp to shard along the height dimension so that each SP rank gets (B, C, H'_local, W') instead of using the token-style (B, S, C) path. """ def shard_latents_for_sp(self, batch, latents): # 4D latents (B, C, H', W') -> shard along H' (dim=2); otherwise fall back to base (B, S, C) sp_world_size = get_sp_world_size() if sp_world_size <= 1: return latents, False if latents.dim() != 4: return super().shard_latents_for_sp(batch, latents) # (B, C, H', W') _, _, h_lat, w_lat = latents.shape if h_lat % sp_world_size != 0: pad_len = sp_world_size - (h_lat % sp_world_size) pad = torch.zeros( (latents.shape[0], latents.shape[1], pad_len, latents.shape[3]), dtype=latents.dtype, device=latents.device, ) latents = torch.cat([latents, pad], dim=2) h_lat = latents.shape[2] rank_in_sp_group = get_sp_parallel_rank() chunk_size = h_lat // sp_world_size h0 = rank_in_sp_group * chunk_size h1 = h0 + chunk_size sharded = latents[:, :, h0:h1, :].contiguous() return sharded, True def gather_latents_for_sp(self, latents): if get_sp_world_size() <= 1: return latents if latents.dim() != 4: return super().gather_latents_for_sp(latents) # Gather along dim=2 (H') to match shard_latents_for_sp return sequence_model_parallel_all_gather(latents, dim=2) @dataclass class SlidingTileAttnConfig(PipelineConfig): """Configuration for sliding tile attention.""" # Override any BaseConfig defaults as needed # Add sliding tile specific parameters window_size: int = 16 stride: int = 8 # You can provide custom defaults for inherited fields height: int = 576 width: int = 1024 # Additional configuration specific to sliding tile attention pad_to_square: bool = False use_overlap_optimization: bool = True def parse_int_list(value: str) -> list[int]: """Parse a comma-separated string of integers into a list.""" if not value: return [] return [int(x.strip()) for x in value.split(",")] ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/diffusers_generic.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Generic pipeline configuration for diffusers backend. This module provides a minimal pipeline configuration that works with the diffusers backend. Since diffusers handles its own model loading and configuration, this config is intentionally minimal. """ from dataclasses import dataclass, field from typing import Any from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, PipelineConfig, ) @dataclass class DiffusersGenericPipelineConfig(PipelineConfig): """ Generic pipeline configuration for diffusers backend. This is a minimal configuration since the diffusers backend handles most configuration internally. It provides sensible defaults for the required fields. """ # default to T2I since it's the most common task_type: ModelTaskType = ModelTaskType.T2I dit_precision: str = "bf16" vae_precision: str = "bf16" should_use_guidance: bool = True embedded_cfg_scale: float = 1.0 flow_shift: float | None = None disable_autocast: bool = True # let diffusers handle dtype # diffusers handles its own loading dit_config: DiTConfig = field(default_factory=DiTConfig) vae_config: VAEConfig = field(default_factory=VAEConfig) image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig) text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (EncoderConfig(),) ) text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp16",)) # VAE settings vae_tiling: bool = False # diffusers handles this vae_slicing: bool = False # slice VAE decode for lower memory usage vae_sp: bool = False # Quantization config for pipeline-level quantization # See: https://huggingface.co/docs/diffusers/main/en/quantization/overview # Use PipelineQuantizationConfig for component-level control: # from diffusers.quantizers import PipelineQuantizationConfig # quantization_config = PipelineQuantizationConfig( # quant_backend="bitsandbytes_4bit", # quant_kwargs={"load_in_4bit": True, "bnb_4bit_compute_dtype": torch.bfloat16}, # components_to_quantize=["transformer", "text_encoder_2"], # ) quantization_config: Any = None def check_pipeline_config(self) -> None: """ Override to skip most validation since diffusers handles its own config. """ pass def adjust_size(self, width, height, image): """ Pass through - diffusers handles size adjustments. """ return width, height def adjust_num_frames(self, num_frames): """ Pass through - diffusers handles frame count. """ return num_frames ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/flux.py ================================================ import math from dataclasses import dataclass, field from typing import Callable, List, Optional import PIL import torch from diffusers.image_processor import VaeImageProcessor from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig from sglang.multimodal_gen.configs.models.encoders import ( BaseEncoderOutput, CLIPTextConfig, T5Config, TextEncoderConfig, ) from sglang.multimodal_gen.configs.models.encoders.base import TextEncoderArchConfig from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig from sglang.multimodal_gen.configs.models.encoders.qwen_image import ( _is_transformer_layer, ) from sglang.multimodal_gen.configs.models.vaes.flux import Flux2VAEConfig, FluxVAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ImagePipelineConfig, ModelTaskType, preprocess_text, shard_rotary_emb_for_sp, ) from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( clip_postprocess_text, clip_preprocess_text, ) from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import _pack_latents from sglang.multimodal_gen.runtime.distributed import get_local_torch_device def t5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: return outputs.last_hidden_state @dataclass class FluxPipelineConfig(ImagePipelineConfig): """Configuration for the FLUX pipeline.""" embedded_cfg_scale: float = 3.5 task_type: ModelTaskType = ModelTaskType.T2I vae_tiling: bool = False vae_sp: bool = False dit_config: DiTConfig = field(default_factory=FluxConfig) # VAE vae_config: VAEConfig = field(default_factory=FluxVAEConfig) enable_autocast: bool = False # Text encoding stage text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (CLIPTextConfig(), T5Config()) ) text_encoder_precisions: tuple[str, ...] = field( default_factory=lambda: ("bf16", "bf16") ) preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( default_factory=lambda: (clip_preprocess_text, preprocess_text), ) postprocess_text_funcs: tuple[Callable[[str], str], ...] = field( default_factory=lambda: (clip_postprocess_text, t5_postprocess_text) ) text_encoder_extra_args: list[dict] = field( default_factory=lambda: [ dict( max_length=77, padding="max_length", truncation=True, return_overflowing_tokens=False, return_length=False, ), None, ] ) def prepare_sigmas(self, sigmas, num_inference_steps): return self._prepare_sigmas(sigmas, num_inference_steps) def prepare_latent_shape(self, batch, batch_size, num_frames): height = 2 * ( batch.height // (self.vae_config.arch_config.vae_scale_factor * 2) ) width = 2 * (batch.width // (self.vae_config.arch_config.vae_scale_factor * 2)) num_channels_latents = self.dit_config.arch_config.in_channels // 4 shape = (batch_size, num_channels_latents, height, width) return shape def maybe_pack_latents(self, latents, batch_size, batch): height = 2 * ( batch.height // (self.vae_config.arch_config.vae_scale_factor * 2) ) width = 2 * (batch.width // (self.vae_config.arch_config.vae_scale_factor * 2)) num_channels_latents = self.dit_config.arch_config.in_channels // 4 # pack latents return _pack_latents(latents, batch_size, num_channels_latents, height, width) def get_pos_prompt_embeds(self, batch): return batch.prompt_embeds[1] def get_neg_prompt_embeds(self, batch): return batch.negative_prompt_embeds[1] def _prepare_latent_image_ids(self, original_height, original_width, device): vae_scale_factor = self.vae_config.arch_config.vae_scale_factor height = int(original_height) // (vae_scale_factor * 2) width = int(original_width) // (vae_scale_factor * 2) latent_image_ids = torch.zeros(height, width, 3, device=device) latent_image_ids[..., 1] = ( latent_image_ids[..., 1] + torch.arange(height, device=device)[:, None] ) latent_image_ids[..., 2] = ( latent_image_ids[..., 2] + torch.arange(width, device=device)[None, :] ) latent_image_id_height, latent_image_id_width, latent_image_id_channels = ( latent_image_ids.shape ) latent_image_ids = latent_image_ids.reshape( latent_image_id_height * latent_image_id_width, latent_image_id_channels ) return latent_image_ids def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch): txt_ids = torch.zeros(prompt_embeds.shape[1], 3, device=device) img_ids = self._prepare_latent_image_ids( original_height=height, original_width=width, device=device, ) # NOTE(mick): prepare it here, to avoid unnecessary computations img_cos, img_sin = rotary_emb.forward(img_ids) img_cos = shard_rotary_emb_for_sp(img_cos) img_sin = shard_rotary_emb_for_sp(img_sin) txt_cos, txt_sin = rotary_emb.forward(txt_ids) cos = torch.cat([txt_cos, img_cos], dim=0).to(device=device) sin = torch.cat([txt_sin, img_sin], dim=0).to(device=device) return cos, sin def post_denoising_loop(self, latents, batch): # unpack latents for flux ( latents, batch_size, channels, height, width, ) = self._unpad_and_unpack_latents(latents, batch) latents = latents.reshape(batch_size, channels // (2 * 2), height, width) return latents def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): return { "freqs_cis": self.get_freqs_cis( batch.prompt_embeds[1], batch.width, batch.height, device, rotary_emb, batch, ), "pooled_projections": ( batch.pooled_embeds[0] if batch.pooled_embeds else None ), } def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): return { "freqs_cis": self.get_freqs_cis( batch.negative_prompt_embeds[1], batch.width, batch.height, device, rotary_emb, batch, ), "pooled_projections": ( batch.neg_pooled_embeds[0] if batch.neg_pooled_embeds else None ), } def _prepare_latent_ids( latents: torch.Tensor, # (B, C, H, W) ): r""" Generates 4D position coordinates (T, H, W, L) for latent tensors. Args: latents (torch.Tensor): Latent tensor of shape (B, C, H, W) Returns: torch.Tensor: Position IDs tensor of shape (B, H*W, 4) All batches share the same coordinate structure: T=0, H=[0..H-1], W=[0..W-1], L=0 """ batch_size, _, height, width = latents.shape t = torch.arange(1) # [0] - time dimension h = torch.arange(height) w = torch.arange(width) layer = torch.arange(1) # [0] - layer dimension # Create position IDs: (H*W, 4) latent_ids = torch.cartesian_prod(t, h, w, layer) # Expand to batch: (B, H*W, 4) latent_ids = latent_ids.unsqueeze(0).expand(batch_size, -1, -1) return latent_ids def _unpack_latents_with_ids( x: torch.Tensor, x_ids: torch.Tensor ) -> list[torch.Tensor]: """ using position ids to scatter tokens into place """ x_list = [] x_ids = x_ids.to(device=x.device) for data, pos in zip(x, x_ids): _, ch = data.shape # noqa: F841 h_ids = pos[:, 1].to(torch.int64) w_ids = pos[:, 2].to(torch.int64) h = torch.max(h_ids) + 1 w = torch.max(w_ids) + 1 flat_ids = h_ids * w + w_ids out = torch.zeros((h * w, ch), device=data.device, dtype=data.dtype) out.scatter_(0, flat_ids.unsqueeze(1).expand(-1, ch), data) # reshape from (H * W, C) to (H, W, C) and permute to (C, H, W) out = out.view(h, w, ch).permute(2, 0, 1) x_list.append(out) return torch.stack(x_list, dim=0) def _patchify_latents(latents): batch_size, num_channels_latents, height, width = latents.shape latents = latents.view( batch_size, num_channels_latents, height // 2, 2, width // 2, 2 ) latents = latents.permute(0, 1, 3, 5, 2, 4) latents = latents.reshape( batch_size, num_channels_latents * 4, height // 2, width // 2 ) return latents def _unpatchify_latents(latents): batch_size, num_channels_latents, height, width = latents.shape latents = latents.reshape( batch_size, num_channels_latents // (2 * 2), 2, 2, height, width ) latents = latents.permute(0, 1, 4, 2, 5, 3) latents = latents.reshape( batch_size, num_channels_latents // (2 * 2), height * 2, width * 2 ) return latents def _prepare_text_ids( x: torch.Tensor, # (B, L, D) or (L, D) t_coord: Optional[torch.Tensor] = None, ): B, L, _ = x.shape out_ids = [] for i in range(B): t = torch.arange(1) if t_coord is None else t_coord[i] h = torch.arange(1) w = torch.arange(1) layer = torch.arange(L) coords = torch.cartesian_prod(t, h, w, layer) out_ids.append(coords) return torch.stack(out_ids) def _prepare_image_ids( image_latents: List[torch.Tensor], # [(1, C, H, W), (1, C, H, W), ...] scale: int = 10, ): if not isinstance(image_latents, list): raise ValueError( f"Expected `image_latents` to be a list, got {type(image_latents)}." ) # create time offset for each reference image t_coords = [scale + scale * t for t in torch.arange(0, len(image_latents))] t_coords = [t.view(-1) for t in t_coords] image_latent_ids = [] for x, t in zip(image_latents, t_coords): x = x.squeeze(0) _, height, width = x.shape x_ids = torch.cartesian_prod( t, torch.arange(height), torch.arange(width), torch.arange(1) ) image_latent_ids.append(x_ids) image_latent_ids = torch.cat(image_latent_ids, dim=0) image_latent_ids = image_latent_ids.unsqueeze(0) return image_latent_ids def flux2_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: hidden_states_layers: list[int] = [10, 20, 30] out = torch.stack([outputs.hidden_states[k] for k in hidden_states_layers], dim=1) batch_size, num_channels, seq_len, hidden_dim = out.shape prompt_embeds = out.permute(0, 2, 1, 3).reshape( batch_size, seq_len, num_channels * hidden_dim ) return prompt_embeds def flux2_klein_postprocess_text( outputs: BaseEncoderOutput, _text_inputs ) -> torch.Tensor: hidden_states_layers: list[int] = [9, 18, 27] out = torch.stack([outputs.hidden_states[k] for k in hidden_states_layers], dim=1) batch_size, num_channels, seq_len, hidden_dim = out.shape prompt_embeds = out.permute(0, 2, 1, 3).reshape( batch_size, seq_len, num_channels * hidden_dim ) return prompt_embeds @dataclass class Flux2MistralTextArchConfig(TextEncoderArchConfig): stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=lambda: [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), ("qkv_proj", "k_proj", "k"), ("qkv_proj", "v_proj", "v"), ] ) _fsdp_shard_conditions: list = field( default_factory=lambda: [_is_transformer_layer] ) def __post_init__(self): self.tokenizer_kwargs = { "padding": "max_length", "truncation": True, "max_length": 512, "add_special_tokens": True, "return_attention_mask": True, "return_tensors": "pt", } @dataclass class Flux2MistralTextConfig(TextEncoderConfig): arch_config: TextEncoderArchConfig = field( default_factory=Flux2MistralTextArchConfig ) def format_text_input(prompts: List[str], system_message: str = None): # Remove [IMG] tokens from prompts to avoid Pixtral validation issues # when truncation is enabled. The processor counts [IMG] tokens and fails # if the count changes after truncation. cleaned_txt = [prompt.replace("[IMG]", "") for prompt in prompts] return [ [ { "role": "system", "content": [{"type": "text", "text": system_message}], }, {"role": "user", "content": [{"type": "text", "text": prompt}]}, ] for prompt in cleaned_txt ] def flux_2_preprocess_text(prompt: str): system_message = "You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object attribution and actions without speculation." return format_text_input([prompt], system_message=system_message) # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._pack_latents def flux2_pack_latents(latents): batch_size, num_channels, height, width = latents.shape latents = latents.reshape(batch_size, num_channels, height * width).permute(0, 2, 1) return latents @dataclass class Flux2PipelineConfig(FluxPipelineConfig): embedded_cfg_scale: float = 4.0 task_type: ModelTaskType = ModelTaskType.TI2I text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",)) text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (Flux2MistralTextConfig(),) ) preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( default_factory=lambda: (flux_2_preprocess_text,), ) postprocess_text_funcs: tuple[Callable[[str], str], ...] = field( default_factory=lambda: (flux2_postprocess_text,) ) vae_config: VAEConfig = field(default_factory=Flux2VAEConfig) def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict: # flatten to 1-d list prompts = [p for prompt in prompts for p in prompt] inputs = tokenizer.apply_chat_template( prompts, add_generation_prompt=False, tokenize=True, return_dict=True, return_tensors="pt", padding="max_length", truncation=True, # 2048 from official github repo, 512 from diffusers max_length=512, ) return inputs def prepare_latent_shape(self, batch, batch_size, num_frames): height = 2 * ( batch.height // (self.vae_config.arch_config.vae_scale_factor * 2) ) width = 2 * (batch.width // (self.vae_config.arch_config.vae_scale_factor * 2)) num_channels_latents = self.dit_config.arch_config.in_channels shape = (batch_size, num_channels_latents, height // 2, width // 2) return shape def get_pos_prompt_embeds(self, batch): return batch.prompt_embeds[0] def get_neg_prompt_embeds(self, batch): return batch.negative_prompt_embeds[0] def calculate_condition_image_size( self, image, width, height ) -> Optional[tuple[int, int]]: vae_scale_factor = self.vae_config.arch_config.vae_scale_factor multiple_of = vae_scale_factor * 2 target_area: int = 1024 * 1024 if width is not None and height is not None: new_width, new_height = width, height if width * height > target_area: scale = math.sqrt(target_area / (width * height)) new_width = int(width * scale) new_height = int(height * scale) # Flux requires multiples of (VAE scale 8 * Patch size 2) new_width = (new_width // multiple_of) * multiple_of new_height = (new_height // multiple_of) * multiple_of if new_width != width or new_height != height: return new_width, new_height return None def preprocess_condition_image( self, image, target_width, target_height, vae_image_processor: VaeImageProcessor ): img = image.resize((target_width, target_height), PIL.Image.Resampling.LANCZOS) image_width, image_height = img.size vae_scale_factor = self.vae_config.arch_config.vae_scale_factor multiple_of = vae_scale_factor * 2 image_width = (image_width // multiple_of) * multiple_of image_height = (image_height // multiple_of) * multiple_of img = vae_image_processor.preprocess( img, height=image_height, width=image_width, resize_mode="crop" ) return img, (image_width, image_height) def postprocess_image_latent(self, latent_condition, batch): batch_size = batch.batch_size # latent: (1, 128, 32, 32) packed = self.maybe_pack_latents( latent_condition, None, batch ) # (1, 1024, 128) packed = packed.squeeze(0) # (1024, 128) - remove batch dim # Concatenate all reference tokens along sequence dimension image_latents = packed.unsqueeze(0) # (1, N*1024, 128) image_latents = image_latents.repeat(batch_size, 1, 1) return image_latents def prepare_condition_image_latent_ids(self, image_latents, batch): image_latent_ids = _prepare_image_ids(image_latents) image_latent_ids = image_latent_ids.repeat(batch.batch_size, 1, 1) batch.condition_image_latent_ids = image_latent_ids.to(get_local_torch_device()) def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch): txt_ids = _prepare_text_ids(prompt_embeds).to(device=device) img_ids = batch.latent_ids if img_ids.ndim == 3: img_ids = img_ids[0] if txt_ids.ndim == 3: txt_ids = txt_ids[0] # NOTE(mick): prepare it here, to avoid unnecessary computations img_cos, img_sin = rotary_emb.forward(img_ids) img_cos = shard_rotary_emb_for_sp(img_cos) img_sin = shard_rotary_emb_for_sp(img_sin) if batch.image_latent is not None: cond_ids = batch.condition_image_latent_ids if cond_ids.ndim == 3: cond_ids = cond_ids[0] cond_cos, cond_sin = rotary_emb.forward(cond_ids) cond_cos = shard_rotary_emb_for_sp(cond_cos) cond_sin = shard_rotary_emb_for_sp(cond_sin) img_cos = torch.cat([img_cos, cond_cos], dim=0) img_sin = torch.cat([img_sin, cond_sin], dim=0) txt_cos, txt_sin = rotary_emb.forward(txt_ids) cos = torch.cat([txt_cos, img_cos], dim=0).to(device=device) sin = torch.cat([txt_sin, img_sin], dim=0).to(device=device) return cos, sin def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): return { "freqs_cis": self.get_freqs_cis( batch.prompt_embeds[0], batch.width, batch.height, device, rotary_emb, batch, ) } def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): return {} def maybe_pack_latents(self, latents, batch_size, batch): return flux2_pack_latents(latents) def maybe_prepare_latent_ids(self, latents): return _prepare_latent_ids(latents) def postprocess_vae_encode(self, image_latents, vae): # patchify image_latents = _patchify_latents(image_latents) return image_latents def _check_vae_has_bn(self, vae): """Check if VAE has bn attribute (cached check to avoid repeated hasattr calls).""" if not hasattr(self, "_vae_has_bn_cache"): self._vae_has_bn_cache = hasattr(vae, "bn") and vae.bn is not None return self._vae_has_bn_cache def preprocess_decoding(self, latents, server_args=None, vae=None): """Preprocess latents before decoding. Dynamically adapts based on VAE type: - Standard Flux2 VAE (has bn): needs unpatchify (128 channels -> 32 channels) - Distilled VAE (no bn): keeps patchified latents (128 channels) """ if vae is not None and self._check_vae_has_bn(vae): return _unpatchify_latents(latents) return latents def get_decode_scale_and_shift(self, device, dtype, vae): """Get scale and shift for decoding. Dynamically adapts based on VAE type: - Standard Flux2 VAE (has bn): uses BatchNorm statistics - Distilled VAE (no bn): uses scaling_factor from config """ vae_arch_config = self.vae_config.arch_config if self._check_vae_has_bn(vae): # Standard Flux2 VAE: use BatchNorm statistics latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device, dtype) latents_bn_std = torch.sqrt( vae.bn.running_var.view(1, -1, 1, 1) + vae_arch_config.batch_norm_eps ).to(device, dtype) return 1 / latents_bn_std, latents_bn_mean # Distilled VAE or unknown: use scaling_factor scaling_factor = ( getattr(vae.config, "scaling_factor", None) if hasattr(vae, "config") else getattr(vae, "scaling_factor", None) ) or getattr(vae_arch_config, "scaling_factor", 0.13025) scale = torch.tensor(scaling_factor, device=device, dtype=dtype).view( 1, 1, 1, 1 ) return 1 / scale, None def post_denoising_loop(self, latents, batch): latent_ids = batch.latent_ids latents = _unpack_latents_with_ids(latents, latent_ids) return latents def slice_noise_pred(self, noise, latents): # remove noise over input image noise = noise[:, : latents.size(1) :] return noise @dataclass class Flux2KleinPipelineConfig(Flux2PipelineConfig): # Klein is distilled, so no guidance embeddings should_use_guidance: bool = False text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",)) text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (Qwen3TextConfig(),) ) preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( default_factory=lambda: (preprocess_text,), ) postprocess_text_funcs: tuple[Callable[[str], str], ...] = field( default_factory=lambda: (flux2_klein_postprocess_text,) ) def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict: if prompts and isinstance(prompts[0], list): prompts = [p for prompt in prompts for p in prompt] def _apply_chat_template(prompt: str) -> str: messages = [{"role": "user", "content": prompt}] try: return tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, enable_thinking=False, ) except TypeError: return tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, ) texts = [_apply_chat_template(prompt) for prompt in prompts] tok_kwargs = dict(tok_kwargs or {}) max_length = tok_kwargs.pop("max_length", 512) padding = tok_kwargs.pop("padding", "max_length") truncation = tok_kwargs.pop("truncation", True) return_tensors = tok_kwargs.pop("return_tensors", "pt") return tokenizer( texts, padding=padding, truncation=truncation, max_length=max_length, return_tensors=return_tensors, **tok_kwargs, ) ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/flux_finetuned.py ================================================ """ Pipeline configuration for Flux fine-tuned/distilled models. This module provides specialized handling for Flux fine-tuned models from HuggingFace, such as fal/FLUX.2-Tiny-AutoEncoder and other community fine-tuned variants. Key differences from standard Flux2PipelineConfig: - Handles custom VAE architectures loaded via auto_map - Supports both patchified (128 channels) and unpatchified (32 channels) latents - Dynamically adapts scale/shift based on VAE type - Properly handles 5D latents (batch, channels, frames, height, width) for decoding """ from dataclasses import dataclass import torch from sglang.multimodal_gen.configs.pipeline_configs.flux import ( Flux2PipelineConfig, _unpatchify_latents, ) @dataclass class Flux2FinetunedPipelineConfig(Flux2PipelineConfig): """ Pipeline configuration for Flux fine-tuned/distilled models. This configuration automatically detects and handles custom VAE architectures (e.g., Flux2TinyAutoEncoder) loaded via HuggingFace's auto_map mechanism. Features: - Automatic VAE type detection (standard vs. distilled) - Proper handling of patchified/unpatchified latents - Support for custom scaling factors from fine-tuned models - 5D latents support for both single-frame and multi-frame generation """ def preprocess_decoding( self, latents: torch.Tensor, server_args=None, vae=None ) -> torch.Tensor: """ Preprocess latents before decoding. Handles both standard Flux2 VAE and fine-tuned/distilled VAEs: - Standard Flux2 VAE (has bn): needs unpatchify (128 channels -> 32 channels) - Distilled/Finetuned VAE (no bn): keeps patchified latents (128 channels) Also handles 5D latents (batch, channels, frames, height, width) by converting to 4D (batch, channels, height, width) for single-frame cases. Args: latents: Input latents tensor, can be 4D or 5D server_args: Server arguments (optional, for compatibility) vae: VAE model instance for dynamic type detection Returns: Preprocessed latents ready for VAE decoding """ # Handle 5D latents (batch, channels, frames, height, width) if latents.ndim == 5: batch_size, channels, frames, height, width = latents.shape if frames == 1: latents = latents.squeeze(2) else: latents = latents.permute(0, 2, 1, 3, 4).contiguous() latents = latents.view(batch_size * frames, channels, height, width) if vae is not None and self._check_vae_has_bn(vae): latents = _unpatchify_latents(latents) return latents def get_decode_scale_and_shift(self, device, dtype, vae): """ Get scale and shift for decoding. Dynamically adapts based on VAE type: - Standard Flux2 VAE (has bn): uses BatchNorm statistics - Distilled/Finetuned VAE (no bn): uses scaling_factor from config Args: device: Target device for tensors dtype: Target dtype for tensors vae: VAE model instance Returns: Tuple of (scaling_factor, shift_factor) - scaling_factor: Tensor or scalar to divide latents by - shift_factor: Tensor or scalar to add to latents (None for distilled VAEs) """ vae_arch_config = self.vae_config.arch_config if self._check_vae_has_bn(vae): # Standard Flux2 VAE: use BatchNorm statistics latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device, dtype) latents_bn_std = torch.sqrt( vae.bn.running_var.view(1, -1, 1, 1) + vae_arch_config.batch_norm_eps ).to(device, dtype) return 1 / latents_bn_std, latents_bn_mean # Distilled/Finetuned VAE: Flux2TinyAutoEncoder doesn't need external scaling scale = torch.tensor(1.0, device=device, dtype=dtype).view(1, 1, 1, 1) return scale, None ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py ================================================ from dataclasses import dataclass, field import torch from diffusers.image_processor import VaeImageProcessor from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig from sglang.multimodal_gen.configs.models.dits.glmimage import GlmImageDitConfig from sglang.multimodal_gen.configs.models.encoders.base import EncoderConfig from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config from sglang.multimodal_gen.configs.models.vaes.glmimage import GlmImageVAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, SpatialImagePipelineConfig, ) @dataclass class GlmImagePipelineConfig(SpatialImagePipelineConfig): """Configuration for the GlmImage pipeline.""" vae_precision: str = "bf16" should_use_guidance: bool = False task_type: ModelTaskType = ModelTaskType.T2I vae_tiling: bool = False vae_sp: bool = False text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (T5Config(),) ) dit_config: DiTConfig = field(default_factory=GlmImageDitConfig) # VAE vae_config: VAEConfig = field(default_factory=GlmImageVAEConfig) # GLM-Image uses T5 text encoder; base default is EncoderConfig() which lacks # parallel_folding and causes AttributeError + fallback to native T5 with missing weights. text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (T5Config(),) ) enable_autocast: bool = False def __post_init__(self): self.vae_scale_factor = self.vae_config.get_vae_scale_factor() self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) def get_freqs_cis(self, batch, device, rotary_emb, dtype): height = batch.height // self.vae_scale_factor width = batch.width // self.vae_scale_factor hidden_states = torch.empty(1, 1, height, width, device=device, dtype=dtype) freqs_cis = rotary_emb(hidden_states) return freqs_cis def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): return { "prior_token_id": batch.prior_token_id, "prior_token_drop": batch.prior_token_drop_cond, "crop_coords": batch.crop_coords, "target_size": batch.target_size, "kv_caches": batch.kv_caches, "kv_caches_mode": "read", "freqs_cis": self.get_freqs_cis(batch, device, rotary_emb, dtype), } def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): return { "prior_token_id": batch.prior_token_id, "prior_token_drop": batch.prior_token_drop_uncond, "crop_coords": batch.crop_coords, "target_size": batch.target_size, "kv_caches": batch.kv_caches, "kv_caches_mode": "skip", "freqs_cis": self.get_freqs_cis(batch, device, rotary_emb, dtype), } def get_decode_scale_and_shift(self, device, dtype, vae): latents_mean = ( torch.tensor(self.vae_config.latents_mean) .view(1, self.vae_config.latent_channels, 1, 1) .to(device, dtype) ) latents_std = ( torch.tensor(self.vae_config.latents_std) .view(1, self.vae_config.latent_channels, 1, 1) .to(device, dtype) ) return 1.0 / latents_std, latents_mean def post_denoising_loop(self, latents, batch): if getattr(batch, "kv_caches", None) is not None: batch.kv_caches.clear() return latents.bfloat16() def post_decoding(self, frames, server_args): return self.image_processor.postprocess(frames, output_type="latent") ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/helios.py ================================================ # SPDX-License-Identifier: Apache-2.0 from collections.abc import Callable from dataclasses import dataclass, field import torch from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput, T5Config from sglang.multimodal_gen.configs.models.encoders.t5 import T5ArchConfig from sglang.multimodal_gen.configs.models.vaes import WanVAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, PipelineConfig, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # Helios UMT5 max sequence length (used for both tokenizer and post-processing padding) # Matches diffusers HeliosPipeline.__call__ default max_sequence_length=512 HELIOS_MAX_SEQUENCE_LENGTH = 512 def umt5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: """Post-process UMT5 text encoder outputs, padding to HELIOS_MAX_SEQUENCE_LENGTH tokens.""" max_seq_len = HELIOS_MAX_SEQUENCE_LENGTH mask: torch.Tensor = outputs.attention_mask hidden_state: torch.Tensor = outputs.last_hidden_state seq_lens = mask.gt(0).sum(dim=1).long() assert torch.isnan(hidden_state).sum() == 0 prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)] prompt_embeds_tensor: torch.Tensor = torch.stack( [ torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in prompt_embeds ], dim=0, ) return prompt_embeds_tensor @dataclass class HeliosT2VConfig(PipelineConfig): """Configuration for the Helios T2V pipeline.""" task_type: ModelTaskType = ModelTaskType.T2V # DiT dit_config: DiTConfig = field(default_factory=HeliosConfig) # VAE (same as Wan) vae_config: VAEConfig = field(default_factory=WanVAEConfig) vae_tiling: bool = False vae_sp: bool = False # Denoising stage flow_shift: float | None = 1.0 # Text encoding stage (UMT5 is T5-compatible) text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: ( T5Config(arch_config=T5ArchConfig(text_len=HELIOS_MAX_SEQUENCE_LENGTH)), ) ) postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = ( field(default_factory=lambda: (umt5_postprocess_text,)) ) # Precision for each component precision: str = "bf16" vae_precision: str = "fp32" text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32",)) # Helios-specific chunked denoising params num_latent_frames_per_chunk: int = 9 history_sizes: list[int] = field(default_factory=lambda: [16, 2, 1]) is_cfg_zero_star: bool = False zero_steps: int = 1 keep_first_frame: bool = True # Stage 2 (Pyramid SR) & Stage 3 (DMD) params is_enable_stage2: bool = False pyramid_num_stages: int = 3 pyramid_num_inference_steps_list: list[int] = field( default_factory=lambda: [10, 10, 10] ) is_distilled: bool = False is_amplify_first_chunk: bool = False scheduler_type: str = "unipc" gamma: float = 1 / 3 def __post_init__(self): self.vae_config.load_encoder = False self.vae_config.load_decoder = True @dataclass class HeliosMidConfig(HeliosT2VConfig): """Configuration for Helios-Mid (Stage 1 + Stage 2 pyramid SR).""" is_enable_stage2: bool = True is_cfg_zero_star: bool = True pyramid_num_inference_steps_list: list[int] = field( default_factory=lambda: [20, 20, 20] ) @dataclass class HeliosDistilledConfig(HeliosT2VConfig): """Configuration for Helios-Distilled (Stage 1 + Stage 2 + Stage 3 DMD).""" is_enable_stage2: bool = True is_distilled: bool = True is_amplify_first_chunk: bool = True scheduler_type: str = "dmd" pyramid_num_inference_steps_list: list[int] = field( default_factory=lambda: [10, 10, 10] ) ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from collections.abc import Callable from dataclasses import dataclass, field from typing import TypedDict import torch from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig from sglang.multimodal_gen.configs.models.dits import HunyuanVideoConfig from sglang.multimodal_gen.configs.models.encoders import ( BaseEncoderOutput, CLIPTextConfig, LlamaConfig, ) from sglang.multimodal_gen.configs.models.vaes import HunyuanVAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, PipelineConfig, ) PROMPT_TEMPLATE_ENCODE_VIDEO = ( "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: " "1. The main content and theme of the video." "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects." "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects." "4. background environment, light, style and atmosphere." "5. camera angles, movements, and transitions used in the video:<|eot_id|>" "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>" ) class PromptTemplate(TypedDict): template: str crop_start: int prompt_template_video: PromptTemplate = { "template": PROMPT_TEMPLATE_ENCODE_VIDEO, "crop_start": 95, } def llama_preprocess_text(prompt: str) -> str: return prompt_template_video["template"].format(prompt) def llama_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.tensor: hidden_state_skip_layer = 2 assert outputs.hidden_states is not None hidden_states: tuple[torch.Tensor, ...] = outputs.hidden_states last_hidden_state: torch.tensor = hidden_states[-(hidden_state_skip_layer + 1)] crop_start = prompt_template_video.get("crop_start", -1) last_hidden_state = last_hidden_state[:, crop_start:] return last_hidden_state def clip_preprocess_text(prompt: str) -> str: return prompt def clip_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.tensor: pooler_output: torch.tensor = outputs.pooler_output return pooler_output @dataclass class HunyuanConfig(PipelineConfig): """Base configuration for HunYuan pipeline architecture.""" task_type: ModelTaskType = ModelTaskType.T2V # HunyuanConfig-specific parameters with defaults # DiT dit_config: DiTConfig = field(default_factory=HunyuanVideoConfig) # VAE vae_config: VAEConfig = field(default_factory=HunyuanVAEConfig) # Denoising stage embedded_cfg_scale: int = 6 flow_shift: int = 7 # Text encoding stage text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (LlamaConfig(), CLIPTextConfig()) ) preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( default_factory=lambda: (llama_preprocess_text, clip_preprocess_text) ) postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor], ...] = ( field(default_factory=lambda: (llama_postprocess_text, clip_postprocess_text)) ) # Precision for each component dit_precision: str = "bf16" vae_precision: str = "fp16" text_encoder_precisions: tuple[str, ...] = field( default_factory=lambda: ("fp16", "fp16") ) def __post_init__(self): self.vae_config.load_encoder = False self.vae_config.load_decoder = True @dataclass class FastHunyuanConfig(HunyuanConfig): """Configuration specifically optimized for FastHunyuan weights.""" # Override HunyuanConfig defaults flow_shift: int = 17 # No need to re-specify guidance_scale or embedded_cfg_scale as they # already have the desired values from HunyuanConfig ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan3d.py ================================================ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from typing import Optional from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConfig from sglang.multimodal_gen.configs.models.vaes.hunyuan3d import Hunyuan3DVAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, PipelineConfig, ) @dataclass class Hunyuan3D2PipelineConfig(PipelineConfig): """Pipeline configuration for Hunyuan3D image-to-mesh generation.""" task_type: ModelTaskType = ModelTaskType.I2M # Subfolder paths shape_subfolder: str = "hunyuan3d-dit-v2-0" paint_subfolder: str = "hunyuan3d-paint-v2-0" delight_subfolder: str = "hunyuan3d-delight-v2-0" # DiT configuration dit_config: DiTConfig = field(default_factory=Hunyuan3DDiTConfig) dit_precision: str = "fp16" # VAE configuration vae_config: VAEConfig = field(default_factory=Hunyuan3DVAEConfig) vae_precision: str = "fp32" # Shape model configuration shape_model_path: Optional[str] = None shape_use_safetensors: bool = True shape_variant: Optional[str] = "fp16" shape_num_inference_steps: int = 50 guidance_scale: float = 5.0 shape_box_v: float = 1.01 shape_octree_resolution: int = 384 shape_mc_level: float = 0.0 shape_mc_algo: Optional[str] = "mc" shape_num_chunks: int = 8000 shape_output_type: str = "trimesh" # Delight model configuration delight_enable: bool = True delight_prompt: str = "" delight_negative_prompt: str = "" delight_strength: float = 1.0 delight_num_inference_steps: int = 50 delight_guidance_scale: float = 1.0 delight_cfg_image: float = 1.5 # Paint model configuration paint_enable: bool = True paint_num_inference_steps: int = 30 paint_guidance_scale: float = 2.0 paint_resolution: int = 512 paint_render_size: int = 2048 paint_texture_size: int = 2048 paint_use_remesh: bool = True paint_save_glb: bool = True paint_turbo_mode: bool = False def __post_init__(self): self.vae_config.load_encoder = False self.vae_config.load_decoder = True def prepare_latent_shape(self, batch, batch_size, num_frames): latent_shape = self.vae_config.arch_config.latent_shape shape = (batch_size, *latent_shape) return shape ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py ================================================ import dataclasses from dataclasses import field from typing import Callable import torch from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2Config from sglang.multimodal_gen.configs.models.encoders import ( BaseEncoderOutput, EncoderConfig, ) from sglang.multimodal_gen.configs.models.encoders.gemma_3 import Gemma3Config from sglang.multimodal_gen.configs.models.vaes.ltx_audio import LTXAudioVAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, PipelineConfig, preprocess_text, ) from sglang.multimodal_gen.runtime.distributed import ( get_sp_parallel_rank, get_sp_world_size, sequence_model_parallel_all_gather, ) def pack_text_embeds( text_hidden_states: torch.Tensor, sequence_lengths: torch.Tensor, padding_side: str = "left", scale_factor: int = 8, eps: float = 1e-6, ) -> torch.Tensor: """ Packs and normalizes text encoder hidden states, respecting padding. Normalization is performed per-batch and per-layer in a masked fashion (only over non-padded positions). Args: text_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_dim, num_layers)`): Per-layer hidden_states from a text encoder (e.g. `Gemma3ForConditionalGeneration`). sequence_lengths (`torch.Tensor of shape `(batch_size,)`): The number of valid (non-padded) tokens for each batch instance. device: (`str` or `torch.device`, *optional*): torch device to place the resulting embeddings on padding_side: (`str`, *optional*, defaults to `"left"`): Whether the text tokenizer performs padding on the `"left"` or `"right"`. scale_factor (`int`, *optional*, defaults to `8`): Scaling factor to multiply the normalized hidden states by. eps (`float`, *optional*, defaults to `1e-6`): A small positive value for numerical stability when performing normalization. Returns: `torch.Tensor` of shape `(batch_size, seq_len, hidden_dim * num_layers)`: Normed and flattened text encoder hidden states. """ batch_size, seq_len, hidden_dim, num_layers = text_hidden_states.shape original_dtype = text_hidden_states.dtype device = text_hidden_states.device # Create padding mask token_indices = torch.arange(seq_len, device=device).unsqueeze(0) if padding_side == "right": mask = token_indices < sequence_lengths[:, None] elif padding_side == "left": start_indices = seq_len - sequence_lengths[:, None] mask = token_indices >= start_indices else: raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}") mask = mask[:, :, None, None] # [batch_size, seq_len, 1, 1] masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0) num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1) masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / ( num_valid_positions + eps ) x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin( dim=(1, 2), keepdim=True ) x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax( dim=(1, 2), keepdim=True ) normalized_hidden_states = (text_hidden_states - masked_mean) / ( x_max - x_min + eps ) normalized_hidden_states = normalized_hidden_states * scale_factor normalized_hidden_states = normalized_hidden_states.flatten(2) mask_flat = mask.squeeze(-1).expand(-1, -1, hidden_dim * num_layers) normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0) normalized_hidden_states = normalized_hidden_states.to(dtype=original_dtype) return normalized_hidden_states def _gemma_postprocess_func( outputs: BaseEncoderOutput, text_inputs: dict ) -> torch.Tensor: # LTX-2 requires all hidden states concatenated for the connector if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: # outputs.hidden_states is a tuple of tensors # We need to stack them along the last dimension and pack them hidden_states = torch.stack(outputs.hidden_states, dim=-1) attention_mask = text_inputs["attention_mask"] sequence_lengths = attention_mask.sum(dim=-1) # Assuming left padding for Gemma as per Diffusers return pack_text_embeds(hidden_states, sequence_lengths, padding_side="left") else: raise AttributeError( "Unsupported text encoder output: expected `hidden_states`." ) @dataclasses.dataclass class LTX2PipelineConfig(PipelineConfig): """Configuration for LTX-Video pipeline.""" task_type: ModelTaskType = ModelTaskType.TI2V skip_input_image_preprocess: bool = True dit_config: LTX2Config = field(default_factory=LTX2Config) # Model architecture in_channels: int = 128 out_channels: int = 128 patch_size: int = 1 patch_size_t: int = 1 # Audio VAE configuration audio_vae_config: LTXAudioVAEConfig = field(default_factory=LTXAudioVAEConfig) audio_vae_precision: str = "fp32" audio_vae_temporal_compression_ratio: int = 4 audio_vae_mel_compression_ratio: int = 4 @property def vae_scale_factor(self): return getattr(self.vae_config.arch_config, "spatial_compression_ratio", 32) @property def vae_temporal_compression(self): return getattr(self.vae_config.arch_config, "temporal_compression_ratio", 8) def prepare_latent_shape(self, batch, batch_size, num_frames): """Return packed latent shape [B, seq, C] directly.""" height = batch.height // self.vae_scale_factor width = batch.width // self.vae_scale_factor post_patch_num_frames = num_frames // self.patch_size_t post_patch_height = height // self.patch_size post_patch_width = width // self.patch_size seq_len = post_patch_num_frames * post_patch_height * post_patch_width num_channels = ( self.in_channels * self.patch_size_t * self.patch_size * self.patch_size ) shape = (batch_size, seq_len, num_channels) return shape def prepare_audio_latent_shape(self, batch, batch_size, num_frames): # Adapted from diffusers pipeline prepare_audio_latents duration_s = num_frames / batch.fps sample_rate = self.audio_vae_config.arch_config.sample_rate hop_length = self.audio_vae_config.arch_config.mel_hop_length temporal_compression = self.audio_vae_temporal_compression_ratio latents_per_second = ( float(sample_rate) / float(hop_length) / float(temporal_compression) ) latent_length = round(duration_s * latents_per_second) num_mel_bins = self.audio_vae_config.arch_config.mel_bins mel_compression_ratio = self.audio_vae_mel_compression_ratio latent_mel_bins = num_mel_bins // mel_compression_ratio # Default to 8 num_channels_latents = self.audio_vae_config.arch_config.latent_channels shape = (batch_size, latent_length, num_channels_latents * latent_mel_bins) return shape # Text encoding stage (Gemma) # LTX-2 needs separate contexts for video/audio streams. We model this as # two logical encoders sharing the same underlying `text_encoder` module. text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (Gemma3Config(),) ) text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",)) text_encoder_extra_args: list[dict] = field(default_factory=lambda: [{}]) preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( default_factory=lambda: (preprocess_text,) ) postprocess_text_funcs: tuple[ Callable[[BaseEncoderOutput, dict], torch.Tensor], ... ] = field(default_factory=lambda: (_gemma_postprocess_func,)) def prepare_sigmas(self, sigmas, num_inference_steps): if sigmas is None: steps = int(num_inference_steps) if steps <= 0: raise ValueError(f"num_inference_steps must be positive, got {steps}") return [1.0 - i / steps for i in range(steps)] return sigmas def tokenize_prompt(self, prompt: list[str], tokenizer, tok_kwargs) -> dict: # Adapted from diffusers_pipeline.py _get_gemma_prompt_embeds # But we only need tokenization here, the embedding happens in TextEncodingStage # Gemma expects left padding for chat-style prompts tokenizer.padding_side = "left" if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token max_sequence_length = tok_kwargs.get( "max_length", 1024 ) # Default from diffusers pipeline text_inputs = tokenizer( prompt, padding="max_length", max_length=max_sequence_length, truncation=True, return_tensors="pt", ) return text_inputs def maybe_pack_latents(self, latents, batch_size, batch): # If already packed (3D shape [B, seq, C]), skip packing if latents.dim() == 3: return latents # Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p]. # The patch dimensions are then permuted and collapsed into the channel dimension of shape: # [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor). # dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features batch_size, num_channels, num_frames, height, width = latents.shape post_patch_num_frames = num_frames // self.patch_size_t post_patch_height = height // self.patch_size post_patch_width = width // self.patch_size latents = latents.reshape( batch_size, -1, post_patch_num_frames, self.patch_size_t, post_patch_height, self.patch_size, post_patch_width, self.patch_size, ) latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) return latents def _infer_video_latent_frames_and_tokens_per_frame( self, batch, seq_len: int ) -> tuple[int, int]: """Infer latent-frame count and tokens-per-frame for packed token latents [B, S, D]. Notes: - This assumes `patch_size_t == 1` (no temporal patching). - Tokens are ordered as (frame, height, width) after packing. """ if int(self.patch_size_t) != 1: raise ValueError( "LTX-2 SP time-sharding for packed token latents currently requires " f"{self.patch_size_t=}. (Expected 1)" ) if int(seq_len) <= 0: raise ValueError(f"Expected {seq_len=} > 0 for packed token latents.") if int(self.vae_scale_factor) <= 0: raise ValueError(f"Invalid {self.vae_scale_factor=}. Must be > 0.") if int(self.patch_size) <= 0: raise ValueError(f"Invalid {self.patch_size=}. Must be > 0.") latent_height = int(batch.height) // int(self.vae_scale_factor) latent_width = int(batch.width) // int(self.vae_scale_factor) if latent_height <= 0 or latent_width <= 0: raise ValueError( "Invalid latent H/W computed from batch.height/width: " f"{batch.height=} {batch.width=} {self.vae_scale_factor=}" ) if (latent_height % int(self.patch_size)) != 0 or ( latent_width % int(self.patch_size) ) != 0: raise ValueError( "Invalid spatial patching for packed token latents. Expected latent H/W " "to be divisible by patch_size, got " f"{latent_height=} {latent_width=} {self.patch_size=}." ) post_patch_h = latent_height // int(self.patch_size) post_patch_w = latent_width // int(self.patch_size) tokens_per_frame = int(post_patch_h) * int(post_patch_w) if tokens_per_frame <= 0: raise ValueError( f"Invalid tokens_per_frame={tokens_per_frame} from " f"{latent_height=} {latent_width=} {self.patch_size=}" ) if int(seq_len) % int(tokens_per_frame) != 0: raise ValueError( f"LTX-2 token latents seq_len={seq_len} is not divisible by " f"tokens_per_frame={tokens_per_frame}. Cannot time-shard for SP." ) latent_num_frames = int(seq_len) // int(tokens_per_frame) return int(latent_num_frames), int(tokens_per_frame) def shard_latents_for_sp(self, batch, latents): """Shard LTX-2 packed token latents across SP ranks by latent time (frame) dimension.""" sp_world_size = get_sp_world_size() if sp_world_size <= 1: return latents, False # Default behavior for 5D latents. if isinstance(latents, torch.Tensor) and latents.ndim == 5: return super().shard_latents_for_sp(batch, latents) # LTX-2 packed token latents [B, S, D] if not (isinstance(latents, torch.Tensor) and latents.ndim == 3): return latents, False sp_rank = get_sp_parallel_rank() seq_len = int(latents.shape[1]) latent_frames, tokens_per_frame = ( self._infer_video_latent_frames_and_tokens_per_frame(batch, seq_len) ) # Pad whole frames so `latent_frames` is divisible by `sp_world_size`. pad_frames = (sp_world_size - (latent_frames % sp_world_size)) % sp_world_size if pad_frames: pad_tokens = int(pad_frames) * int(tokens_per_frame) pad = torch.zeros( (latents.shape[0], pad_tokens, latents.shape[2]), device=latents.device, dtype=latents.dtype, ) latents = torch.cat([latents, pad], dim=1) latent_frames = int(latent_frames) + int(pad_frames) local_frames = int(latent_frames) // int(sp_world_size) start_frame = int(sp_rank) * int(local_frames) start = int(start_frame) * int(tokens_per_frame) end = int(start) + int(local_frames) * int(tokens_per_frame) latents = latents[:, start:end, :] # Store SP metadata for denoising (TI2V gating) and model-side RoPE shift. batch.sp_video_latent_num_frames = int(local_frames) batch.sp_video_start_frame = int(start_frame) batch.sp_video_tokens_per_frame = int(tokens_per_frame) return latents, True def gather_latents_for_sp(self, latents): """Gather latents after SP. For packed token latents [B, S_local, D], gather on dim=1.""" if get_sp_world_size() <= 1: return latents if isinstance(latents, torch.Tensor) and latents.ndim == 3: return sequence_model_parallel_all_gather(latents.contiguous(), dim=1) return super().gather_latents_for_sp(latents) def maybe_pack_audio_latents(self, latents, batch_size, batch): # If already packed (3D shape [B, T, C*F]), skip packing if latents.dim() == 3: return latents # Audio latents shape: [B, C, L, M], where L is the latent audio length and M is the number of mel bins # We need to pack them if patch_size/patch_size_t are defined for audio (not standard DiT patch size) # So for LTX-2 (unless we change patch sizes), we just do: latents = latents.transpose(1, 2).flatten( 2, 3 ) # [B, C, L, M] --> [B, L, C * M] return latents def get_pos_prompt_embeds(self, batch): # LTX-2 returns multiple prompt embed tensors (video/audio contexts). return ( batch.prompt_embeds[0] if isinstance(batch.prompt_embeds, list) else batch.prompt_embeds ) def get_neg_prompt_embeds(self, batch): return ( batch.negative_prompt_embeds[0] if isinstance(batch.negative_prompt_embeds, list) else batch.negative_prompt_embeds ) def get_decode_scale_and_shift(self, device, dtype, vae): latents_mean = getattr(vae, "latents_mean", None) latents_std = getattr(vae, "latents_std", None) scaling_factor = ( getattr(getattr(vae, "config", None), "scaling_factor", None) or getattr(vae, "scaling_factor", None) or getattr(self.vae_config.arch_config, "scaling_factor", None) or 1.0 ) if isinstance(scaling_factor, (int, float)) and float(scaling_factor) == 0.0: scaling_factor = 1.0 if isinstance(latents_mean, torch.Tensor) and isinstance( latents_std, torch.Tensor ): latents_mean = latents_mean.to(device=device, dtype=dtype).view( 1, -1, 1, 1, 1 ) latents_std = latents_std.to(device=device, dtype=dtype).view( 1, -1, 1, 1, 1 ) sf = torch.tensor(float(scaling_factor), device=device, dtype=dtype).view( 1, 1, 1, 1, 1 ) return sf / latents_std, latents_mean sf = torch.tensor(float(scaling_factor), device=device, dtype=dtype).view( 1, 1, 1, 1, 1 ) return sf, None @staticmethod def _unpack_latents( latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1, ) -> torch.Tensor: # Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions) # are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of # what happens in the `_pack_latents` method. batch_size = latents.size(0) latents = latents.reshape( batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size, ) latents = ( latents.permute(0, 4, 1, 5, 2, 6, 3, 7) .flatten(6, 7) .flatten(4, 5) .flatten(2, 3) ) return latents @staticmethod def _denormalize_latents( latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0, ) -> torch.Tensor: # Denormalize latents across the channel dimension [B, C, F, H, W] latents_mean = latents_mean.view(1, -1, 1, 1, 1).to( latents.device, latents.dtype ) latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) latents = latents * latents_std / scaling_factor + latents_mean return latents @staticmethod def _denormalize_audio_latents( latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor ): latents_mean = latents_mean.to(latents.device, latents.dtype) latents_std = latents_std.to(latents.device, latents.dtype) return (latents * latents_std) + latents_mean @staticmethod def _unpack_audio_latents( latents: torch.Tensor, latent_length: int, num_mel_bins: int, patch_size: int | None = None, patch_size_t: int | None = None, ) -> torch.Tensor: # Unpacks an audio patch sequence of shape [B, S, D] into a latent spectrogram tensor of shape [B, C, L, M], # where L is the latent audio length and M is the number of mel bins. if patch_size is not None and patch_size_t is not None: batch_size = latents.size(0) latents = latents.reshape( batch_size, latent_length, num_mel_bins, -1, patch_size_t, patch_size ) latents = latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) else: # Assume [B, S, D] = [B, L, C * M], which implies that patch_size = M and patch_size_t = 1. latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2) return latents def _unpad_and_unpack_latents(self, latents, audio_latents, batch, vae, audio_vae): # Calculate latent dimensions # Assuming batch has height, width, num_frames height = batch.height width = batch.width num_frames = batch.num_frames # Get compression ratios # Default LTX-2 values if not present in config vae_spatial_compression_ratio = getattr( self.vae_config.arch_config, "spatial_compression_ratio", 32 ) vae_temporal_compression_ratio = getattr( self.vae_config.arch_config, "temporal_compression_ratio", 8 ) latent_height = height // vae_spatial_compression_ratio latent_width = width // vae_spatial_compression_ratio latent_num_frames = (num_frames - 1) // vae_temporal_compression_ratio + 1 latents = self._unpack_latents( latents, latent_num_frames, latent_height, latent_width, self.patch_size, self.patch_size_t, ) sample_rate = self.audio_vae_config.arch_config.sample_rate hop_length = self.audio_vae_config.arch_config.mel_hop_length temporal_compression = self.audio_vae_temporal_compression_ratio duration_s = num_frames / batch.fps latents_per_second = ( float(sample_rate) / float(hop_length) / float(temporal_compression) ) audio_num_frames = round(duration_s * latents_per_second) num_mel_bins = self.audio_vae_config.arch_config.mel_bins mel_compression_ratio = self.audio_vae_mel_compression_ratio latent_mel_bins = num_mel_bins // mel_compression_ratio audio_latents_mean = getattr(audio_vae, "latents_mean", None) audio_latents_std = getattr(audio_vae, "latents_std", None) if ( isinstance(audio_latents_mean, torch.Tensor) and isinstance(audio_latents_std, torch.Tensor) and audio_latents_mean.numel() == audio_latents_std.numel() ): audio_latents_mean = audio_latents_mean.to( device=audio_latents.device, dtype=audio_latents.dtype ) audio_latents_std = audio_latents_std.to( device=audio_latents.device, dtype=audio_latents.dtype ) if audio_latents.ndim == 3: if audio_latents.shape[-1] != audio_latents_mean.numel(): raise ValueError( f"audio_latents last dim {audio_latents.shape[-1]} " f"does not match audio_vae stats {audio_latents_mean.numel()}" ) audio_latents = audio_latents * audio_latents_std.view( 1, 1, -1 ) + audio_latents_mean.view(1, 1, -1) elif audio_latents.ndim == 2: if audio_latents.shape[-1] != audio_latents_mean.numel(): raise ValueError( f"audio_latents last dim {audio_latents.shape[-1]} " f"does not match audio_vae stats {audio_latents_mean.numel()}" ) audio_latents = audio_latents * audio_latents_std.view( 1, -1 ) + audio_latents_mean.view(1, -1) else: audio_latents = audio_latents * audio_latents_std + audio_latents_mean audio_latents = self._unpack_audio_latents( audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins ) return latents, audio_latents @dataclasses.dataclass class LTX2I2VPipelineConfig(LTX2PipelineConfig): task_type: ModelTaskType = ModelTaskType.TI2V ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/mova.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ MOVA pipeline configuration. """ from dataclasses import dataclass, field import numpy as np import torch import torch.nn.functional as F from PIL import Image from sglang.multimodal_gen.configs.models.dits import MOVAAudioConfig, MOVAVideoConfig from sglang.multimodal_gen.configs.models.encoders import T5Config from sglang.multimodal_gen.configs.models.vaes import DacVAEConfig, WanVAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, PipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.wan import t5_postprocess_text from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @dataclass class MOVAPipelineConfig(PipelineConfig): """Configuration for MOVA (text+image -> video+audio) pipelines.""" task_type: ModelTaskType = ModelTaskType.I2V # Model configs dit_config: MOVAVideoConfig = field(default_factory=MOVAVideoConfig) audio_dit_config: MOVAAudioConfig = field(default_factory=MOVAAudioConfig) # Video VAE (Wan) + Audio VAE (DAC) vae_config: WanVAEConfig = field(default_factory=WanVAEConfig) audio_vae_config: DacVAEConfig = field(default_factory=DacVAEConfig) audio_vae_precision: str = "fp32" # Text encoder (UMT5 compatible) text_encoder_configs: tuple = field(default_factory=lambda: (T5Config(),)) postprocess_text_funcs: tuple = field( default_factory=lambda: (t5_postprocess_text,) ) # MOVA specific audio_vae_type: str = "dac" boundary_ratio: float | None = 0.9 # temporal alignment: MOVA expects (num_frames - 1) % 4 == 0 time_division_factor: int = 4 time_division_remainder: int = 1 def _center_crop_and_resize( self, image: torch.Tensor | Image.Image, target_height: int, target_width: int ) -> torch.Tensor | Image.Image: if not isinstance(image, (Image.Image, torch.Tensor)): raise TypeError(f"Unsupported image type: {type(image)}") if isinstance(image, Image.Image): image = torch.from_numpy(np.array(image)) if image.ndim == 2: image = image[..., None] if not image.dtype.is_floating_point: image = image.to(torch.float32).div(255.0) if image.ndim == 3: if image.shape[0] in (1, 3, 4) and image.shape[-1] not in (1, 3, 4): image = image.unsqueeze(0) else: image = image.permute(2, 0, 1).unsqueeze(0) elif image.ndim == 4: if image.shape[1] not in (1, 3, 4) and image.shape[-1] in (1, 3, 4): image = image.permute(0, 3, 1, 2) image_height, image_width = image.shape[-2], image.shape[-1] if image_height == target_height and image_width == target_width: return image logger.info( "Center cropping and resizing image to %dx%d", target_width, target_height ) if image_height * target_width < image_width * target_height: cropped_width = (image_height * target_width) // target_height left = (image_width - cropped_width) // 2 image = image[..., :, left : left + cropped_width] else: cropped_height = (image_width * target_height) // target_width top = (image_height - cropped_height) // 2 image = image[..., top : top + cropped_height, :] image = F.interpolate( image, size=(target_height, target_width), mode="bilinear", align_corners=False, antialias=True, ) return image def adjust_num_frames(self, num_frames: int) -> int: if num_frames is None: return num_frames if num_frames % self.time_division_factor != self.time_division_remainder: adjusted = ( (num_frames + self.time_division_factor - 1) // self.time_division_factor * self.time_division_factor + self.time_division_remainder ) logger.warning( "`num_frames` (%s) is not compatible with MOVA temporal constraints. " "Rounding to %s.", num_frames, adjusted, ) return adjusted return num_frames def preprocess_condition_image( self, image, target_width, target_height, _vae_image_processor ): image = self._center_crop_and_resize(image, target_height, target_width) return image, (target_width, target_height) def prepare_latent_shape(self, batch, batch_size, num_frames): spatial = self.vae_config.arch_config.spatial_compression_ratio length = (num_frames - 1) // self.time_division_factor + 1 shape = ( batch_size, self.dit_config.arch_config.out_dim, length, batch.height // spatial, batch.width // spatial, ) return shape def prepare_audio_latent_shape(self, batch_size, num_samples, audio_vae): latent_T = (num_samples + audio_vae.hop_length - 1) // audio_vae.hop_length return (batch_size, audio_vae.latent_dim, latent_T) def normalize_video_latents(self, latents: torch.Tensor, video_vae) -> torch.Tensor: latents_mean = getattr(video_vae.config, "latents_mean", None) latents_std = getattr(video_vae.config, "latents_std", None) if latents_mean is None or latents_std is None: return latents mean = torch.tensor( latents_mean, device=latents.device, dtype=latents.dtype ).view(1, video_vae.config.z_dim, 1, 1, 1) inv_std = ( 1.0 / torch.tensor(latents_std, device=latents.device, dtype=latents.dtype) ).view(1, video_vae.config.z_dim, 1, 1, 1) return (latents - mean) * inv_std def denormalize_video_latents( self, latents: torch.Tensor, video_vae ) -> torch.Tensor: latents_mean = getattr(video_vae.config, "latents_mean", None) latents_std = getattr(video_vae.config, "latents_std", None) if latents_mean is None or latents_std is None: return latents mean = torch.tensor( latents_mean, device=latents.device, dtype=latents.dtype ).view(1, video_vae.config.z_dim, 1, 1, 1) std = torch.tensor( latents_std, device=latents.device, dtype=latents.dtype ).view(1, video_vae.config.z_dim, 1, 1, 1) return latents * std + mean @dataclass class MOVA360PConfig(MOVAPipelineConfig): """Configuration for MOVA 360P (text+image -> video+audio) pipelines.""" max_area: int = 352 * 640 @dataclass class MOVA720PConfig(MOVAPipelineConfig): """Configuration for MOVA 720P (text+image -> video+audio) pipelines.""" max_area: int = 720 * 1280 ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from dataclasses import dataclass, field from typing import Callable import torch from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig from sglang.multimodal_gen.configs.models.dits.qwenimage import ( QwenImageDitConfig, QwenImageEditPlus_2511_DitConfig, ) from sglang.multimodal_gen.configs.models.encoders.qwen_image import Qwen2_5VLConfig from sglang.multimodal_gen.configs.models.vaes.qwenimage import QwenImageVAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ImagePipelineConfig, ModelTaskType, maybe_unpad_latents, shard_rotary_emb_for_sp, ) from sglang.multimodal_gen.runtime.models.vision_utils import resize from sglang.multimodal_gen.utils import calculate_dimensions def _extract_masked_hidden(hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() valid_lengths = bool_mask.sum(dim=1) selected = hidden_states[bool_mask] split_result = torch.split(selected, valid_lengths.tolist(), dim=0) return split_result def qwen_image_preprocess_text(prompt): prompt_template_encode = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" template = prompt_template_encode txt = template.format(prompt) return txt def qwen_image_postprocess_text(outputs, _text_inputs, drop_idx=34): # squeeze the batch dim hidden_states = outputs.hidden_states[-1] split_hidden_states = _extract_masked_hidden( hidden_states, _text_inputs.attention_mask ) split_hidden_states = [e[drop_idx:] for e in split_hidden_states] max_seq_len = max([e.size(0) for e in split_hidden_states]) prompt_embeds = torch.stack( [ torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states ] ) return prompt_embeds def _normalize_prompt_list(prompt): return [prompt] if isinstance(prompt, str) else prompt def _normalize_image_list(images): if images is None: return [] return images if isinstance(images, list) else [images] def _build_qwen_edit_image_prompt(num_images: int) -> str: img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>" return "".join(img_prompt_template.format(i + 1) for i in range(num_images)) def _resolve_qwen_edit_per_prompt_images(prompt_list, image_list): if len(prompt_list) <= 1: return [image_list] if len(image_list) <= 1: return [list(image_list) for _ in prompt_list] if len(image_list) != len(prompt_list): raise ValueError( "QwenImageEditPlus expects either one shared condition image or " "the same number of condition images and prompts." ) return [[image] for image in image_list] # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._pack_latents def _pack_latents(latents, batch_size, num_channels_latents, height, width): latents = latents.view( batch_size, num_channels_latents, height // 2, 2, width // 2, 2 ) latents = latents.permute(0, 2, 4, 1, 3, 5) latents = latents.reshape( batch_size, (height // 2) * (width // 2), num_channels_latents * 4 ) return latents @dataclass class QwenImagePipelineConfig(ImagePipelineConfig): """Configuration for the QwenImage pipeline.""" should_use_guidance: bool = False task_type: ModelTaskType = ModelTaskType.T2I vae_tiling: bool = False vae_sp: bool = False dit_config: DiTConfig = field(default_factory=QwenImageDitConfig) # VAE vae_config: VAEConfig = field(default_factory=QwenImageVAEConfig) enable_autocast: bool = False # Text encoding stage text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (Qwen2_5VLConfig(),) ) text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",)) preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( default_factory=lambda: (qwen_image_preprocess_text,) ) postprocess_text_funcs: tuple[Callable[[str], str], ...] = field( default_factory=lambda: (qwen_image_postprocess_text,) ) text_encoder_extra_args: list[dict] = field( default_factory=lambda: [ dict( padding=True, truncation=True, ), None, ] ) def prepare_sigmas(self, sigmas, num_inference_steps): return self._prepare_sigmas(sigmas, num_inference_steps) def prepare_image_processor_kwargs(self, batch, neg=False): prompt = batch.prompt if not neg else batch.negative_prompt if prompt: prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n" txt = prompt_template_encode.format(batch.prompt) return dict(text=[txt], padding=True) else: return {} def get_vae_scale_factor(self): return self.vae_config.arch_config.vae_scale_factor def prepare_latent_shape(self, batch, batch_size, num_frames): vae_scale_factor = self.vae_config.arch_config.vae_scale_factor height = 2 * (batch.height // (vae_scale_factor * 2)) width = 2 * (batch.width // (vae_scale_factor * 2)) num_channels_latents = self.dit_config.arch_config.in_channels // 4 shape = (batch_size, 1, num_channels_latents, height, width) return shape def maybe_pack_latents(self, latents, batch_size, batch): height = 2 * ( batch.height // (self.vae_config.arch_config.vae_scale_factor * 2) ) width = 2 * (batch.width // (self.vae_config.arch_config.vae_scale_factor * 2)) num_channels_latents = self.dit_config.arch_config.in_channels // 4 # pack latents return _pack_latents(latents, batch_size, num_channels_latents, height, width) def get_decode_scale_and_shift(self, device, dtype, vae): vae_arch_config = self.vae_config.arch_config scaling_factor = 1.0 / torch.tensor( vae_arch_config.latents_std, device=device ).view(1, vae_arch_config.z_dim, 1, 1, 1).to(device, dtype) shift_factor = ( torch.tensor(vae_arch_config.latents_mean) .view(1, vae_arch_config.z_dim, 1, 1, 1) .to(device, dtype) ) return scaling_factor, shift_factor @staticmethod def get_freqs_cis(img_shapes, txt_seq_lens, rotary_emb, device, dtype): # img_shapes: for global entire image img_freqs, txt_freqs = rotary_emb(img_shapes, txt_seq_lens, device=device) # flashinfer RoPE expects a float32 cos/sin cache concatenated on the last dim img_cos_half = img_freqs.real.to(dtype=torch.float32).contiguous() img_sin_half = img_freqs.imag.to(dtype=torch.float32).contiguous() txt_cos_half = txt_freqs.real.to(dtype=torch.float32).contiguous() txt_sin_half = txt_freqs.imag.to(dtype=torch.float32).contiguous() img_cos_sin_cache = torch.cat([img_cos_half, img_sin_half], dim=-1) txt_cos_sin_cache = torch.cat([txt_cos_half, txt_sin_half], dim=-1) return img_cos_sin_cache, txt_cos_sin_cache def _prepare_cond_kwargs(self, batch, prompt_embeds, rotary_emb, device, dtype): batch_size = prompt_embeds[0].shape[0] height = batch.height width = batch.width vae_scale_factor = self.vae_config.arch_config.vae_scale_factor img_shapes = [ [ ( 1, height // vae_scale_factor // 2, width // vae_scale_factor // 2, ) ] ] * batch_size txt_seq_lens = [prompt_embeds[0].shape[1]] if rotary_emb is None: return { "img_shapes": img_shapes, "txt_seq_lens": txt_seq_lens, "freqs_cis": None, } freqs_cis = self.get_freqs_cis( img_shapes, txt_seq_lens, rotary_emb, device, dtype ) img_cache, txt_cache = freqs_cis img_cache = shard_rotary_emb_for_sp(img_cache) return { "txt_seq_lens": txt_seq_lens, "freqs_cis": (img_cache, txt_cache), "img_shapes": img_shapes, } def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): return self._prepare_cond_kwargs( batch, batch.prompt_embeds, rotary_emb, device, dtype ) def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): return self._prepare_cond_kwargs( batch, batch.negative_prompt_embeds, rotary_emb, device, dtype ) def post_denoising_loop(self, latents, batch): # unpack latents for qwen-image ( latents, batch_size, channels, height, width, ) = self._unpad_and_unpack_latents(latents, batch) latents = latents.reshape(batch_size, channels // (2 * 2), 1, height, width) return latents @dataclass class QwenImageEditPipelineConfig(QwenImagePipelineConfig): """Configuration for the QwenImageEdit pipeline.""" task_type: ModelTaskType = ModelTaskType.I2I def _prepare_edit_cond_kwargs( self, batch, prompt_embeds, rotary_emb, device, dtype ): batch_size = batch.latents.shape[0] assert batch_size == 1 height = batch.height width = batch.width image_size = batch.original_condition_image_size edit_width, edit_height, _ = calculate_dimensions( 1024 * 1024, image_size[0] / image_size[1] ) vae_scale_factor = self.get_vae_scale_factor() img_shapes = [ [ ( 1, height // vae_scale_factor // 2, width // vae_scale_factor // 2, ), ( 1, edit_height // vae_scale_factor // 2, edit_width // vae_scale_factor // 2, ), ], ] * batch_size txt_seq_lens = [prompt_embeds[0].shape[1]] if rotary_emb is None: return { "img_shapes": img_shapes, "txt_seq_lens": txt_seq_lens, "freqs_cis": None, } freqs_cis = QwenImagePipelineConfig.get_freqs_cis( img_shapes, txt_seq_lens, rotary_emb, device, dtype ) # perform sp shard on noisy image tokens noisy_img_seq_len = ( 1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2) ) img_cache, txt_cache = freqs_cis noisy_img_cache = shard_rotary_emb_for_sp(img_cache[:noisy_img_seq_len, :]) img_cache = torch.cat( [noisy_img_cache, img_cache[noisy_img_seq_len:, :]], dim=0 ).to(device=device) return { "txt_seq_lens": txt_seq_lens, "freqs_cis": (img_cache, txt_cache), "img_shapes": img_shapes, } def preprocess_condition_image( self, image, target_width, target_height, _vae_image_processor ): return resize(image, target_height, target_width, resize_mode="default"), ( target_width, target_height, ) def postprocess_image_latent(self, latent_condition, batch): batch_size = batch.batch_size if batch_size > latent_condition.shape[0]: if batch_size % latent_condition.shape[0] == 0: # expand init_latents for batch_size additional_image_per_prompt = batch_size // latent_condition.shape[0] image_latents = latent_condition.repeat( additional_image_per_prompt, 1, 1, 1 ) else: raise ValueError( f"Cannot duplicate `image` of batch size {latent_condition.shape[0]} to {batch_size} text prompts." ) else: image_latents = latent_condition image_latent_height, image_latent_width = image_latents.shape[3:] num_channels_latents = self.dit_config.arch_config.in_channels // 4 image_latents = _pack_latents( image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width, ) return image_latents def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): return self._prepare_edit_cond_kwargs( batch, batch.prompt_embeds, rotary_emb, device, dtype ) def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): return self._prepare_edit_cond_kwargs( batch, batch.negative_prompt_embeds, rotary_emb, device, dtype ) def calculate_condition_image_size(self, image, width, height) -> tuple[int, int]: calculated_width, calculated_height, _ = calculate_dimensions( 1024 * 1024, width / height ) return calculated_width, calculated_height def slice_noise_pred(self, noise, latents): # remove noise over input image noise = noise[:, : latents.size(1)] return noise CONDITION_IMAGE_SIZE = 384 * 384 VAE_IMAGE_SIZE = 1024 * 1024 @dataclass class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig): task_type: ModelTaskType = ModelTaskType.I2I def _get_condition_image_sizes(self, batch) -> list[tuple[int, int]]: image = batch.condition_image if not isinstance(image, list): image = [image] condition_image_sizes = [] for img in image: image_width, image_height = img.size edit_width, edit_height, _ = calculate_dimensions( VAE_IMAGE_SIZE, image_width / image_height ) condition_image_sizes.append((edit_width, edit_height)) return condition_image_sizes def prepare_image_processor_kwargs(self, batch, neg=False) -> dict: prompt = batch.prompt if not neg else batch.negative_prompt if not prompt: return {} prompt_list = _normalize_prompt_list(prompt) image_list = _normalize_image_list(batch.condition_image) per_prompt_images = _resolve_qwen_edit_per_prompt_images( prompt_list, image_list ) prompt_template_encode = ( "<|im_start|>system\nDescribe the key features of the input image " "(color, shape, size, texture, objects, background), then explain how " "the user's text instruction should alter or modify the image. Generate " "a new image that meets the user's requirements while maintaining " "consistency with the original input where appropriate.<|im_end|>\n" "<|im_start|>user\n{}<|im_end|>\n" "<|im_start|>assistant\n" ) txt = [ prompt_template_encode.format( _build_qwen_edit_image_prompt(len(prompt_images)) + prompt_text ) for prompt_text, prompt_images in zip(prompt_list, per_prompt_images) ] return dict(text=txt, padding=True, per_prompt_images=per_prompt_images) def prepare_calculated_size(self, image): return self.calculate_vae_image_size(image, image.width, image.height) def resize_condition_image(self, images, target_width, target_height): if not isinstance(images, list): images = [images] new_images = [] for img, width, height in zip(images, target_width, target_height): new_images.append(resize(img, height, width, resize_mode="default")) return new_images def calculate_condition_image_size(self, image, width, height) -> tuple[int, int]: calculated_width, calculated_height, _ = calculate_dimensions( CONDITION_IMAGE_SIZE, width / height ) return calculated_width, calculated_height def calculate_vae_image_size(self, image, width, height) -> tuple[int, int]: calculated_width, calculated_height, _ = calculate_dimensions( VAE_IMAGE_SIZE, width / height ) return calculated_width, calculated_height def preprocess_vae_image(self, batch, vae_image_processor): if not isinstance(batch.condition_image, list): batch.condition_image = [batch.condition_image] new_images = [] vae_image_sizes = [] for img in batch.condition_image: width, height = self.calculate_vae_image_size(img, img.width, img.height) new_images.append(vae_image_processor.preprocess(img, height, width)) vae_image_sizes.append((width, height)) batch.vae_image = new_images batch.vae_image_sizes = vae_image_sizes return batch def _prepare_edit_cond_kwargs( self, batch, prompt_embeds, rotary_emb, device, dtype ): batch_size = batch.latents.shape[0] assert batch_size == 1 height = batch.height width = batch.width vae_scale_factor = self.get_vae_scale_factor() img_shapes = [ [ (1, height // vae_scale_factor // 2, width // vae_scale_factor // 2), *[ ( 1, vae_height // vae_scale_factor // 2, vae_width // vae_scale_factor // 2, ) for vae_width, vae_height in batch.vae_image_sizes ], ], ] * batch_size txt_seq_lens = [prompt_embeds[0].shape[1]] freqs_cis = QwenImageEditPlusPipelineConfig.get_freqs_cis( img_shapes, txt_seq_lens, rotary_emb, device, dtype ) # perform sp shard on noisy image tokens noisy_img_seq_len = ( 1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2) ) if isinstance(freqs_cis[0], torch.Tensor) and freqs_cis[0].dim() == 2: img_cache, txt_cache = freqs_cis noisy_img_cache = shard_rotary_emb_for_sp(img_cache[:noisy_img_seq_len, :]) img_cache = torch.cat( [noisy_img_cache, img_cache[noisy_img_seq_len:, :]], dim=0 ).to(device=device) return { "txt_seq_lens": txt_seq_lens, "freqs_cis": (img_cache, txt_cache), "img_shapes": img_shapes, } (img_cos, img_sin), (txt_cos, txt_sin) = freqs_cis noisy_img_cos = shard_rotary_emb_for_sp(img_cos[:noisy_img_seq_len, :]) noisy_img_sin = shard_rotary_emb_for_sp(img_sin[:noisy_img_seq_len, :]) # concat back the img_cos for input image (since it is not sp-shared later) img_cos = torch.cat([noisy_img_cos, img_cos[noisy_img_seq_len:, :]], dim=0).to( device=device ) img_sin = torch.cat([noisy_img_sin, img_sin[noisy_img_seq_len:, :]], dim=0).to( device=device ) return { "txt_seq_lens": txt_seq_lens, "freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)), "img_shapes": img_shapes, } @dataclass class QwenImageEditPlus_2511_PipelineConfig(QwenImageEditPlusPipelineConfig): dit_config: DiTConfig = field(default_factory=QwenImageEditPlus_2511_DitConfig) @dataclass class QwenImageLayeredPipelineConfig(QwenImageEditPipelineConfig): resolution: int = 640 vae_precision: str = "bf16" def _prepare_edit_cond_kwargs( self, batch, prompt_embeds, rotary_emb, device, dtype ): batch_size = batch.latents.shape[0] assert batch_size == 1 height = batch.height width = batch.width vae_scale_factor = self.get_vae_scale_factor() img_shapes = batch.img_shapes txt_seq_lens = batch.txt_seq_lens freqs_cis = QwenImageEditPlusPipelineConfig.get_freqs_cis( img_shapes, txt_seq_lens, rotary_emb, device, dtype ) # perform sp shard on noisy image tokens noisy_img_seq_len = ( 1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2) ) img_cache, txt_cache = freqs_cis noisy_img_cache = shard_rotary_emb_for_sp(img_cache[:noisy_img_seq_len, :]) img_cache = torch.cat( [noisy_img_cache, img_cache[noisy_img_seq_len:, :]], dim=0 ).to(device=device) return { "txt_seq_lens": txt_seq_lens, "img_shapes": img_shapes, "freqs_cis": (img_cache, txt_cache), "additional_t_cond": torch.tensor([0], device=device, dtype=torch.long), } def _unpad_and_unpack_latents(self, latents, batch): vae_scale_factor = self.vae_config.arch_config.vae_scale_factor channels = self.dit_config.arch_config.in_channels batch_size = latents.shape[0] layers = batch.num_frames height = 2 * (int(batch.height) // (vae_scale_factor * 2)) width = 2 * (int(batch.width) // (vae_scale_factor * 2)) latents = maybe_unpad_latents(latents, batch) latents = latents.view( batch_size, layers + 1, height // 2, width // 2, channels // 4, 2, 2 ) latents = latents.permute(0, 1, 4, 2, 5, 3, 6) latents = latents.reshape( batch_size, layers + 1, channels // (2 * 2), height, width ) latents = latents.permute(0, 2, 1, 3, 4) # (b, c, f, h, w) return latents, batch_size, channels, height, width def allow_set_num_frames(self): return True def post_denoising_loop(self, latents, batch): # unpack latents for qwen-image ( latents, batch_size, channels, height, width, ) = self._unpad_and_unpack_latents(latents, batch) b, c, f, h, w = latents.shape latents = latents[:, :, 1:] # remove the first frame as it is the origin input latents = latents.permute(0, 2, 1, 3, 4).view(-1, c, 1, h, w) # latents = latents.reshape(batch_size, channels // (2 * 2), 1, height, width) return latents ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/sana.py ================================================ # SPDX-License-Identifier: Apache-2.0 # # Pipeline configuration for SANA text-to-image generation. # # SANA produces 4D spatial latents (B, C, H', W') directly — unlike Flux/QwenImage # which use packed token-style latents (B, S, D). This means: # - We inherit SpatialImagePipelineConfig (not ImagePipelineConfig) # - prepare_latent_shape returns 4D, not 5D # - post_denoising_loop is a no-op (no un-packing needed) # - shard_latents_for_sp shards along the H' dimension # # SANA does NOT use rotary position embeddings, so prepare_pos/neg_cond_kwargs # return empty dicts (the DiT only needs hidden_states + encoder_hidden_states + timestep). # # CFG is handled by the denoising stage via guidance_scale in sampling params. # should_use_guidance=False means no embedded guidance (no extra guidance token in forward), # but negative_prompt + guidance_scale > 1.0 still enables standard classifier-free guidance. from collections.abc import Callable from dataclasses import dataclass, field import torch from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig from sglang.multimodal_gen.configs.models.dits.sana import SanaConfig from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput from sglang.multimodal_gen.configs.models.encoders.base import EncoderConfig from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config from sglang.multimodal_gen.configs.models.vaes.sana import SanaVAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, SpatialImagePipelineConfig, preprocess_text, ) def sana_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: # SANA uses the final hidden state from Gemma2 directly as text conditioning. # No intermediate-layer extraction or masking needed (unlike QwenImage/ZImage). return outputs.last_hidden_state @dataclass class SanaPipelineConfig(SpatialImagePipelineConfig): task_type: ModelTaskType = ModelTaskType.T2I # should_use_guidance=False disables *embedded* guidance (timestep-conditioned # guidance token). Standard CFG via guidance_scale is still active. should_use_guidance: bool = False enable_autocast: bool = False # DC-AE does not support tiling or SP VAE decode yet. vae_tiling: bool = False vae_sp: bool = False vae_precision: str = "bf16" dit_config: DiTConfig = field(default_factory=SanaConfig) vae_config: VAEConfig = field(default_factory=SanaVAEConfig) # Single text encoder: Gemma2 (unlike Flux which uses CLIP + T5) text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (Gemma2Config(),) ) text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",)) preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( default_factory=lambda: (preprocess_text,), ) postprocess_text_funcs: tuple[Callable[[str], str], ...] = field( default_factory=lambda: (sana_postprocess_text,) ) def prepare_latent_shape(self, batch, batch_size, num_frames): # 4D latent shape: (B, C, H', W') — no temporal dim for T2I. # DC-AE compresses 1024x1024 -> 32x32 with 32 channels. compression = self.vae_config.arch_config.spatial_compression_ratio height = batch.height // compression width = batch.width // compression num_channels = self.dit_config.arch_config.num_channels_latents shape = (batch_size, num_channels, height, width) return shape def get_pos_prompt_embeds(self, batch): # Single encoder -> index [0] (Flux uses [1] because T5 is encoder #2) return batch.prompt_embeds[0] def get_neg_prompt_embeds(self, batch): return batch.negative_prompt_embeds[0] def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): # encoder_attention_mask: batch stores list-of-tensors; diffusers' SanaTransformer # expects a single tensor (sglang's has list handling). Override with [0]. out = {} m = batch.prompt_attention_mask if isinstance(m, (list, tuple)): out["encoder_attention_mask"] = m[0] if m else None elif m is not None: out["encoder_attention_mask"] = m return out def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): out = {} m = batch.negative_attention_mask if isinstance(m, (list, tuple)): out["encoder_attention_mask"] = m[0] if m else None elif m is not None: out["encoder_attention_mask"] = m return out def post_denoising_loop(self, latents, batch): return latents ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/wan.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from collections.abc import Callable from dataclasses import dataclass, field import torch from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig from sglang.multimodal_gen.configs.models.dits import WanVideoConfig from sglang.multimodal_gen.configs.models.encoders import ( BaseEncoderOutput, CLIPVisionConfig, T5Config, ) from sglang.multimodal_gen.configs.models.vaes import WanVAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, PipelineConfig, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) def t5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: mask: torch.Tensor = outputs.attention_mask hidden_state: torch.Tensor = outputs.last_hidden_state seq_lens = mask.gt(0).sum(dim=1).long() assert torch.isnan(hidden_state).sum() == 0 prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)] prompt_embeds_tensor: torch.Tensor = torch.stack( [ torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))]) for u in prompt_embeds ], dim=0, ) return prompt_embeds_tensor @dataclass class WanI2VCommonConfig(PipelineConfig): # for all wan i2v pipelines def adjust_num_frames(self, num_frames): vae_scale_factor_temporal = self.vae_config.arch_config.scale_factor_temporal if num_frames % vae_scale_factor_temporal != 1: logger.warning( f"`num_frames - 1` has to be divisible by {vae_scale_factor_temporal}. Rounding to the nearest number." ) num_frames = ( num_frames // vae_scale_factor_temporal * vae_scale_factor_temporal + 1 ) return num_frames return num_frames @dataclass class WanT2V480PConfig(PipelineConfig): """Base configuration for Wan T2V 1.3B pipeline architecture.""" task_type: ModelTaskType = ModelTaskType.T2V # WanConfig-specific parameters with defaults # DiT dit_config: DiTConfig = field(default_factory=WanVideoConfig) # VAE vae_config: VAEConfig = field(default_factory=WanVAEConfig) vae_tiling: bool = False vae_sp: bool = False # Denoising stage flow_shift: float | None = 3.0 # Text encoding stage text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (T5Config(),) ) postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = ( field(default_factory=lambda: (t5_postprocess_text,)) ) # Precision for each component precision: str = "bf16" vae_precision: str = "fp32" text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32",)) # WanConfig-specific added parameters def __post_init__(self): self.vae_config.load_encoder = False self.vae_config.load_decoder = True @dataclass class TurboWanT2V480PConfig(WanT2V480PConfig): """Base configuration for Wan T2V 1.3B pipeline architecture.""" flow_shift: float | None = 8.0 dmd_denoising_steps: list[int] | None = field( default_factory=lambda: [988, 932, 852, 608] ) @dataclass class WanT2V720PConfig(WanT2V480PConfig): """Base configuration for Wan T2V 14B 720P pipeline architecture.""" # WanConfig-specific parameters with defaults # Denoising stage flow_shift: float | None = 5.0 @dataclass class WanI2V480PConfig(WanT2V480PConfig, WanI2VCommonConfig): """Base configuration for Wan I2V 14B 480P pipeline architecture.""" max_area: int = 480 * 832 # WanConfig-specific parameters with defaults task_type: ModelTaskType = ModelTaskType.I2V # Precision for each component image_encoder_config: EncoderConfig = field(default_factory=CLIPVisionConfig) image_encoder_precision: str = "fp32" image_encoder_extra_args: dict = field( default_factory=lambda: dict( output_hidden_states=True, ) ) def postprocess_image(self, image): return image.hidden_states[-2] def __post_init__(self) -> None: self.vae_config.load_encoder = True self.vae_config.load_decoder = True @dataclass class WanI2V720PConfig(WanI2V480PConfig): """Base configuration for Wan I2V 14B 720P pipeline architecture.""" max_area: int = 720 * 1280 # WanConfig-specific parameters with defaults # Denoising stage flow_shift: float | None = 5.0 @dataclass class TurboWanI2V720Config(WanI2V720PConfig): flow_shift: float | None = 8.0 dmd_denoising_steps: list[int] | None = field( default_factory=lambda: [996, 932, 852, 608] ) boundary_ratio: float | None = 0.9 def __post_init__(self) -> None: self.dit_config.boundary_ratio = self.boundary_ratio @dataclass class FastWan2_1_T2V_480P_Config(WanT2V480PConfig): """Base configuration for FastWan T2V 1.3B 480P pipeline architecture with DMD""" # WanConfig-specific parameters with defaults # Denoising stage flow_shift: float | None = 8.0 dmd_denoising_steps: list[int] | None = field( default_factory=lambda: [1000, 757, 522] ) @dataclass class Wan2_2_TI2V_5B_Config(WanT2V480PConfig, WanI2VCommonConfig): flow_shift: float | None = 5.0 task_type: ModelTaskType = ModelTaskType.TI2V expand_timesteps: bool = True # ti2v, 5B vae_stride = (4, 16, 16) def prepare_latent_shape(self, batch, batch_size, num_frames): F = num_frames z_dim = self.vae_config.arch_config.z_dim vae_stride = self.vae_stride oh = batch.height ow = batch.width shape = (batch_size, z_dim, F, oh // vae_stride[1], ow // vae_stride[2]) return shape def __post_init__(self) -> None: self.vae_config.load_encoder = True self.vae_config.load_decoder = True self.dit_config.expand_timesteps = self.expand_timesteps @dataclass class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config): flow_shift: float | None = 5.0 dmd_denoising_steps: list[int] | None = field( default_factory=lambda: [1000, 757, 522] ) @dataclass class Wan2_2_T2V_A14B_Config(WanT2V480PConfig): flow_shift: float | None = 12.0 boundary_ratio: float | None = 0.875 def __post_init__(self) -> None: self.dit_config.boundary_ratio = self.boundary_ratio @dataclass class Wan2_2_I2V_A14B_Config(WanI2V480PConfig): flow_shift: float | None = 5.0 boundary_ratio: float | None = 0.900 def __post_init__(self) -> None: super().__post_init__() self.dit_config.boundary_ratio = self.boundary_ratio # ============================================= # ============= Causal Self-Forcing ============= # ============================================= @dataclass class SelfForcingWanT2V480PConfig(WanT2V480PConfig): is_causal: bool = True flow_shift: float | None = 5.0 dmd_denoising_steps: list[int] | None = field( default_factory=lambda: [1000, 750, 500, 250] ) warp_denoising_step: bool = True ================================================ FILE: python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import math from dataclasses import dataclass, field from typing import Callable import torch from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ImagePipelineConfig, ModelTaskType, ) from sglang.multimodal_gen.runtime.distributed.communication_op import ( sequence_model_parallel_all_gather, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_sp_parallel_rank, get_sp_world_size, ) def zimage_preprocess_text(prompt: str): messages = [ {"role": "user", "content": prompt}, ] return messages def zimage_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: device = outputs.hidden_states[-2].device prompt_mask = _text_inputs.attention_mask.to(device).bool() return outputs.hidden_states[-2][0][prompt_mask[0]] class TransformersModelConfig(EncoderConfig): tokenizer_kwargs: dict = field(default_factory=lambda: {}) @dataclass class ZImagePipelineConfig(ImagePipelineConfig): should_use_guidance: bool = False task_type: ModelTaskType = ModelTaskType.T2I dit_config: DiTConfig = field(default_factory=ZImageDitConfig) vae_config: VAEConfig = field(default_factory=FluxVAEConfig) text_encoder_configs: tuple[EncoderConfig, ...] = field( default_factory=lambda: (Qwen3TextConfig(),) ) preprocess_text_funcs: tuple[Callable, ...] = field( default_factory=lambda: (zimage_preprocess_text,) ) postprocess_text_funcs: tuple[Callable, ...] = field( default_factory=lambda: (zimage_postprocess_text,) ) SEQ_LEN_MULTIPLE: int = 32 PATCH_SIZE: int = 2 F_PATCH_SIZE: int = 1 def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict: # flatten to 1-d list inputs = tokenizer.apply_chat_template( prompts, tokenize=True, add_generation_prompt=True, enable_thinking=True, padding="max_length", max_length=512, # TODO (yhyang201): set max length according to config truncation=True, return_tensors="pt", return_dict=True, ) return inputs @staticmethod def _ceil_to_multiple(x: int, m: int) -> int: if m <= 0: return x return int(math.ceil(x / m) * m) def _build_zimage_sp_plan(self, batch) -> dict: """Build a minimal SP plan on batch for zimage (spatial sharding + cap sharding).""" sp_size = get_sp_world_size() rank = get_sp_parallel_rank() raw_latent_shape = getattr(batch, "raw_latent_shape", None) if raw_latent_shape is not None and len(raw_latent_shape) >= 5: H = int(raw_latent_shape[3]) W = int(raw_latent_shape[4]) else: H = int( batch.height // self.vae_config.arch_config.spatial_compression_ratio ) W = int( batch.width // self.vae_config.arch_config.spatial_compression_ratio ) # Rule: shard along the larger spatial dimension (W/H), implemented via optional H/W transpose. # Choose the larger of H and W for sharding, so H_eff = max(H, W). swap_hw = W > H H_eff = W if swap_hw else H W_eff = H if swap_hw else W # ZImage uses PATCH_SIZE=2 for spatial patchify; shard in token space and convert back to latent rows. H_tok = H_eff // self.PATCH_SIZE W_tok = W_eff // self.PATCH_SIZE H_tok_pad = self._ceil_to_multiple(H_tok, sp_size) H_tok_local = H_tok_pad // sp_size h0_tok = rank * H_tok_local # Cap/text sharding: avoid duplicating cap tokens across ranks. cap_len = ( int(batch.prompt_embeds[0].size(0)) if getattr(batch, "prompt_embeds", None) else 0 ) cap_total = self._ceil_to_multiple(cap_len, self.SEQ_LEN_MULTIPLE * sp_size) cap_local = cap_total // sp_size cap_start = rank * cap_local plan = { "sp_size": sp_size, "rank": rank, "swap_hw": swap_hw, "H": H, "W": W, "H_eff": H_eff, "W_eff": W_eff, "H_tok": H_tok, "W_tok": W_tok, "H_tok_pad": H_tok_pad, "H_tok_local": H_tok_local, "h0_tok": h0_tok, "cap_total": cap_total, "cap_local": cap_local, "cap_start": cap_start, } batch._zimage_sp_plan = plan return plan def _get_zimage_sp_plan(self, batch) -> dict: plan = getattr(batch, "_zimage_sp_plan", None) sp_size = get_sp_world_size() if plan is None or plan.get("sp_size") != sp_size: plan = self._build_zimage_sp_plan(batch) return plan def _shard_cap(self, cap: torch.Tensor, plan: dict) -> torch.Tensor: """cap: [L, D] -> [cap_local, D], padded by repeating last token.""" if plan["sp_size"] <= 1: return cap # print(f"cap shape: {cap.shape}") # [L, 2560] for zimage-turbo L = cap.size(0) cap_total = plan["cap_total"] if cap_total > L: cap = torch.cat([cap, cap[-1:].repeat(cap_total - L, 1)], dim=0) start = plan["cap_start"] local = plan["cap_local"] return cap[start : start + local] def get_pos_prompt_embeds(self, batch): # Keep ZImage model signature: encoder_hidden_states is List[Tensor] if get_sp_world_size() <= 1: return batch.prompt_embeds plan = self._get_zimage_sp_plan(batch) return [self._shard_cap(batch.prompt_embeds[0], plan)] def shard_latents_for_sp(self, batch, latents): sp_size = get_sp_world_size() if sp_size <= 1 or latents.dim() != 5: return latents, False plan = self._get_zimage_sp_plan(batch) # Layout: [B, C, T, H, W]. Always shard on dim=3 by optionally swapping H/W. if plan["swap_hw"]: latents = latents.transpose(3, 4).contiguous() # Pad on effective-H so that H_tok is divisible by sp. H_eff = latents.size(3) H_tok = H_eff // self.PATCH_SIZE pad_tok = plan["H_tok_pad"] - H_tok pad_lat = pad_tok * self.PATCH_SIZE if pad_lat > 0: pad = latents[:, :, :, -1:, :].repeat(1, 1, 1, pad_lat, 1) latents = torch.cat([latents, pad], dim=3) h0 = plan["h0_tok"] * self.PATCH_SIZE h1 = (plan["h0_tok"] + plan["H_tok_local"]) * self.PATCH_SIZE latents = latents[:, :, :, h0:h1, :] batch._zimage_sp_swap_hw = plan["swap_hw"] return latents, True def gather_latents_for_sp(self, latents): # Gather on effective-H dim=3 (matches shard_latents_for_sp); swap-back is handled in post_denoising_loop. latents = latents.contiguous() if get_sp_world_size() <= 1 or latents.dim() != 5: return latents return sequence_model_parallel_all_gather(latents, dim=3) def post_denoising_loop(self, latents, batch): # Restore swapped H/W and crop padded spatial dims before final reshape. if latents.dim() == 5 and getattr(batch, "_zimage_sp_swap_hw", False): latents = latents.transpose(3, 4).contiguous() raw_latent_shape = getattr(batch, "raw_latent_shape", None) if raw_latent_shape is not None and latents.dim() == 5: latents = latents[:, :, :, : raw_latent_shape[3], : raw_latent_shape[4]] bs, channels, num_frames, height, width = latents.shape if raw_latent_shape is not None and num_frames > raw_latent_shape[2]: latents = latents[:, :, : raw_latent_shape[2], :, :] num_frames = raw_latent_shape[2] if num_frames != 1: return latents[:, :, 0, :, :] return latents.view(bs, channels, height, width) def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch): def create_coordinate_grid(size, start=None, device=None): if start is None: start = (0 for _ in size) axes = [ torch.arange(x0, x0 + span, dtype=torch.int32, device=device) for x0, span in zip(start, size) ] grids = torch.meshgrid(axes, indexing="ij") return torch.stack(grids, dim=-1) sp_size = get_sp_world_size() if sp_size > 1: # SP path: build local-only freqs_cis matching local cap/x. plan = self._get_zimage_sp_plan(batch) # cap (local) cap_pos_ids = create_coordinate_grid( size=(plan["cap_local"], 1, 1), start=(1 + plan["cap_start"], 0, 0), device=device, ).flatten(0, 2) cap_freqs_cis = rotary_emb(cap_pos_ids) # image (local, effective H-shard). Use cap_total for a stable offset across ranks/passes. F_tokens = 1 H_tokens_local = plan["H_tok_local"] W_tokens = plan["W_tok"] img_pos_ids = create_coordinate_grid( size=(F_tokens, H_tokens_local, W_tokens), start=(plan["cap_total"] + 1, plan["h0_tok"], 0), device=device, ).flatten(0, 2) img_pad_len = (-img_pos_ids.shape[0]) % self.SEQ_LEN_MULTIPLE if img_pad_len: pad_ids = create_coordinate_grid( size=(1, 1, 1), start=(0, 0, 0), device=device ).flatten(0, 2) img_pos_ids = torch.cat( [img_pos_ids, pad_ids.repeat(img_pad_len, 1)], dim=0 ) x_freqs_cis = rotary_emb(img_pos_ids) return (cap_freqs_cis, x_freqs_cis) cap_ori_len = prompt_embeds.size(0) cap_padding_len = (-cap_ori_len) % self.SEQ_LEN_MULTIPLE cap_padded_pos_ids = create_coordinate_grid( size=(cap_ori_len + cap_padding_len, 1, 1), start=(1, 0, 0), device=device, ).flatten(0, 2) F = 1 H = height // self.vae_config.arch_config.spatial_compression_ratio W = width // self.vae_config.arch_config.spatial_compression_ratio pH, pW = self.PATCH_SIZE, self.PATCH_SIZE pF = self.F_PATCH_SIZE F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW image_ori_len = F_tokens * H_tokens * W_tokens image_padding_len = (-image_ori_len) % self.SEQ_LEN_MULTIPLE image_ori_pos_ids = create_coordinate_grid( size=(F_tokens, H_tokens, W_tokens), start=(cap_ori_len + cap_padding_len + 1, 0, 0), device=device, ).flatten(0, 2) image_padding_pos_ids = ( create_coordinate_grid( size=(1, 1, 1), start=(0, 0, 0), device=device, ) .flatten(0, 2) .repeat(image_padding_len, 1) ) image_padded_pos_ids = torch.cat( [image_ori_pos_ids, image_padding_pos_ids], dim=0 ) cap_freqs_cis = rotary_emb(cap_padded_pos_ids) x_freqs_cis = rotary_emb(image_padded_pos_ids) return (cap_freqs_cis, x_freqs_cis) def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): return { "freqs_cis": self.get_freqs_cis( batch.prompt_embeds[0], batch.width, batch.height, device, rotary_emb, batch, ), } def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): return { "freqs_cis": self.get_freqs_cis( batch.prompt_embeds[0], batch.width, batch.height, device, rotary_emb, batch, ), } ================================================ FILE: python/sglang/multimodal_gen/configs/quantization.py ================================================ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import os import re from dataclasses import dataclass from typing import Any import torch from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( is_nunchaku_available, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import StoreBoolean logger = init_logger(__name__) @dataclass class NunchakuSVDQuantArgs: """CLI-facing configuration for Nunchaku (SVDQuant) inference. This is intentionally lightweight and only contains arguments needed to construct `runtime.layers.quantization.nunchaku_config.NunchakuConfig`. """ enable_svdquant: bool = False transformer_weights_path: str | None = None quantization_precision: str | None = None # "int4" or "nvfp4" quantization_rank: int | None = None quantization_act_unsigned: bool = False def _adjust_config(self) -> None: """infer precision and rank from filename if not provided""" if self.transformer_weights_path and not self.enable_svdquant: filename = os.path.basename(self.transformer_weights_path) if re.search(r"svdq-(int4|fp4)_r(\d+)", filename): self.enable_svdquant = True if not self.enable_svdquant or not self.transformer_weights_path: return inferred_precision = None inferred_rank = None filename = os.path.basename(self.transformer_weights_path) # Expected pattern: svdq-{precision}_r{rank}-... # e.g., svdq-int4_r32-qwen-image.safetensors match = re.search(r"svdq-(int4|fp4)_r(\d+)", filename) if match: p_str, r_str = match.groups() inferred_precision = "nvfp4" if p_str == "fp4" else "int4" inferred_rank = int(r_str) if self.quantization_precision is None: self.quantization_precision = inferred_precision or "int4" if inferred_precision: logger.info( f"inferred --quantization-precision: {self.quantization_precision} " f"from --transformer-weights-path: {self.transformer_weights_path}" ) if self.quantization_rank is None: self.quantization_rank = inferred_rank or 32 if inferred_rank: logger.info( f"inferred --quantization-rank: {self.quantization_rank} " f"from --transformer-weights-path: {self.transformer_weights_path}" ) def validate(self) -> None: # TODO: warn if the served model doesn't support nunchaku self._adjust_config() if not self.enable_svdquant: return if not current_platform.is_cuda(): raise ValueError( "Nunchaku SVDQuant is only supported on NVIDIA CUDA GPUs " "(Ampere SM8x or SM12x)." ) device_count = torch.cuda.device_count() unsupported: list[str] = [] for i in range(device_count): major, minor = torch.cuda.get_device_capability(i) if major == 9: unsupported.append(f"cuda:{i} (SM{major}{minor}, Hopper)") elif major not in (8, 12): unsupported.append(f"cuda:{i} (SM{major}{minor})") if unsupported: raise ValueError( "Nunchaku SVDQuant is currently only supported on Ampere (SM8x) or SM12x GPUs; " "Hopper (SM90) is not supported. " f"Unsupported devices: {', '.join(unsupported)}. " "Disable it with --enable-svdquant false." ) if not self.transformer_weights_path: raise ValueError( "--enable-svdquant requires --transformer-weights-path to be set" ) if not is_nunchaku_available(): raise ValueError( "Nunchaku is enabled, but not installed. Please refer to https://nunchaku.tech/docs/nunchaku/installation/installation.html for detailed installation methods." ) if self.quantization_precision not in ("int4", "nvfp4"): raise ValueError( f"Invalid --quantization-precision: {self.quantization_precision}. " "Must be one of: int4, nvfp4" ) if self.quantization_rank <= 0: raise ValueError( f"Invalid --quantization-rank: {self.quantization_rank}. Must be > 0" ) @staticmethod def add_cli_args(parser) -> None: parser.add_argument( "--enable-svdquant", action=StoreBoolean, default=NunchakuSVDQuantArgs.enable_svdquant, help="Enable Nunchaku SVDQuant (W4A4-style) inference.", ) parser.add_argument( "--transformer-weights-path", type=str, default=NunchakuSVDQuantArgs.transformer_weights_path, help=( "Path to pre-quantized transformer weights. Can be a single .safetensors " "file, a directory, or a HuggingFace repo ID. Used by Nunchaku (SVDQuant) and quantized single-file checkpoints." ), ) parser.add_argument( "--quantization-precision", type=str, default=None, help="Quantization precision: int4 or nvfp4. If not specified, inferred from model path or defaults to int4.", ) parser.add_argument( "--quantization-rank", type=int, default=None, help="SVD low-rank dimension (e.g., 32). If not specified, inferred from model path or defaults to 32.", ) parser.add_argument( "--quantization-act-unsigned", action=StoreBoolean, default=NunchakuSVDQuantArgs.quantization_act_unsigned, help="Use unsigned activation quantization (if supported).", ) @classmethod def from_dict(cls, kwargs: dict[str, Any]) -> "NunchakuSVDQuantArgs": # Map CLI/config keys to dataclass fields (keep backwards compatibility). path = ( kwargs.get("transformer_weights_path") or kwargs.get("transformer_quantized_path") or kwargs.get("quantized_model_path") ) return cls( enable_svdquant=bool(kwargs.get("enable_svdquant", cls.enable_svdquant)), transformer_weights_path=path, quantization_precision=kwargs.get("quantization_precision"), quantization_rank=kwargs.get("quantization_rank"), quantization_act_unsigned=bool( kwargs.get("quantization_act_unsigned", cls.quantization_act_unsigned) ), ) ================================================ FILE: python/sglang/multimodal_gen/configs/sample/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from sglang.multimodal_gen.configs.sample.diffusers_generic import ( DiffusersGenericSamplingParams, ) from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams __all__ = ["SamplingParams", "DiffusersGenericSamplingParams"] ================================================ FILE: python/sglang/multimodal_gen/configs/sample/diffusers_generic.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Generic sampling parameters for diffusers backend. This module provides generic sampling parameters that work with any diffusers pipeline. """ from dataclasses import dataclass, field from typing import Any, ClassVar from sglang.multimodal_gen.configs.sample.sampling_params import ( DataType, SamplingParams, ) @dataclass class DiffusersGenericSamplingParams(SamplingParams): """ Generic sampling parameters for diffusers backend. These parameters cover the most common options across different diffusers pipelines. The diffusers pipeline will use whichever parameters it supports. For pipeline-specific parameters, use `diffusers_kwargs` dict which will be passed directly to the diffusers pipeline call. """ _default_height: ClassVar[int] = 1024 _default_width: ClassVar[int] = 1024 # Override defaults with more conservative values that work across pipelines num_frames: int = 1 # default to image generation height: int = 1024 width: int = 1024 num_inference_steps: int = 30 guidance_scale: float = 7.5 negative_prompt: str = "" # extra kwargs to pass directly to the diffusers pipeline # example: {"output_type": "latent", "return_dict": False} diffusers_kwargs: dict[str, Any] = field(default_factory=dict) def __post_init__(self) -> None: if self.num_frames > 1: self.data_type = DataType.VIDEO else: self.data_type = DataType.IMAGE super().__post_init__() ================================================ FILE: python/sglang/multimodal_gen/configs/sample/flux.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass from typing import ClassVar from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams @dataclass class FluxSamplingParams(SamplingParams): _default_height: ClassVar[int] = 128 * 8 # default_sample_size * vae_scale_factor _default_width: ClassVar[int] = 128 * 8 num_frames: int = 1 # Denoising stage guidance_scale: float = 1.0 negative_prompt: str = None num_inference_steps: int = 50 @dataclass class Flux2KleinSamplingParams(FluxSamplingParams): # Klein is step-distilled, so default to 4 steps num_inference_steps: int = 4 ================================================ FILE: python/sglang/multimodal_gen/configs/sample/glmimage.py ================================================ from dataclasses import dataclass from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams @dataclass class GlmImageSamplingParams(SamplingParams): negative_prompt = "" num_frames: int = 1 guidance_scale: float = 1.5 num_inference_steps: int = 30 ================================================ FILE: python/sglang/multimodal_gen/configs/sample/helios.py ================================================ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams @dataclass class HeliosT2VSamplingParams(SamplingParams): # Video parameters height: int = 384 width: int = 640 num_frames: int = 99 fps: int = 24 # Denoising stage guidance_scale: float = 5.0 negative_prompt: str = ( "Bright tones, overexposed, static, blurred details, subtitles, style, " "works, paintings, images, static, overall gray, worst quality, low quality, " "JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, " "poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, " "still picture, messy background, three legs, many people in the background, " "walking backwards" ) num_inference_steps: int = 50 # Helios T2V supported resolutions supported_resolutions: list[tuple[int, int]] | None = field( default_factory=lambda: [ (640, 384), # ~5:3 (384, 640), # ~3:5 (832, 480), # ~16:9-ish (480, 832), # ~9:16-ish ] ) @dataclass class HeliosMidSamplingParams(HeliosT2VSamplingParams): """Sampling params for Helios-Mid (Stage 2 pyramid SR).""" num_inference_steps: int = 20 @dataclass class HeliosDistilledSamplingParams(HeliosT2VSamplingParams): """Sampling params for Helios-Distilled (DMD, no CFG needed).""" guidance_scale: float = 1.0 num_inference_steps: int = 10 ================================================ FILE: python/sglang/multimodal_gen/configs/sample/hunyuan.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams @dataclass class HunyuanSamplingParams(SamplingParams): num_inference_steps: int = 50 num_frames: int = 125 height: int = 720 width: int = 1280 fps: int = 24 guidance_scale: float = 1.0 # HunyuanVideo supported resolutions supported_resolutions: list[tuple[int, int]] | None = field( default_factory=lambda: [ # 540p resolutions (960, 544), # 9:16 (544, 960), # 16:9 (832, 624), # 4:3 (624, 832), # 3:4 (720, 720), # 1:1 # 720p resolutions (recommended) (1280, 720), # 9:16 (720, 1280), # 16:9 (832, 1104), # 4:3 (1104, 832), # 3:4 (960, 960), # 1:1 ] ) teacache_params: TeaCacheParams = field( default_factory=lambda: TeaCacheParams( teacache_thresh=0.15, coefficients=[ 7.33226126e02, -4.01131952e02, 6.75869174e01, -3.14987800e00, 9.61237896e-02, ], ) ) @dataclass class FastHunyuanSamplingParam(HunyuanSamplingParams): num_inference_steps: int = 6 ================================================ FILE: python/sglang/multimodal_gen/configs/sample/hunyuan3d.py ================================================ # SPDX-License-Identifier: Apache-2.0 """Sampling parameters for Hunyuan3D generation.""" from dataclasses import dataclass from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams @dataclass class Hunyuan3DSamplingParams(SamplingParams): """Sampling parameters for Hunyuan3D image-to-mesh generation.""" negative_prompt: str = "" shape_num_inference_steps: int = 50 guidance_scale: float = 5.0 paint_num_inference_steps: int = 30 paint_guidance_scale: float = 2.0 def __post_init__(self): if self.prompt is None: self.prompt = "" if self.num_inference_steps is None: self.num_inference_steps = self.shape_num_inference_steps self.guidance_scale = max(5.0, min(self.guidance_scale, 6.5)) super().__post_init__() ================================================ FILE: python/sglang/multimodal_gen/configs/sample/ltx_2.py ================================================ import dataclasses from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams @dataclasses.dataclass class LTX2SamplingParams(SamplingParams): """Sampling parameters for LTX-2.""" # Match the reference defaults used by ltx-pipelines (one-stage). # See: LTX-2/packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py seed: int = 10 # Video parameters height: int = 512 width: int = 768 num_frames: int = 121 fps: int = 24 # Audio specific generate_audio: bool = True # Denoising parameters guidance_scale: float = 4.0 num_inference_steps: int = 40 # Match ltx-pipelines default negative prompt (covers video + audio artifacts). negative_prompt: str = ( "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, " "grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, " "deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, " "wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of " "field, background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent " "lighting direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny " "valley effect, incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, " "mismatched lip sync, silent or muted audio, distorted voice, robotic voice, echo, background noise, " "off-sync audio, incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward " "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, " "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts." ) ================================================ FILE: python/sglang/multimodal_gen/configs/sample/mova.py ================================================ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams @dataclass class MOVASamplingParams(SamplingParams): # Video parameters (MOVA defaults) height: int = 352 width: int = 640 num_frames: int = 193 fps: int = 24 # Denoising stage guidance_scale: float = 5.0 num_inference_steps: int = 50 sigma_shift: float = 5.0 visual_shift: float = 5.0 audio_shift: float = 5.0 adjust_frames: bool = False negative_prompt: str = ( "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，" "整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，" "画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，" "静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" ) @dataclass class MOVA_360P_SamplingParams(MOVASamplingParams): # Video parameters (MOVA 360P) height: int = 352 width: int = 640 # MOVA 360P supported resolutions supported_resolutions: list[tuple[int, int]] = field( default_factory=lambda: [ (352, 640), (640, 352), ] ) @dataclass class MOVA_720P_SamplingParams(MOVASamplingParams): # Video parameters (MOVA 720P) height: int = 720 width: int = 1280 # MOVA 720P supported resolutions supported_resolutions: list[tuple[int, int]] = field( default_factory=lambda: [ (720, 1280), (1280, 720), ] ) ================================================ FILE: python/sglang/multimodal_gen/configs/sample/qwenimage.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams @dataclass class QwenImageSamplingParams(SamplingParams): negative_prompt: str = " " num_frames: int = 1 # Denoising stage guidance_scale: float = 4.0 num_inference_steps: int = 50 @dataclass class QwenImage2512SamplingParams(QwenImageSamplingParams): negative_prompt: str = ( "低分辨率，低画质，肢体畸形，手指畸形，画面过饱和，蜡像感，人脸无细节，过度光滑，画面具有AI感。构图混乱。文字模糊，扭曲。" ) @dataclass class QwenImageEditPlusSamplingParams(QwenImageSamplingParams): # Denoising stage guidance_scale: float = 4.0 # true_cfg_scale: float = 4.0 num_inference_steps: int = 40 @dataclass class QwenImageLayeredSamplingParams(QwenImageSamplingParams): # num_frames: int = 4 height: int = 640 width: int = 640 prompt: str = " " negative_prompt: str = " " guidance_scale: float = 4.0 num_inference_steps: int = 50 cfg_normalize: bool = True use_en_prompt: bool = True ================================================ FILE: python/sglang/multimodal_gen/configs/sample/sampling_params.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import argparse import dataclasses import hashlib import json import math import os import os.path import re import time import unicodedata import uuid from dataclasses import dataclass from enum import Enum, auto from typing import TYPE_CHECKING, Any, ClassVar from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import StoreBoolean, expand_path_fields logger = init_logger(__name__) if TYPE_CHECKING: from sglang.multimodal_gen.runtime.server_args import ServerArgs def _json_safe(obj: Any): """ Recursively convert objects to JSON-serializable forms. - Enums -> their name - Sets/Tuples -> lists - Dicts/Lists -> recursively processed """ if isinstance(obj, Enum): return obj.name if isinstance(obj, dict): return {k: _json_safe(v) for k, v in obj.items()} if isinstance(obj, (list, tuple, set)): return [_json_safe(v) for v in obj] return obj def generate_request_id() -> str: return str(uuid.uuid4()) def _sanitize_filename(name: str, replacement: str = "_", max_length: int = 150) -> str: """Create a filesystem- and ffmpeg-friendly filename. - Normalize to ASCII (drop accents and unsupported chars) - Replace spaces with underscores - Replace any char not in [A-Za-z0-9_.-] with replacement - Collapse multiple underscores - Trim leading/trailing dots/underscores and limit length """ normalized = unicodedata.normalize("NFKD", name) ascii_name = normalized.encode("ascii", "ignore").decode("ascii") ascii_name = ascii_name.replace(" ", "_") ascii_name = re.sub(r"[^A-Za-z0-9._-]", replacement, ascii_name) ascii_name = re.sub(r"_+", "_", ascii_name).strip("._") if not ascii_name: ascii_name = "output" if max_length and len(ascii_name) > max_length: ascii_name = ascii_name[:max_length] return ascii_name class DataType(Enum): IMAGE = auto() VIDEO = auto() MESH = auto() def get_default_extension(self) -> str: if self == DataType.IMAGE: return "png" if self == DataType.VIDEO: return "mp4" return "glb" @dataclass class SamplingParams: """ Sampling parameters for generation. """ data_type: DataType = DataType.VIDEO request_id: str | None = None # All fields below are copied from ForwardBatch # Image inputs image_path: str | list[str] | None = None # Text inputs prompt: str | list[str] | None = None negative_prompt: str = ( "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" ) prompt_path: str | None = None output_path: str | None = None output_file_name: str | None = None output_quality: str | None = "default" output_compression: int | None = None # Frame interpolation enable_frame_interpolation: bool = False frame_interpolation_exp: int = 1 # 1=2x, 2=4x frame_interpolation_scale: float = 1.0 # RIFE inference scale (0.5 for high-res) frame_interpolation_model_path: str | None = ( None # local dir or HF repo ID with flownet.pkl (default: elfgum/RIFE-4.22.lite) ) # Upscaling enable_upscaling: bool = False upscaling_model_path: str | None = ( None # local .pth, HF repo ID, or repo_id:filename (default: ai-forever/Real-ESRGAN) ) upscaling_scale: int = 4 # Batch info num_outputs_per_prompt: int = 1 seed: int = 42 generator_device: str = "cuda" # Device for random generator: "cuda" or "cpu" # Original dimensions (before VAE scaling) num_frames: int = 1 # Default for image models num_frames_round_down: bool = ( False # Whether to round down num_frames if it's not divisible by num_gpus ) # Subclasses can set these to provide model-specific default resolutions. # The base __post_init__ will apply them when height/width are not provided. _default_height: ClassVar[int | None] = None _default_width: ClassVar[int | None] = None height: int | None = None width: int | None = None fps: int = 24 # Resolution validation supported_resolutions: list[tuple[int, int]] | None = ( None # None means all resolutions allowed ) # Denoising parameters num_inference_steps: int = None guidance_scale: float = 1.0 guidance_scale_2: float = None true_cfg_scale: float = None # for CFG vs guidance distillation (e.g., QwenImage) guidance_rescale: float = 0.0 cfg_normalization: float | bool = 0.0 boundary_ratio: float | None = None # TeaCache parameters enable_teacache: bool = False teacache_params: Any = ( None # TeaCacheParams or WanTeaCacheParams, set by model-specific subclass ) # Profiling profile: bool = False num_profiled_timesteps: int = 5 profile_all_stages: bool = False # Debugging debug: bool = False perf_dump_path: str | None = None # Misc save_output: bool = True return_frames: bool = False return_trajectory_latents: bool = False # returns all latents for each timestep return_trajectory_decoded: bool = False # returns decoded latents for each timestep # if True, disallow user params to override subclass-defined protected fields no_override_protected_fields: bool = False # whether to adjust num_frames for multi-GPU friendly splitting (default: True) adjust_frames: bool = True # if True, suppress verbose logging for this request suppress_logs: bool = False return_file_paths_only: bool = True enable_sequence_shard: bool | None = None def _set_output_file_ext(self): # add extension if needed if not any( self.output_file_name.endswith(ext) for ext in [".mp4", ".jpg", ".png", ".webp", ".obj", ".glb"] ): self.output_file_name = ( f"{self.output_file_name}.{self.data_type.get_default_extension()}" ) def _set_output_file_name(self): # settle output_file_name if ( self.output_file_name is None and self.prompt and isinstance(self.prompt, str) ): # generate a random filename # get a hash of current params params_dict = dataclasses.asdict(self) # Avoid recursion params_dict["output_file_name"] = "" # Convert to a stable JSON string params_str = json.dumps(_json_safe(params_dict), sort_keys=True) # Create a hash hasher = hashlib.sha256() hasher.update(params_str.encode("utf-8")) param_hash = hasher.hexdigest()[:8] timestamp = time.strftime("%Y%m%d-%H%M%S") base = f"{self.prompt[:100]}_{timestamp}_{param_hash}" self.output_file_name = base if self.output_file_name is None: timestamp = time.strftime("%Y%m%d-%H%M%S") self.output_file_name = f"output_{timestamp}" self.output_file_name = _sanitize_filename(self.output_file_name) # Ensure a proper extension is present self._set_output_file_ext() def __post_init__(self) -> None: assert self.num_frames >= 1 if self.width is None and self._default_width is not None: self.width = self._default_width if self.height is None and self._default_height is not None: self.height = self._default_height # Handle output_quality to output_compression conversion if self.output_compression is None and self.output_quality is not None: self.output_compression = self._adjust_output_quality( self.output_quality, self.data_type ) self._validate() # Allow env var to override num_inference_steps (for faster CI testing on AMD) env_steps = os.environ.get("SGLANG_TEST_NUM_INFERENCE_STEPS") if env_steps is not None and self.num_inference_steps is not None: self.num_inference_steps = int(env_steps) def _adjust_output_quality(self, output_quality: str, data_type: DataType) -> int: """Convert output_quality string to compression level.""" output_quality_mapper = {"maximum": 100, "high": 90, "medium": 55, "low": 35} if output_quality == "default": return 50 if data_type == DataType.VIDEO else 75 return output_quality_mapper.get(output_quality) def _validate(self): """ check if the sampling params is correct by itself """ if self.prompt_path and not self.prompt_path.endswith(".txt"): raise ValueError( f"prompt_path must be a txt file, got {self.prompt_path!r}" ) # These are always required to be sane regardless of pipeline. if ( not isinstance(self.num_outputs_per_prompt, int) or self.num_outputs_per_prompt <= 0 ): raise ValueError( f"num_outputs_per_prompt must be a positive int, got {self.num_outputs_per_prompt!r}" ) # Used by seconds() and video writer; fps <= 0 is always invalid. if not isinstance(self.fps, int) or self.fps <= 0: raise ValueError(f"fps must be a positive int, got {self.fps!r}") # num_frames is already asserted in __post_init__, but keep a friendly error here too # (e.g., when validation is triggered from other code paths). if not isinstance(self.num_frames, int) or self.num_frames <= 0: raise ValueError( f"num_frames must be a positive int, got {self.num_frames!r}" ) if self.num_inference_steps is not None: if ( not isinstance(self.num_inference_steps, int) or self.num_inference_steps <= 0 ): raise ValueError( f"num_inference_steps must be a positive int, got {self.num_inference_steps!r}" ) # Numeric hyperparams should not be NaN/Inf and should be within basic ranges. # Note: bool is a subclass of int; reject it explicitly to avoid silent surprises. def _finite_non_negative_float( name: str, value: Any, allow_none: bool = True ) -> None: if value is None and allow_none: return if isinstance(value, bool) or not isinstance(value, (int, float)): raise ValueError(f"{name} must be a number, got {value!r}") if not math.isfinite(float(value)): raise ValueError(f"{name} must be finite, got {value!r}") if float(value) < 0.0: raise ValueError(f"{name} must be non-negative, got {value!r}") _finite_non_negative_float( "guidance_scale", self.guidance_scale, allow_none=True ) _finite_non_negative_float( "guidance_scale_2", self.guidance_scale_2, allow_none=True ) _finite_non_negative_float( "true_cfg_scale", self.true_cfg_scale, allow_none=True ) _finite_non_negative_float( "guidance_rescale", self.guidance_rescale, allow_none=False ) if self.cfg_normalization is None: self.cfg_normalization = 0.0 elif isinstance(self.cfg_normalization, bool): self.cfg_normalization = 1.0 if self.cfg_normalization else 0.0 if self.boundary_ratio is not None: if isinstance(self.boundary_ratio, bool) or not isinstance( self.boundary_ratio, (int, float) ): raise ValueError( f"boundary_ratio must be a number, got {self.boundary_ratio!r}" ) if not math.isfinite(float(self.boundary_ratio)): raise ValueError( f"boundary_ratio must be finite, got {self.boundary_ratio!r}" ) if not (0.0 <= float(self.boundary_ratio) <= 1.0): raise ValueError( f"boundary_ratio must be within [0, 1], got {self.boundary_ratio!r}" ) def check_sampling_param(self): # Keep backward-compatibility for old call sites. self._validate() def _validate_with_pipeline_config(self, pipeline_config): """ check if the sampling params is compatible and valid with server_args """ if pipeline_config.task_type.requires_image_input(): # requires image input if self.image_path is None: raise ValueError( f"Served model with task type '{pipeline_config.task_type.name}' requires an 'image_path' input, but none was provided" ) if not pipeline_config.task_type.accepts_image_input(): # does not support image input if self.image_path is not None: raise ValueError( f"input_reference is not supported for {pipeline_config.task_type.name} models." ) def _adjust( self, server_args, ): """ final adjustment, called after merged with user params """ expand_path_fields(self) # TODO: SamplingParams should not rely on ServerArgs pipeline_config = server_args.pipeline_config if self.guidance_scale is None: try: from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( Hunyuan3D2PipelineConfig, ) if isinstance(pipeline_config, Hunyuan3D2PipelineConfig): self.guidance_scale = pipeline_config.guidance_scale else: self.guidance_scale = 1.0 except ImportError: self.guidance_scale = 1.0 self.data_type = server_args.pipeline_config.task_type.data_type() if self.output_path is None: if server_args.output_path is not None: self.output_path = server_args.output_path logger.debug( f"Overriding output_path with server configuration: {self.output_path}" ) else: self.save_output = False # Process negative prompt if self.negative_prompt is not None and not self.negative_prompt.isspace(): # avoid stripping default negative prompt: ' ' for qwen-image self.negative_prompt = self.negative_prompt.strip() # Validate dimensions if self.num_frames <= 0: raise ValueError( f"height, width, and num_frames must be positive integers, got " f"height={self.height}, width={self.width}, " f"num_frames={self.num_frames}" ) # Validate resolution against pipeline-specific supported resolutions if self.height is None and self.width is None: if self.supported_resolutions is not None: self.width, self.height = self.supported_resolutions[0] logger.info( f"Resolution unspecified, using default: {self.supported_resolutions[0]}" ) if self.height is not None and self.width is not None: if self.supported_resolutions is not None: if (self.width, self.height) not in self.supported_resolutions: supported_str = ", ".join( [f"{w}x{h}" for w, h in self.supported_resolutions] ) error_msg = ( f"Unsupported resolution: {self.width}x{self.height}, output quality may suffer. " f"Supported resolutions: {supported_str}" ) logger.warning(error_msg) pipeline_name_lower = server_args.pipeline_config.__class__.__name__.lower() if ("wan" in pipeline_name_lower or "helios" in pipeline_name_lower) and ( self.enable_sequence_shard is None or self.enable_sequence_shard ): self.enable_sequence_shard = True logger.debug("Automatically enabled enable_sequence_shard") else: self.enable_sequence_shard = False if self.enable_sequence_shard: self.adjust_frames = False logger.info( f"Sequence dimension shard is enabled, disabling frame adjustment for better performance" ) if pipeline_config.task_type.is_image_gen(): # settle num_frames if not server_args.pipeline_config.allow_set_num_frames(): logger.debug(f"Setting `num_frames` to 1 for image generation model") self.num_frames = 1 else: # mandatory frame adjusting logic, mod # NOTE: We must apply adjust_num_frames BEFORE the SP alignment logic below. # If we apply it after, adjust_num_frames might modify the frame count # and break the divisibility constraint (alignment) required by num_gpus. original_num_frames = self.num_frames self.num_frames = server_args.pipeline_config.adjust_num_frames( original_num_frames ) logger.info( "Adjusting number of frames from %s to %s based on model", original_num_frames, self.num_frames, ) if self.adjust_frames: # Adjust number of frames based on number of GPUs for video task use_temporal_scaling_frames = ( pipeline_config.vae_config.use_temporal_scaling_frames ) num_frames = self.num_frames num_gpus = server_args.num_gpus temporal_scale_factor = ( pipeline_config.vae_config.arch_config.temporal_compression_ratio ) if use_temporal_scaling_frames: orig_latent_num_frames = ( num_frames - 1 ) // temporal_scale_factor + 1 else: orig_latent_num_frames = num_frames if orig_latent_num_frames % server_args.num_gpus != 0: # Adjust latent frames to be divisible by number of GPUs if self.num_frames_round_down: # Ensure we have at least 1 batch per GPU new_latent_num_frames = ( max(1, (orig_latent_num_frames // num_gpus)) * num_gpus ) else: new_latent_num_frames = ( math.ceil(orig_latent_num_frames / num_gpus) * num_gpus ) if use_temporal_scaling_frames: # Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor new_num_frames = ( new_latent_num_frames - 1 ) * temporal_scale_factor + 1 else: new_num_frames = new_latent_num_frames logger.info( "Adjusting number of frames from %s to %s based on number of GPUs (%s)", self.num_frames, new_num_frames, server_args.num_gpus, ) self.num_frames = new_num_frames if not server_args.comfyui_mode: self._set_output_file_name() @classmethod def from_pretrained(cls, model_path: str, **kwargs) -> "SamplingParams": from sglang.multimodal_gen.registry import get_model_info backend = kwargs.pop("backend", None) model_id = kwargs.pop("model_id", None) model_info = get_model_info(model_path, backend=backend, model_id=model_id) sampling_params: SamplingParams = model_info.sampling_param_cls(**kwargs) return sampling_params @staticmethod def from_user_sampling_params_args( model_path: str, server_args: "ServerArgs", *args, **kwargs ): try: sampling_params = SamplingParams.from_pretrained( model_path, backend=server_args.backend, model_id=server_args.model_id ) except (AttributeError, ValueError) as e: # Handle safetensors files or other cases where model_index.json is not available # Use appropriate SamplingParams based on pipeline_class_name from registry if os.path.isfile(model_path) and model_path.endswith(".safetensors"): # Determine which sampling params to use based on pipeline_class_name pipeline_class_name = getattr(server_args, "pipeline_class_name", None) # Try to get SamplingParams from registry from sglang.multimodal_gen.registry import get_pipeline_config_classes config_classes = ( get_pipeline_config_classes(pipeline_class_name) if pipeline_class_name else None ) if config_classes is not None: _, sampling_params_cls = config_classes try: sampling_params = sampling_params_cls() logger.info( f"Using {sampling_params_cls.__name__} for {pipeline_class_name} safetensors file (no model_index.json): %s", model_path, ) except Exception as import_error: logger.warning( f"Failed to instantiate {sampling_params_cls.__name__}: {import_error}. " "Using default SamplingParams" ) sampling_params = SamplingParams() else: raise ValueError( f"Could not get pipeline config classes for {pipeline_class_name}" ) else: # Re-raise if it's not a safetensors file issue raise user_kwargs = dict(kwargs) user_kwargs.pop("diffusers_kwargs", None) user_sampling_params = SamplingParams(*args, **user_kwargs) # TODO: refactor sampling_params._merge_with_user_params( user_sampling_params, explicit_fields=set(user_kwargs.keys()) ) sampling_params._adjust(server_args) sampling_params._validate_with_pipeline_config(server_args.pipeline_config) return sampling_params def output_size_str(self) -> str: return f"{self.width}x{self.height}" def seconds(self) -> float: return self.num_frames / self.fps @staticmethod def add_cli_args(parser: Any) -> Any: """Add CLI arguments for SamplingParam fields""" def add_argument(*name_or_flags, **kwargs): kwargs.setdefault("default", argparse.SUPPRESS) return parser.add_argument(*name_or_flags, **kwargs) add_argument("--data-type", type=str, nargs="+") add_argument( "--num-frames-round-down", action="store_true", ) add_argument( "--enable-teacache", action="store_true", ) # profiling add_argument( "--profile", action="store_true", help="Enable torch profiler for denoising stage", ) add_argument( "--num-profiled-timesteps", type=int, help="Number of timesteps to profile after warmup", ) add_argument( "--profile-all-stages", action="store_true", dest="profile_all_stages", help="Used with --profile, profile all pipeline stages", ) add_argument( "--debug", action="store_true", help="", ) add_argument( "--prompt", type=str, nargs="+", help="Text prompt(s) for generation. Use space-separated values for multiple prompts, e.g., --prompt 'prompt 1' 'prompt 2'", ) add_argument( "--negative-prompt", type=str, help="Negative text prompt for generation", ) add_argument( "--prompt-path", type=str, help="Path to a text file containing the prompt", ) add_argument( "--output-file-name", type=str, help="Name of the output file", ) add_argument( "--output-quality", type=str, help="Output quality setting (default, low, medium, high, maximum)", ) add_argument( "--output-compression", type=int, help="Output compression level (0-100, higher means better quality but larger file size)", ) add_argument( "--num-outputs-per-prompt", type=int, help="Number of outputs to generate per prompt", ) add_argument( "--seed", type=int, help="Random seed for generation", ) add_argument( "--generator-device", type=str, choices=["cuda", "musa", "cpu"], help="Device for random generator (cuda, musa or cpu). Default: cuda", ) add_argument( "--num-frames", type=int, help="Number of frames to generate", ) add_argument( "--height", type=int, help="Height of generated output", ) add_argument( "--width", type=int, help="Width of generated output", ) # resolution shortcuts add_argument( "--4k", action="store_true", dest="resolution_4k", help="Set resolution to 4K (3840x2160)", ) add_argument( "--2k", action="store_true", dest="resolution_2k", help="Set resolution to 2K (2560x1440)", ) add_argument( "--1080p", action="store_true", dest="resolution_1080p", help="Set resolution to 1080p (1920x1080)", ) add_argument( "--720p", action="store_true", dest="resolution_720p", help="Set resolution to 720p (1280x720)", ) add_argument( "--fps", type=int, help="Frames per second for saved output", ) add_argument( "--num-inference-steps", type=int, help="Number of denoising steps", ) add_argument( "--guidance-scale", type=float, help="Classifier-free guidance scale", ) add_argument( "--guidance-scale-2", type=float, dest="guidance_scale_2", help="Secondary guidance scale for dual-guidance models (e.g., Wan low-noise expert)", ) add_argument( "--guidance-rescale", type=float, help="Guidance rescale factor", ) add_argument( "--cfg-normalization", type=float, dest="cfg_normalization", help=("CFG renormalization factor (for Z-Image). "), ) add_argument( "--boundary-ratio", type=float, help="Boundary timestep ratio", ) add_argument( "--save-output", action="store_true", help="Whether to save the output to disk", ) add_argument( "--no-save-output", action="store_false", dest="save_output", help="Don't save the output to disk", ) add_argument( "--return-frames", action="store_true", help="Whether to return the raw frames", ) add_argument( "--image-path", type=str, nargs="+", help=( "Path(s) to input image(s) for image-to-image / image-to-video " "generation. For multiple images, pass them as space-separated " "values, e.g.: " '--image-path "img1.png" "img2.png"' ), ) add_argument( "--moba-config-path", type=str, help="Path to a JSON file containing V-MoBA specific configurations.", ) add_argument( "--return-trajectory-latents", action="store_true", help="Whether to return the trajectory", ) add_argument( "--return-trajectory-decoded", action="store_true", help="Whether to return the decoded trajectory", ) add_argument( "--diffusers-kwargs", type=str, help="JSON string of extra kwargs to pass to diffusers pipeline. " 'Example: \'{"output_type": "latent", "clip_skip": 2}\'', ) add_argument( "--no-override-protected-fields", action="store_true", help=( "If set, disallow user params to override fields defined in subclasses." ), ) add_argument( "--adjust-frames", action=StoreBoolean, help=( "Enable/disable adjusting num_frames to evenly split latent frames across GPUs " "and satisfy model temporal constraints. If disabled, tokens might be padded for SP." "Default: true. Examples: --adjust-frames, --adjust-frames true, --adjust-frames false." ), ) add_argument( "--return-file-paths-only", action=StoreBoolean, help="If set, output file will be saved early to get a performance boost, while output tensors will not be returned.", ) add_argument( "--enable-sequence-shard", action=StoreBoolean, help="Enable sequence dimension shard with sequence parallelism.", ) add_argument( "--enable-frame-interpolation", action="store_true", help="Enable post-generation frame interpolation using RIFE 4.22.lite.", ) add_argument( "--frame-interpolation-exp", type=int, help="Frame interpolation exponent: 1=2x, 2=4x (default: 1).", ) add_argument( "--frame-interpolation-scale", type=float, help="RIFE inference scale factor (default: 1.0; use 0.5 for high-res).", ) add_argument( "--frame-interpolation-model-path", type=str, help="Local directory or HuggingFace repo ID containing RIFE flownet.pkl weights " "(default: elfgum/RIFE-4.22.lite, downloaded automatically). " "Only RIFE 4.22.lite architecture is supported; other RIFE versions or " "frame interpolation models are not compatible.", ) add_argument( "--enable-upscaling", action="store_true", help="Enable post-generation upscaling using Real-ESRGAN.", ) add_argument( "--upscaling-model-path", type=str, help="Local .pth file, HuggingFace repo ID, or repo_id:filename for Real-ESRGAN weights " "(default: ai-forever/Real-ESRGAN with RealESRGAN_x4.pth). " "Only RRDBNet (e.g. RealESRGAN_x4plus) and SRVGGNetCompact (e.g. realesr-animevideov3) " "architectures are supported; other super-resolution models are not compatible. " "Use 'repo_id:filename' to specify a custom weight file from a HF repo.", ) add_argument( "--upscaling-scale", type=int, help="Upscaling factor (default: 4).", ) return parser @classmethod def get_cli_args(cls, args: argparse.Namespace): # handle resolution shortcuts if hasattr(args, "resolution_4k") and args.resolution_4k: args.width = 3840 args.height = 2160 elif hasattr(args, "resolution_2k") and args.resolution_2k: args.width = 2560 args.height = 1440 elif hasattr(args, "resolution_1080p") and args.resolution_1080p: args.width = 1920 args.height = 1080 elif hasattr(args, "resolution_720p") and args.resolution_720p: args.width = 1280 args.height = 720 sampling_params_fields = {attr.name for attr in dataclasses.fields(cls)} args_attrs = set(vars(args).keys()) attrs = sampling_params_fields & args_attrs return { attr: getattr(args, attr) for attr in attrs if hasattr(args, attr) and getattr(args, attr) is not None } def output_file_path(self): if self.output_path is None: return None return os.path.join(self.output_path, self.output_file_name) def _merge_with_user_params( self, user_params: "SamplingParams", explicit_fields: set[str] | None = None, ): """ Merges parameters from a user-provided SamplingParams object. Args: explicit_fields: field names explicitly set by the user (e.g. from CLI kwargs). These are always treated as user-modified even when their value matches the base-class default. """ if user_params is None: return predefined_fields = set(type(self).__annotations__.keys()) # global switch: if True, allow overriding protected fields allow_override_protected = not user_params.no_override_protected_fields for field in dataclasses.fields(user_params): field_name = field.name user_value = getattr(user_params, field_name) default_class_value = getattr(SamplingParams, field_name) is_user_modified = user_value != default_class_value or ( explicit_fields is not None and field_name in explicit_fields ) is_protected_field = field_name in predefined_fields if is_user_modified and ( allow_override_protected or not is_protected_field ): setattr(self, field_name, user_value) self.__post_init__() @property def n_tokens(self) -> int: # Calculate latent sizes if self.height and self.width: latents_size = [ (self.num_frames - 1) // 4 + 1, self.height // 8, self.width // 8, ] n_tokens = latents_size[0] * latents_size[1] * latents_size[2] else: n_tokens = -1 return n_tokens @dataclass class CacheParams: cache_type: str = "none" ================================================ FILE: python/sglang/multimodal_gen/configs/sample/sana.py ================================================ # SPDX-License-Identifier: Apache-2.0 """Sampling parameters for SANA image generation (T2I).""" from dataclasses import dataclass from sglang.multimodal_gen.configs.sample.sampling_params import ( DataType, SamplingParams, ) @dataclass class SanaSamplingParams(SamplingParams): """Defaults for SANA 1.5 1024px variant. guidance_scale=4.5 enables standard classifier-free guidance. """ data_type: DataType = DataType.IMAGE num_frames: int = 1 guidance_scale: float = 4.5 num_inference_steps: int = 20 height: int = 1024 width: int = 1024 negative_prompt: str = ( "low quality, low resolution, blurry, overexposed, underexposed, " "distorted, deformed, disfigured, bad anatomy, extra limbs, " "watermark, text, signature, ugly, noisy, artifacts" ) ================================================ FILE: python/sglang/multimodal_gen/configs/sample/teacache.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.sample.sampling_params import CacheParams @dataclass class TeaCacheParams(CacheParams): cache_type: str = "teacache" teacache_thresh: float = 0.0 coefficients: list[float] = field(default_factory=list) @dataclass class WanTeaCacheParams(CacheParams): # Unfortunately, TeaCache is very different for Wan than other models cache_type: str = "teacache" teacache_thresh: float = 0.0 use_ret_steps: bool = True ret_steps_coeffs: list[float] = field(default_factory=list) non_ret_steps_coeffs: list[float] = field(default_factory=list) @property def coefficients(self) -> list[float]: if self.use_ret_steps: return self.ret_steps_coeffs else: return self.non_ret_steps_coeffs @property def ret_steps(self) -> int: if self.use_ret_steps: return 5 * 2 else: return 1 * 2 def get_cutoff_steps(self, num_inference_steps: int) -> int: if self.use_ret_steps: return num_inference_steps * 2 else: return num_inference_steps * 2 - 2 ================================================ FILE: python/sglang/multimodal_gen/configs/sample/wan.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.configs.sample.teacache import WanTeaCacheParams @dataclass class WanT2V_1_3B_SamplingParams(SamplingParams): # Video parameters height: int = 480 width: int = 832 num_frames: int = 81 fps: int = 16 # Denoising stage guidance_scale: float = 3.0 negative_prompt: str = ( "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" ) num_inference_steps: int = 50 # Wan T2V 1.3B supported resolutions supported_resolutions: list[tuple[int, int]] | None = field( default_factory=lambda: [ (832, 480), # 16:9 (480, 832), # 9:16 ] ) teacache_params: WanTeaCacheParams = field( default_factory=lambda: WanTeaCacheParams( teacache_thresh=0.08, ret_steps_coeffs=[ -5.21862437e04, 9.23041404e03, -5.28275948e02, 1.36987616e01, -4.99875664e-02, ], non_ret_steps_coeffs=[ 2.39676752e03, -1.31110545e03, 2.01331979e02, -8.29855975e00, 1.37887774e-01, ], ) ) @dataclass class WanT2V_14B_SamplingParams(SamplingParams): # Video parameters height: int = 720 width: int = 1280 num_frames: int = 81 fps: int = 16 # Denoising stage guidance_scale: float = 5.0 negative_prompt: str = ( "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" ) num_inference_steps: int = 50 # Wan T2V 14B supported resolutions supported_resolutions: list[tuple[int, int]] | None = field( default_factory=lambda: [ (1280, 720), # 16:9 (720, 1280), # 9:16 (832, 480), # 16:9 (480, 832), # 9:16 ] ) teacache_params: WanTeaCacheParams = field( default_factory=lambda: WanTeaCacheParams( teacache_thresh=0.20, use_ret_steps=False, ret_steps_coeffs=[ -3.03318725e05, 4.90537029e04, -2.65530556e03, 5.87365115e01, -3.15583525e-01, ], non_ret_steps_coeffs=[ -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, -13.02252404, ], ) ) @dataclass class WanI2V_14B_480P_SamplingParam(WanT2V_1_3B_SamplingParams): # Denoising stage guidance_scale: float = 5.0 num_inference_steps: int = 50 # num_inference_steps: int = 40 # Wan I2V 480P supported resolutions (override parent) supported_resolutions: list[tuple[int, int]] | None = field( default_factory=lambda: [ (832, 480), # 16:9 (480, 832), # 9:16 ] ) teacache_params: WanTeaCacheParams = field( default_factory=lambda: WanTeaCacheParams( teacache_thresh=0.26, ret_steps_coeffs=[ -3.03318725e05, 4.90537029e04, -2.65530556e03, 5.87365115e01, -3.15583525e-01, ], non_ret_steps_coeffs=[ -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, -13.02252404, ], ) ) @dataclass class WanI2V_14B_720P_SamplingParam(WanT2V_14B_SamplingParams): # Denoising stage guidance_scale: float = 5.0 num_inference_steps: int = 50 # num_inference_steps: int = 40 # Wan I2V 720P supported resolutions (override parent) supported_resolutions: list[tuple[int, int]] | None = field( default_factory=lambda: [ (1280, 720), # 16:9 (720, 1280), # 9:16 (832, 480), # 16:9 (480, 832), # 9:16 ] ) teacache_params: WanTeaCacheParams = field( default_factory=lambda: WanTeaCacheParams( teacache_thresh=0.3, ret_steps_coeffs=[ -3.03318725e05, 4.90537029e04, -2.65530556e03, 5.87365115e01, -3.15583525e-01, ], non_ret_steps_coeffs=[ -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, -13.02252404, ], ) ) @dataclass class FastWanT2V480PConfig(WanT2V_1_3B_SamplingParams): # DMD parameters # dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522]) num_inference_steps: int = 3 num_frames: int = 61 height: int = 448 width: int = 832 fps: int = 16 # ============================================= # ============= Wan2.1 Fun Models ============= # ============================================= @dataclass class Wan2_1_Fun_1_3B_InP_SamplingParams(SamplingParams): """Sampling parameters for Wan2.1 Fun 1.3B InP model.""" height: int = 480 width: int = 832 num_frames: int = 81 fps: int = 16 negative_prompt: str | None = ( "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" ) guidance_scale: float = 6.0 num_inference_steps: int = 50 # ============================================= # ============= Wan2.2 TI2V Models ============= # ============================================= @dataclass class Wan2_2_Base_SamplingParams(SamplingParams): """Sampling parameters for Wan2.2 TI2V 5B model.""" negative_prompt: str | None = ( "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" ) # TODO(Wan2.2): TeaCache coefficients need to be calibrated for Wan2.2 by # profiling L1 distances across timesteps. Until then, teacache_params is None # and enable_teacache will be accepted but silently no-op. # Consider using Cache-DiT (SGLANG_CACHE_DIT_ENABLED=1) as an alternative. @dataclass class Wan2_2_TI2V_5B_SamplingParam(Wan2_2_Base_SamplingParams): """Sampling parameters for Wan2.2 TI2V 5B model.""" height: int = 704 width: int = 1280 num_frames: int = 121 fps: int = 24 guidance_scale: float = 5.0 num_inference_steps: int = 50 # Wan2.2 TI2V 5B supported resolutions supported_resolutions: list[tuple[int, int]] | None = field( default_factory=lambda: [ (1280, 704), # 16:9-ish (704, 1280), # 9:16-ish ] ) @dataclass class Wan2_2_T2V_A14B_SamplingParam(Wan2_2_Base_SamplingParams): guidance_scale: float = 4.0 # high_noise guidance_scale_2: float = 3.0 # low_noise num_inference_steps: int = 40 fps: int = 16 num_frames: int = 81 # Wan2.2 T2V A14B supported resolutions supported_resolutions: list[tuple[int, int]] | None = field( default_factory=lambda: [ (1280, 720), # 16:9 (720, 1280), # 9:16 (832, 480), # 16:9 (480, 832), # 9:16 ] ) @dataclass class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParams): guidance_scale: float = 3.5 # high_noise guidance_scale_2: float = 3.5 # low_noise num_inference_steps: int = 40 fps: int = 16 num_frames: int = 81 # Wan2.2 I2V A14B supported resolutions supported_resolutions: list[tuple[int, int]] | None = field( default_factory=lambda: [ (1280, 720), # 16:9 (720, 1280), # 9:16 (832, 480), # 16:9 (480, 832), # 9:16 ] ) @dataclass class Turbo_Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParams): guidance_scale: float = 3.5 # high_noise guidance_scale_2: float = 3.5 # low_noise num_inference_steps: int = 4 fps: int = 16 # ============================================= # ============= Causal Self-Forcing ============= # ============================================= @dataclass class SelfForcingWanT2V480PConfig(WanT2V_1_3B_SamplingParams): pass ================================================ FILE: python/sglang/multimodal_gen/configs/sample/zimage.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams @dataclass class ZImageTurboSamplingParams(SamplingParams): num_inference_steps: int = 9 num_frames: int = 1 negative_prompt: str = None # height: int = 720 # width: int = 1280 # fps: int = 24 guidance_scale: float = 0.0 cfg_normalization: float | bool = False teacache_params: TeaCacheParams = field( default_factory=lambda: TeaCacheParams( teacache_thresh=0.15, coefficients=[ 7.33226126e02, -4.01131952e02, 6.75869174e01, -3.14987800e00, 9.61237896e-02, ], ) ) @dataclass class ZImageSamplingParams(SamplingParams): num_inference_steps: int = 50 num_frames: int = 1 negative_prompt: str = " " guidance_scale: float = 5.0 cfg_normalization: float | bool = True ================================================ FILE: python/sglang/multimodal_gen/configs/utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import argparse from typing import Any def update_config_from_args( config: Any, args_dict: dict[str, Any], prefix: str = "", pop_args: bool = False ) -> bool: """ Update configuration object from arguments dictionary. Args: config: The configuration object to update args_dict: Dictionary containing arguments prefix: Prefix for the configuration parameters in the args_dict. If None, assumes direct attribute mapping without prefix. """ # Handle top-level attributes (no prefix) args_not_to_remove = [ "model_path", ] args_to_remove = [] if prefix.strip() == "": for key, value in args_dict.items(): if hasattr(config, key) and value is not None: if key == "text_encoder_precisions" and isinstance(value, list): setattr(config, key, tuple(value)) else: setattr(config, key, value) if pop_args: args_to_remove.append(key) else: # Handle nested attributes with prefix prefix_with_dot = f"{prefix}." for key, value in args_dict.items(): if key.startswith(prefix_with_dot) and value is not None: attr_name = key[len(prefix_with_dot) :] if hasattr(config, attr_name): setattr(config, attr_name, value) if pop_args: args_to_remove.append(key) if pop_args: for key in args_to_remove: if key not in args_not_to_remove: args_dict.pop(key) return len(args_to_remove) > 0 def clean_cli_args(args: argparse.Namespace) -> dict[str, Any]: """ Clean the arguments by removing the ones that not explicitly provided by the user. """ provided_args = {} for k, v in vars(args).items(): if v is not None and hasattr(args, "_provided") and k in args._provided: provided_args[k] = v return provided_args ================================================ FILE: python/sglang/multimodal_gen/csrc/attn/vmoba_attn/README.md ================================================ # Attention Kernel Used in SGLang diffusion ## VMoBA: Mixture-of-Block Attention for Video Diffusion Models (VMoBA) ### Installation Please ensure that you have installed FlashAttention version **2.7.1 or higher**, as some interfaces have changed in recent releases. ### Usage You can use `moba_attn_varlen` in the following ways: **Install from source:** ```bash python setup.py install ``` **Import after installation:** ```python from vmoba import moba_attn_varlen ``` **Or import directly from the project root:** ```python from csrc.attn.vmoba_attn.vmoba import moba_attn_varlen ``` ### Verify if you have successfully installed ```bash python csrc/attn/vmoba_attn/vmoba/vmoba.py ``` ================================================ FILE: python/sglang/multimodal_gen/csrc/attn/vmoba_attn/setup.py ================================================ # SPDX-License-Identifier: Apache-2.0 from setuptools import find_packages, setup PACKAGE_NAME = "vmoba" VERSION = "0.0.0" AUTHOR = "JianzongWu" DESCRIPTION = "VMoBA: Mixture-of-Block Attention for Video Diffusion Models" URL = "https://github.com/KwaiVGI/VMoBA" setup( name=PACKAGE_NAME, version=VERSION, author=AUTHOR, description=DESCRIPTION, url=URL, packages=find_packages(), classifiers=[ "Programming Language :: Python :: 3", "License :: OSI Approved :: Apache Software License", ], python_requires=">=3.12", install_requires=[ "flash-attn >= 2.7.1", ], ) ================================================ FILE: python/sglang/multimodal_gen/csrc/attn/vmoba_attn/tests/test_vmoba_attn.py ================================================ # SPDX-License-Identifier: Apache-2.0 import random import pytest import torch from sglang.multimodal_gen.csrc.attn.vmoba_attn.vmoba import moba_attn_varlen def generate_test_data( batch_size, total_seqlen, num_heads, head_dim, dtype, device="cuda" ): """ Generates random data for testing the variable-length attention function. """ torch.manual_seed(42) random.seed(42) torch.cuda.manual_seed_all(42) # Generate sequence lengths for each item in the batch if batch_size > 1: # Ensure sequence lengths are reasonably distributed avg_seqlen = total_seqlen // batch_size seqlens = [ random.randint(avg_seqlen // 2, avg_seqlen + avg_seqlen // 2) for _ in range(batch_size - 1) ] remaining_len = total_seqlen - sum(seqlens) if remaining_len > 0: seqlens.append(remaining_len) else: # Adjust if sum exceeds total_seqlen seqlens.append(avg_seqlen) current_sum = sum(seqlens) seqlens[-1] -= current_sum - total_seqlen # Ensure all lengths are positive seqlens = [max(1, s) for s in seqlens] # Final adjustment to match total_seqlen seqlens[-1] += total_seqlen - sum(seqlens) else: seqlens = [total_seqlen] cu_seqlens = torch.tensor( [0] + list(torch.cumsum(torch.tensor(seqlens), 0)), device=device, dtype=torch.int32, ) max_seqlen = max(seqlens) if seqlens else 0 q = torch.randn( (total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False, ) k = torch.randn( (total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False, ) v = torch.randn( (total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False, ) return q, k, v, cu_seqlens, max_seqlen @pytest.mark.parametrize("batch_size", [1, 2]) @pytest.mark.parametrize("total_seqlen", [512, 1024]) @pytest.mark.parametrize("num_heads", [8]) @pytest.mark.parametrize("head_dim", [64]) @pytest.mark.parametrize("moba_chunk_size", [64]) @pytest.mark.parametrize("moba_topk", [2, 4]) @pytest.mark.parametrize("select_mode", ["topk", "threshold"]) @pytest.mark.parametrize("threshold_type", ["query_head", "head_global", "overall"]) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) def test_moba_attn_varlen_forward( batch_size, total_seqlen, num_heads, head_dim, moba_chunk_size, moba_topk, select_mode, threshold_type, dtype, ): """ Tests the forward pass of moba_attn_varlen for basic correctness. It checks output shape, dtype, and for the presence of NaNs/Infs. """ if dtype == torch.float32: pytest.skip("float32 is not supported in flash attention") q, k, v, cu_seqlens, max_seqlen = generate_test_data( batch_size, total_seqlen, num_heads, head_dim, dtype ) # Ensure chunk size is not larger than the smallest sequence length min_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).min().item() if moba_chunk_size > min_seqlen: pytest.skip( "moba_chunk_size is larger than the minimum sequence length in the batch" ) try: output = moba_attn_varlen( q=q, k=k, v=v, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, moba_chunk_size=moba_chunk_size, moba_topk=moba_topk, select_mode=select_mode, threshold_type=threshold_type, simsum_threshold=0.5, # A reasonable default for threshold mode ) except Exception as e: pytest.fail(f"moba_attn_varlen forward pass failed with exception: {e}") # 1. Check output shape assert ( output.shape == q.shape ), f"Expected output shape {q.shape}, but got {output.shape}" # 2. Check output dtype assert ( output.dtype == q.dtype ), f"Expected output dtype {q.dtype}, but got {output.dtype}" # 3. Check for NaNs or Infs in the output assert torch.all(torch.isfinite(output)), "Output contains NaN or Inf values" ================================================ FILE: python/sglang/multimodal_gen/csrc/attn/vmoba_attn/vmoba/__init__.py ================================================ # SPDX-License-Identifier: Apache-2.0 from .vmoba import moba_attn_varlen, process_moba_input, process_moba_output ================================================ FILE: python/sglang/multimodal_gen/csrc/attn/vmoba_attn/vmoba/vmoba.py ================================================ # SPDX-License-Identifier: Apache-2.0 # Adapt from https://github.com/KwaiVGI/VMoBA/blob/main/src/vmoba.py import random import time from typing import Tuple import torch try: from flash_attn import ( # Use the new flash attention function flash_attn_varlen_func, ) from flash_attn.flash_attn_interface import ( _flash_attn_varlen_backward, _flash_attn_varlen_forward, ) except ImportError: def _unsupported(*args, **kwargs): raise ImportError( "flash-attn is not installed. Please install it, e.g., `pip install flash-attn`." ) _flash_attn_varlen_forward = _unsupported _flash_attn_varlen_backward = _unsupported flash_attn_varlen_func = _unsupported from functools import lru_cache from einops import rearrange @lru_cache(maxsize=16) def calc_chunks(cu_seqlen, moba_chunk_size): """ Calculate chunk boundaries. For vision tasks we include all chunks (even the last one which might be shorter) so that every chunk can be selected. """ batch_sizes = cu_seqlen[1:] - cu_seqlen[:-1] batch_num_chunk = (batch_sizes + (moba_chunk_size - 1)) // moba_chunk_size cu_num_chunk = torch.ones( batch_num_chunk.numel() + 1, device=cu_seqlen.device, dtype=batch_num_chunk.dtype, ) cu_num_chunk[1:] = batch_num_chunk.cumsum(dim=0) num_chunk = cu_num_chunk[-1] chunk_sizes = torch.full( (num_chunk + 1,), moba_chunk_size, dtype=torch.int32, device=cu_seqlen.device ) chunk_sizes[0] = 0 batch_last_chunk_size = batch_sizes - (batch_num_chunk - 1) * moba_chunk_size chunk_sizes[cu_num_chunk[1:]] = batch_last_chunk_size cu_chunk = chunk_sizes.cumsum(dim=-1, dtype=torch.int32) chunk_to_batch = torch.zeros( (num_chunk,), dtype=torch.int32, device=cu_seqlen.device ) chunk_to_batch[cu_num_chunk[1:-1]] = 1 chunk_to_batch = chunk_to_batch.cumsum(dim=0, dtype=torch.int32) # Do not filter out any chunk filtered_chunk_indices = torch.arange( num_chunk, device=cu_seqlen.device, dtype=torch.int32 ) num_filtered_chunk = num_chunk return cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch # --- Threshold Selection Helper Functions --- def _select_threshold_query_head( gate: torch.Tensor, valid_gate_mask: torch.Tensor, gate_self_chunk_mask: torch.Tensor, simsum_threshold: float, ) -> torch.Tensor: """ Selects chunks for each pair based on threshold. Normalization and sorting happen along the chunk dimension (dim=0). """ C, H, S = gate.shape eps = 1e-6 # LSE‐style normalization per (across chunks) gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min row_min = gate_min_val.amin(dim=0) # (H, S) row_max = gate_masked.amax(dim=0) # (H, S) denom = row_max - row_min denom = torch.where( denom <= eps, torch.ones_like(denom), denom ) # avoid divide‑by‑zero gate_norm = (gate - row_min.unsqueeze(0)) / denom.unsqueeze(0) gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S) # 1) pull out the self‐chunk’s normalized weight for each self_norm = (gate_norm * gate_self_chunk_mask).sum(dim=0) # (H, S) # 2) compute how much more normalized weight we need beyond self total_norm_sum = gate_norm.sum(dim=0) # (H, S) remain_ratio = simsum_threshold - self_norm / (total_norm_sum + eps) # (H, S) remain_ratio = torch.clamp( remain_ratio, min=0.0 ) # if already ≥ thresh, no extra needed # 3) zero out the self‐chunk in a copy, so we only sort “others” others_norm = gate_norm.clone() others_norm[gate_self_chunk_mask] = 0.0 # 4) sort the other chunks by descending norm, per sorted_norm, sorted_idx = torch.sort( others_norm, descending=True, dim=0 ) # (C, H, S) # 5) cumulative‑sum the sorted norms per cumsum_others = sorted_norm.cumsum(dim=0) # (C, H, S) # 6) for each , find the smallest k where cumsum_ratio ≥ remain_ratio ratio = cumsum_others / (total_norm_sum.unsqueeze(0) + eps) # (C, H, S) cond = ratio >= remain_ratio.unsqueeze(0) # (C, H, S) boolean mask any_cond = cond.any(dim=0) # (H, S) # Find the index of the first True value along dim 0. If none, use C-1. cutoff = torch.where( any_cond, cond.float().argmax(dim=0), torch.full_like(any_cond, fill_value=C - 1), ) # (H, S) # 7) build a mask in sorted order up to that cutoff idx_range = torch.arange(C, device=gate.device).view(-1, 1, 1) # (C, 1, 1) sorted_mask = idx_range <= cutoff.unsqueeze(0) # (C, H, S) # 8) scatter it back to original chunk order others_mask = torch.zeros_like(gate, dtype=torch.bool) others_mask.scatter_(0, sorted_idx, sorted_mask) # 9) finally, include every self‐chunk plus all selected others final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask) return final_gate_mask def _select_threshold_block( gate: torch.Tensor, valid_gate_mask: torch.Tensor, gate_self_chunk_mask: torch.Tensor, simsum_threshold: float, ) -> torch.Tensor: """ Selects pairs for each block based on threshold. Normalization and sorting happen across the head and sequence dimensions (dim=1, 2). """ C, H, S = gate.shape HS = H * S eps = 1e-6 # LSE‐style normalization per block (across heads and queries) gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min block_max = gate_masked.amax(dim=(1, 2), keepdim=True) # (C, 1, 1) block_min = gate_min_val.amin(dim=(1, 2), keepdim=True) # (C, 1, 1) block_denom = block_max - block_min block_denom = torch.where( block_denom <= eps, torch.ones_like(block_denom), block_denom ) # (C, 1, 1) gate_norm = (gate - block_min) / block_denom # (C, H, S) gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S) # 1) identify normalized weights of entries that *are* self-chunks (from query perspective) self_norm_entries = gate_norm * gate_self_chunk_mask # (C, H, S) # Sum these weights *per block* self_norm_sum_per_block = self_norm_entries.sum(dim=(1, 2)) # (C,) # 2) compute how much more normalized weight each block needs beyond its self-chunk contributions total_norm_sum_per_block = gate_norm.sum(dim=(1, 2)) # (C,) remain_ratio = simsum_threshold - self_norm_sum_per_block / ( total_norm_sum_per_block + eps ) # (C,) remain_ratio = torch.clamp(remain_ratio, min=0.0) # (C,) # 3) zero out the self‐chunk entries in a copy, so we only sort “others” others_norm = gate_norm.clone() others_norm[gate_self_chunk_mask] = 0.0 # Zero out self entries # 4) sort the other pairs by descending norm, per block others_flat = others_norm.contiguous().view(C, HS) # (C, H*S) sorted_others_flat, sorted_indices_flat = torch.sort( others_flat, dim=1, descending=True ) # (C, H*S) # 5) cumulative‑sum the sorted norms per block cumsum_others_flat = sorted_others_flat.cumsum(dim=1) # (C, H*S) # 6) for each block, find the smallest k where cumsum_ratio ≥ remain_ratio ratio_flat = cumsum_others_flat / ( total_norm_sum_per_block.unsqueeze(1) + eps ) # (C, H*S) cond_flat = ratio_flat >= remain_ratio.unsqueeze(1) # (C, H*S) boolean mask any_cond = cond_flat.any(dim=1) # (C,) # Find the index of the first True value along dim 1. If none, use HS-1. cutoff_flat = torch.where( any_cond, cond_flat.float().argmax(dim=1), torch.full_like(any_cond, fill_value=HS - 1), ) # (C,) # 7) build a mask in sorted order up to that cutoff per block idx_range_flat = torch.arange(HS, device=gate.device).unsqueeze(0) # (1, H*S) sorted_mask_flat = idx_range_flat <= cutoff_flat.unsqueeze(1) # (C, H*S) # 8) scatter it back to original order per block others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C, H*S) others_mask_flat.scatter_(1, sorted_indices_flat, sorted_mask_flat) others_mask = others_mask_flat.view(C, H, S) # (C, H, S) # 9) finally, include every self‐chunk entry plus all selected others final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask) return final_gate_mask def _select_threshold_overall( gate: torch.Tensor, valid_gate_mask: torch.Tensor, gate_self_chunk_mask: torch.Tensor, simsum_threshold: float, ) -> torch.Tensor: """ Selects triplets globally based on threshold. Normalization and sorting happen across all valid entries. """ C, H, S = gate.shape CHS = C * H * S eps = 1e-6 # LSE‐style normalization globally across all valid entries gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min overall_max = gate_masked.max() # scalar overall_min = gate_min_val.min() # scalar overall_denom = overall_max - overall_min overall_denom = torch.where( overall_denom <= eps, torch.tensor(1.0, device=gate.device, dtype=gate.dtype), overall_denom, ) gate_norm = (gate - overall_min) / overall_denom # (C, H, S) gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S) # 1) identify normalized weights of entries that *are* self-chunks self_norm_entries = gate_norm * gate_self_chunk_mask # (C, H, S) # Sum these weights globally self_norm_sum_overall = self_norm_entries.sum() # scalar # 2) compute how much more normalized weight is needed globally beyond self-chunk contributions total_norm_sum_overall = gate_norm.sum() # scalar remain_ratio = simsum_threshold - self_norm_sum_overall / ( total_norm_sum_overall + eps ) # scalar remain_ratio = torch.clamp(remain_ratio, min=0.0) # scalar # 3) zero out the self‐chunk entries in a copy, so we only sort “others” others_norm = gate_norm.clone() others_norm[gate_self_chunk_mask] = 0.0 # Zero out self entries # 4) sort all other entries by descending norm, globally others_flat = others_norm.flatten() # (C*H*S,) valid_others_mask_flat = ( valid_gate_mask.flatten() & ~gate_self_chunk_mask.flatten() ) # Mask for valid, non-self entries # Only sort the valid 'other' entries valid_others_indices = torch.where(valid_others_mask_flat)[0] valid_others_values = others_flat[valid_others_indices] sorted_others_values, sort_perm = torch.sort( valid_others_values, descending=True ) # (N_valid_others,) sorted_original_indices = valid_others_indices[ sort_perm ] # Original indices in C*H*S space, sorted by value # 5) cumulative‑sum the sorted valid 'other' norms globally cumsum_others_values = sorted_others_values.cumsum(dim=0) # (N_valid_others,) # 6) find the smallest k where cumsum_ratio ≥ remain_ratio globally ratio_values = cumsum_others_values / ( total_norm_sum_overall + eps ) # (N_valid_others,) cond_values = ratio_values >= remain_ratio # (N_valid_others,) boolean mask any_cond = cond_values.any() # scalar # Find the index of the first True value in the *sorted* list. If none, use all valid others. cutoff_idx_in_sorted = torch.where( any_cond, cond_values.float().argmax(dim=0), torch.tensor( len(sorted_others_values) - 1, device=gate.device, dtype=torch.long ), ) # 7) build a mask selecting the top-k others based on the cutoff # Select the original indices corresponding to the top entries in the sorted list selected_other_indices = sorted_original_indices[: cutoff_idx_in_sorted + 1] # 8) create the mask in the original flat shape others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C*H*S,) if selected_other_indices.numel() > 0: # Check if any 'other' indices were selected others_mask_flat[selected_other_indices] = True others_mask = others_mask_flat.view(C, H, S) # (C, H, S) # 9) finally, include every self‐chunk entry plus all selected others final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask) return final_gate_mask def _select_threshold_head_global( gate: torch.Tensor, valid_gate_mask: torch.Tensor, gate_self_chunk_mask: torch.Tensor, simsum_threshold: float, ) -> torch.Tensor: """ Selects globally for each head based on threshold. """ C, H, S = gate.shape eps = 1e-6 # 1) LSE‐style normalization per head (across chunks and sequence dims) gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) max_per_head = gate_masked.amax(dim=(0, 2), keepdim=True) # (1, H, 1) min_per_head = gate_min_val.amin(dim=(0, 2), keepdim=True) # (1, H, 1) denom = max_per_head - min_per_head denom = torch.where(denom <= eps, torch.ones_like(denom), denom) gate_norm = (gate - min_per_head) / denom gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S) # 2) sum normalized self‐chunk contributions per head self_norm_sum = (gate_norm * gate_self_chunk_mask).sum(dim=(0, 2)) # (H,) # 3) total normalized sum per head total_norm_sum = gate_norm.sum(dim=(0, 2)) # (H,) # 4) how much more normalized weight needed per head remain_ratio = simsum_threshold - self_norm_sum / (total_norm_sum + eps) # (H,) remain_ratio = torch.clamp(remain_ratio, min=0.0) # 5) zero out self‐chunk entries to focus on "others" others_norm = gate_norm.clone() others_norm[gate_self_chunk_mask] = 0.0 # (C, H, S) # 6) flatten chunk and sequence dims, per head CS = C * S others_flat = others_norm.permute(1, 0, 2).reshape(H, CS) # (H, C*S) valid_flat = ( (valid_gate_mask & ~gate_self_chunk_mask).permute(1, 0, 2).reshape(H, CS) ) # (H, C*S) # 7) vectorized selection of “others” per head masked_flat = torch.where(valid_flat, others_flat, torch.zeros_like(others_flat)) sorted_vals, sorted_idx = torch.sort( masked_flat, dim=1, descending=True ) # (H, C*S) cumsum_vals = sorted_vals.cumsum(dim=1) # (H, C*S) ratio_vals = cumsum_vals / (total_norm_sum.unsqueeze(1) + eps) # (H, C*S) cond = ratio_vals >= remain_ratio.unsqueeze(1) # (H, C*S) has_cutoff = cond.any(dim=1) # (H,) default = torch.full((H,), CS - 1, device=gate.device, dtype=torch.long) cutoff = torch.where(has_cutoff, cond.float().argmax(dim=1), default) # (H,) idx_range = torch.arange(CS, device=gate.device).unsqueeze(0) # (1, C*S) sorted_mask = idx_range <= cutoff.unsqueeze(1) # (H, C*S) selected_flat = torch.zeros_like(valid_flat) # (H, C*S) selected_flat.scatter_(1, sorted_idx, sorted_mask) # (H, C*S) # 8) reshape selection mask back to (C, H, S) others_mask = selected_flat.reshape(H, C, S).permute(1, 0, 2) # (C, H, S) # 9) include self‐chunks plus selected others, and obey valid mask final_gate_mask = valid_gate_mask & (gate_self_chunk_mask | others_mask) return final_gate_mask class MixedAttention(torch.autograd.Function): @staticmethod def forward( ctx, q, k, v, self_attn_cu_seqlen, moba_q, moba_kv, moba_cu_seqlen_q, moba_cu_seqlen_kv, max_seqlen, moba_chunk_size, moba_q_sh_indices, ): ctx.max_seqlen = max_seqlen ctx.moba_chunk_size = moba_chunk_size ctx.softmax_scale = softmax_scale = q.shape[-1] ** (-0.5) # Non-causal self-attention branch # return out, softmax_lse, S_dmask, rng_state self_attn_out_sh, self_attn_lse_hs, _, _ = _flash_attn_varlen_forward( q=q, k=k, v=v, cu_seqlens_q=self_attn_cu_seqlen, cu_seqlens_k=self_attn_cu_seqlen, max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen, softmax_scale=softmax_scale, causal=False, dropout_p=0.0, ) # MOBA attention branch (non-causal) moba_attn_out, moba_attn_lse_hs, _, _ = _flash_attn_varlen_forward( q=moba_q, k=moba_kv[:, 0], v=moba_kv[:, 1], cu_seqlens_q=moba_cu_seqlen_q, cu_seqlens_k=moba_cu_seqlen_kv, max_seqlen_q=max_seqlen, max_seqlen_k=moba_chunk_size, softmax_scale=softmax_scale, causal=False, dropout_p=0.0, ) self_attn_lse_sh = self_attn_lse_hs.t().contiguous() moba_attn_lse = moba_attn_lse_hs.t().contiguous() output = torch.zeros( (q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32 ) output_2d = output.view(-1, q.shape[2]) max_lse_1d = self_attn_lse_sh.view(-1) max_lse_1d = max_lse_1d.index_reduce( 0, moba_q_sh_indices, moba_attn_lse.view(-1), "amax" ) self_attn_lse_sh = self_attn_lse_sh - max_lse_1d.view_as(self_attn_lse_sh) moba_attn_lse = ( moba_attn_lse.view(-1) .sub(max_lse_1d.index_select(0, moba_q_sh_indices)) .reshape_as(moba_attn_lse) ) mixed_attn_se_sh = self_attn_lse_sh.exp() moba_attn_se = moba_attn_lse.exp() mixed_attn_se_sh.view(-1).index_add_( 0, moba_q_sh_indices, moba_attn_se.view(-1) ) mixed_attn_lse_sh = mixed_attn_se_sh.log() # Combine self-attention output factor = (self_attn_lse_sh - mixed_attn_lse_sh).exp() # [S, H] self_attn_out_sh = self_attn_out_sh * factor.unsqueeze(-1) output_2d += self_attn_out_sh.reshape_as(output_2d) # Combine MOBA attention output mixed_attn_lse = ( mixed_attn_lse_sh.view(-1) .index_select(0, moba_q_sh_indices) .view_as(moba_attn_lse) ) factor = (moba_attn_lse - mixed_attn_lse).exp() # [S, H] moba_attn_out = moba_attn_out * factor.unsqueeze(-1) raw_attn_out = moba_attn_out.view(-1, moba_attn_out.shape[-1]) output_2d.index_add_(0, moba_q_sh_indices, raw_attn_out) output = output.to(q.dtype) mixed_attn_lse_sh = mixed_attn_lse_sh + max_lse_1d.view_as(mixed_attn_se_sh) ctx.save_for_backward( output, mixed_attn_lse_sh, q, k, v, self_attn_cu_seqlen, moba_q, moba_kv, moba_cu_seqlen_q, moba_cu_seqlen_kv, moba_q_sh_indices, ) return output @staticmethod def backward(ctx, d_output): max_seqlen = ctx.max_seqlen moba_chunk_size = ctx.moba_chunk_size softmax_scale = ctx.softmax_scale ( output, mixed_attn_vlse_sh, q, k, v, self_attn_cu_seqlen, moba_q, moba_kv, moba_cu_seqlen_q, moba_cu_seqlen_kv, moba_q_sh_indices, ) = ctx.saved_tensors d_output = d_output.contiguous() dq = torch.empty_like(q) dk = torch.empty_like(k) dv = torch.empty_like(v) _ = _flash_attn_varlen_backward( dout=d_output, q=q, k=k, v=v, out=output, softmax_lse=mixed_attn_vlse_sh.t().contiguous(), dq=dq, dk=dk, dv=dv, cu_seqlens_q=self_attn_cu_seqlen, cu_seqlens_k=self_attn_cu_seqlen, max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen, softmax_scale=softmax_scale, causal=False, dropout_p=0.0, softcap=0.0, alibi_slopes=None, deterministic=True, window_size_left=-1, window_size_right=-1, ) headdim = q.shape[-1] d_moba_output = ( d_output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1) ) moba_output = ( output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1) ) mixed_attn_vlse = ( mixed_attn_vlse_sh.view(-1).index_select(0, moba_q_sh_indices).view(1, -1) ) dmq = torch.empty_like(moba_q) dmkv = torch.empty_like(moba_kv) _ = _flash_attn_varlen_backward( dout=d_moba_output, q=moba_q, k=moba_kv[:, 0], v=moba_kv[:, 1], out=moba_output, softmax_lse=mixed_attn_vlse, dq=dmq, dk=dmkv[:, 0], dv=dmkv[:, 1], cu_seqlens_q=moba_cu_seqlen_q, cu_seqlens_k=moba_cu_seqlen_kv, max_seqlen_q=max_seqlen, max_seqlen_k=moba_chunk_size, softmax_scale=softmax_scale, causal=False, dropout_p=0.0, softcap=0.0, alibi_slopes=None, deterministic=True, window_size_left=-1, window_size_right=-1, ) return dq, dk, dv, None, dmq, dmkv, None, None, None, None, None def moba_attn_varlen( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, moba_chunk_size: int, moba_topk: int, select_mode: str = "threshold", # "topk" or "threshold" simsum_threshold: float = 0.25, threshold_type: str = "query_head", ) -> torch.Tensor: """ Accelerated MOBA attention for vision tasks with proper LSE normalization. This version: - Splits KV into chunks. - For each query head, selects the top-k relevant KV chunks (including the self chunk) by amplifying the diagonal (self-chunk) logits. - Aggregates the attention outputs from the selected chunks using a log-sum-exp reduction so that attending to each query over the selected chunks is equivalent to the original algorithm. """ # Stack keys and values. kv = torch.stack((k, v), dim=1) seqlen, num_head, head_dim = q.shape # Compute chunk boundaries. cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks( cu_seqlens, moba_chunk_size ) self_attn_cu_seqlen = cu_chunk # Update top-k selection to include the self chunk. moba_topk = min(moba_topk, num_filtered_chunk) # --- Build filtered KV from chunks --- chunk_starts = cu_chunk[filtered_chunk_indices] # [num_filtered_chunk] chunk_ends = cu_chunk[filtered_chunk_indices + 1] # [num_filtered_chunk] chunk_lengths = chunk_ends - chunk_starts # [num_filtered_chunk] max_chunk_len = int(chunk_lengths.max().item()) range_tensor = torch.arange( max_chunk_len, device=kv.device, dtype=chunk_starts.dtype ).unsqueeze(0) indices = chunk_starts.unsqueeze(1) + range_tensor indices = torch.clamp(indices, max=kv.shape[0] - 1) valid_mask = range_tensor < chunk_lengths.unsqueeze(1) gathered = kv[indices.view(-1)].view( num_filtered_chunk, max_chunk_len, *kv.shape[1:] ) gathered = gathered * valid_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).type_as( gathered ) # Compute key_gate_weight over valid tokens. key_values = gathered[ :, :, 0 ].float() # [num_filtered_chunk, max_chunk_len, num_head, head_dim] valid_mask_exp = valid_mask.unsqueeze(-1).unsqueeze(-1) key_sum = (key_values * valid_mask_exp).sum(dim=1) divisor = valid_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1) key_gate_weight = key_sum / divisor # [num_filtered_chunk, num_head, head_dim] # Compute gate logits between key_gate_weight and queries. q_float = q.float() # gate = torch.einsum("nhd,shd->nhs", key_gate_weight, q_float) # [num_filtered_chunk, num_head, seqlen] gate = torch.bmm( key_gate_weight.permute(1, 0, 2), q_float.permute(1, 0, 2).transpose(1, 2) ).permute(1, 0, 2) # Amplify the diagonal (self chunk) contributions. gate_seq_idx = ( torch.arange(seqlen, device=q.device, dtype=torch.int32) .unsqueeze(0) .expand(num_filtered_chunk, seqlen) ) chunk_start = cu_chunk[filtered_chunk_indices] # [num_filtered_chunk] chunk_end = cu_chunk[filtered_chunk_indices + 1] # [num_filtered_chunk] gate_self_chunk_mask = ( ( (gate_seq_idx >= chunk_start.unsqueeze(1)) & (gate_seq_idx < chunk_end.unsqueeze(1)) ) .unsqueeze(1) .expand(-1, num_head, -1) ) amplification_factor = 1e9 # Example factor; adjust as needed. origin_gate = gate.clone() gate = gate.clone() if select_mode == "topk": gate[gate_self_chunk_mask] += amplification_factor # Exclude positions that are outside the valid batch boundaries. batch_starts = cu_seqlens[chunk_to_batch[filtered_chunk_indices]] batch_ends = cu_seqlens[chunk_to_batch[filtered_chunk_indices] + 1] gate_batch_start_mask = gate_seq_idx < batch_starts.unsqueeze(1) gate_batch_end_mask = gate_seq_idx >= batch_ends.unsqueeze(1) gate_inf_mask = gate_batch_start_mask | gate_batch_end_mask gate.masked_fill_(gate_inf_mask.unsqueeze(1), -float("inf")) if select_mode == "topk": # We amplify self‐chunk in gate already, so self entries will rank highest. valid_gate_mask = gate != -float("inf") if threshold_type == "query_head": # === per‐ top-k across chunks (original behavior) === # gate: (C, H, S) _, gate_topk_idx = torch.topk( gate, k=moba_topk, dim=0, largest=True, sorted=False ) gate_idx_mask = torch.zeros_like(gate, dtype=torch.bool) gate_idx_mask.scatter_(0, gate_topk_idx, True) gate_mask = valid_gate_mask & gate_idx_mask elif threshold_type == "overall": # === global top-k across all (chunk, head, seq) entries === C, H, S = gate.shape flat_gate = gate.flatten() flat_mask = valid_gate_mask.flatten() flat_gate_masked = torch.where(flat_mask, flat_gate, -float("inf")) # pick topk global entries vals, idx = torch.topk( flat_gate_masked, k=moba_topk * H * S, largest=True, sorted=False ) others_mask_flat = torch.zeros_like(flat_mask, dtype=torch.bool) others_mask_flat[idx] = True gate_mask = (valid_gate_mask.flatten() & others_mask_flat).view(gate.shape) elif threshold_type == "head_global": # per-head top-k across all chunks and sequence positions C, H, S = gate.shape CS = C * S flat_gate = gate.permute(1, 0, 2).reshape(H, CS) flat_valid = valid_gate_mask.permute(1, 0, 2).reshape(H, CS) flat_gate_masked = torch.where( flat_valid, flat_gate, torch.full_like(flat_gate, -float("inf")) ) # pick top-k indices per head _, topk_idx = torch.topk( flat_gate_masked, k=moba_topk * S, dim=1, largest=True, sorted=False ) gate_idx_flat = torch.zeros_like(flat_valid, dtype=torch.bool) gate_idx_flat.scatter_(1, topk_idx, True) gate_mask = gate_idx_flat.reshape(H, C, S).permute(1, 0, 2) else: raise ValueError( f"Invalid threshold_type for topk: {threshold_type}. " "Choose 'query_head', 'block', or 'overall'." ) elif select_mode == "threshold": # Delegate to the specific thresholding function valid_gate_mask = gate != -float("inf") # (num_chunk, num_head, seqlen) if threshold_type == "query_head": gate_mask = _select_threshold_query_head( gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold ) elif threshold_type == "block": gate_mask = _select_threshold_block( gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold ) elif threshold_type == "overall": gate_mask = _select_threshold_overall( gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold ) elif threshold_type == "head_global": gate_mask = _select_threshold_head_global( gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold ) else: raise ValueError( f"Invalid threshold_type: {threshold_type}. Choose 'query_head', 'block', or 'overall'." ) else: raise ValueError( f"Invalid select_mode: {select_mode}. Choose 'topk' or 'threshold'." ) # eliminate self_chunk in MoBA branch gate_mask = gate_mask & ~gate_self_chunk_mask # if gate_mask is all false, perform flash_attn instead if gate_mask.sum() == 0: return flash_attn_varlen_func( q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, causal=False ) # Determine which query positions are selected. # nonzero_indices has shape [N, 3] where each row is [chunk_index, head_index, seq_index]. moba_q_indices = gate_mask.reshape(gate_mask.shape[0], -1).nonzero(as_tuple=True)[ -1 ] # [(h s k)] moba_q_sh_indices = (moba_q_indices % seqlen) * num_head + ( moba_q_indices // seqlen ) moba_q = ( rearrange(q, "s h d -> (h s) d").index_select(0, moba_q_indices).unsqueeze(1) ) # Build cumulative sequence lengths for the selected queries. moba_seqlen_q = gate_mask.sum(dim=-1).flatten() q_zero_mask = moba_seqlen_q == 0 valid_expert_mask = ~q_zero_mask if q_zero_mask.sum() > 0: moba_seqlen_q = moba_seqlen_q[valid_expert_mask] moba_cu_seqlen_q = torch.cat( ( torch.tensor([0], device=q.device, dtype=moba_seqlen_q.dtype), moba_seqlen_q.cumsum(dim=0), ), dim=0, ).to(torch.int32) # Rearrange gathered KV for the MOBA branch. experts_tensor = rearrange(gathered, "nc cl two h d -> (nc h) cl two d") valid_expert_lengths = ( chunk_lengths.unsqueeze(1) .expand(num_filtered_chunk, num_head) .reshape(-1) .to(torch.int32) ) if q_zero_mask.sum() > 0: experts_tensor = experts_tensor[valid_expert_mask] valid_expert_lengths = valid_expert_lengths[valid_expert_mask] seq_range = torch.arange( experts_tensor.shape[1], device=experts_tensor.device ).unsqueeze(0) mask = seq_range < valid_expert_lengths.unsqueeze(1) moba_kv = experts_tensor[mask] # Shape: ((nc h cl_valid) two d) moba_kv = moba_kv.unsqueeze(2) # Shape: ((nc h cl_valid) two 1 d) moba_cu_seqlen_kv = torch.cat( [ torch.zeros(1, device=experts_tensor.device, dtype=torch.int32), valid_expert_lengths.cumsum(dim=0), ], dim=0, ).to(torch.int32) assert ( moba_cu_seqlen_kv.shape == moba_cu_seqlen_q.shape ), f"Mismatch between moba_cu_seqlen_kv.shape and moba_cu_seqlen_q.shape: {moba_cu_seqlen_kv.shape} vs {moba_cu_seqlen_q.shape}" return MixedAttention.apply( q, k, v, self_attn_cu_seqlen, moba_q, moba_kv, moba_cu_seqlen_q, moba_cu_seqlen_kv, max_seqlen, moba_chunk_size, moba_q_sh_indices, ) def process_moba_input( x, patch_resolution, chunk_size, ): """ Process inputs for the attention function. Args: x (torch.Tensor): Input tensor with shape [batch_size, num_patches, num_heads, head_dim]. patch_resolution (tuple): Tuple containing the patch resolution (t, h, w). chunk_size (int): Size of the chunk. (maybe tuple or int, according to chunk type) Returns: torch.Tensor: Processed input tensor. """ if isinstance(chunk_size, float) or isinstance(chunk_size, int): moba_chunk_size = int(chunk_size * patch_resolution[1] * patch_resolution[2]) else: assert isinstance( chunk_size, (Tuple, list) ), f"chunk_size should be a tuple, list, or int, now it is: {type(chunk_size)}" if len(chunk_size) == 2: assert ( patch_resolution[1] % chunk_size[0] == 0 and patch_resolution[2] % chunk_size[1] == 0 ), f"spatial patch_resolution {patch_resolution[1:]} should be divisible by 2d chunk_size {chunk_size}" nch, ncw = ( patch_resolution[1] // chunk_size[0], patch_resolution[2] // chunk_size[1], ) x = rearrange( x, "b (t nch ch ncw cw) n d -> b (nch ncw t ch cw) n d", t=patch_resolution[0], nch=nch, ncw=ncw, ch=chunk_size[0], cw=chunk_size[1], ) moba_chunk_size = patch_resolution[0] * chunk_size[0] * chunk_size[1] elif len(chunk_size) == 3: assert ( patch_resolution[0] % chunk_size[0] == 0 and patch_resolution[1] % chunk_size[1] == 0 and patch_resolution[2] % chunk_size[2] == 0 ), f"patch_resolution {patch_resolution} should be divisible by 3d chunk_size {chunk_size}" nct, nch, ncw = ( patch_resolution[0] // chunk_size[0], patch_resolution[1] // chunk_size[1], patch_resolution[2] // chunk_size[2], ) x = rearrange( x, "b (nct ct nch ch ncw cw) n d -> b (nct nch ncw ct ch cw) n d", nct=nct, nch=nch, ncw=ncw, ct=chunk_size[0], ch=chunk_size[1], cw=chunk_size[2], ) moba_chunk_size = chunk_size[0] * chunk_size[1] * chunk_size[2] else: raise ValueError( f"chunk_size should be a int, or a tuple of length 2 or 3, now it is: {len(chunk_size)}" ) return x, moba_chunk_size def process_moba_output( x, patch_resolution, chunk_size, ): if isinstance(chunk_size, float) or isinstance(chunk_size, int): pass elif len(chunk_size) == 2: x = rearrange( x, "b (nch ncw t ch cw) n d -> b (t nch ch ncw cw) n d", nch=patch_resolution[1] // chunk_size[0], ncw=patch_resolution[2] // chunk_size[1], t=patch_resolution[0], ch=chunk_size[0], cw=chunk_size[1], ) elif len(chunk_size) == 3: x = rearrange( x, "b (nct nch ncw ct ch cw) n d -> b (nct ct nch ch ncw cw) n d", nct=patch_resolution[0] // chunk_size[0], nch=patch_resolution[1] // chunk_size[1], ncw=patch_resolution[2] // chunk_size[2], ct=chunk_size[0], ch=chunk_size[1], cw=chunk_size[2], ) return x # TEST def generate_data(batch_size, seqlen, num_head, head_dim, dtype): random.seed(0) torch.manual_seed(0) torch.cuda.manual_seed(0) device = torch.cuda.current_device() q = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to( dtype=dtype, device="cuda" ) k = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to( dtype=dtype, device="cuda" ) v = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to( dtype=dtype, device="cuda" ) print(f"q.shape: {q.shape}, k.shape: {k.shape}, v.shape: {v.shape}") cu_seqlens = torch.arange( 0, q.shape[0] * q.shape[1] + 1, q.shape[1], dtype=torch.int32, device="cuda" ) max_seqlen = q.shape[1] q = rearrange(q, "b s ... -> (b s) ...") k = rearrange(k, "b s ... -> (b s) ...") v = rearrange(v, "b s ... -> (b s) ...") return q, k, v, cu_seqlens, max_seqlen def test_attn_varlen_moba_speed( batch, head, seqlen, head_dim, moba_chunk_size, moba_topk, dtype=torch.bfloat16, select_mode="threshold", simsum_threshold=0.25, threshold_type="query_head", ): """Speed test comparing flash_attn vs moba_attention""" # Get data q, k, v, cu_seqlen, max_seqlen = generate_data(batch, seqlen, head, head_dim, dtype) print( f"batch:{batch} head:{head} seqlen:{seqlen} chunk:{moba_chunk_size} topk:{moba_topk} select_mode: {select_mode} simsum_threshold:{simsum_threshold}" ) vo_grad = torch.randn_like(q) # Warmup warmup_iters = 3 perf_test_iters = 10 # Warmup for _ in range(warmup_iters): o = flash_attn_varlen_func( q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False ) torch.autograd.backward(o, vo_grad) torch.cuda.synchronize() start_flash = time.perf_counter() for _ in range(perf_test_iters): o = flash_attn_varlen_func( q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False ) torch.autograd.backward(o, vo_grad) torch.cuda.synchronize() time_flash = (time.perf_counter() - start_flash) / perf_test_iters * 1000 # Warmup for _ in range(warmup_iters): om = moba_attn_varlen( q, k, v, cu_seqlen, max_seqlen, moba_chunk_size=moba_chunk_size, moba_topk=moba_topk, select_mode=select_mode, simsum_threshold=simsum_threshold, threshold_type=threshold_type, ) torch.autograd.backward(om, vo_grad) torch.cuda.synchronize() start_moba = time.perf_counter() for _ in range(perf_test_iters): om = moba_attn_varlen( q, k, v, cu_seqlen, max_seqlen, moba_chunk_size=moba_chunk_size, moba_topk=moba_topk, select_mode=select_mode, simsum_threshold=simsum_threshold, threshold_type=threshold_type, ) torch.autograd.backward(om, vo_grad) torch.cuda.synchronize() time_moba = (time.perf_counter() - start_moba) / perf_test_iters * 1000 print(f"Flash: {time_flash:.2f}ms, MoBA: {time_moba:.2f}ms") print(f"Speedup: {time_flash / time_moba:.2f}x") if __name__ == "__main__": """ CUDA_VISIBLE_DEVICES=1 \ python -u csrc/attn/vmoba_attn/vmoba/vmoba.py """ test_attn_varlen_moba_speed( batch=1, head=12, seqlen=32760, head_dim=128, moba_chunk_size=32760 // 3 // 6 // 4, moba_topk=3, select_mode="threshold", simsum_threshold=0.3, threshold_type="query_head", ) ================================================ FILE: python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Custom CUDA rasterizer for Hunyuan3D texture generation. This module provides JIT-compiled CUDA rasterization for fast mesh rendering. Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2 """ from __future__ import annotations import os from typing import List, Tuple import torch _abs_path = os.path.dirname(os.path.abspath(__file__)) _custom_rasterizer_kernel = None def _load_custom_rasterizer(): """JIT compile and load the custom rasterizer kernel.""" global _custom_rasterizer_kernel if _custom_rasterizer_kernel is not None: return _custom_rasterizer_kernel from torch.utils.cpp_extension import load _custom_rasterizer_kernel = load( name="custom_rasterizer_kernel", sources=[ f"{_abs_path}/rasterizer.cpp", f"{_abs_path}/rasterizer_gpu.cu", ], extra_cflags=["-O3"], extra_cuda_cflags=["-O3", "--use_fast_math"], verbose=False, ) return _custom_rasterizer_kernel def rasterize( pos: torch.Tensor, tri: torch.Tensor, resolution: Tuple[int, int], clamp_depth: torch.Tensor = None, use_depth_prior: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor]: """Rasterize mesh to get face indices and barycentric coordinates.""" kernel = _load_custom_rasterizer() if clamp_depth is None: clamp_depth = torch.zeros(0, device=pos.device) # pos should be [N, 4], remove batch dim if present if pos.dim() == 3: pos = pos[0] findices, barycentric = kernel.rasterize_image( pos, tri, clamp_depth, resolution[1], resolution[0], 1e-6, use_depth_prior ) return findices, barycentric def interpolate( col: torch.Tensor, findices: torch.Tensor, barycentric: torch.Tensor, tri: torch.Tensor, ) -> torch.Tensor: """Interpolate vertex attributes using barycentric coordinates.""" # Handle zero indices (background) f = findices - 1 + (findices == 0) vcol = col[0, tri.long()[f.long()]] result = barycentric.view(*barycentric.shape, 1) * vcol result = torch.sum(result, axis=-2) return result.view(1, *result.shape) __all__ = ["rasterize", "interpolate"] ================================================ FILE: python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.cpp ================================================ // SPDX-License-Identifier: Apache-2.0 // Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2 // Original license: TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT #include "rasterizer.h" void rasterizeTriangleCPU(int idx, float* vt0, float* vt1, float* vt2, int width, int height, INT64* zbuffer, float* d, float occlusion_truncation) { float x_min = std::min(vt0[0], std::min(vt1[0],vt2[0])); float x_max = std::max(vt0[0], std::max(vt1[0],vt2[0])); float y_min = std::min(vt0[1], std::min(vt1[1],vt2[1])); float y_max = std::max(vt0[1], std::max(vt1[1],vt2[1])); for (int px = x_min; px < x_max + 1; ++px) { if (px < 0 || px >= width) continue; for (int py = y_min; py < y_max + 1; ++py) { if (py < 0 || py >= height) continue; float vt[2] = {px + 0.5f, py + 0.5f}; float baryCentricCoordinate[3]; calculateBarycentricCoordinate(vt0, vt1, vt2, vt, baryCentricCoordinate); if (isBarycentricCoordInBounds(baryCentricCoordinate)) { int pixel = py * width + px; if (zbuffer == 0) { zbuffer[pixel] = (INT64)(idx + 1); continue; } float depth = baryCentricCoordinate[0] * vt0[2] + baryCentricCoordinate[1] * vt1[2] + baryCentricCoordinate[2] * vt2[2]; float depth_thres = 0; if (d) { depth_thres = d[pixel] * 0.49999f + 0.5f + occlusion_truncation; } int z_quantize = depth * (2<<17); INT64 token = (INT64)z_quantize * MAXINT + (INT64)(idx + 1); if (depth < depth_thres) continue; zbuffer[pixel] = std::min(zbuffer[pixel], token); } } } } void barycentricFromImgcoordCPU(float* V, int* F, int* findices, INT64* zbuffer, int width, int height, int num_vertices, int num_faces, float* barycentric_map, int pix) { INT64 f = zbuffer[pix] % MAXINT; if (f == (MAXINT-1)) { findices[pix] = 0; barycentric_map[pix * 3] = 0; barycentric_map[pix * 3 + 1] = 0; barycentric_map[pix * 3 + 2] = 0; return; } findices[pix] = f; f -= 1; float barycentric[3] = {0, 0, 0}; if (f >= 0) { float vt[2] = {float(pix % width) + 0.5f, float(pix / width) + 0.5f}; float* vt0_ptr = V + (F[f * 3] * 4); float* vt1_ptr = V + (F[f * 3 + 1] * 4); float* vt2_ptr = V + (F[f * 3 + 2] * 4); float vt0[2] = {(vt0_ptr[0] / vt0_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt0_ptr[1] / vt0_ptr[3]) * (height - 1) + 0.5f}; float vt1[2] = {(vt1_ptr[0] / vt1_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt1_ptr[1] / vt1_ptr[3]) * (height - 1) + 0.5f}; float vt2[2] = {(vt2_ptr[0] / vt2_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt2_ptr[1] / vt2_ptr[3]) * (height - 1) + 0.5f}; calculateBarycentricCoordinate(vt0, vt1, vt2, vt, barycentric); barycentric[0] = barycentric[0] / vt0_ptr[3]; barycentric[1] = barycentric[1] / vt1_ptr[3]; barycentric[2] = barycentric[2] / vt2_ptr[3]; float w = 1.0f / (barycentric[0] + barycentric[1] + barycentric[2]); barycentric[0] *= w; barycentric[1] *= w; barycentric[2] *= w; } barycentric_map[pix * 3] = barycentric[0]; barycentric_map[pix * 3 + 1] = barycentric[1]; barycentric_map[pix * 3 + 2] = barycentric[2]; } void rasterizeImagecoordsKernelCPU(float* V, int* F, float* d, INT64* zbuffer, float occlusion_trunc, int width, int height, int num_vertices, int num_faces, int f) { float* vt0_ptr = V + (F[f * 3] * 4); float* vt1_ptr = V + (F[f * 3 + 1] * 4); float* vt2_ptr = V + (F[f * 3 + 2] * 4); float vt0[3] = {(vt0_ptr[0] / vt0_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt0_ptr[1] / vt0_ptr[3]) * (height - 1) + 0.5f, vt0_ptr[2] / vt0_ptr[3] * 0.49999f + 0.5f}; float vt1[3] = {(vt1_ptr[0] / vt1_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt1_ptr[1] / vt1_ptr[3]) * (height - 1) + 0.5f, vt1_ptr[2] / vt1_ptr[3] * 0.49999f + 0.5f}; float vt2[3] = {(vt2_ptr[0] / vt2_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt2_ptr[1] / vt2_ptr[3]) * (height - 1) + 0.5f, vt2_ptr[2] / vt2_ptr[3] * 0.49999f + 0.5f}; rasterizeTriangleCPU(f, vt0, vt1, vt2, width, height, zbuffer, d, occlusion_trunc); } std::vector rasterize_image_cpu(torch::Tensor V, torch::Tensor F, torch::Tensor D, int width, int height, float occlusion_truncation, int use_depth_prior) { int num_faces = F.size(0); int num_vertices = V.size(0); auto options = torch::TensorOptions().dtype(torch::kInt32).requires_grad(false); auto INT64_options = torch::TensorOptions().dtype(torch::kInt64).requires_grad(false); auto findices = torch::zeros({height, width}, options); INT64 maxint = (INT64)MAXINT * (INT64)MAXINT + (MAXINT - 1); auto z_min = torch::ones({height, width}, INT64_options) * (int64_t)maxint; if (!use_depth_prior) { for (int i = 0; i < num_faces; ++i) { rasterizeImagecoordsKernelCPU(V.data_ptr(), F.data_ptr(), 0, (INT64*)z_min.data_ptr(), occlusion_truncation, width, height, num_vertices, num_faces, i); } } else { for (int i = 0; i < num_faces; ++i) rasterizeImagecoordsKernelCPU(V.data_ptr(), F.data_ptr(), D.data_ptr(), (INT64*)z_min.data_ptr(), occlusion_truncation, width, height, num_vertices, num_faces, i); } auto float_options = torch::TensorOptions().dtype(torch::kFloat32).requires_grad(false); auto barycentric = torch::zeros({height, width, 3}, float_options); for (int i = 0; i < width * height; ++i) barycentricFromImgcoordCPU(V.data_ptr(), F.data_ptr(), findices.data_ptr(), (INT64*)z_min.data_ptr(), width, height, num_vertices, num_faces, barycentric.data_ptr(), i); return {findices, barycentric}; } std::vector rasterize_image(torch::Tensor V, torch::Tensor F, torch::Tensor D, int width, int height, float occlusion_truncation, int use_depth_prior) { int device_id = V.get_device(); if (device_id == -1) return rasterize_image_cpu(V, F, D, width, height, occlusion_truncation, use_depth_prior); else return rasterize_image_gpu(V, F, D, width, height, occlusion_truncation, use_depth_prior); } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("rasterize_image", &rasterize_image, "Custom image rasterization"); } ================================================ FILE: python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.h ================================================ // SPDX-License-Identifier: Apache-2.0 // Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2 // Original license: TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT #ifndef RASTERIZER_H_ #define RASTERIZER_H_ #include #include #include #include #define INT64 unsigned long long #define MAXINT 2147483647 __host__ __device__ inline float calculateSignedArea2(float* a, float* b, float* c) { return ((c[0] - a[0]) * (b[1] - a[1]) - (b[0] - a[0]) * (c[1] - a[1])); } __host__ __device__ inline void calculateBarycentricCoordinate(float* a, float* b, float* c, float* p, float* barycentric) { float beta_tri = calculateSignedArea2(a, p, c); float gamma_tri = calculateSignedArea2(a, b, p); float area = calculateSignedArea2(a, b, c); if (area == 0) { barycentric[0] = -1.0; barycentric[1] = -1.0; barycentric[2] = -1.0; return; } float tri_inv = 1.0 / area; float beta = beta_tri * tri_inv; float gamma = gamma_tri * tri_inv; float alpha = 1.0 - beta - gamma; barycentric[0] = alpha; barycentric[1] = beta; barycentric[2] = gamma; } __host__ __device__ inline bool isBarycentricCoordInBounds(float* barycentricCoord) { return barycentricCoord[0] >= 0.0 && barycentricCoord[0] <= 1.0 && barycentricCoord[1] >= 0.0 && barycentricCoord[1] <= 1.0 && barycentricCoord[2] >= 0.0 && barycentricCoord[2] <= 1.0; } std::vector rasterize_image_gpu(torch::Tensor V, torch::Tensor F, torch::Tensor D, int width, int height, float occlusion_truncation, int use_depth_prior); #endif ================================================ FILE: python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer_gpu.cu ================================================ // SPDX-License-Identifier: Apache-2.0 // Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2 // Original license: TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT #include "rasterizer.h" __device__ void rasterizeTriangleGPU(int idx, float* vt0, float* vt1, float* vt2, int width, int height, INT64* zbuffer, float* d, float occlusion_truncation) { float x_min = std::min(vt0[0], std::min(vt1[0],vt2[0])); float x_max = std::max(vt0[0], std::max(vt1[0],vt2[0])); float y_min = std::min(vt0[1], std::min(vt1[1],vt2[1])); float y_max = std::max(vt0[1], std::max(vt1[1],vt2[1])); for (int px = x_min; px < x_max + 1; ++px) { if (px < 0 || px >= width) continue; for (int py = y_min; py < y_max + 1; ++py) { if (py < 0 || py >= height) continue; float vt[2] = {px + 0.5f, py + 0.5f}; float baryCentricCoordinate[3]; calculateBarycentricCoordinate(vt0, vt1, vt2, vt, baryCentricCoordinate); if (isBarycentricCoordInBounds(baryCentricCoordinate)) { int pixel = py * width + px; if (zbuffer == 0) { atomicExch(&zbuffer[pixel], (INT64)(idx + 1)); continue; } float depth = baryCentricCoordinate[0] * vt0[2] + baryCentricCoordinate[1] * vt1[2] + baryCentricCoordinate[2] * vt2[2]; float depth_thres = 0; if (d) { depth_thres = d[pixel] * 0.49999f + 0.5f + occlusion_truncation; } int z_quantize = depth * (2<<17); INT64 token = (INT64)z_quantize * MAXINT + (INT64)(idx + 1); if (depth < depth_thres) continue; atomicMin(&zbuffer[pixel], token); } } } } __global__ void barycentricFromImgcoordGPU(float* V, int* F, int* findices, INT64* zbuffer, int width, int height, int num_vertices, int num_faces, float* barycentric_map) { int pix = blockIdx.x * blockDim.x + threadIdx.x; if (pix >= width * height) return; INT64 f = zbuffer[pix] % MAXINT; if (f == (MAXINT-1)) { findices[pix] = 0; barycentric_map[pix * 3] = 0; barycentric_map[pix * 3 + 1] = 0; barycentric_map[pix * 3 + 2] = 0; return; } findices[pix] = f; f -= 1; float barycentric[3] = {0, 0, 0}; if (f >= 0) { float vt[2] = {float(pix % width) + 0.5f, float(pix / width) + 0.5f}; float* vt0_ptr = V + (F[f * 3] * 4); float* vt1_ptr = V + (F[f * 3 + 1] * 4); float* vt2_ptr = V + (F[f * 3 + 2] * 4); float vt0[2] = {(vt0_ptr[0] / vt0_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt0_ptr[1] / vt0_ptr[3]) * (height - 1) + 0.5f}; float vt1[2] = {(vt1_ptr[0] / vt1_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt1_ptr[1] / vt1_ptr[3]) * (height - 1) + 0.5f}; float vt2[2] = {(vt2_ptr[0] / vt2_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt2_ptr[1] / vt2_ptr[3]) * (height - 1) + 0.5f}; calculateBarycentricCoordinate(vt0, vt1, vt2, vt, barycentric); barycentric[0] = barycentric[0] / vt0_ptr[3]; barycentric[1] = barycentric[1] / vt1_ptr[3]; barycentric[2] = barycentric[2] / vt2_ptr[3]; float w = 1.0f / (barycentric[0] + barycentric[1] + barycentric[2]); barycentric[0] *= w; barycentric[1] *= w; barycentric[2] *= w; } barycentric_map[pix * 3] = barycentric[0]; barycentric_map[pix * 3 + 1] = barycentric[1]; barycentric_map[pix * 3 + 2] = barycentric[2]; } __global__ void rasterizeImagecoordsKernelGPU(float* V, int* F, float* d, INT64* zbuffer, float occlusion_trunc, int width, int height, int num_vertices, int num_faces) { int f = blockIdx.x * blockDim.x + threadIdx.x; if (f >= num_faces) return; float* vt0_ptr = V + (F[f * 3] * 4); float* vt1_ptr = V + (F[f * 3 + 1] * 4); float* vt2_ptr = V + (F[f * 3 + 2] * 4); float vt0[3] = {(vt0_ptr[0] / vt0_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt0_ptr[1] / vt0_ptr[3]) * (height - 1) + 0.5f, vt0_ptr[2] / vt0_ptr[3] * 0.49999f + 0.5f}; float vt1[3] = {(vt1_ptr[0] / vt1_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt1_ptr[1] / vt1_ptr[3]) * (height - 1) + 0.5f, vt1_ptr[2] / vt1_ptr[3] * 0.49999f + 0.5f}; float vt2[3] = {(vt2_ptr[0] / vt2_ptr[3] * 0.5f + 0.5f) * (width - 1) + 0.5f, (0.5f + 0.5f * vt2_ptr[1] / vt2_ptr[3]) * (height - 1) + 0.5f, vt2_ptr[2] / vt2_ptr[3] * 0.49999f + 0.5f}; rasterizeTriangleGPU(f, vt0, vt1, vt2, width, height, zbuffer, d, occlusion_trunc); } std::vector rasterize_image_gpu(torch::Tensor V, torch::Tensor F, torch::Tensor D, int width, int height, float occlusion_truncation, int use_depth_prior) { int device_id = V.get_device(); cudaSetDevice(device_id); int num_faces = F.size(0); int num_vertices = V.size(0); auto options = torch::TensorOptions().dtype(torch::kInt32).device(torch::kCUDA, device_id).requires_grad(false); auto INT64_options = torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA, device_id).requires_grad(false); auto findices = torch::zeros({height, width}, options); INT64 maxint = (INT64)MAXINT * (INT64)MAXINT + (MAXINT - 1); auto z_min = torch::ones({height, width}, INT64_options) * (int64_t)maxint; if (!use_depth_prior) { rasterizeImagecoordsKernelGPU<<<(num_faces+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(V.data_ptr(), F.data_ptr(), 0, (INT64*)z_min.data_ptr(), occlusion_truncation, width, height, num_vertices, num_faces); } else { rasterizeImagecoordsKernelGPU<<<(num_faces+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(V.data_ptr(), F.data_ptr(), D.data_ptr(), (INT64*)z_min.data_ptr(), occlusion_truncation, width, height, num_vertices, num_faces); } auto float_options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA, device_id).requires_grad(false); auto barycentric = torch::zeros({height, width, 3}, float_options); barycentricFromImgcoordGPU<<<(width * height + 255)/256, 256, 0, at::cuda::getCurrentCUDAStream()>>>(V.data_ptr(), F.data_ptr(), findices.data_ptr(), (INT64*)z_min.data_ptr(), width, height, num_vertices, num_faces, barycentric.data_ptr()); return {findices, barycentric}; } ================================================ FILE: python/sglang/multimodal_gen/csrc/render/mesh_processor/__init__.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Mesh processor C++ extension for texture inpainting. This module provides JIT-compiled C++ mesh processing for fast texture inpainting. Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2 """ from __future__ import annotations import os from typing import Tuple import numpy as np _abs_path = os.path.dirname(os.path.abspath(__file__)) _mesh_processor_kernel = None def _load_mesh_processor(): """JIT compile and load the mesh processor kernel.""" global _mesh_processor_kernel if _mesh_processor_kernel is not None: return _mesh_processor_kernel from torch.utils.cpp_extension import load _mesh_processor_kernel = load( name="mesh_processor_kernel", sources=[ f"{_abs_path}/mesh_processor.cpp", ], extra_cflags=["-O3"], verbose=False, ) return _mesh_processor_kernel def meshVerticeInpaint( texture: np.ndarray, mask: np.ndarray, vtx_pos: np.ndarray, vtx_uv: np.ndarray, pos_idx: np.ndarray, uv_idx: np.ndarray, method: str = "smooth", ) -> Tuple[np.ndarray, np.ndarray]: """Inpaint texture using mesh vertex connectivity.""" kernel = _load_mesh_processor() texture = np.ascontiguousarray(texture, dtype=np.float32) mask = np.ascontiguousarray(mask, dtype=np.uint8) vtx_pos = np.ascontiguousarray(vtx_pos, dtype=np.float32) vtx_uv = np.ascontiguousarray(vtx_uv, dtype=np.float32) pos_idx = np.ascontiguousarray(pos_idx, dtype=np.int32) uv_idx = np.ascontiguousarray(uv_idx, dtype=np.int32) return kernel.meshVerticeInpaint(texture, mask, vtx_pos, vtx_uv, pos_idx, uv_idx, method) __all__ = ["meshVerticeInpaint"] ================================================ FILE: python/sglang/multimodal_gen/csrc/render/mesh_processor/mesh_processor.cpp ================================================ // SPDX-License-Identifier: Apache-2.0 // Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2 // Original license: TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT #include #include #include #include #include #include #include #include namespace py = pybind11; using namespace std; std::pair, py::array_t> meshVerticeInpaint_smooth(py::array_t texture, py::array_t mask, py::array_t vtx_pos, py::array_t vtx_uv, py::array_t pos_idx, py::array_t uv_idx) { auto texture_buf = texture.request(); auto mask_buf = mask.request(); auto vtx_pos_buf = vtx_pos.request(); auto vtx_uv_buf = vtx_uv.request(); auto pos_idx_buf = pos_idx.request(); auto uv_idx_buf = uv_idx.request(); int texture_height = texture_buf.shape[0]; int texture_width = texture_buf.shape[1]; int texture_channel = texture_buf.shape[2]; float* texture_ptr = static_cast(texture_buf.ptr); uint8_t* mask_ptr = static_cast(mask_buf.ptr); int vtx_num = vtx_pos_buf.shape[0]; float* vtx_pos_ptr = static_cast(vtx_pos_buf.ptr); float* vtx_uv_ptr = static_cast(vtx_uv_buf.ptr); int* pos_idx_ptr = static_cast(pos_idx_buf.ptr); int* uv_idx_ptr = static_cast(uv_idx_buf.ptr); vector vtx_mask(vtx_num, 0.0f); vector> vtx_color(vtx_num, vector(texture_channel, 0.0f)); vector uncolored_vtxs; vector> G(vtx_num); for (int i = 0; i < uv_idx_buf.shape[0]; ++i) { for (int k = 0; k < 3; ++k) { int vtx_uv_idx = uv_idx_ptr[i * 3 + k]; int vtx_idx = pos_idx_ptr[i * 3 + k]; int uv_v = round(vtx_uv_ptr[vtx_uv_idx * 2] * (texture_width - 1)); int uv_u = round((1.0 - vtx_uv_ptr[vtx_uv_idx * 2 + 1]) * (texture_height - 1)); if (mask_ptr[uv_u * texture_width + uv_v] > 0) { vtx_mask[vtx_idx] = 1.0f; for (int c = 0; c < texture_channel; ++c) { vtx_color[vtx_idx][c] = texture_ptr[(uv_u * texture_width + uv_v) * texture_channel + c]; } }else{ uncolored_vtxs.push_back(vtx_idx); } G[pos_idx_ptr[i * 3 + k]].push_back(pos_idx_ptr[i * 3 + (k + 1) % 3]); } } int smooth_count = 2; int last_uncolored_vtx_count = 0; while (smooth_count>0) { int uncolored_vtx_count = 0; for (int vtx_idx : uncolored_vtxs) { vector sum_color(texture_channel, 0.0f); float total_weight = 0.0f; array vtx_0 = {vtx_pos_ptr[vtx_idx * 3], vtx_pos_ptr[vtx_idx * 3 + 1], vtx_pos_ptr[vtx_idx * 3 + 2]}; for (int connected_idx : G[vtx_idx]) { if (vtx_mask[connected_idx] > 0) { array vtx1 = {vtx_pos_ptr[connected_idx * 3], vtx_pos_ptr[connected_idx * 3 + 1], vtx_pos_ptr[connected_idx * 3 + 2]}; float dist_weight = 1.0f / max(sqrt(pow(vtx_0[0] - vtx1[0], 2) + pow(vtx_0[1] - vtx1[1], 2) + \ pow(vtx_0[2] - vtx1[2], 2)), 1E-4); dist_weight = dist_weight * dist_weight; for (int c = 0; c < texture_channel; ++c) { sum_color[c] += vtx_color[connected_idx][c] * dist_weight; } total_weight += dist_weight; } } if (total_weight > 0.0f) { for (int c = 0; c < texture_channel; ++c) { vtx_color[vtx_idx][c] = sum_color[c] / total_weight; } vtx_mask[vtx_idx] = 1.0f; } else { uncolored_vtx_count++; } } if(last_uncolored_vtx_count==uncolored_vtx_count){ smooth_count--; }else{ smooth_count++; } last_uncolored_vtx_count = uncolored_vtx_count; } py::array_t new_texture(texture_buf.size); py::array_t new_mask(mask_buf.size); auto new_texture_buf = new_texture.request(); auto new_mask_buf = new_mask.request(); float* new_texture_ptr = static_cast(new_texture_buf.ptr); uint8_t* new_mask_ptr = static_cast(new_mask_buf.ptr); std::copy(texture_ptr, texture_ptr + texture_buf.size, new_texture_ptr); std::copy(mask_ptr, mask_ptr + mask_buf.size, new_mask_ptr); for (int face_idx = 0; face_idx < uv_idx_buf.shape[0]; ++face_idx) { for (int k = 0; k < 3; ++k) { int vtx_uv_idx = uv_idx_ptr[face_idx * 3 + k]; int vtx_idx = pos_idx_ptr[face_idx * 3 + k]; if (vtx_mask[vtx_idx] == 1.0f) { int uv_v = round(vtx_uv_ptr[vtx_uv_idx * 2] * (texture_width - 1)); int uv_u = round((1.0 - vtx_uv_ptr[vtx_uv_idx * 2 + 1]) * (texture_height - 1)); for (int c = 0; c < texture_channel; ++c) { new_texture_ptr[(uv_u * texture_width + uv_v) * texture_channel + c] = vtx_color[vtx_idx][c]; } new_mask_ptr[uv_u * texture_width + uv_v] = 255; } } } new_texture.resize({texture_height, texture_width, 3}); new_mask.resize({texture_height, texture_width}); return std::make_pair(new_texture, new_mask); } std::pair, py::array_t> meshVerticeInpaint(py::array_t texture, py::array_t mask, py::array_t vtx_pos, py::array_t vtx_uv, py::array_t pos_idx, py::array_t uv_idx, const std::string& method = "smooth") { if (method == "smooth") { return meshVerticeInpaint_smooth(texture, mask, vtx_pos, vtx_uv, pos_idx, uv_idx); } else { throw std::invalid_argument("Invalid method. Use 'smooth'."); } } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("meshVerticeInpaint", &meshVerticeInpaint, "Mesh-aware texture inpainting", py::arg("texture"), py::arg("mask"), py::arg("vtx_pos"), py::arg("vtx_uv"), py::arg("pos_idx"), py::arg("uv_idx"), py::arg("method") = "smooth"); } ================================================ FILE: python/sglang/multimodal_gen/docs/quantization.md ================================================ # Quantization This document introduces the model quantization schemes supported in SGLang and how to use them to reduce memory usage and accelerate inference. ## Nunchaku (SVDQuant) ### Introduction **SVDQuant** is a Post-Training Quantization (PTQ) technique for diffusion models that quantizes model weights and activations to 4-bit precision (W4A4) while maintaining high visual quality. This method uses Singular Value Decomposition (SVD) to decompose the weight matrix into low-rank components and residuals, effectively absorbing outliers in activations, making 4-bit quantization possible. **Nunchaku** is a high-performance inference engine that implements SVDQuant, optimized for low-bit neural networks. It is not Quantization-Aware Training (QAT), but directly quantizes pre-trained models. Paper: [SVDQuant: Absorbing Outliers by Low-Rank Components for 4-Bit Diffusion Models](https://arxiv.org/abs/2411.05007) (ICLR 2025 Spotlight) ### Key Features SVDQuant significantly reduces memory usage and accelerates inference while maintaining visual quality: - **Memory Optimization**: Reduces memory usage by **3.6×** compared to BF16 models. - **Inference Acceleration**: - **3.0×** faster than the NF4 (W4A16) baseline on desktop/laptop RTX 4090 GPUs. - **8.7×** speedup on laptop RTX 4090 by eliminating CPU offloading compared to 16-bit models. - **3.1×** faster than BF16 and NF4 models on RTX 5090 GPUs with NVFP4. ### Supported Precisions Nunchaku supports two quantization precisions: - **INT4**: Standard INT4 quantization, supported on NVIDIA GPUs with Compute Capability 7.0+ (RTX 20 series and above). - **NVFP4**: FP4 quantization, providing better image quality on newer cards like the RTX 5090. ### Usage #### 1. Install Nunchaku ```bash pip install nunchaku ``` For more installation information, please refer to the [Nunchaku Official Documentation](https://nunchaku.tech/docs/nunchaku/installation/installation.html). #### 2. Download Quantized Models Nunchaku provides pre-quantized model weights available on Hugging Face: - [nunchaku-ai/nunchaku-qwen-image](https://huggingface.co/nunchaku-ai/nunchaku-qwen-image) - [nunchaku-ai/nunchaku-flux](https://huggingface.co/nunchaku-ai/nunchaku-flux) Taking Qwen-Image as an example, several quantized models with different configurations are provided: | Filename | Precision | Rank | Usage | |----------|-----------|------|-------| | `svdq-int4_r32-qwen-image.safetensors` | INT4 | 32 | Standard Version | | `svdq-int4_r128-qwen-image.safetensors` | INT4 | 128 | High-Quality Version | | `svdq-fp4_r32-qwen-image.safetensors` | NVFP4 | 32 | RTX 5090 Standard Version | | `svdq-fp4_r128-qwen-image.safetensors` | NVFP4 | 128 | RTX 5090 High-Quality Version | | `svdq-int4_r32-qwen-image-lightningv1.0-4steps.safetensors` | INT4 | 32 | Lightning 4-Step Version | | `svdq-int4_r128-qwen-image-lightningv1.1-8steps.safetensors` | INT4 | 128 | Lightning 8-Step Version | > **Note**: Higher Rank usually means better image quality, but with slightly increased memory usage and computation. #### 3. Run Quantized Models SGLang features **smart auto-detection** for Nunchaku models. In most cases, you only need to provide the path to the quantized weights, and the precision and rank will be automatically inferred from the filename. **Simplified Command (Recommended):** ```bash sglang generate \ --model-path Qwen/Qwen-Image \ --prompt "change the raccoon to a cute cat" \ --save-output \ --transformer-weights-path /path/to/svdq-int4_r32-qwen-image.safetensors ``` **Manual Override (If needed):** If your filename doesn't follow the standard naming convention, or you want to force specific settings: - `--enable-svdquant`: Manually enable SVDQuant. - `--quantization-precision`: Set to `int4` or `nvfp4`. - `--quantization-rank`: Set the SVD rank (e.g., 32, 128). - `--quantization-act-unsigned` (Optional): Use unsigned activation quantization. Example with manual overrides: ```bash sglang generate \ --model-path Qwen/Qwen-Image \ --prompt "a beautiful sunset" \ --enable-svdquant \ --transformer-weights-path /path/to/custom_model.safetensors \ --quantization-precision int4 \ --quantization-rank 128 ``` #### 4. Configuration Recommendations Choose the appropriate configuration based on your hardware and requirements: | Scenario | Recommended Config | Description | |----------|-------------------|-------------| | Standard Use (20/30/40 Series GPU) | INT4 + Rank 32 | Balanced performance and quality | | Quality Focus (Sufficient VRAM) | INT4 + Rank 128 | Better image quality | | RTX 5090 Standard Use | NVFP4 + Rank 32 | Utilizes FP4 hardware acceleration | | RTX 5090 Quality Focus | NVFP4 + Rank 128 | Best image quality | | Fast Prototyping/Preview | Lightning 4-Step Version | Extremely fast generation, slightly reduced quality | ### Notes 1. Model Path Correspondence: `--model-path` should point to the original non-quantized model (for loading config and tokenizer, etc.), while `--transformer-weights-path` points to the quantized weight file / folder / Huggingface Repo ID. 2. Auto-Detection Requirements: For auto-detection to work, the filename must contain the pattern `svdq-{precision}_r{rank}` (e.g., `svdq-int4_r32`). 3. GPU Compatibility: - INT4: Supports NVIDIA GPUs with Compute Capability 7.0+ (RTX 20 series and above). - NVFP4: Optimized mainly for newer cards like the RTX 50 series that support FP4. 4. Lightning Models: When using Lightning versions, adjust `--num-inference-steps` accordingly (usually 4 or 8 steps). ### Custom Model Quantization If you want to quantize your own models, you can use the [DeepCompressor](https://github.com/mit-han-lab/deepcompressor) tool. For detailed instructions, please refer to the Nunchaku official documentation. ## Quantization ### Usage #### Option 1: Pre-quantized folder (has `config.json`) For quantized checkpoints that include a `config.json` with a `quantization_config` field (e.g., models converted via `convert_hf_to_fp8.py`), where the transformer's `config.json` already encodes the `quantization_config`, use the component override: ```bash sglang generate \ --model-path /path/to/FLUX.1-dev \ --transformer-path /path/to/FLUX.1-dev/transformer-FP8 \ --prompt "A Logo With Bold Large Text: SGL Diffusion" \ --save-output ``` If you need to convert a model to FP8 format yourself, use the provided conversion script: ```bash # convert transformer to FP8 with block quantization python -m sglang.multimodal_gen.tools.convert_hf_to_fp8 \ --model-dir /path/to/FLUX.1-dev/transformer \ --save-dir /path/to/FLUX.1-dev/transformer-FP8 \ --strategy block \ --block-size 128 128 ``` #### Option 2: Pre-quantized single-file checkpoint (no `config.json`) Some providers (e.g., [black-forest-labs/FLUX.2-klein-9b-fp8](https://huggingface.co/black-forest-labs/FLUX.2-klein-9b-fp8)) distribute a single `.safetensors` file without a companion `config.json`. Use `--transformer-weights-path` to point to this file (or HuggingFace repo ID) while keeping `--model-path` for the base model: ```bash sglang generate \ --model-path black-forest-labs/FLUX.2-klein-9B \ --transformer-weights-path black-forest-labs/FLUX.2-klein-9b-fp8 \ --prompt "A Logo With Bold Large Text: SGL Diffusion" \ --save-output ``` SGLang-Diffusion will automatically read the `quantization_config` metadata embedded in the safetensors file header (if present). For the quant config to be auto-detected, the file's metadata must contain a JSON-encoded `quantization_config` key with at least a `quant_method` field (e.g. `"fp8"`). Note: this feature is a WIP ================================================ FILE: python/sglang/multimodal_gen/envs.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py import logging import os from typing import TYPE_CHECKING, Any, Callable from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var logger = logging.getLogger(__name__) if TYPE_CHECKING: SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL: int = 60 SGLANG_DIFFUSION_NCCL_SO_PATH: str | None = None LD_LIBRARY_PATH: str | None = None LOCAL_RANK: int = 0 CUDA_VISIBLE_DEVICES: str | None = None SGLANG_DIFFUSION_CACHE_ROOT: str = os.path.expanduser("~/.cache/sgl_diffusion") SGLANG_DIFFUSION_CONFIG_ROOT: str = os.path.expanduser("~/.config/sgl_diffusion") SGLANG_DIFFUSION_CONFIGURE_LOGGING: int = 1 SGLANG_DIFFUSION_LOGGING_LEVEL: str = "INFO" SGLANG_DIFFUSION_LOGGING_PREFIX: str = "" SGLANG_DIFFUSION_LOGGING_CONFIG_PATH: str | None = None SGLANG_DIFFUSION_TRACE_FUNCTION: int = 0 SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD: str = "fork" SGLANG_DIFFUSION_TARGET_DEVICE: str = "cuda" MAX_JOBS: str | None = None NVCC_THREADS: str | None = None CMAKE_BUILD_TYPE: str | None = None VERBOSE: bool = False SGLANG_DIFFUSION_SERVER_DEV_MODE: bool = False SGLANG_DIFFUSION_STAGE_LOGGING: bool = False # cache-dit env vars (primary transformer) SGLANG_CACHE_DIT_ENABLED: bool = False SGLANG_CACHE_DIT_FN: int = 1 SGLANG_CACHE_DIT_BN: int = 0 SGLANG_CACHE_DIT_WARMUP: int = 4 SGLANG_CACHE_DIT_RDT: float = 0.24 SGLANG_CACHE_DIT_MC: int = 3 SGLANG_CACHE_DIT_TAYLORSEER: bool = False SGLANG_CACHE_DIT_TS_ORDER: int = 1 SGLANG_CACHE_DIT_SCM_PRESET: str = "none" SGLANG_CACHE_DIT_SCM_COMPUTE_BINS: str | None = None SGLANG_CACHE_DIT_SCM_CACHE_BINS: str | None = None SGLANG_CACHE_DIT_SCM_POLICY: str = "dynamic" # cache-dit env vars (secondary transformer, e.g., Wan2.2 low-noise expert) SGLANG_CACHE_DIT_SECONDARY_FN: int = 1 SGLANG_CACHE_DIT_SECONDARY_BN: int = 0 SGLANG_CACHE_DIT_SECONDARY_WARMUP: int = 4 SGLANG_CACHE_DIT_SECONDARY_RDT: float = 0.24 SGLANG_CACHE_DIT_SECONDARY_MC: int = 3 SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER: bool = False SGLANG_CACHE_DIT_SECONDARY_TS_ORDER: int = 1 # model loading SGLANG_USE_RUNAI_MODEL_STREAMER: bool = True SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D: bool = False SGLANG_USE_ROCM_VAE: bool = False def get_default_cache_root() -> str: return os.getenv( "XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache"), ) def get_default_config_root() -> str: return os.getenv( "XDG_CONFIG_HOME", os.path.join(os.path.expanduser("~"), ".config"), ) def maybe_convert_int(value: str | None) -> int | None: return int(value) if value is not None else None # helpers for environment variable definitions def _lazy_str(key: str, default: str | None = None) -> Callable[[], str | None]: return lambda: os.getenv(key, default) def _lazy_int(key: str, default: str | int | None = None) -> Callable[[], int | None]: def _getter(): val = os.getenv(key) if val is None: return int(default) if default is not None else None return int(val) return _getter def _lazy_float(key: str, default: str | float) -> Callable[[], float]: return lambda: float(os.getenv(key, str(default))) def _lazy_bool(key: str, default: str = "false") -> Callable[[], bool]: return lambda: get_bool_env_var(key, default) def _lazy_bool_any(keys: list[str], default: str = "false") -> Callable[[], bool]: def _getter(): for key in keys: if get_bool_env_var(key, "false"): return True return ( get_bool_env_var("", default) if not keys else get_bool_env_var(keys[0], default) ) return _getter def _lazy_path( key: str, default_func: Callable[[], str] | None = None ) -> Callable[[], str | None]: def _getter(): val = os.getenv(key) if val is None: if default_func is None: return None val = default_func() return os.path.expanduser(val) return _getter # The begin-* and end* here are used by the documentation generator # to extract the used env vars. # begin-env-vars-definition environment_variables: dict[str, Callable[[], Any]] = { # ================== Installation Time Env Vars ================== # Target device of sglang-diffusion, supporting [cuda (by default), # rocm, neuron, cpu, openvino] "SGLANG_DIFFUSION_TARGET_DEVICE": _lazy_str( "SGLANG_DIFFUSION_TARGET_DEVICE", "cuda" ), # Maximum number of compilation jobs to run in parallel. # By default this is the number of CPUs "MAX_JOBS": _lazy_str("MAX_JOBS"), # Number of threads to use for nvcc # By default this is 1. # If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU. "NVCC_THREADS": _lazy_str("NVCC_THREADS"), # If set, sgl_diffusion will use precompiled binaries (*.so) "SGLANG_DIFFUSION_USE_PRECOMPILED": _lazy_bool_any( [ "SGLANG_DIFFUSION_USE_PRECOMPILED", "SGLANG_DIFFUSION_PRECOMPILED_WHEEL_LOCATION", ] ), # CMake build type # If not set, defaults to "Debug" or "RelWithDebInfo" # Available options: "Debug", "Release", "RelWithDebInfo" "CMAKE_BUILD_TYPE": _lazy_str("CMAKE_BUILD_TYPE"), # If set, sgl_diffusion will print verbose logs during installation "VERBOSE": _lazy_bool("VERBOSE"), # Root directory for SGL-diffusion configuration files # Defaults to `~/.config/sgl_diffusion` unless `XDG_CONFIG_HOME` is set # Note that this not only affects how sgl_diffusion finds its configuration files # during runtime, but also affects how sgl_diffusion installs its configuration # files during **installation**. "SGLANG_DIFFUSION_CONFIG_ROOT": _lazy_path( "SGLANG_DIFFUSION_CONFIG_ROOT", lambda: os.path.join(get_default_config_root(), "sgl_diffusion"), ), # ================== Runtime Env Vars ================== # Root directory for SGL-diffusion cache files # Defaults to `~/.cache/sgl_diffusion` unless `XDG_CACHE_HOME` is set "SGLANG_DIFFUSION_CACHE_ROOT": _lazy_path( "SGLANG_DIFFUSION_CACHE_ROOT", lambda: os.path.join(get_default_cache_root(), "sgl_diffusion"), ), # Interval in seconds to log a warning message when the ring buffer is full "SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL": _lazy_int( "SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL", 60 ), # Path to the NCCL library file. It is needed because nccl>=2.19 brought # by PyTorch contains a bug: https://github.com/NVIDIA/nccl/issues/1234 "SGLANG_DIFFUSION_NCCL_SO_PATH": _lazy_str("SGLANG_DIFFUSION_NCCL_SO_PATH"), # when `SGLANG_DIFFUSION_NCCL_SO_PATH` is not set, sgl_diffusion will try to find the nccl # library file in the locations specified by `LD_LIBRARY_PATH` "LD_LIBRARY_PATH": _lazy_str("LD_LIBRARY_PATH"), # Internal flag to enable Dynamo fullgraph capture "SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE": _lazy_bool( "SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1" ), # local rank of the process in the distributed setting, used to determine # the GPU device id "LOCAL_RANK": _lazy_int("LOCAL_RANK", 0), # used to control the visible devices in the distributed setting "CUDA_VISIBLE_DEVICES": _lazy_str("CUDA_VISIBLE_DEVICES"), # timeout for each iteration in the engine "SGLANG_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S": _lazy_int( "SGLANG_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S", 60 ), # Logging configuration # If set to 0, sgl_diffusion will not configure logging # If set to 1, sgl_diffusion will configure logging using the default configuration # or the configuration file specified by SGLANG_DIFFUSION_LOGGING_CONFIG_PATH "SGLANG_DIFFUSION_CONFIGURE_LOGGING": _lazy_int( "SGLANG_DIFFUSION_CONFIGURE_LOGGING", 1 ), "SGLANG_DIFFUSION_LOGGING_CONFIG_PATH": _lazy_str( "SGLANG_DIFFUSION_LOGGING_CONFIG_PATH" ), # this is used for configuring the default logging level "SGLANG_DIFFUSION_LOGGING_LEVEL": _lazy_str( "SGLANG_DIFFUSION_LOGGING_LEVEL", "INFO" ), # if set, SGLANG_DIFFUSION_LOGGING_PREFIX will be prepended to all log messages "SGLANG_DIFFUSION_LOGGING_PREFIX": _lazy_str("SGLANG_DIFFUSION_LOGGING_PREFIX", ""), # Trace function calls # If set to 1, sgl_diffusion will trace function calls # Useful for debugging "SGLANG_DIFFUSION_TRACE_FUNCTION": _lazy_int("SGLANG_DIFFUSION_TRACE_FUNCTION", 0), # Path to the attention configuration file. Only used for sliding tile # attention for now. "SGLANG_DIFFUSION_ATTENTION_CONFIG": _lazy_path( "SGLANG_DIFFUSION_ATTENTION_CONFIG" ), # Optional override to force a specific attention backend (e.g. "aiter") "SGLANG_DIFFUSION_ATTENTION_BACKEND": _lazy_str( "SGLANG_DIFFUSION_ATTENTION_BACKEND" ), # Use dedicated multiprocess context for workers. # Both spawn and fork work "SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD": _lazy_str( "SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD", "fork" ), # Enables torch profiler if set. Path to the directory where torch profiler # traces are saved. Note that it must be an absolute path. "SGLANG_DIFFUSION_TORCH_PROFILER_DIR": _lazy_path( "SGLANG_DIFFUSION_TORCH_PROFILER_DIR" ), # If set, sgl_diffusion will run in development mode, which will enable # some additional endpoints for developing and debugging, # e.g. `/reset_prefix_cache` "SGLANG_DIFFUSION_SERVER_DEV_MODE": _lazy_bool("SGLANG_DIFFUSION_SERVER_DEV_MODE"), # If set, sgl_diffusion will enable stage logging, which will print the time # taken for each stage "SGLANG_DIFFUSION_STAGE_LOGGING": _lazy_bool("SGLANG_DIFFUSION_STAGE_LOGGING"), "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D": _lazy_bool( "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D", "false" ), # ================== cache-dit Env Vars ================== # Enable cache-dit acceleration for DiT inference "SGLANG_CACHE_DIT_ENABLED": _lazy_bool("SGLANG_CACHE_DIT_ENABLED"), # Number of first blocks to always compute (DBCache F parameter) "SGLANG_CACHE_DIT_FN": _lazy_int("SGLANG_CACHE_DIT_FN", 1), # Number of last blocks to always compute (DBCache B parameter) "SGLANG_CACHE_DIT_BN": _lazy_int("SGLANG_CACHE_DIT_BN", 0), # Warmup steps before caching (DBCache W parameter) "SGLANG_CACHE_DIT_WARMUP": _lazy_int("SGLANG_CACHE_DIT_WARMUP", 4), # Residual difference threshold (DBCache R parameter) "SGLANG_CACHE_DIT_RDT": _lazy_float("SGLANG_CACHE_DIT_RDT", 0.24), # Maximum continuous cached steps (DBCache MC parameter) "SGLANG_CACHE_DIT_MC": _lazy_int("SGLANG_CACHE_DIT_MC", 3), # Enable TaylorSeer calibrator "SGLANG_CACHE_DIT_TAYLORSEER": _lazy_bool("SGLANG_CACHE_DIT_TAYLORSEER", "false"), # TaylorSeer order (1 or 2) "SGLANG_CACHE_DIT_TS_ORDER": _lazy_int("SGLANG_CACHE_DIT_TS_ORDER", 1), # SCM preset: none, slow, medium, fast, ultra "SGLANG_CACHE_DIT_SCM_PRESET": _lazy_str("SGLANG_CACHE_DIT_SCM_PRESET", "none"), # SCM custom compute bins (e.g., "8,3,3,2,2") "SGLANG_CACHE_DIT_SCM_COMPUTE_BINS": _lazy_str("SGLANG_CACHE_DIT_SCM_COMPUTE_BINS"), # SCM custom cache bins (e.g., "1,2,2,2,3") "SGLANG_CACHE_DIT_SCM_CACHE_BINS": _lazy_str("SGLANG_CACHE_DIT_SCM_CACHE_BINS"), # SCM policy: dynamic or static "SGLANG_CACHE_DIT_SCM_POLICY": _lazy_str("SGLANG_CACHE_DIT_SCM_POLICY", "dynamic"), # model loading "SGLANG_USE_RUNAI_MODEL_STREAMER": _lazy_bool( "SGLANG_USE_RUNAI_MODEL_STREAMER", "true" ), # ROCm: use AITer GroupNorm in VAE for improved performance "SGLANG_USE_ROCM_VAE": _lazy_bool("SGLANG_USE_ROCM_VAE"), } # Add cache-dit Secondary Transformer Env Vars via programmatic generation to reduce duplication _CACHE_DIT_SECONDARY_CONFIGS = [ ("FN", int, "1"), ("BN", int, "0"), ("WARMUP", int, "4"), ("RDT", float, "0.24"), ("MC", int, "3"), ("TS_ORDER", int, "1"), ] def _create_secondary_getter(suffix, type_func, default_val): primary_key = f"SGLANG_CACHE_DIT_{suffix}" secondary_key = f"SGLANG_CACHE_DIT_SECONDARY_{suffix}" def _getter(): val = os.getenv(secondary_key) if val is not None: return type_func(val) return type_func(os.getenv(primary_key, str(default_val))) return secondary_key, _getter for suffix, type_func, default_val in _CACHE_DIT_SECONDARY_CONFIGS: key, getter = _create_secondary_getter(suffix, type_func, default_val) environment_variables[key] = getter # Special handling for boolean secondary var (TaylorSeer) def _secondary_taylorseer_getter(): return get_bool_env_var( "SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER", default=os.getenv("SGLANG_CACHE_DIT_TAYLORSEER", "false"), ) environment_variables["SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER"] = ( _secondary_taylorseer_getter ) # end-env-vars-definition def __getattr__(name: str): # lazy evaluation of environment variables if name in environment_variables: return environment_variables[name]() raise AttributeError(f"module {__name__!r} has no attribute {name!r}") def __dir__(): return list(environment_variables.keys()) ================================================ FILE: python/sglang/multimodal_gen/registry.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Central registry for multimodal models. This module provides a centralized registry for multimodal models, including pipelines and sampling parameters. It allows for easy registration and retrieval of model information based on model paths or other identifiers. """ import dataclasses import importlib import os import pkgutil from functools import lru_cache from typing import ( TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Type, Union, ) if TYPE_CHECKING: from sglang.multimodal_gen.runtime.server_args import Backend from sglang.multimodal_gen.configs.pipeline_configs import ( FastHunyuanConfig, FluxPipelineConfig, HeliosDistilledConfig, HeliosMidConfig, HeliosT2VConfig, HunyuanConfig, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig, ZImagePipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.flux import ( Flux2KleinPipelineConfig, Flux2PipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.glm_image import ( GlmImagePipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( Hunyuan3D2PipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.mova import ( MOVA360PConfig, MOVA720PConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImageEditPipelineConfig, QwenImageEditPlus_2511_PipelineConfig, QwenImageEditPlusPipelineConfig, QwenImageLayeredPipelineConfig, QwenImagePipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.wan import ( FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config, TurboWanI2V720Config, TurboWanT2V480PConfig, Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config, ) from sglang.multimodal_gen.configs.sample.flux import ( Flux2KleinSamplingParams, FluxSamplingParams, ) from sglang.multimodal_gen.configs.sample.glmimage import GlmImageSamplingParams from sglang.multimodal_gen.configs.sample.helios import ( HeliosDistilledSamplingParams, HeliosMidSamplingParams, HeliosT2VSamplingParams, ) from sglang.multimodal_gen.configs.sample.hunyuan import ( FastHunyuanSamplingParam, HunyuanSamplingParams, ) from sglang.multimodal_gen.configs.sample.hunyuan3d import Hunyuan3DSamplingParams from sglang.multimodal_gen.configs.sample.ltx_2 import LTX2SamplingParams from sglang.multimodal_gen.configs.sample.mova import ( MOVA_360P_SamplingParams, MOVA_720P_SamplingParams, ) from sglang.multimodal_gen.configs.sample.qwenimage import ( QwenImage2512SamplingParams, QwenImageEditPlusSamplingParams, QwenImageLayeredSamplingParams, QwenImageSamplingParams, ) from sglang.multimodal_gen.configs.sample.sana import SanaSamplingParams from sglang.multimodal_gen.configs.sample.wan import ( FastWanT2V480PConfig, Turbo_Wan2_2_I2V_A14B_SamplingParam, Wan2_1_Fun_1_3B_InP_SamplingParams, Wan2_2_I2V_A14B_SamplingParam, Wan2_2_T2V_A14B_SamplingParam, Wan2_2_TI2V_5B_SamplingParam, WanI2V_14B_480P_SamplingParam, WanI2V_14B_720P_SamplingParam, WanT2V_1_3B_SamplingParams, WanT2V_14B_SamplingParams, ) from sglang.multimodal_gen.configs.sample.zimage import ( ZImageSamplingParams, ZImageTurboSamplingParams, ) from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( maybe_download_model_index, verify_model_config_and_directory, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # --- Part 1: Pipeline Discovery --- _PIPELINE_REGISTRY: Dict[str, Type[ComposedPipelineBase]] = {} # Registry for pipeline configuration classes (for safetensors files without model_index.json) # Maps pipeline_class_name -> (PipelineConfig class, SamplingParams class) _PIPELINE_CONFIG_REGISTRY: Dict[str, Tuple[Type[PipelineConfig], Type[Any]]] = {} def _discover_and_register_pipelines(): """ Automatically discover and register all ComposedPipelineBase subclasses. This function scans the 'sglang.multimodal_gen.runtime.pipelines' package, finds modules with an 'EntryClass' attribute, and maps the class's 'pipeline_name' to the class itself in a global registry. """ if _PIPELINE_REGISTRY: # run only once return package_name = "sglang.multimodal_gen.runtime.pipelines" package = importlib.import_module(package_name) for _, module_name, ispkg in pkgutil.walk_packages( package.__path__, package.__name__ + "." ): if not ispkg: pipeline_module = importlib.import_module(module_name) if hasattr(pipeline_module, "EntryClass"): entry_cls = pipeline_module.EntryClass entry_cls_list = ( [entry_cls] if not isinstance(entry_cls, list) else entry_cls ) for cls in entry_cls_list: if not issubclass(cls, ComposedPipelineBase): continue if cls.pipeline_name in _PIPELINE_REGISTRY: logger.warning( f"Duplicate pipeline name '{cls.pipeline_name}' found. Overwriting." ) _PIPELINE_REGISTRY[cls.pipeline_name] = cls # Special handling for ComfyUI Pipelines: # Auto-register config classes if Pipeline class has them defined # since comfyui get model from a single weight file, so we need to register the config classes here if hasattr(cls, "pipeline_config_cls") and hasattr( cls, "sampling_params_cls" ): _PIPELINE_CONFIG_REGISTRY[cls.pipeline_name] = ( cls.pipeline_config_cls, cls.sampling_params_cls, ) logger.debug( f"Auto-registered config classes for pipeline '{cls.pipeline_name}': " f"PipelineConfig={cls.pipeline_config_cls.__name__}, " f"SamplingParams={cls.sampling_params_cls.__name__}" ) logger.debug( f"Registering pipelines complete, {len(_PIPELINE_REGISTRY)} pipelines registered" ) def get_pipeline_config_classes( pipeline_class_name: str, ) -> Tuple[Type[PipelineConfig], Type[Any]] | None: """ Get the configuration classes for a pipeline. """ # Ensure pipelines are discovered first _discover_and_register_pipelines() return _PIPELINE_CONFIG_REGISTRY.get(pipeline_class_name) # --- Part 2: Config Registration --- @dataclasses.dataclass class ConfigInfo: """Encapsulates all configuration information required to register a diffusers model within this framework.""" sampling_param_cls: Any pipeline_config_cls: Type[PipelineConfig] # The central registry mapping a model name to its configuration information _CONFIG_REGISTRY: Dict[str, ConfigInfo] = {} # Mappings from Hugging Face model paths to our internal model names _MODEL_HF_PATH_TO_NAME: Dict[str, str] = {} # Detectors to identify model families from paths or class names _MODEL_NAME_DETECTORS: List[Tuple[str, Callable[[str], bool]]] = [] def register_configs( sampling_param_cls: Any, pipeline_config_cls: Type[PipelineConfig], hf_model_paths: Optional[List[str]] = None, model_detectors: Optional[List[Callable[[str], bool]]] = None, ): """ Registers configuration classes for a new model family. """ model_id = str(len(_CONFIG_REGISTRY)) _CONFIG_REGISTRY[model_id] = ConfigInfo( sampling_param_cls=sampling_param_cls, pipeline_config_cls=pipeline_config_cls, ) if hf_model_paths: for path in hf_model_paths: if path in _MODEL_HF_PATH_TO_NAME: logger.warning( f"Model path '{path}' is already mapped to '{_MODEL_HF_PATH_TO_NAME[path]}' and will be overwritten by '{model_id}'." ) _MODEL_HF_PATH_TO_NAME[path] = model_id if model_detectors: for detector in model_detectors: _MODEL_NAME_DETECTORS.append((model_id, detector)) def get_model_short_name(model_id: str) -> str: if "/" in model_id: return model_id.rstrip("/").split("/")[-1] else: return model_id @lru_cache(maxsize=1) def _get_config_info( model_path: str, model_id: Optional[str] = None ) -> Optional[ConfigInfo]: """ Gets the ConfigInfo for a given model path using mappings and detectors. """ all_model_hf_paths = sorted(_MODEL_HF_PATH_TO_NAME.keys(), key=len, reverse=True) # 0. Explicit model_id override: match by short name if model_id is not None: model_id_lower = model_id.lower() for registered_hf_id in all_model_hf_paths: if get_model_short_name(registered_hf_id).lower() == model_id_lower: logger.debug( f"Resolved model via explicit --model-id '{model_id}' → '{registered_hf_id}'." ) return _CONFIG_REGISTRY.get(_MODEL_HF_PATH_TO_NAME[registered_hf_id]) logger.warning( f"--model-id '{model_id}' did not match any registered model; " "falling back to automatic detection." ) # 1. Exact match if model_path in _MODEL_HF_PATH_TO_NAME: model_id = _MODEL_HF_PATH_TO_NAME[model_path] logger.debug(f"Resolved model path '{model_path}' from exact path match.") return _CONFIG_REGISTRY.get(model_id) # 2. Partial match: find the best (longest) match against all registered model hf paths. model_short_name = get_model_short_name(model_path.lower()) for registered_model_hf_id in all_model_hf_paths: registered_model_name = get_model_short_name(registered_model_hf_id.lower()) if registered_model_name in model_short_name: logger.debug( f"Resolved model name '{registered_model_hf_id}' from partial path match." ) model_id = _MODEL_HF_PATH_TO_NAME[registered_model_hf_id] return _CONFIG_REGISTRY.get(model_id) # 3. Use detectors if os.path.exists(model_path): config = verify_model_config_and_directory(model_path) else: config = maybe_download_model_index(model_path) pipeline_name = config.get("_class_name", "").lower() matched_model_names = [] for model_id, detector in _MODEL_NAME_DETECTORS: if detector(model_path.lower()) or detector(pipeline_name): logger.debug( f"Matched model name '{model_id}' using a registered detector." ) matched_model_names += [model_id] if len(matched_model_names) >= 1: if len(matched_model_names) > 1: logger.warning( f"More than one model name is matched, using the first matched" ) model_id = matched_model_names[0] return _CONFIG_REGISTRY.get(model_id) else: raise RuntimeError(f"No model info found for model path: {model_path}") # --- Part 3: Main Resolver --- @dataclasses.dataclass class ModelInfo: """ Encapsulates all configuration information required to register a diffusers model within this framework. """ pipeline_cls: Type[ComposedPipelineBase] sampling_param_cls: Any pipeline_config_cls: Type[PipelineConfig] def _get_diffusers_model_info( model_path: Optional[str] = None, model_id: Optional[str] = None, ) -> ModelInfo: """ Get model info for diffusers backend. Returns a ModelInfo with DiffusersPipeline and generic configs. When model_path is provided and has a registered native config, inherits task_type from it so that validation (e.g. accepts_image_input) works correctly even under the diffusers backend. """ from sglang.multimodal_gen.configs.pipeline_configs.diffusers_generic import ( DiffusersGenericPipelineConfig, ) from sglang.multimodal_gen.configs.sample.diffusers_generic import ( DiffusersGenericSamplingParams, ) from sglang.multimodal_gen.runtime.pipelines.diffusers_pipeline import ( DiffusersPipeline, ) sampling_param_cls = DiffusersGenericSamplingParams pipeline_config_cls = DiffusersGenericPipelineConfig # If there is a registered native config for this model, inherit its task_type if model_path is not None: config_info = _get_config_info(model_path, model_id=model_id) if config_info is not None: sampling_param_cls = config_info.sampling_param_cls native_task_type = config_info.pipeline_config_cls.task_type if native_task_type != DiffusersGenericPipelineConfig.task_type: pipeline_config_cls = dataclasses.make_dataclass( "DiffusersGenericPipelineConfig", [ ( "task_type", type(native_task_type), dataclasses.field(default=native_task_type), ) ], bases=(DiffusersGenericPipelineConfig,), ) logger.debug( "Inherited task_type=%s from native config for diffusers backend", native_task_type.name, ) return ModelInfo( pipeline_cls=DiffusersPipeline, sampling_param_cls=sampling_param_cls, pipeline_config_cls=pipeline_config_cls, ) @lru_cache(maxsize=1) def get_model_info( model_path: str, backend: Optional[Union[str, "Backend"]] = None, model_id: Optional[str] = None, ) -> Optional[ModelInfo]: """ Resolves all necessary classes (pipeline, sampling, config) for a given model path. This function serves as the main entry point for model resolution. It performs two main tasks: 1. Dynamically resolves the pipeline class by reading 'model_index.json' and matching '_class_name' against an auto-discovered registry of pipeline implementations. 2. Resolves the associated configuration classes (for sampling and pipeline) using a manually registered mapping based on the model path. Args: backend: Backend to use ('auto', 'sglang', 'diffusers'). If None, uses 'auto'. """ # import Backend enum here to avoid circular imports from sglang.multimodal_gen.runtime.server_args import Backend # Normalize backend if backend is None: backend = Backend.AUTO elif isinstance(backend, str): backend = Backend.from_string(backend) # Handle explicit diffusers backend if backend == Backend.DIFFUSERS: logger.info( "Using diffusers backend for model '%s' (explicitly requested)", model_path ) return _get_diffusers_model_info(model_path=model_path, model_id=model_id) # For AUTO or SGLANG backend, try native implementation first # 1. Discover all available pipeline classes and cache them _discover_and_register_pipelines() # Detect quantized models and fallback to diffusers is_quantized = any(q in model_path.lower() for q in ["-4bit", "-awq", "-gptq"]) if is_quantized and backend != Backend.DIFFUSERS: logger.info( "Detected a quantized model format ('%s'). " "The native sglang-diffusion engine currently only supports BF16/FP16. " "Falling back to diffusers backend.", model_path, ) return _get_diffusers_model_info(model_path=model_path, model_id=model_id) # 2. Get pipeline class - check non-diffusers models first pipeline_class_name = get_non_diffusers_pipeline_name(model_path) if pipeline_class_name: # Known non-diffusers model, skip model_index.json download logger.debug( f"Using registered pipeline '{pipeline_class_name}' for non-diffusers model '{model_path}'" ) else: # Try to get from model_index.json try: if os.path.exists(model_path): config = verify_model_config_and_directory(model_path) else: config = maybe_download_model_index(model_path) except Exception as e: logger.error(f"Could not read model config for '{model_path}': {e}") if backend == Backend.AUTO: logger.info("Falling back to diffusers backend") return _get_diffusers_model_info( model_path=model_path, model_id=model_id ) return None pipeline_class_name = config.get("_class_name") if not pipeline_class_name: logger.error( f"'_class_name' not found in model_index.json for '{model_path}'" ) if backend == Backend.AUTO: logger.info("Falling back to diffusers backend") return _get_diffusers_model_info( model_path=model_path, model_id=model_id ) return None pipeline_cls = _PIPELINE_REGISTRY.get(pipeline_class_name) if not pipeline_cls: if backend == Backend.AUTO: logger.warning( f"Pipeline class '{pipeline_class_name}' specified in '{model_path}' has no native sglang support. " f"Falling back to diffusers backend." ) return _get_diffusers_model_info(model_path=model_path, model_id=model_id) else: logger.error( f"Pipeline class '{pipeline_class_name}' specified in '{model_path}' is not a registered EntryClass in the framework. " f"Available pipelines: {list(_PIPELINE_REGISTRY.keys())}. " f"Consider using --backend diffusers to use vanilla diffusers pipeline." ) return None # 3. Get configuration classes (sampling, pipeline config) config_info = _get_config_info(model_path, model_id=model_id) if not config_info: if backend == Backend.AUTO: logger.warning( f"Could not resolve native configuration for model '{model_path}'. " f"Falling back to diffusers backend." ) return _get_diffusers_model_info(model_path=model_path, model_id=model_id) else: logger.error( f"Could not resolve configuration for model '{model_path}'. " "It is not a registered model path or detected by any registered model family detectors. " f"Known model paths: {list(_MODEL_HF_PATH_TO_NAME.keys())}. " f"Consider using --backend diffusers to use vanilla diffusers pipeline." ) return None # 4. Combine and return the complete model info logger.debug("Using native sglang backend for model '%s'", model_path) model_info = ModelInfo( pipeline_cls=pipeline_cls, sampling_param_cls=config_info.sampling_param_cls, pipeline_config_cls=config_info.pipeline_config_cls, ) logger.debug(f"Found model info: {model_info}") return model_info # Registration of model configs def _register_configs(): # LTX-2 register_configs( sampling_param_cls=LTX2SamplingParams, pipeline_config_cls=LTX2PipelineConfig, model_detectors=[ lambda path: "ltx" in path.lower() and "video" in path.lower(), lambda path: "ltx-2" in path.lower(), ], ) # Hunyuan register_configs( sampling_param_cls=HunyuanSamplingParams, pipeline_config_cls=HunyuanConfig, hf_model_paths=[ "hunyuanvideo-community/HunyuanVideo", ], model_detectors=[lambda hf_id: "hunyuanvideo" in hf_id.lower()], ) register_configs( sampling_param_cls=FastHunyuanSamplingParam, pipeline_config_cls=FastHunyuanConfig, hf_model_paths=[ "FastVideo/FastHunyuan-diffusers", ], ) # Wan register_configs( sampling_param_cls=WanT2V_1_3B_SamplingParams, pipeline_config_cls=WanT2V480PConfig, hf_model_paths=[ "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", ], model_detectors=[lambda hf_id: "wanpipeline" in hf_id.lower()], ) register_configs( sampling_param_cls=WanT2V_1_3B_SamplingParams, pipeline_config_cls=TurboWanT2V480PConfig, hf_model_paths=[ "IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers", ], ) register_configs( sampling_param_cls=WanT2V_14B_SamplingParams, pipeline_config_cls=WanT2V720PConfig, hf_model_paths=[ "Wan-AI/Wan2.1-T2V-14B-Diffusers", ], ) register_configs( sampling_param_cls=WanT2V_14B_SamplingParams, pipeline_config_cls=TurboWanT2V480PConfig, hf_model_paths=[ "IPostYellow/TurboWan2.1-T2V-14B-Diffusers", "IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers", ], ) register_configs( sampling_param_cls=WanI2V_14B_480P_SamplingParam, pipeline_config_cls=WanI2V480PConfig, hf_model_paths=[ "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", ], model_detectors=[lambda hf_id: "wanimagetovideo" in hf_id.lower()], ) register_configs( sampling_param_cls=WanI2V_14B_720P_SamplingParam, pipeline_config_cls=WanI2V720PConfig, hf_model_paths=[ "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers", ], ) register_configs( sampling_param_cls=Turbo_Wan2_2_I2V_A14B_SamplingParam, pipeline_config_cls=TurboWanI2V720Config, hf_model_paths=[ "IPostYellow/TurboWan2.2-I2V-A14B-Diffusers", ], ) register_configs( sampling_param_cls=Wan2_1_Fun_1_3B_InP_SamplingParams, pipeline_config_cls=WanI2V480PConfig, hf_model_paths=[ "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers", ], ) register_configs( sampling_param_cls=Wan2_2_TI2V_5B_SamplingParam, pipeline_config_cls=Wan2_2_TI2V_5B_Config, hf_model_paths=[ "Wan-AI/Wan2.2-TI2V-5B-Diffusers", ], ) register_configs( sampling_param_cls=Wan2_2_TI2V_5B_SamplingParam, pipeline_config_cls=FastWan2_2_TI2V_5B_Config, hf_model_paths=[ "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", "FastVideo/FastWan2.2-TI2V-5B-Diffusers", ], ) register_configs( sampling_param_cls=Wan2_2_T2V_A14B_SamplingParam, pipeline_config_cls=Wan2_2_T2V_A14B_Config, hf_model_paths=["Wan-AI/Wan2.2-T2V-A14B-Diffusers"], ) register_configs( sampling_param_cls=Wan2_2_I2V_A14B_SamplingParam, pipeline_config_cls=Wan2_2_I2V_A14B_Config, hf_model_paths=["Wan-AI/Wan2.2-I2V-A14B-Diffusers"], ) register_configs( sampling_param_cls=FastWanT2V480PConfig, pipeline_config_cls=FastWan2_1_T2V_480P_Config, hf_model_paths=[ "FastVideo/FastWan2.1-T2V-1.3B-Diffusers", ], ) # MOVA register_configs( sampling_param_cls=MOVA_360P_SamplingParams, pipeline_config_cls=MOVA360PConfig, model_detectors=[ lambda hf_id: "mova" in hf_id.lower() and "360p" in hf_id.lower() ], ) register_configs( sampling_param_cls=MOVA_720P_SamplingParams, pipeline_config_cls=MOVA720PConfig, model_detectors=[ lambda hf_id: "mova" in hf_id.lower() and "720p" in hf_id.lower() ], ) # FLUX register_configs( sampling_param_cls=FluxSamplingParams, pipeline_config_cls=FluxPipelineConfig, hf_model_paths=[ "black-forest-labs/FLUX.1-dev", ], model_detectors=[lambda hf_id: "flux.1" in hf_id.lower()], ) register_configs( sampling_param_cls=Flux2KleinSamplingParams, pipeline_config_cls=Flux2KleinPipelineConfig, hf_model_paths=[ "black-forest-labs/FLUX.2-klein-4B", "black-forest-labs/FLUX.2-klein-9B", ], model_detectors=[ lambda hf_id: "flux.2-klein" in hf_id.lower() or "flux2-klein" in hf_id.lower() ], ) register_configs( sampling_param_cls=FluxSamplingParams, pipeline_config_cls=Flux2PipelineConfig, hf_model_paths=[ "black-forest-labs/FLUX.2-dev", ], model_detectors=[ lambda hf_id: "flux.2" in hf_id.lower() and "klein" not in hf_id.lower() ], ) register_configs( sampling_param_cls=ZImageTurboSamplingParams, pipeline_config_cls=ZImagePipelineConfig, hf_model_paths=[ "Tongyi-MAI/Z-Image-Turbo", ], model_detectors=[lambda hf_id: "z-image-turbo" in hf_id.lower()], ) register_configs( sampling_param_cls=ZImageSamplingParams, pipeline_config_cls=ZImagePipelineConfig, hf_model_paths=[ "Tongyi-MAI/Z-Image", ], model_detectors=[ lambda hf_id: "z-image" in hf_id.lower() and "turbo" not in hf_id.lower() ], ) # Qwen-Image register_configs( sampling_param_cls=QwenImageSamplingParams, pipeline_config_cls=QwenImagePipelineConfig, hf_model_paths=["Qwen/Qwen-Image"], model_detectors=[ lambda hf_id: "qwen-image" in hf_id.lower() and "edit" not in hf_id.lower() and "layered" not in hf_id.lower() and "2512" not in hf_id.lower() ], ) register_configs( sampling_param_cls=QwenImage2512SamplingParams, pipeline_config_cls=QwenImagePipelineConfig, hf_model_paths=["Qwen/Qwen-Image-2512"], model_detectors=[lambda hf_id: "qwen-image-2512" in hf_id.lower()], ) register_configs( sampling_param_cls=QwenImageSamplingParams, pipeline_config_cls=QwenImageEditPipelineConfig, hf_model_paths=["Qwen/Qwen-Image-Edit"], model_detectors=[ lambda hf_id: "qwen-image-edit" in hf_id.lower() and "2509" not in hf_id.lower() and "2511" not in hf_id.lower() ], ) register_configs( sampling_param_cls=QwenImageEditPlusSamplingParams, pipeline_config_cls=QwenImageEditPlusPipelineConfig, hf_model_paths=["Qwen/Qwen-Image-Edit-2509"], model_detectors=[lambda hf_id: "qwen-image-edit-2509" in hf_id.lower()], ) register_configs( sampling_param_cls=QwenImageEditPlusSamplingParams, pipeline_config_cls=QwenImageEditPlus_2511_PipelineConfig, hf_model_paths=["Qwen/Qwen-Image-Edit-2511"], model_detectors=[lambda hf_id: "qwen-image-edit-2511" in hf_id.lower()], ) register_configs( sampling_param_cls=QwenImageLayeredSamplingParams, pipeline_config_cls=QwenImageLayeredPipelineConfig, hf_model_paths=["Qwen/Qwen-Image-Layered"], model_detectors=[lambda hf_id: "qwen-image-layered" in hf_id.lower()], ) register_configs( sampling_param_cls=GlmImageSamplingParams, pipeline_config_cls=GlmImagePipelineConfig, model_detectors=[lambda hf_id: "glm-image" in hf_id.lower()], ) register_configs( sampling_param_cls=Hunyuan3DSamplingParams, pipeline_config_cls=Hunyuan3D2PipelineConfig, hf_model_paths=[ "tencent/Hunyuan3D-2", ], model_detectors=[lambda hf_id: "hunyuan3d" in hf_id.lower()], ) # Helios register_configs( sampling_param_cls=HeliosT2VSamplingParams, pipeline_config_cls=HeliosT2VConfig, hf_model_paths=[ "BestWishYsh/Helios-Base", ], model_detectors=[ lambda hf_id: "helios" in hf_id.lower() and "mid" not in hf_id.lower() and "distill" not in hf_id.lower() ], ) register_configs( sampling_param_cls=HeliosMidSamplingParams, pipeline_config_cls=HeliosMidConfig, hf_model_paths=[ "BestWishYsh/Helios-Mid", ], ) register_configs( sampling_param_cls=HeliosDistilledSamplingParams, pipeline_config_cls=HeliosDistilledConfig, hf_model_paths=[ "BestWishYsh/Helios-Distilled", ], ) # SANA register_configs( sampling_param_cls=SanaSamplingParams, pipeline_config_cls=SanaPipelineConfig, hf_model_paths=[ "Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers", "Efficient-Large-Model/SANA1.5_4.8B_1024px_diffusers", "Efficient-Large-Model/Sana_1600M_1024px_diffusers", "Efficient-Large-Model/Sana_600M_1024px_diffusers", "Efficient-Large-Model/Sana_1600M_512px_diffusers", "Efficient-Large-Model/Sana_600M_512px_diffusers", ], model_detectors=[lambda hf_id: "sana" in hf_id.lower()], ) _register_configs() # Known non-diffusers multimodal model patterns # Maps pattern -> pipeline_name for models that don't have model_index.json _NON_DIFFUSERS_MULTIMODAL_PATTERNS: Dict[str, str] = { "hunyuan3d": "Hunyuan3D2Pipeline", } def is_known_non_diffusers_multimodal_model(model_path: str) -> bool: model_path_lower = model_path.lower() return any( pattern in model_path_lower for pattern in _NON_DIFFUSERS_MULTIMODAL_PATTERNS ) def get_non_diffusers_pipeline_name(model_path: str) -> Optional[str]: """Get the pipeline name for a known non-diffusers model.""" model_path_lower = model_path.lower() for pattern, pipeline_name in _NON_DIFFUSERS_MULTIMODAL_PATTERNS.items(): if pattern in model_path_lower: return pipeline_name return None ================================================ FILE: python/sglang/multimodal_gen/runtime/cache/__init__.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Cache acceleration module for SGLang-diffusion This module provides various caching strategies to accelerate diffusion transformer (DiT) inference: - TeaCache: Temporal similarity-based caching for diffusion models - cache-dit integration: Block-level caching with DBCache and TaylorSeer """ from sglang.multimodal_gen.runtime.cache.cache_dit_integration import ( CacheDitConfig, enable_cache_on_dual_transformer, enable_cache_on_transformer, get_scm_mask, ) from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheContext, TeaCacheMixin __all__ = [ # TeaCache (always available) "TeaCacheContext", "TeaCacheMixin", # cache-dit integration (lazy-loaded, requires cache-dit package) "CacheDitConfig", "enable_cache_on_transformer", "enable_cache_on_dual_transformer", "get_scm_mask", ] ================================================ FILE: python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ cache-dit integration module for SGLang DiT pipelines. This module provides helper functions to enable cache-dit acceleration on transformer modules in SGLang's modular pipeline architecture. """ from dataclasses import dataclass from typing import List, Optional import torch import torch.distributed as dist from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_ring_parallel_world_size, get_tp_world_size, get_ulysses_parallel_world_size, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) import cache_dit from cache_dit import ( BlockAdapter, DBCacheConfig, ForwardPattern, ParamsModifier, TaylorSeerCalibratorConfig, steps_mask, ) from cache_dit.caching.block_adapters import BlockAdapterRegister from cache_dit.parallelism import ParallelismBackend, ParallelismConfig from sglang.multimodal_gen.runtime.distributed.parallel_state import get_dit_group _original_similarity = None def _patch_cache_dit_similarity(): from cache_dit.caching.cache_contexts import cache_manager global _original_similarity if _original_similarity is not None: return _original_similarity = cache_manager.CachedContextManager.similarity def patched_similarity(self, t1, t2, *, threshold, parallelized=False, prefix="Fn"): if not parallelized: return _original_similarity( self, t1, t2, threshold=threshold, parallelized=parallelized, prefix=prefix, ) sp_group = getattr(self, "_sglang_sp_group", None) tp_group = getattr(self, "_sglang_tp_group", None) tp_sp_group = getattr(self, "_sglang_tp_sp_group", None) target_group = tp_sp_group or sp_group or tp_group if target_group is None: return _original_similarity( self, t1, t2, threshold=threshold, parallelized=parallelized, prefix=prefix, ) # Adapted from https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/caching/cache_contexts/cache_manager.py#L495-L523 condition_thresh = self.get_important_condition_threshold() if condition_thresh > 0.0: raw_diff = (t1 - t2).abs() token_m_df = raw_diff.mean(dim=-1) token_m_t1 = t1.abs().mean(dim=-1) token_diff = token_m_df / token_m_t1 condition = token_diff > condition_thresh if condition.sum() > 0: condition = condition.unsqueeze(-1).expand_as(raw_diff) mean_diff = raw_diff[condition].mean() mean_t1 = t1[condition].abs().mean() else: mean_diff = (t1 - t2).abs().mean() mean_t1 = t1.abs().mean() else: mean_diff = (t1 - t2).abs().mean() mean_t1 = t1.abs().mean() dist.all_reduce(mean_diff, op=dist.ReduceOp.AVG, group=target_group) dist.all_reduce(mean_t1, op=dist.ReduceOp.AVG, group=target_group) diff = (mean_diff / mean_t1).item() self.add_residual_diff(diff) return diff < threshold cache_manager.CachedContextManager.similarity = patched_similarity def _build_parallelism_config( sp_group: Optional[torch.distributed.ProcessGroup], tp_group: Optional[torch.distributed.ProcessGroup], ): if sp_group is None and tp_group is None: return None ulysses_size = None ring_size = None if sp_group is not None: ulysses_size = get_ulysses_parallel_world_size() ring_size = get_ring_parallel_world_size() tp_size = None if tp_group is not None: tp_size = get_tp_world_size() return ParallelismConfig( backend=ParallelismBackend.AUTO, ulysses_size=ulysses_size, ring_size=ring_size, tp_size=tp_size, ) def _mark_transformer_parallelized(transformer, config, sp_group, tp_group): if config is None: return transformer._is_parallelized = True transformer._parallelism_config = config def get_scm_mask( preset: str, num_inference_steps: int, compute_bins: Optional[List[int]] = None, cache_bins: Optional[List[int]] = None, ) -> Optional[List[int]]: """ Get SCM mask using cache-dit's steps_mask(). This is a thin wrapper that delegates to cache-dit's built-in steps_mask() function which handles all presets and scaling logic. Args: preset: Preset name ("none", "slow", "medium", "fast", "ultra"). compute_bins: Custom compute bins (overrides preset). cache_bins: Custom cache bins (overrides preset). Returns: SCM mask list (1=compute, 0=cache), or None if disabled. """ if preset == "none" and not (compute_bins and cache_bins): return None # Use cache-dit's steps_mask() directly mask = steps_mask( compute_bins=compute_bins, cache_bins=cache_bins, total_steps=num_inference_steps, mask_policy=preset if preset != "none" else "medium", ) compute_count = sum(mask) cache_count = len(mask) - compute_count logger.info( "SCM: generated mask with %d compute steps, %d cache steps (preset=%s)", compute_count, cache_count, preset, ) return mask @dataclass class CacheDitConfig: """Configuration for cache-dit integration. Attributes: enabled: Whether to enable cache-dit acceleration. Fn_compute_blocks: Number of first blocks to always compute (DBCache F). Bn_compute_blocks: Number of last blocks to always compute (DBCache B). max_warmup_steps: Number of warmup steps before caching starts (DBCache W). residual_diff_threshold: Threshold for residual difference (DBCache R). max_continuous_cached_steps: Maximum consecutive cached steps (DBCache MC). enable_taylorseer: Whether to enable TaylorSeer calibrator. taylorseer_order: Order of Taylor expansion (1 or 2). num_inference_steps: Total number of inference steps (required for transformer-only mode). steps_computation_mask: Binary mask for step-level caching (1=compute, 0=cache). Generated by get_scm_mask() (wrapper around cache_dit.steps_mask()). steps_computation_policy: Caching policy for SCM ("dynamic" or "static"). """ enabled: bool = False Fn_compute_blocks: int = 1 Bn_compute_blocks: int = 0 # Use 4 as default warmup steps instead of 8 in cache-dit, thus making # DBCache work for few steps distilled models, e.g., Z-Image w/ 8-steps. max_warmup_steps: int = 4 # Use a relatively higher residual diff threshold (namely, 0.24) as default # to allow more aggressive caching due to we have already applied max continuous # cached steps limit, otherwise, we should use a lower threshold here like 0.12. residual_diff_threshold: float = 0.24 max_continuous_cached_steps: int = 3 # TaylorSeer is not suitable for few steps distilled models, so, we choose # to disable it by default. Reference: # - From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers, # https://arxiv.org/pdf/2503.06923 # - FoCa: Forecast then Calibrate: Feature Caching as ODE for Efficient # Diffusion Transformers, https://arxiv.org/pdf/2508.16211 enable_taylorseer: bool = False taylorseer_order: int = 1 num_inference_steps: Optional[int] = None # SCM fields (generated by _maybe_enable_cache_dit from env configuration) steps_computation_mask: Optional[List[int]] = None steps_computation_policy: str = "dynamic" def enable_cache_on_transformer( transformer: torch.nn.Module, config: CacheDitConfig, model_name: str = "transformer", sp_group: Optional[torch.distributed.ProcessGroup] = None, tp_group: Optional[torch.distributed.ProcessGroup] = None, ) -> torch.nn.Module: """Enable cache-dit on a transformer module, by wrapping the module with cache-dit This function enables cache-dit acceleration using the BlockAdapterRegister for pre-registered models Args: model_name: Name of the model for logging purposes. sp_group: Sequence parallel process group (for Ulysses/Ring). tp_group: Tensor parallel process group. """ if not config.enabled: return transformer if config.num_inference_steps is None: raise ValueError( "num_inference_steps is required for transformer-only mode. " "Please provide it in CacheDitConfig." ) # Check if the transformer is pre-registered in cache-dit if not BlockAdapterRegister.is_supported(transformer): transformer_cls_name = transformer.__class__.__name__ raise ValueError( f"{transformer_cls_name} is not officially supported by cache-dit. " "Supported cache-dit DiT families include Flux, QwenImage, HunyuanDiT, " "HunyuanVideo, Wan, CogVideoX, Mochi, and others. " "Please ensure your transformer belongs to one of these families or " "define a custom BlockAdapter." ) # Build cache config (including SCM fields if provided) cache_config = DBCacheConfig( num_inference_steps=config.num_inference_steps, Fn_compute_blocks=config.Fn_compute_blocks, Bn_compute_blocks=config.Bn_compute_blocks, max_warmup_steps=config.max_warmup_steps, residual_diff_threshold=config.residual_diff_threshold, max_continuous_cached_steps=config.max_continuous_cached_steps, # SCM fields steps_computation_mask=config.steps_computation_mask, steps_computation_policy=config.steps_computation_policy, ) # Build calibrator config if TaylorSeer is enabled calibrator_config = None if config.enable_taylorseer: calibrator_config = TaylorSeerCalibratorConfig( taylorseer_order=config.taylorseer_order, ) # Enable cache-dit on the transformer logger.info( "Enabling cache-dit on %s with config: Fn=%d, Bn=%d, W=%d, R=%.2f, MC=%d, " "TaylorSeer=%s (order=%d), steps=%d", model_name, config.Fn_compute_blocks, config.Bn_compute_blocks, config.max_warmup_steps, config.residual_diff_threshold, config.max_continuous_cached_steps, config.enable_taylorseer, config.taylorseer_order, config.num_inference_steps, ) # Log SCM configuration if enabled if config.steps_computation_mask: compute_steps = sum(config.steps_computation_mask) cache_steps = len(config.steps_computation_mask) - compute_steps logger.info( "SCM enabled: %d compute steps, %d cache steps, policy=%s", compute_steps, cache_steps, config.steps_computation_policy, ) parallelism_config = _build_parallelism_config(sp_group, tp_group) if parallelism_config is not None: _patch_cache_dit_similarity() _mark_transformer_parallelized(transformer, parallelism_config, sp_group, tp_group) cache_dit.enable_cache( transformer, cache_config=cache_config, calibrator_config=calibrator_config, parallelism_config=None, ) if parallelism_config is not None: context_manager = getattr(transformer, "_context_manager", None) if context_manager is not None: context_manager._sglang_sp_group = sp_group context_manager._sglang_tp_group = tp_group # In mixed TP + SP (Ulysses/Ring) mode, cache-dit decisions must be consistent # across the full TP×SP model-parallel slice. Prefer using SGLang's DIT group # as a conservative superset group; fallback to None. tp_sp_group = None if sp_group is not None and tp_group is not None: tp_sp_group = get_dit_group() context_manager._sglang_tp_sp_group = tp_sp_group return transformer def enable_cache_on_dual_transformer( transformer: torch.nn.Module, transformer_2: torch.nn.Module, primary_config: CacheDitConfig, secondary_config: CacheDitConfig, model_name: str = "wan2.2", sp_group: Optional[torch.distributed.ProcessGroup] = None, tp_group: Optional[torch.distributed.ProcessGroup] = None, ) -> tuple[torch.nn.Module, torch.nn.Module]: """Enable cache-dit on dual transformers using BlockAdapter. For models with two transformers (high-noise expert and low-noise expert), cache-dit requires enabling cache on both simultaneously via BlockAdapter. This cannot be done by calling enable_cache separately on each transformer. Args: primary_config: CacheDitConfig for primary transformer. secondary_config: CacheDitConfig for secondary transformer. sp_group: Sequence parallel process group (for Ulysses/Ring). tp_group: Tensor parallel process group. """ _supported_dual_transformer_models = [ "wan2.2", # Currently, only Wan2.2 will run into dual-transformer case ] if model_name not in _supported_dual_transformer_models: raise ValueError( f"Dual-transformer cache-dit is only supported for " f"{_supported_dual_transformer_models}, got {model_name}." ) if not primary_config.enabled: return transformer, transformer_2 if primary_config.num_inference_steps is None: raise ValueError( "num_inference_steps is required for dual-transformer mode. " "Please provide it in CacheDitConfig." ) # Build DBCacheConfig for primary transformer primary_cache_config = DBCacheConfig( num_inference_steps=primary_config.num_inference_steps, Fn_compute_blocks=primary_config.Fn_compute_blocks, Bn_compute_blocks=primary_config.Bn_compute_blocks, max_warmup_steps=primary_config.max_warmup_steps, residual_diff_threshold=primary_config.residual_diff_threshold, max_continuous_cached_steps=primary_config.max_continuous_cached_steps, steps_computation_mask=primary_config.steps_computation_mask, steps_computation_policy=primary_config.steps_computation_policy, ) # Build DBCacheConfig for secondary transformer secondary_cache_config = DBCacheConfig( num_inference_steps=secondary_config.num_inference_steps, Fn_compute_blocks=secondary_config.Fn_compute_blocks, Bn_compute_blocks=secondary_config.Bn_compute_blocks, max_warmup_steps=secondary_config.max_warmup_steps, residual_diff_threshold=secondary_config.residual_diff_threshold, max_continuous_cached_steps=secondary_config.max_continuous_cached_steps, steps_computation_mask=secondary_config.steps_computation_mask, steps_computation_policy=secondary_config.steps_computation_policy, ) # Build calibrator configs if TaylorSeer is enabled primary_calibrator = None if primary_config.enable_taylorseer: primary_calibrator = TaylorSeerCalibratorConfig( taylorseer_order=primary_config.taylorseer_order, ) secondary_calibrator = None if secondary_config.enable_taylorseer: secondary_calibrator = TaylorSeerCalibratorConfig( taylorseer_order=secondary_config.taylorseer_order, ) # Build ParamsModifier for each transformer primary_modifier = ParamsModifier( cache_config=primary_cache_config, calibrator_config=primary_calibrator, ) secondary_modifier = ParamsModifier( cache_config=secondary_cache_config, calibrator_config=secondary_calibrator, ) # Log configuration logger.info( "Enabling cache-dit on %s dual transformers with BlockAdapter", model_name, ) logger.info( " Primary (transformer): Fn=%d, Bn=%d, W=%d, R=%.2f, MC=%d, TaylorSeer=%s", primary_config.Fn_compute_blocks, primary_config.Bn_compute_blocks, primary_config.max_warmup_steps, primary_config.residual_diff_threshold, primary_config.max_continuous_cached_steps, primary_config.enable_taylorseer, ) logger.info( " Secondary (transformer_2): Fn=%d, Bn=%d, W=%d, R=%.2f, MC=%d, TaylorSeer=%s", secondary_config.Fn_compute_blocks, secondary_config.Bn_compute_blocks, secondary_config.max_warmup_steps, secondary_config.residual_diff_threshold, secondary_config.max_continuous_cached_steps, secondary_config.enable_taylorseer, ) # Log SCM configuration if enabled if primary_config.steps_computation_mask: compute_steps = sum(primary_config.steps_computation_mask) cache_steps = len(primary_config.steps_computation_mask) - compute_steps logger.info( " SCM enabled for primary transformer: %d compute steps, %d cache steps, policy=%s", compute_steps, cache_steps, primary_config.steps_computation_policy, ) if secondary_config.steps_computation_mask: compute_steps = sum(secondary_config.steps_computation_mask) cache_steps = len(secondary_config.steps_computation_mask) - compute_steps logger.info( " SCM enabled for secondary transformer: %d compute steps, %d cache steps, policy=%s", compute_steps, cache_steps, secondary_config.steps_computation_policy, ) parallelism_config = _build_parallelism_config(sp_group, tp_group) if parallelism_config is not None: _patch_cache_dit_similarity() _mark_transformer_parallelized(transformer, parallelism_config, sp_group, tp_group) _mark_transformer_parallelized( transformer_2, parallelism_config, sp_group, tp_group ) # Get blocks attribute - Wan transformers use 'blocks' attribute transformer_blocks = getattr(transformer, "blocks", None) transformer_2_blocks = getattr(transformer_2, "blocks", None) if transformer_blocks is None or transformer_2_blocks is None: raise ValueError( "Dual transformers must have 'blocks' attribute for cache-dit. " f"transformer has blocks: {transformer_blocks is not None}, " f"transformer_2 has blocks: {transformer_2_blocks is not None}" ) # Enable cache-dit using BlockAdapter for both transformers simultaneously # This is required for Wan2.2 and similar dual-transformer architectures if model_name == "wan2.2": # Use Pattern_2 for Wan2.2 dual-transformer. We should check `model_name` # to ensure we only apply this for supported models. Different models # may require different ForwardPattern. cache_dit.enable_cache( BlockAdapter( transformer=[transformer, transformer_2], blocks=[transformer_blocks, transformer_2_blocks], forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2], params_modifiers=[primary_modifier, secondary_modifier], has_separate_cfg=True, ), parallelism_config=None, ) else: raise ValueError( f"Dual-transformer is not implemented for model {model_name} yet." ) if parallelism_config is not None: for t in [transformer, transformer_2]: context_manager = getattr(t, "_context_manager", None) if context_manager is not None: context_manager._sglang_sp_group = sp_group context_manager._sglang_tp_group = tp_group tp_sp_group = None if sp_group is not None and tp_group is not None: try: tp_sp_group = get_dit_group() except Exception: tp_sp_group = None context_manager._sglang_tp_sp_group = tp_sp_group return transformer, transformer_2 def refresh_context_on_transformer( transformer: torch.nn.Module, num_inference_steps: int, scm_preset: str | None = None, verbose: bool = False, ) -> None: """Refresh cache-dit context for transformer.""" cache_dit.refresh_context( transformer, cache_config=DBCacheConfig().reset( num_inference_steps=num_inference_steps, steps_computation_mask=cache_dit.steps_mask( mask_policy=scm_preset, total_steps=num_inference_steps ), steps_computation_policy=scm_preset, ), verbose=verbose, ) logger.debug(f"cache-dit refreshed on transformer (steps={num_inference_steps})") def refresh_context_on_dual_transformer( transformer: torch.nn.Module, transformer_2: torch.nn.Module, num_high_noise_steps: int, num_low_noise_steps: int, scm_preset: str | None = None, verbose: bool = False, ) -> None: """Refresh cache-dit context for dual transformers.""" cache_dit.refresh_context( transformer, cache_config=DBCacheConfig().reset( num_inference_steps=num_high_noise_steps, steps_computation_mask=cache_dit.steps_mask( mask_policy=scm_preset, total_steps=num_high_noise_steps ), steps_computation_policy=scm_preset, ), verbose=verbose, ) cache_dit.refresh_context( transformer_2, cache_config=DBCacheConfig().reset( num_inference_steps=num_low_noise_steps, steps_computation_mask=cache_dit.steps_mask( mask_policy=scm_preset, total_steps=num_low_noise_steps ), steps_computation_policy=scm_preset, ), verbose=verbose, ) logger.debug( f"cache-dit refreshed on dual transformers (steps={num_high_noise_steps}, {num_low_noise_steps})" ) ================================================ FILE: python/sglang/multimodal_gen/runtime/cache/teacache.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ TeaCache: Temporal similarity-based caching for diffusion models. TeaCache accelerates diffusion inference by selectively skipping redundant computation when consecutive diffusion steps are similar enough. This is achieved by tracking the L1 distance between modulated inputs across timesteps. Key concepts: - Modulated input: The input to transformer blocks after timestep conditioning - L1 distance: Measures how different consecutive timesteps are - Threshold: When accumulated L1 distance exceeds threshold, force computation - CFG support: Separate caches for positive and negative branches References: - TeaCache: Accelerating Diffusion Models with Temporal Similarity https://arxiv.org/abs/2411.14324 """ from dataclasses import dataclass from typing import TYPE_CHECKING, Any import numpy as np import torch from sglang.multimodal_gen.configs.models import DiTConfig if TYPE_CHECKING: from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams @dataclass class TeaCacheContext: """Common context extracted for TeaCache skip decision. This context is populated from the forward_batch and forward_context during each denoising step, providing all information needed to make cache decisions. Attributes: current_timestep: Current denoising timestep index (0-indexed). num_inference_steps: Total number of inference steps. do_cfg: Whether classifier-free guidance is enabled. is_cfg_negative: True if currently processing negative CFG branch. teacache_thresh: Threshold for accumulated L1 distance. coefficients: Polynomial coefficients for L1 rescaling. teacache_params: Full TeaCacheParams for model-specific access. """ current_timestep: int num_inference_steps: int do_cfg: bool is_cfg_negative: bool # For CFG branch selection teacache_thresh: float coefficients: list[float] teacache_params: "TeaCacheParams" # Full params for model-specific access class TeaCacheMixin: """ Mixin class providing TeaCache optimization functionality. TeaCache accelerates diffusion inference by selectively skipping redundant computation when consecutive diffusion steps are similar enough. This mixin should be inherited by DiT model classes that want to support TeaCache optimization. It provides: - State management for tracking L1 distances - CFG-aware caching (separate caches for positive/negative branches) - Decision logic for when to compute vs. use cache Example usage in a DiT model: class MyDiT(TeaCacheMixin, BaseDiT): def __init__(self, config, **kwargs): super().__init__(config, **kwargs) self._init_teacache_state() def forward(self, hidden_states, timestep, ...): ctx = self._get_teacache_context() if ctx is not None: # Compute modulated input (model-specific, e.g., after timestep embedding) modulated_input = self._compute_modulated_input(hidden_states, timestep) is_boundary = (ctx.current_timestep == 0 or ctx.current_timestep >= ctx.num_inference_steps - 1) should_calc = self._compute_teacache_decision( modulated_inp=modulated_input, is_boundary_step=is_boundary, coefficients=ctx.coefficients, teacache_thresh=ctx.teacache_thresh, ) if not should_calc: # Use cached residual (must implement retrieve_cached_states) return self.retrieve_cached_states(hidden_states) # Normal forward pass... output = self._transformer_forward(hidden_states, timestep, ...) # Cache states for next step if ctx is not None: self.maybe_cache_states(output, hidden_states) return output Subclass implementation notes: - `_compute_modulated_input()`: Model-specific method to compute the input after timestep conditioning (used for L1 distance calculation) - `retrieve_cached_states()`: Must be overridden to return cached output - `maybe_cache_states()`: Override to store states for cache retrieval Attributes: cnt: Counter for tracking steps. enable_teacache: Whether TeaCache is enabled. previous_modulated_input: Cached modulated input for positive branch. previous_residual: Cached residual for positive branch. accumulated_rel_l1_distance: Accumulated L1 distance for positive branch. is_cfg_negative: Whether currently processing negative CFG branch. _supports_cfg_cache: Whether this model supports CFG cache separation. CFG-specific attributes (only when _supports_cfg_cache is True): previous_modulated_input_negative: Cached input for negative branch. previous_residual_negative: Cached residual for negative branch. accumulated_rel_l1_distance_negative: L1 distance for negative branch. """ # Models that support CFG cache separation (wan/hunyuan/zimage) # Models not in this set (flux/qwen) auto-disable TeaCache when CFG is enabled _CFG_SUPPORTED_PREFIXES: set[str] = {"wan", "hunyuan", "zimage"} config: DiTConfig def _init_teacache_state(self) -> None: """Initialize TeaCache state. Call this in subclass __init__.""" # Common TeaCache state self.cnt = 0 self.enable_teacache = True # Flag indicating if this model supports CFG cache separation self._supports_cfg_cache = ( self.config.prefix.lower() in self._CFG_SUPPORTED_PREFIXES ) # Always initialize positive cache fields (used in all modes) self.previous_modulated_input: torch.Tensor | None = None self.previous_residual: torch.Tensor | None = None self.accumulated_rel_l1_distance: float = 0.0 self.is_cfg_negative = False # CFG-specific fields initialized to None (created when CFG is used) # These are only used when _supports_cfg_cache is True AND do_cfg is True if self._supports_cfg_cache: self.previous_modulated_input_negative: torch.Tensor | None = None self.previous_residual_negative: torch.Tensor | None = None self.accumulated_rel_l1_distance_negative: float = 0.0 def reset_teacache_state(self) -> None: """Reset all TeaCache state at the start of each generation task.""" self.cnt = 0 # Primary cache fields (always present) self.previous_modulated_input = None self.previous_residual = None self.accumulated_rel_l1_distance = 0.0 self.is_cfg_negative = False self.enable_teacache = True # CFG negative cache fields (always reset, may be unused) if self._supports_cfg_cache: self.previous_modulated_input_negative = None self.previous_residual_negative = None self.accumulated_rel_l1_distance_negative = 0.0 def _compute_l1_and_decide( self, modulated_inp: torch.Tensor, coefficients: list[float], teacache_thresh: float, ) -> tuple[float, bool]: """ Compute L1 distance and decide whether to calculate or use cache. Args: modulated_inp: Current timestep's modulated input. coefficients: Polynomial coefficients for L1 rescaling. teacache_thresh: Threshold for cache decision. Returns: Tuple of (new_accumulated_distance, should_calc). """ prev_modulated_inp = ( self.previous_modulated_input_negative if self.is_cfg_negative else self.previous_modulated_input ) # Defensive check: if previous input is not set, force calculation if prev_modulated_inp is None: return 0.0, True # Compute relative L1 distance diff = modulated_inp - prev_modulated_inp rel_l1 = (diff.abs().mean() / prev_modulated_inp.abs().mean()).cpu().item() # Apply polynomial rescaling rescale_func = np.poly1d(coefficients) accumulated_rel_l1_distance = ( self.accumulated_rel_l1_distance_negative if self.is_cfg_negative else self.accumulated_rel_l1_distance ) accumulated_rel_l1_distance = accumulated_rel_l1_distance + rescale_func(rel_l1) if accumulated_rel_l1_distance >= teacache_thresh: # Threshold exceeded: force compute and reset accumulator return 0.0, True # Cache hit: keep accumulated distance return accumulated_rel_l1_distance, False def _compute_teacache_decision( self, modulated_inp: torch.Tensor, is_boundary_step: bool, coefficients: list[float], teacache_thresh: float, ) -> bool: """ Compute cache decision for TeaCache. Args: modulated_inp: Current timestep's modulated input. is_boundary_step: True for boundary timesteps that always compute. coefficients: Polynomial coefficients for L1 rescaling. teacache_thresh: Threshold for cache decision. Returns: True if forward computation is needed, False to use cache. """ if not self.enable_teacache: return True if is_boundary_step: new_accum, should_calc = 0.0, True else: new_accum, should_calc = self._compute_l1_and_decide( modulated_inp=modulated_inp, coefficients=coefficients, teacache_thresh=teacache_thresh, ) # Advance baseline and accumulator for the active branch if not self.is_cfg_negative: self.previous_modulated_input = modulated_inp.clone() self.accumulated_rel_l1_distance = new_accum elif self._supports_cfg_cache: self.previous_modulated_input_negative = modulated_inp.clone() self.accumulated_rel_l1_distance_negative = new_accum return should_calc def _get_teacache_context(self) -> TeaCacheContext | None: """ Check TeaCache preconditions and extract common context. Returns: TeaCacheContext if TeaCache is enabled and properly configured, None if should skip TeaCache logic entirely. """ from sglang.multimodal_gen.runtime.managers.forward_context import ( get_forward_context, ) forward_context = get_forward_context() forward_batch = forward_context.forward_batch # Early return checks if ( forward_batch is None or not forward_batch.enable_teacache or forward_batch.teacache_params is None ): return None teacache_params = forward_batch.teacache_params # Extract common values current_timestep = forward_context.current_timestep num_inference_steps = forward_batch.num_inference_steps do_cfg = forward_batch.do_classifier_free_guidance is_cfg_negative = forward_batch.is_cfg_negative # Reset at first timestep if current_timestep == 0 and not self.is_cfg_negative: self.reset_teacache_state() return TeaCacheContext( current_timestep=current_timestep, num_inference_steps=num_inference_steps, do_cfg=do_cfg, is_cfg_negative=is_cfg_negative, teacache_thresh=teacache_params.teacache_thresh, coefficients=teacache_params.coefficients, teacache_params=teacache_params, ) def maybe_cache_states( self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor ) -> None: """Cache states for later retrieval. Override in subclass if needed.""" pass def should_skip_forward_for_cached_states(self, **kwargs: dict[str, Any]) -> bool: """Check if forward can be skipped using cached states.""" return False def retrieve_cached_states(self, hidden_states: torch.Tensor) -> torch.Tensor: """Retrieve cached states. Must be implemented by subclass.""" raise NotImplementedError("retrieve_cached_states is not implemented") ================================================ FILE: python/sglang/multimodal_gen/runtime/distributed/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from functools import lru_cache from sglang.multimodal_gen.configs.models.encoders import TextEncoderConfig from sglang.multimodal_gen.runtime.distributed.communication_op import * from sglang.multimodal_gen.runtime.distributed.group_coordinator import ( get_local_torch_device, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( cleanup_dist_env_and_memory, get_dp_group, get_dp_rank, get_dp_world_size, get_sp_group, get_sp_parallel_rank, get_sp_world_size, get_tp_group, get_tp_rank, get_tp_world_size, get_world_group, get_world_rank, get_world_size, init_distributed_environment, initialize_model_parallel, maybe_init_distributed_environment_and_model_parallel, model_parallel_is_initialized, ) from sglang.multimodal_gen.runtime.distributed.utils import * # SPDX-License-Identifier: Apache-2.0 __all__ = [ # Initialization "init_distributed_environment", "initialize_model_parallel", "cleanup_dist_env_and_memory", "model_parallel_is_initialized", "maybe_init_distributed_environment_and_model_parallel", # World group "get_world_group", "get_world_rank", "get_world_size", # Data parallel group "get_dp_group", "get_dp_rank", "get_dp_world_size", # Sequence parallel group "get_sp_group", "get_sp_parallel_rank", "get_sp_world_size", # Tensor parallel group "get_tp_group", "get_tp_rank", "get_tp_world_size", # Get torch device "get_local_torch_device", ] def _get_folding_tp_group( config: TextEncoderConfig, ) -> torch.distributed.ProcessGroup | None: if config.parallel_folding: if config.parallel_folding_mode == "sp": return get_sp_group() elif config.parallel_folding_mode == "ulysses": return get_sp_group().ulysses_group elif config.parallel_folding_mode == "ring": return get_sp_group().ring_group return get_tp_group() ================================================ FILE: python/sglang/multimodal_gen/runtime/distributed/communication_op.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/communication_op.py import torch import torch.distributed as dist from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_cfg_group, get_sp_group, get_tp_group, ) def tensor_model_parallel_all_reduce( input_: torch.Tensor, tp_group: dist.ProcessGroup = None ) -> torch.Tensor: """All-reduce the input tensor across model parallel group.""" tp_group = tp_group or get_tp_group() return tp_group.all_reduce(input_) def tensor_model_parallel_all_gather( input_: torch.Tensor, dim: int = -1, tp_group: dist.ProcessGroup = None ) -> torch.Tensor: """All-gather the input tensor across model parallel group.""" tp_group = tp_group or get_tp_group() return tp_group.all_gather(input_, dim) # TODO: remove model, make it sequence_parallel def sequence_model_parallel_all_to_all_4D( input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1 ) -> torch.Tensor: """All-to-all communication of 4D tensors (e.g. QKV matrices) across sequence parallel group.""" return get_sp_group().all_to_all_4D(input_, scatter_dim, gather_dim) def sequence_model_parallel_all_gather( input_: torch.Tensor, dim: int = -1 ) -> torch.Tensor: """All-gather the input tensor across model parallel group.""" return get_sp_group().all_gather(input_, dim) def cfg_model_parallel_all_gather( input_: torch.Tensor, dim: int = -1, separate_tensors: bool = False ) -> torch.Tensor: """All-gather the input tensor across model parallel group.""" return get_cfg_group().all_gather(input_, dim, separate_tensors) def cfg_model_parallel_all_reduce( input_: torch.Tensor, op: torch._C._distributed_c10d.ReduceOp = torch._C._distributed_c10d.ReduceOp.SUM, ) -> torch.Tensor: """All-reduce the input tensor across CFG parallel group.""" return get_cfg_group().all_reduce(input_, op=op) ================================================ FILE: python/sglang/multimodal_gen/runtime/distributed/device_communicators/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo ================================================ FILE: python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/base_device_communicator.py from typing import Any import torch import torch.distributed as dist from torch import Tensor from torch.distributed import ProcessGroup, ReduceOp class DistributedAutograd: """Collection of autograd functions for distributed operations. This class provides custom autograd functions for distributed operations like all_reduce, all_gather, and all_to_all. Each operation is implemented as a static inner class with proper forward and backward implementations. """ class AllReduce(torch.autograd.Function): """Differentiable all_reduce operation. The gradient of all_reduce is another all_reduce operation since the operation combines values from all ranks equally. """ @staticmethod def forward( ctx: Any, group: ProcessGroup, input_: Tensor, op: dist.ReduceOp | None = None, ) -> Tensor: ctx.group = group ctx.op = op output = input_.clone() dist.all_reduce(output, group=group, op=op) return output @staticmethod def backward(ctx: Any, grad_output: Tensor) -> tuple[None, Tensor, None]: grad_output = grad_output.clone() dist.all_reduce(grad_output, group=ctx.group, op=ctx.op) return None, grad_output, None class AllGather(torch.autograd.Function): """Differentiable all_gather operation. The operation gathers tensors from all ranks and concatenates them along a specified dimension. The backward pass uses reduce_scatter to efficiently distribute gradients back to source ranks. """ @staticmethod def forward( ctx: Any, group: ProcessGroup, input_: Tensor, world_size: int, dim: int ) -> Tensor: ctx.group = group ctx.world_size = world_size ctx.dim = dim ctx.input_shape = input_.shape input_size = input_.size() output_size = (input_size[0] * world_size,) + input_size[1:] output_tensor = torch.empty( output_size, dtype=input_.dtype, device=input_.device ) dist.all_gather_into_tensor(output_tensor, input_, group=group) output_tensor = output_tensor.reshape((world_size,) + input_size) output_tensor = output_tensor.movedim(0, dim) output_tensor = output_tensor.reshape( input_size[:dim] + (world_size * input_size[dim],) + input_size[dim + 1 :] ) return output_tensor @staticmethod def backward(ctx: Any, grad_output: Tensor) -> tuple[None, Tensor, None, None]: # Split the gradient tensor along the gathered dimension dim_size = grad_output.size(ctx.dim) // ctx.world_size grad_chunks = grad_output.reshape( grad_output.shape[: ctx.dim] + (ctx.world_size, dim_size) + grad_output.shape[ctx.dim + 1 :] ) grad_chunks = grad_chunks.movedim(ctx.dim, 0) # Each rank only needs its corresponding gradient grad_input = torch.empty( ctx.input_shape, dtype=grad_output.dtype, device=grad_output.device ) dist.reduce_scatter_tensor( grad_input, grad_chunks.contiguous(), group=ctx.group ) return None, grad_input, None, None class AllToAll4D(torch.autograd.Function): """Differentiable all_to_all operation specialized for 4D tensors. This operation is particularly useful for attention operations where we need to redistribute data across ranks for efficient parallel processing. The operation supports two modes: 1. scatter_dim=2, gather_dim=1: Used for redistributing attention heads 2. scatter_dim=1, gather_dim=2: Used for redistributing sequence dimensions """ @staticmethod def forward( ctx: Any, group: ProcessGroup, input_: Tensor, world_size: int, scatter_dim: int, gather_dim: int, ) -> Tensor: ctx.group = group ctx.world_size = world_size ctx.scatter_dim = scatter_dim ctx.gather_dim = gather_dim if world_size == 1: return input_ assert ( input_.dim() == 4 ), f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}" if scatter_dim == 2 and gather_dim == 1: bs, shard_seqlen, hn, hd = input_.shape seqlen = shard_seqlen * world_size shard_hn = hn // world_size input_ = input_.transpose(0, 2).contiguous() # hn, shard_seqlen, bs, hd output = torch.empty_like(input_) dist.all_to_all_single( output, input_, group=group ) # hn, shard_seqlen, bs, hd output = torch.cat( output.split(shard_hn), dim=1 ) # sharded hn, seqlen, bs, hd output = output.transpose( 0, 2 ).contiguous() # bs, seqlen, sharded_hn, hd return output elif scatter_dim == 1 and gather_dim == 2: bs, seqlen, shard_hn, hd = input_.shape hn = shard_hn * world_size shard_seqlen = seqlen // world_size input_ = input_.transpose(0, 2).contiguous() # shard_hn, seqlen, bs, hd input_ = ( input_.reshape(shard_hn, world_size, shard_seqlen, bs, hd) .transpose(0, 1) .reshape(shard_hn * world_size, shard_seqlen, bs, hd) .contiguous() ) output = torch.empty_like(input_) dist.all_to_all_single(output, input_, group=group) output = output.transpose( 0, 2 ).contiguous() # bs, seqlen, sharded_hn, hd return output else: raise RuntimeError( f"Invalid scatter_dim={scatter_dim}, gather_dim={gather_dim}. " f"Only (scatter_dim=2, gather_dim=1) and (scatter_dim=1, gather_dim=2) are supported." ) @staticmethod def backward( ctx: Any, grad_output: Tensor ) -> tuple[None, Tensor, None, None, None]: if ctx.world_size == 1: return None, grad_output, None, None, None # For backward pass, we swap scatter_dim and gather_dim output = DistributedAutograd.AllToAll4D.apply( ctx.group, grad_output, ctx.world_size, ctx.gather_dim, ctx.scatter_dim ) return None, output, None, None, None class DeviceCommunicatorBase: """ Base class for device-specific communicator with autograd support. It can use the `cpu_group` to initialize the communicator. If the device has PyTorch integration (PyTorch can recognize its communication backend), the `device_group` will also be given. """ def __init__( self, cpu_group: ProcessGroup, device: torch.device | None = None, device_group: ProcessGroup | None = None, unique_name: str = "", ): self.device = device or torch.device("cpu") self.cpu_group = cpu_group self.device_group = device_group self.unique_name = unique_name self.rank = dist.get_rank(cpu_group) self.world_size = dist.get_world_size(cpu_group) self.ranks = dist.get_process_group_ranks(cpu_group) self.global_rank = dist.get_rank() self.global_world_size = dist.get_world_size() self.rank_in_group = dist.get_group_rank(self.cpu_group, self.global_rank) def all_reduce( self, input_: torch.Tensor, op: dist.ReduceOp | None = ReduceOp.SUM ) -> torch.Tensor: """Performs an all_reduce operation with gradient support.""" return DistributedAutograd.AllReduce.apply(self.device_group, input_, op) def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: """Performs an all_gather operation with gradient support.""" if dim < 0: dim += input_.dim() return DistributedAutograd.AllGather.apply( self.device_group, input_, self.world_size, dim ) def all_to_all_4D( self, input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1 ) -> torch.Tensor: """Performs a 4D all-to-all operation with gradient support.""" return DistributedAutograd.AllToAll4D.apply( self.device_group, input_, self.world_size, scatter_dim, gather_dim ) def gather( self, input_: torch.Tensor, dst: int = 0, dim: int = -1 ) -> torch.Tensor | None: """ NOTE: We assume that the input tensor is on the same device across all the ranks. NOTE: `dst` is the local rank of the destination rank. """ world_size = self.world_size assert ( -input_.dim() <= dim < input_.dim() ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" if dim < 0: # Convert negative dim to positive. dim += input_.dim() # Allocate output tensor. if self.rank_in_group == dst: gather_list = [torch.empty_like(input_) for _ in range(world_size)] else: gather_list = None # Gather. torch.distributed.gather( input_, gather_list, dst=self.ranks[dst], group=self.device_group ) if self.rank_in_group == dst: output_tensor = torch.cat(gather_list, dim=dim) else: output_tensor = None return output_tensor def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: """Sends a tensor to the destination rank in a non-blocking way""" """NOTE: `dst` is the local rank of the destination rank.""" if dst is None: dst = (self.rank_in_group + 1) % self.world_size torch.distributed.send(tensor, self.ranks[dst], self.device_group) def recv( self, size: torch.Size, dtype: torch.dtype, src: int | None = None ) -> torch.Tensor: """Receives a tensor from the source rank.""" """NOTE: `src` is the local rank of the source rank.""" if src is None: src = (self.rank_in_group - 1) % self.world_size tensor = torch.empty(size, dtype=dtype, device=self.device) torch.distributed.recv(tensor, self.ranks[src], self.device_group) return tensor def destroy(self) -> None: pass ================================================ FILE: python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from: https://github.com/vllm-project/vllm/blob/main/vllm/distributed/device_communicators/cpu_communicator.py import os import torch from torch.distributed import ProcessGroup from .base_device_communicator import DeviceCommunicatorBase class CpuCommunicator(DeviceCommunicatorBase): def __init__( self, cpu_group: ProcessGroup, device: torch.device | None = None, device_group: ProcessGroup | None = None, unique_name: str = "", ): from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.platforms.interface import CpuArchEnum super().__init__(cpu_group, device, device_group, unique_name) self.dist_module = torch.distributed if ( (current_platform.get_cpu_architecture() == CpuArchEnum.X86) and hasattr(torch.ops._C, "init_shm_manager") and unique_name.startswith("tp") ): self.dist_module = _CPUSHMDistributed(self) def all_reduce( self, input_: torch.Tensor, op: torch.distributed.ReduceOp | None = torch.distributed.ReduceOp.SUM, ) -> torch.Tensor: self.dist_module.all_reduce(input_, group=self.device_group, op=op) return input_ def gather( self, input_: torch.Tensor, dst: int = 0, dim: int = -1 ) -> torch.Tensor | None: """ NOTE: We assume that the input tensor is on the same device across all the ranks. NOTE: `dst` is the local rank of the destination rank. """ world_size = self.world_size assert ( -input_.dim() <= dim < input_.dim() ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" if dim < 0: # Convert negative dim to positive. dim += input_.dim() # Allocate output tensor. if self.rank_in_group == dst: gather_list = [torch.empty_like(input_) for _ in range(world_size)] else: gather_list = None # Gather. self.dist_module.gather( input_, gather_list, dst=self.ranks[dst], group=self.device_group ) if self.rank_in_group == dst: output_tensor = torch.cat(gather_list, dim=dim) else: output_tensor = None return output_tensor def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: if dim < 0: # Convert negative dim to positive. dim += input_.dim() input_size = input_.size() # NOTE: we have to use concat-style all-gather here, # stack-style all-gather has compatibility issues with # torch.compile . see https://github.com/pytorch/pytorch/issues/138795 output_size = (input_size[0] * self.world_size,) + input_size[1:] # Allocate output tensor. output_tensor = torch.empty( output_size, dtype=input_.dtype, device=input_.device ) # All-gather. self.dist_module.all_gather_into_tensor( output_tensor, input_, group=self.device_group ) # Reshape output_tensor = output_tensor.reshape((self.world_size,) + input_size) output_tensor = output_tensor.movedim(0, dim) output_tensor = output_tensor.reshape( input_size[:dim] + (self.world_size * input_size[dim],) + input_size[dim + 1 :] ) return output_tensor class _CPUSHMDistributed: def __init__(self, communicator: CpuCommunicator): instance_identifier = os.environ["VLLM_DIST_IDENT"] unique_name = communicator.unique_name instance_identifier = f"{instance_identifier}-{unique_name}" self.communicator = communicator group_ranks = [str(rank) for rank in self.communicator.ranks] shm_group_identifier = f"[{'-'.join(group_ranks)}]" self.group_name = f"{instance_identifier}-{shm_group_identifier}-cpushm" self.handle = self._init_cpu_shm() def _init_cpu_shm(self) -> int: handle = torch.ops._C.init_shm_manager( self.group_name, self.communicator.world_size, self.communicator.rank, ) torch.distributed.barrier(self.communicator.device_group) torch.ops._C.join_shm_manager( handle, self.group_name, ) torch.distributed.barrier(self.communicator.device_group) return int(handle) def all_reduce( self, input: torch.Tensor, group: ProcessGroup | None = None ) -> None: torch.ops._C.shm_allreduce(self.handle, input) def gather( self, input: torch.Tensor, gather_list: list[torch.Tensor] | None, dst: int = -1, group: ProcessGroup | None = None, ) -> None: # Note: different from the torch gather, here we use local dst rank. torch.ops._C.shm_gather( self.handle, input, gather_list, torch.distributed.get_group_rank(group, dst), ) def all_gather_into_tensor( self, output: torch.Tensor, input: torch.Tensor, group: ProcessGroup | None = None, ) -> None: torch.ops._C.shm_all_gather(self.handle, input, output) ================================================ FILE: python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/cuda_communicator.py import torch from torch.distributed import ProcessGroup from sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_communicator import ( DeviceCommunicatorBase, ) class CudaCommunicator(DeviceCommunicatorBase): def __init__( self, cpu_group: ProcessGroup, device: torch.device | None = None, device_group: ProcessGroup | None = None, unique_name: str = "", ): super().__init__(cpu_group, device, device_group, unique_name) from sglang.multimodal_gen.runtime.distributed.device_communicators.pynccl import ( PyNcclCommunicator, ) self.pynccl_comm: PyNcclCommunicator | None = None if self.world_size > 1: self.pynccl_comm = PyNcclCommunicator( group=self.cpu_group, device=self.device, ) def all_reduce(self, input_, op: torch.distributed.ReduceOp | None = None): pynccl_comm = self.pynccl_comm assert pynccl_comm is not None out = pynccl_comm.all_reduce(input_, op=op) if out is None: # fall back to the default all-reduce using PyTorch. # this usually happens during testing. # when we run the model, allreduce only happens for the TP # group, where we always have either custom allreduce or pynccl. out = input_.clone() torch.distributed.all_reduce(out, group=self.device_group, op=op) return out def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: """Sends a tensor to the destination rank in a non-blocking way""" """NOTE: `dst` is the local rank of the destination rank.""" if dst is None: dst = (self.rank_in_group + 1) % self.world_size pynccl_comm = self.pynccl_comm if pynccl_comm is not None and not pynccl_comm.disabled: pynccl_comm.send(tensor, dst) else: torch.distributed.send(tensor, self.ranks[dst], self.device_group) def recv( self, size: torch.Size, dtype: torch.dtype, src: int | None = None ) -> torch.Tensor: """Receives a tensor from the source rank.""" """NOTE: `src` is the local rank of the source rank.""" if src is None: src = (self.rank_in_group - 1) % self.world_size tensor = torch.empty(size, dtype=dtype, device=self.device) pynccl_comm = self.pynccl_comm if pynccl_comm is not None and not pynccl_comm.disabled: pynccl_comm.recv(tensor, src) else: torch.distributed.recv(tensor, self.ranks[src], self.device_group) return tensor def destroy(self) -> None: if self.pynccl_comm is not None: self.pynccl_comm = None ================================================ FILE: python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl.py # ===================== import region ===================== import torch import torch.distributed as dist from torch.distributed import ProcessGroup, ReduceOp from sglang.multimodal_gen.runtime.distributed.device_communicators.pynccl_wrapper import ( NCCLLibrary, buffer_type, cudaStream_t, ncclComm_t, ncclDataTypeEnum, ncclRedOpTypeEnum, ncclUniqueId, ) from sglang.multimodal_gen.runtime.distributed.utils import StatelessProcessGroup from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import current_stream logger = init_logger(__name__) class PyNcclCommunicator: def __init__( self, group: ProcessGroup | StatelessProcessGroup, device: int | str | torch.device, library_path: str | None = None, ): """ Args: group: the process group to work on. If None, it will use the default process group. device: the device to bind the PyNcclCommunicator to. If None, it will be bind to f"cuda:{local_rank}". library_path: the path to the NCCL library. If None, it will use the default library path. It is the caller's responsibility to make sure each communicator is bind to a unique device. """ if not isinstance(group, StatelessProcessGroup): assert dist.is_initialized() assert ( dist.get_backend(group) != dist.Backend.NCCL ), "PyNcclCommunicator should be attached to a non-NCCL group." # note: this rank is the rank in the group self.rank = dist.get_rank(group) self.world_size = dist.get_world_size(group) else: self.rank = group.rank self.world_size = group.world_size self.group = group # if world_size == 1, no need to create communicator if self.world_size == 1: self.available = False self.disabled = True return try: self.nccl = NCCLLibrary(library_path) except Exception: # disable because of missing NCCL library # e.g. in a non-GPU environment self.available = False self.disabled = True return self.available = True self.disabled = False logger.info("sglang-diffusion is using nccl==%s", self.nccl.ncclGetVersion()) if self.rank == 0: # get the unique id from NCCL self.unique_id = self.nccl.ncclGetUniqueId() else: # construct an empty unique id self.unique_id = ncclUniqueId() if not isinstance(group, StatelessProcessGroup): tensor = torch.ByteTensor(list(self.unique_id.internal)) ranks = dist.get_process_group_ranks(group) # arg `src` in `broadcast` is the global rank dist.broadcast(tensor, src=ranks[0], group=group) byte_list = tensor.tolist() for i, byte in enumerate(byte_list): self.unique_id.internal[i] = byte else: self.unique_id = group.broadcast_obj(self.unique_id, src=0) if isinstance(device, int): device = torch.device(f"cuda:{device}") elif isinstance(device, str): device = torch.device(device) # now `device` is a `torch.device` object assert isinstance(device, torch.device) self.device = device # nccl communicator and stream will use this device # `torch.cuda.device` is a context manager that changes the # current cuda device to the specified one with torch.cuda.device(device): self.comm: ncclComm_t = self.nccl.ncclCommInitRank( self.world_size, self.unique_id, self.rank ) stream = current_stream() # A small all_reduce for warmup. data = torch.zeros(1, device=device) self.all_reduce(data) if stream is not None: stream.synchronize() del data def all_reduce( self, in_tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM, stream=None ) -> torch.Tensor: if self.disabled: return None # nccl communicator created on a specific device # will only work on tensors on the same device # otherwise it will cause "illegal memory access" assert in_tensor.device == self.device, ( f"this nccl communicator is created to work on {self.device}, " f"but the input tensor is on {in_tensor.device}" ) out_tensor = torch.empty_like(in_tensor) if stream is None: stream = current_stream() self.nccl.ncclAllReduce( buffer_type(in_tensor.data_ptr()), buffer_type(out_tensor.data_ptr()), in_tensor.numel(), ncclDataTypeEnum.from_torch(in_tensor.dtype), ncclRedOpTypeEnum.from_torch(op), self.comm, cudaStream_t(stream.cuda_stream), ) return out_tensor def all_gather( self, output_tensor: torch.Tensor, input_tensor: torch.Tensor, stream=None ): if self.disabled: return # nccl communicator created on a specific device # will only work on tensors on the same device # otherwise it will cause "illegal memory access" assert input_tensor.device == self.device, ( f"this nccl communicator is created to work on {self.device}, " f"but the input tensor is on {input_tensor.device}" ) if stream is None: stream = current_stream() self.nccl.ncclAllGather( buffer_type(input_tensor.data_ptr()), buffer_type(output_tensor.data_ptr()), input_tensor.numel(), ncclDataTypeEnum.from_torch(input_tensor.dtype), self.comm, cudaStream_t(stream.cuda_stream), ) def reduce_scatter( self, output_tensor: torch.Tensor, input_tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM, stream=None, ): if self.disabled: return # nccl communicator created on a specific device # will only work on tensors on the same device # otherwise it will cause "illegal memory access" assert input_tensor.device == self.device, ( f"this nccl communicator is created to work on {self.device}, " f"but the input tensor is on {input_tensor.device}" ) if stream is None: stream = current_stream() self.nccl.ncclReduceScatter( buffer_type(input_tensor.data_ptr()), buffer_type(output_tensor.data_ptr()), output_tensor.numel(), ncclDataTypeEnum.from_torch(input_tensor.dtype), ncclRedOpTypeEnum.from_torch(op), self.comm, cudaStream_t(stream.cuda_stream), ) def send(self, tensor: torch.Tensor, dst: int, stream=None): if self.disabled: return assert tensor.device == self.device, ( f"this nccl communicator is created to work on {self.device}, " f"but the input tensor is on {tensor.device}" ) if stream is None: stream = current_stream() self.nccl.ncclSend( buffer_type(tensor.data_ptr()), tensor.numel(), ncclDataTypeEnum.from_torch(tensor.dtype), dst, self.comm, cudaStream_t(stream.cuda_stream), ) def recv(self, tensor: torch.Tensor, src: int, stream=None): if self.disabled: return assert tensor.device == self.device, ( f"this nccl communicator is created to work on {self.device}, " f"but the input tensor is on {tensor.device}" ) if stream is None: stream = current_stream() self.nccl.ncclRecv( buffer_type(tensor.data_ptr()), tensor.numel(), ncclDataTypeEnum.from_torch(tensor.dtype), src, self.comm, cudaStream_t(stream.cuda_stream), ) def broadcast(self, tensor: torch.Tensor, src: int, stream=None): if self.disabled: return assert tensor.device == self.device, ( f"this nccl communicator is created to work on {self.device}, " f"but the input tensor is on {tensor.device}" ) if stream is None: stream = current_stream() if src == self.rank: sendbuff = buffer_type(tensor.data_ptr()) # NCCL requires the sender also to have a receive buffer recvbuff = buffer_type(tensor.data_ptr()) else: sendbuff = buffer_type() recvbuff = buffer_type(tensor.data_ptr()) self.nccl.ncclBroadcast( sendbuff, recvbuff, tensor.numel(), ncclDataTypeEnum.from_torch(tensor.dtype), src, self.comm, cudaStream_t(stream.cuda_stream), ) ================================================ FILE: python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl_wrapper.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl_wrapper.py # This file is a pure Python wrapper for the NCCL library. # The main purpose is to use NCCL combined with CUDA graph. # Before writing this script, we tried the following approach: # 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself # often gets stuck when initializing the NCCL communicator. # 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce` # contains many other potential cuda APIs, that are not allowed during # capturing the CUDA graph. For further details, please check # https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ . # # Another rejected idea is to write a C/C++ binding for NCCL. It is usually # doable, but we often encounter issues related with nccl versions, and need # to switch between different versions of NCCL. See # https://github.com/NVIDIA/nccl/issues/1234 for more details. # A C/C++ binding is not flexible enough to handle this. It requires # recompilation of the code every time we want to switch between different # versions. This current implementation, with a **pure** Python wrapper, is # more flexible. We can easily switch between different versions of NCCL by # changing the environment variable `SGLANG_DIFFUSION_NCCL_SO_PATH`, or the `so_file` # variable in the code. # TODO(will): support SGLANG_DIFFUSION_NCCL_SO_PATH import ctypes import platform from dataclasses import dataclass from typing import Any import torch from torch.distributed import ReduceOp from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import find_nccl_library logger = init_logger(__name__) # === export types and functions from nccl to Python === # for the original nccl definition, please check # https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in ncclResult_t = ctypes.c_int ncclComm_t = ctypes.c_void_p class ncclUniqueId(ctypes.Structure): _fields_ = [("internal", ctypes.c_byte * 128)] cudaStream_t = ctypes.c_void_p buffer_type = ctypes.c_void_p ncclDataType_t = ctypes.c_int class ncclDataTypeEnum: ncclInt8 = 0 ncclChar = 0 ncclUint8 = 1 ncclInt32 = 2 ncclInt = 2 ncclUint32 = 3 ncclInt64 = 4 ncclUint64 = 5 ncclFloat16 = 6 ncclHalf = 6 ncclFloat32 = 7 ncclFloat = 7 ncclFloat64 = 8 ncclDouble = 8 ncclBfloat16 = 9 ncclNumTypes = 10 @classmethod def from_torch(cls, dtype: torch.dtype) -> int: if dtype == torch.int8: return cls.ncclInt8 if dtype == torch.uint8: return cls.ncclUint8 if dtype == torch.int32: return cls.ncclInt32 if dtype == torch.int64: return cls.ncclInt64 if dtype == torch.float16: return cls.ncclFloat16 if dtype == torch.float32: return cls.ncclFloat32 if dtype == torch.float64: return cls.ncclFloat64 if dtype == torch.bfloat16: return cls.ncclBfloat16 raise ValueError(f"Unsupported dtype: {dtype}") ncclRedOp_t = ctypes.c_int class ncclRedOpTypeEnum: ncclSum = 0 ncclProd = 1 ncclMax = 2 ncclMin = 3 ncclAvg = 4 ncclNumOps = 5 @classmethod def from_torch(cls, op: ReduceOp) -> int: if op == ReduceOp.SUM: return cls.ncclSum if op == ReduceOp.PRODUCT: return cls.ncclProd if op == ReduceOp.MAX: return cls.ncclMax if op == ReduceOp.MIN: return cls.ncclMin if op == ReduceOp.AVG: return cls.ncclAvg raise ValueError(f"Unsupported op: {op}") @dataclass class Function: name: str restype: Any argtypes: list[Any] class NCCLLibrary: exported_functions = [ # const char* ncclGetErrorString(ncclResult_t result) Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]), # ncclResult_t ncclGetVersion(int *version); Function("ncclGetVersion", ncclResult_t, [ctypes.POINTER(ctypes.c_int)]), # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId); Function("ncclGetUniqueId", ncclResult_t, [ctypes.POINTER(ncclUniqueId)]), # ncclResult_t ncclCommInitRank( # ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank); # note that ncclComm_t is a pointer type, so the first argument # is a pointer to a pointer Function( "ncclCommInitRank", ncclResult_t, [ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, ctypes.c_int], ), # ncclResult_t ncclAllReduce( # const void* sendbuff, void* recvbuff, size_t count, # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, # cudaStream_t stream); # note that cudaStream_t is a pointer type, so the last argument # is a pointer Function( "ncclAllReduce", ncclResult_t, [ buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t, ncclRedOp_t, ncclComm_t, cudaStream_t, ], ), # ncclResult_t ncclAllGather( # const void* sendbuff, void* recvbuff, size_t count, # ncclDataType_t datatype, ncclComm_t comm, # cudaStream_t stream); # note that cudaStream_t is a pointer type, so the last argument # is a pointer Function( "ncclAllGather", ncclResult_t, [ buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t, ncclComm_t, cudaStream_t, ], ), # ncclResult_t ncclReduceScatter( # const void* sendbuff, void* recvbuff, size_t count, # ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, # cudaStream_t stream); # note that cudaStream_t is a pointer type, so the last argument # is a pointer Function( "ncclReduceScatter", ncclResult_t, [ buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t, ncclRedOp_t, ncclComm_t, cudaStream_t, ], ), # ncclResult_t ncclSend( # const void* sendbuff, size_t count, ncclDataType_t datatype, # int dest, ncclComm_t comm, cudaStream_t stream); Function( "ncclSend", ncclResult_t, [ buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int, ncclComm_t, cudaStream_t, ], ), # ncclResult_t ncclRecv( # void* recvbuff, size_t count, ncclDataType_t datatype, # int src, ncclComm_t comm, cudaStream_t stream); Function( "ncclRecv", ncclResult_t, [ buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int, ncclComm_t, cudaStream_t, ], ), # ncclResult_t ncclBroadcast( # const void* sendbuff, void* recvbuff, size_t count, # ncclDataType_t datatype, int root, ncclComm_t comm, # cudaStream_t stream); Function( "ncclBroadcast", ncclResult_t, [ buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int, ncclComm_t, cudaStream_t, ], ), # be cautious! this is a collective call, it will block until all # processes in the communicator have called this function. # because Python object destruction can happen in random order, # it is better not to call it at all. # ncclResult_t ncclCommDestroy(ncclComm_t comm); Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]), ] # class attribute to store the mapping from the path to the library # to avoid loading the same library multiple times path_to_library_cache: dict[str, Any] = {} # class attribute to store the mapping from library path # to the corresponding dictionary path_to_dict_mapping: dict[str, dict[str, Any]] = {} def __init__(self, so_file: str | None = None): so_file = so_file or find_nccl_library() try: if so_file not in NCCLLibrary.path_to_dict_mapping: lib = ctypes.CDLL(so_file) NCCLLibrary.path_to_library_cache[so_file] = lib self.lib = NCCLLibrary.path_to_library_cache[so_file] except Exception as e: logger.error( "Failed to load NCCL library from %s ." "It is expected if you are not running on NVIDIA/AMD/MTHREADS GPUs." "Otherwise, the nccl library might not exist, be corrupted " "or it does not support the current platform %s." "If you already have the library, please set the " "environment variable SGLANG_DIFFUSION_NCCL_SO_PATH" " to point to the correct nccl library path.", so_file, platform.platform(), ) raise e if so_file not in NCCLLibrary.path_to_dict_mapping: _funcs: dict[str, Any] = {} for func in NCCLLibrary.exported_functions: f = getattr(self.lib, func.name) f.restype = func.restype f.argtypes = func.argtypes _funcs[func.name] = f NCCLLibrary.path_to_dict_mapping[so_file] = _funcs self._funcs = NCCLLibrary.path_to_dict_mapping[so_file] def ncclGetErrorString(self, result: ncclResult_t) -> str: return str(self._funcs["ncclGetErrorString"](result).decode("utf-8")) def NCCL_CHECK(self, result: ncclResult_t) -> None: if result != 0: error_str = self.ncclGetErrorString(result) raise RuntimeError(f"NCCL error: {error_str}") def ncclGetVersion(self) -> str: version = ctypes.c_int() self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version))) version_str = str(version.value) # something like 21903 --> "2.19.3" major = version_str[0].lstrip("0") minor = version_str[1:3].lstrip("0") patch = version_str[3:].lstrip("0") return f"{major}.{minor}.{patch}" def ncclGetUniqueId(self) -> ncclUniqueId: unique_id = ncclUniqueId() self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id))) return unique_id def ncclCommInitRank( self, world_size: int, unique_id: ncclUniqueId, rank: int ) -> ncclComm_t: comm = ncclComm_t() self.NCCL_CHECK( self._funcs["ncclCommInitRank"]( ctypes.byref(comm), world_size, unique_id, rank ) ) return comm def ncclAllReduce( self, sendbuff: buffer_type, recvbuff: buffer_type, count: int, datatype: int, op: int, comm: ncclComm_t, stream: cudaStream_t, ) -> None: # `datatype` actually should be `ncclDataType_t` # and `op` should be `ncclRedOp_t` # both are aliases of `ctypes.c_int` # when we pass int to a function, it will be converted to `ctypes.c_int` # by ctypes automatically self.NCCL_CHECK( self._funcs["ncclAllReduce"]( sendbuff, recvbuff, count, datatype, op, comm, stream ) ) def ncclReduceScatter( self, sendbuff: buffer_type, recvbuff: buffer_type, count: int, datatype: int, op: int, comm: ncclComm_t, stream: cudaStream_t, ) -> None: # `datatype` actually should be `ncclDataType_t` # and `op` should be `ncclRedOp_t` # both are aliases of `ctypes.c_int` # when we pass int to a function, it will be converted to `ctypes.c_int` # by ctypes automatically self.NCCL_CHECK( self._funcs["ncclReduceScatter"]( sendbuff, recvbuff, count, datatype, op, comm, stream ) ) def ncclAllGather( self, sendbuff: buffer_type, recvbuff: buffer_type, count: int, datatype: int, comm: ncclComm_t, stream: cudaStream_t, ) -> None: # `datatype` actually should be `ncclDataType_t` # which is an aliases of `ctypes.c_int` # when we pass int to a function, it will be converted to `ctypes.c_int` # by ctypes automatically self.NCCL_CHECK( self._funcs["ncclAllGather"]( sendbuff, recvbuff, count, datatype, comm, stream ) ) def ncclSend( self, sendbuff: buffer_type, count: int, datatype: int, dest: int, comm: ncclComm_t, stream: cudaStream_t, ) -> None: self.NCCL_CHECK( self._funcs["ncclSend"](sendbuff, count, datatype, dest, comm, stream) ) def ncclRecv( self, recvbuff: buffer_type, count: int, datatype: int, src: int, comm: ncclComm_t, stream: cudaStream_t, ) -> None: self.NCCL_CHECK( self._funcs["ncclRecv"](recvbuff, count, datatype, src, comm, stream) ) def ncclBroadcast( self, sendbuff: buffer_type, recvbuff: buffer_type, count: int, datatype: int, root: int, comm: ncclComm_t, stream: cudaStream_t, ) -> None: self.NCCL_CHECK( self._funcs["ncclBroadcast"]( sendbuff, recvbuff, count, datatype, root, comm, stream ) ) def ncclCommDestroy(self, comm: ncclComm_t) -> None: self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm)) __all__ = [ "NCCLLibrary", "ncclDataTypeEnum", "ncclRedOpTypeEnum", "ncclUniqueId", "ncclComm_t", "cudaStream_t", "buffer_type", ] ================================================ FILE: python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # Copyright 2024 xDiT team. # Adapted from # https://github.com/vllm-project/vllm/blob/main/vllm/distributed/parallel_state.py # Copyright 2023 The vLLM team. # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. import pickle from collections import namedtuple from contextlib import contextmanager from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.distributed from torch.cuda import synchronize from torch.distributed import Backend, ProcessGroup from sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_communicator import ( DeviceCommunicatorBase, ) from sglang.multimodal_gen.runtime.distributed.device_communicators.cpu_communicator import ( CpuCommunicator, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import ( init_logger, suppress_stdout, ) try: import torch_musa # noqa: F401 from torch_musa.core.device import synchronize except ModuleNotFoundError: pass logger = init_logger(__name__) TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) _group_name_counter: dict[str, int] = {} def get_local_torch_device() -> torch.device: """Return the torch device for the current rank.""" return current_platform.get_local_torch_device() def _get_unique_name(name: str) -> str: """Get a unique name for the group. Example: _get_unique_name("tp") -> "tp:0" _get_unique_name("tp") -> "tp:1" """ if name not in _group_name_counter: _group_name_counter[name] = 0 newname = f"{name}:{_group_name_counter[name]}" _group_name_counter[name] += 1 return newname def _split_tensor_dict( tensor_dict: Dict[str, Union[torch.Tensor, Any]], prefix: str = "" ) -> Tuple[List[Tuple[str, Any]], List[torch.Tensor]]: """Split the tensor dictionary into two parts: 1. A list of (key, value) pairs. If the value is a tensor, it is replaced by its metadata. 2. A list of tensors. If the Tensor is nested under `tensor_dict["key1"]["key2"]`, the key of its metadata will be "key1%key2". """ metadata_list: List[Tuple[str, Any]] = [] tensor_list = [] for key, value in tensor_dict.items(): assert "%" not in key, ( "Avoid having '%' in key " "as it is used as a separator for nested entries." ) if isinstance(value, torch.Tensor): # Note: we cannot use `value.device` here, # because it contains not only the device type but also the device # index (e.g. "cuda:0"). We only need the device type. # receiving side will set the device index. device = value.device.type metadata_list.append( ( prefix + key, TensorMetadata(device, value.dtype, value.size()), ) ) tensor_list.append(value) elif isinstance(value, dict): if len(value) == 0: metadata_list.append((prefix + key, value)) inner_metadata_list, inner_tensor_list = _split_tensor_dict( value, prefix + key + "%" ) metadata_list.extend(inner_metadata_list) tensor_list.extend(inner_tensor_list) else: metadata_list.append((prefix + key, value)) return metadata_list, tensor_list def _update_nested_dict(nested_dict, flattened_key, value): key_splits = flattened_key.split("%") cur_dict = nested_dict for k in key_splits[:-1]: if k not in cur_dict: cur_dict[k] = {} cur_dict = cur_dict[k] cur_dict[key_splits[-1]] = value @dataclass class GraphCaptureContext: stream: torch.cuda.Stream | None class GroupCoordinator: """ PyTorch ProcessGroup wrapper for a group of processes. PyTorch ProcessGroup is bound to one specific communication backend, e.g. NCCL, Gloo, MPI, etc. GroupCoordinator takes charge of all the communication operations among the processes in the group. It can route the communication to a specific implementation (e.g. switch allreduce implementation based on the tensor size and cuda graph mode). """ # available attributes: rank: int # global rank ranks: List[int] # global ranks in the group world_size: int # size of the group # difference between `local_rank` and `rank_in_group`: # if we have a group of size 4 across two nodes: # Process | Node | Rank | Local Rank | Rank in Group # 0 | 0 | 0 | 0 | 0 # 1 | 0 | 1 | 1 | 1 # 2 | 1 | 2 | 0 | 2 # 3 | 1 | 3 | 1 | 3 local_rank: int # local rank in the current node, used to assign devices rank_in_group: int # rank inside the group cpu_group: ProcessGroup # group for CPU communication device_group: ProcessGroup # group for device communication use_device_communicator: bool # whether to use device communicator device_communicator: DeviceCommunicatorBase # device communicator def __init__( self, group_ranks: List[List[int]], local_rank: int, torch_distributed_backend: Union[str, Backend], use_device_communicator: bool = True, use_message_queue_broadcaster: bool = False, group_name: str | None = None, ): self.unique_name = _get_unique_name(group_name) self.rank = torch.distributed.get_rank() self.local_rank = local_rank self.device_group = None self.cpu_group = None for ranks in group_ranks: device_group = torch.distributed.new_group( ranks, backend=torch_distributed_backend ) # a group with `gloo` backend, to allow direct coordination between # processes through the CPU. with suppress_stdout(): cpu_group = torch.distributed.new_group(ranks, backend="gloo") if self.rank in ranks: self.ranks = ranks self.world_size = len(ranks) self.rank_in_group = ranks.index(self.rank) self.device_group = device_group self.cpu_group = cpu_group assert self.cpu_group is not None, f"{group_ranks=}, {local_rank=}" assert self.device_group is not None # TODO: fix it for other platforms self.device = get_local_torch_device() self.use_device_communicator = use_device_communicator self.device_communicator: DeviceCommunicatorBase = None # type: ignore if use_device_communicator and self.world_size > 1: # Platform-aware device communicator selection if current_platform.is_cuda_alike(): from sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator import ( CudaCommunicator, ) self.device_communicator = CudaCommunicator( cpu_group=self.cpu_group, device=self.device, device_group=self.device_group, unique_name=self.unique_name, ) else: # For MPS and CPU, use the CPU communicator self.device_communicator = CpuCommunicator( cpu_group=self.cpu_group, device=self.device, device_group=self.device_group, unique_name=self.unique_name, ) self.mq_broadcaster = None # TODO(will): check if this is needed # self.use_custom_op_call = current_platform.is_cuda_alike() self.use_custom_op_call = False @property def first_rank(self): """Return the global rank of the first process in the group""" return self.ranks[0] @property def last_rank(self): """Return the global rank of the last process in the group""" return self.ranks[-1] @property def is_first_rank(self): """Return whether the caller is the first process in the group""" return self.rank == self.first_rank @property def is_last_rank(self): """Return whether the caller is the last process in the group""" return self.rank == self.last_rank @property def next_rank(self): """Return the global rank of the process that follows the caller""" rank_in_group = self.rank_in_group world_size = self.world_size return self.ranks[(rank_in_group + 1) % world_size] @property def prev_rank(self): """Return the global rank of the process that precedes the caller""" rank_in_group = self.rank_in_group world_size = self.world_size return self.ranks[(rank_in_group - 1) % world_size] @property def group_next_rank(self): """Return the group rank of the process that follows the caller""" rank_in_group = self.rank_in_group world_size = self.world_size return (rank_in_group + 1) % world_size @property def group_prev_rank(self): """Return the group rank of the process that precedes the caller""" rank_in_group = self.rank_in_group world_size = self.world_size return (rank_in_group - 1) % world_size @property def skip_rank(self): """Return the global rank of the process that skip connects with the caller""" rank_in_group = self.rank_in_group world_size = self.world_size return self.ranks[(world_size - rank_in_group - 1) % world_size] @property def group_skip_rank(self): """Return the group rank of the process that skip connects with the caller""" rank_in_group = self.rank_in_group world_size = self.world_size return (world_size - rank_in_group - 1) % world_size @contextmanager def graph_capture(self, graph_capture_context: GraphCaptureContext | None = None): if current_platform.is_cuda_alike(): if graph_capture_context is None: stream = torch.cuda.Stream() graph_capture_context = GraphCaptureContext(stream) else: stream = graph_capture_context.stream # ensure all initialization operations complete before attempting to # capture the graph on another stream curr_stream = torch.cuda.current_stream() if curr_stream != stream: stream.wait_stream(curr_stream) with torch.cuda.stream(stream): yield graph_capture_context else: # For non-CUDA platforms (MPS, CPU), just yield the context without stream management if graph_capture_context is None: # Create a dummy context for non-CUDA platforms graph_capture_context = GraphCaptureContext(None) yield graph_capture_context def all_to_all_4D( self, input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1 ) -> torch.Tensor: if self.world_size == 1: return input_ return self.device_communicator.all_to_all_4D(input_, scatter_dim, gather_dim) def all_reduce( self, input_: torch.Tensor, op=torch._C._distributed_c10d.ReduceOp.SUM, async_op: bool = False, ) -> torch.Tensor: """ NOTE: This operation will be applied in-place or out-of-place. Always assume this function modifies its input, but use the return value as the output. """ # Bypass the function if we are using only 1 GPU. if self.world_size == 1: return input_ else: torch.distributed.all_reduce( input_, op=op, group=self.device_group, async_op=async_op ) return input_ def all_gather( self, input_: torch.Tensor, dim: int = 0, separate_tensors: bool = False ) -> Union[torch.Tensor, List[torch.Tensor]]: world_size = self.world_size # Bypass the function if we are using only 1 GPU. if world_size == 1: return input_ assert ( -input_.dim() <= dim < input_.dim() ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" if dim < 0: # Convert negative dim to positive. dim += input_.dim() # Allocate output tensor. input_size = list(input_.size()) input_size[0] *= world_size output_tensor = torch.empty( input_size, dtype=input_.dtype, device=input_.device ) # All-gather. torch.distributed.all_gather_into_tensor( output_tensor, input_, group=self.device_group ) if dim != 0: input_size[0] //= world_size output_tensor = output_tensor.reshape( [ world_size, ] + input_size ) output_tensor = output_tensor.movedim(0, dim) if separate_tensors: tensor_list = [ output_tensor.reshape(-1) .narrow(0, input_.numel() * i, input_.numel()) .view_as(input_) for i in range(world_size) ] return tensor_list else: input_size = list(input_.size()) input_size[dim] = input_size[dim] * world_size # Reshape output_tensor = output_tensor.reshape(input_size) return output_tensor def gather(self, input_: torch.Tensor, dst: int = 0, dim: int = -1) -> torch.Tensor: """ NOTE: We assume that the input tensor is on the same device across all the ranks. NOTE: `dst` is the local rank of the destination rank. """ world_size = self.world_size # Bypass the function if we are using only 1 GPU. if world_size == 1: return input_ assert ( -input_.dim() <= dim < input_.dim() ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" if dim < 0: # Convert negative dim to positive. dim += input_.dim() # Allocate output tensor. if self.rank_in_group == dst: gather_list = [torch.empty_like(input_) for _ in range(world_size)] else: gather_list = None # Gather. torch.distributed.gather( input_, gather_list, dst=self.ranks[dst], group=self.device_group ) if self.rank_in_group == dst: output_tensor = torch.cat(gather_list, dim=dim) else: output_tensor = None return output_tensor def broadcast(self, input_: torch.Tensor, src: int = 0, async_op: bool = False): """Broadcast the input tensor. NOTE: `src` is the local rank of the source rank. """ assert src < self.world_size, f"Invalid src rank ({src})" # Bypass the function if we are using only 1 GPU. if self.world_size == 1: return input_ # Broadcast. torch.distributed.broadcast( input_, src=self.ranks[src], group=self.device_group, async_op=async_op, ) return input_ def broadcast_object(self, obj: Optional[Any] = None, src: int = 0): """Broadcast the input object. NOTE: `src` is the local rank of the source rank. """ assert src < self.world_size, f"Invalid src rank ({src})" # Bypass the function if we are using only 1 GPU. if self.world_size == 1: return obj if self.shm_broadcaster is not None: assert src == 0, "Shared memory broadcaster only supports src=0" return self.shm_broadcaster.broadcast_object(obj) if self.rank_in_group == src: torch.distributed.broadcast_object_list( [obj], src=self.ranks[src], group=self.cpu_group ) return obj else: recv = [None] torch.distributed.broadcast_object_list( recv, src=self.ranks[src], group=self.cpu_group ) return recv[0] def broadcast_object_list( self, obj_list: List[Any], src: int = 0, group: Optional[ProcessGroup] = None, ): """Broadcast the input object list. NOTE: `src` is the local rank of the source rank. """ assert src < self.world_size, f"Invalid src rank ({src})" # Bypass the function if we are using only 1 GPU. if self.world_size == 1: return obj_list # Broadcast. torch.distributed.broadcast_object_list( obj_list, src=self.ranks[src], group=self.device_group ) return obj_list def send_object(self, obj: Any, dst: int) -> None: """Send the input object list to the destination rank.""" """NOTE: `dst` is the local rank of the destination rank.""" assert dst < self.world_size, f"Invalid dst rank ({dst})" assert dst != self.rank, ( "Invalid destination rank. Destination rank is the same " "as the current rank." ) # Serialize object to tensor and get the size as well object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8) size_tensor = torch.tensor( [object_tensor.numel()], dtype=torch.long, device="cpu" ) # Send object size torch.distributed.send(size_tensor, dst=self.ranks[dst], group=self.cpu_group) # Send object torch.distributed.send(object_tensor, dst=self.ranks[dst], group=self.cpu_group) return None def recv_object(self, src: int) -> Any: """Receive the input object list from the source rank.""" """NOTE: `src` is the local rank of the source rank.""" assert src < self.world_size, f"Invalid src rank ({src})" assert ( src != self.rank ), "Invalid source rank. Source rank is the same as the current rank." size_tensor = torch.empty(1, dtype=torch.long, device="cpu") # Receive object size rank_size = torch.distributed.recv( size_tensor, src=self.ranks[src], group=self.cpu_group ) # Tensor to receive serialized objects into. object_tensor = torch.empty( # type: ignore[call-overload] size_tensor.item(), # type: ignore[arg-type] dtype=torch.uint8, device="cpu", ) rank_object = torch.distributed.recv( object_tensor, src=self.ranks[src], group=self.cpu_group ) assert ( rank_object == rank_size ), "Received object sender rank does not match the size sender rank." obj = pickle.loads(object_tensor.numpy().tobytes()) return obj def broadcast_tensor_dict( self, tensor_dict: Optional[Dict[str, Union[torch.Tensor, Any]]] = None, src: int = 0, group: Optional[ProcessGroup] = None, metadata_group: Optional[ProcessGroup] = None, ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]: """Broadcast the input tensor dictionary. NOTE: `src` is the local rank of the source rank. """ # Bypass the function if we are using only 1 GPU. if not torch.distributed.is_initialized() or self.world_size == 1: return tensor_dict group = self.device_group metadata_group = self.cpu_group assert src < self.world_size, f"Invalid src rank ({src})" src = self.ranks[src] rank = self.rank if rank == src: metadata_list: List[Tuple[Any, Any]] = [] assert isinstance( tensor_dict, dict ), f"Expecting a dictionary, got {type(tensor_dict)}" metadata_list, tensor_list = _split_tensor_dict(tensor_dict) # `metadata_list` lives in CPU memory. # `broadcast_object_list` has serialization & deserialization, # all happening on CPU. Therefore, we can use the CPU group. self.broadcast_object(metadata_list, src=src) async_handles = [] for tensor in tensor_list: if tensor.numel() == 0: # Skip broadcasting empty tensors. continue if tensor.is_cpu: # use metadata_group for CPU tensors handle = torch.distributed.broadcast( tensor, src=src, group=metadata_group, async_op=True ) else: # use group for GPU tensors handle = torch.distributed.broadcast( tensor, src=src, group=group, async_op=True ) async_handles.append(handle) for async_handle in async_handles: async_handle.wait() else: metadata_list = self.broadcast_object(None, src=src) tensor_dict = {} async_handles = [] for key, value in metadata_list: if isinstance(value, TensorMetadata): tensor = torch.empty( value.size, dtype=value.dtype, device=value.device ) if tensor.numel() == 0: # Skip broadcasting empty tensors. _update_nested_dict(tensor_dict, key, tensor) continue if tensor.is_cpu: # use metadata_group for CPU tensors handle = torch.distributed.broadcast( tensor, src=src, group=metadata_group, async_op=True ) else: # use group for GPU tensors handle = torch.distributed.broadcast( tensor, src=src, group=group, async_op=True ) async_handles.append(handle) _update_nested_dict(tensor_dict, key, tensor) else: _update_nested_dict(tensor_dict, key, value) for async_handle in async_handles: async_handle.wait() return tensor_dict def send_tensor_dict( self, tensor_dict: Dict[str, Union[torch.Tensor, Any]], dst: Optional[int] = None, ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]: """Send the input tensor dictionary. NOTE: `dst` is the local rank of the source rank. """ # Bypass the function if we are using only 1 GPU. if not torch.distributed.is_initialized() or self.world_size == 1: return tensor_dict group = self.device_group metadata_group = self.cpu_group if dst is None: dst = self.group_next_rank assert dst < self.world_size, f"Invalid dst rank ({dst})" metadata_list: List[Tuple[Any, Any]] = [] assert isinstance( tensor_dict, dict ), f"Expecting a dictionary, got {type(tensor_dict)}" metadata_list, tensor_list = _split_tensor_dict(tensor_dict) # `metadata_list` lives in CPU memory. # `send_object_list` has serialization & deserialization, # all happening on CPU. Therefore, we can use the CPU group. self.send_object(metadata_list, dst=dst) for tensor in tensor_list: if tensor.numel() == 0: # Skip sending empty tensors. continue if tensor.is_cpu: # use metadata_group for CPU tensors torch.distributed.send( tensor, dst=self.ranks[dst], group=metadata_group ) else: # use group for GPU tensors torch.distributed.send(tensor, dst=self.ranks[dst], group=group) return None def recv_tensor_dict( self, src: Optional[int] = None ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]: """Recv the input tensor dictionary. NOTE: `src` is the local rank of the source rank. """ # Bypass the function if we are using only 1 GPU. if not torch.distributed.is_initialized() or self.world_size == 1: return None group = self.device_group metadata_group = self.cpu_group if src is None: src = self.group_prev_rank assert src < self.world_size, f"Invalid src rank ({src})" recv_metadata_list = self.recv_object(src=src) tensor_dict: Dict[str, Any] = {} for key, value in recv_metadata_list: if isinstance(value, TensorMetadata): tensor = torch.empty(value.size, dtype=value.dtype, device=value.device) if tensor.numel() == 0: # Skip broadcasting empty tensors. _update_nested_dict(tensor_dict, key, tensor) continue if tensor.is_cpu: # use metadata_group for CPU tensors torch.distributed.recv( tensor, src=self.ranks[src], group=metadata_group ) else: # use group for GPU tensors torch.distributed.recv(tensor, src=self.ranks[src], group=group) _update_nested_dict(tensor_dict, key, tensor) else: _update_nested_dict(tensor_dict, key, value) return tensor_dict def barrier(self): """Barrier synchronization among the group. NOTE: don't use `device_group` here! `barrier` in NCCL is terrible because it is internally a broadcast operation with secretly created GPU tensors. It is easy to mess up the current device. Use the CPU group instead. """ torch.distributed.barrier(group=self.cpu_group) def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None: """Sends a tensor to the destination rank in a non-blocking way""" """NOTE: `dst` is the rank_in_group of the destination rank.""" if dst is None: dst = self.group_next_rank torch.distributed.send( tensor, self.ranks[dst], group=( self.device_groups[self.rank_in_group % 2] if self.world_size == 2 else self.device_group ), ) def recv( self, size: torch.Size, dtype: torch.dtype, src: Optional[int] = None ) -> torch.Tensor: """Receives a tensor from the src rank.""" """NOTE: `src` is the rank_in_group of the source rank.""" if src is None: src = self.group_prev_rank tensor = torch.empty(size, dtype=dtype, device=self.device) torch.distributed.recv( tensor, self.ranks[src], ( self.device_groups[(self.rank_in_group + 1) % 2] if self.world_size == 2 else self.device_group ), ) return tensor def destroy(self) -> None: if self.device_group is not None: torch.distributed.destroy_process_group(self.device_group) self.device_group = None if self.cpu_group is not None: torch.distributed.destroy_process_group(self.cpu_group) self.cpu_group = None if self.device_communicator is not None: self.device_communicator.destroy() if self.mq_broadcaster is not None: self.mq_broadcaster = None class PipelineGroupCoordinator(GroupCoordinator): """ available attributes: rank: int # global rank ranks: List[int] # global ranks in the group world_size: int # size of the group difference between `local_rank` and `rank_in_group`: if we have a group of size 4 across two nodes: Process | Node | Rank | Local Rank | Rank in Group 0 | 0 | 0 | 0 | 0 1 | 0 | 1 | 1 | 1 2 | 1 | 2 | 0 | 2 3 | 1 | 3 | 1 | 3 local_rank: int # local rank used to assign devices rank_in_group: int # rank inside the group cpu_group: ProcessGroup # group for CPU communication device_group: ProcessGroup # group for device communication """ def __init__( self, group_ranks: List[List[int]], local_rank: int, torch_distributed_backend: Union[str, Backend], group_name: str | None = None, ): super().__init__( group_ranks=group_ranks, local_rank=local_rank, torch_distributed_backend=torch_distributed_backend, group_name=group_name, ) self.rank = torch.distributed.get_rank() self.local_rank = local_rank self.device_group = None self.cpu_group = None self.cpu_groups = [] self.device_groups = [] if len(group_ranks[0]) > 2 or len(group_ranks[0]) == 1: for ranks in group_ranks: device_group = torch.distributed.new_group( ranks, backend=torch_distributed_backend ) # a group with `gloo` backend, to allow direct coordination between # processes through the CPU. with suppress_stdout(): cpu_group = torch.distributed.new_group(ranks, backend="gloo") if self.rank in ranks: self.ranks = ranks self.world_size = len(ranks) self.rank_in_group = ranks.index(self.rank) self.device_group = device_group self.cpu_group = cpu_group # when pipeline parallelism is 2, we need to create two groups to avoid # communication stall. # *_group_0_1 represents the group for communication from device 0 to # device 1. # *_group_1_0 represents the group for communication from device 1 to # device 0. elif len(group_ranks[0]) == 2: for ranks in group_ranks: device_group_0_1 = torch.distributed.new_group( ranks, backend=torch_distributed_backend ) device_group_1_0 = torch.distributed.new_group( ranks, backend=torch_distributed_backend ) # a group with `gloo` backend, to allow direct coordination between # processes through the CPU. with suppress_stdout(): cpu_group_0_1 = torch.distributed.new_group(ranks, backend="gloo") cpu_group_1_0 = torch.distributed.new_group(ranks, backend="gloo") if self.rank in ranks: self.ranks = ranks self.world_size = len(ranks) self.rank_in_group = ranks.index(self.rank) self.device_groups = [device_group_0_1, device_group_1_0] self.cpu_groups = [cpu_group_0_1, cpu_group_1_0] self.device_group = device_group_0_1 self.cpu_group = cpu_group_0_1 assert self.cpu_group is not None assert self.device_group is not None self.device = current_platform.get_device(local_rank) self.recv_buffer_set: bool = False self.recv_tasks_queue: List[Tuple[str, int]] = [] self.receiving_tasks: List[Tuple[torch.distributed.Work, str, int]] = [] self.dtype: Optional[torch.dtype] = None self.num_pipefusion_patches: Optional[int] = None self.recv_shape: Dict[str, Dict[int, torch.Size]] = {} self.send_shape: Dict[str, Dict[int, torch.Size]] = {} self.recv_buffer: Dict[str, Dict[int, torch.Size]] = {} self.skip_tensor_recv_buffer_set: bool = False self.recv_skip_tasks_queue: List[Union[int, Tuple[str, int]]] = [] self.receiving_skip_tasks: List[Tuple[torch.distributed.Work, str, int]] = [] self.skip_tensor_recv_buffer: Optional[ Union[List[torch.Tensor], torch.Tensor] ] = None self.skip_device_group = None for ranks in group_ranks: skip_device_group = torch.distributed.new_group( ranks, backend=torch_distributed_backend ) if self.rank in ranks: self.skip_device_group = skip_device_group assert self.skip_device_group is not None def reset_buffer(self): self.recv_tasks_queue = [] self.receiving_tasks = [] self.recv_shape = {} self.send_shape = {} self.recv_buffer = {} self.recv_skip_tasks_queue = [] self.receiving_skip_tasks = [] self.skip_tensor_recv_buffer = {} def set_config(self, dtype: torch.dtype): self.dtype = dtype def set_recv_buffer( self, num_pipefusion_patches: int, patches_shape_list: List[List[int]], feature_map_shape: List[int], dtype: torch.dtype, ): assert isinstance(dtype, torch.dtype), "dtype must be a torch.dtype object" assert ( isinstance(num_pipefusion_patches, int) and num_pipefusion_patches >= 1 ), "num_pipefusion_patches must be greater than or equal to 1" self.dtype = dtype self.num_pipefusion_patches = num_pipefusion_patches self.recv_buffer = [ torch.zeros(*shape, dtype=self.dtype, device=self.device) for shape in patches_shape_list ] self.recv_buffer.append( torch.zeros(*feature_map_shape, dtype=self.dtype, device=self.device) ) self.recv_buffer_set = True def set_extra_tensors_recv_buffer( self, name: str, shape: List[int], num_buffers: int = 1, dtype: torch.dtype = torch.float16, ): self.extra_tensors_recv_buffer[name] = [ torch.zeros(*shape, dtype=dtype, device=self.device) for _ in range(num_buffers) ] def _check_shape_and_buffer( self, tensor_send_to_next=None, recv_prev=False, name: Optional[str] = None, segment_idx: int = 0, ): send_flag = False name = name or "latent" if tensor_send_to_next is not None: shape_list = self.send_shape.get(name, None) if shape_list is None: self.send_shape[name] = {segment_idx: tensor_send_to_next.shape} send_flag = True elif shape_list.get(segment_idx, None) is None: self.send_shape[name][segment_idx] = tensor_send_to_next.shape send_flag = True recv_flag = False if recv_prev: shape_list = self.recv_shape.get(name, None) if shape_list is None: recv_flag = True elif shape_list.get(segment_idx, None) is None: recv_flag = True recv_prev_shape = self._communicate_shapes( tensor_send_to_next=tensor_send_to_next if send_flag else None, recv_prev=recv_flag, ) if recv_flag: if self.recv_shape.get(name, None) is None: self.recv_shape[name] = {segment_idx: recv_prev_shape} else: self.recv_shape[name][segment_idx] = recv_prev_shape if self.recv_buffer.get(name, None) is None: self.recv_buffer[name] = { segment_idx: torch.zeros( recv_prev_shape, device=self.device, dtype=self.dtype ) } else: if self.recv_buffer[name].get(segment_idx, None) is not None: logger.warning( f"Recv buffer [name: {name}, segment_idx: {segment_idx}] already exist. updating..." ) self.recv_buffer[name][segment_idx] = torch.zeros( recv_prev_shape, device=self.device, dtype=self.dtype ) def _communicate_shapes(self, tensor_send_to_next=None, recv_prev=False): """Communicate tensor shapes between stages. Used to communicate tensor shapes before the actual tensor communication happens. Args: tensor_send_next: tensor to send to next rank (no tensor sent if set to None). recv_prev: boolean for whether tensor should be received from previous rank. """ ops = [] if recv_prev: recv_prev_dim_tensor = torch.empty( (1), device=self.device, dtype=torch.int64 ) recv_prev_dim_op = torch.distributed.P2POp( torch.distributed.irecv, recv_prev_dim_tensor, self.prev_rank, self.device_group, ) ops.append(recv_prev_dim_op) if tensor_send_to_next is not None: send_next_dim_tensor = torch.tensor( tensor_send_to_next.dim(), device=self.device, dtype=torch.int64 ) send_next_dim_op = torch.distributed.P2POp( torch.distributed.isend, send_next_dim_tensor, self.next_rank, self.device_group, ) ops.append(send_next_dim_op) if len(ops) > 0: reqs = torch.distributed.batch_isend_irecv(ops) for req in reqs: req.wait() # To protect against race condition when using batch_isend_irecv(). # should take this out once the bug with batch_isend_irecv is resolved. synchronize() ops = [] recv_prev_shape_tensor = None if recv_prev: recv_prev_shape_tensor = torch.empty( torch.Size(recv_prev_dim_tensor), device=self.device, dtype=torch.int64, ) recv_prev_shape_op = torch.distributed.P2POp( torch.distributed.irecv, recv_prev_shape_tensor, self.prev_rank, self.device_group, ) ops.append(recv_prev_shape_op) if tensor_send_to_next is not None: send_next_shape_tensor = torch.tensor( tensor_send_to_next.size(), device=self.device, dtype=torch.int64, ) send_next_shape_op = torch.distributed.P2POp( torch.distributed.isend, send_next_shape_tensor, self.next_rank, self.device_group, ) ops.append(send_next_shape_op) if len(ops) > 0: reqs = torch.distributed.batch_isend_irecv(ops) for req in reqs: req.wait() synchronize() recv_prev_shape = [0, 0, 0] if recv_prev_shape_tensor is not None: recv_prev_shape = recv_prev_shape_tensor return torch.Size(recv_prev_shape) def pipeline_send( self, tensor: torch.Tensor, name: str = "latent", segment_idx: int = -1 ) -> None: tensor = tensor.contiguous() self._check_shape_and_buffer( tensor_send_to_next=tensor, name=name, segment_idx=segment_idx ) self._pipeline_isend(tensor).wait() def pipeline_isend( self, tensor: torch.Tensor, name: str = "latent", segment_idx: int = -1 ) -> None: tensor = tensor.contiguous() self._check_shape_and_buffer( tensor_send_to_next=tensor, name=name, segment_idx=segment_idx ) self._pipeline_isend(tensor) def pipeline_recv(self, idx: int = -1, name: str = "latent") -> torch.Tensor: name = name or "latent" self._check_shape_and_buffer(recv_prev=True, name=name, segment_idx=idx) self._pipeline_irecv(self.recv_buffer[name][idx]).wait() return self.recv_buffer[name][idx] def add_pipeline_recv_task(self, idx: int = -1, name: str = "latent"): name = name or "latent" self.recv_tasks_queue.append((name, idx)) def recv_next(self): if len(self.recv_tasks_queue) == 0: raise ValueError("No more tasks to receive") elif len(self.recv_tasks_queue) > 0: name, idx = self.recv_tasks_queue.pop(0) self._check_shape_and_buffer(recv_prev=True, name=name, segment_idx=idx) self.receiving_tasks.append( (self._pipeline_irecv(self.recv_buffer[name][idx]), name, idx) ) def get_pipeline_recv_data( self, idx: int = -1, name: str = "latent" ) -> torch.Tensor: assert ( len(self.receiving_tasks) > 0 ), "No tasks to receive, call add_pipeline_recv_task first" receiving_task = self.receiving_tasks.pop(0) receiving_task[0].wait() assert ( receiving_task[1] == name and receiving_task[2] == idx ), "Received tensor does not match the requested" return self.recv_buffer[name][idx] def _pipeline_irecv(self, tensor: torch.tensor): return torch.distributed.irecv( tensor, src=self.prev_rank, group=( self.device_groups[(self.rank_in_group + 1) % 2] if self.world_size == 2 else self.device_group ), ) def _pipeline_isend(self, tensor: torch.tensor): return torch.distributed.isend( tensor, dst=self.next_rank, group=( self.device_groups[self.rank_in_group % 2] if self.world_size == 2 else self.device_group ), ) def set_skip_tensor_recv_buffer( self, patches_shape_list: List[List[int]], feature_map_shape: List[int], ): self.skip_tensor_recv_buffer = [ torch.zeros(*shape, dtype=self.dtype, device=self.device) for shape in patches_shape_list ] self.skip_tensor_recv_buffer.append( torch.zeros(*feature_map_shape, dtype=self.dtype, device=self.device) ) self.skip_tensor_recv_buffer_set = True def pipeline_send_skip(self, tensor: torch.Tensor) -> None: tensor = tensor.contiguous() self._pipeline_isend_skip(tensor).wait() def pipeline_isend_skip(self, tensor: torch.Tensor) -> None: tensor = tensor.contiguous() self._pipeline_isend_skip(tensor) def pipeline_recv_skip(self, idx: int = -1) -> torch.Tensor: self._pipeline_irecv_skip(self.skip_tensor_recv_buffer[idx]).wait() return self.skip_tensor_recv_buffer[idx] def add_pipeline_recv_skip_task(self, idx: int = -1): self.recv_skip_tasks_queue.append(idx) def get_pipeline_recv_skip_data(self, idx: int = -1) -> torch.Tensor: assert ( len(self.receiving_skip_tasks) > 0 ), "No tasks to receive, call add_pipeline_recv_skip_task first" receiving_skip_task = self.receiving_skip_tasks.pop(0) receiving_skip_task[0].wait() assert ( receiving_skip_task[2] == idx ), "Received tensor does not match the requested" return self.skip_tensor_recv_buffer[idx] def recv_skip_next(self): if len(self.recv_skip_tasks_queue) == 0: raise ValueError("No more tasks to receive") elif len(self.recv_skip_tasks_queue) > 0: task = self.recv_skip_tasks_queue.pop(0) idx = task self.receiving_skip_tasks.append( ( self._pipeline_irecv_skip(self.skip_tensor_recv_buffer[idx]), None, idx, ) ) def _pipeline_irecv_skip(self, tensor: torch.tensor): return torch.distributed.irecv( tensor, src=self.skip_rank, group=self.skip_device_group ) def _pipeline_isend_skip(self, tensor: torch.tensor): return torch.distributed.isend( tensor, dst=self.skip_rank, group=self.skip_device_group ) class SequenceParallelGroupCoordinator(GroupCoordinator): def __init__( self, group_ranks: List[List[int]], local_rank: int, torch_distributed_backend: Union[str, Backend], group_name: str | None = None, **kwargs, ): super().__init__( group_ranks=group_ranks, local_rank=local_rank, torch_distributed_backend=torch_distributed_backend, group_name=group_name, ) ulysses_group = kwargs.get("ulysses_group", None) ring_group = kwargs.get("ring_group", None) if ulysses_group is None: raise RuntimeError( f"Please pass argument 'ulysses_group' when calling init func of SequenceParallelGroupCoordinator" ) if ring_group is None: raise RuntimeError( f"Please pass argument 'ring_group' when calling init func of SequenceParallelGroupCoordinator" ) self.ulysses_group = ulysses_group self.ring_group = ring_group self.ulysses_world_size = torch.distributed.get_world_size(self.ulysses_group) self.ulysses_rank = torch.distributed.get_rank(self.ulysses_group) self.ring_world_size = torch.distributed.get_world_size(self.ring_group) self.ring_rank = torch.distributed.get_rank(self.ring_group) ================================================ FILE: python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py ================================================ # Reference: https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py import torch class Singleton: _instance = None def __new__(cls, *args, **kwargs): if not cls._instance: cls._instance = super(Singleton, cls).__new__(cls, *args, **kwargs) return cls._instance class ProcessGroupSingleton(Singleton): def __init__(self): self.ULYSSES_PG = None self.RING_PG = None PROCESS_GROUP = ProcessGroupSingleton() def set_seq_parallel_pg_by_sp_groups( sp_ulysses_degree, sp_ring_degree, rank: int, sp_groups: list[list[int]], use_ulysses_low: bool = True, ): """Create Ulysses/Ring process groups inside each SP group. This is required when TP>1, because SP groups are not necessarily made of consecutive global ranks (e.g., tp-sp order makes SP ranks strided). Args: sp_ulysses_degree: ulysses degree inside SP. sp_ring_degree: ring degree inside SP. rank: global rank of current process. sp_groups: list of global-rank lists for each SP group. use_ulysses_low: keep the same semantics as the original function. """ sp_degree = sp_ring_degree * sp_ulysses_degree assert sp_degree > 0 assert all( len(g) == sp_degree for g in sp_groups ), f"Each SP group must have size {sp_degree}, got sizes {[len(g) for g in sp_groups]}" ulyssess_pg = None ring_pg = None num_ulysses_pgs = sp_ring_degree num_ring_pgs = sp_ulysses_degree def _map_indices_to_ranks(ranks: list[int], indices: list[int]) -> list[int]: return [ranks[i] for i in indices] # Important: call torch.distributed.new_group in the same order on all ranks. for sp_ranks in sp_groups: if use_ulysses_low: for i in range(num_ulysses_pgs): idx = list(range(i * sp_ulysses_degree, (i + 1) * sp_ulysses_degree)) ulysses_ranks = _map_indices_to_ranks(sp_ranks, idx) group = torch.distributed.new_group(ulysses_ranks) if rank in ulysses_ranks: ulyssess_pg = group for i in range(num_ring_pgs): idx = list(range(i, sp_degree, num_ring_pgs)) ring_ranks = _map_indices_to_ranks(sp_ranks, idx) group = torch.distributed.new_group(ring_ranks) if rank in ring_ranks: ring_pg = group else: for i in range(num_ring_pgs): idx = list(range(i * sp_ring_degree, (i + 1) * sp_ring_degree)) ring_ranks = _map_indices_to_ranks(sp_ranks, idx) group = torch.distributed.new_group(ring_ranks) if rank in ring_ranks: ring_pg = group for i in range(num_ulysses_pgs): idx = list(range(i, sp_degree, num_ulysses_pgs)) ulysses_ranks = _map_indices_to_ranks(sp_ranks, idx) group = torch.distributed.new_group(ulysses_ranks) if rank in ulysses_ranks: ulyssess_pg = group PROCESS_GROUP.ULYSSES_PG = ulyssess_pg PROCESS_GROUP.RING_PG = ring_pg ================================================ FILE: python/sglang/multimodal_gen/runtime/distributed/parallel_state.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/parallel_state.py # Copyright 2023 The vLLM team. # Adapted from # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. # Adapted from # Copyright 2024 xDiT team. # Adapted from # https://github.com/vllm-project/vllm/blob/main/vllm/distributed/parallel_state.py # Copyright 2023 The vLLM team. # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. """sglang-diffusion distributed state. It takes over the control of the distributed environment from PyTorch. The typical workflow is: - call `init_distributed_environment` to initialize the distributed environment. - call `initialize_model_parallel` or `ensure_model_parallel_initialized` to initialize the model parallel groups. - any code dealing with the distributed stuff - call `destroy_model_parallel` to destroy the model parallel groups. - call `destroy_distributed_environment` to destroy the distributed environment. If you only need to use the distributed environment without model parallelism, you can skip the model parallel initialization and destruction steps. """ import contextlib import datetime import os import weakref from collections import namedtuple from collections.abc import Callable from contextlib import contextmanager from multiprocessing import shared_memory from typing import Any, List, Optional from unittest.mock import patch import torch import torch.distributed from torch.distributed import ProcessGroup import sglang.multimodal_gen.envs as envs from sglang.multimodal_gen.runtime.distributed.utils import StatelessProcessGroup from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from ..utils.distributed import RankGenerator from .group_coordinator import ( GroupCoordinator, PipelineGroupCoordinator, SequenceParallelGroupCoordinator, get_local_torch_device, ) logger = init_logger(__name__) _WORLD: GroupCoordinator | None = None _TP: GroupCoordinator | None = None _SP: SequenceParallelGroupCoordinator | None = None _PP: PipelineGroupCoordinator | None = None _CFG: GroupCoordinator | None = None _DP: GroupCoordinator | None = None _DIT: ProcessGroup | None = None _VAE: ProcessGroup | None = None TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) def _split_tensor_dict( tensor_dict: dict[str, torch.Tensor | Any], ) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]: """Split the tensor dictionary into two parts: 1. A list of (key, value) pairs. If the value is a tensor, it is replaced by its metadata. 2. A list of tensors. """ metadata_list: list[tuple[str, Any]] = [] tensor_list: list[torch.Tensor] = [] for key, value in tensor_dict.items(): if isinstance(value, torch.Tensor): # Note: we cannot use `value.device` here, # because it contains not only the device type but also the device # index (e.g. "cuda:0"). We only need the device type. # receiving side will set the device index. device = value.device.type metadata_list.append( (key, TensorMetadata(device, value.dtype, value.size())) ) tensor_list.append(value) else: metadata_list.append((key, value)) return metadata_list, tensor_list _groups: dict[str, Callable[[], Optional["GroupCoordinator"]]] = {} def _register_group(group: "GroupCoordinator") -> None: _groups[group.unique_name] = weakref.ref(group) def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor: assert group_name in _groups, f"Group {group_name} is not found." group = _groups[group_name]() if group is None: raise ValueError(f"Group {group_name} is destroyed.") return group._all_reduce_out_place(tensor) def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor: return torch.empty_like(tensor) def get_world_group() -> GroupCoordinator: assert _WORLD is not None, "world group is not initialized" return _WORLD def init_world_group( ranks: list[int], local_rank: int, backend: str ) -> GroupCoordinator: return GroupCoordinator( group_ranks=[ranks], local_rank=local_rank, torch_distributed_backend=backend, use_device_communicator=True, group_name="world", ) def init_parallel_group_coordinator( group_ranks: List[List[int]], local_rank: int, backend: str, parallel_mode: str, **kwargs, ) -> GroupCoordinator: """Return a group coordinator for the given parallel mode.""" assert parallel_mode in [ "data", "pipeline", "tensor", "sequence", "classifier_free_guidance", ], f"parallel_mode {parallel_mode} is not supported" if parallel_mode == "pipeline": return PipelineGroupCoordinator( group_ranks=group_ranks, local_rank=local_rank, torch_distributed_backend=backend, group_name="pp_group", ) elif parallel_mode == "sequence": return SequenceParallelGroupCoordinator( group_ranks=group_ranks, local_rank=local_rank, torch_distributed_backend=backend, group_name="sp_group", **kwargs, ) else: # fallback to GroupCoordinator return GroupCoordinator( group_ranks=group_ranks, local_rank=local_rank, torch_distributed_backend=backend, group_name="cfg_group", ) def get_tp_group() -> GroupCoordinator: assert _TP is not None, "tensor model parallel group is not initialized" return _TP def init_distributed_environment( world_size: int = 1, rank: int = 0, distributed_init_method: str = "env://", local_rank: int = 0, backend: str = "nccl", device_id: torch.device | None = None, timeout: int | None = None, ): # Determine the appropriate backend based on the platform from sglang.multimodal_gen.runtime.platforms import current_platform if backend == "nccl" and not current_platform.is_cuda_alike(): # Use gloo backend for non-CUDA platforms (MPS, CPU) backend = "gloo" logger.info("Using gloo backend for %s platform", current_platform.device_name) logger.debug( "world_size=%d rank=%d local_rank=%d " "distributed_init_method=%s backend=%s timeout=%s", world_size, rank, local_rank, distributed_init_method, backend, timeout, ) if not torch.distributed.is_initialized(): assert distributed_init_method is not None, ( "distributed_init_method must be provided when initializing " "distributed environment" ) # For MPS and MUSA, don't pass device_id as it doesn't support device indices extra_args = ( {} if ( current_platform.is_mps() or current_platform.is_musa() or current_platform.is_npu() ) else dict(device_id=device_id) ) if timeout is not None: extra_args["timeout"] = datetime.timedelta(seconds=timeout) logger.info(f"Setting distributed timeout to {timeout} seconds") torch.distributed.init_process_group( backend=backend, init_method=distributed_init_method, world_size=world_size, rank=rank, **extra_args, ) # set the local rank # local_rank is not available in torch ProcessGroup, # see https://github.com/pytorch/pytorch/issues/122816 if local_rank == -1: # local rank not set, this usually happens in single-node # setting, where we can use rank as local rank if distributed_init_method == "env://": local_rank = envs.LOCAL_RANK else: local_rank = rank global _WORLD if _WORLD is None: ranks = list(range(torch.distributed.get_world_size())) _WORLD = init_world_group(ranks, local_rank, backend) else: assert ( _WORLD.world_size == torch.distributed.get_world_size() ), "world group already initialized with a different world size" def get_sp_group() -> SequenceParallelGroupCoordinator: assert _SP is not None, "sequence parallel group is not initialized" return _SP def get_dp_group() -> GroupCoordinator: assert _DP is not None, "data parallel group is not initialized" return _DP # xDiT def initialize_model_parallel( data_parallel_size: int = 1, classifier_free_guidance_degree: int = 1, sequence_parallel_degree: Optional[int] = None, ulysses_degree: int = 1, ring_degree: int = 1, tensor_parallel_degree: int = 1, pipeline_parallel_degree: int = 1, vae_parallel_size: int = 0, backend: Optional[str] = None, ) -> None: """ Initialize model parallel groups. Arguments: data_parallel_size: number of data parallelism groups. classifier_free_guidance_degree: number of GPUs used for Classifier Free Guidance (CFG) sequence_parallel_degree: number of GPUs used for sequence parallelism. sequence_parallel_degree = ulysses_degree * ring_degree ulysses_degree: number of GPUs used for ulysses sequence parallelism. ring_degree: number of GPUs used for ring sequence parallelism. tensor_parallel_degree: number of GPUs used for tensor parallelism. pipeline_parallel_degree: number of GPUs used for pipeline parallelism. backend: distributed backend of pytorch collective comm. Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we use 2 groups to parallelize the batch dim(dp), 2 groups to parallelize split batch caused by CFG, and 2 GPUs to parallelize sequence. dp_degree (2) * cfg_degree (2) * sp_degree (2) * pp_degree (2) = 16. The present function will create 8 data-parallel groups, 8 CFG group, 8 pipeline-parallel group, and 8 sequence-parallel groups: 8 data-parallel groups: [g0, g8], [g1, g9], [g2, g10], [g3, g11], [g4, g12], [g5, g13], [g6, g14], [g7, g15] 8 CFG-parallel groups: [g0, g4], [g1, g5], [g2, g6], [g3, g7], [g8, g12], [g9, g13], [g10, g14], [g11, g15] 8 sequence-parallel groups: [g0, g1], [g2, g3], [g4, g5], [g6, g7], [g8, g9], [g10, g11], [g12, g13], [g14, g15] 8 pipeline-parallel groups: [g0, g2], [g4, g6], [g8, g10], [g12, g14], [g1, g3], [g5, g7], [g9, g11], [g13, g15] Note that for efficiency, the caller should make sure adjacent ranks are on the same DGX box. For example if we are using 2 DGX-1 boxes with a total of 16 GPUs, rank 0 to 7 belong to the first box and ranks 8 to 15 belong to the second box. """ if backend is None: from sglang.multimodal_gen.runtime.platforms import current_platform backend = current_platform.get_torch_distributed_backend_str() # Get world size and rank. Ensure some consistencies. assert torch.distributed.is_initialized() world_size: int = torch.distributed.get_world_size() backend = backend or torch.distributed.get_backend(get_world_group().device_group) dit_parallel_size = ( data_parallel_size * classifier_free_guidance_degree * sequence_parallel_degree * pipeline_parallel_degree * tensor_parallel_degree ) if world_size < dit_parallel_size: raise RuntimeError( f"world_size ({world_size}) is less than " f"tensor_parallel_degree ({tensor_parallel_degree}) x " f"pipeline_parallel_degree ({pipeline_parallel_degree}) x" f"sequence_parallel_degree ({sequence_parallel_degree}) x" f"classifier_free_guidance_degree " f"({classifier_free_guidance_degree}) x" f"data_parallel_degree ({data_parallel_size})" ) rank_generator: RankGenerator = RankGenerator( tensor_parallel_degree, sequence_parallel_degree, pipeline_parallel_degree, classifier_free_guidance_degree, data_parallel_size, "tp-sp-pp-cfg-dp", ) global _DP assert _DP is None, "data parallel group is already initialized" _DP = init_parallel_group_coordinator( group_ranks=rank_generator.get_ranks("dp"), local_rank=get_world_group().local_rank, backend=backend, parallel_mode="data", ) global _CFG assert _CFG is None, "classifier_free_guidance group is already initialized" _CFG = init_parallel_group_coordinator( group_ranks=rank_generator.get_ranks("cfg"), local_rank=get_world_group().local_rank, backend=backend, parallel_mode="classifier_free_guidance", ) global _PP assert _PP is None, "pipeline model parallel group is already initialized" _PP = init_parallel_group_coordinator( group_ranks=rank_generator.get_ranks("pp"), local_rank=get_world_group().local_rank, backend=backend, parallel_mode="pipeline", ) global _SP assert _SP is None, "sequence parallel group is already initialized" try: from .parallel_groups import PROCESS_GROUP as _YC_PROCESS_GROUP from .parallel_groups import ( set_seq_parallel_pg_by_sp_groups as _set_seq_parallel_pg_by_sp_groups, ) except ImportError: _set_seq_parallel_pg_by_sp_groups = None class _DummyProcessGroup: ULYSSES_PG = torch.distributed.group.WORLD RING_PG = torch.distributed.group.WORLD PROCESS_GROUP = _DummyProcessGroup() else: # Build SGLang Diffusion SP sub-groups based on the true SP groups. This is # critical when TP>1, because SP groups may be strided in global ranks # (e.g., tp-sp order). sp_groups = rank_generator.get_ranks("sp") _set_seq_parallel_pg_by_sp_groups( sp_ulysses_degree=ulysses_degree, sp_ring_degree=ring_degree, rank=get_world_group().rank, sp_groups=sp_groups, ) PROCESS_GROUP = _YC_PROCESS_GROUP _SP = init_parallel_group_coordinator( group_ranks=rank_generator.get_ranks("sp"), local_rank=get_world_group().local_rank, backend=backend, parallel_mode="sequence", ulysses_group=PROCESS_GROUP.ULYSSES_PG, ring_group=PROCESS_GROUP.RING_PG, ) global _TP assert _TP is None, "Tensor parallel group is already initialized" _TP = init_parallel_group_coordinator( group_ranks=rank_generator.get_ranks("tp"), local_rank=get_world_group().local_rank, backend=backend, parallel_mode="tensor", ) if vae_parallel_size > 0: init_vae_group(dit_parallel_size, vae_parallel_size, backend) init_dit_group(dit_parallel_size, backend) def get_sp_world_size() -> int: """Return world size for the sequence model parallel group.""" return get_sp_group().world_size def get_sp_parallel_rank() -> int: """Return my rank for the sequence model parallel group.""" return get_sp_group().rank_in_group def get_world_size() -> int: """Return world size for the world group.""" return get_world_group().world_size def get_world_rank() -> int: """Return my rank for the world group.""" return get_world_group().rank def get_dp_world_size() -> int: """Return world size for the data parallel group.""" return get_dp_group().world_size def get_dp_rank() -> int: """Return my rank for the data parallel group.""" return get_dp_group().rank_in_group def maybe_init_distributed_environment_and_model_parallel( tp_size: int, sp_size: int, enable_cfg_parallel: bool, ulysses_degree: int = 1, ring_degree: int = 1, dp_size: int = 1, distributed_init_method: str = "env://", dist_timeout: int | None = None, ): from sglang.multimodal_gen.runtime.platforms import current_platform if _WORLD is not None and model_parallel_is_initialized(): # make sure the tp and sp sizes are correct assert ( get_tp_world_size() == tp_size ), f"You are trying to initialize model parallel groups with size {tp_size}, but they are already initialized with size {get_tp_world_size()}" assert ( get_sp_world_size() == sp_size ), f"You are trying to initialize model parallel groups with size {sp_size}, but they are already initialized with size {get_sp_world_size()}" return local_rank = int(os.environ.get("LOCAL_RANK", 0)) world_size = int(os.environ.get("WORLD_SIZE", 1)) rank = int(os.environ.get("RANK", 0)) device = get_local_torch_device() logger.info( "Initializing distributed environment with world_size=%d, device=%s, timeout=%s", world_size, device, dist_timeout, main_process_only=False, ) init_distributed_environment( world_size=world_size, rank=rank, local_rank=local_rank, distributed_init_method=distributed_init_method, device_id=device, backend=current_platform.get_torch_distributed_backend_str(), timeout=dist_timeout, ) initialize_model_parallel( data_parallel_size=dp_size, classifier_free_guidance_degree=2 if enable_cfg_parallel else 1, tensor_parallel_degree=tp_size, ulysses_degree=ulysses_degree, ring_degree=ring_degree, sequence_parallel_degree=sp_size, ) # Only set CUDA device if we're on a CUDA platform if current_platform.is_cuda_alike(): device = torch.device(f"cuda:{local_rank}") torch.cuda.set_device(device) elif current_platform.is_npu(): device = torch.device(f"npu:{local_rank}") torch.npu.set_device(device) def model_parallel_is_initialized() -> bool: """Check if model parallel groups are initialized.""" return ( _DP is not None and _CFG is not None and _SP is not None and _PP is not None and _TP is not None ) _TP_STATE_PATCHED = False @contextmanager def patch_tensor_parallel_group(tp_group: GroupCoordinator): """Patch the tp group temporarily until this function ends. This method is for draft workers of speculative decoding to run draft model with different tp degree from that of target model workers. """ global _TP_STATE_PATCHED assert not _TP_STATE_PATCHED, "Should not call when it's already patched" _TP_STATE_PATCHED = True old_tp_group = get_tp_group() global _TP _TP = tp_group try: yield finally: # restore the original state _TP_STATE_PATCHED = False _TP = old_tp_group def get_tp_world_size() -> int: """Return world size for the tensor model parallel group.""" return get_tp_group().world_size def get_tp_rank() -> int: """Return my rank for the tensor model parallel group.""" return get_tp_group().rank_in_group def destroy_distributed_environment() -> None: global _WORLD if _WORLD: _WORLD.destroy() _WORLD = None if torch.distributed.is_initialized(): torch.distributed.destroy_process_group() def cleanup_dist_env_and_memory(shutdown_ray: bool = False): destroy_model_parallel() destroy_distributed_environment() with contextlib.suppress(AssertionError): torch.distributed.destroy_process_group() if shutdown_ray: import ray # Lazy import Ray ray.shutdown() def is_the_same_node_as( pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0 ) -> list[int]: """ This is a collective operation that returns if each rank is in the same node as the source rank. It tests if processes are attached to the same memory system (shared access to shared memory). """ if isinstance(pg, ProcessGroup): assert ( torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL ), "in_the_same_node_as should be tested with a non-NCCL group." # local rank inside the group rank = torch.distributed.get_rank(group=pg) world_size = torch.distributed.get_world_size(group=pg) # global ranks of the processes in the group ranks = torch.distributed.get_process_group_ranks(pg) else: rank = pg.rank world_size = pg.world_size ranks = list(range(world_size)) # local tensor in each process to store the result is_in_the_same_node = torch.tensor([0] * world_size, dtype=torch.int32) magic_message = b"magic_message" shm = None try: with contextlib.suppress(OSError): if rank == source_rank: # create a shared memory segment shm = shared_memory.SharedMemory(create=True, size=128) shm.buf[: len(magic_message)] = magic_message if isinstance(pg, ProcessGroup): torch.distributed.broadcast_object_list( [shm.name], src=ranks[source_rank], group=pg ) else: pg.broadcast_obj(shm.name, src=source_rank) is_in_the_same_node[rank] = 1 else: # try to open the shared memory segment if isinstance(pg, ProcessGroup): recv = [None] torch.distributed.broadcast_object_list( recv, src=ranks[source_rank], group=pg ) name = recv[0] else: name = pg.broadcast_obj(None, src=source_rank) # fix to https://stackoverflow.com/q/62748654/9191338 # Python incorrectly tracks shared memory even if it is not # created by the process. The following patch is a workaround. with patch( "multiprocessing.resource_tracker.register", lambda *args, **kwargs: None, ): shm = shared_memory.SharedMemory(name=name) if shm.buf[: len(magic_message)] == magic_message: is_in_the_same_node[rank] = 1 except Exception as e: logger.error("Error ignored in is_in_the_same_node: %s", e) finally: if shm: shm.close() if isinstance(pg, ProcessGroup): torch.distributed.barrier(group=pg) else: pg.barrier() # clean up the shared memory segment with contextlib.suppress(OSError): if rank == source_rank and shm: shm.unlink() if isinstance(pg, ProcessGroup): torch.distributed.all_reduce(is_in_the_same_node, group=pg) aggregated_data = is_in_the_same_node else: aggregated_data = torch.zeros_like(is_in_the_same_node) for i in range(world_size): rank_data = pg.broadcast_obj(is_in_the_same_node, src=i) aggregated_data += rank_data return [x == 1 for x in aggregated_data.tolist()] def get_tensor_model_parallel_world_size() -> int: """Return world size for the tensor model parallel group.""" return get_tp_world_size() def get_tensor_model_parallel_rank() -> int: """Return my rank for the tensor model parallel group.""" return get_tp_rank() def get_sequence_parallel_world_size() -> int: """Return world size for the sequence parallel group.""" return get_sp_world_size() def get_sequence_parallel_rank() -> int: """Return my rank for the sequence parallel group.""" return get_sp_parallel_rank() def get_ulysses_parallel_world_size() -> int: return get_sp_group().ulysses_world_size def get_ulysses_parallel_rank() -> int: return get_sp_group().ulysses_rank def get_ring_parallel_world_size() -> int: return get_sp_group().ring_world_size def get_ring_parallel_rank() -> int: return get_sp_group().ring_rank # PP def get_pp_group() -> PipelineGroupCoordinator: assert _PP is not None, "pipeline model parallel group is not initialized" return _PP def get_pipeline_parallel_world_size() -> int: """Return world size for the pipeline model parallel group.""" return get_pp_group().world_size def get_pipeline_parallel_rank() -> int: """Return my rank for the pipeline model parallel group.""" return get_pp_group().rank_in_group def is_pipeline_first_stage() -> bool: """Return True if in the first pipeline model parallel stage, False otherwise.""" return get_pipeline_parallel_rank() == 0 def is_pipeline_last_stage() -> bool: """Return True if in the last pipeline model parallel stage, False otherwise.""" return get_pipeline_parallel_rank() == (get_pipeline_parallel_world_size() - 1) # CFG def get_cfg_group() -> GroupCoordinator: assert ( _CFG is not None ), "classifier_free_guidance parallel group is not initialized" return _CFG def get_classifier_free_guidance_world_size() -> int: """Return world size for the classifier_free_guidance parallel group.""" return get_cfg_group().world_size def get_classifier_free_guidance_rank() -> int: """Return my rank for the classifier_free_guidance parallel group.""" return get_cfg_group().rank_in_group def get_data_parallel_world_size() -> int: """Return world size for the data parallel group.""" return get_dp_world_size() def get_data_parallel_rank() -> int: """Return my rank for the data parallel group.""" return get_dp_rank() def is_dp_last_group() -> bool: """Return True if in the last data parallel group, False otherwise.""" return ( get_sequence_parallel_rank() == (get_sequence_parallel_world_size() - 1) and get_classifier_free_guidance_rank() == (get_classifier_free_guidance_world_size() - 1) and get_pipeline_parallel_rank() == (get_pipeline_parallel_world_size() - 1) ) def get_dit_world_size() -> int: """Return world size for the DiT model (excluding VAE).""" return ( get_data_parallel_world_size() * get_classifier_free_guidance_world_size() * get_sequence_parallel_world_size() * get_pipeline_parallel_world_size() * get_tensor_model_parallel_world_size() ) def get_vae_parallel_group() -> ProcessGroup: assert _VAE is not None, "VAE parallel group is not initialized" return _VAE def get_vae_parallel_world_size() -> int: """Return world size for the VAE parallel group.""" return torch.distributed.get_world_size(group=get_vae_parallel_group()) def get_vae_parallel_rank() -> int: """Return my rank for the VAE parallel group.""" return torch.distributed.get_rank(group=get_vae_parallel_group()) def init_dit_group( dit_parallel_size: int, backend: str, ) -> None: global _DIT assert _DIT is None, "DIT group is already initialized" _DIT = torch.distributed.new_group( ranks=list(range(dit_parallel_size)), backend=backend ) def get_dit_group() -> ProcessGroup: assert _DIT is not None, "DIT group is not initialized" return _DIT def init_vae_group( dit_parallel_size: int, vae_parallel_size: int, backend: str, ): # Initialize VAE group first global _VAE assert _VAE is None, "VAE parallel group is already initialized" vae_ranks = list(range(dit_parallel_size, dit_parallel_size + vae_parallel_size)) _VAE = torch.distributed.new_group(ranks=vae_ranks, backend=backend) def destroy_model_parallel() -> None: """Set the groups to none and destroy them.""" global _TP, _SP, _DP, _CFG, _PP, _DIT, _VAE for group in (_TP, _SP, _DP, _CFG, _PP): if group is not None: group.destroy() for group in (_DIT, _VAE): if group is not None: torch.distributed.destroy_process_group(group) _TP, _SP, _DP, _CFG, _PP, _DIT, _VAE = (None,) * 7 ================================================ FILE: python/sglang/multimodal_gen/runtime/distributed/utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/utils.py # Copyright 2023 The vLLM team. # Adapted from # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/utils.py # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. import dataclasses import pickle import time from collections import deque from collections.abc import Sequence from typing import Any import torch from torch.distributed import TCPStore from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) def ensure_divisibility(numerator, denominator) -> None: """Ensure that numerator is divisible by the denominator.""" assert numerator % denominator == 0, "{} is not divisible by {}".format( numerator, denominator ) def divide(numerator: int, denominator: int) -> int: """Ensure that numerator is divisible by the denominator and return the division value.""" ensure_divisibility(numerator, denominator) return numerator // denominator def split_tensor_along_last_dim( tensor: torch.Tensor, num_partitions: int, contiguous_split_chunks: bool = False, ) -> Sequence[torch.Tensor]: """Split a tensor along its last dimension. Arguments: tensor: input tensor. num_partitions: number of partitions to split the tensor contiguous_split_chunks: If True, make each chunk contiguous in memory. Returns: A list of Tensors """ # Get the size and dimension. last_dim = tensor.dim() - 1 last_dim_size = divide(tensor.size()[last_dim], num_partitions) # Split. tensor_list = torch.split(tensor, last_dim_size, dim=last_dim) # NOTE: torch.split does not create contiguous tensors by default. if contiguous_split_chunks: return tuple(chunk.contiguous() for chunk in tensor_list) return tuple(tensor_list) @dataclasses.dataclass class StatelessProcessGroup: """A dataclass to hold a metadata store, and the rank, world_size of the group. Only use it to communicate metadata between processes. For data-plane communication, create NCCL-related objects. """ rank: int world_size: int store: torch._C._distributed_c10d.Store data_expiration_seconds: int = 3600 # 1 hour # dst rank -> counter send_dst_counter: dict[int, int] = dataclasses.field(default_factory=dict) # src rank -> counter recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict) broadcast_send_counter: int = 0 broadcast_recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict) # A deque to store the data entries, with key and timestamp. entries: deque[tuple[str, float]] = dataclasses.field(default_factory=deque) def __post_init__(self): assert self.rank < self.world_size self.send_dst_counter = {i: 0 for i in range(self.world_size)} self.recv_src_counter = {i: 0 for i in range(self.world_size)} self.broadcast_recv_src_counter = {i: 0 for i in range(self.world_size)} def send_obj(self, obj: Any, dst: int): """Send an object to a destination rank.""" self.expire_data() key = f"send_to/{dst}/{self.send_dst_counter[dst]}" self.store.set(key, pickle.dumps(obj)) self.send_dst_counter[dst] += 1 self.entries.append((key, time.perf_counter())) def expire_data(self) -> None: """Expire data that is older than `data_expiration_seconds` seconds.""" while self.entries: # check the oldest entry key, timestamp = self.entries[0] if time.perf_counter() - timestamp > self.data_expiration_seconds: self.store.delete_key(key) self.entries.popleft() else: break def recv_obj(self, src: int) -> Any: """Receive an object from a source rank.""" obj = pickle.loads( self.store.get(f"send_to/{self.rank}/{self.recv_src_counter[src]}") ) self.recv_src_counter[src] += 1 return obj def broadcast_obj(self, obj: Any | None, src: int) -> Any: """Broadcast an object from a source rank to all other ranks. It does not clean up after all ranks have received the object. Use it for limited times, e.g., for initialization. """ if self.rank == src: self.expire_data() key = f"broadcast_from/{src}/" f"{self.broadcast_send_counter}" self.store.set(key, pickle.dumps(obj)) self.broadcast_send_counter += 1 self.entries.append((key, time.perf_counter())) return obj else: key = f"broadcast_from/{src}/" f"{self.broadcast_recv_src_counter[src]}" recv_obj = pickle.loads(self.store.get(key)) self.broadcast_recv_src_counter[src] += 1 return recv_obj def all_gather_obj(self, obj: Any) -> list[Any]: """All gather an object from all ranks.""" gathered_objs = [] for i in range(self.world_size): if i == self.rank: gathered_objs.append(obj) self.broadcast_obj(obj, src=self.rank) else: recv_obj = self.broadcast_obj(None, src=i) gathered_objs.append(recv_obj) return gathered_objs def barrier(self): """A barrier to synchronize all ranks.""" for i in range(self.world_size): if i == self.rank: self.broadcast_obj(None, src=self.rank) else: self.broadcast_obj(None, src=i) @staticmethod def create( host: str, port: int, rank: int, world_size: int, data_expiration_seconds: int = 3600, ) -> "StatelessProcessGroup": """A replacement for `torch.distributed.init_process_group` that does not pollute the global state. If we have process A and process B called `torch.distributed.init_process_group` to form a group, and then we want to form another group with process A, B, C, D, it is not possible in PyTorch, because process A and process B have already formed a group, and process C and process D cannot join that group. This function is a workaround for this issue. `torch.distributed.init_process_group` is a global call, while this function is a stateless call. It will return a `StatelessProcessGroup` object that can be used for exchanging metadata. With this function, process A and process B can call `StatelessProcessGroup.create` to form a group, and then process A, B, C, and D can call `StatelessProcessGroup.create` to form another group. """ # noqa store = TCPStore( host_name=host, port=port, world_size=world_size, is_master=(rank == 0), ) return StatelessProcessGroup( rank=rank, world_size=world_size, store=store, data_expiration_seconds=data_expiration_seconds, ) ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from sglang.multimodal_gen.runtime.utils.logging_utils import globally_suppress_loggers globally_suppress_loggers() ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/cli/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/cli/cli_types.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/types.py import argparse from sglang.multimodal_gen.utils import FlexibleArgumentParser class CLISubcommand: """Base class for CLI subcommands""" name: str def cmd( self, args: argparse.Namespace, unknown_args: list[str] | None = None ) -> None: """Execute the command with the given arguments""" raise NotImplementedError def validate(self, args: argparse.Namespace) -> None: """Validate the arguments for this command""" pass def subparser_init( self, subparsers: argparse._SubParsersAction ) -> FlexibleArgumentParser: """Initialize the subparser for this command""" raise NotImplementedError ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/serve.py import argparse import dataclasses import json import os from typing import cast from sglang.multimodal_gen import DiffGenerator from sglang.multimodal_gen.configs.sample.sampling_params import ( SamplingParams, generate_request_id, ) from sglang.multimodal_gen.runtime.entrypoints.cli.cli_types import CLISubcommand from sglang.multimodal_gen.runtime.entrypoints.cli.utils import ( RaiseNotImplementedAction, ) from sglang.multimodal_gen.runtime.entrypoints.utils import GenerationResult from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import ( MemorySnapshot, PerformanceLogger, RequestMetrics, ) from sglang.multimodal_gen.utils import FlexibleArgumentParser logger = init_logger(__name__) def add_multimodal_gen_generate_args(parser: argparse.ArgumentParser): """Add the arguments for the generate command.""" parser.add_argument( "--config", type=str, default="", required=False, help="Read CLI options from a config JSON or YAML file. If provided, --model-path and --prompt are optional.", ) parser.add_argument( "--perf-dump-path", type=str, default=None, required=False, help="Path to dump the performance metrics (JSON) for the run.", ) parser = ServerArgs.add_cli_args(parser) parser = SamplingParams.add_cli_args(parser) parser.add_argument( "--text-encoder-configs", action=RaiseNotImplementedAction, help="JSON array of text encoder configurations (NOT YET IMPLEMENTED)", ) return parser def maybe_dump_performance( args: argparse.Namespace, server_args, prompt: str, results: GenerationResult | list[GenerationResult] | None, ): """dump performance if necessary""" if not (args.perf_dump_path and results): return if isinstance(results, list): result = results[0] if results else None else: result = results metrics_dict = result.metrics if not (args.perf_dump_path and metrics_dict): return metrics = RequestMetrics(request_id=metrics_dict.get("request_id")) metrics.stages = metrics_dict.get("stages", {}) metrics.steps = metrics_dict.get("steps", []) metrics.total_duration_ms = metrics_dict.get("total_duration_ms", 0) # restore memory snapshots from serialized dict memory_snapshots_dict = metrics_dict.get("memory_snapshots", {}) for checkpoint_name, snapshot_dict in memory_snapshots_dict.items(): snapshot = MemorySnapshot( allocated_mb=snapshot_dict.get("allocated_mb", 0.0), reserved_mb=snapshot_dict.get("reserved_mb", 0.0), peak_allocated_mb=snapshot_dict.get("peak_allocated_mb", 0.0), peak_reserved_mb=snapshot_dict.get("peak_reserved_mb", 0.0), ) metrics.memory_snapshots[checkpoint_name] = snapshot PerformanceLogger.dump_benchmark_report( file_path=args.perf_dump_path, metrics=metrics, meta={ "prompt": prompt, "model": server_args.model_path, }, tag="cli_generate", ) def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None): """The entry point for the generate command.""" args.request_id = "mocked_fake_id_for_offline_generate" server_args = ServerArgs.from_cli_args(args, unknown_args) sampling_params_kwargs = SamplingParams.get_cli_args(args) sampling_params_kwargs["request_id"] = generate_request_id() # Handle diffusers-specific kwargs passed via CLI if hasattr(args, "diffusers_kwargs") and args.diffusers_kwargs: try: sampling_params_kwargs["diffusers_kwargs"] = json.loads( args.diffusers_kwargs ) logger.info( "Parsed diffusers_kwargs: %s", sampling_params_kwargs["diffusers_kwargs"], ) except json.JSONDecodeError as e: logger.error("Failed to parse --diffusers-kwargs as JSON: %s", e) raise ValueError( f"--diffusers-kwargs must be valid JSON. Got: {args.diffusers_kwargs}" ) from e generator = DiffGenerator.from_pretrained( model_path=server_args.model_path, server_args=server_args, local_mode=True ) results = generator.generate(sampling_params_kwargs=sampling_params_kwargs) prompt = sampling_params_kwargs.get("prompt") maybe_dump_performance(args, server_args, prompt, results) class GenerateSubcommand(CLISubcommand): """The `generate` subcommand for the sglang-diffusion CLI""" def __init__(self) -> None: self.name = "generate" super().__init__() self.init_arg_names = self._get_init_arg_names() self.generation_arg_names = self._get_generation_arg_names() def _get_init_arg_names(self) -> list[str]: """Get names of arguments for DiffGenerator initialization""" return ["num_gpus", "tp_size", "sp_size", "model_path"] def _get_generation_arg_names(self) -> list[str]: """Get names of arguments for generate_video method""" return [field.name for field in dataclasses.fields(SamplingParams)] def cmd( self, args: argparse.Namespace, unknown_args: list[str] | None = None ) -> None: generate_cmd(args, unknown_args) def validate(self, args: argparse.Namespace) -> None: """Validate the arguments for this command""" if args.num_gpus is not None and args.num_gpus <= 0: raise ValueError("Number of gpus must be positive") if args.config and not os.path.exists(args.config): raise ValueError(f"Config file not found: {args.config}") def subparser_init( self, subparsers: argparse._SubParsersAction ) -> FlexibleArgumentParser: generate_parser = subparsers.add_parser( "generate", help="Run inference on a model", usage="sglang generate (--model-path MODEL_PATH_OR_ID --prompt PROMPT) | --config CONFIG_FILE [OPTIONS]", ) generate_parser = add_multimodal_gen_generate_args(generate_parser) return cast(FlexibleArgumentParser, generate_parser) ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/cli/main.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/main.py from sglang.multimodal_gen.runtime.entrypoints.cli.cli_types import CLISubcommand from sglang.multimodal_gen.runtime.entrypoints.cli.generate import GenerateSubcommand from sglang.multimodal_gen.runtime.entrypoints.cli.serve import ServeSubcommand from sglang.multimodal_gen.utils import FlexibleArgumentParser def generate_cmd_init() -> list[CLISubcommand]: return [GenerateSubcommand(), ServeSubcommand()] def cmd_init() -> list[CLISubcommand]: """Initialize all commands from separate modules""" commands = [] commands.extend(generate_cmd_init()) return commands def main() -> None: parser = FlexibleArgumentParser(description="sglang-diffusion CLI") parser.add_argument("-v", "--version", action="version", version="0.1.0") subparsers = parser.add_subparsers(required=False, dest="subparser") cmds = {} for cmd in cmd_init(): cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd) cmds[cmd.name] = cmd args, unknown_args = parser.parse_known_args() if args.subparser in cmds: cmds[args.subparser].validate(args) if hasattr(args, "dispatch_function"): args.dispatch_function(args, unknown_args=unknown_args) else: parser.print_help() if __name__ == "__main__": main() ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/cli/serve.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import argparse import os from typing import cast from sglang.multimodal_gen.apps.webui import run_sgl_diffusion_webui from sglang.multimodal_gen.runtime.entrypoints.cli.cli_types import CLISubcommand from sglang.multimodal_gen.runtime.launch_server import launch_server from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import FlexibleArgumentParser logger = init_logger(__name__) def add_multimodal_gen_serve_args(parser: argparse.ArgumentParser): """Add the arguments for the serve command.""" parser.add_argument( "--config", type=str, default="", required=False, help="Read CLI options from a config JSON or YAML file.", ) return ServerArgs.add_cli_args(parser) def execute_serve_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None): """The entry point for the serve command.""" server_args = ServerArgs.from_cli_args(args, unknown_args) launch_server(server_args) if server_args.webui: run_sgl_diffusion_webui(server_args) class ServeSubcommand(CLISubcommand): """The `serve` subcommand for the sglang-diffusion CLI""" def __init__(self) -> None: self.name = "serve" super().__init__() def cmd( self, args: argparse.Namespace, unknown_args: list[str] | None = None ) -> None: execute_serve_cmd(args, unknown_args) def validate(self, args: argparse.Namespace) -> None: """Validate the arguments for this command""" if args.config and not os.path.exists(args.config): raise ValueError(f"Config file not found: {args.config}") def subparser_init( self, subparsers: argparse._SubParsersAction ) -> FlexibleArgumentParser: serve_parser = subparsers.add_parser( "serve", help="Launch the server and start FastAPI listener.", usage="sglang serve --model-path MODEL_PATH_OR_ID [OPTIONS]", ) serve_parser = add_multimodal_gen_serve_args(serve_parser) return cast(FlexibleArgumentParser, serve_parser) def cmd_init() -> list[CLISubcommand]: return [ServeSubcommand()] ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/cli/utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import argparse import os import shlex import subprocess import sys from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class RaiseNotImplementedAction(argparse.Action): def __call__(self, parser, namespace, values, option_string=None): raise NotImplementedError(f"The {option_string} option is not yet implemented") def launch_distributed( num_gpus: int, args: list[str], master_port: int | None = None ) -> int: """ Launch a distributed job with the given arguments Args: num_gpus: Number of GPUs to use args: Arguments to pass to v1_sgl_diffusion_inference.py (defaults to sys.argv[1:]) master_port: Port for the master process (default: random) """ current_env = os.environ.copy() python_executable = sys.executable project_root = os.path.abspath( os.path.join(os.path.dirname(__file__), "../../../..") ) main_script = os.path.join( project_root, "sgl_diffusion/sample/v1_sgl_diffusion_inference.py" ) cmd = [ python_executable, "-m", "torch.distributed.run", f"--nproc_per_node={num_gpus}", ] if master_port is not None: cmd.append(f"--master_port={master_port}") cmd.append(main_script) cmd.extend(args) logger.info("Running inference with %d GPU(s)", num_gpus) logger.info("Launching command: %s", shlex.join(cmd)) current_env["PYTHONIOENCODING"] = "utf-8" process = subprocess.Popen( cmd, env=current_env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True, bufsize=1, encoding="utf-8", errors="replace", ) if process.stdout: for line in iter(process.stdout.readline, ""): print(line.strip()) return process.wait() ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ DiffGenerator module for sglang-diffusion. This module provides a consolidated interface for generating images/videos using diffusion models. """ import dataclasses import multiprocessing as mp import os import time from typing import Any, List, Union from sglang.multimodal_gen.configs.sample.sampling_params import ( DataType, SamplingParams, ) from sglang.multimodal_gen.runtime.entrypoints.utils import ( GenerationResult, ListLorasReq, MergeLoraWeightsReq, SetLoraReq, ShutdownReq, UnmergeLoraWeightsReq, format_lora_message, prepare_request, save_outputs, ) from sglang.multimodal_gen.runtime.launch_server import launch_server from sglang.multimodal_gen.runtime.pipelines_core import Req from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.scheduler_client import sync_scheduler_client from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import ( GREEN, RESET, init_logger, log_batch_completion, log_generation_timer, ) logger = init_logger(__name__) # TODO: move to somewhere appropriate try: # Set the start method to 'spawn' to avoid CUDA errors in forked processes. # This must be done at the top level of the module, before any CUDA context # or other processes are initialized. mp.set_start_method("spawn", force=True) except RuntimeError: # The start method can only be set once per program execution. pass class DiffGenerator: """ A unified class for generating images/videos using diffusion models. This class provides a simple interface for image/video generation with rich customization options, similar to popular frameworks like HF Diffusers. """ def __init__( self, server_args: ServerArgs, ): """ Initialize the generator. Args: server_args: The inference arguments """ self.server_args = server_args self.port_args = PortArgs.from_server_args(server_args) # The executor is now a client to the Scheduler service self.local_scheduler_process: list[mp.Process] | None = None self.owns_scheduler_client: bool = False @classmethod def from_pretrained( cls, local_mode: bool = True, **kwargs, ) -> "DiffGenerator": """ Create a DiffGenerator from a pretrained model. Priority level: Default pipeline config < User's pipeline config < User's kwargs """ # If users also provide some kwargs, it will override the ServerArgs and PipelineConfig. if (server_args := kwargs.get("server_args", None)) is not None: if isinstance(server_args, ServerArgs): pass elif isinstance(server_args, dict): server_args = ServerArgs.from_kwargs(**server_args) else: server_args = ServerArgs.from_kwargs(**kwargs) return cls.from_server_args(server_args, local_mode=local_mode) @classmethod def from_server_args( cls, server_args: ServerArgs, local_mode: bool = True ) -> "DiffGenerator": """ Create a DiffGenerator with the specified arguments. Args: server_args: The inference arguments Returns: The created DiffGenerator """ instance = cls( server_args=server_args, ) logger.info(f"Local mode: {local_mode}") if local_mode: instance.local_scheduler_process = instance._start_local_server_if_needed() else: # In remote mode, we just need to connect and check. sync_scheduler_client.initialize(server_args) instance._check_remote_scheduler() # In both modes, this DiffGenerator instance is responsible for the client's lifecycle. instance.owns_scheduler_client = True return instance def _start_local_server_if_needed( self, ) -> list[mp.Process]: """Check if a local server is running; if not, start it and return the process handles.""" # First, we need a client to test the server. Initialize it temporarily. sync_scheduler_client.initialize(self.server_args) processes = launch_server(self.server_args, launch_http_server=False) return processes def _check_remote_scheduler(self): """Check if the remote scheduler is accessible.""" if not sync_scheduler_client.ping(): raise ConnectionError( f"Could not connect to remote scheduler at " f"{self.server_args.scheduler_endpoint} with `local mode` as False. " "Please ensure the server is running." ) logger.info( f"Successfully connected to remote scheduler at " f"{self.server_args.scheduler_endpoint}." ) @staticmethod def _resolve_image_paths_per_prompt( prompts: list[str], image_paths: str | list[str] | None ) -> list[str | list[str] | None]: if len(prompts) <= 1: return [image_paths] if not isinstance(image_paths, list) or len(image_paths) <= 1: return [image_paths for _ in prompts] if len(image_paths) != len(prompts): raise ValueError( "When using multiple prompts with multiple input images, " "provide either one shared image or exactly one image per prompt." ) return [[image_path] for image_path in image_paths] def generate( self, sampling_params_kwargs: dict | None = None, ) -> GenerationResult | list[GenerationResult] | None: """Generate image(s)/video(s) based on the given prompt(s). Returns a single GenerationResult for a single prompt, a list for multiple prompts, or None when every request failed. """ # 1. prepare requests prompts = self._resolve_prompts(sampling_params_kwargs.get("prompt")) user_output_file_name = sampling_params_kwargs.get("output_file_name") if len(prompts) > 1 and user_output_file_name is not None: raise ValueError( "Cannot use multiple prompts with a fixed output_file_name. " "Either remove --output-file-name or use a single prompt." ) sampling_params_orig = SamplingParams.from_user_sampling_params_args( self.server_args.model_path, server_args=self.server_args, **sampling_params_kwargs, ) requests: list[Req] = [] image_paths_per_prompt = self._resolve_image_paths_per_prompt( prompts, sampling_params_orig.image_path ) for i, p in enumerate(prompts): sampling_params = dataclasses.replace( sampling_params_orig, prompt=p, output_file_name=user_output_file_name, image_path=image_paths_per_prompt[i], ) sampling_params._set_output_file_name() req = prepare_request( server_args=self.server_args, sampling_params=sampling_params, ) requests.append(req) results: list[GenerationResult] = [] total_start_time = time.perf_counter() # 2. send requests to scheduler one at a time # TODO: send batch when supported for request_idx, req in enumerate(requests): try: with log_generation_timer( logger, req.prompt, request_idx + 1, len(requests) ) as timer: output_batch = self._send_to_scheduler_and_wait_for_response([req]) if output_batch.error: raise Exception(f"{output_batch.error}") if ( output_batch.output is None and output_batch.output_file_paths is None ): logger.error( "Received empty output from scheduler for prompt %d", request_idx + 1, ) continue common = dict( prompt=req.prompt, size=(req.height, req.width, req.num_frames), generation_time=timer.duration, peak_memory_mb=output_batch.peak_memory_mb, metrics=( output_batch.metrics.to_dict() if output_batch.metrics else {} ), trajectory_latents=output_batch.trajectory_latents, trajectory_timesteps=output_batch.trajectory_timesteps, trajectory_decoded=output_batch.trajectory_decoded, ) if req.save_output and req.return_file_paths_only: for idx, path in enumerate(output_batch.output_file_paths): results.append( GenerationResult( **common, prompt_index=idx, output_file_path=path, ) ) continue if req.data_type == DataType.MESH: for output_idx, sample in enumerate( output_batch.output_file_paths ): results.append( GenerationResult( **common, prompt_index=output_idx, output_file_path=sample, ) ) continue samples_out: list[Any] = [] audios_out: list[Any] = [] frames_out: list[Any] = [] num_outputs = len(output_batch.output) save_outputs( output_batch.output, req.data_type, req.fps, req.save_output, lambda idx: req.output_file_path(num_outputs, idx), audio=output_batch.audio, audio_sample_rate=output_batch.audio_sample_rate, samples_out=samples_out, audios_out=audios_out, frames_out=frames_out, output_compression=req.output_compression, enable_frame_interpolation=req.enable_frame_interpolation, frame_interpolation_exp=req.frame_interpolation_exp, frame_interpolation_scale=req.frame_interpolation_scale, frame_interpolation_model_path=req.frame_interpolation_model_path, enable_upscaling=req.enable_upscaling, upscaling_model_path=req.upscaling_model_path, upscaling_scale=req.upscaling_scale, ) for idx in range(len(samples_out)): results.append( GenerationResult( **common, samples=samples_out[idx], frames=frames_out[idx], audio=audios_out[idx], prompt_index=idx, output_file_path=req.output_file_path(num_outputs, idx), ) ) except Exception as e: logger.error( "Generation failed for prompt %d/%d: %s", request_idx + 1, len(requests), e, exc_info=True, ) continue total_gen_time = time.perf_counter() - total_start_time log_batch_completion(logger, len(results), total_gen_time) self._log_summary(results) if not results: return None return results[0] if len(results) == 1 else results def _resolve_prompts(self, prompt: str | list[str] | None) -> list[str]: """Collect prompts from the argument or from a prompt file.""" if self.server_args.prompt_file_path is not None: path = self.server_args.prompt_file_path if not os.path.exists(path): raise FileNotFoundError(f"Prompt text file not found: {path}") with open(path, encoding="utf-8") as f: prompts = [line.strip() for line in f if line.strip()] if not prompts: raise ValueError(f"No prompts found in file: {path}") logger.info("Found %d prompts in %s", len(prompts), path) return prompts if prompt is None: return [" "] if isinstance(prompt, str): return [prompt] return list(prompt) def _log_summary(self, results: list[GenerationResult]) -> None: if not results: return if self.server_args.warmup: total_duration_ms = results[0].metrics.get("total_duration_ms", 0) logger.info( f"Warmed-up request processed in {GREEN}%.2f{RESET} seconds (with warmup excluded)", total_duration_ms / 1000.0, ) peak_memories = [r.peak_memory_mb for r in results if r.peak_memory_mb] if peak_memories: logger.info( f"Memory usage - Max peak: {max(peak_memories):.2f} MB, " f"Avg peak: {sum(peak_memories) / len(peak_memories):.2f} MB" ) def _send_to_scheduler_and_wait_for_response(self, batch: list[Req]) -> OutputBatch: """ Sends a request to the scheduler and waits for a response. """ return sync_scheduler_client.forward(batch) # LoRA def _send_lora_request(self, req: Any, success_msg: str, failure_msg: str): response = sync_scheduler_client.forward(req) if response.error is None: logger.info(success_msg) return response else: error_msg = response.error raise RuntimeError(f"{failure_msg}: {error_msg}") def set_lora( self, lora_nickname: Union[str, List[str]], lora_path: Union[str, None, List[Union[str, None]]] = None, target: Union[str, List[str]] = "all", strength: Union[float, List[float]] = 1.0, ) -> None: """ Set LoRA adapter(s) for the specified transformer(s). Supports both single LoRA (backward compatible) and multiple LoRA adapters. Args: lora_nickname: The nickname(s) of the adapter(s). Can be a string or a list of strings. lora_path: Path(s) to the LoRA adapter(s). Can be a string, None, or a list of strings/None. target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings. Valid values: - "all": Apply to all transformers (default) - "transformer": Apply only to the primary transformer (high noise for Wan2.2) - "transformer_2": Apply only to transformer_2 (low noise for Wan2.2) - "critic": Apply only to the critic model strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats. """ req = SetLoraReq( lora_nickname=lora_nickname, lora_path=lora_path, target=target, strength=strength, ) nickname_str, target_str, strength_str = format_lora_message( lora_nickname, target, strength ) self._send_lora_request( req, f"Successfully set LoRA adapter(s): {nickname_str} (target: {target_str}, strength: {strength_str})", "Failed to set LoRA adapter", ) def unmerge_lora_weights(self, target: str = "all") -> None: """ Unmerge LoRA weights from the base model. Args: target: Which transformer(s) to unmerge. """ req = UnmergeLoraWeightsReq(target=target) self._send_lora_request( req, f"Successfully unmerged LoRA weights (target: {target})", "Failed to unmerge LoRA weights", ) def merge_lora_weights(self, target: str = "all", strength: float = 1.0) -> None: """ Merge LoRA weights into the base model. Args: target: Which transformer(s) to merge. strength: LoRA strength for merge, default 1.0. """ req = MergeLoraWeightsReq(target=target, strength=strength) self._send_lora_request( req, f"Successfully merged LoRA weights (target: {target}, strength: {strength})", "Failed to merge LoRA weights", ) def list_loras(self) -> dict: """List loaded LoRA adapters and current application status per module.""" output = self._send_lora_request( req=ListLorasReq(), success_msg="Successfully listed LoRA adapters", failure_msg="Failed to list LoRA adapters", ) # _send_lora_request already raises on error, so output.error is always None here return output.output or {} def _ensure_lora_state( self, lora_path: str | None, lora_nickname: str | None = None, merge_lora: bool = True, ) -> None: """ Ensure the LoRA state matches the desired configuration. Note: This method does not cache client-side state. The server handles idempotent operations, so redundant calls are safe but may have minor overhead. """ if lora_path is None: # Unmerge all LoRA weights when no lora_path is provided self.unmerge_lora_weights() return lora_nickname = lora_nickname or self.server_args.lora_nickname # Set the LoRA adapter (server handles idempotent logic) self.set_lora(lora_nickname, lora_path) # Merge or unmerge based on the merge_lora flag if merge_lora: self.merge_lora_weights() else: self.unmerge_lora_weights() def generate_with_lora( self, prompt: str | list[str] | None = None, sampling_params: SamplingParams | None = None, *, lora_path: str | None = None, lora_nickname: str | None = None, merge_lora: bool = True, **kwargs, ): self._ensure_lora_state( lora_path=lora_path, lora_nickname=lora_nickname, merge_lora=merge_lora ) return self.generate( sampling_params_kwargs=dict( prompt=prompt, sampling_params=sampling_params, **kwargs, ) ) def shutdown(self): """ Shutdown the generator. If in local mode, it also shuts down the scheduler server. """ # sends the shutdown command to the server if self.local_scheduler_process and self.owns_scheduler_client: try: sync_scheduler_client.forward(ShutdownReq()) except Exception: pass if self.local_scheduler_process: for process in self.local_scheduler_process: process.join(timeout=10) if process.is_alive(): logger.warning( f"Local worker {process.name} did not terminate gracefully, forcing." ) process.terminate() self.local_scheduler_process = None if self.owns_scheduler_client: sync_scheduler_client.close() self.owns_scheduler_client = False def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self.shutdown() def __del__(self): if self.owns_scheduler_client: logger.warning( "Generator was garbage collected without being shut down. " "Attempting to shut down the local server and client." ) self.shutdown() elif self.local_scheduler_process: logger.warning( "Generator was garbage collected without being shut down. " "Attempting to shut down the local server." ) self.shutdown() ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/http_server.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import asyncio import base64 import os import uuid from contextlib import asynccontextmanager from typing import TYPE_CHECKING import torch from fastapi import APIRouter, FastAPI, Request from fastapi.responses import ORJSONResponse from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.openai import image_api, video_api from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( VertexGenerateReqInput, ) from sglang.multimodal_gen.runtime.entrypoints.openai.utils import build_sampling_params from sglang.multimodal_gen.runtime.entrypoints.post_training import weights_api from sglang.multimodal_gen.runtime.entrypoints.utils import ( prepare_request, save_outputs, ) from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.version import __version__ if TYPE_CHECKING: from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req logger = init_logger(__name__) DEFAULT_SEED = 1024 VERTEX_ROUTE = os.environ.get("AIP_PREDICT_ROUTE", "/vertex_generate") @asynccontextmanager async def lifespan(app: FastAPI): from sglang.multimodal_gen.runtime.scheduler_client import ( async_scheduler_client, run_zeromq_broker, ) # 1. Initialize the singleton client that connects to the backend Scheduler server_args = app.state.server_args async_scheduler_client.initialize(server_args) # 2. Start the ZMQ Broker in the background to handle offline requests broker_task = asyncio.create_task(run_zeromq_broker(server_args)) yield # On shutdown logger.info("FastAPI app is shutting down...") broker_task.cancel() async_scheduler_client.close() # Health router health_router = APIRouter() @health_router.get("/health") async def health(): return {"status": "ok"} @health_router.get("/models", deprecated=True) async def get_models(request: Request): """ Get information about the model served by this server. .. deprecated:: Use /v1/models instead for OpenAI-compatible model discovery. This endpoint will be removed in a future version. """ from sglang.multimodal_gen.registry import get_model_info server_args: ServerArgs = request.app.state.server_args model_info = get_model_info(server_args.model_path, model_id=server_args.model_id) response = { "model_path": server_args.model_path, "num_gpus": server_args.num_gpus, "task_type": server_args.pipeline_config.task_type.name, "dit_precision": server_args.pipeline_config.dit_precision, "vae_precision": server_args.pipeline_config.vae_precision, } if model_info: response["pipeline_name"] = model_info.pipeline_cls.pipeline_name response["pipeline_class"] = model_info.pipeline_cls.__name__ return response @health_router.get("/server_info") async def server_info_endpoint(request: Request): """Get server information. Returns fields compatible with the LLM engine's /server_info so that the model gateway can discover diffusion workers. """ server_args: ServerArgs = request.app.state.server_args return { "model_path": server_args.model_path, "served_model_name": server_args.model_id or server_args.model_path, "tp_size": server_args.tp_size, "dp_size": server_args.dp_size, "version": __version__, } @health_router.get("/model_info") async def model_info_endpoint(request: Request): """Get model information. Returns fields compatible with the LLM engine's /model_info so that the model gateway can detect capabilities for diffusion workers. """ from sglang.multimodal_gen.registry import get_model_info server_args: ServerArgs = request.app.state.server_args task_type = server_args.pipeline_config.task_type try: registry_info = get_model_info( server_args.model_path, backend=server_args.backend, model_id=server_args.model_id, ) except Exception: logger.warning("Failed to resolve model info from registry", exc_info=True) registry_info = None return { # Fields consumed by the model gateway for worker discovery "model_path": server_args.model_path, "is_generation": True, "model_type": "diffusion", "architectures": ( [registry_info.pipeline_cls.__name__] if registry_info else None ), # Fields matching the LLM engine's /model_info shape "has_image_understanding": task_type.accepts_image_input(), "has_audio_understanding": False, # Diffusion-specific fields "task_type": task_type.name, "is_image_gen": task_type.is_image_gen(), } @health_router.get("/health_generate") async def health_generate(): # TODO : health generate endpoint return {"status": "ok"} def make_serializable(obj): """Recursively converts Tensors to None for JSON serialization.""" if isinstance(obj, torch.Tensor): return None if isinstance(obj, dict): return {k: make_serializable(v) for k, v in obj.items()} if isinstance(obj, list): return [make_serializable(v) for v in obj] return obj def encode_video_to_base64(file_path: str): if not os.path.exists(file_path): return None with open(file_path, "rb") as f: return base64.b64encode(f.read()).decode("utf-8") async def forward_to_scheduler( req_obj: "Req", sp: SamplingParams, ): """Forwards request to scheduler and processes the result.""" try: response = await async_scheduler_client.forward(req_obj) if response.output is None and response.output_file_paths is None: raise RuntimeError("Model generation returned no output.") if response.output_file_paths: output_file_path = response.output_file_paths[0] else: output_file_path = sp.output_file_path() save_outputs( [response.output[0]], sp.data_type, sp.fps, True, lambda _idx: output_file_path, audio=response.audio, audio_sample_rate=response.audio_sample_rate, enable_frame_interpolation=sp.enable_frame_interpolation, frame_interpolation_exp=sp.frame_interpolation_exp, frame_interpolation_scale=sp.frame_interpolation_scale, frame_interpolation_model_path=sp.frame_interpolation_model_path, enable_upscaling=sp.enable_upscaling, upscaling_model_path=sp.upscaling_model_path, upscaling_scale=sp.upscaling_scale, ) if hasattr(response, "model_dump"): data = response.model_dump() else: data = response if isinstance(response, dict) else vars(response) if output_file_path: logger.info("Processing output file: %s", output_file_path) b64_video = encode_video_to_base64(output_file_path) if b64_video: data["output"] = b64_video data.pop("video_data", None) data.pop("video_tensor", None) return make_serializable(data) except Exception as e: logger.error("Error during generation: %s", e, exc_info=True) return {"error": str(e)} vertex_router = APIRouter() @vertex_router.post(VERTEX_ROUTE) async def vertex_generate(vertex_req: VertexGenerateReqInput): if not vertex_req.instances: return ORJSONResponse({"predictions": []}) server_args = get_global_server_args() params = vertex_req.parameters or {} futures = [] for inst in vertex_req.instances: rid = f"vertex_{uuid.uuid4()}" sp = build_sampling_params( rid, prompt=inst.get("prompt") or inst.get("text"), image_path=inst.get("image") or inst.get("image_url"), seed=params.get("seed", DEFAULT_SEED), num_frames=params.get("num_frames"), fps=params.get("fps"), width=params.get("width"), height=params.get("height"), guidance_scale=params.get("guidance_scale"), save_output=params.get("save_output"), ) backend_req = prepare_request(server_args, sampling_params=sp) futures.append(forward_to_scheduler(backend_req, sp)) results = await asyncio.gather(*futures) return ORJSONResponse({"predictions": results}) def create_app(server_args: ServerArgs): """ Create and configure the FastAPI application instance. """ app = FastAPI(lifespan=lifespan) app.include_router(health_router) app.include_router(vertex_router) from sglang.multimodal_gen.runtime.entrypoints.openai import common_api, mesh_api app.include_router(common_api.router) app.include_router(image_api.router) app.include_router(video_api.router) app.include_router(mesh_api.router) app.include_router(weights_api.router) app.state.server_args = server_args return app ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py ================================================ import time from typing import Any, List, Optional, Union from fastapi import APIRouter, Body, HTTPException from fastapi.responses import ORJSONResponse from pydantic import BaseModel, Field from sglang.multimodal_gen.registry import get_model_info from sglang.multimodal_gen.runtime.entrypoints.utils import ( ListLorasReq, MergeLoraWeightsReq, SetLoraReq, UnmergeLoraWeightsReq, format_lora_message, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.multimodal_gen.runtime.server_args import get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger router = APIRouter(prefix="/v1") logger = init_logger(__name__) class ModelCard(BaseModel): """Model cards.""" id: str object: str = "model" created: int = Field(default_factory=lambda: int(time.time())) owned_by: str = "sglang" root: Optional[str] = None parent: Optional[str] = None max_model_len: Optional[int] = None class DiffusionModelCard(ModelCard): """Extended ModelCard with diffusion-specific fields.""" num_gpus: Optional[int] = None task_type: Optional[str] = None dit_precision: Optional[str] = None vae_precision: Optional[str] = None pipeline_name: Optional[str] = None pipeline_class: Optional[str] = None async def _handle_lora_request(req: Any, success_msg: str, failure_msg: str): try: output: OutputBatch = await async_scheduler_client.forward(req) if output.error is None: return {"status": "ok", "message": success_msg} else: error_msg = output.error raise HTTPException(status_code=500, detail=f"{failure_msg}: {error_msg}") except Exception as e: if isinstance(e, HTTPException): raise logger.error(f"Error during '{failure_msg}': {e}", exc_info=True) raise HTTPException(status_code=500, detail=str(e)) @router.post("/set_lora") async def set_lora( lora_nickname: Union[str, List[str]] = Body(..., embed=True), lora_path: Optional[Union[str, List[Optional[str]]]] = Body(None, embed=True), target: Union[str, List[str]] = Body("all", embed=True), strength: Union[float, List[float]] = Body(1.0, embed=True), ): """ Set LoRA adapter(s) for the specified transformer(s). Supports both single LoRA (backward compatible) and multiple LoRA adapters. Args: lora_nickname: The nickname(s) of the adapter(s). Can be a string or a list of strings. lora_path: Path(s) to the LoRA adapter(s) (local path or HF repo id). Can be a string, None, or a list of strings/None. Must match the length of lora_nickname. target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings. If a list, must match the length of lora_nickname. Valid values: - "all": Apply to all transformers (default) - "transformer": Apply only to the primary transformer (high noise for Wan2.2) - "transformer_2": Apply only to transformer_2 (low noise for Wan2.2) - "critic": Apply only to the critic model strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats. If a list, must match the length of lora_nickname. Values < 1.0 reduce the effect, values > 1.0 amplify the effect. """ req = SetLoraReq( lora_nickname=lora_nickname, lora_path=lora_path, target=target, strength=strength, ) nickname_str, target_str, strength_str = format_lora_message( lora_nickname, target, strength ) return await _handle_lora_request( req, f"Successfully set LoRA adapter(s): {nickname_str} (target: {target_str}, strength: {strength_str})", "Failed to set LoRA adapter", ) @router.post("/merge_lora_weights") async def merge_lora_weights( target: str = Body("all", embed=True), strength: float = Body(1.0, embed=True), ): """ Merge LoRA weights into the base model. Args: target: Which transformer(s) to merge. One of "all", "transformer", "transformer_2", "critic". strength: LoRA strength for merge, default 1.0. Values < 1.0 reduce the effect, values > 1.0 amplify the effect. """ req = MergeLoraWeightsReq(target=target, strength=strength) return await _handle_lora_request( req, f"Successfully merged LoRA weights (target: {target}, strength: {strength})", "Failed to merge LoRA weights", ) @router.post("/unmerge_lora_weights") async def unmerge_lora_weights( target: str = Body("all", embed=True), ): """ Unmerge LoRA weights from the base model. Args: target: Which transformer(s) to unmerge. One of "all", "transformer", "transformer_2", "critic". """ req = UnmergeLoraWeightsReq(target=target) return await _handle_lora_request( req, f"Successfully unmerged LoRA weights (target: {target})", "Failed to unmerge LoRA weights", ) @router.get("/model_info") async def model_info(): """Get the model information.""" server_args = get_global_server_args() if not server_args: raise HTTPException(status_code=500, detail="Server args not initialized") result = { "model_path": server_args.model_path, } return result @router.get("/list_loras") async def list_loras(): """List loaded LoRA adapters and current application status per module.""" try: req = ListLorasReq() output: OutputBatch = await async_scheduler_client.forward(req) if output.error is None: return output.output or {} else: raise HTTPException(status_code=500, detail=output.error) except Exception as e: if isinstance(e, HTTPException): raise logger.error(f"Error during 'list_loras': {e}", exc_info=True) raise HTTPException(status_code=500, detail=str(e)) @router.get("/models", response_class=ORJSONResponse) async def available_models(): """Show available models. OpenAI-compatible endpoint with extended diffusion info.""" server_args = get_global_server_args() if not server_args: raise HTTPException(status_code=500, detail="Server args not initialized") model_info = get_model_info( server_args.model_path, backend=server_args.backend, model_id=server_args.model_id, ) card_kwargs = { "id": server_args.model_path, "root": server_args.model_path, # Extended diffusion-specific fields "num_gpus": server_args.num_gpus, "task_type": server_args.pipeline_config.task_type.name, "dit_precision": server_args.pipeline_config.dit_precision, "vae_precision": server_args.pipeline_config.vae_precision, } if model_info: card_kwargs["pipeline_name"] = model_info.pipeline_cls.pipeline_name card_kwargs["pipeline_class"] = model_info.pipeline_cls.__name__ model_card = DiffusionModelCard(**card_kwargs) # Return dict directly to preserve extended fields (ModelList strips them) return {"object": "list", "data": [model_card.model_dump()]} @router.get("/models/{model:path}", response_class=ORJSONResponse) async def retrieve_model(model: str): """Retrieve a model instance. OpenAI-compatible endpoint with extended diffusion info.""" server_args = get_global_server_args() if not server_args: raise HTTPException(status_code=500, detail="Server args not initialized") if model != server_args.model_path: return ORJSONResponse( status_code=404, content={ "error": { "message": f"The model '{model}' does not exist", "type": "invalid_request_error", "param": "model", "code": "model_not_found", } }, ) model_info = get_model_info( server_args.model_path, backend=server_args.backend, model_id=server_args.model_id, ) card_kwargs = { "id": model, "root": model, "num_gpus": server_args.num_gpus, "task_type": server_args.pipeline_config.task_type.name, "dit_precision": server_args.pipeline_config.dit_precision, "vae_precision": server_args.pipeline_config.vae_precision, } if model_info: card_kwargs["pipeline_name"] = model_info.pipeline_cls.pipeline_name card_kwargs["pipeline_class"] = model_info.pipeline_cls.__name__ # Return dict to preserve extended fields return DiffusionModelCard(**card_kwargs).model_dump() ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import base64 import contextlib import os import time from typing import List, Optional from fastapi import APIRouter, File, Form, HTTPException, Path, Query, UploadFile from fastapi.responses import FileResponse from sglang.multimodal_gen.configs.sample.sampling_params import generate_request_id from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( ImageGenerationsRequest, ImageResponse, ImageResponseData, ) from sglang.multimodal_gen.runtime.entrypoints.openai.storage import cloud_storage from sglang.multimodal_gen.runtime.entrypoints.openai.stores import IMAGE_STORE from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( add_common_data_to_response, build_sampling_params, choose_output_image_ext, merge_image_input_list, process_generation_batch, save_image_to_path, temp_dir_if_disabled, ) from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.multimodal_gen.runtime.server_args import get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger router = APIRouter(prefix="/v1/images", tags=["images"]) logger = init_logger(__name__) def _read_b64_for_paths(paths: list[str]) -> list[str]: """Read and base64-encode each file. Must be called before cloud upload deletes them.""" result = [] for path in paths: with open(path, "rb") as f: result.append(base64.b64encode(f.read()).decode("utf-8")) return result def _build_image_response_kwargs( save_file_path_list: list[str], resp_format: str, prompt: str, request_id: str, result: OutputBatch, *, b64_list: list[str] | None = None, cloud_url: str | None = None, fallback_url: str | None = None, is_persistent: bool = True, ) -> dict: """Build ImageResponse data list. For b64_json: uses pre-read b64_list (call _read_b64_for_paths first). For url: uses cloud_url or fallback_url. file_path is omitted when is_persistent=False to avoid exposing stale temp paths. """ ret = None if resp_format == "b64_json": if not b64_list: raise ValueError("b64_list required for b64_json response_format") data = [ ImageResponseData( b64_json=b64, revised_prompt=prompt, file_path=os.path.abspath(path) if is_persistent else None, ) for b64, path in zip(b64_list, save_file_path_list) ] ret = {"data": data} elif resp_format == "url": url = cloud_url or fallback_url if not url: raise HTTPException( status_code=400, detail="response_format='url' requires cloud storage to be configured.", ) ret = { "data": [ ImageResponseData( url=url, revised_prompt=prompt, file_path=( os.path.abspath(save_file_path_list[0]) if is_persistent else None ), ) ], } else: raise HTTPException( status_code=400, detail=f"response_format={resp_format} is not supported" ) ret = add_common_data_to_response(ret, request_id=request_id, result=result) return ret @router.post("/generations", response_model=ImageResponse) async def generations( request: ImageGenerationsRequest, ): request_id = generate_request_id() server_args = get_global_server_args() ext = choose_output_image_ext(request.output_format, request.background) with temp_dir_if_disabled(server_args.output_path) as output_dir: sampling = build_sampling_params( request_id, prompt=request.prompt, size=request.size, width=request.width, height=request.height, num_outputs_per_prompt=max(1, min(int(request.n or 1), 10)), output_file_name=f"{request_id}.{ext}", output_path=output_dir, seed=request.seed, generator_device=request.generator_device, num_inference_steps=request.num_inference_steps, guidance_scale=request.guidance_scale, true_cfg_scale=request.true_cfg_scale, negative_prompt=request.negative_prompt, enable_teacache=request.enable_teacache, output_compression=request.output_compression, output_quality=request.output_quality, enable_upscaling=request.enable_upscaling, upscaling_model_path=request.upscaling_model_path, upscaling_scale=request.upscaling_scale, ) batch = prepare_request( server_args=server_args, sampling_params=sampling, ) # Add diffusers_kwargs if provided if request.diffusers_kwargs: batch.extra["diffusers_kwargs"] = request.diffusers_kwargs save_file_path_list, result = await process_generation_batch( async_scheduler_client, batch ) save_file_path = save_file_path_list[0] resp_format = (request.response_format or "b64_json").lower() # read b64 before cloud upload may delete the local file b64_list = ( _read_b64_for_paths(save_file_path_list) if resp_format == "b64_json" else None ) cloud_url = await cloud_storage.upload_and_cleanup(save_file_path) is_persistent = server_args.output_path is not None await IMAGE_STORE.upsert( request_id, { "id": request_id, "created_at": int(time.time()), "file_path": None if cloud_url or not is_persistent else save_file_path, "url": cloud_url, }, ) response_kwargs = _build_image_response_kwargs( save_file_path_list, resp_format, request.prompt, request_id, result, b64_list=b64_list, cloud_url=cloud_url, fallback_url=f"/v1/images/{request_id}/content" if is_persistent else None, is_persistent=is_persistent, ) return ImageResponse(**response_kwargs) @router.post("/edits", response_model=ImageResponse) async def edits( image: Optional[List[UploadFile]] = File(None), image_array: Optional[List[UploadFile]] = File(None, alias="image[]"), url: Optional[List[str]] = Form(None), url_array: Optional[List[str]] = Form(None, alias="url[]"), prompt: str = Form(...), mask: Optional[UploadFile] = File(None), model: Optional[str] = Form(None), n: Optional[int] = Form(1), response_format: Optional[str] = Form(None), size: Optional[str] = Form(None), output_format: Optional[str] = Form(None), background: Optional[str] = Form("auto"), seed: Optional[int] = Form(1024), generator_device: Optional[str] = Form("cuda"), user: Optional[str] = Form(None), negative_prompt: Optional[str] = Form(None), guidance_scale: Optional[float] = Form(None), true_cfg_scale: Optional[float] = Form(None), num_inference_steps: Optional[int] = Form(None), output_quality: Optional[str] = Form("default"), output_compression: Optional[int] = Form(None), enable_teacache: Optional[bool] = Form(False), enable_upscaling: Optional[bool] = Form(False), upscaling_model_path: Optional[str] = Form(None), upscaling_scale: Optional[int] = Form(4), num_frames: int = Form(1), ): request_id = generate_request_id() server_args = get_global_server_args() # Resolve images from either `image` or `image[]` (OpenAI SDK sends `image[]` when list is provided) images = image or image_array urls = url or url_array if (not images or len(images) == 0) and (not urls or len(urls) == 0): raise HTTPException( status_code=422, detail="Field 'image' or 'url' is required" ) image_list = merge_image_input_list(images, urls) with contextlib.ExitStack() as stack: uploads_dir = stack.enter_context( temp_dir_if_disabled(server_args.input_save_path) ) output_dir = stack.enter_context(temp_dir_if_disabled(server_args.output_path)) input_paths = [] try: for idx, img in enumerate(image_list): filename = img.filename if hasattr(img, "filename") else f"image_{idx}" input_path = await save_image_to_path( img, os.path.join(uploads_dir, f"{request_id}_{idx}_{filename}"), ) input_paths.append(input_path) except Exception as e: raise HTTPException( status_code=400, detail=f"Failed to process image source: {str(e)}", ) ext = choose_output_image_ext(output_format, background) sampling = build_sampling_params( request_id, prompt=prompt, size=size, num_outputs_per_prompt=max(1, min(int(n or 1), 10)), output_file_name=f"{request_id}.{ext}", output_path=output_dir, image_path=input_paths, seed=seed, generator_device=generator_device, negative_prompt=negative_prompt, guidance_scale=guidance_scale, true_cfg_scale=true_cfg_scale, num_inference_steps=num_inference_steps, enable_teacache=enable_teacache, num_frames=num_frames, output_compression=output_compression, output_quality=output_quality, enable_upscaling=enable_upscaling, upscaling_model_path=upscaling_model_path, upscaling_scale=upscaling_scale, ) batch = prepare_request( server_args=server_args, sampling_params=sampling, ) save_file_path_list, result = await process_generation_batch( async_scheduler_client, batch ) save_file_path = save_file_path_list[0] resp_format = (response_format or "b64_json").lower() # read b64 before cloud upload may delete the local file b64_list = ( _read_b64_for_paths(save_file_path_list) if resp_format == "b64_json" else None ) cloud_url = await cloud_storage.upload_and_cleanup(save_file_path) is_persistent = server_args.output_path is not None is_input_persistent = server_args.input_save_path is not None await IMAGE_STORE.upsert( request_id, { "id": request_id, "created_at": int(time.time()), "file_path": None if cloud_url or not is_persistent else save_file_path, "url": cloud_url, "input_image_paths": input_paths if is_input_persistent else None, "num_input_images": len(input_paths), }, ) response_kwargs = _build_image_response_kwargs( save_file_path_list, resp_format, prompt, request_id, result, b64_list=b64_list, cloud_url=cloud_url, fallback_url=f"/v1/images/{request_id}/content" if is_persistent else None, is_persistent=is_persistent, ) return ImageResponse(**response_kwargs) @router.get("/{image_id}/content") async def download_image_content( image_id: str = Path(...), variant: Optional[str] = Query(None) ): item = await IMAGE_STORE.get(image_id) if not item: raise HTTPException(status_code=404, detail="Image not found") if item.get("url"): raise HTTPException( status_code=400, detail=f"Image has been uploaded to cloud storage. Please use the cloud URL: {item.get('url')}", ) file_path = item.get("file_path") if not file_path: raise HTTPException( status_code=404, detail="Image was not persisted on disk (output_path is disabled). Use b64_json response_format or configure cloud storage.", ) if not os.path.exists(file_path): raise HTTPException(status_code=404, detail="Image is still being generated") ext = os.path.splitext(file_path)[1].lower() media_type = "image/jpeg" if ext == ".png": media_type = "image/png" elif ext == ".webp": media_type = "image/webp" return FileResponse( path=file_path, media_type=media_type, filename=os.path.basename(file_path) ) ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/openai/mesh_api.py ================================================ import asyncio import os import time from typing import Any, Dict, List, Optional from fastapi import ( APIRouter, File, Form, HTTPException, Path, Query, Request, UploadFile, ) from fastapi.responses import FileResponse from sglang.multimodal_gen.configs.sample.sampling_params import ( SamplingParams, generate_request_id, ) from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( MeshGenerationsRequest, MeshListResponse, MeshResponse, ) from sglang.multimodal_gen.runtime.entrypoints.openai.storage import cloud_storage from sglang.multimodal_gen.runtime.entrypoints.openai.stores import MESH_STORE from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( add_common_data_to_response, merge_image_input_list, process_generation_batch, save_image_to_path, ) from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.server_args import get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) router = APIRouter(prefix="/v1/meshes", tags=["meshes"]) def _normalize_format(fmt: Optional[str]) -> str: fmt = (fmt or "glb").lower() return fmt if fmt in ("glb", "obj") else "glb" def _build_sampling_params_from_request( request_id: str, req: MeshGenerationsRequest, image_path: Optional[str] = None ) -> SamplingParams: ext = _normalize_format(req.output_format) server_args = get_global_server_args() sampling_kwargs: Dict[str, Any] = { "request_id": request_id, "prompt": req.prompt, "num_frames": 1, "image_path": [image_path] if image_path else None, "save_output": True, "output_file_name": f"{request_id}.{ext}", "seed": req.seed, "generator_device": req.generator_device, } if req.num_inference_steps is not None: sampling_kwargs["num_inference_steps"] = req.num_inference_steps if req.guidance_scale is not None: sampling_kwargs["guidance_scale"] = req.guidance_scale if req.negative_prompt is not None: sampling_kwargs["negative_prompt"] = req.negative_prompt return SamplingParams.from_user_sampling_params_args( model_path=server_args.model_path, server_args=server_args, **sampling_kwargs, ) def _mesh_job_from_sampling( request_id: str, req: MeshGenerationsRequest, sampling: SamplingParams ) -> Dict[str, Any]: return { "id": request_id, "object": "mesh", "model": req.model or "", "status": "queued", "progress": 0, "created_at": int(time.time()), "format": _normalize_format(req.output_format), "file_path": os.path.abspath(sampling.output_file_path()), } async def _dispatch_job_async(job_id: str, batch: Req) -> None: from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client try: save_file_path_list, result = await process_generation_batch( async_scheduler_client, batch ) save_file_path = save_file_path_list[0] file_size = None if os.path.exists(save_file_path): file_size = os.path.getsize(save_file_path) cloud_url = await cloud_storage.upload_and_cleanup(save_file_path) update_fields: Dict[str, Any] = { "status": "completed", "progress": 100, "completed_at": int(time.time()), "url": cloud_url, "file_path": save_file_path if not cloud_url else None, "file_size_bytes": file_size, } update_fields = add_common_data_to_response( update_fields, request_id=job_id, result=result ) await MESH_STORE.update_fields(job_id, update_fields) except Exception as e: logger.error(f"{e}") await MESH_STORE.update_fields( job_id, {"status": "failed", "error": {"message": str(e)}} ) @router.post("", response_model=MeshResponse) async def create_mesh( request: Request, image: Optional[List[UploadFile]] = File(None), image_array: Optional[List[UploadFile]] = File(None, alias="image[]"), url: Optional[List[str]] = Form(None), url_array: Optional[List[str]] = Form(None, alias="url[]"), prompt: Optional[str] = Form("generate 3d mesh"), model: Optional[str] = Form(None), seed: Optional[int] = Form(None), generator_device: Optional[str] = Form("cuda"), guidance_scale: Optional[float] = Form(None), num_inference_steps: Optional[int] = Form(None), negative_prompt: Optional[str] = Form(None), output_format: Optional[str] = Form("glb"), ): content_type = request.headers.get("content-type", "").lower() request_id = generate_request_id() server_args = get_global_server_args() input_path = None if "multipart/form-data" in content_type: images = image or image_array urls = url or url_array image_list = merge_image_input_list(images, urls) if not image_list: raise HTTPException( status_code=422, detail="Field 'image' or 'url' is required for mesh generation", ) uploads_dir = os.path.join("outputs", "uploads") os.makedirs(uploads_dir, exist_ok=True) img = image_list[0] filename = img.filename if hasattr(img, "filename") else "input_image" try: input_path = await save_image_to_path( img, os.path.join(uploads_dir, f"{request_id}_{filename}") ) except Exception as e: raise HTTPException( status_code=400, detail=f"Failed to process image source: {str(e)}" ) req = MeshGenerationsRequest( prompt=prompt or "generate 3d mesh", model=model, seed=seed, generator_device=generator_device, num_inference_steps=num_inference_steps, negative_prompt=negative_prompt, output_format=output_format, **( {"guidance_scale": guidance_scale} if guidance_scale is not None else {} ), ) else: try: body = await request.json() except Exception: body = {} try: payload: Dict[str, Any] = dict(body or {}) if payload.get("input_image"): img_src = payload.pop("input_image") uploads_dir = os.path.join("outputs", "uploads") os.makedirs(uploads_dir, exist_ok=True) input_path = await save_image_to_path( img_src, os.path.join(uploads_dir, f"{request_id}_input_image"), ) req = MeshGenerationsRequest(**payload) except Exception as e: raise HTTPException(status_code=400, detail=f"Invalid request body: {e}") if not input_path: raise HTTPException( status_code=422, detail="An input image is required for mesh generation", ) sampling_params = _build_sampling_params_from_request(request_id, req, input_path) job = _mesh_job_from_sampling(request_id, req, sampling_params) await MESH_STORE.upsert(request_id, job) batch = prepare_request( server_args=server_args, sampling_params=sampling_params, ) asyncio.create_task(_dispatch_job_async(request_id, batch)) return MeshResponse(**job) @router.get("", response_model=MeshListResponse) async def list_meshes( after: Optional[str] = Query(None), limit: Optional[int] = Query(None, ge=1, le=100), order: Optional[str] = Query("desc"), ): order = (order or "desc").lower() if order not in ("asc", "desc"): order = "desc" jobs = await MESH_STORE.list_values() reverse = order != "asc" jobs.sort(key=lambda j: j.get("created_at", 0), reverse=reverse) if after is not None: try: idx = next(i for i, j in enumerate(jobs) if j["id"] == after) jobs = jobs[idx + 1 :] except StopIteration: jobs = [] if limit is not None: jobs = jobs[:limit] items = [MeshResponse(**j) for j in jobs] return MeshListResponse(data=items) @router.get("/{mesh_id}", response_model=MeshResponse) async def retrieve_mesh(mesh_id: str = Path(...)): job = await MESH_STORE.get(mesh_id) if not job: raise HTTPException(status_code=404, detail="Mesh not found") return MeshResponse(**job) @router.delete("/{mesh_id}", response_model=MeshResponse) async def delete_mesh(mesh_id: str = Path(...)): job = await MESH_STORE.pop(mesh_id) if not job: raise HTTPException(status_code=404, detail="Mesh not found") job["status"] = "deleted" return MeshResponse(**job) @router.get("/{mesh_id}/content") async def download_mesh_content( mesh_id: str = Path(...), variant: Optional[str] = Query(None) ): job = await MESH_STORE.get(mesh_id) if not job: raise HTTPException(status_code=404, detail="Mesh not found") if job.get("url"): raise HTTPException( status_code=400, detail=f"Mesh has been uploaded to cloud storage. Please use the cloud URL: {job.get('url')}", ) file_path = job.get("file_path") if not file_path or not os.path.exists(file_path): raise HTTPException(status_code=404, detail="Generation is still in-progress") ext = os.path.splitext(file_path)[1].lower() media_type = { ".glb": "model/gltf-binary", ".obj": "text/plain", }.get(ext, "application/octet-stream") return FileResponse( path=file_path, media_type=media_type, filename=os.path.basename(file_path) ) ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py ================================================ import time import uuid from abc import ABC from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Union from pydantic import BaseModel, Field # Image API protocol models class ImageResponseData(BaseModel): b64_json: Optional[str] = None url: Optional[str] = None revised_prompt: Optional[str] = None file_path: Optional[str] = None class ImageResponse(BaseModel): id: str created: int = Field(default_factory=lambda: int(time.time())) data: List[ImageResponseData] peak_memory_mb: Optional[float] = None inference_time_s: Optional[float] = None class ImageGenerationsRequest(BaseModel): prompt: str model: Optional[str] = None n: Optional[int] = 1 quality: Optional[str] = "auto" response_format: Optional[str] = "url" # url | b64_json size: Optional[str] = "1024x1024" # e.g., 1024x1024 style: Optional[str] = "vivid" background: Optional[str] = "auto" # transparent | opaque | auto output_format: Optional[str] = None # png | jpeg | webp user: Optional[str] = None # SGLang extensions width: Optional[int] = None height: Optional[int] = None num_inference_steps: Optional[int] = None guidance_scale: Optional[float] = None true_cfg_scale: Optional[float] = ( None # for CFG vs guidance distillation (e.g., QwenImage) ) seed: Optional[int] = 1024 generator_device: Optional[str] = "cuda" negative_prompt: Optional[str] = None output_quality: Optional[str] = "default" output_compression: Optional[int] = None enable_teacache: Optional[bool] = False # Upscaling enable_upscaling: Optional[bool] = False upscaling_model_path: Optional[str] = None upscaling_scale: Optional[int] = 4 diffusers_kwargs: Optional[Dict[str, Any]] = None # kwargs for diffusers backend # Video API protocol models class VideoResponse(BaseModel): id: str object: str = "video" model: str = "sora-2" status: str = "queued" progress: int = 0 created_at: int = Field(default_factory=lambda: int(time.time())) size: str = "" seconds: str = "4" quality: str = "standard" url: Optional[str] = None remixed_from_video_id: Optional[str] = None completed_at: Optional[int] = None expires_at: Optional[int] = None error: Optional[Dict[str, Any]] = None file_path: Optional[str] = None peak_memory_mb: Optional[float] = None inference_time_s: Optional[float] = None class VideoGenerationsRequest(BaseModel): prompt: str input_reference: Optional[str] = None reference_url: Optional[str] = None model: Optional[str] = None seconds: Optional[int] = 4 size: Optional[str] = "" fps: Optional[int] = None num_frames: Optional[int] = None seed: Optional[int] = 1024 generator_device: Optional[str] = "cuda" # SGLang extensions num_inference_steps: Optional[int] = None guidance_scale: Optional[float] = None guidance_scale_2: Optional[float] = None true_cfg_scale: Optional[float] = ( None # for CFG vs guidance distillation (e.g., QwenImage) ) negative_prompt: Optional[str] = None enable_teacache: Optional[bool] = False # Frame interpolation enable_frame_interpolation: Optional[bool] = False frame_interpolation_exp: Optional[int] = 1 # 1=2×, 2=4× frame_interpolation_scale: Optional[float] = 1.0 frame_interpolation_model_path: Optional[str] = None # Upscaling enable_upscaling: Optional[bool] = False upscaling_model_path: Optional[str] = None upscaling_scale: Optional[int] = 4 output_quality: Optional[str] = "default" output_compression: Optional[int] = None output_path: Optional[str] = None diffusers_kwargs: Optional[Dict[str, Any]] = None # kwargs for diffusers backend class VideoListResponse(BaseModel): data: List[VideoResponse] object: str = "list" class VideoRemixRequest(BaseModel): prompt: str # Mesh API protocol models class MeshResponse(BaseModel): id: str object: str = "mesh" model: str = "" status: str = "queued" progress: int = 0 created_at: int = Field(default_factory=lambda: int(time.time())) format: str = "glb" url: Optional[str] = None completed_at: Optional[int] = None expires_at: Optional[int] = None error: Optional[Dict[str, Any]] = None file_path: Optional[str] = None file_size_bytes: Optional[int] = None peak_memory_mb: Optional[float] = None inference_time_s: Optional[float] = None class MeshGenerationsRequest(BaseModel): prompt: str = "generate 3d mesh" input_image: Optional[str] = None model: Optional[str] = None seed: Optional[int] = None generator_device: Optional[str] = "cuda" num_inference_steps: Optional[int] = None guidance_scale: Optional[float] = None negative_prompt: Optional[str] = None output_format: Optional[str] = "glb" class MeshListResponse(BaseModel): data: List[MeshResponse] object: str = "list" @dataclass class BaseReq(ABC): rid: Optional[Union[str, List[str]]] = field(default=None, kw_only=True) http_worker_ipc: Optional[str] = field(default=None, kw_only=True) def regenerate_rid(self): """Generate a new request ID and return it.""" if isinstance(self.rid, list): self.rid = [uuid.uuid4().hex for _ in range(len(self.rid))] else: self.rid = uuid.uuid4().hex return self.rid @dataclass class VertexGenerateReqInput(BaseReq): instances: List[dict] parameters: Optional[dict] = None ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/openai/storage.py ================================================ import asyncio import os from typing import Optional from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class CloudStorage: def __init__(self): self.enabled = os.getenv("SGLANG_CLOUD_STORAGE_TYPE", "").lower() == "s3" if not self.enabled: return try: import boto3 except ImportError: logger.error( "boto3 is not installed. Please install it with `pip install boto3` to use cloud storage." ) self.enabled = False return self.bucket_name = os.getenv("SGLANG_S3_BUCKET_NAME") if not self.bucket_name: self.enabled = False return endpoint_url = os.getenv("SGLANG_S3_ENDPOINT_URL") or None region_name = os.getenv("SGLANG_S3_REGION_NAME") or None self.client = boto3.client( "s3", aws_access_key_id=os.getenv("SGLANG_S3_ACCESS_KEY_ID"), aws_secret_access_key=os.getenv("SGLANG_S3_SECRET_ACCESS_KEY"), endpoint_url=endpoint_url, region_name=region_name, ) self.endpoint_url = endpoint_url self.region_name = region_name def is_enabled(self) -> bool: return self.enabled async def upload_file(self, local_path: str, destination_key: str) -> Optional[str]: if not self.is_enabled(): return None def _sync_upload(): """Synchronous part of the upload to run in a thread.""" ext = os.path.splitext(local_path)[1].lower() content_type = { ".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp", ".mp4": "video/mp4", ".glb": "model/gltf-binary", ".obj": "text/plain", }.get(ext, "application/octet-stream") # Use the client created once in __init__ self.client.upload_file( local_path, self.bucket_name, destination_key, ExtraArgs={"ContentType": content_type}, ) try: # Offload the blocking I/O call to a thread executor await asyncio.get_running_loop().run_in_executor(None, _sync_upload) except Exception as e: # If upload fails, log the error and return None for fallback logger.error(f"Upload failed for {destination_key}: {e}") return None # Simplified URL generation with a default region if self.endpoint_url: url = ( f"{self.endpoint_url.rstrip('/')}/{self.bucket_name}/{destination_key}" ) else: region = self.region_name or "us-east-1" url = f"https://{self.bucket_name}.s3.{region}.amazonaws.com/{destination_key}" logger.info(f"Uploaded {local_path} to {url}") return url async def upload_and_cleanup(self, file_path: str) -> Optional[str]: """Helper to upload a file and delete the local copy if successful.""" if not self.is_enabled(): return None key = os.path.basename(file_path) url = await self.upload_file(file_path, key) if url: try: # pass if removal fails os.remove(file_path) except OSError as e: logger.warning(f"Failed to remove temporary file {file_path}: {e}") return url # Global instance cloud_storage = CloudStorage() ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/openai/stores.py ================================================ import asyncio from typing import Any, Dict, List, Optional class AsyncDictStore: """A small async-safe in-memory key-value store for dict items. This encapsulates the usual pattern of a module-level dict guarded by an asyncio.Lock and provides simple CRUD methods that are safe to call concurrently from FastAPI request handlers and background tasks. """ def __init__(self) -> None: self._items: Dict[str, Dict[str, Any]] = {} self._lock = asyncio.Lock() async def upsert(self, key: str, value: Dict[str, Any]) -> None: async with self._lock: self._items[key] = value async def update_fields( self, key: str, updates: Dict[str, Any] ) -> Optional[Dict[str, Any]]: async with self._lock: item = self._items.get(key) if item is None: return None item.update(updates) return item async def get(self, key: str) -> Optional[Dict[str, Any]]: async with self._lock: return self._items.get(key) async def pop(self, key: str) -> Optional[Dict[str, Any]]: async with self._lock: return self._items.pop(key, None) async def list_values(self) -> List[Dict[str, Any]]: async with self._lock: return list(self._items.values()) # Global stores shared by OpenAI entrypoints # [request_id, dict] VIDEO_STORE = AsyncDictStore() IMAGE_STORE = AsyncDictStore() MESH_STORE = AsyncDictStore() ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import base64 import os import re import shutil import tempfile import time from contextlib import contextmanager from typing import Any, Generator, List, Optional, Union import httpx from fastapi import UploadFile from sglang.multimodal_gen.configs.sample.sampling_params import ( DataType, SamplingParams, ) from sglang.multimodal_gen.runtime.entrypoints.utils import ( ListLorasReq, MergeLoraWeightsReq, SetLoraReq, ShutdownReq, UnmergeLoraWeightsReq, format_lora_message, save_outputs, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.scheduler_client import AsyncSchedulerClient from sglang.multimodal_gen.runtime.server_args import get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import ( init_logger, log_batch_completion, log_generation_timer, ) # re-export LoRA protocol types for backward compatibility __all__ = [ "SetLoraReq", "MergeLoraWeightsReq", "UnmergeLoraWeightsReq", "ListLorasReq", "ShutdownReq", "format_lora_message", ] logger = init_logger(__name__) OUTPUT_QUALITY_MAPPER = {"maximum": 100, "high": 90, "medium": 55, "low": 35} DEFAULT_FPS = 24 DEFAULT_VIDEO_SECONDS = 4 @contextmanager def temp_dir_if_disabled( configured_path: str | None, ) -> Generator[str, None, None]: """Yield *configured_path* when it is set, otherwise create a temporary directory that is automatically removed when the context exits.""" if configured_path is not None: os.makedirs(configured_path, exist_ok=True) yield configured_path else: tmp = tempfile.mkdtemp(prefix="sglang_") try: yield tmp finally: shutil.rmtree(tmp, ignore_errors=True) def _parse_size(size: str) -> tuple[int, int] | tuple[None, None]: try: parts = size.lower().replace(" ", "").split("x") if len(parts) != 2: raise ValueError w, h = int(parts[0]), int(parts[1]) return w, h except Exception: return None, None def choose_output_image_ext( output_format: Optional[str], background: Optional[str] ) -> str: fmt = (output_format or "").lower() if fmt in {"png", "webp", "jpeg", "jpg"}: return "jpg" if fmt == "jpeg" else fmt if (background or "auto").lower() == "transparent": return "png" return "jpg" def build_sampling_params(request_id: str, **kwargs) -> SamplingParams: """Build SamplingParams from request parameters. Handles size parsing, output_quality resolution, and None filtering before delegating to SamplingParams.from_user_sampling_params_args. Callers pass only the parameters they have; None values are stripped automatically so that SamplingParams defaults apply. """ server_args = get_global_server_args() # pop HTTP-layer params that aren't SamplingParams fields output_quality = kwargs.pop("output_quality", None) has_explicit_compression = kwargs.get("output_compression") is not None # parse "WxH" size string if provided size = kwargs.pop("size", None) if size: w, h = _parse_size(size) if w is not None: # treat None dimensions as unset so parsed size can fill them if kwargs.get("width") is None: kwargs["width"] = w if kwargs.get("height") is None: kwargs["height"] = h # filter out None values to let SamplingParams defaults apply kwargs = {k: v for k, v in kwargs.items() if v is not None} kwargs.setdefault("save_output", True) sampling_params = SamplingParams.from_user_sampling_params_args( model_path=server_args.model_path, server_args=server_args, request_id=request_id, **kwargs, ) # resolve output_quality → output_compression with the correct data_type. # SamplingParams.__post_init__ may have resolved with the wrong data_type # (default VIDEO) before _adjust() set the correct one. if not has_explicit_compression and output_quality is not None: resolved = adjust_output_quality(output_quality, sampling_params.data_type) if resolved is not None: sampling_params.output_compression = resolved return sampling_params async def save_image_to_path(image: Union[UploadFile, str], target_path: str) -> str: input_path = await _maybe_url_image(image, target_path) if input_path is None: input_path = await _save_upload_to_path(image, target_path) return input_path # Helpers async def _save_upload_to_path(upload: UploadFile, target_path: str) -> str: os.makedirs(os.path.dirname(target_path), exist_ok=True) content = await upload.read() with open(target_path, "wb") as f: f.write(content) return target_path async def _maybe_url_image(img_url: str, target_path: str) -> str | None: if not isinstance(img_url, str): return None if img_url.lower().startswith(("http://", "https://")): # Download image from URL input_path = await _save_url_image_to_path(img_url, target_path) return input_path elif img_url.startswith("data:image"): # encode image base64 url input_path = await _save_base64_image_to_path(img_url, target_path) return input_path else: raise ValueError("Unsupported image url format") async def _save_url_image_to_path(image_url: str, target_path: str) -> str: """Download image from URL and save to target path.""" os.makedirs(os.path.dirname(target_path), exist_ok=True) try: async with httpx.AsyncClient(follow_redirects=True) as client: response = await client.get(image_url, timeout=10.0) response.raise_for_status() # Determine file extension from content type or URL after downloading if not os.path.splitext(target_path)[1]: content_type = response.headers.get("content-type", "").lower() url_path = image_url.split("?")[0] _, url_ext = os.path.splitext(url_path) url_ext = url_ext.lower() if url_ext in {".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp"}: ext = ".jpg" if url_ext == ".jpeg" else url_ext elif content_type.startswith("image/"): if "jpeg" in content_type or "jpg" in content_type: ext = ".jpg" elif "png" in content_type: ext = ".png" elif "webp" in content_type: ext = ".webp" else: ext = ".jpg" # Default to jpg elif content_type == "application/octet-stream": # for octet-stream, if we couldn't get it from URL, default to jpg ext = ".jpg" else: raise ValueError( f"URL does not point to an image. Content-Type: {content_type}" ) target_path = f"{target_path}{ext}" with open(target_path, "wb") as f: f.write(response.content) return target_path except Exception as e: raise Exception(f"Failed to download image from URL: {str(e)}") async def _save_base64_image_to_path(base64_data: str, target_path: str) -> str: """Decode base64 image data and save to target path.""" _B64_FMT_HINT = ( "Failed to decode base64 image. " "Expected format: `data:[];base64,`" ) # split `data:[][;base64],` to media-type base64 data pattern = r"data:(.*?)(;base64)?,(.*)" match = re.match(pattern, base64_data) if not match: raise ValueError(_B64_FMT_HINT) media_type = match.group(1) is_base64 = match.group(2) if not is_base64: raise ValueError(f"{_B64_FMT_HINT} (missing ;base64 marker)") data = match.group(3) if not data: raise ValueError(f"{_B64_FMT_HINT} (empty data payload)") # get ext from url if media_type.startswith("image/"): ext = media_type.split("/")[-1].lower() if ext == "jpeg": ext = "jpg" else: ext = "jpg" target_path = f"{target_path}.{ext}" os.makedirs(os.path.dirname(target_path), exist_ok=True) try: image_data = base64.b64decode(data) with open(target_path, "wb") as f: f.write(image_data) return target_path except Exception as e: raise Exception(f"Failed to decode base64 image: {str(e)}") async def process_generation_batch( scheduler_client: AsyncSchedulerClient, batch, ) -> tuple[list[str], OutputBatch]: total_start_time = time.perf_counter() with log_generation_timer(logger, batch.prompt): result = await scheduler_client.forward([batch]) if result.output is None and result.output_file_paths is None: error_msg = result.error or "Unknown error" raise RuntimeError( f"Model generation returned no output. Error from scheduler: {error_msg}" ) if result.output_file_paths: save_file_path_list = result.output_file_paths else: num_outputs = len(result.output) save_file_path_list = save_outputs( result.output, batch.data_type, batch.fps, batch.save_output, lambda idx: str(batch.output_file_path(num_outputs, idx)), audio=result.audio, audio_sample_rate=result.audio_sample_rate, output_compression=batch.output_compression, enable_frame_interpolation=batch.enable_frame_interpolation, frame_interpolation_exp=batch.frame_interpolation_exp, frame_interpolation_scale=batch.frame_interpolation_scale, frame_interpolation_model_path=batch.frame_interpolation_model_path, enable_upscaling=batch.enable_upscaling, upscaling_model_path=batch.upscaling_model_path, upscaling_scale=batch.upscaling_scale, ) total_time = time.perf_counter() - total_start_time log_batch_completion(logger, 1, total_time) if result.peak_memory_mb and result.peak_memory_mb > 0: logger.info(f"Peak memory usage: {result.peak_memory_mb:.2f} MB") return save_file_path_list, result def merge_image_input_list(*inputs: Union[List, Any, None]) -> List: """ Merge multiple image input sources into a single list. This function handles both single items and lists of items, merging them into a single flattened list. Useful for processing images, URLs, or other multimedia inputs that can come as either single items or lists. Args: *inputs: Variable number of inputs, each can be None, single item, or list Returns: List: Flattened list of all non-None inputs Example: >>> merge_image_input_list(["img1", "img2"], "img3", None) ["img1", "img2", "img3"] """ result = [] for input_item in inputs: if input_item is not None: if isinstance(input_item, list): result.extend(input_item) else: result.append(input_item) return result def add_common_data_to_response( response: dict, request_id: str, result: OutputBatch ) -> dict: if result.peak_memory_mb and result.peak_memory_mb > 0: response["peak_memory_mb"] = result.peak_memory_mb if result.metrics and result.metrics.total_duration_s > 0: response["inference_time_s"] = result.metrics.total_duration_s response["id"] = request_id return response def adjust_output_quality(output_quality: str, data_type: DataType = None) -> int: if output_quality == "default": return 50 if data_type == DataType.VIDEO else 75 return OUTPUT_QUALITY_MAPPER.get(output_quality, None) ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import asyncio import json import os import shutil import tempfile import time from typing import Any, Dict, Optional from fastapi import ( APIRouter, File, Form, HTTPException, Path, Query, Request, UploadFile, ) from fastapi.responses import FileResponse from sglang.multimodal_gen.configs.sample.sampling_params import ( SamplingParams, generate_request_id, ) from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( VideoGenerationsRequest, VideoListResponse, VideoResponse, ) from sglang.multimodal_gen.runtime.entrypoints.openai.storage import cloud_storage from sglang.multimodal_gen.runtime.entrypoints.openai.stores import VIDEO_STORE from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( DEFAULT_FPS, DEFAULT_VIDEO_SECONDS, add_common_data_to_response, build_sampling_params, merge_image_input_list, process_generation_batch, save_image_to_path, ) from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.server_args import get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) router = APIRouter(prefix="/v1/videos", tags=["videos"]) def _build_video_sampling_params(request_id: str, request: VideoGenerationsRequest): """Resolve video-specific defaults (fps, seconds → num_frames) then delegate to the shared build_sampling_params.""" seconds = request.seconds if request.seconds is not None else DEFAULT_VIDEO_SECONDS fps = request.fps if request.fps is not None else DEFAULT_FPS num_frames = request.num_frames if request.num_frames is not None else fps * seconds return build_sampling_params( request_id, prompt=request.prompt, size=request.size, num_frames=num_frames, fps=fps, image_path=request.input_reference, output_file_name=request_id, seed=request.seed, generator_device=request.generator_device, num_inference_steps=request.num_inference_steps, guidance_scale=request.guidance_scale, guidance_scale_2=request.guidance_scale_2, negative_prompt=request.negative_prompt, enable_teacache=request.enable_teacache, enable_frame_interpolation=request.enable_frame_interpolation, frame_interpolation_exp=request.frame_interpolation_exp, frame_interpolation_scale=request.frame_interpolation_scale, frame_interpolation_model_path=request.frame_interpolation_model_path, enable_upscaling=request.enable_upscaling, upscaling_model_path=request.upscaling_model_path, upscaling_scale=request.upscaling_scale, output_path=request.output_path, output_compression=request.output_compression, output_quality=request.output_quality, ) # extract metadata which http_server needs to know def _video_job_from_sampling( request_id: str, req: VideoGenerationsRequest, sampling: SamplingParams ) -> Dict[str, Any]: size_str = f"{sampling.width}x{sampling.height}" seconds = int(round((sampling.num_frames or 0) / float(sampling.fps or 24))) return { "id": request_id, "object": "video", "model": req.model or "sora-2", "status": "queued", "progress": 0, "created_at": int(time.time()), "size": size_str, "seconds": str(seconds), "quality": "standard", "file_path": os.path.abspath(sampling.output_file_path()), } async def _save_first_input_image( image_sources, request_id: str, uploads_dir: str ) -> str | None: """Save the first input image from a list of sources and return its path.""" image_list = merge_image_input_list(image_sources) if not image_list: return None image = image_list[0] os.makedirs(uploads_dir, exist_ok=True) filename = image.filename if hasattr(image, "filename") else "url_image" target_path = os.path.join(uploads_dir, f"{request_id}_{filename}") return await save_image_to_path(image, target_path) async def _dispatch_job_async( job_id: str, batch: Req, *, temp_dirs: list[str] | None = None, output_persistent: bool = True, ) -> None: from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client try: save_file_path_list, result = await process_generation_batch( async_scheduler_client, batch ) save_file_path = save_file_path_list[0] cloud_url = await cloud_storage.upload_and_cleanup(save_file_path) persistent_path = ( save_file_path if not cloud_url and output_persistent else None ) update_fields = { "status": "completed", "progress": 100, "completed_at": int(time.time()), "url": cloud_url, "file_path": persistent_path, } update_fields = add_common_data_to_response( update_fields, request_id=job_id, result=result ) await VIDEO_STORE.update_fields(job_id, update_fields) except Exception as e: logger.error(f"{e}") await VIDEO_STORE.update_fields( job_id, {"status": "failed", "error": {"message": str(e)}} ) finally: for td in temp_dirs or []: shutil.rmtree(td, ignore_errors=True) # TODO: support image to video generation @router.post("", response_model=VideoResponse) async def create_video( request: Request, # multipart/form-data fields (optional; used only when content-type is multipart) prompt: Optional[str] = Form(None), input_reference: Optional[UploadFile] = File(None), reference_url: Optional[str] = Form(None), model: Optional[str] = Form(None), seconds: Optional[int] = Form(None), size: Optional[str] = Form(None), fps: Optional[int] = Form(None), num_frames: Optional[int] = Form(None), seed: Optional[int] = Form(1024), generator_device: Optional[str] = Form("cuda"), negative_prompt: Optional[str] = Form(None), guidance_scale: Optional[float] = Form(None), num_inference_steps: Optional[int] = Form(None), enable_teacache: Optional[bool] = Form(False), enable_frame_interpolation: Optional[bool] = Form(False), frame_interpolation_exp: Optional[int] = Form(1), frame_interpolation_scale: Optional[float] = Form(1.0), frame_interpolation_model_path: Optional[str] = Form(None), enable_upscaling: Optional[bool] = Form(False), upscaling_model_path: Optional[str] = Form(None), upscaling_scale: Optional[int] = Form(4), output_quality: Optional[str] = Form("default"), output_compression: Optional[int] = Form(None), extra_body: Optional[str] = Form(None), ): content_type = request.headers.get("content-type", "").lower() request_id = generate_request_id() server_args = get_global_server_args() task_type = server_args.pipeline_config.task_type # Resolve input upload directory (may be a temp dir when saving is disabled) temp_dirs: list[str] = [] if server_args.input_save_path is not None: uploads_dir = server_args.input_save_path os.makedirs(uploads_dir, exist_ok=True) else: uploads_dir = tempfile.mkdtemp(prefix="sglang_input_") temp_dirs.append(uploads_dir) # Resolve output directory effective_output_path = server_args.output_path output_persistent = True if "multipart/form-data" not in content_type: # JSON body may carry a per-request output_path; checked after parsing below pass if "multipart/form-data" in content_type: if not prompt: raise HTTPException(status_code=400, detail="prompt is required") # Validate image input based on model task type image_sources = merge_image_input_list(input_reference, reference_url) if task_type.requires_image_input() and not image_sources: raise HTTPException( status_code=400, detail="input_reference or reference_url is required for image-to-video generation", ) try: input_path = await _save_first_input_image( image_sources, request_id, uploads_dir ) except Exception as e: raise HTTPException( status_code=400, detail=f"Failed to process image source: {str(e)}" ) # Parse extra_body JSON (if provided in multipart form) to get fps/num_frames overrides extra_from_form: Dict[str, Any] = {} if extra_body: try: extra_from_form = json.loads(extra_body) except Exception: extra_from_form = {} fps_val = fps if fps is not None else extra_from_form.get("fps") num_frames_val = ( num_frames if num_frames is not None else extra_from_form.get("num_frames") ) req = VideoGenerationsRequest( prompt=prompt, input_reference=input_path, model=model, seconds=seconds if seconds is not None else 4, size=size, fps=fps_val, num_frames=num_frames_val, seed=seed, generator_device=generator_device, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, enable_teacache=enable_teacache, enable_frame_interpolation=enable_frame_interpolation, frame_interpolation_exp=frame_interpolation_exp, frame_interpolation_scale=frame_interpolation_scale, frame_interpolation_model_path=frame_interpolation_model_path, enable_upscaling=enable_upscaling, upscaling_model_path=upscaling_model_path, upscaling_scale=upscaling_scale, output_compression=output_compression, output_quality=output_quality, **( {"guidance_scale": guidance_scale} if guidance_scale is not None else {} ), ) else: try: body = await request.json() except Exception: body = {} try: # If client uses extra_body, merge it into the top-level payload payload: Dict[str, Any] = dict(body or {}) extra = payload.pop("extra_body", None) if isinstance(extra, dict): # Shallow-merge: only keys like fps/num_frames are expected payload.update(extra) # openai may turn extra_body to extra_json extra_json = payload.pop("extra_json", None) if isinstance(extra_json, dict): payload.update(extra_json) # Validate image input based on model task type has_image_input = payload.get("reference_url") or payload.get( "input_reference" ) if task_type.requires_image_input() and not has_image_input: raise HTTPException( status_code=400, detail="input_reference or reference_url is required for image-to-video generation", ) # for non-multipart/form-data type if payload.get("reference_url"): try: input_path = await _save_first_input_image( payload.get("reference_url"), request_id, uploads_dir ) except Exception as e: raise HTTPException( status_code=400, detail=f"Failed to process image source: {str(e)}", ) payload["input_reference"] = input_path req = VideoGenerationsRequest(**payload) except Exception as e: raise HTTPException(status_code=400, detail=f"Invalid request body: {e}") # Resolve per-request output_path override effective_output_path = req.output_path or server_args.output_path if effective_output_path is None: output_tmp = tempfile.mkdtemp(prefix="sglang_output_") temp_dirs.append(output_tmp) effective_output_path = output_tmp output_persistent = False # Inject resolved output_path so _build_video_sampling_params picks it up req.output_path = effective_output_path logger.debug(f"Server received from create_video endpoint: req={req}") try: sampling_params = _build_video_sampling_params(request_id, req) except (ValueError, TypeError) as e: raise HTTPException(status_code=400, detail=str(e)) job = _video_job_from_sampling(request_id, req, sampling_params) await VIDEO_STORE.upsert(request_id, job) # Build Req for scheduler batch = prepare_request( server_args=server_args, sampling_params=sampling_params, ) # Add diffusers_kwargs if provided if req.diffusers_kwargs: batch.extra["diffusers_kwargs"] = req.diffusers_kwargs # Enqueue the job asynchronously and return immediately asyncio.create_task( _dispatch_job_async( request_id, batch, temp_dirs=temp_dirs or None, output_persistent=output_persistent, ) ) return VideoResponse(**job) @router.get("", response_model=VideoListResponse) async def list_videos( after: Optional[str] = Query(None), limit: Optional[int] = Query(None, ge=1, le=100), order: Optional[str] = Query("desc"), ): # Normalize order order = (order or "desc").lower() if order not in ("asc", "desc"): order = "desc" jobs = await VIDEO_STORE.list_values() reverse = order != "asc" jobs.sort(key=lambda j: j.get("created_at", 0), reverse=reverse) if after is not None: try: idx = next(i for i, j in enumerate(jobs) if j["id"] == after) jobs = jobs[idx + 1 :] except StopIteration: jobs = [] if limit is not None: jobs = jobs[:limit] items = [VideoResponse(**j) for j in jobs] return VideoListResponse(data=items) @router.get("/{video_id}", response_model=VideoResponse) async def retrieve_video(video_id: str = Path(...)): job = await VIDEO_STORE.get(video_id) if not job: raise HTTPException(status_code=404, detail="Video not found") return VideoResponse(**job) # TODO: support aborting a job. @router.delete("/{video_id}", response_model=VideoResponse) async def delete_video(video_id: str = Path(...)): job = await VIDEO_STORE.pop(video_id) if not job: raise HTTPException(status_code=404, detail="Video not found") # Mark as deleted in response semantics job["status"] = "deleted" return VideoResponse(**job) @router.get("/{video_id}/content") async def download_video_content( video_id: str = Path(...), variant: Optional[str] = Query(None) ): job = await VIDEO_STORE.get(video_id) if not job: raise HTTPException(status_code=404, detail="Video not found") if job.get("url"): raise HTTPException( status_code=400, detail=f"Video has been uploaded to cloud storage. Please use the cloud URL: {job.get('url')}", ) file_path = job.get("file_path") if not file_path or not os.path.exists(file_path): raise HTTPException(status_code=404, detail="Generation is still in-progress") media_type = "video/mp4" # default variant return FileResponse( path=file_path, media_type=media_type, filename=os.path.basename(file_path) ) ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/post_training/__init__.py ================================================ ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py ================================================ """Request/response data structures for post-training APIs.""" from dataclasses import dataclass @dataclass class UpdateWeightFromDiskReqInput: """Request to update model weights from disk for diffusion models.""" model_path: str flush_cache: bool = True target_modules: list[str] | None = None @dataclass class GetWeightsChecksumReqInput: """Compute SHA-256 checksum of loaded module weights for verification.""" module_names: list[str] | None = None ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py ================================================ """Weight update API for the diffusion engine.""" from fastapi import APIRouter, Request from fastapi.responses import ORJSONResponse from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( GetWeightsChecksumReqInput, UpdateWeightFromDiskReqInput, ) from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client router = APIRouter() @router.post("/update_weights_from_disk") async def update_weights_from_disk(request: Request): """Update model weights from disk inplace without restarting the server.""" body = await request.json() model_path = body.get("model_path") if not model_path: return ORJSONResponse( {"success": False, "message": "model_path is required"}, status_code=400, ) req = UpdateWeightFromDiskReqInput( model_path=model_path, flush_cache=body.get("flush_cache", True), target_modules=body.get("target_modules"), ) try: response = await async_scheduler_client.forward(req) except Exception as e: return ORJSONResponse( {"success": False, "message": str(e)}, status_code=500, ) result = response.output success = result.get("success", False) message = result.get("message", "Unknown status") return ORJSONResponse( {"success": success, "message": message}, status_code=200 if success else 400, ) @router.post("/get_weights_checksum") async def get_weights_checksum(request: Request): """Return SHA-256 checksum of each requested module's weights.""" body = await request.json() req = GetWeightsChecksumReqInput( module_names=body.get("module_names"), ) try: response = await async_scheduler_client.forward(req) except Exception as e: return ORJSONResponse({"error": str(e)}, status_code=500) return ORJSONResponse(response.output, status_code=200) ================================================ FILE: python/sglang/multimodal_gen/runtime/entrypoints/utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ DiffGenerator module for sglang-diffusion. This module provides a consolidated interface for generating videos using diffusion models. """ import os import shutil import subprocess import tempfile from dataclasses import dataclass, field from typing import Any, Callable, List, Optional, Sequence, Union import imageio import numpy as np import torch try: import scipy.io.wavfile as scipy_wavfile except ImportError: # pragma: no cover scipy_wavfile = None try: import imageio_ffmpeg as _imageio_ffmpeg except ImportError: # pragma: no cover _imageio_ffmpeg = None from sglang.multimodal_gen.configs.sample.sampling_params import ( DataType, SamplingParams, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_logger logger = init_logger(__name__) @dataclass class SetLoraReq: lora_nickname: Union[str, List[str]] lora_path: Optional[Union[str, List[Optional[str]]]] = None target: Union[str, List[str]] = "all" strength: Union[float, List[float]] = 1.0 @dataclass class MergeLoraWeightsReq: target: str = "all" strength: float = 1.0 @dataclass class UnmergeLoraWeightsReq: target: str = "all" @dataclass class ListLorasReq: pass @dataclass class ShutdownReq: pass def format_lora_message( lora_nickname: Union[str, List[str]], target: Union[str, List[str]], strength: Union[float, List[float]], ) -> tuple[str, str, str]: """Format success message for single or multiple LoRAs.""" if isinstance(lora_nickname, list): nickname_str = ", ".join(lora_nickname) target_str = ", ".join(target) if isinstance(target, list) else target strength_str = ( ", ".join(f"{s:.2f}" for s in strength) if isinstance(strength, list) else f"{strength:.2f}" ) else: nickname_str = lora_nickname target_str = target if isinstance(target, str) else ", ".join(target) strength_str = ( f"{strength:.2f}" if isinstance(strength, (int, float)) else ", ".join(f"{s:.2f}" for s in strength) ) return nickname_str, target_str, strength_str @dataclass class GenerationResult: """Result of a single generation request from DiffGenerator.""" samples: Any = None frames: Any = None audio: Any = None prompt: str | None = None size: tuple | None = None # (height, width, num_frames) generation_time: float = 0.0 peak_memory_mb: float = 0.0 metrics: dict = field(default_factory=dict) trajectory_latents: Any = None trajectory_timesteps: Any = None trajectory_decoded: Any = None prompt_index: int = 0 output_file_path: str | None = None def _normalize_audio_to_numpy(audio: Any) -> np.ndarray | None: """Convert audio (torch / numpy) into a float32 numpy array in [-1, 1], best-effort.""" if audio is None: return None if isinstance(audio, torch.Tensor): audio_np = audio.detach().float().clamp(-1.0, 1.0).cpu().numpy() elif isinstance(audio, np.ndarray): audio_np = audio.astype(np.float32, copy=False) audio_np = np.clip(audio_np, -1.0, 1.0) else: return None # 1. Squeeze leading singleton dimensions (Batch, etc.) while audio_np.ndim > 1 and audio_np.shape[0] == 1: audio_np = audio_np.squeeze(0) # 2. Handle (C, L) -> (L, C) if audio_np.ndim == 2 and audio_np.shape[0] < audio_np.shape[1]: audio_np = audio_np.transpose(1, 0) # 3. Final safety check: if still 2D and channels (dim 1) is huge, something is wrong if audio_np.ndim == 2 and audio_np.shape[1] > 256 and audio_np.shape[0] == 1: audio_np = audio_np.flatten() return audio_np def _pick_audio_sample_rate( *, audio_np: np.ndarray, audio_sample_rate: Optional[int], fps: int, num_frames: int, ) -> int: """Pick a plausible sample rate, falling back to inferring from video duration.""" selected_sr = int(audio_sample_rate) if audio_sample_rate is not None else None if selected_sr is None or not (8000 <= selected_sr <= 192000): selected_sr = 24000 try: duration_s = float(num_frames) / float(fps) if fps else 0.0 if duration_s > 0: audio_len = ( int(audio_np.shape[0]) if audio_np.ndim == 2 else int(audio_np.shape[-1]) ) inferred_sr = int(round(float(audio_len) / duration_s)) if 8000 <= inferred_sr <= 192000: selected_sr = inferred_sr except Exception: pass return selected_sr def _resolve_ffmpeg_exe() -> str: ffmpeg_exe = "ffmpeg" ffmpeg_on_path = shutil.which("ffmpeg") if ffmpeg_on_path: ffmpeg_exe = ffmpeg_on_path try: if _imageio_ffmpeg is not None: ffmpeg_exe = _imageio_ffmpeg.get_ffmpeg_exe() except Exception: pass ffmpeg_ok = False if ffmpeg_exe: if os.path.isabs(ffmpeg_exe): ffmpeg_ok = os.path.exists(ffmpeg_exe) else: ffmpeg_ok = shutil.which(ffmpeg_exe) is not None if not ffmpeg_ok: raise RuntimeError("ffmpeg not found") return ffmpeg_exe def _mux_audio_np_into_mp4( *, save_file_path: str, audio_np: np.ndarray, sample_rate: int, ffmpeg_exe: str, ) -> None: merged_path = save_file_path.rsplit(".", 1)[0] + ".tmp_mux.mp4" tmp_wav_path = None try: if scipy_wavfile is None: raise RuntimeError( "scipy is required to mux audio into mp4 (pip install scipy)" ) with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f: tmp_wav_path = f.name scipy_wavfile.write(tmp_wav_path, sample_rate, audio_np) subprocess.run( [ ffmpeg_exe, "-y", "-i", save_file_path, "-i", tmp_wav_path, "-c:v", "copy", "-c:a", "aac", "-strict", "experimental", merged_path, ], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) os.replace(merged_path, save_file_path) finally: if tmp_wav_path: try: os.remove(tmp_wav_path) except OSError: pass if os.path.exists(merged_path): try: os.remove(merged_path) except OSError: pass def _maybe_mux_audio_into_mp4( *, save_file_path: str, audio: Any, frames: list, fps: int, audio_sample_rate: Optional[int], ) -> None: """Best-effort mux audio into an already-written mp4 at save_file_path. Any failure should keep the silent video and only log a warning. """ audio_np = _normalize_audio_to_numpy(audio) if audio_np is None: return selected_sr = _pick_audio_sample_rate( audio_np=audio_np, audio_sample_rate=audio_sample_rate, fps=fps, num_frames=len(frames), ) try: ffmpeg_exe = _resolve_ffmpeg_exe() _mux_audio_np_into_mp4( save_file_path=save_file_path, audio_np=audio_np, sample_rate=selected_sr, ffmpeg_exe=ffmpeg_exe, ) logger.info(f"Merged video saved to {CYAN}{save_file_path}{RESET}") except Exception as e: logger.warning( "Failed to mux audio into mp4 (saved silent video): %s", str(e), ) def prepare_request( server_args: ServerArgs, sampling_params: SamplingParams, ) -> Req: """ Create a Req object with sampling_params as a parameter. """ req = Req( sampling_params=sampling_params, VSA_sparsity=server_args.attention_backend_config.VSA_sparsity, ) try: diffusers_kwargs = sampling_params.diffusers_kwargs except AttributeError: diffusers_kwargs = None if diffusers_kwargs: req.extra["diffusers_kwargs"] = diffusers_kwargs req.adjust_size(server_args) if not isinstance(req.prompt, str): raise TypeError(f"`prompt` must be a string, but got {type(req.prompt)}") if (req.width is not None and req.width <= 0) or ( req.height is not None and req.height <= 0 ): raise ValueError( f"Height and width must be positive, got height={req.height}, width={req.width}" ) return req def attach_audio_to_video_sample( sample: Any, audio: Any, output_idx: int, ) -> Any: """Attach per-sample audio for video outputs when available.""" if audio is None: return sample if isinstance(audio, torch.Tensor) and audio.ndim >= 2: audio = audio[output_idx] if audio.shape[0] > output_idx else None elif isinstance(audio, np.ndarray) and audio.ndim >= 2: audio = audio[output_idx] if audio.shape[0] > output_idx else None if audio is not None and not ( isinstance(sample, (tuple, list)) and len(sample) == 2 ): return (sample, audio) return sample def save_outputs( outputs: Sequence[Any], data_type: DataType, fps: int, save_output: bool, build_output_path: Callable[[int], str], *, audio: Any = None, audio_sample_rate: Optional[int] = None, samples_out: Optional[list[Any]] = None, audios_out: Optional[list[Any]] = None, frames_out: Optional[list[Any]] = None, output_compression: Optional[int] = None, enable_frame_interpolation: bool = False, frame_interpolation_exp: int = 1, frame_interpolation_scale: float = 1.0, frame_interpolation_model_path: Optional[str] = None, enable_upscaling: bool = False, upscaling_model_path: Optional[str] = None, upscaling_scale: int = 4, ) -> list[str]: """Save outputs to files and return the list of file paths.""" output_paths: list[str] = [] for idx, output in enumerate(outputs): save_file_path = build_output_path(idx) sample = output if data_type == DataType.VIDEO: sample = attach_audio_to_video_sample(sample, audio, idx) frames = post_process_sample( sample, data_type, fps, save_output, save_file_path, audio_sample_rate=audio_sample_rate, output_compression=output_compression, enable_frame_interpolation=enable_frame_interpolation, frame_interpolation_exp=frame_interpolation_exp, frame_interpolation_scale=frame_interpolation_scale, frame_interpolation_model_path=frame_interpolation_model_path, enable_upscaling=enable_upscaling, upscaling_model_path=upscaling_model_path, upscaling_scale=upscaling_scale, ) if samples_out is not None: samples_out.append(sample) if audios_out is not None: if data_type == DataType.VIDEO: audio_item = audio if isinstance(audio, torch.Tensor) and audio.ndim >= 2: audio_item = audio[idx] if audio.shape[0] > idx else None elif isinstance(audio, np.ndarray) and audio.ndim >= 2: audio_item = audio[idx] if audio.shape[0] > idx else None audios_out.append(audio_item) else: audios_out.append(audio) if frames_out is not None: frames_out.append(frames) output_paths.append(save_file_path) return output_paths def post_process_sample( sample: Any, data_type: DataType, fps: int, save_output: bool = True, save_file_path: Optional[str] = None, audio_sample_rate: Optional[int] = None, output_compression: Optional[int] = None, enable_frame_interpolation: bool = False, frame_interpolation_exp: int = 1, frame_interpolation_scale: float = 1.0, frame_interpolation_model_path: Optional[str] = None, enable_upscaling: bool = False, upscaling_model_path: Optional[str] = None, upscaling_scale: int = 4, ): """ Process sample output, optionally interpolate video frames, and save. """ audio = None if isinstance(sample, (tuple, list)) and len(sample) == 2: sample, audio = sample # 1. Convert tensor / array to list of uint8 HWC frames frames = None if isinstance(sample, torch.Tensor): if sample.dim() == 3: sample = sample.unsqueeze(1) sample = (sample * 255).clamp(0, 255).to(torch.uint8) videos = sample.permute(1, 2, 3, 0).cpu().numpy() frames = list(videos) else: if not isinstance(sample, np.ndarray): raise TypeError(f"Unsupported sample type: {type(sample)}") arr = sample if arr.ndim == 3: if arr.shape[-1] in (1, 3, 4): arr = arr[None, ...] else: arr = arr[..., None] if arr.ndim != 4: raise ValueError(f"Unexpected numpy sample shape: {tuple(arr.shape)}") if arr.shape[-1] not in (1, 3, 4) and arr.shape[0] in (1, 3, 4): t = torch.from_numpy(arr) if t.dim() == 3: t = t.unsqueeze(1) t = (t * 255).clamp(0, 255).to(torch.uint8) videos = t.permute(1, 2, 3, 0).cpu().numpy() frames = list(videos) else: if arr.dtype != np.uint8: arr = (np.clip(arr, 0.0, 1.0) * 255.0).astype(np.uint8) frames = list(arr) # 2. Frame interpolation (video only) if enable_frame_interpolation and data_type == DataType.VIDEO and len(frames) > 1: from sglang.multimodal_gen.runtime.postprocess import ( interpolate_video_frames, ) frames, multiplier = interpolate_video_frames( frames, exp=frame_interpolation_exp, scale=frame_interpolation_scale, model_path=frame_interpolation_model_path, ) fps = fps * multiplier # 3. Upscaling (images and videos) if enable_upscaling and frames: from sglang.multimodal_gen.runtime.postprocess import upscale_frames frames = upscale_frames( frames, model_path=upscaling_model_path, scale=upscaling_scale, ) # 4. Save outputs if requested if save_output: if save_file_path: os.makedirs(os.path.dirname(save_file_path), exist_ok=True) if data_type == DataType.VIDEO: quality = ( output_compression / 10 if output_compression is not None else 5 ) imageio.mimsave( save_file_path, frames, fps=fps, format=data_type.get_default_extension(), codec="libx264", quality=quality, ) _maybe_mux_audio_into_mp4( save_file_path=save_file_path, audio=audio, frames=frames, fps=fps, audio_sample_rate=audio_sample_rate, ) else: quality = output_compression if output_compression is not None else 75 if len(frames) > 1: for i, image in enumerate(frames): parts = save_file_path.rsplit(".", 1) if len(parts) == 2: indexed_path = f"{parts[0]}_{i}.{parts[1]}" else: indexed_path = f"{save_file_path}_{i}" imageio.imwrite(indexed_path, image, quality=quality) else: imageio.imwrite(save_file_path, frames[0], quality=quality) logger.info(f"Output saved to {CYAN}{save_file_path}{RESET}") else: logger.info(f"No output path provided, output not saved") return frames ================================================ FILE: python/sglang/multimodal_gen/runtime/launch_server.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import multiprocessing as mp import os import signal import sys import threading import psutil import uvicorn from sglang.multimodal_gen.runtime.entrypoints.http_server import create_app from sglang.multimodal_gen.runtime.managers.gpu_worker import run_scheduler_process from sglang.multimodal_gen.runtime.server_args import ( ServerArgs, prepare_server_args, set_global_server_args, ) from sglang.multimodal_gen.runtime.utils.logging_utils import configure_logger, logger def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None): """Kill the process and all its child processes.""" # Remove sigchld handler to avoid spammy logs. if threading.current_thread() is threading.main_thread(): signal.signal(signal.SIGCHLD, signal.SIG_DFL) if parent_pid is None: parent_pid = os.getpid() include_parent = False try: itself = psutil.Process(parent_pid) except psutil.NoSuchProcess: return children = itself.children(recursive=True) for child in children: if child.pid == skip_pid: continue try: child.kill() except psutil.NoSuchProcess: pass if include_parent: try: if parent_pid == os.getpid(): itself.kill() sys.exit(0) itself.kill() # Sometime processes cannot be killed with SIGKILL (e.g, PID=1 launched by kubernetes), # so we send an additional signal to kill them. itself.send_signal(signal.SIGQUIT) except psutil.NoSuchProcess: pass def launch_server(server_args: ServerArgs, launch_http_server: bool = True): """ Args: launch_http_server: False for offline local mode """ configure_logger(server_args) # Start a new server with multiple worker processes logger.info("Starting server...") num_gpus = server_args.num_gpus processes = [] # Pipes for master to talk to slaves task_pipes_to_slaves_w = [] task_pipes_to_slaves_r = [] for _ in range(num_gpus - 1): r, w = mp.Pipe(duplex=False) task_pipes_to_slaves_r.append(r) task_pipes_to_slaves_w.append(w) # Pipes for slaves to talk to master result_pipes_from_slaves_w = [] result_pipes_from_slaves_r = [] for _ in range(num_gpus - 1): r, w = mp.Pipe(duplex=False) result_pipes_from_slaves_r.append(r) result_pipes_from_slaves_w.append(w) # Launch all worker processes master_port = server_args.master_port or (server_args.master_port + 100) scheduler_pipe_readers = [] scheduler_pipe_writers = [] for i in range(num_gpus): reader, writer = mp.Pipe(duplex=False) scheduler_pipe_writers.append(writer) if i == 0: # Master worker process = mp.Process( target=run_scheduler_process, args=( i, # local_rank i, # rank master_port, server_args, writer, None, # No task pipe to read from master None, # No result pipe to write to master task_pipes_to_slaves_w, result_pipes_from_slaves_r, ), name=f"sglang-diffusionWorker-{i}", daemon=True, ) else: # Slave workers process = mp.Process( target=run_scheduler_process, args=( i, # local_rank i, # rank master_port, server_args, writer, None, # No task pipe to read from master None, # No result pipe to write to master task_pipes_to_slaves_r[i - 1], result_pipes_from_slaves_w[i - 1], ), name=f"sglang-diffusionWorker-{i}", daemon=True, ) scheduler_pipe_readers.append(reader) process.start() processes.append(process) # Wait for all workers to be ready scheduler_infos = [] for writer in scheduler_pipe_writers: writer.close() # Close unused pipe ends in parent process for p in task_pipes_to_slaves_w: p.close() for p in task_pipes_to_slaves_r: p.close() for p in result_pipes_from_slaves_w: p.close() for p in result_pipes_from_slaves_r: p.close() for i, reader in enumerate(scheduler_pipe_readers): try: data = reader.recv() except EOFError: logger.error( f"Rank {i} scheduler is dead. Please check if there are relevant logs." ) processes[i].join() logger.error(f"Exit code: {processes[i].exitcode}") raise if data["status"] != "ready": raise RuntimeError( "Initialization failed. Please see the error messages above." ) scheduler_infos.append(data) reader.close() logger.debug("All workers are ready") if launch_http_server: logger.info("Starting FastAPI server.") if server_args.webui: logger.info("Launch FastAPI server in another process because of webui.") http_server_process = mp.Process( target=launch_http_server_only, args=(server_args,), name=f"sglang-diffusion-webui", daemon=True, ) http_server_process.start() else: launch_http_server_only(server_args) return processes def launch_http_server_only(server_args): # set for endpoints to access global_server_args set_global_server_args(server_args) app = create_app(server_args) uvicorn.run( app, use_colors=True, log_level=server_args.log_level, host=server_args.host, port=server_args.port, reload=False, ) if __name__ == "__main__": server_args = prepare_server_args(sys.argv[1:]) try: launch_server(server_args) finally: kill_process_tree(os.getpid(), include_parent=False) ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/activation.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/activation.py """Custom activation functions.""" import math from typing import Any import torch import torch.nn as nn import torch.nn.functional as F from sglang.multimodal_gen.runtime.platforms import current_platform _is_cuda = current_platform.is_cuda() _is_hip = current_platform.is_hip() _is_npu = current_platform.is_npu() if _is_cuda or _is_hip: from sgl_kernel import silu_and_mul if _is_npu: import torch_npu # TODO (will): remove this dependency from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp @CustomOp.register("silu_and_mul") class SiluAndMul(CustomOp): """An activation function for SwiGLU. The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2. Shapes: x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d) return: (num_tokens, d) or (batch_size, seq_len, d) """ def __init__(self) -> None: super().__init__() def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: d = x.shape[-1] // 2 output_shape = x.shape[:-1] + (d,) out = torch.empty(output_shape, dtype=x.dtype, device=x.device) silu_and_mul(x, out) return out def forward_native(self, x: torch.Tensor) -> torch.Tensor: """PyTorch-native implementation equivalent to forward().""" d = x.shape[-1] // 2 return F.silu(x[..., :d]) * x[..., d:] def forward_npu(self, x: torch.Tensor) -> torch.Tensor: out = torch_npu.npu_swiglu(x) return out def forward_musa(self, x: torch.Tensor) -> torch.Tensor: return nn.SwishGLU()(x) @CustomOp.register("gelu_and_mul") class GeluAndMul(CustomOp): """An activation function for GeGLU. The function computes x -> GELU(x[:d]) * x[d:] where d = x.shape[-1] // 2. Shapes: x: (batch_size, seq_len, 2 * d) or (num_tokens, 2 * d) return: (batch_size, seq_len, d) or (num_tokens, d) """ def __init__(self, approximate: str = "none"): super().__init__() self.approximate = approximate if approximate not in ("none", "tanh"): raise ValueError(f"Unknown approximate mode: {approximate}") def forward_cuda(self, *args, **kwargs) -> Any: return self.forward_native(*args, **kwargs) def forward_native(self, x: torch.Tensor) -> torch.Tensor: """PyTorch-native implementation equivalent to forward().""" d = x.shape[-1] // 2 return F.gelu(x[..., :d], approximate=self.approximate) * x[..., d:] def extra_repr(self) -> str: return f"approximate={repr(self.approximate)}" @CustomOp.register("gelu_new") class NewGELU(CustomOp): def __init__(self): super().__init__() def forward_cuda(self, *args, **kwargs) -> Any: return self.forward_native(*args, **kwargs) def forward_native(self, x: torch.Tensor) -> torch.Tensor: """PyTorch-native implementation equivalent to forward().""" c = math.sqrt(2.0 / math.pi) return 0.5 * x * (1.0 + torch.tanh(c * (x + 0.044715 * torch.pow(x, 3.0)))) @CustomOp.register("quick_gelu") class QuickGELU(CustomOp): # https://github.com/huggingface/transformers/blob/main/src/transformers/activations.py#L90 def __init__(self): super().__init__() def forward_cuda(self, *args, **kwargs) -> Any: return self.forward_native(*args, **kwargs) def forward_native(self, x: torch.Tensor) -> torch.Tensor: """PyTorch-native implementation equivalent to forward().""" return x * torch.sigmoid(1.702 * x) _ACTIVATION_REGISTRY = { "gelu": nn.GELU, "gelu_new": NewGELU, "gelu_pytorch_tanh": lambda: nn.GELU(approximate="tanh"), "relu": nn.ReLU, "silu": nn.SiLU, "quick_gelu": QuickGELU, } def get_act_fn(act_fn_name: str) -> nn.Module: """Get an activation function by name.""" act_fn_name = act_fn_name.lower() if act_fn_name not in _ACTIVATION_REGISTRY: raise ValueError(f"Activation function {act_fn_name!r} is not supported.") return _ACTIVATION_REGISTRY[act_fn_name]() _ACTIVATION_AND_MUL_REGISTRY = { "gelu": GeluAndMul, "silu": SiluAndMul, } def get_act_and_mul_fn(act_fn_name: str) -> nn.Module: """Get an activation-and-mul (i.e. SiluAndMul) function by name.""" act_fn_name = act_fn_name.lower() if act_fn_name not in _ACTIVATION_AND_MUL_REGISTRY: raise ValueError(f"Activation function {act_fn_name!r} is not supported.") return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name]() ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import json import os from collections import defaultdict from typing import Any import numpy as np from sglang.multimodal_gen.utils import dict_to_3d_list def configure_sta( mode: str = "STA_searching", layer_num: int = 40, time_step_num: int = 50, head_num: int = 40, **kwargs, ) -> list[list[list[Any]]]: """ Configure Sliding Tile Attention (STA) parameters based on the specified mode. Parameters: ---------- mode : str The STA mode to use. Options are: - 'STA_searching': Generate a set of mask candidates for initial search - 'STA_tuning': Select best mask strategy based on previously saved results - 'STA_inference': Load and use a previously tuned mask strategy layer_num: int, number of layers time_step_num: int, number of timesteps head_num: int, number of heads **kwargs : dict Mode-specific parameters: For 'STA_searching': - mask_candidates: list of str, optional, mask candidates to use - mask_selected: list of int, optional, indices of selected masks For 'STA_tuning': - mask_search_files_path: str, required, path to mask search results - mask_candidates: list of str, optional, mask candidates to use - mask_selected: list of int, optional, indices of selected masks - skip_time_steps: int, optional, number of time steps to use full attention (default 12) - save_dir: str, optional, directory to save mask strategy (default "mask_candidates") For 'STA_inference': - load_path: str, optional, path to load mask strategy (default "mask_candidates/mask_strategy.json") """ valid_modes = ["STA_searching", "STA_tuning", "STA_inference", "STA_tuning_cfg"] if mode not in valid_modes: raise ValueError(f"Mode must be one of {valid_modes}, got {mode}") if mode == "STA_searching": # Get parameters with defaults mask_candidates: list[str] | None = kwargs.get("mask_candidates") if mask_candidates is None: raise ValueError("mask_candidates is required for STA_searching mode") mask_selected: list[int] = kwargs.get( "mask_selected", list(range(len(mask_candidates))) ) # Parse selected masks selected_masks: list[list[int]] = [] for index in mask_selected: mask = mask_candidates[index] masks_list = [int(x) for x in mask.split(",")] selected_masks.append(masks_list) # Create 3D mask structure with fixed dimensions (t=50, l=60) masks_3d: list[list[list[list[int]]]] = [] for i in range(time_step_num): # Fixed t dimension = 50 row = [] for j in range(layer_num): # Fixed l dimension = 60 row.append(selected_masks) # Add all masks at each position masks_3d.append(row) return masks_3d elif mode == "STA_tuning": # Get required parameters mask_search_files_path: str | None = kwargs.get("mask_search_files_path") if not mask_search_files_path: raise ValueError("mask_search_files_path is required for STA_tuning mode") # Get optional parameters with defaults mask_candidates_tuning: list[str] | None = kwargs.get("mask_candidates") if mask_candidates_tuning is None: raise ValueError("mask_candidates is required for STA_tuning mode") mask_selected_tuning: list[int] = kwargs.get( "mask_selected", list(range(len(mask_candidates_tuning))) ) skip_time_steps_tuning: int | None = kwargs.get("skip_time_steps") save_dir_tuning: str | None = kwargs.get("save_dir", "mask_candidates") # Parse selected masks selected_masks_tuning: list[list[int]] = [] for index in mask_selected_tuning: mask = mask_candidates_tuning[index] masks_list = [int(x) for x in mask.split(",")] selected_masks_tuning.append(masks_list) # Read JSON results results = read_specific_json_files(mask_search_files_path) averaged_results = average_head_losses(results, selected_masks_tuning) # Add full attention mask for specific cases full_attention_mask_tuning: list[int] | None = kwargs.get("full_attention_mask") if full_attention_mask_tuning is not None: selected_masks_tuning.append(full_attention_mask_tuning) # Select best mask strategy timesteps_tuning: int = kwargs.get("timesteps", time_step_num) if skip_time_steps_tuning is None: skip_time_steps_tuning = 12 mask_strategy, sparsity, strategy_counts = select_best_mask_strategy( averaged_results, selected_masks_tuning, skip_time_steps_tuning, timesteps_tuning, head_num, ) # Save mask strategy if save_dir_tuning is not None: os.makedirs(save_dir_tuning, exist_ok=True) file_path = os.path.join( save_dir_tuning, f"mask_strategy_s{skip_time_steps_tuning}.json" ) with open(file_path, "w") as f: json.dump(mask_strategy, f, indent=4) print(f"Successfully saved mask_strategy to {file_path}") # Print sparsity and strategy counts for information print(f"Overall sparsity: {sparsity:.4f}") print("\nStrategy usage counts:") total_heads = time_step_num * layer_num * head_num # Fixed dimensions for strategy, count in strategy_counts.items(): print(f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)") # Convert dictionary to 3D list with fixed dimensions mask_strategy_3d = dict_to_3d_list( mask_strategy, t_max=time_step_num, l_max=layer_num, h_max=head_num ) return mask_strategy_3d elif mode == "STA_tuning_cfg": # Get required parameters for both positive and negative paths mask_search_files_path_pos: str | None = kwargs.get( "mask_search_files_path_pos" ) mask_search_files_path_neg: str | None = kwargs.get( "mask_search_files_path_neg" ) save_dir_cfg: str | None = kwargs.get("save_dir") if ( not mask_search_files_path_pos or not mask_search_files_path_neg or not save_dir_cfg ): raise ValueError( "mask_search_files_path_pos, mask_search_files_path_neg, and save_dir are required for STA_tuning_cfg mode" ) # Get optional parameters with defaults mask_candidates_cfg: list[str] | None = kwargs.get("mask_candidates") if mask_candidates_cfg is None: raise ValueError("mask_candidates is required for STA_tuning_cfg mode") mask_selected_cfg: list[int] = kwargs.get( "mask_selected", list(range(len(mask_candidates_cfg))) ) skip_time_steps_cfg: int | None = kwargs.get("skip_time_steps") # Parse selected masks selected_masks_cfg: list[list[int]] = [] for index in mask_selected_cfg: mask = mask_candidates_cfg[index] masks_list = [int(x) for x in mask.split(",")] selected_masks_cfg.append(masks_list) # Read JSON results for both positive and negative paths pos_results = read_specific_json_files(mask_search_files_path_pos) neg_results = read_specific_json_files(mask_search_files_path_neg) # Combine positive and negative results into one list combined_results = pos_results + neg_results # Average the combined results averaged_results = average_head_losses(combined_results, selected_masks_cfg) # Add full attention mask for specific cases full_attention_mask_cfg: list[int] | None = kwargs.get("full_attention_mask") if full_attention_mask_cfg is not None: selected_masks_cfg.append(full_attention_mask_cfg) timesteps_cfg: int = kwargs.get("timesteps", time_step_num) if skip_time_steps_cfg is None: skip_time_steps_cfg = 12 # Select best mask strategy using combined results mask_strategy, sparsity, strategy_counts = select_best_mask_strategy( averaged_results, selected_masks_cfg, skip_time_steps_cfg, timesteps_cfg, head_num, ) # Save mask strategy os.makedirs(save_dir_cfg, exist_ok=True) file_path = os.path.join( save_dir_cfg, f"mask_strategy_s{skip_time_steps_cfg}.json" ) with open(file_path, "w") as f: json.dump(mask_strategy, f, indent=4) print(f"Successfully saved mask_strategy to {file_path}") # Print sparsity and strategy counts for information print(f"Overall sparsity: {sparsity:.4f}") print("\nStrategy usage counts:") total_heads = time_step_num * layer_num * head_num # Fixed dimensions for strategy, count in strategy_counts.items(): print(f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)") # Convert dictionary to 3D list with fixed dimensions mask_strategy_3d = dict_to_3d_list( mask_strategy, t_max=time_step_num, l_max=layer_num, h_max=head_num ) return mask_strategy_3d else: # STA_inference # Get parameters with defaults load_path: str | None = kwargs.get( "load_path", "mask_candidates/mask_strategy.json" ) if load_path is None: raise ValueError("load_path is required for STA_inference mode") # Load previously saved mask strategy with open(load_path) as f: mask_strategy = json.load(f) # Convert dictionary to 3D list with fixed dimensions mask_strategy_3d = dict_to_3d_list( mask_strategy, t_max=time_step_num, l_max=layer_num, h_max=head_num ) return mask_strategy_3d # Helper functions def read_specific_json_files(folder_path: str) -> list[dict[str, Any]]: """Read and parse JSON files containing mask search results.""" json_contents: list[dict[str, Any]] = [] # List files only in the current directory (no walk) files = os.listdir(folder_path) # Filter files matching_files = [f for f in files if "mask" in f and f.endswith(".json")] print(f"Found {len(matching_files)} matching files: {matching_files}") for file_name in matching_files: file_path = os.path.join(folder_path, file_name) with open(file_path) as file: data = json.load(file) json_contents.append(data) return json_contents def average_head_losses( results: list[dict[str, Any]], selected_masks: list[list[int]] ) -> dict[str, dict[str, np.ndarray]]: """Average losses across all prompts for each mask strategy.""" # Initialize a dictionary to store the averaged results averaged_losses: dict[str, dict[str, np.ndarray]] = {} loss_type = "L2_loss" # Get all loss types (e.g., 'L2_loss') averaged_losses[loss_type] = {} for mask in selected_masks: mask_str = str(mask) data_shape = np.array(results[0][loss_type][mask_str]).shape accumulated_data = np.zeros(data_shape) # Sum across all prompts for prompt_result in results: accumulated_data += np.array(prompt_result[loss_type][mask_str]) # Average by dividing by number of prompts averaged_data = accumulated_data / len(results) averaged_losses[loss_type][mask_str] = averaged_data return averaged_losses def select_best_mask_strategy( averaged_results: dict[str, dict[str, np.ndarray]], selected_masks: list[list[int]], skip_time_steps: int = 12, timesteps: int = 50, head_num: int = 40, ) -> tuple[dict[str, list[int]], float, dict[str, int]]: """Select the best mask strategy for each head based on loss minimization.""" best_mask_strategy: dict[str, list[int]] = {} loss_type = "L2_loss" # Get the shape of time steps and layers layers = len(averaged_results[loss_type][str(selected_masks[0])][0]) # Counter for sparsity calculation total_tokens = 0 # total number of masked tokens total_length = 0 # total sequence length strategy_counts: dict[str, int] = {str(strategy): 0 for strategy in selected_masks} full_attn_strategy = selected_masks[-1] # Last strategy is full attention print(f"Strategy {full_attn_strategy}, skip first {skip_time_steps} steps ") for t in range(timesteps): for layer_idx in range(layers): for h in range(head_num): if t < skip_time_steps: # First steps use full attention strategy = full_attn_strategy else: # Get losses for this head across all strategies head_losses = [] for strategy in selected_masks[:-1]: # Exclude full attention head_losses.append( averaged_results[loss_type][str(strategy)][t][layer_idx][h] ) # Find which strategy gives minimum loss best_strategy_idx = np.argmin(head_losses) strategy = selected_masks[best_strategy_idx] best_mask_strategy[f"{t}_{layer_idx}_{h}"] = strategy # Calculate sparsity nums = strategy # strategy is already a list of numbers total_tokens += ( nums[0] * nums[1] * nums[2] ) # masked tokens for chosen strategy total_length += ( full_attn_strategy[0] * full_attn_strategy[1] * full_attn_strategy[2] ) # Count strategy usage strategy_counts[str(strategy)] += 1 overall_sparsity = 1 - total_tokens / total_length return best_mask_strategy, overall_sparsity, strategy_counts def save_mask_search_results( mask_search_final_result: list[dict[str, list[float]]], prompt: str, mask_strategies: list[str], output_dir: str = "output/mask_search_result/", ) -> str | None: if not mask_search_final_result: print("No mask search results to save") return None # Create result dictionary with defaultdict for nested lists mask_search_dict: dict[str, dict[str, list[list[float]]]] = { "L2_loss": defaultdict(list), "L1_loss": defaultdict(list), } mask_selected = list(range(len(mask_strategies))) selected_masks: list[list[int]] = [] for index in mask_selected: mask = mask_strategies[index] masks_list = [int(x) for x in mask.split(",")] selected_masks.append(masks_list) # Process each mask strategy for i, mask_strategy in enumerate(selected_masks): mask_strategy_str = str(mask_strategy) # Process L2 loss step_results: list[list[float]] = [] for step_data in mask_search_final_result: if isinstance(step_data, dict) and "L2_loss" in step_data: layer_losses = [float(loss) for loss in step_data["L2_loss"]] step_results.append(layer_losses) mask_search_dict["L2_loss"][mask_strategy_str] = step_results step_results = [] for step_data in mask_search_final_result: if isinstance(step_data, dict) and "L1_loss" in step_data: layer_losses = [float(loss) for loss in step_data["L1_loss"]] step_results.append(layer_losses) mask_search_dict["L1_loss"][mask_strategy_str] = step_results # Create the output directory if it doesn't exist os.makedirs(output_dir, exist_ok=True) # Create a filename based on the first 20 characters of the prompt filename = prompt[:50].replace(" ", "_") filepath = os.path.join(output_dir, f"mask_search_{filename}.json") # Save the results to a JSON file with open(filepath, "w") as f: json.dump(mask_search_dict, f, indent=4) print(f"Successfully saved mask research results to {filepath}") return filepath ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionBackend, AttentionMetadata, AttentionMetadataBuilder, ) from sglang.multimodal_gen.runtime.layers.attention.layer import ( LocalAttention, UlyssesAttention, UlyssesAttention_VSA, USPAttention, ) from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend from sglang.multimodal_gen.runtime.layers.attention.turbo_layer import MinimalA2AAttnOp __all__ = [ "USPAttention", "LocalAttention", "UlyssesAttention", "UlyssesAttention_VSA", "MinimalA2AAttnOp", "AttentionBackend", "AttentionMetadata", "AttentionMetadataBuilder", # "AttentionState", "get_attn_backend", ] ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import aiter import torch from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum class AITerBackend(AttentionBackend): """ Backend for AITemplate attention implementation. """ @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.AITER @staticmethod def get_impl_cls() -> type["AITerImpl"]: return AITerImpl @staticmethod def get_metadata_cls() -> type["AttentionMetadata"]: # AITer backend does not require special metadata. return AttentionMetadata @staticmethod def get_builder_cls() -> type["AttentionMetadataBuilder"]: raise NotImplementedError("AITer backend does not have a metadata builder.") class AITerImpl(AttentionImpl): """ Implementation of attention using AITemplate. """ def __init__( self, num_heads: int, head_size: int, softmax_scale: float, causal: bool = False, num_kv_heads: int | None = None, prefix: str = "", dropout_p: float = 0.0, **extra_impl_args, ) -> None: if num_kv_heads is not None and num_kv_heads != num_heads: raise NotImplementedError( "AITer backend does not support Grouped Query Attention yet." ) self.causal = causal self.dropout_p = dropout_p def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: """ Performs attention using aiter.flash_attn_func. Args: query: Query tensor of shape [batch_size, num_heads, seq_len, head_dim] key: Key tensor of shape [batch_size, num_heads, seq_len, head_dim] value: Value tensor of shape [batch_size, num_heads, seq_len, head_dim] attn_metadata: Metadata for the attention operation (unused). Returns: Output tensor of shape [batch_size, num_heads, seq_len, head_dim] """ # aiter.flash_attn_func expects tensors in [B, H, S, D] layout, # which is what ring_attn provides. output, _ = aiter.flash_attn_func( query, key, value, dropout_p=self.dropout_p, causal=self.causal, return_attn_probs=False, return_lse=True, ) return output ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py ================================================ # SPDX-License-Identifier: Apache-2.0 import torch from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum class AITERSageBackend(AttentionBackend): @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.AITER_SAGE @staticmethod def get_impl_cls() -> type["AITERSageImpl"]: return AITERSageImpl @staticmethod def get_metadata_cls() -> type["AttentionMetadata"]: # AITER Sage backend does not require special metadata. return AttentionMetadata @staticmethod def get_builder_cls() -> type["AttentionMetadataBuilder"]: raise NotImplementedError( "AITER Sage backend does not have a metadata builder." ) class AITERSageImpl(AttentionImpl): def __init__( self, num_heads: int, head_size: int, softmax_scale: float, causal: bool = False, num_kv_heads: int | None = None, prefix: str = "", dropout_p: float = 0.0, **extra_impl_args, ) -> None: try: from aiter.ops.triton.attention.fav3_sage import fav3_sage_wrapper_func self.aiter_sage_attn_fn = fav3_sage_wrapper_func except ImportError: raise ImportError( "AITER Sage attention is not available, please update AITER version." ) def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: """ Performs attention using aiter sage backend. Args: query: Query tensor of shape [batch_size, seq_len, head_num, head_dim] key: Key tensor of shape [batch_size, seq_len, head_num, head_dim] value: Value tensor of shape [batch_size, seq_len, head_num, head_dim] attn_metadata: Metadata for the attention operation (unused). Returns: Output tensor of shape [batch_size, seq_len, head_num, head_dim] """ output = self.aiter_sage_attn_fn(query, key, value) return output ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/backends/abstract.py from abc import ABC, abstractmethod from dataclasses import dataclass, fields from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar if TYPE_CHECKING: pass import torch from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum class AttentionBackend(ABC): """Abstract class for attention backends.""" # For some attention backends, we allocate an output tensor before # calling the custom op. When piecewise cudagraph is enabled, this # makes sure the output tensor is allocated inside the cudagraph. accept_output_buffer: bool = False @staticmethod @abstractmethod def get_enum() -> AttentionBackendEnum: raise NotImplementedError @staticmethod @abstractmethod def get_impl_cls() -> type["AttentionImpl"]: raise NotImplementedError @staticmethod @abstractmethod def get_metadata_cls() -> type["AttentionMetadata"]: raise NotImplementedError # @staticmethod # @abstractmethod # def get_state_cls() -> Type["AttentionState"]: # raise NotImplementedError # @classmethod # def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata": # return cls.get_metadata_cls()(*args, **kwargs) @staticmethod @abstractmethod def get_builder_cls() -> type["AttentionMetadataBuilder"]: return None @dataclass class AttentionMetadata: """Attention metadata for prefill and decode batched together.""" # Current step of diffusion process current_timestep: int def asdict_zerocopy(self, skip_fields: set[str] | None = None) -> dict[str, Any]: """Similar to dataclasses.asdict, but avoids deepcopying.""" if skip_fields is None: skip_fields = set() # Note that if we add dataclasses as fields, they will need # similar handling. return { field.name: getattr(self, field.name) for field in fields(self) if field.name not in skip_fields } T = TypeVar("T", bound=AttentionMetadata) class AttentionMetadataBuilder(ABC, Generic[T]): """Abstract class for attention metadata builders.""" @abstractmethod def __init__(self) -> None: """Create the builder, remember some configuration and parameters.""" raise NotImplementedError @abstractmethod def prepare(self) -> None: """Prepare for one batch.""" raise NotImplementedError @abstractmethod def build( self, **kwargs: dict[str, Any], ) -> AttentionMetadata: """Build attention metadata with on-device tensors.""" raise NotImplementedError class AttentionLayer(Protocol): _k_scale: torch.Tensor _v_scale: torch.Tensor _k_scale_float: float _v_scale_float: float def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, kv_cache: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: ... class AttentionImpl(ABC, Generic[T]): @abstractmethod def __init__( self, num_heads: int, head_size: int, softmax_scale: float, causal: bool = False, num_kv_heads: int | None = None, prefix: str = "", **extra_impl_args, ) -> None: raise NotImplementedError def preprocess_qkv(self, qkv: torch.Tensor, attn_metadata: T) -> torch.Tensor: """Preprocess QKV tensor before performing attention operation. Default implementation returns the tensor unchanged. Subclasses can override this to implement custom preprocessing like reshaping, tiling, scaling, or other transformations. Called AFTER all_to_all for distributed attention """ return qkv def postprocess_output( self, output: torch.Tensor, attn_metadata: T, ) -> torch.Tensor: """Postprocess the output tensor after the attention operation. Default implementation returns the tensor unchanged. Subclasses can override this to implement custom postprocessing like untiling, scaling, or other transformations. Called BEFORE all_to_all for distributed attention """ return output @abstractmethod def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: T, ) -> torch.Tensor: raise NotImplementedError ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass from typing import Any, List, Optional, Tuple import torch from sglang.multimodal_gen.runtime.layers.utils import register_custom_op from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, ) try: from sgl_kernel.flash_attn import flash_attn_varlen_func from sglang.jit_kernel.flash_attention_v4 import ( flash_attn_varlen_func as flash_attn_varlen_func_fa4, ) def flash_attn_func(*args, ver: int = 3, **kwargs): if ver == 4: return flash_attn_varlen_func_fa4(*args, **kwargs) return flash_attn_varlen_func(*args, **kwargs) except ImportError as e: raise e def maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]: return x.contiguous() if x is not None and x.stride(-1) != 1 else x # ----------------------------- # Fake implementations for schema / tracing # custom op schema requires FIXED return structure. # We provide TWO ops: # 1) out-only op: always returns Tensor # 2) out+lse op: always returns Tuple[Tensor, Tensor] # ----------------------------- def flash_attn_varlen_func_fake_out( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_k: Optional[torch.Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_k: Optional[int] = None, seqused_q: Optional[torch.Tensor] = None, seqused_k: Optional[torch.Tensor] = None, page_table: Optional[torch.Tensor] = None, softmax_scale: Optional[float] = None, causal: bool = False, qv: Optional[torch.Tensor] = None, q_descale: Optional[torch.Tensor] = None, k_descale: Optional[torch.Tensor] = None, v_descale: Optional[torch.Tensor] = None, window_size: Optional[List[int]] = None, attention_chunk: int = 0, softcap: float = 0.0, num_splits: int = 1, pack_gqa: Optional[bool] = None, sm_margin: int = 0, return_softmax_lse: bool = False, sinks: Optional[torch.Tensor] = None, ver: int = 4, ) -> torch.Tensor: assert ver == 4, "only support flash attention v4" q, k, v = [maybe_contiguous(t) for t in (q, k, v)] num_head, head_dim = q.shape[-2:] if cu_seqlens_q is None: batch_size, seqlen_q = q.shape[:2] else: batch_size = cu_seqlens_q.shape[0] - 1 seqlen_q = None head_dim_v = v.shape[-1] if cu_seqlens_q is not None: assert cu_seqlens_q.shape == ( batch_size + 1, ), "cu_seqlens_q must have shape (batch_size + 1,)" assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32" assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous" assert q.dtype in [ torch.float16, torch.bfloat16, ], "inputs must be float16 or bfloat16" assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype" assert head_dim <= 256, "head_dim must be less than or equal to 256" alignment = 16 // q.element_size() assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}" q_batch_seqlen_shape = ( (batch_size, seqlen_q) if cu_seqlens_q is None else (q.shape[0],) ) out = q.new_empty(*q_batch_seqlen_shape, num_head, head_dim_v) return out def flash_attn_varlen_func_fake_out_lse( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_k: Optional[torch.Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_k: Optional[int] = None, seqused_q: Optional[torch.Tensor] = None, seqused_k: Optional[torch.Tensor] = None, page_table: Optional[torch.Tensor] = None, softmax_scale: Optional[float] = None, causal: bool = False, qv: Optional[torch.Tensor] = None, q_descale: Optional[torch.Tensor] = None, k_descale: Optional[torch.Tensor] = None, v_descale: Optional[torch.Tensor] = None, window_size: Optional[List[int]] = None, attention_chunk: int = 0, softcap: float = 0.0, num_splits: int = 1, pack_gqa: Optional[bool] = None, sm_margin: int = 0, return_softmax_lse: bool = True, sinks: Optional[torch.Tensor] = None, ver: int = 4, ) -> Tuple[torch.Tensor, torch.Tensor]: assert ver == 4, "only support flash attention v4" q, k, v = [maybe_contiguous(t) for t in (q, k, v)] num_head, head_dim = q.shape[-2:] if cu_seqlens_q is None: batch_size, seqlen_q = q.shape[:2] total_q = batch_size * seqlen_q else: batch_size = cu_seqlens_q.shape[0] - 1 seqlen_q = None total_q = q.shape[0] head_dim_v = v.shape[-1] if cu_seqlens_q is not None: assert cu_seqlens_q.shape == ( batch_size + 1, ), "cu_seqlens_q must have shape (batch_size + 1,)" assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32" assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous" assert q.dtype in [ torch.float16, torch.bfloat16, ], "inputs must be float16 or bfloat16" assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype" assert head_dim <= 256, "head_dim must be less than or equal to 256" alignment = 16 // q.element_size() assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}" q_batch_seqlen_shape = ( (batch_size, seqlen_q) if cu_seqlens_q is None else (total_q,) ) lse_shape = ( (batch_size, num_head, seqlen_q) if cu_seqlens_q is None else (num_head, total_q) ) out = q.new_empty(*q_batch_seqlen_shape, num_head, head_dim_v) lse = q.new_empty(lse_shape, dtype=torch.float32) return out, lse # ----------------------------- # Registered custom ops # NOTE: fixed return schemas to avoid: # "Object of type 'Tensor' is not an instance of 'sequence'" # ----------------------------- @register_custom_op(fake_impl=flash_attn_varlen_func_fake_out) def flash_attn_varlen_func_op( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_k: Optional[torch.Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_k: Optional[int] = None, seqused_q: Optional[torch.Tensor] = None, seqused_k: Optional[torch.Tensor] = None, page_table: Optional[torch.Tensor] = None, softmax_scale: Optional[float] = None, causal: bool = False, qv: Optional[torch.Tensor] = None, q_descale: Optional[torch.Tensor] = None, k_descale: Optional[torch.Tensor] = None, v_descale: Optional[torch.Tensor] = None, window_size: Optional[List[int]] = None, attention_chunk: int = 0, softcap: float = 0.0, num_splits: int = 1, pack_gqa: Optional[bool] = None, sm_margin: int = 0, return_softmax_lse: bool = False, sinks: Optional[torch.Tensor] = None, ver: int = 4, ) -> torch.Tensor: if window_size is None: window_size = [-1, -1] if return_softmax_lse: raise ValueError( "flash_attn_varlen_func_op is out-only op; return_softmax_lse must be False. " "Use flash_attn_varlen_func_op_lse for (out, lse)." ) return flash_attn_func( q, k, v, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, seqused_q=seqused_q, seqused_k=seqused_k, page_table=page_table, softmax_scale=softmax_scale, causal=causal, qv=qv, q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, window_size=tuple(window_size), attention_chunk=attention_chunk, softcap=softcap, num_splits=num_splits, pack_gqa=pack_gqa, sm_margin=sm_margin, return_softmax_lse=False, sinks=sinks, ver=ver, ) @register_custom_op(fake_impl=flash_attn_varlen_func_fake_out_lse) def flash_attn_varlen_func_op_lse( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_k: Optional[torch.Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_k: Optional[int] = None, seqused_q: Optional[torch.Tensor] = None, seqused_k: Optional[torch.Tensor] = None, page_table: Optional[torch.Tensor] = None, softmax_scale: Optional[float] = None, causal: bool = False, qv: Optional[torch.Tensor] = None, q_descale: Optional[torch.Tensor] = None, k_descale: Optional[torch.Tensor] = None, v_descale: Optional[torch.Tensor] = None, window_size: Optional[List[int]] = None, attention_chunk: int = 0, softcap: float = 0.0, num_splits: int = 1, pack_gqa: Optional[bool] = None, sm_margin: int = 0, return_softmax_lse: bool = True, sinks: Optional[torch.Tensor] = None, ver: int = 4, ) -> Tuple[torch.Tensor, torch.Tensor]: if window_size is None: window_size = [-1, -1] if not return_softmax_lse: raise ValueError( "flash_attn_varlen_func_op_lse is out+lse op; return_softmax_lse must be True. " "Use flash_attn_varlen_func_op for out-only." ) return flash_attn_func( q, k, v, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, seqused_q=seqused_q, seqused_k=seqused_k, page_table=page_table, softmax_scale=softmax_scale, causal=causal, qv=qv, q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, window_size=tuple(window_size), attention_chunk=attention_chunk, softcap=softcap, num_splits=num_splits, pack_gqa=pack_gqa, sm_margin=sm_margin, return_softmax_lse=True, sinks=sinks, ver=ver, ) from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder, ) fa_ver = 3 def set_fa_ver(ver: int) -> None: global fa_ver fa_ver = ver @dataclass class FlashAttentionMetadata: # Sequence lengths for the forward batch # Maximum sequence length for query max_seqlen_q: int = 1 # Maximum sequence length for key max_seqlen_k: int = 0 # Cumulative sequence lengths for query cu_seqlens_q: torch.Tensor = None # Cumulative sequence lengths for key cu_seqlens_k: torch.Tensor = None class FlashAttentionMetadataBuilder(AttentionMetadataBuilder): def __init__(self) -> None: pass def prepare(self) -> None: pass def build( # type: ignore self, raw_latent_shape=list, **kwargs: dict[str, Any], ) -> FlashAttentionMetadata: # TODO: put empty values here to be set at first-run, since the q_len calculation can be complicated return FlashAttentionMetadata(max_seqlen_q=None, max_seqlen_k=None) class FlashAttentionBackend(AttentionBackend): accept_output_buffer: bool = True @staticmethod def get_supported_head_sizes() -> list[int]: return [32, 64, 96, 128, 160, 192, 224, 256] @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.FA @staticmethod def get_impl_cls() -> type["FlashAttentionImpl"]: return FlashAttentionImpl @staticmethod def get_metadata_cls() -> type["AttentionMetadata"]: raise NotImplementedError @staticmethod def get_builder_cls() -> type["AttentionMetadataBuilder"]: return FlashAttentionMetadataBuilder class FlashAttentionImpl(AttentionImpl): def __init__( self, num_heads: int, head_size: int, causal: bool, softmax_scale: float, num_kv_heads: int | None = None, prefix: str = "", **extra_impl_args, ) -> None: self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.head_size = head_size self.causal = causal self.softmax_scale = softmax_scale self.attention_metadata = FlashAttentionMetadata() def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata = None, *, return_softmax_lse: bool = False, ): attn_metadata: FlashAttentionMetadata = get_forward_context().attn_metadata if attn_metadata is not None and attn_metadata.max_seqlen_q is None: attn_metadata.max_seqlen_q = query.shape[1] attn_metadata.max_seqlen_k = key.shape[1] max_seqlen_q = attn_metadata.max_seqlen_q max_seqlen_k = attn_metadata.max_seqlen_k else: max_seqlen_q = query.shape[1] max_seqlen_k = key.shape[1] # FA version selection: # - fa_ver == 3: call python function (can return Tensor or (Tensor, Tensor) depending on flag) # - fa_ver == 4: call custom ops with FIXED return schema if fa_ver == 3: flash_attn_op = flash_attn_func output = flash_attn_op( q=query, k=key, v=value, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, softmax_scale=self.softmax_scale, causal=self.causal, return_softmax_lse=return_softmax_lse, ver=fa_ver, ) return output if fa_ver == 4: if return_softmax_lse: out_tensor, softmax_lse = flash_attn_varlen_func_op_lse( q=query, k=key, v=value, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, softmax_scale=self.softmax_scale, causal=self.causal, return_softmax_lse=True, ver=fa_ver, ) return out_tensor, softmax_lse out_tensor = flash_attn_varlen_func_op( q=query, k=key, v=value, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, softmax_scale=self.softmax_scale, causal=self.causal, return_softmax_lse=False, ver=fa_ver, ) return out_tensor raise ValueError(f"flash attention version {fa_ver} is not supported.") ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import torch from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder, ) from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( flash_attn_func, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class FlashAttention2Backend(AttentionBackend): accept_output_buffer: bool = True @staticmethod def get_supported_head_sizes() -> list[int]: return [32, 64, 96, 128, 160, 192, 224, 256] @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.FA2 @staticmethod def get_impl_cls() -> type["FlashAttention2Impl"]: return FlashAttention2Impl @staticmethod def get_metadata_cls() -> type["AttentionMetadata"]: raise NotImplementedError @staticmethod def get_builder_cls() -> type["AttentionMetadataBuilder"]: raise NotImplementedError class FlashAttention2Impl(AttentionImpl): def __init__( self, num_heads: int, head_size: int, causal: bool, softmax_scale: float, num_kv_heads: int | None = None, prefix: str = "", **extra_impl_args, ) -> None: self.causal = causal self.softmax_scale = softmax_scale def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata, ): output = flash_attn_func( q=query, # type: ignore[no-untyped-call] k=key, v=value, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=None, max_seqlen_k=None, softmax_scale=self.softmax_scale, causal=self.causal, ) return output ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import torch from sageattention import sageattn from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( # FlashAttentionMetadata, AttentionBackend, AttentionImpl, AttentionMetadata, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class SageAttentionBackend(AttentionBackend): accept_output_buffer: bool = True @staticmethod def get_supported_head_sizes() -> list[int]: return [32, 64, 96, 128, 160, 192, 224, 256] @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.SAGE_ATTN @staticmethod def get_impl_cls() -> type["SageAttentionImpl"]: return SageAttentionImpl class SageAttentionImpl(AttentionImpl): def __init__( self, num_heads: int, head_size: int, causal: bool, softmax_scale: float, num_kv_heads: int | None = None, prefix: str = "", **extra_impl_args, ) -> None: self.causal = causal self.softmax_scale = softmax_scale self.dropout = extra_impl_args.get("dropout_p", 0.0) def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata, *, return_softmax_lse: bool = False, ) -> torch.Tensor: output = sageattn( query, key, value, # since input is (batch_size, seq_len, head_num, head_dim) tensor_layout="NHD", is_causal=self.causal, sm_scale=self.softmax_scale, return_lse=return_softmax_lse, ) if return_softmax_lse: output, softmax_lse = output return output, softmax_lse return output ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import torch import torch.nn.functional as F from sageattn3 import sageattn3_blackwell from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionBackend, AttentionImpl, AttentionMetadata, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class SageAttention3Backend(AttentionBackend): accept_output_buffer: bool = True @staticmethod def get_supported_head_sizes() -> list[int]: return [64, 128, 256] @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.SAGE_ATTN_3 @staticmethod def get_impl_cls() -> type["SageAttention3Impl"]: return SageAttention3Impl @staticmethod def get_metadata_cls() -> type["AttentionMetadata"]: raise NotImplementedError class SageAttention3Impl(AttentionImpl): _warned_gqa_fallback_global: bool = False def __init__( self, num_heads: int, head_size: int, causal: bool, softmax_scale: float, num_kv_heads: int | None = None, prefix: str = "", **extra_impl_args, ) -> None: self.causal = causal self.softmax_scale = softmax_scale self.dropout = extra_impl_args.get("dropout_p", 0.0) def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: query = query.transpose(1, 2) key = key.transpose(1, 2) value = value.transpose(1, 2) # SageAttention3's Blackwell kernel assumes MHA (Hq == Hkv). For GQA/MQA # (Hq != Hkv), fall back to torch SDPA which supports GQA. if key.shape[1] != query.shape[1]: if query.shape[1] % key.shape[1] != 0: raise ValueError( "GQA/MQA requires query heads to be a multiple of KV heads, " f"got q_heads={query.shape[1]} and kv_heads={key.shape[1]}" ) if not type(self)._warned_gqa_fallback_global: logger.warning( "SageAttention3 does not support GQA/MQA (Hq != Hkv); falling back to torch SDPA." ) type(self)._warned_gqa_fallback_global = True output = F.scaled_dot_product_attention( query, key, value, is_causal=self.causal, dropout_p=self.dropout, scale=self.softmax_scale, enable_gqa=True, ) else: output = sageattn3_blackwell(query, key, value, is_causal=self.causal) output = output.transpose(1, 2) return output ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import torch from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( # FlashAttentionMetadata, AttentionBackend, AttentionImpl, AttentionMetadata, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class SDPABackend(AttentionBackend): accept_output_buffer: bool = True @staticmethod def get_supported_head_sizes() -> list[int]: return [32, 64, 96, 128, 160, 192, 224, 256] @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.TORCH_SDPA @staticmethod def get_impl_cls() -> type["SDPAImpl"]: return SDPAImpl # @staticmethod # def get_metadata_cls() -> Type["AttentionMetadata"]: # return FlashAttentionMetadata class SDPAImpl(AttentionImpl): def __init__( self, num_heads: int, head_size: int, causal: bool, softmax_scale: float, num_kv_heads: int | None = None, prefix: str = "", **extra_impl_args, ) -> None: self.causal = causal self.softmax_scale = softmax_scale self.dropout = extra_impl_args.get("dropout_p", 0.0) def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: # transpose to bs, heads, seq_len, head_dim query = query.transpose(1, 2) key = key.transpose(1, 2) value = value.transpose(1, 2) attn_kwargs = { "attn_mask": None, "dropout_p": self.dropout, "is_causal": self.causal, "scale": self.softmax_scale, } if query.shape[1] != key.shape[1]: attn_kwargs["enable_gqa"] = True output = torch.nn.functional.scaled_dot_product_attention( query, key, value, **attn_kwargs ) output = output.transpose(1, 2) return output ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import json from dataclasses import dataclass from typing import Any import torch from einops import rearrange from sglang.multimodal_gen.runtime.distributed import get_sp_group from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder, ) from sglang.multimodal_gen.runtime.managers.forward_context import ( ForwardContext, get_forward_context, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.server_args import get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import dict_to_3d_list try: from st_attn import sliding_tile_attention st_attn_backend_available = True except Exception: st_attn_backend_available = False logger = init_logger(__name__) class RangeDict(dict): def __getitem__(self, item: int) -> str: for key in self.keys(): if isinstance(key, tuple): low, high = key if low <= item <= high: return str(super().__getitem__(key)) elif key == item: return str(super().__getitem__(key)) raise KeyError(f"seq_len {item} not supported for STA") class SlidingTileAttentionBackend(AttentionBackend): accept_output_buffer: bool = True @staticmethod def get_supported_head_sizes() -> list[int]: # TODO(will-refactor): check this return [32, 64, 96, 128, 160, 192, 224, 256] @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.SLIDING_TILE_ATTN @staticmethod def get_impl_cls() -> type["SlidingTileAttentionImpl"]: return SlidingTileAttentionImpl @staticmethod def get_metadata_cls() -> type["SlidingTileAttentionMetadata"]: return SlidingTileAttentionMetadata @staticmethod def get_builder_cls() -> type["SlidingTileAttentionMetadataBuilder"]: return SlidingTileAttentionMetadataBuilder @dataclass class SlidingTileAttentionMetadata(AttentionMetadata): current_timestep: int STA_param: list[ list[Any] ] # each timestep with one metadata, shape [num_layers, num_heads] class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder): def __init__(self): pass def prepare(self): pass def build( # type: ignore self, STA_param: list[list[Any]], current_timestep: int, **kwargs: dict[str, Any], ) -> SlidingTileAttentionMetadata: param = STA_param if param is None: return SlidingTileAttentionMetadata( current_timestep=current_timestep, STA_param=[] ) return SlidingTileAttentionMetadata( current_timestep=current_timestep, STA_param=param[current_timestep] ) class SlidingTileAttentionImpl(AttentionImpl): def __init__( self, num_heads: int, head_size: int, causal: bool, softmax_scale: float, num_kv_heads: int | None = None, prefix: str = "", **extra_impl_args, ) -> None: if not st_attn_backend_available: raise ValueError("st attn not supported") # TODO(will-refactor): for now this is the mask strategy, but maybe we should # have a more general config for STA? mask_strategy_file_path = ( get_global_server_args().attention_backend_config.mask_strategy_file_path ) if mask_strategy_file_path is None: raise ValueError("SGLANG_DIFFUSION_ATTENTION_CONFIG is not set") # TODO(kevin): get mask strategy for different STA modes with open(mask_strategy_file_path) as f: mask_strategy = json.load(f) self.mask_strategy = dict_to_3d_list(mask_strategy) self.prefix = prefix sp_group = get_sp_group() self.sp_size = sp_group.world_size # STA config self.STA_base_tile_size = [6, 8, 8] self.dit_seq_shape_mapping = RangeDict( { (115200, 115456): "30x48x80", 82944: "36x48x48", 69120: "18x48x80", } ) self.full_window_mapping = { "30x48x80": [5, 6, 10], "36x48x48": [6, 6, 6], "18x48x80": [3, 6, 10], } def tile(self, x: torch.Tensor) -> torch.Tensor: return rearrange( x, "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d", n_t=self.full_window_size[0], n_h=self.full_window_size[1], n_w=self.full_window_size[2], ts_t=self.STA_base_tile_size[0], ts_h=self.STA_base_tile_size[1], ts_w=self.STA_base_tile_size[2], ) def untile(self, x: torch.Tensor) -> torch.Tensor: x = rearrange( x, "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d", n_t=self.full_window_size[0], n_h=self.full_window_size[1], n_w=self.full_window_size[2], ts_t=self.STA_base_tile_size[0], ts_h=self.STA_base_tile_size[1], ts_w=self.STA_base_tile_size[2], ) return x def preprocess_qkv( self, qkv: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: img_sequence_length = qkv.shape[1] self.dit_seq_shape_str = self.dit_seq_shape_mapping[img_sequence_length] self.full_window_size = self.full_window_mapping[self.dit_seq_shape_str] self.dit_seq_shape_int = list(map(int, self.dit_seq_shape_str.split("x"))) self.img_seq_length = ( self.dit_seq_shape_int[0] * self.dit_seq_shape_int[1] * self.dit_seq_shape_int[2] ) return self.tile(qkv) def postprocess_output( self, output: torch.Tensor, attn_metadata: SlidingTileAttentionMetadata, ) -> torch.Tensor: return self.untile(output) def forward( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, attn_metadata: SlidingTileAttentionMetadata, ) -> torch.Tensor: if self.mask_strategy is None: raise ValueError("mask_strategy cannot be None for SlidingTileAttention") if self.mask_strategy[0] is None: raise ValueError("mask_strategy[0] cannot be None for SlidingTileAttention") timestep = attn_metadata.current_timestep forward_context: ForwardContext = get_forward_context() forward_batch = forward_context.forward_batch if forward_batch is None: raise ValueError("forward_batch cannot be None") # pattern:'.double_blocks.0.attn.impl' or '.single_blocks.0.attn.impl' layer_idx = int(self.prefix.split(".")[-3]) if attn_metadata.STA_param is None or len(attn_metadata.STA_param) <= layer_idx: raise ValueError("Invalid STA_param") STA_param = attn_metadata.STA_param[layer_idx] text_length = q.shape[1] - self.img_seq_length has_text = text_length > 0 query = q.transpose(1, 2).contiguous() key = k.transpose(1, 2).contiguous() value = v.transpose(1, 2).contiguous() head_num = query.size(1) sp_group = get_sp_group() current_rank = sp_group.rank_in_group start_head = current_rank * head_num # searching or tuning mode if len(STA_param) < head_num * sp_group.world_size: sparse_attn_hidden_states_all = [] full_mask_window = STA_param[-1] for window_size in STA_param[:-1]: sparse_hidden_states = sliding_tile_attention( query, key, value, [window_size] * head_num, text_length, has_text, self.dit_seq_shape_str, ).transpose(1, 2) sparse_attn_hidden_states_all.append(sparse_hidden_states) hidden_states = sliding_tile_attention( query, key, value, [full_mask_window] * head_num, text_length, has_text, self.dit_seq_shape_str, ).transpose(1, 2) attn_L2_loss = [] attn_L1_loss = [] # average loss across all heads for sparse_attn_hidden_states in sparse_attn_hidden_states_all: # L2 loss attn_L2_loss_ = ( torch.mean( (sparse_attn_hidden_states.float() - hidden_states.float()) ** 2, dim=[0, 1, 3], ) .cpu() .numpy() ) attn_L2_loss_ = [round(float(x), 6) for x in attn_L2_loss_] attn_L2_loss.append(attn_L2_loss_) # L1 loss attn_L1_loss_ = ( torch.mean( torch.abs( sparse_attn_hidden_states.float() - hidden_states.float() ), dim=[0, 1, 3], ) .cpu() .numpy() ) attn_L1_loss_ = [round(float(x), 6) for x in attn_L1_loss_] attn_L1_loss.append(attn_L1_loss_) layer_loss_save = {"L2_loss": attn_L2_loss, "L1_loss": attn_L1_loss} if forward_batch.is_cfg_negative: if forward_batch.mask_search_final_result_neg is not None: forward_batch.mask_search_final_result_neg[timestep].append( layer_loss_save ) else: if forward_batch.mask_search_final_result_pos is not None: forward_batch.mask_search_final_result_pos[timestep].append( layer_loss_save ) else: windows = [STA_param[head_idx + start_head] for head_idx in range(head_num)] hidden_states = sliding_tile_attention( query, key, value, windows, text_length, has_text, self.dit_seq_shape_str, ).transpose(1, 2) return hidden_states ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py ================================================ """ Copyright (c) 2025 by SLA team. Licensed under the Apache License, Version 2.0 (the "License"); This implementation is adapted from: from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/core.py and https://github.com/thu-ml/SLA/blob/main/SageSLA/core.py Citation (please cite if you use this code): @article{zhang2025sla, title={SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attention}, author={Jintao Zhang and Haoxu Wang and Kai Jiang and Shuo Yang and Kaiwen Zheng and Haocheng Xi and Ziteng Wang and Hongzhou Zhu and Min Zhao and Ion Stoica and Joseph E. Gonzalez and Jun Zhu and Jianfei Chen}, journal={arXiv preprint arXiv:2509.24006}, year={2025} } """ from collections.abc import Callable from dataclasses import dataclass from typing import Any import torch import torch.nn as nn import torch.nn.functional as F import triton import triton.language as tl from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # ==================================SLA Functions=================================== def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64): arg_k = k - torch.mean( k, dim=-2, keepdim=True ) # smooth-k technique in SageAttention pooled_qblocks = mean_pool(q, BLKQ) pooled_kblocks = mean_pool(arg_k, BLKK) pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2) K = pooled_score.shape[-1] topk = min(K, int(topk_ratio * K)) lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8) sparse_map.scatter_(-1, lut, 1) return sparse_map, lut, topk def mean_pool(x, BLK): assert x.is_contiguous() B, H, L, D = x.shape L_BLOCKS = (L + BLK - 1) // BLK x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=x.dtype) grid = (L_BLOCKS, B * H) compress_kernel[grid](x, x_mean, L, D, BLK) return x_mean @triton.jit def compress_kernel( X, XM, L: tl.constexpr, D: tl.constexpr, BLOCK_L: tl.constexpr, ): idx_l = tl.program_id(0) idx_bh = tl.program_id(1) offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L) offs_d = tl.arange(0, D) x_offset = idx_bh * L * D xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D x = tl.load( X + x_offset + offs_l[:, None] * D + offs_d[None, :], mask=offs_l[:, None] < L ) nx = min(BLOCK_L, L - idx_l * BLOCK_L) x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty)) @triton.jit def _attn_fwd( Q, K, V, qk_scale: tl.constexpr, topk: tl.constexpr, LUT, LSE, OS, L: tl.constexpr, M_BLOCKS: tl.constexpr, D: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, ): idx_m = tl.program_id(0).to(tl.int64) idx_bh = tl.program_id(1).to(tl.int64) qkv_offset = idx_bh * L * D lut_offset = (idx_bh * M_BLOCKS + idx_m) * topk lse_offset = idx_bh * L offs_m = idx_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = tl.arange(0, BLOCK_N) offs_d = tl.arange(0, D) Q_ptrs = Q + qkv_offset + offs_m[:, None] * D + offs_d[None, :] K_ptrs = K + qkv_offset + offs_n[None, :] * D + offs_d[:, None] V_ptrs = V + qkv_offset + offs_n[:, None] * D + offs_d[None, :] OS_ptrs = OS + qkv_offset + offs_m[:, None] * D + offs_d[None, :] LUT_ptr = LUT + lut_offset LSE_ptrs = LSE + lse_offset + offs_m m_i = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32) l_i = tl.zeros([BLOCK_M], dtype=tl.float32) o_s = tl.zeros([BLOCK_M, D], dtype=tl.float32) q = tl.load(Q_ptrs, mask=offs_m[:, None] < L) for block_idx in tl.range(topk): idx_n = tl.load(LUT_ptr + block_idx) n_mask = offs_n < L - idx_n * BLOCK_N k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[None, :]) qk = tl.dot(q, k) * (qk_scale * 1.4426950408889634) # = 1 / ln(2) if L - idx_n * BLOCK_N < BLOCK_N: qk = tl.where(n_mask[None, :], qk, float("-inf")) v = tl.load(V_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None]) local_m = tl.max(qk, 1) new_m = tl.maximum(m_i, local_m) qk = qk - new_m[:, None] p = tl.math.exp2(qk) l_ij = tl.sum(p, 1) alpha = tl.math.exp2(m_i - new_m) o_s = o_s * alpha[:, None] o_s += tl.dot(p.to(v.dtype), v) l_i = l_i * alpha + l_ij m_i = new_m o_s = o_s / l_i[:, None] tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < L) m_i += tl.math.log2(l_i) tl.store(LSE_ptrs, m_i, mask=offs_m < L) def _get_cuda_arch(device_index: int) -> str: """Get CUDA architecture string for the given device.""" major, minor = torch.cuda.get_device_capability(device_index) return f"sm{major}{minor}" # ==================================SLA Class=================================== class SparseLinearAttentionBackend(AttentionBackend): """Sparse Linear Attention Backend for efficient attention computation.""" accept_output_buffer: bool = True @staticmethod def get_supported_head_sizes() -> list[int]: return [64, 128] @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.SLA_ATTN @staticmethod def get_impl_cls() -> type["SparseLinearAttentionImpl"]: return SparseLinearAttentionImpl @staticmethod def get_metadata_cls() -> type["SparseLinearAttentionMetadata"]: return SparseLinearAttentionMetadata @staticmethod def get_builder_cls() -> type["SparseLinearAttentionMetadataBuilder"]: return SparseLinearAttentionMetadataBuilder @dataclass class SparseLinearAttentionMetadata(AttentionMetadata): """Metadata for Sparse Linear Attention computation.""" # Basic attention parameters current_timestep: int # Sparse attention configuration topk_ratio: float = 0.1 class SparseLinearAttentionMetadataBuilder(AttentionMetadataBuilder): """Builder for SparseLinearAttentionMetadata.""" def __init__(self) -> None: pass def prepare(self) -> None: pass def build( self, current_timestep: int, topk_ratio: float = 0.1, **kwargs: dict[str, Any], ) -> SparseLinearAttentionMetadata: return SparseLinearAttentionMetadata( current_timestep=current_timestep, topk_ratio=topk_ratio, ) class SparseLinearAttentionImpl(AttentionImpl, nn.Module): """Implementation of sparse linear attention for the backend.""" def __init__( self, num_heads: int, head_size: int, causal: bool = False, softmax_scale: float | None = None, num_kv_heads: int | None = None, prefix: str = "", # SLA-specific parameters - matched to TurboDiffusion defaults topk_ratio: float = 0.1, # TurboDiffusion uses topk=0.1 feature_map: str = "softmax", BLKQ: int = 128, # TurboDiffusion uses BLKQ=128 BLKK: int = 64, # TurboDiffusion uses BLKK=64 use_bf16: bool = True, **extra_impl_args, ) -> None: nn.Module.__init__(self) # SLA-specific config self.topk_ratio = topk_ratio self.BLKQ = BLKQ self.BLKK = BLKK self.dtype = torch.bfloat16 if use_bf16 else torch.float16 # Learnable linear projection for combining sparse + linear attention self.proj_l = nn.Linear(head_size, head_size, dtype=torch.float32) # Feature map for linear attention # Type annotation for callables self.feature_map_q: Callable[[torch.Tensor], torch.Tensor] self.feature_map_k: Callable[[torch.Tensor], torch.Tensor] if feature_map == "elu": self.feature_map_q = lambda x: F.elu(x) + 1 self.feature_map_k = lambda x: F.elu(x) + 1 elif feature_map == "relu": self.feature_map_q = F.relu self.feature_map_k = F.relu elif feature_map == "softmax": self.feature_map_q = lambda x: F.softmax(x, dim=-1) self.feature_map_k = lambda x: F.softmax(x, dim=-1) else: raise ValueError(f"Unknown feature map: {feature_map}") self._init_weights() def _init_weights(self) -> None: """Initialize projection weights to zero for residual-like behavior.""" with torch.no_grad(): nn.init.zeros_(self.proj_l.weight) nn.init.zeros_(self.proj_l.bias) # type: ignore[arg-type] def _calc_linear_attention_with_torch(self, q, k, v): kv = torch.matmul(k.transpose(-1, -2), v) k_sum = torch.sum(k, dim=-2, keepdim=True) return torch.matmul(q, kv) / (1e-5 + torch.matmul(q, k_sum.transpose(-1, -2))) def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: SparseLinearAttentionMetadata = None, ) -> torch.Tensor: """Forward pass for sparse linear attention. Args: query: query tensor of shape (B, H, L, D) key: key tensor of shape (B, H, L, D) value: value tensor of shape (B, H, L, D) attn_metadata: attention metadata containing configuration Returns: output tensor of shape (B, H, L, D) """ dtype = query.dtype # Transpose for computation query = query.transpose(1, 2).contiguous() key = key.transpose(1, 2).contiguous() value = value.transpose(1, 2).contiguous() # Get sparse attention map sparse_map, lut, real_topk = get_block_map( query, key, topk_ratio=self.topk_ratio, BLKQ=self.BLKQ, BLKK=self.BLKK ) # Convert to computation dtype query = query.to(self.dtype) key = key.to(self.dtype) value = value.to(self.dtype) # Sparse attention computation o_s = _attention.apply( query, key, value, sparse_map, lut, real_topk, self.BLKQ, self.BLKK ) # Apply feature maps query = self.feature_map_q(query).contiguous().to(self.dtype) # c_q key = self.feature_map_k(key).contiguous().to(self.dtype) # c_k # Linear attention computation o_l = self._calc_linear_attention_with_torch(query, key, value) # Apply projection and combine results with torch.amp.autocast("cuda", dtype=self.dtype): o_l = self.proj_l(o_l) # Combine sparse and linear attention output = (o_s + o_l).to(dtype).transpose(1, 2) return output class _attention(torch.autograd.Function): @staticmethod def forward(ctx, q, k, v, k_block_id, lut, topk, BLOCK_M, BLOCK_N, qk_scale=None): assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() assert k_block_id.is_contiguous() and lut.is_contiguous() # We recommend the following two settings assert BLOCK_M == 64 or BLOCK_M == 128 assert BLOCK_N == 64 B, H, L, D = q.shape if qk_scale is None: qk_scale = D**-0.5 M_BLOCKS = triton.cdiv(L, BLOCK_M) o_s = torch.empty_like(v) lse = torch.empty(q.shape[:-1], device=q.device, dtype=torch.float32) grid = (M_BLOCKS, B * H) _attn_fwd[grid]( q, k, v, qk_scale, topk, lut, lse, o_s, L, M_BLOCKS, D, BLOCK_M, BLOCK_N, num_warps=4 if q.shape[-1] == 64 else 8, num_stages=3, ) ctx.save_for_backward(q, k, v, k_block_id, lut, lse, o_s) ctx.qk_scale = qk_scale ctx.topk = topk ctx.BLOCK_M = BLOCK_M ctx.BLOCK_N = BLOCK_N return o_s # ==================================SageSLA Class=================================== SAGESLA_ENABLED = True try: import spas_sage_attn._fused as fused import spas_sage_attn._qattn as qattn from spas_sage_attn.utils import block_map_lut_triton, get_vanilla_qk_quant except ImportError: SAGESLA_ENABLED = False SAGE2PP_ENABLED = True try: from spas_sage_attn._qattn import ( qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold, ) except ImportError: SAGE2PP_ENABLED = False class SageSparseLinearAttentionBackend(AttentionBackend): """Quantized Sparse-Linear Attention backend using SageAttention kernels.""" accept_output_buffer: bool = True @staticmethod def get_supported_head_sizes() -> list[int]: return [64, 128] @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.SAGE_SLA_ATTN @staticmethod def get_impl_cls() -> type["SageSparseLinearAttentionImpl"]: return SageSparseLinearAttentionImpl @staticmethod def get_metadata_cls() -> type["SageSparseLinearAttentionMetadata"]: return SageSparseLinearAttentionMetadata @staticmethod def get_builder_cls() -> type["SageSparseLinearAttentionMetadataBuilder"]: return SageSparseLinearAttentionMetadataBuilder @dataclass class SageSparseLinearAttentionMetadata(AttentionMetadata): """Metadata for Sage Sparse Linear Attention computation.""" # Basic attention parameters current_timestep: int # Sparse attention configuration topk_ratio: float = 0.1 class SageSparseLinearAttentionMetadataBuilder(AttentionMetadataBuilder): """Builder for SageSparseLinearAttentionMetadata.""" def __init__(self) -> None: pass def prepare(self) -> None: pass def build( self, current_timestep: int, topk_ratio: float = 0.1, **kwargs: dict[str, Any], ) -> SageSparseLinearAttentionMetadata: return SageSparseLinearAttentionMetadata( current_timestep=current_timestep, topk_ratio=topk_ratio, ) class SageSparseLinearAttentionImpl(AttentionImpl, nn.Module): def __init__( self, num_heads: int, head_size: int, causal: bool = False, softmax_scale: float | None = None, num_kv_heads: int | None = None, prefix: str = "", topk_ratio: float = 0.5, feature_map: str = "softmax", use_bf16: bool = True, **extra_impl_args, ) -> None: nn.Module.__init__(self) assert ( SAGESLA_ENABLED ), "Install spas_sage_attn(pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation) first to enable SageSLA." self.num_heads = num_heads self.head_size = head_size self.softmax_scale = softmax_scale if softmax_scale else head_size**-0.5 self.causal = causal self.prefix = prefix self.topk_ratio = topk_ratio self.dtype = torch.bfloat16 if use_bf16 else torch.float16 # Learnable linear projection for combining sparse + linear attention self.proj_l = nn.Linear(head_size, head_size, dtype=torch.float32) # Feature map for linear attention # Type annotation for callables self.feature_map_q: Callable[[torch.Tensor], torch.Tensor] self.feature_map_k: Callable[[torch.Tensor], torch.Tensor] if feature_map == "elu": self.feature_map_q = lambda x: F.elu(x) + 1 self.feature_map_k = lambda x: F.elu(x) + 1 elif feature_map == "relu": self.feature_map_q = F.relu self.feature_map_k = F.relu elif feature_map == "softmax": self.feature_map_q = lambda x: F.softmax(x, dim=-1) self.feature_map_k = lambda x: F.softmax(x, dim=-1) else: raise ValueError(f"Unknown feature map: {feature_map}") self._init_weights() def _init_weights(self) -> None: """Initialize projection weights to zero for residual-like behavior.""" with torch.no_grad(): nn.init.zeros_(self.proj_l.weight) nn.init.zeros_(self.proj_l.bias) # type: ignore[arg-type] def _calc_linear_attention_with_torch( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, ): kv = torch.matmul(k.transpose(-1, -2), v) k_sum = torch.sum(k, dim=-2, keepdim=True) return torch.matmul(q, kv) / (1e-5 + torch.matmul(q, k_sum.transpose(-1, -2))) def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: """Forward pass for Sage Sparse Linear attention with quantized kernels. Args: query: query tensor of shape (B, L, H, D) key: key tensor of shape (B, L, H, D) value: value tensor of shape (B, L, H, D) attn_metadata: attention metadata containing configuration Returns: output tensor of shape (B, L, H, D) """ dtype = query.dtype # Transpose from (B, L, H, D) to SLA format (B, H, L, D) q = query.transpose(1, 2).contiguous() k = key.transpose(1, 2).contiguous() v = value.transpose(1, 2).contiguous() # Determine block sizes based on GPU architecture arch = _get_cuda_arch(q.device.index) if arch == "sm90": BLKQ = 64 BLKK = 128 else: BLKQ = 128 BLKK = 64 # Compute block-sparse attention pattern sparse_map, lut, real_topk = get_block_map( q, k, topk_ratio=self.topk_ratio, BLKQ=BLKQ, BLKK=BLKK ) # Convert to compute dtype q = q.to(self.dtype) k = k.to(self.dtype) v = v.to(self.dtype) ########## SPARGE BEGIN ########## km = k.mean(dim=-2, keepdim=True) headdim = q.size(-1) assert headdim in [ 64, 128, ], "headdim should be in [64, 128]. For other headdim, you can use padding and specify the softmax scale." # Quantize Q, K to INT8 q_int8, q_scale, k_int8, k_scale = get_vanilla_qk_quant(q, k, km, BLKQ, BLKK) lut, valid_block_num = block_map_lut_triton(sparse_map) scale = 1.0 / (headdim**0.5) o_s = torch.empty_like(q) if arch in ("sm80", "sm86", "sm87"): pvthreshold = torch.full( (q.shape[-3],), 1e6, dtype=torch.float32, device=q.device ) v_fp16 = v.to(torch.float16) qattn.qk_int8_sv_f16_accum_f16_block_sparse_attn_inst_buf_with_pv_threshold( q_int8, k_int8, v_fp16, o_s, lut, valid_block_num, pvthreshold, q_scale, k_scale, 1, False, 1, scale, 0, ) else: b, h_kv, kv_len, head_dim = v.shape padded_len = (kv_len + 127) // 128 * 128 v_transposed_permutted = torch.empty( (b, h_kv, head_dim, padded_len), dtype=v.dtype, device=v.device ) fused.transpose_pad_permute_cuda(v, v_transposed_permutted, 1) v_fp8 = torch.empty( v_transposed_permutted.shape, dtype=torch.float8_e4m3fn, device=v.device ) v_scale = torch.empty( (b, h_kv, head_dim), dtype=torch.float32, device=v.device ) fused.scale_fuse_quant_cuda( v_transposed_permutted, v_fp8, v_scale, kv_len, 2.25, 1 ) if arch == "sm90": qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90( q_int8, k_int8, v_fp8, o_s, lut, valid_block_num, q_scale, k_scale, v_scale, 1, False, 1, scale, ) else: pvthreshold = torch.full( (q.shape[-3],), 1e6, dtype=torch.float32, device=q.device ) if SAGE2PP_ENABLED: qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold( q_int8, k_int8, v_fp8, o_s, lut, valid_block_num, pvthreshold, q_scale, k_scale, v_scale, 1, False, 1, scale, 0, ) else: qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold( q_int8, k_int8, v_fp8, o_s, lut, valid_block_num, pvthreshold, q_scale, k_scale, v_scale, 1, False, 1, scale, 0, ) ########## SPARGE END ########## # Linear attention with feature maps q_linear = self.feature_map_q(q).contiguous().to(self.dtype) k_linear = self.feature_map_k(k).contiguous().to(self.dtype) o_l = self._calc_linear_attention_with_torch(q_linear, k_linear, v) # Project linear attention output and combine with torch.amp.autocast("cuda", dtype=self.dtype): o_l = self.proj_l(o_l) # Combine sparse and linear outputs output = (o_s + o_l).to(dtype).transpose(1, 2) return output ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_video_gen_2_attn.py ================================================ """ Sparse Video Gen 2 (SAP) attention backend. This is a baseline integration that wires the backend into the attention framework. Adapted from https://github.com/svg-project/Sparse-VideoGen/blob/main/svg/models/wan/attention.py """ from dataclasses import dataclass, field from typing import Any import torch import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel try: from svg.kernels.triton.permute import ( apply_inverse_permutation_triton, permute_tensor_by_labels_triton, ) from svg.kmeans_utils import ( batch_kmeans_Euclid, dynamic_block_sparse_fwd_flashinfer, identify_dynamic_map, ) svg2_available = True except ImportError: svg2_available = False from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class SparseVideoGen2AttentionBackend(AttentionBackend): accept_output_buffer: bool = True @staticmethod def get_supported_head_sizes() -> list[int]: return [64, 128, 256] @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN @staticmethod def get_impl_cls() -> type["SparseVideoGen2AttentionImpl"]: return SparseVideoGen2AttentionImpl @staticmethod def get_metadata_cls() -> type["SparseVideoGen2AttentionMetadata"]: return SparseVideoGen2AttentionMetadata @staticmethod def get_builder_cls() -> type["SparseVideoGen2AttentionMetadataBuilder"]: return SparseVideoGen2AttentionMetadataBuilder @dataclass class Svg2LayerCache: # centroids for kmeans clustering q_centroids: torch.Tensor | None = None k_centroids: torch.Tensor | None = None centroids_initialized: bool = False @dataclass class Svg2Cache: layers: dict[int, Svg2LayerCache] = field(default_factory=dict) def get_layer(self, layer_idx: int) -> Svg2LayerCache: layer_cache = self.layers.get(layer_idx) if layer_cache is None: layer_cache = Svg2LayerCache() self.layers[layer_idx] = layer_cache return layer_cache @dataclass class SparseVideoGen2AttentionMetadata(AttentionMetadata): current_timestep: int num_q_centroids: int num_k_centroids: int top_p_kmeans: float min_kc_ratio: float kmeans_iter_init: int kmeans_iter_step: int zero_step_kmeans_init: bool first_layers_fp: float first_times_fp: float context_length: int num_frame: int frame_size: int cache: Svg2Cache prompt_length: int | None = None max_seqlen_q: int | None = None max_seqlen_k: int | None = None def _require_kwarg(kwargs: dict[str, Any], name: str) -> Any: if name not in kwargs: raise ValueError( f"Missing required argument for SparseVideoGen2Attention: {name}" ) return kwargs[name] class SparseVideoGen2AttentionMetadataBuilder(AttentionMetadataBuilder): def __init__(self) -> None: pass def prepare(self) -> None: pass def build( # type: ignore[override] self, current_timestep: int, raw_latent_shape: tuple[int, ...], patch_size: tuple[int, int, int], cache: Svg2Cache, num_q_centroids: int, num_k_centroids: int, top_p_kmeans: float, min_kc_ratio: float, kmeans_iter_init: int, kmeans_iter_step: int, zero_step_kmeans_init: bool, first_layers_fp: float, first_times_fp: float, context_length: int = 0, prompt_length: int | None = None, **kwargs: dict[str, Any], ) -> SparseVideoGen2AttentionMetadata: raw_shape = tuple(raw_latent_shape) if len(raw_shape) == 5: t, h, w = raw_shape[2:5] elif len(raw_shape) == 3: t, h, w = raw_shape else: raise ValueError( "raw_latent_shape must be (T, H, W) or (B, C, T, H, W) for SAP attention" ) pt, ph, pw = patch_size if t % pt != 0 or h % ph != 0 or w % pw != 0: raise ValueError( "raw_latent_shape must be divisible by patch_size for SAP attention" ) num_frame = t // pt frame_size = (h // ph) * (w // pw) return SparseVideoGen2AttentionMetadata( current_timestep=current_timestep, num_q_centroids=num_q_centroids, num_k_centroids=num_k_centroids, top_p_kmeans=top_p_kmeans, min_kc_ratio=min_kc_ratio, kmeans_iter_init=kmeans_iter_init, kmeans_iter_step=kmeans_iter_step, zero_step_kmeans_init=zero_step_kmeans_init, first_layers_fp=first_layers_fp, first_times_fp=first_times_fp, context_length=context_length, prompt_length=prompt_length, num_frame=num_frame, frame_size=frame_size, cache=cache, ) class SparseVideoGen2AttentionImpl(AttentionImpl): def __init__( self, num_heads: int, head_size: int, causal: bool, softmax_scale: float, num_kv_heads: int | None = None, prefix: str = "", **extra_impl_args, ) -> None: if causal: raise ValueError( "Sparse Video Gen 2 attention does not support causal attention" ) if not svg2_available: raise ImportError( "Sparse Video Gen 2 attention backend requires svg package to be installed" "Please install it by following the instructions at " "https://github.com/svg-project/Sparse-VideoGen" ) self.prefix = prefix self.layer_idx = self._get_layer_idx(prefix) def _get_layer_idx(self, prefix: str) -> int: parts = prefix.split(".") if len(parts) < 3: raise ValueError( f"Invalid prefix for SparseVideoGen2AttentionImpl: {prefix}" ) return int(parts[-3]) def kmeans_init( self, query: torch.Tensor, key: torch.Tensor, attn_metadata: SparseVideoGen2AttentionMetadata, ): cfg, num_heads, seq_len, dim = query.size() qlabels, qcentroids, qcluster_sizes, qiter = batch_kmeans_Euclid( query.reshape(cfg * num_heads, seq_len, dim), n_clusters=attn_metadata.num_q_centroids, max_iters=attn_metadata.kmeans_iter_init, ) klabels, kcentroids, kcluster_sizes, kiter = batch_kmeans_Euclid( key.reshape(cfg * num_heads, seq_len, dim), n_clusters=attn_metadata.num_k_centroids, max_iters=attn_metadata.kmeans_iter_init, ) layer_cache = attn_metadata.cache.get_layer(self.layer_idx) layer_cache.q_centroids = qcentroids layer_cache.k_centroids = kcentroids return ( qlabels, qcentroids, qcluster_sizes, qiter, klabels, kcentroids, kcluster_sizes, kiter, ) def kmeans_step( self, query: torch.Tensor, key: torch.Tensor, attn_metadata: SparseVideoGen2AttentionMetadata, ): cfg, num_heads, seq_len, dim = query.size() layer_cache = attn_metadata.cache.get_layer(self.layer_idx) qlabels, qcentroids, qcluster_sizes, qiter = batch_kmeans_Euclid( query.reshape(cfg * num_heads, seq_len, dim), n_clusters=attn_metadata.num_q_centroids, max_iters=attn_metadata.kmeans_iter_step, init_centroids=layer_cache.q_centroids, ) klabels, kcentroids, kcluster_sizes, kiter = batch_kmeans_Euclid( key.reshape(cfg * num_heads, seq_len, dim), n_clusters=attn_metadata.num_k_centroids, max_iters=attn_metadata.kmeans_iter_step, init_centroids=layer_cache.k_centroids, ) layer_cache.q_centroids = qcentroids layer_cache.k_centroids = kcentroids return ( qlabels, qcentroids, qcluster_sizes, qiter, klabels, kcentroids, kcluster_sizes, kiter, ) def kmeans_clustering( self, query: torch.Tensor, key: torch.Tensor, attn_metadata: SparseVideoGen2AttentionMetadata, ): layer_cache = attn_metadata.cache.get_layer(self.layer_idx) if not layer_cache.centroids_initialized: ( qlabels, qcentroids, qcluster_sizes, qiter, klabels, kcentroids, kcluster_sizes, kiter, ) = self.kmeans_init(query, key, attn_metadata) layer_cache.centroids_initialized = True logger.debug( "Centroids initialized at layer %s (init iters: %s).", self.layer_idx, attn_metadata.kmeans_iter_init, ) else: ( qlabels, qcentroids, qcluster_sizes, qiter, klabels, kcentroids, kcluster_sizes, kiter, ) = self.kmeans_step(query, key, attn_metadata) return ( qlabels, qcentroids, qcluster_sizes, qiter, klabels, kcentroids, kcluster_sizes, kiter, ) def semantic_aware_permutation( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: SparseVideoGen2AttentionMetadata, ): cfg, num_heads, seq_len, dim = query.size() # 1. Kmeans clustering ( qlabels, qcentroids, qcluster_sizes, qiter, klabels, kcentroids, kcluster_sizes, kiter, ) = self.kmeans_clustering(query, key, attn_metadata) # 2. Identify dynamic map q_cluster_sizes = qcluster_sizes.view( cfg, num_heads, attn_metadata.num_q_centroids ) k_cluster_sizes = kcluster_sizes.view( cfg, num_heads, attn_metadata.num_k_centroids ) dynamic_map = identify_dynamic_map( qcentroids.view(cfg, num_heads, attn_metadata.num_q_centroids, dim), kcentroids.view(cfg, num_heads, attn_metadata.num_k_centroids, dim), q_cluster_sizes, k_cluster_sizes, attn_metadata.top_p_kmeans, attn_metadata.min_kc_ratio, ) # 3. Permute the query, key, value q_permuted, q_sorted_indices = permute_tensor_by_labels_triton( query, qlabels, dim=2 ) k_permuted, k_sorted_indices = permute_tensor_by_labels_triton( key, klabels, dim=2 ) v_permuted, v_sorted_indices = permute_tensor_by_labels_triton( value, klabels, dim=2, sorted_indices=k_sorted_indices ) return ( q_permuted, k_permuted, v_permuted, dynamic_map, q_cluster_sizes, k_cluster_sizes, q_sorted_indices, ) def _hunyuan_dynamic_map_post_processing( self, q_perm: torch.Tensor, k_perm: torch.Tensor, v_perm: torch.Tensor, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, dyn_map: torch.Tensor, qc_sz_s: torch.Tensor, kc_sz_s: torch.Tensor, q_sorted_indices: torch.Tensor, video_length: int, context_length: int, prompt_length: int, unprompt_length: int, ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, ]: # Place the permuted video tokens back and keep text tokens at the tail. query[:, :, :-context_length, :] = q_perm key[:, :, :-context_length, :] = k_perm value[:, :, :-context_length, :] = v_perm # Add prompt/unprompt clusters to the dynamic map. dyn_map = F.pad(dyn_map, (0, 2, 0, 2), value=0) dyn_map[:, :, -2, :-1] = True dyn_map[:, :, :-1, -2] = True dyn_map[:, :, -1, -1] = True qc_sz_s = F.pad(qc_sz_s, (0, 2), value=0) qc_sz_s[:, :, -2] = prompt_length qc_sz_s[:, :, -1] = unprompt_length kc_sz_s = F.pad(kc_sz_s, (0, 2), value=0) kc_sz_s[:, :, -2] = prompt_length kc_sz_s[:, :, -1] = unprompt_length q_sorted_indices = F.pad(q_sorted_indices, (0, context_length), value=0) q_sorted_indices[:, video_length:] = torch.arange( video_length, video_length + context_length, device=q_sorted_indices.device, ) return query, key, value, dyn_map, qc_sz_s, kc_sz_s, q_sorted_indices def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: SparseVideoGen2AttentionMetadata, ) -> torch.Tensor: torch.backends.cuda.preferred_linalg_library(backend="magma") res = None # bshd -> bhsd query = query.transpose(1, 2).contiguous() key = key.transpose(1, 2).contiguous() value = value.transpose(1, 2).contiguous() batch_size, num_heads, seq_len, dim = query.size() context_length, num_frame, frame_size = ( attn_metadata.context_length, attn_metadata.num_frame, attn_metadata.frame_size, ) prompt_length = attn_metadata.prompt_length if prompt_length is None: prompt_length = context_length assert ( seq_len == context_length + num_frame * frame_size ), f"Query Shape: {seq_len} is not equivalent to {context_length} + {num_frame} * {frame_size}" # Determine if we use Full Attention to calculate full_attention_flag = False if self.layer_idx < attn_metadata.first_layers_fp: full_attention_flag = True if attn_metadata.current_timestep > attn_metadata.first_times_fp: full_attention_flag = True if full_attention_flag: if attn_metadata.zero_step_kmeans_init: video_length = attn_metadata.num_frame * attn_metadata.frame_size query_video = query[:, :, :video_length, :].contiguous() key_video = key[:, :, :video_length, :].contiguous() self.kmeans_clustering(query_video, key_video, attn_metadata) with sdpa_kernel( SDPBackend.CUDNN_ATTENTION ): # not sure why we need to force cudnn here, but it's faster than flash attention output_hidden_states = torch.nn.functional.scaled_dot_product_attention( query, key, value, dropout_p=0.0, is_causal=False ) res = output_hidden_states.reshape( batch_size, num_heads, seq_len, dim ).transpose(1, 2) else: if context_length > 0: video_length = num_frame * frame_size unprompt_length = max(context_length - prompt_length, 0) query_video = query[:, :, :video_length, :].contiguous() key_video = key[:, :, :video_length, :].contiguous() value_video = value[:, :, :video_length, :].contiguous() ( q_perm, k_perm, v_perm, dyn_map, qc_sz_s, kc_sz_s, q_sorted_indices, ) = self.semantic_aware_permutation( query_video, key_video, value_video, attn_metadata ) ( q_perm, k_perm, v_perm, dyn_map, qc_sz_s, kc_sz_s, q_sorted_indices, ) = self._hunyuan_dynamic_map_post_processing( q_perm, k_perm, v_perm, query, key, value, dyn_map, qc_sz_s, kc_sz_s, q_sorted_indices, video_length, context_length, prompt_length, unprompt_length, ) else: ( q_perm, k_perm, v_perm, dyn_map, qc_sz_s, kc_sz_s, q_sorted_indices, ) = self.semantic_aware_permutation(query, key, value, attn_metadata) output_permuted = dynamic_block_sparse_fwd_flashinfer( q_perm, k_perm, v_perm, dyn_map, qc_sz_s, kc_sz_s, is_cpu=False ) attn_output = apply_inverse_permutation_triton( output_permuted, q_sorted_indices, dim=2 ) res = attn_output.reshape(batch_size, num_heads, seq_len, dim).transpose( 1, 2 ) torch.backends.cuda.preferred_linalg_library( backend="default" ) # reset to default return res.contiguous() ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import functools import math from dataclasses import dataclass import torch try: from vsa import video_sparse_attn except ImportError: video_sparse_attn = None from typing import Any from sglang.multimodal_gen.runtime.distributed import get_sp_group from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) VSA_TILE_SIZE = (4, 4, 4) @functools.lru_cache(maxsize=10) def get_tile_partition_indices( dit_seq_shape: tuple[int, int, int], tile_size: tuple[int, int, int], device: torch.device, ) -> torch.LongTensor: T, H, W = dit_seq_shape ts, hs, ws = tile_size indices = torch.arange(T * H * W, device=device, dtype=torch.long).reshape(T, H, W) ls = [] for t in range(math.ceil(T / ts)): for h in range(math.ceil(H / hs)): for w in range(math.ceil(W / ws)): ls.append( indices[ t * ts : min(t * ts + ts, T), h * hs : min(h * hs + hs, H), w * ws : min(w * ws + ws, W), ].flatten() ) index = torch.cat(ls, dim=0) return index @functools.lru_cache(maxsize=10) def get_reverse_tile_partition_indices( dit_seq_shape: tuple[int, int, int], tile_size: tuple[int, int, int], device: torch.device, ) -> torch.LongTensor: return torch.argsort(get_tile_partition_indices(dit_seq_shape, tile_size, device)) @functools.lru_cache(maxsize=10) def construct_variable_block_sizes( dit_seq_shape: tuple[int, int, int], num_tiles: tuple[int, int, int], device: torch.device, ) -> torch.LongTensor: """ Compute the number of valid (non‑padded) tokens inside every (ts_t × ts_h × ts_w) tile after padding ‑‑ flattened in the order (t‑tile, h‑tile, w‑tile) that `rearrange` uses. Returns ------- torch.LongTensor # shape: [∏ full_window_size] """ # unpack t, h, w = dit_seq_shape ts_t, ts_h, ts_w = VSA_TILE_SIZE n_t, n_h, n_w = num_tiles def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.LongTensor: """Vector with the size of each tile along one dimension.""" sizes = torch.full((n_tiles,), tile, dtype=torch.int, device=device) # size of last (possibly partial) tile remainder = dim_len - (n_tiles - 1) * tile sizes[-1] = remainder if remainder > 0 else tile return sizes t_sizes = _sizes(t, ts_t, n_t) # [n_t] h_sizes = _sizes(h, ts_h, n_h) # [n_h] w_sizes = _sizes(w, ts_w, n_w) # [n_w] # broadcast‑multiply to get voxels per tile, then flatten block_sizes = ( t_sizes[:, None, None] # [n_t, 1, 1] * h_sizes[None, :, None] # [1, n_h, 1] * w_sizes[None, None, :] # [1, 1, n_w] ).reshape( -1 ) # [n_t * n_h * n_w] return block_sizes @functools.lru_cache(maxsize=10) def get_non_pad_index( variable_block_sizes: torch.LongTensor, max_block_size: int, ): n_win = variable_block_sizes.shape[0] device = variable_block_sizes.device starts_pad = torch.arange(n_win, device=device) * max_block_size index_pad = ( starts_pad[:, None] + torch.arange(max_block_size, device=device)[None, :] ) index_mask = ( torch.arange(max_block_size, device=device)[None, :] < variable_block_sizes[:, None] ) return index_pad[index_mask] class VideoSparseAttentionBackend(AttentionBackend): accept_output_buffer: bool = True @staticmethod def get_supported_head_sizes() -> list[int]: return [64, 128] @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.VIDEO_SPARSE_ATTN @staticmethod def get_impl_cls() -> type["VideoSparseAttentionImpl"]: return VideoSparseAttentionImpl @staticmethod def get_metadata_cls() -> type["VideoSparseAttentionMetadata"]: return VideoSparseAttentionMetadata @staticmethod def get_builder_cls() -> type["VideoSparseAttentionMetadataBuilder"]: return VideoSparseAttentionMetadataBuilder @dataclass class VideoSparseAttentionMetadata(AttentionMetadata): current_timestep: int dit_seq_shape: list[int] VSA_sparsity: float num_tiles: list[int] total_seq_length: int tile_partition_indices: torch.LongTensor reverse_tile_partition_indices: torch.LongTensor variable_block_sizes: torch.LongTensor non_pad_index: torch.LongTensor # adaption for FastWan2.1-T2V-1.3B-Diffusers # Sequence lengths for the forward batch # Maximum sequence length for query max_seqlen_q: int = 1 # Maximum sequence length for key max_seqlen_k: int = 0 class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder): def __init__(self): pass def prepare(self): pass def build( # type: ignore self, current_timestep: int, raw_latent_shape: tuple[int, int, int], patch_size: tuple[int, int, int], VSA_sparsity: float, device: torch.device, **kwargs: dict[str, Any], ) -> VideoSparseAttentionMetadata: patch_size = patch_size dit_seq_shape = ( raw_latent_shape[0] // patch_size[0], raw_latent_shape[1] // patch_size[1], raw_latent_shape[2] // patch_size[2], ) num_tiles = ( math.ceil(dit_seq_shape[0] / VSA_TILE_SIZE[0]), math.ceil(dit_seq_shape[1] / VSA_TILE_SIZE[1]), math.ceil(dit_seq_shape[2] / VSA_TILE_SIZE[2]), ) total_seq_length = math.prod(dit_seq_shape) tile_partition_indices = get_tile_partition_indices( dit_seq_shape, VSA_TILE_SIZE, device ) reverse_tile_partition_indices = get_reverse_tile_partition_indices( dit_seq_shape, VSA_TILE_SIZE, device ) variable_block_sizes = construct_variable_block_sizes( dit_seq_shape, num_tiles, device ) non_pad_index = get_non_pad_index( variable_block_sizes, math.prod(VSA_TILE_SIZE) ) return VideoSparseAttentionMetadata( current_timestep=current_timestep, dit_seq_shape=dit_seq_shape, # type: ignore VSA_sparsity=VSA_sparsity, # type: ignore num_tiles=num_tiles, # type: ignore total_seq_length=total_seq_length, # type: ignore tile_partition_indices=tile_partition_indices, # type: ignore reverse_tile_partition_indices=reverse_tile_partition_indices, variable_block_sizes=variable_block_sizes, non_pad_index=non_pad_index, ) class VideoSparseAttentionImpl(AttentionImpl): def __init__( self, num_heads: int, head_size: int, causal: bool, softmax_scale: float, num_kv_heads: int | None = None, prefix: str = "", **extra_impl_args, ) -> None: self.prefix = prefix sp_group = get_sp_group() self.sp_size = sp_group.world_size def tile( self, x: torch.Tensor, num_tiles: list[int], tile_partition_indices: torch.LongTensor, non_pad_index: torch.LongTensor, ) -> torch.Tensor: t_padded_size = num_tiles[0] * VSA_TILE_SIZE[0] h_padded_size = num_tiles[1] * VSA_TILE_SIZE[1] w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2] x_padded = torch.zeros( ( x.shape[0], t_padded_size * h_padded_size * w_padded_size, x.shape[-2], x.shape[-1], ), device=x.device, dtype=x.dtype, ) x_padded[:, non_pad_index] = x[:, tile_partition_indices] return x_padded def untile( self, x: torch.Tensor, reverse_tile_partition_indices: torch.LongTensor, non_pad_index: torch.LongTensor, ) -> torch.Tensor: x = x[:, non_pad_index][:, reverse_tile_partition_indices] return x def preprocess_qkv( self, qkv: torch.Tensor, attn_metadata: VideoSparseAttentionMetadata, ) -> torch.Tensor: return self.tile( qkv, attn_metadata.num_tiles, attn_metadata.tile_partition_indices, attn_metadata.non_pad_index, ) def postprocess_output( self, output: torch.Tensor, attn_metadata: VideoSparseAttentionMetadata, ) -> torch.Tensor: return self.untile( output, attn_metadata.reverse_tile_partition_indices, attn_metadata.non_pad_index, ) def forward( # type: ignore[override] self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, gate_compress: torch.Tensor, attn_metadata: VideoSparseAttentionMetadata, ) -> torch.Tensor: query = query.transpose(1, 2).contiguous() key = key.transpose(1, 2).contiguous() value = value.transpose(1, 2).contiguous() gate_compress = gate_compress.transpose(1, 2).contiguous() VSA_sparsity = attn_metadata.VSA_sparsity cur_topk = math.ceil( (1 - VSA_sparsity) * (attn_metadata.total_seq_length / math.prod(VSA_TILE_SIZE)) ) if video_sparse_attn is None: raise NotImplementedError("video_sparse_attn is not installed") hidden_states = video_sparse_attn( query, key, value, variable_block_sizes=attn_metadata.variable_block_sizes, topk=cur_topk, block_size=VSA_TILE_SIZE, compress_attn_weight=gate_compress, ).transpose(1, 2) return hidden_states ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import re from dataclasses import dataclass import torch from einops import rearrange from kernel.attn.vmoba_attn.vmoba import ( moba_attn_varlen, process_moba_input, process_moba_output, ) from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class VMOBAAttentionBackend(AttentionBackend): accept_output_buffer: bool = True @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.VMOBA_ATTN @staticmethod def get_impl_cls() -> type["VMOBAAttentionImpl"]: return VMOBAAttentionImpl @staticmethod def get_metadata_cls() -> type["VideoMobaAttentionMetadata"]: return VideoMobaAttentionMetadata @staticmethod def get_builder_cls() -> type["VideoMobaAttentionMetadataBuilder"]: return VideoMobaAttentionMetadataBuilder @dataclass class VideoMobaAttentionMetadata(AttentionMetadata): current_timestep: int temporal_chunk_size: int temporal_topk: int spatial_chunk_size: tuple[int, int] spatial_topk: int st_chunk_size: tuple[int, int, int] st_topk: int moba_select_mode: str moba_threshold: float moba_threshold_type: str patch_resolution: list[int] first_full_step: int = 12 first_full_layer: int = 0 # temporal_layer -> spatial_layer -> st_layer temporal_layer: int = 1 spatial_layer: int = 1 st_layer: int = 1 def pad_input(hidden_states, indices, batch, seqlen): """ Arguments: hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. batch: int, batch size for the padded sequence. seqlen: int, maximum sequence length for the padded sequence. Return: hidden_states: (batch, seqlen, ...) """ dim = hidden_states.shape[1:] output = torch.zeros( (batch * seqlen), *dim, device=hidden_states.device, dtype=hidden_states.dtype ) output[indices] = hidden_states return rearrange(output, "(b s) ... -> b s ...", b=batch) class VideoMobaAttentionMetadataBuilder(AttentionMetadataBuilder): def __init__(self): pass def prepare(self): pass def build( # type: ignore self, current_timestep: int, raw_latent_shape: tuple[int, int, int], patch_size: tuple[int, int, int], temporal_chunk_size: int, temporal_topk: int, spatial_chunk_size: tuple[int, int], spatial_topk: int, st_chunk_size: tuple[int, int, int], st_topk: int, moba_select_mode: str = "threshold", moba_threshold: float = 0.25, moba_threshold_type: str = "query_head", device: torch.device = None, first_full_layer: int = 0, first_full_step: int = 12, temporal_layer: int = 1, spatial_layer: int = 1, st_layer: int = 1, **kwargs, ) -> VideoMobaAttentionMetadata: if device is None: device = torch.device("cpu") assert ( raw_latent_shape[0] % patch_size[0] == 0 and raw_latent_shape[1] % patch_size[1] == 0 and raw_latent_shape[2] % patch_size[2] == 0 ), f"spatial patch_resolution {raw_latent_shape} should be divisible by patch_size {patch_size}" patch_resolution = [ t // pt for t, pt in zip(raw_latent_shape, patch_size, strict=False) ] return VideoMobaAttentionMetadata( current_timestep=current_timestep, temporal_chunk_size=temporal_chunk_size, temporal_topk=temporal_topk, spatial_chunk_size=spatial_chunk_size, spatial_topk=spatial_topk, st_chunk_size=st_chunk_size, st_topk=st_topk, moba_select_mode=moba_select_mode, moba_threshold=moba_threshold, moba_threshold_type=moba_threshold_type, patch_resolution=patch_resolution, first_full_layer=first_full_layer, first_full_step=first_full_step, temporal_layer=temporal_layer, spatial_layer=spatial_layer, st_layer=st_layer, ) class VMOBAAttentionImpl(AttentionImpl): def __init__( self, num_heads, head_size, softmax_scale, causal=False, num_kv_heads=None, prefix="", **extra_impl_args, ) -> None: self.prefix = prefix self.layer_idx = self._get_layer_idx(prefix) self.pad_input = pad_input def _get_layer_idx(self, prefix: str) -> int | None: match = re.search(r"blocks\.(\d+)", prefix) if not match: raise ValueError(f"Invalid prefix: {prefix}") return int(match.group(1)) def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: """ query: [B, L, H, D] key: [B, L, H, D] value: [B, L, H, D] attn_metadata: AttentionMetadata """ batch_size, sequence_length, num_heads, head_dim = query.shape # select chunk type according to layer idx: loop_layer_num = ( attn_metadata.temporal_layer + attn_metadata.spatial_layer + attn_metadata.st_layer ) moba_layer = self.layer_idx - attn_metadata.first_full_layer if moba_layer % loop_layer_num < attn_metadata.temporal_layer: moba_chunk_size = attn_metadata.temporal_chunk_size moba_topk = attn_metadata.temporal_topk elif ( moba_layer % loop_layer_num < attn_metadata.temporal_layer + attn_metadata.spatial_layer ): moba_chunk_size = attn_metadata.spatial_chunk_size moba_topk = attn_metadata.spatial_topk elif ( moba_layer % loop_layer_num < attn_metadata.temporal_layer + attn_metadata.spatial_layer + attn_metadata.st_layer ): moba_chunk_size = attn_metadata.st_chunk_size moba_topk = attn_metadata.st_topk query, chunk_size = process_moba_input( query, attn_metadata.patch_resolution, moba_chunk_size ) key, chunk_size = process_moba_input( key, attn_metadata.patch_resolution, moba_chunk_size ) value, chunk_size = process_moba_input( value, attn_metadata.patch_resolution, moba_chunk_size ) max_seqlen = query.shape[1] indices_q = torch.arange( 0, query.shape[0] * query.shape[1], device=query.device ) cu_seqlens = torch.arange( 0, query.shape[0] * query.shape[1] + 1, query.shape[1], dtype=torch.int32, device=query.device, ) query = rearrange(query, "b s ... -> (b s) ...") key = rearrange(key, "b s ... -> (b s) ...") value = rearrange(value, "b s ... -> (b s) ...") # current_timestep=attn_metadata.current_timestep hidden_states = moba_attn_varlen( query, key, value, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, moba_chunk_size=chunk_size, moba_topk=moba_topk, select_mode=attn_metadata.moba_select_mode, simsum_threshold=attn_metadata.moba_threshold, threshold_type=attn_metadata.moba_threshold_type, ) hidden_states = self.pad_input( hidden_states, indices_q, batch_size, sequence_length ) hidden_states = process_moba_output( hidden_states, attn_metadata.patch_resolution, moba_chunk_size ) return hidden_states ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/layer.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from typing import Type import torch import torch.nn as nn from sglang.multimodal_gen.runtime.distributed.communication_op import ( sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_ring_parallel_world_size, get_sequence_parallel_world_size, get_sp_group, get_sp_parallel_rank, get_sp_world_size, get_ulysses_parallel_world_size, ) from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionImpl, ) from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend from sglang.multimodal_gen.runtime.layers.usp import ( _usp_input_all_to_all, _usp_output_all_to_all, ring_attn, ) from sglang.multimodal_gen.runtime.managers.forward_context import ( ForwardContext, get_forward_context, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.utils import get_compute_dtype class UlyssesAttention(nn.Module): """Ulysses-style SequenceParallelism attention layer.""" def __init__( self, num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", **extra_impl_args, ) -> None: super().__init__() if softmax_scale is None: self.softmax_scale = head_size**-0.5 else: self.softmax_scale = softmax_scale if num_kv_heads is None: num_kv_heads = num_heads dtype = get_compute_dtype() attn_backend = get_attn_backend( head_size, dtype, supported_attention_backends=supported_attention_backends ) impl_cls = attn_backend.get_impl_cls() self.attn_impl = impl_cls( num_heads=num_heads, head_size=head_size, causal=causal, softmax_scale=self.softmax_scale, num_kv_heads=num_kv_heads, prefix=f"{prefix}.impl", **extra_impl_args, ) self.num_heads = num_heads self.head_size = head_size self.num_kv_heads = num_kv_heads self.backend = attn_backend.get_enum() self.dtype = dtype def forward( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, replicated_q: torch.Tensor | None = None, replicated_k: torch.Tensor | None = None, replicated_v: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Forward pass for distributed attention. Args: q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim] k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim] v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim] replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens replicated_k (Optional[torch.Tensor]): Replicated key tensor replicated_v (Optional[torch.Tensor]): Replicated value tensor Returns: Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: - o (torch.Tensor): Output tensor after attention for the main sequence - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided """ # Check input shapes assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors" batch_size, seq_len, num_heads, head_dim = q.shape local_rank = get_sp_parallel_rank() world_size = get_sp_world_size() forward_context: ForwardContext = get_forward_context() ctx_attn_metadata = forward_context.attn_metadata # Stack QKV qkv = torch.cat([q, k, v], dim=0) # [3, seq_len, num_heads, head_dim] # Redistribute heads across sequence dimension qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1) # Apply backend-specific preprocess_qkv qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata) # Concatenate with replicated QKV if provided if replicated_q is not None: assert replicated_k is not None and replicated_v is not None replicated_qkv = torch.cat( [replicated_q, replicated_k, replicated_v], dim=0 ) # [3, seq_len, num_heads, head_dim] heads_per_rank = num_heads // world_size replicated_qkv = replicated_qkv[ :, :, local_rank * heads_per_rank : (local_rank + 1) * heads_per_rank ] qkv = torch.cat([qkv, replicated_qkv], dim=1) q, k, v = qkv.chunk(3, dim=0) output = self.attn_impl.forward(q, k, v, ctx_attn_metadata) # Redistribute back if using sequence parallelism replicated_output = None if replicated_q is not None: replicated_output = output[:, seq_len * world_size :] output = output[:, : seq_len * world_size] # TODO: make this asynchronous replicated_output = sequence_model_parallel_all_gather( replicated_output.contiguous(), dim=2 ) # Apply backend-specific postprocess_output output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) output = sequence_model_parallel_all_to_all_4D( output, scatter_dim=1, gather_dim=2 ) return output, replicated_output class UlyssesAttention_VSA(UlyssesAttention): """Distributed attention layer with VSA support.""" def forward( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, replicated_q: torch.Tensor | None = None, replicated_k: torch.Tensor | None = None, replicated_v: torch.Tensor | None = None, gate_compress: torch.Tensor | None = None, ) -> torch.Tensor: """Forward pass for distributed attention. Args: q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim] k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim] v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim] gate_compress (torch.Tensor): Gate compress tensor [batch_size, seq_len, num_heads, head_dim] replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens replicated_k (Optional[torch.Tensor]): Replicated key tensor replicated_v (Optional[torch.Tensor]): Replicated value tensor Returns: Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: - o (torch.Tensor): Output tensor after attention for the main sequence - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided """ # Check text tokens are not supported for VSA now assert ( replicated_q is None and replicated_k is None and replicated_v is None ), "Replicated QKV is not supported for VSA now" # Check input shapes assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors" forward_context: ForwardContext = get_forward_context() ctx_attn_metadata = forward_context.attn_metadata # Stack QKV qkvg = torch.cat( [q, k, v, gate_compress], dim=0 ) # [3, seq_len, num_heads, head_dim] # Redistribute heads across sequence dimension qkvg = sequence_model_parallel_all_to_all_4D(qkvg, scatter_dim=2, gather_dim=1) qkvg = self.attn_impl.preprocess_qkv(qkvg, ctx_attn_metadata) q, k, v, gate_compress = qkvg.chunk(4, dim=0) output = self.attn_impl.forward( q, k, v, gate_compress=gate_compress, attn_metadata=ctx_attn_metadata ) # type: ignore[call-arg] # Apply backend-specific postprocess_output output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) output = sequence_model_parallel_all_to_all_4D( output, scatter_dim=1, gather_dim=2 ) return output class LocalAttention(nn.Module): """Attention layer.""" def __init__( self, num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, **extra_impl_args, ) -> None: super().__init__() if softmax_scale is None: self.softmax_scale = head_size**-0.5 else: self.softmax_scale = softmax_scale if num_kv_heads is None: num_kv_heads = num_heads dtype = get_compute_dtype() attn_backend = get_attn_backend( head_size, dtype, supported_attention_backends=supported_attention_backends ) impl_cls = attn_backend.get_impl_cls() self.attn_impl = impl_cls( num_heads=num_heads, head_size=head_size, softmax_scale=self.softmax_scale, num_kv_heads=num_kv_heads, causal=causal, **extra_impl_args, ) self.num_heads = num_heads self.head_size = head_size self.num_kv_heads = num_kv_heads self.backend = attn_backend.get_enum() self.dtype = dtype def forward( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, ) -> torch.Tensor: """ Apply local attention between query, key and value tensors. Args: q (torch.Tensor): Query tensor of shape [batch_size, seq_len, num_heads, head_dim] k (torch.Tensor): Key tensor of shape [batch_size, seq_len, num_heads, head_dim] v (torch.Tensor): Value tensor of shape [batch_size, seq_len, num_heads, head_dim] Returns: torch.Tensor: Output tensor after local attention """ # Check input shapes assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors" forward_context: ForwardContext = get_forward_context() ctx_attn_metadata = forward_context.attn_metadata output = self.attn_impl.forward(q, k, v, attn_metadata=ctx_attn_metadata) return output class USPAttention(nn.Module): """ Ulysses Sequence Parallelism with Ring Attention. This class implements the USP algorithm, which is a combination of Ulysses-style all-to-all communication for sequence-head dimension sharding and Ring Attention for fine-grained sequence parallelism within subgroups. """ def __init__( self, num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", dropout_rate: float = 0.0, skip_sequence_parallel: bool = False, **extra_impl_args, ) -> None: """ Args: skip_sequence_parallel: when KV is replicated across all SP ranks (e.g. cross-attention to text/image encoder outputs), the full USP pipeline is redundant: each rank's local Q shard can attend directly to the locally-held full KV without any collective communication. """ super().__init__() if softmax_scale is None: self.softmax_scale = head_size**-0.5 else: self.softmax_scale = softmax_scale if num_kv_heads is None: num_kv_heads = num_heads dtype = get_compute_dtype() attn_backend = get_attn_backend( head_size, dtype, supported_attention_backends=supported_attention_backends ) impl_cls: Type["AttentionImpl"] = attn_backend.get_impl_cls() self.attn_impl = impl_cls( num_heads=num_heads, head_size=head_size, causal=causal, softmax_scale=self.softmax_scale, num_kv_heads=num_kv_heads, prefix=f"{prefix}.impl", **extra_impl_args, ) self.num_heads = num_heads self.head_size = head_size self.num_kv_heads = num_kv_heads self.backend = attn_backend.get_enum() self.dtype = dtype self.causal = causal self.dropout_p = dropout_rate self.skip_sequence_parallel = skip_sequence_parallel def forward( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_replicated_prefix: int = 0, ) -> torch.Tensor: """ Forward pass for USPAttention. q, k, v: [B, S_local, H, D] num_replicated_prefix: number of leading tokens in q/k/v that are replicated (identical) across all SP ranks, e.g. text tokens in FLUX joint attention. These tokens are excluded from the Ulysses all-to-all so they appear exactly once in the gathered sequence, preserving correct attention weights. Note: Replicated tensors are not supported in this implementation. When skip_sequence_parallel=True (set at construction time), all SP communication is bypassed — use this for cross-attention where KV content is replicated across ranks (distinct from replicated_k/v args). """ forward_context: ForwardContext = get_forward_context() ctx_attn_metadata = forward_context.attn_metadata if self.skip_sequence_parallel or get_sequence_parallel_world_size() == 1: # No sequence parallelism, just run local attention. out = self.attn_impl.forward(q, k, v, ctx_attn_metadata) return out sp_size = get_ulysses_parallel_world_size() if sp_size > 1 and num_replicated_prefix > 0: return self._forward_with_replicated_prefix( q, k, v, ctx_attn_metadata, num_replicated_prefix ) # Ulysses-style All-to-All for sequence/head sharding if sp_size > 1: # -> [B, S, H_local, D] q = _usp_input_all_to_all(q, head_dim=2) k = _usp_input_all_to_all(k, head_dim=2) v = _usp_input_all_to_all(v, head_dim=2) # Ring Attention within subgroups or local attention if get_ring_parallel_world_size() > 1: out = ring_attn( q, k, v, attn_impl=self.attn_impl, is_causal=self.causal, dropout_p=self.dropout_p, ) else: # -> [B, S, H_local, D] out = self.attn_impl.forward(q, k, v, ctx_attn_metadata) # Ulysses-style All-to-All to restore original sharding if sp_size > 1: # -> [B, S_local, H, D] out = _usp_output_all_to_all(out, head_dim=2) return out def _forward_with_replicated_prefix( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, ctx_attn_metadata, num_rep: int, ) -> torch.Tensor: """Ulysses attention where the first *num_rep* tokens are replicated across SP ranks (e.g. text tokens) and should NOT be duplicated by the all-to-all. Strategy: 1. Split q/k/v into replicated prefix and SP-sharded suffix. 2. All-to-all only the sharded suffix (gathers sequence, shards heads). 3. Locally slice the replicated prefix to the same head shard. 4. Concatenate [prefix_h_local, gathered_suffix] and run attention. 5. Split output, all-to-all back the suffix, all-gather prefix heads. """ sp_size = get_ulysses_parallel_world_size() sp_rank = get_sp_parallel_rank() q_rep, q_shard = q[:, :num_rep], q[:, num_rep:] k_rep, k_shard = k[:, :num_rep], k[:, num_rep:] v_rep, v_shard = v[:, :num_rep], v[:, num_rep:] q_shard = _usp_input_all_to_all(q_shard, head_dim=2) k_shard = _usp_input_all_to_all(k_shard, head_dim=2) v_shard = _usp_input_all_to_all(v_shard, head_dim=2) h_local = q_shard.shape[2] h_start = sp_rank * h_local h_end = h_start + h_local q_rep = q_rep[:, :, h_start:h_end, :].contiguous() k_rep = k_rep[:, :, h_start:h_end, :].contiguous() v_rep = v_rep[:, :, h_start:h_end, :].contiguous() q = torch.cat([q_rep, q_shard], dim=1) k = torch.cat([k_rep, k_shard], dim=1) v = torch.cat([v_rep, v_shard], dim=1) out = self.attn_impl.forward(q, k, v, ctx_attn_metadata) out_rep = out[:, :num_rep] out_shard = out[:, num_rep:] out_shard = _usp_output_all_to_all(out_shard, head_dim=2) gathered = [torch.empty_like(out_rep) for _ in range(sp_size)] torch.distributed.all_gather( gathered, out_rep.contiguous(), group=get_sp_group().ulysses_group, ) out_rep = torch.cat(gathered, dim=2) return torch.cat([out_rep, out_shard], dim=1) ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/selector.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/selector.py import os from collections.abc import Generator from contextlib import contextmanager from functools import cache from typing import cast import torch from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionBackend, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.server_args import get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname logger = init_logger(__name__) def backend_name_to_enum(backend_name: str) -> AttentionBackendEnum | None: """ Convert a string backend name to a _Backend enum value. Returns: * _Backend: enum value if backend_name is a valid in-tree type * None: otherwise it's an invalid in-tree type or an out-of-tree platform is loaded. """ assert backend_name is not None return ( AttentionBackendEnum[backend_name] if backend_name in AttentionBackendEnum.__members__ else None ) def get_env_variable_attn_backend() -> AttentionBackendEnum | None: """ Get the backend override specified by the sglang-diffusion attention backend environment variable, if one is specified. Returns: * _Backend enum value if an override is specified * None otherwise """ backend_name = os.environ.get(STR_BACKEND_ENV_VAR) return None if backend_name is None else backend_name_to_enum(backend_name) # Global state allows a particular choice of backend # to be forced, overriding the logic which auto-selects # a backend based on system & workload configuration # (default behavior if this variable is None) # # THIS SELECTION TAKES PRECEDENCE OVER THE # FASTVIDEO ATTENTION BACKEND ENVIRONMENT VARIABLE forced_attn_backend: AttentionBackendEnum | None = None def global_force_attn_backend(attn_backend: AttentionBackendEnum | None) -> None: """ Force all attention operations to use a specified backend. Passing `None` for the argument re-enables automatic backend selection., Arguments: * attn_backend: backend selection (None to revert to auto) """ global forced_attn_backend forced_attn_backend = attn_backend def get_global_forced_attn_backend() -> AttentionBackendEnum | None: """ Get the currently-forced choice of attention backend, or None if auto-selection is currently enabled. """ return forced_attn_backend def get_attn_backend( head_size: int, dtype: torch.dtype, supported_attention_backends: set[AttentionBackendEnum] | None = None, ) -> type[AttentionBackend]: if supported_attention_backends is None: be_tuple = tuple() else: # Sort the backend names to ensure consistent cache key be_tuple = tuple( sorted(list(supported_attention_backends), key=lambda b: b.name) ) return _cached_get_attn_backend(head_size, dtype, be_tuple) @cache def _cached_get_attn_backend( head_size: int, dtype: torch.dtype, supported_attention_backends: tuple[AttentionBackendEnum], ) -> type[AttentionBackend]: # Check whether a particular choice of backend was # previously forced via global_force_attn_backend() or --attention-backend CLI arg. from sglang.multimodal_gen.runtime.platforms import current_platform supported_attention_backends = set(supported_attention_backends) selected_backend = None backend_by_global_setting: AttentionBackendEnum | None = ( get_global_forced_attn_backend() ) if backend_by_global_setting is not None: selected_backend = backend_by_global_setting else: # Check the server arguments for a backend override server_args = get_global_server_args() if server_args.attention_backend is not None: try: selected_backend = AttentionBackendEnum[ server_args.attention_backend.upper() ] except KeyError: raise ValueError( f"Invalid attention backend '{server_args.attention_backend}' specified via command line. " f"Available options are: {[e.name.lower() for e in AttentionBackendEnum]}" ) # get device-specific attn_backend if len(supported_attention_backends) == 0: # all attention backends are allowed pass elif selected_backend is None: logger.debug(f"Attention backend not specified") elif selected_backend not in supported_attention_backends: supported_attention_backends_str = [ supported_attention_backend.__str__() for supported_attention_backend in supported_attention_backends ] logger.debug( f"Selected attention backend: '{selected_backend}' not in supported attention backends: {supported_attention_backends_str}" ) selected_backend = None attention_cls = current_platform.get_attn_backend_cls_str( selected_backend, head_size, dtype ) if not attention_cls: raise ValueError( f"Invalid attention backend for {current_platform.device_name}" ) return cast(type[AttentionBackend], resolve_obj_by_qualname(attention_cls)) @contextmanager def global_force_attn_backend_context_manager( attn_backend: AttentionBackendEnum, ) -> Generator[None, None, None]: """ Globally force a sglang-diffusion attention backend override within a context manager, reverting the global attention backend override to its prior state upon exiting the context manager. Arguments: * attn_backend: attention backend to force Returns: * Generator """ # Save the current state of the global backend override (if any) original_value = get_global_forced_attn_backend() # Globally force the new backend override global_force_attn_backend(attn_backend) # Yield control back to the enclosed code block try: yield finally: # Revert the original global backend override, if any global_force_attn_backend(original_value) ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py ================================================ # copy and modify from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/rcm/utils/a2a_cp.py and https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/core.py from typing import Any, Callable, List, Tuple, Type, Union import torch import torch.distributed as dist from einops import rearrange from torch import Tensor from torch.distributed import ProcessGroup from torch.nn import Module from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionImpl, ) from sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn import ( SageSparseLinearAttentionBackend, SparseLinearAttentionBackend, ) from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend from sglang.multimodal_gen.runtime.managers.forward_context import ( ForwardContext, get_forward_context, ) from sglang.multimodal_gen.runtime.platforms.interface import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import get_compute_dtype logger = init_logger(__name__) def post_all2all(local_seq_2_local_head, seq_world_size): def post_func(input): # b, s, n, h if local_seq_2_local_head: output = rearrange(input, "w bs seq h d -> bs (w seq) h d") else: output = rearrange(input, "w bs s h d -> bs s (w h) d", w=seq_world_size) return output return post_func def single_all_to_all(input, local_seq_2_local_head, group, async_op=False): seq_world_size = dist.get_world_size(group) # b, s, n, h if local_seq_2_local_head: bs, local_seq_len, num_total_head, head_dim = input.shape assert ( num_total_head % seq_world_size == 0 ), f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!" input_t = rearrange( input, "bs seq_len (w h) d -> w bs seq_len h d", w=seq_world_size, h=num_total_head // seq_world_size, ).contiguous() post_all2all_fun = post_all2all(local_seq_2_local_head, seq_world_size) else: bs, global_seq_len, num_local_head, head_dim = input.shape input_t = rearrange( input, "bs (w s) h d -> w bs s h d", w=seq_world_size, s=global_seq_len // seq_world_size, ).contiguous() post_all2all_fun = post_all2all(local_seq_2_local_head, seq_world_size) output = torch.empty_like(input_t) dist.all_to_all_single(output, input_t, group=group, async_op=async_op) res = post_all2all_fun(output) return res def async_a2a_communicate( a2a_inputs: Union[torch.Tensor, List[torch.Tensor]], cp_size: int, cp_group: ProcessGroup, cp_stream: torch.get_device_module().Stream, local_seq_2_local_head: bool, ) -> Union[torch.Tensor, List[torch.Tensor]]: """ A2A communication for context parallelism. best used in communicate qkv Modified from Nvidia Transformer Engine. """ a2a_inputs = [a2a_inputs] if not isinstance(a2a_inputs, list) else a2a_inputs a2a_outputs, a2a_reqs = [None] * len(a2a_inputs), [None] * len(a2a_inputs) a2a_post_fns = [None] * len(a2a_inputs) if local_seq_2_local_head: for i in range(len(a2a_inputs) + 2): if 0 < i < len(a2a_inputs) + 1: a2a_outputs[i - 1] = torch.empty_like(a2a_inputs[i - 1]) a2a_reqs[i - 1] = torch.distributed.all_to_all_single( a2a_outputs[i - 1], a2a_inputs[i - 1], group=cp_group, async_op=True ) a2a_post_fns[i - 1] = post_all2all(local_seq_2_local_head, cp_size) if i > 1: with torch.get_device_module().stream(cp_stream): a2a_reqs[i - 2].wait() a2a_outputs[i - 2] = a2a_post_fns[i - 2](a2a_outputs[i - 2]) if i < len(a2a_inputs): a2a_inputs[i] = rearrange( a2a_inputs[i], "bs seq_len (w h) d -> w bs seq_len h d", w=cp_size ).contiguous() else: for i in range(len(a2a_inputs) + 2): if 0 < i < len(a2a_inputs) + 1: a2a_outputs[i - 1] = torch.empty_like(a2a_inputs[i - 1]) a2a_reqs[i - 1] = torch.distributed.all_to_all_single( a2a_outputs[i - 1], a2a_inputs[i - 1], group=cp_group, async_op=True ) a2a_post_fns[i - 1] = post_all2all(local_seq_2_local_head, cp_size) if i < len(a2a_inputs): a2a_inputs[i] = rearrange( a2a_inputs[i], "bs (w s) h d -> w bs s h d", w=cp_size ).contiguous() if i > 1: with torch.get_device_module().stream(cp_stream): a2a_reqs[i - 2].wait() a2a_outputs[i - 2] = a2a_post_fns[i - 2](a2a_outputs[i - 2]) torch.get_device_module().current_stream().wait_stream(cp_stream) return a2a_outputs[0] if len(a2a_inputs) == 1 else a2a_outputs class _SeqAllToAll(torch.autograd.Function): @staticmethod def forward( ctx: Any, group: dist.ProcessGroup, input: Tensor, local_seq_2_local_head: bool ) -> Tensor: ctx.group = group res = single_all_to_all(input, local_seq_2_local_head, group, False) ctx.local_seq_2_local_head = local_seq_2_local_head return res @staticmethod def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None]: return ( None, _SeqAllToAll.apply(ctx.group, *grad_output, not ctx.local_seq_2_local_head), None, ) class _SeqAllToAllQKV(torch.autograd.Function): @staticmethod def forward( ctx: Any, group: dist.ProcessGroup, q: Tensor, k: Tensor, v: Tensor, cp_size: int, cp_stream: torch.get_device_module().Stream, local_seq_2_local_head: bool, ) -> Tuple[Tensor, Tensor, Tensor]: ctx.group = group ctx.cp_size = cp_size ctx.cp_stream = cp_stream ctx.local_seq_2_local_head = local_seq_2_local_head q, k, v = async_a2a_communicate( [q, k, v], cp_size, group, cp_stream, local_seq_2_local_head ) return q, k, v @staticmethod def backward( ctx: Any, *grad_output: Tensor ) -> Tuple[None, Tensor, Tensor, Tensor, None, None, None]: q_grad, k_grad, v_grad = _SeqAllToAllQKV.apply( ctx.group, *grad_output, ctx.cp_size, ctx.cp_stream, not ctx.local_seq_2_local_head, ) return (None, q_grad, k_grad, v_grad, None, None, None) class DistributedAttention(torch.nn.Module): """Initialization. Arguments: local_attention (Module): local attention with q,k,v sequence_process_group (ProcessGroup): sequence parallel process group """ def __init__(self, local_attention: Union[Module, Callable]) -> None: super(DistributedAttention, self).__init__() self.local_attn = local_attention self.pg = None self.stream = None def forward( self, query: Tensor, key: Tensor, value: Tensor, ctx_attn_metadata ) -> Tensor: """forward Arguments: query (Tensor): query input to the layer key (Tensor): key input to the layer value (Tensor): value input to the layer Returns: * output (Tensor): context output """ if self.pg is None: return self.local_attn(query, key, value, ctx_attn_metadata) pg_size = dist.get_world_size(self.pg) if pg_size < 2: return self.local_attn(query, key, value, ctx_attn_metadata) query_layer, key_layer, value_layer = _SeqAllToAllQKV.apply( self.pg, query, key, value, pg_size, self.stream, True ) context_layer = self.local_attn( query_layer, key_layer, value_layer, ctx_attn_metadata ) output = _SeqAllToAll.apply(self.pg, context_layer, False) return output def set_context_parallel_group(self, group, stream): self.pg = group self.stream = stream class MinimalA2AAttnOp(DistributedAttention): def __init__( self, num_heads: int, head_size: int, attention_type: str, topk: float, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", ): dtype = get_compute_dtype() attn_backend = get_attn_backend( head_size, dtype, supported_attention_backends=supported_attention_backends ) # Maintained for compatibility purposes; can be removed when CI allows setting Attention_backend or when TurboWan supports FA. if attn_backend not in ( SparseLinearAttentionBackend, SageSparseLinearAttentionBackend, ): logger.warning_once( "TurboWan now only supports `sla_attn` or `sage_sla_attn` and has been automatically set to attention_type. Please set --attention-backend to `sla_attn` or `sage_sla_attn`." ) if attention_type == "sagesla": attn_backend = SageSparseLinearAttentionBackend else: attn_backend = SparseLinearAttentionBackend impl_cls: Type["AttentionImpl"] = attn_backend.get_impl_cls() local_attn = impl_cls( num_heads=num_heads, head_size=head_size, topk_ratio=topk, prefix=f"{prefix}.impl", ) super(MinimalA2AAttnOp, self).__init__(local_attn) def set_context_parallel_group(self, process_group, ranks, stream): del ranks super().set_context_parallel_group(process_group, stream) def forward( self, query: Tensor, key: Tensor, value: Tensor, *args: Any, **kwargs ) -> Tensor: forward_context: ForwardContext = get_forward_context() ctx_attn_metadata = forward_context.attn_metadata results = super().forward(query, key, value, ctx_attn_metadata) return rearrange(results, "b ... h l -> b ... (h l)") ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/custom_op.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/custom_op.py from collections.abc import Callable from typing import Any import torch.nn as nn from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) _is_cuda = current_platform.is_cuda() class CustomOp(nn.Module): """ Base class for custom ops. Dispatches the forward method to the appropriate backend. """ def __init__(self) -> None: super().__init__() self._forward_method = self.dispatch_forward() def forward(self, *args, **kwargs) -> Any: return self._forward_method(*args, **kwargs) def forward_native(self, *args, **kwargs) -> Any: """PyTorch-native implementation of the forward method. This method is optional. If implemented, it can be used with compilers such as torch.compile or PyTorch XLA. Also, it can be used for testing purposes. """ raise NotImplementedError def forward_cuda(self, *args, **kwargs) -> Any: raise NotImplementedError def forward_hip(self, *args, **kwargs) -> Any: # ROCm kernels follow the CUDA path by default. return self.forward_cuda(*args, **kwargs) def forward_cpu(self, *args, **kwargs) -> Any: # By default, we assume that CPU ops are compatible with CUDA ops. return self.forward_cuda(*args, **kwargs) def forward_tpu(self, *args, **kwargs) -> Any: # By default, we assume that TPU ops are compatible with the # PyTorch-native implementation. # NOTE(woosuk): This is a placeholder for future extensions. return self.forward_native(*args, **kwargs) def forward_musa(self, *args, **kwargs) -> Any: # MUSA kernels follow the CUDA path by default. return self.forward_cuda(*args, **kwargs) def forward_oot(self, *args, **kwargs) -> Any: # By default, we assume that OOT ops are compatible with the # PyTorch-native implementation. return self.forward_native(*args, **kwargs) def forward_npu(self, *args, **kwargs) -> Any: # By default, we assume that NPU ops are compatible with the # PyTorch-native implementation. return self.forward_native(*args, **kwargs) def dispatch_forward(self) -> Callable: if _is_cuda: return self.forward_cuda elif current_platform.is_hip(): return self.forward_hip elif current_platform.is_npu(): return self.forward_npu elif current_platform.is_xpu(): return self.forward_xpu elif current_platform.is_musa(): return self.forward_musa else: return self.forward_native @classmethod def enabled(cls) -> bool: # since we are not using Inductor, we always return True return True @staticmethod def default_on() -> bool: """ On by default if level < CompilationLevel.PIECEWISE Specifying 'all' or 'none' in custom_op takes precedence. """ raise NotImplementedError # Dictionary of all custom ops (classes, indexed by registered name). # To check if an op with a name is enabled, call .enabled() on the class. # Examples: # - MyOp.enabled() # - op_registry["my_op"].enabled() op_registry: dict[str, type["CustomOp"]] = {} # Decorator to register custom ops. @classmethod def register(cls, name: str) -> Callable: def decorator(op_cls): assert name not in cls.op_registry, f"Duplicate op name: {name}" op_cls.name = name cls.op_registry[name] = op_cls return op_cls return decorator ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/elementwise.py ================================================ import torch from sglang.jit_kernel.diffusion.triton.scale_shift import fuse_scale_shift_kernel from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp class MulAdd(CustomOp): """ Fuse elementwise mul and add Input: a, b, c, OptionalInt[k] Output: a * (k + b) + c """ def __init__(self, prefix: str = ""): super().__init__() def forward_native( self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0 ) -> torch.Tensor: # a.shape: [batch_size, seq_len, inner_dim] if b.dim() == 4: # b.shape: [batch_size, num_frames, 1, inner_dim] num_frames = b.shape[1] frame_seqlen = a.shape[1] // num_frames return c + ( a.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (k + b) ).flatten(1, 2) else: # b.shape: [batch_size, 1, inner_dim] return c + a * (k + b) def forward_cuda( self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0 ): return fuse_scale_shift_kernel(a, b, c, scale_constant=k) ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/layernorm.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/layernorm.py """Custom normalization layers.""" from typing import Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F from sglang.multimodal_gen.runtime.platforms import current_platform _is_cuda = current_platform.is_cuda() _is_npu = current_platform.is_npu() _is_musa = current_platform.is_musa() if _is_cuda: from sgl_kernel import fused_add_rmsnorm, rmsnorm if _is_npu: import torch_npu if _is_musa: from sgl_kernel import fused_add_rmsnorm from sglang.jit_kernel.diffusion.triton.norm import norm_infer, rms_norm_fn from sglang.jit_kernel.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm from sglang.jit_kernel.diffusion.triton.scale_shift import fuse_scale_shift_kernel from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm, fused_inplace_qknorm from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, get_tp_group, ) from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var # Copied and adapted from sglang @CustomOp.register("rms_norm") class RMSNorm(CustomOp): """Root mean square normalization. Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight. Refer to https://arxiv.org/abs/1910.07467 """ def __init__( self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.float32, var_hidden_size: Optional[int] = None, ) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps self.hidden_size = hidden_size self.variance_size_override = ( None if var_hidden_size == hidden_size else var_hidden_size ) if get_bool_env_var("SGLANG_ENABLE_DETERMINISTIC_INFERENCE"): self._forward_method = self.forward_native def forward_triton(self, x: torch.Tensor, residual: Optional[torch.Tensor] = None): return rms_norm_fn( x, self.weight, bias=None, residual=residual, eps=self.variance_epsilon ) def forward_cuda( self, x: torch.Tensor, residual: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: shape = x.shape device = x.device x = x.reshape(-1, shape[-1]) if residual is not None: residual_shape = residual.shape residual = residual.view(-1, shape[-1]) if x.dtype == torch.float: # fp32 out = self.forward_triton(x, residual) if residual is not None: return out[0].view(shape), out[1].view(residual_shape) out = out.view(shape) return out elif self.variance_size_override is not None: return self.forward_native(x, residual) elif residual is not None: fused_add_rmsnorm(x, residual, self.weight.data, self.variance_epsilon) return x.view(shape), residual.view(residual_shape) else: if x.shape[-1] <= 128: out = triton_one_pass_rms_norm( x, self.weight.data, self.variance_epsilon ) else: out = rmsnorm(x, self.weight.data, self.variance_epsilon) out = out.view(shape) return out def forward_native( self, x: torch.Tensor, residual: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: if not x.is_contiguous(): x = x.contiguous() orig_dtype = x.dtype x = x.to(torch.float32) if residual is not None: x = x + residual.to(torch.float32) residual = x.to(orig_dtype) hidden_size = x.shape[-1] if hidden_size != self.hidden_size: raise ValueError( "Expected hidden_size to be " f"{self.hidden_size}, but found: {hidden_size}" ) if self.variance_size_override is None: x_var = x else: if hidden_size < self.variance_size_override: raise ValueError( "Expected hidden_size to be at least " f"{self.variance_size_override}, but found: {hidden_size}" ) x_var = x[..., : self.variance_size_override] variance = x_var.pow(2).mean(dim=-1, keepdim=True) x = x * torch.rsqrt(variance + self.variance_epsilon) x = (x * self.weight).to(orig_dtype) if residual is None: return x else: return x, residual def forward_cpu( self, x: torch.Tensor, residual: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: return self.forward_native(x, residual) def forward_npu( self, x: torch.Tensor, residual: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: if residual is not None: out, _, residual_out = torch_npu.npu_add_rms_norm( residual, x, self.weight.data, self.variance_epsilon ) return out, residual_out return torch_npu.npu_rms_norm(x, self.weight.data, self.variance_epsilon)[0] def forward_hip( self, x: torch.Tensor, residual: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: # ROCm builds of sgl-kernel do not expose rmsnorm custom ops yet. return self.forward_native(x, residual) def _get_weight(self, dtype: torch.dtype) -> torch.Tensor: """Return weight matched to *dtype*. MUSA kernels require input and weight to share the same dtype, unlike CUDA kernels which may handle mixed dtypes internally. """ weight = self.weight.data if weight.dtype != dtype: weight = weight.to(dtype=dtype) return weight def forward_musa( self, x: torch.Tensor, residual: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: shape = x.shape x = x.reshape(-1, shape[-1]) if residual is not None: residual_shape = residual.shape residual = residual.view(-1, shape[-1]) if self.variance_size_override is not None: return self.forward_native(x, residual) elif residual is not None: # fused_add_rmsnorm requires contiguous inputs. if not x.is_contiguous(): x = x.contiguous() if not residual.is_contiguous(): residual = residual.contiguous() weight = self._get_weight(x.dtype) fused_add_rmsnorm(x, residual, weight, self.variance_epsilon) return x.view(shape), residual.view(residual_shape) else: weight = self._get_weight(x.dtype) out = F.rms_norm(x, (self.hidden_size,), weight, self.variance_epsilon) out = out.view(shape) return out def extra_repr(self) -> str: s = f"hidden_size={self.weight.data.size(0)}" s += f", eps={self.variance_epsilon}" return s # Copied and adapted from sglang @CustomOp.register("layer_norm") class LayerNorm(CustomOp): def __init__( self, hidden_size: int, eps=1e-5, bias: bool = True, elementwise_affine=True, device=None, dtype=None, ) -> None: super().__init__() self.eps = eps factory_kwargs = {"device": device, "dtype": dtype} self.hidden_size = hidden_size if elementwise_affine: self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) self.bias = ( torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) if bias else None ) else: self.register_parameter("weight", None) self.register_parameter("bias", None) # Lazy cache for ones vector (not a registered buffer to avoid FSDP/meta issues) self._weight_fallback_cache = None def _get_weight_fallback(self, x: torch.Tensor) -> torch.Tensor: wf = getattr(self, "_weight_fallback_cache", None) if ( wf is None or wf.device != x.device or wf.dtype != x.dtype or wf.numel() != self.hidden_size ): wf = torch.ones(self.hidden_size, device=x.device, dtype=x.dtype) self._weight_fallback_cache = wf return wf def forward_triton(self, x: torch.Tensor): # Fast inference kernel without residual/dropout branches return norm_infer( x.view(-1, self.hidden_size), self.weight, self.bias, eps=self.eps, is_rms_norm=False, ).view(x.shape) def forward_cuda( self, x: torch.Tensor, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: shape = x.shape x = x.view(-1, self.hidden_size) return self.forward_triton(x).view(shape) @torch.compile(backend="inductor", disable=current_platform.is_npu()) def forward_native( self, x: torch.Tensor, residual: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: input_dtype = x.dtype mean = x.mean(-1, keepdim=True) variance = (x - mean).pow(2).mean(-1, keepdim=True) x = (x - mean) * torch.rsqrt(variance + self.eps) if self.weight is not None: x = self.weight * x # if no affine, this is a no-op if self.bias is not None: x = x + self.bias return x.to(input_dtype) def forward_cpu( self, x: torch.Tensor, residual: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: return self.forward_native(x, residual) def forward_musa(self, x: torch.Tensor): return F.layer_norm(x, (self.hidden_size,), self.weight, self.bias, self.eps) def extra_repr(self) -> str: s = f"hidden_size={self.weight.data.size(0)}" s += f", eps={self.variance_epsilon}" return s # adapted from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py # NOTE(will): Needed to match behavior of diffusers and wan2.1 even while using # FSDP's MixedPrecisionPolicy class FP32LayerNorm(nn.LayerNorm): def forward(self, inputs: torch.Tensor) -> torch.Tensor: origin_dtype = inputs.dtype device = inputs.device return F.layer_norm( inputs.float(), self.normalized_shape, self.weight.float().to(device=device) if self.weight is not None else None, self.bias.float().to(device=device) if self.bias is not None else None, self.eps, ).to(origin_dtype) ################################################################################ # Fused norm kernel ################################################################################ def _ensure_contiguous(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: return tensor.contiguous() if tensor is not None else None class _ScaleResidualNormScaleShift(CustomOp): """ Fused kernel that combines: 1. residual_out = residual + gate * x 2. normed = layernorm(residual_out) or rmsnorm(residual_out) 3. out = normed * (1 + scale) + shift compute_dtype is always fp32 for higher precision. """ norm_type: str def __init__( self, hidden_size: int, eps: float = 1e-6, elementwise_affine: bool = False, dtype: torch.dtype = torch.float32, prefix: str = "", ): super().__init__() self.eps = eps self.dtype = dtype if self.norm_type == "rms": self.norm = RMSNorm(hidden_size, eps=eps, dtype=dtype) elif self.norm_type == "layer": self.norm = FP32LayerNorm( hidden_size, elementwise_affine=elementwise_affine, eps=eps, dtype=dtype ) else: raise NotImplementedError(f"Norm type {self.norm_type} not implemented") def forward_cuda( self, residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor | int, shift: torch.Tensor, scale: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: if x.shape[-1] % 256 != 0 and x.shape[-1] <= 8192: import warnings warnings.warn( "FusedScaleResidualNormScaleShift cuda not available, using native fallback", stacklevel=2, ) return self.forward_native(residual, x, gate, shift, scale) from sglang.jit_kernel.diffusion.cutedsl.scale_residual_norm_scale_shift import ( fused_scale_residual_norm_scale_shift, ) if isinstance(gate, int) and gate != 1: raise ValueError( f"Only gate value of 1 is supported for int type, but got {gate}" ) return fused_scale_residual_norm_scale_shift( residual.contiguous(), x.contiguous(), gate.contiguous() if isinstance(gate, torch.Tensor) else None, _ensure_contiguous(getattr(self.norm, "weight", None)), _ensure_contiguous(getattr(self.norm, "bias", None)), scale.contiguous(), shift.contiguous(), self.norm_type, self.eps, ) def forward_hip(self, *args, **kwargs): # ROCm does not support CUDA/CUTLASS-based fused kernels yet, # so we fall back to the native PyTorch implementation. return self.forward_native(*args, **kwargs) def forward_musa(self, *args, **kwargs): # MUSA does not support CUDA/CUTLASS-based fused kernels yet, # so we fall back to the native PyTorch implementation. return self.forward_native(*args, **kwargs) def forward_native( self, residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor | int, shift: torch.Tensor, scale: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: # x.shape: [batch_size, seq_len, inner_dim] if isinstance(gate, int): # used by cross-attention, should be 1 assert gate == 1 residual_output = residual + x elif isinstance(gate, torch.Tensor): if gate.dim() == 4: # gate.shape: [batch_size, num_frames, 1, inner_dim] num_frames = gate.shape[1] frame_seqlen = x.shape[1] // num_frames residual_output = residual + ( x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate ).flatten(1, 2) else: # gate.shape: [batch_size, 1, inner_dim] residual_output = residual + x * gate else: raise ValueError(f"Gate type {type(gate)} not supported") normalized = self.norm(residual_output) modulated = fuse_scale_shift_kernel(normalized, scale, shift) return modulated, residual_output class ScaleResidualLayerNormScaleShift(_ScaleResidualNormScaleShift): norm_type = "layer" class ScaleResidualRMSNormScaleShift(_ScaleResidualNormScaleShift): norm_type = "rms" class _NormScaleShift(CustomOp): """ Fused kernel that combines: 1. normed = layernorm(x) or rmsnorm(x) 2. out = normed * (1 + scale) + shift compute_dtype is always fp32 for higher precision. """ norm_type: str def __init__( self, hidden_size: int, eps: float = 1e-6, elementwise_affine: bool = False, dtype: torch.dtype = torch.float32, prefix: str = "", ): super().__init__() self.eps = eps if self.norm_type == "rms": self.norm = RMSNorm(hidden_size, eps=eps, dtype=dtype) elif self.norm_type == "layer": self.norm = FP32LayerNorm( hidden_size, elementwise_affine=elementwise_affine, eps=eps, dtype=dtype ) else: raise NotImplementedError(f"Norm type {self.norm_type} not implemented") def forward_cuda( self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor ) -> torch.Tensor: if x.shape[-1] % 256 != 0 and x.shape[-1] <= 8192: import warnings warnings.warn( "FusedNormScaleShift cuda not available, using native fallback", stacklevel=2, ) return self.forward_native(x, shift, scale) from sglang.jit_kernel.diffusion.cutedsl.scale_residual_norm_scale_shift import ( fused_norm_scale_shift, ) return fused_norm_scale_shift( x.contiguous(), _ensure_contiguous(getattr(self.norm, "weight", None)), _ensure_contiguous(getattr(self.norm, "bias", None)), scale.contiguous(), shift.contiguous(), self.norm_type, self.eps, ) def forward_hip(self, *args, **kwargs): # ROCm does not support CUDA/CUTLASS-based fused kernels yet, # so we fall back to the native PyTorch implementation. return self.forward_native(*args, **kwargs) def forward_musa(self, *args, **kwargs): # MUSA does not support CUDA/CUTLASS-based fused kernels yet, # so we fall back to the native PyTorch implementation. return self.forward_native(*args, **kwargs) def forward_native( self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor ) -> torch.Tensor: normalized = self.norm(x) modulated = fuse_scale_shift_kernel(normalized, scale, shift) return modulated.to(x.dtype) class LayerNormScaleShift(_NormScaleShift): norm_type = "layer" class RMSNormScaleShift(_NormScaleShift): norm_type = "rms" def apply_qk_norm( q: torch.Tensor, k: torch.Tensor, q_norm: "RMSNorm", k_norm: "RMSNorm", head_dim: int, allow_inplace: bool = True, ) -> Tuple[torch.Tensor, torch.Tensor]: """Apply QK normalization for query and key tensors. Uses JIT fused inplace kernel when available, falls back to standard RMSNorm. """ batch_size = q.size(0) q_eps = q_norm.variance_epsilon k_eps = k_norm.variance_epsilon # Only try fused path on CUDA and when it won't introduce implicit copies. if ( _is_cuda and allow_inplace and (q_eps == k_eps) and can_use_fused_inplace_qknorm(head_dim, q.dtype) ): fused_inplace_qknorm( q=q.view(batch_size, -1, head_dim), k=k.view(batch_size, -1, head_dim), q_weight=q_norm.weight, k_weight=k_norm.weight, head_dim=head_dim, eps=q_eps, ) return q, k q_shape = q.shape k_shape = k.shape q_out = q_norm(q.view(-1, head_dim)).view(q_shape) k_out = k_norm(k.view(-1, head_dim)).view(k_shape) return q_out, k_out def tensor_parallel_rms_norm(x: torch.Tensor, norm: "RMSNorm") -> torch.Tensor: tp_rank = get_tensor_model_parallel_rank() tp_size = get_tensor_model_parallel_world_size() src_dtype = x.dtype weight = norm.weight.tensor_split(tp_size)[tp_rank].float() x_fp32 = x.float() variance = x_fp32.pow(2).mean(dim=-1, keepdim=True) variance = get_tp_group().all_reduce( variance, op=torch._C._distributed_c10d.ReduceOp.AVG ) output = x_fp32 * torch.rsqrt(variance + norm.variance_epsilon) * weight return output.to(dtype=src_dtype) ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/linear.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/linear.py from abc import abstractmethod import torch import torch.distributed as dist import torch.nn.functional as F from torch.nn.parameter import Parameter from sglang.multimodal_gen.runtime.distributed import ( divide, get_tp_group, split_tensor_along_last_dim, tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, QuantizeMethodBase, ) from sglang.multimodal_gen.runtime.layers.utils import get_group_rank, get_group_size # yapf: disable from sglang.multimodal_gen.runtime.models.parameter import ( BasevLLMParameter, BlockQuantScaleParameter, PackedColumnParameter, PackedvLLMParameter, PerTensorScaleParameter, RowvLLMParameter, ) # yapf: enable from sglang.multimodal_gen.runtime.models.utils import set_weight_attrs from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) WEIGHT_LOADER_V2_SUPPORTED = [ "CompressedTensorsLinearMethod", "AWQMarlinLinearMethod", "AWQLinearMethod", "GPTQMarlinLinearMethod", "Fp8LinearMethod", "MarlinLinearMethod", "QQQLinearMethod", "GPTQMarlin24LinearMethod", "TPUInt8LinearMethod", "GPTQLinearMethod", "FBGEMMFp8LinearMethod", "ModelOptFp8LinearMethod", "IPEXAWQLinearMethod", "IPEXGPTQLinearMethod", "HQQMarlinMethod", "QuarkLinearMethod", ] def adjust_scalar_to_fused_array( param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str | int ) -> tuple[torch.Tensor, torch.Tensor]: """For fused modules (QKV and MLP) we have an array of length N that holds 1 scale for each "logical" matrix. So the param is an array of length N. The loaded_weight corresponds to one of the shards on disk. Here, we slice the param based on the shard_id for loading. """ qkv_idxs = {"q": 0, "k": 1, "v": 2} if isinstance(shard_id, str): shard_id = qkv_idxs[shard_id] elif not isinstance(shard_id, int): raise ValueError(f"Unknown Shard Id {shard_id}") # AutoFP8 scales do not have a shape # compressed-tensors scales do have a shape if len(loaded_weight.shape) != 0: assert loaded_weight.shape[0] == 1 loaded_weight = loaded_weight[0] return param[shard_id], loaded_weight class LinearMethodBase(QuantizeMethodBase): """Base class for different (maybe quantized) linear methods.""" @abstractmethod def create_weights( self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: list[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs, ) -> None: """Create weights for a linear layer. The weights will be set as attributes of the layer. Args: layer: The layer that is using the LinearMethodBase factory. input_size_per_partition: Size of the weight input dim on rank X. output_partition_sizes: Sizes of the output dim of each logical weight on rank X. E.g., output_partition_sizes for QKVLinear is a list contains the width of Wq, Wk, Wv on rank X. input_size: Size of the input dim of the weight across all ranks. output_size: Size of the output dim of the weight across all ranks. params_dtype: Datatype of the parameters. """ raise NotImplementedError @abstractmethod def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: """Apply the weights in layer to the input tensor. Expects create_weights to have been called before on the layer.""" raise NotImplementedError class UnquantizedLinearMethod(LinearMethodBase): """Linear method without quantization.""" def create_weights( self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: list[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs, ) -> None: weight = Parameter( torch.empty( sum(output_partition_sizes), input_size_per_partition, dtype=params_dtype, ), requires_grad=False, ) set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) layer.register_parameter("weight", weight) set_weight_attrs(weight, extra_weight_attrs) def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: output = ( F.linear(x, layer.weight, bias) if current_platform.is_amp_supported() or bias is None else F.linear(x, layer.weight, bias.to(x.dtype)) ) # NOTE: this line assumes that we are using amp when using cuda and is needed to account for the fact that amp isn't supported in mps return output class LinearBase(torch.nn.Module): """Base linear layer. Args: input_size: input dimension of the linear layer. output_size: output dimension of the linear layer. skip_bias_add: If true, skip adding bias but instead return it. params_dtype: Data type for the parameters. quant_config: Quantization configure. """ def __init__( self, input_size: int, output_size: int, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", ): super().__init__() # Keep input parameters self.input_size = input_size self.output_size = output_size self.skip_bias_add = skip_bias_add if params_dtype is None: params_dtype = torch.get_default_dtype() self.params_dtype = params_dtype self.quant_config = quant_config self.prefix = prefix if quant_config is None: self.quant_method: QuantizeMethodBase | None = UnquantizedLinearMethod() else: self.quant_method = quant_config.get_quant_method(self, prefix=prefix) def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]: raise NotImplementedError class ReplicatedLinear(LinearBase): """Replicated linear layer. Args: input_size: input dimension of the linear layer. output_size: output dimension of the linear layer. bias: If true, add bias. skip_bias_add: If true, skip adding bias but instead return it. params_dtype: Data type for the parameters. quant_config: Quantization configure. prefix: The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj) """ def __init__( self, input_size: int, output_size: int, bias: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", ): super().__init__( input_size, output_size, skip_bias_add, params_dtype, quant_config, prefix=prefix, ) # All the linear layer supports quant method. assert self.quant_method is not None self.quant_method.create_weights( self, self.input_size, [self.output_size], self.input_size, self.output_size, self.params_dtype, weight_loader=self.weight_loader, ) if bias: self.bias = Parameter( torch.empty( self.output_size, dtype=self.params_dtype, ) ) set_weight_attrs( self.bias, { "output_dim": 0, "weight_loader": self.weight_loader, }, ) else: self.register_parameter("bias", None) def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor) -> None: # If the weight on disk does not have a shape, give it one # (such scales for AutoFp8). if len(loaded_weight.shape) == 0: loaded_weight = loaded_weight.reshape(1) assert param.size() == loaded_weight.size(), ( f"Tried to load weights of size {loaded_weight.size()}" f"to a parameter of size {param.size()}" ) param.data.copy_(loaded_weight) def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]: bias = self.bias if not self.skip_bias_add else None assert self.quant_method is not None output = self.quant_method.apply(self, x, bias) output_bias = self.bias if self.skip_bias_add else None return output, output_bias def extra_repr(self) -> str: s = f"in_features={self.input_size}" s += f", output_features={self.output_size}" s += f", bias={self.bias is not None}" return s class ColumnParallelLinear(LinearBase): """Linear layer with column parallelism. The linear layer is defined as Y = XA + b. A is parallelized along its second dimension as A = [A_1, ..., A_p]. Args: input_size: first dimension of matrix A. output_size: second dimension of matrix A. bias: If true, add bias. gather_output: If true, call all-gather on output and make Y available to all GPUs, otherwise, every GPU will have its output which is Y_i = XA_i skip_bias_add: This was added to enable performance optimizations where bias can be fused with other element-wise operations. we skip adding bias but instead return it. params_dtype: Data type for the parameters. quant_config: Quantization configure. output_sizes: list of output sizes packed into one output, like for QKV the list would be size 3. prefix: The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj) """ def __init__( self, input_size: int, output_size: int, bias: bool = True, gather_output: bool = False, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, output_sizes: list[int] | None = None, prefix: str = "", tp_group: dist.ProcessGroup = None, ): # Divide the weight matrix along the last dimension. self.tp_group = tp_group or get_tp_group() self.tp_size = get_group_size(self.tp_group) self.tp_rank = get_group_rank(self.tp_group) self.input_size_per_partition = input_size self.output_size_per_partition = divide(output_size, self.tp_size) self.output_partition_sizes = [self.output_size_per_partition] # If QKV or MergedColumn, use output size of each partition. if hasattr(self, "output_sizes"): self.output_partition_sizes = [ divide(output_size, self.tp_size) for output_size in self.output_sizes ] super().__init__( input_size, output_size, skip_bias_add, params_dtype, quant_config, prefix ) self.gather_output = gather_output if output_sizes is None: output_sizes = [output_size] assert self.quant_method is not None self.quant_method.create_weights( layer=self, input_size_per_partition=self.input_size_per_partition, output_partition_sizes=self.output_partition_sizes, input_size=self.input_size, output_size=self.output_size, params_dtype=self.params_dtype, weight_loader=( self.weight_loader_v2 if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED else self.weight_loader ), ) if bias: self.bias = Parameter( torch.empty( self.output_size_per_partition, dtype=params_dtype, ) ) set_weight_attrs( self.bias, { "output_dim": 0, "weight_loader": self.weight_loader, }, ) else: self.register_parameter("bias", None) def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor) -> None: tp_rank = self.tp_rank output_dim = getattr(param, "output_dim", None) is_sharded_weight = getattr(param, "is_sharded_weight", False) is_sharded_weight = is_sharded_weight param_data = param.data if output_dim is not None and not is_sharded_weight: shard_size = param_data.shape[output_dim] start_idx = tp_rank * shard_size loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) # Special case for loading scales off disk, which often do not # have a shape (such as in the case of AutoFP8). if len(loaded_weight.shape) == 0: loaded_weight = loaded_weight.reshape(1) assert param_data.shape == loaded_weight.shape param_data.copy_(loaded_weight) def weight_loader_v2(self, param: Parameter, loaded_weight: torch.Tensor) -> None: # Special case for loading scales off disk, which often do not # have a shape (such as in the case of AutoFP8). if len(loaded_weight.shape) == 0: assert loaded_weight.numel() == 1 loaded_weight = loaded_weight.reshape(1) param.load_column_parallel_weight(loaded_weight=loaded_weight) def forward(self, input_: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]: bias = self.bias if not self.skip_bias_add else None # Matrix multiply. assert self.quant_method is not None output_parallel = self.quant_method.apply(self, input_, bias) if self.gather_output: # All-gather across the partitions. output = tensor_model_parallel_all_gather( output_parallel, tp_group=self.tp_group ) else: output = output_parallel output_bias = self.bias if self.skip_bias_add else None return output, output_bias def extra_repr(self) -> str: s = f"in_features={self.input_size}" s += f", output_features={self.output_size_per_partition}" s += f", bias={self.bias is not None}" s += f", tp_size={self.tp_size}" s += f", gather_output={self.gather_output}" return s class MergedColumnParallelLinear(ColumnParallelLinear): """Packed linear layers with column parallelism. Similar to ColumnParallelLinear, but the weight matrix is concatenated along the output dimension. When the weight matrix is loaded, the different partitions are sharded separately. Args: input_size: input dimension of the linear layer. output_sizes: list of output dimensions of the linear layer. bias: If true, add bias. gather_output: If true, call all-gather on output and make the output available to all GPUs, otherwise, every GPU will have its own output. skip_bias_add: This was added to enable performance optimizations where bias can be fused with other element-wise operations. we skip adding bias but instead return it. params_dtype: Data type for the parameters. quant_config: Quantization configure. prefix: The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj) """ def __init__( self, input_size: int, output_sizes: list[int], bias: bool = True, gather_output: bool = False, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", tp_group: dist.ProcessGroup = None, ): super().__init__( input_size=input_size, output_size=sum(output_sizes), bias=bias, gather_output=gather_output, skip_bias_add=skip_bias_add, params_dtype=params_dtype, quant_config=quant_config, prefix=prefix, tp_group=tp_group, ) self.output_sizes = output_sizes assert all(output_size % self.tp_size == 0 for output_size in output_sizes) def weight_loader( self, param: Parameter, loaded_weight: torch.Tensor, loaded_shard_id: int | None = None, ) -> None: param_data = param.data output_dim = getattr(param, "output_dim", None) # Special case for AQLM codebooks. is_metadata = getattr(param, "is_metadata", False) # Special case for per-tensor scale to load scalar into fused array. needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False) if loaded_shard_id is None: # Loaded weight is already fused on disk (mlp). # (e.g., Phi-3's gate_up_proj). if output_dim is None: if needs_scalar_to_array: param_data, loaded_weight = adjust_scalar_to_fused_array( param_data, loaded_weight, 0 ) assert param_data.shape == loaded_weight.shape param_data.copy_(loaded_weight) return current_shard_offset = 0 shard_offsets: list[tuple[int, int, int]] = [] for i, output_size in enumerate(self.output_sizes): shard_offsets.append((i, current_shard_offset, output_size)) current_shard_offset += output_size for shard_id, shard_offset, shard_size in shard_offsets: loaded_weight_shard = loaded_weight.narrow( output_dim, shard_offset, shard_size ) self.weight_loader(param, loaded_weight_shard, shard_id) return assert loaded_shard_id < len(self.output_sizes) tp_rank = self.tp_rank tp_size = self.tp_size if output_dim is not None: shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size shard_size = self.output_sizes[loaded_shard_id] // tp_size is_sharded_weight = getattr(param, "is_sharded_weight", False) # bitsandbytes loads the weights of the specific portion # no need to narrow is_sharded_weight = is_sharded_weight param_data = param_data.narrow(output_dim, shard_offset, shard_size) start_idx = tp_rank * shard_size if not is_sharded_weight: loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) # Special case for AQLM codebooks. elif is_metadata: # metadata indicates fixed size concatenated along dim 0 shard_size = loaded_weight.shape[0] shard_offset = loaded_shard_id * shard_size param_data = param_data.narrow(0, shard_offset, shard_size) # Special case for per-tensor scales in fused case. elif needs_scalar_to_array: param_data, loaded_weight = adjust_scalar_to_fused_array( param_data, loaded_weight, loaded_shard_id ) else: ignore_warning = getattr(param, "ignore_warning", False) if not ignore_warning: logger.warning( "Loading a weight without `output_dim` attribute in " "MergedColumnParallelLinear, assume the weight is " "the same for all partitions." ) assert param_data.shape == loaded_weight.shape param_data.copy_(loaded_weight) def _load_fused_module_from_checkpoint( self, param: BasevLLMParameter, loaded_weight: torch.Tensor ) -> None: """ Handle special case for models where MLP layers are already fused on disk. In this case, we have no shard id. This function determmines the shard id by splitting these layers and then calls the weight loader using the shard id. An example of a model with these fused layers: https://huggingface.co/microsoft/Phi-3-mini-4k-instruct """ current_shard_offset = 0 shard_offsets: list[tuple[int, int, int]] = [] for i, output_size in enumerate(self.output_sizes): shard_offsets.append((i, current_shard_offset, output_size)) current_shard_offset += output_size for shard_id, shard_offset, shard_size in shard_offsets: # Special case for Quantization. # If quantized, we need to adjust the offset and size to account # for the packing. if ( isinstance(param, PackedColumnParameter | PackedvLLMParameter) and param.packed_dim == param.output_dim ): shard_size, shard_offset = param.adjust_shard_indexes_for_packing( shard_size=shard_size, shard_offset=shard_offset ) loaded_weight_shard = loaded_weight.narrow( param.output_dim, shard_offset, shard_size ) self.weight_loader_v2(param, loaded_weight_shard, shard_id) def weight_loader_v2( self, param: BasevLLMParameter, loaded_weight: torch.Tensor, loaded_shard_id: int | None = None, ) -> None: if loaded_shard_id is None: if isinstance(param, PerTensorScaleParameter): param.load_merged_column_weight(loaded_weight=loaded_weight, shard_id=0) return elif type(param) in (RowvLLMParameter, BasevLLMParameter): param.load_merged_column_weight(loaded_weight=loaded_weight) return # TODO: @dsikka - move to parameter.py self._load_fused_module_from_checkpoint(param, loaded_weight) return assert loaded_shard_id < len(self.output_sizes) tp_size = self.tp_size if isinstance(param, BlockQuantScaleParameter): raise NotImplementedError("FP8 is not implemented yet") # FIXME(will): add fp8 support # from vllm.model_executor.layers.quantization.fp8 import ( # Fp8LinearMethod, Fp8MoEMethod) # assert self.quant_method is not None # assert isinstance(self.quant_method, # (Fp8LinearMethod, Fp8MoEMethod)) # weight_block_size = self.quant_method.quant_config.weight_block_size # assert weight_block_size is not None # block_n, _ = weight_block_size[0], weight_block_size[1] # shard_offset = ( # (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) // # block_n) // tp_size # shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) // # block_n // tp_size) else: shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size shard_size = self.output_sizes[loaded_shard_id] // tp_size param.load_merged_column_weight( loaded_weight=loaded_weight, shard_id=loaded_shard_id, shard_offset=shard_offset, shard_size=shard_size, ) class QKVParallelLinear(ColumnParallelLinear): """Linear layers for the attention's QKV transformation. Linear layers for the linear transformation of the query, key, and value vectors in the attention layer. The weight matrix is concatenated along the output dimension. The layer is parallelized along the head dimension. When the number of key/value heads is smaller than the number of query heads (e.g., multi-query/grouped-query attention), the key/value head may be replicated while the query heads are partitioned. Args: hidden_size: input hidden state size of the transformer. head_size: size of each attention head. total_num_heads: total number of attention query heads. total_num_kv_heads: total number of attention key/value heads. If None, assume total_num_kv_heads = total_num_heads. bias: If true, add bias. skip_bias_add: This was added to enable performance optimizations where bias can be fused with other element-wise operations. we skip adding bias but instead return it. params_dtype: Data type for the parameters. quant_config: Quantization configure. prefix: The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj) """ def __init__( self, hidden_size: int, head_size: int, total_num_heads: int, total_num_kv_heads: int | None = None, bias: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", tp_group: dist.ProcessGroup = None, ): self.hidden_size = hidden_size self.head_size = head_size self.total_num_heads = total_num_heads if total_num_kv_heads is None: total_num_kv_heads = total_num_heads self.total_num_kv_heads = total_num_kv_heads # Divide the weight matrix along the last dimension. tp_group = tp_group or get_tp_group() tp_size = get_group_size(tp_group) self.num_heads = divide(self.total_num_heads, tp_size) if tp_size >= self.total_num_kv_heads: self.num_kv_heads = 1 self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads) else: self.num_kv_heads = divide(self.total_num_kv_heads, tp_size) self.num_kv_head_replicas = 1 input_size = self.hidden_size output_size = ( (self.num_heads + 2 * self.num_kv_heads) * tp_size * self.head_size ) self.output_sizes = [ self.num_heads * self.head_size * tp_size, # q_proj self.num_kv_heads * self.head_size * tp_size, # k_proj self.num_kv_heads * self.head_size * tp_size, # v_proj ] super().__init__( input_size=input_size, output_size=output_size, bias=bias, gather_output=False, skip_bias_add=skip_bias_add, params_dtype=params_dtype, quant_config=quant_config, prefix=prefix, tp_group=tp_group, ) def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int | None: shard_offset_mapping = { "q": 0, "k": self.num_heads * self.head_size, "v": (self.num_heads + self.num_kv_heads) * self.head_size, "total": (self.num_heads + 2 * self.num_kv_heads) * self.head_size, } return shard_offset_mapping.get(loaded_shard_id) def _get_shard_size_mapping(self, loaded_shard_id: str) -> int | None: shard_size_mapping = { "q": self.num_heads * self.head_size, "k": self.num_kv_heads * self.head_size, "v": self.num_kv_heads * self.head_size, } return shard_size_mapping.get(loaded_shard_id) def _load_fused_module_from_checkpoint( self, param: BasevLLMParameter, loaded_weight: torch.Tensor ): """ Handle special case for models where QKV layers are already fused on disk. In this case, we have no shard id. This function determmines the shard id by splitting these layers and then calls the weight loader using the shard id. An example of a model with these fused layers: https://huggingface.co/microsoft/Phi-3-mini-4k-instruct """ shard_offsets = [ # (shard_id, shard_offset, shard_size) ("q", 0, self.total_num_heads * self.head_size), ( "k", self.total_num_heads * self.head_size, self.total_num_kv_heads * self.head_size, ), ( "v", (self.total_num_heads + self.total_num_kv_heads) * self.head_size, self.total_num_kv_heads * self.head_size, ), ] for shard_id, shard_offset, shard_size in shard_offsets: # Special case for Quantization. # If quantized, we need to adjust the offset and size to account # for the packing. if ( isinstance(param, PackedColumnParameter | PackedvLLMParameter) and param.packed_dim == param.output_dim ): shard_size, shard_offset = param.adjust_shard_indexes_for_packing( shard_size=shard_size, shard_offset=shard_offset ) loaded_weight_shard = loaded_weight.narrow( param.output_dim, shard_offset, shard_size ) self.weight_loader_v2(param, loaded_weight_shard, shard_id) def weight_loader_v2( self, param: BasevLLMParameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None, ): if loaded_shard_id is None: # special case for certain models if isinstance(param, PerTensorScaleParameter): param.load_qkv_weight(loaded_weight=loaded_weight, shard_id=0) return elif type(param) in (RowvLLMParameter, BasevLLMParameter): param.load_qkv_weight(loaded_weight=loaded_weight) return # TODO: @dsikka - move to parameter.py self._load_fused_module_from_checkpoint(param, loaded_weight) return assert loaded_shard_id in ["q", "k", "v"] shard_offset = self._get_shard_offset_mapping(loaded_shard_id) shard_size = self._get_shard_size_mapping(loaded_shard_id) param.load_qkv_weight( loaded_weight=loaded_weight, num_heads=self.num_kv_head_replicas, shard_id=loaded_shard_id, shard_offset=shard_offset, shard_size=shard_size, ) def weight_loader( self, param: Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None, ): param_data = param.data output_dim = getattr(param, "output_dim", None) # Special case for AQLM codebooks. is_metadata = getattr(param, "is_metadata", False) # Special case for per-tensor scales in fused case. needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False) if loaded_shard_id is None: # Loaded weight is already fused on disk (qkv). # (e.g., Phi-3's qkv_proj). if output_dim is None: if needs_scalar_to_array: param_data, loaded_weight = adjust_scalar_to_fused_array( param_data, loaded_weight, 0 ) assert param_data.shape == loaded_weight.shape param_data.copy_(loaded_weight) return shard_offsets = [ # (shard_id, shard_offset, shard_size) ("q", 0, self.total_num_heads * self.head_size), ( "k", self.total_num_heads * self.head_size, self.total_num_kv_heads * self.head_size, ), ( "v", (self.total_num_heads + self.total_num_kv_heads) * self.head_size, self.total_num_kv_heads * self.head_size, ), ] for shard_id, shard_offset, shard_size in shard_offsets: loaded_weight_shard = loaded_weight.narrow( output_dim, shard_offset, shard_size ) self.weight_loader(param, loaded_weight_shard, shard_id) return tp_rank = self.tp_rank assert loaded_shard_id in ["q", "k", "v"] # If output dim is defined, use the default loading process. if output_dim is not None: if loaded_shard_id == "q": shard_offset = 0 shard_size = self.num_heads * self.head_size elif loaded_shard_id == "k": shard_offset = self.num_heads * self.head_size shard_size = self.num_kv_heads * self.head_size elif loaded_shard_id == "v": shard_offset = (self.num_heads + self.num_kv_heads) * self.head_size shard_size = self.num_kv_heads * self.head_size is_sharded_weight = getattr(param, "is_sharded_weight", False) # bitsandbytes loads the weights of the specific portion # no need to narrow is_sharded_weight = is_sharded_weight shard_idx = 0 param_data = param_data.narrow(output_dim, shard_offset, shard_size) if loaded_shard_id == "q": shard_idx = tp_rank else: shard_idx = tp_rank // self.num_kv_head_replicas start_idx = shard_idx * shard_size if not is_sharded_weight: loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) # Special case for for AQLM codebooks. elif is_metadata: # metadata indicates fixed size concatenated along dim 0 shard_size = loaded_weight.shape[0] shard_index = ["q", "k", "v"].index(loaded_shard_id) param_data = param_data.narrow(0, shard_index * shard_size, shard_size) # Special case for per-tensor scales in fused case. elif needs_scalar_to_array: param_data, loaded_weight = adjust_scalar_to_fused_array( param_data, loaded_weight, loaded_shard_id ) else: ignore_warning = getattr(param, "ignore_warning", False) if not ignore_warning: logger.warning( "Loading a weight without `output_dim` attribute in " "QKVParallelLinear, assume the weight is the same " "for all partitions." ) assert param_data.shape == loaded_weight.shape param_data.copy_(loaded_weight) class RowParallelLinear(LinearBase): """Linear layer with row parallelism. The linear layer is defined as Y = XA + b. A is parallelized along its first dimension and X along its second dimension as: - - | A_1 | | . | A = | . | X = [X_1, ..., X_p] | . | | A_p | - - Arguments: input_size: first dimension of matrix A. output_size: second dimension of matrix A. bias: If true, add bias. Note that bias is not parallelized. input_is_parallel: If true, we assume that the input is already split across the GPUs and we do not split again. skip_bias_add: This was added to enable performance optimization where bias can be fused with other element-wise operations. We skip adding bias but instead return it. params_dtype: Data type for the parameters. quant_config: Quantization configure. """ def __init__( self, input_size: int, output_size: int, bias: bool = True, input_is_parallel: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, reduce_results: bool = True, quant_config: QuantizationConfig | None = None, prefix: str = "", tp_group: dist.ProcessGroup = None, ): # Divide the weight matrix along the first dimension. self.tp_group = tp_group or get_tp_group() self.tp_rank = get_group_rank(self.tp_group) self.tp_size = get_group_size(self.tp_group) self.input_size_per_partition = divide(input_size, self.tp_size) self.output_size_per_partition = output_size self.output_partition_sizes = [output_size] super().__init__( input_size, output_size, skip_bias_add, params_dtype, quant_config, prefix ) self.input_is_parallel = input_is_parallel self.reduce_results = reduce_results assert self.quant_method is not None self.quant_method.create_weights( layer=self, input_size_per_partition=self.input_size_per_partition, output_partition_sizes=self.output_partition_sizes, input_size=self.input_size, output_size=self.output_size, params_dtype=self.params_dtype, weight_loader=( self.weight_loader_v2 if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED else self.weight_loader ), ) if not reduce_results and (bias and not skip_bias_add): raise ValueError( "When not reduce the results, adding bias to the " "results can lead to incorrect results" ) if bias: self.bias = Parameter(torch.empty(self.output_size, dtype=params_dtype)) set_weight_attrs( self.bias, { "output_dim": 0, "weight_loader": self.weight_loader, }, ) else: self.register_parameter("bias", None) def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): tp_rank = self.tp_rank input_dim = getattr(param, "input_dim", None) is_sharded_weight = getattr(param, "is_sharded_weight", False) # bitsandbytes loads the weights of the specific portion # no need to narrow is_sharded_weight = is_sharded_weight param_data = param.data if input_dim is not None and not is_sharded_weight: shard_size = param_data.shape[input_dim] start_idx = tp_rank * shard_size loaded_weight = loaded_weight.narrow(input_dim, start_idx, shard_size) # Special case for loading scales off disk, which often do not # have a shape (such as in the case of AutoFP8). if len(loaded_weight.shape) == 0: loaded_weight = loaded_weight.reshape(1) assert param_data.shape == loaded_weight.shape param_data.copy_(loaded_weight) def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor): # Special case for loading scales off disk, which often do not # have a shape (such as in the case of AutoFP8). if len(loaded_weight.shape) == 0: assert loaded_weight.numel() == 1 loaded_weight = loaded_weight.reshape(1) param.load_row_parallel_weight(loaded_weight=loaded_weight) def forward(self, input_) -> tuple[torch.Tensor, Parameter | None]: if self.input_is_parallel: input_parallel = input_ else: tp_rank = self.tp_rank splitted_input = split_tensor_along_last_dim( input_, num_partitions=self.tp_size ) input_parallel = splitted_input[tp_rank].contiguous() # Matrix multiply. assert self.quant_method is not None # Only fuse bias add into GEMM for rank 0 (this ensures that # bias will not get added more than once in TP>1 case) bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_) if self.reduce_results and self.tp_size > 1: output = tensor_model_parallel_all_reduce( output_parallel, tp_group=self.tp_group ) else: output = output_parallel output_bias = self.bias if self.skip_bias_add else None return output, output_bias def extra_repr(self) -> str: s = f"input_features={self.input_size_per_partition}" s += f", output_features={self.output_size}" s += f", bias={self.bias is not None}" s += f", tp_size={self.tp_size}" s += f", reduce_results={self.reduce_results}" return s ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/lora/linear.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Code adapted from SGLang https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/lora/layers.py import torch from torch import nn from torch.distributed._composable.fsdp import ( CPUOffloadPolicy, OffloadPolicy, fully_shard, ) from torch.distributed.tensor import DTensor from sglang.multimodal_gen.runtime.distributed import ( get_local_torch_device, get_tp_rank, split_tensor_along_last_dim, tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce, ) from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, LinearBase, MergedColumnParallelLinear, QKVParallelLinear, ReplicatedLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from sglang.multimodal_gen.utils import get_mixed_precision_state torch._dynamo.config.recompile_limit = 16 class BaseLayerWithLoRA(nn.Module): def __init__( self, base_layer: nn.Module, lora_rank: int | None = None, lora_alpha: int | None = None, ): super().__init__() self.base_layer: nn.Module = base_layer self.merged: bool = False # Immutable base-weight snapshot; `to("cpu")` may alias CPU storage. # Use `clone()` so merge updates cannot mutate this backup tensor. self.cpu_weight = base_layer.weight.detach().to("cpu").clone() # indicates adapter weights don't contain this layer # (which shouldn't normally happen, but we want to separate it from the case of erroneous merging) # Default to True to prevent using uninitialized weights; set to False when weights are loaded self.disable_lora: bool = True self.lora_rank = lora_rank self.lora_alpha = lora_alpha self.lora_weights_list: list[ tuple[torch.nn.Parameter, torch.nn.Parameter, str | None, float] ] = [] self.lora_path: str | None = None self.strength: float = 1.0 self.lora_A = None self.lora_B = None @property def weight(self): return self.base_layer.weight @property def bias(self): return getattr(self.base_layer, "bias", None) @torch.compile() def forward(self, x: torch.Tensor) -> torch.Tensor: lora_A = self.lora_A lora_B = self.lora_B if isinstance(self.lora_B, DTensor): lora_B = self.lora_B.to_local() lora_A = self.lora_A.to_local() # TODO: Support multiple LoRA adapters when use not merged mode if not self.merged and not self.disable_lora: lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True)) lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True)) delta = x @ lora_A_sliced.T @ lora_B_sliced.T if self.lora_alpha != self.lora_rank: delta = delta * ( self.lora_alpha / self.lora_rank # type: ignore ) # type: ignore delta = delta * self.strength if delta.dim() > 2: delta = delta.reshape(-1, delta.shape[-1]) out, output_bias = self.base_layer(x) return out + delta, output_bias else: out, output_bias = self.base_layer(x) return out, output_bias def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: return A def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: return B def set_lora_weights( self, A: torch.Tensor, B: torch.Tensor, lora_path: str | None = None, strength: float = 1.0, clear_existing: bool = False, ) -> None: """ Set LoRA weights. Supports multiple LoRA adapters. Args: A: LoRA A weight tensor B: LoRA B weight tensor lora_path: Path to the LoRA adapter (for logging) strength: LoRA strength clear_existing: If True, clear existing LoRA weights before adding new one. If False, append to existing list (for multi-LoRA support). """ lora_A_param = torch.nn.Parameter( A ) # share storage with weights in the pipeline lora_B_param = torch.nn.Parameter(B) if clear_existing: self.lora_weights_list.clear() # Also clear backward compatibility attributes self.lora_A = None self.lora_B = None self.lora_path = None self.strength = 1.0 # Add to list for multi-LoRA support self.lora_weights_list.append((lora_A_param, lora_B_param, lora_path, strength)) # Set backward compatibility attributes to point to the last LoRA (for single LoRA case) # This ensures backward compatibility while supporting multiple LoRA self.lora_A = lora_A_param self.lora_B = lora_B_param self.lora_path = lora_path self.strength = strength self.disable_lora = False self.merge_lora_weights() @torch.no_grad() def _merge_lora_into_data( self, data: torch.Tensor, lora_list: list[ tuple[torch.nn.Parameter, torch.nn.Parameter, str | None, float] ], ) -> None: """ Merge all LoRA adapters into the data tensor in-place. Args: data: The base weight tensor to merge LoRA into (modified in-place) lora_list: List of (lora_A, lora_B, lora_path, lora_strength) tuples """ # Merge all LoRA adapters in order for lora_A, lora_B, _, lora_strength in lora_list: lora_delta = self.slice_lora_b_weights( lora_B.to(data) ) @ self.slice_lora_a_weights(lora_A.to(data)) # Apply lora_alpha / lora_rank scaling for consistency with forward() if self.lora_alpha is not None and self.lora_rank is not None: if self.lora_alpha != self.lora_rank: lora_delta = lora_delta * (self.lora_alpha / self.lora_rank) if lora_delta.dim() > 2: lora_delta = lora_delta.reshape(-1, lora_delta.shape[-1]) data += lora_strength * lora_delta @torch.no_grad() def merge_lora_weights(self, strength: float | None = None) -> None: if strength is not None: self.strength = strength if self.disable_lora: return if self.merged: self.unmerge_lora_weights() # Use lora_weights_list if available, otherwise fall back to single LoRA for backward compatibility lora_list = self.lora_weights_list if self.lora_weights_list else [] if not lora_list and self.lora_A is not None and self.lora_B is not None: lora_list = [(self.lora_A, self.lora_B, self.lora_path, self.strength)] if not lora_list: raise ValueError("LoRA weights not set. Please set them first.") if isinstance(self.base_layer.weight, DTensor): mesh = self.base_layer.weight.data.device_mesh unsharded_base_layer = ReplicatedLinear( input_size=self.base_layer.input_size, output_size=self.base_layer.output_size, bias=getattr(self.base_layer, "bias", None) is not None, skip_bias_add=self.base_layer.skip_bias_add, params_dtype=self.base_layer.params_dtype, quant_config=self.base_layer.quant_config, prefix=self.base_layer.prefix, ) # Using offload param is on CPU, so current_device is for "CPU -> GPU -> merge -> CPU" current_device = self.base_layer.weight.data.device data = self.base_layer.weight.data.to( get_local_torch_device() ).full_tensor() self._merge_lora_into_data(data, lora_list) unsharded_base_layer.weight = nn.Parameter(data.to(current_device)) if isinstance(getattr(self.base_layer, "bias", None), DTensor): unsharded_base_layer.bias = nn.Parameter( self.base_layer.bias.to(get_local_torch_device(), non_blocking=True) .full_tensor() .to(current_device) ) offload_policy = ( CPUOffloadPolicy() if "cpu" in str(current_device) else OffloadPolicy() ) mp_policy = get_mixed_precision_state().mp_policy self.base_layer = fully_shard( unsharded_base_layer, mesh=mesh, mp_policy=mp_policy, offload_policy=offload_policy, ) else: current_device = self.base_layer.weight.data.device data = self.base_layer.weight.data.to(get_local_torch_device()) self._merge_lora_into_data(data, lora_list) self.base_layer.weight.data = data.to(current_device, non_blocking=True) self.merged = True @torch.no_grad() # @torch.compile(dynamic=True) def unmerge_lora_weights(self) -> None: if self.disable_lora: return if not self.merged: raise ValueError( "LoRA weights not merged. Please merge them first before unmerging." ) # avoid precision loss if isinstance(self.base_layer.weight, DTensor): device = self.base_layer.weight.data.device old_weight = self.base_layer.weight new_weight_data = self.cpu_weight.to(device, non_blocking=True) self.base_layer.weight = nn.Parameter(new_weight_data) del old_weight else: current_device = self.base_layer.weight.data.device cpu_weight_on_device = self.cpu_weight.to(current_device, non_blocking=True) self.base_layer.weight.data.copy_(cpu_weight_on_device) if ( cpu_weight_on_device.data_ptr() != self.base_layer.weight.data.data_ptr() ): del cpu_weight_on_device self.merged = False class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA): """ Vocab parallel embedding layer with support for LoRA (Low-Rank Adaptation). Note: The current version does not yet implement the LoRA functionality. This class behaves exactly the same as the base VocabParallelEmbedding. Future versions will integrate LoRA functionality to support efficient parameter fine-tuning. """ def __init__( self, base_layer: VocabParallelEmbedding, ) -> None: super().__init__(base_layer) def forward(self, input_: torch.Tensor) -> torch.Tensor: raise NotImplementedError( "We don't support VocabParallelEmbeddingWithLoRA yet." ) class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA): def __init__( self, base_layer: ColumnParallelLinear, lora_rank: int | None = None, lora_alpha: int | None = None, ) -> None: super().__init__(base_layer, lora_rank, lora_alpha) def forward(self, input_: torch.Tensor) -> torch.Tensor: # duplicate the logic in ColumnParallelLinear bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None output_parallel = self.base_layer.quant_method.apply( self.base_layer, input_, bias ) if self.base_layer.gather_output: output = tensor_model_parallel_all_gather(output_parallel) else: output = output_parallel output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None return output, output_bias def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: return A def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: tp_rank = get_tp_rank() shard_size = self.base_layer.output_partition_sizes[0] start_idx = tp_rank * shard_size end_idx = (tp_rank + 1) * shard_size B = B[start_idx:end_idx, :] return B class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): def __init__( self, base_layer: MergedColumnParallelLinear, lora_rank: int | None = None, lora_alpha: int | None = None, ) -> None: super().__init__(base_layer, lora_rank, lora_alpha) def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: return A def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: tp_rank = get_tp_rank() # Since the outputs for both gate and up are identical, we use a random one. shard_size = self.base_layer.output_partition_sizes[0] start_idx = tp_rank * shard_size end_idx = (tp_rank + 1) * shard_size return B[:, start_idx:end_idx, :] class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): def __init__( self, base_layer: QKVParallelLinear, lora_rank: int | None = None, lora_alpha: int | None = None, ) -> None: super().__init__(base_layer, lora_rank, lora_alpha) def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: return A def slice_lora_b_weights( self, B: list[torch.Tensor] ) -> tuple[torch.Tensor, torch.Tensor]: tp_rank = get_tp_rank() B_q, B_kv = B base_layer = self.base_layer q_proj_shard_size = base_layer.q_proj_shard_size kv_proj_shard_size = base_layer.kv_proj_shard_size num_kv_head_replicas = base_layer.num_kv_head_replicas q_start_idx = q_proj_shard_size * tp_rank q_end_idx = q_start_idx + q_proj_shard_size kv_shard_id = tp_rank // num_kv_head_replicas kv_start_idx = kv_proj_shard_size * kv_shard_id kv_end_idx = kv_start_idx + kv_proj_shard_size return B_q[q_start_idx:q_end_idx, :], B_kv[:, kv_start_idx:kv_end_idx, :] class RowParallelLinearWithLoRA(BaseLayerWithLoRA): def __init__( self, base_layer: RowParallelLinear, lora_rank: int | None = None, lora_alpha: int | None = None, ) -> None: super().__init__(base_layer, lora_rank, lora_alpha) def forward(self, input_: torch.Tensor): # duplicate the logic in RowParallelLinear if self.base_layer.input_is_parallel: input_parallel = input_ else: tp_rank = get_tp_rank() splitted_input = split_tensor_along_last_dim( input_, num_partitions=self.base_layer.tp_size ) input_parallel = splitted_input[tp_rank].contiguous() output_parallel = self.base_layer.quant_method.apply( self.base_layer, input_parallel ) if self.base_layer.reduce_results and self.base_layer.tp_size > 1: output_ = tensor_model_parallel_all_reduce(output_parallel) else: output_ = output_parallel if not self.base_layer.skip_bias_add: output = ( output_ + self.base_layer.bias if self.base_layer.bias is not None else output_ ) output_bias = None else: output = output_ output_bias = self.base_layer.bias return output, output_bias def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: tp_rank = get_tp_rank() shard_size = self.base_layer.input_size_per_partition start_idx = tp_rank * shard_size end_idx = (tp_rank + 1) * shard_size A = A[:, start_idx:end_idx].contiguous() return A def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: return B class LinearWithLoRA(BaseLayerWithLoRA): """ Wrapper for standard torch.nn.Linear to support LoRA. Unlike custom LinearBase classes, nn.Linear.forward() returns a single tensor, not a tuple of (output, bias). """ def __init__( self, base_layer: nn.Linear, lora_rank: int | None = None, lora_alpha: int | None = None, ) -> None: super().__init__(base_layer, lora_rank, lora_alpha) @torch.compile() def forward(self, x: torch.Tensor) -> torch.Tensor: lora_A = self.lora_A lora_B = self.lora_B if isinstance(self.lora_B, DTensor): lora_B = self.lora_B.to_local() lora_A = self.lora_A.to_local() # TODO: Support multiple LoRA adapters when use not merged mode if not self.merged and not self.disable_lora: lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True)) lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True)) delta = x @ lora_A_sliced.T @ lora_B_sliced.T if self.lora_alpha != self.lora_rank: delta = delta * ( self.lora_alpha / self.lora_rank # type: ignore ) # type: ignore delta = delta * self.strength if delta.dim() > 2: delta = delta.reshape(-1, delta.shape[-1]) # nn.Linear.forward() returns a single tensor, not a tuple out = self.base_layer(x) return out + delta else: # nn.Linear.forward() returns a single tensor out = self.base_layer(x) return out def wrap_with_lora_layer( layer: nn.Module, lora_rank: int | None = None, lora_alpha: int | None = None, ) -> BaseLayerWithLoRA | None: """ transform the given layer to its corresponding LoRA layer """ supported_layer_types: dict[ type[LinearBase] | type[nn.Linear], type[BaseLayerWithLoRA] ] = { # the order matters # VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA, QKVParallelLinear: QKVParallelLinearWithLoRA, MergedColumnParallelLinear: MergedColumnParallelLinearWithLoRA, ColumnParallelLinear: ColumnParallelLinearWithLoRA, RowParallelLinear: RowParallelLinearWithLoRA, ReplicatedLinear: BaseLayerWithLoRA, nn.Linear: LinearWithLoRA, } for src_layer_type, lora_layer_type in supported_layer_types.items(): if isinstance(layer, src_layer_type): # type: ignore[arg-type] ret = lora_layer_type( layer, lora_rank=lora_rank, lora_alpha=lora_alpha, ) return ret return None # source: https://github.com/vllm-project/vllm/blob/93b38bea5dd03e1b140ca997dfaadef86f8f1855/vllm/lora/utils.py#L9 def replace_submodule( model: nn.Module, module_name: str, new_module: nn.Module ) -> nn.Module: """Replace a submodule in a model with a new module.""" parent = model.get_submodule(".".join(module_name.split(".")[:-1])) target_name = module_name.split(".")[-1] setattr(parent, target_name, new_module) return new_module ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/mlp.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from typing import Optional import torch import torch.nn as nn from diffusers.models.activations import ( GEGLU, GELU, ApproximateGELU, LinearActivation, SwiGLU, ) from sglang.multimodal_gen.runtime.layers.activation import get_act_fn from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig from sglang.srt.utils import add_prefix class MLP(nn.Module): """ MLP for DiT blocks, NO gated linear units """ def __init__( self, input_dim: int, mlp_hidden_dim: int, output_dim: int | None = None, bias: bool = True, act_type: str = "gelu_pytorch_tanh", dtype: torch.dtype | None = None, prefix: str = "", quant_config: QuantizationConfig = None, ): super().__init__() self.fc_in = ColumnParallelLinear( input_dim, mlp_hidden_dim, bias=True, gather_output=False, quant_config=quant_config, prefix=add_prefix("0.proj", prefix), ) self.act = get_act_fn(act_type) if output_dim is None: output_dim = input_dim self.fc_out = RowParallelLinear( mlp_hidden_dim, output_dim, bias=True, input_is_parallel=True, quant_config=quant_config, prefix=add_prefix("2", prefix), ) def forward(self, x: torch.Tensor) -> torch.Tensor: x, _ = self.fc_in(x) x = self.act(x) x, _ = self.fc_out(x) return x class FeedForward(nn.Module): r""" A feed-forward layer. Parameters: dim (`int`): The number of channels in the input. dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`. mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension. activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. bias (`bool`, defaults to True): Whether to use a bias in the linear layer. """ def __init__( self, dim: int, dim_out: Optional[int] = None, mult: int = 4, activation_fn: str = "geglu", inner_dim=None, bias: bool = True, ): super().__init__() if inner_dim is None: inner_dim = int(dim * mult) dim_out = dim_out if dim_out is not None else dim if activation_fn == "gelu": act_fn = GELU(dim, inner_dim, bias=bias) if activation_fn == "gelu-approximate": act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias) elif activation_fn == "geglu": act_fn = GEGLU(dim, inner_dim, bias=bias) elif activation_fn == "geglu-approximate": act_fn = ApproximateGELU(dim, inner_dim, bias=bias) elif activation_fn == "swiglu": act_fn = SwiGLU(dim, inner_dim, bias=bias) elif activation_fn == "linear-silu": act_fn = LinearActivation(dim, inner_dim, bias=bias, activation="silu") self.net = nn.ModuleList([]) # project in self.net.append(act_fn) # dummy dropout layer to match with checkpoints compatible with diffusers self.net.append(nn.Dropout(0.0)) # project out self.net.append(nn.Linear(inner_dim, dim_out, bias=bias)) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: for module in self.net: hidden_states = module(hidden_states) return hidden_states ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from typing import Literal, get_args from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config from sglang.multimodal_gen.runtime.layers.quantization.modelslim import ModelSlimConfig QuantizationMethods = Literal["fp8", "modelslim"] QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) # The customized quantization methods which will be added to this dict. _CUSTOMIZED_METHOD_TO_QUANT_CONFIG = { "modelslim": ModelSlimConfig, "fp8": Fp8Config, } def register_quantization_config(quantization: str): """Register a customized vllm quantization config. When a quantization method is not supported by vllm, you can register a customized quantization config to support it. Args: quantization (str): The quantization method name. """ # noqa: E501 def _wrapper(quant_config_cls): if quantization in QUANTIZATION_METHODS: raise ValueError( f"The quantization method `{quantization}` is already exists." ) if not issubclass(quant_config_cls, QuantizationConfig): raise ValueError( "The quantization config must be a subclass of " "`QuantizationConfig`." ) _CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls QUANTIZATION_METHODS.append(quantization) return quant_config_cls return _wrapper def get_quantization_config(quantization: str) -> type[QuantizationConfig]: if quantization not in QUANTIZATION_METHODS: raise ValueError(f"Invalid quantization method: {quantization}") method_to_config: dict[str, type[QuantizationConfig]] = {} # Update the `method_to_config` with customized quantization methods. method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG) return method_to_config[quantization] __all__ = [ "QuantizationMethods", "QuantizationConfig", "get_quantization_config", "QUANTIZATION_METHODS", ] ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/quantization/base_config.py import inspect from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any import torch from torch import nn if TYPE_CHECKING: from sglang.multimodal_gen.runtime.layers.quantization import QuantizationMethods else: QuantizationMethods = str class QuantizeMethodBase(ABC): """Base class for different quantized methods.""" @abstractmethod def create_weights( self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs ): """Create weights for a layer. The weights will be set as attributes of the layer.""" raise NotImplementedError @abstractmethod def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: """Apply the weights in layer to the input tensor. Expects create_weights to have been called before on the layer.""" raise NotImplementedError # Not required functions def embedding(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: """Gather embeddings in the layer based on indices in the input tensor. Expects create_weights to have been called before on the layer.""" raise NotImplementedError def process_weights_after_loading(self, layer: nn.Module) -> None: """Process the weight after loading. This can be used for example, to transpose weights for computation. """ return def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> bool: """ Not all quant methods have embedding implemented, so we need to check that it exists for our given method. We check this by making sure the function has been changed from the base implementation. """ base_embedding = inspect.getattr_static(QuantizeMethodBase, "embedding", None) class_embedding = inspect.getattr_static(method_class, "embedding", None) return class_embedding is not None and class_embedding is not base_embedding class QuantizationConfig(ABC): """Base class for quantization configs.""" # for quantization frameworks with a separate quantized model provided, e.g. Nunchaku quantized_model_path: str | None = None def __init__(self): super().__init__() # mapping is updated by models as they initialize self.packed_modules_mapping: dict[str, list[str]] = dict() @abstractmethod def get_name(self) -> QuantizationMethods: """Name of the quantization method.""" raise NotImplementedError @abstractmethod def get_supported_act_dtypes(self) -> list[torch.dtype]: """List of supported activation dtypes.""" raise NotImplementedError @classmethod @abstractmethod def get_min_capability(cls) -> int: """Minimum GPU capability to support the quantization method. E.g., 70 for Volta, 75 for Turing, 80 for Ampere. This requirement is due to the custom CUDA kernels used by the quantization method. """ raise NotImplementedError @staticmethod @abstractmethod def get_config_filenames() -> list[str]: """List of filenames to search for in the model directory.""" raise NotImplementedError @classmethod @abstractmethod def from_config(cls, config: dict[str, Any]) -> "QuantizationConfig": """Create a config class from the model's quantization config.""" raise NotImplementedError @classmethod def override_quantization_method( cls, hf_quant_cfg, user_quant ) -> QuantizationMethods | None: """ Detects if this quantization method can support a given checkpoint format by overriding the user specified quantization method -- this method should only be overwritten by subclasses in exceptional circumstances """ return None @staticmethod def get_from_keys(config: dict[str, Any], keys: list[str]) -> Any: """Get a value from the model's quantization config.""" for key in keys: if key in config: return config[key] raise ValueError( f"Cannot find any of {keys} in the model's " "quantization config." ) @staticmethod def get_from_keys_or(config: dict[str, Any], keys: list[str], default: Any) -> Any: """Get a optional value from the model's quantization config.""" try: return QuantizationConfig.get_from_keys(config, keys) except ValueError: return default @abstractmethod def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> QuantizeMethodBase | None: """Get the quantize method to use for the quantized layer. Args: layer: The layer for the quant method. prefix: The full name of the layer in the state dict Returns: The quantize method. None if the given layer doesn't support quant method. """ raise NotImplementedError def get_cache_scale(self, name: str) -> str | None: return None ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py ================================================ # SPDX-License-Identifier: Apache-2.0 import json import os from dataclasses import dataclass from functools import lru_cache from typing import Any, Optional import torch from safetensors.torch import load_file as safetensors_load_file from torch import nn from sglang.multimodal_gen.runtime.layers.linear import LinearBase from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from .base_config import QuantizationConfig, QuantizeMethodBase logger = init_logger(__name__) @lru_cache(maxsize=1) def is_nunchaku_available() -> bool: try: import nunchaku # noqa logger.debug("Nunchaku package detected") return True except Exception: return False @dataclass class NunchakuConfig(QuantizationConfig): """ Configuration for Nunchaku (SVDQuant) W4A4-style quantization. Attributes: precision: Quantization precision type. Options: - "int4": Standard INT4 quantization - "nvfp4": FP4 quantization rank: SVD low-rank dimension for absorbing outliers group_size: Quantization group size (automatically set based on precision) act_unsigned: Use unsigned activation quantization transformer_weights_path: Path to pre-quantized transformer weights (.safetensors) model_cls: DiT model class that provides quantization rules via get_nunchaku_quant_rules() """ precision: str = "int4" rank: int = 32 group_size: Optional[int] = None act_unsigned: bool = False transformer_weights_path: Optional[str] = None model_cls: Optional[type] = None @classmethod def get_name(cls) -> str: return "svdquant" @classmethod def get_supported_act_dtypes(cls) -> list[torch.dtype]: return [torch.bfloat16, torch.float16] @classmethod def get_min_capability(cls) -> int: return 70 @staticmethod def get_config_filenames() -> list[str]: return ["quantization_config.json", "quant_config.json"] @classmethod def from_config(cls, config: dict[str, Any]) -> "NunchakuConfig": return cls( precision=config.get("precision", "int4"), rank=int(config.get("rank", 32)), group_size=config.get("group_size"), act_unsigned=bool(config.get("act_unsigned", False)), transformer_weights_path=config.get("transformer_weights_path"), ) def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> Optional[QuantizeMethodBase]: if not isinstance(layer, LinearBase): return None # get quantization rules from model class quant_rules = self._get_quant_rules() # priority: skip > awq_w4a16 > svdq_w4a4 > default skip_patterns = quant_rules.get("skip", []) for pattern in skip_patterns: if pattern in prefix.lower(): return None awq_patterns = quant_rules.get("awq_w4a16", []) for pattern in awq_patterns: if pattern in prefix: from ..nunchaku_linear import NunchakuAWQLinearMethod return NunchakuAWQLinearMethod(group_size=64) svdq_patterns = quant_rules.get("svdq_w4a4", []) for pattern in svdq_patterns: if pattern in prefix: from ..nunchaku_linear import NunchakuSVDQLinearMethod return NunchakuSVDQLinearMethod( precision=self.precision, rank=self.rank, act_unsigned=self.act_unsigned, ) # default: apply svdq_w4a4 to all remaining linear layers from ..nunchaku_linear import NunchakuSVDQLinearMethod return NunchakuSVDQLinearMethod( precision=self.precision, rank=self.rank, act_unsigned=self.act_unsigned, ) def _get_quant_rules(self) -> dict[str, list[str]]: if self.model_cls is not None and hasattr( self.model_cls, "get_nunchaku_quant_rules" ): return self.model_cls.get_nunchaku_quant_rules() return {} def __post_init__(self): if self.group_size is None: if self.precision == "nvfp4": self.group_size = 16 elif self.precision == "int4": self.group_size = 64 else: raise ValueError( f"Invalid precision: {self.precision}. Must be 'int4' or 'nvfp4'" ) if self.precision not in ["int4", "nvfp4"]: raise ValueError( f"Invalid precision: {self.precision}. Must be 'int4' or 'nvfp4'" ) if self.rank <= 0: raise ValueError(f"Rank must be positive, got {self.rank}") @classmethod def from_dict(cls, config_dict: dict) -> "NunchakuConfig": """Create configuration from dictionary.""" return cls(**config_dict) def to_dict(self) -> dict: """Convert configuration to dictionary.""" return { "precision": self.precision, "rank": self.rank, "group_size": self.group_size, "act_unsigned": self.act_unsigned, "transformer_weights_path": self.transformer_weights_path, } @classmethod def from_pretrained(cls, model_path: str) -> Optional["NunchakuConfig"]: for filename in cls.get_config_filenames(): config_path = os.path.join(model_path, filename) if os.path.exists(config_path): with open(config_path, "r") as f: config_dict = json.load(f) if config_dict.get("quant_method") == cls.get_name(): return cls.from_config(config_dict) return None def _patch_native_svdq_linear( module: nn.Module, tensor: Any, svdq_linear_cls: type ) -> bool: if ( isinstance(module, svdq_linear_cls) and getattr(module, "wtscale", None) is not None ): module.wtscale = tensor return True return False def _patch_sglang_svdq_linear( module: nn.Module, tensor: Any, svdq_method_cls: type ) -> bool: quant_method = getattr(module, "quant_method", None) if not isinstance(quant_method, svdq_method_cls): return False existing = getattr(module, "wtscale", None) if isinstance(existing, nn.Parameter): with torch.no_grad(): existing.data.copy_(tensor.to(existing.data.dtype)) else: module.wtscale = tensor # Keep alpha in sync (kernel reads `layer._nunchaku_alpha`) try: module._nunchaku_alpha = float(tensor.detach().cpu().item()) except Exception: module._nunchaku_alpha = None return True def _patch_sglang_svdq_wcscales( module: nn.Module, tensor: Any, svdq_method_cls: type ) -> bool: quant_method = getattr(module, "quant_method", None) if not isinstance(quant_method, svdq_method_cls): return False existing = getattr(module, "wcscales", None) if isinstance(existing, nn.Parameter): with torch.no_grad(): existing.data.copy_(tensor.to(existing.data.dtype)) else: module.wcscales = tensor return True def _patch_nunchaku_scales( model: nn.Module, safetensors_list: list[str], ) -> None: """Patch transformer module with Nunchaku scale tensors from safetensors weights. For NVFP4 checkpoints, correctness depends on `wtscale` and attention `wcscales`. The FSDP loader may skip some of these metadata tensors. """ if not safetensors_list: return if len(safetensors_list) != 1: logger.warning( "Nunchaku scale patch expects a single safetensors file, " "but got %d files. Skipping.", len(safetensors_list), ) return from nunchaku.models.linear import SVDQW4A4Linear # type: ignore[import] state_dict = safetensors_load_file(safetensors_list[0]) if state_dict is None: return num_wtscale = 0 num_wcscales = 0 from ..nunchaku_linear import NunchakuSVDQLinearMethod for name, module in model.named_modules(): wt = state_dict.get(f"{name}.wtscale") if wt is not None: if _patch_native_svdq_linear(module, wt, SVDQW4A4Linear): num_wtscale += 1 elif _patch_sglang_svdq_linear(module, wt, NunchakuSVDQLinearMethod): num_wtscale += 1 wc = state_dict.get(f"{name}.wcscales") if wc is not None: # Some modules may have wcscales as a direct attribute/Parameter. existing = getattr(module, "wcscales", None) if isinstance(existing, nn.Parameter): with torch.no_grad(): existing.data.copy_(wc.to(existing.data.dtype)) num_wcscales += 1 elif existing is not None: setattr(module, "wcscales", wc) num_wcscales += 1 elif _patch_sglang_svdq_wcscales(module, wc, NunchakuSVDQLinearMethod): num_wcscales += 1 if num_wtscale > 0: logger.info("Patched wtscale for %d layers", num_wtscale) if num_wcscales > 0: logger.info("Patched wcscales for %d layers", num_wcscales) ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py ================================================ from __future__ import annotations import logging from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union import torch from torch.nn import Module from torch.nn.parameter import Parameter from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_tensor_model_parallel_world_size, ) from sglang.multimodal_gen.runtime.layers.linear import ( LinearMethodBase, UnquantizedLinearMethod, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, QuantizeMethodBase, ) from sglang.multimodal_gen.runtime.models.parameter import ( BlockQuantScaleParameter, ModelWeightParameter, PerTensorScaleParameter, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.common import ( cpu_has_amx_support, get_bool_env_var, use_intel_amx_backend, ) from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading from sglang.srt.layers.quantization.fp8_kernel import ( is_fp8_fnuz, per_token_group_quant_fp8, ) from sglang.srt.layers.quantization.fp8_utils import ( apply_fp8_linear, can_auto_enable_marlin_fp8, cutlass_fp8_supported, dispatch_w8a8_block_fp8_linear, input_to_float8, normalize_e4m3fn_to_e4m3fnuz, requant_weight_ue8m0_inplace, ) from sglang.srt.layers.quantization.marlin_utils_fp8 import ( apply_fp8_marlin_linear, prepare_fp8_layer_for_marlin, ) from sglang.srt.layers.quantization.utils import ( convert_to_channelwise, is_layer_skipped, requantize_with_max_scale, ) if TYPE_CHECKING: from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config _is_hip = current_platform.is_hip() _is_cuda = current_platform.is_cuda() _is_npu = current_platform.is_npu() _is_cpu_amx_available = cpu_has_amx_support() _is_cpu = current_platform.is_cpu() _is_fp8_fnuz = is_fp8_fnuz() _use_hip_int4 = get_bool_env_var("SGLANG_INT4_WEIGHT") and _is_hip _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _use_aiter or _use_hip_int4: pass ACTIVATION_SCHEMES = ["static", "dynamic"] logger = logging.getLogger(__name__) class Fp8Config(QuantizationConfig): """Config class for FP8.""" def __init__( self, is_checkpoint_fp8_serialized: bool = False, activation_scheme: str = "dynamic", ignored_layers: Optional[List[str]] = None, weight_block_size: List[int] = None, ) -> None: self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized if is_checkpoint_fp8_serialized: logger.info("Detected fp8 checkpoint.") if activation_scheme not in ACTIVATION_SCHEMES: raise ValueError(f"Unsupported activation scheme {activation_scheme}") self.activation_scheme = activation_scheme self.ignored_layers = ignored_layers or [] if weight_block_size is not None: if not is_checkpoint_fp8_serialized: raise ValueError( f"The block-wise quantization only supports fp8-serialized checkpoint for now." ) if len(weight_block_size) != 2: raise ValueError( f"The quantization block size of weight must have 2 dimensions, but got {len(weight_block_size)} dimensions." ) if activation_scheme != "dynamic": raise ValueError( f"The block-wise quantization only supports dynamic activation scheme for now, but got {activation_scheme} activation scheme." ) self.weight_block_size = weight_block_size @classmethod def get_name(cls) -> str: return "fp8" @classmethod def get_supported_act_dtypes(cls) -> List[torch.dtype]: return [torch.bfloat16, torch.half] @classmethod def get_min_capability(cls) -> int: return 80 @classmethod def get_config_filenames(cls) -> List[str]: return [] @classmethod def from_config(cls, config: Dict[str, Any]) -> Fp8Config: quant_method = cls.get_from_keys(config, ["quant_method"]) is_checkpoint_fp8_serialized = "fp8" in quant_method activation_scheme = cls.get_from_keys(config, ["activation_scheme"]) ignored_layers = cls.get_from_keys_or( config, ["ignored_layers", "modules_to_not_convert"], None ) if ignored_layers: # hacking ministral ignored_layers = [layer.replace("model.", "") for layer in ignored_layers] weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None) return cls( is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, activation_scheme=activation_scheme, ignored_layers=ignored_layers, weight_block_size=weight_block_size, ) def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> Optional[QuantizeMethodBase]: from sglang.multimodal_gen.runtime.layers.linear import LinearBase if isinstance(layer, LinearBase): if is_layer_skipped(prefix, self.ignored_layers): return UnquantizedLinearMethod() return Fp8LinearMethod(self) return None def get_scaled_act_names(self) -> List[str]: return [] class Fp8LinearMethod(LinearMethodBase): """Linear method for FP8. Supports loading FP8 checkpoints with static weight scale and dynamic/static activation scale. Also supports loading quantized FP16/BF16 model checkpoints with dynamic activation scaling. The weight scaling factor will be initialized after the model weights are loaded. Limitations: 1. Only support per-tensor quantization due to torch._scaled_mm support. 2. Only support float8_e4m3fn data type due to the limitation of torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856) Args: quant_config: The quantization config. """ def __init__(self, quant_config: Union[Fp8Config, W4AFp8Config]): self.quant_config = quant_config self.cutlass_fp8_supported = cutlass_fp8_supported() # For GPUs that lack FP8 hardware support, we can leverage the Marlin # kernel for fast weight-only FP8 quantization self.use_marlin = False if _is_cuda: force_marlin = get_bool_env_var("SGLANG_FORCE_FP8_MARLIN") auto_enable = can_auto_enable_marlin_fp8() self.use_marlin = force_marlin or auto_enable self.block_quant = self.quant_config.weight_block_size is not None self.w8a8_block_fp8_linear = dispatch_w8a8_block_fp8_linear() def create_weights( self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: List[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs, ): output_size_per_partition = sum(output_partition_sizes) weight_loader = extra_weight_attrs.get("weight_loader") tp_size = get_tensor_model_parallel_world_size() if self.block_quant: block_n, block_k = ( self.quant_config.weight_block_size[0], self.quant_config.weight_block_size[1], ) # Required by row parallel if tp_size > 1 and input_size // input_size_per_partition == tp_size: if input_size_per_partition % block_k != 0: raise ValueError( f"Weight input_size_per_partition = " f"{input_size_per_partition} is not divisible by " f"weight quantization block_k = {block_k}." ) # Required by column parallel or enabling merged weights if ( tp_size > 1 and output_size // output_size_per_partition == tp_size ) or len(output_partition_sizes) > 1: for output_partition_size in output_partition_sizes: if output_partition_size % block_n != 0: raise ValueError( f"Weight output_partition_size = " f"{output_partition_size} is not divisible by " f"weight quantization block_n = {block_n}." ) layer.logical_widths = output_partition_sizes layer.input_size_per_partition = input_size_per_partition layer.output_size_per_partition = output_size_per_partition layer.orig_dtype = params_dtype # WEIGHT weight_dtype = ( torch.float8_e4m3fn if self.quant_config.is_checkpoint_fp8_serialized else params_dtype ) weight = ModelWeightParameter( data=torch.empty( output_size_per_partition, input_size_per_partition, dtype=weight_dtype ), input_dim=1, output_dim=0, weight_loader=weight_loader, ) layer.register_parameter("weight", weight) # If checkpoint is serialized fp8, load them. # Otherwise, wait until process_weights_after_loading. if self.quant_config.is_checkpoint_fp8_serialized: # WEIGHT SCALE if self.block_quant: if hasattr(self.quant_config, "activation_scheme"): assert self.quant_config.activation_scheme == "dynamic" elif hasattr(self.quant_config, "linear_activation_scheme"): assert self.quant_config.linear_activation_scheme == "dynamic" scale = BlockQuantScaleParameter( data=torch.empty( (output_size_per_partition + block_n - 1) // block_n, (input_size_per_partition + block_k - 1) // block_k, dtype=torch.float32, ), input_dim=1, output_dim=0, weight_loader=weight_loader, ) scale.format_ue8m0 = False scale[:] = torch.finfo(torch.float32).min layer.register_parameter("weight_scale_inv", scale) else: scale = PerTensorScaleParameter( data=torch.empty(len(output_partition_sizes), dtype=torch.float32), weight_loader=weight_loader, ) scale[:] = torch.finfo(torch.float32).min layer.register_parameter("weight_scale", scale) # INPUT ACTIVATION SCALE if ( hasattr(self.quant_config, "activation_scheme") and self.quant_config.activation_scheme == "static" ) or ( hasattr(self.quant_config, "linear_activation_scheme") and self.quant_config.linear_activation_scheme == "static" ): scale = PerTensorScaleParameter( data=torch.empty(len(output_partition_sizes), dtype=torch.float32), weight_loader=weight_loader, ) scale[:] = torch.finfo(torch.float32).min layer.register_parameter("input_scale", scale) else: layer.register_parameter("input_scale", None) def process_weights_after_loading(self, layer: Module) -> None: if self.block_quant: # If ROCm, normalize the weights and scales to e4m3fnuz if _is_fp8_fnuz: # activation_scheme: dynamic weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( weight=layer.weight, weight_scale=layer.weight_scale_inv, input_scale=None, ) layer.input_scale = None elif _is_cpu: assert ( _is_cpu_amx_available ), "Fp8LinearMethod on CPU requires that CPU has AMX support" _amx_process_weight_after_loading(layer, ["weight"]) layer.weight_scale_inv = torch.nn.Parameter( layer.weight_scale_inv.data, requires_grad=False ) return else: # For fp8 linear weights run with deepgemm, the weights and scales need be requantized to ue8m0 from sglang.srt.layers.quantization.fp8_utils import ( deepgemm_w8a8_block_fp8_linear_with_fallback, ) from sglang.srt.model_loader.utils import ( should_deepgemm_weight_requant_ue8m0, ) if ( should_deepgemm_weight_requant_ue8m0( weight_block_size=getattr( self.quant_config, "weight_block_size", None ), ) and ( self.w8a8_block_fp8_linear is deepgemm_w8a8_block_fp8_linear_with_fallback ) and (not layer.weight_scale_inv.format_ue8m0) ): requant_weight_ue8m0_inplace( layer.weight, layer.weight_scale_inv, self.quant_config.weight_block_size, ) layer.weight_scale_inv.format_ue8m0 = True weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data layer.weight.data = weight.data layer.weight_scale_inv.data = weight_scale.data else: layer.weight = Parameter(layer.weight.data, requires_grad=False) # If checkpoint not serialized fp8, quantize the weights. if not self.quant_config.is_checkpoint_fp8_serialized: if self.cutlass_fp8_supported or self.use_marlin: # apply per-channel quantization default as # cutlass sgl-kernel and marlin only support per-channel scale qweight, weight_scale = per_token_group_quant_fp8( layer.weight, layer.weight.shape[-1] ) weight_scale = weight_scale.t().contiguous() else: # per-tensor quantization qweight, weight_scale = input_to_float8(layer.weight) # Update the layer with the new values. layer.weight = Parameter(qweight.t(), requires_grad=False) layer.weight_scale = Parameter(weight_scale, requires_grad=False) layer.input_scale = None # If checkpoint is fp8, handle that there are N scales for N # shards in a fused module else: layer.weight_scale = Parameter( layer.weight_scale.data, requires_grad=False ) if ( hasattr(self.quant_config, "activation_scheme") and self.quant_config.activation_scheme == "static" ) or ( hasattr(self.quant_config, "linear_activation_scheme") and self.quant_config.linear_activation_scheme == "static" ): layer.input_scale = Parameter( layer.input_scale.data, requires_grad=False ) # cutlass sgl-kernel and marlin only support per-channel scale if self.cutlass_fp8_supported or self.use_marlin: weight = layer.weight weight_scale = convert_to_channelwise( layer.weight_scale, layer.logical_widths ) else: # Dequant -> Quant with max scale so we can run per tensor. weight = layer.weight weight_scale = layer.weight_scale # If ROCm, normalize the weights and scales to e4m3fnuz if _is_fp8_fnuz: weight, weight_scale, input_scale = ( normalize_e4m3fn_to_e4m3fnuz( weight=weight, weight_scale=weight_scale, input_scale=layer.input_scale, ) ) if input_scale is not None: layer.input_scale = Parameter( input_scale, requires_grad=False ) weight_scale, weight = requantize_with_max_scale( weight=weight, weight_scale=weight_scale, logical_widths=layer.logical_widths, ) # Update layer with new values. layer.weight = Parameter(weight.t(), requires_grad=False) layer.weight_scale = Parameter(weight_scale, requires_grad=False) if ( hasattr(self.quant_config, "activation_scheme") and self.quant_config.activation_scheme == "static" ) or ( hasattr(self.quant_config, "linear_activation_scheme") and self.quant_config.linear_activation_scheme == "static" ): layer.input_scale = Parameter( layer.input_scale.max(), requires_grad=False ) if self.use_marlin: if self.block_quant: layer.weight_block_size = self.quant_config.weight_block_size prepare_fp8_layer_for_marlin(layer, not self.block_quant) # Activations not quantized for marlin. del layer.input_scale def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: if self.use_marlin: return apply_fp8_marlin_linear( input=x, weight=layer.weight, weight_scale=layer.weight_scale, workspace=layer.workspace, size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, bias=bias, ) if self.block_quant: if use_intel_amx_backend(layer): return torch.ops.sgl_kernel.fp8_scaled_mm_cpu( x, layer.weight, layer.weight_scale_inv, self.quant_config.weight_block_size, bias, x.dtype, True, # is_vnni ) if isinstance(x, tuple): return self.w8a8_block_fp8_linear( input=x[0], weight=layer.weight, block_size=self.quant_config.weight_block_size, weight_scale=layer.weight_scale_inv, input_scale=x[1], bias=bias, ) return self.w8a8_block_fp8_linear( input=x, weight=layer.weight, block_size=self.quant_config.weight_block_size, weight_scale=layer.weight_scale_inv, input_scale=None, bias=bias, ) return apply_fp8_linear( input=x, weight=layer.weight, weight_scale=layer.weight_scale, input_scale=layer.input_scale, bias=bias, cutlass_fp8_supported=self.cutlass_fp8_supported, use_per_token_if_dynamic=False, ) ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/quantization/modelslim.py ================================================ from __future__ import annotations import logging from types import MappingProxyType from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, cast import torch from sglang.multimodal_gen.runtime.layers.linear import ( LinearMethodBase, UnquantizedLinearMethod, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, QuantizeMethodBase, ) from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer from sglang.srt.layers.quantization.modelslim.schemes import ( ModelSlimW4A4Int4, ModelSlimW8A8Int8, ) if TYPE_CHECKING: from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig from sglang.srt.layers.quantization.modelslim.schemes import ( ModelSlimLinearScheme, ) logger = logging.getLogger(__name__) class ModelSlimConfig(QuantizationConfig): """ Config class for ModelSlim Quantization of Diffusion models https://gitcode.com/Ascend/msmodelslim, a NPU-specific quantization type. The quantization method (W8A8, W4A4, etc.) will be automatically parsed from the `quant_model_description.json` config. ModelSlim for Diffusion models includes support for various quantization schemes, such as: - W4A4 dynamic linear - W8A8 static linear - W8A8 dynamic linear """ def __init__(self, quant_config: Dict[str, Any] = {}): super().__init__() self.quant_description = quant_config ignore = cast(List[str], quant_config.get("ignore", [])) self.ignore = ignore packed_modules_mapping = quant_config.get("packed_modules_mapping", {}) self.packed_modules_mapping = ( packed_modules_mapping if packed_modules_mapping is not None else {} ) def get_linear_method(self) -> ModelSlimLinearMethod: return ModelSlimLinearMethod(self) @classmethod def get_supported_act_dtypes(cls) -> List[torch.dtype]: return [torch.int8, torch.float16, torch.bfloat16] @classmethod def get_min_capability(cls) -> int: return 0 @classmethod def get_name(cls) -> str: return "modelslim" @classmethod def get_config_filenames(cls) -> List[str]: filenames = ["quant_model_description.json"] return filenames @classmethod def from_config(cls, config: Dict[str, Any]) -> ModelSlimConfig: return cls(config) def get_quant_method( self, layer: torch.nn.Module, prefix: str, ) -> Optional[QuantizeMethodBase]: from sglang.multimodal_gen.runtime.layers.linear import LinearBase if isinstance(layer, LinearBase): if should_ignore_layer( prefix, ignore=self.ignore, fused_mapping=self.packed_modules_mapping, ): return UnquantizedLinearMethod() key = "model" packed_modules_mapping_subset = self.packed_modules_mapping.get(key, {}) prefix_in_quant_config = prefix proj_name = prefix.split(".")[-1] if proj_name in packed_modules_mapping_subset: prefix_in_quant_config = prefix.replace( proj_name, packed_modules_mapping_subset[proj_name][0] ) if self.is_layer_skipped(prefix, packed_modules_mapping_subset): return UnquantizedLinearMethod() scheme = self.get_scheme(layer=layer, layer_name=prefix_in_quant_config) layer.scheme = scheme return ModelSlimLinearMethod(self) else: return None def _get_scheme_from_parts( self, layer_name: str, ) -> ModelSlimLinearScheme: quant_type = self.quant_description.get(layer_name + ".weight", "") if quant_type == "W8A8_DYNAMIC" or quant_type == "W8A8": return ModelSlimW8A8Int8( quant_config=self.quant_description, prefix=layer_name ) elif quant_type == "W4A4_DYNAMIC": return ModelSlimW4A4Int4( quant_config=self.quant_description, prefix=layer_name ) raise NotImplementedError("No modelslim compatible scheme was found.") def get_scheme( self, layer: torch.nn.Module, layer_name: Optional[str] = None ) -> Optional[ModelSlimLinearScheme]: """ get_scheme method adjusted for modelslim, taken from python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py """ scheme = self._get_scheme_from_parts( layer_name=layer_name, ) # Ascend doesn't support device capability logger.debug("Using scheme: %s for %s", scheme.__class__.__name__, layer_name) return scheme def is_layer_skipped( self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({}) ): # adapted from vllm.model_executor.layers.quantization.utils.quant_utils.is_layer_skipped proj_name = prefix.split(".")[-1] if proj_name in fused_mapping: shard_prefixes = [ prefix.replace(proj_name, shard_proj_name) for shard_proj_name in fused_mapping[proj_name] ] is_skipped = None for shard_prefix in shard_prefixes: is_shard_skipped = ( self.quant_description.get(shard_prefix + ".weight", "") == "FLOAT" ) if is_skipped is None: is_skipped = is_shard_skipped elif is_shard_skipped != is_skipped: raise ValueError( f"Detected some but not all shards of {prefix} " "are quantized. All shards of fused layers " "to have the same precision." ) else: is_skipped = self.quant_description.get(prefix + ".weight", "") == "FLOAT" assert is_skipped is not None return is_skipped def get_scaled_act_names(self) -> List[str]: return [] class ModelSlimLinearMethod(LinearMethodBase): def __init__(self, quantization_config: ModelSlimConfig): self.quantization_config = quantization_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.scheme.process_weights_after_loading(layer) def create_weights( self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: List[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs, ): """ Use the ModelSlimLinearScheme associated with each layer to create the necessary parameters for the layer. See LinearMethodBase for param details """ weight_loader = extra_weight_attrs.get("weight_loader") layer.scheme.create_weights( layer=layer, input_size=input_size, input_size_per_partition=input_size_per_partition, output_partition_sizes=output_partition_sizes, output_size=output_size, params_dtype=params_dtype, weight_loader=weight_loader, ) def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] = None, ): """ Use the output of create_weights and the CompressedTensorsScheme associated with the layer to apply the forward pass with the layer input. See LinearMethodBase for param details """ scheme = layer.scheme if scheme is None: raise ValueError("A scheme must be defined for each layer") return scheme.apply_weights(layer, x, bias=bias) ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/quantization/nunchaku_linear.py ================================================ # SPDX-License-Identifier: Apache-2.0 from typing import List, Optional import torch import torch.nn as nn from torch.nn.parameter import Parameter from sglang.multimodal_gen.runtime.layers.linear import LinearMethodBase from sglang.multimodal_gen.runtime.models.utils import set_weight_attrs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) try: from nunchaku.ops.gemm import svdq_gemm_w4a4_cuda from nunchaku.ops.gemv import awq_gemv_w4a16_cuda from nunchaku.ops.quantize import svdq_quantize_w4a4_act_fuse_lora_cuda except ImportError: svdq_gemm_w4a4_cuda = None awq_gemv_w4a16_cuda = None svdq_quantize_w4a4_act_fuse_lora_cuda = None class NunchakuSVDQLinearMethod(LinearMethodBase): def __init__( self, precision: str = "int4", rank: int = 32, act_unsigned: bool = False, ): self.precision = precision self.rank = rank self.act_unsigned = act_unsigned if precision == "nvfp4": self.group_size = 16 else: self.group_size = 64 def create_weights( self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: List[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs, ) -> None: output_size_per_partition = sum(output_partition_sizes) qweight = Parameter( torch.empty( output_size_per_partition, input_size_per_partition // 2, dtype=torch.int8, ), requires_grad=False, ) set_weight_attrs(qweight, {"input_dim": 1, "output_dim": 0}) num_groups = input_size_per_partition // self.group_size if self.precision == "nvfp4": scale_dtype = torch.float8_e4m3fn else: scale_dtype = params_dtype wscales = Parameter( torch.empty(num_groups, output_size_per_partition, dtype=scale_dtype), requires_grad=False, ) smooth_factor = Parameter( torch.empty(input_size_per_partition, dtype=params_dtype), requires_grad=False, ) smooth_factor_orig = Parameter( torch.empty(input_size_per_partition, dtype=params_dtype), requires_grad=False, ) proj_down = Parameter( torch.empty(input_size_per_partition, self.rank, dtype=params_dtype), requires_grad=False, ) proj_up = Parameter( torch.empty(output_size_per_partition, self.rank, dtype=params_dtype), requires_grad=False, ) if self.precision == "nvfp4": wcscales = Parameter( torch.empty( output_size_per_partition, dtype=params_dtype, ), requires_grad=False, ) wtscale = Parameter( torch.empty(1, dtype=params_dtype), requires_grad=False, ) else: wcscales = None wtscale = None layer.register_parameter("qweight", qweight) layer.register_parameter("wscales", wscales) layer.register_parameter("smooth_factor", smooth_factor) layer.register_parameter("smooth_factor_orig", smooth_factor_orig) layer.register_parameter("proj_down", proj_down) layer.register_parameter("proj_up", proj_up) if wcscales is not None: layer.register_parameter("wcscales", wcscales) if wtscale is not None: layer.register_parameter("wtscale", wtscale) layer.input_size_per_partition = input_size_per_partition layer.output_size_per_partition = output_size_per_partition layer.precision = self.precision layer.rank = self.rank layer.group_size = self.group_size layer.act_unsigned = self.act_unsigned weight_loader = extra_weight_attrs.get("weight_loader") if weight_loader is not None: set_weight_attrs(qweight, {"weight_loader": weight_loader}) set_weight_attrs(wscales, {"weight_loader": weight_loader}) set_weight_attrs(smooth_factor, {"weight_loader": weight_loader}) set_weight_attrs(smooth_factor_orig, {"weight_loader": weight_loader}) set_weight_attrs(proj_down, {"weight_loader": weight_loader}) set_weight_attrs(proj_up, {"weight_loader": weight_loader}) if wcscales is not None: set_weight_attrs(wcscales, {"weight_loader": weight_loader}) if wtscale is not None: set_weight_attrs(wtscale, {"weight_loader": weight_loader}) def process_weights_after_loading(self, layer: nn.Module) -> None: layer.qweight = Parameter(layer.qweight.data, requires_grad=False) layer.wscales = Parameter(layer.wscales.data, requires_grad=False) layer.smooth_factor = Parameter(layer.smooth_factor.data, requires_grad=False) layer.smooth_factor_orig = Parameter( layer.smooth_factor_orig.data, requires_grad=False ) layer.proj_down = Parameter(layer.proj_down.data, requires_grad=False) layer.proj_up = Parameter(layer.proj_up.data, requires_grad=False) if hasattr(layer, "wcscales") and layer.wcscales is not None: layer.wcscales = Parameter(layer.wcscales.data, requires_grad=False) if hasattr(layer, "wtscale") and layer.wtscale is not None: layer.wtscale = Parameter(layer.wtscale.data, requires_grad=False) alpha: float | None = None wtscale = getattr(layer, "wtscale", None) if wtscale is not None: if isinstance(wtscale, Parameter): wtscale = wtscale.data if isinstance(wtscale, torch.Tensor): alpha = float(wtscale.detach().cpu().item()) else: alpha = float(wtscale) layer._nunchaku_alpha = alpha def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: orig_shape = x.shape x_2d = x.reshape(-1, orig_shape[-1]) quantized_x, ascales, lora_act_out = svdq_quantize_w4a4_act_fuse_lora_cuda( x_2d, lora_down=layer.proj_down, smooth=layer.smooth_factor, fp4=layer.precision == "nvfp4", pad_size=256, ) out_2d = torch.empty( x_2d.shape[0], layer.output_size_per_partition, dtype=x_2d.dtype, device=x_2d.device, ) alpha: float | None = getattr(layer, "_nunchaku_alpha", None) wcscales = getattr(layer, "wcscales", None) svdq_gemm_w4a4_cuda( act=quantized_x, wgt=layer.qweight, out=out_2d, ascales=ascales, wscales=layer.wscales, lora_act_in=lora_act_out, lora_up=layer.proj_up, bias=bias, fp4=layer.precision == "nvfp4", alpha=alpha, wcscales=wcscales, act_unsigned=getattr(layer, "act_unsigned", False), ) out = out_2d.reshape(*orig_shape[:-1], layer.output_size_per_partition) return out class NunchakuAWQLinearMethod(LinearMethodBase): def __init__(self, group_size: int = 64): self.group_size = group_size self.pack_factor = 8 def create_weights( self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: List[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs, ) -> None: output_size_per_partition = sum(output_partition_sizes) qweight = Parameter( torch.empty( output_size_per_partition // 4, input_size_per_partition // 2, dtype=torch.int32, ), requires_grad=False, ) set_weight_attrs(qweight, {"input_dim": 1, "output_dim": 0}) num_groups = input_size_per_partition // self.group_size wscales = Parameter( torch.empty(num_groups, output_size_per_partition, dtype=params_dtype), requires_grad=False, ) wzeros = Parameter( torch.empty(num_groups, output_size_per_partition, dtype=params_dtype), requires_grad=False, ) layer.register_parameter("qweight", qweight) layer.register_parameter("wscales", wscales) layer.register_parameter("wzeros", wzeros) layer.input_size_per_partition = input_size_per_partition layer.output_size_per_partition = output_size_per_partition layer.group_size = self.group_size layer.pack_factor = self.pack_factor weight_loader = extra_weight_attrs.get("weight_loader") if weight_loader is not None: set_weight_attrs(qweight, {"weight_loader": weight_loader}) set_weight_attrs(wscales, {"weight_loader": weight_loader}) set_weight_attrs(wzeros, {"weight_loader": weight_loader}) def process_weights_after_loading(self, layer: nn.Module) -> None: layer.qweight = Parameter(layer.qweight.data, requires_grad=False) layer.wscales = Parameter(layer.wscales.data, requires_grad=False) layer.wzeros = Parameter(layer.wzeros.data, requires_grad=False) def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: orig_shape = x.shape x_2d = x.reshape(-1, orig_shape[-1]) in_features = layer.input_size_per_partition out_features = layer.output_size_per_partition out_2d = awq_gemv_w4a16_cuda( in_feats=x_2d, kernel=layer.qweight, scaling_factors=layer.wscales, zeros=layer.wzeros, m=x_2d.shape[0], n=out_features, k=in_features, group_size=layer.group_size, ) if bias is not None: view_shape = [1] * (out_2d.ndim - 1) + [-1] out_2d.add_(bias.view(view_shape)) out = out_2d.reshape(*orig_shape[:-1], out_features) return out ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/rotary_embedding/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/rotary_embedding.py # Adapted from # https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/llama/modeling_llama.py # Copyright 2023 The vLLM team. # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. # # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX # and OPT implementations in this library. It has been modified from its # original forms to accommodate minor architectural differences compared # to GPT-NeoX and OPT used by the Meta AI team that trained the model. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Rotary Positional Embeddings — unified public API (drop-in replacement).""" from .base import RotaryEmbedding from .factory import get_rope, get_rotary_pos_embed from .mrope import NDRotaryEmbedding from .utils import ( _apply_rotary_emb, apply_flashinfer_rope_qk_inplace, ) __all__ = [ # _utils "_apply_rotary_emb", "apply_flashinfer_rope_qk_inplace", # _base "RotaryEmbedding", # _mrope "NDRotaryEmbedding", # _factory "get_rope", "get_rotary_pos_embed", ] ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/rotary_embedding/base.py ================================================ """RotaryEmbedding base class and LinearScalingRotaryEmbedding variant.""" import torch from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp from .utils import _apply_rotary_emb @CustomOp.register("rotary_embedding") class RotaryEmbedding(CustomOp): """Original rotary positional embedding.""" def __init__( self, head_size: int, rotary_dim: int, max_position_embeddings: int, base: int | float, is_neox_style: bool, dtype: torch.dtype, ) -> None: super().__init__() self.head_size = head_size self.rotary_dim = rotary_dim self.max_position_embeddings = max_position_embeddings self.base = base self.is_neox_style = is_neox_style self.dtype = dtype cache = self._compute_cos_sin_cache() cache = cache.to(dtype) self.cos_sin_cache: torch.Tensor self.register_buffer("cos_sin_cache", cache, persistent=False) def _compute_inv_freq(self, base: int | float) -> torch.Tensor: """Compute the inverse frequency.""" # NOTE(woosuk): To exactly match the HF implementation, we need to # use CPU to compute the cache and then move it to GPU. However, we # create the cache on GPU for faster initialization. This may cause # a slight numerical difference between the HF implementation and ours. inv_freq = 1.0 / ( base ** ( torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim ) ) return inv_freq def _compute_cos_sin_cache(self) -> torch.Tensor: """Compute the cos and sin cache.""" inv_freq = self._compute_inv_freq(self.base) t = torch.arange(self.max_position_embeddings, dtype=torch.float) freqs = torch.einsum("i,j -> ij", t, inv_freq) cos = freqs.cos() sin = freqs.sin() cache = torch.cat((cos, sin), dim=-1) return cache def forward_cuda(self, *args, **kwargs): return self.forward_native(*args, **kwargs) def forward_native( self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor, offsets: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """A PyTorch-native implementation of forward().""" if offsets is not None: positions = positions + offsets positions = positions.flatten() num_tokens = positions.shape[0] cos_sin = self.cos_sin_cache.index_select(0, positions) cos, sin = cos_sin.chunk(2, dim=-1) query_shape = query.shape query = query.view(num_tokens, -1, self.head_size) query_rot = query[..., : self.rotary_dim] query_pass = query[..., self.rotary_dim :] query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style) query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape) key_shape = key.shape key = key.view(num_tokens, -1, self.head_size) key_rot = key[..., : self.rotary_dim] key_pass = key[..., self.rotary_dim :] key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style) key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) return query, key def extra_repr(self) -> str: s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}" s += f", max_position_embeddings={self.max_position_embeddings}" s += f", base={self.base}, is_neox_style={self.is_neox_style}" return s class LinearScalingRotaryEmbedding(RotaryEmbedding): def __init__( self, head_size: int, rotary_dim: int, max_position_embeddings: int, base: int | float, is_neox_style: bool, dtype: torch.dtype, scaling_factor: float, ) -> None: self.scaling_factor = float(scaling_factor) super().__init__( head_size=head_size, rotary_dim=rotary_dim, max_position_embeddings=max_position_embeddings, base=base, is_neox_style=is_neox_style, dtype=dtype, ) def _compute_cos_sin_cache(self) -> torch.Tensor: inv_freq = self._compute_inv_freq(self.base) t = torch.arange(self.max_position_embeddings, dtype=torch.float) t = t / self.scaling_factor freqs = torch.einsum("i,j -> ij", t, inv_freq) cos = freqs.cos() sin = freqs.sin() cache = torch.cat((cos, sin), dim=-1) return cache ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/rotary_embedding/factory.py ================================================ """get_rope / get_rotary_pos_embed factory functions and module-level caches.""" from collections import OrderedDict from typing import Any import torch from .base import LinearScalingRotaryEmbedding, RotaryEmbedding from .mrope import NDRotaryEmbedding, _to_tuple _ROPE_DICT: dict[tuple, RotaryEmbedding] = {} _ND_ROPE_CACHE: "OrderedDict[tuple, NDRotaryEmbedding]" = OrderedDict() _ROPE_3D_CACHE: "OrderedDict[tuple, tuple[torch.Tensor, torch.Tensor]]" = OrderedDict() def get_rope( head_size: int, rotary_dim: int, max_position: int, base: int | float, is_neox_style: bool = True, rope_scaling: dict[str, Any] | None = None, dtype: torch.dtype | None = None, partial_rotary_factor: float = 1.0, ) -> RotaryEmbedding: if dtype is None: dtype = torch.get_default_dtype() if rope_scaling is not None: # Transforms every value that is a list into a tuple for caching calls rope_scaling_tuple = { k: tuple(v) if isinstance(v, list) else v for k, v in rope_scaling.items() } rope_scaling_args = tuple(rope_scaling_tuple.items()) else: rope_scaling_args = None if partial_rotary_factor < 1.0: rotary_dim = int(rotary_dim * partial_rotary_factor) max_position_embeddings = max_position rope_type = None if rope_scaling is not None: rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) if rope_type in (None, "default"): rope_scaling = None elif rope_type == "linear": factor = float(rope_scaling.get("factor", 1.0)) original_max = rope_scaling.get("original_max_position_embeddings", None) if original_max is not None: max_position_embeddings = max( max_position_embeddings, int(float(original_max) * factor) ) key = ( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, rope_scaling_args, dtype, ) if key in _ROPE_DICT: return _ROPE_DICT[key] if rope_scaling is None: rotary_emb = RotaryEmbedding( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype ) else: if rope_type == "linear": factor = float(rope_scaling.get("factor", 1.0)) rotary_emb = LinearScalingRotaryEmbedding( head_size=head_size, rotary_dim=rotary_dim, max_position_embeddings=max_position_embeddings, base=base, is_neox_style=is_neox_style, dtype=dtype, scaling_factor=factor, ) else: raise ValueError(f"Unknown RoPE scaling {rope_scaling}") _ROPE_DICT[key] = rotary_emb return rotary_emb def get_rotary_pos_embed( rope_sizes, hidden_size, heads_num, rope_dim_list, rope_theta, theta_rescale_factor=1.0, interpolation_factor=1.0, shard_dim: int = 0, dtype: torch.dtype = torch.float32, start_frame: int = 0, device: torch.device | str | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Generate rotary positional embeddings for the given sizes. Args: rope_sizes: Tuple of dimensions (t, h, w) hidden_size: Hidden dimension size heads_num: Number of attention heads rope_dim_list: List of dimensions for each axis, or None rope_theta: Base for frequency calculations theta_rescale_factor: Rescale factor for theta. Defaults to 1.0 interpolation_factor: Factor to scale positions. Defaults to 1.0 shard_dim: Which dimension to shard for sequence parallelism. Defaults to 0. Returns: Tuple of (cos, sin) tensors for rotary embeddings """ target_ndim = 3 head_dim = hidden_size // heads_num if rope_dim_list is None: rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] assert ( sum(rope_dim_list) == head_dim ), "sum(rope_dim_list) should equal to head_dim of attention layer" # Get SP info - now handled within NDRotaryEmbedding # sp_group = get_sp_group() # sp_rank = sp_group.rank_in_group # sp_world_size = sp_group.world_size # Simple LRU cache keyed by parameters global _ND_ROPE_CACHE key = ( tuple(rope_dim_list), float(rope_theta), ( tuple(theta_rescale_factor) if isinstance(theta_rescale_factor, list) else float(theta_rescale_factor) ), ( tuple(interpolation_factor) if isinstance(interpolation_factor, list) else float(interpolation_factor) ), dtype, ) cache_hit = key in _ND_ROPE_CACHE if cache_hit: rope_emb = _ND_ROPE_CACHE.pop(key) _ND_ROPE_CACHE[key] = rope_emb # move to end (most-recent) else: rope_emb = NDRotaryEmbedding( rope_dim_list=rope_dim_list, rope_theta=rope_theta, theta_rescale_factor=theta_rescale_factor, interpolation_factor=interpolation_factor, dtype=dtype, ) _ND_ROPE_CACHE[key] = rope_emb if len(_ND_ROPE_CACHE) > 16: # pop least-recently-used _ND_ROPE_CACHE.pop(next(iter(_ND_ROPE_CACHE))) freqs_cos, freqs_sin = rope_emb.forward_from_grid( grid_size=_to_tuple(rope_sizes, dim=3), shard_dim=shard_dim, start_frame=start_frame, device=device, ) return freqs_cos, freqs_sin ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py ================================================ """MRotaryEmbedding, YaRNScalingMRotaryEmbedding, NDRotaryEmbedding, OneDRotaryEmbedding.""" import functools import torch from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_group def _to_tuple(x: int | tuple[int, ...], dim: int = 2) -> tuple[int, ...]: if isinstance(x, int): return (x,) * dim elif len(x) == dim: return x else: raise ValueError(f"Expected length {dim} or int, but got {x}") def get_1d_rotary_pos_embed( dim: int, pos: torch.FloatTensor | int, theta: float = 10000.0, theta_rescale_factor: float = 1.0, interpolation_factor: float = 1.0, dtype: torch.dtype = torch.float32, device: torch.device | str | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Precompute the frequency tensor for complex exponential (cis) with given dimensions. (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.) This function calculates a frequency tensor with complex exponential using the given dimension 'dim' and the end index 'end'. The 'theta' parameter scales the frequencies. Args: dim (int): Dimension of the frequency tensor. pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0. theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0. interpolation_factor (float, optional): Factor to scale positions. Defaults to 1.0. Returns: freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D] """ if isinstance(pos, int): pos = torch.arange(pos, dtype=dtype, device=device) elif ( isinstance(pos, torch.Tensor) and device is not None and pos.device != torch.device(device) ): # Ensure positions are on the requested device to avoid implicit CPU ops. pos = pos.to(device) # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning # has some connection to NTK literature if theta_rescale_factor != 1.0: theta *= theta_rescale_factor ** (dim / (dim - 2)) freqs = 1.0 / ( theta ** (torch.arange(0, dim, 2, device=device)[: (dim // 2)].to(dtype) / dim).to( device=device ) ) # [D/2] freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2] freqs_cos = freqs.cos() # [S, D/2] freqs_sin = freqs.sin() # [S, D/2] return freqs_cos, freqs_sin class OneDRotaryEmbedding(torch.nn.Module): """1D rotary positional embedding with caching.""" def __init__( self, dim: int, theta: float = 10000.0, theta_rescale_factor: float = 1.0, interpolation_factor: float = 1.0, dtype: torch.dtype = torch.float32, use_real: bool = False, repeat_interleave_real: bool = False, ): super().__init__() assert dim % 2 == 0 self.dim = dim self.theta = theta self.theta_rescale_factor = theta_rescale_factor self.interpolation_factor = interpolation_factor # dtype of freqs self.dtype = dtype self.use_real = use_real self.repeat_interleave_real = repeat_interleave_real def build_freqs(self, device): freqs = 1.0 / ( self.theta ** ( torch.arange(0, self.dim, 2, dtype=self.dtype, device=device)[ : (self.dim // 2) ] / self.dim ).to(device=device) ) return freqs def build_freqs_outer(self, pos: torch.Tensor, device): theta = self.theta # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning # has some connection to NTK literature if self.theta_rescale_factor != 1.0: theta *= self.theta_rescale_factor ** (self.dim / (self.dim - 2)) freqs = self.build_freqs(device) freqs = torch.outer(pos * self.interpolation_factor, freqs) freqs_cos = freqs.cos() freqs_sin = freqs.sin() if self.use_real and self.repeat_interleave_real: freqs_cos = freqs_cos.repeat_interleave(2, dim=1) freqs_sin = freqs_sin.repeat_interleave(2, dim=1) return freqs_cos.float(), freqs_sin.float() @functools.lru_cache(maxsize=16) def forward_from_grid( self, seq_len: int, start_pos: int, device_str: str ) -> tuple[torch.Tensor, torch.Tensor]: device = torch.device(device_str) pos = torch.arange( start_pos, start_pos + seq_len, dtype=self.dtype, device=device ) freqs_cos, freqs_sin = self.build_freqs_outer(pos, device) return freqs_cos, freqs_sin def forward(self, pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Calculates 1D rotary embeddings for the given positions. This method converts the input tensor to a hashable representation and calls a cached helper method to perform the computation. """ pos_tuple = tuple(pos.tolist()) device_str = str(pos.device) return self._forward_cached(pos_tuple, device_str) @functools.lru_cache(maxsize=16) def _forward_cached( self, pos_tuple: tuple, device_str: str ) -> tuple[torch.Tensor, torch.Tensor]: """ The core implementation that computes 1D rotary embeddings. This method is wrapped by an LRU cache. """ device = torch.device(device_str) pos = torch.as_tensor(pos_tuple, dtype=self.dtype, device=device) freqs_cos, freqs_sin = self.build_freqs_outer(pos, device) return freqs_cos, freqs_sin class NDRotaryEmbedding(torch.nn.Module): """N-dimensional rotary positional embedding.""" def __init__( self, rope_dim_list: list[int], rope_theta: float, theta_rescale_factor: float | list[float] = 1.0, interpolation_factor: float | list[float] = 1.0, use_real: bool = False, repeat_interleave_real: bool = False, dtype: torch.dtype = torch.float32, ): super().__init__() self.rope_dim_list = rope_dim_list self.ndim = len(rope_dim_list) self.rope_theta = rope_theta # dtype of freqs # does not control the output dtype self.dtype = dtype if isinstance(theta_rescale_factor, (int, float)): self.theta_rescale_factor = [theta_rescale_factor] * self.ndim elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1: self.theta_rescale_factor = [theta_rescale_factor[0]] * self.ndim else: self.theta_rescale_factor = theta_rescale_factor assert ( len(self.theta_rescale_factor) == self.ndim ), "len(theta_rescale_factor) should equal to len(rope_dim_list)" if isinstance(interpolation_factor, (int, float)): self.interpolation_factor = [interpolation_factor] * self.ndim elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1: self.interpolation_factor = [interpolation_factor[0]] * self.ndim else: self.interpolation_factor = interpolation_factor assert ( len(self.interpolation_factor) == self.ndim ), "len(interpolation_factor) should equal to len(rope_dim_list)" self.rope_generators: list[OneDRotaryEmbedding] = torch.nn.ModuleList() _config_to_gen_idx: dict[tuple, int] = {} self.dim_idx_to_gen_idx: list[int] = [] for i in range(self.ndim): dim = self.rope_dim_list[i] rescale = self.theta_rescale_factor[i] interp = self.interpolation_factor[i] config_key = (dim, rescale, interp, use_real, repeat_interleave_real) if config_key not in _config_to_gen_idx: generator = OneDRotaryEmbedding( dim=dim, theta=self.rope_theta, theta_rescale_factor=rescale, interpolation_factor=interp, dtype=self.dtype, use_real=use_real, repeat_interleave_real=repeat_interleave_real, ) _config_to_gen_idx[config_key] = len(self.rope_generators) self.rope_generators.append(generator) gen_idx = _config_to_gen_idx[config_key] self.dim_idx_to_gen_idx.append(gen_idx) def forward(self, positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Calculates n-d rotary embeddings for given absolute positions. Args: positions (torch.Tensor): A tensor of shape `[num_tokens, ndim]` containing the integer coordinates for each token. Returns: A tuple of (cos, sin) tensors. """ # Caching wrapper: convert tensor to a hashable tuple of tuples. pos_tuple = tuple(map(tuple, positions.tolist())) device_str = str(positions.device) return self._forward_cached(pos_tuple, device_str) @functools.lru_cache(maxsize=16) def _forward_cached( self, pos_tuple: tuple[tuple[int, ...], ...], device_str: str ) -> tuple[torch.Tensor, torch.Tensor]: """ The core implementation that computes embeddings from a position tensor. This method is wrapped by an LRU cache. """ device = torch.device(device_str) positions = torch.tensor(pos_tuple, dtype=torch.long, device=device) return self.forward_uncached(pos=positions) def forward_uncached(self, pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ The core implementation that computes embeddings from a position tensor. This method is wrapped by an LRU cache. """ device = pos.device # Pre-allocate the final tensors for efficiency. num_tokens = pos.shape[0] first_generator = self.rope_generators[0] if first_generator.use_real and first_generator.repeat_interleave_real: head_dim = sum(self.rope_dim_list) else: head_dim = sum(self.rope_dim_list) // 2 cos = torch.empty((num_tokens, head_dim), device=device, dtype=self.dtype) sin = torch.empty((num_tokens, head_dim), device=device, dtype=self.dtype) col_offset = 0 for i in range(self.ndim): # Extract position coordinates for the current dimension for all tokens. pos_i = pos[:, i].to(self.dtype) # Get the appropriate 1D generator. gen_idx = self.dim_idx_to_gen_idx[i] generator = self.rope_generators[gen_idx] # Calculate 1D embeddings. cos_1d, sin_1d = generator(pos_i) slice_width = cos_1d.shape[1] cos[:, col_offset : col_offset + slice_width] = cos_1d sin[:, col_offset : col_offset + slice_width] = sin_1d col_offset += slice_width return cos.float(), sin.float() def forward_from_grid( self, grid_size: tuple[int, ...], shard_dim: int = 0, start_frame: int = 0, device: torch.device | str | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Handles sp internally """ # Caching wrapper: use grid parameters directly as the key. # grid_tuple = _to_tuple(grid_size, dim=self.ndim) device_str = str(device) if device is not None else "cpu" return self._forward_cached_from_grid( grid_size, shard_dim, start_frame, device_str ) @functools.lru_cache(maxsize=16) def _forward_cached_from_grid( self, grid_size: tuple[int, ...], shard_dim: int, start_frame: int, device_str: str, ) -> tuple[torch.Tensor, torch.Tensor]: """ Computes embeddings for a structured grid, using a highly efficient implementation that avoids materializing the full position tensor. This method is wrapped by an LRU cache. """ device = torch.device(device_str) sp_group = get_sp_group() sp_rank = sp_group.rank_in_group sp_world_size = sp_group.world_size sizes = _to_tuple(grid_size, dim=self.ndim) starts = (0,) * self.ndim # Apply sequence parallel sharding to the sizes and compute shard offset shard_sizes = list(sizes) shard_offsets = [0] * self.ndim if sp_world_size > 1: assert sizes[shard_dim] % sp_world_size == 0, ( f"Dimension {shard_dim} with size {sizes[shard_dim]} is not divisible " f"by sequence parallel world size {sp_world_size}" ) shard_size = sizes[shard_dim] // sp_world_size shard_offsets[shard_dim] = sp_rank * shard_size shard_sizes[shard_dim] = shard_size # Pre-allocate outputs on the requested device to avoid CPU ops and extra cats num_tokens = 1 for s in shard_sizes: num_tokens *= int(s) head_dim_half = sum(self.rope_dim_list) // 2 cos = torch.empty((num_tokens, head_dim_half), device=device, dtype=self.dtype) sin = torch.empty((num_tokens, head_dim_half), device=device, dtype=self.dtype) # Compute per-axis 1D embeddings once and expand via repeats to [N, d_i/2] col_offset = 0 for i in range(self.ndim): dim_i = self.rope_dim_list[i] dim_i_half = dim_i // 2 size_i = int(shard_sizes[i]) # Starting position for this axis, with optional frame offset for time axis (i==0) base_offset = starts[i] if i == 0 and start_frame > 0: base_offset += start_frame if sp_world_size > 1 and i == shard_dim: base_offset += shard_offsets[i] gen_idx = self.dim_idx_to_gen_idx[i] generator = self.rope_generators[gen_idx] cos_1d, sin_1d = generator.forward_from_grid( size_i, base_offset, device_str ) # Expand to [num_tokens, dim_i/2] matching flatten order (last dims vary fastest) repeats_per_entry = 1 for j in range(i + 1, self.ndim): repeats_per_entry *= int(shard_sizes[j]) tile_count = 1 for j in range(0, i): tile_count *= int(shard_sizes[j]) cos_expanded = cos_1d.repeat_interleave(repeats_per_entry, dim=0) sin_expanded = sin_1d.repeat_interleave(repeats_per_entry, dim=0) if tile_count > 1: cos_expanded = cos_expanded.repeat(tile_count, 1) sin_expanded = sin_expanded.repeat(tile_count, 1) cos[:, col_offset : col_offset + dim_i_half] = cos_expanded sin[:, col_offset : col_offset + dim_i_half] = sin_expanded col_offset += dim_i_half return cos.float(), sin.float() ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py ================================================ """Primitive RoPE ops: rotate helpers and apply_rotary_emb utilities.""" from typing import Optional, Tuple import torch from sglang.jit_kernel.diffusion.triton.rotary import apply_rotary_embedding from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.srt.utils.custom_op import register_custom_op_from_extern _is_cuda = current_platform.is_cuda() if _is_cuda: try: from flashinfer.rope import ( apply_rope_with_cos_sin_cache_inplace as _flashinfer_apply_rope_inplace, ) except Exception: _flashinfer_apply_rope_inplace = None else: _flashinfer_apply_rope_inplace = None if _flashinfer_apply_rope_inplace is not None: flashinfer_apply_rope_inplace = register_custom_op_from_extern( _flashinfer_apply_rope_inplace, op_name="flashinfer_apply_rope_with_cos_sin_cache_inplace", mutates_args=["query", "key"], ) else: flashinfer_apply_rope_inplace = None def _apply_rotary_emb( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, is_neox_style: bool, interleaved: bool = False, ) -> torch.Tensor: """ Args: x: [num_tokens, num_heads, head_size] or [num_tokens, head_size] cos: [num_tokens, head_size // 2] sin: [num_tokens, head_size // 2] is_neox_style: Whether to use the Neox-style or GPT-J-style rotary positional embeddings. """ # cos = cos.unsqueeze(-2).to(x.dtype) # sin = sin.unsqueeze(-2).to(x.dtype) if is_neox_style: cos = cos.unsqueeze(-2) sin = sin.unsqueeze(-2) if is_neox_style: x1, x2 = torch.chunk(x, 2, dim=-1) else: x1 = x[..., ::2] x2 = x[..., 1::2] o1 = (x1.float() * cos - x2.float() * sin).type_as(x) o2 = (x2.float() * cos + x1.float() * sin).type_as(x) return torch.cat((o1, o2), dim=-1) else: return apply_rotary_embedding(x, cos, sin, interleaved) def apply_flashinfer_rope_qk_inplace( q: torch.Tensor, k: torch.Tensor, cos_sin_cache: torch.Tensor, *, head_size: Optional[int] = None, is_neox: bool = False, positions: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: if q.dim() != 4 or k.dim() != 4: raise ValueError( f"Expected q/k to be 4D [bsz, seqlen, nheads, head_size], " f"got q:{tuple(q.shape)} k:{tuple(k.shape)}" ) if q.shape != k.shape: raise ValueError( f"q and k must have the same shape, got {q.shape} vs {k.shape}" ) if not (isinstance(cos_sin_cache, torch.Tensor) and cos_sin_cache.dim() == 2): raise ValueError("cos_sin_cache must be a 2D torch.Tensor") bsz, seqlen, nheads, d = q.shape if head_size is None: head_size = d if head_size != d: raise ValueError(f"head_size mismatch: inferred {d}, but head_size={head_size}") if flashinfer_apply_rope_inplace is None: # Triton fallback for AMD/ROCm where FlashInfer is not available import warnings warnings.warn( "FlashInfer not available, using Triton fallback for RoPE", stacklevel=2, ) half_size = cos_sin_cache.shape[-1] // 2 if positions is None: cos = cos_sin_cache[:seqlen, :half_size].to(q.dtype) sin = cos_sin_cache[:seqlen, half_size:].to(q.dtype) cos = cos.unsqueeze(0).expand(bsz, -1, -1).reshape(bsz * seqlen, -1) sin = sin.unsqueeze(0).expand(bsz, -1, -1).reshape(bsz * seqlen, -1) else: positions = positions.to(cos_sin_cache.device).view(-1) cos = cos_sin_cache[positions, :half_size].to(q.dtype) sin = cos_sin_cache[positions, half_size:].to(q.dtype) q_flat = q.reshape(bsz * seqlen, nheads, d) k_flat = k.reshape(bsz * seqlen, nheads, d) q_rot = apply_rotary_embedding(q_flat, cos, sin, interleaved=not is_neox) k_rot = apply_rotary_embedding(k_flat, cos, sin, interleaved=not is_neox) return q_rot.view(bsz, seqlen, nheads, d), k_rot.view(bsz, seqlen, nheads, d) if positions is None: pos_1d = torch.arange(seqlen, device=q.device, dtype=torch.long) positions = pos_1d if bsz == 1 else pos_1d.repeat(bsz) else: if not ( isinstance(positions, torch.Tensor) and positions.dtype == torch.long and positions.dim() == 1 ): raise ValueError("positions must be a 1D torch.long Tensor") if positions.numel() != bsz * seqlen: raise ValueError( f"positions length must be bsz*seqlen={bsz*seqlen}, got {positions.numel()}" ) q_flat = q.reshape(bsz * seqlen, nheads * d).contiguous() k_flat = k.reshape(bsz * seqlen, nheads * d).contiguous() flashinfer_apply_rope_inplace( positions=positions, query=q_flat, key=k_flat, head_size=d, cos_sin_cache=cos_sin_cache, is_neox=is_neox, ) return q_flat.view(bsz, seqlen, nheads, d), k_flat.view(bsz, seqlen, nheads, d) ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/usp.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import logging from typing import TYPE_CHECKING import torch import torch.distributed._functional_collectives as ft_c from torch.distributed.tensor.experimental._attention import _cp_options from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_sp_group, get_ulysses_parallel_world_size, ) from sglang.srt.utils.common import torch_release _cp_options.enable_load_balance = False if TYPE_CHECKING: from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionImpl, ) logger = logging.getLogger(__name__) def _maybe_wait(tensor: torch.Tensor) -> torch.Tensor: """ When tracing the code, the result tensor is not an AsyncCollectiveTensor, so we cannot call ``wait()``. """ if isinstance(tensor, ft_c.AsyncCollectiveTensor): return tensor.wait() return tensor def _usp_all_to_all_single(x: torch.Tensor) -> torch.Tensor: ulysses_pg = get_sp_group().ulysses_group assert ulysses_pg is not None, "Ulysses process group is not initialized." x_shape = x.shape x = x.flatten() x = ft_c.all_to_all_single( x, output_split_sizes=None, input_split_sizes=None, group=ulysses_pg ) x = _maybe_wait(x) x = x.reshape(x_shape) return x def _usp_input_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor: """ Perform Ulysses-style input all-to-all over the head dimension. Default layout expects heads at dim=1 and sequence at dim=2: [b, h, s_local, d] -> [b, h_local, s_global, d] If heads are at dim=2 (input is [b, s_local, h, d]), set head_dim=2, and the function returns [b, s_global, h_local, d], preserving the original head/sequence dim ordering. Args: x: A 4D tensor with layout [b, *, *, d] where '*' are sequence and heads head_dim: Which dimension index corresponds to heads (1 or 2) Returns: Tensor with the same dim order as input, with heads sharded and sequence gathered. """ world_size = get_ulysses_parallel_world_size() if world_size <= 1: return x assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}" assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}" # Move the dimension to be split (h_global) to dim 0 for all_to_all_single if head_dim == 1: b, h_global, s_local, d = x.shape # Shape transition: [b, h_global, s_local, d] -> [h_global, b, s_local, d] permute_order = (1, 0, 2, 3) else: # head_dim == 2 b, s_local, h_global, d = x.shape # Shape transition: [b, s_local, h_global, d] -> [h_global, b, s_local, d] permute_order = (2, 0, 1, 3) assert ( h_global % world_size == 0 ), f"h_global ({h_global}) must be divisible by world_size ({world_size})" h_local, s_global = h_global // world_size, s_local * world_size x = x.permute(permute_order).contiguous() x = _usp_all_to_all_single(x) x = x.reshape(world_size, h_local, b, s_local, d) # Reorder dims to place 'world_size' adjacent to 's_local' to merge them into 's_global' if head_dim == 1: # Shape transition: [world_size, h_local, b, s_local, d] -> [b, h_local, world_size, s_local, d] x = x.permute(2, 1, 0, 3, 4).contiguous().reshape(b, h_local, s_global, d) else: # head_dim == 2 # Shape transition: [world_size, h_local, b, s_local, d] -> [b, world_size, s_local, h_local, d] x = x.permute(2, 0, 3, 1, 4).contiguous().reshape(b, s_global, h_local, d) return x def _usp_output_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor: """ Perform Ulysses-style output all-to-all over the head dimension (inverse of input). Default layout expects heads at dim=1 and sequence at dim=2: [b, h_local, s, d] -> [b, h, s_local, d] If heads are at dim=2 (input is [b, s_global, h // world_size, d]), set head_dim=2, and the function returns [b, s_local, h, d], preserving the original head/sequence dim ordering. Args: x: A 4D tensor with layout [b, *, *, d] where '*' are sequence and heads head_dim: Which dimension index corresponds to heads (1 or 2) Returns: Tensor with the same dim order as input, with heads gathered and sequence sharded. """ world_size = get_ulysses_parallel_world_size() if world_size <= 1: return x assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}" assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}" # Move the dimension to be split (s_global) to dim 0 for all_to_all_single if head_dim == 1: b, h_local, s_global, d = x.shape # Shape transition: [b, h_local, s_global, d] -> [s_global, b, h_local, d] permute_order = (2, 0, 1, 3) else: # head_dim == 2 b, s_global, h_local, d = x.shape # Shape transition: [b, s_global, h_local, d] -> [s_global, b, h_local, d] permute_order = (1, 0, 2, 3) assert ( s_global % world_size == 0 ), f"s_global ({s_global}) must be divisible by world_size ({world_size})" s_local, h_global = s_global // world_size, h_local * world_size x = x.permute(permute_order).contiguous() x = _usp_all_to_all_single(x) x = x.reshape(world_size, s_local, b, h_local, d) # Reorder dims to place 'world_size' adjacent to 'h_local' to merge them into 'h_global' if head_dim == 1: # Shape transition: [world_size, s_local, b, h_local, d] -> [b, world_size, h_local, s_local, d] x = x.permute(2, 0, 3, 1, 4).contiguous().reshape(b, h_global, s_local, d) else: # head_dim == 2 # Shape transition: [world_size, s_local, b, h_local, d] -> [b, s_local, world_size, h_local, d] x = x.permute(2, 1, 0, 3, 4).contiguous().reshape(b, s_local, h_global, d) return x def ring_attn( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_impl: "AttentionImpl", is_causal: bool = False, dropout_p: float = 0.0, ): """ Ring Attention implementation. This function implements Ring Attention, a strategy for distributed attention computation that reduces peak memory usage. It accepts a generic attention implementation (`attn_impl`) which is called by the underlying PyTorch distributed attention primitive. Args: query, key, value: The input tensors for attention. attn_impl: An instance of an attention implementation backend (e.g., FlashAttentionImpl) whose `forward` method will be used as the computational kernel. is_causal: Whether to apply causal masking. dropout_p: Dropout probability. """ # torch.distributed.tensor.experimental._attention is not a public API, from torch.distributed.tensor.experimental._attention import ( _templated_ring_attention, ) ring_pg = get_sp_group().ring_group assert ring_pg is not None, "Ring process group is not initialized." # Ring attention primitives expect tensors in [B, H, S, D] layout. # We permute the inputs here. query = torch.permute(query, [0, 2, 1, 3]).contiguous() key = torch.permute(key, [0, 2, 1, 3]).contiguous() value = torch.permute(value, [0, 2, 1, 3]).contiguous() # Create an adapter function that matches the signature expected by # _templated_ring_attention. The `attn_impl` already has dropout and # causal settings configured during its initialization. # Note: Please be aware that Attention Backend and Ring Attention may require different QKV tensor shapes. # For example, FlashAttention expects the format to be BSHD. def attn_callable_adapter(q, k, v, *args, **kwargs): # We ignore the dropout_p and is_causal passed by _templated_ring_attention # and rely on the pre-configured attn_impl. # The `attn_metadata` is not available here, so we pass None. # This is a limitation we must accept when using this experimental API. q = torch.permute(q, [0, 2, 1, 3]) k = torch.permute(k, [0, 2, 1, 3]) v = torch.permute(v, [0, 2, 1, 3]) # logger.warning(f"Warning: return_s·oftmax_lse is only supported for FlashAttentionImpl") output, softmax_lse, *rest = attn_impl.forward( q, k, v, attn_metadata=None, return_softmax_lse=True, ) output = torch.permute(output, [0, 2, 1, 3]) return output, softmax_lse, *rest # Starting from torch 2.6.0, _templated_ring_attention expects an integer # segment_id for the attention function. use_segment_id = torch_release >= (2, 6) attn_kwargs = dict( op=attn_callable_adapter, dropout_p=dropout_p, is_causal=is_causal, query=query, key=key, value=value, group=ring_pg, # https://github.com/pytorch/pytorch/blob/c907c778f42ba2fdaf25b733dd25baf9779c6a12/torch/distributed/tensor/experimental/_context_parallel/_attention.py#L309 ) if use_segment_id: # For torch >= 2.6, segment_id is required. The value '1' is a placeholder # as we are not using complex segmentation features. out, *_ = _templated_ring_attention( seq_dim=1, # segment_id **attn_kwargs, ) else: out, *_ = _templated_ring_attention( **attn_kwargs, ) # Permute the output back to [B, S, H, D] layout. output = torch.permute(out, [0, 2, 1, 3]) return output ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/utils.py """Utility methods for model layers.""" import inspect from typing import Any, Callable, List, Optional import torch from torch.library import Library from sglang.multimodal_gen.runtime.platforms import current_platform def get_group_size(group) -> int: if hasattr(group, "world_size"): return group.world_size # GroupCoordinator elif hasattr(group, "size") and callable(getattr(group, "size", None)): return group.size() # ProcessGroup else: raise ValueError(f"Unsupported group type: {type(group)}") def get_group_rank(group) -> int: if hasattr(group, "rank_in_group"): return group.rank_in_group # GroupCoordinator elif hasattr(group, "rank") and callable(getattr(group, "rank", None)): return group.rank() # ProcessGroup else: raise ValueError(f"Unsupported group type: {type(group)}") def get_token_bin_counts_and_mask( tokens: torch.Tensor, vocab_size: int, num_seqs: int, ) -> tuple[torch.Tensor, torch.Tensor]: # Compute the bin counts for the tokens. # vocab_size + 1 for padding. bin_counts = torch.zeros( (num_seqs, vocab_size + 1), dtype=torch.long, device=tokens.device ) bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens)) bin_counts = bin_counts[:, :vocab_size] mask = bin_counts > 0 return bin_counts, mask sglang_lib = Library("sglang", "FRAGMENT") # noqa def direct_register_custom_op( op_name: str, op_func: Callable, mutates_args: List[str], fake_impl: Optional[Callable] = None, target_lib: Optional[Library] = None, ): """ `torch.library.custom_op` can have significant overhead because it needs to consider complicated dispatching logic. This function directly registers a custom op and dispatches it to the CUDA backend. See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5 for more details. By default, the custom op is registered to the vLLM library. If you want to register it to a different library, you can pass the library object to the `target_lib` argument. IMPORTANT: the lifetime of the operator is tied to the lifetime of the library object. If you want to bind the operator to a different library, make sure the library object is alive when the operator is used. Note: This function will silently skip registration if the operator with the same name is already registered to avoid RuntimeError in multi-engine scenarios (e.g., VERL framework). """ import torch.library my_lib = target_lib or sglang_lib # Check if operator is already registered to avoid duplicate registration # This is important for scenarios where multiple SGLang engines run in the same process try: # Try to access the operator to see if it's already registered lib_name = my_lib.m.name if hasattr(my_lib.m, "name") else "sglang" if hasattr(torch.ops, lib_name) and hasattr( getattr(torch.ops, lib_name), op_name ): # Operator already exists, skip registration return except (AttributeError, RuntimeError): # Operator doesn't exist, proceed with registration pass if hasattr(torch.library, "infer_schema"): schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args) else: # for pytorch 2.4 import torch._custom_op.impl schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args) try: my_lib.define(op_name + schema_str) my_lib.impl( op_name, op_func, "CUDA" if not current_platform.is_npu() else "PrivateUse1" ) if fake_impl is not None: my_lib._register_fake(op_name, fake_impl) except RuntimeError as error: if "Tried to register an operator" in str(error) and "multiple times" in str( error ): # Silently ignore duplicate registration errors # This can happen in multi-engine scenarios pass else: # Re-raise other RuntimeErrors raise error except AttributeError as error: # Always re-raise AttributeError as it indicates missing dependencies raise error class CustomOpWrapper: def __init__( self, op_name: str, op_func: Callable, mutates_args: List[str], **extra_kwargs, ): self.op_name = op_name self.op_func = op_func self.mutates_args = mutates_args self.extra_kwargs = extra_kwargs self._impl: Optional[Callable] = None def __call__(self, *args, **kwargs): return self.real_impl(*args, **kwargs) @property def real_impl(self) -> Callable: if self._impl is None: if not hasattr(torch.ops.sglang, self.op_name): # NOTE(dark): if torch compile fail here, mark the decorator as eager # lazy registration does not work with torch compile direct_register_custom_op( op_name=self.op_name, op_func=self.op_func, mutates_args=self.mutates_args, fake_impl=self.fake_impl, ) self._impl = getattr(torch.ops.sglang, self.op_name) assert self._impl is not None return self._impl @property def fake_impl(self) -> Callable: if "fake_impl" in self.extra_kwargs: return self.extra_kwargs["fake_impl"] assert "out_shape" in self.extra_kwargs signature = inspect.signature(self.op_func) out_shape = self.extra_kwargs["out_shape"] # check out_shape in signature def fake_impl(*args, **kwargs): if out_shape is None: return None bound = signature.bind(*args, **kwargs) bound.apply_defaults() try: return torch.empty_like( bound.args[out_shape] if isinstance(out_shape, int) else bound.arguments[out_shape] ) except (IndexError, KeyError): raise RuntimeError( f"Cannot find output argument at position `{out_shape}` for " f"custom operator `{self.op_name}` with signature `{signature}`." ) return fake_impl # Real implementation def register_custom_op( fn: Optional[Callable] = None, *, op_name: Optional[str] = None, mutates_args: Optional[List[str]] = None, eager: bool = True, **extra_kwargs, ) -> Any: """ A decorator to register a custom operator. Example usage: ```python # inplace operator, out_shape is None by default @register_custom_op(mutates_args=["x"]) def add_1_(x: torch.Tensor) -> None: x.add_(1) # operator with output, out_shape indicates the position of output @register_custom_op(mutates_args=["x"], out_shape=0) def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return x.add_(y) ``` :param fn: The function to be registered as a custom operator. If None, return a decorator. :type fn: Callable :param op_name: The name of the operator. If None, use the function name :type op_name: Optional[str] :param mutates_args: A list of argument names that are mutated in-place. :type mutates_args: List[str] :param out_shape: The position (int for positional, str for keyword) of the output-shape tensor. It is used to generate a fake implementation for torch.compile compatibility. If the operator is inplace and has no output, set to None. :type out_shape: Optional[List[Union[int, str]]] :param fake_impl: A fake implementation for the operator. Only one of `out_shape` or `fake_impl` should be provided. :type fake_impl: Optional[Callable] :param eager: Whether to register the operator eagerly. If False, the registration will be deferred until the first call. If you met any issue with torch.compile, try to set eager=True. Currently, to avoid misuse, we set eager=True by default. :type eager: bool :return: The registered JIT custom operator, or a decorator. NOTE: the real register will occur at the first call of the function. :rtype: Callable """ extra_kwarg_keys = set(extra_kwargs.keys()) expected_kwarg_keys = set({"out_shape", "fake_impl"}) assert ( expected_kwarg_keys >= extra_kwarg_keys ), f"Unexpected extra kwargs: {extra_kwarg_keys - expected_kwarg_keys}" has_out_shape = "out_shape" in extra_kwargs has_fake_impl = "fake_impl" in extra_kwargs assert not ( has_out_shape and has_fake_impl ), "Only one of `out_shape` or `fake_impl` should be provided." # Assume inplace if neither out_shape nor fake_impl is provided if not (has_out_shape or has_fake_impl): extra_kwargs["out_shape"] = None def decorator(op_func: Callable) -> Callable: wrapper = CustomOpWrapper( op_name=op_name or op_func.__name__, op_func=op_func, mutates_args=mutates_args or [], **extra_kwargs, ) return wrapper.real_impl if eager else wrapper if fn is not None: return decorator(fn) return decorator ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/visual_embedding.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import math import torch import torch.nn as nn from diffusers.models.embeddings import ( CombinedTimestepGuidanceTextProjEmbeddings as _CombinedTimestepGuidanceTextProjEmbeddings, ) from diffusers.models.embeddings import ( CombinedTimestepTextProjEmbeddings as _CombinedTimestepTextProjEmbeddings, ) from diffusers.models.embeddings import ( PixArtAlphaTextProjection, TimestepEmbedding, ) from diffusers.models.embeddings import Timesteps as _Timesteps from diffusers.models.embeddings import ( get_timestep_embedding as timestep_embedding_diffusers, ) from sglang.jit_kernel.timestep_embedding import ( timestep_embedding as timestep_embedding_cuda, ) from sglang.multimodal_gen.runtime.layers.activation import get_act_fn from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.platforms import current_platform _is_cuda = current_platform.is_cuda() class PatchEmbed(nn.Module): """2D Image to Patch Embedding Image to Patch Embedding using Conv2d A convolution based approach to patchifying a 2D image w/ embedding projection. Based on the impl in https://github.com/google-research/vision_transformer Hacked together by / Copyright 2020 Ross Wightman Remove the _assert function in forward function to be compatible with multi-resolution images. """ def __init__( self, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True, bias=True, dtype=None, prefix: str = "", ): super().__init__() # Convert patch_size to 2-tuple if isinstance(patch_size, list | tuple): if len(patch_size) == 1: patch_size = (patch_size[0], patch_size[0]) else: patch_size = (patch_size, patch_size) self.patch_size = patch_size self.flatten = flatten self.proj = nn.Conv3d( in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias, dtype=dtype, ) self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() def forward(self, x): x = self.proj(x) if self.flatten: x = x.flatten(2).transpose(1, 2) # BCHW -> BNC x = self.norm(x) return x class Timesteps(_Timesteps): def forward(self, timesteps: torch.Tensor) -> torch.Tensor: if _is_cuda: return timestep_embedding_cuda( timesteps, self.num_channels, flip_sin_to_cos=self.flip_sin_to_cos, downscale_freq_shift=self.downscale_freq_shift, scale=self.scale, ) else: return timestep_embedding_diffusers( timesteps, self.num_channels, flip_sin_to_cos=self.flip_sin_to_cos, downscale_freq_shift=self.downscale_freq_shift, scale=self.scale, ) class CombinedTimestepGuidanceTextProjEmbeddings( _CombinedTimestepGuidanceTextProjEmbeddings ): def __init__(self, embedding_dim, pooled_projection_dim): nn.Module.__init__(self) # use sgld op self.time_proj = Timesteps( num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0 ) # use diffusers op self.timestep_embedder = TimestepEmbedding( in_channels=256, time_embed_dim=embedding_dim ) self.guidance_embedder = TimestepEmbedding( in_channels=256, time_embed_dim=embedding_dim ) self.text_embedder = PixArtAlphaTextProjection( pooled_projection_dim, embedding_dim, act_fn="silu" ) class CombinedTimestepTextProjEmbeddings(_CombinedTimestepTextProjEmbeddings): def __init__(self, embedding_dim, pooled_projection_dim): nn.Module.__init__(self) # use sgld op self.time_proj = Timesteps( num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0 ) # use diffusers op self.timestep_embedder = TimestepEmbedding( in_channels=256, time_embed_dim=embedding_dim ) self.text_embedder = PixArtAlphaTextProjection( pooled_projection_dim, embedding_dim, act_fn="silu" ) class TimestepEmbedder(nn.Module): """ Embeds scalar timesteps into vector representations. """ def __init__( self, hidden_size, act_layer="silu", frequency_embedding_size=256, max_period=10000, dtype=None, freq_dtype=torch.float32, prefix: str = "", ): super().__init__() self.frequency_embedding_size = frequency_embedding_size self.max_period = max_period self.mlp = MLP( frequency_embedding_size, hidden_size, hidden_size, act_type=act_layer, dtype=dtype, ) self.freq_dtype = freq_dtype def forward( self, t: torch.Tensor, timestep_seq_len: int | None = None ) -> torch.Tensor: t_freq = timestep_embedding( t, self.frequency_embedding_size, self.max_period, dtype=self.freq_dtype ).to(self.mlp.fc_in.weight.dtype) if timestep_seq_len is not None: assert ( t_freq.shape[0] % timestep_seq_len == 0 ), "timestep length is not divisible by timestep_seq_len" batch_size = t_freq.shape[0] // timestep_seq_len t_freq = t_freq.unflatten(0, (batch_size, timestep_seq_len)) # t_freq = t_freq.to(self.mlp.fc_in.weight.dtype) t_emb = self.mlp(t_freq) return t_emb def timestep_embedding( t: torch.Tensor, dim: int, max_period: int = 10000, dtype: torch.dtype = torch.float32, ) -> torch.Tensor: """ Create sinusoidal timestep embeddings. Args: t: Tensor of shape [B] with timesteps dim: Embedding dimension max_period: Controls the minimum frequency of the embeddings Returns: Tensor of shape [B, dim] with embeddings """ half = dim // 2 freqs = torch.exp( -math.log(max_period) * torch.arange(start=0, end=half, dtype=dtype, device=t.device) / half ) args = t[:, None].float() * freqs[None] embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) if dim % 2: embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) return embedding class ModulateProjection(nn.Module): """Modulation layer for DiT blocks.""" def __init__( self, hidden_size: int, factor: int = 2, act_layer: str = "silu", dtype: torch.dtype | None = None, prefix: str = "", ): super().__init__() self.factor = factor self.hidden_size = hidden_size self.linear = ColumnParallelLinear( hidden_size, hidden_size * factor, bias=True, gather_output=True, params_dtype=dtype, ) self.act = get_act_fn(act_layer) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.act(x) x, _ = self.linear(x) return x def unpatchify(x, t, h, w, patch_size, channels) -> torch.Tensor: """ Convert patched representation back to image space. Args: x: Tensor of shape [B, T*H*W, C*P_t*P_h*P_w] t, h, w: Temporal and spatial dimensions Returns: Unpatchified tensor of shape [B, C, T*P_t, H*P_h, W*P_w] """ assert x.ndim == 3, f"x.ndim: {x.ndim}" assert len(patch_size) == 3, f"patch_size: {patch_size}" assert t * h * w == x.shape[1], f"t * h * w: {t * h * w}, x.shape[1]: {x.shape[1]}" c = channels pt, ph, pw = patch_size x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw)) x = torch.einsum("nthwcopq->nctohpwq", x) imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw)) return imgs ================================================ FILE: python/sglang/multimodal_gen/runtime/layers/vocab_parallel_embedding.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from collections.abc import Sequence from dataclasses import dataclass import torch import torch.distributed as dist import torch.nn.functional as F from torch.nn.parameter import Parameter, UninitializedParameter from sglang.multimodal_gen.runtime.distributed import ( divide, get_tp_group, tensor_model_parallel_all_reduce, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding, ) from sglang.multimodal_gen.runtime.layers.utils import get_group_rank, get_group_size from sglang.multimodal_gen.runtime.models.parameter import BasevLLMParameter from sglang.multimodal_gen.runtime.models.utils import set_weight_attrs from sglang.multimodal_gen.runtime.platforms import current_platform DEFAULT_VOCAB_PADDING_SIZE = 64 class UnquantizedEmbeddingMethod(QuantizeMethodBase): """Unquantized method for embeddings.""" def create_weights( self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: list[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs, ): """Create weights for embedding layer.""" weight = Parameter( torch.empty( sum(output_partition_sizes), input_size_per_partition, dtype=params_dtype, ), requires_grad=False, ) set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) layer.register_parameter("weight", weight) set_weight_attrs(weight, extra_weight_attrs) def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: return F.linear(x, layer.weight, bias) def embedding(self, layer: torch.nn.Module, input_: torch.Tensor) -> torch.Tensor: return F.embedding(input_, layer.weight) def pad_vocab_size(vocab_size: int, pad_to: int = DEFAULT_VOCAB_PADDING_SIZE) -> int: """Pad the vocab size to the given value.""" return ((vocab_size + pad_to - 1) // pad_to) * pad_to def vocab_range_from_per_partition_vocab_size( per_partition_vocab_size: int, rank: int, offset: int = 0 ) -> Sequence[int]: index_f = rank * per_partition_vocab_size index_l = index_f + per_partition_vocab_size return index_f + offset, index_l + offset def vocab_range_from_global_vocab_size( global_vocab_size: int, rank: int, world_size: int, offset: int = 0 ) -> Sequence[int]: per_partition_vocab_size = divide(global_vocab_size, world_size) return vocab_range_from_per_partition_vocab_size( per_partition_vocab_size, rank, offset=offset ) @dataclass class VocabParallelEmbeddingShardIndices: """Indices for a shard of a vocab parallel embedding.""" padded_org_vocab_start_index: int padded_org_vocab_end_index: int padded_added_vocab_start_index: int padded_added_vocab_end_index: int org_vocab_start_index: int org_vocab_end_index: int added_vocab_start_index: int added_vocab_end_index: int @property def num_org_elements(self) -> int: return self.org_vocab_end_index - self.org_vocab_start_index @property def num_added_elements(self) -> int: return self.added_vocab_end_index - self.added_vocab_start_index @property def num_org_elements_padded(self) -> int: return self.padded_org_vocab_end_index - self.padded_org_vocab_start_index @property def num_added_elements_padded(self) -> int: return self.padded_added_vocab_end_index - self.padded_added_vocab_start_index @property def num_org_vocab_padding(self) -> int: return self.num_org_elements_padded - self.num_org_elements @property def num_added_vocab_padding(self) -> int: return self.num_added_elements_padded - self.num_added_elements @property def num_elements_padded(self) -> int: return self.num_org_elements_padded + self.num_added_elements_padded def __post_init__(self): # sanity checks assert self.padded_org_vocab_start_index <= self.padded_org_vocab_end_index assert self.padded_added_vocab_start_index <= self.padded_added_vocab_end_index assert self.org_vocab_start_index <= self.org_vocab_end_index assert self.added_vocab_start_index <= self.added_vocab_end_index assert self.org_vocab_start_index <= self.padded_org_vocab_start_index assert self.added_vocab_start_index <= self.padded_added_vocab_start_index assert self.org_vocab_end_index <= self.padded_org_vocab_end_index assert self.added_vocab_end_index <= self.padded_added_vocab_end_index assert self.num_org_elements <= self.num_org_elements_padded assert self.num_added_elements <= self.num_added_elements_padded @torch.compile( dynamic=True, backend=current_platform.simple_compile_backend, disable=current_platform.is_npu(), ) def get_masked_input_and_mask( input_: torch.Tensor, org_vocab_start_index: int, org_vocab_end_index: int, num_org_vocab_padding: int, added_vocab_start_index: int, added_vocab_end_index: int, ) -> tuple[torch.Tensor, torch.Tensor]: # torch.compile will fuse all of the pointwise ops below # into a single kernel, making it very fast org_vocab_mask = (input_ >= org_vocab_start_index) & (input_ < org_vocab_end_index) added_vocab_mask = (input_ >= added_vocab_start_index) & ( input_ < added_vocab_end_index ) added_offset = ( added_vocab_start_index - (org_vocab_end_index - org_vocab_start_index) - num_org_vocab_padding ) valid_offset = (org_vocab_start_index * org_vocab_mask) + ( added_offset * added_vocab_mask ) vocab_mask = org_vocab_mask | added_vocab_mask input_ = vocab_mask * (input_ - valid_offset) return input_, ~vocab_mask class VocabParallelEmbedding(torch.nn.Module): """Embedding parallelized in the vocabulary dimension. Adapted from torch.nn.Embedding, note that we pad the vocabulary size to make sure it is divisible by the number of model parallel GPUs. In order to support various loading methods, we ensure that LoRA-added embeddings are always at the end of TP-sharded tensors. In other words, we shard base embeddings and LoRA embeddings separately (both padded), and place them in the same tensor. In this example, we will have the original vocab size = 1010, added vocab size = 16 and padding to 64. Therefore, the total vocab size with padding will be 1088 (because we first pad 1010 to 1024, add 16, and then pad to 1088). Therefore, the tensor format looks like the following: TP1, rank 0 (no sharding): |< --------BASE-------- >|< -BASE PADDING-- >|< -----LORA------ >|< -LORA PADDING-- >| corresponding token_id: | 0 | 1 | ... | 1009 | -1 | ... | -1 | 1010 | ... | 1015 | -1 | ... | -1 | index: | 0 | 1 | ... | 1009 | 1010 | ... | 1023 | 1024 | ... | 1039 | 1040 | ... | 1087 | TP2, rank 0: |< --------------------BASE--------------------- >|< -----LORA------ >|< -LORA PADDING- >| corresponding token_id: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 1000 | ... | 1015 | -1 | ... | -1 | index: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 512 | ... | 527 | 520 | ... | 543 | TP2, rank 1: |< -----------BASE----------- >|< -BASE PADDING- >|< -----------LORA PADDING----------- >| corresponding token_id: | 512 | 513 | 514 | ... | 1009 | -1 | ... | -1 | -1 | ... | -1 | -1 | ... | -1 | index: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 512 | ... | 519 | 520 | ... | 543 | Args: num_embeddings: vocabulary size. embedding_dim: size of hidden state. params_dtype: type of the parameters. org_num_embeddings: original vocabulary size (without LoRA). padding_size: padding size for the vocabulary. quant_config: quant config for the layer prefix: full name of the layer in the state dict """ # noqa: E501 def __init__( self, num_embeddings: int, embedding_dim: int, params_dtype: torch.dtype | None = None, org_num_embeddings: int | None = None, padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, quant_config: QuantizationConfig | None = None, prefix: str = "", tp_group: dist.ProcessGroup = None, ): super().__init__() # Keep the input dimensions. tp_group = tp_group or get_tp_group() tp_rank = get_group_rank(tp_group) self.tp_size = get_group_size(tp_group) self.tp_group = tp_group self.num_embeddings = num_embeddings self.padding_size = padding_size self.org_vocab_size = org_num_embeddings or num_embeddings num_added_embeddings = num_embeddings - self.org_vocab_size self.org_vocab_size_padded = pad_vocab_size( self.org_vocab_size, self.padding_size ) self.num_embeddings_padded = pad_vocab_size( self.org_vocab_size_padded + num_added_embeddings, self.padding_size ) assert self.org_vocab_size_padded <= self.num_embeddings_padded self.shard_indices = self._get_indices( self.num_embeddings_padded, self.org_vocab_size_padded, self.num_embeddings, self.org_vocab_size, tp_rank, self.tp_size, ) self.embedding_dim = embedding_dim quant_method = None if quant_config is not None: quant_method = quant_config.get_quant_method(self, prefix=prefix) if quant_method is None: quant_method = UnquantizedEmbeddingMethod() # If we are making an embedding layer, then our quantization linear # method must implement the embedding operation. If we are another # layer type like ParallelLMHead, this is not important. is_embedding_layer = type(self.__class__) is VocabParallelEmbedding quant_method_implements_embedding = method_has_implemented_embedding( type(quant_method) ) if is_embedding_layer and not quant_method_implements_embedding: raise NotImplementedError( f"The class {type(quant_method).__name__} must implement " "the 'embedding' method, see UnquantizedEmbeddingMethod." ) self.quant_method: QuantizeMethodBase = quant_method if params_dtype is None: params_dtype = torch.get_default_dtype() # Divide the weight matrix along the vocaburaly dimension. self.num_added_embeddings = self.num_embeddings - self.org_vocab_size self.num_embeddings_per_partition = divide( self.num_embeddings_padded, self.tp_size ) assert ( self.shard_indices.num_elements_padded == self.num_embeddings_per_partition ) self.num_org_embeddings_per_partition = ( self.shard_indices.org_vocab_end_index - self.shard_indices.org_vocab_start_index ) self.num_added_embeddings_per_partition = ( self.shard_indices.added_vocab_end_index - self.shard_indices.added_vocab_start_index ) self.quant_method.create_weights( self, self.embedding_dim, [self.num_embeddings_per_partition], self.embedding_dim, self.num_embeddings_padded, params_dtype=params_dtype, weight_loader=self.weight_loader, ) @classmethod def _get_indices( cls, vocab_size_padded: int, org_vocab_size_padded: int, vocab_size: int, org_vocab_size: int, tp_rank: int, tp_size: int, ) -> VocabParallelEmbeddingShardIndices: """Get start and end indices for vocab parallel embedding, following the layout outlined in the class docstring, based on the given tp_rank and tp_size.""" num_added_embeddings_padded = vocab_size_padded - org_vocab_size_padded padded_org_vocab_start_index, padded_org_vocab_end_index = ( vocab_range_from_global_vocab_size(org_vocab_size_padded, tp_rank, tp_size) ) padded_added_vocab_start_index, padded_added_vocab_end_index = ( vocab_range_from_global_vocab_size( num_added_embeddings_padded, tp_rank, tp_size, offset=org_vocab_size ) ) # remove padding org_vocab_start_index = min(padded_org_vocab_start_index, org_vocab_size) org_vocab_end_index = min(padded_org_vocab_end_index, org_vocab_size) added_vocab_start_index = min(padded_added_vocab_start_index, vocab_size) added_vocab_end_index = min(padded_added_vocab_end_index, vocab_size) return VocabParallelEmbeddingShardIndices( padded_org_vocab_start_index, padded_org_vocab_end_index, padded_added_vocab_start_index, padded_added_vocab_end_index, org_vocab_start_index, org_vocab_end_index, added_vocab_start_index, added_vocab_end_index, ) def get_sharded_to_full_mapping(self) -> list[int] | None: """Get a mapping that can be used to reindex the gathered logits for sampling. During sampling, we gather logits from all ranks. The relationship of index->token_id will follow the same format as outlined in the class docstring. However, after the gather, we want to reindex the final logits tensor to map index->token_id one-to-one (the index is always equal the token_id it corresponds to). The indices returned by this method allow us to do that. """ if self.tp_size < 2: return None base_embeddings: list[int] = [] added_embeddings: list[int] = [] padding: list[int] = [] for tp_rank in range(self.tp_size): shard_indices = self._get_indices( self.num_embeddings_padded, self.org_vocab_size_padded, self.num_embeddings, self.org_vocab_size, tp_rank, self.tp_size, ) range_start = self.num_embeddings_per_partition * tp_rank range_end = self.num_embeddings_per_partition * (tp_rank + 1) base_embeddings.extend( range(range_start, range_start + shard_indices.num_org_elements) ) padding.extend( range( range_start + shard_indices.num_org_elements, range_start + shard_indices.num_org_elements_padded, ) ) added_embeddings.extend( range( range_start + shard_indices.num_org_elements_padded, range_start + shard_indices.num_org_elements_padded + shard_indices.num_added_elements, ) ) padding.extend( range( range_start + shard_indices.num_org_elements_padded + shard_indices.num_added_elements, range_start + shard_indices.num_org_elements_padded + shard_indices.num_added_elements_padded, ) ) assert ( range_start + shard_indices.num_org_elements_padded + shard_indices.num_added_elements_padded == range_end ) ret = base_embeddings + added_embeddings + padding assert len(ret) == self.num_embeddings_padded return ret def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): output_dim = getattr(param, "output_dim", None) packed_dim = getattr(param, "packed_dim", None) # If the parameter is a gguf weight, then load it directly. if getattr(param, "is_gguf_weight_type", None): param.data.copy_(loaded_weight) param.weight_type = loaded_weight.item() return elif isinstance(param, UninitializedParameter): shape = list(loaded_weight.shape) if output_dim is not None: shape[output_dim] = self.num_embeddings_per_partition param.materialize(tuple(shape), dtype=loaded_weight.dtype) # If parameter does not have output dim, then it should # be copied onto all gpus (e.g. g_idx for act_order gptq). if output_dim is None: assert param.data.shape == loaded_weight.shape param.data.copy_(loaded_weight) return # Shard indexes for loading the weight start_idx = self.shard_indices.org_vocab_start_index shard_size = self.shard_indices.org_vocab_end_index - start_idx # If param packed on the same dim we are sharding on, then # need to adjust offsets of loaded weight by pack_factor. if packed_dim is not None and packed_dim == output_dim: packed_factor = ( param.packed_factor if isinstance(param, BasevLLMParameter) else param.pack_factor ) assert loaded_weight.shape[output_dim] == ( self.org_vocab_size // param.packed_factor ) start_idx = start_idx // packed_factor shard_size = shard_size // packed_factor else: assert loaded_weight.shape[output_dim] == self.org_vocab_size # Copy the data. Select chunk corresponding to current shard. loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) param[: loaded_weight.shape[0]].data.copy_(loaded_weight) param[loaded_weight.shape[0] :].data.fill_(0) def forward(self, input_): if self.tp_size > 1: # Build the mask. masked_input, input_mask = get_masked_input_and_mask( input_, self.shard_indices.org_vocab_start_index, self.shard_indices.org_vocab_end_index, self.shard_indices.num_org_vocab_padding, self.shard_indices.added_vocab_start_index, self.shard_indices.added_vocab_end_index, ) else: masked_input = input_ # Get the embeddings. output_parallel = self.quant_method.embedding(self, masked_input.long()) # Mask the output embedding. if self.tp_size > 1: output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) # Reduce across all the model parallel GPUs. output = tensor_model_parallel_all_reduce( output_parallel, tp_group=self.tp_group ) return output def extra_repr(self) -> str: s = f"num_embeddings={self.num_embeddings_per_partition}" s += f", embedding_dim={self.embedding_dim}" s += f", org_vocab_size={self.org_vocab_size}" s += f", num_embeddings_padded={self.num_embeddings_padded}" s += f", tp_size={self.tp_size}" return s ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py ================================================ from safetensors.torch import load_file as safetensors_load_file from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( ComponentLoader, ) from sglang.multimodal_gen.runtime.loader.utils import ( _list_safetensors_files, set_default_torch_dtype, skip_init_modules, ) from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( get_diffusers_component_config, ) from sglang.multimodal_gen.utils import PRECISION_TO_TYPE class AdapterLoader(ComponentLoader): """Loader for small adapter-style modules (e.g., LTX-2 connectors). This loader intentionally avoids FSDP sharding and just: 1) Instantiates the module from `config.json`. 2) Loads a single safetensors state_dict. """ component_names = ["connectors"] expected_library = "diffusers" def load_customized( self, component_model_path: str, server_args: ServerArgs, *args ): config = get_diffusers_component_config(component_path=component_model_path) cls_name = config.pop("_class_name", None) if cls_name is None: raise ValueError( "Model config does not contain a _class_name attribute. " "Only diffusers format is supported." ) config.pop("_diffusers_version", None) config.pop("_name_or_path", None) server_args.model_paths["connectors"] = component_model_path model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) target_device = get_local_torch_device() default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] from types import SimpleNamespace with set_default_torch_dtype(default_dtype), skip_init_modules(): connector_cfg = SimpleNamespace(**config) model = model_cls(connector_cfg).to( device=target_device, dtype=default_dtype ) safetensors_list = _list_safetensors_files(component_model_path) if not safetensors_list: raise ValueError(f"No safetensors files found in {component_model_path}") if len(safetensors_list) != 1: raise ValueError( f"Found {len(safetensors_list)} safetensors files in {component_model_path}, expected 1" ) loaded = safetensors_load_file(safetensors_list[0]) model.load_state_dict(loaded, strict=False) return model ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py ================================================ from copy import deepcopy import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( ComponentLoader, ) from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( get_diffusers_component_config, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) class BridgeLoader(ComponentLoader): """Loader for MOVA dual tower bridge with FSDP support.""" pipeline_bridge_config_attr: str = "bridge_config" component_names = ["dual_tower_bridge"] expected_library = "diffusers" def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ): config = get_diffusers_component_config(component_path=component_model_path) hf_config = deepcopy(config) class_name = config.pop("_class_name", None) if class_name is None: raise ValueError( "Model config does not contain a _class_name attribute. " "Only diffusers format is supported." ) server_args.model_paths[component_name] = component_model_path # Try to get bridge config from pipeline config, fallback to creating one bridge_config = getattr( server_args.pipeline_config, self.pipeline_bridge_config_attr, None ) if bridge_config is not None: bridge_config.update_model_arch(config) else: # Create a minimal config from hf_config from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import ( MOVADualTowerConfig, ) bridge_config = MOVADualTowerConfig() bridge_config.update_model_arch(config) model_cls, _ = ModelRegistry.resolve_model_cls(class_name) # Find all safetensors files safetensors_list = _list_safetensors_files(component_model_path) if not safetensors_list: raise ValueError(f"No safetensors files found in {component_model_path}") default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] logger.info( "Loading %s from %s safetensors files, default_dtype: %s", class_name, len(safetensors_list), default_dtype, ) # Check if FSDP loading is available if ( server_args.hsdp_shard_dim is not None and hasattr(model_cls, "_fsdp_shard_conditions") and model_cls._fsdp_shard_conditions ): # Load with FSDP support model = maybe_load_fsdp_model( model_cls=model_cls, init_params={"config": bridge_config, "hf_config": hf_config}, weight_dir_list=safetensors_list, device=get_local_torch_device(), hsdp_replicate_dim=server_args.hsdp_replicate_dim, hsdp_shard_dim=server_args.hsdp_shard_dim, cpu_offload=server_args.dit_cpu_offload, pin_cpu_memory=server_args.pin_cpu_memory, fsdp_inference=server_args.use_fsdp_inference, param_dtype=default_dtype, reduce_dtype=torch.float32, output_dtype=None, strict=False, ) else: # Fallback to simple loading (for non-FSDP or legacy models) model = model_cls.from_pretrained( component_model_path, torch_dtype=default_dtype ) model = model.to(device=get_local_torch_device(), dtype=default_dtype) total_params = sum(p.numel() for p in model.parameters()) logger.info("Loaded bridge model with %.2fM parameters", total_params / 1e6) return model ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import importlib import os import pkgutil import traceback from abc import ABC from typing import Any, Type import torch from diffusers import AutoModel from torch import nn from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer from sglang.multimodal_gen.configs.models import ModelConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.utils import ( _normalize_component_type, component_name_to_loader_cls, get_memory_usage_of_component, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import get_hf_config from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class ComponentLoader(ABC): """Base class for loading a specific type of model component.""" # the list of possible name of the component in model_index.json, e.g., scheduler component_names: list[str] = [] # diffusers or transformers expected_library: str = "" _loaders_registered = False def __init_subclass__(cls, **kwargs): """ register loaders, called when subclass is imported """ super().__init_subclass__(**kwargs) for component_name in cls.component_names: component_name_to_loader_cls[component_name] = cls def __init__(self, device=None) -> None: self.device = device def should_offload( self, server_args: ServerArgs, model_config: ModelConfig | None = None ): # not offload by default return False def target_device(self, should_offload): if should_offload: return ( torch.device("mps") if current_platform.is_mps() else torch.device("cpu") ) else: return get_local_torch_device() def load( self, component_model_path: str, server_args: ServerArgs, component_name: str, transformers_or_diffusers: str, ) -> tuple[AutoModel, float]: """ Template method that standardizes logging around the core load implementation. The priority of loading method is: 1. load customized component 2. load native diffusers/transformers component If all of the above methods failed, an error will be thrown """ gpu_mem_before_loading = current_platform.get_available_gpu_memory() logger.info( "Loading %s from %s. avail mem: %.2f GB", component_name, component_model_path, gpu_mem_before_loading, ) try: component = self.load_customized( component_model_path, server_args, component_name ) source = "sgl-diffusion" except Exception as e: if "Unsupported model architecture" in str(e): logger.info( f"Component: {component_name} doesn't have a customized version yet, using native version" ) else: traceback.print_exc() logger.error( f"Error while loading customized {component_name}, falling back to native version" ) # fallback to native version component = self.load_native( component_model_path, server_args, transformers_or_diffusers ) should_offload = self.should_offload(server_args) target_device = self.target_device(should_offload) component = component.to(device=target_device) source = "native" logger.warning( "Native component %s: %s is loaded, performance may be sub-optimal", component_name, component.__class__.__name__, ) if component is None: logger.error("Load %s failed", component_name) consumed = 0.0 else: if isinstance(component, nn.Module): component = component.eval() current_gpu_mem = current_platform.get_available_gpu_memory() model_size = get_memory_usage_of_component(component) or "NA" consumed = gpu_mem_before_loading - current_gpu_mem logger.info( f"Loaded %s: %s ({source} version). model size: %s GB, consumed GPU mem: %.2f GB, avail GPU mem: %.2f GB", component_name, component.__class__.__name__, model_size, consumed, current_gpu_mem, ) return component, consumed def load_native( self, component_model_path: str, server_args: ServerArgs, transformers_or_diffusers: str, ) -> AutoModel: """ Load the component using the native library (transformers/diffusers). """ if transformers_or_diffusers == "transformers": from transformers import AutoModel config = get_hf_config( component_model_path, trust_remote_code=server_args.trust_remote_code, revision=server_args.revision, ) return AutoModel.from_pretrained( component_model_path, config=config, trust_remote_code=server_args.trust_remote_code, revision=server_args.revision, ) elif transformers_or_diffusers == "diffusers": from diffusers import AutoModel return AutoModel.from_pretrained( component_model_path, revision=server_args.revision, trust_remote_code=server_args.trust_remote_code, ) else: raise ValueError(f"Unsupported library: {transformers_or_diffusers}") def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ): """ Load the customized version component, implemented and optimized in SGL-diffusion """ raise NotImplementedError( f"load_customized not implemented for {self.__class__.__name__}" ) @classmethod def _ensure_loaders_registered(cls): """ avoid multiple registration """ if cls._loaders_registered: return package_dir = os.path.dirname(__file__) package_name = ( __package__ or "sglang.multimodal_gen.runtime.loader.component_loaders.component_loaders" ) for _, name, _ in pkgutil.iter_modules([package_dir]): # skip importing self to avoid circular dependency issues if name == "component_loader": continue try: importlib.import_module(f".{name}", package=package_name) except ImportError as e: logger.warning(f"Failed to import loader component {name}: {e}") cls._loaders_registered = True @classmethod def for_component_type( cls, component_name: str, transformers_or_diffusers: str ) -> "ComponentLoader": """ Factory method to create a component loader for a specific component type. Args: component_name: Type of component (e.g., "vae", "text_encoder", "transformer", "scheduler") transformers_or_diffusers: Whether the component is from transformers or diffusers """ cls._ensure_loaders_registered() # Map of component types to their loader classes and expected library component_name = _normalize_component_type(component_name) # NOTE(FlamingoPg): special for LTX-2 models if component_name == "vocoder" or component_name == "connectors": transformers_or_diffusers = "diffusers" # NOTE(CloudRipple): special for MOVA models # TODO(CloudRipple): remove most of these special cases after unifying the loading logic if component_name in [ "audio_vae", "audio_dit", "dual_tower_bridge", "video_dit", ]: transformers_or_diffusers = "diffusers" if ( component_name == "scheduler" and transformers_or_diffusers == "mova.diffusion.schedulers.flow_match_pair" ): transformers_or_diffusers = "diffusers" if component_name in component_name_to_loader_cls: loader_cls: Type[ComponentLoader] = component_name_to_loader_cls[ component_name ] expected_library = loader_cls.expected_library # Assert that the library matches what's expected for this component type assert ( transformers_or_diffusers == expected_library ), f"{component_name} must be loaded from {expected_library}, got {transformers_or_diffusers}" return loader_cls() # For unknown component types, use a generic loader logger.warning( "No specific loader found for component type: %s. Using generic loader.", component_name, ) return GenericComponentLoader(transformers_or_diffusers) class ImageProcessorLoader(ComponentLoader): """Loader for image processor.""" component_names = ["image_processor"] expected_library = "transformers" def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ) -> Any: return AutoImageProcessor.from_pretrained(component_model_path, use_fast=True) class AutoProcessorLoader(ComponentLoader): """Loader for auto processor.""" component_names = ["processor"] expected_library = "transformers" def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ) -> Any: return AutoProcessor.from_pretrained(component_model_path) class TokenizerLoader(ComponentLoader): """Loader for tokenizers.""" component_names = ["tokenizer"] expected_library = "transformers" def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ) -> Any: return AutoTokenizer.from_pretrained( component_model_path, padding_size="right", ) class GenericComponentLoader(ComponentLoader): """Generic loader for components that don't have a specific loader.""" def __init__(self, library="transformers") -> None: super().__init__() self.library = library class PipelineComponentLoader: """ Utility class for loading the components in a pipeline. """ @staticmethod def load_component( component_name: str, component_model_path: str, transformers_or_diffusers: str, server_args: ServerArgs, ): """ Load a pipeline component. Args: component_name: Name of the component (e.g., "vae", "text_encoder", "transformer", "scheduler") component_model_path: Path to the component model transformers_or_diffusers: Whether the component is from transformers or diffusers """ # Get the appropriate loader for this component type loader = ComponentLoader.for_component_type( component_name, transformers_or_diffusers ) try: # Load the component return loader.load( component_model_path, server_args, component_name, transformers_or_diffusers, ) except Exception as e: logger.error( f"Error while loading component: {component_name}, {component_model_path=}" ) raise e ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py ================================================ from sglang.multimodal_gen.configs.models import ModelConfig from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader import ( TextEncoderLoader, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( get_diffusers_component_config, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class ImageEncoderLoader(TextEncoderLoader): component_names = ["image_encoder"] expected_library = "transformers" def should_offload(self, server_args, model_config: ModelConfig | None = None): should_offload = server_args.image_encoder_cpu_offload if not should_offload: return False # _fsdp_shard_conditions is in arch_config, not directly on model_config arch_config = ( getattr(model_config, "arch_config", model_config) if model_config else None ) fsdp_shard_conditions = ( getattr(arch_config, "_fsdp_shard_conditions", []) if arch_config else [] ) use_cpu_offload = should_offload and len(fsdp_shard_conditions) > 0 return use_cpu_offload def load_customized( self, component_model_path: str, server_args: ServerArgs, *args ): """Load the text encoders based on the model path, and inference args.""" # model_config: PretrainedConfig = get_hf_config( # model=model_path, # trust_remote_code=server_args.trust_remote_code, # revision=server_args.revision, # model_override_args=None, # ) model_config = get_diffusers_component_config( component_path=component_model_path ) encoder_config = server_args.pipeline_config.image_encoder_config encoder_config.update_model_arch(model_config) # Always start with local device; load_model will adjust for offload if needed # TODO(will): add support for other dtypes return self.load_model( component_model_path, encoder_config, server_args, server_args.pipeline_config.image_encoder_precision, cpu_offload_flag=server_args.image_encoder_cpu_offload, ) ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py ================================================ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( ComponentLoader, ) from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( get_diffusers_component_config, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class SchedulerLoader(ComponentLoader): """Loader for scheduler.""" component_names = ["scheduler"] expected_library = "diffusers" def load_customized( self, component_model_path: str, server_args: ServerArgs, *args ): """Load the scheduler based on the model path, and inference args.""" config = get_diffusers_component_config(component_path=component_model_path) class_name = config.pop("_class_name") assert ( class_name is not None ), "Model config does not contain a _class_name attribute. Only diffusers format is supported." scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name) scheduler = scheduler_cls(**config) if server_args.pipeline_config.flow_shift is not None: scheduler.set_shift(server_args.pipeline_config.flow_shift) return scheduler ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py ================================================ import dataclasses import glob import os from collections.abc import Generator, Iterable from typing import Generator, Iterable, cast import torch import torch.distributed as dist import torch.nn as nn from torch import nn from torch.distributed import init_device_mesh from transformers.utils import SAFE_WEIGHTS_INDEX_NAME from sglang.multimodal_gen.configs.models import EncoderConfig, ModelConfig from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImageEditPipelineConfig, ) from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( ComponentLoader, ) from sglang.multimodal_gen.runtime.loader.fsdp_load import shard_model from sglang.multimodal_gen.runtime.loader.utils import ( set_default_torch_dtype, skip_init_modules, ) from sglang.multimodal_gen.runtime.loader.weight_utils import ( filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, pt_weights_iterator, safetensors_weights_iterator, ) from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( get_config, get_diffusers_component_config, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE from sglang.srt.environ import envs logger = init_logger(__name__) class TextEncoderLoader(ComponentLoader): """Loader for text encoders.""" component_names = ["text_encoder"] expected_library = "transformers" @dataclasses.dataclass class Source: """A source for weights.""" model_or_path: str """The model ID or path.""" prefix: str = "" """A prefix to prepend to all weights.""" fall_back_to_pt: bool = True """Whether .pt weights can be used.""" allow_patterns_overrides: list[str] | None = None """If defined, weights will load exclusively using these patterns.""" def should_offload(self, server_args, model_config: ModelConfig | None = None): should_offload = server_args.text_encoder_cpu_offload if not should_offload: return False # _fsdp_shard_conditions is in arch_config, not directly on model_config arch_config = ( getattr(model_config, "arch_config", model_config) if model_config else None ) fsdp_shard_conditions = ( getattr(arch_config, "_fsdp_shard_conditions", []) if arch_config else [] ) use_cpu_offload = should_offload and len(fsdp_shard_conditions) > 0 return use_cpu_offload def _prepare_weights( self, model_name_or_path: str, fall_back_to_pt: bool, allow_patterns_overrides: list[str] | None, ) -> tuple[str, list[str], bool]: """Prepare weights for the model. If the model is not local, it will be downloaded.""" # model_name_or_path = (self._maybe_download_from_modelscope( # model_name_or_path, revision) or model_name_or_path) is_local = os.path.isdir(model_name_or_path) assert is_local, "Model path must be a local directory" use_safetensors = False index_file = SAFE_WEIGHTS_INDEX_NAME allow_patterns = ["*.safetensors", "*.bin"] if fall_back_to_pt: allow_patterns += ["*.pt"] if allow_patterns_overrides is not None: allow_patterns = allow_patterns_overrides hf_folder = model_name_or_path hf_weights_files: list[str] = [] for pattern in allow_patterns: hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) if len(hf_weights_files) > 0: if pattern == "*.safetensors": use_safetensors = True break if use_safetensors: hf_weights_files = filter_duplicate_safetensors_files( hf_weights_files, hf_folder, index_file ) else: hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files) if len(hf_weights_files) == 0: raise RuntimeError( f"Cannot find any model weights with `{model_name_or_path}`" ) if envs.SGLANG_SORT_WEIGHT_FILES.get(): hf_weights_files.sort() return hf_folder, hf_weights_files, use_safetensors def _get_weights_iterator( self, source: "Source", to_cpu: bool ) -> Generator[tuple[str, torch.Tensor], None, None]: """get an iterator for the model weights based on the load format.""" hf_folder, hf_weights_files, use_safetensors = self._prepare_weights( source.model_or_path, source.fall_back_to_pt, source.allow_patterns_overrides, ) if use_safetensors: weights_iterator = safetensors_weights_iterator( hf_weights_files, to_cpu=to_cpu ) else: weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu=to_cpu) # apply the prefix. return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) def _get_all_weights( self, model: nn.Module, model_path: str, to_cpu: bool, ) -> Generator[tuple[str, torch.Tensor], None, None]: primary_weights = TextEncoderLoader.Source( model_path, prefix="", fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True), allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None), ) yield from self._get_weights_iterator(primary_weights, to_cpu) secondary_weights = cast( Iterable[TextEncoderLoader.Source], getattr(model, "secondary_weights", ()), ) for source in secondary_weights: yield from self._get_weights_iterator(source, to_cpu) def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ): """Load the text encoders based on the model path, and inference args.""" diffusers_pretrained_config = get_config( component_model_path, trust_remote_code=True ) model_config = get_diffusers_component_config( component_path=component_model_path ) def is_not_first_encoder(module_name): return "2" in module_name # TODO(mick): had to throw an exception for different text-encoder arch if not is_not_first_encoder(component_name): encoder_config = server_args.pipeline_config.text_encoder_configs[0] encoder_config.update_model_arch(model_config) for key, value in diffusers_pretrained_config.__dict__.items(): setattr(encoder_config.arch_config, key, value) encoder_dtype = server_args.pipeline_config.text_encoder_precisions[0] else: assert len(server_args.pipeline_config.text_encoder_configs) == 2 encoder_config = server_args.pipeline_config.text_encoder_configs[1] encoder_config.update_model_arch(model_config) encoder_dtype = server_args.pipeline_config.text_encoder_precisions[1] # TODO(will): add support for other dtypes return self.load_model( component_model_path, encoder_config, server_args, encoder_dtype, ) def load_model( self, model_path: str, model_config: EncoderConfig, server_args: ServerArgs, dtype: str = "fp16", cpu_offload_flag: bool | None = None, ): # Determine CPU offload behavior and target device local_torch_device = get_local_torch_device() should_offload = self.should_offload(server_args, model_config) if should_offload and not current_platform.is_mps(): model_device = torch.device("cpu") else: model_device = local_torch_device with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]): with model_device, skip_init_modules(): architectures = getattr(model_config, "architectures", []) model_cls, _ = ModelRegistry.resolve_model_cls(architectures) enable_image_understanding = ( True if isinstance( server_args.pipeline_config, QwenImageEditPipelineConfig ) else False ) model_config.enable_image_understanding = enable_image_understanding model = model_cls(model_config) weights_to_load = {name for name, _ in model.named_parameters()} loaded_weights = model.load_weights( self._get_all_weights(model, model_path, to_cpu=should_offload) ) # Explicitly move model to target device after loading weights if not should_offload: model = model.to(local_torch_device) if should_offload: # Disable FSDP for MPS as it's not compatible if current_platform.is_mps(): logger.info( "Disabling FSDP sharding for MPS platform as it's not compatible" ) model = model.to(local_torch_device) else: mesh = init_device_mesh( current_platform.device_type, mesh_shape=(1, dist.get_world_size()), mesh_dim_names=("offload", "replicate"), ) shard_model( model, cpu_offload=True, reshard_after_forward=True, mesh=mesh["offload"], fsdp_shard_conditions=model_config.arch_config._fsdp_shard_conditions or getattr(model, "_fsdp_shard_conditions", None), pin_cpu_memory=server_args.pin_cpu_memory, ) else: model = model.to(local_torch_device) # We only enable strict check for non-quantized models # that have loaded weights tracking currently. # if loaded_weights is not None: weights_not_loaded = weights_to_load - loaded_weights if weights_not_loaded: # NOTE: # If we silently continue with uninitialized weights, the text encoder can # produce NaNs/garbage embeddings that later fail stage verification in a # hard-to-debug way (e.g., `prompt_embeds` fails the NaN check). # # We allow a small set of known-optional parameters to be missing, but # default to strict behavior for the rest. allowed_missing_patterns = ( getattr(model, "_allowed_missing_weights_patterns", []) or [] ) unexpected_missing = { n for n in weights_not_loaded if not any(pat in n for pat in allowed_missing_patterns) } if unexpected_missing: raise ValueError( "Following text encoder weights were not initialized from checkpoint: " f"{sorted(unexpected_missing)}. " "This usually indicates a checkpoint/model-arch mismatch or a broken " "weight-name mapping. If these are truly optional, set " "`model._allowed_missing_weights_patterns` to whitelist patterns." ) logger.warning( "Following (allowed) text encoder weights were not initialized from " "checkpoint: %s (allowed patterns: %s)", sorted(weights_not_loaded), allowed_missing_patterns, ) return model ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py ================================================ import json import logging import os from typing import Any, Dict, List, Optional import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( NunchakuConfig, _patch_nunchaku_scales, ) from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( ComponentLoader, ) from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model from sglang.multimodal_gen.runtime.loader.utils import ( _list_safetensors_files, _normalize_component_type, ) from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( get_diffusers_component_config, maybe_download_model, ) from sglang.multimodal_gen.runtime.utils.logging_utils import get_log_level, init_logger from sglang.multimodal_gen.runtime.utils.quantization_utils import ( get_metadata_from_safetensors_file, get_quant_config, get_quant_config_from_safetensors_metadata, ) from sglang.multimodal_gen.utils import PRECISION_TO_TYPE from sglang.srt.utils import is_npu _is_npu = is_npu() logger = init_logger(__name__) class TransformerLoader(ComponentLoader): """Shared loader for (video/audio) DiT transformers.""" component_names = ["transformer", "audio_dit", "video_dit"] expected_library = "diffusers" def get_list_of_safetensors_to_load( self, server_args: ServerArgs, component_model_path: str ) -> list[str]: """ get list of safetensors to load. If --transformer-weights-path is provided, load weights from that path instead of the base model's component directory. """ quantized_path = server_args.transformer_weights_path if quantized_path: quantized_path = maybe_download_model(quantized_path) logger.info("using quantized transformer weights from: %s", quantized_path) if os.path.isfile(quantized_path) and quantized_path.endswith( ".safetensors" ): safetensors_list = [quantized_path] else: safetensors_list = _list_safetensors_files(quantized_path) else: safetensors_list = _list_safetensors_files(component_model_path) if not safetensors_list: raise ValueError( f"no safetensors files found in " f"{quantized_path or component_model_path}" ) return safetensors_list def _resolve_quant_config( self, hf_config: Dict[str, List[str]], server_args: ServerArgs, safetensors_list: list[str], component_model_path: str, ) -> Optional[dict]: # priority: model config.json → safetensors metadata → nunchaku config quant_config = get_quant_config(hf_config, component_model_path) if quant_config is None and server_args.transformer_weights_path: # try to read quantization_config from the safetensors metadata header for safetensors_file in safetensors_list: quant_config = get_quant_config_from_safetensors_metadata( safetensors_file ) if quant_config: break return quant_config def _resolve_target_param_dtype( self, quant_config: Optional[dict], nunchaku_config: Optional[NunchakuConfig], model_cls, server_args: ServerArgs, ) -> Optional[torch.dtype]: if quant_config is not None or nunchaku_config is not None: # TODO: improve the condition # respect dtype from checkpoint param_dtype = None else: param_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] if nunchaku_config is not None: nunchaku_config.model_cls = model_cls # verify that the nunchaku checkpoint matches the selected model class original_dit_cls_name = json.loads( get_metadata_from_safetensors_file( nunchaku_config.transformer_weights_path ).get("config") )["_class_name"] specified_dit_cls_name = str(model_cls.__name__) if original_dit_cls_name != specified_dit_cls_name: raise Exception( f"Class name of DiT specified in nunchaku transformer_weights_path: {original_dit_cls_name} does not match that of specified DiT name: {specified_dit_cls_name}" ) return param_dtype def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ): """Load the transformer based on the model path, and inference args.""" # 1. hf config config = get_diffusers_component_config(component_path=component_model_path) # 2. quant config safetensors_list = self.get_list_of_safetensors_to_load( server_args, component_model_path ) quant_config = self._resolve_quant_config( config, server_args, safetensors_list, component_model_path ) # 3. dit config # Config from Diffusers supersedes sgl_diffusion's model config component_name = _normalize_component_type(component_name) server_args.model_paths[component_name] = component_model_path if component_name in ("transformer", "video_dit"): pipeline_dit_config_attr = "dit_config" elif component_name in ("audio_dit",): pipeline_dit_config_attr = "audio_dit_config" else: raise ValueError(f"Invalid module name: {component_name}") dit_config = getattr(server_args.pipeline_config, pipeline_dit_config_attr) dit_config.update_model_arch(config) cls_name = config.pop("_class_name") model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) nunchaku_config = server_args.nunchaku_config param_dtype = self._resolve_target_param_dtype( quant_config, nunchaku_config, model_cls, server_args ) logger.info( "Loading %s from %s safetensors file(s) %s, param_dtype: %s", cls_name, len(safetensors_list), f": {safetensors_list}" if get_log_level() == logging.DEBUG else "", param_dtype, ) # prepare init_param init_params: dict[str, Any] = { "config": dit_config, "hf_config": config, "quant_config": (quant_config if quant_config else nunchaku_config), } if ( init_params["quant_config"] is None and server_args.transformer_weights_path is not None ): logger.warning( f"transformer_weights_path provided, but quantization config not resolved, which is unexpected and likely to cause errors" ) else: logger.debug("quantization config: %s", init_params["quant_config"]) # Load the model using FSDP loader model = maybe_load_fsdp_model( model_cls=model_cls, init_params=init_params, weight_dir_list=safetensors_list, device=get_local_torch_device(), hsdp_replicate_dim=server_args.hsdp_replicate_dim, hsdp_shard_dim=server_args.hsdp_shard_dim, cpu_offload=server_args.dit_cpu_offload, pin_cpu_memory=server_args.pin_cpu_memory, fsdp_inference=server_args.use_fsdp_inference, # TODO(will): make these configurable param_dtype=param_dtype, reduce_dtype=torch.float32, output_dtype=None, strict=False, ) if nunchaku_config is not None: _patch_nunchaku_scales(model, safetensors_list) total_params = sum(p.numel() for p in model.parameters()) logger.info("Loaded model with %.2fB parameters", total_params / 1e9) # considering the existent of mixed-precision models (e.g., nunchaku) if next(model.parameters()).dtype != param_dtype and param_dtype: logger.warning( f"Model dtype does not match expected param dtype, {next(model.parameters()).dtype} vs {param_dtype}" ) return model ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py ================================================ import importlib.util import os import torch import torch.nn as nn from safetensors.torch import load_file as safetensors_load_file from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.models import ModelConfig from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( ComponentLoader, ) from sglang.multimodal_gen.runtime.loader.utils import ( _list_safetensors_files, set_default_torch_dtype, skip_init_modules, ) from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( get_diffusers_component_config, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) def _convert_conv3d_weights_to_channels_last_3d(module: nn.Module) -> int: """ Convert Conv3d weights to channels_last_3d (NDHWC) memory format. Returns the number of Conv3d modules converted. """ if not hasattr(torch, "channels_last_3d"): return 0 num_converted = 0 for m in module.modules(): if isinstance(m, nn.Conv3d): try: m.weight.data = m.weight.data.to(memory_format=torch.channels_last_3d) num_converted += 1 except Exception: # Best-effort; skip unsupported cases. continue return num_converted class VAELoader(ComponentLoader): """Shared loader for (video/audio) VAE modules.""" component_names = ["vae", "audio_vae", "video_vae"] expected_library = "diffusers" def should_offload( self, server_args: ServerArgs, model_config: ModelConfig | None = None ): return server_args.vae_cpu_offload def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ): """Load the VAE based on the model path, and inference args.""" config = get_diffusers_component_config(component_path=component_model_path) class_name = config.pop("_class_name", None) assert ( class_name is not None ), "Model config does not contain a _class_name attribute. Only diffusers format is supported." server_args.model_paths[component_name] = component_model_path if component_name in ("vae", "video_vae"): pipeline_vae_config_attr = "vae_config" pipeline_vae_precision = "vae_precision" elif component_name in ("audio_vae",): pipeline_vae_config_attr = "audio_vae_config" pipeline_vae_precision = "audio_vae_precision" else: raise ValueError( f"Unsupported module name for VAE loader: {component_name}" ) vae_config = getattr(server_args.pipeline_config, pipeline_vae_config_attr) vae_precision = getattr(server_args.pipeline_config, pipeline_vae_precision) vae_config.update_model_arch(config) if hasattr(vae_config, "post_init"): # NOTE: some post init logics are only available after updated with config vae_config.post_init() should_offload = self.should_offload(server_args) target_device = self.target_device(should_offload) # Check for auto_map first (custom VAE classes) auto_map = config.get("auto_map", {}) auto_model_map = auto_map.get("AutoModel") if auto_model_map: module_path, cls_name = auto_model_map.rsplit(".", 1) custom_module_file = os.path.join(component_model_path, f"{module_path}.py") spec = importlib.util.spec_from_file_location("_custom", custom_module_file) custom_module = importlib.util.module_from_spec(spec) spec.loader.exec_module(custom_module) vae_cls = getattr(custom_module, cls_name) vae_dtype = PRECISION_TO_TYPE[vae_precision] with set_default_torch_dtype(vae_dtype): vae = vae_cls.from_pretrained( component_model_path, revision=server_args.revision, trust_remote_code=server_args.trust_remote_code, ) vae = vae.to(device=target_device, dtype=vae_dtype) if ( component_name in ("vae", "video_vae") and torch.cuda.is_available() and getattr(envs, "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D", False) ): n = _convert_conv3d_weights_to_channels_last_3d(vae) if n > 0: logger.info( "VAE: converted %d Conv3d weights to channels_last_3d", n ) vae = current_platform.optimize_vae(vae) return vae # Load from ModelRegistry (standard VAE classes) with ( set_default_torch_dtype(PRECISION_TO_TYPE[vae_precision]), skip_init_modules(), ): vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) vae = vae_cls(vae_config).to(target_device) safetensors_list = _list_safetensors_files(component_model_path) assert ( len(safetensors_list) >= 1 ), f"Found no safetensors files in {component_model_path}" loaded = {} for sf_path in safetensors_list: loaded.update(safetensors_load_file(sf_path)) vae.load_state_dict(loaded, strict=False) state_keys = set(vae.state_dict().keys()) loaded_keys = set(loaded.keys()) missing_keys = sorted(state_keys - loaded_keys) unexpected_keys = sorted(loaded_keys - state_keys) if missing_keys: logger.warning("VAE missing keys: %s", missing_keys) if unexpected_keys: logger.warning("VAE unexpected keys: %s", unexpected_keys) if ( component_name in ("vae", "video_vae") and torch.cuda.is_available() and getattr(envs, "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D", False) ): n = _convert_conv3d_weights_to_channels_last_3d(vae) if n > 0: logger.info("VAE: converted %d Conv3d weights to channels_last_3d", n) vae = current_platform.optimize_vae(vae) return vae ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/component_loaders/vl_encoder_loader.py ================================================ from typing import Any from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( ComponentLoader, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import get_hf_config class VisionLanguageEncoderLoader(ComponentLoader): """Loader for vision language encoder (typically Causal LM or Vision2Seq).""" component_names = ["vision_language_encoder"] expected_library = "transformers" def load_customized( self, component_model_path: str, server_args: ServerArgs, transformers_or_diffusers: str = "vision_language_encoder", ) -> Any: if transformers_or_diffusers == "vision_language_encoder": from transformers import GlmImageForConditionalGeneration config = get_hf_config( component_model_path, trust_remote_code=server_args.trust_remote_code, revision=server_args.revision, ) model = GlmImageForConditionalGeneration.from_pretrained( component_model_path, config=config, trust_remote_code=server_args.trust_remote_code, revision=server_args.revision, ).to(get_local_torch_device()) return model else: raise ValueError( f"Unsupported library for VisionLanguageEncoder: {transformers_or_diffusers}" ) ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py ================================================ from safetensors.torch import load_file as safetensors_load_file from sglang.multimodal_gen.configs.models import ModelConfig from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( ComponentLoader, ) from sglang.multimodal_gen.runtime.loader.utils import ( _list_safetensors_files, set_default_torch_dtype, skip_init_modules, ) from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( get_diffusers_component_config, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) class VocoderLoader(ComponentLoader): component_names = ["vocoder"] expected_library = "diffusers" def should_offload( self, server_args: ServerArgs, model_config: ModelConfig | None = None ): return server_args.vae_cpu_offload def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ): config = get_diffusers_component_config(component_path=component_model_path) class_name = config.pop("_class_name", None) assert ( class_name is not None ), "Model config does not contain a _class_name attribute. Only diffusers format is supported." server_args.model_paths[component_name] = component_model_path from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import ( LTXVocoderConfig, ) vocoder_config = LTXVocoderConfig() vocoder_config.update_model_arch(config) try: vocoder_precision = server_args.pipeline_config.audio_vae_precision except AttributeError: vocoder_precision = "fp32" vocoder_dtype = PRECISION_TO_TYPE[vocoder_precision] should_offload = self.should_offload(server_args) target_device = self.target_device(should_offload) with set_default_torch_dtype(vocoder_dtype), skip_init_modules(): vocoder_cls, _ = ModelRegistry.resolve_model_cls(class_name) vocoder = vocoder_cls(vocoder_config).to(target_device) safetensors_list = _list_safetensors_files(component_model_path) assert ( len(safetensors_list) == 1 ), f"Found {len(safetensors_list)} safetensors files in {component_model_path}" loaded = safetensors_load_file(safetensors_list[0]) incompatible = vocoder.load_state_dict(loaded, strict=False) missing_keys = [] unexpected_keys = [] try: missing_keys = incompatible.missing_keys unexpected_keys = incompatible.unexpected_keys except AttributeError: # Best-effort fallback in case older torch returns a tuple-like. try: missing_keys = incompatible[0] unexpected_keys = incompatible[1] except Exception: pass if missing_keys or unexpected_keys: logger.warning( "Loaded vocoder with missing_keys=%d unexpected_keys=%d", len(missing_keys), len(unexpected_keys), ) return vocoder ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/fsdp_load.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from torchtune # Copyright 2024 The TorchTune Authors. # Copyright 2025 The sglang-diffusion Authors. from collections.abc import Callable, Generator from itertools import chain from typing import Any import torch from torch import nn from torch.distributed import DeviceMesh, init_device_mesh from torch.distributed._tensor import distribute_tensor from torch.distributed.fsdp import ( CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, fully_shard, ) from torch.nn.modules.module import _IncompatibleKeys from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod from sglang.multimodal_gen.runtime.loader.utils import ( get_param_names_mapping, hf_to_custom_state_dict, set_default_torch_dtype, ) from sglang.multimodal_gen.runtime.loader.weight_utils import ( safetensors_weights_iterator, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import set_mixed_precision_policy from sglang.srt.utils import is_npu _is_npu = is_npu() logger = init_logger(__name__) def _make_param_like( actual_param: torch.nn.Parameter, tensor: torch.Tensor ) -> torch.nn.Parameter: cls = actual_param.__class__ # nn.Parameter defaults to requires_grad=True, which is illegal for non-floating/complex dtypes (e.g., int8/FP8 # quantized weights). try: new_param = cls.__new__(cls, tensor, requires_grad=False) except TypeError: new_param = cls.__new__(cls, tensor) new_param.__dict__.update(actual_param.__dict__) new_param.requires_grad = False return new_param # TODO(PY): add compile option def maybe_load_fsdp_model( model_cls: type[nn.Module], init_params: dict[str, Any], weight_dir_list: list[str], device: torch.device, hsdp_replicate_dim: int, hsdp_shard_dim: int, param_dtype: torch.dtype, reduce_dtype: torch.dtype, cpu_offload: bool = False, fsdp_inference: bool = False, output_dtype: torch.dtype | None = None, pin_cpu_memory: bool = True, strict: bool = True, ) -> torch.nn.Module: """Load a model with optional FSDP (Fully Sharded Data Parallel) support. Args: param_dtype: Data type for model parameters, also used for: - Model initialization context (set_default_torch_dtype) - FSDP mixed precision policy - Weight loading and casting reduce_dtype: Data type for gradient reduction in FSDP mixed precision. strict: If True, enforce strict state dict loading (all keys must match). """ # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are # manually casting the inputs to the model default_torch_dtype = param_dtype if param_dtype else torch.bfloat16 mp_policy = MixedPrecisionPolicy( default_torch_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False ) set_mixed_precision_policy( param_dtype=default_torch_dtype, reduce_dtype=reduce_dtype, output_dtype=output_dtype, mp_policy=mp_policy, ) with set_default_torch_dtype(default_torch_dtype), torch.device("meta"): model = model_cls(**init_params) # Check if we should use FSDP use_fsdp = fsdp_inference # Disable FSDP for MPS as it's not compatible if current_platform.is_mps(): use_fsdp = False logger.info("Disabling FSDP for MPS platform as it's not compatible") if use_fsdp: world_size = hsdp_replicate_dim * hsdp_shard_dim if not fsdp_inference: hsdp_replicate_dim = world_size hsdp_shard_dim = 1 device_mesh = init_device_mesh( current_platform.device_type, # (Replicate(), Shard(dim=0)) mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim), mesh_dim_names=("replicate", "shard"), ) shard_model( model, cpu_offload=cpu_offload, reshard_after_forward=True, mp_policy=mp_policy, mesh=device_mesh, fsdp_shard_conditions=model._fsdp_shard_conditions, pin_cpu_memory=pin_cpu_memory, ) weight_iterator = safetensors_weights_iterator(weight_dir_list) param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping) load_model_from_full_model_state_dict( model, weight_iterator, device, param_dtype, strict=strict, cpu_offload=cpu_offload, param_names_mapping=param_names_mapping_fn, ) for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) if quant_method is not None and hasattr( quant_method, "process_weights_after_loading" ): if _is_npu and not isinstance(quant_method, UnquantizedLinearMethod): # Activate the NZ format for storing weights, # which is a specific optimization for Ascend NPU torch.npu.config.allow_internal_format = True quant_method.process_weights_after_loading(module) if _is_npu: torch.npu.empty_cache() for n, p in chain(model.named_parameters(), model.named_buffers()): if p.is_meta: raise RuntimeError(f"Unexpected param or buffer {n} on meta device.") # Avoid unintended computation graph accumulation during inference if isinstance(p, torch.nn.Parameter): p.requires_grad = False return model def shard_model( model, *, cpu_offload: bool, reshard_after_forward: bool = True, mp_policy: MixedPrecisionPolicy | None = MixedPrecisionPolicy(), # noqa mesh: DeviceMesh | None = None, fsdp_shard_conditions: list[Callable[[str, nn.Module], bool]] = [], # noqa pin_cpu_memory: bool = True, ) -> None: """ Utility to shard a model with FSDP using the PyTorch Distributed fully_shard API. This method will over the model's named modules from the bottom-up and apply shard modules based on whether they meet any of the criteria from shard_conditions. Args: model (TransformerDecoder): Model to shard with FSDP. cpu_offload (bool): If set to True, FSDP will offload parameters, gradients, and optimizer states to CPU. reshard_after_forward (bool): Whether to reshard parameters and buffers after the forward pass. Setting this to True corresponds to the FULL_SHARD sharding strategy from FSDP1, while setting it to False corresponds to the SHARD_GRAD_OP sharding strategy. mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism. Default to None. fsdp_shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine which modules to shard with FSDP. pin_cpu_memory (bool): If set to True, FSDP will pin the CPU memory of the offloaded parameters. """ if fsdp_shard_conditions is None or len(fsdp_shard_conditions) == 0: logger.warning( "The FSDP shard condition list is empty or None. No modules will be sharded in %s", type(model).__name__, ) return fsdp_kwargs = { "reshard_after_forward": reshard_after_forward, "mesh": mesh, "mp_policy": mp_policy, } if cpu_offload: fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(pin_memory=pin_cpu_memory) # iterating in reverse to start with # lowest-level modules first num_layers_sharded = 0 # TODO(will): don't reshard after forward for the last layer to save on the # all-gather that will immediately happen Shard the model with FSDP, for n, m in reversed(list(model.named_modules())): if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]): # type: ignore fully_shard(m, **fsdp_kwargs) num_layers_sharded += 1 if num_layers_sharded == 0: raise ValueError( "No layer modules were sharded. Please check if shard conditions are working as expected." ) # Finally shard the entire model to account for any stragglers fully_shard(model, **fsdp_kwargs) # TODO(PY): device mesh for cfg parallel def load_model_from_full_model_state_dict( model: FSDPModule | torch.nn.Module, full_sd_iterator: Generator[tuple[str, torch.Tensor], None, None], device: torch.device, param_dtype: torch.dtype | None, strict: bool = False, cpu_offload: bool = False, param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None, ) -> _IncompatibleKeys: """ Converting full state dict into a sharded state dict and loading it into FSDP model (if training) or normal huggingface model Args: model (Union[FSDPModule, torch.nn.Module]): Model to generate fully qualified names for cpu_state_dict full_sd_iterator (Generator): an iterator yielding (param_name, tensor) pairs device (torch.device): device used to move full state dict tensors param_dtype (torch.dtype): dtype used to move full state dict tensors. If none, respect original dtype from checkpoint strict (bool): flag to check if to load the model in strict mode cpu_offload (bool): flag to check if FSDP offload is enabled param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name Returns: ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields: * **missing_keys** is a list of str containing the missing keys * **unexpected_keys** is a list of str containing the unexpected keys """ meta_sd = model.state_dict() param_dict = dict(model.named_parameters()) # map names from checkpoint to customized names custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict( full_sd_iterator, param_names_mapping ) # type: ignore is_fsdp_model = isinstance(model, FSDPModule) or any( hasattr(p, "device_mesh") for p in meta_sd.values() ) # sort parameter names to ensure all ranks process parameters in the same order sorted_param_names = sorted(custom_param_sd.keys()) sharded_sd = {} skipped_checkpoint_keys: list[str] = [] # shard from loaded state_dict, custom_param_sd -> sharded_sd for target_param_name in sorted_param_names: full_tensor = custom_param_sd[target_param_name] meta_sharded_param = meta_sd.get(target_param_name) if meta_sharded_param is None: # For FSDP models, ensure all ranks process parameters consistently if strict or is_fsdp_model: raise ValueError( f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect." ) else: skipped_checkpoint_keys.append(target_param_name) continue # use meta param dtype so quantized params (e.g. FP8) keep their dtype; # for non-quantized models meta dtype equals param_dtype anyway if meta_sharded_param is None: # for nunchaku, some scales are patched later target_dtype = full_tensor.dtype else: target_dtype = meta_sharded_param.dtype if not hasattr(meta_sharded_param, "device_mesh"): full_tensor = full_tensor.to(device=device, dtype=target_dtype) actual_param = param_dict.get(target_param_name) weight_loader = ( getattr(actual_param, "weight_loader", None) if actual_param is not None else None ) if weight_loader is not None: assert actual_param is not None sharded_tensor = torch.empty_like( meta_sharded_param, device=device, dtype=target_dtype ) # Preserve requires_grad flag to avoid errors with non-floating dtypes requires_grad = getattr(meta_sharded_param, "requires_grad", False) temp_param = _make_param_like(actual_param, sharded_tensor) if not ( sharded_tensor.is_floating_point() or sharded_tensor.is_complex() ): requires_grad = False temp_param.requires_grad = requires_grad weight_loader(temp_param, full_tensor) sharded_tensor = temp_param.data else: # In cases where parts of the model aren't sharded, some parameters will be plain tensors sharded_tensor = full_tensor # Important: `cpu_offload` is intended for FSDP-managed parameter movement. # If a parameter is not sharded into a DTensor (i.e., no `device_mesh`), FSDP # will NOT manage it. Offloading it here would leave CPU parameters that # later participate in GPU kernels (e.g., conv/embedding), causing device/dtype # mismatches like "Input type (CUDABFloat16Type) and weight type (CPUBFloat16Type)". # # Therefore: # - For non-FSDP models, keep the historical behavior (allow CPU offload). # - For FSDP models, do NOT offload non-sharded parameters here. if cpu_offload and not is_fsdp_model: sharded_tensor = sharded_tensor.cpu() else: full_tensor = full_tensor.to(device=device, dtype=target_dtype) sharded_tensor = distribute_tensor( full_tensor, meta_sharded_param.device_mesh, meta_sharded_param.placements, ) if cpu_offload: sharded_tensor = sharded_tensor.to("cpu") requires_grad = False sharded_sd[target_param_name] = nn.Parameter( sharded_tensor, requires_grad=requires_grad ) model.reverse_param_names_mapping = reverse_param_names_mapping if skipped_checkpoint_keys: logger.warning( "Checkpoint keys not loaded (no matching model parameter) %s", ( skipped_checkpoint_keys[:20] if len(skipped_checkpoint_keys) > 20 else skipped_checkpoint_keys ), ) if len(skipped_checkpoint_keys) > 20: logger.warning( "... and %d more skipped keys.", len(skipped_checkpoint_keys) - 20, ) # parameters in nn.Module that doesn't exist in safetensor files unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys()) if unused_keys: logger.warning("Found unloaded parameters in meta state dict: %s", unused_keys) # for nunchaku; norm_q/norm_k for SANA QK normalization layers ALLOWED_NEW_PARAM_PATTERNS = [ "gate_compress", "wcscales", "wtscale", "bias", "norm_q", "norm_k", ] for new_param_name in unused_keys: if not any(pattern in new_param_name for pattern in ALLOWED_NEW_PARAM_PATTERNS): logger.error( "Unsupported new parameter: %s. Allowed patterns: %s", new_param_name, ALLOWED_NEW_PARAM_PATTERNS, ) raise ValueError( f"New parameter '{new_param_name}' is not supported. " f"Currently only parameters containing {ALLOWED_NEW_PARAM_PATTERNS} are allowed." ) meta_sharded_param = meta_sd.get(new_param_name) meta_sharded_param_dtype = meta_sharded_param.dtype if any( p in new_param_name for p in ("wcscales", "wtscale", "norm_q", "norm_k") ): init_like = torch.ones_like else: init_like = torch.zeros_like if not hasattr(meta_sharded_param, "device_mesh"): sharded_tensor = init_like( meta_sharded_param, device=device, dtype=meta_sharded_param_dtype ) if cpu_offload and not is_fsdp_model: sharded_tensor = sharded_tensor.cpu() else: full_tensor = init_like( meta_sharded_param, device=device, dtype=meta_sharded_param_dtype ) sharded_tensor = distribute_tensor( full_tensor, meta_sharded_param.device_mesh, meta_sharded_param.placements, ) if cpu_offload: sharded_tensor = sharded_tensor.cpu() sharded_sd[new_param_name] = nn.Parameter(sharded_tensor) # choose `assign=True` since we cannot call `copy_` on meta tensor return model.load_state_dict(sharded_sd, strict=strict, assign=True) ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """Utilities for selecting and loading models.""" import contextlib import glob import os import re from collections import defaultdict from collections.abc import Callable, Iterator from typing import Any, Dict, Type import torch from torch import nn from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @contextlib.contextmanager def set_default_torch_dtype(dtype: torch.dtype): """Sets the default torch dtype to the given dtype.""" old_dtype = torch.get_default_dtype() torch.set_default_dtype(dtype) try: yield finally: torch.set_default_dtype(old_dtype) def get_param_names_mapping( mapping_dict: dict[str, str | tuple[str, int, int]], ) -> Callable[[str], tuple[str, Any, Any]]: """ Creates a mapping function that transforms parameter names using regex patterns. Args: mapping_dict (Dict[str, str]): Dictionary mapping regex patterns to replacement patterns Returns: Callable[[str], str]: A function that maps parameter names from source to target format """ def mapping_fn(name: str) -> tuple[str, Any, Any]: # support chained conversions, e.g.: # transformer.xxx.lora_down -> xxx.lora_down -> xxx.proj_down merge_index = None total_split_params = None max_steps = max(8, len(mapping_dict) * 2) applied_patterns: set[str] = set() visited_names: set[str] = {name} for _ in range(max_steps): transformed = False for pattern, replacement in mapping_dict.items(): # avoid re-applying the same rule on its own output if pattern in applied_patterns: continue if re.match(pattern, name) is None: continue curr_merge_index = None curr_total_split_params = None if isinstance(replacement, tuple): curr_merge_index = replacement[1] curr_total_split_params = replacement[2] replacement = replacement[0] new_name = re.sub(pattern, replacement, name) if new_name != name: if curr_merge_index is not None: merge_index = curr_merge_index total_split_params = curr_total_split_params name = new_name applied_patterns.add(pattern) if name in visited_names: transformed = False break visited_names.add(name) transformed = True break if not transformed: break return name, merge_index, total_split_params return mapping_fn def hf_to_custom_state_dict( hf_param_sd: dict[str, torch.Tensor] | Iterator[tuple[str, torch.Tensor]], param_names_mapping: Callable[[str], tuple[str, Any, Any]], ) -> tuple[dict[str, torch.Tensor], dict[str, tuple[str, Any, Any]]]: """ Converts a Hugging Face parameter state dictionary to a custom parameter state dictionary. Args: hf_param_sd (Dict[str, torch.Tensor]): The Hugging Face parameter state dictionary param_names_mapping (Callable[[str], tuple[str, Any, Any]]): A function that maps parameter names from source to target format Returns: custom_param_sd (Dict[str, torch.Tensor]): The custom formatted parameter state dict reverse_param_names_mapping (Dict[str, Tuple[str, Any, Any]]): Maps back from custom to hf """ custom_param_sd = {} to_merge_params = defaultdict(dict) # type: ignore reverse_param_names_mapping = {} if isinstance(hf_param_sd, dict): hf_param_sd = hf_param_sd.items() # type: ignore for source_param_name, full_tensor in hf_param_sd: # type: ignore target_param_name, merge_index, num_params_to_merge = param_names_mapping( source_param_name ) if target_param_name == "" or target_param_name is None: # type: ignore[comparison-overlap] continue reverse_param_names_mapping[target_param_name] = ( source_param_name, merge_index, num_params_to_merge, ) if merge_index is not None: to_merge_params[target_param_name][merge_index] = full_tensor if len(to_merge_params[target_param_name]) == num_params_to_merge: # cat at output dim according to the merge_index order sorted_tensors = [ to_merge_params[target_param_name][i] for i in range(num_params_to_merge) ] full_tensor = torch.cat(sorted_tensors, dim=0) del to_merge_params[target_param_name] else: continue custom_param_sd[target_param_name] = full_tensor return custom_param_sd, reverse_param_names_mapping class skip_init_modules: def __enter__(self): # Save originals self._orig_reset = {} for cls in (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d): self._orig_reset[cls] = cls.reset_parameters cls.reset_parameters = lambda self: None # skip init def __exit__(self, exc_type, exc_value, traceback): # restore originals for cls, orig in self._orig_reset.items(): cls.reset_parameters = orig def _normalize_component_type(module_type: str) -> str: """Normalize module types like 'text_encoder_2' -> 'text_encoder'.""" if module_type.endswith("_2"): return module_type[:-2] return module_type def _clean_hf_config_inplace(model_config: dict) -> None: """Remove common extraneous HF fields if present.""" for key in ( "_name_or_path", "transformers_version", "model_type", "tokenizer_class", "torch_dtype", ): model_config.pop(key, None) def _list_safetensors_files(model_path: str) -> list[str]: """List all .safetensors files under a directory.""" return sorted(glob.glob(os.path.join(str(model_path), "*.safetensors"))) BYTES_PER_GB = 1024**3 def get_memory_usage_of_component(module) -> float | None: """ returned value is in GB, rounded to 2 decimal digits """ if not isinstance(module, nn.Module): return None if hasattr(module, "get_memory_footprint"): usage = module.get_memory_footprint() / BYTES_PER_GB else: # manually param_size = sum(p.numel() * p.element_size() for p in module.parameters()) buffer_size = sum(b.numel() * b.element_size() for b in module.buffers()) total_size_bytes = param_size + buffer_size usage = total_size_bytes / (1024**3) return round(usage, 2) # component name -> ComponentLoader class component_name_to_loader_cls: Dict[str, Type[Any]] = {} ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/weight_utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/model_loader/weight_utils.py """Utilities for downloading, loading, initializing and verifying model weights.""" import hashlib import json import os import tempfile from collections.abc import Generator, Iterable from pathlib import Path import filelock import torch from safetensors.torch import safe_open from torch.distributed.tensor import DTensor from tqdm.auto import tqdm try: from runai_model_streamer import SafetensorsStreamer HAS_RUNAI_MODEL_STREAMER = True except ImportError: HAS_RUNAI_MODEL_STREAMER = False from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # use system-level temp directory for file locks, so that multiple users # can share the same lock without error. # lock files in the temp directory will be automatically deleted when the # system reboots, so users will not complain about annoying lock files temp_dir = tempfile.gettempdir() class DisabledTqdm(tqdm): def __init__(self, *args, **kwargs): kwargs["disable"] = True super().__init__(*args, **kwargs) def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None): lock_dir = cache_dir or temp_dir model_name_or_path = str(model_name_or_path) os.makedirs(os.path.dirname(lock_dir), exist_ok=True) model_name = model_name_or_path.replace("/", "-") hash_name = hashlib.sha256(model_name.encode()).hexdigest() # add hash to avoid conflict with old users' lock files lock_file_name = hash_name + model_name + ".lock" # mode 0o666 is required for the filelock to be shared across users lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666) return lock # For models like Mistral-7B-v0.3, there are both sharded # safetensors files and a consolidated safetensors file. # Passing both of these to the weight loader functionality breaks. # So, we use the index_file to # look up which safetensors files should be used. def filter_duplicate_safetensors_files( hf_weights_files: list[str], hf_folder: str, index_file: str ) -> list[str]: # model.safetensors.index.json is a mapping from keys in the # torch state_dict to safetensors file holding that weight. index_file_name = os.path.join(hf_folder, index_file) if not os.path.isfile(index_file_name): return hf_weights_files # Iterate through the weight_map (weight_name: safetensors files) # to identify weights that we should use. with open(index_file_name) as f: weight_map = json.load(f)["weight_map"] weight_files_in_index = set() for weight_name in weight_map: weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name])) # Filter out any fields that are not found in the index file. hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index] return hf_weights_files def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[str]: """ Exclude files that are not needed for inference. See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233 """ blacklist = [ "training_args.bin", "optimizer.bin", "optimizer.pt", "scheduler.pt", "scaler.pt", ] hf_weights_files = [ f for f in hf_weights_files if not any(f.endswith(x) for x in blacklist) ] return hf_weights_files # explicitly use pure text format, with a newline at the end # this makes it impossible to see the animation in the progress bar # but will avoid messing up with ray or multiprocessing, which wraps # each line of output with some prefix. _BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n" # noqa: E501 def _validate_safetensors_file(file_path: str) -> bool: """ Validate that a safetensors file is readable and not corrupted. Args: file_path: Path to the safetensors file Returns: True if file is valid, False if corrupted """ try: with safe_open(file_path, framework="pt", device="cpu") as f: _ = list(f.keys()) return True except Exception as e: logger.error( "Corrupted safetensors file detected: %s - %s: %s", file_path, type(e).__name__, str(e), ) return False def safetensors_weights_iterator( hf_weights_files: list[str], to_cpu: bool = True, use_runai_model_streamer: bool = HAS_RUNAI_MODEL_STREAMER, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over the weights in the model safetensor files.""" enable_tqdm = ( not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 ) device = "cpu" if to_cpu else str(get_local_torch_device()) # Validate files before loading corrupted_files = [ st_file for st_file in hf_weights_files if not _validate_safetensors_file(st_file) ] if corrupted_files: # Delete corrupted files (both symlink and blob if applicable) for file_path in corrupted_files: try: if os.path.islink(file_path): blob_path = os.path.realpath(file_path) os.remove(file_path) logger.info( "Removed corrupted symlink: %s", os.path.basename(file_path) ) if os.path.exists(blob_path): os.remove(blob_path) logger.info( "Removed corrupted blob: %s", os.path.basename(blob_path) ) elif os.path.isfile(file_path): os.remove(file_path) logger.info( "Removed corrupted file: %s", os.path.basename(file_path) ) except Exception as e: logger.warning("Failed to remove corrupted file %s: %s", file_path, e) raise RuntimeError( f"Found {len(corrupted_files)} corrupted safetensors file(s). " f"Files have been removed: {[os.path.basename(f) for f in corrupted_files]}. " "Please retry - the files will be re-downloaded automatically." ) if use_runai_model_streamer: with SafetensorsStreamer() as streamer: streamer.stream_files(hf_weights_files) for name, tensor in streamer.get_tensors(): if to_cpu: yield name, tensor.clone().detach() else: yield name, tensor.to(device) else: for st_file in tqdm( hf_weights_files, desc="Loading safetensors checkpoint shards", disable=not enable_tqdm, bar_format=_BAR_FORMAT, ): with safe_open(st_file, framework="pt", device=device) as f: for name in f.keys(): # noqa: SIM118 param = f.get_tensor(name) yield name, param def _load_pt_file(bin_file: str, device: str) -> dict: """Load a PyTorch checkpoint file, handling legacy tar format. PyTorch 2.6 changed the default of weights_only from False to True. Legacy tar format files cannot be loaded with weights_only=True. This function tries weights_only=True first, then falls back to False for legacy tar format files from trusted sources (HuggingFace Hub). """ try: return torch.load(bin_file, map_location=device, weights_only=True) except RuntimeError as e: if "legacy .tar format" in str(e): logger.warning( "Loading %s with weights_only=False (legacy tar format)", os.path.basename(bin_file), ) return torch.load(bin_file, map_location=device, weights_only=False) raise def pt_weights_iterator( hf_weights_files: list[str], to_cpu: bool = True, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over the weights in the model bin/pt files.""" device = "cpu" if to_cpu else str(get_local_torch_device()) enable_tqdm = ( not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 ) for bin_file in tqdm( hf_weights_files, desc="Loading pt checkpoint shards", disable=not enable_tqdm, bar_format=_BAR_FORMAT, ): state = _load_pt_file(bin_file, device) yield from state.items() del state def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: """Default weight loader.""" try: if param.numel() == 1 and loaded_weight.numel() == 1: # Sometimes scalar values aren't considered tensors with shapes # so if both param and loaded_weight are a scalar, # "broadcast" instead of copy param.data.fill_(loaded_weight.item()) else: assert param.size() == loaded_weight.size(), ( f"Attempted to load weight ({loaded_weight.size()}) " f"into parameter ({param.size()})" ) param.data.copy_(loaded_weight) except Exception: # NOTE: This exception is added for the purpose of setting breakpoint to # debug weight loading issues. raise def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None: """Remap the name of FP8 k/v_scale parameters. This function handles the remapping of FP8 k/v_scale parameter names. It detects if the given name ends with a suffix and attempts to remap it to the expected name format in the model. If the remapped name is not found in the params_dict, a warning is printed and None is returned. Args: name (str): The original loaded checkpoint parameter name. params_dict (dict): Dictionary containing the model's named parameters. Returns: str: The remapped parameter name if successful, or the original name if no remapping is needed. None: If the remapped name is not found in params_dict. """ if name.endswith(".kv_scale"): logger.warning_once( "DEPRECATED. Found kv_scale in the checkpoint. " "This format is deprecated in favor of separate k_scale and " "v_scale tensors and will be removed in a future release. " "Functionally, we will remap kv_scale to k_scale and duplicate " "k_scale to v_scale" ) # NOTE: we remap the deprecated kv_scale to k_scale remapped_name = name.replace(".kv_scale", ".attn.k_scale") if remapped_name not in params_dict: logger.warning_once( f"Found kv_scale in the checkpoint (e.g. {name}), " "but not found the expected name in the model " f"(e.g. {remapped_name}). kv_scale is " "not loaded." ) return None return remapped_name possible_scale_names = [".k_scale", ".v_scale"] modelopt_scale_names = [".self_attn.k_proj.k_scale", ".self_attn.v_proj.v_scale"] for scale_name in possible_scale_names: if name.endswith(scale_name): if any(mo_scale_name in name for mo_scale_name in modelopt_scale_names): remapped_name = name.replace( f".self_attn.{scale_name[1]}_proj{scale_name}", f".self_attn.attn{scale_name}", ) else: remapped_name = name.replace(scale_name, f".attn{scale_name}") if remapped_name not in params_dict: logger.warning_once( f"Found {scale_name} in the checkpoint (e.g. {name}), " "but not found the expected name in the model " f"(e.g. {remapped_name}). {scale_name} is " "not loaded." ) return None return remapped_name # If there were no matches, return the untouched param name return name def compute_weights_checksum( named_params: Iterable[tuple[str, torch.Tensor]], ) -> str: """Compute a SHA-256 checksum for a set of (name, tensor) pairs. Used to verify the correctness of weight refitting. After a refit, compare the checksum of the in-GPU model weights against the checksum of the on-disk tensors or the tensors in the training engine. """ hasher = hashlib.sha256() for name, tensor in sorted(named_params, key=lambda x: x[0]): hasher.update(name.encode()) t = tensor.detach() # DTensor doesn't support .numpy(); extract the local tensor. if isinstance(t, DTensor): t = t._local_tensor hasher.update(t.cpu().contiguous().reshape(-1).view(torch.uint8).numpy().data) return hasher.hexdigest() ================================================ FILE: python/sglang/multimodal_gen/runtime/loader/weights_updater.py ================================================ """ In-place weight updates for diffusion pipeline modules. This module provides WeightsUpdater, which swaps model weights at runtime without restarting the server. It is the diffusion-engine counterpart of the LLM engine's ModelRunner.update_weights_from_disk. Detailed usage of higher level API can be found in /python/sglang/multimodal_gen/test/server/test_update_weights_from_disk.py Key design decisions: - All-or-nothing with rollback: modules are updated sequentially. If any module fails (shape mismatch, corrupted file, etc.), every module that was already updated is rolled back by reloading its weights from pipeline.model_path (the last successfully-loaded checkpoint). On success, pipeline.model_path is updated to the new model_path so that future rollbacks target the latest good checkpoint, not the originally-launched model. - Rollback failures propagate: if rollback itself fails, the exception is not caught so the caller knows the model is in an inconsistent state. This matches the LLM engine behaviour. - Offload-aware: the diffusion LayerwiseOffloadManager replaces GPU parameters with torch.empty((1,)) placeholders while real weights live in consolidated pinned CPU buffers. A naive param.data.copy_() would fail with a shape mismatch. Instead, the updater dynamically detects active offload managers and writes new weights directly into their CPU buffers via update_cpu_weights(), bypassing the placeholders entirely. For any layer that happens to be prefetched on GPU at update time, the live GPU tensor is also updated so the change takes effect immediately. This requires no extra GPU memory and does not disturb the offload state. - DTensor-aware: parameters that have been distributed via torch.distributed.tensor are updated through distribute_tensor so that each shard is correctly placed on the right device mesh. """ from __future__ import annotations import gc from pathlib import Path import torch from torch.distributed.tensor import DTensor, distribute_tensor from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheMixin from sglang.multimodal_gen.runtime.loader.utils import ( _list_safetensors_files, ) from sglang.multimodal_gen.runtime.loader.weight_utils import ( safetensors_weights_iterator, ) from sglang.multimodal_gen.runtime.pipelines.diffusers_pipeline import DiffusersPipeline from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) def get_updatable_modules(pipeline) -> dict[str, torch.nn.Module]: """Return updatable nn.Module components for the given pipeline. Works with both the native ComposedPipelineBase backend and the DiffusersPipeline wrapper. """ if isinstance(pipeline, DiffusersPipeline): diffusers_pipe = pipeline.get_module("diffusers_pipeline") if diffusers_pipe is not None and diffusers_pipe.components is not None: raw = diffusers_pipe.components else: raw = {} else: raw = pipeline.modules return {n: m for n, m in raw.items() if isinstance(m, torch.nn.Module)} def _get_weights_iter(weights_dir: str): """Return a (name, tensor) iterator over safetensors in weights_dir.""" safetensors_files = _list_safetensors_files(weights_dir) if not safetensors_files: raise FileNotFoundError(f"No safetensors files found in {weights_dir}") return safetensors_weights_iterator(safetensors_files) def _validate_weight_files( local_model_path: str, modules_to_update: list[tuple[str, torch.nn.Module]], ) -> tuple[dict[str, str], list[str]]: """Check that every module has a weights directory with safetensors files. Returns: (weights_map, missing) where weights_map maps module name to its weights directory and missing lists modules without weight files. """ weights_map: dict[str, str] = {} missing: list[str] = [] for module_name, _ in modules_to_update: weights_dir = Path(local_model_path) / module_name if weights_dir.exists() and _list_safetensors_files(str(weights_dir)): weights_map[module_name] = str(weights_dir) else: missing.append(module_name) return weights_map, missing def _load_weights_into_module(module: torch.nn.Module, weights_iter) -> None: """Load weights into a module, handling offload-managed parameters. For offloaded modules, updates CPU buffers directly via update_cpu_weights(); non-offloaded parameters use in-place copy. """ offload_managers: list = [] if isinstance(module, OffloadableDiTMixin) and module.layerwise_offload_managers: offload_managers = [m for m in module.layerwise_offload_managers if m.enabled] if offload_managers: weight_dict = dict(weights_iter) offloaded_names: set[str] = set() for manager in offload_managers: offloaded_names.update(manager.update_cpu_weights(weight_dict)) remaining = ((n, w) for n, w in weight_dict.items() if n not in offloaded_names) load_weights_into_model(remaining, dict(module.named_parameters())) else: load_weights_into_model(weights_iter, dict(module.named_parameters())) def load_weights_into_model(weights_iter, model_params: dict) -> None: """Copy weights from weights_iter into model_params in-place.""" for name, loaded_weight in weights_iter: if name not in model_params: continue param = model_params[name] if param.shape != loaded_weight.shape: raise ValueError( f"Shape mismatch for {name}: model={param.shape}, loaded={loaded_weight.shape}" ) if isinstance(param, DTensor): distributed_weight = distribute_tensor( loaded_weight.to(param.dtype), param.device_mesh, param.placements, ) param._local_tensor.copy_(distributed_weight._local_tensor) else: param.data.copy_(loaded_weight.to(param.dtype)) class WeightsUpdater: """In-place weight updates for diffusion pipeline modules. Args: pipeline: A ComposedPipelineBase (or DiffusersPipeline) instance whose modules will be updated. The pipeline's model_path attribute is used for rollback on failure. """ def __init__(self, pipeline): self.pipeline = pipeline def update_weights_from_disk( self, model_path: str, flush_cache: bool = True, target_modules: list[str] | None = None, ) -> tuple[bool, str]: """Update model weights from disk without restarting the server.""" logger.info(f"Updating weights from disk: {model_path}") try: modules_to_update = self._collect_modules(target_modules) except ValueError as e: logger.error(str(e)) return False, str(e) if not modules_to_update: error_msg = ( f"No matching modules found for update. " f"Requested: {target_modules}. " f"Available nn.Module(s): {list(get_updatable_modules(self.pipeline).keys())}" ) logger.error(error_msg) return False, error_msg try: local_model_path = maybe_download_model(model_path) except Exception as e: return False, f"Failed to download model: {e}" weights_map, missing = _validate_weight_files( local_model_path, modules_to_update ) if missing: error_msg = ( f"Cannot update weights: missing weight files for modules: {missing}. " f"No partial updates allowed." ) logger.error(error_msg) return False, error_msg logger.info( f"Updating {len(weights_map)} modules: " + ", ".join(f"{n} <- {p}" for n, p in weights_map.items()) ) success, message = self._apply_weights(modules_to_update, weights_map) gc.collect() torch.cuda.empty_cache() if success and flush_cache: for _, module in modules_to_update: if isinstance(module, TeaCacheMixin): module.reset_teacache_state() logger.info(message) return success, message def _collect_modules( self, target_modules: list[str] | None ) -> list[tuple[str, torch.nn.Module]]: """Resolve target_modules to (name, module) pairs. Raises: ValueError: If target_modules contains names not found in the pipeline. """ components = get_updatable_modules(self.pipeline) if target_modules is None: names = list(components.keys()) else: unknown = [n for n in target_modules if n not in components] if unknown: raise ValueError( f"Module(s) requested for update not found in pipeline: {unknown}. " f"Available Module(s): {list(components.keys())}" ) names = target_modules return [(name, components[name]) for name in names] def _apply_weights( self, modules_to_update: list[tuple[str, torch.nn.Module]], weights_map: dict[str, str], ) -> tuple[bool, str]: """Load weights into each module; rollback on first failure.""" updated_modules: list[str] = [] for module_name, module in modules_to_update: try: weights_iter = _get_weights_iter(weights_map[module_name]) _load_weights_into_module(module, weights_iter) updated_modules.append(module_name) except Exception as e: rollback_list = updated_modules + [module_name] logger.error( f"Weight update failed for module '{module_name}': {e}. " f"Rolling back {len(rollback_list)} module(s) " f"(including partially-loaded '{module_name}'): " f"{rollback_list}.", exc_info=True, ) self._rollback(rollback_list) return False, ( f"Failed to update module '{module_name}': {e}. " f"All modules rolled back to original weights." ) names = ", ".join(updated_modules) return True, f"Updated {len(updated_modules)} modules ({names})." def _rollback(self, updated_modules: list[str]) -> None: """Restore updated_modules to original weights. If rollback itself fails the exception propagates so the caller knows the model is in an inconsistent state. """ if not updated_modules: return original_path = maybe_download_model(self.pipeline.model_path) for name in updated_modules: module = self.pipeline.get_module(name) if module is None: continue weights_dir = Path(original_path) / name if not weights_dir.exists(): continue weights_iter = _get_weights_iter(str(weights_dir)) _load_weights_into_module(module, weights_iter) ================================================ FILE: python/sglang/multimodal_gen/runtime/managers/forward_context.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/forward_context.py import time from collections import defaultdict from contextlib import contextmanager from dataclasses import dataclass from typing import TYPE_CHECKING, Optional, Type import torch from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger if TYPE_CHECKING: from sglang.multimodal_gen.runtime.layers.attention import AttentionMetadata from sglang.multimodal_gen.runtime.pipelines_core import Req logger = init_logger(__name__) # TODO(will): check if this is needed # track_batchsize: bool = envs.SGLANG_DIFFUSION_LOG_BATCHSIZE_INTERVAL >= 0 track_batchsize: bool = False last_logging_time: float = 0 forward_start_time: float = 0 # batchsize_logging_interval: float = envs.SGLANG_DIFFUSION_LOG_BATCHSIZE_INTERVAL batchsize_logging_interval: float = 1000 batchsize_forward_time: defaultdict = defaultdict(list) @dataclass class ForwardContext: current_timestep: int # TODO(will): check this arg # copy from vllm_config.compilation_config.static_forward_context # attn_layers: Dict[str, Any] # TODO: extend to support per-layer dynamic forward context attn_metadata: "AttentionMetadata" # set dynamically for each forward pass forward_batch: Optional["Req"] = None attention_backend_cls: Optional[Type] = None def set_attn_backend_cls(self, attention_backend_cls: Type): if self.attention_backend_cls: if self.attention_backend_cls != attention_backend_cls: raise RuntimeError( f"Different types of attention backend in a same context detected, previous: {self.attention_backend_cls}, new: {attention_backend_cls}" ) else: self.attention_backend_cls = attention_backend_cls _forward_context: Optional["ForwardContext"] = None def get_forward_context() -> "ForwardContext": """Get the current forward context.""" assert _forward_context is not None, ( "Forward context is not set. " "Please use `set_forward_context` to set the forward context." ) return _forward_context # TODO(will): finalize the interface @contextmanager def set_forward_context( current_timestep, attn_metadata, forward_batch: Optional["Req"] = None ): """A context manager that stores the current forward context, can be attention metadata, etc. Here we can inject common logic for every model forward pass. """ global forward_start_time need_to_track_batchsize = track_batchsize and attn_metadata is not None if need_to_track_batchsize: forward_start_time = time.perf_counter() global _forward_context prev_context = _forward_context _forward_context = ForwardContext( current_timestep=current_timestep, attn_metadata=attn_metadata, forward_batch=forward_batch, ) try: yield finally: global last_logging_time, batchsize_logging_interval if need_to_track_batchsize: if hasattr(attn_metadata, "num_prefill_tokens"): # for v0 attention backends batchsize = ( attn_metadata.num_prefill_tokens + attn_metadata.num_decode_tokens ) else: # for v1 attention backends batchsize = attn_metadata.num_input_tokens now = time.perf_counter() # time measurement is in milliseconds batchsize_forward_time[batchsize].append((now - forward_start_time) * 1000) if now - last_logging_time > batchsize_logging_interval: last_logging_time = now forward_stats = [] for bs, times in batchsize_forward_time.items(): if len(times) <= 1: # can be cudagraph / profiling run continue medium = torch.quantile(torch.tensor(times), q=0.5).item() medium = round(medium, 2) forward_stats.append((bs, len(times), medium)) forward_stats.sort(key=lambda x: x[1], reverse=True) if forward_stats: logger.info( ( "Batchsize forward time stats " "(batchsize, count, median_time(ms)): %s" ), forward_stats, ) _forward_context = prev_context ================================================ FILE: python/sglang/multimodal_gen/runtime/managers/gpu_worker.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import gc import multiprocessing as mp import os import time from typing import List, Union import torch from setproctitle import setproctitle from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.distributed import ( get_sp_group, get_tp_rank, get_tp_world_size, maybe_init_distributed_environment_and_model_parallel, model_parallel_is_initialized, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_cfg_group, get_classifier_free_guidance_rank, get_classifier_free_guidance_world_size, get_ring_parallel_rank, get_ring_parallel_world_size, get_tp_group, get_ulysses_parallel_rank, get_ulysses_parallel_world_size, ) from sglang.multimodal_gen.runtime.entrypoints.utils import save_outputs from sglang.multimodal_gen.runtime.loader.weight_utils import compute_weights_checksum from sglang.multimodal_gen.runtime.loader.weights_updater import ( WeightsUpdater, get_updatable_modules, ) from sglang.multimodal_gen.runtime.pipelines_core import ( ComposedPipelineBase, LoRAPipeline, Req, build_pipeline, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch, set_musa_arch from sglang.multimodal_gen.runtime.utils.layerwise_offload import ( OffloadableDiTMixin, iter_materialized_weights, ) from sglang.multimodal_gen.runtime.utils.logging_utils import ( configure_logger, globally_suppress_loggers, init_logger, ) from sglang.multimodal_gen.runtime.utils.perf_logger import ( PerformanceLogger, capture_memory_snapshot, ) from sglang.srt.utils.network import NetworkAddress logger = init_logger(__name__) class GPUWorker: """ A worker that executes the model on a single GPU. """ def __init__( self, local_rank: int, rank: int, master_port: int, server_args: ServerArgs, ): self.local_rank = local_rank self.rank = rank self.master_port = master_port # FIXME: should we use tcp as distribute init method? self.server_args = server_args self.pipeline: ComposedPipelineBase = None self.init_device_and_model() self.sp_group = get_sp_group() self.sp_cpu_group = self.sp_group.cpu_group self.tp_group = get_tp_group() self.tp_cpu_group = self.tp_group.cpu_group self.cfg_group = get_cfg_group() self.cfg_cpu_group = self.cfg_group.cpu_group def init_device_and_model(self) -> None: """Initialize the device and load the model.""" torch.get_device_module().set_device(self.local_rank) # Set environment variables for distributed initialization os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(self.master_port) os.environ["LOCAL_RANK"] = str(self.local_rank) os.environ["RANK"] = str(self.rank) os.environ["WORLD_SIZE"] = str(self.server_args.num_gpus) # initialize the distributed environment maybe_init_distributed_environment_and_model_parallel( tp_size=self.server_args.tp_size, enable_cfg_parallel=self.server_args.enable_cfg_parallel, ulysses_degree=self.server_args.ulysses_degree, ring_degree=self.server_args.ring_degree, sp_size=self.server_args.sp_degree, dp_size=self.server_args.dp_size, distributed_init_method=NetworkAddress( "127.0.0.1", self.master_port ).to_tcp(), dist_timeout=self.server_args.dist_timeout, ) # set proc title if model_parallel_is_initialized(): suffix = "" if get_tp_world_size() != 1: tp_rank = get_tp_rank() suffix += f"_TP{tp_rank}" if get_ulysses_parallel_world_size() != 1: u_rank = get_ulysses_parallel_rank() suffix += f"_U{u_rank}" if get_ring_parallel_world_size() != 1: r_rank = get_ring_parallel_rank() suffix += f"_R{r_rank}" if get_classifier_free_guidance_world_size() != 1: c_rank = get_classifier_free_guidance_rank() suffix += f"_C{c_rank}" setproctitle(f"sgl_diffusion::scheduler{suffix}") else: setproctitle(f"sgl_diffusion::scheduler_{self.local_rank}") self.pipeline = build_pipeline(self.server_args) # apply layerwise offload after lora is applied while building LoRAPipeline # otherwise empty offloaded weights could fail lora converting if self.server_args.dit_layerwise_offload: # enable layerwise offload if possible for module_name in [ "transformer", "transformer_2", "video_dit", "video_dit_2", "audio_dit", ]: dit = self.pipeline.get_module(module_name) if dit: if isinstance(dit, OffloadableDiTMixin): dit.configure_layerwise_offload(self.server_args) else: logger.info( f"Module {type(dit).__name__} does not support layerwise offload. Skipping." ) logger.info( f"Worker {self.rank}: Initialized device, model, and distributed environment." ) def do_mem_analysis(self, output_batch: OutputBatch): final_snapshot = capture_memory_snapshot() if output_batch.metrics: output_batch.metrics.record_memory_snapshot("mem_analysis", final_snapshot) # for details on max_memory_reserved: https://docs.pytorch.org/docs/stable/generated/torch.cuda.memory.max_memory_reserved.html peak_reserved_bytes = torch.get_device_module().max_memory_reserved() peak_allocated_bytes = torch.get_device_module().max_memory_allocated() output_batch.peak_memory_mb = peak_reserved_bytes / (1024**2) peak_reserved_gb = peak_reserved_bytes / (1024**3) peak_allocated_gb = peak_allocated_bytes / (1024**3) remaining_gpu_mem_gb = ( current_platform.get_device_total_memory() / (1024**3) - peak_reserved_gb ) can_stay_resident = self.get_can_stay_resident_components(remaining_gpu_mem_gb) suggested_args = set() component_to_arg = { "vae": "--vae-cpu-offload", "text_encoder": "--text-encoder-cpu-offload", "text_encoder_2": "--text-encoder-cpu-offload", "image_encoder": "--image-encoder-cpu-offload", } for component in can_stay_resident: if component == "transformer": if self.server_args.dit_layerwise_offload: suggested_args.add("--dit-layerwise-offload") elif self.server_args.dit_cpu_offload: suggested_args.add("--dit-cpu-offload") elif component in component_to_arg: suggested_args.add(component_to_arg[component]) suggested_args_str = ( ", ".join(sorted(suggested_args)) if suggested_args else "None" ) pool_overhead_gb = peak_reserved_gb - peak_allocated_gb logger.info( f"Peak GPU memory: {peak_reserved_gb:.2f} GB, " f"Peak allocated: {peak_allocated_gb:.2f} GB, " f"Memory pool overhead: {pool_overhead_gb:.2f} GB ({pool_overhead_gb / peak_reserved_gb * 100:.1f}%), " f"Remaining GPU memory at peak: {remaining_gpu_mem_gb:.2f} GB. " f"Components that could stay resident (based on the last request workload): {can_stay_resident}. " f"Related offload server args to disable: {suggested_args_str}" ) def execute_forward(self, batch: List[Req]) -> OutputBatch: """ Execute a forward pass. """ assert self.pipeline is not None req = batch[0] output_batch = None try: if self.rank == 0: torch.get_device_module().reset_peak_memory_stats() start_time = time.monotonic() # capture memory baseline before forward if self.rank == 0 and req.metrics: baseline_snapshot = capture_memory_snapshot() req.metrics.record_memory_snapshot("before_forward", baseline_snapshot) req.log(server_args=self.server_args) result = self.pipeline.forward(req, self.server_args) if isinstance(result, Req): output_batch = OutputBatch( output=result.output, audio=getattr(result, "audio", None), audio_sample_rate=getattr(result, "audio_sample_rate", None), metrics=result.metrics, trajectory_timesteps=getattr(result, "trajectory_timesteps", None), trajectory_latents=getattr(result, "trajectory_latents", None), noise_pred=getattr(result, "noise_pred", None), trajectory_decoded=getattr(result, "trajectory_decoded", None), ) else: output_batch = result # capture memory after forward (peak) if self.rank == 0 and output_batch.metrics: peak_snapshot = capture_memory_snapshot() output_batch.metrics.record_memory_snapshot( "after_forward", peak_snapshot ) if self.rank == 0 and not req.suppress_logs: self.do_mem_analysis(output_batch) duration_ms = (time.monotonic() - start_time) * 1000 output_batch.metrics.total_duration_ms = duration_ms # Save output to file and return file path only if requested. Avoid the serialization # and deserialization overhead between scheduler_client and gpu_worker. if req.save_output and req.return_file_paths_only: if self.rank == 0 and output_batch.output is not None: output_paths = save_outputs( output_batch.output, req.data_type, req.fps, True, lambda idx: req.output_file_path(len(output_batch.output), idx), audio=output_batch.audio, audio_sample_rate=output_batch.audio_sample_rate, output_compression=req.output_compression, enable_frame_interpolation=req.enable_frame_interpolation, frame_interpolation_exp=req.frame_interpolation_exp, frame_interpolation_scale=req.frame_interpolation_scale, frame_interpolation_model_path=req.frame_interpolation_model_path, enable_upscaling=req.enable_upscaling, upscaling_model_path=req.upscaling_model_path, upscaling_scale=req.upscaling_scale, ) output_batch.output_file_paths = output_paths # No rank needs to hold on to generated tensors once the file-path # response has been materialized on rank 0 output_batch.output = None output_batch.audio = None output_batch.audio_sample_rate = None if torch.cuda.is_initialized(): torch.cuda.empty_cache() # TODO: extract to avoid duplication if req.perf_dump_path is not None or envs.SGLANG_DIFFUSION_STAGE_LOGGING: # Avoid logging warmup perf records that share the same request_id. if not req.is_warmup: PerformanceLogger.log_request_summary(metrics=output_batch.metrics) except Exception as e: logger.error( f"Error executing request {req.request_id}: {e}", exc_info=True ) if isinstance(e, _oom_exceptions()): logger.warning(OOM_MSG) if output_batch is None: output_batch = OutputBatch() output_batch.error = f"Error executing request {req.request_id}: {e}" return output_batch def get_can_stay_resident_components( self, remaining_gpu_mem_gb: float ) -> List[str]: """ Calculate which components can stay resident on GPU without being offloaded. """ can_stay_resident = [] if not self.pipeline: return can_stay_resident # Map memory_usage keys to server_args offload flags # If the flag is False, the component is ALREADY resident, so we don't suggest it. # If the flag is True, it is currently offloaded, so it's a candidate to "stay resident". offload_flags = { "transformer": self.server_args.dit_cpu_offload or self.server_args.dit_layerwise_offload, "vae": self.server_args.vae_cpu_offload, "text_encoder": self.server_args.text_encoder_cpu_offload, "text_encoder_2": self.server_args.text_encoder_cpu_offload, "image_encoder": self.server_args.image_encoder_cpu_offload, } for name, usage in self.pipeline.memory_usages.items(): # Only consider components that are currently configured to be offloaded is_offload_configured = offload_flags.get(name, False) if not is_offload_configured: continue if usage <= remaining_gpu_mem_gb: can_stay_resident.append(name) remaining_gpu_mem_gb -= usage return can_stay_resident def set_lora( self, lora_nickname: Union[str, List[str]], lora_path: Union[str, None, List[Union[str, None]]] = None, target: Union[str, List[str]] = "all", strength: Union[float, List[float]] = 1.0, ) -> OutputBatch: """ Set the LoRA adapter(s) for the pipeline. Supports both single LoRA (backward compatible) and multiple LoRA adapters. Args: lora_nickname: The nickname(s) of the adapter(s). Can be a string or a list of strings. lora_path: Path(s) to the LoRA adapter(s). Can be a string, None, or a list of strings/None. target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings. strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats. """ if not isinstance(self.pipeline, LoRAPipeline): return OutputBatch(error="Lora is not enabled") self.pipeline.set_lora(lora_nickname, lora_path, target, strength) return OutputBatch() def merge_lora_weights( self, target: str = "all", strength: float = 1.0 ) -> OutputBatch: """ Merge LoRA weights. Args: target: Which transformer(s) to merge. strength: LoRA strength for merge, default 1.0. """ if not isinstance(self.pipeline, LoRAPipeline): return OutputBatch(error="Lora is not enabled") self.pipeline.merge_lora_weights(target, strength) return OutputBatch() def unmerge_lora_weights(self, target: str = "all") -> OutputBatch: """ Unmerge LoRA weights. Args: target: Which transformer(s) to unmerge. """ if not isinstance(self.pipeline, LoRAPipeline): return OutputBatch(error="Lora is not enabled") self.pipeline.unmerge_lora_weights(target) return OutputBatch() def list_loras(self) -> OutputBatch: """ List loaded LoRA adapters and current application status per module. """ from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import ( LoRAPipeline, ) if not isinstance(self.pipeline, LoRAPipeline): return OutputBatch(error="Lora is not enabled") status = self.pipeline.get_lora_status() return OutputBatch(output=status) def update_weights_from_disk( self, model_path: str, flush_cache: bool = True, target_modules: list[str] | None = None, ) -> tuple[bool, str]: """Update model weights from disk inplace without restarting the server.""" if not self.pipeline: return False, "Pipeline is not initialized" updater = WeightsUpdater(self.pipeline) success, message = updater.update_weights_from_disk( model_path, flush_cache=flush_cache, target_modules=target_modules, ) if success: self.server_args.model_path = model_path self.pipeline.model_path = model_path return success, message def get_weights_checksum( self, module_names: list[str] | None = None ) -> dict[str, str]: """Compute SHA-256 checksum of each module's weights.""" if not self.pipeline: return {"error": "Pipeline is not initialized"} all_modules = get_updatable_modules(self.pipeline) names = module_names if module_names is not None else list(all_modules.keys()) checksums: dict[str, str] = {} for name in names: module = all_modules.get(name) if module is None: checksums[name] = "not_found" continue checksums[name] = compute_weights_checksum( iter_materialized_weights(module) ) return checksums OOM_MSG = f""" OOM detected. Possible solutions: - If the OOM occurs during loading: 1. Enable CPU offload for memory-intensive components, or use `--dit-layerwise-offload` for DiT - If the OOM occurs during runtime: 1. Enable SP and/or TP (in a multi-GPU setup) 2. Reduce the number of output tokens by lowering resolution or decreasing `--num-frames` 3. Opt for a sparse-attention backend 4. Enable FSDP by `--use-fsdp-inference` (in a multi-GPU setup) 5. Enable quantization (e.g. nunchaku) Or, open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose """ def _oom_exceptions(): # torch.OutOfMemoryError exists only in some PyTorch builds types = [torch.cuda.OutOfMemoryError] if hasattr(torch, "OutOfMemoryError"): types.append(torch.OutOfMemoryError) return tuple(types) def run_scheduler_process( local_rank: int, rank: int, master_port: int, server_args: ServerArgs, pipe_writer: mp.connection.Connection, # For all workers: pipe to receive tasks from rank 0 task_pipe_r: mp.connection.Connection, # For slave workers: pipe to send results back to rank 0 result_pipe_w: mp.connection.Connection | None, # For rank 0 worker only: pipes to send tasks to slaves task_pipes_to_slaves: list[mp.connection.Connection] | None = None, # For rank 0 worker only: pipes to receive results from slaves result_pipes_from_slaves: list[mp.connection.Connection] | None = None, ) -> None: """ The entry point for the worker process. Rank 0 acts as the master, handling ZMQ requests and coordinating slaves. Ranks > 0 act as slaves, waiting for tasks from the master. """ configure_logger(server_args) globally_suppress_loggers() if current_platform.is_cuda(): set_cuda_arch() elif current_platform.is_musa(): set_musa_arch() port_args = PortArgs.from_server_args(server_args) # start the scheduler event loop assert task_pipes_to_slaves is not None assert result_pipes_from_slaves is not None from sglang.multimodal_gen.runtime.managers.scheduler import Scheduler try: scheduler = Scheduler( server_args, gpu_id=rank, port_args=port_args, task_pipes_to_slaves=task_pipes_to_slaves, result_pipes_from_slaves=result_pipes_from_slaves, ) logger.info(f"Worker {rank}: Scheduler loop started.") pipe_writer.send( { "status": "ready", } ) scheduler.event_loop() except _oom_exceptions() as _e: logger.warning(OOM_MSG) raise finally: # Clean up resources to speed up shutdown if "scheduler" in locals(): del scheduler gc.collect() if torch.cuda.is_initialized(): torch.cuda.empty_cache() if torch.distributed.is_available() and torch.distributed.is_initialized(): torch.distributed.destroy_process_group() logger.info(f"Worker {rank}: Shutdown complete.") ================================================ FILE: python/sglang/multimodal_gen/runtime/managers/scheduler.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import asyncio import os import pickle from collections import deque from typing import Any, List import zmq from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( _parse_size, save_image_to_path, ) from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( GetWeightsChecksumReqInput, UpdateWeightFromDiskReqInput, ) from sglang.multimodal_gen.runtime.entrypoints.utils import ( ListLorasReq, MergeLoraWeightsReq, SetLoraReq, ShutdownReq, UnmergeLoraWeightsReq, ) from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker from sglang.multimodal_gen.runtime.pipelines_core import Req from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.server_args import ( PortArgs, ServerArgs, set_global_server_args, ) from sglang.multimodal_gen.runtime.utils.common import get_zmq_socket from sglang.multimodal_gen.runtime.utils.distributed import broadcast_pyobj from sglang.multimodal_gen.runtime.utils.logging_utils import GREEN, RESET, init_logger logger = init_logger(__name__) MINIMUM_PICTURE_BASE64_FOR_WARMUP = "data:image/jpg;base64,iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAACXBIWXMAAA7EAAAOxAGVKw4bAAAAbUlEQVRYhe3VsQ2AMAxE0Y/lIgNQULD/OqyCMgCihCKSG4yRuKuiNH6JLsoEbMACOGBcua9HOR7Y6w6swBwMy0qLTpkeI77qdEBpBFAHBBDAGH8WrwJKI4AAegUCfAKgEgpQDvh3CR3oQCuav58qlAw73kKCSgAAAABJRU5ErkJggg==" class Scheduler: """ Runs the main event loop for the rank 0 worker. It listens for external requests via ZMQ and coordinates with other workers. This class does NOT manage worker processes. """ def __init__( self, server_args: ServerArgs, gpu_id: int, port_args: PortArgs, task_pipes_to_slaves: list = None, result_pipes_from_slaves: list = None, ): self.server_args = server_args self.port_args = port_args set_global_server_args(server_args=server_args) # Inter-process Communication self.context = zmq.Context(io_threads=2) endpoint = server_args.scheduler_endpoint if gpu_id == 0: # router allocates identify (envelope) for each connection self.receiver, actual_endpoint = get_zmq_socket( self.context, zmq.ROUTER, endpoint, True ) logger.info(f"Scheduler bind at endpoint: {actual_endpoint}") else: self.receiver = None worker = GPUWorker( local_rank=gpu_id, master_port=port_args.master_port, rank=gpu_id, server_args=server_args, ) self.worker = worker self.task_pipes_to_slaves = task_pipes_to_slaves self.result_pipes_from_slaves = result_pipes_from_slaves self.gpu_id = gpu_id self._running = True self.request_handlers = { SetLoraReq: self._handle_set_lora, MergeLoraWeightsReq: self._handle_merge_lora, UnmergeLoraWeightsReq: self._handle_unmerge_lora, Req: self._handle_generation, List[Req]: self._handle_generation, ListLorasReq: self._handle_list_loras, ShutdownReq: self._handle_shutdown, UpdateWeightFromDiskReqInput: self._handle_update_weights_from_disk, GetWeightsChecksumReqInput: self._handle_get_weights_checksum, } # FIFO, new reqs are appended self.waiting_queue: deque[tuple[bytes, Req]] = deque() # whether we've send the necessary warmup reqs self.warmed_up = False # warmup progress tracking self._warmup_total = 0 self._warmup_processed = 0 self.prepare_server_warmup_reqs() # Maximum consecutive errors before terminating the event loop self._max_consecutive_errors = 3 self._consecutive_error_count = 0 def _handle_set_lora(self, reqs: List[Any]) -> OutputBatch: # TODO: return set status # TODO: return with SetLoRAResponse or something more appropriate req = reqs[0] return self.worker.set_lora( req.lora_nickname, req.lora_path, req.target, req.strength ) def _handle_merge_lora(self, reqs: List[Any]): req = reqs[0] return self.worker.merge_lora_weights(req.target, req.strength) def _handle_unmerge_lora(self, reqs: List[Any]) -> OutputBatch: req = reqs[0] return self.worker.unmerge_lora_weights(req.target) def _handle_list_loras(self, _reqs: List[Any]) -> OutputBatch: return self.worker.list_loras() def _handle_shutdown(self, _reqs: List[Any]) -> OutputBatch: self._running = False return OutputBatch() def _handle_update_weights_from_disk(self, reqs: List[Any]) -> OutputBatch: """Handle update_weights_from_disk request for RL workflows.""" req = reqs[0] success, message = self.worker.update_weights_from_disk( model_path=req.model_path, flush_cache=req.flush_cache, target_modules=req.target_modules, ) return OutputBatch( output={"success": success, "message": message}, error=None if success else message, ) def _handle_get_weights_checksum(self, reqs: List[Any]) -> OutputBatch: """Handle get_weights_checksum request.""" req = reqs[0] checksums = self.worker.get_weights_checksum(module_names=req.module_names) return OutputBatch(output=checksums) def _handle_generation(self, reqs: List[Req]): warmup_reqs = [req for req in reqs if req.is_warmup] if warmup_reqs: self._warmup_processed += len(warmup_reqs) if self._warmup_total > 0: logger.info( f"Processing warmup req... ({self._warmup_processed}/{self._warmup_total})" ) else: logger.info("Processing warmup req...") return self.worker.execute_forward(reqs) def return_result( self, output_batch: OutputBatch, identity: bytes | None = None, is_warmup: bool = False, ): """ replies to client, only on rank 0 """ if not is_warmup and self.receiver is not None and identity is not None: self.receiver.send_multipart([identity, b"", pickle.dumps(output_batch)]) def get_next_batch_to_run(self) -> list[tuple[bytes, Req]] | None: """pull a req from waiting_queue""" if not self.waiting_queue: return None # pop the first (earliest) item = self.waiting_queue.popleft() return [item] def prepare_server_warmup_reqs(self): if ( self.server_args.warmup and not self.warmed_up and self.server_args.warmup_resolutions is not None ): # insert warmup reqs constructed with each warmup-resolution self._warmup_total = len(self.server_args.warmup_resolutions) self._warmup_processed = 0 for resolution in self.server_args.warmup_resolutions: width, height = _parse_size(resolution) task_type = self.server_args.pipeline_config.task_type if task_type in ( ModelTaskType.I2I, ModelTaskType.TI2I, ModelTaskType.I2V, ModelTaskType.TI2V, ): uploads_dir = os.path.join("outputs", "uploads") os.makedirs(uploads_dir, exist_ok=True) input_path = asyncio.run( save_image_to_path( MINIMUM_PICTURE_BASE64_FOR_WARMUP, os.path.join(uploads_dir, "warmup_image.jpg"), ) ) req = Req( data_type=task_type.data_type(), width=width, height=height, prompt="", negative_prompt="", image_path=[input_path], ) else: req = Req( data_type=task_type.data_type(), width=width, height=height, prompt="", ) req.set_as_warmup(self.server_args.warmup_steps) self.waiting_queue.append((None, req)) # if server is warmed-up, set this flag to avoid req-based warmup self.warmed_up = True def process_received_reqs_with_req_based_warmup( self, recv_reqs: List[tuple[bytes, Any]] ) -> List[tuple[bytes, Any]]: if ( self.warmed_up or not self.server_args.warmup or not recv_reqs or self.server_args.warmup_resolutions is not None ): return recv_reqs # handle server req-based warmup by inserting an identical req to the beginning of the waiting queue # only the very first req through server's lifetime will be warmed up identity, req = recv_reqs[0] if isinstance(req, Req): warmup_req = req.copy_as_warmup(self.server_args.warmup_steps) recv_reqs.insert(0, (identity, warmup_req)) self._warmup_total = 1 self._warmup_processed = 0 self.warmed_up = True return recv_reqs def recv_reqs(self) -> List[tuple[bytes, Any]]: """ For non-main schedulers, reqs are broadcasted from main using broadcast_pyobj """ if self.receiver is not None: try: try: # Accept valid REQ envelopes only, ignore malformed/probe frames. parts = self.receiver.recv_multipart(zmq.NOBLOCK) identity, payload = parts[0], parts[-1] # Ignore malformed probes or non-pickle data recv_reqs = pickle.loads(payload) if len(parts) > 2 else [] except (zmq.Again, pickle.UnpicklingError, IndexError, EOFError): recv_reqs = [] except zmq.ZMQError: # re-raise or handle appropriately to let the outer loop continue raise if recv_reqs: # Ensure recv_reqs is a list if not isinstance(recv_reqs, list): recv_reqs = [recv_reqs] # Pack with identity for rank 0 recv_reqs = [(identity, req) for req in recv_reqs] else: recv_reqs = None # TODO: fix this condition if self.server_args.sp_degree != 1: recv_reqs = broadcast_pyobj( recv_reqs, self.worker.sp_group.rank, self.worker.sp_cpu_group, src=self.worker.sp_group.ranks[0], ) if self.server_args.enable_cfg_parallel: recv_reqs = broadcast_pyobj( recv_reqs, self.worker.cfg_group.rank, self.worker.cfg_cpu_group, src=self.worker.cfg_group.ranks[0], ) if self.server_args.tp_size > 1: recv_reqs = broadcast_pyobj( recv_reqs, self.worker.tp_group.rank, self.worker.tp_cpu_group, src=self.worker.tp_group.ranks[0], ) assert recv_reqs is not None return recv_reqs def event_loop(self) -> None: """ The main event loop that listens for ZMQ requests. Handles abortion """ logger.debug( f"Rank 0 scheduler listening on tcp://*:{self.server_args.scheduler_port}" ) while self._running: # 1: receive requests try: new_reqs = self.recv_reqs() new_reqs = self.process_received_reqs_with_req_based_warmup(new_reqs) self.waiting_queue.extend(new_reqs) # Reset error count on success self._consecutive_error_count = 0 except Exception as e: self._consecutive_error_count += 1 logger.error( f"Error receiving requests in scheduler event loop " f"(attempt {self._consecutive_error_count}/{self._max_consecutive_errors}): {e}", exc_info=True, ) if self._consecutive_error_count >= self._max_consecutive_errors: logger.error( f"Maximum consecutive errors ({self._max_consecutive_errors}) reached. " "Terminating scheduler event loop." ) raise RuntimeError( f"Scheduler terminated after {self._max_consecutive_errors} " f"consecutive errors. Last error: {e}" ) from e continue # 2: execute, make sure a reply is always sent items = self.get_next_batch_to_run() if not items: continue identities = [item[0] for item in items] reqs = [item[1] for item in items] try: processed_req = reqs[0] handler = self.request_handlers.get(type(processed_req)) if handler: output_batch = handler(reqs) else: output_batch = OutputBatch( error=f"Unknown request type: {type(processed_req)}" ) except Exception as e: logger.error( f"Error executing request in scheduler event loop: {e}", exc_info=True, ) # Determine appropriate error response format output_batch = ( OutputBatch(error=str(e)) if reqs and isinstance(reqs[0], Req) else OutputBatch(error=str(e)) ) # 3. return results try: # log warmup info is_warmup = ( processed_req.is_warmup if isinstance(processed_req, Req) else False ) if is_warmup: if output_batch.error is None: if self._warmup_total > 0: logger.info( f"Warmup req ({self._warmup_processed}/{self._warmup_total}) processed in {GREEN}%.2f{RESET} seconds", output_batch.metrics.total_duration_s, ) else: logger.info( f"Warmup req processed in {GREEN}%.2f{RESET} seconds", output_batch.metrics.total_duration_s, ) else: if self._warmup_total > 0: logger.info( f"Warmup req ({self._warmup_processed}/{self._warmup_total}) processing failed" ) else: logger.info(f"Warmup req processing failed") # TODO: Support sending back to multiple identities if batched self.return_result(output_batch, identities[0], is_warmup=is_warmup) except zmq.ZMQError as e: # Reply failed; log and keep loop alive to accept future requests logger.error(f"ZMQ error sending reply: {e}") continue if self.receiver is not None: self.receiver.close() self.context.destroy(linger=0) def _broadcast_task(self, payload: dict[str, Any]) -> None: """Broadcast a task to all slave worker processes.""" method = payload["method"] kwargs = {k: v for k, v in payload.items() if k != "method"} task = {"method": method, "kwargs": kwargs} for pipe in self.task_pipes_to_slaves: pipe.send(task) def _collect_slave_results(self) -> List[dict[str, Any]]: """Collect results from all slave worker processes.""" results = [] for pipe in self.result_pipes_from_slaves: results.append(pipe.recv()) return results ================================================ FILE: python/sglang/multimodal_gen/runtime/models/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo ================================================ FILE: python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py ================================================ from typing import Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F from diffusers.models.attention import FeedForward from sglang.multimodal_gen.configs.models.adapter.ltx_2_connector import ( LTX2ConnectorConfig, ) from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum def apply_interleaved_rotary_emb( x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor] ) -> torch.Tensor: cos, sin = freqs x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1) # [B, S, C // 2] x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) return out def apply_split_rotary_emb( x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor] ) -> torch.Tensor: cos, sin = freqs x_dtype = x.dtype needs_reshape = False if x.ndim != 4 and cos.ndim == 4: # cos is (#b, h, t, r) -> reshape x to (b, h, t, dim_per_head) # The cos/sin batch dim may only be broadcastable, so take batch size from x b = x.shape[0] _, h, t, _ = cos.shape x = x.reshape(b, t, h, -1).transpose(1, 2) needs_reshape = True # Split last dim (2*r) into (d=2, r) last = x.shape[-1] if last % 2 != 0: raise ValueError( f"Expected x.shape[-1] to be even for split rotary, got {last}." ) r = last // 2 # (..., 2, r) split_x = x.reshape(*x.shape[:-1], 2, r) first_x = split_x[..., :1, :] # (..., 1, r) second_x = split_x[..., 1:, :] # (..., 1, r) cos_u = cos.unsqueeze(-2) # broadcast to (..., 1, r) against (..., 2, r) sin_u = sin.unsqueeze(-2) out = split_x * cos_u first_out = out[..., :1, :] second_out = out[..., 1:, :] first_out.addcmul_(-sin_u, second_x) second_out.addcmul_(sin_u, first_x) out = out.reshape(*out.shape[:-2], last) if needs_reshape: out = out.transpose(1, 2).reshape(b, t, -1) out = out.to(dtype=x_dtype) return out class LTX2Attention(torch.nn.Module): r""" Attention class for all LTX-2.0 attention layers. Compared to LTX-1.0, this supports specifying the query and key RoPE embeddings separately for audio-to-video (a2v) and video-to-audio (v2a) cross-attention. """ def __init__( self, query_dim: int, heads: int = 8, kv_heads: int = 8, dim_head: int = 64, dropout: float = 0.0, bias: bool = True, cross_attention_dim: Optional[int] = None, out_bias: bool = True, qk_norm: str = "rms_norm_across_heads", norm_eps: float = 1e-6, norm_elementwise_affine: bool = True, rope_type: str = "interleaved", processor=None, ): super().__init__() if qk_norm != "rms_norm_across_heads": raise NotImplementedError( "Only 'rms_norm_across_heads' is supported as a valid value for `qk_norm`." ) self.head_dim = dim_head self.inner_dim = dim_head * heads self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads self.query_dim = query_dim self.cross_attention_dim = ( cross_attention_dim if cross_attention_dim is not None else query_dim ) self.use_bias = bias self.dropout = dropout self.out_dim = query_dim self.heads = heads self.rope_type = rope_type self.norm_q = torch.nn.RMSNorm( dim_head * heads, eps=norm_eps, elementwise_affine=norm_elementwise_affine ) self.norm_k = torch.nn.RMSNorm( dim_head * kv_heads, eps=norm_eps, elementwise_affine=norm_elementwise_affine, ) self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) self.to_k = torch.nn.Linear( self.cross_attention_dim, self.inner_kv_dim, bias=bias ) self.to_v = torch.nn.Linear( self.cross_attention_dim, self.inner_kv_dim, bias=bias ) self.to_out = torch.nn.ModuleList([]) self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) self.to_out.append(torch.nn.Dropout(dropout)) # Scaled dot product attention self.attn = USPAttention( num_heads=heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, causal=False, supported_attention_backends={ AttentionBackendEnum.FA, AttentionBackendEnum.AITER, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.SAGE_ATTN_3, }, ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, query_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, key_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: batch_size, sequence_length, _ = ( hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape ) if encoder_hidden_states is None: encoder_hidden_states = hidden_states query = self.to_q(hidden_states) key = self.to_k(encoder_hidden_states) value = self.to_v(encoder_hidden_states) query = self.norm_q(query) key = self.norm_k(key) if query_rotary_emb is not None: if self.rope_type == "interleaved": query = apply_interleaved_rotary_emb(query, query_rotary_emb) key = apply_interleaved_rotary_emb( key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb, ) elif self.rope_type == "split": query = apply_split_rotary_emb(query, query_rotary_emb) key = apply_split_rotary_emb( key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb, ) query = query.unflatten(2, (self.heads, -1)) key = key.unflatten(2, (self.heads, -1)) value = value.unflatten(2, (self.heads, -1)) hidden_states = self.attn( query, key, value, ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) hidden_states = self.to_out[0](hidden_states) hidden_states = self.to_out[1](hidden_states) return hidden_states class LTX2RotaryPosEmbed1d(nn.Module): """ 1D rotary positional embeddings (RoPE) for the LTX 2.0 text encoder connectors. """ def __init__( self, dim: int, base_seq_len: int = 4096, theta: float = 10000.0, double_precision: bool = True, rope_type: str = "interleaved", num_attention_heads: int = 32, ): super().__init__() if rope_type not in ["interleaved", "split"]: raise ValueError( f"{rope_type=} not supported. Choose between 'interleaved' and 'split'." ) self.dim = dim self.base_seq_len = base_seq_len self.theta = theta self.double_precision = double_precision self.rope_type = rope_type self.num_attention_heads = num_attention_heads def forward( self, batch_size: int, pos: int, device: Union[str, torch.device], dtype: Optional[torch.dtype] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: # 1. Get 1D position ids grid_1d = torch.arange(pos, dtype=torch.float32, device=device) # Get fractional indices relative to self.base_seq_len grid_1d = grid_1d / self.base_seq_len grid = grid_1d.unsqueeze(0).repeat(batch_size, 1) # [batch_size, seq_len] # 2. Calculate 1D RoPE frequencies num_rope_elems = 2 # 1 (because 1D) * 2 (for cos, sin) = 2 freqs_dtype = torch.float64 if self.double_precision else torch.float32 pow_indices = torch.pow( self.theta, torch.linspace( start=0.0, end=1.0, steps=self.dim // num_rope_elems, dtype=freqs_dtype, device=device, ), ) freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32) # 3. Matrix-vector outer product between pos ids of shape (batch_size, seq_len) and freqs vector of shape # (self.dim // 2,). freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs # [B, seq_len, self.dim // 2] # 4. Get real, interleaved (cos, sin) frequencies, padded to self.dim if self.rope_type == "interleaved": cos_freqs = freqs.cos().repeat_interleave(2, dim=-1) sin_freqs = freqs.sin().repeat_interleave(2, dim=-1) if self.dim % num_rope_elems != 0: cos_padding = torch.ones_like( cos_freqs[:, :, : self.dim % num_rope_elems] ) sin_padding = torch.zeros_like( sin_freqs[:, :, : self.dim % num_rope_elems] ) cos_freqs = torch.cat([cos_padding, cos_freqs], dim=-1) sin_freqs = torch.cat([sin_padding, sin_freqs], dim=-1) elif self.rope_type == "split": expected_freqs = self.dim // 2 current_freqs = freqs.shape[-1] pad_size = expected_freqs - current_freqs cos_freq = freqs.cos() sin_freq = freqs.sin() if pad_size != 0: cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1) sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1) # Reshape freqs to be compatible with multi-head attention b = cos_freq.shape[0] t = cos_freq.shape[1] cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, -1) sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, -1) cos_freqs = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2) sin_freqs = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2) if dtype is not None: cos_freqs = cos_freqs.to(dtype) sin_freqs = sin_freqs.to(dtype) return cos_freqs, sin_freqs class LTX2TransformerBlock1d(nn.Module): def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, activation_fn: str = "gelu-approximate", eps: float = 1e-6, rope_type: str = "interleaved", ): super().__init__() self.norm1 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False) self.attn1 = LTX2Attention( query_dim=dim, heads=num_attention_heads, kv_heads=num_attention_heads, dim_head=attention_head_dim, rope_type=rope_type, ) self.norm2 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False) self.ff = FeedForward(dim, activation_fn=activation_fn) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, rotary_emb: Optional[torch.Tensor] = None, ) -> torch.Tensor: norm_hidden_states = self.norm1(hidden_states) attn_hidden_states = self.attn1( norm_hidden_states, attention_mask=attention_mask, query_rotary_emb=rotary_emb, ) hidden_states = hidden_states + attn_hidden_states norm_hidden_states = self.norm2(hidden_states) ff_hidden_states = self.ff(norm_hidden_states) hidden_states = hidden_states + ff_hidden_states return hidden_states class LTX2ConnectorTransformer1d(nn.Module): """ A 1D sequence transformer for modalities such as text. In LTX 2.0, this is used to process the text encoder hidden states for each of the video and audio streams. """ _supports_gradient_checkpointing = True def __init__( self, num_attention_heads: int = 30, attention_head_dim: int = 128, num_layers: int = 2, num_learnable_registers: int | None = 128, rope_base_seq_len: int = 4096, rope_theta: float = 10000.0, rope_double_precision: bool = True, eps: float = 1e-6, causal_temporal_positioning: bool = False, rope_type: str = "interleaved", ): super().__init__() self.num_attention_heads = num_attention_heads self.inner_dim = num_attention_heads * attention_head_dim self.causal_temporal_positioning = causal_temporal_positioning self.num_learnable_registers = num_learnable_registers self.learnable_registers = None if num_learnable_registers is not None: init_registers = ( torch.rand(num_learnable_registers, self.inner_dim) * 2.0 - 1.0 ) self.learnable_registers = torch.nn.Parameter(init_registers) self.rope = LTX2RotaryPosEmbed1d( self.inner_dim, base_seq_len=rope_base_seq_len, theta=rope_theta, double_precision=rope_double_precision, rope_type=rope_type, num_attention_heads=num_attention_heads, ) self.transformer_blocks = torch.nn.ModuleList( [ LTX2TransformerBlock1d( dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, rope_type=rope_type, ) for _ in range(num_layers) ] ) self.norm_out = torch.nn.RMSNorm( self.inner_dim, eps=eps, elementwise_affine=False ) self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, attn_mask_binarize_threshold: float = -9000.0, ) -> Tuple[torch.Tensor, torch.Tensor]: # hidden_states shape: [batch_size, seq_len, hidden_dim] # attention_mask shape: [batch_size, seq_len] or [batch_size, 1, 1, seq_len] batch_size, seq_len, _ = hidden_states.shape # 1. Replace padding with learned registers, if using if self.learnable_registers is not None: if seq_len % self.num_learnable_registers != 0: raise ValueError( f"The `hidden_states` sequence length {hidden_states.shape[1]} should be divisible by the number" f" of learnable registers {self.num_learnable_registers}" ) num_register_repeats = seq_len // self.num_learnable_registers registers = torch.tile( self.learnable_registers, (num_register_repeats, 1) ) # [seq_len, inner_dim] binary_attn_mask = (attention_mask >= attn_mask_binarize_threshold).int() if binary_attn_mask.ndim == 4: binary_attn_mask = binary_attn_mask.squeeze(1).squeeze( 1 ) # [B, 1, 1, L] --> [B, L] hidden_states_non_padded = [ hidden_states[i, binary_attn_mask[i].bool(), :] for i in range(batch_size) ] valid_seq_lens = [x.shape[0] for x in hidden_states_non_padded] pad_lengths = [seq_len - valid_seq_len for valid_seq_len in valid_seq_lens] padded_hidden_states = [ F.pad(x, pad=(0, 0, 0, p), value=0) for x, p in zip(hidden_states_non_padded, pad_lengths) ] padded_hidden_states = torch.cat( [x.unsqueeze(0) for x in padded_hidden_states], dim=0 ) # [B, L, D] flipped_mask = torch.flip(binary_attn_mask, dims=[1]).unsqueeze( -1 ) # [B, L, 1] hidden_states = ( flipped_mask * padded_hidden_states + (1 - flipped_mask) * registers ) # Overwrite attention_mask with an all-zeros mask if using registers. attention_mask = torch.zeros_like(attention_mask) # 2. Calculate 1D RoPE positional embeddings rotary_emb = self.rope( batch_size, seq_len, device=hidden_states.device, dtype=hidden_states.dtype ) # 3. Run 1D transformer blocks for block in self.transformer_blocks: if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( block, hidden_states, attention_mask, rotary_emb ) else: hidden_states = block( hidden_states, attention_mask=attention_mask, rotary_emb=rotary_emb ) hidden_states = self.norm_out(hidden_states) return hidden_states, attention_mask class LTX2TextConnectors(nn.Module): """ Text connector stack used by LTX 2.0 to process the packed text encoder hidden states for both the video and audio streams. """ def __init__( self, config: LTX2ConnectorConfig, ): super().__init__() caption_channels = config.caption_channels text_proj_in_factor = config.text_proj_in_factor video_connector_num_attention_heads = config.video_connector_num_attention_heads video_connector_attention_head_dim = config.video_connector_attention_head_dim video_connector_num_layers = config.video_connector_num_layers video_connector_num_learnable_registers = ( config.video_connector_num_learnable_registers ) audio_connector_num_attention_heads = config.audio_connector_num_attention_heads audio_connector_attention_head_dim = config.audio_connector_attention_head_dim audio_connector_num_layers = config.audio_connector_num_layers audio_connector_num_learnable_registers = ( config.audio_connector_num_learnable_registers ) connector_rope_base_seq_len = config.connector_rope_base_seq_len rope_theta = config.rope_theta rope_double_precision = config.rope_double_precision causal_temporal_positioning = config.causal_temporal_positioning rope_type = config.rope_type self.text_proj_in = nn.Linear( caption_channels * text_proj_in_factor, caption_channels, bias=False ) self.video_connector = LTX2ConnectorTransformer1d( num_attention_heads=video_connector_num_attention_heads, attention_head_dim=video_connector_attention_head_dim, num_layers=video_connector_num_layers, num_learnable_registers=video_connector_num_learnable_registers, rope_base_seq_len=connector_rope_base_seq_len, rope_theta=rope_theta, rope_double_precision=rope_double_precision, causal_temporal_positioning=causal_temporal_positioning, rope_type=rope_type, ) self.audio_connector = LTX2ConnectorTransformer1d( num_attention_heads=audio_connector_num_attention_heads, attention_head_dim=audio_connector_attention_head_dim, num_layers=audio_connector_num_layers, num_learnable_registers=audio_connector_num_learnable_registers, rope_base_seq_len=connector_rope_base_seq_len, rope_theta=rope_theta, rope_double_precision=rope_double_precision, causal_temporal_positioning=causal_temporal_positioning, rope_type=rope_type, ) def forward( self, text_encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor, additive_mask: bool = False, ): # Convert to additive attention mask, if necessary if not additive_mask: text_dtype = text_encoder_hidden_states.dtype attention_mask = (attention_mask - 1).reshape( attention_mask.shape[0], 1, -1, attention_mask.shape[-1] ) attention_mask = attention_mask.to(text_dtype) * torch.finfo(text_dtype).max # Ensure input dtype matches the layer's weight dtype if text_encoder_hidden_states.dtype != self.text_proj_in.weight.dtype: text_encoder_hidden_states = text_encoder_hidden_states.to( self.text_proj_in.weight.dtype ) # Ensure sequence length is divisible by num_learnable_registers (128) seq_len = text_encoder_hidden_states.shape[1] num_learnable_registers = self.video_connector.num_learnable_registers if ( num_learnable_registers is not None and seq_len % num_learnable_registers != 0 ): pad_len = num_learnable_registers - (seq_len % num_learnable_registers) text_encoder_hidden_states = F.pad( text_encoder_hidden_states, (0, 0, 0, pad_len), value=0.0 ) if attention_mask.shape[-1] == seq_len: # Pad with a large negative value to mask out the new tokens attention_mask = F.pad(attention_mask, (0, pad_len), value=-1000000.0) text_encoder_hidden_states = self.text_proj_in(text_encoder_hidden_states) video_text_embedding, new_attn_mask = self.video_connector( text_encoder_hidden_states, attention_mask ) attn_mask = (new_attn_mask < 1e-6).to(torch.int64) attn_mask = attn_mask.reshape( video_text_embedding.shape[0], video_text_embedding.shape[1], 1 ) video_text_embedding = video_text_embedding * attn_mask new_attn_mask = attn_mask.squeeze(-1) audio_text_embedding, _ = self.audio_connector( text_encoder_hidden_states, attention_mask ) return video_text_embedding, audio_text_embedding, new_attn_mask EntryClass = LTX2TextConnectors ================================================ FILE: python/sglang/multimodal_gen/runtime/models/bridges/__init__.py ================================================ # SPDX-License-Identifier: Apache-2.0 from sglang.multimodal_gen.runtime.models.bridges.mova_dual_tower import ( DualTowerConditionalBridge, ) __all__ = ["DualTowerConditionalBridge"] ================================================ FILE: python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py ================================================ # SPDX-License-Identifier: Apache-2.0 # Copied and adapted from: mossVG/mova/diffusion/models/interactionv2.py from typing import Any, Dict, List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import ( MOVADualTowerConfig, ) from sglang.multimodal_gen.runtime.distributed import get_tp_world_size from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.layernorm import ( RMSNorm, tensor_parallel_rms_norm, ) from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, ReplicatedLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( apply_flashinfer_rope_qk_inplace, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @torch.no_grad() def compute_rope_cos_sin( position_ids: torch.Tensor, head_dim: int, base: float = 10000.0, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Compute RoPE cos/sin embeddings for given position IDs. This is a functional implementation that doesn't require storing buffers, making it compatible with FSDP meta device initialization. Args: position_ids: Position IDs tensor [B, L] or [1, L] head_dim: Dimension of each attention head base: RoPE base frequency (default: 10000.0) device: Target device dtype: Output dtype Returns: (cos, sin): Each with shape [B, L, head_dim] """ device = device or position_ids.device dtype = dtype or torch.float32 # Compute inverse frequencies inv_freq = 1.0 / ( base ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim) ) # Expand for batch computation: [B, L] -> [B, 1, L] @ [1, head_dim/2, 1] -> [B, head_dim/2, L] inv_freq_expanded = inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1) position_ids_expanded = position_ids[:, None, :].float() # Compute frequencies: [B, head_dim/2, L] -> [B, L, head_dim/2] freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) # Double the frequencies for full head_dim: [B, L, head_dim] emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos().to(dtype=dtype) sin = emb.sin().to(dtype=dtype) return cos, sin class PerFrameAttentionPooling(nn.Module): """Per-frame multi-head attention pooling. Flattens the input sequence [B, L, D] and grid size (T, H, W). Performs single-query attention pooling on the H*W tokens for each time frame. Output shape: [B, T, D]. """ def __init__(self, dim: int, num_heads: int, eps: float = 1e-6): super().__init__() assert dim % num_heads == 0, "dim must be divisible by num_heads" self.dim = dim self.num_heads = num_heads self.probe = nn.Parameter(torch.randn(1, 1, dim)) nn.init.normal_(self.probe, std=0.02) self.attention = nn.MultiheadAttention( embed_dim=dim, num_heads=num_heads, batch_first=True ) self.layernorm = nn.LayerNorm(dim, eps=eps) def forward(self, x: torch.Tensor, grid_size: Tuple[int, int, int]) -> torch.Tensor: """Forward pass. Args: x: Input tensor of shape [B, L, D], where L = T * H * W. grid_size: Tuple of (T, H, W). Returns: Pooled tensor of shape [B, T, D]. """ B, L, D = x.shape T, H, W = grid_size assert ( D == self.dim ), f"Input dimension D={D} does not match module dim={self.dim}" assert L == T * H * W, f"Flattened length L={L} does not match T*H*W={T*H*W}" S = H * W x_bt_s_d = x.view(B, T, S, D).contiguous().view(B * T, S, D) # [B*T, S, D] probe = self.probe.expand(B * T, -1, -1) # [B*T, 1, D] pooled_bt_1_d = self.attention(probe, x_bt_s_d, x_bt_s_d, need_weights=False)[0] pooled_bt_d = pooled_bt_1_d.squeeze(1) # [B*T, D] pooled = pooled_bt_d.view(B, T, D) pooled = self.layernorm(pooled) return pooled class CrossModalInteractionController: """Strategy class to control dual-tower interaction. Manages the interaction mapping between Visual DiT (e.g., 30 layers) and Audio DiT (e.g., 30 layers). """ def __init__(self, visual_layers: int = 30, audio_layers: int = 30): self.visual_layers = visual_layers self.audio_layers = audio_layers self.min_layers = min(visual_layers, audio_layers) def get_interaction_layers( self, strategy: str = "shallow_focus" ) -> Dict[str, List[Tuple[int, int]]]: """Gets the mapping relationship of interaction layers.""" if strategy == "shallow_focus": num_interact = min(10, self.min_layers // 3) interact_layers = list(range(0, num_interact)) elif strategy == "distributed": step = 3 interact_layers = list(range(0, self.min_layers, step)) elif strategy == "progressive": shallow = list(range(0, min(8, self.min_layers))) if self.min_layers > 8: deep = list(range(8, self.min_layers, 3)) interact_layers = shallow + deep else: interact_layers = shallow elif strategy == "custom": interact_layers = [0, 2, 4, 6, 8, 12, 16, 20] interact_layers = [i for i in interact_layers if i < self.min_layers] elif strategy == "full": interact_layers = list(range(0, self.min_layers)) else: raise ValueError(f"Unknown interaction strategy: {strategy}") mapping = { "v2a": [(i, i) for i in interact_layers], "a2v": [(i, i) for i in interact_layers], } return mapping def should_interact( self, layer_idx: int, direction: str, interaction_mapping: Dict ) -> bool: """Determines if the specified layer needs to interact.""" if direction not in interaction_mapping: return False return any(src == layer_idx for src, _ in interaction_mapping[direction]) class ConditionalCrossAttention(nn.Module): """ Cross-modal attention for dual-tower bridge with Tensor Parallel support. This module handles attention between video and audio hidden states, which have different sequence lengths. """ def __init__(self, dim: int, kv_dim: int, num_heads: int, eps: float = 1e-6): super().__init__() self.q_dim = dim self.kv_dim = kv_dim self.num_heads = num_heads self.head_dim = self.q_dim // num_heads self.tp_size = get_tp_world_size() if self.num_heads % self.tp_size != 0: raise ValueError( f"num_heads ({self.num_heads}) must be divisible by tp_size ({self.tp_size})." ) self.num_heads_per_rank = self.num_heads // self.tp_size # TP strategy: shard Q/K/V over heads (column-parallel), then row-parallel output. self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False) self.k = ColumnParallelLinear(kv_dim, dim, bias=True, gather_output=False) self.v = ColumnParallelLinear(kv_dim, dim, bias=True, gather_output=False) self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True) self.norm_q = RMSNorm(dim, eps=eps) self.norm_k = RMSNorm(dim, eps=eps) self.attn = USPAttention( num_heads=self.num_heads_per_rank, head_size=self.head_dim, causal=False, softmax_scale=None, # is_cross_attention=True, ) def forward( self, x: torch.Tensor, y: torch.Tensor, x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ): ctx = y q, _ = self.q(x) k, _ = self.k(ctx) v, _ = self.v(ctx) # RMSNorm over sharded hidden dimension if self.tp_size > 1: q = tensor_parallel_rms_norm(q, self.norm_q) k = tensor_parallel_rms_norm(k, self.norm_k) else: q = self.norm_q(q) k = self.norm_k(k) if x_freqs is not None: x_cos, x_sin = x_freqs q_view = rearrange(q, "b l (h d) -> b l h d", d=self.head_dim) x_cos = x_cos.to(q_view.dtype).to(q_view.device).squeeze(0) x_sin = x_sin.to(q_view.dtype).to(q_view.device).squeeze(0) # FlashInfer expects cos_sin_cache with shape [seqlen, head_dim], # where the first half is cos and the second half is sin, each with # head_dim//2 elements. Since compute_rope_cos_sin duplicates the # frequencies (cat((freqs, freqs))), we only take the first half. half_dim = self.head_dim // 2 cos_sin_cache = torch.cat( [ x_cos[:, :half_dim].to(dtype=torch.float32).contiguous(), x_sin[:, :half_dim].to(dtype=torch.float32).contiguous(), ], dim=-1, ) q_view, _ = apply_flashinfer_rope_qk_inplace( q_view, q_view.clone(), cos_sin_cache, is_neox=True ) q = rearrange(q_view, "b l h d -> b l (h d)") if y_freqs is not None: y_cos, y_sin = y_freqs k_view = rearrange(k, "b l (h d) -> b l h d", d=self.head_dim) y_cos = y_cos.to(k_view.dtype).to(k_view.device).squeeze(0) y_sin = y_sin.to(k_view.dtype).to(k_view.device).squeeze(0) # FlashInfer expects cos_sin_cache with shape [seqlen, head_dim], # where the first half is cos and the second half is sin, each with # head_dim//2 elements. Since compute_rope_cos_sin duplicates the # frequencies (cat((freqs, freqs))), we only take the first half. half_dim = self.head_dim // 2 cos_sin_cache = torch.cat( [ y_cos[:, :half_dim].to(dtype=torch.float32).contiguous(), y_sin[:, :half_dim].to(dtype=torch.float32).contiguous(), ], dim=-1, ) k_view, _ = apply_flashinfer_rope_qk_inplace( k_view, k_view.clone(), cos_sin_cache, is_neox=True ) k = rearrange(k_view, "b l h d -> b l (h d)") q = rearrange(q, "b l (h d) -> b l h d", h=self.num_heads_per_rank) k = rearrange(k, "b l (h d) -> b l h d", h=self.num_heads_per_rank) v = rearrange(v, "b l (h d) -> b l h d", h=self.num_heads_per_rank) x = self.attn(q, k, v) x = rearrange(x, "b l h d -> b l (h d)") x, _ = self.o(x) return x class AdaLayerNorm(nn.Module): """ Norm layer modified to incorporate timestep embeddings. """ def __init__( self, embedding_dim: int, num_embeddings: Optional[int] = None, output_dim: Optional[int] = None, norm_elementwise_affine: bool = False, norm_eps: float = 1e-5, chunk_dim: int = 0, ): super().__init__() self.chunk_dim = chunk_dim output_dim = output_dim or embedding_dim * 2 if num_embeddings is not None: self.emb = nn.Embedding(num_embeddings, embedding_dim) else: self.emb = None self.silu = nn.SiLU() self.linear = ReplicatedLinear(embedding_dim, output_dim) self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine) def forward( self, x: torch.Tensor, timestep: Optional[torch.Tensor] = None, temb: Optional[torch.Tensor] = None, ) -> torch.Tensor: if self.emb is not None: temb = self.emb(timestep) temb, _ = self.linear(self.silu(temb)) if self.chunk_dim == 2: scale, shift = temb.chunk(2, dim=2) elif self.chunk_dim == 1: shift, scale = temb.chunk(2, dim=1) shift = shift[:, None, :] scale = scale[:, None, :] else: scale, shift = temb.chunk(2, dim=0) x = self.norm(x) * (1 + scale) + shift return x class ConditionalCrossAttentionBlock(nn.Module): """A wrapper block for ConditionalCrossAttention that applies LayerNorm to the condition input y.""" def __init__( self, dim: int, kv_dim: int, num_heads: int, eps: float = 1e-6, pooled_adaln: bool = False, ): super().__init__() self.y_norm = nn.LayerNorm(kv_dim, eps=eps) self.inner = ConditionalCrossAttention( dim=dim, kv_dim=kv_dim, num_heads=num_heads, eps=eps ) self.pooled_adaln = pooled_adaln if pooled_adaln: self.per_frame_pooling = PerFrameAttentionPooling( kv_dim, num_heads=num_heads, eps=eps ) self.adaln = AdaLayerNorm(kv_dim, output_dim=dim * 2, chunk_dim=2) def forward( self, x: torch.Tensor, y: torch.Tensor, x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, video_grid_size: Optional[Tuple[int, int, int]] = None, ) -> torch.Tensor: if self.pooled_adaln: assert video_grid_size is not None, "video_grid_size cannot be None" pooled_y = self.per_frame_pooling(y, video_grid_size) if pooled_y.shape[1] != x.shape[1]: pooled_y = F.interpolate( pooled_y.permute(0, 2, 1), size=x.shape[1], mode="linear", align_corners=False, ).permute(0, 2, 1) x = self.adaln(x, temb=pooled_y) y = self.y_norm(y) return self.inner(x=x, y=y, x_freqs=x_freqs, y_freqs=y_freqs) class DualTowerConditionalBridge( CachableDiT, OffloadableDiTMixin, ): """Dual-tower conditional bridge module v2 (SGLang optimized version). Implements the correct architecture: 1. Audio latents -> Audio DiT -> Audio hidden states [B, L, 1536]. 2. Visual latents -> Visual DiT -> Visual hidden states [B, L, 5120]. 3. Cross-attention interaction between the hidden states of the two DiTs. """ _fsdp_shard_conditions = MOVADualTowerConfig()._fsdp_shard_conditions _compile_conditions = MOVADualTowerConfig()._compile_conditions _supported_attention_backends = MOVADualTowerConfig()._supported_attention_backends param_names_mapping = MOVADualTowerConfig().param_names_mapping reverse_param_names_mapping = MOVADualTowerConfig().reverse_param_names_mapping lora_param_names_mapping = MOVADualTowerConfig().lora_param_names_mapping def __init__( self, config: MOVADualTowerConfig | None = None, hf_config: dict[str, Any] | None = None, # Fallback parameters for from_pretrained compatibility visual_layers: int = 40, audio_layers: int = 30, visual_hidden_dim: int = 5120, audio_hidden_dim: int = 1536, audio_fps: float = 50.0, head_dim: int = 128, interaction_strategy: str = "full", apply_cross_rope: bool = True, apply_first_frame_bias_in_rope: bool = False, trainable_condition_scale: bool = False, pooled_adaln: bool = False, ): super().__init__(config=config, hf_config=hf_config) # Use config if provided, otherwise use individual parameters if config is not None: visual_layers = config.visual_layers audio_layers = config.audio_layers visual_hidden_dim = config.visual_hidden_dim audio_hidden_dim = config.audio_hidden_dim audio_fps = config.audio_fps head_dim = config.head_dim interaction_strategy = config.interaction_strategy apply_cross_rope = config.apply_cross_rope apply_first_frame_bias_in_rope = config.apply_first_frame_bias_in_rope trainable_condition_scale = config.trainable_condition_scale pooled_adaln = config.pooled_adaln self.visual_hidden_dim = visual_hidden_dim self.audio_hidden_dim = audio_hidden_dim self.audio_fps = audio_fps self.head_dim = head_dim self.apply_cross_rope = apply_cross_rope self.apply_first_frame_bias_in_rope = apply_first_frame_bias_in_rope self.trainable_condition_scale = trainable_condition_scale self.pooled_adaln = pooled_adaln if self.trainable_condition_scale: self.condition_scale = nn.Parameter( torch.tensor([1.0], dtype=torch.float32) ) else: self.condition_scale = 1.0 self.controller = CrossModalInteractionController(visual_layers, audio_layers) self.interaction_mapping = self.controller.get_interaction_layers( interaction_strategy ) # Cross-modal attention modules - interaction at DiT hidden states level self.audio_to_video_conditioners = nn.ModuleDict() self.video_to_audio_conditioners = nn.ModuleDict() self.rope_base = 10000.0 # RoPE base frequency hardcode. adapted from original mova implementation. # Audio DiT hidden states conditioning Video DiT for v_layer, _ in self.interaction_mapping["a2v"]: self.audio_to_video_conditioners[str(v_layer)] = ( ConditionalCrossAttentionBlock( dim=visual_hidden_dim, kv_dim=audio_hidden_dim, num_heads=visual_hidden_dim // head_dim, pooled_adaln=False, ) ) # Visual DiT hidden states conditioning Audio DiT for a_layer, _ in self.interaction_mapping["v2a"]: self.video_to_audio_conditioners[str(a_layer)] = ( ConditionalCrossAttentionBlock( dim=audio_hidden_dim, kv_dim=visual_hidden_dim, num_heads=audio_hidden_dim // head_dim, pooled_adaln=self.pooled_adaln, ) ) # Required attributes for CachableDiT/BaseDiT self.hidden_size = visual_hidden_dim self.num_attention_heads = visual_hidden_dim // head_dim self.num_channels_latents = ( visual_hidden_dim # Bridge doesn't output latents, but required by BaseDiT ) self.layer_names = [ "audio_to_video_conditioners", "video_to_audio_conditioners", ] self.__post_init__() @torch.no_grad() def build_aligned_freqs( self, video_fps: float, grid_size: Tuple[int, int, int], audio_steps: int, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: """Generates aligned RoPE (cos, sin) based on video FPS, grid size, and audio length. Uses functional RoPE computation to avoid FSDP meta device issues. Args: video_fps: FPS of the video. grid_size: Tuple of (f_v, h, w). audio_steps: Length of the audio sequence. device: Target device. dtype: Output dtype. Returns: A tuple of ((cos_v, sin_v), (cos_a, sin_a)). """ f_v, h, w = grid_size L_v = f_v * h * w L_a = int(audio_steps) device = device or next(self.parameters()).device dtype = dtype or torch.float32 # Audio positions: 0, 1, 2, ..., L_a-1 audio_pos = torch.arange(L_a, device=device, dtype=torch.float32).unsqueeze(0) # Video positions: Align video frames to audio step units if self.apply_first_frame_bias_in_rope: video_effective_fps = float(video_fps) / 4.0 if f_v > 0: t_starts = torch.zeros((f_v,), device=device, dtype=torch.float32) if f_v > 1: t_starts[1:] = (1.0 / float(video_fps)) + torch.arange( f_v - 1, device=device, dtype=torch.float32 ) * (1.0 / video_effective_fps) else: t_starts = torch.zeros((0,), device=device, dtype=torch.float32) video_pos_per_frame = t_starts * float(self.audio_fps) else: scale = float(self.audio_fps) / float(video_fps / 4.0) video_pos_per_frame = ( torch.arange(f_v, device=device, dtype=torch.float32) * scale ) video_pos = video_pos_per_frame.repeat_interleave(h * w).unsqueeze(0) # Use functional RoPE to compute cos/sin cos_v, sin_v = compute_rope_cos_sin( video_pos, self.head_dim, base=self.rope_base, device=device, dtype=dtype ) cos_a, sin_a = compute_rope_cos_sin( audio_pos, self.head_dim, base=self.rope_base, device=device, dtype=dtype ) return (cos_v, sin_v), (cos_a, sin_a) def should_interact(self, layer_idx: int, direction: str) -> bool: return self.controller.should_interact( layer_idx, direction, self.interaction_mapping ) def apply_conditional_control( self, layer_idx: int, direction: str, primary_hidden_states: torch.Tensor, condition_hidden_states: torch.Tensor, x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, condition_scale: Optional[float] = None, video_grid_size: Optional[Tuple[int, int, int]] = None, ) -> torch.Tensor: """Applies conditional control at the DiT hidden states level.""" if not self.controller.should_interact( layer_idx, direction, self.interaction_mapping ): return primary_hidden_states if direction == "a2v": conditioner = self.audio_to_video_conditioners[str(layer_idx)] elif direction == "v2a": conditioner = self.video_to_audio_conditioners[str(layer_idx)] else: raise ValueError(f"Invalid direction: {direction}") conditioned_features = conditioner( x=primary_hidden_states, y=condition_hidden_states, x_freqs=x_freqs, y_freqs=y_freqs, video_grid_size=video_grid_size, ) if self.trainable_condition_scale and condition_scale is not None: logger.warning( "The current model has a trainable condition_scale, but condition_scale " "was passed externally. Ignoring the trainable condition_scale and " "using the external condition_scale=%s.", condition_scale, ) scale = condition_scale if condition_scale is not None else self.condition_scale primary_hidden_states = primary_hidden_states + conditioned_features * scale return primary_hidden_states def forward( self, layer_idx: int, visual_hidden_states: torch.Tensor, audio_hidden_states: torch.Tensor, *, x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, a2v_condition_scale: Optional[float] = None, v2a_condition_scale: Optional[float] = None, condition_scale: Optional[float] = None, video_grid_size: Optional[Tuple[int, int, int]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Performs bidirectional conditional control for both visual and audio towers.""" visual_conditioned = self.apply_conditional_control( layer_idx=layer_idx, direction="a2v", primary_hidden_states=visual_hidden_states, condition_hidden_states=audio_hidden_states, x_freqs=x_freqs, y_freqs=y_freqs, condition_scale=( a2v_condition_scale if a2v_condition_scale is not None else condition_scale ), video_grid_size=video_grid_size, ) audio_conditioned = self.apply_conditional_control( layer_idx=layer_idx, direction="v2a", primary_hidden_states=audio_hidden_states, condition_hidden_states=visual_hidden_states, x_freqs=y_freqs, y_freqs=x_freqs, condition_scale=( v2a_condition_scale if v2a_condition_scale is not None else condition_scale ), video_grid_size=video_grid_size, ) return visual_conditioned, audio_conditioned EntryClass = DualTowerConditionalBridge ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/base.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from abc import ABC, abstractmethod from typing import Any import torch from torch import nn from sglang.multimodal_gen.configs.models import DiTConfig # NOTE: TeaCacheContext and TeaCacheMixin have been moved to # sglang.multimodal_gen.runtime.cache.teacache # For backwards compatibility, re-export from the new location from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheContext # noqa: F401 from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheMixin from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum # TODO class BaseDiT(nn.Module, ABC): _fsdp_shard_conditions: list = [] _compile_conditions: list = [] param_names_mapping: dict reverse_param_names_mapping: dict hidden_size: int num_attention_heads: int num_channels_latents: int # always supports torch_sdpa _supported_attention_backends: set[AttentionBackendEnum] = ( DiTConfig()._supported_attention_backends ) def __init_subclass__(cls) -> None: required_class_attrs = [ "_fsdp_shard_conditions", "param_names_mapping", "_compile_conditions", ] super().__init_subclass__() for attr in required_class_attrs: if not hasattr(cls, attr): raise AttributeError( f"Subclasses of BaseDiT must define '{attr}' class variable" ) def __init__(self, config: DiTConfig, hf_config: dict[str, Any], **kwargs) -> None: super().__init__() self.config = config self.hf_config = hf_config if not self.supported_attention_backends: raise ValueError( f"Subclass {self.__class__.__name__} must define _supported_attention_backends" ) @abstractmethod def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, guidance=None, **kwargs, ) -> torch.Tensor: pass def __post_init__(self) -> None: required_attrs = ["hidden_size", "num_attention_heads", "num_channels_latents"] for attr in required_attrs: if not hasattr(self, attr): raise AttributeError( f"Subclasses of BaseDiT must define '{attr}' instance variable" ) @property def supported_attention_backends(self) -> set[AttentionBackendEnum]: return self._supported_attention_backends @property def device(self) -> torch.device: """Get the device of the model.""" return next(self.parameters()).device class CachableDiT(TeaCacheMixin, BaseDiT): """ An intermediate base class that adds TeaCache optimization functionality to DiT models. Inherits TeaCacheMixin for cache logic and BaseDiT for core DiT functionality. """ # These are required class attributes that should be overridden by concrete implementations _fsdp_shard_conditions = [] param_names_mapping = {} reverse_param_names_mapping = {} lora_param_names_mapping: dict = {} # Ensure these instance attributes are properly defined in subclasses hidden_size: int num_attention_heads: int num_channels_latents: int # always supports torch_sdpa _supported_attention_backends: set[AttentionBackendEnum] = ( DiTConfig()._supported_attention_backends ) def __init__(self, config: DiTConfig, **kwargs) -> None: super().__init__(config, **kwargs) self._init_teacache_state() @classmethod def get_nunchaku_quant_rules(cls) -> dict[str, dict[str, Any]]: """ Get quantization rules for Nunchaku quantization. Returns a dict mapping layer name patterns to quantization configs: { "skip": [list of patterns to skip quantization], "svdq_w4a4": [list of patterns for SVDQ W4A4], "awq_w4a16": [list of patterns for AWQ W4A16], } """ return {} ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import math from typing import Any import torch import torch.nn as nn from torch.nn.attention.flex_attention import ( BlockMask, create_block_mask, flex_attention, ) from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin # wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention # see https://github.com/pytorch/pytorch/issues/133254 # change to default for other models flex_attention = torch.compile( flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs" ) import torch.distributed as dist from sglang.multimodal_gen.configs.models.dits import WanVideoConfig from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size from sglang.multimodal_gen.runtime.layers.attention import LocalAttention from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd from sglang.multimodal_gen.runtime.layers.layernorm import ( FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidualLayerNormScaleShift, ) from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( _apply_rotary_emb, get_rotary_pos_embed, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import PatchEmbed from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT from sglang.multimodal_gen.runtime.models.dits.wanvideo import ( WanT2VCrossAttention, WanTimeTextImageEmbedding, ) from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class CausalWanSelfAttention(nn.Module): def __init__( self, dim: int, num_heads: int, local_attn_size: int = -1, sink_size: int = 0, qk_norm=True, eps=1e-6, parallel_attention=False, ) -> None: assert dim % num_heads == 0 super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads self.local_attn_size = local_attn_size self.sink_size = sink_size self.qk_norm = qk_norm self.eps = eps self.parallel_attention = parallel_attention self.max_attention_size = ( 32760 if local_attn_size == -1 else local_attn_size * 1560 ) # Scaled dot product attention self.attn = LocalAttention( num_heads=num_heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, causal=False, supported_attention_backends=( AttentionBackendEnum.FA, AttentionBackendEnum.AITER, AttentionBackendEnum.TORCH_SDPA, ), ) def forward( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], block_mask: BlockMask, kv_cache: dict | None = None, current_start: int = 0, cache_start: int | None = None, ): r""" Args: x(Tensor): Shape [B, L, num_heads, C / num_heads] seq_lens(Tensor): Shape [B] grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] """ if cache_start is None: cache_start = current_start cos, sin = freqs_cis roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v) roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v) if kv_cache is None: # Padding for flex attention padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] padded_roped_query = torch.cat( [ roped_query, torch.zeros( [q.shape[0], padded_length, q.shape[2], q.shape[3]], device=q.device, dtype=v.dtype, ), ], dim=1, ) padded_roped_key = torch.cat( [ roped_key, torch.zeros( [k.shape[0], padded_length, k.shape[2], k.shape[3]], device=k.device, dtype=v.dtype, ), ], dim=1, ) padded_v = torch.cat( [ v, torch.zeros( [v.shape[0], padded_length, v.shape[2], v.shape[3]], device=v.device, dtype=v.dtype, ), ], dim=1, ) x = flex_attention( query=padded_roped_query.transpose(2, 1), key=padded_roped_key.transpose(2, 1), value=padded_v.transpose(2, 1), block_mask=block_mask, )[:, :, :-padded_length].transpose(2, 1) else: frame_seqlen = q.shape[1] current_end = current_start + roped_query.shape[1] sink_tokens = self.sink_size * frame_seqlen # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache kv_cache_size = kv_cache["k"].shape[1] num_new_tokens = roped_query.shape[1] if ( self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and ( num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size ) ): # Calculate the number of new tokens added in this step # Shift existing cache content left to discard oldest tokens # Clone the source slice to avoid overlapping memory error num_evicted_tokens = ( num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size ) num_rolled_tokens = ( kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens ) kv_cache["k"][ :, sink_tokens : sink_tokens + num_rolled_tokens ] = kv_cache["k"][ :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens, ].clone() kv_cache["v"][ :, sink_tokens : sink_tokens + num_rolled_tokens ] = kv_cache["v"][ :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens, ].clone() # Insert the new keys/values at the end local_end_index = ( kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() - num_evicted_tokens ) local_start_index = local_end_index - num_new_tokens kv_cache["k"][:, local_start_index:local_end_index] = roped_key kv_cache["v"][:, local_start_index:local_end_index] = v else: # Assign new keys/values directly up to current_end local_end_index = ( kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() ) local_start_index = local_end_index - num_new_tokens kv_cache["k"] = kv_cache["k"].detach() kv_cache["v"] = kv_cache["v"].detach() # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None) kv_cache["k"][:, local_start_index:local_end_index] = roped_key kv_cache["v"][:, local_start_index:local_end_index] = v x = self.attn( roped_query, kv_cache["k"][ :, max(0, local_end_index - self.max_attention_size) : local_end_index, ], kv_cache["v"][ :, max(0, local_end_index - self.max_attention_size) : local_end_index, ], ) kv_cache["global_end_index"].fill_(current_end) kv_cache["local_end_index"].fill_(local_end_index) return x class CausalWanTransformerBlock(nn.Module): def __init__( self, dim: int, ffn_dim: int, num_heads: int, local_attn_size: int = -1, sink_size: int = 0, qk_norm: str = "rms_norm_across_heads", cross_attn_norm: bool = False, eps: float = 1e-6, added_kv_proj_dim: int | None = None, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", quant_config: QuantizationConfig | None = None, ): super().__init__() # 1. Self-attention self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config) self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config) self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config) self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config) self.attn1 = CausalWanSelfAttention( dim, num_heads, local_attn_size=local_attn_size, sink_size=sink_size, qk_norm=qk_norm, eps=eps, ) self.hidden_dim = dim self.num_attention_heads = num_heads self.local_attn_size = local_attn_size dim_head = dim // num_heads if qk_norm == "rms_norm": self.norm_q = RMSNorm(dim_head, eps=eps) self.norm_k = RMSNorm(dim_head, eps=eps) elif qk_norm == "rms_norm_across_heads": # LTX applies qk norm across all heads self.norm_q = RMSNorm(dim, eps=eps) self.norm_k = RMSNorm(dim, eps=eps) else: print("QK Norm type not supported") raise Exception assert cross_attn_norm is True self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, eps=eps, elementwise_affine=True, dtype=torch.float32 ) # 2. Cross-attention # Only T2V for now cross_attn_backends = { b for b in supported_attention_backends if not b.is_sparse } self.attn2 = WanT2VCrossAttention( dim, num_heads, qk_norm=qk_norm, eps=eps, supported_attention_backends=cross_attn_backends, quant_config=quant_config, ) self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, eps=eps, elementwise_affine=False, dtype=torch.float32 ) # 3. Feed-forward self.ffn = MLP( dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config ) self.mlp_residual = MulAdd() self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], block_mask: BlockMask, kv_cache: dict | None = None, crossattn_cache: dict | None = None, current_start: int = 0, cache_start: int | None = None, ) -> torch.Tensor: # hidden_states.shape: [batch_size, seq_length, inner_dim] # temb.shape: [batch_size, num_frames, 6, inner_dim] if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) num_frames = temb.shape[1] frame_seqlen = hidden_states.shape[1] // num_frames bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype # assert orig_dtype != torch.float32 e = self.scale_shift_table + temb.float() # e.shape: [batch_size, num_frames, 6, inner_dim] assert e.shape == (bs, num_frames, 6, self.hidden_dim) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( 6, dim=2 ) # *_msa.shape: [batch_size, num_frames, 1, inner_dim] assert shift_msa.dtype == torch.float32 # 1. Self-attention norm_hidden_states = ( ( self.norm1(hidden_states.float()).unflatten( dim=1, sizes=(num_frames, frame_seqlen) ) * (1 + scale_msa) + shift_msa ) .flatten(1, 2) .to(orig_dtype) ) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) if self.norm_q is not None: query = self.norm_q(query) if self.norm_k is not None: key = self.norm_k(key) query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) attn_output = self.attn1( query, key, value, freqs_cis, block_mask, kv_cache, current_start, cache_start, ) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.zeroes( (1,), device=hidden_states.device, dtype=hidden_states.dtype ) norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale ) norm_hidden_states, hidden_states = norm_hidden_states.to( orig_dtype ), hidden_states.to(orig_dtype) # 2. Cross-attention attn_output = self.attn2( norm_hidden_states, context=encoder_hidden_states, context_lens=None, crossattn_cache=crossattn_cache, ) norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) norm_hidden_states, hidden_states = norm_hidden_states.to( orig_dtype ), hidden_states.to(orig_dtype) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states) hidden_states = hidden_states.to(orig_dtype) return hidden_states class CausalWanTransformer3DModel(BaseDiT, OffloadableDiTMixin): _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions _compile_conditions = WanVideoConfig()._compile_conditions _supported_attention_backends = WanVideoConfig()._supported_attention_backends param_names_mapping = WanVideoConfig().param_names_mapping reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping def __init__( self, config: WanVideoConfig, hf_config: dict[str, Any], quant_config: QuantizationConfig | None = None, ) -> None: super().__init__(config=config, hf_config=hf_config) inner_dim = config.num_attention_heads * config.attention_head_dim self.hidden_size = config.hidden_size self.num_attention_heads = config.num_attention_heads self.attention_head_dim = config.attention_head_dim self.in_channels = config.in_channels self.out_channels = config.out_channels self.num_channels_latents = config.num_channels_latents self.patch_size = config.patch_size self.text_len = config.text_len self.local_attn_size = config.local_attn_size # 1. Patch & position embedding self.patch_embedding = PatchEmbed( in_chans=config.in_channels, embed_dim=inner_dim, patch_size=config.patch_size, flatten=False, ) # 2. Condition embeddings self.condition_embedder = WanTimeTextImageEmbedding( dim=inner_dim, time_freq_dim=config.freq_dim, text_embed_dim=config.text_dim, image_embed_dim=config.image_dim, ) # 3. Transformer blocks self.blocks = nn.ModuleList( [ CausalWanTransformerBlock( inner_dim, config.ffn_dim, config.num_attention_heads, config.local_attn_size, config.sink_size, config.qk_norm, config.cross_attn_norm, config.eps, config.added_kv_proj_dim, self._supported_attention_backends, prefix=f"{config.prefix}.blocks.{i}", quant_config=quant_config, ) for i in range(config.num_layers) ] ) # 4. Output norm & projection self.norm_out = LayerNormScaleShift( inner_dim, eps=config.eps, elementwise_affine=False, dtype=torch.float32, ) self.proj_out = nn.Linear( inner_dim, config.out_channels * math.prod(config.patch_size) ) self.scale_shift_table = nn.Parameter( torch.randn(1, 2, inner_dim) / inner_dim**0.5 ) self.gradient_checkpointing = False # Causal-specific self.block_mask = None self.num_frame_per_block = config.arch_config.num_frames_per_block assert self.num_frame_per_block <= 3 self.independent_first_frame = False self.__post_init__() self.layer_names = [ "blocks", ] @staticmethod def _prepare_blockwise_causal_attn_mask( device: torch.device | str, num_frames: int = 21, frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1, ) -> BlockMask: """ we will divide the token sequence into the following format [1 latent frame] [1 latent frame] ... [1 latent frame] We use flexattention to construct the attention mask """ total_length = num_frames * frame_seqlen # we do right padding to get to a multiple of 128 padded_length = math.ceil(total_length / 128) * 128 - total_length ends = torch.zeros( total_length + padded_length, device=device, dtype=torch.long ) # Block-wise causal mask will attend to all elements that are before the end of the current chunk frame_indices = torch.arange( start=0, end=total_length, step=frame_seqlen * num_frame_per_block, device=device, ) for tmp in frame_indices: ends[tmp : tmp + frame_seqlen * num_frame_per_block] = ( tmp + frame_seqlen * num_frame_per_block ) def attention_mask(b, h, q_idx, kv_idx): if local_attn_size == -1: return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) else: return ( (kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen)) ) | (q_idx == kv_idx) # return ((kv_idx < total_length) & (q_idx < total_length)) | (q_idx == kv_idx) # bidirectional mask block_mask = create_block_mask( attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, KV_LEN=total_length + padded_length, _compile=False, device=device, ) if not dist.is_initialized() or dist.get_rank() == 0: print( f" cache a block wise causal mask with block size of {num_frame_per_block} frames" ) print(block_mask) # import imageio # import numpy as np # from torch.nn.attention.flex_attention import create_mask # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + # padded_length, KV_LEN=total_length + padded_length, device=device) # import cv2 # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) return block_mask def _forward_inference( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, kv_cache: dict = None, crossattn_cache: dict = None, current_start: int = 0, cache_start: int = 0, start_frame: int = 0, **kwargs, ) -> torch.Tensor: r""" Run the diffusion model with kv caching. See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. This function will be run for num_frame times. Process the latent frames one by one (1560 tokens each) """ orig_dtype = hidden_states.dtype if not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] if ( isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0 ): encoder_hidden_states_image = encoder_hidden_states_image[0] else: encoder_hidden_states_image = None batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.patch_size post_patch_num_frames = num_frames // p_t post_patch_height = height // p_h post_patch_width = width // p_w # Get rotary embeddings d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed( ( post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width, ), self.hidden_size, self.num_attention_heads, rope_dim_list, dtype=( torch.float32 if current_platform.is_mps() or current_platform.is_musa() else torch.float64 ), rope_theta=10000, start_frame=start_frame, # Assume that start_frame is 0 when kv_cache is None ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) freqs_cis = ( (freqs_cos.float(), freqs_sin.float()) if freqs_cos is not None else None ) hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = ( self.condition_embedder( timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image ) ) timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten( dim=0, sizes=timestep.shape ) if encoder_hidden_states_image is not None: encoder_hidden_states = torch.concat( [encoder_hidden_states_image, encoder_hidden_states], dim=1 ) encoder_hidden_states = ( encoder_hidden_states.to(orig_dtype) if current_platform.is_mps() else encoder_hidden_states ) # cast to orig_dtype for MPS assert encoder_hidden_states.dtype == orig_dtype # 4. Transformer blocks for block_index, block in enumerate(self.blocks): if torch.is_grad_enabled() and self.gradient_checkpointing: causal_kwargs = { "kv_cache": kv_cache[block_index], "current_start": current_start, "cache_start": cache_start, "block_mask": self.block_mask, } hidden_states = self._gradient_checkpointing_func( block, hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, **causal_kwargs, ) else: causal_kwargs = { "kv_cache": kv_cache[block_index], "crossattn_cache": crossattn_cache[block_index], "current_start": current_start, "cache_start": cache_start, "block_mask": self.block_mask, } hidden_states = block( hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, **causal_kwargs, ) # 5. Output norm, projection & unpatchify temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) hidden_states = hidden_states.reshape( batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1, ) hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) return output def _forward_train( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, start_frame: int = 0, **kwargs, ) -> torch.Tensor: orig_dtype = hidden_states.dtype if not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] if ( isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0 ): encoder_hidden_states_image = encoder_hidden_states_image[0] else: encoder_hidden_states_image = None batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.patch_size post_patch_num_frames = num_frames // p_t post_patch_height = height // p_h post_patch_width = width // p_w # Get rotary embeddings d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed( ( post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width, ), self.hidden_size, self.num_attention_heads, rope_dim_list, dtype=( torch.float32 if current_platform.is_mps() or current_platform.is_musa() else torch.float64 ), rope_theta=10000, start_frame=start_frame, ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) freqs_cis = ( (freqs_cos.float(), freqs_sin.float()) if freqs_cos is not None else None ) # Construct blockwise causal attn mask if self.block_mask is None: self.block_mask = self._prepare_blockwise_causal_attn_mask( device=hidden_states.device, num_frames=num_frames, frame_seqlen=post_patch_height * post_patch_width, num_frame_per_block=self.num_frame_per_block, local_attn_size=self.local_attn_size, ) hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = ( self.condition_embedder( timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image ) ) timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten( dim=0, sizes=timestep.shape ) if encoder_hidden_states_image is not None: encoder_hidden_states = torch.concat( [encoder_hidden_states_image, encoder_hidden_states], dim=1 ) encoder_hidden_states = ( encoder_hidden_states.to(orig_dtype) if current_platform.is_mps() else encoder_hidden_states ) # cast to orig_dtype for MPS assert encoder_hidden_states.dtype == orig_dtype # 4. Transformer blocks if torch.is_grad_enabled() and self.gradient_checkpointing: for block in self.blocks: hidden_states = self._gradient_checkpointing_func( block, hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, block_mask=self.block_mask, ) else: for block in self.blocks: hidden_states = block( hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, block_mask=self.block_mask, ) # 5. Output norm, projection & unpatchify temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) hidden_states = hidden_states.reshape( batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1, ) hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) return output def forward(self, *args, **kwargs): if kwargs.get("kv_cache") is not None: return self._forward_inference(*args, **kwargs) else: return self._forward_train(*args, **kwargs) EntryClass = CausalWanTransformer3DModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/flux.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # Copyright 2025 Black Forest Labs, The HuggingFace Team and The InstantX Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.nn as nn from diffusers.models.attention import AttentionModuleMixin from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import ( AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle, ) from torch.nn import LayerNorm as LayerNorm from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, ) from sglang.multimodal_gen.runtime.layers.mlp import FeedForward from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( NunchakuConfig, is_nunchaku_available, ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( NDRotaryEmbedding, apply_flashinfer_rope_qk_inplace, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import ( CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name try: from nunchaku.models.attention import NunchakuFeedForward # type: ignore[import] from nunchaku.models.normalization import ( # type: ignore[import] NunchakuAdaLayerNormZero, NunchakuAdaLayerNormZeroSingle, ) from nunchaku.ops.gemm import ( svdq_gemm_w4a4_cuda as _svdq_gemm_w4a4, # type: ignore[import] ) from nunchaku.ops.quantize import ( svdq_quantize_w4a4_act_fuse_lora_cuda as _svdq_quantize_w4a4, # type: ignore[import] ) _nunchaku_fused_ops_available = True except Exception: NunchakuFeedForward = None NunchakuAdaLayerNormZero = None NunchakuAdaLayerNormZeroSingle = None _svdq_gemm_w4a4 = None _svdq_quantize_w4a4 = None _nunchaku_fused_ops_available = False def _fused_gelu_mlp( x: torch.Tensor, fc1, fc2, pad_size: int = 256, ) -> torch.Tensor: """ Fused GELU MLP matching nunchaku's fused_gelu_mlp kernel path. nunchaku's single-block MLP checkpoint is calibrated for the fused path where: 1. fc1 GEMM + GELU + 0.171875 shift + unsigned re-quantization + fc2.lora_down are all done in a single fused kernel call 2. fc2 GEMM then receives unsigned INT4 activations (act_unsigned=True) Using the sequential path (fc1 → GELU → fc2 with symmetric quantization) is fundamentally incompatible with these wscales, causing visually wrong outputs. """ batch_size, seq_len, channels = x.shape x_2d = x.view(batch_size * seq_len, channels) quantized_x, ascales, lora_act = _svdq_quantize_w4a4( x_2d, lora_down=fc1.proj_down, smooth=fc1.smooth_factor, fp4=fc1.precision == "nvfp4", pad_size=pad_size, ) batch_size_pad = (batch_size * seq_len + pad_size - 1) // pad_size * pad_size is_fp4 = fc2.precision == "nvfp4" qout_act = torch.empty( batch_size_pad, fc1.output_size_per_partition // 2, dtype=torch.uint8, device=x_2d.device, ) if is_fp4: qout_ascales = torch.empty( fc1.output_size_per_partition // 16, batch_size_pad, dtype=torch.float8_e4m3fn, device=x_2d.device, ) else: qout_ascales = torch.empty( fc1.output_size_per_partition // 64, batch_size_pad, dtype=x_2d.dtype, device=x_2d.device, ) qout_lora_act = torch.empty( batch_size_pad, fc2.proj_down.shape[1], dtype=torch.float32, device=x_2d.device ) # fused: fc1 GEMM + GELU + shift + unsigned quantize + fc2.lora_down _svdq_gemm_w4a4( act=quantized_x, wgt=fc1.qweight, qout=qout_act, ascales=ascales, wscales=fc1.wscales, oscales=qout_ascales, lora_act_in=lora_act, lora_up=fc1.proj_up, lora_down=fc2.proj_down, lora_act_out=qout_lora_act, bias=fc1.bias, smooth_factor=fc2.smooth_factor, fp4=is_fp4, alpha=getattr(fc1, "_nunchaku_alpha", None), wcscales=getattr(fc1, "wcscales", None), ) output = torch.empty( batch_size * seq_len, fc2.output_size_per_partition, dtype=x_2d.dtype, device=x_2d.device, ) # fc2 GEMM with unsigned INT4 activations (fused kernel shifted by 0.171875) _svdq_gemm_w4a4( act=qout_act, wgt=fc2.qweight, out=output, ascales=qout_ascales, wscales=fc2.wscales, lora_act_in=qout_lora_act, lora_up=fc2.proj_up, bias=fc2.bias, fp4=is_fp4, alpha=getattr(fc2, "_nunchaku_alpha", None), wcscales=getattr(fc2, "wcscales", None), act_unsigned=True, ) return output.view(batch_size, seq_len, -1) def _get_qkv_projections( attn: "FluxAttention", hidden_states, encoder_hidden_states=None ): if getattr(attn, "use_fused_qkv", False): qkv, _ = attn.to_qkv(hidden_states) query, key, value = [x.contiguous() for x in qkv.chunk(3, dim=-1)] else: query, _ = attn.to_q(hidden_states) key, _ = attn.to_k(hidden_states) value, _ = attn.to_v(hidden_states) encoder_query = encoder_key = encoder_value = None if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None: if getattr(attn, "use_fused_added_qkv", False): added_qkv, _ = attn.to_added_qkv(encoder_hidden_states) encoder_query, encoder_key, encoder_value = [ x.contiguous() for x in added_qkv.chunk(3, dim=-1) ] else: encoder_query, _ = attn.add_q_proj(encoder_hidden_states) encoder_key, _ = attn.add_k_proj(encoder_hidden_states) encoder_value, _ = attn.add_v_proj(encoder_hidden_states) return query, key, value, encoder_query, encoder_key, encoder_value class FluxAttention(torch.nn.Module, AttentionModuleMixin): def __init__( self, query_dim: int, num_heads: int = 8, dim_head: int = 64, dropout: float = 0.0, bias: bool = False, added_kv_proj_dim: Optional[int] = None, added_proj_bias: Optional[bool] = True, out_bias: bool = True, eps: float = 1e-5, out_dim: int = None, context_pre_only: Optional[bool] = None, pre_only: bool = False, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.head_dim = dim_head self.inner_dim = out_dim if out_dim is not None else dim_head * num_heads self.query_dim = query_dim self.use_bias = bias self.dropout = dropout self.out_dim = out_dim if out_dim is not None else query_dim self.context_pre_only = context_pre_only self.pre_only = pre_only self.heads = out_dim // dim_head if out_dim is not None else num_heads self.added_kv_proj_dim = added_kv_proj_dim self.added_proj_bias = added_proj_bias self.use_fused_qkv = isinstance(quant_config, NunchakuConfig) self.use_fused_added_qkv = isinstance(quant_config, NunchakuConfig) self.norm_q = RMSNorm(dim_head, eps=eps) self.norm_k = RMSNorm(dim_head, eps=eps) if self.use_fused_qkv: self.to_qkv = MergedColumnParallelLinear( query_dim, [self.inner_dim] * 3, bias=bias, gather_output=True, quant_config=quant_config, prefix=f"{prefix}.to_qkv" if prefix else "to_qkv", ) else: self.to_q = ColumnParallelLinear( query_dim, self.inner_dim, bias=bias, gather_output=True, quant_config=quant_config, ) self.to_k = ColumnParallelLinear( query_dim, self.inner_dim, bias=bias, gather_output=True, quant_config=quant_config, ) self.to_v = ColumnParallelLinear( query_dim, self.inner_dim, bias=bias, gather_output=True, quant_config=quant_config, ) if not self.pre_only: self.to_out = torch.nn.ModuleList([]) self.to_out.append( ColumnParallelLinear( self.inner_dim, self.out_dim, bias=out_bias, gather_output=True, quant_config=quant_config, prefix=f"{prefix}.to_out.0" if prefix else "", ) ) if dropout != 0.0: self.to_out.append(torch.nn.Dropout(dropout)) if added_kv_proj_dim is not None: self.norm_added_q = RMSNorm(dim_head, eps=eps) self.norm_added_k = RMSNorm(dim_head, eps=eps) if self.use_fused_added_qkv: self.to_added_qkv = MergedColumnParallelLinear( added_kv_proj_dim, [self.inner_dim] * 3, bias=added_proj_bias, gather_output=True, quant_config=quant_config, prefix=f"{prefix}.to_added_qkv" if prefix else "to_added_qkv", ) else: self.add_q_proj = ColumnParallelLinear( added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, gather_output=True, quant_config=quant_config, ) self.add_k_proj = ColumnParallelLinear( added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, gather_output=True, quant_config=quant_config, ) self.add_v_proj = ColumnParallelLinear( added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, gather_output=True, quant_config=quant_config, ) self.to_add_out = ColumnParallelLinear( self.inner_dim, query_dim, bias=out_bias, gather_output=True, quant_config=quant_config, prefix=f"{prefix}.to_add_out" if prefix else "", ) self.attn = USPAttention( num_heads=num_heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, causal=False, ) def forward( self, x: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor] = None, freqs_cis=None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: query, key, value, encoder_query, encoder_key, encoder_value = ( _get_qkv_projections(self, x, encoder_hidden_states) ) query = query.unflatten(-1, (self.heads, -1)) key = key.unflatten(-1, (self.heads, -1)) value = value.unflatten(-1, (self.heads, -1)) query, key = apply_qk_norm( q=query, k=key, q_norm=self.norm_q, k_norm=self.norm_k, head_dim=self.head_dim, allow_inplace=True, ) if self.added_kv_proj_dim is not None: encoder_query = encoder_query.unflatten(-1, (self.heads, -1)) encoder_key = encoder_key.unflatten(-1, (self.heads, -1)) encoder_value = encoder_value.unflatten(-1, (self.heads, -1)) encoder_query, encoder_key = apply_qk_norm( q=encoder_query, k=encoder_key, q_norm=self.norm_added_q, k_norm=self.norm_added_k, head_dim=self.head_dim, allow_inplace=True, ) bsz, seq_len, _, _ = query.shape query = torch.cat([encoder_query, query], dim=1) key = torch.cat([encoder_key, key], dim=1) value = torch.cat([encoder_value, value], dim=1) if freqs_cis is not None: cos, sin = freqs_cis cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(), sin.to(dtype=torch.float32).contiguous(), ], dim=-1, ) query, key = apply_flashinfer_rope_qk_inplace( query, key, cos_sin_cache, is_neox=False ) x = self.attn(query, key, value) x = x.flatten(2, 3) x = x.to(query.dtype) if encoder_hidden_states is not None: encoder_hidden_states, x = x.split_with_sizes( [ encoder_hidden_states.shape[1], x.shape[1] - encoder_hidden_states.shape[1], ], dim=1, ) if not self.pre_only: x, _ = self.to_out[0](x) if len(self.to_out) == 2: x = self.to_out[1](x) encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states) return x, encoder_hidden_states else: if not self.pre_only: x, _ = self.to_out[0](x) if len(self.to_out) == 2: x = self.to_out[1](x) return x class FluxSingleTransformerBlock(nn.Module): def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 4.0, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.mlp_hidden_dim = int(dim * mlp_ratio) self.use_nunchaku_structure = isinstance(quant_config, NunchakuConfig) self.norm = AdaLayerNormZeroSingle(dim) if self.use_nunchaku_structure: self.mlp_fc1 = ColumnParallelLinear( dim, self.mlp_hidden_dim, bias=True, gather_output=True, quant_config=quant_config, prefix=f"{prefix}.mlp_fc1" if prefix else "mlp_fc1", ) self.act_mlp = nn.GELU(approximate="tanh") self.mlp_fc2 = ColumnParallelLinear( self.mlp_hidden_dim, dim, bias=True, gather_output=True, quant_config=quant_config, prefix=f"{prefix}.mlp_fc2" if prefix else "mlp_fc2", ) self.attn = FluxAttention( query_dim=dim, dim_head=attention_head_dim, num_heads=num_attention_heads, out_dim=dim, bias=True, eps=1e-6, pre_only=False, quant_config=quant_config, prefix=f"{prefix}.attn" if prefix else "attn", ) if is_nunchaku_available(): self.norm = NunchakuAdaLayerNormZeroSingle(self.norm, scale_shift=0) else: self.proj_mlp = ColumnParallelLinear( dim, self.mlp_hidden_dim, bias=True, gather_output=True, quant_config=quant_config, ) self.act_mlp = nn.GELU(approximate="tanh") self.proj_out = ColumnParallelLinear( dim + self.mlp_hidden_dim, dim, bias=True, gather_output=True, quant_config=quant_config, ) self.attn = FluxAttention( query_dim=dim, dim_head=attention_head_dim, num_heads=num_attention_heads, out_dim=dim, bias=True, eps=1e-6, pre_only=True, quant_config=quant_config, ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: text_seq_len = encoder_hidden_states.shape[1] hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) residual = hidden_states norm_hidden_states, gate = self.norm(hidden_states, emb=temb) joint_attention_kwargs = joint_attention_kwargs or {} if self.use_nunchaku_structure: if _nunchaku_fused_ops_available: mlp_hidden_states = _fused_gelu_mlp( norm_hidden_states, self.mlp_fc1, self.mlp_fc2 ) else: mlp_out, _ = self.mlp_fc1(norm_hidden_states) mlp_hidden_states = self.act_mlp(mlp_out) mlp_hidden_states, _ = self.mlp_fc2(mlp_hidden_states) attn_output = self.attn( x=norm_hidden_states, freqs_cis=freqs_cis, **joint_attention_kwargs, ) if isinstance(attn_output, tuple): attn_output = attn_output[0] hidden_states = attn_output + mlp_hidden_states gate = gate.unsqueeze(1) hidden_states = gate * hidden_states hidden_states = residual + hidden_states else: proj_hidden_states, _ = self.proj_mlp(norm_hidden_states) mlp_hidden_states = self.act_mlp(proj_hidden_states) attn_output = self.attn( x=norm_hidden_states, freqs_cis=freqs_cis, **joint_attention_kwargs, ) hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) gate = gate.unsqueeze(1) proj_out, _ = self.proj_out(hidden_states) hidden_states = gate * proj_out hidden_states = residual + hidden_states if hidden_states.dtype == torch.float16: hidden_states = hidden_states.clip(-65504, 65504) encoder_hidden_states, hidden_states = ( hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:], ) return encoder_hidden_states, hidden_states class FluxTransformerBlock(nn.Module): def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.norm1 = AdaLayerNormZero(dim) self.norm1_context = AdaLayerNormZero(dim) self.attn = FluxAttention( query_dim=dim, added_kv_proj_dim=dim, dim_head=attention_head_dim, num_heads=num_attention_heads, out_dim=dim, context_pre_only=False, bias=True, eps=eps, quant_config=quant_config, prefix=f"{prefix}.attn" if prefix else "attn", ) self.norm2 = LayerNorm(dim, eps=1e-6, elementwise_affine=False) self.norm2_context = LayerNorm(dim, eps=1e-6, elementwise_affine=False) nunchaku_enabled = ( quant_config is not None and hasattr(quant_config, "get_name") and quant_config.get_name() == "svdquant" and is_nunchaku_available() ) self.use_nunchaku_structure = nunchaku_enabled self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") self.ff_context = FeedForward( dim=dim, dim_out=dim, activation_fn="gelu-approximate", ) if nunchaku_enabled: nunchaku_kwargs = { "precision": quant_config.precision, "rank": quant_config.rank, "act_unsigned": quant_config.act_unsigned, } self.ff = NunchakuFeedForward(self.ff, **nunchaku_kwargs) self.ff_context = NunchakuFeedForward(self.ff_context, **nunchaku_kwargs) self.norm1 = NunchakuAdaLayerNormZero(self.norm1, scale_shift=0) self.norm1_context = NunchakuAdaLayerNormZero( self.norm1_context, scale_shift=0 ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( hidden_states, emb=temb ) norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = ( self.norm1_context(encoder_hidden_states, emb=temb) ) joint_attention_kwargs = joint_attention_kwargs or {} # Attention. attention_outputs = self.attn( x=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, freqs_cis=freqs_cis, **joint_attention_kwargs, ) if len(attention_outputs) == 2: attn_output, context_attn_output = attention_outputs elif len(attention_outputs) == 3: attn_output, context_attn_output, ip_attn_output = attention_outputs # Process attention outputs for the `hidden_states`. attn_output = gate_msa.unsqueeze(1) * attn_output hidden_states = hidden_states + attn_output norm_hidden_states = self.norm2(hidden_states) if self.use_nunchaku_structure: norm_hidden_states = ( norm_hidden_states * scale_mlp[:, None] + shift_mlp[:, None] ) else: norm_hidden_states = ( norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] ) ff_output = self.ff(norm_hidden_states) ff_output = gate_mlp.unsqueeze(1) * ff_output hidden_states = hidden_states + ff_output if len(attention_outputs) == 3: hidden_states = hidden_states + ip_attn_output # Process attention outputs for the `encoder_hidden_states`. context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output encoder_hidden_states = encoder_hidden_states + context_attn_output norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) if self.use_nunchaku_structure: norm_encoder_hidden_states = ( norm_encoder_hidden_states * c_scale_mlp[:, None] + c_shift_mlp[:, None] ) else: norm_encoder_hidden_states = ( norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] ) context_ff_output = self.ff_context(norm_encoder_hidden_states) encoder_hidden_states = ( encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output ) if encoder_hidden_states.dtype == torch.float16: encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) return encoder_hidden_states, hidden_states class FluxPosEmbed(nn.Module): # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11 def __init__(self, theta: int, axes_dim: List[int]): super().__init__() self.rope = NDRotaryEmbedding( rope_dim_list=axes_dim, rope_theta=theta, use_real=False, repeat_interleave_real=False, dtype=( torch.float32 if current_platform.is_mps() or current_platform.is_musa() else torch.float64 ), ) def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: pos = ids.float() # TODO: potential error: flux use n_axes = ids.shape[-1] # see: https://github.com/huggingface/diffusers/blob/17c0e79dbdf53fb6705e9c09cc1a854b84c39249/src/diffusers/models/transformers/transformer_flux.py#L509 freqs_cos, freqs_sin = self.rope.forward_uncached(pos=pos) return freqs_cos.contiguous().float(), freqs_sin.contiguous().float() class FluxTransformer2DModel(CachableDiT, OffloadableDiTMixin): """ The Transformer model introduced in Flux. Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ """ param_names_mapping = FluxConfig().arch_config.param_names_mapping @classmethod def get_nunchaku_quant_rules(cls) -> dict[str, list[str]]: return { "skip": [ "norm", "embed", "rotary", "pos_embed", ], "svdq_w4a4": [ "attn.to_qkv", "attn.to_out", "attn.add_qkv_proj", "attn.to_added_qkv", "attn.to_add_out", "img_mlp", "txt_mlp", "attention.to_qkv", "attention.to_out", "proj_mlp", "proj_out", "mlp_fc1", "mlp_fc2", "ff.net", "ff_context.net", ], "awq_w4a16": [ "img_mod", "txt_mod", ], } def __init__( self, config: FluxConfig, hf_config: dict[str, Any], quant_config: Optional[QuantizationConfig] = None, ) -> None: super().__init__(config=config, hf_config=hf_config) self.config = config.arch_config self.out_channels = ( getattr(self.config, "out_channels", None) or self.config.in_channels ) self.inner_dim = ( self.config.num_attention_heads * self.config.attention_head_dim ) self.rotary_emb = FluxPosEmbed(theta=10000, axes_dim=self.config.axes_dims_rope) text_time_guidance_cls = ( CombinedTimestepGuidanceTextProjEmbeddings if self.config.guidance_embeds else CombinedTimestepTextProjEmbeddings ) self.time_text_embed = text_time_guidance_cls( embedding_dim=self.inner_dim, pooled_projection_dim=self.config.pooled_projection_dim, ) self.context_embedder = ColumnParallelLinear( self.config.joint_attention_dim, self.inner_dim, bias=True, gather_output=True, ) self.x_embedder = ColumnParallelLinear( self.config.in_channels, self.inner_dim, bias=True, gather_output=True ) self.transformer_blocks = nn.ModuleList( [ FluxTransformerBlock( dim=self.inner_dim, num_attention_heads=self.config.num_attention_heads, attention_head_dim=self.config.attention_head_dim, quant_config=quant_config, prefix=f"transformer_blocks.{i}", ) for i in range(self.config.num_layers) ] ) self.single_transformer_blocks = nn.ModuleList( [ FluxSingleTransformerBlock( dim=self.inner_dim, num_attention_heads=self.config.num_attention_heads, attention_head_dim=self.config.attention_head_dim, quant_config=quant_config, prefix=f"single_transformer_blocks.{i}", ) for i in range(self.config.num_single_layers) ] ) self.norm_out = AdaLayerNormContinuous( self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6 ) self.proj_out = ColumnParallelLinear( self.inner_dim, self.config.patch_size * self.config.patch_size * self.out_channels, bias=True, gather_output=True, ) self.layer_names = [ "transformer_blocks", "single_transformer_blocks", ] def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, pooled_projections: torch.Tensor = None, timestep: torch.LongTensor = None, guidance: torch.Tensor = None, freqs_cis: torch.Tensor = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, ) -> Union[torch.Tensor, Transformer2DModelOutput]: """ The [`FluxTransformer2DModel`] forward method. Args: hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`): Input `hidden_states`. encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`): Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. pooled_projections (`torch.Tensor` of shape `(batch_size, projection_dim)`): Embeddings projected from the embeddings of input conditions. timestep ( `torch.LongTensor`): Used to indicate denoising step. guidance (`torch.Tensor`): Guidance embeddings. joint_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). """ if ( joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None ): logger.warning( "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." ) hidden_states, _ = self.x_embedder(hidden_states) # Only pass guidance to time_text_embed if the model supports it if self.config.guidance_embeds and guidance is not None: temb = self.time_text_embed(timestep, guidance, pooled_projections) else: temb = self.time_text_embed(timestep, pooled_projections) encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states) if ( joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs ): ip_adapter_image_embeds = joint_attention_kwargs.pop( "ip_adapter_image_embeds" ) ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds) joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states}) for block in self.transformer_blocks: encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb, freqs_cis=freqs_cis, joint_attention_kwargs=joint_attention_kwargs, ) for block in self.single_transformer_blocks: encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb, freqs_cis=freqs_cis, joint_attention_kwargs=joint_attention_kwargs, ) hidden_states = self.norm_out(hidden_states, temb) output, _ = self.proj_out(hidden_states) return output EntryClass = FluxTransformer2DModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/flux_2.py ================================================ # Copyright 2025 Black Forest Labs, The HuggingFace Team and The InstantX Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from typing import Any, Dict, List, Optional, Tuple import torch import torch.nn as nn from diffusers.models.attention import AttentionModuleMixin from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.normalization import AdaLayerNormContinuous from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( NDRotaryEmbedding, apply_flashinfer_rope_qk_inplace, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name def _get_qkv_projections( attn: "Flux2Attention", hidden_states, encoder_hidden_states=None ): query, _ = attn.to_q(hidden_states) key, _ = attn.to_k(hidden_states) value, _ = attn.to_v(hidden_states) encoder_query = encoder_key = encoder_value = None if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None: encoder_query, _ = attn.add_q_proj(encoder_hidden_states) encoder_key, _ = attn.add_k_proj(encoder_hidden_states) encoder_value, _ = attn.add_v_proj(encoder_hidden_states) return query, key, value, encoder_query, encoder_key, encoder_value class Flux2SwiGLU(nn.Module): """ Flux 2 uses a SwiGLU-style activation in the transformer feedforward sub-blocks, but with the linear projection layer fused into the first linear layer of the FF sub-block. Thus, this module has no trainable parameters. """ def __init__(self): super().__init__() self.gate_fn = nn.SiLU() def forward(self, x: torch.Tensor) -> torch.Tensor: x1, x2 = x.chunk(2, dim=-1) x = self.gate_fn(x1) * x2 return x class Flux2FeedForward(nn.Module): def __init__( self, dim: int, dim_out: Optional[int] = None, mult: float = 3.0, inner_dim: Optional[int] = None, bias: bool = False, quant_config: Optional[QuantizationConfig] = None, ): super().__init__() if inner_dim is None: inner_dim = int(dim * mult) dim_out = dim_out or dim # Flux2SwiGLU will reduce the dimension by half self.linear_in = ColumnParallelLinear( dim, inner_dim * 2, bias=bias, gather_output=True, quant_config=quant_config ) self.act_fn = Flux2SwiGLU() self.linear_out = ColumnParallelLinear( inner_dim, dim_out, bias=bias, gather_output=True, quant_config=quant_config ) def forward(self, x: torch.Tensor) -> torch.Tensor: x, _ = self.linear_in(x) x = self.act_fn(x) x, _ = self.linear_out(x) return x class Flux2Attention(torch.nn.Module, AttentionModuleMixin): def __init__( self, query_dim: int, num_heads: int = 8, dim_head: int = 64, dropout: float = 0.0, bias: bool = False, added_kv_proj_dim: Optional[int] = None, added_proj_bias: Optional[bool] = True, out_bias: bool = True, eps: float = 1e-5, out_dim: int = None, elementwise_affine: bool = True, quant_config: Optional[QuantizationConfig] = None, ): super().__init__() self.head_dim = dim_head self.inner_dim = out_dim if out_dim is not None else dim_head * num_heads self.query_dim = query_dim self.out_dim = out_dim if out_dim is not None else query_dim self.heads = out_dim // dim_head if out_dim is not None else num_heads self.use_bias = bias self.dropout = dropout self.added_kv_proj_dim = added_kv_proj_dim self.added_proj_bias = added_proj_bias self.to_q = ColumnParallelLinear( query_dim, self.inner_dim, bias=bias, gather_output=True, quant_config=quant_config, ) self.to_k = ColumnParallelLinear( query_dim, self.inner_dim, bias=bias, gather_output=True, quant_config=quant_config, ) self.to_v = ColumnParallelLinear( query_dim, self.inner_dim, bias=bias, gather_output=True, quant_config=quant_config, ) # QK Norm self.norm_q = RMSNorm(dim_head, eps=eps) self.norm_k = RMSNorm(dim_head, eps=eps) self.to_out = torch.nn.ModuleList([]) self.to_out.append( ColumnParallelLinear( self.inner_dim, self.out_dim, bias=out_bias, gather_output=True, quant_config=quant_config, ) ) self.to_out.append(torch.nn.Dropout(dropout)) if added_kv_proj_dim is not None: self.norm_added_q = RMSNorm(dim_head, eps=eps) self.norm_added_k = RMSNorm(dim_head, eps=eps) self.add_q_proj = ColumnParallelLinear( added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, gather_output=True, quant_config=quant_config, ) self.add_k_proj = ColumnParallelLinear( added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, gather_output=True, quant_config=quant_config, ) self.add_v_proj = ColumnParallelLinear( added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, gather_output=True, quant_config=quant_config, ) self.to_add_out = ColumnParallelLinear( self.inner_dim, query_dim, bias=out_bias, gather_output=True, quant_config=quant_config, ) self.attn = USPAttention( num_heads=num_heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, causal=False, ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor] = None, freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: query, key, value, encoder_query, encoder_key, encoder_value = ( _get_qkv_projections(self, hidden_states, encoder_hidden_states) ) query = query.unflatten(-1, (self.heads, -1)) key = key.unflatten(-1, (self.heads, -1)) value = value.unflatten(-1, (self.heads, -1)) query, key = apply_qk_norm( q=query, k=key, q_norm=self.norm_q, k_norm=self.norm_k, head_dim=self.head_dim, allow_inplace=True, ) if self.added_kv_proj_dim is not None: encoder_query = encoder_query.unflatten(-1, (self.heads, -1)) encoder_key = encoder_key.unflatten(-1, (self.heads, -1)) encoder_value = encoder_value.unflatten(-1, (self.heads, -1)) encoder_query, encoder_key = apply_qk_norm( q=encoder_query, k=encoder_key, q_norm=self.norm_added_q, k_norm=self.norm_added_k, head_dim=self.head_dim, allow_inplace=True, ) query = torch.cat([encoder_query, query], dim=1) key = torch.cat([encoder_key, key], dim=1) value = torch.cat([encoder_value, value], dim=1) if freqs_cis is not None: cos, sin = freqs_cis cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(), sin.to(dtype=torch.float32).contiguous(), ], dim=-1, ) query, key = apply_flashinfer_rope_qk_inplace( query, key, cos_sin_cache, is_neox=False ) num_rep = ( encoder_hidden_states.shape[1] if encoder_hidden_states is not None else 0 ) hidden_states = self.attn(query, key, value, num_replicated_prefix=num_rep) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) if encoder_hidden_states is not None: encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( [ encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1], ], dim=1, ) encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states) hidden_states, _ = self.to_out[0](hidden_states) hidden_states = self.to_out[1](hidden_states) if encoder_hidden_states is not None: return hidden_states, encoder_hidden_states else: return hidden_states class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): """ Flux 2 parallel self-attention for the Flux 2 single-stream transformer blocks. This implements a parallel transformer block, where the attention QKV projections are fused to the feedforward (FF) input projections, and the attention output projections are fused to the FF output projections. See the [ViT-22B paper](https://arxiv.org/abs/2302.05442) for a visual depiction of this type of transformer block. """ # Does not support QKV fusion as the QKV projections are always fused _supports_qkv_fusion = False def __init__( self, query_dim: int, num_heads: int = 8, dim_head: int = 64, dropout: float = 0.0, bias: bool = False, out_bias: bool = True, eps: float = 1e-5, out_dim: int = None, elementwise_affine: bool = True, mlp_ratio: float = 4.0, mlp_mult_factor: int = 2, quant_config: Optional[QuantizationConfig] = None, ): super().__init__() self.head_dim = dim_head self.inner_dim = out_dim if out_dim is not None else dim_head * num_heads self.query_dim = query_dim self.out_dim = out_dim if out_dim is not None else query_dim self.heads = out_dim // dim_head if out_dim is not None else num_heads self.use_bias = bias self.dropout = dropout self.mlp_ratio = mlp_ratio self.mlp_hidden_dim = int(query_dim * self.mlp_ratio) self.mlp_mult_factor = mlp_mult_factor # Fused QKV projections + MLP input projection self.to_qkv_mlp_proj = ColumnParallelLinear( self.query_dim, self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor, bias=bias, gather_output=True, quant_config=quant_config, ) self.mlp_act_fn = Flux2SwiGLU() # QK Norm self.norm_q = RMSNorm(dim_head, eps=eps) self.norm_k = RMSNorm(dim_head, eps=eps) # Fused attention output projection + MLP output projection self.to_out = ColumnParallelLinear( self.inner_dim + self.mlp_hidden_dim, self.out_dim, bias=out_bias, gather_output=True, quant_config=quant_config, ) self.attn = USPAttention( num_heads=num_heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, causal=False, ) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, num_replicated_prefix: int = 0, **kwargs, ) -> torch.Tensor: # Parallel in (QKV + MLP in) projection hidden_states, _ = self.to_qkv_mlp_proj(hidden_states) qkv, mlp_hidden_states = torch.split( hidden_states, [3 * self.inner_dim, self.mlp_hidden_dim * self.mlp_mult_factor], dim=-1, ) # Handle the attention logic query, key, value = qkv.chunk(3, dim=-1) query = query.unflatten(-1, (self.heads, -1)) key = key.unflatten(-1, (self.heads, -1)) value = value.unflatten(-1, (self.heads, -1)) query = self.norm_q(query) key = self.norm_k(key) if freqs_cis is not None: cos, sin = freqs_cis cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(), sin.to(dtype=torch.float32).contiguous(), ], dim=-1, ) query, key = apply_flashinfer_rope_qk_inplace( query, key, cos_sin_cache, is_neox=False ) hidden_states = self.attn( query, key, value, num_replicated_prefix=num_replicated_prefix ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) # Handle the feedforward (FF) logic mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states) # Concatenate and parallel output projection hidden_states = torch.cat([hidden_states, mlp_hidden_states], dim=-1) hidden_states, _ = self.to_out(hidden_states) return hidden_states class Flux2SingleTransformerBlock(nn.Module): def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 3.0, eps: float = 1e-6, bias: bool = False, quant_config: Optional[QuantizationConfig] = None, ): super().__init__() self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) # Note that the MLP in/out linear layers are fused with the attention QKV/out projections, respectively; this # is often called a "parallel" transformer block. See the [ViT-22B paper](https://arxiv.org/abs/2302.05442) # for a visual depiction of this type of transformer block. self.attn = Flux2ParallelSelfAttention( query_dim=dim, dim_head=attention_head_dim, num_heads=num_attention_heads, out_dim=dim, bias=bias, out_bias=bias, eps=eps, mlp_ratio=mlp_ratio, mlp_mult_factor=2, quant_config=quant_config, ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor], temb_mod_params: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, split_hidden_states: bool = False, text_seq_len: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: # If encoder_hidden_states is None, hidden_states is assumed to have encoder_hidden_states already # concatenated if encoder_hidden_states is not None: text_seq_len = encoder_hidden_states.shape[1] hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) mod_shift, mod_scale, mod_gate = temb_mod_params norm_hidden_states = self.norm(hidden_states) norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift joint_attention_kwargs = joint_attention_kwargs or {} attn_output = self.attn( hidden_states=norm_hidden_states, freqs_cis=freqs_cis, num_replicated_prefix=text_seq_len or 0, **joint_attention_kwargs, ) hidden_states = hidden_states + mod_gate * attn_output if hidden_states.dtype == torch.float16: hidden_states = hidden_states.clip(-65504, 65504) if split_hidden_states: encoder_hidden_states, hidden_states = ( hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:], ) return encoder_hidden_states, hidden_states else: return hidden_states class Flux2TransformerBlock(nn.Module): def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 3.0, eps: float = 1e-6, bias: bool = False, quant_config: Optional[QuantizationConfig] = None, ): super().__init__() self.mlp_hidden_dim = int(dim * mlp_ratio) self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) self.attn = Flux2Attention( query_dim=dim, added_kv_proj_dim=dim, dim_head=attention_head_dim, num_heads=num_attention_heads, out_dim=dim, bias=bias, added_proj_bias=bias, out_bias=bias, eps=eps, quant_config=quant_config, ) self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) self.ff = Flux2FeedForward( dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias, quant_config=quant_config ) self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) self.ff_context = Flux2FeedForward( dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias, quant_config=quant_config ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb_mod_params_img: Tuple[ Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ... ], temb_mod_params_txt: Tuple[ Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ... ], freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: joint_attention_kwargs = joint_attention_kwargs or {} # Modulation parameters shape: [1, 1, self.dim] (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = ( temb_mod_params_img ) (c_shift_msa, c_scale_msa, c_gate_msa), ( c_shift_mlp, c_scale_mlp, c_gate_mlp, ) = temb_mod_params_txt # Img stream norm_hidden_states = self.norm1(hidden_states) norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa # Conditioning txt stream norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states) norm_encoder_hidden_states = ( 1 + c_scale_msa ) * norm_encoder_hidden_states + c_shift_msa # Attention on concatenated img + txt stream attention_outputs = self.attn( hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, freqs_cis=freqs_cis, **joint_attention_kwargs, ) attn_output, context_attn_output = attention_outputs # Process attention outputs for the image stream (`hidden_states`). attn_output = gate_msa * attn_output hidden_states = hidden_states + attn_output norm_hidden_states = self.norm2(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp ff_output = self.ff(norm_hidden_states) hidden_states = hidden_states + gate_mlp * ff_output # Process attention outputs for the text stream (`encoder_hidden_states`). context_attn_output = c_gate_msa * context_attn_output encoder_hidden_states = encoder_hidden_states + context_attn_output norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) norm_encoder_hidden_states = ( norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp ) context_ff_output = self.ff_context(norm_encoder_hidden_states) encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output if encoder_hidden_states.dtype == torch.float16: encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) return encoder_hidden_states, hidden_states class Flux2TimestepGuidanceEmbeddings(nn.Module): def __init__( self, in_channels: int = 256, embedding_dim: int = 6144, bias: bool = False, guidance_embeds: bool = True, ): super().__init__() self.time_proj = Timesteps( num_channels=in_channels, flip_sin_to_cos=True, downscale_freq_shift=0 ) self.timestep_embedder = TimestepEmbedding( in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias ) if guidance_embeds: self.guidance_embedder = TimestepEmbedding( in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias, ) else: self.guidance_embedder = None def forward( self, timestep: torch.Tensor, guidance: Optional[torch.Tensor] = None ) -> torch.Tensor: timesteps_proj = self.time_proj(timestep) timesteps_emb = self.timestep_embedder( timesteps_proj.to(timestep.dtype) ) # (N, D) if guidance is not None and self.guidance_embedder is not None: guidance_proj = self.time_proj(guidance) guidance_emb = self.guidance_embedder( guidance_proj.to(guidance.dtype) ) # (N, D) time_guidance_emb = timesteps_emb + guidance_emb return time_guidance_emb else: return timesteps_emb class Flux2Modulation(nn.Module): def __init__(self, dim: int, mod_param_sets: int = 2, bias: bool = False): super().__init__() self.mod_param_sets = mod_param_sets self.linear = ColumnParallelLinear( dim, dim * 3 * self.mod_param_sets, bias=bias, gather_output=True ) self.act_fn = nn.SiLU() def forward( self, temb: torch.Tensor ) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]: mod = self.act_fn(temb) mod, _ = self.linear(mod) if mod.ndim == 2: mod = mod.unsqueeze(1) mod_params = torch.chunk(mod, 3 * self.mod_param_sets, dim=-1) # Return tuple of 3-tuples of modulation params shift/scale/gate return tuple( mod_params[3 * i : 3 * (i + 1)] for i in range(self.mod_param_sets) ) class Flux2PosEmbed(nn.Module): def __init__(self, theta: int, axes_dim: List[int]): super().__init__() self.rope = NDRotaryEmbedding( rope_dim_list=axes_dim, rope_theta=theta, use_real=False, repeat_interleave_real=False, dtype=( torch.float32 if current_platform.is_mps() or current_platform.is_musa() else torch.float64 ), ) def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: pos = ids.float() # TODO: potential error: flux use n_axes = ids.shape[-1] # see: https://github.com/huggingface/diffusers/blob/17c0e79dbdf53fb6705e9c09cc1a854b84c39249/src/diffusers/models/transformers/transformer_flux.py#L509 freqs_cos, freqs_sin = self.rope.forward_uncached(pos=pos) return freqs_cos.contiguous().float(), freqs_sin.contiguous().float() class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin): """ The Transformer model introduced in Flux 2. Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ """ param_names_mapping = FluxConfig().arch_config.param_names_mapping def __init__( self, config: FluxConfig, hf_config: dict[str, Any], quant_config: Optional[QuantizationConfig] = None, ): super().__init__(config=config, hf_config=hf_config) patch_size: int = config.patch_size in_channels: int = config.in_channels out_channels: Optional[int] = config.out_channels num_layers: int = config.num_layers num_single_layers: int = config.num_single_layers attention_head_dim: int = config.attention_head_dim num_attention_heads: int = config.num_attention_heads joint_attention_dim: int = config.joint_attention_dim timestep_guidance_channels: int = config.timestep_guidance_channels mlp_ratio: float = config.mlp_ratio axes_dims_rope: Tuple[int, ...] = config.axes_dims_rope rope_theta: int = config.rope_theta eps: float = config.eps guidance_embeds: bool = getattr(config, "guidance_embeds", True) self.out_channels = out_channels or in_channels self.inner_dim = num_attention_heads * attention_head_dim self.guidance_embeds = guidance_embeds # 1. Sinusoidal positional embedding for RoPE on image and text tokens self.rotary_emb = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope) # 2. Combined timestep + guidance embedding self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings( in_channels=timestep_guidance_channels, embedding_dim=self.inner_dim, bias=False, guidance_embeds=guidance_embeds, ) # 3. Modulation (double stream and single stream blocks share modulation parameters, resp.) # Two sets of shift/scale/gate modulation parameters for the double stream attn and FF sub-blocks self.double_stream_modulation_img = Flux2Modulation( self.inner_dim, mod_param_sets=2, bias=False ) self.double_stream_modulation_txt = Flux2Modulation( self.inner_dim, mod_param_sets=2, bias=False ) # Only one set of modulation parameters as the attn and FF sub-blocks are run in parallel for single stream self.single_stream_modulation = Flux2Modulation( self.inner_dim, mod_param_sets=1, bias=False ) # 4. Input projections self.x_embedder = ColumnParallelLinear( in_channels, self.inner_dim, bias=False, gather_output=True ) self.context_embedder = ColumnParallelLinear( joint_attention_dim, self.inner_dim, bias=False, gather_output=True ) # 5. Double Stream Transformer Blocks self.transformer_blocks = nn.ModuleList( [ Flux2TransformerBlock( dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, mlp_ratio=mlp_ratio, eps=eps, bias=False, quant_config=quant_config, ) for _ in range(num_layers) ] ) # 6. Single Stream Transformer Blocks self.single_transformer_blocks = nn.ModuleList( [ Flux2SingleTransformerBlock( dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, mlp_ratio=mlp_ratio, eps=eps, bias=False, quant_config=quant_config, ) for _ in range(num_single_layers) ] ) # 7. Output layers self.norm_out = AdaLayerNormContinuous( self.inner_dim, self.inner_dim, elementwise_affine=False, eps=eps, bias=False, ) self.proj_out = ColumnParallelLinear( self.inner_dim, patch_size * patch_size * self.out_channels, bias=False, gather_output=True, ) self.layer_names = ["transformer_blocks", "single_transformer_blocks"] def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, timestep: torch.LongTensor = None, guidance: torch.Tensor = None, freqs_cis: torch.Tensor = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, ) -> torch.Tensor: """ The [`FluxTransformer2DModel`] forward method. Args: hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`): Input `hidden_states`. encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`): Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. timestep ( `torch.LongTensor`): Used to indicate denoising step. joint_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). """ # 0. Handle input arguments if joint_attention_kwargs is not None: joint_attention_kwargs = joint_attention_kwargs.copy() lora_scale = joint_attention_kwargs.pop("scale", 1.0) else: lora_scale = 1.0 num_txt_tokens = encoder_hidden_states.shape[1] # 1. Calculate timestep embedding and modulation parameters timestep = timestep.to(hidden_states.dtype) if guidance is not None: guidance = guidance.to(hidden_states.dtype) temb = self.time_guidance_embed(timestep, guidance) double_stream_mod_img = self.double_stream_modulation_img(temb) double_stream_mod_txt = self.double_stream_modulation_txt(temb) single_stream_mod = self.single_stream_modulation(temb)[0] # 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states) hidden_states, _ = self.x_embedder(hidden_states) encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states) # 3. Calculate RoPE embeddings from image and text tokens # NOTE: the below logic means that we can't support batched inference with images of different resolutions or # text prompts of different lengths. Is this a use case we want to support? # 4. Double Stream Transformer Blocks for index_block, block in enumerate(self.transformer_blocks): encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb_mod_params_img=double_stream_mod_img, temb_mod_params_txt=double_stream_mod_txt, freqs_cis=freqs_cis, joint_attention_kwargs=joint_attention_kwargs, ) # Concatenate text and image streams for single-block inference hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) # 5. Single Stream Transformer Blocks for index_block, block in enumerate(self.single_transformer_blocks): hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=None, temb_mod_params=single_stream_mod, freqs_cis=freqs_cis, joint_attention_kwargs=joint_attention_kwargs, text_seq_len=num_txt_tokens, ) # Remove text tokens from concatenated stream hidden_states = hidden_states[:, num_txt_tokens:, ...] # 6. Output layers hidden_states = self.norm_out(hidden_states, temb) output, _ = self.proj_out(hidden_states) return output EntryClass = Flux2Transformer2DModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/glm_image.py ================================================ # Copyright 2025 The CogView team, Tsinghua University & ZhipuAI and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F from sglang.multimodal_gen.configs.models.dits.glmimage import GlmImageDitConfig from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_sp_parallel_rank, get_sp_world_size, ) from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.layernorm import ( ScaleResidualLayerNormScaleShift, ) from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear from sglang.multimodal_gen.runtime.layers.mlp import FeedForward from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( _apply_rotary_emb, apply_flashinfer_rope_qk_inplace, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import Timesteps from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) _is_cuda = current_platform.is_cuda() class GlmImageLayerKVCache: """KV cache for GlmImage model.""" def __init__(self): self.k_cache = None self.v_cache = None self.mode: Optional[str] = None # "write", "read", "skip" def store(self, k: torch.Tensor, v: torch.Tensor): if self.k_cache is None: self.k_cache = k self.v_cache = v else: self.k_cache = torch.cat([self.k_cache, k], dim=2) self.v_cache = torch.cat([self.v_cache, v], dim=2) def get(self): return self.k_cache, self.v_cache def clear(self): self.k_cache = None self.v_cache = None self.mode = None class GlmImageKVCache: """Container for all layers' KV caches.""" def __init__(self, num_layers: int): self.num_layers = num_layers self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)] def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache: return self.caches[layer_idx] def set_mode(self, mode: Optional[str]): if mode is not None and mode not in ["write", "read", "skip"]: raise ValueError( f"Invalid mode: {mode}, must be one of 'write', 'read', 'skip'" ) for cache in self.caches: cache.mode = mode def clear(self): for cache in self.caches: cache.clear() class GlmImageTimestepEmbedding(nn.Module): """ Replacement for diffusers TimestepEmbedding using ReplicatedLinear. Structure: linear_1 -> act(silu) -> linear_2 """ def __init__( self, in_channels: int, time_embed_dim: int, act_fn: str = "silu", out_dim: int = None, ): super().__init__() if out_dim is None: out_dim = time_embed_dim self.linear_1 = ReplicatedLinear(in_channels, time_embed_dim, bias=True) if act_fn == "silu": self.act = nn.SiLU() elif act_fn == "gelu": self.act = nn.GELU(approximate="tanh") else: self.act = nn.SiLU() self.linear_2 = ReplicatedLinear(time_embed_dim, out_dim, bias=True) def forward(self, sample: torch.Tensor) -> torch.Tensor: sample, _ = self.linear_1(sample) sample = self.act(sample) sample, _ = self.linear_2(sample) return sample class GlmImageTextProjection(nn.Module): """ Replacement for diffusers PixArtAlphaTextProjection using ReplicatedLinear. Structure: linear_1 -> act_1 -> linear_2 """ def __init__( self, in_features: int, hidden_size: int, out_features: int = None, act_fn: str = "silu", ): super().__init__() if out_features is None: out_features = hidden_size self.linear_1 = ReplicatedLinear(in_features, hidden_size, bias=True) if act_fn == "silu": self.act_1 = nn.SiLU() elif act_fn == "gelu_tanh": self.act_1 = nn.GELU(approximate="tanh") else: self.act_1 = nn.SiLU() self.linear_2 = ReplicatedLinear(hidden_size, out_features, bias=True) def forward(self, caption: torch.Tensor) -> torch.Tensor: hidden_states, _ = self.linear_1(caption) hidden_states = self.act_1(hidden_states) hidden_states, _ = self.linear_2(hidden_states) return hidden_states class GlmImageCombinedTimestepSizeEmbeddings(nn.Module): def __init__( self, embedding_dim: int, condition_dim: int, pooled_projection_dim: int, timesteps_dim: int = 256, ): super().__init__() self.time_proj = Timesteps( num_channels=timesteps_dim, flip_sin_to_cos=True, downscale_freq_shift=0 ) self.condition_proj = Timesteps( num_channels=condition_dim, flip_sin_to_cos=True, downscale_freq_shift=0 ) self.timestep_embedder = GlmImageTimestepEmbedding( in_channels=timesteps_dim, time_embed_dim=embedding_dim ) self.condition_embedder = GlmImageTextProjection( pooled_projection_dim, embedding_dim, act_fn="silu" ) def forward( self, timestep: torch.Tensor, target_size: torch.Tensor, crop_coords: torch.Tensor, hidden_dtype: torch.dtype, ) -> torch.Tensor: timesteps_proj = self.time_proj(timestep) crop_coords_proj = self.condition_proj(crop_coords.flatten()).view( crop_coords.size(0), -1 ) target_size_proj = self.condition_proj(target_size.flatten()).view( target_size.size(0), -1 ) # (B, 2 * condition_dim) condition_proj = torch.cat([crop_coords_proj, target_size_proj], dim=1) timesteps_emb = self.timestep_embedder( timesteps_proj.to(dtype=hidden_dtype) ) # (B, embedding_dim) condition_emb = self.condition_embedder( condition_proj.to(dtype=hidden_dtype) ) # (B, embedding_dim) conditioning = timesteps_emb + condition_emb return conditioning class GlmImageImageProjector(nn.Module): def __init__( self, in_channels: int = 16, hidden_size: int = 2560, patch_size: int = 2, ): super().__init__() self.patch_size = patch_size self.proj = nn.Linear(in_channels * patch_size**2, hidden_size) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: batch_size, channel, height, width = hidden_states.shape post_patch_height = height // self.patch_size post_patch_width = width // self.patch_size hidden_states = hidden_states.reshape( batch_size, channel, post_patch_height, self.patch_size, post_patch_width, self.patch_size, ) hidden_states = ( hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) ) hidden_states = self.proj(hidden_states) return hidden_states class GlmImageAdaLayerNormZero(nn.Module): def __init__(self, embedding_dim: int, dim: int) -> None: super().__init__() self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) self.linear = ReplicatedLinear(embedding_dim, 12 * dim, bias=True) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: dtype = hidden_states.dtype norm_hidden_states = self.norm(hidden_states).to(dtype=dtype) norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to( dtype=dtype ) emb, _ = self.linear(temb) ( shift_msa, c_shift_msa, scale_msa, c_scale_msa, gate_msa, c_gate_msa, shift_mlp, c_shift_mlp, scale_mlp, c_scale_mlp, gate_mlp, c_gate_mlp, ) = emb.chunk(12, dim=1) hidden_states = norm_hidden_states * ( 1 + scale_msa.unsqueeze(1) ) + shift_msa.unsqueeze(1) encoder_hidden_states = norm_encoder_hidden_states * ( 1 + c_scale_msa.unsqueeze(1) ) + c_shift_msa.unsqueeze(1) return ( hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp, ) class GlmImageAttention(torch.nn.Module): def __init__( self, query_dim, heads, dim_head, out_dim, bias, qk_norm, elementwise_affine, eps, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", quant_config: QuantizationConfig | None = None, ): super().__init__() self.k_cache = None self.v_cache = None self.heads = out_dim // dim_head if out_dim is not None else heads self.dim_head = dim_head self.inner_dim = out_dim if out_dim is not None else dim_head * heads self.inner_kv_dim = self.inner_dim self.out_dim = out_dim if out_dim is not None else query_dim self.num_kv_heads = self.dim_head // self.inner_kv_dim self.to_q = ReplicatedLinear( query_dim, self.inner_dim, bias=bias, quant_config=quant_config ) self.to_k = ReplicatedLinear( query_dim, self.inner_kv_dim, bias=bias, quant_config=quant_config ) self.to_v = ReplicatedLinear( query_dim, self.inner_kv_dim, bias=bias, quant_config=quant_config ) # (dropout omitted) self.to_out = nn.ModuleList( [ ReplicatedLinear( self.inner_dim, self.out_dim, bias=True, quant_config=quant_config ) ] ) if qk_norm is None: self.norm_q = None self.norm_k = None elif qk_norm == "layer_norm": self.norm_q = nn.LayerNorm( dim_head, eps=eps, elementwise_affine=elementwise_affine ) self.norm_k = nn.LayerNorm( dim_head, eps=eps, elementwise_affine=elementwise_affine ) else: raise ValueError( f"unknown qk_norm: {qk_norm}. Should be one of None, 'layer_norm', 'fp32_layer_norm', 'layer_norm_across_heads', 'rms_norm', 'rms_norm_across_heads', 'l2'." ) self.attn = USPAttention( num_heads=self.heads, head_size=dim_head, num_kv_heads=self.num_kv_heads, dropout_rate=0, softmax_scale=None, causal=False, ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, kv_cache: Optional[GlmImageLayerKVCache] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: dtype = encoder_hidden_states.dtype batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape batch_size, image_seq_length, embed_dim = hidden_states.shape hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) # 1. QKV projections query, _ = self.to_q(hidden_states) key, _ = self.to_k(hidden_states) value, _ = self.to_v(hidden_states) query = query.unflatten(2, (self.heads, -1)) key = key.unflatten(2, (self.heads, -1)) value = value.unflatten(2, (self.heads, -1)) # 2. QK normalization if self.norm_q is not None: query = self.norm_q(query).to(dtype=dtype) if self.norm_k is not None: key = self.norm_k(key).to(dtype=dtype) # 3. Rotational positional embeddings applied to latent stream if image_rotary_emb is not None: cos, sin = image_rotary_emb if _is_cuda and cos.dim() == 2: q_img = query[:, text_seq_length:, :, :] k_img = key[:, text_seq_length:, :, :] cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(), sin.to(dtype=torch.float32).contiguous(), ], dim=-1, ) # apply_flashinfer_rope_qk_inplace is inplace kernel and q_img/k_img are views of query/key, so we need not copy back q_out, k_out = apply_flashinfer_rope_qk_inplace( q_img, k_img, cos_sin_cache, is_neox=True ) else: query[:, text_seq_length:, :, :] = _apply_rotary_emb( query[:, text_seq_length:, :, :], cos, sin, is_neox_style=True ) key[:, text_seq_length:, :, :] = _apply_rotary_emb( key[:, text_seq_length:, :, :], cos, sin, is_neox_style=True ) if kv_cache is not None: if kv_cache.mode == "write": kv_cache.store(key, value) elif kv_cache.mode == "read": k_cache, v_cache = kv_cache.get() key = torch.cat([k_cache, key], dim=1) if k_cache is not None else key value = ( torch.cat([v_cache, value], dim=1) if v_cache is not None else value ) elif kv_cache.mode == "skip": pass # 4. Attention if attention_mask is not None: text_attn_mask = attention_mask assert ( text_attn_mask.dim() == 2 ), "the shape of text_attn_mask should be (batch_size, text_seq_length)" text_attn_mask = text_attn_mask.float().to(query.device) mix_attn_mask = torch.ones( (batch_size, text_seq_length + image_seq_length), device=query.device ) mix_attn_mask[:, :text_seq_length] = text_attn_mask mix_attn_mask = mix_attn_mask.unsqueeze(2) attn_mask_matrix = mix_attn_mask @ mix_attn_mask.transpose(1, 2) attention_mask = (attn_mask_matrix > 0).unsqueeze(1).to(query.dtype) hidden_states = self.attn(query, key, value) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) # 5. Output projection hidden_states, _ = self.to_out[0](hidden_states) # hidden_states = self.to_out[1](hidden_states) # (dropout omitted) encoder_hidden_states, hidden_states = hidden_states.split( [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1 ) return hidden_states, encoder_hidden_states class GlmImageTransformerBlock(nn.Module): def __init__( self, dim: int = 2560, num_attention_heads: int = 64, attention_head_dim: int = 40, time_embed_dim: int = 512, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", quant_config: QuantizationConfig | None = None, ) -> None: super().__init__() # 1. Attention self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim) self.attn1 = GlmImageAttention( query_dim=dim, heads=num_attention_heads, dim_head=attention_head_dim, out_dim=dim, bias=True, qk_norm="layer_norm", elementwise_affine=False, eps=1e-5, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn1", quant_config=quant_config, ) # 2. Feedforward self.norm2 = ScaleResidualLayerNormScaleShift( dim, eps=1e-5, elementwise_affine=False ) self.norm2_context = ScaleResidualLayerNormScaleShift( dim, eps=1e-5, elementwise_affine=False ) self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, image_rotary_emb: Optional[ Union[ Tuple[torch.Tensor, torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]], ] ] = None, attention_mask: Optional[Dict[str, torch.Tensor]] = None, attention_kwargs: Optional[Dict[str, Any]] = None, kv_cache: Optional[GlmImageLayerKVCache] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: # 1. Timestep conditioning ( norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp, ) = self.norm1(hidden_states, encoder_hidden_states, temb) # 2. Attention if attention_kwargs is None: attention_kwargs = {} attn_hidden_states, attn_encoder_hidden_states = self.attn1( hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, image_rotary_emb=image_rotary_emb, attention_mask=attention_mask, kv_cache=kv_cache, **attention_kwargs, ) # 3. Feedforward (fused residual + norm + scale/shift) norm_hidden_states, hidden_states = self.norm2( hidden_states, attn_hidden_states, gate_msa.unsqueeze(1), shift_mlp.unsqueeze(1), scale_mlp.unsqueeze(1), ) norm_encoder_hidden_states, encoder_hidden_states = self.norm2_context( encoder_hidden_states, attn_encoder_hidden_states, c_gate_msa.unsqueeze(1), c_shift_mlp.unsqueeze(1), c_scale_mlp.unsqueeze(1), ) ff_output = self.ff(norm_hidden_states) ff_output_context = self.ff(norm_encoder_hidden_states) hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1) encoder_hidden_states = ( encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1) ) return hidden_states, encoder_hidden_states class GlmImageRotaryPosEmbed(nn.Module): def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None: super().__init__() self.dim = dim self.patch_size = patch_size self.theta = theta def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: batch_size, num_channels, height, width = hidden_states.shape height, width = height // self.patch_size, width // self.patch_size device = hidden_states.device dim_h, dim_w = self.dim // 2, self.dim // 2 h_inv_freq = 1.0 / ( self.theta ** ( torch.arange(0, dim_h, 2, dtype=torch.float32, device=device)[ : (dim_h // 2) ].float() / dim_h ) ) w_inv_freq = 1.0 / ( self.theta ** ( torch.arange(0, dim_w, 2, dtype=torch.float32, device=device)[ : (dim_w // 2) ].float() / dim_w ) ) h_seq = torch.arange(height, device=device) w_seq = torch.arange(width, device=device) freqs_h = torch.outer(h_seq, h_inv_freq) freqs_w = torch.outer(w_seq, w_inv_freq) # Create position matrices for height and width # [height, 1, dim//4] and [1, width, dim//4] freqs_h = freqs_h.unsqueeze(1) freqs_w = freqs_w.unsqueeze(0) # Broadcast freqs_h and freqs_w to [height, width, dim//4] freqs_h = freqs_h.expand(height, width, -1) freqs_w = freqs_w.expand(height, width, -1) # Concatenate along last dimension to get [height, width, dim//2] freqs = torch.cat([freqs_h, freqs_w], dim=-1) freqs = freqs.reshape(height * width, -1) # [height * width, dim//2] return (freqs.cos(), freqs.sin()) class GlmImageAdaLayerNormContinuous(nn.Module): """ GlmImage-only final AdaLN: LN(x) -> Linear(cond) -> chunk -> affine. Matches Megatron: **no activation** before the Linear on conditioning embedding. """ def __init__( self, embedding_dim: int, conditioning_embedding_dim: int, elementwise_affine: bool = True, eps: float = 1e-5, bias: bool = True, norm_type: str = "layer_norm", ): super().__init__() self.linear = nn.Linear( conditioning_embedding_dim, embedding_dim * 2, bias=bias ) if norm_type == "layer_norm": self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias) # For now, don’t replace this with sglang’s LayerNorm # because the model doesn’t have this parameter and it will break model loading elif norm_type == "rms_norm": self.norm = nn.RMSNorm(embedding_dim, eps, elementwise_affine) else: raise ValueError(f"unknown norm_type {norm_type}") def forward( self, x: torch.Tensor, conditioning_embedding: torch.Tensor ) -> torch.Tensor: # *** NO SiLU here *** emb = self.linear(conditioning_embedding.to(x.dtype)) scale, shift = torch.chunk(emb, 2, dim=1) x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] return x class GlmImageTransformer2DModel(CachableDiT, OffloadableDiTMixin): r""" Args: patch_size (`int`, defaults to `2`): The size of the patches to use in the patch embedding layer. in_channels (`int`, defaults to `16`): The number of channels in the input. num_layers (`int`, defaults to `30`): The number of layers of Transformer blocks to use. attention_head_dim (`int`, defaults to `40`): The number of channels in each head. num_attention_heads (`int`, defaults to `64`): The number of heads to use for multi-head attention. out_channels (`int`, defaults to `16`): The number of channels in the output. text_embed_dim (`int`, defaults to `1472`): Input dimension of text embeddings from the text encoder. time_embed_dim (`int`, defaults to `512`): Output dimension of timestep embeddings. condition_dim (`int`, defaults to `256`): The embedding dimension of the input SDXL-style resolution conditions (original_size, target_size, crop_coords). pos_embed_max_size (`int`, defaults to `128`): The maximum resolution of the positional embeddings, from which slices of shape `H x W` are taken and added to input patched latents, where `H` and `W` are the latent height and width respectively. A value of 128 means that the maximum supported height and width for image generation is `128 * vae_scale_factor * patch_size => 128 * 8 * 2 => 2048`. sample_size (`int`, defaults to `128`): The base resolution of input latents. If height/width is not provided during generation, this value is used to determine the resolution as `sample_size * vae_scale_factor => 128 * 8 => 1024` """ def __init__( self, config: GlmImageDitConfig, hf_config: dict[str, Any], quant_config: QuantizationConfig | None = None, ): super().__init__(config=config, hf_config=hf_config) self.config_data = config # Store config arch_config = config.arch_config self.in_channels = arch_config.in_channels self.out_channels = arch_config.out_channels self.patch_size = arch_config.patch_size self.num_layers = arch_config.num_layers self.attention_head_dim = arch_config.attention_head_dim self.num_attention_heads = arch_config.num_attention_heads self.text_embed_dim = arch_config.text_embed_dim self.time_embed_dim = arch_config.time_embed_dim # GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords # Each of these are sincos embeddings of shape 2 * condition_dim pooled_projection_dim = 2 * 2 * arch_config.condition_dim inner_dim = arch_config.num_attention_heads * arch_config.attention_head_dim # 1. RoPE self.rotary_emb = GlmImageRotaryPosEmbed( arch_config.attention_head_dim, arch_config.patch_size, theta=10000.0 ) # 2. Patch & Text-timestep embedding self.image_projector = GlmImageImageProjector( arch_config.in_channels, inner_dim, arch_config.patch_size ) self.glyph_projector = FeedForward( arch_config.text_embed_dim, inner_dim, inner_dim=inner_dim, activation_fn="gelu", ) self.prior_token_embedding = nn.Embedding( arch_config.prior_vq_quantizer_codebook_size, inner_dim ) self.prior_projector = FeedForward( inner_dim, inner_dim, inner_dim=inner_dim, activation_fn="linear-silu" ) self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings( embedding_dim=arch_config.time_embed_dim, condition_dim=arch_config.condition_dim, pooled_projection_dim=pooled_projection_dim, timesteps_dim=arch_config.time_embed_dim, ) # 3. Transformer blocks self._supported_attention_backends = arch_config._supported_attention_backends self.transformer_blocks = nn.ModuleList( [ GlmImageTransformerBlock( inner_dim, arch_config.num_attention_heads, arch_config.attention_head_dim, arch_config.time_embed_dim, supported_attention_backends=self._supported_attention_backends, prefix=f"transformer_blocks.{i}", quant_config=quant_config, ) for i in range(arch_config.num_layers) ] ) # 4. Output projection self.norm_out = GlmImageAdaLayerNormContinuous( inner_dim, arch_config.time_embed_dim, elementwise_affine=False ) self.proj_out = nn.Linear( inner_dim, arch_config.patch_size * arch_config.patch_size * arch_config.out_channels, bias=True, ) self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, prior_token_id: torch.Tensor, prior_token_drop: torch.Tensor, timestep: torch.LongTensor, target_size: torch.Tensor, crop_coords: torch.Tensor, attention_kwargs: Optional[Dict[str, Any]] = None, attention_mask: Optional[torch.Tensor] = None, kv_caches: Optional[GlmImageKVCache] = None, kv_caches_mode: Optional[str] = None, freqs_cis: Optional[ Union[ Tuple[torch.Tensor, torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]], ] ] = None, ### guidance: torch.Tensor = None, # TODO: this should probably be removed ) -> Tuple[torch.Tensor]: if kv_caches is not None: kv_caches.set_mode(kv_caches_mode) batch_size, num_channels, height, width = hidden_states.shape timestep -= 1.0 if isinstance(encoder_hidden_states, list): encoder_hidden_states = encoder_hidden_states[0] # 1. RoPE image_rotary_emb = freqs_cis if image_rotary_emb is None: image_rotary_emb = self.rotary_emb(hidden_states) # 2. Patch & Timestep embeddings p = self.config.patch_size post_patch_height = height // p post_patch_width = width // p hidden_states = self.image_projector(hidden_states) encoder_hidden_states = self.glyph_projector(encoder_hidden_states) prior_embedding = self.prior_token_embedding(prior_token_id) prior_embedding[prior_token_drop] *= 0.0 prior_hidden_states = self.prior_projector(prior_embedding) # SP: when latents are H-sharded, hidden_states has fewer patches than prior_hidden_states. # Shard prior_hidden_states along seq dim to match (prior is row-major, same as latent patches). if ( get_sp_world_size() > 1 and prior_hidden_states.shape[1] != hidden_states.shape[1] ): rank = get_sp_parallel_rank() sp_world_size = get_sp_world_size() chunk = prior_hidden_states.shape[1] // sp_world_size prior_hidden_states = prior_hidden_states[ :, rank * chunk : (rank + 1) * chunk, : ] hidden_states = hidden_states + prior_hidden_states temb = self.time_condition_embed( timestep, target_size, crop_coords, hidden_states.dtype ) temb = F.silu(temb) # 3. Transformer blocks for idx, block in enumerate(self.transformer_blocks): hidden_states, encoder_hidden_states = block( hidden_states, encoder_hidden_states, temb, image_rotary_emb, attention_mask, attention_kwargs, kv_cache=kv_caches[idx] if kv_caches is not None else None, ) # 4. Output norm & projection hidden_states = self.norm_out(hidden_states, temb) hidden_states = self.proj_out(hidden_states) # 5. Unpatchify hidden_states = hidden_states.reshape( batch_size, post_patch_height, post_patch_width, -1, p, p ) output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) return output.float() # float() # reference: https://github.com/zRzRzRzRzRzRzR/diffusers/blob/6cfc83b4abc5b083fef56a18ec4700f48ba3aaba/src/diffusers/pipelines/glm_image/pipeline_glm_image.py#L737 EntryClass = GlmImageTransformer2DModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/helios.py ================================================ # SPDX-License-Identifier: Apache-2.0 # Adapted from Helios diffusers transformer: # https://github.com/BestWishYsh/Helios """ Helios Transformer 3D model for video generation. Implements the HeliosTransformer3DModel with multi-term memory patches, 3D rotary position embeddings, and per-block scale-shift modulation. """ import math from functools import lru_cache from typing import Any import torch import torch.nn as nn import torch.nn.functional as F from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig from sglang.multimodal_gen.runtime.distributed import ( divide, get_sp_world_size, get_tp_world_size, ) from sglang.multimodal_gen.runtime.distributed.communication_op import ( sequence_model_parallel_all_gather, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_group from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.layernorm import ( FP32LayerNorm, RMSNorm, tensor_parallel_rms_norm, ) from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import ( ModulateProjection, PatchEmbed, TimestepEmbedder, ) from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # --------------------------------------------------------------------------- # Utility functions # --------------------------------------------------------------------------- def pad_for_3d_conv(x, kernel_size): """Pad input to make it divisible by kernel_size using replicate mode.""" b, c, t, h, w = x.shape pt, ph, pw = kernel_size pad_t = (pt - (t % pt)) % pt pad_h = (ph - (h % ph)) % ph pad_w = (pw - (w % pw)) % pw return F.pad(x, (0, pad_w, 0, pad_h, 0, pad_t), mode="replicate") def center_down_sample_3d(x, kernel_size): """Average pooling for 3D downsampling.""" return F.avg_pool3d(x, kernel_size, stride=kernel_size) def apply_rotary_emb_transposed(hidden_states, freqs_cis): """Apply rotary positional embeddings with transposed cos/sin format.""" x_1, x_2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1) out = torch.empty_like(hidden_states) out[..., 0::2] = x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2] out[..., 1::2] = x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2] return out.type_as(hidden_states) # --------------------------------------------------------------------------- # Output norm # --------------------------------------------------------------------------- class HeliosOutputNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.scale_shift_table = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) self.norm = FP32LayerNorm(dim, eps, elementwise_affine=False) def forward(self, hidden_states, temb, original_context_length): temb = temb[:, -original_context_length:, :] shift, scale = ( self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2) ).chunk(2, dim=2) shift = shift.squeeze(2).to(hidden_states.device) scale = scale.squeeze(2).to(hidden_states.device) hidden_states = hidden_states[:, -original_context_length:, :] hidden_states = ( self.norm(hidden_states.float()) * (1 + scale) + shift ).type_as(hidden_states) return hidden_states # --------------------------------------------------------------------------- # Rotary Positional Embedding (3D) # --------------------------------------------------------------------------- class HeliosRotaryPosEmbed(nn.Module): """3D rotary position embeddings for (time, height, width).""" def __init__(self, rope_dim, theta): super().__init__() self.DT, self.DY, self.DX = rope_dim self.theta = theta # Store as plain attributes (not buffers) to avoid meta-device issues # during FSDP loading. They'll be re-created on the correct device in forward. self._freqs_base_t = None self._freqs_base_y = None self._freqs_base_x = None def _get_freqs_base(self, dim): return 1.0 / ( self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32)[: (dim // 2)] / dim) ) def _ensure_freqs_base(self, device): """Lazily create frequency bases on the correct device.""" if self._freqs_base_t is None or self._freqs_base_t.device != device: self._freqs_base_t = self._get_freqs_base(self.DT).to(device) self._freqs_base_y = self._get_freqs_base(self.DY).to(device) self._freqs_base_x = self._get_freqs_base(self.DX).to(device) @torch.no_grad() def get_frequency_batched(self, freqs_base, pos): freqs = torch.einsum("d,bthw->dbthw", freqs_base, pos) freqs = freqs.repeat_interleave(2, dim=0) return freqs.cos(), freqs.sin() @torch.no_grad() @lru_cache(maxsize=32) def _get_spatial_meshgrid(self, height, width, device_str): device = torch.device(device_str) grid_y_coords = torch.arange(height, device=device, dtype=torch.float32) grid_x_coords = torch.arange(width, device=device, dtype=torch.float32) grid_y, grid_x = torch.meshgrid(grid_y_coords, grid_x_coords, indexing="ij") return grid_y, grid_x @torch.no_grad() def forward(self, frame_indices, height, width, device): self._ensure_freqs_base(device) batch_size = frame_indices.shape[0] num_frames = frame_indices.shape[1] frame_indices = frame_indices.to(device=device, dtype=torch.float32) grid_y, grid_x = self._get_spatial_meshgrid(height, width, str(device)) grid_t = frame_indices[:, :, None, None].expand( batch_size, num_frames, height, width ) grid_y_batch = grid_y[None, None, :, :].expand(batch_size, num_frames, -1, -1) grid_x_batch = grid_x[None, None, :, :].expand(batch_size, num_frames, -1, -1) freqs_cos_t, freqs_sin_t = self.get_frequency_batched( self._freqs_base_t, grid_t ) freqs_cos_y, freqs_sin_y = self.get_frequency_batched( self._freqs_base_y, grid_y_batch ) freqs_cos_x, freqs_sin_x = self.get_frequency_batched( self._freqs_base_x, grid_x_batch ) result = torch.cat( [ freqs_cos_t, freqs_cos_y, freqs_cos_x, freqs_sin_t, freqs_sin_y, freqs_sin_x, ], dim=0, ) return result.permute(1, 0, 2, 3, 4) # --------------------------------------------------------------------------- # Condition Embedder # --------------------------------------------------------------------------- class HeliosTimeTextEmbedding(nn.Module): """Condition embedder combining timestep and text embeddings.""" def __init__(self, dim, time_freq_dim, time_proj_dim, text_embed_dim): super().__init__() self.time_embedder = TimestepEmbedder( dim, frequency_embedding_size=time_freq_dim, act_layer="silu" ) self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu") self.text_embedder = MLP( text_embed_dim, dim, dim, bias=True, act_type="gelu_pytorch_tanh" ) def forward( self, timestep, encoder_hidden_states, is_return_encoder_hidden_states=True ): temb = self.time_embedder(timestep) timestep_proj = self.time_modulation(temb) if encoder_hidden_states is not None and is_return_encoder_hidden_states: encoder_hidden_states = self.text_embedder(encoder_hidden_states) return temb, timestep_proj, encoder_hidden_states # --------------------------------------------------------------------------- # Self-Attention for Helios # --------------------------------------------------------------------------- class HeliosSelfAttention(nn.Module): """Self-attention with RMSNorm Q/K, optional history key amplification.""" def __init__( self, dim: int, num_heads: int, eps: float = 1e-6, is_amplify_history: bool = False, history_scale_mode: str = "per_head", quant_config: QuantizationConfig | None = None, ): super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads tp_size = get_tp_world_size() self.local_num_heads = divide(num_heads, tp_size) self.to_q = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config ) self.to_k = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config ) self.to_v = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config ) self.to_out = RowParallelLinear( dim, dim, bias=True, reduce_results=True, quant_config=quant_config ) self.norm_q = RMSNorm(dim, eps=eps) self.norm_k = RMSNorm(dim, eps=eps) self.tp_rmsnorm = tp_size > 1 self.attn = USPAttention( num_heads=self.local_num_heads, head_size=self.head_dim, causal=False, is_cross_attention=False, ) self.is_amplify_history = is_amplify_history if is_amplify_history: if history_scale_mode == "scalar": self.history_key_scale = nn.Parameter(torch.ones(1)) elif history_scale_mode == "per_head": self.history_key_scale = nn.Parameter(torch.ones(num_heads)) else: raise ValueError(f"Unknown history_scale_mode: {history_scale_mode}") self.history_scale_mode = history_scale_mode self.max_scale = 10.0 def forward(self, hidden_states, rotary_emb=None, original_context_length=None): q, _ = self.to_q(hidden_states) k, _ = self.to_k(hidden_states) v, _ = self.to_v(hidden_states) if self.tp_rmsnorm: q = tensor_parallel_rms_norm(q, self.norm_q) k = tensor_parallel_rms_norm(k, self.norm_k) else: q = self.norm_q(q) k = self.norm_k(k) q = q.unflatten(2, (self.local_num_heads, self.head_dim)) k = k.unflatten(2, (self.local_num_heads, self.head_dim)) v = v.unflatten(2, (self.local_num_heads, self.head_dim)) if rotary_emb is not None: q = apply_rotary_emb_transposed(q, rotary_emb) k = apply_rotary_emb_transposed(k, rotary_emb) history_seq_len = ( hidden_states.shape[1] - original_context_length if original_context_length is not None else 0 ) if self.is_amplify_history and original_context_length is not None: if history_seq_len > 0: scale_key = 1.0 + torch.sigmoid(self.history_key_scale) * ( self.max_scale - 1.0 ) if self.history_scale_mode == "per_head": scale_key = scale_key.view(1, 1, -1, 1) k = torch.cat( [k[:, :history_seq_len] * scale_key, k[:, history_seq_len:]], dim=1, ) x = self.attn(q, k, v, num_replicated_prefix=history_seq_len) x = x.flatten(2) x, _ = self.to_out(x) return x # --------------------------------------------------------------------------- # Cross-Attention for Helios # --------------------------------------------------------------------------- class HeliosCrossAttention(nn.Module): """Cross-attention with RMSNorm Q/K normalization.""" def __init__( self, dim: int, num_heads: int, eps: float = 1e-6, quant_config: QuantizationConfig | None = None, ): super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads tp_size = get_tp_world_size() self.local_num_heads = divide(num_heads, tp_size) self.to_q = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config ) self.to_k = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config ) self.to_v = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config ) self.to_out = RowParallelLinear( dim, dim, bias=True, reduce_results=True, quant_config=quant_config ) self.norm_q = RMSNorm(dim, eps=eps) self.norm_k = RMSNorm(dim, eps=eps) self.tp_rmsnorm = tp_size > 1 self.attn = USPAttention( num_heads=self.local_num_heads, head_size=self.head_dim, causal=False, skip_sequence_parallel=True, ) def forward(self, hidden_states, encoder_hidden_states): q, _ = self.to_q(hidden_states) k, _ = self.to_k(encoder_hidden_states) v, _ = self.to_v(encoder_hidden_states) if self.tp_rmsnorm: q = tensor_parallel_rms_norm(q, self.norm_q) k = tensor_parallel_rms_norm(k, self.norm_k) else: q = self.norm_q(q) k = self.norm_k(k) q = q.unflatten(2, (self.local_num_heads, self.head_dim)) k = k.unflatten(2, (self.local_num_heads, self.head_dim)) v = v.unflatten(2, (self.local_num_heads, self.head_dim)) x = self.attn(q, k, v) x = x.flatten(2) x, _ = self.to_out(x) return x # --------------------------------------------------------------------------- # Transformer Block # --------------------------------------------------------------------------- class HeliosTransformerBlock(nn.Module): """ Single transformer block with self-attention, cross-attention, FFN, and scale-shift modulation from timestep embeddings. """ def __init__( self, dim: int, ffn_dim: int, num_heads: int, cross_attn_norm: bool = True, eps: float = 1e-6, guidance_cross_attn: bool = True, is_amplify_history: bool = False, history_scale_mode: str = "per_head", quant_config: QuantizationConfig | None = None, ): super().__init__() # 1. Self-attention self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) self.attn1 = HeliosSelfAttention( dim=dim, num_heads=num_heads, eps=eps, is_amplify_history=is_amplify_history, history_scale_mode=history_scale_mode, quant_config=quant_config, ) # 2. Cross-attention self.attn2 = HeliosCrossAttention( dim=dim, num_heads=num_heads, eps=eps, quant_config=quant_config, ) self.self_attn_residual_norm = ( FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() ) # 3. Feed-forward self.ffn = MLP( dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config ) self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) # 4. Guidance cross-attention flag self.guidance_cross_attn = guidance_cross_attn def forward( self, hidden_states, encoder_hidden_states, temb, rotary_emb, original_context_length=None, ): if temb.ndim == 4: shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( self.scale_shift_table.unsqueeze(0) + temb.float() ).chunk(6, dim=2) shift_msa = shift_msa.squeeze(2) scale_msa = scale_msa.squeeze(2) gate_msa = gate_msa.squeeze(2) c_shift_msa = c_shift_msa.squeeze(2) c_scale_msa = c_scale_msa.squeeze(2) c_gate_msa = c_gate_msa.squeeze(2) else: shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( self.scale_shift_table + temb.float() ).chunk(6, dim=1) # 1. Self-attention norm_hidden_states = ( self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa ).type_as(hidden_states) attn_output = self.attn1( norm_hidden_states, rotary_emb, original_context_length ) hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as( hidden_states ) # 2. Cross-attention if self.guidance_cross_attn: history_seq_len = hidden_states.shape[1] - original_context_length history_hidden_states, current_hidden_states = torch.split( hidden_states, [history_seq_len, original_context_length], dim=1 ) norm_hidden_states = self.self_attn_residual_norm( current_hidden_states.float() ).type_as(current_hidden_states) attn_output = self.attn2(norm_hidden_states, encoder_hidden_states) current_hidden_states = current_hidden_states + attn_output hidden_states = torch.cat( [history_hidden_states, current_hidden_states], dim=1 ) else: norm_hidden_states = self.self_attn_residual_norm( hidden_states.float() ).type_as(hidden_states) attn_output = self.attn2(norm_hidden_states, encoder_hidden_states) hidden_states = hidden_states + attn_output # 3. Feed-forward norm_hidden_states = ( self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa ).type_as(hidden_states) ff_output = self.ffn(norm_hidden_states) hidden_states = ( hidden_states.float() + ff_output.float() * c_gate_msa ).type_as(hidden_states) return hidden_states # --------------------------------------------------------------------------- # Main model # --------------------------------------------------------------------------- class HeliosTransformer3DModel(CachableDiT, OffloadableDiTMixin): """ Helios Transformer 3D model for video generation. Implements multi-scale history patches, 3D RoPE, and chunked denoising with zero_history_timestep and guidance_cross_attn. """ _fsdp_shard_conditions = HeliosConfig()._fsdp_shard_conditions _compile_conditions = HeliosConfig()._compile_conditions _supported_attention_backends = HeliosConfig()._supported_attention_backends param_names_mapping = HeliosConfig().param_names_mapping reverse_param_names_mapping = HeliosConfig().reverse_param_names_mapping lora_param_names_mapping = HeliosConfig().lora_param_names_mapping def __init__( self, config: HeliosConfig, hf_config: dict[str, Any], quant_config: QuantizationConfig | None = None, ) -> None: super().__init__(config=config, hf_config=hf_config) inner_dim = config.num_attention_heads * config.attention_head_dim self.hidden_size = config.hidden_size self.num_attention_heads = config.num_attention_heads self.in_channels = config.in_channels self.out_channels = config.out_channels self.num_channels_latents = config.num_channels_latents self.patch_size = config.patch_size self.text_len = config.text_len self.inner_dim = inner_dim # Helios-specific config self.zero_history_timestep = config.zero_history_timestep self.has_multi_term_memory_patch = config.has_multi_term_memory_patch self.guidance_cross_attn = config.guidance_cross_attn # 1. Patch & position embedding self.patch_embedding = PatchEmbed( in_chans=config.in_channels, embed_dim=inner_dim, patch_size=config.patch_size, flatten=False, ) # 2. Rotary position embeddings self.rope = HeliosRotaryPosEmbed( rope_dim=config.rope_dim, theta=config.rope_theta ) # 3. Multi-term memory patches if self.has_multi_term_memory_patch: self.patch_short = nn.Conv3d( config.in_channels, inner_dim, kernel_size=config.patch_size, stride=config.patch_size, ) self.patch_mid = nn.Conv3d( config.in_channels, inner_dim, kernel_size=tuple(2 * p for p in config.patch_size), stride=tuple(2 * p for p in config.patch_size), ) self.patch_long = nn.Conv3d( config.in_channels, inner_dim, kernel_size=tuple(4 * p for p in config.patch_size), stride=tuple(4 * p for p in config.patch_size), ) # 4. Condition embeddings self.condition_embedder = HeliosTimeTextEmbedding( dim=inner_dim, time_freq_dim=config.freq_dim, time_proj_dim=inner_dim * 6, text_embed_dim=config.text_dim, ) # 5. Transformer blocks self.blocks = nn.ModuleList( [ HeliosTransformerBlock( dim=inner_dim, ffn_dim=config.ffn_dim, num_heads=config.num_attention_heads, cross_attn_norm=config.cross_attn_norm, eps=config.eps, guidance_cross_attn=config.guidance_cross_attn, is_amplify_history=config.is_amplify_history, history_scale_mode=config.history_scale_mode, quant_config=quant_config, ) for _ in range(config.num_layers) ] ) # 6. Output norm & projection self.norm_out = HeliosOutputNorm(inner_dim, config.eps) self.proj_out = ColumnParallelLinear( inner_dim, config.out_channels * math.prod(config.patch_size), bias=True, gather_output=True, quant_config=quant_config, ) self.cnt = 0 self.__post_init__() self.layer_names = ["blocks"] self.sp_size = get_sp_world_size() def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, # Stage 1 history inputs indices_hidden_states=None, indices_latents_history_short=None, indices_latents_history_mid=None, indices_latents_history_long=None, latents_history_short=None, latents_history_mid=None, latents_history_long=None, **kwargs, ) -> torch.Tensor: orig_dtype = hidden_states.dtype if not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] # Check if sequence parallelism is enabled forward_batch = get_forward_context().forward_batch if forward_batch is not None: sequence_shard_enabled = ( forward_batch.enable_sequence_shard and self.sp_size > 1 ) else: sequence_shard_enabled = False batch_size = hidden_states.shape[0] p_t, p_h, p_w = self.patch_size # 1. Patch embed the noisy latents hidden_states = self.patch_embedding(hidden_states) _, _, post_patch_num_frames, post_patch_height, post_patch_width = ( hidden_states.shape ) if indices_hidden_states is None: indices_hidden_states = ( torch.arange(0, post_patch_num_frames) .unsqueeze(0) .expand(batch_size, -1) ) hidden_states = hidden_states.flatten(2).transpose(1, 2) # 2. Compute rotary embeddings rotary_emb = self.rope( frame_indices=indices_hidden_states, height=post_patch_height, width=post_patch_width, device=hidden_states.device, ) rotary_emb = rotary_emb.flatten(2).transpose(1, 2) original_context_length = hidden_states.shape[1] # Sequence parallelism: shard current tokens and RoPE across SP ranks seq_shard_pad = 0 if sequence_shard_enabled: sp_rank = get_sp_group().rank_in_group seq_len = hidden_states.shape[1] if seq_len % self.sp_size != 0: seq_shard_pad = self.sp_size - (seq_len % self.sp_size) hs_pad = torch.zeros( batch_size, seq_shard_pad, hidden_states.shape[2], dtype=hidden_states.dtype, device=hidden_states.device, ) re_pad = torch.zeros( batch_size, seq_shard_pad, rotary_emb.shape[2], dtype=rotary_emb.dtype, device=rotary_emb.device, ) hidden_states = torch.cat([hidden_states, hs_pad], dim=1) rotary_emb = torch.cat([rotary_emb, re_pad], dim=1) local_seq_len = hidden_states.shape[1] // self.sp_size hidden_states = hidden_states.view( batch_size, self.sp_size, local_seq_len, -1 )[:, sp_rank, :, :].contiguous() rotary_emb = rotary_emb.view(batch_size, self.sp_size, local_seq_len, -1)[ :, sp_rank, :, : ].contiguous() effective_context_length = local_seq_len else: effective_context_length = original_context_length # 3. Process short history if ( latents_history_short is not None and indices_latents_history_short is not None ): latents_history_short = latents_history_short.to(hidden_states) latents_history_short = self.patch_short(latents_history_short) _, _, _, H1, W1 = latents_history_short.shape latents_history_short = latents_history_short.flatten(2).transpose(1, 2) rotary_emb_history_short = self.rope( frame_indices=indices_latents_history_short, height=H1, width=W1, device=latents_history_short.device, ) rotary_emb_history_short = rotary_emb_history_short.flatten(2).transpose( 1, 2 ) hidden_states = torch.cat([latents_history_short, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_short, rotary_emb], dim=1) # 4. Process mid history if latents_history_mid is not None and indices_latents_history_mid is not None: latents_history_mid = latents_history_mid.to(hidden_states) latents_history_mid = pad_for_3d_conv(latents_history_mid, (2, 4, 4)) latents_history_mid = self.patch_mid(latents_history_mid) latents_history_mid = latents_history_mid.flatten(2).transpose(1, 2) rotary_emb_history_mid = self.rope( frame_indices=indices_latents_history_mid, height=H1, width=W1, device=latents_history_mid.device, ) rotary_emb_history_mid = pad_for_3d_conv(rotary_emb_history_mid, (2, 2, 2)) rotary_emb_history_mid = center_down_sample_3d( rotary_emb_history_mid, (2, 2, 2) ) rotary_emb_history_mid = rotary_emb_history_mid.flatten(2).transpose(1, 2) hidden_states = torch.cat([latents_history_mid, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_mid, rotary_emb], dim=1) # 5. Process long history if ( latents_history_long is not None and indices_latents_history_long is not None ): latents_history_long = latents_history_long.to(hidden_states) latents_history_long = pad_for_3d_conv(latents_history_long, (4, 8, 8)) latents_history_long = self.patch_long(latents_history_long) latents_history_long = latents_history_long.flatten(2).transpose(1, 2) rotary_emb_history_long = self.rope( frame_indices=indices_latents_history_long, height=H1, width=W1, device=latents_history_long.device, ) rotary_emb_history_long = pad_for_3d_conv( rotary_emb_history_long, (4, 4, 4) ) rotary_emb_history_long = center_down_sample_3d( rotary_emb_history_long, (4, 4, 4) ) rotary_emb_history_long = rotary_emb_history_long.flatten(2).transpose(1, 2) hidden_states = torch.cat([latents_history_long, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_long, rotary_emb], dim=1) history_context_length = hidden_states.shape[1] - effective_context_length # 6. Compute condition embeddings if indices_hidden_states is not None and self.zero_history_timestep: timestep_t0 = torch.zeros( (1,), dtype=timestep.dtype, device=timestep.device ) temb_t0, timestep_proj_t0, _ = self.condition_embedder( timestep_t0, encoder_hidden_states, is_return_encoder_hidden_states=False, ) temb_t0 = temb_t0.unsqueeze(1).expand( batch_size, history_context_length, -1 ) timestep_proj_t0 = ( timestep_proj_t0.unflatten(-1, (6, -1)) .view(1, 6, 1, -1) .expand(batch_size, -1, history_context_length, -1) ) temb, timestep_proj, encoder_hidden_states = self.condition_embedder( timestep, encoder_hidden_states ) timestep_proj = timestep_proj.unflatten(-1, (6, -1)) if indices_hidden_states is not None and not self.zero_history_timestep: main_repeat_size = hidden_states.shape[1] else: main_repeat_size = effective_context_length temb = temb.view(batch_size, 1, -1).expand(batch_size, main_repeat_size, -1) timestep_proj = timestep_proj.view(batch_size, 6, 1, -1).expand( batch_size, 6, main_repeat_size, -1 ) if indices_hidden_states is not None and self.zero_history_timestep: temb = torch.cat([temb_t0, temb], dim=1) timestep_proj = torch.cat([timestep_proj_t0, timestep_proj], dim=2) if timestep_proj.ndim == 4: timestep_proj = timestep_proj.permute(0, 2, 1, 3) # 7. Transformer blocks hidden_states = hidden_states.contiguous() encoder_hidden_states = encoder_hidden_states.contiguous() rotary_emb = rotary_emb.contiguous() for block in self.blocks: hidden_states = block( hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, effective_context_length, ) self.cnt += 1 # SP: all-gather current tokens before output if sequence_shard_enabled: current_tokens = hidden_states[:, -local_seq_len:, :].contiguous() current_tokens = sequence_model_parallel_all_gather(current_tokens, dim=1) if seq_shard_pad > 0: current_tokens = current_tokens[:, :original_context_length, :] hidden_states = current_tokens # Re-create temb for norm_out (all current tokens share same timestep) temb = temb[:, :1, :].expand(batch_size, original_context_length, -1) # 8. Output norm & projection hidden_states = self.norm_out(hidden_states, temb, original_context_length) hidden_states, _ = self.proj_out(hidden_states) # 9. Unpatchify hidden_states = hidden_states.reshape( batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1, ) hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) return output EntryClass = HeliosTransformer3DModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py ================================================ # Copied and adapted from: https://github.com/Tencent-Hunyuan/Hunyuan3D-2 from __future__ import annotations import math from dataclasses import dataclass from typing import List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange from sglang.multimodal_gen.configs.models.dits.hunyuan3d import ( Hunyuan3DDiTArchConfig, Hunyuan3DDiTConfig, ) from sglang.multimodal_gen.runtime.distributed import divide from sglang.multimodal_gen.runtime.distributed.parallel_state import get_tp_world_size from sglang.multimodal_gen.runtime.layers.attention import LocalAttention from sglang.multimodal_gen.runtime.layers.linear import ( MergedColumnParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class MixedRowParallelLinear(RowParallelLinear): """RowParallel for inputs concatenated from multiple separately-sharded sources.""" def __init__(self, input_sizes: list[int], output_size: int, **kwargs): self.input_sizes = input_sizes super().__init__(sum(input_sizes), output_size, **kwargs) def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): input_dim = getattr(param, "input_dim", None) if input_dim is not None: shards = [] offset = 0 for sz in self.input_sizes: part = loaded_weight.narrow(input_dim, offset, sz) per_rank = sz // self.tp_size shard = part.narrow(input_dim, self.tp_rank * per_rank, per_rank) shards.append(shard) offset += sz param.data.copy_(torch.cat(shards, dim=input_dim)) else: param.data.copy_(loaded_weight) def _flux_timestep_embedding( t: torch.Tensor, dim, max_period=10000, time_factor: float = 1000.0 ): """Create sinusoidal timestep embeddings for Flux-style model.""" t = time_factor * t half = dim // 2 freqs = torch.exp( -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half ).to(t.device) args = t[:, None].float() * freqs[None] embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) if dim % 2: embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) if torch.is_floating_point(t): embedding = embedding.to(t) return embedding class _FluxGELU(nn.Module): def __init__(self, approximate="tanh"): super().__init__() self.approximate = approximate def forward(self, x: torch.Tensor) -> torch.Tensor: return F.gelu(x, approximate=self.approximate) class _FluxMLPEmbedder(nn.Module): def __init__(self, in_dim: int, hidden_dim: int): super().__init__() self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True) self.silu = nn.SiLU() self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.out_layer(self.silu(self.in_layer(x))) class _FluxRMSNorm(nn.Module): def __init__(self, dim: int): super().__init__() self.scale = nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor): x_dtype = x.dtype x = x.float() rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6) return (x * rrms).to(dtype=x_dtype) * self.scale class _FluxQKNorm(nn.Module): def __init__(self, dim: int): super().__init__() self.query_norm = _FluxRMSNorm(dim) self.key_norm = _FluxRMSNorm(dim) def forward( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: q = self.query_norm(q) k = self.key_norm(k) return q.to(v), k.to(v) class _FluxSelfAttention(nn.Module): def __init__( self, dim: int, num_heads: int = 8, qkv_bias: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, ): super().__init__() tp_size = get_tp_world_size() self.num_heads = num_heads self.local_num_heads = divide(num_heads, tp_size) self.head_dim = dim // num_heads self.qkv = MergedColumnParallelLinear( dim, [dim, dim, dim], bias=qkv_bias, gather_output=False ) self.norm = _FluxQKNorm(self.head_dim) self.proj = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True) if supported_attention_backends is None: supported_attention_backends = { AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA, } self.local_attn = LocalAttention( num_heads=self.local_num_heads, head_size=self.head_dim, causal=False, supported_attention_backends=supported_attention_backends, ) def forward(self, x: torch.Tensor, pe: torch.Tensor) -> torch.Tensor: qkv, _ = self.qkv(x) B, L, _ = qkv.shape qkv = qkv.view(B, L, 3, self.local_num_heads, self.head_dim) q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] q = q.transpose(1, 2) k = k.transpose(1, 2) v_for_norm = v.transpose(1, 2) q, k = self.norm(q, k, v_for_norm) q = q.transpose(1, 2) k = k.transpose(1, 2) x = self.local_attn(q, k, v) x = x.flatten(2) x, _ = self.proj(x) return x @dataclass class _FluxModulationOut: shift: torch.Tensor scale: torch.Tensor gate: torch.Tensor class _FluxModulation(nn.Module): def __init__(self, dim: int, double: bool): super().__init__() self.is_double = double self.multiplier = 6 if double else 3 self.lin = nn.Linear(dim, self.multiplier * dim, bias=True) def forward( self, vec: torch.Tensor ) -> Tuple[_FluxModulationOut, Optional[_FluxModulationOut]]: out = self.lin(F.silu(vec))[:, None, :] out = out.chunk(self.multiplier, dim=-1) return ( _FluxModulationOut(*out[:3]), _FluxModulationOut(*out[3:]) if self.is_double else None, ) class _FluxDoubleStreamBlock(nn.Module): def __init__( self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, ): super().__init__() mlp_hidden_dim = int(hidden_size * mlp_ratio) tp_size = get_tp_world_size() self.num_heads = num_heads self.local_num_heads = divide(num_heads, tp_size) self.hidden_size = hidden_size self.head_dim = hidden_size // num_heads self.img_mod = _FluxModulation(hidden_size, double=True) self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.img_attn = _FluxSelfAttention( dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, supported_attention_backends=supported_attention_backends, ) self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.img_mlp = MLP(hidden_size, mlp_hidden_dim, act_type="gelu_pytorch_tanh") self.txt_mod = _FluxModulation(hidden_size, double=True) self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.txt_attn = _FluxSelfAttention( dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, supported_attention_backends=supported_attention_backends, ) self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, act_type="gelu_pytorch_tanh") if supported_attention_backends is None: supported_attention_backends = { AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA, } self.local_attn_joint = LocalAttention( num_heads=self.local_num_heads, head_size=self.head_dim, causal=False, supported_attention_backends=supported_attention_backends, ) def forward( self, img: torch.Tensor, txt: torch.Tensor, vec: torch.Tensor, pe: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: img_mod1, img_mod2 = self.img_mod(vec) txt_mod1, txt_mod2 = self.txt_mod(vec) img_modulated = self.img_norm1(img) img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift B, img_L, _ = img_modulated.shape img_qkv, _ = self.img_attn.qkv(img_modulated) img_qkv = img_qkv.view(B, img_L, 3, self.local_num_heads, self.head_dim) img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2] img_q_t = img_q.transpose(1, 2) img_k_t = img_k.transpose(1, 2) img_v_t = img_v.transpose(1, 2) img_q_t, img_k_t = self.img_attn.norm(img_q_t, img_k_t, img_v_t) img_q = img_q_t.transpose(1, 2) img_k = img_k_t.transpose(1, 2) txt_modulated = self.txt_norm1(txt) txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift txt_L = txt_modulated.shape[1] txt_qkv, _ = self.txt_attn.qkv(txt_modulated) txt_qkv = txt_qkv.view(B, txt_L, 3, self.local_num_heads, self.head_dim) txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2] txt_q_t = txt_q.transpose(1, 2) txt_k_t = txt_k.transpose(1, 2) txt_v_t = txt_v.transpose(1, 2) txt_q_t, txt_k_t = self.txt_attn.norm(txt_q_t, txt_k_t, txt_v_t) txt_q = txt_q_t.transpose(1, 2) txt_k = txt_k_t.transpose(1, 2) q = torch.cat((txt_q, img_q), dim=1) k = torch.cat((txt_k, img_k), dim=1) v = torch.cat((txt_v, img_v), dim=1) attn = self.local_attn_joint(q, k, v) attn = attn.flatten(2) txt_attn, img_attn = attn[:, :txt_L], attn[:, txt_L:] img_proj, _ = self.img_attn.proj(img_attn) img = img + img_mod1.gate * img_proj img = img + img_mod2.gate * self.img_mlp( (1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift ) txt_proj, _ = self.txt_attn.proj(txt_attn) txt = txt + txt_mod1.gate * txt_proj txt = txt + txt_mod2.gate * self.txt_mlp( (1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift ) return img, txt class _FluxSingleStreamBlock(nn.Module): """ A DiT block with parallel linear layers as described in https://arxiv.org/abs/2302.05442 and adapted modulation interface. """ def __init__( self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0, qk_scale: Optional[float] = None, supported_attention_backends: set[AttentionBackendEnum] | None = None, ): super().__init__() tp_size = get_tp_world_size() self.hidden_dim = hidden_size self.num_heads = num_heads self.local_num_heads = divide(num_heads, tp_size) self.head_dim = hidden_size // num_heads self.tp_size = tp_size self.mlp_hidden_dim = int(hidden_size * mlp_ratio) self.linear1 = MergedColumnParallelLinear( hidden_size, [hidden_size, hidden_size, hidden_size, self.mlp_hidden_dim], bias=True, gather_output=False, ) self.linear2 = MixedRowParallelLinear( [hidden_size, self.mlp_hidden_dim], hidden_size, bias=True, input_is_parallel=True, ) self.norm = _FluxQKNorm(self.head_dim) self.hidden_size = hidden_size self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.mlp_act = _FluxGELU(approximate="tanh") self.modulation = _FluxModulation(hidden_size, double=False) if supported_attention_backends is None: supported_attention_backends = { AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA, } self.local_attn = LocalAttention( num_heads=self.local_num_heads, head_size=self.head_dim, causal=False, supported_attention_backends=supported_attention_backends, ) def forward( self, x: torch.Tensor, vec: torch.Tensor, pe: torch.Tensor ) -> torch.Tensor: mod, _ = self.modulation(vec) x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift linear1_out, _ = self.linear1(x_mod) local_qkv_dim = 3 * self.head_dim * self.local_num_heads local_mlp_dim = self.mlp_hidden_dim // self.tp_size qkv, mlp = torch.split(linear1_out, [local_qkv_dim, local_mlp_dim], dim=-1) B, L, _ = qkv.shape qkv = qkv.view(B, L, 3, self.local_num_heads, self.head_dim) q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] q_t = q.transpose(1, 2) k_t = k.transpose(1, 2) v_t = v.transpose(1, 2) q_t, k_t = self.norm(q_t, k_t, v_t) q = q_t.transpose(1, 2) k = k_t.transpose(1, 2) attn = self.local_attn(q, k, v) attn = attn.flatten(2) output, _ = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) return x + mod.gate * output class _FluxLastLayer(nn.Module): def __init__(self, hidden_size: int, patch_size: int, out_channels: int): super().__init__() self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.linear = nn.Linear( hidden_size, patch_size * patch_size * out_channels, bias=True ) self.adaLN_modulation = nn.Sequential( nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True) ) def forward(self, x: torch.Tensor, vec: torch.Tensor) -> torch.Tensor: shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1) x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :] x = self.linear(x) return x class Hunyuan3D2DiT(CachableDiT, OffloadableDiTMixin): """Hunyuan3D DiT model (Flux-style architecture for Hunyuan3D-2.0).""" _aliases = ["hy3dgen.shapegen.models.Hunyuan3DDiT"] param_names_mapping = Hunyuan3DDiTConfig().param_names_mapping @classmethod def build_config_from_params(cls, params: dict) -> Hunyuan3DDiTConfig: """Build a DiTConfig from YAML-style parameter dict.""" field_mapping = { "num_heads": "num_attention_heads", "depth": "num_layers", "depth_single_blocks": "num_single_layers", } arch_kwargs = {} for k, v in params.items(): if k in ("ckpt_path", "supported_attention_backends"): continue mapped = field_mapping.get(k, k) if k == "axes_dim" and isinstance(v, list): v = tuple(v) arch_kwargs[mapped] = v return Hunyuan3DDiTConfig(arch_config=Hunyuan3DDiTArchConfig(**arch_kwargs)) def __init__( self, config: Hunyuan3DDiTConfig, hf_config: dict | None = None, **kwargs, ): super().__init__(config=config, hf_config=hf_config or {}, **kwargs) arch = config.arch_config in_channels = arch.in_channels context_in_dim = arch.context_in_dim hidden_size = arch.hidden_size mlp_ratio = arch.mlp_ratio num_heads = arch.num_attention_heads depth = arch.num_layers depth_single_blocks = arch.num_single_layers axes_dim = list(arch.axes_dim) theta = arch.theta qkv_bias = arch.qkv_bias time_factor = arch.time_factor guidance_embed = arch.guidance_embed supported_attention_backends = arch._supported_attention_backends self.in_channels = in_channels self.context_in_dim = context_in_dim self.hidden_size = hidden_size self.mlp_ratio = mlp_ratio self.num_heads = num_heads self.num_attention_heads = num_heads self.depth = depth self.depth_single_blocks = depth_single_blocks self.axes_dim = axes_dim self.theta = theta self.qkv_bias = qkv_bias self.time_factor = time_factor self.out_channels = self.in_channels self.num_channels_latents = self.in_channels self.guidance_embed = guidance_embed if hidden_size % num_heads != 0: raise ValueError( f"Hidden size {hidden_size} must be divisible by num_heads {num_heads}" ) pe_dim = hidden_size // num_heads if sum(axes_dim) != pe_dim: raise ValueError(f"Got {axes_dim} but expected positional dim {pe_dim}") self.latent_in = nn.Linear(self.in_channels, self.hidden_size, bias=True) self.time_in = _FluxMLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) self.cond_in = nn.Linear(context_in_dim, self.hidden_size) self.guidance_in = ( _FluxMLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if guidance_embed else nn.Identity() ) self.double_blocks = nn.ModuleList( [ _FluxDoubleStreamBlock( self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, supported_attention_backends=supported_attention_backends, ) for _ in range(depth) ] ) self.single_blocks = nn.ModuleList( [ _FluxSingleStreamBlock( self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio, supported_attention_backends=supported_attention_backends, ) for _ in range(depth_single_blocks) ] ) self.final_layer = _FluxLastLayer(self.hidden_size, 1, self.out_channels) # OffloadableDiTMixin self.layer_names = ["double_blocks", "single_blocks"] def forward( self, x, t, contexts, **kwargs, ) -> torch.Tensor: """Forward pass for denoising.""" cond = contexts["main"] latent = self.latent_in(x) t_emb = _flux_timestep_embedding(t, 256, self.time_factor).to( dtype=latent.dtype ) vec = self.time_in(t_emb) if self.guidance_embed: guidance = kwargs.get("guidance", None) if guidance is None: raise ValueError( "Didn't get guidance strength for guidance distilled model." ) vec = vec + self.guidance_in( _flux_timestep_embedding(guidance, 256, self.time_factor) ) cond = self.cond_in(cond) pe = None # Double blocks for i, block in enumerate(self.double_blocks): latent, cond = block(img=latent, txt=cond, vec=vec, pe=pe) latent = torch.cat((cond, latent), 1) # Single blocks for i, block in enumerate(self.single_blocks): latent = block(latent, vec=vec, pe=pe) latent = latent[:, cond.shape[1] :, ...] latent = self.final_layer(latent, vec) return latent import copy import json import os as _os from diffusers.models import UNet2DConditionModel from diffusers.models.attention_processor import Attention as DiffusersAttention from diffusers.models.transformers.transformer_2d import BasicTransformerBlock def _chunked_feed_forward( ff: nn.Module, hidden_states: torch.Tensor, chunk_dim: int, chunk_size: int ): """Feed forward with chunking to save memory.""" if hidden_states.shape[chunk_dim] % chunk_size != 0: raise ValueError( f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]}" f"has to be divisible by chunk size: {chunk_size}." f" Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." ) num_chunks = hidden_states.shape[chunk_dim] // chunk_size ff_output = torch.cat( [ff(hid_slice) for hid_slice in hidden_states.chunk(num_chunks, dim=chunk_dim)], dim=chunk_dim, ) return ff_output class SGLangAttentionWrapper(torch.nn.Module): """Drop-in replacement for DiffusersAttention that uses sglang's attention backend.""" _SUPPORTED_BACKENDS = {AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA} def __init__( self, query_dim: int, heads: int = 8, dim_head: int = 64, dropout: float = 0.0, bias: bool = False, cross_attention_dim: int | None = None, out_bias: bool = True, ) -> None: super().__init__() self.inner_dim = dim_head * heads self.heads = heads self.dim_head = dim_head self.query_dim = query_dim cross_attention_dim = cross_attention_dim or query_dim self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias) self.to_k = nn.Linear(cross_attention_dim, self.inner_dim, bias=bias) self.to_v = nn.Linear(cross_attention_dim, self.inner_dim, bias=bias) self.to_out = nn.ModuleList( [nn.Linear(self.inner_dim, query_dim, bias=out_bias), nn.Dropout(dropout)] ) from sglang.multimodal_gen.runtime.layers.attention.selector import ( get_attn_backend, ) attn_backend = get_attn_backend( dim_head, torch.float16, self._SUPPORTED_BACKENDS ) impl_cls = attn_backend.get_impl_cls() self.attn_impl = impl_cls( num_heads=heads, head_size=dim_head, softmax_scale=dim_head**-0.5, num_kv_heads=heads, causal=False, ) self._attn_backend_name = attn_backend.get_enum().name def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: if encoder_hidden_states is None: encoder_hidden_states = hidden_states B, N_q, _ = hidden_states.shape _, N_kv, _ = encoder_hidden_states.shape q = self.to_q(hidden_states).view(B, N_q, self.heads, self.dim_head) k = self.to_k(encoder_hidden_states).view(B, N_kv, self.heads, self.dim_head) v = self.to_v(encoder_hidden_states).view(B, N_kv, self.heads, self.dim_head) from sglang.multimodal_gen.runtime.managers.forward_context import ( get_forward_context, ) ctx = get_forward_context() out = self.attn_impl.forward(q, k, v, attn_metadata=ctx.attn_metadata) out = out.reshape(B, N_q, self.inner_dim) out = self.to_out[0](out) out = self.to_out[1](out) return out class Basic2p5DTransformerBlock(torch.nn.Module): """2.5D Transformer block with Multiview Attention (MVA) and Reference View Attention (RVA).""" def __init__( self, transformer: BasicTransformerBlock, layer_name: str, use_ma: bool = True, use_ra: bool = True, is_turbo: bool = False, use_sglang_attn: bool = True, ) -> None: super().__init__() self.transformer = transformer self.layer_name = layer_name self.use_ma = use_ma self.use_ra = use_ra self.is_turbo = is_turbo self.use_sglang_attn = use_sglang_attn and not is_turbo attn_cls = ( SGLangAttentionWrapper if self.use_sglang_attn else DiffusersAttention ) attn_kwargs = dict( query_dim=self.dim, heads=self.num_attention_heads, dim_head=self.attention_head_dim, dropout=self.dropout, bias=self.attention_bias, cross_attention_dim=None, upcast_attention=self.attn1.upcast_attention, out_bias=True, ) if self.use_sglang_attn: attn_kwargs.pop("upcast_attention") if self.use_ma: self.attn_multiview = attn_cls(**attn_kwargs) if self.use_ra: self.attn_refview = attn_cls(**attn_kwargs) if self.is_turbo: self._initialize_attn_weights() def _initialize_attn_weights(self): """Initialize attention weights for turbo mode.""" if self.use_ma: self.attn_multiview.load_state_dict(self.attn1.state_dict()) with torch.no_grad(): for layer in self.attn_multiview.to_out: for param in layer.parameters(): param.zero_() if self.use_ra: self.attn_refview.load_state_dict(self.attn1.state_dict()) with torch.no_grad(): for layer in self.attn_refview.to_out: for param in layer.parameters(): param.zero_() def __getattr__(self, name: str): try: return super().__getattr__(name) except AttributeError: return getattr(self.transformer, name) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, timestep: Optional[torch.LongTensor] = None, cross_attention_kwargs: dict = None, class_labels: Optional[torch.LongTensor] = None, added_cond_kwargs: Optional[dict] = None, ) -> torch.Tensor: """Forward pass with MVA and RVA support.""" batch_size = hidden_states.shape[0] cross_attention_kwargs = ( cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} ) num_in_batch = cross_attention_kwargs.pop("num_in_batch", 1) mode = cross_attention_kwargs.pop("mode", None) if not self.is_turbo: mva_scale = cross_attention_kwargs.pop("mva_scale", 1.0) ref_scale = cross_attention_kwargs.pop("ref_scale", 1.0) else: position_attn_mask = cross_attention_kwargs.pop("position_attn_mask", None) position_voxel_indices = cross_attention_kwargs.pop( "position_voxel_indices", None ) mva_scale = 1.0 ref_scale = 1.0 condition_embed_dict = cross_attention_kwargs.pop("condition_embed_dict", None) # Normalization if self.norm_type == "ada_norm": norm_hidden_states = self.norm1(hidden_states, timestep) elif self.norm_type == "ada_norm_zero": norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype ) elif self.norm_type in ["layer_norm", "layer_norm_i2vgen"]: norm_hidden_states = self.norm1(hidden_states) elif self.norm_type == "ada_norm_continuous": norm_hidden_states = self.norm1( hidden_states, added_cond_kwargs["pooled_text_emb"] ) elif self.norm_type == "ada_norm_single": shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) ).chunk(6, dim=1) norm_hidden_states = self.norm1(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa else: raise ValueError("Incorrect norm used") if self.pos_embed is not None: norm_hidden_states = self.pos_embed(norm_hidden_states) # Prepare GLIGEN inputs cross_attention_kwargs = ( cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} ) gligen_kwargs = cross_attention_kwargs.pop("gligen", None) # Self-attention attn_output = self.attn1( norm_hidden_states, encoder_hidden_states=( encoder_hidden_states if self.only_cross_attention else None ), attention_mask=attention_mask, **cross_attention_kwargs, ) if self.norm_type == "ada_norm_zero": attn_output = gate_msa.unsqueeze(1) * attn_output elif self.norm_type == "ada_norm_single": attn_output = gate_msa * attn_output hidden_states = attn_output + hidden_states if hidden_states.ndim == 4: hidden_states = hidden_states.squeeze(1) # Reference Attention - Write mode if mode is not None and "w" in mode: condition_embed_dict[self.layer_name] = rearrange( norm_hidden_states, "(b n) l c -> b (n l) c", n=num_in_batch ) # Reference Attention - Read mode if mode is not None and "r" in mode and self.use_ra: condition_embed = ( condition_embed_dict[self.layer_name] .unsqueeze(1) .repeat(1, num_in_batch, 1, 1) ) condition_embed = rearrange(condition_embed, "b n l c -> (b n) l c") attn_output = self.attn_refview( norm_hidden_states, encoder_hidden_states=condition_embed, attention_mask=None, **cross_attention_kwargs, ) if not self.is_turbo: ref_scale_timing = ref_scale if isinstance(ref_scale, torch.Tensor): ref_scale_timing = ( ref_scale.unsqueeze(1).repeat(1, num_in_batch).view(-1) ) for _ in range(attn_output.ndim - 1): ref_scale_timing = ref_scale_timing.unsqueeze(-1) hidden_states = ref_scale_timing * attn_output + hidden_states if hidden_states.ndim == 4: hidden_states = hidden_states.squeeze(1) # Multiview Attention if num_in_batch > 1 and self.use_ma: multivew_hidden_states = rearrange( norm_hidden_states, "(b n) l c -> b (n l) c", n=num_in_batch ) if self.is_turbo: position_mask = None if position_attn_mask is not None: if multivew_hidden_states.shape[1] in position_attn_mask: position_mask = position_attn_mask[ multivew_hidden_states.shape[1] ] position_indices = None if position_voxel_indices is not None: if multivew_hidden_states.shape[1] in position_voxel_indices: position_indices = position_voxel_indices[ multivew_hidden_states.shape[1] ] attn_output = self.attn_multiview( multivew_hidden_states, encoder_hidden_states=multivew_hidden_states, attention_mask=position_mask, position_indices=position_indices, **cross_attention_kwargs, ) else: attn_output = self.attn_multiview( multivew_hidden_states, encoder_hidden_states=multivew_hidden_states, **cross_attention_kwargs, ) attn_output = rearrange( attn_output, "b (n l) c -> (b n) l c", n=num_in_batch ) hidden_states = mva_scale * attn_output + hidden_states if hidden_states.ndim == 4: hidden_states = hidden_states.squeeze(1) # GLIGEN Control if gligen_kwargs is not None: hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) # Cross-Attention if self.attn2 is not None: if self.norm_type == "ada_norm": norm_hidden_states = self.norm2(hidden_states, timestep) elif self.norm_type in ["ada_norm_zero", "layer_norm", "layer_norm_i2vgen"]: norm_hidden_states = self.norm2(hidden_states) elif self.norm_type == "ada_norm_single": norm_hidden_states = hidden_states elif self.norm_type == "ada_norm_continuous": norm_hidden_states = self.norm2( hidden_states, added_cond_kwargs["pooled_text_emb"] ) else: raise ValueError("Incorrect norm") if self.pos_embed is not None and self.norm_type != "ada_norm_single": norm_hidden_states = self.pos_embed(norm_hidden_states) attn_output = self.attn2( norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=encoder_attention_mask, **cross_attention_kwargs, ) hidden_states = attn_output + hidden_states # Feed-forward if self.norm_type == "ada_norm_continuous": norm_hidden_states = self.norm3( hidden_states, added_cond_kwargs["pooled_text_emb"] ) elif not self.norm_type == "ada_norm_single": norm_hidden_states = self.norm3(hidden_states) if self.norm_type == "ada_norm_zero": norm_hidden_states = ( norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] ) if self.norm_type == "ada_norm_single": norm_hidden_states = self.norm2(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp if self._chunk_size is not None: ff_output = _chunked_feed_forward( self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size ) else: ff_output = self.ff(norm_hidden_states) if self.norm_type == "ada_norm_zero": ff_output = gate_mlp.unsqueeze(1) * ff_output elif self.norm_type == "ada_norm_single": ff_output = gate_mlp * ff_output hidden_states = ff_output + hidden_states if hidden_states.ndim == 4: hidden_states = hidden_states.squeeze(1) return hidden_states @torch.no_grad() def compute_voxel_grid_mask(position: torch.Tensor, grid_resolution: int = 8): """Compute voxel grid mask for position-aware attention.""" position = position.half() B, N, _, H, W = position.shape assert H % grid_resolution == 0 and W % grid_resolution == 0 valid_mask = (position != 1).all(dim=2, keepdim=True) valid_mask = valid_mask.expand_as(position) position[valid_mask == False] = 0 position = rearrange( position, "b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w", num_h=grid_resolution, num_w=grid_resolution, ) valid_mask = rearrange( valid_mask, "b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w", num_h=grid_resolution, num_w=grid_resolution, ) grid_position = position.sum(dim=(-2, -1)) count_masked = valid_mask.sum(dim=(-2, -1)) grid_position = grid_position / count_masked.clamp(min=1) grid_position[count_masked < 5] = 0 grid_position = grid_position.permute(0, 1, 4, 2, 3) grid_position = rearrange(grid_position, "b n c h w -> b n (h w) c") grid_position_expanded_1 = grid_position.unsqueeze(2).unsqueeze(4) grid_position_expanded_2 = grid_position.unsqueeze(1).unsqueeze(3) distances = torch.norm(grid_position_expanded_1 - grid_position_expanded_2, dim=-1) weights = distances grid_distance = 1.73 / grid_resolution weights = weights < grid_distance return weights def compute_multi_resolution_mask( position_maps: torch.Tensor, grid_resolutions: List[int] = [32, 16, 8] ) -> dict: """Compute multi-resolution position attention masks.""" position_attn_mask = {} with torch.no_grad(): for grid_resolution in grid_resolutions: position_mask = compute_voxel_grid_mask(position_maps, grid_resolution) position_mask = rearrange( position_mask, "b ni nj li lj -> b (ni li) (nj lj)" ) position_attn_mask[position_mask.shape[1]] = position_mask return position_attn_mask @torch.no_grad() def compute_discrete_voxel_indice( position: torch.Tensor, grid_resolution: int = 8, voxel_resolution: int = 128 ): """Compute discrete voxel indices for position encoding.""" position = position.half() B, N, _, H, W = position.shape assert H % grid_resolution == 0 and W % grid_resolution == 0 valid_mask = (position != 1).all(dim=2, keepdim=True) valid_mask = valid_mask.expand_as(position) position[valid_mask == False] = 0 position = rearrange( position, "b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w", num_h=grid_resolution, num_w=grid_resolution, ) valid_mask = rearrange( valid_mask, "b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w", num_h=grid_resolution, num_w=grid_resolution, ) grid_position = position.sum(dim=(-2, -1)) count_masked = valid_mask.sum(dim=(-2, -1)) grid_position = grid_position / count_masked.clamp(min=1) grid_position[count_masked < 5] = 0 grid_position = grid_position.permute(0, 1, 4, 2, 3).clamp(0, 1) voxel_indices = grid_position * (voxel_resolution - 1) voxel_indices = torch.round(voxel_indices).long() return voxel_indices def compute_multi_resolution_discrete_voxel_indice( position_maps: torch.Tensor, grid_resolutions: List[int] = [64, 32, 16, 8], voxel_resolutions: List[int] = [512, 256, 128, 64], ) -> dict: """Compute multi-resolution discrete voxel indices.""" voxel_indices = {} with torch.no_grad(): for grid_resolution, voxel_resolution in zip( grid_resolutions, voxel_resolutions ): voxel_indice = compute_discrete_voxel_indice( position_maps, grid_resolution, voxel_resolution ) voxel_indice = rearrange(voxel_indice, "b n c h w -> b (n h w) c") voxel_indices[voxel_indice.shape[1]] = { "voxel_indices": voxel_indice, "voxel_resolution": voxel_resolution, } return voxel_indices class UNet2p5DConditionModel(torch.nn.Module): """2.5D UNet for multi-view texture generation.""" def __init__(self, unet: UNet2DConditionModel) -> None: super().__init__() self.unet = unet self.use_ma = True self.use_ra = True self.use_camera_embedding = True self.use_dual_stream = True self.is_turbo = False if self.use_dual_stream: self.unet_dual = copy.deepcopy(unet) self.init_attention(self.unet_dual) self.init_attention( self.unet, use_ma=self.use_ma, use_ra=self.use_ra, is_turbo=self.is_turbo ) self.init_condition() self.init_camera_embedding() @staticmethod def from_pretrained(pretrained_model_name_or_path: str, **kwargs): """Load a pretrained UNet2p5DConditionModel.""" torch_dtype = kwargs.pop("dtype", kwargs.pop("torch_dtype", torch.float32)) config_path = _os.path.join(pretrained_model_name_or_path, "config.json") unet_ckpt_path = _os.path.join( pretrained_model_name_or_path, "diffusion_pytorch_model.bin" ) with open(config_path, "r", encoding="utf-8") as file: config = json.load(file) unet = UNet2DConditionModel(**config) unet = UNet2p5DConditionModel(unet) unet_ckpt = torch.load(unet_ckpt_path, map_location="cpu", weights_only=True) unet.load_state_dict(unet_ckpt, strict=True) unet = unet.to(torch_dtype) return unet def init_condition(self): """Initialize condition-related modules.""" self.unet.conv_in = torch.nn.Conv2d( 12, # 4 (latent) + 4 (normal) + 4 (position) self.unet.conv_in.out_channels, kernel_size=self.unet.conv_in.kernel_size, stride=self.unet.conv_in.stride, padding=self.unet.conv_in.padding, dilation=self.unet.conv_in.dilation, groups=self.unet.conv_in.groups, bias=self.unet.conv_in.bias is not None, ) self.unet.learned_text_clip_gen = nn.Parameter(torch.randn(1, 77, 1024)) self.unet.learned_text_clip_ref = nn.Parameter(torch.randn(1, 77, 1024)) def init_camera_embedding(self): """Initialize camera embedding module.""" if self.use_camera_embedding: time_embed_dim = 1280 self.max_num_ref_image = 5 self.max_num_gen_image = 12 * 3 + 4 * 2 self.unet.class_embedding = nn.Embedding( self.max_num_ref_image + self.max_num_gen_image, time_embed_dim ) def init_attention( self, unet: UNet2DConditionModel, use_ma: bool = False, use_ra: bool = False, is_turbo: bool = False, use_sglang_attn: bool = True, ): """Initialize attention blocks with MVA and RVA support.""" block_kwargs = dict( use_ma=use_ma, use_ra=use_ra, is_turbo=is_turbo, use_sglang_attn=use_sglang_attn, ) # Down blocks for down_block_i, down_block in enumerate(unet.down_blocks): if ( hasattr(down_block, "has_cross_attention") and down_block.has_cross_attention ): for attn_i, attn in enumerate(down_block.attentions): for transformer_i, transformer in enumerate( attn.transformer_blocks ): if isinstance(transformer, BasicTransformerBlock): attn.transformer_blocks[transformer_i] = ( Basic2p5DTransformerBlock( transformer, f"down_{down_block_i}_{attn_i}_{transformer_i}", **block_kwargs, ) ) # Mid block if ( hasattr(unet.mid_block, "has_cross_attention") and unet.mid_block.has_cross_attention ): for attn_i, attn in enumerate(unet.mid_block.attentions): for transformer_i, transformer in enumerate(attn.transformer_blocks): if isinstance(transformer, BasicTransformerBlock): attn.transformer_blocks[transformer_i] = ( Basic2p5DTransformerBlock( transformer, f"mid_{attn_i}_{transformer_i}", **block_kwargs, ) ) # Up blocks for up_block_i, up_block in enumerate(unet.up_blocks): if ( hasattr(up_block, "has_cross_attention") and up_block.has_cross_attention ): for attn_i, attn in enumerate(up_block.attentions): for transformer_i, transformer in enumerate( attn.transformer_blocks ): if isinstance(transformer, BasicTransformerBlock): attn.transformer_blocks[transformer_i] = ( Basic2p5DTransformerBlock( transformer, f"up_{up_block_i}_{attn_i}_{transformer_i}", **block_kwargs, ) ) if use_sglang_attn and (use_ma or use_ra): backend = "unknown" for block in self._iter_2p5d_blocks(unet): for attr in ("attn_multiview", "attn_refview"): wrapper = getattr(block, attr, None) if isinstance(wrapper, SGLangAttentionWrapper): backend = wrapper._attn_backend_name break if backend != "unknown": break count = sum(1 for _ in self._iter_2p5d_blocks(unet)) logger.info( "Initialized %d Basic2p5DTransformerBlocks with sglang %s attention", count, backend, ) @staticmethod def _iter_2p5d_blocks(unet): """Yield all Basic2p5DTransformerBlock instances in a UNet.""" for block_group in (unet.down_blocks, [unet.mid_block], unet.up_blocks): for block in block_group: if not hasattr(block, "attentions"): continue for attn in block.attentions: for tb in attn.transformer_blocks: if isinstance(tb, Basic2p5DTransformerBlock): yield tb def __getattr__(self, name: str): try: return super().__getattr__(name) except AttributeError: return getattr(self.unet, name) def forward( self, sample: torch.Tensor, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, *args, down_intrablock_additional_residuals=None, down_block_res_samples=None, mid_block_res_sample=None, **cached_condition, ): """Forward pass for multi-view texture generation.""" B, N_gen, _, H, W = sample.shape assert H == W if self.use_camera_embedding: camera_info_gen = ( cached_condition["camera_info_gen"] + self.max_num_ref_image ) camera_info_gen = rearrange(camera_info_gen, "b n -> (b n)") else: camera_info_gen = None # Concatenate latents with normal and position maps sample = [sample] if "normal_imgs" in cached_condition: sample.append(cached_condition["normal_imgs"]) if "position_imgs" in cached_condition: sample.append(cached_condition["position_imgs"]) sample = torch.cat(sample, dim=2) sample = rearrange(sample, "b n c h w -> (b n) c h w") encoder_hidden_states_gen = encoder_hidden_states.unsqueeze(1).repeat( 1, N_gen, 1, 1 ) encoder_hidden_states_gen = rearrange( encoder_hidden_states_gen, "b n l c -> (b n) l c" ) # Process reference images for RVA if self.use_ra: if "condition_embed_dict" in cached_condition: condition_embed_dict = cached_condition["condition_embed_dict"] else: condition_embed_dict = {} ref_latents = cached_condition["ref_latents"] N_ref = ref_latents.shape[1] if self.use_camera_embedding: camera_info_ref = cached_condition["camera_info_ref"] camera_info_ref = rearrange(camera_info_ref, "b n -> (b n)") else: camera_info_ref = None ref_latents = rearrange(ref_latents, "b n c h w -> (b n) c h w") encoder_hidden_states_ref = self.unet.learned_text_clip_ref.unsqueeze( 1 ).repeat(B, N_ref, 1, 1) encoder_hidden_states_ref = rearrange( encoder_hidden_states_ref, "b n l c -> (b n) l c" ) noisy_ref_latents = ref_latents timestep_ref = 0 if self.use_dual_stream: unet_ref = self.unet_dual else: unet_ref = self.unet unet_ref( noisy_ref_latents, timestep_ref, encoder_hidden_states=encoder_hidden_states_ref, class_labels=camera_info_ref, return_dict=False, cross_attention_kwargs={ "mode": "w", "num_in_batch": N_ref, "condition_embed_dict": condition_embed_dict, }, ) cached_condition["condition_embed_dict"] = condition_embed_dict else: condition_embed_dict = None mva_scale = cached_condition.get("mva_scale", 1.0) ref_scale = cached_condition.get("ref_scale", 1.0) if self.is_turbo: position_attn_mask = cached_condition.get("position_attn_mask", None) position_voxel_indices = cached_condition.get( "position_voxel_indices", None ) cross_attention_kwargs_ = { "mode": "r", "num_in_batch": N_gen, "condition_embed_dict": condition_embed_dict, "position_attn_mask": position_attn_mask, "position_voxel_indices": position_voxel_indices, "mva_scale": mva_scale, "ref_scale": ref_scale, } else: cross_attention_kwargs_ = { "mode": "r", "num_in_batch": N_gen, "condition_embed_dict": condition_embed_dict, "mva_scale": mva_scale, "ref_scale": ref_scale, } return self.unet( sample, timestep, encoder_hidden_states_gen, *args, class_labels=camera_info_gen, down_intrablock_additional_residuals=( [ s.to(dtype=self.unet.dtype) for s in down_intrablock_additional_residuals ] if down_intrablock_additional_residuals is not None else None ), down_block_additional_residuals=( [s.to(dtype=self.unet.dtype) for s in down_block_res_samples] if down_block_res_samples is not None else None ), mid_block_additional_residual=( mid_block_res_sample.to(dtype=self.unet.dtype) if mid_block_res_sample is not None else None ), return_dict=False, cross_attention_kwargs=cross_attention_kwargs_, ) # Entry class for model registry EntryClass = [Hunyuan3D2DiT, UNet2p5DConditionModel] ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from typing import Any import numpy as np import torch import torch.nn as nn from sglang.multimodal_gen.configs.models.dits import HunyuanVideoConfig from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size from sglang.multimodal_gen.runtime.layers.attention import ( LocalAttention, UlyssesAttention, ) from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd from sglang.multimodal_gen.runtime.layers.layernorm import ( LayerNormScaleShift, RMSNorm, ScaleResidualLayerNormScaleShift, ) from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( _apply_rotary_emb, get_rotary_pos_embed, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import ( ModulateProjection, PatchEmbed, TimestepEmbedder, unpatchify, ) from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.models.utils import modulate from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin class MMDoubleStreamBlock(nn.Module): """ A multimodal DiT block with separate modulation for text and image/video, using distributed attention and linear layers. """ def __init__( self, hidden_size: int, num_attention_heads: int, mlp_ratio: float, dtype: torch.dtype | None = None, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", quant_config: QuantizationConfig | None = None, ): super().__init__() self.deterministic = False self.num_attention_heads = num_attention_heads head_dim = hidden_size // num_attention_heads mlp_hidden_dim = int(hidden_size * mlp_ratio) # Image modulation components self.img_mod = ModulateProjection( hidden_size, factor=6, act_layer="silu", dtype=dtype, prefix=f"{prefix}.img_mod", ) # Fused operations for image stream self.img_attn_norm = LayerNormScaleShift( hidden_size, elementwise_affine=False, dtype=dtype ) self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift( hidden_size, elementwise_affine=False, dtype=dtype ) self.img_mlp_residual = MulAdd() # Image attention components self.img_attn_qkv = ReplicatedLinear( hidden_size, hidden_size * 3, bias=True, params_dtype=dtype, prefix=f"{prefix}.img_attn_qkv", quant_config=quant_config, ) self.img_attn_q_norm = RMSNorm(head_dim, eps=1e-6, dtype=dtype) self.img_attn_k_norm = RMSNorm(head_dim, eps=1e-6, dtype=dtype) self.img_attn_proj = ReplicatedLinear( hidden_size, hidden_size, bias=True, params_dtype=dtype, prefix=f"{prefix}.img_attn_proj", quant_config=quant_config, ) self.img_mlp = MLP( hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, prefix=f"{prefix}.img_mlp", quant_config=quant_config, ) # Text modulation components self.txt_mod = ModulateProjection( hidden_size, factor=6, act_layer="silu", dtype=dtype, prefix=f"{prefix}.txt_mod", ) # Fused operations for text stream self.txt_attn_norm = LayerNormScaleShift( hidden_size, elementwise_affine=False, dtype=dtype ) self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift( hidden_size, elementwise_affine=False, dtype=dtype ) self.txt_mlp_residual = MulAdd() # Text attention components self.txt_attn_qkv = ReplicatedLinear( hidden_size, hidden_size * 3, bias=True, params_dtype=dtype, quant_config=quant_config, ) # QK norm layers for text self.txt_attn_q_norm = RMSNorm(head_dim, eps=1e-6, dtype=dtype) self.txt_attn_k_norm = RMSNorm(head_dim, eps=1e-6, dtype=dtype) self.txt_attn_proj = ReplicatedLinear( hidden_size, hidden_size, bias=True, params_dtype=dtype, quant_config=quant_config, ) self.txt_mlp = MLP( hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, quant_config=quant_config, ) # Use UlyssesAttention to replace Distributed attention self.attn = UlyssesAttention( num_heads=num_attention_heads, head_size=head_dim, causal=False, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn", ) def forward( self, img: torch.Tensor, txt: torch.Tensor, vec: torch.Tensor, freqs_cis: tuple, ) -> tuple[torch.Tensor, torch.Tensor]: # Process modulation vectors img_mod_outputs = self.img_mod(vec) ( img_attn_shift, img_attn_scale, img_attn_gate, img_mlp_shift, img_mlp_scale, img_mlp_gate, ) = torch.chunk(img_mod_outputs, 6, dim=-1) txt_mod_outputs = self.txt_mod(vec) ( txt_attn_shift, txt_attn_scale, txt_attn_gate, txt_mlp_shift, txt_mlp_scale, txt_mlp_gate, ) = torch.chunk(txt_mod_outputs, 6, dim=-1) # Prepare image for attention using fused operation img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale) # Get QKV for image img_qkv, _ = self.img_attn_qkv(img_attn_input) batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1] # Split QKV img_qkv = img_qkv.view( batch_size, image_seq_len, 3, self.num_attention_heads, -1 ) img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2] # Apply QK-Norm if needed img_q = self.img_attn_q_norm(img_q.contiguous()).to(img_v) img_k = self.img_attn_k_norm(img_k.contiguous()).to(img_v) # Apply rotary embeddings cos, sin = freqs_cis img_q, img_k = _apply_rotary_emb( img_q, cos, sin, is_neox_style=False ), _apply_rotary_emb(img_k, cos, sin, is_neox_style=False) # Prepare text for attention using fused operation txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale) # Get QKV for text txt_qkv, _ = self.txt_attn_qkv(txt_attn_input) batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1] # Split QKV txt_qkv = txt_qkv.view( batch_size, text_seq_len, 3, self.num_attention_heads, -1 ) txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2] # Apply QK-Norm if needed txt_q = self.txt_attn_q_norm(txt_q.contiguous()).to(txt_q.dtype) txt_k = self.txt_attn_k_norm(txt_k.contiguous()).to(txt_k.dtype) # Run distributed attention img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v) img_attn_out, _ = self.img_attn_proj( img_attn.view(batch_size, image_seq_len, -1) ) # Use fused operation for residual connection, normalization, and modulation img_mlp_input, img_residual = self.img_attn_residual_mlp_norm( img, img_attn_out, img_attn_gate, img_mlp_shift, img_mlp_scale ) # Process image MLP img_mlp_out = self.img_mlp(img_mlp_input) img = self.img_mlp_residual(img_mlp_out, img_mlp_gate, img_residual) # Process text attention output txt_attn_out, _ = self.txt_attn_proj( txt_attn.reshape(batch_size, text_seq_len, -1) ) # Use fused operation for residual connection, normalization, and modulation txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm( txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale ) # Process text MLP txt_mlp_out = self.txt_mlp(txt_mlp_input) txt = self.txt_mlp_residual(txt_mlp_out, txt_mlp_gate, txt_residual) return img, txt class MMSingleStreamBlock(nn.Module): """ A DiT block with parallel linear layers using distributed attention and tensor parallelism. """ def __init__( self, hidden_size: int, num_attention_heads: int, mlp_ratio: float = 4.0, dtype: torch.dtype | None = None, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", quant_config: QuantizationConfig | None = None, ): super().__init__() self.deterministic = False self.hidden_size = hidden_size self.num_attention_heads = num_attention_heads head_dim = hidden_size // num_attention_heads mlp_hidden_dim = int(hidden_size * mlp_ratio) self.mlp_hidden_dim = mlp_hidden_dim # Combined QKV and MLP input projection self.linear1 = ReplicatedLinear( hidden_size, hidden_size * 3 + mlp_hidden_dim, bias=True, params_dtype=dtype, prefix=f"{prefix}.linear1", quant_config=quant_config, ) # Combined projection and MLP output self.linear2 = ReplicatedLinear( hidden_size + mlp_hidden_dim, hidden_size, bias=True, params_dtype=dtype, prefix=f"{prefix}.linear2", quant_config=quant_config, ) # QK norm layers self.q_norm = RMSNorm(head_dim, eps=1e-6, dtype=dtype) self.k_norm = RMSNorm(head_dim, eps=1e-6, dtype=dtype) # Fused operations with better naming self.input_norm_scale_shift = LayerNormScaleShift( hidden_size, eps=1e-6, elementwise_affine=False, dtype=dtype, ) self.output_residual = MulAdd() # Activation function self.mlp_act = nn.GELU(approximate="tanh") # Modulation self.modulation = ModulateProjection( hidden_size, factor=3, act_layer="silu", dtype=dtype, prefix=f"{prefix}.modulation", ) # Use UlyssesAttention to replace Distributed attention self.attn = UlyssesAttention( num_heads=num_attention_heads, head_size=head_dim, causal=False, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn", ) def forward( self, x: torch.Tensor, vec: torch.Tensor, txt_len: int, freqs_cis: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: # Process modulation mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1) # Apply pre-norm and modulation using fused operation x_mod = self.input_norm_scale_shift(x, mod_shift, mod_scale) # Get combined projections linear1_out, _ = self.linear1(x_mod) # Split into QKV and MLP parts qkv, mlp = torch.split( linear1_out, [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1 ) # Process QKV batch_size, seq_len = qkv.shape[0], qkv.shape[1] qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1) q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] # Apply QK-Norm q = self.q_norm(q.contiguous()).to(v.dtype) k = self.k_norm(k.contiguous()).to(v.dtype) # Split into image and text parts img_q, txt_q = q[:, :-txt_len], q[:, -txt_len:] img_k, txt_k = k[:, :-txt_len], k[:, -txt_len:] img_v, txt_v = v[:, :-txt_len], v[:, -txt_len:] # Apply rotary embeddings to image parts cos, sin = freqs_cis img_q, img_k = _apply_rotary_emb( img_q, cos, sin, is_neox_style=False ), _apply_rotary_emb(img_k, cos, sin, is_neox_style=False) # Run distributed attention img_attn_output, txt_attn_output = self.attn( img_q, img_k, img_v, txt_q, txt_k, txt_v ) attn_output = torch.cat((img_attn_output, txt_attn_output), dim=1).view( batch_size, seq_len, -1 ) # Process MLP activation mlp_output = self.mlp_act(mlp) # Combine attention and MLP outputs combined = torch.cat((attn_output, mlp_output), dim=-1) # Final projection output, _ = self.linear2(combined) # Apply residual connection with gating using fused operation return self.output_residual(output, mod_gate, x) class HunyuanVideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): """ HunyuanVideo Transformer backbone adapted for distributed training. This implementation uses distributed attention and linear layers for efficient parallel processing across multiple GPUs. Based on the architecture from: - Flux.1: https://github.com/black-forest-labs/flux - MMDiT: http://arxiv.org/abs/2403.03206 """ # PY: we make the input args the same as HF config # shard single stream, double stream blocks, and refiner_blocks _fsdp_shard_conditions = HunyuanVideoConfig()._fsdp_shard_conditions _compile_conditions = HunyuanVideoConfig()._compile_conditions _supported_attention_backends = HunyuanVideoConfig()._supported_attention_backends param_names_mapping = HunyuanVideoConfig().param_names_mapping reverse_param_names_mapping = HunyuanVideoConfig().reverse_param_names_mapping lora_param_names_mapping = HunyuanVideoConfig().lora_param_names_mapping def __init__( self, config: HunyuanVideoConfig, hf_config: dict[str, Any], quant_config: QuantizationConfig | None = None, ): super().__init__(config=config, hf_config=hf_config) self.patch_size = [config.patch_size_t, config.patch_size, config.patch_size] self.in_channels = config.in_channels self.num_channels_latents = config.num_channels_latents self.out_channels = ( config.in_channels if config.out_channels is None else config.out_channels ) self.unpatchify_channels = self.out_channels self.guidance_embeds = config.guidance_embeds self.rope_dim_list = list(config.rope_axes_dim) self.rope_theta = config.rope_theta self.text_states_dim = config.text_embed_dim self.text_states_dim_2 = config.pooled_projection_dim # TODO(will): hack? self.dtype = config.dtype pe_dim = config.hidden_size // config.num_attention_heads if sum(config.rope_axes_dim) != pe_dim: raise ValueError( f"Got {config.rope_axes_dim} but expected positional dim {pe_dim}" ) self.hidden_size = config.hidden_size self.num_attention_heads = config.num_attention_heads self.num_channels_latents = config.num_channels_latents # Image projection self.img_in = PatchEmbed( self.patch_size, self.in_channels, self.hidden_size, dtype=config.dtype, prefix=f"{config.prefix}.img_in", ) self.txt_in = SingleTokenRefiner( self.text_states_dim, config.hidden_size, config.num_attention_heads, depth=config.num_refiner_layers, dtype=config.dtype, prefix=f"{config.prefix}.txt_in", ) # Time modulation self.time_in = TimestepEmbedder( self.hidden_size, act_layer="silu", dtype=config.dtype, prefix=f"{config.prefix}.time_in", ) # Text modulation self.vector_in = MLP( self.text_states_dim_2, self.hidden_size, self.hidden_size, act_type="silu", dtype=config.dtype, prefix=f"{config.prefix}.vector_in", ) # Guidance modulation self.guidance_in = ( TimestepEmbedder( self.hidden_size, act_layer="silu", dtype=config.dtype, prefix=f"{config.prefix}.guidance_in", ) if self.guidance_embeds else None ) # Double blocks self.double_blocks = nn.ModuleList( [ MMDoubleStreamBlock( config.hidden_size, config.num_attention_heads, mlp_ratio=config.mlp_ratio, dtype=config.dtype, supported_attention_backends=self._supported_attention_backends, prefix=f"{config.prefix}.double_blocks.{i}", quant_config=quant_config, ) for i in range(config.num_layers) ] ) # Single blocks self.single_blocks = nn.ModuleList( [ MMSingleStreamBlock( config.hidden_size, config.num_attention_heads, mlp_ratio=config.mlp_ratio, dtype=config.dtype, supported_attention_backends=self._supported_attention_backends, prefix=f"{config.prefix}.single_blocks.{i + config.num_layers}", quant_config=quant_config, ) for i in range(config.num_single_layers) ] ) self.final_layer = FinalLayer( config.hidden_size, self.patch_size, self.out_channels, dtype=config.dtype, prefix=f"{config.prefix}.final_layer", ) self.__post_init__() self.layer_names = ["double_blocks", "single_blocks"] # TODO: change the input the FORWARD_BATCH Dict # TODO: change output to a dict def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, guidance=None, **kwargs, ): """ Forward pass of the HunyuanDiT model. Args: hidden_states: Input image/video latents [B, C, T, H, W] encoder_hidden_states: Text embeddings [B, L, D] timestep: Diffusion timestep guidance: Guidance scale for CFG Returns: Tuple of (output) """ forward_context = get_forward_context() forward_batch = forward_context.forward_batch enable_teacache = forward_batch is not None and forward_batch.enable_teacache if guidance is None: guidance = torch.tensor( [6016.0], device=hidden_states.device, dtype=hidden_states.dtype ) img = x = hidden_states t = timestep # Split text embeddings - first token is global, rest are per-token if isinstance(encoder_hidden_states, torch.Tensor): txt = encoder_hidden_states[:, 1:] text_states_2 = encoder_hidden_states[:, 0, : self.text_states_dim_2] else: txt = encoder_hidden_states[0] text_states_2 = encoder_hidden_states[1] # Get spatial dimensions _, _, ot, oh, ow = x.shape # codespell:ignore tt, th, tw = ( ot // self.patch_size[0], # codespell:ignore oh // self.patch_size[1], ow // self.patch_size[2], ) # Get rotary embeddings freqs_cos, freqs_sin = get_rotary_pos_embed( (tt * get_sp_world_size(), th, tw), self.hidden_size, self.num_attention_heads, self.rope_dim_list, self.rope_theta, ) freqs_cos = freqs_cos.to(x.device) freqs_sin = freqs_sin.to(x.device) # Prepare modulation vectors vec = self.time_in(t) # Add text modulation vec = vec + self.vector_in(text_states_2) # Add guidance modulation if needed if self.guidance_in and guidance is not None: vec = vec + self.guidance_in(guidance) # Embed image and text img = self.img_in(img) txt = self.txt_in(txt, t) txt_seq_len = txt.shape[1] img_seq_len = img.shape[1] freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None should_skip_forward = self.should_skip_forward_for_cached_states( img=img, vec=vec ) if should_skip_forward: img = self.retrieve_cached_states(img) else: if enable_teacache: original_img = img.clone() # Process through double stream blocks for index, block in enumerate(self.double_blocks): double_block_args = [img, txt, vec, freqs_cis] img, txt = block(*double_block_args) # Merge txt and img to pass through single stream blocks x = torch.cat((img, txt), 1) # Process through single stream blocks if len(self.single_blocks) > 0: for index, block in enumerate(self.single_blocks): single_block_args = [ x, vec, txt_seq_len, freqs_cis, ] x = block(*single_block_args) # Extract image features img = x[:, :img_seq_len, ...] if enable_teacache: self.maybe_cache_states(img, original_img) # Final layer processing img = self.final_layer(img, vec) # Unpatchify to get original shape img = unpatchify(img, tt, th, tw, self.patch_size, self.out_channels) return img def maybe_cache_states( self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor ) -> None: self.previous_residual = hidden_states - original_hidden_states def should_skip_forward_for_cached_states(self, **kwargs) -> bool: forward_context = get_forward_context() forward_batch = forward_context.forward_batch if forward_batch is None: return False current_timestep = forward_context.current_timestep enable_teacache = forward_batch.enable_teacache if not enable_teacache: return False raise NotImplementedError("teacache is not supported yet for HunyuanVideo") teacache_params = forward_batch.teacache_params assert teacache_params is not None, "teacache_params is not initialized" assert isinstance( teacache_params, TeaCacheParams ), "teacache_params is not a TeaCacheParams" num_inference_steps = forward_batch.num_inference_steps teache_thresh = teacache_params.teacache_thresh coefficients = teacache_params.coefficients if current_timestep == 0: self.cnt = 0 inp = kwargs["img"].clone() vec_ = kwargs["vec"].clone() # convert to DTensor vec_ = torch.distributed.tensor.DTensor.from_local( vec_, torch.distributed.DeviceMesh( current_platform.device_type, list(range(get_sp_world_size())), mesh_dim_names=("dp",), ), [torch.distributed.tensor.Replicate()], ) inp = torch.distributed.tensor.DTensor.from_local( inp, torch.distributed.DeviceMesh( current_platform.device_type, list(range(get_sp_world_size())), mesh_dim_names=("dp",), ), [torch.distributed.tensor.Replicate()], ) # txt_ = kwargs["txt"].clone() # inp = img.clone() # vec_ = vec.clone() # txt_ = txt.clone() ( img_mod1_shift, img_mod1_scale, img_mod1_gate, img_mod2_shift, img_mod2_scale, img_mod2_gate, ) = ( self.double_blocks[0].img_mod(vec_).chunk(6, dim=-1) ) normed_inp = self.double_blocks[0].img_attn_norm.norm(inp) modulated_inp = modulate(normed_inp, shift=img_mod1_shift, scale=img_mod1_scale) if self.cnt == 0 or self.cnt == num_inference_steps - 1: should_calc = True self.accumulated_rel_l1_distance = 0 else: coefficients = [ 7.33226126e02, -4.01131952e02, 6.75869174e01, -3.14987800e00, 9.61237896e-02, ] rescale_func = np.poly1d(coefficients) assert ( self.previous_modulated_input is not None ), "previous_modulated_input is not initialized" self.accumulated_rel_l1_distance += rescale_func( ( (modulated_inp - self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean() ) .cpu() .item() ) if self.accumulated_rel_l1_distance < teache_thresh: should_calc = False else: should_calc = True self.accumulated_rel_l1_distance = 0 self.previous_modulated_input = modulated_inp self.cnt += 1 return not should_calc def retrieve_cached_states(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states + self.previous_residual class SingleTokenRefiner(nn.Module): """ A token refiner that processes text embeddings with attention to improve their representation for cross-attention with image features. """ def __init__( self, in_channels, hidden_size, num_attention_heads, depth=2, qkv_bias=True, dtype=None, prefix: str = "", ) -> None: super().__init__() # Input projection self.input_embedder = ReplicatedLinear( in_channels, hidden_size, bias=True, params_dtype=dtype, prefix=f"{prefix}.input_embedder", ) # Timestep embedding self.t_embedder = TimestepEmbedder( hidden_size, act_layer="silu", dtype=dtype, prefix=f"{prefix}.t_embedder" ) # Context embedding self.c_embedder = MLP( in_channels, hidden_size, hidden_size, act_type="silu", dtype=dtype, prefix=f"{prefix}.c_embedder", ) # Refiner blocks self.refiner_blocks = nn.ModuleList( [ IndividualTokenRefinerBlock( hidden_size, num_attention_heads, qkv_bias=qkv_bias, dtype=dtype, prefix=f"{prefix}.refiner_blocks.{i}", ) for i in range(depth) ] ) def forward(self, x, t): # Get timestep embeddings timestep_aware_representations = self.t_embedder(t) # Get context-aware representations context_aware_representations = torch.mean(x, dim=1) context_aware_representations = self.c_embedder(context_aware_representations) c = timestep_aware_representations + context_aware_representations # Project input x, _ = self.input_embedder(x) # Process through refiner blocks for block in self.refiner_blocks: x = block(x, c) return x class IndividualTokenRefinerBlock(nn.Module): """ A transformer block for refining individual tokens with self-attention. """ def __init__( self, hidden_size, num_attention_heads, mlp_ratio=4.0, qkv_bias=True, dtype=None, prefix: str = "", ) -> None: super().__init__() self.num_attention_heads = num_attention_heads mlp_hidden_dim = int(hidden_size * mlp_ratio) # Normalization and attention self.norm1 = nn.LayerNorm( hidden_size, eps=1e-6, elementwise_affine=True, dtype=dtype ) self.self_attn_qkv = ReplicatedLinear( hidden_size, hidden_size * 3, bias=qkv_bias, params_dtype=dtype, prefix=f"{prefix}.self_attn_qkv", ) self.self_attn_proj = ReplicatedLinear( hidden_size, hidden_size, bias=qkv_bias, params_dtype=dtype, prefix=f"{prefix}.self_attn_proj", ) # MLP self.norm2 = nn.LayerNorm( hidden_size, eps=1e-6, elementwise_affine=True, dtype=dtype ) self.mlp = MLP( hidden_size, mlp_hidden_dim, bias=True, act_type="silu", dtype=dtype, prefix=f"{prefix}.mlp", ) # Modulation self.adaLN_modulation = ModulateProjection( hidden_size, factor=2, act_layer="silu", dtype=dtype, prefix=f"{prefix}.adaLN_modulation", ) # Scaled dot product attention self.attn = LocalAttention( num_heads=num_attention_heads, head_size=hidden_size // num_attention_heads, # TODO: remove hardcode; remove STA supported_attention_backends=( AttentionBackendEnum.FA, AttentionBackendEnum.AITER, AttentionBackendEnum.TORCH_SDPA, ), ) def forward(self, x, c): # Get modulation parameters gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=-1) # Self-attention norm_x = self.norm1(x) qkv, _ = self.self_attn_qkv(norm_x) batch_size, seq_len = qkv.shape[0], qkv.shape[1] qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1) q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] # Run scaled dot product attention attn_output = self.attn(q, k, v) # [B, L, H, D] attn_output = attn_output.reshape(batch_size, seq_len, -1) # [B, L, H*D] # Project and apply residual connection with gating attn_out, _ = self.self_attn_proj(attn_output) x = x + attn_out * gate_msa.unsqueeze(1) # MLP mlp_out = self.mlp(self.norm2(x)) x = x + mlp_out * gate_mlp.unsqueeze(1) return x class FinalLayer(nn.Module): """ The final layer of DiT that projects features to pixel space. """ def __init__( self, hidden_size, patch_size, out_channels, dtype=None, prefix: str = "" ) -> None: super().__init__() # Normalization self.norm_final = nn.LayerNorm( hidden_size, eps=1e-6, elementwise_affine=False, dtype=dtype ) output_dim = patch_size[0] * patch_size[1] * patch_size[2] * out_channels self.linear = ReplicatedLinear( hidden_size, output_dim, bias=True, params_dtype=dtype, prefix=f"{prefix}.linear", ) # Modulation self.adaLN_modulation = ModulateProjection( hidden_size, factor=2, act_layer="silu", dtype=dtype, prefix=f"{prefix}.adaLN_modulation", ) def forward(self, x, c): # What the heck HF? Why you change the scale and shift order here??? scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1) x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1) x, _ = self.linear(x) return x EntryClass = HunyuanVideoTransformer3DModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py ================================================ # Copied and adapted from LTX-2 and WanVideo implementations. # # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations from typing import Any, Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2ArchConfig, LTX2Config from sglang.multimodal_gen.runtime.distributed import ( get_sp_parallel_rank, get_sp_world_size, get_tp_rank, get_tp_world_size, model_parallel_is_initialized, ) from sglang.multimodal_gen.runtime.distributed.communication_op import ( tensor_model_parallel_all_reduce, ) from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) def apply_interleaved_rotary_emb( x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor] ) -> torch.Tensor: cos, sin = freqs x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1) x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) return (x.float() * cos + x_rotated.float() * sin).to(x.dtype) def apply_split_rotary_emb( x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor] ) -> torch.Tensor: cos, sin = freqs x_dtype = x.dtype needs_reshape = False if x.ndim != 4 and cos.ndim == 4: b = x.shape[0] _, h, t, _ = cos.shape x = x.reshape(b, t, h, -1).swapaxes(1, 2) needs_reshape = True last = x.shape[-1] if last % 2 != 0: raise ValueError( f"Expected x.shape[-1] to be even for split rotary, got {last}." ) r = last // 2 split_x = x.reshape(*x.shape[:-1], 2, r) first_x = split_x[..., :1, :] second_x = split_x[..., 1:, :] cos_u = cos.unsqueeze(-2) sin_u = sin.unsqueeze(-2) out = split_x * cos_u first_out = out[..., :1, :] second_out = out[..., 1:, :] first_out.addcmul_(-sin_u, second_x) second_out.addcmul_(sin_u, first_x) out = out.reshape(*out.shape[:-2], last) if needs_reshape: out = out.swapaxes(1, 2).reshape(b, t, -1) return out.to(dtype=x_dtype) # ============================================================================== # Layers and Embeddings # ============================================================================== class LTX2AudioVideoRotaryPosEmbed(nn.Module): def __init__( self, dim: int, patch_size: int = 1, patch_size_t: int = 1, base_num_frames: int = 20, base_height: int = 2048, base_width: int = 2048, sampling_rate: int = 16000, hop_length: int = 160, scale_factors: Tuple[int, ...] = (8, 32, 32), theta: float = 10000.0, causal_offset: int = 1, modality: str = "video", double_precision: bool = True, rope_type: str = "interleaved", num_attention_heads: int = 32, ) -> None: super().__init__() self.dim = int(dim) self.patch_size = int(patch_size) self.patch_size_t = int(patch_size_t) if rope_type not in ["interleaved", "split"]: raise ValueError( f"{rope_type=} not supported. Choose between 'interleaved' and 'split'." ) self.rope_type = rope_type self.base_num_frames = int(base_num_frames) self.num_attention_heads = int(num_attention_heads) self.base_height = int(base_height) self.base_width = int(base_width) self.sampling_rate = int(sampling_rate) self.hop_length = int(hop_length) self.audio_latents_per_second = ( float(self.sampling_rate) / float(self.hop_length) / float(scale_factors[0]) ) self.scale_factors = tuple(int(x) for x in scale_factors) self.theta = float(theta) self.causal_offset = int(causal_offset) self.modality = modality self.coords_dtype = torch.bfloat16 if modality == "video" else torch.float32 if self.modality not in ["video", "audio"]: raise ValueError( f"Modality {modality} is not supported. Supported modalities are `video` and `audio`." ) self.double_precision = bool(double_precision) def prepare_video_coords( self, batch_size: int, num_frames: int, height: int, width: int, device: torch.device, fps: float = 24.0, *, start_frame: int = 0, ) -> torch.Tensor: grid_f = torch.arange( start=int(start_frame), end=int(num_frames) + int(start_frame), step=self.patch_size_t, dtype=torch.float32, device=device, ) grid_h = torch.arange( start=0, end=height, step=self.patch_size, dtype=torch.float32, device=device, ) grid_w = torch.arange( start=0, end=width, step=self.patch_size, dtype=torch.float32, device=device, ) grid = torch.meshgrid(grid_f, grid_h, grid_w, indexing="ij") grid = torch.stack(grid, dim=0) patch_size = (self.patch_size_t, self.patch_size, self.patch_size) patch_size_delta = torch.tensor( patch_size, dtype=grid.dtype, device=grid.device ) patch_ends = grid + patch_size_delta.view(3, 1, 1, 1) latent_coords = torch.stack([grid, patch_ends], dim=-1) latent_coords = latent_coords.flatten(1, 3) latent_coords = latent_coords.unsqueeze(0).repeat(batch_size, 1, 1, 1) scale_tensor = torch.tensor(self.scale_factors, device=latent_coords.device) broadcast_shape = [1] * latent_coords.ndim broadcast_shape[1] = -1 pixel_coords = latent_coords * scale_tensor.view(*broadcast_shape) pixel_coords[:, 0, ...] = ( pixel_coords[:, 0, ...] + self.causal_offset - self.scale_factors[0] ).clamp(min=0) pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps return pixel_coords def prepare_audio_coords( self, batch_size: int, num_frames: int, device: torch.device, *, start_frame: int = 0, ) -> torch.Tensor: grid_f = torch.arange( start=int(start_frame), end=int(num_frames) + int(start_frame), step=self.patch_size_t, dtype=torch.float32, device=device, ) audio_scale_factor = self.scale_factors[0] grid_start_mel = grid_f * audio_scale_factor grid_start_mel = ( grid_start_mel + self.causal_offset - audio_scale_factor ).clip(min=0) grid_start_s = grid_start_mel * self.hop_length / self.sampling_rate grid_end_mel = (grid_f + self.patch_size_t) * audio_scale_factor grid_end_mel = (grid_end_mel + self.causal_offset - audio_scale_factor).clip( min=0 ) grid_end_s = grid_end_mel * self.hop_length / self.sampling_rate audio_coords = torch.stack([grid_start_s, grid_end_s], dim=-1) audio_coords = audio_coords.unsqueeze(0).expand(batch_size, -1, -1) audio_coords = audio_coords.unsqueeze(1) return audio_coords def prepare_coords(self, *args, **kwargs): if self.modality == "video": return self.prepare_video_coords(*args, **kwargs) return self.prepare_audio_coords(*args, **kwargs) def forward( self, coords: torch.Tensor, device: Optional[Union[str, torch.device]] = None ) -> Tuple[torch.Tensor, torch.Tensor]: device = device or coords.device num_pos_dims = coords.shape[1] coords = coords.to(self.coords_dtype) if coords.ndim == 4: coords_start, coords_end = coords.chunk(2, dim=-1) coords = (coords_start + coords_end) / 2.0 coords = coords.squeeze(-1) if self.modality == "video": max_positions = (self.base_num_frames, self.base_height, self.base_width) else: max_positions = (self.base_num_frames,) grid = torch.stack( [coords[:, i] / max_positions[i] for i in range(num_pos_dims)], dim=-1 ).to(device) num_rope_elems = num_pos_dims * 2 freqs_dtype = torch.float64 if self.double_precision else torch.float32 pow_indices = torch.pow( self.theta, torch.linspace( start=0.0, end=1.0, steps=self.dim // num_rope_elems, dtype=freqs_dtype, device=device, ), ) freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32) freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs freqs = freqs.transpose(-1, -2).flatten(2) if self.rope_type == "interleaved": cos_freqs = freqs.cos().repeat_interleave(2, dim=-1) sin_freqs = freqs.sin().repeat_interleave(2, dim=-1) if self.dim % num_rope_elems != 0: cos_padding = torch.ones_like( cos_freqs[:, :, : self.dim % num_rope_elems] ) sin_padding = torch.zeros_like( cos_freqs[:, :, : self.dim % num_rope_elems] ) cos_freqs = torch.cat([cos_padding, cos_freqs], dim=-1) sin_freqs = torch.cat([sin_padding, sin_freqs], dim=-1) else: expected_freqs = self.dim // 2 current_freqs = freqs.shape[-1] pad_size = expected_freqs - current_freqs cos_freq = freqs.cos() sin_freq = freqs.sin() if pad_size != 0: cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) cos_freq = torch.cat([cos_padding, cos_freq], dim=-1) sin_freq = torch.cat([sin_padding, sin_freq], dim=-1) b = cos_freq.shape[0] t = cos_freq.shape[1] cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, -1) sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, -1) cos_freqs = torch.swapaxes(cos_freq, 1, 2) sin_freqs = torch.swapaxes(sin_freq, 1, 2) # Cast to bf16 to match model weights dtype. coords_dtype controls # intermediate coordinate precision (fp32 for audio) and differs. return cos_freqs.to(torch.bfloat16), sin_freqs.to(torch.bfloat16) def rms_norm(x: torch.Tensor, eps: float) -> torch.Tensor: return F.rms_norm(x, normalized_shape=(x.shape[-1],), eps=eps) class LTX2TextProjection(nn.Module): def __init__( self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh", ) -> None: super().__init__() if out_features is None: out_features = hidden_size self.linear_1 = ColumnParallelLinear( in_features, hidden_size, bias=True, gather_output=True ) if act_fn == "gelu_tanh": self.act_1 = nn.GELU(approximate="tanh") elif act_fn == "silu": self.act_1 = nn.SiLU() else: raise ValueError(f"Unknown activation function: {act_fn}") self.linear_2 = ColumnParallelLinear( hidden_size, out_features, bias=True, gather_output=True ) def forward(self, caption: torch.Tensor) -> torch.Tensor: hidden_states, _ = self.linear_1(caption) hidden_states = self.act_1(hidden_states) hidden_states, _ = self.linear_2(hidden_states) return hidden_states class LTX2TimestepEmbedder(nn.Module): def __init__(self, embedding_dim: int, in_channels: int = 256) -> None: super().__init__() self.linear_1 = ColumnParallelLinear( in_channels, embedding_dim, bias=True, gather_output=True ) self.linear_2 = ColumnParallelLinear( embedding_dim, embedding_dim, bias=True, gather_output=True ) def forward(self, t_emb: torch.Tensor) -> torch.Tensor: x, _ = self.linear_1(t_emb) x = F.silu(x) x, _ = self.linear_2(x) return x class LTX2PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module): def __init__(self, embedding_dim: int) -> None: super().__init__() self.timestep_embedder = LTX2TimestepEmbedder(embedding_dim, in_channels=256) def forward( self, timestep: torch.Tensor, hidden_dtype: torch.dtype | None = None ) -> torch.Tensor: t = timestep.reshape(-1).to(dtype=torch.float32) t_emb = timestep_embedding(t, dim=256, max_period=10000, dtype=torch.float32) if hidden_dtype is not None: t_emb = t_emb.to(dtype=hidden_dtype) return self.timestep_embedder(t_emb) class LTX2AdaLayerNormSingle(nn.Module): def __init__(self, embedding_dim: int, embedding_coefficient: int = 6) -> None: super().__init__() self.emb = LTX2PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim) self.silu = nn.SiLU() self.linear = ColumnParallelLinear( embedding_dim, embedding_coefficient * embedding_dim, bias=True, gather_output=True, ) def forward( self, timestep: torch.Tensor, hidden_dtype: torch.dtype | None = None ) -> tuple[torch.Tensor, torch.Tensor]: embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype).to( dtype=self.linear.weight.dtype ) out, _ = self.linear(self.silu(embedded_timestep)) return out, embedded_timestep class LTX2TPRMSNormAcrossHeads(nn.Module): def __init__( self, full_hidden_size: int, local_hidden_size: int, eps: float ) -> None: super().__init__() self.full_hidden_size = full_hidden_size self.local_hidden_size = local_hidden_size self.eps = eps self.weight = nn.Parameter(torch.ones(local_hidden_size)) tp_rank = get_tp_rank() def _weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: shard = loaded_weight.narrow( 0, tp_rank * local_hidden_size, local_hidden_size ) param.data.copy_(shard.to(dtype=param.dtype, device=param.device)) setattr(self.weight, "weight_loader", _weight_loader) def forward(self, x: torch.Tensor) -> torch.Tensor: # Keep track of the original dtype. We do the statistics in fp32 for # numerical stability, but cast the output back to the input dtype to orig_dtype = x.dtype if get_tp_world_size() == 1: var = x.float().pow(2).mean(dim=-1, keepdim=True) else: local_sumsq = x.float().pow(2).sum(dim=-1, keepdim=True) global_sumsq = tensor_model_parallel_all_reduce(local_sumsq) var = global_sumsq / float(self.full_hidden_size) inv_rms_fp32 = torch.rsqrt(var + self.eps) y = (x.float() * inv_rms_fp32).to(dtype=orig_dtype) return y * self.weight.to(dtype=orig_dtype) class LTX2Attention(nn.Module): def __init__( self, query_dim: int, context_dim: int | None = None, heads: int = 8, dim_head: int = 64, norm_eps: float = 1e-6, qk_norm: bool = True, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", quant_config: QuantizationConfig | None = None, ) -> None: super().__init__() self.query_dim = int(query_dim) self.context_dim = int(query_dim if context_dim is None else context_dim) self.heads = int(heads) self.dim_head = int(dim_head) self.inner_dim = self.heads * self.dim_head self.norm_eps = float(norm_eps) self.qk_norm = bool(qk_norm) tp_size = get_tp_world_size() if tp_size <= 0: raise ValueError(f"Invalid {tp_size=}. Expected tp_size >= 1.") if self.heads % tp_size != 0: raise ValueError( f"LTX2Attention requires heads divisible by tp_size, got " f"{self.heads=} {tp_size=}." ) if self.inner_dim % tp_size != 0: # This should follow from heads % tp_size, but keep explicit for clarity. raise ValueError( f"LTX2Attention requires inner_dim divisible by tp_size, got " f"{self.inner_dim=} {tp_size=}." ) self.local_heads = self.heads // tp_size self.to_q = ColumnParallelLinear( self.query_dim, self.inner_dim, bias=True, gather_output=False, quant_config=quant_config, ) self.to_k = ColumnParallelLinear( self.context_dim, self.inner_dim, bias=True, gather_output=False, quant_config=quant_config, ) self.to_v = ColumnParallelLinear( self.context_dim, self.inner_dim, bias=True, gather_output=False, quant_config=quant_config, ) self.q_norm: nn.Module | None = None self.k_norm: nn.Module | None = None if self.qk_norm: if tp_size == 1: self.q_norm = torch.nn.RMSNorm(self.inner_dim, eps=self.norm_eps) self.k_norm = torch.nn.RMSNorm(self.inner_dim, eps=self.norm_eps) else: self.q_norm = LTX2TPRMSNormAcrossHeads( full_hidden_size=self.inner_dim, local_hidden_size=self.inner_dim // tp_size, eps=self.norm_eps, ) self.k_norm = LTX2TPRMSNormAcrossHeads( full_hidden_size=self.inner_dim, local_hidden_size=self.inner_dim // tp_size, eps=self.norm_eps, ) self.to_out = nn.Sequential( RowParallelLinear( self.inner_dim, self.query_dim, bias=True, input_is_parallel=True, quant_config=quant_config, ), nn.Identity(), ) self.attn = USPAttention( num_heads=self.local_heads, head_size=self.dim_head, num_kv_heads=self.local_heads, dropout_rate=0, softmax_scale=None, causal=False, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn", ) def forward( self, x: torch.Tensor, context: torch.Tensor | None = None, mask: torch.Tensor | None = None, pe: tuple[torch.Tensor, torch.Tensor] | None = None, k_pe: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: q, _ = self.to_q(x) context_ = x if context is None else context k, _ = self.to_k(context_) v, _ = self.to_v(context_) if self.qk_norm: assert self.q_norm is not None and self.k_norm is not None q = self.q_norm(q) k = self.k_norm(k) if pe is not None: cos, sin = pe k_cos, k_sin = pe if k_pe is None else k_pe tp_size = get_tp_world_size() if tp_size > 1: tp_rank = get_tp_rank() cos, sin = self._slice_rope_for_tp( cos, sin, tp_rank=tp_rank, tp_size=tp_size ) k_cos, k_sin = self._slice_rope_for_tp( k_cos, k_sin, tp_rank=tp_rank, tp_size=tp_size ) if cos.dim() == 3: q = apply_interleaved_rotary_emb(q, (cos, sin)) k = apply_interleaved_rotary_emb(k, (k_cos, k_sin)) else: q = apply_split_rotary_emb(q, (cos, sin)) k = apply_split_rotary_emb(k, (k_cos, k_sin)) q = q.view(*q.shape[:-1], self.local_heads, self.dim_head) k = k.view(*k.shape[:-1], self.local_heads, self.dim_head) v = v.view(*v.shape[:-1], self.local_heads, self.dim_head) if mask is not None: # Fallback to SDPA for masked attention q_ = q.transpose(1, 2) k_ = k.transpose(1, 2) v_ = v.transpose(1, 2) if torch.is_floating_point(mask): m = mask if m.dim() == 2: m = m[:, None, None, :] elif m.dim() == 3: m = m[:, None, :, :] sdpa_mask = m.to(dtype=q_.dtype, device=q_.device) else: m = mask.to(dtype=q_.dtype, device=q_.device) if m.dim() == 2: m = m[:, None, None, :] elif m.dim() == 3: m = m[:, None, :, :] sdpa_mask = (m - 1.0) * torch.finfo(q_.dtype).max out = torch.nn.functional.scaled_dot_product_attention( q_, k_, v_, attn_mask=sdpa_mask, dropout_p=0.0, is_causal=False ).transpose(1, 2) else: out = self.attn(q, k, v) out = out.flatten(2) out, _ = self.to_out[0](out) return out def _slice_rope_for_tp( self, cos: torch.Tensor, sin: torch.Tensor, *, tp_rank: int, tp_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Slice RoPE tensors to the local TP shard. - split-rope: cos/sin are shaped [B, H, T, R] (head-major), slice by heads. - interleaved-rope: cos/sin are shaped [B, T, D], where D matches the projected feature dimension and is sharded by TP. """ if cos.ndim == 4: # [B, H, T, R] start = tp_rank * self.local_heads end = start + self.local_heads return cos[:, start:end, :, :], sin[:, start:end, :, :] elif cos.ndim == 3: # [B, T, D] d = cos.shape[-1] if d % tp_size != 0: raise ValueError( f"RoPE dim must be divisible by tp_size, got {d=} {tp_size=}." ) local_d = d // tp_size start = tp_rank * local_d end = start + local_d return cos[:, :, start:end], sin[:, :, start:end] raise ValueError(f"Unexpected RoPE tensor rank: {cos.ndim}. Expected 3 or 4.") class LTX2FeedForward(nn.Module): def __init__( self, dim: int, dim_out: int | None = None, mult: int = 4, quant_config: QuantizationConfig | None = None, ) -> None: super().__init__() if dim_out is None: dim_out = dim inner_dim = int(dim * mult) self.proj_in = ColumnParallelLinear( dim, inner_dim, bias=True, gather_output=True, quant_config=quant_config ) self.act = nn.GELU(approximate="tanh") self.proj_out = ColumnParallelLinear( inner_dim, dim_out, bias=True, gather_output=True, quant_config=quant_config ) def forward(self, x: torch.Tensor) -> torch.Tensor: x, _ = self.proj_in(x) x = self.act(x) x, _ = self.proj_out(x) return x class LTX2TransformerBlock(nn.Module): def __init__( self, idx: int, dim: int, num_attention_heads: int, attention_head_dim: int, cross_attention_dim: int, audio_dim: int, audio_num_attention_heads: int, audio_attention_head_dim: int, audio_cross_attention_dim: int, qk_norm: bool = True, norm_eps: float = 1e-6, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", quant_config: QuantizationConfig | None = None, ): super().__init__() self.idx = idx self.norm_eps = norm_eps # 1. Self-Attention (video and audio) self.attn1 = LTX2Attention( query_dim=dim, heads=num_attention_heads, dim_head=attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn1", quant_config=quant_config, ) self.audio_attn1 = LTX2Attention( query_dim=audio_dim, heads=audio_num_attention_heads, dim_head=audio_attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.audio_attn1", quant_config=quant_config, ) # 2. Prompt Cross-Attention self.attn2 = LTX2Attention( query_dim=dim, context_dim=cross_attention_dim, heads=num_attention_heads, dim_head=attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn2", quant_config=quant_config, ) self.audio_attn2 = LTX2Attention( query_dim=audio_dim, context_dim=audio_cross_attention_dim, heads=audio_num_attention_heads, dim_head=audio_attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.audio_attn2", quant_config=quant_config, ) # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention self.audio_to_video_attn = LTX2Attention( query_dim=dim, context_dim=audio_dim, heads=audio_num_attention_heads, dim_head=audio_attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.audio_to_video_attn", quant_config=quant_config, ) self.video_to_audio_attn = LTX2Attention( query_dim=audio_dim, context_dim=dim, heads=audio_num_attention_heads, dim_head=audio_attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.video_to_audio_attn", quant_config=quant_config, ) # 4. Feedforward layers self.ff = LTX2FeedForward(dim, dim_out=dim, quant_config=quant_config) self.audio_ff = LTX2FeedForward( audio_dim, dim_out=audio_dim, quant_config=quant_config ) # 5. Modulation Parameters self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) self.audio_scale_shift_table = nn.Parameter( torch.randn(6, audio_dim) / audio_dim**0.5 ) self.video_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, dim)) self.audio_a2v_cross_attn_scale_shift_table = nn.Parameter( torch.randn(5, audio_dim) ) def get_ada_values( self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice, ) -> tuple[torch.Tensor, ...]: num_ada_params = int(scale_shift_table.shape[0]) ada_values = ( scale_shift_table[indices] .unsqueeze(0) .unsqueeze(0) .to(device=timestep.device, dtype=timestep.dtype) + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[ :, :, indices, : ] ).unbind(dim=2) return [t.squeeze(2) for t in ada_values] def forward( self, hidden_states: torch.Tensor, audio_hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, audio_encoder_hidden_states: torch.Tensor, temb: torch.Tensor, temb_audio: torch.Tensor, temb_ca_scale_shift: torch.Tensor, temb_ca_audio_scale_shift: torch.Tensor, temb_ca_gate: torch.Tensor, temb_ca_audio_gate: torch.Tensor, video_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, audio_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ca_video_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ca_audio_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, encoder_attention_mask: Optional[torch.Tensor] = None, audio_encoder_attention_mask: Optional[torch.Tensor] = None, a2v_cross_attention_mask: Optional[torch.Tensor] = None, v2a_cross_attention_mask: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor]: batch_size = hidden_states.size(0) # 1. Video and Audio Self-Attention vshift_msa, vscale_msa, vgate_msa = self.get_ada_values( self.scale_shift_table, batch_size, temb, slice(0, 3) ) norm_hidden_states = ( rms_norm(hidden_states, self.norm_eps) * (1 + vscale_msa) + vshift_msa ) attn_hidden_states = self.attn1(norm_hidden_states, pe=video_rotary_emb) hidden_states = hidden_states + attn_hidden_states * vgate_msa ashift_msa, ascale_msa, agate_msa = self.get_ada_values( self.audio_scale_shift_table, batch_size, temb_audio, slice(0, 3) ) norm_audio_hidden_states = ( rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_msa) + ashift_msa ) attn_audio_hidden_states = self.audio_attn1( norm_audio_hidden_states, pe=audio_rotary_emb ) audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * agate_msa # 2. Prompt Cross-Attention norm_hidden_states = rms_norm(hidden_states, self.norm_eps) attn_hidden_states = self.attn2( norm_hidden_states, context=encoder_hidden_states, mask=encoder_attention_mask, ) hidden_states = hidden_states + attn_hidden_states norm_audio_hidden_states = rms_norm(audio_hidden_states, self.norm_eps) attn_audio_hidden_states = self.audio_attn2( norm_audio_hidden_states, context=audio_encoder_hidden_states, mask=audio_encoder_attention_mask, ) audio_hidden_states = audio_hidden_states + attn_audio_hidden_states # 3. Audio-to-Video and Video-to-Audio Cross-Attention norm_hidden_states = rms_norm(hidden_states, self.norm_eps) norm_audio_hidden_states = rms_norm(audio_hidden_states, self.norm_eps) # Compute combined ada params video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[ :4, : ] video_per_layer_ca_gate = self.video_a2v_cross_attn_scale_shift_table[4:, :] video_ca_scale_shift_table = ( video_per_layer_ca_scale_shift[None, None, :, :].to( dtype=temb_ca_scale_shift.dtype, device=temb_ca_scale_shift.device ) + temb_ca_scale_shift.reshape( batch_size, temb_ca_scale_shift.shape[1], 4, -1 ) ).unbind(dim=2) video_ca_gate = ( video_per_layer_ca_gate[None, None, :, :].to( dtype=temb_ca_gate.dtype, device=temb_ca_gate.device ) + temb_ca_gate.reshape(batch_size, temb_ca_gate.shape[1], 1, -1) ).unbind(dim=2) ( video_a2v_ca_scale, video_a2v_ca_shift, video_v2a_ca_scale, video_v2a_ca_shift, ) = [t.squeeze(2) for t in video_ca_scale_shift_table] a2v_gate = video_ca_gate[0].squeeze(2) audio_per_layer_ca_scale_shift = self.audio_a2v_cross_attn_scale_shift_table[ :4, : ] audio_per_layer_ca_gate = self.audio_a2v_cross_attn_scale_shift_table[4:, :] audio_ca_scale_shift_table = ( audio_per_layer_ca_scale_shift[None, None, :, :].to( dtype=temb_ca_audio_scale_shift.dtype, device=temb_ca_audio_scale_shift.device, ) + temb_ca_audio_scale_shift.reshape( batch_size, temb_ca_audio_scale_shift.shape[1], 4, -1 ) ).unbind(dim=2) audio_ca_gate = ( audio_per_layer_ca_gate[None, None, :, :].to( dtype=temb_ca_audio_gate.dtype, device=temb_ca_audio_gate.device ) + temb_ca_audio_gate.reshape(batch_size, temb_ca_audio_gate.shape[1], 1, -1) ).unbind(dim=2) ( audio_a2v_ca_scale, audio_a2v_ca_shift, audio_v2a_ca_scale, audio_v2a_ca_shift, ) = [t.squeeze(2) for t in audio_ca_scale_shift_table] v2a_gate = audio_ca_gate[0].squeeze(2) # A2V mod_norm_hidden_states = ( norm_hidden_states * (1 + video_a2v_ca_scale) + video_a2v_ca_shift ) mod_norm_audio_hidden_states = ( norm_audio_hidden_states * (1 + audio_a2v_ca_scale) + audio_a2v_ca_shift ) a2v_attn_hidden_states = self.audio_to_video_attn( mod_norm_hidden_states, context=mod_norm_audio_hidden_states, pe=ca_video_rotary_emb, k_pe=ca_audio_rotary_emb, mask=a2v_cross_attention_mask, ) hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states # V2A mod_norm_hidden_states = ( norm_hidden_states * (1 + video_v2a_ca_scale) + video_v2a_ca_shift ) mod_norm_audio_hidden_states = ( norm_audio_hidden_states * (1 + audio_v2a_ca_scale) + audio_v2a_ca_shift ) v2a_attn_hidden_states = self.video_to_audio_attn( mod_norm_audio_hidden_states, context=mod_norm_hidden_states, pe=ca_audio_rotary_emb, k_pe=ca_video_rotary_emb, mask=v2a_cross_attention_mask, ) audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states # 4. Feedforward vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values( self.scale_shift_table, batch_size, temb, slice(3, None) ) norm_hidden_states = ( rms_norm(hidden_states, self.norm_eps) * (1 + vscale_mlp) + vshift_mlp ) ff_output = self.ff(norm_hidden_states) hidden_states = hidden_states + ff_output * vgate_mlp ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values( self.audio_scale_shift_table, batch_size, temb_audio, slice(3, None) ) norm_audio_hidden_states = ( rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_mlp) + ashift_mlp ) audio_ff_output = self.audio_ff(norm_audio_hidden_states) audio_hidden_states = audio_hidden_states + audio_ff_output * agate_mlp return hidden_states, audio_hidden_states class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): _fsdp_shard_conditions = LTX2ArchConfig()._fsdp_shard_conditions _compile_conditions = LTX2ArchConfig()._compile_conditions _supported_attention_backends = LTX2ArchConfig()._supported_attention_backends param_names_mapping = LTX2ArchConfig().param_names_mapping reverse_param_names_mapping = LTX2ArchConfig().reverse_param_names_mapping lora_param_names_mapping = LTX2ArchConfig().lora_param_names_mapping def _validate_tp_config(self, *, arch: LTX2ArchConfig, tp_size: int) -> None: """Validate TP-related dimension constraints (fail-fast).""" if tp_size < 1: raise ValueError(f"Invalid tp_size={tp_size}. Expected tp_size >= 1.") if self.hidden_size % self.num_attention_heads != 0: raise ValueError( "video hidden_size must be divisible by num_attention_heads, got " f"{self.hidden_size=} {self.num_attention_heads=}." ) if self.audio_hidden_size % self.audio_num_attention_heads != 0: raise ValueError( "audio_hidden_size must be divisible by audio_num_attention_heads, got " f"{self.audio_hidden_size=} {self.audio_num_attention_heads=}." ) if tp_size == 1: return if self.num_attention_heads % tp_size != 0: raise ValueError( "num_attention_heads must be divisible by tp_size, got " f"{self.num_attention_heads=} {tp_size=}." ) if self.audio_num_attention_heads % tp_size != 0: raise ValueError( "audio_num_attention_heads must be divisible by tp_size, got " f"{self.audio_num_attention_heads=} {tp_size=}." ) if self.hidden_size % tp_size != 0: raise ValueError( "hidden_size must be divisible by tp_size for TP-sharded projections, got " f"{self.hidden_size=} {tp_size=}." ) if self.audio_hidden_size % tp_size != 0: raise ValueError( "audio_hidden_size must be divisible by tp_size for TP-sharded projections, got " f"{self.audio_hidden_size=} {tp_size=}." ) if int(arch.out_channels) % tp_size != 0: raise ValueError( "out_channels must be divisible by tp_size for TP-sharded output projection, got " f"{arch.out_channels=} {tp_size=}." ) if int(arch.audio_out_channels) % tp_size != 0: raise ValueError( "audio_out_channels must be divisible by tp_size for TP-sharded output projection, got " f"{arch.audio_out_channels=} {tp_size=}." ) def __init__( self, config: LTX2Config, hf_config: dict[str, Any], quant_config: QuantizationConfig | None = None, ) -> None: super().__init__(config=config, hf_config=hf_config) arch = config.arch_config self.hidden_size = arch.hidden_size self.num_attention_heads = arch.num_attention_heads self.audio_hidden_size = arch.audio_hidden_size self.audio_num_attention_heads = arch.audio_num_attention_heads self.norm_eps = arch.norm_eps tp_size = get_tp_world_size() self._validate_tp_config(arch=arch, tp_size=tp_size) # 1. Patchification input projections # Matches LTX2Config().param_names_mapping self.patchify_proj = ColumnParallelLinear( arch.in_channels, self.hidden_size, bias=True, gather_output=True, quant_config=quant_config, ) self.audio_patchify_proj = ColumnParallelLinear( arch.audio_in_channels, self.audio_hidden_size, bias=True, gather_output=True, quant_config=quant_config, ) # 2. Prompt embeddings self.caption_projection = LTX2TextProjection( in_features=arch.caption_channels, hidden_size=self.hidden_size ) self.audio_caption_projection = LTX2TextProjection( in_features=arch.caption_channels, hidden_size=self.audio_hidden_size ) # 3. Timestep Modulation Params and Embedding self.adaln_single = LTX2AdaLayerNormSingle( self.hidden_size, embedding_coefficient=6 ) self.audio_adaln_single = LTX2AdaLayerNormSingle( self.audio_hidden_size, embedding_coefficient=6 ) # Global Cross Attention Modulation Parameters self.av_ca_video_scale_shift_adaln_single = LTX2AdaLayerNormSingle( self.hidden_size, embedding_coefficient=4 ) self.av_ca_a2v_gate_adaln_single = LTX2AdaLayerNormSingle( self.hidden_size, embedding_coefficient=1 ) self.av_ca_audio_scale_shift_adaln_single = LTX2AdaLayerNormSingle( self.audio_hidden_size, embedding_coefficient=4 ) self.av_ca_v2a_gate_adaln_single = LTX2AdaLayerNormSingle( self.audio_hidden_size, embedding_coefficient=1 ) # Output Layer Scale/Shift Modulation parameters self.scale_shift_table = nn.Parameter( torch.randn(2, self.hidden_size) / self.hidden_size**0.5 ) self.audio_scale_shift_table = nn.Parameter( torch.randn(2, self.audio_hidden_size) / self.audio_hidden_size**0.5 ) hf_patch_size = int(hf_config.get("patch_size", 1)) hf_patch_size_t = int(hf_config.get("patch_size_t", 1)) self.patch_size = (hf_patch_size_t, hf_patch_size, hf_patch_size) hf_audio_patch_size = int(hf_config.get("audio_patch_size", 1)) hf_audio_patch_size_t = int(hf_config.get("audio_patch_size_t", 1)) rope_type = ( arch.rope_type.value if hasattr(arch.rope_type, "value") else str(arch.rope_type) ) rope_double_precision = bool( hf_config.get("rope_double_precision", arch.double_precision_rope) ) causal_offset = int(hf_config.get("causal_offset", 1)) pos_embed_max_pos = int(arch.positional_embedding_max_pos[0]) base_height = int(arch.positional_embedding_max_pos[1]) base_width = int(arch.positional_embedding_max_pos[2]) audio_pos_embed_max_pos = int(arch.audio_positional_embedding_max_pos[0]) self.video_scale_factors = (8, 32, 32) self.audio_scale_factors = (4,) self.rope = LTX2AudioVideoRotaryPosEmbed( dim=self.hidden_size, patch_size=hf_patch_size, patch_size_t=hf_patch_size_t, base_num_frames=pos_embed_max_pos, base_height=base_height, base_width=base_width, scale_factors=self.video_scale_factors, theta=float(arch.positional_embedding_theta), causal_offset=causal_offset, modality="video", double_precision=rope_double_precision, rope_type=rope_type, num_attention_heads=self.num_attention_heads, ) self.audio_rope = LTX2AudioVideoRotaryPosEmbed( dim=self.audio_hidden_size, patch_size=hf_audio_patch_size, patch_size_t=hf_audio_patch_size_t, base_num_frames=audio_pos_embed_max_pos, sampling_rate=16000, hop_length=160, scale_factors=self.audio_scale_factors, theta=float(arch.positional_embedding_theta), causal_offset=causal_offset, modality="audio", double_precision=rope_double_precision, rope_type=rope_type, num_attention_heads=self.audio_num_attention_heads, ) cross_attn_pos_embed_max_pos = max(pos_embed_max_pos, audio_pos_embed_max_pos) self.cross_attn_rope = LTX2AudioVideoRotaryPosEmbed( dim=int(arch.audio_cross_attention_dim), patch_size=hf_patch_size, patch_size_t=hf_patch_size_t, base_num_frames=cross_attn_pos_embed_max_pos, base_height=base_height, base_width=base_width, theta=float(arch.positional_embedding_theta), causal_offset=causal_offset, modality="video", double_precision=rope_double_precision, rope_type=rope_type, num_attention_heads=self.num_attention_heads, ) self.cross_attn_audio_rope = LTX2AudioVideoRotaryPosEmbed( dim=int(arch.audio_cross_attention_dim), patch_size=hf_audio_patch_size, patch_size_t=hf_audio_patch_size_t, base_num_frames=cross_attn_pos_embed_max_pos, sampling_rate=16000, hop_length=160, theta=float(arch.positional_embedding_theta), causal_offset=causal_offset, modality="audio", double_precision=rope_double_precision, rope_type=rope_type, num_attention_heads=self.audio_num_attention_heads, ) self.cross_pe_max_pos = cross_attn_pos_embed_max_pos # 5. Transformer Blocks self.transformer_blocks = nn.ModuleList( [ LTX2TransformerBlock( idx=idx, dim=self.hidden_size, num_attention_heads=self.num_attention_heads, attention_head_dim=self.hidden_size // self.num_attention_heads, cross_attention_dim=arch.cross_attention_dim, audio_dim=self.audio_hidden_size, audio_num_attention_heads=self.audio_num_attention_heads, audio_attention_head_dim=self.audio_hidden_size // self.audio_num_attention_heads, audio_cross_attention_dim=arch.audio_cross_attention_dim, norm_eps=self.norm_eps, qk_norm=True, # Always True in LTX2 supported_attention_backends=self._supported_attention_backends, prefix=config.prefix, quant_config=quant_config, ) for idx in range(arch.num_layers) ] ) # 6. Output layers self.norm_out = nn.LayerNorm( self.hidden_size, eps=self.norm_eps, elementwise_affine=False ) self.proj_out = ColumnParallelLinear( self.hidden_size, arch.out_channels, bias=True, gather_output=True, quant_config=quant_config, ) self.audio_norm_out = nn.LayerNorm( self.audio_hidden_size, eps=self.norm_eps, elementwise_affine=False ) self.audio_proj_out = ColumnParallelLinear( self.audio_hidden_size, arch.audio_out_channels, bias=True, gather_output=True, quant_config=quant_config, ) self.out_channels_raw = arch.out_channels // ( self.patch_size[0] * self.patch_size[1] * self.patch_size[2] ) self.audio_out_channels = arch.audio_out_channels self.timestep_scale_multiplier = arch.timestep_scale_multiplier self.av_ca_timestep_scale_multiplier = arch.av_ca_timestep_scale_multiplier self.layer_names = ["transformer_blocks"] def forward( self, hidden_states: torch.Tensor, audio_hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, audio_encoder_hidden_states: torch.Tensor, timestep: torch.LongTensor, audio_timestep: Optional[torch.LongTensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, audio_encoder_attention_mask: Optional[torch.Tensor] = None, num_frames: Optional[int] = None, height: Optional[int] = None, width: Optional[int] = None, fps: float = 24.0, audio_num_frames: Optional[int] = None, video_coords: Optional[torch.Tensor] = None, audio_coords: Optional[torch.Tensor] = None, **kwargs, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: batch_size = hidden_states.size(0) audio_timestep = audio_timestep if audio_timestep is not None else timestep if num_frames is None or height is None or width is None: raise ValueError( "num_frames/height/width must be provided for RoPE coordinate generation." ) if audio_num_frames is None: raise ValueError( "audio_num_frames must be provided for RoPE coordinate generation." ) if video_coords is None: # Wan-style SP-RoPE: when SP is enabled, each rank runs on its local # time shard but RoPE positions must be offset to global time. # # We assume equal time sharding across SP ranks. if model_parallel_is_initialized(): sp_world_size = get_sp_world_size() sp_rank = get_sp_parallel_rank() else: sp_world_size = 1 sp_rank = 0 video_shift = int(sp_rank) * int(num_frames) if sp_world_size > 1 else 0 video_coords = self.rope.prepare_video_coords( batch_size=batch_size, num_frames=num_frames, height=height, width=width, device=hidden_states.device, fps=fps, start_frame=video_shift, ) if audio_coords is None: audio_coords = self.audio_rope.prepare_audio_coords( batch_size=batch_size, num_frames=audio_num_frames, device=audio_hidden_states.device, ) video_rotary_emb = self.rope(video_coords, device=hidden_states.device) audio_rotary_emb = self.audio_rope( audio_coords, device=audio_hidden_states.device ) ca_video_rotary_emb = self.cross_attn_rope( video_coords[:, 0:1, :], device=hidden_states.device ) ca_audio_rotary_emb = self.cross_attn_audio_rope( audio_coords[:, 0:1, :], device=audio_hidden_states.device ) # 2. Patchify input projections hidden_states, _ = self.patchify_proj(hidden_states) audio_hidden_states, _ = self.audio_patchify_proj(audio_hidden_states) # 3. Prepare timestep embeddings # 3.1. Prepare global modality (video and audio) timestep embedding and modulation parameters temb, embedded_timestep = self.adaln_single( timestep.flatten(), ) temb = temb.view(batch_size, -1, temb.size(-1)) embedded_timestep = embedded_timestep.view( batch_size, -1, embedded_timestep.size(-1) ) temb_audio, audio_embedded_timestep = self.audio_adaln_single( audio_timestep.flatten() ) temb_audio = temb_audio.view(batch_size, -1, temb_audio.size(-1)) audio_embedded_timestep = audio_embedded_timestep.view( batch_size, -1, audio_embedded_timestep.size(-1) ) # 3.2. Prepare global modality cross attention modulation parameters ts_ca_mult = ( self.av_ca_timestep_scale_multiplier / self.timestep_scale_multiplier ) hidden_dtype = hidden_states.dtype temb_ca_scale_shift, _ = self.av_ca_video_scale_shift_adaln_single( timestep.flatten(), hidden_dtype=hidden_dtype ) temb_ca_scale_shift = temb_ca_scale_shift.view( batch_size, -1, temb_ca_scale_shift.shape[-1] ) temb_ca_gate, _ = self.av_ca_a2v_gate_adaln_single( timestep.flatten() * self.av_ca_timestep_scale_multiplier, hidden_dtype=hidden_dtype, ) temb_ca_gate = temb_ca_gate.view(batch_size, -1, temb_ca_gate.shape[-1]) temb_ca_audio_scale_shift, _ = self.av_ca_audio_scale_shift_adaln_single( audio_timestep.flatten(), hidden_dtype=audio_hidden_states.dtype ) temb_ca_audio_scale_shift = temb_ca_audio_scale_shift.view( batch_size, -1, temb_ca_audio_scale_shift.shape[-1] ) temb_ca_audio_gate, _ = self.av_ca_v2a_gate_adaln_single( audio_timestep.flatten() * self.av_ca_timestep_scale_multiplier, hidden_dtype=audio_hidden_states.dtype, ) temb_ca_audio_gate = temb_ca_audio_gate.view( batch_size, -1, temb_ca_audio_gate.shape[-1] ) # 4. Prepare prompt embeddings encoder_hidden_states = self.caption_projection(encoder_hidden_states) audio_encoder_hidden_states = self.audio_caption_projection( audio_encoder_hidden_states ) # 5. Run blocks for block in self.transformer_blocks: hidden_states, audio_hidden_states = block( hidden_states, audio_hidden_states, encoder_hidden_states, audio_encoder_hidden_states, # Keep the first 4 args positional to stay compatible with cache-dit's # LTX2 adapter, which treats `audio_hidden_states` as `encoder_hidden_states` # under ForwardPattern.Pattern_0. temb=temb, temb_audio=temb_audio, temb_ca_scale_shift=temb_ca_scale_shift, temb_ca_audio_scale_shift=temb_ca_audio_scale_shift, temb_ca_gate=temb_ca_gate, temb_ca_audio_gate=temb_ca_audio_gate, video_rotary_emb=video_rotary_emb, audio_rotary_emb=audio_rotary_emb, ca_video_rotary_emb=ca_video_rotary_emb, ca_audio_rotary_emb=ca_audio_rotary_emb, encoder_attention_mask=encoder_attention_mask, audio_encoder_attention_mask=audio_encoder_attention_mask, ) # 6. Output layers # Video scale_shift_values = self.scale_shift_table[None, None].to( device=hidden_states.device, dtype=hidden_states.dtype ) + embedded_timestep[:, :, None].to(dtype=hidden_states.dtype) shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] with torch.autocast(device_type=hidden_states.device.type, enabled=False): hidden_states = self.norm_out(hidden_states) hidden_states = hidden_states * (1 + scale) + shift hidden_states, _ = self.proj_out(hidden_states) # Audio audio_scale_shift_values = self.audio_scale_shift_table[None, None].to( device=audio_hidden_states.device, dtype=audio_hidden_states.dtype ) + audio_embedded_timestep[:, :, None].to(dtype=audio_hidden_states.dtype) audio_shift, audio_scale = ( audio_scale_shift_values[:, :, 0], audio_scale_shift_values[:, :, 1], ) with torch.autocast(device_type=audio_hidden_states.device.type, enabled=False): audio_hidden_states = self.audio_norm_out(audio_hidden_states) audio_hidden_states = audio_hidden_states * (1 + audio_scale) + audio_shift audio_hidden_states, _ = self.audio_proj_out(audio_hidden_states) # Unpatchify if requested (default True for pipeline compatibility) return_latents = kwargs.get("return_latents", True) if return_latents: # Unpatchify Video # [B, N, C_out_raw*patch_vol] -> [B, C_out_raw, T, H, W] # Requires num_frames, height, width to be known if num_frames is not None and height is not None and width is not None: p_t, p_h, p_w = self.patch_size post_t, post_h, post_w = num_frames // p_t, height // p_h, width // p_w b = batch_size hidden_states = hidden_states.reshape( b, post_t, post_h, post_w, self.out_channels_raw, p_t, p_h, p_w ) hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape( b, self.out_channels_raw, num_frames, height, width ) # Unpatchify Audio # [B, N, C_out] -> [B, C_out, T] (or 4D/5D) if audio_num_frames is not None: b = batch_size # simple reshape for 1D patch audio_hidden_states = audio_hidden_states.permute(0, 2, 1) # [B, C, T] return hidden_states, audio_hidden_states # Backward-compatible alias (older internal name). LTXModel = LTX2VideoTransformer3DModel EntryClass = LTX2VideoTransformer3DModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py ================================================ # Copied and adapted from: mossVG/mova/diffusion/models/wan_audio_dit.py # SPDX-License-Identifier: Apache-2.0 # # NOTE: This module reuses common functions from mova_video_dit.py to reduce code duplication. # Audio-specific functions (precompute_freqs_cis_1d, legacy_precompute_freqs_cis_1d) are kept here. import math from typing import Any, Optional, Tuple import torch import torch.nn as nn from einops import rearrange from torch.distributed.tensor import DTensor from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin # Reuse common functions and classes from mova_video_dit from .mova_video_dit import DiTBlock, precompute_freqs_cis, sinusoidal_embedding_1d # Audio-specific positional encoding functions def legacy_precompute_freqs_cis_1d( dim: int, end: int = 16384, theta: float = 10000.0, base_tps=4.0, target_tps=44100 / 2048, ): s = float(base_tps) / float(target_tps) # 1d rope precompute f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta, s) # No positional encoding is applied to the remaining dimensions no_freqs_cis = precompute_freqs_cis(dim // 3, end, theta, s) no_freqs_cis = torch.ones_like(no_freqs_cis) return f_freqs_cis, no_freqs_cis, no_freqs_cis def precompute_freqs_cis_1d(dim: int, end: int = 16384, theta: float = 10000.0): f_freqs_cis = precompute_freqs_cis(dim, end, theta) return f_freqs_cis.chunk(3, dim=-1) class Head(nn.Module): def __init__( self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float ): super().__init__() self.dim = dim self.patch_size = patch_size self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) self.head = ReplicatedLinear(dim, out_dim * math.prod(patch_size)) self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) def forward(self, x, t_mod): if len(t_mod.shape) == 3: shift, scale = ( self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(2) ).chunk(2, dim=2) x, _ = self.head(self.norm(x) * (1 + scale.squeeze(2)) + shift.squeeze(2)) else: # NOTE: t_mod was originally [B, C]. This works correctly with broadcasting when B=1, but it won't match [1, 2, C] when B > 1. shift, scale = ( self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(1) ).chunk(2, dim=1) x, _ = self.head(self.norm(x) * (1 + scale) + shift) return x class Conv1dLocalIsland(nn.Conv1d): """Inherits from Conv1d and overrides forward. - Parameters remain as DTensors (optimizer consistency is maintained). - In the forward pass, x, weight, and bias are aggregated as Replicate, and then local convolution is performed via to_local. - The output is then redistributed as a DTensor (default is Replicate, placements can be customized). """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def forward(self, input): if isinstance(input, DTensor): x_local = input.to_local() # type: ignore[attr-defined] w_local = self.weight.to_local() # type: ignore[attr-defined] b_local = ( self.bias.to_local() if self.bias is not None else None # type: ignore[attr-defined] ) return self._conv_forward(x_local, w_local, b_local) else: return super().forward(input) class WanAudioModel(CachableDiT, OffloadableDiTMixin): _fsdp_shard_conditions = MOVAAudioConfig()._fsdp_shard_conditions _compile_conditions = MOVAAudioConfig()._compile_conditions _supported_attention_backends = MOVAAudioConfig()._supported_attention_backends param_names_mapping = MOVAAudioConfig().param_names_mapping reverse_param_names_mapping = MOVAAudioConfig().reverse_param_names_mapping lora_param_names_mapping = MOVAAudioConfig().lora_param_names_mapping def __init__( self, config: MOVAAudioConfig, hf_config: dict[str, Any], quant_config: QuantizationConfig | None = None, ) -> None: super().__init__(config=config, hf_config=hf_config) # Extract parameters from config dim = config.dim in_dim = config.in_dim ffn_dim = config.ffn_dim out_dim = config.out_dim text_dim = config.text_dim freq_dim = config.freq_dim eps = config.eps patch_size = config.patch_size num_heads = config.num_heads num_layers = config.num_layers has_image_pos_emb = config.has_image_pos_emb has_ref_conv = config.has_ref_conv separated_timestep = config.separated_timestep require_vae_embedding = config.require_vae_embedding require_clip_embedding = config.require_clip_embedding fuse_vae_embedding_in_latents = config.fuse_vae_embedding_in_latents vae_type = config.vae_type self.dim = dim self.freq_dim = freq_dim self.patch_size = patch_size self.separated_timestep = separated_timestep self.require_vae_embedding = require_vae_embedding self.require_clip_embedding = require_clip_embedding self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents self.vae_type = vae_type # self.patch_embedding = nn.Conv3d( # in_dim, dim, kernel_size=patch_size, stride=patch_size) self.patch_embedding = Conv1dLocalIsland( in_dim, dim, kernel_size=patch_size, stride=patch_size ) self.text_embedding = MLP( text_dim, dim, output_dim=dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, ) self.time_embedding = MLP( freq_dim, dim, output_dim=dim, act_type="silu", quant_config=quant_config ) # Preserve state_dict keys (time_projection.1.weight/bias). self.time_projection = nn.Sequential( nn.SiLU(), ReplicatedLinear(dim, dim * 6, quant_config=quant_config) ) self.blocks = nn.ModuleList( [ DiTBlock(dim, num_heads, ffn_dim, eps, quant_config=quant_config) for _ in range(num_layers) ] ) self.head = Head(dim, out_dim, patch_size, eps) self.num_heads = num_heads self.freqs = None self.img_pos_emb = None if has_ref_conv: self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2)) self.has_image_pos_emb = has_image_pos_emb self.has_ref_conv = has_ref_conv self.hidden_size = dim self.num_attention_heads = num_heads self.num_channels_latents = out_dim self.layer_names = ["blocks"] self.cnt = 0 self.teacache_thresh = 0 self.coefficients = [] self.accumulated_rel_l1_distance = 0 self.previous_modulated_input = None self.previous_resiual = None self.previous_e0_even = None self.previous_e0_odd = None self.previous_residual_even = None self.previous_residual_odd = None self.is_even = False self.should_calc_even = True self.should_calc_odd = True self.accumulated_rel_l1_distance_even = 0 self.accumulated_rel_l1_distance_odd = 0 self.__post_init__() def _init_freqs(self): if self.freqs is not None: return head_dim = self.dim // self.num_heads if self.vae_type == "dac": self.freqs = precompute_freqs_cis_1d(head_dim) else: raise ValueError(f"Invalid VAE type: {self.vae_type}") def patchify( self, x: torch.Tensor, control_camera_latents_input: Optional[torch.Tensor] = None, ): x = self.patch_embedding(x) grid_size = x.shape[2:] x = rearrange(x, "b c f -> b f c").contiguous() return x, grid_size # x, grid_size: (f) def unpatchify(self, x: torch.Tensor, grid_size: tuple[int]): return rearrange( x, "b f (p c) -> b c (f p)", f=grid_size[0], p=self.patch_size[0] ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, ) -> torch.Tensor: # MOVA audio uses x/context naming historically. x = hidden_states context = ( encoder_hidden_states[0] if isinstance(encoder_hidden_states, list) else encoder_hidden_states ) t = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep)) t_proj, _ = self.time_projection(t) t_mod = t_proj.unflatten(1, (6, self.dim)) context = self.text_embedding(context) x, (f,) = self.patchify(x) freqs = ( torch.cat( [ self.freqs[0][:f].view(f, -1).expand(f, -1), self.freqs[1][:f].view(f, -1).expand(f, -1), self.freqs[2][:f].view(f, -1).expand(f, -1), ], dim=-1, ) .reshape(f, 1, -1) .to(x.device) ) for block in self.blocks: x = block(x, context, t_mod, freqs) x = self.head(x, t) x = self.unpatchify(x, (f,)) return x EntryClass = WanAudioModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py ================================================ # Copied and adapted from: mossVG/mova/diffusion/models/wan_video_dit.py # SPDX-License-Identifier: Apache-2.0 # # NOTE: This module shares common functions (sinusoidal_embedding_1d, precompute_freqs_cis, etc.) # with wanvideo.py. These functions are kept here for MOVA-specific model architecture, # but could be refactored to a common module in the future. import math from typing import Any, Tuple import torch import torch.nn as nn from einops import rearrange from torch.distributed.tensor import DTensor from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig from sglang.multimodal_gen.runtime.distributed import get_tp_world_size from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention # Reuse SGLang's optimized RMSNorm instead of torch.nn.RMSNorm or custom SlowRMSNorm from sglang.multimodal_gen.runtime.layers.layernorm import ( LayerNormScaleShift, RMSNorm, ScaleResidualLayerNormScaleShift, tensor_parallel_rms_norm, ) from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, ReplicatedLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # @torch.compile(fullgraph=True) def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor): return x * (1 + scale) + shift def sinusoidal_embedding_1d(dim, position): sinusoid = torch.outer( position.type(torch.float64), torch.pow( 10000, -torch.arange(dim // 2, dtype=torch.float64, device=position.device).div( dim // 2 ), ), ) x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) return x.to(position.dtype) def precompute_freqs_cis_3d(dim: int, end: int = 1024, theta: float = 10000.0): # 3d rope precompute f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta) h_freqs_cis = precompute_freqs_cis(dim // 3, end, theta) w_freqs_cis = precompute_freqs_cis(dim // 3, end, theta) return f_freqs_cis, h_freqs_cis, w_freqs_cis def precompute_freqs_cis( dim: int, end: int = 1024, theta: float = 10000.0, s: float = 1.0 ): # 1d rope precompute # Note: s parameter is used for audio-specific scaling (e.g., tps adjustment) freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].double() / dim)) pos = torch.arange(end, dtype=torch.float64, device=freqs.device) * s freqs = torch.outer(pos, freqs) freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 return freqs_cis def rope_apply(x, freqs, num_heads): x = rearrange(x, "b s (n d) -> b s n d", n=num_heads) x_out = torch.view_as_complex( x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2) ) x_out = torch.view_as_real(x_out * freqs).flatten(2) return x_out.to(x.dtype) def rope_apply_head_dim(x, freqs, head_dim): x = rearrange(x, "b s (n d) -> b s n d", d=head_dim) x_out = torch.view_as_complex( x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2) ) # print(f"{x_out.shape = }, {freqs.shape = }") x_out = torch.view_as_real(x_out * freqs).flatten(2) return x_out.to(x.dtype) class SelfAttention(nn.Module): """ Self-Attention module for MOVA DiT with Sequence Parallelism support. SP is handled at the pipeline level (latents are pre-sharded before DiT forward). USPAttention internally handles the all-to-all communication for distributed attention. Input x should already be the local shard [B, S_local, D] when SP is enabled. """ def __init__( self, dim: int, num_heads: int, eps: float = 1e-6, quant_config: QuantizationConfig | None = None, ): super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads self.tp_size = get_tp_world_size() if self.num_heads % self.tp_size != 0: raise ValueError( f"num_heads ({self.num_heads}) must be divisible by tp_size ({self.tp_size})." ) self.num_heads_per_rank = self.num_heads // self.tp_size # TP strategy: shard Q/K/V over heads (column-parallel), then row-parallel output. self.q = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config ) self.k = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config ) self.v = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config ) self.o = RowParallelLinear( dim, dim, bias=True, input_is_parallel=True, quant_config=quant_config ) self.norm_q = RMSNorm(dim, eps=eps) self.norm_k = RMSNorm(dim, eps=eps) self.attn = USPAttention( # Local heads per TP rank. num_heads=self.num_heads_per_rank, head_size=self.head_dim, causal=False, softmax_scale=None, ) def forward(self, x, freqs): """ Forward pass for self-attention. Args: x: Input tensor [B, S_local, D] - already sharded by SP when SP > 1 freqs: RoPE frequencies [S_local, 1, head_dim] - should match x's sequence length Returns: Output tensor [B, S_local, D] """ if isinstance(freqs, DTensor): freqs = freqs.to_local() # Compute Q, K, V on local sequence q, _ = self.q(x) k, _ = self.k(x) v, _ = self.v(x) # RMSNorm over sharded hidden dimension. if self.tp_size > 1: q = tensor_parallel_rms_norm(q, self.norm_q) k = tensor_parallel_rms_norm(k, self.norm_k) else: q = self.norm_q(q) k = self.norm_k(k) # Apply RoPE q = rope_apply_head_dim(q, freqs, self.head_dim) k = rope_apply_head_dim(k, freqs, self.head_dim) # USPAttention expects [B, S_local, H, D] format q = rearrange(q, "b s (n d) -> b s n d", n=self.num_heads_per_rank) k = rearrange(k, "b s (n d) -> b s n d", n=self.num_heads_per_rank) v = rearrange(v, "b s (n d) -> b s n d", n=self.num_heads_per_rank) # USPAttention handles SP communication internally out = self.attn(q, k, v) out = rearrange(out, "b s n d -> b s (n d)") out, _ = self.o(out) return out class CrossAttention(nn.Module): """ Cross-Attention module for MOVA DiT. Cross-attention does NOT require SP communication because: - Query comes from the main sequence (already sharded by SP) - Key/Value come from context (text embeddings, which are replicated across all ranks) Uses LocalAttention instead of USPAttention for efficiency. """ def __init__( self, dim: int, num_heads: int, eps: float = 1e-6, quant_config: QuantizationConfig | None = None, ): super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads self.tp_size = get_tp_world_size() if self.num_heads % self.tp_size != 0: raise ValueError( f"num_heads ({self.num_heads}) must be divisible by tp_size ({self.tp_size})." ) self.num_heads_per_rank = self.num_heads // self.tp_size self.q = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config ) self.k = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config ) self.v = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config ) self.o = RowParallelLinear( dim, dim, bias=True, input_is_parallel=True, quant_config=quant_config ) self.norm_q = RMSNorm(dim, eps=eps) self.norm_k = RMSNorm(dim, eps=eps) # Use LocalAttention for cross-attention (no SP communication needed) self.attn = LocalAttention( num_heads=self.num_heads_per_rank, head_size=self.head_dim, causal=False, softmax_scale=None, ) def forward(self, x: torch.Tensor, y: torch.Tensor): """ Forward pass for cross-attention. Args: x: Query tensor [B, S_local, D] - the main sequence (sharded by SP) y: Context tensor [B, S_ctx, D] - text/image embeddings (replicated) Returns: Output tensor [B, S_local, D] """ ctx = y q, _ = self.q(x) k, _ = self.k(ctx) v, _ = self.v(ctx) if self.tp_size > 1: q = tensor_parallel_rms_norm(q, self.norm_q) k = tensor_parallel_rms_norm(k, self.norm_k) else: q = self.norm_q(q) k = self.norm_k(k) q = rearrange(q, "b s (n d) -> b s n d", n=self.num_heads_per_rank) k = rearrange(k, "b s (n d) -> b s n d", n=self.num_heads_per_rank) v = rearrange(v, "b s (n d) -> b s n d", n=self.num_heads_per_rank) x = self.attn(q, k, v) x = rearrange(x, "b s n d -> b s (n d)") x, _ = self.o(x) return x class MulAdd(nn.Module): def __init__(self): super().__init__() def forward(self, x, gate, residual): return residual + gate * x class DiTBlock(nn.Module): def __init__( self, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6, quant_config: QuantizationConfig | None = None, ): super().__init__() self.dim = dim self.num_heads = num_heads self.ffn_dim = ffn_dim self.self_attn = SelfAttention(dim, num_heads, eps, quant_config=quant_config) self.cross_attn = CrossAttention(dim, num_heads, eps, quant_config=quant_config) self.norm1 = LayerNormScaleShift( dim, eps=eps, elementwise_affine=False, dtype=torch.float32 ) self.self_attn_norm = nn.LayerNorm(dim, eps=eps) # Fused: residual + 1 * cross_attn_out → layernorm + scale/shift # Replaces the old norm2 (LayerNormScaleShift) + residual add for cross-attention self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, eps=eps, elementwise_affine=False, dtype=torch.float32 ) self.ffn = MLP( dim, ffn_dim, output_dim=dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, ) self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) self.mlp_residual = MulAdd() def forward(self, x, context, t_mod, freqs): has_seq = len(t_mod.shape) == 4 chunk_dim = 2 if has_seq else 1 # msa: multi-head self-attention mlp: multi-layer perceptron shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod ).chunk(6, dim=chunk_dim) if has_seq: shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2), shift_mlp.squeeze(2), scale_mlp.squeeze(2), gate_mlp.squeeze(2), ) orig_dtype = x.dtype # 1. Self-attention, fuse: # - layernorm(x) * (1 + scale_msa) + shift_msa input_x = self.norm1(x, shift_msa, scale_msa) # 2. torch.compile may fuse mlp_residual and self_attn_norm x = self.mlp_residual(self.self_attn(input_x, freqs), gate_msa, x) norm_x = self.self_attn_norm(x) # 3. Cross-attention, fuse: # - x = x + 1 * cross_output # - input_x = layernorm(x) * (1 + scale_mlp) + shift_mlp cross_output = self.cross_attn(norm_x, context) input_x, x = self.cross_attn_residual_norm( x, cross_output, 1, shift_mlp, scale_mlp ) # 4. Feed-forward x = self.mlp_residual(self.ffn(input_x), gate_mlp, x) x = x.to(orig_dtype) return x class Head(nn.Module): def __init__( self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float ): super().__init__() self.dim = dim self.patch_size = patch_size self.norm = LayerNormScaleShift( dim, eps=eps, elementwise_affine=False, dtype=torch.float32 ) # Output dim is small for MOVA; replicate to avoid TP shape coupling. self.head = ReplicatedLinear(dim, out_dim * math.prod(patch_size)) self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) def forward(self, x, t_mod): if len(t_mod.shape) == 3: shift, scale = ( self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(2) ).chunk(2, dim=2) x, _ = self.head(self.norm(x, shift.squeeze(2), scale.squeeze(2))) else: shift, scale = ( self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod ).chunk(2, dim=1) x, _ = self.head(self.norm(x, shift, scale)) return x class Conv3dLocalIsland(nn.Conv3d): """ Inherits from Conv3d and overrides the forward method. Key behaviors: - Parameters are kept as DTensor to maintain optimizer consistency. - The forward pass aggregates input, weight, and bias into a Replicate state, then performs the convolution locally using to_local(). - The output is then redistributed as a DTensor (defaults to Replicate, but placements can be customized). """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def forward(self, input): if isinstance(input, DTensor): # NOTE: DTensor typing stubs are incomplete; at runtime DTensor has # to_local() and parameters may also be DTensor. x_local = input.to_local() # type: ignore[attr-defined] w_local = self.weight.to_local() # type: ignore[attr-defined] b_local = ( self.bias.to_local() if self.bias is not None else None # type: ignore[attr-defined] ) return self._conv_forward(x_local, w_local, b_local) else: return super().forward(input) class WanModel(CachableDiT, OffloadableDiTMixin): _fsdp_shard_conditions = MOVAVideoConfig()._fsdp_shard_conditions _compile_conditions = MOVAVideoConfig()._compile_conditions _supported_attention_backends = MOVAVideoConfig()._supported_attention_backends param_names_mapping = MOVAVideoConfig().param_names_mapping reverse_param_names_mapping = MOVAVideoConfig().reverse_param_names_mapping lora_param_names_mapping = MOVAVideoConfig().lora_param_names_mapping def __init__( self, config: MOVAVideoConfig, hf_config: dict[str, Any], quant_config: QuantizationConfig | None = None, ) -> None: super().__init__(config=config, hf_config=hf_config) # Extract parameters from config dim = config.dim in_dim = config.in_dim ffn_dim = config.ffn_dim out_dim = config.out_dim text_dim = config.text_dim freq_dim = config.freq_dim eps = config.eps patch_size = config.patch_size num_heads = config.num_heads num_layers = config.num_layers has_image_pos_emb = config.has_image_pos_emb has_ref_conv = config.has_ref_conv separated_timestep = config.separated_timestep require_vae_embedding = config.require_vae_embedding require_clip_embedding = config.require_clip_embedding fuse_vae_embedding_in_latents = config.fuse_vae_embedding_in_latents self.dim = dim self.freq_dim = freq_dim self.patch_size = patch_size self.separated_timestep = separated_timestep self.require_vae_embedding = require_vae_embedding self.require_clip_embedding = require_clip_embedding self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents self.patch_embedding = Conv3dLocalIsland( in_dim, dim, kernel_size=patch_size, stride=patch_size ) self.text_embedding = MLP( text_dim, dim, output_dim=dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, ) self.time_embedding = MLP( freq_dim, dim, output_dim=dim, act_type="silu", quant_config=quant_config ) # Preserve state_dict keys (time_projection.1.weight/bias). self.time_projection = nn.Sequential( nn.SiLU(), ReplicatedLinear(dim, dim * 6, quant_config=quant_config) ) self.blocks = nn.ModuleList( [ DiTBlock(dim, num_heads, ffn_dim, eps, quant_config=quant_config) for _ in range(num_layers) ] ) self.head = Head(dim, out_dim, patch_size, eps) self.num_heads = num_heads self.freqs = None if has_ref_conv: self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2)) self.has_image_pos_emb = has_image_pos_emb self.has_ref_conv = has_ref_conv self.hidden_size = dim self.num_attention_heads = num_heads self.num_channels_latents = out_dim self.layer_names = ["blocks"] self.cnt = 0 self.teacache_thresh = 0 self.coefficients = [] self.accumulated_rel_l1_distance = 0 self.previous_modulated_input = None self.previous_resiual = None self.previous_e0_even = None self.previous_e0_odd = None self.previous_residual_even = None self.previous_residual_odd = None self.is_even = False self.should_calc_even = True self.should_calc_odd = True self.accumulated_rel_l1_distance_even = 0 self.accumulated_rel_l1_distance_odd = 0 self.__post_init__() def _init_freqs(self): if self.freqs is not None: return head_dim = self.dim // self.num_heads self.freqs = precompute_freqs_cis_3d(head_dim) def patchify( self, x: torch.Tensor, control_camera_latents_input: torch.Tensor | None = None ): # NOTE(dhyu): avoid slow_conv x = x.contiguous(memory_format=torch.channels_last_3d) x = self.patch_embedding(x) grid_size = x.shape[2:] x = rearrange(x, "b c f h w -> b (f h w) c").contiguous() return x, grid_size # x, grid_size: (f, h, w) def unpatchify(self, x: torch.Tensor, grid_size: tuple[int, int, int]): return rearrange( x, "b (f h w) (x y z c) -> b c (f x) (h y) (w z)", f=grid_size[0], h=grid_size[1], w=grid_size[2], x=self.patch_size[0], y=self.patch_size[1], z=self.patch_size[2], ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, ) -> torch.Tensor: # MOVA code historically uses x/context/y/clip_feature naming. x = hidden_states context = ( encoder_hidden_states[0] if isinstance(encoder_hidden_states, list) else encoder_hidden_states ) t = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep)) t_proj, _ = self.time_projection(t) t_mod = t_proj.unflatten(1, (6, self.dim)) context = self.text_embedding(context) x, (f, h, w) = self.patchify(x) freqs = ( torch.cat( [ self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), ], dim=-1, ) .reshape(f * h * w, 1, -1) .to(x.device) ) for block in self.blocks: x = block(x, context, t_mod, freqs) x = self.head(x, t) x = self.unpatchify(x, (f, h, w)) return x EntryClass = WanModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import functools from typing import Any, Dict, List, Optional, Tuple, Union import diffusers import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from diffusers.models.attention import FeedForward from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import AdaLayerNormContinuous from sglang.jit_kernel.diffusion.triton.scale_shift import ( fuse_layernorm_scale_shift_gate_select01_kernel, fuse_residual_layernorm_scale_shift_gate_select01_kernel, ) from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd from sglang.multimodal_gen.runtime.layers.layernorm import ( LayerNormScaleShift, RMSNorm, ScaleResidualLayerNormScaleShift, apply_qk_norm, ) from sglang.multimodal_gen.runtime.layers.linear import ( MergedColumnParallelLinear, ReplicatedLinear, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( NunchakuConfig, is_nunchaku_available, ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( apply_flashinfer_rope_qk_inplace, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name try: from nunchaku.models.attention import NunchakuFeedForward # type: ignore[import] except Exception: NunchakuFeedForward = None def _get_qkv_projections( attn: "QwenImageCrossAttention", hidden_states, encoder_hidden_states=None ): if attn.use_fused_qkv: img_qkv, _ = attn.to_qkv(hidden_states) img_query, img_key, img_value = [ x.contiguous() for x in img_qkv.chunk(3, dim=-1) ] else: img_query, _ = attn.to_q(hidden_states) img_key, _ = attn.to_k(hidden_states) img_value, _ = attn.to_v(hidden_states) txt_query = txt_key = txt_value = None if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None: if attn.use_fused_added_qkv: txt_qkv, _ = attn.to_added_qkv(encoder_hidden_states) txt_query, txt_key, txt_value = [ x.contiguous() for x in txt_qkv.chunk(3, dim=-1) ] else: txt_query, _ = attn.add_q_proj(encoder_hidden_states) txt_key, _ = attn.add_k_proj(encoder_hidden_states) txt_value, _ = attn.add_v_proj(encoder_hidden_states) return img_query, img_key, img_value, txt_query, txt_key, txt_value class QwenTimestepProjEmbeddings(nn.Module): def __init__(self, embedding_dim, use_additional_t_cond=False): super().__init__() self.time_proj = Timesteps( num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000 ) self.timestep_embedder = TimestepEmbedding( in_channels=256, time_embed_dim=embedding_dim ) self.use_additional_t_cond = use_additional_t_cond if use_additional_t_cond: self.addition_t_embedding = nn.Embedding(2, embedding_dim) def forward(self, timestep, hidden_states, addition_t_cond=None): timesteps_proj = self.time_proj(timestep) timesteps_emb = self.timestep_embedder( timesteps_proj.to(dtype=hidden_states.dtype) ) # (N, D) conditioning = timesteps_emb if self.use_additional_t_cond: if addition_t_cond is None: raise ValueError( "When additional_t_cond is True, addition_t_cond must be provided." ) addition_t_emb = self.addition_t_embedding(addition_t_cond) addition_t_emb = addition_t_emb.to(dtype=hidden_states.dtype) conditioning = conditioning + addition_t_emb return conditioning class QwenEmbedRope(nn.Module): def __init__(self, theta: int, axes_dim: List[int], scale_rope=False): super().__init__() self.theta = theta self.axes_dim = axes_dim pos_index = torch.arange(4096) neg_index = torch.arange(4096).flip(0) * -1 - 1 self.pos_freqs = torch.cat( [ self.rope_params(pos_index, self.axes_dim[0], self.theta), self.rope_params(pos_index, self.axes_dim[1], self.theta), self.rope_params(pos_index, self.axes_dim[2], self.theta), ], dim=1, ) self.neg_freqs = torch.cat( [ self.rope_params(neg_index, self.axes_dim[0], self.theta), self.rope_params(neg_index, self.axes_dim[1], self.theta), self.rope_params(neg_index, self.axes_dim[2], self.theta), ], dim=1, ) # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART self.scale_rope = scale_rope def rope_params(self, index, dim, theta=10000): """ Args: index: [0, 1, 2, 3] 1D Tensor representing the position index of the token """ device = index.device assert dim % 2 == 0 freqs = torch.outer( index, ( 1.0 / torch.pow( theta, torch.arange(0, dim, 2, device=device).to(torch.float32).div(dim), ) ).to(device=device), ) freqs = torch.polar(torch.ones_like(freqs), freqs) return freqs def forward( self, video_fhw: Union[Tuple[int, int, int], List[Tuple[int, int, int]]], txt_seq_lens: List[int], device: torch.device, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Args: video_fhw (`Tuple[int, int, int]` or `List[Tuple[int, int, int]]`): A list of 3 integers [frame, height, width] representing the shape of the video. txt_seq_lens (`List[int]`): A list of integers of length batch_size representing the length of each text prompt. device: (`torch.device`): The device on which to perform the RoPE computation. """ # When models are initialized under a "meta" device context (e.g. init_empty_weights), # tensors created during __init__ become meta tensors. Calling .to(...) on a meta tensor # raises "Cannot copy out of meta tensor". Rebuild the frequencies on the target device # in that case; otherwise move them if just on a different device. if getattr(self.pos_freqs, "device", torch.device("meta")).type == "meta": pos_index = torch.arange(4096, device=device) neg_index = torch.arange(4096, device=device).flip(0) * -1 - 1 self.pos_freqs = torch.cat( [ self.rope_params(pos_index, self.axes_dim[0], self.theta), self.rope_params(pos_index, self.axes_dim[1], self.theta), self.rope_params(pos_index, self.axes_dim[2], self.theta), ], dim=1, ).to(device=device) self.neg_freqs = torch.cat( [ self.rope_params(neg_index, self.axes_dim[0], self.theta), self.rope_params(neg_index, self.axes_dim[1], self.theta), self.rope_params(neg_index, self.axes_dim[2], self.theta), ], dim=1, ).to(device=device) elif self.pos_freqs.device != device: self.pos_freqs = self.pos_freqs.to(device) self.neg_freqs = self.neg_freqs.to(device) if isinstance(video_fhw, list): video_fhw = video_fhw[0] if not isinstance(video_fhw, list): video_fhw = [video_fhw] vid_freqs = [] max_vid_index = 0 for idx, fhw in enumerate(video_fhw): frame, height, width = fhw # RoPE frequencies are cached via a lru_cache decorator on _compute_video_freqs video_freq = self._compute_video_freqs(frame, height, width, idx) video_freq = video_freq.to(device) vid_freqs.append(video_freq) if self.scale_rope: max_vid_index = max(height // 2, width // 2, max_vid_index) else: max_vid_index = max(height, width, max_vid_index) max_len = max(txt_seq_lens) txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...] vid_freqs = torch.cat(vid_freqs, dim=0).to(device=device) return vid_freqs, txt_freqs @functools.lru_cache(maxsize=128) def _compute_video_freqs( self, frame: int, height: int, width: int, idx: int = 0 ) -> torch.Tensor: seq_lens = frame * height * width freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) freqs_frame = ( freqs_pos[0][idx : idx + frame] .view(frame, 1, 1, -1) .expand(frame, height, width, -1) ) if self.scale_rope: freqs_height = torch.cat( [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], dim=0, ) freqs_height = freqs_height.view(1, height, 1, -1).expand( frame, height, width, -1 ) freqs_width = torch.cat( [freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], dim=0, ) freqs_width = freqs_width.view(1, 1, width, -1).expand( frame, height, width, -1 ) else: freqs_height = ( freqs_pos[1][:height] .view(1, height, 1, -1) .expand(frame, height, width, -1) ) freqs_width = ( freqs_pos[2][:width] .view(1, 1, width, -1) .expand(frame, height, width, -1) ) freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape( seq_lens, -1 ) return freqs.clone().contiguous() class QwenEmbedLayer3DRope(nn.Module): def __init__(self, theta: int, axes_dim: List[int], scale_rope=False): super().__init__() self.theta = theta self.axes_dim = axes_dim pos_index = torch.arange(4096) neg_index = torch.arange(4096).flip(0) * -1 - 1 self.pos_freqs = torch.cat( [ self.rope_params(pos_index, self.axes_dim[0], self.theta), self.rope_params(pos_index, self.axes_dim[1], self.theta), self.rope_params(pos_index, self.axes_dim[2], self.theta), ], dim=1, ) self.neg_freqs = torch.cat( [ self.rope_params(neg_index, self.axes_dim[0], self.theta), self.rope_params(neg_index, self.axes_dim[1], self.theta), self.rope_params(neg_index, self.axes_dim[2], self.theta), ], dim=1, ) self.scale_rope = scale_rope def rope_params(self, index, dim, theta=10000): """ Args: index: [0, 1, 2, 3] 1D Tensor representing the position index of the token """ device = index.device assert dim % 2 == 0 freqs = torch.outer( index, ( 1.0 / torch.pow( theta, torch.arange(0, dim, 2, device=device).to(torch.float32).div(dim), ) ).to(device=device), ) freqs = torch.polar(torch.ones_like(freqs), freqs) return freqs def forward(self, video_fhw, txt_seq_lens, device): """ Args: video_fhw: [frame, height, width] a list of 3 integers representing the shape of the video Args: txt_length: [bs] a list of 1 integers representing the length of the text """ # When models are initialized under a "meta" device context (e.g. init_empty_weights), # tensors created during __init__ become meta tensors. Calling .to(...) on a meta tensor # raises "Cannot copy out of meta tensor". Rebuild the frequencies on the target device # in that case; otherwise move them if just on a different device. if getattr(self.pos_freqs, "device", torch.device("meta")).type == "meta": pos_index = torch.arange(4096, device=device) neg_index = torch.arange(4096, device=device).flip(0) * -1 - 1 self.pos_freqs = torch.cat( [ self.rope_params(pos_index, self.axes_dim[0], self.theta), self.rope_params(pos_index, self.axes_dim[1], self.theta), self.rope_params(pos_index, self.axes_dim[2], self.theta), ], dim=1, ).to(device=device) self.neg_freqs = torch.cat( [ self.rope_params(neg_index, self.axes_dim[0], self.theta), self.rope_params(neg_index, self.axes_dim[1], self.theta), self.rope_params(neg_index, self.axes_dim[2], self.theta), ], dim=1, ).to(device=device) elif self.pos_freqs.device != device: self.pos_freqs = self.pos_freqs.to(device) self.neg_freqs = self.neg_freqs.to(device) if isinstance(video_fhw, list): video_fhw = video_fhw[0] if not isinstance(video_fhw, list): video_fhw = [video_fhw] vid_freqs = [] max_vid_index = 0 layer_num = len(video_fhw) - 1 for idx, fhw in enumerate(video_fhw): frame, height, width = fhw if idx != layer_num: video_freq = self._compute_video_freqs(frame, height, width, idx) else: # For the condition image, we set the layer index to -1 video_freq = self._compute_condition_freqs(frame, height, width) video_freq = video_freq.to(device) vid_freqs.append(video_freq) if self.scale_rope: max_vid_index = max(height // 2, width // 2, max_vid_index) else: max_vid_index = max(height, width, max_vid_index) max_vid_index = max(max_vid_index, layer_num) max_len = max(txt_seq_lens) txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...] vid_freqs = torch.cat(vid_freqs, dim=0) return vid_freqs, txt_freqs @functools.lru_cache(maxsize=None) def _compute_video_freqs(self, frame, height, width, idx=0): seq_lens = frame * height * width freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) freqs_frame = ( freqs_pos[0][idx : idx + frame] .view(frame, 1, 1, -1) .expand(frame, height, width, -1) ) if self.scale_rope: freqs_height = torch.cat( [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], dim=0, ) freqs_height = freqs_height.view(1, height, 1, -1).expand( frame, height, width, -1 ) freqs_width = torch.cat( [freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], dim=0, ) freqs_width = freqs_width.view(1, 1, width, -1).expand( frame, height, width, -1 ) else: freqs_height = ( freqs_pos[1][:height] .view(1, height, 1, -1) .expand(frame, height, width, -1) ) freqs_width = ( freqs_pos[2][:width] .view(1, 1, width, -1) .expand(frame, height, width, -1) ) freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape( seq_lens, -1 ) return freqs.clone().contiguous() @functools.lru_cache(maxsize=None) def _compute_condition_freqs(self, frame, height, width): seq_lens = frame * height * width freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) freqs_frame = ( freqs_neg[0][-1:].view(frame, 1, 1, -1).expand(frame, height, width, -1) ) if self.scale_rope: freqs_height = torch.cat( [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], dim=0, ) freqs_height = freqs_height.view(1, height, 1, -1).expand( frame, height, width, -1 ) freqs_width = torch.cat( [freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], dim=0, ) freqs_width = freqs_width.view(1, 1, width, -1).expand( frame, height, width, -1 ) else: freqs_height = ( freqs_pos[1][:height] .view(1, height, 1, -1) .expand(frame, height, width, -1) ) freqs_width = ( freqs_pos[2][:width] .view(1, 1, width, -1) .expand(frame, height, width, -1) ) freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape( seq_lens, -1 ) return freqs.clone().contiguous() class QwenImageCrossAttention(nn.Module): def __init__( self, dim: int, # query_dim num_heads: int, head_dim: int, window_size=(-1, -1), added_kv_proj_dim: int = None, out_bias: bool = True, qk_norm=True, # rmsnorm eps=1e-6, pre_only=False, context_pre_only: bool = False, parallel_attention=False, out_dim: int = None, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: assert dim % num_heads == 0 super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads self.window_size = window_size self.qk_norm = qk_norm self.eps = eps self.parallel_attention = parallel_attention self.added_kv_proj_dim = added_kv_proj_dim self.prefix = prefix self.use_fused_qkv = isinstance(quant_config, NunchakuConfig) self.inner_dim = out_dim if out_dim is not None else head_dim * num_heads self.inner_kv_dim = self.inner_dim if self.use_fused_qkv: # Use fused QKV projection for nunchaku quantization self.to_qkv = MergedColumnParallelLinear( dim, [self.inner_dim] * 3, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_qkv", ) else: # Use separate Q/K/V projections for non-quantized models self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True) self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True) self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True) if self.qk_norm: self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() if added_kv_proj_dim is not None: self.use_fused_added_qkv = isinstance(quant_config, NunchakuConfig) if self.use_fused_added_qkv: self.to_added_qkv = MergedColumnParallelLinear( added_kv_proj_dim, [self.inner_dim] * 3, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_added_qkv", ) else: # Use separate Q/K/V projections for non-quantized models self.add_q_proj = ReplicatedLinear( added_kv_proj_dim, self.inner_dim, bias=True ) self.add_k_proj = ReplicatedLinear( added_kv_proj_dim, self.inner_dim, bias=True ) self.add_v_proj = ReplicatedLinear( added_kv_proj_dim, self.inner_dim, bias=True ) if context_pre_only is not None and not context_pre_only: self.to_add_out = ReplicatedLinear(self.inner_dim, self.dim, bias=out_bias) else: self.to_add_out = None if not pre_only: self.to_out = nn.ModuleList([]) self.to_out.append( ReplicatedLinear(self.inner_dim, self.dim, bias=out_bias) ) else: self.to_out = None self.norm_added_q = RMSNorm(head_dim, eps=eps) self.norm_added_k = RMSNorm(head_dim, eps=eps) # Scaled dot product attention self.attn = USPAttention( num_heads=num_heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, causal=False, supported_attention_backends={ AttentionBackendEnum.FA, AttentionBackendEnum.AITER, AttentionBackendEnum.AITER_SAGE, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.SAGE_ATTN_3, }, ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, image_rotary_emb: tuple[torch.Tensor, torch.Tensor], **cross_attention_kwargs, ): seq_len_txt = encoder_hidden_states.shape[1] img_query, img_key, img_value, txt_query, txt_key, txt_value = ( _get_qkv_projections(self, hidden_states, encoder_hidden_states) ) # Reshape for multi-head attention img_query = img_query.unflatten(-1, (self.num_heads, -1)) img_key = img_key.unflatten(-1, (self.num_heads, -1)) img_value = img_value.unflatten(-1, (self.num_heads, -1)) txt_query = txt_query.unflatten(-1, (self.num_heads, -1)) txt_key = txt_key.unflatten(-1, (self.num_heads, -1)) txt_value = txt_value.unflatten(-1, (self.num_heads, -1)) # Apply QK normalization if self.qk_norm: img_query, img_key = apply_qk_norm( q=img_query, k=img_key, q_norm=self.norm_q, k_norm=self.norm_k, head_dim=img_query.shape[-1], allow_inplace=True, ) txt_query, txt_key = apply_qk_norm( q=txt_query, k=txt_key, q_norm=self.norm_added_q, k_norm=self.norm_added_k, head_dim=txt_query.shape[-1], allow_inplace=True, ) # Apply RoPE if image_rotary_emb is not None: if not ( isinstance(image_rotary_emb[0], torch.Tensor) and image_rotary_emb[0].dim() == 2 ): raise RuntimeError("image_rotary_emb must be cos_sin_cache tensors") img_cache, txt_cache = image_rotary_emb img_query, img_key = apply_flashinfer_rope_qk_inplace( img_query, img_key, img_cache, is_neox=False ) txt_query, txt_key = apply_flashinfer_rope_qk_inplace( txt_query, txt_key, txt_cache, is_neox=False ) # Concatenate for joint attention # Order: [text, image] joint_query = torch.cat([txt_query, img_query], dim=1) joint_key = torch.cat([txt_key, img_key], dim=1) joint_value = torch.cat([txt_value, img_value], dim=1) # Compute joint attention joint_hidden_states = self.attn( joint_query, joint_key, joint_value, ) # Reshape back joint_hidden_states = joint_hidden_states.flatten(2, 3) joint_hidden_states = joint_hidden_states.to(joint_query.dtype) # Split attention outputs back txt_attn_output = joint_hidden_states[:, :seq_len_txt, :] # Text part img_attn_output = joint_hidden_states[:, seq_len_txt:, :] # Image part # Apply output projections img_attn_output, _ = self.to_out[0](img_attn_output) if len(self.to_out) > 1: (img_attn_output,) = self.to_out[1](img_attn_output) # dropout txt_attn_output, _ = self.to_add_out(txt_attn_output) return img_attn_output, txt_attn_output class QwenImageTransformerBlock(nn.Module): def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6, quant_config: Optional[QuantizationConfig] | NunchakuConfig = None, prefix: str = "", zero_cond_t: bool = False, ): super().__init__() self.prefix = prefix self.dim = dim self.num_attention_heads = num_attention_heads self.attention_head_dim = attention_head_dim self.quant_config = quant_config self.zero_cond_t = zero_cond_t mod_quant_config = ( quant_config if (quant_config is not None and quant_config.get_name() == "svdquant") else None ) # Image processing modules self.img_mod = nn.Sequential( nn.SiLU(), nn.Linear( dim, 6 * dim, bias=True ), # For scale, shift, gate for norm1 and norm2 ) self.img_norm1 = LayerNormScaleShift( hidden_size=dim, eps=eps, elementwise_affine=False ) self.attn = QwenImageCrossAttention( dim=dim, num_heads=num_attention_heads, added_kv_proj_dim=dim, context_pre_only=False, head_dim=attention_head_dim, quant_config=quant_config, prefix=f"{prefix}.attn", ) self.img_norm2 = ScaleResidualLayerNormScaleShift( dim, eps=eps, elementwise_affine=False ) # Text processing modules self.txt_mod = nn.Sequential( nn.SiLU(), nn.Linear( dim, 6 * dim, bias=True ), # For scale, shift, gate for norm1 and norm2 ) self.txt_norm1 = LayerNormScaleShift( hidden_size=dim, eps=eps, elementwise_affine=False ) # Text doesn't need separate attention - it's handled by img_attn joint computation self.txt_norm2 = ScaleResidualLayerNormScaleShift( hidden_size=dim, eps=eps, elementwise_affine=False ) # Utils self.fuse_mul_add = MulAdd() nunchaku_enabled = ( quant_config is not None and hasattr(quant_config, "get_name") and quant_config.get_name() == "svdquant" and is_nunchaku_available() ) if nunchaku_enabled: ff_class = diffusers.models.attention.FeedForward self.img_mlp = ff_class( dim=dim, dim_out=dim, activation_fn="gelu-approximate", ) self.txt_mlp = ff_class( dim=dim, dim_out=dim, activation_fn="gelu-approximate", ) else: self.img_mlp = FeedForward( dim=dim, dim_out=dim, activation_fn="gelu-approximate", ) self.txt_mlp = FeedForward( dim=dim, dim_out=dim, activation_fn="gelu-approximate", ) if nunchaku_enabled: nunchaku_kwargs = { "precision": quant_config.precision, "rank": quant_config.rank, "act_unsigned": quant_config.act_unsigned, } self.img_mlp = NunchakuFeedForward(self.img_mlp, **nunchaku_kwargs) self.txt_mlp = NunchakuFeedForward(self.txt_mlp, **nunchaku_kwargs) def _modulate( self, x: torch.Tensor, mod_params: torch.Tensor, norm_module: Union[LayerNormScaleShift, ScaleResidualLayerNormScaleShift], index: Optional[torch.Tensor] = None, gate_x: Optional[torch.Tensor] = None, residual_x: Optional[torch.Tensor] = None, ) -> Union[ Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ]: # Apply attention gates and add residual (like in Megatron) # - residual_out = gate_x * x + residual_x # - x = norm(residual_out) * (1 + scale) + shift # TODO: clean code here is_scale_residual = isinstance(norm_module, ScaleResidualLayerNormScaleShift) shift, scale, gate = mod_params.chunk(3, dim=-1) if index is not None: actual_batch = x.shape[0] shift0, shift1 = ( shift[:actual_batch], shift[actual_batch : 2 * actual_batch], ) scale0, scale1 = ( scale[:actual_batch], scale[actual_batch : 2 * actual_batch], ) gate0, gate1 = gate[:actual_batch], gate[actual_batch : 2 * actual_batch] if not x.is_contiguous(): x = x.contiguous() if not index.is_contiguous(): index = index.contiguous() if is_scale_residual: if not residual_x.is_contiguous(): residual_x = residual_x.contiguous() if not gate_x.is_contiguous(): gate_x = gate_x.contiguous() x, residual_out, gate_result = ( fuse_residual_layernorm_scale_shift_gate_select01_kernel( x, residual=residual_x, residual_gate=gate_x, weight=getattr(norm_module.norm, "weight", None), bias=getattr(norm_module.norm, "bias", None), scale0=scale0.contiguous(), shift0=shift0.contiguous(), gate0=gate0.contiguous(), scale1=scale1.contiguous(), shift1=shift1.contiguous(), gate1=gate1.contiguous(), index=index, eps=norm_module.eps, ) ) return x, residual_out, gate_result else: x, gate_result = fuse_layernorm_scale_shift_gate_select01_kernel( x, weight=getattr(norm_module.norm, "weight", None), bias=getattr(norm_module.norm, "bias", None), scale0=scale0.contiguous(), shift0=shift0.contiguous(), gate0=gate0.contiguous(), scale1=scale1.contiguous(), shift1=shift1.contiguous(), gate1=gate1.contiguous(), index=index, eps=norm_module.eps, ) return x, gate_result else: shift_result = shift.unsqueeze(1) scale_result = scale.unsqueeze(1) gate_result = gate.unsqueeze(1) if is_scale_residual: modulated, residual_out = norm_module( residual=residual_x, x=x, gate=gate_x, shift=shift_result, scale=scale_result, ) return modulated, residual_out, gate_result else: modulated = norm_module(x=x, shift=shift_result, scale=scale_result) return modulated, gate_result def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, encoder_hidden_states_mask: torch.Tensor, temb_img_silu: torch.Tensor, temb_txt_silu: torch.Tensor, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, modulate_index: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: # Get modulation parameters for both streams img_mod_params = self.img_mod[1](temb_img_silu) # [B, 6*dim] txt_mod_params = self.txt_mod[1](temb_txt_silu) # [B, 6*dim] if ( self.quant_config is not None and hasattr(self.quant_config, "get_name") and self.quant_config.get_name() == "svdquant" ): # When NOT using nunchaku, reshape mod_params from [B, 6*dim] to [B, dim*6] # When using nunchaku (svdquant), keep original format img_mod_params = ( img_mod_params.view(img_mod_params.shape[0], -1, 6) .transpose(1, 2) .reshape(img_mod_params.shape[0], -1) ) txt_mod_params = ( txt_mod_params.view(txt_mod_params.shape[0], -1, 6) .transpose(1, 2) .reshape(txt_mod_params.shape[0], -1) ) # Split modulation parameters for norm1 and norm2 img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1) # Each [B, 3*dim] txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1) # Each [B, 3*dim] # Process image stream - norm1 + modulation img_modulated, img_gate1 = self._modulate( hidden_states, img_mod1, self.img_norm1, modulate_index ) # Process text stream - norm1 + modulation txt_shift1, txt_scale1, txt_gate1_raw = txt_mod1.chunk(3, dim=-1) txt_modulated = self.txt_norm1( encoder_hidden_states, shift=txt_shift1, scale=txt_scale1 ) txt_gate1 = txt_gate1_raw.unsqueeze(1) # Use QwenAttnProcessor2_0 for joint attention computation # This directly implements the DoubleStreamLayerMegatron logic: # 1. Computes QKV for both streams # 2. Applies QK normalization and RoPE # 3. Concatenates and runs joint attention # 4. Splits results back to separate streams joint_attention_kwargs = joint_attention_kwargs or {} attn_output = self.attn( # Image stream (will be processed as "sample") hidden_states=img_modulated, # Text stream (will be processed as "context") encoder_hidden_states=txt_modulated, encoder_hidden_states_mask=encoder_hidden_states_mask, image_rotary_emb=image_rotary_emb, **joint_attention_kwargs, ) # QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided img_attn_output, txt_attn_output = attn_output # Process image stream - norm2 + MLP img_modulated2, hidden_states, img_gate2 = self._modulate( img_attn_output, img_mod2, self.img_norm2, modulate_index, gate_x=img_gate1, residual_x=hidden_states, ) img_mlp_output = self.img_mlp(img_modulated2) if img_mlp_output.dim() == 2: img_mlp_output = img_mlp_output.unsqueeze(0) hidden_states = self.fuse_mul_add(img_mlp_output, img_gate2, hidden_states) # Process text stream - norm2 + MLP txt_shift2, txt_scale2, txt_gate2_raw = txt_mod2.chunk(3, dim=-1) txt_modulated2, encoder_hidden_states = self.txt_norm2( residual=encoder_hidden_states, x=txt_attn_output, gate=txt_gate1, shift=txt_shift2, scale=txt_scale2, ) txt_gate2 = txt_gate2_raw.unsqueeze(1) txt_mlp_output = self.txt_mlp(txt_modulated2) if txt_mlp_output.dim() == 2: txt_mlp_output = txt_mlp_output.unsqueeze(0) encoder_hidden_states = self.fuse_mul_add( txt_mlp_output, txt_gate2, encoder_hidden_states ) # Clip to prevent overflow for fp16 if encoder_hidden_states.dtype == torch.float16: encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) if hidden_states.dtype == torch.float16: hidden_states = hidden_states.clip(-65504, 65504) return encoder_hidden_states, hidden_states def to_hashable(obj): if isinstance(obj, list): return tuple(to_hashable(x) for x in obj) return obj class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin): """ The Transformer model introduced in Qwen. """ _supports_gradient_checkpointing = True _no_split_modules = ["QwenImageTransformerBlock"] _skip_layerwise_casting_patterns = ["pos_embed", "norm"] _repeated_blocks = ["QwenImageTransformerBlock"] param_names_mapping = QwenImageDitConfig().arch_config.param_names_mapping @classmethod def get_nunchaku_quant_rules(cls) -> dict[str, list[str]]: return { "skip": [ "norm", "embed", "rotary", "pos_embed", ], "svdq_w4a4": [ "attn.to_qkv", "attn.to_out", "attn.add_qkv_proj", "attn.to_add_out", "img_mlp", "txt_mlp", ], "awq_w4a16": [ "img_mod", "txt_mod", ], } def __init__( self, config: QwenImageDitConfig, hf_config: dict[str, Any], quant_config: Optional[QuantizationConfig] = None, ): super().__init__(config=config, hf_config=hf_config) patch_size = config.arch_config.patch_size in_channels = config.arch_config.in_channels out_channels = config.arch_config.out_channels num_layers = config.arch_config.num_layers attention_head_dim = config.arch_config.attention_head_dim num_attention_heads = config.arch_config.num_attention_heads joint_attention_dim = config.arch_config.joint_attention_dim axes_dims_rope = config.arch_config.axes_dims_rope self.zero_cond_t = getattr(config.arch_config, "zero_cond_t", False) self.out_channels = out_channels or in_channels self.inner_dim = num_attention_heads * attention_head_dim self.use_additional_t_cond: bool = getattr( config.arch_config, "use_additional_t_cond", False ) # For qwen-image-layered now self.use_layer3d_rope: bool = getattr( config.arch_config, "use_layer3d_rope", False ) # For qwen-image-layered now if not self.use_layer3d_rope: self.rotary_emb = QwenEmbedRope( theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True ) else: self.rotary_emb = QwenEmbedLayer3DRope( theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True ) self.time_text_embed = QwenTimestepProjEmbeddings( embedding_dim=self.inner_dim, use_additional_t_cond=self.use_additional_t_cond, ) self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6) self.img_in = nn.Linear(in_channels, self.inner_dim) self.txt_in = nn.Linear(joint_attention_dim, self.inner_dim) self.transformer_blocks = nn.ModuleList( [ QwenImageTransformerBlock( dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, quant_config=quant_config, prefix=f"transformer_blocks.{layer_idx}", zero_cond_t=self.zero_cond_t, ) for layer_idx in range(num_layers) ] ) self.norm_out = AdaLayerNormContinuous( self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6 ) self.proj_out = nn.Linear( self.inner_dim, patch_size * patch_size * self.out_channels, bias=True ) self.timestep_zero = torch.zeros( (1,), dtype=torch.int, device=get_local_torch_device() ) self.layer_names = ["transformer_blocks"] @functools.lru_cache(maxsize=50) def build_modulate_index(self, img_shapes: tuple[int, int, int], device): modulate_index_list = [] for sample in img_shapes: first_size = sample[0][0] * sample[0][1] * sample[0][2] total_size = sum(s[0] * s[1] * s[2] for s in sample) idx = (torch.arange(total_size, device=device) >= first_size).int() modulate_index_list.append(idx) modulate_index = torch.stack(modulate_index_list) return modulate_index def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, encoder_hidden_states_mask: torch.Tensor = None, timestep: torch.LongTensor = None, img_shapes: Optional[List[Tuple[int, int, int]]] = None, txt_seq_lens: Optional[List[int]] = None, freqs_cis: tuple[torch.Tensor, torch.Tensor] = None, additional_t_cond: Optional[torch.Tensor] = None, guidance: torch.Tensor = None, # TODO: this should probably be removed attention_kwargs: Optional[Dict[str, Any]] = None, controlnet_block_samples=None, return_dict: bool = True, ) -> Union[torch.Tensor, Transformer2DModelOutput]: """ The [`QwenTransformer2DModel`] forward method. Args: hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`): Input `hidden_states`. encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`): Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`): Mask of the input conditions. timestep ( `torch.LongTensor`): Used to indicate denoising step. attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain tuple. Returns: If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a `tuple` where the first element is the sample tensor. """ if ( attention_kwargs is not None and attention_kwargs.get("scale", None) is not None ): logger.warning( "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." ) if isinstance(encoder_hidden_states, list): encoder_hidden_states = encoder_hidden_states[0] hidden_states = self.img_in(hidden_states) timestep = (timestep / 1000).to(hidden_states.dtype) if self.zero_cond_t: timestep = torch.cat([timestep, self.timestep_zero], dim=0) device = timestep.device modulate_index = self.build_modulate_index(to_hashable(img_shapes), device) else: modulate_index = None encoder_hidden_states = self.txt_norm(encoder_hidden_states) encoder_hidden_states = self.txt_in(encoder_hidden_states) temb = self.time_text_embed(timestep, hidden_states, additional_t_cond) temb_img_silu = F.silu(temb) if self.zero_cond_t: temb_txt = temb.chunk(2, dim=0)[0] temb_txt_silu = temb_img_silu.chunk(2, dim=0)[0] else: temb_txt = temb temb_txt_silu = temb_img_silu image_rotary_emb = freqs_cis for index_block, block in enumerate(self.transformer_blocks): encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_mask=encoder_hidden_states_mask, temb_img_silu=temb_img_silu, temb_txt_silu=temb_txt_silu, image_rotary_emb=image_rotary_emb, joint_attention_kwargs=attention_kwargs, modulate_index=modulate_index, ) # controlnet residual if controlnet_block_samples is not None: interval_control = len(self.transformer_blocks) / len( controlnet_block_samples ) interval_control = int(np.ceil(interval_control)) hidden_states = ( hidden_states + controlnet_block_samples[index_block // interval_control] ) # Use only the image part (hidden_states) from the dual-stream blocks hidden_states = self.norm_out(hidden_states, temb_txt) output = self.proj_out(hidden_states) return output EntryClass = QwenImageTransformer2DModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/sana.py ================================================ # SPDX-License-Identifier: Apache-2.0 import torch import torch.nn as nn import torch.nn.functional as F from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding from sglang.multimodal_gen.configs.models.dits.sana import SanaConfig from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm from sglang.multimodal_gen.runtime.layers.visual_embedding import Timesteps from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class SanaCombinedTimestepSizeEmbeddings(nn.Module): def __init__(self, embedding_dim): super().__init__() self.time_proj = Timesteps( num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0 ) self.timestep_embedder = TimestepEmbedding( in_channels=256, time_embed_dim=embedding_dim ) def forward(self, timestep, hidden_dtype=None): timesteps_proj = self.time_proj(timestep) if hidden_dtype is not None: timesteps_proj = timesteps_proj.to(dtype=hidden_dtype) timesteps_emb = self.timestep_embedder(timesteps_proj) return timesteps_emb class SanaAdaLayerNormSingle(nn.Module): def __init__(self, embedding_dim): super().__init__() self.emb = SanaCombinedTimestepSizeEmbeddings(embedding_dim) self.silu = nn.SiLU() self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) def forward(self, timestep, hidden_dtype=None): embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype) out = self.linear(self.silu(embedded_timestep)) return out, embedded_timestep class SanaModulatedNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) def forward(self, x, temb, scale_shift_table): x = self.norm(x) shift, scale = (scale_shift_table[None] + temb[:, None]).chunk(2, dim=1) x = x * (1 + scale) + shift return x class GLUMBConv(nn.Module): """Gated Linear Unit with Multi-Branch Convolution.""" def __init__(self, in_channels, out_channels, expand_ratio=2.5): super().__init__() hidden_channels = int(expand_ratio * in_channels) self.nonlinearity = nn.SiLU() self.conv_inverted = nn.Conv2d(in_channels, hidden_channels * 2, 1, 1, 0) self.conv_depth = nn.Conv2d( hidden_channels * 2, hidden_channels * 2, 3, 1, 1, groups=hidden_channels * 2, ) self.conv_point = nn.Conv2d(hidden_channels, out_channels, 1, 1, 0, bias=False) def forward(self, hidden_states): hidden_states = self.conv_inverted(hidden_states) hidden_states = self.nonlinearity(hidden_states) hidden_states = self.conv_depth(hidden_states) hidden_states, gate = torch.chunk(hidden_states, 2, dim=1) hidden_states = hidden_states * self.nonlinearity(gate) hidden_states = self.conv_point(hidden_states) return hidden_states class SanaLinearAttention(nn.Module): """Linear attention with O(N*D^2) complexity instead of O(N^2*D).""" def __init__(self, query_dim, num_heads, head_dim, qk_norm_dim, bias=False): super().__init__() inner_dim = num_heads * head_dim self.num_heads = num_heads self.head_dim = head_dim self.to_q = nn.Linear(query_dim, inner_dim, bias=bias) self.to_k = nn.Linear(query_dim, inner_dim, bias=bias) self.to_v = nn.Linear(query_dim, inner_dim, bias=bias) self.to_out = nn.ModuleList( [nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()] ) self.norm_q = RMSNorm(qk_norm_dim) self.norm_k = RMSNorm(qk_norm_dim) def forward(self, hidden_states): B, S, _ = hidden_states.shape query = self.to_q(hidden_states) key = self.to_k(hidden_states) value = self.to_v(hidden_states) query = self.norm_q(query) key = self.norm_k(key) query = query.view(B, S, self.num_heads, self.head_dim).transpose(1, 2) key = key.view(B, S, self.num_heads, self.head_dim).transpose(1, 2) value = value.view(B, S, self.num_heads, self.head_dim).transpose(1, 2) query = F.relu(query) key = F.relu(key) kv = torch.matmul(key.transpose(-2, -1), value) # (B, H, D, D) qkv = torch.matmul(query, kv) # (B, H, S, D) key_sum = key.sum(dim=-2, keepdim=True) # (B, H, 1, D) normalizer = torch.matmul(query, key_sum.transpose(-2, -1)).clamp(min=1e-6) hidden_states = qkv / normalizer hidden_states = hidden_states.transpose(1, 2).reshape(B, S, -1) hidden_states = self.to_out[0](hidden_states) return hidden_states class SanaCrossAttention(nn.Module): def __init__(self, query_dim, cross_attention_dim, num_heads, head_dim, bias=False): super().__init__() inner_dim = num_heads * head_dim self.num_heads = num_heads self.head_dim = head_dim self.to_q = nn.Linear(query_dim, inner_dim, bias=bias) self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias) self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=bias) self.to_out = nn.ModuleList( [nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()] ) self.norm_q = RMSNorm(inner_dim) self.norm_k = RMSNorm(inner_dim) def forward( self, hidden_states, encoder_hidden_states, encoder_attention_mask=None ): B, S, _ = hidden_states.shape T = encoder_hidden_states.shape[1] query = self.to_q(hidden_states) key = self.to_k(encoder_hidden_states) value = self.to_v(encoder_hidden_states) query = self.norm_q(query) key = self.norm_k(key) query = query.view(B, S, self.num_heads, self.head_dim).transpose(1, 2) key = key.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) value = value.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) attn_mask = None if encoder_attention_mask is not None: attn_mask = encoder_attention_mask.bool() attn_mask = attn_mask[:, None, None, :].expand(B, self.num_heads, S, T) hidden_states = F.scaled_dot_product_attention( query, key, value, attn_mask=attn_mask ) hidden_states = hidden_states.transpose(1, 2).reshape(B, S, -1) hidden_states = self.to_out[0](hidden_states) return hidden_states class SanaTransformerBlock(nn.Module): def __init__( self, dim, num_attention_heads, attention_head_dim, num_cross_attention_heads, cross_attention_head_dim, cross_attention_dim, mlp_ratio, norm_eps, attention_bias=False, ): super().__init__() self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps) self.attn1 = SanaLinearAttention( query_dim=dim, num_heads=num_attention_heads, head_dim=attention_head_dim, qk_norm_dim=num_attention_heads * attention_head_dim, bias=attention_bias, ) self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps) self.attn2 = SanaCrossAttention( query_dim=dim, cross_attention_dim=cross_attention_dim, num_heads=num_cross_attention_heads, head_dim=cross_attention_head_dim, bias=True, ) self.ff = GLUMBConv(in_channels=dim, out_channels=dim, expand_ratio=mlp_ratio) def forward( self, hidden_states, encoder_hidden_states, timestep, height, width, encoder_attention_mask=None, ): batch_size = hidden_states.shape[0] shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) ).chunk(6, dim=1) norm_hidden = self.norm1(hidden_states) norm_hidden = norm_hidden * (1 + scale_msa) + shift_msa attn_output = self.attn1(norm_hidden) hidden_states = hidden_states + gate_msa * attn_output attn_output = self.attn2( hidden_states, encoder_hidden_states, encoder_attention_mask ) hidden_states = hidden_states + attn_output norm_hidden = self.norm2(hidden_states) norm_hidden = norm_hidden * (1 + scale_mlp) + shift_mlp norm_hidden = norm_hidden.unflatten(1, (height, width)).permute(0, 3, 1, 2) ff_output = self.ff(norm_hidden) ff_output = ff_output.flatten(2, 3).permute(0, 2, 1) hidden_states = hidden_states + gate_mlp * ff_output return hidden_states class SanaTransformer2DModel(CachableDiT, OffloadableDiTMixin): _fsdp_shard_conditions = [ lambda n, m: isinstance(m, SanaTransformerBlock), ] _compile_conditions = [ lambda n, m: isinstance(m, SanaTransformerBlock), ] param_names_mapping = SanaConfig().arch_config.param_names_mapping reverse_param_names_mapping = {} def __init__(self, config: SanaConfig, hf_config=None, **kwargs): super().__init__(config, hf_config=hf_config or {}, **kwargs) arch = config.arch_config self.out_channels = arch.out_channels self.patch_size = arch.patch_size self.inner_dim = arch.num_attention_heads * arch.attention_head_dim self.hidden_size = self.inner_dim self.num_attention_heads = arch.num_attention_heads self.num_channels_latents = arch.num_channels_latents self.patch_embed = nn.ModuleDict( { "proj": nn.Conv2d( arch.in_channels, self.inner_dim, kernel_size=arch.patch_size, stride=arch.patch_size, bias=True, ), } ) self.time_embed = SanaAdaLayerNormSingle(self.inner_dim) self.caption_projection = PixArtAlphaTextProjection( in_features=arch.caption_channels, hidden_size=self.inner_dim, ) self.caption_norm = RMSNorm(self.inner_dim) self.transformer_blocks = nn.ModuleList( [ SanaTransformerBlock( dim=self.inner_dim, num_attention_heads=arch.num_attention_heads, attention_head_dim=arch.attention_head_dim, num_cross_attention_heads=arch.num_cross_attention_heads, cross_attention_head_dim=arch.cross_attention_head_dim, cross_attention_dim=arch.cross_attention_dim, mlp_ratio=arch.mlp_ratio, norm_eps=arch.norm_eps, attention_bias=False, ) for _ in range(arch.num_layers) ] ) self.scale_shift_table = nn.Parameter( torch.randn(2, self.inner_dim) / self.inner_dim**0.5 ) self.norm_out = SanaModulatedNorm(self.inner_dim, eps=arch.norm_eps) self.proj_out = nn.Linear( self.inner_dim, arch.patch_size * arch.patch_size * self.out_channels, bias=True, ) self.layer_names = ["transformer_blocks"] def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, timestep: torch.LongTensor = None, guidance: torch.Tensor = None, encoder_attention_mask: torch.Tensor = None, **kwargs, ) -> torch.Tensor: # Input validation - fail fast if encoder_hidden_states is None: raise ValueError("SANA forward pass requires encoder_hidden_states") batch_size, channels, height, width = hidden_states.shape p = self.patch_size post_patch_height = height // p post_patch_width = width // p hidden_states = self.patch_embed["proj"](hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) timestep_emb, embedded_timestep = self.time_embed( timestep, hidden_dtype=hidden_states.dtype ) if isinstance(encoder_attention_mask, (list, tuple)): encoder_attention_mask = encoder_attention_mask[0] encoder_hidden_states = self.caption_projection(encoder_hidden_states) if encoder_hidden_states.shape[0] != batch_size: encoder_hidden_states = encoder_hidden_states.expand( batch_size, -1, -1 ).contiguous() encoder_hidden_states = encoder_hidden_states.view( batch_size, -1, hidden_states.shape[-1] ) encoder_hidden_states = self.caption_norm(encoder_hidden_states) if ( encoder_attention_mask is not None and encoder_attention_mask.shape[0] != batch_size ): encoder_attention_mask = encoder_attention_mask.expand( batch_size, -1 ).contiguous() for block in self.transformer_blocks: hidden_states = block( hidden_states, encoder_hidden_states, timestep_emb, post_patch_height, post_patch_width, encoder_attention_mask=encoder_attention_mask, ) hidden_states = self.norm_out( hidden_states, embedded_timestep, self.scale_shift_table ) hidden_states = self.proj_out(hidden_states) hidden_states = hidden_states.reshape( batch_size, post_patch_height, post_patch_width, p, p, self.out_channels ) hidden_states = hidden_states.permute(0, 5, 1, 3, 2, 4) hidden_states = hidden_states.reshape( batch_size, self.out_channels, height, width ) return hidden_states EntryClass = SanaTransformer2DModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import math from functools import lru_cache from typing import Any import torch import torch.nn as nn from sglang.multimodal_gen.configs.models.dits import WanVideoConfig from sglang.multimodal_gen.configs.sample.wan import WanTeaCacheParams from sglang.multimodal_gen.runtime.distributed import ( divide, get_sp_group, get_sp_world_size, get_tp_world_size, sequence_model_parallel_all_gather, ) from sglang.multimodal_gen.runtime.layers.attention import ( MinimalA2AAttnOp, UlyssesAttention_VSA, USPAttention, ) from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd from sglang.multimodal_gen.runtime.layers.layernorm import ( FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidualLayerNormScaleShift, tensor_parallel_rms_norm, ) from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( NDRotaryEmbedding, _apply_rotary_emb, apply_flashinfer_rope_qk_inplace, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import ( ModulateProjection, PatchEmbed, TimestepEmbedder, ) from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.server_args import get_global_server_args from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.srt.utils import add_prefix logger = init_logger(__name__) _is_cuda = current_platform.is_cuda() class WanImageEmbedding(torch.nn.Module): def __init__(self, in_features: int, out_features: int): super().__init__() self.norm1 = FP32LayerNorm(in_features) self.ff = MLP(in_features, in_features, out_features, act_type="gelu") self.norm2 = FP32LayerNorm(out_features) def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: dtype = encoder_hidden_states_image.dtype hidden_states = self.norm1(encoder_hidden_states_image) hidden_states = self.ff(hidden_states) hidden_states = self.norm2(hidden_states).to(dtype) return hidden_states class WanTimeTextImageEmbedding(nn.Module): def __init__( self, dim: int, time_freq_dim: int, text_embed_dim: int, image_embed_dim: int | None = None, ): super().__init__() self.time_embedder = TimestepEmbedder( dim, frequency_embedding_size=time_freq_dim, act_layer="silu" ) self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu") self.text_embedder = MLP( text_embed_dim, dim, dim, bias=True, act_type="gelu_pytorch_tanh" ) self.image_embedder = None if image_embed_dim is not None: self.image_embedder = WanImageEmbedding(image_embed_dim, dim) def forward( self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, encoder_hidden_states_image: torch.Tensor | None = None, timestep_seq_len: int | None = None, ): temb = self.time_embedder(timestep, timestep_seq_len) timestep_proj = self.time_modulation(temb) encoder_hidden_states = self.text_embedder(encoder_hidden_states) if encoder_hidden_states_image is not None: assert self.image_embedder is not None encoder_hidden_states_image = self.image_embedder( encoder_hidden_states_image ) return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image class WanSelfAttention(nn.Module): def __init__( self, dim: int, num_heads: int, window_size=(-1, -1), qk_norm=True, eps=1e-6, parallel_attention=False, prefix: str = "", supported_attention_backends: set[AttentionBackendEnum] | None = None, is_cross_attention: bool = False, quant_config: QuantizationConfig | None = None, ) -> None: assert dim % num_heads == 0 super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads self.window_size = window_size self.qk_norm = qk_norm self.eps = eps self.parallel_attention = parallel_attention tp_size = get_tp_world_size() # layers self.to_q = ColumnParallelLinear( dim, dim, gather_output=False, quant_config=quant_config, prefix=add_prefix("to_q", prefix), ) self.to_k = ColumnParallelLinear( dim, dim, gather_output=False, quant_config=quant_config, prefix=add_prefix("to_k", prefix), ) self.to_v = ColumnParallelLinear( dim, dim, gather_output=False, quant_config=quant_config, prefix=add_prefix("to_v", prefix), ) self.to_out = RowParallelLinear( dim, dim, input_is_parallel=True, quant_config=quant_config, prefix=add_prefix("to_out.0", prefix), ) self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() self.tp_rmsnorm = tp_size > 1 and qk_norm self.local_num_heads = divide(num_heads, tp_size) # Scaled dot product attention self.attn = USPAttention( num_heads=self.local_num_heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, causal=False, supported_attention_backends=supported_attention_backends, skip_sequence_parallel=is_cross_attention, quant_config=quant_config, ) def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: int): r""" Args: x(Tensor): Shape [B, L, num_heads, C / num_heads] """ pass class WanT2VCrossAttention(WanSelfAttention): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs, is_cross_attention=True) def forward(self, x, context, context_lens): r""" Args: x(Tensor): Shape [B, L1, C] context(Tensor): Shape [B, L2, C] context_lens(Tensor): Shape [B] """ q, _ = self.to_q(x) if self.tp_rmsnorm: q = tensor_parallel_rms_norm(q, self.norm_q) else: q = self.norm_q(q) q = q.unflatten(2, (self.local_num_heads, self.head_dim)) k, _ = self.to_k(context) if self.tp_rmsnorm: k = tensor_parallel_rms_norm(k, self.norm_k) else: k = self.norm_k(k) k = k.unflatten(2, (self.local_num_heads, self.head_dim)) v, _ = self.to_v(context) v = v.unflatten(2, (self.local_num_heads, self.head_dim)) # compute attention x = self.attn(q, k, v) # output x = x.flatten(2) x, _ = self.to_out(x) return x class WanI2VCrossAttention(WanSelfAttention): def __init__( self, dim: int, num_heads: int, window_size=(-1, -1), qk_norm=True, eps=1e-6, prefix: str = "", supported_attention_backends: set[AttentionBackendEnum] | None = None, quant_config: QuantizationConfig | None = None, ) -> None: super().__init__( dim, num_heads, window_size, qk_norm, eps, supported_attention_backends=supported_attention_backends, is_cross_attention=True, quant_config=quant_config, ) self.add_k_proj = ColumnParallelLinear( dim, dim, gather_output=False, quant_config=quant_config, prefix=add_prefix("add_k_proj", prefix), ) self.add_v_proj = ColumnParallelLinear( dim, dim, gather_output=False, quant_config=quant_config, prefix=add_prefix("add_v_proj", prefix), ) self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() def forward(self, x, context, context_lens): r""" Args: x(Tensor): Shape [B, L1, C] context(Tensor): Shape [B, L2, C] context_lens(Tensor): Shape [B] """ context_img = context[:, :257] context = context[:, 257:] q, _ = self.to_q(x) if self.tp_rmsnorm: q = tensor_parallel_rms_norm(q, self.norm_q) else: q = self.norm_q(q) q = q.unflatten(2, (self.local_num_heads, self.head_dim)) k, _ = self.to_k(context) if self.tp_rmsnorm: k = tensor_parallel_rms_norm(k, self.norm_k) else: k = self.norm_k(k) k = k.unflatten(2, (self.local_num_heads, self.head_dim)) v, _ = self.to_v(context) v = v.unflatten(2, (self.local_num_heads, self.head_dim)) k_img, _ = self.add_k_proj(context_img) if self.tp_rmsnorm: k_img = tensor_parallel_rms_norm(k_img, self.norm_added_k) else: k_img = self.norm_added_k(k_img) k_img = k_img.unflatten(2, (self.local_num_heads, self.head_dim)) v_img, _ = self.add_v_proj(context_img) v_img = v_img.unflatten(2, (self.local_num_heads, self.head_dim)) img_x = self.attn(q, k_img, v_img) x = self.attn(q, k, v) # output x = x.flatten(2) img_x = img_x.flatten(2) x = x + img_x x, _ = self.to_out(x) return x class WanTransformerBlock(nn.Module): def __init__( self, dim: int, ffn_dim: int, num_heads: int, qk_norm: str = "rms_norm_across_heads", cross_attn_norm: bool = False, eps: float = 1e-6, added_kv_proj_dim: int | None = None, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", attention_type: str = "original", sla_topk: float = 0.1, quant_config: QuantizationConfig | None = None, ): super().__init__() # 1. Self-attention self.norm1 = LayerNormScaleShift( dim, eps=eps, elementwise_affine=False, dtype=torch.float32, ) self.to_q = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config, prefix=add_prefix("attn1.to_q", prefix), ) self.to_k = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config, prefix=add_prefix("attn1.to_k", prefix), ) self.to_v = ColumnParallelLinear( dim, dim, bias=True, gather_output=False, quant_config=quant_config, prefix=add_prefix("attn1.to_v", prefix), ) self.to_out = RowParallelLinear( dim, dim, bias=True, reduce_results=True, quant_config=quant_config, prefix=add_prefix("attn1.to_out.0", prefix), ) tp_size = get_tp_world_size() self.local_num_heads = divide(num_heads, tp_size) self_attn_backends = supported_attention_backends if attention_type in ("sla", "sagesla"): self.attn1 = MinimalA2AAttnOp( num_heads=self.local_num_heads, head_size=dim // num_heads, attention_type=attention_type, topk=sla_topk, supported_attention_backends={ AttentionBackendEnum.SLA_ATTN, AttentionBackendEnum.SAGE_SLA_ATTN, }, prefix=add_prefix("attn1", prefix), ) else: self.attn1 = USPAttention( num_heads=self.local_num_heads, head_size=dim // num_heads, causal=False, supported_attention_backends=self_attn_backends, prefix=add_prefix("attn1", prefix), quant_config=quant_config, is_cross_attention=False, ) self.hidden_dim = dim self.num_attention_heads = num_heads self.dim_head = dim // num_heads if qk_norm == "rms_norm": self.norm_q = RMSNorm(self.dim_head, eps=eps) self.norm_k = RMSNorm(self.dim_head, eps=eps) elif qk_norm == "rms_norm_across_heads": # LTX applies qk norm across all heads self.norm_q = RMSNorm(dim, eps=eps) self.norm_k = RMSNorm(dim, eps=eps) else: logger.error("QK Norm type not supported") raise Exception assert cross_attn_norm is True self.qk_norm = qk_norm self.tp_rmsnorm = qk_norm == "rms_norm_across_heads" and tp_size > 1 self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, eps=eps, elementwise_affine=True, dtype=torch.float32, ) # 2. Cross-attention cross_attn_backends = { b for b in supported_attention_backends if not b.is_sparse } if added_kv_proj_dim is not None: # I2V self.attn2 = WanI2VCrossAttention( dim, num_heads, qk_norm=qk_norm, eps=eps, prefix=add_prefix("attn2", prefix), supported_attention_backends=cross_attn_backends, quant_config=quant_config, ) else: # T2V self.attn2 = WanT2VCrossAttention( dim, num_heads, qk_norm=qk_norm, eps=eps, prefix=add_prefix("attn2", prefix), supported_attention_backends=cross_attn_backends, quant_config=quant_config, ) self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, eps=eps, elementwise_affine=False, dtype=torch.float32, ) # 3. Feed-forward self.ffn = MLP( dim, ffn_dim, act_type="gelu_pytorch_tanh", prefix=add_prefix("ffn.net", prefix), quant_config=quant_config, ) self.mlp_residual = MulAdd() self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype if temb.dim() == 4: # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( self.scale_shift_table.unsqueeze(0) + temb.float() ).chunk(6, dim=2) # batch_size, seq_len, 1, inner_dim shift_msa = shift_msa.squeeze(2) scale_msa = scale_msa.squeeze(2) gate_msa = gate_msa.squeeze(2) c_shift_msa = c_shift_msa.squeeze(2) c_scale_msa = c_scale_msa.squeeze(2) c_gate_msa = c_gate_msa.squeeze(2) else: # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B) e = self.scale_shift_table + temb.float() shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( e.chunk(6, dim=1) ) assert shift_msa.dtype == torch.float32 # 1. Self-attention norm_hidden_states = self.norm1(hidden_states, shift_msa, scale_msa) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) if self.norm_q is not None: if self.tp_rmsnorm: query = tensor_parallel_rms_norm(query, self.norm_q) else: query = self.norm_q(query) if self.norm_k is not None: if self.tp_rmsnorm: key = tensor_parallel_rms_norm(key, self.norm_k) else: key = self.norm_k(key) query = query.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head)) key = key.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head)) value = value.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head)) # Apply rotary embeddings cos, sin = freqs_cis if _is_cuda and query.shape == key.shape: cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(), sin.to(dtype=torch.float32).contiguous(), ], dim=-1, ) query, key = apply_flashinfer_rope_qk_inplace( query, key, cos_sin_cache, is_neox=False ) else: query, key = _apply_rotary_emb( query, cos, sin, is_neox_style=False ), _apply_rotary_emb(key, cos, sin, is_neox_style=False) attn_output = self.attn1(query, key, value) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.zeros( (1,), device=hidden_states.device, dtype=hidden_states.dtype ) norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale ) norm_hidden_states, hidden_states = norm_hidden_states.to( orig_dtype ), hidden_states.to(orig_dtype) # 2. Cross-attention attn_output = self.attn2( norm_hidden_states, context=encoder_hidden_states, context_lens=None ) norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) norm_hidden_states, hidden_states = norm_hidden_states.to( orig_dtype ), hidden_states.to(orig_dtype) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states) hidden_states = hidden_states.to(orig_dtype) return hidden_states class WanTransformerBlock_VSA(nn.Module): def __init__( self, dim: int, ffn_dim: int, num_heads: int, qk_norm: str = "rms_norm_across_heads", cross_attn_norm: bool = False, eps: float = 1e-6, added_kv_proj_dim: int | None = None, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", quant_config: QuantizationConfig | None = None, ): super().__init__() # 1. Self-attention self.norm1 = LayerNormScaleShift( dim, eps=eps, elementwise_affine=False, dtype=torch.float32, ) self.to_q = ColumnParallelLinear( dim, dim, bias=True, gather_output=True, quant_config=quant_config, prefix=add_prefix("attn1.to_q", prefix), ) self.to_k = ColumnParallelLinear( dim, dim, bias=True, gather_output=True, quant_config=quant_config, prefix=add_prefix("attn1.to_k", prefix), ) self.to_v = ColumnParallelLinear( dim, dim, bias=True, gather_output=True, quant_config=quant_config, prefix=add_prefix("attn1.to_v", prefix), ) self.to_gate_compress = ColumnParallelLinear( dim, dim, bias=True, gather_output=True, quant_config=quant_config, prefix=add_prefix("attn1.to_gate_compress", prefix), ) self.to_out = ColumnParallelLinear( dim, dim, bias=True, gather_output=True, quant_config=quant_config, prefix=add_prefix("attn1.to_out.0", prefix), ) self.attn1 = UlyssesAttention_VSA( num_heads=num_heads, head_size=dim // num_heads, causal=False, supported_attention_backends=supported_attention_backends, prefix=add_prefix("attn1", prefix), quant_config=quant_config, ) self.hidden_dim = dim self.num_attention_heads = num_heads dim_head = dim // num_heads if qk_norm == "rms_norm": self.norm_q = RMSNorm(dim_head, eps=eps) self.norm_k = RMSNorm(dim_head, eps=eps) elif qk_norm == "rms_norm_across_heads": # LTX applies qk norm across all heads self.norm_q = RMSNorm(dim, eps=eps) self.norm_k = RMSNorm(dim, eps=eps) else: logger.error("QK Norm type not supported") raise Exception assert cross_attn_norm is True self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, eps=eps, elementwise_affine=True, dtype=torch.float32, ) # 2. Cross-attention cross_attn_backends = { b for b in supported_attention_backends if not b.is_sparse } if added_kv_proj_dim is not None: # I2V self.attn2 = WanI2VCrossAttention( dim, num_heads, qk_norm=qk_norm, eps=eps, prefix=add_prefix("attn2", prefix), supported_attention_backends=cross_attn_backends, quant_config=quant_config, ) else: # T2V self.attn2 = WanT2VCrossAttention( dim, num_heads, qk_norm=qk_norm, eps=eps, prefix=add_prefix("attn2", prefix), supported_attention_backends=cross_attn_backends, quant_config=quant_config, ) self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, eps=eps, elementwise_affine=False, dtype=torch.float32, ) # 3. Feed-forward self.ffn = MLP( dim, ffn_dim, act_type="gelu_pytorch_tanh", prefix=add_prefix("ffn.net", prefix), quant_config=quant_config, ) self.mlp_residual = MulAdd() self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype # assert orig_dtype != torch.float32 e = self.scale_shift_table + temb.float() shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( 6, dim=1 ) assert shift_msa.dtype == torch.float32 # 1. Self-attention norm_hidden_states = self.norm1(hidden_states, shift_msa, scale_msa) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) gate_compress, _ = self.to_gate_compress(norm_hidden_states) if self.norm_q is not None: query = self.norm_q(query) if self.norm_k is not None: key = self.norm_k(key) query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) gate_compress = gate_compress.squeeze(1).unflatten( 2, (self.num_attention_heads, -1) ) # Apply rotary embeddings cos, sin = freqs_cis if _is_cuda and query.shape == key.shape: cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(), sin.to(dtype=torch.float32).contiguous(), ], dim=-1, ) query, key = apply_flashinfer_rope_qk_inplace( query, key, cos_sin_cache, is_neox=False ) else: query, key = _apply_rotary_emb( query, cos, sin, is_neox_style=False ), _apply_rotary_emb(key, cos, sin, is_neox_style=False) attn_output = self.attn1(query, key, value, gate_compress=gate_compress) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.zeros((1,), device=hidden_states.device) norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale ) norm_hidden_states, hidden_states = norm_hidden_states.to( orig_dtype ), hidden_states.to(orig_dtype) # 2. Cross-attention attn_output = self.attn2( norm_hidden_states, context=encoder_hidden_states, context_lens=None ) norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) norm_hidden_states, hidden_states = norm_hidden_states.to( orig_dtype ), hidden_states.to(orig_dtype) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states) hidden_states = hidden_states.to(orig_dtype) return hidden_states class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin): _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions _compile_conditions = WanVideoConfig()._compile_conditions _supported_attention_backends = WanVideoConfig()._supported_attention_backends param_names_mapping = WanVideoConfig().param_names_mapping reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping def __init__( self, config: WanVideoConfig, hf_config: dict[str, Any], quant_config: QuantizationConfig | None = None, ) -> None: super().__init__(config=config, hf_config=hf_config) inner_dim = config.num_attention_heads * config.attention_head_dim self.hidden_size = config.hidden_size self.num_attention_heads = config.num_attention_heads self.in_channels = config.in_channels self.out_channels = config.out_channels self.num_channels_latents = config.num_channels_latents self.patch_size = config.patch_size self.text_len = config.text_len # 1. Patch & position embedding self.patch_embedding = PatchEmbed( in_chans=config.in_channels, embed_dim=inner_dim, patch_size=config.patch_size, flatten=False, ) # 2. Condition embeddings self.condition_embedder = WanTimeTextImageEmbedding( dim=inner_dim, time_freq_dim=config.freq_dim, text_embed_dim=config.text_dim, image_embed_dim=config.image_dim, ) # 3. Transformer blocks attn_backend = get_global_server_args().attention_backend transformer_block = ( WanTransformerBlock_VSA if (attn_backend and attn_backend.lower() == "video_sparse_attn") else WanTransformerBlock ) self.blocks = nn.ModuleList( [ transformer_block( inner_dim, config.ffn_dim, config.num_attention_heads, config.qk_norm, config.cross_attn_norm, config.eps, config.added_kv_proj_dim, self._supported_attention_backends | {AttentionBackendEnum.VIDEO_SPARSE_ATTN}, prefix=f"blocks.{i}", attention_type=config.attention_type, sla_topk=config.sla_topk, quant_config=quant_config, ) for i in range(config.num_layers) ] ) # 4. Output norm & projection self.norm_out = LayerNormScaleShift( inner_dim, eps=config.eps, elementwise_affine=False, dtype=torch.float32, ) self.proj_out = ColumnParallelLinear( inner_dim, config.out_channels * math.prod(config.patch_size), bias=True, gather_output=True, prefix=f"proj_out", quant_config=quant_config, ) self.scale_shift_table = nn.Parameter( torch.randn(1, 2, inner_dim) / inner_dim**0.5 ) # For type checking self.cnt = 0 self.__post_init__() # misc self.sp_size = get_sp_world_size() # Get rotary embeddings d = self.hidden_size // self.num_attention_heads self.rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] self.rotary_emb = NDRotaryEmbedding( rope_dim_list=self.rope_dim_list, rope_theta=10000, dtype=( torch.float32 if current_platform.is_mps() or current_platform.is_musa() else torch.float64 ), ) self.layer_names = ["blocks"] @lru_cache(maxsize=1) def _compute_rope_for_sequence_shard( self, local_len: int, rank: int, frame_stride_local: int, width_local: int, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor]: token_start = rank * local_len token_indices = torch.arange( token_start, token_start + local_len, device=device, dtype=torch.long, ) t_idx = token_indices // frame_stride_local rem = token_indices % frame_stride_local h_idx = rem // width_local w_idx = rem % width_local positions = torch.stack((t_idx, h_idx, w_idx), dim=1) return self.rotary_emb.forward_uncached(positions) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, guidance=None, **kwargs, ) -> torch.Tensor: forward_batch = get_forward_context().forward_batch if forward_batch is not None: sequence_shard_enabled = ( forward_batch.enable_sequence_shard and self.sp_size > 1 ) else: sequence_shard_enabled = False self.enable_teacache = ( forward_batch is not None and forward_batch.enable_teacache ) orig_dtype = hidden_states.dtype if not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] if ( isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0 ): encoder_hidden_states_image = encoder_hidden_states_image[0] else: encoder_hidden_states_image = None batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.patch_size post_patch_num_frames = num_frames // p_t post_patch_height = height // p_h post_patch_width = width // p_w if not sequence_shard_enabled: # The rotary embedding layer correctly handles SP offsets internally. freqs_cos, freqs_sin = self.rotary_emb.forward_from_grid( ( post_patch_num_frames * self.sp_size, post_patch_height, post_patch_width, ), shard_dim=0, start_frame=0, device=hidden_states.device, ) assert freqs_cos.dtype == torch.float32 assert freqs_cos.device == hidden_states.device freqs_cis = ( (freqs_cos.float(), freqs_sin.float()) if freqs_cos is not None else None ) hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) # shape is [B, T' * H' * W', C] seq_len_orig = hidden_states.shape[1] seq_shard_pad = 0 if sequence_shard_enabled: if seq_len_orig % self.sp_size != 0: seq_shard_pad = self.sp_size - (seq_len_orig % self.sp_size) pad = torch.zeros( (batch_size, seq_shard_pad, hidden_states.shape[2]), dtype=hidden_states.dtype, device=hidden_states.device, ) hidden_states = torch.cat([hidden_states, pad], dim=1) sp_rank = get_sp_group().rank_in_group local_seq_len = hidden_states.shape[1] // self.sp_size hidden_states = hidden_states.view( batch_size, self.sp_size, local_seq_len, hidden_states.shape[2] ) hidden_states = hidden_states[:, sp_rank, :, :] frame_stride = post_patch_height * post_patch_width freqs_cos, freqs_sin = self._compute_rope_for_sequence_shard( local_seq_len, sp_rank, frame_stride, post_patch_width, hidden_states.device, ) freqs_cis = (freqs_cos.float(), freqs_sin.float()) # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v) if timestep.dim() == 2: # ti2v ts_seq_len = timestep.shape[1] timestep = timestep.flatten() # batch_size * seq_len else: ts_seq_len = None temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = ( self.condition_embedder( timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len, ) ) if ts_seq_len is not None: # batch_size, seq_len, 6, inner_dim timestep_proj = timestep_proj.unflatten(2, (6, -1)) else: # batch_size, 6, inner_dim timestep_proj = timestep_proj.unflatten(1, (6, -1)) if sequence_shard_enabled and ts_seq_len is not None: if seq_shard_pad > 0: pad = torch.zeros( ( batch_size, seq_shard_pad, timestep_proj.shape[2], timestep_proj.shape[3], ), dtype=timestep_proj.dtype, device=timestep_proj.device, ) timestep_proj = torch.cat([timestep_proj, pad], dim=1) timestep_proj = timestep_proj.view( batch_size, self.sp_size, local_seq_len, timestep_proj.shape[2], timestep_proj.shape[3], ) timestep_proj = timestep_proj[:, sp_rank, :, :, :] if encoder_hidden_states_image is not None: encoder_hidden_states = torch.concat( [encoder_hidden_states_image, encoder_hidden_states], dim=1 ) encoder_hidden_states = ( encoder_hidden_states.to(orig_dtype) if not current_platform.is_amp_supported() else encoder_hidden_states ) # cast to orig_dtype for MPS assert encoder_hidden_states.dtype == orig_dtype # 4. Transformer blocks # if caching is enabled, we might be able to skip the forward pass should_skip_forward = self.should_skip_forward_for_cached_states( timestep_proj=timestep_proj, temb=temb ) if should_skip_forward: hidden_states = self.retrieve_cached_states(hidden_states) else: # if teacache is enabled, we need to cache the original hidden states if self.enable_teacache: original_hidden_states = hidden_states.clone() for block in self.blocks: hidden_states = block( hidden_states, encoder_hidden_states, timestep_proj, freqs_cis ) # if teacache is enabled, we need to cache the original hidden states if self.enable_teacache: self.maybe_cache_states(hidden_states, original_hidden_states) self.cnt += 1 if sequence_shard_enabled: hidden_states = hidden_states.contiguous() hidden_states = sequence_model_parallel_all_gather(hidden_states, dim=1) if seq_shard_pad > 0: hidden_states = hidden_states[:, :seq_len_orig, :] # 5. Output norm, projection & unpatchify if temb.dim() == 3: # batch_size, seq_len, inner_dim (wan 2.2 ti2v) shift, scale = ( self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2) ).chunk(2, dim=2) shift = shift.squeeze(2) scale = scale.squeeze(2) else: # batch_size, inner_dim shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states, _ = self.proj_out(hidden_states) hidden_states = hidden_states.reshape( batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1, ) hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) return output def maybe_cache_states( self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor ) -> None: """Cache residual with CFG positive/negative separation.""" residual = hidden_states.squeeze(0) - original_hidden_states if not self.is_cfg_negative: self.previous_residual = residual else: self.previous_residual_negative = residual def should_skip_forward_for_cached_states(self, **kwargs) -> bool: if not self.enable_teacache: return False ctx = self._get_teacache_context() if ctx is None: return False # Wan uses WanTeaCacheParams with additional fields teacache_params = ctx.teacache_params assert isinstance( teacache_params, WanTeaCacheParams ), "teacache_params is not a WanTeaCacheParams" # Initialize Wan-specific parameters use_ret_steps = teacache_params.use_ret_steps cutoff_steps = teacache_params.get_cutoff_steps(ctx.num_inference_steps) ret_steps = teacache_params.ret_steps # Adjust ret_steps and cutoff_steps for non-CFG mode # (WanTeaCacheParams uses *2 factor assuming CFG) if not ctx.do_cfg: ret_steps = ret_steps // 2 cutoff_steps = cutoff_steps // 2 timestep_proj = kwargs["timestep_proj"] temb = kwargs["temb"] modulated_inp = timestep_proj if use_ret_steps else temb self.is_cfg_negative = ctx.is_cfg_negative # Wan uses ret_steps/cutoff_steps for boundary detection is_boundary_step = self.cnt < ret_steps or self.cnt >= cutoff_steps # Use shared helper to compute cache decision should_calc = self._compute_teacache_decision( modulated_inp=modulated_inp, is_boundary_step=is_boundary_step, coefficients=ctx.coefficients, teacache_thresh=ctx.teacache_thresh, ) return not should_calc def retrieve_cached_states(self, hidden_states: torch.Tensor) -> torch.Tensor: """Retrieve cached residual with CFG positive/negative separation.""" if not self.is_cfg_negative: return hidden_states + self.previous_residual else: return hidden_states + self.previous_residual_negative EntryClass = WanTransformer3DModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/dits/zimage.py ================================================ import math from typing import Any, List, Optional, Tuple import torch import torch.nn as nn from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig from sglang.multimodal_gen.runtime.distributed import get_tp_world_size from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, ReplicatedLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( NunchakuConfig, is_nunchaku_available, ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( _apply_rotary_emb, apply_flashinfer_rope_qk_inplace, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger try: from nunchaku.models.attention import NunchakuFeedForward # type: ignore[import] except Exception: NunchakuFeedForward = None logger = init_logger(__name__) _is_cuda = current_platform.is_cuda() ADALN_EMBED_DIM = 256 SEQ_MULTI_OF = 32 class SelectFirstElement(nn.Module): def __init__(self): super().__init__() def forward(self, x): return x[0] class TimestepEmbedder(nn.Module): def __init__(self, out_size, mid_size=None, frequency_embedding_size=256): super().__init__() if mid_size is None: mid_size = out_size self.mlp = nn.ModuleList( [ ColumnParallelLinear( frequency_embedding_size, mid_size, bias=True, gather_output=False ), nn.SiLU(), RowParallelLinear( mid_size, out_size, bias=True, input_is_parallel=True ), ] ) self.frequency_embedding_size = frequency_embedding_size @staticmethod def timestep_embedding(t, dim, max_period=10000): with torch.amp.autocast(current_platform.device_type, enabled=False): half = dim // 2 freqs = torch.exp( -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half ) args = t[:, None].float() * freqs[None] embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) if dim % 2: embedding = torch.cat( [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 ) return embedding def forward(self, t): t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to( self.mlp[0].weight.dtype ) t_emb, _ = self.mlp[0](t_freq) t_emb = self.mlp[1](t_emb) t_emb, _ = self.mlp[2](t_emb) return t_emb class FeedForward(nn.Module): def __init__(self, dim: int, hidden_dim: int): super().__init__() # Use MergedColumnParallelLinear for gate and up projection (fused) self.w13 = MergedColumnParallelLinear( dim, [hidden_dim, hidden_dim], bias=False, gather_output=False ) self.w2 = RowParallelLinear(hidden_dim, dim, bias=False, input_is_parallel=True) self.act = SiluAndMul() def forward(self, x): x13, _ = self.w13(x) x = self.act(x13) out, _ = self.w2(x) return out class ZImageAttention(nn.Module): def __init__( self, dim: int, num_heads: int, num_kv_heads: int, qk_norm: bool = True, eps: float = 1e-6, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.dim = dim self.head_dim = dim // num_heads self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.qk_norm = qk_norm tp_size = get_tp_world_size() assert ( num_heads % tp_size == 0 ), f"num_heads {num_heads} must be divisible by tp world size {tp_size}" assert ( num_kv_heads % tp_size == 0 ), f"num_kv_heads {num_kv_heads} must be divisible by tp world size {tp_size}" self.local_num_heads = num_heads // tp_size self.local_num_kv_heads = num_kv_heads // tp_size kv_dim = self.head_dim * num_kv_heads self.use_fused_qkv = isinstance(quant_config, NunchakuConfig) if self.use_fused_qkv: self.to_qkv = MergedColumnParallelLinear( dim, [dim, kv_dim, kv_dim], bias=False, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.to_qkv", ) else: self.to_q = ColumnParallelLinear( dim, dim, bias=False, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.to_q", ) self.to_k = ColumnParallelLinear( dim, kv_dim, bias=False, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.to_k", ) self.to_v = ColumnParallelLinear( dim, kv_dim, bias=False, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.to_v", ) if self.qk_norm: self.norm_q = RMSNorm(self.head_dim, eps=eps) self.norm_k = RMSNorm(self.head_dim, eps=eps) else: self.norm_q = None self.norm_k = None self.to_out = nn.ModuleList( [ RowParallelLinear( dim, dim, bias=False, input_is_parallel=True, quant_config=quant_config, prefix=f"{prefix}.to_out.0", ) ] ) self.attn = USPAttention( num_heads=self.local_num_heads, head_size=self.head_dim, num_kv_heads=self.local_num_kv_heads, dropout_rate=0, softmax_scale=None, causal=False, ) def forward( self, hidden_states: torch.Tensor, freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ): if self.use_fused_qkv: qkv, _ = self.to_qkv(hidden_states) q, k, v = qkv.split( [ self.local_num_heads * self.head_dim, self.local_num_kv_heads * self.head_dim, self.local_num_kv_heads * self.head_dim, ], dim=-1, ) q = q.contiguous() k = k.contiguous() v = v.contiguous() else: q, _ = self.to_q(hidden_states) k, _ = self.to_k(hidden_states) v, _ = self.to_v(hidden_states) q = q.view(*q.shape[:-1], self.local_num_heads, self.head_dim) k = k.view(*k.shape[:-1], self.local_num_kv_heads, self.head_dim) v = v.view(*v.shape[:-1], self.local_num_kv_heads, self.head_dim) if self.qk_norm: q, k = apply_qk_norm( q=q, k=k, q_norm=self.norm_q, k_norm=self.norm_k, head_dim=self.head_dim, allow_inplace=True, ) if freqs_cis is not None: cos, sin = freqs_cis if _is_cuda and q.shape == k.shape: cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(), sin.to(dtype=torch.float32).contiguous(), ], dim=-1, ) q, k = apply_flashinfer_rope_qk_inplace( q, k, cos_sin_cache, is_neox=False ) else: q = _apply_rotary_emb(q, cos, sin, is_neox_style=False) k = _apply_rotary_emb(k, cos, sin, is_neox_style=False) hidden_states = self.attn(q, k, v) hidden_states = hidden_states.flatten(2) hidden_states, _ = self.to_out[0](hidden_states) return hidden_states class ZImageTransformerBlock(nn.Module): def __init__( self, layer_id: int, dim: int, n_heads: int, n_kv_heads: int, norm_eps: float, qk_norm: bool, modulation=True, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.dim = dim self.head_dim = dim // n_heads self.layer_id = layer_id self.modulation = modulation self.attention = ZImageAttention( dim=dim, num_heads=n_heads, num_kv_heads=n_kv_heads, qk_norm=qk_norm, eps=1e-5, quant_config=quant_config, prefix=f"{prefix}.attention", ) hidden_dim = int(dim / 3 * 8) nunchaku_enabled = ( isinstance(quant_config, NunchakuConfig) and is_nunchaku_available() ) if nunchaku_enabled: import diffusers ff = diffusers.models.attention.FeedForward( dim=dim, dim_out=dim, activation_fn="swiglu", inner_dim=hidden_dim, bias=False, ) nunchaku_kwargs = { "precision": quant_config.precision, "rank": quant_config.rank, "act_unsigned": quant_config.act_unsigned, } self.feed_forward = NunchakuFeedForward(ff, **nunchaku_kwargs) # NunchakuFeedForward overrides net[2].act_unsigned=True for int4 (GELU-specific # optimization for non-negative activations). Z-Image uses SwiGLU whose output # can be negative, so we must restore the original act_unsigned value. if hasattr(self.feed_forward, "net") and len(self.feed_forward.net) > 2: self.feed_forward.net[2].act_unsigned = quant_config.act_unsigned else: self.feed_forward = FeedForward(dim=dim, hidden_dim=hidden_dim) self.attention_norm1 = RMSNorm(dim, eps=norm_eps) self.ffn_norm1 = RMSNorm(dim, eps=norm_eps) self.attention_norm2 = RMSNorm(dim, eps=norm_eps) self.ffn_norm2 = RMSNorm(dim, eps=norm_eps) if modulation: self.adaLN_modulation = nn.Sequential( ReplicatedLinear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True) ) def forward( self, x: torch.Tensor, freqs_cis: Tuple[torch.Tensor, torch.Tensor], adaln_input: Optional[torch.Tensor] = None, ): if self.modulation: assert adaln_input is not None scale_msa_gate, _ = self.adaLN_modulation(adaln_input) scale_msa, gate_msa, scale_mlp, gate_mlp = scale_msa_gate.unsqueeze( 1 ).chunk(4, dim=2) gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp # Attention block attn_out = self.attention( self.attention_norm1(x) * scale_msa, freqs_cis=freqs_cis, ) x = x + gate_msa * self.attention_norm2(attn_out) # FFN block x = x + gate_mlp * self.ffn_norm2( self.feed_forward( self.ffn_norm1(x) * scale_mlp, ) ) else: # Attention block attn_out = self.attention( self.attention_norm1(x), freqs_cis=freqs_cis, ) x = x + self.attention_norm2(attn_out) # FFN block x = x + self.ffn_norm2( self.feed_forward( self.ffn_norm1(x), ) ) return x class FinalLayer(nn.Module): def __init__(self, hidden_size, out_channels): super().__init__() self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.linear = ColumnParallelLinear( hidden_size, out_channels, bias=True, gather_output=True ) self.act = nn.SiLU() self.adaLN_modulation = nn.Sequential( nn.SiLU(), ReplicatedLinear(min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True), ) def forward(self, x, c): scale, _ = self.adaLN_modulation(c) scale = 1.0 + scale x = self.norm_final(x) * scale.unsqueeze(1) x, _ = self.linear(x) return x class RopeEmbedder: def __init__( self, theta: float = 256.0, axes_dims: List[int] = (16, 56, 56), axes_lens: List[int] = (64, 128, 128), ): self.theta = theta self.axes_dims = axes_dims self.axes_lens = axes_lens assert len(axes_dims) == len( axes_lens ), "axes_dims and axes_lens must have the same length" self.cos_cached = None self.sin_cached = None @staticmethod def precompute_freqs(dim: List[int], end: List[int], theta: float = 256.0): with torch.device("cpu"): cos_list = [] sin_list = [] for i, (d, e) in enumerate(zip(dim, end)): freqs = 1.0 / ( theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d) ) timestep = torch.arange(e, device=freqs.device, dtype=torch.float64) freqs = torch.outer(timestep, freqs).float() cos_list.append(torch.cos(freqs)) sin_list.append(torch.sin(freqs)) return cos_list, sin_list def __call__(self, ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """ Args: ids: [batch, len(axes_dims)] or [seq_len, len(axes_dims)] Returns: cos: [batch/seq, head_dim // 2] sin: [batch/seq, head_dim // 2] """ assert ids.ndim == 2 assert ids.shape[-1] == len(self.axes_dims) device = ids.device if self.cos_cached is None: self.cos_cached, self.sin_cached = self.precompute_freqs( self.axes_dims, self.axes_lens, theta=self.theta ) self.cos_cached = [c.to(device) for c in self.cos_cached] self.sin_cached = [s.to(device) for s in self.sin_cached] else: if self.cos_cached[0].device != device: self.cos_cached = [c.to(device) for c in self.cos_cached] self.sin_cached = [s.to(device) for s in self.sin_cached] cos_out = [] sin_out = [] for i in range(len(self.axes_dims)): index = ids[:, i] cos_out.append(self.cos_cached[i][index]) sin_out.append(self.sin_cached[i][index]) return torch.cat(cos_out, dim=-1), torch.cat(sin_out, dim=-1) class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin): _supports_gradient_checkpointing = True _no_split_modules = ["ZImageTransformerBlock"] _fsdp_shard_conditions = ZImageDitConfig().arch_config._fsdp_shard_conditions param_names_mapping = ZImageDitConfig().arch_config.param_names_mapping param_names_mapping = ZImageDitConfig().arch_config.param_names_mapping reverse_param_names_mapping = ( ZImageDitConfig().arch_config.reverse_param_names_mapping ) @classmethod def get_nunchaku_quant_rules(cls) -> dict[str, list[str]]: return { "skip": [ "norm", "embed", "rotary", "pos_embed", ], "svdq_w4a4": [ "attention.to_qkv", "attention.to_out", "img_mlp", "txt_mlp", ], "awq_w4a16": [ "img_mod", "txt_mod", ], } def __init__( self, config: ZImageDitConfig, hf_config: dict[str, Any], quant_config: Optional[QuantizationConfig] = None, ) -> None: super().__init__(config=config, hf_config=hf_config) self.config_data = config # Store config arch_config = config.arch_config self.in_channels = arch_config.in_channels self.out_channels = arch_config.out_channels self.all_patch_size = arch_config.all_patch_size self.all_f_patch_size = arch_config.all_f_patch_size self.dim = arch_config.dim self.n_heads = arch_config.num_attention_heads self.rope_theta = arch_config.rope_theta self.t_scale = arch_config.t_scale self.gradient_checkpointing = False assert len(self.all_patch_size) == len(self.all_f_patch_size) all_x_embedder = {} all_final_layer = {} for patch_idx, (patch_size, f_patch_size) in enumerate( zip(self.all_patch_size, self.all_f_patch_size) ): x_embedder = ColumnParallelLinear( f_patch_size * patch_size * patch_size * self.in_channels, self.dim, bias=True, gather_output=True, ) all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder final_layer = FinalLayer( self.dim, patch_size * patch_size * f_patch_size * self.out_channels ) all_final_layer[f"{patch_size}-{f_patch_size}"] = final_layer self.all_x_embedder = nn.ModuleDict(all_x_embedder) self.all_final_layer = nn.ModuleDict(all_final_layer) self.noise_refiner = nn.ModuleList( [ ZImageTransformerBlock( 1000 + layer_id, self.dim, self.n_heads, arch_config.n_kv_heads, arch_config.norm_eps, arch_config.qk_norm, modulation=True, quant_config=quant_config, prefix=f"noise_refiner.{layer_id}", ) for layer_id in range(arch_config.n_refiner_layers) ] ) self.context_refiner = nn.ModuleList( [ ZImageTransformerBlock( layer_id, self.dim, self.n_heads, arch_config.n_kv_heads, arch_config.norm_eps, arch_config.qk_norm, modulation=False, quant_config=quant_config, prefix=f"context_refiner.{layer_id}", ) for layer_id in range(arch_config.n_refiner_layers) ] ) self.t_embedder = TimestepEmbedder( min(self.dim, ADALN_EMBED_DIM), mid_size=1024 ) self.cap_embedder = nn.Sequential( RMSNorm(arch_config.cap_feat_dim, eps=arch_config.norm_eps), ReplicatedLinear(arch_config.cap_feat_dim, self.dim, bias=True), ) self.x_pad_token = nn.Parameter(torch.empty((1, self.dim))) self.cap_pad_token = nn.Parameter(torch.empty((1, self.dim))) self.layers = nn.ModuleList( [ ZImageTransformerBlock( layer_id, self.dim, self.n_heads, arch_config.n_kv_heads, arch_config.norm_eps, arch_config.qk_norm, quant_config=quant_config, prefix=f"layers.{layer_id}", ) for layer_id in range(arch_config.num_layers) ] ) head_dim = self.dim // self.n_heads assert head_dim == sum(arch_config.axes_dims) self.axes_dims = arch_config.axes_dims self.axes_lens = arch_config.axes_lens self.rotary_emb = RopeEmbedder( theta=self.rope_theta, axes_dims=self.axes_dims, axes_lens=self.axes_lens ) self.layer_names = ["layers"] def unpatchify( self, x: List[torch.Tensor], size: List[Tuple], patch_size, f_patch_size ) -> List[torch.Tensor]: pH = pW = patch_size pF = f_patch_size bsz = len(x) assert len(size) == bsz for i in range(bsz): F, H, W = size[i] ori_len = (F // pF) * (H // pH) * (W // pW) # "f h w pf ph pw c -> c (f pf) (h ph) (w pw)" x[i] = ( x[i][:ori_len] .view(F // pF, H // pH, W // pW, pF, pH, pW, self.out_channels) .permute(6, 0, 3, 1, 4, 2, 5) .reshape(self.out_channels, F, H, W) ) return x @staticmethod def create_coordinate_grid(size, start=None, device=None): if start is None: start = (0 for _ in size) axes = [ torch.arange(x0, x0 + span, dtype=torch.int32, device=device) for x0, span in zip(start, size) ] grids = torch.meshgrid(axes, indexing="ij") return torch.stack(grids, dim=-1) def patchify_and_embed( self, all_image: List[torch.Tensor], all_cap_feats: List[torch.Tensor], patch_size: int, f_patch_size: int, ): assert len(all_image) == len(all_cap_feats) == 1 image = all_image[0] # C, F, H, W cap_feat = all_cap_feats[0] # L, D pH = pW = patch_size pF = f_patch_size device = image.device all_image_out = [] all_image_size = [] all_cap_feats_out = [] # ------------ Process Caption ------------ cap_ori_len = cap_feat.size(0) cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF # padded feature cap_padded_feat = torch.cat( [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)], dim=0, ) all_cap_feats_out.append(cap_padded_feat) # ------------ Process Image ------------ C, F, H, W = image.size() all_image_size.append((F, H, W)) F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW) # "c f pf h ph w pw -> (f h w) (pf ph pw c)" image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape( F_tokens * H_tokens * W_tokens, pF * pH * pW * C ) image_ori_len = image.size(0) image_padding_len = (-image_ori_len) % SEQ_MULTI_OF # padded feature image_padded_feat = torch.cat( [image, image[-1:].repeat(image_padding_len, 1)], dim=0, ) all_image_out.append(image_padded_feat) return ( all_image_out, all_cap_feats_out, all_image_size, ) def forward( self, hidden_states: List[torch.Tensor], encoder_hidden_states: List[torch.Tensor], timestep, guidance=0, patch_size=2, f_patch_size=1, freqs_cis=None, **kwargs, ): assert patch_size in self.all_patch_size assert f_patch_size in self.all_f_patch_size x = hidden_states cap_feats = encoder_hidden_states timestep = 1000.0 - timestep t = timestep bsz = 1 device = x[0].device t = self.t_embedder(t) adaln_input = t.type_as(x) ( x, cap_feats, x_size, ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) x = torch.cat(x, dim=0) x, _ = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x) x_freqs_cis = freqs_cis[1] x = x.unsqueeze(0) x_freqs_cis = x_freqs_cis for layer in self.noise_refiner: x = layer(x, x_freqs_cis, adaln_input) cap_feats = torch.cat(cap_feats, dim=0) cap_feats, _ = self.cap_embedder(cap_feats) cap_freqs_cis = freqs_cis[0] cap_feats = cap_feats.unsqueeze(0) for layer in self.context_refiner: cap_feats = layer(cap_feats, cap_freqs_cis) unified = torch.cat([x, cap_feats], dim=1) unified_freqs_cis = ( torch.cat([x_freqs_cis[0], cap_freqs_cis[0]], dim=0), torch.cat([x_freqs_cis[1], cap_freqs_cis[1]], dim=0), ) for layer in self.layers: unified = layer(unified, unified_freqs_cis, adaln_input) unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"]( unified, adaln_input ) unified = list(unified.unbind(dim=0)) x = self.unpatchify(unified, x_size, patch_size, f_patch_size) return -x[0] EntryClass = ZImageTransformer2DModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/encoders/base.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from abc import ABC, abstractmethod from dataclasses import field import torch from torch import nn from sglang.multimodal_gen.configs.models.encoders import ( BaseEncoderOutput, ImageEncoderConfig, TextEncoderConfig, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum class TextEncoder(nn.Module, ABC): _fsdp_shard_conditions: list = field(default_factory=lambda: []) _stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list) _supported_attention_backends: set[AttentionBackendEnum] = ( TextEncoderConfig()._supported_attention_backends ) def __init__(self, config: TextEncoderConfig) -> None: super().__init__() self.config = config self._fsdp_shard_conditions = config.arch_config._fsdp_shard_conditions self._stacked_params_mapping = config.arch_config.stacked_params_mapping if not self.supported_attention_backends: raise ValueError( f"Subclass {self.__class__.__name__} must define _supported_attention_backends" ) @abstractmethod def forward( self, input_ids: torch.Tensor | None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs, ) -> BaseEncoderOutput: pass @property def supported_attention_backends(self) -> set[AttentionBackendEnum]: return self._supported_attention_backends class ImageEncoder(nn.Module, ABC): _supported_attention_backends: set[AttentionBackendEnum] = ( ImageEncoderConfig()._supported_attention_backends ) def __init__(self, config: ImageEncoderConfig) -> None: super().__init__() self.config = config if not self.supported_attention_backends: raise ValueError( f"Subclass {self.__class__.__name__} must define _supported_attention_backends" ) @abstractmethod def forward(self, pixel_values: torch.Tensor, **kwargs) -> BaseEncoderOutput: pass @property def supported_attention_backends(self) -> set[AttentionBackendEnum]: return self._supported_attention_backends ================================================ FILE: python/sglang/multimodal_gen/runtime/models/encoders/bert.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # type: ignore import os import torch import torch.nn as nn from transformers import BertModel, BertTokenizer class HunyuanClip(nn.Module): """ Hunyuan clip code copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py hunyuan's clip used BertModel and BertTokenizer, so we copy it. """ def __init__(self, model_dir, max_length=77): super().__init__() self.max_length = max_length self.tokenizer = BertTokenizer.from_pretrained( os.path.join(model_dir, "tokenizer") ) self.text_encoder = BertModel.from_pretrained( os.path.join(model_dir, "clip_text_encoder") ) @torch.no_grad def forward(self, prompts, with_mask=True): self.device = next(self.text_encoder.parameters()).device text_inputs = self.tokenizer( prompts, padding="max_length", max_length=self.max_length, truncation=True, return_attention_mask=True, return_tensors="pt", ) prompt_embeds = self.text_encoder( text_inputs.input_ids.to(self.device), attention_mask=( text_inputs.attention_mask.to(self.device) if with_mask else None ), ) return prompt_embeds.last_hidden_state, prompt_embeds.pooler_output ================================================ FILE: python/sglang/multimodal_gen/runtime/models/encoders/clip.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/clip.py # Adapted from transformers: https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/clip/modeling_clip.py """Minimal implementation of CLIPVisionModel intended to be only used within a vision language model.""" from collections.abc import Iterable from typing import Optional import torch import torch.nn as nn from sglang.multimodal_gen.configs.models.encoders import ( BaseEncoderOutput, CLIPTextConfig, CLIPVisionConfig, ) from sglang.multimodal_gen.runtime.distributed import divide, get_tp_world_size from sglang.multimodal_gen.runtime.layers.activation import get_act_fn from sglang.multimodal_gen.runtime.layers.attention import LocalAttention from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, QKVParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig # TODO: support quantization # from vllm.model_executor.layers.quantization import QuantizationConfig from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader from sglang.multimodal_gen.runtime.models.encoders.base import ImageEncoder, TextEncoder from sglang.multimodal_gen.runtime.models.encoders.vision import ( resolve_visual_encoder_outputs, ) from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # Adapted from https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/clip/modeling_clip.py#L164 # noqa class CLIPVisionEmbeddings(nn.Module): def __init__(self, config: CLIPVisionConfig): super().__init__() self.config = config self.embed_dim = config.hidden_size self.image_size = config.image_size self.patch_size = config.patch_size assert self.image_size % self.patch_size == 0 self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) self.patch_embedding = nn.Conv2d( in_channels=config.num_channels, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size, bias=False, ) self.num_patches = (self.image_size // self.patch_size) ** 2 self.num_positions = self.num_patches + 1 self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) self.register_buffer( "position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False, ) def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: batch_size = pixel_values.shape[0] target_dtype = self.patch_embedding.weight.dtype patch_embeds = self.patch_embedding( pixel_values.to(dtype=target_dtype) ) # shape = [*, width, grid, grid] patch_embeds = patch_embeds.flatten(2).transpose(1, 2) class_embeds = self.class_embedding.expand(batch_size, 1, -1) embeddings = torch.cat([class_embeds, patch_embeds], dim=1) embeddings = embeddings + self.position_embedding(self.position_ids) return embeddings class CLIPTextEmbeddings(nn.Module): def __init__(self, config: CLIPTextConfig): super().__init__() self.config = config embed_dim = config.hidden_size self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) self.position_embedding = nn.Embedding( config.max_position_embeddings, embed_dim ) # position_ids (1, len position emb) is contiguous in memory and exported when serialized self.register_buffer( "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False, ) def forward( self, input_ids: torch.LongTensor | None = None, position_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, ) -> torch.Tensor: if input_ids is not None: seq_length = input_ids.shape[-1] elif inputs_embeds is not None: seq_length = inputs_embeds.shape[-2] else: raise ValueError("Either input_ids or inputs_embeds must be provided.") max_position_embedding = self.position_embedding.weight.shape[0] if seq_length > max_position_embedding: raise ValueError( f"Sequence length must be less than max_position_embeddings (got `sequence length`: " f"{seq_length} and max_position_embeddings: {max_position_embedding}" ) if position_ids is None: position_ids = self.position_ids[:, :seq_length] if inputs_embeds is None: inputs_embeds = self.token_embedding(input_ids) position_embeddings = self.position_embedding(position_ids) embeddings = inputs_embeds + position_embeddings return embeddings class CLIPAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__( self, config: CLIPVisionConfig | CLIPTextConfig, quant_config: QuantizationConfig | None = None, prefix: str = "", ): super().__init__() self.config = config self.embed_dim = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.embed_dim // self.num_heads if self.head_dim * self.num_heads != self.embed_dim: raise ValueError( "embed_dim must be divisible by num_heads " f"(got `embed_dim`: {self.embed_dim} and `num_heads`:" f" {self.num_heads})." ) self.scale = self.head_dim**-0.5 self.dropout = config.attention_dropout self.qkv_proj = QKVParallelLinear( hidden_size=self.embed_dim, head_size=self.head_dim, total_num_heads=self.num_heads, quant_config=quant_config, prefix=f"{prefix}.qkv_proj", ) self.out_proj = RowParallelLinear( input_size=self.embed_dim, output_size=self.embed_dim, quant_config=quant_config, prefix=f"{prefix}.out_proj", ) self.tp_size = get_tp_world_size() self.num_heads_per_partition = divide(self.num_heads, self.tp_size) self.attn = LocalAttention( self.num_heads_per_partition, self.head_dim, self.num_heads_per_partition, softmax_scale=self.scale, causal=True, supported_attention_backends=config._supported_attention_backends, ) def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): return ( tensor.view(bsz, seq_len, self.num_heads, self.head_dim) .transpose(1, 2) .contiguous() ) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, ): """Input shape: Batch x Time x Channel""" qkv_states, _ = self.qkv_proj(hidden_states) query_states, key_states, value_states = qkv_states.chunk(3, dim=-1) # use flash_attn_func query_states = query_states.reshape( query_states.shape[0], query_states.shape[1], self.num_heads_per_partition, self.head_dim, ) key_states = key_states.reshape( key_states.shape[0], key_states.shape[1], self.num_heads_per_partition, self.head_dim, ) value_states = value_states.reshape( value_states.shape[0], value_states.shape[1], self.num_heads_per_partition, self.head_dim, ) if self.attn.backend == AttentionBackendEnum.TORCH_SDPA: query_states = query_states.transpose(1, 2) # [B, H, S, D] key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) if current_platform.is_rocm() or current_platform.is_musa(): # ROCm: Using both is_causal=True and attn_mask causes NaN. # Use is_causal=True alone (padding mask not needed for CLIP # since pooler_output comes from EOS token before padding). # XXX (MUSA): Torch SDPA on MUSA currently does not support # using both `attn_mask` and `is_causal=True` simultaneously. attn_output = torch.nn.functional.scaled_dot_product_attention( query_states, key_states, value_states, attn_mask=None, is_causal=True, scale=self.scale, ) else: if attention_mask is not None: # SDPA requires [B, 1, 1, S] or [B, S, S] format mask if attention_mask.dim() == 2: attn_mask = attention_mask[:, None, None, :].to( dtype=query_states.dtype ) attn_mask = (1.0 - attn_mask) * torch.finfo( query_states.dtype ).min else: attn_mask = attention_mask else: attn_mask = None attn_output = torch.nn.functional.scaled_dot_product_attention( query_states, key_states, value_states, attn_mask=attn_mask, is_causal=attention_mask is None, scale=self.scale, ) attn_output = attn_output.transpose(1, 2) else: # Use LocalAttention (doesn't support attention_mask, but maintains compatibility) attn_output = self.attn(query_states, key_states, value_states) attn_output = attn_output.reshape( attn_output.shape[0], attn_output.shape[1], self.num_heads_per_partition * self.head_dim, ) attn_output, _ = self.out_proj(attn_output) return attn_output, None class CLIPMLP(nn.Module): def __init__( self, config: CLIPVisionConfig | CLIPTextConfig, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() self.config = config self.activation_fn = get_act_fn(config.hidden_act) self.fc1 = ColumnParallelLinear( config.hidden_size, config.intermediate_size, bias=True, quant_config=quant_config, prefix=f"{prefix}.fc1", ) self.fc2 = RowParallelLinear( config.intermediate_size, config.hidden_size, bias=True, quant_config=quant_config, prefix=f"{prefix}.fc2", ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states, _ = self.fc1(hidden_states) hidden_states = self.activation_fn(hidden_states) hidden_states, _ = self.fc2(hidden_states) return hidden_states class CLIPEncoderLayer(nn.Module): def __init__( self, config: CLIPTextConfig | CLIPVisionConfig, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() self.self_attn = CLIPAttention( config, quant_config=quant_config, prefix=f"{prefix}.self_attn", ) self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.mlp = CLIPMLP(config, quant_config=quant_config, prefix=f"{prefix}.mlp") self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: residual = hidden_states hidden_states = self.layer_norm1(hidden_states) hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.layer_norm2(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states class CLIPEncoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`CLIPEncoderLayer`]. Args: config: CLIPConfig """ def __init__( self, config: CLIPVisionConfig | CLIPTextConfig, quant_config: QuantizationConfig | None = None, num_hidden_layers_override: int | None = None, prefix: str = "", ) -> None: super().__init__() self.config = config if num_hidden_layers_override is None: num_hidden_layers = config.num_hidden_layers else: num_hidden_layers = num_hidden_layers_override self.layers = nn.ModuleList( [ CLIPEncoderLayer( config=config, quant_config=quant_config, prefix=f"{prefix}.layers.{layer_idx}", ) for layer_idx in range(num_hidden_layers) ] ) def forward( self, inputs_embeds: torch.Tensor, return_all_hidden_states: bool, attention_mask: torch.Tensor | None = None, ) -> torch.Tensor | list[torch.Tensor]: hidden_states_pool = [inputs_embeds] hidden_states = inputs_embeds for idx, encoder_layer in enumerate(self.layers): hidden_states = encoder_layer( hidden_states, attention_mask=attention_mask, ) if return_all_hidden_states: hidden_states_pool.append(hidden_states) # If we have multiple feature sample layers, we return all hidden # states in order and grab the ones we need by index. if return_all_hidden_states: return hidden_states_pool return [hidden_states] class CLIPTextTransformer(nn.Module): def __init__( self, config: CLIPTextConfig, quant_config: QuantizationConfig | None = None, num_hidden_layers_override: int | None = None, prefix: str = "", ): super().__init__() self.config = config embed_dim = config.hidden_size self.embeddings = CLIPTextEmbeddings(config) self.encoder = CLIPEncoder( config, quant_config=quant_config, num_hidden_layers_override=num_hidden_layers_override, prefix=prefix, ) self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) # For `pooled_output` computation self.eos_token_id = config.eos_token_id def forward( self, input_ids: torch.Tensor | None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, ) -> BaseEncoderOutput: output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) if input_ids is None: raise ValueError("You have to specify input_ids") input_shape = input_ids.size() input_ids = input_ids.view(-1, input_shape[-1]) hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) # CLIP's text model uses causal mask, prepare it here. # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 # causal_attention_mask = _create_4d_causal_attention_mask( # input_shape, hidden_states.dtype, device=hidden_states.device # ) # # expand attention_mask # if attention_mask is not None and not self._use_flash_attention_2: # raise NotImplementedError("attention_mask is not supported for CLIPTextTransformer") # # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] # attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) encoder_outputs = self.encoder( inputs_embeds=hidden_states, return_all_hidden_states=output_hidden_states, attention_mask=attention_mask, ) last_hidden_state = encoder_outputs[-1] last_hidden_state = self.final_layer_norm(last_hidden_state) if self.eos_token_id == 2: # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here. # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added # ------------------------------------------------------------ # text_embeds.shape = [batch_size, sequence_length, transformer.width] # take features from the eot embedding (eot_token is the highest number in each sequence) # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14 pooled_output = last_hidden_state[ torch.arange( last_hidden_state.shape[0], device=last_hidden_state.device ), input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax( dim=-1 ), ] else: # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible) pooled_output = last_hidden_state[ torch.arange( last_hidden_state.shape[0], device=last_hidden_state.device ), # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`) # Note: we assume each sequence (along batch dim.) contains an `eos_token_id` (e.g. prepared by the tokenizer) ( input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.eos_token_id ) .int() .argmax(dim=-1), ] return BaseEncoderOutput( last_hidden_state=last_hidden_state, pooler_output=pooled_output, hidden_states=encoder_outputs, # attentions=encoder_outputs.attentions, ) class CLIPTextModel(TextEncoder): def __init__( self, config: CLIPTextConfig, ) -> None: super().__init__(config) self.text_model = CLIPTextTransformer( config=config, quant_config=config.quant_config, prefix=config.prefix ) def forward( self, input_ids: torch.Tensor | None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs, ) -> BaseEncoderOutput: outputs: BaseEncoderOutput = self.text_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, output_hidden_states=output_hidden_states, ) return outputs def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Define mapping for stacked parameters stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), ("qkv_proj", "k_proj", "k"), ("qkv_proj", "v_proj", "v"), ] params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: # Handle q_proj, k_proj, v_proj -> qkv_proj mapping for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name in name: # Replace the weight name with the parameter name model_param_name = name.replace(weight_name, param_name) if model_param_name in params_dict: param = params_dict[model_param_name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) loaded_params.add(model_param_name) break else: # Use default weight loader for all other parameters if name in params_dict: param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params class CLIPVisionTransformer(nn.Module): def __init__( self, config: CLIPVisionConfig, quant_config: QuantizationConfig | None = None, num_hidden_layers_override: int | None = None, require_post_norm: bool | None = None, prefix: str = "", ) -> None: super().__init__() self.config = config embed_dim = config.hidden_size self.embeddings = CLIPVisionEmbeddings(config) # NOTE: This typo of "layrnorm" is not fixed on purpose to match # the original transformers code and name of the model weights. self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) self.encoder = CLIPEncoder( config=config, quant_config=quant_config, num_hidden_layers_override=num_hidden_layers_override, prefix=f"{prefix}.encoder", ) num_hidden_layers = config.num_hidden_layers if len(self.encoder.layers) > config.num_hidden_layers: raise ValueError( f"The original encoder only has {num_hidden_layers} " f"layers, but you requested {len(self.encoder.layers)} layers." ) # If possible, skip post_layernorm to conserve memory if require_post_norm is None: require_post_norm = len(self.encoder.layers) == num_hidden_layers if require_post_norm: self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) else: self.post_layernorm = None def forward( self, pixel_values: torch.Tensor, output_hidden_states: Optional[bool] = None, feature_sample_layers: list[int] | None = None, ) -> BaseEncoderOutput: hidden_states = self.embeddings(pixel_values) hidden_states = self.pre_layrnorm(hidden_states) return_all_hidden_states = output_hidden_states or ( feature_sample_layers is not None ) # Produces either the last layer output or all of the hidden states, # depending on if we have feature_sample_layers or not encoder_outputs = self.encoder( inputs_embeds=hidden_states, return_all_hidden_states=return_all_hidden_states, ) if not return_all_hidden_states: encoder_outputs = encoder_outputs[0] # Handle post-norm (if applicable) and stacks feature layers if needed encoder_outputs = resolve_visual_encoder_outputs( encoder_outputs, feature_sample_layers, self.post_layernorm, self.config.num_hidden_layers, ) if return_all_hidden_states: return BaseEncoderOutput(hidden_states=encoder_outputs) return BaseEncoderOutput(last_hidden_state=encoder_outputs) class CLIPVisionModel(ImageEncoder): config_class = CLIPVisionConfig main_input_name = "pixel_values" packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} def __init__(self, config: CLIPVisionConfig) -> None: super().__init__(config) self.vision_model = CLIPVisionTransformer( config=config, quant_config=config.quant_config, num_hidden_layers_override=config.num_hidden_layers_override, require_post_norm=config.require_post_norm, prefix=f"{config.prefix}.vision_model", ) def forward( self, pixel_values: torch.Tensor, feature_sample_layers: list[int] | None = None, output_hidden_states: Optional[bool] = None, **kwargs, ) -> BaseEncoderOutput: base_encoder_output = self.vision_model( pixel_values, output_hidden_states=output_hidden_states, feature_sample_layers=feature_sample_layers, ) return base_encoder_output @property def device(self): return next(self.parameters()).device # (TODO) Add prefix argument for filtering out weights to be loaded # ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() layer_count = len(self.vision_model.encoder.layers) for name, loaded_weight in weights: if name.startswith("visual_projection"): continue # post_layernorm is not needed in CLIPVisionModel if ( name.startswith("vision_model.post_layernorm") and self.vision_model.post_layernorm is None ): continue # omit layers when num_hidden_layers_override is set if name.startswith("vision_model.encoder.layers"): layer_idx = int(name.split(".")[3]) if layer_idx >= layer_count: continue for ( param_name, weight_name, shard_id, ) in self.config.arch_config.stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) break else: param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params class BertModel(CLIPTextModel): pass EntryClass = [CLIPTextModel, CLIPVisionModel] ================================================ FILE: python/sglang/multimodal_gen/runtime/models/encoders/gemma2.py ================================================ # SPDX-License-Identifier: Apache-2.0 # # Gemma2 2B text encoder for SANA. # # This is a decoder-only language model used as a text encoder: we feed # in tokenized text and extract the final hidden states (not logits) as # the conditioning signal for SANA's cross-attention layers. # # Architecture follows google/gemma-2-2b-it: # - 26 layers, alternating global / sliding-window attention # - GQA with 8 query heads, 4 KV heads, head_dim=256 # - Pre/post attention + pre/post feedforward LayerNorm (Gemma2-style) # - GeGLU activation (gelu_pytorch_tanh) # # Adapted from the Gemma3 text model implementation in this codebase. import logging from typing import Any, Iterable import torch from torch import nn from sglang.multimodal_gen.configs.models.encoders.base import BaseEncoderOutput from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config from sglang.multimodal_gen.runtime.distributed import get_tp_world_size from sglang.multimodal_gen.runtime.layers.activation import GeluAndMul from sglang.multimodal_gen.runtime.layers.linear import ( MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig from sglang.multimodal_gen.runtime.layers.rotary_embedding import get_rope from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader logger = logging.getLogger(__name__) class Gemma2RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.zeros(dim)) def _norm(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) def forward(self, x): output = self._norm(x.float()) output = output * (1.0 + self.weight.float()) return output.type_as(x) class Gemma2MLP(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, hidden_act: str, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() self.gate_up_proj = MergedColumnParallelLinear( input_size=hidden_size, output_sizes=[intermediate_size] * 2, bias=False, quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", ) self.down_proj = RowParallelLinear( input_size=intermediate_size, output_size=hidden_size, bias=False, quant_config=quant_config, prefix=f"{prefix}.down_proj", ) if hidden_act != "gelu_pytorch_tanh": raise ValueError( "Gemma2 uses `gelu_pytorch_tanh` as the hidden activation. " f"Got: {hidden_act}" ) self.act_fn = GeluAndMul(approximate="tanh") def forward(self, x): x, _ = self.gate_up_proj(x) x = self.act_fn(x) x, _ = self.down_proj(x) return x class Gemma2Attention(nn.Module): def __init__( self, layer_id: int, config: Gemma2Config, hidden_size: int, num_heads: int, num_kv_heads: int, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() self.layer_id = layer_id self.hidden_size = hidden_size tp_size = get_tp_world_size() self.total_num_heads = num_heads assert self.total_num_heads % tp_size == 0 self.num_heads = self.total_num_heads // tp_size self.total_num_kv_heads = num_kv_heads if self.total_num_kv_heads >= tp_size: assert self.total_num_kv_heads % tp_size == 0 else: assert tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) arch = config.arch_config self.head_dim = arch.head_dim self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = arch.query_pre_attn_scalar**-0.5 self.qkv_proj = QKVParallelLinear( hidden_size=hidden_size, head_size=self.head_dim, total_num_heads=self.total_num_heads, total_num_kv_heads=self.total_num_kv_heads, bias=arch.attention_bias, quant_config=quant_config, prefix=f"{prefix}.qkv_proj", ) self.o_proj = RowParallelLinear( input_size=self.total_num_heads * self.head_dim, output_size=hidden_size, bias=arch.attention_bias, quant_config=quant_config, prefix=f"{prefix}.o_proj", ) # Gemma2 interleaves global (even layers) and sliding-window (odd layers) # attention. This pattern reduces memory for long sequences while # maintaining global context every other layer. self.is_sliding = (layer_id % 2) == 1 if self.is_sliding: self.sliding_window = arch.sliding_window else: self.sliding_window = None self.rotary_emb = get_rope( self.head_dim, rotary_dim=self.head_dim, max_position=arch.max_position_embeddings, base=arch.rope_theta, is_neox_style=True, ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) batch_size, seq_len, _ = q.shape q = q.view(batch_size, seq_len, self.num_heads, self.head_dim) k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) q, k = self.rotary_emb(positions, q, k) query = q.transpose(1, 2) key = k.transpose(1, 2) value = v.transpose(1, 2) attn_mask = torch.zeros( (seq_len, seq_len), device=hidden_states.device, dtype=torch.float32 ) causal = torch.triu( torch.ones( (seq_len, seq_len), device=hidden_states.device, dtype=torch.bool ), diagonal=1, ) attn_mask = attn_mask.masked_fill(causal, float("-inf")) if self.is_sliding and self.sliding_window is not None: idx = torch.arange(seq_len, device=hidden_states.device) dist = idx[None, :] - idx[:, None] too_far = dist > self.sliding_window attn_mask = attn_mask.masked_fill(too_far, float("-inf")) if attention_mask is not None: key_pad = ~attention_mask.to(torch.bool) attn_mask = attn_mask[None, None, :, :].expand( batch_size, 1, seq_len, seq_len ) attn_mask = attn_mask.masked_fill( key_pad[:, None, None, :].expand(batch_size, 1, seq_len, seq_len), float("-inf"), ) attn_kwargs = { "attn_mask": attn_mask, "dropout_p": 0.0, "is_causal": False, "scale": self.scaling, } if query.shape[1] != key.shape[1]: attn_kwargs["enable_gqa"] = True attn_output = torch.nn.functional.scaled_dot_product_attention( query, key, value, **attn_kwargs ) # NOTE: Gemma2 specifies attn_logit_softcapping (tanh(logits/cap)*cap) but # PyTorch's scaled_dot_product_attention does not support it natively. # For short text-encoder sequences (~300 tokens), the quality impact is # negligible. A custom attention kernel would be needed for full fidelity. attn_output = attn_output.transpose(1, 2) attn_output = attn_output.reshape( batch_size, seq_len, self.num_heads * self.head_dim ) output, _ = self.o_proj(attn_output) return output class Gemma2DecoderLayer(nn.Module): def __init__( self, layer_id: int, config: Gemma2Config, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() arch = config.arch_config self.hidden_size = arch.hidden_size self.self_attn = Gemma2Attention( layer_id=layer_id, config=config, hidden_size=self.hidden_size, num_heads=arch.num_attention_heads, num_kv_heads=arch.num_key_value_heads, quant_config=quant_config, prefix=f"{prefix}.self_attn", ) self.mlp = Gemma2MLP( hidden_size=self.hidden_size, intermediate_size=arch.intermediate_size, hidden_act=arch.hidden_activation, quant_config=quant_config, prefix=f"{prefix}.mlp", ) self.input_layernorm = Gemma2RMSNorm(self.hidden_size, eps=arch.rms_norm_eps) self.post_attention_layernorm = Gemma2RMSNorm( self.hidden_size, eps=arch.rms_norm_eps ) self.pre_feedforward_layernorm = Gemma2RMSNorm( self.hidden_size, eps=arch.rms_norm_eps ) self.post_feedforward_layernorm = Gemma2RMSNorm( self.hidden_size, eps=arch.rms_norm_eps ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states = self.self_attn(positions, hidden_states, attention_mask) hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.pre_feedforward_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = self.post_feedforward_layernorm(hidden_states) hidden_states = residual + hidden_states return hidden_states class Gemma2Model(nn.Module): """Gemma2 text encoder model for SANA pipeline.""" _fsdp_shard_conditions = [] def __init__(self, config: Gemma2Config, **kwargs): super().__init__() self.config = config arch = config.arch_config self.quant_config = None self.vocab_size = arch.vocab_size self.embed_tokens = VocabParallelEmbedding( self.vocab_size, arch.hidden_size, org_num_embeddings=arch.vocab_size, quant_config=self.quant_config, ) self.embed_scale = arch.hidden_size**0.5 self.layers = nn.ModuleList( [ Gemma2DecoderLayer( layer_id=i, config=config, quant_config=self.quant_config, prefix=f"model.layers.{i}", ) for i in range(arch.num_hidden_layers) ] ) self.norm = Gemma2RMSNorm(arch.hidden_size, eps=arch.rms_norm_eps) def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) * self.embed_scale def forward( self, input_ids: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs, ) -> BaseEncoderOutput: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( "You must specify exactly one of input_ids or inputs_embeds" ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else getattr(self.config.arch_config, "output_hidden_states", False) ) if inputs_embeds is not None: hidden_states = inputs_embeds else: hidden_states = self.get_input_embeddings(input_ids) if position_ids is None: position_ids = torch.arange( 0, hidden_states.shape[1], device=hidden_states.device ).unsqueeze(0) all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None for layer in self.layers: if all_hidden_states is not None: all_hidden_states += (hidden_states,) hidden_states = layer(position_ids, hidden_states, attention_mask) hidden_states = self.norm(hidden_states) if all_hidden_states is not None: all_hidden_states += (hidden_states,) return BaseEncoderOutput( last_hidden_state=hidden_states, hidden_states=all_hidden_states, ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() stacked_params_mapping = getattr( self.config.arch_config, "stacked_params_mapping", None ) if stacked_params_mapping is None: stacked_params_mapping = [ (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), (".gate_up_proj", ".gate_proj", "0"), (".gate_up_proj", ".up_proj", "1"), ] for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue # HF Gemma2Model stores weights as model.layers.X... / model.embed_tokens... # Strip "model." prefix if present to match our naming if name.startswith("model."): name = name[len("model.") :] for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) if name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader self._load_with_shard_id(weight_loader, param, loaded_weight, shard_id) break else: if name not in params_dict: continue param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params @staticmethod def _load_with_shard_id(weight_loader, param, loaded_weight, shard_id): try: weight_loader(param, loaded_weight, shard_id) return except (AssertionError, TypeError): pass if isinstance(shard_id, str): mapping = {"q": 0, "k": 1, "v": 2} if shard_id in mapping: weight_loader(param, loaded_weight, mapping[shard_id]) return if shard_id.isdigit(): weight_loader(param, loaded_weight, int(shard_id)) return elif isinstance(shard_id, int): mapping = {0: "q", 1: "k", 2: "v"} if shard_id in mapping: weight_loader(param, loaded_weight, mapping[shard_id]) return raise TypeError( f"Unsupported shard_id={shard_id!r} for weight_loader={weight_loader}" ) EntryClass = Gemma2Model ================================================ FILE: python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from sglang: python/sglang/srt/models/gemma3_causal.py import logging from functools import partial from typing import Any, Iterable, Optional, Set, Tuple import torch from torch import nn from sglang.multimodal_gen.configs.models.encoders.base import BaseEncoderOutput from sglang.multimodal_gen.configs.models.encoders.gemma_3 import Gemma3Config from sglang.multimodal_gen.runtime.distributed import get_tp_world_size from sglang.multimodal_gen.runtime.layers.activation import GeluAndMul from sglang.multimodal_gen.runtime.layers.attention import LocalAttention from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig from sglang.multimodal_gen.runtime.layers.rotary_embedding import get_rope from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader from sglang.multimodal_gen.runtime.utils.common import add_prefix logger = logging.getLogger(__name__) def get_attention_sliding_window_size(config): return config.sliding_window - 1 class Gemma3RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.zeros(dim)) def _norm(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) def forward(self, x): output = self._norm(x.float()) output = output * (1.0 + self.weight.float()) return output.type_as(x) def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.eps}" class Gemma3MLP(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, hidden_act: str, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() self.gate_up_proj = MergedColumnParallelLinear( input_size=hidden_size, output_sizes=[intermediate_size] * 2, bias=False, quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", ) self.down_proj = RowParallelLinear( input_size=intermediate_size, output_size=hidden_size, bias=False, quant_config=quant_config, prefix=f"{prefix}.down_proj", ) if hidden_act != "gelu_pytorch_tanh": raise ValueError( "Gemma3 uses `gelu_pytorch_tanh` as the hidden activation " "function. Please set `hidden_activation` to " "`gelu_pytorch_tanh`." ) self.act_fn = GeluAndMul(approximate="tanh") def forward(self, x): x, _ = self.gate_up_proj(x) x = self.act_fn(x) x, _ = self.down_proj(x) return x def _rotate_half(x: torch.Tensor) -> torch.Tensor: x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) class Gemma3Attention(nn.Module): def __init__( self, layer_id: int, config: Gemma3Config, hidden_size: int, num_heads: int, num_kv_heads: int, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() self.layer_id = layer_id self.hidden_size = hidden_size tp_size = get_tp_world_size() self.total_num_heads = num_heads assert self.total_num_heads % tp_size == 0 self.num_heads = self.total_num_heads // tp_size self.total_num_kv_heads = num_kv_heads if self.total_num_kv_heads >= tp_size: assert self.total_num_kv_heads % tp_size == 0 else: assert tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) self.head_dim = getattr( config.text_config, "head_dim", self.hidden_size // self.total_num_heads ) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = config.text_config.query_pre_attn_scalar**-0.5 self.qkv_proj = QKVParallelLinear( hidden_size=hidden_size, head_size=self.head_dim, total_num_heads=self.total_num_heads, total_num_kv_heads=self.total_num_kv_heads, bias=config.text_config.attention_bias, quant_config=quant_config, prefix=f"{prefix}.qkv_proj", ) self.o_proj = RowParallelLinear( input_size=self.total_num_heads * self.head_dim, output_size=hidden_size, bias=config.text_config.attention_bias, quant_config=quant_config, prefix=f"{prefix}.o_proj", ) self.is_sliding = ( config.text_config.layer_types[layer_id] == "sliding_attention" ) # Initialize the rotary embedding. if self.is_sliding: # Local attention. self.rope_theta = config.text_config.rope_local_base_freq rope_scaling = None # Default # sliding window self.sliding_window = get_attention_sliding_window_size(config.text_config) # (left, right) = (window, 0) effectively for causal self.window_size = (self.sliding_window, 0) else: # Global attention. self.rope_theta = config.text_config.rope_theta rope_scaling = config.text_config.rope_scaling self.sliding_window = None self.window_size = (-1, -1) self.rotary_emb = get_rope( self.head_dim, rotary_dim=self.head_dim, max_position=config.text_config.max_position_embeddings, base=self.rope_theta, rope_scaling=rope_scaling, is_neox_style=True, ) # NOTE(gmixiaojin): The shared RotaryEmbedding above computes inv_freq on # GPU and uses the x1*cos - x2*sin formula, which causes slight # numerical differences vs HuggingFace (see the NOTE in # rotary_embedding.py:_compute_inv_freq). For HF-exact alignment we # precompute inv_freq on CPU and use rotate_half in self.rotary_emb(). freq_indices = ( torch.arange(0, self.head_dim, 2, dtype=torch.int64).float() / self.head_dim ) inv_freq = 1.0 / (self.rope_theta**freq_indices) if rope_scaling and rope_scaling.get("factor"): inv_freq = inv_freq / float(rope_scaling["factor"]) self.register_buffer("_hf_inv_freq", inv_freq, persistent=False) # Local Attention not support attention mask, we use global attention instead. # self.attn = LocalAttention( # self.num_heads, # self.head_dim, # self.num_kv_heads, # softmax_scale=self.scaling, # causal=True, # supported_attention_backends=config._supported_attention_backends, # window_size=self.window_size, # ) # Gemma3 adds normalization for q and k self.q_norm = Gemma3RMSNorm( dim=self.head_dim, eps=config.text_config.rms_norm_eps ) self.k_norm = Gemma3RMSNorm( dim=self.head_dim, eps=config.text_config.rms_norm_eps ) def rotary_emb(self, positions, q, k): """Apply RoPE using HF-exact formula with precomputed inv_freq.""" positions_flat = positions.flatten().float() num_tokens = positions_flat.shape[0] with torch.autocast(device_type=q.device.type, enabled=False): freqs = torch.outer(positions_flat, self._hf_inv_freq.float()) emb = freqs.repeat(1, 2) cos = emb.cos().to(q.dtype).unsqueeze(1) sin = emb.sin().to(q.dtype).unsqueeze(1) q = q.reshape(num_tokens, -1, self.head_dim) k = k.reshape(num_tokens, -1, self.head_dim) q = q * cos + _rotate_half(q) * sin k = k * cos + _rotate_half(k) * sin return q, k def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) batch_size, seq_len, _ = q.shape q = q.view(batch_size, seq_len, self.num_heads, self.head_dim) k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) # Apply QK Norm q = self.q_norm(q) k = self.k_norm(k) # Apply RoPE q, k = self.rotary_emb(positions, q, k) q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim) k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) # TODO(FlamingoPg): Support LocalAttention query = q.transpose(1, 2) key = k.transpose(1, 2) value = v.transpose(1, 2) min_val = torch.finfo(query.dtype).min attn_mask = torch.zeros( (seq_len, seq_len), device=hidden_states.device, dtype=query.dtype, ) causal = torch.triu( torch.ones( (seq_len, seq_len), device=hidden_states.device, dtype=torch.bool ), diagonal=1, ) attn_mask = attn_mask.masked_fill(causal, min_val) if self.is_sliding and self.sliding_window is not None: idx = torch.arange(seq_len, device=hidden_states.device) dist = idx[None, :] - idx[:, None] too_far = dist > self.sliding_window attn_mask = attn_mask.masked_fill(too_far, min_val) key_pad = ~attention_mask.to(torch.bool) attn_mask = attn_mask[None, None, :, :].expand(batch_size, 1, seq_len, seq_len) attn_mask = attn_mask.masked_fill( key_pad[:, None, None, :].expand(batch_size, 1, seq_len, seq_len), min_val, ) attn_kwargs = { "attn_mask": attn_mask, "dropout_p": 0.0, "is_causal": False, "scale": self.scaling, } if query.shape[1] != key.shape[1]: attn_kwargs["enable_gqa"] = True attn_output = torch.nn.functional.scaled_dot_product_attention( query, key, value, **attn_kwargs ) attn_output = attn_output.transpose(1, 2) attn_output = attn_output.reshape( batch_size, seq_len, self.num_heads * self.head_dim ) output, _ = self.o_proj(attn_output) return output class Gemma3DecoderLayer(nn.Module): def __init__( self, layer_id: int, config: Gemma3Config, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() self.hidden_size = config.text_config.hidden_size self.self_attn = Gemma3Attention( layer_id=layer_id, config=config, hidden_size=self.hidden_size, num_heads=config.text_config.num_attention_heads, num_kv_heads=getattr( config.text_config, "num_key_value_heads", config.text_config.num_attention_heads, ), quant_config=quant_config, prefix=f"{prefix}.self_attn", ) self.mlp = Gemma3MLP( hidden_size=self.hidden_size, intermediate_size=config.text_config.intermediate_size, hidden_act=config.text_config.hidden_activation, quant_config=quant_config, prefix=f"{prefix}.mlp", ) self.input_layernorm = Gemma3RMSNorm( config.text_config.hidden_size, eps=config.text_config.rms_norm_eps ) self.post_attention_layernorm = Gemma3RMSNorm( config.text_config.hidden_size, eps=config.text_config.rms_norm_eps ) self.pre_feedforward_layernorm = Gemma3RMSNorm( config.text_config.hidden_size, eps=config.text_config.rms_norm_eps ) self.post_feedforward_layernorm = Gemma3RMSNorm( config.text_config.hidden_size, eps=config.text_config.rms_norm_eps ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None, attention_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: # Self Attention # Gemma3 uses "sandwich norm": # x = x + norm(attn(norm(x))) # So we treat input hidden_states as the residual base. if residual is not None: hidden_states = hidden_states + residual residual = None residual_input = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, attention_mask=attention_mask, ) hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = residual_input + hidden_states # MLP residual_mlp = hidden_states hidden_states = self.pre_feedforward_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = self.post_feedforward_layernorm(hidden_states) hidden_states = residual_mlp + hidden_states return hidden_states, None class Gemma3TextScaledWordEmbedding(nn.Embedding): def __init__( self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: Optional[float] = 1.0, ): super().__init__(num_embeddings, embedding_dim, padding_idx) self.embed_scale = embed_scale def forward(self, input_ids: torch.Tensor): return super().forward(input_ids) * self.embed_scale # --- Siglip Vision Model Implementation --- class QuickGELU(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: return x * torch.sigmoid(1.702 * x) class SiglipVisionEmbeddings(nn.Module): def __init__(self, config): super().__init__() self.config = config self.embed_dim = config.hidden_size self.image_size = config.image_size self.patch_size = config.patch_size self.patch_embedding = nn.Conv2d( in_channels=config.num_channels, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size, padding="valid", ) self.num_patches = (self.image_size // self.patch_size) ** 2 self.num_positions = self.num_patches # Use simple Embedding for position embeddings (usually small enough) self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) self.register_buffer( "position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False, ) def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: target_dtype = self.patch_embedding.weight.dtype patch_embeds = self.patch_embedding( pixel_values.to(dtype=target_dtype) ) # shape = [*, width, grid, grid] embeddings = patch_embeds.flatten(2).transpose(1, 2) embeddings = embeddings + self.position_embedding(self.position_ids) return embeddings class SiglipMLP(nn.Module): def __init__( self, config, act_layer: type[nn.Module] = QuickGELU, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.fc1 = ColumnParallelLinear( config.hidden_size, config.intermediate_size, quant_config=quant_config, prefix=add_prefix("fc1", prefix), ) self.act = act_layer() self.fc2 = RowParallelLinear( config.intermediate_size, config.hidden_size, quant_config=quant_config, prefix=add_prefix("fc2", prefix), ) def forward(self, x: torch.Tensor) -> torch.Tensor: x_parallel, _ = self.fc1(x) x_parallel = self.act(x_parallel) x, _ = self.fc2(x_parallel) return x class SiglipAttention(nn.Module): def __init__( self, hidden_size: int, num_heads: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.hidden_size = hidden_size self.num_heads = num_heads tp_size = get_tp_world_size() self.head_dim = hidden_size // num_heads self.num_heads_per_partition = num_heads // tp_size self.scaling = self.head_dim**-0.5 self.qkv_proj = QKVParallelLinear( hidden_size=hidden_size, head_size=self.head_dim, total_num_heads=num_heads, total_num_kv_heads=num_heads, bias=True, quant_config=quant_config, prefix=add_prefix("qkv_proj", prefix), ) self.out_proj = RowParallelLinear( input_size=hidden_size, output_size=hidden_size, bias=True, quant_config=quant_config, prefix=add_prefix("out_proj", prefix), ) self.attn = LocalAttention( num_heads=self.num_heads_per_partition, head_size=self.head_dim, num_kv_heads=self.num_heads_per_partition, softmax_scale=self.scaling, causal=False, # Bidirectional for Vision ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.hidden_size // get_tp_world_size()] * 3, dim=-1) batch_size, seq_len, _ = q.shape q = q.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim) k = k.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim) v = v.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim) attn_output = self.attn(q, k, v) attn_output = attn_output.reshape( batch_size, seq_len, self.hidden_size // get_tp_world_size() ) output, _ = self.out_proj(attn_output) return output class SiglipEncoderLayer(nn.Module): def __init__( self, config, act_layer: type[nn.Module] = QuickGELU, norm_layer: type[nn.Module] = None, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() if norm_layer is None: norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps) self.layer_norm1 = norm_layer(config.hidden_size) self.layer_norm2 = norm_layer(config.hidden_size) self.self_attn = SiglipAttention( hidden_size=config.hidden_size, num_heads=config.num_attention_heads, quant_config=quant_config, prefix=add_prefix("self_attn", prefix), ) self.mlp = SiglipMLP( config, act_layer=act_layer, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) def forward( self, hidden_states: torch.Tensor, ) -> torch.Tensor: residual = hidden_states hidden_states = self.layer_norm1(hidden_states) hidden_states = self.self_attn(hidden_states) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.layer_norm2(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states class SiglipEncoder(nn.Module): def __init__( self, config, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.config = config num_hidden_layers = config.num_hidden_layers norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps) self.layers = nn.ModuleList( [ SiglipEncoderLayer( config=config, norm_layer=norm_layer, quant_config=quant_config, prefix=add_prefix(f"layers.{layer_idx}", prefix), ) for layer_idx in range(num_hidden_layers) ] ) def forward( self, inputs_embeds: torch.Tensor, ) -> torch.Tensor: hidden_states = inputs_embeds for encoder_layer in self.layers: hidden_states = encoder_layer(hidden_states) return hidden_states class SiglipVisionTransformer(nn.Module): def __init__( self, config, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.config = config embed_dim = config.hidden_size self.embeddings = SiglipVisionEmbeddings(config) self.encoder = SiglipEncoder( config=config, quant_config=quant_config, prefix=add_prefix("encoder", prefix), ) self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) @property def device(self) -> torch.device: return self.encoder.layers[0].layer_norm1.weight.device def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: hidden_states = self.embeddings(pixel_values.to(self.device)) last_hidden_state = self.encoder(inputs_embeds=hidden_states) last_hidden_state = self.post_layernorm(last_hidden_state) return last_hidden_state class SiglipVisionModel(nn.Module): def __init__( self, config, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.vision_model = SiglipVisionTransformer( config, quant_config, prefix=add_prefix("vision_model", prefix) ) @property def device(self) -> torch.device: return self.vision_model.device def forward(self, pixel_values: torch.Tensor): return self.vision_model(pixel_values) class Gemma3MultiModalProjector(nn.Module): """Projector for Gemma3 multimodal.""" def __init__(self, config: Gemma3Config): super().__init__() self.mm_input_projection_weight = nn.Parameter( torch.zeros( config.vision_config.hidden_size, config.text_config.hidden_size ) ) self.mm_soft_emb_norm = Gemma3RMSNorm( config.vision_config.hidden_size, eps=config.vision_config.layer_norm_eps ) self.patches_per_image = int( config.vision_config.image_size // config.vision_config.patch_size ) self.tokens_per_side = int(config.mm_tokens_per_image**0.5) self.kernel_size = self.patches_per_image // self.tokens_per_side self.avg_pool = nn.AvgPool2d( kernel_size=self.kernel_size, stride=self.kernel_size ) def forward(self, vision_outputs: torch.Tensor) -> torch.Tensor: batch_size, seq_length, hidden_size = vision_outputs.shape # Reshape for pooling reshaped_vision_outputs = vision_outputs.transpose(1, 2) reshaped_vision_outputs = reshaped_vision_outputs.reshape( batch_size, hidden_size, self.patches_per_image, self.patches_per_image ) reshaped_vision_outputs = reshaped_vision_outputs.contiguous() # Apply pooling pooled_vision_outputs = self.avg_pool(reshaped_vision_outputs) pooled_vision_outputs = pooled_vision_outputs.flatten(2) pooled_vision_outputs = pooled_vision_outputs.transpose(1, 2) # Apply normalization normed_vision_outputs = self.mm_soft_emb_norm(pooled_vision_outputs) # Project to text embedding space projected_vision_outputs = torch.matmul( normed_vision_outputs, self.mm_input_projection_weight ) return projected_vision_outputs.type_as(vision_outputs) class Gemma3TextModel(nn.Module): def __init__(self, config: Gemma3Config): super().__init__() self.config = config # TODO(yinfan.1024) support text encoding model quant later self.quant_config = None # Use VocabParallelEmbedding from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) self.vocab_size = config.text_config.vocab_size self.embed_tokens = VocabParallelEmbedding( self.vocab_size, config.text_config.hidden_size, org_num_embeddings=config.text_config.vocab_size, quant_config=self.quant_config, ) self.embed_scale = config.text_config.hidden_size**0.5 self.layers = nn.ModuleList( [ Gemma3DecoderLayer( layer_id=i, config=config, quant_config=self.quant_config, prefix=f"{config.text_config.prefix}.layers.{i}", ) for i in range(config.text_config.num_hidden_layers) ] ) self.norm = Gemma3RMSNorm( config.text_config.hidden_size, eps=config.text_config.rms_norm_eps ) def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: out = self.embed_tokens(input_ids) return out * torch.tensor(self.embed_scale, device=out.device, dtype=out.dtype) def forward( self, input_ids: torch.Tensor | None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs, ) -> BaseEncoderOutput: output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) if inputs_embeds is not None: hidden_states = inputs_embeds else: hidden_states = self.get_input_embeddings(input_ids) residual = None if position_ids is None: position_ids = torch.arange( 0, hidden_states.shape[1], device=hidden_states.device ).unsqueeze(0) all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None for layer in self.layers: if all_hidden_states is not None: all_hidden_states += (hidden_states,) hidden_states, residual = layer( position_ids, hidden_states, residual, attention_mask=attention_mask, ) hidden_states = self.norm(hidden_states) if all_hidden_states is not None: all_hidden_states += (hidden_states,) output = BaseEncoderOutput( last_hidden_state=hidden_states, hidden_states=all_hidden_states, ) return output def load_weights(self, weights: Any) -> set[str]: # Copied from LlamaModel.load_weights but adapted params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() def _load_with_shard_id( weight_loader, param, loaded_weight: torch.Tensor, shard_id ) -> None: """Call param.weight_loader with best-effort shard_id normalization. Different fused-QKV implementations expect different shard_id types: - Some expect strings: "q"/"k"/"v" - Some expect integer indices: 0/1/2 We try the provided shard_id first, then fall back between str/int forms. """ try: weight_loader(param, loaded_weight, shard_id) return except (AssertionError, TypeError): pass # Fall back between common representations. if isinstance(shard_id, str): mapping = {"q": 0, "k": 1, "v": 2} if shard_id in mapping: weight_loader(param, loaded_weight, mapping[shard_id]) return if shard_id.isdigit(): weight_loader(param, loaded_weight, int(shard_id)) return elif isinstance(shard_id, int): mapping = {0: "q", 1: "k", 2: "v"} if shard_id in mapping: weight_loader(param, loaded_weight, mapping[shard_id]) return # Re-raise with a clearer message. raise TypeError( f"Unsupported shard_id={shard_id!r} for weight_loader={weight_loader} " f"(param={getattr(param, 'name', '')})." ) stacked_params_mapping = getattr( getattr(self.config, "arch_config", object()), "stacked_params_mapping", None, ) if stacked_params_mapping is None: stacked_params_mapping = [ # Fused QKV shards; downstream loaders may want "q/k/v" or 0/1/2. (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), (".gate_up_proj", ".gate_proj", 0), (".gate_up_proj", ".up_proj", 1), ] for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue # The config has stacked_params_mapping for ( param_name, weight_name, shard_id, ) in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) if name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader _load_with_shard_id(weight_loader, param, loaded_weight, shard_id) break else: if name not in params_dict: continue param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params class Gemma3ForConditionalGeneration(nn.Module): def __init__( self, config: Gemma3Config, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.config = config self.quant_config = quant_config self.text_config = config.text_config # Vision Tower self.vision_tower = SiglipVisionModel( config=config.vision_config, quant_config=quant_config, prefix=add_prefix("vision_tower", prefix), ) # Projector self.multi_modal_projector = Gemma3MultiModalProjector(config) # Text Model self.language_model = Gemma3TextModel(config) def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor, ) -> torch.Tensor: image_token_index = int(getattr(self.config, "image_token_index", -1)) if image_token_index < 0: image_token_index = int(getattr(self.text_config, "image_token_index", -1)) special_image_mask = input_ids == image_token_index n_image_tokens = int(special_image_mask.sum().item()) special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds) n_image_features = int(image_features.shape[0] * image_features.shape[1]) if inputs_embeds[special_image_mask].numel() != image_features.numel(): raise ValueError( f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" ) return special_image_mask def forward( self, input_ids: torch.LongTensor, pixel_values: torch.FloatTensor | None = None, **kwargs, ): vocab_size = int(self.language_model.vocab_size) image_token_index = int(getattr(self.config, "image_token_index", -1)) if image_token_index < 0: image_token_index = int(getattr(self.text_config, "image_token_index", -1)) if input_ids is not None and image_token_index >= vocab_size: special_image_mask = input_ids == image_token_index llm_input_ids = input_ids.clone() llm_input_ids[special_image_mask] = 0 else: llm_input_ids = input_ids inputs_embeds = self.language_model.get_input_embeddings(llm_input_ids) if pixel_values is not None: if pixel_values.dim() == 5: pixel_values = pixel_values.reshape( -1, pixel_values.shape[2], pixel_values.shape[3], pixel_values.shape[4], ) elif pixel_values.dim() == 3: pixel_values = pixel_values.unsqueeze(0) elif pixel_values.dim() != 4: raise ValueError(f"Unexpected pixel_values shape: {pixel_values.shape}") vision_outputs = self.vision_tower(pixel_values) image_features = self.multi_modal_projector(vision_outputs) image_features = image_features.to( device=inputs_embeds.device, dtype=inputs_embeds.dtype ) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features ) inputs_embeds = inputs_embeds.masked_scatter( special_image_mask, image_features ) return self.language_model.forward( llm_input_ids, inputs_embeds=inputs_embeds, **kwargs ) def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]: loaded_params: Set[str] = set() params_dict = dict(self.named_parameters()) def _load_with_shard_id( weight_loader, param, loaded_weight: torch.Tensor, shard_id ) -> None: """Call param.weight_loader with best-effort shard_id normalization. Different fused-QKV implementations expect different shard_id types: - Some expect strings: "q"/"k"/"v" - Some expect integer indices: 0/1/2 We try the provided shard_id first, then fall back between str/int forms. """ try: weight_loader(param, loaded_weight, shard_id) return except (AssertionError, TypeError): pass # Fall back between common representations. if isinstance(shard_id, str): mapping = {"q": 0, "k": 1, "v": 2} if shard_id in mapping: weight_loader(param, loaded_weight, mapping[shard_id]) return if shard_id.isdigit(): weight_loader(param, loaded_weight, int(shard_id)) return elif isinstance(shard_id, int): mapping = {0: "q", 1: "k", 2: "v"} if shard_id in mapping: weight_loader(param, loaded_weight, mapping[shard_id]) return raise TypeError( f"Unsupported shard_id={shard_id!r} for weight_loader={weight_loader} " f"(param={getattr(param, 'name', '')})." ) # Separate weights language_model_weights: list[tuple[str, torch.Tensor]] = [] other_weights: list[tuple[str, torch.Tensor]] = [] for name, loaded_weight in weights: # Handle prefix mapping if needed # HF weights might be "model.vision_tower...", "model.language_model..." if "vision_tower" in name or "vision_model" in name: # Load vision tower weights # Map name to local name local_name = name if "model.vision_tower" in name: local_name = name.replace("model.vision_tower", "vision_tower") elif "vision_tower" in name: pass # already correct prefix if matching self.vision_tower elif local_name.startswith("vision_model."): local_name = ( "vision_tower.vision_model." + local_name[len("vision_model.") :] ) # We need to map HF Siglip names to our Siglip implementation # Our Siglip: vision_tower.vision_model.encoder.layers... # HF Siglip: vision_model.encoder.layers... # If loading from Gemma3 checkpoint, it usually has "model.vision_tower.vision_model..." if local_name in params_dict: param = params_dict[local_name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) loaded_params.add(local_name) else: qkv_shard_id = None fused_name = None if ".self_attn.q_proj." in local_name: fused_name = local_name.replace( ".self_attn.q_proj.", ".self_attn.qkv_proj." ) qkv_shard_id = "q" elif ".self_attn.k_proj." in local_name: fused_name = local_name.replace( ".self_attn.k_proj.", ".self_attn.qkv_proj." ) qkv_shard_id = "k" elif ".self_attn.v_proj." in local_name: fused_name = local_name.replace( ".self_attn.v_proj.", ".self_attn.qkv_proj." ) qkv_shard_id = "v" if fused_name is not None and fused_name in params_dict: param = params_dict[fused_name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) _load_with_shard_id( weight_loader, param, loaded_weight, qkv_shard_id ) loaded_params.add(fused_name) continue if ".self_attn.proj." in local_name: candidate = local_name.replace( ".self_attn.proj.", ".self_attn.out_proj." ) if candidate in params_dict: param = params_dict[candidate] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) loaded_params.add(candidate) continue if ".self_attn.out_proj." in local_name: candidate = local_name.replace( ".self_attn.out_proj.", ".self_attn.proj." ) if candidate in params_dict: param = params_dict[candidate] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) loaded_params.add(candidate) continue # Try to find match suffix = local_name.split("vision_tower.")[-1] # Try adding vision_model candidate = f"vision_tower.vision_model.{suffix}" if candidate in params_dict: param = params_dict[candidate] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) loaded_params.add(candidate) elif "multi_modal_projector" in name: local_name = name if "model.multi_modal_projector" in name: local_name = name.replace( "model.multi_modal_projector", "multi_modal_projector" ) if local_name in params_dict: param = params_dict[local_name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) loaded_params.add(local_name) elif "language_model" in name or "model.language_model" in name: # Strip prefix for language model # If name is "model.language_model.model.layers.0...", we want "model.layers.0..." for Gemma3ForCausalLM # Gemma3ForCausalLM has .model (Gemma3TextModel) and .lm_head # HF: model.language_model.model.layers... # Ours: language_model.model.layers... # We pass (name, weight) to language_model.load_weights # We should strip "model.language_model." or "language_model." suffix = name if "model.language_model." in name: suffix = name.replace("model.language_model.", "") elif "language_model." in name: suffix = name.replace("language_model.", "") if suffix.startswith("model."): suffix = suffix[len("model.") :] language_model_weights.append((suffix, loaded_weight)) else: # Fallback for other weights (maybe direct lm_head if not nested?) other_weights.append((name, loaded_weight)) if language_model_weights: lm_loaded = self.language_model.load_weights(language_model_weights) loaded_params.update({f"language_model.{n}" for n in lm_loaded}) return loaded_params def get_attention_sliding_window_size(self): if self.text_config is not None and hasattr( self.text_config, "get_attention_sliding_window_size" ): return self.text_config.get_attention_sliding_window_size() sliding_window = getattr(self.text_config, "sliding_window", None) if sliding_window is None: sliding_window = getattr(self.config, "sliding_window", None) if sliding_window is None: return None return int(sliding_window) - 1 EntryClass = Gemma3ForConditionalGeneration ================================================ FILE: python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py ================================================ # Copied and adapted from: https://github.com/Tencent-Hunyuan/Hunyuan3D-2 import numpy as np import torch import torch.nn as nn from torchvision import transforms from transformers import ( CLIPVisionConfig, CLIPVisionModelWithProjection, Dinov2Config, Dinov2Model, ) def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): assert embed_dim % 2 == 0 omega = np.arange(embed_dim // 2, dtype=np.float64) omega /= embed_dim / 2.0 omega = 1.0 / 10000**omega # (D/2,) pos = pos.reshape(-1) # (M,) out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product emb_sin = np.sin(out) # (M, D/2) emb_cos = np.cos(out) # (M, D/2) return np.concatenate([emb_sin, emb_cos], axis=1) class ImageEncoder(nn.Module): MODEL_CLASS = None MODEL_CONFIG_CLASS = None mean = [] std = [] def __init__( self, version=None, config=None, use_cls_token=True, image_size=224, **kwargs, ): super().__init__() if config is None: self.model = self.MODEL_CLASS.from_pretrained(version) else: self.model = self.MODEL_CLASS(self.MODEL_CONFIG_CLASS.from_dict(config)) self.model.eval() self.model.requires_grad_(False) self.use_cls_token = use_cls_token self.size = image_size // 14 self.num_patches = (image_size // 14) ** 2 if self.use_cls_token: self.num_patches += 1 self.transform = transforms.Compose( [ transforms.Resize( image_size, transforms.InterpolationMode.BILINEAR, antialias=True ), transforms.CenterCrop(image_size), transforms.Normalize( mean=self.mean, std=self.std, ), ] ) def forward(self, image, mask=None, value_range=(-1, 1), **kwargs): if value_range is not None: low, high = value_range image = (image - low) / (high - low) image = image.to(self.model.device, dtype=self.model.dtype) inputs = self.transform(image) outputs = self.model(inputs) last_hidden_state = outputs.last_hidden_state if not self.use_cls_token: last_hidden_state = last_hidden_state[:, 1:, :] return last_hidden_state def unconditional_embedding(self, batch_size, **kwargs): device = next(self.model.parameters()).device dtype = next(self.model.parameters()).dtype zero = torch.zeros( batch_size, self.num_patches, self.model.config.hidden_size, device=device, dtype=dtype, ) return zero class CLIPImageEncoder(ImageEncoder): MODEL_CLASS = CLIPVisionModelWithProjection MODEL_CONFIG_CLASS = CLIPVisionConfig mean = [0.48145466, 0.4578275, 0.40821073] std = [0.26862954, 0.26130258, 0.27577711] class DinoImageEncoder(ImageEncoder): MODEL_CLASS = Dinov2Model MODEL_CONFIG_CLASS = Dinov2Config mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] class DinoImageEncoderMV(DinoImageEncoder): _aliases = [ "hy3dshape.models.conditioner.DinoImageEncoderMV", ] def __init__( self, version=None, config=None, use_cls_token=True, image_size=224, view_num=4, **kwargs, ): super().__init__(version, config, use_cls_token, image_size, **kwargs) self.view_num = view_num self.num_patches = self.num_patches pos = np.arange(self.view_num, dtype=np.float32) view_embedding = torch.from_numpy( get_1d_sincos_pos_embed_from_grid(self.model.config.hidden_size, pos) ).float() view_embedding = view_embedding.unsqueeze(1).repeat(1, self.num_patches, 1) self.view_embed = view_embedding.unsqueeze(0) def forward(self, image, mask=None, value_range=(-1, 1), view_idxs=None, **kwargs): if value_range is not None: low, high = value_range image = (image - low) / (high - low) image = image.to(self.model.device, dtype=self.model.dtype) bs, num_views, c, h, w = image.shape image = image.view(bs * num_views, c, h, w) inputs = self.transform(image) outputs = self.model(inputs) last_hidden_state = outputs.last_hidden_state last_hidden_state = last_hidden_state.view( bs, num_views, last_hidden_state.shape[-2], last_hidden_state.shape[-1] ) view_embedding = self.view_embed.to(last_hidden_state.dtype).to( last_hidden_state.device ) if view_idxs is not None: assert len(view_idxs) == bs view_embeddings = [] for i in range(bs): view_idx = view_idxs[i] assert num_views == len(view_idx) view_embeddings.append(self.view_embed[:, view_idx, ...]) view_embedding = ( torch.cat(view_embeddings, 0) .to(last_hidden_state.dtype) .to(last_hidden_state.device) ) if num_views != self.view_num: view_embedding = view_embedding[:, :num_views, ...] last_hidden_state = last_hidden_state + view_embedding last_hidden_state = last_hidden_state.view( bs, num_views * last_hidden_state.shape[-2], last_hidden_state.shape[-1] ) return last_hidden_state def unconditional_embedding(self, batch_size, view_idxs, **kwargs): device = next(self.model.parameters()).device dtype = next(self.model.parameters()).dtype zero = torch.zeros( batch_size, self.num_patches * len(view_idxs[0]), self.model.config.hidden_size, device=device, dtype=dtype, ) return zero def build_image_encoder(config): if config["type"] == "CLIPImageEncoder": return CLIPImageEncoder(**config["kwargs"]) elif config["type"] == "DinoImageEncoder": return DinoImageEncoder(**config["kwargs"]) elif config["type"] == "DinoImageEncoderMV": return DinoImageEncoderMV(**config["kwargs"]) else: raise ValueError(f'Unknown image encoder type: {config["type"]}') class DualImageEncoder(nn.Module): def __init__( self, main_image_encoder, additional_image_encoder, ): super().__init__() self.main_image_encoder = build_image_encoder(main_image_encoder) self.additional_image_encoder = build_image_encoder(additional_image_encoder) def forward(self, image, mask=None, **kwargs): outputs = { "main": self.main_image_encoder(image, mask=mask, **kwargs), "additional": self.additional_image_encoder(image, mask=mask, **kwargs), } return outputs def unconditional_embedding(self, batch_size, **kwargs): outputs = { "main": self.main_image_encoder.unconditional_embedding( batch_size, **kwargs ), "additional": self.additional_image_encoder.unconditional_embedding( batch_size, **kwargs ), } return outputs class SingleImageEncoder(nn.Module): def __init__( self, main_image_encoder, ): super().__init__() self.main_image_encoder = build_image_encoder(main_image_encoder) def forward(self, image, mask=None, **kwargs): outputs = { "main": self.main_image_encoder(image, mask=mask, **kwargs), } return outputs def unconditional_embedding(self, batch_size, **kwargs): outputs = { "main": self.main_image_encoder.unconditional_embedding( batch_size, **kwargs ), } return outputs # Entry class for model registry EntryClass = [ SingleImageEncoder, DualImageEncoder, DinoImageEncoder, DinoImageEncoderMV, ] ================================================ FILE: python/sglang/multimodal_gen/runtime/models/encoders/llama.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/llama.py # Adapted from # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py # Copyright 2023 The vLLM team. # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. # # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX # and OPT implementations in this library. It has been modified from its # original forms to accommodate minor architectural differences compared # to GPT-NeoX and OPT used by the Meta AI team that trained the model. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Inference-only LLaMA model compatible with HuggingFace weights.""" from collections.abc import Iterable from typing import Any import torch from torch import nn # from ..utils import (extract_layer_index) from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput, LlamaConfig from sglang.multimodal_gen.runtime.distributed import get_tp_world_size from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul # from vllm.model_executor.layers.quantization import QuantizationConfig from sglang.multimodal_gen.runtime.layers.attention import LocalAttention from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm from sglang.multimodal_gen.runtime.layers.linear import ( MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig from sglang.multimodal_gen.runtime.layers.rotary_embedding import get_rope from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from sglang.multimodal_gen.runtime.loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, ) from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder class LlamaMLP(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, hidden_act: str, quant_config: QuantizationConfig | None = None, bias: bool = False, prefix: str = "", ) -> None: super().__init__() self.gate_up_proj = MergedColumnParallelLinear( input_size=hidden_size, output_sizes=[intermediate_size] * 2, # output_size=intermediate_size, bias=bias, quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", ) self.down_proj = RowParallelLinear( input_size=intermediate_size, output_size=hidden_size, bias=bias, quant_config=quant_config, prefix=f"{prefix}.down_proj", ) if hidden_act != "silu": raise ValueError( f"Unsupported activation: {hidden_act}. " "Only silu is supported for now." ) self.act_fn = SiluAndMul() def forward(self, x): x, _ = self.gate_up_proj(x) x = self.act_fn(x) x, _ = self.down_proj(x) return x class LlamaAttention(nn.Module): def __init__( self, config: LlamaConfig, hidden_size: int, num_heads: int, num_kv_heads: int, rope_theta: float = 10000, rope_scaling: dict[str, Any] | None = None, max_position_embeddings: int = 8192, quant_config: QuantizationConfig | None = None, bias: bool = False, bias_o_proj: bool = False, prefix: str = "", ) -> None: super().__init__() # layer_idx = extract_layer_index(prefix) self.hidden_size = hidden_size tp_size = get_tp_world_size() self.total_num_heads = num_heads assert self.total_num_heads % tp_size == 0 self.num_heads = self.total_num_heads // tp_size self.total_num_kv_heads = num_kv_heads if self.total_num_kv_heads >= tp_size: # Number of KV heads is greater than TP size, so we partition # the KV heads across multiple tensor parallel GPUs. assert self.total_num_kv_heads % tp_size == 0 else: # Number of KV heads is less than TP size, so we replicate # the KV heads across multiple tensor parallel GPUs. assert tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) # MistralConfig has an optional head_dim introduced by Mistral-Nemo self.head_dim = getattr( config, "head_dim", self.hidden_size // self.total_num_heads ) # Phi models introduced a partial_rotary_factor parameter in the config partial_rotary_factor = getattr(config, "partial_rotary_factor", 1) self.rotary_dim = int(partial_rotary_factor * self.head_dim) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 self.rope_theta = rope_theta self.max_position_embeddings = max_position_embeddings self.qkv_proj = QKVParallelLinear( hidden_size=hidden_size, head_size=self.head_dim, total_num_heads=self.total_num_heads, total_num_kv_heads=self.total_num_kv_heads, bias=bias, quant_config=quant_config, prefix=f"{prefix}.qkv_proj", ) self.o_proj = RowParallelLinear( input_size=self.total_num_heads * self.head_dim, output_size=hidden_size, bias=bias_o_proj, quant_config=quant_config, prefix=f"{prefix}.o_proj", ) is_neox_style = True is_gguf = ( quant_config and hasattr(quant_config, "get_name") and quant_config.get_name() == "gguf" ) if is_gguf and config.model_type == "llama": is_neox_style = False self.rotary_emb = get_rope( self.head_dim, rotary_dim=self.rotary_dim, max_position=max_position_embeddings, base=int(rope_theta), rope_scaling=rope_scaling, is_neox_style=is_neox_style, ) self.attn = LocalAttention( self.num_heads, self.head_dim, self.num_kv_heads, softmax_scale=self.scaling, causal=True, supported_attention_backends=config._supported_attention_backends, ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) q, k = self.rotary_emb(positions, q, k) # attn_output = self.attn(q, k, v) # use flash_attn_func # TODO (Attn abstraction and backend) # reshape q, k, v to (batch_size, seq_len, num_heads, head_dim) batch_size = q.shape[0] seq_len = q.shape[1] q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim) k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) # import pdb; pdb.set_trace() # attn_output = flash_attn_varlen_func(q, k, v, softmax_scale=self.scaling, causal=True) attn_output = self.attn(q, k, v) attn_output = attn_output.reshape( batch_size, seq_len, self.num_heads * self.head_dim ) output, _ = self.o_proj(attn_output) return output class LlamaDecoderLayer(nn.Module): def __init__( self, config: LlamaConfig, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() self.hidden_size = config.hidden_size rope_theta = config.rope_parameters["rope_theta"] rope_scaling = config.rope_parameters if rope_scaling is not None and getattr( config, "original_max_position_embeddings", None ): rope_scaling["original_max_position_embeddings"] = ( config.original_max_position_embeddings ) max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) bias_o_proj = attention_bias # support internlm/internlm3-8b with qkv_bias if hasattr(config, "qkv_bias"): attention_bias = config.qkv_bias self.self_attn = LlamaAttention( config=config, hidden_size=self.hidden_size, num_heads=config.num_attention_heads, num_kv_heads=getattr( config, "num_key_value_heads", config.num_attention_heads ), rope_theta=rope_theta, rope_scaling=rope_scaling, max_position_embeddings=max_position_embeddings, quant_config=quant_config, bias=attention_bias, bias_o_proj=bias_o_proj, prefix=f"{prefix}.self_attn", ) self.mlp = LlamaMLP( hidden_size=self.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, bias=getattr(config, "mlp_bias", False), prefix=f"{prefix}.mlp", ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: # Self Attention if residual is None: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: hidden_states, residual = self.input_layernorm(hidden_states, residual) hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) # Fully Connected hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) hidden_states = self.mlp(hidden_states) return hidden_states, residual class LlamaModel(TextEncoder): def __init__( self, config: LlamaConfig, ): super().__init__(config) self.config = config self.quant_config = self.config.quant_config if config.lora_config is not None: max_loras = 1 lora_vocab_size = 1 if hasattr(config.lora_config, "max_loras"): max_loras = config.lora_config.max_loras if hasattr(config.lora_config, "lora_extra_vocab_size"): lora_vocab_size = config.lora_config.lora_extra_vocab_size lora_vocab = lora_vocab_size * max_loras else: lora_vocab = 0 self.vocab_size = config.vocab_size + lora_vocab self.org_vocab_size = config.vocab_size self.embed_tokens = VocabParallelEmbedding( self.vocab_size, config.hidden_size, org_num_embeddings=config.vocab_size, quant_config=config.quant_config, ) self.layers = nn.ModuleList( [ LlamaDecoderLayer( config=config, quant_config=config.quant_config, prefix=f"{config.prefix}.layers.{i}", ) for i in range(config.num_hidden_layers) ] ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) def forward( self, input_ids: torch.Tensor | None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs, ) -> BaseEncoderOutput: output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) if inputs_embeds is not None: hidden_states = inputs_embeds else: hidden_states = self.get_input_embeddings(input_ids) residual = None if position_ids is None: position_ids = torch.arange( 0, hidden_states.shape[1], device=hidden_states.device ).unsqueeze(0) all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None for layer in self.layers: if all_hidden_states is not None: # TODO all_hidden_states += ( (hidden_states,) if residual is None else (hidden_states + residual,) ) hidden_states, residual = layer(position_ids, hidden_states, residual) hidden_states, _ = self.norm(hidden_states, residual) # add hidden states from the last decoder layer if all_hidden_states is not None: all_hidden_states += (hidden_states,) # TODO(will): maybe unify the output format with other models and use # our own class output = BaseEncoderOutput( last_hidden_state=hidden_states, # past_key_values=past_key_values if use_cache else None, hidden_states=all_hidden_states, # attentions=all_self_attns, ) return output def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: # Models trained using ColossalAI may include these tensors in # the checkpoint. Skip them. continue # if (self.quant_config is not None and # (scale_name := self.quant_config.get_cache_scale(name))): # # Loading kv cache quantization scales # param = params_dict[scale_name] # weight_loader = getattr(param, "weight_loader", # default_weight_loader) # loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else # loaded_weight[0]) # weight_loader(param, loaded_weight) # loaded_params.add(scale_name) # continue if "scale" in name: # Remapping the name of FP8 kv-scale. kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict) if kv_scale_name is None: continue else: name = kv_scale_name for ( param_name, weight_name, shard_id, ) in self.config.arch_config.stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) break else: # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: continue param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params EntryClass = LlamaModel ================================================ FILE: python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py ================================================ # coding=utf-8 # Copyright 2025 HuggingFace Inc. team. All rights reserved. # # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from typing import Iterable, Optional, Union import torch from torch import nn from transformers import Cache, DynamicCache, LlavaConfig, Mistral3Config, MistralConfig from transformers.integrations.sdpa_attention import sdpa_attention_forward from transformers.masking_utils import create_causal_mask from transformers.modeling_outputs import BaseModelOutputWithPast from transformers.models.mistral3.modeling_mistral3 import ( Mistral3CausalLMOutputWithPast, Mistral3ModelOutputWithPast, ) from transformers.models.mistral.modeling_mistral import ( MistralMLP, MistralRMSNorm, MistralRotaryEmbedding, apply_rotary_pos_emb, ) from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) """ batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand( batch, num_key_value_heads, n_rep, slen, head_dim ) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) class MistralAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__(self, config: MistralConfig, layer_idx: int): super().__init__() self.config = config self.layer_idx = layer_idx self.num_key_value_groups = ( config.num_attention_heads // config.num_key_value_heads ) self.head_dim = ( getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads ) self.num_key_value_groups = ( config.num_attention_heads // config.num_key_value_heads ) self.scaling = self.head_dim**-0.5 self.attention_dropout = config.attention_dropout self.is_causal = True self.q_proj = nn.Linear( config.hidden_size, config.num_attention_heads * self.head_dim, bias=False ) self.k_proj = nn.Linear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False ) self.v_proj = nn.Linear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False ) self.o_proj = nn.Linear( config.num_attention_heads * self.head_dim, config.hidden_size, bias=False ) self.is_causal = True self.num_heads = config.num_attention_heads self.num_key_value_heads = config.num_key_value_heads self.attn = USPAttention( num_heads=self.num_heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, causal=False, supported_attention_backends={ AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA, }, ) def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_values: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb( query_states, key_states, cos, sin ) if past_key_values is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_values.update( key_states, value_states, self.layer_idx, cache_kwargs ) attention_interface = sdpa_attention_forward attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0, scaling=self.scaling, sliding_window=getattr( self.config, "sliding_window", None ), # main diff with Llama **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output class MistralDecoderLayer(nn.Module): def __init__(self, config: MistralConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.self_attn = MistralAttention(config=config, layer_idx=layer_idx) self.mlp = MistralMLP(config) self.input_layernorm = MistralRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.post_attention_layernorm = MistralRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[ tuple[torch.Tensor, torch.Tensor] ] = None, # necessary, but kept here for BC **kwargs, ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, **kwargs, ) hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states class MistralModel(nn.Module): def __init__(self, config: MistralConfig): super().__init__() self.config = config self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, self.padding_idx ) self.layers = nn.ModuleList( [ MistralDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers) ] ) self.norm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = MistralRotaryEmbedding(config=config) self.gradient_checkpointing = False self.config._attn_implementation = "sdpa" def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, output_hidden_states: Optional[bool] = None, **kwargs, ) -> BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( "You must specify exactly one of input_ids or inputs_embeds" ) if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) if cache_position is None: past_seen_tokens = ( past_key_values.get_seq_length() if past_key_values is not None else 0 ) cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device, ) if position_ids is None: position_ids = cache_position.unsqueeze(0) mask_function = create_causal_mask causal_mask = mask_function( config=self.config, input_embeds=inputs_embeds, attention_mask=attention_mask, cache_position=cache_position, past_key_values=past_key_values, position_ids=position_ids, ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids) hidden_states_pool = [] for decoder_layer in self.layers[: self.config.num_hidden_layers]: hidden_states = decoder_layer( hidden_states, attention_mask=causal_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, **kwargs, ) if output_hidden_states: hidden_states_pool.append(hidden_states) hidden_states = self.norm(hidden_states) if output_hidden_states: hidden_states_pool.append(hidden_states) return BaseModelOutputWithPast( hidden_states=hidden_states_pool, last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None, ) class Mistral3Model(nn.Module): _checkpoint_conversion_mapping = {"language_model.model": "language_model"} def __init__(self, config: Mistral3Config): super().__init__() self.language_model = MistralModel(config.text_config) self.config = config def get_input_embeddings(self): return self.language_model.embed_tokens def set_decoder(self, decoder): self.language_model = decoder def get_decoder(self): return self.language_model def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, output_hidden_states: Optional[bool] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidoutput_hidden_statesden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, image_sizes: Optional[torch.Tensor] = None, **kwargs, ) -> Union[tuple, Mistral3ModelOutputWithPast]: output_attentions = False output_hidden_states = True if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( "You must specify exactly one of input_ids or inputs_embeds" ) if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) outputs: BaseModelOutputWithPast = self.language_model( attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_position, **kwargs, ) return Mistral3ModelOutputWithPast( last_hidden_state=outputs.last_hidden_state, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) class Mistral3ForConditionalGeneration(nn.Module): _checkpoint_conversion_mapping = { "^language_model.model": "model.language_model", "^multi_modal_projector": "model.multi_modal_projector", "^language_model.lm_head": "lm_head", } _tied_weights_keys = ["lm_head.weight"] def __init__(self, config: LlavaConfig): super().__init__() self.model = Mistral3Model(config.arch_config) def get_input_embeddings(self): return self.model.get_input_embeddings() def set_decoder(self, decoder): self.model.set_decoder(decoder) def get_decoder(self): return self.model.get_decoder() # Make modules available through conditional class for BC @property def language_model(self): return self.model.language_model def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, output_hidden_states: Optional[bool] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, image_sizes: Optional[torch.Tensor] = None, **kwargs, ) -> Union[tuple, Mistral3CausalLMOutputWithPast]: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Example: """ output_hidden_states = True outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_position, image_sizes=image_sizes, **kwargs, ) return Mistral3CausalLMOutputWithPast( hidden_states=outputs.hidden_states, ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Define mapping for stacked parameters params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: name_lower = name.lower() if ( "vision" in name_lower or "multi" in name_lower or "lm_head" in name_lower ): continue final_name = name.replace("language_model.model.", "model.language_model.") if final_name in params_dict: param = params_dict[final_name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(final_name) else: logger.warning(f"Param {name=} {final_name=} from weight is not loaded") return loaded_params EntryClass = Mistral3ForConditionalGeneration ================================================ FILE: python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from transformers import ( Cache, DynamicCache, PretrainedConfig, Qwen2_5_VLTextConfig, Qwen2RMSNorm, ) from transformers.masking_utils import ( create_causal_mask, create_sliding_window_causal_mask, ) from transformers.modeling_flash_attention_utils import FlashAttentionKwargs from transformers.modeling_outputs import BaseModelOutputWithPast from transformers.utils import TransformersKwargs, is_torchdynamo_compiling from sglang.multimodal_gen.configs.models.encoders.qwen_image import Qwen2_5VLConfig from sglang.multimodal_gen.runtime.layers.attention import LocalAttention from sglang.multimodal_gen.runtime.layers.linear import ( MergedColumnParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.common import add_prefix # coding=utf-8 # Adapted from # https://github.com/huggingface/transformers/blob/19e6e80e10118f855137b90740936c0b11ac397f/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py # Copyright 2024 The Qwen team. # Copyright 2023 The vLLM team. # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. # # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX # and OPT implementations in this library. It has been modified from its # original forms to accommodate minor architectural differences compared # to GPT-NeoX and OPT used by the Meta AI team that trained the model. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Inference-only Qwen2-VL model compatible with HuggingFace weights.""" import logging from typing import Callable, Iterable, Optional, Tuple, Union try: from typing import Unpack # type: ignore[attr-defined] except ImportError: # Python 3.10 and below from typing_extensions import Unpack import torch import torch.nn as nn from transformers.activations import ACT2FN from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( Qwen2_5_VisionTransformerPretrainedModel, Qwen2_5_VLAttention, Qwen2_5_VLCausalLMOutputWithPast, Qwen2_5_VLModelOutputWithPast, Qwen2_5_VLRotaryEmbedding, Qwen2MLP, apply_multimodal_rotary_pos_emb, eager_attention_forward, ) logger = logging.getLogger(__name__) class Qwen2_5_VLAttention(nn.Module): """ Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer and "Generating Long Sequences with Sparse Transformers". """ def __init__(self, config: Qwen2_5_VLTextConfig, layer_idx: Optional[int] = None): super().__init__() self.config = config self.layer_idx = layer_idx if layer_idx is None: logger.warn( f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " "when creating this class." ) self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.hidden_size // self.num_heads self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.is_causal = True self.attention_dropout = config.attention_dropout self.rope_scaling = config.rope_scaling self.scaling = self.head_dim**-0.5 if (self.head_dim * self.num_heads) != self.hidden_size: raise ValueError( f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" f" and `num_heads`: {self.num_heads})." ) self.q_proj = nn.Linear( self.hidden_size, self.num_heads * self.head_dim, bias=True ) self.k_proj = nn.Linear( self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True ) self.v_proj = nn.Linear( self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True ) self.o_proj = nn.Linear( self.num_heads * self.head_dim, self.hidden_size, bias=False ) self.sliding_window = ( config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None ) self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) self.attn = LocalAttention( num_heads=self.num_heads, head_size=self.head_dim, num_kv_heads=self.num_key_value_heads, softmax_scale=self.scaling, causal=True, supported_attention_backends=( AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA, ), ) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[ tuple[torch.Tensor, torch.Tensor] ] = None, # necessary, but kept here for BC **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_multimodal_rotary_pos_emb( query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] ) if past_key_values is not None: cache_kwargs = { "sin": sin, "cos": cos, "cache_position": cache_position, } # Specific to RoPE models key_states, value_states = past_key_values.update( key_states, value_states, self.layer_idx, cache_kwargs ) attention_interface: Callable = eager_attention_forward # if self.config._attn_implementation != "eager": # attention_interface = ALL_ATTENTION_FUNCTIONS["sdpa"] query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) attn_output = self.attn(query_states, key_states, value_states) # # attn_output, attn_weights = attention_interface( # self, # query_states, # key_states, # value_states, # attention_mask, # dropout=0.0 if not self.training else self.attention_dropout, # scaling=self.scaling, # sliding_window=self.sliding_window, # position_ids=position_ids, # pass positions for FA2 # **kwargs, # ) attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output class Qwen2_5_VLDecoderLayer(nn.Module): def __init__(self, config: Qwen2_5_VLTextConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size if ( config.use_sliding_window and config._attn_implementation != "flash_attention_2" ): logger.warning( f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " "unexpected results may be encountered." ) self.self_attn = Qwen2_5_VLAttention(config, layer_idx) self.mlp = Qwen2MLP(config) self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = Qwen2RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.attention_type = config.layer_types[layer_idx] def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[ tuple[torch.Tensor, torch.Tensor] ] = None, # necessary, but kept here for BC **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple[ torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]] ]: """ Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch, sequence_length)` where padding elements are indicated by 0. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). past_key_values (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, with `head_dim` being the embedding dimension of each attention head. kwargs (`dict`, *optional*): Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code into the model """ residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, **kwargs, ) hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states class Qwen2_5_VLMLP(nn.Module): def __init__( self, in_features: int, hidden_features: int = None, bias: bool = True, hidden_act="silu", quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.gate_up_proj = MergedColumnParallelLinear( input_size=in_features, output_sizes=[hidden_features] * 2, # [gate_proj, up_proj] bias=bias, quant_config=quant_config, prefix=add_prefix("gate_up_proj", prefix), ) self.down_proj = RowParallelLinear( hidden_features, in_features, bias=bias, quant_config=quant_config, prefix=add_prefix("down_proj", prefix), ) self.act = ACT2FN[hidden_act] def forward(self, x: torch.Tensor) -> torch.Tensor: gate_up, _ = self.gate_up_proj(x) gate, up = gate_up.chunk(2, dim=-1) x = self.act(gate) * up x_down, _ = self.down_proj(x) return x_down class Qwen2_5_VLTextModel(nn.Module): def __init__(self, config: PretrainedConfig): super().__init__() self.config = config self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, self.padding_idx ) self.layers = nn.ModuleList( [ Qwen2_5_VLDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers) ] ) self._attn_implementation = config._attn_implementation self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) self.has_sliding_layers = "sliding_attention" in self.config.layer_types self.gradient_checkpointing = False # Initialize weights and apply final processing # self.post_init() def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> Union[tuple, BaseModelOutputWithPast]: output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( "You must specify exactly one of input_ids or inputs_embeds" ) # torch.jit.trace() doesn't support cache objects in the output if use_cache and past_key_values is None and not torch.jit.is_tracing(): past_key_values = DynamicCache(config=self.config) if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) if cache_position is None: past_seen_tokens = ( past_key_values.get_seq_length() if past_key_values is not None else 0 ) cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device, ) # the hard coded `3` is for temporal, height and width. if position_ids is None: position_ids = cache_position.view(1, 1, -1).expand( 3, inputs_embeds.shape[0], -1 ) elif position_ids.ndim == 2: position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) # NOTE: we need to pass text position ids for packing. Qwen2-VL uses 3D positions # where each dim indicates visual spatial positions for temporal/height/width grids. # There are two scenarios when FA2-like packed masking might be activated. # 1. User specifically passed packed `position_ids` and no attention mask. # In this case we expect the user to create correct position ids for all 3 grids # and prepend text-only position ids to it. The final tensor will be [4, bs, seq-len] # 2. User runs forward with no attention mask and no position ids. In this case, position ids # are prepared by the model (`get_rope_index`) as `[4, bs, seq-len]` tensor. Text-only positions are # prepended by us when creating positions so that the mask is constructed correctly. NOTE: failing to pass # text-only positions will cause incorrect mask construction, do not change `prepare_input_for_generation` if position_ids.ndim == 3 and position_ids.shape[0] == 4: text_position_ids = position_ids[0] position_ids = position_ids[1:] else: text_position_ids = position_ids[0] # It may already have been prepared by e.g. `generate` if not isinstance(causal_mask_mapping := attention_mask, dict): # Prepare mask arguments mask_kwargs = { "config": self.config, "input_embeds": inputs_embeds, "attention_mask": attention_mask, "cache_position": cache_position, "past_key_values": past_key_values, "position_ids": text_position_ids, } # Create the masks causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), } # The sliding window alternating layers are not always activated depending on the config if self.has_sliding_layers: causal_mask_mapping["sliding_attention"] = ( create_sliding_window_causal_mask(**mask_kwargs) ) hidden_states = inputs_embeds # create position embeddings to be shared across the decoder layers position_embeddings = self.rotary_emb(hidden_states, position_ids) # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None for decoder_layer in self.layers: if output_hidden_states: all_hidden_states += (hidden_states,) hidden_states = decoder_layer( hidden_states, attention_mask=causal_mask_mapping[decoder_layer.attention_type], position_ids=text_position_ids, past_key_values=past_key_values, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, **kwargs, ) hidden_states = self.norm(hidden_states) # add hidden states from the last decoder layer if output_hidden_states: all_hidden_states += (hidden_states,) if not return_dict: return tuple( v for v in [ hidden_states, past_key_values, all_hidden_states, all_self_attns, ] if v is not None ) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, hidden_states=all_hidden_states, attentions=all_self_attns, ) class Qwen2_5_VLModel(nn.Module): base_model_prefix = "" _checkpoint_conversion_mapping = {"^model": "language_model"} # Reference: fix gemma3 grad acc #37208 accepts_loss_kwargs = False _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] def __init__(self, config, enable_image_understanding: bool = False): super().__init__() self.language_model = Qwen2_5_VLTextModel(config.text_config) if enable_image_understanding: self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config( config.vision_config ) self.visual.to(torch.get_default_dtype()) self.rope_deltas = None # cache rope_deltas here self.config = config # Initialize weights and apply final processing # self.post_init() def get_input_embeddings(self): return self.language_model.embed_tokens def set_input_embeddings(self, value): self.language_model.embed_tokens = value def set_decoder(self, decoder): self.language_model = decoder def get_decoder(self): return self.language_model def get_rope_index( self, input_ids: Optional[torch.LongTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Calculate the 3D rope index based on image and video's temporal, height and width in LLM. Explanation: Each embedding sequence contains vision embedding and text embedding or just contains text embedding. For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. Examples: input_ids: [T T T T T], here T is for text. temporal position_ids: [0, 1, 2, 3, 4] height position_ids: [0, 1, 2, 3, 4] width position_ids: [0, 1, 2, 3, 4] For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part and 1D rotary position embedding for text part. Examples: Temporal (Time): 3 patches, representing different segments of the video in time. Height: 2 patches, dividing each frame vertically. Width: 2 patches, dividing each frame horizontally. We also have some important parameters: fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity. temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs. input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] text temporal position_ids: [101, 102, 103, 104, 105] text height position_ids: [101, 102, 103, 104, 105] text width position_ids: [101, 102, 103, 104, 105] Here we calculate the text start position_ids as the max vision position_ids plus 1. Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): The temporal, height and width of feature shape of each video in LLM. second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. Returns: position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) """ spatial_merge_size = self.config.vision_config.spatial_merge_size image_token_id = self.config.image_token_id video_token_id = self.config.video_token_id vision_start_token_id = self.config.vision_start_token_id mrope_position_deltas = [] if input_ids is not None and ( image_grid_thw is not None or video_grid_thw is not None ): total_input_ids = input_ids if attention_mask is None: attention_mask = torch.ones_like(total_input_ids) position_ids = torch.ones( 3, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device, ) image_index, video_index = 0, 0 attention_mask = attention_mask.to(total_input_ids.device) for i, input_ids in enumerate(total_input_ids): input_ids = input_ids[attention_mask[i] == 1] image_nums, video_nums = 0, 0 vision_start_indices = torch.argwhere( input_ids == vision_start_token_id ).squeeze(1) vision_tokens = input_ids[vision_start_indices + 1] image_nums = (vision_tokens == image_token_id).sum() video_nums = (vision_tokens == video_token_id).sum() input_tokens = input_ids.tolist() llm_pos_ids_list: list = [] st = 0 remain_images, remain_videos = image_nums, video_nums for _ in range(image_nums + video_nums): if image_token_id in input_tokens and remain_images > 0: ed_image = input_tokens.index(image_token_id, st) else: ed_image = len(input_tokens) + 1 if video_token_id in input_tokens and remain_videos > 0: ed_video = input_tokens.index(video_token_id, st) else: ed_video = len(input_tokens) + 1 if ed_image < ed_video: t, h, w = ( image_grid_thw[image_index][0], image_grid_thw[image_index][1], image_grid_thw[image_index][2], ) second_per_grid_t = 0 image_index += 1 remain_images -= 1 ed = ed_image else: t, h, w = ( video_grid_thw[video_index][0], video_grid_thw[video_index][1], video_grid_thw[video_index][2], ) if second_per_grid_ts is not None: second_per_grid_t = second_per_grid_ts[video_index] else: second_per_grid_t = 1.0 video_index += 1 remain_videos -= 1 ed = ed_video llm_grid_t, llm_grid_h, llm_grid_w = ( t.item(), h.item() // spatial_merge_size, w.item() // spatial_merge_size, ) text_len = ed - st st_idx = ( llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 ) llm_pos_ids_list.append( torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx ) range_tensor = torch.arange(llm_grid_t).view(-1, 1) expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) ## normalize type, send to device. second_per_grid_t = torch.as_tensor( second_per_grid_t, dtype=range_tensor.dtype, device=range_tensor.device, ) time_tensor = ( expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second ) time_tensor_long = time_tensor.long() t_index = time_tensor_long.flatten() h_index = ( torch.arange(llm_grid_h) .view(1, -1, 1) .expand(llm_grid_t, -1, llm_grid_w) .flatten() ) w_index = ( torch.arange(llm_grid_w) .view(1, 1, -1) .expand(llm_grid_t, llm_grid_h, -1) .flatten() ) llm_pos_ids_list.append( torch.stack([t_index, h_index, w_index]) + text_len + st_idx ) st = ed + llm_grid_t * llm_grid_h * llm_grid_w if st < len(input_tokens): st_idx = ( llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 ) text_len = len(input_tokens) - st llm_pos_ids_list.append( torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx ) llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( position_ids.device ) mrope_position_deltas.append( llm_positions.max() + 1 - len(total_input_ids[i]) ) mrope_position_deltas = torch.tensor( mrope_position_deltas, device=input_ids.device ).unsqueeze(1) return position_ids, mrope_position_deltas else: if attention_mask is not None: position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) position_ids = ( position_ids.unsqueeze(0) .expand(3, -1, -1) .to(attention_mask.device) ) max_position_ids = position_ids.max(0, keepdim=False)[0].max( -1, keepdim=True )[0] mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] else: position_ids = ( torch.arange(input_ids.shape[1], device=input_ids.device) .view(1, 1, -1) .expand(3, input_ids.shape[0], -1) ) mrope_position_deltas = torch.zeros( [input_ids.shape[0], 1], device=input_ids.device, dtype=input_ids.dtype, ) return position_ids, mrope_position_deltas def get_video_features( self, pixel_values_videos: torch.FloatTensor, video_grid_thw: Optional[torch.LongTensor] = None, ): """ Encodes videos into continuous embeddings that can be forwarded to the language model. Args: pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): The tensors corresponding to the input videos. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): The temporal, height and width of feature shape of each video in LLM. """ pixel_values_videos = pixel_values_videos.type(self.visual.dtype) video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) split_sizes = ( video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2 ).tolist() video_embeds = torch.split(video_embeds, split_sizes) return video_embeds def get_image_features( self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None, ): """ Encodes images into continuous embeddings that can be forwarded to the language model. Args: pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): The tensors corresponding to the input images. image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype) image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) if not isinstance(image_embeds, torch.Tensor): # In transformers v5, the visual encoder returns BaseModelOutputWithPooling. # pooler_output contains the spatially merged embeddings (what we need), # while last_hidden_state contains the raw unmerged output. image_embeds = image_embeds.pooler_output split_sizes = ( image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2 ).tolist() image_embeds = torch.split(image_embeds, split_sizes) return image_embeds def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor = None, video_features: torch.FloatTensor = None, ): """ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is equal to the length of multimodal features. If the lengths are different, an error is raised. """ if input_ids is None: special_image_mask = inputs_embeds == self.get_input_embeddings()( torch.tensor( self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device, ) ) special_image_mask = special_image_mask.all(-1) special_video_mask = inputs_embeds == self.get_input_embeddings()( torch.tensor( self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device, ) ) special_video_mask = special_video_mask.all(-1) else: special_image_mask = input_ids == self.config.image_token_id special_video_mask = input_ids == self.config.video_token_id n_image_tokens = special_image_mask.sum() special_image_mask = ( special_image_mask.unsqueeze(-1) .expand_as(inputs_embeds) .to(inputs_embeds.device) ) if ( image_features is not None and inputs_embeds[special_image_mask].numel() != image_features.numel() ): raise ValueError( f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}" ) n_video_tokens = special_video_mask.sum() special_video_mask = ( special_video_mask.unsqueeze(-1) .expand_as(inputs_embeds) .to(inputs_embeds.device) ) if ( video_features is not None and inputs_embeds[special_video_mask].numel() != video_features.numel() ): raise ValueError( f"Videos features and video tokens do not match: tokens: {n_video_tokens}, features {video_features.shape[0]}" ) return special_image_mask, special_video_mask def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, pixel_values: Optional[torch.Tensor] = None, pixel_values_videos: Optional[torch.FloatTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, rope_deltas: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None, **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple, Qwen2_5_VLModelOutputWithPast]: r""" image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): The temporal, height and width of feature shape of each video in LLM. rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): The rope index difference between sequence length and multimodal rope. second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. """ output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: image_embeds = self.get_image_features(pixel_values, image_grid_thw) image_embeds = torch.cat(image_embeds, dim=0).to( inputs_embeds.device, inputs_embeds.dtype ) image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds ) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) video_embeds = torch.cat(video_embeds, dim=0).to( inputs_embeds.device, inputs_embeds.dtype ) _, video_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds ) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) if position_ids is None: # Calculate RoPE index once per generation in the pre-fill stage only. # When compiling, we can't check tensor values thus we check only input length # It is safe to assume that `length!=1` means we're in pre-fill because compiled # models currently cannot do asssisted decoding prefill_compiled_stage = is_torchdynamo_compiling() and ( (input_ids is not None and input_ids.shape[1] != 1) or (inputs_embeds is not None and inputs_embeds.shape[1] != 1) ) prefill_noncompiled_stage = not is_torchdynamo_compiling() and ( (cache_position is not None and cache_position[0] == 0) or (past_key_values is None or past_key_values.get_seq_length() == 0) ) if ( prefill_compiled_stage or prefill_noncompiled_stage ) or self.rope_deltas is None: position_ids, rope_deltas = self.get_rope_index( input_ids, image_grid_thw, video_grid_thw, second_per_grid_ts=second_per_grid_ts, attention_mask=attention_mask, ) self.rope_deltas = rope_deltas else: batch_size, seq_length, _ = inputs_embeds.shape position_ids = torch.arange(seq_length, device=inputs_embeds.device) position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1) if cache_position is not None: delta = (cache_position[0] + self.rope_deltas).to( inputs_embeds.device ) else: delta = torch.zeros( (batch_size, seq_length), device=inputs_embeds.device ) delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=1) position_ids += delta.to(position_ids.device) outputs = self.language_model( input_ids=None, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_position, **kwargs, ) output = Qwen2_5_VLModelOutputWithPast( last_hidden_state=outputs.last_hidden_state, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, rope_deltas=self.rope_deltas, ) return output if return_dict else output.to_tuple() class Qwen2_5_VLForConditionalGeneration(TextEncoder): # BitandBytes specific attributes default_bitsandbytes_target_modules = [ ".gate_up_proj.", ".down_proj.", ".q_proj.", ".k_proj.", ".v_proj.", ".o_proj.", ] bitsandbytes_stacked_params_mapping = { # shard_name, weight_name, index "q_proj": ("qkv_proj", 0), "k_proj": ("qkv_proj", 1), "v_proj": ("qkv_proj", 2), "gate_proj": ("gate_up_proj", 0), "up_proj": ("gate_up_proj", 1), } def __init__( self, config: Qwen2_5VLConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__(config) enable_image_understanding = config.enable_image_understanding config = config.arch_config self.model = Qwen2_5_VLModel( config, enable_image_understanding=enable_image_understanding ) self.lm_head = nn.Linear( config.text_config.hidden_size, config.text_config.vocab_size, bias=False ) self.enable_image_understanding = enable_image_understanding self.config = config def get_input_embeddings(self): return self.model.embed_tokens @torch.no_grad() def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, pixel_values: Optional[torch.Tensor] = None, pixel_values_videos: Optional[torch.FloatTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, rope_deltas: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, **kwargs: Unpack[TransformersKwargs], ): """Run forward pass for Qwen2_5-VL. Args: input_ids: Flattened (concatenated) input_ids corresponding to a batch. positions: Flattened (concatenated) position ids corresponding to a batch. **NOTE**: If mrope is enabled (default setting for Qwen2-VL opensource models), the shape will be `(3, seq_len)`, otherwise it will be `(seq_len,). (Use input_metadata.mrope_positions to replace it) """ output_attentions = False output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, pixel_values_videos=pixel_values_videos, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, second_per_grid_ts=second_per_grid_ts, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_position, **kwargs, ) hidden_states = outputs[0] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss slice_indices = ( slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep ) logits = self.lm_head(hidden_states[:, slice_indices, :]) return Qwen2_5_VLCausalLMOutputWithPast( loss=None, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, rope_deltas=outputs.rope_deltas, ) def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): loaded_params: set[str] = set() params_dict = dict(self.named_parameters(remove_duplicate=False)) for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue name = name.replace("model.", "model.language_model.") if "visual." in name: if not self.enable_image_understanding: continue name = name.replace("visual.", "model.visual.") try: # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue param = params_dict[name] except KeyError: raise weight_loader = getattr(param, "weight_loader", default_weight_loader) loaded_weight = loaded_weight.to(param.dtype) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params def get_embed_and_head(self): return self.model.embed_tokens.weight, self.lm_head.weight EntryClass = Qwen2_5_VLForConditionalGeneration ================================================ FILE: python/sglang/multimodal_gen/runtime/models/encoders/qwen3.py ================================================ from collections.abc import Iterable from typing import Any import torch from torch import nn from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig from sglang.multimodal_gen.runtime.distributed import get_tp_world_size from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul from sglang.multimodal_gen.runtime.layers.attention import LocalAttention from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm from sglang.multimodal_gen.runtime.layers.linear import ( MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig from sglang.multimodal_gen.runtime.layers.rotary_embedding import get_rope from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from sglang.multimodal_gen.runtime.loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, ) from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder class Qwen3MLP(nn.Module): """Qwen3 MLP with SwiGLU activation and tensor parallelism.""" def __init__( self, hidden_size: int, intermediate_size: int, hidden_act: str, quant_config: QuantizationConfig | None = None, bias: bool = False, prefix: str = "", ) -> None: super().__init__() self.gate_up_proj = MergedColumnParallelLinear( input_size=hidden_size, output_sizes=[intermediate_size] * 2, bias=bias, quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", ) self.down_proj = RowParallelLinear( input_size=intermediate_size, output_size=hidden_size, bias=bias, quant_config=quant_config, prefix=f"{prefix}.down_proj", ) if hidden_act != "silu": raise ValueError( f"Unsupported activation: {hidden_act}. Only silu is supported." ) self.act_fn = SiluAndMul() def forward(self, x: torch.Tensor) -> torch.Tensor: x, _ = self.gate_up_proj(x) x = self.act_fn(x) x, _ = self.down_proj(x) return x class Qwen3Attention(nn.Module): """Qwen3 attention with QK-Norm and tensor parallelism. Key difference from LLaMA: RMSNorm is applied to Q and K before attention. """ def __init__( self, config: Qwen3TextConfig, hidden_size: int, num_heads: int, num_kv_heads: int, rope_theta: float = 1000000.0, rope_scaling: dict[str, Any] | None = None, max_position_embeddings: int = 40960, quant_config: QuantizationConfig | None = None, bias: bool = False, prefix: str = "", ) -> None: super().__init__() self.hidden_size = hidden_size tp_size = get_tp_world_size() self.total_num_heads = num_heads assert self.total_num_heads % tp_size == 0 self.num_heads = self.total_num_heads // tp_size self.total_num_kv_heads = num_kv_heads if self.total_num_kv_heads >= tp_size: assert self.total_num_kv_heads % tp_size == 0 else: assert tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) self.head_dim = getattr( config, "head_dim", self.hidden_size // self.total_num_heads ) self.rotary_dim = self.head_dim self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 self.rope_theta = rope_theta self.max_position_embeddings = max_position_embeddings # QKV projection with tensor parallelism self.qkv_proj = QKVParallelLinear( hidden_size=hidden_size, head_size=self.head_dim, total_num_heads=self.total_num_heads, total_num_kv_heads=self.total_num_kv_heads, bias=bias, quant_config=quant_config, prefix=f"{prefix}.qkv_proj", ) # Output projection self.o_proj = RowParallelLinear( input_size=self.total_num_heads * self.head_dim, output_size=hidden_size, bias=bias, quant_config=quant_config, prefix=f"{prefix}.o_proj", ) # QK-Norm: Key difference from LLaMA rms_norm_eps = getattr(config, "rms_norm_eps", 1e-6) self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) # Rotary embeddings self.rotary_emb = get_rope( self.head_dim, rotary_dim=self.rotary_dim, max_position=max_position_embeddings, base=int(rope_theta), rope_scaling=rope_scaling, is_neox_style=True, ) # Attention with FlashAttention/SageAttn support self.attn = LocalAttention( self.num_heads, self.head_dim, self.num_kv_heads, softmax_scale=self.scaling, causal=True, supported_attention_backends=config._supported_attention_backends, ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, ) -> torch.Tensor: # QKV projection qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # Reshape for QK-norm batch_size, seq_len = q.shape[0], q.shape[1] q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim) k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) # Apply QK-Norm (key difference from LLaMA) q = self.q_norm(q) k = self.k_norm(k) # Reshape back for rotary embeddings q = q.reshape(batch_size, seq_len, -1) k = k.reshape(batch_size, seq_len, -1) # Apply rotary embeddings q, k = self.rotary_emb(positions, q, k) # Reshape for attention q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim) k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) # Attention attn_output = self.attn(q, k, v) attn_output = attn_output.reshape(batch_size, seq_len, -1) # Output projection output, _ = self.o_proj(attn_output) return output class Qwen3DecoderLayer(nn.Module): """Qwen3 transformer decoder layer.""" def __init__( self, config: Qwen3TextConfig, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() self.hidden_size = config.hidden_size rope_theta = config.rope_parameters["rope_theta"] rope_scaling = config.rope_parameters max_position_embeddings = getattr(config, "max_position_embeddings", 40960) attention_bias = getattr(config, "attention_bias", False) self.self_attn = Qwen3Attention( config=config, hidden_size=self.hidden_size, num_heads=config.num_attention_heads, num_kv_heads=getattr( config, "num_key_value_heads", config.num_attention_heads ), rope_theta=rope_theta, rope_scaling=rope_scaling, max_position_embeddings=max_position_embeddings, quant_config=quant_config, bias=attention_bias, prefix=f"{prefix}.self_attn", ) self.mlp = Qwen3MLP( hidden_size=self.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, bias=getattr(config, "mlp_bias", False), prefix=f"{prefix}.mlp", ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: # Self Attention if residual is None: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: hidden_states, residual = self.input_layernorm(hidden_states, residual) hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) # MLP hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) hidden_states = self.mlp(hidden_states) return hidden_states, residual class Qwen3ForCausalLM(TextEncoder): """Qwen3 causal language model for text encoding in diffusion models. Features: - Tensor parallelism support - FlashAttention/SageAttn/SDPA support via LocalAttention - QK-Norm for better training stability - FSDP sharding for CPU offload """ def __init__(self, config: Qwen3TextConfig) -> None: super().__init__(config) self.config = config self.quant_config = config.quant_config # Embedding layer with tensor parallelism if config.lora_config is not None: max_loras = getattr(config.lora_config, "max_loras", 1) lora_vocab_size = getattr(config.lora_config, "lora_extra_vocab_size", 1) lora_vocab = lora_vocab_size * max_loras else: lora_vocab = 0 self.vocab_size = config.vocab_size + lora_vocab self.org_vocab_size = config.vocab_size self.embed_tokens = VocabParallelEmbedding( self.vocab_size, config.hidden_size, org_num_embeddings=config.vocab_size, quant_config=config.quant_config, ) # Transformer layers self.layers = nn.ModuleList( [ Qwen3DecoderLayer( config=config, quant_config=config.quant_config, prefix=f"{config.prefix}.layers.{i}", ) for i in range(config.num_hidden_layers) ] ) # Final layer norm self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) def forward( self, input_ids: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs, ) -> BaseEncoderOutput: output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) if inputs_embeds is not None: hidden_states = inputs_embeds else: hidden_states = self.get_input_embeddings(input_ids) residual = None if position_ids is None: position_ids = torch.arange( 0, hidden_states.shape[1], device=hidden_states.device ).unsqueeze(0) all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None for layer in self.layers: if all_hidden_states is not None: all_hidden_states += ( (hidden_states,) if residual is None else (hidden_states + residual,) ) hidden_states, residual = layer(position_ids, hidden_states, residual) hidden_states, _ = self.norm(hidden_states, residual) # Add hidden states from the last decoder layer if all_hidden_states is not None: all_hidden_states += (hidden_states,) return BaseEncoderOutput( last_hidden_state=hidden_states, hidden_states=all_hidden_states, ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights with support for tensor parallelism and weight remapping.""" params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: # Strip 'model.' prefix from HuggingFace Qwen3 weights if name.startswith("model."): name = name[6:] # len("model.") == 6 # Skip rotary embedding weights if "rotary_emb.inv_freq" in name: continue if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: continue # Handle KV scale remapping if "scale" in name: kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict) if kv_scale_name is None: continue else: name = kv_scale_name # Handle stacked params mapping (qkv_proj, gate_up_proj) for ( param_name, weight_name, shard_id, ) in self.config.arch_config.stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) break else: # Skip loading extra bias for GPTQ models if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: continue param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params EntryClass = Qwen3ForCausalLM ================================================ FILE: python/sglang/multimodal_gen/runtime/models/encoders/t5.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from transformers: https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/t5/modeling_t5.py # Derived from T5 implementation posted on HuggingFace; license below: # # coding=utf-8 # Copyright 2018 Mesh TensorFlow authors, T5 Authors and HuggingFace Inc. team. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """PyTorch T5 & UMT5 model.""" import math from collections.abc import Iterable from dataclasses import dataclass import torch import torch.nn.functional as F from torch import nn from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput, T5Config from sglang.multimodal_gen.runtime.distributed import _get_folding_tp_group from sglang.multimodal_gen.runtime.layers.activation import get_act_fn from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm from sglang.multimodal_gen.runtime.layers.linear import ( MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig from sglang.multimodal_gen.runtime.layers.utils import get_group_rank, get_group_size from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder from sglang.multimodal_gen.runtime.platforms import current_platform class AttentionType: """ Attention type. Use string to be compatible with `torch.compile`. """ # Decoder attention between previous layer Q/K/V DECODER = "decoder" # Encoder attention between previous layer Q/K/V for encoder-decoder ENCODER = "encoder" # Encoder attention between previous layer Q/K/V ENCODER_ONLY = "encoder_only" # Attention between dec. Q and enc. K/V for encoder-decoder ENCODER_DECODER = "encoder_decoder" @dataclass class AttentionMetadata: attn_bias: torch.Tensor class T5DenseActDense(nn.Module): def __init__( self, config: T5Config, quant_config: QuantizationConfig | None = None ): super().__init__() tp_group = _get_folding_tp_group(config) self.wi = MergedColumnParallelLinear( config.d_model, [config.d_ff], bias=False, tp_group=tp_group ) self.wo = RowParallelLinear( config.d_ff, config.d_model, bias=False, quant_config=quant_config, tp_group=tp_group, ) self.act = get_act_fn(config.dense_act_fn) def forward(self, hidden_states) -> torch.Tensor: hidden_states, _ = self.wi(hidden_states) hidden_states = self.act(hidden_states) hidden_states, _ = self.wo(hidden_states) return hidden_states class T5DenseGatedActDense(nn.Module): def __init__( self, config: T5Config, quant_config: QuantizationConfig | None = None ): super().__init__() tp_group = _get_folding_tp_group(config) self.wi_0 = MergedColumnParallelLinear( config.d_model, [config.d_ff], bias=False, quant_config=quant_config, tp_group=tp_group, ) self.wi_1 = MergedColumnParallelLinear( config.d_model, [config.d_ff], bias=False, quant_config=quant_config, tp_group=tp_group, ) # Should not run in fp16 unless mixed-precision is used, # see https://github.com/huggingface/transformers/issues/20287. self.wo = RowParallelLinear( config.d_ff, config.d_model, bias=False, quant_config=quant_config, tp_group=tp_group, ) self.act = get_act_fn(config.dense_act_fn) def forward(self, hidden_states) -> torch.Tensor: hidden_gelu = self.act(self.wi_0(hidden_states)[0]) hidden_linear, _ = self.wi_1(hidden_states) hidden_states = hidden_gelu * hidden_linear hidden_states, _ = self.wo(hidden_states) return hidden_states class T5LayerFF(nn.Module): def __init__( self, config: T5Config, quant_config: QuantizationConfig | None = None ): super().__init__() if config.is_gated_act: self.DenseReluDense = T5DenseGatedActDense( config, quant_config=quant_config ) else: self.DenseReluDense = T5DenseActDense(config, quant_config=quant_config) self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon) def forward(self, hidden_states) -> torch.Tensor: forwarded_states = self.layer_norm(hidden_states) forwarded_states = self.DenseReluDense(forwarded_states) hidden_states = hidden_states + forwarded_states return hidden_states # T5 has attn_bias and does not use softmax scaling class T5MultiHeadAttention(nn.Module): def __init__(self) -> None: super().__init__() def forward(self, q, k, v, attn_bias=None): b, _, n, c = q.shape attn = torch.einsum("binc,bjnc->bnij", q, k) if attn_bias is not None: attn += attn_bias attn = F.softmax(attn.float(), dim=-1).type_as(attn) x = torch.einsum("bnij,bjnc->binc", attn, v) x = x.reshape(b, -1, n * c) return x class T5Attention(nn.Module): def __init__( self, config: T5Config, attn_type: str, has_relative_attention_bias=False, quant_config: QuantizationConfig | None = None, prefix: str = "", ): super().__init__() self.attn_type = attn_type # Cross-attention has no relative pos encoding anyway self.is_decoder = attn_type == AttentionType.DECODER self.has_relative_attention_bias = has_relative_attention_bias self.relative_attention_num_buckets = config.relative_attention_num_buckets self.relative_attention_max_distance = config.relative_attention_max_distance self.d_model = config.d_model self.key_value_proj_dim = config.d_kv self.total_num_heads = self.total_num_kv_heads = config.num_heads # Partition heads across multiple tensor parallel GPUs. self.tp_group = _get_folding_tp_group(config) self.tp_world_size = get_group_size(self.tp_group) assert config.num_heads % self.tp_world_size == 0 self.n_heads = config.num_heads // self.tp_world_size self.inner_dim = self.n_heads * self.key_value_proj_dim # No GQA in t5. # self.n_kv_heads = self.n_heads self.qkv_proj = QKVParallelLinear( self.d_model, self.key_value_proj_dim, self.total_num_heads, self.total_num_kv_heads, bias=False, quant_config=quant_config, prefix=f"{prefix}.qkv_proj", tp_group=self.tp_group, ) self.attn = T5MultiHeadAttention() if self.has_relative_attention_bias: self.relative_attention_bias = VocabParallelEmbedding( self.relative_attention_num_buckets, self.total_num_heads, org_num_embeddings=self.relative_attention_num_buckets, padding_size=self.relative_attention_num_buckets, quant_config=quant_config, tp_group=self.tp_group, ) self.o = RowParallelLinear( self.total_num_heads * self.key_value_proj_dim, self.d_model, bias=False, quant_config=quant_config, prefix=f"{prefix}.o_proj", tp_group=self.tp_group, ) @staticmethod def _relative_position_bucket( relative_position, bidirectional=True, num_buckets=32, max_distance=128 ) -> torch.Tensor: """ Adapted from Mesh Tensorflow: https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 Translate relative position to a bucket number for relative attention. The relative position is defined as memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for small absolute relative_position and larger buckets for larger absolute relative_positions. All relative positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. This should allow for more graceful generalization to longer sequences than the model has been trained on Args: relative_position: an int32 Tensor bidirectional: a boolean - whether the attention is bidirectional num_buckets: an integer max_distance: an integer Returns: a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) """ # noqa: E501 relative_buckets = 0 if bidirectional: num_buckets //= 2 relative_buckets += (relative_position > 0).to(torch.long) * num_buckets relative_position = torch.abs(relative_position) else: relative_position = -torch.min( relative_position, torch.zeros_like(relative_position) ) # now relative_position is in the range [0, inf) # half of the buckets are for exact increments in positions max_exact = num_buckets // 2 is_small = relative_position < max_exact # The other half of the buckets are for logarithmically bigger bins # in positions up to max_distance relative_position_if_large = max_exact + ( torch.log(relative_position.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact) ).to(torch.long) relative_position_if_large = torch.min( relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1), ) relative_buckets += torch.where( is_small, relative_position, relative_position_if_large ) return relative_buckets def compute_bias(self, query_length, key_length, device=None) -> torch.Tensor: """Compute binned relative position bias""" if device is None: device = self.relative_attention_bias.weight.device context_position = torch.arange(query_length, dtype=torch.long, device=device)[ :, None ] memory_position = torch.arange(key_length, dtype=torch.long, device=device)[ None, : ] # max_seq_len, nh relative_position = memory_position - context_position relative_position_bucket = self._relative_position_bucket( relative_position, # shape (query_length, key_length) bidirectional=(not self.is_decoder), num_buckets=self.relative_attention_num_buckets, max_distance=self.relative_attention_max_distance, ) values = self.relative_attention_bias( relative_position_bucket ) # shape (query_length, key_length, num_heads) x = values.permute([2, 0, 1]).unsqueeze( 0 ) # shape (1, num_heads, query_length, key_length) return x def forward( self, hidden_states: torch.Tensor, # (num_tokens, d_model) attention_mask: torch.Tensor, attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: bs, seq_len, _ = hidden_states.shape num_seqs = bs n, c = ( self.n_heads, self.key_value_proj_dim, ) qkv, _ = self.qkv_proj(hidden_states) # Projection of 'own' hidden state (self-attention). No GQA here. q, k, v = qkv.split(self.inner_dim, dim=-1) q = q.reshape(bs, seq_len, n, c) k = k.reshape(bs, seq_len, n, c) v = v.reshape(bs, seq_len, n, c) assert attn_metadata is not None attn_bias = attn_metadata.attn_bias # Not compatible with CP here (as all encoder-decoder models), # as it assumes homogeneous batch (prefills or decodes). if self.has_relative_attention_bias: # Self-attention. Compute T5 relative positional encoding. # The bias term is computed on longest sequence in batch. Biases # for shorter sequences are slices of the longest. assert self.attn_type == AttentionType.ENCODER attn_bias = self.compute_bias(seq_len, seq_len).repeat(num_seqs, 1, 1, 1) attn_metadata.attn_bias = attn_bias else: # Encoder/Decoder Self-Attention Layer, attn bias already cached. assert attn_bias is not None if attention_mask is not None: attention_mask = ( attention_mask.view(bs, 1, 1, -1) if attention_mask.ndim == 2 else attention_mask.unsqueeze(1) ) mask_val = -1e4 if current_platform.is_mps() else torch.finfo(q.dtype).min attn_bias.masked_fill_(attention_mask == 0, mask_val) if self.tp_world_size > 1: rank = get_group_rank(self.tp_group) attn_bias = attn_bias[ :, rank * self.n_heads : (rank + 1) * self.n_heads, :, : ] attn_output = self.attn(q, k, v, attn_bias) output, _ = self.o(attn_output) return output class T5LayerSelfAttention(nn.Module): def __init__( self, config, has_relative_attention_bias=False, quant_config: QuantizationConfig | None = None, prefix: str = "", ): super().__init__() self.SelfAttention = T5Attention( config, AttentionType.DECODER if "decoder" in prefix else AttentionType.ENCODER, has_relative_attention_bias=has_relative_attention_bias, quant_config=quant_config, prefix=f"{prefix}.SelfAttention", ) self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: normed_hidden_states = self.layer_norm(hidden_states) attention_output = self.SelfAttention( hidden_states=normed_hidden_states, attention_mask=attention_mask, attn_metadata=attn_metadata, ) hidden_states = hidden_states + attention_output return hidden_states class T5LayerCrossAttention(nn.Module): def __init__( self, config, quant_config: QuantizationConfig | None = None, prefix: str = "" ): super().__init__() self.EncDecAttention = T5Attention( config, AttentionType.ENCODER_DECODER, has_relative_attention_bias=False, quant_config=quant_config, prefix=f"{prefix}.EncDecAttention", ) self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon) def forward( self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: normed_hidden_states = self.layer_norm(hidden_states) attention_output = self.EncDecAttention( hidden_states=normed_hidden_states, attn_metadata=attn_metadata, ) hidden_states = hidden_states + attention_output return hidden_states class T5Block(nn.Module): def __init__( self, config: T5Config, is_decoder: bool, has_relative_attention_bias=False, quant_config: QuantizationConfig | None = None, prefix: str = "", ): super().__init__() self.is_decoder = is_decoder self.layer = nn.ModuleList() self.layer.append( T5LayerSelfAttention( config, has_relative_attention_bias=has_relative_attention_bias, quant_config=quant_config, prefix=f"{prefix}.self_attn", ) ) if self.is_decoder: self.layer.append( T5LayerCrossAttention( config, quant_config=quant_config, prefix=f"{prefix}.cross_attn" ) ) self.layer.append(T5LayerFF(config, quant_config=quant_config)) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: if attention_mask is None: attention_mask = torch.ones( hidden_states.shape[:2], device=hidden_states.device ) hidden_states = self.layer[0]( hidden_states=hidden_states, attention_mask=attention_mask, attn_metadata=attn_metadata, ) if self.is_decoder: hidden_states = self.layer[1]( hidden_states=hidden_states, attn_metadata=attn_metadata ) # Apply Feed Forward layer hidden_states = self.layer[-1](hidden_states) return hidden_states class T5Stack(nn.Module): def __init__( self, config: T5Config, is_decoder: bool, n_layers: int, embed_tokens=None, quant_config: QuantizationConfig | None = None, prefix: str = "", is_umt5: bool = False, ): super().__init__() self.embed_tokens = embed_tokens self.is_umt5 = is_umt5 if is_umt5: self.block = nn.ModuleList( [ T5Block( config, is_decoder=is_decoder, has_relative_attention_bias=True, quant_config=quant_config, prefix=f"{prefix}.blocks.{i}", ) for i in range(n_layers) ] ) else: # Only the first block has relative positional encoding. self.block = nn.ModuleList( [ T5Block( config, is_decoder=is_decoder, has_relative_attention_bias=i == 0, quant_config=quant_config, prefix=f"{prefix}.blocks.{i}", ) for i in range(n_layers) ] ) self.final_layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon) def forward( self, input_ids: torch.Tensor, attention_mask: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: hidden_states = self.embed_tokens(input_ids) for idx, block in enumerate(self.block): hidden_states = block( hidden_states=hidden_states, attention_mask=attention_mask, attn_metadata=attn_metadata, ) hidden_states = self.final_layer_norm(hidden_states) return hidden_states class T5EncoderModel(TextEncoder): def __init__(self, config: T5Config, prefix: str = ""): super().__init__(config) quant_config = None tp_group = _get_folding_tp_group(config) self.shared = VocabParallelEmbedding( config.vocab_size, config.d_model, org_num_embeddings=config.vocab_size, tp_group=tp_group, ) self.encoder = T5Stack( config, False, config.num_layers, self.shared, quant_config=quant_config, prefix=f"{prefix}.encoder", is_umt5=False, ) def get_input_embeddings(self): return self.shared def forward( self, input_ids: torch.Tensor | None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs, ) -> BaseEncoderOutput: attn_metadata = AttentionMetadata(None) hidden_states = self.encoder( input_ids=input_ids, attention_mask=attention_mask, attn_metadata=attn_metadata, ) return BaseEncoderOutput(last_hidden_state=hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = [ # (param_name, shard_name, shard_id) (".qkv_proj", ".q", "q"), (".qkv_proj", ".k", "k"), (".qkv_proj", ".v", "v"), ] params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: loaded = False if "decoder" in name or "lm_head" in name: continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) loaded = True break if not loaded: # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: continue param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params class UMT5EncoderModel(TextEncoder): def __init__(self, config: T5Config, prefix: str = ""): super().__init__(config) quant_config = None tp_group = _get_folding_tp_group(config) self.shared = VocabParallelEmbedding( config.vocab_size, config.d_model, org_num_embeddings=config.vocab_size, tp_group=tp_group, ) self.encoder = T5Stack( config, False, config.num_layers, self.shared, quant_config=quant_config, prefix=f"{prefix}.encoder", is_umt5=True, ) def get_input_embeddings(self): return self.shared def forward( self, input_ids: torch.Tensor | None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs, ) -> BaseEncoderOutput: attn_metadata = AttentionMetadata(None) hidden_states = self.encoder( input_ids=input_ids, attention_mask=attention_mask, attn_metadata=attn_metadata, ) return BaseEncoderOutput( last_hidden_state=hidden_states, attention_mask=attention_mask, ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: loaded = False if "decoder" in name or "lm_head" in name: continue for ( param_name, weight_name, shard_id, ) in self.config.arch_config.stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) loaded = True break if not loaded: # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: continue param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params EntryClass = [T5EncoderModel, UMT5EncoderModel] ================================================ FILE: python/sglang/multimodal_gen/runtime/models/encoders/vision.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/vision.py from abc import ABC, abstractmethod from typing import Generic, TypeVar import torch from transformers import PretrainedConfig from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) _C = TypeVar("_C", bound=PretrainedConfig) class VisionEncoderInfo(ABC, Generic[_C]): def __init__(self, vision_config: _C) -> None: super().__init__() self.vision_config = vision_config @abstractmethod def get_num_image_tokens( self, *, image_width: int, image_height: int, ) -> int: raise NotImplementedError @abstractmethod def get_max_image_tokens(self) -> int: raise NotImplementedError @abstractmethod def get_image_size(self) -> int: raise NotImplementedError @abstractmethod def get_patch_size(self) -> int: raise NotImplementedError @abstractmethod def get_patch_grid_length(self) -> int: raise NotImplementedError def resolve_visual_encoder_outputs( encoder_outputs: torch.Tensor | list[torch.Tensor], feature_sample_layers: list[int] | None, post_layer_norm: torch.nn.LayerNorm | None, max_possible_layers: int, ) -> torch.Tensor: """Given the outputs a visual encoder module that may correspond to the output of the last layer, or a list of hidden states to be stacked, handle post normalization and resolve it into a single output tensor. Args: encoder_outputs: Output of encoder's last layer or all hidden states. feature_sample_layers: Optional layer indices to grab from the encoder outputs; if provided, encoder outputs must be a list. post_layer_norm: Post norm to apply to the output of the encoder. max_possible_layers: Total layers in the fully loaded visual encoder. """ if feature_sample_layers is None: if post_layer_norm is not None: return post_layer_norm(encoder_outputs) return encoder_outputs # Get the hidden states corresponding to the layer indices. # Negative values are relative to the full visual encoder, # so offset them depending on how many layers were loaded. # NOTE: this assumes that encoder_outputs is a list containing # the inputs to the visual encoder, followed by the hidden states # of each layer. num_loaded_layers = len(encoder_outputs) - 1 offset = max_possible_layers - num_loaded_layers hs_pool = [ ( encoder_outputs[layer_idx] if layer_idx >= 0 else encoder_outputs[layer_idx + offset] ) for layer_idx in feature_sample_layers ] # Apply post-norm on the final hidden state if we are using it uses_last_layer = feature_sample_layers[-1] in (len(hs_pool) - 1, -1) if post_layer_norm is not None and uses_last_layer: hs_pool[-1] = post_layer_norm(encoder_outputs) return torch.cat(hs_pool, dim=-1) ================================================ FILE: python/sglang/multimodal_gen/runtime/models/parameter.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/parameter.py from collections.abc import Callable from fractions import Fraction from typing import Any import torch from torch.nn import Parameter from sglang.multimodal_gen.runtime.distributed import get_tp_rank from sglang.multimodal_gen.runtime.models.utils import _make_synced_weight_loader from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class BasevLLMParameter(Parameter): """ Base parameter for vLLM linear layers. Extends the torch.nn.parameter by taking in a linear weight loader. Will copy the loaded weight into the parameter when the provided weight loader is called. """ def __new__(cls, data: torch.Tensor, **kwargs): return super().__new__(cls, data=data, requires_grad=False) def __init__(self, data: torch.Tensor, weight_loader: Callable): """ Initialize the BasevLLMParameter :param data: torch tensor with the parameter data :param weight_loader: weight loader callable :returns: a torch.nn.parameter """ # During weight loading, we often do something like: # narrowed_tensor = param.data.narrow(0, offset, len) # narrowed_tensor.copy_(real_weight) # expecting narrowed_tensor and param.data to share the same storage. # However, on TPUs, narrowed_tensor will lazily propagate to the base # tensor, which is param.data, leading to the redundant memory usage. # This sometimes causes OOM errors during model loading. To avoid this, # we sync the param tensor after its weight loader is called. from sglang.multimodal_gen.runtime.platforms import current_platform if current_platform.is_tpu(): weight_loader = _make_synced_weight_loader(weight_loader) self._weight_loader = weight_loader @property def weight_loader(self): return self._weight_loader def _is_1d_and_scalar(self, loaded_weight: torch.Tensor): cond1 = self.data.ndim == 1 and self.data.numel() == 1 cond2 = loaded_weight.ndim == 0 and loaded_weight.numel() == 1 return cond1 and cond2 def _assert_and_load(self, loaded_weight: torch.Tensor) -> None: assert self.data.shape == loaded_weight.shape or self._is_1d_and_scalar( loaded_weight ) self.data.copy_(loaded_weight) def load_column_parallel_weight(self, loaded_weight: torch.Tensor) -> None: self._assert_and_load(loaded_weight) def load_row_parallel_weight(self, loaded_weight: torch.Tensor) -> None: self._assert_and_load(loaded_weight) def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None: self._assert_and_load(loaded_weight) def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None: self._assert_and_load(loaded_weight) class _ColumnvLLMParameter(BasevLLMParameter): """ Private class defining weight loading functionality (load_merged_column_weight, load_qkv_weight) for parameters being loaded into linear layers with column parallelism. This includes QKV and MLP layers which are not already fused on disk. Requires an output dimension to be defined. Called within the weight loader of each of the column parallel linear layers. """ def __init__(self, output_dim: int, **kwargs): self._output_dim = output_dim super().__init__(**kwargs) @property def output_dim(self): return self._output_dim def load_column_parallel_weight(self, loaded_weight: torch.Tensor) -> None: tp_rank = get_tp_rank() shard_size = self.data.shape[self.output_dim] loaded_weight = loaded_weight.narrow( self.output_dim, tp_rank * shard_size, shard_size ) assert self.data.shape == loaded_weight.shape self.data.copy_(loaded_weight) def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None: shard_offset = kwargs.get("shard_offset") shard_size = kwargs.get("shard_size") if shard_offset is None or shard_size is None: raise ValueError("shard_offset and shard_size must be provided") if ( isinstance(self, PackedColumnParameter | PackedvLLMParameter) and self.packed_dim == self.output_dim ): shard_size, shard_offset = self.adjust_shard_indexes_for_packing( shard_offset=shard_offset, shard_size=shard_size ) param_data = self.data tp_rank = get_tp_rank() param_data = param_data.narrow(self.output_dim, shard_offset, shard_size) loaded_weight = loaded_weight.narrow( self.output_dim, tp_rank * shard_size, shard_size ) assert param_data.shape == loaded_weight.shape param_data.copy_(loaded_weight) def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs) -> None: shard_offset = kwargs.get("shard_offset") shard_size = kwargs.get("shard_size") shard_id = kwargs.get("shard_id") num_heads = kwargs.get("num_heads") assert shard_offset is not None assert shard_size is not None assert shard_id is not None assert num_heads is not None if ( isinstance(self, PackedColumnParameter | PackedvLLMParameter) and self.output_dim == self.packed_dim ): shard_size, shard_offset = self.adjust_shard_indexes_for_packing( shard_offset=shard_offset, shard_size=shard_size ) param_data = self.data tp_rank = get_tp_rank() shard_id = tp_rank if shard_id == "q" else tp_rank // num_heads param_data = param_data.narrow(self.output_dim, shard_offset, shard_size) loaded_weight = loaded_weight.narrow( self.output_dim, shard_id * shard_size, shard_size ) assert param_data.shape == loaded_weight.shape param_data.copy_(loaded_weight) class RowvLLMParameter(BasevLLMParameter): """ Parameter class defining weight_loading functionality (load_row_parallel_weight) for parameters being loaded into linear layers with row parallel functionality. Requires an input_dim to be defined. """ def __init__(self, input_dim: int, **kwargs): self._input_dim = input_dim super().__init__(**kwargs) @property def input_dim(self): return self._input_dim def load_row_parallel_weight(self, loaded_weight: torch.Tensor) -> None: tp_rank = get_tp_rank() shard_size = self.data.shape[self.input_dim] loaded_weight = loaded_weight.narrow( self.input_dim, tp_rank * shard_size, shard_size ) if len(loaded_weight.shape) == 0: loaded_weight = loaded_weight.reshape(1) assert self.data.shape == loaded_weight.shape self.data.copy_(loaded_weight) class ModelWeightParameter(_ColumnvLLMParameter, RowvLLMParameter): """ Parameter class for linear layer weights. Uses both column and row parallelism. """ pass class GroupQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): """ Parameter class for weight scales loaded for weights with grouped quantization. Uses both column and row parallelism. """ pass class ChannelQuantScaleParameter(_ColumnvLLMParameter): """ Parameter class for weight scales loaded for weights with channel-wise quantization. Equivalent to _ColumnvLLMParameter. """ pass class PerTensorScaleParameter(BasevLLMParameter): """ Parameter class for scales where the number of scales is equivalent to the number of logical matrices in fused linear layers (e.g. for QKV, there are 3 scales loaded from disk). This is relevant to weights with per-tensor quantization. Adds functionality to map the scalers to a shard during weight loading. Note: additional parameter manipulation may be handled for each quantization config specifically, within process_weights_after_loading """ def __init__(self, **kwargs): self.qkv_idxs = {"q": 0, "k": 1, "v": 2} super().__init__(**kwargs) def _shard_id_as_int(self, shard_id: str | int) -> int: if isinstance(shard_id, int): return shard_id # if not int, assume shard_id for qkv # map to int and return assert isinstance(shard_id, str) assert shard_id in self.qkv_idxs return self.qkv_idxs[shard_id] # For row parallel layers, no sharding needed # load weight into parameter as is def load_row_parallel_weight(self, *args, **kwargs) -> None: super().load_row_parallel_weight(*args, **kwargs) def load_merged_column_weight(self, *args, **kwargs) -> None: self._load_into_shard_id(*args, **kwargs) def load_qkv_weight(self, *args, **kwargs) -> None: self._load_into_shard_id(*args, **kwargs) def load_column_parallel_weight(self, *args, **kwargs) -> None: super().load_row_parallel_weight(*args, **kwargs) def _load_into_shard_id( self, loaded_weight: torch.Tensor, shard_id: str | int, **kwargs ): """ Slice the parameter data based on the shard id for loading. """ param_data = self.data shard_id = self._shard_id_as_int(shard_id) # AutoFP8 scales do not have a shape # compressed-tensors scales do have a shape if len(loaded_weight.shape) != 0: assert loaded_weight.shape[0] == 1 loaded_weight = loaded_weight[0] param_data = param_data[shard_id] assert param_data.shape == loaded_weight.shape param_data.copy_(loaded_weight) class PackedColumnParameter(_ColumnvLLMParameter): """ Parameter for model parameters which are packed on disk and support column parallelism only. See PackedvLLMParameter for more details on the packed properties. """ def __init__(self, packed_factor: int | Fraction, packed_dim: int, **kwargs): self._packed_factor = packed_factor self._packed_dim = packed_dim super().__init__(**kwargs) @property def packed_dim(self): return self._packed_dim @property def packed_factor(self): return self._packed_factor def adjust_shard_indexes_for_packing( self, shard_size, shard_offset ) -> tuple[Any, Any]: return _adjust_shard_indexes_for_packing( shard_size=shard_size, shard_offset=shard_offset, packed_factor=self.packed_factor, ) class PackedvLLMParameter(ModelWeightParameter): """ Parameter for model weights which are packed on disk. Example: GPTQ Marlin weights are int4 or int8, packed into int32. Extends the ModelWeightParameter to take in the packed factor, the packed dimension, and optionally, marlin tile size for marlin kernels. Adjusts the shard_size and shard_offset for fused linear layers model weight loading by accounting for packing and optionally, marlin tile size. """ def __init__(self, packed_factor: int | Fraction, packed_dim: int, **kwargs): self._packed_factor = packed_factor self._packed_dim = packed_dim super().__init__(**kwargs) @property def packed_dim(self): return self._packed_dim @property def packed_factor(self): return self._packed_factor def adjust_shard_indexes_for_packing(self, shard_size, shard_offset): return _adjust_shard_indexes_for_packing( shard_size=shard_size, shard_offset=shard_offset, packed_factor=self.packed_factor, ) class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter): """ Parameter class for weight scales loaded for weights with block-wise quantization. Uses both column and row parallelism. """ pass def permute_param_layout_( param: BasevLLMParameter, input_dim: int, output_dim: int, **kwargs ) -> BasevLLMParameter: """ Permute a parameter's layout to the specified input and output dimensions, useful for forcing the parameter into a known layout, for example, if I need a packed (quantized) weight matrix to be in the layout {input_dim = 0, output_dim = 1, packed_dim = 0} then I can call: permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) to ensure x is in the correct layout (permuting it to the correct layout if required, asserting if it cannot get it to the correct layout) """ curr_input_dim = getattr(param, "input_dim", None) curr_output_dim = getattr(param, "output_dim", None) if curr_input_dim is None or curr_output_dim is None: assert param.data.dim() == 2, ( "permute_param_layout_ only supports 2D parameters when either " "input_dim or output_dim is not set" ) # if one of the dimensions is not set, set it to the opposite of the other # we can only do this since we asserted the parameter is 2D above if curr_input_dim is None: assert curr_output_dim is not None, "either input or output dim must be set" curr_input_dim = (curr_output_dim + 1) % 2 if curr_output_dim is None: assert curr_input_dim is not None, "either input or output dim must be set" curr_output_dim = (curr_input_dim + 1) % 2 # create permutation from the current layout to the layout with # self.input_dim at input_dim and self.output_dim at output_dim preserving # other dimensions perm = [ i for i in range(param.data.dim()) if i not in [curr_input_dim, curr_output_dim] ] perm.insert(input_dim, curr_input_dim) perm.insert(output_dim, curr_output_dim) if "packed_dim" in kwargs: assert ( hasattr(param, "packed_dim") and param.packed_dim == perm[kwargs["packed_dim"]] ), "permute_param_layout_ currently doesn't support repacking" param.data = param.data.permute(*perm) if hasattr(param, "_input_dim"): param._input_dim = input_dim if hasattr(param, "_output_dim"): param._output_dim = output_dim if "packed_dim" in kwargs and hasattr(param, "_packed_dim"): param._packed_dim = kwargs["packed_dim"] return param def _adjust_shard_indexes_for_packing( shard_size, shard_offset, packed_factor ) -> tuple[Any, Any]: shard_size = shard_size // packed_factor shard_offset = shard_offset // packed_factor return shard_size, shard_offset ================================================ FILE: python/sglang/multimodal_gen/runtime/models/registry.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/registry.py import ast import importlib import os import pickle import subprocess import sys import tempfile from abc import ABC, abstractmethod from collections.abc import Callable, Set from dataclasses import dataclass, field from functools import lru_cache from typing import NoReturn, TypeVar, cast import cloudpickle from torch import nn from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) MODELS_PATH = os.path.dirname(__file__) COMPONENT_DIRS = [ d for d in os.listdir(MODELS_PATH) if os.path.isdir(os.path.join(MODELS_PATH, d)) and not d.startswith("__") and not d.startswith(".") ] _IMAGE_ENCODER_MODELS: dict[str, tuple] = { # "HunyuanVideoTransformer3DModel": ("image_encoder", "hunyuanvideo", "HunyuanVideoImageEncoder"), "CLIPVisionModelWithProjection": ("encoders", "clip", "CLIPVisionModel"), } # Global alias mapping: external_path -> canonical_class_name _ALIAS_TO_MODEL: dict[str, str] = {} def _parse_aliases_from_ast(value_node: ast.expr) -> list[str]: """Parse _aliases list from AST node.""" aliases = [] if isinstance(value_node, (ast.List, ast.Tuple)): for elt in value_node.elts: if isinstance(elt, ast.Constant) and isinstance(elt.value, str): aliases.append(elt.value) return aliases @lru_cache(maxsize=None) def _discover_and_register_models() -> dict[str, tuple[str, str, str]]: discovered_models = dict(_IMAGE_ENCODER_MODELS) # Collect class definitions with their _aliases class_aliases: dict[str, list[str]] = {} for component in COMPONENT_DIRS: component_path = os.path.join(MODELS_PATH, component) for filename in os.listdir(component_path): if not filename.endswith(".py"): continue mod_relname = filename[:-3] filepath = os.path.join(component_path, filename) try: with open(filepath, "r", encoding="utf-8") as f: source = f.read() tree = ast.parse(source, filename=filename) entry_class_node = None first_class_def = None # Collect all class definitions and their _aliases file_class_aliases: dict[str, list[str]] = {} for node in ast.walk(tree): if isinstance(node, ast.ClassDef): if first_class_def is None: first_class_def = node # Look for _aliases in the class body for class_body_node in node.body: if isinstance(class_body_node, ast.Assign): for target in class_body_node.targets: if ( isinstance(target, ast.Name) and target.id == "_aliases" ): aliases = _parse_aliases_from_ast( class_body_node.value ) if aliases: file_class_aliases[node.name] = aliases if isinstance(node, ast.Assign): for target in node.targets: if ( isinstance(target, ast.Name) and target.id == "EntryClass" ): entry_class_node = node break if entry_class_node and first_class_def: model_cls_name_list = [] value_node = entry_class_node.value # EntryClass = ClassName if isinstance(value_node, ast.Name): model_cls_name_list.append(value_node.id) # EntryClass = ["...", ClassName, ...] elif isinstance(value_node, (ast.List, ast.Tuple)): for elt in value_node.elts: if isinstance(elt, ast.Constant): model_cls_name_list.append(elt.value) elif isinstance(elt, ast.Name): model_cls_name_list.append(elt.id) if model_cls_name_list: for model_cls_str in model_cls_name_list: if model_cls_str in discovered_models: logger.warning( f"Duplicate architecture found: {model_cls_str}. It will be overwritten." ) model_arch = model_cls_str discovered_models[model_arch] = ( component, mod_relname, model_cls_str, ) # Collect aliases for this class if model_cls_str in file_class_aliases: class_aliases[model_cls_str] = file_class_aliases[ model_cls_str ] except Exception as e: logger.warning(f"Could not parse {filepath} to find models: {e}") # Build alias -> canonical class name mapping for class_name, aliases in class_aliases.items(): for alias in aliases: if alias in _ALIAS_TO_MODEL: logger.warning( f"Alias '{alias}' already registered for '{_ALIAS_TO_MODEL[alias]}', " f"will be overwritten by '{class_name}'" ) _ALIAS_TO_MODEL[alias] = class_name return discovered_models _SGLANG_DIFFUSION_MODELS = _discover_and_register_models() _SUBPROCESS_COMMAND = [ sys.executable, "-m", "sglang.multimodal_gen.runtime.models.dits.registry", ] _T = TypeVar("_T") @dataclass(frozen=True) class _ModelInfo: architecture: str @staticmethod def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": return _ModelInfo( architecture=model.__name__, ) class _BaseRegisteredModel(ABC): @abstractmethod def inspect_model_cls(self) -> _ModelInfo: raise NotImplementedError @abstractmethod def load_model_cls(self) -> type[nn.Module]: raise NotImplementedError @dataclass(frozen=True) class _RegisteredModel(_BaseRegisteredModel): """ Represents a model that has already been imported in the main process. """ interfaces: _ModelInfo model_cls: type[nn.Module] @staticmethod def from_model_cls(model_cls: type[nn.Module]): return _RegisteredModel( interfaces=_ModelInfo.from_model_cls(model_cls), model_cls=model_cls, ) def inspect_model_cls(self) -> _ModelInfo: return self.interfaces def load_model_cls(self) -> type[nn.Module]: return self.model_cls def _run_in_subprocess(fn: Callable[[], _T]) -> _T: # NOTE: We use a temporary directory instead of a temporary file to avoid # issues like https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file with tempfile.TemporaryDirectory() as tempdir: output_filepath = os.path.join(tempdir, "registry_output.tmp") # `cloudpickle` allows pickling lambda functions directly input_bytes = cloudpickle.dumps((fn, output_filepath)) # cannot use `sys.executable __file__` here because the script # contains relative imports returned = subprocess.run( _SUBPROCESS_COMMAND, input=input_bytes, capture_output=True ) # check if the subprocess is successful try: returned.check_returncode() except Exception as e: # wrap raised exception to provide more information raise RuntimeError( f"Error raised in subprocess:\n" f"{returned.stderr.decode()}" ) from e with open(output_filepath, "rb") as f: return cast(_T, pickle.load(f)) @dataclass(frozen=True) class _LazyRegisteredModel(_BaseRegisteredModel): """ Represents a model that has not been imported in the main process. """ module_name: str component_name: str class_name: str # Performed in another process to avoid initializing CUDA def inspect_model_cls(self) -> _ModelInfo: return _run_in_subprocess( lambda: _ModelInfo.from_model_cls(self.load_model_cls()) ) def load_model_cls(self) -> type[nn.Module]: mod = importlib.import_module(self.module_name) return cast(type[nn.Module], getattr(mod, self.class_name)) @lru_cache(maxsize=128) def _try_load_model_cls( model_arch: str, model: _BaseRegisteredModel, ) -> type[nn.Module] | None: from sglang.multimodal_gen.runtime.platforms import current_platform current_platform.verify_model_arch(model_arch) try: return model.load_model_cls() except Exception: logger.exception("Ignore import error when loading '%s'", model_arch) return None @lru_cache(maxsize=128) def _try_inspect_model_cls( model_arch: str, model: _BaseRegisteredModel, ) -> _ModelInfo | None: try: return model.inspect_model_cls() except Exception: logger.exception("Error in inspecting model architecture '%s'", model_arch) return None @dataclass class _ModelRegistry: # Keyed by model_arch registered_models: dict[str, _BaseRegisteredModel] = field(default_factory=dict) def get_supported_archs(self) -> Set[str]: return self.registered_models.keys() def resolve_by_alias(self, alias: str) -> type[nn.Module] | None: """Resolve a model class by its alias (external module path).""" if alias in _ALIAS_TO_MODEL: canonical_name = _ALIAS_TO_MODEL[alias] return self._try_load_model_cls(canonical_name) return None def register_model( self, model_arch: str, model_cls: type[nn.Module] | str, ) -> None: """ Register an external model to be used in vLLM. :code:`model_cls` can be either: - A :class:`torch.nn.Module` class directly referencing the model. - A string in the format :code:`:` which can be used to lazily import the model. This is useful to avoid initializing CUDA when importing the model and thus the related error :code:`RuntimeError: Cannot re-initialize CUDA in forked subprocess`. """ if model_arch in self.registered_models: logger.warning( "Model architecture %s is already registered, and will be " "overwritten by the new model class %s.", model_arch, model_cls, ) if isinstance(model_cls, str): split_str = model_cls.split(":") if len(split_str) != 2: msg = "Expected a string in the format `:`" raise ValueError(msg) model = _LazyRegisteredModel(*split_str) else: model = _RegisteredModel.from_model_cls(model_cls) self.registered_models[model_arch] = model def _raise_for_unsupported(self, architectures: list[str]) -> NoReturn: all_supported_archs = self.get_supported_archs() if any(arch in all_supported_archs for arch in architectures): raise ValueError( f"Model architectures {architectures} failed " "to be inspected. Please check the logs for more details." ) raise ValueError( f"Model architectures {architectures} are not supported for now. " f"Supported architectures: {all_supported_archs}" ) def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None: if model_arch not in self.registered_models: return None return _try_load_model_cls(model_arch, self.registered_models[model_arch]) def _try_inspect_model_cls(self, model_arch: str) -> _ModelInfo | None: if model_arch not in self.registered_models: return None return _try_inspect_model_cls(model_arch, self.registered_models[model_arch]) def _normalize_archs( self, architectures: str | list[str], ) -> list[str]: if isinstance(architectures, str): architectures = [architectures] if not architectures: logger.warning("No model architectures are specified") normalized_arch = [] for arch in architectures: if arch not in self.registered_models: registered_models = list(self.registered_models.keys()) raise Exception( f"Unsupported model architecture: {arch}. Registered architectures: {registered_models}" ) normalized_arch.append(arch) return normalized_arch def inspect_model_cls( self, architectures: str | list[str], ) -> tuple[_ModelInfo, str]: architectures = self._normalize_archs(architectures) for arch in architectures: model_info = self._try_inspect_model_cls(arch) if model_info is not None: return (model_info, arch) return self._raise_for_unsupported(architectures) def resolve_model_cls( self, architectures: str | list[str], ) -> tuple[type[nn.Module], str]: architectures = self._normalize_archs(architectures) for arch in architectures: model_cls = self._try_load_model_cls(arch) if model_cls is not None: return (model_cls, arch) return self._raise_for_unsupported(architectures) ModelRegistry = _ModelRegistry( { model_arch: _LazyRegisteredModel( module_name=f"sglang.multimodal_gen.runtime.models.{component_name}.{mod_relname}", component_name=component_name, class_name=cls_name, ) for model_arch, ( component_name, mod_relname, cls_name, ) in _SGLANG_DIFFUSION_MODELS.items() } ) ================================================ FILE: python/sglang/multimodal_gen/runtime/models/schedulers/base.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from abc import ABC, abstractmethod import torch class BaseScheduler(ABC): timesteps: torch.Tensor order: int num_train_timesteps: int def __init__(self, *args, **kwargs) -> None: # Check if subclass has defined all required properties required_attributes = ["timesteps", "order", "num_train_timesteps"] for attr in required_attributes: if not hasattr(self, attr): raise AttributeError( f"Subclasses of BaseScheduler must define '{attr}' property" ) @abstractmethod def set_shift(self, shift: float) -> None: pass @abstractmethod def set_timesteps(self, *args, **kwargs) -> None: pass @abstractmethod def scale_model_input( self, sample: torch.Tensor, timestep: int | None = None ) -> torch.Tensor: pass ================================================ FILE: python/sglang/multimodal_gen/runtime/models/schedulers/flow_match_pair.py ================================================ # Copied and adapted from: https://github.com/OpenMOSS/MOVA/tree/main/mova/diffusion/schedulers/flow_match.py and flow_match_pair.py # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import math import torch from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler class FlowMatchScheduler(BaseScheduler): def __init__( self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False, exponential_shift=False, exponential_shift_mu=None, shift_terminal=None, ): self.order = 1 self.num_train_timesteps = num_train_timesteps self.shift = shift self.sigma_max = sigma_max self.sigma_min = sigma_min self.inverse_timesteps = inverse_timesteps self.extra_one_step = extra_one_step self.reverse_sigmas = reverse_sigmas self.exponential_shift = exponential_shift self.exponential_shift_mu = exponential_shift_mu self.shift_terminal = shift_terminal self.train_timesteps = None self.train_sigmas = None self.set_timesteps(num_train_timesteps) self.set_timesteps(num_inference_steps) BaseScheduler.__init__(self) def set_shift(self, shift: float) -> None: self.shift = shift def set_timesteps( self, num_inference_steps=100, denoising_strength=1.0, training=False, shift=None, dynamic_shift_len=None, ): if shift is not None: self.shift = shift sigma_start = ( self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength ) if self.extra_one_step: self.sigmas = torch.linspace( sigma_start, self.sigma_min, num_inference_steps + 1 )[:-1] else: self.sigmas = torch.linspace( sigma_start, self.sigma_min, num_inference_steps ) if self.inverse_timesteps: self.sigmas = torch.flip(self.sigmas, dims=[0]) if self.exponential_shift: mu = ( self.calculate_shift(dynamic_shift_len) if dynamic_shift_len is not None else self.exponential_shift_mu ) self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1)) else: self.sigmas = ( self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) ) if self.shift_terminal is not None: one_minus_z = 1 - self.sigmas scale_factor = one_minus_z[-1] / (1 - self.shift_terminal) self.sigmas = 1 - (one_minus_z / scale_factor) if self.reverse_sigmas: self.sigmas = 1 - self.sigmas self.timesteps = self.sigmas * self.num_train_timesteps # Initialize train_timesteps on first set. if self.train_timesteps is None: self.train_timesteps = self.timesteps self.train_sigmas = self.sigmas if training: x = self.timesteps y = torch.exp( -2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2 ) y_shifted = y - y.min() bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum()) self.linear_timesteps_weights = bsmntw_weighing self.training = True else: self.training = False def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None): return sample def step(self, model_output, timestep, sample, to_final=False, **kwargs): if isinstance(timestep, torch.Tensor): timestep = timestep.cpu() timestep_id = torch.argmin((self.timesteps - timestep).abs()) sigma = self.sigmas[timestep_id] if to_final or timestep_id + 1 >= len(self.timesteps): sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0 else: sigma_ = self.sigmas[timestep_id + 1] prev_sample = sample + model_output * (sigma_ - sigma) return prev_sample def return_to_timestep(self, timestep, sample, sample_stablized): if isinstance(timestep, torch.Tensor): timestep = timestep.cpu() timestep_id = torch.argmin((self.timesteps - timestep).abs()) sigma = self.sigmas[timestep_id] model_output = (sample - sample_stablized) / sigma return model_output def add_noise(self, original_samples, noise, timestep): if isinstance(timestep, torch.Tensor): timestep = timestep.cpu() timestep_id = torch.argmin((self.timesteps - timestep).abs()) sigma = self.sigmas[timestep_id] sample = (1 - sigma) * original_samples + sigma * noise return sample def training_target(self, sample, noise, timestep): target = noise - sample return target def training_weight(self, timestep): timestep_id = torch.argmin( (self.timesteps - timestep.to(self.timesteps.device)).abs() ) weights = self.linear_timesteps_weights[timestep_id] return weights def calculate_shift( self, image_seq_len, base_seq_len: int = 256, max_seq_len: int = 8192, base_shift: float = 0.5, max_shift: float = 0.9, ): m = (max_shift - base_shift) / (max_seq_len - base_seq_len) b = base_shift - m * base_seq_len mu = image_seq_len * m + b return mu class FlowMatchPairScheduler(FlowMatchScheduler): """Pairing scheduler built on FlowMatchScheduler. Provides a convenient pairing interface for timesteps or sigmas. Attributes: pair_timesteps: Cached timestep pairs of shape [num_timesteps, 2]. pair_sigmas: Cached sigma pairs of shape [num_timesteps, 2]. """ def __init__( self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False, exponential_shift=False, exponential_shift_mu=None, shift_terminal=None, ): self._pair_postprocess_fn = None self._pair_postprocess_requires_source = False self.pair_timesteps: torch.Tensor | None = None self.pair_sigmas: torch.Tensor | None = None self.timesteps: torch.Tensor | None = None self.sigmas: torch.Tensor | None = None super().__init__( num_inference_steps=num_inference_steps, num_train_timesteps=num_train_timesteps, shift=shift, sigma_max=sigma_max, sigma_min=sigma_min, inverse_timesteps=inverse_timesteps, extra_one_step=extra_one_step, reverse_sigmas=reverse_sigmas, exponential_shift=exponential_shift, exponential_shift_mu=exponential_shift_mu, shift_terminal=shift_terminal, ) def set_pair_postprocess(self, fn): """Set a postprocess function to customize pairs after construction. Args: fn: Callable with signature fn(pairs: torch.Tensor) -> torch.Tensor. The returned tensor must have the same shape as input pairs. Raises: TypeError: If fn is not callable or None. RuntimeError: If scheduler is not initialized. """ if fn is not None and not callable(fn): raise TypeError("pair_postprocess must be callable or None") self._pair_postprocess_fn = fn self._pair_postprocess_requires_source = ( False if fn is None else bool(getattr(fn, "_requires_source", False)) ) if self.timesteps is None or self.sigmas is None: raise RuntimeError("Scheduler not initialized; call set_timesteps() first") self._refresh_pair_cache() def set_pair_postprocess_by_name(self, name: str | None, **kwargs): """Configure a postprocess function by name. Supported names: - None/"none"/"off"/"false"/"no": disable - "quadratic_perp_bulge_swap": x2=x+d, y2=x-d, where d=4*amp*s*(1-s), s=t/T - "v2a_sequential": assume pairs are (t,t); sample half sequence from column 0 with stride 2, then let column 0 follow this sequence first, followed by column 1 - "a2v_sequential": same as above, but column 1 first then column 0 - "dual_sigma_shift": use only timestep count; rebuild two columns independently using FlowMatchScheduler sigma transform logic; configurable visual_shift/audio_shift Args: name: Postprocess name or None to disable. **kwargs: Extra parameters for the named postprocess. For example: - amp: Float amplitude, default 150.0. Raises: ValueError: If name is unknown. """ if name is None or str(name).lower() in ("none", "off", "false", "no"): self.set_pair_postprocess(None) return if name == "quadratic_perp_bulge_swap": amp = float(kwargs.get("amp", 150.0)) def _quadratic_perp_bulge_swap(pairs: torch.Tensor): if ( not isinstance(pairs, torch.Tensor) or pairs.ndim != 2 or pairs.shape[1] != 2 ): raise ValueError("pairs must be a torch.Tensor of shape [N, 2]") x = pairs[:, 0] T = float(self.num_train_timesteps) s = x / T d = 4.0 * amp * s * (1.0 - s) x2 = x + d y2 = x - d return torch.stack([x2, y2], dim=1) self.set_pair_postprocess(_quadratic_perp_bulge_swap) return if name == "v2a_sequential": def _v2a(pairs: torch.Tensor): if ( not isinstance(pairs, torch.Tensor) or pairs.ndim != 2 or pairs.shape[1] != 2 ): raise ValueError("pairs must be a torch.Tensor of shape [N, 2]") N = pairs.shape[0] base = pairs[:, 0] seq_half = base[::2] m = int(seq_half.shape[0]) col0 = torch.cat([seq_half, seq_half[-1:].repeat(m)], dim=0)[:N] col1 = torch.cat([seq_half[0:1].repeat(m), seq_half], dim=0)[:N] return torch.stack( [ col0.to(dtype=pairs.dtype, device=pairs.device), col1.to(dtype=pairs.dtype, device=pairs.device), ], dim=1, ) self.set_pair_postprocess(_v2a) return if name == "a2v_sequential": def _a2v(pairs: torch.Tensor): if ( not isinstance(pairs, torch.Tensor) or pairs.ndim != 2 or pairs.shape[1] != 2 ): raise ValueError("pairs must be a torch.Tensor of shape [N, 2]") N = pairs.shape[0] base = pairs[:, 0] seq_half = base[::2] m = int(seq_half.shape[0]) col0 = torch.cat([seq_half[0:1].repeat(m), seq_half], dim=0)[:N] col1 = torch.cat([seq_half, seq_half[-1:].repeat(m)], dim=0)[:N] return torch.stack( [ col0.to(dtype=pairs.dtype, device=pairs.device), col1.to(dtype=pairs.dtype, device=pairs.device), ], dim=1, ) self.set_pair_postprocess(_a2v) return if name == "v2a": def _v2a_classic(pairs: torch.Tensor): if ( not isinstance(pairs, torch.Tensor) or pairs.ndim != 2 or pairs.shape[1] != 2 ): raise ValueError("pairs must be a torch.Tensor of shape [N, 2]") zeros = torch.zeros_like(pairs[:, 0]) return torch.stack([zeros, pairs[:, 1]], dim=1) self.set_pair_postprocess(_v2a_classic) return if name == "a2v": def _a2v_classic(pairs: torch.Tensor): if ( not isinstance(pairs, torch.Tensor) or pairs.ndim != 2 or pairs.shape[1] != 2 ): raise ValueError("pairs must be a torch.Tensor of shape [N, 2]") zeros = torch.zeros_like(pairs[:, 1]) return torch.stack([pairs[:, 0], zeros], dim=1) self.set_pair_postprocess(_a2v_classic) return if name == "dual_sigma_shift": visual_shift = float(kwargs.get("visual_shift", self.shift)) audio_shift = float(kwargs.get("audio_shift", self.shift)) visual_denoising_strength = float( kwargs.get("visual_denoising_strength", 1.0) ) audio_denoising_strength = float( kwargs.get("audio_denoising_strength", 1.0) ) visual_mu = kwargs.get( "visual_exponential_shift_mu", self.exponential_shift_mu ) audio_mu = kwargs.get( "audio_exponential_shift_mu", self.exponential_shift_mu ) def _dual_sigma_shift(pairs: torch.Tensor, *, source: str): if not isinstance(pairs, torch.Tensor): raise TypeError("pairs must be a torch.Tensor") if pairs.ndim != 2 or pairs.shape[1] != 2: raise ValueError("pairs must be a torch.Tensor of shape [N, 2]") if pairs.shape[0] == 0: raise ValueError("pairs length must be greater than 0") if source not in ("timesteps", "sigmas"): raise ValueError("source must be 'timesteps' or 'sigmas'") num_steps = pairs.shape[0] device = pairs.device dtype = pairs.dtype def _build_column( shift_value: float, denoising_strength: float, mu_override ): if shift_value <= 0: raise ValueError("shift must be positive") if denoising_strength <= 0: raise ValueError("denoising_strength must be positive") sigma_start = ( self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength ) if self.extra_one_step: base = torch.linspace( sigma_start, self.sigma_min, num_steps + 1, device=device, dtype=dtype, )[:-1] else: base = torch.linspace( sigma_start, self.sigma_min, num_steps, device=device, dtype=dtype, ) if self.inverse_timesteps: base = torch.flip(base, dims=[0]) if self.exponential_shift: mu_value = mu_override if mu_value is None: raise RuntimeError( "exponential_shift enabled but exponential_shift_mu is missing" ) exp_mu = math.exp(float(mu_value)) base = exp_mu / (exp_mu + (1 / base - 1)) else: base = shift_value * base / (1 + (shift_value - 1) * base) if self.shift_terminal is not None: one_minus_z = 1 - base scale_factor = one_minus_z[-1] / (1 - self.shift_terminal) base = 1 - (one_minus_z / scale_factor) if self.reverse_sigmas: base = 1 - base if source == "timesteps": return base * self.num_train_timesteps return base col0 = _build_column(visual_shift, visual_denoising_strength, visual_mu) col1 = _build_column(audio_shift, audio_denoising_strength, audio_mu) return torch.stack([col0, col1], dim=1) _dual_sigma_shift._requires_source = True self.set_pair_postprocess(_dual_sigma_shift) return raise ValueError(f"Unknown pair_postprocess name: {name}") def _make_pairs_from_vector(self, vec: torch.Tensor) -> torch.Tensor: if vec.ndim != 1: raise ValueError("vec must be 1D") return torch.stack([vec, vec], dim=1) def get_pairs(self, source: str = "timesteps") -> torch.Tensor: if source == "timesteps": if self.pair_timesteps is None: self._refresh_pair_cache() return self.pair_timesteps if source == "sigmas": if self.pair_sigmas is None: self._refresh_pair_cache() return self.pair_sigmas raise ValueError("source must be 'timesteps' or 'sigmas'") def timestep_to_sigma(self, timestep: torch.Tensor | float) -> torch.Tensor: """Return sigma for a scalar timestep via nearest neighbor lookup. Args: timestep: Scalar timestep value. Returns: Sigma corresponding to the nearest timestep. """ t_value = float(timestep) t_cpu = torch.tensor(t_value) idx = torch.argmin((self.train_timesteps - t_cpu).abs()) return self.train_sigmas[idx] def step_from_to( self, model_output: torch.Tensor, timestep_from: torch.Tensor, timestep_to: torch.Tensor | None, sample: torch.Tensor, ) -> torch.Tensor: """Advance one step using an explicit (from, to) timestep pair. The update rule is: x_{to} = x_{from} + model_output * (sigma(to) - sigma(from)) Args: model_output: Predicted model output. timestep_from: Source timestep. timestep_to: Target timestep or None for terminal. sample: Current sample at timestep_from. Returns: Updated sample at timestep_to. """ sigma_from = self.timestep_to_sigma(timestep_from) if timestep_to is None: sigma_to = torch.tensor( 1.0 if (self.inverse_timesteps or self.reverse_sigmas) else 0.0, device=sigma_from.device, dtype=sigma_from.dtype, ) else: sigma_to = self.timestep_to_sigma(timestep_to) prev_sample = sample + model_output * (sigma_to - sigma_from) return prev_sample def _refresh_pair_cache(self) -> None: if self.timesteps is None or self.sigmas is None: raise RuntimeError("Scheduler not initialized; call set_timesteps() first") def _apply_postprocess(pairs: torch.Tensor, source: str) -> torch.Tensor: if self._pair_postprocess_fn is None: return pairs if self._pair_postprocess_requires_source: modified = self._pair_postprocess_fn(pairs, source=source) else: modified = self._pair_postprocess_fn(pairs) if not isinstance(modified, torch.Tensor): raise TypeError("pair_postprocess must return a torch.Tensor") if modified.shape != pairs.shape: raise ValueError("pair_postprocess must return the same shape as input") return modified base_pairs_timesteps = self._make_pairs_from_vector(self.timesteps) base_pairs_sigmas = self._make_pairs_from_vector(self.sigmas) self.pair_timesteps = _apply_postprocess(base_pairs_timesteps, "timesteps") self.pair_sigmas = _apply_postprocess(base_pairs_sigmas, "sigmas") EntryClass = FlowMatchPairScheduler ================================================ FILE: python/sglang/multimodal_gen/runtime/models/schedulers/hunyuan3d_scheduler.py ================================================ # Copied and adapted from: https://github.com/Tencent-Hunyuan/Hunyuan3D-2 from __future__ import annotations import math from dataclasses import dataclass from typing import List, Optional, Tuple, Union import numpy as np import torch from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.schedulers.scheduling_utils import SchedulerMixin from diffusers.utils import BaseOutput @dataclass class Hunyuan3DFlowMatchSchedulerOutput(BaseOutput): """Output class for the scheduler's step function.""" prev_sample: torch.FloatTensor class Hunyuan3DFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): """Euler discrete scheduler for flow matching.""" # External module path aliases for compatibility with Hunyuan3D configs _aliases = [ "hy3dgen.shapegen.schedulers.FlowMatchEulerDiscreteScheduler", "hy3dshape.schedulers.FlowMatchEulerDiscreteScheduler", ] _compatibles = [] order = 1 @register_to_config def __init__( self, num_train_timesteps: int = 1000, shift: float = 1.0, use_dynamic_shifting: bool = False, ): timesteps = np.linspace( 1, num_train_timesteps, num_train_timesteps, dtype=np.float32 ).copy() timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) sigmas = timesteps / num_train_timesteps if not use_dynamic_shifting: sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) self.timesteps = sigmas * num_train_timesteps self._step_index = None self._begin_index = None self.sigmas = sigmas.to("cpu") self.sigma_min = self.sigmas[-1].item() self.sigma_max = self.sigmas[0].item() @property def step_index(self) -> Optional[int]: """The index counter for current timestep.""" return self._step_index @property def begin_index(self) -> Optional[int]: """The index for the first timestep.""" return self._begin_index def set_begin_index(self, begin_index: int = 0): """Set the begin index for the scheduler. Args: begin_index: The begin index for the scheduler. """ self._begin_index = begin_index def scale_model_input( self, sample: torch.FloatTensor, timestep: Optional[Union[float, torch.FloatTensor]] = None, ) -> torch.FloatTensor: """Identity operation for flow matching (no input scaling needed).""" return sample def scale_noise( self, sample: torch.FloatTensor, timestep: Union[float, torch.FloatTensor], noise: Optional[torch.FloatTensor] = None, ) -> torch.FloatTensor: """Forward process in flow-matching (add noise to sample).""" sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype) if sample.device.type == "mps" and torch.is_floating_point(timestep): schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) timestep = timestep.to(sample.device, dtype=torch.float32) else: schedule_timesteps = self.timesteps.to(sample.device) timestep = timestep.to(sample.device) if self.begin_index is None: step_indices = [ self.index_for_timestep(t, schedule_timesteps) for t in timestep ] elif self.step_index is not None: step_indices = [self.step_index] * timestep.shape[0] else: step_indices = [self.begin_index] * timestep.shape[0] sigma = sigmas[step_indices].flatten() while len(sigma.shape) < len(sample.shape): sigma = sigma.unsqueeze(-1) sample = sigma * noise + (1.0 - sigma) * sample return sample def _sigma_to_t(self, sigma: float) -> float: """Convert sigma to timestep.""" return sigma * self.config.num_train_timesteps def time_shift(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor: """Apply time shift transformation.""" return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) def set_timesteps( self, num_inference_steps: int = None, device: Union[str, torch.device] = None, sigmas: Optional[List[float]] = None, mu: Optional[float] = None, ): """Set the discrete timesteps for the diffusion chain.""" if self.config.use_dynamic_shifting and mu is None: raise ValueError( "Must pass a value for `mu` when `use_dynamic_shifting` is True" ) if sigmas is None: self.num_inference_steps = num_inference_steps timesteps = np.linspace( self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps, ) sigmas = timesteps / self.config.num_train_timesteps if self.config.use_dynamic_shifting: sigmas = self.time_shift(mu, 1.0, sigmas) else: sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas) sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) timesteps = sigmas * self.config.num_train_timesteps self.timesteps = timesteps.to(device=device) self.sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)]) self._step_index = None self._begin_index = None def index_for_timestep( self, timestep: float, schedule_timesteps: Optional[torch.Tensor] = None ) -> int: """Find the index for a given timestep.""" if schedule_timesteps is None: schedule_timesteps = self.timesteps indices = (schedule_timesteps == timestep).nonzero() pos = 1 if len(indices) > 1 else 0 return indices[pos].item() def _init_step_index(self, timestep: Union[float, torch.Tensor]): """Initialize step index from timestep.""" if self.begin_index is None: if isinstance(timestep, torch.Tensor): timestep = timestep.to(self.timesteps.device) self._step_index = self.index_for_timestep(timestep) else: self._step_index = self._begin_index def step( self, model_output: torch.FloatTensor, timestep: Union[float, torch.FloatTensor], sample: torch.FloatTensor, s_churn: float = 0.0, s_tmin: float = 0.0, s_tmax: float = float("inf"), s_noise: float = 1.0, generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[Hunyuan3DFlowMatchSchedulerOutput, Tuple]: """Predict the sample from the previous timestep.""" if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): raise ValueError( "Passing integer indices as timesteps is not supported. " "Pass one of `scheduler.timesteps` as a timestep." ) if self.step_index is None: self._init_step_index(timestep) # Upcast to avoid precision issues sample = sample.to(torch.float32) sigma = self.sigmas[self.step_index] sigma_next = self.sigmas[self.step_index + 1] prev_sample = sample + (sigma_next - sigma) * model_output prev_sample = prev_sample.to(model_output.dtype) self._step_index += 1 if not return_dict: return (prev_sample,) return Hunyuan3DFlowMatchSchedulerOutput(prev_sample=prev_sample) def __len__(self) -> int: return self.config.num_train_timesteps @dataclass class Hunyuan3DConsistencyFlowMatchSchedulerOutput(BaseOutput): """Output for consistency flow matching scheduler.""" prev_sample: torch.FloatTensor pred_original_sample: torch.FloatTensor class Hunyuan3DConsistencyFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): """Consistency Flow Matching Euler Discrete Scheduler.""" # External module path aliases for compatibility with Hunyuan3D configs _aliases = [ "hy3dshape.schedulers.Hunyuan3DConsistencyFlowMatchEulerDiscreteScheduler", ] _compatibles = [] order = 1 @register_to_config def __init__( self, num_train_timesteps: int = 1000, pcm_timesteps: int = 50, ): sigmas = np.linspace(0, 1, num_train_timesteps) step_ratio = num_train_timesteps // pcm_timesteps euler_timesteps = (np.arange(1, pcm_timesteps) * step_ratio).round().astype( np.int64 ) - 1 euler_timesteps = np.asarray([0] + euler_timesteps.tolist()) self.euler_timesteps = euler_timesteps self.sigmas = sigmas[self.euler_timesteps] self.sigmas = torch.from_numpy(self.sigmas.copy()).to(dtype=torch.float32) self.timesteps = self.sigmas * num_train_timesteps self._step_index = None self._begin_index = None self.sigmas = self.sigmas.to("cpu") @property def step_index(self) -> Optional[int]: return self._step_index @property def begin_index(self) -> Optional[int]: return self._begin_index def set_begin_index(self, begin_index: int = 0): self._begin_index = begin_index def scale_model_input( self, sample: torch.FloatTensor, timestep: Optional[Union[float, torch.FloatTensor]] = None, ) -> torch.FloatTensor: """Identity operation for flow matching (no input scaling needed).""" return sample def _sigma_to_t(self, sigma: float) -> float: return sigma * self.config.num_train_timesteps def set_timesteps( self, num_inference_steps: int = None, device: Union[str, torch.device] = None, sigmas: Optional[List[float]] = None, ): """Set timesteps for inference.""" self.num_inference_steps = ( num_inference_steps if num_inference_steps is not None else len(sigmas) ) inference_indices = np.linspace( 0, self.config.pcm_timesteps, num=self.num_inference_steps, endpoint=False ) inference_indices = np.floor(inference_indices).astype(np.int64) inference_indices = torch.from_numpy(inference_indices).long() self.sigmas_ = self.sigmas[inference_indices] timesteps = self.sigmas_ * self.config.num_train_timesteps self.timesteps = timesteps.to(device=device) self.sigmas_ = torch.cat( [self.sigmas_, torch.ones(1, device=self.sigmas_.device)] ) self._step_index = None self._begin_index = None def index_for_timestep( self, timestep: float, schedule_timesteps: Optional[torch.Tensor] = None ) -> int: if schedule_timesteps is None: schedule_timesteps = self.timesteps indices = (schedule_timesteps == timestep).nonzero() pos = 1 if len(indices) > 1 else 0 return indices[pos].item() def _init_step_index(self, timestep: Union[float, torch.Tensor]): if self.begin_index is None: if isinstance(timestep, torch.Tensor): timestep = timestep.to(self.timesteps.device) self._step_index = self.index_for_timestep(timestep) else: self._step_index = self._begin_index def step( self, model_output: torch.FloatTensor, timestep: Union[float, torch.FloatTensor], sample: torch.FloatTensor, generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[Hunyuan3DConsistencyFlowMatchSchedulerOutput, Tuple]: """Perform one step of the consistency flow matching scheduler.""" if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): raise ValueError("Passing integer indices as timesteps is not supported.") if self.step_index is None: self._init_step_index(timestep) sample = sample.to(torch.float32) sigma = self.sigmas_[self.step_index] sigma_next = self.sigmas_[self.step_index + 1] prev_sample = sample + (sigma_next - sigma) * model_output prev_sample = prev_sample.to(model_output.dtype) pred_original_sample = sample + (1.0 - sigma) * model_output pred_original_sample = pred_original_sample.to(model_output.dtype) self._step_index += 1 if not return_dict: return (prev_sample,) return Hunyuan3DConsistencyFlowMatchSchedulerOutput( prev_sample=prev_sample, pred_original_sample=pred_original_sample ) def __len__(self) -> int: return self.config.num_train_timesteps # Entry class for model registry EntryClass = [ Hunyuan3DFlowMatchEulerDiscreteScheduler, Hunyuan3DConsistencyFlowMatchEulerDiscreteScheduler, ] ================================================ FILE: python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_comfyui_passthrough.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Pass-through scheduler for ComfyUI integration. This scheduler does not modify latents - it simply returns the input sample unchanged. The actual denoising logic is handled by ComfyUI. """ import torch from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.schedulers.scheduling_utils import SchedulerMixin from diffusers.utils import BaseOutput from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class ComfyUIPassThroughSchedulerOutput(BaseOutput): """ Output class for the scheduler's `step` function output. Args: prev_sample (`torch.FloatTensor`): The input sample unchanged (pass-through). """ prev_sample: torch.FloatTensor class ComfyUIPassThroughScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): """ Pass-through scheduler for ComfyUI integration. This scheduler does not modify latents. It is used when the denoising logic is handled externally by ComfyUI. The scheduler simply returns the input sample unchanged, allowing ComfyUI to manage the denoising process. Usage: - num_inference_steps is always 1 (each step is handled separately) - timesteps are provided externally by ComfyUI - step() returns the input sample unchanged """ config_name = "scheduler_config.json" order = 1 @register_to_config def __init__( self, num_train_timesteps=1000, *args, **kwargs, ): self.num_train_timesteps = num_train_timesteps # Initialize timesteps as empty - will be set externally self.timesteps = torch.tensor([], dtype=torch.long) self.shift = 0.0 self._step_index = 0 # Track current step index self._begin_index: int | None = None # For compatibility with DenoisingStage def set_timesteps( self, num_inference_steps=1, # Always 1 for ComfyUI timesteps=None, # Can be provided externally device=None, **kwargs, ): """ Set timesteps. For ComfyUI, timesteps are provided externally. Args: num_inference_steps: Ignored (always 1 for ComfyUI) timesteps: External timesteps provided by ComfyUI device: Device to place timesteps on """ if timesteps is not None: # Use externally provided timesteps if isinstance(timesteps, torch.Tensor): self.timesteps = timesteps else: self.timesteps = torch.tensor(timesteps, dtype=torch.long) if device is not None: self.timesteps = self.timesteps.to(device) else: # Create a single timestep if none provided if device is None: device = torch.device("cpu") self.timesteps = torch.tensor([0], dtype=torch.long, device=device) def step( self, model_output: torch.FloatTensor, timestep: torch.FloatTensor | int, sample: torch.FloatTensor, return_dict: bool = False, **kwargs, ) -> tuple | ComfyUIPassThroughSchedulerOutput: """ Pass-through step: returns the input sample unchanged. This scheduler does not modify latents. The actual denoising is handled by ComfyUI, so we simply return the input sample as-is. Args: model_output: Predicted noise (ignored, but kept for API compatibility) timestep: Current timestep (ignored, but kept for API compatibility) sample: Input latents (returned unchanged) return_dict: Whether to return a dict or tuple Returns: The input sample unchanged (prev_sample = sample) """ # Increment step index for tracking self._step_index += 1 # Simply return the input sample unchanged prev_sample = sample if not return_dict: return (prev_sample,) return ComfyUIPassThroughSchedulerOutput(prev_sample=prev_sample) def scale_model_input( self, sample: torch.Tensor, timestep: int | None = None ) -> torch.Tensor: """ Scale model input. For pass-through scheduler, returns input unchanged. Args: sample: Input sample timestep: Timestep (ignored) Returns: Input sample unchanged """ return sample def set_shift(self, shift: float) -> None: """ Set shift parameter (no-op for pass-through scheduler). Args: shift: Shift value (ignored) """ self.shift = shift def set_begin_index(self, begin_index: int = 0) -> None: """ Sets the begin index for the scheduler. This function should be run from pipeline before the inference. Args: begin_index: The begin index for the scheduler. """ self._begin_index = begin_index @property def begin_index(self) -> int | None: """ The index for the first timestep. """ return self._begin_index @property def step_index(self) -> int: """ The index counter for current timestep. """ return self._step_index def add_noise( self, original_samples: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor, ) -> torch.Tensor: """ Add noise to samples. For pass-through scheduler, returns original samples. Args: original_samples: Original clean samples noise: Noise to add (ignored) timestep: Timestep (ignored) Returns: Original samples unchanged """ return original_samples EntryClass = ComfyUIPassThroughScheduler ================================================ FILE: python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_dpm_solver_multistep.py ================================================ # SPDX-License-Identifier: Apache-2.0 # # DPM-Solver++ multistep scheduler wrapper for SANA. # # SANA uses DPM-Solver++ (Lu et al., 2022) as its noise scheduler, which # is a high-order ODE solver that converges in fewer steps than DDIM. # With solver_order=2 and 20 steps, SANA achieves high-quality results. # # This wrapper delegates all numerical work to diffusers' implementation # and only adapts the interface for sglang's denoising stage. import torch from diffusers import ( DPMSolverMultistepScheduler as DiffusersDPMSolverMultistepScheduler, ) from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.schedulers.scheduling_utils import SchedulerMixin from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """DPM-Solver++ multistep scheduler wrapper for sglang's BaseScheduler interface.""" order = 1 num_train_timesteps = 1000 @register_to_config def __init__( self, num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: str = "scaled_linear", trained_betas=None, solver_order: int = 2, prediction_type: str = "epsilon", thresholding: bool = False, dynamic_thresholding_ratio: float = 0.995, sample_max_value: float = 1.0, algorithm_type: str = "dpmsolver++", solver_type: str = "midpoint", lower_order_final: bool = True, euler_at_final: bool = False, use_karras_sigmas: bool = False, use_lu_lambdas: bool = False, use_exponential_sigmas: bool = False, use_beta_sigmas: bool = False, use_flow_sigmas: bool = False, final_sigmas_type: str = "zero", lambda_min_clipped: float = -float("inf"), variance_type: str | None = None, timestep_spacing: str = "linspace", steps_offset: int = 0, rescale_betas_zero_snr: bool = False, flow_shift: float | None = None, **kwargs, ): self.num_train_timesteps = num_train_timesteps self._inner = DiffusersDPMSolverMultistepScheduler( num_train_timesteps=num_train_timesteps, beta_start=beta_start, beta_end=beta_end, beta_schedule=beta_schedule, trained_betas=trained_betas, solver_order=solver_order, prediction_type=prediction_type, thresholding=thresholding, dynamic_thresholding_ratio=dynamic_thresholding_ratio, sample_max_value=sample_max_value, algorithm_type=algorithm_type, solver_type=solver_type, lower_order_final=lower_order_final, euler_at_final=euler_at_final, use_karras_sigmas=use_karras_sigmas, use_lu_lambdas=use_lu_lambdas, use_exponential_sigmas=use_exponential_sigmas, use_beta_sigmas=use_beta_sigmas, use_flow_sigmas=use_flow_sigmas, flow_shift=flow_shift, final_sigmas_type=final_sigmas_type, lambda_min_clipped=lambda_min_clipped, variance_type=variance_type, timestep_spacing=timestep_spacing, steps_offset=steps_offset, rescale_betas_zero_snr=rescale_betas_zero_snr, ) self.timesteps = self._inner.timesteps self.order = solver_order self._flow_shift = flow_shift self._begin_index: int | None = None BaseScheduler.__init__(self) def set_shift(self, shift: float) -> None: self._flow_shift = shift def set_begin_index(self, begin_index: int = 0) -> None: self._begin_index = begin_index @property def begin_index(self) -> int | None: return self._begin_index def set_timesteps(self, num_inference_steps: int, device=None, **kwargs): self._inner.set_timesteps(num_inference_steps, device=device, **kwargs) self.timesteps = self._inner.timesteps def scale_model_input( self, sample: torch.Tensor, timestep: int | None = None ) -> torch.Tensor: return self._inner.scale_model_input(sample, timestep) def step( self, model_output: torch.Tensor, timestep: int, sample: torch.Tensor, **kwargs, ): return self._inner.step(model_output, timestep, sample, **kwargs) @property def sigmas(self): return getattr(self._inner, "sigmas", None) @property def init_noise_sigma(self): return self._inner.init_noise_sigma def add_noise( self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.Tensor, ) -> torch.Tensor: return self._inner.add_noise(original_samples, noise, timesteps) EntryClass = DPMSolverMultistepScheduler ================================================ FILE: python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== # # Modified from diffusers==0.29.2 # # ============================================================================== import math from dataclasses import dataclass from typing import Any import numpy as np import scipy.stats import torch from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.schedulers.scheduling_utils import SchedulerMixin from diffusers.utils import BaseOutput from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @dataclass class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput): """ Output class for the scheduler's `step` function output. Args: prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the denoising loop. """ prev_sample: torch.FloatTensor class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """ Euler scheduler. This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic methods the library implements for all schedulers such as loading and saving. Args: num_train_timesteps (`int`, defaults to 1000): The number of diffusion steps to train the model. shift (`float`, defaults to 1.0): The shift value for the timestep schedule. use_dynamic_shifting (`bool`, defaults to False): Whether to apply timestep shifting on-the-fly based on the image resolution. base_shift (`float`, defaults to 0.5): Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent with desired output. max_shift (`float`, defaults to 1.15): Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be more exaggerated or stylized. base_image_seq_len (`int`, defaults to 256): The base image sequence length. max_image_seq_len (`int`, defaults to 4096): The maximum image sequence length. invert_sigmas (`bool`, defaults to False): Whether to invert the sigmas. shift_terminal (`float`, defaults to None): The end value of the shifted timestep schedule. use_karras_sigmas (`bool`, defaults to False): Whether to use Karras sigmas for step sizes in the noise schedule during sampling. use_exponential_sigmas (`bool`, defaults to False): Whether to use exponential sigmas for step sizes in the noise schedule during sampling. use_beta_sigmas (`bool`, defaults to False): Whether to use beta sigmas for step sizes in the noise schedule during sampling. time_shift_type (`str`, defaults to "exponential"): The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear". stochastic_sampling (`bool`, defaults to False): Whether to use stochastic sampling. """ _compatibles: list[Any] = [] order = 1 @register_to_config def __init__( self, num_train_timesteps: int = 1000, shift: float = 1.0, use_dynamic_shifting: bool = False, base_shift: float | None = 0.5, max_shift: float | None = 1.15, base_image_seq_len: int | None = 256, max_image_seq_len: int | None = 4096, invert_sigmas: bool = False, shift_terminal: float | None = None, use_karras_sigmas: bool | None = False, use_exponential_sigmas: bool | None = False, use_beta_sigmas: bool | None = False, time_shift_type: str = "exponential", stochastic_sampling: bool = False, ): if ( sum( [ self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas, ] ) > 1 ): raise ValueError( "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." ) if time_shift_type not in {"exponential", "linear"}: raise ValueError( "`time_shift_type` must either be 'exponential' or 'linear'." ) timesteps = np.linspace( 1, num_train_timesteps, num_train_timesteps, dtype=np.float32 )[::-1].copy() timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) sigmas = timesteps / num_train_timesteps if not use_dynamic_shifting: # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) self.timesteps = sigmas * num_train_timesteps self.num_train_timesteps = num_train_timesteps self._step_index: int | None = None self._begin_index: int | None = None self._shift = shift self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication self.sigma_min = self.sigmas[-1].item() self.sigma_max = self.sigmas[0].item() BaseScheduler.__init__(self) @property def shift(self) -> float: """ The value used for shifting. """ return self._shift @property def step_index(self) -> int | None: """ The index counter for current timestep. It will increase 1 after each scheduler step. """ return self._step_index @property def begin_index(self) -> int | None: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. """ return self._begin_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index def set_begin_index(self, begin_index: int = 0) -> None: """ Sets the begin index for the scheduler. This function should be run from pipeline before the inference. Args: begin_index (`int`): The begin index for the scheduler. """ self._begin_index = begin_index def set_shift(self, shift: float) -> None: self._shift = shift def scale_noise( self, sample: torch.FloatTensor, timestep: float | torch.FloatTensor, noise: torch.FloatTensor | None = None, ) -> torch.FloatTensor: """ Forward process in flow-matching """ # Make sure sigmas and timesteps have the same device and dtype as original_samples sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype) if sample.device.type == "mps" and torch.is_floating_point(timestep): # mps does not support float64 schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) assert isinstance(timestep, torch.Tensor) timestep = timestep.to(sample.device, dtype=torch.float32) else: schedule_timesteps = self.timesteps.to(sample.device) assert isinstance(timestep, torch.Tensor) timestep = timestep.to(sample.device) # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index if self.begin_index is None: step_indices = [ self.index_for_timestep(t, schedule_timesteps) for t in timestep ] elif self.step_index is not None: # add_noise is called after first denoising step (for inpainting) step_indices = [self.step_index] * timestep.shape[0] else: # add noise is called before first denoising step to create initial latent(img2img) step_indices = [self.begin_index] * timestep.shape[0] sigma = sigmas[step_indices].flatten() while len(sigma.shape) < len(sample.shape): sigma = sigma.unsqueeze(-1) sample = sigma * noise + (1.0 - sigma) * sample return sample def _sigma_to_t(self, sigma: float) -> float: return sigma * self.config.num_train_timesteps def time_shift( self, mu: float, sigma: float, t: torch.Tensor | np.ndarray ) -> torch.Tensor | np.ndarray: if self.config.time_shift_type == "exponential": return self._time_shift_exponential(mu, sigma, t) elif self.config.time_shift_type == "linear": return self._time_shift_linear(mu, sigma, t) else: raise ValueError(f"Unknown time_shift_type: {self.config.time_shift_type}") def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor: r""" Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config value. Reference: https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51 Args: t (`torch.Tensor`): A tensor of timesteps to be stretched and shifted. Returns: `torch.Tensor`: A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`. """ one_minus_z = 1 - t scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal) stretched_t = 1 - (one_minus_z / scale_factor) return stretched_t def set_timesteps( self, num_inference_steps: int | None = None, device: str | torch.device = None, sigmas: list[float] | None = None, mu: float | None = None, timesteps: list[float] | None = None, ) -> None: """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (`int`, *optional*): The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. sigmas (`List[float]`, *optional*): Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed automatically. mu (`float`, *optional*): Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep shifting. timesteps (`List[float]`, *optional*): Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed automatically. """ if self.config.use_dynamic_shifting and mu is None: raise ValueError( "`mu` must be passed when `use_dynamic_shifting` is set to be `True`" ) if ( sigmas is not None and timesteps is not None and len(sigmas) != len(timesteps) ): raise ValueError("`sigmas` and `timesteps` should have the same length") if num_inference_steps is not None: if (sigmas is not None and len(sigmas) != num_inference_steps) or ( timesteps is not None and len(timesteps) != num_inference_steps ): raise ValueError( "`sigmas` and `timesteps` should have the same length as num_inference_steps, if `num_inference_steps` is provided" ) else: if sigmas is not None: num_inference_steps = len(sigmas) elif timesteps is not None: num_inference_steps = len(timesteps) else: raise ValueError( "Either num_inference_steps, sigmas, or timesteps must be provided" ) self.num_inference_steps = num_inference_steps # 1. Prepare default sigmas is_timesteps_provided = timesteps is not None timesteps_array: np.ndarray | None = None if is_timesteps_provided: assert timesteps is not None timesteps_array = np.array(timesteps).astype(np.float32) sigmas_array: np.ndarray if sigmas is None: if timesteps_array is None: timesteps_array = np.linspace( self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps, ) sigmas_array = timesteps_array / self.config.num_train_timesteps else: sigmas_array = np.array(sigmas).astype(np.float32) num_inference_steps = len(sigmas_array) # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of # "exponential" or "linear" type is applied if self.config.use_dynamic_shifting: assert mu is not None, "mu cannot be None when use_dynamic_shifting is True" sigmas_array = self.time_shift(mu, 1.0, sigmas_array) else: sigmas_array = ( self.shift * sigmas_array / (1 + (self.shift - 1) * sigmas_array) ) # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value if self.config.shift_terminal: sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32) sigmas_tensor = self.stretch_shift_to_terminal(sigmas_tensor) sigmas_array = sigmas_tensor.numpy() # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules if self.config.use_karras_sigmas: sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32) sigmas_tensor = self._convert_to_karras( in_sigmas=sigmas_tensor, num_inference_steps=num_inference_steps ) sigmas_array = sigmas_tensor.numpy() elif self.config.use_exponential_sigmas: sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32) sigmas_tensor = self._convert_to_exponential( in_sigmas=sigmas_tensor, num_inference_steps=num_inference_steps ) sigmas_array = sigmas_tensor.numpy() elif self.config.use_beta_sigmas: sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32) sigmas_tensor = self._convert_to_beta( in_sigmas=sigmas_tensor, num_inference_steps=num_inference_steps ) sigmas_array = sigmas_tensor.numpy() # 5. Convert sigmas and timesteps to tensors and move to specified device sigmas_tensor = torch.from_numpy(sigmas_array).to( dtype=torch.float32, device=device ) if not is_timesteps_provided: timesteps_tensor = sigmas_tensor * self.config.num_train_timesteps else: assert timesteps_array is not None timesteps_tensor = torch.from_numpy(timesteps_array).to( dtype=torch.float32, device=device ) # 6. Append the terminal sigma value. # If a model requires inverted sigma schedule for denoising but timesteps without inversion, the # `invert_sigmas` flag can be set to `True`. This case is only required in Mochi if self.config.invert_sigmas: sigmas_tensor = 1.0 - sigmas_tensor timesteps_tensor = sigmas_tensor * self.config.num_train_timesteps sigmas_tensor = torch.cat( [sigmas_tensor, torch.ones(1, device=sigmas_tensor.device)] ) else: sigmas_tensor = torch.cat( [sigmas_tensor, torch.zeros(1, device=sigmas_tensor.device)] ) self.timesteps = timesteps_tensor self.sigmas = sigmas_tensor self._step_index = None self._begin_index = None def index_for_timestep( self, timestep: float | torch.FloatTensor, schedule_timesteps: torch.Tensor | None = None, ) -> int: if schedule_timesteps is None: schedule_timesteps = self.timesteps indices = (schedule_timesteps == timestep).nonzero() # The sigma index that is taken for the **very** first `step` # is always the second index (or the last index if there is only 1) # This way we can ensure we don't accidentally skip a sigma in # case we start in the middle of the denoising schedule (e.g. for image-to-image) pos = 1 if len(indices) > 1 else 0 return indices[pos].item() def _init_step_index(self, timestep: float | torch.FloatTensor) -> None: if self.begin_index is None: if isinstance(timestep, torch.Tensor): timestep = timestep.to(self.timesteps.device) self._step_index = self.index_for_timestep(timestep) else: self._step_index = self._begin_index def step( self, model_output: torch.FloatTensor, timestep: int | torch.Tensor, sample: torch.FloatTensor, s_churn: float = 0.0, s_tmin: float = 0.0, s_tmax: float = float("inf"), s_noise: float = 1.0, generator: torch.Generator | None = None, per_token_timesteps: torch.Tensor | None = None, return_dict: bool = True, ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor, ...]: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). Args: model_output (`torch.FloatTensor`): The direct output from learned diffusion model. timestep (`int` or `torch.Tensor`): The current discrete timestep in the diffusion chain. sample (`torch.FloatTensor`): A current instance of a sample created by the diffusion process. s_churn (`float`): s_tmin (`float`): s_tmax (`float`): s_noise (`float`, defaults to 1.0): Scaling factor for noise added to the sample. generator (`torch.Generator`, *optional*): A random number generator. per_token_timesteps (`torch.Tensor`, *optional*): The timesteps for each token in the sample. return_dict (`bool`): Whether or not to return a [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple. Returns: [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`: If return_dict is `True`, [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned, otherwise a tuple is returned where the first element is the sample tensor. """ if isinstance(timestep, int | torch.IntTensor | torch.LongTensor): raise ValueError( ( "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass" " one of the `scheduler.timesteps` as a timestep." ), ) if self.step_index is None: self._init_step_index(timestep) # Upcast to avoid precision issues when computing prev_sample sample = sample.to(torch.float32) if per_token_timesteps is not None: per_token_sigmas = per_token_timesteps / self.config.num_train_timesteps sigmas = self.sigmas[:, None, None] lower_mask = sigmas < per_token_sigmas[None] - 1e-6 lower_sigmas = lower_mask * sigmas lower_sigmas, _ = lower_sigmas.max(dim=0) current_sigma = per_token_sigmas[..., None] next_sigma = lower_sigmas[..., None] dt = current_sigma - next_sigma else: assert self.step_index is not None, "step_index should not be None" sigma_idx = self.step_index sigma = self.sigmas[sigma_idx] sigma_next = self.sigmas[sigma_idx + 1] current_sigma = sigma next_sigma = sigma_next dt = sigma_next - sigma if self.config.stochastic_sampling: x0 = sample - current_sigma * model_output noise = torch.randn_like(sample) prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise else: prev_sample = sample + dt * model_output # upon completion increase step index by one assert self._step_index is not None, "_step_index should not be None" self._step_index += 1 if per_token_timesteps is None: # Cast sample back to model compatible dtype prev_sample = prev_sample.to(model_output.dtype) if isinstance(prev_sample, torch.Tensor | float) and not return_dict: return (prev_sample,) return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample) # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras def _convert_to_karras( self, in_sigmas: torch.Tensor, num_inference_steps: int ) -> torch.Tensor: """Constructs the noise schedule of Karras et al. (2022).""" # Hack to make sure that other schedulers which copy this function don't break # TODO: Add this logic to the other schedulers if hasattr(self.config, "sigma_min"): sigma_min = self.config.sigma_min else: sigma_min = None if hasattr(self.config, "sigma_max"): sigma_max = self.config.sigma_max else: sigma_max = None sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() rho = 7.0 # 7.0 is the value used in the paper ramp = np.linspace(0, 1, num_inference_steps) min_inv_rho = sigma_min ** (1 / rho) max_inv_rho = sigma_max ** (1 / rho) sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho return sigmas # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential def _convert_to_exponential( self, in_sigmas: torch.Tensor, num_inference_steps: int ) -> torch.Tensor: """Constructs an exponential noise schedule.""" # Hack to make sure that other schedulers which copy this function don't break # TODO: Add this logic to the other schedulers if hasattr(self.config, "sigma_min"): sigma_min = self.config.sigma_min else: sigma_min = None if hasattr(self.config, "sigma_max"): sigma_max = self.config.sigma_max else: sigma_max = None sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() sigmas = np.exp( np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps) ) return sigmas # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta def _convert_to_beta( self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6, ) -> torch.Tensor: """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" # Hack to make sure that other schedulers which copy this function don't break # TODO: Add this logic to the other schedulers if hasattr(self.config, "sigma_min"): sigma_min = self.config.sigma_min else: sigma_min = None if hasattr(self.config, "sigma_max"): sigma_max = self.config.sigma_max else: sigma_max = None sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() sigmas = np.array( [ sigma_min + (ppf * (sigma_max - sigma_min)) for ppf in [ scipy.stats.beta.ppf(timestep, alpha, beta) for timestep in 1 - np.linspace(0, 1, num_inference_steps) ] ] ) return sigmas def _time_shift_exponential( self, mu: float, sigma: float, t: torch.Tensor | np.ndarray ) -> torch.Tensor | np.ndarray: if isinstance(t, np.ndarray): return np.exp(mu) / (np.exp(mu) + (1 / t - 1) ** sigma) else: return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) def _time_shift_linear( self, mu: float, sigma: float, t: torch.Tensor | np.ndarray ) -> torch.Tensor | np.ndarray: return mu / (mu + (1 / t - 1) ** sigma) def add_noise( self, clean_latent: torch.Tensor, noise: torch.Tensor, timestep: torch.IntTensor, ) -> torch.Tensor: """ Args: clean_latent: the clean latent with shape [B, C, H, W], where B is batch_size or batch_size * num_frames noise: the noise with shape [B, C, H, W] timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames] Returns: the corrupted latent with shape [B, C, H, W] """ # If timestep is [bs, num_frames] if timestep.ndim == 2: timestep = timestep.flatten(0, 1) assert timestep.numel() == clean_latent.shape[0] elif timestep.ndim == 1: # If timestep is [1] if timestep.shape[0] == 1: timestep = timestep.expand(clean_latent.shape[0]) else: assert timestep.numel() == clean_latent.shape[0] else: raise ValueError(f"[add_noise] Invalid timestep shape: {timestep.shape}") # timestep shape should be [B] self.sigmas = self.sigmas.to(noise.device) self.timesteps = self.timesteps.to(noise.device) timestep_id = torch.argmin( (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1 ) sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) sample = (1 - sigma) * clean_latent + sigma * noise return sample.type_as(noise) def scale_model_input( self, sample: torch.Tensor, timestep: int | None = None ) -> torch.Tensor: return sample def __len__(self) -> int: return 0 EntryClass = FlowMatchEulerDiscreteScheduler ================================================ FILE: python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py # Convert unipc for flow matching # Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. import math from typing import Any import numpy as np import torch from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.schedulers.scheduling_utils import ( KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput, ) from diffusers.utils import deprecate from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """ `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models. This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic methods the library implements for all schedulers such as loading and saving. Args: num_train_timesteps (`int`, defaults to 1000): The number of diffusion steps to train the model. solver_order (`int`, default `2`): The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1` due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for unconditional sampling. prediction_type (`str`, defaults to "flow_prediction"): Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts the flow of the diffusion process. thresholding (`bool`, defaults to `False`): Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such as Stable Diffusion. dynamic_thresholding_ratio (`float`, defaults to 0.995): The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. sample_max_value (`float`, defaults to 1.0): The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`. predict_x0 (`bool`, defaults to `True`): Whether to use the updating algorithm on the predicted x0. solver_type (`str`, default `bh2`): Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2` otherwise. lower_order_final (`bool`, default `True`): Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10. disable_corrector (`list`, default `[]`): Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)` and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is usually disabled during the first few steps. solver_p (`SchedulerMixin`, default `None`): Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`. use_karras_sigmas (`bool`, *optional*, defaults to `False`): Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, the sigmas are determined according to a sequence of noise levels {σi}. use_exponential_sigmas (`bool`, *optional*, defaults to `False`): Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. timestep_spacing (`str`, defaults to `"linspace"`): The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. steps_offset (`int`, defaults to 0): An offset added to the inference steps, as required by some model families. final_sigmas_type (`str`, defaults to `"zero"`): The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0. """ _compatibles = [e.name for e in KarrasDiffusionSchedulers] order = 1 @register_to_config def __init__( self, num_train_timesteps: int = 1000, solver_order: int = 2, prediction_type: str = "flow_prediction", shift: float | None = 1.0, use_dynamic_shifting=False, thresholding: bool = False, dynamic_thresholding_ratio: float = 0.995, sample_max_value: float = 1.0, predict_x0: bool = True, solver_type: str = "bh2", lower_order_final: bool = True, disable_corrector: tuple = (), solver_p: SchedulerMixin = None, timestep_spacing: str = "linspace", steps_offset: int = 0, final_sigmas_type: str | None = "zero", # "zero", "sigma_min" **kwargs, ): if solver_type not in ["bh1", "bh2"]: if solver_type in ["midpoint", "heun", "logrho"]: self.register_to_config(solver_type="bh2") else: raise NotImplementedError( f"{solver_type} is not implemented for {self.__class__}" ) self.predict_x0 = predict_x0 # setable values self.num_inference_steps: int | None = None alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps)[ ::-1 ].copy() sigmas = 1.0 - alphas sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32) if not use_dynamic_shifting: # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution assert shift is not None sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) # pyright: ignore self.sigmas = sigmas self.sigma_min = self.sigmas[-1].item() self.sigma_max = self.sigmas[0].item() self.timesteps = sigmas * num_train_timesteps self.num_train_timesteps = num_train_timesteps self.model_outputs = [None] * solver_order self.timestep_list: list[Any | None] = [None] * solver_order self.lower_order_nums = 0 self.disable_corrector = list(disable_corrector) self.solver_p = solver_p self.last_sample = None self._step_index: int | None = None self._begin_index: int | None = None BaseScheduler.__init__(self) @property def step_index(self): """ The index counter for current timestep. It will increase 1 after each scheduler step. """ return self._step_index @property def begin_index(self): """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. """ return self._begin_index def set_shift(self, shift: float) -> None: self.config.shift = shift # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index def set_begin_index(self, begin_index: int = 0): """ Sets the begin index for the scheduler. This function should be run from pipeline before the inference. Args: begin_index (`int`): The begin index for the scheduler. """ self._begin_index = begin_index # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps def set_timesteps( self, num_inference_steps: int | None = None, device: str | torch.device = None, sigmas: list[float] | None = None, mu: float | None | None = None, shift: float | None | None = None, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (`int`): Total number of the spacing of the time steps. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. """ if self.config.use_dynamic_shifting and mu is None: raise ValueError( " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`" ) if sigmas is None: assert num_inference_steps is not None sigmas = np.linspace( self.sigma_max, self.sigma_min, num_inference_steps + 1 ).copy()[ :-1 ] # pyright: ignore if self.config.use_dynamic_shifting: assert mu is not None sigmas = self.time_shift(mu, 1.0, sigmas) # pyright: ignore else: if shift is None: shift = self.config.shift assert isinstance(sigmas, np.ndarray) sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) # pyright: ignore if self.config.final_sigmas_type == "sigma_min": sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5 elif self.config.final_sigmas_type == "zero": sigma_last = 0 else: raise ValueError( f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" ) timesteps = sigmas * self.config.num_train_timesteps sigmas = np.concatenate([sigmas, [sigma_last]]).astype( np.float32 ) # pyright: ignore self.sigmas = torch.from_numpy(sigmas).to(device=device) self.timesteps = torch.from_numpy(timesteps).to( device=device, dtype=torch.int64 ) self.num_inference_steps = len(timesteps) self.model_outputs = [ None, ] * self.config.solver_order self.lower_order_nums = 0 self.last_sample = None if self.solver_p: self.solver_p.set_timesteps(self.num_inference_steps, device=device) # add an index counter for schedulers that allow duplicated timesteps self._step_index = None self._begin_index = None # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: """ "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing pixels from saturation at each step. We find that dynamic thresholding results in significantly better photorealism as well as better image-text alignment, especially when using very large guidance weights." https://arxiv.org/abs/2205.11487 """ dtype = sample.dtype batch_size, channels, *remaining_dims = sample.shape if dtype not in (torch.float32, torch.float64): sample = ( sample.float() ) # upcast for quantile calculation, and clamp not implemented for cpu half # Flatten sample for doing quantile calculation along each image sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) abs_sample = sample.abs() # "a certain percentile absolute pixel value" s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) s = torch.clamp( s, min=1, max=self.config.sample_max_value ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 sample = ( torch.clamp(sample, -s, s) / s ) # "we threshold xt0 to the range [-s, s] and then divide by s" sample = sample.reshape(batch_size, channels, *remaining_dims) sample = sample.to(dtype) return sample def _sigma_to_alpha_sigma_t(self, sigma) -> tuple[Any, Any]: return 1 - sigma, sigma # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps def time_shift(self, mu: float, sigma: float, t: torch.Tensor): return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) def convert_model_output( self, model_output: torch.Tensor, *args, sample: torch.Tensor = None, **kwargs, ) -> torch.Tensor: r""" Convert the model output to the corresponding type the UniPC algorithm needs. Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model. timestep (`int`): The current discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. Returns: `torch.Tensor`: The converted model output. """ timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) if sample is None: if len(args) > 1: sample = args[1] else: raise ValueError("missing `sample` as a required keyword argument") if timestep is not None: deprecate( "timesteps", "1.0.0", "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", ) if self.predict_x0: if self.config.prediction_type == "flow_prediction": sigma_t = self.sigmas[self.step_index] x0_pred = sample - sigma_t * model_output else: raise ValueError( f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." ) if self.config.thresholding: x0_pred = self._threshold_sample(x0_pred) return x0_pred else: if self.config.prediction_type == "flow_prediction": sigma_t = self.sigmas[self.step_index] epsilon = sample - (1 - sigma_t) * model_output else: raise ValueError( f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." ) if self.config.thresholding: sigma_t = self.sigmas[self.step_index] x0_pred = sample - sigma_t * model_output x0_pred = self._threshold_sample(x0_pred) epsilon = model_output + x0_pred return epsilon def multistep_uni_p_bh_update( self, model_output: torch.Tensor, *args, sample: torch.Tensor = None, order: int | None = None, # pyright: ignore **kwargs, ) -> torch.Tensor: """ One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified. Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model at the current timestep. prev_timestep (`int`): The previous discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. order (`int`): The order of UniP at this timestep (corresponds to the *p* in UniPC-p). Returns: `torch.Tensor`: The sample tensor at the previous timestep. """ prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None) if sample is None: if len(args) > 1: sample = args[1] else: raise ValueError(" missing `sample` as a required keyword argument") if order is None: if len(args) > 2: order = args[2] else: raise ValueError(" missing `order` as a required keyword argument") if prev_timestep is not None: deprecate( "prev_timestep", "1.0.0", "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", ) model_output_list = self.model_outputs s0 = self.timestep_list[-1] m0 = model_output_list[-1] x = sample if self.solver_p: x_t = self.solver_p.step(model_output, s0, x).prev_sample return x_t sigma_t, sigma_s0 = ( self.sigmas[self.step_index + 1], self.sigmas[self.step_index], ) # pyright: ignore alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) lambda_t = torch.log(alpha_t) - torch.log(sigma_t) lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) h = lambda_t - lambda_s0 device = sample.device rks = [] D1s: list[Any] | None = [] sigmas = self.sigmas.to(device=device) for i in range(1, order): si = self.step_index - i # pyright: ignore mi = model_output_list[-(i + 1)] alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(sigmas[si]) lambda_si = torch.log(alpha_si) - torch.log(sigma_si) rk = (lambda_si - lambda_s0) / h rks.append(rk) assert mi is not None D1s.append((mi - m0) / rk) # pyright: ignore if len(rks) > 0: rks = torch.stack(rks) one = torch.ones(1, device=device, dtype=rks.dtype) rks = torch.cat([rks, one]) else: rks = torch.ones(1, device=device, dtype=h.dtype) R = [] b = [] hh = -h if self.predict_x0 else h h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 h_phi_k = h_phi_1 / hh - 1 factorial_i = 1 if self.config.solver_type == "bh1": B_h = hh elif self.config.solver_type == "bh2": B_h = torch.expm1(hh) else: raise NotImplementedError() for i in range(1, order + 1): R.append(torch.pow(rks, i - 1)) b.append(h_phi_k * factorial_i / B_h) factorial_i *= i + 1 h_phi_k = h_phi_k / hh - 1 / factorial_i R = torch.stack(R) b = torch.stack(b) if D1s is not None and len(D1s) > 0: D1s = torch.stack(D1s, dim=1) # (B, K) # for order 2, we use a simplified version if order == 2: rhos_p = 0.5 * torch.ones(1, dtype=x.dtype, device=device) else: assert isinstance(R, torch.Tensor) rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype) else: D1s = None if self.predict_x0: x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 if D1s is not None: pred_res = torch.einsum( "k,bkc...->bc...", rhos_p, D1s ) # pyright: ignore else: pred_res = 0 x_t = x_t_ - alpha_t * B_h * pred_res else: x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 if D1s is not None: pred_res = torch.einsum( "k,bkc...->bc...", rhos_p, D1s ) # pyright: ignore else: pred_res = 0 x_t = x_t_ - sigma_t * B_h * pred_res x_t = x_t.to(x.dtype) return x_t def multistep_uni_c_bh_update( self, this_model_output: torch.Tensor, *args, last_sample: torch.Tensor = None, this_sample: torch.Tensor = None, order: int | None = None, # pyright: ignore **kwargs, ) -> torch.Tensor: """ One step for the UniC (B(h) version). Args: this_model_output (`torch.Tensor`): The model outputs at `x_t`. this_timestep (`int`): The current timestep `t`. last_sample (`torch.Tensor`): The generated sample before the last predictor `x_{t-1}`. this_sample (`torch.Tensor`): The generated sample after the last predictor `x_{t}`. order (`int`): The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`. Returns: `torch.Tensor`: The corrected sample tensor at the current timestep. """ this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None) if last_sample is None: if len(args) > 1: last_sample = args[1] else: raise ValueError(" missing`last_sample` as a required keyword argument") if this_sample is None: if len(args) > 2: this_sample = args[2] else: raise ValueError(" missing`this_sample` as a required keyword argument") if order is None: if len(args) > 3: order = args[3] else: raise ValueError(" missing`order` as a required keyword argument") if this_timestep is not None: deprecate( "this_timestep", "1.0.0", "Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", ) model_output_list = self.model_outputs m0 = model_output_list[-1] x = last_sample x_t = this_sample model_t = this_model_output sigma_t, sigma_s0 = ( self.sigmas[self.step_index], self.sigmas[self.step_index - 1], ) # pyright: ignore alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) lambda_t = torch.log(alpha_t) - torch.log(sigma_t) lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) h = lambda_t - lambda_s0 device = this_sample.device # Build rks and D1s fully on device to avoid any host-device sync # Fast paths for small orders (common cases: 1 or 2) if order == 1: rks = torch.ones(1, device=device, dtype=h.dtype) D1s = None elif order == 2: # order == 2 -> only one historical point is used si = self.step_index - 2 # i = 1 mi = model_output_list[-2] alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) lambda_si = torch.log(alpha_si) - torch.log(sigma_si) rk = (lambda_si - lambda_s0) / h # 0-dim tensor on device # rks = [rk, 1.0] but keep it on device without list->tensor sync rks = torch.stack((rk, torch.ones_like(rk))) assert mi is not None # D1s shape: (B, K=1, C, ...) to match later einsum over K D1s = ((mi - m0) / rk).unsqueeze(1) # pyright: ignore else: rks_list = [] D1s_list = [] for i in range(1, order): si = self.step_index - (i + 1) mi = model_output_list[-(i + 1)] alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) lambda_si = torch.log(alpha_si) - torch.log(sigma_si) rk = (lambda_si - lambda_s0) / h rks_list.append(rk) assert mi is not None D1s_list.append((mi - m0) / rk) # pyright: ignore # Append 1.0 as a device tensor to rks rks = torch.stack(rks_list + [torch.ones_like(rks_list[0])]) D1s = torch.stack(D1s_list, dim=1) if len(D1s_list) > 0 else None R = [] b = [] hh = -h if self.predict_x0 else h h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 h_phi_k = h_phi_1 / hh - 1 factorial_i = 1 if self.config.solver_type == "bh1": B_h = hh elif self.config.solver_type == "bh2": B_h = torch.expm1(hh) else: raise NotImplementedError() for i in range(1, order + 1): R.append(torch.pow(rks, i - 1)) b.append(h_phi_k * factorial_i / B_h) factorial_i *= i + 1 h_phi_k = h_phi_k / hh - 1 / factorial_i R = torch.stack(R) # Avoid torch.tensor(list_of_gpu_scalars) which syncs to host b = torch.stack(b) # D1s is already prepared above for order==2; remains None for order==1 # for order 1, we use a simplified version if order == 1: rhos_c = 0.5 * torch.ones(1, dtype=x.dtype, device=device) elif order == 2: # Manually solve the 2x2 linear system to avoid device synchronization from torch.linalg.solve # R = [[1, 1], [rk, 1]], where rk = rks[0] rk = rks[0] det = 1 - rk # Using Cramer's rule to solve for rhos_c = [x0, x1] # x0 = (b0 - b1) / det # x1 = (b1 - rk * b0) / det rhos_c_0 = (b[0] - b[1]) / det rhos_c_1 = (b[1] - rk * b[0]) / det rhos_c = torch.stack([rhos_c_0, rhos_c_1]) else: rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype) if self.predict_x0: x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 if D1s is not None: corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) else: corr_res = 0 D1_t = model_t - m0 x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) else: x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 if D1s is not None: corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) else: corr_res = 0 D1_t = model_t - m0 x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) x_t = x_t.to(x.dtype) return x_t def index_for_timestep(self, timestep, schedule_timesteps=None) -> int: if schedule_timesteps is None: schedule_timesteps = self.timesteps indices = (schedule_timesteps == timestep).nonzero() # The sigma index that is taken for the **very** first `step` # is always the second index (or the last index if there is only 1) # This way we can ensure we don't accidentally skip a sigma in # case we start in the middle of the denoising schedule (e.g. for image-to-image) pos = 1 if len(indices) > 1 else 0 step_index: int = indices[pos].item() return step_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index def _init_step_index(self, timestep) -> None: """ Initialize the step_index counter for the scheduler. """ if self.begin_index is None: if isinstance(timestep, torch.Tensor): timestep = timestep.to(self.timesteps.device) self._step_index = self.index_for_timestep(timestep) else: self._step_index = self._begin_index def step( self, model_output: torch.Tensor, timestep: int | torch.Tensor, sample: torch.Tensor, return_dict: bool = True, generator=None, ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep UniPC. Args: model_output (`torch.Tensor`): The direct output from learned diffusion model. timestep (`int`): The current discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. return_dict (`bool`): Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`. Returns: [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`: If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a tuple is returned where the first element is the sample tensor. """ if self.num_inference_steps is None: raise ValueError( "Number of inference steps is 'None', you need to call 'set_timesteps' after creating the scheduler" ) if self.step_index is None: self._init_step_index(timestep) use_corrector = ( self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None # pyright: ignore ) sample = sample.to(model_output.device) model_output_convert = self.convert_model_output(model_output, sample=sample) if use_corrector: sample = self.multistep_uni_c_bh_update( this_model_output=model_output_convert, last_sample=self.last_sample, this_sample=sample, order=self.this_order, ) for i in range(self.config.solver_order - 1): self.model_outputs[i] = self.model_outputs[i + 1] self.timestep_list[i] = self.timestep_list[i + 1] self.model_outputs[-1] = model_output_convert self.timestep_list[-1] = timestep # pyright: ignore if self.config.lower_order_final: this_order = min( self.config.solver_order, len(self.timesteps) - self.step_index ) # pyright: ignore else: this_order = self.config.solver_order self.this_order: int = min( this_order, self.lower_order_nums + 1 ) # warmup for multistep assert self.this_order > 0 self.last_sample = sample prev_sample = self.multistep_uni_p_bh_update( model_output=model_output, # pass the original non-converted model output, in case solver-p is used sample=sample, order=self.this_order, ) if self.lower_order_nums < self.config.solver_order: self.lower_order_nums += 1 # upon completion increase step index by one assert self._step_index is not None self._step_index += 1 # pyright: ignore if not return_dict: return (prev_sample,) return SchedulerOutput(prev_sample=prev_sample) def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. Args: sample (`torch.Tensor`): The input sample. Returns: `torch.Tensor`: A scaled input sample. """ return sample # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise def add_noise( self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor, ) -> torch.Tensor: # Make sure sigmas and timesteps have the same device and dtype as original_samples sigmas = self.sigmas.to( device=original_samples.device, dtype=original_samples.dtype ) if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): # mps does not support float64 schedule_timesteps = self.timesteps.to( original_samples.device, dtype=torch.float32 ) timesteps = timesteps.to(original_samples.device, dtype=torch.float32) else: schedule_timesteps = self.timesteps.to(original_samples.device) timesteps = timesteps.to(original_samples.device) # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index if self.begin_index is None: step_indices = [ self.index_for_timestep(t, schedule_timesteps) for t in timesteps ] elif self.step_index is not None: # add_noise is called after first denoising step (for inpainting) step_indices = [self.step_index] * timesteps.shape[0] else: # add noise is called before first denoising step to create initial latent(img2img) step_indices = [self.begin_index] * timesteps.shape[0] sigma = sigmas[step_indices].flatten() while len(sigma.shape) < len(original_samples.shape): sigma = sigma.unsqueeze(-1) alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) noisy_samples = alpha_t * original_samples + sigma_t * noise return noisy_samples EntryClass = FlowUniPCMultistepScheduler ================================================ FILE: python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py ================================================ # SPDX-License-Identifier: Apache-2.0 # Adapted from Helios diffusers scheduler: # https://github.com/BestWishYsh/Helios """ Helios scheduler implementing flow-matching with UniPC/Euler solvers. For Phase 1 T2V (stages=1), this simplifies to standard flow-matching with dynamic shifting and UniPC multistep solver. """ import math from dataclasses import dataclass import numpy as np import torch @dataclass class HeliosSchedulerOutput: prev_sample: torch.FloatTensor model_outputs: torch.FloatTensor | None = None last_sample: torch.FloatTensor | None = None this_order: int | None = None class HeliosSchedulerConfig: """Mimics diffusers config interface for scheduler parameters.""" def __init__(self, **kwargs): for k, v in kwargs.items(): setattr(self, k, v) def get(self, key, default=None): return getattr(self, key, default) class HeliosScheduler: """ Helios multi-stage scheduler supporting Euler, UniPC, and DMD solvers. For Phase 1 T2V with stages=1, this is a standard flow-matching scheduler with optional time shifting and UniPC multistep updates. """ order = 1 def __init__( self, num_train_timesteps: int = 1000, shift: float = 1.0, stages: int = 1, stage_range: list | None = None, gamma: float = 1 / 3, thresholding: bool = False, prediction_type: str = "flow_prediction", solver_order: int = 2, predict_x0: bool = True, solver_type: str = "bh2", lower_order_final: bool = True, disable_corrector: list[int] | None = None, use_flow_sigmas: bool = True, scheduler_type: str = "unipc", use_dynamic_shifting: bool = False, time_shift_type: str = "linear", **kwargs, ): if stage_range is None: # Evenly divide [0, 1] into 3 stages for pyramid SR stage_range = [0, 1 / 3, 2 / 3, 1] if disable_corrector is None: disable_corrector = [] self.config = HeliosSchedulerConfig( num_train_timesteps=num_train_timesteps, shift=shift, stages=stages, stage_range=stage_range, gamma=gamma, thresholding=thresholding, prediction_type=prediction_type, solver_order=solver_order, predict_x0=predict_x0, solver_type=solver_type, lower_order_final=lower_order_final, disable_corrector=disable_corrector, use_flow_sigmas=use_flow_sigmas, scheduler_type=scheduler_type, use_dynamic_shifting=use_dynamic_shifting, time_shift_type=time_shift_type, ) self.timestep_ratios = {} self.timesteps_per_stage = {} self.sigmas_per_stage = {} self.start_sigmas = {} self.end_sigmas = {} self.ori_start_sigmas = {} self.init_sigmas_for_each_stage() self.sigma_min = self.sigmas[-1].item() self.sigma_max = self.sigmas[0].item() self.gamma = gamma if solver_type not in ["bh1", "bh2"]: raise NotImplementedError(f"{solver_type} is not implemented") self.predict_x0 = predict_x0 self.model_outputs = [None] * solver_order self.timestep_list = [None] * solver_order self.lower_order_nums = 0 self.disable_corrector = disable_corrector self.solver_p = None self.last_sample = None self._step_index = None self._begin_index = None self.num_inference_steps = None def init_sigmas(self): num_train_timesteps = self.config.num_train_timesteps shift = self.config.shift alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps + 1) sigmas = 1.0 - alphas sigmas = np.flip(shift * sigmas / (1 + (shift - 1) * sigmas))[:-1].copy() sigmas = torch.from_numpy(sigmas) timesteps = (sigmas * num_train_timesteps).clone() self._step_index = None self._begin_index = None self.timesteps = timesteps self.sigmas = sigmas.to("cpu") def init_sigmas_for_each_stage(self): self.init_sigmas() stage_distance = [] stages = self.config.stages training_steps = self.config.num_train_timesteps stage_range = self.config.stage_range for i_s in range(stages): start_indice = int(stage_range[i_s] * training_steps) start_indice = max(start_indice, 0) end_indice = int(stage_range[i_s + 1] * training_steps) end_indice = min(end_indice, training_steps) start_sigma = self.sigmas[start_indice].item() end_sigma = ( self.sigmas[end_indice].item() if end_indice < training_steps else 0.0 ) self.ori_start_sigmas[i_s] = start_sigma if i_s != 0: ori_sigma = 1 - start_sigma gamma = self.config.gamma corrected_sigma = ( 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma) ) * ori_sigma start_sigma = 1 - corrected_sigma stage_distance.append(start_sigma - end_sigma) self.start_sigmas[i_s] = start_sigma self.end_sigmas[i_s] = end_sigma tot_distance = sum(stage_distance) for i_s in range(stages): if i_s == 0: start_ratio = 0.0 else: start_ratio = sum(stage_distance[:i_s]) / tot_distance if i_s == stages - 1: # Use value just below 1.0 to avoid out-of-bounds indexing end_ratio = 1.0 - 1e-16 else: end_ratio = sum(stage_distance[: i_s + 1]) / tot_distance self.timestep_ratios[i_s] = (start_ratio, end_ratio) for i_s in range(stages): timestep_ratio = self.timestep_ratios[i_s] # Clamp to max valid timestep (num_train_timesteps - 1) timestep_max = min( self.timesteps[int(timestep_ratio[0] * training_steps)], 999 ) timestep_min = self.timesteps[ min(int(timestep_ratio[1] * training_steps), training_steps - 1) ] timesteps = np.linspace(timestep_max, timestep_min, training_steps + 1) self.timesteps_per_stage[i_s] = ( timesteps[:-1] if isinstance(timesteps, torch.Tensor) else torch.from_numpy(timesteps[:-1]) ) # Sigma range [0.999, 0]: start just below 1.0 to avoid singularity stage_sigmas = np.linspace(0.999, 0, training_steps + 1) self.sigmas_per_stage[i_s] = torch.from_numpy(stage_sigmas[:-1]) @property def step_index(self): return self._step_index @property def begin_index(self): return self._begin_index def set_begin_index(self, begin_index: int = 0): self._begin_index = begin_index def time_shift(self, mu, sigma, t): if self.config.time_shift_type == "exponential": return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) elif self.config.time_shift_type == "linear": return mu / (mu + (1 / t - 1) ** sigma) def set_timesteps( self, num_inference_steps: int, stage_index: int | None = None, device: str | torch.device = None, sigmas=None, mu=None, is_amplify_first_chunk: bool = False, ): if self.config.scheduler_type == "dmd": if is_amplify_first_chunk: num_inference_steps = num_inference_steps * 2 + 1 else: num_inference_steps = num_inference_steps + 1 self.num_inference_steps = num_inference_steps self.init_sigmas() if self.config.stages == 1: if sigmas is None: sigmas = np.linspace( 1, 1 / self.config.num_train_timesteps, num_inference_steps + 1, )[:-1].astype(np.float32) if self.config.shift != 1.0: assert not self.config.use_dynamic_shifting sigmas = self.time_shift(self.config.shift, 1.0, sigmas) timesteps = (sigmas * self.config.num_train_timesteps).copy() sigmas = torch.from_numpy(sigmas) else: stage_timesteps = self.timesteps_per_stage[stage_index] timesteps = np.linspace( stage_timesteps[0].item(), stage_timesteps[-1].item(), num_inference_steps, ) stage_sigmas = self.sigmas_per_stage[stage_index] ratios = np.linspace( stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps ) sigmas = torch.from_numpy(ratios) self.timesteps = torch.from_numpy(timesteps).to(device=device) self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device) self._step_index = None self.reset_scheduler_history() if self.config.scheduler_type == "dmd": self.timesteps = self.timesteps[:-1] self.sigmas = torch.cat([self.sigmas[:-2], self.sigmas[-1:]]) if self.config.use_dynamic_shifting: assert self.config.shift == 1.0 self.sigmas = self.time_shift(mu, 1.0, self.sigmas) if self.config.stages == 1: self.timesteps = self.sigmas[:-1] * self.config.num_train_timesteps else: self.timesteps = self.timesteps_per_stage[ stage_index ].min() + self.sigmas[:-1] * ( self.timesteps_per_stage[stage_index].max() - self.timesteps_per_stage[stage_index].min() ) # ---------------------------------- Euler ---------------------------------- def index_for_timestep(self, timestep, schedule_timesteps=None): if schedule_timesteps is None: schedule_timesteps = self.timesteps indices = (schedule_timesteps == timestep).nonzero() pos = 1 if len(indices) > 1 else 0 return indices[pos].item() def _init_step_index(self, timestep): if self.begin_index is None: if isinstance(timestep, torch.Tensor): timestep = timestep.to(self.timesteps.device) self._step_index = self.index_for_timestep(timestep) else: self._step_index = self._begin_index def step_euler( self, model_output: torch.FloatTensor, timestep=None, sample: torch.FloatTensor = None, return_dict: bool = True, **kwargs, ) -> HeliosSchedulerOutput | tuple: if self.step_index is None: self._step_index = 0 sample = sample.to(torch.float32) sigma = self.sigmas[self.step_index] sigma_next = self.sigmas[self.step_index + 1] prev_sample = sample + (sigma_next - sigma) * model_output prev_sample = prev_sample.to(model_output.dtype) self._step_index += 1 if not return_dict: return (prev_sample,) return HeliosSchedulerOutput(prev_sample=prev_sample) # ---------------------------------- UniPC ---------------------------------- def _sigma_to_alpha_sigma_t(self, sigma): if self.config.use_flow_sigmas: alpha_t = 1 - sigma sigma_t = torch.clamp(sigma, min=1e-8) else: alpha_t = 1 / ((sigma**2 + 1) ** 0.5) sigma_t = sigma * alpha_t return alpha_t, sigma_t def convert_model_output(self, model_output, sample=None, sigma=None, **kwargs): flag = False if sigma is None: flag = True sigma = self.sigmas[self.step_index] alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) if self.predict_x0: if self.config.prediction_type == "flow_prediction": if flag: sigma_t = self.sigmas[self.step_index] else: sigma_t = sigma x0_pred = sample - sigma_t * model_output elif self.config.prediction_type == "epsilon": x0_pred = (sample - sigma_t * model_output) / alpha_t elif self.config.prediction_type == "sample": x0_pred = model_output elif self.config.prediction_type == "v_prediction": x0_pred = alpha_t * sample - sigma_t * model_output else: raise ValueError( f"prediction_type {self.config.prediction_type} not supported" ) return x0_pred else: if self.config.prediction_type == "epsilon": return model_output elif self.config.prediction_type == "sample": return (sample - alpha_t * model_output) / sigma_t elif self.config.prediction_type == "v_prediction": return alpha_t * model_output + sigma_t * sample else: raise ValueError( f"prediction_type {self.config.prediction_type} not supported" ) def multistep_uni_p_bh_update( self, model_output, sample=None, order=None, sigma=None, sigma_next=None ): model_output_list = self.model_outputs m0 = model_output_list[-1] x = sample if sigma_next is None and sigma is None: sigma_t, sigma_s0 = ( self.sigmas[self.step_index + 1], self.sigmas[self.step_index], ) else: sigma_t, sigma_s0 = sigma_next, sigma alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) lambda_t = torch.log(alpha_t) - torch.log(sigma_t) lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) h = lambda_t - lambda_s0 device = sample.device rks = [] D1s = [] for i in range(1, order): si = self.step_index - i mi = model_output_list[-(i + 1)] alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) lambda_si = torch.log(alpha_si) - torch.log(sigma_si) rk = (lambda_si - lambda_s0) / h rks.append(rk) D1s.append((mi - m0) / rk) rks.append(1.0) rks = torch.tensor(rks, device=device) R = [] b = [] hh = -h if self.predict_x0 else h h_phi_1 = torch.expm1(hh) h_phi_k = h_phi_1 / hh - 1 factorial_i = 1 if self.config.solver_type == "bh1": B_h = hh elif self.config.solver_type == "bh2": B_h = torch.expm1(hh) else: raise NotImplementedError() for i in range(1, order + 1): R.append(torch.pow(rks, i - 1)) b.append(h_phi_k * factorial_i / B_h) factorial_i *= i + 1 h_phi_k = h_phi_k / hh - 1 / factorial_i R = torch.stack(R) b = torch.tensor(b, device=device) if len(D1s) > 0: D1s = torch.stack(D1s, dim=1) if order == 2: rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device) else: rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype) else: D1s = None if self.predict_x0: x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 pred_res = ( torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0 ) x_t = x_t_ - alpha_t * B_h * pred_res else: x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 pred_res = ( torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0 ) x_t = x_t_ - sigma_t * B_h * pred_res return x_t.to(x.dtype) def multistep_uni_c_bh_update( self, this_model_output, last_sample=None, this_sample=None, order=None, sigma_before=None, sigma=None, ): model_output_list = self.model_outputs m0 = model_output_list[-1] x = last_sample model_t = this_model_output if sigma_before is None and sigma is None: sigma_t, sigma_s0 = ( self.sigmas[self.step_index], self.sigmas[self.step_index - 1], ) else: sigma_t, sigma_s0 = sigma, sigma_before alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) lambda_t = torch.log(alpha_t) - torch.log(sigma_t) lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) h = lambda_t - lambda_s0 device = this_sample.device rks = [] D1s = [] for i in range(1, order): si = self.step_index - (i + 1) mi = model_output_list[-(i + 1)] alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) lambda_si = torch.log(alpha_si) - torch.log(sigma_si) rk = (lambda_si - lambda_s0) / h rks.append(rk) D1s.append((mi - m0) / rk) rks.append(1.0) rks = torch.tensor(rks, device=device) R = [] b = [] hh = -h if self.predict_x0 else h h_phi_1 = torch.expm1(hh) h_phi_k = h_phi_1 / hh - 1 factorial_i = 1 if self.config.solver_type == "bh1": B_h = hh elif self.config.solver_type == "bh2": B_h = torch.expm1(hh) else: raise NotImplementedError() for i in range(1, order + 1): R.append(torch.pow(rks, i - 1)) b.append(h_phi_k * factorial_i / B_h) factorial_i *= i + 1 h_phi_k = h_phi_k / hh - 1 / factorial_i R = torch.stack(R) b = torch.tensor(b, device=device) if len(D1s) > 0: D1s = torch.stack(D1s, dim=1) else: D1s = None if order == 1: rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device) else: rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype) if self.predict_x0: x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 corr_res = ( torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) if D1s is not None else 0 ) D1_t = model_t - m0 x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) else: x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 corr_res = ( torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) if D1s is not None else 0 ) D1_t = model_t - m0 x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) return x_t.to(x.dtype) def step_unipc( self, model_output, timestep=None, sample=None, return_dict: bool = True, **kwargs, ) -> HeliosSchedulerOutput | tuple: if self.num_inference_steps is None: raise ValueError( "Number of inference steps is 'None', run 'set_timesteps' first" ) if self.step_index is None: self._step_index = 0 use_corrector = ( self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None ) model_output_convert = self.convert_model_output(model_output, sample=sample) if use_corrector: sample = self.multistep_uni_c_bh_update( this_model_output=model_output_convert, last_sample=self.last_sample, this_sample=sample, order=self.this_order, ) for i in range(self.config.solver_order - 1): self.model_outputs[i] = self.model_outputs[i + 1] self.timestep_list[i] = self.timestep_list[i + 1] self.model_outputs[-1] = model_output_convert self.timestep_list[-1] = timestep if self.config.lower_order_final: this_order = min( self.config.solver_order, len(self.timesteps) - self.step_index ) else: this_order = self.config.solver_order self.this_order = min(this_order, self.lower_order_nums + 1) assert self.this_order > 0 self.last_sample = sample prev_sample = self.multistep_uni_p_bh_update( model_output=model_output, sample=sample, order=self.this_order, ) if self.lower_order_nums < self.config.solver_order: self.lower_order_nums += 1 self._step_index += 1 if not return_dict: return (prev_sample,) return HeliosSchedulerOutput(prev_sample=prev_sample) # ---------------------------------- DMD ---------------------------------- def add_noise(self, original_samples, noise, timestep, sigmas, timesteps): sigmas = sigmas.to(noise.device) timesteps = timesteps.to(noise.device) timestep_id = torch.argmin( (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1 ) sigma = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1) sample = (1 - sigma) * original_samples + sigma * noise return sample.type_as(noise) def convert_flow_pred_to_x0(self, flow_pred, xt, timestep, sigmas, timesteps): original_dtype = flow_pred.dtype device = flow_pred.device flow_pred, xt, sigmas, timesteps = ( x.double().to(device) for x in (flow_pred, xt, sigmas, timesteps) ) timestep_id = torch.argmin( (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1 ) sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1) x0_pred = xt - sigma_t * flow_pred return x0_pred.to(original_dtype) def step_dmd( self, model_output: torch.FloatTensor, timestep=None, sample: torch.FloatTensor = None, return_dict: bool = True, cur_sampling_step: int = 0, dmd_noisy_tensor: torch.FloatTensor | None = None, dmd_sigmas: torch.FloatTensor | None = None, dmd_timesteps: torch.FloatTensor | None = None, all_timesteps: torch.FloatTensor | None = None, **kwargs, ) -> HeliosSchedulerOutput | tuple: pred_image_or_video = self.convert_flow_pred_to_x0( flow_pred=model_output, xt=sample, timestep=torch.full( (model_output.shape[0],), timestep, dtype=torch.long, device=model_output.device, ), sigmas=dmd_sigmas, timesteps=dmd_timesteps, ) if cur_sampling_step < len(all_timesteps) - 1: prev_sample = self.add_noise( pred_image_or_video, dmd_noisy_tensor, torch.full( (model_output.shape[0],), all_timesteps[cur_sampling_step + 1], dtype=torch.long, device=model_output.device, ), sigmas=dmd_sigmas, timesteps=dmd_timesteps, ) else: prev_sample = pred_image_or_video if not return_dict: return (prev_sample,) return HeliosSchedulerOutput(prev_sample=prev_sample) # ---------------------------------- Main step ---------------------------------- def step( self, model_output, timestep=None, sample=None, return_dict: bool = True, **kwargs, ) -> HeliosSchedulerOutput | tuple: if self.config.scheduler_type == "euler": return self.step_euler( model_output=model_output, timestep=timestep, sample=sample, return_dict=return_dict, ) elif self.config.scheduler_type == "unipc": return self.step_unipc( model_output=model_output, timestep=timestep, sample=sample, return_dict=return_dict, ) elif self.config.scheduler_type == "dmd": return self.step_dmd( model_output=model_output, timestep=timestep, sample=sample, return_dict=return_dict, **kwargs, ) else: raise NotImplementedError( f"Scheduler type '{self.config.scheduler_type}' not implemented" ) def reset_scheduler_history(self): self.model_outputs = [None] * self.config.solver_order self.timestep_list = [None] * self.config.solver_order self.lower_order_nums = 0 self.disable_corrector = self.config.disable_corrector self.solver_p = None self.last_sample = None self._step_index = None self._begin_index = None def set_shift(self, shift: float): """Update the shift parameter (called by SchedulerLoader after loading).""" self.config.shift = shift self.shift = shift def __len__(self): return self.config.num_train_timesteps # Alias for Helios-Distilled which uses "HeliosDMDScheduler" in scheduler_config.json HeliosDMDScheduler = HeliosScheduler EntryClass = [HeliosScheduler, "HeliosDMDScheduler"] ================================================ FILE: python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_self_forcing_flow_match.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import torch from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.schedulers.scheduling_utils import SchedulerMixin from diffusers.utils import BaseOutput from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class SelfForcingFlowMatchSchedulerOutput(BaseOutput): """ Output class for the scheduler's `step` function output. Args: prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the denoising loop. """ prev_sample: torch.FloatTensor class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): config_name = "scheduler_config.json" order = 1 @register_to_config def __init__( self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False, *args, **kwargs, ): self.num_train_timesteps = num_train_timesteps self.shift = shift self.sigma_max = sigma_max self.sigma_min = sigma_min self.inverse_timesteps = inverse_timesteps self.extra_one_step = extra_one_step self.reverse_sigmas = reverse_sigmas self.set_timesteps(num_inference_steps) def set_timesteps( self, num_inference_steps=100, denoising_strength=1.0, return_dict=False, **kwargs, ): sigma_start = ( self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength ) if self.extra_one_step: self.sigmas = torch.linspace( sigma_start, self.sigma_min, num_inference_steps + 1 )[:-1] else: self.sigmas = torch.linspace( sigma_start, self.sigma_min, num_inference_steps ) if self.inverse_timesteps: self.sigmas = torch.flip(self.sigmas, dims=[0]) self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) if self.reverse_sigmas: self.sigmas = 1 - self.sigmas self.timesteps = self.sigmas * self.num_train_timesteps def step( self, model_output: torch.FloatTensor, timestep: torch.FloatTensor, sample: torch.FloatTensor, to_final=False, return_dict=False, **kwargs, ): if timestep.ndim == 2: timestep = timestep.flatten(0, 1) elif timestep.ndim == 0: # handles the case where timestep is a scalar, this occurs when we # use this scheduler for ODE trajectory timestep = timestep.unsqueeze(0) self.sigmas = self.sigmas.to(model_output.device) self.timesteps = self.timesteps.to(model_output.device) timestep = timestep.to(model_output.device) timestep_id = torch.argmin( (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1 ) sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) if to_final or (timestep_id + 1 >= len(self.timesteps)).any(): sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0 else: sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1) prev_sample = sample + model_output * (sigma_ - sigma) if isinstance(prev_sample, torch.Tensor | float) and not return_dict: return (prev_sample,) return SelfForcingFlowMatchSchedulerOutput(prev_sample=prev_sample) def add_noise(self, original_samples, noise, timestep): """ Diffusion forward corruption process. Input: - clean_latent: the clean latent with shape [B*T, C, H, W] - noise: the noise with shape [B*T, C, H, W] - timestep: the timestep with shape [B*T] Output: the corrupted latent with shape [B*T, C, H, W] """ if timestep.ndim == 2: timestep = timestep.flatten(0, 1) self.sigmas = self.sigmas.to(noise.device) self.timesteps = self.timesteps.to(noise.device) timestep_id = torch.argmin( (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1 ) sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) sample = (1 - sigma) * original_samples + sigma * noise return sample.type_as(noise) def scale_model_input( self, sample: torch.Tensor, timestep: int | None = None ) -> torch.Tensor: return sample def set_shift(self, shift: float) -> None: self.shift = shift EntryClass = SelfForcingFlowMatchScheduler ================================================ FILE: python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # Copyright 2025 TSAIL Team and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # DISCLAIMER: check https://huggingface.co/papers/2302.04867 and https://github.com/wl-zhao/UniPC for more info # The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py # ============================================================================== # # Modified from diffusers==0.35.0.dev0 # # ============================================================================== import math import numpy as np import torch from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.schedulers.scheduling_utils import ( KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput, ) from diffusers.utils import deprecate, is_scipy_available from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler if is_scipy_available(): import scipy.stats # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar def betas_for_alpha_bar( num_diffusion_timesteps, max_beta=0.999, alpha_transform_type="cosine", ): """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of (1-beta) over time from t = [0,1]. Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up to that part of the diffusion process. Args: num_diffusion_timesteps (`int`): the number of betas to produce. max_beta (`float`): the maximum beta to use; use values lower than 1 to prevent singularities. alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. Choose from `cosine` or `exp` Returns: betas (`np.ndarray`): the betas used by the scheduler to step the model outputs """ if alpha_transform_type == "cosine": def alpha_bar_fn(t): return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 elif alpha_transform_type == "exp": def alpha_bar_fn(t): return math.exp(t * -12.0) else: raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}") betas = [] for i in range(num_diffusion_timesteps): t1 = i / num_diffusion_timesteps t2 = (i + 1) / num_diffusion_timesteps betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) return torch.tensor(betas, dtype=torch.float32) # Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr def rescale_zero_terminal_snr(betas): """ Rescales betas to have zero terminal SNR Based on https://huggingface.co/papers/2305.08891 (Algorithm 1) Args: betas (`torch.Tensor`): the betas that the scheduler is being initialized with. Returns: `torch.Tensor`: rescaled betas with zero terminal SNR """ # Convert betas to alphas_bar_sqrt alphas = 1.0 - betas alphas_cumprod = torch.cumprod(alphas, dim=0) alphas_bar_sqrt = alphas_cumprod.sqrt() # Store old values. alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() # Shift so the last timestep is zero. alphas_bar_sqrt -= alphas_bar_sqrt_T # Scale so the first timestep is back to the old value. alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) # Convert alphas_bar_sqrt to betas alphas_bar = alphas_bar_sqrt**2 # Revert sqrt alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod alphas = torch.cat([alphas_bar[0:1], alphas]) betas = 1 - alphas return betas class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """ `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models. This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic methods the library implements for all schedulers such as loading and saving. Args: num_train_timesteps (`int`, defaults to 1000): The number of diffusion steps to train the model. beta_start (`float`, defaults to 0.0001): The starting `beta` value of inference. beta_end (`float`, defaults to 0.02): The final `beta` value. beta_schedule (`str`, defaults to `"linear"`): The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from `linear`, `scaled_linear`, or `squaredcos_cap_v2`. trained_betas (`np.ndarray`, *optional*): Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. solver_order (`int`, default `2`): The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1` due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for unconditional sampling. prediction_type (`str`, defaults to `epsilon`, *optional*): Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process), `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen Video](https://imagen.research.google/video/paper.pdf) paper). thresholding (`bool`, defaults to `False`): Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such as Stable Diffusion. dynamic_thresholding_ratio (`float`, defaults to 0.995): The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. sample_max_value (`float`, defaults to 1.0): The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`. predict_x0 (`bool`, defaults to `True`): Whether to use the updating algorithm on the predicted x0. solver_type (`str`, default `bh2`): Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2` otherwise. lower_order_final (`bool`, default `True`): Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10. disable_corrector (`list`, default `[]`): Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)` and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is usually disabled during the first few steps. solver_p (`SchedulerMixin`, default `None`): Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`. use_karras_sigmas (`bool`, *optional*, defaults to `False`): Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, the sigmas are determined according to a sequence of noise levels {σi}. use_exponential_sigmas (`bool`, *optional*, defaults to `False`): Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. use_beta_sigmas (`bool`, *optional*, defaults to `False`): Whether to use beta sigmas for step sizes in the noise schedule during the sampling process. Refer to [Beta Sampling is All You Need](https://huggingface.co/papers/2407.12173) for more information. use_flow_sigmas (`bool`, *optional*, defaults to `False`): Whether to use flow sigmas for step sizes in the noise schedule during the sampling process. timestep_spacing (`str`, defaults to `"linspace"`): The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. steps_offset (`int`, defaults to 0): An offset added to the inference steps, as required by some model families. final_sigmas_type (`str`, defaults to `"zero"`): The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0. rescale_betas_zero_snr (`bool`, defaults to `False`): Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and dark samples instead of limiting it to samples with medium brightness. Loosely related to [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506). """ _compatibles = [e.name for e in KarrasDiffusionSchedulers] order = 1 @register_to_config def __init__( self, num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, beta_schedule: str = "linear", trained_betas: np.ndarray | list[float] | None = None, solver_order: int = 2, prediction_type: str = "epsilon", thresholding: bool = False, dynamic_thresholding_ratio: float = 0.995, sample_max_value: float = 1.0, predict_x0: bool = True, solver_type: str = "bh2", lower_order_final: bool = True, disable_corrector: list[int] = [], solver_p: SchedulerMixin = None, use_karras_sigmas: bool | None = False, use_exponential_sigmas: bool | None = False, use_beta_sigmas: bool | None = False, use_flow_sigmas: bool | None = False, flow_shift: float | None = 1.0, timestep_spacing: str = "linspace", steps_offset: int = 0, final_sigmas_type: str | None = "zero", # "zero", "sigma_min" rescale_betas_zero_snr: bool = False, use_dynamic_shifting: bool = False, time_shift_type: str = "exponential", ): if self.config.use_beta_sigmas and not is_scipy_available(): raise ImportError( "Make sure to install scipy if you want to use beta sigmas." ) if ( sum( [ self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas, ] ) > 1 ): raise ValueError( "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." ) if trained_betas is not None: self.betas = torch.tensor(trained_betas, dtype=torch.float32) elif beta_schedule == "linear": self.betas = torch.linspace( beta_start, beta_end, num_train_timesteps, dtype=torch.float32 ) elif beta_schedule == "scaled_linear": # this schedule is very specific to the latent diffusion model. self.betas = ( torch.linspace( beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32, ) ** 2 ) elif beta_schedule == "squaredcos_cap_v2": # Glide cosine schedule self.betas = betas_for_alpha_bar(num_train_timesteps) else: raise NotImplementedError( f"{beta_schedule} is not implemented for {self.__class__}" ) if rescale_betas_zero_snr: self.betas = rescale_zero_terminal_snr(self.betas) self.alphas = 1.0 - self.betas self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) if rescale_betas_zero_snr: # Close to 0 without being 0 so first sigma is not inf # FP16 smallest positive subnormal works well here self.alphas_cumprod[-1] = 2**-24 # Currently we only support VP-type noise schedule self.alpha_t = torch.sqrt(self.alphas_cumprod) self.sigma_t = torch.sqrt(1 - self.alphas_cumprod) self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t) self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5 # standard deviation of the initial noise distribution self.init_noise_sigma = 1.0 if solver_type not in ["bh1", "bh2"]: if solver_type in ["midpoint", "heun", "logrho"]: self.register_to_config(solver_type="bh2") else: raise NotImplementedError( f"{solver_type} is not implemented for {self.__class__}" ) self.predict_x0 = predict_x0 # setable values self.num_inference_steps = None timesteps = np.linspace( 0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32 )[::-1].copy() self.timesteps = torch.from_numpy(timesteps) self.num_train_timesteps = num_train_timesteps self.model_outputs = [None] * solver_order self.timestep_list = [None] * solver_order self.lower_order_nums = 0 self.disable_corrector = disable_corrector self.solver_p = solver_p self.last_sample = None self._step_index = None self._begin_index = None self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication BaseScheduler.__init__(self) @property def step_index(self): """ The index counter for current timestep. It will increase 1 after each scheduler step. """ return self._step_index @property def begin_index(self): """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. """ return self._begin_index def set_shift(self, shift: float) -> None: self.config.flow_shift = shift # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index def set_begin_index(self, begin_index: int = 0): """ Sets the begin index for the scheduler. This function should be run from pipeline before the inference. Args: begin_index (`int`): The begin index for the scheduler. """ self._begin_index = begin_index def set_timesteps( self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (`int`): The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. """ # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891 if mu is not None: assert ( self.config.use_dynamic_shifting and self.config.time_shift_type == "exponential" ) self.config.flow_shift = np.exp(mu) if self.config.timestep_spacing == "linspace": timesteps = ( np.linspace( 0, self.config.num_train_timesteps - 1, num_inference_steps + 1 ) .round()[::-1][:-1] .copy() .astype(np.int64) ) elif self.config.timestep_spacing == "leading": step_ratio = self.config.num_train_timesteps // (num_inference_steps + 1) # creates integer timesteps by multiplying by ratio # casting to int to avoid issues when num_inference_step is power of 3 timesteps = ( (np.arange(0, num_inference_steps + 1) * step_ratio) .round()[::-1][:-1] .copy() .astype(np.int64) ) timesteps += self.config.steps_offset elif self.config.timestep_spacing == "trailing": step_ratio = self.config.num_train_timesteps / num_inference_steps # creates integer timesteps by multiplying by ratio # casting to int to avoid issues when num_inference_step is power of 3 timesteps = ( np.arange(self.config.num_train_timesteps, 0, -step_ratio) .round() .copy() .astype(np.int64) ) timesteps -= 1 else: raise ValueError( f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'." ) sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) if self.config.use_karras_sigmas: log_sigmas = np.log(sigmas) sigmas = np.flip(sigmas).copy() sigmas = self._convert_to_karras( in_sigmas=sigmas, num_inference_steps=num_inference_steps ) timesteps = np.array( [self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas] ).round() if self.config.final_sigmas_type == "sigma_min": sigma_last = sigmas[-1] elif self.config.final_sigmas_type == "zero": sigma_last = 0 else: raise ValueError( f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" ) sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) elif self.config.use_exponential_sigmas: log_sigmas = np.log(sigmas) sigmas = np.flip(sigmas).copy() sigmas = self._convert_to_exponential( in_sigmas=sigmas, num_inference_steps=num_inference_steps ) timesteps = np.array( [self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas] ) if self.config.final_sigmas_type == "sigma_min": sigma_last = sigmas[-1] elif self.config.final_sigmas_type == "zero": sigma_last = 0 else: raise ValueError( f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" ) sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) elif self.config.use_beta_sigmas: log_sigmas = np.log(sigmas) sigmas = np.flip(sigmas).copy() sigmas = self._convert_to_beta( in_sigmas=sigmas, num_inference_steps=num_inference_steps ) timesteps = np.array( [self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas] ) if self.config.final_sigmas_type == "sigma_min": sigma_last = sigmas[-1] elif self.config.final_sigmas_type == "zero": sigma_last = 0 else: raise ValueError( f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" ) sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) elif self.config.use_flow_sigmas: alphas = np.linspace( 1, 1 / self.config.num_train_timesteps, num_inference_steps + 1 ) sigmas = 1.0 - alphas sigmas = np.flip( self.config.flow_shift * sigmas / (1 + (self.config.flow_shift - 1) * sigmas) )[:-1].copy() timesteps = (sigmas * self.config.num_train_timesteps).copy() if self.config.final_sigmas_type == "sigma_min": sigma_last = sigmas[-1] elif self.config.final_sigmas_type == "zero": sigma_last = 0 else: raise ValueError( f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" ) sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) else: sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) if self.config.final_sigmas_type == "sigma_min": sigma_last = ( (1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0] ) ** 0.5 elif self.config.final_sigmas_type == "zero": sigma_last = 0 else: raise ValueError( f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" ) sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) self.sigmas = torch.from_numpy(sigmas) self.timesteps = torch.from_numpy(timesteps).to( device=device, dtype=torch.int64 ) self.num_inference_steps = len(timesteps) self.model_outputs = [ None, ] * self.config.solver_order self.lower_order_nums = 0 self.last_sample = None if self.solver_p: self.solver_p.set_timesteps(self.num_inference_steps, device=device) # add an index counter for schedulers that allow duplicated timesteps self._step_index = None self._begin_index = None self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: """ "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing pixels from saturation at each step. We find that dynamic thresholding results in significantly better photorealism as well as better image-text alignment, especially when using very large guidance weights." https://huggingface.co/papers/2205.11487 """ dtype = sample.dtype batch_size, channels, *remaining_dims = sample.shape if dtype not in (torch.float32, torch.float64): sample = ( sample.float() ) # upcast for quantile calculation, and clamp not implemented for cpu half # Flatten sample for doing quantile calculation along each image sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) abs_sample = sample.abs() # "a certain percentile absolute pixel value" s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) s = torch.clamp( s, min=1, max=self.config.sample_max_value ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 sample = ( torch.clamp(sample, -s, s) / s ) # "we threshold xt0 to the range [-s, s] and then divide by s" sample = sample.reshape(batch_size, channels, *remaining_dims) sample = sample.to(dtype) return sample # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t def _sigma_to_t(self, sigma, log_sigmas): # get log sigma log_sigma = np.log(np.maximum(sigma, 1e-10)) # get distribution dists = log_sigma - log_sigmas[:, np.newaxis] # get sigmas range low_idx = ( np.cumsum((dists >= 0), axis=0) .argmax(axis=0) .clip(max=log_sigmas.shape[0] - 2) ) high_idx = low_idx + 1 low = log_sigmas[low_idx] high = log_sigmas[high_idx] # interpolate sigmas w = (low - log_sigma) / (low - high) w = np.clip(w, 0, 1) # transform interpolation to time range t = (1 - w) * low_idx + w * high_idx t = t.reshape(sigma.shape) return t # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t def _sigma_to_alpha_sigma_t(self, sigma): if self.config.use_flow_sigmas: alpha_t = 1 - sigma sigma_t = sigma else: alpha_t = 1 / ((sigma**2 + 1) ** 0.5) sigma_t = sigma * alpha_t return alpha_t, sigma_t # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras def _convert_to_karras( self, in_sigmas: torch.Tensor, num_inference_steps ) -> torch.Tensor: """Constructs the noise schedule of Karras et al. (2022).""" # Hack to make sure that other schedulers which copy this function don't break # TODO: Add this logic to the other schedulers if hasattr(self.config, "sigma_min"): sigma_min = self.config.sigma_min else: sigma_min = None if hasattr(self.config, "sigma_max"): sigma_max = self.config.sigma_max else: sigma_max = None sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() rho = 7.0 # 7.0 is the value used in the paper ramp = np.linspace(0, 1, num_inference_steps) min_inv_rho = sigma_min ** (1 / rho) max_inv_rho = sigma_max ** (1 / rho) sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho return sigmas # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential def _convert_to_exponential( self, in_sigmas: torch.Tensor, num_inference_steps: int ) -> torch.Tensor: """Constructs an exponential noise schedule.""" # Hack to make sure that other schedulers which copy this function don't break # TODO: Add this logic to the other schedulers if hasattr(self.config, "sigma_min"): sigma_min = self.config.sigma_min else: sigma_min = None if hasattr(self.config, "sigma_max"): sigma_max = self.config.sigma_max else: sigma_max = None sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() sigmas = np.exp( np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps) ) return sigmas # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta def _convert_to_beta( self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6, ) -> torch.Tensor: """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" # Hack to make sure that other schedulers which copy this function don't break # TODO: Add this logic to the other schedulers if hasattr(self.config, "sigma_min"): sigma_min = self.config.sigma_min else: sigma_min = None if hasattr(self.config, "sigma_max"): sigma_max = self.config.sigma_max else: sigma_max = None sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() sigmas = np.array( [ sigma_min + (ppf * (sigma_max - sigma_min)) for ppf in [ scipy.stats.beta.ppf(timestep, alpha, beta) for timestep in 1 - np.linspace(0, 1, num_inference_steps) ] ] ) return sigmas def convert_model_output( self, model_output: torch.Tensor, *args, sample: torch.Tensor = None, **kwargs, ) -> torch.Tensor: r""" Convert the model output to the corresponding type the UniPC algorithm needs. Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model. timestep (`int`): The current discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. Returns: `torch.Tensor`: The converted model output. """ timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) if sample is None: if len(args) > 1: sample = args[1] else: raise ValueError("missing `sample` as a required keyword argument") if timestep is not None: deprecate( "timesteps", "1.0.0", "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", ) sigma = self.sigmas[self.step_index] alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) if self.predict_x0: if self.config.prediction_type == "epsilon": x0_pred = (sample - sigma_t * model_output) / alpha_t elif self.config.prediction_type == "sample": x0_pred = model_output elif self.config.prediction_type == "v_prediction": x0_pred = alpha_t * sample - sigma_t * model_output elif self.config.prediction_type == "flow_prediction": sigma_t = self.sigmas[self.step_index] x0_pred = sample - sigma_t * model_output else: raise ValueError( f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, " "`v_prediction`, or `flow_prediction` for the UniPCMultistepScheduler." ) if self.config.thresholding: x0_pred = self._threshold_sample(x0_pred) return x0_pred else: if self.config.prediction_type == "epsilon": return model_output elif self.config.prediction_type == "sample": epsilon = (sample - alpha_t * model_output) / sigma_t return epsilon elif self.config.prediction_type == "v_prediction": epsilon = alpha_t * model_output + sigma_t * sample return epsilon else: raise ValueError( f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" " `v_prediction` for the UniPCMultistepScheduler." ) def multistep_uni_p_bh_update( self, model_output: torch.Tensor, *args, sample: torch.Tensor = None, order: int = None, **kwargs, ) -> torch.Tensor: """ One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified. Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model at the current timestep. prev_timestep (`int`): The previous discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. order (`int`): The order of UniP at this timestep (corresponds to the *p* in UniPC-p). Returns: `torch.Tensor`: The sample tensor at the previous timestep. """ prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None) if sample is None: if len(args) > 1: sample = args[1] else: raise ValueError("missing `sample` as a required keyword argument") if order is None: if len(args) > 2: order = args[2] else: raise ValueError("missing `order` as a required keyword argument") if prev_timestep is not None: deprecate( "prev_timestep", "1.0.0", "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", ) model_output_list = self.model_outputs s0 = self.timestep_list[-1] m0 = model_output_list[-1] x = sample if self.solver_p: x_t = self.solver_p.step(model_output, s0, x).prev_sample return x_t sigma_t, sigma_s0 = ( self.sigmas[self.step_index + 1], self.sigmas[self.step_index], ) alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) lambda_t = torch.log(alpha_t) - torch.log(sigma_t) lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) h = lambda_t - lambda_s0 device = sample.device rks = [] D1s = [] for i in range(1, order): si = self.step_index - i mi = model_output_list[-(i + 1)] alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) lambda_si = torch.log(alpha_si) - torch.log(sigma_si) rk = (lambda_si - lambda_s0) / h rks.append(rk) D1s.append((mi - m0) / rk) rks.append(1.0) rks = torch.tensor(rks, device=device) R = [] b = [] hh = -h if self.predict_x0 else h h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 h_phi_k = h_phi_1 / hh - 1 factorial_i = 1 if self.config.solver_type == "bh1": B_h = hh elif self.config.solver_type == "bh2": B_h = torch.expm1(hh) else: raise NotImplementedError() for i in range(1, order + 1): R.append(torch.pow(rks, i - 1)) b.append(h_phi_k * factorial_i / B_h) factorial_i *= i + 1 h_phi_k = h_phi_k / hh - 1 / factorial_i R = torch.stack(R) b = torch.tensor(b, device=device) if len(D1s) > 0: D1s = torch.stack(D1s, dim=1) # (B, K) # for order 2, we use a simplified version if order == 2: rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device) else: rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype) else: D1s = None if self.predict_x0: x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 if D1s is not None: pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) else: pred_res = 0 x_t = x_t_ - alpha_t * B_h * pred_res else: x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 if D1s is not None: pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) else: pred_res = 0 x_t = x_t_ - sigma_t * B_h * pred_res x_t = x_t.to(x.dtype) return x_t def multistep_uni_c_bh_update( self, this_model_output: torch.Tensor, *args, last_sample: torch.Tensor = None, this_sample: torch.Tensor = None, order: int = None, **kwargs, ) -> torch.Tensor: """ One step for the UniC (B(h) version). Args: this_model_output (`torch.Tensor`): The model outputs at `x_t`. this_timestep (`int`): The current timestep `t`. last_sample (`torch.Tensor`): The generated sample before the last predictor `x_{t-1}`. this_sample (`torch.Tensor`): The generated sample after the last predictor `x_{t}`. order (`int`): The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`. Returns: `torch.Tensor`: The corrected sample tensor at the current timestep. """ this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None) if last_sample is None: if len(args) > 1: last_sample = args[1] else: raise ValueError("missing `last_sample` as a required keyword argument") if this_sample is None: if len(args) > 2: this_sample = args[2] else: raise ValueError("missing `this_sample` as a required keyword argument") if order is None: if len(args) > 3: order = args[3] else: raise ValueError("missing `order` as a required keyword argument") if this_timestep is not None: deprecate( "this_timestep", "1.0.0", "Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", ) model_output_list = self.model_outputs m0 = model_output_list[-1] x = last_sample x_t = this_sample model_t = this_model_output sigma_t, sigma_s0 = ( self.sigmas[self.step_index], self.sigmas[self.step_index - 1], ) alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) lambda_t = torch.log(alpha_t) - torch.log(sigma_t) lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) h = lambda_t - lambda_s0 device = this_sample.device rks = [] D1s = [] for i in range(1, order): si = self.step_index - (i + 1) mi = model_output_list[-(i + 1)] alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) lambda_si = torch.log(alpha_si) - torch.log(sigma_si) rk = (lambda_si - lambda_s0) / h rks.append(rk) D1s.append((mi - m0) / rk) rks.append(1.0) rks = torch.tensor(rks, device=device) R = [] b = [] hh = -h if self.predict_x0 else h h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 h_phi_k = h_phi_1 / hh - 1 factorial_i = 1 if self.config.solver_type == "bh1": B_h = hh elif self.config.solver_type == "bh2": B_h = torch.expm1(hh) else: raise NotImplementedError() for i in range(1, order + 1): R.append(torch.pow(rks, i - 1)) b.append(h_phi_k * factorial_i / B_h) factorial_i *= i + 1 h_phi_k = h_phi_k / hh - 1 / factorial_i R = torch.stack(R) b = torch.tensor(b, device=device) if len(D1s) > 0: D1s = torch.stack(D1s, dim=1) else: D1s = None # for order 1, we use a simplified version if order == 1: rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device) else: rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype) if self.predict_x0: x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 if D1s is not None: corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) else: corr_res = 0 D1_t = model_t - m0 x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) else: x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 if D1s is not None: corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) else: corr_res = 0 D1_t = model_t - m0 x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) x_t = x_t.to(x.dtype) return x_t # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep def index_for_timestep(self, timestep, schedule_timesteps=None): if schedule_timesteps is None: schedule_timesteps = self.timesteps index_candidates = (schedule_timesteps == timestep).nonzero() if len(index_candidates) == 0: step_index = len(self.timesteps) - 1 # The sigma index that is taken for the **very** first `step` # is always the second index (or the last index if there is only 1) # This way we can ensure we don't accidentally skip a sigma in # case we start in the middle of the denoising schedule (e.g. for image-to-image) elif len(index_candidates) > 1: step_index = index_candidates[1].item() else: step_index = index_candidates[0].item() return step_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index def _init_step_index(self, timestep): """ Initialize the step_index counter for the scheduler. """ if self.begin_index is None: if isinstance(timestep, torch.Tensor): timestep = timestep.to(self.timesteps.device) self._step_index = self.index_for_timestep(timestep) else: self._step_index = self._begin_index def step( self, model_output: torch.Tensor, timestep: int | torch.Tensor, sample: torch.Tensor, return_dict: bool = True, ) -> SchedulerOutput | tuple: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep UniPC. Args: model_output (`torch.Tensor`): The direct output from learned diffusion model. timestep (`int`): The current discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. return_dict (`bool`): Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`. Returns: [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`: If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a tuple is returned where the first element is the sample tensor. """ if self.num_inference_steps is None: raise ValueError( "Number of inference steps is 'None', you need to call 'set_timesteps' after creating the scheduler" ) if self.step_index is None: self._init_step_index(timestep) use_corrector = ( self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None ) model_output_convert = self.convert_model_output(model_output, sample=sample) if use_corrector: sample = self.multistep_uni_c_bh_update( this_model_output=model_output_convert, last_sample=self.last_sample, this_sample=sample, order=self.this_order, ) for i in range(self.config.solver_order - 1): self.model_outputs[i] = self.model_outputs[i + 1] self.timestep_list[i] = self.timestep_list[i + 1] self.model_outputs[-1] = model_output_convert self.timestep_list[-1] = timestep if self.config.lower_order_final: this_order = min( self.config.solver_order, len(self.timesteps) - self.step_index ) else: this_order = self.config.solver_order self.this_order = min( this_order, self.lower_order_nums + 1 ) # warmup for multistep assert self.this_order > 0 self.last_sample = sample prev_sample = self.multistep_uni_p_bh_update( model_output=model_output, # pass the original non-converted model output, in case solver-p is used sample=sample, order=self.this_order, ) if self.lower_order_nums < self.config.solver_order: self.lower_order_nums += 1 # upon completion increase step index by one self._step_index += 1 if not return_dict: return (prev_sample,) return SchedulerOutput(prev_sample=prev_sample) def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: """ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. Args: sample (`torch.Tensor`): The input sample. Returns: `torch.Tensor`: A scaled input sample. """ return sample # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise def add_noise( self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor, ) -> torch.Tensor: # Make sure sigmas and timesteps have the same device and dtype as original_samples sigmas = self.sigmas.to( device=original_samples.device, dtype=original_samples.dtype ) if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): # mps does not support float64 schedule_timesteps = self.timesteps.to( original_samples.device, dtype=torch.float32 ) timesteps = timesteps.to(original_samples.device, dtype=torch.float32) else: schedule_timesteps = self.timesteps.to(original_samples.device) timesteps = timesteps.to(original_samples.device) # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index if self.begin_index is None: step_indices = [ self.index_for_timestep(t, schedule_timesteps) for t in timesteps ] elif self.step_index is not None: # add_noise is called after first denoising step (for inpainting) step_indices = [self.step_index] * timesteps.shape[0] else: # add noise is called before first denoising step to create initial latent(img2img) step_indices = [self.begin_index] * timesteps.shape[0] sigma = sigmas[step_indices].flatten() while len(sigma.shape) < len(original_samples.shape): sigma = sigma.unsqueeze(-1) alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) noisy_samples = alpha_t * original_samples + sigma_t * noise return noisy_samples def __len__(self): return self.config.num_train_timesteps EntryClass = UniPCMultistepScheduler ================================================ FILE: python/sglang/multimodal_gen/runtime/models/utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/utils.py """Utils for model executor.""" from typing import Any import torch def set_weight_attrs( weight: torch.Tensor, weight_attrs: dict[str, Any] | None, ): """Set attributes on a weight tensor. This method is used to set attributes on a weight tensor. This method will not overwrite existing attributes. Args: weight: The weight tensor. weight_attrs: A dictionary of attributes to set on the weight tensor. """ if weight_attrs is None: return for key, value in weight_attrs.items(): assert not hasattr(weight, key), f"Overwriting existing tensor attribute: {key}" # NOTE(woosuk): During weight loading, we often do something like: # narrowed_tensor = param.data.narrow(0, offset, len) # narrowed_tensor.copy_(real_weight) # expecting narrowed_tensor and param.data to share the same storage. # However, on TPUs, narrowed_tensor will lazily propagate to the base # tensor, which is param.data, leading to the redundant memory usage. # This sometimes causes OOM errors during model loading. To avoid this, # we sync the param tensor after its weight loader is called. # TODO(woosuk): Remove this hack once we have a better solution. from sglang.multimodal_gen.runtime.platforms import current_platform if current_platform.is_tpu() and key == "weight_loader": value = _make_synced_weight_loader(value) setattr(weight, key, value) def _make_synced_weight_loader(original_weight_loader) -> Any: def _synced_weight_loader(param, *args, **kwargs): original_weight_loader(param, *args, **kwargs) torch._sync(param) return _synced_weight_loader def extract_layer_index(layer_name: str) -> int: """ Extract the layer index from the module name. Examples: - "encoder.layers.0" -> 0 - "encoder.layers.1.self_attn" -> 1 - "2.self_attn" -> 2 - "model.encoder.layers.0.sub.1" -> ValueError """ subnames = layer_name.split(".") int_vals: list[int] = [] for subname in subnames: try: int_vals.append(int(subname)) except ValueError: continue assert len(int_vals) == 1, ( f"layer name {layer_name} should" " only contain one integer" ) return int_vals[0] def modulate( x: torch.Tensor, shift: torch.Tensor | None = None, scale: torch.Tensor | None = None, ) -> torch.Tensor: """modulate by shift and scale""" if scale is None and shift is None: return x elif shift is None: return x * (1 + scale.unsqueeze(1)) # type: ignore[union-attr] elif scale is None: return x + shift.unsqueeze(1) # type: ignore[union-attr] else: return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze( 1 ) # type: ignore[union-attr] def pred_noise_to_pred_video( pred_noise: torch.Tensor, noise_input_latent: torch.Tensor, timestep: torch.Tensor, scheduler: Any, ) -> torch.Tensor: """ Convert predicted noise to clean latent. Args: pred_noise: the predicted noise with shape [B, C, H, W] where B is batch_size or batch_size * num_frames noise_input_latent: the noisy latent with shape [B, C, H, W], timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames] scheduler: the scheduler Returns: the predicted video with shape [B, C, H, W] """ # If timestep is [bs, num_frames] if timestep.ndim == 2: timestep = timestep.flatten(0, 1) assert timestep.numel() == noise_input_latent.shape[0] elif timestep.ndim == 1: # If timestep is [1] if timestep.shape[0] == 1: timestep = timestep.expand(noise_input_latent.shape[0]) else: assert timestep.numel() == noise_input_latent.shape[0] else: raise ValueError( f"[pred_noise_to_pred_video] Invalid timestep shape: {timestep.shape}" ) # timestep shape should be [B] dtype = pred_noise.dtype device = pred_noise.device pred_noise = pred_noise.double().to(device) noise_input_latent = noise_input_latent.double().to(device) sigmas = scheduler.sigmas.double().to(device) timesteps = scheduler.timesteps.double().to(device) timestep_id = torch.argmin( (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1 ) sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) pred_video = noise_input_latent - sigma_t * pred_noise return pred_video.to(dtype) ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vaes/autoencoder.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from typing import Dict, Optional, Tuple, Union import torch from diffusers.models.attention_processor import ( ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, Attention, AttentionProcessor, AttnAddedKVProcessor, AttnProcessor, FusedAttnProcessor2_0, ) from diffusers.models.autoencoders.vae import ( Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder, ) from diffusers.models.modeling_outputs import AutoencoderKLOutput from torch import nn from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig class AutoencoderKL(nn.Module): r""" A VAE model with KL loss for encoding images into latents and decoding latent representations into images. This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented for all models (such as downloading or saving). Parameters: in_channels (int, *optional*, defaults to 3): Number of channels in the input image. out_channels (int, *optional*, defaults to 3): Number of channels in the output. down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): Tuple of downsample block types. up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): Tuple of upsample block types. block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): Tuple of block output channels. act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space. sample_size (`int`, *optional*, defaults to `32`): Sample input size. scaling_factor (`float`, *optional*, defaults to 0.18215): The component-wise standard deviation of the trained latent space computed using the first batch of the training set. This is used to scale the latent space to have unit variance when training the diffusion model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1 / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image Synthesis with Latent Diffusion Models](https://huggingface.co/papers/2112.10752) paper. force_upcast (`bool`, *optional*, default to `True`): If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix mid_block_add_attention (`bool`, *optional*, default to `True`): If enabled, the mid_block of the Encoder and Decoder will have attention blocks. If set to false, the mid_block will only have resnet blocks """ _supports_gradient_checkpointing = True _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"] def __init__( self, config: FluxVAEConfig, ): super().__init__() self.config = config arch_config = config.arch_config in_channels = arch_config.in_channels out_channels = arch_config.out_channels down_block_types = arch_config.down_block_types up_block_types = arch_config.up_block_types block_out_channels = arch_config.block_out_channels layers_per_block = arch_config.layers_per_block act_fn = arch_config.act_fn latent_channels = arch_config.latent_channels norm_num_groups = arch_config.norm_num_groups sample_size = arch_config.sample_size use_quant_conv = arch_config.use_quant_conv use_post_quant_conv = arch_config.use_post_quant_conv mid_block_add_attention = arch_config.mid_block_add_attention # pass init params to Encoder self.encoder = Encoder( in_channels=in_channels, out_channels=latent_channels, down_block_types=down_block_types, block_out_channels=block_out_channels, layers_per_block=layers_per_block, act_fn=act_fn, norm_num_groups=norm_num_groups, double_z=True, mid_block_add_attention=mid_block_add_attention, ) # pass init params to Decoder self.decoder = Decoder( in_channels=latent_channels, out_channels=out_channels, up_block_types=up_block_types, block_out_channels=block_out_channels, layers_per_block=layers_per_block, norm_num_groups=norm_num_groups, act_fn=act_fn, mid_block_add_attention=mid_block_add_attention, ) self.quant_conv = ( nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) if use_quant_conv else None ) self.post_quant_conv = ( nn.Conv2d(latent_channels, latent_channels, 1) if use_post_quant_conv else None ) self.use_slicing = False self.use_tiling = False # only relevant if vae tiling is enabled self.tile_sample_min_size = sample_size sample_size = ( self.config.sample_size[0] if isinstance(self.config.sample_size, (list, tuple)) else self.config.sample_size ) self.tile_latent_min_size = int( sample_size / (2 ** (len(self.config.block_out_channels) - 1)) ) self.tile_overlap_factor = 0.25 def enable_tiling(self, use_tiling: bool = True): r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow processing larger images. """ self.use_tiling = use_tiling def disable_tiling(self): r""" Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing decoding in one step. """ self.enable_tiling(False) def enable_slicing(self): r""" Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. """ self.use_slicing = True def disable_slicing(self): r""" Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing decoding in one step. """ self.use_slicing = False @property # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors def attn_processors(self) -> Dict[str, AttentionProcessor]: r""" Returns: `dict` of attention processors: A dictionary containing all attention processors used in the model with indexed by its weight name. """ # set recursively processors = {} def fn_recursive_add_processors( name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor], ): if hasattr(module, "get_processor"): processors[f"{name}.processor"] = module.get_processor() for sub_name, child in module.named_children(): fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) return processors for name, module in self.named_children(): fn_recursive_add_processors(name, module, processors) return processors # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor def set_attn_processor( self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]] ): r""" Sets the attention processor to use to compute attention. Parameters: processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): The instantiated processor class or a dictionary of processor classes that will be set as the processor for **all** `Attention` layers. If `processor` is a dict, the key needs to define the path to the corresponding cross attention processor. This is strongly recommended when setting trainable attention processors. """ count = len(self.attn_processors.keys()) if isinstance(processor, dict) and len(processor) != count: raise ValueError( f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" f" number of attention layers: {count}. Please make sure to pass {count} processor classes." ) def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): module.set_processor(processor) else: module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) for name, module in self.named_children(): fn_recursive_attn_processor(name, module, processor) # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor def set_default_attn_processor(self): """ Disables custom attention processors and sets the default attention implementation. """ if all( proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values() ): processor = AttnAddedKVProcessor() elif all( proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values() ): processor = AttnProcessor() else: raise ValueError( f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) self.set_attn_processor(processor) def _encode(self, x: torch.Tensor) -> torch.Tensor: batch_size, num_channels, height, width = x.shape if self.use_tiling and ( width > self.tile_sample_min_size or height > self.tile_sample_min_size ): return self._tiled_encode(x) enc = self.encoder(x) if self.quant_conv is not None: enc = self.quant_conv(enc) return enc def encode( self, x: torch.Tensor, return_dict: bool = True ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: """ Encode a batch of images into latents. Args: x (`torch.Tensor`): Input batch of images. return_dict (`bool`, *optional*, defaults to `True`): Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. Returns: The latent representations of the encoded images. If `return_dict` is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. """ if self.use_slicing and x.shape[0] > 1: encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] h = torch.cat(encoded_slices) else: h = self._encode(x) posterior = DiagonalGaussianDistribution(h) if not return_dict: return (posterior,) return AutoencoderKLOutput(latent_dist=posterior) def _decode( self, z: torch.Tensor, return_dict: bool = True ) -> Union[DecoderOutput, torch.Tensor]: if self.use_tiling and ( z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size ): return self.tiled_decode(z, return_dict=return_dict) if self.post_quant_conv is not None: z = self.post_quant_conv(z) dec = self.decoder(z) if not return_dict: return (dec,) return DecoderOutput(sample=dec) def decode(self, z: torch.FloatTensor) -> Union[DecoderOutput, torch.FloatTensor]: """ Decode a batch of images. Args: z (`torch.Tensor`): Input batch of latent vectors. Returns: [`~models.vae.DecoderOutput`] or `tuple`: If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is returned. """ if self.use_slicing and z.shape[0] > 1: decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] decoded = torch.cat(decoded_slices) else: decoded = self._decode(z).sample return decoded def blend_v( self, a: torch.Tensor, b: torch.Tensor, blend_extent: int ) -> torch.Tensor: blend_extent = min(a.shape[2], b.shape[2], blend_extent) for y in range(blend_extent): b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[ :, :, y, : ] * (y / blend_extent) return b def blend_h( self, a: torch.Tensor, b: torch.Tensor, blend_extent: int ) -> torch.Tensor: blend_extent = min(a.shape[3], b.shape[3], blend_extent) for x in range(blend_extent): b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[ :, :, :, x ] * (x / blend_extent) return b def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor: r"""Encode a batch of images using a tiled encoder. When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the output, but they should be much less noticeable. Args: x (`torch.Tensor`): Input batch of images. Returns: `torch.Tensor`: The latent representation of the encoded videos. """ overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) row_limit = self.tile_latent_min_size - blend_extent # Split the image into 512x512 tiles and encode them separately. rows = [] for i in range(0, x.shape[2], overlap_size): row = [] for j in range(0, x.shape[3], overlap_size): tile = x[ :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size, ] tile = self.encoder(tile) if self.config.use_quant_conv: tile = self.quant_conv(tile) row.append(tile) rows.append(row) result_rows = [] for i, row in enumerate(rows): result_row = [] for j, tile in enumerate(row): # blend the above tile and the left tile # to the current tile and add the current tile to the result row if i > 0: tile = self.blend_v(rows[i - 1][j], tile, blend_extent) if j > 0: tile = self.blend_h(row[j - 1], tile, blend_extent) result_row.append(tile[:, :, :row_limit, :row_limit]) result_rows.append(torch.cat(result_row, dim=3)) enc = torch.cat(result_rows, dim=2) return enc def tiled_encode( self, x: torch.Tensor, return_dict: bool = True ) -> AutoencoderKLOutput: r"""Encode a batch of images using a tiled encoder. When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the output, but they should be much less noticeable. Args: x (`torch.Tensor`): Input batch of images. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. Returns: [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`: If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. """ deprecation_message = ( "The tiled_encode implementation supporting the `return_dict` parameter is deprecated. In the future, the " "implementation of this method will be replaced with that of `_tiled_encode` and you will no longer be able " "to pass `return_dict`. You will also have to create a `DiagonalGaussianDistribution()` from the returned value." ) # deprecate("tiled_encode", "1.0.0", deprecation_message, standard_warn=False) overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) row_limit = self.tile_latent_min_size - blend_extent # Split the image into 512x512 tiles and encode them separately. rows = [] for i in range(0, x.shape[2], overlap_size): row = [] for j in range(0, x.shape[3], overlap_size): tile = x[ :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size, ] tile = self.encoder(tile) if self.config.use_quant_conv: tile = self.quant_conv(tile) row.append(tile) rows.append(row) result_rows = [] for i, row in enumerate(rows): result_row = [] for j, tile in enumerate(row): # blend the above tile and the left tile # to the current tile and add the current tile to the result row if i > 0: tile = self.blend_v(rows[i - 1][j], tile, blend_extent) if j > 0: tile = self.blend_h(row[j - 1], tile, blend_extent) result_row.append(tile[:, :, :row_limit, :row_limit]) result_rows.append(torch.cat(result_row, dim=3)) moments = torch.cat(result_rows, dim=2) posterior = DiagonalGaussianDistribution(moments) if not return_dict: return (posterior,) return AutoencoderKLOutput(latent_dist=posterior) def tiled_decode( self, z: torch.Tensor, return_dict: bool = True ) -> Union[DecoderOutput, torch.Tensor]: r""" Decode a batch of images using a tiled decoder. Args: z (`torch.Tensor`): Input batch of latent vectors. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. Returns: [`~models.vae.DecoderOutput`] or `tuple`: If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is returned. """ overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) row_limit = self.tile_sample_min_size - blend_extent # Split z into overlapping 64x64 tiles and decode them separately. # The tiles have an overlap to avoid seams between tiles. rows = [] for i in range(0, z.shape[2], overlap_size): row = [] for j in range(0, z.shape[3], overlap_size): tile = z[ :, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size, ] if self.config.use_post_quant_conv: tile = self.post_quant_conv(tile) decoded = self.decoder(tile) row.append(decoded) rows.append(row) result_rows = [] for i, row in enumerate(rows): result_row = [] for j, tile in enumerate(row): # blend the above tile and the left tile # to the current tile and add the current tile to the result row if i > 0: tile = self.blend_v(rows[i - 1][j], tile, blend_extent) if j > 0: tile = self.blend_h(row[j - 1], tile, blend_extent) result_row.append(tile[:, :, :row_limit, :row_limit]) result_rows.append(torch.cat(result_row, dim=3)) dec = torch.cat(result_rows, dim=2) if not return_dict: return (dec,) return DecoderOutput(sample=dec) def forward( self, sample: torch.Tensor, sample_posterior: bool = False, generator: Optional[torch.Generator] = None, ) -> Union[DecoderOutput, torch.Tensor]: r""" Args: sample (`torch.Tensor`): Input sample. sample_posterior (`bool`, *optional*, defaults to `False`): Whether to sample from the posterior. """ x = sample posterior = self.encode(x).latent_dist if sample_posterior: z = posterior.sample(generator=generator) else: z = posterior.mode() dec = self.decode(z).sample return dec # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections def fuse_qkv_projections(self): """ Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused. > [!WARNING] > This API is 🧪 experimental. """ self.original_attn_processors = None for _, attn_processor in self.attn_processors.items(): if "Added" in str(attn_processor.__class__.__name__): raise ValueError( "`fuse_qkv_projections()` is not supported for models having added KV projections." ) self.original_attn_processors = self.attn_processors for module in self.modules(): if isinstance(module, Attention): module.fuse_projections(fuse=True) self.set_attn_processor(FusedAttnProcessor2_0()) EntryClass = AutoencoderKL ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_dc.py ================================================ # SPDX-License-Identifier: Apache-2.0 from collections.abc import Iterable import torch from torch import nn from sglang.multimodal_gen.configs.models.vaes.sana import SanaVAEConfig from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class AutoencoderDC(nn.Module): """Deep Compression Autoencoder wrapper with 32x spatial compression.""" def __init__(self, config: SanaVAEConfig = None, **kwargs): super().__init__() self._config = config self._inner_model = None self._loaded_state_dict: dict[str, torch.Tensor] = {} def _ensure_inner_model(self, state_dict: dict[str, torch.Tensor] | None = None): if self._inner_model is not None: return from diffusers import AutoencoderDC as DiffusersAutoencoderDC device = "cpu" state_to_load = ( state_dict if state_dict is not None else self._loaded_state_dict ) if state_to_load: first_tensor = next(iter(state_to_load.values())) device = first_tensor.device hf_config = {} if self._config is not None: arch = self._config.arch_config for key, value in vars(arch).items(): if key == "extra_attrs" and isinstance(value, dict): for ek, ev in value.items(): hf_config[ek] = ev elif not key.startswith("_") and not callable(value): hf_config[key] = value self._inner_model = DiffusersAutoencoderDC.from_config(hf_config) if state_to_load: missing, unexpected = self._inner_model.load_state_dict( state_to_load, strict=False ) if missing: logger.warning( "AutoencoderDC missing keys when loading: %d keys", len(missing) ) if len(missing) > 10: logger.debug("First 10 missing keys: %s", list(missing)[:10]) else: logger.debug("Missing keys: %s", list(missing)) if unexpected: logger.debug( "AutoencoderDC unexpected keys when loading: %d keys", len(unexpected), ) if state_dict is None: self._loaded_state_dict.clear() self._inner_model = self._inner_model.to(device) @property def config(self): if self._inner_model is not None: return self._inner_model.config return self._config @property def dtype(self): if self._inner_model is not None: return next(self._inner_model.parameters()).dtype return torch.float32 @property def device(self): if self._inner_model is not None: return next(self._inner_model.parameters()).device return torch.device("cpu") def encode(self, x: torch.Tensor, **kwargs): self._ensure_inner_model() return self._inner_model.encode(x, **kwargs) def decode(self, z: torch.Tensor, **kwargs): self._ensure_inner_model() z = z.to(dtype=self.dtype) return self._inner_model.decode(z, **kwargs) def forward(self, x: torch.Tensor, **kwargs): self._ensure_inner_model() return self._inner_model(x, **kwargs) def load_state_dict( self, state_dict: dict[str, torch.Tensor], strict: bool = True, assign: bool = False, ): """Intercept load_state_dict to route weights into the inner diffusers model.""" self._ensure_inner_model(state_dict=state_dict) def state_dict(self, *args, **kwargs) -> dict[str, torch.Tensor]: self._ensure_inner_model() return self._inner_model.state_dict(*args, **kwargs) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Buffer weights for deferred loading. The inner model is built lazily.""" loaded_params: set[str] = set() for name, weight in weights: self._loaded_state_dict[name] = weight loaded_params.add(name) return loaded_params def to(self, *args, **kwargs): if self._inner_model is not None: self._inner_model = self._inner_model.to(*args, **kwargs) return super().to(*args, **kwargs) EntryClass = AutoencoderDC ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_flux2.py ================================================ import math from typing import Dict, Optional, Tuple, Union import torch import torch.nn as nn from diffusers.models.attention_processor import ( ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnAddedKVProcessor, AttnProcessor, ) from diffusers.models.autoencoders.vae import ( Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder, ) from diffusers.models.modeling_outputs import AutoencoderKLOutput from sglang.multimodal_gen.configs.models.vaes.flux import Flux2VAEConfig from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE class AutoencoderKLFlux2(ParallelTiledVAE): r""" A VAE model with KL loss for encoding images into latents and decoding latent representations into images. This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented for all models (such as downloading or saving). Parameters: """ _supports_gradient_checkpointing = True _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"] def __init__( self, config: Flux2VAEConfig, ): super().__init__(config=config) self.config = config arch_config = config.arch_config in_channels: int = arch_config.in_channels out_channels: int = arch_config.out_channels down_block_types: Tuple[str, ...] = arch_config.down_block_types up_block_types: Tuple[str, ...] = arch_config.up_block_types block_out_channels: Tuple[int, ...] = arch_config.block_out_channels layers_per_block: int = arch_config.layers_per_block act_fn: str = arch_config.act_fn latent_channels: int = arch_config.latent_channels norm_num_groups: int = arch_config.norm_num_groups sample_size: int = arch_config.sample_size force_upcast: bool = arch_config.force_upcast use_quant_conv: bool = arch_config.use_quant_conv use_post_quant_conv: bool = arch_config.use_post_quant_conv mid_block_add_attention: bool = arch_config.mid_block_add_attention batch_norm_eps: float = arch_config.batch_norm_eps batch_norm_momentum: float = arch_config.batch_norm_momentum patch_size: Tuple[int, int] = arch_config.patch_size # pass init params to Encoder self.encoder = Encoder( in_channels=in_channels, out_channels=latent_channels, down_block_types=down_block_types, block_out_channels=block_out_channels, layers_per_block=layers_per_block, act_fn=act_fn, norm_num_groups=norm_num_groups, double_z=True, mid_block_add_attention=mid_block_add_attention, ) # pass init params to Decoder self.decoder = Decoder( in_channels=latent_channels, out_channels=out_channels, up_block_types=up_block_types, block_out_channels=block_out_channels, layers_per_block=layers_per_block, norm_num_groups=norm_num_groups, act_fn=act_fn, mid_block_add_attention=mid_block_add_attention, ) self.quant_conv = ( nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) if use_quant_conv else None ) self.post_quant_conv = ( nn.Conv2d(latent_channels, latent_channels, 1) if use_post_quant_conv else None ) self.bn = nn.BatchNorm2d( math.prod(patch_size) * latent_channels, eps=batch_norm_eps, momentum=batch_norm_momentum, affine=False, track_running_stats=True, ) self.use_slicing = False self.use_tiling = False # only relevant if vae tiling is enabled self.tile_sample_min_size = self.config.sample_size sample_size = ( self.config.sample_size[0] if isinstance(self.config.sample_size, (list, tuple)) else self.config.sample_size ) self.tile_latent_min_size = int( sample_size / (2 ** (len(self.config.block_out_channels) - 1)) ) self.tile_overlap_factor = 0.25 @property # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors def attn_processors(self) -> Dict[str, AttentionProcessor]: r""" Returns: `dict` of attention processors: A dictionary containing all attention processors used in the model with indexed by its weight name. """ # set recursively processors = {} def fn_recursive_add_processors( name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor], ): if hasattr(module, "get_processor"): processors[f"{name}.processor"] = module.get_processor() for sub_name, child in module.named_children(): fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) return processors for name, module in self.named_children(): fn_recursive_add_processors(name, module, processors) return processors # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor def set_attn_processor( self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]] ): r""" Sets the attention processor to use to compute attention. Parameters: processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): The instantiated processor class or a dictionary of processor classes that will be set as the processor for **all** `Attention` layers. If `processor` is a dict, the key needs to define the path to the corresponding cross attention processor. This is strongly recommended when setting trainable attention processors. """ count = len(self.attn_processors.keys()) if isinstance(processor, dict) and len(processor) != count: raise ValueError( f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" f" number of attention layers: {count}. Please make sure to pass {count} processor classes." ) def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): if hasattr(module, "set_processor"): if not isinstance(processor, dict): module.set_processor(processor) else: module.set_processor(processor.pop(f"{name}.processor")) for sub_name, child in module.named_children(): fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) for name, module in self.named_children(): fn_recursive_attn_processor(name, module, processor) # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor def set_default_attn_processor(self): """ Disables custom attention processors and sets the default attention implementation. """ if all( proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values() ): processor = AttnAddedKVProcessor() elif all( proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values() ): processor = AttnProcessor() else: raise ValueError( f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" ) self.set_attn_processor(processor) def _encode(self, x: torch.Tensor) -> torch.Tensor: batch_size, num_channels, height, width = x.shape if self.use_tiling and ( width > self.tile_sample_min_size or height > self.tile_sample_min_size ): return self._tiled_encode(x) enc = self.encoder(x) if self.quant_conv is not None: enc = self.quant_conv(enc) return enc def encode( self, x: torch.Tensor, return_dict: bool = True ) -> Union[DiagonalGaussianDistribution]: """ Encode a batch of images into latents. Args: x (`torch.Tensor`): Input batch of images. return_dict (`bool`, *optional*, defaults to `True`): Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. Returns: The latent representations of the encoded images. If `return_dict` is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. """ if x.ndim == 5: assert x.shape[2] == 1 x = x.squeeze(2) if self.use_slicing and x.shape[0] > 1: encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] h = torch.cat(encoded_slices) else: h = self._encode(x) posterior = DiagonalGaussianDistribution(h) return posterior def _decode( self, z: torch.Tensor, return_dict: bool = True ) -> Union[DecoderOutput, torch.Tensor]: if self.use_tiling and ( z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size ): return self.tiled_decode(z, return_dict=return_dict) if self.post_quant_conv is not None: z = self.post_quant_conv(z) dec = self.decoder(z) if not return_dict: return (dec,) return DecoderOutput(sample=dec) def decode( self, z: torch.FloatTensor, return_dict: bool = True, generator=None ) -> Union[DecoderOutput, torch.FloatTensor]: """ Decode a batch of images. Args: z (`torch.Tensor`): Input batch of latent vectors. return_dict (`bool`, *optional*, defaults to `True`): Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. Returns: [`~models.vae.DecoderOutput`] or `tuple`: If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is returned. """ if self.use_slicing and z.shape[0] > 1: decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] decoded = torch.cat(decoded_slices) else: decoded = self._decode(z).sample return decoded def blend_v( self, a: torch.Tensor, b: torch.Tensor, blend_extent: int ) -> torch.Tensor: blend_extent = min(a.shape[2], b.shape[2], blend_extent) for y in range(blend_extent): b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[ :, :, y, : ] * (y / blend_extent) return b def blend_h( self, a: torch.Tensor, b: torch.Tensor, blend_extent: int ) -> torch.Tensor: blend_extent = min(a.shape[3], b.shape[3], blend_extent) for x in range(blend_extent): b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[ :, :, :, x ] * (x / blend_extent) return b def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor: r"""Encode a batch of images using a tiled encoder. When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the output, but they should be much less noticeable. Args: x (`torch.Tensor`): Input batch of images. Returns: `torch.Tensor`: The latent representation of the encoded videos. """ overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) row_limit = self.tile_latent_min_size - blend_extent # Split the image into 512x512 tiles and encode them separately. rows = [] for i in range(0, x.shape[2], overlap_size): row = [] for j in range(0, x.shape[3], overlap_size): tile = x[ :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size, ] tile = self.encoder(tile) if self.config.use_quant_conv: tile = self.quant_conv(tile) row.append(tile) rows.append(row) result_rows = [] for i, row in enumerate(rows): result_row = [] for j, tile in enumerate(row): # blend the above tile and the left tile # to the current tile and add the current tile to the result row if i > 0: tile = self.blend_v(rows[i - 1][j], tile, blend_extent) if j > 0: tile = self.blend_h(row[j - 1], tile, blend_extent) result_row.append(tile[:, :, :row_limit, :row_limit]) result_rows.append(torch.cat(result_row, dim=3)) enc = torch.cat(result_rows, dim=2) return enc def tiled_encode( self, x: torch.Tensor, return_dict: bool = True ) -> AutoencoderKLOutput: r"""Encode a batch of images using a tiled encoder. When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the output, but they should be much less noticeable. Args: x (`torch.Tensor`): Input batch of images. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. Returns: [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`: If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. """ deprecation_message = ( "The tiled_encode implementation supporting the `return_dict` parameter is deprecated. In the future, the " "implementation of this method will be replaced with that of `_tiled_encode` and you will no longer be able " "to pass `return_dict`. You will also have to create a `DiagonalGaussianDistribution()` from the returned value." ) overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) row_limit = self.tile_latent_min_size - blend_extent # Split the image into 512x512 tiles and encode them separately. rows = [] for i in range(0, x.shape[2], overlap_size): row = [] for j in range(0, x.shape[3], overlap_size): tile = x[ :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size, ] tile = self.encoder(tile) if self.config.use_quant_conv: tile = self.quant_conv(tile) row.append(tile) rows.append(row) result_rows = [] for i, row in enumerate(rows): result_row = [] for j, tile in enumerate(row): # blend the above tile and the left tile # to the current tile and add the current tile to the result row if i > 0: tile = self.blend_v(rows[i - 1][j], tile, blend_extent) if j > 0: tile = self.blend_h(row[j - 1], tile, blend_extent) result_row.append(tile[:, :, :row_limit, :row_limit]) result_rows.append(torch.cat(result_row, dim=3)) moments = torch.cat(result_rows, dim=2) posterior = DiagonalGaussianDistribution(moments) if not return_dict: return (posterior,) return AutoencoderKLOutput(latent_dist=posterior) def tiled_decode( self, z: torch.Tensor, return_dict: bool = True ) -> Union[DecoderOutput, torch.Tensor]: r""" Decode a batch of images using a tiled decoder. Args: z (`torch.Tensor`): Input batch of latent vectors. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. Returns: [`~models.vae.DecoderOutput`] or `tuple`: If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is returned. """ overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) row_limit = self.tile_sample_min_size - blend_extent # Split z into overlapping 64x64 tiles and decode them separately. # The tiles have an overlap to avoid seams between tiles. rows = [] for i in range(0, z.shape[2], overlap_size): row = [] for j in range(0, z.shape[3], overlap_size): tile = z[ :, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size, ] if self.config.use_post_quant_conv: tile = self.post_quant_conv(tile) decoded = self.decoder(tile) row.append(decoded) rows.append(row) result_rows = [] for i, row in enumerate(rows): result_row = [] for j, tile in enumerate(row): # blend the above tile and the left tile # to the current tile and add the current tile to the result row if i > 0: tile = self.blend_v(rows[i - 1][j], tile, blend_extent) if j > 0: tile = self.blend_h(row[j - 1], tile, blend_extent) result_row.append(tile[:, :, :row_limit, :row_limit]) result_rows.append(torch.cat(result_row, dim=3)) dec = torch.cat(result_rows, dim=2) if not return_dict: return (dec,) return DecoderOutput(sample=dec) def forward( self, sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, generator: Optional[torch.Generator] = None, ) -> Union[DecoderOutput, torch.Tensor]: r""" Args: sample (`torch.Tensor`): Input sample. sample_posterior (`bool`, *optional*, defaults to `False`): Whether to sample from the posterior. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`DecoderOutput`] instead of a plain tuple. """ x = sample posterior = self.encode(x).latent_dist if sample_posterior: z = posterior.sample(generator=generator) else: z = posterior.mode() dec = self.decode(z).sample if not return_dict: return (dec,) return DecoderOutput(sample=dec) EntryClass = AutoencoderKLFlux2 ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from typing import Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F from diffusers.models.activations import get_activation from diffusers.models.autoencoders.vae import ( DecoderOutput, DiagonalGaussianDistribution, ) from diffusers.models.modeling_outputs import AutoencoderKLOutput from sglang.multimodal_gen.configs.models.vaes.qwenimage import QwenImageVAEConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name CACHE_T = 2 class QwenImageCausalConv3d(nn.Conv3d): r""" A custom 3D causal convolution layer with feature caching support. This layer extends the standard Conv3D layer by ensuring causality in the time dimension and handling feature caching for efficient inference. Args: in_channels (int): Number of channels in the input image out_channels (int): Number of channels produced by the convolution kernel_size (int or tuple): Size of the convolving kernel stride (int or tuple, optional): Stride of the convolution. Default: 1 padding (int or tuple, optional): Zero-padding added to all three sides of the input. Default: 0 """ def __init__( self, in_channels: int, out_channels: int, kernel_size: Union[int, Tuple[int, int, int]], stride: Union[int, Tuple[int, int, int]] = 1, padding: Union[int, Tuple[int, int, int]] = 0, ) -> None: super().__init__( in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, ) # Set up causal padding self._padding = ( self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0, ) self.padding = (0, 0, 0) def forward(self, x, cache_x=None): padding = list(self._padding) if cache_x is not None and self._padding[4] > 0: cache_x = cache_x.to(x.device) x = torch.cat([cache_x, x], dim=2) padding[4] -= cache_x.shape[2] x = F.pad(x, padding) return super().forward(x) class QwenImageRMS_norm(nn.Module): r""" A custom RMS normalization layer. Args: dim (int): The number of dimensions to normalize over. channel_first (bool, optional): Whether the input tensor has channels as the first dimension. Default is True. images (bool, optional): Whether the input represents image data. Default is True. bias (bool, optional): Whether to include a learnable bias term. Default is False. """ def __init__( self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False, ) -> None: super().__init__() broadcastable_dims = (1, 1, 1) if not images else (1, 1) shape = (dim, *broadcastable_dims) if channel_first else (dim,) self.channel_first = channel_first self.scale = dim**0.5 self.gamma = nn.Parameter(torch.ones(shape)) self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 def forward(self, x): return ( F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias ) class QwenImageUpsample(nn.Upsample): r""" Perform upsampling while ensuring the output tensor has the same data type as the input. Returns: torch.Tensor: Upsampled tensor with the same data type as the input. """ def forward(self, x): return super().forward(x.float()).type_as(x) class QwenImageResample(nn.Module): r""" A custom resampling module for 2D and 3D data. Args: dim (int): The number of input/output channels. mode (str): The resampling mode. Must be one of: - 'none': No resampling (identity operation). - 'upsample2d': 2D upsampling with nearest-exact interpolation and convolution. - 'upsample3d': 3D upsampling with nearest-exact interpolation, convolution, and causal 3D convolution. - 'downsample2d': 2D downsampling with zero-padding and convolution. - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution. """ def __init__(self, dim: int, mode: str) -> None: super().__init__() self.dim = dim self.mode = mode # layers if mode == "upsample2d": self.resample = nn.Sequential( QwenImageUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim // 2, 3, padding=1), ) elif mode == "upsample3d": self.resample = nn.Sequential( QwenImageUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim // 2, 3, padding=1), ) self.time_conv = QwenImageCausalConv3d( dim, dim * 2, (3, 1, 1), padding=(1, 0, 0) ) elif mode == "downsample2d": self.resample = nn.Sequential( nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) ) elif mode == "downsample3d": self.resample = nn.Sequential( nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) ) self.time_conv = QwenImageCausalConv3d( dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0) ) else: self.resample = nn.Identity() def forward(self, x, feat_cache=None, feat_idx=[0]): b, c, t, h, w = x.size() if self.mode == "upsample3d": if feat_cache is not None: idx = feat_idx[0] if feat_cache[idx] is None: feat_cache[idx] = "Rep" feat_idx[0] += 1 else: cache_x = x[:, :, -CACHE_T:, :, :].clone() if ( cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep" ): # cache last frame of last two chunk cache_x = torch.cat( [ feat_cache[idx][:, :, -1, :, :] .unsqueeze(2) .to(cache_x.device), cache_x, ], dim=2, ) if ( cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep" ): cache_x = torch.cat( [torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2, ) if feat_cache[idx] == "Rep": x = self.time_conv(x) else: x = self.time_conv(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 x = x.reshape(b, 2, c, t, h, w) x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) x = x.reshape(b, c, t * 2, h, w) t = x.shape[2] x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) x = self.resample(x) x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) if self.mode == "downsample3d": if feat_cache is not None: idx = feat_idx[0] if feat_cache[idx] is None: feat_cache[idx] = x.clone() feat_idx[0] += 1 else: cache_x = x[:, :, -1:, :, :].clone() x = self.time_conv( torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2) ) feat_cache[idx] = cache_x feat_idx[0] += 1 return x class QwenImageResidualBlock(nn.Module): r""" A custom residual block module. Args: in_dim (int): Number of input channels. out_dim (int): Number of output channels. dropout (float, optional): Dropout rate for the dropout layer. Default is 0.0. non_linearity (str, optional): Type of non-linearity to use. Default is "silu". """ def __init__( self, in_dim: int, out_dim: int, dropout: float = 0.0, non_linearity: str = "silu", ) -> None: super().__init__() self.in_dim = in_dim self.out_dim = out_dim self.nonlinearity = get_activation(non_linearity) # layers self.norm1 = QwenImageRMS_norm(in_dim, images=False) self.conv1 = QwenImageCausalConv3d(in_dim, out_dim, 3, padding=1) self.norm2 = QwenImageRMS_norm(out_dim, images=False) self.dropout = nn.Dropout(dropout) self.conv2 = QwenImageCausalConv3d(out_dim, out_dim, 3, padding=1) self.conv_shortcut = ( QwenImageCausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() ) def forward(self, x, feat_cache=None, feat_idx=[0]): # Apply shortcut connection h = self.conv_shortcut(x) # First normalization and activation x = self.norm1(x) x = self.nonlinearity(x) if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: cache_x = torch.cat( [ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x, ], dim=2, ) x = self.conv1(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv1(x) # Second normalization and activation x = self.norm2(x) x = self.nonlinearity(x) # Dropout x = self.dropout(x) if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: cache_x = torch.cat( [ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x, ], dim=2, ) x = self.conv2(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv2(x) # Add residual connection return x + h class QwenImageAttentionBlock(nn.Module): r""" Causal self-attention with a single head. Args: dim (int): The number of channels in the input tensor. """ def __init__(self, dim): super().__init__() self.dim = dim # layers self.norm = QwenImageRMS_norm(dim) self.to_qkv = nn.Conv2d(dim, dim * 3, 1) self.proj = nn.Conv2d(dim, dim, 1) def forward(self, x): identity = x batch_size, channels, time, height, width = x.size() x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, height, width) x = self.norm(x) # compute query, key, value qkv = self.to_qkv(x) qkv = qkv.reshape(batch_size * time, 1, channels * 3, -1) qkv = qkv.permute(0, 1, 3, 2).contiguous() q, k, v = qkv.chunk(3, dim=-1) # apply attention x = F.scaled_dot_product_attention(q, k, v) x = ( x.squeeze(1) .permute(0, 2, 1) .reshape(batch_size * time, channels, height, width) ) # output projection x = self.proj(x) # Reshape back: [(b*t), c, h, w] -> [b, c, t, h, w] x = x.view(batch_size, time, channels, height, width) x = x.permute(0, 2, 1, 3, 4) return x + identity class QwenImageMidBlock(nn.Module): """ Middle block for QwenImageVAE encoder and decoder. Args: dim (int): Number of input/output channels. dropout (float): Dropout rate. non_linearity (str): Type of non-linearity to use. """ def __init__( self, dim: int, dropout: float = 0.0, non_linearity: str = "silu", num_layers: int = 1, ): super().__init__() self.dim = dim # Create the components resnets = [QwenImageResidualBlock(dim, dim, dropout, non_linearity)] attentions = [] for _ in range(num_layers): attentions.append(QwenImageAttentionBlock(dim)) resnets.append(QwenImageResidualBlock(dim, dim, dropout, non_linearity)) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) self.gradient_checkpointing = False def forward(self, x, feat_cache=None, feat_idx=[0]): # First residual block x = self.resnets[0](x, feat_cache, feat_idx) # Process through attention and residual blocks for attn, resnet in zip(self.attentions, self.resnets[1:]): if attn is not None: x = attn(x) x = resnet(x, feat_cache, feat_idx) return x class QwenImageEncoder3d(nn.Module): r""" A 3D encoder module. Args: dim (int): The base number of channels in the first layer. z_dim (int): The dimensionality of the latent space. dim_mult (list of int): Multipliers for the number of channels in each block. num_res_blocks (int): Number of residual blocks in each block. attn_scales (list of float): Scales at which to apply attention mechanisms. temperal_downsample (list of bool): Whether to downsample temporally in each block. dropout (float): Dropout rate for the dropout layers. non_linearity (str): Type of non-linearity to use. """ def __init__( self, dim=128, z_dim=4, dim_mult=[1, 2, 4, 4], num_res_blocks=2, attn_scales=[], temperal_downsample=[True, True, False], dropout=0.0, non_linearity: str = "silu", input_channels: int = 3, ): super().__init__() # dim = config.arch_config.dim # z_dim = config.arch_config.z_dim # dim_mult = config.arch_config.dim_mult # num_res_blocks = config.arch_config.num_res_blocks # attn_scales = config.arch_config.attn_scales # temperal_downsample = config.arch_config.temperal_downsample # dropout = config.arch_config.dropout # non_linearity = config.arch_config.non_linearity self.dim = dim self.z_dim = z_dim self.dim_mult = dim_mult self.num_res_blocks = num_res_blocks self.attn_scales = attn_scales self.temperal_downsample = temperal_downsample self.nonlinearity = get_activation(non_linearity) # dimensions dims = [dim * u for u in [1] + dim_mult] scale = 1.0 # init block self.conv_in = QwenImageCausalConv3d(input_channels, dims[0], 3, padding=1) # downsample blocks self.down_blocks = nn.ModuleList([]) for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): # residual (+attention) blocks for _ in range(num_res_blocks): self.down_blocks.append( QwenImageResidualBlock(in_dim, out_dim, dropout) ) if scale in attn_scales: self.down_blocks.append(QwenImageAttentionBlock(out_dim)) in_dim = out_dim # downsample block if i != len(dim_mult) - 1: mode = "downsample3d" if temperal_downsample[i] else "downsample2d" self.down_blocks.append(QwenImageResample(out_dim, mode=mode)) scale /= 2.0 # middle blocks self.mid_block = QwenImageMidBlock( out_dim, dropout, non_linearity, num_layers=1 ) # output blocks self.norm_out = QwenImageRMS_norm(out_dim, images=False) self.conv_out = QwenImageCausalConv3d(out_dim, z_dim, 3, padding=1) self.gradient_checkpointing = False def forward(self, x, feat_cache=None, feat_idx=[0]): if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat( [ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x, ], dim=2, ) x = self.conv_in(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv_in(x) ## downsamples for layer in self.down_blocks: if feat_cache is not None: x = layer(x, feat_cache, feat_idx) else: x = layer(x) ## middle x = self.mid_block(x, feat_cache, feat_idx) ## head x = self.norm_out(x) x = self.nonlinearity(x) if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat( [ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x, ], dim=2, ) x = self.conv_out(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv_out(x) return x class QwenImageUpBlock(nn.Module): """ A block that handles upsampling for the QwenImageVAE decoder. Args: in_dim (int): Input dimension out_dim (int): Output dimension num_res_blocks (int): Number of residual blocks dropout (float): Dropout rate upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d') non_linearity (str): Type of non-linearity to use """ def __init__( self, in_dim: int, out_dim: int, num_res_blocks: int, dropout: float = 0.0, upsample_mode: Optional[str] = None, non_linearity: str = "silu", ): super().__init__() self.in_dim = in_dim self.out_dim = out_dim # Create layers list resnets = [] # Add residual blocks and attention if needed current_dim = in_dim for _ in range(num_res_blocks + 1): resnets.append( QwenImageResidualBlock(current_dim, out_dim, dropout, non_linearity) ) current_dim = out_dim self.resnets = nn.ModuleList(resnets) # Add upsampling layer if needed self.upsamplers = None if upsample_mode is not None: self.upsamplers = nn.ModuleList( [QwenImageResample(out_dim, mode=upsample_mode)] ) self.gradient_checkpointing = False def forward(self, x, feat_cache=None, feat_idx=[0]): """ Forward pass through the upsampling block. Args: x (torch.Tensor): Input tensor feat_cache (list, optional): Feature cache for causal convolutions feat_idx (list, optional): Feature index for cache management Returns: torch.Tensor: Output tensor """ for resnet in self.resnets: if feat_cache is not None: x = resnet(x, feat_cache, feat_idx) else: x = resnet(x) if self.upsamplers is not None: if feat_cache is not None: x = self.upsamplers[0](x, feat_cache, feat_idx) else: x = self.upsamplers[0](x) return x class QwenImageDecoder3d(nn.Module): r""" A 3D decoder module. Args: dim (int): The base number of channels in the first layer. z_dim (int): The dimensionality of the latent space. dim_mult (list of int): Multipliers for the number of channels in each block. num_res_blocks (int): Number of residual blocks in each block. attn_scales (list of float): Scales at which to apply attention mechanisms. temperal_upsample (list of bool): Whether to upsample temporally in each block. dropout (float): Dropout rate for the dropout layers. non_linearity (str): Type of non-linearity to use. """ def __init__( self, dim=128, z_dim=4, dim_mult=[1, 2, 4, 4], num_res_blocks=2, attn_scales=[], temperal_upsample=[False, True, True], dropout=0.0, non_linearity: str = "silu", input_channels=3, ): super().__init__() self.dim = dim self.z_dim = z_dim self.dim_mult = dim_mult self.num_res_blocks = num_res_blocks self.attn_scales = attn_scales self.temperal_upsample = temperal_upsample self.nonlinearity = get_activation(non_linearity) # dimensions dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] scale = 1.0 / 2 ** (len(dim_mult) - 2) # init block self.conv_in = QwenImageCausalConv3d(z_dim, dims[0], 3, padding=1) # middle blocks self.mid_block = QwenImageMidBlock( dims[0], dropout, non_linearity, num_layers=1 ) # upsample blocks self.up_blocks = nn.ModuleList([]) for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): # residual (+attention) blocks if i > 0: in_dim = in_dim // 2 # Determine if we need upsampling upsample_mode = None if i != len(dim_mult) - 1: upsample_mode = "upsample3d" if temperal_upsample[i] else "upsample2d" # Create and add the upsampling block up_block = QwenImageUpBlock( in_dim=in_dim, out_dim=out_dim, num_res_blocks=num_res_blocks, dropout=dropout, upsample_mode=upsample_mode, non_linearity=non_linearity, ) self.up_blocks.append(up_block) # Update scale for next iteration if upsample_mode is not None: scale *= 2.0 # output blocks self.norm_out = QwenImageRMS_norm(out_dim, images=False) self.conv_out = QwenImageCausalConv3d(out_dim, input_channels, 3, padding=1) self.gradient_checkpointing = False def forward(self, x, feat_cache=None, feat_idx=[0]): ## conv1 if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat( [ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x, ], dim=2, ) x = self.conv_in(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv_in(x) ## middle x = self.mid_block(x, feat_cache, feat_idx) ## upsamples for up_block in self.up_blocks: x = up_block(x, feat_cache, feat_idx) ## head x = self.norm_out(x) x = self.nonlinearity(x) if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat( [ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x, ], dim=2, ) x = self.conv_out(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv_out(x) return x class AutoencoderKLQwenImage(ParallelTiledVAE): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented for all models (such as downloading or saving). """ _supports_gradient_checkpointing = False # fmt: off def __init__( self, config: QwenImageVAEConfig, ) -> None: # fmt: on super().__init__(config=config) base_dim = config.arch_config.base_dim z_dim = config.arch_config.z_dim dim_mult = config.arch_config.dim_mult num_res_blocks = config.arch_config.num_res_blocks attn_scales = config.arch_config.attn_scales temperal_downsample = config.arch_config.temperal_downsample dropout = config.arch_config.dropout # non_linearity = config.arch_config.non_linearity self.z_dim = z_dim self.temperal_downsample = temperal_downsample self.temperal_upsample = temperal_downsample[::-1] self.input_channels = config.arch_config.input_channels self.latents_mean = config.arch_config.latents_mean self.config = config.arch_config self.encoder = QwenImageEncoder3d( base_dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout, input_channels=self.input_channels ) self.quant_conv = QwenImageCausalConv3d(z_dim * 2, z_dim * 2, 1) self.post_quant_conv = QwenImageCausalConv3d(z_dim, z_dim, 1) self.decoder = QwenImageDecoder3d( base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout, input_channels=self.input_channels ) # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension # to perform decoding of a single video latent at a time. self.use_slicing = False # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the # intermediate tiles together, the memory requirement can be lowered. self.use_tiling = False # The minimal tile height and width for spatial tiling to be used self.tile_sample_min_height = 256 self.tile_sample_min_width = 256 # The minimal distance between two spatial tiles self.tile_sample_stride_height = 192 self.tile_sample_stride_width = 192 # Precompute and cache conv counts for encoder and decoder for clear_cache speedup self._cached_conv_counts = { "decoder": sum(isinstance(m, QwenImageCausalConv3d) for m in self.decoder.modules()) if self.decoder is not None else 0, "encoder": sum(isinstance(m, QwenImageCausalConv3d) for m in self.encoder.modules()) if self.encoder is not None else 0, } cuda_device = get_local_torch_device() # FIXME: hardcode dtype = torch.bfloat16 latent_channels = config.arch_config.z_dim self.shift_factor = ( torch.tensor( config.arch_config.latents_mean ) .view(1, latent_channels, 1, 1, 1) .to(cuda_device, dtype) ) def enable_tiling( self, tile_sample_min_height: Optional[int] = None, tile_sample_min_width: Optional[int] = None, tile_sample_stride_height: Optional[float] = None, tile_sample_stride_width: Optional[float] = None, ) -> None: r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow processing larger images. Args: tile_sample_min_height (`int`, *optional*): The minimum height required for a sample to be separated into tiles across the height dimension. tile_sample_min_width (`int`, *optional*): The minimum width required for a sample to be separated into tiles across the width dimension. tile_sample_stride_height (`int`, *optional*): The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are no tiling artifacts produced across the height dimension. tile_sample_stride_width (`int`, *optional*): The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling artifacts produced across the width dimension. """ self.use_tiling = True self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width def disable_tiling(self) -> None: r""" Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing decoding in one step. """ self.use_tiling = False def enable_slicing(self) -> None: r""" Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. """ self.use_slicing = True def disable_slicing(self) -> None: r""" Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing decoding in one step. """ self.use_slicing = False def clear_cache(self): def _count_conv3d(model): count = 0 for m in model.modules(): if isinstance(m, QwenImageCausalConv3d): count += 1 return count self._conv_num = _count_conv3d(self.decoder) self._conv_idx = [0] self._feat_map = [None] * self._conv_num # cache encode self._enc_conv_num = _count_conv3d(self.encoder) self._enc_conv_idx = [0] self._enc_feat_map = [None] * self._enc_conv_num def _encode(self, x: torch.Tensor): _, _, num_frame, height, width = x.shape if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): return self.tiled_encode(x) self.clear_cache() iter_ = 1 + (num_frame - 1) // 4 for i in range(iter_): self._enc_conv_idx = [0] if i == 0: out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) else: out_ = self.encoder( x[:, :, 1 + 4 * (i - 1): 1 + 4 * i, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx, ) out = torch.cat([out, out_], 2) enc = self.quant_conv(out) self.clear_cache() return enc def encode( self, x: torch.Tensor, return_dict: bool = True ) -> DiagonalGaussianDistribution: r""" Encode a batch of images into latents. Args: x (`torch.Tensor`): Input batch of images. return_dict (`bool`, *optional*, defaults to `True`): Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. Returns: The latent representations of the encoded videos. If `return_dict` is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. """ if self.use_slicing and x.shape[0] > 1: encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] h = torch.cat(encoded_slices) else: h = self._encode(x) posterior = DiagonalGaussianDistribution(h) return posterior def _decode(self, z: torch.Tensor, return_dict: bool = True): _, _, num_frame, height, width = z.shape tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height): return self.tiled_decode(z, return_dict=return_dict) self.clear_cache() x = self.post_quant_conv(z) for i in range(num_frame): self._conv_idx = [0] if i == 0: out = self.decoder(x[:, :, i: i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) else: out_ = self.decoder(x[:, :, i: i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) out = torch.cat([out, out_], 2) out = torch.clamp(out, min=-1.0, max=1.0) self.clear_cache() if not return_dict: return (out,) return DecoderOutput(sample=out) def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: r""" Decode a batch of images. Args: z (`torch.Tensor`): Input batch of latent vectors. return_dict (`bool`, *optional*, defaults to `True`): Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. Returns: [`~models.vae.DecoderOutput`] or `tuple`: If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is returned. """ if self.use_slicing and z.shape[0] > 1: decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] decoded = torch.cat(decoded_slices) else: decoded = self._decode(z).sample return decoded def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) for y in range(blend_extent): b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( y / blend_extent ) return b def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) for x in range(blend_extent): b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( x / blend_extent ) return b def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput: r"""Encode a batch of images using a tiled encoder. Args: x (`torch.Tensor`): Input batch of videos. Returns: `torch.Tensor`: The latent representation of the encoded videos. """ _, _, num_frames, height, width = x.shape latent_height = height // self.spatial_compression_ratio latent_width = width // self.spatial_compression_ratio tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio blend_height = tile_latent_min_height - tile_latent_stride_height blend_width = tile_latent_min_width - tile_latent_stride_width # Split x into overlapping tiles and encode them separately. # The tiles have an overlap to avoid seams between tiles. rows = [] for i in range(0, height, self.tile_sample_stride_height): row = [] for j in range(0, width, self.tile_sample_stride_width): self.clear_cache() time = [] frame_range = 1 + (num_frames - 1) // 4 for k in range(frame_range): self._enc_conv_idx = [0] if k == 0: tile = x[:, :, :1, i: i + self.tile_sample_min_height, j: j + self.tile_sample_min_width] else: tile = x[ :, :, 1 + 4 * (k - 1): 1 + 4 * k, i: i + self.tile_sample_min_height, j: j + self.tile_sample_min_width, ] tile = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) tile = self.quant_conv(tile) time.append(tile) row.append(torch.cat(time, dim=2)) rows.append(row) self.clear_cache() result_rows = [] for i, row in enumerate(rows): result_row = [] for j, tile in enumerate(row): # blend the above tile and the left tile # to the current tile and add the current tile to the result row if i > 0: tile = self.blend_v(rows[i - 1][j], tile, blend_height) if j > 0: tile = self.blend_h(row[j - 1], tile, blend_width) result_row.append(tile[:, :, :, :tile_latent_stride_height, :tile_latent_stride_width]) result_rows.append(torch.cat(result_row, dim=-1)) enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width] return enc def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]: r""" Decode a batch of images using a tiled decoder. Args: z (`torch.Tensor`): Input batch of latent vectors. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. Returns: [`~models.vae.DecoderOutput`] or `tuple`: If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is returned. """ _, _, num_frames, height, width = z.shape sample_height = height * self.spatial_compression_ratio sample_width = width * self.spatial_compression_ratio tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio blend_height = self.tile_sample_min_height - self.tile_sample_stride_height blend_width = self.tile_sample_min_width - self.tile_sample_stride_width # Split z into overlapping tiles and decode them separately. # The tiles have an overlap to avoid seams between tiles. rows = [] for i in range(0, height, tile_latent_stride_height): row = [] for j in range(0, width, tile_latent_stride_width): self.clear_cache() time = [] for k in range(num_frames): self._conv_idx = [0] tile = z[:, :, k: k + 1, i: i + tile_latent_min_height, j: j + tile_latent_min_width] tile = self.post_quant_conv(tile) decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx) time.append(decoded) row.append(torch.cat(time, dim=2)) rows.append(row) self.clear_cache() result_rows = [] for i, row in enumerate(rows): result_row = [] for j, tile in enumerate(row): # blend the above tile and the left tile # to the current tile and add the current tile to the result row if i > 0: tile = self.blend_v(rows[i - 1][j], tile, blend_height) if j > 0: tile = self.blend_h(row[j - 1], tile, blend_width) result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width]) result_rows.append(torch.cat(result_row, dim=-1)) dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width] if not return_dict: return (dec,) return DecoderOutput(sample=dec) def forward( self, sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, generator: Optional[torch.Generator] = None, ) -> Union[DecoderOutput, torch.Tensor]: """ Args: sample (`torch.Tensor`): Input sample. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`DecoderOutput`] instead of a plain tuple. """ x = sample posterior = self.encode(x).latent_dist if sample_posterior: z = posterior.sample(generator=generator) else: z = posterior.mode() dec = self.decode(z, return_dict=return_dict) return dec EntryClass = AutoencoderKLQwenImage ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vaes/common.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from abc import ABC, abstractmethod from collections.abc import Iterator from math import prod from typing import Optional, cast import numpy as np import torch import torch.distributed as dist from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution from diffusers.utils.torch_utils import randn_tensor from torch import nn from sglang.multimodal_gen.configs.models import VAEConfig from sglang.multimodal_gen.runtime.distributed import ( get_sp_parallel_rank, get_sp_world_size, ) class ParallelTiledVAE(ABC, nn.Module): tile_sample_min_height: int tile_sample_min_width: int tile_sample_min_num_frames: int tile_sample_stride_height: int tile_sample_stride_width: int tile_sample_stride_num_frames: int blend_num_frames: int use_tiling: bool use_temporal_tiling: bool use_parallel_tiling: bool def __init__(self, config: VAEConfig, **kwargs) -> None: super().__init__() self.config = config self.tile_sample_min_height = config.tile_sample_min_height self.tile_sample_min_width = config.tile_sample_min_width self.tile_sample_min_num_frames = config.tile_sample_min_num_frames self.tile_sample_stride_height = config.tile_sample_stride_height self.tile_sample_stride_width = config.tile_sample_stride_width self.tile_sample_stride_num_frames = config.tile_sample_stride_num_frames self.blend_num_frames = config.blend_num_frames self.use_tiling = config.use_tiling self.use_temporal_tiling = config.use_temporal_tiling self.use_parallel_tiling = config.use_parallel_tiling @property def device(self): return next(self.parameters()).device @property def temporal_compression_ratio(self) -> int: return cast(int, self.config.temporal_compression_ratio) @property def spatial_compression_ratio(self) -> int: return cast(int, self.config.spatial_compression_ratio) @property def scaling_factor(self) -> float | torch.Tensor: return cast(float | torch.Tensor, self.config.scaling_factor) @abstractmethod def _encode(self, *args, **kwargs) -> torch.Tensor: pass @abstractmethod def _decode(self, *args, **kwargs) -> torch.Tensor: pass def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution: batch_size, num_channels, num_frames, height, width = x.shape latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1 if ( self.use_tiling and self.use_temporal_tiling and num_frames > self.tile_sample_min_num_frames ): latents = self.tiled_encode(x)[:, :, :latent_num_frames] elif self.use_tiling and ( width > self.tile_sample_min_width or height > self.tile_sample_min_height ): latents = self.spatial_tiled_encode(x)[:, :, :latent_num_frames] else: latents = self._encode(x)[:, :, :latent_num_frames] return DiagonalGaussianDistribution(latents) def decode(self, z: torch.Tensor) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = z.shape tile_latent_min_height = ( self.tile_sample_min_height // self.spatial_compression_ratio ) tile_latent_min_width = ( self.tile_sample_stride_width // self.spatial_compression_ratio ) tile_latent_min_num_frames = ( self.tile_sample_min_num_frames // self.temporal_compression_ratio ) num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1 if self.use_tiling and self.use_parallel_tiling and get_sp_world_size() > 1: return self.parallel_tiled_decode(z)[:, :, :num_sample_frames] if ( self.use_tiling and self.use_temporal_tiling and num_frames > tile_latent_min_num_frames ): return self.tiled_decode(z)[:, :, :num_sample_frames] if self.use_tiling and ( width > tile_latent_min_width or height > tile_latent_min_height ): return self.spatial_tiled_decode(z)[:, :, :num_sample_frames] return self._decode(z)[:, :, :num_sample_frames] def blend_v( self, a: torch.Tensor, b: torch.Tensor, blend_extent: int ) -> torch.Tensor: blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) for y in range(blend_extent): b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * ( 1 - y / blend_extent ) + b[:, :, :, y, :] * (y / blend_extent) return b def blend_h( self, a: torch.Tensor, b: torch.Tensor, blend_extent: int ) -> torch.Tensor: blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) for x in range(blend_extent): b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * ( 1 - x / blend_extent ) + b[:, :, :, :, x] * (x / blend_extent) return b def blend_t( self, a: torch.Tensor, b: torch.Tensor, blend_extent: int ) -> torch.Tensor: blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) for x in range(blend_extent): b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * ( 1 - x / blend_extent ) + b[:, :, x, :, :] * (x / blend_extent) return b def spatial_tiled_encode(self, x: torch.Tensor) -> torch.Tensor: r"""Encode a batch of images using a tiled encoder. Args: x (`torch.Tensor`): Input batch of videos. Returns: `torch.Tensor`: The latent representation of the encoded videos. """ _, _, _, height, width = x.shape # latent_height = height // self.spatial_compression_ratio # latent_width = width // self.spatial_compression_ratio tile_latent_min_height = ( self.tile_sample_min_height // self.spatial_compression_ratio ) tile_latent_min_width = ( self.tile_sample_min_width // self.spatial_compression_ratio ) tile_latent_stride_height = ( self.tile_sample_stride_height // self.spatial_compression_ratio ) tile_latent_stride_width = ( self.tile_sample_stride_width // self.spatial_compression_ratio ) blend_height = tile_latent_min_height - tile_latent_stride_height blend_width = tile_latent_min_width - tile_latent_stride_width # Split x into overlapping tiles and encode them separately. # The tiles have an overlap to avoid seams between tiles. rows = [] for i in range(0, height, self.tile_sample_stride_height): row = [] for j in range(0, width, self.tile_sample_stride_width): tile = x[ :, :, :, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width, ] tile = self._encode(tile) row.append(tile) rows.append(row) return self._merge_spatial_tiles( rows, blend_height, blend_width, tile_latent_stride_height, tile_latent_stride_width, ) def _parallel_data_generator( self, gathered_results, gathered_dim_metadata ) -> Iterator[tuple[torch.Tensor, int]]: global_idx = 0 for i, per_rank_metadata in enumerate(gathered_dim_metadata): _start_shape = 0 for shape in per_rank_metadata: mul_shape = prod(shape) yield ( gathered_results[ i, _start_shape : _start_shape + mul_shape ].reshape(shape), global_idx, ) _start_shape += mul_shape global_idx += 1 def parallel_tiled_decode(self, z: torch.FloatTensor) -> torch.FloatTensor: """ Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs """ world_size, rank = get_sp_world_size(), get_sp_parallel_rank() B, C, T, H, W = z.shape # Calculate parameters tile_latent_min_height = ( self.tile_sample_min_height // self.spatial_compression_ratio ) tile_latent_min_width = ( self.tile_sample_min_width // self.spatial_compression_ratio ) tile_latent_min_num_frames = ( self.tile_sample_min_num_frames // self.temporal_compression_ratio ) tile_latent_stride_height = ( self.tile_sample_stride_height // self.spatial_compression_ratio ) tile_latent_stride_width = ( self.tile_sample_stride_width // self.spatial_compression_ratio ) tile_latent_stride_num_frames = ( self.tile_sample_stride_num_frames // self.temporal_compression_ratio ) blend_height = self.tile_sample_min_height - self.tile_sample_stride_height blend_width = self.tile_sample_min_width - self.tile_sample_stride_width # Calculate tile dimensions num_t_tiles = ( T + tile_latent_stride_num_frames - 1 ) // tile_latent_stride_num_frames num_h_tiles = (H + tile_latent_stride_height - 1) // tile_latent_stride_height num_w_tiles = (W + tile_latent_stride_width - 1) // tile_latent_stride_width total_spatial_tiles = num_h_tiles * num_w_tiles total_tiles = num_t_tiles * total_spatial_tiles # Calculate tiles per rank and padding tiles_per_rank = (total_tiles + world_size - 1) // world_size start_tile_idx = rank * tiles_per_rank end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles) local_results = [] local_dim_metadata = [] # Process assigned tiles for local_idx, global_idx in enumerate(range(start_tile_idx, end_tile_idx)): t_idx = global_idx // total_spatial_tiles spatial_idx = global_idx % total_spatial_tiles h_idx = spatial_idx // num_w_tiles w_idx = spatial_idx % num_w_tiles # Calculate positions t_start = t_idx * tile_latent_stride_num_frames h_start = h_idx * tile_latent_stride_height w_start = w_idx * tile_latent_stride_width # Extract and process tile tile = z[ :, :, t_start : t_start + tile_latent_min_num_frames + 1, h_start : h_start + tile_latent_min_height, w_start : w_start + tile_latent_min_width, ] # Process tile tile = self._decode(tile) if t_start > 0: tile = tile[:, :, 1:, :, :] # Store metadata shape = tile.shape # Store decoded data (flattened) decoded_flat = tile.reshape(-1) local_results.append(decoded_flat) local_dim_metadata.append(shape) results = torch.cat(local_results, dim=0).contiguous() del local_results # first gather size to pad the results local_size = torch.tensor( [results.size(0)], device=results.device, dtype=torch.int64 ) all_sizes = [ torch.zeros(1, device=results.device, dtype=torch.int64) for _ in range(world_size) ] dist.all_gather(all_sizes, local_size) max_size = max(size.item() for size in all_sizes) padded_results = torch.zeros(max_size, device=results.device) padded_results[: results.size(0)] = results del results # Gather all results gathered_dim_metadata = [None] * world_size gathered_results = ( torch.zeros_like(padded_results) .repeat(world_size, *[1] * len(padded_results.shape)) .contiguous() ) # use contiguous to make sure it won't copy data in the following operations # TODO (PY): use sgl_diffusion distributed methods dist.all_gather_into_tensor(gathered_results, padded_results) dist.all_gather_object(gathered_dim_metadata, local_dim_metadata) # Process gathered results data: list = [ [[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)] for _ in range(num_t_tiles) ] for current_data, global_idx in self._parallel_data_generator( gathered_results, gathered_dim_metadata ): t_idx = global_idx // total_spatial_tiles spatial_idx = global_idx % total_spatial_tiles h_idx = spatial_idx // num_w_tiles w_idx = spatial_idx % num_w_tiles data[t_idx][h_idx][w_idx] = current_data # Merge results result_slices = [] last_slice_data = None for i, tem_data in enumerate(data): slice_data = self._merge_spatial_tiles( tem_data, blend_height, blend_width, self.tile_sample_stride_height, self.tile_sample_stride_width, ) if i > 0: slice_data = self.blend_t( last_slice_data, slice_data, self.blend_num_frames ) result_slices.append( slice_data[:, :, : self.tile_sample_stride_num_frames, :, :] ) else: result_slices.append( slice_data[:, :, : self.tile_sample_stride_num_frames + 1, :, :] ) last_slice_data = slice_data dec = torch.cat(result_slices, dim=2) return dec def _merge_spatial_tiles( self, tiles, blend_height, blend_width, stride_height, stride_width ) -> torch.Tensor: """Helper function to merge spatial tiles with blending""" result_rows = [] for i, row in enumerate(tiles): result_row = [] for j, tile in enumerate(row): if i > 0: tile = self.blend_v(tiles[i - 1][j], tile, blend_height) if j > 0: tile = self.blend_h(row[j - 1], tile, blend_width) result_row.append(tile[:, :, :, :stride_height, :stride_width]) result_rows.append(torch.cat(result_row, dim=-1)) return torch.cat(result_rows, dim=-2) def spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor: r""" Decode a batch of images using a tiled decoder. Args: z (`torch.Tensor`): Input batch of latent vectors. Returns: `torch.Tensor`: The decoded images. """ _, _, _, height, width = z.shape # sample_height = height * self.spatial_compression_ratio # sample_width = width * self.spatial_compression_ratio tile_latent_min_height = ( self.tile_sample_min_height // self.spatial_compression_ratio ) tile_latent_min_width = ( self.tile_sample_min_width // self.spatial_compression_ratio ) tile_latent_stride_height = ( self.tile_sample_stride_height // self.spatial_compression_ratio ) tile_latent_stride_width = ( self.tile_sample_stride_width // self.spatial_compression_ratio ) blend_height = self.tile_sample_min_height - self.tile_sample_stride_height blend_width = self.tile_sample_min_width - self.tile_sample_stride_width # Split z into overlapping tiles and decode them separately. # The tiles have an overlap to avoid seams between tiles. rows = [] for i in range(0, height, tile_latent_stride_height): row = [] for j in range(0, width, tile_latent_stride_width): tile = z[ :, :, :, i : i + tile_latent_min_height, j : j + tile_latent_min_width, ] decoded = self._decode(tile) row.append(decoded) rows.append(row) return self._merge_spatial_tiles( rows, blend_height, blend_width, self.tile_sample_stride_height, self.tile_sample_stride_width, ) def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: _, _, num_frames, height, width = x.shape # tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio tile_latent_stride_num_frames = ( self.tile_sample_stride_num_frames // self.temporal_compression_ratio ) row = [] for i in range(0, num_frames, self.tile_sample_stride_num_frames): tile = x[:, :, i : i + self.tile_sample_min_num_frames + 1, :, :] if self.use_tiling and ( height > self.tile_sample_min_height or width > self.tile_sample_min_width ): tile = self.spatial_tiled_encode(tile) else: tile = self._encode(tile) if i > 0: tile = tile[:, :, 1:, :, :] row.append(tile) result_row = [] for i, tile in enumerate(row): if i > 0: tile = self.blend_t(row[i - 1], tile, self.blend_num_frames) result_row.append(tile[:, :, :tile_latent_stride_num_frames, :, :]) else: result_row.append(tile[:, :, : tile_latent_stride_num_frames + 1, :, :]) enc = torch.cat(result_row, dim=2) return enc def tiled_decode(self, z: torch.Tensor) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = z.shape tile_latent_min_height = ( self.tile_sample_min_height // self.spatial_compression_ratio ) tile_latent_min_width = ( self.tile_sample_min_width // self.spatial_compression_ratio ) tile_latent_min_num_frames = ( self.tile_sample_min_num_frames // self.temporal_compression_ratio ) tile_latent_stride_num_frames = ( self.tile_sample_stride_num_frames // self.temporal_compression_ratio ) row = [] for i in range(0, num_frames, tile_latent_stride_num_frames): tile = z[:, :, i : i + tile_latent_min_num_frames + 1, :, :] if self.use_tiling and ( tile.shape[-1] > tile_latent_min_width or tile.shape[-2] > tile_latent_min_height ): decoded = self.spatial_tiled_decode(tile) else: decoded = self._decode(tile) if i > 0: decoded = decoded[:, :, 1:, :, :] row.append(decoded) result_row = [] for i, tile in enumerate(row): if i > 0: tile = self.blend_t(row[i - 1], tile, self.blend_num_frames) result_row.append( tile[:, :, : self.tile_sample_stride_num_frames, :, :] ) else: result_row.append( tile[:, :, : self.tile_sample_stride_num_frames + 1, :, :] ) dec = torch.cat(result_row, dim=2) return dec def enable_tiling( self, tile_sample_min_height: int | None = None, tile_sample_min_width: int | None = None, tile_sample_min_num_frames: int | None = None, tile_sample_stride_height: int | None = None, tile_sample_stride_width: int | None = None, tile_sample_stride_num_frames: int | None = None, blend_num_frames: int | None = None, use_tiling: bool | None = None, use_temporal_tiling: bool | None = None, use_parallel_tiling: bool | None = None, ) -> None: r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow processing larger images. Args: tile_sample_min_height (`int`, *optional*): The minimum height required for a sample to be separated into tiles across the height dimension. tile_sample_min_width (`int`, *optional*): The minimum width required for a sample to be separated into tiles across the width dimension. tile_sample_min_num_frames (`int`, *optional*): The minimum number of frames required for a sample to be separated into tiles across the frame dimension. tile_sample_stride_height (`int`, *optional*): The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are no tiling artifacts produced across the height dimension. tile_sample_stride_width (`int`, *optional*): The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling artifacts produced across the width dimension. tile_sample_stride_num_frames (`int`, *optional*): The stride between two consecutive frame tiles. This is to ensure that there are no tiling artifacts produced across the frame dimension. """ self.use_tiling = True self.tile_sample_min_height = ( tile_sample_min_height or self.tile_sample_min_height ) self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width self.tile_sample_min_num_frames = ( tile_sample_min_num_frames or self.tile_sample_min_num_frames ) self.tile_sample_stride_height = ( tile_sample_stride_height or self.tile_sample_stride_height ) self.tile_sample_stride_width = ( tile_sample_stride_width or self.tile_sample_stride_width ) self.tile_sample_stride_num_frames = ( tile_sample_stride_num_frames or self.tile_sample_stride_num_frames ) if blend_num_frames is not None: self.blend_num_frames = blend_num_frames else: self.blend_num_frames = ( self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames ) self.use_tiling = use_tiling or self.use_tiling self.use_temporal_tiling = use_temporal_tiling or self.use_temporal_tiling self.use_parallel_tiling = use_parallel_tiling or self.use_parallel_tiling def disable_tiling(self) -> None: r""" Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing decoding in one step. """ self.use_tiling = False # adapted from https://github.com/huggingface/diffusers/blob/e7ffeae0a191f710881d1fbde00cd6ff025e81f2/src/diffusers/models/autoencoders/vae.py#L691 class DiagonalGaussianDistribution: def __init__(self, parameters: torch.Tensor, deterministic: bool = False): self.parameters = parameters self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) self.logvar = torch.clamp(self.logvar, -30.0, 20.0) self.deterministic = deterministic self.std = torch.exp(0.5 * self.logvar) self.var = torch.exp(self.logvar) if self.deterministic: self.var = self.std = torch.zeros_like( self.mean, device=self.parameters.device, dtype=self.parameters.dtype ) def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: # make sure sample is on the same device as the parameters and has same dtype sample = randn_tensor( self.mean.shape, generator=generator, device=self.parameters.device, dtype=self.parameters.dtype, ) x = self.mean + self.std * sample return x def kl( self, other: Optional["DiagonalGaussianDistribution"] = None, dims: tuple[int, ...] = (1, 2, 3), ) -> torch.Tensor: if self.deterministic: return torch.Tensor([0.0]) else: if other is None: return 0.5 * torch.sum( torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, dim=dims, ) else: return 0.5 * torch.sum( torch.pow(self.mean - other.mean, 2) / other.var + self.var / other.var - 1.0 - self.logvar + other.logvar, dim=dims, ) def nll( self, sample: torch.Tensor, dims: tuple[int, ...] = (1, 2, 3) ) -> torch.Tensor: if self.deterministic: return torch.Tensor([0.0]) logtwopi = np.log(2.0 * np.pi) return 0.5 * torch.sum( logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, dim=dims, ) def mode(self) -> torch.Tensor: return self.mean ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vaes/dac.py ================================================ # Copied and adapted from: https://github.com/descriptinc/descript-audio-codec # SPDX-License-Identifier: MIT import math from bisect import bisect_right from typing import Union import torch import torch.nn.functional as F from einops import rearrange from torch import nn from sglang.multimodal_gen.configs.models.vaes.dac import DacVAEConfig from sglang.multimodal_gen.runtime.models.vaes.common import ( DiagonalGaussianDistribution, ) # Scripting this brings model speed up 1.4x @torch.jit.script def snake(x, alpha): shape = x.shape x = x.reshape(shape[0], shape[1], -1) x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) x = x.reshape(shape) return x class Snake1d(nn.Module): def __init__(self, channels): super().__init__() self.alpha = nn.Parameter(torch.ones(1, channels, 1)) def forward(self, x): return snake(x, self.alpha) class VectorQuantize(nn.Module): """ Implementation of VQ similar to Karpathy's repo: https://github.com/karpathy/deep-vector-quantization Additionally uses following tricks from Improved VQGAN (https://arxiv.org/pdf/2110.04627.pdf): 1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space for improved codebook usage 2. l2-normalized codes: Converts euclidean distance to cosine similarity which improves training stability """ def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int): super().__init__() self.codebook_size = codebook_size self.codebook_dim = codebook_dim self.in_proj = nn.Conv1d(input_dim, codebook_dim, kernel_size=1) self.out_proj = nn.Conv1d(codebook_dim, input_dim, kernel_size=1) self.codebook = nn.Embedding(codebook_size, codebook_dim) def forward(self, z): """Quantize the input tensor using a fixed codebook and return the corresponding codebook vectors. Args: z (torch.Tensor): Input tensor with shape ``[B, D, T]``. Returns: tuple: A tuple containing: - z_q (torch.Tensor): Quantized continuous representation with shape ``[B, D, T]``. - commitment_loss (torch.Tensor): Commitment loss scalar to train encoder to predict vectors closer to codebook entries. - codebook_loss (torch.Tensor): Codebook loss scalar to update the codebook. - indices (torch.Tensor): Codebook indices (quantized discrete representation) with shape ``[B, T]``. - z_e (torch.Tensor): Projected latents (continuous representation before quantization) with shape ``[B, D, T]``. """ # Factorized codes (ViT-VQGAN) Project input into low-dimensional space z_e = self.in_proj(z) # z_e : (B x D x T) z_q, indices = self.decode_latents(z_e) commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2]) codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2]) z_q = ( z_e + (z_q - z_e).detach() ) # noop in forward pass, straight-through gradient estimator in backward pass z_q = self.out_proj(z_q) return z_q, commitment_loss, codebook_loss, indices, z_e def embed_code(self, embed_id): return F.embedding(embed_id, self.codebook.weight) def decode_code(self, embed_id): return self.embed_code(embed_id).transpose(1, 2) def decode_latents(self, latents): encodings = rearrange(latents, "b d t -> (b t) d") codebook = self.codebook.weight # codebook: (N x D) # L2 normalize encodings and codebook (ViT-VQGAN) encodings = F.normalize(encodings) codebook = F.normalize(codebook) # Compute euclidean distance with codebook dist = ( encodings.pow(2).sum(1, keepdim=True) - 2 * encodings @ codebook.t() + codebook.pow(2).sum(1, keepdim=True).t() ) indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0)) z_q = self.decode_code(indices) return z_q, indices class ResidualVectorQuantize(nn.Module): """ Introduced in SoundStream: An end2end neural audio codec https://arxiv.org/abs/2107.03312 """ def __init__( self, input_dim: int = 512, n_codebooks: int = 9, codebook_size: int = 1024, codebook_dim: Union[int, list] = 8, quantizer_dropout: float = 0.0, ): super().__init__() if isinstance(codebook_dim, int): codebook_dim = [codebook_dim for _ in range(n_codebooks)] self.n_codebooks = n_codebooks self.codebook_dim = codebook_dim self.codebook_size = codebook_size dim_offsets = [0] for dim in self.codebook_dim: dim_offsets.append(dim_offsets[-1] + dim) self._codebook_dim_offsets = tuple(dim_offsets) self.quantizers = nn.ModuleList( [ VectorQuantize(input_dim, codebook_size, codebook_dim[i]) for i in range(n_codebooks) ] ) self.quantizer_dropout = quantizer_dropout def forward(self, z, n_quantizers: int = None): """Quantize the input tensor using a fixed set of codebooks and return the corresponding codebook vectors. Args: z (torch.Tensor): Input tensor with shape ``[B, D, T]``. n_quantizers (int, optional): Number of quantizers to use. If ``None``, all quantizers are used. When ``n_quantizers`` < ``self.n_codebooks``, quantizer dropout is applied. Note: if ``self.quantizer_dropout`` > 0 and in training mode, this argument is ignored and a random number of quantizers is used. Returns: tuple: A tuple containing: - z_q (torch.Tensor): Quantized continuous representation with shape ``[B, D, T]``. - codes (torch.Tensor): Codebook indices for each codebook with shape ``[B, N, T]`` (quantized discrete representation of input). - latents (torch.Tensor): Projected latents with shape ``[B, N*D, T]`` (continuous representation before quantization). - commitment_loss (torch.Tensor): Commitment loss scalar to train encoder to predict vectors closer to codebook entries. - codebook_loss (torch.Tensor): Codebook loss scalar to update the codebook. """ z_q = 0 residual = z commitment_loss = 0 codebook_loss = 0 codebook_indices = [] latents = [] if n_quantizers is None: n_quantizers = self.n_codebooks quantizers = self.quantizers if self.training: batch_size = z.shape[0] device = z.device n_quantizers = torch.full( (batch_size,), self.n_codebooks + 1, device=device, dtype=torch.long, ) if self.quantizer_dropout > 0: dropout = torch.randint( 1, self.n_codebooks + 1, (batch_size,), device=device, ) n_dropout = int(batch_size * self.quantizer_dropout) if n_dropout > 0: n_quantizers[:n_dropout] = dropout[:n_dropout] for i, quantizer in enumerate(quantizers): z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer( residual ) # Create mask to apply quantizer dropout mask = i < n_quantizers z_q = z_q + z_q_i * mask[:, None, None] residual = residual - z_q_i # Sum losses commitment_loss += (commitment_loss_i * mask).mean() codebook_loss += (codebook_loss_i * mask).mean() codebook_indices.append(indices_i) latents.append(z_e_i) else: for i, quantizer in enumerate(quantizers): if i >= n_quantizers: break z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer( residual ) z_q = z_q + z_q_i residual = residual - z_q_i commitment_loss += commitment_loss_i.mean() codebook_loss += codebook_loss_i.mean() codebook_indices.append(indices_i) latents.append(z_e_i) codes = torch.stack(codebook_indices, dim=1) latents = torch.cat(latents, dim=1) return z_q, codes, latents, commitment_loss, codebook_loss def from_codes(self, codes: torch.Tensor): """Reconstruct the continuous representation from quantized codes. Args: codes (torch.Tensor): Quantized discrete representation with shape ``[B, N, T]``. Returns: tuple: A tuple containing: - z_q (torch.Tensor): Quantized continuous representation with shape ``[B, D, T]``. - z_p (torch.Tensor): Concatenated latent space representation with shape ``[B, N*D, T]``. - codes (torch.Tensor): Original input codebook indices with shape ``[B, N, T]``. """ z_q = 0.0 z_p = [] n_codebooks = codes.shape[1] for i in range(n_codebooks): z_p_i = self.quantizers[i].decode_code(codes[:, i, :]) z_p.append(z_p_i) z_q_i = self.quantizers[i].out_proj(z_p_i) z_q = z_q + z_q_i return z_q, torch.cat(z_p, dim=1), codes def from_latents(self, latents: torch.Tensor): """Reconstruct the continuous representation from unquantized latents. Args: latents (torch.Tensor): Continuous representation after projection with shape ``[B, N*D, T]``. Returns: tuple: A tuple containing: - z_q (torch.Tensor): Quantized representation of full-projected space with shape ``[B, D, T]``. - z_p (torch.Tensor): Quantized representation of latent space with shape ``[B, N*D, T]``. - codes (torch.Tensor): Codebook indices with shape ``[B, N, T]``. """ z_q = 0 z_p = [] codes = [] dims = self._codebook_dim_offsets n_codebooks = bisect_right(dims, latents.shape[1]) - 1 for i in range(n_codebooks): j, k = dims[i], dims[i + 1] z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :]) z_p.append(z_p_i) codes.append(codes_i) z_q_i = self.quantizers[i].out_proj(z_p_i) z_q = z_q + z_q_i return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1) class ResidualUnit(nn.Module): def __init__(self, dim: int = 16, dilation: int = 1): super().__init__() pad = ((7 - 1) * dilation) // 2 self.block = nn.Sequential( Snake1d(dim), nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad), Snake1d(dim), nn.Conv1d(dim, dim, kernel_size=1), ) def forward(self, x): y = self.block(x) pad = (x.shape[-1] - y.shape[-1]) // 2 if pad > 0: x = x[..., pad:-pad] return x + y class EncoderBlock(nn.Module): def __init__(self, dim: int = 16, stride: int = 1): super().__init__() self.block = nn.Sequential( ResidualUnit(dim // 2, dilation=1), ResidualUnit(dim // 2, dilation=3), ResidualUnit(dim // 2, dilation=9), Snake1d(dim // 2), nn.Conv1d( dim // 2, dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2), ), ) def forward(self, x): return self.block(x) class Encoder(nn.Module): def __init__( self, d_model: int = 64, strides: list = [2, 4, 8, 8], d_latent: int = 64, ): super().__init__() # Create first convolution self.block = [nn.Conv1d(1, d_model, kernel_size=7, padding=3)] # Create EncoderBlocks that double channels as they downsample by `stride` for stride in strides: d_model *= 2 self.block += [EncoderBlock(d_model, stride=stride)] # Create last convolution self.block += [ Snake1d(d_model), nn.Conv1d(d_model, d_latent, kernel_size=3, padding=1), ] # Wrap black into nn.Sequential self.block = nn.Sequential(*self.block) self.enc_dim = d_model def forward(self, x): return self.block(x) class DecoderBlock(nn.Module): def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1): super().__init__() self.block = nn.Sequential( Snake1d(input_dim), nn.ConvTranspose1d( input_dim, output_dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2), output_padding=stride % 2, ), ResidualUnit(output_dim, dilation=1), ResidualUnit(output_dim, dilation=3), ResidualUnit(output_dim, dilation=9), ) def forward(self, x): return self.block(x) class Decoder(nn.Module): def __init__( self, input_channel, channels, rates, d_out: int = 1, ): super().__init__() # Add first conv layer layers = [nn.Conv1d(input_channel, channels, kernel_size=7, padding=3)] # Add upsampling + MRF blocks for i, stride in enumerate(rates): input_dim = channels // 2**i output_dim = channels // 2 ** (i + 1) layers += [DecoderBlock(input_dim, output_dim, stride)] # Add final conv layer layers += [ Snake1d(output_dim), nn.Conv1d(output_dim, d_out, kernel_size=7, padding=3), nn.Tanh(), ] self.model = nn.Sequential(*layers) def forward(self, x): return self.model(x) class DAC(nn.Module): def __init__( self, config: DacVAEConfig, ): super().__init__() self.continuous = config.continuous self.decoder_dim = config.decoder_dim self.decoder_rates = config.decoder_rates self.encoder_dim = config.encoder_dim self.encoder_rates = config.encoder_rates self.hop_length = math.prod(config.encoder_rates) self.sample_rate = config.sample_rate if config.latent_dim is None: latent_dim = config.encoder_dim * (2 ** len(config.encoder_rates)) else: latent_dim = config.latent_dim self.latent_dim = latent_dim if config.load_encoder: self.encoder = Encoder(config.encoder_dim, config.encoder_rates, latent_dim) if not config.continuous: self.n_codebooks = config.n_codebooks self.codebook_size = config.codebook_size self.codebook_dim = config.codebook_dim self.quantizer = ResidualVectorQuantize( input_dim=latent_dim, n_codebooks=config.n_codebooks, codebook_size=config.codebook_size, codebook_dim=config.codebook_dim, quantizer_dropout=config.quantizer_dropout, ) else: self.quant_conv = torch.nn.Conv1d(latent_dim, 2 * latent_dim, 1) self.post_quant_conv = torch.nn.Conv1d(latent_dim, latent_dim, 1) if config.load_decoder: self.decoder = Decoder( latent_dim, config.decoder_dim, config.decoder_rates, ) self.apply(self.init_weights) @staticmethod def init_weights(m): if isinstance(m, nn.Conv1d): nn.init.trunc_normal_(m.weight, std=0.02) nn.init.constant_(m.bias, 0) @property def dtype(self): return next(self.parameters()).dtype @property def device(self): return next(self.parameters()).device def preprocess(self, audio_data, sample_rate): if sample_rate is None: sample_rate = self.sample_rate assert sample_rate == self.sample_rate length = audio_data.shape[-1] right_pad = math.ceil(length / self.hop_length) * self.hop_length - length audio_data = nn.functional.pad(audio_data, (0, right_pad)) return audio_data def encode( self, audio_data: torch.Tensor, n_quantizers: int = None, ): """Encode audio data into latent representations. This method processes audio through the encoder network and optionally applies vector quantization (in VQ mode) or projects to a Gaussian distribution (in continuous mode) to produce latent representations. Args: audio_data (torch.Tensor): Audio data to encode, with shape ``[B, 1, T]``. n_quantizers (int, optional): Number of quantizers to use. If ``None``, all quantizers are used. Only applicable in VQ mode (``continuous=False``). Returns: tuple: A tuple containing: - z (torch.Tensor): Encoded representation. In VQ mode, this is the quantized continuous representation with shape ``[B, D, T]``. In continuous mode, this is a ``DiagonalGaussianDistribution`` object. - codes (torch.Tensor or None): Codebook indices with shape ``[B, N, T]`` in VQ mode, ``None`` in continuous mode. - latents (torch.Tensor or None): Projected latents with shape ``[B, N*D, T]`` in VQ mode, ``None`` in continuous mode. - commitment_loss (torch.Tensor): Commitment loss scalar. - codebook_loss (torch.Tensor): Codebook loss scalar. Note: In continuous mode, the encoded representation is projected through a quantization convolution layer and wrapped in a ``DiagonalGaussianDistribution`` for VAE training. """ z = self.encoder(audio_data) # [B x D x T] if not self.continuous: z, codes, latents, commitment_loss, codebook_loss = self.quantizer( z, n_quantizers ) else: z = self.quant_conv(z) # [B x 2D x T] z = DiagonalGaussianDistribution(z) codes, latents, commitment_loss, codebook_loss = None, None, 0, 0 return z, codes, latents, commitment_loss, codebook_loss def decode(self, z: torch.Tensor): """Decode latent representations back to audio waveforms. This method takes latent representations (either quantized from VQ mode or sampled from the posterior in continuous mode) and reconstructs the corresponding audio through the decoder network. Args: z (torch.Tensor): Latent representation to decode, with shape ``[B, D, T]``. In VQ mode (``continuous=False``), this is the quantized continuous representation. In continuous mode (``continuous=True``), this is sampled from the posterior distribution. Returns: torch.Tensor: Decoded audio data with shape ``[B, 1, T']``. The output length T' is determined by the decoder's upsampling rates and may differ from the input temporal dimension T. Note: In continuous mode (``continuous=True``), the input is first passed through a post-quantization convolution layer before being fed to the decoder. """ if not self.continuous: audio = self.decoder(z) else: z = self.post_quant_conv(z) audio = self.decoder(z) return audio def forward( self, audio_data: torch.Tensor, sample_rate: int = None, n_quantizers: int = None, ): """Model forward pass. Args: audio_data (torch.Tensor): Audio to encode, shape [B, 1, T]. sample_rate (int, optional): Sample rate in Hz. Defaults to ``self.sample_rate`` when ``None``. n_quantizers (int, optional): Number of quantizers to use. When ``None``, all quantizers are used. Only used in VQ mode (``continuous=False``). Returns: dict: A dictionary containing different keys depending on the mode: **VQ Mode (``continuous=False``):** - "audio" (torch.Tensor): Decoded audio, shape [B, 1, length]. - "z" (torch.Tensor): Quantized continuous representation, shape [B, D, T]. - "codes" (torch.Tensor): Codebook indices, shape [B, N, T]. - "latents" (torch.Tensor): Projected latents, shape [B, N*D, T]. - "vq/commitment_loss" (torch.Tensor): Commitment loss. - "vq/codebook_loss" (torch.Tensor): Codebook loss. **Continuous Mode (``continuous=True``):** - "audio" (torch.Tensor): Decoded audio, shape [B, 1, length]. - "z" (torch.Tensor): Latent representation, shape [B, D, T]. - "kl_loss" (torch.Tensor): KL divergence loss (for VAE training). """ length = audio_data.shape[-1] audio_data = self.preprocess(audio_data, sample_rate) if not self.continuous: z, codes, latents, commitment_loss, codebook_loss = self.encode( audio_data, n_quantizers ) x = self.decode(z) return { "audio": x[..., :length], "z": z, "codes": codes, "latents": latents, "vq/commitment_loss": commitment_loss, "vq/codebook_loss": codebook_loss, } else: posterior, _, _, _, _ = self.encode(audio_data, n_quantizers) z = posterior.sample() x = self.decode(z) kl_loss = posterior.kl(dims=(1, 2)) kl_loss = kl_loss.mean() return { "audio": x[..., :length], "z": z, "kl_loss": kl_loss, } EntryClass = DAC ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py ================================================ # Copied and adapted from: https://github.com/Tencent-Hunyuan/Hunyuan3D-2 from __future__ import annotations from typing import Callable, List, Optional, Tuple, Union import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange, repeat from tqdm import tqdm from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # Attention backend selection scaled_dot_product_attention = F.scaled_dot_product_attention class CrossAttentionProcessor: def __call__(self, attn, q, k, v): out = scaled_dot_product_attention(q, k, v) return out class FlashVDMCrossAttentionProcessor: def __init__(self, topk=None): self.topk = topk def __call__(self, attn, q, k, v): if k.shape[-2] == 3072: topk = 1024 elif k.shape[-2] == 512: topk = 256 else: topk = k.shape[-2] // 3 if self.topk is True: q1 = q[:, :, ::100, :] sim = q1 @ k.transpose(-1, -2) sim = torch.mean(sim, -2) topk_ind = torch.topk(sim, dim=-1, k=topk).indices.squeeze(-2).unsqueeze(-1) topk_ind = topk_ind.expand(-1, -1, -1, v.shape[-1]) v0 = torch.gather(v, dim=-2, index=topk_ind) k0 = torch.gather(k, dim=-2, index=topk_ind) out = scaled_dot_product_attention(q, k0, v0) elif self.topk is False: out = scaled_dot_product_attention(q, k, v) else: idx, counts = self.topk start = 0 outs = [] for grid_coord, count in zip(idx, counts): end = start + count q_chunk = q[:, :, start:end, :] k0, v0 = self.select_topkv(q_chunk, k, v, topk) out = scaled_dot_product_attention(q_chunk, k0, v0) outs.append(out) start += count out = torch.cat(outs, dim=-2) self.topk = False return out def select_topkv(self, q_chunk, k, v, topk): q1 = q_chunk[:, :, ::50, :] sim = q1 @ k.transpose(-1, -2) sim = torch.mean(sim, -2) topk_ind = torch.topk(sim, dim=-1, k=topk).indices.squeeze(-2).unsqueeze(-1) topk_ind = topk_ind.expand(-1, -1, -1, v.shape[-1]) v0 = torch.gather(v, dim=-2, index=topk_ind) k0 = torch.gather(k, dim=-2, index=topk_ind) return k0, v0 class FlashVDMTopMCrossAttentionProcessor(FlashVDMCrossAttentionProcessor): def select_topkv(self, q_chunk, k, v, topk): q1 = q_chunk[:, :, ::30, :] sim = q1 @ k.transpose(-1, -2) # sim = sim.to(torch.float32) sim = sim.softmax(-1) sim = torch.mean(sim, 1) activated_token = torch.where(sim > 1e-6)[2] index = ( torch.unique(activated_token, return_counts=True)[0] .unsqueeze(0) .unsqueeze(0) .unsqueeze(-1) ) index = index.expand(-1, v.shape[1], -1, v.shape[-1]) v0 = torch.gather(v, dim=-2, index=index) k0 = torch.gather(k, dim=-2, index=index) return k0, v0 class FourierEmbedder(nn.Module): def __init__( self, num_freqs: int = 6, logspace: bool = True, input_dim: int = 3, include_input: bool = True, include_pi: bool = True, ) -> None: """The initialization""" super().__init__() if logspace: frequencies = 2.0 ** torch.arange(num_freqs, dtype=torch.float32) else: frequencies = torch.linspace( 1.0, 2.0 ** (num_freqs - 1), num_freqs, dtype=torch.float32 ) if include_pi: frequencies *= torch.pi self.register_buffer("frequencies", frequencies, persistent=False) self.include_input = include_input self.num_freqs = num_freqs self.out_dim = self.get_dims(input_dim) def get_dims(self, input_dim): temp = 1 if self.include_input or self.num_freqs == 0 else 0 out_dim = input_dim * (self.num_freqs * 2 + temp) return out_dim def forward(self, x: torch.Tensor) -> torch.Tensor: """Forward process.""" if self.num_freqs > 0: embed = (x[..., None].contiguous() * self.frequencies).view( *x.shape[:-1], -1 ) if self.include_input: return torch.cat((x, embed.sin(), embed.cos()), dim=-1) else: return torch.cat((embed.sin(), embed.cos()), dim=-1) else: return x class DropPath(nn.Module): """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True): super(DropPath, self).__init__() self.drop_prob = drop_prob self.scale_by_keep = scale_by_keep def forward(self, x): """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" if self.drop_prob == 0.0 or not self.training: return x keep_prob = 1 - self.drop_prob shape = (x.shape[0],) + (1,) * ( x.ndim - 1 ) # work with diff dim tensors, not just 2D ConvNets random_tensor = x.new_empty(shape).bernoulli_(keep_prob) if keep_prob > 0.0 and self.scale_by_keep: random_tensor.div_(keep_prob) return x * random_tensor def extra_repr(self): return f"drop_prob={round(self.drop_prob, 3):0.3f}" class MLP(nn.Module): def __init__( self, *, width: int, expand_ratio: int = 4, output_width: int = None, drop_path_rate: float = 0.0, ): super().__init__() self.width = width self.c_fc = nn.Linear(width, width * expand_ratio) self.c_proj = nn.Linear( width * expand_ratio, output_width if output_width is not None else width ) self.gelu = nn.GELU() self.drop_path = ( DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() ) def forward(self, x): return self.drop_path(self.c_proj(self.gelu(self.c_fc(x)))) class QKVMultiheadCrossAttention(nn.Module): def __init__( self, *, heads: int, n_data: Optional[int] = None, width=None, qk_norm=False, norm_layer=nn.LayerNorm, ): super().__init__() self.heads = heads self.n_data = n_data self.q_norm = ( norm_layer(width // heads, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity() ) self.k_norm = ( norm_layer(width // heads, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity() ) self.attn_processor = CrossAttentionProcessor() def forward(self, q, kv): _, n_ctx, _ = q.shape bs, n_data, width = kv.shape attn_ch = width // self.heads // 2 q = q.view(bs, n_ctx, self.heads, -1) kv = kv.view(bs, n_data, self.heads, -1) k, v = torch.split(kv, attn_ch, dim=-1) q = self.q_norm(q) k = self.k_norm(k) q, k, v = map( lambda t: rearrange(t, "b n h d -> b h n d", h=self.heads), (q, k, v) ) out = self.attn_processor(self, q, k, v) out = out.transpose(1, 2).reshape(bs, n_ctx, -1) return out class MultiheadCrossAttention(nn.Module): def __init__( self, *, width: int, heads: int, qkv_bias: bool = True, n_data: Optional[int] = None, data_width: Optional[int] = None, norm_layer=nn.LayerNorm, qk_norm: bool = False, kv_cache: bool = False, ): super().__init__() self.n_data = n_data self.width = width self.heads = heads self.data_width = width if data_width is None else data_width self.c_q = nn.Linear(width, width, bias=qkv_bias) self.c_kv = nn.Linear(self.data_width, width * 2, bias=qkv_bias) self.c_proj = nn.Linear(width, width) self.attention = QKVMultiheadCrossAttention( heads=heads, n_data=n_data, width=width, norm_layer=norm_layer, qk_norm=qk_norm, ) self.kv_cache = kv_cache self.data = None def forward(self, x, data): x = self.c_q(x) if self.kv_cache: if self.data is None: self.data = self.c_kv(data) logger.info( "Save kv cache,this should be called only once for one mesh" ) data = self.data else: data = self.c_kv(data) x = self.attention(x, data) x = self.c_proj(x) return x class ResidualCrossAttentionBlock(nn.Module): def __init__( self, *, n_data: Optional[int] = None, width: int, heads: int, mlp_expand_ratio: int = 4, data_width: Optional[int] = None, qkv_bias: bool = True, norm_layer=nn.LayerNorm, qk_norm: bool = False, ): super().__init__() if data_width is None: data_width = width self.attn = MultiheadCrossAttention( n_data=n_data, width=width, heads=heads, data_width=data_width, qkv_bias=qkv_bias, norm_layer=norm_layer, qk_norm=qk_norm, ) self.ln_1 = norm_layer(width, elementwise_affine=True, eps=1e-6) self.ln_2 = norm_layer(data_width, elementwise_affine=True, eps=1e-6) self.ln_3 = norm_layer(width, elementwise_affine=True, eps=1e-6) self.mlp = MLP(width=width, expand_ratio=mlp_expand_ratio) def forward(self, x: torch.Tensor, data: torch.Tensor): x = x + self.attn(self.ln_1(x), self.ln_2(data)) x = x + self.mlp(self.ln_3(x)) return x class QKVMultiheadAttention(nn.Module): def __init__( self, *, heads: int, n_ctx: int, width=None, qk_norm=False, norm_layer=nn.LayerNorm, ): super().__init__() self.heads = heads self.n_ctx = n_ctx self.q_norm = ( norm_layer(width // heads, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity() ) self.k_norm = ( norm_layer(width // heads, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity() ) def forward(self, qkv): bs, n_ctx, width = qkv.shape attn_ch = width // self.heads // 3 qkv = qkv.view(bs, n_ctx, self.heads, -1) q, k, v = torch.split(qkv, attn_ch, dim=-1) q = self.q_norm(q) k = self.k_norm(k) q, k, v = map( lambda t: rearrange(t, "b n h d -> b h n d", h=self.heads), (q, k, v) ) out = ( scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(bs, n_ctx, -1) ) return out class MultiheadAttention(nn.Module): def __init__( self, *, n_ctx: int, width: int, heads: int, qkv_bias: bool, norm_layer=nn.LayerNorm, qk_norm: bool = False, drop_path_rate: float = 0.0, ): super().__init__() self.n_ctx = n_ctx self.width = width self.heads = heads self.c_qkv = nn.Linear(width, width * 3, bias=qkv_bias) self.c_proj = nn.Linear(width, width) self.attention = QKVMultiheadAttention( heads=heads, n_ctx=n_ctx, width=width, norm_layer=norm_layer, qk_norm=qk_norm, ) self.drop_path = ( DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() ) def forward(self, x): x = self.c_qkv(x) x = self.attention(x) x = self.drop_path(self.c_proj(x)) return x class ResidualAttentionBlock(nn.Module): def __init__( self, *, n_ctx: int, width: int, heads: int, qkv_bias: bool = True, norm_layer=nn.LayerNorm, qk_norm: bool = False, drop_path_rate: float = 0.0, ): super().__init__() self.attn = MultiheadAttention( n_ctx=n_ctx, width=width, heads=heads, qkv_bias=qkv_bias, norm_layer=norm_layer, qk_norm=qk_norm, drop_path_rate=drop_path_rate, ) self.ln_1 = norm_layer(width, elementwise_affine=True, eps=1e-6) self.mlp = MLP(width=width, drop_path_rate=drop_path_rate) self.ln_2 = norm_layer(width, elementwise_affine=True, eps=1e-6) def forward(self, x: torch.Tensor): x = x + self.attn(self.ln_1(x)) x = x + self.mlp(self.ln_2(x)) return x class Transformer(nn.Module): def __init__( self, *, n_ctx: int, width: int, layers: int, heads: int, qkv_bias: bool = True, norm_layer=nn.LayerNorm, qk_norm: bool = False, drop_path_rate: float = 0.0, ): super().__init__() self.n_ctx = n_ctx self.width = width self.layers = layers self.resblocks = nn.ModuleList( [ ResidualAttentionBlock( n_ctx=n_ctx, width=width, heads=heads, qkv_bias=qkv_bias, norm_layer=norm_layer, qk_norm=qk_norm, drop_path_rate=drop_path_rate, ) for _ in range(layers) ] ) def forward(self, x: torch.Tensor): for block in self.resblocks: x = block(x) return x class CrossAttentionDecoder(nn.Module): def __init__( self, *, num_latents: int, out_channels: int, fourier_embedder: FourierEmbedder, width: int, heads: int, mlp_expand_ratio: int = 4, downsample_ratio: int = 1, enable_ln_post: bool = True, qkv_bias: bool = True, qk_norm: bool = False, label_type: str = "binary", ): super().__init__() self.enable_ln_post = enable_ln_post self.fourier_embedder = fourier_embedder self.downsample_ratio = downsample_ratio self.query_proj = nn.Linear(self.fourier_embedder.out_dim, width) if self.downsample_ratio != 1: self.latents_proj = nn.Linear(width * downsample_ratio, width) if self.enable_ln_post == False: qk_norm = False self.cross_attn_decoder = ResidualCrossAttentionBlock( n_data=num_latents, width=width, mlp_expand_ratio=mlp_expand_ratio, heads=heads, qkv_bias=qkv_bias, qk_norm=qk_norm, ) if self.enable_ln_post: self.ln_post = nn.LayerNorm(width) self.output_proj = nn.Linear(width, out_channels) self.label_type = label_type def set_cross_attention_processor(self, processor): self.cross_attn_decoder.attn.attention.attn_processor = processor def forward(self, queries=None, query_embeddings=None, latents=None): if query_embeddings is None: fourier_out = self.fourier_embedder(queries) query_embeddings = self.query_proj(fourier_out.to(latents.dtype)) if self.downsample_ratio != 1: latents = self.latents_proj(latents) x = self.cross_attn_decoder(query_embeddings, latents) if self.enable_ln_post: x = self.ln_post(x) occ = self.output_proj(x) return occ def generate_dense_grid_points( bbox_min: np.ndarray, bbox_max: np.ndarray, octree_resolution: int, indexing: str = "ij", ): length = bbox_max - bbox_min num_cells = octree_resolution x = np.linspace(bbox_min[0], bbox_max[0], int(num_cells) + 1, dtype=np.float32) y = np.linspace(bbox_min[1], bbox_max[1], int(num_cells) + 1, dtype=np.float32) z = np.linspace(bbox_min[2], bbox_max[2], int(num_cells) + 1, dtype=np.float32) [xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing) xyz = np.stack((xs, ys, zs), axis=-1) grid_size = [int(num_cells) + 1, int(num_cells) + 1, int(num_cells) + 1] return xyz, grid_size, length def extract_near_surface_volume_fn(input_tensor: torch.Tensor, alpha: float): """Extract near-surface voxels for hierarchical decoding.""" device = input_tensor.device val = input_tensor + alpha valid_mask = val > -9000 def get_neighbor(t, shift, axis): if shift == 0: return t.clone() pad_dims = [0, 0, 0, 0, 0, 0] if axis == 0: pad_idx = 0 if shift > 0 else 1 pad_dims[pad_idx] = abs(shift) elif axis == 1: pad_idx = 2 if shift > 0 else 3 pad_dims[pad_idx] = abs(shift) elif axis == 2: pad_idx = 4 if shift > 0 else 5 pad_dims[pad_idx] = abs(shift) padded = F.pad(t.unsqueeze(0).unsqueeze(0), pad_dims[::-1], mode="replicate") slice_dims = [slice(None)] * 3 if axis == 0: slice_dims[0] = slice(shift, None) if shift > 0 else slice(None, shift) elif axis == 1: slice_dims[1] = slice(shift, None) if shift > 0 else slice(None, shift) elif axis == 2: slice_dims[2] = slice(shift, None) if shift > 0 else slice(None, shift) padded = padded.squeeze(0).squeeze(0) return padded[slice_dims] left = get_neighbor(val, 1, axis=0) right = get_neighbor(val, -1, axis=0) back = get_neighbor(val, 1, axis=1) front = get_neighbor(val, -1, axis=1) down = get_neighbor(val, 1, axis=2) up = get_neighbor(val, -1, axis=2) def safe_where(neighbor): return torch.where(neighbor > -9000, neighbor, val) left, right = safe_where(left), safe_where(right) back, front = safe_where(back), safe_where(front) down, up = safe_where(down), safe_where(up) sign = torch.sign(val.to(torch.float32)) neighbors_sign = torch.stack( [ torch.sign(left.to(torch.float32)), torch.sign(right.to(torch.float32)), torch.sign(back.to(torch.float32)), torch.sign(front.to(torch.float32)), torch.sign(down.to(torch.float32)), torch.sign(up.to(torch.float32)), ], dim=0, ) same_sign = torch.all(neighbors_sign == sign, dim=0) mask = (~same_sign).to(torch.int32) return mask * valid_mask.to(torch.int32) class VanillaVolumeDecoder: """Standard volume decoder using dense grid evaluation.""" @torch.no_grad() def __call__( self, latents: torch.FloatTensor, geo_decoder: Callable, bounds: Union[Tuple[float], List[float], float] = 1.01, num_chunks: int = 10000, octree_resolution: int = None, enable_pbar: bool = True, **kwargs, ): device = latents.device dtype = latents.dtype batch_size = latents.shape[0] if isinstance(bounds, float): bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds] bbox_min, bbox_max = np.array(bounds[0:3]), np.array(bounds[3:6]) xyz_samples, grid_size, length = generate_dense_grid_points( bbox_min=bbox_min, bbox_max=bbox_max, octree_resolution=octree_resolution, indexing="ij", ) xyz_samples = ( torch.from_numpy(xyz_samples) .to(device, dtype=dtype) .contiguous() .reshape(-1, 3) ) batch_logits = [] for start in tqdm( range(0, xyz_samples.shape[0], num_chunks), desc="Volume Decoding", disable=not enable_pbar, ): chunk_queries = xyz_samples[start : start + num_chunks, :] chunk_queries = repeat(chunk_queries, "p c -> b p c", b=batch_size) logits = geo_decoder(queries=chunk_queries, latents=latents) batch_logits.append(logits) grid_logits = torch.cat(batch_logits, dim=1) grid_logits = grid_logits.view((batch_size, *grid_size)).float() return grid_logits class HierarchicalVolumeDecoding: """Hierarchical volume decoder with multi-resolution refinement.""" @torch.no_grad() def __call__( self, latents: torch.FloatTensor, geo_decoder: Callable, bounds: Union[Tuple[float], List[float], float] = 1.01, num_chunks: int = 10000, mc_level: float = 0.0, octree_resolution: int = None, min_resolution: int = 63, enable_pbar: bool = True, **kwargs, ): device = latents.device dtype = latents.dtype resolutions = [] if octree_resolution < min_resolution: resolutions.append(octree_resolution) while octree_resolution >= min_resolution: resolutions.append(octree_resolution) octree_resolution = octree_resolution // 2 resolutions.reverse() # 1. generate query points if isinstance(bounds, float): bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds] bbox_min = np.array(bounds[0:3]) bbox_max = np.array(bounds[3:6]) bbox_size = bbox_max - bbox_min xyz_samples, grid_size, length = generate_dense_grid_points( bbox_min=bbox_min, bbox_max=bbox_max, octree_resolution=resolutions[0], indexing="ij", ) dilate = nn.Conv3d(1, 1, 3, padding=1, bias=False, device=device, dtype=dtype) dilate.weight = torch.nn.Parameter( torch.ones(dilate.weight.shape, dtype=dtype, device=device) ) grid_size = np.array(grid_size) xyz_samples = ( torch.from_numpy(xyz_samples) .to(device, dtype=dtype) .contiguous() .reshape(-1, 3) ) # 2. latents to 3d volume batch_logits = [] batch_size = latents.shape[0] for start in tqdm( range(0, xyz_samples.shape[0], num_chunks), desc=f"Hierarchical Volume Decoding [r{resolutions[0] + 1}]", disable=not enable_pbar, ): queries = xyz_samples[start : start + num_chunks, :] batch_queries = repeat(queries, "p c -> b p c", b=batch_size) logits = geo_decoder(queries=batch_queries, latents=latents) batch_logits.append(logits) grid_logits = torch.cat(batch_logits, dim=1).view( (batch_size, grid_size[0], grid_size[1], grid_size[2]) ) for octree_depth_now in resolutions[1:]: grid_size = np.array([octree_depth_now + 1] * 3) resolution = bbox_size / octree_depth_now next_index = torch.zeros(tuple(grid_size), dtype=dtype, device=device) next_logits = torch.full( next_index.shape, -10000.0, dtype=dtype, device=device ) curr_points = extract_near_surface_volume_fn( grid_logits.squeeze(0), mc_level ) curr_points += grid_logits.squeeze(0).abs() < 0.95 if octree_depth_now == resolutions[-1]: expand_num = 0 else: expand_num = 1 for i in range(expand_num): curr_points = dilate(curr_points.unsqueeze(0).to(dtype)).squeeze(0) cidx_x, cidx_y, cidx_z = torch.where(curr_points > 0) next_index[cidx_x * 2, cidx_y * 2, cidx_z * 2] = 1 for i in range(2 - expand_num): next_index = dilate(next_index.unsqueeze(0)).squeeze(0) nidx = torch.where(next_index > 0) next_points = torch.stack(nidx, dim=1) next_points = next_points * torch.tensor( resolution, dtype=next_points.dtype, device=device ) + torch.tensor(bbox_min, dtype=next_points.dtype, device=device) # Check if next_points is empty if next_points.shape[0] == 0: logger.warning( f"No valid surface points found at resolution {octree_depth_now}, " f"skipping this level" ) continue batch_logits = [] for start in tqdm( range(0, next_points.shape[0], num_chunks), desc=f"Hierarchical Volume Decoding [r{octree_depth_now + 1}]", disable=not enable_pbar, ): queries = next_points[start : start + num_chunks, :] batch_queries = repeat(queries, "p c -> b p c", b=batch_size) logits = geo_decoder( queries=batch_queries.to(latents.dtype), latents=latents ) batch_logits.append(logits) grid_logits = torch.cat(batch_logits, dim=1) next_logits[nidx] = grid_logits[0, ..., 0] grid_logits = next_logits.unsqueeze(0) grid_logits[grid_logits == -10000.0] = float("nan") return grid_logits class FlashVDMVolumeDecoding: """Flash VDM volume decoder with adaptive KV selection.""" def __init__(self, topk_mode="mean"): if topk_mode not in ["mean", "merge"]: raise ValueError(f"Unsupported topk_mode {topk_mode}") if topk_mode == "mean": self.processor = FlashVDMCrossAttentionProcessor() else: self.processor = FlashVDMTopMCrossAttentionProcessor() @torch.no_grad() def __call__( self, latents: torch.FloatTensor, geo_decoder: CrossAttentionDecoder, bounds: Union[Tuple[float], List[float], float] = 1.01, num_chunks: int = 10000, mc_level: float = 0.0, octree_resolution: int = None, min_resolution: int = 63, mini_grid_num: int = 4, enable_pbar: bool = True, **kwargs, ): processor = self.processor geo_decoder.set_cross_attention_processor(processor) device = latents.device dtype = latents.dtype resolutions = [] orig_resolution = octree_resolution if octree_resolution < min_resolution: resolutions.append(octree_resolution) while octree_resolution >= min_resolution: resolutions.append(octree_resolution) octree_resolution = octree_resolution // 2 resolutions.reverse() resolutions[0] = round(resolutions[0] / mini_grid_num) * mini_grid_num - 1 for i, resolution in enumerate(resolutions[1:]): resolutions[i + 1] = resolutions[0] * 2 ** (i + 1) if isinstance(bounds, float): bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds] bbox_min = np.array(bounds[0:3]) bbox_max = np.array(bounds[3:6]) bbox_size = bbox_max - bbox_min xyz_samples, grid_size, length = generate_dense_grid_points( bbox_min=bbox_min, bbox_max=bbox_max, octree_resolution=resolutions[0], indexing="ij", ) logger.info(f"FlashVDMVolumeDecoding Resolution: {resolutions}") dilate = nn.Conv3d(1, 1, 3, padding=1, bias=False, device=device, dtype=dtype) dilate.weight = torch.nn.Parameter( torch.ones(dilate.weight.shape, dtype=dtype, device=device) ) grid_size = np.array(grid_size) xyz_samples = torch.from_numpy(xyz_samples).to(device, dtype=dtype) batch_size = latents.shape[0] mini_grid_size = xyz_samples.shape[0] // mini_grid_num xyz_samples = ( xyz_samples.view( mini_grid_num, mini_grid_size, mini_grid_num, mini_grid_size, mini_grid_num, mini_grid_size, 3, ) .permute(0, 2, 4, 1, 3, 5, 6) .reshape(-1, mini_grid_size * mini_grid_size * mini_grid_size, 3) ) batch_logits = [] num_batchs = max(num_chunks // xyz_samples.shape[1], 1) for start in tqdm( range(0, xyz_samples.shape[0], num_batchs), desc="FlashVDM Volume Decoding", disable=not enable_pbar, ): queries = xyz_samples[start : start + num_batchs, :] batch = queries.shape[0] batch_latents = repeat(latents.squeeze(0), "p c -> b p c", b=batch) processor.topk = True logits = geo_decoder(queries=queries, latents=batch_latents) batch_logits.append(logits) grid_logits = ( torch.cat(batch_logits, dim=0) .reshape( mini_grid_num, mini_grid_num, mini_grid_num, mini_grid_size, mini_grid_size, mini_grid_size, ) .permute(0, 3, 1, 4, 2, 5) .contiguous() .view((batch_size, grid_size[0], grid_size[1], grid_size[2])) ) for octree_depth_now in resolutions[1:]: grid_size = np.array([octree_depth_now + 1] * 3) resolution = bbox_size / octree_depth_now next_index = torch.zeros(tuple(grid_size), dtype=dtype, device=device) next_logits = torch.full( next_index.shape, -10000.0, dtype=dtype, device=device ) curr_points = extract_near_surface_volume_fn( grid_logits.squeeze(0), mc_level ) curr_points += grid_logits.squeeze(0).abs() < 0.95 expand_num = 0 if octree_depth_now == resolutions[-1] else 1 for _ in range(expand_num): curr_points = dilate(curr_points.unsqueeze(0).to(dtype)).squeeze(0) cidx_x, cidx_y, cidx_z = torch.where(curr_points > 0) next_index[cidx_x * 2, cidx_y * 2, cidx_z * 2] = 1 for _ in range(2 - expand_num): next_index = dilate(next_index.unsqueeze(0)).squeeze(0) nidx = torch.where(next_index > 0) next_points = torch.stack(nidx, dim=1) next_points = next_points * torch.tensor( resolution, dtype=torch.float32, device=device ) + torch.tensor(bbox_min, dtype=torch.float32, device=device) # Check if next_points is empty (no valid surface points found) if next_points.shape[0] == 0: # Skip this resolution level if no points found # Use the previous grid_logits as fallback logger.warning( f"No valid surface points found at resolution {octree_depth_now}, " f"skipping this level and using previous resolution grid_logits" ) continue query_grid_num = 6 min_val = next_points.min(axis=0).values max_val = next_points.max(axis=0).values vol_queries_index = ( (next_points - min_val) / (max_val - min_val) * (query_grid_num - 0.001) ) index = torch.floor(vol_queries_index).long() index = ( index[..., 0] * (query_grid_num**2) + index[..., 1] * query_grid_num + index[..., 2] ) index = index.sort() next_points = next_points[index.indices].unsqueeze(0).contiguous() unique_values = torch.unique(index.values, return_counts=True) grid_logits_flat = torch.zeros( (next_points.shape[1]), dtype=latents.dtype, device=latents.device ) input_grid = [[], []] logits_grid_list = [] start_num = 0 sum_num = 0 for grid_index, count in zip( unique_values[0].cpu().tolist(), unique_values[1].cpu().tolist() ): if sum_num + count < num_chunks or sum_num == 0: sum_num += count input_grid[0].append(grid_index) input_grid[1].append(count) else: processor.topk = input_grid logits_grid = geo_decoder( queries=next_points[:, start_num : start_num + sum_num], latents=latents, ) start_num = start_num + sum_num logits_grid_list.append(logits_grid) input_grid = [[grid_index], [count]] sum_num = count if sum_num > 0: processor.topk = input_grid logits_grid = geo_decoder( queries=next_points[:, start_num : start_num + sum_num], latents=latents, ) logits_grid_list.append(logits_grid) logits_grid = torch.cat(logits_grid_list, dim=1) grid_logits_flat[index.indices] = logits_grid.squeeze(0).squeeze(-1) next_logits[nidx] = grid_logits_flat grid_logits = next_logits.unsqueeze(0) grid_logits[grid_logits == -10000.0] = float("nan") return grid_logits class Latent2MeshOutput: """Container for mesh output from VAE decoder.""" def __init__(self, mesh_v=None, mesh_f=None): self.mesh_v = mesh_v self.mesh_f = mesh_f def center_vertices(vertices): """Translate vertices so bounding box is centered at zero.""" vert_min = vertices.min(dim=0)[0] vert_max = vertices.max(dim=0)[0] vert_center = 0.5 * (vert_min + vert_max) return vertices - vert_center class SurfaceExtractor: """Base class for surface extraction algorithms.""" def _compute_box_stat( self, bounds: Union[Tuple[float], List[float], float], octree_resolution: int ): if isinstance(bounds, float): bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds] bbox_min, bbox_max = np.array(bounds[0:3]), np.array(bounds[3:6]) bbox_size = bbox_max - bbox_min grid_size = [ int(octree_resolution) + 1, int(octree_resolution) + 1, int(octree_resolution) + 1, ] return grid_size, bbox_min, bbox_size def run(self, *args, **kwargs): raise NotImplementedError def __call__(self, grid_logits, **kwargs): outputs = [] for i in range(grid_logits.shape[0]): try: vertices, faces = self.run(grid_logits[i], **kwargs) vertices = vertices.astype(np.float32) faces = np.ascontiguousarray(faces) outputs.append(Latent2MeshOutput(mesh_v=vertices, mesh_f=faces)) except Exception: import traceback traceback.print_exc() outputs.append(None) return outputs class MCSurfaceExtractor(SurfaceExtractor): """Marching Cubes surface extractor.""" def run(self, grid_logit, *, mc_level, bounds, octree_resolution, **kwargs): from skimage import measure vertices, faces, normals, _ = measure.marching_cubes( grid_logit.cpu().numpy(), mc_level, method="lewiner" ) grid_size, bbox_min, bbox_size = self._compute_box_stat( bounds, octree_resolution ) vertices = vertices / grid_size * bbox_size + bbox_min return vertices, faces class DMCSurfaceExtractor(SurfaceExtractor): """Differentiable Marching Cubes surface extractor.""" def run(self, grid_logit, *, octree_resolution, **kwargs): device = grid_logit.device if not hasattr(self, "dmc"): try: from diso import DiffDMC self.dmc = DiffDMC(dtype=torch.float32).to(device) except ImportError: raise ImportError( "Please install diso via `pip install diso`, or set mc_algo to 'mc'" ) sdf = -grid_logit / octree_resolution sdf = sdf.to(torch.float32).contiguous() verts, faces = self.dmc(sdf, deform=None, return_quads=False, normalize=True) verts = center_vertices(verts) vertices = verts.detach().cpu().numpy() faces = faces.detach().cpu().numpy()[:, ::-1] return vertices, faces SurfaceExtractors = { "mc": MCSurfaceExtractor, "dmc": DMCSurfaceExtractor, } class VectsetVAE(nn.Module): """Base VAE class for vector set encoding.""" def __init__(self, volume_decoder=None, surface_extractor=None): super().__init__() if volume_decoder is None: volume_decoder = VanillaVolumeDecoder() if surface_extractor is None: surface_extractor = MCSurfaceExtractor() self.volume_decoder = volume_decoder self.surface_extractor = surface_extractor def latents2mesh(self, latents: torch.FloatTensor, **kwargs): """Convert latents to mesh.""" grid_logits = self.volume_decoder(latents, self.geo_decoder, **kwargs) outputs = self.surface_extractor(grid_logits, **kwargs) return outputs def enable_flashvdm_decoder( self, enabled: bool = True, adaptive_kv_selection=True, topk_mode="mean", mc_algo="dmc", ): """Enable or disable FlashVDM decoder for faster inference.""" if enabled: if adaptive_kv_selection: self.volume_decoder = FlashVDMVolumeDecoding(topk_mode) else: self.volume_decoder = HierarchicalVolumeDecoding() if mc_algo not in SurfaceExtractors: raise ValueError( f"Unsupported mc_algo {mc_algo}, available: {list(SurfaceExtractors.keys())}" ) self.surface_extractor = SurfaceExtractors[mc_algo]() else: self.volume_decoder = VanillaVolumeDecoder() self.surface_extractor = MCSurfaceExtractor() class ShapeVAE(VectsetVAE): """Shape VAE for 3D mesh generation from latent codes.""" _aliases = ["hy3dgen.shapegen.models.ShapeVAE"] def __init__( self, *, num_latents: int, embed_dim: int, width: int, heads: int, num_decoder_layers: int, num_encoder_layers: int = 8, pc_size: int = 5120, pc_sharpedge_size: int = 5120, point_feats: int = 3, downsample_ratio: int = 20, geo_decoder_downsample_ratio: int = 1, geo_decoder_mlp_expand_ratio: int = 4, geo_decoder_ln_post: bool = True, num_freqs: int = 8, include_pi: bool = True, qkv_bias: bool = True, qk_norm: bool = False, label_type: str = "binary", drop_path_rate: float = 0.0, scale_factor: float = 1.0, use_ln_post: bool = True, ckpt_path=None, ): super().__init__() self.geo_decoder_ln_post = geo_decoder_ln_post self.downsample_ratio = downsample_ratio self.fourier_embedder = FourierEmbedder( num_freqs=num_freqs, include_pi=include_pi ) self.post_kl = nn.Linear(embed_dim, width) self.transformer = Transformer( n_ctx=num_latents, width=width, layers=num_decoder_layers, heads=heads, qkv_bias=qkv_bias, qk_norm=qk_norm, drop_path_rate=drop_path_rate, ) self.geo_decoder = CrossAttentionDecoder( fourier_embedder=self.fourier_embedder, out_channels=1, num_latents=num_latents, mlp_expand_ratio=geo_decoder_mlp_expand_ratio, downsample_ratio=geo_decoder_downsample_ratio, enable_ln_post=self.geo_decoder_ln_post, width=width // geo_decoder_downsample_ratio, heads=heads // geo_decoder_downsample_ratio, qkv_bias=qkv_bias, qk_norm=qk_norm, label_type=label_type, ) self.scale_factor = scale_factor self.latent_shape = (num_latents, embed_dim) def forward(self, latents): latents = self.post_kl(latents) latents = self.transformer(latents) return latents def decode(self, latents): """Decode latents to features.""" latents = self.post_kl(latents) latents = self.transformer(latents) return latents # Entry class for model registry EntryClass = ShapeVAE ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from diffusers # Copyright 2024 The Hunyuan Team, The HuggingFace Team and The sglang-diffusion Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from sglang.multimodal_gen.configs.models.vaes import HunyuanVAEConfig from sglang.multimodal_gen.runtime.layers.activation import get_act_fn from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE def prepare_causal_attention_mask( num_frames: int, height_width: int, dtype: torch.dtype, device: torch.device, batch_size: int | None = None, ) -> torch.Tensor: indices = torch.arange(1, num_frames + 1, dtype=torch.int32, device=device) indices_blocks = indices.repeat_interleave(height_width) x, y = torch.meshgrid(indices_blocks, indices_blocks, indexing="xy") mask = torch.where(x <= y, 0, -float("inf")).to(dtype=dtype) if batch_size is not None: mask = mask.unsqueeze(0).expand(batch_size, -1, -1) return mask class HunyuanVAEAttention(nn.Module): def __init__( self, in_channels, heads, dim_head, eps, norm_num_groups, bias ) -> None: super().__init__() self.in_channels = in_channels self.heads = heads self.dim_head = dim_head self.eps = eps self.norm_num_groups = norm_num_groups self.bias = bias inner_dim = heads * dim_head # Define the projection layers self.to_q = nn.Linear(in_channels, inner_dim, bias=bias) self.to_k = nn.Linear(in_channels, inner_dim, bias=bias) self.to_v = nn.Linear(in_channels, inner_dim, bias=bias) self.to_out = nn.Sequential(nn.Linear(inner_dim, in_channels, bias=bias)) # Optional normalization layers self.group_norm = nn.GroupNorm( norm_num_groups, in_channels, eps=eps, affine=True ) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None ) -> torch.Tensor: residual = hidden_states batch_size, sequence_length, _ = hidden_states.shape hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) # Project to query, key, value query = self.to_q(hidden_states) key = self.to_k(hidden_states) value = self.to_v(hidden_states) # Reshape for multi-head attention head_dim = self.dim_head query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) # Perform scaled dot-product attention hidden_states = F.scaled_dot_product_attention( query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False ) # Reshape back hidden_states = hidden_states.transpose(1, 2).reshape( batch_size, -1, self.heads * head_dim ) hidden_states = hidden_states.to(query.dtype) # Linear projection hidden_states = self.to_out(hidden_states) # Residual connection and rescale hidden_states = hidden_states + residual return hidden_states class HunyuanVideoCausalConv3d(nn.Module): def __init__( self, in_channels: int, out_channels: int, kernel_size: int | tuple[int, int, int] = 3, stride: int | tuple[int, int, int] = 1, padding: int | tuple[int, int, int] = 0, dilation: int | tuple[int, int, int] = 1, bias: bool = True, pad_mode: str = "replicate", ) -> None: super().__init__() kernel_size = ( (kernel_size, kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size ) self.pad_mode = pad_mode self.time_causal_padding = ( kernel_size[0] // 2, kernel_size[0] // 2, kernel_size[1] // 2, kernel_size[1] // 2, kernel_size[2] - 1, 0, ) self.conv = nn.Conv3d( in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = F.pad( hidden_states, self.time_causal_padding, mode=self.pad_mode ) return self.conv(hidden_states) class HunyuanVideoUpsampleCausal3D(nn.Module): def __init__( self, in_channels: int, out_channels: int | None = None, kernel_size: int = 3, stride: int = 1, bias: bool = True, upsample_factor: tuple[int, ...] = (2, 2, 2), ) -> None: super().__init__() out_channels = out_channels or in_channels self.upsample_factor = upsample_factor self.conv = HunyuanVideoCausalConv3d( in_channels, out_channels, kernel_size, stride, bias=bias ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_frames = hidden_states.size(2) first_frame, other_frames = hidden_states.split((1, num_frames - 1), dim=2) first_frame = F.interpolate( first_frame.squeeze(2), scale_factor=self.upsample_factor[1:], mode="nearest", ).unsqueeze(2) if num_frames > 1: # See: https://github.com/pytorch/pytorch/issues/81665 # Unless you have a version of pytorch where non-contiguous implementation of F.interpolate # is fixed, this will raise either a runtime error, or fail silently with bad outputs. # If you are encountering an error here, make sure to try running encoding/decoding with # `vae.enable_tiling()` first. If that doesn't work, open an issue at: # https://github.com/huggingface/diffusers/issues other_frames = other_frames.contiguous() other_frames = F.interpolate( other_frames, scale_factor=self.upsample_factor, mode="nearest" ) hidden_states = torch.cat((first_frame, other_frames), dim=2) else: hidden_states = first_frame hidden_states = self.conv(hidden_states) return hidden_states class HunyuanVideoDownsampleCausal3D(nn.Module): def __init__( self, channels: int, out_channels: int | None = None, padding: int = 1, kernel_size: int = 3, bias: bool = True, stride=2, ) -> None: super().__init__() out_channels = out_channels or channels self.conv = HunyuanVideoCausalConv3d( channels, out_channels, kernel_size, stride, padding, bias=bias ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.conv(hidden_states) return hidden_states class HunyuanVideoResnetBlockCausal3D(nn.Module): def __init__( self, in_channels: int, out_channels: int | None = None, dropout: float = 0.0, groups: int = 32, eps: float = 1e-6, non_linearity: str = "silu", ) -> None: super().__init__() out_channels = out_channels or in_channels self.nonlinearity = get_act_fn(non_linearity) self.norm1 = nn.GroupNorm(groups, in_channels, eps=eps, affine=True) self.conv1 = HunyuanVideoCausalConv3d(in_channels, out_channels, 3, 1, 0) self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps, affine=True) self.dropout = nn.Dropout(dropout) self.conv2 = HunyuanVideoCausalConv3d(out_channels, out_channels, 3, 1, 0) self.conv_shortcut = None if in_channels != out_channels: self.conv_shortcut = HunyuanVideoCausalConv3d( in_channels, out_channels, 1, 1, 0 ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = hidden_states.contiguous() residual = hidden_states hidden_states = self.norm1(hidden_states) hidden_states = self.nonlinearity(hidden_states) hidden_states = self.conv1(hidden_states) hidden_states = self.norm2(hidden_states) hidden_states = self.nonlinearity(hidden_states) hidden_states = self.dropout(hidden_states) hidden_states = self.conv2(hidden_states) if self.conv_shortcut is not None: residual = self.conv_shortcut(residual) hidden_states = hidden_states + residual return hidden_states class HunyuanVideoMidBlock3D(nn.Module): def __init__( self, in_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, resnet_act_fn: str = "silu", resnet_groups: int = 32, add_attention: bool = True, attention_head_dim: int = 1, ) -> None: super().__init__() resnet_groups = ( resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) ) self.add_attention = add_attention # There is always at least one resnet resnets = [ HunyuanVideoResnetBlockCausal3D( in_channels=in_channels, out_channels=in_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, non_linearity=resnet_act_fn, ) ] attentions: list[HunyuanVAEAttention | None] = [] for _ in range(num_layers): if self.add_attention: attentions.append( HunyuanVAEAttention( in_channels, heads=in_channels // attention_head_dim, dim_head=attention_head_dim, eps=resnet_eps, norm_num_groups=resnet_groups, bias=True, ) ) else: attentions.append(None) resnets.append( HunyuanVideoResnetBlockCausal3D( in_channels=in_channels, out_channels=in_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, non_linearity=resnet_act_fn, ) ) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) self.gradient_checkpointing = False def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( self.resnets[0], hidden_states ) for attn, resnet in zip(self.attentions, self.resnets[1:], strict=True): if attn is not None: batch_size, num_channels, num_frames, height, width = ( hidden_states.shape ) hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3) attention_mask = prepare_causal_attention_mask( num_frames, height * width, hidden_states.dtype, hidden_states.device, batch_size=batch_size, ) hidden_states = attn(hidden_states, attention_mask=attention_mask) hidden_states = hidden_states.unflatten( 1, (num_frames, height, width) ).permute(0, 4, 1, 2, 3) hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) else: hidden_states = self.resnets[0](hidden_states) for attn, resnet in zip(self.attentions, self.resnets[1:], strict=True): if attn is not None: batch_size, num_channels, num_frames, height, width = ( hidden_states.shape ) hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3) attention_mask = prepare_causal_attention_mask( num_frames, height * width, hidden_states.dtype, hidden_states.device, batch_size=batch_size, ) hidden_states = attn(hidden_states, attention_mask=attention_mask) hidden_states = hidden_states.unflatten( 1, (num_frames, height, width) ).permute(0, 4, 1, 2, 3) hidden_states = resnet(hidden_states) return hidden_states class HunyuanVideoDownBlock3D(nn.Module): def __init__( self, in_channels: int, out_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, resnet_act_fn: str = "silu", resnet_groups: int = 32, add_downsample: bool = True, downsample_stride: tuple[int, ...] | int = 2, downsample_padding: int = 1, ) -> None: super().__init__() resnets = [] for i in range(num_layers): in_channels = in_channels if i == 0 else out_channels resnets.append( HunyuanVideoResnetBlockCausal3D( in_channels=in_channels, out_channels=out_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, non_linearity=resnet_act_fn, ) ) self.resnets = nn.ModuleList(resnets) if add_downsample: self.downsamplers = nn.ModuleList( [ HunyuanVideoDownsampleCausal3D( out_channels, out_channels=out_channels, padding=downsample_padding, stride=downsample_stride, ) ] ) else: self.downsamplers = None self.gradient_checkpointing = False def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if torch.is_grad_enabled() and self.gradient_checkpointing: for resnet in self.resnets: hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) else: for resnet in self.resnets: hidden_states = resnet(hidden_states) if self.downsamplers is not None: for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states) return hidden_states class HunyuanVideoUpBlock3D(nn.Module): def __init__( self, in_channels: int, out_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, resnet_act_fn: str = "silu", resnet_groups: int = 32, add_upsample: bool = True, upsample_scale_factor: tuple[int, ...] = (2, 2, 2), ) -> None: super().__init__() resnets = [] for i in range(num_layers): input_channels = in_channels if i == 0 else out_channels resnets.append( HunyuanVideoResnetBlockCausal3D( in_channels=input_channels, out_channels=out_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, non_linearity=resnet_act_fn, ) ) self.resnets = nn.ModuleList(resnets) if add_upsample: self.upsamplers = nn.ModuleList( [ HunyuanVideoUpsampleCausal3D( out_channels, out_channels=out_channels, upsample_factor=upsample_scale_factor, ) ] ) else: self.upsamplers = None self.gradient_checkpointing = False def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if torch.is_grad_enabled() and self.gradient_checkpointing: for resnet in self.resnets: hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) else: for resnet in self.resnets: hidden_states = resnet(hidden_states) if self.upsamplers is not None: for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states) return hidden_states class HunyuanVideoEncoder3D(nn.Module): r""" Causal encoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603). """ def __init__( self, in_channels: int = 3, out_channels: int = 3, down_block_types: tuple[str, ...] = ( "HunyuanVideoDownBlock3D", "HunyuanVideoDownBlock3D", "HunyuanVideoDownBlock3D", "HunyuanVideoDownBlock3D", ), block_out_channels: tuple[int, ...] = (128, 256, 512, 512), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", double_z: bool = True, mid_block_add_attention=True, temporal_compression_ratio: int = 4, spatial_compression_ratio: int = 8, ) -> None: super().__init__() self.conv_in = HunyuanVideoCausalConv3d( in_channels, block_out_channels[0], kernel_size=3, stride=1 ) self.mid_block: HunyuanVideoMidBlock3D | None = None self.down_blocks = nn.ModuleList([]) output_channel = block_out_channels[0] for i, down_block_type in enumerate(down_block_types): if down_block_type != "HunyuanVideoDownBlock3D": raise ValueError(f"Unsupported down_block_type: {down_block_type}") input_channel = output_channel output_channel = block_out_channels[i] is_final_block = i == len(block_out_channels) - 1 num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio)) num_time_downsample_layers = int(np.log2(temporal_compression_ratio)) if temporal_compression_ratio == 4: add_spatial_downsample = bool(i < num_spatial_downsample_layers) add_time_downsample = bool( i >= (len(block_out_channels) - 1 - num_time_downsample_layers) and not is_final_block ) elif temporal_compression_ratio == 8: add_spatial_downsample = bool(i < num_spatial_downsample_layers) add_time_downsample = bool(i < num_time_downsample_layers) else: raise ValueError( f"Unsupported time_compression_ratio: {temporal_compression_ratio}" ) downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1) downsample_stride_T = (2,) if add_time_downsample else (1,) downsample_stride = tuple(downsample_stride_T + downsample_stride_HW) down_block = HunyuanVideoDownBlock3D( num_layers=layers_per_block, in_channels=input_channel, out_channels=output_channel, add_downsample=bool(add_spatial_downsample or add_time_downsample), resnet_eps=1e-6, resnet_act_fn=act_fn, resnet_groups=norm_num_groups, downsample_stride=downsample_stride, downsample_padding=0, ) self.down_blocks.append(down_block) self.mid_block = HunyuanVideoMidBlock3D( in_channels=block_out_channels[-1], resnet_eps=1e-6, resnet_act_fn=act_fn, attention_head_dim=block_out_channels[-1], resnet_groups=norm_num_groups, add_attention=mid_block_add_attention, ) self.conv_norm_out = nn.GroupNorm( num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6 ) self.conv_act = nn.SiLU() conv_out_channels = 2 * out_channels if double_z else out_channels self.conv_out = HunyuanVideoCausalConv3d( block_out_channels[-1], conv_out_channels, kernel_size=3 ) self.gradient_checkpointing = False def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.conv_in(hidden_states) if torch.is_grad_enabled() and self.gradient_checkpointing: for down_block in self.down_blocks: hidden_states = self._gradient_checkpointing_func( down_block, hidden_states ) hidden_states = self._gradient_checkpointing_func( self.mid_block, hidden_states ) else: for down_block in self.down_blocks: hidden_states = down_block(hidden_states) assert self.mid_block is not None hidden_states = self.mid_block(hidden_states) hidden_states = self.conv_norm_out(hidden_states) hidden_states = self.conv_act(hidden_states) hidden_states = self.conv_out(hidden_states) return hidden_states class HunyuanVideoDecoder3D(nn.Module): r""" Causal decoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603). """ def __init__( self, in_channels: int = 3, out_channels: int = 3, up_block_types: tuple[str, ...] = ( "HunyuanVideoUpBlock3D", "HunyuanVideoUpBlock3D", "HunyuanVideoUpBlock3D", "HunyuanVideoUpBlock3D", ), block_out_channels: tuple[int, ...] = (128, 256, 512, 512), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", mid_block_add_attention=True, time_compression_ratio: int = 4, spatial_compression_ratio: int = 8, ): super().__init__() self.layers_per_block = layers_per_block self.conv_in = HunyuanVideoCausalConv3d( in_channels, block_out_channels[-1], kernel_size=3, stride=1 ) self.up_blocks = nn.ModuleList([]) # mid self.mid_block = HunyuanVideoMidBlock3D( in_channels=block_out_channels[-1], resnet_eps=1e-6, resnet_act_fn=act_fn, attention_head_dim=block_out_channels[-1], resnet_groups=norm_num_groups, add_attention=mid_block_add_attention, ) # up reversed_block_out_channels = list(reversed(block_out_channels)) output_channel = reversed_block_out_channels[0] for i, up_block_type in enumerate(up_block_types): if up_block_type != "HunyuanVideoUpBlock3D": raise ValueError(f"Unsupported up_block_type: {up_block_type}") prev_output_channel = output_channel output_channel = reversed_block_out_channels[i] is_final_block = i == len(block_out_channels) - 1 num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio)) num_time_upsample_layers = int(np.log2(time_compression_ratio)) if time_compression_ratio == 4: add_spatial_upsample = bool(i < num_spatial_upsample_layers) add_time_upsample = bool( i >= len(block_out_channels) - 1 - num_time_upsample_layers and not is_final_block ) else: raise ValueError( f"Unsupported time_compression_ratio: {time_compression_ratio}" ) upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1) upsample_scale_factor_T = (2,) if add_time_upsample else (1,) upsample_scale_factor = tuple( upsample_scale_factor_T + upsample_scale_factor_HW ) up_block = HunyuanVideoUpBlock3D( num_layers=self.layers_per_block + 1, in_channels=prev_output_channel, out_channels=output_channel, add_upsample=bool(add_spatial_upsample or add_time_upsample), upsample_scale_factor=upsample_scale_factor, resnet_eps=1e-6, resnet_act_fn=act_fn, resnet_groups=norm_num_groups, ) self.up_blocks.append(up_block) prev_output_channel = output_channel # out self.conv_norm_out = nn.GroupNorm( num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6 ) self.conv_act = nn.SiLU() self.conv_out = HunyuanVideoCausalConv3d( block_out_channels[0], out_channels, kernel_size=3 ) self.gradient_checkpointing = False def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.conv_in(hidden_states) if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( self.mid_block, hidden_states ) for up_block in self.up_blocks: hidden_states = self._gradient_checkpointing_func( up_block, hidden_states ) else: hidden_states = self.mid_block(hidden_states) for up_block in self.up_blocks: hidden_states = up_block(hidden_states) # post-process hidden_states = self.conv_norm_out(hidden_states) hidden_states = self.conv_act(hidden_states) hidden_states = self.conv_out(hidden_states) return hidden_states class AutoencoderKLHunyuanVideo(ParallelTiledVAE): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Introduced in [HunyuanVideo](https://huggingface.co/papers/2412.03603). This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented for all models (such as downloading or saving). """ _supports_gradient_checkpointing = True def __init__( self, config: HunyuanVAEConfig, ) -> None: nn.Module.__init__(self) ParallelTiledVAE.__init__(self, config) # TODO(will): only pass in config. We do this by manually defining a # config for hunyuan vae self.block_out_channels = config.block_out_channels if config.load_encoder: self.encoder = HunyuanVideoEncoder3D( in_channels=config.in_channels, out_channels=config.latent_channels, down_block_types=config.down_block_types, block_out_channels=config.block_out_channels, layers_per_block=config.layers_per_block, norm_num_groups=config.norm_num_groups, act_fn=config.act_fn, double_z=True, mid_block_add_attention=config.mid_block_add_attention, temporal_compression_ratio=config.temporal_compression_ratio, spatial_compression_ratio=config.spatial_compression_ratio, ) self.quant_conv = nn.Conv3d( 2 * config.latent_channels, 2 * config.latent_channels, kernel_size=1 ) if config.load_decoder: self.decoder = HunyuanVideoDecoder3D( in_channels=config.latent_channels, out_channels=config.out_channels, up_block_types=config.up_block_types, block_out_channels=config.block_out_channels, layers_per_block=config.layers_per_block, norm_num_groups=config.norm_num_groups, act_fn=config.act_fn, time_compression_ratio=config.temporal_compression_ratio, spatial_compression_ratio=config.spatial_compression_ratio, mid_block_add_attention=config.mid_block_add_attention, ) self.post_quant_conv = nn.Conv3d( config.latent_channels, config.latent_channels, kernel_size=1 ) def _encode(self, x: torch.Tensor) -> torch.Tensor: x = self.encoder(x) enc = self.quant_conv(x) return enc def _decode(self, z: torch.Tensor) -> torch.Tensor: z = self.post_quant_conv(z) dec = self.decoder(z) return dec def forward( self, sample: torch.Tensor, sample_posterior: bool = False, generator: torch.Generator | None = None, ) -> torch.Tensor: r""" Args: sample (`torch.Tensor`): Input sample. sample_posterior (`bool`, *optional*, defaults to `False`): Whether to sample from the posterior. """ x = sample posterior = self.encode(x).latent_dist if sample_posterior: z = posterior.sample(generator=generator) else: z = posterior.mode() dec = self.decode(z) return dec EntryClass = AutoencoderKLHunyuanVideo ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_audio.py ================================================ from typing import Optional, Tuple, Union import torch import torch.nn.functional as F from diffusers.models.autoencoders.vae import ( DecoderOutput, DiagonalGaussianDistribution, ) from diffusers.models.modeling_outputs import AutoencoderKLOutput from torch import nn from sglang.multimodal_gen.configs.models.vaes.ltx_audio import LTXAudioVAEConfig from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE LATENT_DOWNSAMPLE_FACTOR = 4 class LTX2AudioCausalConv2d(nn.Module): """ A causal 2D convolution that pads asymmetrically along the causal axis. """ def __init__( self, in_channels: int, out_channels: int, kernel_size: Union[int, Tuple[int, int]], stride: int = 1, dilation: Union[int, Tuple[int, int]] = 1, groups: int = 1, bias: bool = True, causality_axis: str = "height", ) -> None: super().__init__() self.causality_axis = causality_axis kernel_size = ( (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size ) dilation = (dilation, dilation) if isinstance(dilation, int) else dilation pad_h = (kernel_size[0] - 1) * dilation[0] pad_w = (kernel_size[1] - 1) * dilation[1] if self.causality_axis == "none": padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2) elif self.causality_axis in {"width", "width-compatibility"}: padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2) elif self.causality_axis == "height": padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0) else: raise ValueError(f"Invalid causality_axis: {causality_axis}") self.padding = padding self.conv = nn.Conv2d( in_channels, out_channels, kernel_size, stride=stride, padding=0, dilation=dilation, groups=groups, bias=bias, ) def forward(self, x: torch.Tensor) -> torch.Tensor: x = F.pad(x, self.padding) return self.conv(x) class LTX2AudioPixelNorm(nn.Module): """ Per-pixel (per-location) RMS normalization layer. """ def __init__(self, dim: int = 1, eps: float = 1e-8) -> None: super().__init__() self.dim = dim self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True) rms = torch.sqrt(mean_sq + self.eps) return x / rms class LTX2AudioAttnBlock(nn.Module): def __init__( self, in_channels: int, norm_type: str = "group", ) -> None: super().__init__() self.in_channels = in_channels if norm_type == "group": self.norm = nn.GroupNorm( num_groups=32, num_channels=in_channels, eps=1e-6, affine=True ) elif norm_type == "pixel": self.norm = LTX2AudioPixelNorm(dim=1, eps=1e-6) else: raise ValueError(f"Invalid normalization type: {norm_type}") self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) self.proj_out = nn.Conv2d( in_channels, in_channels, kernel_size=1, stride=1, padding=0 ) def forward(self, x: torch.Tensor) -> torch.Tensor: h_ = self.norm(x) q = self.q(h_) k = self.k(h_) v = self.v(h_) batch, channels, height, width = q.shape q = q.reshape(batch, channels, height * width).permute(0, 2, 1).contiguous() k = k.reshape(batch, channels, height * width).contiguous() attn = torch.bmm(q, k) * (int(channels) ** (-0.5)) attn = torch.nn.functional.softmax(attn, dim=2) v = v.reshape(batch, channels, height * width) attn = attn.permute(0, 2, 1).contiguous() h_ = torch.bmm(v, attn).reshape(batch, channels, height, width) h_ = self.proj_out(h_) return x + h_ class LTX2AudioResnetBlock(nn.Module): def __init__( self, in_channels: int, out_channels: Optional[int] = None, conv_shortcut: bool = False, dropout: float = 0.0, temb_channels: int = 512, norm_type: str = "group", causality_axis: str = "height", ) -> None: super().__init__() self.causality_axis = causality_axis if ( self.causality_axis is not None and self.causality_axis != "none" and norm_type == "group" ): raise ValueError("Causal ResnetBlock with GroupNorm is not supported.") self.in_channels = in_channels out_channels = in_channels if out_channels is None else out_channels self.out_channels = out_channels self.use_conv_shortcut = conv_shortcut if norm_type == "group": self.norm1 = nn.GroupNorm( num_groups=32, num_channels=in_channels, eps=1e-6, affine=True ) elif norm_type == "pixel": self.norm1 = LTX2AudioPixelNorm(dim=1, eps=1e-6) else: raise ValueError(f"Invalid normalization type: {norm_type}") self.non_linearity = nn.SiLU() if causality_axis is not None: self.conv1 = LTX2AudioCausalConv2d( in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis, ) else: self.conv1 = nn.Conv2d( in_channels, out_channels, kernel_size=3, stride=1, padding=1 ) if temb_channels > 0: self.temb_proj = nn.Linear(temb_channels, out_channels) if norm_type == "group": self.norm2 = nn.GroupNorm( num_groups=32, num_channels=out_channels, eps=1e-6, affine=True ) elif norm_type == "pixel": self.norm2 = LTX2AudioPixelNorm(dim=1, eps=1e-6) else: raise ValueError(f"Invalid normalization type: {norm_type}") self.dropout = nn.Dropout(dropout) if causality_axis is not None: self.conv2 = LTX2AudioCausalConv2d( out_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis, ) else: self.conv2 = nn.Conv2d( out_channels, out_channels, kernel_size=3, stride=1, padding=1 ) if self.in_channels != self.out_channels: if self.use_conv_shortcut: if causality_axis is not None: self.conv_shortcut = LTX2AudioCausalConv2d( in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis, ) else: self.conv_shortcut = nn.Conv2d( in_channels, out_channels, kernel_size=3, stride=1, padding=1 ) else: if causality_axis is not None: self.nin_shortcut = LTX2AudioCausalConv2d( in_channels, out_channels, kernel_size=1, stride=1, causality_axis=causality_axis, ) else: self.nin_shortcut = nn.Conv2d( in_channels, out_channels, kernel_size=1, stride=1, padding=0 ) def forward( self, x: torch.Tensor, temb: Optional[torch.Tensor] = None ) -> torch.Tensor: h = self.norm1(x) h = self.non_linearity(h) h = self.conv1(h) if temb is not None: h = h + self.temb_proj(self.non_linearity(temb))[:, :, None, None] h = self.norm2(h) h = self.non_linearity(h) h = self.dropout(h) h = self.conv2(h) if self.in_channels != self.out_channels: x = ( self.conv_shortcut(x) if self.use_conv_shortcut else self.nin_shortcut(x) ) return x + h class LTX2AudioDownsample(nn.Module): def __init__( self, in_channels: int, with_conv: bool, causality_axis: Optional[str] = "height", ) -> None: super().__init__() self.with_conv = with_conv self.causality_axis = causality_axis if self.with_conv: self.conv = torch.nn.Conv2d( in_channels, in_channels, kernel_size=3, stride=2, padding=0 ) def forward(self, x: torch.Tensor) -> torch.Tensor: if self.with_conv: # Padding tuple is in the order: (left, right, top, bottom). if self.causality_axis == "none": pad = (0, 1, 0, 1) elif self.causality_axis == "width": pad = (2, 0, 0, 1) elif self.causality_axis == "height": pad = (0, 1, 2, 0) elif self.causality_axis == "width-compatibility": pad = (1, 0, 0, 1) else: raise ValueError( f"Invalid `causality_axis` {self.causality_axis}; supported values are `none`, `width`, `height`," f" and `width-compatibility`." ) x = F.pad(x, pad, mode="constant", value=0) x = self.conv(x) else: # with_conv=False implies that causality_axis is "none" x = F.avg_pool2d(x, kernel_size=2, stride=2) return x class LTX2AudioUpsample(nn.Module): def __init__( self, in_channels: int, with_conv: bool, causality_axis: Optional[str] = "height", ) -> None: super().__init__() self.with_conv = with_conv self.causality_axis = causality_axis if self.with_conv: if causality_axis is not None: self.conv = LTX2AudioCausalConv2d( in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis, ) else: self.conv = nn.Conv2d( in_channels, in_channels, kernel_size=3, stride=1, padding=1 ) def forward(self, x: torch.Tensor) -> torch.Tensor: x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") if self.with_conv: x = self.conv(x) if self.causality_axis is None or self.causality_axis == "none": pass elif self.causality_axis == "height": x = x[:, :, 1:, :] elif self.causality_axis == "width": x = x[:, :, :, 1:] elif self.causality_axis == "width-compatibility": pass else: raise ValueError(f"Invalid causality_axis: {self.causality_axis}") return x class LTX2AudioAudioPatchifier: """ Patchifier for spectrogram/audio latents. """ def __init__( self, patch_size: int, sample_rate: int = 16000, hop_length: int = 160, audio_latent_downsample_factor: int = 4, is_causal: bool = True, ): self.hop_length = hop_length self.sample_rate = sample_rate self.audio_latent_downsample_factor = audio_latent_downsample_factor self.is_causal = is_causal self._patch_size = (1, patch_size, patch_size) def patchify(self, audio_latents: torch.Tensor) -> torch.Tensor: batch, channels, time, freq = audio_latents.shape return audio_latents.permute(0, 2, 1, 3).reshape(batch, time, channels * freq) def unpatchify( self, audio_latents: torch.Tensor, channels: int, mel_bins: int ) -> torch.Tensor: batch, time, _ = audio_latents.shape return audio_latents.view(batch, time, channels, mel_bins).permute(0, 2, 1, 3) @property def patch_size(self) -> Tuple[int, int, int]: return self._patch_size class LTX2AudioEncoder(nn.Module): def __init__( self, base_channels: int = 128, output_channels: int = 1, num_res_blocks: int = 2, attn_resolutions: Optional[Tuple[int, ...]] = None, in_channels: int = 2, resolution: int = 256, latent_channels: int = 8, ch_mult: Tuple[int, ...] = (1, 2, 4), norm_type: str = "group", causality_axis: Optional[str] = "width", dropout: float = 0.0, mid_block_add_attention: bool = False, sample_rate: int = 16000, mel_hop_length: int = 160, is_causal: bool = True, mel_bins: Optional[int] = 64, double_z: bool = True, ): super().__init__() self.sample_rate = sample_rate self.mel_hop_length = mel_hop_length self.is_causal = is_causal self.mel_bins = mel_bins self.base_channels = base_channels self.temb_ch = 0 self.num_resolutions = len(ch_mult) self.num_res_blocks = num_res_blocks self.resolution = resolution self.in_channels = in_channels self.out_ch = output_channels self.give_pre_end = False self.tanh_out = False self.norm_type = norm_type self.latent_channels = latent_channels self.channel_multipliers = ch_mult self.attn_resolutions = attn_resolutions self.causality_axis = causality_axis base_block_channels = base_channels base_resolution = resolution self.z_shape = (1, latent_channels, base_resolution, base_resolution) if self.causality_axis is not None: self.conv_in = LTX2AudioCausalConv2d( in_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis, ) else: self.conv_in = nn.Conv2d( in_channels, base_block_channels, kernel_size=3, stride=1, padding=1 ) self.down = nn.ModuleList() block_in = base_block_channels curr_res = self.resolution for level in range(self.num_resolutions): stage = nn.Module() stage.block = nn.ModuleList() stage.attn = nn.ModuleList() block_out = self.base_channels * self.channel_multipliers[level] for _ in range(self.num_res_blocks): stage.block.append( LTX2AudioResnetBlock( in_channels=block_in, out_channels=block_out, temb_channels=self.temb_ch, dropout=dropout, norm_type=self.norm_type, causality_axis=self.causality_axis, ) ) block_in = block_out if self.attn_resolutions: if curr_res in self.attn_resolutions: stage.attn.append( LTX2AudioAttnBlock(block_in, norm_type=self.norm_type) ) if level != self.num_resolutions - 1: stage.downsample = LTX2AudioDownsample( block_in, True, causality_axis=self.causality_axis ) curr_res = curr_res // 2 self.down.append(stage) self.mid = nn.Module() self.mid.block_1 = LTX2AudioResnetBlock( in_channels=block_in, out_channels=block_in, temb_channels=self.temb_ch, dropout=dropout, norm_type=self.norm_type, causality_axis=self.causality_axis, ) if mid_block_add_attention: self.mid.attn_1 = LTX2AudioAttnBlock(block_in, norm_type=self.norm_type) else: self.mid.attn_1 = nn.Identity() self.mid.block_2 = LTX2AudioResnetBlock( in_channels=block_in, out_channels=block_in, temb_channels=self.temb_ch, dropout=dropout, norm_type=self.norm_type, causality_axis=self.causality_axis, ) final_block_channels = block_in z_channels = 2 * latent_channels if double_z else latent_channels if self.norm_type == "group": self.norm_out = nn.GroupNorm( num_groups=32, num_channels=final_block_channels, eps=1e-6, affine=True ) elif self.norm_type == "pixel": self.norm_out = LTX2AudioPixelNorm(dim=1, eps=1e-6) else: raise ValueError(f"Invalid normalization type: {self.norm_type}") self.non_linearity = nn.SiLU() if self.causality_axis is not None: self.conv_out = LTX2AudioCausalConv2d( final_block_channels, z_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis, ) else: self.conv_out = nn.Conv2d( final_block_channels, z_channels, kernel_size=3, stride=1, padding=1 ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # hidden_states expected shape: (batch_size, channels, time, num_mel_bins) hidden_states = self.conv_in(hidden_states) for level in range(self.num_resolutions): stage = self.down[level] for block_idx, block in enumerate(stage.block): hidden_states = block(hidden_states, temb=None) if stage.attn: hidden_states = stage.attn[block_idx](hidden_states) if level != self.num_resolutions - 1 and hasattr(stage, "downsample"): hidden_states = stage.downsample(hidden_states) hidden_states = self.mid.block_1(hidden_states, temb=None) hidden_states = self.mid.attn_1(hidden_states) hidden_states = self.mid.block_2(hidden_states, temb=None) hidden_states = self.norm_out(hidden_states) hidden_states = self.non_linearity(hidden_states) hidden_states = self.conv_out(hidden_states) return hidden_states class LTX2AudioDecoder(nn.Module): """ Symmetric decoder that reconstructs audio spectrograms from latent features. The decoder mirrors the encoder structure with configurable channel multipliers, attention resolutions, and causal convolutions. """ def __init__( self, base_channels: int = 128, output_channels: int = 1, num_res_blocks: int = 2, attn_resolutions: Optional[Tuple[int, ...]] = None, in_channels: int = 2, resolution: int = 256, latent_channels: int = 8, ch_mult: Tuple[int, ...] = (1, 2, 4), norm_type: str = "group", causality_axis: Optional[str] = "width", dropout: float = 0.0, mid_block_add_attention: bool = False, sample_rate: int = 16000, mel_hop_length: int = 160, is_causal: bool = True, mel_bins: Optional[int] = 64, ) -> None: super().__init__() self.sample_rate = sample_rate self.mel_hop_length = mel_hop_length self.is_causal = is_causal self.mel_bins = mel_bins self.patchifier = LTX2AudioAudioPatchifier( patch_size=1, audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR, sample_rate=sample_rate, hop_length=mel_hop_length, is_causal=is_causal, ) self.base_channels = base_channels self.temb_ch = 0 self.num_resolutions = len(ch_mult) self.num_res_blocks = num_res_blocks self.resolution = resolution self.in_channels = in_channels self.out_ch = output_channels self.give_pre_end = False self.tanh_out = False self.norm_type = norm_type self.latent_channels = latent_channels self.channel_multipliers = ch_mult self.attn_resolutions = attn_resolutions self.causality_axis = causality_axis base_block_channels = base_channels * self.channel_multipliers[-1] base_resolution = resolution // (2 ** (self.num_resolutions - 1)) self.z_shape = (1, latent_channels, base_resolution, base_resolution) if self.causality_axis is not None: self.conv_in = LTX2AudioCausalConv2d( latent_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis, ) else: self.conv_in = nn.Conv2d( latent_channels, base_block_channels, kernel_size=3, stride=1, padding=1 ) self.non_linearity = nn.SiLU() self.mid = nn.Module() self.mid.block_1 = LTX2AudioResnetBlock( in_channels=base_block_channels, out_channels=base_block_channels, temb_channels=self.temb_ch, dropout=dropout, norm_type=self.norm_type, causality_axis=self.causality_axis, ) if mid_block_add_attention: self.mid.attn_1 = LTX2AudioAttnBlock( base_block_channels, norm_type=self.norm_type ) else: self.mid.attn_1 = nn.Identity() self.mid.block_2 = LTX2AudioResnetBlock( in_channels=base_block_channels, out_channels=base_block_channels, temb_channels=self.temb_ch, dropout=dropout, norm_type=self.norm_type, causality_axis=self.causality_axis, ) self.up = nn.ModuleList() block_in = base_block_channels curr_res = self.resolution // (2 ** (self.num_resolutions - 1)) for level in reversed(range(self.num_resolutions)): stage = nn.Module() stage.block = nn.ModuleList() stage.attn = nn.ModuleList() block_out = self.base_channels * self.channel_multipliers[level] for _ in range(self.num_res_blocks + 1): stage.block.append( LTX2AudioResnetBlock( in_channels=block_in, out_channels=block_out, temb_channels=self.temb_ch, dropout=dropout, norm_type=self.norm_type, causality_axis=self.causality_axis, ) ) block_in = block_out if self.attn_resolutions: if curr_res in self.attn_resolutions: stage.attn.append( LTX2AudioAttnBlock(block_in, norm_type=self.norm_type) ) if level != 0: stage.upsample = LTX2AudioUpsample( block_in, True, causality_axis=self.causality_axis ) curr_res *= 2 self.up.insert(0, stage) final_block_channels = block_in if self.norm_type == "group": self.norm_out = nn.GroupNorm( num_groups=32, num_channels=final_block_channels, eps=1e-6, affine=True ) elif self.norm_type == "pixel": self.norm_out = LTX2AudioPixelNorm(dim=1, eps=1e-6) else: raise ValueError(f"Invalid normalization type: {self.norm_type}") if self.causality_axis is not None: self.conv_out = LTX2AudioCausalConv2d( final_block_channels, output_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis, ) else: self.conv_out = nn.Conv2d( final_block_channels, output_channels, kernel_size=3, stride=1, padding=1, ) def forward( self, sample: torch.Tensor, ) -> torch.Tensor: _, _, frames, mel_bins = sample.shape target_frames = frames * LATENT_DOWNSAMPLE_FACTOR if self.causality_axis is not None: target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1) target_channels = self.out_ch target_mel_bins = self.mel_bins if self.mel_bins is not None else mel_bins hidden_features = self.conv_in(sample) hidden_features = self.mid.block_1(hidden_features, temb=None) hidden_features = self.mid.attn_1(hidden_features) hidden_features = self.mid.block_2(hidden_features, temb=None) for level in reversed(range(self.num_resolutions)): stage = self.up[level] for block_idx, block in enumerate(stage.block): hidden_features = block(hidden_features, temb=None) if stage.attn: hidden_features = stage.attn[block_idx](hidden_features) if level != 0 and hasattr(stage, "upsample"): hidden_features = stage.upsample(hidden_features) if self.give_pre_end: return hidden_features hidden = self.norm_out(hidden_features) hidden = self.non_linearity(hidden) decoded_output = self.conv_out(hidden) decoded_output = torch.tanh(decoded_output) if self.tanh_out else decoded_output _, _, current_time, current_freq = decoded_output.shape target_time = target_frames target_freq = target_mel_bins decoded_output = decoded_output[ :, :target_channels, : min(current_time, target_time), : min(current_freq, target_freq), ] time_padding_needed = target_time - decoded_output.shape[2] freq_padding_needed = target_freq - decoded_output.shape[3] if time_padding_needed > 0 or freq_padding_needed > 0: padding = ( 0, max(freq_padding_needed, 0), 0, max(time_padding_needed, 0), ) decoded_output = F.pad(decoded_output, padding) decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq] return decoded_output class AutoencoderKLLTX2Audio(ParallelTiledVAE): r""" LTX2 audio VAE for encoding and decoding audio latent representations. """ _supports_gradient_checkpointing = False def __init__( self, config: LTXAudioVAEConfig, ) -> None: super().__init__(config=config) causality_axis = config.arch_config.causality_axis attn_resolutions = config.arch_config.attn_resolutions base_channels = config.arch_config.base_channels output_channels = config.arch_config.output_channels ch_mult = config.arch_config.ch_mult num_res_blocks = config.arch_config.num_res_blocks in_channels = config.arch_config.in_channels resolution = config.arch_config.resolution latent_channels = config.arch_config.latent_channels norm_type = config.arch_config.norm_type dropout = config.arch_config.dropout mid_block_add_attention = config.arch_config.mid_block_add_attention sample_rate = config.arch_config.sample_rate mel_hop_length = config.arch_config.mel_hop_length is_causal = config.arch_config.is_causal mel_bins = config.arch_config.mel_bins double_z = config.arch_config.double_z supported_causality_axes = {"none", "width", "height", "width-compatibility"} if causality_axis not in supported_causality_axes: raise ValueError( f"{causality_axis=} is not valid. Supported values: {supported_causality_axes}" ) attn_resolution_set = ( set(attn_resolutions) if attn_resolutions else attn_resolutions ) self.encoder = LTX2AudioEncoder( base_channels=base_channels, output_channels=output_channels, ch_mult=ch_mult, num_res_blocks=num_res_blocks, attn_resolutions=attn_resolution_set, in_channels=in_channels, resolution=resolution, latent_channels=latent_channels, norm_type=norm_type, causality_axis=causality_axis, dropout=dropout, mid_block_add_attention=mid_block_add_attention, sample_rate=sample_rate, mel_hop_length=mel_hop_length, is_causal=is_causal, mel_bins=mel_bins, double_z=double_z, ) self.decoder = LTX2AudioDecoder( base_channels=base_channels, output_channels=output_channels, ch_mult=ch_mult, num_res_blocks=num_res_blocks, attn_resolutions=attn_resolution_set, in_channels=in_channels, resolution=resolution, latent_channels=latent_channels, norm_type=norm_type, causality_axis=causality_axis, dropout=dropout, mid_block_add_attention=mid_block_add_attention, sample_rate=sample_rate, mel_hop_length=mel_hop_length, is_causal=is_causal, mel_bins=mel_bins, ) # Per-channel statistics for normalizing and denormalizing the latent representation. This statistics is computed over # the entire dataset and stored in model's checkpoint under AudioVAE state_dict latents_std = torch.zeros((base_channels,)) latents_mean = torch.ones((base_channels,)) self.register_buffer("latents_mean", latents_mean, persistent=True) self.register_buffer("latents_std", latents_std, persistent=True) # TODO: confirm whether the mel compression ratio below is correct self.mel_compression_ratio = LATENT_DOWNSAMPLE_FACTOR self.use_slicing = False def _encode(self, x: torch.Tensor) -> torch.Tensor: return self.encoder(x) def encode(self, x: torch.Tensor, return_dict: bool = True): if self.use_slicing and x.shape[0] > 1: encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] h = torch.cat(encoded_slices) else: h = self._encode(x) posterior = DiagonalGaussianDistribution(h) if not return_dict: return (posterior,) return AutoencoderKLOutput(latent_dist=posterior) def _decode(self, z: torch.Tensor) -> torch.Tensor: return self.decoder(z) def decode( self, z: torch.Tensor, return_dict: bool = True ) -> Union[DecoderOutput, torch.Tensor]: if self.use_slicing and z.shape[0] > 1: decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)] decoded = torch.cat(decoded_slices) else: decoded = self._decode(z) if not return_dict: return (decoded,) return DecoderOutput(sample=decoded) def forward( self, sample: torch.Tensor, sample_posterior: bool = False, return_dict: bool = True, generator: Optional[torch.Generator] = None, ) -> Union[DecoderOutput, torch.Tensor]: posterior = self.encode(sample).latent_dist if sample_posterior: z = posterior.sample(generator=generator) else: z = posterior.mode() dec = self.decode(z) if not return_dict: return (dec.sample,) return dec EntryClass = AutoencoderKLLTX2Audio ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py ================================================ from typing import Optional, Tuple, Union import torch import torch.nn as nn from diffusers.models.activations import get_activation from diffusers.models.autoencoders.vae import ( DecoderOutput, DiagonalGaussianDistribution, ) from diffusers.models.embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings from diffusers.models.modeling_outputs import AutoencoderKLOutput from sglang.multimodal_gen.configs.models.vaes.ltx_video import LTXVideoVAEConfig from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE class PerChannelRMSNorm(nn.Module): """ Per-pixel (per-location) RMS normalization layer. For each element along the chosen dimension, this layer normalizes the tensor by the root-mean-square of its values across that dimension: y = x / sqrt(mean(x^2, dim=dim, keepdim=True) + eps) """ def __init__(self, channel_dim: int = 1, eps: float = 1e-8) -> None: """ Args: dim: Dimension along which to compute the RMS (typically channels). eps: Small constant added for numerical stability. """ super().__init__() self.channel_dim = channel_dim self.eps = eps def forward( self, x: torch.Tensor, channel_dim: Optional[int] = None ) -> torch.Tensor: """ Apply RMS normalization along the configured dimension. """ channel_dim = channel_dim or self.channel_dim # Compute mean of squared values along `dim`, keep dimensions for broadcasting. mean_sq = torch.mean(x**2, dim=self.channel_dim, keepdim=True) # Normalize by the root-mean-square (RMS). rms = torch.sqrt(mean_sq + self.eps) return x / rms # Like LTXCausalConv3d, but whether causal inference is performed can be specified at runtime class LTX2VideoCausalConv3d(nn.Module): def __init__( self, in_channels: int, out_channels: int, kernel_size: Union[int, Tuple[int, int, int]] = 3, stride: Union[int, Tuple[int, int, int]] = 1, dilation: Union[int, Tuple[int, int, int]] = 1, groups: int = 1, spatial_padding_mode: str = "zeros", ): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.kernel_size = ( kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size, kernel_size) ) dilation = dilation if isinstance(dilation, tuple) else (dilation, 1, 1) stride = stride if isinstance(stride, tuple) else (stride, stride, stride) height_pad = self.kernel_size[1] // 2 width_pad = self.kernel_size[2] // 2 padding = (0, height_pad, width_pad) self.conv = nn.Conv3d( in_channels, out_channels, self.kernel_size, stride=stride, dilation=dilation, groups=groups, padding=padding, padding_mode=spatial_padding_mode, ) def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Tensor: time_kernel_size = self.kernel_size[0] if causal: pad_left = hidden_states[:, :, :1, :, :].repeat( (1, 1, time_kernel_size - 1, 1, 1) ) hidden_states = torch.concatenate([pad_left, hidden_states], dim=2) else: pad_left = hidden_states[:, :, :1, :, :].repeat( (1, 1, (time_kernel_size - 1) // 2, 1, 1) ) pad_right = hidden_states[:, :, -1:, :, :].repeat( (1, 1, (time_kernel_size - 1) // 2, 1, 1) ) hidden_states = torch.concatenate( [pad_left, hidden_states, pad_right], dim=2 ) hidden_states = self.conv(hidden_states) return hidden_states # Like LTXVideoResnetBlock3d, but uses new causal Conv3d, normal Conv3d for the conv_shortcut, and the spatial padding # mode is configurable class LTX2VideoResnetBlock3d(nn.Module): r""" A 3D ResNet block used in the LTX 2.0 audiovisual model. Args: in_channels (`int`): Number of input channels. out_channels (`int`, *optional*): Number of output channels. If None, defaults to `in_channels`. dropout (`float`, defaults to `0.0`): Dropout rate. eps (`float`, defaults to `1e-6`): Epsilon value for normalization layers. elementwise_affine (`bool`, defaults to `False`): Whether to enable elementwise affinity in the normalization layers. non_linearity (`str`, defaults to `"swish"`): Activation function to use. conv_shortcut (bool, defaults to `False`): Whether or not to use a convolution shortcut. """ def __init__( self, in_channels: int, out_channels: Optional[int] = None, dropout: float = 0.0, eps: float = 1e-6, elementwise_affine: bool = False, non_linearity: str = "swish", inject_noise: bool = False, timestep_conditioning: bool = False, spatial_padding_mode: str = "zeros", ) -> None: super().__init__() out_channels = out_channels or in_channels self.nonlinearity = get_activation(non_linearity) self.norm1 = PerChannelRMSNorm() self.conv1 = LTX2VideoCausalConv3d( in_channels=in_channels, out_channels=out_channels, kernel_size=3, spatial_padding_mode=spatial_padding_mode, ) self.norm2 = PerChannelRMSNorm() self.dropout = nn.Dropout(dropout) self.conv2 = LTX2VideoCausalConv3d( in_channels=out_channels, out_channels=out_channels, kernel_size=3, spatial_padding_mode=spatial_padding_mode, ) self.norm3 = None self.conv_shortcut = None if in_channels != out_channels: self.norm3 = nn.LayerNorm( in_channels, eps=eps, elementwise_affine=True, bias=True ) # LTX 2.0 uses a normal nn.Conv3d here rather than LTXVideoCausalConv3d self.conv_shortcut = nn.Conv3d( in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, ) self.per_channel_scale1 = None self.per_channel_scale2 = None if inject_noise: self.per_channel_scale1 = nn.Parameter(torch.zeros(in_channels, 1, 1)) self.per_channel_scale2 = nn.Parameter(torch.zeros(in_channels, 1, 1)) self.scale_shift_table = None if timestep_conditioning: self.scale_shift_table = nn.Parameter( torch.randn(4, in_channels) / in_channels**0.5 ) def forward( self, inputs: torch.Tensor, temb: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None, causal: bool = True, ) -> torch.Tensor: hidden_states = inputs hidden_states = self.norm1(hidden_states) if self.scale_shift_table is not None: temb = ( temb.unflatten(1, (4, -1)) + self.scale_shift_table[None, ..., None, None, None] ) shift_1, scale_1, shift_2, scale_2 = temb.unbind(dim=1) hidden_states = hidden_states * (1 + scale_1) + shift_1 hidden_states = self.nonlinearity(hidden_states) hidden_states = self.conv1(hidden_states, causal=causal) if self.per_channel_scale1 is not None: spatial_shape = hidden_states.shape[-2:] spatial_noise = torch.randn( spatial_shape, generator=generator, device=hidden_states.device, dtype=hidden_states.dtype, )[None] hidden_states = ( hidden_states + (spatial_noise * self.per_channel_scale1)[None, :, None, ...] ) hidden_states = self.norm2(hidden_states) if self.scale_shift_table is not None: hidden_states = hidden_states * (1 + scale_2) + shift_2 hidden_states = self.nonlinearity(hidden_states) hidden_states = self.dropout(hidden_states) hidden_states = self.conv2(hidden_states, causal=causal) if self.per_channel_scale2 is not None: spatial_shape = hidden_states.shape[-2:] spatial_noise = torch.randn( spatial_shape, generator=generator, device=hidden_states.device, dtype=hidden_states.dtype, )[None] hidden_states = ( hidden_states + (spatial_noise * self.per_channel_scale2)[None, :, None, ...] ) if self.norm3 is not None: inputs = self.norm3(inputs.movedim(1, -1)).movedim(-1, 1) if self.conv_shortcut is not None: inputs = self.conv_shortcut(inputs) hidden_states = hidden_states + inputs return hidden_states # Like LTX 1.0 LTXVideoDownsampler3d, but uses new causal Conv3d class LTXVideoDownsampler3d(nn.Module): def __init__( self, in_channels: int, out_channels: int, stride: Union[int, Tuple[int, int, int]] = 1, spatial_padding_mode: str = "zeros", ) -> None: super().__init__() self.stride = stride if isinstance(stride, tuple) else (stride, stride, stride) self.group_size = ( in_channels * stride[0] * stride[1] * stride[2] ) // out_channels out_channels = out_channels // ( self.stride[0] * self.stride[1] * self.stride[2] ) self.conv = LTX2VideoCausalConv3d( in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, spatial_padding_mode=spatial_padding_mode, ) def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Tensor: hidden_states = torch.cat( [hidden_states[:, :, : self.stride[0] - 1], hidden_states], dim=2 ) residual = ( hidden_states.unflatten(4, (-1, self.stride[2])) .unflatten(3, (-1, self.stride[1])) .unflatten(2, (-1, self.stride[0])) ) residual = residual.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(1, 4) residual = residual.unflatten(1, (-1, self.group_size)) residual = residual.mean(dim=2) hidden_states = self.conv(hidden_states, causal=causal) hidden_states = ( hidden_states.unflatten(4, (-1, self.stride[2])) .unflatten(3, (-1, self.stride[1])) .unflatten(2, (-1, self.stride[0])) ) hidden_states = hidden_states.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(1, 4) hidden_states = hidden_states + residual return hidden_states # Like LTX 1.0 LTXVideoUpsampler3d, but uses new causal Conv3d class LTXVideoUpsampler3d(nn.Module): def __init__( self, in_channels: int, stride: Union[int, Tuple[int, int, int]] = 1, residual: bool = False, upscale_factor: int = 1, spatial_padding_mode: str = "zeros", ) -> None: super().__init__() self.stride = stride if isinstance(stride, tuple) else (stride, stride, stride) self.residual = residual self.upscale_factor = upscale_factor out_channels = ( in_channels * stride[0] * stride[1] * stride[2] ) // upscale_factor self.conv = LTX2VideoCausalConv3d( in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, spatial_padding_mode=spatial_padding_mode, ) def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = hidden_states.shape if self.residual: residual = hidden_states.reshape( batch_size, -1, self.stride[0], self.stride[1], self.stride[2], num_frames, height, width, ) residual = ( residual.permute(0, 1, 5, 2, 6, 3, 7, 4) .flatten(6, 7) .flatten(4, 5) .flatten(2, 3) ) repeats = ( self.stride[0] * self.stride[1] * self.stride[2] ) // self.upscale_factor residual = residual.repeat(1, repeats, 1, 1, 1) residual = residual[:, :, self.stride[0] - 1 :] hidden_states = self.conv(hidden_states, causal=causal) hidden_states = hidden_states.reshape( batch_size, -1, self.stride[0], self.stride[1], self.stride[2], num_frames, height, width, ) hidden_states = ( hidden_states.permute(0, 1, 5, 2, 6, 3, 7, 4) .flatten(6, 7) .flatten(4, 5) .flatten(2, 3) ) hidden_states = hidden_states[:, :, self.stride[0] - 1 :] if self.residual: hidden_states = hidden_states + residual return hidden_states # Like LTX 1.0 LTXVideo095DownBlock3D, but with the updated LTX2VideoResnetBlock3d class LTX2VideoDownBlock3D(nn.Module): r""" Down block used in the LTXVideo model. Args: in_channels (`int`): Number of input channels. out_channels (`int`, *optional*): Number of output channels. If None, defaults to `in_channels`. num_layers (`int`, defaults to `1`): Number of resnet layers. dropout (`float`, defaults to `0.0`): Dropout rate. resnet_eps (`float`, defaults to `1e-6`): Epsilon value for normalization layers. resnet_act_fn (`str`, defaults to `"swish"`): Activation function to use. spatio_temporal_scale (`bool`, defaults to `True`): Whether or not to use a downsampling layer. If not used, output dimension would be same as input dimension. Whether or not to downsample across temporal dimension. is_causal (`bool`, defaults to `True`): Whether this layer behaves causally (future frames depend only on past frames) or not. """ _supports_gradient_checkpointing = True def __init__( self, in_channels: int, out_channels: Optional[int] = None, num_layers: int = 1, dropout: float = 0.0, resnet_eps: float = 1e-6, resnet_act_fn: str = "swish", spatio_temporal_scale: bool = True, downsample_type: str = "conv", spatial_padding_mode: str = "zeros", ): super().__init__() out_channels = out_channels or in_channels resnets = [] for _ in range(num_layers): resnets.append( LTX2VideoResnetBlock3d( in_channels=in_channels, out_channels=in_channels, dropout=dropout, eps=resnet_eps, non_linearity=resnet_act_fn, spatial_padding_mode=spatial_padding_mode, ) ) self.resnets = nn.ModuleList(resnets) self.downsamplers = None if spatio_temporal_scale: self.downsamplers = nn.ModuleList() if downsample_type == "conv": self.downsamplers.append( LTX2VideoCausalConv3d( in_channels=in_channels, out_channels=in_channels, kernel_size=3, stride=(2, 2, 2), spatial_padding_mode=spatial_padding_mode, ) ) elif downsample_type == "spatial": self.downsamplers.append( LTXVideoDownsampler3d( in_channels=in_channels, out_channels=out_channels, stride=(1, 2, 2), spatial_padding_mode=spatial_padding_mode, ) ) elif downsample_type == "temporal": self.downsamplers.append( LTXVideoDownsampler3d( in_channels=in_channels, out_channels=out_channels, stride=(2, 1, 1), spatial_padding_mode=spatial_padding_mode, ) ) elif downsample_type == "spatiotemporal": self.downsamplers.append( LTXVideoDownsampler3d( in_channels=in_channels, out_channels=out_channels, stride=(2, 2, 2), spatial_padding_mode=spatial_padding_mode, ) ) self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None, causal: bool = True, ) -> torch.Tensor: r"""Forward method of the `LTXDownBlock3D` class.""" for i, resnet in enumerate(self.resnets): if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( resnet, hidden_states, temb, generator, causal ) else: hidden_states = resnet(hidden_states, temb, generator, causal=causal) if self.downsamplers is not None: for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states, causal=causal) return hidden_states # Adapted from diffusers.models.autoencoders.autoencoder_kl_cogvideox.CogVideoMidBlock3d # Like LTX 1.0 LTXVideoMidBlock3d, but with the updated LTX2VideoResnetBlock3d class LTX2VideoMidBlock3d(nn.Module): r""" A middle block used in the LTXVideo model. Args: in_channels (`int`): Number of input channels. num_layers (`int`, defaults to `1`): Number of resnet layers. dropout (`float`, defaults to `0.0`): Dropout rate. resnet_eps (`float`, defaults to `1e-6`): Epsilon value for normalization layers. resnet_act_fn (`str`, defaults to `"swish"`): Activation function to use. is_causal (`bool`, defaults to `True`): Whether this layer behaves causally (future frames depend only on past frames) or not. """ _supports_gradient_checkpointing = True def __init__( self, in_channels: int, num_layers: int = 1, dropout: float = 0.0, resnet_eps: float = 1e-6, resnet_act_fn: str = "swish", inject_noise: bool = False, timestep_conditioning: bool = False, spatial_padding_mode: str = "zeros", ) -> None: super().__init__() self.time_embedder = None if timestep_conditioning: self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings( in_channels * 4, 0 ) resnets = [] for _ in range(num_layers): resnets.append( LTX2VideoResnetBlock3d( in_channels=in_channels, out_channels=in_channels, dropout=dropout, eps=resnet_eps, non_linearity=resnet_act_fn, inject_noise=inject_noise, timestep_conditioning=timestep_conditioning, spatial_padding_mode=spatial_padding_mode, ) ) self.resnets = nn.ModuleList(resnets) self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None, causal: bool = True, ) -> torch.Tensor: r"""Forward method of the `LTXMidBlock3D` class.""" if self.time_embedder is not None: temb = self.time_embedder( timestep=temb.flatten(), resolution=None, aspect_ratio=None, batch_size=hidden_states.size(0), hidden_dtype=hidden_states.dtype, ) temb = temb.view(hidden_states.size(0), -1, 1, 1, 1) for i, resnet in enumerate(self.resnets): if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( resnet, hidden_states, temb, generator, causal ) else: hidden_states = resnet(hidden_states, temb, generator, causal=causal) return hidden_states # Like LTXVideoUpBlock3d but with no conv_in and the updated LTX2VideoResnetBlock3d class LTX2VideoUpBlock3d(nn.Module): r""" Up block used in the LTXVideo model. Args: in_channels (`int`): Number of input channels. out_channels (`int`, *optional*): Number of output channels. If None, defaults to `in_channels`. num_layers (`int`, defaults to `1`): Number of resnet layers. dropout (`float`, defaults to `0.0`): Dropout rate. resnet_eps (`float`, defaults to `1e-6`): Epsilon value for normalization layers. resnet_act_fn (`str`, defaults to `"swish"`): Activation function to use. spatio_temporal_scale (`bool`, defaults to `True`): Whether or not to use a downsampling layer. If not used, output dimension would be same as input dimension. Whether or not to downsample across temporal dimension. is_causal (`bool`, defaults to `True`): Whether this layer behaves causally (future frames depend only on past frames) or not. """ _supports_gradient_checkpointing = True def __init__( self, in_channels: int, out_channels: Optional[int] = None, num_layers: int = 1, dropout: float = 0.0, resnet_eps: float = 1e-6, resnet_act_fn: str = "swish", spatio_temporal_scale: bool = True, inject_noise: bool = False, timestep_conditioning: bool = False, upsample_residual: bool = False, upscale_factor: int = 1, spatial_padding_mode: str = "zeros", ): super().__init__() out_channels = out_channels or in_channels self.time_embedder = None if timestep_conditioning: self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings( in_channels * 4, 0 ) self.conv_in = None if in_channels != out_channels: self.conv_in = LTX2VideoResnetBlock3d( in_channels=in_channels, out_channels=out_channels, dropout=dropout, eps=resnet_eps, non_linearity=resnet_act_fn, inject_noise=inject_noise, timestep_conditioning=timestep_conditioning, spatial_padding_mode=spatial_padding_mode, ) self.upsamplers = None if spatio_temporal_scale: self.upsamplers = nn.ModuleList( [ LTXVideoUpsampler3d( out_channels * upscale_factor, stride=(2, 2, 2), residual=upsample_residual, upscale_factor=upscale_factor, spatial_padding_mode=spatial_padding_mode, ) ] ) resnets = [] for _ in range(num_layers): resnets.append( LTX2VideoResnetBlock3d( in_channels=out_channels, out_channels=out_channels, dropout=dropout, eps=resnet_eps, non_linearity=resnet_act_fn, inject_noise=inject_noise, timestep_conditioning=timestep_conditioning, spatial_padding_mode=spatial_padding_mode, ) ) self.resnets = nn.ModuleList(resnets) self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None, causal: bool = True, ) -> torch.Tensor: if self.conv_in is not None: hidden_states = self.conv_in(hidden_states, temb, generator, causal=causal) if self.time_embedder is not None: temb = self.time_embedder( timestep=temb.flatten(), resolution=None, aspect_ratio=None, batch_size=hidden_states.size(0), hidden_dtype=hidden_states.dtype, ) temb = temb.view(hidden_states.size(0), -1, 1, 1, 1) if self.upsamplers is not None: for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states, causal=causal) for i, resnet in enumerate(self.resnets): if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( resnet, hidden_states, temb, generator, causal ) else: hidden_states = resnet(hidden_states, temb, generator, causal=causal) return hidden_states # Like LTX 1.0 LTXVideoEncoder3d but with different default args - the spatiotemporal downsampling pattern is # different, as is the layers_per_block (the 2.0 VAE is bigger) class LTX2VideoEncoder3d(nn.Module): r""" The `LTXVideoEncoder3d` layer of a variational autoencoder that encodes input video samples to its latent representation. Args: in_channels (`int`, defaults to 3): Number of input channels. out_channels (`int`, defaults to 128): Number of latent channels. block_out_channels (`Tuple[int, ...]`, defaults to `(256, 512, 1024, 2048)`): The number of output channels for each block. spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, True)`: Whether a block should contain spatio-temporal downscaling layers or not. layers_per_block (`Tuple[int, ...]`, defaults to `(4, 6, 6, 2, 2)`): The number of layers per block. downsample_type (`Tuple[str, ...]`, defaults to `("spatial", "temporal", "spatiotemporal", "spatiotemporal")`): The spatiotemporal downsampling pattern per block. Per-layer values can be - `"spatial"` (downsample spatial dims by 2x) - `"temporal"` (downsample temporal dim by 2x) - `"spatiotemporal"` (downsample both spatial and temporal dims by 2x) patch_size (`int`, defaults to `4`): The size of spatial patches. patch_size_t (`int`, defaults to `1`): The size of temporal patches. resnet_norm_eps (`float`, defaults to `1e-6`): Epsilon value for ResNet normalization layers. is_causal (`bool`, defaults to `True`): Whether this layer behaves causally (future frames depend only on past frames) or not. """ def __init__( self, in_channels: int = 3, out_channels: int = 128, block_out_channels: Tuple[int, ...] = (256, 512, 1024, 2048), down_block_types: Tuple[str, ...] = ( "LTX2VideoDownBlock3D", "LTX2VideoDownBlock3D", "LTX2VideoDownBlock3D", "LTX2VideoDownBlock3D", ), spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, True), layers_per_block: Tuple[int, ...] = (4, 6, 6, 2, 2), downsample_type: Tuple[str, ...] = ( "spatial", "temporal", "spatiotemporal", "spatiotemporal", ), patch_size: int = 4, patch_size_t: int = 1, resnet_norm_eps: float = 1e-6, is_causal: bool = True, spatial_padding_mode: str = "zeros", ): super().__init__() self.patch_size = patch_size self.patch_size_t = patch_size_t self.in_channels = in_channels * patch_size**2 self.is_causal = is_causal output_channel = out_channels self.conv_in = LTX2VideoCausalConv3d( in_channels=self.in_channels, out_channels=output_channel, kernel_size=3, stride=1, spatial_padding_mode=spatial_padding_mode, ) # down blocks num_block_out_channels = len(block_out_channels) self.down_blocks = nn.ModuleList([]) for i in range(num_block_out_channels): input_channel = output_channel output_channel = block_out_channels[i] if down_block_types[i] == "LTX2VideoDownBlock3D": down_block = LTX2VideoDownBlock3D( in_channels=input_channel, out_channels=output_channel, num_layers=layers_per_block[i], resnet_eps=resnet_norm_eps, spatio_temporal_scale=spatio_temporal_scaling[i], downsample_type=downsample_type[i], spatial_padding_mode=spatial_padding_mode, ) else: raise ValueError(f"Unknown down block type: {down_block_types[i]}") self.down_blocks.append(down_block) # mid block self.mid_block = LTX2VideoMidBlock3d( in_channels=output_channel, num_layers=layers_per_block[-1], resnet_eps=resnet_norm_eps, spatial_padding_mode=spatial_padding_mode, ) # out self.norm_out = PerChannelRMSNorm() self.conv_act = nn.SiLU() self.conv_out = LTX2VideoCausalConv3d( in_channels=output_channel, out_channels=out_channels + 1, kernel_size=3, stride=1, spatial_padding_mode=spatial_padding_mode, ) self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, causal: Optional[bool] = None ) -> torch.Tensor: r"""The forward method of the `LTXVideoEncoder3d` class.""" p = self.patch_size p_t = self.patch_size_t batch_size, num_channels, num_frames, height, width = hidden_states.shape post_patch_num_frames = num_frames // p_t post_patch_height = height // p post_patch_width = width // p causal = causal or self.is_causal hidden_states = hidden_states.reshape( batch_size, num_channels, post_patch_num_frames, p_t, post_patch_height, p, post_patch_width, p, ) # Thanks for driving me insane with the weird patching order :( hidden_states = hidden_states.permute(0, 1, 3, 7, 5, 2, 4, 6).flatten(1, 4) hidden_states = self.conv_in(hidden_states, causal=causal) if torch.is_grad_enabled() and self.gradient_checkpointing: for down_block in self.down_blocks: hidden_states = self._gradient_checkpointing_func( down_block, hidden_states, None, None, causal ) hidden_states = self._gradient_checkpointing_func( self.mid_block, hidden_states, None, None, causal ) else: for down_block in self.down_blocks: hidden_states = down_block(hidden_states, causal=causal) hidden_states = self.mid_block(hidden_states, causal=causal) hidden_states = self.norm_out(hidden_states) hidden_states = self.conv_act(hidden_states) hidden_states = self.conv_out(hidden_states, causal=causal) last_channel = hidden_states[:, -1:] last_channel = last_channel.repeat(1, hidden_states.size(1) - 2, 1, 1, 1) hidden_states = torch.cat([hidden_states, last_channel], dim=1) return hidden_states # Like LTX 1.0 LTXVideoDecoder3d, but has only 3 symmetric up blocks which are causal and residual with upsample_factor 2 class LTX2VideoDecoder3d(nn.Module): r""" The `LTXVideoDecoder3d` layer of a variational autoencoder that decodes its latent representation into an output sample. Args: in_channels (`int`, defaults to 128): Number of latent channels. out_channels (`int`, defaults to 3): Number of output channels. block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`): The number of output channels for each block. spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, False)`: Whether a block should contain spatio-temporal upscaling layers or not. layers_per_block (`Tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`): The number of layers per block. patch_size (`int`, defaults to `4`): The size of spatial patches. patch_size_t (`int`, defaults to `1`): The size of temporal patches. resnet_norm_eps (`float`, defaults to `1e-6`): Epsilon value for ResNet normalization layers. is_causal (`bool`, defaults to `False`): Whether this layer behaves causally (future frames depend only on past frames) or not. timestep_conditioning (`bool`, defaults to `False`): Whether to condition the model on timesteps. """ def __init__( self, in_channels: int = 128, out_channels: int = 3, block_out_channels: Tuple[int, ...] = (256, 512, 1024), spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True), layers_per_block: Tuple[int, ...] = (5, 5, 5, 5), patch_size: int = 4, patch_size_t: int = 1, resnet_norm_eps: float = 1e-6, is_causal: bool = False, inject_noise: Tuple[bool, ...] = (False, False, False), timestep_conditioning: bool = False, upsample_residual: Tuple[bool, ...] = (True, True, True), upsample_factor: Tuple[bool, ...] = (2, 2, 2), spatial_padding_mode: str = "reflect", ) -> None: super().__init__() self.patch_size = patch_size self.patch_size_t = patch_size_t self.out_channels = out_channels * patch_size**2 self.is_causal = is_causal block_out_channels = tuple(reversed(block_out_channels)) spatio_temporal_scaling = tuple(reversed(spatio_temporal_scaling)) layers_per_block = tuple(reversed(layers_per_block)) inject_noise = tuple(reversed(inject_noise)) upsample_residual = tuple(reversed(upsample_residual)) upsample_factor = tuple(reversed(upsample_factor)) output_channel = block_out_channels[0] self.conv_in = LTX2VideoCausalConv3d( in_channels=in_channels, out_channels=output_channel, kernel_size=3, stride=1, spatial_padding_mode=spatial_padding_mode, ) self.mid_block = LTX2VideoMidBlock3d( in_channels=output_channel, num_layers=layers_per_block[0], resnet_eps=resnet_norm_eps, inject_noise=inject_noise[0], timestep_conditioning=timestep_conditioning, spatial_padding_mode=spatial_padding_mode, ) # up blocks num_block_out_channels = len(block_out_channels) self.up_blocks = nn.ModuleList([]) for i in range(num_block_out_channels): input_channel = output_channel // upsample_factor[i] output_channel = block_out_channels[i] // upsample_factor[i] up_block = LTX2VideoUpBlock3d( in_channels=input_channel, out_channels=output_channel, num_layers=layers_per_block[i + 1], resnet_eps=resnet_norm_eps, spatio_temporal_scale=spatio_temporal_scaling[i], inject_noise=inject_noise[i + 1], timestep_conditioning=timestep_conditioning, upsample_residual=upsample_residual[i], upscale_factor=upsample_factor[i], spatial_padding_mode=spatial_padding_mode, ) self.up_blocks.append(up_block) # out self.norm_out = PerChannelRMSNorm() self.conv_act = nn.SiLU() self.conv_out = LTX2VideoCausalConv3d( in_channels=output_channel, out_channels=self.out_channels, kernel_size=3, stride=1, spatial_padding_mode=spatial_padding_mode, ) # timestep embedding self.time_embedder = None self.scale_shift_table = None self.timestep_scale_multiplier = None if timestep_conditioning: self.timestep_scale_multiplier = nn.Parameter( torch.tensor(1000.0, dtype=torch.float32) ) self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings( output_channel * 2, 0 ) self.scale_shift_table = nn.Parameter( torch.randn(2, output_channel) / output_channel**0.5 ) self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, causal: Optional[bool] = None, ) -> torch.Tensor: causal = causal or self.is_causal hidden_states = self.conv_in(hidden_states, causal=causal) if self.timestep_scale_multiplier is not None: temb = temb * self.timestep_scale_multiplier if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( self.mid_block, hidden_states, temb, None, causal ) for up_block in self.up_blocks: hidden_states = self._gradient_checkpointing_func( up_block, hidden_states, temb, None, causal ) else: hidden_states = self.mid_block(hidden_states, temb, causal=causal) for up_block in self.up_blocks: hidden_states = up_block(hidden_states, temb, causal=causal) hidden_states = self.norm_out(hidden_states) if self.time_embedder is not None: temb = self.time_embedder( timestep=temb.flatten(), resolution=None, aspect_ratio=None, batch_size=hidden_states.size(0), hidden_dtype=hidden_states.dtype, ) temb = temb.view(hidden_states.size(0), -1, 1, 1, 1).unflatten(1, (2, -1)) temb = temb + self.scale_shift_table[None, ..., None, None, None] shift, scale = temb.unbind(dim=1) hidden_states = hidden_states * (1 + scale) + shift hidden_states = self.conv_act(hidden_states) hidden_states = self.conv_out(hidden_states, causal=causal) p = self.patch_size p_t = self.patch_size_t batch_size, num_channels, num_frames, height, width = hidden_states.shape hidden_states = hidden_states.reshape( batch_size, -1, p_t, p, p, num_frames, height, width ) hidden_states = ( hidden_states.permute(0, 1, 5, 2, 6, 4, 7, 3) .flatten(6, 7) .flatten(4, 5) .flatten(2, 3) ) return hidden_states class AutoencoderKLLTX2Video(ParallelTiledVAE): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented for all models (such as downloading or saving). """ _supports_gradient_checkpointing = False def __init__(self, config: LTXVideoVAEConfig): super().__init__(config=config) in_channels = config.arch_config.in_channels latent_channels = config.arch_config.latent_channels out_channels = config.arch_config.out_channels block_out_channels = config.arch_config.block_out_channels down_block_types = config.arch_config.down_block_types spatio_temporal_scaling = config.arch_config.spatio_temporal_scaling layers_per_block = config.arch_config.layers_per_block downsample_type = config.arch_config.downsample_type patch_size = config.arch_config.patch_size patch_size_t = config.arch_config.patch_size_t resnet_norm_eps = config.arch_config.resnet_norm_eps encoder_causal = config.arch_config.encoder_causal encoder_spatial_padding_mode = config.arch_config.encoder_spatial_padding_mode decoder_block_out_channels = config.arch_config.decoder_block_out_channels decoder_spatio_temporal_scaling = ( config.arch_config.decoder_spatio_temporal_scaling ) decoder_layers_per_block = config.arch_config.decoder_layers_per_block decoder_causal = config.arch_config.decoder_causal decoder_spatial_padding_mode = config.arch_config.decoder_spatial_padding_mode self.encoder = LTX2VideoEncoder3d( in_channels, latent_channels, block_out_channels, down_block_types, spatio_temporal_scaling, layers_per_block, downsample_type, patch_size, patch_size_t, resnet_norm_eps, encoder_causal, encoder_spatial_padding_mode, ) self.decoder = LTX2VideoDecoder3d( latent_channels, out_channels, decoder_block_out_channels, decoder_spatio_temporal_scaling, decoder_layers_per_block, patch_size, patch_size_t, resnet_norm_eps, decoder_causal, decoder_spatial_padding_mode, ) latents_mean = torch.zeros((latent_channels,), requires_grad=False) latents_std = torch.ones((latent_channels,), requires_grad=False) self.register_buffer("latents_mean", latents_mean, persistent=True) self.register_buffer("latents_std", latents_std, persistent=True) # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension # to perform decoding of a single video latent at a time. self.use_slicing = False # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the # intermediate tiles together, the memory requirement can be lowered. self.use_tiling = False # When decoding temporally long video latents, the memory requirement is very high. By decoding latent frames # at a fixed frame batch size (based on `self.num_latent_frames_batch_sizes`), the memory requirement can be lowered. self.use_framewise_encoding = False self.use_framewise_decoding = False # This can be configured based on the amount of GPU memory available. # `16` for sample frames and `2` for latent frames are sensible defaults for consumer GPUs. # Setting it to higher values results in higher memory usage. self.num_sample_frames_batch_size = 16 self.num_latent_frames_batch_size = 2 # The minimal tile height and width for spatial tiling to be used self.tile_sample_min_height = 512 self.tile_sample_min_width = 512 self.tile_sample_min_num_frames = 16 # The minimal distance between two spatial tiles self.tile_sample_stride_height = 448 self.tile_sample_stride_width = 448 self.tile_sample_stride_num_frames = 8 def enable_tiling( self, tile_sample_min_height: Optional[int] = None, tile_sample_min_width: Optional[int] = None, tile_sample_min_num_frames: Optional[int] = None, tile_sample_stride_height: Optional[float] = None, tile_sample_stride_width: Optional[float] = None, tile_sample_stride_num_frames: Optional[float] = None, ) -> None: r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow processing larger images. Args: tile_sample_min_height (`int`, *optional*): The minimum height required for a sample to be separated into tiles across the height dimension. tile_sample_min_width (`int`, *optional*): The minimum width required for a sample to be separated into tiles across the width dimension. tile_sample_stride_height (`int`, *optional*): The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are no tiling artifacts produced across the height dimension. tile_sample_stride_width (`int`, *optional*): The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling artifacts produced across the width dimension. """ self.use_tiling = True self.tile_sample_min_height = ( tile_sample_min_height or self.tile_sample_min_height ) self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width self.tile_sample_min_num_frames = ( tile_sample_min_num_frames or self.tile_sample_min_num_frames ) self.tile_sample_stride_height = ( tile_sample_stride_height or self.tile_sample_stride_height ) self.tile_sample_stride_width = ( tile_sample_stride_width or self.tile_sample_stride_width ) self.tile_sample_stride_num_frames = ( tile_sample_stride_num_frames or self.tile_sample_stride_num_frames ) def _encode(self, x: torch.Tensor, causal: Optional[bool] = None) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = x.shape if self.use_framewise_decoding and num_frames > self.tile_sample_min_num_frames: return self._temporal_tiled_encode(x, causal=causal) if self.use_tiling and ( width > self.tile_sample_min_width or height > self.tile_sample_min_height ): return self.tiled_encode(x, causal=causal) enc = self.encoder(x, causal=causal) return enc def encode( self, x: torch.Tensor, causal: Optional[bool] = None, return_dict: bool = True ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: """ Encode a batch of images into latents. Args: x (`torch.Tensor`): Input batch of images. return_dict (`bool`, *optional*, defaults to `True`): Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. Returns: The latent representations of the encoded videos. If `return_dict` is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. """ if self.use_slicing and x.shape[0] > 1: encoded_slices = [ self._encode(x_slice, causal=causal) for x_slice in x.split(1) ] h = torch.cat(encoded_slices) else: h = self._encode(x, causal=causal) posterior = DiagonalGaussianDistribution(h) if not return_dict: return (posterior,) return AutoencoderKLOutput(latent_dist=posterior) def _decode( self, z: torch.Tensor, temb: Optional[torch.Tensor] = None, causal: Optional[bool] = None, return_dict: bool = True, ) -> Union[DecoderOutput, torch.Tensor]: batch_size, num_channels, num_frames, height, width = z.shape tile_latent_min_height = ( self.tile_sample_min_height // self.spatial_compression_ratio ) tile_latent_min_width = ( self.tile_sample_min_width // self.spatial_compression_ratio ) tile_latent_min_num_frames = ( self.tile_sample_min_num_frames // self.temporal_compression_ratio ) if self.use_framewise_decoding and num_frames > tile_latent_min_num_frames: return self._temporal_tiled_decode( z, temb, causal=causal, return_dict=return_dict ) if self.use_tiling and ( width > tile_latent_min_width or height > tile_latent_min_height ): return self.tiled_decode(z, temb, causal=causal, return_dict=return_dict) dec = self.decoder(z, temb, causal=causal) if not return_dict: return (dec,) return DecoderOutput(sample=dec) def decode( self, z: torch.Tensor, temb: Optional[torch.Tensor] = None, causal: Optional[bool] = None, return_dict: bool = True, ) -> Union[DecoderOutput, torch.Tensor]: """ Decode a batch of images. Args: z (`torch.Tensor`): Input batch of latent vectors. return_dict (`bool`, *optional*, defaults to `True`): Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. Returns: [`~models.vae.DecoderOutput`] or `tuple`: If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is returned. """ if self.use_slicing and z.shape[0] > 1: if temb is not None: decoded_slices = [ self._decode(z_slice, t_slice, causal=causal).sample for z_slice, t_slice in (z.split(1), temb.split(1)) ] else: decoded_slices = [ self._decode(z_slice, causal=causal).sample for z_slice in z.split(1) ] decoded = torch.cat(decoded_slices) else: decoded = self._decode(z, temb, causal=causal).sample if not return_dict: return (decoded,) return DecoderOutput(sample=decoded) def blend_v( self, a: torch.Tensor, b: torch.Tensor, blend_extent: int ) -> torch.Tensor: blend_extent = min(a.shape[3], b.shape[3], blend_extent) for y in range(blend_extent): b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * ( 1 - y / blend_extent ) + b[:, :, :, y, :] * (y / blend_extent) return b def blend_h( self, a: torch.Tensor, b: torch.Tensor, blend_extent: int ) -> torch.Tensor: blend_extent = min(a.shape[4], b.shape[4], blend_extent) for x in range(blend_extent): b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * ( 1 - x / blend_extent ) + b[:, :, :, :, x] * (x / blend_extent) return b def blend_t( self, a: torch.Tensor, b: torch.Tensor, blend_extent: int ) -> torch.Tensor: blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) for x in range(blend_extent): b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * ( 1 - x / blend_extent ) + b[:, :, x, :, :] * (x / blend_extent) return b def tiled_encode( self, x: torch.Tensor, causal: Optional[bool] = None ) -> torch.Tensor: r"""Encode a batch of images using a tiled encoder. Args: x (`torch.Tensor`): Input batch of videos. Returns: `torch.Tensor`: The latent representation of the encoded videos. """ batch_size, num_channels, num_frames, height, width = x.shape latent_height = height // self.spatial_compression_ratio latent_width = width // self.spatial_compression_ratio tile_latent_min_height = ( self.tile_sample_min_height // self.spatial_compression_ratio ) tile_latent_min_width = ( self.tile_sample_min_width // self.spatial_compression_ratio ) tile_latent_stride_height = ( self.tile_sample_stride_height // self.spatial_compression_ratio ) tile_latent_stride_width = ( self.tile_sample_stride_width // self.spatial_compression_ratio ) blend_height = tile_latent_min_height - tile_latent_stride_height blend_width = tile_latent_min_width - tile_latent_stride_width # Split x into overlapping tiles and encode them separately. # The tiles have an overlap to avoid seams between tiles. rows = [] for i in range(0, height, self.tile_sample_stride_height): row = [] for j in range(0, width, self.tile_sample_stride_width): time = self.encoder( x[ :, :, :, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width, ], causal=causal, ) row.append(time) rows.append(row) result_rows = [] for i, row in enumerate(rows): result_row = [] for j, tile in enumerate(row): # blend the above tile and the left tile # to the current tile and add the current tile to the result row if i > 0: tile = self.blend_v(rows[i - 1][j], tile, blend_height) if j > 0: tile = self.blend_h(row[j - 1], tile, blend_width) result_row.append( tile[:, :, :, :tile_latent_stride_height, :tile_latent_stride_width] ) result_rows.append(torch.cat(result_row, dim=4)) enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width] return enc def tiled_decode( self, z: torch.Tensor, temb: Optional[torch.Tensor], causal: Optional[bool] = None, return_dict: bool = True, ) -> Union[DecoderOutput, torch.Tensor]: r""" Decode a batch of images using a tiled decoder. Args: z (`torch.Tensor`): Input batch of latent vectors. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. Returns: [`~models.vae.DecoderOutput`] or `tuple`: If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is returned. """ batch_size, num_channels, num_frames, height, width = z.shape sample_height = height * self.spatial_compression_ratio sample_width = width * self.spatial_compression_ratio tile_latent_min_height = ( self.tile_sample_min_height // self.spatial_compression_ratio ) tile_latent_min_width = ( self.tile_sample_min_width // self.spatial_compression_ratio ) tile_latent_stride_height = ( self.tile_sample_stride_height // self.spatial_compression_ratio ) tile_latent_stride_width = ( self.tile_sample_stride_width // self.spatial_compression_ratio ) blend_height = self.tile_sample_min_height - self.tile_sample_stride_height blend_width = self.tile_sample_min_width - self.tile_sample_stride_width # Split z into overlapping tiles and decode them separately. # The tiles have an overlap to avoid seams between tiles. rows = [] for i in range(0, height, tile_latent_stride_height): row = [] for j in range(0, width, tile_latent_stride_width): time = self.decoder( z[ :, :, :, i : i + tile_latent_min_height, j : j + tile_latent_min_width, ], temb, causal=causal, ) row.append(time) rows.append(row) result_rows = [] for i, row in enumerate(rows): result_row = [] for j, tile in enumerate(row): # blend the above tile and the left tile # to the current tile and add the current tile to the result row if i > 0: tile = self.blend_v(rows[i - 1][j], tile, blend_height) if j > 0: tile = self.blend_h(row[j - 1], tile, blend_width) result_row.append( tile[ :, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width, ] ) result_rows.append(torch.cat(result_row, dim=4)) dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width] if not return_dict: return (dec,) return DecoderOutput(sample=dec) def _temporal_tiled_encode( self, x: torch.Tensor, causal: Optional[bool] = None ) -> AutoencoderKLOutput: batch_size, num_channels, num_frames, height, width = x.shape latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1 tile_latent_min_num_frames = ( self.tile_sample_min_num_frames // self.temporal_compression_ratio ) tile_latent_stride_num_frames = ( self.tile_sample_stride_num_frames // self.temporal_compression_ratio ) blend_num_frames = tile_latent_min_num_frames - tile_latent_stride_num_frames row = [] for i in range(0, num_frames, self.tile_sample_stride_num_frames): tile = x[:, :, i : i + self.tile_sample_min_num_frames + 1, :, :] if self.use_tiling and ( height > self.tile_sample_min_height or width > self.tile_sample_min_width ): tile = self.tiled_encode(tile, causal=causal) else: tile = self.encoder(tile, causal=causal) if i > 0: tile = tile[:, :, 1:, :, :] row.append(tile) result_row = [] for i, tile in enumerate(row): if i > 0: tile = self.blend_t(row[i - 1], tile, blend_num_frames) result_row.append(tile[:, :, :tile_latent_stride_num_frames, :, :]) else: result_row.append(tile[:, :, : tile_latent_stride_num_frames + 1, :, :]) enc = torch.cat(result_row, dim=2)[:, :, :latent_num_frames] return enc def _temporal_tiled_decode( self, z: torch.Tensor, temb: Optional[torch.Tensor], causal: Optional[bool] = None, return_dict: bool = True, ) -> Union[DecoderOutput, torch.Tensor]: batch_size, num_channels, num_frames, height, width = z.shape num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1 tile_latent_min_height = ( self.tile_sample_min_height // self.spatial_compression_ratio ) tile_latent_min_width = ( self.tile_sample_min_width // self.spatial_compression_ratio ) tile_latent_min_num_frames = ( self.tile_sample_min_num_frames // self.temporal_compression_ratio ) tile_latent_stride_num_frames = ( self.tile_sample_stride_num_frames // self.temporal_compression_ratio ) blend_num_frames = ( self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames ) row = [] for i in range(0, num_frames, tile_latent_stride_num_frames): tile = z[:, :, i : i + tile_latent_min_num_frames + 1, :, :] if self.use_tiling and ( tile.shape[-1] > tile_latent_min_width or tile.shape[-2] > tile_latent_min_height ): decoded = self.tiled_decode( tile, temb, causal=causal, return_dict=True ).sample else: decoded = self.decoder(tile, temb, causal=causal) if i > 0: decoded = decoded[:, :, :-1, :, :] row.append(decoded) result_row = [] for i, tile in enumerate(row): if i > 0: tile = self.blend_t(row[i - 1], tile, blend_num_frames) tile = tile[:, :, : self.tile_sample_stride_num_frames, :, :] result_row.append(tile) else: result_row.append( tile[:, :, : self.tile_sample_stride_num_frames + 1, :, :] ) dec = torch.cat(result_row, dim=2)[:, :, :num_sample_frames] if not return_dict: return (dec,) return DecoderOutput(sample=dec) def forward( self, sample: torch.Tensor, temb: Optional[torch.Tensor] = None, sample_posterior: bool = False, encoder_causal: Optional[bool] = None, decoder_causal: Optional[bool] = None, return_dict: bool = True, generator: Optional[torch.Generator] = None, ) -> Union[torch.Tensor, torch.Tensor]: x = sample posterior = self.encode(x, causal=encoder_causal).latent_dist if sample_posterior: z = posterior.sample(generator=generator) else: z = posterior.mode() dec = self.decode(z, temb, causal=decoder_causal) if not return_dict: return (dec.sample,) return dec EntryClass = AutoencoderKLLTX2Video ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_common_utils.py ================================================ from __future__ import annotations import torch import torch.nn as nn import torch.nn.functional as F from sglang.multimodal_gen.runtime.platforms import current_platform class AvgDown3D(nn.Module): def __init__( self, in_channels, out_channels, factor_t, factor_s=1, ): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.factor_t = factor_t self.factor_s = factor_s self.factor = self.factor_t * self.factor_s * self.factor_s assert in_channels * self.factor % out_channels == 0 self.group_size = in_channels * self.factor // out_channels def forward(self, x: torch.Tensor) -> torch.Tensor: pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t pad = (0, 0, 0, 0, pad_t, 0) x = F.pad(x, pad) B, C, T, H, W = x.shape x = x.view( B, C, T // self.factor_t, self.factor_t, H // self.factor_s, self.factor_s, W // self.factor_s, self.factor_s, ) x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() x = x.view( B, C * self.factor, T // self.factor_t, H // self.factor_s, W // self.factor_s, ) x = x.view( B, self.out_channels, self.group_size, T // self.factor_t, H // self.factor_s, W // self.factor_s, ) x = x.mean(dim=2) return x class DupUp3D(nn.Module): def __init__( self, in_channels: int, out_channels: int, factor_t, factor_s=1, ): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.factor_t = factor_t self.factor_s = factor_s self.factor = self.factor_t * self.factor_s * self.factor_s assert out_channels * self.factor % in_channels == 0 self.repeats = out_channels * self.factor // in_channels def forward(self, x: torch.Tensor) -> torch.Tensor: x = x.repeat_interleave(self.repeats, dim=1) x = x.view( x.size(0), self.out_channels, self.factor_t, self.factor_s, self.factor_s, x.size(2), x.size(3), x.size(4), ) x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() x = x.view( x.size(0), self.out_channels, x.size(2) * self.factor_t, x.size(4) * self.factor_s, x.size(6) * self.factor_s, ) _first_chunk = first_chunk.get() if first_chunk is not None else None if _first_chunk: x = x[:, :, self.factor_t - 1 :, :, :] return x class WanCausalConv3d(nn.Conv3d): r""" A custom 3D causal convolution layer with feature caching support. This layer extends the standard Conv3D layer by ensuring causality in the time dimension and handling feature caching for efficient inference. """ def __init__( self, in_channels: int, out_channels: int, kernel_size: int | tuple[int, int, int], stride: int | tuple[int, int, int] = 1, padding: int | tuple[int, int, int] = 0, ) -> None: super().__init__( in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, ) self.padding: tuple[int, int, int] # Set up causal padding self._padding: tuple[int, ...] = ( self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0, ) self.padding = (0, 0, 0) def forward(self, x, cache_x=None): padding = list(self._padding) if cache_x is not None and self._padding[4] > 0: cache_x = cache_x.to(x.device) x = torch.cat([cache_x, x], dim=2) padding[4] -= cache_x.shape[2] x = F.pad(x, padding) x = ( x.to(self.weight.dtype) if current_platform.is_mps() else x ) # casting needed for mps since amp isn't supported return super().forward(x) class WanRMS_norm(nn.Module): r""" A custom RMS normalization layer. """ def __init__( self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False, ) -> None: super().__init__() broadcastable_dims = (1, 1, 1) if not images else (1, 1) shape = (dim, *broadcastable_dims) if channel_first else (dim,) self.channel_first = channel_first self.scale = dim**0.5 self.gamma = nn.Parameter(torch.ones(shape)) self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 def forward(self, x): return ( F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias ) class WanUpsample(nn.Upsample): r""" Perform upsampling while ensuring the output tensor has the same data type as the input. """ def forward(self, x): return super().forward(x.float()).type_as(x) is_first_frame = None feat_cache = None feat_idx = None cache_t = None first_chunk = None def bind_context( is_first_frame_var, feat_cache_var, feat_idx_var, cache_t_value, first_chunk_var, ): global is_first_frame global feat_cache global feat_idx global cache_t global first_chunk is_first_frame = is_first_frame_var feat_cache = feat_cache_var feat_idx = feat_idx_var cache_t = cache_t_value first_chunk = first_chunk_var def _ensure_bound(): if ( is_first_frame is None or feat_cache is None or feat_idx is None or cache_t is None or first_chunk is None ): raise RuntimeError("common_utils.bind_context() must be called before use.") def resample_forward(self, x): _ensure_bound() b, c, t, h, w = x.size() first_frame = is_first_frame.get() if first_frame: assert t == 1 _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if self.mode == "upsample3d": if _feat_cache is not None: idx = _feat_idx if _feat_cache[idx] is None: _feat_cache[idx] = "Rep" _feat_idx += 1 else: cache_x = x[:, :, -cache_t:, :, :].clone() if ( cache_x.shape[2] < 2 and _feat_cache[idx] is not None and _feat_cache[idx] != "Rep" ): # cache last frame of last two chunk cache_x = torch.cat( [ _feat_cache[idx][:, :, -1, :, :] .unsqueeze(2) .to(cache_x.device), cache_x, ], dim=2, ) if ( cache_x.shape[2] < 2 and _feat_cache[idx] is not None and _feat_cache[idx] == "Rep" ): cache_x = torch.cat( [torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2, ) if _feat_cache[idx] == "Rep": x = self.time_conv(x) else: x = self.time_conv(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 x = x.reshape(b, 2, c, t, h, w) x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) x = x.reshape(b, c, t * 2, h, w) feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) elif not first_frame and hasattr(self, "time_conv"): x = self.time_conv(x) x = x.reshape(b, 2, c, t, h, w) x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) x = x.reshape(b, c, t * 2, h, w) t = x.shape[2] x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) x = self.resample(x) x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if self.mode == "downsample3d": if _feat_cache is not None: idx = _feat_idx if _feat_cache[idx] is None: _feat_cache[idx] = x.clone() _feat_idx += 1 else: cache_x = x[:, :, -1:, :, :].clone() x = self.time_conv(torch.cat([_feat_cache[idx][:, :, -1:, :, :], x], 2)) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) elif not first_frame and hasattr(self, "time_conv"): x = self.time_conv(x) return x def residual_block_forward(self, x): _ensure_bound() # Apply shortcut connection h = self.conv_shortcut(x) # First normalization and activation x = self.norm1(x) x = self.nonlinearity(x) _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx cache_x = x[:, :, -cache_t:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: cache_x = torch.cat( [ _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x, ], dim=2, ) x = self.conv1(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) else: x = self.conv1(x) # Second normalization and activation x = self.norm2(x) x = self.nonlinearity(x) # Dropout x = self.dropout(x) _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx cache_x = x[:, :, -cache_t:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: cache_x = torch.cat( [ _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x, ], dim=2, ) x = self.conv2(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) else: x = self.conv2(x) # Add residual connection return x + h def attention_block_forward(self, x): identity = x batch_size, channels, num_frames, height, width = x.size() x = x.permute(0, 2, 1, 3, 4).reshape( batch_size * num_frames, channels, height, width ) x = self.norm(x) # compute query, key, value qkv = self.to_qkv(x) qkv = qkv.reshape(batch_size * num_frames, 1, channels * 3, -1) qkv = qkv.permute(0, 1, 3, 2).contiguous() q, k, v = qkv.chunk(3, dim=-1) x = torch.nn.functional.scaled_dot_product_attention(q, k, v) x = ( x.squeeze(1) .permute(0, 2, 1) .reshape(batch_size * num_frames, channels, height, width) ) # output projection x = self.proj(x) # Reshape back: [(b*t), c, h, w] -> [b, c, t, h, w] x = x.view(batch_size, num_frames, channels, height, width) x = x.permute(0, 2, 1, 3, 4) return x + identity def mid_block_forward(self, x): # First residual block x = self.resnets[0](x) # Process through attention and residual blocks for attn, resnet in zip(self.attentions, self.resnets[1:], strict=True): if attn is not None: x = attn(x) x = resnet(x) return x def residual_down_block_forward(self, x): x_copy = x for resnet in self.resnets: x = resnet(x) if self.downsampler is not None: x = self.downsampler(x) return x + self.avg_shortcut(x_copy) def residual_up_block_forward(self, x): if self.avg_shortcut is not None: x_copy = x for resnet in self.resnets: x = resnet(x) if self.upsampler is not None: x = self.upsampler(x) if self.avg_shortcut is not None: x = x + self.avg_shortcut(x_copy) return x def up_block_forward(self, x): for resnet in self.resnets: x = resnet(x) if self.upsamplers is not None: x = self.upsamplers[0](x) return x ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_dist_utils.py ================================================ import math import torch import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_sp_group, get_sp_parallel_rank, get_sp_world_size, ) from sglang.multimodal_gen.runtime.layers.activation import get_act_fn from sglang.multimodal_gen.runtime.models.vaes.parallel.wan_common_utils import ( AvgDown3D, DupUp3D, WanCausalConv3d, WanRMS_norm, WanUpsample, attention_block_forward, mid_block_forward, resample_forward, residual_block_forward, residual_down_block_forward, residual_up_block_forward, up_block_forward, ) from sglang.multimodal_gen.runtime.platforms import current_platform def tensor_pad(x: torch.Tensor, len_to_pad: int, dim: int = -2): x = torch.cat( [ x, torch.zeros( *x.shape[:dim], len_to_pad, *x.shape[dim + 1 :], dtype=x.dtype, device=x.device, ), ], dim=dim, ) return x def tensor_chunk(x: torch.Tensor, dim: int = -2, world_size: int = 1, rank: int = 0): if x is None: return None if world_size <= 1: return x len_to_padding = (int(math.ceil(x.shape[dim] / world_size)) * world_size) - x.shape[ dim ] if len_to_padding != 0: x = tensor_pad(x, len_to_padding, dim=dim) return torch.chunk(x, world_size, dim=dim)[rank] def split_for_parallel_encode( x: torch.Tensor, downsample_count: int, world_size: int, rank: int ): orig_height = x.shape[-2] expected_height = orig_height // (2**downsample_count) factor = world_size * (2**downsample_count) pad_h = (factor - orig_height % factor) % factor if pad_h: x = F.pad(x, (0, 0, 0, pad_h, 0, 0)) expected_local_height = (orig_height + pad_h) // (2**downsample_count) // world_size x = tensor_chunk(x, dim=-2, world_size=world_size, rank=rank) return x, expected_height, expected_local_height def ensure_local_height(x: torch.Tensor, expected_local_height: int | None): if expected_local_height is None: return x if x.shape[-2] < expected_local_height: pad = expected_local_height - x.shape[-2] return F.pad(x, (0, 0, 0, pad, 0, 0)) if x.shape[-2] > expected_local_height: return x[..., :expected_local_height, :].contiguous() return x def split_for_parallel_decode( x: torch.Tensor, upsample_count: int, world_size: int, rank: int ): expected_height = x.shape[-2] * (2**upsample_count) x = tensor_chunk(x, dim=-2, world_size=world_size, rank=rank) return x, expected_height def gather_and_trim_height(x: torch.Tensor, expected_height: int | None): if expected_height is None: return x x = get_sp_group().all_gather(x, dim=-2) if x.shape[-2] != expected_height: x = x[..., :expected_height, :].contiguous() return x def _ensure_recv_buf( recv_buf: torch.Tensor | None, reference: torch.Tensor ) -> torch.Tensor: if ( recv_buf is None or recv_buf.shape != reference.shape or recv_buf.dtype != reference.dtype or recv_buf.device != reference.device ): return torch.empty_like(reference) return recv_buf def halo_exchange( x: torch.Tensor, height_halo_size: int = 1, recv_top_buf: torch.Tensor | None = None, recv_bottom_buf: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if height_halo_size == 0: return x, recv_top_buf, recv_bottom_buf sp_group = get_sp_group() rank = get_sp_parallel_rank() world_size = get_sp_world_size() group = sp_group.device_group group_ranks = sp_group.ranks top_row = x[..., :height_halo_size, :].contiguous() bottom_row = x[..., -height_halo_size:, :].contiguous() recv_top_buf = _ensure_recv_buf(recv_top_buf, top_row) recv_bottom_buf = _ensure_recv_buf(recv_bottom_buf, bottom_row) # use batched P2P operations p2p_ops = [] if rank > 0: # has previous neighbor, recv previous rank's data to recv_top_buf and send top_row to it. prev_rank = group_ranks[rank - 1] p2p_ops.append(dist.P2POp(dist.irecv, recv_top_buf, prev_rank, group)) p2p_ops.append(dist.P2POp(dist.isend, top_row, prev_rank, group)) if rank < world_size - 1: # has next neighbor, send bottom_row to next rank and recv next rank's data to recv_bottom_buf. next_rank = group_ranks[rank + 1] p2p_ops.append(dist.P2POp(dist.isend, bottom_row, next_rank, group)) p2p_ops.append(dist.P2POp(dist.irecv, recv_bottom_buf, next_rank, group)) if rank == 0: recv_top_buf.zero_() if rank == world_size - 1: recv_bottom_buf.zero_() if p2p_ops: reqs = dist.batch_isend_irecv(p2p_ops) for req in reqs: req.wait() return ( torch.concat([recv_top_buf, x, recv_bottom_buf], dim=-2), recv_top_buf, recv_bottom_buf, ) class WanDistConv2d(nn.Conv2d): def __init__( self, in_channels: int, out_channels: int, kernel_size: int | tuple[int, int, int], stride: int | tuple[int, int, int] = 1, padding: int | tuple[int, int, int] = 0, height_padding: tuple[int, int] | None = None, ): super().__init__( in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, ) self.height_halo_size = (self.kernel_size[-2] - 1) // 2 if height_padding is None: height_padding = (self.padding[-2], self.padding[-2]) self.height_pad_top, self.height_pad_bottom = height_padding self.padding: tuple[int, int] if self.height_halo_size > 0: self._padding = (self.padding[1], self.padding[1], 0, 0) else: self._padding = ( self.padding[1], self.padding[1], self.padding[0], self.padding[0], ) self.padding = (0, 0) self._halo_recv_top_buf: torch.Tensor | None = None self._halo_recv_bottom_buf: torch.Tensor | None = None self.rank = get_sp_parallel_rank() self.world_size = get_sp_world_size() def forward(self, x): x = F.pad(x, self._padding) x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange( x, height_halo_size=self.height_halo_size, recv_top_buf=self._halo_recv_top_buf, recv_bottom_buf=self._halo_recv_bottom_buf, ) pad_top = self.height_pad_top stride = self.stride[-2] global_start = self.rank * x.shape[-2] if self.height_halo_size > 0 and stride > 1: shift = (global_start - self.height_halo_size + pad_top) % stride if shift: x_padded = x_padded[..., shift:, :] global_start += shift out = super().forward(x_padded) if self.height_halo_size == 0: return out local_height = x.shape[-2] global_height = local_height * self.world_size halo = self.height_halo_size pad_bottom = self.height_pad_bottom kernel = self.kernel_size[-2] min_i = math.ceil(((-pad_top) - (global_start - halo)) / stride) max_i = math.floor( ((global_height - 1 + pad_bottom) - (kernel - 1) - (global_start - halo)) / stride ) start = max(min_i, 0) end = min(max_i + 1, out.shape[-2]) if start != 0 or end != out.shape[-2]: out = out[..., start:end, :] return out class WanDistCausalConv3d(nn.Conv3d): def __init__( self, in_channels: int, out_channels: int, kernel_size: int | tuple[int, int, int], stride: int | tuple[int, int, int] = 1, padding: int | tuple[int, int, int] = 0, ): super().__init__( in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, ) self.height_pad_top = self.padding[1] self.height_pad_bottom = self.padding[1] self.height_halo_size = (self.kernel_size[-2] - 1) // 2 self.padding: tuple[int, int, int] # Set up causal padding, let the halo to control height padding if self.height_halo_size > 0: self._padding: tuple[int, ...] = ( self.padding[2], self.padding[2], 0, 0, 2 * self.padding[0], 0, ) else: self._padding: tuple[int, ...] = ( self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0, ) self.padding = (0, 0, 0) self._halo_recv_top_buf: torch.Tensor | None = None self._halo_recv_bottom_buf: torch.Tensor | None = None self.rank = get_sp_parallel_rank() self.world_size = get_sp_world_size() def forward(self, x, cache_x=None): padding = list(self._padding) if cache_x is not None and self._padding[4] > 0: cache_x = cache_x.to(x.device) x = torch.cat([cache_x, x], dim=2) padding[4] -= cache_x.shape[2] x = F.pad(x, padding) x = ( x.to(self.weight.dtype) if current_platform.is_mps() else x ) # casting needed for mps since amp isn't supported x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange( x, height_halo_size=self.height_halo_size, recv_top_buf=self._halo_recv_top_buf, recv_bottom_buf=self._halo_recv_bottom_buf, ) pad_top = self.height_pad_top stride = self.stride[-2] global_start = self.rank * x.shape[-2] if self.height_halo_size > 0 and stride > 1: shift = (global_start - self.height_halo_size + pad_top) % stride if shift: x_padded = x_padded[..., shift:, :] global_start += shift out = super().forward(x_padded) if self.height_halo_size == 0: return out local_height = x.shape[-2] global_height = local_height * self.world_size halo = self.height_halo_size pad_bottom = self.height_pad_bottom kernel = self.kernel_size[-2] min_i = math.ceil(((-pad_top) - (global_start - halo)) / stride) max_i = math.floor( ((global_height - 1 + pad_bottom) - (kernel - 1) - (global_start - halo)) / stride ) start = max(min_i, 0) end = min(max_i + 1, out.shape[-2]) if start != 0 or end != out.shape[-2]: out = out[..., start:end, :] return out class WanDistZeroPad2d(nn.Module): """Apply 2D padding once globally across sequence-parallel height splits.""" def __init__(self, padding: tuple[int, int, int, int]) -> None: super().__init__() self.padding = padding # (left, right, top, bottom) self.rank = get_sp_parallel_rank() self.world_size = get_sp_world_size() def forward(self, x: torch.Tensor) -> torch.Tensor: left, right, top, bottom = self.padding if self.world_size <= 1: return F.pad(x, (left, right, top, bottom)) # Only the first/last rank should contribute global top/bottom padding. top = top if self.rank == 0 else 0 bottom = bottom if self.rank == self.world_size - 1 else 0 return F.pad(x, (left, right, top, bottom)) class WanDistResample(nn.Module): r""" A custom resampling module for 2D and 3D data used for parallel decoding. Args: dim (int): The number of input/output channels. mode (str): The resampling mode. Must be one of: - 'none': No resampling (identity operation). - 'upsample2d': 2D upsampling with nearest-exact interpolation and convolution. - 'upsample3d': 3D upsampling with nearest-exact interpolation, convolution, and causal 3D convolution. - 'downsample2d': 2D downsampling with zero-padding and convolution. - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution. """ def __init__(self, dim: int, mode: str, upsample_out_dim: int = None) -> None: super().__init__() self.dim = dim self.mode = mode # default to dim //2 if upsample_out_dim is None: upsample_out_dim = dim // 2 # layers # We support parallel encode/decode; downsample uses halo exchange as well. if mode == "upsample2d": self.resample = nn.Sequential( WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), WanDistConv2d(dim, upsample_out_dim, 3, padding=1), ) elif mode == "upsample3d": self.resample = nn.Sequential( WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), WanDistConv2d(dim, upsample_out_dim, 3, padding=1), ) self.time_conv = WanCausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) elif mode == "downsample2d": self.resample = nn.Sequential( WanDistZeroPad2d((0, 1, 0, 0)), WanDistConv2d(dim, dim, 3, stride=(2, 2), height_padding=(0, 1)), ) elif mode == "downsample3d": self.resample = nn.Sequential( WanDistZeroPad2d((0, 1, 0, 0)), WanDistConv2d(dim, dim, 3, stride=(2, 2), height_padding=(0, 1)), ) self.time_conv = WanCausalConv3d( dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0) ) else: self.resample = nn.Identity() def forward(self, x): return resample_forward(self, x) class WanDistResidualBlock(nn.Module): r""" A custom residual block module. Args: in_dim (int): Number of input channels. out_dim (int): Number of output channels. dropout (float, optional): Dropout rate for the dropout layer. Default is 0.0. non_linearity (str, optional): Type of non-linearity to use. Default is "silu". """ def __init__( self, in_dim: int, out_dim: int, dropout: float = 0.0, non_linearity: str = "silu", ) -> None: super().__init__() self.in_dim = in_dim self.out_dim = out_dim self.nonlinearity = get_act_fn(non_linearity) # layers self.norm1 = WanRMS_norm(in_dim, images=False) self.conv1 = WanDistCausalConv3d(in_dim, out_dim, 3, padding=1) self.norm2 = WanRMS_norm(out_dim, images=False) self.dropout = nn.Dropout(dropout) self.conv2 = WanDistCausalConv3d(out_dim, out_dim, 3, padding=1) self.conv_shortcut = ( WanDistCausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() ) def forward(self, x): return residual_block_forward(self, x) class WanDistAttentionBlock(nn.Module): r""" Causal self-attention with a single head. Args: dim (int): The number of channels in the input tensor. """ def __init__(self, dim) -> None: super().__init__() self.dim = dim # layers self.norm = WanRMS_norm(dim) self.to_qkv = nn.Conv2d(dim, dim * 3, 1) self.proj = nn.Conv2d(dim, dim, 1) self.rank = get_sp_parallel_rank() self.world_size = get_sp_world_size() self.sp_group = get_sp_group() def forward(self, x): if self.world_size > 1: x = self.sp_group.all_gather(x, dim=-2) x = x.contiguous() x = attention_block_forward(self, x) if self.world_size > 1: x = torch.chunk(x, self.world_size, dim=-2)[self.rank] return x class WanDistMidBlock(nn.Module): """ Middle block for WanVAE encoder and decoder. Args: dim (int): Number of input/output channels. dropout (float): Dropout rate. non_linearity (str): Type of non-linearity to use. """ def __init__( self, dim: int, dropout: float = 0.0, non_linearity: str = "silu", num_layers: int = 1, ): super().__init__() self.dim = dim # Create the components resnets = [WanDistResidualBlock(dim, dim, dropout, non_linearity)] attentions = [] for _ in range(num_layers): attentions.append(WanDistAttentionBlock(dim)) resnets.append(WanDistResidualBlock(dim, dim, dropout, non_linearity)) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) self.gradient_checkpointing = False def forward(self, x): return mid_block_forward(self, x) class WanDistResidualDownBlock(nn.Module): def __init__( self, in_dim, out_dim, dropout, num_res_blocks, temperal_downsample=False, down_flag=False, ): super().__init__() # Shortcut path with downsample self.avg_shortcut = AvgDown3D( in_dim, out_dim, factor_t=2 if temperal_downsample else 1, factor_s=2 if down_flag else 1, ) # Main path with residual blocks and downsample resnets = [] for _ in range(num_res_blocks): resnets.append(WanDistResidualBlock(in_dim, out_dim, dropout)) in_dim = out_dim self.resnets = nn.ModuleList(resnets) # Add the final downsample block if down_flag: mode = "downsample3d" if temperal_downsample else "downsample2d" self.downsampler = WanDistResample(out_dim, mode=mode) else: self.downsampler = None def forward(self, x): return residual_down_block_forward(self, x) class WanDistResidualUpBlock(nn.Module): """ A block that handles upsampling for the WanVAE decoder. Args: in_dim (int): Input dimension out_dim (int): Output dimension num_res_blocks (int): Number of residual blocks dropout (float): Dropout rate temperal_upsample (bool): Whether to upsample on temporal dimension up_flag (bool): Whether to upsample or not non_linearity (str): Type of non-linearity to use """ def __init__( self, in_dim: int, out_dim: int, num_res_blocks: int, dropout: float = 0.0, temperal_upsample: bool = False, up_flag: bool = False, non_linearity: str = "silu", ): super().__init__() self.in_dim = in_dim self.out_dim = out_dim if up_flag: self.avg_shortcut = DupUp3D( in_dim, out_dim, factor_t=2 if temperal_upsample else 1, factor_s=2, ) else: self.avg_shortcut = None # create residual blocks resnets = [] current_dim = in_dim for _ in range(num_res_blocks + 1): resnets.append( WanDistResidualBlock(current_dim, out_dim, dropout, non_linearity) ) current_dim = out_dim self.resnets = nn.ModuleList(resnets) # Add upsampling layer if needed if up_flag: upsample_mode = "upsample3d" if temperal_upsample else "upsample2d" self.upsampler = WanDistResample( out_dim, mode=upsample_mode, upsample_out_dim=out_dim ) else: self.upsampler = None self.gradient_checkpointing = False def forward(self, x): return residual_up_block_forward(self, x) class WanDistUpBlock(nn.Module): """ A block that handles upsampling for the WanVAE decoder. Args: in_dim (int): Input dimension out_dim (int): Output dimension num_res_blocks (int): Number of residual blocks dropout (float): Dropout rate upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d') non_linearity (str): Type of non-linearity to use """ def __init__( self, in_dim: int, out_dim: int, num_res_blocks: int, dropout: float = 0.0, upsample_mode: str | None = None, non_linearity: str = "silu", ): super().__init__() self.in_dim = in_dim self.out_dim = out_dim # Create layers list resnets = [] # Add residual blocks and attention if needed current_dim = in_dim for _ in range(num_res_blocks + 1): resnets.append( WanDistResidualBlock(current_dim, out_dim, dropout, non_linearity) ) current_dim = out_dim self.resnets = nn.ModuleList(resnets) # Add upsampling layer if needed self.upsamplers = None if upsample_mode is not None: self.upsamplers = nn.ModuleList( [WanDistResample(out_dim, mode=upsample_mode)] ) self.gradient_checkpointing = False def forward(self, x): return up_block_forward(self, x) ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Copyright 2025 The Wan Team and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import contextvars from contextlib import contextmanager import torch import torch.distributed as dist import torch.nn as nn from einops import rearrange from sglang.multimodal_gen.configs.models.vaes import WanVAEConfig from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_sp_parallel_rank, get_sp_world_size, ) from sglang.multimodal_gen.runtime.layers.activation import get_act_fn from sglang.multimodal_gen.runtime.models.vaes.common import ( DiagonalGaussianDistribution, ParallelTiledVAE, ) from sglang.multimodal_gen.runtime.models.vaes.parallel.wan_common_utils import ( AvgDown3D, DupUp3D, WanCausalConv3d, WanRMS_norm, WanUpsample, attention_block_forward, bind_context, mid_block_forward, resample_forward, residual_block_forward, residual_down_block_forward, residual_up_block_forward, up_block_forward, ) from sglang.multimodal_gen.runtime.models.vaes.parallel.wan_dist_utils import ( WanDistAttentionBlock, WanDistCausalConv3d, WanDistMidBlock, WanDistResample, WanDistResidualBlock, WanDistResidualDownBlock, WanDistResidualUpBlock, WanDistUpBlock, ensure_local_height, gather_and_trim_height, split_for_parallel_decode, split_for_parallel_encode, ) CACHE_T = 2 is_first_frame = contextvars.ContextVar("is_first_frame", default=False) feat_cache = contextvars.ContextVar("feat_cache", default=None) feat_idx = contextvars.ContextVar("feat_idx", default=0) first_chunk = contextvars.ContextVar("first_chunk", default=None) bind_context(is_first_frame, feat_cache, feat_idx, CACHE_T, first_chunk) @contextmanager def forward_context( first_frame_arg=False, feat_cache_arg=None, feat_idx_arg=None, first_chunk_arg=None ): is_first_frame_token = is_first_frame.set(first_frame_arg) feat_cache_token = feat_cache.set(feat_cache_arg) feat_idx_token = feat_idx.set(feat_idx_arg) first_chunk_token = first_chunk.set(first_chunk_arg) try: yield finally: is_first_frame.reset(is_first_frame_token) feat_cache.reset(feat_cache_token) feat_idx.reset(feat_idx_token) first_chunk.reset(first_chunk_token) class WanResample(nn.Module): r""" A custom resampling module for 2D and 3D data. Args: dim (int): The number of input/output channels. mode (str): The resampling mode. Must be one of: - 'none': No resampling (identity operation). - 'upsample2d': 2D upsampling with nearest-exact interpolation and convolution. - 'upsample3d': 3D upsampling with nearest-exact interpolation, convolution, and causal 3D convolution. - 'downsample2d': 2D downsampling with zero-padding and convolution. - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution. """ def __init__(self, dim: int, mode: str, upsample_out_dim: int = None) -> None: super().__init__() self.dim = dim self.mode = mode # default to dim //2 if upsample_out_dim is None: upsample_out_dim = dim // 2 # layers if mode == "upsample2d": self.resample = nn.Sequential( WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, upsample_out_dim, 3, padding=1), ) elif mode == "upsample3d": self.resample = nn.Sequential( WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, upsample_out_dim, 3, padding=1), ) self.time_conv = WanCausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) elif mode == "downsample2d": self.resample = nn.Sequential( nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) ) elif mode == "downsample3d": self.resample = nn.Sequential( nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) ) self.time_conv = WanCausalConv3d( dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0) ) else: self.resample = nn.Identity() def forward(self, x): return resample_forward(self, x) class WanResidualBlock(nn.Module): r""" A custom residual block module. Args: in_dim (int): Number of input channels. out_dim (int): Number of output channels. dropout (float, optional): Dropout rate for the dropout layer. Default is 0.0. non_linearity (str, optional): Type of non-linearity to use. Default is "silu". """ def __init__( self, in_dim: int, out_dim: int, dropout: float = 0.0, non_linearity: str = "silu", ) -> None: super().__init__() self.in_dim = in_dim self.out_dim = out_dim self.nonlinearity = get_act_fn(non_linearity) # layers self.norm1 = WanRMS_norm(in_dim, images=False) self.conv1 = WanCausalConv3d(in_dim, out_dim, 3, padding=1) self.norm2 = WanRMS_norm(out_dim, images=False) self.dropout = nn.Dropout(dropout) self.conv2 = WanCausalConv3d(out_dim, out_dim, 3, padding=1) self.conv_shortcut = ( WanCausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() ) def forward(self, x): return residual_block_forward(self, x) class WanAttentionBlock(nn.Module): r""" Causal self-attention with a single head. Args: dim (int): The number of channels in the input tensor. """ def __init__(self, dim) -> None: super().__init__() self.dim = dim # layers self.norm = WanRMS_norm(dim) self.to_qkv = nn.Conv2d(dim, dim * 3, 1) self.proj = nn.Conv2d(dim, dim, 1) def forward(self, x): return attention_block_forward(self, x) class WanMidBlock(nn.Module): """ Middle block for WanVAE encoder and decoder. Args: dim (int): Number of input/output channels. dropout (float): Dropout rate. non_linearity (str): Type of non-linearity to use. """ def __init__( self, dim: int, dropout: float = 0.0, non_linearity: str = "silu", num_layers: int = 1, ): super().__init__() self.dim = dim # Create the components resnets = [WanResidualBlock(dim, dim, dropout, non_linearity)] attentions = [] for _ in range(num_layers): attentions.append(WanAttentionBlock(dim)) resnets.append(WanResidualBlock(dim, dim, dropout, non_linearity)) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) self.gradient_checkpointing = False def forward(self, x): return mid_block_forward(self, x) class WanResidualDownBlock(nn.Module): def __init__( self, in_dim, out_dim, dropout, num_res_blocks, temperal_downsample=False, down_flag=False, ): super().__init__() # Shortcut path with downsample self.avg_shortcut = AvgDown3D( in_dim, out_dim, factor_t=2 if temperal_downsample else 1, factor_s=2 if down_flag else 1, ) # Main path with residual blocks and downsample resnets = [] for _ in range(num_res_blocks): resnets.append(WanResidualBlock(in_dim, out_dim, dropout)) in_dim = out_dim self.resnets = nn.ModuleList(resnets) # Add the final downsample block if down_flag: mode = "downsample3d" if temperal_downsample else "downsample2d" self.downsampler = WanResample(out_dim, mode=mode) else: self.downsampler = None def forward(self, x): return residual_down_block_forward(self, x) class WanEncoder3d(nn.Module): r""" A 3D encoder module. Args: dim (int): The base number of channels in the first layer. z_dim (int): The dimensionality of the latent space. dim_mult (list of int): Multipliers for the number of channels in each block. num_res_blocks (int): Number of residual blocks in each block. attn_scales (list of float): Scales at which to apply attention mechanisms. temperal_downsample (list of bool): Whether to downsample temporally in each block. dropout (float): Dropout rate for the dropout layers. non_linearity (str): Type of non-linearity to use. """ def __init__( self, in_channels: int = 3, dim=128, z_dim=4, dim_mult=(1, 2, 4, 4), num_res_blocks=2, attn_scales=(), temperal_downsample=(True, True, False), dropout=0.0, non_linearity: str = "silu", is_residual: bool = False, # wan 2.2 vae use a residual downblock use_parallel_encode: bool = False, ): super().__init__() self.dim = dim self.z_dim = z_dim dim_mult = list(dim_mult) self.dim_mult = dim_mult self.num_res_blocks = num_res_blocks self.attn_scales = list(attn_scales) self.temperal_downsample = list(temperal_downsample) self.nonlinearity = get_act_fn(non_linearity) self.use_parallel_encode = use_parallel_encode self.downsample_count = max(len(dim_mult) - 1, 0) # dimensions dims = [dim * u for u in [1] + dim_mult] scale = 1.0 world_size = 1 if dist.is_initialized(): world_size = get_sp_world_size() if use_parallel_encode and world_size > 1: CausalConv3d = WanDistCausalConv3d ResidualDownBlock = WanDistResidualDownBlock ResidualBlock = WanDistResidualBlock AttentionBlock = WanDistAttentionBlock Resample = WanDistResample MidBlock = WanDistMidBlock else: CausalConv3d = WanCausalConv3d ResidualDownBlock = WanResidualDownBlock ResidualBlock = WanResidualBlock AttentionBlock = WanAttentionBlock Resample = WanResample MidBlock = WanMidBlock # init block self.conv_in = CausalConv3d(in_channels, dims[0], 3, padding=1) # downsample blocks self.down_blocks = nn.ModuleList([]) for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)): # residual (+attention) blocks if is_residual: self.down_blocks.append( ResidualDownBlock( in_dim, out_dim, dropout, num_res_blocks, temperal_downsample=( temperal_downsample[i] if i != len(dim_mult) - 1 else False ), down_flag=i != len(dim_mult) - 1, ) ) else: for _ in range(num_res_blocks): self.down_blocks.append(ResidualBlock(in_dim, out_dim, dropout)) if scale in attn_scales: self.down_blocks.append(AttentionBlock(out_dim)) in_dim = out_dim # downsample block if i != len(dim_mult) - 1: mode = "downsample3d" if temperal_downsample[i] else "downsample2d" self.down_blocks.append(Resample(out_dim, mode=mode)) scale /= 2.0 # middle blocks self.mid_block = MidBlock(out_dim, dropout, non_linearity, num_layers=1) # output blocks self.norm_out = WanRMS_norm(out_dim, images=False) self.conv_out = CausalConv3d(out_dim, z_dim, 3, padding=1) self.gradient_checkpointing = False self.world_size = 1 self.rank = 0 if dist.is_initialized(): self.world_size = get_sp_world_size() self.rank = get_sp_parallel_rank() def forward(self, x): expected_local_height = None expected_height = None if self.use_parallel_encode and self.world_size > 1: x, expected_height, expected_local_height = split_for_parallel_encode( x, self.downsample_count, self.world_size, self.rank ) _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat( [ _feat_cache[idx][:, :, -1, :, :] .unsqueeze(2) .to(cache_x.device), cache_x, ], dim=2, ) x = self.conv_in(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) else: x = self.conv_in(x) ## downsamples for layer in self.down_blocks: x = layer(x) ## middle if self.use_parallel_encode and self.world_size > 1: x = ensure_local_height(x, expected_local_height) x = self.mid_block(x) ## head x = self.norm_out(x) x = self.nonlinearity(x) _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat( [ _feat_cache[idx][:, :, -1, :, :] .unsqueeze(2) .to(cache_x.device), cache_x, ], dim=2, ) x = self.conv_out(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) else: x = self.conv_out(x) if self.use_parallel_encode and self.world_size > 1: x = gather_and_trim_height(x, expected_height) return x # adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/autoencoder_kl_wan.py class WanResidualUpBlock(nn.Module): """ A block that handles upsampling for the WanVAE decoder. Args: in_dim (int): Input dimension out_dim (int): Output dimension num_res_blocks (int): Number of residual blocks dropout (float): Dropout rate temperal_upsample (bool): Whether to upsample on temporal dimension up_flag (bool): Whether to upsample or not non_linearity (str): Type of non-linearity to use """ def __init__( self, in_dim: int, out_dim: int, num_res_blocks: int, dropout: float = 0.0, temperal_upsample: bool = False, up_flag: bool = False, non_linearity: str = "silu", ): super().__init__() self.in_dim = in_dim self.out_dim = out_dim if up_flag: self.avg_shortcut = DupUp3D( in_dim, out_dim, factor_t=2 if temperal_upsample else 1, factor_s=2, ) else: self.avg_shortcut = None # create residual blocks resnets = [] current_dim = in_dim for _ in range(num_res_blocks + 1): resnets.append( WanResidualBlock(current_dim, out_dim, dropout, non_linearity) ) current_dim = out_dim self.resnets = nn.ModuleList(resnets) # Add upsampling layer if needed if up_flag: upsample_mode = "upsample3d" if temperal_upsample else "upsample2d" self.upsampler = WanResample( out_dim, mode=upsample_mode, upsample_out_dim=out_dim ) else: self.upsampler = None self.gradient_checkpointing = False def forward(self, x): return residual_up_block_forward(self, x) class WanUpBlock(nn.Module): """ A block that handles upsampling for the WanVAE decoder. Args: in_dim (int): Input dimension out_dim (int): Output dimension num_res_blocks (int): Number of residual blocks dropout (float): Dropout rate upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d') non_linearity (str): Type of non-linearity to use """ def __init__( self, in_dim: int, out_dim: int, num_res_blocks: int, dropout: float = 0.0, upsample_mode: str | None = None, non_linearity: str = "silu", ): super().__init__() self.in_dim = in_dim self.out_dim = out_dim # Create layers list resnets = [] # Add residual blocks and attention if needed current_dim = in_dim for _ in range(num_res_blocks + 1): resnets.append( WanResidualBlock(current_dim, out_dim, dropout, non_linearity) ) current_dim = out_dim self.resnets = nn.ModuleList(resnets) # Add upsampling layer if needed self.upsamplers = None if upsample_mode is not None: self.upsamplers = nn.ModuleList([WanResample(out_dim, mode=upsample_mode)]) self.gradient_checkpointing = False def forward(self, x): return up_block_forward(self, x) class WanDecoder3d(nn.Module): r""" A 3D decoder module. Args: dim (int): The base number of channels in the first layer. z_dim (int): The dimensionality of the latent space. dim_mult (list of int): Multipliers for the number of channels in each block. num_res_blocks (int): Number of residual blocks in each block. attn_scales (list of float): Scales at which to apply attention mechanisms. temperal_upsample (list of bool): Whether to upsample temporally in each block. dropout (float): Dropout rate for the dropout layers. non_linearity (str): Type of non-linearity to use. """ def __init__( self, dim=128, z_dim=4, dim_mult=(1, 2, 4, 4), num_res_blocks=2, attn_scales=(), temperal_upsample=(False, True, True), dropout=0.0, non_linearity: str = "silu", out_channels: int = 3, is_residual: bool = False, use_parallel_decode: bool = False, ): super().__init__() self.dim = dim self.z_dim = z_dim dim_mult = list(dim_mult) self.dim_mult = dim_mult self.num_res_blocks = num_res_blocks self.attn_scales = list(attn_scales) self.temperal_upsample = list(temperal_upsample) self.nonlinearity = get_act_fn(non_linearity) self.use_parallel_decode = use_parallel_decode self.upsample_count = 0 # dimensions dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] world_size = 1 if dist.is_initialized(): world_size = get_sp_world_size() if use_parallel_decode and world_size > 1: CausalConv3d = WanDistCausalConv3d MidBlock = WanDistMidBlock ResidualUpBlock = WanDistResidualUpBlock UpBlock = WanDistUpBlock else: CausalConv3d = WanCausalConv3d MidBlock = WanMidBlock ResidualUpBlock = WanResidualUpBlock UpBlock = WanUpBlock # init block self.conv_in = CausalConv3d(z_dim, dims[0], 3, padding=1) # middle blocks self.mid_block = MidBlock(dims[0], dropout, non_linearity, num_layers=1) # upsample blocks self.upsample_count = 0 self.up_blocks = nn.ModuleList([]) for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)): # residual (+attention) blocks if i > 0 and not is_residual: # wan vae 2.1 in_dim = in_dim // 2 # determine if we need upsampling up_flag = i != len(dim_mult) - 1 # determine upsampling mode, if not upsampling, set to None upsample_mode = None if up_flag and temperal_upsample[i]: upsample_mode = "upsample3d" elif up_flag: upsample_mode = "upsample2d" # Create and add the upsampling block if is_residual: up_block = ResidualUpBlock( in_dim=in_dim, out_dim=out_dim, num_res_blocks=num_res_blocks, dropout=dropout, temperal_upsample=temperal_upsample[i] if up_flag else False, up_flag=up_flag, non_linearity=non_linearity, ) else: up_block = UpBlock( in_dim=in_dim, out_dim=out_dim, num_res_blocks=num_res_blocks, dropout=dropout, upsample_mode=upsample_mode, non_linearity=non_linearity, ) self.up_blocks.append(up_block) if up_flag: self.upsample_count += 1 # output blocks self.norm_out = WanRMS_norm(out_dim, images=False) self.conv_out = CausalConv3d(out_dim, out_channels, 3, padding=1) self.gradient_checkpointing = False self.world_size = 1 self.rank = 0 if dist.is_initialized(): self.world_size = get_sp_world_size() self.rank = get_sp_parallel_rank() def forward(self, x): expected_height = None if self.use_parallel_decode and self.world_size > 1: x, expected_height = split_for_parallel_decode( x, self.upsample_count, self.world_size, self.rank ) ## conv1 _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat( [ _feat_cache[idx][:, :, -1, :, :] .unsqueeze(2) .to(cache_x.device), cache_x, ], dim=2, ) x = self.conv_in(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) else: x = self.conv_in(x) ## middle x = self.mid_block(x) ## upsamples for up_block in self.up_blocks: x = up_block(x) ## head x = self.norm_out(x) x = self.nonlinearity(x) _feat_cache = feat_cache.get() _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat( [ _feat_cache[idx][:, :, -1, :, :] .unsqueeze(2) .to(cache_x.device), cache_x, ], dim=2, ) x = self.conv_out(x, _feat_cache[idx]) _feat_cache[idx] = cache_x _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) else: x = self.conv_out(x) if self.use_parallel_decode and self.world_size > 1: x = gather_and_trim_height(x, expected_height) return x def patchify(x, patch_size): if patch_size == 1: return x if x.dim() == 4: x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) elif x.dim() == 5: x = rearrange( x, "b c f (h q) (w r) -> b (c r q) f h w", q=patch_size, r=patch_size, ) else: raise ValueError(f"Invalid input shape: {x.shape}") return x def unpatchify(x, patch_size): if patch_size == 1: return x if x.dim() == 4: x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) elif x.dim() == 5: x = rearrange( x, "b (c r q) f h w -> b c f (h q) (w r)", q=patch_size, r=patch_size, ) return x class AutoencoderKLWan(ParallelTiledVAE): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Introduced in [Wan 2.1]. """ _supports_gradient_checkpointing = False def __init__( self, config: WanVAEConfig, ) -> None: nn.Module.__init__(self) ParallelTiledVAE.__init__(self, config) self.z_dim = config.z_dim self.temperal_downsample = list(config.temperal_downsample) self.temperal_upsample = list(config.temperal_downsample)[::-1] if config.decoder_base_dim is None: decoder_base_dim = config.base_dim else: decoder_base_dim = config.decoder_base_dim self.latents_mean = list(config.latents_mean) self.latents_std = list(config.latents_std) self.shift_factor = config.shift_factor self.use_parallel_encode = getattr(config, "use_parallel_encode", False) self.use_parallel_decode = getattr(config, "use_parallel_decode", False) if config.load_encoder: self.encoder = WanEncoder3d( in_channels=config.in_channels, dim=config.base_dim, z_dim=self.z_dim * 2, dim_mult=config.dim_mult, num_res_blocks=config.num_res_blocks, attn_scales=config.attn_scales, temperal_downsample=self.temperal_downsample, dropout=config.dropout, is_residual=config.is_residual, use_parallel_encode=self.use_parallel_encode, ) self.quant_conv = WanCausalConv3d(self.z_dim * 2, self.z_dim * 2, 1) self.post_quant_conv = WanCausalConv3d(self.z_dim, self.z_dim, 1) if config.load_decoder: self.decoder = WanDecoder3d( dim=decoder_base_dim, z_dim=self.z_dim, dim_mult=config.dim_mult, num_res_blocks=config.num_res_blocks, attn_scales=config.attn_scales, temperal_upsample=self.temperal_upsample, dropout=config.dropout, out_channels=config.out_channels, is_residual=config.is_residual, use_parallel_decode=self.use_parallel_decode, ) self.use_feature_cache = config.use_feature_cache def clear_cache(self) -> None: def _count_conv3d(model) -> int: count = 0 for m in model.modules(): if isinstance(m, WanCausalConv3d) or isinstance(m, WanDistCausalConv3d): count += 1 return count if self.config.load_decoder: self._conv_num = _count_conv3d(self.decoder) self._conv_idx = 0 self._feat_map = [None] * self._conv_num # cache encode if self.config.load_encoder: self._enc_conv_num = _count_conv3d(self.encoder) self._enc_conv_idx = 0 self._enc_feat_map = [None] * self._enc_conv_num def encode(self, x: torch.Tensor) -> torch.Tensor: if self.use_feature_cache: self.clear_cache() if self.config.patch_size is not None: x = patchify(x, patch_size=self.config.patch_size) with forward_context( feat_cache_arg=self._enc_feat_map, feat_idx_arg=self._enc_conv_idx ): t = x.shape[2] iter_ = 1 + (t - 1) // 4 for i in range(iter_): feat_idx.set(0) if i == 0: out = self.encoder(x[:, :, :1, :, :]) else: out_ = self.encoder(x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :]) out = torch.cat([out, out_], 2) enc = self.quant_conv(out) mu, logvar = enc[:, : self.z_dim, :, :, :], enc[:, self.z_dim :, :, :, :] enc = torch.cat([mu, logvar], dim=1) enc = DiagonalGaussianDistribution(enc) self.clear_cache() else: for block in self.encoder.down_blocks: if isinstance(block, WanResample) and block.mode == "downsample3d": _padding = list(block.time_conv._padding) _padding[4] = 2 block.time_conv._padding = tuple(_padding) enc = ParallelTiledVAE.encode(self, x) return enc def _encode(self, x: torch.Tensor, first_frame=False) -> torch.Tensor: with forward_context(first_frame_arg=first_frame): out = self.encoder(x) enc = self.quant_conv(out) mu, logvar = enc[:, : self.z_dim, :, :, :], enc[:, self.z_dim :, :, :, :] enc = torch.cat([mu, logvar], dim=1) return enc def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: first_frame = x[:, :, 0, :, :].unsqueeze(2) first_frame = self._encode(first_frame, first_frame=True) enc = ParallelTiledVAE.tiled_encode(self, x) enc = enc[:, :, 1:] enc = torch.cat([first_frame, enc], dim=2) return enc def spatial_tiled_encode(self, x: torch.Tensor) -> torch.Tensor: first_frame = x[:, :, 0, :, :].unsqueeze(2) first_frame = self._encode(first_frame, first_frame=True) enc = ParallelTiledVAE.spatial_tiled_encode(self, x) enc = enc[:, :, 1:] enc = torch.cat([first_frame, enc], dim=2) return enc def decode(self, z: torch.Tensor) -> torch.Tensor: if self.use_feature_cache: self.clear_cache() iter_ = z.shape[2] x = self.post_quant_conv(z) with forward_context( feat_cache_arg=self._feat_map, feat_idx_arg=self._conv_idx ): for i in range(iter_): feat_idx.set(0) if i == 0: first_chunk.set(True) out = self.decoder(x[:, :, i : i + 1, :, :]) else: first_chunk.set(False) out_ = self.decoder(x[:, :, i : i + 1, :, :]) out = torch.cat([out, out_], 2) if self.config.patch_size is not None: out = unpatchify(out, patch_size=self.config.patch_size) out = out.float() out = torch.clamp(out, min=-1.0, max=1.0) self.clear_cache() else: out = ParallelTiledVAE.decode(self, z) return out def _decode(self, z: torch.Tensor, first_frame=False) -> torch.Tensor: x = self.post_quant_conv(z) with forward_context(first_frame_arg=first_frame): out = self.decoder(x) out = torch.clamp(out, min=-1.0, max=1.0) return out def tiled_decode(self, z: torch.Tensor) -> torch.Tensor: self.blend_num_frames *= 2 dec = ParallelTiledVAE.tiled_decode(self, z) start_frame_idx = self.temporal_compression_ratio - 1 dec = dec[:, :, start_frame_idx:] return dec def spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor: dec = ParallelTiledVAE.spatial_tiled_decode(self, z) start_frame_idx = self.temporal_compression_ratio - 1 dec = dec[:, :, start_frame_idx:] return dec def parallel_tiled_decode(self, z: torch.FloatTensor) -> torch.FloatTensor: self.blend_num_frames *= 2 dec = ParallelTiledVAE.parallel_tiled_decode(self, z) start_frame_idx = self.temporal_compression_ratio - 1 dec = dec[:, :, start_frame_idx:] return dec def forward( self, sample: torch.Tensor, sample_posterior: bool = False, generator: torch.Generator | None = None, ) -> torch.Tensor: """ Args: sample (`torch.Tensor`): Input sample. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`DecoderOutput`] instead of a plain tuple. """ x = sample posterior = self.encode(x).latent_dist if sample_posterior: z = posterior.sample(generator=generator) else: z = posterior.mode() dec = self.decode(z) return dec EntryClass = AutoencoderKLWan ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vision_utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import os import tempfile from collections.abc import Callable from urllib.parse import unquote, urlparse import imageio import numpy as np import PIL.Image import PIL.ImageOps import requests import torch from packaging import version if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"): PIL_INTERPOLATION = { "linear": PIL.Image.Resampling.BILINEAR, "bilinear": PIL.Image.Resampling.BILINEAR, "bicubic": PIL.Image.Resampling.BICUBIC, "lanczos": PIL.Image.Resampling.LANCZOS, "nearest": PIL.Image.Resampling.NEAREST, } else: PIL_INTERPOLATION = { "linear": PIL.Image.LINEAR, "bilinear": PIL.Image.BILINEAR, "bicubic": PIL.Image.BICUBIC, "lanczos": PIL.Image.LANCZOS, "nearest": PIL.Image.NEAREST, } def pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray: r""" Convert a PIL image or a list of PIL images to NumPy arrays. Args: images (`PIL.Image.Image` or `List[PIL.Image.Image]`): The PIL image or list of images to convert to NumPy format. Returns: `np.ndarray`: A NumPy array representation of the images. """ if not isinstance(images, list): images = [images] images = [np.array(image).astype(np.float32) / 255.0 for image in images] images_arr: np.ndarray = np.stack(images, axis=0) return images_arr def numpy_to_pt(images: np.ndarray) -> torch.Tensor: r""" Convert a NumPy image to a PyTorch tensor. Args: images (`np.ndarray`): The NumPy image array to convert to PyTorch format. Returns: `torch.Tensor`: A PyTorch tensor representation of the images. """ if images.ndim == 3: images = images[..., None] images = torch.from_numpy(images.transpose(0, 3, 1, 2)) return images def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: r""" Normalize an image array to [-1,1]. Args: images (`np.ndarray` or `torch.Tensor`): The image array to normalize. Returns: `np.ndarray` or `torch.Tensor`: The normalized image array. """ return 2.0 * images - 1.0 # adapted from diffusers.utils import load_image def load_image( image: str | PIL.Image.Image, convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None, ) -> PIL.Image.Image: """ Loads `image` to a PIL Image. Args: image (`str` or `PIL.Image.Image`): The image to convert to the PIL Image format. convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], *optional*): A conversion method to apply to the image after loading it. When set to `None` the image will be converted "RGB". """ if isinstance(image, str): if image.startswith("http://") or image.startswith("https://"): image = PIL.Image.open(requests.get(image, stream=True).raw) elif os.path.isfile(image): image = PIL.Image.open(image) else: raise ValueError( f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {image} is not a valid path." ) elif isinstance(image, PIL.Image.Image): image = image else: raise ValueError( "Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image." ) image = PIL.ImageOps.exif_transpose(image) if convert_method is not None: image = convert_method(image) else: image = image.convert("RGB") return image # adapted from diffusers.utils import load_video def load_video( video: str, convert_method: ( Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]] | None ) = None, ) -> list[PIL.Image.Image]: """ Loads `video` to a list of PIL Image. Args: video (`str`): A URL or Path to a video to convert to a list of PIL Image format. convert_method (Callable[[List[PIL.Image.Image]], List[PIL.Image.Image]], *optional*): A conversion method to apply to the video after loading it. When set to `None` the images will be converted to "RGB". Returns: `List[PIL.Image.Image]`: The video as a list of PIL images. """ is_url = video.startswith("http://") or video.startswith("https://") is_file = os.path.isfile(video) was_tempfile_created = False if not (is_url or is_file): raise ValueError( f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {video} is not a valid path." ) if is_url: response = requests.get(video, stream=True) if response.status_code != 200: raise ValueError( f"Failed to download video. Status code: {response.status_code}" ) parsed_url = urlparse(video) file_name = os.path.basename(unquote(parsed_url.path)) suffix = os.path.splitext(file_name)[1] or ".mp4" with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as temp_file: video_path = temp_file.name video_data = response.iter_content(chunk_size=8192) for chunk in video_data: temp_file.write(chunk) video = video_path pil_images = [] if video.endswith(".gif"): gif = PIL.Image.open(video) try: while True: pil_images.append(gif.copy()) gif.seek(gif.tell() + 1) except EOFError: pass else: try: imageio.plugins.ffmpeg.get_exe() except AttributeError: raise AttributeError( "`Unable to find an ffmpeg installation on your machine. Please install via `pip install imageio-ffmpeg" ) from None with imageio.get_reader(video) as reader: # Read all frames for frame in reader: pil_images.append(PIL.Image.fromarray(frame)) if was_tempfile_created: os.remove(video_path) if convert_method is not None: pil_images = convert_method(pil_images) return pil_images def get_default_height_width( image: PIL.Image.Image | np.ndarray | torch.Tensor, vae_scale_factor: int, height: int | None = None, width: int | None = None, ) -> tuple[int, int]: r""" Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`. Args: image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor]`): The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch tensor, it should have shape `[batch, channels, height, width]`. height (`Optional[int]`, *optional*, defaults to `None`): The height of the preprocessed image. If `None`, the height of the `image` input will be used. width (`Optional[int]`, *optional*, defaults to `None`): The width of the preprocessed image. If `None`, the width of the `image` input will be used. Returns: `Tuple[int, int]`: A tuple containing the height and width, both resized to the nearest integer multiple of `vae_scale_factor`. """ if height is None: if isinstance(image, PIL.Image.Image): height = image.height elif isinstance(image, torch.Tensor): height = image.shape[2] else: height = image.shape[1] if width is None: if isinstance(image, PIL.Image.Image): width = image.width elif isinstance(image, torch.Tensor): width = image.shape[3] else: width = image.shape[2] width, height = ( x - x % vae_scale_factor for x in (width, height) ) # resize to integer multiple of vae_scale_factor return height, width def resize( image: PIL.Image.Image | np.ndarray | torch.Tensor, height: int, width: int, resize_mode: str = "default", # "default", "fill", "crop" resample: str = "lanczos", ) -> PIL.Image.Image | np.ndarray | torch.Tensor: """ Resize image. Args: image (`PIL.Image.Image`, `np.ndarray` or `torch.Tensor`): The image input, can be a PIL image, numpy array or pytorch tensor. height (`int`): The height to resize to. width (`int`): The width to resize to. resize_mode (`str`, *optional*, defaults to `default`): The resize mode to use, can be one of `default` or `fill`. If `default`, will resize the image to fit within the specified width and height, and it may not maintaining the original aspect ratio. If `fill`, will resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, filling empty with data from image. If `crop`, will resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only supported for PIL image input. Returns: `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`: The resized image. """ if resize_mode != "default" and not isinstance(image, PIL.Image.Image): raise ValueError( f"Only PIL image input is supported for resize_mode {resize_mode}" ) assert isinstance(image, PIL.Image.Image) if resize_mode == "default": image = image.resize((width, height), resample=PIL_INTERPOLATION[resample]) else: raise ValueError(f"resize_mode {resize_mode} is not supported") return image ================================================ FILE: python/sglang/multimodal_gen/runtime/models/vocoder/ltx_2_vocoder.py ================================================ import math from abc import ABC from typing import Tuple import torch import torch.nn as nn import torch.nn.functional as F from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import LTXVocoderConfig class ResBlock(nn.Module): def __init__( self, channels: int, kernel_size: int = 3, stride: int = 1, dilations: Tuple[int, ...] = (1, 3, 5), leaky_relu_negative_slope: float = 0.1, padding_mode: str = "same", ): super().__init__() self.dilations = dilations self.negative_slope = leaky_relu_negative_slope self.convs1 = nn.ModuleList( [ nn.Conv1d( channels, channels, kernel_size, stride=stride, dilation=dilation, padding=padding_mode, ) for dilation in dilations ] ) self.convs2 = nn.ModuleList( [ nn.Conv1d( channels, channels, kernel_size, stride=stride, dilation=1, padding=padding_mode, ) for _ in range(len(dilations)) ] ) def forward(self, x: torch.Tensor) -> torch.Tensor: for conv1, conv2 in zip(self.convs1, self.convs2): xt = F.leaky_relu(x, negative_slope=self.negative_slope) xt = conv1(xt) xt = F.leaky_relu(xt, negative_slope=self.negative_slope) xt = conv2(xt) x = x + xt return x class LTX2Vocoder(ABC, nn.Module): r""" LTX 2.0 vocoder for converting generated mel spectrograms back to audio waveforms. """ def __init__( self, config: LTXVocoderConfig, ): super().__init__() self.config = config self.sample_rate = ( getattr(config.arch_config, "sample_rate", None) or getattr(config.arch_config, "sampling_rate", None) or getattr(config.arch_config, "audio_sample_rate", None) ) in_channels = config.arch_config.in_channels hidden_channels = config.arch_config.hidden_channels out_channels = config.arch_config.out_channels upsample_kernel_sizes = config.arch_config.upsample_kernel_sizes upsample_factors = config.arch_config.upsample_factors resnet_kernel_sizes = config.arch_config.resnet_kernel_sizes resnet_dilations = config.arch_config.resnet_dilations leaky_relu_negative_slope = config.arch_config.leaky_relu_negative_slope self.num_upsample_layers = len(upsample_kernel_sizes) self.resnets_per_upsample = len(resnet_kernel_sizes) self.out_channels = out_channels self.total_upsample_factor = math.prod(upsample_factors) self.negative_slope = leaky_relu_negative_slope if self.num_upsample_layers != len(upsample_factors): raise ValueError( f"`upsample_kernel_sizes` and `upsample_factors` should be lists of the same length but are length" f" {self.num_upsample_layers} and {len(upsample_factors)}, respectively." ) if self.resnets_per_upsample != len(resnet_dilations): raise ValueError( f"`resnet_kernel_sizes` and `resnet_dilations` should be lists of the same length but are length" f" {len(self.resnets_per_upsample)} and {len(resnet_dilations)}, respectively." ) self.conv_in = nn.Conv1d( in_channels, hidden_channels, kernel_size=7, stride=1, padding=3 ) self.upsamplers = nn.ModuleList() self.resnets = nn.ModuleList() input_channels = hidden_channels for i, (stride, kernel_size) in enumerate( zip(upsample_factors, upsample_kernel_sizes) ): output_channels = input_channels // 2 self.upsamplers.append( nn.ConvTranspose1d( input_channels, # hidden_channels // (2 ** i) output_channels, # hidden_channels // (2 ** (i + 1)) kernel_size, stride=stride, padding=(kernel_size - stride) // 2, ) ) for kernel_size, dilations in zip(resnet_kernel_sizes, resnet_dilations): self.resnets.append( ResBlock( output_channels, kernel_size, dilations=dilations, leaky_relu_negative_slope=leaky_relu_negative_slope, ) ) input_channels = output_channels self.conv_out = nn.Conv1d(output_channels, out_channels, 7, stride=1, padding=3) def forward( self, hidden_states: torch.Tensor, time_last: bool = False ) -> torch.Tensor: r""" Forward pass of the vocoder. Args: hidden_states (`torch.Tensor`): Input Mel spectrogram tensor of shape `(batch_size, num_channels, time, num_mel_bins)` if `time_last` is `False` (the default) or shape `(batch_size, num_channels, num_mel_bins, time)` if `time_last` is `True`. time_last (`bool`, *optional*, defaults to `False`): Whether the last dimension of the input is the time/frame dimension or the Mel bins dimension. Returns: `torch.Tensor`: Audio waveform tensor of shape (batch_size, out_channels, audio_length) """ # Ensure that the time/frame dimension is last if not time_last: hidden_states = hidden_states.transpose(2, 3) # Combine channels and frequency (mel bins) dimensions hidden_states = hidden_states.flatten(1, 2) hidden_states = self.conv_in(hidden_states) for i in range(self.num_upsample_layers): hidden_states = F.leaky_relu( hidden_states, negative_slope=self.negative_slope ) hidden_states = self.upsamplers[i](hidden_states) # Run all resnets in parallel on hidden_states start = i * self.resnets_per_upsample end = (i + 1) * self.resnets_per_upsample resnet_outputs = torch.stack( [self.resnets[j](hidden_states) for j in range(start, end)], dim=0 ) hidden_states = torch.mean(resnet_outputs, dim=0) # NOTE: unlike the first leaky ReLU, this leaky ReLU is set to use the default F.leaky_relu negative slope of # 0.01 (whereas the others usually use a slope of 0.1). Not sure if this is intended hidden_states = F.leaky_relu(hidden_states, negative_slope=0.01) hidden_states = self.conv_out(hidden_states) hidden_states = torch.tanh(hidden_states) return hidden_states EntryClass = LTX2Vocoder ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import os import re from typing import Any, Generator import torch from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.models.schedulers.scheduling_comfyui_passthrough import ( ComfyUIPassThroughScheduler, ) from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.stages import ( ComfyUILatentPreparationStage, DenoisingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) class ComfyUIFluxPipeline(LoRAPipeline, ComposedPipelineBase): """ Simplified pipeline for ComfyUI integration with only denoising stage. This pipeline requires pre-processed inputs: - prompt_embeds: Pre-encoded text embeddings (list of tensors) - negative_prompt_embeds: Pre-encoded negative prompt embeddings (if using CFG) - latents: Optional initial noise latents (will be generated if not provided) Usage: generator = DiffGenerator.from_pretrained( model_path="path/to/model", pipeline_class_name="ComfyUIFluxPipeline", device="cuda", ) """ pipeline_name = "ComfyUIFluxPipeline" # Configuration classes for safetensors files without model_index.json from sglang.multimodal_gen.configs.pipeline_configs.flux import FluxPipelineConfig from sglang.multimodal_gen.configs.sample.flux import FluxSamplingParams pipeline_config_cls = FluxPipelineConfig sampling_params_cls = FluxSamplingParams _required_config_modules = [ "transformer", "scheduler", ] def initialize_pipeline(self, server_args: ServerArgs): """ Initialize the pipeline with ComfyUI pass-through scheduler. This scheduler does not modify latents, allowing ComfyUI to handle denoising. """ self.modules["scheduler"] = ComfyUIPassThroughScheduler( num_train_timesteps=1000 ) if hasattr(server_args.pipeline_config, "vae_config"): vae_config = server_args.pipeline_config.vae_config if hasattr(vae_config, "post_init") and not hasattr( vae_config, "_post_init_called" ): vae_config.post_init() logger.info( "Called vae_config.post_init() to set spatial_compression_ratio. " f"spatial_compression_ratio={vae_config.arch_config.spatial_compression_ratio}" ) def load_modules( self, server_args: ServerArgs, loaded_modules: dict[str, torch.nn.Module] | None = None, ) -> dict[str, Any]: """ Load modules for ComfyUIFluxPipeline. If model_path is a safetensors file, load transformer directly from it without requiring model_index.json. Otherwise, fall back to default loading. """ if os.path.isfile(self.model_path) and self.model_path.endswith(".safetensors"): logger.info( "Detected safetensors file, loading transformer directly from: %s", self.model_path, ) return self._load_transformer_from_safetensors(server_args, loaded_modules) else: logger.info( "Model path is a directory, using default loading method: %s", self.model_path, ) return super().load_modules(server_args, loaded_modules) def _load_and_convert_weights_from_safetensors( self, model_cls: type, dit_config: FluxConfig, hf_config: dict, safetensors_list: list[str], updated_mapping: dict, qkv_size: int, mlp_hidden_dim: int, has_guidance_embeds: bool, default_dtype: torch.dtype, ) -> tuple[torch.nn.Module, dict]: """ Load and convert weights from safetensors file, then load them into the model. """ from sglang.multimodal_gen.runtime.loader.utils import ( get_param_names_mapping, set_default_torch_dtype, ) from sglang.multimodal_gen.runtime.loader.weight_utils import ( safetensors_weights_iterator, ) logger.info( "Converting ComfyUI Flux weights to SGLang format and loading model..." ) # Create model on target device device = get_local_torch_device() with set_default_torch_dtype(default_dtype): model = model_cls(**{"config": dit_config, "hf_config": hf_config}) model = model.to(device) # Verify model has guidance_embedder if config says it should has_guidance_embedder = hasattr(model.time_text_embed, "guidance_embedder") if has_guidance_embeds and not has_guidance_embedder: logger.warning( "Config has guidance_embeds=True but model doesn't have guidance_embedder. " "This may indicate a configuration mismatch." ) elif not has_guidance_embeds and has_guidance_embedder: logger.warning( "Config has guidance_embeds=False but model has guidance_embedder. " "This may indicate a configuration mismatch." ) # Note: guidance_in mappings are already included in comfyui_flux_mappings above. # If model doesn't support guidance embeddings, the weights will be filtered out # in _convert_comfyui_weights() based on has_guidance_embeds flag. param_names_mapping_fn = get_param_names_mapping(updated_mapping) weight_iterator = safetensors_weights_iterator(safetensors_list) converted_weights = self._convert_comfyui_weights( weight_iterator=weight_iterator, qkv_size=qkv_size, mlp_hidden_dim=mlp_hidden_dim, has_guidance_embeds=has_guidance_embeds, ) model_state_dict = model.state_dict() missing_keys = set(model_state_dict.keys()) unexpected_keys = [] loaded_count = 0 reverse_param_names_mapping = {} # Handle merged parameters (collect all parts before merging) from collections import defaultdict to_merge_params = defaultdict(dict) # Process weights incrementally: load immediately after conversion for source_name, tensor in converted_weights: target_name, merge_index, num_params_to_merge = param_names_mapping_fn( source_name ) reverse_param_names_mapping[target_name] = ( source_name, merge_index, num_params_to_merge, ) if merge_index is not None: # Collect parts for merging to_merge_params[target_name][merge_index] = tensor if len(to_merge_params[target_name]) == num_params_to_merge: # All parts collected, merge them sorted_tensors = [ to_merge_params[target_name][i] for i in range(num_params_to_merge) ] merged_tensor = torch.cat(sorted_tensors, dim=0) # Load immediately after merging if target_name in model_state_dict: param = model_state_dict[target_name] loaded_tensor = merged_tensor.to( device=param.device, dtype=param.dtype ) param.data.copy_(loaded_tensor) missing_keys.discard(target_name) loaded_count += 1 del merged_tensor, loaded_tensor else: unexpected_keys.append(target_name) # Clear merged parts del to_merge_params[target_name] for t in sorted_tensors: del t else: # Direct mapping, load immediately if target_name in model_state_dict: param = model_state_dict[target_name] # Check shape compatibility if tensor.shape != param.shape: logger.warning( f"Shape mismatch for {target_name}: " f"loaded {tensor.shape} vs model {param.shape}, skipping. " f"Source: {source_name}" ) unexpected_keys.append(target_name) del tensor continue # Debug logging for norm_out.linear to verify mapping if ( "norm_out.linear" in target_name or "final_layer.adaLN_modulation" in source_name ): logger.info( f"Loading norm_out.linear: {source_name} -> {target_name}, " f"shape: {tensor.shape}" ) loaded_tensor = tensor.to(device=param.device, dtype=param.dtype) param.data.copy_(loaded_tensor) missing_keys.discard(target_name) loaded_count += 1 del tensor, loaded_tensor else: # Debug logging for unmapped parameters if "norm_out.linear" in target_name: logger.warning( f"norm_out.linear parameter {target_name} not found in model state_dict. " f"Source: {source_name}" ) unexpected_keys.append(target_name) optional_missing_keys = [] required_missing_keys = [] for key in missing_keys: if key.endswith(".bias"): # Check if corresponding weight exists (if weight exists but bias doesn't, it's optional) weight_key = key.replace(".bias", ".weight") if weight_key not in missing_keys: optional_missing_keys.append(key) else: required_missing_keys.append(key) else: required_missing_keys.append(key) if required_missing_keys: logger.warning( f"Required missing keys (first 10): {required_missing_keys[:10]}..." ) if optional_missing_keys: logger.info( f"Optional missing keys (bias parameters, {len(optional_missing_keys)} total): " f"These will use default values (zeros)" ) if unexpected_keys: logger.warning(f"Unexpected keys (first 10): {unexpected_keys[:10]}...") logger.info(f"Successfully loaded {loaded_count} weight tensors") return model, reverse_param_names_mapping def _convert_comfyui_weights( self, weight_iterator: Generator[tuple[str, torch.Tensor], None, None], qkv_size: int, mlp_hidden_dim: int, has_guidance_embeds: bool, ) -> Generator[tuple[str, torch.Tensor], None, None]: """ Convert ComfyUI Flux weights to SGLang format. Splits fused qkv weights into to_q/to_k/to_v plus proj_mlp. Filters out guidance_in weights if model doesn't support guidance embeddings. Handles scale/shift order difference between ComfyUI and AdaLayerNormContinuous. """ for name, tensor in weight_iterator: if not has_guidance_embeds and name.startswith("guidance_in."): logger.debug( f"Skipping {name} (model doesn't support guidance embeddings)" ) continue # Split fused qkv in double blocks into separate q/k/v projections match = re.match( r"double_blocks\.(\d+)\.(img_attn|txt_attn)\.qkv\.(weight|bias)$", name ) if match: block_idx, attn_type, param_type = match.groups() hidden_size = qkv_size // 3 if tensor.shape[0] < 3 * hidden_size: logger.warning( f"{name} shape {tensor.shape} smaller than expected qkv size {3 * hidden_size}, skipping" ) continue if param_type == "bias": q_tensor = tensor[:hidden_size] k_tensor = tensor[hidden_size : 2 * hidden_size] v_tensor = tensor[2 * hidden_size : 3 * hidden_size] else: q_tensor = tensor[:hidden_size, :] k_tensor = tensor[hidden_size : 2 * hidden_size, :] v_tensor = tensor[2 * hidden_size : 3 * hidden_size, :] target_prefix = f"transformer_blocks.{block_idx}.attn" if attn_type == "img_attn": yield f"{target_prefix}.to_q.{param_type}", q_tensor yield f"{target_prefix}.to_k.{param_type}", k_tensor yield f"{target_prefix}.to_v.{param_type}", v_tensor else: # txt_attn corresponds to encoder projections yield f"{target_prefix}.add_q_proj.{param_type}", q_tensor yield f"{target_prefix}.add_k_proj.{param_type}", k_tensor yield f"{target_prefix}.add_v_proj.{param_type}", v_tensor continue match = re.match(r"single_blocks\.(\d+)\.linear1\.(weight|bias)$", name) if match: block_idx, param_type = match.groups() expected_size = qkv_size + mlp_hidden_dim if tensor.shape[0] < expected_size: logger.warning( f"linear1.{param_type} shape {tensor.shape} doesn't match " f"expected size {expected_size}, skipping" ) continue # Split tensor qkv_tensor = ( tensor[:qkv_size] if param_type == "bias" else tensor[:qkv_size, :] ) mlp_tensor = ( tensor[qkv_size:] if param_type == "bias" else tensor[qkv_size:, :] ) # Split qkv into q/k/v for single blocks hidden_size = qkv_size // 3 if param_type == "bias": q_tensor = qkv_tensor[:hidden_size] k_tensor = qkv_tensor[hidden_size : 2 * hidden_size] v_tensor = qkv_tensor[2 * hidden_size : 3 * hidden_size] else: q_tensor = qkv_tensor[:hidden_size, :] k_tensor = qkv_tensor[hidden_size : 2 * hidden_size, :] v_tensor = qkv_tensor[2 * hidden_size : 3 * hidden_size, :] yield f"single_transformer_blocks.{block_idx}.attn.to_q.{param_type}", q_tensor yield f"single_transformer_blocks.{block_idx}.attn.to_k.{param_type}", k_tensor yield f"single_transformer_blocks.{block_idx}.attn.to_v.{param_type}", v_tensor yield f"single_transformer_blocks.{block_idx}.proj_mlp.{param_type}", mlp_tensor elif name == "final_layer.adaLN_modulation.1.weight": # ComfyUI: output order is [shift, scale] # AdaLayerNormContinuous: expects [scale, shift] # Need to swap the first half and second half of the weight matrix # Weight shape: (2 * hidden_size, hidden_size) # Split into two halves and swap them half_size = tensor.shape[0] // 2 shift_weights = tensor[:half_size, :] scale_weights = tensor[half_size:, :] # Swap: put scale first, then shift swapped_tensor = torch.cat([scale_weights, shift_weights], dim=0) logger.info( f"Swapped scale/shift order for {name}: " f"shape {tensor.shape} -> {swapped_tensor.shape}" ) yield name, swapped_tensor elif name == "final_layer.adaLN_modulation.1.bias": # Same swap for bias: (2 * hidden_size,) half_size = tensor.shape[0] // 2 shift_bias = tensor[:half_size] scale_bias = tensor[half_size:] swapped_tensor = torch.cat([scale_bias, shift_bias], dim=0) logger.info( f"Swapped scale/shift order for {name}: " f"shape {tensor.shape} -> {swapped_tensor.shape}" ) yield name, swapped_tensor else: # Other weights pass through (handled by param_names_mapping) yield name, tensor def _load_transformer_from_safetensors( self, server_args: ServerArgs, loaded_modules: dict[str, torch.nn.Module] | None = None, ) -> dict[str, Any]: """ Load transformer directly from safetensors file without model_index.json. """ if loaded_modules is not None and "transformer" in loaded_modules: logger.info("Using provided transformer module") components = { "transformer": loaded_modules["transformer"], "scheduler": self.modules.get("scheduler"), } return components if hasattr(server_args.pipeline_config, "dit_config"): dit_config = server_args.pipeline_config.dit_config if not isinstance(dit_config, FluxConfig): logger.warning("dit_config is not FluxConfig, creating new FluxConfig") dit_config = FluxConfig() server_args.pipeline_config.dit_config = dit_config else: logger.info("Creating default FluxConfig") dit_config = FluxConfig() server_args.pipeline_config.dit_config = dit_config # Set guidance_embeds to True for ComfyUI Flux models dit_config.arch_config.guidance_embeds = True logger.info("Set guidance_embeds=True for ComfyUI Flux model") if dit_config.arch_config.param_names_mapping is None: dit_config.arch_config.param_names_mapping = {} # ComfyUI Flux uses different parameter names than SGLang Flux # Key differences: # - ComfyUI: single_blocks.{i}.linear1 (fused QKV + MLP input) # - SGLang: single_transformer_blocks.{i}.attn.to_qkv + proj_mlp (separate) # - ComfyUI: single_blocks.{i}.linear2 # - SGLang: single_transformer_blocks.{i}.proj_out # - ComfyUI: double_blocks.{i}.img_attn.qkv / txt_attn.qkv # - SGLang: transformer_blocks.{i}.attn.to_qkv / attn.to_added_qkv # Note: For fused layers like linear1, we need custom weight splitting logic # which will be handled in the weight conversion function below comfyui_flux_mappings = { # Double stream blocks - attention layers r"double_blocks\.(\d+)\.img_attn\.qkv\.(weight|bias)$": ( r"transformer_blocks.\1.attn.to_qkv.\2", None, None, ), r"double_blocks\.(\d+)\.txt_attn\.qkv\.(weight|bias)$": ( r"transformer_blocks.\1.attn.to_added_qkv.\2", None, None, ), r"double_blocks\.(\d+)\.img_attn\.proj\.(weight|bias)$": ( r"transformer_blocks.\1.attn.to_out.0.\2", None, None, ), r"double_blocks\.(\d+)\.txt_attn\.proj\.(weight|bias)$": ( r"transformer_blocks.\1.attn.to_add_out.\2", None, None, ), r"double_blocks\.(\d+)\.img_attn\.norm\.query_norm\.scale$": ( r"transformer_blocks.\1.attn.norm_q.weight", None, None, ), r"double_blocks\.(\d+)\.img_attn\.norm\.key_norm\.scale$": ( r"transformer_blocks.\1.attn.norm_k.weight", None, None, ), r"double_blocks\.(\d+)\.txt_attn\.norm\.query_norm\.scale$": ( r"transformer_blocks.\1.attn.norm_added_q.weight", None, None, ), r"double_blocks\.(\d+)\.txt_attn\.norm\.key_norm\.scale$": ( r"transformer_blocks.\1.attn.norm_added_k.weight", None, None, ), # Double stream blocks - MLP layers (map to net structure) r"double_blocks\.(\d+)\.img_mlp\.0\.(weight|bias)$": ( r"transformer_blocks.\1.ff.net.0.proj.\2", None, None, ), r"double_blocks\.(\d+)\.img_mlp\.2\.(weight|bias)$": ( r"transformer_blocks.\1.ff.net.2.\2", None, None, ), r"double_blocks\.(\d+)\.txt_mlp\.0\.(weight|bias)$": ( r"transformer_blocks.\1.ff_context.net.0.proj.\2", None, None, ), r"double_blocks\.(\d+)\.txt_mlp\.2\.(weight|bias)$": ( r"transformer_blocks.\1.ff_context.net.2.\2", None, None, ), # Double stream blocks - modulation layers r"double_blocks\.(\d+)\.img_mod\.lin\.(weight|bias)$": ( r"transformer_blocks.\1.norm1.linear.\2", None, None, ), r"double_blocks\.(\d+)\.txt_mod\.lin\.(weight|bias)$": ( r"transformer_blocks.\1.norm1_context.linear.\2", None, None, ), # Single stream blocks - linear2 maps to proj_out r"single_blocks\.(\d+)\.linear2\.(weight|bias)$": ( r"single_transformer_blocks.\1.proj_out.\2", None, None, ), # Single stream blocks - norm layers (scale -> weight) r"single_blocks\.(\d+)\.norm\.query_norm\.scale$": ( r"single_transformer_blocks.\1.attn.norm_q.weight", None, None, ), r"single_blocks\.(\d+)\.norm\.key_norm\.scale$": ( r"single_transformer_blocks.\1.attn.norm_k.weight", None, None, ), # Single stream blocks - modulation (maps to norm.linear) r"single_blocks\.(\d+)\.modulation\.lin\.(weight|bias)$": ( r"single_transformer_blocks.\1.norm.linear.\2", None, None, ), # Time and guidance embeddings r"^time_in\.in_layer\.(weight|bias)$": ( r"time_text_embed.timestep_embedder.linear_1.\1", None, None, ), r"^time_in\.out_layer\.(weight|bias)$": ( r"time_text_embed.timestep_embedder.linear_2.\1", None, None, ), r"^txt_in\.(weight|bias)$": (r"context_embedder.\1", None, None), r"^vector_in\.in_layer\.(weight|bias)$": ( r"time_text_embed.text_embedder.linear_1.\1", None, None, ), r"^vector_in\.out_layer\.(weight|bias)$": ( r"time_text_embed.text_embedder.linear_2.\1", None, None, ), # Final layer mappings r"^final_layer\.linear\.(weight|bias)$": (r"proj_out.\1", None, None), r"^final_layer\.norm_final\.(weight|bias)$": (r"norm_out.\1", None, None), r"^final_layer\.adaLN_modulation\.1\.(weight|bias)$": ( r"norm_out.linear.\1", None, None, ), # Image input embedding r"^img_in\.(weight|bias)$": (r"x_embedder.\1", None, None), # Guidance embeddings (if model supports guidance) r"^guidance_in\.in_layer\.(weight|bias)$": ( r"time_text_embed.guidance_embedder.linear_1.\1", None, None, ), r"^guidance_in\.out_layer\.(weight|bias)$": ( r"time_text_embed.guidance_embedder.linear_2.\1", None, None, ), } # Merge ComfyUI mappings with existing mappings (ComfyUI mappings take precedence) updated_mapping = { **dit_config.arch_config.param_names_mapping, **comfyui_flux_mappings, } dit_config.arch_config.param_names_mapping = updated_mapping logger.info( "Added ComfyUI weight name mappings for Flux model. " f"Total mappings: {len(updated_mapping)}" ) cls_name = "FluxTransformer2DModel" model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) logger.info("Resolved transformer class: %s", cls_name) original_mapping = None if comfyui_flux_mappings: original_mapping = model_cls.param_names_mapping model_cls.param_names_mapping = updated_mapping logger.info( "Temporarily updated model class param_names_mapping with ComfyUI mappings. " f"Total mappings: {len(updated_mapping)}" ) safetensors_list = [self.model_path] logger.info("Loading weights from: %s", safetensors_list) default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] server_args.model_paths["transformer"] = os.path.dirname(self.model_path) or "." hf_config = {} hidden_size = ( dit_config.arch_config.num_attention_heads * dit_config.arch_config.attention_head_dim ) mlp_ratio = getattr(dit_config.arch_config, "mlp_ratio", 4.0) mlp_hidden_dim = int(hidden_size * mlp_ratio) qkv_size = 3 * hidden_size has_guidance_embeds = True # Load and convert weights from safetensors file model, reverse_param_names_mapping = ( self._load_and_convert_weights_from_safetensors( model_cls=model_cls, dit_config=dit_config, hf_config=hf_config, safetensors_list=safetensors_list, updated_mapping=updated_mapping, qkv_size=qkv_size, mlp_hidden_dim=mlp_hidden_dim, has_guidance_embeds=has_guidance_embeds, default_dtype=default_dtype, ) ) model = model.eval() for param in model.parameters(): param.requires_grad = False model.reverse_param_names_mapping = reverse_param_names_mapping if original_mapping is not None: model_cls.param_names_mapping = original_mapping total_params = sum(p.numel() for p in model.parameters()) logger.info("Loaded transformer with %.2fB parameters", total_params / 1e9) components = { "transformer": model, "scheduler": self.modules.get("scheduler"), } logger.info("Successfully loaded modules: %s", list(components.keys())) return components def create_pipeline_stages(self, server_args: ServerArgs): logger.info( "ComfyUIFluxPipeline.create_pipeline_stages() called - creating latent_preparation_stage and denoising_stage" ) self.add_stages( [ ComfyUILatentPreparationStage( scheduler=self.get_module("scheduler"), transformer=self.get_module("transformer"), ), DenoisingStage( transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), ), ] ) logger.info( f"ComfyUIFluxPipeline stages created: {list(self._stage_name_mapping.keys())}" ) EntryClass = ComfyUIFluxPipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/comfyui_qwen_image_pipeline.py ================================================ # SPDX-License-Identifier: Apache-2.0 import os from itertools import chain from typing import Any import torch from torch.distributed import init_device_mesh from torch.distributed.fsdp import MixedPrecisionPolicy from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.fsdp_load import ( load_model_from_full_model_state_dict, shard_model, ) from sglang.multimodal_gen.runtime.loader.utils import ( get_param_names_mapping, set_default_torch_dtype, ) from sglang.multimodal_gen.runtime.loader.weight_utils import ( safetensors_weights_iterator, ) from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.models.schedulers.scheduling_comfyui_passthrough import ( ComfyUIPassThroughScheduler, ) from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.stages import ( ComfyUILatentPreparationStage, DenoisingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE, set_mixed_precision_policy logger = init_logger(__name__) class ComfyUIQwenImagePipelineBase(LoRAPipeline, ComposedPipelineBase): """ Base pipeline for ComfyUI QwenImage integration with only denoising stage. This pipeline requires pre-processed inputs: - prompt_embeds: Pre-encoded text embeddings (list of tensors) - latents: Pre-processed image latents in sequence format [B, S, D] Usage: generator = DiffGenerator.from_pretrained( model_path="path/to/model", pipeline_class_name="ComfyUIQwenImagePipeline", device="cuda", ) """ # Subclasses should override this zero_cond_t: bool = False pipeline_name = "ComfyUIQwenImagePipeline" _required_config_modules = [ "transformer", "scheduler", ] def initialize_pipeline(self, server_args: ServerArgs): """ Initialize the pipeline with ComfyUI pass-through scheduler. This scheduler does not modify latents, allowing ComfyUI to handle denoising. """ self.modules["scheduler"] = ComfyUIPassThroughScheduler( num_train_timesteps=1000 ) # Ensure VAE config is properly initialized even though we don't load the VAE model vae_config = server_args.pipeline_config.vae_config vae_config.post_init() logger.info( "Called vae_config.post_init() to set vae_scale_factor. " f"vae_scale_factor={vae_config.arch_config.vae_scale_factor}" ) def load_modules( self, server_args: ServerArgs, loaded_modules: dict[str, torch.nn.Module] | None = None, ) -> dict[str, Any]: """ Load modules for ComfyUIQwenImagePipeline. If model_path is a safetensors file, load transformer directly from it without requiring model_index.json. Otherwise, fall back to default loading. """ if os.path.isfile(self.model_path) and self.model_path.endswith(".safetensors"): logger.info( "Detected safetensors file, loading transformer directly from: %s", self.model_path, ) return self._load_transformer_from_safetensors(server_args, loaded_modules) else: logger.info( "Model path is a directory, using default loading method: %s", self.model_path, ) return super().load_modules(server_args, loaded_modules) def _load_transformer_from_safetensors( self, server_args: ServerArgs, loaded_modules: dict[str, torch.nn.Module] | None = None, ) -> dict[str, Any]: """Load transformer directly from safetensors without model_index.json.""" # 1) Fast path: use provided module if loaded_modules is not None and "transformer" in loaded_modules: logger.info("Using provided transformer module") return { "transformer": loaded_modules["transformer"], "scheduler": self.modules.get("scheduler"), } # 2) Build config and mappings dit_config, updated_mapping, model_cls, default_dtype = ( self._prepare_dit_config_and_mapping(server_args) ) safetensors_list = [self.model_path] logger.info("Loading weights from: %s", safetensors_list) # 3) Instantiate model (meta) and optionally shard model = self._instantiate_model( model_cls, dit_config, default_dtype, updated_mapping, server_args ) # 4) Load weights self._load_weights_into_model( model, safetensors_list, default_dtype, updated_mapping, server_args ) components = { "transformer": model, "scheduler": self.modules.get("scheduler"), } logger.info("Successfully loaded modules: %s", list(components.keys())) return components def _prepare_dit_config_and_mapping(self, server_args: ServerArgs): from sglang.multimodal_gen.configs.models.dits.qwenimage import ( QwenImageArchConfig, ) comfyui_arch_config = QwenImageArchConfig( patch_size=2, in_channels=64, out_channels=16, num_layers=60, attention_head_dim=128, num_attention_heads=24, joint_attention_dim=3584, pooled_projection_dim=768, guidance_embeds=False, axes_dims_rope=(16, 56, 56), zero_cond_t=self.zero_cond_t, ) dit_config = QwenImageDitConfig(arch_config=comfyui_arch_config) server_args.pipeline_config.dit_config = dit_config if dit_config.arch_config.param_names_mapping is None: dit_config.arch_config.param_names_mapping = {} comfyui_qwen_mappings = {r"^model\.diffusion_model\.(.*)$": r"\1"} updated_mapping = { **dit_config.arch_config.param_names_mapping, **comfyui_qwen_mappings, } dit_config.arch_config.param_names_mapping = updated_mapping logger.info( "Added ComfyUI weight name mappings to param_names_mapping. " f"Total mappings: {len(updated_mapping)}" ) cls_name = "QwenImageTransformer2DModel" model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) logger.info("Resolved transformer class: %s", cls_name) default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] server_args.model_paths["transformer"] = os.path.dirname(self.model_path) or "." assert server_args.hsdp_shard_dim is not None, "hsdp_shard_dim must be set" logger.info( "Loading %s from safetensors file, default_dtype: %s", cls_name, default_dtype, ) return dit_config, updated_mapping, model_cls, default_dtype def _instantiate_model( self, model_cls, dit_config, default_dtype, updated_mapping, server_args: ServerArgs, ): from sglang.multimodal_gen.runtime.platforms import current_platform hf_config = {} original_mapping = model_cls.param_names_mapping model_cls.param_names_mapping = updated_mapping logger.info( "Temporarily updated model class param_names_mapping with ComfyUI mappings. " f"Total mappings: {len(updated_mapping)}" ) try: mp_policy = MixedPrecisionPolicy( torch.bfloat16, torch.float32, None, cast_forward_inputs=False ) set_mixed_precision_policy( param_dtype=torch.bfloat16, reduce_dtype=torch.float32, output_dtype=None, mp_policy=mp_policy, ) with set_default_torch_dtype(default_dtype), torch.device("meta"): model = model_cls(**{"config": dit_config, "hf_config": hf_config}) use_fsdp = server_args.use_fsdp_inference if current_platform.is_mps(): use_fsdp = False logger.info("Disabling FSDP for MPS platform as it's not compatible") if use_fsdp: device_mesh = init_device_mesh( current_platform.device_type, mesh_shape=( server_args.hsdp_replicate_dim, server_args.hsdp_shard_dim, ), mesh_dim_names=("replicate", "shard"), ) shard_model( model, cpu_offload=server_args.dit_cpu_offload, reshard_after_forward=True, mp_policy=mp_policy, mesh=device_mesh, fsdp_shard_conditions=model._fsdp_shard_conditions, pin_cpu_memory=server_args.pin_cpu_memory, ) finally: model_cls.param_names_mapping = original_mapping return model def _load_weights_into_model( self, model, safetensors_list, default_dtype, updated_mapping, server_args: ServerArgs, ): # Create weight iterator for loading weight_iterator = safetensors_weights_iterator(safetensors_list) # Load weights param_names_mapping_fn = get_param_names_mapping(updated_mapping) load_model_from_full_model_state_dict( model, weight_iterator, get_local_torch_device(), default_dtype, strict=True, cpu_offload=server_args.dit_cpu_offload, param_names_mapping=param_names_mapping_fn, ) # Check for meta parameters for n, p in chain(model.named_parameters(), model.named_buffers()): if p.is_meta: raise RuntimeError(f"Unexpected param or buffer {n} on meta device.") if isinstance(p, torch.nn.Parameter): p.requires_grad = False total_params = sum(p.numel() for p in model.parameters()) logger.info("Loaded transformer with %.2fB parameters", total_params / 1e9) def create_pipeline_stages(self, server_args: ServerArgs): logger.info( f"{self.__class__.__name__}.create_pipeline_stages() called - creating latent_preparation_stage and denoising_stage" ) self.add_stages( [ ComfyUILatentPreparationStage( scheduler=self.get_module("scheduler"), transformer=self.get_module("transformer"), ), DenoisingStage( transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), ), ] ) logger.info( f"{self.__class__.__name__} stages created: {list(self._stage_name_mapping.keys())}" ) class ComfyUIQwenImagePipeline(ComfyUIQwenImagePipelineBase): """ComfyUI QwenImage pipeline for text-to-image generation.""" pipeline_name = "ComfyUIQwenImagePipeline" zero_cond_t = False from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImagePipelineConfig, ) from sglang.multimodal_gen.configs.sample.qwenimage import QwenImageSamplingParams pipeline_config_cls = QwenImagePipelineConfig sampling_params_cls = QwenImageSamplingParams class ComfyUIQwenImageEditPipeline(ComfyUIQwenImagePipelineBase): """ComfyUI QwenImage pipeline for image-to-image editing.""" pipeline_name = "ComfyUIQwenImageEditPipeline" zero_cond_t = True from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImageEditPlusPipelineConfig, ) from sglang.multimodal_gen.configs.sample.qwenimage import ( QwenImageEditPlusSamplingParams, ) pipeline_config_cls = QwenImageEditPlusPipelineConfig sampling_params_cls = QwenImageEditPlusSamplingParams EntryClass = [ComfyUIQwenImagePipeline, ComfyUIQwenImageEditPipeline] ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/comfyui_zimage_pipeline.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import os import re from collections.abc import Generator from itertools import chain from typing import Any import torch from torch.distributed import init_device_mesh from torch.distributed.fsdp import MixedPrecisionPolicy from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.fsdp_load import ( load_model_from_full_model_state_dict, shard_model, ) from sglang.multimodal_gen.runtime.loader.utils import ( get_param_names_mapping, set_default_torch_dtype, ) from sglang.multimodal_gen.runtime.loader.weight_utils import ( safetensors_weights_iterator, ) from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.models.schedulers.scheduling_comfyui_passthrough import ( ComfyUIPassThroughScheduler, ) from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.stages import ( ComfyUILatentPreparationStage, DenoisingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE, set_mixed_precision_policy logger = init_logger(__name__) class ComfyUIZImagePipeline(LoRAPipeline, ComposedPipelineBase): """ Simplified pipeline for ComfyUI integration with only denoising stage. This pipeline requires pre-processed inputs: - prompt_embeds: Pre-encoded text embeddings (list of tensors) - negative_prompt_embeds: Pre-encoded negative prompt embeddings (if using CFG) - latents: Optional initial noise latents (will be generated if not provided) Usage: generator = DiffGenerator.from_pretrained( model_path="path/to/model", pipeline_class_name="ComfyUIZImagePipeline", device="cuda", ) """ pipeline_name = "ComfyUIZImagePipeline" from sglang.multimodal_gen.configs.pipeline_configs.zimage import ( ZImagePipelineConfig, ) from sglang.multimodal_gen.configs.sample.zimage import ZImageSamplingParams pipeline_config_cls = ZImagePipelineConfig sampling_params_cls = ZImageSamplingParams _required_config_modules = [ "transformer", "scheduler", ] def initialize_pipeline(self, server_args: ServerArgs): """ Initialize the pipeline with ComfyUI pass-through scheduler. This scheduler does not modify latents, allowing ComfyUI to handle denoising. """ self.modules["scheduler"] = ComfyUIPassThroughScheduler( num_train_timesteps=1000 ) # Ensure VAE config is properly initialized even though we don't load the VAE model # This is necessary because get_freqs_cis uses spatial_compression_ratio if hasattr(server_args.pipeline_config, "vae_config"): vae_config = server_args.pipeline_config.vae_config if hasattr(vae_config, "post_init") and not hasattr( vae_config, "_post_init_called" ): vae_config.post_init() logger.info( "Called vae_config.post_init() to set spatial_compression_ratio. " f"spatial_compression_ratio={vae_config.arch_config.spatial_compression_ratio}" ) def load_modules( self, server_args: ServerArgs, loaded_modules: dict[str, torch.nn.Module] | None = None, ) -> dict[str, Any]: """ Load modules for ComfyUIZImagePipeline. If model_path is a safetensors file, load transformer directly from it without requiring model_index.json. Otherwise, fall back to default loading. """ if os.path.isfile(self.model_path) and self.model_path.endswith(".safetensors"): logger.info( "Detected safetensors file, loading transformer directly from: %s", self.model_path, ) return self._load_transformer_from_safetensors(server_args, loaded_modules) else: logger.info( "Model path is a directory, using default loading method: %s", self.model_path, ) return super().load_modules(server_args, loaded_modules) def _convert_comfyui_qkv_weights( self, weight_iterator: Generator[tuple[str, torch.Tensor], None, None], dim: int, num_heads: int, num_kv_heads: int, ) -> Generator[tuple[str, torch.Tensor], None, None]: """ Convert ComfyUI zimage qkv weights to SGLang format. Splits merged qkv.weight into separate to_q, to_k, to_v weights. Args: weight_iterator: Iterator yielding (name, tensor) pairs from safetensors dim: Model dimension num_heads: Number of attention heads num_kv_heads: Number of key-value heads Yields: (name, tensor) pairs with qkv weights split into to_q, to_k, to_v """ head_dim = dim // num_heads q_size = dim k_size = head_dim * num_kv_heads v_size = head_dim * num_kv_heads for name, tensor in weight_iterator: # Match qkv weights in layers, noise_refiner, or context_refiner # Pattern: (layers|noise_refiner|context_refiner).{i}.attention.qkv.(weight|bias) match = re.match( r"(layers|noise_refiner|context_refiner)\.(\d+)\.attention\.qkv\.(weight|bias)$", name, ) if match: module_name, layer_idx, param_type = match.groups() base_name = f"{module_name}.{layer_idx}.attention" if param_type == "weight": # Weight shape: (q_size + k_size + v_size, dim) # Split into q, k, v q_weight = tensor[:q_size, :] k_weight = tensor[q_size : q_size + k_size, :] v_weight = tensor[q_size + k_size :, :] logger.debug( f"Splitting {name} (shape {tensor.shape}) into " f"to_q ({q_weight.shape}), to_k ({k_weight.shape}), to_v ({v_weight.shape})" ) yield f"{base_name}.to_q.weight", q_weight yield f"{base_name}.to_k.weight", k_weight yield f"{base_name}.to_v.weight", v_weight else: # bias # Bias shape: (q_size + k_size + v_size,) # Split into q, k, v q_bias = tensor[:q_size] k_bias = tensor[q_size : q_size + k_size] v_bias = tensor[q_size + k_size :] logger.debug( f"Splitting {name} (shape {tensor.shape}) into " f"to_q ({q_bias.shape}), to_k ({k_bias.shape}), to_v ({v_bias.shape})" ) yield f"{base_name}.to_q.bias", q_bias yield f"{base_name}.to_k.bias", k_bias yield f"{base_name}.to_v.bias", v_bias else: # Pass through other weights unchanged yield name, tensor def _load_transformer_from_safetensors( self, server_args: ServerArgs, loaded_modules: dict[str, torch.nn.Module] | None = None, ) -> dict[str, Any]: """ Load transformer directly from safetensors file without model_index.json. This method: 1. Uses hardcoded ZImageDitConfig for zimage model 2. Loads transformer from the safetensors file 3. Uses ComfyUIPassThroughScheduler (already created in initialize_pipeline) """ # Check if transformer is already provided if loaded_modules is not None and "transformer" in loaded_modules: logger.info("Using provided transformer module") components = { "transformer": loaded_modules["transformer"], "scheduler": self.modules.get("scheduler"), } return components if hasattr(server_args.pipeline_config, "dit_config"): dit_config = server_args.pipeline_config.dit_config if not isinstance(dit_config, ZImageDitConfig): logger.warning( "dit_config is not ZImageDitConfig, creating new ZImageDitConfig" ) dit_config = ZImageDitConfig() server_args.pipeline_config.dit_config = dit_config else: logger.info("Creating default ZImageDitConfig") dit_config = ZImageDitConfig() server_args.pipeline_config.dit_config = dit_config if dit_config.arch_config.param_names_mapping is None: dit_config.arch_config.param_names_mapping = {} # Add mappings for norm layers: map from ComfyUI format (k_norm/q_norm) to SGLang format (norm_k/norm_q) # The regex matches the source name from safetensors, and the tuple specifies the target name in the model # Note: qkv weights are handled separately by _convert_comfyui_qkv_weights function comfyui_norm_mappings = { r"(.*)\.attention\.k_norm\.weight$": ( r"\1.attention.norm_k.weight", None, None, ), r"(.*)\.attention\.q_norm\.weight$": ( r"\1.attention.norm_q.weight", None, None, ), r"(.*)\.attention\.out\.weight$": ( r"\1.attention.to_out.0.weight", None, None, ), r"^final_layer\.(.*)$": (r"all_final_layer.2-1.\1", None, None), r"^x_embedder\.(.*)$": (r"all_x_embedder.2-1.\1", None, None), } # Merge ComfyUI mappings with existing mappings (ComfyUI mappings take precedence) updated_mapping = { **dit_config.arch_config.param_names_mapping, **comfyui_norm_mappings, } dit_config.arch_config.param_names_mapping = updated_mapping logger.info( "Added ComfyUI weight name mappings (k_norm/q_norm -> norm_k/norm_q) to param_names_mapping. " f"Total mappings: {len(updated_mapping)}" ) cls_name = "ZImageTransformer2DModel" model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) logger.info("Resolved transformer class: %s", cls_name) safetensors_list = [self.model_path] logger.info("Loading weights from: %s", safetensors_list) default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] server_args.model_paths["transformer"] = os.path.dirname(self.model_path) or "." hf_config = {} assert server_args.hsdp_shard_dim is not None, "hsdp_shard_dim must be set" logger.info( "Loading %s from safetensors file, default_dtype: %s", cls_name, default_dtype, ) original_mapping = model_cls.param_names_mapping model_cls.param_names_mapping = updated_mapping logger.info( "Temporarily updated model class param_names_mapping with ComfyUI mappings. " f"Total mappings: {len(updated_mapping)}" ) try: # Create model first (same as maybe_load_fsdp_model) from sglang.multimodal_gen.runtime.platforms import current_platform mp_policy = MixedPrecisionPolicy( torch.bfloat16, torch.float32, None, cast_forward_inputs=False ) set_mixed_precision_policy( param_dtype=torch.bfloat16, reduce_dtype=torch.float32, output_dtype=None, mp_policy=mp_policy, ) with set_default_torch_dtype(default_dtype), torch.device("meta"): model = model_cls(**{"config": dit_config, "hf_config": hf_config}) # Check if we should use FSDP use_fsdp = server_args.use_fsdp_inference if current_platform.is_mps(): use_fsdp = False logger.info("Disabling FSDP for MPS platform as it's not compatible") if use_fsdp: world_size = server_args.hsdp_replicate_dim * server_args.hsdp_shard_dim device_mesh = init_device_mesh( current_platform.device_type, mesh_shape=( server_args.hsdp_replicate_dim, server_args.hsdp_shard_dim, ), mesh_dim_names=("replicate", "shard"), ) shard_model( model, cpu_offload=server_args.dit_cpu_offload, reshard_after_forward=True, mp_policy=mp_policy, mesh=device_mesh, fsdp_shard_conditions=model._fsdp_shard_conditions, pin_cpu_memory=server_args.pin_cpu_memory, ) # Get model dimensions for qkv splitting arch_config = dit_config.arch_config dim = arch_config.dim num_heads = arch_config.num_attention_heads num_kv_heads = arch_config.n_kv_heads # Create weight iterator with qkv conversion base_weight_iterator = safetensors_weights_iterator(safetensors_list) converted_weight_iterator = self._convert_comfyui_qkv_weights( base_weight_iterator, dim, num_heads, num_kv_heads ) # Load weights param_names_mapping_fn = get_param_names_mapping(updated_mapping) load_model_from_full_model_state_dict( model, converted_weight_iterator, get_local_torch_device(), default_dtype, strict=True, cpu_offload=server_args.dit_cpu_offload, param_names_mapping=param_names_mapping_fn, ) # Check for meta parameters for n, p in chain(model.named_parameters(), model.named_buffers()): if p.is_meta: raise RuntimeError( f"Unexpected param or buffer {n} on meta device." ) if isinstance(p, torch.nn.Parameter): p.requires_grad = False finally: model_cls.param_names_mapping = original_mapping total_params = sum(p.numel() for p in model.parameters()) logger.info("Loaded transformer with %.2fB parameters", total_params / 1e9) components = { "transformer": model, "scheduler": self.modules.get("scheduler"), } logger.info("Successfully loaded modules: %s", list(components.keys())) return components def create_pipeline_stages(self, server_args: ServerArgs): logger.info( "ComfyUIZImagePipeline.create_pipeline_stages() called - creating latent_preparation_stage and denoising_stage" ) self.add_stages( [ ComfyUILatentPreparationStage( scheduler=self.get_module("scheduler"), transformer=self.get_module("transformer"), ), DenoisingStage( transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), ), ] ) logger.info( f"ComfyUIZImagePipeline stages created: {list(self._stage_name_mapping.keys())}" ) EntryClass = ComfyUIZImagePipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Diffusers backend pipeline wrapper. This module provides a wrapper that allows running any diffusers-supported model through sglang's infrastructure using vanilla diffusers pipelines. """ import argparse import inspect import re import warnings from io import BytesIO from typing import Any import numpy as np import requests import torch import torchvision.transforms as T from diffusers import DiffusionPipeline from PIL import Image from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import ( PipelineExecutor, ) from sglang.multimodal_gen.runtime.pipelines_core.executors.sync_executor import ( SyncExecutor, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages import PipelineStage from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class DiffusersExecutionStage(PipelineStage): """Pipeline stage that wraps diffusers pipeline execution.""" def __init__(self, diffusers_pipe: DiffusionPipeline): super().__init__() self.diffusers_pipe = diffusers_pipe def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Execute the diffusers pipeline.""" kwargs = self._build_pipeline_kwargs(batch) # Filter kwargs to only those supported by the pipeline, warn about ignored args kwargs, _ = self._filter_pipeline_kwargs(kwargs) # Request tensor output for cleaner handling if "output_type" not in kwargs: kwargs["output_type"] = "pt" with torch.no_grad(), warnings.catch_warnings(record=True): warnings.simplefilter("always") try: output = self.diffusers_pipe(**kwargs) except TypeError as e: # Some pipelines don't support output_type="pt" if "output_type" in str(e): kwargs.pop("output_type", None) output = self.diffusers_pipe(**kwargs) else: raise batch.output = self._extract_output(output) if batch.output is not None: batch.output = self._postprocess_output(batch.output) return batch def _filter_pipeline_kwargs( self, kwargs: dict[str, Any], *, strict: bool = False ) -> tuple[dict[str, Any], list[str]]: """Filter kwargs to those accepted by the pipeline's __call__. Args: kwargs: Arguments to filter strict: If True, raise ValueError on unsupported args; otherwise warn Returns: Tuple of (filtered_kwargs, ignored_keys) """ try: sig = inspect.signature(self.diffusers_pipe.__call__) except (ValueError, TypeError): return kwargs, [] params = sig.parameters accepts_var_kwargs = any( p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values() ) if accepts_var_kwargs: return kwargs, [] valid = set(params.keys()) - {"self"} filtered = {} ignored = [] for k, v in kwargs.items(): if k in valid: filtered[k] = v else: ignored.append(k) if ignored: pipe_name = type(self.diffusers_pipe).__name__ msg = ( f"Pipeline '{pipe_name}' does not support: {', '.join(sorted(ignored))}. " "These arguments will be ignored." ) if strict: raise ValueError(msg) logger.warning(msg) return filtered, ignored def _extract_output(self, output: Any) -> torch.Tensor | None: """Extract tensor output from pipeline result.""" for attr in ["images", "frames", "video", "sample", "pred_original_sample"]: data = getattr(output, attr, None) if data is None: continue result = self._convert_to_tensor(data) if result is not None: logger.debug( "Extracted output from '%s': shape=%s, dtype=%s", attr, result.shape, result.dtype, ) return result logger.warning("Could not extract output from pipeline result") return None def _convert_to_tensor(self, data: Any) -> torch.Tensor | None: """Convert various data formats to a tensor.""" if isinstance(data, torch.Tensor): return data if isinstance(data, np.ndarray): tensor = torch.from_numpy(data).float() if tensor.max() > 1.0: tensor = tensor / 255.0 # (B, H, W, C) -> (B, C, H, W) or (B, T, H, W, C) -> (B, C, T, H, W) if tensor.ndim == 4: tensor = tensor.permute(0, 3, 1, 2) elif tensor.ndim == 5: tensor = tensor.permute(0, 4, 1, 2, 3) return tensor if isinstance(data, Image.Image): return T.ToTensor()(data) if isinstance(data, list) and len(data) > 0: return self._convert_list_to_tensor(data) return None def _convert_list_to_tensor(self, data: list) -> torch.Tensor | None: """Convert a list of items to a tensor.""" first = data[0] # Nested list (e.g., [[frame1, frame2, ...]] for video batches) if isinstance(first, list) and len(first) > 0: data = first first = data[0] if isinstance(first, Image.Image): tensors = [T.ToTensor()(img) for img in data] stacked = torch.stack(tensors) if len(tensors) > 1: return stacked.permute(1, 0, 2, 3) # (T, C, H, W) -> (C, T, H, W) return stacked[0] if isinstance(first, torch.Tensor): stacked = torch.stack(data) if len(data) > 1: return stacked.permute(1, 0, 2, 3) return stacked[0] if isinstance(first, np.ndarray): tensors = [torch.from_numpy(arr).float() for arr in data] if tensors[0].max() > 1.0: tensors = [t / 255.0 for t in tensors] if tensors[0].ndim == 3: tensors = [t.permute(2, 0, 1) for t in tensors] stacked = torch.stack(tensors) if len(data) > 1: return stacked.permute(1, 0, 2, 3) return stacked[0] return None def _postprocess_output(self, output: torch.Tensor) -> torch.Tensor: """Post-process output tensor to ensure valid values and correct shape.""" output = output.cpu().float() # Handle NaN or Inf values if torch.isnan(output).any() or torch.isinf(output).any(): logger.warning("Output contains invalid values, fixing...") output = torch.nan_to_num(output, nan=0.5, posinf=1.0, neginf=0.0) # Normalize to [0, 1] range if needed min_val, max_val = output.min().item(), output.max().item() if min_val < -0.5 or max_val > 1.5: output = (output + 1) / 2 output = output.clamp(0, 1) # Ensure correct shape for downstream processing output = self._fix_output_shape(output) logger.debug("Final output tensor shape: %s", output.shape) return output def _fix_output_shape(self, output: torch.Tensor) -> torch.Tensor: """Fix tensor shape for downstream processing. Expected: (B, C, H, W) for images or (B, C, T, H, W) for videos. """ if output.dim() == 5: # Video: (B, T, C, H, W) -> (B, C, T, H, W) return output.permute(0, 2, 1, 3, 4) if output.dim() == 4: if output.shape[0] == 1 or output.shape[1] in [1, 3, 4]: return output # Already (B, C, H, W) # (T, C, H, W) -> (1, C, T, H, W) return output.unsqueeze(0).permute(0, 2, 1, 3, 4) if output.dim() == 3: c, h, w = output.shape if c > 4 and w <= 4: output = output.permute(2, 0, 1) if output.shape[0] == 1: output = output.repeat(3, 1, 1) return output.unsqueeze(0) if output.dim() == 2: return output.unsqueeze(0).repeat(3, 1, 1).unsqueeze(0) return output def _build_pipeline_kwargs(self, batch: Req) -> dict[str, Any]: """Build kwargs dict for diffusers pipeline call.""" kwargs = {} if batch.prompt is not None: kwargs["prompt"] = batch.prompt if batch.negative_prompt: kwargs["negative_prompt"] = batch.negative_prompt if batch.num_inference_steps is not None: kwargs["num_inference_steps"] = batch.num_inference_steps if batch.guidance_scale is not None: kwargs["guidance_scale"] = batch.guidance_scale if batch.true_cfg_scale is not None: kwargs["true_cfg_scale"] = batch.true_cfg_scale if batch.height is not None: kwargs["height"] = batch.height if batch.width is not None: kwargs["width"] = batch.width if batch.num_frames is not None and batch.num_frames > 1: kwargs["num_frames"] = batch.num_frames # Generator for reproducibility if batch.generator is not None: kwargs["generator"] = batch.generator elif batch.seed is not None: device = self._get_generator_device(batch) kwargs["generator"] = torch.Generator(device=device).manual_seed(batch.seed) # Image input for img2img or inpainting image = self._load_input_image(batch) if image is not None: kwargs["image"] = image if batch.num_outputs_per_prompt > 1: kwargs["num_images_per_prompt"] = batch.num_outputs_per_prompt # Extra diffusers-specific kwargs if batch.extra: diffusers_kwargs = batch.extra.get("diffusers_kwargs", {}) if diffusers_kwargs: kwargs.update(diffusers_kwargs) return kwargs def _get_generator_device(self, batch: Req) -> str: """Resolve RNG device consistently with the non-diffusers path. Diffusers CPU offload can temporarily park modules on CPU, but that should not silently switch a CUDA request to CPU RNG, otherwise the same seed produces different outputs depending on runtime placement. """ if batch.generator_device == "cpu": return "cpu" return current_platform.device_type def _load_input_image(self, batch: Req) -> Image.Image | None: """Load input image from batch.""" # Check for PIL image in condition_image or pixel_values if batch.condition_image is not None and isinstance( batch.condition_image, Image.Image ): return batch.condition_image if batch.pixel_values is not None and isinstance( batch.pixel_values, Image.Image ): return batch.pixel_values if not batch.image_path: return None if isinstance(batch.image_path, list): batch.image_path = batch.image_path[0] try: if batch.image_path.startswith(("http://", "https://")): response = requests.get(batch.image_path, timeout=30) response.raise_for_status() return Image.open(BytesIO(response.content)).convert("RGB") return Image.open(batch.image_path).convert("RGB") except Exception as e: logger.error("Failed to load image from %s: %s", batch.image_path, e) return None class DiffusersPipeline(ComposedPipelineBase): """ Pipeline wrapper that uses vanilla diffusers pipelines. This allows running any diffusers-supported model through sglang's infrastructure without requiring native sglang implementation. """ pipeline_name = "DiffusersPipeline" is_video_pipeline = False _required_config_modules: list[str] = [] def __init__( self, model_path: str, server_args: ServerArgs, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None, executor: PipelineExecutor | None = None, ): self.server_args = server_args self.model_path = model_path self._stages: list[PipelineStage] = [] self._stage_name_mapping: dict[str, PipelineStage] = {} self.modules: dict[str, Any] = {} self.memory_usages: dict[str, float] = {} self.post_init_called = False self.executor = executor or SyncExecutor(server_args=server_args) self._cache_dit_enabled = False logger.info("Loading diffusers pipeline from %s", model_path) self.diffusers_pipe = self._load_diffusers_pipeline(model_path, server_args) self._detect_pipeline_type() def _load_diffusers_pipeline( self, model_path: str, server_args: ServerArgs ) -> DiffusionPipeline: """Load the diffusers pipeline. Optimizations applied: - device_map: Loads models directly to GPU, warming up CUDA caching allocator to avoid small tensor allocations during inference. - Parallel shard loading: When using device_map with accelerate, model shards are loaded in parallel for faster initialization. """ original_model_path = model_path # Keep original for custom_pipeline model_path = maybe_download_model(model_path, force_diffusers_model=True) self.model_path = model_path dtype = self._get_dtype(server_args) logger.info("Loading diffusers pipeline with dtype=%s", dtype) # Build common kwargs for from_pretrained load_kwargs = { "torch_dtype": dtype, "trust_remote_code": server_args.trust_remote_code, "revision": server_args.revision, } # Add quantization config if provided (e.g., BitsAndBytesConfig for 4/8-bit) quant_config = getattr(server_args.pipeline_config, "quantization_config", None) if quant_config is not None: load_kwargs["quantization_config"] = quant_config logger.info("Using quantization config: %s", type(quant_config).__name__) try: pipe = DiffusionPipeline.from_pretrained(model_path, **load_kwargs) except AttributeError as e: if "has no attribute" in str(e): # Custom pipeline class not in diffusers - try loading with custom_pipeline logger.info( "Pipeline class not found in diffusers, trying custom_pipeline from repo..." ) try: custom_kwargs = { **load_kwargs, "custom_pipeline": original_model_path, } custom_kwargs["trust_remote_code"] = True pipe = DiffusionPipeline.from_pretrained( model_path, **custom_kwargs ) except Exception as e2: match = re.search(r"has no attribute (\w+)", str(e)) class_name = match.group(1) if match else "unknown" raise RuntimeError( f"Pipeline class '{class_name}' not found in diffusers and no custom pipeline.py in repo. " f"Try: pip install --upgrade diffusers (some pipelines require latest version). " f"Original error: {e}" ) from e2 else: raise except Exception as e: # Only retry with float32 for dtype-related errors if "dtype" in str(e).lower() or "float" in str(e).lower(): logger.warning( "Failed with dtype=%s, falling back to float32: %s", dtype, e ) load_kwargs["torch_dtype"] = torch.float32 pipe = DiffusionPipeline.from_pretrained(model_path, **load_kwargs) else: raise # Use CPU offload (all-or-nothing in diffusers) if any component offload is requested. any_offload = ( server_args.dit_cpu_offload or server_args.text_encoder_cpu_offload or server_args.image_encoder_cpu_offload or server_args.vae_cpu_offload ) if any_offload: device = get_local_torch_device() gpu_id = device.index if device.index is not None else 0 pipe.enable_model_cpu_offload(gpu_id=gpu_id) logger.info( "Enabled model CPU offload for diffusers pipeline (gpu_id=%d)", gpu_id ) else: pipe = pipe.to(get_local_torch_device()) # Apply VAE memory optimizations from pipeline config self._apply_vae_optimizations(pipe, server_args) # Apply attention backend if specified self._apply_attention_backend(pipe, server_args) # Apply cache-dit acceleration if configured pipe = self._apply_cache_dit(pipe, server_args) # Apply torch.compile if enabled and supported pipe = self._apply_torch_compile(pipe, server_args) logger.info("Loaded diffusers pipeline: %s", pipe.__class__.__name__) return pipe def _apply_vae_optimizations( self, pipe: DiffusionPipeline, server_args: ServerArgs ) -> None: """Apply VAE memory optimizations (tiling, slicing) from pipeline config.""" config = server_args.pipeline_config # VAE slicing: decode latents slice-by-slice for lower peak memory # https://huggingface.co/docs/diffusers/optimization/memory#vae-slicing if config.vae_slicing: if hasattr(pipe, "vae") and hasattr(pipe.vae, "enable_slicing"): pipe.vae.enable_slicing() logger.info("Enabled VAE slicing for lower memory usage") elif hasattr(pipe, "enable_vae_slicing"): pipe.enable_vae_slicing() logger.info("Enabled VAE slicing for lower memory usage") else: logger.warning( "VAE slicing is not available: neither " "`pipe.vae.enable_slicing()` nor `pipe.enable_vae_slicing()` was found." ) # VAE tiling: decode latents tile-by-tile for large images # https://huggingface.co/docs/diffusers/optimization/memory#vae-tiling if config.vae_tiling: if hasattr(pipe, "vae") and hasattr(pipe.vae, "enable_tiling"): pipe.vae.enable_tiling() logger.info("Enabled VAE tiling for large image support") elif hasattr(pipe, "enable_vae_tiling"): pipe.enable_vae_tiling() logger.info("Enabled VAE tiling for large image support") else: logger.warning( "VAE tiling is not available: neither " "`pipe.vae.enable_tiling()` nor `pipe.enable_vae_tiling()` was found." ) def _apply_attention_backend( self, pipe: DiffusionPipeline, server_args: ServerArgs ) -> None: """Apply attention backend setting from pipeline config or server_args. See: https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends Available backends: flash, _flash_3_hub, sage, xformers, native, etc. """ backend = server_args.attention_backend if backend is None: backend = getattr( server_args.pipeline_config, "diffusers_attention_backend", None ) if backend is None: return backend = backend.lower() sglang_backends = {e.name.lower() for e in AttentionBackendEnum} | { "fa3", "fa4", } if backend in sglang_backends: logger.debug( "Skipping diffusers attention backend '%s' because it matches a " "SGLang backend name. Use diffusers backend names when running " "the diffusers backend.", backend, ) return for component_name in ["transformer", "unet"]: component = getattr(pipe, component_name, None) if component is not None and hasattr(component, "set_attention_backend"): try: component.set_attention_backend(backend) logger.info( "Set attention backend '%s' on %s", backend, component_name ) except Exception as e: logger.warning( "Failed to set attention backend '%s' on %s: %s", backend, component_name, e, ) def _apply_cache_dit( self, pipe: DiffusionPipeline, server_args: ServerArgs ) -> DiffusionPipeline: """Enable cache-dit for diffusers pipeline if configured.""" cache_dit_config = server_args.cache_dit_config if not cache_dit_config: return pipe try: import cache_dit except ImportError as e: raise RuntimeError( "cache-dit is required for --cache-dit-config. " "Install it with `pip install cache-dit`." ) from e if not hasattr(cache_dit, "load_configs"): raise RuntimeError( "cache-dit>=1.2.0 is required for --cache-dit-config. " "Please upgrade cache-dit." ) try: cache_options = cache_dit.load_configs(cache_dit_config) except Exception as e: raise ValueError( "Failed to load cache-dit config. Provide a YAML/JSON path (or a dict " "supported by cache-dit>=1.2.0)." ) from e try: pipe = cache_dit.enable_cache(pipe, **cache_options) except Exception: # cache-dit is an external integration and can raise a variety of errors. logger.exception("Failed to enable cache-dit for diffusers pipeline") raise logger.info("Enabled cache-dit for diffusers pipeline") self._cache_dit_enabled = True return pipe def _apply_torch_compile(self, pipe: Any, server_args: ServerArgs) -> Any: """Apply torch.compile to the pipeline if configured and supported.""" if not server_args.enable_torch_compile: return pipe # check if the pipeline has 'transformer' or 'unet' components which are # typically the most expensive parts to compile. 'transformer_2' for some # video pipelines, e.g, Wan 2.2 series, also check for that. compilable_components = ["transformer", "transformer_2", "unet"] if not any(hasattr(pipe, comp) for comp in compilable_components): logger.warning( "Pipeline does not have 'transformer' or 'unet' components. " "torch.compile may not provide significant benefits and could increase latency." ) return pipe if self._cache_dit_enabled: try: import cache_dit if hasattr(cache_dit, "set_compile_configs"): cache_dit.set_compile_configs() except Exception as e: logger.warning( f"Failed to set torch_compile configs for cache-dit: {e}" ) for comp in compilable_components: if hasattr(pipe, comp): try: component = getattr(pipe, comp) # TODO(DefTruth): Add support for 'compile_repeated_blocks' for 'transformer' # modules which can significantly reduce compilation time for large models # with repeated blocks. if isinstance(component, torch.nn.Module) and hasattr( component, "compile" ): # Prefer in-place compilation if supported. According to PyTorch documentation: # https://docs.pytorch.org/docs/stable/generated/torch.compile.html component.compile() else: compiled_component = torch.compile(component) setattr(pipe, comp, compiled_component) logger.info( f"Applied torch.compile to {comp} component of the pipeline" ) except Exception as e: logger.warning(f"Failed to apply torch.compile to {comp}: {e}") return pipe def _get_dtype(self, server_args: ServerArgs) -> torch.dtype: dtype = ( torch.bfloat16 if torch.get_device_module().is_bf16_supported() else torch.float16 ) dit_precision = server_args.pipeline_config.dit_precision if dit_precision == "fp16": dtype = torch.float16 elif dit_precision == "bf16": dtype = torch.bfloat16 elif dit_precision == "fp32": dtype = torch.float32 return dtype def _detect_pipeline_type(self) -> None: """Detect if this is an image or video pipeline.""" pipe_class_name = self.diffusers_pipe.__class__.__name__.lower() video_indicators = ["video", "animat", "cogvideo", "wan", "hunyuan"] self.is_video_pipeline = any(ind in pipe_class_name for ind in video_indicators) logger.debug( "Detected pipeline type: %s", "video" if self.is_video_pipeline else "image", ) def load_modules( self, server_args: ServerArgs, loaded_modules: dict[str, torch.nn.Module] | None = None, ) -> dict[str, Any]: """Skip sglang's module loading - diffusers handles it.""" return {"diffusers_pipeline": self.diffusers_pipe} def create_pipeline_stages(self, server_args: ServerArgs) -> None: """Create the execution stage wrapping the diffusers pipeline.""" self.add_stage( stage_name="diffusers_execution", stage=DiffusersExecutionStage(self.diffusers_pipe), ) def initialize_pipeline(self, server_args: ServerArgs) -> None: pass def post_init(self) -> None: """Post initialization hook.""" if self.post_init_called: return self.post_init_called = True self.initialize_pipeline(self.server_args) self.create_pipeline_stages(self.server_args) def add_stage(self, stage_name: str, stage: PipelineStage) -> None: """Add a stage to the pipeline.""" if stage_name is None: stage_name = self._infer_stage_name(stage) if stage_name in self._stage_name_mapping: raise ValueError(f"Duplicate stage name detected: {stage_name}") self._stages.append(stage) self._stage_name_mapping[stage_name] = stage return self @property def stages(self) -> list[PipelineStage]: """List of stages in the pipeline.""" return self._stages @torch.no_grad() def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Execute the pipeline on the given batch.""" if not self.post_init_called: self.post_init() return self.executor.execute(self.stages, batch, server_args) @classmethod def from_pretrained( cls, model_path: str, device: str | None = None, torch_dtype: torch.dtype | None = None, pipeline_config: str | PipelineConfig | None = None, args: argparse.Namespace | None = None, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None, **kwargs, ) -> "DiffusersPipeline": """Load a pipeline from a pretrained model using diffusers backend.""" kwargs["model_path"] = model_path server_args = ServerArgs.from_kwargs(**kwargs) pipe = cls( model_path, server_args, required_config_modules=required_config_modules, loaded_modules=loaded_modules, ) pipe.post_init() return pipe def get_module(self, module_name: str, default_value: Any = None) -> Any: """Get a module by name.""" if module_name == "diffusers_pipeline": return self.diffusers_pipe return self.modules.get(module_name, default_value) EntryClass = DiffusersPipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/flux.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages import ( InputValidationStage, TextEncodingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger # TODO(will): move PRECISION_TO_TYPE to better place logger = init_logger(__name__) def calculate_shift( image_seq_len, base_seq_len: int = 256, max_seq_len: int = 4096, base_shift: float = 0.5, max_shift: float = 1.15, ): m = (max_shift - base_shift) / (max_seq_len - base_seq_len) b = base_shift - m * base_seq_len mu = image_seq_len * m + b return mu def prepare_mu(batch: Req, server_args: ServerArgs): height = batch.height width = batch.width vae_scale_factor = ( server_args.pipeline_config.vae_config.arch_config.vae_scale_factor ) image_seq_len = (int(height) // vae_scale_factor) * (int(width) // vae_scale_factor) mu = calculate_shift( image_seq_len, # hard code, since scheduler_config is not in PipelineConfig now 256, 4096, 0.5, 1.15, ) return "mu", mu class FluxPipeline(LoRAPipeline, ComposedPipelineBase): pipeline_name = "FluxPipeline" _required_config_modules = [ "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae", "transformer", "scheduler", ] def create_pipeline_stages(self, server_args: ServerArgs): self.add_stage(InputValidationStage()) self.add_stage( TextEncodingStage( text_encoders=[ self.get_module("text_encoder"), self.get_module("text_encoder_2"), ], tokenizers=[ self.get_module("tokenizer"), self.get_module("tokenizer_2"), ], ), "prompt_encoding_stage_primary", ) self.add_standard_timestep_preparation_stage(prepare_extra_kwargs=[prepare_mu]) self.add_standard_latent_preparation_stage() self.add_standard_denoising_stage() self.add_standard_decoding_stage() EntryClass = FluxPipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/flux_2.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from diffusers.image_processor import VaeImageProcessor from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline, Req from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) def compute_empirical_mu(batch: Req, server_args: ServerArgs): num_steps = batch.num_inference_steps image_seq_len = batch.raw_latent_shape[1] a1, b1 = 8.73809524e-05, 1.89833333 a2, b2 = 0.00016927, 0.45666666 if image_seq_len > 4300: mu = a2 * image_seq_len + b2 return "mu", float(mu) m_200 = a2 * image_seq_len + b2 m_10 = a1 * image_seq_len + b1 a = (m_200 - m_10) / 190.0 b = m_200 - 200.0 * a mu = a * num_steps + b return "mu", float(mu) class Flux2Pipeline(LoRAPipeline, ComposedPipelineBase): pipeline_name = "Flux2Pipeline" _required_config_modules = [ "text_encoder", "tokenizer", "vae", "transformer", "scheduler", ] def create_pipeline_stages(self, server_args: ServerArgs): vae_image_processor = VaeImageProcessor( vae_scale_factor=server_args.pipeline_config.vae_config.arch_config.vae_scale_factor * 2 ) self.add_standard_ti2i_stages( include_input_validation=True, vae_image_processor=vae_image_processor, prompt_encoding="text", image_vae_stage_kwargs={"vae_image_processor": vae_image_processor}, prepare_extra_timestep_kwargs=[compute_empirical_mu], ) EntryClass = Flux2Pipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/flux_2_klein.py ================================================ from sglang.multimodal_gen.runtime.pipelines.flux_2 import Flux2Pipeline class Flux2KleinPipeline(Flux2Pipeline): pipeline_name = "Flux2KleinPipeline" EntryClass = Flux2KleinPipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/glm_image.py ================================================ from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.stages import DenoisingStage from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.glm_image import ( GlmImageBeforeDenoisingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase): pipeline_name = "GlmImagePipeline" _required_config_modules = [ "text_encoder", "tokenizer", "vae", "vision_language_encoder", "processor", "transformer", "scheduler", ] def create_pipeline_stages(self, server_args: ServerArgs): self.add_stage( GlmImageBeforeDenoisingStage( vae=self.get_module("vae"), text_encoder=self.get_module("text_encoder"), tokenizer=self.get_module("tokenizer"), processor=self.get_module("processor"), transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), vision_language_encoder=self.get_module("vision_language_encoder"), ), "glm_image_before_denoising_stage", ) self.add_stage( DenoisingStage( transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), ), ) self.add_standard_decoding_stage() EntryClass = [GlmImagePipeline] ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/helios_pipeline.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Helios video diffusion pipeline implementation. This module contains an implementation of the Helios video diffusion pipeline using the modular pipeline architecture. Phase 1: T2V only. """ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.stages import ( InputValidationStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.helios_decoding import ( HeliosDecodingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.helios_denoising import ( HeliosChunkedDenoisingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class HeliosPipeline(LoRAPipeline, ComposedPipelineBase): """ Helios video diffusion pipeline with LoRA support. Implements the Helios T2V pipeline with chunked denoising, multi-term memory history, and CFG Zero Star guidance. """ pipeline_name = "HeliosPipeline" _required_config_modules = [ "text_encoder", "tokenizer", "vae", "transformer", "scheduler", ] def initialize_pipeline(self, server_args: ServerArgs): # Use the scheduler loaded from model's scheduler_config.json as-is. # It contains critical config: use_dynamic_shifting=true, # time_shift_type="exponential", etc. scheduler = self.modules.get("scheduler") if scheduler is not None and server_args.pipeline_config.flow_shift is not None: scheduler.set_shift(server_args.pipeline_config.flow_shift) # Configure scheduler for Stage 2/3 if enabled pipeline_config = server_args.pipeline_config if scheduler is not None and pipeline_config.is_enable_stage2: scheduler.config.stages = pipeline_config.pyramid_num_stages scheduler.config.scheduler_type = pipeline_config.scheduler_type scheduler.config.gamma = pipeline_config.gamma scheduler.init_sigmas_for_each_stage() def create_pipeline_stages(self, server_args: ServerArgs) -> None: self.add_stage(InputValidationStage()) self.add_standard_text_encoding_stage() self.add_standard_latent_preparation_stage() # Skip standard timestep preparation — the Helios denoising stage # handles scheduler.set_timesteps internally per-chunk with mu. self.add_stage( HeliosChunkedDenoisingStage( transformer=self.get_module("transformer"), scheduler=self.modules["scheduler"], ), "helios_chunked_denoising_stage", ) # Helios-specific decoding: decode each chunk's latents separately # to avoid temporal artifacts from Wan VAE causal convolutions self.add_stage( HeliosDecodingStage(vae=self.get_module("vae"), pipeline=self), "helios_decoding_stage", ) class HeliosPyramidPipeline(HeliosPipeline): """Helios pyramid SR pipeline (used by Helios-Mid and Helios-Distilled).""" pipeline_name = "HeliosPyramidPipeline" EntryClass = [HeliosPipeline, HeliosPyramidPipeline] ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py ================================================ """ Hunyuan3D image-to-mesh pipeline implementation. Shape pipeline: BeforeDenoising -> Denoising -> Export -> Save Paint pipeline (optional): Preprocess -> TexGen -> Postprocess """ from __future__ import annotations import glob import importlib import os from itertools import chain from typing import Any import torch import torch.nn as nn from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( Hunyuan3D2PipelineConfig, ) from sglang.multimodal_gen.runtime.loader.fsdp_load import ( load_model_from_full_model_state_dict, set_default_torch_dtype, ) from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.stages import ( Hunyuan3DPaintPostprocessStage, Hunyuan3DPaintPreprocessStage, Hunyuan3DPaintTexGenStage, Hunyuan3DShapeBeforeDenoisingStage, Hunyuan3DShapeDenoisingStage, Hunyuan3DShapeExportStage, Hunyuan3DShapeSaveStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class Hunyuan3D2Pipeline(ComposedPipelineBase): """Hunyuan3D 2.0 image-to-mesh pipeline. Shape pipeline: BeforeDenoising -> Denoising -> Export -> Save Paint pipeline (optional): Preprocess -> TexGen -> Postprocess """ pipeline_name = "Hunyuan3D2Pipeline" _required_config_modules = [ "hy3dshape_model", "hy3dshape_vae", "hy3dshape_scheduler", "hy3dshape_conditioner", "hy3dshape_image_processor", ] def _load_config(self) -> dict[str, Any]: return { "_class_name": self.pipeline_name, "_diffusers_version": "0.0.0", "hy3dshape_model": ["diffusers", "Hunyuan3DShapeModel"], "hy3dshape_vae": ["diffusers", "Hunyuan3DShapeVAE"], "hy3dshape_scheduler": ["diffusers", "Hunyuan3DShapeScheduler"], "hy3dshape_conditioner": ["diffusers", "Hunyuan3DShapeConditioner"], "hy3dshape_image_processor": ["diffusers", "Hunyuan3DShapeImageProcessor"], } # Class resolution @staticmethod def _resolve_class(target: str) -> Any: """Resolve a YAML target string to a Python class.""" from sglang.multimodal_gen.runtime.models.registry import ModelRegistry cls = ModelRegistry.resolve_by_alias(target) if cls is not None: return cls class_name = target.rsplit(".", 1)[-1] try: cls, _ = ModelRegistry.resolve_model_cls(class_name) return cls except Exception: pass from sglang.multimodal_gen.runtime.utils.mesh3d_utils import ( resolve_hunyuan3d_tool, ) for name in (target, class_name): tool_cls = resolve_hunyuan3d_tool(name) if tool_cls is not None: return tool_cls module, cls_name = target.rsplit(".", 1) return getattr(importlib.import_module(module, package=None), cls_name) # Path / checkpoint resolution @staticmethod def _resolve_shape_dir( model_path: str, subfolder: str, use_safetensors: bool, variant: str | None, ) -> tuple[str, str]: """Locate (or download) the shape subfolder and return (config_path, ckpt_path).""" local_path = os.path.join(model_path, subfolder) if not os.path.exists(local_path): local_path = os.path.expanduser(local_path) if not os.path.exists(local_path): logger.info( "Local path %s not found, downloading from HuggingFace Hub", local_path, ) from huggingface_hub import snapshot_download downloaded = snapshot_download( repo_id=model_path, allow_patterns=[f"{subfolder}/*"], ) local_path = os.path.join(downloaded, subfolder) config_path = os.path.join(local_path, "config.yaml") if not os.path.exists(config_path): for alt in ("config.yml", "model_config.yaml"): alt_path = os.path.join(local_path, alt) if os.path.exists(alt_path): config_path = alt_path break if use_safetensors: ckpt_name = ( f"model.{variant}.safetensors" if variant else "model.safetensors" ) else: ckpt_name = f"model-{variant}.ckpt" if variant else "model.ckpt" ckpt_path = os.path.join(local_path, ckpt_name) if not os.path.exists(ckpt_path): pattern = "*.safetensors" if use_safetensors else "*.ckpt" files = glob.glob(os.path.join(local_path, pattern)) if files: ckpt_path = files[0] logger.info("Config path: %s", config_path) logger.info("Checkpoint path: %s", ckpt_path) return config_path, ckpt_path @staticmethod def _resolve_paint_dir(model_path: str, subfolder: str) -> str: """Locate (or download) the paint subfolder and return its local path.""" local_path = os.path.join(model_path, subfolder) if not os.path.exists(local_path): local_path = os.path.expanduser(local_path) if not os.path.exists(local_path): logger.info( "Local path %s not found, downloading from HuggingFace Hub", local_path, ) from huggingface_hub import snapshot_download downloaded = snapshot_download( repo_id=model_path, allow_patterns=[f"{subfolder}/*"], ) local_path = os.path.join(downloaded, subfolder) for subdir in ("vae", "unet"): config_file = os.path.join(local_path, subdir, "config.json") if not os.path.exists(config_file): raise FileNotFoundError( f"Paint model incomplete: {config_file} not found. " "Download the model or check network connectivity." ) logger.info("Resolved paint model directory: %s", local_path) return local_path @staticmethod def _load_and_split_checkpoint( ckpt_path: str, use_safetensors: bool ) -> dict[str, dict[str, torch.Tensor]]: """Load a bundled checkpoint and split by the first '.' in each key.""" if use_safetensors: import safetensors.torch flat = safetensors.torch.load_file(ckpt_path, device="cpu") ckpt: dict[str, dict[str, torch.Tensor]] = {} for key, value in flat.items(): component = key.split(".")[0] sub_key = key[len(component) + 1 :] ckpt.setdefault(component, {})[sub_key] = value return ckpt else: return torch.load(ckpt_path, map_location="cpu", weights_only=True) # Component loading helpers @classmethod def _load_dit_model( cls, cfg: dict[str, Any], weights: dict[str, torch.Tensor], device: torch.device, dtype: torch.dtype, ) -> nn.Module: """Load the DiT model using meta-device instantiation + standard weight loading.""" if "target" not in cfg: raise KeyError("Expected key 'target' in model config.") target_cls = cls._resolve_class(cfg["target"]) params = cfg.get("params", {}) if hasattr(target_cls, "build_config_from_params"): dit_config = target_cls.build_config_from_params(params) init_kwargs: dict[str, Any] = {"config": dit_config, "hf_config": {}} else: init_kwargs = params with set_default_torch_dtype(dtype), torch.device("meta"): model = target_cls(**init_kwargs) weight_iterator = ((k, v) for k, v in weights.items()) param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping) load_model_from_full_model_state_dict( model, weight_iterator, device, dtype, strict=False, param_names_mapping=param_names_mapping_fn, ) for name, p in chain(model.named_parameters(), model.named_buffers()): if p.is_meta: raise RuntimeError(f"Unexpected param or buffer {name} on meta device.") if isinstance(p, nn.Parameter): p.requires_grad = False return model.eval() @classmethod def _load_simple_component( cls, cfg: dict[str, Any], weights: dict[str, torch.Tensor] | None, device: torch.device, dtype: torch.dtype, ) -> nn.Module: """Load a component (VAE / conditioner) with direct instantiation + state_dict.""" if "target" not in cfg: raise KeyError("Expected key 'target' in component config.") target_cls = cls._resolve_class(cfg["target"]) params = cfg.get("params", {}) with set_default_torch_dtype(dtype): component = target_cls(**params) if weights is not None: component.load_state_dict(weights, strict=False) component.to(device=device, dtype=dtype) return component.eval() @classmethod def _instantiate_component(cls, cfg: dict[str, Any]) -> Any: """Instantiate a lightweight component (scheduler / image_processor) without weights.""" if "target" not in cfg: raise KeyError("Expected key 'target' in component config.") target_cls = cls._resolve_class(cfg["target"]) params = cfg.get("params", {}) return target_cls(**params) # Module loading override def load_modules( self, server_args: ServerArgs, loaded_modules: dict[str, torch.nn.Module] | None = None, ) -> dict[str, Any]: """Load all Hunyuan3D shape components from a bundled checkpoint.""" import yaml from sglang.multimodal_gen.runtime.distributed import get_local_torch_device config = server_args.pipeline_config if not isinstance(config, Hunyuan3D2PipelineConfig): raise TypeError(f"Expected Hunyuan3D2PipelineConfig, got {type(config)}") model_path = config.shape_model_path or server_args.model_path logger.info("Loading Hunyuan3D shape models from %s", model_path) config_path, ckpt_path = self._resolve_shape_dir( model_path, config.shape_subfolder, config.shape_use_safetensors, config.shape_variant, ) with open(config_path, "r") as f: model_config = yaml.safe_load(f) ckpt = self._load_and_split_checkpoint(ckpt_path, config.shape_use_safetensors) dtype = torch.float16 if config.shape_variant and "bf16" in config.shape_variant: dtype = torch.bfloat16 device = get_local_torch_device() components: dict[str, Any] = {} components["hy3dshape_model"] = self._load_dit_model( model_config["model"], ckpt["model"], device, dtype ) components["hy3dshape_vae"] = self._load_simple_component( model_config["vae"], ckpt.get("vae"), device, dtype ) components["hy3dshape_conditioner"] = self._load_simple_component( model_config["conditioner"], ckpt.get("conditioner"), device, dtype ) components["hy3dshape_scheduler"] = self._instantiate_component( model_config["scheduler"] ) components["hy3dshape_image_processor"] = self._instantiate_component( model_config["image_processor"] ) logger.info("All Hunyuan3D shape components loaded successfully") if config.paint_enable: try: paint_dir = self._resolve_paint_dir( server_args.model_path, config.paint_subfolder ) components["hy3dpaint_dir"] = paint_dir except Exception as e: logger.warning("Failed to resolve paint model path: %s", e) return components # Pipeline lifecycle def initialize_pipeline(self, server_args: ServerArgs): config = server_args.pipeline_config if not isinstance(config, Hunyuan3D2PipelineConfig): raise TypeError( "Hunyuan3D2Pipeline requires Hunyuan3D2PipelineConfig, " f"got {type(config)}" ) def create_pipeline_stages(self, server_args: ServerArgs): config = server_args.pipeline_config assert isinstance(config, Hunyuan3D2PipelineConfig) # Shape: 4 stages self.add_stage( stage_name="shape_before_denoising", stage=Hunyuan3DShapeBeforeDenoisingStage( image_processor=self.get_module("hy3dshape_image_processor"), conditioner=self.get_module("hy3dshape_conditioner"), vae=self.get_module("hy3dshape_vae"), model=self.get_module("hy3dshape_model"), scheduler=self.get_module("hy3dshape_scheduler"), config=config, ), ) self.add_stage( stage_name="shape_denoising", stage=Hunyuan3DShapeDenoisingStage( transformer=self.get_module("hy3dshape_model"), scheduler=self.get_module("hy3dshape_scheduler"), ), ) self.add_stage( stage_name="shape_export", stage=Hunyuan3DShapeExportStage( vae=self.get_module("hy3dshape_vae"), config=config, ), ) self.add_stage( stage_name="shape_save", stage=Hunyuan3DShapeSaveStage(config=config), ) # Paint: 3 stages (optional) if config.paint_enable: self.add_stage( stage_name="paint_preprocess", stage=Hunyuan3DPaintPreprocessStage(config=config), ) self.add_stage( stage_name="paint_texgen", stage=Hunyuan3DPaintTexGenStage( config=config, paint_dir=self.get_module("hy3dpaint_dir"), ), ) self.add_stage( stage_name="paint_postprocess", stage=Hunyuan3DPaintPostprocessStage(config=config), ) EntryClass = Hunyuan3D2Pipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Hunyuan video diffusion pipeline implementation. This module contains an implementation of the Hunyuan video diffusion pipeline using the modular pipeline architecture. """ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.stages import ( InputValidationStage, TextEncodingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger # TODO(will): move PRECISION_TO_TYPE to better place logger = init_logger(__name__) class HunyuanVideoPipeline(ComposedPipelineBase): pipeline_name = "HunyuanVideoPipeline" _required_config_modules = [ "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae", "transformer", "scheduler", ] def create_pipeline_stages(self, server_args: ServerArgs): self.add_stage(InputValidationStage()) self.add_stage( TextEncodingStage( text_encoders=[ self.get_module("text_encoder"), self.get_module("text_encoder_2"), ], tokenizers=[ self.get_module("tokenizer"), self.get_module("tokenizer_2"), ], ), "prompt_encoding_stage_primary", ) self.add_standard_timestep_preparation_stage() self.add_standard_latent_preparation_stage() self.add_standard_denoising_stage() self.add_standard_decoding_stage() EntryClass = HunyuanVideoPipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py ================================================ import inspect import json import math import os import numpy as np import torch from diffusers import FlowMatchEulerDiscreteScheduler from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages import ( InputValidationStage, LTX2AVDecodingStage, LTX2AVDenoisingStage, LTX2AVLatentPreparationStage, LTX2TextConnectorStage, TextEncodingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) def calculate_shift( image_seq_len, base_seq_len: int = 256, max_seq_len: int = 4096, base_shift: float = 0.5, max_shift: float = 1.15, ): m = (max_shift - base_shift) / (max_seq_len - base_seq_len) b = base_shift - m * base_seq_len mu = image_seq_len * m + b return mu def prepare_mu(batch: Req, server_args: ServerArgs): height = batch.height width = batch.width num_frames = batch.num_frames vae_arch = getattr( getattr(server_args.pipeline_config, "vae_config", None), "arch_config", None ) vae_scale_factor = ( getattr(vae_arch, "spatial_compression_ratio", None) or getattr(vae_arch, "vae_scale_factor", None) or getattr(server_args.pipeline_config, "vae_scale_factor", None) ) vae_temporal_compression = getattr( vae_arch, "temporal_compression_ratio", None ) or getattr(server_args.pipeline_config, "vae_temporal_compression", None) # Values from LTX2Pipeline in diffusers mu = calculate_shift( 4096, base_seq_len=1024, max_seq_len=4096, base_shift=0.95, max_shift=2.05, ) return "mu", mu def _load_component_config(model_path: str, component_name: str): """Helper to load component config from model_index.json or config.json""" try: # Try loading model_index.json first index_path = os.path.join(model_path, "model_index.json") if os.path.exists(index_path): with open(index_path, "r") as f: index = json.load(f) if component_name in index: # It's a subfolder subfolder = index[component_name][1] config_path = os.path.join(model_path, subfolder, "config.json") if os.path.exists(config_path): with open(config_path, "r") as f: return json.load(f) # Fallback to direct config.json in subfolder if standard structure config_path = os.path.join(model_path, component_name, "config.json") if os.path.exists(config_path): with open(config_path, "r") as f: return json.load(f) except Exception as e: logger.warning(f"Failed to load config for {component_name}: {e}") return {} def _filter_kwargs_for_cls(cls, kwargs): """Filter kwargs to only include those accepted by cls.__init__""" sig = inspect.signature(cls.__init__) return {k: v for k, v in kwargs.items() if k in sig.parameters} class LTX2FlowMatchScheduler(FlowMatchEulerDiscreteScheduler): """Override ``_time_shift_exponential`` to use torch f32 instead of numpy f64.""" def _time_shift_exponential(self, mu, sigma, t): if isinstance(t, np.ndarray): t_torch = torch.from_numpy(t).to(torch.float32) result = math.exp(mu) / (math.exp(mu) + (1 / t_torch - 1) ** sigma) return result.numpy() return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) class LTX2Pipeline(ComposedPipelineBase): # NOTE: must match `model_index.json`'s `_class_name` for native dispatch. pipeline_name = "LTX2Pipeline" _required_config_modules = [ "transformer", "text_encoder", "tokenizer", "scheduler", "vae", "audio_vae", "vocoder", "connectors", ] def initialize_pipeline(self, server_args: ServerArgs): orig = self.get_module("scheduler") self.modules["scheduler"] = LTX2FlowMatchScheduler.from_config(orig.config) def create_pipeline_stages(self, server_args: ServerArgs): self.add_stages( [ InputValidationStage(), TextEncodingStage( text_encoders=[self.get_module("text_encoder")], tokenizers=[self.get_module("tokenizer")], ), LTX2TextConnectorStage(connectors=self.get_module("connectors")), ] ) self.add_standard_timestep_preparation_stage(prepare_extra_kwargs=[prepare_mu]) self.add_stages( [ LTX2AVLatentPreparationStage( scheduler=self.get_module("scheduler"), transformer=self.get_module("transformer"), audio_vae=self.get_module("audio_vae"), ), LTX2AVDenoisingStage( transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), vae=self.get_module("vae"), audio_vae=self.get_module("audio_vae"), ), LTX2AVDecodingStage( vae=self.get_module("vae"), audio_vae=self.get_module("audio_vae"), vocoder=self.get_module("vocoder"), pipeline=self, ), ] ) EntryClass = LTX2Pipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/mova_pipeline.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ MOVA pipeline integration (native SGLang pipeline). """ from __future__ import annotations from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig from sglang.multimodal_gen.configs.sample.mova import MOVASamplingParams from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.stages import ( ImageVAEEncodingStage, InputValidationStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.mova import ( MOVADecodingStage, MOVADenoisingStage, MOVALatentPreparationStage, MOVATimestepPreparationStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class MOVAPipeline(ComposedPipelineBase): """MOVA pipeline with SGLang stage orchestration.""" pipeline_name = "MOVA" is_video_pipeline = True _required_config_modules = [ "video_vae", "audio_vae", "text_encoder", "tokenizer", "scheduler", "video_dit", "video_dit_2", "audio_dit", "dual_tower_bridge", ] pipeline_config_cls = MOVAPipelineConfig sampling_params_cls = MOVASamplingParams def initialize_pipeline(self, server_args: ServerArgs) -> None: """ Initialize the pipeline. MOVA supports Context Parallel (sequence parallel) through USPAttention, which uses Ulysses-style all-to-all communication for distributed attention. """ if server_args.sp_degree > 1: logger.info( "MOVA Context Parallel enabled with sp_degree=%d. " "Using USPAttention for distributed self-attention.", server_args.sp_degree, ) def create_pipeline_stages(self, server_args: ServerArgs) -> None: self.add_stage(InputValidationStage()) self.add_standard_text_encoding_stage() if getattr(self.get_module("video_dit"), "require_vae_embedding", True): self.add_stage(ImageVAEEncodingStage(vae=self.get_module("video_vae"))) self.add_stage( MOVALatentPreparationStage( audio_vae=self.get_module("audio_vae"), require_vae_embedding=getattr( self.get_module("video_dit"), "require_vae_embedding", True ), ), "mova_latent_preparation_stage", ) self.add_stage( MOVATimestepPreparationStage( scheduler=self.get_module("scheduler"), ), "mova_timestep_preparation_stage", ) self.add_stage( MOVADenoisingStage( video_dit=self.get_module("video_dit"), video_dit_2=self.get_module("video_dit_2"), audio_dit=self.get_module("audio_dit"), dual_tower_bridge=self.get_module("dual_tower_bridge"), scheduler=self.get_module("scheduler"), ), "mova_denoising_stage", ) self.add_stage( MOVADecodingStage( video_vae=self.get_module("video_vae"), audio_vae=self.get_module("audio_vae"), ), "mova_decoding_stage", ) class MOVAPipelineAlias(MOVAPipeline): pipeline_name = "MOVAPipeline" EntryClass = [MOVAPipeline, MOVAPipelineAlias] ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from diffusers.image_processor import VaeImageProcessor from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.qwen_image_layered import ( QwenImageLayeredBeforeDenoisingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger # TODO(will): move PRECISION_TO_TYPE to better place logger = init_logger(__name__) def calculate_shift( image_seq_len, base_seq_len: int = 256, max_seq_len: int = 4096, base_shift: float = 0.5, max_shift: float = 1.15, ): m = (max_shift - base_shift) / (max_seq_len - base_seq_len) b = base_shift - m * base_seq_len mu = image_seq_len * m + b return mu def prepare_mu(batch: Req, server_args: ServerArgs): height = batch.height width = batch.width vae_scale_factor = server_args.pipeline_config.vae_config.vae_scale_factor image_seq_len = (int(height) // vae_scale_factor // 2) * ( int(width) // vae_scale_factor // 2 ) mu = calculate_shift( image_seq_len, # hard code, since scheduler_config is not in PipelineConfig now 256, 8192, 0.5, 0.9, ) return "mu", mu class QwenImagePipeline(LoRAPipeline, ComposedPipelineBase): pipeline_name = "QwenImagePipeline" _required_config_modules = [ "text_encoder", "tokenizer", "vae", "transformer", "scheduler", ] def create_pipeline_stages(self, server_args: ServerArgs): self.add_standard_t2i_stages(prepare_extra_timestep_kwargs=[prepare_mu]) class QwenImageEditPipeline(LoRAPipeline, ComposedPipelineBase): pipeline_name = "QwenImageEditPipeline" _required_config_modules = [ "processor", "scheduler", "text_encoder", "tokenizer", "transformer", "vae", ] def create_pipeline_stages(self, server_args: ServerArgs): vae_image_processor = VaeImageProcessor( vae_scale_factor=server_args.pipeline_config.vae_config.arch_config.vae_scale_factor * 2 ) self.add_standard_ti2i_stages( vae_image_processor=vae_image_processor, prompt_encoding="image_encoding", image_processor_key="processor", prompt_text_encoder_key="text_encoder", prepare_extra_timestep_kwargs=[prepare_mu], ) class QwenImageEditPlusPipeline(QwenImageEditPipeline): pipeline_name = "QwenImageEditPlusPipeline" def prepare_mu_layered(batch: Req, server_args: ServerArgs): base_seqlen = 256 * 256 / 16 / 16 mu = (batch.image_latent.shape[1] / base_seqlen) ** 0.5 return "mu", mu class QwenImageLayeredPipeline(QwenImageEditPipeline): pipeline_name = "QwenImageLayeredPipeline" _required_config_modules = [ "vae", "tokenizer", "processor", "transformer", "scheduler", ] def create_pipeline_stages(self, server_args: ServerArgs): self.add_stage( QwenImageLayeredBeforeDenoisingStage( vae=self.get_module("vae"), tokenizer=self.get_module("tokenizer"), processor=self.get_module("processor"), transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), model_path=self.model_path, ) ) self.add_standard_timestep_preparation_stage( prepare_extra_kwargs=[prepare_mu_layered] ) self.add_standard_denoising_stage() self.add_standard_decoding_stage() EntryClass = [ QwenImagePipeline, QwenImageEditPipeline, QwenImageEditPlusPipeline, QwenImageLayeredPipeline, ] ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/sana.py ================================================ # SPDX-License-Identifier: Apache-2.0 # # SANA text-to-image pipeline. # # Stage order matches Flux (InputValidation -> TextEncoding -> TimestepPrep -> # LatentPrep -> Denoising -> Decoding) rather than the add_standard_t2i_stages # helper (which puts LatentPrep before TimestepPrep). Both orderings are # functionally equivalent since these stages are independent. # # SANA uses a single text encoder (Gemma2), so only one text_encoder + tokenizer # pair is registered — unlike Flux which has text_encoder + text_encoder_2. # The pipeline_name must match the _class_name in HF model_index.json. from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.stages import ( InputValidationStage, TextEncodingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class SanaPipeline(LoRAPipeline, ComposedPipelineBase): pipeline_name = "SanaPipeline" _required_config_modules = [ "text_encoder", "tokenizer", "vae", "transformer", "scheduler", ] def create_pipeline_stages(self, server_args: ServerArgs): self.add_stage(InputValidationStage()) self.add_stage( TextEncodingStage( text_encoders=[self.get_module("text_encoder")], tokenizers=[self.get_module("tokenizer")], ), "prompt_encoding_stage_primary", ) self.add_standard_timestep_preparation_stage() self.add_standard_latent_preparation_stage() self.add_standard_denoising_stage() self.add_standard_decoding_stage() EntryClass = SanaPipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/wan_causal_dmd_pipeline.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Wan causal DMD pipeline implementation. This module wires the causal DMD denoising stage into the modular pipeline. """ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline # isort: off from sglang.multimodal_gen.runtime.pipelines_core.stages import ( CausalDMDDenoisingStage, InputValidationStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger # isort: on logger = init_logger(__name__) class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase): pipeline_name = "WanCausalDMDPipeline" _required_config_modules = [ "text_encoder", "tokenizer", "vae", "transformer", "scheduler", ] def create_pipeline_stages(self, server_args: ServerArgs) -> None: self.add_stage(InputValidationStage()) self.add_standard_text_encoding_stage() self.add_standard_latent_preparation_stage() self.add_stage( CausalDMDDenoisingStage( transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), ), ) self.add_standard_decoding_stage() EntryClass = WanCausalDMDPipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/wan_dmd_pipeline.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Wan video diffusion pipeline implementation. This module contains an implementation of the Wan video diffusion pipeline using the modular pipeline architecture. """ from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger # isort: off from sglang.multimodal_gen.runtime.pipelines_core.stages import ( DmdDenoisingStage, InputValidationStage, ) # isort: on logger = init_logger(__name__) class WanDMDPipeline(LoRAPipeline, ComposedPipelineBase): """ Wan video diffusion pipeline with LoRA support. """ pipeline_name = "WanDMDPipeline" _required_config_modules = [ "text_encoder", "tokenizer", "vae", "transformer", "scheduler", ] def initialize_pipeline(self, server_args: ServerArgs): self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler( shift=server_args.pipeline_config.flow_shift ) def create_pipeline_stages(self, server_args: ServerArgs) -> None: self.add_stages( [ InputValidationStage(), ] ) self.add_standard_text_encoding_stage() self.add_standard_timestep_preparation_stage() self.add_standard_latent_preparation_stage() self.add_stages( [ DmdDenoisingStage( transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), ), ] ) self.add_standard_decoding_stage() EntryClass = WanDMDPipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Wan video diffusion pipeline implementation. This module contains an implementation of the Wan video diffusion pipeline using the modular pipeline architecture. """ from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.stages import DmdDenoisingStage from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase): pipeline_name = "WanImageToVideoDmdPipeline" _required_config_modules = [ "text_encoder", "tokenizer", "vae", "transformer", "scheduler", "image_encoder", "image_processor", ] def initialize_pipeline(self, server_args: ServerArgs): self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler( shift=server_args.pipeline_config.flow_shift ) def create_pipeline_stages(self, server_args: ServerArgs): self.add_standard_ti2v_stages( image_vae_encoding_position="after_latent", denoising_stage_factory=lambda: DmdDenoisingStage( transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), transformer_2=self.get_module("transformer_2"), ), ) EntryClass = WanImageToVideoDmdPipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_pipeline.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Wan video diffusion pipeline implementation. This module contains an implementation of the Wan video diffusion pipeline using the modular pipeline architecture. """ from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_unipc_multistep import ( FlowUniPCMultistepScheduler, ) from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase): pipeline_name = "WanImageToVideoPipeline" _required_config_modules = [ "text_encoder", "tokenizer", "vae", "transformer", "scheduler", "image_encoder", "image_processor", ] def initialize_pipeline(self, server_args: ServerArgs): self.modules["scheduler"] = FlowUniPCMultistepScheduler( shift=server_args.pipeline_config.flow_shift ) def create_pipeline_stages(self, server_args: ServerArgs): self.add_standard_ti2v_stages() EntryClass = WanImageToVideoPipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Wan video diffusion pipeline implementation. This module contains an implementation of the Wan video diffusion pipeline using the modular pipeline architecture. """ from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_unipc_multistep import ( FlowUniPCMultistepScheduler, ) from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class WanPipeline(LoRAPipeline, ComposedPipelineBase): """ Wan video diffusion pipeline with LoRA support. """ pipeline_name = "WanPipeline" _required_config_modules = [ "text_encoder", "tokenizer", "vae", "transformer", "scheduler", ] def initialize_pipeline(self, server_args: ServerArgs): # We use UniPCMScheduler from Wan2.1 official repo, not the one in diffusers. self.modules["scheduler"] = FlowUniPCMultistepScheduler( shift=server_args.pipeline_config.flow_shift ) def create_pipeline_stages(self, server_args: ServerArgs) -> None: self.add_standard_t2i_stages() EntryClass = WanPipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline, Req from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) def calculate_shift( image_seq_len, base_seq_len: int = 256, max_seq_len: int = 4096, base_shift: float = 0.5, max_shift: float = 1.15, ): m = (max_shift - base_shift) / (max_seq_len - base_seq_len) b = base_shift - m * base_seq_len mu = image_seq_len * m + b return mu def prepare_mu(batch: Req, server_args: ServerArgs): height = batch.height width = batch.width vae_scale_factor = server_args.pipeline_config.vae_config.vae_scale_factor image_seq_len = ((int(height) // vae_scale_factor) // 2) * ( (int(width) // vae_scale_factor) // 2 ) mu = calculate_shift( image_seq_len, # hard code, since scheduler_config is not in PipelineConfig now 256, 4096, 0.5, 1.15, ) return "mu", mu class ZImagePipeline(LoRAPipeline, ComposedPipelineBase): pipeline_name = "ZImagePipeline" _required_config_modules = [ "text_encoder", "tokenizer", "vae", "transformer", "scheduler", ] def create_pipeline_stages(self, server_args: ServerArgs): self.add_standard_t2i_stages(prepare_extra_timestep_kwargs=[prepare_mu]) EntryClass = ZImagePipeline ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Diffusion pipelines for sglang.multimodal_gen. This package contains diffusion pipelines for generating videos and images. """ from typing import cast from sglang.multimodal_gen.registry import get_model_info from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( maybe_download_model, verify_model_config_and_directory, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class PipelineWithLoRA(LoRAPipeline, ComposedPipelineBase): """Type for a pipeline that has both ComposedPipelineBase and LoRAPipeline functionality.""" pass def build_pipeline( server_args: ServerArgs, ) -> PipelineWithLoRA: """ Only works with valid hf diffusers configs. (model_index.json) We want to build a pipeline based on the inference args mode_path: 1. download the model from the hub if it's not already downloaded 2. verify the model config and directory 3. based on the config, determine the pipeline class """ model_path = server_args.model_path # Check if pipeline class is explicitly specified if server_args.pipeline_class_name: from sglang.multimodal_gen.registry import ( _PIPELINE_REGISTRY, _discover_and_register_pipelines, ) _discover_and_register_pipelines() logger.info(f"Requested pipeline_class_name: {server_args.pipeline_class_name}") logger.info( f"Available pipelines in registry: {list(_PIPELINE_REGISTRY.keys())}" ) pipeline_cls = _PIPELINE_REGISTRY.get(server_args.pipeline_class_name) if pipeline_cls is None: raise ValueError( f"Pipeline class '{server_args.pipeline_class_name}' not found in registry. " f"Available pipelines: {list(_PIPELINE_REGISTRY.keys())}" ) logger.info( f"✓ Using explicitly specified pipeline: {server_args.pipeline_class_name} (class: {pipeline_cls.__name__})" ) else: logger.info("No pipeline_class_name specified, using model_index.json") model_info = get_model_info( model_path, backend=server_args.backend, model_id=server_args.model_id, ) pipeline_cls = model_info.pipeline_cls logger.info(f"Using pipeline from model_index.json: {pipeline_cls.__name__}") # instantiate the pipelines pipeline = pipeline_cls(model_path, server_args) logger.info("Pipeline instantiated") return cast(PipelineWithLoRA, pipeline) __all__ = [ "build_pipeline", "ComposedPipelineBase", "Req", "LoRAPipeline", ] ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Base class for composed pipelines. This module defines the base class for pipelines that are composed of multiple stages. """ import os import re from abc import ABC, abstractmethod from typing import Any, Callable, Literal, cast import torch from tqdm import tqdm from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PipelineComponentLoader, ) from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import ( PipelineExecutor, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages import ( DecodingStage, DenoisingStage, ImageEncodingStage, ImageVAEEncodingStage, InputValidationStage, LatentPreparationStage, PipelineStage, TextEncodingStage, TimestepPreparationStage, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( maybe_download_model, verify_model_config_and_directory, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class ComposedPipelineBase(ABC): """ Base class for pipelines composed of multiple stages. This class provides the framework for creating pipelines by composing multiple stages together. Each stage is responsible for a specific part of the diffusion process, and the pipeline orchestrates the execution of these stages. """ is_video_pipeline: bool = False # To be overridden by video pipelines # should contains only the modules to be loaded _required_config_modules: list[str] = [] _extra_config_module_map: dict[str, str] = {} server_args: ServerArgs | None = None modules: dict[str, Any] = {} executor: PipelineExecutor | None = None # the name of the pipeline it associated with, in diffusers pipeline_name: str def is_lora_effective(self): return False def is_lora_set(self): return False def __init__( self, model_path: str, server_args: ServerArgs, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None, executor: PipelineExecutor | None = None, ): """ Initialize the pipeline. After __init__, the pipeline should be ready to use. The pipeline should be stateless and not hold any batch state. """ self.server_args = server_args self.model_path: str = model_path self._stages: list[PipelineStage] = [] self._stage_name_mapping: dict[str, PipelineStage] = {} self.executor = executor or self.build_executor(server_args=server_args) if required_config_modules is not None: self._required_config_modules = required_config_modules if self._required_config_modules is None: raise NotImplementedError("Subclass must set _required_config_modules") # [module_name, gpu memory usage] self.memory_usages: dict[str, float] = {} # Load modules directly in initialization logger.info("Loading pipeline modules...") self.modules = self.load_modules(server_args, loaded_modules) self.__post_init__() def build_executor(self, server_args: ServerArgs): # TODO from sglang.multimodal_gen.runtime.pipelines_core.executors.parallel_executor import ( ParallelExecutor, ) # return SyncExecutor(server_args=server_args) return ParallelExecutor(server_args=server_args) def __post_init__(self) -> None: assert self.server_args is not None, "server_args must be set" self.initialize_pipeline(self.server_args) logger.info("Creating pipeline stages...") self.create_pipeline_stages(self.server_args) def get_module(self, module_name: str, default_value: Any = None) -> Any: return self.modules.get(module_name, default_value) def add_module(self, module_name: str, module: Any): self.modules[module_name] = module def _load_config(self) -> dict[str, Any]: model_path = maybe_download_model(self.model_path, force_diffusers_model=True) self.model_path = model_path logger.info("Model path: %s", model_path) config = verify_model_config_and_directory(model_path) return cast(dict[str, Any], config) @property def required_config_modules(self) -> list[str]: """ List of modules that are required by the pipeline. The names should match the diffusers directory and model_index.json file. These modules will be loaded using the PipelineComponentLoader and made available in the modules dictionary. Access these modules using the get_module method. class ConcretePipeline(ComposedPipelineBase): _required_config_modules = ["vae", "text_encoder", "transformer", "scheduler", "tokenizer"] @property def required_config_modules(self): return self._required_config_modules """ return self._required_config_modules @property def stages(self) -> list[PipelineStage]: """ List of stages in the pipeline. """ return self._stages @abstractmethod def create_pipeline_stages(self, server_args: ServerArgs): """ Create the inference pipeline stages. """ raise NotImplementedError def initialize_pipeline(self, server_args: ServerArgs): """ Initialize the pipeline. """ return def _resolve_component_path( self, server_args: ServerArgs, module_name: str, load_module_name: str ) -> str: override_path = server_args.component_paths.get(module_name) if override_path is not None: # overridden with args like --vae-path component_model_path = maybe_download_model(override_path) else: component_model_path = os.path.join(self.model_path, load_module_name) logger.debug("Resolved component path: %s", component_model_path) return component_model_path def load_modules( self, server_args: ServerArgs, loaded_modules: dict[str, torch.nn.Module] | None = None, ) -> dict[str, Any]: """ Load the modules from the config. loaded_modules: Optional[Dict[str, torch.nn.Module]] = None, If provided, loaded_modules will be used instead of loading from config/pretrained weights. """ model_index = self._load_config() logger.info("Loading pipeline modules from config: %s", model_index) # remove keys that are not pipeline modules model_index.pop("_class_name") model_index.pop("_diffusers_version") if ( "boundary_ratio" in model_index and model_index["boundary_ratio"] is not None ): has_transformer = ( "transformer" in model_index or "transformer_2" in model_index or "transformer" in self.required_config_modules or "transformer_2" in self.required_config_modules ) if has_transformer: logger.info( "MoE pipeline detected. Adding transformer_2 to self.required_config_modules..." ) if "transformer_2" not in self.required_config_modules: self.required_config_modules.append("transformer_2") else: logger.info( "Boundary ratio found in model_index.json without transformers; " "using it for pipeline config only." ) logger.info( "Setting boundary ratio to %s", model_index["boundary_ratio"], ) server_args.pipeline_config.dit_config.boundary_ratio = model_index[ "boundary_ratio" ] model_index.pop("boundary_ratio", None) # used by Wan2.2 ti2v model_index.pop("expand_timesteps", None) # some sanity checks assert ( len(model_index) > 1 ), "model_index.json must contain at least one pipeline module" model_index = { required_module: model_index[required_module] for required_module in self.required_config_modules } for module_name in self.required_config_modules: if ( module_name not in model_index and module_name in self._extra_config_module_map ): extra_module_value = self._extra_config_module_map[module_name] logger.warning( "model_index.json does not contain a %s module, but found {%s: %s} in _extra_config_module_map, adding to model_index.", module_name, module_name, extra_module_value, ) if extra_module_value in model_index: logger.info( "Using module %s for %s", extra_module_value, module_name ) model_index[module_name] = model_index[extra_module_value] continue else: raise ValueError( f"Required module key: {module_name} value: {model_index.get(module_name)} was not found in loaded modules {model_index.keys()}" ) # all the component models used by the pipeline required_modules = self.required_config_modules logger.info("Loading required components: %s", required_modules) loaded_components = {} for module_name, ( transformers_or_diffusers, architecture, ) in tqdm(iterable=model_index.items(), desc="Loading required modules"): if transformers_or_diffusers is None: logger.warning( "Module %s in model_index.json has null value, removing from required_config_modules", module_name, ) if module_name in self.required_config_modules: self.required_config_modules.remove(module_name) continue if module_name not in required_modules: logger.info("Skipping module %s", module_name) continue if loaded_modules is not None and module_name in loaded_modules: logger.info("Using module %s already provided", module_name) loaded_components[module_name] = loaded_modules[module_name] continue # we load the module from the extra config module map if it exists if module_name in self._extra_config_module_map: load_module_name = self._extra_config_module_map[module_name] else: load_module_name = module_name component_model_path = self._resolve_component_path( server_args, module_name, load_module_name ) module, memory_usage = PipelineComponentLoader.load_component( component_name=load_module_name, component_model_path=component_model_path, transformers_or_diffusers=transformers_or_diffusers, server_args=server_args, ) self.memory_usages[load_module_name] = memory_usage if module_name in loaded_components: logger.warning("Overwriting module %s", module_name) loaded_components[module_name] = module # Check if all required modules were loaded for module_name in required_modules: if ( module_name not in loaded_components or loaded_components[module_name] is None ): raise ValueError( f"Required module: {module_name} was not found in loaded modules: {list(loaded_components.keys())}" ) logger.debug( "Memory usage of loaded modules (GiB): %s. avail mem: %s GB", self.memory_usages, round(current_platform.get_available_gpu_memory(), 2), ) return loaded_components @staticmethod def _infer_stage_name(stage: PipelineStage) -> str: class_name = stage.__class__.__name__ # snake_case name = re.sub(r"(? "ComposedPipelineBase": assert self.modules is not None, "No modules are registered" if stage_name is None: stage_name = self._infer_stage_name(stage) if stage_name in self._stage_name_mapping: raise ValueError(f"Duplicate stage name detected: {stage_name}") self._stages.append(stage) self._stage_name_mapping[stage_name] = stage return self def add_stages( self, stages: list[PipelineStage | tuple[PipelineStage, str]] ) -> "ComposedPipelineBase": for item in stages: if isinstance(item, tuple): stage, name = item self.add_stage(stage, name) else: self.add_stage(item) return self def add_stage_if( self, condition: bool | Callable[[], bool], stage: PipelineStage, ) -> "ComposedPipelineBase": should_add = condition() if callable(condition) else condition if should_add: self.add_stage(stage) return self def get_stage(self, stage_name: str) -> PipelineStage | None: """Get a stage by name.""" return self._stage_name_mapping.get(stage_name) def add_standard_text_encoding_stage( self, text_encoder_key: str = "text_encoder", tokenizer_key: str = "tokenizer", ) -> "ComposedPipelineBase": return self.add_stage( TextEncodingStage( text_encoders=[self.get_module(text_encoder_key)], tokenizers=[self.get_module(tokenizer_key)], ), ) def add_standard_timestep_preparation_stage( self, scheduler_key: str = "scheduler", prepare_extra_kwargs: list[Callable] | None = [], ) -> "ComposedPipelineBase": return self.add_stage( TimestepPreparationStage( scheduler=self.get_module(scheduler_key), prepare_extra_set_timesteps_kwargs=prepare_extra_kwargs, ), ) def add_standard_latent_preparation_stage( self, scheduler_key: str = "scheduler", transformer_key: str = "transformer", ) -> "ComposedPipelineBase": return self.add_stage( LatentPreparationStage( scheduler=self.get_module(scheduler_key), transformer=self.get_module(transformer_key), ), ) def add_standard_denoising_stage( self, transformer_key: str = "transformer", transformer_2_key: str | None = "transformer_2", scheduler_key: str = "scheduler", vae_key: str | None = "vae", ) -> "ComposedPipelineBase": kwargs = { "transformer": self.get_module(transformer_key), "scheduler": self.get_module(scheduler_key), } if transformer_2_key: transformer_2 = self.get_module(transformer_2_key, None) if transformer_2 is not None: kwargs["transformer_2"] = transformer_2 if vae_key: vae = self.get_module(vae_key, None) if vae is not None: kwargs["vae"] = vae kwargs["pipeline"] = self return self.add_stage(DenoisingStage(**kwargs)) def add_standard_decoding_stage( self, vae_key: str = "vae", ) -> "ComposedPipelineBase": return self.add_stage( DecodingStage( vae=self.get_module(vae_key), pipeline=self, component_name=vae_key, ), ) def add_standard_t2i_stages( self, include_input_validation: bool = True, prepare_extra_timestep_kwargs: list[Callable] | None = [], ) -> "ComposedPipelineBase": if include_input_validation: self.add_stage(InputValidationStage()) self.add_standard_text_encoding_stage() self.add_standard_latent_preparation_stage() self.add_standard_timestep_preparation_stage( prepare_extra_kwargs=prepare_extra_timestep_kwargs ) self.add_standard_denoising_stage() self.add_standard_decoding_stage() return self def add_standard_ti2i_stages( self, *, include_input_validation: bool = True, vae_image_processor: Any | None = None, prompt_encoding: Literal["text", "image_encoding"] = "text", text_encoder_key: str = "text_encoder", tokenizer_key: str = "tokenizer", image_processor_key: str = "processor", prompt_text_encoder_key: str = "text_encoder", image_vae_key: str = "vae", image_vae_stage_kwargs: dict[str, Any] | None = None, prepare_extra_timestep_kwargs: list[Callable] | None = [], ) -> "ComposedPipelineBase": if include_input_validation: self.add_stage( InputValidationStage(vae_image_processor=vae_image_processor) ) if prompt_encoding == "text": self.add_standard_text_encoding_stage( text_encoder_key=text_encoder_key, tokenizer_key=tokenizer_key, ) elif prompt_encoding == "image_encoding": self.add_stage( ImageEncodingStage( image_processor=self.get_module(image_processor_key), text_encoder=self.get_module(prompt_text_encoder_key), ), ) else: raise ValueError(f"Unknown prompt_encoding: {prompt_encoding}") self.add_stage( ImageVAEEncodingStage( vae=self.get_module(image_vae_key), **(image_vae_stage_kwargs or {}), ), ) self.add_standard_latent_preparation_stage() self.add_standard_timestep_preparation_stage( prepare_extra_kwargs=prepare_extra_timestep_kwargs ) self.add_standard_denoising_stage() self.add_standard_decoding_stage() return self def add_standard_ti2v_stages( self, *, include_input_validation: bool = True, vae_image_processor: Any | None = None, text_encoder_key: str = "text_encoder", tokenizer_key: str = "tokenizer", image_encoder_key: str = "image_encoder", image_processor_key: str = "image_processor", image_vae_key: str = "vae", image_vae_stage_kwargs: dict[str, Any] | None = None, image_vae_encoding_position: Literal[ "before_timestep", "after_latent" ] = "before_timestep", prepare_extra_timestep_kwargs: list[Callable] | None = [], denoising_stage_factory: Callable[[], PipelineStage] | None = None, ) -> "ComposedPipelineBase": if include_input_validation: self.add_stage( InputValidationStage(vae_image_processor=vae_image_processor) ) self.add_standard_text_encoding_stage( text_encoder_key=text_encoder_key, tokenizer_key=tokenizer_key, ) image_encoder = self.get_module(image_encoder_key, None) image_processor = self.get_module(image_processor_key, None) self.add_stage_if( image_encoder is not None and image_processor is not None, ImageEncodingStage( image_encoder=image_encoder, image_processor=image_processor, ), ) if image_vae_encoding_position == "before_timestep": self.add_stage( ImageVAEEncodingStage( vae=self.get_module(image_vae_key), **(image_vae_stage_kwargs or {}), ) ) self.add_standard_latent_preparation_stage() self.add_standard_timestep_preparation_stage( prepare_extra_kwargs=prepare_extra_timestep_kwargs ) if image_vae_encoding_position == "after_latent": self.add_stage( ImageVAEEncodingStage( vae=self.get_module(image_vae_key), **(image_vae_stage_kwargs or {}), ) ) elif image_vae_encoding_position != "before_timestep": raise ValueError( f"Unknown image_vae_encoding_position: {image_vae_encoding_position}" ) if denoising_stage_factory is None: self.add_standard_denoising_stage() else: self.add_stage(denoising_stage_factory()) self.add_standard_decoding_stage() return self # TODO(will): don't hardcode no_grad @torch.no_grad() def forward( self, batch: Req, server_args: ServerArgs, ) -> OutputBatch: """ Generate a video or image using the pipeline. Args: batch: The batch to generate from. server_args: The inference arguments. Returns: Req: The batch with the generated video or image. """ if self.is_lora_set() and not self.is_lora_effective(): logger.warning( "LoRA adapter is set, but not effective. Please make sure the LoRA weights are merged" ) # Execute each stage if not batch.is_warmup and not batch.suppress_logs: logger.info( "Running pipeline stages: %s", list(self._stage_name_mapping.keys()), main_process_only=True, ) return self.executor.execute_with_profiling(self.stages, batch, server_args) ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from typing import List import torch from sglang.multimodal_gen.runtime.distributed import get_sp_group from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_cfg_group, get_classifier_free_guidance_rank, get_world_rank, ) from sglang.multimodal_gen.runtime.pipelines_core import Req from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import ( PipelineExecutor, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( PipelineStage, StageParallelismType, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.distributed import broadcast_pyobj from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class ParallelExecutor(PipelineExecutor): """ The correctness of the execution relies on the parallelism_type declared by stages """ def collect_from_main(self, batches: list[Req]): # TODO: fix this condition if self.server_args.sp_degree != 1: sp_group = get_sp_group() batches = broadcast_pyobj( batches, sp_group.rank, sp_group.cpu_group, src=sp_group.ranks[0], ) if self.server_args.enable_cfg_parallel: batches = broadcast_pyobj( batches, self.worker.cfg_group.rank, self.worker.cfg_cpu_group, src=self.worker.cfg_group.ranks[0], ) def _execute( self, stages: List[PipelineStage], batch: Req, server_args: ServerArgs, ) -> OutputBatch: """ Execute all pipeline stages respecting their declared parallelism type. """ if server_args.enable_cfg_parallel: rank = get_classifier_free_guidance_rank() else: rank = get_world_rank() cfg_group = get_cfg_group() # TODO: decide when to gather on main when CFG_PARALLEL -> MAIN_RANK_ONLY for stage in stages: paradigm = stage.parallelism_type if paradigm == StageParallelismType.MAIN_RANK_ONLY: if rank == 0: # Only main rank executes, others just wait batch = stage(batch, server_args) torch.distributed.barrier() elif paradigm == StageParallelismType.CFG_PARALLEL: obj_list = [batch] if rank == 0 else [] broadcasted_list = broadcast_pyobj( obj_list, rank=rank, dist_group=cfg_group.cpu_group, src=0 ) if rank != 0: batch = broadcasted_list[0] batch = stage(batch, server_args) torch.distributed.barrier() elif paradigm == StageParallelismType.REPLICATED: batch = stage(batch, server_args) return batch def execute( self, stages: List[PipelineStage], batch: Req, server_args: ServerArgs, ) -> OutputBatch: batch = self._execute(stages, batch, server_args) return batch ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Base class for all pipeline executors. """ import contextlib from abc import ABC, abstractmethod from typing import TYPE_CHECKING, List from sglang.multimodal_gen.runtime.distributed import get_world_rank from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler if TYPE_CHECKING: # Only for type checkers; avoids runtime circular import from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage logger = init_logger(__name__) class Timer(StageProfiler): """ A wrapper around StageProfiler to maintain backward compatibility. It forces simple logging behavior (log start/end) regardless of env vars. """ def __init__(self, name="Stage"): super().__init__( stage_name=name, logger=logger, metrics=None, log_stage_start_end=True ) class PipelineExecutor(ABC): """ Abstract base class for all pipeline executors. Executors orchestrate the execution of pipeline, with managing the parallel and communications required by stages """ def __init__(self, server_args): self.server_args = server_args def execute_with_profiling( self, stages: List["PipelineStage"], batch: Req, server_args: ServerArgs, ) -> OutputBatch: with self.profile_execution(batch, dump_rank=0): batch = self.execute(stages, batch, server_args) return batch @abstractmethod def execute( self, stages: List["PipelineStage"], batch: Req, server_args: ServerArgs, ) -> OutputBatch: """ Execute the pipeline stages. Args: stages: A list of pipeline stages to execute. batch: The batch to process. server_args: The server arguments. Returns: The processed batch. """ raise NotImplementedError @contextlib.contextmanager def profile_execution(self, batch: Req, dump_rank: int = 0): """ Context manager for profiling execution. """ do_profile = batch.profile and not batch.is_warmup if not do_profile: # fast forward yield return request_id = batch.request_id rank = get_world_rank() profiler = SGLDiffusionProfiler( request_id=request_id, rank=rank, full_profile=batch.profile_all_stages, num_steps=batch.num_profiled_timesteps, num_inference_steps=batch.num_inference_steps, ) try: yield finally: profiler.stop(dump_rank=dump_rank) ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Synchronous pipeline executor implementation. """ from typing import List from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import ( PipelineExecutor, SGLDiffusionProfiler, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs class SyncExecutor(PipelineExecutor): """ A simple synchronous executor that runs stages sequentially. """ def run_profile_all_stages( self, stages: List[PipelineStage], batch: Req, server_args: ServerArgs, ) -> OutputBatch: """ Execute all pipeline stages sequentially. """ for stage in stages: batch = stage(batch, server_args) profiler = SGLDiffusionProfiler.get_instance() if profiler: profiler.step_stage() return batch def execute( self, stages: List[PipelineStage], batch: Req, server_args: ServerArgs, ) -> OutputBatch: """ Execute the pipeline stages sequentially. """ batch = self.run_profile_all_stages(stages, batch, server_args) return batch ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/lora_format_adapter.py ================================================ from __future__ import annotations import logging from enum import Enum from typing import Dict, Iterable, Mapping, Optional import torch from diffusers.loaders import lora_conversion_utils as lcu logger = logging.getLogger("LoRAFormatAdapter") class LoRAFormat(str, Enum): """Supported external LoRA formats before normalization.""" STANDARD = "standard" NON_DIFFUSERS_SD = "non-diffusers-sd" QWEN_IMAGE_STANDARD = "qwen-image-standard" XLABS_FLUX = "xlabs-ai" KOHYA_FLUX = "kohya-flux" WAN = "wan" AI_TOOLKIT_FLUX = "ai-toolkit-flux" def _sample_keys(keys: Iterable[str], k: int = 20) -> list[str]: out = [] for i, key in enumerate(keys): if i >= k: break out.append(key) return out def _has_substring_key(keys: Iterable[str], substr: str) -> bool: return any(substr in k for k in keys) def _has_prefix_key(keys: Iterable[str], prefix: str) -> bool: return any(k.startswith(prefix) for k in keys) def _looks_like_xlabs_flux_key(k: str) -> bool: """XLabs FLUX-style keys under double_blocks/single_blocks with lora down/up.""" if not (k.endswith(".down.weight") or k.endswith(".up.weight")): return False if not k.startswith( ( "double_blocks.", "single_blocks.", "diffusion_model.double_blocks", "diffusion_model.single_blocks", ) ): return False return ".processor." in k or ".proj_lora" in k or ".qkv_lora" in k def _looks_like_kohya_flux(state_dict: Mapping[str, torch.Tensor]) -> bool: """Kohya FLUX LoRA (flux_lora.py) under lora_unet_double/single_blocks_ prefixes.""" if not state_dict: return False keys = state_dict.keys() return any( k.startswith("lora_unet_double_blocks_") or k.startswith("lora_unet_single_blocks_") for k in keys ) def _looks_like_non_diffusers_sd(state_dict: Mapping[str, torch.Tensor]) -> bool: """Classic non-diffusers SD LoRA (Kohya/A1111/sd-scripts).""" if not state_dict: return False keys = state_dict.keys() return all( k.startswith(("lora_unet_", "lora_te_", "lora_te1_", "lora_te2_")) for k in keys ) def _looks_like_wan_lora(state_dict: Mapping[str, torch.Tensor]) -> bool: """Wan2.2 distill LoRAs (Wan-AI / Wan2.2-Distill-Loras style).""" if not state_dict: return False for k in state_dict.keys(): if not k.startswith("diffusion_model.blocks."): continue if ".lora_down" not in k and ".lora_up" not in k: continue if ".cross_attn." in k or ".self_attn." in k or ".ffn." in k or ".norm3." in k: return True return False def _looks_like_qwen_image(state_dict: Mapping[str, torch.Tensor]) -> bool: keys = list(state_dict.keys()) if not keys: return False return _has_prefix_key(keys, "transformer.transformer_blocks.") and ( _has_substring_key(keys, ".lora.down.weight") or _has_substring_key(keys, ".lora.up.weight") ) def _looks_like_ai_toolkit_flux_lora(state_dict: Mapping[str, torch.Tensor]) -> bool: """Detect ai-toolkit/ComfyUI trained Flux LoRA with double_blocks/single_blocks naming. Key patterns: double_blocks.{N}.img_attn.proj.lora_A.weight """ keys = list(state_dict.keys()) if not keys: return False has_double_blocks = any( k.startswith("double_blocks.") or k.startswith("base_model.model.double_blocks.") for k in keys ) has_single_blocks = any( k.startswith("single_blocks.") or k.startswith("base_model.model.single_blocks.") for k in keys ) has_lora_ab = _has_substring_key(keys, ".lora_A") or _has_substring_key( keys, ".lora_B" ) return (has_double_blocks or has_single_blocks) and has_lora_ab def detect_lora_format_from_state_dict( state_dict: Mapping[str, torch.Tensor], ) -> LoRAFormat: """Classify LoRA format by key patterns only.""" keys = list(state_dict.keys()) if not keys: return LoRAFormat.STANDARD if _looks_like_ai_toolkit_flux_lora(state_dict): return LoRAFormat.AI_TOOLKIT_FLUX if _has_substring_key(keys, ".lora_A") or _has_substring_key(keys, ".lora_B"): return LoRAFormat.STANDARD if any(_looks_like_xlabs_flux_key(k) for k in keys): return LoRAFormat.XLABS_FLUX if _looks_like_kohya_flux(state_dict): return LoRAFormat.KOHYA_FLUX if _looks_like_wan_lora(state_dict): return LoRAFormat.WAN if _looks_like_qwen_image(state_dict): return LoRAFormat.STANDARD if _looks_like_non_diffusers_sd(state_dict): return LoRAFormat.NON_DIFFUSERS_SD if _has_substring_key(keys, ".lora.down") or _has_substring_key(keys, ".lora_up"): return LoRAFormat.NON_DIFFUSERS_SD return LoRAFormat.STANDARD def _convert_qwen_image_standard( state_dict: Mapping[str, torch.Tensor], log: logging.Logger, ) -> Dict[str, torch.Tensor]: """Qwen-Image: transformer.*.lora.down/up -> transformer_blocks.*.lora_A/B.""" out: Dict[str, torch.Tensor] = {} for name, tensor in state_dict.items(): new_name = name if new_name.startswith("transformer."): new_name = new_name[len("transformer.") :] if new_name.endswith(".lora.down.weight"): new_name = new_name.replace(".lora.down.weight", ".lora_A.weight") elif new_name.endswith(".lora.up.weight"): new_name = new_name.replace(".lora.up.weight", ".lora_B.weight") out[new_name] = tensor return out def _convert_non_diffusers_sd_simple( state_dict: Mapping[str, torch.Tensor], log: logging.Logger, ) -> Dict[str, torch.Tensor]: """Generic down/up -> A/B conversion for non-diffusers SD-like formats.""" out: Dict[str, torch.Tensor] = {} for name, tensor in state_dict.items(): new_name = name if "lora_down.weight" in new_name: new_name = new_name.replace("lora_down.weight", "lora_A.weight") elif "lora_up.weight" in new_name: new_name = new_name.replace("lora_up.weight", "lora_B.weight") elif new_name.endswith(".lora_down"): new_name = new_name.replace(".lora_down", ".lora_A") elif new_name.endswith(".lora_up"): new_name = new_name.replace(".lora_up", ".lora_B") out[new_name] = tensor sample = _sample_keys(out.keys(), 20) log.info( "[LoRAFormatAdapter] after NON_DIFFUSERS_SD simple conversion, " "sample keys (<=20): %s", ", ".join(sample), ) return out def _convert_with_diffusers_utils_if_available( state_dict: Mapping[str, torch.Tensor], log: logging.Logger, ) -> Optional[Dict[str, torch.Tensor]]: """Use diffusers.lora_conversion_utils if available.""" try: if hasattr(lcu, "maybe_convert_state_dict"): converted = lcu.maybe_convert_state_dict( # type: ignore[attr-defined] state_dict ) else: converted = dict(state_dict) if not isinstance(converted, dict): converted = dict(converted) sample = _sample_keys(converted.keys(), 20) log.info( "[LoRAFormatAdapter] diffusers.lora_conversion_utils converted keys, " "sample keys (<=20): %s", ", ".join(sample), ) return converted except Exception as exc: # pragma: no cover log.warning( "[LoRAFormatAdapter] diffusers lora_conversion_utils failed, " "falling back to internal converters. Error: %s", exc, ) return None def _convert_via_diffusers_candidates( state_dict: Mapping[str, torch.Tensor], candidate_names: tuple[str, ...], log: logging.Logger, unavailable_warning: str, no_converter_warning: str, success_info: str, all_failed_warning: str, ) -> Dict[str, torch.Tensor]: """Try multiple named converters in lora_conversion_utils, use the first that works.""" converters = [ (n, getattr(lcu, n)) for n in candidate_names if callable(getattr(lcu, n, None)) ] if not converters: log.warning(no_converter_warning) return dict(state_dict) last_err: Optional[Exception] = None for name, fn in converters: try: sd_copy = dict(state_dict) out = fn(sd_copy) if isinstance(out, tuple) and isinstance(out[0], dict): out = out[0] if not isinstance(out, dict): raise TypeError(f"Converter {name} returned {type(out)}") log.info(success_info.format(name=name)) return out except Exception as exc: last_err = exc log.warning(all_failed_warning.format(last_err=last_err)) return dict(state_dict) def _convert_xlabs_ai_via_diffusers( state_dict: Mapping[str, torch.Tensor], log: logging.Logger, ) -> Dict[str, torch.Tensor]: """Convert XLabs FLUX LoRA via diffusers helpers.""" return _convert_via_diffusers_candidates( state_dict, ( "_convert_xlabs_flux_lora_to_diffusers", "convert_xlabs_lora_state_dict_to_diffusers", "convert_xlabs_lora_to_diffusers", "convert_xlabs_flux_lora_to_diffusers", ), log=log, unavailable_warning=( "[LoRAFormatAdapter] XLabs FLUX detected but diffusers is unavailable." ), no_converter_warning=( "[LoRAFormatAdapter] No XLabs FLUX converter found in diffusers." ), success_info="[LoRAFormatAdapter] Converted XLabs FLUX LoRA using {name}", all_failed_warning=( "[LoRAFormatAdapter] All XLabs FLUX converters failed; " "last error: {last_err}" ), ) def _convert_kohya_flux_via_diffusers( state_dict: Mapping[str, torch.Tensor], log: logging.Logger, ) -> Dict[str, torch.Tensor]: """Convert Kohya FLUX LoRA via diffusers helpers.""" return _convert_via_diffusers_candidates( state_dict, ( "_convert_kohya_flux_lora_to_diffusers", "convert_kohya_flux_lora_to_diffusers", ), log=log, unavailable_warning=( "[LoRAFormatAdapter] Kohya FLUX detected but diffusers is unavailable." ), no_converter_warning="[LoRAFormatAdapter] No Kohya FLUX converter found.", success_info="[LoRAFormatAdapter] Converted Kohya FLUX LoRA using {name}", all_failed_warning=( "[LoRAFormatAdapter] Kohya FLUX conversion failed; " "last error: {last_err}" ), ) def _convert_ai_toolkit_flux_lora( state_dict: Mapping[str, torch.Tensor], log: logging.Logger, ) -> Dict[str, torch.Tensor]: """Convert ai-toolkit/ComfyUI trained Flux LoRA to SGLang format. Handles the naming convention conversion: - double_blocks.{N}.img_attn.qkv -> transformer_blocks.{N}.attn.to_q/k/v - double_blocks.{N}.txt_attn.qkv -> transformer_blocks.{N}.attn.add_q/k/v_proj - double_blocks.{N}.img_attn.proj -> transformer_blocks.{N}.attn.to_out.0 - double_blocks.{N}.txt_attn.proj -> transformer_blocks.{N}.attn.to_add_out - double_blocks -> transformer_blocks - single_blocks -> single_transformer_blocks """ out: Dict[str, torch.Tensor] = {} original_state_dict: Dict[str, torch.Tensor] = {} for name, tensor in state_dict.items(): new_name = name if new_name.startswith("diffusion_model."): new_name = new_name[len("diffusion_model.") :] if new_name.startswith("base_model.model."): new_name = new_name[len("base_model.model.") :] original_state_dict[new_name] = tensor num_double_layers = 0 num_single_layers = 0 for key in original_state_dict.keys(): if key.startswith("single_blocks."): parts = key.split(".") if len(parts) > 1 and parts[1].isdigit(): num_single_layers = max(num_single_layers, int(parts[1]) + 1) elif key.startswith("double_blocks."): parts = key.split(".") if len(parts) > 1 and parts[1].isdigit(): num_double_layers = max(num_double_layers, int(parts[1]) + 1) lora_keys = ("lora_A", "lora_B") attn_types = ("img_attn", "txt_attn") for sl in range(num_single_layers): single_block_prefix = f"single_blocks.{sl}" attn_prefix = f"single_transformer_blocks.{sl}.attn" for lora_key in lora_keys: linear1_key = f"{single_block_prefix}.linear1.{lora_key}.weight" if linear1_key in original_state_dict: out[f"{attn_prefix}.to_qkv_mlp_proj.{lora_key}.weight"] = ( original_state_dict.pop(linear1_key) ) linear2_key = f"{single_block_prefix}.linear2.{lora_key}.weight" if linear2_key in original_state_dict: out[f"{attn_prefix}.to_out.{lora_key}.weight"] = ( original_state_dict.pop(linear2_key) ) for dl in range(num_double_layers): transformer_block_prefix = f"transformer_blocks.{dl}" for lora_key in lora_keys: for attn_type in attn_types: attn_prefix = f"{transformer_block_prefix}.attn" qkv_key = f"double_blocks.{dl}.{attn_type}.qkv.{lora_key}.weight" if qkv_key not in original_state_dict: continue fused_qkv_weight = original_state_dict.pop(qkv_key) if lora_key == "lora_A": diff_attn_proj_keys = ( ["to_q", "to_k", "to_v"] if attn_type == "img_attn" else ["add_q_proj", "add_k_proj", "add_v_proj"] ) for proj_key in diff_attn_proj_keys: out[f"{attn_prefix}.{proj_key}.{lora_key}.weight"] = ( fused_qkv_weight ) else: if fused_qkv_weight.shape[0] % 3 != 0: log.warning( "[LoRAFormatAdapter] QKV weight shape %s not divisible by 3, " "may cause shape mismatch for %s", fused_qkv_weight.shape, qkv_key, ) sample_q, sample_k, sample_v = torch.chunk( fused_qkv_weight, 3, dim=0 ) if attn_type == "img_attn": out[f"{attn_prefix}.to_q.{lora_key}.weight"] = sample_q out[f"{attn_prefix}.to_k.{lora_key}.weight"] = sample_k out[f"{attn_prefix}.to_v.{lora_key}.weight"] = sample_v else: out[f"{attn_prefix}.add_q_proj.{lora_key}.weight"] = sample_q out[f"{attn_prefix}.add_k_proj.{lora_key}.weight"] = sample_k out[f"{attn_prefix}.add_v_proj.{lora_key}.weight"] = sample_v proj_mappings = [ ("img_attn.proj", "attn.to_out.0"), ("txt_attn.proj", "attn.to_add_out"), ] for org_proj, diff_proj in proj_mappings: for lora_key in lora_keys: original_key = f"double_blocks.{dl}.{org_proj}.{lora_key}.weight" if original_key in original_state_dict: diffusers_key = ( f"{transformer_block_prefix}.{diff_proj}.{lora_key}.weight" ) out[diffusers_key] = original_state_dict.pop(original_key) for key, tensor in original_state_dict.items(): new_key = key.replace("double_blocks.", "transformer_blocks.") new_key = new_key.replace("single_blocks.", "single_transformer_blocks.") out[new_key] = tensor extra_mappings = { "img_in": "x_embedder", "txt_in": "context_embedder", "time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1", "time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2", "final_layer.linear": "proj_out", "final_layer.adaLN_modulation.1": "norm_out.linear", "single_stream_modulation.lin": "single_stream_modulation.linear", "double_stream_modulation_img.lin": "double_stream_modulation_img.linear", "double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear", } final_out: Dict[str, torch.Tensor] = {} for key, tensor in out.items(): new_key = key for org_key, diff_key in extra_mappings.items(): if key.startswith(org_key): new_key = key.replace(org_key, diff_key, 1) break final_out[new_key] = tensor sample = _sample_keys(final_out.keys(), 20) log.info( "[LoRAFormatAdapter] after AI_TOOLKIT_FLUX conversion, " "sample keys (<=20): %s", ", ".join(sample), ) return final_out def convert_lora_state_dict_by_format( state_dict: Mapping[str, torch.Tensor], fmt: LoRAFormat, log: logging.Logger, ) -> Dict[str, torch.Tensor]: """Normalize a raw LoRA state_dict into A/B + .weight naming.""" if fmt == LoRAFormat.QWEN_IMAGE_STANDARD: return _convert_qwen_image_standard(state_dict, log) if fmt == LoRAFormat.AI_TOOLKIT_FLUX: return _convert_ai_toolkit_flux_lora(state_dict, log) if fmt == LoRAFormat.XLABS_FLUX: converted = _convert_xlabs_ai_via_diffusers(state_dict, log) return _convert_non_diffusers_sd_simple(converted, log) if fmt == LoRAFormat.KOHYA_FLUX: converted = _convert_kohya_flux_via_diffusers(state_dict, log) return _convert_non_diffusers_sd_simple(converted, log) if fmt == LoRAFormat.WAN: maybe = _convert_with_diffusers_utils_if_available(state_dict, log) if maybe is None: maybe = dict(state_dict) return _convert_non_diffusers_sd_simple(maybe, log) if fmt == LoRAFormat.STANDARD: maybe = _convert_with_diffusers_utils_if_available(state_dict, log) if maybe is None: maybe = dict(state_dict) if _looks_like_qwen_image(maybe): return _convert_qwen_image_standard(maybe, log) return maybe if fmt == LoRAFormat.NON_DIFFUSERS_SD: maybe = _convert_with_diffusers_utils_if_available(state_dict, log) if maybe is None: maybe = dict(state_dict) return _convert_non_diffusers_sd_simple(maybe, log) log.info( "[LoRAFormatAdapter] format %s not handled specially, returning as-is", fmt, ) return dict(state_dict) def normalize_lora_state_dict( state_dict: Mapping[str, torch.Tensor], logger: Optional[logging.Logger] = None, ) -> Dict[str, torch.Tensor]: """Normalize any supported LoRA format into a single canonical layout.""" log = logger or globals()["logger"] keys = list(state_dict.keys()) log.info( "[LoRAFormatAdapter] normalize_lora_state_dict called, #keys=%d", len(keys), ) if keys: log.info( "[LoRAFormatAdapter] before convert, sample keys (<=20): %s", ", ".join(_sample_keys(keys, 20)), ) fmt = detect_lora_format_from_state_dict(state_dict) log.info("[LoRAFormatAdapter] detected format: %s", fmt) normalized = convert_lora_state_dict_by_format(state_dict, fmt, log) norm_keys = list(normalized.keys()) if norm_keys: log.info( "[LoRAFormatAdapter] after convert, sample keys (<=20): %s", ", ".join(_sample_keys(norm_keys, 20)), ) return normalized ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import os from collections import defaultdict from collections.abc import Hashable from contextlib import contextmanager from typing import Any import torch import torch.distributed as dist from safetensors.torch import load_file from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.layers.lora.linear import ( BaseLayerWithLoRA, replace_submodule, wrap_with_lora_layer, ) from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.lora_format_adapter import ( normalize_lora_state_dict, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_lora from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger # to avoid deadlocks when forking os.environ["TOKENIZERS_PARALLELISM"] = "false" logger = init_logger(__name__) class LoRAPipeline(ComposedPipelineBase): """ Pipeline that supports injecting LoRA adapters into the diffusion transformer. """ # Type annotations for instance attributes (initialized in __init__) # [lora_nickname][target_LoRA_weight_name_in_SGLang_dit] = weight # e.g., [jinx][transformer_blocks.0.attn.to_v.lora_A] lora_adapters: dict[str, dict[str, torch.Tensor]] loaded_adapter_paths: dict[str, str] # nickname -> lora_path # Track current adapter per module: {"transformer": "high_lora", "transformer_2": "low_lora"} cur_adapter_name: dict[str, str] cur_adapter_path: dict[str, str] cur_adapter_strength: dict[str, float] # Track current strength per module cur_adapter_config: dict[str, tuple[list[str], list[float]]] # [dit_layer_name] = wrapped_lora_layer lora_layers: dict[str, BaseLayerWithLoRA] lora_layers_critic: dict[str, BaseLayerWithLoRA] lora_layers_transformer_2: dict[str, BaseLayerWithLoRA] server_args: ServerArgs exclude_lora_layers: list[str] device: torch.device lora_target_modules: list[str] | None lora_path: str | None lora_nickname: str lora_rank: int | None lora_alpha: int | None lora_initialized: bool # Track merge status per module: {"transformer": True, "transformer_2": False} is_lora_merged: dict[str, bool] # Valid target values for set_lora (class constant, immutable) VALID_TARGETS: list[str] = ["all", "transformer", "transformer_2", "critic"] def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) # Initialize all mutable instance attributes to avoid sharing across instances self.lora_adapters = defaultdict(dict) self.loaded_adapter_paths = {} self.cur_adapter_name = {} self.cur_adapter_path = {} self.cur_adapter_strength = {} # Track full LoRA config: {module_name: (nickname_list, strength_list)} self.cur_adapter_config = {} self.lora_layers = {} self.lora_layers_critic = {} self.lora_layers_transformer_2 = {} self.is_lora_merged = {} self.lora_initialized = False self.lora_rank = None self.lora_alpha = None self.lora_path = None self.lora_nickname = "default" # Initialize from server_args self.device = get_local_torch_device() self.exclude_lora_layers = ( self.server_args.pipeline_config.dit_config.arch_config.exclude_lora_layers ) self.lora_target_modules = self.server_args.lora_target_modules self.lora_path = self.server_args.lora_path self.lora_nickname = self.server_args.lora_nickname if self.lora_path is not None: self.convert_to_lora_layers() self.set_lora( self.lora_nickname, self.lora_path, strength=self.server_args.lora_scale # type: ignore ) # type: ignore def is_target_layer(self, module_name: str) -> bool: if self.lora_target_modules is None: return True return any( target_name in module_name for target_name in self.lora_target_modules ) def _get_target_lora_layers( self, target: str ) -> tuple[list[tuple[str, dict[str, BaseLayerWithLoRA]]], str | None]: """ Return a list of (module_name, lora_layers_dict) based on the target. Args: target: One of "all", "transformer", "transformer_2", "critic". Returns: A tuple of (result, error_message): - result: List of tuples (module_name, lora_layers_dict) to operate on. - error_message: Error description if target is invalid or module doesn't exist, None otherwise. """ if target == "all": result: list[tuple[str, dict[str, BaseLayerWithLoRA]]] = [ ("transformer", self.lora_layers) ] if self.lora_layers_transformer_2: result.append(("transformer_2", self.lora_layers_transformer_2)) if self.lora_layers_critic: result.append(("critic", self.lora_layers_critic)) return result, None elif target == "transformer": return [("transformer", self.lora_layers)], None elif target == "transformer_2": if not self.lora_layers_transformer_2: return [], "transformer_2 does not exist in this pipeline" return [("transformer_2", self.lora_layers_transformer_2)], None elif target == "critic": if not self.lora_layers_critic: return ( [], "critic (fake_score_transformer) does not exist in this pipeline", ) return [("critic", self.lora_layers_critic)], None else: return [], f"Invalid target: {target}. Valid targets: {self.VALID_TARGETS}" @contextmanager def _temporarily_disable_offload( self, target_modules: list[tuple[str, dict[str, BaseLayerWithLoRA]]] | None = None, target: str | None = None, use_module_names_only: bool = False, ): """ Context manager to temporarily disable layerwise offload for the given modules. Args: target_modules: List of (module_name, lora_layers_dict) tuples. If None, will be determined from target. target: Target string ("all", "transformer", etc.). Used if target_modules is None. use_module_names_only: If True, determine module names directly from target without requiring LoRA initialization. Used for convert_to_lora_layers scenario. Yields: List of modules that had offload disabled. """ from sglang.multimodal_gen.runtime.utils.layerwise_offload import ( OffloadableDiTMixin, ) module_names = [] if target_modules is not None: # Extract module names from target_modules module_names = [module_name for module_name, _ in target_modules] elif target is not None: if use_module_names_only: if target == "all": module_names = ["transformer", "transformer_2"] elif target in ["transformer", "transformer_2", "critic"]: module_names = [target] else: target_modules, _ = self._get_target_lora_layers(target) if target_modules: module_names = [module_name for module_name, _ in target_modules] else: yield [] return if not module_names: yield [] return # clear device cache to free up unused memory if torch.get_device_module().is_available(): torch.get_device_module().synchronize() torch.get_device_module().empty_cache() offload_disabled_modules = [] for module_name in module_names: module = self.modules.get(module_name) if module is not None and isinstance(module, OffloadableDiTMixin): if module.layerwise_offload_managers is not None: module.disable_offload() offload_disabled_modules.append(module) try: yield offload_disabled_modules finally: # Re-enable layerwise offload: sync weights to CPU and restore hooks for module in offload_disabled_modules: module.enable_offload() def convert_module_lora_layers( self, module: torch.nn.Module, module_name: str, target_lora_layers: dict[str, BaseLayerWithLoRA], check_exclude: bool = True, ) -> int: """ Convert layers in a module to LoRA layers. Args: module: The module to convert. module_name: The name of the module (for replace_submodule). target_lora_layers: The dictionary to store the converted LoRA layers. check_exclude: Whether to check the exclude_lora_layers list. Returns: The number of layers converted. """ converted_count = 0 for name, layer in module.named_modules(): if not self.is_target_layer(name): continue if check_exclude: excluded = any( exclude_layer in name for exclude_layer in self.exclude_lora_layers ) if excluded: continue lora_layer = wrap_with_lora_layer( layer, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, ) if lora_layer is not None: target_lora_layers[name] = lora_layer replace_submodule(self.modules[module_name], name, lora_layer) converted_count += 1 return converted_count def convert_to_lora_layers(self) -> None: """ Unified method to convert the transformer to a LoRA transformer. """ if self.lora_initialized: return self.lora_initialized = True # Convert transformer converted_count = self.convert_module_lora_layers( self.modules["transformer"], "transformer", self.lora_layers, check_exclude=True, ) logger.info("Converted %d layers to LoRA layers", converted_count) # Convert transformer_2 if exists (e.g., Wan2.2 A14B dual-transformer) if ( "transformer_2" in self.modules and self.modules["transformer_2"] is not None ): converted_count_2 = self.convert_module_lora_layers( self.modules["transformer_2"], "transformer_2", self.lora_layers_transformer_2, check_exclude=True, ) logger.info( "Converted %d layers to LoRA layers in transformer_2", converted_count_2 ) # Convert fake_score_transformer if exists if "fake_score_transformer" in self.modules: converted_count_critic = self.convert_module_lora_layers( self.modules["fake_score_transformer"], "fake_score_transformer", self.lora_layers_critic, check_exclude=False, ) logger.info( "Converted %d layers to LoRA layers in the critic model", converted_count_critic, ) def _normalize_lora_params( self, lora_nickname: str | list[str], lora_path: str | None | list[str | None], strength: float | list[float], target: str | list[str], ) -> tuple[list[str], list[str | None], list[float], list[str]]: """ Normalize LoRA parameters to lists for multi-LoRA support. Requirements: - each nickname must have a corresponding lora_path (no implicit repeat) - strength / target if scalar broadcast, else length must match nickname """ # nickname if isinstance(lora_nickname, str): lora_nicknames = [lora_nickname] else: lora_nicknames = lora_nickname # lora_path: require 1:1 mapping with nickname (no implicit repeat) if isinstance(lora_path, list): lora_paths = lora_path else: lora_paths = [lora_path] if len(lora_paths) != len(lora_nicknames): raise ValueError( f"Length mismatch: lora_nickname has {len(lora_nicknames)} items, " f"but lora_path has {len(lora_paths)} items. " "Provide one path per nickname." ) # strength and target: allow scalar broadcast, else length must match if isinstance(strength, (int, float)): strengths = [float(strength)] * len(lora_nicknames) else: strengths = [float(s) for s in strength] if len(strengths) != len(lora_nicknames): raise ValueError( f"Length mismatch: lora_nickname has {len(lora_nicknames)} items, " f"but strength has {len(strengths)} items" ) if isinstance(target, str): targets = [target] * len(lora_nicknames) else: targets = target if len(targets) != len(lora_nicknames): raise ValueError( f"Length mismatch: lora_nickname has {len(lora_nicknames)} items, " f"but target has {len(targets)} items" ) return lora_nicknames, lora_paths, strengths, targets def _check_lora_config_matches( self, module_name: str, target_nicknames: list[str], target_strengths: list[float], adapter_updated: bool, ) -> bool: """ Check if current LoRA configuration matches the target configuration. Args: module_name: The name of the module to check. target_nicknames: List of LoRA nicknames to apply. target_strengths: List of LoRA strengths to apply. adapter_updated: Whether any adapter was updated/loaded. Returns: True if the configuration matches exactly (including order and strength), False otherwise. """ if not self.is_lora_merged.get(module_name, False): return False if adapter_updated: return False # Adapter was updated, need to reapply stored_config = self.cur_adapter_config.get(module_name) if stored_config is None: return False stored_nicknames, stored_strengths = stored_config # Compare: nickname list and strength list must match exactly (including order) return ( stored_nicknames == target_nicknames and stored_strengths == target_strengths ) def _apply_lora_to_layers( self, lora_layers: dict[str, BaseLayerWithLoRA], lora_nicknames: list[str], lora_paths: list[str | None], rank: int, strengths: list[float], clear_existing: bool = False, ) -> int: """ Apply LoRA weights to the given lora_layers. Supports multiple LoRA adapters. Args: lora_layers: The dictionary of LoRA layers to apply weights to. lora_nicknames: The list of nicknames of the LoRA adapters. lora_paths: The list of paths to the LoRA adapters. Must match length of lora_nicknames. rank: The distributed rank (for logging). strengths: The list of LoRA strengths for merge. Must match length of lora_nicknames. clear_existing: If True, clear existing LoRA weights before adding new ones. Returns: The number of layers that had LoRA weights applied. """ if len(lora_paths) != len(lora_nicknames): raise ValueError( f"Length mismatch: lora_nicknames has {len(lora_nicknames)} items, " f"but lora_paths has {len(lora_paths)} items" ) if len(strengths) != len(lora_nicknames): raise ValueError( f"Length mismatch: lora_nicknames has {len(lora_nicknames)} items, " f"but strengths has {len(strengths)} items" ) adapted_count = 0 for name, layer in lora_layers.items(): # Apply all LoRA adapters in order for idx, (nickname, path, lora_strength) in enumerate( zip(lora_nicknames, lora_paths, strengths) ): lora_A_name = name + ".lora_A" lora_B_name = name + ".lora_B" if ( lora_A_name in self.lora_adapters[nickname] and lora_B_name in self.lora_adapters[nickname] ): # Some LoRA checkpoints (e.g. Lightning distill) store per-layer alpha as ".alpha". # If present, we must apply the standard LoRA scaling: scale = alpha / rank. try: inferred_rank = int( self.lora_adapters[nickname][lora_A_name].shape[0] ) except Exception: inferred_rank = None # Default to None for some checkpoints without ".alpha" inferred_alpha: int | None = None alpha_key = name + ".alpha" if alpha_key in self.lora_adapters[nickname]: try: inferred_alpha = int( self.lora_adapters[nickname][alpha_key].item() ) except Exception: inferred_alpha = None if inferred_rank is not None: layer.lora_rank = inferred_rank layer.lora_alpha = ( inferred_alpha if inferred_alpha is not None else inferred_rank ) layer.set_lora_weights( self.lora_adapters[nickname][lora_A_name], self.lora_adapters[nickname][lora_B_name], lora_path=path, strength=lora_strength, clear_existing=( clear_existing and idx == 0 ), # Only clear on first LoRA ) adapted_count += 1 else: if rank == 0 and idx == 0: # Only warn for first missing LoRA logger.warning( "LoRA adapter %s does not contain the weights for layer '%s'. LoRA will not be applied to it.", path, name, ) # Only disable if no LoRA was applied at all if idx == len(lora_nicknames) - 1: has_any_lora = any( name + ".lora_A" in self.lora_adapters[n] and name + ".lora_B" in self.lora_adapters[n] for n in lora_nicknames ) if not has_any_lora: layer.disable_lora = True return adapted_count def is_lora_effective(self, target: str = "all") -> bool: """ Check if LoRA is currently effective (merged) for the specified target. Args: target: Which transformer to check. "all" returns True if any is merged. """ if target == "all": return any(self.is_lora_merged.values()) return self.is_lora_merged.get(target, False) def is_lora_set(self, target: str = "all") -> bool: """ Check if LoRA has been set for the specified target. Args: target: Which transformer to check. "all" returns True if any is set. """ if not self.lora_initialized: return False if target == "all": return bool(self.cur_adapter_name) return target in self.cur_adapter_name def load_lora_adapter(self, lora_path: str, lora_nickname: str, rank: int): """ Load the LoRA, and setup the lora_adapters for later weight replacement """ assert lora_path is not None # Only rank 0 downloads to avoid race conditions where other ranks # try to load incomplete downloads if rank == 0: lora_local_path = maybe_download_lora(lora_path) else: lora_local_path = None # Synchronize all ranks after download completes if dist.is_initialized(): dist.barrier() # Non-rank-0 workers now download (will hit cache since rank 0 completed) if rank != 0: lora_local_path = maybe_download_lora(lora_path) raw_state_dict = load_file(lora_local_path) lora_state_dict = normalize_lora_state_dict(raw_state_dict, logger=logger) if lora_nickname in self.lora_adapters: self.lora_adapters[lora_nickname].clear() config = self.server_args.pipeline_config.dit_config.arch_config param_names_mapping_fn = get_param_names_mapping( config.param_names_mapping or self.modules["transformer"].param_names_mapping ) lora_param_names_mapping_fn = get_param_names_mapping( config.lora_param_names_mapping or self.modules["transformer"].lora_param_names_mapping ) to_merge_params: defaultdict[Hashable, dict[Any, Any]] = defaultdict(dict) for name, weight in lora_state_dict.items(): name = name.replace("diffusion_model.", "") name = name.replace(".weight", "") # misc-format -> HF-format name, _, _ = lora_param_names_mapping_fn(name) # HF-format (LoRA) -> SGLang-dit-format target_name, merge_index, num_params_to_merge = param_names_mapping_fn(name) # for fuse B(out_dim, r) @ A(r, in_dim) -> (N, out_dim, r) @ (N, r, in_dim) # see param mapping in HunyuanVideoArchConfig if merge_index is not None: to_merge_params[target_name][merge_index] = weight if len(to_merge_params[target_name]) == num_params_to_merge: sorted_tensors = [ to_merge_params[target_name][i] for i in range(num_params_to_merge) ] # Use stack instead of cat because it needs to be compatible with TP. weight = torch.stack(sorted_tensors, dim=0) del to_merge_params[target_name] else: continue if target_name in self.lora_adapters[lora_nickname]: raise ValueError( f"Dit target weight name {target_name} already exists in lora_adapters[{lora_nickname}]" ) self.lora_adapters[lora_nickname][target_name] = weight.to(self.device) self.loaded_adapter_paths[lora_nickname] = lora_path logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path) def set_lora( self, lora_nickname: str | list[str], lora_path: str | None | list[str | None] = None, target: str | list[str] = "all", strength: float | list[float] = 1.0, ): # type: ignore """ Load LoRA adapter(s) into the pipeline and apply them to the specified transformer(s). Supports both single LoRA (backward compatible) and multiple LoRA adapters. """ # Normalize inputs to lists for multi-LoRA support lora_nicknames, lora_paths, strengths, targets = self._normalize_lora_params( lora_nickname, lora_path, strength, target ) # Validate targets invalid_targets = [t for t in targets if t not in self.VALID_TARGETS] if invalid_targets: raise ValueError( f"Invalid target(s): {invalid_targets}. Valid targets: {self.VALID_TARGETS}" ) # Disable layerwise offload before convert_to_lora_layers to ensure weights are accessible # This is critical because convert_to_lora_layers needs to save cpu_weight from actual weights, # not from offloaded placeholder tensors if not self.lora_initialized: with self._temporarily_disable_offload( target="all", use_module_names_only=True ): self.convert_to_lora_layers() # Check adapter presence and load missing adapters adapter_updated = False rank = dist.get_rank() # load required adapters for nickname, path in zip(lora_nicknames, lora_paths): if nickname not in self.lora_adapters and path is None: raise ValueError( f"Adapter {nickname} not found in the pipeline. Please provide lora_path to load it." ) # Check if adapter needs to be loaded should_load = False if path is not None: if nickname not in self.loaded_adapter_paths: should_load = True elif self.loaded_adapter_paths[nickname] != path: should_load = True if should_load: adapter_updated = True self.load_lora_adapter(path, nickname, rank) # Group by target to apply separately target_to_indices = {} for idx, tgt in enumerate(targets): if tgt not in target_to_indices: target_to_indices[tgt] = [] target_to_indices[tgt].append(idx) adapted_count = 0 for tgt, idx_list in target_to_indices.items(): target_modules, error = self._get_target_lora_layers(tgt) if error: logger.warning("set_lora: %s", error) if not target_modules: continue # Disable layerwise offload if enabled: load all layers to GPU # the LoRA weights merging process requires weights being on device with self._temporarily_disable_offload(target_modules=target_modules): tgt_nicknames = [lora_nicknames[i] for i in idx_list] tgt_paths = [lora_paths[i] for i in idx_list] tgt_strengths = [strengths[i] for i in idx_list] merged_name = ( ",".join(tgt_nicknames) if len(tgt_nicknames) > 1 else tgt_nicknames[0] ) # Skip if LoRA configuration matches exactly (including order and strength) # Since all modules for the same target apply the same config, checking one is sufficient first_module_name, _ = target_modules[0] if self._check_lora_config_matches( first_module_name, tgt_nicknames, tgt_strengths, adapter_updated ): logger.info("LoRA configuration matches exactly, skipping") continue # Apply LoRA to modules for this target for module_name, lora_layers_dict in target_modules: count = self._apply_lora_to_layers( lora_layers_dict, tgt_nicknames, tgt_paths, rank, tgt_strengths, clear_existing=True, ) adapted_count += count self.cur_adapter_name[module_name] = merged_name self.cur_adapter_path[module_name] = ",".join( str(p or self.loaded_adapter_paths.get(n, "")) for n, p in zip(tgt_nicknames, tgt_paths) ) self.is_lora_merged[module_name] = True self.cur_adapter_strength[module_name] = tgt_strengths[0] # Store full configuration for multi-LoRA support (preserves order and all strengths) self.cur_adapter_config[module_name] = ( tgt_nicknames.copy(), tgt_strengths.copy(), ) logger.info( "Rank %d: LoRA adapter(s) %s applied to %d layers (targets: %s, strengths: %s)", rank, ", ".join(map(str, lora_paths)) if lora_paths else None, adapted_count, ", ".join(targets) if len(set(targets)) > 1 else targets[0], ( ", ".join(f"{s:.2f}" for s in strengths) if len(strengths) > 1 else f"{strengths[0]:.2f}" ), ) def merge_lora_weights(self, target: str = "all", strength: float = 1.0) -> None: """ Merge LoRA weights into the base model for the specified target. This operation is idempotent - calling it when LoRA is already merged is safe. Args: target: Which transformer(s) to merge. One of "all", "transformer", "transformer_2", "critic". strength: LoRA strength for merge, default 1.0. """ target_modules, error = self._get_target_lora_layers(target) if error: logger.warning("merge_lora_weights: %s", error) if not target_modules: return # Disable layerwise offload if enabled: load all layers to GPU with self._temporarily_disable_offload(target_modules=target_modules): for module_name, lora_layers_dict in target_modules: if self.is_lora_merged.get(module_name, False): # Check if strength is the same - if so, skip (idempotent) if self.cur_adapter_strength.get(module_name) == strength: logger.warning( "LoRA weights are already merged for %s with same strength", module_name, ) continue # Different strength requested - allow re-merge (layer handles unmerge internally) logger.info( "Re-merging LoRA weights for %s with new strength %s", module_name, strength, ) for name, layer in lora_layers_dict.items(): # Only re-enable LoRA for layers that actually have LoRA weights has_lora_weights = ( hasattr(layer, "lora_A") and layer.lora_A is not None ) if not has_lora_weights: continue if hasattr(layer, "disable_lora"): layer.disable_lora = False try: layer.merge_lora_weights(strength=strength) except Exception as e: logger.warning("Could not merge layer %s: %s", name, e) continue self.is_lora_merged[module_name] = True self.cur_adapter_strength[module_name] = strength logger.info( "LoRA weights merged for %s (strength: %s)", module_name, strength ) def unmerge_lora_weights(self, target: str = "all") -> None: """ Unmerge LoRA weights from the base model for the specified target. This also disables LoRA so it won't be computed on-the-fly. This operation is idempotent - calling it when LoRA is not merged is safe. Args: target: Which transformer(s) to unmerge. One of "all", "transformer", "transformer_2", "critic". """ target_modules, error = self._get_target_lora_layers(target) if error: logger.warning("unmerge_lora_weights: %s", error) if not target_modules: return # Disable layerwise offload if enabled: load all layers to GPU for module_name, lora_layers_dict in target_modules: if not self.is_lora_merged.get(module_name, False): logger.warning( "LoRA weights are not merged for %s, skipping", module_name ) continue with self._temporarily_disable_offload(target_modules=target_modules): for name, layer in lora_layers_dict.items(): # Check layer-level state to avoid raising exception if hasattr(layer, "merged") and not layer.merged: logger.warning("Layer %s is not merged, skipping", name) # Still disable LoRA to prevent on-the-fly computation if hasattr(layer, "disable_lora"): layer.disable_lora = True continue try: layer.unmerge_lora_weights() # Disable LoRA after unmerge to prevent on-the-fly computation if hasattr(layer, "disable_lora"): layer.disable_lora = True except ValueError as e: logger.warning("Could not unmerge layer %s: %s", name, e) # Still disable LoRA even if unmerge failed if hasattr(layer, "disable_lora"): layer.disable_lora = True continue self.is_lora_merged[module_name] = False self.cur_adapter_strength.pop(module_name, None) self.cur_adapter_config.pop(module_name, None) logger.info("LoRA weights unmerged for %s", module_name) def get_lora_status(self) -> dict[str, Any]: """ Summarize loaded LoRA adapters and current application status per module. Returns a plain Python dict with no tensor values to allow safe JSON serialization. """ # Loaded adapters: list of {nickname, path} loaded_adapters = [ {"nickname": nickname, "path": path} for nickname, path in self.loaded_adapter_paths.items() ] def _module_status(module_name: str) -> list[dict] | None: # return list of dict to support multi-lora in the future if not self.is_lora_merged.get(module_name, False): return None else: return [ { "nickname": self.cur_adapter_name.get(module_name, None), "path": self.cur_adapter_path.get(module_name, None), "merged": self.is_lora_merged.get(module_name, False), "strength": self.cur_adapter_strength.get(module_name, None), } ] # Build active usage per module only for modules that exist in this pipeline active: dict[str, Any] = {} if ( "transformer" in self.modules and self.modules["transformer"] is not None and (status := _module_status("transformer")) is not None ): active["transformer"] = status if ( "transformer_2" in self.modules and self.modules["transformer_2"] is not None and (status := _module_status("transformer_2")) is not None ): active["transformer_2"] = status if ( "fake_score_transformer" in self.modules and self.modules["fake_score_transformer"] is not None and (status := _module_status("critic")) is not None ): active["critic"] = status return { "loaded_adapters": loaded_adapters, "active": active, } ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/model_executor/forward_batch_info.py """ Data structures for functional pipeline processing. This module defines the dataclasses used to pass state between pipeline components in a functional manner, reducing the need for explicit parameter passing. """ from __future__ import annotations import logging import os import pprint from copy import deepcopy from dataclasses import MISSING, asdict, dataclass, field, fields from typing import Any, Optional import PIL.Image import torch from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import ( _sanitize_for_logging, init_logger, ) from sglang.multimodal_gen.runtime.utils.perf_logger import RequestMetrics from sglang.multimodal_gen.utils import align_to logger = init_logger(__name__) SAMPLING_PARAMS_FIELDS = {f.name for f in fields(SamplingParams)} @dataclass(init=False) class Req: """ Complete state passed through the pipeline execution. This dataclass contains all information needed during the diffusion pipeline execution, allowing methods to update specific components without needing to manage numerous individual parameters. [IMPORTANT] Fields that overlap with SamplingParams are automatically delegated to the sampling_params member via __getattr__ and __setattr__. """ sampling_params: SamplingParams | None = None generator: torch.Generator | list[torch.Generator] | None = None # Image encoder hidden states image_embeds: list[torch.Tensor] = field(default_factory=list) original_condition_image_size: tuple[int, int] = None condition_image: torch.Tensor | PIL.Image.Image | None = None vae_image: torch.Tensor | PIL.Image.Image | None = None pixel_values: torch.Tensor | PIL.Image.Image | None = None preprocessed_image: torch.Tensor | None = None output_file_ext: str | None = None # Primary encoder embeddings prompt_embeds: list[torch.Tensor] | torch.Tensor = field(default_factory=list) negative_prompt_embeds: list[torch.Tensor] | None = None prompt_attention_mask: list[torch.Tensor] | None = None negative_attention_mask: list[torch.Tensor] | None = None clip_embedding_pos: list[torch.Tensor] | None = None clip_embedding_neg: list[torch.Tensor] | None = None pooled_embeds: list[torch.Tensor] = field(default_factory=list) neg_pooled_embeds: list[torch.Tensor] = field(default_factory=list) # Additional text-related parameters max_sequence_length: int | None = None prompt_template: dict[str, Any] | None = None do_classifier_free_guidance: bool = False seeds: list[int] | None = None # Tracking if embeddings are already processed is_prompt_processed: bool = False # Audio Embeddings (LTX-2) audio_prompt_embeds: list[torch.Tensor] | torch.Tensor = field(default_factory=list) negative_audio_prompt_embeds: list[torch.Tensor] | torch.Tensor = field( default_factory=list ) # Latent tensors latents: torch.Tensor | None = None y: torch.Tensor | None = None # Flux-2 latent_ids: torch.Tensor | None = None # Audio Latents audio_latents: torch.Tensor | None = None audio_noise: torch.Tensor | None = None raw_audio_latent_shape: tuple[int, ...] | None = None # Audio Parameters generate_audio: bool = True raw_latent_shape: torch.Tensor | None = None noise_pred: torch.Tensor | None = None # vae-encoded condition image image_latent: torch.Tensor | list[torch.Tensor] | None = None condition_image_latent_ids: torch.Tensor | list[torch.Tensor] | None = None vae_image_sizes: list[tuple[int, int]] | None = None # Latent dimensions height_latents: list[int] | int | None = None width_latents: list[int] | int | None = None # Timesteps timesteps: torch.Tensor | None = None paired_timesteps: torch.Tensor | None = None timestep: torch.Tensor | float | int | None = None step_index: int | None = None eta: float = 0.0 sigmas: list[float] | None = None n_tokens: int | None = None # Other parameters that may be needed by specific schedulers extra_step_kwargs: dict[str, Any] = field(default_factory=dict) # Component modules (populated by the pipeline) modules: dict[str, Any] = field(default_factory=dict) trajectory_timesteps: list[torch.Tensor] | None = None trajectory_latents: torch.Tensor | None = None trajectory_audio_latents: torch.Tensor | None = None # Extra parameters that might be needed by specific pipeline implementations extra: dict[str, Any] = field(default_factory=dict) is_warmup: bool = False # STA parameters STA_param: list | None = None is_cfg_negative: bool = False mask_search_final_result_pos: list[list] | None = None mask_search_final_result_neg: list[list] | None = None # VSA parameters VSA_sparsity: float = 0.0 # stage logging metrics: Optional["RequestMetrics"] = None # results output: torch.Tensor | None = None audio: torch.Tensor | None = None audio_sample_rate: int | None = None def __init__(self, **kwargs): # Initialize dataclass fields for name, field in self.__class__.__dataclass_fields__.items(): if name in kwargs: object.__setattr__(self, name, kwargs.pop(name)) elif field.default is not MISSING: object.__setattr__(self, name, field.default) elif field.default_factory is not MISSING: object.__setattr__(self, name, field.default_factory()) for name, value in kwargs.items(): setattr(self, name, value) self.validate() def __getattr__(self, name: str) -> Any: """ Delegate attribute access to sampling_params if not found in Req. This is only called when the attribute is not found in the instance. """ if name == "sampling_params": raise AttributeError( f"'{type(self).__name__}' object has no attribute '{name}'" ) sampling_params = object.__getattribute__(self, "sampling_params") if sampling_params is not None and hasattr(sampling_params, name): return getattr(sampling_params, name) raise AttributeError( f"'{type(self).__name__}' object has no attribute '{name}'" ) def __setattr__(self, name: str, value: Any) -> None: """ Smart attribute setting: 1. If field exists in Req, set it in Req 2. Else if field exists in sampling_params, set it in sampling_params 3. Else set it in Req (for dynamic attributes) """ if name == "sampling_params": object.__setattr__(self, name, value) return if name in self.__class__.__dataclass_fields__: object.__setattr__(self, name, value) return try: sampling_params = object.__getattribute__(self, "sampling_params") except AttributeError: sampling_params = None if sampling_params is not None and hasattr(sampling_params, name): setattr(sampling_params, name, value) return if sampling_params is None and name in SAMPLING_PARAMS_FIELDS: new_sp = SamplingParams() object.__setattr__(self, "sampling_params", new_sp) setattr(new_sp, name, value) return object.__setattr__(self, name, value) @property def batch_size(self): # Determine batch size if isinstance(self.prompt, list): batch_size = len(self.prompt) elif self.prompt is not None: batch_size = 1 else: batch_size = self.prompt_embeds[0].shape[0] # Adjust batch size for number of videos per prompt batch_size *= self.num_outputs_per_prompt return batch_size def output_file_path(self, num_outputs=1, output_idx=None): output_file_name = self.output_file_name if num_outputs > 1 and output_file_name: base, ext = os.path.splitext(output_file_name) output_file_name = f"{base}_{output_idx}{ext}" if self.output_path is None or not output_file_name: return None return os.path.join(self.output_path, output_file_name) def set_as_warmup(self, warmup_steps: int = 1): self.is_warmup = True self.save_output = False self.suppress_logs = True self.extra["cache_dit_num_inference_steps"] = self.num_inference_steps self.num_inference_steps = warmup_steps def copy_as_warmup(self, warmup_steps: int = 1) -> "Req": req = deepcopy(self) req.set_as_warmup(warmup_steps) return req def validate(self): """Initialize dependent fields after dataclass initialization.""" # Set do_classifier_free_guidance based on guidance scale and negative prompt if self.guidance_scale > 1.0 and self.negative_prompt is not None: self.do_classifier_free_guidance = True if self.negative_prompt_embeds is None: self.negative_prompt_embeds = [] if self.guidance_scale_2 is None: self.guidance_scale_2 = self.guidance_scale self.metrics = RequestMetrics(request_id=self.request_id) def adjust_size(self, server_args: ServerArgs): pass def __str__(self): return pprint.pformat(asdict(self), indent=2, width=120) def log(self, server_args: ServerArgs): if self.is_warmup or self.suppress_logs: return # TODO: in some cases (e.g., TI2I), height and weight might be undecided at this moment if self.height: target_height = align_to(self.height, 16) else: target_height = -1 if self.width: target_width = align_to(self.width, 16) else: target_width = -1 if logger.isEnabledFor(logging.DEBUG): display_prompt = self.prompt display_neg_prompt = self.negative_prompt else: display_prompt = _sanitize_for_logging(self.prompt, key_hint="prompt") display_neg_prompt = _sanitize_for_logging( self.negative_prompt, key_hint="negative_prompt" ) debug_str = f"""Sampling params: width: {target_width} height: {target_height} num_frames: {self.num_frames} fps: {self.fps} prompt: {display_prompt} neg_prompt: {display_neg_prompt} seed: {self.seed} infer_steps: {self.num_inference_steps} num_outputs_per_prompt: {self.num_outputs_per_prompt} guidance_scale: {self.guidance_scale} embedded_guidance_scale: {server_args.pipeline_config.embedded_cfg_scale} n_tokens: {self.n_tokens} flow_shift: {server_args.pipeline_config.flow_shift} image_path: {self.image_path} save_output: {self.save_output} output_file_path: {self.output_file_path()} """ # type: ignore[attr-defined] logger.info(debug_str) @dataclass class OutputBatch: """ Final output (after pipeline completion) """ output: torch.Tensor | None = None audio: torch.Tensor | None = None audio_sample_rate: int | None = None trajectory_timesteps: list[torch.Tensor] | None = None trajectory_latents: torch.Tensor | None = None trajectory_decoded: list[torch.Tensor] | None = None error: str | None = None output_file_paths: list[str] | None = None # logged metrics info, directly from Req.timings metrics: Optional["RequestMetrics"] = None # For ComfyUI integration: noise prediction from denoising stage noise_pred: torch.Tensor | None = None peak_memory_mb: float = 0.0 ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Pipeline stages for diffusion models. This package contains the various stages that can be composed to create complete diffusion pipelines. """ from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines_core.stages.causal_denoising import ( CausalDMDDenoisingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.comfyui_latent_preparation import ( ComfyUILatentPreparationStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding_av import ( LTX2AVDecodingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_av import ( LTX2AVDenoisingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_dmd import ( DmdDenoisingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.encoding import EncodingStage # Hunyuan3D paint stages from sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_paint import ( Hunyuan3DPaintPostprocessStage, Hunyuan3DPaintPreprocessStage, Hunyuan3DPaintTexGenStage, ) # Hunyuan3D shape stages from sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_shape import ( Hunyuan3DShapeBeforeDenoisingStage, Hunyuan3DShapeDenoisingStage, Hunyuan3DShapeExportStage, Hunyuan3DShapeSaveStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.image_encoding import ( ImageEncodingStage, ImageVAEEncodingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import ( InputValidationStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import ( LatentPreparationStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation_av import ( LTX2AVLatentPreparationStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.text_connector import ( LTX2TextConnectorStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import ( TextEncodingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation import ( TimestepPreparationStage, ) __all__ = [ "PipelineStage", "InputValidationStage", "TimestepPreparationStage", "LatentPreparationStage", "ComfyUILatentPreparationStage", "LTX2AVLatentPreparationStage", "DenoisingStage", "DmdDenoisingStage", "LTX2AVDenoisingStage", "CausalDMDDenoisingStage", "EncodingStage", "DecodingStage", "LTX2AVDecodingStage", "ImageEncodingStage", "ImageVAEEncodingStage", "TextEncodingStage", "LTX2TextConnectorStage", # Hunyuan3D shape stages "Hunyuan3DShapeBeforeDenoisingStage", "Hunyuan3DShapeDenoisingStage", "Hunyuan3DShapeExportStage", "Hunyuan3DShapeSaveStage", # Hunyuan3D paint stages "Hunyuan3DPaintPreprocessStage", "Hunyuan3DPaintTexGenStage", "Hunyuan3DPaintPostprocessStage", ] ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Base classes for pipeline stages. This module defines the abstract base classes for pipeline stages that can be composed to create complete diffusion pipelines. """ from abc import ABC, abstractmethod from enum import Enum, auto import torch from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler logger = init_logger(__name__) class StageParallelismType(Enum): # execute on all gpus REPLICATED = auto() # executed on main rank only MAIN_RANK_ONLY = auto() # this stage requires a cfg-parallel CFG_PARALLEL = auto() class StageVerificationError(Exception): """Exception raised when stage verification fails.""" pass class PipelineStage(ABC): """ Abstract base class for all pipeline stages. A pipeline stage represents a discrete step in the diffusion process that can be composed with other stages to create a complete pipeline. Each stage is responsible for a specific part of the process, such as prompt encoding, latent preparation, etc. """ def __init__(self): self.server_args = get_global_server_args() def log_info(self, msg, *args): """Logs an informational message with the stage name as a prefix.""" if self.server_args.comfyui_mode: return logger.info(f"[{self.__class__.__name__}] {msg}", *args) def log_warning(self, msg, *args): """Logs a warning message with the stage name as a prefix.""" logger.warning(f"[{self.__class__.__name__}] {msg}", *args) def log_error(self, msg, *args): """Logs an error message with the stage name as a prefix.""" logger.error(f"[{self.__class__.__name__}] {msg}", *args) def log_debug(self, msg, *args): """Logs a debug message with the stage name as a prefix.""" logger.debug(f"[{self.__class__.__name__}] {msg}", *args) def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """ Verify the input for the stage. Example: from sglang.multimodal_gen.runtime.pipelines.stages.validators import V, VerificationResult def verify_input(self, batch, server_args): result = VerificationResult() result.add_check("height", batch.height, V.positive_int_divisible(8)) result.add_check("width", batch.width, V.positive_int_divisible(8)) result.add_check("image_latent", batch.image_latent, V.is_tensor) return result """ # Default implementation - no verification return VerificationResult() def maybe_free_model_hooks(self): pass def load_model(self): """ Load the model for the stage. """ pass def offload_model(self): """ Offload the model for the stage. """ pass # execute on all ranks by default @property def parallelism_type(self) -> StageParallelismType: # if get_global_server_args().enable_cfg_parallel: # return StageParallelismType.MAIN_RANK_ONLY return StageParallelismType.REPLICATED def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """ Verify the output for the stage. Returns: A VerificationResult containing the verification status. """ # Default implementation - no verification return VerificationResult() def _run_verification( self, verification_result: VerificationResult, stage_name: str, verification_type: str, ) -> None: """ Run verification and raise errors if any checks fail. Args: verification_result: Results from verify_input or verify_output stage_name: Name of the current stage verification_type: "input" or "output" """ if not verification_result.is_valid(): failed_fields = verification_result.get_failed_fields() if failed_fields: # Get detailed failure information detailed_summary = verification_result.get_failure_summary() failed_fields_str = ", ".join(failed_fields) error_msg = ( f"{verification_type.capitalize()} verification failed for {stage_name}: " f"Failed fields: {failed_fields_str}\n" f"Details: {detailed_summary}" ) raise StageVerificationError(error_msg) @property def device(self) -> torch.device: """Get the device for this stage.""" return torch.device( current_platform.device_type, ) def set_logging(self, enable: bool): """ Enable or disable logging for this stage. Args: enable: Whether to enable logging. """ self._enable_logging = enable def __call__( self, batch: Req, server_args: ServerArgs, ) -> Req: """ Execute the stage's processing on the batch with optional verification and logging. Should not be overridden by subclasses. Returns: The updated batch information after this stage's processing. """ stage_name = self.__class__.__name__ # Check if verification is enabled (simple approach for prototype) # Pre-execution input verification try: input_result = self.verify_input(batch, server_args) self._run_verification(input_result, stage_name, "input") except Exception as e: logger.error("Input verification failed for %s: %s", stage_name, str(e)) raise # Execute the actual stage logic with unified profiling with StageProfiler( stage_name, logger=logger, metrics=batch.metrics, log_stage_start_end=not batch.is_warmup and not (self.server_args and self.server_args.comfyui_mode), perf_dump_path_provided=batch.perf_dump_path is not None, ): result = self.forward(batch, server_args) # Post-execution output verification try: output_result = self.verify_output(result, server_args) self._run_verification(output_result, stage_name, "output") except Exception as e: logger.error("Output verification failed for %s: %s", stage_name, str(e)) raise return result @abstractmethod def forward( self, batch: Req, server_args: ServerArgs, ) -> Req: """ Forward pass of the stage's processing. This method should be implemented by subclasses to provide the forward processing logic for the stage. Returns: The updated batch information after this stage's processing. """ raise NotImplementedError def backward( self, batch: Req, server_args: ServerArgs, ) -> Req: raise NotImplementedError ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import torch # type: ignore from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.models.utils import pred_noise_to_pred_video from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class CausalDMDDenoisingStage(DenoisingStage): """ Denoising stage for causal diffusion. """ def __init__(self, transformer, scheduler) -> None: super().__init__(transformer, scheduler) # KV and cross-attention cache state (initialized on first forward) self.kv_cache1: list | None = None self.crossattn_cache: list | None = None # Model-dependent constants (aligned with causal_inference.py assumptions) self.num_transformer_blocks = self.transformer.config.arch_config.num_layers self.num_frames_per_block = ( self.transformer.config.arch_config.num_frames_per_block ) self.sliding_window_num_frames = ( self.transformer.config.arch_config.sliding_window_num_frames ) try: self.local_attn_size = getattr( self.transformer.model, "local_attn_size", -1 ) # type: ignore except Exception: self.local_attn_size = -1 def forward( self, batch: Req, server_args: ServerArgs, ) -> Req: target_dtype = torch.bfloat16 autocast_enabled = ( target_dtype != torch.float32 ) and not server_args.disable_autocast latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2] patch_ratio = ( self.transformer.config.arch_config.patch_size[-1] * self.transformer.config.arch_config.patch_size[-2] ) self.frame_seq_length = latent_seq_length // patch_ratio # TODO(will): make this a parameter once we add i2v support independent_first_frame = self.transformer.independent_first_frame # Timesteps for DMD timesteps = torch.tensor( server_args.pipeline_config.dmd_denoising_steps, dtype=torch.long ).cpu() if server_args.pipeline_config.warp_denoising_step: logger.info("Warping timesteps...") scheduler_timesteps = torch.cat( (self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)) ) timesteps = scheduler_timesteps[1000 - timesteps] timesteps = timesteps.to(get_local_torch_device()) logger.info("Using timesteps: %s", timesteps) # Image kwargs (kept empty unless caller provides compatible args) image_kwargs: dict = {} pos_cond_kwargs = self.prepare_extra_func_kwargs( self.transformer.forward, { # "encoder_hidden_states_2": batch.clip_embedding_pos, "encoder_attention_mask": batch.prompt_attention_mask, }, ) # STA if self.attn_backend.get_enum() == AttentionBackendEnum.SLIDING_TILE_ATTN: self.prepare_sta_param(batch, server_args) # Latents and prompts assert batch.latents is not None, "latents must be provided" latents = batch.latents # [B, C, T, H, W] b, c, t, h, w = latents.shape prompt_embeds = batch.prompt_embeds assert torch.isnan(prompt_embeds[0]).sum() == 0 # Initialize or reset caches if self.kv_cache1 is None: self._initialize_kv_cache( batch_size=latents.shape[0], dtype=target_dtype, device=latents.device ) self._initialize_crossattn_cache( batch_size=latents.shape[0], max_text_len=server_args.pipeline_config.text_encoder_configs[ 0 ].arch_config.text_len, dtype=target_dtype, device=latents.device, ) else: assert self.crossattn_cache is not None # reset cross-attention cache for block_index in range(self.num_transformer_blocks): self.crossattn_cache[block_index]["is_init"] = False # type: ignore # reset kv cache pointers for block_index in range(len(self.kv_cache1)): self.kv_cache1[block_index]["global_end_index"] = ( torch.tensor( # type: ignore [0], dtype=torch.long, device=latents.device ) ) self.kv_cache1[block_index]["local_end_index"] = ( torch.tensor( # type: ignore [0], dtype=torch.long, device=latents.device ) ) # Optional: cache context features from provided image latents prior to generation current_start_frame = 0 if getattr(batch, "image_latent", None) is not None: image_latent = batch.image_latent assert image_latent is not None input_frames = image_latent.shape[2] # timestep zero (or configured context noise) for cache warm-up t_zero = torch.zeros( [latents.shape[0]], device=latents.device, dtype=torch.long ) if independent_first_frame and input_frames >= 1: # warm-up with the very first frame independently image_first_btchw = ( image_latent[:, :, :1, :, :].to(target_dtype).permute(0, 2, 1, 3, 4) ) with torch.autocast( device_type=current_platform.device_type, dtype=target_dtype, enabled=autocast_enabled, ): _ = self.transformer( image_first_btchw, prompt_embeds, t_zero, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=current_start_frame * self.frame_seq_length, **image_kwargs, **pos_cond_kwargs, ) current_start_frame += 1 remaining_frames = input_frames - 1 else: remaining_frames = input_frames # process remaining input frames in blocks of num_frame_per_block while remaining_frames > 0: block = min(self.num_frames_per_block, remaining_frames) ref_btchw = ( image_latent[ :, :, current_start_frame : current_start_frame + block, :, : ] .to(target_dtype) .permute(0, 2, 1, 3, 4) ) with torch.autocast( device_type=current_platform.device_type, dtype=target_dtype, enabled=autocast_enabled, ): _ = self.transformer( ref_btchw, prompt_embeds, t_zero, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=current_start_frame * self.frame_seq_length, **image_kwargs, **pos_cond_kwargs, ) current_start_frame += block remaining_frames -= block # Base position offset from any cache warm-up pos_start_base = current_start_frame # Determine block sizes if not independent_first_frame or ( independent_first_frame and batch.image_latent is not None ): if t % self.num_frames_per_block != 0: raise ValueError( "num_frames must be divisible by num_frames_per_block for causal DMD denoising" ) num_blocks = t // self.num_frames_per_block block_sizes = [self.num_frames_per_block] * num_blocks start_index = 0 else: if (t - 1) % self.num_frames_per_block != 0: raise ValueError( "(num_frames - 1) must be divisible by num_frame_per_block when independent_first_frame=True" ) num_blocks = (t - 1) // self.num_frames_per_block block_sizes = [1] + [self.num_frames_per_block] * num_blocks start_index = 0 # DMD loop in causal blocks with self.progress_bar(total=len(block_sizes) * len(timesteps)) as progress_bar: for current_num_frames in block_sizes: current_latents = latents[ :, :, start_index : start_index + current_num_frames, :, : ] # use BTCHW for DMD conversion routines noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4) video_raw_latent_shape = noise_latents_btchw.shape for i, t_cur in enumerate(timesteps): # Copy for pred conversion noise_latents = noise_latents_btchw.clone() latent_model_input = current_latents.to(target_dtype) if ( batch.image_latent is not None and independent_first_frame and start_index == 0 ): latent_model_input = torch.cat( [latent_model_input, batch.image_latent.to(target_dtype)], dim=2, ) # Prepare inputs t_expand = t_cur.repeat(latent_model_input.shape[0]) # Attention metadata if needed if ( self.attn_backend.get_enum() == AttentionBackendEnum.VIDEO_SPARSE_ATTN ): self.attn_metadata_builder_cls = ( self.attn_backend.get_builder_cls() ) if self.attn_metadata_builder_cls is not None: self.attn_metadata_builder = ( self.attn_metadata_builder_cls() ) attn_metadata = self.attn_metadata_builder.build( # type: ignore current_timestep=i, # type: ignore raw_latent_shape=( current_num_frames, h, w, ), # type: ignore patch_size=server_args.pipeline_config.dit_config.patch_size, # type: ignore STA_param=batch.STA_param, # type: ignore VSA_sparsity=server_args.attention_backend_config.VSA_sparsity, # type: ignore device=get_local_torch_device(), # type: ignore ) # type: ignore assert ( attn_metadata is not None ), "attn_metadata cannot be None" else: attn_metadata = None else: attn_metadata = None with ( torch.autocast( device_type=current_platform.device_type, dtype=target_dtype, enabled=autocast_enabled, ), set_forward_context( current_timestep=i, attn_metadata=attn_metadata, forward_batch=batch, ), ): # Run transformer; follow DMD stage pattern t_expanded_noise = t_cur * torch.ones( (latent_model_input.shape[0], 1), device=latent_model_input.device, dtype=torch.long, ) pred_noise_btchw = self.transformer( latent_model_input, prompt_embeds, t_expanded_noise, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=(pos_start_base + start_index) * self.frame_seq_length, start_frame=start_index, **image_kwargs, **pos_cond_kwargs, ).permute(0, 2, 1, 3, 4) # Convert pred noise to pred video with FM Euler scheduler utilities pred_video_btchw = pred_noise_to_pred_video( pred_noise=pred_noise_btchw.flatten(0, 1), noise_input_latent=noise_latents.flatten(0, 1), timestep=t_expand, scheduler=self.scheduler, ).unflatten(0, pred_noise_btchw.shape[:2]) if i < len(timesteps) - 1: next_timestep = timesteps[i + 1] * torch.ones( [1], dtype=torch.long, device=pred_video_btchw.device ) noise = torch.randn( video_raw_latent_shape, dtype=pred_video_btchw.dtype, generator=( batch.generator[0] if isinstance(batch.generator, list) else batch.generator ), device=self.device, ) noise_btchw = noise noise_latents_btchw = self.scheduler.add_noise( pred_video_btchw.flatten(0, 1), noise_btchw.flatten(0, 1), next_timestep, ).unflatten(0, pred_video_btchw.shape[:2]) current_latents = noise_latents_btchw.permute(0, 2, 1, 3, 4) else: current_latents = pred_video_btchw.permute(0, 2, 1, 3, 4) if progress_bar is not None: progress_bar.update() # Write back and advance latents[:, :, start_index : start_index + current_num_frames, :, :] = ( current_latents ) # Re-run with context timestep to update KV cache using clean context context_noise = getattr(server_args.pipeline_config, "context_noise", 0) t_context = torch.ones( [latents.shape[0]], device=latents.device, dtype=torch.long ) * int(context_noise) context_bcthw = current_latents.to(target_dtype) with ( torch.autocast( device_type=current_platform.device_type, dtype=target_dtype, enabled=autocast_enabled, ), set_forward_context( current_timestep=0, attn_metadata=attn_metadata, forward_batch=batch, ), ): t_expanded_context = t_context.unsqueeze(1) _ = self.transformer( context_bcthw, prompt_embeds, t_expanded_context, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=(pos_start_base + start_index) * self.frame_seq_length, start_frame=start_index, **image_kwargs, **pos_cond_kwargs, ) start_index += current_num_frames batch.latents = latents return batch def _initialize_kv_cache(self, batch_size, dtype, device) -> None: """ Initialize a Per-GPU KV cache aligned with the Wan model assumptions. """ kv_cache1 = [] num_attention_heads = self.transformer.num_attention_heads attention_head_dim = self.transformer.attention_head_dim if self.local_attn_size != -1: kv_cache_size = self.local_attn_size * self.frame_seq_length else: kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames for _ in range(self.num_transformer_blocks): kv_cache1.append( { "k": torch.zeros( [ batch_size, kv_cache_size, num_attention_heads, attention_head_dim, ], dtype=dtype, device=device, ), "v": torch.zeros( [ batch_size, kv_cache_size, num_attention_heads, attention_head_dim, ], dtype=dtype, device=device, ), "global_end_index": torch.tensor( [0], dtype=torch.long, device=device ), "local_end_index": torch.tensor( [0], dtype=torch.long, device=device ), } ) self.kv_cache1 = kv_cache1 def _initialize_crossattn_cache( self, batch_size, max_text_len, dtype, device ) -> None: """ Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions. """ crossattn_cache = [] num_attention_heads = self.transformer.num_attention_heads attention_head_dim = self.transformer.attention_head_dim for _ in range(self.num_transformer_blocks): crossattn_cache.append( { "k": torch.zeros( [ batch_size, max_text_len, num_attention_heads, attention_head_dim, ], dtype=dtype, device=device, ), "v": torch.zeros( [ batch_size, max_text_len, num_attention_heads, attention_head_dim, ], dtype=dtype, device=device, ), "is_init": False, } ) self.crossattn_cache = crossattn_cache def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify denoising stage inputs.""" result = VerificationResult() result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) result.add_check("image_embeds", batch.image_embeds, V.is_list) result.add_check( "image_latent", batch.image_latent, V.none_or_tensor_with_dims(5) ) result.add_check( "num_inference_steps", batch.num_inference_steps, V.positive_int ) result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float) result.add_check("eta", batch.eta, V.non_negative_float) result.add_check("generator", batch.generator, V.generator_or_list_generators) result.add_check( "do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value, ) result.add_check( "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x), ) return result ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/comfyui_latent_preparation.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ ComfyUI latent preparation stage with device mismatch fix. This stage extends LatentPreparationStage to handle device mismatch issues that occur when tensors are pickled and unpickled via broadcast_pyobj in multi-GPU scenarios. """ import dataclasses import torch from sglang.multimodal_gen.runtime.distributed import ( get_local_torch_device, get_sp_group, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import ( LatentPreparationStage, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class ComfyUILatentPreparationStage(LatentPreparationStage): """ ComfyUI-specific latent preparation stage with device mismatch fix. This stage extends LatentPreparationStage to automatically fix device mismatches for tensor fields on non-source ranks in multi-GPU scenarios. """ @staticmethod def _fix_tensor_device(value, target_device): """Recursively fix tensor device, handling single tensors, lists, and tuples.""" if isinstance(value, torch.Tensor): if value.device != target_device: return value.detach().clone().to(target_device) return value elif isinstance(value, list): return [ ComfyUILatentPreparationStage._fix_tensor_device(v, target_device) for v in value ] elif isinstance(value, tuple): return tuple( ComfyUILatentPreparationStage._fix_tensor_device(v, target_device) for v in value ) return value @staticmethod def _has_tensor(value): """Check if value contains any tensor.""" if isinstance(value, torch.Tensor): return True elif isinstance(value, (list, tuple)): return any(ComfyUILatentPreparationStage._has_tensor(v) for v in value) return False def forward(self, batch, server_args): """ Prepare latents with device mismatch fix for ComfyUI pipelines. This method first fixes device mismatches for all tensor fields, then calls the parent class's forward method, and ensures raw_latent_shape is set correctly (before packing, for proper unpadding later). """ # Fix device mismatch for tensor fields on non-source ranks if get_sp_world_size() > 1: sp_group = get_sp_group() target_device = get_local_torch_device() if sp_group.rank != 0: logger.debug( f"[ComfyUILatentPreparationStage] Fixing tensor device on rank={sp_group.rank} " f"target_device={target_device}" ) if dataclasses.is_dataclass(batch): for field in dataclasses.fields(batch): value = getattr(batch, field.name, None) if value is not None and self._has_tensor(value): fixed_value = self._fix_tensor_device(value, target_device) setattr(batch, field.name, fixed_value) else: for attr_name in dir(batch): if not attr_name.startswith("_") and not callable( getattr(batch, attr_name, None) ): try: value = getattr(batch, attr_name, None) if value is not None and self._has_tensor(value): fixed_value = self._fix_tensor_device( value, target_device ) setattr(batch, attr_name, fixed_value) except (AttributeError, TypeError): continue original_latents_shape = None if batch.latents is not None: original_latents_shape = batch.latents.shape # Call parent class's forward method result = super().forward(batch, server_args) if original_latents_shape is not None: # Preserve the original shape before any potential packing/conversion # (e.g., 4D spatial -> 3D sequence) to ensure proper unpadding later. result.raw_latent_shape = original_latents_shape return result ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Decoding stage for diffusion pipelines. """ import weakref import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import VAELoader from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( PipelineStage, StageParallelismType, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) def _ensure_tensor_decode_output(decode_output): """ Ensure VAE decode output is a tensor. Some VAE implementations return DecoderOutput objects with a .sample attribute, tuples, or tensors directly. This function normalizes the output to always be a tensor. Args: decode_output: Output from VAE.decode(), can be DecoderOutput, tuple, or torch.Tensor Returns: torch.Tensor: The decoded image tensor """ if isinstance(decode_output, tuple): return decode_output[0] if hasattr(decode_output, "sample"): return decode_output.sample return decode_output class DecodingStage(PipelineStage): """ Stage for decoding latent representations into pixel space. This stage handles the decoding of latent representations into the final output format (e.g., pixel values). """ def __init__(self, vae, pipeline=None, component_name: str = "vae") -> None: super().__init__() self.vae: ParallelTiledVAE = vae self.pipeline = weakref.ref(pipeline) if pipeline else None self.component_name = component_name @property def parallelism_type(self) -> StageParallelismType: if get_global_server_args().enable_cfg_parallel: return StageParallelismType.MAIN_RANK_ONLY return StageParallelismType.REPLICATED def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify decoding stage inputs.""" result = VerificationResult() # Denoised latents for VAE decoding: [batch_size, channels, frames, height_latents, width_latents] # result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify decoding stage outputs.""" result = VerificationResult() # Decoded video/images: [batch_size, channels, frames, height, width] # result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)]) return result def scale_and_shift(self, latents: torch.Tensor, server_args): scaling_factor, shift_factor = ( server_args.pipeline_config.get_decode_scale_and_shift( latents.device, latents.dtype, self.vae ) ) # 1. scale if isinstance(scaling_factor, torch.Tensor): latents = latents / scaling_factor.to(latents.device, latents.dtype) else: latents = latents / scaling_factor # 2. apply shifting if needed if shift_factor is not None: if isinstance(shift_factor, torch.Tensor): latents += shift_factor.to(latents.device, latents.dtype) else: latents += shift_factor return latents @torch.no_grad() def decode(self, latents: torch.Tensor, server_args: ServerArgs) -> torch.Tensor: """ Decode latent representations into pixel space using VAE. Args: latents: Input latent tensor with shape (batch, channels, frames, height_latents, width_latents) server_args: Configuration containing: - disable_autocast: Whether to disable automatic mixed precision (default: False) - pipeline_config.vae_precision: VAE computation precision ("fp32", "fp16", "bf16") - pipeline_config.vae_tiling: Whether to enable VAE tiling for memory efficiency Returns: Decoded video tensor with shape (batch, channels, frames, height, width), normalized to [0, 1] range and moved to CPU as float32 """ vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] self.vae = self.vae.to(device=get_local_torch_device(), dtype=vae_dtype) latents = latents.to(get_local_torch_device()) vae_autocast_enabled = ( vae_dtype != torch.float32 ) and not server_args.disable_autocast # scale and shift latents = self.scale_and_shift(latents, server_args) # Preprocess latents before decoding (e.g., unpatchify for standard Flux2 VAE) latents = server_args.pipeline_config.preprocess_decoding( latents, server_args, vae=self.vae ) # Decode latents with torch.autocast( device_type=current_platform.device_type, dtype=vae_dtype, enabled=vae_autocast_enabled, ): try: # TODO: make it more specific if server_args.pipeline_config.vae_tiling: self.vae.enable_tiling() except Exception: pass if not vae_autocast_enabled: latents = latents.to(vae_dtype) decode_output = self.vae.decode(latents) image = _ensure_tensor_decode_output(decode_output) # De-normalize image to [0, 1] range image = (image / 2 + 0.5).clamp(0, 1) return image def load_model(self): # load vae if not already loaded (used for memory constrained devices) pipeline = self.pipeline() if self.pipeline else None if not self.server_args.model_loaded[self.component_name]: loader = VAELoader() self.vae, _ = loader.load( self.server_args.model_paths[self.component_name], self.server_args, component_name=self.component_name, transformers_or_diffusers=loader.expected_library, ) if pipeline: pipeline.add_module(self.component_name, self.vae) self.server_args.model_loaded[self.component_name] = True def offload_model(self): # Offload models if needed self.maybe_free_model_hooks() if self.server_args.vae_cpu_offload: self.vae.to("cpu", non_blocking=True) if torch.backends.mps.is_available(): # Flush lazy MPS kernels before freeing weights to avoid hangs. torch.mps.synchronize() del self.vae pipeline = self.pipeline() if self.pipeline else None if pipeline is not None and self.component_name in pipeline.modules: del pipeline.modules[self.component_name] self.server_args.model_loaded[self.component_name] = False @torch.no_grad() def forward( self, batch: Req, server_args: ServerArgs, ) -> OutputBatch: """ Decode latent representations into pixel space. This method processes the batch through the VAE decoder, converting latent representations to pixel-space video/images. It also optionally decodes trajectory latents for visualization purposes. """ # load vae if not already loaded (used for memory constrained devices) self.load_model() frames = self.decode(batch.latents, server_args) # decode trajectory latents if needed if batch.return_trajectory_decoded: assert ( batch.trajectory_latents is not None ), "batch should have trajectory latents" # 1. Batch trajectory decoding to improve GPU utilization # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width] B, T, C, F, H, W = batch.trajectory_latents.shape flat_latents = batch.trajectory_latents.view(B * T, C, F, H, W) logger.info("decoding %s trajectory latents in batch", B * T) # Use the optimized batch decode all_decoded = self.decode(flat_latents, server_args) # 2. Reshape back # Keep on GPU to allow faster vectorized post-processing decoded_tensor = all_decoded.view(B, T, *all_decoded.shape[1:]) # Convert to list of tensors (per timestep) as expected by OutputBatch # Each element in list is [B, channels, frames, H_out, W_out] trajectory_decoded = [decoded_tensor[:, i] for i in range(T)] else: trajectory_decoded = None frames = server_args.pipeline_config.post_decoding(frames, server_args) # Update batch with decoded image output_batch = OutputBatch( output=frames, trajectory_timesteps=batch.trajectory_timesteps, trajectory_latents=batch.trajectory_latents, trajectory_decoded=trajectory_decoded, metrics=batch.metrics, ) # Keep VAE resident during warmup; the real request needs it next. if not getattr(batch, "is_warmup", False): self.offload_model() return output_batch ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py ================================================ import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) class LTX2AVDecodingStage(DecodingStage): """ LTX-2 specific decoding stage that handles both video and audio decoding. """ def __init__(self, vae, audio_vae, vocoder, pipeline=None): super().__init__(vae, pipeline) self.audio_vae = audio_vae self.vocoder = vocoder # Add video processor for postprocessing from diffusers.video_processor import VideoProcessor self.video_processor = VideoProcessor(vae_scale_factor=32) def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: self.load_model() self.vae = self.vae.to(get_local_torch_device()) self.vae.eval() latents = batch.latents.to(get_local_torch_device()) vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] vae_autocast_enabled = ( vae_dtype != torch.float32 ) and not server_args.disable_autocast original_dtype = vae_dtype self.vae.to(torch.bfloat16) latents = latents.to(torch.bfloat16) std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latents) mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latents) latents = latents * std + mean latents = server_args.pipeline_config.preprocess_decoding( latents, server_args, vae=self.vae ) with torch.autocast( device_type=current_platform.device_type, dtype=vae_dtype, enabled=vae_autocast_enabled, ): try: if server_args.pipeline_config.vae_tiling: self.vae.enable_tiling() except Exception: pass decode_output = self.vae.decode(latents) if isinstance(decode_output, tuple): video = decode_output[0] elif hasattr(decode_output, "sample"): video = decode_output.sample else: video = decode_output self.vae.to(original_dtype) video = self.video_processor.postprocess_video(video, output_type="np") output_batch = OutputBatch( output=video, trajectory_timesteps=batch.trajectory_timesteps, trajectory_latents=batch.trajectory_latents, trajectory_decoded=None, metrics=batch.metrics, ) # 2. Decode Audio try: audio_latents = batch.audio_latents except AttributeError: audio_latents = None if audio_latents is not None: # Ensure device/dtype device = get_local_torch_device() self.audio_vae = self.audio_vae.to(device) self.vocoder = self.vocoder.to(device) self.audio_vae.eval() self.vocoder.eval() try: dtype = self.audio_vae.dtype except AttributeError: dtype = None if dtype is None: try: dtype = next(self.audio_vae.parameters()).dtype except StopIteration: dtype = torch.float32 audio_latents = audio_latents.to(device, dtype=dtype) try: latents_std = self.audio_vae.latents_std except AttributeError: latents_std = None if isinstance(latents_std, torch.Tensor) and torch.all(latents_std == 0): logger.warning( "audio_vae.latents_std is all zeros; audio denorm may be incorrect." ) with torch.no_grad(): # Decode latents to spectrogram spectrogram = self.audio_vae.decode(audio_latents, return_dict=False)[0] if hasattr(self.vocoder, "conv_in") and hasattr( self.vocoder.conv_in, "in_channels" ): expected_in = int(self.vocoder.conv_in.in_channels) actual_in = int(spectrogram.shape[1]) * int(spectrogram.shape[3]) if actual_in != expected_in: raise ValueError( f"Vocoder expects channels*mel_bins={expected_in}, got {actual_in} from spectrogram shape {tuple(spectrogram.shape)}" ) # Decode spectrogram to waveform waveform = self.vocoder(spectrogram) output_batch.audio = waveform.cpu().float() try: pipeline_audio_cfg = server_args.pipeline_config.audio_vae_config except AttributeError: pipeline_audio_cfg = None try: pipeline_audio_arch = pipeline_audio_cfg.arch_config # type: ignore[union-attr] except AttributeError: pipeline_audio_arch = None try: pipeline_audio_sr = pipeline_audio_arch.sample_rate # type: ignore[union-attr] except AttributeError: pipeline_audio_sr = None try: vocoder_sr = self.vocoder.sample_rate except AttributeError: vocoder_sr = None try: audio_vae_sr = self.audio_vae.sample_rate except AttributeError: audio_vae_sr = None output_batch.audio_sample_rate = ( vocoder_sr or audio_vae_sr or pipeline_audio_sr ) self.offload_model() return output_batch ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Denoising stage for diffusion pipelines. """ import inspect import math import os import time import weakref from collections.abc import Iterable from functools import lru_cache from typing import Any import torch import torch.nn as nn from einops import rearrange from tqdm.auto import tqdm from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType, STA_Mode from sglang.multimodal_gen.configs.pipeline_configs.wan import ( Wan2_2_TI2V_5B_Config, ) from sglang.multimodal_gen.runtime.cache.cache_dit_integration import ( CacheDitConfig, enable_cache_on_dual_transformer, enable_cache_on_transformer, get_scm_mask, refresh_context_on_dual_transformer, refresh_context_on_transformer, ) from sglang.multimodal_gen.runtime.distributed import ( cfg_model_parallel_all_reduce, get_local_torch_device, get_sp_group, get_sp_parallel_rank, get_sp_world_size, get_tp_group, get_world_group, get_world_size, ) from sglang.multimodal_gen.runtime.distributed.communication_op import ( sequence_model_parallel_all_gather, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_cfg_group, get_classifier_free_guidance_rank, ) from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend from sglang.multimodal_gen.runtime.layers.attention.STA_configuration import ( configure_sta, save_mask_search_results, ) from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import ( TransformerLoader, ) from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( PipelineStage, StageParallelismType, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler from sglang.multimodal_gen.utils import dict_to_3d_list, masks_like from sglang.srt.utils.common import get_compiler_backend logger = init_logger(__name__) class DenoisingStage(PipelineStage): """ Stage for running the denoising loop in diffusion pipelines. This stage handles the iterative denoising process that transforms the initial noise into the final output. """ def __init__( self, transformer, scheduler, pipeline=None, transformer_2=None, vae=None ) -> None: super().__init__() self.transformer = transformer self.transformer_2 = transformer_2 hidden_size = self.server_args.pipeline_config.dit_config.hidden_size num_attention_heads = ( self.server_args.pipeline_config.dit_config.num_attention_heads ) attn_head_size = hidden_size // num_attention_heads # torch compile for transformer in filter(None, [self.transformer, self.transformer_2]): self._maybe_enable_torch_compile(transformer) self.scheduler = scheduler self.vae = vae self.pipeline = weakref.ref(pipeline) if pipeline else None # TODO(will): hack, should use the actual one in dit self.attn_backend = get_attn_backend( head_size=attn_head_size, dtype=torch.float16, ) # cfg self.guidance = None # misc self.profiler = None # cache-dit state (for delayed mounting and idempotent control) self._cache_dit_enabled = False self._cached_num_steps = None self._is_warmed_up = False def _maybe_enable_torch_compile(self, module: object) -> None: """ Compile a module with torch.compile, and enable inductor overlap tweak if available. No-op if torch compile is disabled or the object is not a nn.Module. """ if not self.server_args.enable_torch_compile or not isinstance( module, nn.Module ): return compile_kwargs: dict[str, Any] = {"fullgraph": False, "dynamic": None} if current_platform.is_npu(): backend = get_compiler_backend() compile_kwargs["backend"] = backend compile_kwargs["dynamic"] = False logger.info("Compiling transformer with torchair backend on NPU") else: try: import torch._inductor.config as _inductor_cfg _inductor_cfg.reorder_for_compute_comm_overlap = True except ImportError: pass mode = os.environ.get( "SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs" ) compile_kwargs["mode"] = mode logger.info(f"Compiling transformer with mode: {mode}") # TODO(triple-mu): support customized fullgraph and dynamic in the future module.compile(**compile_kwargs) def _maybe_enable_cache_dit( self, num_inference_steps: int | tuple[int, int], batch: Req ) -> None: """Enable cache-dit on the transformers if configured (idempotent). This method should be called after the transformer is fully loaded and before torch.compile is applied. For dual-transformer models (e.g., Wan2.2), this enables cache-dit on both transformers with (potentially) different configurations. """ if isinstance(num_inference_steps, tuple): num_high_noise_steps, num_low_noise_steps = num_inference_steps # NOTE: When a new request arrives, we need to refresh the cache-dit context. if self._cache_dit_enabled: scm_preset = envs.SGLANG_CACHE_DIT_SCM_PRESET scm_preset = None if scm_preset == "none" else scm_preset if isinstance(num_inference_steps, tuple): refresh_context_on_dual_transformer( self.transformer, self.transformer_2, num_high_noise_steps, num_low_noise_steps, scm_preset=scm_preset, ) else: refresh_context_on_transformer( self.transformer, num_inference_steps, scm_preset=scm_preset, ) return # check if cache-dit is enabled in config if not envs.SGLANG_CACHE_DIT_ENABLED or batch.is_warmup: return world_size = get_world_size() parallelized = world_size > 1 sp_group = None tp_group = None if parallelized: sp_group_candidate = get_sp_group() tp_group_candidate = get_tp_group() sp_world_size = sp_group_candidate.world_size if sp_group_candidate else 1 tp_world_size = tp_group_candidate.world_size if tp_group_candidate else 1 has_sp = sp_world_size > 1 has_tp = tp_world_size > 1 sp_group = sp_group_candidate.device_group if has_sp else None tp_group = tp_group_candidate.device_group if has_tp else None logger.info( "cache-dit enabled in distributed environment (world_size=%d, has_sp=%s, has_tp=%s)", world_size, has_sp, has_tp, ) # === Parse SCM configuration from envs === # SCM is shared between primary and secondary transformers scm_preset = envs.SGLANG_CACHE_DIT_SCM_PRESET scm_compute_bins_str = envs.SGLANG_CACHE_DIT_SCM_COMPUTE_BINS scm_cache_bins_str = envs.SGLANG_CACHE_DIT_SCM_CACHE_BINS scm_policy = envs.SGLANG_CACHE_DIT_SCM_POLICY # parse custom bins if provided (both must be set together) scm_compute_bins = None scm_cache_bins = None if scm_compute_bins_str and scm_cache_bins_str: try: scm_compute_bins = [ int(x.strip()) for x in scm_compute_bins_str.split(",") ] scm_cache_bins = [int(x.strip()) for x in scm_cache_bins_str.split(",")] except ValueError as e: logger.warning("Failed to parse SCM bins: %s. SCM disabled.", e) scm_preset = "none" elif scm_compute_bins_str or scm_cache_bins_str: # Only one of the bins was provided - warn user logger.warning( "SCM custom bins require both compute_bins and cache_bins. " "Only one was provided (compute=%s, cache=%s). Falling back to preset '%s'.", scm_compute_bins_str, scm_cache_bins_str, scm_preset, ) # generate SCM mask using cache-dit's steps_mask() # cache-dit handles step count validation and scaling internally steps_computation_mask = get_scm_mask( preset=scm_preset, num_inference_steps=( num_inference_steps if isinstance(num_inference_steps, int) else num_high_noise_steps ), compute_bins=scm_compute_bins, cache_bins=scm_cache_bins, ) if isinstance(num_inference_steps, tuple): steps_computation_mask_2 = get_scm_mask( preset=scm_preset, num_inference_steps=num_low_noise_steps, compute_bins=scm_compute_bins, cache_bins=scm_cache_bins, ) # build config for primary transformer (high-noise expert) primary_config = CacheDitConfig( enabled=True, Fn_compute_blocks=envs.SGLANG_CACHE_DIT_FN, Bn_compute_blocks=envs.SGLANG_CACHE_DIT_BN, max_warmup_steps=envs.SGLANG_CACHE_DIT_WARMUP, residual_diff_threshold=envs.SGLANG_CACHE_DIT_RDT, max_continuous_cached_steps=envs.SGLANG_CACHE_DIT_MC, enable_taylorseer=envs.SGLANG_CACHE_DIT_TAYLORSEER, taylorseer_order=envs.SGLANG_CACHE_DIT_TS_ORDER, num_inference_steps=( num_inference_steps if isinstance(num_inference_steps, int) else num_high_noise_steps ), # SCM fields steps_computation_mask=steps_computation_mask, steps_computation_policy=scm_policy, ) if self.transformer_2 is not None: # dual transformer # build config for secondary transformer (low-noise expert) # uses secondary parameters which inherit from primary if not explicitly set secondary_config = CacheDitConfig( enabled=True, Fn_compute_blocks=envs.SGLANG_CACHE_DIT_SECONDARY_FN, Bn_compute_blocks=envs.SGLANG_CACHE_DIT_SECONDARY_BN, max_warmup_steps=envs.SGLANG_CACHE_DIT_SECONDARY_WARMUP, residual_diff_threshold=envs.SGLANG_CACHE_DIT_SECONDARY_RDT, max_continuous_cached_steps=envs.SGLANG_CACHE_DIT_SECONDARY_MC, enable_taylorseer=envs.SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER, taylorseer_order=envs.SGLANG_CACHE_DIT_SECONDARY_TS_ORDER, num_inference_steps=num_low_noise_steps, # SCM fields - shared with primary steps_computation_mask=steps_computation_mask_2, steps_computation_policy=scm_policy, ) # for dual transformers, must use BlockAdapter to enable cache on both simultaneously. # Don't call enable_cache separately on each transformer. self.transformer, self.transformer_2 = enable_cache_on_dual_transformer( self.transformer, self.transformer_2, primary_config, secondary_config, model_name="wan2.2", sp_group=sp_group, tp_group=tp_group, ) logger.info( "cache-dit enabled on dual transformers (steps=%d, %d)", num_high_noise_steps, num_low_noise_steps, ) else: # single transformer self.transformer = enable_cache_on_transformer( self.transformer, primary_config, model_name="transformer", sp_group=sp_group, tp_group=tp_group, ) logger.info( "cache-dit enabled on transformer (steps=%d, Fn=%d, Bn=%d, rdt=%.3f)", num_inference_steps, envs.SGLANG_CACHE_DIT_FN, envs.SGLANG_CACHE_DIT_BN, envs.SGLANG_CACHE_DIT_RDT, ) self._cache_dit_enabled = True self._cached_num_steps = num_inference_steps @lru_cache(maxsize=8) def _build_guidance(self, batch_size, target_dtype, device, guidance_val): """Builds a guidance tensor. This method is cached.""" return ( torch.full( (batch_size,), guidance_val, dtype=target_dtype, device=device, ) * 1000.0 ) def get_or_build_guidance(self, bsz: int, dtype, device): """ Get the guidance tensor, using a cached version if available. This method retrieves a cached guidance tensor using `_build_guidance`. The caching is based on batch size, dtype, device, and the guidance value, preventing repeated tensor creation within the denoising loop. """ if self.server_args.pipeline_config.should_use_guidance: # TODO: should the guidance_scale be picked-up from sampling_params? guidance_val = self.server_args.pipeline_config.embedded_cfg_scale return self._build_guidance(bsz, dtype, device, guidance_val) else: return None @property def parallelism_type(self) -> StageParallelismType: # return StageParallelismType.CFG_PARALLEL if get_global_server_args().enable_cfg_parallel else StageParallelismType.REPLICATED return StageParallelismType.REPLICATED def _preprocess_latents_for_ti2v( self, latents, target_dtype, batch, server_args: ServerArgs ): # FIXME: should probably move to latent preparation stage, to handle with offload # Wan2.2 TI2V directly replaces the first frame of the latent with # the image latent instead of appending along the channel dim assert batch.image_latent is None, "TI2V task should not have image latents" assert self.vae is not None, "VAE is not provided for TI2V task" self.vae = self.vae.to(batch.condition_image.device) z = self.vae.encode(batch.condition_image).mean.float() if self.vae.device != "cpu" and server_args.vae_cpu_offload: self.vae = self.vae.to("cpu") if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None: if isinstance(self.vae.shift_factor, torch.Tensor): z -= self.vae.shift_factor.to(z.device, z.dtype) else: z -= self.vae.shift_factor if isinstance(self.vae.scaling_factor, torch.Tensor): z = z * self.vae.scaling_factor.to(z.device, z.dtype) else: z = z * self.vae.scaling_factor # z: [B, C, 1, H, W] latent_model_input = latents.to(target_dtype) # Keep as [B, C, T, H, W] for proper broadcasting assert latent_model_input.ndim == 5 # Create mask with proper shape [B, C, T, H, W] latent_for_mask = latent_model_input.squeeze(0) # [C, T, H, W] _, reserved_frames_masks = masks_like([latent_for_mask], zero=True) reserved_frames_mask = reserved_frames_masks[0].unsqueeze(0) # [1, C, T, H, W] # replace GLOBAL first frame with image - proper broadcasting # z: [B, C, 1, H, W], reserved_frames_mask: [1, C, T, H, W] # Both will broadcast correctly latents = ( 1.0 - reserved_frames_mask ) * z + reserved_frames_mask * latent_model_input assert latents.ndim == 5 latents = latents.to(get_local_torch_device()) batch.latents = latents F = batch.num_frames temporal_scale = ( server_args.pipeline_config.vae_config.arch_config.scale_factor_temporal ) spatial_scale = ( server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial ) patch_size = server_args.pipeline_config.dit_config.arch_config.patch_size seq_len = ( ((F - 1) // temporal_scale + 1) * (batch.height // spatial_scale) * (batch.width // spatial_scale) // (patch_size[1] * patch_size[2]) ) seq_len = int(math.ceil(seq_len / get_sp_world_size())) * get_sp_world_size() return seq_len, z, reserved_frames_masks def _postprocess_latents_for_ti2v(self, z, reserved_frames_masks, batch): rank_in_sp_group = get_sp_parallel_rank() sp_world_size = get_sp_world_size() if getattr(batch, "did_sp_shard_latents", False): # Shard z (image latent) along time dimension # z shape: [1, C, 1, H, W] - only first frame # Only rank 0 has the first frame after sharding if z.shape[2] == 1: # z is single frame, only rank 0 needs it if rank_in_sp_group == 0: z_sp = z else: # Other ranks don't have the first frame z_sp = None else: # Should not happen for TI2V z_sp = z # Shard reserved_frames_mask along time dimension to match sharded latents # reserved_frames_mask is a list from masks_like, extract reserved_frames_mask[0] first # reserved_frames_mask[0] shape: [C, T, H, W] # All ranks need their portion of reserved_frames_mask for timestep calculation if reserved_frames_masks is not None: reserved_frames_mask = reserved_frames_masks[ 0 ] # Extract tensor from list time_dim = reserved_frames_mask.shape[1] # [C, T, H, W] if time_dim > 0 and time_dim % sp_world_size == 0: reserved_frames_mask_sp_tensor = rearrange( reserved_frames_mask, "c (n t) h w -> c n t h w", n=sp_world_size, ).contiguous() reserved_frames_mask_sp_tensor = reserved_frames_mask_sp_tensor[ :, rank_in_sp_group, :, :, : ] reserved_frames_mask_sp = ( reserved_frames_mask_sp_tensor # Store as tensor, not list ) else: reserved_frames_mask_sp = reserved_frames_mask else: reserved_frames_mask_sp = None else: # SP not enabled or latents not sharded z_sp = z reserved_frames_mask_sp = ( reserved_frames_masks[0] if reserved_frames_masks is not None else None ) # Extract tensor return reserved_frames_mask_sp, z_sp def _handle_boundary_ratio( self, server_args, batch, ): """ (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert """ boundary_ratio = server_args.pipeline_config.dit_config.boundary_ratio if batch.boundary_ratio is not None: logger.info( "Overriding boundary ratio from %s to %s", boundary_ratio, batch.boundary_ratio, ) boundary_ratio = batch.boundary_ratio if boundary_ratio is not None: boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps else: boundary_timestep = None return boundary_timestep def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): """ Prepare all necessary invariant variables for the denoising loop. Returns: A dictionary containing all the prepared variables for the denoising loop. """ assert self.transformer is not None pipeline = self.pipeline() if self.pipeline else None boundary_timestep = self._handle_boundary_ratio(server_args, batch) # Get timesteps and calculate warmup steps timesteps = batch.timesteps num_inference_steps = batch.num_inference_steps num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order if self.transformer_2 is not None: assert boundary_timestep is not None, "boundary_timestep must be provided" num_high_noise_steps = (timesteps >= boundary_timestep).sum().item() num_low_noise_steps = num_inference_steps - num_high_noise_steps cache_dit_num_inference_steps = (num_high_noise_steps, num_low_noise_steps) else: cache_dit_num_inference_steps = num_inference_steps if not server_args.model_loaded["transformer"]: # FIXME: reuse more code loader = TransformerLoader() self.transformer = loader.load( server_args.model_paths["transformer"], server_args, "transformer" ) # enable cache-dit before torch.compile (delayed mounting) self._maybe_enable_cache_dit(cache_dit_num_inference_steps, batch) self._maybe_enable_torch_compile(self.transformer) if pipeline: pipeline.add_module("transformer", self.transformer) server_args.model_loaded["transformer"] = True else: self._maybe_enable_cache_dit(cache_dit_num_inference_steps, batch) # Prepare extra step kwargs for scheduler extra_step_kwargs = self.prepare_extra_func_kwargs( self.scheduler.step, {"generator": batch.generator, "eta": batch.eta}, ) # Setup precision and autocast settings target_dtype = torch.bfloat16 autocast_enabled = ( target_dtype != torch.float32 ) and not server_args.disable_autocast # Prepare image latents and embeddings for I2V generation image_embeds = batch.image_embeds if len(image_embeds) > 0: image_embeds = [ image_embed.to(target_dtype) for image_embed in image_embeds ] # Prepare STA parameters if self.attn_backend.get_enum() == AttentionBackendEnum.SLIDING_TILE_ATTN: self.prepare_sta_param(batch, server_args) # Get latents and embeddings latents = batch.latents prompt_embeds = batch.prompt_embeds # Removed Tensor truthiness assert to avoid GPU sync neg_prompt_embeds = None if batch.do_classifier_free_guidance: neg_prompt_embeds = batch.negative_prompt_embeds assert neg_prompt_embeds is not None # Removed Tensor truthiness assert to avoid GPU sync # specifically for Wan2_2_TI2V_5B_Config, not applicable for FastWan2_2_TI2V_5B_Config should_preprocess_for_wan_ti2v = ( server_args.pipeline_config.task_type == ModelTaskType.TI2V and batch.condition_image is not None and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config ) # TI2V specific preparations - before SP sharding if should_preprocess_for_wan_ti2v: seq_len, z, reserved_frames_masks = self._preprocess_latents_for_ti2v( latents, target_dtype, batch, server_args ) else: seq_len, z, reserved_frames_masks = ( None, None, None, ) # Handle sequence parallelism after TI2V processing self._preprocess_sp_latents(batch, server_args) latents = batch.latents # Shard z and reserved_frames_mask for TI2V if SP is enabled if should_preprocess_for_wan_ti2v: reserved_frames_mask_sp, z_sp = self._postprocess_latents_for_ti2v( z, reserved_frames_masks, batch ) else: reserved_frames_mask_sp, z_sp = ( reserved_frames_masks[0] if reserved_frames_masks is not None else None ), z guidance = self.get_or_build_guidance( # TODO: replace with raw_latent_shape? latents.shape[0], latents.dtype, latents.device, ) image_kwargs = self.prepare_extra_func_kwargs( getattr(self.transformer, "forward", self.transformer), { # TODO: make sure on-device "encoder_hidden_states_image": image_embeds, "mask_strategy": dict_to_3d_list(None, t_max=50, l_max=60, h_max=24), }, ) pos_cond_kwargs = self.prepare_extra_func_kwargs( getattr(self.transformer, "forward", self.transformer), { "encoder_hidden_states_2": batch.clip_embedding_pos, "encoder_attention_mask": batch.prompt_attention_mask, } | server_args.pipeline_config.prepare_pos_cond_kwargs( batch, self.device, getattr(self.transformer, "rotary_emb", None), dtype=target_dtype, ) | dict( encoder_hidden_states=server_args.pipeline_config.get_pos_prompt_embeds( batch ) ), ) if batch.do_classifier_free_guidance: neg_cond_kwargs = self.prepare_extra_func_kwargs( getattr(self.transformer, "forward", self.transformer), { "encoder_hidden_states_2": batch.clip_embedding_neg, "encoder_attention_mask": batch.negative_attention_mask, } | server_args.pipeline_config.prepare_neg_cond_kwargs( batch, self.device, getattr(self.transformer, "rotary_emb", None), dtype=target_dtype, ) | dict( encoder_hidden_states=server_args.pipeline_config.get_neg_prompt_embeds( batch ) ), ) else: neg_cond_kwargs = {} return { "extra_step_kwargs": extra_step_kwargs, "target_dtype": target_dtype, "autocast_enabled": autocast_enabled, "timesteps": timesteps, "num_inference_steps": num_inference_steps, "num_warmup_steps": num_warmup_steps, "image_kwargs": image_kwargs, "pos_cond_kwargs": pos_cond_kwargs, "neg_cond_kwargs": neg_cond_kwargs, "latents": latents, "prompt_embeds": prompt_embeds, "neg_prompt_embeds": neg_prompt_embeds, "boundary_timestep": boundary_timestep, "z": z_sp, # Use SP-sharded version # ndim == 5 "reserved_frames_mask": reserved_frames_mask_sp, # Use SP-sharded version "seq_len": seq_len, "guidance": guidance, } def _post_denoising_loop( self, batch: Req, latents: torch.Tensor, trajectory_latents: list, trajectory_timesteps: list, server_args: ServerArgs, is_warmup: bool = False, ): # Gather results if using sequence parallelism if trajectory_latents: trajectory_tensor = torch.stack(trajectory_latents, dim=1) trajectory_timesteps_tensor = torch.stack(trajectory_timesteps, dim=0) else: trajectory_tensor = None trajectory_timesteps_tensor = None # Gather results if using sequence parallelism latents, trajectory_tensor = self._postprocess_sp_latents( batch, latents, trajectory_tensor ) # Gather noise_pred if using sequence parallelism # noise_pred has the same shape as latents (sharded along sequence dimension) if ( get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False) and server_args.comfyui_mode and hasattr(batch, "noise_pred") and batch.noise_pred is not None ): batch.noise_pred = server_args.pipeline_config.gather_latents_for_sp( batch.noise_pred ) if hasattr(batch, "raw_latent_shape"): orig_s = batch.raw_latent_shape[1] if batch.noise_pred.shape[1] > orig_s: batch.noise_pred = batch.noise_pred[:, :orig_s, :] if trajectory_tensor is not None and trajectory_timesteps_tensor is not None: batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu() batch.trajectory_latents = trajectory_tensor.cpu() # Update batch with final latents batch.latents = self.server_args.pipeline_config.post_denoising_loop( latents, batch ) # Save STA mask search results if needed if ( not is_warmup and self.attn_backend.get_enum() == AttentionBackendEnum.SLIDING_TILE_ATTN and server_args.attention_backend_config.STA_mode == "STA_SEARCHING" ): self.save_sta_search_results(batch) # Capture references before potential deletion on MPS dits = list(filter(None, [self.transformer, self.transformer_2])) # deallocate transformer if on mps pipeline = self.pipeline() if self.pipeline else None if torch.backends.mps.is_available() and not is_warmup: logger.info( "Memory before deallocating transformer: %s", torch.mps.current_allocated_memory(), ) del self.transformer if pipeline is not None and "transformer" in pipeline.modules: del pipeline.modules["transformer"] server_args.model_loaded["transformer"] = False logger.info( "Memory after deallocating transformer: %s", torch.mps.current_allocated_memory(), ) # reset offload managers with prefetching first layer for next forward for dit in dits: if isinstance(dit, OffloadableDiTMixin): # release all DiT weights to avoid peak VRAM usage, which may increasing the latency for next req # TODO: should be make this an option? for manager in dit.layerwise_offload_managers: manager.release_all() def _preprocess_sp_latents(self, batch: Req, server_args: ServerArgs): """Shard latents for Sequence Parallelism if applicable.""" if get_sp_world_size() <= 1: return if batch.latents is not None: ( batch.latents, did_shard, ) = server_args.pipeline_config.shard_latents_for_sp(batch, batch.latents) batch.did_sp_shard_latents = did_shard else: batch.did_sp_shard_latents = False # image_latent must be sharded consistently with latents when it is # concatenated along the sequence dimension in the denoising loop. if batch.image_latent is not None: batch.image_latent, _ = server_args.pipeline_config.shard_latents_for_sp( batch, batch.image_latent ) def _postprocess_sp_latents( self, batch: Req, latents: torch.Tensor, trajectory_tensor: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Gather latents after Sequence Parallelism if they were sharded.""" if get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False): latents = self.server_args.pipeline_config.gather_latents_for_sp(latents) if trajectory_tensor is not None: # trajectory_tensor shapes: # - video: [b, num_steps, c, t_local, h, w] -> gather on dim=3 # - image: [b, num_steps, s_local, d] -> gather on dim=2 trajectory_tensor = trajectory_tensor.to(get_local_torch_device()) gather_dim = 3 if trajectory_tensor.dim() >= 5 else 2 trajectory_tensor = sequence_model_parallel_all_gather( trajectory_tensor, dim=gather_dim ) if gather_dim == 2 and hasattr(batch, "raw_latent_shape"): orig_s = batch.raw_latent_shape[1] if trajectory_tensor.shape[2] > orig_s: trajectory_tensor = trajectory_tensor[:, :, :orig_s, :] return latents, trajectory_tensor def step_profile(self): profiler = SGLDiffusionProfiler.get_instance() if profiler: profiler.step_denoising_step() def _manage_device_placement( self, model_to_use: nn.Module, model_to_offload: nn.Module | None, server_args: ServerArgs, ): """ Manages the offload / load behavior of dit """ if not server_args.dit_cpu_offload: return # FSDP manages offloading internally if server_args.use_fsdp_inference: return # Offload the unused model if it's on CUDA if ( model_to_offload is not None and next(model_to_offload.parameters()).device.type == "cuda" ): model_to_offload.to("cpu") # Load the model to use if it's on CPU if ( model_to_use is not None and next(model_to_use.parameters()).device.type == "cpu" ): model_to_use.to(get_local_torch_device()) def _select_and_manage_model( self, t_int: int, boundary_timestep: float | None, server_args: ServerArgs, batch: Req, ): if boundary_timestep is None or t_int >= boundary_timestep: # High-noise stage current_model = self.transformer model_to_offload = self.transformer_2 current_guidance_scale = batch.guidance_scale else: # Low-noise stage current_model = self.transformer_2 model_to_offload = self.transformer current_guidance_scale = batch.guidance_scale_2 self._manage_device_placement(current_model, model_to_offload, server_args) assert current_model is not None, "The model for the current step is not set." return current_model, current_guidance_scale def expand_timestep_before_forward( self, batch: Req, server_args: ServerArgs, t_device, target_dtype, seq_len: int | None, reserved_frames_mask, ): bsz = batch.raw_latent_shape[0] should_preprocess_for_wan_ti2v = ( server_args.pipeline_config.task_type == ModelTaskType.TI2V and batch.condition_image is not None and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config ) # expand timestep if should_preprocess_for_wan_ti2v: # Explicitly cast t_device to the target float type at the beginning. # This ensures any precision-based rounding (e.g., float32(999.0) -> bfloat16(1000.0)) # is applied consistently *before* it's used by any rank. t_device_rounded = t_device.to(target_dtype) local_seq_len = seq_len if get_sp_world_size() > 1 and getattr( batch, "did_sp_shard_latents", False ): local_seq_len = seq_len // get_sp_world_size() if get_sp_parallel_rank() == 0 and reserved_frames_mask is not None: # Rank 0 has the first frame, create a special timestep tensor # NOTE: The spatial downsampling in the next line is suspicious but kept # to match original model's potential training configuration. temp_ts = ( reserved_frames_mask[0][:, ::2, ::2] * t_device_rounded ).flatten() # Pad to full local sequence length temp_ts = torch.cat( [ temp_ts, temp_ts.new_ones(local_seq_len - temp_ts.size(0)) * t_device_rounded, ] ) timestep = temp_ts.unsqueeze(0).repeat(bsz, 1) else: # Other ranks get a uniform timestep tensor of the correct shape [B, local_seq_len] timestep = t_device.repeat(bsz, local_seq_len) else: timestep = t_device.repeat(bsz) return timestep def post_forward_for_ti2v_task( self, batch: Req, server_args: ServerArgs, reserved_frames_mask, latents, z ): """ For Wan2.2 ti2v task, global first frame should be replaced with encoded image after each timestep """ should_preprocess_for_wan_ti2v = ( server_args.pipeline_config.task_type == ModelTaskType.TI2V and batch.condition_image is not None and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config ) if should_preprocess_for_wan_ti2v: # Apply TI2V mask blending with SP-aware z and reserved_frames_mask. # This ensures the first frame is always the condition image after each step. # This is only applied on rank 0, where z is not None. if z is not None and reserved_frames_mask is not None: # z: [1, C, 1, H, W] # latents: [1, C, T_local, H, W] # reserved_frames_mask: [C, T_local, H, W] # Unsqueeze mask to [1, C, T_local, H, W] for broadcasting. # z will broadcast along the time dimension. latents = ( 1.0 - reserved_frames_mask.unsqueeze(0) ) * z + reserved_frames_mask.unsqueeze(0) * latents return latents @torch.no_grad() def forward( self, batch: Req, server_args: ServerArgs, ) -> Req: """ Run the denoising loop. """ # Prepare variables for the denoising loop prepared_vars = self._prepare_denoising_loop(batch, server_args) extra_step_kwargs = prepared_vars["extra_step_kwargs"] target_dtype = prepared_vars["target_dtype"] autocast_enabled = prepared_vars["autocast_enabled"] timesteps = prepared_vars["timesteps"] num_inference_steps = prepared_vars["num_inference_steps"] num_warmup_steps = prepared_vars["num_warmup_steps"] image_kwargs = prepared_vars["image_kwargs"] pos_cond_kwargs = prepared_vars["pos_cond_kwargs"] neg_cond_kwargs = prepared_vars["neg_cond_kwargs"] latents = prepared_vars["latents"] boundary_timestep = prepared_vars["boundary_timestep"] z = prepared_vars["z"] reserved_frames_mask = prepared_vars["reserved_frames_mask"] seq_len = prepared_vars["seq_len"] guidance = prepared_vars["guidance"] # Initialize lists for ODE trajectory trajectory_timesteps: list[torch.Tensor] = [] trajectory_latents: list[torch.Tensor] = [] # Run denoising loop denoising_start_time = time.time() # to avoid device-sync caused by timestep comparison is_warmup = batch.is_warmup self.scheduler.set_begin_index(0) timesteps_cpu = timesteps.cpu() num_timesteps = timesteps_cpu.shape[0] with torch.autocast( device_type=current_platform.device_type, dtype=target_dtype, enabled=autocast_enabled, ): with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t_host in enumerate(timesteps_cpu): with StageProfiler( f"denoising_step_{i}", logger=logger, metrics=batch.metrics, perf_dump_path_provided=batch.perf_dump_path is not None, ): t_int = int(t_host.item()) t_device = timesteps[i] current_model, current_guidance_scale = ( self._select_and_manage_model( t_int=t_int, boundary_timestep=boundary_timestep, server_args=server_args, batch=batch, ) ) # Expand latents for I2V latent_model_input = latents.to(target_dtype) if batch.image_latent is not None: assert ( not server_args.pipeline_config.task_type == ModelTaskType.TI2V ), "image latents should not be provided for TI2V task" latent_model_input = torch.cat( [latent_model_input, batch.image_latent], dim=1 ).to(target_dtype) timestep = self.expand_timestep_before_forward( batch, server_args, t_device, target_dtype, seq_len, reserved_frames_mask, ) latent_model_input = self.scheduler.scale_model_input( latent_model_input, t_device ) # Predict noise residual attn_metadata = self._build_attn_metadata( i, batch, server_args, timestep_value=t_int, timesteps=timesteps_cpu, ) noise_pred = self._predict_noise_with_cfg( current_model=current_model, latent_model_input=latent_model_input, timestep=timestep, batch=batch, timestep_index=i, attn_metadata=attn_metadata, target_dtype=target_dtype, current_guidance_scale=current_guidance_scale, image_kwargs=image_kwargs, pos_cond_kwargs=pos_cond_kwargs, neg_cond_kwargs=neg_cond_kwargs, server_args=server_args, guidance=guidance, latents=latents, ) # Save noise_pred to batch for external access (e.g., ComfyUI) if server_args.comfyui_mode: batch.noise_pred = noise_pred # Compute the previous noisy sample latents = self.scheduler.step( model_output=noise_pred, timestep=t_device, sample=latents, **extra_step_kwargs, return_dict=False, )[0] latents = self.post_forward_for_ti2v_task( batch, server_args, reserved_frames_mask, latents, z ) # save trajectory latents if needed if batch.return_trajectory_latents: trajectory_timesteps.append(t_host) trajectory_latents.append(latents) # Update progress bar if i == num_timesteps - 1 or ( (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 and progress_bar is not None ): progress_bar.update() if not is_warmup: self.step_profile() denoising_end_time = time.time() if num_timesteps > 0 and not is_warmup: self.log_info( "average time per step: %.4f seconds", (denoising_end_time - denoising_start_time) / len(timesteps), ) self._post_denoising_loop( batch=batch, latents=latents, trajectory_latents=trajectory_latents, trajectory_timesteps=trajectory_timesteps, server_args=server_args, is_warmup=is_warmup, ) return batch # TODO: this will extends the preparation stage, should let subclass/passed-in variables decide which to prepare def prepare_extra_func_kwargs(self, func, kwargs) -> dict[str, Any]: """ Prepare extra kwargs for the scheduler step / denoise step. Args: func: The function to prepare kwargs for. kwargs: The kwargs to prepare. """ import functools # Handle cache-dit's partial wrapping logic. # Cache-dit wraps the forward method with functools.partial where args[0] is the instance. # We access `_original_forward` if available to inspect the underlying signature. # See: https://github.com/vipshop/cache-dit if isinstance(func, functools.partial) and func.args: func = getattr(func.args[0], "_original_forward", func) # Unwrap any decorators (e.g. functools.wraps) target_func = inspect.unwrap(func) # Filter kwargs based on the signature params = inspect.signature(target_func).parameters return {k: v for k, v in kwargs.items() if k in params} def progress_bar( self, iterable: Iterable | None = None, total: int | None = None ) -> tqdm: """ Create a progress bar for the denoising process. """ local_rank = get_world_group().local_rank disable = local_rank != 0 return tqdm(iterable=iterable, total=total, disable=disable) def rescale_noise_cfg( self, noise_cfg, noise_pred_text, guidance_rescale=0.0 ) -> torch.Tensor: """ Rescale noise prediction according to guidance_rescale. Based on findings of "Common Diffusion Noise Schedules and Sample Steps are Flawed" (https://arxiv.org/pdf/2305.08891.pdf), Section 3.4. Args: noise_cfg: The noise prediction with guidance. noise_pred_text: The text-conditioned noise prediction. guidance_rescale: The guidance rescale factor. Returns: The rescaled noise prediction. """ std_text = noise_pred_text.std( dim=list(range(1, noise_pred_text.ndim)), keepdim=True ) std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) # Rescale the results from guidance (fixes overexposure) noise_pred_rescaled = noise_cfg * (std_text / std_cfg) # Mix with the original results from guidance by factor guidance_rescale noise_cfg = ( guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg ) return noise_cfg def _build_attn_metadata( self, i: int, batch: Req, server_args: ServerArgs, *, timestep_value: int | None = None, timesteps: torch.Tensor | None = None, ) -> Any | None: """ Build attention metadata for custom attention backends. Args: i: The current timestep index. """ attn_metadata = None self.attn_metadata_builder = None try: self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls() except NotImplementedError: self.attn_metadata_builder_cls = None if self.attn_metadata_builder_cls: self.attn_metadata_builder = self.attn_metadata_builder_cls() if ( self.attn_backend.get_enum() == AttentionBackendEnum.SLIDING_TILE_ATTN or self.attn_backend.get_enum() == AttentionBackendEnum.VIDEO_SPARSE_ATTN ): attn_metadata = self.attn_metadata_builder.build( current_timestep=i, raw_latent_shape=batch.raw_latent_shape[2:5], patch_size=server_args.pipeline_config.dit_config.patch_size, STA_param=batch.STA_param, VSA_sparsity=server_args.attention_backend_config.VSA_sparsity, device=get_local_torch_device(), ) elif ( self.attn_backend.get_enum() == AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN ): if timestep_value is None or timesteps is None: raise ValueError( "timestep_value and timesteps must be provided for SVG2 attention metadata" ) svg2_cfg = server_args.attention_backend_config or {} num_layers = server_args.pipeline_config.dit_config.num_layers if ( server_args.pipeline_config.dit_config.prefix.lower() == "hunyuan" and hasattr(server_args.pipeline_config.dit_config, "num_single_layers") ): num_layers += server_args.pipeline_config.dit_config.num_single_layers first_layers_fp = svg2_cfg.get("svg2_first_layers_fp", 0.03) if first_layers_fp <= 1.0: first_layers_fp = math.floor(first_layers_fp * num_layers) first_layers_fp = max(0, min(int(first_layers_fp), num_layers)) first_times_fp = svg2_cfg.get("svg2_first_times_fp", 0.2) if first_times_fp <= 1.0: num_fp_steps = math.floor(first_times_fp * len(timesteps)) if num_fp_steps > 0: first_times_fp = float(timesteps[num_fp_steps - 1].item() - 1) else: first_times_fp = float(timesteps.max().item() + 1) current_timestep = int(timestep_value) cache = batch.extra.get("svg2_cache") if cache is None: from sglang.multimodal_gen.runtime.layers.attention.backends.sparse_video_gen_2_attn import ( Svg2Cache, ) cache = Svg2Cache() batch.extra["svg2_cache"] = cache patch_size = server_args.pipeline_config.dit_config.patch_size if isinstance(patch_size, list): patch_size = tuple(patch_size) if isinstance(patch_size, int): patch_size_t = getattr( server_args.pipeline_config.dit_config, "patch_size_t", None ) if patch_size_t is not None: patch_size = (patch_size_t, patch_size, patch_size) context_length = 0 prompt_length = None if server_args.pipeline_config.dit_config.prefix.lower() == "hunyuan": prompt_embeds = server_args.pipeline_config.get_pos_prompt_embeds(batch) if isinstance(prompt_embeds, list): text_embeds = prompt_embeds[0] if prompt_embeds else None else: text_embeds = prompt_embeds if isinstance(text_embeds, torch.Tensor) and text_embeds.ndim >= 2: context_length = int(text_embeds.shape[1]) if context_length > 0 and batch.prompt_attention_mask: mask = batch.prompt_attention_mask[0] if isinstance(mask, torch.Tensor): if mask.shape[-1] > context_length: mask = mask[:, -context_length:] prompt_length = int(mask[0].sum().item()) if prompt_length is None: prompt_length = context_length attn_metadata = self.attn_metadata_builder.build( current_timestep=current_timestep, raw_latent_shape=batch.raw_latent_shape, patch_size=patch_size, num_q_centroids=svg2_cfg.get("svg2_num_q_centroids", 300), num_k_centroids=svg2_cfg.get("svg2_num_k_centroids", 1000), top_p_kmeans=svg2_cfg.get("svg2_top_p_kmeans", 0.9), min_kc_ratio=svg2_cfg.get("svg2_min_kc_ratio", 0.1), kmeans_iter_init=svg2_cfg.get("svg2_kmeans_iter_init", 50), kmeans_iter_step=svg2_cfg.get("svg2_kmeans_iter_step", 2), zero_step_kmeans_init=svg2_cfg.get("svg2_zero_step_kmeans_init", False), first_layers_fp=first_layers_fp, first_times_fp=first_times_fp, context_length=context_length, prompt_length=prompt_length, cache=cache, calculate_density=False, # only need density when doing head load balancing ) elif self.attn_backend.get_enum() == AttentionBackendEnum.VMOBA_ATTN: moba_params = server_args.attention_backend_config.moba_config.copy() moba_params.update( { "current_timestep": i, "raw_latent_shape": batch.raw_latent_shape[2:5], "patch_size": server_args.pipeline_config.dit_config.patch_size, "device": get_local_torch_device(), } ) elif self.attn_backend.get_enum() == AttentionBackendEnum.FA: attn_metadata = self.attn_metadata_builder.build( raw_latent_shape=batch.raw_latent_shape ) else: # attn_metadata can be None for SDPA attention backend return None return attn_metadata def _predict_noise( self, current_model, latent_model_input, timestep, target_dtype, guidance: torch.Tensor, **kwargs, ): return current_model( hidden_states=latent_model_input, timestep=timestep, guidance=guidance, **kwargs, ) def _predict_noise_with_cfg( self, current_model: nn.Module, latent_model_input: torch.Tensor, timestep, batch: Req, timestep_index: int, attn_metadata, target_dtype, current_guidance_scale, image_kwargs: dict[str, Any], pos_cond_kwargs: dict[str, Any], neg_cond_kwargs: dict[str, Any], server_args, guidance, latents, ): """ Predict the noise residual with classifier-free guidance. Args: current_model: The transformer model to use for the current step. latent_model_input: The input latents for the model. timestep: The expanded timestep tensor. batch: The current batch information. timestep_index: The current timestep index. attn_metadata: Attention metadata for custom backends. target_dtype: The target data type for autocasting. current_guidance_scale: The guidance scale for the current step. image_kwargs: Keyword arguments for image conditioning. pos_cond_kwargs: Keyword arguments for positive prompt conditioning. neg_cond_kwargs: Keyword arguments for negative prompt conditioning. Returns: The predicted noise. """ noise_pred_cond: torch.Tensor | None = None noise_pred_uncond: torch.Tensor | None = None cfg_rank = get_classifier_free_guidance_rank() # positive pass if not (server_args.enable_cfg_parallel and cfg_rank != 0): batch.is_cfg_negative = False with set_forward_context( current_timestep=timestep_index, attn_metadata=attn_metadata, forward_batch=batch, ): noise_pred_cond = self._predict_noise( current_model=current_model, latent_model_input=latent_model_input, timestep=timestep, target_dtype=target_dtype, guidance=guidance, **image_kwargs, **pos_cond_kwargs, ) # TODO: can it be moved to after _predict_noise_with_cfg? noise_pred_cond = server_args.pipeline_config.slice_noise_pred( noise_pred_cond, latents ) if not batch.do_classifier_free_guidance: # If CFG is disabled, we are done. Return the conditional prediction. return noise_pred_cond # negative pass if not server_args.enable_cfg_parallel or cfg_rank != 0: batch.is_cfg_negative = True with set_forward_context( current_timestep=timestep_index, attn_metadata=attn_metadata, forward_batch=batch, ): noise_pred_uncond = self._predict_noise( current_model=current_model, latent_model_input=latent_model_input, timestep=timestep, target_dtype=target_dtype, guidance=guidance, **image_kwargs, **neg_cond_kwargs, ) noise_pred_uncond = server_args.pipeline_config.slice_noise_pred( noise_pred_uncond, latents ) # Combine predictions if server_args.enable_cfg_parallel: # Each rank computes its partial contribution and we sum via all-reduce: # final = s*cond + (1-s)*uncond if cfg_rank == 0: assert noise_pred_cond is not None partial = current_guidance_scale * noise_pred_cond else: assert noise_pred_uncond is not None partial = (1 - current_guidance_scale) * noise_pred_uncond noise_pred = cfg_model_parallel_all_reduce(partial) if batch.cfg_normalization and float(batch.cfg_normalization) > 0: factor = float(batch.cfg_normalization) pred_f = noise_pred.float() new_norm = torch.linalg.vector_norm(pred_f) if cfg_rank == 0: cond_f = noise_pred_cond.float() ori_norm = torch.linalg.vector_norm(cond_f) else: ori_norm = torch.empty_like(new_norm) ori_norm = get_cfg_group().broadcast(ori_norm, src=0) max_norm = ori_norm * factor if new_norm > max_norm: noise_pred = noise_pred * (max_norm / new_norm) # Guidance rescale: broadcast std(cond) from rank 0, compute std(cfg) locally if batch.guidance_rescale > 0.0: std_cfg = noise_pred.std( dim=list(range(1, noise_pred.ndim)), keepdim=True ) if cfg_rank == 0: assert noise_pred_cond is not None std_text = noise_pred_cond.std( dim=list(range(1, noise_pred_cond.ndim)), keepdim=True ) else: std_text = torch.empty_like(std_cfg) # Broadcast std_text from local src=0 to all ranks in CFG group std_text = get_cfg_group().broadcast(std_text, src=0) noise_pred_rescaled = noise_pred * (std_text / std_cfg) noise_pred = ( batch.guidance_rescale * noise_pred_rescaled + (1 - batch.guidance_rescale) * noise_pred ) return noise_pred else: # Serial CFG: both cond and uncond are available locally assert noise_pred_cond is not None and noise_pred_uncond is not None noise_pred = noise_pred_uncond + current_guidance_scale * ( noise_pred_cond - noise_pred_uncond ) if batch.cfg_normalization and float(batch.cfg_normalization) > 0: factor = float(batch.cfg_normalization) cond_f = noise_pred_cond.float() pred_f = noise_pred.float() ori_norm = torch.linalg.vector_norm(cond_f) new_norm = torch.linalg.vector_norm(pred_f) max_norm = ori_norm * factor if new_norm > max_norm: noise_pred = noise_pred * (max_norm / new_norm) if batch.guidance_rescale > 0.0: noise_pred = self.rescale_noise_cfg( noise_pred, noise_pred_cond, guidance_rescale=batch.guidance_rescale, ) return noise_pred def prepare_sta_param(self, batch: Req, server_args: ServerArgs): """ Prepare Sliding Tile Attention (STA) parameters and settings. """ # TODO(kevin): STA mask search, currently only support Wan2.1 with 69x768x1280 try: STA_mode = STA_Mode[server_args.attention_backend_config.STA_mode] except Exception as e: logger.error(f"Passed STA_mode: {STA_mode} doesn't exist") raise e skip_time_steps = server_args.attention_backend_config.skip_time_steps if batch.timesteps is None: raise ValueError("Timesteps must be provided") timesteps_num = batch.timesteps.shape[0] logger.info("STA_mode: %s", STA_mode) if (batch.num_frames, batch.height, batch.width) != ( 69, 768, 1280, ) and STA_mode != "STA_inference": raise NotImplementedError( "STA mask search/tuning is not supported for this resolution" ) if ( STA_mode == STA_Mode.STA_SEARCHING or STA_mode == STA_Mode.STA_TUNING or STA_mode == STA_Mode.STA_TUNING_CFG ): size = (batch.width, batch.height) if size == (1280, 768): # TODO: make it configurable sparse_mask_candidates_searching = [ "3, 1, 10", "1, 5, 7", "3, 3, 3", "1, 6, 5", "1, 3, 10", "3, 6, 1", ] sparse_mask_candidates_tuning = [ "3, 1, 10", "1, 5, 7", "3, 3, 3", "1, 6, 5", "1, 3, 10", "3, 6, 1", ] full_mask = ["3,6,10"] else: raise NotImplementedError( "STA mask search is not supported for this resolution" ) layer_num = self.transformer.config.num_layers # specific for HunyuanVideo if hasattr(self.transformer.config, "num_single_layers"): layer_num += self.transformer.config.num_single_layers head_num = self.transformer.config.num_attention_heads if STA_mode == STA_Mode.STA_SEARCHING: STA_param = configure_sta( mode=STA_Mode.STA_SEARCHING, layer_num=layer_num, head_num=head_num, time_step_num=timesteps_num, mask_candidates=sparse_mask_candidates_searching + full_mask, # last is full mask; Can add more sparse masks while keep last one as full mask ) elif STA_mode == STA_Mode.STA_TUNING: STA_param = configure_sta( mode=STA_Mode.STA_TUNING, layer_num=layer_num, head_num=head_num, time_step_num=timesteps_num, mask_search_files_path=f"output/mask_search_result_pos_{size[0]}x{size[1]}/", mask_candidates=sparse_mask_candidates_tuning, full_attention_mask=[int(x) for x in full_mask[0].split(",")], skip_time_steps=skip_time_steps, # Use full attention for first 12 steps save_dir=f"output/mask_search_strategy_{size[0]}x{size[1]}/", # Custom save directory timesteps=timesteps_num, ) elif STA_mode == STA_Mode.STA_TUNING_CFG: STA_param = configure_sta( mode=STA_Mode.STA_TUNING_CFG, layer_num=layer_num, head_num=head_num, time_step_num=timesteps_num, mask_search_files_path_pos=f"output/mask_search_result_pos_{size[0]}x{size[1]}/", mask_search_files_path_neg=f"output/mask_search_result_neg_{size[0]}x{size[1]}/", mask_candidates=sparse_mask_candidates_tuning, full_attention_mask=[int(x) for x in full_mask[0].split(",")], skip_time_steps=skip_time_steps, save_dir=f"output/mask_search_strategy_{size[0]}x{size[1]}/", timesteps=timesteps_num, ) elif STA_mode == STA_Mode.STA_INFERENCE: import sglang.multimodal_gen.envs as envs config_file = envs.SGLANG_DIFFUSION_ATTENTION_CONFIG if config_file is None: raise ValueError("SGLANG_DIFFUSION_ATTENTION_CONFIG is not set") STA_param = configure_sta( mode=STA_Mode.STA_INFERENCE, layer_num=layer_num, head_num=head_num, time_step_num=timesteps_num, load_path=config_file, ) batch.STA_param = STA_param batch.mask_search_final_result_pos = [[] for _ in range(timesteps_num)] batch.mask_search_final_result_neg = [[] for _ in range(timesteps_num)] def save_sta_search_results(self, batch: Req): """ Save the STA mask search results. Args: batch: The current batch information. """ size = (batch.width, batch.height) if size == (1280, 768): # TODO: make it configurable sparse_mask_candidates_searching = [ "3, 1, 10", "1, 5, 7", "3, 3, 3", "1, 6, 5", "1, 3, 10", "3, 6, 1", ] else: raise NotImplementedError( "STA mask search is not supported for this resolution" ) if batch.mask_search_final_result_pos is not None and batch.prompt is not None: save_mask_search_results( [dict(layer_data) for layer_data in batch.mask_search_final_result_pos], prompt=str(batch.prompt), mask_strategies=sparse_mask_candidates_searching, output_dir=f"output/mask_search_result_pos_{size[0]}x{size[1]}/", ) if batch.mask_search_final_result_neg is not None and batch.prompt is not None: save_mask_search_results( [dict(layer_data) for layer_data in batch.mask_search_final_result_neg], prompt=str(batch.prompt), mask_strategies=sparse_mask_candidates_searching, output_dir=f"output/mask_search_result_neg_{size[0]}x{size[1]}/", ) def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify denoising stage inputs.""" result = VerificationResult() result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)]) # disable temporarily for image-generation models # result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) result.add_check("image_embeds", batch.image_embeds, V.is_list) # result.add_check( # "image_latent", batch.image_latent, V.none_or_tensor_with_dims(5) # ) result.add_check( "num_inference_steps", batch.num_inference_steps, V.positive_int ) result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float) result.add_check("eta", batch.eta, V.non_negative_float) result.add_check("generator", batch.generator, V.generator_or_list_generators) result.add_check( "do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value, ) result.add_check( "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x), ) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify denoising stage outputs.""" result = VerificationResult() # result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) return result ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py ================================================ import copy import math import time from io import BytesIO import av import numpy as np import PIL.Image import torch from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution from diffusers.models.modeling_outputs import AutoencoderKLOutput from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.models.vision_utils import ( load_image, normalize, numpy_to_pt, pil_to_numpy, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) class LTX2AVDenoisingStage(DenoisingStage): """ LTX-2 specific denoising stage that handles joint video and audio generation. """ def __init__(self, transformer, scheduler, vae=None, audio_vae=None, **kwargs): super().__init__( transformer=transformer, scheduler=scheduler, vae=vae, **kwargs ) self.audio_vae = audio_vae @staticmethod def _get_video_latent_num_frames_for_model( batch: Req, server_args: ServerArgs, latents: torch.Tensor ) -> int: """Return the latent-frame length the DiT model should see. - If video latents were time-sharded for SP and are packed as token latents ([B, S, D]), the model only sees the local shard and must use the local latent-frame count (stored on the batch during SP sharding). - Otherwise, fall back to the global latent-frame count inferred from the requested output frames and the VAE temporal compression ratio. """ did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) is_token_latents = isinstance(latents, torch.Tensor) and latents.ndim == 3 if did_sp_shard and is_token_latents: if not hasattr(batch, "sp_video_latent_num_frames"): raise ValueError( "SP-sharded LTX2 token latents require `batch.sp_video_latent_num_frames` " "to be set by `LTX2PipelineConfig.shard_latents_for_sp()`." ) return int(batch.sp_video_latent_num_frames) pc = server_args.pipeline_config return int((batch.num_frames - 1) // int(pc.vae_temporal_compression) + 1) @staticmethod def _truncate_sp_padded_token_latents( batch: Req, latents: torch.Tensor ) -> torch.Tensor: """Remove token padding introduced by SP time-sharding (if applicable).""" did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) if not did_sp_shard or not ( isinstance(latents, torch.Tensor) and latents.ndim == 3 ): return latents raw_shape = getattr(batch, "raw_latent_shape", None) if not (isinstance(raw_shape, tuple) and len(raw_shape) == 3): return latents orig_s = int(raw_shape[1]) cur_s = int(latents.shape[1]) if cur_s == orig_s: return latents if cur_s < orig_s: raise ValueError( f"Unexpected gathered token-latents seq_len {cur_s} < original seq_len {orig_s}." ) return latents[:, :orig_s, :].contiguous() def _maybe_enable_cache_dit(self, num_inference_steps: int, batch: Req) -> None: """Disable cache-dit for TI2V-style requests (image-conditioned), to avoid stale activations. NOTE: base denoising stage calls this hook with (num_inference_steps, batch). """ if getattr(self, "_disable_cache_dit_for_request", False): return return super()._maybe_enable_cache_dit(num_inference_steps, batch) @staticmethod def _resize_center_crop( img: PIL.Image.Image, *, width: int, height: int ) -> PIL.Image.Image: return img.resize((width, height), resample=PIL.Image.Resampling.BILINEAR) @staticmethod def _apply_video_codec_compression( img_array: np.ndarray, crf: int = 33 ) -> np.ndarray: """Encode as a single H.264 frame and decode back to simulate compression artifacts.""" if crf == 0: return img_array height, width = img_array.shape[0] // 2 * 2, img_array.shape[1] // 2 * 2 img_array = img_array[:height, :width] buffer = BytesIO() container = av.open(buffer, mode="w", format="mp4") stream = container.add_stream( "libx264", rate=1, options={"crf": str(crf), "preset": "veryfast"} ) stream.height, stream.width = height, width frame = av.VideoFrame.from_ndarray(img_array, format="rgb24").reformat( format="yuv420p" ) container.mux(stream.encode(frame)) container.mux(stream.encode()) container.close() buffer.seek(0) container = av.open(buffer) decoded = next(container.decode(container.streams.video[0])) container.close() return decoded.to_ndarray(format="rgb24") @staticmethod def _resize_center_crop_tensor( img: PIL.Image.Image, *, width: int, height: int, device: torch.device, dtype: torch.dtype, apply_codec_compression: bool = True, codec_crf: int = 33, ) -> torch.Tensor: """Resize, center-crop, and normalize to [1, C, 1, H, W] tensor in [-1, 1].""" img_array = np.array(img).astype(np.uint8)[..., :3] if apply_codec_compression: img_array = LTX2AVDenoisingStage._apply_video_codec_compression( img_array, crf=codec_crf ) tensor = ( torch.from_numpy(img_array.astype(np.float32)) .permute(2, 0, 1) .unsqueeze(0) .to(device=device) ) src_h, src_w = tensor.shape[2], tensor.shape[3] scale = max(height / src_h, width / src_w) new_h, new_w = math.ceil(src_h * scale), math.ceil(src_w * scale) tensor = torch.nn.functional.interpolate( tensor, size=(new_h, new_w), mode="bilinear", align_corners=False ) top, left = (new_h - height) // 2, (new_w - width) // 2 tensor = tensor[:, :, top : top + height, left : left + width] return ((tensor / 127.5 - 1.0).to(dtype=dtype)).unsqueeze(2) @staticmethod def _pil_to_normed_tensor(img: PIL.Image.Image) -> torch.Tensor: # PIL -> numpy [0,1] -> torch [B,C,H,W], then [-1,1] arr = pil_to_numpy(img) t = numpy_to_pt(arr) return normalize(t) @staticmethod def _should_apply_ltx2_ti2v(batch: Req) -> bool: """True if we have an image-latent token prefix to condition with. SP note: when token latents are time-sharded, only the rank that owns the *global* first latent frame should apply TI2V conditioning (rank with start_frame==0). """ if ( batch.image_latent is None or int(getattr(batch, "ltx2_num_image_tokens", 0)) <= 0 ): return False did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) if not did_sp_shard: return True return int(getattr(batch, "sp_video_start_frame", 0)) == 0 def _prepare_ltx2_image_latent(self, batch: Req, server_args: ServerArgs) -> None: """Encode `batch.image_path` into packed token latents for LTX-2 TI2V.""" if ( batch.image_latent is not None and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0 ): return batch.ltx2_num_image_tokens = 0 batch.image_latent = None if batch.image_path is None: return if batch.width is None or batch.height is None: raise ValueError("width/height must be provided for LTX-2 TI2V.") if self.vae is None: raise ValueError("VAE must be provided for LTX-2 TI2V.") image_path = ( batch.image_path[0] if isinstance(batch.image_path, list) else batch.image_path ) img = load_image(image_path) batch.condition_image = self._resize_center_crop( img, width=int(batch.width), height=int(batch.height) ) latents_device = ( batch.latents.device if isinstance(batch.latents, torch.Tensor) else torch.device("cpu") ) encode_dtype = batch.latents.dtype original_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] self.vae = self.vae.to(device=latents_device, dtype=encode_dtype) vae_autocast_enabled = ( original_dtype != torch.float32 ) and not server_args.disable_autocast video_condition = self._resize_center_crop_tensor( img, width=int(batch.width), height=int(batch.height), device=latents_device, dtype=encode_dtype, ) with torch.autocast( device_type=current_platform.device_type, dtype=original_dtype, enabled=vae_autocast_enabled, ): try: if server_args.pipeline_config.vae_tiling: self.vae.enable_tiling() except Exception: pass if not vae_autocast_enabled: video_condition = video_condition.to(encode_dtype) latent_dist: DiagonalGaussianDistribution = self.vae.encode(video_condition) if isinstance(latent_dist, AutoencoderKLOutput): latent_dist = latent_dist.latent_dist mode = server_args.pipeline_config.vae_config.encode_sample_mode() if mode == "argmax": latent = latent_dist.mode() elif mode == "sample": if batch.generator is None: raise ValueError("Generator must be provided for VAE sampling.") latent = latent_dist.sample(batch.generator) else: raise ValueError(f"Unsupported encode_sample_mode: {mode}") # Per-channel normalization: normalized = (x - mean) / std mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latent) std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latent) latent = (latent - mean) / std packed = server_args.pipeline_config.maybe_pack_latents( latent, latent.shape[0], batch ) if not (isinstance(packed, torch.Tensor) and packed.ndim == 3): raise ValueError("Expected packed image latents [B, S0, D].") # Fail-fast token count: must match one latent frame's tokens. vae_sf = int(server_args.pipeline_config.vae_scale_factor) patch = int(server_args.pipeline_config.patch_size) latent_h = int(batch.height) // vae_sf latent_w = int(batch.width) // vae_sf expected_tokens = (latent_h // patch) * (latent_w // patch) if int(packed.shape[1]) != int(expected_tokens): raise ValueError( "LTX-2 conditioning token count mismatch: " f"{int(packed.shape[1])=} {int(expected_tokens)=}." ) batch.image_latent = packed batch.ltx2_num_image_tokens = int(packed.shape[1]) if batch.debug: logger.info( "LTX2 TI2V conditioning prepared: %d tokens (shape=%s) for %sx%s", batch.ltx2_num_image_tokens, tuple(batch.image_latent.shape), batch.width, batch.height, ) self.vae.to(original_dtype) if server_args.vae_cpu_offload: self.vae = self.vae.to("cpu") @torch.no_grad() def forward(self, batch: Req, server_args: ServerArgs) -> Req: """ Run the denoising loop. Args: batch: The current batch information. server_args: The inference arguments. Returns: The batch with denoised latents. """ # Disable cache-dit for image-conditioned requests (TI2V-style) for correctness/debuggability. self._disable_cache_dit_for_request = batch.image_path is not None # Prepare variables for the denoising loop prepared_vars = self._prepare_denoising_loop(batch, server_args) extra_step_kwargs = prepared_vars["extra_step_kwargs"] target_dtype = prepared_vars["target_dtype"] autocast_enabled = prepared_vars["autocast_enabled"] timesteps = prepared_vars["timesteps"] num_inference_steps = prepared_vars["num_inference_steps"] num_warmup_steps = prepared_vars["num_warmup_steps"] image_kwargs = prepared_vars["image_kwargs"] pos_cond_kwargs = prepared_vars["pos_cond_kwargs"] neg_cond_kwargs = prepared_vars["neg_cond_kwargs"] latents = prepared_vars["latents"] boundary_timestep = prepared_vars["boundary_timestep"] z = prepared_vars["z"] reserved_frames_mask = prepared_vars["reserved_frames_mask"] seq_len = prepared_vars["seq_len"] guidance = prepared_vars["guidance"] audio_latents = batch.audio_latents audio_scheduler = copy.deepcopy(self.scheduler) # Prepare TI2V conditioning once (encode image -> patchify tokens). self._prepare_ltx2_image_latent(batch, server_args) # For LTX-2 packed token latents, SP sharding happens on the time dimension # (frames). The model must see local latent frames (RoPE offset is applied # inside the model using SP rank). latent_num_frames_for_model = self._get_video_latent_num_frames_for_model( batch=batch, server_args=server_args, latents=latents ) latent_height = batch.height // server_args.pipeline_config.vae_scale_factor latent_width = batch.width // server_args.pipeline_config.vae_scale_factor # Initialize lists for ODE trajectory trajectory_timesteps: list[torch.Tensor] = [] trajectory_latents: list[torch.Tensor] = [] trajectory_audio_latents: list[torch.Tensor] = [] # Run denoising loop denoising_start_time = time.time() # to avoid device-sync caused by timestep comparison is_warmup = batch.is_warmup self.scheduler.set_begin_index(0) audio_scheduler.set_begin_index(0) timesteps_cpu = timesteps.cpu() num_timesteps = timesteps_cpu.shape[0] do_ti2v = self._should_apply_ltx2_ti2v(batch) num_img_tokens = int(getattr(batch, "ltx2_num_image_tokens", 0)) denoise_mask = None clean_latent = None if do_ti2v: if not (isinstance(latents, torch.Tensor) and latents.ndim == 3): raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].") latents[:, :num_img_tokens, :] = batch.image_latent[ :, :num_img_tokens, : ].to(device=latents.device, dtype=latents.dtype) denoise_mask = torch.ones( (latents.shape[0], latents.shape[1], 1), device=latents.device, dtype=torch.float32, ) denoise_mask[:, :num_img_tokens, :] = 0.0 clean_latent = latents.detach().clone() clean_latent[:, :num_img_tokens, :] = batch.image_latent[ :, :num_img_tokens, : ].to(device=latents.device, dtype=latents.dtype) with torch.autocast( device_type=current_platform.device_type, dtype=target_dtype, enabled=autocast_enabled, ): with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t_host in enumerate(timesteps_cpu): with StageProfiler( f"denoising_step_{i}", logger=logger, metrics=batch.metrics, perf_dump_path_provided=batch.perf_dump_path is not None, ): t_int = int(t_host.item()) t_device = timesteps[i] current_model, current_guidance_scale = ( self._select_and_manage_model( t_int=t_int, boundary_timestep=boundary_timestep, server_args=server_args, batch=batch, ) ) # Predict noise residual attn_metadata = self._build_attn_metadata(i, batch, server_args) # === LTX-2 sigma-space Euler step (flow matching) === # Use scheduler-generated sigmas (includes terminal sigma=0). sigmas = getattr(self.scheduler, "sigmas", None) if sigmas is None or not isinstance(sigmas, torch.Tensor): raise ValueError( "Expected scheduler.sigmas to be a tensor for LTX-2." ) sigma = sigmas[i].to(device=latents.device, dtype=torch.float32) sigma_next = sigmas[i + 1].to( device=latents.device, dtype=torch.float32 ) dt = sigma_next - sigma latent_model_input = latents.to(target_dtype) audio_latent_model_input = audio_latents.to(target_dtype) latent_num_frames = latent_num_frames_for_model # Audio latent dims if audio_latent_model_input.ndim == 3: audio_num_frames_latent = int( audio_latent_model_input.shape[1] ) elif audio_latent_model_input.ndim == 4: audio_num_frames_latent = int( audio_latent_model_input.shape[2] ) else: raise ValueError( f"Unexpected audio latents rank: {audio_latent_model_input.ndim}, shape={tuple(audio_latent_model_input.shape)}" ) # LTX-2 model can generate coords internally. video_coords = None audio_coords = None timestep = t_device.expand(int(latent_model_input.shape[0])) if do_ti2v and denoise_mask is not None: timestep_video = timestep.unsqueeze( -1 ) * denoise_mask.squeeze(-1) else: timestep_video = timestep timestep_audio = timestep # Conditions encoder_hidden_states = batch.prompt_embeds[0] audio_encoder_hidden_states = batch.audio_prompt_embeds[0] encoder_attention_mask = batch.prompt_attention_mask # Follow ltx-pipelines structure: separate pos/neg forward passes, # then apply CFG on denoised (x0) predictions. with set_forward_context( current_timestep=i, attn_metadata=attn_metadata ): v_pos, a_v_pos = current_model( hidden_states=latent_model_input, audio_hidden_states=audio_latent_model_input, encoder_hidden_states=encoder_hidden_states, audio_encoder_hidden_states=audio_encoder_hidden_states, timestep=timestep_video, audio_timestep=timestep_audio, encoder_attention_mask=encoder_attention_mask, audio_encoder_attention_mask=encoder_attention_mask, num_frames=latent_num_frames, height=latent_height, width=latent_width, fps=batch.fps, audio_num_frames=audio_num_frames_latent, video_coords=video_coords, audio_coords=audio_coords, return_latents=False, return_dict=False, ) if batch.do_classifier_free_guidance: neg_encoder_hidden_states = ( batch.negative_prompt_embeds[0] ) neg_audio_encoder_hidden_states = ( batch.negative_audio_prompt_embeds[0] ) neg_encoder_attention_mask = ( batch.negative_attention_mask ) v_neg, a_v_neg = current_model( hidden_states=latent_model_input, audio_hidden_states=audio_latent_model_input, encoder_hidden_states=neg_encoder_hidden_states, audio_encoder_hidden_states=neg_audio_encoder_hidden_states, timestep=timestep_video, audio_timestep=timestep_audio, encoder_attention_mask=neg_encoder_attention_mask, audio_encoder_attention_mask=neg_encoder_attention_mask, num_frames=latent_num_frames, height=latent_height, width=latent_width, fps=batch.fps, audio_num_frames=audio_num_frames_latent, video_coords=video_coords, audio_coords=audio_coords, return_latents=False, return_dict=False, ) else: v_neg = None a_v_neg = None v_pos = v_pos.float() a_v_pos = a_v_pos.float() if v_neg is not None: v_neg = v_neg.float() if a_v_neg is not None: a_v_neg = a_v_neg.float() # Velocity -> denoised (x0): x0 = x - sigma * v sigma_val = float(sigma.item()) denoised_video = (latents.float() - sigma_val * v_pos).to( latents.dtype ) denoised_audio = ( audio_latents.float() - sigma_val * a_v_pos ).to(audio_latents.dtype) if ( batch.do_classifier_free_guidance and v_neg is not None and a_v_neg is not None ): denoised_video_neg = ( latents.float() - sigma_val * v_neg ).to(latents.dtype) denoised_audio_neg = ( audio_latents.float() - sigma_val * a_v_neg ).to(audio_latents.dtype) denoised_video = denoised_video + ( batch.guidance_scale - 1.0 ) * (denoised_video - denoised_video_neg) denoised_audio = denoised_audio + ( batch.guidance_scale - 1.0 ) * (denoised_audio - denoised_audio_neg) # Apply conditioning mask (keep conditioned tokens clean). if ( do_ti2v and denoise_mask is not None and clean_latent is not None ): denoised_video = ( denoised_video * denoise_mask + clean_latent.float() * (1.0 - denoise_mask) ) # Euler step in sigma space: x_next = x + (sigma_next - sigma) * v, # where v = (x - x0) / sigma. if sigma_val == 0.0: v_video = torch.zeros_like(denoised_video) v_audio = torch.zeros_like(denoised_audio) else: v_video = ( (latents.float() - denoised_video.float()) / sigma_val ).to(latents.dtype) v_audio = ( (audio_latents.float() - denoised_audio.float()) / sigma_val ).to(audio_latents.dtype) latents = (latents.float() + v_video.float() * dt).to( dtype=latents.dtype ) audio_latents = ( audio_latents.float() + v_audio.float() * dt ).to(dtype=audio_latents.dtype) if do_ti2v: latents[:, :num_img_tokens, :] = batch.image_latent[ :, :num_img_tokens, : ].to(device=latents.device, dtype=latents.dtype) latents = self.post_forward_for_ti2v_task( batch, server_args, reserved_frames_mask, latents, z ) # save trajectory latents if needed if batch.return_trajectory_latents: trajectory_timesteps.append(t_host) trajectory_latents.append(latents) if audio_latents is not None: trajectory_audio_latents.append(audio_latents) # Update progress bar if i == num_timesteps - 1 or ( (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 and progress_bar is not None ): progress_bar.update() if not is_warmup: self.step_profile() denoising_end_time = time.time() if num_timesteps > 0 and not is_warmup: self.log_info( "average time per step: %.4f seconds", (denoising_end_time - denoising_start_time) / len(timesteps), ) batch.audio_latents = audio_latents self._post_denoising_loop( batch=batch, latents=latents, trajectory_latents=trajectory_latents, trajectory_timesteps=trajectory_timesteps, trajectory_audio_latents=trajectory_audio_latents, server_args=server_args, is_warmup=is_warmup, ) return batch def _post_denoising_loop( self, batch: Req, latents: torch.Tensor, trajectory_latents: list, trajectory_timesteps: list, trajectory_audio_latents: list, server_args: ServerArgs, is_warmup: bool = False, ): # 1. Handle Trajectory (Video) - Copy from base if trajectory_latents: trajectory_tensor = torch.stack(trajectory_latents, dim=1) trajectory_timesteps_tensor = torch.stack(trajectory_timesteps, dim=0) else: trajectory_tensor = None trajectory_timesteps_tensor = None latents, trajectory_tensor = self._postprocess_sp_latents( batch, latents, trajectory_tensor ) # If SP time-sharding padded whole frames worth of tokens, remove padding # after gather and before unpacking. latents = self._truncate_sp_padded_token_latents(batch, latents) if trajectory_tensor is not None and trajectory_timesteps_tensor is not None: batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu() batch.trajectory_latents = trajectory_tensor.cpu() # 2. Handle Trajectory (Audio) - LTX-2 specific if trajectory_audio_latents: trajectory_audio_tensor = torch.stack(trajectory_audio_latents, dim=1) # We don't have SP support for audio latents yet (or needed?) batch.trajectory_audio_latents = trajectory_audio_tensor.cpu() # 3. Unpack and Denormalize # Call pipeline_config._unpad_and_unpack_latents # latents is video latents. # batch.audio_latents is audio latents. audio_latents = batch.audio_latents # NOTE: self.vae and self.audio_vae should be populated via __init__ or manual setting if self.vae is None or self.audio_vae is None: logger.warning( "VAE or Audio VAE not found in DenoisingStage. Skipping unpack and denormalize." ) batch.latents = latents batch.audio_latents = audio_latents else: latents, audio_latents = ( server_args.pipeline_config._unpad_and_unpack_latents( latents, audio_latents, batch, self.vae, self.audio_vae ) ) batch.latents = latents batch.audio_latents = audio_latents if isinstance(self.transformer, OffloadableDiTMixin): for manager in self.transformer.layerwise_offload_managers: manager.release_all() def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify denoising stage inputs. Note: LTX-2 connector stage converts `prompt_embeds`/`negative_prompt_embeds` from list-of-tensors to a single tensor (video context) and stores audio context separately. """ result = VerificationResult() result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)]) # LTX-2 may carry prompt embeddings as either a tensor (preferred) or legacy list. result.add_check( "prompt_embeds", batch.prompt_embeds, lambda x: V.is_tensor(x) or V.list_not_empty(x), ) # Keep base expectation: image_embeds is always a list (may be empty). result.add_check("image_embeds", batch.image_embeds, V.is_list) result.add_check( "num_inference_steps", batch.num_inference_steps, V.positive_int ) result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float) result.add_check("eta", batch.eta, V.non_negative_float) result.add_check("generator", batch.generator, V.generator_or_list_generators) result.add_check( "do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value, ) # When CFG is enabled, negative prompt embeddings must exist (tensor or legacy list). result.add_check( "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x: (not batch.do_classifier_free_guidance) or V.is_tensor(x) or V.list_not_empty(x), ) return result def do_classifier_free_guidance(self, batch: Req) -> bool: return batch.guidance_scale > 1.0 class LTX2RefinementStage(LTX2AVDenoisingStage): def __init__( self, transformer, scheduler, distilled_sigmas, vae=None, audio_vae=None ): super().__init__(transformer, scheduler, vae, audio_vae) self.distilled_sigmas = torch.tensor(distilled_sigmas) def forward(self, batch: Req, server_args: ServerArgs) -> Req: # 1. Add noise to latents noise_scale = self.distilled_sigmas[0].to(batch.latents.device) noise = torch.randn_like(batch.latents) batch.latents = batch.latents + noise * noise_scale # 2. Run denoising loop with distilled_sigmas # Save original sigmas original_sigmas = self.scheduler.sigmas original_timesteps = self.scheduler.timesteps original_num_inference_steps = self.scheduler.num_inference_steps # Set distilled sigmas self.scheduler.sigmas = self.distilled_sigmas.to(self.scheduler.sigmas.device) # Approximation for timesteps self.scheduler.timesteps = self.scheduler.sigmas * 1000 self.scheduler.num_inference_steps = len(self.distilled_sigmas) - 1 # Call parent forward try: batch = super().forward(batch, server_args) finally: # Restore original sigmas self.scheduler.sigmas = original_sigmas self.scheduler.timesteps = original_timesteps self.scheduler.num_inference_steps = original_num_inference_steps return batch def do_classifier_free_guidance(self, batch: Req) -> bool: return False # Stage 2 uses simple denoising (no CFG) ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import time import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) from sglang.multimodal_gen.runtime.models.utils import pred_noise_to_pred_video from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages import DenoisingStage from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.utils import dict_to_3d_list logger = init_logger(__name__) class DmdDenoisingStage(DenoisingStage): """ Denoising stage for DMD. """ def __init__(self, transformer, scheduler, transformer_2=None) -> None: super().__init__( transformer=transformer, scheduler=scheduler, transformer_2=transformer_2 ) self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0) def _preprocess_sp_latents(self, batch: Req, server_args: ServerArgs): # 1. to shard latents (B, C, T, H, W) along dim 2 super()._preprocess_sp_latents(batch, server_args) # 2. DMD expects (B, T, C, H, W) for the main latents in the loop if batch.latents is not None: batch.latents = batch.latents.permute(0, 2, 1, 3, 4) # Note: batch.image_latent is kept as (B, C, T, H, W) here def _postprocess_sp_latents( self, batch: Req, latents: torch.Tensor, trajectory_tensor: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor | None]: # 1. convert back from DMD's (B, T, C, H, W) to standard (B, C, T, H, W) # this is because base gather_latents_for_sp expects dim=2 for T latents = latents.permute(0, 2, 1, 3, 4) # 2. use base method to gather return super()._postprocess_sp_latents(batch, latents, trajectory_tensor) def forward( self, batch: Req, server_args: ServerArgs, ) -> Req: """ Run the denoising loop. """ prepared_vars = self._prepare_denoising_loop(batch, server_args) target_dtype = prepared_vars["target_dtype"] autocast_enabled = prepared_vars["autocast_enabled"] num_warmup_steps = prepared_vars["num_warmup_steps"] latents = prepared_vars["latents"] video_raw_latent_shape = latents.shape timesteps = torch.tensor( server_args.pipeline_config.dmd_denoising_steps, dtype=torch.long, device=get_local_torch_device(), ) # prepare image_kwargs image_embeds = batch.image_embeds if len(image_embeds) > 0: image_embeds = [img.to(target_dtype) for img in image_embeds] image_kwargs = self.prepare_extra_func_kwargs( self.transformer.forward, { "encoder_hidden_states_image": image_embeds, "mask_strategy": dict_to_3d_list(None, t_max=50, l_max=60, h_max=24), }, ) pos_cond_kwargs = prepared_vars["pos_cond_kwargs"] denoising_loop_start_time = time.time() with self.progress_bar(total=len(timesteps)) as progress_bar: for i, t in enumerate(timesteps): # Skip if interrupted if hasattr(self, "interrupt") and self.interrupt: continue with StageProfiler( f"denoising_step_{i}", logger=logger, metrics=batch.metrics, perf_dump_path_provided=batch.perf_dump_path is not None, ): t_int = int(t.item()) if self.transformer_2 is not None: current_model, current_guidance_scale = ( self._select_and_manage_model( t_int=t_int, boundary_timestep=self._handle_boundary_ratio( server_args, batch ), server_args=server_args, batch=batch, ) ) else: current_model = self.transformer self._manage_device_placement(current_model, None, server_args) # Expand latents for I2V noise_latents = latents.clone() latent_model_input = latents.to(target_dtype) if batch.image_latent is not None: latent_model_input = torch.cat( [ latent_model_input, batch.image_latent.permute(0, 2, 1, 3, 4), ], dim=2, ).to(target_dtype) assert not torch.isnan( latent_model_input ).any(), "latent_model_input contains nan" # Prepare inputs for transformer t_expand = t.repeat(latent_model_input.shape[0]) guidance_expand = self.get_or_build_guidance( latent_model_input.shape[0], target_dtype, get_local_torch_device(), ) # Predict noise residual with torch.autocast( device_type=current_platform.device_type, dtype=target_dtype, enabled=autocast_enabled, ): attn_metadata = self._build_attn_metadata(i, batch, server_args) batch.is_cfg_negative = False with set_forward_context( current_timestep=i, attn_metadata=attn_metadata, forward_batch=batch, ): # Run transformer pred_noise = current_model( hidden_states=latent_model_input.permute(0, 2, 1, 3, 4), timestep=t_expand, guidance=guidance_expand, **image_kwargs, **pos_cond_kwargs, ).permute(0, 2, 1, 3, 4) pred_video = pred_noise_to_pred_video( pred_noise=pred_noise.flatten(0, 1), noise_input_latent=noise_latents.flatten(0, 1), timestep=t_expand, scheduler=self.scheduler, ).unflatten(0, pred_noise.shape[:2]) if i < len(timesteps) - 1: next_timestep = timesteps[i + 1] * torch.ones( [1], dtype=torch.long, device=pred_video.device ) noise = torch.randn( video_raw_latent_shape, dtype=pred_video.dtype, generator=batch.generator[0], device=self.device, ) latents = self.scheduler.add_noise( pred_video.flatten(0, 1), noise.flatten(0, 1), next_timestep, ).unflatten(0, pred_video.shape[:2]) else: latents = pred_video # Update progress bar if i == len(timesteps) - 1 or ( (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 and progress_bar is not None ): progress_bar.update() self.step_profile() denoising_loop_end_time = time.time() if len(timesteps) > 0: self.log_info( "average time per step: %.4f seconds", (denoising_loop_end_time - denoising_loop_start_time) / len(timesteps), ) self._post_denoising_loop( batch=batch, latents=latents, trajectory_latents=[], trajectory_timesteps=[], server_args=server_args, ) return batch def _select_and_manage_model( self, t_int: int, boundary_timestep: float | None, server_args: ServerArgs, batch: Req, ): if boundary_timestep is None or t_int >= boundary_timestep: # High-noise stage current_model = self.transformer model_to_offload = self.transformer_2 current_guidance_scale = batch.guidance_scale else: # Low-noise stage current_model = self.transformer_2 model_to_offload = self.transformer current_guidance_scale = batch.guidance_scale_2 self._manage_device_placement(current_model, model_to_offload, server_args) assert current_model is not None, "The model for the current step is not set." return current_model, current_guidance_scale def _manage_device_placement( self, model_to_use: torch.nn.Module, model_to_offload: torch.nn.Module | None, server_args: ServerArgs, ): """ Manages the offload / load behavior of dit """ if not server_args.dit_cpu_offload: return # Offload the unused model if it's on CUDA if ( model_to_offload is not None and next(model_to_offload.parameters()).device.type == "cuda" ): model_to_offload.to("cpu") # Load the model to use if it's on CPU if ( model_to_use is not None and next(model_to_use.parameters()).device.type == "cpu" ): model_to_use.to(get_local_torch_device()) def _handle_boundary_ratio( self, server_args, batch, ): """ (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert """ boundary_ratio = server_args.pipeline_config.dit_config.boundary_ratio if batch.boundary_ratio is not None: logger.info( "Overriding boundary ratio from %s to %s", boundary_ratio, batch.boundary_ratio, ) boundary_ratio = batch.boundary_ratio if boundary_ratio is not None: boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps else: boundary_timestep = None return boundary_timestep ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Encoding stage for diffusion pipelines. """ import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( V, # Import validators ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) class EncodingStage(PipelineStage): """ Stage for encoding pixel space representations into latent space. This stage handles the encoding of pixel-space video/images into latent representations for further processing in the diffusion pipeline. """ def __init__(self, vae: ParallelTiledVAE) -> None: super().__init__() self.vae: ParallelTiledVAE = vae @torch.no_grad() def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify encoding stage inputs.""" result = VerificationResult() # Input video/images for VAE encoding: [batch_size, channels, frames, height, width] result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify encoding stage outputs.""" result = VerificationResult() # Encoded latents: [batch_size, channels, frames, height_latents, width_latents] result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) return result def forward( self, batch: Req, server_args: ServerArgs, ) -> Req: """ Encode pixel space representations into latent space. Returns: The batch with encoded latents. """ assert batch.latents is not None and isinstance(batch.latents, torch.Tensor) self.vae = self.vae.to(get_local_torch_device()) # Setup VAE precision vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] vae_autocast_enabled = ( vae_dtype != torch.float32 ) and not server_args.disable_autocast # Normalize input to [-1, 1] range (reverse of decoding normalization) latents = (batch.latents * 2.0 - 1.0).clamp(-1, 1) # Move to appropriate device and dtype latents = latents.to(get_local_torch_device()) # Encode image to latents with torch.autocast( device_type=current_platform.device_type, dtype=vae_dtype, enabled=vae_autocast_enabled, ): if server_args.pipeline_config.vae_tiling: self.vae.enable_tiling() # if server_args.vae_sp: # self.vae.enable_parallel() if not vae_autocast_enabled: latents = latents.to(vae_dtype) latents = self.vae.encode(latents).mean # Update batch with encoded latents batch.latents = latents # Offload models if needed self.maybe_free_model_hooks() if server_args.vae_cpu_offload: self.vae.to("cpu") return batch ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_paint.py ================================================ """ Hunyuan3D paint/texture generation stages. Three-stage pipeline: Preprocess -> TexGen -> Postprocess. """ from __future__ import annotations import os from typing import Any import numpy as np import torch from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import ( retrieve_timesteps, ) from einops import rearrange from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( Hunyuan3D2PipelineConfig, ) from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( PipelineStage, StageParallelismType, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # Utility functions def guidance_scale_embedding( w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 ) -> torch.Tensor: """Generate guidance scale embeddings.""" assert len(w.shape) == 1 w = w * 1000.0 half_dim = embedding_dim // 2 emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) emb = w.to(dtype)[:, None] * emb[None, :] emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) if embedding_dim % 2 == 1: emb = torch.nn.functional.pad(emb, (0, 1)) assert emb.shape == (w.shape[0], embedding_dim) return emb def extract_into_tensor( a: torch.Tensor, t: torch.Tensor, x_shape: tuple, n_gen: int ) -> torch.Tensor: """Extract values from tensor and reshape for multi-view generation.""" out = a.gather(-1, t) out = out.repeat(n_gen) out = rearrange(out, "(b n) -> b n", n=n_gen) b, c, *_ = out.shape return out.reshape(b, c, *((1,) * (len(x_shape) - 2))) def get_predicted_original_sample( model_output: torch.Tensor, timesteps: torch.Tensor, sample: torch.Tensor, prediction_type: str, alphas: torch.Tensor, sigmas: torch.Tensor, n_gen: int, ) -> torch.Tensor: """Get predicted original sample from model output.""" alphas = extract_into_tensor(alphas, timesteps, sample.shape, n_gen) sigmas = extract_into_tensor(sigmas, timesteps, sample.shape, n_gen) model_output = rearrange(model_output, "(b n) c h w -> b n c h w", n=n_gen) if prediction_type == "epsilon": pred_x_0 = (sample - sigmas * model_output) / alphas elif prediction_type == "sample": pred_x_0 = model_output elif prediction_type == "v_prediction": pred_x_0 = alphas * sample - sigmas * model_output else: raise ValueError( f"Prediction type {prediction_type} is not supported; " "currently, `epsilon`, `sample`, and `v_prediction` are supported." ) return pred_x_0 def get_predicted_noise( model_output: torch.Tensor, timesteps: torch.Tensor, sample: torch.Tensor, prediction_type: str, alphas: torch.Tensor, sigmas: torch.Tensor, n_gen: int, ) -> torch.Tensor: """Get predicted noise from model output.""" alphas = extract_into_tensor(alphas, timesteps, sample.shape, n_gen) sigmas = extract_into_tensor(sigmas, timesteps, sample.shape, n_gen) model_output = rearrange(model_output, "(b n) c h w -> b n c h w", n=n_gen) if prediction_type == "epsilon": pred_epsilon = model_output elif prediction_type == "sample": pred_epsilon = (sample - alphas * model_output) / sigmas elif prediction_type == "v_prediction": pred_epsilon = alphas * model_output + sigmas * sample else: raise ValueError( f"Prediction type {prediction_type} is not supported; " "currently, `epsilon`, `sample`, and `v_prediction` are supported." ) return pred_epsilon def to_rgb_image(maybe_rgba): """Convert RGBA image to RGB.""" from PIL import Image if maybe_rgba.mode == "RGB": return maybe_rgba if maybe_rgba.mode == "RGBA": rgba = maybe_rgba img = np.random.randint( 127, 128, size=[rgba.size[1], rgba.size[0], 3], dtype=np.uint8 ) img = Image.fromarray(img, "RGB") img.paste(rgba, mask=rgba.getchannel("A")) return img raise ValueError(f"Unsupported image type: {maybe_rgba.mode}") class DDIMSolver: """DDIM solver for fast sampling.""" def __init__( self, alpha_cumprods: np.ndarray, timesteps: int = 1000, ddim_timesteps: int = 50, ): step_ratio = timesteps // ddim_timesteps self.ddim_timesteps = ( np.arange(1, ddim_timesteps + 1) * step_ratio ).round().astype(np.int64) - 1 self.ddim_alpha_cumprods = alpha_cumprods[self.ddim_timesteps] self.ddim_alpha_cumprods_prev = np.asarray( [alpha_cumprods[0]] + alpha_cumprods[self.ddim_timesteps[:-1]].tolist() ) self.ddim_timesteps = torch.from_numpy(self.ddim_timesteps).long() self.ddim_alpha_cumprods = torch.from_numpy(self.ddim_alpha_cumprods) self.ddim_alpha_cumprods_prev = torch.from_numpy(self.ddim_alpha_cumprods_prev) def to(self, device: torch.device) -> "DDIMSolver": self.ddim_timesteps = self.ddim_timesteps.to(device) self.ddim_alpha_cumprods = self.ddim_alpha_cumprods.to(device) self.ddim_alpha_cumprods_prev = self.ddim_alpha_cumprods_prev.to(device) return self def ddim_step( self, pred_x0: torch.Tensor, pred_noise: torch.Tensor, timestep_index: torch.Tensor, n_gen: int, ) -> torch.Tensor: alpha_cumprod_prev = extract_into_tensor( self.ddim_alpha_cumprods_prev, timestep_index, pred_x0.shape, n_gen ) dir_xt = (1.0 - alpha_cumprod_prev).sqrt() * pred_noise x_prev = alpha_cumprod_prev.sqrt() * pred_x0 + dir_xt return x_prev def _recorrect_rgb( src_image: torch.Tensor, target_image: torch.Tensor, alpha_channel: torch.Tensor, scale: float = 0.95, ) -> torch.Tensor: """Correct RGB values to match target color distribution.""" def flat_and_mask(bgr, a): mask = torch.where(a > 0.5, True, False) bgr_flat = bgr.reshape(-1, bgr.shape[-1]) mask_flat = mask.reshape(-1) bgr_flat_masked = bgr_flat[mask_flat, :] return bgr_flat_masked src_flat = flat_and_mask(src_image, alpha_channel) target_flat = flat_and_mask(target_image, alpha_channel) corrected_bgr = torch.zeros_like(src_image) for i in range(3): src_mean, src_stddev = torch.mean(src_flat[:, i]), torch.std(src_flat[:, i]) target_mean, target_stddev = torch.mean(target_flat[:, i]), torch.std( target_flat[:, i] ) corrected_bgr[:, :, i] = torch.clamp( (src_image[:, :, i] - scale * src_mean) * (target_stddev / src_stddev) + scale * target_mean, 0, 1, ) src_mse = torch.mean((src_image - target_image) ** 2) modify_mse = torch.mean((corrected_bgr - target_image) ** 2) if src_mse < modify_mse: corrected_bgr = torch.cat([src_image, alpha_channel], dim=-1) else: corrected_bgr = torch.cat([corrected_bgr, alpha_channel], dim=-1) return corrected_bgr # Stage 1: Preprocess (UV unwrap + delight + multi-view rendering) class Hunyuan3DPaintPreprocessStage(PipelineStage): """Preprocessing: UV unwrap + delight in parallel, then multi-view rendering.""" CAMERA_AZIMS = [0, 90, 180, 270, 0, 180] CAMERA_ELEVS = [0, 0, 0, 0, 90, -90] VIEW_WEIGHTS = [1, 0.1, 0.5, 0.1, 0.05, 0.05] @property def parallelism_type(self) -> StageParallelismType: return StageParallelismType.MAIN_RANK_ONLY def __init__(self, config: Hunyuan3D2PipelineConfig) -> None: super().__init__() self.config = config self._delight_pipeline = None self._delight_loaded = False self._renderer = None self._renderer_loaded = False # --- UV unwrap --- def _do_uv_unwrap(self, batch: Req, server_args: ServerArgs) -> Req: import time from sglang.multimodal_gen.runtime.utils.mesh3d_utils import mesh_uv_wrap mesh = batch.extra["shape_meshes"] if isinstance(mesh, list): mesh = mesh[0] try: start_time = time.time() mesh = mesh_uv_wrap(mesh) elapsed = time.time() - start_time logger.info(f"UV unwrapping completed in {elapsed:.2f}s") except Exception as e: logger.warning(f"UV unwrapping failed: {e}") batch.extra["paint_mesh"] = mesh return batch # --- Delight --- def _load_delight_model(self, server_args: ServerArgs): if self._delight_loaded: return from diffusers import ( EulerAncestralDiscreteScheduler, StableDiffusionInstructPix2PixPipeline, ) from huggingface_hub import snapshot_download model_path = server_args.model_path delight_subfolder = getattr( self.config, "delight_subfolder", "hunyuan3d-delight-v2-0" ) local_path = os.path.join(model_path, delight_subfolder) if not os.path.exists(local_path): local_path = os.path.expanduser(local_path) if not os.path.exists(local_path): try: downloaded = snapshot_download( repo_id=model_path, allow_patterns=[f"{delight_subfolder}/*"], ) local_path = os.path.join(downloaded, delight_subfolder) except Exception as e: logger.warning("Could not download delight model: %s", e) local_path = None if local_path and os.path.exists(local_path): pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained( local_path, torch_dtype=torch.float16, safety_checker=None, ) pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config( pipeline.scheduler.config ) pipeline.set_progress_bar_config(disable=True) self._delight_pipeline = pipeline.to(self.device, torch.float16) logger.info("Delight model loaded successfully") else: logger.warning( "Delight model not available, skipping delight preprocessing" ) self._delight_loaded = True @torch.no_grad() def _run_delight(self, image): import cv2 from PIL import Image as PILImage image = image.resize((512, 512)) if image.mode == "RGBA": image_array = np.array(image) alpha_channel = image_array[:, :, 3] erosion_size = 3 kernel = np.ones((erosion_size, erosion_size), np.uint8) alpha_channel = cv2.erode(alpha_channel, kernel, iterations=1) image_array[alpha_channel == 0, :3] = 255 image_array[:, :, 3] = alpha_channel image = PILImage.fromarray(image_array) image_tensor = torch.tensor(np.array(image) / 255.0).to(self.device) alpha = image_tensor[:, :, 3:] rgb_target = image_tensor[:, :, :3] else: image_tensor = torch.tensor(np.array(image) / 255.0).to(self.device) alpha = torch.ones_like(image_tensor)[:, :, :1] rgb_target = image_tensor[:, :, :3] image = image.convert("RGB") image = self._delight_pipeline( prompt=self.config.delight_prompt, image=image, generator=torch.manual_seed(42), height=512, width=512, num_inference_steps=self.config.delight_num_inference_steps, image_guidance_scale=self.config.delight_cfg_image, guidance_scale=self.config.delight_guidance_scale, ).images[0] image_tensor = torch.tensor(np.array(image) / 255.0).to(self.device) rgb_src = image_tensor[:, :, :3] image = _recorrect_rgb(rgb_src, rgb_target, alpha) image = image[:, :, :3] * image[:, :, 3:] + torch.ones_like(image[:, :, :3]) * ( 1.0 - image[:, :, 3:] ) image = PILImage.fromarray((image.cpu().numpy() * 255).astype(np.uint8)) return image def _do_delight(self, batch: Req, server_args: ServerArgs) -> Req: from PIL import Image from sglang.multimodal_gen.runtime.utils.mesh3d_utils import recenter_image image = Image.open(batch.image_path) image = recenter_image(image) if not self.config.delight_enable: logger.info("Delight preprocessing disabled, using original image") batch.extra["delighted_image"] = image return batch self._load_delight_model(server_args) if self._delight_pipeline is not None: try: image = self._run_delight(image) logger.info("Image delight completed") except Exception as e: logger.warning(f"Image delight failed: {e}") batch.extra["delighted_image"] = image return batch # --- Multi-view rendering --- def _init_renderer(self): if self._renderer_loaded: return from sglang.multimodal_gen.runtime.utils.mesh3d_utils import MeshRender self._renderer = MeshRender( default_resolution=self.config.paint_render_size, texture_size=self.config.paint_texture_size, ) self._renderer_loaded = True logger.info("Mesh renderer initialized") def _render_multiview(self, mesh) -> tuple: self._init_renderer() self._renderer.load_mesh(mesh) normal_maps = self._renderer.render_normal_multiview( self.CAMERA_ELEVS, self.CAMERA_AZIMS, use_abs_coor=True ) position_maps = self._renderer.render_position_multiview( self.CAMERA_ELEVS, self.CAMERA_AZIMS ) logger.info(f"Rendered {len(normal_maps)} views for texture generation") return normal_maps, position_maps # --- Forward --- def forward(self, batch: Req, server_args: ServerArgs) -> Req: if batch.extra.get("_mesh_failed"): logger.warning("Mesh generation failed, skipping paint preprocessing") batch.extra["paint_mesh"] = None batch.extra["delighted_image"] = None batch.extra["normal_maps"] = [] batch.extra["position_maps"] = [] batch.extra["camera_azims"] = self.CAMERA_AZIMS batch.extra["camera_elevs"] = self.CAMERA_ELEVS batch.extra["view_weights"] = self.VIEW_WEIGHTS batch.extra["renderer"] = None return batch import concurrent.futures import copy # 1. UV unwrap + delight in parallel batch_for_uv = batch batch_for_delight = copy.copy(batch) batch_for_delight.extra = batch.extra.copy() with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: uv_future = executor.submit(self._do_uv_unwrap, batch_for_uv, server_args) delight_future = executor.submit( self._do_delight, batch_for_delight, server_args ) uv_future.result() delight_future.result() batch.extra["paint_mesh"] = batch_for_uv.extra.get("paint_mesh") batch.extra["delighted_image"] = batch_for_delight.extra.get("delighted_image") # 2. Multi-view rendering normal_maps, position_maps = self._render_multiview(batch.extra["paint_mesh"]) batch.extra["normal_maps"] = normal_maps batch.extra["position_maps"] = position_maps batch.extra["camera_azims"] = self.CAMERA_AZIMS batch.extra["camera_elevs"] = self.CAMERA_ELEVS batch.extra["view_weights"] = self.VIEW_WEIGHTS batch.extra["renderer"] = self._renderer return batch def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: result = VerificationResult() result.add_check("shape_meshes", batch.extra.get("shape_meshes"), V.not_none) result.add_check("image_path", batch.image_path, V.not_none) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: result = VerificationResult() result.add_check("paint_mesh", batch.extra.get("paint_mesh"), V.not_none) result.add_check( "delighted_image", batch.extra.get("delighted_image"), V.not_none ) result.add_check("normal_maps", batch.extra.get("normal_maps"), V.is_list) result.add_check("position_maps", batch.extra.get("position_maps"), V.is_list) result.add_check("renderer", batch.extra.get("renderer"), V.not_none) return result # Stage 2: TexGen (model loading + input prep + denoising + decode) class Hunyuan3DPaintTexGenStage(PipelineStage): def __init__( self, config: Hunyuan3D2PipelineConfig, paint_dir: str | None = None, transformer: Any = None, scheduler: Any = None, vae: Any = None, vae_scale_factor: int = 8, image_processor: Any = None, solver: Any = None, is_turbo: bool = False, ) -> None: super().__init__() self.config = config self.paint_dir = paint_dir self.transformer = transformer self.scheduler = scheduler self.vae = vae self.vae_scale_factor = vae_scale_factor self.image_processor = image_processor self.solver = solver self.is_turbo = is_turbo self._loaded = transformer is not None @property def parallelism_type(self) -> StageParallelismType: return StageParallelismType.MAIN_RANK_ONLY def _load_paint_models(self, server_args: ServerArgs) -> None: """Load paint models from pre-resolved local path (no network).""" if self._loaded: return if self.paint_dir is None: logger.warning("No paint model directory resolved, skipping") self._loaded = True return try: self._do_load_paint(server_args) logger.info("Paint pipeline loaded successfully") except Exception as e: logger.warning("Failed to load paint pipeline: %s", e) self.vae = None self.transformer = None self.scheduler = None self._loaded = True def _do_load_paint(self, server_args: ServerArgs) -> None: import json from diffusers import AutoencoderKL from diffusers.image_processor import VaeImageProcessor from sglang.multimodal_gen.runtime.models.dits.hunyuan3d import ( UNet2p5DConditionModel, ) local_path = self.paint_dir logger.info("Loading paint model from %s", local_path) vae_dir = os.path.join(local_path, "vae") with open(os.path.join(vae_dir, "config.json"), "r") as f: vae_config = json.load(f) vae_config = {k: v for k, v in vae_config.items() if not k.startswith("_")} self.vae = AutoencoderKL(**vae_config) st_path = os.path.join(vae_dir, "diffusion_pytorch_model.safetensors") bin_path = os.path.join(vae_dir, "diffusion_pytorch_model.bin") if os.path.exists(st_path): from safetensors.torch import load_file state_dict = load_file(st_path) elif os.path.exists(bin_path): state_dict = torch.load(bin_path, map_location="cpu", weights_only=True) else: raise FileNotFoundError(f"No VAE weights in {vae_dir}") self.vae.load_state_dict(state_dict) self.vae = self.vae.to(device=self.device, dtype=torch.float16).eval() self.transformer = UNet2p5DConditionModel.from_pretrained( os.path.join(local_path, "unet"), torch_dtype=torch.float16, ).to(self.device) self.is_turbo = bool(getattr(self.config, "paint_turbo_mode", False)) sched_path = os.path.join(local_path, "scheduler", "scheduler_config.json") with open(sched_path, "r") as f: sched_cfg = json.load(f) if self.is_turbo: from diffusers import LCMScheduler self.scheduler = LCMScheduler.from_config(sched_cfg) else: from diffusers import EulerAncestralDiscreteScheduler self.scheduler = EulerAncestralDiscreteScheduler.from_config( sched_cfg, timestep_spacing="trailing" ) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) self.solver = DDIMSolver( self.scheduler.alphas_cumprod.cpu().numpy(), timesteps=self.scheduler.config.num_train_timesteps, ddim_timesteps=30, ).to(self.device) if server_args.enable_torch_compile: compile_mode = os.environ.get( "SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs" ) logger.info("Compiling paint transformer with mode: %s", compile_mode) self.transformer.compile(mode=compile_mode, fullgraph=False, dynamic=None) def _convert_pil_list_to_tensor( self, images: list, device: torch.device ) -> torch.Tensor: bg_c = [1.0, 1.0, 1.0] images_tensor = [] for batch_imgs in images: view_imgs = [] for pil_img in batch_imgs: if pil_img.mode == "L": pil_img = pil_img.point( lambda x: 255 if x > 1 else 0, mode="1" ).convert("RGB") img = np.asarray(pil_img, dtype=np.float32) / 255.0 if img.shape[2] > 3: alpha = img[:, :, 3:] img = img[:, :, :3] * alpha + bg_c * (1 - alpha) img = ( torch.from_numpy(img) .permute(2, 0, 1) .unsqueeze(0) .contiguous() .to(device=device, dtype=self.vae.dtype) ) view_imgs.append(img) view_imgs = torch.cat(view_imgs, dim=0) images_tensor.append(view_imgs.unsqueeze(0)) return torch.cat(images_tensor, dim=0) @torch.no_grad() def _encode_images(self, images: torch.Tensor) -> torch.Tensor: batch_size = images.shape[0] images = rearrange(images, "b n c h w -> (b n) c h w") dtype = next(self.vae.parameters()).dtype images = (images - 0.5) * 2.0 posterior = self.vae.encode(images.to(dtype)).latent_dist latents = posterior.sample() * self.vae.config.scaling_factor return rearrange(latents, "(b n) c h w -> b n c h w", b=batch_size) @staticmethod def _compute_camera_index(azim: float, elev: float) -> int: base_idx = int(((azim // 30) + 9) % 12) if elev == 0: base, divisor = 12, 1 elif elev == 20: base, divisor = 24, 1 elif elev == -20: base, divisor = 0, 1 elif elev == 90: base, divisor = 40, 3 elif elev == -90: base, divisor = 36, 3 else: base, divisor = 12, 1 return base + (base_idx // divisor) def _prepare_denoising_inputs( self, batch: Req, server_args: ServerArgs, ) -> dict[str, Any]: import random from diffusers.utils.torch_utils import randn_tensor device = self.device normal_maps = batch.extra["normal_maps"] position_maps = batch.extra["position_maps"] camera_azims = batch.extra["camera_azims"] camera_elevs = batch.extra["camera_elevs"] num_steps = self.config.paint_num_inference_steps guidance_scale = self.config.paint_guidance_scale render_size = self.config.paint_resolution num_in_batch = len(normal_maps) seed = 0 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) generator = torch.Generator(device=device).manual_seed(seed) image = batch.extra["delighted_image"] if not isinstance(image, list): image = [image] image = [to_rgb_image(img) for img in image] image_vae = [ torch.tensor(np.array(img, dtype=np.float32) / 255.0) for img in image ] image_vae = [ iv.unsqueeze(0).permute(0, 3, 1, 2).unsqueeze(0) for iv in image_vae ] image_vae = torch.cat(image_vae, dim=1).to(device=device, dtype=self.vae.dtype) ref_latents = self._encode_images(image_vae) target_size = render_size if isinstance(normal_maps, list): normal_maps = [ ( img.resize((target_size, target_size)) if hasattr(img, "resize") else img ) for img in normal_maps ] normal_maps = self._convert_pil_list_to_tensor([normal_maps], device) if isinstance(position_maps, list): position_maps = [ ( img.resize((target_size, target_size)) if hasattr(img, "resize") else img ) for img in position_maps ] position_maps = self._convert_pil_list_to_tensor([position_maps], device) normal_imgs = ( self._encode_images(normal_maps) if normal_maps is not None else None ) position_imgs = ( self._encode_images(position_maps) if position_maps is not None else None ) camera_info = [ self._compute_camera_index(azim, elev) for azim, elev in zip(camera_azims, camera_elevs) ] camera_info_gen = torch.tensor([camera_info], device=device, dtype=torch.int64) camera_info_ref = torch.tensor([[0]], device=device, dtype=torch.int64) do_cfg = guidance_scale > 1 and not self.is_turbo if self.is_turbo and position_maps is not None: from sglang.multimodal_gen.runtime.models.dits.hunyuan3d import ( compute_multi_resolution_discrete_voxel_indice, compute_multi_resolution_mask, ) position_attn_mask = compute_multi_resolution_mask(position_maps) position_voxel_indices = compute_multi_resolution_discrete_voxel_indice( position_maps ) else: position_attn_mask = None position_voxel_indices = None if do_cfg: negative_ref_latents = torch.zeros_like(ref_latents) ref_latents = torch.cat([negative_ref_latents, ref_latents]) ref_scale = torch.as_tensor([0.0, 1.0]).to(ref_latents) if normal_imgs is not None: normal_imgs = torch.cat((normal_imgs, normal_imgs)) if position_imgs is not None: position_imgs = torch.cat((position_imgs, position_imgs)) if position_maps is not None: position_maps = torch.cat((position_maps, position_maps)) camera_info_gen = torch.cat((camera_info_gen, camera_info_gen)) camera_info_ref = torch.cat((camera_info_ref, camera_info_ref)) else: ref_scale = None model_kwargs = { "ref_latents": ref_latents, "num_in_batch": num_in_batch, } if ref_scale is not None: model_kwargs["ref_scale"] = ref_scale if normal_imgs is not None: model_kwargs["normal_imgs"] = normal_imgs if position_imgs is not None: model_kwargs["position_imgs"] = position_imgs if position_maps is not None: model_kwargs["position_maps"] = position_maps model_kwargs["camera_info_gen"] = camera_info_gen model_kwargs["camera_info_ref"] = camera_info_ref if position_attn_mask is not None: model_kwargs["position_attn_mask"] = position_attn_mask if position_voxel_indices is not None: model_kwargs["position_voxel_indices"] = position_voxel_indices prompt_embeds = self.transformer.learned_text_clip_gen.repeat(1, 1, 1) negative_prompt_embeds = torch.zeros_like(prompt_embeds) if self.is_turbo: bsz = 3 index = torch.arange(29, -1, -bsz, device=device).long() timesteps = self.solver.ddim_timesteps[index] self.scheduler.set_timesteps(timesteps=timesteps.cpu(), device=device) timesteps = self.scheduler.timesteps else: timesteps, num_steps = retrieve_timesteps( self.scheduler, num_steps, device, None, None ) num_channels_latents = self.transformer.config.in_channels latent_shape = ( num_in_batch, num_channels_latents, render_size // self.vae_scale_factor, render_size // self.vae_scale_factor, ) latents = randn_tensor( latent_shape, generator=generator, device=device, dtype=prompt_embeds.dtype ) latents = latents * self.scheduler.init_noise_sigma return { "timesteps": timesteps, "latents": latents, "prompt_embeds": prompt_embeds, "negative_prompt_embeds": negative_prompt_embeds, "model_kwargs": model_kwargs, "num_in_batch": num_in_batch, "num_inference_steps": num_steps, "guidance_scale": guidance_scale, "do_cfg": do_cfg, "generator": generator, "num_channels_latents": num_channels_latents, } @torch.no_grad() def _denoise_loop( self, timesteps: torch.Tensor, latents: torch.Tensor, prompt_embeds: torch.Tensor, negative_prompt_embeds: torch.Tensor, model_kwargs: dict[str, Any], num_in_batch: int, guidance_scale: float, do_cfg: bool, generator: torch.Generator, num_channels_latents: int, ) -> torch.Tensor: import inspect if do_cfg: prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) extra_step_kwargs = {} if "eta" in inspect.signature(self.scheduler.step).parameters: extra_step_kwargs["eta"] = 0.0 if "generator" in inspect.signature(self.scheduler.step).parameters: extra_step_kwargs["generator"] = generator for step_idx, t in enumerate(timesteps): latents = rearrange(latents, "(b n) c h w -> b n c h w", n=num_in_batch) latent_model_input = torch.cat([latents] * 2) if do_cfg else latents latent_model_input = rearrange( latent_model_input, "b n c h w -> (b n) c h w" ) latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) latent_model_input = rearrange( latent_model_input, "(b n) c h w -> b n c h w", n=num_in_batch ) with set_forward_context( current_timestep=step_idx, attn_metadata=None, ): noise_pred = self.transformer( latent_model_input, t, encoder_hidden_states=prompt_embeds, timestep_cond=None, cross_attention_kwargs=None, added_cond_kwargs=None, return_dict=False, **model_kwargs, )[0] latents = rearrange(latents, "b n c h w -> (b n) c h w") if do_cfg: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * ( noise_pred_text - noise_pred_uncond ) latents = self.scheduler.step( noise_pred, t, latents[:, :num_channels_latents, :, :], **extra_step_kwargs, return_dict=False, )[0] return latents @torch.no_grad() def _decode_latents(self, latents: torch.Tensor) -> list: image = self.vae.decode( latents / self.vae.config.scaling_factor, return_dict=False )[0] return self.image_processor.postprocess(image, output_type="pil") def forward(self, batch: Req, server_args: ServerArgs) -> Req: if batch.extra.get("_mesh_failed"): logger.warning("Mesh generation failed, skipping paint texgen") batch.extra["multiview_textures"] = [] return batch self._load_paint_models(server_args) delighted_image = batch.extra["delighted_image"] normal_maps = batch.extra["normal_maps"] if self.transformer is not None: try: prepared = self._prepare_denoising_inputs(batch, server_args) latents = self._denoise_loop( timesteps=prepared["timesteps"], latents=prepared["latents"], prompt_embeds=prepared["prompt_embeds"], negative_prompt_embeds=prepared["negative_prompt_embeds"], model_kwargs=prepared["model_kwargs"], num_in_batch=prepared["num_in_batch"], guidance_scale=prepared["guidance_scale"], do_cfg=prepared["do_cfg"], generator=prepared["generator"], num_channels_latents=prepared["num_channels_latents"], ) multiview_textures = self._decode_latents(latents) logger.info( "Paint pipeline generated %d textures", len(multiview_textures) ) except Exception as e: logger.error(f"Paint pipeline execution failed: {e}") import traceback traceback.print_exc() render_size = self.config.paint_resolution multiview_textures = [ delighted_image.resize((render_size, render_size)) for _ in range(len(normal_maps)) ] else: logger.warning( "Paint pipeline not available, using reference image for all views" ) render_size = self.config.paint_resolution multiview_textures = [ delighted_image.resize((render_size, render_size)) for _ in range(len(normal_maps)) ] batch.extra["multiview_textures"] = multiview_textures logger.info(f"Generated {len(multiview_textures)} texture views") return batch def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: if batch.extra.get("_mesh_failed"): return VerificationResult() result = VerificationResult() result.add_check( "delighted_image", batch.extra.get("delighted_image"), V.not_none ) result.add_check("normal_maps", batch.extra.get("normal_maps"), V.is_list) result.add_check("position_maps", batch.extra.get("position_maps"), V.is_list) result.add_check("camera_azims", batch.extra.get("camera_azims"), V.is_list) result.add_check("camera_elevs", batch.extra.get("camera_elevs"), V.is_list) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: result = VerificationResult() result.add_check( "multiview_textures", batch.extra.get("multiview_textures"), V.is_list ) return result # Stage 3: Postprocess (texture baking + mesh export) class Hunyuan3DPaintPostprocessStage(PipelineStage): """Texture baking from multi-view images and final mesh export.""" @property def parallelism_type(self) -> StageParallelismType: return StageParallelismType.MAIN_RANK_ONLY def __init__(self, config: Hunyuan3D2PipelineConfig) -> None: super().__init__() self.config = config def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: if batch.extra.get("_mesh_failed"): logger.warning("Mesh generation failed, skipping paint postprocess") return OutputBatch(output_file_paths=[], metrics=batch.metrics) renderer = batch.extra["renderer"] multiview_textures = batch.extra["multiview_textures"] camera_elevs = batch.extra["camera_elevs"] camera_azims = batch.extra["camera_azims"] view_weights = batch.extra["view_weights"] render_size = getattr(self.config, "paint_render_size", 2048) resized_textures = [] for tex in multiview_textures: if hasattr(tex, "resize"): resized_textures.append(tex.resize((render_size, render_size))) else: resized_textures.append(tex) try: texture, mask = renderer.bake_from_multiview( resized_textures, camera_elevs, camera_azims, view_weights, method="fast", ) mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype("uint8") texture = renderer.texture_inpaint(texture, mask_np) renderer.set_texture(texture) textured_mesh = renderer.save_mesh() logger.info("Texture baking completed") except Exception as e: logger.error(f"Texture baking failed: {e}") textured_mesh = batch.extra["paint_mesh"] obj_path = batch.extra["shape_obj_path"] return_path = batch.extra["shape_return_path"] try: textured_mesh.export(obj_path) if self.config.paint_save_glb: glb_path = obj_path[:-4] + ".glb" textured_mesh.export(glb_path) return_path = glb_path self._cleanup_obj_artifacts(obj_path) except Exception as e: logger.error(f"Mesh export failed: {e}") return OutputBatch(output_file_paths=[return_path], metrics=batch.metrics) @staticmethod def _cleanup_obj_artifacts(obj_path: str) -> None: """Remove OBJ file and trimesh-generated material artifacts.""" obj_dir = os.path.dirname(obj_path) or "." targets = [obj_path] for f in os.listdir(obj_dir): if f.endswith(".mtl") or (f.startswith("material") and f.endswith(".png")): targets.append(os.path.join(obj_dir, f)) for path in targets: try: os.remove(path) except OSError: pass def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: if batch.extra.get("_mesh_failed"): return VerificationResult() result = VerificationResult() result.add_check("renderer", batch.extra.get("renderer"), V.not_none) result.add_check( "multiview_textures", batch.extra.get("multiview_textures"), V.is_list ) result.add_check("camera_elevs", batch.extra.get("camera_elevs"), V.is_list) result.add_check("camera_azims", batch.extra.get("camera_azims"), V.is_list) result.add_check("view_weights", batch.extra.get("view_weights"), V.is_list) return result __all__ = [ "Hunyuan3DPaintPreprocessStage", "Hunyuan3DPaintTexGenStage", "Hunyuan3DPaintPostprocessStage", ] ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_shape.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Hunyuan3D shape generation stages. Four-stage pipeline: BeforeDenoising -> Denoising -> Export -> Save. """ from __future__ import annotations import os from typing import Any import numpy as np import torch from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import ( Hunyuan3D2PipelineConfig, ) from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import ( TransformerLoader, ) from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.mesh3d_utils import export_to_trimesh logger = init_logger(__name__) def retrieve_timesteps( scheduler, num_inference_steps=None, device=None, timesteps=None, sigmas=None, **kwargs, ): """Retrieve timesteps from scheduler.""" import inspect if timesteps is not None and sigmas is not None: raise ValueError("Only one of timesteps or sigmas can be passed.") if timesteps is not None: accepts_timesteps = "timesteps" in set( inspect.signature(scheduler.set_timesteps).parameters.keys() ) if not accepts_timesteps: raise ValueError( f"Scheduler {scheduler.__class__} doesn't support custom timesteps." ) scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) elif sigmas is not None: accepts_sigmas = "sigmas" in set( inspect.signature(scheduler.set_timesteps).parameters.keys() ) if not accepts_sigmas: raise ValueError( f"Scheduler {scheduler.__class__} doesn't support custom sigmas." ) scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) else: scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) timesteps = scheduler.timesteps return timesteps, num_inference_steps def _prepare_shape_image(image_processor, image, mask=None) -> dict: """Prepare shape image for conditioning.""" if isinstance(image, torch.Tensor) and isinstance(mask, torch.Tensor): return {"image": image, "mask": mask} if isinstance(image, str) and not os.path.exists(image): raise FileNotFoundError(f"Couldn't find image at path {image}") if not isinstance(image, list): image = [image] outputs = [image_processor(img) for img in image] cond_input = {k: [] for k in outputs[0].keys()} for output in outputs: for key, value in output.items(): cond_input[key].append(value) for key, value in cond_input.items(): if isinstance(value[0], torch.Tensor): cond_input[key] = torch.cat(value, dim=0) return cond_input def _move_to_device(payload, device, dtype): """Recursively move tensors in payload to specified device and dtype.""" if isinstance(payload, torch.Tensor): return payload.to(device=device, dtype=dtype) if isinstance(payload, dict): return {k: _move_to_device(v, device, dtype) for k, v in payload.items()} if isinstance(payload, list): return [_move_to_device(v, device, dtype) for v in payload] return payload class Hunyuan3DShapeBeforeDenoisingStage(PipelineStage): """Monolithic pre-processing stage for Hunyuan3D shape generation. Consolidates input validation, image preprocessing, conditioning, and latent/timestep preparation into a single stage. """ def __init__( self, image_processor: Any, conditioner: Any, vae: Any, model: Any, scheduler: Any, config: Hunyuan3D2PipelineConfig, ) -> None: super().__init__() self.image_processor = image_processor self.conditioner = conditioner self.vae = vae self.model = model self.scheduler = scheduler self.config = config def _validate_input(self, batch: Req, server_args: ServerArgs) -> None: if batch.image_path is None: raise ValueError("Hunyuan3D requires 'image_path' input.") if isinstance(batch.image_path, list): if len(batch.image_path) != 1: raise ValueError("Hunyuan3D only supports a single image input.") batch.image_path = batch.image_path[0] if not isinstance(batch.image_path, str): raise ValueError( f"Hunyuan3D expects image_path as str, got {type(batch.image_path)}" ) if not os.path.exists(batch.image_path): raise FileNotFoundError(f"Image path not found: {batch.image_path}") if batch.num_outputs_per_prompt != 1: raise ValueError("Hunyuan3D only supports num_outputs_per_prompt=1.") def _prepare_latents(self, batch_size, dtype, device, generator): from diffusers.utils.torch_utils import randn_tensor shape = (batch_size, *self.vae.latent_shape) latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) return latents * getattr(self.scheduler, "init_noise_sigma", 1.0) def forward(self, batch: Req, server_args: ServerArgs) -> Req: # 1. Input validation self._validate_input(batch, server_args) # 2. Image preprocessing cond_inputs = _prepare_shape_image(self.image_processor, batch.image_path) image = cond_inputs.pop("image") device = self.device dtype = next(self.model.parameters()).dtype image = _move_to_device(image, device, dtype) cond_inputs = _move_to_device(cond_inputs, device, dtype) # 3. Conditioning with CFG do_cfg = batch.guidance_scale >= 0 and not ( hasattr(self.model, "guidance_embed") and self.model.guidance_embed is True ) cond = self.conditioner(image=image, **cond_inputs) if do_cfg: un_cond = self.conditioner.unconditional_embedding( image.shape[0], **cond_inputs ) def cat_recursive(a, b): if isinstance(a, torch.Tensor): return torch.cat([a, b], dim=0).to(dtype) out = {} for key in a.keys(): out[key] = cat_recursive(a[key], b[key]) return out cond = cat_recursive(cond, un_cond) # 4. Latent and timestep preparation batch_size = image.shape[0] sigmas = np.linspace(0, 1, batch.num_inference_steps) timesteps, _ = retrieve_timesteps( self.scheduler, batch.num_inference_steps, device, sigmas=sigmas, ) generator = batch.generator if generator is None and batch.seed is not None: generator = torch.Generator(device=device).manual_seed(batch.seed) latents = self._prepare_latents(batch_size, dtype, device, generator) guidance = None if hasattr(self.model, "guidance_embed") and self.model.guidance_embed is True: guidance = torch.tensor( [batch.guidance_scale] * batch_size, device=device, dtype=dtype ) # 5. Populate batch batch.prompt_embeds = [cond] batch.do_classifier_free_guidance = do_cfg batch.timesteps = timesteps batch.latents = latents batch.extra["shape_guidance"] = guidance batch.extra["shape_image"] = image return batch def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: result = VerificationResult() result.add_check("image_path", batch.image_path, V.not_none) result.add_check( "num_inference_steps", batch.num_inference_steps, V.positive_int ) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: result = VerificationResult() result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)]) result.add_check("latents", batch.latents, V.is_tensor) result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) return result class Hunyuan3DShapeDenoisingStage(DenoisingStage): """Denoising stage for Hunyuan3D shape generation.""" def __init__(self, transformer: Any, scheduler: Any, **kwargs) -> None: super().__init__(transformer=transformer, scheduler=scheduler, **kwargs) def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): """Prepare Hunyuan3D-specific variables for the base denoising loop.""" assert self.transformer is not None pipeline = self.pipeline() if self.pipeline else None cache_dit_num_inference_steps = batch.extra.get( "cache_dit_num_inference_steps", batch.num_inference_steps ) if not server_args.model_loaded["transformer"]: loader = TransformerLoader() self.transformer = loader.load( server_args.model_paths["transformer"], server_args, "transformer" ) self._maybe_enable_cache_dit(cache_dit_num_inference_steps, batch) self._maybe_enable_torch_compile(self.transformer) if pipeline: pipeline.add_module("transformer", self.transformer) server_args.model_loaded["transformer"] = True else: self._maybe_enable_cache_dit(cache_dit_num_inference_steps, batch) timesteps = batch.timesteps if timesteps is None: raise ValueError("Timesteps must be provided") latents = batch.latents if latents is None: raise ValueError("Latents must be provided") cond = batch.prompt_embeds[0] if batch.prompt_embeds else None if cond is None: raise ValueError("Conditioning (prompt_embeds) must be provided") if batch.raw_latent_shape is None: batch.raw_latent_shape = latents.shape guidance = batch.extra.get("shape_guidance") num_inference_steps = batch.num_inference_steps num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order extra_step_kwargs = self.prepare_extra_func_kwargs( self.scheduler.step, {"generator": batch.generator, "eta": batch.eta}, ) target_dtype = next(self.transformer.parameters()).dtype autocast_enabled = False pos_cond_kwargs = {"encoder_hidden_states": cond} neg_cond_kwargs = {} return { "extra_step_kwargs": extra_step_kwargs, "target_dtype": target_dtype, "autocast_enabled": autocast_enabled, "timesteps": timesteps, "num_inference_steps": num_inference_steps, "num_warmup_steps": num_warmup_steps, "image_kwargs": {}, "pos_cond_kwargs": pos_cond_kwargs, "neg_cond_kwargs": neg_cond_kwargs, "latents": latents, "prompt_embeds": batch.prompt_embeds, "neg_prompt_embeds": None, "boundary_timestep": None, "z": None, "reserved_frames_mask": None, "seq_len": None, "guidance": guidance, } def _predict_noise( self, current_model, latent_model_input, timestep, target_dtype, guidance: torch.Tensor, **kwargs, ): """Hunyuan3D-specific noise prediction with normalized timestep.""" cond = kwargs.get("encoder_hidden_states") timestep_norm = timestep / self.scheduler.config.num_train_timesteps return current_model(latent_model_input, timestep_norm, cond, guidance=guidance) def _predict_noise_with_cfg( self, current_model, latent_model_input: torch.Tensor, timestep, batch: Req, timestep_index: int, attn_metadata, target_dtype, current_guidance_scale, image_kwargs: dict[str, Any], pos_cond_kwargs: dict[str, Any], neg_cond_kwargs: dict[str, Any], server_args, guidance, latents, ): """Hunyuan3D-specific CFG: concat latents, single forward, then split.""" cond = pos_cond_kwargs.get("encoder_hidden_states") do_cfg = batch.do_classifier_free_guidance if do_cfg: latent_input = torch.cat([latent_model_input] * 2) else: latent_input = latent_model_input timestep_expanded = timestep.expand(latent_input.shape[0]).to(latents.dtype) with set_forward_context( current_timestep=timestep_index, attn_metadata=attn_metadata, forward_batch=batch, ): noise_pred = self._predict_noise( current_model=current_model, latent_model_input=latent_input, timestep=timestep_expanded, target_dtype=target_dtype, guidance=guidance, encoder_hidden_states=cond, ) if do_cfg: noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2) noise_pred = noise_pred_uncond + current_guidance_scale * ( noise_pred_cond - noise_pred_uncond ) return noise_pred def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: result = VerificationResult() result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)]) result.add_check("latents", batch.latents, V.is_tensor) result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) result.add_check( "num_inference_steps", batch.num_inference_steps, V.positive_int ) result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: result = VerificationResult() result.add_check("latents", batch.latents, V.is_tensor) return result class Hunyuan3DShapeExportStage(PipelineStage): """VAE decoding and mesh extraction stage.""" def __init__(self, vae: Any, config: Hunyuan3D2PipelineConfig) -> None: super().__init__() self.vae = vae self.config = config def forward(self, batch: Req, server_args: ServerArgs) -> Req: if self.config.shape_mc_algo is not None: try: from sglang.multimodal_gen.runtime.models.vaes.hunyuan3d_vae import ( SurfaceExtractors, ) self.vae.surface_extractor = SurfaceExtractors[ self.config.shape_mc_algo ]() except ImportError: logger.warning( f"Could not load SurfaceExtractors for mc_algo={self.config.shape_mc_algo}" ) latents = batch.latents if self.config.shape_output_type != "latent": latents = 1.0 / self.vae.scale_factor * latents latents = self.vae(latents) outputs = self.vae.latents2mesh( latents, bounds=self.config.shape_box_v, mc_level=self.config.shape_mc_level, num_chunks=self.config.shape_num_chunks, octree_resolution=self.config.shape_octree_resolution, mc_algo=self.config.shape_mc_algo, enable_pbar=False, ) else: outputs = latents if self.config.shape_output_type == "trimesh": outputs = export_to_trimesh(outputs) batch.extra["shape_meshes"] = outputs return batch def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: result = VerificationResult() result.add_check("latents", batch.latents, V.is_tensor) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: result = VerificationResult() result.add_check("shape_meshes", batch.extra.get("shape_meshes"), V.not_none) return result class Hunyuan3DShapeSaveStage(PipelineStage): """Mesh file export and output decision stage.""" def __init__(self, config: Hunyuan3D2PipelineConfig) -> None: super().__init__() self.config = config def _get_output_paths(self, batch: Req) -> tuple[str, str]: output_path = batch.output_file_path() or os.path.join( batch.output_path, "output.obj" ) if output_path.endswith(".glb"): obj_path = output_path[:-4] + ".obj" return obj_path, output_path if output_path.endswith(".obj"): return output_path, output_path return output_path + ".obj", output_path + ".obj" def forward(self, batch: Req, server_args: ServerArgs) -> Req | OutputBatch: mesh_outputs = batch.extra["shape_meshes"] mesh = mesh_outputs[0] if isinstance(mesh_outputs, list) else mesh_outputs if isinstance(mesh, list): mesh = mesh[0] if mesh is None: if batch.is_warmup: logger.info( "Skipping mesh export during warmup " "(surface extraction returned None)" ) batch.extra["_mesh_failed"] = True if self.config.paint_enable: return batch return OutputBatch(output_file_paths=[], metrics=batch.metrics) raise RuntimeError( "Mesh generation failed: surface extraction returned None. " "The surface level may be outside the volume data range." ) obj_path, return_path = self._get_output_paths(batch) output_dir = os.path.dirname(obj_path) if output_dir: os.makedirs(output_dir, exist_ok=True) mesh.export(obj_path) batch.extra["shape_obj_path"] = obj_path batch.extra["shape_return_path"] = return_path if self.config.paint_enable: return batch if return_path.endswith(".glb"): return_path = obj_path return OutputBatch(output_file_paths=[return_path], timings=batch.timings) def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: result = VerificationResult() result.add_check("shape_meshes", batch.extra.get("shape_meshes"), V.not_none) return result __all__ = [ "retrieve_timesteps", "Hunyuan3DShapeBeforeDenoisingStage", "Hunyuan3DShapeDenoisingStage", "Hunyuan3DShapeExportStage", "Hunyuan3DShapeSaveStage", ] ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Image encoding stages for I2V diffusion pipelines. This module contains implementations of image encoding stages for diffusion pipelines. """ import inspect import PIL import torch from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution from diffusers.models.modeling_outputs import AutoencoderKLOutput from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( qwen_image_postprocess_text, ) from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE from sglang.multimodal_gen.runtime.models.vision_utils import ( normalize, numpy_to_pt, pil_to_numpy, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) class ImageEncodingStage(PipelineStage): """ Stage for encoding image prompts into embeddings for diffusion models. This stage handles the encoding of image prompts into the embedding space expected by the diffusion model. """ def __init__( self, image_processor, image_encoder=None, text_encoder=None, ) -> None: """ Initialize the prompt encoding stage. Args: text_encoder: An encoder to encode input_ids and pixel values """ super().__init__() self.image_processor = image_processor self.image_encoder = image_encoder self.text_encoder = text_encoder def load_model(self): if self.server_args.image_encoder_cpu_offload: device = get_local_torch_device() self.move_to_device(device) def offload_model(self): if self.server_args.image_encoder_cpu_offload: self.move_to_device("cpu") def move_to_device(self, device): if self.server_args.use_fsdp_inference: return fields = [ "image_processor", "image_encoder", ] for field in fields: processor = getattr(self, field, None) if processor and hasattr(processor, "to"): setattr(self, field, processor.to(device)) def encoding_qwen_image_edit(self, outputs, image_inputs): # encoder hidden state prompt_embeds = qwen_image_postprocess_text(outputs, image_inputs, 64) return prompt_embeds @torch.no_grad() def forward( self, batch: Req, server_args: ServerArgs, ) -> Req: """ Encode the prompt into image encoder hidden states. """ if batch.condition_image is None: return batch cuda_device = get_local_torch_device() self.load_model() image_processor_kwargs = ( server_args.pipeline_config.prepare_image_processor_kwargs(batch) ) per_prompt_images = image_processor_kwargs.pop("per_prompt_images", None) texts = image_processor_kwargs.pop("text", None) if per_prompt_images is None: per_prompt_images = [batch.condition_image] texts = [None] if texts is None else texts all_prompt_embeds = [] all_neg_prompt_embeds = [] image_processor_call_params = inspect.signature( self.image_processor.__call__ ).parameters image_processor_kwargs = { k: v for k, v in image_processor_kwargs.items() if k in image_processor_call_params } for idx, prompt_images in enumerate(per_prompt_images): if not prompt_images: continue cur_kwargs = image_processor_kwargs.copy() if texts and idx < len(texts) and "text" in image_processor_call_params: cur_kwargs["text"] = [texts[idx]] image_inputs = self.image_processor( images=prompt_images, return_tensors="pt", **cur_kwargs ).to(cuda_device) if self.image_encoder: # if an image encoder is provided with set_forward_context(current_timestep=0, attn_metadata=None): outputs = self.image_encoder( **image_inputs, **server_args.pipeline_config.image_encoder_extra_args, ) image_embeds = server_args.pipeline_config.postprocess_image( outputs ) batch.image_embeds.append(image_embeds) elif self.text_encoder: # if a text encoder is provided, e.g. Qwen-Image-Edit # 1. neg prompt embeds if batch.do_classifier_free_guidance: neg_image_processor_kwargs = ( server_args.pipeline_config.prepare_image_processor_kwargs( batch, neg=True ) ) neg_image_processor_kwargs.pop("per_prompt_images", None) neg_texts = neg_image_processor_kwargs.pop("text", None) if neg_texts and idx < len(neg_texts): neg_image_processor_kwargs["text"] = [neg_texts[idx]] neg_image_inputs = self.image_processor( images=prompt_images, return_tensors="pt", **neg_image_processor_kwargs, ).to(cuda_device) with set_forward_context(current_timestep=0, attn_metadata=None): outputs = self.text_encoder( input_ids=image_inputs.input_ids, attention_mask=image_inputs.attention_mask, pixel_values=image_inputs.pixel_values, image_grid_thw=image_inputs.image_grid_thw, output_hidden_states=True, ) if batch.do_classifier_free_guidance: neg_outputs = self.text_encoder( input_ids=neg_image_inputs.input_ids, attention_mask=neg_image_inputs.attention_mask, pixel_values=neg_image_inputs.pixel_values, image_grid_thw=neg_image_inputs.image_grid_thw, output_hidden_states=True, ) all_prompt_embeds.append( self.encoding_qwen_image_edit(outputs, image_inputs) ) if batch.do_classifier_free_guidance: all_neg_prompt_embeds.append( self.encoding_qwen_image_edit(neg_outputs, neg_image_inputs) ) if all_prompt_embeds: batch.prompt_embeds.append(torch.cat(all_prompt_embeds, dim=0)) if all_neg_prompt_embeds: batch.negative_prompt_embeds.append(torch.cat(all_neg_prompt_embeds, dim=0)) self.offload_model() return batch def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify image encoding stage inputs.""" result = VerificationResult() if batch.debug: logger.debug(f"{batch.condition_image=}") logger.debug(f"{batch.image_embeds=}") result.add_check("pil_image", batch.condition_image, V.not_none) result.add_check("image_embeds", batch.image_embeds, V.is_list) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify image encoding stage outputs.""" result = VerificationResult() # result.add_check("image_embeds", batch.image_embeds, V.list_of_tensors_dims(3)) return result class ImageVAEEncodingStage(PipelineStage): """ Stage for encoding pixel representations into latent space. This stage handles the encoding of pixel representations into the final input format (e.g., image_latents). """ def __init__(self, vae: ParallelTiledVAE, **kwargs) -> None: super().__init__() self.vae: ParallelTiledVAE = vae def load_model(self): self.vae = self.vae.to(get_local_torch_device()) def offload_model(self): if self.server_args.vae_cpu_offload: self.vae = self.vae.to("cpu") def forward( self, batch: Req, server_args: ServerArgs, ) -> Req: """ Encode pixel representations into latent space. """ if batch.condition_image is None: return batch self.load_model() num_frames = batch.num_frames images = ( batch.vae_image if batch.vae_image is not None else batch.condition_image ) if not isinstance(images, list): images = [images] all_image_latents = [] prepare_condition_image_latent_ids = getattr( server_args.pipeline_config, "prepare_condition_image_latent_ids", None ) condition_latents = [] if callable(prepare_condition_image_latent_ids) else None for image in images: image = self.preprocess( image, ).to(get_local_torch_device(), dtype=torch.float32) # (B, C, H, W) -> (B, C, 1, H, W) image = image.unsqueeze(2) if num_frames == 1: video_condition = image else: video_condition = torch.cat( [ image, image.new_zeros( image.shape[0], image.shape[1], num_frames - 1, image.shape[3], image.shape[4], ), ], dim=2, ) video_condition = video_condition.to( device=get_local_torch_device(), dtype=torch.float32 ) # Setup VAE precision vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] vae_autocast_enabled = ( vae_dtype != torch.float32 ) and not server_args.disable_autocast # Encode Image with torch.autocast( device_type=current_platform.device_type, dtype=vae_dtype, enabled=vae_autocast_enabled, ): if server_args.pipeline_config.vae_tiling: self.vae.enable_tiling() # if server_args.vae_sp: # self.vae.enable_parallel() if not vae_autocast_enabled: video_condition = video_condition.to(vae_dtype) latent_dist: DiagonalGaussianDistribution = self.vae.encode( video_condition ) # for auto_encoder from diffusers if isinstance(latent_dist, AutoencoderKLOutput): latent_dist = latent_dist.latent_dist generator = batch.generator if generator is None: raise ValueError("Generator must be provided") sample_mode = server_args.pipeline_config.vae_config.encode_sample_mode() latent_condition = self.retrieve_latents( latent_dist, generator, sample_mode=sample_mode ) latent_condition = server_args.pipeline_config.postprocess_vae_encode( latent_condition, self.vae ) scaling_factor, shift_factor = ( server_args.pipeline_config.get_decode_scale_and_shift( device=latent_condition.device, dtype=latent_condition.dtype, vae=self.vae, ) ) # apply shift & scale if needed if isinstance(shift_factor, torch.Tensor): shift_factor = shift_factor.to(latent_condition.device) if isinstance(scaling_factor, torch.Tensor): scaling_factor = scaling_factor.to(latent_condition.device) latent_condition -= shift_factor latent_condition = latent_condition * scaling_factor if condition_latents is not None: condition_latents.append(latent_condition) image_latent = server_args.pipeline_config.postprocess_image_latent( latent_condition, batch ) all_image_latents.append(image_latent) batch.image_latent = torch.cat(all_image_latents, dim=1) if condition_latents is not None: prepare_condition_image_latent_ids(condition_latents, batch) self.offload_model() return batch def retrieve_latents( self, encoder_output: DiagonalGaussianDistribution, generator: torch.Generator | None = None, sample_mode: str = "sample", ): if sample_mode == "sample": return encoder_output.sample(generator) elif sample_mode == "argmax": return encoder_output.mode() else: raise AttributeError("Could not access latents of provided encoder_output") def preprocess( self, image: torch.Tensor | PIL.Image.Image, ) -> torch.Tensor: if isinstance(image, PIL.Image.Image): image = pil_to_numpy(image) # to np image = numpy_to_pt(image) # to pt do_normalize = True if image.min() < 0: do_normalize = False if do_normalize: image = normalize(image) return image def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify encoding stage inputs.""" result = VerificationResult() assert batch.condition_image is None or ( isinstance(batch.condition_image, PIL.Image.Image) or isinstance(batch.condition_image, torch.Tensor) or isinstance(batch.condition_image, list) ) assert batch.height is not None and isinstance(batch.height, int) assert batch.width is not None and isinstance(batch.width, int) assert batch.num_frames is not None and isinstance(batch.num_frames, int) result.add_check("generator", batch.generator, V.generator_or_list_generators) result.add_check("height", batch.height, V.positive_int) result.add_check("width", batch.width, V.positive_int) result.add_check("num_frames", batch.num_frames, V.positive_int) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify encoding stage outputs.""" result = VerificationResult() # result.add_check( # "image_latent", batch.image_latent, [V.is_tensor, V.with_dims(5)] # ) return result ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Input validation stage for diffusion pipelines. """ import numpy as np import torch import torchvision.transforms.functional as TF from PIL import Image from sglang.multimodal_gen.configs.pipeline_configs import WanI2V480PConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig from sglang.multimodal_gen.runtime.models.vision_utils import load_image, load_video from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators, VerificationResult, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import best_output_size logger = init_logger(__name__) # Alias for convenience V = StageValidators # TODO: since this might change sampling params after logging, should be do this beforehand? class InputValidationStage(PipelineStage): """ Stage for validating and preparing inputs for diffusion pipelines. This stage validates that all required inputs are present and properly formatted before proceeding with the diffusion process. In this stage, input image and output image may be resized """ def __init__(self, vae_image_processor=None): super().__init__() self.vae_image_processor = vae_image_processor @staticmethod def _calculate_dimensions_from_area( max_area: float, aspect_ratio: float, mod_value: int ) -> tuple[int, int]: """ Calculate output dimensions based on maximum area and aspect ratio. Args: max_area: Maximum area constraint for the output aspect_ratio: Target aspect ratio (height/width) mod_value: Value to round dimensions to (typically vae_scale * patch_size) Returns: Tuple of (width, height) rounded to multiples of mod_value """ height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value return width, height def _generate_seeds(self, batch: Req, server_args: ServerArgs): """Generate seeds for the inference""" seed = batch.seed num_videos_per_prompt = batch.num_outputs_per_prompt assert seed is not None seeds = [seed + i for i in range(num_videos_per_prompt)] batch.seeds = seeds # Create generators based on generator_device parameter # Note: This will overwrite any existing batch.generator generator_device = batch.generator_device if generator_device == "cpu": device_str = "cpu" else: device_str = current_platform.device_type batch.generator = [ torch.Generator(device_str).manual_seed(seed) for seed in seeds ] def preprocess_condition_image( self, batch: Req, server_args: ServerArgs, condition_image_width, condition_image_height, ): """ preprocess condition image NOTE: condition image resizing is only allowed in InputValidationStage """ if batch.condition_image is not None and ( server_args.pipeline_config.task_type == ModelTaskType.I2I or server_args.pipeline_config.task_type == ModelTaskType.TI2I ): # calculate new condition image size if not isinstance(batch.condition_image, list): batch.condition_image = [batch.condition_image] processed_images = [] final_image = batch.condition_image[-1] config = server_args.pipeline_config config.preprocess_vae_image(batch, self.vae_image_processor) for img in batch.condition_image: size = config.calculate_condition_image_size(img, img.width, img.height) if size is not None: width, height = size img, _ = config.preprocess_condition_image( img, width, height, self.vae_image_processor ) processed_images.append(img) batch.condition_image = processed_images calculated_size = config.prepare_calculated_size(final_image) # adjust output image size if calculated_size is not None: calculated_width, calculated_height = calculated_size width = batch.width or calculated_width height = batch.height or calculated_height multiple_of = ( server_args.pipeline_config.vae_config.get_vae_scale_factor() * 2 ) width = width // multiple_of * multiple_of height = height // multiple_of * multiple_of batch.width = width batch.height = height elif server_args.pipeline_config.task_type == ModelTaskType.TI2V: if server_args.pipeline_config.skip_input_image_preprocess: return # duplicate with vae_image_processor # further processing for ti2v task if isinstance( batch.condition_image, list ): # not support multi image input yet. batch.condition_image = batch.condition_image[0] img = batch.condition_image ih, iw = img.height, img.width patch_size = server_args.pipeline_config.dit_config.arch_config.patch_size vae_stride = ( server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial ) dh, dw = patch_size[1] * vae_stride, patch_size[2] * vae_stride max_area = 704 * 1280 ow, oh = best_output_size(iw, ih, dw, dh, max_area) scale = max(ow / iw, oh / ih) img = img.resize((round(iw * scale), round(ih * scale)), Image.LANCZOS) logger.debug("resized condition image to: %sx%s", img.height, img.width) # center-crop x1 = (img.width - ow) // 2 y1 = (img.height - oh) // 2 img = img.crop((x1, y1, x1 + ow, y1 + oh)) assert img.width == ow and img.height == oh # to tensor img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device).unsqueeze(1) img = img.unsqueeze(0) batch.height = oh batch.width = ow # TODO: should we store in a new field: pixel values? batch.condition_image = img elif isinstance(server_args.pipeline_config, WanI2V480PConfig): # TODO: could we merge with above? # resize image only, Wan2.1 I2V if isinstance(batch.condition_image, list): batch.condition_image = batch.condition_image[ 0 ] # not support multi image input yet. max_area = server_args.pipeline_config.max_area aspect_ratio = condition_image_height / condition_image_width mod_value = ( server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial * server_args.pipeline_config.dit_config.arch_config.patch_size[1] ) width, height = self._calculate_dimensions_from_area( max_area, aspect_ratio, mod_value ) batch.condition_image = batch.condition_image.resize((width, height)) batch.height = height batch.width = width elif issubclass(type(server_args.pipeline_config), MOVAPipelineConfig): # resize image only, MOVA image = batch.condition_image if isinstance(image, list): image = image[0] # not support multi image input yet. max_area = server_args.pipeline_config.max_area if hasattr(batch, "height") and hasattr(batch, "width"): aspect_ratio = batch.height / batch.width else: aspect_ratio = ( batch.sampling_params.height / batch.sampling_params.width ) mod_value = ( server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial * server_args.pipeline_config.dit_config.arch_config.patch_size[1] ) width, height = self._calculate_dimensions_from_area( max_area, aspect_ratio, mod_value ) config = server_args.pipeline_config image, (final_w, final_h) = ( server_args.pipeline_config.preprocess_condition_image( image, width, height, self.vae_image_processor ) ) batch.condition_image = image batch.width = final_w batch.height = final_h def forward( self, batch: Req, server_args: ServerArgs, ) -> Req: """ Validate and prepare inputs. """ self._generate_seeds(batch, server_args) if ( server_args.pipeline_config.task_type == ModelTaskType.I2M and batch.num_inference_steps is None and hasattr(server_args.pipeline_config, "shape_num_inference_steps") ): batch.num_inference_steps = ( server_args.pipeline_config.shape_num_inference_steps ) # Ensure prompt is properly formatted (I2M can be image-only) if ( server_args.pipeline_config.task_type != ModelTaskType.I2M and batch.prompt is None and batch.prompt_embeds is None ): raise ValueError("Either `prompt` or `prompt_embeds` must be provided") # Ensure negative prompt is properly formatted if using classifier-free guidance if ( batch.do_classifier_free_guidance and batch.negative_prompt is None and batch.negative_prompt_embeds is None ): raise ValueError( "For classifier-free guidance, either `negative_prompt` or " "`negative_prompt_embeds` must be provided" ) # Validate number of inference steps if batch.num_inference_steps <= 0: raise ValueError( f"Number of inference steps must be positive, but got {batch.num_inference_steps}" ) # Validate guidance scale if using classifier-free guidance if batch.do_classifier_free_guidance and batch.guidance_scale < 0: raise ValueError( f"Guidance scale must be positive, but got {batch.guidance_scale}" ) # for i2v, get image from image_path # @TODO(Wei) hard-coded for wan2.2 5b ti2v for now. Should put this in image_encoding stage if batch.image_path is not None: if isinstance(batch.image_path, list): batch.condition_image = [] for path in batch.image_path: if path.endswith(".mp4"): image = load_video(path)[0] else: image = load_image(path) batch.condition_image.append(image) # Use the first image for size reference condition_image_width = batch.condition_image[0].width condition_image_height = batch.condition_image[0].height batch.original_condition_image_size = ( condition_image_width, condition_image_height, ) else: if batch.image_path.endswith(".mp4"): image = load_video(batch.image_path)[0] else: image = load_image(batch.image_path) batch.condition_image = image condition_image_width, condition_image_height = ( image.width, image.height, ) batch.original_condition_image_size = image.size if server_args.pipeline_config.task_type != ModelTaskType.I2M: self.preprocess_condition_image( batch, server_args, condition_image_width, condition_image_height ) # if height or width is not specified at this point, set default to 720p default_height = 720 default_width = 1280 if batch.height is None and batch.width is None: batch.height = default_height batch.width = default_width elif batch.height is None: batch.height = batch.width * default_height // default_width elif batch.width is None: batch.width = batch.height * default_width // default_height return batch def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify input validation stage inputs.""" result = VerificationResult() result.add_check("seed", batch.seed, [V.not_none, V.non_negative_int]) result.add_check( "num_videos_per_prompt", batch.num_outputs_per_prompt, V.positive_int ) if server_args.pipeline_config.task_type != ModelTaskType.I2M: result.add_check( "prompt_or_embeds", None, lambda _: V.string_or_list_strings(batch.prompt) or V.list_not_empty(batch.prompt_embeds), ) if server_args.pipeline_config.task_type != ModelTaskType.I2M: result.add_check( "num_inference_steps", batch.num_inference_steps, V.positive_int ) else: result.add_check( "num_inference_steps", batch.num_inference_steps, lambda x: x is None or V.positive_int(x), ) result.add_check( "guidance_scale", batch.guidance_scale, lambda x: not batch.do_classifier_free_guidance or V.non_negative_float(x), ) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify input validation stage outputs.""" result = VerificationResult() result.add_check("height", batch.height, V.positive_int) result.add_check("width", batch.width, V.positive_int) result.add_check("seeds", batch.seeds, V.list_not_empty) result.add_check("generator", batch.generator, V.generator_or_list_generators) return result ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Latent preparation stage for diffusion pipelines. """ from diffusers.utils.torch_utils import randn_tensor from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class LatentPreparationStage(PipelineStage): """ Stage for preparing initial latent variables for the diffusion process. This stage handles the preparation of the initial latent variables that will be denoised during the diffusion process. """ def __init__(self, scheduler, transformer) -> None: super().__init__() self.scheduler = scheduler self.transformer = transformer def forward( self, batch: Req, server_args: ServerArgs, ) -> Req: """ Prepare initial latent variables for the diffusion process. Returns: The batch with prepared latent variables. """ # Adjust video length based on VAE version if needed latent_num_frames = self.adjust_video_length(batch, server_args) batch_size = batch.batch_size # Get required parameters dtype = batch.prompt_embeds[0].dtype device = get_local_torch_device() generator = batch.generator latents = batch.latents num_frames = ( latent_num_frames if latent_num_frames is not None else batch.num_frames ) height = batch.height width = batch.width # TODO(will): remove this once we add input/output validation for stages if height is None or width is None: raise ValueError("Height and width must be provided") # Validate generator if it's a list if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) # Generate or use provided latents if latents is None: shape = server_args.pipeline_config.prepare_latent_shape( batch, batch_size, num_frames ) latents = randn_tensor( shape, generator=generator, device=device, dtype=dtype ) latent_ids = server_args.pipeline_config.maybe_prepare_latent_ids(latents) if latent_ids is not None: batch.latent_ids = latent_ids.to(device=device) latents = server_args.pipeline_config.maybe_pack_latents( latents, batch_size, batch ) else: latents = latents.to(device) # Scale the initial noise if needed if hasattr(self.scheduler, "init_noise_sigma"): latents = latents * self.scheduler.init_noise_sigma # Update batch with prepared latents batch.latents = latents batch.raw_latent_shape = latents.shape return batch def adjust_video_length(self, batch: Req, server_args: ServerArgs) -> int: """ Adjust video length based on VAE version. """ video_length = batch.num_frames latent_num_frames = video_length use_temporal_scaling_frames = ( server_args.pipeline_config.vae_config.use_temporal_scaling_frames ) if use_temporal_scaling_frames: temporal_scale_factor = ( server_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio ) latent_num_frames = (video_length - 1) // temporal_scale_factor + 1 return int(latent_num_frames) def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify latent preparation stage inputs.""" result = VerificationResult() result.add_check( "prompt_or_embeds", None, lambda _: V.string_or_list_strings(batch.prompt) or V.list_not_empty(batch.prompt_embeds), ) result.add_check("prompt_embeds", batch.prompt_embeds, V.list_of_tensors) result.add_check( "num_videos_per_prompt", batch.num_outputs_per_prompt, V.positive_int ) result.add_check("generator", batch.generator, V.generator_or_list_generators) result.add_check("num_frames", batch.num_frames, V.positive_int) result.add_check("height", batch.height, V.positive_int) result.add_check("width", batch.width, V.positive_int) result.add_check("latents", batch.latents, V.none_or_tensor) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify latent preparation stage outputs.""" result = VerificationResult() if batch.debug: logger.debug(f"{batch.raw_latent_shape=}") # disable temporarily for image-generation models # result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) result.add_check("raw_latent_shape", batch.raw_latent_shape, V.is_tuple) return result ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py ================================================ import torch from diffusers.utils.torch_utils import randn_tensor from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import ( LatentPreparationStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class LTX2AVLatentPreparationStage(LatentPreparationStage): """ LTX-2 specific latent preparation stage that handles both video and audio latents. """ def __init__(self, scheduler, transformer=None, audio_vae=None): super().__init__(scheduler, transformer) self.audio_vae = audio_vae def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify latent preparation stage inputs.""" result = VerificationResult() result.add_check( "prompt_or_embeds", None, lambda _: V.string_or_list_strings(batch.prompt) or V.list_not_empty(batch.prompt_embeds) or V.is_tensor(batch.prompt_embeds), ) if isinstance(batch.prompt_embeds, list): result.add_check("prompt_embeds", batch.prompt_embeds, V.list_of_tensors) else: result.add_check("prompt_embeds", batch.prompt_embeds, V.is_tensor) result.add_check( "num_videos_per_prompt", batch.num_outputs_per_prompt, V.positive_int ) result.add_check("generator", batch.generator, V.generator_or_list_generators) result.add_check("num_frames", batch.num_frames, V.positive_int) result.add_check("height", batch.height, V.positive_int) result.add_check("width", batch.width, V.positive_int) result.add_check("latents", batch.latents, V.none_or_tensor) return result def forward(self, batch: Req, server_args: ServerArgs) -> Req: # 1. Prepare Video Latents using base class logic # This sets batch.latents and batch.raw_latent_shape batch = super().forward(batch, server_args) # 2. Prepare Audio Latents (optional) # Default to True if not specified try: generate_audio = batch.generate_audio except AttributeError: generate_audio = True if not generate_audio: batch.audio_latents = None batch.raw_audio_latent_shape = None return batch device = get_local_torch_device() if isinstance(batch.prompt_embeds, list) and batch.prompt_embeds: dtype = batch.prompt_embeds[0].dtype elif isinstance(batch.prompt_embeds, torch.Tensor): dtype = batch.prompt_embeds.dtype else: dtype = torch.float16 generator = batch.generator audio_latents = batch.audio_latents batch_size = batch.batch_size num_frames = batch.num_frames if audio_latents is None: shape = server_args.pipeline_config.prepare_audio_latent_shape( batch, batch_size, num_frames ) audio_latents = randn_tensor( shape, generator=generator, device=device, dtype=dtype ) else: audio_latents = audio_latents.to(device) audio_latents = server_args.pipeline_config.maybe_pack_audio_latents( audio_latents, batch_size, batch ) # Store in batch batch.audio_latents = audio_latents batch.raw_audio_latent_shape = audio_latents.shape return batch ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/glm_image.py ================================================ import inspect import re import time from math import sqrt from typing import List, Optional, Tuple, Union import numpy as np import PIL import torch from diffusers.image_processor import VaeImageProcessor from diffusers.utils.torch_utils import randn_tensor from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.models.dits.glm_image import GlmImageKVCache from sglang.multimodal_gen.runtime.models.vision_utils import load_image from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) def calculate_shift( image_seq_len, base_seq_len: int = 256, base_shift: float = 0.25, max_shift: float = 0.75, ) -> float: m = (image_seq_len / base_seq_len) ** 0.5 mu = m * max_shift + base_shift return mu def retrieve_timesteps( scheduler, num_inference_steps: Optional[int] = None, device: Optional[Union[str, torch.device]] = None, timesteps: Optional[List[int]] = None, sigmas: Optional[List[float]] = None, **kwargs, ): r""" Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. """ accepts_timesteps = "timesteps" in set( inspect.signature(scheduler.set_timesteps).parameters.keys() ) accepts_sigmas = "sigmas" in set( inspect.signature(scheduler.set_timesteps).parameters.keys() ) if timesteps is not None and sigmas is not None: if not accepts_timesteps and not accepts_sigmas: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" timestep or sigma schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps( timesteps=timesteps, sigmas=sigmas, device=device, **kwargs ) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) elif timesteps is not None and sigmas is None: if not accepts_timesteps: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" timestep schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) elif timesteps is None and sigmas is not None: if not accepts_sigmas: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" sigmas schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) else: scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) timesteps = scheduler.timesteps return timesteps, num_inference_steps # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample", ): if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": return encoder_output.latent_dist.mode() elif hasattr(encoder_output, "latents"): return encoder_output.latents else: raise AttributeError("Could not access latents of provided encoder_output") class GlmImageBeforeDenoisingStage(PipelineStage): r""" Pipeline for text-to-image generation using GLM-Image. This pipeline integrates both the AR (autoregressive) model for token generation and the DiT (diffusion transformer) model for image decoding. Args: vae ([`AutoencoderKL`]): Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. text_encoder ([`T5EncoderModel`]): Frozen text-encoder for glyph embeddings. tokenizer (`PreTrainedTokenizer`): Tokenizer for the text encoder. processor (`AutoProcessor`): Processor for the AR model to handle chat templates and tokenization. vision_language_encoder ([`GlmImageForConditionalGeneration`]): The AR model that generates image tokens from text prompts. transformer ([`GlmImageTransformer2DModel`]): A text conditioned transformer to denoise the encoded image latents (DiT). scheduler ([`SchedulerMixin`]): A scheduler to be used in combination with `transformer` to denoise the encoded image latents. """ def __init__( self, tokenizer, processor, text_encoder, vision_language_encoder, vae, transformer, scheduler, ): super().__init__() self.tokenizer = tokenizer self.processor = processor self.text_encoder = text_encoder self.vision_language_encoder = vision_language_encoder self.vae = vae self.transformer = transformer self.scheduler = scheduler self.vae_scale_factor = ( 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 ) self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) self.default_sample_size = ( self.transformer.config.sample_size if hasattr(self, "transformer") and self.transformer is not None and hasattr(self.transformer.config, "sample_size") else 128 ) def _parse_and_expand_shape_info( self, prompt: str ) -> Tuple[str, int, int, int, int]: """ Parse the shape info from prompt and expand it for AR model. Args: prompt: The prompt containing H W shape specification Returns: Tuple of (expanded_prompt, token_h, token_w, prev_token_h, prev_token_w) """ match = re.search(r"(\d+)\s+(\d+)", prompt) if match is None: raise ValueError( f"Prompt must contain shape info in format 'H W', got: {prompt}" ) token_h, token_w = int(match.group(1)), int(match.group(2)) ratio = token_h / token_w prev_token_h = int(sqrt(ratio) * 16) prev_token_w = int(sqrt(1 / ratio) * 16) old_shape = f"{token_h} {token_w}" new_shape = ( f"{token_h} {token_w}{prev_token_h} {prev_token_w}" ) expanded_prompt = prompt.replace(old_shape, new_shape) return expanded_prompt, token_h, token_w, prev_token_h, prev_token_w def _build_image_grid_thw( self, token_h: int, token_w: int, prev_token_h: int, prev_token_w: int, existing_grid: Optional[torch.Tensor] = None, device: Optional[torch.device] = None, ) -> torch.Tensor: """ Build image grid tensor for AR model. For text-to-image: creates grid for large image + small image For image-to-image: appends new image to existing grid """ if existing_grid is None or existing_grid.numel() == 0: # Text-to-image: large image + small image return torch.tensor( [ [1, token_h, token_w], [1, prev_token_h, prev_token_w], ], device=device, ) else: # Image-to-image: append to existing return torch.cat( [existing_grid, torch.tensor([[1, token_h, token_w]], device=device)], dim=0, ) def _calculate_ar_generation_params( self, token_h: int, token_w: int, prev_token_h: int, prev_token_w: int, is_text_to_image: bool, ) -> Tuple[int, int]: """ Calculate max_new_tokens and large_image_start_offset for AR generation. """ large_image_tokens = token_h * token_w small_image_tokens = prev_token_h * prev_token_w if is_text_to_image: max_new_tokens = small_image_tokens + large_image_tokens + 1 large_image_start_offset = small_image_tokens else: max_new_tokens = large_image_tokens + 1 large_image_start_offset = 0 return max_new_tokens, large_image_start_offset def _extract_large_image_tokens( self, outputs: torch.Tensor, input_length: int, large_image_start_offset: int, large_image_tokens: int, ) -> torch.Tensor: """ Extract the large image tokens from AR model output. """ generated_tokens = outputs[0][input_length:] large_image_start = large_image_start_offset large_image_end = large_image_start + large_image_tokens return generated_tokens[large_image_start:large_image_end] def _upsample_d32_to_d16( self, token_ids: torch.Tensor, token_h: int, token_w: int ) -> torch.Tensor: """ Upsample token IDs from d32 format to d16 format. AR model generates tokens at d32 resolution (each token = 32x32 pixels). DiT expects tokens at d16 resolution (each token = 16x16 pixels). This function performs 2x nearest-neighbor upsampling. Args: token_ids: Token IDs of shape [N] where N = token_h * token_w token_h: Height in d32 token units token_w: Width in d32 token units Returns: Upsampled token IDs of shape [1, N*4] where N*4 = (token_h*2) * (token_w*2) """ # Reshape to spatial format: [1, 1, H, W] token_ids = token_ids.view(1, 1, token_h, token_w) # 2x nearest-neighbor upsampling token_ids = torch.nn.functional.interpolate( token_ids.float(), scale_factor=2, mode="nearest" ).to(dtype=torch.long) # Flatten back to [1, H*W*4] token_ids = token_ids.view(1, -1) return token_ids @staticmethod def _compute_generation_params( image_grid_thw, is_text_to_image: bool, ): grid_sizes = [] grid_hw = [] for i in range(image_grid_thw.shape[0]): t, h, w = image_grid_thw[i].tolist() grid_sizes.append(int(h * w)) grid_hw.append((int(h), int(w))) if not is_text_to_image: max_new_tokens = grid_sizes[-1] + 1 large_image_start_offset = 0 target_grid_h, target_grid_w = grid_hw[-1] else: total_tokens = sum(grid_sizes) max_new_tokens = total_tokens + 1 large_image_start_offset = sum(grid_sizes[1:]) target_grid_h, target_grid_w = grid_hw[0] return max_new_tokens, large_image_start_offset, target_grid_h, target_grid_w @staticmethod def _upsample_token_ids( token_ids: torch.Tensor, token_h: int, token_w: int ) -> torch.Tensor: token_ids = token_ids.view(1, 1, token_h, token_w) token_ids = torch.nn.functional.interpolate( token_ids.float(), scale_factor=2, mode="nearest" ).to(dtype=torch.long) token_ids = token_ids.view(1, -1) return token_ids def generate_prior_tokens( self, prompt: str, height: int, width: int, image: Optional[List[PIL.Image.Image]] = None, factor: int = 32, ) -> Tuple[torch.Tensor, int, int]: """ Generate prior tokens using the AR (vision_language_encoder) model. Args: prompt: The text prompt with shape info (e.g., "description36 24") condition_images: Optional list of condition images for i2i Returns: Tuple of (prior_token_ids, pixel_height, pixel_width) - prior_token_ids: Upsampled to d16 format, shape [1, token_h*token_w*4] - pixel_height: Image height in pixels - pixel_width: Image width in pixels """ device = self.vision_language_encoder.device height = (height // factor) * factor width = (width // factor) * factor is_text_to_image = image is None or len(image) == 0 # Build messages for processor content = [] if image is not None: for img in image: content.append({"type": "image", "image": img}) content.append({"type": "text", "text": prompt}) messages = [{"role": "user", "content": content}] inputs = self.processor.apply_chat_template( messages, tokenize=True, target_h=height, target_w=width, return_dict=True, return_tensors="pt", ).to(device) image_grid_thw = inputs.get("image_grid_thw") max_new_tokens, large_image_offset, token_h, token_w = ( self._compute_generation_params( image_grid_thw=image_grid_thw, is_text_to_image=is_text_to_image ) ) prior_token_image_ids = None if image is not None: prior_token_image_embed = self.vision_language_encoder.get_image_features( inputs["pixel_values"], image_grid_thw[:-1] ) prior_token_image_embed = torch.cat(prior_token_image_embed, dim=0) prior_token_image_ids = self.vision_language_encoder.get_image_tokens( prior_token_image_embed, image_grid_thw[:-1] ) # For GLM-Image, greedy decoding is not allowed; it may cause repetitive outputs. # max_new_tokens must be exactly grid_h * grid_w + 1 (the +1 is for EOS). outputs = self.vision_language_encoder.generate( **inputs, max_new_tokens=max_new_tokens, do_sample=True, ) prior_token_ids_d32 = self._extract_large_image_tokens( outputs, inputs["input_ids"].shape[-1], large_image_offset, token_h * token_w, ) prior_token_ids = self._upsample_token_ids( prior_token_ids_d32, token_h, token_w ) return prior_token_ids, prior_token_image_ids def get_glyph_texts(self, prompt): prompt = prompt[0] if isinstance(prompt, list) else prompt ocr_texts = ( re.findall(r"'([^']*)'", prompt) + re.findall(r"“([^“”]*)”", prompt) + re.findall(r'"([^"]*)"', prompt) + re.findall(r"「([^「」]*)」", prompt) ) return ocr_texts def _get_glyph_embeds( self, prompt: Union[str, List[str]] = None, max_sequence_length: int = 2048, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, ): device = device or self._execution_device dtype = dtype or self.text_encoder.dtype glyph_texts = self.get_glyph_texts(prompt) input_ids = self.tokenizer( glyph_texts if len(glyph_texts) > 0 else [""], max_length=max_sequence_length, truncation=True, ).input_ids input_ids = [ [self.tokenizer.pad_token_id] * ((len(input_ids) + 1) % 2) + input_ids_ for input_ids_ in input_ids ] max_length = max(len(input_ids_) for input_ids_ in input_ids) attention_mask = torch.tensor( [ [1] * len(input_ids_) + [0] * (max_length - len(input_ids_)) for input_ids_ in input_ids ], device=device, ) input_ids = torch.tensor( [ input_ids_ + [self.tokenizer.pad_token_id] * (max_length - len(input_ids_)) for input_ids_ in input_ids ], device=device, ) outputs = self.text_encoder(input_ids, attention_mask=attention_mask) glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0) return glyph_embeds.to(device=device, dtype=dtype) def encode_prompt( self, prompt: Union[str, List[str]], do_classifier_free_guidance: bool = True, prompt_embeds: Optional[torch.Tensor] = None, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, max_sequence_length: int = 2048, ): r""" Encodes the prompt into text encoder hidden states. Args: prompt (`str` or `List[str]`, *optional*): prompt to be encoded do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): Whether to use classifier free guidance or not. prompt_embeds (`torch.Tensor`, *optional*): Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument. device: (`torch.device`, *optional*): torch device dtype: (`torch.dtype`, *optional*): torch dtype max_sequence_length (`int`, defaults to `2048`): Maximum sequence length in encoded prompt. Can be set to other values but may lead to poorer results. """ device = device or self._execution_device prompt = [prompt] if isinstance(prompt, str) else prompt if prompt is not None: batch_size = len(prompt) else: batch_size = prompt_embeds.shape[0] if prompt_embeds is None: prompt_embeds = self._get_glyph_embeds( prompt, max_sequence_length, device, dtype ) seq_len = prompt_embeds.size(1) prompt_embeds = prompt_embeds.repeat(1, 1, 1) prompt_embeds = prompt_embeds.view(1, seq_len, -1) negative_prompt_embeds = None if do_classifier_free_guidance: negative_prompt = "" negative_prompt = ( batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt ) if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" f" {type(prompt)}." ) elif batch_size != len(negative_prompt): raise ValueError( f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" " the batch size of `prompt`." ) negative_prompt_embeds = self._get_glyph_embeds( negative_prompt, max_sequence_length, device, dtype ) seq_len = negative_prompt_embeds.size(1) negative_prompt_embeds = negative_prompt_embeds.repeat(1, 1, 1) negative_prompt_embeds = negative_prompt_embeds.view(1, seq_len, -1) return prompt_embeds, negative_prompt_embeds def prepare_latents( self, batch_size, num_channels_latents, height, width, dtype, device, generator, ): shape = ( batch_size, num_channels_latents, int(height) // self.vae_scale_factor, int(width) // self.vae_scale_factor, ) if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) return latents def check_inputs( self, prompt, height, width, callback_on_step_end_tensor_inputs, prompt_embeds=None, ): if ( height is not None and height % (self.vae_scale_factor * self.transformer.config.patch_size) != 0 or width is not None and width % (self.transformer.config.patch_size) != 0 ): logger.warning( f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly" ) if callback_on_step_end_tensor_inputs is not None and not all( k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs ): raise ValueError( f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" ) if prompt is not None and prompt_embeds is not None: raise ValueError( f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" " only forward one of the two." ) elif prompt is None and prompt_embeds is None: raise ValueError( "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." ) elif prompt is not None and ( not isinstance(prompt, str) and not isinstance(prompt, list) ): raise ValueError( f"`prompt` has to be of type `str` or `list` but is {type(prompt)}" ) @property def guidance_scale(self): return self._guidance_scale @property def do_classifier_free_guidance(self): return self._guidance_scale > 1 @property def num_timesteps(self): return self._num_timesteps @property def attention_kwargs(self): return self._attention_kwargs @property def current_timestep(self): return self._current_timestep @property def interrupt(self): return self._interrupt @torch.no_grad() def forward( self, batch: Req, server_args: ServerArgs, ) -> Req: guidance_scale = batch.guidance_scale prompt = batch.prompt num_inference_steps = batch.num_inference_steps if batch.image_path is not None: ar_condition_images = [ load_image(img_path) for img_path in batch.image_path ] else: ar_condition_images = None height = batch.height width = batch.width device = get_local_torch_device() max_sequence_length = 1024 generator = torch.Generator(device=device).manual_seed(batch.seed) attention_kwargs = {} prompt_embeds = None do_classifier_free_guidance = True dtype = torch.bfloat16 self._guidance_scale = guidance_scale self._current_timestep = None self._interrupt = False batch_size = 1 device = get_local_torch_device() if ar_condition_images is not None: height = height or ar_condition_images[0].height width = width or ar_condition_images[0].width time_start = time.time() prior_token_id, prior_token_image_ids = self.generate_prior_tokens( prompt=prompt, image=ar_condition_images, height=height, width=width, ) prior_token_id = prior_token_id.to(device=device) time_end = time.time() logger.info(f"generate_prior_tokens time: {time_end - time_start}") # 3. Encode input prompt prompt_embeds, negative_prompt_embeds = self.encode_prompt( prompt, do_classifier_free_guidance, prompt_embeds=prompt_embeds, max_sequence_length=max_sequence_length, device=device, dtype=dtype, ) # 4. process images if ar_condition_images is not None: preprocessed_condition_images = [] for img in ar_condition_images: image_height, image_width = ( img.size[::-1] if isinstance(img, PIL.Image.Image) else img.shape[:2] ) multiple_of = self.vae_scale_factor * self.transformer.config.patch_size image_height = (image_height // multiple_of) * multiple_of image_width = (image_width // multiple_of) * multiple_of img = self.image_processor.preprocess( img, height=image_height, width=image_width ) preprocessed_condition_images.append(img) ar_condition_images = preprocessed_condition_images # 5. Prepare latents and (optional) condition_images kv cache latent_channels = self.transformer.config.in_channels latents = self.prepare_latents( batch_size=1, num_channels_latents=latent_channels, height=height, width=width, dtype=torch.float32, device=device, generator=generator, ) kv_caches = GlmImageKVCache(num_layers=self.transformer.config.num_layers) if ar_condition_images is not None: latents_mean = torch.tensor(self.vae.config.latents_mean).view( 1, self.vae.config.latent_channels, 1, 1 ) latents_std = torch.tensor(self.vae.config.latents_std).view( 1, self.vae.config.latent_channels, 1, 1 ) latents_mean = latents_mean.to(device=device, dtype=prompt_embeds.dtype) latents_std = latents_std.to(device=device, dtype=prompt_embeds.dtype) for condition_image, condition_image_prior_token_id in zip( ar_condition_images, prior_token_image_ids ): condition_image = condition_image.to( device=device, dtype=prompt_embeds.dtype ) condition_latent = retrieve_latents( self.vae.encode(condition_image), generator=generator, sample_mode="argmax", ) condition_latent = (condition_latent - latents_mean) / latents_std # Do not remove. # It would be use to run the reference image through a # forward pass at timestep 0 and keep the KV cache. with set_forward_context(current_timestep=1, attn_metadata=None): _ = self.transformer( hidden_states=condition_latent, encoder_hidden_states=torch.zeros_like(prompt_embeds)[ :1, :0, ... ], prior_token_id=condition_image_prior_token_id, prior_token_drop=torch.full_like( condition_image_prior_token_id, False, dtype=torch.bool ), timestep=torch.zeros((1,), device=device), target_size=torch.tensor( [condition_image.shape[-2:]], device=device ), crop_coords=torch.zeros((1, 2), device=device), attention_kwargs=attention_kwargs, kv_caches=kv_caches, kv_caches_mode="write", ) # 6. Prepare additional timestep conditions target_size = (height, width) target_size = torch.tensor( [target_size], dtype=prompt_embeds.dtype, device=device ) crops_coords_top_left = torch.tensor( [(0, 0)], dtype=prompt_embeds.dtype, device=device ) # Prepare timesteps image_seq_len = ( (height // self.vae_scale_factor) * (width // self.vae_scale_factor) ) // (self.transformer.config.patch_size**2) timesteps = np.linspace( self.scheduler.config.num_train_timesteps, 1.0, num_inference_steps + 1 )[:-1] timesteps = timesteps.astype(np.int64).astype(np.float32) sigmas = timesteps / self.scheduler.config.num_train_timesteps mu = calculate_shift( image_seq_len, self.scheduler.config.get("base_image_seq_len", 256), self.scheduler.config.get("base_shift", 0.25), self.scheduler.config.get("max_shift", 0.75), ) timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, device, timesteps, sigmas, mu=mu ) self._num_timesteps = len(timesteps) # 7. Prepare for denoising loop batch.prompt_embeds = [prompt_embeds] batch.negative_prompt_embeds = [negative_prompt_embeds] batch.latents = latents batch.timesteps = timesteps batch.num_inference_steps = num_inference_steps batch.sigmas = sigmas.tolist() # Convert numpy array to list for validation batch.generator = generator batch.raw_latent_shape = latents.shape batch.prior_token_id = prior_token_id batch.prior_token_drop_cond = torch.full_like( prior_token_id, False, dtype=torch.bool ) batch.prior_token_drop_uncond = torch.full_like( prior_token_id, True, dtype=torch.bool ) batch.target_size = target_size batch.crop_coords = crops_coords_top_left batch.kv_caches = kv_caches batch.height = height batch.width = width return batch ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_decoding.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Helios-specific decoding stage. Decodes latent chunks one at a time (matching diffusers HeliosPipeline behavior) to avoid temporal artifacts at chunk boundaries caused by Wan VAE's causal convolutions. """ import torch from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import ( DecodingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class HeliosDecodingStage(DecodingStage): """ Helios-specific decoding stage that decodes latent chunks independently. The Wan VAE uses causal 3D convolutions with feature caching. When decoding the full latent sequence at once, the causal conv processes all frames with continuous context, producing a different number of output frames per latent frame compared to chunk-by-chunk decoding. This causes temporal misalignment and visible seams at chunk boundaries. This stage decodes each chunk's latents separately (matching diffusers' HeliosPipeline behavior) and concatenates the results in pixel space. """ @torch.no_grad() def forward( self, batch: Req, server_args: ServerArgs, ) -> OutputBatch: latent_chunks = getattr(batch, "latent_chunks", None) if latent_chunks is None or len(latent_chunks) <= 1: # No chunked latents or single chunk — use standard decode return super().forward(batch, server_args) # Load VAE if needed self.load_model() # Decode each chunk separately and concatenate in pixel space video_chunks = [] for chunk_latents in latent_chunks: chunk_video = self.decode(chunk_latents, server_args) video_chunks.append(chunk_video) frames = torch.cat(video_chunks, dim=2) frames = server_args.pipeline_config.post_decoding(frames, server_args) output_batch = OutputBatch( output=frames, trajectory_timesteps=batch.trajectory_timesteps, trajectory_latents=batch.trajectory_latents, trajectory_decoded=None, metrics=batch.metrics, ) self.offload_model() return output_batch ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Helios-specific chunked denoising stage. Implements Stage 1 chunked denoising with multi-term memory history and CFG Zero Star guidance. VAE decoding is handled by the standard DecodingStage downstream. """ import math import numpy as np import torch import torch.nn.functional as F from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( PipelineStage, StageParallelismType, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) def optimized_scale(positive_flat, negative_flat): """CFG Zero Star: compute optimal guidance scale.""" positive_flat = positive_flat.float() negative_flat = negative_flat.float() dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8 return dot_product / squared_norm def calculate_shift( image_seq_len, base_seq_len: int = 256, max_seq_len: int = 4096, base_shift: float = 0.5, max_shift: float = 1.15, ): m = (max_shift - base_shift) / (max_seq_len - base_seq_len) b = base_shift - m * base_seq_len mu = image_seq_len * m + b return mu def sample_block_noise( batch_size, channel, num_frames, height, width, gamma, patch_size=(1, 2, 2) ): """Generate spatially-correlated block noise for pyramid SR.""" _, ph, pw = patch_size block_size = ph * pw # Explicitly use CPU to avoid requiring MAGMA for cholesky on ROCm/CUDA cov = ( torch.eye(block_size, device="cpu") * (1 + gamma) - torch.ones(block_size, block_size, device="cpu") * gamma ) dist = torch.distributions.MultivariateNormal( torch.zeros(block_size, device="cpu"), covariance_matrix=cov ) block_number = batch_size * channel * num_frames * (height // ph) * (width // pw) noise = dist.sample((block_number,)) noise = noise.view( batch_size, channel, num_frames, height // ph, width // pw, ph, pw ) noise = noise.permute(0, 1, 2, 3, 5, 4, 6).reshape( batch_size, channel, num_frames, height, width ) return noise class HeliosChunkedDenoisingStage(PipelineStage): """ Helios chunked denoising stage implementing Stage 1 loop. Iterates over video chunks, manages history buffers (short/mid/long), runs transformer per chunk with CFG guidance, scheduler step, and accumulates denoised latents. VAE decoding is left to DecodingStage. """ def __init__(self, transformer, scheduler): super().__init__() self.transformer = transformer self.scheduler = scheduler @property def parallelism_type(self): return StageParallelismType.REPLICATED def _denoise_one_chunk( self, latents, prompt_embeds, negative_prompt_embeds, timesteps, guidance_scale, indices_hidden_states, indices_latents_history_short, indices_latents_history_mid, indices_latents_history_long, latents_history_short, latents_history_mid, latents_history_long, target_dtype, device, is_cfg_zero_star=True, use_zero_init=True, zero_steps=1, batch=None, server_args=None, global_step_offset=0, ): """Denoise a single chunk with full timestep loop.""" batch_size = latents.shape[0] do_cfg = guidance_scale > 1.0 for i, t in enumerate(timesteps): with StageProfiler( f"denoising_step_{global_step_offset + i}", logger=logger, metrics=batch.metrics if batch is not None else None, perf_dump_path_provided=( batch.perf_dump_path is not None if batch is not None else False ), ): timestep = t.expand(batch_size) latent_model_input = latents.to(target_dtype) with set_forward_context( current_timestep=t, forward_batch=batch, attn_metadata=None, ): noise_pred = self.transformer( hidden_states=latent_model_input, timestep=timestep, encoder_hidden_states=prompt_embeds, indices_hidden_states=indices_hidden_states, indices_latents_history_short=indices_latents_history_short, indices_latents_history_mid=indices_latents_history_mid, indices_latents_history_long=indices_latents_history_long, latents_history_short=( latents_history_short.to(target_dtype) if latents_history_short is not None else None ), latents_history_mid=( latents_history_mid.to(target_dtype) if latents_history_mid is not None else None ), latents_history_long=( latents_history_long.to(target_dtype) if latents_history_long is not None else None ), ) if do_cfg: with set_forward_context( current_timestep=t, forward_batch=batch, attn_metadata=None, ): noise_uncond = self.transformer( hidden_states=latent_model_input, timestep=timestep, encoder_hidden_states=negative_prompt_embeds, indices_hidden_states=indices_hidden_states, indices_latents_history_short=indices_latents_history_short, indices_latents_history_mid=indices_latents_history_mid, indices_latents_history_long=indices_latents_history_long, latents_history_short=( latents_history_short.to(target_dtype) if latents_history_short is not None else None ), latents_history_mid=( latents_history_mid.to(target_dtype) if latents_history_mid is not None else None ), latents_history_long=( latents_history_long.to(target_dtype) if latents_history_long is not None else None ), ) if is_cfg_zero_star: noise_pred_text = noise_pred positive_flat = noise_pred_text.reshape(batch_size, -1) negative_flat = noise_uncond.reshape(batch_size, -1) alpha = optimized_scale(positive_flat, negative_flat) alpha = alpha.view( batch_size, *([1] * (len(noise_pred_text.shape) - 1)) ) alpha = alpha.to(noise_pred_text.dtype) if (i <= zero_steps) and use_zero_init: noise_pred = noise_pred_text * 0.0 else: noise_pred = noise_uncond * alpha + guidance_scale * ( noise_pred_text - noise_uncond * alpha ) else: noise_pred = noise_uncond + guidance_scale * ( noise_pred - noise_uncond ) latents = self.scheduler.step( noise_pred, t, latents, return_dict=False )[0] return latents def _denoise_one_chunk_stage2( self, latents, prompt_embeds, negative_prompt_embeds, guidance_scale, indices_hidden_states, indices_latents_history_short, indices_latents_history_mid, indices_latents_history_long, latents_history_short, latents_history_mid, latents_history_long, target_dtype, device, pyramid_num_stages, pyramid_num_inference_steps_list, is_distilled, is_amplify_first_chunk, gamma, is_cfg_zero_star=True, use_zero_init=True, zero_steps=1, batch=None, server_args=None, global_step_offset=0, ): """Denoise a single chunk using pyramid super-resolution (Stage 2).""" batch_size, num_channel, num_frames, height, width = latents.shape patch_size = self.transformer.patch_size # Downsample to lowest pyramid level latents = latents.permute(0, 2, 1, 3, 4).reshape( batch_size * num_frames, num_channel, height, width ) for _ in range(pyramid_num_stages - 1): height //= 2 width //= 2 latents = F.interpolate(latents, size=(height, width), mode="bilinear") * 2 latents = latents.reshape( batch_size, num_frames, num_channel, height, width ).permute(0, 2, 1, 3, 4) start_point_list = None if is_distilled: start_point_list = [latents] do_cfg = guidance_scale > 1.0 step_counter = global_step_offset for i_s in range(pyramid_num_stages): # Compute mu for current resolution image_seq_len = ( latents.shape[-1] * latents.shape[-2] * latents.shape[-3] // (patch_size[0] * patch_size[1] * patch_size[2]) ) mu = calculate_shift(image_seq_len) self.scheduler.set_timesteps( pyramid_num_inference_steps_list[i_s], i_s, device=device, mu=mu, is_amplify_first_chunk=is_amplify_first_chunk, ) timesteps = self.scheduler.timesteps if i_s > 0: # Upsample 2x nearest-neighbor height *= 2 width *= 2 latents = latents.permute(0, 2, 1, 3, 4).reshape( batch_size * num_frames, num_channel, height // 2, width // 2, ) latents = F.interpolate(latents, size=(height, width), mode="nearest") latents = latents.reshape( batch_size, num_frames, num_channel, height, width ).permute(0, 2, 1, 3, 4) # Renoise with correlated block noise ori_sigma = 1 - self.scheduler.ori_start_sigmas[i_s] alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma) beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) bs, ch, nf, h, w = latents.shape noise = sample_block_noise(bs, ch, nf, h, w, gamma, patch_size) noise = noise.to(device=device, dtype=target_dtype) latents = alpha * latents + beta * noise if is_distilled: start_point_list.append(latents) # Denoising loop for this pyramid stage for idx, t in enumerate(timesteps): with StageProfiler( f"denoising_step_{step_counter}", logger=logger, metrics=batch.metrics if batch is not None else None, perf_dump_path_provided=( batch.perf_dump_path is not None if batch is not None else False ), ): timestep = t.expand(batch_size) latent_model_input = latents.to(target_dtype) with set_forward_context( current_timestep=t, forward_batch=batch, attn_metadata=None, ): noise_pred = self.transformer( hidden_states=latent_model_input, timestep=timestep, encoder_hidden_states=prompt_embeds, indices_hidden_states=indices_hidden_states, indices_latents_history_short=indices_latents_history_short, indices_latents_history_mid=indices_latents_history_mid, indices_latents_history_long=indices_latents_history_long, latents_history_short=( latents_history_short.to(target_dtype) if latents_history_short is not None else None ), latents_history_mid=( latents_history_mid.to(target_dtype) if latents_history_mid is not None else None ), latents_history_long=( latents_history_long.to(target_dtype) if latents_history_long is not None else None ), ) if do_cfg: with set_forward_context( current_timestep=t, forward_batch=batch, attn_metadata=None, ): noise_uncond = self.transformer( hidden_states=latent_model_input, timestep=timestep, encoder_hidden_states=negative_prompt_embeds, indices_hidden_states=indices_hidden_states, indices_latents_history_short=indices_latents_history_short, indices_latents_history_mid=indices_latents_history_mid, indices_latents_history_long=indices_latents_history_long, latents_history_short=( latents_history_short.to(target_dtype) if latents_history_short is not None else None ), latents_history_mid=( latents_history_mid.to(target_dtype) if latents_history_mid is not None else None ), latents_history_long=( latents_history_long.to(target_dtype) if latents_history_long is not None else None ), ) if is_cfg_zero_star: noise_pred_text = noise_pred positive_flat = noise_pred_text.reshape(batch_size, -1) negative_flat = noise_uncond.reshape(batch_size, -1) alpha_cfg = optimized_scale(positive_flat, negative_flat) alpha_cfg = alpha_cfg.view( batch_size, *([1] * (len(noise_pred_text.shape) - 1)), ) alpha_cfg = alpha_cfg.to(noise_pred_text.dtype) if (i_s == 0 and idx <= zero_steps) and use_zero_init: noise_pred = noise_pred_text * 0.0 else: noise_pred = ( noise_uncond * alpha_cfg + guidance_scale * (noise_pred_text - noise_uncond * alpha_cfg) ) else: noise_pred = noise_uncond + guidance_scale * ( noise_pred - noise_uncond ) latents = self.scheduler.step( noise_pred, t, latents, return_dict=False, cur_sampling_step=idx, dmd_noisy_tensor=( start_point_list[i_s] if start_point_list is not None else None ), dmd_sigmas=self.scheduler.sigmas, dmd_timesteps=self.scheduler.timesteps, all_timesteps=timesteps, )[0] step_counter += 1 return latents, step_counter def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Run the Helios chunked denoising loop.""" pipeline_config = server_args.pipeline_config device = ( batch.latents.device if hasattr(batch, "latents") and batch.latents is not None else torch.device("cuda") ) target_dtype = PRECISION_TO_TYPE.get( server_args.pipeline_config.precision, torch.bfloat16 ) # Get config params num_latent_frames_per_chunk = pipeline_config.num_latent_frames_per_chunk history_sizes = sorted(list(pipeline_config.history_sizes), reverse=True) is_cfg_zero_star = pipeline_config.is_cfg_zero_star zero_steps = pipeline_config.zero_steps keep_first_frame = pipeline_config.keep_first_frame guidance_scale = batch.guidance_scale num_inference_steps = batch.num_inference_steps # Stage 2 params is_enable_stage2 = pipeline_config.is_enable_stage2 pyramid_num_stages = pipeline_config.pyramid_num_stages pyramid_num_inference_steps_list = ( pipeline_config.pyramid_num_inference_steps_list ) is_distilled = pipeline_config.is_distilled is_amplify_first_chunk = pipeline_config.is_amplify_first_chunk gamma = pipeline_config.gamma # Move transformer to GPU if CPU-offloaded if server_args.dit_cpu_offload and not server_args.use_fsdp_inference: if next(self.transformer.parameters()).device.type == "cpu": self.transformer.to(get_local_torch_device()) # Get encoder outputs (prompt_embeds is a list of tensors, one per encoder) prompt_embeds = batch.prompt_embeds if isinstance(prompt_embeds, list): prompt_embeds = prompt_embeds[0] prompt_embeds = prompt_embeds.to(target_dtype) negative_prompt_embeds = batch.negative_prompt_embeds if isinstance(negative_prompt_embeds, list): negative_prompt_embeds = ( negative_prompt_embeds[0] if negative_prompt_embeds else None ) if negative_prompt_embeds is not None: negative_prompt_embeds = negative_prompt_embeds.to(target_dtype) # Scale factors inherited from the Wan VAE used by Helios # (AutoencoderKLWan: temporal_compression_ratio=4, spatial_compression_ratio=8) vae_scale_factor_temporal = 4 vae_scale_factor_spatial = 8 # Compute chunking height = batch.height width = batch.width num_frames = batch.num_frames num_channels_latents = self.transformer.in_channels window_num_frames = ( num_latent_frames_per_chunk - 1 ) * vae_scale_factor_temporal + 1 num_latent_chunk = max( 1, (num_frames + window_num_frames - 1) // window_num_frames ) num_history_latent_frames = sum(history_sizes) batch_size = 1 # Helios processes one video at a time # Prepare history latents if not keep_first_frame: history_sizes[-1] = history_sizes[-1] + 1 history_latents = torch.zeros( batch_size, num_channels_latents, num_history_latent_frames, height // vae_scale_factor_spatial, width // vae_scale_factor_spatial, device=device, dtype=torch.float32, ) # Build frame indices if keep_first_frame: indices = torch.arange( 0, sum([1, *history_sizes, num_latent_frames_per_chunk]) ) ( indices_prefix, indices_latents_history_long, indices_latents_history_mid, indices_latents_history_1x, indices_hidden_states, ) = indices.split([1, *history_sizes, num_latent_frames_per_chunk], dim=0) indices_latents_history_short = torch.cat( [indices_prefix, indices_latents_history_1x], dim=0 ) else: indices = torch.arange( 0, sum([*history_sizes, num_latent_frames_per_chunk]) ) ( indices_latents_history_long, indices_latents_history_mid, indices_latents_history_short, indices_hidden_states, ) = indices.split([*history_sizes, num_latent_frames_per_chunk], dim=0) indices_hidden_states = indices_hidden_states.unsqueeze(0) indices_latents_history_short = indices_latents_history_short.unsqueeze(0) indices_latents_history_mid = indices_latents_history_mid.unsqueeze(0) indices_latents_history_long = indices_latents_history_long.unsqueeze(0) # Set up scheduler patch_size = self.transformer.patch_size image_seq_len = ( num_latent_frames_per_chunk * (height // vae_scale_factor_spatial) * (width // vae_scale_factor_spatial) // (patch_size[0] * patch_size[1] * patch_size[2]) ) # Sigma schedule from near-1.0 (pure noise) to 0.0 (clean); 0.999 avoids singularity sigmas = np.linspace(0.999, 0.0, num_inference_steps + 1)[:-1] mu = calculate_shift(image_seq_len) # Chunk loop image_latents = None total_generated_latent_frames = 0 chunk_latents_list = [] # Store per-chunk latents for chunk-by-chunk decode global_step_offset = 0 # Track step index across chunks for perf logging self.log_info( f"Starting chunked denoising: {num_latent_chunk} chunks, " f"{num_inference_steps} steps each" ) for k in range(num_latent_chunk): is_first_chunk = k == 0 # Extract history if keep_first_frame: ( latents_history_long, latents_history_mid, latents_history_1x, ) = history_latents[:, :, -num_history_latent_frames:].split( history_sizes, dim=2 ) if image_latents is None and is_first_chunk: latents_prefix = torch.zeros( ( batch_size, num_channels_latents, 1, latents_history_1x.shape[-2], latents_history_1x.shape[-1], ), device=device, dtype=latents_history_1x.dtype, ) else: latents_prefix = image_latents latents_history_short = torch.cat( [latents_prefix, latents_history_1x], dim=2 ) else: ( latents_history_long, latents_history_mid, latents_history_short, ) = history_latents[:, :, -num_history_latent_frames:].split( history_sizes, dim=2 ) # Generate noise latents for this chunk # Use batch.generator to ensure identical noise across SP ranks latent_shape = ( batch_size, num_channels_latents, (window_num_frames - 1) // vae_scale_factor_temporal + 1, height // vae_scale_factor_spatial, width // vae_scale_factor_spatial, ) generator = batch.generator if isinstance(generator, list): generator = generator[0] if len(generator) > 0 else None gen_device = generator.device if generator is not None else device latents = torch.randn( latent_shape, generator=generator, device=gen_device, dtype=torch.float32, ) if latents.device != device: latents = latents.to(device) if is_enable_stage2: # Stage 2: Pyramid SR denoising (handles scheduler internally) latents, global_step_offset = self._denoise_one_chunk_stage2( latents=latents, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, guidance_scale=guidance_scale, indices_hidden_states=indices_hidden_states, indices_latents_history_short=indices_latents_history_short, indices_latents_history_mid=indices_latents_history_mid, indices_latents_history_long=indices_latents_history_long, latents_history_short=latents_history_short, latents_history_mid=latents_history_mid, latents_history_long=latents_history_long, target_dtype=target_dtype, device=device, pyramid_num_stages=pyramid_num_stages, pyramid_num_inference_steps_list=pyramid_num_inference_steps_list, is_distilled=is_distilled, is_amplify_first_chunk=(is_amplify_first_chunk and is_first_chunk), gamma=gamma, is_cfg_zero_star=is_cfg_zero_star, use_zero_init=True, zero_steps=zero_steps, batch=batch, server_args=server_args, global_step_offset=global_step_offset, ) else: # Stage 1: Standard flat denoising self.scheduler.set_timesteps( num_inference_steps, device=device, sigmas=sigmas, mu=mu ) timesteps = self.scheduler.timesteps latents = self._denoise_one_chunk( latents=latents, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, timesteps=timesteps, guidance_scale=guidance_scale, indices_hidden_states=indices_hidden_states, indices_latents_history_short=indices_latents_history_short, indices_latents_history_mid=indices_latents_history_mid, indices_latents_history_long=indices_latents_history_long, latents_history_short=latents_history_short, latents_history_mid=latents_history_mid, latents_history_long=latents_history_long, target_dtype=target_dtype, device=device, is_cfg_zero_star=is_cfg_zero_star, use_zero_init=True, zero_steps=zero_steps, batch=batch, server_args=server_args, global_step_offset=global_step_offset, ) global_step_offset += num_inference_steps # Extract first frame as image_latents for subsequent chunks if keep_first_frame and is_first_chunk and image_latents is None: image_latents = latents[:, :, 0:1, :, :] # Update history total_generated_latent_frames += latents.shape[2] history_latents = torch.cat([history_latents, latents], dim=2) chunk_latents_list.append(latents) # Move transformer back to CPU after denoising if server_args.dit_cpu_offload and not server_args.use_fsdp_inference: if next(self.transformer.parameters()).device.type != "cpu": self.transformer.to("cpu") torch.cuda.empty_cache() # Store per-chunk latents for chunk-by-chunk VAE decode (matches diffusers behavior). # The standard DecodingStage will check for this attribute and decode each chunk # separately to avoid temporal artifacts at chunk boundaries. batch.latent_chunks = chunk_latents_list batch.latents = history_latents[:, :, -total_generated_latent_frames:] return batch ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ MOVA-specific pipeline stages. Sequence Parallelism (SP) Support: - Video latents are sharded along the sequence dimension (T*H*W) after patchify - Audio latents are sharded along the sequence dimension (L) after patchify - USPAttention handles all-to-all communication internally - Latents are gathered before unpatchify to restore full sequence """ from __future__ import annotations import functools import inspect import os from collections.abc import Iterable import torch import torch.nn as nn from diffusers.utils.torch_utils import randn_tensor from tqdm.auto import tqdm from sglang.multimodal_gen.runtime.distributed import ( get_local_torch_device, get_world_group, ) from sglang.multimodal_gen.runtime.distributed.communication_op import ( cfg_model_parallel_all_reduce, sequence_model_parallel_all_gather, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_cfg_group, get_classifier_free_guidance_rank, get_sp_parallel_rank, get_sp_world_size, ) from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context # Both audio and video DiT use the same sinusoidal_embedding_1d function # Import from mova_video_dit where it's defined (mova_audio_dit re-exports it) from sglang.multimodal_gen.runtime.models.dits.mova_video_dit import ( sinusoidal_embedding_1d, ) # Create aliases for backward compatibility video_sinusoidal_embedding_1d = sinusoidal_embedding_1d audio_sinusoidal_embedding_1d = sinusoidal_embedding_1d from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( PipelineStage, StageParallelismType, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import ( _ensure_tensor_decode_output, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler from sglang.multimodal_gen.utils import PRECISION_TO_TYPE from sglang.srt.utils.common import get_compiler_backend logger = init_logger(__name__) class MOVALatentPreparationStage(PipelineStage): """Prepare video/audio noise latents for MOVA.""" def __init__(self, audio_vae, require_vae_embedding: bool = True) -> None: super().__init__() self.audio_vae = audio_vae self.require_vae_embedding = require_vae_embedding def forward(self, batch: Req, server_args: ServerArgs) -> Req: batch_size = batch.batch_size num_frames = batch.num_frames if num_frames is None: raise ValueError("num_frames is required for MOVA") audio_num_samples = int(self.audio_vae.sample_rate * num_frames / batch.fps) video_shape = server_args.pipeline_config.prepare_latent_shape( batch, batch_size, num_frames ) audio_shape = server_args.pipeline_config.prepare_audio_latent_shape( batch_size, audio_num_samples, self.audio_vae ) device = get_local_torch_device() generator = batch.generator if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) dit_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] batch.latents = randn_tensor( video_shape, generator=generator, device=device, dtype=dit_dtype ) batch.audio_latents = randn_tensor( audio_shape, generator=generator, device=device, dtype=dit_dtype ) if batch.image_latent is not None: batch.y = batch.image_latent.to(device=device, dtype=dit_dtype) elif self.require_vae_embedding: raise ValueError("MOVA requires reference image latents for denoising") return batch class MOVATimestepPreparationStage(PipelineStage): """Prepare paired timesteps for MOVA.""" def __init__(self, scheduler) -> None: super().__init__() self.scheduler = scheduler def forward(self, batch: Req, server_args: ServerArgs) -> Req: self.scheduler.set_timesteps( batch.num_inference_steps, denoising_strength=1.0, shift=getattr(batch, "sigma_shift", self.scheduler.shift), ) self.scheduler.set_pair_postprocess_by_name( "dual_sigma_shift", visual_shift=getattr(batch, "visual_shift", 5.0), audio_shift=getattr(batch, "audio_shift", 5.0), ) paired = self.scheduler.get_pairs() batch.paired_timesteps = paired batch.timesteps = paired return batch class MOVADenoisingStage(PipelineStage): """Run MOVA dual-tower denoising loop.""" def __init__(self, video_dit, video_dit_2, audio_dit, dual_tower_bridge, scheduler): super().__init__() self.video_dit = video_dit self.video_dit_2 = video_dit_2 self.audio_dit = audio_dit self.dual_tower_bridge = dual_tower_bridge self.scheduler = scheduler self._cache_dit_enabled = False self._cached_num_steps = None self._torch_compiled = False @property def parallelism_type(self) -> StageParallelismType: if get_global_server_args().enable_cfg_parallel: return StageParallelismType.CFG_PARALLEL return StageParallelismType.REPLICATED def _predict( self, visual_dit, visual_latents, audio_latents, y, context, timestep, audio_timestep, video_fps, timestep_index: int, attn_metadata, forward_batch: Req | None = None, ): # Set forward context for distributed attention (USPAttention) with set_forward_context( current_timestep=timestep_index, attn_metadata=attn_metadata, forward_batch=forward_batch, ): return self.inference_single_step( visual_dit=visual_dit, visual_latents=visual_latents, audio_latents=audio_latents, y=y, context=context, timestep=timestep, audio_timestep=audio_timestep, video_fps=video_fps, ) def _cfg_combine(self, pos, neg, guidance_scale, cfg_rank, enable_cfg_parallel): if not enable_cfg_parallel: return neg + guidance_scale * (pos - neg) if cfg_rank == 0: partial = guidance_scale * pos else: partial = (1 - guidance_scale) * neg return cfg_model_parallel_all_reduce(partial) def _maybe_enable_torch_compile(self, module: nn.Module, server_args: ServerArgs): """ Compile a module with torch.compile, and enable inductor overlap tweak if available. No-op if torch compile is disabled or the object is not a nn.Module. """ if not server_args.enable_torch_compile or not isinstance(module, nn.Module): return compile_kwargs: dict[str, object] = {"fullgraph": False, "dynamic": None} if current_platform.is_npu(): backend = get_compiler_backend() compile_kwargs["backend"] = backend compile_kwargs["dynamic"] = False logger.info( "Compiling %s with torchair backend on NPU", module.__class__.__name__, ) else: try: import torch._inductor.config as _inductor_cfg _inductor_cfg.reorder_for_compute_comm_overlap = True except ImportError: pass mode = os.environ.get( "SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs" ) compile_kwargs["mode"] = mode logger.info("Compiling %s with mode: %s", module.__class__.__name__, mode) # TODO(triple-mu): support customized fullgraph and dynamic in the future module.compile(**compile_kwargs) def _maybe_compile_dits(self, server_args: ServerArgs): if self._torch_compiled or not server_args.enable_torch_compile: return for module in filter(None, [self.video_dit, self.video_dit_2, self.audio_dit]): self._maybe_enable_torch_compile(module, server_args) self._torch_compiled = True def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify denoising stage inputs.""" result = VerificationResult() result.add_check("y", batch.y, V.is_tensor) result.add_check("paired_timesteps", batch.paired_timesteps, V.is_tensor) result.add_check("latents", batch.latents, V.is_tensor) result.add_check("audio_latents", batch.audio_latents, V.is_tensor) result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) result.add_check( "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x), ) result.add_check( "num_inference_steps", batch.num_inference_steps, V.positive_int ) result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float) result.add_check( "guidance_rescale", batch.guidance_rescale, V.non_negative_float ) result.add_check( "do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value, ) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify denoising stage outputs.""" result = VerificationResult() result.add_check("latents", batch.latents, V.is_tensor) result.add_check("audio_latents", batch.audio_latents, V.is_tensor) return result def progress_bar( self, iterable: Iterable | None = None, total: int | None = None ) -> tqdm: """ Create a progress bar for the denoising process. """ local_rank = get_world_group().local_rank disable = local_rank != 0 return tqdm(iterable=iterable, total=total, disable=disable) def step_profile(self): profiler = SGLDiffusionProfiler.get_instance() if profiler: profiler.step_denoising_step() def rescale_noise_cfg( self, noise_cfg, noise_pred_text, guidance_rescale=0.0 ) -> torch.Tensor: """ Rescale noise prediction according to guidance_rescale. Based on findings of "Common Diffusion Noise Schedules and Sample Steps are Flawed" (https://arxiv.org/pdf/2305.08891.pdf), Section 3.4. """ std_text = noise_pred_text.std( dim=list(range(1, noise_pred_text.ndim)), keepdim=True ) std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) noise_pred_rescaled = noise_cfg * (std_text / std_cfg) noise_cfg = ( guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg ) return noise_cfg def prepare_extra_func_kwargs(self, func, kwargs) -> dict[str, object]: if not kwargs: return {} if isinstance(func, functools.partial) and func.args: func = getattr(func.args[0], "_original_forward", func) target_func = inspect.unwrap(func) params = inspect.signature(target_func).parameters return {k: v for k, v in kwargs.items() if k in params} def _build_attn_metadata( self, i: int, batch: Req, server_args: ServerArgs ) -> object | None: return None def _manage_device_placement( self, model_to_use: nn.Module | None, model_to_offload: nn.Module | None, server_args: ServerArgs, ): if not server_args.dit_cpu_offload: return if ( model_to_offload is not None and next(model_to_offload.parameters()).device.type == "cuda" ): model_to_offload.to("cpu") if ( model_to_use is not None and next(model_to_use.parameters()).device.type == "cpu" ): model_to_use.to(get_local_torch_device()) def _select_visual_dit( self, timestep: float, boundary_ratio: float | None, server_args: ServerArgs ): if boundary_ratio is None or self.video_dit_2 is None: self._manage_device_placement(self.video_dit, None, server_args) return self.video_dit boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps if timestep >= boundary_timestep: current_model = self.video_dit model_to_offload = self.video_dit_2 else: current_model = self.video_dit_2 model_to_offload = self.video_dit self._manage_device_placement(current_model, model_to_offload, server_args) return current_model def _ensure_shared_models_on_device(self, server_args: ServerArgs): """Ensure shared denoising modules are on the active device when cpu offload is enabled.""" self._manage_device_placement(self.audio_dit, None, server_args) self._manage_device_placement(self.dual_tower_bridge, None, server_args) def _apply_guidance_rescale( self, noise_pred, noise_pred_text, guidance_rescale, cfg_rank, enable_cfg_parallel, ): if guidance_rescale <= 0.0: return noise_pred if enable_cfg_parallel: std_cfg = noise_pred.std(dim=list(range(1, noise_pred.ndim)), keepdim=True) if cfg_rank == 0: assert noise_pred_text is not None std_text = noise_pred_text.std( dim=list(range(1, noise_pred_text.ndim)), keepdim=True ) else: std_text = torch.empty_like(std_cfg) std_text = get_cfg_group().broadcast(std_text, src=0) noise_pred_rescaled = noise_pred * (std_text / std_cfg) return guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * ( noise_pred ) return self.rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale) @torch.no_grad() def forward(self, batch: Req, server_args: ServerArgs) -> Req: self._maybe_compile_dits(server_args) self._ensure_shared_models_on_device(server_args) paired_timesteps = batch.paired_timesteps if paired_timesteps is None: raise ValueError("paired_timesteps must be set for MOVA") y = batch.y if batch.y is not None else batch.image_latent if getattr(self.video_dit, "require_vae_embedding", False) and y is None: raise ValueError("MOVA requires reference image latents for denoising") boundary_ratio = server_args.pipeline_config.boundary_ratio total_steps = paired_timesteps.shape[0] cfg_rank = get_classifier_free_guidance_rank() enable_cfg_parallel = server_args.enable_cfg_parallel is_warmup = batch.is_warmup extra_step_kwargs = self.prepare_extra_func_kwargs( self.scheduler.step_from_to, getattr(batch, "extra_step_kwargs", None) or {}, ) metrics = getattr(batch, "metrics", None) perf_dump_path_provided = getattr(batch, "perf_dump_path", None) is not None with self.progress_bar(total=total_steps) as progress_bar: for idx_step in range(total_steps): with StageProfiler( f"denoising_step_{idx_step}", logger=logger, metrics=metrics, perf_dump_path_provided=perf_dump_path_provided, ): pair_t = paired_timesteps[idx_step] if getattr(pair_t, "shape", None) == (2,): timestep, audio_timestep = pair_t else: timestep = pair_t audio_timestep = pair_t cur_visual_dit = self._select_visual_dit( timestep.item(), boundary_ratio, server_args ) timestep = timestep.unsqueeze(0).to(device=get_local_torch_device()) audio_timestep = audio_timestep.unsqueeze(0).to( device=get_local_torch_device() ) attn_metadata = self._build_attn_metadata( idx_step, batch, server_args ) if not batch.do_classifier_free_guidance: visual_noise_pred, audio_noise_pred = self._predict( cur_visual_dit, batch.latents, batch.audio_latents, y, batch.prompt_embeds[0], timestep, audio_timestep, batch.fps, idx_step, attn_metadata, batch, ) else: if enable_cfg_parallel: if cfg_rank == 0: pos = self._predict( cur_visual_dit, batch.latents, batch.audio_latents, y, batch.prompt_embeds[0], timestep, audio_timestep, batch.fps, idx_step, attn_metadata, batch, ) neg = (None, None) else: pos = (None, None) neg = self._predict( cur_visual_dit, batch.latents, batch.audio_latents, y, batch.negative_prompt_embeds[0], timestep, audio_timestep, batch.fps, idx_step, attn_metadata, batch, ) else: pos = self._predict( cur_visual_dit, batch.latents, batch.audio_latents, y, batch.prompt_embeds[0], timestep, audio_timestep, batch.fps, idx_step, attn_metadata, batch, ) neg = self._predict( cur_visual_dit, batch.latents, batch.audio_latents, y, batch.negative_prompt_embeds[0], timestep, audio_timestep, batch.fps, idx_step, attn_metadata, batch, ) visual_noise_pred = self._cfg_combine( pos[0] if pos[0] is not None else neg[0], neg[0] if neg[0] is not None else pos[0], batch.guidance_scale, cfg_rank, enable_cfg_parallel, ) audio_noise_pred = self._cfg_combine( pos[1] if pos[1] is not None else neg[1], neg[1] if neg[1] is not None else pos[1], batch.guidance_scale, cfg_rank, enable_cfg_parallel, ) if batch.guidance_rescale > 0.0: visual_noise_pred = self._apply_guidance_rescale( visual_noise_pred, pos[0] if pos[0] is not None else None, batch.guidance_rescale, cfg_rank, enable_cfg_parallel, ) audio_noise_pred = self._apply_guidance_rescale( audio_noise_pred, pos[1] if pos[1] is not None else None, batch.guidance_rescale, cfg_rank, enable_cfg_parallel, ) if idx_step + 1 < total_steps: next_pair_t = paired_timesteps[idx_step + 1] if getattr(next_pair_t, "shape", None) == (2,): next_timestep, next_audio_timestep = next_pair_t else: next_timestep = next_pair_t next_audio_timestep = next_pair_t else: next_timestep = None next_audio_timestep = None batch.latents = self.scheduler.step_from_to( visual_noise_pred, timestep, next_timestep, batch.latents, **extra_step_kwargs, ) batch.audio_latents = self.scheduler.step_from_to( audio_noise_pred, audio_timestep, next_audio_timestep, batch.audio_latents, **extra_step_kwargs, ) if progress_bar is not None: progress_bar.update() if not is_warmup and hasattr(self, "step_profile"): self.step_profile() for dit in filter(None, [self.video_dit, self.video_dit_2, self.audio_dit]): if isinstance(dit, OffloadableDiTMixin): dit.prepare_for_next_req() return batch def _shard_sequence_for_sp( self, x: torch.Tensor, dim: int = 1 ) -> tuple[torch.Tensor, int]: """ Shard tensor along sequence dimension for Sequence Parallelism. Args: x: Input tensor dim: Dimension to shard along Returns: (sharded_tensor, pad_len) """ sp_size = get_sp_world_size() if sp_size <= 1: return x, 0 sp_rank = get_sp_parallel_rank() seq_len = x.shape[dim] # Pad if needed pad_len = (sp_size - (seq_len % sp_size)) % sp_size if pad_len > 0: pad_shape = list(x.shape) pad_shape[dim] = pad_len pad = torch.zeros(pad_shape, dtype=x.dtype, device=x.device) x = torch.cat([x, pad], dim=dim) # Shard chunk_size = x.shape[dim] // sp_size start = sp_rank * chunk_size end = start + chunk_size idx = [slice(None)] * x.dim() idx[dim] = slice(start, end) return x[tuple(idx)], pad_len def _gather_sequence_from_sp( self, x: torch.Tensor, pad_len: int, dim: int = 1 ) -> torch.Tensor: """ Gather tensor along sequence dimension after Sequence Parallelism. Args: x: Sharded tensor pad_len: Padding length that was added during sharding dim: Dimension to gather along Returns: Gathered tensor with padding removed """ sp_size = get_sp_world_size() if sp_size <= 1: return x gathered = sequence_model_parallel_all_gather(x, dim=dim) if pad_len > 0: idx = [slice(None)] * gathered.dim() idx[dim] = slice(0, gathered.shape[dim] - pad_len) gathered = gathered[tuple(idx)] return gathered def inference_single_step( self, visual_dit, visual_latents: torch.Tensor, audio_latents: torch.Tensor, y, context: torch.Tensor, timestep: torch.Tensor, audio_timestep: torch.Tensor, video_fps: float, ): """ Single inference step for MOVA dual-tower denoising. Supports Sequence Parallelism (SP): - After patchify, sequences are sharded across SP ranks - USPAttention handles distributed attention communication - Before unpatchify, sequences are gathered back """ model_dtype = visual_dit.time_embedding.fc_in.weight.dtype device = visual_latents.device visual_context = context.to(device=device, dtype=model_dtype) audio_context = context.to(device=device, dtype=model_dtype) with torch.autocast( device_type=current_platform.device_type, dtype=torch.float32 ): visual_t = visual_dit.time_embedding( video_sinusoidal_embedding_1d(visual_dit.freq_dim, timestep) ) visual_t_mod, _ = visual_dit.time_projection(visual_t) visual_t_mod = visual_t_mod.unflatten(1, (6, visual_dit.dim)) audio_t = self.audio_dit.time_embedding( audio_sinusoidal_embedding_1d(self.audio_dit.freq_dim, audio_timestep) ) audio_t_mod, _ = self.audio_dit.time_projection(audio_t) audio_t_mod = audio_t_mod.unflatten(1, (6, self.audio_dit.dim)) visual_t = visual_t.to(model_dtype) visual_t_mod = visual_t_mod.to(model_dtype) audio_t = audio_t.to(model_dtype) audio_t_mod = audio_t_mod.to(model_dtype) visual_context_emb = visual_dit.text_embedding(visual_context) audio_context_emb = self.audio_dit.text_embedding(audio_context) visual_x = visual_latents.to(model_dtype) audio_x = audio_latents.to(model_dtype) if getattr(visual_dit, "require_vae_embedding", False): visual_x = torch.cat([visual_x, y], dim=1) # Patchify visual latents visual_x, (t, h, w) = visual_dit.patchify(visual_x) grid_size = (t, h, w) full_visual_seq_len = t * h * w # Build visual freqs for full sequence visual_dit._init_freqs() visual_freqs = tuple(freq.to(visual_x.device) for freq in visual_dit.freqs) visual_freqs = ( torch.cat( [ visual_freqs[0][:t].view(t, 1, 1, -1).expand(t, h, w, -1), visual_freqs[1][:h].view(1, h, 1, -1).expand(t, h, w, -1), visual_freqs[2][:w].view(1, 1, w, -1).expand(t, h, w, -1), ], dim=-1, ) .reshape(full_visual_seq_len, 1, -1) .to(visual_x.device) ) # Patchify audio latents audio_x, (f,) = self.audio_dit.patchify(audio_x, None) full_audio_seq_len = f # Build audio freqs for full sequence self.audio_dit._init_freqs() audio_freqs = ( torch.cat( [ self.audio_dit.freqs[0][:f].view(f, -1).expand(f, -1), self.audio_dit.freqs[1][:f].view(f, -1).expand(f, -1), self.audio_dit.freqs[2][:f].view(f, -1).expand(f, -1), ], dim=-1, ) .reshape(full_audio_seq_len, 1, -1) .to(audio_x.device) ) # Shard sequences for SP visual_x, visual_pad_len = self._shard_sequence_for_sp(visual_x, dim=1) audio_x, audio_pad_len = self._shard_sequence_for_sp(audio_x, dim=1) # Shard freqs to match local sequence length visual_freqs, _ = self._shard_sequence_for_sp(visual_freqs, dim=0) audio_freqs, _ = self._shard_sequence_for_sp(audio_freqs, dim=0) # Forward through dual-tower DiT visual_x, audio_x = self.forward_dual_tower_dit( visual_dit=visual_dit, visual_x=visual_x, audio_x=audio_x, visual_context=visual_context_emb, audio_context=audio_context_emb, visual_t_mod=visual_t_mod, audio_t_mod=audio_t_mod, visual_freqs=visual_freqs, audio_freqs=audio_freqs, grid_size=grid_size, video_fps=video_fps, full_visual_seq_len=full_visual_seq_len, full_audio_seq_len=full_audio_seq_len, ) # Gather sequences back from SP before head/unpatchify visual_x = self._gather_sequence_from_sp(visual_x, visual_pad_len, dim=1) audio_x = self._gather_sequence_from_sp(audio_x, audio_pad_len, dim=1) visual_output = visual_dit.head(visual_x, visual_t) visual_output = visual_dit.unpatchify(visual_output, grid_size) audio_output = self.audio_dit.head(audio_x, audio_t) audio_output = self.audio_dit.unpatchify(audio_output, (f,)) return visual_output.float(), audio_output.float() def forward_dual_tower_dit( self, visual_dit, visual_x: torch.Tensor, audio_x: torch.Tensor, visual_context: torch.Tensor, audio_context: torch.Tensor, visual_t_mod: torch.Tensor, audio_t_mod: torch.Tensor, visual_freqs: torch.Tensor, audio_freqs: torch.Tensor, grid_size: tuple[int, int, int], video_fps: float, full_visual_seq_len: int, full_audio_seq_len: int, condition_scale: float | None = 1.0, a2v_condition_scale: float | None = None, v2a_condition_scale: float | None = None, ): """ Forward pass through dual-tower DiT with cross-modal interaction. Sequence Parallelism (SP) Support: - visual_x and audio_x are already sharded along sequence dimension - visual_freqs and audio_freqs match the local sequence length - USPAttention in self-attention handles distributed communication - LocalAttention in cross-attention operates on local sequence vs replicated context - Cross-modal attention (dual_tower_bridge) uses LocalAttention (no SP communication) Args: full_visual_seq_len: Full visual sequence length before SP sharding full_audio_seq_len: Full audio sequence length before SP sharding """ min_layers = min(len(visual_dit.blocks), len(self.audio_dit.blocks)) visual_layers = len(visual_dit.blocks) sp_size = get_sp_world_size() # Build RoPE frequencies for cross-attention if needed (only used when SP == 1) # When SP > 1, we rebuild freqs inside the loop after gathering full sequences visual_rope_cos_sin, audio_rope_cos_sin = ( self.dual_tower_bridge.build_aligned_freqs( video_fps=video_fps, grid_size=grid_size, audio_steps=full_audio_seq_len, device=visual_x.device, dtype=visual_x.dtype, ) ) if visual_rope_cos_sin is not None: visual_rope_cos_sin = [ self._shard_sequence_for_sp(rope_cos_sin, dim=1)[0] for rope_cos_sin in visual_rope_cos_sin ] if audio_rope_cos_sin is not None: audio_rope_cos_sin = [ self._shard_sequence_for_sp(rope_cos_sin, dim=1)[0] for rope_cos_sin in audio_rope_cos_sin ] for layer_idx in range(min_layers): visual_block = visual_dit.blocks[layer_idx] audio_block = self.audio_dit.blocks[layer_idx] # Cross-modal interaction via dual tower bridge # Bridge operations (PerFrameAttentionPooling, RoPE) expect full sequences # When SP is enabled, we need to gather before bridge and shard after if self.dual_tower_bridge.should_interact(layer_idx, "a2v"): visual_x, audio_x = self.dual_tower_bridge( layer_idx, visual_x, audio_x, x_freqs=visual_rope_cos_sin, y_freqs=audio_rope_cos_sin, a2v_condition_scale=a2v_condition_scale, v2a_condition_scale=v2a_condition_scale, condition_scale=condition_scale, video_grid_size=grid_size, ) # Self-attention and FFN in DiT blocks visual_x = visual_block( visual_x, visual_context, visual_t_mod, visual_freqs ) audio_x = audio_block(audio_x, audio_context, audio_t_mod, audio_freqs) # Process remaining visual layers (if visual has more layers than audio) for layer_idx in range(min_layers, visual_layers): visual_block = visual_dit.blocks[layer_idx] visual_x = visual_block( visual_x, visual_context, visual_t_mod, visual_freqs ) return visual_x, audio_x class MOVADecodingStage(PipelineStage): """Decode video and audio outputs for MOVA.""" def __init__(self, video_vae, audio_vae) -> None: super().__init__() self.video_vae = video_vae self.audio_vae = audio_vae @property def parallelism_type(self) -> StageParallelismType: if get_global_server_args().enable_cfg_parallel: return StageParallelismType.MAIN_RANK_ONLY return StageParallelismType.REPLICATED @torch.no_grad() def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: self.video_vae = self.video_vae.to(get_local_torch_device()) self.audio_vae = self.audio_vae.to(get_local_torch_device()) video_latents = server_args.pipeline_config.denormalize_video_latents( batch.latents, self.video_vae ) vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] vae_autocast_enabled = ( vae_dtype != torch.float32 ) and not server_args.disable_autocast with torch.autocast( device_type=current_platform.device_type, dtype=vae_dtype, enabled=vae_autocast_enabled, ): if server_args.pipeline_config.vae_tiling: self.video_vae.enable_tiling() if not vae_autocast_enabled: video_latents = video_latents.to(vae_dtype) decode_output = self.video_vae.decode(video_latents) video = _ensure_tensor_decode_output(decode_output) video = (video / 2 + 0.5).clamp(0, 1) with torch.autocast( device_type=current_platform.device_type, dtype=torch.float32 ): audio = self.audio_vae.decode(batch.audio_latents) output_batch = OutputBatch( output=video, audio=audio, audio_sample_rate=getattr(self.audio_vae, "sample_rate", None), metrics=batch.metrics, ) return output_batch ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image_layered.py ================================================ import inspect import math from typing import List, Optional, Union import numpy as np import torch from diffusers.image_processor import VaeImageProcessor from diffusers.utils.torch_utils import randn_tensor from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.models.vision_utils import load_image from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit_plus.calculate_dimensions def calculate_dimensions(target_area, ratio): width = math.sqrt(target_area * ratio) height = width / ratio width = round(width / 32) * 32 height = round(height / 32) * 32 return width, height # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample", ): if sample_mode == "sample": return encoder_output.sample(generator) elif sample_mode == "argmax": return encoder_output.mode() else: return encoder_output # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, num_inference_steps: Optional[int] = None, device: Optional[Union[str, torch.device]] = None, timesteps: Optional[List[int]] = None, sigmas: Optional[List[float]] = None, **kwargs, ): r""" Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. Args: scheduler (`SchedulerMixin`): The scheduler to get timesteps from. num_inference_steps (`int`): The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. timesteps (`List[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. sigmas (`List[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: raise ValueError( "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values" ) if timesteps is not None: accepts_timesteps = "timesteps" in set( inspect.signature(scheduler.set_timesteps).parameters.keys() ) if not accepts_timesteps: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" timestep schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) elif sigmas is not None: accept_sigmas = "sigmas" in set( inspect.signature(scheduler.set_timesteps).parameters.keys() ) if not accept_sigmas: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" sigmas schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) else: scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) timesteps = scheduler.timesteps return timesteps, num_inference_steps class QwenImageLayeredBeforeDenoisingStage(PipelineStage): def __init__( self, vae, tokenizer, processor, transformer, scheduler, model_path ) -> None: super().__init__() self.vae = vae.to(torch.bfloat16) from transformers import Qwen2_5_VLForConditionalGeneration self.text_encoder = ( Qwen2_5_VLForConditionalGeneration.from_pretrained( model_path, subfolder="text_encoder" ) .to(get_local_torch_device()) .to(torch.bfloat16) ) self.tokenizer = tokenizer self.processor = processor self.transformer = transformer self.scheduler = scheduler self.vae_scale_factor = ( 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8 ) self.image_processor = VaeImageProcessor( vae_scale_factor=self.vae_scale_factor * 2 ) self.vl_processor = processor self.tokenizer_max_length = 1024 self.latent_channels = self.vae.z_dim if getattr(self, "vae", None) else 16 self.prompt_template_encode = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" self.prompt_template_encode_start_idx = 34 self.image_caption_prompt_cn = """<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n# 图像标注器\n你是一个专业的图像标注器。请基于输入图像，撰写图注:\n1. 使用自然、描述性的语言撰写图注，不要使用结构化形式或富文本形式。\n2. 通过加入以下内容，丰富图注细节：\n - 对象的属性：如数量、颜色、形状、大小、位置、材质、状态、动作等\n - 对象间的视觉关系：如空间关系、功能关系、动作关系、从属关系、比较关系、因果关系等\n - 环境细节：例如天气、光照、颜色、纹理、气氛等\n - 文字内容：识别图像中清晰可见的文字，不做翻译和解释，用引号在图注中强调\n3. 保持真实性与准确性：\n - 不要使用笼统的描述\n - 描述图像中所有可见的信息，但不要加入没有在图像中出现的内容\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>assistant\n""" self.image_caption_prompt_en = """<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n# Image Annotator\nYou are a professional image annotator. Please write an image caption based on the input image:\n1. Write the caption using natural, descriptive language without structured formats or rich text.\n2. Enrich caption details by including: \n - Object attributes, such as quantity, color, shape, size, material, state, position, actions, and so on\n - Vision Relations between objects, such as spatial relations, functional relations, possessive relations, attachment relations, action relations, comparative relations, causal relations, and so on\n - Environmental details, such as weather, lighting, colors, textures, atmosphere, and so on\n - Identify the text clearly visible in the image, without translation or explanation, and highlight it in the caption with quotation marks\n3. Maintain authenticity and accuracy:\n - Avoid generalizations\n - Describe all visible information in the image, while do not add information not explicitly shown in the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>assistant\n""" self.default_sample_size = 128 # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._extract_masked_hidden def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() valid_lengths = bool_mask.sum(dim=1) selected = hidden_states[bool_mask] split_result = torch.split(selected, valid_lengths.tolist(), dim=0) return split_result def get_image_caption(self, prompt_image, use_en_prompt=True, device=None): if use_en_prompt: prompt = self.image_caption_prompt_en else: prompt = self.image_caption_prompt_cn model_inputs = self.vl_processor( text=prompt, images=prompt_image, padding=True, return_tensors="pt", ).to(device) with set_forward_context(current_timestep=0, attn_metadata=None): generated_ids = self.text_encoder.generate( **model_inputs, max_new_tokens=512 ) generated_ids_trimmed = [ out_ids[len(in_ids) :] for in_ids, out_ids in zip(model_inputs.input_ids, generated_ids) ] output_text = self.vl_processor.batch_decode( generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False, )[0] return output_text.strip() def _get_qwen_prompt_embeds( self, prompt: Union[str, List[str]] = None, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, ): dtype = dtype or self.text_encoder.dtype prompt = [prompt] if isinstance(prompt, str) else prompt template = self.prompt_template_encode drop_idx = self.prompt_template_encode_start_idx txt = [template.format(e) for e in prompt] txt_tokens = self.tokenizer( txt, padding=True, return_tensors="pt", ).to(device) encoder_hidden_states = self.text_encoder( input_ids=txt_tokens.input_ids, attention_mask=txt_tokens.attention_mask, output_hidden_states=True, ) hidden_states = encoder_hidden_states.hidden_states[-1] split_hidden_states = self._extract_masked_hidden( hidden_states, txt_tokens.attention_mask ) split_hidden_states = [e[drop_idx:] for e in split_hidden_states] attn_mask_list = [ torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states ] max_seq_len = max([e.size(0) for e in split_hidden_states]) prompt_embeds = torch.stack( [ torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states ] ) encoder_attention_mask = torch.stack( [ torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list ] ) prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) return prompt_embeds, encoder_attention_mask @staticmethod def _pack_latents(latents, batch_size, num_channels_latents, height, width, layers): latents = latents.view( batch_size, layers, num_channels_latents, height // 2, 2, width // 2, 2 ) latents = latents.permute(0, 1, 3, 5, 2, 4, 6) latents = latents.reshape( batch_size, layers * (height // 2) * (width // 2), num_channels_latents * 4 ) return latents # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline.encode_prompt def encode_prompt( self, prompt: Union[str, List[str]], device: Optional[torch.device] = None, num_images_per_prompt: int = 1, prompt_embeds: Optional[torch.Tensor] = None, prompt_embeds_mask: Optional[torch.Tensor] = None, max_sequence_length: int = 1024, ): r""" Args: prompt (`str` or `List[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device num_images_per_prompt (`int`): number of images that should be generated per prompt prompt_embeds (`torch.Tensor`, *optional*): Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument. """ prompt = [prompt] if isinstance(prompt, str) else prompt if prompt_embeds is None: prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds( prompt, device ) prompt_embeds = prompt_embeds[:, :max_sequence_length] prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(num_images_per_prompt, seq_len, -1) prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) prompt_embeds_mask = prompt_embeds_mask.view(num_images_per_prompt, seq_len) return prompt_embeds, prompt_embeds_mask # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline._encode_vae_image def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): self.vae = self.vae.to(get_local_torch_device()) if isinstance(generator, list): image_latents = [ retrieve_latents( self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax", ) for i in range(image.shape[0]) ] image_latents = torch.cat(image_latents, dim=0) else: image_latents = retrieve_latents( self.vae.encode(image), generator=generator, sample_mode="argmax" ) latents_mean = ( torch.tensor(self.vae.config.latents_mean) .view(1, self.latent_channels, 1, 1, 1) .to(image_latents.device, image_latents.dtype) ) latents_std = ( torch.tensor(self.vae.config.latents_std) .view(1, self.latent_channels, 1, 1, 1) .to(image_latents.device, image_latents.dtype) ) image_latents = (image_latents - latents_mean) / latents_std self.vae.to("cpu") return image_latents def prepare_latents( self, image, batch_size, num_channels_latents, height, width, layers, dtype, device, generator, latents=None, ): # VAE applies 8x compression on images but we must also account for packing which requires # latent height and width to be divisible by 2. height = 2 * (int(height) // (self.vae_scale_factor * 2)) width = 2 * (int(width) // (self.vae_scale_factor * 2)) shape = ( batch_size, layers + 1, num_channels_latents, height, width, ) ### the generated first image is combined image image_latents = None if image is not None: image = image.to(device=device, dtype=dtype) if image.shape[1] != self.latent_channels: image_latents = self._encode_vae_image(image=image, generator=generator) else: image_latents = image if ( batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0 ): # expand init_latents for batch_size additional_image_per_prompt = batch_size // image_latents.shape[0] image_latents = torch.cat( [image_latents] * additional_image_per_prompt, dim=0 ) elif ( batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0 ): raise ValueError( f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts." ) else: image_latents = torch.cat([image_latents], dim=0) image_latent_height, image_latent_width = image_latents.shape[3:] image_latents = image_latents.permute( 0, 2, 1, 3, 4 ) # (b, c, f, h, w) -> (b, f, c, h, w) image_latents = self._pack_latents( image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width, 1, ) if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) if latents is None: latents = randn_tensor( shape, generator=generator, device=device, dtype=dtype ) latents = self._pack_latents( latents, batch_size, num_channels_latents, height, width, layers + 1 ) else: latents = latents.to(device=device, dtype=dtype) return latents, image_latents def forward( self, batch: Req, server_args: ServerArgs, ) -> Req: use_en_prompt = True device = get_local_torch_device() layers = batch.num_frames num_inference_steps = batch.num_inference_steps generator = batch.generator assert batch.image_path is not None image = load_image(batch.image_path[0]) image = image.convert("RGBA") image_size = image.size resolution = server_args.pipeline_config.resolution calculated_width, calculated_height = calculate_dimensions( resolution * resolution, image_size[0] / image_size[1] ) height = calculated_height width = calculated_width multiple_of = self.vae_scale_factor * 2 width = width // multiple_of * multiple_of height = height // multiple_of * multiple_of # if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): image = self.image_processor.resize(image, calculated_height, calculated_width) prompt_image = image image = self.image_processor.preprocess( image, calculated_height, calculated_width ) image = image.unsqueeze(2) image = image.to(dtype=torch.bfloat16) prompt = self.get_image_caption( prompt_image, use_en_prompt=use_en_prompt, device=device ) prompt_embeds, prompt_embeds_mask = self.encode_prompt( prompt=prompt, device=device, ) negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt( prompt=batch.negative_prompt, device=device, ) num_channels_latents = self.transformer.config.in_channels // 4 latents, image_latents = self.prepare_latents( image, 1, num_channels_latents, height, width, layers, prompt_embeds.dtype, device, generator, ) img_shapes = [ [ *[ ( 1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2, ) for _ in range(layers + 1) ], ( 1, calculated_height // self.vae_scale_factor // 2, calculated_width // self.vae_scale_factor // 2, ), ] ] # 5. Prepare timesteps sigmas = np.linspace(1.0, 0, num_inference_steps + 1)[:-1] image_seq_len = latents.shape[1] base_seqlen = 256 * 256 / 16 / 16 mu = (image_latents.shape[1] / base_seqlen) ** 0.5 timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu, ) txt_seq_lens = ( prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None ) negative_txt_seq_lens = ( negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None ) is_rgb = torch.tensor([0]).to(device=device, dtype=torch.long) batch.prompt_embeds = [prompt_embeds] batch.prompt_embeds_mask = [prompt_embeds_mask] batch.negative_prompt_embeds = [negative_prompt_embeds] batch.negative_prompt_embeds_mask = [negative_prompt_embeds_mask] batch.latents = latents batch.image_latent = image_latents batch.num_inference_steps = num_inference_steps batch.sigmas = sigmas.tolist() # Convert numpy array to list for validation batch.generator = torch.manual_seed(0) batch.original_condition_image_size = image_size batch.raw_latent_shape = latents.shape batch.txt_seq_lens = txt_seq_lens batch.img_shapes = img_shapes return batch ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py ================================================ import torch from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs class LTX2TextConnectorStage(PipelineStage): """ Stage for applying LTX-2 Text Connectors to split/transform text embeddings into video and audio contexts. """ def __init__(self, connectors): super().__init__() self.connectors = connectors def forward(self, batch: Req, server_args: ServerArgs) -> Req: # Input: batch.prompt_embeds (from Gemma, [B, S, D]) # Output: batch.prompt_embeds (Video Context), batch.audio_prompt_embeds (Audio Context) prompt_embeds = batch.prompt_embeds prompt_attention_mask = batch.prompt_attention_mask neg_prompt_embeds = batch.negative_prompt_embeds neg_prompt_attention_mask = batch.negative_attention_mask if isinstance(prompt_embeds, list): prompt_embeds = prompt_embeds[0] if len(prompt_embeds) > 0 else None if isinstance(prompt_attention_mask, list): prompt_attention_mask = ( prompt_attention_mask[0] if len(prompt_attention_mask) > 0 else None ) if isinstance(neg_prompt_embeds, list): neg_prompt_embeds = ( neg_prompt_embeds[0] if len(neg_prompt_embeds) > 0 else None ) if isinstance(neg_prompt_attention_mask, list): neg_prompt_attention_mask = ( neg_prompt_attention_mask[0] if len(neg_prompt_attention_mask) > 0 else None ) # Handle CFG: Concatenate negative and positive inputs if batch.do_classifier_free_guidance: # Concatenate: [Negative, Positive] prompt_embeds = torch.cat([neg_prompt_embeds, prompt_embeds], dim=0) prompt_attention_mask = torch.cat( [neg_prompt_attention_mask, prompt_attention_mask], dim=0 ) # Prepare additive mask for connectors (as per Diffusers implementation) dtype = prompt_embeds.dtype additive_attention_mask = (1 - prompt_attention_mask.to(dtype)) * -1000000.0 # Call connectors # Expects: prompt_embeds, attention_mask, additive_mask=True with set_forward_context(current_timestep=None, attn_metadata=None): connector_prompt_embeds, connector_audio_prompt_embeds, connector_mask = ( self.connectors( prompt_embeds, additive_attention_mask, additive_mask=True ) ) # Split results if CFG was enabled if batch.do_classifier_free_guidance: neg_embeds, pos_embeds = connector_prompt_embeds.chunk(2, dim=0) neg_audio_embeds, pos_audio_embeds = connector_audio_prompt_embeds.chunk( 2, dim=0 ) neg_mask, pos_mask = connector_mask.chunk(2, dim=0) batch.prompt_embeds = [pos_embeds] batch.audio_prompt_embeds = [pos_audio_embeds] batch.prompt_attention_mask = pos_mask batch.negative_prompt_embeds = [neg_embeds] batch.negative_audio_prompt_embeds = [neg_audio_embeds] batch.negative_attention_mask = neg_mask else: # Update positive fields batch.prompt_embeds = [connector_prompt_embeds] batch.audio_prompt_embeds = [connector_audio_prompt_embeds] batch.prompt_attention_mask = connector_mask return batch ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Prompt encoding stages for diffusion pipelines. This module contains implementations of prompt encoding stages for diffusion pipelines. """ import torch from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput from sglang.multimodal_gen.configs.pipeline_configs import FluxPipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.flux import Flux2PipelineConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class TextEncodingStage(PipelineStage): """ Stage for encoding text prompts into embeddings for diffusion models. This stage handles the encoding of text prompts into the embedding space expected by the diffusion model. """ def __init__(self, text_encoders, tokenizers) -> None: """ Initialize the prompt encoding stage. """ super().__init__() self.tokenizers = tokenizers self.text_encoders = text_encoders @torch.no_grad() def forward( self, batch: Req, server_args: ServerArgs, ) -> Req: """ Encode the prompt into text encoder hidden states. """ assert len(self.tokenizers) == len(self.text_encoders) assert len(self.text_encoders) == len( server_args.pipeline_config.text_encoder_configs ) # Encode positive prompt with all available encoders assert batch.prompt is not None prompt_text: str | list[str] = batch.prompt all_indices: list[int] = list(range(len(self.text_encoders))) prompt_embeds_list, prompt_masks_list, pooler_embeds_list = self.encode_text( prompt_text, server_args, encoder_index=all_indices, return_attention_mask=True, ) for pe in prompt_embeds_list: batch.prompt_embeds.append(pe) for pe in pooler_embeds_list: batch.pooled_embeds.append(pe) if batch.prompt_attention_mask is None: batch.prompt_attention_mask = [] for am in prompt_masks_list: batch.prompt_attention_mask.append(am) # Encode negative prompt if CFG is enabled if batch.do_classifier_free_guidance: assert isinstance(batch.negative_prompt, str) neg_embeds_list, neg_masks_list, neg_pooler_embeds_list = self.encode_text( batch.negative_prompt, server_args, encoder_index=all_indices, return_attention_mask=True, ) assert batch.negative_prompt_embeds is not None for ne in neg_embeds_list: batch.negative_prompt_embeds.append(ne) for pe in neg_pooler_embeds_list: batch.neg_pooled_embeds.append(pe) if batch.negative_attention_mask is None: batch.negative_attention_mask = [] for nm in neg_masks_list: batch.negative_attention_mask.append(nm) return batch def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify text encoding stage inputs.""" result = VerificationResult() result.add_check("prompt", batch.prompt, V.string_or_list_strings) result.add_check( "negative_prompt", batch.negative_prompt, lambda x: not batch.do_classifier_free_guidance or V.string_not_none(x), ) result.add_check( "do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value, ) result.add_check("prompt_embeds", batch.prompt_embeds, V.is_list) result.add_check( "negative_prompt_embeds", batch.negative_prompt_embeds, V.none_or_list ) return result def prepare_tokenizer_kwargs(self, tokenizer_kwargs, **kwargs): tok_kwargs = tokenizer_kwargs | kwargs return tok_kwargs @torch.no_grad() def encode_text( self, text: str | list[str], server_args: ServerArgs, encoder_index: int | list[int] | None = None, return_attention_mask: bool = False, return_type: str = "list", # one of: "list", "dict", "stack" device: torch.device | str | None = None, dtype: torch.dtype | None = None, max_length: int | None = None, truncation: bool | None = None, padding: bool | str | None = None, return_overflowing_tokens=None, return_length=None, ): """ Encode plain text using selected text encoder(s) and return embeddings. Args: text: A single string or a list of strings to encode. server_args: The inference arguments providing pipeline config, including tokenizer and encoder settings, preprocess and postprocess functions. encoder_index: Encoder selector by index. Accepts an int or list of ints. return_attention_mask: If True, also return attention masks for each selected encoder. return_type: "list" (default) returns a list aligned with selection; "dict" returns a dict keyed by encoder index as a string; "stack" stacks along a new first dimension (requires matching shapes). device: Optional device override for inputs; defaults to local torch device. dtype: Optional dtype to cast returned embeddings to. max_length: Optional per-call tokenizer override. truncation: Optional per-call tokenizer override. padding: Optional per-call tokenizer override. Returns: Depending on return_type and return_attention_mask: - list: List[Tensor] or (List[Tensor], List[Tensor]) - dict: Dict[str, Tensor] or (Dict[str, Tensor], Dict[str, Tensor]) - stack: Tensor of shape [num_encoders, ...] or a tuple with stacked attention masks """ assert len(self.tokenizers) == len(self.text_encoders) assert len(self.text_encoders) == len( server_args.pipeline_config.text_encoder_configs ) # Resolve selection into indices encoder_cfgs = server_args.pipeline_config.text_encoder_configs if encoder_index is None: indices: list[int] = [0] elif isinstance(encoder_index, int): indices = [encoder_index] else: indices = list(encoder_index) # validate range num_encoders = len(self.text_encoders) for idx in indices: if idx < 0 or idx >= num_encoders: raise IndexError( f"encoder index {idx} out of range [0, {num_encoders - 1}]" ) # Validate indices are within range num_encoders = len(self.text_encoders) # Normalize input to list[str] assert isinstance(text, str | list) if isinstance(text, str): texts: list[str] = [text] else: texts = text embeds_list: list[torch.Tensor] = [] pooled_embeds_list: list[torch.Tensor] = [] attn_masks_list: list[torch.Tensor] = [] preprocess_funcs = server_args.pipeline_config.preprocess_text_funcs postprocess_funcs = server_args.pipeline_config.postprocess_text_funcs text_encoder_extra_args = server_args.pipeline_config.text_encoder_extra_args encoder_cfgs = server_args.pipeline_config.text_encoder_configs if return_type not in ("list", "dict", "stack"): raise ValueError( f"Invalid return_type '{return_type}'. Expected one of: 'list', 'dict', 'stack'" ) target_device = device if device is not None else get_local_torch_device() for i in indices: tokenizer = self.tokenizers[i] text_encoder = self.text_encoders[i] encoder_config = encoder_cfgs[i] preprocess_func = preprocess_funcs[i] postprocess_func = postprocess_funcs[i] text_encoder_extra_arg = ( text_encoder_extra_args[i] if i < len(text_encoder_extra_args) and text_encoder_extra_args[i] else {} ) processed_text_list: list[str] = [] for prompt_str in texts: preprocessed = preprocess_func(prompt_str) processed_text_list.append(preprocessed) # Prepare tokenizer args tok_kwargs = self.prepare_tokenizer_kwargs( encoder_config.tokenizer_kwargs, **text_encoder_extra_arg, ) text_inputs: dict = server_args.pipeline_config.tokenize_prompt( processed_text_list, tokenizer, tok_kwargs ).to(target_device) input_ids = text_inputs["input_ids"] is_flux_v1 = isinstance( server_args.pipeline_config, FluxPipelineConfig ) and not isinstance(server_args.pipeline_config, Flux2PipelineConfig) is_flux_t5 = is_flux_v1 and i == 1 if is_flux_t5: attention_mask = torch.ones(input_ids.shape[:2], device=target_device) else: attention_mask = text_inputs["attention_mask"] with set_forward_context(current_timestep=0, attn_metadata=None): outputs: BaseEncoderOutput = text_encoder( input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True, use_cache=False, ) prompt_embeds = postprocess_func(outputs, text_inputs) if dtype is not None: prompt_embeds = prompt_embeds.to(dtype=dtype) embeds_list.append(prompt_embeds) if is_flux_v1: pooled_embeds_list.append(outputs.pooler_output) if return_attention_mask: attn_masks_list.append(attention_mask) # Shape results according to return_type if return_type == "list": if return_attention_mask: return embeds_list, attn_masks_list, pooled_embeds_list return embeds_list, pooled_embeds_list if return_type == "dict": key_strs = [str(i) for i in indices] embeds_dict = {k: v for k, v in zip(key_strs, embeds_list, strict=False)} if return_attention_mask: attn_dict = { k: v for k, v in zip(key_strs, attn_masks_list, strict=False) } return embeds_dict, attn_dict return embeds_dict # return_type == "stack" # Validate shapes are compatible base_shape = list(embeds_list[0].shape) for t in embeds_list[1:]: if list(t.shape) != base_shape: raise ValueError( f"Cannot stack embeddings with differing shapes: {[list(t.shape) for t in embeds_list]}" ) stacked_embeds = torch.stack(embeds_list, dim=0) if return_attention_mask: base_mask_shape = list(attn_masks_list[0].shape) for m in attn_masks_list[1:]: if list(m.shape) != base_mask_shape: raise ValueError( f"Cannot stack attention masks with differing shapes: {[list(m.shape) for m in attn_masks_list]}" ) stacked_masks = torch.stack(attn_masks_list, dim=0) return stacked_embeds, stacked_masks return stacked_embeds def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify text encoding stage outputs.""" result = VerificationResult() result.add_check( "prompt_embeds", batch.prompt_embeds, V.list_of_tensors_min_dims(2) ) result.add_check( "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x: not batch.do_classifier_free_guidance or V.list_of_tensors_with_min_dims(x, 2), ) if batch.debug: logger.debug(f"{batch.prompt_embeds=}") logger.debug(f"{batch.negative_prompt_embeds=}") return result ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Timestep preparation stages for diffusion pipelines. This module contains implementations of timestep preparation stages for diffusion pipelines. """ import inspect from typing import Any, Callable, Tuple import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( PipelineStage, StageParallelismType, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class TimestepPreparationStage(PipelineStage): """ Stage for preparing timesteps for the diffusion process. This stage handles the preparation of the timestep sequence that will be used during the diffusion process. """ def __init__( self, scheduler, prepare_extra_set_timesteps_kwargs: list[ Callable[[Req, ServerArgs], Tuple[str, Any]] ] = [], ) -> None: super().__init__() self.scheduler = scheduler self.prepare_extra_set_timesteps_kwargs = ( prepare_extra_set_timesteps_kwargs or [] ) @property def parallelism_type(self) -> StageParallelismType: return StageParallelismType.REPLICATED def forward( self, batch: Req, server_args: ServerArgs, ) -> Req: """ Prepare timesteps for the diffusion process. Returns: The batch with prepared timesteps. """ scheduler = self.scheduler device = get_local_torch_device() num_inference_steps = batch.num_inference_steps timesteps = batch.timesteps sigmas = batch.sigmas n_tokens = batch.n_tokens sigmas = server_args.pipeline_config.prepare_sigmas(sigmas, num_inference_steps) batch.sigmas = sigmas # Prepare extra kwargs for set_timesteps extra_set_timesteps_kwargs = {} if ( n_tokens is not None and "n_tokens" in inspect.signature(scheduler.set_timesteps).parameters ): extra_set_timesteps_kwargs["n_tokens"] = n_tokens for callee in self.prepare_extra_set_timesteps_kwargs: key, value = callee(batch, server_args) assert isinstance(key, str) extra_set_timesteps_kwargs[key] = value if key == "mu": batch.extra["mu"] = value # Handle custom timesteps or sigmas if timesteps is not None and sigmas is not None: raise ValueError( "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values" ) if timesteps is not None: accepts_timesteps = ( "timesteps" in inspect.signature(scheduler.set_timesteps).parameters ) if not accepts_timesteps: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" timestep schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps( timesteps=timesteps, device=device, **extra_set_timesteps_kwargs ) timesteps = scheduler.timesteps elif sigmas is not None: accept_sigmas = ( "sigmas" in inspect.signature(scheduler.set_timesteps).parameters ) if not accept_sigmas: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" sigmas schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps( sigmas=sigmas, device=device, **extra_set_timesteps_kwargs ) timesteps = scheduler.timesteps else: scheduler.set_timesteps( num_inference_steps, device=device, **extra_set_timesteps_kwargs ) timesteps = scheduler.timesteps # Update batch with prepared timesteps batch.timesteps = timesteps if not batch.is_warmup: self.log_debug("timesteps: %s", timesteps) return batch def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify timestep preparation stage inputs.""" result = VerificationResult() result.add_check( "num_inference_steps", batch.num_inference_steps, V.positive_int ) result.add_check("timesteps", batch.timesteps, V.none_or_tensor) result.add_check("sigmas", batch.sigmas, V.none_or_list) result.add_check("n_tokens", batch.n_tokens, V.none_or_positive_int) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: """Verify timestep preparation stage outputs.""" if ( batch.is_warmup and isinstance(batch.timesteps, torch.Tensor) and torch.isnan(batch.timesteps).any() ): # when num-inference-steps == 1, the last sigma being 1, the 1 / last_sigma could be nan # this a workaround for warmup req only batch.timesteps = torch.ones( (1,), dtype=torch.float32, device=get_local_torch_device() ) result = VerificationResult() result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.with_dims(1)]) return result ================================================ FILE: python/sglang/multimodal_gen/runtime/pipelines_core/stages/validators.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 """ Common validators for pipeline stage verification. This module provides reusable validation functions that can be used across all pipeline stages for input/output verification. """ from collections.abc import Callable from typing import Any import torch class StageValidators: """Common validators for pipeline stages.""" @staticmethod def not_none(value: Any) -> bool: """Check if value is not None.""" return value is not None @staticmethod def positive_int(value: Any) -> bool: """Check if value is a positive integer.""" return isinstance(value, int) and value > 0 @staticmethod def non_negative_int(value: Any) -> bool: """Check if value is a non-negative float.""" return isinstance(value, int | float) and value >= 0 @staticmethod def positive_float(value: Any) -> bool: """Check if value is a positive float.""" return isinstance(value, int | float) and value > 0 @staticmethod def non_negative_float(value: Any) -> bool: """Check if value is a non-negative float.""" return isinstance(value, int | float) and value >= 0 @staticmethod def divisible_by(value: Any, divisor: int) -> bool: """Check if value is divisible by divisor.""" return value is not None and isinstance(value, int) and value % divisor == 0 @staticmethod def is_tensor(value: Any) -> bool: """Check if value is a torch tensor and doesn't contain NaN values.""" if not isinstance(value, torch.Tensor): return False return not torch.isnan(value).any().item() @staticmethod def tensor_with_dims(value: Any, dims: int) -> bool: """Check if value is a tensor with specific dimensions and no NaN values.""" if not isinstance(value, torch.Tensor): return False if value.dim() != dims: return False return not torch.isnan(value).any().item() @staticmethod def tensor_min_dims(value: Any, min_dims: int) -> bool: """Check if value is a tensor with at least min_dims dimensions and no NaN values.""" if not isinstance(value, torch.Tensor): return False if value.dim() < min_dims: return False return not torch.isnan(value).any().item() @staticmethod def tensor_shape_matches(value: Any, expected_shape: tuple) -> bool: """Check if tensor shape matches expected shape (None for any size) and no NaN values.""" if not isinstance(value, torch.Tensor): return False if len(value.shape) != len(expected_shape): return False for actual, expected in zip(value.shape, expected_shape, strict=True): if expected is not None and actual != expected: return False return not torch.isnan(value).any().item() @staticmethod def list_not_empty(value: Any) -> bool: """Check if value is a non-empty list.""" return isinstance(value, list) and len(value) > 0 @staticmethod def list_length(value: Any, length: int) -> bool: """Check if list has specific length.""" return isinstance(value, list) and len(value) == length @staticmethod def list_min_length(value: Any, min_length: int) -> bool: """Check if list has at least min_length items.""" return isinstance(value, list) and len(value) >= min_length @staticmethod def string_not_empty(value: Any) -> bool: """Check if value is a non-empty string.""" return isinstance(value, str) and len(value.strip()) > 0 @staticmethod def string_not_none(value: Any) -> bool: """Check if value is a non-empty string.""" return isinstance(value, str) and len(value) > 0 @staticmethod def string_or_list_strings(value: Any) -> bool: """Check if value is a string or list of strings.""" if isinstance(value, str): return True if isinstance(value, list): return all(isinstance(item, str) for item in value) return False @staticmethod def bool_value(value: Any) -> bool: """Check if value is a boolean.""" return isinstance(value, bool) @staticmethod def generator_or_list_generators(value: Any) -> bool: """Check if value is a Generator or list of Generators.""" if isinstance(value, torch.Generator): return True if isinstance(value, list): return all(isinstance(item, torch.Generator) for item in value) return False @staticmethod def is_list(value: Any) -> bool: """Check if value is a list (can be empty).""" return isinstance(value, list) @staticmethod def is_tuple(value: Any) -> bool: """Check if value is a tuple.""" return isinstance(value, tuple) @staticmethod def none_or_tensor(value: Any) -> bool: """Check if value is None or a tensor without NaN values.""" if value is None: return True if not isinstance(value, torch.Tensor): return False return not torch.isnan(value).any().item() @staticmethod def list_of_tensors_with_dims(value: Any, dims: int) -> bool: """Check if value is a non-empty list where all items are tensors with specific dimensions and no NaN values.""" if not isinstance(value, list) or len(value) == 0: return False for item in value: if not isinstance(item, torch.Tensor): return False if item.dim() != dims: return False if torch.isnan(item).any().item(): return False return True @staticmethod def list_of_tensors(value: Any) -> bool: """Check if value is a non-empty list where all items are tensors without NaN values.""" if not isinstance(value, list) or len(value) == 0: return False for item in value: if not isinstance(item, torch.Tensor): return False if torch.isnan(item).any().item(): return False return True @staticmethod def list_of_tensors_with_min_dims(value: Any, min_dims: int) -> bool: """Check if value is a non-empty list where all items are tensors with at least min_dims dimensions and no NaN values.""" if not isinstance(value, list) or len(value) == 0: return False for item in value: if not isinstance(item, torch.Tensor): return False if item.dim() < min_dims: return False if torch.isnan(item).any().item(): return False return True @staticmethod def none_or_tensor_with_dims(dims: int) -> Callable[[Any], bool]: """Return a validator that checks if value is None or a tensor with specific dimensions and no NaN values.""" def validator(value: Any) -> bool: if value is None: return True if not isinstance(value, torch.Tensor): return False if value.dim() != dims: return False return not torch.isnan(value).any().item() return validator @staticmethod def none_or_list(value: Any) -> bool: """Check if value is None or a list.""" return value is None or isinstance(value, list) @staticmethod def none_or_positive_int(value: Any) -> bool: """Check if value is None or a positive integer.""" return value is None or (isinstance(value, int) and value > 0) # Helper methods that return functions for common patterns @staticmethod def with_dims(dims: int) -> Callable[[Any], bool]: """Return a validator that checks if tensor has specific dimensions and no NaN values.""" def validator(value: Any) -> bool: return StageValidators.tensor_with_dims(value, dims) return validator @staticmethod def min_dims(min_dims: int) -> Callable[[Any], bool]: """Return a validator that checks if tensor has at least min_dims dimensions and no NaN values.""" def validator(value: Any) -> bool: return StageValidators.tensor_min_dims(value, min_dims) return validator @staticmethod def divisible(divisor: int) -> Callable[[Any], bool]: """Return a validator that checks if value is divisible by divisor.""" def validator(value: Any) -> bool: return StageValidators.divisible_by(value, divisor) return validator @staticmethod def positive_int_divisible(divisor: int) -> Callable[[Any], bool]: """Return a validator that checks if value is a positive integer divisible by divisor.""" def validator(value: Any) -> bool: return ( isinstance(value, int) and value > 0 and StageValidators.divisible_by(value, divisor) ) return validator @staticmethod def list_of_tensors_dims(dims: int) -> Callable[[Any], bool]: """Return a validator that checks if value is a list of tensors with specific dimensions and no NaN values.""" def validator(value: Any) -> bool: return StageValidators.list_of_tensors_with_dims(value, dims) return validator @staticmethod def list_of_tensors_min_dims(min_dims: int) -> Callable[[Any], bool]: """Return a validator that checks if value is a list of tensors with at least min_dims dimensions and no NaN values.""" def validator(value: Any) -> bool: return StageValidators.list_of_tensors_with_min_dims(value, min_dims) return validator class ValidationFailure: """Details about a specific validation failure.""" def __init__( self, validator_name: str, actual_value: Any, expected: str | None = None, error_msg: str | None = None, ): self.validator_name = validator_name self.actual_value = actual_value self.expected = expected self.error_msg = error_msg def __str__(self) -> str: parts = [f"Validator '{self.validator_name}' failed"] if self.error_msg: parts.append(f"Error: {self.error_msg}") # Add actual value info (but limit very long representations) actual_str = self._format_value(self.actual_value) parts.append(f"Actual: {actual_str}") if self.expected: parts.append(f"Expected: {self.expected}") return ". ".join(parts) def _format_value(self, value: Any) -> str: """Format a value for display in error messages.""" if value is None: return "None" elif isinstance(value, torch.Tensor): return f"tensor(shape={list(value.shape)}, dtype={value.dtype})" elif isinstance(value, list): if len(value) == 0: return "[]" elif len(value) <= 3: item_strs = [self._format_value(item) for item in value] return f"[{', '.join(item_strs)}]" else: return f"list(length={len(value)}, first_item={self._format_value(value[0])})" elif isinstance(value, str): if len(value) > 50: return f"'{value[:47]}...'" else: return f"'{value}'" else: return f"{type(value).__name__}({value})" class VerificationResult: """Wrapper class for stage verification results.""" def __init__(self) -> None: self._checks: dict[str, bool] = {} self._failures: dict[str, list[ValidationFailure]] = {} def add_check( self, field_name: str, value: Any, validators: Callable[[Any], bool] | list[Callable[[Any], bool]], ) -> "VerificationResult": """ Add a validation check for a field. Args: field_name: Name of the field being checked value: The actual value to validate validators: Single validation function or list of validation functions. Each function will be called with the value as its first argument. Returns: Self for method chaining Examples: # Single validator result.add_check("tensor", my_tensor, V.is_tensor) # Multiple validators (all must pass) result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) # Using partial functions for parameters result.add_check("height", batch.height, [V.not_none, V.divisible(8)]) """ if not isinstance(validators, list): validators = [validators] failures = [] all_passed = True # Apply all validators and collect detailed failure info for validator in validators: try: passed = validator(value) if not passed: all_passed = False failure = self._create_validation_failure(validator, value) failures.append(failure) except Exception as e: # If any validator raises an exception, consider the check failed all_passed = False validator_name = getattr(validator, "__name__", str(validator)) failure = ValidationFailure( validator_name=validator_name, actual_value=value, error_msg=f"Exception during validation: {str(e)}", ) failures.append(failure) self._checks[field_name] = all_passed if not all_passed: self._failures[field_name] = failures return self def _create_validation_failure( self, validator: Callable, value: Any ) -> ValidationFailure: """Create a ValidationFailure with detailed information.""" validator_name = getattr(validator, "__name__", str(validator)) # Try to extract meaningful expected value info based on validator type expected = None error_msg = None # Handle common validator patterns if hasattr(validator, "__closure__") and validator.__closure__: # This is likely a closure (like our helper functions) if "dims" in validator_name or "with_dims" in str(validator): if isinstance(value, torch.Tensor): expected = f"tensor with {validator.__closure__[0].cell_contents} dimensions" else: expected = "tensor with specific dimensions" elif "divisible" in str(validator): expected = ( f"integer divisible by {validator.__closure__[0].cell_contents}" ) # Handle specific validator types and check for NaN values if validator_name == "is_tensor": expected = "torch.Tensor without NaN values" if isinstance(value, torch.Tensor) and torch.isnan(value).any().item(): error_msg = ( f"tensor contains {torch.isnan(value).sum().item()} NaN values" ) elif validator_name == "positive_int": expected = "positive integer" elif validator_name == "not_none": expected = "non-None value" elif validator_name == "list_not_empty": expected = "non-empty list" elif validator_name == "bool_value": expected = "boolean value" elif ( "tensor_with_dims" in validator_name or "tensor_min_dims" in validator_name ): if isinstance(value, torch.Tensor): if torch.isnan(value).any().item(): error_msg = f"tensor has {value.dim()} dimensions but contains {torch.isnan(value).sum().item()} NaN values" else: error_msg = f"tensor has {value.dim()} dimensions" elif validator_name == "is_list": expected = "list" elif validator_name == "none_or_tensor": expected = "None or tensor without NaN values" if isinstance(value, torch.Tensor) and torch.isnan(value).any().item(): error_msg = ( f"tensor contains {torch.isnan(value).sum().item()} NaN values" ) elif validator_name == "list_of_tensors": expected = "non-empty list of tensors without NaN values" if isinstance(value, list) and len(value) > 0: nan_count = 0 for item in value: if ( isinstance(item, torch.Tensor) and torch.isnan(item).any().item() ): nan_count += torch.isnan(item).sum().item() if nan_count > 0: error_msg = ( f"list contains tensors with total {nan_count} NaN values" ) elif "list_of_tensors_with_dims" in validator_name: expected = ( "non-empty list of tensors with specific dimensions and no NaN values" ) if isinstance(value, list) and len(value) > 0: nan_count = 0 for item in value: if ( isinstance(item, torch.Tensor) and torch.isnan(item).any().item() ): nan_count += torch.isnan(item).sum().item() if nan_count > 0: error_msg = ( f"list contains tensors with total {nan_count} NaN values" ) return ValidationFailure( validator_name=validator_name, actual_value=value, expected=expected, error_msg=error_msg, ) def is_valid(self) -> bool: """Check if all validations passed.""" return all(self._checks.values()) def get_failed_fields(self) -> list[str]: """Get list of fields that failed validation.""" return [field for field, passed in self._checks.items() if not passed] def get_detailed_failures(self) -> dict[str, list[ValidationFailure]]: """Get detailed failure information for each failed field.""" return self._failures.copy() def get_failure_summary(self) -> str: """Get a comprehensive summary of all validation failures.""" if self.is_valid(): return "All validations passed" summary_parts = [] for field_name, failures in self._failures.items(): field_summary = f"\n Field '{field_name}':" for i, failure in enumerate(failures, 1): field_summary += f"\n {i}. {failure}" summary_parts.append(field_summary) return "Validation failures:" + "".join(summary_parts) def to_dict(self) -> dict: """Convert to dictionary for backward compatibility.""" return self._checks.copy() # Alias for convenience V = StageValidators ================================================ FILE: python/sglang/multimodal_gen/runtime/platforms/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/__init__.py import traceback from typing import TYPE_CHECKING # imported by other files, do not remove from sglang.multimodal_gen.runtime.platforms.interface import ( # noqa: F401 AttentionBackendEnum, Platform, PlatformEnum, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import resolve_obj_by_qualname logger = init_logger(__name__) def cuda_platform_plugin() -> str | None: is_cuda = False try: from sglang.multimodal_gen.utils import import_pynvml pynvml = import_pynvml() # type: ignore[no-untyped-call] pynvml.nvmlInit() try: # NOTE: Edge case: sgl_diffusion cpu build on a GPU machine. # Third-party pynvml can be imported in cpu build, # we need to check if sgl_diffusion is built with cpu too. # Otherwise, sgl_diffusion will always activate cuda plugin # on a GPU machine, even if in a cpu build. is_cuda = pynvml.nvmlDeviceGetCount() > 0 finally: pynvml.nvmlShutdown() except Exception as e: if "nvml" not in e.__class__.__name__.lower(): # If the error is not related to NVML, re-raise it. raise e # CUDA is supported on Jetson, but NVML may not be. import os def cuda_is_jetson() -> bool: return os.path.isfile("/etc/nv_tegra_release") or os.path.exists( "/sys/class/tegra-firmware" ) if cuda_is_jetson(): is_cuda = True if is_cuda: logger.debug("CUDA is available") return ( "sglang.multimodal_gen.runtime.platforms.cuda.CudaPlatform" if is_cuda else None ) def mps_platform_plugin() -> str | None: """Detect if MPS (Metal Performance Shaders) is available on macOS.""" is_mps = False try: import torch if torch.backends.mps.is_available(): is_mps = True logger.debug("MPS (Metal Performance Shaders) is available") except Exception as e: logger.debug("MPS detection failed: %s", e) return "sglang.multimodal_gen.runtime.platforms.mps.MpsPlatform" if is_mps else None def cpu_platform_plugin() -> str | None: """Detect if CPU platform should be used.""" # CPU is always available as a fallback return "sglang.multimodal_gen.runtime.platforms.cpu.CpuPlatform" def rocm_platform_plugin() -> str | None: is_rocm = False try: import amdsmi amdsmi.amdsmi_init() try: if len(amdsmi.amdsmi_get_processor_handles()) > 0: is_rocm = True logger.debug("ROCm platform is available") finally: amdsmi.amdsmi_shut_down() except Exception as e: logger.debug("ROCm platform is unavailable: %s", e) return ( "sglang.multimodal_gen.runtime.platforms.rocm.RocmPlatform" if is_rocm else None ) def npu_platform_plugin() -> str | None: is_npu = False try: import torch if torch.npu.is_available(): is_npu = True logger.debug("NPU is available") except Exception as e: logger.debug("NPU detection failed: %s", e) return ( "sglang.multimodal_gen.runtime.platforms.npu.NPUPlatformBase" if is_npu else None ) def musa_platform_plugin() -> str | None: is_musa = False try: import pymtml pymtml.mtmlLibraryInit() try: is_musa = pymtml.mtmlLibraryCountDevice() > 0 finally: pymtml.mtmlLibraryShutDown() except Exception as e: logger.debug("MUSA platform is unavailable: %s", e) return ( "sglang.multimodal_gen.runtime.platforms.musa.MusaPlatform" if is_musa else None ) builtin_platform_plugins = { "cuda": cuda_platform_plugin, "rocm": rocm_platform_plugin, "mps": mps_platform_plugin, "cpu": cpu_platform_plugin, "npu": npu_platform_plugin, "musa": musa_platform_plugin, } def resolve_current_platform_cls_qualname() -> str: # TODO(will): if we need to support other platforms, we should consider if # vLLM's plugin architecture is suitable for our needs. # Try MPS first on macOS platform_cls_qualname = mps_platform_plugin() if platform_cls_qualname is not None: return platform_cls_qualname # Fall back to ROCm platform_cls_qualname = rocm_platform_plugin() if platform_cls_qualname is not None: return platform_cls_qualname # Fall back to CUDA platform_cls_qualname = cuda_platform_plugin() if platform_cls_qualname is not None: return platform_cls_qualname # Fall back to NPU platform_cls_qualname = npu_platform_plugin() if platform_cls_qualname is not None: return platform_cls_qualname # Fall back to MUSA platform_cls_qualname = musa_platform_plugin() if platform_cls_qualname is not None: return platform_cls_qualname # Fall back to CPU as last resort platform_cls_qualname = cpu_platform_plugin() if platform_cls_qualname is not None: return platform_cls_qualname raise RuntimeError("No platform plugin found. Please check your " "installation.") _current_platform: Platform | None = None _init_trace: str = "" current_platform: Platform def __getattr__(name: str): if name == "current_platform": # lazy init current_platform. # 1. out-of-tree platform plugins need `from sglang.multimodal_gen.runtime.platforms import # Platform` so that they can inherit `Platform` class. Therefore, # we cannot resolve `current_platform` during the import of # `sglang.multimodal_gen.runtime.platforms`. global _current_platform if _current_platform is None: platform_cls_qualname = resolve_current_platform_cls_qualname() _current_platform = resolve_obj_by_qualname(platform_cls_qualname)() global _init_trace _init_trace = "".join(traceback.format_stack()) return _current_platform elif name in globals(): return globals()[name] else: raise AttributeError(f"No attribute named '{name}' exists in {__name__}.") __all__ = ["Platform", "PlatformEnum", "current_platform", "_init_trace"] ================================================ FILE: python/sglang/multimodal_gen/runtime/platforms/cpu.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/cpu.py import platform from functools import lru_cache from typing import Any import psutil import torch from sglang.multimodal_gen.runtime.platforms.interface import ( CpuArchEnum, Platform, PlatformEnum, ) class CpuPlatform(Platform): _enum = PlatformEnum.CPU device_name = "CPU" device_type = "cpu" dispatch_key = "CPU" @classmethod def get_cpu_architecture(cls) -> CpuArchEnum: """Get the CPU architecture.""" machine = platform.machine().lower() if machine in ("x86_64", "amd64", "i386", "i686"): return CpuArchEnum.X86 elif machine in ("arm64", "aarch64"): return CpuArchEnum.ARM else: return CpuArchEnum.UNSPECIFIED @classmethod def get_device_name(cls, device_id: int = 0) -> str: return platform.processor() @classmethod def get_device_uuid(cls, device_id: int = 0) -> str: return platform.machine() @classmethod @lru_cache(maxsize=1) def get_device_total_memory(cls, device_id: int = 0) -> int: return psutil.virtual_memory().total @classmethod def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: return True @classmethod def get_current_memory_usage( cls, device: torch.types.Device | None = None ) -> float: # For CPU, we can't easily get memory usage without additional libraries return 0.0 @classmethod def get_available_gpu_memory( cls, device_id: int = 0, distributed: bool = False, empty_cache: bool = True, cpu_group: Any = None, ) -> float: total_free_memory = psutil.virtual_memory().available # For simplicity, we assume 1 NUMA node for now in this platform abstraction # as get_cpu_ids_by_node is not available in multimodal_gen.runtime.utils n_numa_node = 1 free_memory = total_free_memory / n_numa_node if distributed: import torch.distributed as dist tensor = torch.tensor(free_memory, dtype=torch.float32) dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group) free_memory = float(tensor.item()) return free_memory / (1 << 30) @classmethod def get_device_communicator_cls(cls) -> str: return "sglang.multimodal_gen.runtime.distributed.device_communicators.cpu_communicator.CpuCommunicator" ================================================ FILE: python/sglang/multimodal_gen/runtime/platforms/cuda.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/cuda.py """Code inside this file can safely assume cuda platform, e.g. importing pynvml. However, it should not initialize cuda context. """ import os from collections.abc import Callable from functools import lru_cache, wraps from typing import Any, TypeVar import psutil import torch from typing_extensions import ParamSpec from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.platforms.interface import ( AttentionBackendEnum, DeviceCapability, Platform, PlatformEnum, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import import_pynvml logger = init_logger(__name__) _P = ParamSpec("_P") _R = TypeVar("_R") pynvml = import_pynvml() # type: ignore[no-untyped-call] # pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models # see https://github.com/huggingface/diffusers/issues/9704 for details torch.backends.cuda.enable_cudnn_sdp(False) def device_id_to_physical_device_id(device_id: int) -> int: if "CUDA_VISIBLE_DEVICES" in os.environ: device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",") if device_ids == [""]: msg = ( "CUDA_VISIBLE_DEVICES is set to empty string, which means" " GPU support is disabled. If you are using ray, please unset" " the environment variable `CUDA_VISIBLE_DEVICES` inside the" " worker/actor. " "Check https://github.com/vllm-project/vllm/issues/8402 for" " more information." ) raise RuntimeError(msg) physical_device_id = device_ids[device_id] return int(physical_device_id) else: return device_id def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]: @wraps(fn) def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: pynvml.nvmlInit() try: return fn(*args, **kwargs) finally: pynvml.nvmlShutdown() return wrapper class CudaPlatformBase(Platform): _enum = PlatformEnum.CUDA device_name: str = "cuda" device_type: str = "cuda" dispatch_key: str = "CUDA" device_control_env_var: str = "CUDA_VISIBLE_DEVICES" @classmethod def get_local_torch_device(cls) -> torch.device: return torch.device(f"cuda:{envs.LOCAL_RANK}") @classmethod def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: raise NotImplementedError @classmethod def get_device_name(cls, device_id: int = 0) -> str: raise NotImplementedError @classmethod @lru_cache(maxsize=1) def get_device_total_memory(cls, device_id: int = 0) -> int: raise NotImplementedError @classmethod def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: if enforce_eager: logger.warning( "To see benefits of async output processing, enable CUDA " "graph. Since, enforce-eager is enabled, async output " "processor cannot be used" ) return False return True @classmethod def is_full_nvlink(cls, device_ids: list[int]) -> bool: raise NotImplementedError @classmethod def log_warnings(cls) -> None: pass @classmethod def get_current_memory_usage( cls, device: torch.types.Device | None = None ) -> float: torch.cuda.reset_peak_memory_stats(device) return float(torch.cuda.max_memory_allocated(device)) @classmethod def get_available_gpu_memory( cls, device_id: int = 0, distributed: bool = False, empty_cache: bool = True, cpu_group: Any = None, ) -> float: if empty_cache: torch.cuda.empty_cache() if torch.distributed.is_initialized(): device_id = torch.distributed.get_rank() device_props = torch.cuda.get_device_properties(device_id) if device_props.is_integrated: free_gpu_memory = psutil.virtual_memory().available else: free_gpu_memory, _ = torch.cuda.mem_get_info(device_id) if distributed: import torch.distributed as dist tensor = torch.tensor(free_gpu_memory, dtype=torch.float32, device="cuda") dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group) free_gpu_memory = float(tensor.item()) return free_gpu_memory / (1 << 30) @classmethod def get_attn_backend_cls_str( cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype, ) -> str: target_backend: AttentionBackendEnum | None = None # TODO(will): maybe come up with a more general interface for local attention # if distributed is False, we always try to use Flash attn if selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN: try: from st_attn import sliding_tile_attention # noqa: F401 from sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn import ( # noqa: F401 SlidingTileAttentionBackend, ) logger.info("Using Sliding Tile Attention backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend" except ImportError as e: logger.error( "Failed to import Sliding Tile Attention backend: %s", str(e) ) raise ImportError( "Sliding Tile Attention backend is not installed. " ) from e elif selected_backend == AttentionBackendEnum.SAGE_ATTN: try: from sageattention import sageattn # noqa: F401 from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn import ( # noqa: F401 SageAttentionBackend, ) logger.info("Using Sage Attention backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn.SageAttentionBackend" except ImportError as e: logger.info(e) logger.info( "Sage Attention backend is not installed (To install it, run `pip install sageattention==2.2.0 --no-build-isolation`). Falling back to Flash Attention." ) target_backend = AttentionBackendEnum.FA elif selected_backend == AttentionBackendEnum.SAGE_ATTN_3: try: from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3 import ( # noqa: F401 SageAttention3Backend, ) logger.info("Using Sage Attention 3 backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3.SageAttention3Backend" except ImportError as e: logger.info(e) logger.info( "Sage Attention 3 backend is not installed (To install it, see https://github.com/thu-ml/SageAttention/tree/main/sageattention3_blackwell#installation). Falling back to Torch SDPA." ) target_backend = AttentionBackendEnum.TORCH_SDPA elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: try: from vsa import block_sparse_attn # noqa: F401 from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn import ( # noqa: F401 VideoSparseAttentionBackend, ) logger.info("Using Video Sparse Attention backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn.VideoSparseAttentionBackend" except ImportError as e: logger.error( "Failed to import Video Sparse Attention backend: %s", str(e) ) raise ImportError( "Video Sparse Attention backend is not installed." ) from e elif selected_backend == AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN: try: from svg.kernels.triton.permute import ( # noqa: F401 apply_inverse_permutation_triton, permute_tensor_by_labels_triton, ) from svg.kmeans_utils import ( # noqa: F401 batch_kmeans_Euclid, density_calculation, dynamic_block_sparse_fwd_flashinfer, identify_dynamic_map, ) from sglang.multimodal_gen.runtime.layers.attention.backends.sparse_video_gen_2_attn import ( # noqa: F401 SparseVideoGen2AttentionBackend, ) logger.info("Using Sparse Video Gen 2 (SAP) Attention backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_video_gen_2_attn.SparseVideoGen2AttentionBackend" except ImportError as e: logger.error( "Failed to import Sparse Video Gen 2 (SAP) Attention backend: %s", str(e), ) raise ImportError( "Sparse Video Gen 2 (SAP) Attention backend is not installed. " "Please install it by following the instructions at " "https://github.com/svg-project/Sparse-VideoGen" ) from e elif selected_backend == AttentionBackendEnum.VMOBA_ATTN: try: from kernel.attn.vmoba_attn.vmoba import moba_attn_varlen # noqa: F401 from sglang.multimodal_gen.runtime.layers.attention.backends.vmoba import ( # noqa: F401 VMOBAAttentionBackend, ) logger.info("Using Video MOBA Attention backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.vmoba.VMOBAAttentionBackend" except ImportError as e: logger.error( "Failed to import Video MoBA Attention backend: %s", str(e) ) raise ImportError( "Video MoBA Attention backend is not installed. " ) from e elif selected_backend == AttentionBackendEnum.AITER: logger.info("Using AITer backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.aiter.AITerBackend" elif selected_backend == AttentionBackendEnum.TORCH_SDPA: logger.info("Using Torch SDPA backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" elif selected_backend == AttentionBackendEnum.SLA_ATTN: logger.info("Using Sparse Linear Attention backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SparseLinearAttentionBackend" elif selected_backend == AttentionBackendEnum.SAGE_SLA_ATTN: logger.info("Using Sage Sparse Linear Attention backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SageSparseLinearAttentionBackend" elif selected_backend == AttentionBackendEnum.FA2: from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn_2 import ( # noqa: F401 FlashAttention2Backend, ) logger.info("Using FlashAttention2 backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn_2.FlashAttention2Backend" elif selected_backend in [ AttentionBackendEnum.FA, ]: if cls.is_sm120(): logger.info( "FlashAttention is not supported on SM12.x in this build; falling back to Torch SDPA." ) target_backend = AttentionBackendEnum.TORCH_SDPA elif cls.is_blackwell(): from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( set_fa_ver, ) set_fa_ver(4) target_backend = AttentionBackendEnum.FA else: target_backend = AttentionBackendEnum.FA elif selected_backend: raise ValueError(f"Invalid attention backend for {cls.device_name}") else: if cls.is_sm120(): # On SM12.x, the sgl-kernel FlashAttention wheels may not include # support yet. Default to Torch SDPA for correctness. logger.info("Defaulting to Torch SDPA backend on SM12.x") target_backend = AttentionBackendEnum.TORCH_SDPA elif cls.is_blackwell(): from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( set_fa_ver, ) set_fa_ver(4) target_backend = AttentionBackendEnum.FA else: target_backend = AttentionBackendEnum.FA # Ensure we have a target backend selected before validation/fallback. if target_backend is None: target_backend = AttentionBackendEnum.FA if target_backend == AttentionBackendEnum.FA and cls.is_blackwell(): from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( set_fa_ver, ) set_fa_ver(4) if not cls.has_device_capability(80): logger.info("Cannot use FlashAttention backend for Volta and Turing GPUs.") target_backend = AttentionBackendEnum.TORCH_SDPA elif dtype not in (torch.float16, torch.bfloat16): logger.info( "Cannot use FlashAttention backend for dtype other than " "torch.float16 or torch.bfloat16." ) target_backend = AttentionBackendEnum.TORCH_SDPA # FlashAttn is valid for the model, checking if the package is # installed. if target_backend == AttentionBackendEnum.FA: try: from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( # noqa: F401 FlashAttentionBackend, ) supported_sizes = FlashAttentionBackend.get_supported_head_sizes() if head_size not in supported_sizes: logger.info( "Cannot use FlashAttention backend for head size %d.", head_size, ) target_backend = AttentionBackendEnum.TORCH_SDPA except ImportError: logger.info( "Cannot use FlashAttention backend because the " "flash_attn package is not found. " "Make sure that flash_attn was built and installed " "(on by default)." ) target_backend = AttentionBackendEnum.TORCH_SDPA if target_backend == AttentionBackendEnum.TORCH_SDPA: logger.info("Using Torch SDPA backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" logger.info("Using FlashAttention (FA3 for hopper, FA4 for blackwell) backend") return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn.FlashAttentionBackend" @classmethod def get_device_communicator_cls(cls) -> str: return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa # NVML utils # Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`, # all the related functions work on real physical device ids. # the major benefit of using NVML is that it will not initialize CUDA class NvmlCudaPlatform(CudaPlatformBase): @classmethod @lru_cache(maxsize=8) @with_nvml_context def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: try: physical_device_id = device_id_to_physical_device_id(device_id) handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle) return DeviceCapability(major=major, minor=minor) except RuntimeError: return None @classmethod @lru_cache(maxsize=8) @with_nvml_context def has_device_capability( cls, capability: tuple[int, int] | int, device_id: int = 0, ) -> bool: try: return bool(super().has_device_capability(capability, device_id)) except RuntimeError: return False @classmethod @lru_cache(maxsize=8) @with_nvml_context def get_device_name(cls, device_id: int = 0) -> str: physical_device_id = device_id_to_physical_device_id(device_id) return cls._get_physical_device_name(physical_device_id) @classmethod @lru_cache(maxsize=8) @with_nvml_context def get_device_uuid(cls, device_id: int = 0) -> str: physical_device_id = device_id_to_physical_device_id(device_id) handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) return str(pynvml.nvmlDeviceGetUUID(handle)) @classmethod @lru_cache(maxsize=8) @with_nvml_context def get_device_total_memory(cls, device_id: int = 0) -> int: physical_device_id = device_id_to_physical_device_id(device_id) handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total) @classmethod @with_nvml_context def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool: """ query if the set of gpus are fully connected by nvlink (1 hop) """ handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids] for i, handle in enumerate(handles): for j, peer_handle in enumerate(handles): if i < j: try: p2p_status = pynvml.nvmlDeviceGetP2PStatus( handle, peer_handle, pynvml.NVML_P2P_CAPS_INDEX_NVLINK, ) if p2p_status != pynvml.NVML_P2P_STATUS_OK: return False except pynvml.NVMLError: logger.exception( "NVLink detection failed. This is normal if" " your machine has no NVLink equipped." ) return False return True @classmethod def _get_physical_device_name(cls, device_id: int = 0) -> str: handle = pynvml.nvmlDeviceGetHandleByIndex(device_id) return str(pynvml.nvmlDeviceGetName(handle)) @classmethod @with_nvml_context def log_warnings(cls) -> None: device_ids: int = pynvml.nvmlDeviceGetCount() if device_ids > 1: device_names = [cls._get_physical_device_name(i) for i in range(device_ids)] if ( len(set(device_names)) > 1 and os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID" ): logger.warning( "Detected different devices in the system: %s. Please" " make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to " "avoid unexpected behavior.", ", ".join(device_names), ) class NonNvmlCudaPlatform(CudaPlatformBase): @classmethod def get_device_capability(cls, device_id: int = 0) -> DeviceCapability: major, minor = torch.cuda.get_device_capability(device_id) return DeviceCapability(major=major, minor=minor) @classmethod def get_device_name(cls, device_id: int = 0) -> str: return str(torch.cuda.get_device_name(device_id)) @classmethod @lru_cache(maxsize=1) def get_device_total_memory(cls, device_id: int = 0) -> int: device_props = torch.cuda.get_device_properties(device_id) return int(device_props.total_memory) @classmethod def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool: logger.exception( "NVLink detection not possible, as context support was" " not found. Assuming no NVLink available." ) return False # Autodetect either NVML-enabled or non-NVML platform # based on whether NVML is available. nvml_available = False try: try: pynvml.nvmlInit() nvml_available = True except Exception: # On Jetson, NVML is not supported. nvml_available = False finally: if nvml_available: pynvml.nvmlShutdown() CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform try: from sphinx.ext.autodoc.mock import _MockModule if not isinstance(pynvml, _MockModule): CudaPlatform.log_warnings() except ModuleNotFoundError: CudaPlatform.log_warnings() ================================================ FILE: python/sglang/multimodal_gen/runtime/platforms/interface.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/interface.py from __future__ import annotations import enum import random from functools import lru_cache from typing import TYPE_CHECKING, Any, NamedTuple import numpy as np import torch from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import resolve_obj_by_qualname if TYPE_CHECKING: from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionImpl, ) logger = init_logger(__name__) class AttentionBackendEnum(enum.Enum): FA2 = enum.auto() FA = enum.auto() SLIDING_TILE_ATTN = enum.auto() TORCH_SDPA = enum.auto() SAGE_ATTN = enum.auto() SAGE_ATTN_3 = enum.auto() VIDEO_SPARSE_ATTN = enum.auto() SPARSE_VIDEO_GEN_2_ATTN = enum.auto() VMOBA_ATTN = enum.auto() AITER = enum.auto() AITER_SAGE = enum.auto() SLA_ATTN = enum.auto() SAGE_SLA_ATTN = enum.auto() NO_ATTENTION = enum.auto() def __str__(self): return self.name.lower() @property def is_sparse(self) -> bool: return self in { AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SLA_ATTN, AttentionBackendEnum.SAGE_SLA_ATTN, } class PlatformEnum(enum.Enum): CUDA = enum.auto() ROCM = enum.auto() TPU = enum.auto() CPU = enum.auto() MPS = enum.auto() NPU = enum.auto() MUSA = enum.auto() OOT = enum.auto() UNSPECIFIED = enum.auto() class CpuArchEnum(enum.Enum): X86 = enum.auto() ARM = enum.auto() UNSPECIFIED = enum.auto() class DeviceCapability(NamedTuple): major: int minor: int def as_version_str(self) -> str: return f"{self.major}.{self.minor}" def to_int(self) -> int: """ Express device capability as an integer ````. It is assumed that the minor version is always a single digit. """ assert 0 <= self.minor < 10 return self.major * 10 + self.minor class Platform: _enum: PlatformEnum device_name: str device_type: str device: torch.device | None = None # Dummy attribute for compatibility # available dispatch keys: # check https://github.com/pytorch/pytorch/blob/313dac6c1ca0fa0cde32477509cce32089f8532a/torchgen/model.py#L134 # noqa # use "CPU" as a fallback for platforms not registered in PyTorch dispatch_key: str = "CPU" # The torch.compile backend for compiling simple and # standalone functions. The default value is "inductor" to keep # the same behavior as PyTorch. # NOTE: for the forward part of the model, vLLM has another separate # compilation strategy. simple_compile_backend: str = "inductor" supported_quantization: list[str] = [] @lru_cache(maxsize=1) def is_cuda(self) -> bool: return self.is_cuda_static() @lru_cache(maxsize=1) def is_npu(self) -> bool: return self._enum == PlatformEnum.NPU @lru_cache(maxsize=1) def is_rocm(self) -> bool: return self.is_rocm_static() @lru_cache(maxsize=1) def is_tpu(self) -> bool: return self._enum == PlatformEnum.TPU @lru_cache(maxsize=1) def is_cpu(self) -> bool: return self._enum == PlatformEnum.CPU @classmethod @lru_cache(maxsize=1) def is_blackwell(cls): if not cls.is_cuda_static(): return False return torch.cuda.get_device_capability()[0] == 10 @classmethod @lru_cache(maxsize=1) def is_hopper(cls): if not cls.is_cuda_static(): return False return torch.cuda.get_device_capability() == (9, 0) @classmethod @lru_cache(maxsize=1) def is_sm120(cls): if not cls.is_cuda_static(): return False return torch.cuda.get_device_capability()[0] == 12 @classmethod def is_cuda_static(cls) -> bool: return getattr(cls, "_enum", None) == PlatformEnum.CUDA @classmethod def is_rocm_static(cls) -> bool: return getattr(cls, "_enum", None) == PlatformEnum.ROCM @lru_cache(maxsize=1) def is_hpu(self) -> bool: return hasattr(torch, "hpu") and torch.hpu.is_available() @lru_cache(maxsize=1) def is_xpu(self) -> bool: return hasattr(torch, "xpu") and torch.xpu.is_available() @lru_cache(maxsize=1) def is_npu(self) -> bool: return hasattr(torch, "npu") and torch.npu.is_available() def is_out_of_tree(self) -> bool: return self._enum == PlatformEnum.OOT @lru_cache(maxsize=1) def is_cuda_alike(self) -> bool: """Stateless version of :func:`torch.cuda.is_available`.""" return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM, PlatformEnum.MUSA) @lru_cache(maxsize=1) def is_mps(self) -> bool: return self._enum == PlatformEnum.MPS @lru_cache(maxsize=1) def is_musa(self): try: return hasattr(torch, "musa") and torch.musa.is_available() except ModuleNotFoundError: return False @lru_cache(maxsize=1) def is_hip(self) -> bool: return self.is_rocm() @classmethod @lru_cache(maxsize=1) def is_amp_supported(cls) -> bool: return True @classmethod def get_local_torch_device(cls) -> torch.device: raise NotImplementedError @classmethod def get_attn_backend_cls_str( cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype, ) -> str: """Get the attention backend class of a device.""" return "" @classmethod def get_device_capability( cls, device_id: int = 0, ) -> DeviceCapability | None: """Stateless version of :func:`torch.cuda.get_device_capability`.""" return None @classmethod def has_device_capability( cls, capability: tuple[int, int] | int, device_id: int = 0, ) -> bool: """ Test whether this platform is compatible with a device capability. The ``capability`` argument can either be: - A tuple ``(major, minor)``. - An integer ````. (See :meth:`DeviceCapability.to_int`) """ current_capability = cls.get_device_capability(device_id=device_id) if current_capability is None: return False if isinstance(capability, tuple): return current_capability >= capability return current_capability.to_int() >= capability @classmethod def get_device_name(cls, device_id: int = 0) -> str: """Get the name of a device.""" raise NotImplementedError @classmethod def get_device_uuid(cls, device_id: int = 0) -> str: """Get the uuid of a device, e.g. the PCI bus ID.""" raise NotImplementedError @classmethod @lru_cache(maxsize=1) def get_device_total_memory(cls, device_id: int = 0) -> int: """Get the total memory of a device in bytes.""" raise NotImplementedError @lru_cache(maxsize=1) def get_device(self, local_rank: int) -> torch.device: if self.is_cuda() or self.is_rocm(): return torch.device("cuda", local_rank) elif self.is_npu(): return torch.device("npu", local_rank) elif self.is_musa(): return torch.device("musa", local_rank) elif self.is_mps(): return torch.device("mps") else: return torch.device("cpu") @lru_cache(maxsize=1) def get_torch_distributed_backend_str(self) -> str: if self.is_cuda_alike(): return "nccl" elif self.is_npu(): return "hccl" elif self.is_musa(): return "mccl" elif self.is_mps(): return "gloo" else: raise NotImplementedError( "No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available" ) @classmethod def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: """ Check if the current platform supports async output. """ raise NotImplementedError @classmethod def inference_mode(cls): """A device-specific wrapper of `torch.inference_mode`. This wrapper is recommended because some hardware backends such as TPU do not support `torch.inference_mode`. In such a case, they will fall back to `torch.no_grad` by overriding this method. """ return torch.inference_mode(mode=True) @classmethod def seed_everything(cls, seed: int | None = None) -> None: """ Set the seed of each random module. `torch.manual_seed` will set seed on all devices. Loosely based on: https://github.com/Lightning-AI/pytorch-lightning/blob/2.4.0/src/lightning/fabric/utilities/seed.py#L20 """ if seed is not None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.get_device_module().manual_seed_all(seed) @classmethod def verify_model_arch(cls, model_arch: str) -> None: """ Verify whether the current platform supports the specified model architecture. - This will raise an Error or Warning based on the model support on the current platform. - By default all models are considered supported. """ pass @classmethod def verify_quantization(cls, quant: str) -> None: """ Verify whether the quantization is supported by the current platform. """ if cls.supported_quantization and quant not in cls.supported_quantization: raise ValueError( f"{quant} quantization is currently not supported in " f"{cls.device_name}." ) @classmethod def get_current_memory_usage( cls, device: torch.types.Device | None = None ) -> float: """ Return the memory usage in bytes. """ raise NotImplementedError @classmethod def get_available_gpu_memory( cls, device_id: int = 0, distributed: bool = False, empty_cache: bool = True, cpu_group: Any = None, ) -> float: """ Return the available memory in GiB. """ raise NotImplementedError @classmethod def get_device_communicator_cls(cls) -> str: """ Get device specific communicator class for distributed communication. """ return "sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" # noqa @classmethod def get_cpu_architecture(cls) -> CpuArchEnum: """Get the CPU architecture of the current platform.""" return CpuArchEnum.UNSPECIFIED @classmethod def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool: """Whether to enable DIT layerwise offload by default on the current platform.""" return True @classmethod def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module: """Apply platform-specific optimizations to VAE after loading.""" return vae def get_attn_backend(self, *args, **kwargs) -> AttentionImpl: attention_cls_str = self.get_attn_backend_cls_str(*args, **kwargs) return resolve_obj_by_qualname(attention_cls_str) class UnspecifiedPlatform(Platform): _enum = PlatformEnum.UNSPECIFIED device_type = "" ================================================ FILE: python/sglang/multimodal_gen/runtime/platforms/mps.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from functools import lru_cache from typing import Any import psutil import torch from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, Platform, PlatformEnum, ) from sglang.multimodal_gen.runtime.platforms.interface import DeviceCapability from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger # SPDX-License-Identifier: Apache-2.0 logger = init_logger(__name__) class MpsPlatform(Platform): _enum = PlatformEnum.MPS device_name: str = "mps" device_type: str = "mps" dispatch_key: str = "MPS" device_control_env_var: str = "MPS_VISIBLE_DEVICES" @classmethod @lru_cache(maxsize=1) def is_amp_supported(cls) -> bool: return False @classmethod def get_local_torch_device(cls) -> torch.device: return torch.device("mps") @classmethod def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: raise NotImplementedError @classmethod def get_device_name(cls, device_id: int = 0) -> str: raise NotImplementedError @classmethod def get_device_uuid(cls, device_id: int = 0) -> str: raise NotImplementedError @classmethod @lru_cache(maxsize=1) def get_device_total_memory(cls, device_id: int = 0) -> int: return psutil.virtual_memory().total @classmethod def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: if enforce_eager: logger.warning( "To see benefits of async output processing, enable MPS " "graph. Since, enforce-eager is enabled, async output " "processor cannot be used" ) return False return True @classmethod def get_current_memory_usage( cls, device: torch.types.Device | None = None ) -> float: return 0.0 @classmethod def get_available_gpu_memory( cls, device_id: int = 0, distributed: bool = False, empty_cache: bool = True, cpu_group: Any = None, ) -> float: if empty_cache: torch.mps.empty_cache() # For MPS, available memory is essentially the system available memory free_memory = psutil.virtual_memory().available if distributed: import torch.distributed as dist tensor = torch.tensor(free_memory, dtype=torch.float32) dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group) free_memory = float(tensor.item()) return free_memory / (1 << 30) @classmethod def get_attn_backend_cls_str( cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype, ) -> str: # MPS supports SDPA (Scaled Dot-Product Attention) which is the most compatible logger.info("Using Torch SDPA backend for MPS.") return ( "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" ) @classmethod def get_device_communicator_cls(cls) -> str: # Use base communicator for MPS return "sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" @classmethod def seed_everything(cls, seed: int | None = None) -> None: """Set the seed for MPS device.""" if seed is not None: import random import numpy as np random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) # MPS doesn't have manual_seed_all like CUDA # The manual_seed above should be sufficient ================================================ FILE: python/sglang/multimodal_gen/runtime/platforms/musa.py ================================================ """ This file is a platform abstraction for MThreads (MUSA) GPUs, adjusted to match the structure and interface of `cuda.py`. """ import os from collections.abc import Callable from functools import lru_cache, wraps from typing import Any, TypeVar import psutil import pymtml # isort: off import torch import torchada # noqa: F401 # isort: on from typing_extensions import ParamSpec from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.platforms.interface import ( AttentionBackendEnum, DeviceCapability, Platform, PlatformEnum, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) _P = ParamSpec("_P") _R = TypeVar("_R") def device_id_to_physical_device_id(device_id: int) -> int: if "MUSA_VISIBLE_DEVICES" in os.environ: device_ids = os.environ["MUSA_VISIBLE_DEVICES"].split(",") if device_ids == [""]: msg = ( "MUSA_VISIBLE_DEVICES is set to empty string, which means" " GPU support is disabled. If you are using ray, please unset" " the environment variable `MUSA_VISIBLE_DEVICES` inside the" " worker/actor. " "Check https://github.com/vllm-project/vllm/issues/8402 for" " more information." ) raise RuntimeError(msg) physical_device_id = device_ids[device_id] return int(physical_device_id) else: return device_id def with_mtml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]: @wraps(fn) def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: pymtml.nvmlInit() try: return fn(*args, **kwargs) finally: pymtml.nvmlShutdown() return wrapper class MusaPlatformBase(Platform): _enum = PlatformEnum.MUSA device_name: str = "musa" device_type: str = "musa" dispatch_key: str = "MUSA" device_control_env_var: str = "MUSA_VISIBLE_DEVICES" @classmethod def get_local_torch_device(cls) -> torch.device: return torch.device(f"musa:{envs.LOCAL_RANK}") @classmethod def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: raise NotImplementedError @classmethod def get_device_name(cls, device_id: int = 0) -> str: raise NotImplementedError @classmethod @lru_cache(maxsize=1) def get_device_total_memory(cls, device_id: int = 0) -> int: raise NotImplementedError @classmethod def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: if enforce_eager: logger.warning( "To see benefits of async output processing, enable MUSA " "graph. Since, enforce-eager is enabled, async output " "processor cannot be used" ) return False return True @classmethod def is_full_mtlink(cls, device_ids: list[int]) -> bool: raise NotImplementedError @classmethod def log_warnings(cls) -> None: pass @classmethod def get_current_memory_usage( cls, device: torch.types.Device | None = None ) -> float: torch.cuda.reset_peak_memory_stats(device) return float(torch.cuda.max_memory_allocated(device)) @classmethod def get_available_gpu_memory( cls, device_id: int = 0, distributed: bool = False, empty_cache: bool = True, cpu_group: Any = None, ) -> float: if empty_cache: torch.cuda.empty_cache() if torch.distributed.is_initialized(): device_id = torch.distributed.get_rank() device_props = torch.cuda.get_device_properties(device_id) if device_props.is_integrated: free_gpu_memory = psutil.virtual_memory().available else: free_gpu_memory, _ = torch.cuda.mem_get_info(device_id) if distributed: import torch.distributed as dist tensor = torch.tensor(free_gpu_memory, dtype=torch.float32, device="musa") dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group) free_gpu_memory = float(tensor.item()) return free_gpu_memory / (1 << 30) @classmethod def get_attn_backend_cls_str( cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype, ) -> str: logger.info("Using Torch SDPA backend.") return ( "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" ) @classmethod def get_device_communicator_cls(cls) -> str: return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa # MTML utils # Note that MTML is not affected by `MUSA_VISIBLE_DEVICES`, # all the related functions work on real physical device ids. # the major benefit of using MTML is that it will not initialize MUSA class MtmlMusaPlatform(MusaPlatformBase): @classmethod @lru_cache(maxsize=8) @with_mtml_context def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: try: physical_device_id = device_id_to_physical_device_id(device_id) handle = pymtml.nvmlDeviceGetHandleByIndex(physical_device_id) major, minor = pymtml.nvmlDeviceGetCudaComputeCapability(handle) return DeviceCapability(major=major, minor=minor) except RuntimeError: return None @classmethod @lru_cache(maxsize=8) @with_mtml_context def has_device_capability( cls, capability: tuple[int, int] | int, device_id: int = 0, ) -> bool: try: return bool(super().has_device_capability(capability, device_id)) except RuntimeError: return False @classmethod @lru_cache(maxsize=8) @with_mtml_context def get_device_name(cls, device_id: int = 0) -> str: physical_device_id = device_id_to_physical_device_id(device_id) return cls._get_physical_device_name(physical_device_id) @classmethod @lru_cache(maxsize=8) @with_mtml_context def get_device_uuid(cls, device_id: int = 0) -> str: physical_device_id = device_id_to_physical_device_id(device_id) handle = pymtml.nvmlDeviceGetHandleByIndex(physical_device_id) return str(pymtml.nvmlDeviceGetUUID(handle)) @classmethod @lru_cache(maxsize=8) @with_mtml_context def get_device_total_memory(cls, device_id: int = 0) -> int: physical_device_id = device_id_to_physical_device_id(device_id) handle = pymtml.nvmlDeviceGetHandleByIndex(physical_device_id) return int(pymtml.nvmlDeviceGetMemoryInfo(handle).total) @classmethod @with_mtml_context def is_full_mtlink(cls, physical_device_ids: list[int]) -> bool: """ query if the set of gpus are fully connected by mtlink (1 hop) """ handles = [pymtml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids] for i, handle in enumerate(handles): for j, peer_handle in enumerate(handles): if i < j: try: p2p_status = pymtml.nvmlDeviceGetP2PStatus( handle, peer_handle, pymtml.NVML_P2P_CAPS_INDEX_NVLINK, ) if p2p_status != pymtml.NVML_P2P_STATUS_OK: return False except pymtml.NVMLError: logger.exception( "MTLink detection failed. This is normal if" " your machine has no MTLink equipped." ) return False return True @classmethod def _get_physical_device_name(cls, device_id: int = 0) -> str: handle = pymtml.nvmlDeviceGetHandleByIndex(device_id) return str(pymtml.nvmlDeviceGetName(handle)) @classmethod @with_mtml_context def log_warnings(cls) -> None: device_ids: int = pymtml.nvmlDeviceGetCount() if device_ids > 1: device_names = [cls._get_physical_device_name(i) for i in range(device_ids)] if ( len(set(device_names)) > 1 and os.environ.get("MUSA_DEVICE_ORDER") != "PCI_BUS_ID" ): logger.warning( "Detected different devices in the system: %s. Please" " make sure to set `MUSA_DEVICE_ORDER=PCI_BUS_ID` to " "avoid unexpected behavior.", ", ".join(device_names), ) class NonMtmlMusaPlatform(MusaPlatformBase): @classmethod def get_device_capability(cls, device_id: int = 0) -> DeviceCapability: major, minor = torch.cuda.get_device_capability(device_id) return DeviceCapability(major=major, minor=minor) @classmethod def get_device_name(cls, device_id: int = 0) -> str: return str(torch.cuda.get_device_name(device_id)) @classmethod @lru_cache(maxsize=1) def get_device_total_memory(cls, device_id: int = 0) -> int: device_props = torch.cuda.get_device_properties(device_id) return int(device_props.total_memory) @classmethod def is_full_mtlink(cls, physical_device_ids: list[int]) -> bool: logger.error( "MTLink detection not possible, as context support was" " not found. Assuming no MTLink available." ) return False # Autodetect either MTML-enabled or non-MTML platform # based on whether MTML is available. mtml_available = False if "MUSA_DISABLE_MTML" not in os.environ: try: try: pymtml.nvmlInit() mtml_available = True except Exception: mtml_available = False finally: if mtml_available: pymtml.nvmlShutdown() MusaPlatform = MtmlMusaPlatform if mtml_available else NonMtmlMusaPlatform try: from sphinx.ext.autodoc.mock import _MockModule if not isinstance(pymtml, _MockModule): MusaPlatform.log_warnings() except ModuleNotFoundError: MusaPlatform.log_warnings() if __name__ == "__main__": print(MusaPlatform.__name__) print(MusaPlatform.get_device_name()) print(MusaPlatform.get_device_capability()) print(MusaPlatform.get_device_total_memory()) print(MusaPlatform.is_full_mtlink([0, 1, 2, 3, 4, 5, 6, 7])) ================================================ FILE: python/sglang/multimodal_gen/runtime/platforms/npu.py ================================================ # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm-ascend: https://github.com/vllm-project/vllm-ascend/blob/main/vllm_ascend/platform.py import os from typing import Any import torch from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.platforms.interface import ( AttentionBackendEnum, DeviceCapability, Platform, PlatformEnum, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) def device_id_to_physical_device_id(device_id: int) -> int: if "ASCEND_RT_VISIBLE_DEVICES" in os.environ: device_ids = os.environ["ASCEND_RT_VISIBLE_DEVICES"].split(",") if device_ids == [""]: msg = ( "ASCEND_RT_VISIBLE_DEVICES is set to empty string, which means" " NPU support is disabled" ) raise RuntimeError(msg) physical_device_id = device_ids[device_id] return int(physical_device_id) else: return device_id class NPUPlatformBase(Platform): _enum = PlatformEnum.NPU device_name: str = "npu" device_type: str = "npu" dispatch_key: str = "NPU" device_control_env_var: str = "ASCEND_RT_VISIBLE_DEVICES" @classmethod def get_local_torch_device(cls) -> torch.device: return torch.device(f"npu:{envs.LOCAL_RANK}") @classmethod def get_device_capability(cls, device_id: int = 0) -> DeviceCapability: return None @classmethod def get_device_name(cls, device_id: int = 0) -> str: return str(torch.npu.get_device_name(device_id)) @classmethod def get_device_total_memory(cls, device_id: int = 0) -> int: device_props = torch.npu.get_device_properties(device_id) return int(device_props.total_memory) @classmethod def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: if enforce_eager: logger.warning( "To see benefits of async output processing, enable NPU " "graph. Since, enforce-eager is enabled, async output " "processor cannot be used" ) return False return True @classmethod def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool: logger.exception( "NVLink detection not possible, as context support was" " not found. Assuming no NVLink available." ) return False @classmethod def get_available_gpu_memory( cls, device_id: int = 0, distributed: bool = False, empty_cache: bool = True, cpu_group: Any = None, ) -> float: if empty_cache: torch.npu.empty_cache() free_gpu_memory, _ = torch.npu.mem_get_info(device_id) if distributed: import torch.distributed as dist tensor = torch.tensor(free_gpu_memory, dtype=torch.float32, device="npu") dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group) free_gpu_memory = float(tensor.item()) return free_gpu_memory / (1 << 30) @classmethod def log_warnings(cls) -> None: pass @classmethod def get_current_memory_usage( cls, device: torch.types.Device | None = None ) -> float: torch.npu.reset_peak_memory_stats(device) return float(torch.npu.max_memory_allocated(device)) @classmethod def get_attn_backend_cls_str( cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype, ) -> str: logger.info("Using Torch SDPA backend.") return ( "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" ) @classmethod def get_device_communicator_cls(cls) -> str: return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa @classmethod def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool: """The performance of the layerwise_offload feature depends on the device's memory size and the memory size occupied by the model. Use --dit-layerwise-offload True if it suitable for your case.""" return False ================================================ FILE: python/sglang/multimodal_gen/runtime/platforms/rocm.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from rocm/vllm: https://github.com/ROCm/vllm/blob/v0.7.3%2Brocm/vllm/platforms/rocm.py """ This file is a platform abstraction for ROCm GPUs, adjusted to match the structure and interface of `cuda.py`. """ from functools import lru_cache from typing import Any import torch import sglang.multimodal_gen.envs as envs from sglang.multimodal_gen.runtime.platforms.interface import ( AttentionBackendEnum, DeviceCapability, Platform, PlatformEnum, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # ROCm uses the same torch.cuda interface class RocmPlatform(Platform): _enum = PlatformEnum.ROCM device_name: str = "rocm" device_type: str = "cuda" # torch uses 'cuda' backend string dispatch_key: str = "CUDA" device_control_env_var: str = "CUDA_VISIBLE_DEVICES" @classmethod def get_local_torch_device(cls) -> torch.device: return torch.device(f"cuda:{envs.LOCAL_RANK}") @classmethod def get_device_capability(cls, device_id: int = 0) -> DeviceCapability: major, minor = torch.cuda.get_device_capability(device_id) return DeviceCapability(major=major, minor=minor) @classmethod def get_device_name(cls, device_id: int = 0) -> str: return str(torch.cuda.get_device_name(device_id)) @classmethod @lru_cache(maxsize=1) def get_device_total_memory(cls, device_id: int = 0) -> int: return torch.cuda.get_device_properties(device_id).total_memory @classmethod def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: if enforce_eager: logger.warning( "To see benefits of async output processing, enable CUDA graph. " "Since enforce-eager is enabled, async output processor cannot be used" ) return False return True @classmethod def log_warnings(cls) -> None: pass # ROCm-specific warnings can be added here @classmethod def get_current_memory_usage(cls, device: torch.device | None = None) -> float: torch.cuda.reset_peak_memory_stats(device) return float(torch.cuda.max_memory_allocated(device)) @classmethod def get_available_gpu_memory( cls, device_id: int = 0, distributed: bool = False, empty_cache: bool = True, cpu_group: Any = None, ) -> float: if empty_cache: torch.cuda.empty_cache() free_gpu_memory, _ = torch.cuda.mem_get_info(device_id) if distributed: import torch.distributed as dist tensor = torch.tensor(free_gpu_memory, dtype=torch.float32, device="cuda") dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group) free_gpu_memory = float(tensor.item()) return free_gpu_memory / (1 << 30) @classmethod def get_attn_backend_cls_str( cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype, ) -> str: if selected_backend == AttentionBackendEnum.TORCH_SDPA: logger.info("Using Torch SDPA backend.") return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" elif selected_backend in (AttentionBackendEnum.FA, None): pass elif selected_backend == AttentionBackendEnum.AITER: if dtype not in (torch.float16, torch.bfloat16): logger.warning( "AITer backend works best with fp16/bf16 inputs but got dtype=%s. " "Proceeding with AITer anyway.", dtype, ) logger.info("Using AITer backend on ROCm.") return "sglang.multimodal_gen.runtime.layers.attention.backends.aiter.AITerBackend" elif selected_backend == AttentionBackendEnum.AITER_SAGE: if dtype in (torch.float16, torch.bfloat16): logger.info("Using AITER Sage backend on ROCm.") return "sglang.multimodal_gen.runtime.layers.attention.backends.aiter_sage.AITERSageBackend" else: logger.warning( "AITER Sage backend only supports bf16/fp16 inputs but got dtype=%s.", dtype, ) elif selected_backend in ( AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN, ): raise ValueError( f"{selected_backend.name} is not supported on {cls.device_name}." ) elif selected_backend: raise ValueError( f"Invalid attention backend for {cls.device_name}: {selected_backend}" ) target_backend = AttentionBackendEnum.FA if dtype not in (torch.float16, torch.bfloat16): logger.info( "Cannot use FlashAttention backend for dtype other than " "torch.float16 or torch.bfloat16." ) target_backend = AttentionBackendEnum.TORCH_SDPA if target_backend == AttentionBackendEnum.FA: try: import flash_attn # noqa: F401 from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( # noqa: F401 FlashAttentionBackend, ) supported_sizes = FlashAttentionBackend.get_supported_head_sizes() if head_size not in supported_sizes: logger.info( "Cannot use FlashAttention-2 backend for head size %d.", head_size, ) target_backend = AttentionBackendEnum.TORCH_SDPA except ImportError: logger.info( "Cannot use FlashAttention backend because the " "flash_attn package is not found. " "Make sure that flash_attn was built and installed " "(on by default)." ) target_backend = AttentionBackendEnum.TORCH_SDPA if target_backend == AttentionBackendEnum.TORCH_SDPA: logger.info("Using Torch SDPA backend.") return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" logger.info("Using Flash Attention backend.") return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn.FlashAttentionBackend" @classmethod def get_device_communicator_cls(cls) -> str: return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator" # works for ROCm too @classmethod def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module: """Replace nn.GroupNorm with AITer GroupNorm for improved ROCm VAE performance.""" if not envs.SGLANG_USE_ROCM_VAE: return vae try: from aiter.ops.groupnorm import GroupNorm as AiterGroupNorm count = cls._replace_groupnorm(vae, AiterGroupNorm) if count > 0: logger.info( "Replaced %d nn.GroupNorm modules with AITer GroupNorm in VAE", count, ) except Exception: logger.warning( "Failed to apply AITer GroupNorm to VAE.", exc_info=True, ) return vae @staticmethod def _replace_groupnorm(module: torch.nn.Module, aiter_gn_cls: type) -> int: count = 0 for name, child in module.named_children(): if isinstance(child, torch.nn.GroupNorm) and child.affine: replacement = aiter_gn_cls( num_groups=child.num_groups, num_channels=child.num_channels, eps=child.eps, affine=True, device=child.weight.device, dtype=child.weight.dtype, ) replacement.weight = child.weight replacement.bias = child.bias setattr(module, name, replacement) count += 1 else: count += RocmPlatform._replace_groupnorm(child, aiter_gn_cls) return count @classmethod def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool: """ROCm performs better without DIT layerwise offload on Wan.""" return False ================================================ FILE: python/sglang/multimodal_gen/runtime/postprocess/__init__.py ================================================ # SPDX-License-Identifier: Apache-2.0 """Frame interpolation and upscaling support for SGLang diffusion pipelines.""" from sglang.multimodal_gen.runtime.postprocess.realesrgan_upscaler import ( ImageUpscaler, upscale_frames, ) from sglang.multimodal_gen.runtime.postprocess.rife_interpolator import ( FrameInterpolator, interpolate_video_frames, ) __all__ = [ "FrameInterpolator", "interpolate_video_frames", "ImageUpscaler", "upscale_frames", ] ================================================ FILE: python/sglang/multimodal_gen/runtime/postprocess/realesrgan_upscaler.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ Real-ESRGAN upscaling for SGLang diffusion pipelines. Real-ESRGAN model code is vendored and adapted from: - https://github.com/xinntao/Real-ESRGAN (BSD-3-Clause License) Copyright (c) 2021 xinntao The ImageUpscaler wrapper and integration code are original work. """ import math import os from typing import Optional import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # Default HuggingFace repo and filename for Real-ESRGAN weights _DEFAULT_REALESRGAN_HF_REPO = "ai-forever/Real-ESRGAN" _DEFAULT_REALESRGAN_FILENAME = "RealESRGAN_x4.pth" # Module-level cache: model_path -> UpscalerModel instance _MODEL_CACHE: dict[str, "UpscalerModel"] = {} # --------------------------------------------------------------------------- # Vendored Real-ESRGAN architecture code # (SRVGGNetCompact, ResidualDenseBlock, RRDB, RRDBNet) # --------------------------------------------------------------------------- class SRVGGNetCompact(nn.Module): """Compact VGG-style network for super resolution. Corresponds to ``realesr-animevideov3`` and ``realesr-general-x4v3``. Reference: xinntao/Real-ESRGAN (BSD-3-Clause). """ def __init__( self, num_in_ch: int = 3, num_out_ch: int = 3, num_feat: int = 64, num_conv: int = 16, upscale: int = 4, act_type: str = "prelu", ): super().__init__() self.num_in_ch = num_in_ch self.num_out_ch = num_out_ch self.num_feat = num_feat self.num_conv = num_conv self.upscale = upscale self.act_type = act_type self.body = nn.ModuleList() # first conv self.body.append(nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)) # first activation self.body.append(self._make_act(act_type, num_feat)) # body convs + activations for _ in range(num_conv): self.body.append(nn.Conv2d(num_feat, num_feat, 3, 1, 1)) self.body.append(self._make_act(act_type, num_feat)) # last conv: maps to out_ch * upscale^2 for pixel shuffle self.body.append(nn.Conv2d(num_feat, num_out_ch * upscale * upscale, 3, 1, 1)) self.upsampler = nn.PixelShuffle(upscale) @staticmethod def _make_act(act_type: str, num_feat: int) -> nn.Module: if act_type == "relu": return nn.ReLU(inplace=True) elif act_type == "prelu": return nn.PReLU(num_parameters=num_feat) elif act_type == "leakyrelu": return nn.LeakyReLU(negative_slope=0.1, inplace=True) else: raise ValueError(f"Unsupported activation type: {act_type}") def forward(self, x: torch.Tensor) -> torch.Tensor: out = x for layer in self.body: out = layer(out) out = self.upsampler(out) # residual addition with nearest upsampled input base = F.interpolate(x, scale_factor=self.upscale, mode="nearest") return out + base class ResidualDenseBlock(nn.Module): """Residual Dense Block used in RRDB (RealESRGAN_x4plus).""" def __init__(self, num_feat: int = 64, num_grow_ch: int = 32): super().__init__() self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1) self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1) self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1) self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1) self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1) self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) def forward(self, x: torch.Tensor) -> torch.Tensor: x1 = self.lrelu(self.conv1(x)) x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1))) x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1))) x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1))) x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1)) return x5 * 0.2 + x class RRDB(nn.Module): """Residual in Residual Dense Block.""" def __init__(self, num_feat: int, num_grow_ch: int = 32): super().__init__() self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch) self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch) self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch) def forward(self, x: torch.Tensor) -> torch.Tensor: out = self.rdb1(x) out = self.rdb2(out) out = self.rdb3(out) return out * 0.2 + x class RRDBNet(nn.Module): """RRDB network for RealESRGAN_x4plus (heavier, higher quality for photos).""" def __init__( self, num_in_ch: int = 3, num_out_ch: int = 3, scale: int = 4, num_feat: int = 64, num_block: int = 23, num_grow_ch: int = 32, ): super().__init__() self.scale = scale in_ch = num_in_ch if scale == 2: in_ch = num_in_ch * 4 elif scale == 1: in_ch = num_in_ch * 16 self.conv_first = nn.Conv2d(in_ch, num_feat, 3, 1, 1) self.body = nn.Sequential( *[RRDB(num_feat, num_grow_ch) for _ in range(num_block)] ) self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1) # upsample self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1) self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) def forward(self, x: torch.Tensor) -> torch.Tensor: if self.scale == 2: feat = F.pixel_unshuffle(x, 2) elif self.scale == 1: feat = F.pixel_unshuffle(x, 4) else: feat = x feat = self.conv_first(feat) body_feat = self.conv_body(self.body(feat)) feat = feat + body_feat feat = self.lrelu( self.conv_up1(F.interpolate(feat, scale_factor=2, mode="nearest")) ) feat = self.lrelu( self.conv_up2(F.interpolate(feat, scale_factor=2, mode="nearest")) ) return self.conv_last(self.lrelu(self.conv_hr(feat))) # --------------------------------------------------------------------------- # Architecture auto-detection # --------------------------------------------------------------------------- def _build_net_from_state_dict(state_dict: dict) -> nn.Module: """Detect architecture from checkpoint keys and return an unloaded network.""" if "conv_first.weight" in state_dict: # RRDBNet (e.g., RealESRGAN_x4plus) num_feat = state_dict["conv_first.weight"].shape[0] num_block = sum( 1 for k in state_dict if k.startswith("body.") and k.endswith(".rdb1.conv1.weight") ) num_grow_ch = state_dict["body.0.rdb1.conv1.weight"].shape[0] logger.info( "Detected RRDBNet: num_feat=%d, num_block=%d, num_grow_ch=%d", num_feat, num_block, num_grow_ch, ) return RRDBNet( num_in_ch=3, num_out_ch=3, scale=4, num_feat=num_feat, num_block=num_block, num_grow_ch=num_grow_ch, ) else: # SRVGGNetCompact (e.g., realesr-animevideov3) num_feat = state_dict["body.0.weight"].shape[0] # body layout: [first_conv, first_act, (conv, act)*num_conv, last_conv] # count 4-D weight tensors = first_conv + loop_convs + last_conv = num_conv + 2 conv_keys = sorted( [ k for k in state_dict if k.startswith("body.") and k.endswith(".weight") and state_dict[k].dim() == 4 ], key=lambda k: int(k.split(".")[1]), ) num_conv = len(conv_keys) - 2 # subtract first and last # upscale from last conv output channels: out_ch = num_out_ch * upscale^2 last_out_ch = state_dict[conv_keys[-1]].shape[0] upscale = int(math.sqrt(last_out_ch / 3)) logger.info( "Detected SRVGGNetCompact: num_feat=%d, num_conv=%d, upscale=%d", num_feat, num_conv, upscale, ) return SRVGGNetCompact( num_in_ch=3, num_out_ch=3, num_feat=num_feat, num_conv=num_conv, upscale=upscale, act_type="prelu", ) # --------------------------------------------------------------------------- # UpscalerModel # --------------------------------------------------------------------------- class UpscalerModel: """Wraps a Real-ESRGAN network, provides load() and upscale() API.""" def __init__(self, net: nn.Module, scale: int): self.net = net self.scale = scale # the model's native upscaling factor (e.g. 4) @property def device(self) -> torch.device: return next(self.net.parameters()).device def upscale(self, frame: np.ndarray, outscale: float | None = None) -> np.ndarray: """Upscale a single HWC uint8 frame → HWC uint8 frame. Args: frame: Input HWC uint8 numpy array. outscale: Desired final upscaling factor. If different from the model's native scale, a cheap resize is applied after the network output (same approach as the official Real-ESRGAN ``inference_realesrgan.py --outscale``). ``None`` means use the model's native scale as-is. """ h, w = frame.shape[:2] img = frame.astype(np.float32) / 255.0 img_t = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).to(self.device) with torch.no_grad(): out = self.net(img_t) # If the desired outscale differs from the model's native scale, # resize to (h * outscale, w * outscale). if outscale is not None and outscale != self.scale: target_h = int(h * outscale) target_w = int(w * outscale) out = F.interpolate( out, size=(target_h, target_w), mode="bicubic", align_corners=False ) out_np = out.squeeze(0).permute(1, 2, 0).clamp(0.0, 1.0).cpu().numpy() return (out_np * 255.0).astype(np.uint8) # --------------------------------------------------------------------------- # ImageUpscaler public class # --------------------------------------------------------------------------- class ImageUpscaler: """ Lazy-loaded Real-ESRGAN upscaler. Weights are downloaded and cached on first call to `.upscale()`. Supports both SRVGGNetCompact (lightweight, default) and RRDBNet (heavier). """ def __init__( self, model_path: Optional[str] = None, scale: int = 4, half_precision: bool = False, ): self._model_path = model_path self._scale = scale self._half_precision = half_precision def _ensure_model_loaded(self) -> UpscalerModel: """Download/load Real-ESRGAN weights, detect arch, and cache globally.""" model_path = self._model_path or _DEFAULT_REALESRGAN_HF_REPO # Resolve: local .pth pass-through, or HF repo → download single file resolved_path = _resolve_model_path(model_path) if resolved_path in _MODEL_CACHE: return _MODEL_CACHE[resolved_path] logger.info("Loading Real-ESRGAN weights from %s", resolved_path) try: state_dict = torch.load( resolved_path, map_location="cpu", weights_only=True ) except Exception as e: raise RuntimeError( f"Failed to load Real-ESRGAN checkpoint from '{resolved_path}'. " f"The file may be corrupted or not a valid PyTorch checkpoint. " f"Original error: {e}" ) from e # Some checkpoints wrap weights under a 'params' or 'params_ema' key if "params_ema" in state_dict: state_dict = state_dict["params_ema"] elif "params" in state_dict: state_dict = state_dict["params"] try: net = _build_net_from_state_dict(state_dict) net.load_state_dict(state_dict, strict=True) except (RuntimeError, KeyError) as e: raise RuntimeError( f"Real-ESRGAN weight file '{resolved_path}' is not compatible " f"with the supported architectures (SRVGGNetCompact / RRDBNet). " f"Please ensure you are using a valid Real-ESRGAN checkpoint. " f"Original error: {e}" ) from e net.eval() device = current_platform.get_local_torch_device() if self._half_precision: net = net.half() net = net.to(device) # Detect the model's native scale from network architecture native_scale = 4 # sensible default if hasattr(net, "upscale"): native_scale = net.upscale elif hasattr(net, "scale"): native_scale = net.scale model = UpscalerModel(net=net, scale=native_scale) _MODEL_CACHE[resolved_path] = model logger.info( "Real-ESRGAN model loaded on device: %s (native_scale=%dx, outscale=%s)", device, native_scale, f"{self._scale}x" if self._scale != native_scale else "native", ) return model def upscale(self, frames: list[np.ndarray]) -> list[np.ndarray]: """Upscale a list of HWC uint8 frames. Uses the model's native scale for super-resolution, then resizes to the desired ``outscale`` if it differs (cheap bicubic resize). """ if not frames: return frames model = self._ensure_model_loaded() outscale = self._scale if self._scale != model.scale else None return [model.upscale(frame, outscale=outscale) for frame in frames] # --------------------------------------------------------------------------- # HF download helper # --------------------------------------------------------------------------- def _resolve_model_path(model_path: str) -> str: """Return a local .pth file path. Accepts: - An existing local file path (pass-through). - A HuggingFace ``repo_id`` → downloads the default weight file (``RealESRGAN_x4.pth``). - A HuggingFace ``repo_id:filename`` → downloads *filename* from *repo_id*, allowing users to specify custom weight files hosted on HF. """ if os.path.isfile(model_path): return model_path # Parse optional "repo_id:filename" syntax; fall back to default filename. if ":" in model_path and not model_path.startswith("/"): repo_id, filename = model_path.split(":", 1) else: repo_id = model_path filename = _DEFAULT_REALESRGAN_FILENAME try: from huggingface_hub import hf_hub_download except ImportError as e: raise ImportError( "huggingface_hub is required to download Real-ESRGAN weights. " "Install it with: pip install huggingface_hub" ) from e logger.info( "Downloading Real-ESRGAN weights from HF repo %s (file: %s)", repo_id, filename, ) try: local_path = hf_hub_download( repo_id=repo_id, filename=filename, ) except Exception as e: raise FileNotFoundError( f"Failed to download Real-ESRGAN weights from HuggingFace repo " f"'{repo_id}' (file: '{filename}'). If you are using a custom " f"model, provide either a local .pth file path or use the " f"'repo_id:filename' format (e.g. 'my-org/my-esrgan:weights.pth'). " f"Original error: {e}" ) from e return local_path # --------------------------------------------------------------------------- # Module-level convenience function # --------------------------------------------------------------------------- def upscale_frames( frames: list[np.ndarray], model_path: Optional[str] = None, scale: int = 4, half_precision: bool = False, ) -> list[np.ndarray]: """ Convenience wrapper around ImageUpscaler. The model always runs at its native resolution (e.g. 4× for ``RealESRGAN_x4.pth``). If *scale* differs from the native factor, a cheap bicubic resize is applied after the network output – the same approach used by the official Real-ESRGAN ``--outscale`` flag. Args: frames: List of uint8 HWC numpy frames. model_path: Local .pth file, HuggingFace repo ID, or ``repo_id:filename`` for a custom weight file. None → default ``ai-forever/Real-ESRGAN`` with ``RealESRGAN_x4.pth``. scale: Desired final upscaling factor (e.g. 2, 3, 4). The 4× model is used internally; the output is resized to match *scale* when it differs. half_precision: Use fp16 inference (faster on supported GPUs). Returns: List of upscaled uint8 HWC numpy frames. """ upscaler = ImageUpscaler( model_path=model_path, scale=scale, half_precision=half_precision ) return upscaler.upscale(frames) ================================================ FILE: python/sglang/multimodal_gen/runtime/postprocess/rife_interpolator.py ================================================ # SPDX-License-Identifier: Apache-2.0 """ RIFE 4.22.lite frame interpolation for SGLang diffusion pipelines. RIFE model code is vendored and adapted from: - https://github.com/hzwer/ECCV2022-RIFE (MIT License) - https://github.com/hzwer/Practical-RIFE (MIT License) Copyright (c) 2021 Zhewei Huang The FrameInterpolator wrapper and integration code are original work. """ import os from typing import Optional import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # Default HuggingFace repo for RIFE 4.22.lite weights _DEFAULT_RIFE_HF_REPO = "elfgum/RIFE-4.22.lite" # Module-level cache: model_path -> Model instance _MODEL_CACHE: dict[str, "Model"] = {} # --------------------------------------------------------------------------- # Vendored RIFE 4.22.lite model code # (IFBlock, IFNet_HDv3 backbone, Model wrapper) # --------------------------------------------------------------------------- def warp(tenInput: torch.Tensor, tenFlow: torch.Tensor) -> torch.Tensor: """Warp tenInput by tenFlow using grid_sample.""" # Build base grid for the current size tenHorizontal = ( torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=tenFlow.device) .view(1, 1, 1, tenFlow.shape[3]) .expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1) ) tenVertical = ( torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=tenFlow.device) .view(1, 1, tenFlow.shape[2], 1) .expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3]) ) tenGrid = torch.cat([tenHorizontal, tenVertical], dim=1) tenFlow = torch.cat( [ tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0), ], dim=1, ) grid = (tenGrid + tenFlow).permute(0, 2, 3, 1) return F.grid_sample( input=tenInput, grid=grid, mode="bilinear", padding_mode="border", align_corners=True, ) def _conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): """Conv2d + LeakyReLU helper (matches RIFE 4.22 conv()).""" return nn.Sequential( nn.Conv2d( in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=True, ), nn.LeakyReLU(0.2, True), ) class ResConv(nn.Module): """Residual convolution block with learnable beta scaling (RIFE 4.22).""" def __init__(self, c: int, dilation: int = 1): super().__init__() self.conv = nn.Conv2d(c, c, 3, 1, dilation, dilation=dilation, groups=1) self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True) self.relu = nn.LeakyReLU(0.2, True) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.relu(self.conv(x) * self.beta + x) class IFBlock(nn.Module): """Single-scale optical flow + mask + feature block (RIFE 4.22).""" def __init__(self, in_planes: int, c: int = 64): super().__init__() self.conv0 = nn.Sequential( _conv(in_planes, c // 2, 3, 2, 1), _conv(c // 2, c, 3, 2, 1), ) self.convblock = nn.Sequential( ResConv(c), ResConv(c), ResConv(c), ResConv(c), ResConv(c), ResConv(c), ResConv(c), ResConv(c), ) self.lastconv = nn.Sequential( nn.ConvTranspose2d(c, 4 * 13, 4, 2, 1), nn.PixelShuffle(2), ) def forward( self, x: torch.Tensor, flow: Optional[torch.Tensor] = None, scale: float = 1.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: x = F.interpolate( x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False ) if flow is not None: flow = ( F.interpolate( flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False, ) * 1.0 / scale ) x = torch.cat((x, flow), 1) feat = self.conv0(x) feat = self.convblock(feat) tmp = self.lastconv(feat) tmp = F.interpolate( tmp, scale_factor=scale, mode="bilinear", align_corners=False ) flow = tmp[:, :4] * scale mask = tmp[:, 4:5] feat = tmp[:, 5:] return flow, mask, feat class Head(nn.Module): """Feature encoder producing 4-channel features at full resolution (RIFE 4.22).""" def __init__(self): super().__init__() self.cnn0 = nn.Conv2d(3, 16, 3, 2, 1) self.cnn1 = nn.Conv2d(16, 16, 3, 1, 1) self.cnn2 = nn.Conv2d(16, 16, 3, 1, 1) self.cnn3 = nn.ConvTranspose2d(16, 4, 4, 2, 1) self.relu = nn.LeakyReLU(0.2, True) def forward(self, x: torch.Tensor) -> torch.Tensor: x0 = self.cnn0(x) x = self.relu(x0) x1 = self.cnn1(x) x = self.relu(x1) x2 = self.cnn2(x) x = self.relu(x2) x3 = self.cnn3(x) return x3 class IFNet(nn.Module): """4-scale IFNet optical flow network (RIFE 4.22 backbone).""" def __init__(self): super().__init__() self.block0 = IFBlock(7 + 8, c=192) self.block1 = IFBlock(8 + 4 + 8 + 8, c=128) self.block2 = IFBlock(8 + 4 + 8 + 8, c=64) self.block3 = IFBlock(8 + 4 + 8 + 8, c=32) self.encode = Head() def forward( self, x: torch.Tensor, timestep: float = 0.5, scale_list: Optional[list] = None, ) -> tuple[list, torch.Tensor, list]: if scale_list is None: scale_list = [8, 4, 2, 1] channel = x.shape[1] // 2 img0 = x[:, :channel] img1 = x[:, channel:] if not torch.is_tensor(timestep): timestep = (x[:, :1].clone() * 0 + 1) * timestep else: timestep = timestep.repeat(1, 1, img0.shape[2], img0.shape[3]) f0 = self.encode(img0[:, :3]) f1 = self.encode(img1[:, :3]) flow_list = [] merged = [] mask_list = [] warped_img0 = img0 warped_img1 = img1 flow = None mask = None block = [self.block0, self.block1, self.block2, self.block3] for i in range(4): if flow is None: flow, mask, feat = block[i]( torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1), None, scale=scale_list[i], ) else: wf0 = warp(f0, flow[:, :2]) wf1 = warp(f1, flow[:, 2:4]) fd, m0, feat = block[i]( torch.cat( ( warped_img0[:, :3], warped_img1[:, :3], wf0, wf1, timestep, mask, feat, ), 1, ), flow, scale=scale_list[i], ) mask = m0 flow = flow + fd mask_list.append(mask) flow_list.append(flow) warped_img0 = warp(img0, flow[:, :2]) warped_img1 = warp(img1, flow[:, 2:4]) merged.append((warped_img0, warped_img1)) mask = torch.sigmoid(mask) merged[3] = warped_img0 * mask + warped_img1 * (1 - mask) return flow_list, mask_list[3], merged class Model: """Wraps IFNet, provides load_model() and inference() API.""" def __init__(self): self.flownet = IFNet() self.device_type: str = "cpu" def eval(self) -> "Model": self.flownet.eval() return self def device(self) -> torch.device: return next(self.flownet.parameters()).device def load_model(self, path: str, strip_module_prefix: bool = True) -> None: """Load weights from {path}/flownet.pkl. Args: path: Directory containing ``flownet.pkl``. strip_module_prefix: If True, strip the ``module.`` prefix that ``DataParallel`` / ``DistributedDataParallel`` adds to keys. """ flownet_path = os.path.join(path, "flownet.pkl") if not os.path.isfile(flownet_path): raise FileNotFoundError( f"RIFE weight file not found: {flownet_path}\n" "Expected layout: /flownet.pkl" ) def convert(param): if strip_module_prefix: return { k.replace("module.", ""): v for k, v in param.items() if "module." in k } else: return {k: v for k, v in param.items() if "module." not in k} state = torch.load(flownet_path, map_location="cpu", weights_only=False) self.flownet.load_state_dict(convert(state), strict=False) logger.info("Loaded RIFE weights from %s", flownet_path) def inference( self, img0: torch.Tensor, img1: torch.Tensor, scale: float = 1.0, timestep: float = 0.5, ) -> torch.Tensor: """Interpolate a single intermediate frame between img0 and img1.""" n, c, h, w = img0.shape # Pad to multiples of 32 so that RIFE's downsample/upsample round-trips # preserve spatial dimensions exactly. ph = ((h - 1) // 32 + 1) * 32 pw = ((w - 1) // 32 + 1) * 32 pad = (0, pw - w, 0, ph - h) img0 = F.pad(img0, pad) img1 = F.pad(img1, pad) imgs = torch.cat((img0, img1), 1) scale_list = [8 / scale, 4 / scale, 2 / scale, 1 / scale] with torch.no_grad(): flow_list, mask, merged = self.flownet( imgs, timestep=timestep, scale_list=scale_list, ) # Crop back to original resolution return merged[3][:, :, :h, :w] # --------------------------------------------------------------------------- # FrameInterpolator public class # --------------------------------------------------------------------------- class FrameInterpolator: """ Lazy-loaded RIFE 4.22.lite frame interpolator. Weights are loaded on first call to `.interpolate()` and cached globally per model_path to avoid reloading across requests. """ def __init__(self, model_path: Optional[str] = None): self._model_path = model_path self._resolved_path: Optional[str] = None def _ensure_model_loaded(self) -> Model: """Load RIFE model weights. Accepts a local directory **or** a HuggingFace repo ID. When *None* (the default) the weights are downloaded (and cached) automatically from ``elfgum/RIFE-4.22.lite`` via ``maybe_download_model()``. """ from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( maybe_download_model, ) model_path = self._model_path or _DEFAULT_RIFE_HF_REPO # Resolve: local path pass-through, HF repo ID → download & cache model_path = maybe_download_model(model_path) self._resolved_path = model_path if model_path in _MODEL_CACHE: return _MODEL_CACHE[model_path] device = current_platform.get_local_torch_device() model = Model() model.load_model(model_path, strip_module_prefix=True) model.eval() model.flownet = model.flownet.to(device) _MODEL_CACHE[model_path] = model logger.info("RIFE model loaded on device: %s", device) return model @staticmethod def _frame_to_tensor(frame: np.ndarray, device: torch.device) -> torch.Tensor: """Convert uint8 HWC numpy frame to float32 CHW tensor on device.""" t = torch.from_numpy(frame).permute(2, 0, 1).unsqueeze(0).float() / 255.0 return t.to(device) @staticmethod def _tensor_to_frame(t: torch.Tensor) -> np.ndarray: """Convert float32 CHW tensor (batch=1) to uint8 HWC numpy frame.""" arr = t.squeeze(0).permute(1, 2, 0).clamp(0.0, 1.0).cpu().numpy() return (arr * 255.0).astype(np.uint8) def _make_inference( self, model: Model, I0: torch.Tensor, I1: torch.Tensor, n: int, scale: float ) -> list[torch.Tensor]: """ Recursively generate n-1 intermediate frames between I0 and I1. Returns a list of intermediate frame tensors (not including I0 or I1). """ if n == 1: return [model.inference(I0, I1, scale=scale)] mid = model.inference(I0, I1, scale=scale) return ( self._make_inference(model, I0, mid, n // 2, scale) + [mid] + self._make_inference(model, mid, I1, n // 2, scale) ) def interpolate( self, frames: list[np.ndarray], exp: int = 1, scale: float = 1.0, ) -> tuple[list[np.ndarray], int]: """ Interpolate frames using RIFE. Args: frames: List of uint8 numpy arrays with shape [H, W, 3]. exp: Exponent for interpolation factor. 1 → 2×, 2 → 4×. scale: RIFE inference scale. Use 0.5 for high-resolution inputs. Returns: (interpolated_frames, multiplier) where multiplier = 2**exp. """ if len(frames) < 2: logger.warning( "Frame interpolation requires at least 2 frames; returning input unchanged." ) return frames, 1 model = self._ensure_model_loaded() device = model.device() n_intermediate = 2**exp // 2 # intermediates per adjacent pair result: list[np.ndarray] = [] for i in range(len(frames) - 1): I0 = self._frame_to_tensor(frames[i], device) I1 = self._frame_to_tensor(frames[i + 1], device) intermediate_tensors = self._make_inference( model, I0, I1, n_intermediate, scale ) result.append(frames[i]) for t in intermediate_tensors: result.append(self._tensor_to_frame(t)) result.append(frames[-1]) multiplier = 2**exp return result, multiplier # --------------------------------------------------------------------------- # Module-level convenience function # --------------------------------------------------------------------------- def interpolate_video_frames( frames: list[np.ndarray], exp: int = 1, scale: float = 1.0, model_path: Optional[str] = None, ) -> tuple[list[np.ndarray], int]: """ Convenience wrapper around FrameInterpolator. Args: frames: List of uint8 HWC numpy frames. exp: Interpolation exponent (1=2×, 2=4×). scale: RIFE inference scale (default 1.0; use 0.5 for high-res). model_path: Local directory or HuggingFace repo ID containing ``flownet.pkl``. *None* → default ``elfgum/RIFE-4.22.lite``. Returns: (interpolated_frames, multiplier) """ interpolator = FrameInterpolator(model_path=model_path) return interpolator.interpolate(frames, exp=exp, scale=scale) ================================================ FILE: python/sglang/multimodal_gen/runtime/scheduler_client.py ================================================ import pickle from typing import Any import zmq import zmq.asyncio from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) async def run_zeromq_broker(server_args: ServerArgs): """ This function runs as a background task in the FastAPI process. It listens for TCP requests from offline clients (e.g., DiffGenerator). """ ctx = zmq.asyncio.Context() # This is the REP socket that listens for requests from DiffGenerator socket = ctx.socket(zmq.REP) broker_endpoint = f"tcp://*:{server_args.broker_port}" socket.bind(broker_endpoint) logger.info(f"ZMQ Broker is listening for offline jobs on {broker_endpoint}") while True: try: # 1. Receive a request from an offline client payload = await socket.recv() request_batch = pickle.loads(payload) logger.info("Broker received an offline job from a client.") # 2. Forward the request to the main Scheduler via the shared client response_batch = await async_scheduler_client.forward(request_batch) # 3. Send the Scheduler's reply back to the offline client await socket.send(pickle.dumps(response_batch)) except Exception as e: logger.error(f"Error in ZMQ Broker: {e}", exc_info=True) # A reply must be sent to prevent the client from hanging try: await socket.send(pickle.dumps({"status": "error", "message": str(e)})) except Exception: pass class SchedulerClient: """ A synchronous, singleton client for communicating with the Scheduler service. Designed for use in DiffGenerator, where synchronous usage is preferred """ def __init__(self): self.context = None self.scheduler_socket = None self.server_args = None def initialize(self, server_args: ServerArgs): if self.context is not None and not self.context.closed: logger.warning("SchedulerClient is already initialized. Re-initializing.") self.close() self.server_args = server_args self.context = zmq.Context() self.scheduler_socket = self.context.socket(zmq.REQ) # Set socket options for the main communication socket self.scheduler_socket.setsockopt(zmq.LINGER, 0) # 100 minute timeout for generation self.scheduler_socket.setsockopt(zmq.RCVTIMEO, 6000000) scheduler_endpoint = self.server_args.scheduler_endpoint self.scheduler_socket.connect(scheduler_endpoint) logger.debug( f"SchedulerClient connected to backend scheduler at {scheduler_endpoint}" ) def forward(self, batch: Any) -> Any: """Sends a batch or request to the scheduler and waits for the response.""" try: self.scheduler_socket.send_pyobj(batch) output_batch = self.scheduler_socket.recv_pyobj() return output_batch except zmq.error.Again: logger.error("Timeout waiting for response from scheduler.") raise TimeoutError("Scheduler did not respond in time.") def ping(self) -> bool: """ Checks if the scheduler server is alive using a temporary socket. """ if self.context is None or self.context.closed: logger.error("Cannot ping: client is not initialized.") return False ping_socket = self.context.socket(zmq.REQ) ping_socket.setsockopt(zmq.LINGER, 0) ping_socket.setsockopt(zmq.RCVTIMEO, 2000) # 2-second timeout for pings endpoint = self.server_args.scheduler_endpoint try: ping_socket.connect(endpoint) ping_socket.send_pyobj({"method": "ping"}) ping_socket.recv_pyobj() return True except zmq.error.Again: return False finally: ping_socket.close() def close(self): """Closes the socket and terminates the context.""" if self.scheduler_socket: self.scheduler_socket.close() self.scheduler_socket = None if self.context: self.context.term() self.context = None class AsyncSchedulerClient: """ An asynchronous, singleton client for communicating with the Scheduler service. Designed for use in asynchronous environments like FastAPI entrypoints. To support high concurrency, it creates a new REQ socket for each request rather than sharing a single one (which would cause ZMQ state errors). """ def __init__(self): self.context = None self.server_args = None def initialize(self, server_args: ServerArgs): if self.context is not None and not self.context.closed: logger.warning( "AsyncSchedulerClient is already initialized. Re-initializing." ) self.close() self.server_args = server_args self.context = zmq.asyncio.Context() logger.debug("AsyncSchedulerClient initialized with zmq.asyncio.Context") async def forward(self, batch: Any) -> Any: """Sends a batch or request to the scheduler and waits for the response.""" if self.context is None: raise RuntimeError( "AsyncSchedulerClient is not initialized. Call initialize() first." ) # Create a temporary REQ socket for this request to allow concurrency socket = self.context.socket(zmq.REQ) socket.setsockopt(zmq.LINGER, 0) # 100 minute timeout socket.setsockopt(zmq.RCVTIMEO, 6000000) endpoint = self.server_args.scheduler_endpoint socket.connect(endpoint) try: await socket.send(pickle.dumps(batch)) payload = await socket.recv() return pickle.loads(payload) except zmq.error.Again: logger.error("Timeout waiting for response from scheduler.") raise TimeoutError("Scheduler did not respond in time.") finally: socket.close() async def ping(self) -> bool: """ Checks if the scheduler server is alive using a temporary socket. """ if self.context is None or self.context.closed: logger.error("Cannot ping: client is not initialized.") return False ping_socket = self.context.socket(zmq.REQ) ping_socket.setsockopt(zmq.LINGER, 0) ping_socket.setsockopt(zmq.RCVTIMEO, 2000) endpoint = self.server_args.scheduler_endpoint try: ping_socket.connect(endpoint) await ping_socket.send(pickle.dumps({"method": "ping"})) await ping_socket.recv() return True except zmq.error.Again: return False finally: ping_socket.close() def close(self): """Closes the socket and terminates the context.""" if self.context: self.context.term() self.context = None # Singleton instances for easy access async_scheduler_client = AsyncSchedulerClient() sync_scheduler_client = SchedulerClient() ================================================ FILE: python/sglang/multimodal_gen/runtime/server_args.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py """The arguments of sglang-diffusion Inference.""" import argparse import dataclasses import json import math import os import random import sys import tempfile from dataclasses import field from enum import Enum from typing import Any, Optional import addict import yaml from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.models.encoders import T5Config from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig from sglang.multimodal_gen.configs.quantization import NunchakuSVDQuantArgs from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( NunchakuConfig, ) from sglang.multimodal_gen.runtime.loader.utils import BYTES_PER_GB from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.utils.common import ( is_port_available, is_valid_ipv6_address, ) from sglang.multimodal_gen.runtime.utils.logging_utils import ( _sanitize_for_logging, configure_logger, init_logger, ) from sglang.multimodal_gen.utils import ( FlexibleArgumentParser, StoreBoolean, expand_path_fields, expand_path_kwargs, ) logger = init_logger(__name__) class Backend(str, Enum): """ Enumeration for different model backends. - AUTO: Automatically select backend (prefer sglang native, fallback to diffusers) - SGLANG: Use sglang's native optimized implementation - DIFFUSERS: Use vanilla diffusers pipeline (supports all diffusers models) """ AUTO = "auto" SGLANG = "sglang" DIFFUSERS = "diffusers" @classmethod def from_string(cls, value: str) -> "Backend": """Convert string to Backend enum.""" try: return cls(value.lower()) except ValueError: raise ValueError( f"Invalid backend: {value}. Must be one of: {', '.join([m.value for m in cls])}" ) from None @classmethod def choices(cls) -> list[str]: """Get all available choices as strings for argparse.""" return [backend.value for backend in cls] @dataclasses.dataclass class ServerArgs: # Model and path configuration (for convenience) model_path: str # explicit model ID override (e.g. "Qwen-Image") model_id: str | None = None # Model backend (sglang native or diffusers) backend: Backend = Backend.AUTO # Attention attention_backend: str = None attention_backend_config: addict.Dict | None = None cache_dit_config: str | dict[str, Any] | None = ( None # cache-dit config for diffusers ) # Distributed executor backend nccl_port: Optional[int] = None # HuggingFace specific parameters trust_remote_code: bool = False revision: str | None = None # Parallelism num_gpus: int = 1 tp_size: Optional[int] = None sp_degree: Optional[int] = None # sequence parallelism ulysses_degree: Optional[int] = None ring_degree: Optional[int] = None # data parallelism # number of data parallelism groups dp_size: int = 1 # number of gpu in a dp group dp_degree: int = 1 # cfg parallel enable_cfg_parallel: bool = False hsdp_replicate_dim: int = 1 hsdp_shard_dim: Optional[int] = None dist_timeout: int | None = 3600 # 1 hour pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False) # Pipeline override pipeline_class_name: str | None = ( None # Override pipeline class from model_index.json ) # LoRA parameters # (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated. lora_path: str | None = None lora_nickname: str = "default" # for swapping adapters in the pipeline lora_scale: float = 1.0 # LoRA scale for merging (e.g., 0.125 for Hyper-SD) # Component path overrides (key = model_index.json component name, value = path) component_paths: dict[str, str] = field(default_factory=dict) # path to pre-quantized transformer weights (single .safetensors or directory). transformer_weights_path: str | None = None # can restrict layers to adapt, e.g. ["q_proj"] # Will adapt only q, k, v, o by default. lora_target_modules: list[str] | None = None # CPU offload parameters dit_cpu_offload: bool | None = None dit_layerwise_offload: bool | None = None dit_offload_prefetch_size: float = 0.0 text_encoder_cpu_offload: bool | None = None image_encoder_cpu_offload: bool | None = None vae_cpu_offload: bool | None = None use_fsdp_inference: bool = False pin_cpu_memory: bool = True # ComfyUI integration comfyui_mode: bool = False # Compilation enable_torch_compile: bool = False # warmup warmup: bool = False warmup_resolutions: list[str] = None warmup_steps: int = 1 disable_autocast: bool | None = None # Quantization / Nunchaku SVDQuant configuration nunchaku_config: NunchakuSVDQuantArgs | NunchakuConfig | None = field( default_factory=NunchakuSVDQuantArgs, repr=False ) # Master port for distributed inference # TODO: do not hard code master_port: int | None = None # http server endpoint config host: str | None = "127.0.0.1" port: int | None = 30000 # TODO: webui and their endpoint, check if webui_port is available. webui: bool = False webui_port: int | None = 12312 scheduler_port: int = 5555 output_path: str | None = "outputs/" input_save_path: str | None = "inputs/uploads" # Prompt text file for batch processing prompt_file_path: str | None = None # model paths for correct deallocation model_paths: dict[str, str] = field(default_factory=dict) model_loaded: dict[str, bool] = field( default_factory=lambda: { "transformer": True, "vae": True, "video_vae": True, "audio_vae": True, "video_dit": True, "audio_dit": True, "dual_tower_bridge": True, } ) # # DMD parameters # dmd_denoising_steps: List[int] | None = field(default=None) # MoE parameters used by Wan2.2 boundary_ratio: float | None = None # Logging log_level: str = "info" @property def broker_port(self) -> int: return self.port + 1 @property def is_local_mode(self) -> bool: """ If no server is running when a generation task begins, 'local_mode' will be enabled: a dedicated server will be launched """ return self.host is None or self.port is None def _adjust_path(self): expand_path_fields(self) self._adjust_save_paths() def _adjust_parameters(self): """set defaults and normalize values.""" self._adjust_offload() self._adjust_path() self._adjust_quant_config() self._adjust_warmup() self._adjust_network_ports() # adjust parallelism before attention backend self._adjust_parallelism() self._adjust_attention_backend() self._adjust_platform_specific() self._adjust_autocast() self.adjust_pipeline_config() def _validate_parameters(self): """check consistency and raise errors for invalid configs""" self._validate_pipeline() self._validate_offload() self._validate_parallelism() self._validate_cfg_parallel() def _adjust_save_paths(self): """Normalize empty-string save paths to None (disabled).""" if self.output_path is not None and self.output_path.strip() == "": self.output_path = None if self.input_save_path is not None and self.input_save_path.strip() == "": self.input_save_path = None def _adjust_quant_config(self): """validate and adjust""" # nunchaku ncfg = self.nunchaku_config if ncfg is None or isinstance(ncfg, NunchakuConfig): return ncfg.validate() # propagate the path to server_args if ncfg.transformer_weights_path: self.transformer_weights_path = ncfg.transformer_weights_path if not ncfg.enable_svdquant or not ncfg.transformer_weights_path: self.nunchaku_config = None else: self.nunchaku_config = NunchakuConfig( precision=self.nunchaku_config.quantization_precision, rank=self.nunchaku_config.quantization_rank, act_unsigned=self.nunchaku_config.quantization_act_unsigned, transformer_weights_path=self.nunchaku_config.transformer_weights_path, ) def adjust_pipeline_config(self): # enable parallel folding when SP is enabled if self.tp_size != 1 or self.sp_degree <= 1: return enabled = False for text_encoder_config in self.pipeline_config.text_encoder_configs: if isinstance(text_encoder_config, T5Config): text_encoder_config.parallel_folding = True enabled = True text_encoder_config.parallel_folding_mode = "sp" if enabled: logger.info( "Enabled T5 text encoder parallel folding (mode=sp) for %s (tp_size=%s, sp_degree=%s).", self.__class__.__name__, self.tp_size, self.sp_degree, ) def _adjust_offload(self): # TODO: to be handled by each platform if current_platform.get_device_total_memory() / BYTES_PER_GB < 30: logger.info("Enabling all offloading for GPU with low device memory") if self.dit_cpu_offload is None: self.dit_cpu_offload = True if self.text_encoder_cpu_offload is None: self.text_encoder_cpu_offload = True if self.image_encoder_cpu_offload is None: self.image_encoder_cpu_offload = True if self.vae_cpu_offload is None: self.vae_cpu_offload = True elif self.pipeline_config.task_type.is_image_gen(): logger.info( "Disabling some offloading (except dit, text_encoder) for image generation model" ) if self.dit_cpu_offload is None: self.dit_cpu_offload = True if self.text_encoder_cpu_offload is None: self.text_encoder_cpu_offload = True if self.image_encoder_cpu_offload is None: self.image_encoder_cpu_offload = False if self.vae_cpu_offload is None: self.vae_cpu_offload = False else: if self.dit_cpu_offload is None: self.dit_cpu_offload = True if self.text_encoder_cpu_offload is None: self.text_encoder_cpu_offload = True if self.image_encoder_cpu_offload is None: self.image_encoder_cpu_offload = True if self.vae_cpu_offload is None: self.vae_cpu_offload = True def _adjust_attention_backend(self): if self.attention_backend in ["fa3", "fa4"]: self.attention_backend = "fa" # attention_backend_config if self.attention_backend_config is None: self.attention_backend_config = addict.Dict() elif isinstance(self.attention_backend_config, str): self.attention_backend_config = addict.Dict( self._parse_attention_backend_config(self.attention_backend_config) ) if self.ring_degree > 1: if self.attention_backend is not None and self.attention_backend not in ( "fa", "sage_attn", ): raise ValueError( "Ring Attention is only supported for flash attention or sage attention backend for now" ) if self.attention_backend is None: self.attention_backend = "fa" logger.info( "Ring Attention is currently only supported for flash attention or sage attention; " "attention_backend has been automatically set to flash attention" ) if self.attention_backend is None and self.backend != Backend.DIFFUSERS: self._set_default_attention_backend() def _adjust_warmup(self): if self.warmup_resolutions is not None: self.warmup = True if self.warmup: logger.info( "Warmup enabled, the launch time is expected to be longer than usual" ) def _adjust_network_ports(self): self.port = self.settle_port(self.port) initial_scheduler_port = self.scheduler_port + ( random.randint(0, 100) if self.scheduler_port == 5555 else 0 ) self.scheduler_port = self.settle_port(initial_scheduler_port) initial_master_port = ( self.master_port if self.master_port is not None else (30005 + random.randint(0, 100)) ) self.master_port = self.settle_port(initial_master_port, 37) def _adjust_parallelism(self): if self.tp_size is None: self.tp_size = 1 if self.hsdp_shard_dim is None: self.hsdp_shard_dim = self.num_gpus # adjust sp_degree: allocate all remaining GPUs after TP and DP if self.sp_degree is None: num_gpus_per_group = self.dp_size * self.tp_size if self.enable_cfg_parallel: num_gpus_per_group *= 2 if self.num_gpus % num_gpus_per_group == 0: self.sp_degree = self.num_gpus // num_gpus_per_group else: # Will be validated later self.sp_degree = 1 if ( self.ulysses_degree is None and self.ring_degree is None and self.sp_degree != 1 ): self.ulysses_degree = self.sp_degree logger.info( f"Automatically set ulysses_degree=sp_degree={self.ulysses_degree} for best performance" ) if self.ulysses_degree is None: self.ulysses_degree = 1 logger.debug( f"Ulysses degree not set, using default value {self.ulysses_degree}" ) if self.ring_degree is None: self.ring_degree = 1 logger.debug(f"Ring degree not set, using default value {self.ring_degree}") def _adjust_platform_specific(self): if current_platform.is_mps(): self.use_fsdp_inference = False self.dit_layerwise_offload = False # automatically enable dit_layerwise_offload for Wan/MOVA models if appropriate if not envs.SGLANG_CACHE_DIT_ENABLED: pipeline_name_lower = self.pipeline_config.__class__.__name__.lower() if ( ("wan" in pipeline_name_lower or "mova" in pipeline_name_lower) and self.dit_layerwise_offload is None and current_platform.enable_dit_layerwise_offload_for_wan_by_default() ): logger.info( f"Automatically enable dit_layerwise_offload for {self.pipeline_config.__class__.__name__} " "for low memory and performance balance" ) self.dit_layerwise_offload = True def _adjust_autocast(self): if self.disable_autocast is None: self.disable_autocast = not self.pipeline_config.enable_autocast def _parse_attention_backend_config(self, config_str: str) -> dict[str, Any]: """parse attention backend config from string.""" if not config_str: return {} # 1. treat as file path if os.path.exists(config_str): if config_str.endswith((".yaml", ".yml")): with open(config_str, "r") as f: return yaml.safe_load(f) elif config_str.endswith(".json"): with open(config_str, "r") as f: return json.load(f) # 2. treat as JSON string try: return json.loads(config_str) except json.JSONDecodeError: pass # 3. treat as k=v pairs (simple implementation). e.g., "sparsity=0.5,enable_x=true" try: config = {} pairs = config_str.split(",") for pair in pairs: k, v = pair.split("=", 1) k = k.strip() v = v.strip() if v.lower() == "true": v = True elif v.lower() == "false": v = False elif v.replace(".", "", 1).isdigit(): v = float(v) if "." in v else int(v) config[k] = v return config except Exception: raise ValueError(f"Could not parse attention backend config: {config_str}") def __post_init__(self): # configure logger before use configure_logger(server_args=self) # 1. adjust parameters self._adjust_parameters() # 2. Validate parameters self._validate_parameters() # log clean server_args try: safe_args = _sanitize_for_logging(self, key_hint="server_args") logger.info("server_args: %s", json.dumps(safe_args, ensure_ascii=False)) except Exception: # Fallback to default repr if sanitization fails logger.info(f"server_args: {self}") @staticmethod def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: # Model and path configuration parser.add_argument( "--model-path", type=str, help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.", ) parser.add_argument( "--model-id", type=str, default=ServerArgs.model_id, help=( "Override the model ID used for config resolution. " "Useful when --model-path is a local directory whose name does not match " "any registered HF repo name. Should be the repo name portion of the HF ID " "(e.g. 'Qwen-Image' for 'Qwen/Qwen-Image')." ), ) # attention parser.add_argument( "--attention-backend", type=str, default=None, help=( "The attention backend to use. For SGLang-native pipelines, use " "values like fa, torch_sdpa, sage_attn, etc. For diffusers pipelines, " "use diffusers attention backend names such as flash, _flash_3_hub, " "sage, or xformers." ), ) parser.add_argument( "--attention-backend-config", type=str, default=None, help="Configuration for the attention backend. Can be a JSON string, a path to a JSON/YAML file, or key=value pairs.", ) parser.add_argument( "--cache-dit-config", type=str, default=ServerArgs.cache_dit_config, help="Path to a Cache-DiT YAML/JSON config. Enables cache-dit for diffusers backend.", ) # HuggingFace specific parameters parser.add_argument( "--trust-remote-code", action=StoreBoolean, default=ServerArgs.trust_remote_code, help="Trust remote code when loading HuggingFace models", ) parser.add_argument( "--revision", type=str, default=ServerArgs.revision, help="The specific model version to use (can be a branch name, tag name, or commit id)", ) # Parallelism parser.add_argument( "--num-gpus", type=int, default=ServerArgs.num_gpus, help="The number of GPUs to use.", ) parser.add_argument( "--tp-size", type=int, default=None, help="The tensor parallelism size. Defaults to 1 if not specified.", ) parser.add_argument( "--sp-degree", type=int, default=None, help="The sequence parallelism size. If not specified, will use all remaining GPUs after accounting for TP and DP.", ) parser.add_argument( "--ulysses-degree", type=int, default=ServerArgs.ulysses_degree, help="Ulysses sequence parallel degree. Used in attention layer.", ) parser.add_argument( "--ring-degree", type=int, default=ServerArgs.ring_degree, help="Ring sequence parallel degree. Used in attention layer.", ) parser.add_argument( "--enable-cfg-parallel", action="store_true", default=ServerArgs.enable_cfg_parallel, help="Enable cfg parallel.", ) parser.add_argument( "--data-parallel-size", "--dp-size", "--dp", type=int, default=ServerArgs.dp_size, help="The data parallelism size.", ) parser.add_argument( "--hsdp-replicate-dim", type=int, default=ServerArgs.hsdp_replicate_dim, help="The data parallelism size.", ) parser.add_argument( "--hsdp-shard-dim", type=int, default=None, help="The data parallelism shards. Defaults to num_gpus if not specified.", ) parser.add_argument( "--dist-timeout", type=int, default=ServerArgs.dist_timeout, help="Timeout for torch.distributed operations in seconds. " "Increase this value if you encounter 'Connection closed by peer' errors after the service is idle. ", ) # Prompt text file for batch processing parser.add_argument( "--prompt-file-path", type=str, default=ServerArgs.prompt_file_path, help="Path to a text file containing prompts (one per line) for batch processing", ) parser.add_argument( "--mask-strategy-file-path", type=str, help="Path to mask strategy JSON file for STA", ) parser.add_argument( "--enable-torch-compile", action=StoreBoolean, default=ServerArgs.enable_torch_compile, help="Use torch.compile to speed up DiT inference." + "However, will likely cause precision drifts. See (https://github.com/pytorch/pytorch/issues/145213)", ) # warmup parser.add_argument( "--warmup", action=StoreBoolean, default=ServerArgs.warmup, help="Perform some warmup after server starts (if `--warmup-resolutions` is specified) or before processing the first request (if `--warmup-resolutions` is not specified)." "Recommended to enable when benchmarking to ensure fair comparison and best performance." "When enabled with `--warmup-resolutions` unspecified, look for the line ending with `(with warmup excluded)` for actual processing time.", ) parser.add_argument( "--warmup-resolutions", type=str, nargs="+", default=ServerArgs.warmup_resolutions, help="Specify resolutions for server to warmup. e.g., `--warmup-resolutions 256x256, 720x720`", ) parser.add_argument( "--warmup-steps", type=int, default=ServerArgs.warmup_steps, help="The number of warmup steps to perform for each resolution.", ) parser.add_argument( "--dit-cpu-offload", action=StoreBoolean, help="Use CPU offload for DiT inference. Enable if run out of memory with FSDP.", ) parser.add_argument( "--dit-layerwise-offload", action=StoreBoolean, default=ServerArgs.dit_layerwise_offload, help="Enable layerwise CPU offload with async H2D prefetch overlap for supported DiT models (e.g., Wan, MOVA). " "Cannot be used together with cache-dit (SGLANG_CACHE_DIT_ENABLED), dit_cpu_offload, or use_fsdp_inference.", ) parser.add_argument( "--dit-offload-prefetch-size", type=float, default=ServerArgs.dit_offload_prefetch_size, help="The size of prefetch for dit-layerwise-offload. If the value is between 0.0 and 1.0, it is treated as a ratio of the total number of layers. If the value is >= 1, it is treated as the absolute number of layers. 0.0 means prefetch 1 layer (lowest memory). Values above 0.5 might have peak memory close to no offload but worse performance.", ) parser.add_argument( "--use-fsdp-inference", action=StoreBoolean, help="Use FSDP for inference by sharding the model weights. Latency is very low due to prefetch--enable if run out of memory.", ) parser.add_argument( "--text-encoder-cpu-offload", action=StoreBoolean, help="Use CPU offload for text encoder. Enable if run out of memory.", ) parser.add_argument( "--image-encoder-cpu-offload", action=StoreBoolean, help="Use CPU offload for image encoder. Enable if run out of memory.", ) parser.add_argument( "--vae-cpu-offload", action=StoreBoolean, help="Use CPU offload for VAE. Enable if run out of memory.", ) parser.add_argument( "--pin-cpu-memory", action=StoreBoolean, help='Pin memory for CPU offload. Only added as a temp workaround if it throws "CUDA error: invalid argument". ' "Should be enabled in almost all cases", ) parser.add_argument( "--disable-autocast", action=StoreBoolean, help="Disable autocast for denoising loop and vae decoding in pipeline sampling", ) # Nunchaku SVDQuant quantization parameters NunchakuSVDQuantArgs.add_cli_args(parser) # Master port for distributed inference parser.add_argument( "--master-port", type=int, default=ServerArgs.master_port, help="Master port for distributed inference. If not set, a random free port will be used.", ) parser.add_argument( "--scheduler-port", type=int, default=ServerArgs.scheduler_port, help="Port for the scheduler server.", ) parser.add_argument( "--host", type=str, default=ServerArgs.host, help="Host for the HTTP API server.", ) parser.add_argument( "--port", type=int, default=ServerArgs.port, help="Port for the HTTP API server.", ) parser.add_argument( "--webui", action=StoreBoolean, default=ServerArgs.webui, help="Whether to use webui for better display", ) parser.add_argument( "--webui-port", type=int, default=ServerArgs.webui_port, help="Whether to use webui for better display", ) parser.add_argument( "--output-path", type=str, default=ServerArgs.output_path, help='Directory path to save generated images/videos. Set to "" to disable persistent saving.', ) parser.add_argument( "--input-save-path", type=str, default=ServerArgs.input_save_path, help='Directory path to save uploaded input images/videos. Set to "" to disable persistent saving.', ) # LoRA parser.add_argument( "--lora-path", type=str, default=ServerArgs.lora_path, help="The path to the LoRA adapter weights (can be local file path or HF hub id) to launch with", ) parser.add_argument( "--lora-nickname", type=str, default=ServerArgs.lora_nickname, help="The nickname for the LoRA adapter to launch with", ) parser.add_argument( "--lora-scale", type=float, default=ServerArgs.lora_scale, help="LoRA scale for merging (e.g., 0.125 for Hyper-SD). Same as lora_scale in Diffusers", ) # Add pipeline configuration arguments PipelineConfig.add_cli_args(parser) # Logging parser.add_argument( "--log-level", type=str, default=ServerArgs.log_level, help="The logging level of all loggers.", ) parser.add_argument( "--backend", type=str, choices=Backend.choices(), default=ServerArgs.backend.value, help="The model backend to use. 'auto' prefers sglang native and falls back to diffusers. " "'sglang' uses native optimized implementation. 'diffusers' uses vanilla diffusers pipeline.", ) return parser def url(self): host = self.host if not host or host == "0.0.0.0": host = "127.0.0.1" elif host == "::": host = "::1" if is_valid_ipv6_address(host): return f"http://[{host}]:{self.port}" else: return f"http://{host}:{self.port}" @property def scheduler_endpoint(self): """ Internal endpoint for scheduler. Prefers the configured host but normalizes localhost -> 127.0.0.1 to avoid ZMQ issues. """ scheduler_host = self.host if scheduler_host is None or scheduler_host == "localhost": scheduler_host = "127.0.0.1" return f"tcp://{scheduler_host}:{self.scheduler_port}" def settle_port( self, port: int, port_inc: int = 42, max_attempts: int = 100 ) -> int: """ Find an available port with retry logic. """ attempts = 0 original_port = port while attempts < max_attempts: if is_port_available(port): if attempts > 0: logger.info( f"Port {original_port} was unavailable, using port {port} instead" ) return port attempts += 1 if port < 60000: port += port_inc else: # Wrap around with randomization to avoid collision port = 5000 + random.randint(0, 1000) raise RuntimeError( f"Failed to find available port after {max_attempts} attempts " f"(started from port {original_port})" ) @staticmethod def _extract_component_paths( unknown_args: list[str], ) -> tuple[dict[str, str], list[str]]: """ Extract dynamic ``---path`` args from unrecognised CLI args. """ component_paths: dict[str, str] = {} remaining: list[str] = [] i = 0 while i < len(unknown_args): arg = unknown_args[i] key_part = arg.split("=", 1)[0] if "=" in arg else arg if key_part.startswith("--") and key_part.endswith("-path"): component = key_part[2:-5].replace("-", "_") if "=" in arg: component_paths[component] = arg.split("=", 1)[1] elif i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith( "-" ): i += 1 component_paths[component] = unknown_args[i] else: remaining.append(arg) i += 1 continue else: remaining.append(arg) i += 1 # canonicalize and validate for component, path in component_paths.items(): path = os.path.expanduser(path) component_paths[component] = path return component_paths, remaining @classmethod def from_cli_args( cls, args: argparse.Namespace, unknown_args: list[str] | None = None ) -> "ServerArgs": if unknown_args is None: unknown_args = [] # extract dynamic ---path from unknown args dynamic_paths, remaining = cls._extract_component_paths(unknown_args) if remaining: raise SystemExit(f"error: unrecognized arguments: {' '.join(remaining)}") provided_args = cls.get_provided_args(args, unknown_args) # Handle config file config_file = provided_args.get("config") if config_file: config_args = cls.load_config_file(config_file) provided_args = {**config_args, **provided_args} if dynamic_paths: existing = dict(provided_args.get("component_paths") or {}) existing.update(dynamic_paths) provided_args["component_paths"] = existing return cls.from_dict(provided_args) @classmethod def from_dict(cls, kwargs: dict[str, Any]) -> "ServerArgs": """Create a ServerArgs object from a dictionary.""" kwargs = expand_path_kwargs(dict(kwargs)) attrs = [attr.name for attr in dataclasses.fields(cls)] server_args_kwargs: dict[str, Any] = {} component_paths = dict(kwargs.get("component_paths") or {}) if component_paths: server_args_kwargs["component_paths"] = component_paths for attr in attrs: if attr == "pipeline_config": pipeline_config = PipelineConfig.from_kwargs(kwargs) logger.debug(f"Using PipelineConfig: {type(pipeline_config)}") server_args_kwargs["pipeline_config"] = pipeline_config elif attr == "nunchaku_config": nunchaku_config = NunchakuSVDQuantArgs.from_dict(kwargs) server_args_kwargs["nunchaku_config"] = nunchaku_config elif attr in kwargs: server_args_kwargs[attr] = kwargs[attr] return cls(**server_args_kwargs) @staticmethod def load_config_file(config_file: str) -> dict[str, Any]: """Load a config file.""" if config_file.endswith(".json"): with open(config_file, "r") as f: return json.load(f) elif config_file.endswith((".yaml", ".yml")): try: import yaml except ImportError: raise ImportError( "Please install PyYAML to use YAML config files. " "`pip install pyyaml`" ) with open(config_file, "r") as f: return yaml.safe_load(f) else: raise ValueError(f"Unsupported config file format: {config_file}") @classmethod def from_kwargs(cls, **kwargs: Any) -> "ServerArgs": # Convert backend string to enum if necessary if "backend" in kwargs and isinstance(kwargs["backend"], str): kwargs["backend"] = Backend.from_string(kwargs["backend"]) kwargs["pipeline_config"] = PipelineConfig.from_kwargs(kwargs) return cls(**kwargs) @staticmethod def get_provided_args( args: argparse.Namespace, unknown_args: list[str] ) -> dict[str, Any]: """Get the arguments provided by the user.""" provided_args = {} # We need to check against the raw command-line arguments to see what was # explicitly provided by the user, vs. what's a default value from argparse. raw_argv = sys.argv + unknown_args # Create a set of argument names that were present on the command line. # This handles both styles: '--arg=value' and '--arg value'. provided_arg_names = set() for arg in raw_argv: if arg.startswith("--"): # For '--arg=value', this gets 'arg'; for '--arg', this also gets 'arg'. arg_name = arg.split("=", 1)[0].replace("-", "_").lstrip("_") provided_arg_names.add(arg_name) # Populate provided_args if the argument from the namespace was on the command line. for k, v in vars(args).items(): if k in provided_arg_names: provided_args[k] = v return provided_args def _validate_pipeline(self): if self.pipeline_config is None: raise ValueError("pipeline_config is not set in ServerArgs") self.pipeline_config.check_pipeline_config() def _validate_offload(self): # validate dit_offload_prefetch_size if self.dit_offload_prefetch_size > 1 and ( isinstance(self.dit_offload_prefetch_size, float) and not self.dit_offload_prefetch_size.is_integer() ): self.dit_offload_prefetch_size = int( math.floor(self.dit_offload_prefetch_size) ) logger.info( f"Invalid --dit-offload-prefetch-size value passed, truncated to: {self.dit_offload_prefetch_size}" ) if 0.5 <= self.dit_offload_prefetch_size < 1.0: logger.info( "We do not recommend --dit-offload-prefetch-size to be between 0.5 and 1.0" ) # validate dit_layerwise_offload conflicts if self.dit_layerwise_offload: if self.dit_offload_prefetch_size < 0.0: raise ValueError("dit_offload_prefetch_size must be non-negative") if self.use_fsdp_inference: logger.warning( "dit_layerwise_offload is enabled, automatically disabling use_fsdp_inference." ) self.use_fsdp_inference = False if self.dit_cpu_offload is None: logger.warning( "dit_layerwise_offload is enabled, automatically disabling dit_cpu_offload." ) self.dit_cpu_offload = False if envs.SGLANG_CACHE_DIT_ENABLED: raise ValueError( "dit_layerwise_offload cannot be enabled together with cache-dit. " "cache-dit may reuse skipped blocks whose weights have been released by layerwise offload, " "causing shape mismatch errors. " "Please disable either --dit-layerwise-offload or SGLANG_CACHE_DIT_ENABLED." ) def _validate_parallelism(self): if self.sp_degree > self.num_gpus or self.num_gpus % self.sp_degree != 0: raise ValueError( f"num_gpus ({self.num_gpus}) must be >= and divisible by sp_degree ({self.sp_degree})" ) if ( self.hsdp_replicate_dim > self.num_gpus or self.num_gpus % self.hsdp_replicate_dim != 0 ): raise ValueError( f"num_gpus ({self.num_gpus}) must be >= and divisible by hsdp_replicate_dim ({self.hsdp_replicate_dim})" ) if ( self.hsdp_shard_dim > self.num_gpus or self.num_gpus % self.hsdp_shard_dim != 0 ): raise ValueError( f"num_gpus ({self.num_gpus}) must be >= and divisible by hsdp_shard_dim ({self.hsdp_shard_dim})" ) if self.num_gpus % self.dp_size != 0: raise ValueError( f"num_gpus ({self.num_gpus}) must be divisible by dp_size ({self.dp_size})" ) if self.dp_size < 1: raise ValueError("--dp-size must be a natural number") if self.dp_size > 1: raise ValueError("DP is not yet supported") num_gpus_per_group = self.dp_size * self.tp_size if self.enable_cfg_parallel: num_gpus_per_group *= 2 if self.num_gpus % num_gpus_per_group != 0: raise ValueError( f"num_gpus ({self.num_gpus}) must be divisible by (dp_size * tp_size{' * 2' if self.enable_cfg_parallel else ''}) = {num_gpus_per_group}" ) if self.sp_degree != self.ring_degree * self.ulysses_degree: raise ValueError( f"sp_degree ({self.sp_degree}) must equal ring_degree * ulysses_degree " f"({self.ring_degree} * {self.ulysses_degree} = {self.ring_degree * self.ulysses_degree})" ) if os.getenv("SGLANG_CACHE_DIT_ENABLED", "").lower() == "true": has_sp = self.sp_degree > 1 has_tp = self.tp_size > 1 if has_sp and has_tp: logger.warning( "cache-dit is enabled with hybrid parallelism (SP + TP). " "Proceeding anyway (SGLang integration may support this mode)." ) def _validate_cfg_parallel(self): if self.enable_cfg_parallel and self.num_gpus == 1: raise ValueError( "CFG Parallelism is enabled via `--enable-cfg-parallel`, but num_gpus == 1" ) def _set_default_attention_backend(self) -> None: """Configure ROCm defaults when users do not specify an attention backend.""" if current_platform.is_rocm(): default_backend = AttentionBackendEnum.AITER.name.lower() self.attention_backend = default_backend logger.info( "Attention backend not specified. Using '%s' by default on ROCm " "to match SGLang SRT defaults.", default_backend, ) @dataclasses.dataclass class PortArgs: # The ipc filename for scheduler (rank 0) to receive inputs from tokenizer (zmq) scheduler_input_ipc_name: str # The port for nccl initialization (torch.dist) nccl_port: int # The ipc filename for rpc call between Engine and Scheduler rpc_ipc_name: str # The ipc filename for Scheduler to send metrics metrics_ipc_name: str # Master port for distributed inference master_port: int | None = None @staticmethod def from_server_args( server_args: ServerArgs, dp_rank: Optional[int] = None ) -> "PortArgs": if server_args.nccl_port is None: nccl_port = server_args.scheduler_port + random.randint(100, 1000) while True: if is_port_available(nccl_port): break if nccl_port < 60000: nccl_port += 42 else: nccl_port -= 43 else: nccl_port = server_args.nccl_port # Normal case, use IPC within a single node return PortArgs( scheduler_input_ipc_name=f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}", nccl_port=nccl_port, rpc_ipc_name=f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}", metrics_ipc_name=f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}", master_port=server_args.master_port, ) _global_server_args = None def prepare_server_args(argv: list[str]) -> ServerArgs: """ Prepare the inference arguments from the command line arguments. """ parser = FlexibleArgumentParser() ServerArgs.add_cli_args(parser) raw_args, unknown_args = parser.parse_known_args(argv) server_args = ServerArgs.from_cli_args(raw_args, unknown_args) return server_args def set_global_server_args(server_args: ServerArgs): """ Set the global sgl_diffusion config for each process """ global _global_server_args _global_server_args = server_args def get_global_server_args() -> ServerArgs: if _global_server_args is None: # in ci, usually when we test custom ops/modules directly, # we don't set the sgl_diffusion config. In that case, we set a default # config. # TODO(will): may need to handle this for CI. raise ValueError("Global sgl_diffusion args is not set.") return _global_server_args ================================================ FILE: python/sglang/multimodal_gen/runtime/utils/common.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import ipaddress import logging import os import platform import signal import socket import sys import threading from functools import lru_cache import psutil import torch import zmq # use the native logger to avoid circular import logger = logging.getLogger(__name__) def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None): """Kill the process and all its child processes.""" # Remove sigchld handler to avoid spammy logs. if threading.current_thread() is threading.main_thread(): signal.signal(signal.SIGCHLD, signal.SIG_DFL) if parent_pid is None: parent_pid = os.getpid() include_parent = False try: itself = psutil.Process(parent_pid) except psutil.NoSuchProcess: return children = itself.children(recursive=True) for child in children: if child.pid == skip_pid: continue try: child.kill() except psutil.NoSuchProcess: pass if include_parent: try: if parent_pid == os.getpid(): itself.kill() sys.exit(0) itself.kill() # Sometime processes cannot be killed with SIGKILL (e.g, PID=1 launched by kubernetes), # so we send an additional signal to kill them. itself.send_signal(signal.SIGQUIT) except psutil.NoSuchProcess: pass def add_prefix(name: str, prefix: str) -> str: """Add a weight path prefix to a module name. Args: name: base module name. prefix: weight prefix str to added to the front of `name` concatenated with `.`. Returns: The string `prefix.name` if prefix is non-empty, otherwise just `name`. """ return name if not prefix else f"{prefix}.{name}" def is_valid_ipv6_address(address: str) -> bool: try: ipaddress.IPv6Address(address) return True except ValueError: return False def configure_ipv6(dist_init_addr): addr = dist_init_addr end = addr.find("]") if end == -1: raise ValueError("invalid IPv6 address format: missing ']'") host = addr[: end + 1] # this only validates the address without brackets: we still need the below checks. # if it's invalid, immediately raise an error so we know it's not formatting issues. if not is_valid_ipv6_address(host[1:end]): raise ValueError(f"invalid IPv6 address: {host}") port_str = None if len(addr) > end + 1: if addr[end + 1] == ":": port_str = addr[end + 2 :] else: raise ValueError("received IPv6 address format: expected ':' after ']'") if not port_str: raise ValueError( "a port must be specified in IPv6 address (format: [ipv6]:port)" ) try: port = int(port_str) except ValueError: raise ValueError(f"invalid port in IPv6 address: '{port_str}'") return port, host def is_port_available(port): """Return whether a port is available.""" with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: try: s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) s.bind(("", port)) s.listen(1) return True except socket.error: return False except OverflowError: return False def get_zmq_socket( context: zmq.Context, socket_type: zmq.SocketType, endpoint: str, bind: bool, max_bind_retries: int = 10, ) -> tuple[zmq.Socket, str]: """ Create and configure a ZMQ socket. Args: context: ZMQ context socket_type: Type of ZMQ socket endpoint: Endpoint string (e.g., "tcp://localhost:5555") bind: Whether to bind (True) or connect (False) max_bind_retries: Maximum number of retries if bind fails due to address already in use Returns: A tuple of (socket, actual_endpoint). The actual_endpoint may differ from the requested endpoint if bind retry was needed. """ mem = psutil.virtual_memory() total_mem = mem.total / 1024**3 available_mem = mem.available / 1024**3 if total_mem > 32 and available_mem > 16: buf_size = int(0.5 * 1024**3) else: buf_size = -1 socket = context.socket(socket_type) if endpoint.find("[") != -1: socket.setsockopt(zmq.IPV6, 1) def set_send_opt(): socket.setsockopt(zmq.SNDHWM, 0) socket.setsockopt(zmq.SNDBUF, buf_size) def set_recv_opt(): socket.setsockopt(zmq.RCVHWM, 0) socket.setsockopt(zmq.RCVBUF, buf_size) if socket_type == zmq.PUSH: set_send_opt() elif socket_type == zmq.PULL: set_recv_opt() elif socket_type in [zmq.DEALER, zmq.REQ, zmq.REP, zmq.ROUTER]: set_send_opt() set_recv_opt() else: raise ValueError(f"Unsupported socket type: {socket_type}") if bind: # Parse port from endpoint for retry logic import re port_match = re.search(r":(\d+)$", endpoint) if port_match and max_bind_retries > 1: original_port = int(port_match.group(1)) last_exception = None for attempt in range(max_bind_retries): try: current_endpoint = endpoint if attempt > 0: # Try next port (increment by 42 to match settle_port logic) current_port = original_port + attempt * 42 current_endpoint = re.sub( r":(\d+)$", f":{current_port}", endpoint ) logger.info( f"ZMQ bind failed for port {original_port + (attempt - 1) * 42}, " f"retrying with port {current_port} (attempt {attempt + 1}/{max_bind_retries})" ) socket.bind(current_endpoint) if attempt > 0: logger.warning( f"Successfully bound ZMQ socket to {current_endpoint} after {attempt + 1} attempts. " f"Original port {original_port} was unavailable." ) return socket, current_endpoint except zmq.ZMQError as e: last_exception = e if e.errno == zmq.EADDRINUSE and attempt < max_bind_retries - 1: # Address already in use, try next port continue elif attempt == max_bind_retries - 1: # Last attempt failed logger.error( f"Failed to bind ZMQ socket after {max_bind_retries} attempts. " f"Original endpoint: {endpoint}, Last tried port: {original_port + attempt * 42}" ) raise else: # Different error, raise immediately raise # Should not reach here, but just in case if last_exception: raise last_exception else: # No retry logic needed (either no port in endpoint or max_bind_retries == 1) socket.bind(endpoint) return socket, endpoint else: socket.connect(endpoint) return socket, endpoint return socket, endpoint # https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip @lru_cache(maxsize=1) def is_host_cpu_x86() -> bool: machine = platform.machine().lower() return ( machine in ("x86_64", "amd64", "i386", "i686") and hasattr(torch, "cpu") and torch.cpu.is_available() ) # cuda def set_cuda_arch(): capability = torch.cuda.get_device_capability() arch = f"{capability[0]}.{capability[1]}" os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}" # musa def set_musa_arch(): capability = torch.cuda.get_device_capability() arch = f"{capability[0]}{capability[1]}" os.environ["TORCH_MUSA_ARCH_LIST"] = f"{arch}" # env var managements _warned_bool_env_var_keys = set() def get_bool_env_var(name: str, default: str = "false") -> bool: value = os.getenv(name, default) value = str(value).strip().lower() truthy_values = {"1", "true", "yes", "y", "t", "on"} falsy_values = {"0", "false", "no", "n", "f", "off", ""} if (value not in truthy_values) and (value not in falsy_values): if value not in _warned_bool_env_var_keys: logger.warning( f"get_bool_env_var({name}) see non-understandable value={value} and treat as false" ) _warned_bool_env_var_keys.add(value) return value in truthy_values try: import sgl_kernel # noqa: F401 is_intel_amx_backend_available = hasattr( torch.ops.sgl_kernel, "convert_weight_packed" ) except: is_intel_amx_backend_available = False try: # move torch._C._cpu._is_amx_tile_supported() from cpu_has_amx_support # to support torch compile is_amx_tile_supported = torch._C._cpu._is_amx_tile_supported() except: is_amx_tile_supported = False def cpu_has_amx_support(): return is_amx_tile_supported and is_intel_amx_backend_available def use_intel_amx_backend(layer): return getattr(layer, "use_intel_amx_backend", False) ================================================ FILE: python/sglang/multimodal_gen/runtime/utils/distributed.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import pickle from typing import Any, List, Optional import numpy as np import torch import torch.distributed as dist from sglang.multimodal_gen.runtime.platforms import current_platform def broadcast_pyobj( data: List[Any], rank: int, dist_group: Optional[torch.distributed.ProcessGroup] = None, src: int = 0, force_cpu_device: bool = True, ): """Broadcast inputs from src rank to all other ranks with torch.dist backend. The `rank` here refer to the source rank on global process group (regardless of dist_group argument). """ device = torch.device( current_platform.device_type if not force_cpu_device else "cpu" ) if rank == src: if data is None or len(data) == 0: tensor_size = torch.tensor([0], dtype=torch.long, device=device) dist.broadcast(tensor_size, src=src, group=dist_group) else: serialized_data = pickle.dumps(data) size = len(serialized_data) tensor_data = torch.ByteTensor( np.frombuffer(serialized_data, dtype=np.uint8).copy() ).to(device) tensor_size = torch.tensor([size], dtype=torch.long, device=device) dist.broadcast(tensor_size, src=src, group=dist_group) dist.broadcast(tensor_data, src=src, group=dist_group) return data else: tensor_size = torch.tensor([0], dtype=torch.long, device=device) dist.broadcast(tensor_size, src=src, group=dist_group) size = tensor_size.item() if size == 0: return [] tensor_data = torch.empty(size, dtype=torch.uint8, device=device) dist.broadcast(tensor_data, src=src, group=dist_group) serialized_data = bytes(tensor_data.cpu().numpy()) data = pickle.loads(serialized_data) return data def generate_masked_orthogonal_rank_groups( world_size: int, parallel_size: list[int], mask: list[bool] ) -> list[list[int]]: """Generate orthogonal parallel groups based on the parallel size and mask. Arguments: world_size (int): world size parallel_size (List[int]): The parallel size of each orthogonal parallel type. For example, if tensor_parallel_size = 2, pipeline_model_parallel_group = 3, data_parallel_size = 4, and the parallel mapping order is tp-pp-dp, then the parallel_size = [2, 3, 4]. mask (List[bool]): The mask controls which parallel methods the generated groups represent. If mask[i] is True, it means the generated group contains the i-th parallelism method. For example, if parallel_size = [tp_size, pp_size, dp_size], and mask = [True, False , True], then the generated group is the `tp-dp` group, if the mask = [False, True, False], then the generated group is the `pp` group. Algorithm: For orthogonal parallelism, such as tp/dp/pp/cp, the global_rank and If we want to get the `dp_group` (tp_size * pp_size groups of dp_size ranks each. For example, if the gpu size is 8 and order is 'tp-pp-dp', size is '2-2-2', and the dp_group here is [[0, 4], [1, 5], [2, 6], [3, 7]].) The tp_rank and pp_rank will be combined to form the `dp_group_index`. dp_group_index = tp_rank + pp_rank * tp_size (2) So, Given that tp_rank and pp_rank satisfy equation (2), and dp_rank in range(0, dp_size), the ranks in dp_group[dp_group_index] satisfies the equation (1). This function solve this math problem. For example, if the parallel_size = [tp_size, dp_size, pp_size] = [2, 3, 4], and the mask = [False, True, False]. Then, dp_group_index(0) = tp_rank(0) + pp_rank(0) * 2 dp_group_index(1) = tp_rank(1) + pp_rank(0) * 2 ... dp_group_index(7) = tp_rank(1) + pp_rank(3) * 2 dp_group[0] = 0 + range(0, 3) * 2 + 0 = [0, 2, 4] dp_group[1] = 1 + range(0, 3) * 2 + 0 = [1, 3, 5] ... dp_group[7] = 1 + range(0, 3) * 2 + 3 * 2 * 3 = [19, 21, 23] """ def prefix_product(a: List[int], init=1) -> List[int]: r = [init] for v in a: init = init * v r.append(init) return r def inner_product(a: List[int], b: List[int]) -> int: return sum([x * y for x, y in zip(a, b)]) def decompose(index, shape, stride=None): """ This function solve the math problem below: There is an equation: index = sum(idx[i] * stride[i]) And given the value of index, stride. Return the idx. This function will used to get the pp/dp/pp_rank from group_index and rank_in_group. """ if stride is None: stride = prefix_product(shape) idx = [(index // d) % s for s, d in zip(shape, stride)] # stride is a prefix_product result. And the value of stride[-1] # is not used. assert ( sum([x * y for x, y in zip(idx, stride[:-1])]) == index ), "idx {} with shape {} mismatch the return idx {}".format(index, shape, idx) return idx masked_shape = [s for s, m in zip(parallel_size, mask) if m] unmasked_shape = [s for s, m in zip(parallel_size, mask) if not m] global_stride = prefix_product(parallel_size) masked_stride = [d for d, m in zip(global_stride, mask) if m] unmasked_stride = [d for d, m in zip(global_stride, mask) if not m] group_size = prefix_product(masked_shape)[-1] num_of_group = world_size // group_size ranks = [] for group_index in range(num_of_group): # get indices from unmaksed for group_index. decomposed_group_idx = decompose(group_index, unmasked_shape) rank = [] for rank_in_group in range(group_size): # get indices from masked for rank_in_group. decomposed_rank_idx = decompose(rank_in_group, masked_shape) rank.append( inner_product(decomposed_rank_idx, masked_stride) + inner_product(decomposed_group_idx, unmasked_stride) ) ranks.append(rank) return ranks class RankGenerator(object): def __init__( self, tp: int, sp: int, pp: int, cfg: int, dp: int, order: str, rank_offset: int = 0, ) -> None: self.tp = tp self.sp = sp self.pp = pp self.cfg = cfg self.dp = dp self.rank_offset = rank_offset self.world_size = tp * sp * pp * cfg * dp self.name_to_size = { "tp": self.tp, "sp": self.sp, "pp": self.pp, "cfg": self.cfg, "dp": self.dp, } order = order.lower() for name in self.name_to_size.keys(): if name not in order and self.name_to_size[name] != 1: raise RuntimeError( f"The size of ({name}) is ({self.name_to_size[name]}), but you haven't specified the order ({self.order})." ) elif name not in order: order = order + "-" + name self.order = order self.ordered_size = [] for token in order.split("-"): self.ordered_size.append(self.name_to_size[token]) def get_mask(self, order: str, token: str): ordered_token = order.split("-") token = token.split("-") mask = [False] * len(ordered_token) for t in token: mask[ordered_token.index(t)] = True return mask def get_ranks(self, token): """Get rank group by input token. Arguments: token (str): Specify the ranks type that want to get. If we want to obtain multiple parallel types, we can use a hyphen '-' to separate them. For example, if we want to obtain the TP_DP group, the token should be 'tp-dp'. """ mask = self.get_mask(self.order, token) ranks = generate_masked_orthogonal_rank_groups( self.world_size, self.ordered_size, mask ) if self.rank_offset > 0: for rank_group in ranks: for i in range(len(rank_group)): rank_group[i] += self.rank_offset return ranks ================================================ FILE: python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/hf_transformers_utils.py # Copyright 2023-2024 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Utilities for Huggingface Transformers.""" import contextlib import glob import json import os import shutil import time from functools import reduce from pathlib import Path from typing import Any, Optional, Union, cast from diffusers.loaders.lora_base import ( _best_guess_weight_name, # watch out for potetential removal from diffusers ) from huggingface_hub.errors import ( LocalEntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError, ) from requests.exceptions import ConnectionError as RequestsConnectionError from requests.exceptions import RequestException from transformers import AutoConfig, PretrainedConfig from sglang.multimodal_gen.runtime.loader.utils import _clean_hf_config_inplace from sglang.multimodal_gen.runtime.loader.weight_utils import get_lock from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.srt.environ import envs from sglang.utils import is_in_ci logger = init_logger(__name__) def _check_index_files_for_missing_shards( model_path: str, ) -> tuple[bool, list[str], list[str]]: """ Check all subdirectories for missing shards based on index files. This catches cases where a model download was interrupted, leaving some safetensors shards missing while the index file exists. Args: model_path: Path to the model directory Returns: Tuple of (all_valid, missing_files, checked_subdirs) """ missing_files = [] checked_subdirs = [] # Add common subdirectories for diffusers models try: subdirs = os.listdir(model_path) except OSError as e: logger.warning("Failed to list model directory %s: %s", model_path, e) return True, [], [] # Assume valid if we can't check # Check the root directory and all subdirectories that might contain model weights dirs_to_check = [model_path] for subdir in subdirs: subdir_path = os.path.join(model_path, subdir) if os.path.isdir(subdir_path): dirs_to_check.append(subdir_path) for dir_path in dirs_to_check: # Find all safetensors index files index_files = glob.glob(os.path.join(dir_path, "*.safetensors.index.json")) for index_file in index_files: checked_subdirs.append(os.path.basename(dir_path)) try: with open(index_file) as f: index_data = json.load(f) weight_map = index_data.get("weight_map", {}) if not weight_map: continue # Get unique files referenced in weight_map required_files = set(weight_map.values()) for file_name in required_files: file_path = os.path.join(dir_path, file_name) if not os.path.exists(file_path): relative_path = os.path.relpath(file_path, model_path) missing_files.append(relative_path) except Exception as e: logger.warning("Failed to read index file %s: %s", index_file, e) continue return len(missing_files) == 0, missing_files, checked_subdirs def _cleanup_model_cache(model_path: str, reason: str) -> bool: """ Remove the model cache directory to force a clean re-download. Args: model_path: Path to the model directory (snapshot path) reason: Reason for cleanup (for logging) Returns: True if cleanup was performed, False otherwise """ # Navigate up to the model root directory: snapshots/hash -> snapshots -> model_root # HF cache structure: models--org--name/snapshots/hash/ try: snapshot_dir = os.path.abspath(model_path) snapshots_dir = os.path.dirname(snapshot_dir) repo_folder = os.path.dirname(snapshots_dir) # Verify this looks like an HF cache structure if os.path.basename(snapshots_dir) != "snapshots": logger.warning( "Model path %s doesn't appear to be in HF cache structure, skipping cleanup", model_path, ) return False logger.warning( "Removing model cache at %s. Reason: %s", repo_folder, reason, ) shutil.rmtree(repo_folder) logger.info("Successfully removed corrupted cache directory") return True except Exception as e: logger.error( "Failed to remove corrupted cache directory %s: %s. " "Manual cleanup may be required.", model_path, e, ) return False def _ci_validate_diffusers_model(model_path: str) -> tuple[bool, bool]: """ CI-specific validation for diffusers models. Checks all subdirectories (transformer, transformer_2, vae, etc.) for missing shards based on their index files. If issues are found in CI, cleans up the cache to force re-download. Args: model_path: Path to the model directory Returns: Tuple of (is_valid, cleanup_performed) - is_valid: True if the model is valid - cleanup_performed: True if cleanup was performed (only relevant when is_valid=False) """ if not is_in_ci(): return True, False is_valid, missing_files, checked_subdirs = _check_index_files_for_missing_shards( model_path ) if not is_valid: logger.error( "CI validation failed for %s. Missing %d file(s): %s. " "Checked subdirectories: %s", model_path, len(missing_files), missing_files[:5] if len(missing_files) > 5 else missing_files, checked_subdirs, ) cleanup_performed = _cleanup_model_cache( model_path, f"Missing {len(missing_files)} shard file(s): {missing_files[:3]}", ) return False, cleanup_performed if checked_subdirs: logger.info( "CI validation passed for %s. Checked subdirectories: %s", model_path, checked_subdirs, ) return True, False def _verify_diffusers_model_complete(path: str) -> bool: """Check if a diffusers model directory has all required component subdirectories.""" config_path = os.path.join(path, "model_index.json") if not os.path.exists(config_path): return False try: with open(config_path) as config_file: model_index = json.load(config_file) except Exception as exc: logger.warning("Failed to read model_index.json at %s: %s", config_path, exc) return False component_keys = [ key for key, value in model_index.items() if isinstance(value, (list, tuple)) and len(value) == 2 and all(isinstance(item, str) for item in value) ] if component_keys: return all(os.path.exists(os.path.join(path, key)) for key in component_keys) return os.path.exists(os.path.join(path, "transformer")) and os.path.exists( os.path.join(path, "vae") ) _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = { # ChatGLMConfig.model_type: ChatGLMConfig, # DbrxConfig.model_type: DbrxConfig, # ExaoneConfig.model_type: ExaoneConfig, # Qwen2_5_VLConfig.model_type: Qwen2_5_VLConfig, } for name, cls in _CONFIG_REGISTRY.items(): with contextlib.suppress(ValueError): AutoConfig.register(name, cls) def download_from_hf(model_path: str): if os.path.exists(model_path): return model_path return snapshot_download(model_path, allow_patterns=["*.json", "*.bin", "*.model"]) def get_hf_config( component_model_path: str, trust_remote_code: bool, revision: str | None = None, model_override_args: dict | None = None, **kwargs, ) -> PretrainedConfig: if check_gguf_file(component_model_path): raise NotImplementedError("GGUF models are not supported.") config = AutoConfig.from_pretrained( component_model_path, trust_remote_code=trust_remote_code, revision=revision, **kwargs, ) if config.model_type in _CONFIG_REGISTRY: config_class = _CONFIG_REGISTRY[config.model_type] config = config_class.from_pretrained(component_model_path, revision=revision) # NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`. config._name_or_path = component_model_path if model_override_args: config.update(model_override_args) return config def get_config( model: str, trust_remote_code: bool, revision: Optional[str] = None, model_override_args: Optional[dict] = None, **kwargs, ): return AutoConfig.from_pretrained( model, trust_remote_code=trust_remote_code, revision=revision, **kwargs ) def load_dict(file_path): if not os.path.exists(file_path): return {} try: # Load the config directly from the file with open(file_path) as f: config_dict: dict[str, Any] = json.load(f) if "_diffusers_version" in config_dict: config_dict.pop("_diffusers_version") # TODO(will): apply any overrides from inference args return config_dict except Exception as e: raise RuntimeError( f"Failed to load diffusers config from {file_path}: {e}" ) from e def get_diffusers_component_config( component_path: str, ) -> dict[str, Any]: """Gets a configuration of a submodule for the given diffusers model.""" # Download from HuggingFace Hub if path doesn't exist locally if not os.path.exists(component_path): component_path = maybe_download_model(component_path) config_names = ["generation_config.json"] # By default, we load config.json, but scheduler_config.json for scheduler if "scheduler" in component_path: config_names.append("scheduler_config.json") else: config_names.append("config.json") config_file_paths = [ os.path.join(component_path, config_name) for config_name in config_names ] combined_config = reduce( lambda acc, path: acc | load_dict(path), config_file_paths, {} ) _clean_hf_config_inplace(combined_config) logger.debug("HF model config: %s", combined_config) return combined_config # Models don't use the same configuration key for determining the maximum # context length. Store them here so we can sanely check them. # NOTE: The ordering here is important. Some models have two of these and we # have a preference for which value gets used. CONTEXT_LENGTH_KEYS = [ "max_sequence_length", "seq_length", "max_seq_len", "model_max_length", "max_position_embeddings", ] def attach_additional_stop_token_ids(tokenizer): # Special handling for stop token <|eom_id|> generated by llama 3 tool use. if "<|eom_id|>" in tokenizer.get_added_vocab(): tokenizer.additional_stop_token_ids = { tokenizer.get_added_vocab()["<|eom_id|>"] } else: tokenizer.additional_stop_token_ids = None def check_gguf_file(model: str | os.PathLike) -> bool: """Check if the file is a GGUF model.""" model = Path(model) if not model.is_file(): return False elif model.suffix == ".gguf": return True with open(model, "rb") as f: header = f.read(4) return header == b"GGUF" def maybe_download_lora( model_name_or_path: str, local_dir: str | None = None, download: bool = True ) -> str: """ Check if the model path is a Hugging Face Hub model ID and download it if needed. Args: model_name_or_path: Local path or Hugging Face Hub model ID local_dir: Local directory to save the model download: Whether to download the model from Hugging Face Hub Returns: Local path to the model """ allow_patterns = ["*.json", "*.safetensors", "*.bin"] local_path = maybe_download_model( model_name_or_path, local_dir, download, is_lora=True, allow_patterns=allow_patterns, ) # return directly if local_path is a file if os.path.isfile(local_path): return local_path weight_name = _best_guess_weight_name(local_path, file_extension=".safetensors") # AMD workaround: PR 15813 changed from model_name_or_path to local_path, # which can return None. Fall back to original behavior on ROCm. if weight_name is None and current_platform.is_rocm(): weight_name = _best_guess_weight_name( model_name_or_path, file_extension=".safetensors" ) return os.path.join(local_path, weight_name) def verify_model_config_and_directory(model_path: str) -> dict[str, Any]: """ Verify that the model directory contains a valid diffusers configuration. Args: model_path: Path to the model directory Returns: The loaded model configuration as a dictionary """ # Check for model_index.json which is required for diffusers models config_path = os.path.join(model_path, "model_index.json") if not os.path.exists(config_path): raise ValueError( f"Model directory {model_path} does not contain model_index.json. " "Only HuggingFace diffusers format is supported." ) # Load the config with open(config_path) as f: config = json.load(f) # Verify diffusers version exists if "_diffusers_version" not in config: raise ValueError("model_index.json does not contain _diffusers_version") logger.info("Diffusers version: %s", config["_diffusers_version"]) component_keys = [ key for key, value in config.items() if isinstance(value, (list, tuple)) and len(value) == 2 and all(isinstance(item, str) for item in value) ] if component_keys: missing_components = [ component_key for component_key in component_keys if not os.path.exists(os.path.join(model_path, component_key)) ] if missing_components: missing_str = ", ".join(missing_components) raise ValueError( f"Model directory {model_path} is missing required component " f"directories: {missing_str}." ) else: transformer_dir = os.path.join(model_path, "transformer") vae_dir = os.path.join(model_path, "vae") if not os.path.exists(transformer_dir): raise ValueError( f"Model directory {model_path} does not contain a transformer/ directory." ) if not os.path.exists(vae_dir): raise ValueError( f"Model directory {model_path} does not contain a vae/ directory." ) return cast(dict[str, Any], config) def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]: """ Download and extract just the model_index.json for a Hugging Face model. Args: model_name_or_path: Path or HF Hub model ID Returns: The parsed model_index.json as a dictionary """ import tempfile from huggingface_hub.errors import EntryNotFoundError # If it's a local path, verify it directly if os.path.exists(model_name_or_path): try: return verify_model_config_and_directory(model_name_or_path) except ValueError: # Not a pipeline, maybe a single model. config_path = os.path.join(model_name_or_path, "config.json") if os.path.exists(config_path): with open(config_path) as f: config = json.load(f) return config raise # For remote models, download just the model_index.json try: with tempfile.TemporaryDirectory() as tmp_dir: # Download just the model_index.json file model_index_path = hf_hub_download( repo_id=model_name_or_path, filename="model_index.json", local_dir=tmp_dir, ) # Load the model_index.json with open(model_index_path) as f: config: dict[str, Any] = json.load(f) # Verify it has the required fields if "_class_name" not in config: raise ValueError( f"model_index.json for {model_name_or_path} does not contain _class_name field" ) if "_diffusers_version" not in config: raise ValueError( f"model_index.json for {model_name_or_path} does not contain _diffusers_version field" ) # Add the pipeline name for downstream use config["pipeline_name"] = config["_class_name"] logger.debug( "Downloaded model_index.json for %s, pipeline: %s", model_name_or_path, config["_class_name"], ) return config except EntryNotFoundError: logger.warning( "model_index.json not found for %s. Assuming it is a single model and downloading it.", model_name_or_path, ) local_path = maybe_download_model(model_name_or_path) config_path = os.path.join(local_path, "config.json") if not os.path.exists(config_path): raise ValueError( f"Failed to find config.json for {model_name_or_path} after failing to find model_index.json" f"You might be looking for models ending with '-Diffusers'" ) with open(config_path) as f: config = json.load(f) return config except Exception as e: raise ValueError( f"Failed to download or parse model_index.json for {model_name_or_path}: {e}" ) from e def maybe_download_model( model_name_or_path: str, local_dir: str | None = None, download: bool = True, is_lora: bool = False, allow_patterns: list[str] | None = None, force_diffusers_model: bool = False, ) -> str: """ Check if the model path is a Hugging Face Hub model ID and download it if needed. Args: model_name_or_path: Local path or Hugging Face Hub model ID local_dir: Local directory to save the model download: Whether to download the model from Hugging Face Hub is_lora: If True, skip model completeness verification (LoRA models don't have transformer/vae directories) force_diffusers_model: If True, apply diffusers model check. Otherwise it should be a component model Returns: Local path to the model """ # 1. Local path check: if path exists locally, verify it's complete (skip for LoRA) if os.path.exists(model_name_or_path): if not force_diffusers_model: return model_name_or_path if is_lora or _verify_diffusers_model_complete(model_name_or_path): if not is_lora: is_valid, cleanup_performed = _ci_validate_diffusers_model( model_name_or_path ) if not is_valid: if cleanup_performed: logger.warning( "CI validation failed for local model at %s, " "cache has been cleaned up, will re-download", model_name_or_path, ) # Fall through to download else: raise ValueError( f"CI validation failed for local model at {model_name_or_path}. " "Some safetensors shards are missing. " "Please manually delete the model directory and retry." ) else: logger.info("Model already exists locally and is complete") return model_name_or_path else: logger.info("Model already exists locally and is complete") return model_name_or_path else: logger.warning( "Local model at %s appears incomplete (missing required components), " "will attempt re-download", model_name_or_path, ) # 2. Cache-first strategy (Fast Path) # Try to read from HF cache without network access try: logger.info( "Checking for cached model in HF Hub cache for %s...", model_name_or_path ) local_path = snapshot_download( repo_id=model_name_or_path, ignore_patterns=["*.onnx", "*.msgpack"], local_dir=local_dir, local_files_only=True, max_workers=8, ) if not force_diffusers_model: return str(local_path) if is_lora or _verify_diffusers_model_complete(local_path): if not is_lora: is_valid, cleanup_performed = _ci_validate_diffusers_model(local_path) if not is_valid: logger.warning( "CI validation failed for cached model at %s, " "%s, will re-download", local_path, ( "cache has been cleaned up" if cleanup_performed else "cleanup was not performed" ), ) # Fall through to download else: logger.info("Found complete model in cache at %s", local_path) return str(local_path) else: logger.info("Found complete model in cache at %s", local_path) return str(local_path) else: if not download: raise ValueError( f"Model {model_name_or_path} found in cache but is incomplete and download=False." ) logger.info( "Model found in cache but incomplete, will download from HF Hub" ) except LocalEntryNotFoundError: if not download: raise ValueError( f"Model {model_name_or_path} not found in local cache and download=False." ) logger.info("Model not found in cache, will download from HF Hub") except Exception as e: logger.warning( "Unexpected error while checking cache for %s: %s, will attempt download", model_name_or_path, e, ) if not download: raise ValueError( f"Error checking cache for {model_name_or_path} and download=False: {e}" ) from e # 3. Download strategy (with retry mechanism) MAX_RETRIES = 5 for attempt in range(MAX_RETRIES): try: logger.info( "Downloading model snapshot from HF Hub for %s (attempt %d/%d)...", model_name_or_path, attempt + 1, MAX_RETRIES, ) with get_lock(model_name_or_path).acquire(poll_interval=2): local_path = snapshot_download( repo_id=model_name_or_path, ignore_patterns=["*.onnx", "*.msgpack"], allow_patterns=allow_patterns, local_dir=local_dir, max_workers=8, ) if not force_diffusers_model: return str(local_path) # Verify downloaded model is complete (skip for LoRA) elif not is_lora and not _verify_diffusers_model_complete(local_path): logger.warning( "Downloaded model at %s is incomplete, retrying with force_download=True", local_path, ) with get_lock(model_name_or_path).acquire(poll_interval=2): local_path = snapshot_download( repo_id=model_name_or_path, ignore_patterns=["*.onnx", "*.msgpack"], local_dir=local_dir, max_workers=8, force_download=True, ) if not _verify_diffusers_model_complete(local_path): raise ValueError( f"Downloaded model at {local_path} is still incomplete after forced re-download. " "The model repository may be missing required components (model_index.json, transformer/, or vae/)." ) # CI validation: check all subdirectories for missing shards after download if not is_lora: is_valid, cleanup_performed = _ci_validate_diffusers_model(local_path) if not is_valid: # In CI, if validation fails after download, we have a serious issue # If cleanup was performed, the next retry should get a fresh download raise ValueError( f"CI validation failed for downloaded model at {local_path}. " f"Some safetensors shards are missing. Cleanup performed: {cleanup_performed}." ) logger.info("Downloaded model to %s", local_path) return str(local_path) except (RepositoryNotFoundError, RevisionNotFoundError) as e: raise ValueError( f"Model or revision not found at {model_name_or_path}. " f"Please check the model ID or ensure you have access to the repository. Error: {e}" ) from e except (RequestException, RequestsConnectionError) as e: if attempt == MAX_RETRIES - 1: raise ValueError( f"Could not find model at {model_name_or_path} and failed to download from HF Hub " f"after {MAX_RETRIES} attempts due to network error: {e}" ) from e wait_time = 2**attempt logger.warning( "Download failed (attempt %d/%d) due to network error: %s. " "Retrying in %d seconds...", attempt + 1, MAX_RETRIES, e, wait_time, ) time.sleep(wait_time) except Exception as e: raise ValueError( f"Could not find model at {model_name_or_path} and failed to download from HF Hub: {e}" ) from e # Unified download functions with Hugging Face-compatible names def hf_hub_download( repo_id: str, filename: str, local_dir: Optional[Union[str, Path]] = None, **kwargs, ) -> str: """Unified hf_hub_download that supports both Hugging Face Hub and ModelScope.""" if envs.SGLANG_USE_MODELSCOPE.get(): from modelscope import model_file_download return model_file_download( model_id=repo_id, file_path=filename, cache_dir=local_dir, **kwargs, ) else: from huggingface_hub import hf_hub_download as _hf_hub_download return _hf_hub_download( repo_id=repo_id, filename=filename, local_dir=local_dir, **kwargs, ) def snapshot_download( repo_id: str, local_dir: Optional[Union[str, Path]] = None, ignore_patterns: Optional[Union[list[str], str]] = None, allow_patterns: Optional[Union[list[str], str]] = None, local_files_only: bool = False, max_workers: int = 8, **kwargs, ) -> str: """Unified snapshot_download that supports both Hugging Face Hub and ModelScope.""" if envs.SGLANG_USE_MODELSCOPE.get(): from modelscope import snapshot_download as _ms_snapshot_download ms_kwargs = { "model_id": repo_id, "local_dir": local_dir, "ignore_patterns": ignore_patterns, "allow_patterns": allow_patterns, "local_files_only": local_files_only, "max_workers": max_workers, } ms_kwargs.update(kwargs) return _ms_snapshot_download(**ms_kwargs) else: from huggingface_hub import snapshot_download as _hf_snapshot_download hf_kwargs = { "repo_id": repo_id, "local_dir": local_dir, "ignore_patterns": ignore_patterns, "allow_patterns": allow_patterns, "local_files_only": local_files_only, "max_workers": max_workers, "etag_timeout": 60, } hf_kwargs.update(kwargs) return _hf_snapshot_download(**hf_kwargs) ================================================ FILE: python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py ================================================ import re from itertools import chain from typing import Any, Dict, List, Set, Tuple import torch from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # Adapted from skywork AI Infra diffusion optimize class LayerwiseOffloadManager: """A lightweight layerwise CPU offload manager. This utility offloads per-layer parameters/buffers from GPU to CPU, and supports async H2D prefetch using a dedicated CUDA stream. Typical usage: - Construct the manager with the target model and the list-like module attribute that represents transformer blocks (e.g. ``blocks``). - Call :meth:`initialize` once to offload weights and prefetch layer 0. - During forward, call :meth:`prefetch_layer` for the next layer and :meth:`release_layer` for the finished layer. """ def __init__( self, model: torch.nn.Module, *, layers_attr_str: str, num_layers: int, enabled: bool, pin_cpu_memory: bool = True, prefetch_size: int = 1, ) -> None: self.model = model self.layers_attr_str = layers_attr_str self.num_layers = num_layers self.pin_cpu_memory = pin_cpu_memory self.prefetch_size = min(max(1, prefetch_size), self.num_layers) self.enabled = bool(enabled and torch.get_device_module().is_available()) if not self.enabled: return self.device = torch.device( current_platform.device_type, torch.get_device_module().current_device() ) self.copy_stream = torch.get_device_module().Stream() self._layer_name_re = re.compile( rf"(^|\.){re.escape(layers_attr_str)}\.(\d+)(\.|$)" ) # layer_idx -> {dtype: consolidated_pinned_cpu_tensor} # stores the consolidated weight from a same layer, of same dtype self._consolidated_cpu_weights: Dict[int, Dict[torch.dtype, torch.Tensor]] = {} # layer_idx -> {name: {dtype, offset, numel, shape}} # stores the offset and numel of each weight from a same layer, of same dtype self._weight_metadata: Dict[int, Dict[str, Dict[str, Any]]] = {} # layer indices that are already in gpu self._gpu_layers: Set[int] = set() # layer_idx -> torch.get_device_module().Event for fine-grained sync, to make sure the weight is resident in pre-hook self._prefetch_events: Dict[int, torch.get_device_module().Event] = {} self._named_parameters: Dict[str, torch.nn.Parameter] = {} self._named_buffers: Dict[str, torch.Tensor] = {} self._offload_placeholders: Dict[torch.dtype, torch.Tensor] = {} # Store forward hooks for removal self._forward_hooks: List[Any] = [] self._initialize() def _match_layer_idx(self, name: str) -> int | None: m = self._layer_name_re.search(name) if not m: return None try: return int(m.group(2)) except Exception: return None def _get_shared_empty_tensor(self, dtype: torch.dtype) -> torch.Tensor: placeholder = self._offload_placeholders.get(dtype) if placeholder is None: placeholder = torch.empty((1,), device=self.device, dtype=dtype) self._offload_placeholders[dtype] = placeholder return placeholder @torch.compiler.disable def _initialize(self) -> None: if not self.enabled: return self._named_parameters = dict(self.model.named_parameters()) self._named_buffers = dict(self.model.named_buffers()) # 1. collect and group tensors by layer and dtype layer_groups: Dict[int, Dict[torch.dtype, List[Tuple[str, torch.Tensor]]]] = {} all_tensors = chain(self._named_parameters.items(), self._named_buffers.items()) for name, tensor in all_tensors: layer_idx = self._match_layer_idx(name) if layer_idx is None or layer_idx >= self.num_layers: continue layer_groups.setdefault(layer_idx, {}).setdefault(tensor.dtype, []).append( (name, tensor) ) # 2. concat and offload (in pinned memory) for layer_idx, dtype_to_params in layer_groups.items(): self._consolidated_cpu_weights[layer_idx] = {} self._weight_metadata[layer_idx] = {} for dtype, weights in dtype_to_params.items(): total_numel = sum(t.numel() for _, t in weights) # create concatenated CPU buffer (in pinned memory) cpu_buffer = torch.empty( total_numel, dtype=dtype, pin_memory=self.pin_cpu_memory ) # offload weights to the buffer current_offset = 0 for name, weight in weights: numel = weight.numel() cpu_buffer[current_offset : current_offset + numel].copy_( weight.flatten() ) self._weight_metadata[layer_idx][name] = { "dtype": dtype, "offset": current_offset, "numel": numel, "shape": weight.shape, } weight.data = self._get_shared_empty_tensor(dtype) current_offset += numel self._consolidated_cpu_weights[layer_idx][dtype] = cpu_buffer # prefetch the first layer for warm-up self.prepare_for_next_req(non_blocking=False) self.register_forward_hooks() logger.info( f"LayerwiseOffloadManager initialized with num prefetched layer: {self.prefetch_size}, total num layers: {self.num_layers}" ) def prepare_for_next_req(self, non_blocking=True): """ Prepare for the next round of denoising loop with prefetching the necessary layers """ for i in range(self.prefetch_size): self.prefetch_layer(i, non_blocking=non_blocking) if not non_blocking and self.copy_stream is not None: torch.get_device_module().current_stream().wait_stream(self.copy_stream) def get_target_with_name(self, name: str) -> torch.Tensor: """get the target model weight/buffer to be replaced""" if name in self._named_parameters: target = self._named_parameters[name] else: target = self._named_buffers[name] return target @torch.compiler.disable def prefetch_layer(self, layer_idx: int, non_blocking: bool = True) -> None: """ idempotent """ if not self.enabled or self.device is None or self.copy_stream is None: return if layer_idx < 0 or layer_idx >= self.num_layers: return if layer_idx in self._gpu_layers: return if layer_idx not in self._consolidated_cpu_weights: return self.copy_stream.wait_stream(torch.get_device_module().current_stream()) # create gpu buffer and load from CPU buffer gpu_buffers: Dict[torch.dtype, torch.Tensor] = {} with torch.get_device_module().stream(self.copy_stream): for dtype, cpu_buffer in self._consolidated_cpu_weights[layer_idx].items(): gpu_buffer = torch.empty( cpu_buffer.shape, dtype=dtype, device=self.device ) gpu_buffer.copy_(cpu_buffer, non_blocking=non_blocking) gpu_buffers[dtype] = gpu_buffer # record the prefetch event of this layer event = torch.get_device_module().Event() event.record(self.copy_stream) self._prefetch_events[layer_idx] = event # restore model's weights by their metadata using gpu buffer for name, meta in self._weight_metadata[layer_idx].items(): dtype = meta["dtype"] gpu_buffer = gpu_buffers[dtype] # map the parameter's data to the correct slice of the GPU buffer target = self.get_target_with_name(name) target.data = gpu_buffer[ meta["offset"] : meta["offset"] + meta["numel"] ].view(meta["shape"]) self._gpu_layers.add(layer_idx) @torch.compiler.disable def release_layer(self, layer_idx: int) -> None: """ lightweight release layer weights Basically set the reference count to the gpu weight tensor to zero. The weights on cpu is untouched """ if not self.enabled or self.device is None: return # clear prefetch event, since it's useless and needs to be reset self._prefetch_events.pop(layer_idx, None) if layer_idx not in self._gpu_layers: return for name, meta in self._weight_metadata.get(layer_idx, {}).items(): target = self.get_target_with_name(name) # Wraparound prefetch will reload the layer when it is needed again target.data = self._get_shared_empty_tensor(meta["dtype"]) self._gpu_layers.discard(layer_idx) @torch.compiler.disable def release_all(self) -> None: if not self.enabled or self.device is None: return if self.copy_stream is not None: torch.get_device_module().current_stream().wait_stream(self.copy_stream) for layer_idx in list(self._gpu_layers): self.release_layer(layer_idx) @torch.compiler.disable def load_all_layers(self) -> None: """Load all layers from CPU to GPU.""" if not self.enabled or self.device is None: return if self.copy_stream is not None: torch.get_device_module().current_stream().wait_stream(self.copy_stream) for layer_idx in range(self.num_layers): if layer_idx not in self._gpu_layers: self.prefetch_layer(layer_idx, non_blocking=False) @torch.compiler.disable def sync_layer_to_cpu(self, layer_idx: int) -> None: """Sync a layer's weights from GPU back to CPU.""" if not self.enabled or layer_idx not in self._gpu_layers: return if layer_idx not in self._consolidated_cpu_weights: return if self.copy_stream is not None: torch.get_device_module().current_stream().wait_stream(self.copy_stream) # Collect current GPU weights and write back to CPU buffer for name, meta in self._weight_metadata.get(layer_idx, {}).items(): target = self.get_target_with_name(name) gpu_weight = target.data.flatten().cpu() dtype = meta["dtype"] cpu_buffer = self._consolidated_cpu_weights[layer_idx][dtype] offset = meta["offset"] numel = meta["numel"] cpu_buffer[offset : offset + numel].copy_(gpu_weight) @torch.compiler.disable def sync_all_layers_to_cpu(self) -> None: """Sync all loaded layers' weights from GPU back to CPU.""" if not self.enabled or self.device is None: return if self.copy_stream is not None: torch.get_device_module().current_stream().wait_stream(self.copy_stream) for layer_idx in list(self._gpu_layers): self.sync_layer_to_cpu(layer_idx) @torch.compiler.disable def update_cpu_weights( self, weight_dict: Dict[str, torch.Tensor] ) -> Set[str] | None: """Update consolidated CPU buffers with new weights. When layerwise offload (--dit-layerwise-offload) is enabled, the offload manager replaces GPU parameters with small torch.empty((1,)) placeholders while real weights live in consolidated pinned CPU buffers. The refit process writes new weights directly into the CPU buffers, bypassing the placeholders. For any layer that happens to be resident on the GPU at update time, the live GPU tensor is also updated. Args: weight_dict: Mapping of parameter name to new weight tensor. Returns: Set of parameter names that were successfully updated. Raises: ValueError: If a weight's shape does not match the recorded metadata (i.e., the real shape, not the placeholder shape). """ if not self.enabled: return None updated_names: Set[str] = set() for name, loaded_weight in weight_dict.items(): layer_idx = self._match_layer_idx(name) if layer_idx is None: continue meta_layer = self._weight_metadata.get(layer_idx) if meta_layer is None or name not in meta_layer: continue meta = meta_layer[name] if tuple(meta["shape"]) != tuple(loaded_weight.shape): raise ValueError( f"Shape mismatch for {name}: " f"expected={tuple(meta['shape'])}, " f"loaded={tuple(loaded_weight.shape)}" ) dtype = meta["dtype"] offset = meta["offset"] numel = meta["numel"] cpu_buffer = self._consolidated_cpu_weights[layer_idx][dtype] cpu_buffer[offset : offset + numel].copy_( loaded_weight.to(dtype=dtype).flatten() ) # If this layer is currently on GPU, update the live parameter. if layer_idx in self._gpu_layers: target = self.get_target_with_name(name) target.data.copy_(loaded_weight.to(dtype=target.dtype)) updated_names.add(name) return updated_names def iter_cpu_weights(self): """Yield (name, tensor) pairs from consolidated CPU buffers. This reconstructs the original weight tensors (with correct shapes) from the flat CPU buffers using stored metadata. Unlike model.named_parameters(), which returns (1,) placeholders when offload is enabled, this method returns the real weights and can be used for checksum computation. """ for layer_idx in sorted(self._weight_metadata): for name, meta in self._weight_metadata[layer_idx].items(): dtype = meta["dtype"] offset = meta["offset"] numel = meta["numel"] shape = meta["shape"] cpu_buffer = self._consolidated_cpu_weights[layer_idx][dtype] yield name, cpu_buffer[offset : offset + numel].reshape(shape) def register_forward_hooks(self) -> None: if not self.enabled: return layers = getattr(self.model, self.layers_attr_str) def make_pre_hook(i): def hook(module, input): # wait only for the current layer if it's being prefetched if i == 0: self.prepare_for_next_req(non_blocking=False) if i in self._prefetch_events: torch.get_device_module().current_stream().wait_event( self._prefetch_events[i] ) # trigger batch prefetch (i + prefetch_size ~ i + 2 * prefetch_size) if needed if i % self.prefetch_size == 0: for j in range(i + self.prefetch_size, i + 2 * self.prefetch_size): layer_to_prefetch = j % self.num_layers self.prefetch_layer(layer_to_prefetch, non_blocking=True) return hook def make_post_hook(i): def hook(module, input, output): # previous, we wait here, until the copy stream for next layer is finished, # now with any prefetch_size, only wait for the copy stream, when the copy stream is for the next layer self.release_layer(i) return hook # register prefetch & release hooks for each layer self._forward_hooks.clear() for i, layer in enumerate(layers): pre_hook_handle = layer.register_forward_pre_hook(make_pre_hook(i)) post_hook_handle = layer.register_forward_hook(make_post_hook(i)) self._forward_hooks.extend([pre_hook_handle, post_hook_handle]) def remove_forward_hooks(self) -> None: """Remove all registered forward hooks.""" for hook_handle in self._forward_hooks: hook_handle.remove() self._forward_hooks.clear() class OffloadableDiTMixin: """ A mixin that registers forward hooks for a DiT to enable layerwise offload """ # the list of names of a DiT's layers/blocks layer_names: List[str] layerwise_offload_managers: list[LayerwiseOffloadManager] = [] def configure_layerwise_offload(self, server_args: ServerArgs): self.layerwise_offload_managers = [] for layer_name in self.layer_names: # a manager per layer-list module_list = getattr(self, layer_name, None) if module_list is None or not isinstance(module_list, torch.nn.ModuleList): continue num_layers = len(module_list) if server_args.dit_offload_prefetch_size < 1.0: prefetch_size = 1 + int( round(server_args.dit_offload_prefetch_size * (num_layers - 1)) ) else: prefetch_size = int(server_args.dit_offload_prefetch_size) manager = LayerwiseOffloadManager( model=self, layers_attr_str=layer_name, num_layers=num_layers, enabled=True, pin_cpu_memory=server_args.pin_cpu_memory, prefetch_size=prefetch_size, ) self.layerwise_offload_managers.append(manager) logger.info( f"Enabled layerwise offload for {self.__class__.__name__} on modules: {self.layer_names}" ) def prepare_for_next_req(self): if self.layerwise_offload_managers is None: return for manager in self.layerwise_offload_managers: manager.prepare_for_next_req(non_blocking=True) def disable_offload(self) -> None: """Disable layerwise offload: load all layers to GPU and remove hooks.""" if self.layerwise_offload_managers is None: return for manager in self.layerwise_offload_managers: if manager.enabled: manager.remove_forward_hooks() manager.load_all_layers() def enable_offload(self) -> None: """Re-enable layerwise offload: sync weights to CPU, release layers, and restore hooks.""" if self.layerwise_offload_managers is None: return for manager in self.layerwise_offload_managers: if manager.enabled: manager.sync_all_layers_to_cpu() manager.release_all() manager.register_forward_hooks() def iter_materialized_weights(module: torch.nn.Module): """Yield (name, tensor) pairs with materialized weights, even under offload. When layerwise offload is active, module.named_parameters() returns (1,) placeholders for offloaded layers. This function reads the actual data from the offload manager's CPU buffers and chains it with the non-offloaded parameters. """ offload_managers: list = [] if isinstance(module, OffloadableDiTMixin) and module.layerwise_offload_managers: offload_managers = [m for m in module.layerwise_offload_managers if m.enabled] if not offload_managers: yield from module.named_parameters() return # Collect offloaded names and their real tensors from CPU buffers. offloaded_names: set[str] = set() for manager in offload_managers: for name, tensor in manager.iter_cpu_weights(): offloaded_names.add(name) yield name, tensor # Yield non-offloaded parameters (e.g. final norms, embeddings). for name, param in module.named_parameters(): if name not in offloaded_names: yield name, param ================================================ FILE: python/sglang/multimodal_gen/runtime/utils/logging_utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/logger.py """Logging configuration for sglang.multimodal_gen.""" import argparse import contextlib import dataclasses import datetime import inspect import logging import os import sys import time from contextlib import contextmanager from enum import Enum from functools import lru_cache, partial from logging import Logger from types import MethodType from typing import Any, cast import sglang.multimodal_gen.envs as envs SGLANG_DIFFUSION_LOGGING_LEVEL = envs.SGLANG_DIFFUSION_LOGGING_LEVEL SGLANG_DIFFUSION_LOGGING_PREFIX = envs.SGLANG_DIFFUSION_LOGGING_PREFIX # color CYAN = "\033[1;36m" RED = "\033[91m" GREEN = "\033[92m" YELLOW = "\033[93m" RESET = "\033[0;0m" _FORMAT = ( f"{SGLANG_DIFFUSION_LOGGING_PREFIX}%(levelname)s %(asctime)s " "[%(filename)s: %(lineno)d] %(message)s" ) # _FORMAT = "[%(asctime)s] %(message)s" _DATE_FORMAT = "%m-%d %H:%M:%S" DEFAULT_LOGGING_CONFIG = { "formatters": { "sgl_diffusion": { "class": "sglang.multimodal_gen.runtime.utils.logging_utils.ColoredFormatter", "datefmt": _DATE_FORMAT, "format": _FORMAT, }, }, "handlers": { "sgl_diffusion": { "class": "logging.StreamHandler", "formatter": "sgl_diffusion", "level": SGLANG_DIFFUSION_LOGGING_LEVEL, "stream": "ext://sys.stdout", }, }, "loggers": { "sgl_diffusion": { "handlers": ["sgl_diffusion"], "level": "WARNING", "propagate": False, }, }, "root": { "handlers": ["sgl_diffusion"], "level": "DEBUG", }, "version": 1, "disable_existing_loggers": False, } class ColoredFormatter(logging.Formatter): """A logging formatter that adds color to log levels.""" LEVEL_COLORS = { logging.ERROR: RED, logging.WARNING: YELLOW, } def format(self, record: logging.LogRecord) -> str: """Adds color to the log""" formatted_message = super().format(record) color = self.LEVEL_COLORS.get(record.levelno) if color: formatted_message = f"{color}{formatted_message}{RESET}" return formatted_message class SortedHelpFormatter(argparse.HelpFormatter): """SortedHelpFormatter that sorts arguments by their option strings.""" def add_arguments(self, actions): actions = sorted(actions, key=lambda x: x.option_strings) super().add_arguments(actions) @lru_cache def _print_info_once(logger: Logger, msg: str) -> None: # Set the stacklevel to 2 to print the original caller's line info logger.info(msg, stacklevel=2) @lru_cache def _print_warning_once(logger: Logger, msg: str) -> None: # Set the stacklevel to 2 to print the original caller's line info logger.warning(msg, stacklevel=2) def get_is_main_process(): try: rank = int(os.environ["RANK"]) except (KeyError, ValueError): rank = 0 return rank == 0 def get_is_local_main_process(): try: rank = int(os.environ["LOCAL_RANK"]) except (KeyError, ValueError): rank = 0 return rank == 0 def _log_process_aware( server_log_level: int, level: int, logger_self: Logger, msg: object, *args: Any, main_process_only: bool, local_main_process_only: bool, **kwargs: Any, ) -> None: """Helper function to log a message if the process rank matches the criteria.""" is_main_process = get_is_main_process() is_local_main_process = get_is_local_main_process() should_log = ( not main_process_only and not local_main_process_only or (main_process_only and is_main_process) or (local_main_process_only and is_local_main_process) or server_log_level <= logging.DEBUG ) if should_log: # stacklevel=3 to show the original caller's location, # as this function is called by the patched methods. if "stacklevel" in kwargs: logger_self.log(level, msg, *args, **kwargs) else: logger_self.log(level, msg, *args, stacklevel=3, **kwargs) class _SGLDiffusionLogger(Logger): """ Note: This class is just to provide type information. We actually patch the methods directly on the :class:`logging.Logger` instance to avoid conflicting with other libraries such as `intel_extension_for_pytorch.utils._logger`. """ def info_once(self, msg: str) -> None: """ As :meth:`info`, but subsequent calls with the same message are silently dropped. """ _print_info_once(self, msg) def warning_once(self, msg: str) -> None: """ As :meth:`warning`, but subsequent calls with the same message are silently dropped. """ _print_warning_once(self, msg) def info( # type: ignore[override] self, msg: object, *args: Any, main_process_only: bool = True, local_main_process_only: bool = True, **kwargs: Any, ) -> None: ... def debug( # type: ignore[override] self, msg: object, *args: Any, main_process_only: bool = True, local_main_process_only: bool = True, **kwargs: Any, ) -> None: ... def warning( # type: ignore[override] self, msg: object, *args: Any, main_process_only: bool = False, local_main_process_only: bool = True, **kwargs: Any, ) -> None: ... def error( # type: ignore[override] self, msg: object, *args: Any, main_process_only: bool = False, local_main_process_only: bool = True, **kwargs: Any, ) -> None: ... def init_logger(name: str) -> _SGLDiffusionLogger: """The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root sgl_diffusion logger has already been configured.""" logger = logging.getLogger(name) server_log_level = logger.getEffectiveLevel() # Patch instance methods setattr(logger, "info_once", MethodType(_print_info_once, logger)) setattr(logger, "warning_once", MethodType(_print_warning_once, logger)) def _create_patched_method( level: int, main_process_only_default: bool, local_main_process_only_default: bool, ): def _method( self: Logger, msg: object, *args: Any, main_process_only: bool = main_process_only_default, local_main_process_only: bool = local_main_process_only_default, **kwargs: Any, ) -> None: _log_process_aware( server_log_level, level, self, msg, *args, main_process_only=main_process_only, local_main_process_only=local_main_process_only, **kwargs, ) return _method setattr( logger, "info", MethodType(_create_patched_method(logging.INFO, True, True), logger), ) setattr( logger, "debug", MethodType(_create_patched_method(logging.DEBUG, True, True), logger), ) setattr( logger, "warning", MethodType(_create_patched_method(logging.WARNING, False, True), logger), ) setattr( logger, "error", MethodType(_create_patched_method(logging.ERROR, False, False), logger), ) return cast(_SGLDiffusionLogger, logger) logger = init_logger(__name__) def _is_torch_tensor(obj: Any) -> tuple[bool, Any]: """Return (is_tensor, torch_module_or_None) without importing torch at module import time.""" try: import torch # type: ignore return isinstance(obj, torch.Tensor), torch except Exception: return False, None def _sanitize_for_logging(obj: Any, key_hint: str | None = None) -> Any: """Recursively convert objects to JSON-serializable forms for concise logging. Rules: - Drop any field/dict key named 'param_names_mapping'. - Render Enums using their value. - Render torch.Tensor as a compact summary; if key name is 'scaling_factor', include stats. - Dataclasses are expanded to dicts and sanitized recursively. - Callables/functions are rendered as their qualified name. - Redact sensitive fields like 'prompt' and 'negative_prompt' (only show length). - Fallback to str(...) for unknown types. """ if obj is None or isinstance(obj, (str, int, float, bool)): if key_hint in ("prompt", "negative_prompt"): if isinstance(obj, str): return f"" return obj if isinstance(obj, Enum): return obj.value is_tensor, torch_mod = _is_torch_tensor(obj) if is_tensor: try: ten = obj.detach().cpu() if key_hint == "scaling_factor": stats = { "shape": list(ten.shape), "dtype": str(ten.dtype), } try: stats["min"] = float(ten.min().item()) except Exception: pass try: stats["max"] = float(ten.max().item()) except Exception: pass try: stats["mean"] = float(ten.float().mean().item()) except Exception: pass return {"tensor": "scaling_factor", **stats} return {"tensor": True, "shape": list(ten.shape), "dtype": str(ten.dtype)} except Exception: return "" if dataclasses.is_dataclass(obj): result: dict[str, Any] = {} for f in dataclasses.fields(obj): if not f.repr: continue name = f.name if "names_mapping" in name: continue try: value = getattr(obj, name) except Exception: continue result[name] = _sanitize_for_logging(value, key_hint=name) return result if isinstance(obj, dict): result_dict: dict[str, Any] = {} for k, v in obj.items(): try: key_str = str(k) except Exception: key_str = "" if key_str == "param_names_mapping": continue result_dict[key_str] = _sanitize_for_logging(v, key_hint=key_str) return result_dict if isinstance(obj, (list, tuple, set)): return [_sanitize_for_logging(x, key_hint=key_hint) for x in obj] try: if inspect.isroutine(obj) or inspect.isclass(obj): module = getattr(obj, "__module__", "") qn = getattr(obj, "__qualname__", getattr(obj, "__name__", "")) return f"{module}.{qn}" if module else qn except Exception: pass try: return str(obj) except Exception: return "" def _trace_calls(log_path, root_dir, frame, event, arg=None): if event in ["call", "return"]: # Extract the filename, line number, function name, and the code object filename = frame.f_code.co_filename lineno = frame.f_lineno func_name = frame.f_code.co_name if not filename.startswith(root_dir): # only log the functions in the sgl_diffusion root_dir return # Log every function call or return try: last_frame = frame.f_back if last_frame is not None: last_filename = last_frame.f_code.co_filename last_lineno = last_frame.f_lineno last_func_name = last_frame.f_code.co_name else: # initial frame last_filename = "" last_lineno = 0 last_func_name = "" with open(log_path, "a") as f: ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f") if event == "call": f.write( f"{ts} Call to" f" {func_name} in {filename}:{lineno}" f" from {last_func_name} in {last_filename}:" f"{last_lineno}\n" ) else: f.write( f"{ts} Return from" f" {func_name} in {filename}:{lineno}" f" to {last_func_name} in {last_filename}:" f"{last_lineno}\n" ) except NameError: # modules are deleted during shutdown pass return partial(_trace_calls, log_path, root_dir) def enable_trace_function_call(log_file_path: str, root_dir: str | None = None): """ Enable tracing of every function call in code under `root_dir`. This is useful for debugging hangs or crashes. `log_file_path` is the path to the log file. `root_dir` is the root directory of the code to trace. If None, it is the sgl_diffusion root directory. Note that this call is thread-level, any threads calling this function will have the trace enabled. Other threads will not be affected. """ logger.warning( "SGLANG_DIFFUSION_TRACE_FUNCTION is enabled. It will record every" " function executed by Python. This will slow down the code. It " "is suggested to be used for debugging hang or crashes only." ) logger.info("Trace frame log is saved to %s", log_file_path) if root_dir is None: # by default, this is the sgl_diffusion root directory root_dir = os.path.dirname(os.path.dirname(__file__)) sys.settrace(partial(_trace_calls, log_file_path, root_dir)) def set_uvicorn_logging_configs(): from uvicorn.config import LOGGING_CONFIG LOGGING_CONFIG["formatters"]["default"][ "fmt" ] = "[%(asctime)s] %(levelprefix)s %(message)s" LOGGING_CONFIG["formatters"]["default"]["datefmt"] = "%Y-%m-%d %H:%M:%S" LOGGING_CONFIG["formatters"]["access"][ "fmt" ] = '[%(asctime)s] %(levelprefix)s %(client_addr)s - "%(request_line)s" %(status_code)s' LOGGING_CONFIG["formatters"]["access"]["datefmt"] = "%Y-%m-%d %H:%M:%S" def configure_logger(server_args, prefix: str = ""): log_format = f"[%(asctime)s{prefix}] %(message)s" datefmt = "%m-%d %H:%M:%S" formatter = ColoredFormatter(log_format, datefmt=datefmt) handler = logging.StreamHandler(sys.stdout) handler.setFormatter(formatter) root = logging.getLogger() root.handlers.clear() root.addHandler(handler) root.setLevel(getattr(logging, server_args.log_level.upper())) set_uvicorn_logging_configs() @lru_cache(maxsize=1) def get_log_level() -> int: root = logging.getLogger() return root.level def suppress_loggers(loggers_to_suppress: list[str], level: int = logging.WARNING): original_levels = {} for logger_name in loggers_to_suppress: logger = logging.getLogger(logger_name) original_levels[logger_name] = logger.level logger.setLevel(level) return original_levels def globally_suppress_loggers(): # globally suppress some obsessive loggers target_names = [ "imageio", "imageio_ffmpeg", "PIL", "PIL_Image", "python_multipart.multipart", "filelock", "urllib3", "httpx", "httpcore", ] for name in target_names: logging.getLogger(name).setLevel(logging.ERROR) # source: https://github.com/vllm-project/vllm/blob/a11f4a81e027efd9ef783b943489c222950ac989/vllm/utils/system_utils.py#L60 @contextlib.contextmanager def suppress_stdout(): """ Suppress stdout from C libraries at the file descriptor level. Only suppresses stdout, not stderr, to preserve error messages. Example: with suppress_stdout(): # C library calls that would normally print to stdout torch.distributed.new_group(ranks, backend="gloo") """ # Don't suppress if logging level is DEBUG stdout_fd = sys.stdout.fileno() stdout_dup = os.dup(stdout_fd) devnull_fd = os.open(os.devnull, os.O_WRONLY) try: sys.stdout.flush() os.dup2(devnull_fd, stdout_fd) yield finally: sys.stdout.flush() os.dup2(stdout_dup, stdout_fd) os.close(stdout_dup) os.close(devnull_fd) class GenerationTimer: def __init__(self): self.start_time = 0.0 self.end_time = 0.0 self.duration = 0.0 @contextmanager def log_generation_timer( logger: logging.Logger, prompt: str, request_idx: int | None = None, total_requests: int | None = None, ): if request_idx is not None and total_requests is not None: logger.info( "Processing prompt %d/%d: %s", request_idx, total_requests, _sanitize_for_logging(prompt, key_hint="prompt"), ) timer = GenerationTimer() timer.start_time = time.perf_counter() try: yield timer timer.end_time = time.perf_counter() timer.duration = timer.end_time - timer.start_time logger.info( f"Pixel data generated successfully in {GREEN}%.2f{RESET} seconds", timer.duration, ) except Exception as e: if request_idx is not None: logger.error( "Failed to generate output for prompt %d: %s", request_idx, e, exc_info=True, ) else: logger.error( f"Failed to generate output for prompt: {e}", exc_info=True, ) raise def log_batch_completion( logger: logging.Logger, num_outputs: int, total_time: float ) -> None: logger.info( f"Completed batch processing. Generated %d outputs in {GREEN}%.2f{RESET} seconds", num_outputs, total_time, ) ================================================ FILE: python/sglang/multimodal_gen/runtime/utils/mesh3d_utils.py ================================================ """Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2""" from __future__ import annotations import math from typing import Any, List, Optional, Tuple, Union import cv2 import numpy as np import torch import torch.nn.functional as F import trimesh from einops import rearrange, repeat from PIL import Image from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # Import C++ mesh processor extension from sglang.multimodal_gen.csrc.render.mesh_processor import meshVerticeInpaint def transform_pos( mtx: Union[np.ndarray, torch.Tensor], pos: torch.Tensor, keepdim: bool = False, ) -> torch.Tensor: """Transform positions by a matrix.""" t_mtx = torch.from_numpy(mtx).to(pos.device) if isinstance(mtx, np.ndarray) else mtx if pos.shape[-1] == 3: posw = torch.cat([pos, torch.ones([pos.shape[0], 1]).to(pos.device)], axis=1) else: posw = pos if keepdim: return torch.matmul(posw, t_mtx.t())[...] else: return torch.matmul(posw, t_mtx.t())[None, ...] def get_mv_matrix( elev: float, azim: float, camera_distance: float, center: Optional[np.ndarray] = None, ) -> np.ndarray: """Compute model-view matrix from camera parameters.""" elev = -elev azim += 90 elev_rad = math.radians(elev) azim_rad = math.radians(azim) camera_position = np.array( [ camera_distance * math.cos(elev_rad) * math.cos(azim_rad), camera_distance * math.cos(elev_rad) * math.sin(azim_rad), camera_distance * math.sin(elev_rad), ] ) if center is None: center = np.array([0, 0, 0]) else: center = np.array(center) lookat = center - camera_position lookat = lookat / np.linalg.norm(lookat) up = np.array([0, 0, 1.0]) right = np.cross(lookat, up) right = right / np.linalg.norm(right) up = np.cross(right, lookat) up = up / np.linalg.norm(up) c2w = np.concatenate( [np.stack([right, up, -lookat], axis=-1), camera_position[:, None]], axis=-1 ) w2c = np.zeros((4, 4)) w2c[:3, :3] = np.transpose(c2w[:3, :3], (1, 0)) w2c[:3, 3:] = -np.matmul(np.transpose(c2w[:3, :3], (1, 0)), c2w[:3, 3:]) w2c[3, 3] = 1.0 return w2c.astype(np.float32) def get_orthographic_projection_matrix( left: float = -1, right: float = 1, bottom: float = -1, top: float = 1, near: float = 0, far: float = 2, ) -> np.ndarray: """Compute orthographic projection matrix.""" ortho_matrix = np.eye(4, dtype=np.float32) ortho_matrix[0, 0] = 2 / (right - left) ortho_matrix[1, 1] = 2 / (top - bottom) ortho_matrix[2, 2] = -2 / (far - near) ortho_matrix[0, 3] = -(right + left) / (right - left) ortho_matrix[1, 3] = -(top + bottom) / (top - bottom) ortho_matrix[2, 3] = -(far + near) / (far - near) return ortho_matrix def get_perspective_projection_matrix( fovy: float, aspect_wh: float, near: float, far: float, ) -> np.ndarray: """Compute perspective projection matrix.""" fovy_rad = math.radians(fovy) return np.array( [ [1.0 / (math.tan(fovy_rad / 2.0) * aspect_wh), 0, 0, 0], [0, 1.0 / math.tan(fovy_rad / 2.0), 0, 0], [0, 0, -(far + near) / (far - near), -2.0 * far * near / (far - near)], [0, 0, -1, 0], ] ).astype(np.float32) def export_to_trimesh(mesh_output: Any) -> Any: """Convert mesh output to trimesh format.""" if isinstance(mesh_output, list): outputs = [] for mesh in mesh_output: if mesh is None: outputs.append(None) else: # Reverse face winding mesh.mesh_f = mesh.mesh_f[:, ::-1] mesh_obj = trimesh.Trimesh(mesh.mesh_v, mesh.mesh_f) outputs.append(mesh_obj) return outputs else: mesh_output.mesh_f = mesh_output.mesh_f[:, ::-1] return trimesh.Trimesh(mesh_output.mesh_v, mesh_output.mesh_f) def mesh_uv_wrap(mesh: Any) -> Any: """Apply UV unwrapping to mesh. In-place like native Hunyuan3D-2 for same layout.""" try: import xatlas except ImportError: logger.warning("xatlas not available, skipping UV unwrap") return mesh if isinstance(mesh, trimesh.Scene): mesh = mesh.dump(concatenate=True) if len(mesh.faces) > 500000000: raise ValueError( "The mesh has more than 500,000,000 faces, which is not supported." ) vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces) mesh.vertices = mesh.vertices[vmapping] mesh.faces = indices if not hasattr(mesh.visual, "uv"): mesh.visual = trimesh.visual.TextureVisuals( uv=uvs, material=trimesh.visual.material.SimpleMaterial() ) else: mesh.visual.uv = uvs return mesh def stride_from_shape(shape: Tuple[int, ...]) -> List[int]: """Compute stride from shape for scatter operations.""" stride = [1] for x in reversed(shape[1:]): stride.append(stride[-1] * x) return list(reversed(stride)) def scatter_add_nd_with_count( input: torch.Tensor, count: torch.Tensor, indices: torch.Tensor, values: torch.Tensor, weights: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Scatter add with counting for texture baking.""" D = indices.shape[-1] C = input.shape[-1] size = input.shape[:-1] stride = stride_from_shape(size) assert len(size) == D input = input.view(-1, C) count = count.view(-1, 1) flatten_indices = ( indices * torch.tensor(stride, dtype=torch.long, device=indices.device) ).sum(-1) if weights is None: weights = torch.ones_like(values[..., :1]) input.scatter_add_(0, flatten_indices.unsqueeze(1).repeat(1, C), values) count.scatter_add_(0, flatten_indices.unsqueeze(1), weights) return input.view(*size, C), count.view(*size, 1) def linear_grid_put_2d( H: int, W: int, coords: torch.Tensor, values: torch.Tensor, return_count: bool = False, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: """Put values into a 2D grid using linear interpolation.""" C = values.shape[-1] indices = coords * torch.tensor( [H - 1, W - 1], dtype=torch.float32, device=coords.device ) indices_00 = indices.floor().long() indices_00[:, 0].clamp_(0, H - 2) indices_00[:, 1].clamp_(0, W - 2) indices_01 = indices_00 + torch.tensor( [0, 1], dtype=torch.long, device=indices.device ) indices_10 = indices_00 + torch.tensor( [1, 0], dtype=torch.long, device=indices.device ) indices_11 = indices_00 + torch.tensor( [1, 1], dtype=torch.long, device=indices.device ) h = indices[..., 0] - indices_00[..., 0].float() w = indices[..., 1] - indices_00[..., 1].float() w_00 = (1 - h) * (1 - w) w_01 = (1 - h) * w w_10 = h * (1 - w) w_11 = h * w result = torch.zeros(H, W, C, device=values.device, dtype=values.dtype) count = torch.zeros(H, W, 1, device=values.device, dtype=values.dtype) weights = torch.ones_like(values[..., :1]) result, count = scatter_add_nd_with_count( result, count, indices_00, values * w_00.unsqueeze(1), weights * w_00.unsqueeze(1), ) result, count = scatter_add_nd_with_count( result, count, indices_01, values * w_01.unsqueeze(1), weights * w_01.unsqueeze(1), ) result, count = scatter_add_nd_with_count( result, count, indices_10, values * w_10.unsqueeze(1), weights * w_10.unsqueeze(1), ) result, count = scatter_add_nd_with_count( result, count, indices_11, values * w_11.unsqueeze(1), weights * w_11.unsqueeze(1), ) if return_count: return result, count mask = count.squeeze(-1) > 0 result[mask] = result[mask] / count[mask].repeat(1, C) return result class MeshRender: """Mesh renderer using CUDA rasterization for texture generation.""" def __init__( self, camera_distance: float = 1.45, camera_type: str = "orth", default_resolution: int = 1024, texture_size: int = 1024, bake_mode: str = "linear", device: str = "cuda", ): """Initialize the mesh renderer.""" self.device = device self.set_default_render_resolution(default_resolution) self.set_default_texture_resolution(texture_size) self.camera_distance = camera_distance self.camera_type = camera_type self.bake_angle_thres = 75 self.bake_unreliable_kernel_size = int( (2 / 512) * max(self.default_resolution[0], self.default_resolution[1]) ) self.bake_mode = bake_mode # Set up camera projection matrix if camera_type == "orth": self.ortho_scale = 1.2 self.camera_proj_mat = get_orthographic_projection_matrix( left=-self.ortho_scale * 0.5, right=self.ortho_scale * 0.5, bottom=-self.ortho_scale * 0.5, top=self.ortho_scale * 0.5, near=0.1, far=100, ) elif camera_type == "perspective": self.camera_proj_mat = get_perspective_projection_matrix( 49.13, self.default_resolution[1] / self.default_resolution[0], 0.01, 100.0, ) else: raise ValueError(f"Unknown camera type: {camera_type}") # Mesh data self.vtx_pos = None self.pos_idx = None self.vtx_uv = None self.uv_idx = None self.tex = None self.mesh_copy = None self.scale_factor = 1.0 def set_default_render_resolution( self, default_resolution: Union[int, Tuple[int, int]] ): """Set default rendering resolution.""" if isinstance(default_resolution, int): default_resolution = (default_resolution, default_resolution) self.default_resolution = default_resolution def set_default_texture_resolution(self, texture_size: Union[int, Tuple[int, int]]): """Set default texture resolution.""" if isinstance(texture_size, int): texture_size = (texture_size, texture_size) self.texture_size = texture_size def _rasterize( self, pos_clip: torch.Tensor, tri: torch.Tensor, resolution: Tuple[int, int], ) -> torch.Tensor: """Rasterize using CUDA rasterizer.""" from sglang.multimodal_gen.csrc.render.hunyuan3d_rasterizer import rasterize if pos_clip.dim() == 2: pos_clip = pos_clip.unsqueeze(0) findices, barycentric = rasterize(pos_clip, tri, resolution) rast_out = torch.cat((barycentric, findices.unsqueeze(-1).float()), dim=-1) rast_out = rast_out.unsqueeze(0) return rast_out def _interpolate( self, attr: torch.Tensor, rast_out: torch.Tensor, tri: torch.Tensor, ) -> torch.Tensor: """Interpolate vertex attributes.""" from sglang.multimodal_gen.csrc.render.hunyuan3d_rasterizer import interpolate barycentric = rast_out[0, ..., :-1] findices = rast_out[0, ..., -1].int() if attr.dim() == 2: attr = attr.unsqueeze(0) result = interpolate(attr, findices, barycentric, tri) return result def load_mesh( self, mesh: Union[trimesh.Trimesh, trimesh.Scene], scale_factor: float = 1.15, auto_center: bool = True, ): """Load a mesh for rendering.""" if isinstance(mesh, trimesh.Scene): mesh = mesh.dump(concatenate=True) self.mesh_copy = mesh.copy() vtx_pos = mesh.vertices.astype(np.float32) pos_idx = mesh.faces.astype(np.int32) # Get UV coordinates if available if hasattr(mesh.visual, "uv") and mesh.visual.uv is not None: vtx_uv = mesh.visual.uv.astype(np.float32) uv_idx = pos_idx.copy() else: vtx_uv = None uv_idx = None self.vtx_pos = torch.from_numpy(vtx_pos).to(self.device).float() self.pos_idx = torch.from_numpy(pos_idx).to(self.device).to(torch.int32) if vtx_uv is not None and uv_idx is not None: self.vtx_uv = torch.from_numpy(vtx_uv).to(self.device).float() self.uv_idx = torch.from_numpy(uv_idx).to(self.device).to(torch.int32) else: self.vtx_uv = None self.uv_idx = None # Coordinate transformation (Y-up to Z-up) self.vtx_pos[:, [0, 1]] = -self.vtx_pos[:, [0, 1]] self.vtx_pos[:, [1, 2]] = self.vtx_pos[:, [2, 1]] if self.vtx_uv is not None: self.vtx_uv[:, 1] = 1.0 - self.vtx_uv[:, 1] if auto_center: max_bb = (self.vtx_pos - 0).max(0)[0] min_bb = (self.vtx_pos - 0).min(0)[0] center = (max_bb + min_bb) / 2 scale = torch.norm(self.vtx_pos - center, dim=1).max() * 2.0 self.vtx_pos = (self.vtx_pos - center) * (scale_factor / float(scale)) self.scale_factor = scale_factor def save_mesh(self) -> trimesh.Trimesh: """Save mesh with current texture, reusing the original mesh object.""" texture_data = self.get_texture() texture_img = Image.fromarray((texture_data * 255).astype(np.uint8)) material = trimesh.visual.material.SimpleMaterial( image=texture_img, diffuse=(255, 255, 255) ) self.mesh_copy.visual = trimesh.visual.TextureVisuals( uv=self.mesh_copy.visual.uv, image=texture_img, material=material ) return self.mesh_copy def get_mesh(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Get mesh data with inverse coordinate transformation.""" vtx_pos = self.vtx_pos.cpu().numpy().copy() pos_idx = self.pos_idx.cpu().numpy() # Inverse coordinate transformation vtx_pos[:, [1, 2]] = vtx_pos[:, [2, 1]] vtx_pos[:, [0, 1]] = -vtx_pos[:, [0, 1]] if self.vtx_uv is not None: vtx_uv = self.vtx_uv.cpu().numpy().copy() vtx_uv[:, 1] = 1.0 - vtx_uv[:, 1] uv_idx = self.uv_idx.cpu().numpy() else: vtx_uv = None uv_idx = None return vtx_pos, pos_idx, vtx_uv, uv_idx def set_texture(self, tex: Union[np.ndarray, torch.Tensor, Image.Image]): """Set texture for the mesh.""" if isinstance(tex, np.ndarray): if tex.max() <= 1.0: tex = (tex * 255).astype(np.uint8) tex = Image.fromarray(tex.astype(np.uint8)) elif isinstance(tex, torch.Tensor): tex_np = tex.cpu().numpy() if tex_np.max() <= 1.0: tex_np = (tex_np * 255).astype(np.uint8) tex = Image.fromarray(tex_np.astype(np.uint8)) tex = tex.resize(self.texture_size).convert("RGB") tex = np.array(tex) / 255.0 self.tex = torch.from_numpy(tex).to(self.device).float() def get_texture(self) -> np.ndarray: """Get current texture as numpy array.""" if self.tex is None: return np.ones((*self.texture_size, 3), dtype=np.float32) return self.tex.cpu().numpy() def _get_pos_from_mvp( self, elev: float, azim: float, camera_distance: Optional[float] = None, center: Optional[np.ndarray] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Get camera-space and clip-space positions.""" proj = self.camera_proj_mat r_mv = get_mv_matrix( elev=elev, azim=azim, camera_distance=( self.camera_distance if camera_distance is None else camera_distance ), center=center, ) pos_camera = transform_pos(r_mv, self.vtx_pos, keepdim=True) pos_clip = transform_pos(proj, pos_camera) return pos_camera, pos_clip def render_normal( self, elev: float, azim: float, camera_distance: Optional[float] = None, center: Optional[np.ndarray] = None, resolution: Optional[Tuple[int, int]] = None, bg_color: List[float] = [1, 1, 1], use_abs_coor: bool = False, normalize_rgb: bool = True, return_type: str = "th", ) -> Union[torch.Tensor, np.ndarray, Image.Image]: """Render normal map from a viewpoint.""" pos_camera, pos_clip = self._get_pos_from_mvp( elev, azim, camera_distance, center ) if resolution is None: resolution = self.default_resolution if isinstance(resolution, (int, float)): resolution = (int(resolution), int(resolution)) rast_out = self._rasterize(pos_clip, self.pos_idx, resolution) # Compute face normals if use_abs_coor: mesh_triangles = self.vtx_pos[self.pos_idx[:, :3].long(), :] else: pos_camera_3d = pos_camera[:, :3] / pos_camera[:, 3:4] mesh_triangles = pos_camera_3d[self.pos_idx[:, :3].long(), :] face_normals = F.normalize( torch.cross( mesh_triangles[:, 1, :] - mesh_triangles[:, 0, :], mesh_triangles[:, 2, :] - mesh_triangles[:, 0, :], dim=-1, ), dim=-1, ) # Compute vertex normals vertex_normals = trimesh.geometry.mean_vertex_normals( vertex_count=self.vtx_pos.shape[0], faces=self.pos_idx.cpu().numpy(), face_normals=face_normals.cpu().numpy(), ) vertex_normals = ( torch.from_numpy(vertex_normals).float().to(self.device).contiguous() ) # Interpolate normals normal = self._interpolate(vertex_normals[None, ...], rast_out, self.pos_idx) # Apply visibility mask visible_mask = torch.clamp(rast_out[..., -1:], 0, 1) bg_tensor = torch.tensor(bg_color, dtype=torch.float32, device=self.device) normal = normal * visible_mask + bg_tensor * (1 - visible_mask) if normalize_rgb: normal = (normal + 1) * 0.5 image = normal[0, ...] if return_type == "np": image = image.cpu().numpy() elif return_type == "pl": image = image.cpu().numpy() * 255 image = Image.fromarray(image.astype(np.uint8)) return image def render_position( self, elev: float, azim: float, camera_distance: Optional[float] = None, center: Optional[np.ndarray] = None, resolution: Optional[Tuple[int, int]] = None, bg_color: List[float] = [1, 1, 1], return_type: str = "th", ) -> Union[torch.Tensor, np.ndarray, Image.Image]: """Render position map from a viewpoint.""" pos_camera, pos_clip = self._get_pos_from_mvp( elev, azim, camera_distance, center ) if resolution is None: resolution = self.default_resolution if isinstance(resolution, (int, float)): resolution = (int(resolution), int(resolution)) rast_out = self._rasterize(pos_clip, self.pos_idx, resolution) # Position colors (normalized vertex positions) tex_position = 0.5 - self.vtx_pos[:, :3] / self.scale_factor tex_position = tex_position.contiguous() # Interpolate positions position = self._interpolate(tex_position[None, ...], rast_out, self.pos_idx) # Apply visibility mask visible_mask = torch.clamp(rast_out[..., -1:], 0, 1) bg_tensor = torch.tensor(bg_color, dtype=torch.float32, device=self.device) position = position * visible_mask + bg_tensor * (1 - visible_mask) image = position[0, ...] if return_type == "np": image = image.cpu().numpy() elif return_type == "pl": image = image.cpu().numpy() * 255 image = Image.fromarray(image.astype(np.uint8)) return image def render_normal_multiview( self, camera_elevs: List[float], camera_azims: List[float], use_abs_coor: bool = True, ) -> List[Image.Image]: """Render normal maps from multiple viewpoints.""" normal_maps = [] for elev, azim in zip(camera_elevs, camera_azims): normal_map = self.render_normal( elev, azim, use_abs_coor=use_abs_coor, return_type="pl" ) normal_maps.append(normal_map) return normal_maps def render_position_multiview( self, camera_elevs: List[float], camera_azims: List[float], ) -> List[Image.Image]: """Render position maps from multiple viewpoints.""" position_maps = [] for elev, azim in zip(camera_elevs, camera_azims): position_map = self.render_position(elev, azim, return_type="pl") position_maps.append(position_map) return position_maps def _render_sketch_from_depth(self, depth_image: torch.Tensor) -> torch.Tensor: """Render sketch from depth using edge detection.""" depth_image_np = depth_image.cpu().numpy() depth_image_np = (depth_image_np * 255).astype(np.uint8) depth_edges = cv2.Canny(depth_image_np, 30, 80) sketch_image = ( torch.from_numpy(depth_edges).to(depth_image.device).float() / 255.0 ) sketch_image = sketch_image.unsqueeze(-1) return sketch_image def back_project( self, image: Union[Image.Image, np.ndarray, torch.Tensor], elev: float, azim: float, camera_distance: Optional[float] = None, center: Optional[np.ndarray] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Back-project an image onto mesh UV space.""" if isinstance(image, Image.Image): image = torch.tensor(np.array(image) / 255.0) elif isinstance(image, np.ndarray): image = torch.tensor(image) if image.dim() == 2: image = image.unsqueeze(-1) image = image.float().to(self.device) resolution = image.shape[:2] channel = image.shape[-1] pos_camera, pos_clip = self._get_pos_from_mvp( elev, azim, camera_distance, center ) rast_out = self._rasterize(pos_clip, self.pos_idx, resolution) visible_mask = torch.clamp(rast_out[..., -1:], 0, 1)[0, ...] # Compute vertex normals for angle-based weighting pos_camera_3d = pos_camera[:, :3] / pos_camera[:, 3:4] v0 = pos_camera_3d[self.pos_idx[:, 0].long(), :] v1 = pos_camera_3d[self.pos_idx[:, 1].long(), :] v2 = pos_camera_3d[self.pos_idx[:, 2].long(), :] face_normals = F.normalize(torch.cross(v1 - v0, v2 - v0, dim=-1), dim=-1) vertex_normals = trimesh.geometry.mean_vertex_normals( vertex_count=self.vtx_pos.shape[0], faces=self.pos_idx.cpu().numpy(), face_normals=face_normals.cpu().numpy(), ) vertex_normals = ( torch.from_numpy(vertex_normals).float().to(self.device).contiguous() ) # Interpolate normals and UVs normal = self._interpolate(vertex_normals[None, ...], rast_out, self.pos_idx) normal = normal[0, ...] if self.vtx_uv is not None: uv = self._interpolate(self.vtx_uv[None, ...], rast_out, self.uv_idx) else: # No UV coordinates texture = torch.zeros( self.texture_size[1], self.texture_size[0], channel, device=self.device ) cos_map = torch.zeros( self.texture_size[1], self.texture_size[0], 1, device=self.device ) boundary_map = torch.zeros_like(cos_map) return texture, cos_map, boundary_map # Compute depth for sketch tex_depth = pos_camera_3d[:, 2].reshape(1, -1, 1).contiguous() depth = self._interpolate(tex_depth, rast_out, self.pos_idx)[0, ...] depth_masked = depth[visible_mask > 0] if depth_masked.numel() > 0: depth_max, depth_min = depth_masked.max(), depth_masked.min() depth_normalized = (depth - depth_min) / (depth_max - depth_min + 1e-8) else: depth_normalized = depth depth_image = depth_normalized * visible_mask sketch_image = self._render_sketch_from_depth(depth_image) # Cosine weighting lookat = torch.tensor([[0, 0, -1]], device=self.device) cos_image = torch.nn.functional.cosine_similarity(lookat, normal.view(-1, 3)) cos_image = cos_image.view(normal.shape[0], normal.shape[1], 1) cos_thres = np.cos(self.bake_angle_thres / 180 * np.pi) cos_image[cos_image < cos_thres] = 0 # Shrink visible mask kernel_size = self.bake_unreliable_kernel_size * 2 + 1 kernel = torch.ones((1, 1, kernel_size, kernel_size), dtype=torch.float32).to( sketch_image.device ) visible_mask_proc = visible_mask.permute(2, 0, 1).unsqueeze(0).float() visible_mask_proc = F.conv2d( 1.0 - visible_mask_proc, kernel, padding=kernel_size // 2 ) visible_mask_proc = 1.0 - (visible_mask_proc > 0).float() visible_mask_proc = visible_mask_proc.squeeze(0).permute(1, 2, 0) sketch_proc = sketch_image.permute(2, 0, 1).unsqueeze(0) sketch_proc = F.conv2d(sketch_proc, kernel, padding=kernel_size // 2) sketch_proc = (sketch_proc > 0).float() sketch_proc = sketch_proc.squeeze(0).permute(1, 2, 0) visible_mask_proc = visible_mask_proc * (sketch_proc < 0.5) cos_image[visible_mask_proc == 0] = 0 # Linear baking proj_mask = (visible_mask_proc != 0).view(-1) uv_flat = uv.squeeze(0).contiguous().view(-1, 2)[proj_mask] image_flat = image.squeeze(0).contiguous().view(-1, channel)[proj_mask] cos_flat = cos_image.contiguous().view(-1, 1)[proj_mask] sketch_flat = sketch_image.contiguous().view(-1, 1)[proj_mask] texture = linear_grid_put_2d( self.texture_size[1], self.texture_size[0], uv_flat[..., [1, 0]], image_flat ) cos_map = linear_grid_put_2d( self.texture_size[1], self.texture_size[0], uv_flat[..., [1, 0]], cos_flat ) boundary_map = linear_grid_put_2d( self.texture_size[1], self.texture_size[0], uv_flat[..., [1, 0]], sketch_flat, ) return texture, cos_map, boundary_map def bake_from_multiview( self, views: List[Image.Image], camera_elevs: List[float], camera_azims: List[float], view_weights: List[float], method: str = "fast", ) -> Tuple[torch.Tensor, torch.Tensor]: """Bake texture from multiple views.""" project_textures, project_weighted_cos_maps = [], [] bake_exp = 4 for view, camera_elev, camera_azim, weight in zip( views, camera_elevs, camera_azims, view_weights ): project_texture, project_cos_map, _ = self.back_project( view, camera_elev, camera_azim ) project_cos_map = weight * (project_cos_map**bake_exp) project_textures.append(project_texture) project_weighted_cos_maps.append(project_cos_map) if method == "fast": texture, ori_trust_map = self.fast_bake_texture( project_textures, project_weighted_cos_maps ) else: raise ValueError(f"Unknown bake method: {method}") return texture, ori_trust_map > 1e-8 @torch.no_grad() def fast_bake_texture( self, textures: List[torch.Tensor], cos_maps: List[torch.Tensor], ) -> Tuple[torch.Tensor, torch.Tensor]: """Fast texture baking by weighted averaging.""" channel = textures[0].shape[-1] texture_merge = torch.zeros(self.texture_size + (channel,)).to(self.device) trust_map_merge = torch.zeros(self.texture_size + (1,)).to(self.device) for texture, cos_map in zip(textures, cos_maps): view_sum = (cos_map > 0).sum() painted_sum = ((cos_map > 0) * (trust_map_merge > 0)).sum() if view_sum > 0 and painted_sum / view_sum > 0.99: continue texture_merge += texture * cos_map trust_map_merge += cos_map texture_merge = texture_merge / torch.clamp(trust_map_merge, min=1e-8) texture_merge = texture_merge.clamp(0.0, 1.0) return texture_merge, trust_map_merge > 1e-8 def texture_inpaint( self, texture: torch.Tensor, mask: Union[torch.Tensor, np.ndarray], ) -> torch.Tensor: """Inpaint missing regions in UV texture using mesh-aware method.""" if isinstance(texture, torch.Tensor): texture_np = texture.cpu().numpy() else: texture_np = texture if isinstance(mask, torch.Tensor): mask_np = mask.cpu().numpy() else: mask_np = mask # Ensure proper format if texture_np.max() <= 1.0: texture_np = texture_np.astype(np.float32) else: texture_np = (texture_np / 255.0).astype(np.float32) if mask_np.ndim == 3: mask_np = mask_np.squeeze(-1) if mask_np.dtype == np.uint8: mask_uint8 = mask_np else: mask_uint8 = ((mask_np > 0) * 255).astype(np.uint8) # Get mesh data for mesh-aware inpainting vtx_pos, pos_idx, vtx_uv, uv_idx = self.get_mesh() if vtx_uv is not None and uv_idx is not None: texture_np, mask_uint8 = meshVerticeInpaint( texture_np, mask_uint8, vtx_pos, vtx_uv, pos_idx, uv_idx ) # Final OpenCV inpainting for remaining holes texture_uint8 = (texture_np * 255).astype(np.uint8) inpaint_mask = 255 - mask_uint8 texture_inpainted = cv2.inpaint(texture_uint8, inpaint_mask, 3, cv2.INPAINT_NS) return torch.from_numpy(texture_inpainted / 255.0).float().to(self.device) # Alias for compatibility uv_inpaint = texture_inpaint def array_to_tensor(np_array): """Convert numpy array to normalized tensor.""" image_pt = torch.tensor(np_array).float() image_pt = image_pt / 255 * 2 - 1 image_pt = rearrange(image_pt, "h w c -> c h w") image_pts = repeat(image_pt, "c h w -> b c h w", b=1) return image_pts def recenter_image(image, border_ratio=0.2): """Recenter a PIL image, cropping to non-transparent content with a border.""" from PIL import Image as PILImage if image.mode == "RGB": return image elif image.mode == "L": return image.convert("RGB") if image.mode != "RGBA": image = image.convert("RGBA") alpha_channel = np.array(image)[:, :, 3] non_zero_indices = np.argwhere(alpha_channel > 0) if non_zero_indices.size == 0: raise ValueError("Image is fully transparent") min_row, min_col = non_zero_indices.min(axis=0) max_row, max_col = non_zero_indices.max(axis=0) cropped_image = image.crop((min_col, min_row, max_col + 1, max_row + 1)) width, height = cropped_image.size border_width = int(width * border_ratio) border_height = int(height * border_ratio) new_width = width + 2 * border_width new_height = height + 2 * border_height square_size = max(new_width, new_height) new_image = PILImage.new("RGBA", (square_size, square_size), (255, 255, 255, 0)) paste_x = (square_size - new_width) // 2 + border_width paste_y = (square_size - new_height) // 2 + border_height new_image.paste(cropped_image, (paste_x, paste_y)) return new_image class ImageProcessorV2: """Image processor for Hunyuan3D single-view input.""" # External module path aliases for compatibility with Hunyuan3D configs _aliases = [ "hy3dshape.preprocessors.ImageProcessorV2", "hy3dgen.shapegen.preprocessors.ImageProcessorV2", ] def __init__(self, size=512, border_ratio=None): self.size = size self.border_ratio = border_ratio @staticmethod def recenter(image, border_ratio: float = 0.2): """recenter an image to leave some empty space at the image border.""" if image.shape[-1] == 4: mask = image[..., 3] else: mask = np.ones_like(image[..., 0:1]) * 255 image = np.concatenate([image, mask], axis=-1) mask = mask[..., 0] height, width, channels = image.shape size = max(height, width) result = np.zeros((size, size, channels), dtype=np.uint8) coords = np.nonzero(mask) x_min, x_max = coords[0].min(), coords[0].max() y_min, y_max = coords[1].min(), coords[1].max() crop_h = x_max - x_min crop_w = y_max - y_min if crop_h == 0 or crop_w == 0: raise ValueError("input image is empty") desired_size = int(size * (1 - border_ratio)) scale = desired_size / max(crop_h, crop_w) scaled_h = int(crop_h * scale) scaled_w = int(crop_w * scale) x2_min = (size - scaled_h) // 2 x2_max = x2_min + scaled_h y2_min = (size - scaled_w) // 2 y2_max = y2_min + scaled_w result[x2_min:x2_max, y2_min:y2_max] = cv2.resize( image[x_min:x_max, y_min:y_max], (scaled_w, scaled_h), interpolation=cv2.INTER_AREA, ) bg = np.ones((result.shape[0], result.shape[1], 3), dtype=np.uint8) * 255 mask = result[..., 3:].astype(np.float32) / 255 result = result[..., :3] * mask + bg * (1 - mask) mask = mask * 255 result = result.clip(0, 255).astype(np.uint8) mask = mask.clip(0, 255).astype(np.uint8) return result, mask def load_image(self, image, border_ratio=0.15, to_tensor=True): if isinstance(image, str): image = cv2.imread(image, cv2.IMREAD_UNCHANGED) image, mask = self.recenter(image, border_ratio=border_ratio) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) elif isinstance(image, Image.Image): image = image.convert("RGBA") image = np.asarray(image) image, mask = self.recenter(image, border_ratio=border_ratio) image = cv2.resize(image, (self.size, self.size), interpolation=cv2.INTER_CUBIC) mask = cv2.resize(mask, (self.size, self.size), interpolation=cv2.INTER_NEAREST) mask = mask[..., np.newaxis] if to_tensor: image = array_to_tensor(image) mask = array_to_tensor(mask) return image, mask def __call__(self, image, border_ratio=0.15, to_tensor=True, **kwargs): if self.border_ratio is not None: border_ratio = self.border_ratio image, mask = self.load_image( image, border_ratio=border_ratio, to_tensor=to_tensor ) outputs = {"image": image, "mask": mask} return outputs class MVImageProcessorV2(ImageProcessorV2): """Multi-view image processor for Hunyuan3D.""" # External module path aliases for compatibility with Hunyuan3D configs _aliases = [ "hy3dshape.preprocessors.MVImageProcessorV2", ] return_view_idx = True def __init__(self, size=512, border_ratio=None): super().__init__(size, border_ratio) self.view2idx = {"front": 0, "left": 1, "back": 2, "right": 3} def __call__(self, image_dict, border_ratio=0.15, to_tensor=True, **kwargs): if self.border_ratio is not None: border_ratio = self.border_ratio images = [] masks = [] view_idxs = [] for view_tag, image in image_dict.items(): view_idxs.append(self.view2idx[view_tag]) image, mask = self.load_image( image, border_ratio=border_ratio, to_tensor=to_tensor ) images.append(image) masks.append(mask) zipped_lists = zip(view_idxs, images, masks) sorted_zipped_lists = sorted(zipped_lists) view_idxs, images, masks = zip(*sorted_zipped_lists) image = torch.cat(images, 0).unsqueeze(0) mask = torch.cat(masks, 0).unsqueeze(0) outputs = {"image": image, "mask": mask, "view_idxs": view_idxs} return outputs # All tool classes available in this module for resolution TOOL_CLASSES = ( ImageProcessorV2, MVImageProcessorV2, ) def resolve_hunyuan3d_tool(target: str): """Resolve a Hunyuan3D tool class by target string.""" # First, try to match against _aliases for cls in TOOL_CLASSES: aliases = getattr(cls, "_aliases", []) if target in aliases: return cls # Then, try to match against class names for cls in TOOL_CLASSES: if cls.__name__ == target: return cls return None __all__ = [ "transform_pos", "get_mv_matrix", "get_orthographic_projection_matrix", "get_perspective_projection_matrix", "export_to_trimesh", "mesh_uv_wrap", "meshVerticeInpaint", "stride_from_shape", "scatter_add_nd_with_count", "linear_grid_put_2d", "MeshRender", "recenter_image", "array_to_tensor", "ImageProcessorV2", "MVImageProcessorV2", "TOOL_CLASSES", "resolve_hunyuan3d_tool", ] ================================================ FILE: python/sglang/multimodal_gen/runtime/utils/perf_logger.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import dataclasses import json import logging import os import subprocess import sys import time from datetime import datetime from functools import lru_cache from pathlib import Path from typing import Any, Dict, Optional import torch from dateutil.tz import UTC import sglang import sglang.multimodal_gen.envs as envs from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import ( CYAN, RESET, _SGLDiffusionLogger, get_is_main_process, init_logger, ) logger = init_logger(__name__) @dataclasses.dataclass class MemorySnapshot: allocated_mb: float # current allocated memory reserved_mb: float # current reserved memory (actual VRAM) peak_allocated_mb: float # peak allocated since last reset peak_reserved_mb: float # peak reserved since last reset def to_dict(self) -> Dict[str, Any]: return { "allocated_mb": round(self.allocated_mb, 2), "reserved_mb": round(self.reserved_mb, 2), "peak_allocated_mb": round(self.peak_allocated_mb, 2), "peak_reserved_mb": round(self.peak_reserved_mb, 2), } @dataclasses.dataclass class RequestMetrics: """Performance metrics for a single request, including timings and memory snapshots.""" def __init__(self, request_id: str): self.request_id = request_id self.stages: Dict[str, float] = {} self.steps: list[float] = [] self.total_duration_ms: float = 0.0 # memory tracking: {checkpoint_name: MemorySnapshot} self.memory_snapshots: Dict[str, MemorySnapshot] = {} @property def total_duration_s(self) -> float: return self.total_duration_ms / 1000.0 def record_stage(self, stage_name: str, duration_s: float): """Records the duration of a pipeline stage""" self.stages[stage_name] = duration_s * 1000 # Store as milliseconds def record_steps(self, index: int, duration_s: float): """Records the duration of a denoising step""" assert index == len(self.steps) self.steps.append(duration_s * 1000) def record_memory_snapshot(self, checkpoint_name: str, snapshot: MemorySnapshot): self.memory_snapshots[checkpoint_name] = snapshot def to_dict(self) -> Dict[str, Any]: """Serializes the metrics data to a dictionary.""" return { "request_id": self.request_id, "stages": self.stages, "steps": self.steps, "total_duration_ms": self.total_duration_ms, "memory_snapshots": { name: snapshot.to_dict() for name, snapshot in self.memory_snapshots.items() }, } def get_diffusion_perf_log_dir() -> str: """ Determines the directory for performance logs. """ log_dir = os.environ.get("SGLANG_PERF_LOG_DIR") if log_dir: return os.path.abspath(log_dir) if log_dir is None: sglang_path = Path(sglang.__file__).resolve() target_path = (sglang_path.parent / "../../.cache/logs").resolve() return str(target_path) return "" @lru_cache(maxsize=1) def get_git_commit_hash() -> str: try: commit_hash = os.environ.get("SGLANG_GIT_COMMIT") if not commit_hash: commit_hash = ( subprocess.check_output( ["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL ) .strip() .decode("utf-8") ) _CACHED_COMMIT_HASH = commit_hash return commit_hash except (subprocess.CalledProcessError, FileNotFoundError): _CACHED_COMMIT_HASH = "N/A" return "N/A" def capture_memory_snapshot() -> MemorySnapshot: if not torch.get_device_module().is_available(): return MemorySnapshot( allocated_mb=0.0, reserved_mb=0.0, peak_allocated_mb=0.0, peak_reserved_mb=0.0, ) allocated = torch.get_device_module().memory_allocated() reserved = torch.get_device_module().memory_reserved() peak_allocated = torch.get_device_module().max_memory_allocated() peak_reserved = torch.get_device_module().max_memory_reserved() return MemorySnapshot( allocated_mb=allocated / (1024**2), reserved_mb=reserved / (1024**2), peak_allocated_mb=peak_allocated / (1024**2), peak_reserved_mb=peak_reserved / (1024**2), ) @dataclasses.dataclass class RequestPerfRecord: request_id: str timestamp: str commit_hash: str tag: str stages: list[dict] steps: list[float] total_duration_ms: float memory_snapshots: dict[str, dict] = dataclasses.field(default_factory=dict) def __init__( self, request_id, commit_hash, tag, stages, steps, total_duration_ms, memory_snapshots=None, timestamp=None, ): self.request_id = request_id if timestamp is not None: self.timestamp = timestamp else: self.timestamp = datetime.now(UTC).isoformat() self.commit_hash = commit_hash self.tag = tag self.stages = stages self.steps = steps self.total_duration_ms = total_duration_ms self.memory_snapshots = memory_snapshots or {} class StageProfiler: """ A unified context manager, records performance metrics (usually of a single Stage or a step) into a provided RequestMetrics object (usually from a Req). """ def __init__( self, stage_name: str, logger: _SGLDiffusionLogger, metrics: Optional["RequestMetrics"], log_stage_start_end: bool = False, perf_dump_path_provided: bool = False, capture_memory: bool = False, ): self.stage_name = stage_name self.metrics = metrics self.logger = logger self.start_time = 0.0 self.log_timing = perf_dump_path_provided or envs.SGLANG_DIFFUSION_STAGE_LOGGING self.log_stage_start_end = log_stage_start_end self.capture_memory = capture_memory def __enter__(self): if self.log_stage_start_end: msg = f"[{self.stage_name}] started..." if self.logger.isEnabledFor(logging.DEBUG): msg += f" ({round(current_platform.get_available_gpu_memory(), 2)} GB left)" self.logger.info(msg) if (self.log_timing and self.metrics) or self.log_stage_start_end: if ( os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1" and self.stage_name.startswith("denoising_step_") and torch.get_device_module().is_available() ): torch.get_device_module().synchronize() self.start_time = time.perf_counter() return self def __exit__(self, exc_type, exc_val, exc_tb): if not ((self.log_timing and self.metrics) or self.log_stage_start_end): return False if ( os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1" and self.stage_name.startswith("denoising_step_") and torch.get_device_module().is_available() ): torch.get_device_module().synchronize() execution_time_s = time.perf_counter() - self.start_time if exc_type: self.logger.error( "[%s] Error during execution after %.4f ms: %s", self.stage_name, execution_time_s * 1000, exc_val, exc_info=True, ) return False if self.log_stage_start_end: self.logger.info( f"[{self.stage_name}] finished in {execution_time_s:.4f} seconds", ) if self.log_timing and self.metrics: if "denoising_step_" in self.stage_name: index = int(self.stage_name[len("denoising_step_") :]) self.metrics.record_steps(index, execution_time_s) else: self.metrics.record_stage(self.stage_name, execution_time_s) # capture memory snapshot after stage if requested if self.capture_memory and torch.get_device_module().is_available(): snapshot = capture_memory_snapshot() self.metrics.record_memory_snapshot( f"after_{self.stage_name}", snapshot ) return False class PerformanceLogger: """ A global utility class for logging performance metrics for all request, categorized by request-id. Serves both as a runtime logger (stream to file) and a dump utility. Notice that RequestMetrics stores the performance metrics of a single request """ @classmethod def dump_benchmark_report( cls, file_path: str, metrics: "RequestMetrics", meta: Optional[Dict[str, Any]] = None, tag: str = "benchmark_dump", ): """ Static method to dump a standardized benchmark report to a file. Eliminates duplicate logic in CLI/Client code. """ formatted_steps = [ {"name": name, "duration_ms": duration_ms} for name, duration_ms in metrics.stages.items() ] denoise_steps_ms = [ {"step": idx, "duration_ms": duration_ms} for idx, duration_ms in enumerate(metrics.steps) ] memory_checkpoints = { name: snapshot.to_dict() for name, snapshot in metrics.memory_snapshots.items() } report = { "timestamp": datetime.now(UTC).isoformat(), "request_id": metrics.request_id, "commit_hash": get_git_commit_hash(), "tag": tag, "total_duration_ms": metrics.total_duration_ms, "steps": formatted_steps, "denoise_steps_ms": denoise_steps_ms, "memory_checkpoints": memory_checkpoints, "meta": meta or {}, } try: abs_path = os.path.abspath(file_path) os.makedirs(os.path.dirname(abs_path), exist_ok=True) with open(abs_path, "w", encoding="utf-8") as f: json.dump(report, f, indent=2) logger.info(f"Metrics dumped to: {CYAN}{abs_path}{RESET}") except IOError as e: logger.error(f"Failed to dump metrics to {abs_path}: {e}") @classmethod def log_request_summary( cls, metrics: "RequestMetrics", tag: str = "total_inference_time", ): """logs the stage metrics and total duration for a completed request to the performance_log file. Note that this accords to the time spent internally in server, postprocess is not included """ formatted_stages = [ {"name": name, "execution_time_ms": duration_ms} for name, duration_ms in metrics.stages.items() ] memory_checkpoints = { name: snapshot.to_dict() for name, snapshot in metrics.memory_snapshots.items() } record = RequestPerfRecord( metrics.request_id, commit_hash=get_git_commit_hash(), tag="pipeline_stage_metrics", stages=formatted_stages, steps=metrics.steps, total_duration_ms=metrics.total_duration_ms, memory_snapshots=memory_checkpoints, ) try: if get_is_main_process(): log_dir = get_diffusion_perf_log_dir() if not os.path.exists(log_dir): os.makedirs(log_dir, exist_ok=True) log_file = os.path.join(log_dir, "performance.log") with open(log_file, "a", encoding="utf-8") as f: f.write(json.dumps(dataclasses.asdict(record)) + "\n") except (OSError, PermissionError) as e: print(f"WARNING: Failed to log performance record: {e}", file=sys.stderr) ================================================ FILE: python/sglang/multimodal_gen/runtime/utils/profiler.py ================================================ import gzip import os import torch from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_logger if current_platform.is_npu(): import torch_npu patches = [ ["profiler.profile", torch_npu.profiler.profile], ["profiler.schedule", torch_npu.profiler.schedule], ] torch_npu._apply_patches(patches) logger = init_logger(__name__) class SGLDiffusionProfiler: """ A wrapper around torch.profiler to simplify usage in pipelines. Supports both full profiling and scheduled profiling. 1. if profile_all_stages is on: profile all stages, including all denoising steps 2. otherwise, if num_profiled_timesteps is specified: profile {num_profiled_timesteps} denoising steps. profile all steps if num_profiled_timesteps==-1 """ _instance = None def __init__( self, request_id: str | None = None, rank: int = 0, full_profile: bool = False, num_steps: int | None = None, num_inference_steps: int | None = None, log_dir: str | None = None, ): self.request_id = request_id or "profile_trace" self.rank = rank self.full_profile = full_profile self.log_dir = ( log_dir if log_dir is not None else os.getenv("SGLANG_TORCH_PROFILER_DIR", "./logs") ) try: os.makedirs(self.log_dir, exist_ok=True) except OSError: pass activities = [torch.profiler.ProfilerActivity.CPU] if torch.cuda.is_available() or ( hasattr(torch, "musa") and torch.musa.is_available() ): activities.append(torch.profiler.ProfilerActivity.CUDA) if current_platform.is_npu(): activities.append(torch_npu.profiler.ProfilerActivity.NPU) common_torch_profiler_args = dict( activities=activities, record_shapes=True, with_stack=True, on_trace_ready=( None if not current_platform.is_npu() else torch_npu.profiler.tensorboard_trace_handler(self.log_dir) ), ) if self.full_profile: # profile all stages self.profiler = torch.profiler.profile(**common_torch_profiler_args) self.profile_mode_id = "full stages" else: # profile denoising stage only warmup = 1 num_actual_steps = num_inference_steps if num_steps == -1 else num_steps self.num_active_steps = num_actual_steps + warmup self.profiler = torch.profiler.profile( **common_torch_profiler_args, schedule=torch.profiler.schedule( skip_first=0, wait=0, warmup=warmup, active=self.num_active_steps, repeat=1, ), ) self.profile_mode_id = f"{num_actual_steps} steps" logger.info(f"Profiling request: {request_id} for {self.profile_mode_id}...") self.has_stopped = False SGLDiffusionProfiler._instance = self self.start() def start(self): logger.info("Starting Profiler...") self.profiler.start() def _step(self): self.profiler.step() def step_stage(self): if self.full_profile: self._step() def step_denoising_step(self): if not self.full_profile: if self.num_active_steps >= 0: self._step() self.num_active_steps -= 1 else: # early exit when enough steps are captured, to reduce the trace file size self.stop(dump_rank=0) @classmethod def get_instance(cls) -> "SGLDiffusionProfiler": return cls._instance def stop(self, export_trace: bool = True, dump_rank: int | None = None): if self.has_stopped: return self.has_stopped = True logger.info("Stopping Profiler...") if torch.cuda.is_available() or ( hasattr(torch, "musa") and torch.musa.is_available() ): torch.cuda.synchronize() if current_platform.is_npu(): torch.npu.synchronize() export_trace = False # set to false because our internal torch_npu.profiler will generate trace file self.profiler.stop() if export_trace: if dump_rank is not None and dump_rank != self.rank: pass else: self._export_trace() SGLDiffusionProfiler._instance = None def _export_trace(self): try: os.makedirs(self.log_dir, exist_ok=True) sanitized_profile_mode_id = self.profile_mode_id.replace(" ", "_") trace_path = os.path.abspath( os.path.join( self.log_dir, f"{self.request_id}-{sanitized_profile_mode_id}-global-rank{self.rank}.trace.json.gz", ) ) self.profiler.export_chrome_trace(trace_path) if self._check_trace_integrity(trace_path): logger.info(f"Saved profiler traces to: {CYAN}{trace_path}{RESET}") else: logger.warning(f"Trace file may be corrupted: {trace_path}") except Exception as e: logger.error(f"Failed to save trace: {e}") def _check_trace_integrity(self, trace_path: str) -> bool: try: if not os.path.exists(trace_path) or os.path.getsize(trace_path) == 0: return False with gzip.open(trace_path, "rb") as f: content = f.read() if content.count(b"\x1f\x8b") > 1: logger.warning("Multiple gzip headers detected") return False return True except Exception as e: logger.warning(f"Trace file integrity check failed: {e}") return False ================================================ FILE: python/sglang/multimodal_gen/runtime/utils/quantization_utils.py ================================================ import glob import json import os from pathlib import Path from typing import Dict, List, Optional from safetensors import safe_open from sglang.multimodal_gen.runtime.layers.quantization import ( QuantizationConfig, get_quantization_config, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) def find_quant_modelslim_config(model_config, component_model_path): quant_config_file = Path(component_model_path, "quant_model_description.json") quant_cfg = None if quant_config_file.is_file(): with open(quant_config_file) as f: quant_cfg = json.load(f) # This field is required for flagless model loading but is not present in # modelslim model description, so we're adding it here manually. quant_cfg["quant_method"] = "modelslim" return quant_cfg def replace_prefix(key: str, prefix_mapping: dict[str, str]) -> str: for prefix, new_prefix in prefix_mapping.items(): if key.startswith(prefix): key = key.replace(prefix, new_prefix, 1) return key def get_quant_config( model_config, component_model_path: str, packed_modules_mapping: Dict[str, List[str]] = {}, remap_prefix: Dict[str, str] | None = None, ) -> QuantizationConfig: quant_cfg = find_quant_modelslim_config(model_config, component_model_path) if quant_cfg is not None: quant_cls = get_quantization_config(quant_cfg["quant_method"]) return quant_cls.from_config(quant_cfg) else: if "quantization_config" not in model_config: return None quant_cls = get_quantization_config( model_config["quantization_config"]["quant_method"] ) # GGUF doesn't have config file if model_config["quantization_config"]["quant_method"] == "gguf": return quant_cls.from_config({}) # Read the quantization config from the HF model config, if available. hf_quant_config = model_config["quantization_config"] # some vision model may keep quantization_config in their text_config hf_text_config = getattr(model_config, "text_config", None) if hf_quant_config is None and hf_text_config is not None: hf_quant_config = getattr(hf_text_config, "quantization_config", None) if hf_quant_config is None: # compressed-tensors uses a compressions_config hf_quant_config = getattr(model_config, "compression_config", None) if hf_quant_config is not None: hf_quant_config["packed_modules_mapping"] = packed_modules_mapping return quant_cls.from_config(hf_quant_config) # In case of bitsandbytes/QLoRA, get quant config from the adapter model. else: model_name_or_path = model_config["model_path"] is_local = os.path.isdir(model_name_or_path) hf_folder = model_name_or_path possible_config_filenames = quant_cls.get_config_filenames() # If the quantization config is not found, use the default config. if not possible_config_filenames: return quant_cls() config_files = glob.glob(os.path.join(hf_folder, "*.json")) quant_config_files = [ f for f in config_files if any(f.endswith(x) for x in possible_config_filenames) ] if len(quant_config_files) == 0: raise ValueError( f"Cannot find the config file for {model_config['quantization_config']['quant_method']}" ) if len(quant_config_files) > 1: raise ValueError( f"Found multiple config files for {model_config['quantization_config']['quant_method']}: " f"{quant_config_files}" ) quant_config_file = quant_config_files[0] with open(quant_config_file) as f: config = json.load(f) if remap_prefix is not None: exclude_modules = [ replace_prefix(key, remap_prefix) for key in config["quantization"]["exclude_modules"] ] config["quantization"]["exclude_modules"] = exclude_modules config["packed_modules_mapping"] = packed_modules_mapping return quant_cls.from_config(config) def handle_fp8_metadata_format(quant_config_dict): layers = quant_config_dict.get("layers", {}) if any( isinstance(v, dict) and "float8" in v.get("format", "") for v in layers.values() ): quant_config_dict["quant_method"] = "fp8" quant_config_dict["activation_scheme"] = "dynamic" return quant_config_dict def get_quant_config_from_safetensors_metadata( file_path: str, ) -> Optional[QuantizationConfig]: """Extract quantization config from a safetensors file's metadata header. Returns None if no recognizable quantization metadata is found. """ metadata = get_metadata_from_safetensors_file(file_path) if not metadata: return None quant_config_str = metadata.get("_quantization_metadata") if not quant_config_str: return None try: quant_config_dict = json.loads(quant_config_str) except Exception as _e: return None # handle diffusers fp8 safetensors metadata format if ( "quant_method" not in quant_config_dict and "format_version" in quant_config_dict and "layers" in quant_config_dict ): quant_config_dict = handle_fp8_metadata_format(quant_config_dict) quant_method = quant_config_dict.get("quant_method") if not quant_method: return None try: quant_cls = get_quantization_config(quant_method) config = quant_cls.from_config(quant_config_dict) logger.debug(f"Get quantization config from safetensors file: {file_path}") return config except Exception as _e: return None def get_metadata_from_safetensors_file(file_path: str): try: with safe_open(file_path, framework="pt", device="cpu") as f: metadata = f.metadata() return metadata except Exception as e: logger.warning(e) ================================================ FILE: python/sglang/multimodal_gen/test/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo ================================================ FILE: python/sglang/multimodal_gen/test/cli/test_generate_common.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo """ Common generate cli test, one test for image and video each """ import dataclasses import os import shlex import subprocess import sys import unittest from typing import Optional from PIL import Image from sglang.multimodal_gen.configs.sample.sampling_params import DataType from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.test.test_utils import check_image_size logger = init_logger(__name__) @dataclasses.dataclass class TestResult: name: str key: str succeed: bool def run_command(command) -> Optional[float]: """Runs a command and returns the execution time and status.""" print(f"Running command: {shlex.join(command)}") with subprocess.Popen( command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, encoding="utf-8", ) as process: for line in process.stdout: sys.stdout.write(line) process.wait() if process.returncode == 0: return True print(f"Command failed with exit code {process.returncode}") return False class CLIBase(unittest.TestCase): model_path: str = None extra_args = [] data_type: DataType = None # tested on h100 width: int = 720 height: int = 720 output_path: str = "test_outputs" def setUp(self): super().setUp() if not os.path.exists(self.output_path): os.makedirs(self.output_path, exist_ok=True) if os.path.exists(self.output_path): for f in os.listdir(self.output_path): path = os.path.join(self.output_path, f) if os.path.isfile(path): os.remove(path) def tearDown(self): super().tearDown() if os.path.exists(self.output_path): for f in os.listdir(self.output_path): path = os.path.join(self.output_path, f) if os.path.isfile(path): os.remove(path) def get_base_command(self): return [ "sglang", "generate", "--prompt", "A curious raccoon", "--save-output", "--log-level=debug", f"--width={self.width}", f"--height={self.height}", f"--output-path={self.output_path}", ] def _run_command(self, name: str, model_path: str, args=[]): command = ( self.get_base_command() + [f"--model-path={model_path}"] + shlex.split(args or "") + ["--output-file-name", f"{name}"] + self.extra_args ) succeed = run_command(command) status = "Success" if succeed else "Failed" return name, status def _run_test(self, name: str, args, model_path: str, test_key: str): name, status = self._run_command(name, args=args, model_path=model_path) self.verify(status, name) def verify(self, status, name): print("-" * 80) print("\n" * 3) # test task status self.assertEqual(status, "Success", f"{name} command failed") # test output file path = os.path.join( self.output_path, f"{name}.{self.data_type.get_default_extension()}" ) self.assertTrue(os.path.exists(path), f"Output file not exist for {path}") if self.data_type == DataType.IMAGE: with Image.open(path) as image: check_image_size(self, image, self.width, self.height) def model_name(self): return self.model_path.split("/")[-1] def test_single_gpu(self): """single gpu""" self._run_test( name=f"{self.model_name()}_single_gpu", args=None, model_path=self.model_path, test_key="test_single_gpu", ) ================================================ FILE: python/sglang/multimodal_gen/test/cli/test_generate_i2i.py ================================================ import os import unittest from PIL import Image from sglang.multimodal_gen.configs.sample.sampling_params import DataType from sglang.multimodal_gen.test.cli.test_generate_common import CLIBase, run_command from sglang.multimodal_gen.test.test_utils import ( DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST, check_image_size, ) class TestQwenImageEditI2I(CLIBase): model_path: str = DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST data_type: DataType = DataType.IMAGE width: int = 512 height: int = 512 test_image_urls = [ "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_1.jpg", "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_2.jpg", ] def get_base_command(self): return [ "sglang", "generate", "--save-output", "--log-level=info", f"--width={self.width}", f"--height={self.height}", f"--output-path={self.output_path}", ] def verify_multi_output(self, name: str, num_outputs: int): output_files = [] try: all_files = os.listdir(self.output_path) ext = self.data_type.get_default_extension() for f in all_files: if f.endswith(f".{ext}"): output_files.append(f) self.assertEqual( len(output_files), num_outputs, f"Expected {num_outputs} output files, found {len(output_files)}: {output_files}", ) for f in output_files: path = os.path.join(self.output_path, f) with Image.open(path) as image: check_image_size(self, image, self.width, self.height) finally: for f in output_files: path = os.path.join(self.output_path, f) if os.path.exists(path): os.remove(path) def test_single_prompt_single_image(self): """Case 1: Single prompt + single image.""" name = "single_prompt_single_image" command = self.get_base_command() + [ f"--model-path={self.model_path}", "--prompt", "Add a red hat", "--image-path", self.test_image_urls[0], ] succeed = run_command(command) self.assertTrue(succeed, f"{name} command failed") self.verify_multi_output(name, 1) def test_single_prompt_multi_image(self): """Case 2: Single prompt + multiple images (image composition).""" name = "single_prompt_multi_image" command = self.get_base_command() + [ f"--model-path={self.model_path}", "--prompt", "Combine both images", "--image-path", *self.test_image_urls, ] succeed = run_command(command) self.assertTrue(succeed, f"{name} command failed") self.verify_multi_output(name, 1) def test_multi_prompt_multi_image(self): """Case 3: Multiple prompts + multiple images (image editing).""" name = "multi_prompt_multi_image" command = self.get_base_command() + [ f"--model-path={self.model_path}", "--prompt", "Convert to oil painting style", "Convert to watercolor style", "--image-path", *self.test_image_urls, ] succeed = run_command(command) self.assertTrue(succeed, f"{name} command failed") self.verify_multi_output(name, 2) def test_multi_prompt_single_image(self): """Case 4: Multiple prompts + single image (image editing).""" name = "multi_prompt_single_image" command = self.get_base_command() + [ f"--model-path={self.model_path}", "--prompt", "Add a red hat", "Change to blue background", "--image-path", self.test_image_urls[0], ] succeed = run_command(command) self.assertTrue(succeed, f"{name} command failed") self.verify_multi_output(name, 2) del CLIBase if __name__ == "__main__": unittest.main() ================================================ FILE: python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import unittest from sglang.multimodal_gen.configs.sample.sampling_params import DataType from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.test.cli.test_generate_common import CLIBase from sglang.multimodal_gen.test.test_utils import DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST logger = init_logger(__name__) class TestFlux_T2V(CLIBase): model_path = DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST extra_args = [] data_type: DataType = DataType.IMAGE del CLIBase if __name__ == "__main__": unittest.main() ================================================ FILE: python/sglang/multimodal_gen/test/run_suite.py ================================================ """ Test runner for multimodal_gen that manages test suites and parallel execution. Usage: python3 run_suite.py --suite --partition-id --total-partitions Example: python3 run_suite.py --suite 1-gpu --partition-id 0 --total-partitions 4 """ import argparse import os import random import subprocess import sys from pathlib import Path import tabulate from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) _UPDATE_WEIGHTS_FROM_DISK_TEST_FILE = "test_update_weights_from_disk.py" _UPDATE_WEIGHTS_MODEL_PAIR_ENV = "SGLANG_MMGEN_UPDATE_WEIGHTS_PAIR" _UPDATE_WEIGHTS_MODEL_PAIR_IDS = ( "FLUX.2-klein-base-4B", "Qwen-Image", ) SUITES = { # no GPU required; safe to run on any CPU-only runner "unit": [ "../unit/test_sampling_params.py", "../unit/test_storage.py", "../unit/test_lora_format_adapter.py", "../unit/test_server_args.py", # add new unit tests here ], "1-gpu": [ "test_server_a.py", "test_server_b.py", # cli test "../cli/test_generate_t2i_perf.py", "test_update_weights_from_disk.py", # add new 1-gpu test files here ], "2-gpu": [ "test_server_2_gpu_a.py", "test_server_2_gpu_b.py", # add new 2-gpu test files here ], } suites_ascend = { "1-npu": [ "ascend/test_server_1_npu.py", # add new 1-npu test files here ], "2-npu": [ "ascend/test_server_2_npu.py", # add new 2-npu test files here ], "8-npu": [ "ascend/test_server_8_npu.py", # add new 8-npu test files here ], } SUITES.update(suites_ascend) STRICT_SUITES = {"unit"} def parse_args(): parser = argparse.ArgumentParser(description="Run multimodal_gen test suite") parser.add_argument( "--suite", type=str, required=True, choices=list(SUITES.keys()), help="The test suite to run (e.g., 1-gpu, 2-gpu)", ) parser.add_argument( "--partition-id", type=int, default=0, help="Index of the current partition (for parallel execution)", ) parser.add_argument( "--total-partitions", type=int, default=1, help="Total number of partitions", ) parser.add_argument( "--base-dir", type=str, default="server", help="Base directory for tests relative to this script's parent", ) parser.add_argument( "-k", "--filter", type=str, default=None, help="Pytest filter expression (passed to pytest -k)", ) parser.add_argument( "--continue-on-error", action="store_true", default=False, help="Continue running remaining tests even if one fails (for CI consistency; pytest already continues by default)", ) return parser.parse_args() def collect_test_items(files, filter_expr=None): """Collect test item node IDs from the given files using pytest --collect-only.""" cmd = [sys.executable, "-m", "pytest", "--collect-only", "-q"] if filter_expr: cmd.extend(["-k", filter_expr]) cmd.extend(files) print(f"Collecting tests with command: {' '.join(cmd)}") result = subprocess.run(cmd, capture_output=True, text=True) # Check for collection errors # pytest exit codes: # 0: success # 1: tests collected but some had errors during collection # 2: test execution interrupted # 3: internal error # 4: command line usage error # 5: no tests collected (may be expected with filters) if result.returncode not in (0, 5): error_msg = ( f"pytest --collect-only failed with exit code {result.returncode}\n" f"Command: {' '.join(cmd)}\n" ) if result.stderr: error_msg += f"stderr:\n{result.stderr}\n" if result.stdout: error_msg += f"stdout:\n{result.stdout}\n" logger.error(error_msg) raise RuntimeError(error_msg) if result.returncode == 5: print( "No tests were collected (exit code 5). This may be expected with filters." ) # Parse the output to extract test node IDs # pytest -q outputs lines like: test_file.py::TestClass::test_method[param] test_items = [] for line in result.stdout.strip().split("\n"): line = line.strip() # Skip empty lines and summary lines if line and "::" in line and not line.startswith(("=", "-", " ")): # Handle lines that might have extra info after the test ID test_id = line.split()[0] if " " in line else line if "::" in test_id: test_items.append(test_id) print(f"Collected {len(test_items)} test items") return test_items def run_pytest(files, filter_expr=None): if not files: print("No files to run.") return 0 base_cmd = [sys.executable, "-m", "pytest", "-s", "-v"] # Add pytest -k filter if provided if filter_expr: base_cmd.extend(["-k", filter_expr]) max_retries = 6 # retry if the perf assertion failed, for {max_retries} times for i in range(max_retries + 1): cmd = list(base_cmd) if i > 0: cmd.append("--last-failed") # Always include files to constrain test discovery scope # This prevents pytest from scanning the entire rootdir and # discovering unrelated tests that may have missing dependencies cmd.extend(files) if i > 0: print( f"Performance assertion failed. Retrying ({i}/{max_retries}) with --last-failed..." ) print(f"Running command: {' '.join(cmd)}") process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=0, ) output_bytes = bytearray() while True: chunk = process.stdout.read(4096) if not chunk: break sys.stdout.buffer.write(chunk) sys.stdout.buffer.flush() output_bytes.extend(chunk) process.wait() returncode = process.returncode if returncode == 0: return 0 # Exit code 5 means no tests were collected/selected - treat as success # when using filters, since some partitions may have all tests filtered out if returncode == 5: print( "No tests collected (exit code 5). This is expected when filters " "deselect all tests in a partition. Treating as success." ) return 0 # check if the failure is due to an assertion in test_server_utils.py full_output = output_bytes.decode("utf-8", errors="replace") is_perf_assertion = ( "multimodal_gen/test/server/test_server_utils.py" in full_output and "AssertionError" in full_output ) is_flaky_ci_assertion = ( "SafetensorError" in full_output or "FileNotFoundError" in full_output ) is_oom_error = ( "out of memory" in full_output.lower() or "oom killer" in full_output.lower() ) if not (is_perf_assertion or is_flaky_ci_assertion or is_oom_error): return returncode print(f"Max retry exceeded") return returncode def _is_in_ci() -> bool: return os.environ.get("SGLANG_IS_IN_CI", "").lower() in ("1", "true", "yes", "on") def _maybe_pin_update_weights_model_pair(suite_files_rel: list[str]) -> None: if not _is_in_ci(): return if _UPDATE_WEIGHTS_FROM_DISK_TEST_FILE not in suite_files_rel: return if os.environ.get(_UPDATE_WEIGHTS_MODEL_PAIR_ENV): print( f"Using preset {_UPDATE_WEIGHTS_MODEL_PAIR_ENV}=" f"{os.environ[_UPDATE_WEIGHTS_MODEL_PAIR_ENV]}" ) return selected_pair = random.choice(_UPDATE_WEIGHTS_MODEL_PAIR_IDS) os.environ[_UPDATE_WEIGHTS_MODEL_PAIR_ENV] = selected_pair print(f"Selected {_UPDATE_WEIGHTS_MODEL_PAIR_ENV}={selected_pair} for this CI run") def main(): args = parse_args() # 1. resolve base path current_file_path = Path(__file__).resolve() test_root_dir = current_file_path.parent target_dir = test_root_dir / args.base_dir if not target_dir.exists(): print(f"Error: Target directory {target_dir} does not exist.") sys.exit(1) # 2. get files from suite suite_files_rel = SUITES[args.suite] _maybe_pin_update_weights_model_pair(suite_files_rel) suite_files_abs = [] for f_rel in suite_files_rel: f_abs = target_dir / f_rel if not f_abs.exists(): msg = f"Test file {f_rel} not found in {target_dir}." if args.suite in STRICT_SUITES: print(f"Error: {msg}") sys.exit(1) print(f"Warning: {msg} Skipping.") continue suite_files_abs.append(str(f_abs)) if not suite_files_abs: print(f"No valid test files found for suite '{args.suite}'.") sys.exit(1 if args.suite in STRICT_SUITES else 0) # 3. collect all test items and partition by items (not files) all_test_items = collect_test_items(suite_files_abs, filter_expr=args.filter) if not all_test_items: print(f"No test items found for suite '{args.suite}'.") sys.exit(0) # Partition by test items my_items = [ item for i, item in enumerate(all_test_items) if i % args.total_partitions == args.partition_id ] # Print test info at beginning (similar to test/run_suite.py pretty_print_tests) partition_info = f"{args.partition_id + 1}/{args.total_partitions} (0-based id={args.partition_id})" headers = ["Suite", "Partition"] rows = [[args.suite, partition_info]] msg = tabulate.tabulate(rows, headers=headers, tablefmt="psql") + "\n" msg += f"✅ Enabled {len(my_items)} test(s):\n" for item in my_items: msg += f" - {item}\n" print(msg, flush=True) print( f"Suite: {args.suite} | Partition: {args.partition_id}/{args.total_partitions}" ) print(f"Selected {len(suite_files_abs)} files:") for f in suite_files_abs: print(f" - {os.path.basename(f)}") if not my_items: print("No items assigned to this partition. Exiting success.") sys.exit(0) print(f"Running {len(my_items)} items in this shard: {', '.join(my_items)}") # 4. execute with the specific test items exit_code = run_pytest(my_items) # Print tests again at the end for visibility msg = "\n" + tabulate.tabulate(rows, headers=headers, tablefmt="psql") + "\n" msg += f"✅ Executed {len(my_items)} test(s):\n" for item in my_items: msg += f" - {item}\n" print(msg, flush=True) sys.exit(exit_code) if __name__ == "__main__": main() ================================================ FILE: python/sglang/multimodal_gen/test/scripts/gen_diffusion_ci_outputs.py ================================================ #!/usr/bin/env python3 """ Generate diffusion CI outputs for consistency testing. This script reuses the CI test code by calling run_suite.py with SGLANG_GEN_GT=1, ensuring that GT generation uses exactly the same code path as CI tests. Usage: python gen_diffusion_ci_outputs.py --suite 1-gpu --partition-id 0 --total-partitions 2 --out-dir ./output python gen_diffusion_ci_outputs.py --suite 1-gpu --case-ids qwen_image_t2i flux_image_t2i --out-dir ./output """ import argparse import os import sys from pathlib import Path from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.test.run_suite import SUITES, collect_test_items, run_pytest logger = init_logger(__name__) def main(): """Main entry point.""" parser = argparse.ArgumentParser(description="Generate diffusion CI outputs") parser.add_argument( "--suite", type=str, choices=["1-gpu", "2-gpu"], required=True, help="Test suite to run (1-gpu or 2-gpu)", ) parser.add_argument( "--partition-id", type=int, required=False, help="Partition ID for matrix partitioning (0-based)", ) parser.add_argument( "--total-partitions", type=int, required=False, help="Total number of partitions", ) parser.add_argument( "--out-dir", type=str, required=True, help="Output directory for generated files", ) parser.add_argument( "--continue-on-error", action="store_true", help="Continue processing other cases if one fails", ) parser.add_argument( "--case-ids", type=str, nargs="*", required=False, help="Specific case IDs to run (space-separated). If provided, only these cases will be run.", ) args = parser.parse_args() # Validate partition arguments if args.partition_id is not None and args.total_partitions is not None: if args.partition_id < 0 or args.partition_id >= args.total_partitions: parser.error(f"partition-id must be in range [0, {args.total_partitions})") elif args.partition_id is not None or args.total_partitions is not None: parser.error( "Both --partition-id and --total-partitions must be provided together" ) # Create output directory out_dir = Path(args.out_dir) out_dir.mkdir(parents=True, exist_ok=True) # Set environment variables for GT generation mode os.environ["SGLANG_GEN_GT"] = "1" os.environ["SGLANG_GT_OUTPUT_DIR"] = str(out_dir.absolute()) os.environ["SGLANG_SKIP_CONSISTENCY"] = ( "1" # Skip consistency checks in GT gen mode ) logger.info(f"GT generation mode enabled") logger.info(f"Output directory: {out_dir}") # Resolve test files path (same as run_suite.py) current_file_path = Path(__file__).resolve() test_root_dir = current_file_path.parent.parent # scripts -> test target_dir = test_root_dir / "server" # Get files from suite (same as run_suite.py) suite_files_rel = SUITES[args.suite] suite_files_abs = [] for f_rel in suite_files_rel: f_abs = target_dir / f_rel if not f_abs.exists(): logger.warning(f"Test file {f_rel} not found in {target_dir}. Skipping.") continue suite_files_abs.append(str(f_abs)) if not suite_files_abs: logger.error(f"No valid test files found for suite '{args.suite}'.") sys.exit(1) # Build pytest filter for case_ids if provided filter_expr = None if args.case_ids: # pytest parametrized test format: test_diffusion_generation[case_id] filters = [f"test_diffusion_generation[{case_id}]" for case_id in args.case_ids] filter_expr = " or ".join(filters) logger.info(f"Filtering by case IDs: {args.case_ids}") # Collect all test items (same as run_suite.py) all_test_items = collect_test_items(suite_files_abs, filter_expr=filter_expr) if not all_test_items: logger.warning(f"No test items found for suite '{args.suite}'.") sys.exit(0) # Partition by test items (same as run_suite.py) partition_id = args.partition_id if args.partition_id is not None else 0 total_partitions = args.total_partitions if args.total_partitions is not None else 1 my_items = [ item for i, item in enumerate(all_test_items) if i % total_partitions == partition_id ] logger.info( f"Partition {partition_id}/{total_partitions}: " f"running {len(my_items)} of {len(all_test_items)} test items" ) if not my_items: logger.warning("No items assigned to this partition. Exiting success.") sys.exit(0) # Run pytest with the specific test items (same as run_suite.py) exit_code = run_pytest(my_items) if exit_code != 0: if args.continue_on_error: logger.warning(f"pytest exited with code {exit_code}") else: sys.exit(exit_code) if __name__ == "__main__": main() ================================================ FILE: python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py ================================================ import argparse import inspect import json import os import re import sys from pathlib import Path from openai import OpenAI from sglang.multimodal_gen.test.server.test_server_utils import ( ServerManager, get_generate_fn, ) from sglang.multimodal_gen.test.server.testcase_configs import ( BASELINE_CONFIG, DiffusionTestCase, ) from sglang.multimodal_gen.test.test_utils import ( get_dynamic_server_port, wait_for_req_perf_record, ) def _all_cases() -> list[DiffusionTestCase]: import sglang.multimodal_gen.test.server.testcase_configs as cfg cases: list[DiffusionTestCase] = [] for _, v in inspect.getmembers(cfg): if isinstance(v, list) and v and isinstance(v[0], DiffusionTestCase): cases.extend(v) seen: set[str] = set() out: list[DiffusionTestCase] = [] for c in cases: if c.id not in seen: seen.add(c.id) out.append(c) return out def _baseline_path() -> Path: import sglang.multimodal_gen.test.server.testcase_configs as cfg return Path(cfg.__file__).with_name("perf_baselines.json") def _openai_client(port: int) -> OpenAI: return OpenAI(api_key="sglang-anything", base_url=f"http://localhost:{port}/v1") def _build_server_extra_args(case: DiffusionTestCase) -> str: server_args = case.server_args a = os.environ.get("SGLANG_TEST_SERVE_ARGS", "") a += f" --num-gpus {server_args.num_gpus}" if server_args.tp_size is not None: a += f" --tp-size {server_args.tp_size}" if server_args.ulysses_degree is not None: a += f" --ulysses-degree {server_args.ulysses_degree}" if server_args.dit_layerwise_offload: a += " --dit-layerwise-offload true" if server_args.dit_offload_prefetch_size: a += f" --dit-offload-prefetch-size {server_args.dit_offload_prefetch_size}" if server_args.text_encoder_cpu_offload: a += " --text-encoder-cpu-offload" if server_args.ring_degree is not None: a += f" --ring-degree {server_args.ring_degree}" if server_args.lora_path: a += f" --lora-path {server_args.lora_path}" # default warmup a += " --warmup" for extra_arg in server_args.extras: a += f" {extra_arg}" return a def _build_env_vars(case: DiffusionTestCase) -> dict[str, str]: if case.server_args.enable_cache_dit: return {"SGLANG_CACHE_DIT_ENABLED": "true"} return {} def _torch_cleanup() -> None: try: import gc gc.collect() except Exception: pass try: import torch if torch.get_device_module().is_available(): torch.get_device_module().synchronize() torch.get_device_module().empty_cache() except Exception: pass def _run_case(case: DiffusionTestCase) -> dict: default_port = get_dynamic_server_port() port = int(os.environ.get("SGLANG_TEST_SERVER_PORT", default_port)) mgr = ServerManager( model=case.server_args.model_path, port=port, wait_deadline=float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200")), extra_args=_build_server_extra_args(case), env_vars=_build_env_vars(case), ) ctx = mgr.start() try: sp = case.sampling_params output_size = os.environ.get("SGLANG_TEST_OUTPUT_SIZE", sp.output_size) client = _openai_client(ctx.port) gen = get_generate_fn( model_path=case.server_args.model_path, modality=case.server_args.modality, sampling_params=sp, ) rid, _ = gen(case.id, client) rec = wait_for_req_perf_record( rid, ctx.perf_log_path, timeout=float(os.environ.get("SGLANG_PERF_TIMEOUT", "300")), ) if rec is None: raise RuntimeError(f"missing perf record: {case.id}") from sglang.multimodal_gen.test.server.testcase_configs import ( PerformanceSummary, ) perf = PerformanceSummary.from_req_perf_record( rec, BASELINE_CONFIG.step_fractions ) if case.server_args.modality == "video" and sp.num_frames and sp.num_frames > 0: if "per_frame_generation" not in perf.stage_metrics: perf.stage_metrics["per_frame_generation"] = perf.e2e_ms / sp.num_frames return { "stages_ms": {k: round(v, 2) for k, v in perf.stage_metrics.items()}, "denoise_step_ms": { str(k): round(v, 2) for k, v in perf.all_denoise_steps.items() }, "expected_e2e_ms": round(perf.e2e_ms, 2), "expected_avg_denoise_ms": round(perf.avg_denoise_ms, 2), "expected_median_denoise_ms": round(perf.median_denoise_ms, 2), } finally: ctx.cleanup() def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--baseline", default="") ap.add_argument("--out", default="") ap.add_argument("--match", default="") ap.add_argument("--case", action="append", default=[]) ap.add_argument("--all-from-baseline", action="store_true") ap.add_argument("--timeout", type=float, default=300.0) args = ap.parse_args() os.environ.setdefault("SGLANG_GEN_BASELINE", "1") os.environ["SGLANG_PERF_TIMEOUT"] = str(args.timeout) baseline_path = Path(args.baseline) if args.baseline else _baseline_path() out_path = Path(args.out) if args.out else baseline_path data = json.loads(baseline_path.read_text(encoding="utf-8")) scenarios = data.setdefault("scenarios", {}) ids = set(args.case) if args.case else None pat = re.compile(args.match) if args.match else None if args.all_from_baseline: ids = set(scenarios.keys()) pat = None all_cases = _all_cases() cases = [] for c in all_cases: if ids and c.id not in ids: continue if pat and not pat.search(c.id): continue cases.append(c) if args.all_from_baseline and ids: case_ids = {c.id for c in all_cases} missing = sorted([i for i in ids if i not in case_ids]) if missing: sys.stderr.write(f"missing cases in testcase_configs.py: {len(missing)}\n") if not cases: return 0 for c in cases: prev = scenarios.get(c.id, {}) note = prev.get("notes") baseline = _run_case(c) if note is not None: baseline["notes"] = note scenarios[c.id] = baseline sys.stdout.write(f"{c.id}\n") sys.stdout.flush() _torch_cleanup() out_path.write_text(json.dumps(data, indent=4) + "\n", encoding="utf-8") return 0 if __name__ == "__main__": sys.exit(main()) ================================================ FILE: python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json ================================================ { "metadata": { "model": "Diffusion Server", "hardware": "CI A2 64GB pool", "description": "Reference numbers captured from the CI diffusion server baseline run" }, "scenarios": { "flux_image_t2i_npu": { "stages_ms": { "InputValidationStage": 0.07, "TextEncodingStage": 154.51, "TimestepPreparationStage": 53.52, "LatentPreparationStage": 0.39, "DenoisingStage": 19423.39, "DecodingStage": 40.14 }, "denoise_step_ms": { "0": 123.16, "1": 91.7, "2": 265.62, "3": 402.68, "4": 402.86, "5": 402.78, "6": 402.99, "7": 402.77, "8": 402.59, "9": 402.93, "10": 402.05, "11": 402.99, "12": 402.29, "13": 403.07, "14": 402.62, "15": 402.99, "16": 402.68, "17": 403.0, "18": 402.74, "19": 402.85, "20": 402.83, "21": 403.03, "22": 402.56, "23": 402.84, "24": 402.79, "25": 402.95, "26": 402.65, "27": 403.01, "28": 402.66, "29": 402.92, "30": 402.75, "31": 403.0, "32": 402.9, "33": 402.48, "34": 402.85, "35": 402.03, "36": 402.93, "37": 402.3, "38": 403.12, "39": 402.83, "40": 402.84, "41": 402.75, "42": 402.97, "43": 402.62, "44": 402.91, "45": 402.81, "46": 402.97, "47": 402.57, "48": 403.0, "49": 402.75 }, "expected_e2e_ms": 23819.1, "expected_avg_denoise_ms": 388.22, "expected_median_denoise_ms": 402.82 }, "flux_2_image_t2i_2npu": { "stages_ms": { "InputValidationStage": 0.06, "TextEncodingStage": 5628.31, "ImageVAEEncodingStage": 0.01, "LatentPreparationStage": 0.75, "TimestepPreparationStage": 30.68, "DenoisingStage": 55002.26, "DecodingStage": 43.73 }, "denoise_step_ms": { "0": 110.35, "1": 301.82, "2": 1139.81, "3": 1114.17, "4": 1099.34, "5": 1099.12, "6": 1100.16, "7": 1099.67, "8": 1099.09, "9": 1089.81, "10": 1109.73, "11": 1099.97, "12": 1100.26, "13": 1099.67, "14": 1099.79, "15": 1099.6, "16": 1100.16, "17": 1099.87, "18": 1100.02, "19": 1099.34, "20": 1099.6, "21": 1099.45, "22": 1100.2, "23": 1099.29, "24": 1098.86, "25": 1090.38, "26": 1109.19, "27": 1099.67, "28": 1100.06, "29": 1099.22, "30": 1100.08, "31": 1098.86, "32": 1099.73, "33": 1099.11, "34": 1100.13, "35": 1103.97, "36": 1095.26, "37": 1099.38, "38": 1099.34, "39": 1099.17, "40": 1100.08, "41": 1089.89, "42": 1106.69, "43": 1102.57, "44": 1100.17, "45": 1099.21, "46": 1100.42, "47": 1099.38, "48": 1099.59, "49": 1099.47 }, "expected_e2e_ms": 64195.08, "expected_avg_denoise_ms": 1065.0, "expected_median_denoise_ms": 1099.63 }, "wan2_1_t2v_1.3b_1_npu": { "stages_ms": { "InputValidationStage": 0.07, "TextEncodingStage": 876.11, "LatentPreparationStage": 0.25, "TimestepPreparationStage": 2.9, "DenoisingStage": 26188.0, "DecodingStage": 320.03, "per_frame_generation": null }, "denoise_step_ms": { "0": 103.56, "1": 329.59, "2": 545.23, "3": 537.0, "4": 536.27, "5": 536.29, "6": 536.33, "7": 536.0, "8": 536.17, "9": 536.28, "10": 535.53, "11": 536.04, "12": 536.42, "13": 536.09, "14": 536.32, "15": 536.25, "16": 536.36, "17": 536.21, "18": 536.29, "19": 536.15, "20": 536.28, "21": 536.5, "22": 536.46, "23": 536.06, "24": 536.45, "25": 536.24, "26": 536.14, "27": 536.13, "28": 536.22, "29": 536.15, "30": 535.94, "31": 536.1, "32": 536.13, "33": 536.2, "34": 536.24, "35": 536.34, "36": 536.54, "37": 536.42, "38": 536.41, "39": 536.42, "40": 536.13, "41": 536.32, "42": 536.23, "43": 536.16, "44": 536.05, "45": 536.18, "46": 536.08, "47": 536.34, "48": 536.26, "49": 535.41 }, "expected_e2e_ms": 38738.17, "expected_avg_denoise_ms": 523.62, "expected_median_denoise_ms": 536.23 }, "wan2_2_t2v_14b_w8a8_8npu": { "stages_ms": { "InputValidationStage": 0.07, "TextEncodingStage": 301.21, "LatentPreparationStage": 0.2, "TimestepPreparationStage": 2.68, "DenoisingStage": 83661.46, "DecodingStage": 232.94, "per_frame_generation": null }, "denoise_step_ms": { "0": 1919.92, "1": 2099.45, "2": 2092.11, "3": 2090.84, "4": 2089.89, "5": 2090.6, "6": 2090.77, "7": 2091.43, "8": 2091.24, "9": 2067.83, "10": 2078.02, "11": 2090.75, "12": 2108.36, "13": 2096.16, "14": 2091.74, "15": 2091.47, "16": 2091.6, "17": 2091.94, "18": 2091.39, "19": 2090.69, "20": 2090.27, "21": 2090.77, "22": 2090.24, "23": 2091.65, "24": 2091.21, "25": 2126.82, "26": 2338.39, "27": 2085.18, "28": 2084.68, "29": 2084.71, "30": 2051.48, "31": 2104.3, "32": 2084.58, "33": 2085.04, "34": 2085.03, "35": 2084.58, "36": 2084.41, "37": 2085.16, "38": 2084.88, "39": 2083.54 }, "expected_e2e_ms": 91733.92, "expected_avg_denoise_ms": 2091.33, "expected_median_denoise_ms": 2090.72 } } } ================================================ FILE: python/sglang/multimodal_gen/test/server/ascend/test_server_1_npu.py ================================================ """ Config-driven diffusion performance test with pytest parametrization. If the actual run is significantly better than the baseline, the improved cases with their updated baseline will be printed """ from __future__ import annotations import pytest from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.test.server.ascend.testcase_configs_npu import ONE_NPU_CASES from sglang.multimodal_gen.test.server.test_server_common import ( # noqa: F401 DiffusionServerBase, diffusion_server, ) from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase logger = init_logger(__name__) class TestDiffusionServerOneNpu(DiffusionServerBase): """Performance tests for 1-NPU diffusion cases.""" @pytest.fixture(params=ONE_NPU_CASES, ids=lambda c: c.id) def case(self, request) -> DiffusionTestCase: """Provide a DiffusionTestCase for each 1-NPU test.""" return request.param ================================================ FILE: python/sglang/multimodal_gen/test/server/ascend/test_server_2_npu.py ================================================ """ Config-driven diffusion performance test with pytest parametrization. If the actual run is significantly better than the baseline, the improved cases with their updated baseline will be printed """ from __future__ import annotations import pytest from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.test.server.ascend.testcase_configs_npu import TWO_NPU_CASES from sglang.multimodal_gen.test.server.test_server_common import ( # noqa: F401 DiffusionServerBase, diffusion_server, ) from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase logger = init_logger(__name__) class TestDiffusionServerTwoNpu(DiffusionServerBase): """Performance tests for 2-NPU diffusion cases.""" @pytest.fixture(params=TWO_NPU_CASES, ids=lambda c: c.id) def case(self, request) -> DiffusionTestCase: """Provide a DiffusionTestCase for each 2-NPU test.""" return request.param ================================================ FILE: python/sglang/multimodal_gen/test/server/ascend/test_server_8_npu.py ================================================ """ Config-driven diffusion performance test with pytest parametrization. If the actual run is significantly better than the baseline, the improved cases with their updated baseline will be printed """ from __future__ import annotations import pytest from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.test.server.ascend.testcase_configs_npu import ( EIGHT_NPU_CASES, ) from sglang.multimodal_gen.test.server.test_server_common import ( # noqa: F401 DiffusionServerBase, diffusion_server, ) from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase logger = init_logger(__name__) class TestDiffusionServerEightNpu(DiffusionServerBase): """Performance tests for 8-NPU diffusion cases.""" @pytest.fixture(params=EIGHT_NPU_CASES, ids=lambda c: c.id) def case(self, request) -> DiffusionTestCase: """Provide a DiffusionTestCase for each 8-NPU test.""" return request.param ================================================ FILE: python/sglang/multimodal_gen/test/server/ascend/testcase_configs_npu.py ================================================ from sglang.multimodal_gen.test.server.testcase_configs import ( T2V_PROMPT, DiffusionSamplingParams, DiffusionServerArgs, DiffusionTestCase, T2I_sampling_params, ) ONE_NPU_CASES: list[DiffusionTestCase] = [ # === Text to Image (T2I) === DiffusionTestCase( "flux_image_t2i_npu", DiffusionServerArgs( model_path="/root/.cache/modelscope/hub/models/black-forest-labs/FLUX.1-dev", modality="image", ), T2I_sampling_params, ), # === Text to Video (T2V) === DiffusionTestCase( "wan2_1_t2v_1.3b_1_npu", DiffusionServerArgs( model_path="/root/.cache/modelscope/hub/models/Wan-AI/Wan2.1-T2V-1.3B-Diffusers", modality="video", custom_validator="video", ), DiffusionSamplingParams( prompt=T2V_PROMPT, ), ), ] TWO_NPU_CASES: list[DiffusionTestCase] = [ # === Text to Image (T2I) === DiffusionTestCase( "flux_2_image_t2i_2npu", DiffusionServerArgs( model_path="/root/.cache/modelscope/hub/models/black-forest-labs/FLUX.2-dev", modality="image", num_gpus=2, tp_size=2, ), T2I_sampling_params, ), ] EIGHT_NPU_CASES: list[DiffusionTestCase] = [ # === Text to Video (T2V) === DiffusionTestCase( "wan2_2_t2v_14b_w8a8_8npu", DiffusionServerArgs( model_path="/root/.cache/modelscope/hub/models/Eco-Tech/Wan2.2-T2V-A14B-Diffusers-w8a8", modality="video", custom_validator="video", num_gpus=8, tp_size=4, ), DiffusionSamplingParams( prompt=T2V_PROMPT, ), ), ] ================================================ FILE: python/sglang/multimodal_gen/test/server/conftest.py ================================================ import os import pytest print("[CONFTEST] Loading conftest.py at import time") def pytest_configure(config): """ Create the perf results StashKey once and store it in config. This hook runs once per test session, before module double-import issues. """ if not hasattr(config, "_diffusion_perf_key"): config._diffusion_perf_key = pytest.StashKey[list]() print(f"[CONFTEST] Created perf_results_key: {config._diffusion_perf_key}") def add_perf_results(config, results: list): """Add performance results to the shared stash.""" # Get the shared key from config (created once in pytest_configure) key = config._diffusion_perf_key existing = config.stash.get(key, []) existing.extend(results) config.stash[key] = existing print(f"[CONFTEST] Added {len(results)} results, total now: {len(existing)}") @pytest.fixture(scope="session") def perf_config(request): """Provide access to pytest config for storing perf results.""" return request.config def _write_github_step_summary(content: str): """Write content to GitHub Step Summary if available.""" summary_file = os.environ.get("GITHUB_STEP_SUMMARY") if summary_file: with open(summary_file, "a") as f: f.write(content) def _write_results_json(results: list, output_path: str = "diffusion-results.json"): """Write performance results to JSON file for CI artifact collection.""" import json try: with open(output_path, "w") as f: json.dump(results, f, indent=2) print(f"[CONFTEST] Wrote results to {output_path}") except Exception as e: print(f"[CONFTEST] Failed to write results JSON: {e}") def _generate_diffusion_markdown_report(results: list) -> str: """Generate a markdown report for diffusion performance results.""" if not results: return "" gpu_config = os.environ.get("GPU_CONFIG", "") header = "## Diffusion Performance Summary" if gpu_config: header += f" [{gpu_config}]" header += "\n\n" # Main performance table markdown = header markdown += "| Test Suite | Test Name | Modality | E2E (ms) | Avg Denoise (ms) | Median Denoise (ms) |\n" markdown += "| ---------- | --------- | -------- | -------- | ---------------- | ------------------- |\n" for entry in sorted(results, key=lambda x: (x["class_name"], x["test_name"])): modality = entry.get("modality", "image") markdown += ( f"| {entry['class_name']} | {entry['test_name']} | {modality} | " f"{entry['e2e_ms']:.2f} | {entry['avg_denoise_ms']:.2f} | " f"{entry['median_denoise_ms']:.2f} |\n" ) # Video-specific metrics table (if any video tests) video_results = [r for r in results if r.get("modality") == "video"] if video_results: markdown += "\n### Video Generation Metrics\n\n" markdown += "| Test Name | FPS | Total Frames | Avg Frame Time (ms) |\n" markdown += "| --------- | --- | ------------ | ------------------- |\n" for entry in video_results: fps = entry.get("frames_per_second", "N/A") frames = entry.get("total_frames", "N/A") avg_frame = entry.get("avg_frame_time_ms", "N/A") if isinstance(fps, float): fps = f"{fps:.2f}" if isinstance(avg_frame, float): avg_frame = f"{avg_frame:.2f}" markdown += f"| {entry['test_name']} | {fps} | {frames} | {avg_frame} |\n" return markdown def pytest_sessionfinish(session): """ This hook is called by pytest at the end of the entire test session. It prints a consolidated summary of all performance results. """ # Get results from stash using the shared key from config key = session.config._diffusion_perf_key results = session.config.stash.get(key, []) print(f"\n[DEBUG] pytest_sessionfinish called, has {len(results)} entries") if not results: print("[DEBUG] No results collected, skipping summary output") return # Print to stdout (existing behavior) print("\n\n" + "=" * 35 + " Performance Summary " + "=" * 35) print( f"{'Test Suite':<30} | {'Test Name':<20} | {'E2E (ms)':>12} | {'Avg Denoise (ms)':>18} | {'Median Denoise (ms)':>20}" ) print( "-" * 30 + "-+-" + "-" * 20 + "-+-" + "-" * 12 + "-+-" + "-" * 18 + "-+-" + "-" * 20 ) for entry in sorted(results, key=lambda x: x["class_name"]): print( f"{entry['class_name']:<30} | {entry['test_name']:<20} | {entry['e2e_ms']:>12.2f} | " f"{entry['avg_denoise_ms']:>18.2f} | {entry['median_denoise_ms']:>20.2f}" ) print("=" * 91) print("\n\n" + "=" * 36 + " Detailed Reports " + "=" * 37) for entry in sorted(results, key=lambda x: x["class_name"]): print(f"\n--- Details for {entry['class_name']} / {entry['test_name']} ---") stage_report = ", ".join( f"{name}:{duration:.2f}ms" for name, duration in entry.get("stage_metrics", {}).items() ) if stage_report: print(f" Stages: {stage_report}") sampled_steps = entry.get("sampled_steps") or {} if sampled_steps: step_report = ", ".join( f"{idx}:{duration:.2f}ms" for idx, duration in sorted(sampled_steps.items()) ) print(f" Sampled Steps: {step_report}") print("=" * 91) # Write to GitHub Step Summary (new behavior for CI monitoring) markdown_report = _generate_diffusion_markdown_report(results) if markdown_report: _write_github_step_summary(markdown_report) # Write results to JSON file for CI artifact collection _write_results_json(results) ================================================ FILE: python/sglang/multimodal_gen/test/server/perf_baselines.json ================================================ { "metadata": { "model": "Diffusion Server", "hardware": "CI H100 80GB pool", "description": "Reference numbers captured from the CI diffusion server baseline run" }, "tolerances": { "long_term": { "e2e": 0.1, "denoise_stage": 0.05, "non_denoise_stage": 0.4, "denoise_step": 0.2, "denoise_agg": 0.1 }, "pr_test": { "e2e": 0.15, "denoise_stage": 0.1, "non_denoise_stage": 0.6, "denoise_step": 0.25, "denoise_agg": 0.15 } }, "improvement_reporting": { "threshold": 0.2 }, "sampling": { "step_fractions": [ 0.0, 0.2, 0.4, 0.6, 0.8, 1.0 ] }, "scenarios": { "qwen_image_t2i": { "notes": "Single-image generation using the default prompt", "stages_ms": { "DecodingStage": 51.86, "TextEncodingStage": 611.83, "InputValidationStage": 0.05, "DenoisingStage": 14289.46, "LatentPreparationStage": 0.2, "TimestepPreparationStage": 3.34 }, "denoise_step_ms": { "0": 240.5, "1": 279.1, "2": 283.29, "3": 296.63, "4": 287.72, "5": 283.39, "6": 283.98, "7": 291.82, "8": 283.1, "9": 284.43, "10": 288.95, "11": 285.6, "12": 285.99, "13": 285.47, "14": 289.66, "15": 285.74, "16": 284.15, "17": 290.27, "18": 288.04, "19": 284.57, "20": 286.69, "21": 288.95, "22": 287.09, "23": 285.6, "24": 289.31, "25": 285.48, "26": 285.53, "27": 288.13, "28": 287.65, "29": 285.97, "30": 288.9, "31": 287.97, "32": 286.48, "33": 285.38, "34": 286.62, "35": 288.22, "36": 285.6, "37": 286.61, "38": 287.06, "39": 286.2, "40": 284.6, "41": 285.69, "42": 288.46, "43": 285.53, "44": 285.34, "45": 285.74, "46": 287.25, "47": 285.0, "48": 286.82, "49": 287.19 }, "expected_e2e_ms": 14959.11, "expected_avg_denoise_ms": 285.67, "expected_median_denoise_ms": 286.1 }, "qwen_image_t2i_2_gpus": { "stages_ms": { "InputValidationStage": 0.04, "TextEncodingStage": 693.2, "TimestepPreparationStage": 2.84, "LatentPreparationStage": 9.13, "DenoisingStage": 24529.77, "DecodingStage": 612.79 }, "denoise_step_ms": { "0": 405.94, "1": 420.06, "2": 414.79, "3": 392.4, "4": 408.14, "5": 605.0, "6": 469.39, "7": 574.04, "8": 539.61, "9": 452.93, "10": 279.36, "11": 271.8, "12": 438.26, "13": 552.65, "14": 576.1, "15": 679.84, "16": 543.0, "17": 512.81, "18": 522.27, "19": 545.06, "20": 545.85, "21": 523.83, "22": 519.36, "23": 513.78, "24": 532.54, "25": 524.94, "26": 542.59, "27": 570.91, "28": 568.73, "29": 564.52, "30": 564.57, "31": 544.94, "32": 496.81, "33": 488.98, "34": 457.18, "35": 441.42, "36": 437.44, "37": 477.6, "38": 429.17, "39": 465.55, "40": 448.25, "41": 511.83, "42": 450.6, "43": 375.78, "44": 504.4, "45": 524.44, "46": 535.22, "47": 514.52, "48": 431.58, "49": 410.68 }, "expected_e2e_ms": 25850.45, "expected_avg_denoise_ms": 490.43, "expected_median_denoise_ms": 512.32 }, "flux_image_t2i": { "stages_ms": { "DecodingStage": 32.72, "TextEncodingStage": 51.96, "InputValidationStage": 0.03, "DenoisingStage": 7545.16, "LatentPreparationStage": 0.2, "TimestepPreparationStage": 2.43 }, "denoise_step_ms": { "0": 50.06, "1": 58.88, "2": 151.24, "3": 150.97, "4": 151.23, "5": 151.63, "6": 159.11, "7": 158.31, "8": 153.42, "9": 151.42, "10": 151.91, "11": 151.05, "12": 151.52, "13": 157.2, "14": 152.76, "15": 153.85, "16": 153.02, "17": 151.09, "18": 151.49, "19": 155.13, "20": 155.2, "21": 152.82, "22": 152.2, "23": 150.99, "24": 152.74, "25": 153.45, "26": 153.63, "27": 154.92, "28": 152.72, "29": 151.84, "30": 151.84, "31": 152.44, "32": 153.03, "33": 154.07, "34": 152.36, "35": 153.48, "36": 152.05, "37": 152.45, "38": 152.42, "39": 154.91, "40": 152.68, "41": 153.43, "42": 151.62, "43": 153.52, "44": 153.13, "45": 152.85, "46": 152.33, "47": 151.61, "48": 152.4, "49": 152.33 }, "expected_e2e_ms": 7798.99, "expected_avg_denoise_ms": 150.77, "expected_median_denoise_ms": 152.45 }, "flux_2_image_t2i": { "stages_ms": { "LatentPreparationStage": 0.52, "TimestepPreparationStage": 2.91, "TextEncodingStage": 518.54, "ImageVAEEncodingStage": 0.0, "InputValidationStage": 0.05, "DenoisingStage": 24901.97, "DecodingStage": 8.98 }, "denoise_step_ms": { "0": 69.14, "1": 132.57, "2": 508.67, "3": 493.52, "4": 504.31, "5": 492.99, "6": 501.91, "7": 495.18, "8": 500.87, "9": 497.36, "10": 498.74, "11": 497.46, "12": 499.08, "13": 494.65, "14": 500.35, "15": 496.89, "16": 500.23, "17": 497.01, "18": 501.68, "19": 493.8, "20": 501.1, "21": 494.81, "22": 501.04, "23": 499.27, "24": 500.04, "25": 497.14, "26": 499.05, "27": 494.91, "28": 496.89, "29": 498.53, "30": 497.94, "31": 497.09, "32": 497.7, "33": 497.58, "34": 496.43, "35": 497.7, "36": 497.37, "37": 497.17, "38": 499.27, "39": 495.52, "40": 501.67, "41": 495.11, "42": 500.69, "43": 501.61, "44": 501.91, "45": 495.58, "46": 499.37, "47": 496.8, "48": 497.49, "49": 495.69 }, "expected_e2e_ms": 25832.82, "expected_avg_denoise_ms": 489.43, "expected_median_denoise_ms": 497.53 }, "flux_2_klein_image_t2i": { "stages_ms": { "DecodingStage": 9.27, "TextEncodingStage": 92.17, "InputValidationStage": 0.05, "ImageVAEEncodingStage": 0.0, "DenoisingStage": 252.01, "LatentPreparationStage": 0.42, "TimestepPreparationStage": 1.5 }, "denoise_step_ms": { "0": 19.91, "1": 19.32, "2": 51.99, "3": 61.78 }, "expected_e2e_ms": 430.73, "expected_avg_denoise_ms": 38.25, "expected_median_denoise_ms": 35.95 }, "layerwise_offload": { "stages_ms": { "InputValidationStage": 0.06, "TextEncodingStage": 513.58, "LatentPreparationStage": 0.46, "TimestepPreparationStage": 2.38, "DenoisingStage": 52187.62, "DecodingStage": 190.31 }, "denoise_step_ms": { "0": 1033.45, "1": 137.03, "2": 1046.96, "3": 1039.28, "4": 1039.05, "5": 1043.91, "6": 1041.75, "7": 1037.6, "8": 1043.54, "9": 1048.63, "10": 1039.8, "11": 1042.25, "12": 1041.54, "13": 1045.89, "14": 1038.99, "15": 1041.82, "16": 1038.32, "17": 1045.53, "18": 1046.54, "19": 1041.22, "20": 1044.55, "21": 1041.31, "22": 1051.28, "23": 1043.12, "24": 1044.65, "25": 1042.25, "26": 1046.47, "27": 1052.9, "28": 1039.04, "29": 1042.39, "30": 1045.33, "31": 1038.05, "32": 1037.76, "33": 1037.93, "34": 1052.85, "35": 1045.59, "36": 1054.32, "37": 1044.59, "38": 1043.57, "39": 1041.93, "40": 1043.59, "41": 1046.17, "42": 1046.92, "43": 1047.04, "44": 1046.8, "45": 1041.86, "46": 1041.05, "47": 1044.04, "48": 1039.77, "49": 1047.12 }, "expected_e2e_ms": 53290.15, "expected_avg_denoise_ms": 1025.35, "expected_median_denoise_ms": 1043.33 }, "flux_2_ti2i": { "stages_ms": { "InputValidationStage": 99.82, "TextEncodingStage": 519.88, "ImageVAEEncodingStage": 254.56, "LatentPreparationStage": 12.4, "TimestepPreparationStage": 2.71, "DenoisingStage": 54705.41, "DecodingStage": 469.47 }, "denoise_step_ms": { "0": 1067.03, "1": 271.58, "2": 1073.07, "3": 1071.93, "4": 1100.0, "5": 1102.28, "6": 1088.3, "7": 1089.09, "8": 1086.95, "9": 1089.33, "10": 1089.28, "11": 1096.51, "12": 1098.88, "13": 1080.84, "14": 1098.44, "15": 1100.88, "16": 1086.83, "17": 1090.58, "18": 1096.35, "19": 1086.25, "20": 1082.71, "21": 1097.6, "22": 1098.72, "23": 1100.9, "24": 1099.02, "25": 1101.52, "26": 1098.75, "27": 1101.41, "28": 1091.75, "29": 1087.2, "30": 1101.33, "31": 1098.14, "32": 1100.14, "33": 1098.91, "34": 1100.05, "35": 1099.12, "36": 1100.22, "37": 1103.29, "38": 1092.79, "39": 1086.59, "40": 1094.81, "41": 1105.6, "42": 1100.54, "43": 1099.95, "44": 1096.5, "45": 1086.69, "46": 1095.85, "47": 1092.85, "48": 1086.17, "49": 1099.67 }, "expected_e2e_ms": 56308.23, "expected_avg_denoise_ms": 1077.26, "expected_median_denoise_ms": 1096.5 }, "flux_2_ti2i_multi_image_cache_dit": { "stages_ms": { "ImageVAEEncodingStage": 282.83, "DenoisingStage": 26936.93, "DecodingStage": 129.33, "TextEncodingStage": 737.01, "LatentPreparationStage": 0.84, "TimestepPreparationStage": 20.57, "InputValidationStage": 84.29 }, "denoise_step_ms": { "0": 1846.55, "1": 232.6, "2": 1621.05, "3": 1606.1, "4": 1441.78, "5": 59.2, "6": 60.32, "7": 235.36, "8": 1443.21, "9": 59.71, "10": 59.79, "11": 236.12, "12": 1439.11, "13": 59.97, "14": 60.46, "15": 239.29, "16": 1441.32, "17": 60.76, "18": 60.68, "19": 240.16, "20": 1442.46, "21": 60.14, "22": 61.1, "23": 239.26, "24": 1443.28, "25": 59.41, "26": 60.74, "27": 238.02, "28": 1444.38, "29": 59.09, "30": 59.12, "31": 241.69, "32": 1441.99, "33": 60.44, "34": 61.93, "35": 241.0, "36": 1443.56, "37": 60.55, "38": 61.07, "39": 238.11, "40": 1443.33, "41": 59.08, "42": 60.43, "43": 239.23, "44": 1444.53, "45": 61.77, "46": 61.65, "47": 239.84, "48": 1443.89, "49": 59.92 }, "expected_e2e_ms": 28591.93, "expected_avg_denoise_ms": 533.42, "expected_median_denoise_ms": 235.74 }, "flux_image_t2i_2_gpus": { "stages_ms": { "InputValidationStage": 0.03, "TextEncodingStage": 74.47, "TimestepPreparationStage": 2.23, "LatentPreparationStage": 6.17, "DenoisingStage": 8400.49, "DecodingStage": 381.56 }, "denoise_step_ms": { "0": 73.27, "1": 166.6, "2": 167.31, "3": 168.7, "4": 168.83, "5": 171.05, "6": 174.64, "7": 170.92, "8": 169.69, "9": 169.21, "10": 167.71, "11": 177.62, "12": 166.44, "13": 174.61, "14": 170.43, "15": 169.47, "16": 167.24, "17": 169.15, "18": 169.51, "19": 172.3, "20": 172.19, "21": 172.36, "22": 168.39, "23": 168.47, "24": 170.55, "25": 170.96, "26": 168.43, "27": 169.01, "28": 169.62, "29": 170.95, "30": 171.83, "31": 171.92, "32": 170.1, "33": 170.46, "34": 169.91, "35": 168.91, "36": 170.27, "37": 170.23, "38": 169.62, "39": 169.66, "40": 169.57, "41": 169.42, "42": 168.59, "43": 171.12, "44": 169.6, "45": 169.93, "46": 171.23, "47": 171.03, "48": 170.14, "49": 169.4 }, "expected_e2e_ms": 9006.3, "expected_avg_denoise_ms": 167.89, "expected_median_denoise_ms": 169.67 }, "zimage_image_t2i": { "stages_ms": { "InputValidationStage": 0.03, "TextEncodingStage": 403.47, "TimestepPreparationStage": 1.41, "LatentPreparationStage": 0.11, "DenoisingStage": 756.21, "DecodingStage": 29.41 }, "denoise_step_ms": { "0": 22.29, "1": 75.04, "2": 93.82, "3": 93.34, "4": 93.38, "5": 93.58, "6": 94.01, "7": 93.97, "8": 94.32 }, "expected_e2e_ms": 1292.92, "expected_avg_denoise_ms": 83.75, "expected_median_denoise_ms": 93.58 }, "zimage_image_t2i_fp8": { "stages_ms": { "InputValidationStage": 0.04, "TextEncodingStage": 428.59, "LatentPreparationStage": 0.14, "TimestepPreparationStage": 47.26, "DenoisingStage": 778.56, "DecodingStage": 10.39 }, "denoise_step_ms": { "0": 40.9, "1": 61.08, "2": 95.65, "3": 95.83, "4": 95.65, "5": 96.09, "6": 96.23, "7": 96.04, "8": 96.29 }, "expected_e2e_ms": 1370.28, "expected_avg_denoise_ms": 85.97, "expected_median_denoise_ms": 95.83 }, "zimage_image_t2i_multi_lora": { "stages_ms": { "InputValidationStage": 0.04, "TextEncodingStage": 413.69, "TimestepPreparationStage": 1.3, "LatentPreparationStage": 0.11, "DenoisingStage": 813.7, "DecodingStage": 34.51 }, "denoise_step_ms": { "0": 30.35, "1": 74.53, "2": 99.34, "3": 100.92, "4": 99.46, "5": 100.57, "6": 99.72, "7": 100.86, "8": 103.87 }, "expected_e2e_ms": 1464.31, "expected_avg_denoise_ms": 89.96, "expected_median_denoise_ms": 99.72 }, "zimage_image_t2i_2_gpus": { "stages_ms": { "InputValidationStage": 0.08, "TextEncodingStage": 420.74, "TimestepPreparationStage": 1.5, "LatentPreparationStage": 0.12, "DenoisingStage": 1304.07, "DecodingStage": 37.83 }, "denoise_step_ms": { "0": 49.76, "1": 155.22, "2": 155.98, "3": 156.16, "4": 157.04, "5": 156.54, "6": 156.29, "7": 157.36, "8": 156.05 }, "expected_e2e_ms": 1464.87, "expected_avg_denoise_ms": 144.49, "expected_median_denoise_ms": 156.16 }, "qwen_image_edit_ti2i": { "stages_ms": { "LatentPreparationStage": 0.16, "TimestepPreparationStage": 2.62, "ImageEncodingStage": 1174.26, "ImageVAEEncodingStage": 132.67, "InputValidationStage": 38.1, "DenoisingStage": 38135.64, "DecodingStage": 139.72 }, "denoise_step_ms": { "0": 618.31, "1": 769.07, "2": 766.91, "3": 762.77, "4": 764.26, "5": 765.27, "6": 767.35, "7": 764.18, "8": 766.16, "9": 766.89, "10": 766.1, "11": 764.96, "12": 763.52, "13": 765.22, "14": 765.44, "15": 763.9, "16": 763.19, "17": 764.83, "18": 765.36, "19": 765.19, "20": 765.96, "21": 765.74, "22": 765.87, "23": 764.85, "24": 765.44, "25": 765.95, "26": 766.21, "27": 767.91, "28": 765.45, "29": 764.81, "30": 766.26, "31": 765.37, "32": 766.71, "33": 765.67, "34": 766.64, "35": 765.98, "36": 766.04, "37": 764.19, "38": 765.15, "39": 766.33, "40": 767.68, "41": 765.36, "42": 766.61, "43": 766.06, "44": 765.26, "45": 765.29, "46": 764.64, "47": 766.07, "48": 762.89, "49": 763.01 }, "expected_e2e_ms": 39706.9, "expected_avg_denoise_ms": 762.57, "expected_median_denoise_ms": 765.44 }, "qwen_image_t2i_cache_dit_enabled": { "stages_ms": { "InputValidationStage": 0.05, "TextEncodingStage": 675.95, "TimestepPreparationStage": 3.21, "LatentPreparationStage": 0.2, "DenoisingStage": 5248.83, "DecodingStage": 52.24 }, "denoise_step_ms": { "0": 227.68, "1": 277.41, "2": 276.7, "3": 291.52, "4": 52.8, "5": 6.58, "6": 231.58, "7": 52.55, "8": 7.69, "9": 230.59, "10": 52.58, "11": 7.14, "12": 6.95, "13": 234.71, "14": 53.28, "15": 7.09, "16": 6.63, "17": 233.93, "18": 52.71, "19": 6.64, "20": 6.5, "21": 231.37, "22": 52.28, "23": 6.61, "24": 6.48, "25": 232.86, "26": 54.92, "27": 7.51, "28": 7.19, "29": 233.51, "30": 52.97, "31": 6.72, "32": 7.02, "33": 233.14, "34": 52.47, "35": 6.66, "36": 6.52, "37": 233.84, "38": 51.49, "39": 6.87, "40": 6.74, "41": 233.75, "42": 52.65, "43": 6.62, "44": 6.55, "45": 233.45, "46": 52.33, "47": 6.55, "48": 232.58, "49": 52.84 }, "expected_e2e_ms": 5982.78, "expected_avg_denoise_ms": 104.84, "expected_median_denoise_ms": 102.01 }, "wan2_1_t2v_1.3b_teacache_enabled": { "stages_ms": { "DenoisingStage": 4598.36, "InputValidationStage": 0.07, "DecodingStage": 552.92, "LatentPreparationStage": 0.26, "per_frame_generation": null, "TextEncodingStage": 1114.01, "TimestepPreparationStage": 2.1 }, "denoise_step_ms": { "0": 94.24, "1": 172.68, "2": 169.48, "3": 169.08, "4": 168.38, "5": 167.27, "6": 62.95, "7": 119.56, "8": 53.34, "9": 121.85, "10": 47.64, "11": 125.75, "12": 3.24, "13": 48.21, "14": 125.17, "15": 3.71, "16": 48.15, "17": 124.61, "18": 3.3, "19": 47.25, "20": 129.33, "21": 3.11, "22": 48.03, "23": 127.46, "24": 3.37, "25": 45.6, "26": 127.17, "27": 3.35, "28": 49.83, "29": 125.42, "30": 3.19, "31": 42.76, "32": 131.19, "33": 2.93, "34": 130.04, "35": 44.77, "36": 131.45, "37": 44.06, "38": 131.02, "39": 43.48, "40": 130.42, "41": 45.24, "42": 129.46, "43": 44.6, "44": 130.33, "45": 173.84, "46": 175.58, "47": 168.16, "48": 173.85, "49": 177.56 }, "expected_e2e_ms": 6497.84, "expected_avg_denoise_ms": 91.85, "expected_median_denoise_ms": 120.7 }, "wan2_1_t2v_1.3b": { "stages_ms": { "InputValidationStage": 0.07, "TextEncodingStage": 2237.78, "TimestepPreparationStage": 2.1, "LatentPreparationStage": 0.84, "DenoisingStage": 13041.23, "DecodingStage": 1274.63, "per_frame_generation": null }, "denoise_step_ms": { "0": 224.71, "1": 248.13, "2": 246.48, "3": 247.87, "4": 249.38, "5": 246.76, "6": 250.42, "7": 250.81, "8": 250.98, "9": 249.9, "10": 246.72, "11": 249.79, "12": 250.46, "13": 249.19, "14": 247.55, "15": 250.12, "16": 247.57, "17": 247.21, "18": 247.32, "19": 247.42, "20": 248.21, "21": 247.19, "22": 247.72, "23": 247.45, "24": 247.9, "25": 247.87, "26": 247.18, "27": 247.65, "28": 246.91, "29": 248.26, "30": 247.82, "31": 247.73, "32": 247.38, "33": 247.84, "34": 247.46, "35": 247.52, "36": 247.94, "37": 248.76, "38": 248.01, "39": 247.45, "40": 247.84, "41": 248.33, "42": 247.41, "43": 248.16, "44": 248.18, "45": 248.44, "46": 248.65, "47": 247.73, "48": 247.48, "49": 247.54 }, "expected_e2e_ms": 18382.19, "expected_avg_denoise_ms": 260.76, "expected_median_denoise_ms": 247.84 }, "wan2_1_t2v_1.3b_text_encoder_cpu_offload": { "stages_ms": { "InputValidationStage": 0.09, "TextEncodingStage": 2480.54, "TimestepPreparationStage": 3.73, "LatentPreparationStage": 1.34, "DenoisingStage": 12514.88, "DecodingStage": 1147.6, "per_frame_generation": null }, "denoise_step_ms": { "0": 487.21, "1": 243.47, "2": 244.28, "3": 244.06, "4": 244.77, "5": 245.86, "6": 245.38, "7": 246.74, "8": 246.28, "9": 245.58, "10": 245.6, "11": 245.21, "12": 245.08, "13": 245.03, "14": 245.53, "15": 245.36, "16": 246.17, "17": 245.32, "18": 244.37, "19": 246.83, "20": 245.87, "21": 244.93, "22": 245.11, "23": 245.23, "24": 245.76, "25": 245.44, "26": 246.47, "27": 244.56, "28": 244.76, "29": 244.79, "30": 244.76, "31": 244.8, "32": 245.11, "33": 245.27, "34": 245.37, "35": 245.3, "36": 244.84, "37": 245.26, "38": 245.38, "39": 245.31, "40": 244.7, "41": 245.84, "42": 245.66, "43": 246.68, "44": 245.38, "45": 245.98, "46": 246.02, "47": 245.96, "48": 245.31, "49": 244.99 }, "expected_e2e_ms": 16161.11, "expected_avg_denoise_ms": 250.18, "expected_median_denoise_ms": 245.32 }, "wan2_1_t2v_1.3b_cfg_parallel": { "stages_ms": { "InputValidationStage": 0.08, "TextEncodingStage": 2700.44, "TimestepPreparationStage": 2.82, "LatentPreparationStage": 2.0, "DenoisingStage": 11640.75, "DecodingStage": 890.63, "per_frame_generation": null }, "denoise_step_ms": { "0": 266.91, "1": 211.32, "2": 206.59, "3": 208.12, "4": 210.68, "5": 210.28, "6": 213.92, "7": 211.25, "8": 212.89, "9": 205.35, "10": 205.92, "11": 208.99, "12": 207.1, "13": 208.1, "14": 206.52, "15": 205.5, "16": 205.24, "17": 204.93, "18": 207.05, "19": 203.78, "20": 205.23, "21": 203.87, "22": 204.28, "23": 203.8, "24": 206.02, "25": 207.2, "26": 209.53, "27": 207.46, "28": 206.77, "29": 208.14, "30": 208.05, "31": 208.78, "32": 209.23, "33": 209.72, "34": 208.26, "35": 208.55, "36": 205.24, "37": 204.96, "38": 203.77, "39": 210.2, "40": 202.57, "41": 204.77, "42": 204.96, "43": 203.8, "44": 203.9, "45": 204.49, "46": 207.75, "47": 209.09, "48": 207.51, "49": 207.38 }, "expected_e2e_ms": 15245.6, "expected_avg_denoise_ms": 224.37, "expected_median_denoise_ms": 207.15 }, "turbo_wan2_1_t2v_1.3b": { "stages_ms": { "InputValidationStage": 0.06, "TextEncodingStage": 2508.95, "TimestepPreparationStage": 73.51, "LatentPreparationStage": 1.34, "DmdDenoisingStage": 1285.25, "DecodingStage": 805.04, "per_frame_generation": null }, "denoise_step_ms": { "0": 897.62, "1": 126.04, "2": 126.52, "3": 128.26 }, "expected_e2e_ms": 4686.66, "expected_avg_denoise_ms": 319.61, "expected_median_denoise_ms": 127.39 }, "wan2_2_ti2v_5b": { "stages_ms": { "InputValidationStage": 96.27, "TextEncodingStage": 2238.81, "TimestepPreparationStage": 2.39, "LatentPreparationStage": 27.62, "DenoisingStage": 134069.79, "DecodingStage": 13559.79, "per_frame_generation": null }, "denoise_step_ms": { "0": 3181.0, "1": 2561.67, "2": 2578.49, "3": 2582.1, "4": 2572.24, "5": 2577.72, "6": 2581.35, "7": 2578.79, "8": 2584.98, "9": 2588.49, "10": 2594.37, "11": 2591.19, "12": 2591.32, "13": 2595.35, "14": 2594.35, "15": 2595.62, "16": 2596.35, "17": 2596.11, "18": 2597.24, "19": 2603.13, "20": 2599.9, "21": 2601.48, "22": 2603.58, "23": 2601.13, "24": 2600.47, "25": 2604.13, "26": 2606.04, "27": 2605.3, "28": 2602.02, "29": 2601.83, "30": 2603.57, "31": 2606.63, "32": 2606.1, "33": 2602.24, "34": 2603.29, "35": 2602.34, "36": 2602.16, "37": 2608.14, "38": 2603.48, "39": 2601.7, "40": 2603.96, "41": 2604.58, "42": 2606.67, "43": 2603.52, "44": 2599.88, "45": 2598.66, "46": 2600.74, "47": 2602.31, "48": 2608.4, "49": 2606.02 }, "expected_e2e_ms": 150004.2, "expected_avg_denoise_ms": 2608.84, "expected_median_denoise_ms": 2601.59 }, "qwen_image_edit_2509_ti2i": { "stages_ms": { "InputValidationStage": 213.24, "ImageEncodingStage": 1089.12, "ImageVAEEncodingStage": 304.56, "TimestepPreparationStage": 2.94, "LatentPreparationStage": 0.2, "DenoisingStage": 50724.5, "DecodingStage": 601.02 }, "denoise_step_ms": { "0": 1057.09, "1": 1267.06, "2": 1268.33, "3": 1268.94, "4": 1270.36, "5": 1270.44, "6": 1268.61, "7": 1270.21, "8": 1274.98, "9": 1271.57, "10": 1273.15, "11": 1271.56, "12": 1272.69, "13": 1271.62, "14": 1274.04, "15": 1276.81, "16": 1272.2, "17": 1269.33, "18": 1275.96, "19": 1274.43, "20": 1272.57, "21": 1275.28, "22": 1273.63, "23": 1275.06, "24": 1277.39, "25": 1277.27, "26": 1274.74, "27": 1273.38, "28": 1276.77, "29": 1275.59, "30": 1275.51, "31": 1274.9, "32": 1274.8, "33": 1279.03, "34": 1272.9, "35": 1274.67, "36": 1272.61, "37": 1272.82, "38": 1276.41, "39": 1273.55 }, "expected_e2e_ms": 52938.04, "expected_avg_denoise_ms": 1267.96, "expected_median_denoise_ms": 1273.46 }, "qwen_image_layered_i2i": { "stages_ms": { "QwenImageLayeredBeforeDenoisingStage": 2897.28, "DecodingStage": 312.93, "DenoisingStage": 39417.66, "TimestepPreparationStage": 2.29 }, "denoise_step_ms": { "0": 657.28, "1": 799.2, "2": 790.35, "3": 785.79, "4": 792.9, "5": 795.78, "6": 791.28, "7": 790.87, "8": 786.47, "9": 791.03, "10": 788.77, "11": 790.57, "12": 788.7, "13": 786.01, "14": 791.43, "15": 789.88, "16": 791.18, "17": 792.78, "18": 792.06, "19": 790.47, "20": 792.48, "21": 789.13, "22": 792.12, "23": 789.36, "24": 790.2, "25": 790.87, "26": 792.37, "27": 794.92, "28": 792.9, "29": 791.43, "30": 793.01, "31": 793.71, "32": 794.15, "33": 787.93, "34": 792.12, "35": 794.01, "36": 789.05, "37": 790.51, "38": 793.29, "39": 791.94, "40": 788.94, "41": 788.85, "42": 789.76, "43": 788.89, "44": 791.62, "45": 788.04, "46": 790.03, "47": 786.82, "48": 789.75, "49": 789.0 }, "expected_e2e_ms": 42660.88, "expected_avg_denoise_ms": 788.2, "expected_median_denoise_ms": 790.72 }, "fastwan2_2_ti2v_5b": { "stages_ms": { "InputValidationStage": 300.0, "TextEncodingStage": 843.86, "TimestepPreparationStage": 58.66, "LatentPreparationStage": 28.55, "DmdDenoisingStage": 499.34, "DecodingStage": 1924.01, "per_frame_generation": null }, "denoise_step_ms": { "0": 164.76, "1": 165.6, "2": 165.84 }, "expected_e2e_ms": 7722.91, "expected_avg_denoise_ms": 165.42, "expected_median_denoise_ms": 165.66 }, "fast_hunyuan_video": { "stages_ms": { "InputValidationStage": 0.34, "TextEncodingStage": 550.63, "TimestepPreparationStage": 44.28, "LatentPreparationStage": 0.29, "DenoisingStage": 9154.39, "DecodingStage": 5995.09, "per_frame_generation": null }, "denoise_step_ms": { "0": 485.99, "1": 1399.84, "2": 1399.91, "3": 1397.79, "4": 1400.61, "5": 1402.53 }, "expected_e2e_ms": 16672.15, "expected_avg_denoise_ms": 1608.46, "expected_median_denoise_ms": 1488.48 }, "wan2_2_i2v_a14b_2gpu": { "stages_ms": { "InputValidationStage": 18.45, "TextEncodingStage": 3337.77, "TimestepPreparationStage": 2.9, "LatentPreparationStage": 1.25, "ImageVAEEncodingStage": 1655.89, "DenoisingStage": 106972.82, "DecodingStage": 1355.52, "per_frame_generation": null }, "denoise_step_ms": { "0": 1525.6, "1": 1582.6, "2": 1597.84, "3": 1601.34, "4": 1600.86, "5": 1598.32, "6": 1600.93, "7": 1599.88, "8": 1600.0, "9": 1600.55, "10": 1599.27, "11": 1600.59, "12": 1600.17, "13": 1599.72, "14": 1599.76, "15": 24098.85, "16": 1601.29, "17": 1598.89, "18": 1600.12, "19": 1600.52, "20": 1599.59, "21": 1600.37, "22": 1600.35, "23": 1599.7, "24": 1599.92, "25": 1599.75, "26": 1600.2, "27": 1600.06, "28": 1600.41, "29": 1599.35, "30": 1600.69, "31": 1600.15, "32": 1599.33, "33": 1599.86, "34": 1600.52, "35": 1599.84, "36": 1600.38, "37": 1599.23, "38": 1600.27, "39": 1599.78 }, "expected_e2e_ms": 123182.9887, "expected_avg_denoise_ms": 2831.0, "expected_median_denoise_ms": 1600.09 }, "turbo_wan2_2_i2v_a14b_2gpu": { "stages_ms": { "InputValidationStage": 25.01, "TextEncodingStage": 5198.6, "TimestepPreparationStage": 56.26, "LatentPreparationStage": 1.4, "ImageVAEEncodingStage": 1001.89, "DmdDenoisingStage": 4487.79, "DecodingStage": 821.01, "per_frame_generation": null }, "denoise_step_ms": { "0": 3042.56, "1": 485.88, "2": 721.2, "3": 475.58 }, "expected_e2e_ms": 11605.97, "expected_avg_denoise_ms": 1120.4, "expected_median_denoise_ms": 481.74 }, "wan2_1_i2v_14b_480P_2gpu": { "stages_ms": { "InputValidationStage": 38.23, "TextEncodingStage": 3550.36, "ImageEncodingStage": 3462.55, "TimestepPreparationStage": 2.6, "LatentPreparationStage": 9.73, "ImageVAEEncodingStage": 2290.98, "DenoisingStage": 415021.17, "DecodingStage": 3016.1, "per_frame_generation": null }, "denoise_step_ms": { "0": 10200.25, "1": 8222.39, "2": 8279.38, "3": 8301.48, "4": 8338.87, "5": 8352.39, "6": 8354.64, "7": 8353.64, "8": 8315.58, "9": 8308.48, "10": 8299.65, "11": 8292.7, "12": 8292.73, "13": 8285.21, "14": 8276.06, "15": 8270.41, "16": 8273.04, "17": 8266.04, "18": 8267.7, "19": 8264.06, "20": 8259.32, "21": 8257.26, "22": 8253.02, "23": 8251.77, "24": 8260.97, "25": 8251.39, "26": 8237.43, "27": 8241.33, "28": 8235.96, "29": 8240.6, "30": 8232.48, "31": 8237.85, "32": 8244.3, "33": 8236.79, "34": 8239.83, "35": 8239.89, "36": 8239.12, "37": 8246.74, "38": 8235.67, "39": 8242.77, "40": 8241.17, "41": 8240.24, "42": 8237.01, "43": 8231.26, "44": 8232.85, "45": 8226.56, "46": 8236.98, "47": 8226.73, "48": 8220.49, "49": 8217.04 }, "expected_e2e_ms": 426697.37, "expected_avg_denoise_ms": 8300.19, "expected_median_denoise_ms": 8267.01 }, "wan2_1_i2v_14b_720P_2gpu": { "stages_ms": { "InputValidationStage": 53.67, "TextEncodingStage": 2838, "ImageEncodingStage": 3123.99, "TimestepPreparationStage": 3.39, "LatentPreparationStage": 8.41, "ImageVAEEncodingStage": 2261.05, "DenoisingStage": 417418.12, "DecodingStage": 2968.35 }, "denoise_step_ms": { "0": 11848.08, "1": 8220.3, "2": 8274.3, "3": 8298.9, "4": 8303.34, "5": 8322.44, "6": 8314.37, "7": 8318.54, "8": 8304.94, "9": 8303.04, "10": 8305.22, "11": 8296.22, "12": 8289.2, "13": 8294.19, "14": 8294.87, "15": 8285.96, "16": 8284.98, "17": 8281.61, "18": 8277.35, "19": 8287.46, "20": 8280.3, "21": 8279.18, "22": 8279.37, "23": 8280.16, "24": 8282.67, "25": 8272.14, "26": 8279.37, "27": 8271.66, "28": 8274.6, "29": 8272.88, "30": 8273.76, "31": 8266.17, "32": 8267.77, "33": 8266.88, "34": 8263.14, "35": 8265.97, "36": 8267.76, "37": 8268.03, "38": 8262.24, "39": 8261.4, "40": 8263.65, "41": 8272.46, "42": 8254.9, "43": 8261.03, "44": 8252.92, "45": 8262.49, "46": 8253.67, "47": 8254.92, "48": 8257.08, "49": 8236.56 }, "expected_e2e_ms": 427536.9, "expected_avg_denoise_ms": 8348.21, "expected_median_denoise_ms": 8274.45 }, "wan2_2_t2v_a14b_2gpu": { "stages_ms": { "InputValidationStage": 0.07, "TextEncodingStage": 2575.3, "TimestepPreparationStage": 1.99, "LatentPreparationStage": 1.26, "DenoisingStage": 156678.8406, "DecodingStage": 2702.7, "per_frame_generation": null }, "denoise_step_ms": { "0": 17908.3, "1": 2379.69, "2": 2393.59, "3": 2400.91, "4": 2398.76, "5": 2403.1, "6": 2403.26, "7": 2399.48, "8": 2401.33, "9": 2398.4, "10": 2401.14, "11": 2409.1, "12": 2401.16, "13": 2408.74, "14": 2404.97, "15": 2400.51, "16": 2402.84, "17": 2401.87, "18": 2399.67, "19": 2400.71, "20": 2399.23, "21": 2400.13, "22": 2400.64, "23": 2399.15, "24": 2399.58, "25": 2400.26, "26": 35247.02, "27": 2390.25, "28": 2398.42, "29": 2399.8, "30": 2400.08, "31": 2400.58, "32": 2403.68, "33": 2399.37, "34": 2401.53, "35": 2399.69, "36": 2399.9, "37": 2400.75, "38": 2398.97, "39": 2399.12 }, "expected_e2e_ms": 149864.99, "expected_avg_denoise_ms": 3608.89, "expected_median_denoise_ms": 2400.38 }, "wan2_1_t2v_14b_2gpu": { "stages_ms": { "InputValidationStage": 0.05, "TextEncodingStage": 2310.34, "TimestepPreparationStage": 2.42, "LatentPreparationStage": 27.7, "DenoisingStage": 803631.52, "DecodingStage": 8898.74, "per_frame_generation": null }, "denoise_step_ms": { "0": 17347.88, "1": 15956.93, "2": 16027.54, "3": 16054.15, "4": 16081.46, "5": 16062.7, "6": 16058.56, "7": 16057.58, "8": 16061.04, "9": 16120.97, "10": 16036.84, "11": 16019.6, "12": 16042.29, "13": 16039.87, "14": 16063.0, "15": 16036.16, "16": 16079.82, "17": 16019.7, "18": 16061.5, "19": 16039.95, "20": 16009.42, "21": 16051.01, "22": 16039.31, "23": 16048.22, "24": 16071.41, "25": 16078.75, "26": 16061.78, "27": 16018.39, "28": 16041.44, "29": 16039.64, "30": 16041.89, "31": 16039.6, "32": 16038.97, "33": 15999.48, "34": 16019.93, "35": 16040.27, "36": 16020.3, "37": 16039.38, "38": 15999.4, "39": 16022.15, "40": 16042.32, "41": 16016.62, "42": 15998.92, "43": 16041.48, "44": 15999.63, "45": 16003.21, "46": 15995.91, "47": 16023.52, "48": 16016.64, "49": 16019.6 }, "expected_e2e_ms": 814884.71, "expected_avg_denoise_ms": 16062.92, "expected_median_denoise_ms": 16039.62 }, "wan2_2_t2v_a14b_lora_2gpu": { "stages_ms": { "InputValidationStage": 0.09, "TextEncodingStage": 2552.97, "TimestepPreparationStage": 1.99, "LatentPreparationStage": 1.29, "DenoisingStage": 154340.69, "DecodingStage": 2730.86, "per_frame_generation": null }, "denoise_step_ms": { "0": 26510.7, "1": 2381.25, "2": 2396.9, "3": 2400.96, "4": 2402.47, "5": 2399.6, "6": 2400.5, "7": 2401.13, "8": 2399.32, "9": 2400.0, "10": 2401.35, "11": 2400.04, "12": 2408.27, "13": 2407.08, "14": 2405.92, "15": 2403.99, "16": 2402.12, "17": 2402.52, "18": 2398.08, "19": 2399.9, "20": 2400.14, "21": 2398.64, "22": 2401.32, "23": 2400.75, "24": 2399.27, "25": 2400.21, "26": 36387.55, "27": 2399.77, "28": 2398.09, "29": 2404.64, "30": 2400.68, "31": 2404.3, "32": 2392.44, "33": 2390.56, "34": 2396.05, "35": 2394.86, "36": 2396.07, "37": 2398.49, "38": 2394.77, "39": 2394.19 }, "expected_e2e_ms": 159643.06, "expected_avg_denoise_ms": 3851.87, "expected_median_denoise_ms": 2400.09 }, "wan2_1_t2v_1_3b_lora_1gpu": { "stages_ms": { "InputValidationStage": 0.06, "TextEncodingStage": 2467.44, "TimestepPreparationStage": 2.96, "LatentPreparationStage": 1.87, "DenoisingStage": 14859.47, "DecodingStage": 1199.31, "per_frame_generation": null }, "denoise_step_ms": { "0": 233.29, "1": 265.02, "2": 257.83, "3": 260.27, "4": 261.43, "5": 258.58, "6": 256.64, "7": 256.91, "8": 258.41, "9": 257.84, "10": 257.08, "11": 257.0, "12": 258.44, "13": 257.1, "14": 256.95, "15": 257.2, "16": 256.84, "17": 257.64, "18": 257.22, "19": 257.42, "20": 256.91, "21": 256.99, "22": 257.17, "23": 257.63, "24": 258.89, "25": 257.46, "26": 257.3, "27": 257.42, "28": 257.19, "29": 257.65, "30": 257.39, "31": 256.93, "32": 258.23, "33": 257.62, "34": 281.86, "35": 295.86, "36": 296.73, "37": 287.21, "38": 300.87, "39": 303.47, "40": 294.09, "41": 270.52, "42": 256.53, "43": 256.58, "44": 256.29, "45": 255.81, "46": 256.34, "47": 256.08, "48": 255.92, "49": 255.87 }, "expected_e2e_ms": 18547.46, "expected_avg_denoise_ms": 297.09, "expected_median_denoise_ms": 257.42 }, "wan2_1_i2v_14b_lora_2gpu": { "stages_ms": { "InputValidationStage": 23.97, "TextEncodingStage": 2485.39, "ImageEncodingStage": 2372.07, "TimestepPreparationStage": 2.6, "LatentPreparationStage": 0.18, "ImageVAEEncodingStage": 2500.13, "DenoisingStage": 193514.04, "DecodingStage": 3341.78, "per_frame_generation": null }, "denoise_step_ms": { "0": 6680.62, "1": 3765.8, "2": 3774.63, "3": 3772.93, "4": 3781.13, "5": 3778.22, "6": 3776.41, "7": 3772.02, "8": 3776.15, "9": 3768.82, "10": 3775.31, "11": 3771.32, "12": 3774.33, "13": 3772.5, "14": 3778.41, "15": 3775.31, "16": 3771.38, "17": 3774.87, "18": 3780.01, "19": 3772.85, "20": 3773.65, "21": 3774.47, "22": 3774.39, "23": 3773.08, "24": 3776.71, "25": 3780.01, "26": 3774.83, "27": 3773.27, "28": 3773.76, "29": 3772.75, "30": 3773.01, "31": 3773.34, "32": 3773.13, "33": 3774.12, "34": 3772.19, "35": 3774.7, "36": 3773.98, "37": 3772.47, "38": 3771.72, "39": 3774.07, "40": 3773.71, "41": 3773.6, "42": 3772.12, "43": 3773.75, "44": 3782.43, "45": 3779.66, "46": 3779.86, "47": 3774.58, "48": 3770.54, "49": 3776.76 }, "expected_e2e_ms": 204257.12, "expected_avg_denoise_ms": 3855.55, "expected_median_denoise_ms": 3774.03 }, "flux_2_image_t2i_2_gpus": { "stages_ms": { "InputValidationStage": 0.05, "TextEncodingStage": 518.88, "ImageVAEEncodingStage": 0.0, "LatentPreparationStage": 0.45, "TimestepPreparationStage": 3.41, "DenoisingStage": 26377.63, "DecodingStage": 321.94 }, "denoise_step_ms": { "0": 129.07, "1": 437.16, "2": 437.7, "3": 437.67, "4": 437.84, "5": 438.03, "6": 438.09, "7": 437.65, "8": 437.95, "9": 438.31, "10": 437.99, "11": 438.54, "12": 438.47, "13": 438.2, "14": 438.56, "15": 438.69, "16": 438.69, "17": 438.98, "18": 437.96, "19": 438.9, "20": 438.87, "21": 438.04, "22": 437.88, "23": 439.09, "24": 438.61, "25": 437.68, "26": 439.2, "27": 439.63, "28": 438.65, "29": 439.32, "30": 439.01, "31": 438.84, "32": 438.72, "33": 439.09, "34": 438.3, "35": 439.48, "36": 438.2, "37": 439.67, "38": 440.65, "39": 439.96, "40": 439.0, "41": 439.2, "42": 439.37, "43": 439.98, "44": 438.6, "45": 439.58, "46": 440.23, "47": 440.1, "48": 440.21, "49": 439.22 }, "expected_e2e_ms": 27624.8, "expected_avg_denoise_ms": 518.23, "expected_median_denoise_ms": 528.06 }, "qwen_image_edit_2511_ti2i": { "stages_ms": { "InputValidationStage": 55.15, "ImageEncodingStage": 770.33, "ImageVAEEncodingStage": 88.06, "TimestepPreparationStage": 2.12, "LatentPreparationStage": 0.14, "DenoisingStage": 23869.32, "DecodingStage": 108.23 }, "denoise_step_ms": { "0": 478.35, "1": 608.56, "2": 588.51, "3": 607.26, "4": 599.37, "5": 595.19, "6": 603.22, "7": 594.48, "8": 605.06, "9": 597.63, "10": 601.03, "11": 597.18, "12": 598.82, "13": 600.05, "14": 598.57, "15": 601.4, "16": 595.17, "17": 599.21, "18": 600.86, "19": 600.93, "20": 600.35, "21": 600.63, "22": 597.58, "23": 600.73, "24": 599.36, "25": 600.48, "26": 600.33, "27": 599.34, "28": 599.61, "29": 599.71, "30": 596.03, "31": 599.85, "32": 599.36, "33": 601.58, "34": 597.91, "35": 600.79, "36": 599.29, "37": 601.64, "38": 598.24, "39": 599.87 }, "expected_e2e_ms": 24895.28, "expected_avg_denoise_ms": 596.59, "expected_median_denoise_ms": 599.66 }, "fsdp-inference": { "stages_ms": { "InputValidationStage": 0.04, "TextEncodingStage": 411.12, "TimestepPreparationStage": 1.44, "LatentPreparationStage": 0.1, "DenoisingStage": 1569.61, "DecodingStage": 41.43 }, "denoise_step_ms": { "0": 165.33, "1": 158.34, "2": 167.65, "3": 179.11, "4": 183.98, "5": 175.08, "6": 178.34, "7": 178.53, "8": 178.08 }, "expected_e2e_ms": 2103.05, "expected_avg_denoise_ms": 173.83, "expected_median_denoise_ms": 178.08 }, "hunyuan3d_shape_gen": { "stages_ms": { "Hunyuan3DShapeBeforeDenoisingStage": 31.42, "Hunyuan3DShapeDenoisingStage": 3259.83, "Hunyuan3DShapeExportStage": 8735.55, "Hunyuan3DShapeSaveStage": 981.64, "Hunyuan3DPaintPreprocessStage": 226071.67, "Hunyuan3DPaintTexGenStage": 11083.05, "Hunyuan3DPaintPostprocessStage": 7469.29 }, "denoise_step_ms": { "0": 32.26, "1": 63.34, "2": 65.44, "3": 65.44, "4": 65.6, "5": 65.81, "6": 65.82, "7": 65.48, "8": 65.9, "9": 65.77, "10": 65.54, "11": 65.68, "12": 65.85, "13": 65.77, "14": 65.7, "15": 65.78, "16": 66.0, "17": 66.15, "18": 65.91, "19": 66.5, "20": 65.76, "21": 66.08, "22": 66.06, "23": 66.23, "24": 65.79, "25": 65.58, "26": 65.88, "27": 65.67, "28": 65.87, "29": 66.09, "30": 65.81, "31": 65.91, "32": 66.18, "33": 65.93, "34": 66.26, "35": 66.26, "36": 66.27, "37": 65.57, "38": 66.02, "39": 66.19, "40": 65.23, "41": 66.11, "42": 66.18, "43": 65.86, "44": 65.86, "45": 65.92, "46": 65.65, "47": 65.78, "48": 66.01, "49": 66.08 }, "expected_e2e_ms": 257696.97, "expected_avg_denoise_ms": 65.16, "expected_median_denoise_ms": 65.86 }, "wan2_1_t2v_1.3b_frame_interp_2x": { "stages_ms": { "TextEncodingStage": 1104.4, "TimestepPreparationStage": 2.19, "LatentPreparationStage": 0.15, "DenoisingStage": 8502.22, "DecodingStage": 498.36, "InputValidationStage": 0.07 }, "denoise_step_ms": { "0": 91.83, "1": 174.57, "2": 170.48, "3": 169.33, "4": 169.24, "5": 177.43, "6": 173.73, "7": 171.67, "8": 170.98, "9": 168.61, "10": 169.96, "11": 174.75, "12": 172.33, "13": 170.62, "14": 169.84, "15": 168.86, "16": 171.32, "17": 174.7, "18": 172.31, "19": 171.71, "20": 170.98, "21": 169.83, "22": 170.54, "23": 173.08, "24": 172.11, "25": 171.49, "26": 171.0, "27": 170.9, "28": 171.78, "29": 173.44, "30": 171.14, "31": 170.72, "32": 170.64, "33": 170.58, "34": 172.51, "35": 171.74, "36": 171.57, "37": 170.73, "38": 171.49, "39": 170.98, "40": 172.63, "41": 171.88, "42": 171.71, "43": 170.94, "44": 170.31, "45": 171.25, "46": 171.43, "47": 171.55, "48": 172.08, "49": 169.92 }, "expected_e2e_ms": 10464.97, "expected_avg_denoise_ms": 169.92, "expected_median_denoise_ms": 171.37 }, "flux_2_klein_ti2i_2_gpus": { "stages_ms": { "InputValidationStage": 40.19, "TextEncodingStage": 88.84, "ImageVAEEncodingStage": 80.81, "LatentPreparationStage": 1.05, "TimestepPreparationStage": 28.64, "DenoisingStage": 354.04, "DecodingStage": 11.11 }, "denoise_step_ms": { "0": 33.54, "1": 61.3, "2": 86.9, "3": 87.55 }, "expected_e2e_ms": 716.81, "expected_avg_denoise_ms": 67.32, "expected_median_denoise_ms": 74.1 }, "flux_2_image_t2i_upscaling_4x": { "stages_ms": { "InputValidationStage": 0.15, "TextEncodingStage": 537.4, "ImageVAEEncodingStage": 0.01, "LatentPreparationStage": 1.11, "TimestepPreparationStage": 39.57, "DenoisingStage": 24738.35, "DecodingStage": 14.04 }, "denoise_step_ms": { "0": 70.71, "1": 476.99, "2": 502.01, "3": 483.39, "4": 500.65, "5": 487.43, "6": 502.39, "7": 484.76, "8": 498.5, "9": 489.98, "10": 499.22, "11": 490.47, "12": 498.98, "13": 491.42, "14": 495.17, "15": 492.24, "16": 494.69, "17": 491.68, "18": 497.13, "19": 493.63, "20": 495.29, "21": 496.19, "22": 496.52, "23": 496.31, "24": 493.64, "25": 494.35, "26": 493.27, "27": 495.52, "28": 493.06, "29": 494.66, "30": 494.08, "31": 496.01, "32": 494.79, "33": 495.81, "34": 493.94, "35": 495.56, "36": 493.5, "37": 495.98, "38": 495.82, "39": 496.7, "40": 495.29, "41": 496.84, "42": 495.67, "43": 495.32, "44": 496.48, "45": 496.03, "46": 495.65, "47": 498.27, "48": 496.44, "49": 496.79 }, "expected_e2e_ms": 25735.08, "expected_avg_denoise_ms": 486.1, "expected_median_denoise_ms": 495.42 }, "wan2_1_t2v_1.3b_upscaling_4x": { "stages_ms": { "InputValidationStage": 0.08, "TextEncodingStage": 1164.21, "LatentPreparationStage": 0.26, "TimestepPreparationStage": 3.86, "DenoisingStage": 10234.35, "DecodingStage": 499.44 }, "denoise_step_ms": { "0": 200.87, "1": 202.85, "2": 203.32, "3": 206.02, "4": 205.75, "5": 204.14, "6": 205.12, "7": 204.65, "8": 203.99, "9": 204.96, "10": 204.35, "11": 206.89, "12": 200.69, "13": 209.67, "14": 204.91, "15": 203.5, "16": 206.73, "17": 202.43, "18": 205.92, "19": 204.61, "20": 211.47, "21": 197.43, "22": 203.58, "23": 205.82, "24": 204.01, "25": 205.06, "26": 204.86, "27": 206.03, "28": 200.78, "29": 206.99, "30": 206.58, "31": 202.84, "32": 204.51, "33": 204.19, "34": 202.89, "35": 204.55, "36": 205.03, "37": 204.2, "38": 203.92, "39": 204.9, "40": 203.24, "41": 204.21, "42": 205.76, "43": 205.32, "44": 202.63, "45": 205.67, "46": 204.55, "47": 202.89, "48": 205.29, "49": 203.87 }, "expected_e2e_ms": 12021.58, "expected_avg_denoise_ms": 204.49, "expected_median_denoise_ms": 204.55 }, "wan2_1_t2v_1.3b_frame_interp_2x_upscaling_4x": { "stages_ms": { "InputValidationStage": 0.03, "TextEncodingStage": 1089.94, "LatentPreparationStage": 0.12, "TimestepPreparationStage": 1.99, "DenoisingStage": 8617.56, "DecodingStage": 469.68 }, "denoise_step_ms": { "0": 122.51, "1": 172.31, "2": 169.58, "3": 171.76, "4": 171.5, "5": 174.24, "6": 175.85, "7": 171.2, "8": 172.0, "9": 172.13, "10": 171.22, "11": 174.8, "12": 174.01, "13": 172.28, "14": 172.77, "15": 173.52, "16": 172.19, "17": 175.12, "18": 172.78, "19": 175.1, "20": 171.54, "21": 173.38, "22": 171.61, "23": 174.14, "24": 174.06, "25": 172.32, "26": 173.08, "27": 173.94, "28": 173.32, "29": 174.3, "30": 173.63, "31": 172.21, "32": 174.4, "33": 173.25, "34": 173.54, "35": 175.12, "36": 172.93, "37": 172.76, "38": 174.73, "39": 174.46, "40": 172.66, "41": 174.58, "42": 173.9, "43": 174.88, "44": 172.35, "45": 173.52, "46": 175.94, "47": 172.88, "48": 174.97, "49": 172.94 }, "expected_e2e_ms": 10425.77, "expected_avg_denoise_ms": 172.28, "expected_median_denoise_ms": 173.29 }, "helios_base_t2v": { "stages_ms": { "InputValidationStage": 0.07, "TextEncodingStage": 1103.97, "LatentPreparationStage": 0.24, "HeliosChunkedDenoisingStage": 118580.37, "HeliosDecodingStage": 664.79, "per_frame_generation": null }, "denoise_step_ms": {}, "expected_e2e_ms": 120413.51, "expected_avg_denoise_ms": 0.0, "expected_median_denoise_ms": 0.0 }, "helios_distilled_t2v": { "stages_ms": { "InputValidationStage": 0.13, "TextEncodingStage": 581.79, "LatentPreparationStage": 0.18, "HeliosChunkedDenoisingStage": 49752.88, "HeliosDecodingStage": 666.69, "per_frame_generation": null }, "denoise_step_ms": {}, "expected_e2e_ms": 51038.66, "expected_avg_denoise_ms": 0.0, "expected_median_denoise_ms": 0.0 }, "helios_mid_t2v": { "stages_ms": { "InputValidationStage": 0.05, "TextEncodingStage": 1101.99, "LatentPreparationStage": 0.16, "HeliosChunkedDenoisingStage": 77728.72, "HeliosDecodingStage": 661.23, "per_frame_generation": null }, "denoise_step_ms": {}, "expected_e2e_ms": 79600.62, "expected_avg_denoise_ms": 0.0, "expected_median_denoise_ms": 0.0 }, "helios_base_t2v": { "stages_ms": { "InputValidationStage": 0.04, "TextEncodingStage": 1102.45, "LatentPreparationStage": 0.14, "HeliosChunkedDenoisingStage": 116964.69, "HeliosDecodingStage": 664.76, "per_frame_generation": null }, "denoise_step_ms": { "0": 1893.3, "1": 1900.93, "2": 1934.08, "3": 1897.65, "4": 1907.59, "5": 1909.1, "6": 1911.51, "7": 1909.25, "8": 1911.69, "9": 1911.77, "10": 1913.35, "11": 1915.44, "12": 1912.11, "13": 1910.08, "14": 1911.77, "15": 1908.22, "16": 1908.83, "17": 1910.11, "18": 1908.19, "19": 1911.99, "20": 1909.96, "21": 1910.32, "22": 1911.76, "23": 1911.87, "24": 1908.91, "25": 1912.41, "26": 1913.15, "27": 1908.34, "28": 1913.21, "29": 1911.98, "30": 1912.16, "31": 1914.17, "32": 1911.45, "33": 1912.5, "34": 1914.48, "35": 1912.64, "36": 1912.24, "37": 1914.48, "38": 1911.06, "39": 1915.45, "40": 1914.0, "41": 1912.99, "42": 1913.68, "43": 1914.09, "44": 1915.83, "45": 1913.36, "46": 1914.84, "47": 1915.31, "48": 1915.58, "49": 1912.63 }, "expected_e2e_ms": 118821.41, "expected_avg_denoise_ms": 1911.64, "expected_median_denoise_ms": 1912.05 }, "helios_mid_t2v": { "stages_ms": { "InputValidationStage": 0.09, "TextEncodingStage": 1102.28, "LatentPreparationStage": 0.23, "HeliosChunkedDenoisingStage": 77947.9, "HeliosDecodingStage": 664.96, "per_frame_generation": null }, "denoise_step_ms": { "0": 404.46, "1": 404.88, "2": 405.35, "3": 406.01, "4": 404.97, "5": 405.07, "6": 405.06, "7": 404.98, "8": 405.39, "9": 405.52, "10": 405.76, "11": 405.53, "12": 405.16, "13": 405.46, "14": 405.75, "15": 405.69, "16": 405.26, "17": 405.23, "18": 405.42, "19": 405.99, "20": 663.39, "21": 666.6, "22": 665.73, "23": 666.37, "24": 667.43, "25": 668.28, "26": 667.96, "27": 668.93, "28": 667.78, "29": 668.15, "30": 668.91, "31": 667.22, "32": 669.31, "33": 666.57, "34": 669.78, "35": 668.38, "36": 669.95, "37": 668.76, "38": 667.82, "39": 668.98, "40": 1891.05, "41": 1893.52, "42": 1893.48, "43": 1892.79, "44": 1892.03, "45": 1892.87, "46": 1895.55, "47": 1892.19, "48": 1892.89, "49": 1892.32, "50": 1890.25, "51": 1894.1, "52": 1890.67, "53": 1892.09, "54": 1892.64, "55": 1891.91, "56": 1894.27, "57": 1893.62, "58": 1892.65, "59": 1891.9 }, "expected_e2e_ms": 79824.32, "expected_avg_denoise_ms": 988.6, "expected_median_denoise_ms": 668.05 }, "helios_distilled_t2v": { "stages_ms": { "InputValidationStage": 0.05, "TextEncodingStage": 552.02, "LatentPreparationStage": 0.13, "HeliosChunkedDenoisingStage": 57879.88, "HeliosDecodingStage": 663.31, "per_frame_generation": null }, "denoise_step_ms": { "0": 207.03, "1": 204.36, "2": 203.87, "3": 204.51, "4": 206.21, "5": 205.54, "6": 205.06, "7": 205.45, "8": 205.96, "9": 205.95, "10": 205.22, "11": 204.43, "12": 205.14, "13": 205.06, "14": 205.11, "15": 206.09, "16": 205.1, "17": 204.99, "18": 204.55, "19": 205.14, "20": 337.47, "21": 337.06, "22": 337.68, "23": 336.58, "24": 335.98, "25": 335.84, "26": 336.01, "27": 335.61, "28": 335.79, "29": 335.62, "30": 336.69, "31": 335.98, "32": 336.15, "33": 336.55, "34": 336.98, "35": 337.33, "36": 336.34, "37": 335.94, "38": 336.69, "39": 336.14, "40": 954.88, "41": 956.2, "42": 953.9, "43": 953.49, "44": 957.1, "45": 956.95, "46": 955.02, "47": 954.98, "48": 956.0, "49": 956.63, "50": 958.66, "51": 957.26, "52": 956.73, "53": 955.06, "54": 957.04, "55": 958.07, "56": 958.28, "57": 957.99, "58": 957.61, "59": 956.98 }, "expected_e2e_ms": 59168.9, "expected_avg_denoise_ms": 499.37, "expected_median_denoise_ms": 336.25 } } } ================================================ FILE: python/sglang/multimodal_gen/test/server/test_server_2_gpu_a.py ================================================ """ 2 GPU tests """ from __future__ import annotations import pytest from sglang.multimodal_gen.test.server.test_server_common import ( # noqa: F401 DiffusionServerBase, diffusion_server, ) from sglang.multimodal_gen.test.server.testcase_configs import ( TWO_GPU_CASES_A, DiffusionTestCase, ) class TestDiffusionServerTwoGpu(DiffusionServerBase): """Performance tests for 2-GPU diffusion cases.""" @pytest.fixture(params=TWO_GPU_CASES_A, ids=lambda c: c.id) def case(self, request) -> DiffusionTestCase: """Provide a DiffusionTestCase for each 2-GPU test.""" return request.param ================================================ FILE: python/sglang/multimodal_gen/test/server/test_server_2_gpu_b.py ================================================ """ 2 GPU tests """ from __future__ import annotations import pytest from sglang.multimodal_gen.test.server.test_server_common import ( # noqa: F401 DiffusionServerBase, diffusion_server, ) from sglang.multimodal_gen.test.server.testcase_configs import ( TWO_GPU_CASES_B, DiffusionTestCase, ) class TestDiffusionServerTwoGpu(DiffusionServerBase): """Performance tests for 2-GPU diffusion cases.""" @pytest.fixture(params=TWO_GPU_CASES_B, ids=lambda c: c.id) def case(self, request) -> DiffusionTestCase: """Provide a DiffusionTestCase for each 2-GPU test.""" return request.param ================================================ FILE: python/sglang/multimodal_gen/test/server/test_server_a.py ================================================ """ Config-driven diffusion performance test with pytest parametrization. If the actual run is significantly better than the baseline, the improved cases with their updated baseline will be printed """ from __future__ import annotations import pytest from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.test.server.test_server_common import ( # noqa: F401 DiffusionServerBase, diffusion_server, ) from sglang.multimodal_gen.test.server.testcase_configs import ( ONE_GPU_CASES_A, DiffusionTestCase, ) logger = init_logger(__name__) class TestDiffusionServerOneGpu(DiffusionServerBase): """Performance tests for 1-GPU diffusion cases.""" @pytest.fixture(params=ONE_GPU_CASES_A, ids=lambda c: c.id) def case(self, request) -> DiffusionTestCase: """Provide a DiffusionTestCase for each 1-GPU test.""" return request.param ================================================ FILE: python/sglang/multimodal_gen/test/server/test_server_b.py ================================================ """ Config-driven diffusion performance test with pytest parametrization. If the actual run is significantly better than the baseline, the improved cases with their updated baseline will be printed """ from __future__ import annotations import pytest from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.test.server.test_server_common import ( # noqa: F401 DiffusionServerBase, diffusion_server, ) from sglang.multimodal_gen.test.server.testcase_configs import ( ONE_GPU_CASES_B, DiffusionTestCase, ) logger = init_logger(__name__) class TestDiffusionServerOneGpu(DiffusionServerBase): """Performance tests for 1-GPU diffusion cases.""" @pytest.fixture(params=ONE_GPU_CASES_B, ids=lambda c: c.id) def case(self, request) -> DiffusionTestCase: """Provide a DiffusionTestCase for each 1-GPU test.""" return request.param ================================================ FILE: python/sglang/multimodal_gen/test/server/test_server_common.py ================================================ """ Config-driven diffusion generation test with pytest parametrization. If the actual run is significantly better than the baseline, the improved cases with their updated baseline will be printed """ from __future__ import annotations import os from pathlib import Path from typing import Any, Callable import openai import pytest import requests from openai import OpenAI from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord from sglang.multimodal_gen.test.server import conftest from sglang.multimodal_gen.test.server.test_server_utils import ( VALIDATOR_REGISTRY, PerformanceValidator, ServerContext, ServerManager, get_generate_fn, ) from sglang.multimodal_gen.test.server.testcase_configs import ( BASELINE_CONFIG, DiffusionTestCase, PerformanceSummary, ScenarioConfig, ) from sglang.multimodal_gen.test.test_utils import ( _consistency_gt_filenames, extract_key_frames_from_video, get_dynamic_server_port, wait_for_req_perf_record, ) logger = init_logger(__name__) @pytest.fixture def diffusion_server(case: DiffusionTestCase) -> ServerContext: """Start a diffusion server for a single case and tear it down afterwards.""" server_args = case.server_args # Skip ring attention tests on AMD/ROCm - Ring Attention requires Flash Attention # which is not available on AMD. Use Ulysses parallelism instead. if ( current_platform.is_hip() and server_args.ring_degree is not None and server_args.ring_degree > 1 ): pytest.skip( f"Skipping {case.id}: Ring Attention (ring_degree={server_args.ring_degree}) " "requires Flash Attention which is not available on AMD/ROCm" ) default_port = get_dynamic_server_port() port = int(os.environ.get("SGLANG_TEST_SERVER_PORT", default_port)) sampling_params = case.sampling_params extra_args = os.environ.get("SGLANG_TEST_SERVE_ARGS", "") # In GT generation mode, force --backend diffusers if os.environ.get("SGLANG_GEN_GT", "0") == "1": if "--backend" not in extra_args: extra_args = "--backend diffusers " + extra_args.strip() extra_args += f" --num-gpus {server_args.num_gpus}" if server_args.tp_size is not None: extra_args += f" --tp-size {server_args.tp_size}" if server_args.ulysses_degree is not None: extra_args += f" --ulysses-degree {server_args.ulysses_degree}" if server_args.dit_layerwise_offload: extra_args += f" --dit-layerwise-offload true" if server_args.dit_offload_prefetch_size: extra_args += ( f" --dit-offload-prefetch-size {server_args.dit_offload_prefetch_size}" ) if server_args.text_encoder_cpu_offload: extra_args += f" --text-encoder-cpu-offload" if server_args.ring_degree is not None: extra_args += f" --ring-degree {server_args.ring_degree}" # LoRA support if server_args.lora_path: extra_args += f" --lora-path {server_args.lora_path}" # default warmup extra_args += f" --warmup" for arg in server_args.extras: extra_args += f" {arg}" # Build custom environment variables env_vars = {} if server_args.enable_cache_dit: env_vars["SGLANG_CACHE_DIT_ENABLED"] = "true" # start server manager = ServerManager( model=server_args.model_path, port=port, wait_deadline=float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200")), extra_args=extra_args, env_vars=env_vars, ) ctx = manager.start() try: # Reconstruct output size for OpenAI API # Allow override via environment variable (useful for AMD where large resolutions can cause GPU hang) output_size = os.environ.get( "SGLANG_TEST_OUTPUT_SIZE", sampling_params.output_size ) except Exception as exc: logger.error("Warm-up failed for %s: %s", case.id, exc) ctx.cleanup() raise try: yield ctx finally: ctx.cleanup() class DiffusionServerBase: """Performance tests for all diffusion models/scenarios. This single test class runs against all cases defined in ONE_GPU_CASES. Each case gets its own server instance via the parametrized fixture. """ _perf_results: list[dict[str, Any]] = [] _improved_baselines: list[dict[str, Any]] = [] _pytest_config = None # Store pytest config for stash access @classmethod def setup_class(cls): cls._perf_results = [] cls._improved_baselines = [] @classmethod def teardown_class(cls): print( f"\n[DEBUG teardown_class] Called for {cls.__name__}, _perf_results has {len(cls._perf_results)} entries" ) if cls._pytest_config: # Add results to pytest stash (shared across all import contexts) for result in cls._perf_results: result["class_name"] = cls.__name__ conftest.add_perf_results(cls._pytest_config, cls._perf_results) print( f"[DEBUG teardown_class] Added {len(cls._perf_results)} results to stash" ) else: print( "[DEBUG teardown_class] No pytest_config available, skipping stash update" ) if cls._improved_baselines: import json output = """ --- POTENTIAL BASELINE IMPROVEMENTS DETECTED --- The following test cases performed significantly better than their baselines. Consider updating perf_baselines.json with the snippets below: """ for item in cls._improved_baselines: output += ( f'\n"{item["id"]}": {json.dumps(item["baseline"], indent=4)},\n' ) print(output) @pytest.fixture(autouse=True) def _capture_pytest_config(self, request): """Capture pytest config for use in teardown_class.""" self.__class__._pytest_config = request.config def _client(self, ctx: ServerContext) -> OpenAI: """Get OpenAI client for the server.""" return OpenAI( api_key="sglang-anything", base_url=f"http://localhost:{ctx.port}/v1", ) def run_and_collect( self, ctx: ServerContext, case_id: str, generate_fn: Callable[[str, openai.Client], tuple[str, bytes]], ) -> tuple[RequestPerfRecord, bytes]: """Run generation and collect performance records. Returns: Tuple of (performance_record, content_bytes) """ log_path = ctx.perf_log_path log_wait_timeout = 30 client = self._client(ctx) rid, content = generate_fn(case_id, client) req_perf_record = wait_for_req_perf_record( rid, log_path, timeout=log_wait_timeout, ) return (req_perf_record, content) def _validate_and_record( self, case: DiffusionTestCase, perf_record: RequestPerfRecord, ) -> None: """Validate metrics and record results.""" is_baseline_generation_mode = os.environ.get("SGLANG_GEN_BASELINE", "0") == "1" scenario = BASELINE_CONFIG.scenarios.get(case.id) missing_scenario = False if scenario is None: # Create dummy scenario to allow metric collection scenario = type( "DummyScenario", (), { "expected_e2e_ms": 0, "expected_avg_denoise_ms": 0, "expected_median_denoise_ms": 0, "stages_ms": {}, "denoise_step_ms": {}, }, )() if not is_baseline_generation_mode: missing_scenario = True validator_name = case.server_args.custom_validator or "default" validator_class = VALIDATOR_REGISTRY.get(validator_name, PerformanceValidator) validator = validator_class( scenario=scenario, tolerances=BASELINE_CONFIG.tolerances, step_fractions=BASELINE_CONFIG.step_fractions, ) summary = validator.collect_metrics(perf_record) if case.run_perf_check: if is_baseline_generation_mode or missing_scenario: self._dump_baseline_for_testcase(case, summary, missing_scenario) if missing_scenario: pytest.fail( f"Testcase '{case.id}' not found in perf_baselines.json" ) return self._check_for_improvement(case, summary, scenario) # only run performance validation if run_perf_check is True try: validator.validate(perf_record, case.sampling_params.num_frames) except AssertionError as e: logger.error(f"Performance validation failed for {case.id}:\n{e}") self._dump_baseline_for_testcase(case, summary, missing_scenario) raise result = { "test_name": case.id, "modality": case.server_args.modality, "e2e_ms": summary.e2e_ms, "avg_denoise_ms": summary.avg_denoise_ms, "median_denoise_ms": summary.median_denoise_ms, "stage_metrics": summary.stage_metrics, "sampled_steps": summary.sampled_steps, } # video-specific metrics if summary.frames_per_second: result.update( { "frames_per_second": summary.frames_per_second, "total_frames": summary.total_frames, "avg_frame_time_ms": summary.avg_frame_time_ms, } ) self.__class__._perf_results.append(result) print( f"[DEBUG _validate_and_record] Appended result for {case.id}, class {self.__class__.__name__} now has {len(self.__class__._perf_results)} results" ) def _check_for_improvement( self, case: DiffusionTestCase, summary: PerformanceSummary, scenario: "ScenarioConfig", ) -> None: """Check for potential significant performance improvements and record them.""" is_improved = False threshold = BASELINE_CONFIG.improvement_threshold def is_sig_faster(actual, expected): if expected == 0 or expected is None: return False return actual < expected * (1 - threshold) def safe_get_metric(metric_dict, key): val = metric_dict.get(key) return val if val is not None else float("inf") # Check for any significant improvement if ( is_sig_faster(summary.e2e_ms, scenario.expected_e2e_ms) or is_sig_faster(summary.avg_denoise_ms, scenario.expected_avg_denoise_ms) or is_sig_faster( summary.median_denoise_ms, scenario.expected_median_denoise_ms ) ): is_improved = True # Combine metrics, always taking the better (lower) value new_stages = { stage: min( safe_get_metric(summary.stage_metrics, stage), safe_get_metric(scenario.stages_ms, stage), ) for stage in set(summary.stage_metrics) | set(scenario.stages_ms) } new_denoise_steps = { step: min( safe_get_metric(summary.all_denoise_steps, step), safe_get_metric(scenario.denoise_step_ms, step), ) for step in set(summary.all_denoise_steps.keys()) | set(scenario.denoise_step_ms) } # Check for stage-level improvements if not is_improved: for stage, new_val in new_stages.items(): if is_sig_faster(new_val, scenario.stages_ms.get(stage, float("inf"))): is_improved = True break if not is_improved: for step, new_val in new_denoise_steps.items(): if is_sig_faster( new_val, scenario.denoise_step_ms.get(step, float("inf")) ): is_improved = True break if is_improved: new_baseline = { "stages_ms": {k: round(v, 2) for k, v in new_stages.items()}, "denoise_step_ms": { str(k): round(v, 2) for k, v in new_denoise_steps.items() }, "expected_e2e_ms": round( min(summary.e2e_ms, scenario.expected_e2e_ms), 2 ), "expected_avg_denoise_ms": round( min(summary.avg_denoise_ms, scenario.expected_avg_denoise_ms), 2 ), "expected_median_denoise_ms": round( min(summary.median_denoise_ms, scenario.expected_median_denoise_ms), 2, ), } self._improved_baselines.append({"id": case.id, "baseline": new_baseline}) def _dump_baseline_for_testcase( self, case: DiffusionTestCase, summary: "PerformanceSummary", missing_scenario: bool = False, ) -> None: """Dump performance metrics as a JSON scenario for baselines.""" import json denoise_steps_formatted = { str(k): round(v, 2) for k, v in summary.all_denoise_steps.items() } stages_formatted = {k: round(v, 2) for k, v in summary.stage_metrics.items()} baseline = { "stages_ms": stages_formatted, "denoise_step_ms": denoise_steps_formatted, "expected_e2e_ms": round(summary.e2e_ms, 2), "expected_avg_denoise_ms": round(summary.avg_denoise_ms, 2), "expected_median_denoise_ms": round(summary.median_denoise_ms, 2), } # Video-specific metrics if case.server_args.modality == "video": if "per_frame_generation" not in baseline["stages_ms"]: baseline["stages_ms"]["per_frame_generation"] = ( round(summary.avg_frame_time_ms, 2) if summary.avg_frame_time_ms else None ) action = "add" if missing_scenario else "update" output = f""" {action} this baseline in the "scenarios" section of perf_baselines.json: "{case.id}": {json.dumps(baseline, indent=4)} """ logger.error(output) def _save_gt_output( self, case: DiffusionTestCase, content: bytes, ) -> None: """Save generated content as ground truth files. Args: case: Test case configuration content: Generated content bytes (image or video) """ gt_output_dir = os.environ.get("SGLANG_GT_OUTPUT_DIR") if not gt_output_dir: logger.error("SGLANG_GT_OUTPUT_DIR not set, cannot save GT output") return out_dir = Path(gt_output_dir) out_dir.mkdir(parents=True, exist_ok=True) num_gpus = case.server_args.num_gpus is_video = case.server_args.modality == "video" if is_video: # Extract key frames from video frames = extract_key_frames_from_video( content, num_frames=case.sampling_params.num_frames ) if len(frames) != 3: logger.warning( f"{case.id}: expected 3 frames, got {len(frames)}, skipping frame save" ) return # Save frames (reuse naming from _consistency_gt_filenames) filenames = _consistency_gt_filenames(case.id, num_gpus, is_video=True) from PIL import Image for frame, fn in zip(frames, filenames): frame_path = out_dir / fn Image.fromarray(frame).save(frame_path) logger.info(f"Saved GT frame: {frame_path}") else: # Save image from sglang.multimodal_gen.test.test_utils import detect_image_format detected_format = detect_image_format(content) filenames = _consistency_gt_filenames( case.id, num_gpus, is_video=False, output_format=detected_format ) output_path = out_dir / filenames[0] output_path.write_bytes(content) logger.info(f"Saved GT image: {output_path} (format: {detected_format})") def _test_lora_api_functionality( self, ctx: ServerContext, case: DiffusionTestCase, generate_fn: Callable[[str, openai.Client], tuple[str, bytes]], ) -> None: """ Test LoRA API functionality with end-to-end validation: merge, unmerge, and set_lora. This test verifies that each API call succeeds AND that generation works after each operation. """ base_url = f"http://localhost:{ctx.port}/v1" client = OpenAI(base_url=base_url, api_key="dummy") # Test 1: unmerge_lora_weights - API should succeed and generation should work logger.info("[LoRA E2E] Testing unmerge_lora_weights for %s", case.id) resp = requests.post(f"{base_url}/unmerge_lora_weights") assert resp.status_code == 200, f"unmerge_lora_weights failed: {resp.text}" logger.info("[LoRA E2E] Verifying generation after unmerge for %s", case.id) rid_after_unmerge, _ = generate_fn(case.id, client) assert rid_after_unmerge is not None, "Generation after unmerge failed" logger.info("[LoRA E2E] Generation after unmerge succeeded") # Test 2: merge_lora_weights - API should succeed and generation should work logger.info("[LoRA E2E] Testing merge_lora_weights for %s", case.id) resp = requests.post(f"{base_url}/merge_lora_weights") assert resp.status_code == 200, f"merge_lora_weights failed: {resp.text}" logger.info("[LoRA E2E] Verifying generation after re-merge for %s", case.id) rid_after_merge, _ = generate_fn(case.id, client) assert rid_after_merge is not None, "Generation after merge failed" logger.info("[LoRA E2E] Generation after merge succeeded") # Test 3: set_lora (re-set the same adapter) - API should succeed and generation should work logger.info("[LoRA E2E] Testing set_lora for %s", case.id) resp = requests.post(f"{base_url}/set_lora", json={"lora_nickname": "default"}) assert resp.status_code == 200, f"set_lora failed: {resp.text}" logger.info("[LoRA E2E] Verifying generation after set_lora for %s", case.id) rid_after_set, _ = generate_fn(case.id, client) assert rid_after_set is not None, "Generation after set_lora failed" logger.info("[LoRA E2E] Generation after set_lora succeeded") # Test 4: list_loras - API should return the expected list of LoRA adapters logger.info("[LoRA E2E] Testing list_loras for %s", case.id) resp = requests.get(f"{base_url}/list_loras") assert resp.status_code == 200, f"list_loras failed: {resp.text}" lora_info = resp.json() logger.info("[LoRA E2E] list_loras returned %s", lora_info) assert ( isinstance(lora_info["loaded_adapters"], list) and len(lora_info["loaded_adapters"]) > 0 ), "loaded_adapters should be a non-empty list" assert any( a.get("nickname") == "default" for a in lora_info["loaded_adapters"] ), f"nickname 'default' not found in loaded_adapters: {lora_info['loaded_adapters']}" logger.info("[LoRA E2E] list_loras returned expected LoRA adapters") logger.info("[LoRA E2E] All LoRA API E2E tests passed for %s", case.id) def _test_lora_dynamic_switch_e2e( self, ctx: ServerContext, case: DiffusionTestCase, generate_fn: Callable[[str, openai.Client], tuple[str, bytes]], second_lora_path: str, ) -> None: """ Test dynamic LoRA switching with end-to-end validation. This test verifies that switching between LoRA adapters works correctly and generation succeeds after each switch. """ base_url = f"http://localhost:{ctx.port}/v1" client = OpenAI(base_url=base_url, api_key="dummy") # Test 1: Generate with initial LoRA logger.info( "[LoRA Switch E2E] Testing generation with initial LoRA for %s", case.id ) rid_initial, _ = generate_fn(case.id, client) assert rid_initial is not None, "Generation with initial LoRA failed" logger.info("[LoRA Switch E2E] Generation with initial LoRA succeeded") # Test 2: Switch to second LoRA and generate logger.info( "[LoRA Switch E2E] Switching to second LoRA adapter for %s", case.id ) resp = requests.post( f"{base_url}/set_lora", json={"lora_nickname": "lora2", "lora_path": second_lora_path}, ) assert ( resp.status_code == 200 ), f"set_lora to second adapter failed: {resp.text}" logger.info( "[LoRA Switch E2E] Verifying generation with second LoRA for %s", case.id ) rid_second, _ = generate_fn(case.id, client) assert rid_second is not None, "Generation with second LoRA failed" logger.info("[LoRA Switch E2E] Generation with second LoRA succeeded") # Test 3: Switch back to original LoRA and generate logger.info("[LoRA Switch E2E] Switching back to original LoRA for %s", case.id) resp = requests.post(f"{base_url}/set_lora", json={"lora_nickname": "default"}) assert resp.status_code == 200, f"set_lora back to default failed: {resp.text}" logger.info( "[LoRA Switch E2E] Verifying generation after switching back for %s", case.id, ) rid_switched_back, _ = generate_fn(case.id, client) assert rid_switched_back is not None, "Generation after switching back failed" logger.info("[LoRA Switch E2E] Generation after switching back succeeded") logger.info( "[LoRA Switch E2E] All dynamic switch E2E tests passed for %s", case.id ) def _test_dynamic_lora_loading( self, ctx: ServerContext, case: DiffusionTestCase, ) -> None: """ Test dynamic LoRA loading after server startup. This test reproduces the LayerwiseOffload + set_lora issue: - Server starts WITHOUT lora_path (LayerwiseOffloadManager initializes first) - Then set_lora is called via API to load LoRA dynamically - This tests the interaction between layerwise offload and dynamic LoRA loading """ base_url = f"http://localhost:{ctx.port}/v1" dynamic_lora_path = case.server_args.dynamic_lora_path # Call set_lora to load LoRA dynamically after server startup logger.info( "[Dynamic LoRA] Loading LoRA dynamically via set_lora API for %s", case.id ) logger.info("[Dynamic LoRA] LoRA path: %s", dynamic_lora_path) resp = requests.post( f"{base_url}/set_lora", json={"lora_nickname": "default", "lora_path": dynamic_lora_path}, ) assert resp.status_code == 200, f"Dynamic set_lora failed: {resp.text}" logger.info("[Dynamic LoRA] set_lora succeeded for %s", case.id) def _test_multi_lora_e2e( self, ctx: ServerContext, case: DiffusionTestCase, generate_fn: Callable[[str, openai.Client], tuple[str, bytes]], first_lora_path: str, second_lora_path: str, ) -> None: """ Test multiple LoRA adapters with different set_lora input scenarios. Tests: basic multi-LoRA, different strengths, cached adapters, switch back to single. """ base_url = f"http://localhost:{ctx.port}/v1" client = OpenAI(base_url=base_url, api_key="dummy") # Test 1: Basic multi-LoRA with list format resp = requests.post( f"{base_url}/set_lora", json={ "lora_nickname": ["default", "lora2"], "lora_path": [first_lora_path, second_lora_path], "target": "all", "strength": [1.0, 1.0], }, ) assert ( resp.status_code == 200 ), f"set_lora with multiple adapters failed: {resp.text}" rid, _ = generate_fn(case.id, client) assert rid is not None # Test 2: Different strengths resp = requests.post( f"{base_url}/set_lora", json={ "lora_nickname": ["default", "lora2"], "lora_path": [first_lora_path, second_lora_path], "target": "all", "strength": [0.8, 0.5], }, ) assert ( resp.status_code == 200 ), f"set_lora with different strengths failed: {resp.text}" rid, _ = generate_fn(case.id, client) assert rid is not None # Test 3: Different targets requests.post(f"{base_url}/set_lora", json={"lora_nickname": "default"}) resp = requests.post( f"{base_url}/set_lora", json={ "lora_nickname": ["default", "lora2"], "lora_path": [first_lora_path, second_lora_path], "target": ["transformer", "transformer_2"], "strength": [0.8, 0.5], }, ) assert ( resp.status_code == 200 ), f"set_lora with cached adapters failed: {resp.text}" rid, _ = generate_fn(case.id, client) assert rid is not None # Test 4: Switch back to single LoRA resp = requests.post(f"{base_url}/set_lora", json={"lora_nickname": "default"}) assert ( resp.status_code == 200 ), f"set_lora back to single adapter failed: {resp.text}" rid, _ = generate_fn(case.id, client) assert rid is not None logger.info("[Multi-LoRA] All multi-LoRA tests passed for %s", case.id) def _test_v1_models_endpoint( self, ctx: ServerContext, case: DiffusionTestCase ) -> None: """ Test /v1/models endpoint returns OpenAI-compatible response. This endpoint is required for sgl-model-gateway router compatibility. """ base_url = f"http://localhost:{ctx.port}" # Test GET /v1/models logger.info("[Models API] Testing GET /v1/models for %s", case.id) resp = requests.get(f"{base_url}/v1/models") assert resp.status_code == 200, f"/v1/models failed: {resp.text}" data = resp.json() assert ( data["object"] == "list" ), f"Expected object='list', got {data.get('object')}" assert len(data["data"]) >= 1, "Expected at least one model in response" model = data["data"][0] assert "id" in model, "Model missing 'id' field" assert ( model["object"] == "model" ), f"Expected object='model', got {model.get('object')}" assert ( model["id"] == case.server_args.model_path ), f"Model ID mismatch: expected {case.server_args.model_path}, got {model['id']}" # Verify extended diffusion-specific fields assert "num_gpus" in model, "Model missing 'num_gpus' field" assert "task_type" in model, "Model missing 'task_type' field" assert "dit_precision" in model, "Model missing 'dit_precision' field" assert "vae_precision" in model, "Model missing 'vae_precision' field" assert ( model["num_gpus"] == case.server_args.num_gpus ), f"num_gpus mismatch: expected {case.server_args.num_gpus}, got {model['num_gpus']}" # Verify task_type is consistent with the modality specified in the test config. # We can't access pipeline_config from test config, but we can validate against modality. modality_to_valid_task_types = { "image": {"T2I", "I2I", "TI2I"}, "video": {"T2V", "I2V", "TI2V"}, "3d": {"I2M"}, } valid_task_types = modality_to_valid_task_types.get( case.server_args.modality, set() ) assert model["task_type"] in valid_task_types, ( f"task_type '{model['task_type']}' not valid for modality " f"'{case.server_args.modality}'. Expected one of: {valid_task_types}" ) logger.info( "[Models API] GET /v1/models returned valid response with extended fields" ) # Test GET /v1/models/{model_path} model_path = model["id"] logger.info("[Models API] Testing GET /v1/models/%s", model_path) resp = requests.get(f"{base_url}/v1/models/{model_path}") assert resp.status_code == 200, f"/v1/models/{model_path} failed: {resp.text}" single_model = resp.json() assert single_model["id"] == model_path, "Single model ID mismatch" assert single_model["object"] == "model", "Single model object type mismatch" # Verify extended fields on single model endpoint too assert "num_gpus" in single_model, "Single model missing 'num_gpus' field" assert "task_type" in single_model, "Single model missing 'task_type' field" assert single_model["task_type"] in valid_task_types, ( f"Single model task_type '{single_model['task_type']}' not valid for modality " f"'{case.server_args.modality}'. Expected one of: {valid_task_types}" ) logger.info( "[Models API] GET /v1/models/{model_path} returned valid response with extended fields" ) # Test GET /v1/models/{non_existent_model} returns 404 logger.info("[Models API] Testing GET /v1/models/non_existent_model") resp = requests.get(f"{base_url}/v1/models/non_existent_model") assert resp.status_code == 404, f"Expected 404, got {resp.status_code}" error_data = resp.json() assert "error" in error_data, "404 response missing 'error' field" assert ( error_data["error"]["code"] == "model_not_found" ), f"Incorrect error code: {error_data['error'].get('code')}" logger.info("[Models API] GET /v1/models/non_existent returns 404 as expected") logger.info("[Models API] All /v1/models tests passed for %s", case.id) def _test_t2v_rejects_input_reference( self, ctx: ServerContext, case: DiffusionTestCase ) -> None: if case.server_args.modality != "video": return base_url = f"http://localhost:{ctx.port}" resp = requests.get(f"{base_url}/v1/models") assert resp.status_code == 200, f"/v1/models failed: {resp.text}" data = resp.json().get("data", []) if not data: pytest.fail("/v1/models returned empty model list") task_type = data[0].get("task_type") if task_type != "T2V": return prompt = case.sampling_params.prompt or "test" payload = {"prompt": prompt, "input_reference": "dummy"} if case.sampling_params.output_size: payload["size"] = case.sampling_params.output_size resp = requests.post(f"{base_url}/v1/videos", json=payload) assert ( resp.status_code == 400 ), f"Expected 400 for T2V input_reference, got {resp.status_code}: {resp.text}" detail = resp.json().get("detail", "") assert ( "input_reference is not supported" in detail ), f"Unexpected error detail for T2V input_reference: {detail}" def test_diffusion_generation( self, case: DiffusionTestCase, diffusion_server: ServerContext, ): """Single parametrized test that runs for all cases. This test performs: 1. Generation 2. Performance validation against baselines 3. Consistency validation against ground truth Pytest will execute this test once per case in ONE_GPU_CASES, with test IDs like: - test_diffusion_generation[qwen_image_text] - test_diffusion_generation[qwen_image_edit] - etc. """ # Check if we're in GT generation mode is_gt_gen_mode = os.environ.get("SGLANG_GEN_GT", "0") == "1" # Dynamic LoRA loading test - tests LayerwiseOffload + set_lora interaction # Server starts WITHOUT lora_path, then set_lora is called after startup if case.server_args.dynamic_lora_path and not is_gt_gen_mode: self._test_dynamic_lora_loading(diffusion_server, case) generate_fn = get_generate_fn( model_path=case.server_args.model_path, modality=case.server_args.modality, sampling_params=case.sampling_params, ) # Single generation - output is reused for both validations perf_record, content = self.run_and_collect( diffusion_server, case.id, generate_fn, ) if is_gt_gen_mode: # GT generation mode: save output and skip all validations/tests self._save_gt_output(case, content) return # Validation 1: Performance self._validate_and_record(case, perf_record) # Mesh correctness check (Chamfer Distance) for 3D models if case.server_args.custom_validator == "mesh": from sglang.multimodal_gen.test.server.test_server_utils import ( MESH_OUTPUT_PATHS, validate_mesh_correctness, ) mesh_path = MESH_OUTPUT_PATHS.pop(case.id, None) if mesh_path: validate_mesh_correctness(mesh_path) # Test /v1/models endpoint for router compatibility self._test_v1_models_endpoint(diffusion_server, case) self._test_t2v_rejects_input_reference(diffusion_server, case) # LoRA API functionality test with E2E validation (only for LoRA-enabled cases) if case.server_args.lora_path or case.server_args.dynamic_lora_path: self._test_lora_api_functionality(diffusion_server, case, generate_fn) # Test dynamic LoRA switching (requires a second LoRA adapter) if case.server_args.second_lora_path: self._test_lora_dynamic_switch_e2e( diffusion_server, case, generate_fn, case.server_args.second_lora_path, ) # Test multi-LoRA functionality self._test_multi_lora_e2e( diffusion_server, case, generate_fn, case.server_args.lora_path, case.server_args.second_lora_path, ) ================================================ FILE: python/sglang/multimodal_gen/test/server/test_server_utils.py ================================================ """ Server management and performance validation for diffusion tests. """ from __future__ import annotations import base64 import os import shlex import subprocess import sys import tempfile import threading import time from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable, Sequence from urllib.request import urlopen import pytest from openai import Client from sglang.multimodal_gen.benchmarks.compare_perf import calculate_upper_bound from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.common import kill_process_tree from sglang.multimodal_gen.runtime.utils.logging_utils import ( globally_suppress_loggers, init_logger, ) from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord from sglang.multimodal_gen.test.server.testcase_configs import ( DiffusionSamplingParams, PerformanceSummary, ScenarioConfig, ToleranceConfig, ) from sglang.multimodal_gen.test.slack_utils import upload_file_to_slack from sglang.multimodal_gen.test.test_utils import ( get_expected_image_format, get_video_frame_count, is_image_url, prepare_perf_log, validate_image, validate_image_file, validate_openai_video, validate_video_file, ) logger = init_logger(__name__) globally_suppress_loggers() # Tracks mesh output file paths from generate_mesh for later correctness validation. # Keyed by case_id, cleaned up after use. MESH_OUTPUT_PATHS: dict[str, str] = {} def download_image_from_url(url: str) -> Path: """Download an image from a URL to a temporary file. Args: url: The URL of the image to download Returns: Path to the downloaded temporary file """ logger.info(f"Downloading image from URL: {url}") # Determine file extension from URL ext = ".jpg" # default if url.lower().endswith((".png", ".jpeg", ".jpg", ".webp", ".gif")): ext = url[url.rfind(".") :] # Create temporary file temp_file = ( Path(tempfile.gettempdir()) / f"diffusion_test_image_{int(time.time())}{ext}" ) try: with urlopen(url, timeout=30) as response: temp_file.write_bytes(response.read()) logger.info(f"Downloaded image to: {temp_file}") return temp_file except Exception as e: logger.error(f"Failed to download image from {url}: {e}") raise def parse_dimensions(size_string: str | None) -> tuple[int | None, int | None]: """Parse a size string in "widthxheight" format to (width, height) tuple. Args: size_string: Size string in "widthxheight" format (e.g., "1024x1024") or None. Spaces are automatically stripped. Returns: Tuple of (width, height) as integers if parsing succeeds, (None, None) otherwise. """ if not size_string: return (None, None) # Strip spaces from the entire string size_string = size_string.strip() if not size_string: return (None, None) # Split by "x" parts = size_string.split("x") if len(parts) != 2: return (None, None) # Strip spaces from each part and try to convert to int try: width_str = parts[0].strip() height_str = parts[1].strip() if not width_str or not height_str: return (None, None) width = int(width_str) height = int(height_str) # Validate that both are positive if width <= 0 or height <= 0: return (None, None) return (width, height) except ValueError: return (None, None) @dataclass class ServerContext: """Context for a running diffusion server.""" port: int process: subprocess.Popen model: str stdout_file: Path perf_log_path: Path log_dir: Path _stdout_fh: Any = field(repr=False) _log_thread: threading.Thread | None = field(default=None, repr=False) def cleanup(self) -> None: """Clean up server resources.""" try: kill_process_tree(self.process.pid) except Exception: pass try: self._stdout_fh.flush() self._stdout_fh.close() except Exception: pass # ROCm/AMD: Extra cleanup to ensure GPU memory is released between tests # This is needed because ROCm memory release can be slower than CUDA if current_platform.is_hip(): self._cleanup_rocm_gpu_memory() # Clean up downloaded models if HF cache is not persistent # This prevents disk exhaustion in CI when cache is not mounted self._cleanup_hf_cache_if_not_persistent() def _cleanup_hf_cache_if_not_persistent(self) -> None: """Clean up HF cache if it's not on a persistent volume. When running in CI without persistent cache, downloaded models accumulate and can cause disk/memory exhaustion. This cleans up the model after each test if the cache is not persistent. """ import shutil hf_home = os.environ.get("HF_HOME", "") if not hf_home: return hf_hub_cache = os.path.join(hf_home, "hub") # Check if HF cache is on a persistent volume by looking for a marker file # or checking if the directory existed before this test run persistent_marker = os.path.join(hf_home, ".persistent_cache") if os.path.exists(persistent_marker): logger.info("HF cache is persistent, skipping cleanup") return # Check if the cache directory is empty or was just created # If it has very few models, it's likely not persistent if not os.path.exists(hf_hub_cache): return try: # Get model cache directories model_dirs = [ d for d in os.listdir(hf_hub_cache) if d.startswith("models--") and os.path.isdir(os.path.join(hf_hub_cache, d)) ] # If there are cached models but no persistent marker, clean up # to prevent disk exhaustion in CI if model_dirs: logger.info( "HF cache appears non-persistent (no .persistent_cache marker), " "cleaning up %d model(s) to prevent disk exhaustion", len(model_dirs), ) for model_dir in model_dirs: model_path = os.path.join(hf_hub_cache, model_dir) try: shutil.rmtree(model_path) logger.info("Cleaned up model cache: %s", model_dir) except Exception as e: logger.warning("Failed to clean up %s: %s", model_dir, e) except Exception as e: logger.warning("Error during HF cache cleanup: %s", e) def _cleanup_rocm_gpu_memory(self) -> None: """ROCm-specific cleanup to ensure GPU memory is fully released.""" import gc # Wait for process to fully terminate try: self.process.wait(timeout=30) except Exception: pass # Force garbage collection multiple times for _ in range(3): gc.collect() # Clear HIP memory on all GPUs try: import torch for i in range(torch.cuda.device_count()): with torch.cuda.device(i): torch.cuda.empty_cache() torch.cuda.synchronize() except Exception: pass # Wait for GPU memory to be released (ROCm can be much slower than CUDA) # The GPU driver needs time to reclaim memory from killed processes time.sleep(15) class ServerManager: """Manages diffusion server lifecycle.""" def __init__( self, model: str, port: int, wait_deadline: float = 1200.0, extra_args: str = "", env_vars: dict[str, str] | None = None, ): self.model = model self.port = port self.wait_deadline = wait_deadline self.extra_args = extra_args self.env_vars = env_vars or {} def _wait_for_rocm_gpu_memory_clear(self, max_wait: float = 60.0) -> None: """ROCm-specific: Wait for GPU memory to be mostly free before starting. ROCm GPU memory release from killed processes can be significantly slower than CUDA, so we need to wait longer and be more patient. """ try: import torch if not torch.cuda.is_available(): return start_time = time.time() last_total_used = float("inf") while time.time() - start_time < max_wait: # Check GPU memory usage total_used = 0 for i in range(torch.cuda.device_count()): mem_info = torch.cuda.mem_get_info(i) free, total = mem_info used = total - free total_used += used # If less than 5GB is used across all GPUs, we're good if total_used < 5 * 1024 * 1024 * 1024: # 5GB logger.info( "[server-test] ROCm GPU memory is clear (used: %.2f GB)", total_used / (1024**3), ) return # Log progress elapsed = int(time.time() - start_time) if total_used < last_total_used: logger.info( "[server-test] ROCm: GPU memory clearing (used: %.2f GB, elapsed: %ds)", total_used / (1024**3), elapsed, ) else: logger.info( "[server-test] ROCm: Waiting for GPU memory (used: %.2f GB, elapsed: %ds)", total_used / (1024**3), elapsed, ) last_total_used = total_used time.sleep(3) # Final warning with detailed GPU info logger.warning( "[server-test] ROCm GPU memory not fully cleared after %.0fs (used: %.2f GB). " "Proceeding anyway - this may cause OOM.", max_wait, total_used / (1024**3), ) except Exception as e: logger.debug("[server-test] Could not check ROCm GPU memory: %s", e) def start(self) -> ServerContext: """Start the diffusion server and wait for readiness.""" # ROCm/AMD: Wait for GPU memory to be clear before starting # This prevents OOM when running sequential tests on ROCm if current_platform.is_hip(): self._wait_for_rocm_gpu_memory_clear() log_dir, perf_log_path = prepare_perf_log() safe_model_name = self.model.replace("/", "_") stdout_path = ( Path(tempfile.gettempdir()) / f"sgl_server_{self.port}_{safe_model_name}.log" ) stdout_path.unlink(missing_ok=True) command = [ "sglang", "serve", "--model-path", self.model, "--port", str(self.port), "--log-level=debug", ] if self.extra_args.strip(): command.extend(self.extra_args.strip().split()) env = os.environ.copy() env["SGLANG_DIFFUSION_STAGE_LOGGING"] = "1" env["SGLANG_PERF_LOG_DIR"] = log_dir.as_posix() # Apply custom environment variables env.update(self.env_vars) # TODO: unify with run_command logger.info(f"Running command: {shlex.join(command)}") process = subprocess.Popen( command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, env=env, ) log_thread = None stdout_fh = stdout_path.open("w", encoding="utf-8", buffering=1) if process.stdout: def _log_pipe(pipe: Any, file: Any) -> None: """Read from pipe and write to file and stdout.""" try: with pipe: for line in iter(pipe.readline, ""): sys.stdout.write(line) sys.stdout.flush() file.write(line) file.flush() except Exception as e: logger.error("Log pipe thread error: %s", e) finally: file.close() logger.debug("Log pipe thread finished.") log_thread = threading.Thread( target=_log_pipe, args=(process.stdout, stdout_fh) ) log_thread.daemon = True log_thread.start() logger.info( "[server-test] Starting server pid=%s, model=%s, log=%s", process.pid, self.model, stdout_path, ) self._wait_for_ready(process, stdout_path) return ServerContext( port=self.port, process=process, model=self.model, stdout_file=stdout_path, perf_log_path=perf_log_path, log_dir=log_dir, _stdout_fh=stdout_fh, _log_thread=log_thread, ) def _wait_for_ready(self, process: subprocess.Popen, stdout_path: Path) -> None: """Wait for server to become ready.""" start = time.time() ready_message = "Application startup complete." log_period = 30 prev_log_period_count = 0 while time.time() - start < self.wait_deadline: if process.poll() is not None: tail = self._get_log_tail(stdout_path) raise RuntimeError( f"Server exited early (code {process.returncode}).\n{tail}" ) if stdout_path.exists(): try: content = stdout_path.read_text(encoding="utf-8", errors="ignore") if ready_message in content: logger.info("[server-test] Server ready") return except Exception as e: logger.debug("Could not read log yet: %s", e) elapsed = int(time.time() - start) if (elapsed // log_period) > prev_log_period_count: prev_log_period_count = elapsed // log_period logger.info("[server-test] Waiting for server... elapsed=%ss", elapsed) time.sleep(1) tail = self._get_log_tail(stdout_path) raise TimeoutError(f"Server not ready within {self.wait_deadline}s.\n{tail}") @staticmethod def _get_log_tail(path: Path, lines: int = 200) -> str: """Get the last N lines from a log file.""" try: content = path.read_text(encoding="utf-8", errors="ignore") return "\n".join(content.splitlines()[-lines:]) except Exception: return "" class PerformanceValidator: """Validates performance metrics against expectations.""" is_video_gen: bool = False def __init__( self, scenario: ScenarioConfig, tolerances: ToleranceConfig, step_fractions: Sequence[float], ): self.scenario = scenario self.tolerances = tolerances self.step_fractions = step_fractions self.is_baseline_generation_mode = ( os.environ.get("SGLANG_GEN_BASELINE", "0") == "1" ) def _assert_le( self, name: str, actual: float, expected: float, tolerance: float, min_abs_tolerance_ms: float = 20.0, ): """Assert that actual is less than or equal to expected within a tolerance. Uses the larger of relative tolerance or absolute tolerance to prevent flaky failures on very fast operations. For AMD GPUs, uses 100% higher tolerance and issues warning instead of assertion. """ # Check if running on AMD GPU is_amd = current_platform.is_hip() if is_amd: # Use 100% higher tolerance for AMD (2x the expected value) amd_tolerance = 1.0 # 100% upper_bound = calculate_upper_bound( expected, amd_tolerance, min_abs_tolerance_ms ) if actual > upper_bound: logger.warning( f"[AMD PERF WARNING] Validation would fail for '{name}'.\n" f" Actual: {actual:.4f}ms\n" f" Expected: {expected:.4f}ms\n" f" AMD Limit: {upper_bound:.4f}ms " f"(rel_tol: {amd_tolerance:.1%}, abs_pad: {min_abs_tolerance_ms}ms)\n" f" Original tolerance was: {tolerance:.1%}" ) else: upper_bound = calculate_upper_bound( expected, tolerance, min_abs_tolerance_ms ) assert actual <= upper_bound, ( f"Validation failed for '{name}'.\n" f" Actual: {actual:.4f}ms\n" f" Expected: {expected:.4f}ms\n" f" Limit: {upper_bound:.4f}ms " f"(rel_tol: {tolerance:.1%}, abs_pad: {min_abs_tolerance_ms}ms)" ) def validate( self, perf_record: RequestPerfRecord, *args, **kwargs ) -> PerformanceSummary: """Validate all performance metrics and return summary.""" summary = self.collect_metrics(perf_record) if self.is_baseline_generation_mode: return summary self._validate_e2e(summary) self._validate_denoise_agg(summary) self._validate_denoise_steps(summary) self._validate_stages(summary) return summary def collect_metrics( self, perf_record: RequestPerfRecord, ) -> PerformanceSummary: return PerformanceSummary.from_req_perf_record(perf_record, self.step_fractions) def _validate_e2e(self, summary: PerformanceSummary) -> None: """Validate end-to-end performance.""" assert summary.e2e_ms > 0, "E2E duration missing" self._assert_le( "E2E Latency", summary.e2e_ms, self.scenario.expected_e2e_ms, self.tolerances.e2e, ) def _validate_denoise_agg(self, summary: PerformanceSummary) -> None: """Validate aggregate denoising metrics.""" assert summary.avg_denoise_ms > 0, "Denoising step timings missing" self._assert_le( "Average Denoise Step", summary.avg_denoise_ms, self.scenario.expected_avg_denoise_ms, self.tolerances.denoise_agg, ) self._assert_le( "Median Denoise Step", summary.median_denoise_ms, self.scenario.expected_median_denoise_ms, self.tolerances.denoise_agg, ) def _validate_denoise_steps(self, summary: PerformanceSummary) -> None: """Validate individual denoising steps.""" for idx, actual in summary.sampled_steps.items(): expected = self.scenario.denoise_step_ms.get(idx) if expected is None: continue # FIXME: hardcode, looser for first step tolerance = 0.4 if idx == 0 else self.tolerances.denoise_step self._assert_le( f"Denoise Step {idx}", actual, expected, tolerance, ) def _validate_stages(self, summary: PerformanceSummary) -> None: """Validate stage-level metrics.""" assert summary.stage_metrics, "Stage metrics missing" for stage, expected in self.scenario.stages_ms.items(): if stage == "per_frame_generation" and self.is_video_gen: continue actual = summary.stage_metrics.get(stage) assert actual is not None, f"Stage {stage} timing missing" tolerance = ( self.tolerances.denoise_stage if stage == "DenoisingStage" else self.tolerances.non_denoise_stage ) self._assert_le( f"Stage '{stage}'", actual, expected, tolerance, min_abs_tolerance_ms=120.0, # relax absolute tolerance for non-denoising stages ) class VideoPerformanceValidator(PerformanceValidator): """Extended validator for video diffusion with frame-level metrics.""" is_video_gen = True def validate( self, perf_record: RequestPerfRecord, num_frames: int | None = None, ) -> PerformanceSummary: """Validate video metrics including frame generation rates.""" summary = super().validate(perf_record) if num_frames and summary.e2e_ms > 0: summary.total_frames = num_frames summary.avg_frame_time_ms = summary.e2e_ms / num_frames summary.frames_per_second = 1000.0 / summary.avg_frame_time_ms if not self.is_baseline_generation_mode: self._validate_frame_rate(summary) return summary def _validate_frame_rate(self, summary: PerformanceSummary) -> None: """Validate frame generation performance.""" expected_frame_time = self.scenario.stages_ms.get("per_frame_generation") if expected_frame_time and summary.avg_frame_time_ms: self._assert_le( "Average Frame Time", summary.avg_frame_time_ms, expected_frame_time, self.tolerances.denoise_stage, ) class MeshValidator(PerformanceValidator): """Validator for 3D mesh generation. Inherits perf validation from PerformanceValidator.""" pass HUNYUAN3D_REFERENCE_URL = ( "https://raw.githubusercontent.com/sgl-project/sgl-test-files/" "main/diffusion-ci/consistency_gt/1-gpu/hunyuan3d_2_0/hunyuan3d.glb" ) def _download_reference_mesh(url: str) -> Path: """Download a reference mesh from URL, caching in temp dir.""" import hashlib cache_name = f"ref_mesh_{hashlib.md5(url.encode()).hexdigest()}.glb" cache_path = Path(tempfile.gettempdir()) / cache_name if cache_path.exists(): logger.info(f"Using cached reference mesh: {cache_path}") return cache_path logger.info(f"Downloading reference mesh from: {url}") with urlopen(url, timeout=60) as resp: cache_path.write_bytes(resp.read()) logger.info(f"Reference mesh cached at: {cache_path}") return cache_path def validate_mesh_correctness( generated_mesh_path: str, reference_url: str = HUNYUAN3D_REFERENCE_URL, num_sample_points: int = 4096, cd_threshold_ratio: float = 0.01, random_seed: int = 42, ): """Validate mesh geometric similarity against a reference via Chamfer Distance. Downloads the reference mesh from a URL (cached), samples point clouds from both meshes, and asserts Chamfer Distance is within threshold. """ import numpy as np try: import trimesh except ImportError: pytest.fail("trimesh is required for mesh validation: pip install trimesh") from scipy.spatial import cKDTree # Load generated mesh generated_mesh = trimesh.load(generated_mesh_path) if isinstance(generated_mesh, trimesh.Scene): generated_mesh = generated_mesh.dump(concatenate=True) # Download and load reference mesh ref_path = _download_reference_mesh(reference_url) reference_mesh = trimesh.load(str(ref_path)) if isinstance(reference_mesh, trimesh.Scene): reference_mesh = reference_mesh.dump(concatenate=True) # Bounding box diagonal for threshold normalization ref_bbox = reference_mesh.bounding_box.bounds bbox_diagonal = float(np.linalg.norm(ref_bbox[1] - ref_bbox[0])) cd_threshold = cd_threshold_ratio * bbox_diagonal # Sample point clouds np.random.seed(random_seed) gen_points = np.array( generated_mesh.sample(num_sample_points, return_index=True)[0] ) ref_points = np.array( reference_mesh.sample(num_sample_points, return_index=True)[0] ) # Bidirectional Chamfer Distance tree1 = cKDTree(gen_points) tree2 = cKDTree(ref_points) forward_cd = float(np.mean(tree2.query(gen_points)[0] ** 2)) backward_cd = float(np.mean(tree1.query(ref_points)[0] ** 2)) total_cd = forward_cd + backward_cd assert total_cd <= cd_threshold, ( f"Chamfer Distance check failed: total_cd={total_cd:.6f}, " f"threshold={cd_threshold:.6f} ({cd_threshold_ratio * 100:.2f}% of bbox diagonal {bbox_diagonal:.4f})" ) # Registry of validators by name VALIDATOR_REGISTRY = { "default": PerformanceValidator, "video": VideoPerformanceValidator, "mesh": MeshValidator, } def get_generate_fn( model_path: str, modality: str, sampling_params: DiffusionSamplingParams, ) -> Callable[[str, Client], tuple[str, bytes]]: """Return appropriate generation function for the case.""" # Allow override via environment variable (useful for AMD where large resolutions cause slow VAE) output_size = os.environ.get("SGLANG_TEST_OUTPUT_SIZE", sampling_params.output_size) n = sampling_params.num_outputs_per_prompt def _create_and_download_video( client, case_id, *, model: str, size: str, prompt: str | None = None, seconds: int | None = None, input_reference: Any | None = None, extra_body: dict[Any] | None = None, expected_frame_count: int | None = None, ) -> str: """ Create a video job via /v1/videos, poll until completion, then download the binary content and validate it. Returns request-id """ create_kwargs: dict[str, Any] = { "model": model, "size": size, } if prompt is not None: create_kwargs["prompt"] = prompt if seconds is not None: create_kwargs["seconds"] = seconds if input_reference is not None: create_kwargs["input_reference"] = input_reference # triggers multipart if extra_body is not None: create_kwargs["extra_body"] = extra_body job = client.videos.create(**create_kwargs) # type: ignore[attr-defined] video_id = job.id job_completed = False is_baseline_generation_mode = os.environ.get("SGLANG_GEN_BASELINE", "0") == "1" # Check if running on AMD GPU - use longer timeout is_amd = current_platform.is_hip() if is_baseline_generation_mode: timeout = 3600.0 elif is_amd: timeout = 2400.0 # 40 minutes for AMD else: timeout = 1200.0 deadline = time.time() + timeout while True: page = client.videos.list() # type: ignore[attr-defined] item = next((v for v in page.data if v.id == video_id), None) if item and getattr(item, "status", None) == "completed": job_completed = True break if time.time() > deadline: break time.sleep(1) if not job_completed: if is_baseline_generation_mode: logger.warning( f"{case_id}: video job {video_id} timed out during baseline generation. " "Attempting to collect performance data anyway." ) return (video_id, b"") if is_amd: logger.warning( f"[AMD TIMEOUT WARNING] {case_id}: video job {video_id} did not complete " f"within {timeout}s timeout. This may indicate performance issues on AMD." ) pytest.skip( f"{case_id}: video job timed out on AMD after {timeout}s - skipping" ) pytest.fail(f"{case_id}: video job {video_id} did not complete in time") # download video resp = client.videos.download_content(video_id=video_id) # type: ignore[attr-defined] content = resp.read() validate_openai_video(content) expected_filename = f"{video_id}.mp4" tmp_path = expected_filename with open(tmp_path, "wb") as f: f.write(content) # Validate output file expected_width, expected_height = parse_dimensions(size) validate_video_file( tmp_path, expected_filename, expected_width, expected_height ) if expected_frame_count is not None: actual_count = get_video_frame_count(tmp_path) assert actual_count == expected_frame_count, ( f"{case_id}: frame count mismatch after interpolation — " f"expected {expected_frame_count}, got {actual_count}" ) upload_file_to_slack( case_id=case_id, model=model_path, prompt=sampling_params.prompt, file_path=tmp_path, origin_file_path=sampling_params.image_path, ) os.remove(tmp_path) return (video_id, content) video_seconds = sampling_params.seconds or 4 def generate_image(case_id, client) -> tuple[str, bytes]: """T2I: Text to Image generation.""" if not sampling_params.prompt: pytest.skip(f"{case_id}: no text prompt configured") # Request parameters that affect output format req_output_format = None # Not specified in current request req_background = None # Not specified in current request # Build extra_body for optional features extra_body = dict(sampling_params.extras) response = client.images.with_raw_response.generate( model=model_path, prompt=sampling_params.prompt, n=n, size=output_size, response_format="b64_json", extra_body=extra_body if extra_body else None, ) result = response.parse() validate_image(result.data[0].b64_json) rid = result.id img_data = base64.b64decode(result.data[0].b64_json) # Infer expected format from request parameters expected_ext = get_expected_image_format(req_output_format, req_background) expected_filename = f"{result.created}.{expected_ext}" tmp_path = expected_filename with open(tmp_path, "wb") as f: f.write(img_data) # Validate output file expected_width, expected_height = parse_dimensions(output_size) if ( sampling_params.extras.get("enable_upscaling") and expected_width and expected_height ): expected_width *= sampling_params.extras.get("upscaling_scale", 4) expected_height *= sampling_params.extras.get("upscaling_scale", 4) validate_image_file( tmp_path, expected_filename, expected_width, expected_height, output_format=req_output_format, background=req_background, ) upload_file_to_slack( case_id=case_id, model=model_path, prompt=sampling_params.prompt, file_path=tmp_path, ) os.remove(tmp_path) return (rid, img_data) def generate_image_edit(case_id, client) -> tuple[str, bytes]: """TI2I: Text + Image -> Image edit.""" if not sampling_params.prompt or not sampling_params.image_path: pytest.skip(f"{case_id}: no edit config") image_paths = sampling_params.image_path if not isinstance(image_paths, list): image_paths = [image_paths] new_image_paths = [] for image_path in image_paths: if is_image_url(image_path): new_image_paths.append(download_image_from_url(str(image_path))) else: local_path = Path(image_path) new_image_paths.append(local_path) if not local_path.exists(): pytest.skip(f"{case_id}: file missing: {image_path}") image_paths = new_image_paths # Request parameters that affect output format req_output_format = ( sampling_params.output_format ) # Not specified in current request req_background = None # Not specified in current request # Build extra_body for optional features extra_body = {"num_frames": sampling_params.num_frames} extra_body.update(sampling_params.extras) images = [open(image_path, "rb") for image_path in image_paths] try: response = client.images.with_raw_response.edit( model=model_path, image=images, prompt=sampling_params.prompt, n=n, size=output_size, response_format="b64_json", output_format=req_output_format, extra_body=extra_body, ) finally: for img in images: img.close() result = response.parse() validate_image(result.data[0].b64_json) img_data = base64.b64decode(result.data[0].b64_json) rid = result.id # Infer expected format from request parameters expected_ext = get_expected_image_format(req_output_format, req_background) expected_filename = f"{rid}.{expected_ext}" tmp_path = expected_filename with open(tmp_path, "wb") as f: f.write(img_data) # Validate output file expected_width, expected_height = parse_dimensions(output_size) validate_image_file( tmp_path, expected_filename, expected_width, expected_height, output_format=req_output_format, background=req_background, ) upload_file_to_slack( case_id=case_id, model=model_path, prompt=sampling_params.prompt, file_path=tmp_path, origin_file_path=sampling_params.image_path, ) os.remove(tmp_path) return (rid, img_data) def generate_image_edit_url(case_id, client) -> tuple[str, bytes]: """TI2I: Text + Image ? Image edit using direct URL transfer (no pre-download).""" if not sampling_params.prompt or not sampling_params.image_path: pytest.skip(f"{case_id}: no edit config") # Handle both single URL and list of URLs image_urls = sampling_params.image_path if not isinstance(image_urls, list): image_urls = [image_urls] # Validate all URLs for url in image_urls: if not is_image_url(url): pytest.skip( f"{case_id}: image_path must be a URL for URL direct test: {url}" ) # Request parameters that affect output format req_output_format = ( sampling_params.output_format ) # Not specified in current request req_background = None # Not specified in current request response = client.images.with_raw_response.edit( model=model_path, prompt=sampling_params.prompt, image=[], # Only for OpenAI verification n=n, size=sampling_params.output_size, response_format="b64_json", output_format=req_output_format, extra_body={"url": image_urls, "num_frames": sampling_params.num_frames}, ) result = response.parse() rid = result.id validate_image(result.data[0].b64_json) # Save and upload result for verification img_data = base64.b64decode(result.data[0].b64_json) # Infer expected format from request parameters expected_ext = get_expected_image_format(req_output_format, req_background) expected_filename = f"{rid}.{expected_ext}" tmp_path = expected_filename with open(tmp_path, "wb") as f: f.write(img_data) # Validate output file expected_width, expected_height = parse_dimensions(sampling_params.output_size) validate_image_file( tmp_path, expected_filename, expected_width, expected_height, output_format=req_output_format, background=req_background, ) upload_file_to_slack( case_id=case_id, model=model_path, prompt=sampling_params.prompt, file_path=tmp_path, origin_file_path=str(sampling_params.image_path), ) os.remove(tmp_path) return (rid, img_data) def generate_video(case_id, client) -> tuple[str, bytes]: """T2V: Text ? Video.""" if not sampling_params.prompt: pytest.skip(f"{case_id}: no text prompt configured") # Build extra_body for optional features extra_body = dict(sampling_params.extras) if sampling_params.num_frames: extra_body["num_frames"] = sampling_params.num_frames # Compute expected output frame count for validation expected_frame_count = None if ( sampling_params.extras.get("enable_frame_interpolation") and sampling_params.num_frames ): n = sampling_params.num_frames exp = sampling_params.extras.get("frame_interpolation_exp", 1) expected_frame_count = (n - 1) * (2**exp) + 1 return _create_and_download_video( client, case_id, model=model_path, prompt=sampling_params.prompt, size=output_size, seconds=video_seconds, extra_body=extra_body if extra_body else None, expected_frame_count=expected_frame_count, ) def generate_image_to_video(case_id, client) -> tuple[str, bytes]: """I2V: Image -> Video (optional prompt).""" if not sampling_params.image_path: pytest.skip(f"{case_id}: no input image configured") if is_image_url(sampling_params.image_path): image_path = download_image_from_url(str(sampling_params.image_path)) else: image_path = Path(sampling_params.image_path) if not image_path.exists(): pytest.skip(f"{case_id}: file missing: {image_path}") # Build extra_body for optional features extra_body = dict(sampling_params.extras) with image_path.open("rb") as fh: return _create_and_download_video( client, case_id, model=model_path, prompt=sampling_params.prompt, size=output_size, seconds=video_seconds, input_reference=fh, extra_body=extra_body if extra_body else None, ) def generate_text_url_image_to_video(case_id, client) -> tuple[str, bytes]: if not sampling_params.prompt or not sampling_params.image_path: pytest.skip(f"{case_id}: no edit config") # Build extra_body for optional features extra_body = {"reference_url": sampling_params.image_path} extra_body.update(sampling_params.extras) return _create_and_download_video( client, case_id, model=model_path, prompt=sampling_params.prompt, size=sampling_params.output_size, seconds=video_seconds, extra_body={ "reference_url": sampling_params.image_path, "fps": sampling_params.fps, "num_frames": sampling_params.num_frames, }, ) def generate_text_image_to_video(case_id, client) -> tuple[str, bytes]: """TI2V: Text + Image -> Video.""" if not sampling_params.prompt or not sampling_params.image_path: pytest.skip(f"{case_id}: no edit config") if is_image_url(sampling_params.image_path): image_path = download_image_from_url(str(sampling_params.image_path)) else: image_path = Path(sampling_params.image_path) if not image_path.exists(): pytest.skip(f"{case_id}: file missing: {image_path}") # Build extra_body for optional features extra_body = dict(sampling_params.extras) with image_path.open("rb") as fh: return _create_and_download_video( client, case_id, model=model_path, prompt=sampling_params.prompt, size=output_size, seconds=video_seconds, input_reference=fh, extra_body={ "fps": sampling_params.fps, "num_frames": sampling_params.num_frames, **extra_body, }, ) def generate_mesh(case_id, client) -> tuple[str, bytes]: """I2M: Image to Mesh generation using async /v1/meshes API.""" import requests as http_requests if not sampling_params.image_path: pytest.skip(f"{case_id}: no input image configured for mesh generation") image_path = sampling_params.image_path if isinstance(image_path, str) and is_image_url(image_path): image_path = download_image_from_url(image_path) elif isinstance(image_path, Path): if not image_path.exists(): pytest.skip(f"{case_id}: image file missing: {image_path}") else: image_path = Path(str(image_path)) if not image_path.exists(): pytest.skip(f"{case_id}: image file missing: {image_path}") base_url = str(client.base_url).rstrip("/") if base_url.endswith("/v1"): base_url = base_url[:-3] create_url = f"{base_url}/v1/meshes" with open(str(image_path), "rb") as img_file: files = {"image": (Path(str(image_path)).name, img_file, "image/png")} data = { "prompt": "generate 3d mesh", "model": model_path, "seed": "0", "guidance_scale": "5.0", "num_inference_steps": "50", } logger.info(f"[Mesh Gen] Sending request to {create_url}") try: response = http_requests.post( create_url, files=files, data=data, timeout=60 ) except Exception as e: pytest.fail(f"{case_id}: mesh creation request failed: {e}") if response.status_code != 200: pytest.fail(f"{case_id}: mesh creation failed: {response.text}") job = response.json() mesh_id = job.get("id") if not mesh_id: pytest.fail(f"{case_id}: no mesh id in response: {job}") poll_url = f"{base_url}/v1/meshes/{mesh_id}" poll_interval = 5 max_wait = 1200 elapsed = 0 while elapsed < max_wait: time.sleep(poll_interval) elapsed += poll_interval try: poll_resp = http_requests.get(poll_url, timeout=30) except Exception as e: logger.warning(f"[Mesh Gen] Poll failed: {e}") continue if poll_resp.status_code != 200: continue status_data = poll_resp.json() status = status_data.get("status", "") if status == "completed": content_url = f"{base_url}/v1/meshes/{mesh_id}/content" try: content_resp = http_requests.get(content_url, timeout=60) except Exception as e: pytest.fail(f"{case_id}: mesh download failed: {e}") if content_resp.status_code != 200: pytest.fail(f"{case_id}: mesh download failed: {content_resp.text}") temp_path = Path(tempfile.gettempdir()) / f"mesh_test_{mesh_id}.glb" temp_path.write_bytes(content_resp.content) MESH_OUTPUT_PATHS[case_id] = str(temp_path) logger.info(f"[Mesh Gen] Mesh downloaded to {temp_path}") return (mesh_id, b"") elif status == "failed": error = status_data.get("error", {}) pytest.fail(f"{case_id}: mesh generation failed: {error}") pytest.fail(f"{case_id}: mesh generation timed out after {max_wait}s") if modality == "3d": fn = generate_mesh elif modality == "video": if sampling_params.image_path and sampling_params.prompt: if getattr(sampling_params, "direct_url_test", False): fn = generate_text_url_image_to_video else: fn = generate_text_image_to_video elif sampling_params.image_path: fn = generate_image_to_video else: fn = generate_video elif sampling_params.prompt and sampling_params.image_path: if getattr(sampling_params, "direct_url_test", False): fn = generate_image_edit_url else: fn = generate_image_edit else: fn = generate_image return fn ================================================ FILE: python/sglang/multimodal_gen/test/server/test_update_weights_from_disk.py ================================================ """Tests for diffusion `update_weights_from_disk`. This module verifies the ability to update model weights in place without restarting the server, which is critical for RL workflows and iterative fine-tuning scenarios. Author: Menyang Liu, https://github.com/dreamyang-liu Chenyang Zhao, https://github.com/zhaochenyang20 We use two model pairs for testing (base model / instruct model pairs): - FLUX.2-klein-base-4B / FLUX.2-klein-4B - Qwen/Qwen-Image / Qwen/Qwen-Image-2512 These model pairs share the same architecture but differ in transformer weights. The basic testing logic is to refit the instruct model into the base model and verify the checksum of the transformer weights are the same, which simulates the real-world RL scenario. However, since these two model pairs only differ in transformer weights, and we want to verify update a specific module with update_weights_from_disk API, we need to create a perturbed instruct model that adds noise to the vae weights. In this sense, the instruct model differs from the base model in vae and transformer weights, the text encoder are still the same. To strictly verify the correctness of the refit API, we compare the checksum in SHA-256 on the disk and the server. NOTE and TODO: In the refit a specific module test, we randomly select one module from the transformer and vae to refit the server and keep other modules the same. As described above, the vae's weights are perturbed. If we select the vae to be the target module, ideally speaking, we should assert that the refitted vae's checksum is the same as directly computed from the perturbed vae weights in the disk. However, since the there is complex weight-name remapping and QKV merge during model loading, it is not easy to compare the server-disk checksum for vae and text encoder directly. Therefore, if the target module is vae, we only verify that the refitted vae's checksum is different from the base model's vae's checksum. It should be good issue to solve for the community to adds comparison the server-disk checksum for vae and text encoder in this test. ============================================================================= Test organization: 7 test cases in 2 classes; two model pairs are tested locally, one in CI. ============================================================================= Class 1: TestUpdateWeightsFromDisk (6 tests) — API contract, checksum & rollback Class 2: TestUpdateWeightsFromDiskWithOffload (1 test) — Offload-aware update + checksum ----------------------------------------------------------------------------- Class 1: TestUpdateWeightsFromDisk Validate the update_weights_from_disk API contract, request/response shape, error handling, checksum verification, and corrupted-weight rollback. All tests share one class-scoped server (same process, same in-memory weights). Tests that require "base model then update" should be explicitly reset to base model first so behavior is order-independent and updates are real (base -> perturbed), not no-ops (perturbed -> perturbed). • test_update_weights_from_disk_default base model -> perturbed model with flush_cache=True. Verifies after-update transformer checksum == perturbed model's transformer disk checksum • test_update_weights_specific_modules base -> perturbed with flush_cache=False. Randomly selects one module from _DIFFERING_MODULES (transformer and vae) as target_modules, updates only that module. Verifies that: (1) targeted module's in-memory checksum changed; (2) non-targeted modules' in-memory checksums are unchanged. • test_update_weights_nonexistent_model model_path set to a non-existent path; must fail (400, success=False). Ensure server is healthy after failed update and server's transformer checksums equal base model's transformer disk checksum. • test_update_weights_missing_model_path Request body empty (no model_path); must fail (400, success=False). Ensure server is healthy after failed update and server's transformer checksums equal base model's transformer disk checksum. • test_update_weights_nonexistent_module target_modules=["nonexistent_module"]; must fail (400, success=False). Verify server is healthy after failed update and server's checksums equal base model's transformer disk checksum. • test_corrupted_weights_rollback All-or-nothing rollback: We first refit the server from base model -> perturbed model. We manually truncate the vae weights of the base model to get a corrupted model. We then call the refit to update the server from the perturbed model -> corrupted model. Verify that: 1. The update fails due to truncated vae, server should roll back to the perturbed model, i.e., server's transformer weights == perturbed model's transformer weights != base model's transformer weights. 2. After the rollback, server's vae weights == perturbed model's vae weights != base model's vae weights. 3. After the rollback, server's text encoder weights == base model's text encoder weights == perturbed model's text encoder weights. ----------------------------------------------------------------------------- Class 2: TestUpdateWeightsFromDiskWithOffload Ensure weight updates and checksum verification work when layerwise offload is enabled (--dit-layerwise-offload). With offload, parameters live in CPU buffers and only left small torch.empty((1,)) as placeholders on GPU; the updater must write into CPU buffers and update prefetched GPU tensors without shape mismatch. • test_update_weights_with_offload_enabled Server with --dit-layerwise-offload (base). Load perturbed checkpoint; must succeed (200, success=True), no "Shape mismatch". server's transformer checksum matches perturbed model's transformer disk checksum. """ from __future__ import annotations import functools import os import random import shutil import tempfile import threading from collections.abc import Callable import pytest import requests from safetensors.torch import load_file, save_file from sglang.multimodal_gen.runtime.loader.utils import ( _list_safetensors_files, ) from sglang.multimodal_gen.runtime.loader.weight_utils import ( compute_weights_checksum, safetensors_weights_iterator, ) from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.test.server.test_server_utils import ( ServerManager, ) from sglang.multimodal_gen.test.test_utils import ( DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST, DEFAULT_FLUX_2_KLEIN_BASE_4B_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_2512_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST, get_dynamic_server_port, is_in_ci, ) logger = init_logger(__name__) _TRANSFORMER_MODULE = "transformer" _VAE_MODULE = "vae" _TEXT_ENCODER_MODULE_PREFIX = "text_encoder" # Modules whose weights differ between the base model and the perturbed # perturbed checkpoint _DIFFERING_MODULES: list[str] = [_TRANSFORMER_MODULE, _VAE_MODULE] _ALL_MODEL_PAIRS: list[tuple[str, str]] = [ ( DEFAULT_FLUX_2_KLEIN_BASE_4B_MODEL_NAME_FOR_TEST, DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST, ), ( DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_2512_MODEL_NAME_FOR_TEST, ), ] _CI_MODEL_PAIR_ENV = "SGLANG_MMGEN_UPDATE_WEIGHTS_PAIR" def _resolve_active_model_pairs() -> list[tuple[str, str]]: if not is_in_ci(): return _ALL_MODEL_PAIRS pair_by_id = {pair[0].split("/")[-1]: pair for pair in _ALL_MODEL_PAIRS} selected_pair_id = os.environ.get(_CI_MODEL_PAIR_ENV) if selected_pair_id is None: return [random.choice(_ALL_MODEL_PAIRS)] selected_pair = pair_by_id.get(selected_pair_id) if selected_pair is None: valid_ids = ", ".join(sorted(pair_by_id)) raise ValueError( f"Invalid {_CI_MODEL_PAIR_ENV}={selected_pair_id!r}. " f"Expected one of: {valid_ids}." ) return [selected_pair] _ACTIVE_MODEL_PAIRS = _resolve_active_model_pairs() _PAIR_IDS = [p[0].split("/")[-1] for p in _ACTIVE_MODEL_PAIRS] @functools.lru_cache(maxsize=None) def _compute_checksum_from_disk(model_path: str, module_name: str) -> str: """Compute SHA-256 checksum from safetensors files on disk. Uses the same compute_weights_checksum function as the server, so the checksums are directly comparable. Results are cached (keyed on model_path and module_name) because the same disk checksum is requested multiple times across tests. """ local_path = maybe_download_model(model_path) weights_dir = os.path.join(local_path, module_name) assert os.path.exists( weights_dir ), f"No weights dir for {module_name} in {local_path}" safetensors_files = _list_safetensors_files(weights_dir) assert safetensors_files, f"No safetensors files in {weights_dir}" return compute_weights_checksum(safetensors_weights_iterator(safetensors_files)) def _clone_model_with_modified_module( src_model: str, dst_model: str, target_module: str, transform_safetensor: Callable[[str, str], None], ) -> None: # Symlink root-level files (model_index.json, etc.). for fname in os.listdir(src_model): src_path = os.path.join(src_model, fname) dst_path = os.path.join(dst_model, fname) if os.path.isfile(src_path) and not os.path.exists(dst_path): os.symlink(src_path, dst_path) for module_dir in sorted(os.listdir(src_model)): src_dir = os.path.join(src_model, module_dir) dst_dir = os.path.join(dst_model, module_dir) if not os.path.isdir(src_dir): continue if module_dir != target_module: if not os.path.exists(dst_dir): os.symlink(src_dir, dst_dir) continue os.makedirs(dst_dir, exist_ok=True) transformed = False for fname in sorted(os.listdir(src_dir)): src_file = os.path.join(src_dir, fname) dst_file = os.path.join(dst_dir, fname) if not os.path.isfile(src_file): continue if not fname.endswith(".safetensors") or transformed: if not os.path.exists(dst_file): os.symlink(src_file, dst_file) continue transform_safetensor(src_file, dst_file) transformed = True def _truncate_safetensor(src_file: str, dst_file: str) -> None: shutil.copy2(src_file, dst_file) size = os.path.getsize(dst_file) with open(dst_file, "r+b") as f: f.truncate(size - 2) logger.info( "Created corrupted safetensors: %s (%d -> %d bytes)", dst_file, size, size - 2, ) def _perturb_safetensor(src_file: str, dst_file: str) -> None: tensors = load_file(src_file) perturbed = { k: (t + 0.01 if t.is_floating_point() else t) for k, t in tensors.items() } save_file(perturbed, dst_file) logger.info("Created perturbed safetensors: %s", dst_file) class _UpdateWeightsApiMixin: def _update_weights( self, base_url: str, model_path: str, flush_cache: bool = True, target_modules: list[str] | None = None, timeout: int = 300, ) -> tuple[dict, int]: payload = {"model_path": model_path, "flush_cache": flush_cache} if target_modules is not None: payload["target_modules"] = target_modules response = requests.post( f"{base_url}/update_weights_from_disk", json=payload, timeout=timeout, ) return response.json(), response.status_code def _get_weights_checksum( self, base_url: str, module_names: list[str] | None = None, timeout: int = 300, ) -> dict: payload = {} if module_names is not None: payload["module_names"] = module_names response = requests.post( f"{base_url}/get_weights_checksum", json=payload, timeout=timeout, ) assert ( response.status_code == 200 ), f"get_weights_checksum failed: {response.status_code} {response.text}" return response.json() def _assert_server_matches_model( self, base_url: str, expected_model: str, ) -> None: server_checksums = self._get_weights_checksum( base_url, module_names=[_TRANSFORMER_MODULE] ) expected_cs = _compute_checksum_from_disk(expected_model, _TRANSFORMER_MODULE) server_cs = server_checksums.get(_TRANSFORMER_MODULE) assert server_cs == expected_cs, ( f"Checksum mismatch on '{_TRANSFORMER_MODULE}'\n" f" expected({expected_model}): {expected_cs}\n" f" server: {server_cs}" ) class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin): @pytest.fixture( scope="class", params=_ACTIVE_MODEL_PAIRS, ids=_PAIR_IDS, ) def diffusion_server_no_offload(self, request): default_model, source_model = request.param port = get_dynamic_server_port() wait_deadline = float(os.environ.get("SGLANG_TEST_WAIT_SECS", "600")) manager = ServerManager( model=default_model, port=port, wait_deadline=wait_deadline, extra_args="--num-gpus 1", ) # Ensure models are local before spawning threads that need the paths. local_default = maybe_download_model(default_model) local_source = maybe_download_model(source_model) perturbed_vae_model_dir = tempfile.mkdtemp(prefix="sglang_perturbed_vae_") corrupted_vae_model_dir = tempfile.mkdtemp(prefix="sglang_corrupted_") # Run all disk I/O in background while the server boots. bg_threads = [ threading.Thread( target=_compute_checksum_from_disk, args=(default_model, module) ) for module in _DIFFERING_MODULES ] + [ threading.Thread( target=_clone_model_with_modified_module, args=( local_source, perturbed_vae_model_dir, _VAE_MODULE, _perturb_safetensor, ), ), threading.Thread( target=_clone_model_with_modified_module, args=( local_default, corrupted_vae_model_dir, _VAE_MODULE, _truncate_safetensor, ), ), ] for t in bg_threads: t.start() ctx = manager.start() for t in bg_threads: t.join() # Sanity: all _DIFFERING_MODULES should differ between base and perturbed. for module in _DIFFERING_MODULES: assert _compute_checksum_from_disk( default_model, module ) != _compute_checksum_from_disk(perturbed_vae_model_dir, module), ( f"Assumption violated: {module} should differ between " f"{default_model} and {perturbed_vae_model_dir}" ) try: yield ctx, default_model, perturbed_vae_model_dir, corrupted_vae_model_dir finally: ctx.cleanup() shutil.rmtree(perturbed_vae_model_dir, ignore_errors=True) shutil.rmtree(corrupted_vae_model_dir, ignore_errors=True) def test_update_weights_from_disk_default(self, diffusion_server_no_offload): """Default update (target_modules=None, flush_cache=True): all changed modules updated.""" ctx, default_model, perturbed_model_dir, _ = diffusion_server_no_offload base_url = f"http://localhost:{ctx.port}" self._update_weights(base_url, default_model, flush_cache=True) result, status_code = self._update_weights( base_url, perturbed_model_dir, flush_cache=True ) assert status_code == 200 assert result.get("success", False), f"Update failed: {result.get('message')}" self._assert_server_matches_model(base_url, perturbed_model_dir) def test_update_weights_specific_modules(self, diffusion_server_no_offload): ctx, default_model, perturbed_model_dir, _ = diffusion_server_no_offload base_url = f"http://localhost:{ctx.port}" # Reset server to default_model. self._update_weights(base_url, default_model) before_checksums = self._get_weights_checksum( base_url, module_names=_DIFFERING_MODULES ) target_modules = [random.choice(_DIFFERING_MODULES)] result, status_code = self._update_weights( base_url, perturbed_model_dir, target_modules=target_modules, flush_cache=False, ) assert status_code == 200, f"Update failed: {result}" assert result.get("success", False), f"Update failed: {result.get('message')}" after_checksums = self._get_weights_checksum( base_url, module_names=_DIFFERING_MODULES ) # Targeted module should have changed. for name in target_modules: assert after_checksums.get(name) != before_checksums.get(name), ( f"Targeted module '{name}' checksum should change after update\n" f" before: {before_checksums.get(name)}\n" f" after: {after_checksums.get(name)}" ) # Non-targeted modules should be unchanged. for name, cs in after_checksums.items(): if name in target_modules or cs == "not_found": continue assert cs == before_checksums.get(name), ( f"Non-targeted module '{name}' should be unchanged\n" f" before: {before_checksums.get(name)}\n" f" after: {cs}" ) def test_update_weights_nonexistent_model(self, diffusion_server_no_offload): """Nonexistent model path must fail (400). Server healthy, checksums == base disk.""" ctx, default_model, _, _ = diffusion_server_no_offload base_url = f"http://localhost:{ctx.port}" self._update_weights(base_url, default_model) result, status_code = self._update_weights( base_url, "/nonexistent/path/to/model", timeout=60, ) logger.info(f"Update result for nonexistent model: {result}") assert status_code == 400, f"Expected 400, got {status_code}" assert not result.get("success", True), "Should fail for nonexistent model" self._assert_server_matches_model(base_url, default_model) def test_update_weights_missing_model_path(self, diffusion_server_no_offload): """Request without model_path must fail (400). Server healthy, checksums == base disk.""" ctx, default_model, _, _ = diffusion_server_no_offload base_url = f"http://localhost:{ctx.port}" self._update_weights(base_url, default_model) response = requests.post( f"{base_url}/update_weights_from_disk", json={}, timeout=30, ) assert response.status_code == 400, f"Expected 400, got {response.status_code}" result = response.json() assert not result.get("success", True), "Should fail when model_path is missing" self._assert_server_matches_model(base_url, default_model) def test_update_weights_nonexistent_module(self, diffusion_server_no_offload): """Nonexistent module must fail (400). Server healthy, checksums == base disk.""" ctx, default_model, perturbed_model_dir, _ = diffusion_server_no_offload base_url = f"http://localhost:{ctx.port}" self._update_weights(base_url, default_model) result, status_code = self._update_weights( base_url, perturbed_model_dir, target_modules=["nonexistent_module"], timeout=60, ) logger.info(f"Update nonexistent module result: {result}") assert status_code == 400, f"Expected 400, got {status_code}" assert not result.get("success", True), "Should fail for nonexistent module" assert "not found in pipeline" in result.get("message", "") self._assert_server_matches_model(base_url, default_model) def test_corrupted_weights_rollback(self, diffusion_server_no_offload): ctx, default_model, perturbed_model_dir, corrupted_vae_model_dir = ( diffusion_server_no_offload ) base_url = f"http://localhost:{ctx.port}" # base → perturbed self._update_weights(base_url, default_model) base_checksums = self._get_weights_checksum(base_url) result, status_code = self._update_weights(base_url, perturbed_model_dir) assert status_code == 200 and result.get("success") perturbed_checksums = self._get_weights_checksum(base_url) text_encoder_modules = sorted( name for name in perturbed_checksums if _TEXT_ENCODER_MODULE_PREFIX in name and perturbed_checksums.get(name) != "not_found" and base_checksums.get(name) != "not_found" ) assert ( text_encoder_modules ), "Expected at least one text encoder module checksum" # perturbed → corrupted (should fail and rollback) rollback_targets = [_TRANSFORMER_MODULE, _VAE_MODULE] result, status_code = self._update_weights( base_url, corrupted_vae_model_dir, target_modules=rollback_targets, ) assert ( status_code == 400 ), f"Expected 400 on corrupted weights, got {status_code}" assert not result.get("success", True) message = result.get("message", "") assert "rolled back" in message.lower() # The updater reports the first failing module in the error message. # With ordered target_modules=[transformer, vae], this makes the # failure point explicit: transformer is processed first, then vae fails. assert ( "Failed to update module 'vae'" in message ), f"Expected vae to be the explicit failure point, got: {message}" rolled_back_checksums = self._get_weights_checksum(base_url) # 1) transformer: server == perturbed != base transformer_base = base_checksums.get(_TRANSFORMER_MODULE) transformer_perturbed = perturbed_checksums.get(_TRANSFORMER_MODULE) transformer_rolled_back = rolled_back_checksums.get(_TRANSFORMER_MODULE) assert transformer_rolled_back == transformer_perturbed assert transformer_rolled_back != transformer_base # 2) vae: server == perturbed != base vae_base = base_checksums.get(_VAE_MODULE) vae_perturbed = perturbed_checksums.get(_VAE_MODULE) vae_rolled_back = rolled_back_checksums.get(_VAE_MODULE) assert vae_rolled_back == vae_perturbed assert vae_rolled_back != vae_base # 3) text encoder(s): server == base == perturbed for name in text_encoder_modules: assert rolled_back_checksums.get(name) == perturbed_checksums.get( name ), f"Text encoder module '{name}' should stay equal to perturbed" assert rolled_back_checksums.get(name) == base_checksums.get( name ), f"Text encoder module '{name}' should stay equal to base" class TestUpdateWeightsFromDiskWithOffload(_UpdateWeightsApiMixin): """Test update_weights_from_disk with layerwise offload enabled.""" @pytest.fixture(scope="class", params=_ACTIVE_MODEL_PAIRS, ids=_PAIR_IDS) def diffusion_server_with_offload(self, request): default_model, source_model = request.param port = get_dynamic_server_port() wait_deadline = float(os.environ.get("SGLANG_TEST_WAIT_SECS", "600")) local_source = maybe_download_model(source_model) perturbed_vae_model_dir = tempfile.mkdtemp(prefix="sglang_perturbed_vae_") clone_thread = threading.Thread( target=_clone_model_with_modified_module, args=( local_source, perturbed_vae_model_dir, _VAE_MODULE, _perturb_safetensor, ), ) clone_thread.start() manager = ServerManager( model=default_model, port=port, wait_deadline=wait_deadline, extra_args="--num-gpus 1 --dit-layerwise-offload true", ) ctx = manager.start() clone_thread.join() try: yield ctx, default_model, perturbed_vae_model_dir finally: ctx.cleanup() shutil.rmtree(perturbed_vae_model_dir, ignore_errors=True) def test_update_weights_with_offload_enabled(self, diffusion_server_with_offload): ctx, _, perturbed_model_dir = diffusion_server_with_offload base_url = f"http://localhost:{ctx.port}" result, status_code = self._update_weights(base_url, perturbed_model_dir) assert status_code == 200, f"Expected 200, got {status_code}" assert result.get("success", False), f"Update failed: {result.get('message')}" message = result.get("message", "") assert "Shape mismatch" not in message, f"Shape mismatch detected: {message}" self._assert_server_matches_model(base_url, perturbed_model_dir) if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) ================================================ FILE: python/sglang/multimodal_gen/test/server/testcase_configs.py ================================================ """ Configuration and data structures for diffusion performance tests. Usage: pytest python/sglang/multimodal_gen/test/server/test_server_a.py # for a single testcase, look for the name of the testcases in DIFFUSION_CASES pytest python/sglang/multimodal_gen/test/server/test_server_a.py -k qwen_image_t2i To add a new testcase: 1. add your testcase with case-id: `my_new_test_case_id` to DIFFUSION_CASES 2. run `SGLANG_GEN_BASELINE=1 pytest -s python/sglang/multimodal_gen/test/server/ -k my_new_test_case_id` 3. insert or override the corresponding scenario in `scenarios` section of perf_baselines.json with the output baseline of step-2 """ from __future__ import annotations import json import os import statistics from dataclasses import dataclass, field from pathlib import Path from typing import Sequence from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord from sglang.multimodal_gen.test.test_utils import ( DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST, DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST, DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_EDIT_2509_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_EDIT_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_LAYERED_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, DEFAULT_WAN_2_1_I2V_14B_480P_MODEL_NAME_FOR_TEST, DEFAULT_WAN_2_1_I2V_14B_720P_MODEL_NAME_FOR_TEST, DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST, DEFAULT_WAN_2_1_T2V_14B_MODEL_NAME_FOR_TEST, DEFAULT_WAN_2_2_I2V_A14B_MODEL_NAME_FOR_TEST, DEFAULT_WAN_2_2_T2V_A14B_MODEL_NAME_FOR_TEST, DEFAULT_WAN_2_2_TI2V_5B_MODEL_NAME_FOR_TEST, ) @dataclass class ToleranceConfig: """Tolerance ratios for performance validation.""" e2e: float denoise_stage: float non_denoise_stage: float denoise_step: float denoise_agg: float @classmethod def load_profile(cls, all_tolerances: dict, profile_name: str) -> ToleranceConfig: """Load a specific tolerance profile from a dictionary of profiles.""" # Support both flat structure (backward compatibility) and profiled structure if "e2e" in all_tolerances and not isinstance(all_tolerances["e2e"], dict): tol_data = all_tolerances actual_profile = "legacy/flat" else: tol_data = all_tolerances.get( profile_name, all_tolerances.get("pr_test", {}) ) actual_profile = ( profile_name if profile_name in all_tolerances else "pr_test" ) if not tol_data: raise ValueError( f"No tolerance profile found for '{profile_name}' and no default 'pr_test' profile exists." ) print(f"--- Performance Tolerance Profile: {actual_profile} ---") return cls( e2e=float(os.getenv("SGLANG_E2E_TOLERANCE", tol_data["e2e"])), denoise_stage=float( os.getenv("SGLANG_STAGE_TIME_TOLERANCE", tol_data["denoise_stage"]) ), non_denoise_stage=float( os.getenv( "SGLANG_NON_DENOISE_STAGE_TIME_TOLERANCE", tol_data["non_denoise_stage"], ) ), denoise_step=float( os.getenv("SGLANG_DENOISE_STEP_TOLERANCE", tol_data["denoise_step"]) ), denoise_agg=float( os.getenv("SGLANG_DENOISE_AGG_TOLERANCE", tol_data["denoise_agg"]) ), ) @dataclass class ScenarioConfig: """Expected performance metrics for a test scenario.""" stages_ms: dict[str, float] denoise_step_ms: dict[int, float] expected_e2e_ms: float expected_avg_denoise_ms: float expected_median_denoise_ms: float @dataclass class BaselineConfig: """Full baseline configuration.""" scenarios: dict[str, ScenarioConfig] step_fractions: Sequence[float] tolerances: ToleranceConfig improvement_threshold: float @classmethod def load(cls, path: Path) -> BaselineConfig: """Load baseline configuration from JSON file.""" with path.open("r", encoding="utf-8") as fh: data = json.load(fh) # Get tolerance profile, defaulting to 'pr_test' profile_name = "pr_test" tolerances = ToleranceConfig.load_profile( data.get("tolerances", {}), profile_name ) scenarios = {} for name, cfg in data["scenarios"].items(): scenarios[name] = ScenarioConfig( stages_ms=cfg["stages_ms"], denoise_step_ms={int(k): v for k, v in cfg["denoise_step_ms"].items()}, expected_e2e_ms=float(cfg["expected_e2e_ms"]), expected_avg_denoise_ms=float(cfg["expected_avg_denoise_ms"]), expected_median_denoise_ms=float(cfg["expected_median_denoise_ms"]), ) return cls( scenarios=scenarios, step_fractions=tuple(data["sampling"]["step_fractions"]), tolerances=tolerances, improvement_threshold=data.get("improvement_reporting", {}).get( "threshold", 0.2 ), ) def update(self, path: Path): """Load baseline configuration from JSON file.""" with path.open("r", encoding="utf-8") as fh: data = json.load(fh) scenarios_new = {} for name, cfg in data["scenarios"].items(): scenarios_new[name] = ScenarioConfig( stages_ms=cfg["stages_ms"], denoise_step_ms={int(k): v for k, v in cfg["denoise_step_ms"].items()}, expected_e2e_ms=float(cfg["expected_e2e_ms"]), expected_avg_denoise_ms=float(cfg["expected_avg_denoise_ms"]), expected_median_denoise_ms=float(cfg["expected_median_denoise_ms"]), ) self.scenarios.update(scenarios_new) return self @dataclass class DiffusionServerArgs: """Configuration for a single model/scenario test case.""" model_path: str # HF repo or local path modality: str = "image" # "image" or "video" or "3d" custom_validator: str | None = None # optional custom validator name # resources num_gpus: int = 1 tp_size: int | None = None ulysses_degree: int | None = None ring_degree: int | None = None cfg_parallel: bool | None = None # LoRA lora_path: str | None = ( None # LoRA adapter path (HF repo or local path, loaded at startup) ) dynamic_lora_path: str | None = ( None # LoRA path for dynamic loading test (loaded via set_lora after startup) ) second_lora_path: str | None = ( None # Second LoRA adapter path for multi-LoRA testing ) dit_layerwise_offload: bool = False dit_offload_prefetch_size: int | float | None = None enable_cache_dit: bool = False text_encoder_cpu_offload: bool = False extras: list[str] = field(default_factory=lambda: []) def __post_init__(self): if self.modality == "image": self.custom_validator = "image" elif self.modality == "video": self.custom_validator = "video" elif self.modality == "3d": self.custom_validator = "mesh" @dataclass(frozen=True) class DiffusionSamplingParams: """Configuration for a single model/scenario test case.""" output_size: str = "" # inputs and conditioning prompt: str | None = None # text prompt for generation image_path: Path | str | None = None # input image/video for editing (Path or URL) # duration seconds: int = 1 # for video: duration in seconds num_frames: int | None = None # for video: number of frames fps: int | None = None # for video: frames per second # URL direct test flag - if True, don't pre-download URL images direct_url_test: bool = False # output format output_format: str | None = None # "png", "jpeg", "mp4", etc. num_outputs_per_prompt: int = 1 # Additional request-level parameters (e.g. enable_teacache, enable_upscaling, …) # merged directly into the OpenAI extra_body dict. extras: dict = field(default_factory=dict) @dataclass(frozen=True) class DiffusionTestCase: """Configuration for a single model/scenario test case.""" id: str # pytest test id and scenario name server_args: DiffusionServerArgs sampling_params: DiffusionSamplingParams run_perf_check: bool = True def sample_step_indices( step_map: dict[int, float], fractions: Sequence[float] ) -> list[int]: if not step_map: return [] max_idx = max(step_map.keys()) indices = set() for fraction in fractions: idx = min(max_idx, max(0, int(round(fraction * max_idx)))) if idx in step_map: indices.add(idx) return sorted(indices) @dataclass class PerformanceSummary: """Summary of performance of a request, built from RequestPerfRecord""" e2e_ms: float avg_denoise_ms: float median_denoise_ms: float # { "stage_1": time_1, "stage_2": time_2 } stage_metrics: dict[str, float] step_metrics: list[float] sampled_steps: dict[int, float] all_denoise_steps: dict[int, float] frames_per_second: float | None = None total_frames: int | None = None avg_frame_time_ms: float | None = None @staticmethod def from_req_perf_record( record: RequestPerfRecord, step_fractions: Sequence[float] ): """Collect all performance metrics into a summary without validation.""" e2e_ms = record.total_duration_ms step_durations = record.steps avg_denoise = 0.0 median_denoise = 0.0 if step_durations: avg_denoise = sum(step_durations) / len(step_durations) median_denoise = statistics.median(step_durations) per_step = {index: s for index, s in enumerate(step_durations)} sample_indices = sample_step_indices(per_step, step_fractions) sampled_steps = {idx: per_step[idx] for idx in sample_indices} # convert from list to dict stage_metrics = {} for item in record.stages: if isinstance(item, dict) and "name" in item: val = item.get("execution_time_ms", 0.0) stage_metrics[item["name"]] = val return PerformanceSummary( e2e_ms=e2e_ms, avg_denoise_ms=avg_denoise, median_denoise_ms=median_denoise, stage_metrics=stage_metrics, step_metrics=step_durations, sampled_steps=sampled_steps, all_denoise_steps=per_step, ) T2I_sampling_params = DiffusionSamplingParams( prompt="Doraemon is eating dorayaki", output_size="1024x1024", ) TI2I_sampling_params = DiffusionSamplingParams( prompt="Convert 2D style to 3D style", image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg", ) MULTI_IMAGE_TI2I_sampling_params = DiffusionSamplingParams( prompt="The magician bear is on the left, the alchemist bear is on the right, facing each other in the central park square.", image_path=[ "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_1.jpg", "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_2.jpg", ], direct_url_test=True, ) MULTI_IMAGE_TI2I_UPLOAD_sampling_params = DiffusionSamplingParams( prompt="The magician bear is on the left, the alchemist bear is on the right, facing each other in the central park square.", image_path=[ "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_1.jpg", "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_2.jpg", ], ) MULTI_FRAME_I2I_sampling_params = DiffusionSamplingParams( prompt="a high quality, cute halloween themed illustration, consistent style and lighting", image_path=[ "https://raw.githubusercontent.com/QwenLM/Qwen-Image-Layered/main/assets/test_images/4.png" ], num_frames=4, direct_url_test=True, output_format="png", ) T2V_PROMPT = "A curious raccoon" TI2V_sampling_params = DiffusionSamplingParams( prompt="The man in the picture slowly turns his head, his expression enigmatic and otherworldly. The camera performs a slow, cinematic dolly out, focusing on his face. Moody lighting, neon signs glowing in the background, shallow depth of field.", image_path="https://is1-ssl.mzstatic.com/image/thumb/Music114/v4/5f/fa/56/5ffa56c2-ea1f-7a17-6bad-192ff9b6476d/825646124206.jpg/600x600bb.jpg", direct_url_test=True, ) TURBOWAN_I2V_sampling_params = DiffusionSamplingParams( prompt="The man in the picture slowly turns his head, his expression enigmatic and otherworldly. The camera performs a slow, cinematic dolly out, focusing on his face. Moody lighting, neon signs glowing in the background, shallow depth of field.", image_path="https://is1-ssl.mzstatic.com/image/thumb/Music114/v4/5f/fa/56/5ffa56c2-ea1f-7a17-6bad-192ff9b6476d/825646124206.jpg/600x600bb.jpg", direct_url_test=True, output_size="960x960", num_frames=4, fps=4, ) # All test cases with clean default values # To test different models, simply add more DiffusionCase entries ONE_GPU_CASES_A: list[DiffusionTestCase] = [ # === Text to Image (T2I) === DiffusionTestCase( "qwen_image_t2i", DiffusionServerArgs( model_path=DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST, modality="image", ), T2I_sampling_params, ), DiffusionTestCase( "qwen_image_t2i_cache_dit_enabled", DiffusionServerArgs( model_path=DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST, modality="image", enable_cache_dit=True, ), T2I_sampling_params, ), DiffusionTestCase( "flux_image_t2i", DiffusionServerArgs( model_path=DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST, modality="image" ), T2I_sampling_params, ), # TODO: modeling of flux different from official flux, so weights can't be loaded # consider opting for a different quantized hf-repo # DiffusionTestCase( # "flux_image_t2i_override_transformer_weights_path_fp8", # DiffusionServerArgs( # model_path="black-forest-labs/FLUX.1-dev", modality="image", # extras=["--transformer-weights-path black-forest-labs/FLUX.1-dev-FP8"] # ), # T2I_sampling_params, # ), DiffusionTestCase( "flux_2_image_t2i", DiffusionServerArgs( model_path=DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST, modality="image" ), T2I_sampling_params, ), DiffusionTestCase( "flux_2_klein_image_t2i", DiffusionServerArgs( model_path=DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST, modality="image", ), T2I_sampling_params, ), # TODO: replace with a faster model to test the --dit-layerwise-offload # TODO: currently, we don't support sending more than one request in test, and setting `num_outputs_per_prompt` to 2 doesn't guarantee the denoising be executed twice, # so we do one warmup and send one request instead DiffusionTestCase( "layerwise_offload", DiffusionServerArgs( model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, modality="image", dit_layerwise_offload=True, dit_offload_prefetch_size=2, ), T2I_sampling_params, ), DiffusionTestCase( "zimage_image_t2i", DiffusionServerArgs( model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, modality="image" ), T2I_sampling_params, ), DiffusionTestCase( "zimage_image_t2i_fp8", DiffusionServerArgs( model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, modality="image", extras=["--transformer-path MickJ/Z-Image-Turbo-fp8"], ), T2I_sampling_params, ), # Multi-LoRA test case for Z-Image-Turbo DiffusionTestCase( "zimage_image_t2i_multi_lora", DiffusionServerArgs( model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, modality="image", lora_path="reverentelusarca/elusarca-anime-style-lora-z-image-turbo", second_lora_path="tarn59/pixel_art_style_lora_z_image_turbo", ), T2I_sampling_params, ), DiffusionTestCase( "sana_image_t2i", DiffusionServerArgs( model_path="Efficient-Large-Model/Sana_600M_1024px_diffusers", modality="image", ), T2I_sampling_params, run_perf_check=False, ), # === Text and Image to Image (TI2I) === DiffusionTestCase( "qwen_image_edit_ti2i", DiffusionServerArgs( model_path=DEFAULT_QWEN_IMAGE_EDIT_MODEL_NAME_FOR_TEST, modality="image" ), TI2I_sampling_params, ), DiffusionTestCase( "qwen_image_edit_2509_ti2i", DiffusionServerArgs( model_path=DEFAULT_QWEN_IMAGE_EDIT_2509_MODEL_NAME_FOR_TEST, modality="image", ), MULTI_IMAGE_TI2I_sampling_params, ), DiffusionTestCase( "qwen_image_edit_2511_ti2i", DiffusionServerArgs( model_path=DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST, modality="image", ), TI2I_sampling_params, ), DiffusionTestCase( "qwen_image_layered_i2i", DiffusionServerArgs( model_path=DEFAULT_QWEN_IMAGE_LAYERED_MODEL_NAME_FOR_TEST, modality="image", ), MULTI_FRAME_I2I_sampling_params, ), # Upscaling (Real-ESRGAN 4×) for T2I DiffusionTestCase( "flux_2_image_t2i_upscaling_4x", DiffusionServerArgs( model_path="black-forest-labs/FLUX.2-dev", modality="image", ), DiffusionSamplingParams( prompt="Doraemon is eating dorayaki", output_size="1024x1024", extras={"enable_upscaling": True, "upscaling_scale": 4}, ), ), ] HUNYUAN3D_SHAPE_sampling_params = DiffusionSamplingParams( prompt="", image_path="https://raw.githubusercontent.com/sgl-project/sgl-test-files/main/diffusion-ci/consistency_gt/1-gpu/hunyuan3d_2_0/hunyuan3d.png", ) ONE_GPU_CASES_B: list[DiffusionTestCase] = [ # === Text to Video (T2V) === DiffusionTestCase( "wan2_1_t2v_1.3b", DiffusionServerArgs( model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST, modality="video", custom_validator="video", ), DiffusionSamplingParams( prompt=T2V_PROMPT, ), ), DiffusionTestCase( "wan2_1_t2v_1.3b_text_encoder_cpu_offload", DiffusionServerArgs( model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST, modality="video", custom_validator="video", text_encoder_cpu_offload=True, ), DiffusionSamplingParams( prompt=T2V_PROMPT, ), ), # TeaCache acceleration test for Wan video model DiffusionTestCase( "wan2_1_t2v_1.3b_teacache_enabled", DiffusionServerArgs( model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST, modality="video", custom_validator="video", ), DiffusionSamplingParams( prompt=T2V_PROMPT, extras={"enable_teacache": True}, ), ), # Frame interpolation (2× / exp=1) # Uses the same 1.3B model already in the suite; DiffusionTestCase( "wan2_1_t2v_1.3b_frame_interp_2x", DiffusionServerArgs( model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", modality="video", custom_validator="video", ), DiffusionSamplingParams( prompt=T2V_PROMPT, extras={"enable_frame_interpolation": True, "frame_interpolation_exp": 1}, ), ), # Upscaling (Real-ESRGAN 4×) # Uses the same 1.3B model already in the suite; DiffusionTestCase( "wan2_1_t2v_1.3b_upscaling_4x", DiffusionServerArgs( model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", modality="video", custom_validator="video", ), DiffusionSamplingParams( prompt=T2V_PROMPT, extras={"enable_upscaling": True, "upscaling_scale": 4}, ), ), # Combined: Frame interpolation (2×) + Upscaling (4×) # Verifies that both post-processing steps compose correctly. DiffusionTestCase( "wan2_1_t2v_1.3b_frame_interp_2x_upscaling_4x", DiffusionServerArgs( model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", modality="video", custom_validator="video", ), DiffusionSamplingParams( prompt=T2V_PROMPT, extras={ "enable_frame_interpolation": True, "frame_interpolation_exp": 1, "enable_upscaling": True, "upscaling_scale": 4, }, ), ), # LoRA test case for single transformer + merge/unmerge API test # Note: Uses dynamic_lora_path instead of lora_path to test LayerwiseOffload + set_lora interaction # Server starts WITHOUT LoRA, then set_lora is called after startup (Wan models auto-enable layerwise offload) DiffusionTestCase( "wan2_1_t2v_1_3b_lora_1gpu", DiffusionServerArgs( model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST, modality="video", custom_validator="video", num_gpus=1, dynamic_lora_path="Cseti/Wan-LoRA-Arcane-Jinx-v1", ), DiffusionSamplingParams( prompt="csetiarcane Nfj1nx with blue hair, a woman walking in a cyberpunk city at night", ), ), # NOTE(mick): flaky # DiffusionTestCase( # "hunyuan_video", # DiffusionServerArgs( # model_path="hunyuanvideo-community/HunyuanVideo", # modality="video", # ), # DiffusionSamplingParams( # prompt=T2V_PROMPT, # ), # ), DiffusionTestCase( "flux_2_ti2i", DiffusionServerArgs( model_path=DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST, modality="image" ), TI2I_sampling_params, ), DiffusionTestCase( "flux_2_t2i_customized_vae_path", DiffusionServerArgs( model_path=DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST, modality="image", extras=["--vae-path=fal/FLUX.2-Tiny-AutoEncoder"], ), T2I_sampling_params, run_perf_check=False, ), DiffusionTestCase( "fast_hunyuan_video", DiffusionServerArgs( model_path="FastVideo/FastHunyuan-diffusers", modality="video", custom_validator="video", ), DiffusionSamplingParams( prompt=T2V_PROMPT, ), ), # === Text and Image to Video (TI2V) === DiffusionTestCase( "wan2_2_ti2v_5b", DiffusionServerArgs( model_path=DEFAULT_WAN_2_2_TI2V_5B_MODEL_NAME_FOR_TEST, modality="video", custom_validator="video", ), TI2V_sampling_params, ), DiffusionTestCase( "fastwan2_2_ti2v_5b", DiffusionServerArgs( model_path="FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", modality="video", custom_validator="video", ), TI2V_sampling_params, ), # === Helios T2V === DiffusionTestCase( "helios_base_t2v", DiffusionServerArgs( model_path="BestWishYsh/Helios-Base", modality="video", ), DiffusionSamplingParams( prompt=T2V_PROMPT, output_size="640x384", num_frames=33, ), ), DiffusionTestCase( "helios_mid_t2v", DiffusionServerArgs( model_path="BestWishYsh/Helios-Mid", modality="video", ), DiffusionSamplingParams( prompt=T2V_PROMPT, output_size="640x384", num_frames=33, ), ), DiffusionTestCase( "helios_distilled_t2v", DiffusionServerArgs( model_path="BestWishYsh/Helios-Distilled", modality="video", ), DiffusionSamplingParams( prompt=T2V_PROMPT, output_size="640x384", num_frames=33, ), ), ] # Skip hunyuan3d on AMD: marching_cubes surface extraction produces invalid SDF on ROCm. if not current_platform.is_hip(): ONE_GPU_CASES_B.append( DiffusionTestCase( "hunyuan3d_shape_gen", DiffusionServerArgs( model_path="tencent/Hunyuan3D-2", modality="3d", ), HUNYUAN3D_SHAPE_sampling_params, ), ) # Skip turbowan on AMD: Triton requires 81920 shared memory, but AMD only has 65536. if not current_platform.is_hip(): ONE_GPU_CASES_B.append( DiffusionTestCase( "turbo_wan2_1_t2v_1.3b", DiffusionServerArgs( model_path="IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers", modality="video", custom_validator="video", ), DiffusionSamplingParams( prompt=T2V_PROMPT, ), ) ) TWO_GPU_CASES_A = [ DiffusionTestCase( "wan2_2_i2v_a14b_2gpu", DiffusionServerArgs( model_path=DEFAULT_WAN_2_2_I2V_A14B_MODEL_NAME_FOR_TEST, modality="video", custom_validator="video", ), TI2V_sampling_params, ), DiffusionTestCase( "wan2_2_t2v_a14b_2gpu", DiffusionServerArgs( model_path=DEFAULT_WAN_2_2_T2V_A14B_MODEL_NAME_FOR_TEST, modality="video", custom_validator="video", num_gpus=2, ), DiffusionSamplingParams( prompt=T2V_PROMPT, ), ), # TeaCache smoke test for Wan2.2 T2V A14B — verifies enable_teacache=True # doesn't crash. Perf check disabled because Wan2.2-specific TeaCache # coefficients are not yet calibrated (teacache_params=None, so no speedup). DiffusionTestCase( "wan2_2_t2v_a14b_teacache_2gpu", DiffusionServerArgs( model_path=DEFAULT_WAN_2_2_T2V_A14B_MODEL_NAME_FOR_TEST, modality="video", custom_validator="video", num_gpus=2, ), DiffusionSamplingParams( prompt=T2V_PROMPT, extras={"enable_teacache": True}, ), run_perf_check=False, ), # LoRA test case for transformer_2 support DiffusionTestCase( "wan2_2_t2v_a14b_lora_2gpu", DiffusionServerArgs( model_path=DEFAULT_WAN_2_2_T2V_A14B_MODEL_NAME_FOR_TEST, modality="video", custom_validator="video", num_gpus=2, lora_path="Cseti/wan2.2-14B-Arcane_Jinx-lora-v1", ), DiffusionSamplingParams( prompt="Nfj1nx with blue hair, a woman walking in a cyberpunk city at night", ), ), DiffusionTestCase( "wan2_1_t2v_14b_2gpu", DiffusionServerArgs( model_path=DEFAULT_WAN_2_1_T2V_14B_MODEL_NAME_FOR_TEST, modality="video", num_gpus=2, custom_validator="video", ), DiffusionSamplingParams( prompt=T2V_PROMPT, output_size="832x480", ), ), DiffusionTestCase( "wan2_1_t2v_1.3b_cfg_parallel", DiffusionServerArgs( model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST, modality="video", custom_validator="video", num_gpus=2, cfg_parallel=True, ), DiffusionSamplingParams( prompt=T2V_PROMPT, ), ), DiffusionTestCase( "fsdp-inference", DiffusionServerArgs( model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, modality="image", num_gpus=2, extras=["--use-fsdp-inference"], ), T2I_sampling_params, ), ] TWO_GPU_CASES_B = [ DiffusionTestCase( "wan2_1_i2v_14b_480P_2gpu", DiffusionServerArgs( model_path=DEFAULT_WAN_2_1_I2V_14B_480P_MODEL_NAME_FOR_TEST, modality="video", custom_validator="video", num_gpus=2, ), TI2V_sampling_params, ), # I2V LoRA test case DiffusionTestCase( "wan2_1_i2v_14b_lora_2gpu", DiffusionServerArgs( model_path=DEFAULT_WAN_2_1_I2V_14B_720P_MODEL_NAME_FOR_TEST, modality="video", custom_validator="video", num_gpus=2, lora_path="starsfriday/Wan2.1-Divine-Power-LoRA", ), TI2V_sampling_params, ), DiffusionTestCase( "wan2_1_i2v_14b_720P_2gpu", DiffusionServerArgs( model_path=DEFAULT_WAN_2_1_I2V_14B_720P_MODEL_NAME_FOR_TEST, modality="video", custom_validator="video", num_gpus=2, ), TI2V_sampling_params, ), DiffusionTestCase( "qwen_image_t2i_2_gpus", DiffusionServerArgs( model_path=DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST, modality="image", num_gpus=2, # test ring attn ulysses_degree=1, ring_degree=2, ), T2I_sampling_params, ), DiffusionTestCase( "zimage_image_t2i_2_gpus", DiffusionServerArgs( model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, modality="image", num_gpus=2, ulysses_degree=2, ), T2I_sampling_params, ), DiffusionTestCase( "flux_image_t2i_2_gpus", DiffusionServerArgs( model_path=DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST, modality="image", num_gpus=2, ), T2I_sampling_params, ), DiffusionTestCase( "flux_2_image_t2i_2_gpus", DiffusionServerArgs( model_path=DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST, modality="image", num_gpus=2, tp_size=2, ), T2I_sampling_params, ), DiffusionTestCase( "flux_2_klein_ti2i_2_gpus", DiffusionServerArgs( model_path="black-forest-labs/FLUX.2-klein-4B", modality="image", num_gpus=2, ), TI2I_sampling_params, ), ] if not current_platform.is_hip(): # Flux2 multi-image edit with cache-dit, regression test ONE_GPU_CASES_B.append( DiffusionTestCase( "flux_2_ti2i_multi_image_cache_dit", DiffusionServerArgs( model_path="black-forest-labs/FLUX.2-dev", modality="image", enable_cache_dit=True, ), MULTI_IMAGE_TI2I_UPLOAD_sampling_params, ) ) # Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536. ONE_GPU_CASES_B.append( DiffusionTestCase( "turbo_wan2_1_t2v_1.3b", DiffusionServerArgs( model_path="IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers", modality="video", custom_validator="video", ), DiffusionSamplingParams( prompt=T2V_PROMPT, ), ) ) # Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536. TWO_GPU_CASES_A.append( DiffusionTestCase( "turbo_wan2_2_i2v_a14b_2gpu", DiffusionServerArgs( model_path="IPostYellow/TurboWan2.2-I2V-A14B-Diffusers", modality="video", custom_validator="video", num_gpus=2, tp_size=2, ), TURBOWAN_I2V_sampling_params, ) ) # Load global configuration BASELINE_CONFIG = BaselineConfig.load( Path(__file__).with_name("perf_baselines.json") ).update(Path(__file__).parent / "ascend" / "perf_baselines_npu.json") ================================================ FILE: python/sglang/multimodal_gen/test/slack_utils.py ================================================ """ This file upload the media generated in diffusion-nightly-test to a slack channel of SGLang """ import logging import os import tempfile from datetime import datetime from typing import List, Union from urllib.parse import urlparse from urllib.request import urlopen from sglang.multimodal_gen.runtime.utils.perf_logger import get_git_commit_hash logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) import inspect try: import sglang.multimodal_gen.test.server.testcase_configs as configs from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase ALL_CASES = [] for name, value in inspect.getmembers(configs): if name.endswith("_CASES") or "_CASES_" in name: if ( isinstance(value, list) and len(value) > 0 and isinstance(value[0], DiffusionTestCase) ): ALL_CASES.extend(value) elif isinstance(value, list) and len(value) == 0: # Assume empty list with matching name is a valid case list container pass # Deduplicate cases by ID seen_ids = set() unique_cases = [] for c in ALL_CASES: if c.id not in seen_ids: seen_ids.add(c.id) unique_cases.append(c) ALL_CASES = unique_cases except Exception as e: logger.warning(f"Failed to import test cases: {e}") ALL_CASES = [] def _get_status_message(run_id, current_case_id, thread_messages=None): date_str = datetime.now().strftime("%d/%m") base_header = f"""🧵 for nightly test of {date_str} *Git Revision:* {get_git_commit_hash()} *GitHub Run ID:* {run_id} *Total Tasks:* {len(ALL_CASES)} """ if not ALL_CASES: return base_header default_emoji_for_case_in_progress = "⏳" status_map = {c.id: default_emoji_for_case_in_progress for c in ALL_CASES} if thread_messages: for msg in thread_messages: text = msg.get("text", "") # Look for case_id in the message (format: *Case ID:* `case_id`) for c in ALL_CASES: if f"*Case ID:* `{c.id}`" in text: status_map[c.id] = "✅" if current_case_id: status_map[current_case_id] = "✅" lines = [base_header, "", "*Tasks Status:*"] # Calculate padding max_len = max(len(c.id) for c in ALL_CASES) if ALL_CASES else 10 max_len = max(max_len, len("Case ID")) # Build markdown table inside a code block table_lines = ["```"] table_lines.append(f"| {'Case ID'.ljust(max_len)} | Status |") table_lines.append(f"| {'-' * max_len} | :----: |") for c in ALL_CASES: mark = status_map.get(c.id, default_emoji_for_case_in_progress) table_lines.append(f"| {c.id.ljust(max_len)} | {mark} |") table_lines.append("```") lines.extend(table_lines) return "\n".join(lines) def upload_file_to_slack( case_id: str = None, model: str = None, prompt: str = None, file_path: str = None, origin_file_path: Union[str, List[str]] = None, ) -> bool: temp_paths = [] try: from slack_sdk import WebClient run_id = os.getenv("GITHUB_RUN_ID", "local") token = os.environ.get("SGLANG_DIFFUSION_SLACK_TOKEN") if not token: logger.info(f"Slack upload failed: no token") return False if not file_path or not os.path.exists(file_path): logger.info(f"Slack upload failed: no file path") return False origin_paths = [] if isinstance(origin_file_path, str): if origin_file_path: origin_paths.append(origin_file_path) elif isinstance(origin_file_path, list): origin_paths = [p for p in origin_file_path if p] final_origin_paths = [] for path in origin_paths: if path.startswith(("http", "https")): try: suffix = os.path.splitext(urlparse(path).path)[1] or ".tmp" with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tf: with urlopen(path) as response: tf.write(response.read()) temp_paths.append(tf.name) final_origin_paths.append(tf.name) except Exception as e: logger.warning(f"Failed to download {path}: {e}") else: final_origin_paths.append(path) uploads = [] for i, path in enumerate(final_origin_paths): if os.path.exists(path): title = ( "Original Image" if len(final_origin_paths) == 1 else f"Original Image {i+1}" ) uploads.append({"file": path, "title": title}) uploads.append({"file": file_path, "title": "Generated Image"}) message = ( f"*Case ID:* `{case_id}`\n" f"*Model:* `{model}`\n" f"*Prompt:* {prompt}" ) client = WebClient(token=token) channel_id = "C0A02NDF7UY" thread_ts = None parent_msg_text = None try: history = client.conversations_history(channel=channel_id, limit=100) for msg in history.get("messages", []): if f"*GitHub Run ID:* {run_id}" in msg.get("text", ""): # Use thread_ts if it exists (msg is a reply), otherwise use ts (msg is a parent) thread_ts = msg.get("thread_ts") or msg.get("ts") parent_msg_text = msg.get("text", "") logger.info(f"Found thread_ts: {thread_ts}") break except Exception as e: logger.warning(f"Failed to search slack history: {e}") if not thread_ts: try: text = _get_status_message(run_id, case_id) response = client.chat_postMessage(channel=channel_id, text=text) thread_ts = response["ts"] except Exception as e: logger.warning(f"Failed to create parent thread: {e}") # Upload first to ensure it's in history client.files_upload_v2( channel=channel_id, file_uploads=uploads, initial_comment=message, thread_ts=thread_ts, ) # Then update status based on thread replies if thread_ts: try: replies = client.conversations_replies( channel=channel_id, ts=thread_ts, limit=200 ) messages = replies.get("messages", []) new_text = _get_status_message(run_id, case_id, messages) # Only update if changed significantly (ignoring timestamp diffs if any) # But here we just check text content if new_text != parent_msg_text: client.chat_update(channel=channel_id, ts=thread_ts, text=new_text) except Exception as e: logger.warning(f"Failed to update parent message: {e}") logger.info(f"File uploaded successfully: {os.path.basename(file_path)}") return True except Exception as e: logger.info(f"Slack upload failed: {e}") return False finally: for p in temp_paths: if os.path.exists(p): os.remove(p) ================================================ FILE: python/sglang/multimodal_gen/test/test_files/launch_flux.json ================================================ { "model_path": "black-forest-labs/FLUX.1-dev", "prompt": "A beautiful woman in a red dress walking down a street", "text_encoder_cpu_offload": true, "pin_cpu_memory": true, "save_output": true, "width": 720, "height": 720, "output_path": "outputs", "output_file_name": "FLUX.1-dev, single gpu" } ================================================ FILE: python/sglang/multimodal_gen/test/test_files/launch_wan.json ================================================ { "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "prompt": "A beautiful woman in a red dress walking down a street", "text_encoder_cpu_offload": true, "pin_cpu_memory": true, "save_output": true, "width": 720, "height": 720, "output_path": "outputs", "output_file_name": "Wan2.1-T2V-1.3B-Diffusers, single gpu" } ================================================ FILE: python/sglang/multimodal_gen/test/test_utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import base64 import io import json import os import socket import subprocess import tempfile import time from pathlib import Path from urllib.parse import urljoin import cv2 import httpx import numpy as np from PIL import Image from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import ( RequestPerfRecord, get_diffusion_perf_log_dir, ) logger = init_logger(__name__) # --------------------------------------------------------------------------- # Common model IDs for diffusion tests # # Centralised here so every test file references the same constants instead # of scattering hard-coded strings. When adding a new model that will be # reused across tests, define it here. # --------------------------------------------------------------------------- DEFAULT_SMALL_MODEL_NAME_FOR_TEST = "Tongyi-MAI/Z-Image-Turbo" # Qwen image generation models DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image" DEFAULT_QWEN_IMAGE_2512_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-2512" DEFAULT_QWEN_IMAGE_EDIT_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-Edit" DEFAULT_QWEN_IMAGE_EDIT_2509_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-Edit-2509" DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-Edit-2511" DEFAULT_QWEN_IMAGE_LAYERED_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-Layered" # FLUX image generation models DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST = "black-forest-labs/FLUX.1-dev" DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST = "black-forest-labs/FLUX.2-dev" DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST = "black-forest-labs/FLUX.2-klein-4B" DEFAULT_FLUX_2_KLEIN_BASE_4B_MODEL_NAME_FOR_TEST = ( "black-forest-labs/FLUX.2-klein-base-4B" ) # Wan video generation models DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DEFAULT_WAN_2_1_T2V_14B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.1-T2V-14B-Diffusers" DEFAULT_WAN_2_1_I2V_14B_480P_MODEL_NAME_FOR_TEST = ( "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" ) DEFAULT_WAN_2_1_I2V_14B_720P_MODEL_NAME_FOR_TEST = ( "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers" ) DEFAULT_WAN_2_2_TI2V_5B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" DEFAULT_WAN_2_2_T2V_A14B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" DEFAULT_WAN_2_2_I2V_A14B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.2-I2V-A14B-Diffusers" def print_value_formatted(description: str, value: int | float | str): """Helper function to print a metric value formatted.""" if isinstance(value, int): if value >= 1e6: value_str = f"{value / 1e6:<30.2f}M" elif value >= 1e3: value_str = f"{value / 1e3:<30.2f}K" else: value_str = f"{value:<30}" elif isinstance(value, float): value_str = f"{value:<30.2f}" else: value_str = f"{value:<30}" print(f"{description:<45} {value_str}") def print_divider(length: int, char: str = "-"): """Helper function to print a divider line.""" print(char * length) def is_image_url(image_path: str | Path | None) -> bool: """Check if image_path is a URL.""" if image_path is None: return False return isinstance(image_path, str) and ( image_path.startswith("http://") or image_path.startswith("https://") ) def probe_port(host="127.0.0.1", port=30010, timeout=2.0) -> bool: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.settimeout(timeout) try: s.connect((host, port)) return True except OSError: return False def is_in_ci() -> bool: return get_bool_env_var("SGLANG_IS_IN_CI") def get_dynamic_server_port() -> int: cuda_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "0") if not cuda_devices: cuda_devices = "0" try: first_device_id = int(cuda_devices.split(",")[0].strip()[0]) except (ValueError, IndexError): first_device_id = 0 if is_in_ci(): base_port = 10000 + first_device_id * 2000 else: base_port = 20000 + first_device_id * 1000 return base_port + 1000 def find_free_port(host: str = "127.0.0.1") -> int: """Bind to port 0 and let the OS assign an available port.""" with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind((host, 0)) return s.getsockname()[1] def wait_for_server_health( base_url: str, path: str = "/health", timeout: float = 180.0, interval: float = 1.0, ) -> None: """Poll ``GET `` until it returns HTTP 200.""" deadline = time.time() + timeout last_err: httpx.RequestError | None = None last_status: int | None = None while time.time() < deadline: try: r = httpx.get(urljoin(base_url, path), timeout=5.0) last_status = r.status_code if r.status_code == 200: return except httpx.RequestError as e: last_err = e time.sleep(interval) raise TimeoutError( f"Server at {urljoin(base_url, path)} not healthy after {timeout}s. " f"{last_status=} {last_err=}" ) def post_json( base_url: str, path: str, payload: dict, timeout: float = 300.0, ) -> httpx.Response: """POST JSON to ```` and return the response.""" return httpx.post(urljoin(base_url, path), json=payload, timeout=timeout) # --------------------------------------------------------------------------- # GPU memory helpers (nvidia-smi) # --------------------------------------------------------------------------- def query_gpu_mem_used_mib(gpu_index: int = 0, required: bool = False) -> int | None: """Return GPU memory usage in MiB via ``nvidia-smi``, or *None* on failure. When *required* is ``True`` the function raises instead of returning ``None``. """ try: out = subprocess.check_output( [ "nvidia-smi", f"--id={gpu_index}", "--query-gpu=memory.used", "--format=csv,noheader,nounits", ], text=True, ).strip() return int(out.splitlines()[0].strip()) except Exception as e: logger.warning(f"nvidia-smi memory query failed: {type(e).__name__}: {e}") assert not required, ( "nvidia-smi memory query is unavailable; " "cannot enforce GPU memory assertions." ) return None def require_gpu_mem_query(gpu_index: int = 0) -> int: """Same as :func:`query_gpu_mem_used_mib` but asserts availability. Raises ``AssertionError`` when ``nvidia-smi`` is unavailable instead of returning ``None``, so callers can rely on a valid ``int`` result. """ mem = query_gpu_mem_used_mib(gpu_index, required=True) assert mem is not None return mem def assert_gpu_mem_changed( label: str, before_mib: int, after_mib: int, min_delta_mib: int, ) -> None: """Assert that GPU memory changed by at least *min_delta_mib* MiB.""" delta = abs(after_mib - before_mib) logger.debug( f"[MEM] {label}: before={before_mib} MiB after={after_mib} MiB |delta|={delta} MiB" ) assert delta >= min_delta_mib, ( f"GPU memory change too small for '{label}': " f"|after-before|={delta} MiB < {min_delta_mib} MiB " f"(before={before_mib} MiB, after={after_mib} MiB)" ) def is_mp4(data: bytes) -> bool: """Check if data represents a valid MP4 file by magic bytes.""" if len(data) < 8: return False return data[4:8] == b"ftyp" def is_jpeg(data: bytes) -> bool: # JPEG files start with: FF D8 FF return data.startswith(b"\xff\xd8\xff") def is_png(data): # PNG files start with: 89 50 4E 47 0D 0A 1A 0A return data.startswith(b"\x89PNG\r\n\x1a\n") def is_webp(data: bytes) -> bool: # WebP files start with: RIFF....WEBP return data[:4] == b"RIFF" and data[8:12] == b"WEBP" def detect_image_format(data: bytes) -> str: """Detect image format from bytes (magic). Returns 'png'|'jpeg'|'webp'; default 'png'.""" if len(data) < 12: return "png" if is_png(data): return "png" if is_jpeg(data): return "jpeg" if is_webp(data): return "webp" return "png" def get_expected_image_format( output_format: str | None = None, background: str | None = None, ) -> str: """Infer expected image format based on request parameters. Args: output_format: The output_format parameter from the request (png/jpeg/webp/jpg) background: The background parameter from the request (transparent/opaque/auto) Returns: Expected file extension: "jpg", "png", or "webp" """ fmt = (output_format or "").lower() if fmt in {"png", "webp", "jpeg", "jpg"}: return "jpg" if fmt == "jpeg" else fmt if (background or "auto").lower() == "transparent": return "png" return "jpg" # Default def wait_for_port(host="127.0.0.1", port=30010, deadline=300.0, interval=0.5): end = time.time() + deadline last_err = None while time.time() < end: if probe_port(host, port, timeout=interval): return True time.sleep(interval) raise TimeoutError(f"Port {host}:{port} not ready. Last error: {last_err}") def check_image_size(ut, image, width, height): # check image size ut.assertEqual(image.size, (width, height)) def get_perf_log_dir() -> Path: """Gets the performance log directory from the centralized sglang utility.""" log_dir_str = get_diffusion_perf_log_dir() if not log_dir_str: raise RuntimeError( "Performance logging is disabled (SGLANG_PERF_LOG_DIR is empty), " "but a test tried to access the log directory." ) return Path(log_dir_str) def _ensure_log_path(log_dir: Path) -> Path: log_dir.mkdir(parents=True, exist_ok=True) return log_dir / "performance.log" def clear_perf_log(log_dir: Path) -> Path: """Delete the perf log file so tests can watch for fresh entries.""" log_path = _ensure_log_path(log_dir) if log_path.exists(): log_path.unlink() logger.info("[server-test] Monitoring perf log at %s", log_path.as_posix()) return log_path def prepare_perf_log() -> tuple[Path, Path]: """Convenience helper to resolve and clear the perf log in one call.""" log_dir = get_perf_log_dir() log_path = clear_perf_log(log_dir) return log_dir, log_path def read_perf_logs(log_path: Path) -> list[RequestPerfRecord]: if not log_path.exists(): return [] records: list[RequestPerfRecord] = [] with log_path.open("r", encoding="utf-8") as fh: for line in fh: line = line.strip() if not line: continue try: record_dict = json.loads(line) records.append(RequestPerfRecord(**record_dict)) except json.JSONDecodeError: continue return records def wait_for_req_perf_record( request_id: str, log_path: Path, timeout: float = 30.0, ) -> RequestPerfRecord | None: """ the stage metrics of this request should be in the performance_log file with {request-id} """ logger.info(f"Waiting for req perf record with request id: {request_id}") deadline = time.time() + timeout while time.time() < deadline: records = read_perf_logs(log_path) for record in records: if record.request_id == request_id: return record time.sleep(0.5) if os.environ.get("SGLANG_GEN_BASELINE", "0") == "1": return None logger.error(f"record: {records}") raise AssertionError(f"Timeout waiting for stage metrics for request {request_id} ") def validate_image(b64_json: str) -> None: """Decode and validate that image is PNG or JPEG.""" image_bytes = base64.b64decode(b64_json) assert is_png(image_bytes) or is_jpeg(image_bytes), "Image must be PNG or JPEG" def validate_video(b64_json: str) -> None: """Decode and validate that video is a valid format.""" video_bytes = base64.b64decode(b64_json) is_webm = video_bytes[:4] == b"\x1a\x45\xdf\xa3" assert is_mp4(video_bytes) or is_webm, "Video must be MP4 or WebM" def validate_openai_video(video_bytes: bytes) -> None: """Validate that video is MP4 or WebM by magic bytes.""" is_webm = video_bytes.startswith(b"\x1a\x45\xdf\xa3") assert is_mp4(video_bytes) or is_webm, "Video must be MP4 or WebM" def validate_image_file( file_path: str, expected_filename: str, expected_width: int | None = None, expected_height: int | None = None, output_format: str | None = None, background: str | None = None, ) -> None: """Validate image output file: existence, extension, size, filename, format, dimensions.""" # Infer expected format from request parameters expected_ext = get_expected_image_format(output_format, background) # 1. File existence assert os.path.exists(file_path), f"Image file does not exist: {file_path}" # 2. Extension check assert file_path.endswith( f".{expected_ext}" ), f"Expected .{expected_ext} extension, got: {file_path}" # 3. File size > 0 file_size = os.path.getsize(file_path) assert file_size > 0, f"Image file is empty: {file_path}" # 4. Filename validation actual_filename = os.path.basename(file_path) assert ( actual_filename == expected_filename ), f"Filename mismatch: expected '{expected_filename}', got '{actual_filename}'" # 5. Image format validation (magic bytes check based on expected format) with open(file_path, "rb") as f: header = f.read(12) # Read enough bytes for webp detection if expected_ext == "png": assert is_png(header), f"File is not a valid PNG: {file_path}" elif expected_ext == "jpg": assert is_jpeg(header), f"File is not a valid JPEG: {file_path}" elif expected_ext == "webp": assert is_webp(header), f"File is not a valid WebP: {file_path}" # 6. Image dimension validation (reuse PIL) if expected_width is not None and expected_height is not None: with Image.open(file_path) as img: width, height = img.size assert ( width == expected_width ), f"Width mismatch: expected {expected_width}, got {width}" assert ( height == expected_height ), f"Height mismatch: expected {expected_height}, got {height}" def _get_video_dimensions_from_metadata( cap: cv2.VideoCapture, ) -> tuple[int, int] | None: """Get video dimensions from metadata properties. Args: cap: OpenCV VideoCapture object Returns: Tuple of (width, height) if successful, None if metadata is invalid """ width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) if width == 0 or height == 0: return None return int(width), int(height) def _get_video_dimensions_from_frame(cap: cv2.VideoCapture) -> tuple[int, int]: """Get video dimensions by reading the first frame. Args: cap: OpenCV VideoCapture object Returns: Tuple of (width, height) """ ret, frame = cap.read() if not ret or frame is None: raise ValueError("Unable to read video frame to get dimensions") # frame.shape is (height, width, channels) height, width = frame.shape[:2] return int(width), int(height) def get_video_dimensions(file_path: str) -> tuple[int, int]: """Get video dimensions (width, height) from a video file. Tries to get dimensions from metadata first, falls back to reading first frame. Returns: Tuple of (width, height) """ cap = cv2.VideoCapture(file_path) try: # Try to get dimensions from metadata first dimensions = _get_video_dimensions_from_metadata(cap) if dimensions is not None: return dimensions # Fall back to reading first frame return _get_video_dimensions_from_frame(cap) finally: cap.release() def get_video_frame_count(file_path: str) -> int: """Return the number of frames in a video file using OpenCV.""" cap = cv2.VideoCapture(file_path) try: count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) if count > 0: return count # Fallback: count frames manually n = 0 while cap.read()[0]: n += 1 return n finally: cap.release() def validate_video_file( file_path: str, expected_filename: str, expected_width: int | None = None, expected_height: int | None = None, ) -> None: """Validate video output file: existence, extension, size, filename, format, dimensions.""" # 1. File existence assert os.path.exists(file_path), f"Video file does not exist: {file_path}" # 2. Extension check assert file_path.endswith(".mp4"), f"Expected .mp4 extension, got: {file_path}" # 3. File size > 0 file_size = os.path.getsize(file_path) assert file_size > 0, f"Video file is empty: {file_path}" # 4. Filename validation actual_filename = os.path.basename(file_path) assert ( actual_filename == expected_filename ), f"Filename mismatch: expected '{expected_filename}', got '{actual_filename}'" # 5. Video format validation (reuse is_mp4) with open(file_path, "rb") as f: header = f.read(32) assert is_mp4(header), f"File is not a valid MP4: {file_path}" # 6. Video dimension validation (using OpenCV) if expected_width is not None and expected_height is not None: actual_width, actual_height = get_video_dimensions(file_path) assert ( actual_width == expected_width ), f"Video width mismatch: expected {expected_width}, got {actual_width}" assert ( actual_height == expected_height ), f"Video height mismatch: expected {expected_height}, got {actual_height}" def output_format_to_ext(output_format: str | None) -> str: """Map output_format to file extension. Used by GT naming and consistency check.""" if not output_format: return "png" of = output_format.lower() if of == "jpeg": return "jpg" if of in ("png", "webp", "jpg"): return of return "png" def _consistency_gt_filenames( case_id: str, num_gpus: int, is_video: bool, output_format: str | None = None ) -> list[str]: """Return the list of GT image filenames for a case. Reused by GT generation and consistency check.""" n = num_gpus if is_video: return [ f"{case_id}_{n}gpu_frame_0.png", f"{case_id}_{n}gpu_frame_mid.png", f"{case_id}_{n}gpu_frame_last.png", ] ext = output_format_to_ext(output_format) return [f"{case_id}_{n}gpu.{ext}"] def extract_key_frames_from_video( video_bytes: bytes, num_frames: int | None = None, ) -> list[np.ndarray]: """ Extract key frames (first, middle, last) from video bytes. Args: video_bytes: Raw video bytes (MP4 format) num_frames: Total number of frames (if known), used for validation Returns: List of numpy arrays [first_frame, middle_frame, last_frame]. """ with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: tmp.write(video_bytes) tmp_path = tmp.name try: cap = cv2.VideoCapture(tmp_path) if not cap.isOpened(): raise ValueError("Failed to open video file") total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) if total_frames < 1: raise ValueError("Video has no frames") first_idx = 0 mid_idx = total_frames // 2 last_idx = total_frames - 1 key_indices = [first_idx, mid_idx, last_idx] frames = [] for idx in key_indices: cap.set(cv2.CAP_PROP_POS_FRAMES, idx) ret, frame = cap.read() if not ret: raise ValueError(f"Failed to read frame at index {idx}") frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frames.append(frame_rgb) cap.release() logger.info( f"Extracted {len(frames)} key frames from video " f"(total: {total_frames}, indices: {key_indices})" ) return frames finally: os.unlink(tmp_path) def image_bytes_to_numpy(image_bytes: bytes) -> np.ndarray: """Convert image bytes to numpy array.""" img = Image.open(io.BytesIO(image_bytes)).convert("RGB") return np.array(img) ================================================ FILE: python/sglang/multimodal_gen/test/unit/test_lora_format_adapter.py ================================================ """ test_lora_format_adapter.py Small regression test for the LoRA format adapter. It downloads several public LoRA checkpoints from Hugging Face, runs format detection and normalization, and prints a compact summary table. """ import logging import os import tempfile from typing import Dict, List import torch from huggingface_hub import hf_hub_download from safetensors.torch import load_file from sglang.multimodal_gen.runtime.pipelines_core.lora_format_adapter import ( LoRAFormat, detect_lora_format_from_state_dict, normalize_lora_state_dict, ) logging.basicConfig(level=logging.INFO, force=True) logger = logging.getLogger("lora_test") ROOT_DIR = os.path.join(tempfile.gettempdir(), "sglang_lora_tests") os.makedirs(ROOT_DIR, exist_ok=True) def download_lora( repo_id: str, filename: str, local_name: str, ) -> str: """ Download a LoRA safetensors file into ROOT_DIR and return its local path. """ print(f"=== Downloading LoRA from {repo_id} ({filename}) ===") path = hf_hub_download( repo_id=repo_id, filename=filename, local_dir=ROOT_DIR, local_dir_use_symlinks=False, ) dst = os.path.join(ROOT_DIR, local_name) if os.path.abspath(path) != os.path.abspath(dst): try: import shutil shutil.copy2(path, dst) except Exception: dst = path print(f"Saved to: {dst}") return dst def is_diffusers_style_keys( sd: Dict[str, torch.Tensor], debug_name: str = "", ) -> bool: """ Relaxed structural check that a state_dict looks like diffusers-style LoRA. The check verifies: 1) No known non-diffusers prefixes. 2) No non-diffusers suffixes such as alpha / dora_scale / magnitude vectors. 3) Most top-level roots match common diffusers module namespaces. """ if not sd: print(f"[{debug_name}] diffusers-style check: EMPTY state_dict") return False keys: List[str] = list(sd.keys()) total = len(keys) banned_prefixes = ( "lora_unet_", "lora_te_", "lora_te1_", "lora_te2_", "lora_unet_double_blocks_", "lora_unet_single_blocks_", ) bad_prefix_keys = [k for k in keys if k.startswith(banned_prefixes)] cond1 = len(bad_prefix_keys) == 0 banned_suffixes = ( ".alpha", ".dora_scale", ".lora_magnitude_vector", ) bad_suffix_keys = [k for k in keys if k.endswith(banned_suffixes)] cond2 = len(bad_suffix_keys) == 0 allowed_roots = { "unet", "text_encoder", "text_encoder_2", "transformer", "prior", "image_encoder", "vae", "diffusion_model", } root_names = [k.split(".", 1)[0] for k in keys] root_ok_count = sum(r in allowed_roots for r in root_names) cond3 = root_ok_count >= 0.6 * total ok = cond1 and cond2 and cond3 if not ok: print(f"[{debug_name}] diffusers-style check FAILED (relaxed):") print(f" total keys = {total}") print( f" cond1(no banned prefixes) = {cond1}, bad_prefix_keys={len(bad_prefix_keys)}" ) if not cond1 and bad_prefix_keys: print(" example bad prefix key:", bad_prefix_keys[0]) print( f" cond2(no banned suffixes) = {cond2}, bad_suffix_keys={len(bad_suffix_keys)}" ) if not cond2 and bad_suffix_keys: print(" example bad suffix key:", bad_suffix_keys[0]) print(f" cond3(allowed roots>=60%) = {cond3}, root_ok_count={root_ok_count}") return ok def run_single_test( name: str, repo_id: str, filename: str, local_name: str, expected_before: LoRAFormat, expected_after: LoRAFormat = LoRAFormat.STANDARD, ): """ Run a single end-to-end test for one LoRA checkpoint. Steps: 1) Download. 2) Detect format on raw keys. 3) Normalize via lora_format_adapter. 4) Detect again on the normalized dict. 5) Optionally check for diffusers-style key structure. """ logger.info(f"=== Running test: {name} ===") local_path = download_lora(repo_id, filename, local_name) raw_state = load_file(local_path) detected_before = detect_lora_format_from_state_dict(raw_state) norm_state = normalize_lora_state_dict(raw_state, logger=logger) detected_after = detect_lora_format_from_state_dict(norm_state) standard_like = is_diffusers_style_keys(norm_state, debug_name=name) passed = detected_before == expected_before and detected_after == expected_after return { "name": name, "expected_before": expected_before.value, "detected_before": detected_before.value, "expected_after": expected_after.value, "detected_after": detected_after.value, "standard_like_keys": standard_like, "pass": passed, "num_keys_raw": len(raw_state), "num_keys_norm": len(norm_state), } def _run_all_tests() -> List[Dict]: results: List[Dict] = [] # SDXL LoRA that is already in diffusers/PEFT format. results.append( run_single_test( name="HF standard SDXL LoRA", repo_id="jbilcke-hf/sdxl-cinematic-1", filename="pytorch_lora_weights.safetensors", local_name="sdxl_cinematic1_pytorch_lora_weights.safetensors", expected_before=LoRAFormat.STANDARD, expected_after=LoRAFormat.STANDARD, ) ) # XLabs FLUX LoRA (non-diffusers → diffusers). results.append( run_single_test( name="XLabs FLUX Realism LoRA", repo_id="XLabs-AI/flux-RealismLora", filename="lora.safetensors", local_name="flux_realism_lora.safetensors", expected_before=LoRAFormat.XLABS_FLUX, expected_after=LoRAFormat.STANDARD, ) ) # Kohya-style FLUX LoRA (sd-scripts flux_lora.py → diffusers). results.append( run_single_test( name="Kohya-style Flux LoRA", repo_id="kohya-ss/misc-models", filename="flux-hasui-lora-d4-sigmoid-raw-gs1.0.safetensors", local_name="flux_hasui_lora_d4_sigmoid_raw_gs1_0.safetensors", expected_before=LoRAFormat.KOHYA_FLUX, expected_after=LoRAFormat.STANDARD, ) ) # AI-Toolkit Flux LoRA (non-diffusers → diffusers). results.append( run_single_test( name="AI-Toolkit Flux LoRA", repo_id="fal/flux-2-klein-4b-spritesheet-lora", filename="flux-spritesheet-lora.safetensors", local_name="flux_spritesheet_lora.safetensors", expected_before=LoRAFormat.AI_TOOLKIT_FLUX, expected_after=LoRAFormat.STANDARD, ) ) # Classic Kohya/A1111 SD LoRA (non-diffusers SD → diffusers). results.append( run_single_test( name="Kohya-style SD LoRA", repo_id="kohya-ss/misc-models", filename="fp-1f-chibi-1024.safetensors", local_name="fp_1f_chibi_1024.safetensors", expected_before=LoRAFormat.NON_DIFFUSERS_SD, expected_after=LoRAFormat.STANDARD, ) ) # Wan2.1 Fun Reward LoRA (ComfyUI format → diffusers). results.append( run_single_test( name="Wan2.1 Fun Reward LoRA (Comfy)", repo_id="alibaba-pai/Wan2.1-Fun-Reward-LoRAs", filename="Wan2.1-Fun-1.3B-InP-MPS.safetensors", local_name="wan21_fun_1_3b_inp_mps.safetensors", expected_before=LoRAFormat.NON_DIFFUSERS_SD, expected_after=LoRAFormat.STANDARD, ) ) # Qwen-Image EVA LoRA (already diffusers/PEFT-style). results.append( run_single_test( name="Qwen-Image EVA LoRA", repo_id="starsfriday/Qwen-Image-EVA-LoRA", filename="qwen_image_eva.safetensors", local_name="qwen_image_eva.safetensors", expected_before=LoRAFormat.STANDARD, expected_after=LoRAFormat.STANDARD, ) ) # Qwen-Image Lightning LoRA (non-diffusers Qwen → diffusers). results.append( run_single_test( name="Qwen-Image Lightning LoRA", repo_id="lightx2v/Qwen-Image-Lightning", filename="Qwen-Image-Lightning-4steps-V1.0-bf16.safetensors", local_name="qwen_image_lightning_4steps_v1_bf16.safetensors", expected_before=LoRAFormat.NON_DIFFUSERS_SD, expected_after=LoRAFormat.STANDARD, ) ) # Classic Painting Z-Image Turbo LoRA (Z-Image family). results.append( run_single_test( name="Classic Painting Z-Image LoRA", repo_id="renderartist/Classic-Painting-Z-Image-Turbo-LoRA", filename="Classic_Painting_Z_Image_Turbo_v1_renderartist_1750.safetensors", local_name="classic_painting_z_image_turbo_v1_renderartist_1750.safetensors", expected_before=LoRAFormat.STANDARD, expected_after=LoRAFormat.STANDARD, ) ) return results def _print_summary(results: List[Dict]) -> None: print("\n================ LoRA format adapter test ================") header = ( f"{'Test Name':30} " f"{'Exp(b)':12} " f"{'Act(b)':12} " f"{'Exp(a)':12} " f"{'Act(a)':12} " f"{'StdLike':8} " f"{'#Raw':7} " f"{'#Norm':7} " f"{'PASS':5}" ) print(header) print("-" * len(header)) for r in results: print( f"{r['name'][:30]:30} " f"{r['expected_before'][:12]:12} " f"{r['detected_before'][:12]:12} " f"{r['expected_after'][:12]:12} " f"{r['detected_after'][:12]:12} " f"{str(r['standard_like_keys']):8} " f"{r['num_keys_raw']:7d} " f"{r['num_keys_norm']:7d} " f"{str(r['pass']):5}" ) print("=========================================================\n") def main() -> None: results = _run_all_tests() _print_summary(results) if not all(r["pass"] for r in results): raise SystemExit(1) class TestLoRAFormatAdapter: def test_lora_format_adapter_all_formats(self): results = _run_all_tests() assert all( r["pass"] for r in results ), "At least one LoRA format adapter case failed" if __name__ == "__main__": main() ================================================ FILE: python/sglang/multimodal_gen/test/unit/test_sampling_params.py ================================================ import argparse import math import unittest from sglang.multimodal_gen.configs.sample.diffusers_generic import ( DiffusersGenericSamplingParams, ) from sglang.multimodal_gen.configs.sample.flux import FluxSamplingParams from sglang.multimodal_gen.configs.sample.qwenimage import QwenImageSamplingParams from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams class TestSamplingParamsValidate(unittest.TestCase): def test_prompt_path_suffix(self): with self.assertRaisesRegex(ValueError, r"prompt_path"): SamplingParams(prompt_path="bad.png") def test_num_outputs_per_prompt_must_be_positive(self): with self.assertRaisesRegex(ValueError, r"num_outputs_per_prompt"): SamplingParams(num_outputs_per_prompt=0) def test_fps_must_be_positive_int(self): with self.assertRaisesRegex(ValueError, r"\bfps\b"): SamplingParams(fps=0) with self.assertRaisesRegex(ValueError, r"\bfps\b"): SamplingParams(fps=None) # type: ignore[arg-type] def test_num_inference_steps_optional_but_if_set_must_be_positive(self): SamplingParams(num_inference_steps=None) with self.assertRaisesRegex(ValueError, r"num_inference_steps"): SamplingParams(num_inference_steps=-1) def test_guidance_scale_must_be_finite_non_negative_if_set(self): SamplingParams(guidance_scale=None) with self.assertRaisesRegex(ValueError, r"guidance_scale"): SamplingParams(guidance_scale=math.nan) with self.assertRaisesRegex(ValueError, r"guidance_scale"): SamplingParams(guidance_scale=-0.1) def test_guidance_rescale_must_be_finite_non_negative(self): with self.assertRaisesRegex(ValueError, r"guidance_rescale"): SamplingParams(guidance_rescale=-1.0) with self.assertRaisesRegex(ValueError, r"guidance_rescale"): SamplingParams(guidance_rescale=math.inf) def test_boundary_ratio_range(self): SamplingParams(boundary_ratio=None) with self.assertRaisesRegex(ValueError, r"boundary_ratio"): SamplingParams(boundary_ratio=1.5) with self.assertRaisesRegex(ValueError, r"boundary_ratio"): SamplingParams(boundary_ratio=math.nan) class TestSamplingParamsSubclass(unittest.TestCase): def test_flux_defaults_resolution_when_not_provided(self): params = FluxSamplingParams() self.assertEqual(params.height, 1024) self.assertEqual(params.width, 1024) def test_flux_preserves_user_resolution(self): params = FluxSamplingParams(height=640, width=768) self.assertEqual(params.height, 640) self.assertEqual(params.width, 768) def test_diffusers_generic_calls_base_post_init(self): with self.assertRaises(AssertionError): DiffusersGenericSamplingParams(num_frames=0) class TestSamplingParamsCliArgs(unittest.TestCase): def _parse_cli_kwargs(self, argv: list[str]) -> dict: parser = argparse.ArgumentParser() SamplingParams.add_cli_args(parser) args = parser.parse_args(argv) return SamplingParams.get_cli_args(args) def _make_qwen_image_params(self, argv: list[str]) -> QwenImageSamplingParams: return QwenImageSamplingParams(**self._parse_cli_kwargs(argv)) def test_get_cli_args_drops_unset_sampling_params(self): self.assertEqual(self._parse_cli_kwargs([]), {}) def test_get_cli_args_keeps_explicit_sampling_params(self): kwargs = self._parse_cli_kwargs( [ "--guidance-scale", str(SamplingParams.guidance_scale), "--negative-prompt", SamplingParams.negative_prompt, "--save-output", ] ) self.assertEqual(kwargs["guidance_scale"], SamplingParams.guidance_scale) self.assertEqual(kwargs["negative_prompt"], SamplingParams.negative_prompt) self.assertTrue(kwargs["save_output"]) def test_qwen_image_cli_path_preserves_model_defaults(self): params = self._make_qwen_image_params([]) self.assertEqual(params.negative_prompt, " ") self.assertEqual(params.guidance_scale, 4.0) def test_qwen_image_cli_path_allows_explicit_override_to_base_defaults(self): params = self._make_qwen_image_params( [ "--guidance-scale", str(SamplingParams.guidance_scale), "--negative-prompt", SamplingParams.negative_prompt, ] ) self.assertEqual(params.guidance_scale, SamplingParams.guidance_scale) self.assertEqual(params.negative_prompt, SamplingParams.negative_prompt) def test_merge_allows_explicit_field_matching_base_default(self): target = DiffusersGenericSamplingParams() user = SamplingParams(negative_prompt=SamplingParams.negative_prompt) target._merge_with_user_params(user, explicit_fields={"negative_prompt"}) self.assertEqual(target.negative_prompt, SamplingParams.negative_prompt) if __name__ == "__main__": unittest.main() ================================================ FILE: python/sglang/multimodal_gen/test/unit/test_server_args.py ================================================ import os import sys import unittest from unittest.mock import patch from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImagePipelineConfig, ) from sglang.multimodal_gen.registry import _get_config_info from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.utils import FlexibleArgumentParser class TestServerArgsPathExpansion(unittest.TestCase): def _from_dict_without_model_resolution(self, kwargs): with patch.object( PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig() ): return ServerArgs.from_dict(kwargs) def test_tilde_model_path_is_expanded(self): args = self._from_dict_without_model_resolution( {"model_path": "~/fake/local/model"} ) expected = os.path.expanduser("~/fake/local/model") self.assertEqual(args.model_path, expected) self.assertFalse(args.model_path.startswith("~")) def test_absolute_path_is_unchanged(self): args = self._from_dict_without_model_resolution( {"model_path": "/data/my-model"} ) self.assertEqual(args.model_path, "/data/my-model") def test_component_paths_are_expanded_before_pipeline_resolution(self): args = self._from_dict_without_model_resolution( { "model_path": "/data/my-model", "component_paths": {"vae": "~/fake/local/vae"}, } ) self.assertEqual( args.component_paths["vae"], os.path.expanduser("~/fake/local/vae") ) class TestModelIdResolution(unittest.TestCase): def setUp(self): _get_config_info.cache_clear() def test_model_id_overrides_arbitrary_local_path(self): # a local path whose directory name does not match any HF repo name; # --model-id tells the engine which config to use info = _get_config_info("/data/my-custom-qwen", model_id="Qwen-Image") self.assertIsNotNone(info) self.assertIs(info.pipeline_config_cls, QwenImagePipelineConfig) def test_model_id_works_after_tilde_expansion(self): # simulate the full flow: user passes ~/..., engine expands and resolves expanded = os.path.expanduser("~/.cache/huggingface/hub/bbb/snapshots/ccc") _get_config_info.cache_clear() info = _get_config_info(expanded, model_id="Qwen-Image") self.assertIsNotNone(info) def test_model_id_unknown_falls_back_without_crash(self): # unrecognized model_id: should warn and fall back to path-based detection # with an unresolvable path, expect RuntimeError from the detector step with self.assertRaises((RuntimeError, Exception)): _get_config_info("/data/no-such-model", model_id="NonExistentModelXYZ") class TestPipelineResolutionCliOverride(unittest.TestCase): def setUp(self): _get_config_info.cache_clear() def test_resolution_flag_overrides_qwen_image_layered_pipeline_config(self): parser = FlexibleArgumentParser() ServerArgs.add_cli_args(parser) argv = [ "--model-path", "Qwen/Qwen-Image-Layered", "--resolution", "768", ] with patch.object(sys, "argv", ["sglang"] + argv): args, unknown_args = parser.parse_known_args(argv) server_args = ServerArgs.from_cli_args(args, unknown_args) self.assertEqual(server_args.pipeline_config.resolution, 768) if __name__ == "__main__": unittest.main() ================================================ FILE: python/sglang/multimodal_gen/test/unit/test_storage.py ================================================ """ Test suite for S3 CloudStorage integration. Tests verify file upload, cleanup, URL generation, and error handling. """ import asyncio import importlib import os from types import SimpleNamespace import pytest import sglang.multimodal_gen.runtime.entrypoints.openai.storage as storage_mod from sglang.multimodal_gen.runtime.entrypoints.openai.storage import CloudStorage def _create_temp_file(tmp_path, name="test.png", content=b"\x89PNG\r\n\x1a\nfake"): """Create a temporary test file.""" p = tmp_path / name p.write_bytes(content) return str(p) # UNIT TESTS def test_upload_file_success(tmp_path): """Test successful upload with correct URL generation.""" file_path = _create_temp_file(tmp_path, "image.png") storage_mod.cloud_storage.enabled = True storage_mod.cloud_storage.bucket_name = "my-bucket" storage_mod.cloud_storage.endpoint_url = "https://s3.example.com" storage_mod.cloud_storage.region_name = None called = {} def fake_upload(local_path, bucket, key, ExtraArgs=None): called["local_path"] = local_path called["bucket"] = bucket called["key"] = key called["extra"] = ExtraArgs storage_mod.cloud_storage.client = SimpleNamespace(upload_file=fake_upload) url = asyncio.run(storage_mod.cloud_storage.upload_file(file_path, "image.png")) assert url == "https://s3.example.com/my-bucket/image.png" assert called["local_path"] == file_path assert called["bucket"] == "my-bucket" assert called["key"] == "image.png" assert called["extra"]["ContentType"] == "image/png" def test_upload_and_cleanup(tmp_path): """Test that local file is deleted after successful upload.""" file_path = _create_temp_file(tmp_path, "cleanup.png") storage_mod.cloud_storage.enabled = True storage_mod.cloud_storage.bucket_name = "my-bucket" storage_mod.cloud_storage.endpoint_url = "https://s3.example.com" storage_mod.cloud_storage.client = SimpleNamespace( upload_file=lambda *args, **kwargs: None ) assert os.path.exists(file_path) url = asyncio.run(storage_mod.cloud_storage.upload_and_cleanup(file_path)) assert url == "https://s3.example.com/my-bucket/cleanup.png" assert not os.path.exists(file_path) def test_upload_failure_preserves_file(tmp_path): """Test that file is preserved when upload fails.""" file_path = _create_temp_file(tmp_path, "preserve.png") storage_mod.cloud_storage.enabled = True storage_mod.cloud_storage.bucket_name = "my-bucket" storage_mod.cloud_storage.endpoint_url = "https://s3.example.com" def fake_upload_raises(*args, **kwargs): raise RuntimeError("simulated failure") storage_mod.cloud_storage.client = SimpleNamespace(upload_file=fake_upload_raises) result = asyncio.run(storage_mod.cloud_storage.upload_and_cleanup(file_path)) assert result is None assert os.path.exists(file_path) def test_disabled_storage_returns_none(tmp_path): """Test that disabled storage returns None.""" file_path = _create_temp_file(tmp_path, "test.png") prev_enabled = storage_mod.cloud_storage.enabled storage_mod.cloud_storage.enabled = False try: result = asyncio.run( storage_mod.cloud_storage.upload_file(file_path, "test.png") ) assert result is None finally: storage_mod.cloud_storage.enabled = prev_enabled def test_aws_url_with_region(tmp_path): """Test AWS S3 URL generation with specific region.""" file_path = _create_temp_file(tmp_path, "aws.png") storage_mod.cloud_storage.enabled = True storage_mod.cloud_storage.bucket_name = "aws-bucket" storage_mod.cloud_storage.endpoint_url = None storage_mod.cloud_storage.region_name = "us-west-2" storage_mod.cloud_storage.client = SimpleNamespace( upload_file=lambda *args, **kwargs: None ) url = asyncio.run(storage_mod.cloud_storage.upload_file(file_path, "aws.png")) assert url == "https://aws-bucket.s3.us-west-2.amazonaws.com/aws.png" def test_aws_url_default_region(tmp_path): """Test AWS S3 URL defaults to us-east-1 when region not specified.""" file_path = _create_temp_file(tmp_path, "default.png") storage_mod.cloud_storage.enabled = True storage_mod.cloud_storage.bucket_name = "default-bucket" storage_mod.cloud_storage.endpoint_url = None storage_mod.cloud_storage.region_name = None storage_mod.cloud_storage.client = SimpleNamespace( upload_file=lambda *args, **kwargs: None ) url = asyncio.run(storage_mod.cloud_storage.upload_file(file_path, "default.png")) assert url == "https://default-bucket.s3.us-east-1.amazonaws.com/default.png" def test_custom_endpoint_url(tmp_path): """Test URL generation with custom endpoint (MinIO/OSS/COS).""" file_path = _create_temp_file(tmp_path, "custom.png") storage_mod.cloud_storage.enabled = True storage_mod.cloud_storage.bucket_name = "custom-bucket" storage_mod.cloud_storage.endpoint_url = "https://minio.example.com/" storage_mod.cloud_storage.region_name = None storage_mod.cloud_storage.client = SimpleNamespace( upload_file=lambda *args, **kwargs: None ) url = asyncio.run(storage_mod.cloud_storage.upload_file(file_path, "custom.png")) # Verify trailing slash is stripped assert url == "https://minio.example.com/custom-bucket/custom.png" def test_content_type_detection(tmp_path): """Test Content-Type header for different file extensions.""" storage_mod.cloud_storage.enabled = True storage_mod.cloud_storage.bucket_name = "test-bucket" storage_mod.cloud_storage.endpoint_url = "https://s3.test" test_cases = [ ("image.png", "image/png"), ("image.jpg", "image/jpeg"), ("image.jpeg", "image/jpeg"), ("image.webp", "image/webp"), ("video.mp4", "video/mp4"), ("file.bin", "application/octet-stream"), ] for filename, expected_type in test_cases: called = {} def fake_upload(local_path, bucket, key, ExtraArgs=None): called["content_type"] = ExtraArgs.get("ContentType") storage_mod.cloud_storage.client = SimpleNamespace(upload_file=fake_upload) file_path = _create_temp_file(tmp_path, filename) asyncio.run(storage_mod.cloud_storage.upload_file(file_path, filename)) assert called["content_type"] == expected_type # requires moto and boto3 has_moto = ( importlib.util.find_spec("moto") is not None and importlib.util.find_spec("boto3") is not None ) @pytest.mark.skipif(not has_moto, reason="moto/boto3 not installed") def test_integration_with_moto(tmp_path): """Integration test using moto to mock real S3 service.""" import boto3 from moto import mock_aws os.environ["SGLANG_CLOUD_STORAGE_TYPE"] = "s3" os.environ["SGLANG_S3_BUCKET_NAME"] = "integration-test" os.environ["SGLANG_S3_REGION_NAME"] = "us-east-1" with mock_aws(): s3 = boto3.client("s3", region_name="us-east-1") s3.create_bucket(Bucket="integration-test") storage = CloudStorage() assert storage.is_enabled() file_path = _create_temp_file(tmp_path, "integration.png", b"test_data") url = asyncio.run(storage.upload_and_cleanup(file_path)) assert url is not None assert "integration-test" in url assert "integration.png" in url assert not os.path.exists(file_path) obj = s3.get_object(Bucket="integration-test", Key="integration.png") assert obj["Body"].read() == b"test_data" for key in [ "SGLANG_CLOUD_STORAGE_TYPE", "SGLANG_S3_BUCKET_NAME", "SGLANG_S3_REGION_NAME", ]: os.environ.pop(key, None) ================================================ FILE: python/sglang/multimodal_gen/third_party/__init__.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo ================================================ FILE: python/sglang/multimodal_gen/third_party/pynvml.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # copied from https://pypi.org/project/nvidia-ml-py # version 12.570.86 ##### # Copyright (c) 2011-2023, NVIDIA Corporation. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are met: # # * Redistributions of source code must retain the above copyright notice, # this list of conditions and the following disclaimer. # * Redistributions in binary form must reproduce the above copyright # notice, this list of conditions and the following disclaimer in the # documentation and/or other materials provided with the distribution. # * Neither the name of the NVIDIA Corporation nor the names of its # contributors may be used to endorse or promote products derived from # this software without specific prior written permission. # # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF # THE POSSIBILITY OF SUCH DAMAGE. ##### import os import string import sys import threading ## # Python bindings for the NVML library ## from ctypes import * from functools import wraps ## C Type mappings ## ## Enums _nvmlEnableState_t = c_uint NVML_FEATURE_DISABLED = 0 NVML_FEATURE_ENABLED = 1 _nvmlBrandType_t = c_uint NVML_BRAND_UNKNOWN = 0 NVML_BRAND_QUADRO = 1 NVML_BRAND_TESLA = 2 NVML_BRAND_NVS = 3 NVML_BRAND_GRID = ( 4 # Deprecated from API reporting. Keeping definition for backward compatibility. ) NVML_BRAND_GEFORCE = 5 NVML_BRAND_TITAN = 6 NVML_BRAND_NVIDIA_VAPPS = 7 # NVIDIA Virtual Applications NVML_BRAND_NVIDIA_VPC = 8 # NVIDIA Virtual PC NVML_BRAND_NVIDIA_VCS = 9 # NVIDIA Virtual Compute Server NVML_BRAND_NVIDIA_VWS = 10 # NVIDIA RTX Virtual Workstation NVML_BRAND_NVIDIA_CLOUD_GAMING = 11 # NVIDIA Cloud Gaming NVML_BRAND_NVIDIA_VGAMING = NVML_BRAND_NVIDIA_CLOUD_GAMING # Deprecated from API reporting. Keeping definition for backward compatibility. NVML_BRAND_QUADRO_RTX = 12 NVML_BRAND_NVIDIA_RTX = 13 NVML_BRAND_NVIDIA = 14 NVML_BRAND_GEFORCE_RTX = 15 # Unused NVML_BRAND_TITAN_RTX = 16 # Unused NVML_BRAND_COUNT = 17 _nvmlTemperatureThresholds_t = c_uint NVML_TEMPERATURE_THRESHOLD_SHUTDOWN = 0 NVML_TEMPERATURE_THRESHOLD_SLOWDOWN = 1 NVML_TEMPERATURE_THRESHOLD_MEM_MAX = 2 NVML_TEMPERATURE_THRESHOLD_GPU_MAX = 3 NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MIN = 4 NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_CURR = 5 NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MAX = 6 NVML_TEMPERATURE_THRESHOLD_GPS_CURR = 7 NVML_TEMPERATURE_THRESHOLD_COUNT = 8 _nvmlTemperatureSensors_t = c_uint NVML_TEMPERATURE_GPU = 0 NVML_TEMPERATURE_COUNT = 1 _nvmlComputeMode_t = c_uint NVML_COMPUTEMODE_DEFAULT = 0 NVML_COMPUTEMODE_EXCLUSIVE_THREAD = 1 ## Support Removed NVML_COMPUTEMODE_PROHIBITED = 2 NVML_COMPUTEMODE_EXCLUSIVE_PROCESS = 3 NVML_COMPUTEMODE_COUNT = 4 _nvmlMemoryLocation_t = c_uint NVML_MEMORY_LOCATION_L1_CACHE = 0 NVML_MEMORY_LOCATION_L2_CACHE = 1 NVML_MEMORY_LOCATION_DEVICE_MEMORY = 2 NVML_MEMORY_LOCATION_DRAM = 2 NVML_MEMORY_LOCATION_REGISTER_FILE = 3 NVML_MEMORY_LOCATION_TEXTURE_MEMORY = 4 NVML_MEMORY_LOCATION_TEXTURE_SHM = 5 NVML_MEMORY_LOCATION_CBU = 6 NVML_MEMORY_LOCATION_SRAM = 7 NVML_MEMORY_LOCATION_COUNT = 8 NVML_NVLINK_MAX_LINKS = 18 # For backwards compatibility, maintain the incorrectly-named "LANES" define NVML_NVLINK_MAX_LANES = NVML_NVLINK_MAX_LINKS _nvmlNvLinkErrorCounter_t = c_uint NVML_NVLINK_ERROR_DL_REPLAY = 0 NVML_NVLINK_ERROR_DL_RECOVERY = 1 NVML_NVLINK_ERROR_DL_CRC_FLIT = 2 NVML_NVLINK_ERROR_DL_CRC_DATA = 3 NVML_NVLINK_ERROR_DL_ECC_DATA = 4 NVML_NVLINK_ERROR_COUNT = 5 _nvmlNvLinkEccLaneErrorCounter_t = c_uint NVML_NVLINK_ERROR_DL_ECC_LANE0 = 0 NVML_NVLINK_ERROR_DL_ECC_LANE1 = 1 NVML_NVLINK_ERROR_DL_ECC_LANE2 = 2 NVML_NVLINK_ERROR_DL_ECC_LANE3 = 3 NVML_NVLINK_ERROR_DL_ECC_COUNT = 5 _nvmlNvLinkCapability_t = c_uint NVML_NVLINK_CAP_P2P_SUPPORTED = 0 NVML_NVLINK_CAP_SYSMEM_ACCESS = 1 NVML_NVLINK_CAP_P2P_ATOMICS = 2 NVML_NVLINK_CAP_SYSMEM_ATOMICS = 3 NVML_NVLINK_CAP_SLI_BRIDGE = 4 NVML_NVLINK_CAP_VALID = 5 NVML_NVLINK_CAP_COUNT = 6 _nvmlNvLinkUtilizationCountPktTypes_t = c_uint NVML_NVLINK_COUNTER_PKTFILTER_NOP = 0x1 NVML_NVLINK_COUNTER_PKTFILTER_READ = 0x2 NVML_NVLINK_COUNTER_PKTFILTER_WRITE = 0x4 NVML_NVLINK_COUNTER_PKTFILTER_RATOM = 0x8 NVML_NVLINK_COUNTER_PKTFILTER_NRATOM = 0x10 NVML_NVLINK_COUNTER_PKTFILTER_FLUSH = 0x20 NVML_NVLINK_COUNTER_PKTFILTER_RESPDATA = 0x40 NVML_NVLINK_COUNTER_PKTFILTER_RESPNODATA = 0x80 NVML_NVLINK_COUNTER_PKTFILTER_ALL = 0xFF _nvmlNvLinkUtilizationCountUnits_t = c_uint NVML_NVLINK_COUNTER_UNIT_CYCLES = 0 NVML_NVLINK_COUNTER_UNIT_PACKETS = 1 NVML_NVLINK_COUNTER_UNIT_BYTES = 2 NVML_NVLINK_COUNTER_UNIT_RESERVED = 3 NVML_NVLINK_COUNTER_UNIT_COUNT = 4 _nvmlNvLinkDeviceType_t = c_uint NVML_NVLINK_DEVICE_TYPE_GPU = 0x00 NVML_NVLINK_DEVICE_TYPE_IBMNPU = 0x01 NVML_NVLINK_DEVICE_TYPE_SWITCH = 0x02 NVML_NVLINK_DEVICE_TYPE_UNKNOWN = 0xFF # These are deprecated, instead use _nvmlMemoryErrorType_t _nvmlEccBitType_t = c_uint NVML_SINGLE_BIT_ECC = 0 NVML_DOUBLE_BIT_ECC = 1 NVML_ECC_ERROR_TYPE_COUNT = 2 _nvmlEccCounterType_t = c_uint NVML_VOLATILE_ECC = 0 NVML_AGGREGATE_ECC = 1 NVML_ECC_COUNTER_TYPE_COUNT = 2 _nvmlMemoryErrorType_t = c_uint NVML_MEMORY_ERROR_TYPE_CORRECTED = 0 NVML_MEMORY_ERROR_TYPE_UNCORRECTED = 1 NVML_MEMORY_ERROR_TYPE_COUNT = 2 _nvmlClockType_t = c_uint NVML_CLOCK_GRAPHICS = 0 NVML_CLOCK_SM = 1 NVML_CLOCK_MEM = 2 NVML_CLOCK_VIDEO = 3 NVML_CLOCK_COUNT = 4 _nvmlClockId_t = c_uint NVML_CLOCK_ID_CURRENT = 0 NVML_CLOCK_ID_APP_CLOCK_TARGET = 1 NVML_CLOCK_ID_APP_CLOCK_DEFAULT = 2 NVML_CLOCK_ID_CUSTOMER_BOOST_MAX = 3 NVML_CLOCK_ID_COUNT = 4 _nvmlDriverModel_t = c_uint NVML_DRIVER_WDDM = 0 NVML_DRIVER_WDM = 1 NVML_DRIVER_MCDM = 2 NVML_MAX_GPU_PERF_PSTATES = 16 _nvmlPstates_t = c_uint NVML_PSTATE_0 = 0 NVML_PSTATE_1 = 1 NVML_PSTATE_2 = 2 NVML_PSTATE_3 = 3 NVML_PSTATE_4 = 4 NVML_PSTATE_5 = 5 NVML_PSTATE_6 = 6 NVML_PSTATE_7 = 7 NVML_PSTATE_8 = 8 NVML_PSTATE_9 = 9 NVML_PSTATE_10 = 10 NVML_PSTATE_11 = 11 NVML_PSTATE_12 = 12 NVML_PSTATE_13 = 13 NVML_PSTATE_14 = 14 NVML_PSTATE_15 = 15 NVML_PSTATE_UNKNOWN = 32 _nvmlInforomObject_t = c_uint NVML_INFOROM_OEM = 0 NVML_INFOROM_ECC = 1 NVML_INFOROM_POWER = 2 NVML_INFOROM_DEN = 3 NVML_INFOROM_COUNT = 4 _nvmlReturn_t = c_uint NVML_SUCCESS = 0 NVML_ERROR_UNINITIALIZED = 1 NVML_ERROR_INVALID_ARGUMENT = 2 NVML_ERROR_NOT_SUPPORTED = 3 NVML_ERROR_NO_PERMISSION = 4 NVML_ERROR_ALREADY_INITIALIZED = 5 NVML_ERROR_NOT_FOUND = 6 NVML_ERROR_INSUFFICIENT_SIZE = 7 NVML_ERROR_INSUFFICIENT_POWER = 8 NVML_ERROR_DRIVER_NOT_LOADED = 9 NVML_ERROR_TIMEOUT = 10 NVML_ERROR_IRQ_ISSUE = 11 NVML_ERROR_LIBRARY_NOT_FOUND = 12 NVML_ERROR_FUNCTION_NOT_FOUND = 13 NVML_ERROR_CORRUPTED_INFOROM = 14 NVML_ERROR_GPU_IS_LOST = 15 NVML_ERROR_RESET_REQUIRED = 16 NVML_ERROR_OPERATING_SYSTEM = 17 NVML_ERROR_LIB_RM_VERSION_MISMATCH = 18 NVML_ERROR_IN_USE = 19 NVML_ERROR_MEMORY = 20 NVML_ERROR_NO_DATA = 21 NVML_ERROR_VGPU_ECC_NOT_SUPPORTED = 22 NVML_ERROR_INSUFFICIENT_RESOURCES = 23 NVML_ERROR_FREQ_NOT_SUPPORTED = 24 NVML_ERROR_ARGUMENT_VERSION_MISMATCH = 25 NVML_ERROR_DEPRECATED = 26 NVML_ERROR_NOT_READY = 27 NVML_ERROR_GPU_NOT_FOUND = 28 NVML_ERROR_INVALID_STATE = 29 NVML_ERROR_UNKNOWN = 999 _nvmlFanState_t = c_uint NVML_FAN_NORMAL = 0 NVML_FAN_FAILED = 1 _nvmlFanControlPolicy_t = c_uint NVML_FAN_POLICY_TEMPERATURE_CONTINOUS_SW = 0 NVML_FAN_POLICY_MANUAL = 1 _nvmlLedColor_t = c_uint NVML_LED_COLOR_GREEN = 0 NVML_LED_COLOR_AMBER = 1 _nvmlGpuOperationMode_t = c_uint NVML_GOM_ALL_ON = 0 NVML_GOM_COMPUTE = 1 NVML_GOM_LOW_DP = 2 _nvmlPageRetirementCause_t = c_uint NVML_PAGE_RETIREMENT_CAUSE_MULTIPLE_SINGLE_BIT_ECC_ERRORS = 0 NVML_PAGE_RETIREMENT_CAUSE_DOUBLE_BIT_ECC_ERROR = 1 NVML_PAGE_RETIREMENT_CAUSE_COUNT = 2 _nvmlRestrictedAPI_t = c_uint NVML_RESTRICTED_API_SET_APPLICATION_CLOCKS = 0 NVML_RESTRICTED_API_SET_AUTO_BOOSTED_CLOCKS = 1 NVML_RESTRICTED_API_COUNT = 2 _nvmlBridgeChipType_t = c_uint NVML_BRIDGE_CHIP_PLX = 0 NVML_BRIDGE_CHIP_BRO4 = 1 NVML_MAX_PHYSICAL_BRIDGE = 128 _nvmlValueType_t = c_uint NVML_VALUE_TYPE_DOUBLE = 0 NVML_VALUE_TYPE_UNSIGNED_INT = 1 NVML_VALUE_TYPE_UNSIGNED_LONG = 2 NVML_VALUE_TYPE_UNSIGNED_LONG_LONG = 3 NVML_VALUE_TYPE_SIGNED_LONG_LONG = 4 NVML_VALUE_TYPE_SIGNED_INT = 5 NVML_VALUE_TYPE_UNSIGNED_SHORT = 6 NVML_VALUE_TYPE_COUNT = 7 _nvmlNvlinkVersion_t = c_uint NVML_NVLINK_VERSION_INVALID = 0 NVML_NVLINK_VERSION_1_0 = 1 NVML_NVLINK_VERSION_2_0 = 2 NVML_NVLINK_VERSION_2_2 = 3 NVML_NVLINK_VERSION_3_0 = 4 NVML_NVLINK_VERSION_3_1 = 5 NVML_NVLINK_VERSION_4_0 = 6 NVML_NVLINK_VERSION_5_0 = 7 _nvmlPerfPolicyType_t = c_uint NVML_PERF_POLICY_POWER = 0 NVML_PERF_POLICY_THERMAL = 1 NVML_PERF_POLICY_SYNC_BOOST = 2 NVML_PERF_POLICY_BOARD_LIMIT = 3 NVML_PERF_POLICY_LOW_UTILIZATION = 4 NVML_PERF_POLICY_RELIABILITY = 5 NVML_PERF_POLICY_TOTAL_APP_CLOCKS = 10 NVML_PERF_POLICY_TOTAL_BASE_CLOCKS = 11 NVML_PERF_POLICY_COUNT = 12 _nvmlEncoderQueryType_t = c_uint NVML_ENCODER_QUERY_H264 = 0 NVML_ENCODER_QUERY_HEVC = 1 NVML_ENCODER_QUERY_AV1 = 2 NVML_ENCODER_QUERY_UNKNOWN = 255 _nvmlFBCSessionType_t = c_uint NVML_FBC_SESSION_TYPE_UNKNOWN = 0 NVML_FBC_SESSION_TYPE_TOSYS = 1 NVML_FBC_SESSION_TYPE_CUDA = 2 NVML_FBC_SESSION_TYPE_VID = 3 NVML_FBC_SESSION_TYPE_HWENC = 4 _nvmlDetachGpuState_t = c_uint NVML_DETACH_GPU_KEEP = 0 NVML_DETACH_GPU_REMOVE = 1 _nvmlPcieLinkState_t = c_uint NVML_PCIE_LINK_KEEP = 0 NVML_PCIE_LINK_SHUT_DOWN = 1 _nvmlSamplingType_t = c_uint NVML_TOTAL_POWER_SAMPLES = 0 NVML_GPU_UTILIZATION_SAMPLES = 1 NVML_MEMORY_UTILIZATION_SAMPLES = 2 NVML_ENC_UTILIZATION_SAMPLES = 3 NVML_DEC_UTILIZATION_SAMPLES = 4 NVML_PROCESSOR_CLK_SAMPLES = 5 NVML_MEMORY_CLK_SAMPLES = 6 NVML_MODULE_POWER_SAMPLES = 7 NVML_JPG_UTILIZATION_SAMPLES = 8 NVML_OFA_UTILIZATION_SAMPLES = 9 NVML_SAMPLINGTYPE_COUNT = 10 _nvmlPcieUtilCounter_t = c_uint NVML_PCIE_UTIL_TX_BYTES = 0 NVML_PCIE_UTIL_RX_BYTES = 1 NVML_PCIE_UTIL_COUNT = 2 _nvmlGpuTopologyLevel_t = c_uint NVML_TOPOLOGY_INTERNAL = 0 NVML_TOPOLOGY_SINGLE = 10 NVML_TOPOLOGY_MULTIPLE = 20 NVML_TOPOLOGY_HOSTBRIDGE = 30 NVML_TOPOLOGY_NODE = 40 NVML_TOPOLOGY_CPU = NVML_TOPOLOGY_NODE NVML_TOPOLOGY_SYSTEM = 50 _nvmlGpuP2PCapsIndex_t = c_uint NVML_P2P_CAPS_INDEX_READ = (0,) NVML_P2P_CAPS_INDEX_WRITE = 1 NVML_P2P_CAPS_INDEX_NVLINK = 2 NVML_P2P_CAPS_INDEX_ATOMICS = 3 # # NVML_P2P_CAPS_INDEX_PROP is deprecated. # Use NVML_P2P_CAPS_INDEX_PCI instead. # NVML_P2P_CAPS_INDEX_PROP = 4 NVML_P2P_CAPS_INDEX_PCI = 4 NVML_P2P_CAPS_INDEX_UNKNOWN = 5 _nvmlGpuP2PStatus_t = c_uint NVML_P2P_STATUS_OK = 0 NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED = 1 NVML_P2P_STATUS_CHIPSET_NOT_SUPPORTED = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED NVML_P2P_STATUS_GPU_NOT_SUPPORTED = 2 NVML_P2P_STATUS_IOH_TOPOLOGY_NOT_SUPPORTED = 3 NVML_P2P_STATUS_DISABLED_BY_REGKEY = 4 NVML_P2P_STATUS_NOT_SUPPORTED = 5 NVML_P2P_STATUS_UNKNOWN = 6 _nvmlDeviceArchitecture_t = c_uint NVML_DEVICE_ARCH_KEPLER = 2 NVML_DEVICE_ARCH_MAXWELL = 3 NVML_DEVICE_ARCH_PASCAL = 4 NVML_DEVICE_ARCH_VOLTA = 5 NVML_DEVICE_ARCH_TURING = 6 NVML_DEVICE_ARCH_AMPERE = 7 NVML_DEVICE_ARCH_ADA = 8 NVML_DEVICE_ARCH_HOPPER = 9 NVML_DEVICE_ARCH_BLACKWELL = 10 NVML_DEVICE_ARCH_T23X = 11 NVML_DEVICE_ARCH_UNKNOWN = 0xFFFFFFFF # PCI bus Types _nvmlBusType_t = c_uint NVML_BUS_TYPE_UNKNOWN = 0 NVML_BUS_TYPE_PCI = 1 NVML_BUS_TYPE_PCIE = 2 NVML_BUS_TYPE_FPCI = 3 NVML_BUS_TYPE_AGP = 4 _nvmlPowerSource_t = c_uint NVML_POWER_SOURCE_AC = 0x00000000 NVML_POWER_SOURCE_BATTERY = 0x00000001 NVML_POWER_SOURCE_UNDERSIZED = 0x00000002 _nvmlAdaptiveClockInfoStatus_t = c_uint NVML_ADAPTIVE_CLOCKING_INFO_STATUS_DISABLED = 0x00000000 NVML_ADAPTIVE_CLOCKING_INFO_STATUS_ENABLED = 0x00000001 _nvmlClockLimitId_t = c_uint NVML_CLOCK_LIMIT_ID_RANGE_START = 0xFFFFFF00 NVML_CLOCK_LIMIT_ID_TDP = 0xFFFFFF01 NVML_CLOCK_LIMIT_ID_UNLIMITED = 0xFFFFFF02 _nvmlPcieLinkMaxSpeed_t = c_uint NVML_PCIE_LINK_MAX_SPEED_INVALID = 0x00000000 NVML_PCIE_LINK_MAX_SPEED_2500MBPS = 0x00000001 NVML_PCIE_LINK_MAX_SPEED_5000MBPS = 0x00000002 NVML_PCIE_LINK_MAX_SPEED_8000MBPS = 0x00000003 NVML_PCIE_LINK_MAX_SPEED_16000MBPS = 0x00000004 NVML_PCIE_LINK_MAX_SPEED_32000MBPS = 0x00000005 NVML_PCIE_LINK_MAX_SPEED_64000MBPS = 0x00000006 _nvmlPcieAtomicsCapability_t = c_uint NVML_PCIE_ATOMICS_CAP_FETCHADD32 = 0x01 NVML_PCIE_ATOMICS_CAP_FETCHADD64 = 0x02 NVML_PCIE_ATOMICS_CAP_SWAP32 = 0x04 NVML_PCIE_ATOMICS_CAP_SWAP64 = 0x08 NVML_PCIE_ATOMICS_CAP_CAS32 = 0x10 NVML_PCIE_ATOMICS_CAP_CAS64 = 0x20 NVML_PCIE_ATOMICS_CAP_CAS128 = 0x40 NVML_PCIE_ATOMICS_OPS_MAX = 7 _nvmlAffinityScope_t = c_uint NVML_AFFINITY_SCOPE_NODE = 0 NVML_AFFINITY_SCOPE_SOCKET = 1 _nvmlDeviceGpuRecoveryAction_t = c_uint NVML_GPU_RECOVERY_ACTION_NONE = 0 NVML_GPU_RECOVERY_ACTION_GPU_RESET = 1 NVML_GPU_RECOVERY_ACTION_NODE_REBOOT = 2 NVML_GPU_RECOVERY_ACTION_DRAIN_P2P = 3 NVML_GPU_RECOVERY_ACTION_DRAIN_AND_RESET = 4 # C preprocessor defined values nvmlFlagDefault = 0 nvmlFlagForce = 1 NVML_INIT_FLAG_NO_GPUS = 1 NVML_INIT_FLAG_NO_ATTACH = 2 NVML_MAX_GPC_COUNT = 32 # buffer size NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE = 16 NVML_DEVICE_UUID_BUFFER_SIZE = 80 NVML_DEVICE_UUID_V2_BUFFER_SIZE = 96 NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE = 80 NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE = 80 NVML_DEVICE_NAME_BUFFER_SIZE = 64 NVML_DEVICE_NAME_V2_BUFFER_SIZE = 96 NVML_DEVICE_SERIAL_BUFFER_SIZE = 30 NVML_DEVICE_PART_NUMBER_BUFFER_SIZE = 80 NVML_DEVICE_GPU_PART_NUMBER_BUFFER_SIZE = 80 NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE = 32 NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE = 32 NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE = 16 NVML_GRID_LICENSE_BUFFER_SIZE = 128 NVML_VGPU_NAME_BUFFER_SIZE = 64 NVML_GRID_LICENSE_FEATURE_MAX_COUNT = 3 NVML_VGPU_METADATA_OPAQUE_DATA_SIZE = sizeof(c_uint) + 256 NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE = 256 NVML_DEVICE_GPU_FRU_PART_NUMBER_BUFFER_SIZE = ( 0x14 # NV2080_GPU_MAX_PRODUCT_PART_NUMBER_LENGTH ) NVML_PERF_MODES_BUFFER_SIZE = 2048 # Format strings NVML_DEVICE_PCI_BUS_ID_LEGACY_FMT = "%04X:%02X:%02X.0" NVML_DEVICE_PCI_BUS_ID_FMT = "%08X:%02X:%02X.0" NVML_VALUE_NOT_AVAILABLE_ulonglong = c_ulonglong(-1) NVML_VALUE_NOT_AVAILABLE_uint = c_uint(-1) """ Field Identifiers. All Identifiers pertain to a device. Each ID is only used once and is guaranteed never to change. """ NVML_FI_DEV_ECC_CURRENT = 1 # Current ECC mode. 1=Active. 0=Inactive NVML_FI_DEV_ECC_PENDING = 2 # Pending ECC mode. 1=Active. 0=Inactive # ECC Count Totals NVML_FI_DEV_ECC_SBE_VOL_TOTAL = 3 # Total single bit volatile ECC errors NVML_FI_DEV_ECC_DBE_VOL_TOTAL = 4 # Total double bit volatile ECC errors NVML_FI_DEV_ECC_SBE_AGG_TOTAL = 5 # Total single bit aggregate (persistent) ECC errors NVML_FI_DEV_ECC_DBE_AGG_TOTAL = 6 # Total double bit aggregate (persistent) ECC errors # Individual ECC locations NVML_FI_DEV_ECC_SBE_VOL_L1 = 7 # L1 cache single bit volatile ECC errors NVML_FI_DEV_ECC_DBE_VOL_L1 = 8 # L1 cache double bit volatile ECC errors NVML_FI_DEV_ECC_SBE_VOL_L2 = 9 # L2 cache single bit volatile ECC errors NVML_FI_DEV_ECC_DBE_VOL_L2 = 10 # L2 cache double bit volatile ECC errors NVML_FI_DEV_ECC_SBE_VOL_DEV = 11 # Device memory single bit volatile ECC errors NVML_FI_DEV_ECC_DBE_VOL_DEV = 12 # Device memory double bit volatile ECC errors NVML_FI_DEV_ECC_SBE_VOL_REG = 13 # Register file single bit volatile ECC errors NVML_FI_DEV_ECC_DBE_VOL_REG = 14 # Register file double bit volatile ECC errors NVML_FI_DEV_ECC_SBE_VOL_TEX = 15 # Texture memory single bit volatile ECC errors NVML_FI_DEV_ECC_DBE_VOL_TEX = 16 # Texture memory double bit volatile ECC errors NVML_FI_DEV_ECC_DBE_VOL_CBU = 17 # CBU double bit volatile ECC errors NVML_FI_DEV_ECC_SBE_AGG_L1 = 18 # L1 cache single bit aggregate (persistent) ECC errors NVML_FI_DEV_ECC_DBE_AGG_L1 = 19 # L1 cache double bit aggregate (persistent) ECC errors NVML_FI_DEV_ECC_SBE_AGG_L2 = 20 # L2 cache single bit aggregate (persistent) ECC errors NVML_FI_DEV_ECC_DBE_AGG_L2 = 21 # L2 cache double bit aggregate (persistent) ECC errors NVML_FI_DEV_ECC_SBE_AGG_DEV = ( 22 # Device memory single bit aggregate (persistent) ECC errors ) NVML_FI_DEV_ECC_DBE_AGG_DEV = ( 23 # Device memory double bit aggregate (persistent) ECC errors ) NVML_FI_DEV_ECC_SBE_AGG_REG = ( 24 # Register File single bit aggregate (persistent) ECC errors ) NVML_FI_DEV_ECC_DBE_AGG_REG = ( 25 # Register File double bit aggregate (persistent) ECC errors ) NVML_FI_DEV_ECC_SBE_AGG_TEX = ( 26 # Texture memory single bit aggregate (persistent) ECC errors ) NVML_FI_DEV_ECC_DBE_AGG_TEX = ( 27 # Texture memory double bit aggregate (persistent) ECC errors ) NVML_FI_DEV_ECC_DBE_AGG_CBU = 28 # CBU double bit aggregate ECC errors # Page Retirement NVML_FI_DEV_RETIRED_SBE = 29 # Number of retired pages because of single bit errors NVML_FI_DEV_RETIRED_DBE = 30 # Number of retired pages because of double bit errors NVML_FI_DEV_RETIRED_PENDING = 31 # If any pages are pending retirement. 1=yes. 0=no. # NvLink Flit Error Counters NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L0 = ( 32 # NVLink flow control CRC Error Counter for Lane 0 ) NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L1 = ( 33 # NVLink flow control CRC Error Counter for Lane 1 ) NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L2 = ( 34 # NVLink flow control CRC Error Counter for Lane 2 ) NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L3 = ( 35 # NVLink flow control CRC Error Counter for Lane 3 ) NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L4 = ( 36 # NVLink flow control CRC Error Counter for Lane 4 ) NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L5 = ( 37 # NVLink flow control CRC Error Counter for Lane 5 ) NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_TOTAL = ( 38 # NVLink flow control CRC Error Counter total for all Lanes ) # NvLink CRC Data Error Counters NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L0 = ( 39 # NVLink data CRC Error Counter for Lane 0 ) NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L1 = ( 40 # NVLink data CRC Error Counter for Lane 1 ) NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L2 = ( 41 # NVLink data CRC Error Counter for Lane 2 ) NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L3 = ( 42 # NVLink data CRC Error Counter for Lane 3 ) NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L4 = ( 43 # NVLink data CRC Error Counter for Lane 4 ) NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L5 = ( 44 # NVLink data CRC Error Counter for Lane 5 ) NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_TOTAL = ( 45 # NvLink data CRC Error Counter total for all Lanes ) # NvLink Replay Error Counters NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L0 = 46 # NVLink Replay Error Counter for Lane 0 NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L1 = 47 # NVLink Replay Error Counter for Lane 1 NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L2 = 48 # NVLink Replay Error Counter for Lane 2 NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L3 = 49 # NVLink Replay Error Counter for Lane 3 NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L4 = 50 # NVLink Replay Error Counter for Lane 4 NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L5 = 51 # NVLink Replay Error Counter for Lane 5 NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_TOTAL = ( 52 # NVLink Replay Error Counter total for all Lanes ) # NvLink Recovery Error Counters NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L0 = ( 53 # NVLink Recovery Error Counter for Lane 0 ) NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L1 = ( 54 # NVLink Recovery Error Counter for Lane 1 ) NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L2 = ( 55 # NVLink Recovery Error Counter for Lane 2 ) NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L3 = ( 56 # NVLink Recovery Error Counter for Lane 3 ) NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L4 = ( 57 # NVLink Recovery Error Counter for Lane 4 ) NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L5 = ( 58 # NVLink Recovery Error Counter for Lane 5 ) NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_TOTAL = ( 59 # NVLink Recovery Error Counter total for all Lanes ) # NvLink Bandwidth Counters NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L0 = ( 60 # NVLink Bandwidth Counter for Counter Set 0, Lane 0 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L1 = ( 61 # NVLink Bandwidth Counter for Counter Set 0, Lane 1 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L2 = ( 62 # NVLink Bandwidth Counter for Counter Set 0, Lane 2 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L3 = ( 63 # NVLink Bandwidth Counter for Counter Set 0, Lane 3 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L4 = ( 64 # NVLink Bandwidth Counter for Counter Set 0, Lane 4 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L5 = ( 65 # NVLink Bandwidth Counter for Counter Set 0, Lane 5 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C0_TOTAL = ( 66 # NVLink Bandwidth Counter Total for Counter Set 0, All Lanes ) # NvLink Bandwidth Counters NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L0 = ( 67 # NVLink Bandwidth Counter for Counter Set 1, Lane 0 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L1 = ( 68 # NVLink Bandwidth Counter for Counter Set 1, Lane 1 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L2 = ( 69 # NVLink Bandwidth Counter for Counter Set 1, Lane 2 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L3 = ( 70 # NVLink Bandwidth Counter for Counter Set 1, Lane 3 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L4 = ( 71 # NVLink Bandwidth Counter for Counter Set 1, Lane 4 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L5 = ( 72 # NVLink Bandwidth Counter for Counter Set 1, Lane 5 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C1_TOTAL = ( 73 # NVLink Bandwidth Counter Total for Counter Set 1, All Lanes ) # Perf Policy Counters NVML_FI_DEV_PERF_POLICY_POWER = 74 # Perf Policy Counter for Power Policy NVML_FI_DEV_PERF_POLICY_THERMAL = 75 # Perf Policy Counter for Thermal Policy NVML_FI_DEV_PERF_POLICY_SYNC_BOOST = 76 # Perf Policy Counter for Sync boost Policy NVML_FI_DEV_PERF_POLICY_BOARD_LIMIT = 77 # Perf Policy Counter for Board Limit NVML_FI_DEV_PERF_POLICY_LOW_UTILIZATION = ( 78 # Perf Policy Counter for Low GPU Utilization Policy ) NVML_FI_DEV_PERF_POLICY_RELIABILITY = 79 # Perf Policy Counter for Reliability Policy NVML_FI_DEV_PERF_POLICY_TOTAL_APP_CLOCKS = ( 80 # Perf Policy Counter for Total App Clock Policy ) NVML_FI_DEV_PERF_POLICY_TOTAL_BASE_CLOCKS = ( 81 # Perf Policy Counter for Total Base Clocks Policy ) # Memory temperatures NVML_FI_DEV_MEMORY_TEMP = 82 # Memory temperature for the device # Energy Counter NVML_FI_DEV_TOTAL_ENERGY_CONSUMPTION = ( 83 # Total energy consumption for the GPU in mJ since the driver was last reloaded ) # NVLink Speed NVML_FI_DEV_NVLINK_SPEED_MBPS_L0 = 84 NVML_FI_DEV_NVLINK_SPEED_MBPS_L1 = 85 NVML_FI_DEV_NVLINK_SPEED_MBPS_L2 = 86 NVML_FI_DEV_NVLINK_SPEED_MBPS_L3 = 87 NVML_FI_DEV_NVLINK_SPEED_MBPS_L4 = 88 NVML_FI_DEV_NVLINK_SPEED_MBPS_L5 = 89 NVML_FI_DEV_NVLINK_SPEED_MBPS_COMMON = 90 # NVLink Link Count NVML_FI_DEV_NVLINK_LINK_COUNT = 91 # Page Retirement pending fields NVML_FI_DEV_RETIRED_PENDING_SBE = 92 NVML_FI_DEV_RETIRED_PENDING_DBE = 93 # PCIe replay and replay rollover counters NVML_FI_DEV_PCIE_REPLAY_COUNTER = 94 NVML_FI_DEV_PCIE_REPLAY_ROLLOVER_COUNTER = 95 # NvLink Flit Error Counters NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L6 = ( 96 # NVLink flow control CRC Error Counter for Lane 6 ) NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L7 = ( 97 # NVLink flow control CRC Error Counter for Lane 7 ) NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L8 = ( 98 # NVLink flow control CRC Error Counter for Lane 8 ) NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L9 = ( 99 # NVLink flow control CRC Error Counter for Lane 9 ) NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L10 = ( 100 # NVLink flow control CRC Error Counter for Lane 10 ) NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L11 = ( 101 # NVLink flow control CRC Error Counter for Lane 11 ) # NvLink CRC Data Error Counters NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L6 = ( 102 # NVLink data CRC Error Counter for Lane 6 ) NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L7 = ( 103 # NVLink data CRC Error Counter for Lane 7 ) NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L8 = ( 104 # NVLink data CRC Error Counter for Lane 8 ) NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L9 = ( 105 # NVLink data CRC Error Counter for Lane 9 ) NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L10 = ( 106 # NVLink data CRC Error Counter for Lane 10 ) NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L11 = ( 107 # NVLink data CRC Error Counter for Lane 11 ) # NvLink Replay Error Counters NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L6 = 108 # NVLink Replay Error Counter for Lane 6 NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L7 = 109 # NVLink Replay Error Counter for Lane 7 NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L8 = 110 # NVLink Replay Error Counter for Lane 8 NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L9 = 111 # NVLink Replay Error Counter for Lane 9 NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L10 = ( 112 # NVLink Replay Error Counter for Lane 10 ) NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L11 = ( 113 # NVLink Replay Error Counter for Lane 11 ) # NvLink Recovery Error Counters NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L6 = ( 114 # NVLink Recovery Error Counter for Lane 6 ) NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L7 = ( 115 # NVLink Recovery Error Counter for Lane 7 ) NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L8 = ( 116 # NVLink Recovery Error Counter for Lane 8 ) NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L9 = ( 117 # NVLink Recovery Error Counter for Lane 9 ) NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L10 = ( 118 # NVLink Recovery Error Counter for Lane 10 ) NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L11 = ( 119 # NVLink Recovery Error Counter for Lane 11 ) # NvLink Bandwidth Counters NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L6 = ( 120 # NVLink Bandwidth Counter for Counter Set 0, Lane 6 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L7 = ( 121 # NVLink Bandwidth Counter for Counter Set 0, Lane 7 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L8 = ( 122 # NVLink Bandwidth Counter for Counter Set 0, Lane 8 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L9 = ( 123 # NVLink Bandwidth Counter for Counter Set 0, Lane 9 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L10 = ( 124 # NVLink Bandwidth Counter for Counter Set 0, Lane 10 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L11 = ( 125 # NVLink Bandwidth Counter for Counter Set 0, Lane 11 ) # NvLink Bandwidth Counters NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L6 = ( 126 # NVLink Bandwidth Counter for Counter Set 1, Lane 6 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L7 = ( 127 # NVLink Bandwidth Counter for Counter Set 1, Lane 7 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L8 = ( 128 # NVLink Bandwidth Counter for Counter Set 1, Lane 8 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L9 = ( 129 # NVLink Bandwidth Counter for Counter Set 1, Lane 9 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L10 = ( 130 # NVLink Bandwidth Counter for Counter Set 1, Lane 10 ) NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L11 = ( 131 # NVLink Bandwidth Counter for Counter Set 1, Lane 11 ) # NVLink Speed NVML_FI_DEV_NVLINK_SPEED_MBPS_L6 = 132 NVML_FI_DEV_NVLINK_SPEED_MBPS_L7 = 133 NVML_FI_DEV_NVLINK_SPEED_MBPS_L8 = 134 NVML_FI_DEV_NVLINK_SPEED_MBPS_L9 = 135 NVML_FI_DEV_NVLINK_SPEED_MBPS_L10 = 136 NVML_FI_DEV_NVLINK_SPEED_MBPS_L11 = 137 # NVLink Throughput Counters NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_TX = 138 # NVLink TX Data throughput in KiB NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_RX = 139 # NVLink RX Data throughput in KiB NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_TX = 140 # NVLink TX Data + protocol overhead in KiB NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_RX = 141 # NVLink RX Data + protocol overhead in KiB # Row Remapper NVML_FI_DEV_REMAPPED_COR = 142 NVML_FI_DEV_REMAPPED_UNC = 143 NVML_FI_DEV_REMAPPED_PENDING = 144 NVML_FI_DEV_REMAPPED_FAILURE = 145 # Remote device NVLink ID NVML_FI_DEV_NVLINK_REMOTE_NVLINK_ID = 146 # Number of NVLinks connected to NVSwitch NVML_FI_DEV_NVSWITCH_CONNECTED_LINK_COUNT = 147 # NvLink ECC Data Error Counters NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L0 = ( 148 # < NVLink data ECC Error Counter for Link 0 ) NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L1 = ( 149 # < NVLink data ECC Error Counter for Link 1 ) NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L2 = ( 150 # < NVLink data ECC Error Counter for Link 2 ) NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L3 = ( 151 # < NVLink data ECC Error Counter for Link 3 ) NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L4 = ( 152 # < NVLink data ECC Error Counter for Link 4 ) NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L5 = ( 153 # < NVLink data ECC Error Counter for Link 5 ) NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L6 = ( 154 # < NVLink data ECC Error Counter for Link 6 ) NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L7 = ( 155 # < NVLink data ECC Error Counter for Link 7 ) NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L8 = ( 156 # < NVLink data ECC Error Counter for Link 8 ) NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L9 = ( 157 # < NVLink data ECC Error Counter for Link 9 ) NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L10 = ( 158 # < NVLink data ECC Error Counter for Link 10 ) NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L11 = ( 159 # < NVLink data ECC Error Counter for Link 11 ) NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_TOTAL = ( 160 # < NvLink data ECC Error Counter total for all Links ) NVML_FI_DEV_NVLINK_ERROR_DL_REPLAY = 161 NVML_FI_DEV_NVLINK_ERROR_DL_RECOVERY = 162 NVML_FI_DEV_NVLINK_ERROR_DL_CRC = 163 NVML_FI_DEV_NVLINK_GET_SPEED = 164 NVML_FI_DEV_NVLINK_GET_STATE = 165 NVML_FI_DEV_NVLINK_GET_VERSION = 166 NVML_FI_DEV_NVLINK_GET_POWER_STATE = 167 NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD = 168 NVML_FI_DEV_PCIE_L0_TO_RECOVERY_COUNTER = 169 NVML_FI_DEV_C2C_LINK_COUNT = 170 NVML_FI_DEV_C2C_LINK_GET_STATUS = 171 NVML_FI_DEV_C2C_LINK_GET_MAX_BW = 172 NVML_FI_DEV_PCIE_COUNT_CORRECTABLE_ERRORS = 173 NVML_FI_DEV_PCIE_COUNT_NAKS_RECEIVED = 174 NVML_FI_DEV_PCIE_COUNT_RECEIVER_ERROR = 175 NVML_FI_DEV_PCIE_COUNT_BAD_TLP = 176 NVML_FI_DEV_PCIE_COUNT_NAKS_SENT = 177 NVML_FI_DEV_PCIE_COUNT_BAD_DLLP = 178 NVML_FI_DEV_PCIE_COUNT_NON_FATAL_ERROR = 179 NVML_FI_DEV_PCIE_COUNT_FATAL_ERROR = 180 NVML_FI_DEV_PCIE_COUNT_UNSUPPORTED_REQ = 181 NVML_FI_DEV_PCIE_COUNT_LCRC_ERROR = 182 NVML_FI_DEV_PCIE_COUNT_LANE_ERROR = 183 NVML_FI_DEV_IS_RESETLESS_MIG_SUPPORTED = 184 NVML_FI_DEV_POWER_AVERAGE = 185 NVML_FI_DEV_POWER_INSTANT = 186 NVML_FI_DEV_POWER_MIN_LIMIT = 187 NVML_FI_DEV_POWER_MAX_LIMIT = 188 NVML_FI_DEV_POWER_DEFAULT_LIMIT = 189 NVML_FI_DEV_POWER_CURRENT_LIMIT = 190 NVML_FI_DEV_ENERGY = 191 NVML_FI_DEV_POWER_REQUESTED_LIMIT = 192 NVML_FI_DEV_TEMPERATURE_SHUTDOWN_TLIMIT = 193 NVML_FI_DEV_TEMPERATURE_SLOWDOWN_TLIMIT = 194 NVML_FI_DEV_TEMPERATURE_MEM_MAX_TLIMIT = 195 NVML_FI_DEV_TEMPERATURE_GPU_MAX_TLIMIT = 196 NVML_FI_DEV_PCIE_COUNT_TX_BYTES = 197 NVML_FI_DEV_PCIE_COUNT_RX_BYTES = 198 NVML_FI_DEV_IS_MIG_MODE_INDEPENDENT_MIG_QUERY_CAPABLE = 199 NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MAX = 200 NVML_FI_DEV_NVLINK_COUNT_XMIT_PACKETS = 201 NVML_FI_DEV_NVLINK_COUNT_XMIT_BYTES = 202 NVML_FI_DEV_NVLINK_COUNT_RCV_PACKETS = 203 NVML_FI_DEV_NVLINK_COUNT_RCV_BYTES = 204 NVML_FI_DEV_NVLINK_COUNT_VL15_DROPPED = 205 # Deprecated, do not use NVML_FI_DEV_NVLINK_COUNT_MALFORMED_PACKET_ERRORS = 206 NVML_FI_DEV_NVLINK_COUNT_BUFFER_OVERRUN_ERRORS = 207 NVML_FI_DEV_NVLINK_COUNT_RCV_ERRORS = 208 NVML_FI_DEV_NVLINK_COUNT_RCV_REMOTE_ERRORS = 209 NVML_FI_DEV_NVLINK_COUNT_RCV_GENERAL_ERRORS = 210 NVML_FI_DEV_NVLINK_COUNT_LOCAL_LINK_INTEGRITY_ERRORS = 211 NVML_FI_DEV_NVLINK_COUNT_XMIT_DISCARDS = 212 NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_SUCCESSFUL_EVENTS = 213 NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_FAILED_EVENTS = 214 NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_EVENTS = 215 NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE0 = 216 # Deprecated, do not use NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE1 = 217 # Deprecated, do not use NVML_FI_DEV_NVLINK_COUNT_RAW_BER = 218 # Deprecated, do not use NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_ERRORS = 219 NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_BER = 220 NVML_FI_DEV_NVLINK_COUNT_SYMBOL_ERRORS = 221 NVML_FI_DEV_NVLINK_COUNT_SYMBOL_BER = 222 NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MIN = 223 NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS = ( 224 # Values are in the form NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_* ) NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_SUPPORTED = 225 NVML_FI_DEV_RESET_STATUS = ( 226 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead ) NVML_FI_DEV_DRAIN_AND_RESET_STATUS = ( 227 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead ) NVML_FI_DEV_PCIE_OUTBOUND_ATOMICS_MASK = 228 NVML_FI_DEV_PCIE_INBOUND_ATOMICS_MASK = 229 NVML_FI_DEV_GET_GPU_RECOVERY_ACTION = 230 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_0 = 235 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_1 = 236 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_2 = 237 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_3 = 238 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_4 = 239 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_5 = 240 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_6 = 241 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_7 = 242 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_8 = 243 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_9 = 244 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_10 = 245 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_11 = 246 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_12 = 247 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_13 = 248 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_14 = 249 NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_15 = 250 NVML_FI_PWR_SMOOTHING_ENABLED = 251 # Enablement (0/DISABLED or 1/ENABLED) NVML_FI_PWR_SMOOTHING_PRIV_LVL = 252 # Current privilege level NVML_FI_PWR_SMOOTHING_IMM_RAMP_DOWN_ENABLED = ( 253 # Immediate ramp down enablement (0/DISABLED or 1/ENABLED) ) NVML_FI_PWR_SMOOTHING_APPLIED_TMP_CEIL = 254 # Applied TMP ceiling value NVML_FI_PWR_SMOOTHING_APPLIED_TMP_FLOOR = 255 # Applied TMP floor value NVML_FI_PWR_SMOOTHING_MAX_PERCENT_TMP_FLOOR_SETTING = 256 # Max % TMP Floor value NVML_FI_PWR_SMOOTHING_MIN_PERCENT_TMP_FLOOR_SETTING = 257 # Min % TMP Floor value NVML_FI_PWR_SMOOTHING_HW_CIRCUITRY_PERCENT_LIFETIME_REMAINING = ( 258 # HW Circuitry % lifetime remaining ) NVML_FI_PWR_SMOOTHING_MAX_NUM_PRESET_PROFILES = 259 # Max number of preset profiles NVML_FI_PWR_SMOOTHING_PROFILE_PERCENT_TMP_FLOOR = 260 # % TMP floor for a given profile NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_UP_RATE = ( 261 # Ramp up rate in mW/s for a given profile ) NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_RATE = ( 262 # Ramp down rate in mW/s for a given profile ) NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_HYST_VAL = ( 263 # Ramp down hysteresis value in ms for a given profile ) NVML_FI_PWR_SMOOTHING_ACTIVE_PRESET_PROFILE = 264 # Active preset profile number NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_PERCENT_TMP_FLOOR = ( 265 # % TMP floor for a given profile ) NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_UP_RATE = ( 266 # Ramp up rate in mW/s for a given profile ) NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_RATE = ( 267 # Ramp down rate in mW/s for a given profile ) NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_HYST_VAL = ( 268 # Ramp down hysteresis value in ms for a given profile ) NVML_FI_MAX = 269 # One greater than the largest field ID defined above # NVML_FI_DEV_NVLINK_GET_STATE state enums NVML_NVLINK_STATE_INACTIVE = 0x0 NVML_NVLINK_STATE_ACTIVE = 0x1 NVML_NVLINK_STATE_SLEEP = 0x2 NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_100US = ( 0 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS ) NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_50US = ( 1 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS ) ## Enums needed for the method nvmlDeviceGetVirtualizationMode and nvmlDeviceSetVirtualizationMode NVML_GPU_VIRTUALIZATION_MODE_NONE = 0 # Represents Bare Metal GPU NVML_GPU_VIRTUALIZATION_MODE_PASSTHROUGH = ( 1 # Device is associated with GPU-Passthorugh ) NVML_GPU_VIRTUALIZATION_MODE_VGPU = ( 2 # Device is associated with vGPU inside virtual machine. ) NVML_GPU_VIRTUALIZATION_MODE_HOST_VGPU = ( 3 # Device is associated with VGX hypervisor in vGPU mode ) NVML_GPU_VIRTUALIZATION_MODE_HOST_VSGA = ( 4 # Device is associated with VGX hypervisor in vSGA mode ) ## Lib loading ## nvmlLib = None libLoadLock = threading.Lock() _nvmlLib_refcount = 0 # Incremented on each nvmlInit and decremented on nvmlShutdown ## vGPU Management _nvmlVgpuTypeId_t = c_uint _nvmlVgpuInstance_t = c_uint _nvmlVgpuVmIdType_t = c_uint NVML_VGPU_VM_ID_DOMAIN_ID = 0 NVML_VGPU_VM_ID_UUID = 1 _nvmlGridLicenseFeatureCode_t = c_uint NVML_GRID_LICENSE_FEATURE_CODE_UNKNOWN = 0 NVML_GRID_LICENSE_FEATURE_CODE_VGPU = 1 NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX = 2 NVML_GRID_LICENSE_FEATURE_CODE_VWORKSTATION = ( 2 # deprecated, use NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX. ) NVML_GRID_LICENSE_FEATURE_CODE_GAMING = 3 NVML_GRID_LICENSE_FEATURE_CODE_COMPUTE = 4 _nvmlGridLicenseExpiryStatus_t = c_uint8 NVML_GRID_LICENSE_EXPIRY_NOT_AVAILABLE = (0,) # Expiry information not available NVML_GRID_LICENSE_EXPIRY_INVALID = (1,) # Invalid expiry or error fetching expiry NVML_GRID_LICENSE_EXPIRY_VALID = (2,) # Valid expiry NVML_GRID_LICENSE_EXPIRY_NOT_APPLICABLE = (3,) # Expiry not applicable NVML_GRID_LICENSE_EXPIRY_PERMANENT = (4,) # Permanent expiry _nvmlVgpuCapability_t = c_uint NVML_VGPU_CAP_NVLINK_P2P = 0 # vGPU P2P over NVLink is supported NVML_VGPU_CAP_GPUDIRECT = 1 # GPUDirect capability is supported NVML_VGPU_CAP_MULTI_VGPU_EXCLUSIVE = ( 2 # vGPU profile cannot be mixed with other vGPU profiles in same VM ) NVML_VGPU_CAP_EXCLUSIVE_TYPE = ( 3 # vGPU profile cannot run on a GPU alongside other profiles of different type ) NVML_VGPU_CAP_EXCLUSIVE_SIZE = ( 4 # vGPU profile cannot run on a GPU alongside other profiles of different size ) NVML_VGPU_CAP_COUNT = 5 _nvmlVgpuDriverCapability_t = c_uint NVML_VGPU_DRIVER_CAP_HETEROGENEOUS_MULTI_VGPU = ( 0 # Supports mixing of different vGPU profiles within one guest VM ) NVML_VGPU_DRIVER_CAP_WARM_UPDATE = 1 # Supports FSR and warm update of vGPU host driver without terminating the running guest VM NVML_VGPU_DRIVER_CAP_COUNT = 2 _nvmlDeviceVgpuCapability_t = c_uint NVML_DEVICE_VGPU_CAP_FRACTIONAL_MULTI_VGPU = 0 # Query whether the fractional vGPU profiles on this GPU can be used in multi-vGPU configurations NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_PROFILES = 1 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing types NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_SIZES = 2 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing framebuffer sizes NVML_DEVICE_VGPU_CAP_READ_DEVICE_BUFFER_BW = 3 # Query the GPU's read_device_buffer expected bandwidth capacity in megabytes per second NVML_DEVICE_VGPU_CAP_WRITE_DEVICE_BUFFER_BW = 4 # Query the GPU's write_device_buffer expected bandwidth capacity in megabytes per second NVML_DEVICE_VGPU_CAP_DEVICE_STREAMING = ( 5 # Query whether the vGPU profiles on the GPU supports migration data streaming ) NVML_DEVICE_VGPU_CAP_MINI_QUARTER_GPU = ( 6 # Set/Get support of mini-quarter vGPU profiles ) NVML_DEVICE_VGPU_CAP_COMPUTE_MEDIA_ENGINE_GPU = ( 7 # Set/Get support for compute media engine vGPU profiles ) NVML_DEVICE_VGPU_CAP_WARM_UPDATE = ( 8 # Query whether the GPU supports FSR and warm update ) NVML_DEVICE_VGPU_CAP_HOMOGENEOUS_PLACEMENTS = 9 # Query whether the GPU supports reporting of placements of timesliced vGPU profiles with identical framebuffer sizes NVML_DEVICE_VGPU_CAP_COUNT = 10 _nvmlVgpuGuestInfoState_t = c_uint NVML_VGPU_INSTANCE_GUEST_INFO_STATE_UNINITIALIZED = 0 NVML_VGPU_INSTANCE_GUEST_INFO_STATE_INITIALIZED = 1 _nvmlVgpuVmCompatibility_t = c_uint NVML_VGPU_VM_COMPATIBILITY_NONE = 0x0 NVML_VGPU_VM_COMPATIBILITY_COLD = 0x1 NVML_VGPU_VM_COMPATIBILITY_HIBERNATE = 0x2 NVML_VGPU_VM_COMPATIBILITY_SLEEP = 0x4 NVML_VGPU_VM_COMPATIBILITY_LIVE = 0x8 _nvmlVgpuPgpuCompatibilityLimitCode_t = c_uint NVML_VGPU_COMPATIBILITY_LIMIT_NONE = 0x0 NVML_VGPU_COMPATIBILITY_LIMIT_HOST_DRIVER = 0x1 NVML_VGPU_COMPATIBILITY_LIMIT_GUEST_DRIVER = 0x2 NVML_VGPU_COMPATIBILITY_LIMIT_GPU = 0x4 NVML_VGPU_COMPATIBILITY_LIMIT_OTHER = 0x80000000 _nvmlHostVgpuMode_t = c_uint NVML_HOST_VGPU_MODE_NON_SRIOV = 0 NVML_HOST_VGPU_MODE_SRIOV = 1 _nvmlConfComputeGpusReadyState_t = c_uint NVML_CC_ACCEPTING_CLIENT_REQUESTS_FALSE = 0 NVML_CC_ACCEPTING_CLIENT_REQUESTS_TRUE = 1 _nvmlConfComputeGpuCaps_t = c_uint NVML_CC_SYSTEM_GPUS_CC_NOT_CAPABLE = 0 NVML_CC_SYSTEM_GPUS_CC_CAPABLE = 1 _nvmlConfComputeCpuCaps_t = c_uint NVML_CC_SYSTEM_CPU_CAPS_NONE = 0 NVML_CC_SYSTEM_CPU_CAPS_AMD_SEV = 1 NVML_CC_SYSTEM_CPU_CAPS_INTEL_TDX = 2 NVML_CC_SYSTEM_CPU_CAPS_AMD_SEV_SNP = 3 NVML_CC_SYSTEM_CPU_CAPS_AMD_SNP_VTOM = 4 _nvmlConfComputeDevToolsMode_t = c_uint NVML_CC_SYSTEM_DEVTOOLS_MODE_OFF = 0 NVML_CC_SYSTEM_DEVTOOLS_MODE_ON = 1 NVML_CC_SYSTEM_MULTIGPU_NONE = 0 NVML_CC_SYSTEM_MULTIGPU_PROTECTED_PCIE = 1 NVML_CC_SYSTEM_ENVIRONMENT_UNAVAILABLE = 0 NVML_CC_SYSTEM_ENVIRONMENT_SIM = 1 NVML_CC_SYSTEM_ENVIRONMENT_PROD = 2 _nvmlConfComputeCcFeature_t = c_uint NVML_CC_SYSTEM_FEATURE_DISABLED = 0 NVML_CC_SYSTEM_FEATURE_ENABLED = 1 _nvmlConfComputeCcKeyRotationThreshAttackerAdv_t = c_uint NVML_CC_KEY_ROTATION_THRESH_ATTACKER_ADVANTAGE_MIN = 50 NVML_CC_KEY_ROTATION_THRESH_ATTACKER_ADVANTAGE_MAX = 65 # GSP firmware NVML_GSP_FIRMWARE_VERSION_BUF_SIZE = 0x40 class NVMLLibraryMismatchError(Exception): pass ## Error Checking ## class NVMLError(Exception): _valClassMapping = dict() # List of currently known error codes _errcode_to_string = { NVML_ERROR_UNINITIALIZED: "Uninitialized", NVML_ERROR_INVALID_ARGUMENT: "Invalid Argument", NVML_ERROR_NOT_SUPPORTED: "Not Supported", NVML_ERROR_NO_PERMISSION: "Insufficient Permissions", NVML_ERROR_ALREADY_INITIALIZED: "Already Initialized", NVML_ERROR_NOT_FOUND: "Not Found", NVML_ERROR_INSUFFICIENT_SIZE: "Insufficient Size", NVML_ERROR_INSUFFICIENT_POWER: "Insufficient External Power", NVML_ERROR_DRIVER_NOT_LOADED: "Driver Not Loaded", NVML_ERROR_TIMEOUT: "Timeout", NVML_ERROR_IRQ_ISSUE: "Interrupt Request Issue", NVML_ERROR_LIBRARY_NOT_FOUND: "NVML Shared Library Not Found", NVML_ERROR_FUNCTION_NOT_FOUND: "Function Not Found", NVML_ERROR_CORRUPTED_INFOROM: "Corrupted infoROM", NVML_ERROR_GPU_IS_LOST: "GPU is lost", NVML_ERROR_RESET_REQUIRED: "GPU requires restart", NVML_ERROR_OPERATING_SYSTEM: "The operating system has blocked the request.", NVML_ERROR_LIB_RM_VERSION_MISMATCH: "RM has detected an NVML/RM version mismatch.", NVML_ERROR_MEMORY: "Insufficient Memory", NVML_ERROR_UNKNOWN: "Unknown Error", } def __new__(typ, value): """ Maps value to a proper subclass of NVMLError. See _extractNVMLErrorsAsClasses function for more details """ if typ == NVMLError: typ = NVMLError._valClassMapping.get(value, typ) obj = Exception.__new__(typ) obj.value = value return obj def __str__(self): try: if self.value not in NVMLError._errcode_to_string: NVMLError._errcode_to_string[self.value] = str( nvmlErrorString(self.value) ) return NVMLError._errcode_to_string[self.value] except NVMLError: return "NVML Error with code %d" % self.value def __eq__(self, other): return self.value == other.value def nvmlExceptionClass(nvmlErrorCode): if nvmlErrorCode not in NVMLError._valClassMapping: raise ValueError("nvmlErrorCode %s is not valid" % nvmlErrorCode) return NVMLError._valClassMapping[nvmlErrorCode] def _extractNVMLErrorsAsClasses(): """ Generates a hierarchy of classes on top of NVMLError class. Each NVML Error gets a new NVMLError subclass. This way try,except blocks can filter appropriate exceptions more easily. NVMLError is a parent class. Each NVML_ERROR_* gets it's own subclass. e.g. NVML_ERROR_ALREADY_INITIALIZED will be turned into NVMLError_AlreadyInitialized """ this_module = sys.modules[__name__] nvmlErrorsNames = [x for x in dir(this_module) if x.startswith("NVML_ERROR_")] for err_name in nvmlErrorsNames: # e.g. Turn NVML_ERROR_ALREADY_INITIALIZED into NVMLError_AlreadyInitialized class_name = "NVMLError_" + string.capwords( err_name.replace("NVML_ERROR_", ""), "_" ).replace("_", "") err_val = getattr(this_module, err_name) def gen_new(val): def new(typ): obj = NVMLError.__new__(typ, val) return obj return new new_error_class = type(class_name, (NVMLError,), {"__new__": gen_new(err_val)}) new_error_class.__module__ = __name__ setattr(this_module, class_name, new_error_class) NVMLError._valClassMapping[err_val] = new_error_class _extractNVMLErrorsAsClasses() def _nvmlCheckReturn(ret): if ret != NVML_SUCCESS: raise NVMLError(ret) return ret ## Function access ## _nvmlGetFunctionPointer_cache = ( dict() ) # function pointers are cached to prevent unnecessary libLoadLock locking def _nvmlGetFunctionPointer(name): global nvmlLib if name in _nvmlGetFunctionPointer_cache: return _nvmlGetFunctionPointer_cache[name] libLoadLock.acquire() try: # ensure library was loaded if nvmlLib is None: raise NVMLError(NVML_ERROR_UNINITIALIZED) try: _nvmlGetFunctionPointer_cache[name] = getattr(nvmlLib, name) return _nvmlGetFunctionPointer_cache[name] except AttributeError: raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) finally: # lock is always freed libLoadLock.release() ## Alternative object # Allows the object to be printed # Allows mismatched types to be assigned # - like None when the Structure variant requires c_uint class nvmlFriendlyObject(object): def __init__(self, dictionary): for x in dictionary: setattr(self, x, dictionary[x]) def __str__(self): return self.__dict__.__str__() def nvmlStructToFriendlyObject(struct): d = {} for x in struct._fields_: key = x[0] value = getattr(struct, key) # only need to convert from bytes if bytes, no need to check python version. d[key] = value.decode() if isinstance(value, bytes) else value obj = nvmlFriendlyObject(d) return obj # pack the object so it can be passed to the NVML library def nvmlFriendlyObjectToStruct(obj, model): for x in model._fields_: key = x[0] value = obj.__dict__[key] # any c_char_p in python3 needs to be bytes, default encoding works fine. if sys.version_info >= (3,): setattr(model, key, value.encode()) else: setattr(model, key, value) return model ## Unit structures class struct_c_nvmlUnit_t(Structure): pass # opaque handle c_nvmlUnit_t = POINTER(struct_c_nvmlUnit_t) class _PrintableStructure(Structure): """ Abstract class that produces nicer __str__ output than ctypes.Structure. e.g. instead of: >>> print str(obj) this class will print class_name(field_name: formatted_value, field_name: formatted_value) _fmt_ dictionary of -> e.g. class that has _field_ 'hex_value', c_uint could be formatted with _fmt_ = {"hex_value" : "%08X"} to produce nicer output. Default formatting string for all fields can be set with key "" like: _fmt_ = {"" : "%d MHz"} # e.g all values are numbers in MHz. If not set it's assumed to be just "%s" Exact format of returned str from this class is subject to change in the future. """ _fmt_ = {} def __str__(self): result = [] for x in self._fields_: key = x[0] value = getattr(self, key) fmt = "%s" if key in self._fmt_: fmt = self._fmt_[key] elif "" in self._fmt_: fmt = self._fmt_[""] result.append(("%s: " + fmt) % (key, value)) return self.__class__.__name__ + "(" + ", ".join(result) + ")" def __getattribute__(self, name): res = super(_PrintableStructure, self).__getattribute__(name) # need to convert bytes to unicode for python3 don't need to for python2 # Python 2 strings are of both str and bytes # Python 3 strings are not of type bytes # ctypes should convert everything to the correct values otherwise if isinstance(res, bytes): if isinstance(res, str): return res return res.decode() return res def __setattr__(self, name, value): if isinstance(value, str): # encoding a python2 string returns the same value, since python2 strings are bytes already # bytes passed in python3 will be ignored. value = value.encode() super(_PrintableStructure, self).__setattr__(name, value) class c_nvmlUnitInfo_t(_PrintableStructure): _fields_ = [ ("name", c_char * 96), ("id", c_char * 96), ("serial", c_char * 96), ("firmwareVersion", c_char * 96), ] class c_nvmlC2cModeInfo_v1_t(_PrintableStructure): _fields_ = [("isC2cEnabled", c_uint)] nvmlC2cModeInfo_v1 = 0x1000008 class c_nvmlLedState_t(_PrintableStructure): _fields_ = [ ("cause", c_char * 256), ("color", _nvmlLedColor_t), ] class c_nvmlPSUInfo_t(_PrintableStructure): _fields_ = [ ("state", c_char * 256), ("current", c_uint), ("voltage", c_uint), ("power", c_uint), ] class c_nvmlUnitFanInfo_t(_PrintableStructure): _fields_ = [ ("speed", c_uint), ("state", _nvmlFanState_t), ] class c_nvmlUnitFanSpeeds_t(_PrintableStructure): _fields_ = [("fans", c_nvmlUnitFanInfo_t * 24), ("count", c_uint)] ## Device structures class struct_c_nvmlDevice_t(Structure): pass # opaque handle c_nvmlDevice_t = POINTER(struct_c_nvmlDevice_t) class nvmlPciInfoExt_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("domain", c_uint), ("bus", c_uint), ("device", c_uint), ("pciDeviceId", c_uint), ("pciSubSystemId", c_uint), ("baseClass", c_uint), ("subClass", c_uint), ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), ] _fmt_ = { "version": "0x%04X", "domain": "0x%04X", "bus": "0x%02X", "device": "0x%02X", "pciDeviceId": "0x%08X", "pciSubSystemId": "0x%08X", "baseClass": "0x%01X", "subClass": "0x%01X", } nvmlPciInfoExt_v1 = 0x1000040 # Legacy pciInfo used for _v1 and _v2 class nvmlPciInfo_v2_t(_PrintableStructure): _fields_ = [ ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), ("domain", c_uint), ("bus", c_uint), ("device", c_uint), ("pciDeviceId", c_uint), # Added in 2.285 ("pciSubSystemId", c_uint), ("reserved0", c_uint), ("reserved1", c_uint), ("reserved2", c_uint), ("reserved3", c_uint), ] _fmt_ = { "domain": "0x%04X", "bus": "0x%02X", "device": "0x%02X", "pciDeviceId": "0x%08X", "pciSubSystemId": "0x%08X", } class nvmlPciInfo_t(_PrintableStructure): _fields_ = [ # Moved to the new busId location below ("busIdLegacy", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), ("domain", c_uint), ("bus", c_uint), ("device", c_uint), ("pciDeviceId", c_uint), # Added in 2.285 ("pciSubSystemId", c_uint), # New busId replaced the long deprecated and reserved fields with a # field of the same size in 9.0 ("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), ] _fmt_ = { "domain": "0x%08X", "bus": "0x%02X", "device": "0x%02X", "pciDeviceId": "0x%08X", "pciSubSystemId": "0x%08X", } class c_nvmlSystemDriverBranchInfo_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("branch", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), ] SystemDriverBranchInfo_v1 = 0x1000054 class c_nvmlExcludedDeviceInfo_t(_PrintableStructure): _fields_ = [("pci", nvmlPciInfo_t), ("uuid", c_char * NVML_DEVICE_UUID_BUFFER_SIZE)] class nvmlNvLinkUtilizationControl_t(_PrintableStructure): _fields_ = [ ("units", _nvmlNvLinkUtilizationCountUnits_t), ("pktfilter", _nvmlNvLinkUtilizationCountPktTypes_t), ] class c_nvmlMemory_t(_PrintableStructure): _fields_ = [ ("total", c_ulonglong), ("free", c_ulonglong), ("used", c_ulonglong), ] _fmt_ = {"": "%d B"} class c_nvmlMemory_v2_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("total", c_ulonglong), ("reserved", c_ulonglong), ("free", c_ulonglong), ("used", c_ulonglong), ] _fmt_ = {"": "%d B"} nvmlMemory_v2 = 0x02000028 class c_nvmlBAR1Memory_t(_PrintableStructure): _fields_ = [ ("bar1Total", c_ulonglong), ("bar1Free", c_ulonglong), ("bar1Used", c_ulonglong), ] _fmt_ = {"": "%d B"} class nvmlClkMonFaultInfo_t(Structure): _fields_ = [("clkApiDomain", c_uint), ("clkDomainFaultMask", c_uint)] MAX_CLK_DOMAINS = 32 class nvmlClkMonStatus_t(Structure): _fields_ = [ ("bGlobalStatus", c_uint), ("clkMonListSize", c_uint), ("clkMonList", nvmlClkMonFaultInfo_t * MAX_CLK_DOMAINS), ] # On Windows with the WDDM driver, usedGpuMemory is reported as None # Code that processes this structure should check for None, I.E. # # if (info.usedGpuMemory is None): # # TODO handle the error # pass # else: # print("Using %d MiB of memory" % (info.usedGpuMemory / 1024 / 1024)) # endif # # See NVML documentation for more information class c_nvmlProcessInfo_v2_t(_PrintableStructure): _fields_ = [ ("pid", c_uint), ("usedGpuMemory", c_ulonglong), ("gpuInstanceId", c_uint), ("computeInstanceId", c_uint), ] _fmt_ = {"usedGpuMemory": "%d B"} c_nvmlProcessInfo_v3_t = c_nvmlProcessInfo_v2_t c_nvmlProcessInfo_t = c_nvmlProcessInfo_v3_t _nvmlProcessMode_t = c_uint NVML_PROCESS_MODE_COMPUTE = 0 NVML_PROCESS_MODE_GRAPHICS = 1 NVML_PROCESS_MODE_MPS = 2 class c_nvmlProcessDetail_v1_t(Structure): _fields_ = [ ("pid", c_uint), ("usedGpuMemory", c_ulonglong), ("gpuInstanceId", c_uint), ("computeInstanceId", c_uint), ("usedGpuCcProtectedMemory", c_ulonglong), ] class c_nvmlProcessDetailList_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("mode", _nvmlProcessMode_t), ("numProcArrayEntries", c_uint), ("procArray", POINTER(c_nvmlProcessDetail_v1_t)), ] _fmt_ = {"numProcArrayEntries": "%d B"} c_nvmlProcessDetailList_t = c_nvmlProcessDetailList_v1_t nvmlProcessDetailList_v1 = 0x1000018 class c_nvmlBridgeChipInfo_t(_PrintableStructure): _fields_ = [ ("type", _nvmlBridgeChipType_t), ("fwVersion", c_uint), ] class c_nvmlBridgeChipHierarchy_t(_PrintableStructure): _fields_ = [ ("bridgeCount", c_uint), ("bridgeChipInfo", c_nvmlBridgeChipInfo_t * 128), ] class c_nvmlEccErrorCounts_t(_PrintableStructure): _fields_ = [ ("l1Cache", c_ulonglong), ("l2Cache", c_ulonglong), ("deviceMemory", c_ulonglong), ("registerFile", c_ulonglong), ] class c_nvmlUtilization_t(_PrintableStructure): _fields_ = [ ("gpu", c_uint), ("memory", c_uint), ] _fmt_ = {"": "%d %%"} # Added in 2.285 class c_nvmlHwbcEntry_t(_PrintableStructure): _fields_ = [ ("hwbcId", c_uint), ("firmwareVersion", c_char * 32), ] class c_nvmlValue_t(Union): _fields_ = [ ("dVal", c_double), ("uiVal", c_uint), ("ulVal", c_ulong), ("ullVal", c_ulonglong), ("sllVal", c_longlong), ("siVal", c_int), ("usVal", c_ushort), ] class c_nvmlSample_t(_PrintableStructure): _fields_ = [ ("timeStamp", c_ulonglong), ("sampleValue", c_nvmlValue_t), ] class c_nvmlViolationTime_t(_PrintableStructure): _fields_ = [ ("referenceTime", c_ulonglong), ("violationTime", c_ulonglong), ] class c_nvmlFieldValue_t(_PrintableStructure): _fields_ = [ ("fieldId", c_uint32), ("scopeId", c_uint32), ("timestamp", c_int64), ("latencyUsec", c_int64), ("valueType", _nvmlValueType_t), ("nvmlReturn", _nvmlReturn_t), ("value", c_nvmlValue_t), ] NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES = 23 nvmlNvlinkSupportedBwModes_v1 = 0x100001C class c_nvmlNvlinkSupportedBwModes_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("bwModes", c_uint8 * NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES), ("totalBwModes", c_uint8), ] def __init__(self): super(c_nvmlNvlinkSupportedBwModes_v1_t, self).__init__( version=nvmlNvlinkSupportedBwModes_v1 ) nvmlNvlinkGetBwMode_v1 = 0x100000C class c_nvmlNvlinkGetBwMode_v1_t(_PrintableStructure): _fields_ = [("version", c_uint), ("bIsBest", c_uint), ("bwMode", c_uint8)] def __init__(self): super(c_nvmlNvlinkGetBwMode_v1_t, self).__init__(version=nvmlNvlinkGetBwMode_v1) nvmlNvlinkSetBwMode_v1 = 0x100000C class c_nvmlNvlinkSetBwMode_v1_t(_PrintableStructure): _fields_ = [("version", c_uint), ("bSetBest", c_uint), ("bwMode", c_uint8)] def __init__(self): super(c_nvmlNvlinkSetBwMode_v1_t, self).__init__(version=nvmlNvlinkSetBwMode_v1) class c_nvmlVgpuHeterogeneousMode_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("mode", c_uint), ] VgpuHeterogeneousMode_v1 = 0x1000008 class c_nvmlVgpuPlacementId_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("placementId", c_uint), ] VgpuPlacementId_v1 = 0x1000008 class c_nvmlVgpuPlacementList_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("count", c_uint), ("placementSize", c_uint), ("placementIds", POINTER(c_uint)), ] VgpuPlacementList_v1 = 0x1000018 NVML_VGPU_PGPU_HETEROGENEOUS_MODE = 0 NVML_VGPU_PGPU_HOMOGENEOUS_MODE = 1 class c_nvmlVgpuPlacementList_v2_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("placementSize", c_uint), ("count", c_uint), ("placementIds", POINTER(c_uint)), ("mode", c_uint), ] VgpuPlacementList_v2 = 0x2000020 class c_nvmlVgpuTypeBar1Info_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("bar1Size", c_ulonglong), ] VgpuTypeBar1Info_v1 = 0x1000010 class c_nvmlVgpuInstanceUtilizationSample_t(_PrintableStructure): _fields_ = [ ("vgpuInstance", _nvmlVgpuInstance_t), ("timeStamp", c_ulonglong), ("smUtil", c_nvmlValue_t), ("memUtil", c_nvmlValue_t), ("encUtil", c_nvmlValue_t), ("decUtil", c_nvmlValue_t), ] class c_nvmlVgpuInstanceUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ ("timeStamp", c_ulonglong), ("vgpuInstance", _nvmlVgpuInstance_t), ("smUtil", c_nvmlValue_t), ("memUtil", c_nvmlValue_t), ("encUtil", c_nvmlValue_t), ("decUtil", c_nvmlValue_t), ("jpgUtil", c_nvmlValue_t), ("ofaUtil", c_nvmlValue_t), ] class c_nvmlVgpuInstancesUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("sampleValType", _nvmlValueType_t), ("vgpuInstanceCount", c_uint), ("lastSeenTimeStamp", c_ulonglong), ("vgpuUtilArray", POINTER(c_nvmlVgpuInstanceUtilizationInfo_v1_t)), ] VgpuInstancesUtilizationInfo_v1 = 0x01000020 class c_nvmlVgpuProcessUtilizationSample_t(_PrintableStructure): _fields_ = [ ("vgpuInstance", _nvmlVgpuInstance_t), ("pid", c_uint), ("processName", c_char * NVML_VGPU_NAME_BUFFER_SIZE), ("timeStamp", c_ulonglong), ("smUtil", c_uint), ("memUtil", c_uint), ("encUtil", c_uint), ("decUtil", c_uint), ] class c_nvmlVgpuProcessUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ ("processName", c_char * NVML_VGPU_NAME_BUFFER_SIZE), ("timeStamp", c_ulonglong), ("vgpuInstance", _nvmlVgpuInstance_t), ("pid", c_uint), ("smUtil", c_uint), ("memUtil", c_uint), ("encUtil", c_uint), ("decUtil", c_uint), ("jpgUtil", c_uint), ("ofaUtil", c_uint), ] class c_nvmlVgpuProcessesUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("vgpuProcessCount", c_uint), ("lastSeenTimeStamp", c_ulonglong), ("vgpuProcUtilArray", POINTER(c_nvmlVgpuProcessUtilizationInfo_v1_t)), ] VgpuProcessesUtilizationInfo_v1 = 0x01000018 class nvmlVgpuRuntimeState_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("size", c_ulonglong), ] VgpuRuntimeState_v1 = 0x1000010 class c_nvmlVgpuLicenseExpiry_t(_PrintableStructure): _fields_ = [ ("year", c_uint32), ("month", c_uint16), ("day", c_uint16), ("hour", c_uint16), ("min", c_uint16), ("sec", c_uint16), ("status", c_uint8), ] NVML_GRID_LICENSE_STATE_UNKNOWN = 0 NVML_GRID_LICENSE_STATE_UNINITIALIZED = 1 NVML_GRID_LICENSE_STATE_UNLICENSED_UNRESTRICTED = 2 NVML_GRID_LICENSE_STATE_UNLICENSED_RESTRICTED = 3 NVML_GRID_LICENSE_STATE_UNLICENSED = 4 NVML_GRID_LICENSE_STATE_LICENSED = 5 class c_nvmlVgpuLicenseInfo_t(_PrintableStructure): _fields_ = [ ("isLicensed", c_uint8), ("licenseExpiry", c_nvmlVgpuLicenseExpiry_t), ("currentState", c_uint), ] class c_nvmlEncoderSession_t(_PrintableStructure): _fields_ = [ ("sessionId", c_uint), ("pid", c_uint), ("vgpuInstance", _nvmlVgpuInstance_t), ("codecType", c_uint), ("hResolution", c_uint), ("vResolution", c_uint), ("averageFps", c_uint), ("encodeLatency", c_uint), ] class c_nvmlProcessUtilizationSample_t(_PrintableStructure): _fields_ = [ ("pid", c_uint), ("timeStamp", c_ulonglong), ("smUtil", c_uint), ("memUtil", c_uint), ("encUtil", c_uint), ("decUtil", c_uint), ] class c_nvmlProcessUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ ("timeStamp", c_ulonglong), ("pid", c_uint), ("smUtil", c_uint), ("memUtil", c_uint), ("encUtil", c_uint), ("decUtil", c_uint), ("jpgUtil", c_uint), ("ofaUtil", c_uint), ] class c_nvmlProcessesUtilizationInfo_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("processSamplesCount", c_uint), ("lastSeenTimeStamp", c_ulonglong), ("procUtilArray", POINTER(c_nvmlProcessUtilizationInfo_v1_t)), ] ProcessesUtilizationInfo_v1 = 0x01000018 class c_nvmlGridLicenseExpiry_t(_PrintableStructure): _fields_ = [ ("year", c_uint32), ("month", c_uint16), ("day", c_uint16), ("hour", c_uint16), ("min", c_uint16), ("sec", c_uint16), ("status", c_uint8), ] class c_nvmlGridLicensableFeature_v4_t(_PrintableStructure): _fields_ = [ ("featureCode", _nvmlGridLicenseFeatureCode_t), ("featureState", c_uint), ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), ("featureEnabled", c_uint), ("licenseExpiry", c_nvmlGridLicenseExpiry_t), ] class c_nvmlGridLicensableFeatures_v4_t(_PrintableStructure): _fields_ = [ ("isGridLicenseSupported", c_int), ("licensableFeaturesCount", c_uint), ( "gridLicensableFeatures", c_nvmlGridLicensableFeature_v4_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, ), ] class c_nvmlGridLicensableFeature_v3_t(_PrintableStructure): _fields_ = [ ("featureCode", _nvmlGridLicenseFeatureCode_t), ("featureState", c_uint), ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), ("featureEnabled", c_uint), ] class c_nvmlGridLicensableFeatures_v3_t(_PrintableStructure): _fields_ = [ ("isGridLicenseSupported", c_int), ("licensableFeaturesCount", c_uint), ( "gridLicensableFeatures", c_nvmlGridLicensableFeature_v3_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, ), ] class c_nvmlGridLicensableFeature_v2_t(_PrintableStructure): _fields_ = [ ("featureCode", _nvmlGridLicenseFeatureCode_t), ("featureState", c_uint), ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), ("productName", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), ] class c_nvmlGridLicensableFeatures_v2_t(_PrintableStructure): _fields_ = [ ("isGridLicenseSupported", c_int), ("licensableFeaturesCount", c_uint), ( "gridLicensableFeatures", c_nvmlGridLicensableFeature_v2_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, ), ] class c_nvmlGridLicensableFeature_t(_PrintableStructure): _fields_ = [ ("featureCode", _nvmlGridLicenseFeatureCode_t), ("featureState", c_uint), ("licenseInfo", c_char * NVML_GRID_LICENSE_BUFFER_SIZE), ] class c_nvmlGridLicensableFeatures_t(_PrintableStructure): _fields_ = [ ("isGridLicenseSupported", c_int), ("licensableFeaturesCount", c_uint), ( "gridLicensableFeatures", c_nvmlGridLicensableFeature_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT, ), ] class c_nvmlMarginTemperature_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("marginTemperature", c_int), ] nvmlMarginTemperature_v1 = 0x1000008 ## Event structures class struct_c_nvmlEventSet_t(Structure): pass # opaque handle c_nvmlEventSet_t = POINTER(struct_c_nvmlEventSet_t) nvmlEventTypeSingleBitEccError = 0x0000000000000001 nvmlEventTypeDoubleBitEccError = 0x0000000000000002 nvmlEventTypePState = 0x0000000000000004 nvmlEventTypeXidCriticalError = 0x0000000000000008 nvmlEventTypeClock = 0x0000000000000010 nvmlEventTypePowerSourceChange = 0x0000000000000080 nvmlEventMigConfigChange = 0x0000000000000100 nvmlEventTypeSingleBitEccErrorStorm = 0x0000000000000200 nvmlEventTypeDramRetirementEvent = 0x0000000000000400 nvmlEventTypeDramRetirementFailure = 0x0000000000000800 nvmlEventTypeNonFatalPoisonError = 0x0000000000001000 nvmlEventTypeFatalPoisonError = 0x0000000000002000 nvmlEventTypeGpuUnavailableError = 0x0000000000004000 nvmlEventTypeGpuRecoveryAction = 0x0000000000008000 nvmlEventTypeNone = 0x0000000000000000 nvmlEventTypeAll = ( nvmlEventTypeNone | nvmlEventTypeSingleBitEccError | nvmlEventTypeDoubleBitEccError | nvmlEventTypePState | nvmlEventTypeClock | nvmlEventTypePowerSourceChange | nvmlEventTypeXidCriticalError | nvmlEventMigConfigChange | nvmlEventTypeSingleBitEccErrorStorm | nvmlEventTypeDramRetirementEvent | nvmlEventTypeDramRetirementFailure | nvmlEventTypeNonFatalPoisonError | nvmlEventTypeFatalPoisonError | nvmlEventTypeGpuUnavailableError | nvmlEventTypeGpuRecoveryAction ) ## Clock Event Reasons defines nvmlClocksEventReasonGpuIdle = 0x0000000000000001 nvmlClocksEventReasonApplicationsClocksSetting = 0x0000000000000002 nvmlClocksEventReasonUserDefinedClocks = nvmlClocksEventReasonApplicationsClocksSetting # deprecated, use nvmlClocksEventReasonApplicationsClocksSetting nvmlClocksEventReasonSwPowerCap = 0x0000000000000004 nvmlClocksEventReasonHwSlowdown = 0x0000000000000008 nvmlClocksEventReasonSyncBoost = 0x0000000000000010 nvmlClocksEventReasonSwThermalSlowdown = 0x0000000000000020 nvmlClocksEventReasonHwThermalSlowdown = 0x0000000000000040 nvmlClocksEventReasonHwPowerBrakeSlowdown = 0x0000000000000080 nvmlClocksEventReasonDisplayClockSetting = 0x0000000000000100 nvmlClocksEventReasonNone = 0x0000000000000000 nvmlClocksEventReasonAll = ( nvmlClocksEventReasonNone | nvmlClocksEventReasonGpuIdle | nvmlClocksEventReasonApplicationsClocksSetting | nvmlClocksEventReasonSwPowerCap | nvmlClocksEventReasonHwSlowdown | nvmlClocksEventReasonSyncBoost | nvmlClocksEventReasonSwThermalSlowdown | nvmlClocksEventReasonHwThermalSlowdown | nvmlClocksEventReasonHwPowerBrakeSlowdown | nvmlClocksEventReasonDisplayClockSetting ) ## Following have been deprecated nvmlClocksThrottleReasonGpuIdle = 0x0000000000000001 nvmlClocksThrottleReasonApplicationsClocksSetting = 0x0000000000000002 nvmlClocksThrottleReasonUserDefinedClocks = nvmlClocksThrottleReasonApplicationsClocksSetting # deprecated, use nvmlClocksThrottleReasonApplicationsClocksSetting nvmlClocksThrottleReasonSwPowerCap = 0x0000000000000004 nvmlClocksThrottleReasonHwSlowdown = 0x0000000000000008 nvmlClocksThrottleReasonSyncBoost = 0x0000000000000010 nvmlClocksThrottleReasonSwThermalSlowdown = 0x0000000000000020 nvmlClocksThrottleReasonHwThermalSlowdown = 0x0000000000000040 nvmlClocksThrottleReasonHwPowerBrakeSlowdown = 0x0000000000000080 nvmlClocksThrottleReasonDisplayClockSetting = 0x0000000000000100 nvmlClocksThrottleReasonNone = 0x0000000000000000 nvmlClocksThrottleReasonAll = ( nvmlClocksThrottleReasonNone | nvmlClocksThrottleReasonGpuIdle | nvmlClocksThrottleReasonApplicationsClocksSetting | nvmlClocksThrottleReasonSwPowerCap | nvmlClocksThrottleReasonHwSlowdown | nvmlClocksThrottleReasonSyncBoost | nvmlClocksThrottleReasonSwThermalSlowdown | nvmlClocksThrottleReasonHwThermalSlowdown | nvmlClocksThrottleReasonHwPowerBrakeSlowdown | nvmlClocksThrottleReasonDisplayClockSetting ) class c_nvmlEventData_t(_PrintableStructure): _fields_ = [ ("device", c_nvmlDevice_t), ("eventType", c_ulonglong), ("eventData", c_ulonglong), ("gpuInstanceId", c_uint), ("computeInstanceId", c_uint), ] _fmt_ = {"eventType": "0x%08X"} class c_nvmlAccountingStats_t(_PrintableStructure): _fields_ = [ ("gpuUtilization", c_uint), ("memoryUtilization", c_uint), ("maxMemoryUsage", c_ulonglong), ("time", c_ulonglong), ("startTime", c_ulonglong), ("isRunning", c_uint), ("reserved", c_uint * 5), ] class c_nvmlVgpuVersion_t(Structure): _fields_ = [("minVersion", c_uint), ("maxVersion", c_uint)] class c_nvmlVgpuMetadata_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("revision", c_uint), ("guestInfoState", _nvmlVgpuGuestInfoState_t), ("guestDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), ("reserved", c_uint * 6), ("vgpuVirtualizationCaps", c_uint), ("guestVgpuVersion", c_uint), ("opaqueDataSize", c_uint), ("opaqueData", c_char * NVML_VGPU_METADATA_OPAQUE_DATA_SIZE), ] class c_nvmlVgpuPgpuMetadata_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("revision", c_uint), ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), ("pgpuVirtualizationCaps", c_uint), ("reserved", c_uint * 5), ("hostSupportedVgpuRange", c_nvmlVgpuVersion_t), ("opaqueDataSize", c_uint), ("opaqueData", c_char * NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE), ] class c_nvmlVgpuPgpuCompatibility_t(Structure): _fields_ = [ ("vgpuVmCompatibility", _nvmlVgpuVmCompatibility_t), ("compatibilityLimitCode", _nvmlVgpuPgpuCompatibilityLimitCode_t), ] ## vGPU scheduler policy defines NVML_VGPU_SCHEDULER_POLICY_UNKNOWN = 0 NVML_VGPU_SCHEDULER_POLICY_BEST_EFFORT = 1 NVML_VGPU_SCHEDULER_POLICY_EQUAL_SHARE = 2 NVML_VGPU_SCHEDULER_POLICY_FIXED_SHARE = 3 ## Supported vGPU scheduler policy count NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT = 3 NVML_SCHEDULER_SW_MAX_LOG_ENTRIES = 200 NVML_VGPU_SCHEDULER_ARR_DEFAULT = 0 NVML_VGPU_SCHEDULER_ARR_DISABLE = 1 NVML_VGPU_SCHEDULER_ARR_ENABLE = 2 class c_nvmlVgpuSchedDataWithARR_t(_PrintableStructure): _fields_ = [ ("avgFactor", c_uint), ("timeslice", c_uint), ] class c_nvmlVgpuSchedData_t(_PrintableStructure): _fields_ = [ ("timeslice", c_uint), ] class c_nvmlVgpuSchedulerParams_t(Union): _fields_ = [ ("vgpuSchedDataWithARR", c_nvmlVgpuSchedDataWithARR_t), ("vgpuSchedData", c_nvmlVgpuSchedData_t), ] class c_nvmlVgpuSchedulerLogEntry_t(_PrintableStructure): _fields_ = [ ("timestamp", c_ulonglong), ("timeRunTotal", c_ulonglong), ("timeRun", c_ulonglong), ("swRunlistId", c_uint), ("targetTimeSlice", c_ulonglong), ("cumulativePreemptionTime", c_ulonglong), ] class c_nvmlVgpuSchedulerLog_t(_PrintableStructure): _fields_ = [ ("engineId", c_uint), ("schedulerPolicy", c_uint), ("arrMode", c_uint), ("schedulerParams", c_nvmlVgpuSchedulerParams_t), ("entriesCount", c_uint), ( "logEntries", c_nvmlVgpuSchedulerLogEntry_t * NVML_SCHEDULER_SW_MAX_LOG_ENTRIES, ), ] class c_nvmlVgpuSchedulerGetState_t(_PrintableStructure): _fields_ = [ ("schedulerPolicy", c_uint), ("arrMode", c_uint), ("schedulerParams", c_nvmlVgpuSchedulerParams_t), ] class c_nvmlVgpuSchedSetDataWithARR_t(_PrintableStructure): _fields_ = [ ("avgFactor", c_uint), ("frequency", c_uint), ] class c_nvmlVgpuSchedSetData_t(_PrintableStructure): _fields_ = [ ("timeslice", c_uint), ] class c_nvmlVgpuSchedulerSetParams_t(Union): _fields_ = [ ("vgpuSchedDataWithARR", c_nvmlVgpuSchedSetDataWithARR_t), ("vgpuSchedData", c_nvmlVgpuSchedSetData_t), ] class c_nvmlVgpuSchedulerSetState_t(_PrintableStructure): _fields_ = [ ("schedulerPolicy", c_uint), ("enableARRMode", c_uint), ("schedulerParams", c_nvmlVgpuSchedulerSetParams_t), ] class c_nvmlVgpuSchedulerCapabilities_t(_PrintableStructure): _fields_ = [ ("supportedSchedulers", c_uint * NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT), ("maxTimeslice", c_uint), ("minTimeslice", c_uint), ("isArrModeSupported", c_uint), ("maxFrequencyForARR", c_uint), ("minFrequencyForARR", c_uint), ("maxAvgFactorForARR", c_uint), ("minAvgFactorForARR", c_uint), ] class c_nvmlFBCStats_t(Structure): _fields_ = [ ("sessionsCount", c_uint), ("averageFPS", c_uint), ("averageLatency", c_uint), ] class c_nvmlFBCSession_t(_PrintableStructure): _fields_ = [ ("sessionId", c_uint), ("pid", c_uint), ("vgpuInstance", _nvmlVgpuInstance_t), ("displayOrdinal", c_uint), ("sessionType", c_uint), ("sessionFlags", c_uint), ("hMaxResolution", c_uint), ("vMaxResolution", c_uint), ("hResolution", c_uint), ("vResolution", c_uint), ("averageFPS", c_uint), ("averageLatency", c_uint), ] NVML_DEVICE_MIG_DISABLE = 0x0 NVML_DEVICE_MIG_ENABLE = 0x1 NVML_GPU_INSTANCE_PROFILE_1_SLICE = 0x0 NVML_GPU_INSTANCE_PROFILE_2_SLICE = 0x1 NVML_GPU_INSTANCE_PROFILE_3_SLICE = 0x2 NVML_GPU_INSTANCE_PROFILE_4_SLICE = 0x3 NVML_GPU_INSTANCE_PROFILE_7_SLICE = 0x4 NVML_GPU_INSTANCE_PROFILE_8_SLICE = 0x5 NVML_GPU_INSTANCE_PROFILE_6_SLICE = 0x6 NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 NVML_GPU_INSTANCE_PROFILE_2_SLICE_REV1 = 0x8 NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV2 = 0x9 NVML_GPU_INSTANCE_PROFILE_1_SLICE_GFX = 0xA NVML_GPU_INSTANCE_PROFILE_2_SLICE_GFX = 0xB NVML_GPU_INSTANCE_PROFILE_4_SLICE_GFX = 0xC NVML_GPU_INSTANCE_PROFILE_COUNT = 0xD class c_nvmlGpuInstancePlacement_t(Structure): _fields_ = [("start", c_uint), ("size", c_uint)] class c_nvmlGpuInstanceProfileInfo_t(Structure): _fields_ = [ ("id", c_uint), ("isP2pSupported", c_uint), ("sliceCount", c_uint), ("instanceCount", c_uint), ("multiprocessorCount", c_uint), ("copyEngineCount", c_uint), ("decoderCount", c_uint), ("encoderCount", c_uint), ("jpegCount", c_uint), ("ofaCount", c_uint), ("memorySizeMB", c_ulonglong), ] nvmlGpuInstanceProfileInfo_v2 = 0x02000098 class c_nvmlGpuInstanceProfileInfo_v2_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("id", c_uint), ("isP2pSupported", c_uint), ("sliceCount", c_uint), ("instanceCount", c_uint), ("multiprocessorCount", c_uint), ("copyEngineCount", c_uint), ("decoderCount", c_uint), ("encoderCount", c_uint), ("jpegCount", c_uint), ("ofaCount", c_uint), ("memorySizeMB", c_ulonglong), ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE), ] def __init__(self): super(c_nvmlGpuInstanceProfileInfo_v2_t, self).__init__( version=nvmlGpuInstanceProfileInfo_v2 ) class c_nvmlGpuInstanceInfo_t(Structure): _fields_ = [ ("device", c_nvmlDevice_t), ("id", c_uint), ("profileId", c_uint), ("placement", c_nvmlGpuInstancePlacement_t), ] class struct_c_nvmlGpuInstance_t(Structure): pass # opaque handle c_nvmlGpuInstance_t = POINTER(struct_c_nvmlGpuInstance_t) NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE = 0x0 NVML_COMPUTE_INSTANCE_PROFILE_2_SLICE = 0x1 NVML_COMPUTE_INSTANCE_PROFILE_3_SLICE = 0x2 NVML_COMPUTE_INSTANCE_PROFILE_4_SLICE = 0x3 NVML_COMPUTE_INSTANCE_PROFILE_7_SLICE = 0x4 NVML_COMPUTE_INSTANCE_PROFILE_8_SLICE = 0x5 NVML_COMPUTE_INSTANCE_PROFILE_6_SLICE = 0x6 NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 NVML_COMPUTE_INSTANCE_PROFILE_COUNT = 0x8 NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED = 0x0 NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_COUNT = 0x1 class c_nvmlComputeInstancePlacement_t(Structure): _fields_ = [("start", c_uint), ("size", c_uint)] class c_nvmlComputeInstanceProfileInfo_t(Structure): _fields_ = [ ("id", c_uint), ("sliceCount", c_uint), ("instanceCount", c_uint), ("multiprocessorCount", c_uint), ("sharedCopyEngineCount", c_uint), ("sharedDecoderCount", c_uint), ("sharedEncoderCount", c_uint), ("sharedJpegCount", c_uint), ("sharedOfaCount", c_uint), ] nvmlComputeInstanceProfileInfo_v2 = 0x02000088 class c_nvmlComputeInstanceProfileInfo_v2_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("id", c_uint), ("sliceCount", c_uint), ("instanceCount", c_uint), ("multiprocessorCount", c_uint), ("sharedCopyEngineCount", c_uint), ("sharedDecoderCount", c_uint), ("sharedEncoderCount", c_uint), ("sharedJpegCount", c_uint), ("sharedOfaCount", c_uint), ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE), ] def __init__(self): super(c_nvmlComputeInstanceProfileInfo_v2_t, self).__init__( version=nvmlComputeInstanceProfileInfo_v2 ) class c_nvmlComputeInstanceInfo_t(Structure): _fields_ = [ ("device", c_nvmlDevice_t), ("gpuInstance", c_nvmlGpuInstance_t), ("id", c_uint), ("profileId", c_uint), ("placement", c_nvmlComputeInstancePlacement_t), ] NVML_MAX_GPU_UTILIZATIONS = 8 NVML_GPU_UTILIZATION_DOMAIN_GPU = 0 NVML_GPU_UTILIZATION_DOMAIN_FB = 1 NVML_GPU_UTILIZATION_DOMAIN_VID = 2 NVML_GPU_UTILIZATION_DOMAIN_BUS = 3 class c_nvmlGpuDynamicPstatesUtilization_t(Structure): _fields_ = [ ("bIsPresent", c_uint, 1), ("percentage", c_uint), ("incThreshold", c_uint), ("decThreshold", c_uint), ] class c_nvmlGpuDynamicPstatesInfo_t(Structure): _fields_ = [ ("flags", c_uint), ( "utilization", c_nvmlGpuDynamicPstatesUtilization_t * NVML_MAX_GPU_UTILIZATIONS, ), ] NVML_MAX_THERMAL_SENSORS_PER_GPU = 3 NVML_THERMAL_TARGET_NONE = 0 NVML_THERMAL_TARGET_GPU = 1 NVML_THERMAL_TARGET_MEMORY = 2 NVML_THERMAL_TARGET_POWER_SUPPLY = 4 NVML_THERMAL_TARGET_BOARD = 8 NVML_THERMAL_TARGET_VCD_BOARD = 9 NVML_THERMAL_TARGET_VCD_INLET = 10 NVML_THERMAL_TARGET_VCD_OUTLET = 11 NVML_THERMAL_TARGET_ALL = 15 NVML_THERMAL_TARGET_UNKNOWN = -1 NVML_THERMAL_CONTROLLER_NONE = 0 NVML_THERMAL_CONTROLLER_GPU_INTERNAL = 1 NVML_THERMAL_CONTROLLER_ADM1032 = 2 NVML_THERMAL_CONTROLLER_ADT7461 = 3 NVML_THERMAL_CONTROLLER_MAX6649 = 4 NVML_THERMAL_CONTROLLER_MAX1617 = 5 NVML_THERMAL_CONTROLLER_LM99 = 6 NVML_THERMAL_CONTROLLER_LM89 = 7 NVML_THERMAL_CONTROLLER_LM64 = 8 NVML_THERMAL_CONTROLLER_G781 = 9 NVML_THERMAL_CONTROLLER_ADT7473 = 10 NVML_THERMAL_CONTROLLER_SBMAX6649 = 11 NVML_THERMAL_CONTROLLER_VBIOSEVT = 12 NVML_THERMAL_CONTROLLER_OS = 13 NVML_THERMAL_CONTROLLER_NVSYSCON_CANOAS = 14 NVML_THERMAL_CONTROLLER_NVSYSCON_E551 = 15 NVML_THERMAL_CONTROLLER_MAX6649R = 16 NVML_THERMAL_CONTROLLER_ADT7473S = 17 NVML_THERMAL_CONTROLLER_UNKNOWN = -1 class c_nvmlGpuThermalSensor_t(Structure): _fields_ = [ ("controller", c_int), ("defaultMinTemp", c_int), ("defaultMaxTemp", c_int), ("currentTemp", c_int), ("target", c_int), ] class c_nvmlGpuThermalSettings_t(Structure): _fields_ = [ ("count", c_uint), ("sensor", c_nvmlGpuThermalSensor_t * NVML_MAX_THERMAL_SENSORS_PER_GPU), ] _nvmlCoolerControl_t = c_uint NVML_THERMAL_COOLER_SIGNAL_NONE = 0 NVML_THERMAL_COOLER_SIGNAL_TOGGLE = 1 NVML_THERMAL_COOLER_SIGNAL_VARIABLE = 2 NVML_THERMAL_COOLER_SIGNAL_COUNT = 3 _nvmlCoolerTarget_t = c_uint NVML_THERMAL_COOLER_TARGET_NONE = 1 << 0 NVML_THERMAL_COOLER_TARGET_GPU = 1 << 1 NVML_THERMAL_COOLER_TARGET_MEMORY = 1 << 2 NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY = 1 << 3 NVML_THERMAL_COOLER_TARGET_GPU_RELATED = ( NVML_THERMAL_COOLER_TARGET_GPU | NVML_THERMAL_COOLER_TARGET_MEMORY | NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY ) class c_nvmlCoolerInfo_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("index", c_uint), ("coolerControlType", _nvmlCoolerControl_t), ("coolerTarget", _nvmlCoolerTarget_t), ] nvmlCoolerInfo_v1 = 0x1000010 def nvmlDeviceGetCoolerInfo(handle): c_coolerInfo = c_nvmlCoolerInfo_t() c_coolerInfo.version = nvmlCoolerInfo_v1 c_coolerInfo.index = 0 fn = _nvmlGetFunctionPointer("nvmlDeviceGetCoolerInfo") ret = fn(handle, byref(c_coolerInfo)) _nvmlCheckReturn(ret) return [c_coolerInfo.coolerControlType, c_coolerInfo.coolerTarget] class struct_c_nvmlComputeInstance_t(Structure): pass # opaque handle c_nvmlComputeInstance_t = POINTER(struct_c_nvmlComputeInstance_t) class c_nvmlDeviceAttributes(Structure): _fields_ = [ ("multiprocessorCount", c_uint), ("sharedCopyEngineCount", c_uint), ("sharedDecoderCount", c_uint), ("sharedEncoderCount", c_uint), ("sharedJpegCount", c_uint), ("sharedOfaCount", c_uint), ("gpuInstanceSliceCount", c_uint), ("computeInstanceSliceCount", c_uint), ("memorySizeMB", c_ulonglong), ] class c_nvmlRowRemapperHistogramValues(Structure): _fields_ = [ ("max", c_uint), ("high", c_uint), ("partial", c_uint), ("low", c_uint), ("none", c_uint), ] NVML_GPU_CERT_CHAIN_SIZE = 0x1000 NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE = 0x1400 NVML_CC_GPU_CEC_NONCE_SIZE = 0x20 NVML_CC_GPU_ATTESTATION_REPORT_SIZE = 0x2000 NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE = 0x1000 NVML_CC_CEC_ATTESTATION_REPORT_NOT_PRESENT = 0 NVML_CC_CEC_ATTESTATION_REPORT_PRESENT = 1 class c_nvmlConfComputeSystemState_t(Structure): _fields_ = [ ("environment", c_uint), ("ccFeature", c_uint), ("devToolsMode", c_uint), ] nvmlSystemConfComputeSettings_v1 = 0x1000014 class c_nvmlSystemConfComputeSettings_v1_t(Structure): _fields_ = [ ("version", c_uint), ("environment", c_uint), ("ccFeature", c_uint), ("devToolsMode", c_uint), ("multiGpuMode", c_uint), ] def __init__(self): super(c_nvmlSystemConfComputeSettings_v1_t, self).__init__( version=nvmlSystemConfComputeSettings_v1 ) class c_nvmlConfComputeSystemCaps_t(Structure): _fields_ = [ ("cpuCaps", c_uint), ("gpusCaps", c_uint), ] class c_nvmlConfComputeMemSizeInfo_t(Structure): _fields_ = [ ("protectedMemSizeKib", c_ulonglong), ("unprotectedMemSizeKib", c_ulonglong), ] class c_nvmlConfComputeGpuCertificate_t(Structure): _fields_ = [ ("certChainSize", c_uint), ("attestationCertChainSize", c_uint), ("certChain", c_uint8 * NVML_GPU_CERT_CHAIN_SIZE), ("attestationCertChain", c_uint8 * NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE), ] class c_nvmlConfComputeGpuAttestationReport_t(Structure): _fields_ = [ ("isCecAttestationReportPresent", c_uint), ("attestationReportSize", c_uint), ("cecAttestationReportSize", c_uint), ("nonce", c_uint8 * NVML_CC_GPU_CEC_NONCE_SIZE), ("attestationReport", c_uint8 * NVML_CC_GPU_ATTESTATION_REPORT_SIZE), ("cecAttestationReport", c_uint8 * NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE), ] class c_nvmlConfComputeSetKeyRotationThresholdInfo_t(Structure): _fields_ = [ ("version", c_uint), ("maxAttackerAdvantage", c_ulong), ] ConfComputeSetKeyRotationThresholdInfo_v1 = 0x1000010 class c_nvmlConfComputeGetKeyRotationThresholdInfo_t(Structure): _fields_ = [ ("version", c_uint), ("attackerAdvantage", c_ulong), ] ConfComputeGetKeyRotationThresholdInfo_v1 = 0x1000010 ## string/bytes conversion for ease of use def convertStrBytes(func): """ In python 3, strings are unicode instead of bytes, and need to be converted for ctypes Args from caller: (1, 'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF>) Args passed to function: (1, b'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF)> ---- Returned from function: b'returned string' Returned to caller: 'returned string' """ @wraps(func) def wrapper(*args, **kwargs): # encoding a str returns bytes in python 2 and 3 args = [arg.encode() if isinstance(arg, str) else arg for arg in args] res = func(*args, **kwargs) # In python 2, str and bytes are the same # In python 3, str is unicode and should be decoded. # Ctypes handles most conversions, this only effects c_char and char arrays. if isinstance(res, bytes): if isinstance(res, str): return res return res.decode() return res if sys.version_info >= (3,): return wrapper return func def throwOnVersionMismatch(func): @wraps(func) def wrapper(*args, **kwargs): try: return func(*args, **kwargs) except NVMLError_FunctionNotFound: raise NVMLLibraryMismatchError( "Unversioned function called and the " "pyNVML version does not match the NVML lib version. " "Either use matching pyNVML and NVML lib versions or " "use a versioned function such as " + func.__name__ + "_v2" ) return wrapper ## C function wrappers ## def nvmlInitWithFlags(flags): _LoadNvmlLibrary() # # Initialize the library # fn = _nvmlGetFunctionPointer("nvmlInitWithFlags") ret = fn(flags) _nvmlCheckReturn(ret) # Atomically update refcount global _nvmlLib_refcount libLoadLock.acquire() _nvmlLib_refcount += 1 libLoadLock.release() return None def nvmlInit(): nvmlInitWithFlags(0) return None def _LoadNvmlLibrary(): """ Load the library if it isn't loaded already """ global nvmlLib if nvmlLib is None: # lock to ensure only one caller loads the library libLoadLock.acquire() try: # ensure the library still isn't loaded if nvmlLib is None: try: if sys.platform[:3] == "win": # cdecl calling convention try: # Check for nvml.dll in System32 first for DCH drivers nvmlLib = CDLL( os.path.join( os.getenv("WINDIR", "C:/Windows"), "System32/nvml.dll", ) ) except OSError as ose: # If nvml.dll is not found in System32, it should be in ProgramFiles # load nvml.dll from %ProgramFiles%/NVIDIA Corporation/NVSMI/nvml.dll nvmlLib = CDLL( os.path.join( os.getenv("ProgramFiles", "C:/Program Files"), "NVIDIA Corporation/NVSMI/nvml.dll", ) ) else: # assume linux nvmlLib = CDLL("libnvidia-ml.so.1") except OSError as ose: _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) if nvmlLib is None: _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) finally: # lock is always freed libLoadLock.release() def nvmlShutdown(): # # Leave the library loaded, but shutdown the interface # fn = _nvmlGetFunctionPointer("nvmlShutdown") ret = fn() _nvmlCheckReturn(ret) # Atomically update refcount global _nvmlLib_refcount libLoadLock.acquire() if 0 < _nvmlLib_refcount: _nvmlLib_refcount -= 1 libLoadLock.release() return None # Added in 2.285 @convertStrBytes def nvmlErrorString(result): fn = _nvmlGetFunctionPointer("nvmlErrorString") fn.restype = c_char_p # otherwise return is an int ret = fn(result) return ret # Added in 2.285 @convertStrBytes def nvmlSystemGetNVMLVersion(): c_version = create_string_buffer(NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlSystemGetNVMLVersion") ret = fn(c_version, c_uint(NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE)) _nvmlCheckReturn(ret) return c_version.value def nvmlSystemGetCudaDriverVersion(): c_cuda_version = c_int() fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion") ret = fn(byref(c_cuda_version)) _nvmlCheckReturn(ret) return c_cuda_version.value def nvmlSystemGetCudaDriverVersion_v2(): c_cuda_version = c_int() fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion_v2") ret = fn(byref(c_cuda_version)) _nvmlCheckReturn(ret) return c_cuda_version.value # Added in 2.285 @convertStrBytes def nvmlSystemGetProcessName(pid): c_name = create_string_buffer(1024) fn = _nvmlGetFunctionPointer("nvmlSystemGetProcessName") ret = fn(c_uint(pid), c_name, c_uint(1024)) _nvmlCheckReturn(ret) return c_name.value @convertStrBytes def nvmlSystemGetDriverVersion(): c_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverVersion") ret = fn(c_version, c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)) _nvmlCheckReturn(ret) return c_version.value # Added in 2.285 def nvmlSystemGetHicVersion(): c_count = c_uint(0) hics = None fn = _nvmlGetFunctionPointer("nvmlSystemGetHicVersion") # get the count ret = fn(byref(c_count), None) # this should only fail with insufficient size if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE): raise NVMLError(ret) # If there are no hics if c_count.value == 0: return [] hic_array = c_nvmlHwbcEntry_t * c_count.value hics = hic_array() ret = fn(byref(c_count), hics) _nvmlCheckReturn(ret) return hics def nvmlSystemGetDriverBranch(): c_branchInfo = c_nvmlSystemDriverBranchInfo_v1_t(0) c_branchInfo.version = SystemDriverBranchInfo_v1 fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverBranch") ret = fn(byref(c_branchInfo), c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)) _nvmlCheckReturn(ret) return c_branchInfo ## Unit get functions def nvmlUnitGetCount(): c_count = c_uint() fn = _nvmlGetFunctionPointer("nvmlUnitGetCount") ret = fn(byref(c_count)) _nvmlCheckReturn(ret) return c_count.value def nvmlUnitGetHandleByIndex(index): c_index = c_uint(index) unit = c_nvmlUnit_t() fn = _nvmlGetFunctionPointer("nvmlUnitGetHandleByIndex") ret = fn(c_index, byref(unit)) _nvmlCheckReturn(ret) return unit def nvmlUnitGetUnitInfo(unit): c_info = c_nvmlUnitInfo_t() fn = _nvmlGetFunctionPointer("nvmlUnitGetUnitInfo") ret = fn(unit, byref(c_info)) _nvmlCheckReturn(ret) return c_info def nvmlUnitGetLedState(unit): c_state = c_nvmlLedState_t() fn = _nvmlGetFunctionPointer("nvmlUnitGetLedState") ret = fn(unit, byref(c_state)) _nvmlCheckReturn(ret) return c_state def nvmlUnitGetPsuInfo(unit): c_info = c_nvmlPSUInfo_t() fn = _nvmlGetFunctionPointer("nvmlUnitGetPsuInfo") ret = fn(unit, byref(c_info)) _nvmlCheckReturn(ret) return c_info def nvmlUnitGetTemperature(unit, type): c_temp = c_uint() fn = _nvmlGetFunctionPointer("nvmlUnitGetTemperature") ret = fn(unit, c_uint(type), byref(c_temp)) _nvmlCheckReturn(ret) return c_temp.value def nvmlUnitGetFanSpeedInfo(unit): c_speeds = c_nvmlUnitFanSpeeds_t() fn = _nvmlGetFunctionPointer("nvmlUnitGetFanSpeedInfo") ret = fn(unit, byref(c_speeds)) _nvmlCheckReturn(ret) return c_speeds # added to API def nvmlUnitGetDeviceCount(unit): c_count = c_uint(0) # query the unit to determine device count fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices") ret = fn(unit, byref(c_count), None) if ret == NVML_ERROR_INSUFFICIENT_SIZE: ret = NVML_SUCCESS _nvmlCheckReturn(ret) return c_count.value def nvmlUnitGetDevices(unit): c_count = c_uint(nvmlUnitGetDeviceCount(unit)) device_array = c_nvmlDevice_t * c_count.value c_devices = device_array() fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices") ret = fn(unit, byref(c_count), c_devices) _nvmlCheckReturn(ret) return c_devices ## Device get functions def nvmlDeviceGetCount(): c_count = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetCount_v2") ret = fn(byref(c_count)) _nvmlCheckReturn(ret) return c_count.value def nvmlDeviceGetHandleByIndex(index): c_index = c_uint(index) device = c_nvmlDevice_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByIndex_v2") ret = fn(c_index, byref(device)) _nvmlCheckReturn(ret) return device @convertStrBytes def nvmlDeviceGetHandleBySerial(serial): c_serial = c_char_p(serial) device = c_nvmlDevice_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleBySerial") ret = fn(c_serial, byref(device)) _nvmlCheckReturn(ret) return device @convertStrBytes def nvmlDeviceGetHandleByUUID(uuid): c_uuid = c_char_p(uuid) device = c_nvmlDevice_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByUUID") ret = fn(c_uuid, byref(device)) _nvmlCheckReturn(ret) return device @convertStrBytes def nvmlDeviceGetHandleByPciBusId(pciBusId): c_busId = c_char_p(pciBusId) device = c_nvmlDevice_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByPciBusId_v2") ret = fn(c_busId, byref(device)) _nvmlCheckReturn(ret) return device @convertStrBytes def nvmlDeviceGetName(handle): c_name = create_string_buffer(NVML_DEVICE_NAME_V2_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlDeviceGetName") ret = fn(handle, c_name, c_uint(NVML_DEVICE_NAME_V2_BUFFER_SIZE)) _nvmlCheckReturn(ret) return c_name.value class c_nvmlDevicePerfModes_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("str", c_char * NVML_PERF_MODES_BUFFER_SIZE), ] nvmlDevicePerfModes_v1 = 0x1000804 @convertStrBytes def nvmlDeviceGetPerformanceModes(handle): perfModes = c_nvmlDevicePerfModes_v1_t() perfModes.version = nvmlDevicePerfModes_v1 fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceModes") ret = fn(handle, byref(perfModes)) _nvmlCheckReturn(ret) return perfModes.str class c_nvmlDeviceCurrentClockFreqs_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("str", c_char * NVML_PERF_MODES_BUFFER_SIZE), ] nvmlDeviceCurrentClockFreqs_v1 = 0x1000804 @convertStrBytes def nvmlDeviceGetCurrentClockFreqs(handle): currentClockFreqs = c_nvmlDeviceCurrentClockFreqs_v1_t() currentClockFreqs.version = nvmlDeviceCurrentClockFreqs_v1 fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClockFreqs") ret = fn(handle, byref(currentClockFreqs)) _nvmlCheckReturn(ret) return currentClockFreqs.str def nvmlDeviceGetBoardId(handle): c_id = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardId") ret = fn(handle, byref(c_id)) _nvmlCheckReturn(ret) return c_id.value def nvmlDeviceGetMultiGpuBoard(handle): c_multiGpu = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMultiGpuBoard") ret = fn(handle, byref(c_multiGpu)) _nvmlCheckReturn(ret) return c_multiGpu.value def nvmlDeviceGetBrand(handle): c_type = _nvmlBrandType_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBrand") ret = fn(handle, byref(c_type)) _nvmlCheckReturn(ret) return c_type.value def nvmlDeviceGetC2cModeInfoV1(handle): c_info = c_nvmlC2cModeInfo_v1_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetC2cModeInfoV") ret = fn(handle, byref(c_info)) _nvmlCheckReturn(ret) return c_info def nvmlDeviceGetC2cModeInfoV(handle): return nvmlDeviceGetC2cModeInfoV1(handle) @convertStrBytes def nvmlDeviceGetBoardPartNumber(handle): c_part_number = create_string_buffer(NVML_DEVICE_PART_NUMBER_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardPartNumber") ret = fn(handle, c_part_number, c_uint(NVML_DEVICE_PART_NUMBER_BUFFER_SIZE)) _nvmlCheckReturn(ret) return c_part_number.value @convertStrBytes def nvmlDeviceGetSerial(handle): c_serial = create_string_buffer(NVML_DEVICE_SERIAL_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlDeviceGetSerial") ret = fn(handle, c_serial, c_uint(NVML_DEVICE_SERIAL_BUFFER_SIZE)) _nvmlCheckReturn(ret) return c_serial.value def nvmlDeviceGetModuleId(handle, moduleId=c_uint()): isReference = type(moduleId) is not c_uint moduleIdRef = moduleId if isReference else byref(moduleId) fn = _nvmlGetFunctionPointer("nvmlDeviceGetModuleId") ret = fn(handle, moduleIdRef) if isReference: return ret else: _nvmlCheckReturn(ret) return moduleId.value def nvmlDeviceGetMemoryAffinity(handle, nodeSetSize, scope): affinity_array = c_ulonglong * nodeSetSize c_affinity = affinity_array() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryAffinity") ret = fn(handle, nodeSetSize, byref(c_affinity), _nvmlAffinityScope_t(scope)) _nvmlCheckReturn(ret) return c_affinity def nvmlDeviceGetCpuAffinityWithinScope(handle, cpuSetSize, scope): affinity_array = c_ulonglong * cpuSetSize c_affinity = affinity_array() fn = _nvmlGetFunctionPointer("nvmlDeviceGetCpuAffinityWithinScope") ret = fn(handle, cpuSetSize, byref(c_affinity), _nvmlAffinityScope_t(scope)) _nvmlCheckReturn(ret) return c_affinity def nvmlDeviceGetCpuAffinity(handle, cpuSetSize): affinity_array = c_ulonglong * cpuSetSize c_affinity = affinity_array() fn = _nvmlGetFunctionPointer("nvmlDeviceGetCpuAffinity") ret = fn(handle, cpuSetSize, byref(c_affinity)) _nvmlCheckReturn(ret) return c_affinity def nvmlDeviceSetCpuAffinity(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceSetCpuAffinity") ret = fn(handle) _nvmlCheckReturn(ret) return None def nvmlDeviceClearCpuAffinity(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceClearCpuAffinity") ret = fn(handle) _nvmlCheckReturn(ret) return None def nvmlDeviceGetNumaNodeId(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumaNodeId") node = c_int() ret = fn(handle, byref(node)) _nvmlCheckReturn(ret) return node.value def nvmlDeviceGetMinorNumber(handle): c_minor_number = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinorNumber") ret = fn(handle, byref(c_minor_number)) _nvmlCheckReturn(ret) return c_minor_number.value @convertStrBytes def nvmlDeviceGetUUID(handle): c_uuid = create_string_buffer(NVML_DEVICE_UUID_V2_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlDeviceGetUUID") ret = fn(handle, c_uuid, c_uint(NVML_DEVICE_UUID_V2_BUFFER_SIZE)) _nvmlCheckReturn(ret) return c_uuid.value @convertStrBytes def nvmlDeviceGetInforomVersion(handle, infoRomObject): c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomVersion") ret = fn( handle, _nvmlInforomObject_t(infoRomObject), c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE), ) _nvmlCheckReturn(ret) return c_version.value # Added in 4.304 @convertStrBytes def nvmlDeviceGetInforomImageVersion(handle): c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomImageVersion") ret = fn(handle, c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)) _nvmlCheckReturn(ret) return c_version.value # Added in 4.304 def nvmlDeviceGetInforomConfigurationChecksum(handle): c_checksum = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomConfigurationChecksum") ret = fn(handle, byref(c_checksum)) _nvmlCheckReturn(ret) return c_checksum.value # Added in 4.304 def nvmlDeviceValidateInforom(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceValidateInforom") ret = fn(handle) _nvmlCheckReturn(ret) return None def nvmlDeviceGetLastBBXFlushTime(handle): c_timestamp = c_ulonglong() c_durationUs = c_ulong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetLastBBXFlushTime") ret = fn(handle, byref(c_timestamp), byref(c_durationUs)) _nvmlCheckReturn(ret) return [c_timestamp.value, c_durationUs.value] def nvmlDeviceGetDisplayMode(handle): c_mode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayMode") ret = fn(handle, byref(c_mode)) _nvmlCheckReturn(ret) return c_mode.value def nvmlDeviceGetDisplayActive(handle): c_mode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayActive") ret = fn(handle, byref(c_mode)) _nvmlCheckReturn(ret) return c_mode.value def nvmlDeviceGetPersistenceMode(handle): c_state = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPersistenceMode") ret = fn(handle, byref(c_state)) _nvmlCheckReturn(ret) return c_state.value def nvmlDeviceGetPciInfoExt(handle, c_info): fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfoExt") ret = fn(handle, c_info) _nvmlCheckReturn(ret) return None def nvmlDeviceGetPciInfo_v3(handle): c_info = nvmlPciInfo_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfo_v3") ret = fn(handle, byref(c_info)) _nvmlCheckReturn(ret) return c_info def nvmlDeviceGetPciInfo(handle): return nvmlDeviceGetPciInfo_v3(handle) def nvmlDeviceGetClockInfo(handle, type): c_clock = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockInfo") ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) _nvmlCheckReturn(ret) return c_clock.value # Added in 2.285 def nvmlDeviceGetMaxClockInfo(handle, type): c_clock = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxClockInfo") ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) _nvmlCheckReturn(ret) return c_clock.value # Added in 4.304 def nvmlDeviceGetApplicationsClock(handle, type): c_clock = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetApplicationsClock") ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) _nvmlCheckReturn(ret) return c_clock.value def nvmlDeviceGetMaxCustomerBoostClock(handle, type): c_clock = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxCustomerBoostClock") ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) _nvmlCheckReturn(ret) return c_clock.value def nvmlDeviceGetClock(handle, type, id): c_clock = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetClock") ret = fn(handle, _nvmlClockType_t(type), _nvmlClockId_t(id), byref(c_clock)) _nvmlCheckReturn(ret) return c_clock.value # Added in 5.319 def nvmlDeviceGetDefaultApplicationsClock(handle, type): c_clock = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultApplicationsClock") ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) _nvmlCheckReturn(ret) return c_clock.value # Added in 4.304 def nvmlDeviceGetSupportedMemoryClocks(handle): # first call to get the size c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedMemoryClocks") ret = fn(handle, byref(c_count), None) if ret == NVML_SUCCESS: # special case, no clocks return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case clocks_array = c_uint * c_count.value c_clocks = clocks_array() # make the call again ret = fn(handle, byref(c_count), c_clocks) _nvmlCheckReturn(ret) procs = [] for i in range(c_count.value): procs.append(c_clocks[i]) return procs else: # error case raise NVMLError(ret) # Added in 4.304 def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz): # first call to get the size c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedGraphicsClocks") ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), None) if ret == NVML_SUCCESS: # special case, no clocks return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case clocks_array = c_uint * c_count.value c_clocks = clocks_array() # make the call again ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), c_clocks) _nvmlCheckReturn(ret) procs = [] for i in range(c_count.value): procs.append(c_clocks[i]) return procs else: # error case raise NVMLError(ret) def nvmlDeviceGetFanSpeed(handle): c_speed = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed") ret = fn(handle, byref(c_speed)) _nvmlCheckReturn(ret) return c_speed.value def nvmlDeviceGetFanSpeed_v2(handle, fan): c_speed = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed_v2") ret = fn(handle, fan, byref(c_speed)) _nvmlCheckReturn(ret) return c_speed.value class c_nvmlFanSpeedInfo_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("fan", c_uint), ("speed", c_uint), ] nvmlFanSpeedInfo_v1 = 0x100000C def nvmlDeviceGetFanSpeedRPM(handle): c_fanSpeed = c_nvmlFanSpeedInfo_t() c_fanSpeed.fan = 0 c_fanSpeed.version = nvmlFanSpeedInfo_v1 fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeedRPM") ret = fn(handle, byref(c_fanSpeed)) _nvmlCheckReturn(ret) return c_fanSpeed.speed def nvmlDeviceGetTargetFanSpeed(handle, fan): c_speed = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTargetFanSpeed") ret = fn(handle, fan, byref(c_speed)) _nvmlCheckReturn(ret) return c_speed.value def nvmlDeviceGetNumFans(device): c_numFans = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumFans") ret = fn(device, byref(c_numFans)) _nvmlCheckReturn(ret) return c_numFans.value def nvmlDeviceSetDefaultFanSpeed_v2(handle, index): fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultFanSpeed_v2") ret = fn(handle, index) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceGetMinMaxFanSpeed(handle, minSpeed=c_uint(), maxSpeed=c_uint()): isReference = (type(minSpeed) is not c_uint) or (type(maxSpeed) is not c_uint) minSpeedRef = minSpeed if isReference else byref(minSpeed) maxSpeedRef = maxSpeed if isReference else byref(maxSpeed) fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxFanSpeed") ret = fn(handle, minSpeedRef, maxSpeedRef) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else [minSpeed.value, maxSpeed.value] def nvmlDeviceGetFanControlPolicy_v2(handle, fan, fanControlPolicy=c_uint()): isReference = type(fanControlPolicy) is not c_uint fanControlPolicyRef = fanControlPolicy if isReference else byref(fanControlPolicy) fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanControlPolicy_v2") ret = fn(handle, fan, fanControlPolicyRef) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else fanControlPolicy.value def nvmlDeviceSetFanControlPolicy(handle, fan, fanControlPolicy): fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanControlPolicy") ret = fn(handle, fan, _nvmlFanControlPolicy_t(fanControlPolicy)) _nvmlCheckReturn(ret) return NVML_SUCCESS class c_nvmlTemperature_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("sensorType", _nvmlTemperatureSensors_t), ("temperature", c_int), ] nvmlTemperature_v1 = 0x100000C def nvmlDeviceGetTemperatureV1(handle, sensor): c_temp = c_nvmlTemperature_v1_t() c_temp.version = nvmlTemperature_v1 c_temp.sensorType = _nvmlTemperatureSensors_t(sensor) fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureV") ret = fn(handle, byref(c_temp)) _nvmlCheckReturn(ret) return c_temp.temperature def nvmlDeviceGetTemperatureV(handle, sensor, version=nvmlTemperature_v1): if version == nvmlTemperature_v1: return nvmlDeviceGetTemperatureV1(handle, sensor) else: raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) # DEPRECATED use nvmlDeviceGetTemperatureV instead def nvmlDeviceGetTemperature(handle, sensor): c_temp = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperature") ret = fn(handle, _nvmlTemperatureSensors_t(sensor), byref(c_temp)) _nvmlCheckReturn(ret) return c_temp.value def nvmlDeviceGetTemperatureThreshold(handle, threshold): c_temp = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureThreshold") ret = fn(handle, _nvmlTemperatureThresholds_t(threshold), byref(c_temp)) _nvmlCheckReturn(ret) return c_temp.value def nvmlDeviceSetTemperatureThreshold(handle, threshold, temp): c_temp = c_uint() c_temp.value = temp fn = _nvmlGetFunctionPointer("nvmlDeviceSetTemperatureThreshold") ret = fn(handle, _nvmlTemperatureThresholds_t(threshold), byref(c_temp)) _nvmlCheckReturn(ret) return None def nvmlDeviceGetMarginTemperature(handle): c_marginTempInfo = c_nvmlMarginTemperature_v1_t() c_marginTempInfo.version = nvmlMarginTemperature_v1 fn = _nvmlGetFunctionPointer("nvmlDeviceGetMarginTemperature") ret = fn(handle, byref(c_marginTempInfo)) _nvmlCheckReturn(ret) return c_marginTempInfo.marginTemperature # DEPRECATED use nvmlDeviceGetPerformanceState def nvmlDeviceGetPowerState(handle): c_pstate = _nvmlPstates_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerState") ret = fn(handle, byref(c_pstate)) _nvmlCheckReturn(ret) return c_pstate.value def nvmlDeviceGetPerformanceState(handle): c_pstate = _nvmlPstates_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceState") ret = fn(handle, byref(c_pstate)) _nvmlCheckReturn(ret) return c_pstate.value def nvmlDeviceGetPowerManagementMode(handle): c_pcapMode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementMode") ret = fn(handle, byref(c_pcapMode)) _nvmlCheckReturn(ret) return c_pcapMode.value def nvmlDeviceGetPowerManagementLimit(handle): c_limit = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimit") ret = fn(handle, byref(c_limit)) _nvmlCheckReturn(ret) return c_limit.value # Added in 4.304 def nvmlDeviceGetPowerManagementLimitConstraints(handle): c_minLimit = c_uint() c_maxLimit = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimitConstraints") ret = fn(handle, byref(c_minLimit), byref(c_maxLimit)) _nvmlCheckReturn(ret) return [c_minLimit.value, c_maxLimit.value] # Added in 4.304 def nvmlDeviceGetPowerManagementDefaultLimit(handle): c_limit = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementDefaultLimit") ret = fn(handle, byref(c_limit)) _nvmlCheckReturn(ret) return c_limit.value # Added in 331 def nvmlDeviceGetEnforcedPowerLimit(handle): c_limit = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetEnforcedPowerLimit") ret = fn(handle, byref(c_limit)) _nvmlCheckReturn(ret) return c_limit.value def nvmlDeviceGetPowerUsage(handle): c_watts = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerUsage") ret = fn(handle, byref(c_watts)) _nvmlCheckReturn(ret) return c_watts.value def nvmlDeviceGetTotalEnergyConsumption(handle): c_millijoules = c_uint64() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEnergyConsumption") ret = fn(handle, byref(c_millijoules)) _nvmlCheckReturn(ret) return c_millijoules.value # Added in 4.304 def nvmlDeviceGetGpuOperationMode(handle): c_currState = _nvmlGpuOperationMode_t() c_pendingState = _nvmlGpuOperationMode_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuOperationMode") ret = fn(handle, byref(c_currState), byref(c_pendingState)) _nvmlCheckReturn(ret) return [c_currState.value, c_pendingState.value] # Added in 4.304 def nvmlDeviceGetCurrentGpuOperationMode(handle): return nvmlDeviceGetGpuOperationMode(handle)[0] # Added in 4.304 def nvmlDeviceGetPendingGpuOperationMode(handle): return nvmlDeviceGetGpuOperationMode(handle)[1] def nvmlDeviceGetMemoryInfo(handle, version=None): if not version: c_memory = c_nvmlMemory_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryInfo") else: c_memory = c_nvmlMemory_v2_t() c_memory.version = version fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryInfo_v2") ret = fn(handle, byref(c_memory)) _nvmlCheckReturn(ret) return c_memory def nvmlDeviceGetBAR1MemoryInfo(handle): c_bar1_memory = c_nvmlBAR1Memory_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBAR1MemoryInfo") ret = fn(handle, byref(c_bar1_memory)) _nvmlCheckReturn(ret) return c_bar1_memory def nvmlDeviceGetComputeMode(handle): c_mode = _nvmlComputeMode_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeMode") ret = fn(handle, byref(c_mode)) _nvmlCheckReturn(ret) return c_mode.value def nvmlDeviceGetCudaComputeCapability(handle): c_major = c_int() c_minor = c_int() fn = _nvmlGetFunctionPointer("nvmlDeviceGetCudaComputeCapability") ret = fn(handle, byref(c_major), byref(c_minor)) _nvmlCheckReturn(ret) return (c_major.value, c_minor.value) def nvmlDeviceGetEccMode(handle): c_currState = _nvmlEnableState_t() c_pendingState = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetEccMode") ret = fn(handle, byref(c_currState), byref(c_pendingState)) _nvmlCheckReturn(ret) return [c_currState.value, c_pendingState.value] # added to API def nvmlDeviceGetCurrentEccMode(handle): return nvmlDeviceGetEccMode(handle)[0] # added to API def nvmlDeviceGetPendingEccMode(handle): return nvmlDeviceGetEccMode(handle)[1] def nvmlDeviceGetDefaultEccMode(handle): c_defaultState = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultEccMode") ret = fn(handle, byref(c_defaultState)) _nvmlCheckReturn(ret) return [c_defaultState.value] def nvmlDeviceGetTotalEccErrors(handle, errorType, counterType): c_count = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEccErrors") ret = fn( handle, _nvmlMemoryErrorType_t(errorType), _nvmlEccCounterType_t(counterType), byref(c_count), ) _nvmlCheckReturn(ret) return c_count.value # This is deprecated, instead use nvmlDeviceGetMemoryErrorCounter def nvmlDeviceGetDetailedEccErrors(handle, errorType, counterType): c_counts = c_nvmlEccErrorCounts_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDetailedEccErrors") ret = fn( handle, _nvmlMemoryErrorType_t(errorType), _nvmlEccCounterType_t(counterType), byref(c_counts), ) _nvmlCheckReturn(ret) return c_counts # Added in 4.304 def nvmlDeviceGetMemoryErrorCounter(handle, errorType, counterType, locationType): c_count = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryErrorCounter") ret = fn( handle, _nvmlMemoryErrorType_t(errorType), _nvmlEccCounterType_t(counterType), _nvmlMemoryLocation_t(locationType), byref(c_count), ) _nvmlCheckReturn(ret) return c_count.value def nvmlDeviceGetUtilizationRates(handle): c_util = c_nvmlUtilization_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetUtilizationRates") ret = fn(handle, byref(c_util)) _nvmlCheckReturn(ret) return c_util def nvmlDeviceGetEncoderUtilization(handle): c_util = c_uint() c_samplingPeriod = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderUtilization") ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) _nvmlCheckReturn(ret) return [c_util.value, c_samplingPeriod.value] def nvmlDeviceGetDecoderUtilization(handle): c_util = c_uint() c_samplingPeriod = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDecoderUtilization") ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) _nvmlCheckReturn(ret) return [c_util.value, c_samplingPeriod.value] def nvmlDeviceGetJpgUtilization(handle): c_util = c_uint() c_samplingPeriod = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetJpgUtilization") ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) _nvmlCheckReturn(ret) return [c_util.value, c_samplingPeriod.value] def nvmlDeviceGetOfaUtilization(handle): c_util = c_uint() c_samplingPeriod = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetOfaUtilization") ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) _nvmlCheckReturn(ret) return [c_util.value, c_samplingPeriod.value] def nvmlDeviceGetPcieReplayCounter(handle): c_replay = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieReplayCounter") ret = fn(handle, byref(c_replay)) _nvmlCheckReturn(ret) return c_replay.value def nvmlDeviceGetDriverModel(handle): c_currModel = _nvmlDriverModel_t() c_pendingModel = _nvmlDriverModel_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDriverModel") ret = fn(handle, byref(c_currModel), byref(c_pendingModel)) _nvmlCheckReturn(ret) return [c_currModel.value, c_pendingModel.value] # added to API def nvmlDeviceGetCurrentDriverModel(handle): return nvmlDeviceGetDriverModel(handle)[0] # added to API def nvmlDeviceGetPendingDriverModel(handle): return nvmlDeviceGetDriverModel(handle)[1] # Added in 2.285 @convertStrBytes def nvmlDeviceGetVbiosVersion(handle): c_version = create_string_buffer(NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlDeviceGetVbiosVersion") ret = fn(handle, c_version, c_uint(NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE)) _nvmlCheckReturn(ret) return c_version.value # Added in 2.285 def nvmlDeviceGetComputeRunningProcesses_v2(handle): # first call to get the size c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v2") ret = fn(handle, byref(c_count), None) if ret == NVML_SUCCESS: # special case, no running processes return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case # oversize the array in case more processes are created c_count.value = c_count.value * 2 + 5 proc_array = c_nvmlProcessInfo_v2_t * c_count.value c_procs = proc_array() # make the call again ret = fn(handle, byref(c_count), c_procs) _nvmlCheckReturn(ret) procs = [] for i in range(c_count.value): # use an alternative struct for this object obj = nvmlStructToFriendlyObject(c_procs[i]) if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above obj.usedGpuMemory = None procs.append(obj) return procs else: # error case raise NVMLError(ret) # Added in 2.285 def nvmlDeviceGetComputeRunningProcesses_v3(handle): # first call to get the size c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v3") ret = fn(handle, byref(c_count), None) if ret == NVML_SUCCESS: # special case, no running processes return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case # oversize the array in case more processes are created c_count.value = c_count.value * 2 + 5 proc_array = c_nvmlProcessInfo_v3_t * c_count.value c_procs = proc_array() # make the call again ret = fn(handle, byref(c_count), c_procs) _nvmlCheckReturn(ret) procs = [] for i in range(c_count.value): # use an alternative struct for this object obj = nvmlStructToFriendlyObject(c_procs[i]) if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above obj.usedGpuMemory = None procs.append(obj) return procs else: # error case raise NVMLError(ret) @throwOnVersionMismatch def nvmlDeviceGetComputeRunningProcesses(handle): return nvmlDeviceGetComputeRunningProcesses_v3(handle) def nvmlDeviceGetGraphicsRunningProcesses_v2(handle): # first call to get the size c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v2") ret = fn(handle, byref(c_count), None) if ret == NVML_SUCCESS: # special case, no running processes return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case # oversize the array in case more processes are created c_count.value = c_count.value * 2 + 5 proc_array = c_nvmlProcessInfo_v2_t * c_count.value c_procs = proc_array() # make the call again ret = fn(handle, byref(c_count), c_procs) _nvmlCheckReturn(ret) procs = [] for i in range(c_count.value): # use an alternative struct for this object obj = nvmlStructToFriendlyObject(c_procs[i]) if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above obj.usedGpuMemory = None procs.append(obj) return procs else: # error case raise NVMLError(ret) def nvmlDeviceGetGraphicsRunningProcesses_v3(handle): # first call to get the size c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v3") ret = fn(handle, byref(c_count), None) if ret == NVML_SUCCESS: # special case, no running processes return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case # oversize the array in case more processes are created c_count.value = c_count.value * 2 + 5 proc_array = c_nvmlProcessInfo_v3_t * c_count.value c_procs = proc_array() # make the call again ret = fn(handle, byref(c_count), c_procs) _nvmlCheckReturn(ret) procs = [] for i in range(c_count.value): # use an alternative struct for this object obj = nvmlStructToFriendlyObject(c_procs[i]) if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above obj.usedGpuMemory = None procs.append(obj) return procs else: # error case raise NVMLError(ret) @throwOnVersionMismatch def nvmlDeviceGetGraphicsRunningProcesses(handle): return nvmlDeviceGetGraphicsRunningProcesses_v3(handle) @throwOnVersionMismatch def nvmlDeviceGetMPSComputeRunningProcesses(handle): return nvmlDeviceGetMPSComputeRunningProcesses_v3(handle) def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle): # first call to get the size c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v2") ret = fn(handle, byref(c_count), None) if ret == NVML_SUCCESS: # special case, no running processes return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case # oversize the array in case more processes are created c_count.value = c_count.value * 2 + 5 proc_array = c_nvmlProcessInfo_v2_t * c_count.value c_procs = proc_array() # make the call again ret = fn(handle, byref(c_count), c_procs) _nvmlCheckReturn(ret) procs = [] for i in range(c_count.value): # use an alternative struct for this object obj = nvmlStructToFriendlyObject(c_procs[i]) if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above obj.usedGpuMemory = None procs.append(obj) return procs else: # error case raise NVMLError(ret) def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle): # first call to get the size c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v3") ret = fn(handle, byref(c_count), None) if ret == NVML_SUCCESS: # special case, no running processes return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case # oversize the array in case more processes are created c_count.value = c_count.value * 2 + 5 proc_array = c_nvmlProcessInfo_v3_t * c_count.value c_procs = proc_array() # make the call again ret = fn(handle, byref(c_count), c_procs) _nvmlCheckReturn(ret) procs = [] for i in range(c_count.value): # use an alternative struct for this object obj = nvmlStructToFriendlyObject(c_procs[i]) if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above obj.usedGpuMemory = None procs.append(obj) return procs else: # error case raise NVMLError(ret) def nvmlDeviceGetRunningProcessDetailList(handle, version, mode): c_processDetailList = c_nvmlProcessDetailList_t() c_processDetailList.version = version c_processDetailList.mode = mode fn = _nvmlGetFunctionPointer("nvmlDeviceGetRunningProcessDetailList") # first call to get the size ret = fn(handle, byref(c_processDetailList)) if ret == NVML_SUCCESS: # special case, no running processes return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: c_procs = c_nvmlProcessDetail_v1_t * c_processDetailList.numProcArrayEntries c_processDetailList.procArray = cast( (c_procs)(), POINTER(c_nvmlProcessDetail_v1_t) ) # make the call again ret = fn(handle, byref(c_processDetailList)) _nvmlCheckReturn(ret) procs = [] for i in range(c_processDetailList.numProcArrayEntries): # use an alternative struct for this object obj = c_processDetailList.procArray[i] if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: obj.usedGpuMemory = None if obj.usedGpuCcProtectedMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: obj.usedGpuCcProtectedMemory = None procs.append(obj) return procs else: # error case raise NVMLError(ret) def nvmlDeviceGetAutoBoostedClocksEnabled(handle): c_isEnabled = _nvmlEnableState_t() c_defaultIsEnabled = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAutoBoostedClocksEnabled") ret = fn(handle, byref(c_isEnabled), byref(c_defaultIsEnabled)) _nvmlCheckReturn(ret) return [c_isEnabled.value, c_defaultIsEnabled.value] # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks ## Set functions def nvmlUnitSetLedState(unit, color): fn = _nvmlGetFunctionPointer("nvmlUnitSetLedState") ret = fn(unit, _nvmlLedColor_t(color)) _nvmlCheckReturn(ret) return None def nvmlDeviceSetPersistenceMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetPersistenceMode") ret = fn(handle, _nvmlEnableState_t(mode)) _nvmlCheckReturn(ret) return None def nvmlDeviceSetComputeMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetComputeMode") ret = fn(handle, _nvmlComputeMode_t(mode)) _nvmlCheckReturn(ret) return None def nvmlDeviceSetEccMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetEccMode") ret = fn(handle, _nvmlEnableState_t(mode)) _nvmlCheckReturn(ret) return None def nvmlDeviceClearEccErrorCounts(handle, counterType): fn = _nvmlGetFunctionPointer("nvmlDeviceClearEccErrorCounts") ret = fn(handle, _nvmlEccCounterType_t(counterType)) _nvmlCheckReturn(ret) return None def nvmlDeviceSetDriverModel(handle, model): fn = _nvmlGetFunctionPointer("nvmlDeviceSetDriverModel") ret = fn(handle, _nvmlDriverModel_t(model)) _nvmlCheckReturn(ret) return None def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled): fn = _nvmlGetFunctionPointer("nvmlDeviceSetAutoBoostedClocksEnabled") ret = fn(handle, _nvmlEnableState_t(enabled)) _nvmlCheckReturn(ret) return None # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags): fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultAutoBoostedClocksEnabled") ret = fn(handle, _nvmlEnableState_t(enabled), c_uint(flags)) _nvmlCheckReturn(ret) return None # Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks def nvmlDeviceSetGpuLockedClocks(handle, minGpuClockMHz, maxGpuClockMHz): fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuLockedClocks") ret = fn(handle, c_uint(minGpuClockMHz), c_uint(maxGpuClockMHz)) _nvmlCheckReturn(ret) return None def nvmlDeviceResetGpuLockedClocks(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceResetGpuLockedClocks") ret = fn(handle) _nvmlCheckReturn(ret) return None def nvmlDeviceSetMemoryLockedClocks(handle, minMemClockMHz, maxMemClockMHz): fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemoryLockedClocks") ret = fn(handle, c_uint(minMemClockMHz), c_uint(maxMemClockMHz)) _nvmlCheckReturn(ret) return None def nvmlDeviceResetMemoryLockedClocks(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceResetMemoryLockedClocks") ret = fn(handle) _nvmlCheckReturn(ret) return None def nvmlDeviceGetClkMonStatus(handle, c_clkMonInfo=nvmlClkMonStatus_t()): isReference = type(c_clkMonInfo) is not nvmlClkMonStatus_t c_clkMonInfoRef = c_clkMonInfo if isReference else byref(c_clkMonInfo) fn = _nvmlGetFunctionPointer("nvmlDeviceGetClkMonStatus") ret = fn(handle, c_clkMonInfoRef) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else c_clkMonInfo # Added in 4.304 def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz): fn = _nvmlGetFunctionPointer("nvmlDeviceSetApplicationsClocks") ret = fn(handle, c_uint(maxMemClockMHz), c_uint(maxGraphicsClockMHz)) _nvmlCheckReturn(ret) return None # Added in 4.304 def nvmlDeviceResetApplicationsClocks(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceResetApplicationsClocks") ret = fn(handle) _nvmlCheckReturn(ret) return None # Added in 4.304 def nvmlDeviceSetPowerManagementLimit(handle, limit): fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit") ret = fn(handle, c_uint(limit)) _nvmlCheckReturn(ret) return None # Added in 4.304 def nvmlDeviceSetGpuOperationMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuOperationMode") ret = fn(handle, _nvmlGpuOperationMode_t(mode)) _nvmlCheckReturn(ret) return None # Added in 2.285 def nvmlEventSetCreate(): fn = _nvmlGetFunctionPointer("nvmlEventSetCreate") eventSet = c_nvmlEventSet_t() ret = fn(byref(eventSet)) _nvmlCheckReturn(ret) return eventSet # Added in 2.285 def nvmlDeviceRegisterEvents(handle, eventTypes, eventSet): fn = _nvmlGetFunctionPointer("nvmlDeviceRegisterEvents") ret = fn(handle, c_ulonglong(eventTypes), eventSet) _nvmlCheckReturn(ret) return None # Added in 2.285 def nvmlDeviceGetSupportedEventTypes(handle): c_eventTypes = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedEventTypes") ret = fn(handle, byref(c_eventTypes)) _nvmlCheckReturn(ret) return c_eventTypes.value # raises NVML_ERROR_TIMEOUT exception on timeout def nvmlEventSetWait_v2(eventSet, timeoutms): fn = _nvmlGetFunctionPointer("nvmlEventSetWait_v2") data = c_nvmlEventData_t() ret = fn(eventSet, byref(data), c_uint(timeoutms)) _nvmlCheckReturn(ret) return data def nvmlEventSetWait(eventSet, timeoutms): return nvmlEventSetWait_v2(eventSet, timeoutms) # Added in 2.285 def nvmlEventSetFree(eventSet): fn = _nvmlGetFunctionPointer("nvmlEventSetFree") ret = fn(eventSet) _nvmlCheckReturn(ret) return None # Added in 3.295 def nvmlDeviceOnSameBoard(handle1, handle2): fn = _nvmlGetFunctionPointer("nvmlDeviceOnSameBoard") onSameBoard = c_int() ret = fn(handle1, handle2, byref(onSameBoard)) _nvmlCheckReturn(ret) return onSameBoard.value != 0 # Added in 3.295 def nvmlDeviceGetCurrPcieLinkGeneration(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkGeneration") gen = c_uint() ret = fn(handle, byref(gen)) _nvmlCheckReturn(ret) return gen.value # Added in 3.295 def nvmlDeviceGetMaxPcieLinkGeneration(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkGeneration") gen = c_uint() ret = fn(handle, byref(gen)) _nvmlCheckReturn(ret) return gen.value # Added in 3.295 def nvmlDeviceGetCurrPcieLinkWidth(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkWidth") width = c_uint() ret = fn(handle, byref(width)) _nvmlCheckReturn(ret) return width.value # Added in 3.295 def nvmlDeviceGetMaxPcieLinkWidth(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkWidth") width = c_uint() ret = fn(handle, byref(width)) _nvmlCheckReturn(ret) return width.value def nvmlDeviceGetGpuMaxPcieLinkGeneration(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuMaxPcieLinkGeneration") gen = c_uint() ret = fn(handle, byref(gen)) _nvmlCheckReturn(ret) return gen.value # Added in 4.304 def nvmlDeviceGetSupportedClocksThrottleReasons(handle): c_reasons = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksThrottleReasons") ret = fn(handle, byref(c_reasons)) _nvmlCheckReturn(ret) return c_reasons.value def nvmlDeviceGetSupportedClocksEventReasons(handle): c_reasons = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksEventReasons") ret = fn(handle, byref(c_reasons)) _nvmlCheckReturn(ret) return c_reasons.value # Added in 4.304 def nvmlDeviceGetCurrentClocksThrottleReasons(handle): c_reasons = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksThrottleReasons") ret = fn(handle, byref(c_reasons)) _nvmlCheckReturn(ret) return c_reasons.value def nvmlDeviceGetCurrentClocksEventReasons(handle): c_reasons = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksEventReasons") ret = fn(handle, byref(c_reasons)) _nvmlCheckReturn(ret) return c_reasons.value # Added in 5.319 def nvmlDeviceGetIndex(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetIndex") c_index = c_uint() ret = fn(handle, byref(c_index)) _nvmlCheckReturn(ret) return c_index.value # Added in 5.319 def nvmlDeviceGetAccountingMode(handle): c_mode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingMode") ret = fn(handle, byref(c_mode)) _nvmlCheckReturn(ret) return c_mode.value def nvmlDeviceSetAccountingMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetAccountingMode") ret = fn(handle, _nvmlEnableState_t(mode)) _nvmlCheckReturn(ret) return None def nvmlDeviceClearAccountingPids(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceClearAccountingPids") ret = fn(handle) _nvmlCheckReturn(ret) return None def nvmlDeviceGetAccountingStats(handle, pid): stats = c_nvmlAccountingStats_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingStats") ret = fn(handle, c_uint(pid), byref(stats)) _nvmlCheckReturn(ret) if stats.maxMemoryUsage == NVML_VALUE_NOT_AVAILABLE_ulonglong.value: # special case for WDDM on Windows, see comment above stats.maxMemoryUsage = None return stats def nvmlDeviceGetAccountingPids(handle): count = c_uint(nvmlDeviceGetAccountingBufferSize(handle)) pids = (c_uint * count.value)() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingPids") ret = fn(handle, byref(count), pids) _nvmlCheckReturn(ret) return list(map(int, pids[0 : count.value])) def nvmlDeviceGetAccountingBufferSize(handle): bufferSize = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingBufferSize") ret = fn(handle, byref(bufferSize)) _nvmlCheckReturn(ret) return int(bufferSize.value) def nvmlDeviceGetRetiredPages(device, sourceFilter): c_source = _nvmlPageRetirementCause_t(sourceFilter) c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPages") # First call will get the size ret = fn(device, c_source, byref(c_count), None) # this should only fail with insufficient size if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE): raise NVMLError(ret) # call again with a buffer # oversize the array for the rare cases where additional pages # are retired between NVML calls c_count.value = c_count.value * 2 + 5 page_array = c_ulonglong * c_count.value c_pages = page_array() ret = fn(device, c_source, byref(c_count), c_pages) _nvmlCheckReturn(ret) return list(map(int, c_pages[0 : c_count.value])) def nvmlDeviceGetRetiredPages_v2(device, sourceFilter): c_source = _nvmlPageRetirementCause_t(sourceFilter) c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPages_v2") # First call will get the size ret = fn(device, c_source, byref(c_count), None) # this should only fail with insufficient size if (ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE): raise NVMLError(ret) # call again with a buffer # oversize the array for the rare cases where additional pages # are retired between NVML calls c_count.value = c_count.value * 2 + 5 page_array = c_ulonglong * c_count.value c_pages = page_array() times_array = c_ulonglong * c_count.value c_times = times_array() ret = fn(device, c_source, byref(c_count), c_pages, c_times) _nvmlCheckReturn(ret) return [ {"address": int(c_pages[i]), "timestamp": int(c_times[i])} for i in range(c_count.value) ] def nvmlDeviceGetRetiredPagesPendingStatus(device): c_pending = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPagesPendingStatus") ret = fn(device, byref(c_pending)) _nvmlCheckReturn(ret) return int(c_pending.value) def nvmlDeviceGetAPIRestriction(device, apiType): c_permission = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAPIRestriction") ret = fn(device, _nvmlRestrictedAPI_t(apiType), byref(c_permission)) _nvmlCheckReturn(ret) return int(c_permission.value) def nvmlDeviceSetAPIRestriction(handle, apiType, isRestricted): fn = _nvmlGetFunctionPointer("nvmlDeviceSetAPIRestriction") ret = fn(handle, _nvmlRestrictedAPI_t(apiType), _nvmlEnableState_t(isRestricted)) _nvmlCheckReturn(ret) return None def nvmlDeviceGetBridgeChipInfo(handle): bridgeHierarchy = c_nvmlBridgeChipHierarchy_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBridgeChipInfo") ret = fn(handle, byref(bridgeHierarchy)) _nvmlCheckReturn(ret) return bridgeHierarchy def nvmlDeviceGetSamples(device, sampling_type, timeStamp): c_sampling_type = _nvmlSamplingType_t(sampling_type) c_time_stamp = c_ulonglong(timeStamp) c_sample_count = c_uint(0) c_sample_value_type = _nvmlValueType_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetSamples") ## First Call gets the size ret = fn( device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), None, ) # Stop if this fails if ret != NVML_SUCCESS: raise NVMLError(ret) sampleArray = c_sample_count.value * c_nvmlSample_t c_samples = sampleArray() ret = fn( device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), c_samples, ) _nvmlCheckReturn(ret) return (c_sample_value_type.value, c_samples[0 : c_sample_count.value]) def nvmlDeviceGetViolationStatus(device, perfPolicyType): c_perfPolicy_type = _nvmlPerfPolicyType_t(perfPolicyType) c_violTime = c_nvmlViolationTime_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetViolationStatus") ## Invoke the method to get violation time ret = fn(device, c_perfPolicy_type, byref(c_violTime)) _nvmlCheckReturn(ret) return c_violTime def nvmlDeviceGetPcieThroughput(device, counter): c_util = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieThroughput") ret = fn(device, _nvmlPcieUtilCounter_t(counter), byref(c_util)) _nvmlCheckReturn(ret) return c_util.value def nvmlSystemGetTopologyGpuSet(cpuNumber): c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlSystemGetTopologyGpuSet") # First call will get the size ret = fn(cpuNumber, byref(c_count), None) if ret != NVML_SUCCESS: raise NVMLError(ret) # call again with a buffer device_array = c_nvmlDevice_t * c_count.value c_devices = device_array() ret = fn(cpuNumber, byref(c_count), c_devices) _nvmlCheckReturn(ret) return list(c_devices[0 : c_count.value]) def nvmlDeviceGetTopologyNearestGpus(device, level): c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyNearestGpus") # First call will get the size ret = fn(device, level, byref(c_count), None) if ret != NVML_SUCCESS: raise NVMLError(ret) # call again with a buffer device_array = c_nvmlDevice_t * c_count.value c_devices = device_array() ret = fn(device, level, byref(c_count), c_devices) _nvmlCheckReturn(ret) return list(c_devices[0 : c_count.value]) def nvmlDeviceGetTopologyCommonAncestor(device1, device2): c_level = _nvmlGpuTopologyLevel_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyCommonAncestor") ret = fn(device1, device2, byref(c_level)) _nvmlCheckReturn(ret) return c_level.value def nvmlDeviceGetNvLinkUtilizationCounter(device, link, counter): c_rxcounter = c_ulonglong() c_txcounter = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkUtilizationCounter") ret = fn(device, link, counter, byref(c_rxcounter), byref(c_txcounter)) _nvmlCheckReturn(ret) return (c_rxcounter.value, c_txcounter.value) def nvmlDeviceFreezeNvLinkUtilizationCounter(device, link, counter, freeze): fn = _nvmlGetFunctionPointer("nvmlDeviceFreezeNvLinkUtilizationCounter") ret = fn(device, link, counter, freeze) _nvmlCheckReturn(ret) return None def nvmlDeviceResetNvLinkUtilizationCounter(device, link, counter): fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkUtilizationCounter") ret = fn(device, link, counter) _nvmlCheckReturn(ret) return None def nvmlDeviceSetNvLinkUtilizationControl(device, link, counter, control, reset): fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvLinkUtilizationControl") ret = fn(device, link, counter, byref(control), reset) _nvmlCheckReturn(ret) return None def nvmlDeviceGetNvLinkUtilizationControl(device, link, counter): c_control = nvmlNvLinkUtilizationControl_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkUtilizationControl") ret = fn(device, link, counter, byref(c_control)) _nvmlCheckReturn(ret) return c_control def nvmlDeviceGetNvLinkCapability(device, link, capability): c_capResult = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkCapability") ret = fn(device, link, capability, byref(c_capResult)) _nvmlCheckReturn(ret) return c_capResult.value def nvmlDeviceGetNvLinkErrorCounter(device, link, counter): c_result = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkErrorCounter") ret = fn(device, link, counter, byref(c_result)) _nvmlCheckReturn(ret) return c_result.value def nvmlDeviceResetNvLinkErrorCounters(device, link): fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkErrorCounters") ret = fn(device, link) _nvmlCheckReturn(ret) return None def nvmlDeviceGetNvLinkRemotePciInfo(device, link): c_pci = nvmlPciInfo_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemotePciInfo_v2") ret = fn(device, link, byref(c_pci)) _nvmlCheckReturn(ret) return c_pci def nvmlDeviceGetNvLinkRemoteDeviceType(handle, link): c_type = _nvmlNvLinkDeviceType_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemoteDeviceType") ret = fn(handle, link, byref(c_type)) _nvmlCheckReturn(ret) return c_type.value def nvmlDeviceGetNvLinkState(device, link): c_isActive = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkState") ret = fn(device, link, byref(c_isActive)) _nvmlCheckReturn(ret) return c_isActive.value def nvmlDeviceGetNvLinkVersion(device, link): c_version = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkVersion") ret = fn(device, link, byref(c_version)) _nvmlCheckReturn(ret) return c_version.value def nvmlDeviceModifyDrainState(pciInfo, newState): fn = _nvmlGetFunctionPointer("nvmlDeviceModifyDrainState") ret = fn(pointer(pciInfo), newState) _nvmlCheckReturn(ret) return None def nvmlDeviceQueryDrainState(pciInfo): c_newState = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceQueryDrainState") ret = fn(pointer(pciInfo), byref(c_newState)) _nvmlCheckReturn(ret) return c_newState.value def nvmlDeviceRemoveGpu(pciInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceRemoveGpu") ret = fn(pointer(pciInfo)) _nvmlCheckReturn(ret) return None def nvmlDeviceDiscoverGpus(pciInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceDiscoverGpus") ret = fn(pointer(pciInfo)) _nvmlCheckReturn(ret) return None def nvmlDeviceGetFieldValues(handle, fieldIds): values_arr = c_nvmlFieldValue_t * len(fieldIds) values = values_arr() fn = _nvmlGetFunctionPointer("nvmlDeviceGetFieldValues") for i, fieldId in enumerate(fieldIds): try: values[i].fieldId, values[i].scopeId = fieldId except TypeError: values[i].fieldId = fieldId ret = fn(handle, c_int32(len(fieldIds)), byref(values)) _nvmlCheckReturn(ret) return values def nvmlDeviceClearFieldValues(handle, fieldIds): values_arr = c_nvmlFieldValue_t * len(fieldIds) values = values_arr() fn = _nvmlGetFunctionPointer("nvmlDeviceClearFieldValues") for i, fieldId in enumerate(fieldIds): try: values[i].fieldId, values[i].scopeId = fieldId except TypeError: values[i].fieldId = fieldId ret = fn(handle, c_int32(len(fieldIds)), byref(values)) _nvmlCheckReturn(ret) return values def nvmlDeviceGetVirtualizationMode(handle): c_virtualization_mode = c_ulonglong() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVirtualizationMode") ret = fn(handle, byref(c_virtualization_mode)) _nvmlCheckReturn(ret) return c_virtualization_mode.value def nvmlDeviceSetVirtualizationMode(handle, virtualization_mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetVirtualizationMode") return fn(handle, virtualization_mode) def nvmlDeviceGetVgpuHeterogeneousMode(handle): c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuHeterogeneousMode") ret = fn(handle, byref(c_vgpuHeterogeneousMode)) _nvmlCheckReturn(ret) return c_vgpuHeterogeneousMode.mode def nvmlDeviceSetVgpuHeterogeneousMode(handle, heterogeneous_mode): c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 c_vgpuHeterogeneousMode.mode = heterogeneous_mode fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuHeterogeneousMode") ret = fn(handle, byref(c_vgpuHeterogeneousMode)) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlVgpuInstanceGetPlacementId(vgpuInstance): c_placement = c_nvmlVgpuPlacementId_v1_t(0) c_placement.version = VgpuPlacementId_v1 fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetPlacementId") ret = fn(vgpuInstance, byref(c_placement)) _nvmlCheckReturn(ret) return c_placement.placementId def nvmlDeviceGetVgpuTypeSupportedPlacements(handle, vgpuTypeId, mode=0, version=1): c_max_instances = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") ret = fn(handle, vgpuTypeId, byref(c_max_instances)) _nvmlCheckReturn(ret) if version == 2: c_vgpu_placements = c_nvmlVgpuPlacementList_v2_t() c_vgpu_placements.version = VgpuPlacementList_v2 c_vgpu_placements.count = c_max_instances.value c_vgpu_placements.mode = mode elif version == 1: c_vgpu_placements = c_nvmlVgpuPlacementList_v1_t() c_vgpu_placements.version = VgpuPlacementList_v1 else: raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) c_placements = c_uint * c_max_instances.value c_vgpu_placements.placementIds = c_placements() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeSupportedPlacements") ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) _nvmlCheckReturn(ret) return c_vgpu_placements def nvmlDeviceGetVgpuTypeCreatablePlacements(handle, vgpuTypeId, version=1): c_max_instances = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") ret = fn(handle, vgpuTypeId, byref(c_max_instances)) _nvmlCheckReturn(ret) if version == 2: c_vgpu_placements = c_nvmlVgpuPlacementList_v2_t() c_vgpu_placements.version = VgpuPlacementList_v2 c_vgpu_placements.count = c_max_instances.value elif version == 1: c_vgpu_placements = c_nvmlVgpuPlacementList_v1_t() c_vgpu_placements.version = VgpuPlacementList_v1 c_placements = c_uint * c_max_instances.value c_vgpu_placements.placementIds = c_placements() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeCreatablePlacements") ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) _nvmlCheckReturn(ret) return c_vgpu_placements def nvmlGetVgpuDriverCapabilities(capability): c_capResult = c_uint() fn = _nvmlGetFunctionPointer("nvmlGetVgpuDriverCapabilities") ret = fn(_nvmlVgpuDriverCapability_t(capability), byref(c_capResult)) _nvmlCheckReturn(ret) return c_capResult.value def nvmlDeviceGetVgpuCapabilities(handle, capability): c_capResult = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuCapabilities") ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), byref(c_capResult)) _nvmlCheckReturn(ret) return c_capResult.value def nvmlDeviceSetVgpuCapabilities(handle, capability, state): fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuCapabilities") ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), state) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceGetSupportedVgpus(handle): # first call to get the size c_vgpu_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedVgpus") ret = fn(handle, byref(c_vgpu_count), None) if ret == NVML_SUCCESS: # special case, no supported vGPUs return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value c_vgpu_type_ids = vgpu_type_ids_array() # make the call again ret = fn(handle, byref(c_vgpu_count), c_vgpu_type_ids) _nvmlCheckReturn(ret) vgpus = [] for i in range(c_vgpu_count.value): vgpus.append(c_vgpu_type_ids[i]) return vgpus else: # error case raise NVMLError(ret) def nvmlDeviceGetCreatableVgpus(handle): # first call to get the size c_vgpu_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetCreatableVgpus") ret = fn(handle, byref(c_vgpu_count), None) if ret == NVML_SUCCESS: # special case, no supported vGPUs return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value c_vgpu_type_ids = vgpu_type_ids_array() # make the call again ret = fn(handle, byref(c_vgpu_count), c_vgpu_type_ids) _nvmlCheckReturn(ret) vgpus = [] for i in range(c_vgpu_count.value): vgpus.append(c_vgpu_type_ids[i]) return vgpus else: # error case raise NVMLError(ret) def nvmlVgpuTypeGetGpuInstanceProfileId(vgpuTypeId): c_profile_id = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGpuInstanceProfileId") ret = fn(vgpuTypeId, byref(c_profile_id)) _nvmlCheckReturn(ret) return c_profile_id.value @convertStrBytes def nvmlVgpuTypeGetClass(vgpuTypeId): c_class = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetClass") ret = fn(vgpuTypeId, c_class, byref(c_buffer_size)) _nvmlCheckReturn(ret) return c_class.value @convertStrBytes def nvmlVgpuTypeGetName(vgpuTypeId): c_name = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetName") ret = fn(vgpuTypeId, c_name, byref(c_buffer_size)) _nvmlCheckReturn(ret) return c_name.value def nvmlVgpuTypeGetDeviceID(vgpuTypeId): c_device_id = c_ulonglong(0) c_subsystem_id = c_ulonglong(0) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetDeviceID") ret = fn(vgpuTypeId, byref(c_device_id), byref(c_subsystem_id)) _nvmlCheckReturn(ret) return (c_device_id.value, c_subsystem_id.value) def nvmlVgpuTypeGetFramebufferSize(vgpuTypeId): c_fb_size = c_ulonglong(0) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFramebufferSize") ret = fn(vgpuTypeId, byref(c_fb_size)) _nvmlCheckReturn(ret) return c_fb_size.value def nvmlVgpuTypeGetNumDisplayHeads(vgpuTypeId): c_num_heads = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetNumDisplayHeads") ret = fn(vgpuTypeId, byref(c_num_heads)) _nvmlCheckReturn(ret) return c_num_heads.value def nvmlVgpuTypeGetResolution(vgpuTypeId): c_xdim = c_uint(0) c_ydim = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetResolution") ret = fn(vgpuTypeId, 0, byref(c_xdim), byref(c_ydim)) _nvmlCheckReturn(ret) return (c_xdim.value, c_ydim.value) @convertStrBytes def nvmlVgpuTypeGetLicense(vgpuTypeId): c_license = create_string_buffer(NVML_GRID_LICENSE_BUFFER_SIZE) c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetLicense") ret = fn(vgpuTypeId, c_license, c_buffer_size) _nvmlCheckReturn(ret) return c_license.value def nvmlVgpuTypeGetFrameRateLimit(vgpuTypeId): c_frl_config = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFrameRateLimit") ret = fn(vgpuTypeId, byref(c_frl_config)) _nvmlCheckReturn(ret) return c_frl_config.value def nvmlVgpuTypeGetGspHeapSize(vgpuTypeId): c_gsp_heap = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGspHeapSize") ret = fn(vgpuTypeId, byref(c_gsp_heap)) _nvmlCheckReturn(ret) return c_gsp_heap.value def nvmlVgpuTypeGetFbReservation(vgpuTypeId): c_fb_reservation = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFbReservation") ret = fn(vgpuTypeId, byref(c_fb_reservation)) _nvmlCheckReturn(ret) return c_fb_reservation.value def nvmlVgpuInstanceGetRuntimeStateSize(vgpuInstance): c_runtime_state = nvmlVgpuRuntimeState_v1_t() c_runtime_state.version = VgpuRuntimeState_v1 fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetRuntimeStateSize") ret = fn(vgpuInstance, byref(c_runtime_state)) _nvmlCheckReturn(ret) return c_runtime_state def nvmlVgpuTypeGetMaxInstances(handle, vgpuTypeId): c_max_instances = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") ret = fn(handle, vgpuTypeId, byref(c_max_instances)) _nvmlCheckReturn(ret) return c_max_instances.value def nvmlVgpuTypeGetMaxInstancesPerVm(vgpuTypeId): c_max_instances_per_vm = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstancesPerVm") ret = fn(vgpuTypeId, byref(c_max_instances_per_vm)) _nvmlCheckReturn(ret) return c_max_instances_per_vm.value def nvmlVgpuTypeGetBAR1Info(vgpuTypeId): c_bar1Info = c_nvmlVgpuTypeBar1Info_v1_t(0) c_bar1Info.version = VgpuTypeBar1Info_v1 fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetBAR1Info") ret = fn(vgpuTypeId, byref(c_bar1Info)) _nvmlCheckReturn(ret) return c_bar1Info def nvmlDeviceGetActiveVgpus(handle): # first call to get the size c_vgpu_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetActiveVgpus") ret = fn(handle, byref(c_vgpu_count), None) if ret == NVML_SUCCESS: # special case, no active vGPUs return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case vgpu_instance_array = _nvmlVgpuInstance_t * c_vgpu_count.value c_vgpu_instances = vgpu_instance_array() # make the call again ret = fn(handle, byref(c_vgpu_count), c_vgpu_instances) _nvmlCheckReturn(ret) vgpus = [] for i in range(c_vgpu_count.value): vgpus.append(c_vgpu_instances[i]) return vgpus else: # error case raise NVMLError(ret) @convertStrBytes def nvmlVgpuInstanceGetVmID(vgpuInstance): c_vm_id = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) c_vm_id_type = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmID") ret = fn(vgpuInstance, byref(c_vm_id), c_buffer_size, byref(c_vm_id_type)) _nvmlCheckReturn(ret) return (c_vm_id.value, c_vm_id_type.value) @convertStrBytes def nvmlVgpuInstanceGetUUID(vgpuInstance): c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetUUID") ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) _nvmlCheckReturn(ret) return c_uuid.value @convertStrBytes def nvmlVgpuInstanceGetMdevUUID(vgpuInstance): c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMdevUUID") ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) _nvmlCheckReturn(ret) return c_uuid.value @convertStrBytes def nvmlVgpuInstanceGetVmDriverVersion(vgpuInstance): c_driver_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) c_buffer_size = c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmDriverVersion") ret = fn(vgpuInstance, byref(c_driver_version), c_buffer_size) _nvmlCheckReturn(ret) return c_driver_version.value def nvmlVgpuInstanceGetLicenseStatus(vgpuInstance): c_license_status = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseStatus") ret = fn(vgpuInstance, byref(c_license_status)) _nvmlCheckReturn(ret) return c_license_status.value def nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance): fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseInfo_v2") c_license_info = c_nvmlVgpuLicenseInfo_t() ret = fn(vgpuInstance, byref(c_license_info)) _nvmlCheckReturn(ret) return c_license_info def nvmlVgpuInstanceGetLicenseInfo(vgpuInstance): return nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance) def nvmlVgpuInstanceGetFrameRateLimit(vgpuInstance): c_frl = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFrameRateLimit") ret = fn(vgpuInstance, byref(c_frl)) _nvmlCheckReturn(ret) return c_frl.value def nvmlVgpuInstanceGetEccMode(vgpuInstance): c_mode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEccMode") ret = fn(vgpuInstance, byref(c_mode)) _nvmlCheckReturn(ret) return c_mode.value def nvmlVgpuInstanceGetType(vgpuInstance): c_vgpu_type = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetType") ret = fn(vgpuInstance, byref(c_vgpu_type)) _nvmlCheckReturn(ret) return c_vgpu_type.value def nvmlVgpuInstanceGetEncoderCapacity(vgpuInstance): c_encoder_capacity = c_ulonglong(0) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderCapacity") ret = fn(vgpuInstance, byref(c_encoder_capacity)) _nvmlCheckReturn(ret) return c_encoder_capacity.value def nvmlVgpuInstanceSetEncoderCapacity(vgpuInstance, encoder_capacity): fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceSetEncoderCapacity") return fn(vgpuInstance, encoder_capacity) def nvmlVgpuInstanceGetFbUsage(vgpuInstance): c_fb_usage = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFbUsage") ret = fn(vgpuInstance, byref(c_fb_usage)) _nvmlCheckReturn(ret) return c_fb_usage.value def nvmlVgpuTypeGetCapabilities(vgpuTypeId, capability): c_cap_result = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetCapabilities") ret = fn(vgpuTypeId, _nvmlVgpuCapability_t(capability), byref(c_cap_result)) _nvmlCheckReturn(ret) return c_cap_result.value def nvmlVgpuInstanceGetGpuInstanceId(vgpuInstance): c_id = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuInstanceId") ret = fn(vgpuInstance, byref(c_id)) _nvmlCheckReturn(ret) return c_id.value @convertStrBytes def nvmlVgpuInstanceGetGpuPciId(vgpuInstance): c_vgpuPciId = create_string_buffer(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuPciId") ret = fn( vgpuInstance, c_vgpuPciId, byref(c_uint(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE)) ) _nvmlCheckReturn(ret) return c_vgpuPciId.value def nvmlDeviceGetVgpuUtilization(handle, timeStamp): # first call to get the size c_vgpu_count = c_uint(0) c_time_stamp = c_ulonglong(timeStamp) c_sample_value_type = _nvmlValueType_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuUtilization") ret = fn( handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), None ) if ret == NVML_SUCCESS: # special case, no active vGPUs return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case sampleArray = c_vgpu_count.value * c_nvmlVgpuInstanceUtilizationSample_t c_samples = sampleArray() # make the call again ret = fn( handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), c_samples, ) _nvmlCheckReturn(ret) return c_samples[0 : c_vgpu_count.value] else: # error case raise NVMLError(ret) def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp): # first call to get the size c_time_stamp = c_ulonglong(timeStamp) c_vgpuUtilInfo = c_nvmlVgpuInstancesUtilizationInfo_v1_t(0) c_vgpuUtilInfo.version = VgpuInstancesUtilizationInfo_v1 c_vgpuUtilInfo.sampleValType = _nvmlValueType_t() c_vgpuUtilInfo.vgpuInstanceCount = c_uint(0) c_vgpuUtilInfo.lastSeenTimeStamp = c_time_stamp fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuInstancesUtilizationInfo") ret = fn(handle, byref(c_vgpuUtilInfo)) if ret == NVML_SUCCESS: # special case, no active vGPUs return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case sampleArray = ( c_vgpuUtilInfo.vgpuInstanceCount * c_nvmlVgpuInstanceUtilizationInfo_v1_t ) c_samples = sampleArray() c_vgpuUtilInfo.vgpuUtilArray = c_samples # make the call again ret = fn(handle, byref(c_vgpuUtilInfo)) _nvmlCheckReturn(ret) return c_samples[0 : c_vgpuUtilInfo.vgpuInstanceCount] else: # error case raise NVMLError(ret) def nvmlDeviceGetP2PStatus(device1, device2, p2pIndex): c_p2pstatus = _nvmlGpuP2PStatus_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetP2PStatus") ret = fn(device1, device2, p2pIndex, byref(c_p2pstatus)) _nvmlCheckReturn(ret) return c_p2pstatus.value def nvmlDeviceGetGridLicensableFeatures_v4(handle): c_get_grid_licensable_features = c_nvmlGridLicensableFeatures_v4_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGridLicensableFeatures_v4") ret = fn(handle, byref(c_get_grid_licensable_features)) _nvmlCheckReturn(ret) return c_get_grid_licensable_features def nvmlDeviceGetGridLicensableFeatures(handle): return nvmlDeviceGetGridLicensableFeatures_v4(handle) def nvmlDeviceGetGspFirmwareVersion(handle, version=None): isUserDefined = version is not None if not isUserDefined: version = (c_char * NVML_GSP_FIRMWARE_VERSION_BUF_SIZE)() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGspFirmwareVersion") ret = fn(handle, version) _nvmlCheckReturn(ret) return NVML_SUCCESS if isUserDefined else version.value def nvmlDeviceGetGspFirmwareMode(handle, isEnabled=c_uint(), defaultMode=c_uint()): isReference = type(isEnabled) is not c_uint isEnabledRef = isEnabled if isReference else byref(isEnabled) defaultModeRef = defaultMode if isReference else byref(defaultMode) fn = _nvmlGetFunctionPointer("nvmlDeviceGetGspFirmwareMode") ret = fn(handle, isEnabledRef, defaultModeRef) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else [isEnabled.value, defaultMode.value] def nvmlDeviceGetEncoderCapacity(handle, encoderQueryType): c_encoder_capacity = c_ulonglong(0) c_encoderQuery_type = _nvmlEncoderQueryType_t(encoderQueryType) fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderCapacity") ret = fn(handle, c_encoderQuery_type, byref(c_encoder_capacity)) _nvmlCheckReturn(ret) return c_encoder_capacity.value def nvmlDeviceGetVgpuProcessUtilization(handle, timeStamp): # first call to get the size c_vgpu_count = c_uint(0) c_time_stamp = c_ulonglong(timeStamp) fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessUtilization") ret = fn(handle, c_time_stamp, byref(c_vgpu_count), None) if ret == NVML_SUCCESS: # special case, no active vGPUs return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case sampleArray = c_vgpu_count.value * c_nvmlVgpuProcessUtilizationSample_t c_samples = sampleArray() # make the call again ret = fn(handle, c_time_stamp, byref(c_vgpu_count), c_samples) _nvmlCheckReturn(ret) return c_samples[0 : c_vgpu_count.value] else: # error case raise NVMLError(ret) def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp): # first call to get the size c_time_stamp = c_ulonglong(timeStamp) c_vgpuProcUtilInfo = c_nvmlVgpuProcessesUtilizationInfo_v1_t(0) c_vgpuProcUtilInfo.version = VgpuProcessesUtilizationInfo_v1 c_vgpuProcUtilInfo.vgpuProcessCount = c_uint(0) c_vgpuProcUtilInfo.lastSeenTimeStamp = c_time_stamp fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessesUtilizationInfo") ret = fn(handle, byref(c_vgpuProcUtilInfo)) if ret == NVML_SUCCESS: # special case, no active vGPUs return [] elif ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case sampleArray = ( c_vgpuProcUtilInfo.vgpuProcessCount * c_nvmlVgpuProcessUtilizationInfo_v1_t ) c_samples = sampleArray() c_vgpuProcUtilInfo.vgpuProcUtilArray = c_samples # make the call again ret = fn(handle, byref(c_vgpuProcUtilInfo)) _nvmlCheckReturn(ret) return c_samples[0 : c_vgpuProcUtilInfo.vgpuProcessCount] else: # error case raise NVMLError(ret) def nvmlDeviceGetEncoderStats(handle): c_encoderCount = c_ulonglong(0) c_encodeFps = c_ulonglong(0) c_encoderLatency = c_ulonglong(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderStats") ret = fn(handle, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency)) _nvmlCheckReturn(ret) return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) def nvmlDeviceGetEncoderSessions(handle): # first call to get the size c_session_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderSessions") ret = fn(handle, byref(c_session_count), None) if ret == NVML_SUCCESS: if c_session_count.value != 0: # typical case session_array = c_nvmlEncoderSession_t * c_session_count.value c_sessions = session_array() # make the call again ret = fn(handle, byref(c_session_count), c_sessions) _nvmlCheckReturn(ret) sessions = [] for i in range(c_session_count.value): sessions.append(c_sessions[i]) return sessions else: return [] # no active sessions else: # error case raise NVMLError(ret) def nvmlDeviceGetFBCStats(handle): c_fbcStats = c_nvmlFBCStats_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCStats") ret = fn(handle, byref(c_fbcStats)) _nvmlCheckReturn(ret) return c_fbcStats def nvmlDeviceGetFBCSessions(handle): # first call to get the size c_session_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCSessions") ret = fn(handle, byref(c_session_count), None) if ret == NVML_SUCCESS: if c_session_count.value != 0: # typical case session_array = c_nvmlFBCSession_t * c_session_count.value c_sessions = session_array() # make the call again ret = fn(handle, byref(c_session_count), c_sessions) _nvmlCheckReturn(ret) sessions = [] for i in range(c_session_count.value): sessions.append(c_sessions[i]) return sessions else: return [] # no active sessions else: # error case raise NVMLError(ret) def nvmlVgpuInstanceGetEncoderStats(vgpuInstance): c_encoderCount = c_ulonglong(0) c_encodeFps = c_ulonglong(0) c_encoderLatency = c_ulonglong(0) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderStats") ret = fn( vgpuInstance, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency) ) _nvmlCheckReturn(ret) return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) def nvmlVgpuInstanceGetEncoderSessions(vgpuInstance): # first call to get the size c_session_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderSessions") ret = fn(vgpuInstance, byref(c_session_count), None) if ret == NVML_SUCCESS: if c_session_count.value != 0: # typical case session_array = c_nvmlEncoderSession_t * c_session_count.value c_sessions = session_array() # make the call again ret = fn(vgpuInstance, byref(c_session_count), c_sessions) _nvmlCheckReturn(ret) sessions = [] for i in range(c_session_count.value): sessions.append(c_sessions[i]) return sessions else: return [] # no active sessions else: # error case raise NVMLError(ret) def nvmlVgpuInstanceGetFBCStats(vgpuInstance): c_fbcStats = c_nvmlFBCStats_t() fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCStats") ret = fn(vgpuInstance, byref(c_fbcStats)) _nvmlCheckReturn(ret) return c_fbcStats def nvmlVgpuInstanceGetFBCSessions(vgpuInstance): # first call to get the size c_session_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCSessions") ret = fn(vgpuInstance, byref(c_session_count), None) if ret == NVML_SUCCESS: if c_session_count.value != 0: # typical case session_array = c_nvmlFBCSession_t * c_session_count.value c_sessions = session_array() # make the call again ret = fn(vgpuInstance, byref(c_session_count), c_sessions) _nvmlCheckReturn(ret) sessions = [] for i in range(c_session_count.value): sessions.append(c_sessions[i]) return sessions else: return [] # no active sessions else: # error case raise NVMLError(ret) def nvmlDeviceGetProcessUtilization(handle, timeStamp): # first call to get the size c_count = c_uint(0) c_time_stamp = c_ulonglong(timeStamp) fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessUtilization") ret = fn(handle, None, byref(c_count), c_time_stamp) if ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case sampleArray = c_count.value * c_nvmlProcessUtilizationSample_t c_samples = sampleArray() # make the call again ret = fn(handle, c_samples, byref(c_count), c_time_stamp) _nvmlCheckReturn(ret) return c_samples[0 : c_count.value] else: # error case raise NVMLError(ret) def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp): # first call to get the size c_time_stamp = c_ulonglong(timeStamp) c_processesUtilInfo = c_nvmlProcessesUtilizationInfo_v1_t(0) c_processesUtilInfo.version = ProcessesUtilizationInfo_v1 c_processesUtilInfo.processSamplesCount = c_uint(0) c_processesUtilInfo.lastSeenTimeStamp = c_time_stamp fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessesUtilizationInfo") ret = fn(handle, byref(c_processesUtilInfo)) if ret == NVML_ERROR_INSUFFICIENT_SIZE: # typical case sampleArray = ( c_processesUtilInfo.processSamplesCount * c_nvmlProcessUtilizationInfo_v1_t ) c_samples = sampleArray() c_processesUtilInfo.procUtilArray = c_samples # make the call again ret = fn(handle, byref(c_processesUtilInfo)) _nvmlCheckReturn(ret) return c_samples[0 : c_processesUtilInfo.processSamplesCount] else: # error case raise NVMLError(ret) def nvmlVgpuInstanceGetMetadata(vgpuInstance): fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMetadata") c_vgpuMetadata = c_nvmlVgpuMetadata_t() c_bufferSize = c_uint(0) # Make the first NVML API call to get the c_bufferSize value. # We have already allocated required buffer above. ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize)) if ret == NVML_ERROR_INSUFFICIENT_SIZE: ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize)) _nvmlCheckReturn(ret) else: raise NVMLError(ret) return c_vgpuMetadata def nvmlDeviceGetVgpuMetadata(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuMetadata") c_vgpuPgpuMetadata = c_nvmlVgpuPgpuMetadata_t() c_bufferSize = c_uint(0) # Make the first NVML API call to get the c_bufferSize value. # We have already allocated required buffer above. ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize)) if ret == NVML_ERROR_INSUFFICIENT_SIZE: ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize)) _nvmlCheckReturn(ret) else: raise NVMLError(ret) return c_vgpuPgpuMetadata def nvmlGetVgpuCompatibility(vgpuMetadata, pgpuMetadata): fn = _nvmlGetFunctionPointer("nvmlGetVgpuCompatibility") c_vgpuPgpuCompatibility = c_nvmlVgpuPgpuCompatibility_t() ret = fn(byref(vgpuMetadata), byref(pgpuMetadata), byref(c_vgpuPgpuCompatibility)) _nvmlCheckReturn(ret) return c_vgpuPgpuCompatibility @convertStrBytes def nvmlDeviceGetPgpuMetadataString(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetPgpuMetadataString") c_pgpuMetadata = create_string_buffer(NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE) c_bufferSize = c_uint(0) # Make the first NVML API call to get the c_bufferSize value. # We have already allocated required buffer above. ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize)) if ret == NVML_ERROR_INSUFFICIENT_SIZE: ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize)) _nvmlCheckReturn(ret) else: raise NVMLError(ret) return (c_pgpuMetadata.value, c_bufferSize.value) def nvmlDeviceGetVgpuSchedulerLog(handle): c_vgpu_sched_log = c_nvmlVgpuSchedulerLog_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerLog") ret = fn(handle, byref(c_vgpu_sched_log)) _nvmlCheckReturn(ret) return c_vgpu_sched_log def nvmlDeviceGetVgpuSchedulerState(handle): c_vgpu_sched_state = c_nvmlVgpuSchedulerGetState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerState") ret = fn(handle, byref(c_vgpu_sched_state)) _nvmlCheckReturn(ret) return c_vgpu_sched_state def nvmlDeviceGetVgpuSchedulerCapabilities(handle): c_vgpu_sched_caps = c_nvmlVgpuSchedulerCapabilities_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerCapabilities") ret = fn(handle, byref(c_vgpu_sched_caps)) _nvmlCheckReturn(ret) return c_vgpu_sched_caps def nvmlDeviceSetVgpuSchedulerState(handle, sched_state): fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuSchedulerState") ret = fn(handle, byref(sched_state)) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlSetVgpuVersion(vgpuVersion): fn = _nvmlGetFunctionPointer("nvmlSetVgpuVersion") ret = fn(byref(vgpuVersion)) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlGetVgpuVersion(supported=None, current=None): isUserDefined = (supported is not None) or (current is not None) if not isUserDefined: supported = c_nvmlVgpuVersion_t() current = c_nvmlVgpuVersion_t() fn = _nvmlGetFunctionPointer("nvmlGetVgpuVersion") ret = fn(byref(supported), byref(current)) _nvmlCheckReturn(ret) return ( NVML_SUCCESS if isUserDefined else [ (supported.minVersion, supported.maxVersion), (current.minVersion, current.maxVersion), ] ) def nvmlVgpuInstanceGetAccountingMode(vgpuInstance): c_mode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingMode") ret = fn(vgpuInstance, byref(c_mode)) _nvmlCheckReturn(ret) return c_mode.value def nvmlVgpuInstanceGetAccountingPids(vgpuInstance): c_pidCount = c_uint() fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingPids") ret = fn(vgpuInstance, byref(c_pidCount), None) if ret == NVML_ERROR_INSUFFICIENT_SIZE: sampleArray = c_pidCount.value * c_uint c_pidArray = sampleArray() ret = fn(vgpuInstance, byref(c_pidCount), byref(c_pidArray)) _nvmlCheckReturn(ret) else: raise NVMLError(ret) return (c_pidCount, c_pidArray) def nvmlVgpuInstanceGetAccountingStats(vgpuInstance, pid): c_accountingStats = c_nvmlAccountingStats_t() fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingStats") ret = fn(vgpuInstance, pid, byref(c_accountingStats)) _nvmlCheckReturn(ret) return c_accountingStats def nvmlVgpuInstanceClearAccountingPids(vgpuInstance): fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceClearAccountingPids") ret = fn(vgpuInstance) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlGetExcludedDeviceCount(): c_count = c_uint() fn = _nvmlGetFunctionPointer("nvmlGetExcludedDeviceCount") ret = fn(byref(c_count)) _nvmlCheckReturn(ret) return c_count.value def nvmlGetExcludedDeviceInfoByIndex(index): c_index = c_uint(index) info = c_nvmlExcludedDeviceInfo_t() fn = _nvmlGetFunctionPointer("nvmlGetExcludedDeviceInfoByIndex") ret = fn(c_index, byref(info)) _nvmlCheckReturn(ret) return info def nvmlDeviceGetHostVgpuMode(handle): c_host_vgpu_mode = _nvmlHostVgpuMode_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetHostVgpuMode") ret = fn(handle, byref(c_host_vgpu_mode)) _nvmlCheckReturn(ret) return c_host_vgpu_mode.value def nvmlDeviceSetMigMode(device, mode): c_activationStatus = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceSetMigMode") ret = fn(device, mode, byref(c_activationStatus)) _nvmlCheckReturn(ret) return c_activationStatus.value def nvmlDeviceGetMigMode(device): c_currentMode = c_uint() c_pendingMode = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMigMode") ret = fn(device, byref(c_currentMode), byref(c_pendingMode)) _nvmlCheckReturn(ret) return [c_currentMode.value, c_pendingMode.value] def nvmlDeviceGetGpuInstanceProfileInfo(device, profile, version=2): if version == 2: c_info = c_nvmlGpuInstanceProfileInfo_v2_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceProfileInfoV") elif version == 1: c_info = c_nvmlGpuInstanceProfileInfo_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceProfileInfo") else: raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) ret = fn(device, profile, byref(c_info)) _nvmlCheckReturn(ret) return c_info # Define function alias for the API exposed by NVML nvmlDeviceGetGpuInstanceProfileInfoV = nvmlDeviceGetGpuInstanceProfileInfo def nvmlDeviceGetGpuInstanceRemainingCapacity(device, profileId): c_count = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceRemainingCapacity") ret = fn(device, profileId, byref(c_count)) _nvmlCheckReturn(ret) return c_count.value def nvmlDeviceGetGpuInstancePossiblePlacements( device, profileId, placementsRef, countRef ): fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstancePossiblePlacements_v2") ret = fn(device, profileId, placementsRef, countRef) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceCreateGpuInstance(device, profileId): c_instance = c_nvmlGpuInstance_t() fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstance") ret = fn(device, profileId, byref(c_instance)) _nvmlCheckReturn(ret) return c_instance def nvmlDeviceCreateGpuInstanceWithPlacement(device, profileId, placement): c_instance = c_nvmlGpuInstance_t() fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstanceWithPlacement") ret = fn(device, profileId, placement, byref(c_instance)) _nvmlCheckReturn(ret) return c_instance def nvmlGpuInstanceDestroy(gpuInstance): fn = _nvmlGetFunctionPointer("nvmlGpuInstanceDestroy") ret = fn(gpuInstance) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceGetGpuInstances(device, profileId, gpuInstancesRef, countRef): fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstances") ret = fn(device, profileId, gpuInstancesRef, countRef) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceGetGpuInstanceById(device, gpuInstanceId): c_instance = c_nvmlGpuInstance_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceById") ret = fn(device, gpuInstanceId, byref(c_instance)) _nvmlCheckReturn(ret) return c_instance def nvmlGpuInstanceGetInfo(gpuInstance): c_info = c_nvmlGpuInstanceInfo_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetInfo") ret = fn(gpuInstance, byref(c_info)) _nvmlCheckReturn(ret) return c_info def nvmlGpuInstanceGetComputeInstanceProfileInfo( device, profile, engProfile, version=2 ): if version == 2: c_info = c_nvmlComputeInstanceProfileInfo_v2_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceProfileInfoV") elif version == 1: c_info = c_nvmlComputeInstanceProfileInfo_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceProfileInfo") else: raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) ret = fn(device, profile, engProfile, byref(c_info)) _nvmlCheckReturn(ret) return c_info # Define function alias for the API exposed by NVML nvmlGpuInstanceGetComputeInstanceProfileInfoV = ( nvmlGpuInstanceGetComputeInstanceProfileInfo ) def nvmlGpuInstanceGetComputeInstanceRemainingCapacity(gpuInstance, profileId): c_count = c_uint() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceRemainingCapacity") ret = fn(gpuInstance, profileId, byref(c_count)) _nvmlCheckReturn(ret) return c_count.value def nvmlGpuInstanceGetComputeInstancePossiblePlacements( gpuInstance, profileId, placementsRef, countRef ): fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstancePossiblePlacements") ret = fn(gpuInstance, profileId, placementsRef, countRef) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlGpuInstanceCreateComputeInstance(gpuInstance, profileId): c_instance = c_nvmlComputeInstance_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstance") ret = fn(gpuInstance, profileId, byref(c_instance)) _nvmlCheckReturn(ret) return c_instance def nvmlGpuInstanceCreateComputeInstanceWithPlacement( gpuInstance, profileId, placement ): c_instance = c_nvmlComputeInstance_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstanceWithPlacement") ret = fn(gpuInstance, profileId, placement, byref(c_instance)) _nvmlCheckReturn(ret) return c_instance def nvmlComputeInstanceDestroy(computeInstance): fn = _nvmlGetFunctionPointer("nvmlComputeInstanceDestroy") ret = fn(computeInstance) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlGpuInstanceGetComputeInstances( gpuInstance, profileId, computeInstancesRef, countRef ): fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstances") ret = fn(gpuInstance, profileId, computeInstancesRef, countRef) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlGpuInstanceGetComputeInstanceById(gpuInstance, computeInstanceId): c_instance = c_nvmlComputeInstance_t() fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceById") ret = fn(gpuInstance, computeInstanceId, byref(c_instance)) _nvmlCheckReturn(ret) return c_instance def nvmlComputeInstanceGetInfo_v2(computeInstance): c_info = c_nvmlComputeInstanceInfo_t() fn = _nvmlGetFunctionPointer("nvmlComputeInstanceGetInfo_v2") ret = fn(computeInstance, byref(c_info)) _nvmlCheckReturn(ret) return c_info def nvmlComputeInstanceGetInfo(computeInstance): return nvmlComputeInstanceGetInfo_v2(computeInstance) def nvmlDeviceIsMigDeviceHandle(device): c_isMigDevice = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceIsMigDeviceHandle") ret = fn(device, byref(c_isMigDevice)) _nvmlCheckReturn(ret) return c_isMigDevice def nvmlDeviceGetGpuInstanceId(device): c_gpuInstanceId = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceId") ret = fn(device, byref(c_gpuInstanceId)) _nvmlCheckReturn(ret) return c_gpuInstanceId.value def nvmlDeviceGetComputeInstanceId(device): c_computeInstanceId = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeInstanceId") ret = fn(device, byref(c_computeInstanceId)) _nvmlCheckReturn(ret) return c_computeInstanceId.value def nvmlDeviceGetMaxMigDeviceCount(device): c_count = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxMigDeviceCount") ret = fn(device, byref(c_count)) _nvmlCheckReturn(ret) return c_count.value def nvmlDeviceGetMigDeviceHandleByIndex(device, index): c_index = c_uint(index) migDevice = c_nvmlDevice_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMigDeviceHandleByIndex") ret = fn(device, c_index, byref(migDevice)) _nvmlCheckReturn(ret) return migDevice def nvmlDeviceGetDeviceHandleFromMigDeviceHandle(migDevice): device = c_nvmlDevice_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDeviceHandleFromMigDeviceHandle") ret = fn(migDevice, byref(device)) _nvmlCheckReturn(ret) return device def nvmlDeviceGetAttributes_v2(device): c_attrs = c_nvmlDeviceAttributes() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAttributes_v2") ret = fn(device, byref(c_attrs)) _nvmlCheckReturn(ret) return c_attrs def nvmlDeviceGetAttributes(device): return nvmlDeviceGetAttributes_v2(device) def nvmlDeviceGetRemappedRows(device): fn = _nvmlGetFunctionPointer("nvmlDeviceGetRemappedRows") c_corr = c_uint() c_unc = c_uint() c_bpending = c_uint() c_bfailure = c_uint() ret = fn(device, byref(c_corr), byref(c_unc), byref(c_bpending), byref(c_bfailure)) _nvmlCheckReturn(ret) return (c_corr.value, c_unc.value, c_bpending.value, c_bfailure.value) def nvmlDeviceGetRowRemapperHistogram(device): c_vals = c_nvmlRowRemapperHistogramValues() fn = _nvmlGetFunctionPointer("nvmlDeviceGetRowRemapperHistogram") ret = fn(device, byref(c_vals)) _nvmlCheckReturn(ret) return c_vals def nvmlDeviceGetArchitecture(device): arch = _nvmlDeviceArchitecture_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetArchitecture") ret = fn(device, byref(arch)) _nvmlCheckReturn(ret) return arch.value def nvmlDeviceGetBusType(device): c_busType = _nvmlBusType_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetBusType") ret = fn(device, byref(c_busType)) _nvmlCheckReturn(ret) return c_busType.value def nvmlDeviceGetIrqNum(device): c_irqNum = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetIrqNum") ret = fn(device, byref(c_irqNum)) _nvmlCheckReturn(ret) return c_irqNum.value def nvmlDeviceGetNumGpuCores(device): c_numCores = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumGpuCores") ret = fn(device, byref(c_numCores)) _nvmlCheckReturn(ret) return c_numCores.value def nvmlDeviceGetPowerSource(device): c_powerSource = _nvmlPowerSource_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerSource") ret = fn(device, byref(c_powerSource)) _nvmlCheckReturn(ret) return c_powerSource.value def nvmlDeviceGetMemoryBusWidth(device): c_memBusWidth = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryBusWidth") ret = fn(device, byref(c_memBusWidth)) _nvmlCheckReturn(ret) return c_memBusWidth.value def nvmlDeviceGetPcieLinkMaxSpeed(device): c_speed = _nvmlPcieLinkMaxSpeed_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieLinkMaxSpeed") ret = fn(device, byref(c_speed)) _nvmlCheckReturn(ret) return c_speed.value def nvmlDeviceGetAdaptiveClockInfoStatus(device): c_adaptiveClockInfoStatus = _nvmlAdaptiveClockInfoStatus_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetAdaptiveClockInfoStatus") ret = fn(device, byref(c_adaptiveClockInfoStatus)) _nvmlCheckReturn(ret) return c_adaptiveClockInfoStatus.value def nvmlDeviceGetPcieSpeed(device): c_speed = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieSpeed") ret = fn(device, byref(c_speed)) _nvmlCheckReturn(ret) return c_speed.value def nvmlDeviceGetDynamicPstatesInfo( device, c_dynamicpstatesinfo=c_nvmlGpuDynamicPstatesInfo_t() ): isReference = type(c_dynamicpstatesinfo) is not c_nvmlGpuDynamicPstatesInfo_t dynamicpstatesinfoRef = ( c_dynamicpstatesinfo if isReference else byref(c_dynamicpstatesinfo) ) fn = _nvmlGetFunctionPointer("nvmlDeviceGetDynamicPstatesInfo") ret = fn(device, dynamicpstatesinfoRef) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else c_dynamicpstatesinfo def nvmlDeviceSetFanSpeed_v2(handle, index, speed): fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanSpeed_v2") ret = fn(handle, index, speed) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceGetThermalSettings( device, sensorindex, c_thermalsettings=c_nvmlGpuThermalSettings_t() ): isReference = type(c_thermalsettings) is not c_nvmlGpuThermalSettings_t thermalsettingsRef = c_thermalsettings if isReference else byref(c_thermalsettings) fn = _nvmlGetFunctionPointer("nvmlDeviceGetThermalSettings") ret = fn(device, sensorindex, thermalsettingsRef) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else c_thermalsettings.sensor[:] def nvmlDeviceGetMinMaxClockOfPState( device, clockType, pstate, minClockMHz=c_uint(), maxClockMHz=c_uint() ): isReference = (type(minClockMHz) is not c_uint) or (type(maxClockMHz) is not c_uint) minClockMHzRef = minClockMHz if isReference else byref(minClockMHz) maxClockMHzRef = maxClockMHz if isReference else byref(maxClockMHz) fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxClockOfPState") ret = fn( device, _nvmlClockType_t(clockType), _nvmlClockType_t(pstate), minClockMHzRef, maxClockMHzRef, ) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else (minClockMHz.value, maxClockMHz.value) class c_nvmlClockOffset_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("type", _nvmlClockType_t), ("pstate", _nvmlPstates_t), ("clockOffsetMHz", c_int), ("minClockOffsetMHz", c_int), ("maxClockOffsetMHz", c_int), ] nvmlClockOffset_v1 = 0x1000018 def nvmlDeviceGetClockOffsets(device, info): fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockOffsets") ret = fn(device, info) return NVML_SUCCESS def nvmlDeviceSetClockOffsets(device, info): fn = _nvmlGetFunctionPointer("nvmlDeviceSetClockOffsets") ret = fn(device, info) return NVML_SUCCESS def nvmlDeviceGetSupportedPerformanceStates(device): pstates = [] c_count = c_uint(NVML_MAX_GPU_PERF_PSTATES) c_size = sizeof(c_uint) * c_count.value # NOTE: use 'c_uint' to represent the size of the nvmlPstate_t enumeration. pstates_array = _nvmlPstates_t * c_count.value c_pstates = pstates_array() fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedPerformanceStates") ret = fn(device, c_pstates, c_size) _nvmlCheckReturn(ret) for value in c_pstates: if value != NVML_PSTATE_UNKNOWN: pstates.append(value) return pstates def nvmlDeviceGetGpcClkVfOffset(device): offset = c_int32() fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpcClkVfOffset") ret = fn(device, byref(offset)) _nvmlCheckReturn(ret) return offset.value def nvmlDeviceSetGpcClkVfOffset(device, offset): c_offset = c_int32(offset) fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpcClkVfOffset") ret = fn(device, c_offset) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceGetGpcClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()): isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int) minOffsetRef = minOffset if isReference else byref(minOffset) maxOffsetRef = maxOffset if isReference else byref(maxOffset) fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpcClkMinMaxVfOffset") ret = fn(device, minOffsetRef, maxOffsetRef) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) def nvmlDeviceGetMemClkVfOffset(device): offset = c_int32() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemClkVfOffset") ret = fn(device, byref(offset)) _nvmlCheckReturn(ret) return offset.value def nvmlDeviceSetMemClkVfOffset(device, offset): c_offset = c_int32(offset) fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemClkVfOffset") ret = fn(device, c_offset) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceGetMemClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()): isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int) minOffsetRef = minOffset if isReference else byref(minOffset) maxOffsetRef = maxOffset if isReference else byref(maxOffset) fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemClkMinMaxVfOffset") ret = fn(device, minOffsetRef, maxOffsetRef) _nvmlCheckReturn(ret) return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) def nvmlSystemSetConfComputeGpusReadyState(state): c_state = c_uint(state) fn = _nvmlGetFunctionPointer("nvmlSystemSetConfComputeGpusReadyState") ret = fn(c_state) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlSystemGetConfComputeGpusReadyState(): c_state = c_uint() fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeGpusReadyState") ret = fn(byref(c_state)) _nvmlCheckReturn(ret) return c_state.value def nvmlSystemGetConfComputeCapabilities(): c_ccSysCaps = c_nvmlConfComputeSystemCaps_t() fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeCapabilities") ret = fn(byref(c_ccSysCaps)) _nvmlCheckReturn(ret) return c_ccSysCaps def nvmlSystemGetConfComputeState(): c_state = c_nvmlConfComputeSystemState_t() fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeState") ret = fn(byref(c_state)) _nvmlCheckReturn(ret) return c_state def nvmlSystemGetConfComputeSettings(settings): fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeSettings") return fn(settings) def nvmlDeviceSetConfComputeUnprotectedMemSize(device, c_ccMemSize): fn = _nvmlGetFunctionPointer("nvmlDeviceSetConfComputeUnprotectedMemSize") ret = fn(device, c_ccMemSize) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceGetConfComputeMemSizeInfo(device): c_ccMemSize = c_nvmlConfComputeMemSizeInfo_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeMemSizeInfo") ret = fn(device, byref(c_ccMemSize)) _nvmlCheckReturn(ret) return c_ccMemSize def nvmlDeviceGetConfComputeProtectedMemoryUsage(device): c_memory = c_nvmlMemory_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeProtectedMemoryUsage") ret = fn(device, byref(c_memory)) _nvmlCheckReturn(ret) return c_memory def nvmlDeviceGetConfComputeGpuCertificate(device): c_cert = c_nvmlConfComputeGpuCertificate_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuCertificate") ret = fn(device, byref(c_cert)) _nvmlCheckReturn(ret) return c_cert def nvmlDeviceGetConfComputeGpuAttestationReport(device, c_nonce): c_attestReport = c_nvmlConfComputeGpuAttestationReport_t() c_nonce_arr = (c_uint8 * len(c_nonce))(*(c_nonce)) setattr(c_attestReport, "nonce", c_nonce_arr) fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuAttestationReport") ret = fn(device, byref(c_attestReport)) _nvmlCheckReturn(ret) return c_attestReport def nvmlSystemSetConfComputeKeyRotationThresholdInfo(max_atk_adv): c_keyRotationThrInfo = c_nvmlConfComputeSetKeyRotationThresholdInfo_t(0) c_keyRotationThrInfo.version = ConfComputeSetKeyRotationThresholdInfo_v1 c_keyRotationThrInfo.maxAttackerAdvantage = max_atk_adv fn = _nvmlGetFunctionPointer("nvmlSystemSetConfComputeKeyRotationThresholdInfo") ret = fn(byref(c_keyRotationThrInfo)) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlSystemGetConfComputeKeyRotationThresholdInfo(): c_keyRotationThrInfo = c_nvmlConfComputeGetKeyRotationThresholdInfo_t(0) c_keyRotationThrInfo.version = ConfComputeGetKeyRotationThresholdInfo_v1 fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeKeyRotationThresholdInfo") ret = fn(byref(c_keyRotationThrInfo)) _nvmlCheckReturn(ret) return c_keyRotationThrInfo ## GPM ## ######### ## Enums/defines #### GPM Metric Identifiers NVML_GPM_METRIC_GRAPHICS_UTIL = ( 1 # Percentage of time any compute/graphics app was active on the GPU. 0.0 - 100.0 ) NVML_GPM_METRIC_SM_UTIL = 2 # Percentage of SMs that were busy. 0.0 - 100.0 NVML_GPM_METRIC_SM_OCCUPANCY = ( 3 # Percentage of warps that were active vs theoretical maximum. 0.0 - 100.0 ) NVML_GPM_METRIC_INTEGER_UTIL = ( 4 # Percentage of time the GPU's SMs were doing integer operations. 0.0 - 100.0 ) NVML_GPM_METRIC_ANY_TENSOR_UTIL = ( 5 # Percentage of time the GPU's SMs were doing ANY tensor operations. 0.0 - 100.0 ) NVML_GPM_METRIC_DFMA_TENSOR_UTIL = ( 6 # Percentage of time the GPU's SMs were doing DFMA tensor operations. 0.0 - 100.0 ) NVML_GPM_METRIC_HMMA_TENSOR_UTIL = ( 7 # Percentage of time the GPU's SMs were doing HMMA tensor operations. 0.0 - 100.0 ) NVML_GPM_METRIC_IMMA_TENSOR_UTIL = ( 9 # Percentage of time the GPU's SMs were doing IMMA tensor operations. 0.0 - 100.0 ) NVML_GPM_METRIC_DRAM_BW_UTIL = ( 10 # Percentage of DRAM bw used vs theoretical maximum. 0.0 - 100.0 ) NVML_GPM_METRIC_FP64_UTIL = ( 11 # Percentage of time the GPU's SMs were doing non-tensor FP64 math. 0.0 - 100.0 ) NVML_GPM_METRIC_FP32_UTIL = ( 12 # Percentage of time the GPU's SMs were doing non-tensor FP32 math. 0.0 - 100.0 ) NVML_GPM_METRIC_FP16_UTIL = ( 13 # Percentage of time the GPU's SMs were doing non-tensor FP16 math. 0.0 - 100.0 ) NVML_GPM_METRIC_PCIE_TX_PER_SEC = 20 # PCIe traffic from this GPU in MiB/sec NVML_GPM_METRIC_PCIE_RX_PER_SEC = 21 # PCIe traffic to this GPU in MiB/sec NVML_GPM_METRIC_NVDEC_0_UTIL = 30 # Percent utilization of NVDEC 0. 0.0 - 100.0 NVML_GPM_METRIC_NVDEC_1_UTIL = 31 # Percent utilization of NVDEC 1. 0.0 - 100.0 NVML_GPM_METRIC_NVDEC_2_UTIL = 32 # Percent utilization of NVDEC 2. 0.0 - 100.0 NVML_GPM_METRIC_NVDEC_3_UTIL = 33 # Percent utilization of NVDEC 3. 0.0 - 100.0 NVML_GPM_METRIC_NVDEC_4_UTIL = 34 # Percent utilization of NVDEC 4. 0.0 - 100.0 NVML_GPM_METRIC_NVDEC_5_UTIL = 35 # Percent utilization of NVDEC 5. 0.0 - 100.0 NVML_GPM_METRIC_NVDEC_6_UTIL = 36 # Percent utilization of NVDEC 6. 0.0 - 100.0 NVML_GPM_METRIC_NVDEC_7_UTIL = 37 # Percent utilization of NVDEC 7. 0.0 - 100.0 NVML_GPM_METRIC_NVJPG_0_UTIL = 40 # Percent utilization of NVJPG 0. 0.0 - 100.0 NVML_GPM_METRIC_NVJPG_1_UTIL = 41 # Percent utilization of NVJPG 1. 0.0 - 100.0 NVML_GPM_METRIC_NVJPG_2_UTIL = 42 # Percent utilization of NVJPG 2. 0.0 - 100.0 NVML_GPM_METRIC_NVJPG_3_UTIL = 43 # Percent utilization of NVJPG 3. 0.0 - 100.0 NVML_GPM_METRIC_NVJPG_4_UTIL = 44 # Percent utilization of NVJPG 4. 0.0 - 100.0 NVML_GPM_METRIC_NVJPG_5_UTIL = 45 # Percent utilization of NVJPG 5. 0.0 - 100.0 NVML_GPM_METRIC_NVJPG_6_UTIL = 46 # Percent utilization of NVJPG 6. 0.0 - 100.0 NVML_GPM_METRIC_NVJPG_7_UTIL = 47 # Percent utilization of NVJPG 7. 0.0 - 100.0 NVML_GPM_METRIC_NVOFA_0_UTIL = 50 # Percent utilization of NVOFA 0. 0.0 - 100.0 NVML_GPM_METRIC_NVOFA_1_UTIL = 51 # Percent utilization of NVOFA 1. 0.0 - 100.0 NVML_GPM_METRIC_NVLINK_TOTAL_RX_PER_SEC = ( 60 # NvLink read bandwidth for all links in MiB/sec ) NVML_GPM_METRIC_NVLINK_TOTAL_TX_PER_SEC = ( 61 # NvLink write bandwidth for all links in MiB/sec ) NVML_GPM_METRIC_NVLINK_L0_RX_PER_SEC = 62 # NvLink read bandwidth for link 0 in MiB/sec NVML_GPM_METRIC_NVLINK_L0_TX_PER_SEC = ( 63 # NvLink write bandwidth for link 0 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L1_RX_PER_SEC = 64 # NvLink read bandwidth for link 1 in MiB/sec NVML_GPM_METRIC_NVLINK_L1_TX_PER_SEC = ( 65 # NvLink write bandwidth for link 1 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L2_RX_PER_SEC = 66 # NvLink read bandwidth for link 2 in MiB/sec NVML_GPM_METRIC_NVLINK_L2_TX_PER_SEC = ( 67 # NvLink write bandwidth for link 2 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L3_RX_PER_SEC = 68 # NvLink read bandwidth for link 3 in MiB/sec NVML_GPM_METRIC_NVLINK_L3_TX_PER_SEC = ( 69 # NvLink write bandwidth for link 3 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L4_RX_PER_SEC = 70 # NvLink read bandwidth for link 4 in MiB/sec NVML_GPM_METRIC_NVLINK_L4_TX_PER_SEC = ( 71 # NvLink write bandwidth for link 4 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L5_RX_PER_SEC = 72 # NvLink read bandwidth for link 5 in MiB/sec NVML_GPM_METRIC_NVLINK_L5_TX_PER_SEC = ( 73 # NvLink write bandwidth for link 5 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L6_RX_PER_SEC = 74 # NvLink read bandwidth for link 6 in MiB/sec NVML_GPM_METRIC_NVLINK_L6_TX_PER_SEC = ( 75 # NvLink write bandwidth for link 6 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L7_RX_PER_SEC = 76 # NvLink read bandwidth for link 7 in MiB/sec NVML_GPM_METRIC_NVLINK_L7_TX_PER_SEC = ( 77 # NvLink write bandwidth for link 7 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L8_RX_PER_SEC = 78 # NvLink read bandwidth for link 8 in MiB/sec NVML_GPM_METRIC_NVLINK_L8_TX_PER_SEC = ( 79 # NvLink write bandwidth for link 8 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L9_RX_PER_SEC = 80 # NvLink read bandwidth for link 9 in MiB/sec NVML_GPM_METRIC_NVLINK_L9_TX_PER_SEC = ( 81 # NvLink write bandwidth for link 9 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L10_RX_PER_SEC = ( 82 # NvLink read bandwidth for link 10 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L10_TX_PER_SEC = ( 83 # NvLink write bandwidth for link 10 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L11_RX_PER_SEC = ( 84 # NvLink read bandwidth for link 11 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L11_TX_PER_SEC = ( 85 # NvLink write bandwidth for link 11 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L12_RX_PER_SEC = ( 86 # NvLink read bandwidth for link 12 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L12_TX_PER_SEC = ( 87 # NvLink write bandwidth for link 12 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L13_RX_PER_SEC = ( 88 # NvLink read bandwidth for link 13 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L13_TX_PER_SEC = ( 89 # NvLink write bandwidth for link 13 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L14_RX_PER_SEC = ( 90 # NvLink read bandwidth for link 14 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L14_TX_PER_SEC = ( 91 # NvLink write bandwidth for link 14 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L15_RX_PER_SEC = ( 92 # NvLink read bandwidth for link 15 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L15_TX_PER_SEC = ( 93 # NvLink write bandwidth for link 15 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L16_RX_PER_SEC = ( 94 # NvLink read bandwidth for link 16 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L16_TX_PER_SEC = ( 95 # NvLink write bandwidth for link 16 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L17_RX_PER_SEC = ( 96 # NvLink read bandwidth for link 17 in MiB/sec ) NVML_GPM_METRIC_NVLINK_L17_TX_PER_SEC = ( 97 # NvLink write bandwidth for link 17 in MiB/sec ) NVML_GPM_METRIC_MAX = 98 ## Structs class c_nvmlUnitInfo_t(_PrintableStructure): _fields_ = [ ("name", c_char * 96), ("id", c_char * 96), ("serial", c_char * 96), ("firmwareVersion", c_char * 96), ] class struct_c_nvmlGpmSample_t(Structure): pass # opaque handle c_nvmlGpmSample_t = POINTER(struct_c_nvmlGpmSample_t) class c_metricInfo_t(Structure): _fields_ = [ ("shortName", c_char_p), ("longName", c_char_p), ("unit", c_char_p), ] class c_nvmlGpmMetric_t(_PrintableStructure): _fields_ = [ ("metricId", c_uint), ("nvmlReturn", _nvmlReturn_t), ("value", c_double), ("metricInfo", c_metricInfo_t), ] class c_nvmlGpmMetricsGet_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("numMetrics", c_uint), ("sample1", c_nvmlGpmSample_t), ("sample2", c_nvmlGpmSample_t), ("metrics", c_nvmlGpmMetric_t * NVML_GPM_METRIC_MAX), ] NVML_GPM_METRICS_GET_VERSION = 1 class c_nvmlGpmSupport_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("isSupportedDevice", c_uint), ] NVML_GPM_SUPPORT_VERSION = 1 ## Functions def nvmlGpmMetricsGet(metricsGet): fn = _nvmlGetFunctionPointer("nvmlGpmMetricsGet") ret = fn(byref(metricsGet)) _nvmlCheckReturn(ret) return metricsGet def nvmlGpmSampleFree(gpmSample): fn = _nvmlGetFunctionPointer("nvmlGpmSampleFree") ret = fn(gpmSample) _nvmlCheckReturn(ret) return def nvmlGpmSampleAlloc(): gpmSample = c_nvmlGpmSample_t() fn = _nvmlGetFunctionPointer("nvmlGpmSampleAlloc") ret = fn(byref(gpmSample)) _nvmlCheckReturn(ret) return gpmSample def nvmlGpmSampleGet(device, gpmSample): fn = _nvmlGetFunctionPointer("nvmlGpmSampleGet") ret = fn(device, gpmSample) _nvmlCheckReturn(ret) return gpmSample def nvmlGpmMigSampleGet(device, gpuInstanceId, gpmSample): fn = _nvmlGetFunctionPointer("nvmlGpmMigSampleGet") ret = fn(device, gpuInstanceId, gpmSample) _nvmlCheckReturn(ret) return gpmSample def nvmlGpmQueryDeviceSupport(device): gpmSupport = c_nvmlGpmSupport_t() gpmSupport.version = NVML_GPM_SUPPORT_VERSION fn = _nvmlGetFunctionPointer("nvmlGpmQueryDeviceSupport") ret = fn(device, byref(gpmSupport)) _nvmlCheckReturn(ret) return gpmSupport def nvmlGpmSetStreamingEnabled(device, state): c_state = c_uint(state) fn = _nvmlGetFunctionPointer("nvmlGpmSetStreamingEnabled") ret = fn(device, c_state) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlGpmQueryIfStreamingEnabled(device): c_state = c_uint() fn = _nvmlGetFunctionPointer("nvmlGpmQueryIfStreamingEnabled") ret = fn(device, byref(c_state)) _nvmlCheckReturn(ret) return c_state.value # Low Power Structure and Function NVML_NVLINK_POWER_STATE_HIGH_SPEED = 0x0 NVML_NVLINK_POWER_STATE_LOW = 0x1 NVML_NVLINK_LOW_POWER_THRESHOLD_MIN = 0x1 NVML_NVLINK_LOW_POWER_THRESHOLD_MAX = 0x1FFF NVML_NVLINK_LOW_POWER_THRESHOLD_RESET = 0xFFFFFFFF NVML_NVLINK_LOW_POWER_THRESHOLD_DEFAULT = NVML_NVLINK_LOW_POWER_THRESHOLD_RESET class c_nvmlNvLinkPowerThres_t(Structure): _fields_ = [ ("lowPwrThreshold", c_uint), ] def nvmlDeviceSetNvLinkDeviceLowPowerThreshold(device, l1threshold): c_info = c_nvmlNvLinkPowerThres_t() c_info.lowPwrThreshold = l1threshold fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvLinkDeviceLowPowerThreshold") ret = fn(device, byref(c_info)) _nvmlCheckReturn(ret) return NVML_SUCCESS NVML_GPU_FABRIC_UUID_LEN = 16 _nvmlGpuFabricState_t = c_uint NVML_GPU_FABRIC_STATE_NOT_SUPPORTED = 0 NVML_GPU_FABRIC_STATE_NOT_STARTED = 1 NVML_GPU_FABRIC_STATE_IN_PROGRESS = 2 NVML_GPU_FABRIC_STATE_COMPLETED = 3 class c_nvmlGpuFabricInfo_t(_PrintableStructure): _fields_ = [ ("clusterUuid", c_char * NVML_DEVICE_UUID_BUFFER_SIZE), ("status", _nvmlReturn_t), ("cliqueId", c_uint32), ("state", _nvmlGpuFabricState_t), ] NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_NOT_SUPPORTED = 0 NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_TRUE = 1 NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_FALSE = 2 NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_DEGRADED_BW = 0 NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_DEGRADED_BW = 0x11 NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_NOT_SUPPORTED = 0 NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_TRUE = 1 NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_FALSE = 2 NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_RECOVERY = 2 NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_RECOVERY = 0x11 NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_NOT_SUPPORTED = 0 NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_TRUE = 1 NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_FALSE = 2 NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_UNHEALTHY = 4 NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_UNHEALTHY = 0x11 NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_NOT_SUPPORTED = 0 NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_TRUE = 1 NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_FALSE = 2 NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ACCESS_TIMEOUT_RECOVERY = 6 NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ACCESS_TIMEOUT_RECOVERY = 0x11 nvmlGpuFabricInfo_v2 = 0x02000024 class c_nvmlGpuFabricInfoV_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("clusterUuid", c_char * NVML_GPU_FABRIC_UUID_LEN), ("status", _nvmlReturn_t), ("cliqueId", c_uint32), ("state", _nvmlGpuFabricState_t), ("healthMask", c_uint32), ] def __init__(self): super(c_nvmlGpuFabricInfoV_t, self).__init__(version=nvmlGpuFabricInfo_v2) def nvmlDeviceGetGpuFabricInfo(device, gpuFabricInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfo") ret = fn(device, gpuFabricInfo) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceGetGpuFabricInfoV(device, gpuFabricInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfoV") ret = fn(device, gpuFabricInfo) _nvmlCheckReturn(ret) return NVML_SUCCESS ###################### ## Enums/defines #### NVML GPU NVLINK BW MODE NVML_GPU_NVLINK_BW_MODE_FULL = 0x0 NVML_GPU_NVLINK_BW_MODE_OFF = 0x1 NVML_GPU_NVLINK_BW_MODE_MIN = 0x2 NVML_GPU_NVLINK_BW_MODE_HALF = 0x3 NVML_GPU_NVLINK_BW_MODE_3QUARTER = 0x4 NVML_GPU_NVLINK_BW_MODE_COUNT = 0x5 def nvmlSystemSetNvlinkBwMode(mode): fn = _nvmlGetFunctionPointer("nvmlSystemSetNvlinkBwMode") ret = fn(mode) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlSystemGetNvlinkBwMode(): mode = c_uint() fn = _nvmlGetFunctionPointer("nvmlSystemGetNvlinkBwMode") ret = fn(byref(mode)) _nvmlCheckReturn(ret) return mode.value _nvmlPowerScopeType_t = c_uint NVML_POWER_SCOPE_GPU = 0 NVML_POWER_SCOPE_MODULE = 1 NVML_POWER_SCOPE_MEMORY = 2 class c_nvmlPowerValue_v2_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("powerScope", _nvmlPowerScopeType_t), ("powerValueMw", c_uint), ] _fmt_ = {"": "%d B"} nvmlPowerValue_v2 = 0x0200000C def nvmlDeviceSetPowerManagementLimit_v2( device, powerScope, powerLimit, version=nvmlPowerValue_v2 ): c_powerScope = _nvmlPowerScopeType_t(powerScope) c_powerValue = c_nvmlPowerValue_v2_t() c_powerValue.version = c_uint(version) c_powerValue.powerScope = c_powerScope c_powerValue.powerValueMw = c_uint(powerLimit) fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit_v2") ret = fn(device, byref(c_powerValue)) return NVML_SUCCESS class c_nvmlEccSramErrorStatus_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("aggregateUncParity", c_ulonglong), ("aggregateUncSecDed", c_ulonglong), ("aggregateCor", c_ulonglong), ("volatileUncParity", c_ulonglong), ("volatileUncSecDed", c_ulonglong), ("volatileCor", c_ulonglong), ("aggregateUncBucketL2", c_ulonglong), ("aggregateUncBucketSm", c_ulonglong), ("aggregateUncBucketPcie", c_ulonglong), ("aggregateUncBucketMcu", c_ulonglong), ("aggregateUncBucketOther", c_ulonglong), ("bThresholdExceeded", c_uint), ] def __init__(self): super(c_nvmlEccSramErrorStatus_v1_t, self).__init__( version=nvmlEccSramErrorStatus_v1 ) nvmlEccSramErrorStatus_v1 = 0x1000068 def nvmlDeviceGetSramEccErrorStatus(device, status): fn = _nvmlGetFunctionPointer("nvmlDeviceGetSramEccErrorStatus") ret = fn(device, status) _nvmlCheckReturn(ret) return NVML_SUCCESS NVML_DEV_CAP_EGM = 1 << 0 nvmlDeviceCapabilities_v1 = 0x1000008 class c_nvmlDeviceCapabilities_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("capMask", c_uint), ] def __init__(self): super(c_nvmlDeviceCapabilities_v1_t, self).__init__( version=nvmlDeviceCapabilities_v1 ) def nvmlDeviceGetCapabilities(device, caps): fn = _nvmlGetFunctionPointer("nvmlDeviceGetCapabilities") return fn(device, caps) class c_nvmlPlatformInfo_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("ibGuid", c_char * 16), ("rackGuid", c_char * 16), ("chassisPhysicalSlotNumber", c_char), ("computeSlotIndex", c_char), ("nodeIndex", c_char), ("peerType", c_char), ("moduleId", c_char), ] def __init__(self): super(c_nvmlPlatformInfo_v1_t, self).__init__(version=nvmlPlatformInfo_v1) nvmlPlatformInfo_v1 = 0x100002C def nvmlDeviceGetPlatformInfo(device, platformInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceGetPlatformInfo") ret = fn(device, platformInfo) _nvmlCheckReturn(ret) return NVML_SUCCESS class c_nvmlMask255_t(_PrintableStructure): _fields_ = [ ("mask", c_uint * 8), ] NVML_WORKLOAD_POWER_MAX_PROFILES = 255 NVML_POWER_PROFILE_MAX_P = 0 NVML_POWER_PROFILE_MAX_Q = 1 NVML_POWER_PROFILE_COMPUTE = 2 NVML_POWER_PROFILE_MEMORY_BOUND = 3 NVML_POWER_PROFILE_NETWORK = 4 NVML_POWER_PROFILE_BALANCED = 5 NVML_POWER_PROFILE_LLM_INFERENCE = 6 NVML_POWER_PROFILE_LLM_TRAINING = 7 NVML_POWER_PROFILE_RBM = 8 NVML_POWER_PROFILE_DCPCIE = 9 NVML_POWER_PROFILE_HMMA_SPARSE = 10 NVML_POWER_PROFILE_HMMA_DENSE = 11 NVML_POWER_PROFILE_SYNC_BALANCED = 12 NVML_POWER_PROFILE_HPC = 13 NVML_POWER_PROFILE_MIG = 14 NVML_POWER_PROFILE_MAX = 15 nvmlWorkloadPowerProfileInfo_v1 = 0x100002C class c_nvmlWorkloadPowerProfileInfo_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("profileId", c_uint), ("priority", c_uint), ("conflictingmask", c_nvmlMask255_t), ] def __init__(self): super(c_nvmlWorkloadPowerProfileInfo_v1_t, self).__init__( version=nvmlWorkloadPowerProfileInfo_v1 ) nvmlWorkloadPowerProfileProfilesInfo_v1 = 0x1002BF8 class c_nvmlWorkloadPowerProfileProfilesInfo_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("perfProfilesMask", c_nvmlMask255_t), ( "perfProfile", c_nvmlWorkloadPowerProfileInfo_v1_t * NVML_WORKLOAD_POWER_MAX_PROFILES, ), ] def __init__(self): super(c_nvmlWorkloadPowerProfileProfilesInfo_v1_t, self).__init__( version=nvmlWorkloadPowerProfileProfilesInfo_v1 ) nvmlWorkloadPowerProfileCurrentProfiles_v1 = 0x1000064 class c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("perfProfilesMask", c_nvmlMask255_t), ("requestedProfilesMask", c_nvmlMask255_t), ("enforcedProfilesMask", c_nvmlMask255_t), ] def __init__(self): super(c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t, self).__init__( version=nvmlWorkloadPowerProfileCurrentProfiles_v1 ) nvmlWorkloadPowerProfileRequestedProfiles_v1 = 0x1000024 class c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("requestedProfilesMask", c_nvmlMask255_t), ] def __init__(self): super(c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t, self).__init__( version=nvmlWorkloadPowerProfileRequestedProfiles_v1 ) def nvmlDeviceWorkloadPowerProfileGetProfilesInfo(device, profilesInfo): fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetProfilesInfo") ret = fn(device, profilesInfo) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceWorkloadPowerProfileGetCurrentProfiles(device, currentProfiles): fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetCurrentProfiles") ret = fn(device, currentProfiles) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceWorkloadPowerProfileSetRequestedProfiles(device, requestedProfiles): fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileSetRequestedProfiles") ret = fn(device, requestedProfiles) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceWorkloadPowerProfileClearRequestedProfiles(device, requestedProfiles): fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileClearRequestedProfiles") ret = fn(device, requestedProfiles) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceGetNvlinkSupportedBwModes(device, supportedBwModes): fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkSupportedBwModes") ret = fn(device, supportedBwModes) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceGetNvlinkBwMode(device, getBwMode): fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkBwMode") ret = fn(device, getBwMode) _nvmlCheckReturn(ret) return NVML_SUCCESS def nvmlDeviceSetNvlinkBwMode(device, setBwMode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvlinkBwMode") ret = fn(device, setBwMode) _nvmlCheckReturn(ret) return NVML_SUCCESS nvmlDramEncryptionInfo_v1 = 0x01000008 class c_nvmlDramEncryptionInfo_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("encryptionState", _nvmlEnableState_t), ] def __init__(self): super(c_nvmlDramEncryptionInfo_t, self).__init__( version=nvmlDramEncryptionInfo_v1 ) def nvmlDeviceGetDramEncryptionMode(handle): c_currState = c_nvmlDramEncryptionInfo_t() c_pendingState = c_nvmlDramEncryptionInfo_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDramEncryptionMode") ret = fn(handle, byref(c_currState), byref(c_pendingState)) _nvmlCheckReturn(ret) return [c_currState.encryptionState, c_pendingState.encryptionState] # added to API def nvmlDeviceGetCurrentDramEncryptionMode(handle): return nvmlDeviceGetDramEncryptionMode(handle)[0] # added to API def nvmlDeviceGetPendingDramEncryptionMode(handle): return nvmlDeviceGetDramEncryptionMode(handle)[1] def nvmlDeviceSetDramEncryptionMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetDramEncryptionMode") c_dramEncryptionMode = c_nvmlDramEncryptionInfo_t() c_dramEncryptionMode.encryptionState = mode ret = fn(handle, byref(c_dramEncryptionMode)) _nvmlCheckReturn(ret) return None # Power Smoothing defines NVML_POWER_SMOOTHING_MAX_NUM_PROFILES = 5 NVML_POWER_SMOOTHING_ADMIN_OVERRIDE_NOT_SET = 0xFFFFFFFF NVML_POWER_SMOOTHING_PROFILE_PARAM_PERCENT_TMP_FLOOR = 0 NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_UP_RATE = 1 NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_RATE = 2 NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_HYSTERESIS = 3 nvmlPowerSmoothingState_v1 = 0x1000008 class c_nvmlPowerSmoothingState_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("state", c_uint), ] def __init__(self): super(c_nvmlPowerSmoothingState_v1_t, self).__init__( version=nvmlPowerSmoothingState_v1 ) nvmlPowerSmoothingProfile_v1 = 0x1000018 class c_nvmlPowerSmoothingProfile_v1_t(_PrintableStructure): _fields_ = [ ("version", c_uint), ("profileId", c_uint), ("paramId", c_uint), ("value", c_double), ] def __init__(self): super(c_nvmlPowerSmoothingProfile_v1_t, self).__init__( version=nvmlPowerSmoothingProfile_v1 ) def nvmlDevicePowerSmoothingActivatePresetProfile(device, profile): fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingActivatePresetProfile") ret = fn(device, profile) _nvmlCheckReturn(ret) def nvmlDevicePowerSmoothingUpdatePresetProfileParam(device, profile): fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingUpdatePresetProfileParam") ret = fn(device, profile) _nvmlCheckReturn(ret) def nvmlDevicePowerSmoothingSetState(device, state): fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingSetState") ret = fn(device, state) _nvmlCheckReturn(ret) ================================================ FILE: python/sglang/multimodal_gen/tools/convert_hf_to_fp8.py ================================================ # copied and adapted from Slime """ Convert HuggingFace safetensors model to FP8 format for efficient inference. Example usage: # convert FLUX.1-dev transformer to FP8 python -m sglang.multimodal_gen.tools.convert_hf_to_fp8 \ --model-dir /path/to/FLUX.1-dev/transformer \ --save-dir /path/to/FLUX.1-dev/transformer-FP8 \ --strategy block \ --block-size 128 128 Options: --model-dir MODEL_DIR path to the directory of the HF safetensors model (e.g., transformer subfolder) --save-dir SAVE_DIR path to the directory to save the converted FP8 model --strategy {block,channel,tensor} quantization strategy (default: block) --block-size [BLOCK_SIZE ...] block size for block quantization, e.g., --block-size 128 128 --max-workers MAX_WORKERS number of worker threads for parallel processing (default: 1) """ import argparse import gc import json import os import shutil import threading from concurrent.futures import ThreadPoolExecutor import safetensors import safetensors.torch import torch import torch.nn.functional as F from tqdm import tqdm FP8_INFO = torch.finfo(torch.float8_e4m3fn) FP8_MAX, FP8_MIN = FP8_INFO.max, FP8_INFO.min def ceildiv(a, b): return -(-a // b) def block_fp8(weight, block_size): # per block quant block_n, block_k = block_size[0], block_size[1] shape_0, shape_1 = weight.shape n_tiles = ceildiv(shape_0, block_n) k_tiles = ceildiv(shape_1, block_k) q_weight = F.pad( weight, (0, k_tiles * block_k - shape_1, 0, n_tiles * block_n - shape_0), mode="constant", value=0.0, ) qweight = q_weight.reshape(n_tiles, block_n, k_tiles, block_k) block_max = torch.max(torch.abs(qweight), dim=1, keepdim=True)[0] block_max = torch.max(block_max, dim=3, keepdim=True)[0] scale = block_max.to(torch.float32) / FP8_MAX qweight = ( (qweight / scale) .clamp(min=FP8_MIN, max=FP8_MAX) .reshape((n_tiles * block_n, k_tiles * block_k)) .to(torch.float8_e4m3fn) ) qweight = qweight[:shape_0, :shape_1].clone().detach() scale = scale.squeeze() return qweight, scale def channel_fp8(weight): channel_max = torch.max(weight.abs(), dim=-1, keepdim=True)[0] scale = channel_max.clamp(min=1e-12).to(torch.float32) / FP8_MAX qweight = (weight / scale).clamp(min=FP8_MIN, max=FP8_MAX) qweight = qweight.to(torch.float8_e4m3fn) return qweight, scale def tensor_fp8(weight): scale = weight.abs().max().clamp(min=1e-12).to(torch.float32) / FP8_MAX qweight = (weight / scale).clamp(min=FP8_MIN, max=FP8_MAX) qweight = qweight.to(torch.float8_e4m3fn) scale = scale.view(1) return qweight, scale def quant_fp8(weight, strategy, block_size=None): if strategy == "tensor": return tensor_fp8(weight) elif strategy == "channel": return channel_fp8(weight) else: return block_fp8(weight, block_size) class ConversionResult: def __init__(self): self.lock = threading.Lock() self.weight_map = {} self.param_count = 0 self.modules_to_not_convert = [] def add_result(self, filename, q_weights, module_names): with self.lock: for k, v in q_weights.items(): self.weight_map[k] = filename self.param_count += v.numel() self.modules_to_not_convert.extend(module_names) def process_file( input_path, output_path, filename, strategy, block_size, result_collector ): if not filename.endswith(".safetensors"): return print(f"Processing {filename}, memory usage: {torch.cuda.memory_allocated()}") weights = {} q_weights = {} with safetensors.safe_open( os.path.join(input_path, filename), framework="pt", device="cuda" ) as f: for k in f.keys(): weights[k] = f.get_tensor(k) modules_to_not_convert = [] for key in weights.keys(): if ( "weight" in key and "layernorm" not in key and "embed" not in key and "router" not in key and "mlp.gate." not in key and "norm" not in key and "lm_head" not in key and "eh_proj" not in key and "net" not in key and "txt_mod" not in key and "img_mod" not in key and "modulation" not in key and "img_in" not in key and "txt_in" not in key and "time_in" not in key and "vector_in" not in key and "adaLN_modulation" not in key and "all_final_layer" not in key and "feed_forward" not in key and "proj_out.weight" != key ): qw, s = quant_fp8(weights[key], strategy, block_size) q_weights[key] = qw if block_size: scale_name = key.replace(".weight", ".weight_scale_inv") else: scale_name = key.replace(".weight", ".weight_scale") q_weights[scale_name] = s else: modules_to_not_convert.append(key.replace(".weight", "")) q_weights[key] = weights[key] safetensors.torch.save_file( q_weights, os.path.join(output_path, filename), metadata={"format": "pt"} ) result_collector.add_result(filename, q_weights, modules_to_not_convert) def convert_fp8(input_path, output_path, strategy, block_size=None, max_workers=4): input_path = os.path.abspath(input_path) os.makedirs(output_path, exist_ok=True) for filename in os.listdir(input_path): if not filename.endswith(".safetensors") and not os.path.isdir( os.path.join(input_path, filename) ): shutil.copyfile( os.path.join(input_path, filename), os.path.join(output_path, filename) ) safetensors_files = [ f for f in os.listdir(input_path) if f.endswith(".safetensors") ] result_collector = ConversionResult() with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [] for filename in safetensors_files: future = executor.submit( process_file, input_path, output_path, filename, strategy, block_size, result_collector, ) futures.append(future) for future in tqdm(futures, desc="Processing files"): future.result() if strategy == "block" or strategy == "tensor": quantization_config = { "activation_scheme": "dynamic", "fmt": "e4m3", "quant_method": "fp8", } if block_size: quantization_config["weight_block_size"] = block_size if len(result_collector.modules_to_not_convert) > 0: quantization_config["modules_to_not_convert"] = list( set(result_collector.modules_to_not_convert) ) else: quant_group = { "group_0": { "input_activations": { "actorder": None, "block_structure": None, "dynamic": True, "group_size": None, "num_bits": 8, "observer": None, "observer_kwargs": {}, "strategy": "token", "symmetric": True, "type": "float", }, "output_activations": None, "targets": ["Linear"], "weights": { "actorder": None, "block_structure": None, "dynamic": False, "group_size": None, "num_bits": 8, "observer": "minmax", "observer_kwargs": {}, "strategy": strategy, "symmetric": True, "type": "float", }, }, } quantization_config = { "config_groups": quant_group, "format": "float-quantized", "ignore": list(set(result_collector.modules_to_not_convert)), "quant_method": "compressed-tensors", "quantization_status": "compressed", } config_path = os.path.join(input_path, "config.json") if os.path.exists(config_path): cfg = json.load(open(config_path)) cfg["quantization_config"] = quantization_config json.dump(cfg, open(os.path.join(output_path, "config.json"), "w"), indent=2) index_dict = { "weight_map": result_collector.weight_map, "metadata": {"total_size": result_collector.param_count}, } json.dump( index_dict, open(os.path.join(output_path, "model.safetensors.index.json"), "w"), indent=2, ) gc.collect() torch.cuda.empty_cache() if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--model-dir", type=str, help="Path to the directory of the HF safetensors model.", ) parser.add_argument( "--save-dir", type=str, help="Path to the directory to save the converted model.", ) parser.add_argument( "--strategy", type=str, default="block", choices=["block", "channel", "tensor"] ) parser.add_argument( "--block-size", type=int, nargs="*", default=None, help="eg. --block-size 32 32" ) parser.add_argument( "--max-workers", type=int, default=8, help="Number of worker threads for parallel processing", ) args = parser.parse_args() if not os.path.exists(args.save_dir): print(f"Creating directory {args.save_dir}") os.makedirs(args.save_dir) elif not os.path.isdir(args.save_dir): raise ValueError("The save_dir should be a directory.") convert_fp8( args.model_dir, args.save_dir, args.strategy, args.block_size, args.max_workers ) ================================================ FILE: python/sglang/multimodal_gen/tools/wan_repack.py ================================================ ### Based on https://github.com/huggingface/diffusers/blob/main/scripts/convert_wan_to_diffusers.py import argparse import json import pathlib from typing import Any, Dict, Tuple from safetensors.torch import load_file, save_file TRANSFORMER_KEYS_RENAME_DICT = { "time_embedding.0": "condition_embedder.time_embedder.linear_1", "time_embedding.2": "condition_embedder.time_embedder.linear_2", "text_embedding.0": "condition_embedder.text_embedder.linear_1", "text_embedding.2": "condition_embedder.text_embedder.linear_2", "time_projection.1": "condition_embedder.time_proj", "head.modulation": "scale_shift_table", "head.head": "proj_out", "modulation": "scale_shift_table", "ffn.0": "ffn.net.0.proj", "ffn.2": "ffn.net.2", # Hack to swap the layer names # The original model calls the norms in following order: norm1, norm3, norm2 # We convert it to: norm1, norm2, norm3 "norm2": "norm__placeholder", "norm3": "norm2", "norm__placeholder": "norm3", # For the I2V model "img_emb.proj.0": "condition_embedder.image_embedder.norm1", "img_emb.proj.1": "condition_embedder.image_embedder.ff.net.0.proj", "img_emb.proj.3": "condition_embedder.image_embedder.ff.net.2", "img_emb.proj.4": "condition_embedder.image_embedder.norm2", # for the FLF2V model "img_emb.emb_pos": "condition_embedder.image_embedder.pos_embed", # Add attention component mappings "self_attn.q": "attn1.to_q", "self_attn.k": "attn1.to_k", "self_attn.v": "attn1.to_v", "self_attn.o": "attn1.to_out.0", "self_attn.norm_q": "attn1.norm_q", "self_attn.norm_k": "attn1.norm_k", "cross_attn.q": "attn2.to_q", "cross_attn.k": "attn2.to_k", "cross_attn.v": "attn2.to_v", "cross_attn.o": "attn2.to_out.0", "cross_attn.norm_q": "attn2.norm_q", "cross_attn.norm_k": "attn2.norm_k", "attn2.to_k_img": "attn2.add_k_proj", "attn2.to_v_img": "attn2.add_v_proj", "attn2.norm_k_img": "attn2.norm_added_k", } def get_transformer_config(model_type: str) -> Tuple[Dict[str, Any], ...]: if model_type == "Wan-T2V-14B": RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT return RENAME_DICT def update_dict_(dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]: dict[new_key] = dict.pop(old_key) def load_sharded_safetensors(path: pathlib.Path): file_path = path state_dict = {} state_dict.update(load_file(file_path)) return state_dict def convert_transformer(model_type: str, model_dir: str, output_dir: str): pathlib.Path(output_dir).mkdir(parents=True, exist_ok=True) RENAME_DICT = get_transformer_config(model_type) original_state_dict = load_sharded_safetensors( pathlib.Path(model_dir, "*model*.safetensors") ) with open(pathlib.Path(model_dir, "*quant_model_description*.json")) as f: original_quant_config = json.load(f) for key in list(original_state_dict.keys()): new_key = key[:] for replace_key, rename_key in RENAME_DICT.items(): new_key = new_key.replace(replace_key, rename_key) update_dict_(original_state_dict, key, new_key) update_dict_(original_quant_config, key, new_key) save_file( original_state_dict, pathlib.Path(output_dir, "diffusion_pytorch_model.safetensors"), ) with open(pathlib.Path(output_dir, "quant_model_description.json"), "w") as f: json.dump(original_quant_config, f) def get_args(): parser = argparse.ArgumentParser() parser.add_argument("--input-path", type=str, required=True) parser.add_argument("--output-path", type=str, required=True) return parser.parse_args() if __name__ == "__main__": args = get_args() convert_transformer( "Wan-T2V-14B", model_dir=pathlib.Path(args.input_path, "high_noise_model"), output_dir=pathlib.Path(args.output_path, "transformer"), ) convert_transformer( "Wan-T2V-14B", model_dir=pathlib.Path(args.input_path, "low_noise_model"), output_dir=pathlib.Path(args.output_path, "transformer_2"), ) ================================================ FILE: python/sglang/multimodal_gen/utils.py ================================================ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py import argparse import ctypes import importlib import importlib.util import inspect import math import os import signal import sys import threading import traceback from collections.abc import Callable from dataclasses import dataclass, fields, is_dataclass from functools import lru_cache, partial, wraps from typing import Any, TypeVar, cast import cloudpickle import torch import yaml from torch.distributed.fsdp import MixedPrecisionPolicy import sglang.multimodal_gen.envs as envs from sglang.multimodal_gen.runtime.utils.logging_utils import ( SortedHelpFormatter, init_logger, ) logger = init_logger(__name__) T = TypeVar("T") def _expand_path_value(field_name: str, value: Any) -> Any: eu = os.path.expanduser if field_name.endswith("_path") and isinstance(value, str): return eu(value) if field_name.endswith("_path") and isinstance(value, list): return [eu(x) if isinstance(x, str) else x for x in value] if field_name.endswith("_paths") and isinstance(value, dict): return {k: eu(p) if isinstance(p, str) else p for k, p in value.items()} return value def expand_path_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: return {key: _expand_path_value(key, value) for key, value in kwargs.items()} def expand_path_fields(obj) -> None: """In-place expanduser on all dataclass fields whose name ends with '_path' or '_paths'.""" for f in fields(obj): setattr(obj, f.name, _expand_path_value(f.name, getattr(obj, f.name))) # TODO(will): used to convert server_args.precision to torch.dtype. Find a # cleaner way to do this. PRECISION_TO_TYPE = { "fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16, } STR_BACKEND_ENV_VAR: str = "SGLANG_DIFFUSION_ATTENTION_BACKEND" STR_ATTN_CONFIG_ENV_VAR: str = "SGLANG_DIFFUSION_ATTENTION_CONFIG" def find_nccl_library() -> str: """ We either use the library file specified by the `VLLM_NCCL_SO_PATH` environment variable, or we find the library file brought by PyTorch. After importing `torch`, `libnccl.so.2`, `librccl.so.1` or `libmccl.so.2` can be found by `ctypes` automatically. """ so_file = envs.SGLANG_DIFFUSION_NCCL_SO_PATH # manually load the nccl library if so_file: logger.info( "Found nccl from environment variable SGLANG_DIFFUSION_NCCL_SO_PATH=%s", so_file, ) else: if torch.version.cuda is not None: so_file = "libnccl.so.2" elif torch.version.hip is not None: so_file = "librccl.so.1" elif hasattr(torch.version, "musa") and torch.version.musa is not None: so_file = "libmccl.so.2" else: raise ValueError("NCCL only supports CUDA, ROCm and MUSA backends.") logger.info("Found nccl from library %s", so_file) return str(so_file) prev_set_stream = torch.cuda.set_stream _current_stream = None def _patched_set_stream(stream: torch.cuda.Stream | None) -> None: global _current_stream _current_stream = stream if stream is not None: prev_set_stream(stream) torch.cuda.set_stream = _patched_set_stream def current_stream() -> torch.cuda.Stream | None: """ replace `torch.cuda.current_stream()` with `sglang.multimodal_gen.utils.current_stream()`. it turns out that `torch.cuda.current_stream()` is quite expensive, as it will construct a new stream object at each call. here we patch `torch.cuda.set_stream` to keep track of the current stream directly, so that we can avoid calling `torch.cuda.current_stream()`. the underlying hypothesis is that we do not call `torch._C._cuda_setStream` from C/C++ code. """ from sglang.multimodal_gen.runtime.platforms import current_platform # For non-CUDA platforms, return None if not current_platform.is_cuda_alike(): return None global _current_stream if _current_stream is None: # when this function is called before any stream is set, # we return the default stream. # On ROCm using the default 0 stream in combination with RCCL # is hurting performance. Therefore creating a dedicated stream # per process _current_stream = ( torch.cuda.Stream() if current_platform.is_rocm() else torch.cuda.current_stream() ) return _current_stream class StoreBoolean(argparse.Action): def __init__(self, option_strings, dest, default=False, required=False, help=None): super().__init__( option_strings=option_strings, dest=dest, nargs="?", const=True, default=default, required=required, help=help, ) def __call__(self, parser, namespace, values, option_string=None): if values is None: setattr(namespace, self.dest, True) elif isinstance(values, str): if values.lower() == "true": setattr(namespace, self.dest, True) elif values.lower() == "false": setattr(namespace, self.dest, False) else: raise ValueError( f"Invalid boolean value: {values}. " "Expected 'true' or 'false'." ) else: setattr(namespace, self.dest, bool(values)) class FlexibleArgumentParser(argparse.ArgumentParser): """ArgumentParser that allows both underscore and dash in names.""" def __init__(self, *args, **kwargs) -> None: # Set the default 'formatter_class' to SortedHelpFormatter if "formatter_class" not in kwargs: kwargs["formatter_class"] = SortedHelpFormatter super().__init__(*args, **kwargs) def parse_args( # type: ignore[override] self, args=None, namespace=None ) -> argparse.Namespace: if args is None: args = sys.argv[1:] if any(arg.startswith("--config") for arg in args): args = self._pull_args_from_config(args) # Convert underscores to dashes and vice versa in argument names processed_args = [] for arg in args: if arg.startswith("--"): if "=" in arg: key, value = arg.split("=", 1) key = "--" + key[len("--") :].replace("_", "-") processed_args.append(f"{key}={value}") else: processed_args.append("--" + arg[len("--") :].replace("_", "-")) elif arg.startswith("-O") and arg != "-O" and len(arg) == 2: # allow -O flag to be used without space, e.g. -O3 processed_args.append("-O") processed_args.append(arg[2:]) else: processed_args.append(arg) namespace = super().parse_args(processed_args, namespace) # Track which arguments were explicitly provided namespace._provided = set() i = 0 while i < len(args): arg = args[i] if arg.startswith("--"): # Handle --key=value format if "=" in arg: key = arg.split("=")[0][2:].replace("-", "_") namespace._provided.add(key) i += 1 # Handle --key value format else: key = arg[2:].replace("-", "_") namespace._provided.add(key) # Skip the value if there is one if i + 1 < len(args) and not args[i + 1].startswith("-"): i += 2 else: i += 1 else: i += 1 return namespace # type: ignore[no-any-return] def _pull_args_from_config(self, args: list[str]) -> list[str]: """Method to pull arguments specified in the config file into the command-line args variable. The arguments in config file will be inserted between the argument list. example: ```yaml port: 12323 tensor-parallel-size: 4 ``` ```python $: vllm {serve,chat,complete} "facebook/opt-12B" \ --config config.yaml -tp 2 $: args = [ "serve,chat,complete", "facebook/opt-12B", '--config', 'config.yaml', '-tp', '2' ] $: args = [ "serve,chat,complete", "facebook/opt-12B", '--port', '12323', '--tp-size', '4', '-tp', '2' ] ``` Please note how the config args are inserted after the sub command. this way the order of priorities is maintained when these are args parsed by super(). """ index = -1 config_arg = None for i, arg in enumerate(args): if arg.startswith("--config"): if index != -1: raise ValueError("More than one config file specified!") index = i config_arg = arg if config_arg is None: return args args_before_config = args[:index] if "=" in config_arg: file_path = config_arg.split("=", 1)[1] args_after_config = args[index + 1 :] else: if index == len(args) - 1: raise ValueError( "No config file specified! " "Please check your command-line arguments." ) file_path = args[index + 1] args_after_config = args[index + 2 :] config_args = self._load_config_file(file_path) # 0th index is for {serve,chat,complete} # followed by model_tag (only for serve) # followed by config args # followed by rest of cli args. # maintaining this order will enforce the precedence # of cli > config > defaults if args[0] == "serve": if index == 1: raise ValueError( "No model_tag specified! Please check your command-line" " arguments." ) command = args_before_config[0] model_tag = args_before_config[1] other_args_before = args_before_config[2:] args = ( [command, model_tag] + config_args + other_args_before + args_after_config ) else: command = args_before_config[0] other_args_before = args_before_config[1:] args = [command] + config_args + other_args_before + args_after_config return args def _load_config_file(self, file_path: str) -> list[str]: """Loads a yaml file and returns the key value pairs as a flattened list with argparse like pattern ```yaml port: 12323 tensor-parallel-size: 4 vae_config: load_encoder: false load_decoder: true ``` returns: processed_args: list[str] = [ '--port': '12323', '--tp-size': '4', '--vae-config.load-encoder': 'false', '--vae-config.load-decoder': 'true' ] """ extension: str = file_path.split(".")[-1] if extension not in ("yaml", "yml", "json"): raise ValueError( "Config file must be of a yaml/yml/json type.\ %s supplied", extension, ) processed_args: list[str] = [] config: dict[str, Any] = {} try: with open(file_path) as config_file: config = yaml.safe_load(config_file) except Exception as ex: logger.error( "Unable to read the config file at %s. \ Make sure path is correct", file_path, ) raise ex store_boolean_arguments = [ action.dest for action in self._actions if isinstance(action, StoreBoolean) ] def process_dict(prefix: str, d: dict[str, Any]): for key, value in d.items(): full_key = f"{prefix}.{key}" if prefix else key if isinstance(value, bool) and full_key not in store_boolean_arguments: if value: processed_args.append("--" + full_key) else: processed_args.append("--" + full_key) processed_args.append("false") elif isinstance(value, list): processed_args.append("--" + full_key) for item in value: processed_args.append(str(item)) elif isinstance(value, dict): process_dict(full_key, value) else: processed_args.append("--" + full_key) processed_args.append(str(value)) process_dict("", config) return processed_args def warn_for_unimplemented_methods(cls: type[T]) -> type[T]: """ A replacement for `abc.ABC`. When we use `abc.ABC`, subclasses will fail to instantiate if they do not implement all abstract methods. Here, we only require `raise NotImplementedError` in the base class, and log a warning if the method is not implemented in the subclass. """ original_init = cls.__init__ def find_unimplemented_methods(self: object): unimplemented_methods = [] for attr_name in dir(self): # bypass inner method if attr_name.startswith("_"): continue try: attr = getattr(self, attr_name) # get the func of callable method if callable(attr): attr_func = attr.__func__ except AttributeError: continue src = inspect.getsource(attr_func) if "NotImplementedError" in src: unimplemented_methods.append(attr_name) if unimplemented_methods: method_names = ",".join(unimplemented_methods) msg = f"Methods {method_names} not implemented in {self}" logger.warning(msg) @wraps(original_init) def wrapped_init(self, *args, **kwargs) -> None: original_init(self, *args, **kwargs) find_unimplemented_methods(self) type.__setattr__(cls, "__init__", wrapped_init) return cls def align_to(value: int, alignment: int) -> int: """align height, width according to alignment Args: value (int): height or width alignment (int): target alignment factor Returns: int: the aligned value """ return int(math.ceil(value / alignment) * alignment) def resolve_obj_by_qualname(qualname: str) -> Any: """ Resolve an object by its fully qualified name. """ module_name, obj_name = qualname.rsplit(".", 1) module = importlib.import_module(module_name) return getattr(module, obj_name) # From vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py def import_pynvml(): """ Historical comments: libnvml.so is the library behind nvidia-smi, and pynvml is a Python wrapper around it. We use it to get GPU status without initializing CUDA context in the current process. Historically, there are two packages that provide pynvml: - `nvidia-ml-py` (https://pypi.org/project/nvidia-ml-py/): The official wrapper. It is a dependency of sglang-diffusion, and is installed when users install sglang-diffusion. It provides a Python module named `pynvml`. - `pynvml` (https://pypi.org/project/pynvml/): An unofficial wrapper. Prior to version 12.0, it also provides a Python module `pynvml`, and therefore conflicts with the official one which is a standalone Python file. This causes errors when both of them are installed. Starting from version 12.0, it migrates to a new module named `pynvml_utils` to avoid the conflict. It is so confusing that many packages in the community use the unofficial one by mistake, and we have to handle this case. For example, `nvcr.io/nvidia/pytorch:24.12-py3` uses the unofficial one, and it will cause errors, see the issue https://github.com/vllm-project/vllm/issues/12847 for example. After all the troubles, we decide to copy the official `pynvml` module to our codebase, and use it directly. """ import sglang.multimodal_gen.third_party.pynvml as pynvml return pynvml def update_environment_variables(envs: dict[str, str]): for k, v in envs.items(): if k in os.environ and os.environ[k] != v: logger.warning( "Overwriting environment variable %s " "from '%s' to '%s'", k, os.environ[k], v, ) os.environ[k] = v def run_method( obj: Any, method: str | bytes | Callable, args: tuple[Any], kwargs: dict[str, Any] ) -> Any: """ Run a method of an object with the given arguments and keyword arguments. If the method is string, it will be converted to a method using getattr. If the method is serialized bytes and will be deserialized using cloudpickle. If the method is a callable, it will be called directly. """ if isinstance(method, bytes): func = partial(cloudpickle.loads(method), obj) elif isinstance(method, str): try: func = getattr(obj, method) except AttributeError: raise NotImplementedError( f"Method {method!r} is not" " implemented." ) from None else: func = partial(method, obj) # type: ignore return func(*args, **kwargs) def shallow_asdict(obj) -> dict[str, Any]: if not is_dataclass(obj): raise TypeError("Expected dataclass instance") return {f.name: getattr(obj, f.name) for f in fields(obj)} # TODO: validate that this is fine def kill_itself_when_parent_died() -> None: # if sys.platform == "linux": # sigkill this process when parent worker manager dies PR_SET_PDEATHSIG = 1 import platform if platform.system() == "Linux": libc = ctypes.CDLL("libc.so.6") libc.prctl(PR_SET_PDEATHSIG, signal.SIGKILL) # elif platform.system() == "Darwin": # libc = ctypes.CDLL("libc.dylib") # logger.warning("kill_itself_when_parent_died is only supported in linux.") else: logger.warning("kill_itself_when_parent_died is only supported in linux.") def get_exception_traceback() -> str: etype, value, tb = sys.exc_info() err_str = "".join(traceback.format_exception(etype, value, tb)) return err_str class TypeBasedDispatcher: def __init__(self, mapping: list[tuple[type, Callable]]): self._mapping = mapping def __call__(self, obj: Any): for ty, fn in self._mapping: if isinstance(obj, ty): return fn(obj) raise ValueError(f"Invalid object: {obj}") @dataclass class MixedPrecisionState: param_dtype: torch.dtype | None = None reduce_dtype: torch.dtype | None = None output_dtype: torch.dtype | None = None compute_dtype: torch.dtype | None = None mp_policy: MixedPrecisionPolicy | None = None # Thread-local storage for mixed precision state _mixed_precision_state = threading.local() def get_mixed_precision_state() -> MixedPrecisionState: """Get the current mixed precision state.""" if not hasattr(_mixed_precision_state, "state"): raise ValueError("Mixed precision state not set") return cast(MixedPrecisionState, _mixed_precision_state.state) def set_mixed_precision_policy( param_dtype: torch.dtype, reduce_dtype: torch.dtype, output_dtype: torch.dtype | None = None, mp_policy: MixedPrecisionPolicy | None = None, ): """Set mixed precision policy globally. Args: param_dtype: Parameter dtype used for training reduce_dtype: Reduction dtype used for gradients output_dtype: Optional output dtype """ state = MixedPrecisionState( param_dtype=param_dtype, reduce_dtype=reduce_dtype, output_dtype=output_dtype, mp_policy=mp_policy, ) _mixed_precision_state.state = state def get_compute_dtype() -> torch.dtype: """Get the current compute dtype from mixed precision policy.""" if not hasattr(_mixed_precision_state, "state"): return torch.get_default_dtype() else: state = get_mixed_precision_state() return state.param_dtype def dict_to_3d_list( mask_strategy: dict[str, Any] | None = None, t_max: int | None = None, l_max: int | None = None, h_max: int | None = None, ) -> list[list[list[torch.Tensor | None]]]: """ Convert a dictionary of mask indices to a 3D list of tensors. Args: mask_strategy: keys are "t_l_h", values are torch.Tensor masks. t_max, l_max, h_max: if provided (all three), force the output shape to (t_max, l_max, h_max). If all three are None, infer shape from the data. """ # Case 1: no data, but fixed shape requested if mask_strategy is None: assert ( t_max is not None and l_max is not None and h_max is not None ), "If mask_strategy is None, you must provide t_max, l_max, and h_max" return [ [[None for _ in range(h_max)] for _ in range(l_max)] for _ in range(t_max) ] # Parse all keys into integer tuples indices = [tuple(map(int, key.split("_"))) for key in mask_strategy] # Decide on dimensions if t_max is None and l_max is None and h_max is None: # fully dynamic: infer from data max_timesteps_idx = max(t for t, _, _ in indices) + 1 max_layer_idx = max(l for _, l, _ in indices) + 1 # noqa: E741 max_head_idx = max(h for _, _, h in indices) + 1 else: # require all three to be provided assert t_max is not None and l_max is not None and h_max is not None, ( "Either supply none of (t_max, l_max, h_max) to infer dimensions, " "or supply all three to fix the shape." ) max_timesteps_idx = t_max max_layer_idx = l_max max_head_idx = h_max # Preallocate result = [ [[None for _ in range(max_head_idx)] for _ in range(max_layer_idx)] for _ in range(max_timesteps_idx) ] # Fill in, skipping any out-of-bounds entries for key, value in mask_strategy.items(): t, l, h = map(int, key.split("_")) # noqa: E741 if ( 0 <= t < max_timesteps_idx and 0 <= l < max_layer_idx and 0 <= h < max_head_idx ): result[t][l][h] = value # else: silently ignore any key that doesn't fit return result def set_random_seed(seed: int) -> None: from sglang.multimodal_gen.runtime.platforms import current_platform current_platform.seed_everything(seed) @lru_cache(maxsize=1) def is_vsa_available() -> bool: return importlib.util.find_spec("vsa") is not None @lru_cache(maxsize=1) def is_vmoba_available() -> bool: if importlib.util.find_spec("kernel.csrc.attn.vmoba_attn.vmoba") is None: return False try: import flash_attn return flash_attn.__version__ >= "2.7.4" except Exception: return False # adapted from: https://github.com/Wan-Video/Wan2.2/blob/main/wan/utils/utils.py def masks_like( tensors, zero=False, generator=None, p=0.2 ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: """ Generate binary masks for Text-to-Image-to-Video (TI2V) tasks. Creates masks to control which frames should be preserved vs replaced. Primarily used to fix the first frame to the input image while generating other frames. Args: tensors: List of tensors with shape [C, T, H, W] zero: If True, set first frame (dim 1, index 0) to zero. Default: False generator: Optional random generator for stochastic masking p: Probability of applying special noise when generator is provided. Default: 0.2 Returns: Tuple of two lists of tensors: - When zero=False: Both lists contain all-ones tensors - When zero=True (no generator): First frame set to 0, others to 1 - When zero=True (with generator): First frame set to small random values with probability p Example: >>> latent = torch.randn(48, 69, 96, 160) # [C, T, H, W] >>> _, mask = masks_like([latent], zero=True) >>> # mask[0][:, 0] == 0 (first frame) >>> # mask[0][:, 1:] == 1 (other frames) >>> blended = (1.0 - mask[0]) * image + mask[0] * latent >>> # Result: first frame = image, other frames = latent """ assert isinstance(tensors, list) out1 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensors] out2 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensors] if zero: if generator is not None: for u, v in zip(out1, out2, strict=False): random_num = torch.rand( 1, generator=generator, device=generator.device ).item() if random_num < p: u[:, 0] = ( torch.normal( mean=-3.5, std=0.5, size=(1,), device=u.device, generator=generator, ) .expand_as(u[:, 0]) .exp() ) v[:, 0] = torch.zeros_like(v[:, 0]) else: u[:, 0] = u[:, 0] v[:, 0] = v[:, 0] else: for u, v in zip(out1, out2, strict=False): u[:, 0] = torch.zeros_like(u[:, 0]) v[:, 0] = torch.zeros_like(v[:, 0]) return out1, out2 # adapted from: https://github.com/Wan-Video/Wan2.2/blob/main/wan/utils/utils.py def best_output_size(w, h, dw, dh, expected_area): # float output size ratio = w / h ow = (expected_area * ratio) ** 0.5 oh = expected_area / ow # process width first ow1 = int(ow // dw * dw) oh1 = int(expected_area / ow1 // dh * dh) assert ow1 % dw == 0 and oh1 % dh == 0 and ow1 * oh1 <= expected_area ratio1 = ow1 / oh1 # process height first oh2 = int(oh // dh * dh) ow2 = int(expected_area / oh2 // dw * dw) assert oh2 % dh == 0 and ow2 % dw == 0 and ow2 * oh2 <= expected_area ratio2 = ow2 / oh2 # compare ratios if max(ratio / ratio1, ratio1 / ratio) < max(ratio / ratio2, ratio2 / ratio): return ow1, oh1 else: return ow2, oh2 def calculate_dimensions(target_area, ratio): width = math.sqrt(target_area * ratio) height = width / ratio width = round(width / 32) * 32 height = round(height / 32) * 32 return width, height, None ================================================ FILE: python/sglang/profiler.py ================================================ """ Run live profiling. Usage: python3 -m sglang.profiler """ import argparse import json import os import time from argparse import ArgumentParser from pathlib import Path from typing import List, Optional import requests PROFILER_DIR = os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp") def run_profile( url: Optional[str], num_steps: int, activities: List[str], output_dir: Optional[str] = None, profile_by_stage: bool = False, merge_profiles: bool = False, profile_prefix: Optional[str] = None, start_step: Optional[int] = None, ) -> str: if output_dir is None: output_dir = PROFILER_DIR output_dir = Path(os.path.abspath(os.path.normpath(output_dir))) / str(time.time()) output_dir.mkdir(exist_ok=True, parents=True) print(f"Dump profiling traces to {output_dir}") print( f"Waiting for {num_steps} steps and the trace to be flushed.... ({profile_by_stage=})" ) # Dump server args. file_path = Path(output_dir) / "server_args.json" if not file_path.exists(): response = requests.get(url + "/get_server_info") response.raise_for_status() server_args_data = response.json() with open(file_path, "w") as file: file.write(json.dumps(server_args_data)) # Start profiler. The API replies when all steps are processed # and files are generated. json_data = { "output_dir": str(output_dir), "num_steps": str(num_steps), "activities": activities, "profile_by_stage": profile_by_stage, "merge_profiles": merge_profiles, "profile_prefix": profile_prefix, } if start_step is not None: json_data["start_step"] = str(start_step) response = requests.post(url=url + "/start_profile", json=json_data) response.raise_for_status() trace_link = str(output_dir) return trace_link if __name__ == "__main__": parser = ArgumentParser(description="Benchmark the online serving throughput.") parser.add_argument( "--url", type=str, default="http://localhost:30000", help="Server or API base url if not using http host and port.", ) parser.add_argument( "--output-dir", type=str, default=None, help="Profile directory to dump profile traces.", ) parser.add_argument( "--num-steps", type=int, default=5, help="The number of forward steps to profile.", ) parser.add_argument( "--profile-by-stage", action=argparse.BooleanOptionalAction, type=bool, default=False, help="Whether to profile prefill and decode separately", ) parser.add_argument( "--profile-prefix", type=str, help="The prefix of this profiler file.", ) parser.add_argument( "--cpu", action=argparse.BooleanOptionalAction, type=bool, default=True, help="Whether to profile CPU activity", ) parser.add_argument( "--gpu", action=argparse.BooleanOptionalAction, type=bool, default=True, help="Whether to profile GPU activity", ) parser.add_argument( "--mem", action=argparse.BooleanOptionalAction, type=bool, default=False, help="Whether to profile memory usage (https://pytorch.org/memory_viz)", ) parser.add_argument( "--rpd", action=argparse.BooleanOptionalAction, type=bool, default=False, help="Whether to use ROCM rpd profiler (https://github.com/ROCm/rocmProfileData)", ) parser.add_argument( "--merge-profiles", action=argparse.BooleanOptionalAction, type=bool, default=False, help="Whether to merge profiles from all ranks into a single trace file", ) args = parser.parse_args() activities = [] if args.cpu: activities.append("CPU") if args.gpu: activities.append("GPU") if args.mem: activities.append("MEM") if args.rpd: activities.append("RPD") run_profile( url=args.url, num_steps=args.num_steps, activities=activities, output_dir=args.output_dir, profile_by_stage=args.profile_by_stage, profile_prefix=args.profile_prefix, merge_profiles=args.merge_profiles, ) ================================================ FILE: python/sglang/srt/batch_invariant_ops/__init__.py ================================================ # Adapted from https://github.com/thinking-machines-lab/batch_invariant_ops/blob/main/batch_invariant_ops/__init__.py from .batch_invariant_ops import ( AttentionBlockSize, disable_batch_invariant_mode, enable_batch_invariant_mode, get_batch_invariant_attention_block_size, is_batch_invariant_mode_enabled, log_softmax, matmul_persistent, mean_dim, rms_norm_batch_invariant, set_batch_invariant_mode, ) __version__ = "0.1.0" __all__ = [ "set_batch_invariant_mode", "is_batch_invariant_mode_enabled", "disable_batch_invariant_mode", "enable_batch_invariant_mode", "matmul_persistent", "log_softmax", "mean_dim", "get_batch_invariant_attention_block_size", "AttentionBlockSize", "rms_norm_batch_invariant", ] ================================================ FILE: python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py ================================================ # Adapted from https://github.com/thinking-machines-lab/batch_invariant_ops/blob/main/batch_invariant_ops/batch_invariant_ops.py import contextlib from collections import namedtuple from collections.abc import Callable from typing import Any, Dict import torch import triton import triton.language as tl from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM from sglang.srt.utils.common import calc_diff, get_bool_env_var if ENABLE_JIT_DEEPGEMM: import deep_gemm _ENABLE_MM_DEEPGEMM = get_bool_env_var( "SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM", "1" ) # If true, allows to fallback to batch variant gemm when the shape cannot be run in DeepGEMM _ENABLE_MM_FALLBACK_VARIANT = get_bool_env_var( "SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_FALLBACK_VARIANT", "0" ) _ENABLE_MM_COMPARISON_TEST = get_bool_env_var( "SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_COMPARISON_TEST" ) if not _ENABLE_MM_DEEPGEMM: print("Disable DeepGEMM in batch invariant ops. Performance may be suboptimal.") __all__ = [ "set_batch_invariant_mode", "is_batch_invariant_mode_enabled", "disable_batch_invariant_mode", "enable_batch_invariant_mode", ] def _matmul_launch_metadata( grid: Callable[..., Any], kernel: Any, args: Dict[str, Any] ) -> Dict[str, Any]: ret = {} m, n, k = args["M"], args["N"], args["K"] ret["name"] = f"{kernel.name} [M={m}, N={n}, K={k}]" if "tiles_per_update" in args: ret["name"] = ( f"{kernel.name} [M={m}, N={n}, K={k}, tiles_per_update={args['tiles_per_update']:02}]" ) if "c_ptr" in args: bytes_per_elem = args["c_ptr"].element_size() else: bytes_per_elem = 1 if args["FP8_OUTPUT"] else 2 ret[f"flops{bytes_per_elem * 8}"] = 2.0 * m * n * k ret["bytes"] = bytes_per_elem * (m * k + n * k + m * n) return ret @triton.jit def _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS): group_id = tile_id // num_pid_in_group first_pid_m = group_id * GROUP_SIZE_M group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) pid_m = first_pid_m + (tile_id % group_size_m) pid_n = (tile_id % num_pid_in_group) // group_size_m return pid_m, pid_n @triton.jit(launch_metadata=_matmul_launch_metadata) def matmul_kernel_persistent( a_ptr, b_ptr, c_ptr, # bias_ptr, M, N, K, # stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_SIZE_M: tl.constexpr, # BLOCK_SIZE_N: tl.constexpr, # BLOCK_SIZE_K: tl.constexpr, # GROUP_SIZE_M: tl.constexpr, # NUM_SMS: tl.constexpr, # A_LARGE: tl.constexpr, B_LARGE: tl.constexpr, C_LARGE: tl.constexpr, HAS_BIAS: tl.constexpr, ): start_pid = tl.program_id(axis=0) num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) k_tiles = tl.cdiv(K, BLOCK_SIZE_K) num_tiles = num_pid_m * num_pid_n offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) num_pid_in_group = GROUP_SIZE_M * num_pid_n for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True): pid_m, pid_n = _compute_pid( tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS ) start_m = pid_m * BLOCK_SIZE_M start_n = pid_n * BLOCK_SIZE_N offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) if A_LARGE: offs_am = offs_am.to(tl.int64) if B_LARGE: offs_bn = offs_bn.to(tl.int64) offs_am = tl.where(offs_am < M, offs_am, 0) offs_bn = tl.where(offs_bn < N, offs_bn, 0) offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for ki in range(k_tiles): if A_LARGE or B_LARGE: offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64) else: offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) a_ptrs = a_ptr + ( offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak ) b_ptrs = b_ptr + ( offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn ) a = tl.load( a_ptrs, mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K, other=0.0 ) b = tl.load( b_ptrs, mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K, other=0.0 ) accumulator = tl.dot(a, b, accumulator) offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) if C_LARGE: offs_cm = offs_cm.to(tl.int64) offs_cn = offs_cn.to(tl.int64) c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) if HAS_BIAS: bias_ptrs = bias_ptr + offs_cn bias = tl.load(bias_ptrs, mask=offs_cn < N, other=0.0).to(tl.float32) accumulator += bias if c_ptr.dtype.element_ty == tl.float8e4nv: c = accumulator.to(tl.float8e4nv) elif c_ptr.dtype.element_ty == tl.bfloat16: c = accumulator.to(tl.bfloat16) elif c_ptr.dtype.element_ty == tl.float32: c = accumulator.to(tl.float32) else: c = accumulator.to(tl.float16) tl.store(c_ptrs, c, mask=c_mask) def _matmul_persistent_triton( a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None ): # Check constraints. assert a.shape[1] == b.shape[0], "Incompatible dimensions" assert a.dtype == b.dtype, "Incompatible dtypes" assert ( bias is None or bias.dim() == 1 ), "Currently assuming bias is 1D, let Horace know if you run into this" NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count M, K = a.shape K, N = b.shape dtype = a.dtype # Allocates output. c = torch.empty((M, N), device=a.device, dtype=dtype) # 1D launch kernel where each block gets its own program. def grid(META): return ( min( NUM_SMS, triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), ), ) configs = { torch.bfloat16: { "BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8, }, torch.float16: { "BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8, }, torch.float32: { "BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8, }, } # print(a.device, b.device, c.device) matmul_kernel_persistent[grid]( a, b, c, # bias, M, N, K, # a.stride(0), a.stride(1), # b.stride(0), b.stride(1), # c.stride(0), c.stride(1), # NUM_SMS=NUM_SMS, # A_LARGE=a.numel() > 2**31, B_LARGE=b.numel() > 2**31, C_LARGE=c.numel() > 2**31, HAS_BIAS=bias is not None, **configs[dtype], ) return c def _matmul_persistent_deepgemm( a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None ): M, K = a.shape K, N = b.shape dtype = a.dtype out = torch.empty((M, N), device=a.device, dtype=dtype) try: deep_gemm.bf16_gemm_nn(a, b, out) except RuntimeError as e: raise RuntimeError( f"DeepGEMM failed for matrix shapes M={M}, N={N}, K={K}. " f"This typically occurs when dimensions are too small for DeepGEMM's TMA descriptors. " f"Consider increasing MIN_DEEPGEMM_DIM in matmul_persistent() or disabling DeepGEMM " f"for small matrices. Original error: {e}" ) from e # TODO can this be put in DeepGEMM's `c`? if bias is not None: out += bias return out def matmul_persistent( a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None ): K, N = b.shape # DeepGEMM has minimum dimension requirements for TMA descriptors MIN_DEEPGEMM_DIM = 16 if ( _ENABLE_MM_DEEPGEMM and ENABLE_JIT_DEEPGEMM and (a.dtype == torch.bfloat16) and (b.dtype == torch.bfloat16) and a.is_contiguous() and b.transpose(0, 1).is_contiguous() and N >= MIN_DEEPGEMM_DIM ): if _ENABLE_MM_COMPARISON_TEST: out_triton = _matmul_persistent_triton(a=a, b=b, bias=bias) out_deepgemm = _matmul_persistent_deepgemm(a=a, b=b, bias=bias) diff = calc_diff(out_triton, out_deepgemm) assert diff < 0.0001, f"{diff=} {out_triton=} {out_deepgemm=}" # can be enabled for debugging # print( # f"{diff=} " # f"{(out_triton - out_deepgemm).abs().mean()=} " # f"{(out_triton - out_deepgemm).abs().sum()=} " # f"{torch.sum(out_triton != out_deepgemm)=} " # ) # print(f"{a=} {b=} {bias=} {out_triton=} {out_deepgemm=}") return out_deepgemm return _matmul_persistent_deepgemm(a=a, b=b, bias=bias) if _ENABLE_MM_FALLBACK_VARIANT: out = torch.einsum("ik,kj->ij", a, b) if bias is not None: out += bias return out return _matmul_persistent_triton(a=a, b=b, bias=bias) @triton.jit def _log_softmax_kernel( input_ptr, output_ptr, input_row_stride: tl.constexpr, output_row_stride: tl.constexpr, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): """ Compute log_softmax along the last dimension of a 2D tensor. Each block handles one row of the input tensor. """ # Get the row index for this block row_idx = tl.program_id(0).to(tl.int64) # Compute base pointers for input and output rows row_start_ptr = input_ptr + row_idx * input_row_stride output_row_start_ptr = output_ptr + row_idx * output_row_stride # Step 1: Find maximum value in the row for numerical stability # Load first block to infer dtype and initialize max_val with correct type col_idx_init = tl.arange(0, BLOCK_SIZE) mask_init = col_idx_init < n_cols vals_init = tl.load( row_start_ptr + col_idx_init, mask=mask_init, other=-float("inf") ) max_val = tl.max(vals_init) # Continue with remaining blocks for col_offset in range(BLOCK_SIZE, n_cols, BLOCK_SIZE): col_idx = col_offset + tl.arange(0, BLOCK_SIZE) mask = col_idx < n_cols # Load values vals = tl.load(row_start_ptr + col_idx, mask=mask, other=-float("inf")) # Update maximum max_val = tl.max(tl.maximum(vals, max_val)) # Step 2: Compute sum of exp(x - max_val) # Initialize sum_exp with correct dtype by using tl.sum on a zero vector sum_exp = tl.sum(tl.zeros([1], dtype=max_val.dtype)) for col_offset in range(0, n_cols, BLOCK_SIZE): col_idx = col_offset + tl.arange(0, BLOCK_SIZE) mask = col_idx < n_cols # Load values vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) # Compute exp(x - max_val) and accumulate exp_vals = tl.exp(vals - max_val) sum_exp += tl.sum(tl.where(mask, exp_vals, 0.0)) # Compute log(sum_exp) log_sum_exp = tl.log(sum_exp) # Step 3: Compute final log_softmax values: x - max_val - log_sum_exp for col_offset in range(0, n_cols, BLOCK_SIZE): col_idx = col_offset + tl.arange(0, BLOCK_SIZE) mask = col_idx < n_cols # Load values vals = tl.load(row_start_ptr + col_idx, mask=mask) # Compute log_softmax output = vals - max_val - log_sum_exp # Store results tl.store(output_row_start_ptr + col_idx, output, mask=mask) def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: """ Compute log_softmax using Triton kernel. Args: input: Input tensor dim: Dimension along which to compute log_softmax (only -1 or last dim supported) >> Stashed changes Returns: Tensor with log_softmax applied along the specified dimension """ if dim != -1 and dim != input.ndim - 1: raise ValueError( "This implementation only supports log_softmax along the last dimension" ) # Flatten all dimensions except the last one original_shape = input.shape input_2d = input.reshape(-1, input.shape[-1]) input_2d = input_2d.contiguous() n_rows, n_cols = input_2d.shape # Allocate output tensor output = torch.empty_like(input_2d) # Choose block size based on the number of columns BLOCK_SIZE = 1024 # Launch kernel with one block per row grid = (n_rows,) _log_softmax_kernel[grid]( input_2d, output, input_2d.stride(0), output.stride(0), n_cols, BLOCK_SIZE=BLOCK_SIZE, ) # Reshape output back to original shape return output.reshape(original_shape) @triton.jit def mean_kernel( input_ptr, output_ptr, input_stride0, input_stride1, input_stride2, output_stride0, output_stride1, M, # size before reduction dim N, # size of reduction dim K, # size after reduction dim BLOCK_SIZE: tl.constexpr, ): """ Kernel for computing mean along a single dimension. Input is viewed as (M, N, K) where N is the dimension being reduced. """ # Program ID gives us which output element we're computing pid = tl.program_id(0) # Compute output indices m_idx = pid // K k_idx = pid % K # Bounds check if m_idx >= M or k_idx >= K: return # Accumulate sum across reduction dimension acc = 0.0 for n_start in range(0, N, BLOCK_SIZE): n_offsets = n_start + tl.arange(0, BLOCK_SIZE) mask = n_offsets < N # Calculate input indices input_idx = ( m_idx * input_stride0 + n_offsets * input_stride1 + k_idx * input_stride2 ) # Load and accumulate vals = tl.load(input_ptr + input_idx, mask=mask, other=0.0) acc += tl.sum(vals) # Compute mean and store mean_val = acc / N output_idx = m_idx * output_stride0 + k_idx * output_stride1 tl.store(output_ptr + output_idx, mean_val) def mean_dim( input: torch.Tensor, dim: int, keepdim: bool = False, dtype: torch.dtype | None = None, ) -> torch.Tensor: """ Triton implementation of torch.mean with single dimension reduction. Args: input: Input tensor dim: Single dimension along which to compute mean keepdim: Whether to keep the reduced dimension dtype: Output dtype. If None, uses input dtype (or float32 for integer inputs) Returns: Tensor with mean values along specified dimension """ # Validate inputs assert input.is_cuda, "Input must be a CUDA tensor" assert ( -input.ndim <= dim < input.ndim ), f"Invalid dimension {dim} for tensor with {input.ndim} dimensions" # Handle negative dim if dim < 0: dim = dim + input.ndim # Handle dtype if dtype is None: if input.dtype in [torch.int8, torch.int16, torch.int32, torch.int64]: dtype = torch.float32 else: dtype = input.dtype # Convert input to appropriate dtype if needed if input.dtype != dtype: input = input.to(dtype) # Get input shape and strides shape = list(input.shape) # Calculate dimensions for kernel M = 1 for i in range(dim): M *= shape[i] N = shape[dim] K = 1 for i in range(dim + 1, len(shape)): K *= shape[i] # Reshape input to 3D view (M, N, K) input_3d = input.reshape(M, N, K) # Create output shape if keepdim: output_shape = shape.copy() output_shape[dim] = 1 else: output_shape = shape[:dim] + shape[dim + 1 :] # Create output tensor output = torch.empty(output_shape, dtype=dtype, device=input.device) # Reshape output for kernel if keepdim: output_2d = output.reshape(M, 1, K).squeeze(1) else: output_2d = output.reshape(M, K) # Launch kernel grid = (M * K,) BLOCK_SIZE = 1024 mean_kernel[grid]( input_3d, output_2d, input_3d.stride(0), input_3d.stride(1), input_3d.stride(2), output_2d.stride(0), output_2d.stride(1) if output_2d.ndim > 1 else 0, M, N, K, BLOCK_SIZE, ) return output def mm_batch_invariant(a, b): return matmul_persistent(a, b) def addmm_batch_invariant(bias, a, b): return matmul_persistent(a, b, bias=bias) def _log_softmax_batch_invariant(input, dim, _half_to_float): assert not _half_to_float, "not implemented" return log_softmax(input, dim=dim) def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None = None): assert dtype is None or dtype == torch.float32, f"unsupported dtype: {dtype}" if len(dim) == 1: return mean_dim(input, dim[0], keepdim=keepdim) else: assert input.dtype in { torch.float16, torch.bfloat16, torch.float32, }, "only float types supported for now" n_elems = 1 for d in dim: n_elems *= input.shape[d] return torch.sum(input, dim=dim, keepdim=keepdim, dtype=torch.float32) / n_elems @triton.jit def bmm_kernel_persistent( a_ptr, b_ptr, c_ptr, # B, M, N, K, # stride_ab, stride_am, stride_ak, stride_bb, stride_bk, stride_bn, stride_cb, stride_cm, stride_cn, BLOCK_SIZE_M: tl.constexpr, # BLOCK_SIZE_N: tl.constexpr, # BLOCK_SIZE_K: tl.constexpr, # GROUP_SIZE_M: tl.constexpr, # NUM_SMS: tl.constexpr, # A_LARGE: tl.constexpr, B_LARGE: tl.constexpr, C_LARGE: tl.constexpr, ): """ Batched matrix multiplication kernel that processes batches in parallel. Each tile processes a (BLOCK_SIZE_M, BLOCK_SIZE_N) output block for a specific batch. """ start_pid = tl.program_id(axis=0) num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) k_tiles = tl.cdiv(K, BLOCK_SIZE_K) num_tiles_per_batch = num_pid_m * num_pid_n num_tiles_total = B * num_tiles_per_batch offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) num_pid_in_group = GROUP_SIZE_M * num_pid_n # Process tiles in a deterministic order: batch-major ordering for tile_id in tl.range(start_pid, num_tiles_total, NUM_SMS, flatten=True): # Decompose tile_id into batch and within-batch tile batch_idx = tile_id // num_tiles_per_batch tile_in_batch = tile_id % num_tiles_per_batch pid_m, pid_n = _compute_pid( tile_in_batch, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS ) start_m = pid_m * BLOCK_SIZE_M start_n = pid_n * BLOCK_SIZE_N offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) if A_LARGE: offs_am = offs_am.to(tl.int64) if B_LARGE: offs_bn = offs_bn.to(tl.int64) offs_am = tl.where(offs_am < M, offs_am, 0) offs_bn = tl.where(offs_bn < N, offs_bn, 0) offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) # Add batch offset if A_LARGE or B_LARGE: batch_idx_typed = batch_idx.to(tl.int64) else: batch_idx_typed = batch_idx accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for ki in range(k_tiles): if A_LARGE or B_LARGE: offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64) else: offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) a_ptrs = a_ptr + ( batch_idx_typed * stride_ab + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak ) b_ptrs = b_ptr + ( batch_idx_typed * stride_bb + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn ) a = tl.load( a_ptrs, mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K, other=0.0 ) b = tl.load( b_ptrs, mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K, other=0.0 ) accumulator = tl.dot(a, b, accumulator) offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) if C_LARGE: offs_cm = offs_cm.to(tl.int64) offs_cn = offs_cn.to(tl.int64) c_ptrs = ( c_ptr + batch_idx_typed * stride_cb + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] ) c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) if c_ptr.dtype.element_ty == tl.float8e4nv: c = accumulator.to(tl.float8e4nv) elif c_ptr.dtype.element_ty == tl.bfloat16: c = accumulator.to(tl.bfloat16) elif c_ptr.dtype.element_ty == tl.float32: c = accumulator.to(tl.float32) else: c = accumulator.to(tl.float16) tl.store(c_ptrs, c, mask=c_mask) def bmm_batch_invariant(a, b, *, out=None): # Batched matrix multiply: (B, M, K) x (B, K, N) -> (B, M, N) # Process batches in parallel with our persistent kernel if a.ndim == 3 and b.ndim == 3: # Check constraints assert a.shape[0] == b.shape[0], "Batch sizes must match" assert a.shape[2] == b.shape[1], "Incompatible dimensions" assert a.dtype == b.dtype, "Incompatible dtypes" B = a.shape[0] M = a.shape[1] K = a.shape[2] N = b.shape[2] dtype = a.dtype # Allocate output if out is None: c = torch.empty((B, M, N), device=a.device, dtype=dtype) else: c = out NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count # Use fixed kernel configuration for determinism configs = { torch.bfloat16: { "BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8, }, torch.float16: { "BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8, }, torch.float32: { "BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8, }, } config = configs.get(dtype) if config is None: raise ValueError( f"Unsupported dtype {dtype} for bmm_batch_invariant. " f"Supported dtypes are: {list(configs.keys())}" ) # Grid: limit by NUM_SMS for persistent kernel approach num_tiles_per_batch = triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv( N, config["BLOCK_SIZE_N"] ) num_tiles_total = B * num_tiles_per_batch grid = (min(NUM_SMS, num_tiles_total),) bmm_kernel_persistent[grid]( a, b, c, # B, M, N, K, # a.stride(0), a.stride(1), a.stride(2), # b.stride(0), b.stride(1), b.stride(2), # c.stride(0), c.stride(1), c.stride(2), # NUM_SMS=NUM_SMS, # A_LARGE=a.numel() > 2**31, B_LARGE=b.numel() > 2**31, C_LARGE=c.numel() > 2**31, **config, ) return c else: raise ValueError( f"bmm_batch_invariant expects 3D tensors, " f"got shapes {a.shape} and {b.shape}" ) @triton.jit def _rms_norm_kernel( input_ptr, weight_ptr, output_ptr, input_row_stride: tl.constexpr, output_row_stride: tl.constexpr, n_cols: tl.constexpr, eps, BLOCK_SIZE: tl.constexpr, ): """ Compute RMS normalization along the last dimension of a 2D tensor. RMS Norm: y = x / sqrt(mean(x^2) + eps) * weight Each block handles one row of the input tensor. """ row_idx = tl.program_id(0).to(tl.int64) row_start_ptr = input_ptr + row_idx * input_row_stride output_row_start_ptr = output_ptr + row_idx * output_row_stride # Step 1: Compute sum of squares in float32 to avoid overflow sum_sq = tl.zeros([1], dtype=tl.float32) for col_offset in range(0, n_cols, BLOCK_SIZE): col_idx = col_offset + tl.arange(0, BLOCK_SIZE) mask = col_idx < n_cols vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) # Convert to float32 for accumulation to prevent overflow vals_f32 = vals.to(tl.float32) sq_vals = vals_f32 * vals_f32 sum_sq += tl.sum(tl.where(mask, sq_vals, 0.0)) # Step 2: Compute RMS (root mean square) in float32 mean_sq = sum_sq / n_cols rms = tl.sqrt(mean_sq + eps) inv_rms = 1.0 / rms # Step 3: Normalize and apply weight for col_offset in range(0, n_cols, BLOCK_SIZE): col_idx = col_offset + tl.arange(0, BLOCK_SIZE) mask = col_idx < n_cols vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0) # Compute in float32 then convert back to input dtype vals_f32 = vals.to(tl.float32) weight_f32 = weight.to(tl.float32) output_f32 = vals_f32 * inv_rms * weight_f32 output = output_f32.to(vals.dtype) tl.store(output_row_start_ptr + col_idx, output, mask=mask) def rms_norm( input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6 ) -> torch.Tensor: """ Compute RMS normalization using Triton kernel. RMS Norm normalizes the input by the root mean square and scales by weight: output = input / sqrt(mean(input^2) + eps) * weight Args: input: Input tensor of shape (..., hidden_size) weight: Weight tensor of shape (hidden_size,) eps: Small constant for numerical stability Returns: Tensor with RMS normalization applied along the last dimension """ assert weight.dim() == 1, "Weight must be 1-dimensional" assert input.shape[-1] == weight.shape[0], ( f"Input last dimension ({input.shape[-1]}) must match " f"weight dimension ({weight.shape[0]})" ) # Flatten all dimensions except the last one original_shape = input.shape input_2d = input.reshape(-1, input.shape[-1]) input_2d = input_2d.contiguous() weight = weight.contiguous() n_rows, n_cols = input_2d.shape output = torch.empty_like(input_2d) BLOCK_SIZE = 1024 grid = (n_rows,) _rms_norm_kernel[grid]( input_2d, weight, output, input_2d.stride(0), output.stride(0), n_cols, eps, BLOCK_SIZE=BLOCK_SIZE, ) return output.reshape(original_shape) def rms_norm_batch_invariant( input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6 ) -> torch.Tensor: """ Batch-invariant wrapper for RMS normalization. This function provides a deterministic, batch-invariant implementation of RMS normalization for use with the batch_invariant mode. Adapted from @https://github.com/vllm-project/vllm/blob/66a168a197ba214a5b70a74fa2e713c9eeb3251a/vllm/model_executor/layers/batch_invariant.py#L649 Args: input: Input tensor of shape (..., hidden_size) weight: Weight tensor of shape (hidden_size,) eps: Small constant for numerical stability Returns: RMS normalized tensor """ return rms_norm(input, weight, eps=eps) _batch_invariant_MODE = False _batch_invariant_LIB = None _original_torch_bmm = None def is_batch_invariant_mode_enabled(): return _batch_invariant_MODE def enable_batch_invariant_mode( enable_bmm: bool = True, ): global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm if _batch_invariant_MODE: return _batch_invariant_MODE = True _batch_invariant_LIB = torch.library.Library("aten", "IMPL") _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA") _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA") _batch_invariant_LIB.impl( "aten::_log_softmax", _log_softmax_batch_invariant, "CUDA" ) _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA") if enable_bmm: _batch_invariant_LIB.impl("aten::bmm", bmm_batch_invariant, "CUDA") # Also monkeypatch torch.bmm directly as a fallback _original_torch_bmm = torch.bmm torch.bmm = bmm_batch_invariant def disable_batch_invariant_mode(): global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm if _batch_invariant_LIB is not None: _batch_invariant_LIB._destroy() if _original_torch_bmm is not None: torch.bmm = _original_torch_bmm _original_torch_bmm = None _batch_invariant_MODE = False _batch_invariant_LIB = None @contextlib.contextmanager def set_batch_invariant_mode(enabled: bool = True): global _batch_invariant_MODE, _batch_invariant_LIB old_data = (_batch_invariant_MODE, _batch_invariant_LIB) if enabled: enable_batch_invariant_mode() else: disable_batch_invariant_mode() yield if _batch_invariant_LIB is not None: _batch_invariant_LIB._destroy() _batch_invariant_MODE, _batch_invariant_LIB = old_data AttentionBlockSize = namedtuple("AttentionBlockSize", ["block_m", "block_n"]) def get_batch_invariant_attention_block_size() -> AttentionBlockSize: return AttentionBlockSize(block_m=16, block_n=16) ================================================ FILE: python/sglang/srt/batch_overlap/operations.py ================================================ from __future__ import annotations import os from contextlib import contextmanager from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Callable, Dict, Generator, List, Sequence, Union import torch from sglang.srt.layers.dp_attention import set_dp_buffer_len if TYPE_CHECKING: from sglang.srt.model_executor.forward_batch_info import ForwardBatch _ENABLE_PROFILE = bool(int(os.environ.get("SGLANG_OPERATIONS_ENABLE_PROFILE", "0"))) if _ENABLE_PROFILE: import nvtx def execute_operations(inputs, operations): stages = _convert_operations_to_stages(operations) executor = _StageExecutor("primary", stages, inputs=inputs) for _ in range(executor.num_stages): executor.next() assert executor.done return executor.output def execute_overlapped_operations( inputs_arr: Sequence, operations_arr: Sequence, delta_stages: Sequence[int], ) -> Sequence: # Make it explicit for clarity; if we need multi-batch overlap, this can be generalized inputs_a, inputs_b = inputs_arr operations_a, operations_b = operations_arr delta_stage_a, delta_stage_b = delta_stages assert delta_stage_a == 0 delta_stage = delta_stage_b stages_a = _convert_operations_to_stages(operations_a) stages_b = _convert_operations_to_stages(operations_b) executor_a = _StageExecutor("a", stages_a, inputs=inputs_a) executor_b = _StageExecutor("b", stages_b, inputs=inputs_b) for _ in range(delta_stage): executor_a.next() for _ in range(executor_a.num_stages - delta_stage): executor_a.next() executor_b.next() for _ in range(delta_stage): executor_b.next() assert executor_a.done and executor_b.done return [executor_a.output, executor_b.output] class YieldOperation: pass @dataclass class ExecutionOperation: debug_name: str fn: Callable Operation = Union[YieldOperation, ExecutionOperation, Callable] Stage = List[ExecutionOperation] class _StageExecutor: def __init__(self, debug_name: str, stages: List[Stage], inputs: dict): self._debug_name = debug_name self._stages = stages self._index = 0 self._stage_state = _StateDict() self._stage_output = inputs # handling DP attention forward_batch: ForwardBatch = inputs["forward_batch"] self._global_dp_buffer_len = forward_batch.global_dp_buffer_len self._local_dp_buffer_len = forward_batch.tbo_padded_len self._global_num_tokens = forward_batch.global_num_tokens_cpu self._is_dp_max_padding = forward_batch.dp_padding_mode.is_max_len() def next(self): assert not self.done stage = self._stages[self._index] # TODO: We currently always call set_dp_buffer_len here because sub-batches # may have different padded lengths. It can likely be removed after TBO slice & # pad logic is refactored. set_dp_buffer_len( self._global_dp_buffer_len, self._local_dp_buffer_len, self._is_dp_max_padding, self._global_num_tokens, ) with _annotate_region(debug_name=f"{self._debug_name}{self._index}"): for op in stage: with _annotate_region(debug_name=op.debug_name): self._stage_output = op.fn( state=self._stage_state, **( self._stage_output if self._stage_output is not None else {} ), ) self._index += 1 @property def output(self): assert self.done return self._stage_output @property def done(self): return self._index >= self.num_stages @property def num_stages(self): return len(self._stages) @contextmanager def _annotate_region(debug_name): if _ENABLE_PROFILE: with torch.autograd.profiler.record_function(debug_name): with nvtx.annotate(debug_name): yield else: yield class _StateDict: def __init__(self): self._data = {} def __setattr__(self, key, value): if key == "_data": super().__setattr__(key, value) return assert ( key not in self._data ), f"`{key}` already exist, are you sure you want to override it?" self._data[key] = value def __getattr__(self, item): return self._data[item] def __delattr__(self, item): del self._data[item] def pop(self, item): return self._data.pop(item) def update(self, values: Dict[str, Any]): for k, v in values.items(): setattr(self, k, v) def get(self, item): return self._data.get(item) def clear(self, expect_keys: Sequence[str]): if set(self._data.keys()) != set(expect_keys): raise Exception( f"Unexpected keys when clearing. This may indicate you do not release memory early enough but leave it until here. {list(self._data.keys())=} {expect_keys=}" ) self._data.clear() def _convert_operations_to_stages(operations: List[Operation]) -> List[Stage]: operations = _decorate_operations(operations) operation_chunks = list( _chunk_by_separator(operations, lambda op: isinstance(op, YieldOperation)) ) assert all(len(chunk) > 0 for chunk in operation_chunks) return operation_chunks def _chunk_by_separator( items: List[Any], is_separator: Callable[[Any], bool] ) -> Generator[List[Any], None, None]: pending_items = [] for item in items: if is_separator(item): yield pending_items pending_items = [] else: pending_items.append(item) if len(pending_items) > 0: yield pending_items def _decorate_operations(operations: List[Operation], debug_name_prefix: str = ""): return [_decorate_operation(op, debug_name_prefix) for op in operations] def _decorate_operation(operation: Operation, debug_name_prefix: str): if isinstance(operation, YieldOperation): return operation return ExecutionOperation( debug_name=debug_name_prefix + getattr(operation, "__name__", "unknown").replace("op_", ""), fn=operation, ) ================================================ FILE: python/sglang/srt/batch_overlap/operations_strategy.py ================================================ from dataclasses import dataclass from typing import List, Optional import torch from sglang.srt.batch_overlap import operations from sglang.srt.batch_overlap.operations import Operation from sglang.srt.layers.moe.token_dispatcher import DeepEPConfig from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.utils import is_hip _is_hip = is_hip() @dataclass class OperationsStrategy: operations: List[Operation] deep_gemm_num_sms: Optional[int] = None tbo_delta_stages: Optional[int] = None @classmethod def concat(cls, items: List["OperationsStrategy"]) -> "OperationsStrategy": return OperationsStrategy( operations=[x for item in items for x in item.operations], deep_gemm_num_sms=_assert_all_same( [item.deep_gemm_num_sms for item in items] ), tbo_delta_stages=_assert_all_same( [item.tbo_delta_stages for item in items] ), ) @staticmethod def init_new_tbo( layers: torch.nn.ModuleList, forward_mode: ForwardMode, ) -> "OperationsStrategy": layer_name = layers[0].__class__.__name__ if layer_name == "DeepseekV2DecoderLayer": return OperationsStrategy.concat( [ _compute_moe_deepseek_layer_operations_strategy_tbo( layer, forward_mode ) for layer in layers ] ) elif layer_name == "Qwen3MoeDecoderLayer": return OperationsStrategy.concat( [ _compute_moe_qwen3_layer_operations_strategy_tbo( layer, forward_mode ) for layer in layers ] ) elif layer_name == "MiMoV2DecoderLayer": return OperationsStrategy.concat( [ _compute_moe_mimov2_layer_operations_strategy_tbo( layer, forward_mode ) for layer in layers ] ) else: raise NotImplementedError def _assert_all_same(items: List): assert all(item == items[0] for item in items) return items[0] # -------------------------------- Strategy for DeepSeek --------------------------------------- # TODO can refactor to make it more fancy if we have more complex strategies def _compute_moe_deepseek_layer_operations_strategy_tbo( layer: torch.nn.Module, forward_mode: ForwardMode, ) -> OperationsStrategy: assert layer.is_layer_sparse, "dense layer TBO not yet implemented" if forward_mode == ForwardMode.EXTEND: return _compute_moe_deepseek_blog_prefill(layer) elif ( forward_mode == ForwardMode.DECODE or forward_mode == ForwardMode.TARGET_VERIFY ): return _compute_moe_deepseek_blog_decode(layer) else: raise NotImplementedError(f"Unsupported {forward_mode=}") def _compute_moe_deepseek_blog_prefill(layer): device_properties = torch.cuda.get_device_properties(device="cuda") total_num_sms = device_properties.multi_processor_count deep_gemm_num_sms = None if not _is_hip: deep_gemm_num_sms = total_num_sms - DeepEPConfig.get_instance().num_sms return OperationsStrategy( deep_gemm_num_sms=deep_gemm_num_sms, tbo_delta_stages=0, operations=[ layer.op_comm_prepare_attn, layer.self_attn.op_prepare, layer.self_attn.op_core, layer.op_comm_prepare_mlp, layer.mlp.op_gate, layer.mlp.op_select_experts, layer.mlp.op_dispatch_a, operations.YieldOperation(), layer.mlp.op_dispatch_b, layer.mlp.op_experts, layer.mlp.op_combine_a, operations.YieldOperation(), layer.mlp.op_shared_experts, layer.mlp.op_combine_b, layer.mlp.op_output, layer.op_comm_postprocess_layer, ], ) def _compute_moe_deepseek_blog_decode(layer): return OperationsStrategy( deep_gemm_num_sms=None, tbo_delta_stages=2, operations=[ layer.op_comm_prepare_attn, layer.self_attn.op_prepare, operations.YieldOperation(), layer.self_attn.op_core, layer.op_comm_prepare_mlp, layer.mlp.op_gate, layer.mlp.op_select_experts, operations.YieldOperation(), layer.mlp.op_dispatch_a, layer.mlp.op_shared_experts, operations.YieldOperation(), layer.mlp.op_dispatch_b, layer.mlp.op_experts, layer.mlp.op_combine_a, operations.YieldOperation(), layer.mlp.op_combine_b, operations.YieldOperation(), layer.mlp.op_output, layer.op_comm_postprocess_layer, ], ) # -------------------------------- Strategy for Qwen3 --------------------------------------- # TODO: unstable, current strategy is almost the same as DeepSeek, keep redundant code here for # convenience to adjust strategy def _compute_moe_qwen3_layer_operations_strategy_tbo( layer: torch.nn.Module, forward_mode: ForwardMode, ) -> OperationsStrategy: assert layer.is_layer_sparse, "qwen3 moe only support sparse layers" if forward_mode == ForwardMode.EXTEND: return _compute_moe_qwen3_prefill(layer) elif ( forward_mode == ForwardMode.DECODE or forward_mode == ForwardMode.TARGET_VERIFY ): return _compute_moe_qwen3_decode(layer) else: raise NotImplementedError(f"Unsupported {forward_mode=}") def _compute_moe_qwen3_prefill(layer): device_properties = torch.cuda.get_device_properties(device="cuda") total_num_sms = device_properties.multi_processor_count deep_gemm_num_sms = None if not _is_hip: deep_gemm_num_sms = total_num_sms - DeepEPConfig.get_instance().num_sms return OperationsStrategy( deep_gemm_num_sms=deep_gemm_num_sms, tbo_delta_stages=0, operations=[ layer.op_comm_prepare_attn, layer.self_attn.op_prepare, layer.self_attn.op_core, layer.op_comm_prepare_mlp, layer.mlp.op_gate, layer.mlp.op_select_experts, layer.mlp.op_dispatch_a, operations.YieldOperation(), layer.mlp.op_dispatch_b, layer.mlp.op_experts, layer.mlp.op_combine_a, operations.YieldOperation(), layer.mlp.op_combine_b, layer.mlp.op_output, layer.op_comm_postprocess_layer, ], ) def _compute_moe_qwen3_decode(layer): return OperationsStrategy( deep_gemm_num_sms=None, tbo_delta_stages=2, operations=[ layer.op_comm_prepare_attn, layer.self_attn.op_prepare, operations.YieldOperation(), layer.self_attn.op_core, layer.op_comm_prepare_mlp, layer.mlp.op_gate, layer.mlp.op_select_experts, operations.YieldOperation(), layer.mlp.op_dispatch_a, operations.YieldOperation(), layer.mlp.op_dispatch_b, layer.mlp.op_experts, layer.mlp.op_combine_a, operations.YieldOperation(), layer.mlp.op_combine_b, layer.mlp.op_output, layer.op_comm_postprocess_layer, operations.YieldOperation(), ], ) # -------------------------------- Strategy for MiMoV2DecoderLayer --------------------------------------- # TODO: unstable; current strategy matches DeepSeek for the common operations (MiMoV2 has no op_shared_experts), # so we keep this redundant code here for convenience when adjusting the strategy def _compute_moe_mimov2_layer_operations_strategy_tbo( layer: torch.nn.Module, forward_mode: ForwardMode, ) -> OperationsStrategy: assert layer.is_layer_sparse, "MiMoV2DecoderLayer moe only support sparse layers" if forward_mode == ForwardMode.EXTEND: return _compute_moe_mimov2_prefill(layer) elif ( forward_mode == ForwardMode.DECODE or forward_mode == ForwardMode.TARGET_VERIFY ): return _compute_moe_mimov2_decode(layer) else: raise NotImplementedError(f"Unsupported {forward_mode=}") def _compute_moe_mimov2_prefill(layer): device_properties = torch.cuda.get_device_properties(device="cuda") total_num_sms = device_properties.multi_processor_count deep_gemm_num_sms = total_num_sms - DeepEPConfig.get_instance().num_sms return OperationsStrategy( deep_gemm_num_sms=deep_gemm_num_sms, tbo_delta_stages=0, operations=[ layer.op_comm_prepare_attn, layer.self_attn.op_prepare, layer.self_attn.op_core, layer.op_comm_prepare_mlp, layer.mlp.op_gate, layer.mlp.op_select_experts, layer.mlp.op_dispatch_a, operations.YieldOperation(), layer.mlp.op_dispatch_b, layer.mlp.op_experts, layer.mlp.op_combine_a, operations.YieldOperation(), layer.mlp.op_combine_b, layer.mlp.op_output, layer.op_comm_postprocess_layer, ], ) def _compute_moe_mimov2_decode(layer): return OperationsStrategy( deep_gemm_num_sms=None, tbo_delta_stages=2, operations=[ layer.op_comm_prepare_attn, layer.self_attn.op_prepare, operations.YieldOperation(), layer.self_attn.op_core, layer.op_comm_prepare_mlp, layer.mlp.op_gate, layer.mlp.op_select_experts, operations.YieldOperation(), layer.mlp.op_dispatch_a, operations.YieldOperation(), layer.mlp.op_dispatch_b, layer.mlp.op_experts, layer.mlp.op_combine_a, operations.YieldOperation(), layer.mlp.op_combine_b, layer.mlp.op_output, layer.op_comm_postprocess_layer, operations.YieldOperation(), ], ) ================================================ FILE: python/sglang/srt/batch_overlap/single_batch_overlap.py ================================================ # Copyright 2025 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== from __future__ import annotations from dataclasses import dataclass from typing import Optional import torch from sglang.srt.environ import envs from sglang.srt.layers.moe import get_moe_runner_backend from sglang.srt.layers.moe.utils import is_sbo_enabled from sglang.srt.utils import is_blackwell class SboFlags: # TODO may have: "enable_dispatch_gateup_gemm_two_stream_overlap", ... @classmethod def enable_combine_down_gemm_two_stream_overlap(cls): return ( is_sbo_enabled() # currently only cutedsl backend supports it and ( get_moe_runner_backend().is_flashinfer_cutedsl() or (get_moe_runner_backend().is_deep_gemm() and not is_blackwell()) ) ) @classmethod def enable_combine_shared_two_stream_overlap(cls): return ( is_sbo_enabled() and not cls.enable_dispatch_shared_one_stream_overlap() and not envs.SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO.get() ) @classmethod def enable_dispatch_shared_one_stream_overlap(cls): return is_sbo_enabled() and not is_blackwell() @classmethod def fuse_shared_experts_inside_sbo(cls): return ( cls.enable_combine_shared_two_stream_overlap() or cls.enable_dispatch_shared_one_stream_overlap() ) @dataclass class CombineOverlapArgs: # this "overlap" flag means overlapping with down gemm, not the general two-stream overlap overlap: bool stream: torch.cuda.Stream wait_event: torch.cuda.Event num_sms: Optional[int] = None signal: Optional[torch.Tensor] = None block_m: Optional[int] = 64 threshold: Optional[int] = 0 @dataclass class DownGemmOverlapArgs: num_sms: int signal: torch.Tensor start_event: torch.cuda.Event def compute_overlap_args(dispatch_output, alt_stream): if not ( SboFlags.enable_combine_down_gemm_two_stream_overlap() or SboFlags.enable_combine_shared_two_stream_overlap() ): return None, None, {} hidden_states = dispatch_output.hidden_states num_local_experts, num_tokens_static, hidden_dim = hidden_states.shape total_num_sms = torch.cuda.get_device_properties( device="cuda" ).multi_processor_count if envs.SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS.is_set(): communicate_num_sms = envs.SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS.get() else: communicate_num_sms = 32 if is_blackwell() else 3 compute_num_sms = total_num_sms - communicate_num_sms assert alt_stream is not None combine_wait_event = torch.cuda.Event() combine_overlap_args = CombineOverlapArgs( overlap=False, num_sms=communicate_num_sms, stream=alt_stream, wait_event=combine_wait_event, ) meta_overlap_args = dict( compute_num_sms=compute_num_sms, ) down_gemm_overlap_args = None if SboFlags.enable_combine_down_gemm_two_stream_overlap(): # TODO use zero_allocator to remove this `torch.zeros` call # NOTE ours v2 use uint32 not int32 currently if is_blackwell(): combine_signal = torch.zeros( num_local_experts, dtype=torch.uint32, device=hidden_states.device ) else: MIN_BLOCK_M = 64 combine_signal_size = num_local_experts * ( (num_tokens_static + MIN_BLOCK_M - 1) // MIN_BLOCK_M ) combine_signal = torch.zeros( combine_signal_size, dtype=torch.int32, device=hidden_states.device ) down_gemm_overlap_args = DownGemmOverlapArgs( signal=combine_signal, start_event=combine_wait_event, num_sms=compute_num_sms, ) combine_overlap_args.overlap = True combine_overlap_args.signal = combine_signal combine_overlap_args.threshold = compute_num_sms else: meta_overlap_args |= dict( record_event_after_down=combine_wait_event, ) return combine_overlap_args, down_gemm_overlap_args, meta_overlap_args ================================================ FILE: python/sglang/srt/batch_overlap/two_batch_overlap.py ================================================ from __future__ import annotations import copy import dataclasses import logging from dataclasses import replace from typing import TYPE_CHECKING, Dict, List, Optional, Sequence import torch from sglang.srt.batch_overlap.operations import ( execute_operations, execute_overlapped_operations, ) from sglang.srt.batch_overlap.operations_strategy import OperationsStrategy from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.communicator import ( CommunicateContext, CommunicateSummableTensorPairFn, ScatterMode, ) from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.layers.moe import ( get_deepep_mode, get_moe_a2a_backend, get_tbo_token_distribution_threshold, is_tbo_enabled, ) from sglang.srt.layers.moe.token_dispatcher import ( DeepEPDispatcher, MooncakeEPDispatcher, MoriEPDispatcher, NixlEPDispatcher, ) from sglang.srt.layers.moe.token_dispatcher.base import BaseDispatcher from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.model_executor.forward_batch_info import ( ForwardBatch, ForwardMode, compute_position, ) from sglang.srt.server_args import get_global_server_args from sglang.srt.speculative.spec_info import SpecInput from sglang.srt.utils import BumpAllocator, empty_context, get_bool_env_var, is_hip if TYPE_CHECKING: from sglang.srt.batch_overlap.single_batch_overlap import CombineOverlapArgs from sglang.srt.layers.moe.token_dispatcher import DispatchOutput from sglang.srt.speculative.eagle_info import EagleVerifyInput _is_hip = is_hip() _tbo_debug = get_bool_env_var("SGLANG_TBO_DEBUG") logger = logging.getLogger(__name__) # -------------------------------- Compute Basic Info --------------------------------------- def get_token_num_per_seq( forward_mode: ForwardMode, spec_info: Optional[SpecInput] = None, ): if forward_mode.is_target_verify(): return spec_info.draft_token_num elif forward_mode.is_decode(): return 1 elif forward_mode.is_idle(): return 0 else: # For extend, we should not use `token_num_per_seq`. return None # TODO: may smartly disable TBO when batch size is too small b/c it will slow down def compute_split_seq_index( forward_mode: ForwardMode, num_tokens: int, extend_lens: Optional[Sequence[int]], token_num_per_seq: Optional[int], ) -> Optional[int]: if forward_mode == ForwardMode.EXTEND: assert extend_lens is not None return _split_extend_seqs(extend_lens) elif forward_mode.is_target_verify() or forward_mode.is_decode(): assert token_num_per_seq is not None return (num_tokens // token_num_per_seq) // 2 elif forward_mode.is_idle() or forward_mode.is_prebuilt(): assert num_tokens == 0 return 0 else: raise NotImplementedError() def _is_two_chunk_split_enabled(extend_lens: Sequence[int]) -> bool: if extend_lens is None: return False vanilla_split_seq_index = _split_array_by_balanced_sum(extend_lens) left_sum = sum(extend_lens[:vanilla_split_seq_index]) overall_sum = sum(extend_lens) threshold = get_tbo_token_distribution_threshold() assert threshold <= 0.5, f"{threshold=}" return left_sum < overall_sum * threshold or left_sum > overall_sum * ( 1 - threshold ) def _split_extend_seqs(arr: Sequence[int]) -> int: if _is_two_chunk_split_enabled(arr): return _split_array_by_cum_less_than_half(arr) return _split_array_by_balanced_sum(arr) def _split_array_by_cum_less_than_half(arr: Sequence[int]) -> int: left_sum = 0 overall_sum = sum(arr) half_sum = overall_sum // 2 chosen_index = 0 for i in range(len(arr)): left_sum += arr[i] if left_sum > half_sum: chosen_index = i break return chosen_index def _split_array_by_balanced_sum(arr: Sequence[int]) -> int: overall_sum = sum(arr) left_sum = 0 min_diff = float("inf") best_index = 0 for i in range(1, len(arr)): left_sum += arr[i - 1] right_sum = overall_sum - left_sum diff = abs(left_sum - right_sum) if diff <= min_diff: min_diff = diff best_index = i else: break return best_index def _update_device_and_sum_field_from_cpu_field( batch: ForwardBatch, cpu_field: str, device_field: str, sum_field: str = None ): cpu_value = getattr(batch, cpu_field, None) old_device_value = getattr(batch, device_field, None) if ( cpu_value is None or old_device_value is None or not (isinstance(cpu_value, torch.Tensor) or isinstance(cpu_value, list)) ): return new_device_value = ( cpu_value if isinstance(cpu_value, torch.Tensor) else torch.tensor(cpu_value, dtype=old_device_value.dtype) ).to(device=get_global_server_args().device, non_blocking=True) setattr(batch, device_field, new_device_value) if sum_field is not None: sum_value = ( cpu_value.sum().item() if isinstance(cpu_value, torch.Tensor) else sum(cpu_value) ) setattr(batch, sum_field, sum_value) def _compute_mask_offset(seq_index: int, spec_info: Optional[EagleVerifyInput]) -> int: if seq_index == 0: return 0 offset = 0 max_seq_len = min(seq_index, spec_info.seq_lens_cpu.shape[0]) for i in range(max_seq_len): offset += ( spec_info.seq_lens_cpu[i] + spec_info.draft_token_num ) * spec_info.draft_token_num return offset def split_spec_info( spec_info: Optional[EagleVerifyInput], start_seq_index: int, end_seq_index: int, start_token_index: int, end_token_index: int, ): if spec_info is None: return None if spec_info.draft_token is not None: draft_token = spec_info.draft_token[start_token_index:end_token_index] else: draft_token = None if spec_info.custom_mask is not None and spec_info.draft_token is not None: custom_mask_start = _compute_mask_offset(start_seq_index, spec_info) if end_seq_index == spec_info.seq_lens_cpu.shape[0]: custom_mask_end = spec_info.custom_mask.shape[0] else: custom_mask_end = _compute_mask_offset(end_seq_index, spec_info) if custom_mask_end > custom_mask_start: custom_mask = spec_info.custom_mask[custom_mask_start:custom_mask_end] else: custom_mask = spec_info.custom_mask else: custom_mask = spec_info.custom_mask if spec_info.positions is not None: positions = spec_info.positions[start_token_index:end_token_index] else: positions = None if spec_info.retrive_index is not None: retrive_index = spec_info.retrive_index[start_seq_index:end_seq_index] else: retrive_index = None if spec_info.retrive_next_token is not None: retrive_next_token = spec_info.retrive_next_token[start_seq_index:end_seq_index] else: retrive_next_token = None if spec_info.retrive_next_sibling is not None: retrive_next_sibling = spec_info.retrive_next_sibling[ start_seq_index:end_seq_index ] else: retrive_next_sibling = None if spec_info.retrive_cum_len is not None: retrive_cum_len = spec_info.retrive_cum_len[start_seq_index:end_seq_index] else: retrive_cum_len = None if spec_info.seq_lens_cpu is not None: seq_lens_cpu = spec_info.seq_lens_cpu[start_seq_index:end_seq_index] else: seq_lens_cpu = None if seq_lens_cpu is not None: seq_lens_sum = seq_lens_cpu.sum() else: seq_lens_sum = None output_spec_info = replace( spec_info, custom_mask=custom_mask, draft_token=draft_token, positions=positions, retrive_index=retrive_index, retrive_next_token=retrive_next_token, retrive_next_sibling=retrive_next_sibling, retrive_cum_len=retrive_cum_len, seq_lens_cpu=seq_lens_cpu, seq_lens_sum=seq_lens_sum, ) return output_spec_info def compute_split_token_index( split_seq_index: int, forward_mode: "ForwardMode", extend_seq_lens: Optional[Sequence[int]], token_num_per_seq: Optional[int], ) -> int: if forward_mode == ForwardMode.EXTEND: assert extend_seq_lens is not None if _is_two_chunk_split_enabled(extend_seq_lens): return sum(extend_seq_lens) // 2 return sum(extend_seq_lens[:split_seq_index]) elif forward_mode.is_target_verify() or forward_mode.is_decode(): assert token_num_per_seq is not None return split_seq_index * token_num_per_seq elif forward_mode.is_idle(): assert split_seq_index == 0 return 0 else: raise NotImplementedError def compute_split_indices_for_cuda_graph_replay( forward_mode: ForwardMode, cuda_graph_num_tokens: int, spec_info: Optional[SpecInput], ): forward_mode_for_tbo_split = ( forward_mode if forward_mode != ForwardMode.IDLE else ForwardMode.DECODE ) token_num_per_seq = get_token_num_per_seq( forward_mode=forward_mode, spec_info=spec_info ) tbo_split_seq_index = compute_split_seq_index( forward_mode=forward_mode_for_tbo_split, num_tokens=cuda_graph_num_tokens, extend_lens=None, token_num_per_seq=token_num_per_seq, ) tbo_split_token_index = compute_split_token_index( split_seq_index=tbo_split_seq_index, forward_mode=forward_mode_for_tbo_split, extend_seq_lens=None, token_num_per_seq=token_num_per_seq, ) return tbo_split_seq_index, tbo_split_token_index # -------------------------------- Preparation --------------------------------------- class TboCudaGraphRunnerPlugin: def __init__(self): self._tbo_children_num_token_non_padded = torch.zeros((2,), dtype=torch.int32) def capture_one_batch_size(self, batch: ForwardBatch, num_tokens: int): if not is_tbo_enabled(): return token_num_per_seq = get_token_num_per_seq( forward_mode=batch.forward_mode, spec_info=batch.spec_info ) batch.tbo_split_seq_index = compute_split_seq_index( forward_mode=batch.forward_mode, num_tokens=num_tokens, extend_lens=None, token_num_per_seq=token_num_per_seq, ) # For simplicity, when two_batch_overlap is enabled, we only capture CUDA Graph for tbo=true assert batch.tbo_split_seq_index is not None, f"{num_tokens=}" self._tbo_children_num_token_non_padded[...] = ( TboForwardBatchPreparer.compute_tbo_children_num_token_non_padded(batch) ) TboForwardBatchPreparer.prepare_raw( batch, tbo_children_num_token_non_padded=self._tbo_children_num_token_non_padded, ) def replay_prepare( self, forward_mode: ForwardMode, bs: int, num_token_non_padded: int, spec_info: Optional[SpecInput], ): token_num_per_seq = get_token_num_per_seq( forward_mode=forward_mode, spec_info=spec_info ) tbo_split_seq_index, tbo_split_token_index = ( compute_split_indices_for_cuda_graph_replay( forward_mode=forward_mode, cuda_graph_num_tokens=bs * token_num_per_seq, spec_info=spec_info, ) ) self._tbo_children_num_token_non_padded[...] = ( TboForwardBatchPreparer.compute_tbo_children_num_token_non_padded_raw( tbo_split_token_index=tbo_split_token_index, num_token_non_padded=num_token_non_padded, ) ) class TboDPAttentionPreparer: def prepare_all_gather( self, local_batch: ScheduleBatch, ): deepep_mode = get_deepep_mode() enable_a2a_moe = not get_moe_a2a_backend().is_none() enable_two_batch_overlap = is_tbo_enabled() self.enable_two_batch_overlap = enable_two_batch_overlap if local_batch is not None: token_num_per_seq = get_token_num_per_seq( forward_mode=local_batch.forward_mode, spec_info=local_batch.spec_info ) if ( local_batch.forward_mode.is_target_verify() or local_batch.forward_mode.is_decode() ): num_tokens = local_batch.batch_size() * token_num_per_seq elif local_batch.forward_mode.is_prebuilt(): num_tokens = 0 else: num_tokens = local_batch.extend_num_tokens self.local_tbo_split_seq_index = compute_split_seq_index( forward_mode=local_batch.forward_mode, num_tokens=num_tokens, extend_lens=local_batch.extend_lens, token_num_per_seq=token_num_per_seq, ) resolved_deepep_mode = deepep_mode.resolve(local_batch.is_extend_in_batch) local_can_run_tbo = (self.local_tbo_split_seq_index is not None) and not ( ( local_batch.forward_mode.is_extend() and not local_batch.forward_mode.is_target_verify() ) and enable_a2a_moe and (resolved_deepep_mode.is_low_latency()) ) else: self.local_tbo_split_seq_index = 0 local_can_run_tbo = True local_forward_mode = self._compute_local_forward_mode(local_batch) return local_can_run_tbo, local_forward_mode def compute_output(self, partial_global_info): # Perform only one Device-to-Host (D2H) memory copy cpu_data = partial_global_info[:, :2].cpu() local_can_run_tbo_aggregated = min(cpu_data[:, 0].tolist()) forward_modes = cpu_data[:, 1].tolist() global_forward_mode, forward_mode_agree = self._compute_global_forward_mode( forward_modes ) can_run_tbo = ( self.enable_two_batch_overlap and local_can_run_tbo_aggregated and forward_mode_agree ) tbo_split_seq_index = self.local_tbo_split_seq_index if can_run_tbo else None global_forward_mode = global_forward_mode if can_run_tbo else None return tbo_split_seq_index, global_forward_mode @staticmethod def _compute_local_forward_mode(local_batch): return ( local_batch.forward_mode if local_batch is not None else ForwardMode.IDLE ).value @staticmethod def _compute_global_forward_mode(forward_modes): forward_modes_excluding_idle_and_prebuilt = [ x for x in forward_modes if x != ForwardMode.IDLE.value and x != ForwardMode.PREBUILT.value ] if not forward_modes_excluding_idle_and_prebuilt: return ForwardMode.IDLE, False forward_mode_agree = TboDPAttentionPreparer._is_all_same( forward_modes_excluding_idle_and_prebuilt ) global_forward_mode = ( ForwardMode(forward_modes_excluding_idle_and_prebuilt[0]) if forward_mode_agree else None ) return global_forward_mode, forward_mode_agree @staticmethod def _is_all_same(x): return all(value == x[0] for value in x) class TboForwardBatchPreparer: @classmethod def prepare(cls, batch: ForwardBatch, is_draft_worker: bool = False): if batch.tbo_split_seq_index is None or is_draft_worker: return tbo_children_num_token_non_padded = ( cls.compute_tbo_children_num_token_non_padded(batch) ) cls.prepare_raw( batch, tbo_children_num_token_non_padded=tbo_children_num_token_non_padded ) @classmethod def prepare_raw( cls, batch: ForwardBatch, tbo_children_num_token_non_padded: torch.Tensor ): from sglang.srt.layers.attention.tbo_backend import TboAttnBackend tbo_split_token_index = cls._compute_split_token_index(batch) is_enable_two_chunk = ( batch.forward_mode == ForwardMode.EXTEND and _is_two_chunk_split_enabled(batch.extend_seq_lens_cpu) ) if _tbo_debug: logger.info( f"TboForwardBatchPreparer.prepare " f"is_enable_two_chunk={is_enable_two_chunk} " f"tbo_split_seq_index={batch.tbo_split_seq_index} " f"tbo_split_token_index={tbo_split_token_index} " f"extend_seq_lens={batch.extend_seq_lens_cpu} " f"bs={batch.batch_size} " f"forward_mode={batch.forward_mode}" ) assert isinstance(batch.attn_backend, TboAttnBackend) attn_backend_child_a, attn_backend_child_b = batch.attn_backend.children [out_num_token_non_padded_a, out_num_token_non_padded_b] = ( tbo_children_num_token_non_padded ) child_a = cls.filter_batch( batch, start_token_index=0, end_token_index=tbo_split_token_index, start_seq_index=0, end_seq_index=( batch.tbo_split_seq_index + 1 if is_enable_two_chunk else batch.tbo_split_seq_index ), output_attn_backend=attn_backend_child_a, out_num_token_non_padded=out_num_token_non_padded_a, ) child_b = cls.filter_batch( batch, start_token_index=tbo_split_token_index, end_token_index=batch.input_ids.shape[0], start_seq_index=batch.tbo_split_seq_index, end_seq_index=batch.batch_size, output_attn_backend=attn_backend_child_b, out_num_token_non_padded=out_num_token_non_padded_b, ) if is_enable_two_chunk: cls.derive_fields_related_to_seq_len_for_two_chunk( batch, child_a=child_a, child_b=child_b, tbo_split_seq_index=batch.tbo_split_seq_index, ) assert batch.tbo_children is None batch.tbo_children = [child_a, child_b] @classmethod def derive_fields_related_to_seq_len_for_two_chunk( cls, batch: ForwardBatch, *, child_a: ForwardBatch, child_b: ForwardBatch, tbo_split_seq_index: int, ): extend_seq_lens_cpu = batch.extend_seq_lens_cpu overall_seq_lens_sum = sum(extend_seq_lens_cpu) half_seq_lens_sum = overall_seq_lens_sum // 2 left_last_seq_token_num = half_seq_lens_sum - sum( extend_seq_lens_cpu[:tbo_split_seq_index] ) right_first_seq_token_num = ( extend_seq_lens_cpu[tbo_split_seq_index] - left_last_seq_token_num ) # making deepcopy to be extra safe child_a.extend_seq_lens_cpu = copy.deepcopy(child_a.extend_seq_lens_cpu) child_a.extend_seq_lens_cpu[-1] = left_last_seq_token_num child_b.extend_seq_lens_cpu = copy.deepcopy(child_b.extend_seq_lens_cpu) child_b.extend_seq_lens_cpu[0] = right_first_seq_token_num for child in [child_a, child_b]: _update_device_and_sum_field_from_cpu_field( batch=child, cpu_field="extend_seq_lens_cpu", device_field="extend_seq_lens", sum_field="extend_num_tokens", ) assert ( child_a.extend_num_tokens == half_seq_lens_sum ), f"{child_a.extend_num_tokens=}, {half_seq_lens_sum=}" child_a.seq_lens_cpu = copy.deepcopy(child_a.seq_lens_cpu) child_a.seq_lens_cpu[-1] = ( child_a.extend_seq_lens_cpu[-1] + child_a.extend_prefix_lens_cpu[-1] ) _update_device_and_sum_field_from_cpu_field( batch=child_a, cpu_field="seq_lens_cpu", device_field="seq_lens", sum_field="seq_lens_sum", ) child_b.extend_prefix_lens_cpu = copy.deepcopy(child_b.extend_prefix_lens_cpu) child_b.extend_prefix_lens_cpu[0] += left_last_seq_token_num _update_device_and_sum_field_from_cpu_field( batch=child_b, cpu_field="extend_prefix_lens_cpu", device_field="extend_prefix_lens", sum_field=None, ) _, child_b.extend_start_loc = compute_position( get_global_server_args().attention_backend, child_b.extend_prefix_lens, child_b.extend_seq_lens, child_b.extend_num_tokens, ) @classmethod def filter_batch( cls, batch: ForwardBatch, *, start_token_index: int, end_token_index: int, start_seq_index: int, end_seq_index: int, output_attn_backend: AttentionBackend, out_num_token_non_padded: torch.Tensor, ): assert ( end_token_index >= start_token_index ), f"{end_token_index=}, {start_token_index=}, batch={batch}" num_tokens = batch.input_ids.shape[0] num_seqs = batch.batch_size output_dict = dict() for key in [ "input_ids", "positions", "out_cache_loc", ]: old_value = getattr(batch, key) assert ( old_value.shape[0] == num_tokens ), f"{key=} {old_value=} {num_tokens=} {batch=}" output_dict[key] = old_value[start_token_index:end_token_index] attention_tp_size = get_attention_tp_size() output_dict["tbo_padded_len"] = ( (end_token_index - start_token_index - 1) // attention_tp_size + 1 ) * attention_tp_size for key in [ "req_pool_indices", "seq_lens", "seq_lens_cpu", "extend_seq_lens", "extend_prefix_lens", "extend_start_loc", "extend_prefix_lens_cpu", "extend_seq_lens_cpu", "extend_logprob_start_lens_cpu", "lora_ids", "rids", ]: old_value = getattr(batch, key) if old_value is None: continue elif batch.forward_mode.is_target_verify() and ( key == "extend_seq_lens" or key == "extend_prefix_lens" or key == "extend_start_loc" or key == "extend_prefix_lens_cpu" or key == "extend_seq_lens_cpu" or key == "extend_logprob_start_lens_cpu" ): output_dict[key] = None continue assert ( len(old_value) == num_seqs ), f"{key=} {old_value=} {num_seqs=} {batch=}" output_dict[key] = old_value[start_seq_index:end_seq_index] spec_info = getattr(batch, "spec_info") output_spec_info = split_spec_info( spec_info=spec_info, start_token_index=start_token_index, end_token_index=end_token_index, start_seq_index=start_seq_index, end_seq_index=end_seq_index, ) output_dict["spec_info"] = output_spec_info for key in [ "forward_mode", "is_extend_in_batch", "all_extend_in_batch", "return_logprob", "req_to_token_pool", "token_to_kv_pool", "can_run_dp_cuda_graph", "dp_padding_mode", "global_forward_mode", "is_prefill_only", "spec_algorithm", "capture_hidden_mode", "padded_static_len", "mrope_positions", # only used by qwen2-vl, thus not care "split_index", # for split prefill "orig_seq_lens", # only used by qwen-1m, thus not care ]: output_dict[key] = getattr(batch, key) if not batch.forward_mode.is_target_verify(): assert ( _compute_extend_num_tokens(batch.input_ids, batch.forward_mode) == batch.extend_num_tokens ), f"{batch=}" extend_num_tokens = _compute_extend_num_tokens( output_dict["input_ids"], output_dict["forward_mode"] ) # TODO improve, e.g. unify w/ `init_raw` if ( get_global_server_args().moe_dense_tp_size == 1 and batch.global_dp_buffer_len is not None ): sum_len = end_token_index - start_token_index global_dp_buffer_len = sum_len else: global_dp_buffer_len = None output_dict.update( dict( batch_size=end_seq_index - start_seq_index, seq_lens_sum=( output_dict["seq_lens_cpu"].sum() if "seq_lens_cpu" in output_dict else None ), extend_num_tokens=extend_num_tokens, attn_backend=output_attn_backend, num_token_non_padded=out_num_token_non_padded, # TODO: handle it when we need TBO + DeepSeek V3.2 num_token_non_padded_cpu=None, tbo_split_seq_index=None, tbo_parent_token_range=(start_token_index, end_token_index), tbo_children=None, original_global_num_tokens_cpu=None, global_num_tokens_gpu=None, global_num_tokens_cpu=None, global_dp_buffer_len=global_dp_buffer_len, global_num_tokens_for_logprob_gpu=None, global_num_tokens_for_logprob_cpu=None, sampling_info=None, # For logits and logprobs post processing, thus we do not care temp_scaled_logprobs=False, temperature=None, top_p_normalized_logprobs=False, top_p=None, mm_inputs=None, top_logprobs_nums=None, token_ids_logprobs=None, next_token_logits_buffer=None, return_hidden_states_before_norm=False, ) ) errors = [] for field in dataclasses.fields(ForwardBatch): if getattr(batch, field.name) is not None and field.name not in output_dict: errors.append( f"Field {field.name} has value, but is not yet supported (value={getattr(batch, field.name)} batch={batch})" ) if len(errors) > 0: raise Exception(f"{len(errors)} errors happen:\n" + "\n\n".join(errors)) return ForwardBatch(**output_dict) @classmethod def compute_tbo_children_num_token_non_padded(cls, batch: ForwardBatch): return cls.compute_tbo_children_num_token_non_padded_raw( tbo_split_token_index=cls._compute_split_token_index(batch), num_token_non_padded=len(batch.input_ids), ) @classmethod def compute_tbo_children_num_token_non_padded_raw( cls, tbo_split_token_index: int, num_token_non_padded: int ): # TODO we may make padding on both sub-batches to make it slightly more balanced value_a = min(tbo_split_token_index, num_token_non_padded) value_b = max(0, num_token_non_padded - tbo_split_token_index) return torch.tensor([value_a, value_b], dtype=torch.int32).to( device=get_global_server_args().device, non_blocking=True ) @classmethod def _compute_split_token_index(cls, batch: ForwardBatch): token_num_per_seq = get_token_num_per_seq( forward_mode=batch.forward_mode, spec_info=batch.spec_info ) return compute_split_token_index( split_seq_index=batch.tbo_split_seq_index, forward_mode=batch.forward_mode, extend_seq_lens=batch.extend_seq_lens_cpu, token_num_per_seq=token_num_per_seq, ) def _compute_extend_num_tokens(input_ids, forward_mode: ForwardMode): if ( forward_mode.is_decode() or forward_mode.is_idle() or forward_mode.is_target_verify() ): return None elif forward_mode.is_extend(): return input_ids.shape[0] raise NotImplementedError # -------------------------------- Execution --------------------------------------- def model_forward_maybe_tbo( layers, enable_tbo: bool, positions: torch.Tensor, forward_batch: ForwardBatch, hidden_states: torch.Tensor, input_data_scatter_mode: ScatterMode, residual: Optional[torch.Tensor], zero_allocator: Optional[BumpAllocator] = None, ): inputs = dict( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, residual=residual, zero_allocator=zero_allocator, ) layer_input_scatter_mode = layers[0].layer_scatter_modes.layer_input_mode operations_strategy = OperationsStrategy.init_new_tbo( layers, forward_batch.global_forward_mode ) if enable_tbo: return _model_forward_tbo( inputs=inputs, operations_strategy=operations_strategy, input_data_scatter_mode=input_data_scatter_mode, layer_input_scatter_mode=layer_input_scatter_mode, ) else: return _model_forward_non_tbo(inputs, operations_strategy) def _model_forward_tbo( inputs, operations_strategy: OperationsStrategy, input_data_scatter_mode: ScatterMode, layer_input_scatter_mode: ScatterMode, ): inputs_arr = _model_forward_tbo_split_inputs( **inputs, input_data_scatter_mode=input_data_scatter_mode, layer_input_scatter_mode=layer_input_scatter_mode, ) original_hidden_states_len = inputs["hidden_states"].shape[0] del inputs context = ( empty_context() if _is_hip else deep_gemm_wrapper.configure_deep_gemm_num_sms( operations_strategy.deep_gemm_num_sms ) ) with context: outputs_arr = execute_overlapped_operations( inputs_arr=inputs_arr, operations_arr=[operations_strategy.operations] * 2, delta_stages=[0, operations_strategy.tbo_delta_stages], ) return _model_forward_tbo_merge_outputs(*outputs_arr, original_hidden_states_len) def _model_forward_non_tbo(inputs, operations_strategy: OperationsStrategy): outputs = execute_operations(inputs, operations_strategy.operations) return outputs["hidden_states"], outputs["residual"] def _model_forward_tbo_split_inputs( hidden_states: torch.Tensor, residual: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, zero_allocator: Optional[BumpAllocator], input_data_scatter_mode: ScatterMode, layer_input_scatter_mode: ScatterMode, ) -> List[Dict]: tbo_splitter_scatter_mode = ScatterMode.TP_ATTN_FULL context = CommunicateContext.init_new() hidden_states, residual = CommunicateSummableTensorPairFn.execute( hidden_states_input_mode=input_data_scatter_mode, residual_input_mode=input_data_scatter_mode, output_mode=tbo_splitter_scatter_mode, hidden_states=hidden_states, residual=residual, forward_batch=forward_batch, context=context, ) inputs_arr = _model_forward_tbo_split_inputs_raw( hidden_states=hidden_states, residual=residual, positions=positions, forward_batch=forward_batch, zero_allocator=zero_allocator, ) def _post_transform(hidden_states, residual, forward_batch, **kwargs): hidden_states, residual = CommunicateSummableTensorPairFn.execute( hidden_states_input_mode=tbo_splitter_scatter_mode, residual_input_mode=tbo_splitter_scatter_mode, output_mode=layer_input_scatter_mode, hidden_states=hidden_states, residual=residual, forward_batch=forward_batch, context=context, ) return dict( hidden_states=hidden_states, residual=residual, forward_batch=forward_batch, **kwargs, ) return [_post_transform(**inputs) for inputs in inputs_arr] def _model_forward_tbo_split_inputs_raw( hidden_states: torch.Tensor, residual: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, zero_allocator: Optional[BumpAllocator], ) -> List[Dict]: return [ dict( **_model_forward_filter_inputs( hidden_states=hidden_states, residual=residual, positions=positions, output_forward_batch=output_forward_batch, tbo_subbatch_index=tbo_subbatch_index, ), **( dict(zero_allocator=zero_allocator) if zero_allocator is not None else {} ), ) for tbo_subbatch_index, output_forward_batch in enumerate( forward_batch.tbo_children ) ] def _model_forward_filter_inputs( hidden_states: torch.Tensor, residual: torch.Tensor, positions: torch.Tensor, output_forward_batch: ForwardBatch, tbo_subbatch_index: int, ) -> Dict: token_slice = slice(*output_forward_batch.tbo_parent_token_range) hidden_states = hidden_states[token_slice] residual = None if residual is None else residual[token_slice] positions = positions[token_slice] assert output_forward_batch.tbo_padded_len is not None padded_len = output_forward_batch.tbo_padded_len def _pad(x): nonlocal padded_len if x is None: return None if x.shape[0] == padded_len: return x res = torch.zeros((padded_len, *x.shape[1:]), dtype=x.dtype, device=x.device) res[: x.shape[0]] = x return res return dict( hidden_states=_pad(hidden_states), residual=_pad(residual), positions=_pad(positions), forward_batch=output_forward_batch, tbo_subbatch_index=tbo_subbatch_index, ) def _model_forward_tbo_merge_outputs(output_a, output_b, original_len): def _handle_key(name): value_a = output_a[name] value_b = output_b[name] assert (value_a is None) == (value_b is None) if value_a is None: return None s0, t0 = output_a["forward_batch"].tbo_parent_token_range s1, t1 = output_b["forward_batch"].tbo_parent_token_range res = torch.zeros( (original_len, *value_a.shape[1:]), dtype=value_a.dtype, device=value_a.device, ) res[slice(s0, t0)] = value_a[: t0 - s0] res[slice(s1, t1)] = value_b[: t1 - s1] return res return _handle_key("hidden_states"), _handle_key("residual") # -------------------------------- Utilities and wrappers --------------------------------------- class MaybeTboDeepEPDispatcher(BaseDispatcher): def __init__(self, **kwargs): super().__init__() num_inner_dispatchers = 2 if is_tbo_enabled() else 1 if get_moe_a2a_backend().is_deepep(): self._inners = [ DeepEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers) ] elif get_moe_a2a_backend().is_mooncake(): self._inners = [ MooncakeEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers) ] elif get_moe_a2a_backend().is_mori(): self._inners = [ MoriEPDispatcher(instance_id=i, **kwargs) for i in range(num_inner_dispatchers) ] elif get_moe_a2a_backend().is_nixl(): self._inners = [ NixlEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers) ] def _execute(self, name, tbo_subbatch_index: Optional[int] = None, **kwargs): return getattr(self._inners[tbo_subbatch_index or 0], name)(**kwargs) def dispatch(self, **kwargs) -> DispatchOutput: return self._execute("dispatch", **kwargs) def dispatch_a(self, **kwargs): return self._execute("dispatch_a", **kwargs) def dispatch_b(self, **kwargs): return self._execute("dispatch_b", **kwargs) def combine(self, **kwargs) -> torch.Tensor: return self._execute("combine", **kwargs) def combine_a(self, **kwargs): return self._execute("combine_a", **kwargs) def combine_b(self, **kwargs): return self._execute("combine_b", **kwargs) def register_deepep_dispatch_hook(self, hook): handle_list = [] for inner in self._inners: handle_list.append(inner.register_deepep_dispatch_hook(hook)) return handle_list def set_quant_config(self, quant_config: dict): super().set_quant_config(quant_config) for inner in self._inners: inner.set_quant_config(quant_config) def set_overlap_args( self, combine_overlap_args: CombineOverlapArgs, meta_overlap_args: dict ): super().set_overlap_args(combine_overlap_args, meta_overlap_args) for inner in self._inners: inner.set_overlap_args(combine_overlap_args, meta_overlap_args) def clear_overlap_args(self): super().clear_overlap_args() for inner in self._inners: inner.clear_overlap_args() ================================================ FILE: python/sglang/srt/checkpoint_engine/__init__.py ================================================ """ Checkpoint engine module for SGLang. This module provides functionality for updating model weights via checkpoint engine. """ from sglang.srt.checkpoint_engine.update import main __all__ = ["main"] ================================================ FILE: python/sglang/srt/checkpoint_engine/checkpoint_engine_worker.py ================================================ # Copyright 2023-2024 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """ Checkpoint-engine integration for SGLang. This module provides weight update functionality via IPC for checkpoint-engine compatibility. """ import logging from typing import Callable, Dict, Optional import torch import zmq try: from checkpoint_engine.worker import update_weights_from_ipc except ImportError: raise ImportError( "checkpoint-engine is not installed. " "Please install it with: pip install sglang[checkpoint-engine]" ) logger = logging.getLogger(__name__) class SGLangCheckpointEngineWorkerExtension: """ Worker extension for SGLang to support checkpoint-engine IPC weight updates. This class provides the interface needed for checkpoint-engine integration. """ def __init__(self): self._zmq_ctx: Optional[zmq.Context] = None def get_device_uuid(self) -> str: """Get the UUID of current device.""" # We need to implement this to get the device UUID # This will be overridden when integrated into SGLang's worker raise NotImplementedError( "This method should be overridden by SGLang integration" ) def get_device_id(self) -> int: """Get the device ID.""" raise NotImplementedError( "This method should be overridden by SGLang integration" ) def get_model_loader(self) -> Callable: """Get the model weight loader function.""" raise NotImplementedError( "This method should be overridden by SGLang integration" ) def get_post_hook(self) -> Optional[Callable]: """Get the post-processing hook after weight loading.""" return None def update_weights_from_ipc(self, zmq_handles: Dict[str, str]): """ Update weights from IPC communication. Args: zmq_handles: Dict mapping device UUID to ZMQ socket path """ if self._zmq_ctx is None: self._zmq_ctx = zmq.Context() device_uuid = self.get_device_uuid() device_id = self.get_device_id() if device_uuid not in zmq_handles: raise ValueError( f"Device UUID {device_uuid} not found in zmq_handles: {list(zmq_handles.keys())}" ) update_weights_from_ipc( self._zmq_ctx, zmq_handles[device_uuid], device_id=device_id, run=self.get_model_loader(), post_hook=self.get_post_hook(), ) class SGLangCheckpointEngineWorkerExtensionImpl(SGLangCheckpointEngineWorkerExtension): """ Implementation of SGLangCheckpointEngineWorkerExtension that integrates with SGLang's model runner. This class provides the concrete implementation for checkpoint-engine IPC weight updates. """ def __init__(self, model_runner): super().__init__() self.model_runner = model_runner def get_device_uuid(self) -> str: """Get the UUID of current device.""" # Get device UUID for current device device_id = torch.cuda.current_device() try: return f"GPU-{torch.cuda.get_device_properties(device_id).uuid!s}" except AssertionError as e: raise ValueError(f"Failed to get GPU UUID for device {device_id}") from e def get_device_id(self) -> int: """Get the device ID.""" return torch.cuda.current_device() def get_model_loader(self) -> Callable: """Get the model weight loader function.""" return self.model_runner.model.load_weights def get_post_hook(self) -> Optional[Callable]: """Get the post-processing hook after weight loading.""" def post_hook(): # Perform post-processing after weight loading similar to DefaultModelLoader try: from sglang.srt.model_loader.loader import device_loading_context # Process quantization methods after loading weights for _, module in self.model_runner.model.named_modules(): quant_method = getattr(module, "quant_method", None) if quant_method is not None: # Move parameters to device if needed for quantization processing target_device = torch.device( "cuda", torch.cuda.current_device() ) with device_loading_context(module, target_device): quant_method.process_weights_after_loading(module) # Call model-specific post-loading hook if available if hasattr(self.model_runner.model, "post_load_weights"): self.model_runner.model.post_load_weights() except Exception as e: logger.warning(f"Post-hook processing failed: {e}") return post_hook ================================================ FILE: python/sglang/srt/checkpoint_engine/update.py ================================================ """ Usage: 1) Launch the server with wait-for-initial-weights option in one terminal: python -m sglang.launch_server --model-path /workspace/Qwen/Qwen3-4B/ --tensor-parallel-size 2 --port 19730 --load-format dummy --checkpoint-engine-wait-weights-before-ready --mem-fraction-static 0.7 2) Torchrun this script in another terminal: torchrun --nproc-per-node 2 update.py --update-method broadcast --checkpoint-path /workspace/Qwen/Qwen3-4B/ --inference-parallel-size 2 Or use the integrated entry point: python -m sglang.srt.checkpoint_engine.update --update-method broadcast --checkpoint-path /workspace/Qwen/Qwen3-4B/ --inference-parallel-size 2 """ import argparse import json import os import pickle import subprocess import sys import time from collections import defaultdict from collections.abc import Callable from contextlib import contextmanager from typing import Literal import httpx import torch import torch.distributed as dist from safetensors import safe_open try: from checkpoint_engine.ps import ParameterServer from loguru import logger except ImportError: # Fallback for when checkpoint_engine is not available ParameterServer = None import logging logger = logging.getLogger(__name__) @contextmanager def timer(msg: str): start = time.perf_counter() yield end = time.perf_counter() logger.info(f"{msg} duration: {end - start:.2f} seconds") def check_sglang_ready( endpoint: str, inference_parallel_size: int, uds: str | None = None ): rank = int(os.getenv("RANK", 0)) if rank != rank // inference_parallel_size * inference_parallel_size: return retry_num = 0 transport = None if uds is not None: transport = httpx.HTTPTransport(uds=uds) with httpx.Client(transport=transport) as client: while True: try: response = client.get(f"{endpoint}/ping", timeout=10) response.raise_for_status() break except (httpx.ConnectError, httpx.HTTPStatusError) as e: if retry_num % 10 == 0: logger.warning( f"fail to check sglang ready, retry {retry_num} times, error: {e}" ) retry_num += 1 time.sleep(0.1) def split_checkpoint_files( checkpoint_path: str, rank: int, world_size: int ) -> list[str]: checkpoint_files = [ os.path.join(checkpoint_path, f) for f in filter( lambda x: x.endswith(".safetensors"), os.listdir(checkpoint_path) ) ] files_per_rank = (len(checkpoint_files) + world_size - 1) // world_size return checkpoint_files[rank * files_per_rank : (rank + 1) * files_per_rank] def split_tensors( checkpoint_path: str, rank: int, world_size: int ) -> dict[str, torch.Tensor]: index_fn = os.path.join(checkpoint_path, "model.safetensors.index.json") with open(index_fn) as f: weight_map: dict[str, str] = json.load(f)["weight_map"] weights_per_rank = (len(weight_map) + world_size - 1) // world_size fn_tensors: dict[str, list[str]] = defaultdict(list) weight_keys = list(weight_map.items()) for name, file in weight_keys[ rank * weights_per_rank : (rank + 1) * weights_per_rank ]: fn_tensors[file].append(name) named_tensors = {} for file, names in fn_tensors.items(): with safe_open(os.path.join(checkpoint_path, file), framework="pt") as f: for name in names: named_tensors[name] = f.get_tensor(name) return named_tensors def req_inference( endpoint: str, inference_parallel_size: int, timeout: float = 300.0, uds: str | None = None, weight_version: str | None = None, ) -> Callable[[list[tuple[str, str]]], None]: rank = int(os.getenv("RANK", 0)) src = rank // inference_parallel_size * inference_parallel_size def req_func(socket_paths: list[tuple[str, str]]): if rank == src: with httpx.Client(transport=httpx.HTTPTransport(uds=uds)) as client: resp = client.post( f"{endpoint}/update_weights_from_ipc", json={ "zmq_handles": dict( socket_paths[src : src + inference_parallel_size] ), "flush_cache": True, "weight_version": weight_version, }, timeout=timeout, ) resp.raise_for_status() return req_func def update_weights( ps, checkpoint_name: str, checkpoint_files: list[str], named_tensors: dict[str, torch.Tensor], req_func: Callable[[list[tuple[str, str]]], None], inference_parallel_size: int, endpoint: str, save_metas_file: str | None = None, update_method: Literal["broadcast", "p2p", "all"] = "broadcast", uds: str | None = None, ): ps.register_checkpoint( checkpoint_name, files=checkpoint_files, named_tensors=named_tensors ) ps.init_process_group() check_sglang_ready(endpoint, inference_parallel_size, uds) dist.barrier() with timer("Gather metas"): ps.gather_metas(checkpoint_name) if save_metas_file and int(os.getenv("RANK")) == 0: with open(save_metas_file, "wb") as f: pickle.dump(ps.get_metas(), f) if update_method == "broadcast" or update_method == "all": with timer("Update weights without setting ranks"): ps.update(checkpoint_name, req_func) if update_method == "p2p" or update_method == "all": if update_method: # sleep 2s to wait destroy process group time.sleep(2) with timer("Update weights with setting ranks"): ps.update( checkpoint_name, req_func, ranks=list(range(inference_parallel_size)) ) def join( ps: ParameterServer, checkpoint_name: str, load_metas_file: str, req_func: Callable[[list[tuple[str, str]]], None], inference_parallel_size: int, endpoint: str, uds: str | None = None, ): assert load_metas_file, "load_metas_file is required" with open(load_metas_file, "rb") as f: metas = pickle.load(f) ps.init_process_group() check_sglang_ready(endpoint, inference_parallel_size, uds) dist.barrier() with timer("Gather metas before join"): ps.gather_metas(checkpoint_name) ps.load_metas(metas) with timer( f"Update weights with setting ranks as range(0, {inference_parallel_size}) by using p2p" ): ps.update(checkpoint_name, req_func, ranks=list(range(inference_parallel_size))) def run_with_torchrun(): """Run the update script with torchrun automatically.""" # Parse inference_parallel_size from command line arguments to determine nproc-per-node inference_parallel_size = 8 # default args = sys.argv[1:] # Skip the script name # Look for --inference-parallel-size in arguments for i, arg in enumerate(args): if arg == "--inference-parallel-size" and i + 1 < len(args): try: inference_parallel_size = int(args[i + 1]) except ValueError: pass break elif arg.startswith("--inference-parallel-size="): try: inference_parallel_size = int(arg.split("=", 1)[1]) except ValueError: pass break # Build torchrun command cmd = ["torchrun", f"--nproc-per-node={inference_parallel_size}", __file__] + args print(f"Running: {' '.join(cmd)}", file=sys.stderr) # Execute torchrun with the original script try: result = subprocess.run(cmd, check=False) sys.exit(result.returncode) except FileNotFoundError: print( "Error: torchrun command not found. Please ensure PyTorch is installed.", file=sys.stderr, ) sys.exit(1) except KeyboardInterrupt: print("\nInterrupted by user", file=sys.stderr) sys.exit(130) def main(): # Check if we're running under torchrun or need to invoke it if os.getenv("RANK") is None: # Not running under torchrun, so invoke it run_with_torchrun() return # Running under torchrun, proceed with normal execution parser = argparse.ArgumentParser(description="Update weights example") parser.add_argument("--checkpoint-path", type=str, default=None) parser.add_argument("--save-metas-file", type=str, default=None) parser.add_argument("--load-metas-file", type=str, default=None) parser.add_argument("--sleep-time", type=int, default=0) parser.add_argument("--endpoint", type=str, default="http://localhost:19730") parser.add_argument("--inference-parallel-size", type=int, default=8) parser.add_argument("--checkpoint-name", type=str, default="my-checkpoint-iter-0") parser.add_argument("--update-method", type=str, default="broadcast") parser.add_argument("--uds", type=str, default=None) parser.add_argument("--weight-version", type=str, default=None) args = parser.parse_args() # Get rank and world_size from environment (set by torchrun) rank = int(os.getenv("RANK", 0)) world_size = int(os.getenv("WORLD_SIZE", 1)) req_func = req_inference( args.endpoint, args.inference_parallel_size, uds=args.uds, weight_version=args.weight_version, ) if ParameterServer is None: print("Error: checkpoint_engine package not available", file=sys.stderr) sys.exit(1) ps = ParameterServer(auto_pg=True) ps._p2p_store = None if args.load_metas_file: join( ps, args.checkpoint_name, args.load_metas_file, req_func, args.inference_parallel_size, args.endpoint, args.uds, ) else: if args.checkpoint_path and os.path.exists( os.path.join(args.checkpoint_path, "model.safetensors.index.json") ): named_tensors = split_tensors(args.checkpoint_path, rank, world_size) checkpoint_files = [] else: checkpoint_files = ( split_checkpoint_files(args.checkpoint_path, rank, world_size) if args.checkpoint_path else [] ) named_tensors = {} update_weights( ps, args.checkpoint_name, checkpoint_files, named_tensors, req_func, args.inference_parallel_size, args.endpoint, args.save_metas_file, args.update_method, args.uds, ) time.sleep(args.sleep_time) if __name__ == "__main__": main() ================================================ FILE: python/sglang/srt/compilation/backend.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/backend.py import ast import dataclasses import logging import os import pprint import time from collections.abc import Sequence from contextlib import contextmanager from typing import Any, Callable, Optional import torch import torch.fx as fx from torch._dispatch.python import enable_python_dispatcher from sglang.srt.compilation.compilation_config import CompilationConfig from sglang.srt.compilation.compilation_counter import compilation_counter from sglang.srt.compilation.compiler_interface import EagerAdapter, InductorAdaptor from sglang.srt.compilation.cuda_piecewise_backend import CUDAPiecewiseBackend from sglang.srt.compilation.npu_piecewise_backend import NPUPiecewiseBackend from sglang.srt.compilation.pass_manager import PostGradPassManager from sglang.srt.utils.common import is_npu logger = logging.getLogger(__name__) def make_compiler(config: CompilationConfig): if config.compiler == "eager": return EagerAdapter() elif config.compiler == "inductor": return InductorAdaptor() else: raise ValueError(f"Unknown compiler: {config.compiler}") def make_backend( graph: fx.GraphModule, compile_config: CompilationConfig, inductor_config: dict[str, Any], graph_pool: Any, piecewise_compile_index: int, total_piecewise_compiles: int, sym_shape_indices: list[int], compiled_graph_for_general_shape: Callable, sglang_backend, ): backend_cls = CUDAPiecewiseBackend if not is_npu() else NPUPiecewiseBackend return backend_cls( graph, compile_config, inductor_config, graph_pool, piecewise_compile_index, total_piecewise_compiles, sym_shape_indices, compiled_graph_for_general_shape, sglang_backend, ) class CompilerManager: def __init__( self, config: CompilationConfig, ): self.cache = dict() self.is_cache_updated = False self.compiler = make_compiler(config) def compute_hash(self): return self.compiler.compute_hash() def initialize_cache( self, cache_dir: str, disable_cache: bool = False, prefix: str = "" ): self.disable_cache = disable_cache self.cache_dir = cache_dir self.cache_file_path = os.path.join(cache_dir, "sglang_compile_cache.py") if not disable_cache and os.path.exists(self.cache_file_path): with open(self.cache_file_path) as f: self.cache = ast.literal_eval(f.read()) self.compiler.initialize_cache( cache_dir=cache_dir, disable_cache=disable_cache, prefix=prefix ) def save_to_file(self): if self.disable_cache or not self.is_cache_updated: return printer = pprint.PrettyPrinter(indent=4) data = printer.pformat(self.cache) with open(self.cache_file_path, "w") as f: f.write(data) def load( self, graph: fx.GraphModule, example_inputs: list[Any], graph_index: int, runtime_shape: Optional[int] = None, ) -> Optional[Callable]: handle = self.cache[(runtime_shape, graph_index, self.compiler.name)] compiled_graph = self.compiler.load( handle, graph, example_inputs, graph_index, runtime_shape ) if runtime_shape is None: logger.debug( "Directly load the %s-th graph for dynamic shape from %s via " "handle %s", graph_index, self.compiler.name, handle, ) else: logger.debug( "Directly load the %s-th graph for shape %s from %s via " "handle %s", graph_index, str(runtime_shape), self.compiler.name, handle, ) return compiled_graph def compile( self, graph: fx.GraphModule, example_inputs, inductor_config: dict[str, Any], graph_index: int = 0, num_graphs: int = 1, runtime_shape: Optional[int] = None, ) -> Any: if graph_index == 0: # before compiling the first graph, record the start time global compilation_start_time compilation_start_time = time.time() compilation_counter.num_backend_compilations += 1 compiled_graph = None # TODO(Yuwei): support cache loading # no compiler cached the graph, or the cache is disabled, # we need to compile it if isinstance(self.compiler, InductorAdaptor): maybe_key = None else: maybe_key = f"artifact_shape_{runtime_shape}_subgraph_{graph_index}" compiled_graph, handle = self.compiler.compile( graph, example_inputs, inductor_config, runtime_shape, maybe_key ) assert compiled_graph is not None, "Failed to compile the graph" # store the artifact in the cache if handle is not None: self.cache[(runtime_shape, graph_index, self.compiler.name)] = handle compilation_counter.num_cache_entries_updated += 1 self.is_cache_updated = True if graph_index == 0: # adds some info logging for the first graph if runtime_shape is None: logger.info("Cache the graph for dynamic shape for later use") else: logger.info( "Cache the graph of shape %s for later use", str(runtime_shape) ) if runtime_shape is None: logger.debug( "Store the %s-th graph for dynamic shape from %s via " "handle %s", graph_index, self.compiler.name, handle, ) else: logger.debug( "Store the %s-th graph for shape %s from %s via handle %s", graph_index, str(runtime_shape), self.compiler.name, handle, ) # after compiling the last graph, record the end time if graph_index == num_graphs - 1: now = time.time() elapsed = now - compilation_start_time if runtime_shape is None: logger.info("Compiling a graph for dynamic shape takes %.2f s", elapsed) else: logger.info( "Compiling a graph for shape %s takes %.2f s", runtime_shape, elapsed, ) return compiled_graph @dataclasses.dataclass class SplitItem: submod_name: str graph_id: int is_splitting_graph: bool graph: fx.GraphModule def split_graph( graph: fx.GraphModule, ops: list[str] ) -> tuple[fx.GraphModule, list[SplitItem]]: # split graph by ops subgraph_id = 0 node_to_subgraph_id = {} split_op_graphs = [] for node in graph.graph.nodes: if node.op in ("output", "placeholder"): continue if node.op == "call_function" and str(node.target) in ops: subgraph_id += 1 node_to_subgraph_id[node] = subgraph_id split_op_graphs.append(subgraph_id) subgraph_id += 1 else: node_to_subgraph_id[node] = subgraph_id # `keep_original_order` is important! # otherwise pytorch might reorder the nodes and # the semantics of the graph will change when we # have mutations in the graph split_gm = torch.fx.passes.split_module.split_module( graph, None, lambda node: node_to_subgraph_id[node], keep_original_order=True ) outputs = [] names = [name for (name, module) in split_gm.named_modules()] for name in names: if "." in name or name == "": # recursive child module or the root module continue module = getattr(split_gm, name) graph_id = int(name.replace("submod_", "")) outputs.append(SplitItem(name, graph_id, (graph_id in split_op_graphs), module)) # sort by intetger graph_id, rather than string name outputs.sort(key=lambda x: x.graph_id) return split_gm, outputs # we share the global graph pool among all the backends global_graph_pool = None compilation_start_time = 0.0 class PiecewiseCompileInterpreter(torch.fx.Interpreter): def __init__( self, module: torch.fx.GraphModule, compile_submod_names: list[str], inductor_config: dict[str, Any], graph_pool, compile_config: CompilationConfig, sglang_backend: "SGLangBackend", ): super().__init__(module) from torch._guards import detect_fake_mode self.fake_mode = detect_fake_mode() self.compile_submod_names = compile_submod_names self.graph_pool = graph_pool self.sglang_backend = sglang_backend # When True, it annoyingly dumps the torch.fx.Graph on errors. self.extra_traceback = False self.inductor_config = inductor_config self.compile_config = compile_config def run(self, *args): fake_args = [ self.fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t for t in args ] with self.fake_mode, enable_python_dispatcher(): return super().run(*fake_args) def call_module( self, target: torch.fx.node.Target, args: tuple[torch.fx.node.Argument, ...], kwargs: dict[str, Any], ) -> Any: assert isinstance(target, str) output = super().call_module(target, args, kwargs) if target in self.compile_submod_names: index = self.compile_submod_names.index(target) submod = self.fetch_attr(target) sym_shape_indices = [ i for i, x in enumerate(args) if isinstance(x, torch.SymInt) ] global compilation_start_time compiled_graph_for_dynamic_shape = ( self.sglang_backend.compiler_manager.compile( submod, args, self.inductor_config, graph_index=index, num_graphs=len(self.compile_submod_names), runtime_shape=None, ) ) self.module.__dict__[target] = make_backend( submod, self.compile_config, self.inductor_config, self.graph_pool, index, len(self.compile_submod_names), sym_shape_indices, compiled_graph_for_dynamic_shape, self.sglang_backend, ) compilation_counter.num_piecewise_capturable_graphs_seen += 1 return output model_tag: str = "backbone" @contextmanager def set_model_tag(tag: str): """Context manager to set the model tag.""" global model_tag assert ( tag != model_tag ), f"Model tag {tag} is the same as the current tag {model_tag}." old_tag = model_tag model_tag = tag try: yield finally: model_tag = old_tag class SGLangBackend: graph_pool: Any _called: bool = False # the graph we compiled graph: fx.GraphModule # the stiching graph module for all the piecewise graphs split_gm: fx.GraphModule piecewise_graphs: list[SplitItem] returned_callable: Callable # Inductor passes to run on the graph pre-defunctionalization post_grad_passes: Sequence[Callable] sym_tensor_indices: list[int] input_buffers: list[torch.Tensor] compiler_manager: CompilerManager def __init__( self, config: CompilationConfig, graph_pool: Any, ): assert graph_pool is not None self.graph_pool = graph_pool self.post_grad_pass_manager = PostGradPassManager() self.sym_tensor_indices = [] self.input_buffers = [] self.compiler_manager = CompilerManager(config) self.inductor_config = { "enable_auto_functionalized_v2": False, } self.compile_config = config def configure_post_pass(self): self.post_grad_pass_manager.configure() self.inductor_config["post_grad_custom_post_pass"] = self.post_grad_pass_manager def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable: base_cache_dir = os.path.expanduser( os.getenv("SGLANG_CACHE_DIR", "~/.cache/sglang/") ) cache_hash = self.compiler_manager.compute_hash() cache_dir = os.path.join( base_cache_dir, "torch_compile_cache", cache_hash, ) os.makedirs(cache_dir, exist_ok=True) rank = 0 dp_rank = 0 local_cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}", model_tag) os.makedirs(local_cache_dir, exist_ok=True) self.compiler_manager.initialize_cache( local_cache_dir, disable_cache=False, prefix="" ) compilation_counter.num_graphs_seen += 1 assert not self._called, "SGLangBackend can only be called once" self.graph = graph self.configure_post_pass() self.split_gm, self.piecewise_graphs = split_graph( graph, self.compile_config.split_ops, ) from torch._dynamo.utils import lazy_format_graph_code # depyf will hook lazy_format_graph_code and dump the graph # for debugging, no need to print the graph here lazy_format_graph_code("before split", self.graph) lazy_format_graph_code("after split", self.split_gm) compilation_counter.num_piecewise_graphs_seen += len(self.piecewise_graphs) submod_names_to_compile = [ item.submod_name for item in self.piecewise_graphs if not item.is_splitting_graph ] PiecewiseCompileInterpreter( self.split_gm, submod_names_to_compile, self.inductor_config, self.graph_pool, self.compile_config, self, ).run(*example_inputs) rank = torch.distributed.get_rank() if rank == 0: graph_path = os.path.join( local_cache_dir, f"computation_graph_{time.time()}.py" ) if not os.path.exists(graph_path): # code adapted from https://github.com/thuml/depyf/blob/dab831108a752d1facc00acdd6d4243891845c37/depyf/explain/patched_lazy_format_graph_code.py#L30 # noqa # use `print_readable` because it can include submodules src = ( "from __future__ import annotations\nimport torch\n" + self.split_gm.print_readable(print_output=False) ) src = src.replace("", "GraphModule") with open(graph_path, "w") as f: f.write(src) self._called = True return self.split_gm ================================================ FILE: python/sglang/srt/compilation/compilation_config.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/compilation_config.py from typing import Callable, List, Optional SPLIT_OPS = [] def register_split_op(op_name: Optional[str] = None): def decorator(op_func: Callable): name = op_name or op_func.__name__ SPLIT_OPS.append(f"sglang.{name}") return op_func return decorator # TODO(Yuwei): support better compile config support class CompilationConfig: def __init__( self, capture_sizes: List[int], compiler: str = "eager", enable_debug_mode: bool = False, ): self.traced_files = set() self.capture_sizes = capture_sizes self.compiler = compiler self.enable_debug_mode = enable_debug_mode self.split_ops = [] self.split_ops.extend(SPLIT_OPS) def add_split_op(self, op: str): self.split_ops.append(op) def add_traced_file(self, file_path: str): self.traced_files.add(file_path) def get_traced_files(self): return self.traced_files def get_capture_sizes(self): return self.capture_sizes def get_enable_debug_mode(self): return self.enable_debug_mode ================================================ FILE: python/sglang/srt/compilation/compilation_counter.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/compilation_counter.py import copy import dataclasses from contextlib import contextmanager @dataclasses.dataclass class CompilationCounter: num_models_seen: int = 0 num_graphs_seen: int = 0 # including the splitting ops num_piecewise_graphs_seen: int = 0 # not including the splitting ops num_piecewise_capturable_graphs_seen: int = 0 num_backend_compilations: int = 0 # Number of gpu_model_runner attempts to trigger CUDAGraphs capture num_gpu_runner_capture_triggers: int = 0 # Number of CUDAGraphs captured num_cudagraph_captured: int = 0 # InductorAdapter.compile calls num_inductor_compiles: int = 0 # EagerAdapter.compile calls num_eager_compiles: int = 0 # The number of time vLLM's compiler cache entry was updated num_cache_entries_updated: int = 0 # The number of standalone_compile compiled artifacts saved num_compiled_artifacts_saved: int = 0 # Number of times a model was loaded with CompilationLevel.DYNAMO_AS_IS dynamo_as_is_count: int = 0 def clone(self) -> "CompilationCounter": return copy.deepcopy(self) @contextmanager def expect(self, **kwargs): old = self.clone() yield for k, v in kwargs.items(): assert getattr(self, k) - getattr(old, k) == v, ( f"{k} not as expected, before it is {getattr(old, k)}" f", after it is {getattr(self, k)}, " f"expected diff is {v}" ) compilation_counter = CompilationCounter() ================================================ FILE: python/sglang/srt/compilation/compile.py ================================================ import inspect import logging import os import sys import types from dataclasses import dataclass from typing import Any, Callable, Optional, Union import torch from sglang.srt.compilation.compilation_config import CompilationConfig from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph logger = logging.getLogger(__name__) @dataclass class IntermediateTensors: """For all pipeline stages except the last, we need to return the hidden states and residuals to be sent to the next stage. This data structure contains the hidden states and residuals for a request. Each stage also needs to handle its own finished_sending and finished_recving in case of kv transfer. """ tensors: dict[str, torch.Tensor] # [req_ids] finished_sending: Optional[set[str]] = None finished_recving: Optional[set[str]] = None def __init__(self, tensors): # manually define this function, so that # Dynamo knows `IntermediateTensors()` comes from this file. # Otherwise, dataclass will generate this function by evaluating # a string, and we will lose the information about the source file. self.tensors = tensors def __getitem__(self, key: Union[str, slice]): if isinstance(key, str): return self.tensors[key] elif isinstance(key, slice): return self.__class__({k: v[key] for k, v in self.tensors.items()}) def __setitem__(self, key: str, value: torch.Tensor): self.tensors[key] = value def items(self): return self.tensors.items() def __len__(self): return len(self.tensors) def __eq__(self, other: object): return isinstance(other, self.__class__) and self def __repr__(self) -> str: return f"IntermediateTensors(tensors={self.tensors})" def _normalize_dims(dims, ndim: int): dims = [dims] if isinstance(dims, int) else list(dims) return [d if d >= 0 else ndim + d for d in dims] class _MaybeIntermediateTensors: """Duck-typed check to support your IntermediateTensors without importing.""" def __init__(self, obj): self.is_intermediate = hasattr(obj, "tensors") and isinstance( getattr(obj, "tensors"), dict ) self.obj = obj def _mark_dynamic_on_value(val, dims): if isinstance(val, torch.Tensor): torch._dynamo.maybe_mark_dynamic(val, _normalize_dims(dims, val.ndim)) else: mit = _MaybeIntermediateTensors(val) if mit.is_intermediate: for t in mit.obj.tensors.values(): torch._dynamo.maybe_mark_dynamic(t, _normalize_dims(dims, t.ndim)) # else: ignore (None or non-tensor) def _infer_dynamic_arg_dims_from_annotations(forward_fn): sig = inspect.signature(forward_fn) dyn = {} for name, p in sig.parameters.items(): ann = p.annotation # Accept torch.Tensor / Optional[torch.Tensor] / your IntermediateTensors types by name if ( ann is torch.Tensor or getattr(getattr(ann, "__args__", [None])[0], "__name__", "") == "Tensor" ): dyn[name] = 0 elif getattr(ann, "__name__", "") in ("IntermediateTensors",) or any( getattr(a, "__name__", "") == "IntermediateTensors" for a in getattr(ann, "__args__", []) ): dyn[name] = 0 elif ann == "torch.Tensor" or ann == "Optional[torch.Tensor]": # For future import annotations (e.g. from __future__ import annotations), the annotation is a string dyn[name] = 0 if not dyn: raise ValueError("No dynamic dims inferred; pass dynamic_arg_dims explicitly.") return dyn def install_torch_compiled( module: torch.nn.Module, *, dynamic_arg_dims: dict[str, Union[int, list[int]]] | None = None, backend_factory: Optional[Callable[[torch.fx.GraphModule, list], Callable]] = None, compile_config: CompilationConfig = None, fullgraph: bool = True, graph_pool: Any = None, ): unbound_fwd = module.__class__.forward if not callable(unbound_fwd): raise TypeError("module.__class__.forward must be callable") original_code = unbound_fwd.__code__ dyn_map = dynamic_arg_dims or _infer_dynamic_arg_dims_from_annotations(unbound_fwd) if backend_factory is None: from sglang.srt.compilation.backend import SGLangBackend backend_factory = lambda gm, ex: SGLangBackend(compile_config, graph_pool)( gm, ex ) compiled_codes: list[type(original_code)] = [] state = {"compiled": False, "compiled_callable": None} def bytecode_hook(old_code, new_code): if old_code is not original_code: return frame = sys._getframe() while frame and frame.f_back: frame = frame.f_back if ( frame.f_code.co_name == "_compile" and os.path.basename(frame.f_code.co_filename) == "convert_frame.py" ): break try: dynamo_frame = frame.f_locals["frame"] except Exception: return if dynamo_frame.f_code is not old_code: return if dynamo_frame.f_locals.get("self") is not module: return compiled_codes.append(new_code) torch._dynamo.convert_frame.register_bytecode_hook(bytecode_hook) def _ensure_compiled(self, *args, **kwargs): """Compile on first use (with flag ON).""" if state["compiled"]: return # Mark dynamic dims only when we are about to compile sig = inspect.signature(unbound_fwd) ba = sig.bind(self, *args, **kwargs) ba.apply_defaults() for name, dims in (dyn_map or {}).items(): if name in ba.arguments: val = ba.arguments[name] if val is not None: _mark_dynamic_on_value(val, dims) # Avoid cross-instance cache reuse torch._dynamo.eval_frame.remove_from_cache(unbound_fwd.__code__) bound = types.MethodType(unbound_fwd, self) compiled_callable = torch.compile( bound, fullgraph=fullgraph, backend=backend_factory ) # Trigger Dynamo so bytecode hook can capture compiled_callable(*args, **kwargs) state["compiled"] = True state["compiled_callable"] = compiled_callable def trampoline(self, *args, **kwargs): use_compiled = is_in_piecewise_cuda_graph() if use_compiled: if not state["compiled"]: _ensure_compiled(self, *args, **kwargs) compiled_callable = state["compiled_callable"] return compiled_callable(*args, **kwargs) else: # Explicitly run the original uncompiled forward return unbound_fwd(self, *args, **kwargs) module.forward = types.MethodType(trampoline, module) return module ================================================ FILE: python/sglang/srt/compilation/compiler_interface.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/compiler_interface.py import contextlib import copy import hashlib import os from contextlib import ExitStack from typing import Any, Callable, Optional from unittest.mock import patch import torch import torch._inductor.compile_fx import torch.fx as fx from sglang.srt.compilation.compilation_counter import compilation_counter from sglang.srt.compilation.inductor_pass import pass_context from sglang.srt.utils.common import torch_release class CompilerInterface: """ The interface for a compiler that can be used by vLLM. """ # The name of the compiler, e.g. inductor. # This is a class-level attribute. name: str def initialize_cache( self, cache_dir: str, disable_cache: bool = False, prefix: str = "" ): """ when the vLLM process uses `cache_dir` as the cache directory, the compiler should initialize itself with the cache directory, e.g. by re-directing its own cache directory to a sub-directory. prefix can be used in combination with cache_dir to figure out the base cache directory, e.g. there're multiple parts of model being compiled, but we want to share the same cache directory for all of them. e.g. cache_dir = "/path/to/dir/backbone", prefix = "backbone" cache_dir = "/path/to/dir/eagle_head", prefix = "eagle_head" """ pass def compute_hash(self) -> str: """ Gather all the relevant information from the vLLM config, to compute a hash so that we can cache the compiled model. See [`VllmConfig.compute_hash`][vllm.config.VllmConfig.compute_hash] to check what information is already considered by default. This function should only consider the information that is specific to the compiler. """ return "" def compile( self, graph: fx.GraphModule, example_inputs: list[Any], compiler_config: dict[str, Any], runtime_shape: Optional[int] = None, key: Optional[str] = None, ) -> tuple[Optional[Callable], Optional[Any]]: """ Compile the graph with the given example inputs and compiler config, with a runtime shape. If the `runtime_shape` is None, it means the `example_inputs` have a dynamic shape. Otherwise, the `runtime_shape` specifies the shape of the inputs. Right now we only support one variable shape for all inputs, which is the batchsize (number of tokens) during inference. Dynamo will make sure `graph(*example_inputs)` is valid. The function should return a compiled callable function, as well as a handle that can be used to directly load the compiled function. The handle should be a plain Python object, preferably a string or a file path for readability. If the compiler doesn't support caching, it should return None for the handle. If the compiler fails to compile the graph, it should return None for the compiled function as well. `key` is required for StandaloneInductorAdapter, it specifies where to save the compiled artifact. The compiled artifact gets saved to `cache_dir/key`. """ return None, None def load( self, handle: Any, graph: fx.GraphModule, example_inputs: list[Any], graph_index: int, runtime_shape: Optional[int] = None, ) -> Callable: """ Load the compiled function from the handle. Raises an error if the handle is invalid. The handle is the second return value of the `compile` function. """ raise NotImplementedError("caching is not supported") def get_inductor_factors() -> list[Any]: factors: list[Any] = [] # summarize system state from torch._inductor.codecache import CacheBase system_factors = CacheBase.get_system() factors.append(system_factors) # summarize pytorch state from torch._inductor.codecache import torch_key torch_factors = torch_key() factors.append(torch_factors) return factors class AlwaysHitShapeEnv: """ Why do we need this class: For normal `torch.compile` usage, every compilation will have one Dynamo bytecode compilation and one Inductor compilation. The Inductor compilation happens under the context of the Dynamo bytecode compilation, and that context is used to determine the dynamic shape information, etc. For our use case, we only run Dynamo bytecode compilation once, and run Inductor compilation multiple times with different shapes plus a general shape. The compilation for specific shapes happens outside of the context of the Dynamo bytecode compilation. At that time, we don't have shape environment to provide to Inductor, and it will fail the Inductor code cache lookup. By providing a dummy shape environment that always hits, we can make the Inductor code cache lookup always hit, and we can compile the graph for different shapes as needed. The following dummy methods are obtained by trial-and-error until it works. """ def __init__(self) -> None: self.guards: list[Any] = [] def evaluate_guards_expression(self, *args, **kwargs): return True def get_pruned_guards(self, *args, **kwargs): return [] def produce_guards_expression(self, *args, **kwargs): return "" class InductorAdaptor(CompilerInterface): """ The adaptor for the Inductor compiler, version 2.5, 2.6, 2.7. """ name = "inductor" def compute_hash(self) -> str: factors = get_inductor_factors() hash_str = hashlib.md5( str(factors).encode(), usedforsecurity=False ).hexdigest()[:10] return hash_str def initialize_cache( self, cache_dir: str, disable_cache: bool = False, prefix: str = "" ): self.cache_dir = cache_dir self.prefix = prefix self.base_cache_dir = cache_dir[: -len(prefix)] if prefix else cache_dir if disable_cache: return # redirect the cache directory to a sub-directory # set flags so that Inductor and Triton store their cache # in the cache_dir, then users only need to copy the cache_dir # to another machine to reuse the cache. inductor_cache = os.path.join(self.base_cache_dir, "inductor_cache") os.makedirs(inductor_cache, exist_ok=True) os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache triton_cache = os.path.join(self.base_cache_dir, "triton_cache") os.makedirs(triton_cache, exist_ok=True) os.environ["TRITON_CACHE_DIR"] = triton_cache def compile( self, graph: fx.GraphModule, example_inputs: list[Any], compiler_config: dict[str, Any], runtime_shape: Optional[int] = None, key: Optional[str] = None, ) -> tuple[Optional[Callable], Optional[Any]]: compilation_counter.num_inductor_compiles += 1 from torch._inductor.compile_fx import compile_fx current_config = {} if compiler_config is not None: current_config.update(compiler_config) # disable remote cache current_config["fx_graph_cache"] = True current_config["fx_graph_remote_cache"] = False set_inductor_config(current_config, runtime_shape) # inductor can inplace modify the graph, so we need to copy it # see https://github.com/pytorch/pytorch/issues/138980 graph = copy.deepcopy(graph) # it's the first time we compile this graph # the assumption is that we don't have nested Inductor compilation. # compiled_fx_graph_hash will only be called once, and we can hook # it to get the hash of the compiled graph directly. hash_str, file_path = None, None from torch._inductor.codecache import FxGraphCache, compiled_fx_graph_hash if torch_release[:2] == (2, 5): original_load = FxGraphCache.load original_load_name = "torch._inductor.codecache.FxGraphCache.load" def hijack_load(*args, **kwargs): inductor_compiled_graph = original_load(*args, **kwargs) nonlocal file_path compiled_fn = inductor_compiled_graph.current_callable file_path = compiled_fn.__code__.co_filename # noqa if not file_path.startswith(self.base_cache_dir): # hooked in the align_inputs_from_check_idxs function # in torch/_inductor/utils.py for cell in compiled_fn.__closure__: if not callable(cell.cell_contents): continue if cell.cell_contents.__code__.co_filename.startswith( self.base_cache_dir ): # this is the real file path compiled from Inductor file_path = cell.cell_contents.__code__.co_filename break return inductor_compiled_graph hijacked_compile_fx_inner = ( torch._inductor.compile_fx.compile_fx_inner ) # noqa elif torch_release >= (2, 6): # function renamed in 2.6 original_load_name = None def hijacked_compile_fx_inner(*args, **kwargs): output = torch._inductor.compile_fx.compile_fx_inner(*args, **kwargs) nonlocal hash_str inductor_compiled_graph = output if inductor_compiled_graph is not None: nonlocal file_path compiled_fn = inductor_compiled_graph.current_callable file_path = compiled_fn.__code__.co_filename # noqa if not file_path.startswith(self.base_cache_dir): # hooked in the align_inputs_from_check_idxs function # in torch/_inductor/utils.py for cell in compiled_fn.__closure__: if not callable(cell.cell_contents): continue code = cell.cell_contents.__code__ if code.co_filename.startswith(self.base_cache_dir): # this is the real file path # compiled from Inductor file_path = code.co_filename break hash_str = inductor_compiled_graph._fx_graph_cache_key return output def hijack_compiled_fx_graph_hash(*args, **kwargs): out = compiled_fx_graph_hash(*args, **kwargs) nonlocal hash_str hash_str = out[0] return out def _check_can_cache(*args, **kwargs): # no error means it can be cached. # Inductor refuses to cache the graph outside of Dynamo # tracing context, and also disables caching for graphs # with high-order ops. # For vLLM, in either case, we want to cache the graph. # see https://github.com/pytorch/pytorch/blob/9f5ebf3fc609105a74eab4ccc24932d6353ff566/torch/_inductor/codecache.py#L1221 # noqa return def _get_shape_env() -> AlwaysHitShapeEnv: return AlwaysHitShapeEnv() with ExitStack() as stack: # hijack to get the compiled graph itself if original_load_name is not None: stack.enter_context(patch(original_load_name, hijack_load)) # for hijacking the hash of the compiled graph stack.enter_context( patch( "torch._inductor.codecache.compiled_fx_graph_hash", hijack_compiled_fx_graph_hash, ) ) # for providing a dummy shape environment stack.enter_context( patch( "torch._inductor.codecache.FxGraphCache._get_shape_env", _get_shape_env, ) ) from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCache # torch 2.8+ on main uses _get_shape_env in AOTAutogradCache if hasattr(AOTAutogradCache, "_get_shape_env"): stack.enter_context( patch( "torch._functorch._aot_autograd.autograd_cache.AOTAutogradCache._get_shape_env", _get_shape_env, ) ) # for forcing the graph to be cached stack.enter_context( patch( "torch._inductor.codecache.FxGraphCache._check_can_cache", _check_can_cache, ) ) # Dynamo metrics context, see method for more details. stack.enter_context(self.metrics_context()) # Disable remote caching. When these are on, on remote cache-hit, # the monkey-patched functions never actually get called. # vLLM today assumes and requires the monkey-patched functions to # get hit. # TODO(zou3519): we're going to replace this all with # standalone_compile sometime. stack.enter_context( torch._inductor.config.patch(fx_graph_remote_cache=False) ) # InductorAdaptor (unfortunately) requires AOTAutogradCache # to be turned off to run. It will fail to acquire the hash_str # and error if not. # StandaloneInductorAdaptor (PyTorch 2.8+) fixes this problem. stack.enter_context( torch._functorch.config.patch(enable_autograd_cache=False) ) stack.enter_context( torch._functorch.config.patch(enable_remote_autograd_cache=False) ) with pass_context(runtime_shape): compiled_graph = compile_fx( graph, example_inputs, inner_compile=hijacked_compile_fx_inner, config_patches=current_config, ) return compiled_graph, (hash_str, file_path) def load( self, handle: Any, graph: fx.GraphModule, example_inputs: list[Any], graph_index: int, runtime_shape: Optional[int] = None, ) -> Callable: assert isinstance(handle, tuple) assert isinstance(handle[0], str) assert isinstance(handle[1], str) hash_str = handle[0] from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCache from torch._inductor.codecache import FxGraphCache with ExitStack() as exit_stack: exit_stack.enter_context( patch( "torch._inductor.codecache.FxGraphCache._get_shape_env", lambda *args, **kwargs: AlwaysHitShapeEnv(), ) ) # torch 2.8+ on main uses _get_shape_env in AOTAutogradCache if hasattr(AOTAutogradCache, "_get_shape_env"): exit_stack.enter_context( patch( "torch._functorch._aot_autograd.autograd_cache.AOTAutogradCache._get_shape_env", lambda *args, **kwargs: AlwaysHitShapeEnv(), ) ) # Dynamo metrics context, see method for more details. exit_stack.enter_context(self.metrics_context()) if torch_release[:2] == (2, 5): inductor_compiled_graph = FxGraphCache._lookup_graph( hash_str, example_inputs, True, False ) assert inductor_compiled_graph is not None, ( "Inductor cache lookup failed. Please remove" f"the cache directory and try again." # noqa ) elif torch_release >= (2, 6): from torch._inductor.output_code import CompiledFxGraphConstantsWithGm constants = CompiledFxGraphConstantsWithGm(graph) inductor_compiled_graph, _ = FxGraphCache._lookup_graph( hash_str, example_inputs, True, None, constants ) assert inductor_compiled_graph is not None, ( "Inductor cache lookup failed. Please remove" f"the cache directory and try again." # noqa ) # Inductor calling convention (function signature): # f(list) -> tuple # Dynamo calling convention (function signature): # f(*args) -> Any # need to know if the graph returns a tuple from torch._inductor.compile_fx import graph_returns_tuple returns_tuple = graph_returns_tuple(graph) # this is the callable we return to Dynamo to run def compiled_graph(*args): # convert args to list list_args = list(args) graph_output = inductor_compiled_graph(list_args) # unpack the tuple if needed if returns_tuple: return graph_output else: return graph_output[0] return compiled_graph def metrics_context(self) -> contextlib.AbstractContextManager: """ This method returns the Dynamo metrics context (if it exists, otherwise a null context). It is used by various compile components. Present in torch>=2.6, it's used inside FxGraphCache in torch==2.6 (but not after). It might also be used in various other torch.compile internal functions. Because it is re-entrant, we always set it (even if entering via Dynamo and the context was already entered). We might want to revisit if it should be set at a different level of compilation. This is likely a bug in PyTorch: public APIs should not rely on manually setting up internal contexts. But we also rely on non-public APIs which might not provide these guarantees. """ import torch._dynamo.utils return torch._dynamo.utils.get_metrics_context() def set_inductor_config(config, runtime_shape): if isinstance(runtime_shape, int): # for a specific batchsize, tuning triton kernel parameters # can be beneficial config["max_autotune"] = True config["coordinate_descent_tuning"] = True class EagerAdapter(CompilerInterface): name = "eager" def compile( self, graph: fx.GraphModule, example_inputs: list[Any], compiler_config: dict[str, Any], runtime_shape: Optional[int] = None, key: Optional[str] = None, num_graphs: int = 1, ) -> tuple[Optional[Callable], Optional[Any]]: return graph, None def load( self, handle: Any, graph: fx.GraphModule, example_inputs: list[Any], graph_index: int, runtime_shape: Optional[int] = None, num_graphs: int = 1, ) -> Callable: raise NotImplementedError("eager compilation is not supported") ================================================ FILE: python/sglang/srt/compilation/cuda_piecewise_backend.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/cuda_piecewise_backend.py import dataclasses import logging from contextlib import ExitStack from typing import Any, Callable, Optional from unittest.mock import patch import torch import torch.fx as fx from sglang.srt.compilation.compilation_config import CompilationConfig from sglang.srt.compilation.compilation_counter import compilation_counter from sglang.srt.compilation.piecewise_context_manager import ( get_pcg_capture_stream, is_in_pcg_torch_compile, ) from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors logger = logging.getLogger(__name__) @dataclasses.dataclass class ConcreteSizeEntry: runtime_shape: int need_to_compile: bool # the size is in compile_sizes use_cudagraph: bool # the size is in cudagraph_capture_sizes compiled: bool = False runnable: Callable = None # type: ignore num_finished_warmup: int = 0 cudagraph: Optional[torch.cuda.CUDAGraph] = None output: Optional[Any] = None # for cudagraph debugging, track the input addresses # during capture, and check if they are the same during replay input_addresses: Optional[list[int]] = None class CUDAPiecewiseBackend: def __init__( self, graph: fx.GraphModule, compile_config: CompilationConfig, inductor_config: dict[str, Any], graph_pool: Any, piecewise_compile_index: int, total_piecewise_compiles: int, sym_shape_indices: list[int], compiled_graph_for_general_shape: Callable, sglang_backend, ): """ The backend for piecewise compilation. It mainly handles the compilation and cudagraph capturing. We will compile `self.graph` once for the general shape, and then compile for different shapes specified in `compilation_config.compile_sizes`. Independently, we will capture cudagraph for different shapes. If a shape needs both compilation and cudagraph, we will compile it first, and then capture cudagraph. """ self.graph = graph self.inductor_config = inductor_config self.graph_pool = graph_pool self.piecewise_compile_index = piecewise_compile_index self.total_piecewise_compiles = total_piecewise_compiles self.sglang_backend = sglang_backend self.is_first_graph = piecewise_compile_index == 0 self.is_last_graph = piecewise_compile_index == total_piecewise_compiles - 1 self.compile_sizes: set[int] = set([]) self.compile_config = compile_config self.cudagraph_capture_sizes: set[int] = set(compile_config.get_capture_sizes()) self.first_run_finished = False self.compiled_graph_for_general_shape = compiled_graph_for_general_shape # noqa self.sym_shape_indices = sym_shape_indices # the entries for different shapes that we need to either # compile or capture cudagraph self.concrete_size_entries: dict[int, ConcreteSizeEntry] = {} # to_be_compiled_sizes tracks the remaining sizes to compile, # and updates during the compilation process, so we need to copy it self.to_be_compiled_sizes: set[int] = self.compile_sizes.copy() for shape in self.compile_sizes.union(self.cudagraph_capture_sizes): self.concrete_size_entries[shape] = ConcreteSizeEntry( runtime_shape=shape, need_to_compile=shape in self.compile_sizes, use_cudagraph=shape in self.cudagraph_capture_sizes, ) def check_for_ending_compilation(self): if self.is_last_graph and not self.to_be_compiled_sizes: # no specific sizes to compile # save the hash of the inductor graph for the next run self.sglang_backend.compiler_manager.save_to_file() def __call__(self, *args) -> Any: if not self.first_run_finished: self.first_run_finished = True self.check_for_ending_compilation() return self.compiled_graph_for_general_shape(*args) if len(self.sym_shape_indices) == 0: return self.compiled_graph_for_general_shape(*args) runtime_shape = args[self.sym_shape_indices[0]] if runtime_shape not in self.concrete_size_entries: # we don't need to do anything for this shape return self.compiled_graph_for_general_shape(*args) entry = self.concrete_size_entries[runtime_shape] if entry.runnable is None: entry.runnable = self.compiled_graph_for_general_shape if entry.need_to_compile and not entry.compiled: entry.compiled = True self.to_be_compiled_sizes.remove(runtime_shape) # args are real arguments entry.runnable = self.sglang_backend.compiler_manager.compile( self.graph, args, self.inductor_config, graph_index=self.piecewise_compile_index, num_graphs=self.total_piecewise_compiles, runtime_shape=runtime_shape, ) # finished compilations for all required shapes if self.is_last_graph and not self.to_be_compiled_sizes: self.check_for_ending_compilation() if is_in_pcg_torch_compile(): return entry.runnable(*args) if entry.cudagraph is None: if entry.num_finished_warmup < 1: # noqa entry.num_finished_warmup += 1 return entry.runnable(*args) if self.compile_config.get_enable_debug_mode(): input_addresses = [ x.data_ptr() for x in args if isinstance(x, torch.Tensor) ] entry.input_addresses = input_addresses cudagraph = torch.cuda.CUDAGraph() with ExitStack() as stack: if not self.is_first_graph: # during every model forward, we will capture # many pieces of cudagraphs (roughly one per layer). # running gc again and again across layers will # make the cudagraph capture very slow. # therefore, we only run gc for the first graph, # and disable gc for the rest of the graphs. stack.enter_context(patch("gc.collect", lambda: None)) stack.enter_context(patch("torch.cuda.empty_cache", lambda: None)) # mind-exploding: carefully manage the reference and memory. stream = get_pcg_capture_stream() assert ( stream is not None ), "PCG capture stream is not set, please check if runtime recompilation happened" with torch.cuda.graph(cudagraph, pool=self.graph_pool, stream=stream): # `output` is managed by pytorch's cudagraph pool output = entry.runnable(*args) if self.is_last_graph: # by converting it to weak ref, # the original `output` will immediately be released # to save memory. It is only safe to do this for # the last graph, because the output of the last graph # will not be used by any other cuda graph. output = weak_ref_tensors(output) # here we always use weak ref for the output # to save memory entry.output = weak_ref_tensors(output) entry.cudagraph = cudagraph compilation_counter.num_cudagraph_captured += 1 # important: we need to return the output, rather than # the weak ref of the output, so that pytorch can correctly # manage the memory during cuda graph capture return output if self.compile_config.get_enable_debug_mode(): # check if the input addresses are the same new_input_addresses = [ x.data_ptr() for x in args if isinstance(x, torch.Tensor) ] assert new_input_addresses == entry.input_addresses, ( "Input addresses for cudagraphs are different during replay." f" Expected {entry.input_addresses}, got {new_input_addresses}" ) entry.cudagraph.replay() return entry.output ================================================ FILE: python/sglang/srt/compilation/fix_functionalization.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/fix_functionalization.py import logging import operator from collections.abc import Iterable from typing import Optional, Union import torch from torch._higher_order_ops.auto_functionalize import auto_functionalized from sglang.srt.compilation.fx_utils import is_func from sglang.srt.compilation.inductor_pass import SGLangInductorPass logger = logging.getLogger(__name__) class FixFunctionalizationPass(SGLangInductorPass): """ This pass defunctionalizes certain nodes to avoid redundant tensor copies. After this pass, DCE (dead-code elimination) should never be run, as de-functionalized nodes may appear as dead code. To add new nodes to defunctionalize, add to the if-elif chain in __call__. """ def __call__(self, graph: torch.fx.Graph): self.begin() self.dump_graph(graph, "before_fix_functionalization") self.nodes_to_remove: list[torch.fx.Node] = [] count = 0 for node in graph.nodes: if not is_func(node, auto_functionalized): continue # Avoid deep if-elif nesting count += 1 self.dump_graph(graph, "before_fix_functionalization_cleanup") # Remove the nodes all at once count_removed = len(self.nodes_to_remove) for node in self.nodes_to_remove: graph.erase_node(node) logger.debug( "De-functionalized %s nodes, removed %s nodes", count, count_removed ) self.dump_graph(graph, "after_fix_functionalization") self.end_and_log() def _remove(self, node_or_nodes: Union[torch.fx.Node, Iterable[torch.fx.Node]]): """ Stage a node (or nodes) for removal at the end of the pass. """ if isinstance(node_or_nodes, torch.fx.Node): self.nodes_to_remove.append(node_or_nodes) else: self.nodes_to_remove.extend(node_or_nodes) def defunctionalize( self, graph: torch.fx.Graph, node: torch.fx.Node, mutated_args: dict[int, Union[torch.fx.Node, str]], args: Optional[tuple[Union[torch.fx.Node, str], ...]] = None, ): """ De-functionalize a node by replacing it with a call to the original. It also replaces the getitem users with the mutated arguments. See replace_users_with_mutated_args and insert_defunctionalized. """ self.replace_users_with_mutated_args(node, mutated_args) self.insert_defunctionalized(graph, node, args=args) self._remove(node) def replace_users_with_mutated_args( self, node: torch.fx.Node, mutated_args: dict[int, Union[torch.fx.Node, str]] ): """ Replace all getitem users of the auto-functionalized node with the mutated arguments. :param node: The auto-functionalized node :param mutated_args: The mutated arguments, indexed by getitem index. If the value of an arg is a string, `node.kwargs[arg]` is used. """ for idx, user in self.getitem_users(node).items(): arg = mutated_args[idx] arg = node.kwargs[arg] if isinstance(arg, str) else arg user.replace_all_uses_with(arg) self._remove(user) def getitem_users(self, node: torch.fx.Node) -> dict[int, torch.fx.Node]: """ Returns the operator.getitem users of the auto-functionalized node, indexed by the index they are getting. """ users = {} for user in node.users: if is_func(user, operator.getitem): idx = user.args[1] users[idx] = user return users def insert_defunctionalized( self, graph: torch.fx.Graph, node: torch.fx.Node, args: Optional[tuple[Union[torch.fx.Node, str], ...]] = None, ): """ Insert a new defunctionalized node into the graph before node. If one of the kwargs is 'out', provide args directly, as node.kwargs cannot be used. See https://github.com/pytorch/pytorch/blob/a00faf440888ffb724bad413f329a49e2b6388e7/torch/_inductor/lowering.py#L351 :param graph: Graph to insert the defunctionalized node into :param node: The auto-functionalized node to defunctionalize :param args: If we cannot use kwargs, specify args directly. If an arg is a string, `node.kwargs[arg]` is used. """ # noqa: E501 assert is_func( node, auto_functionalized ), f"node must be auto-functionalized, is {node} instead" # Create a new call to the original function with graph.inserting_before(node): function = node.args[0] if args is None: graph.call_function(function, kwargs=node.kwargs) else: # Args passed as strings refer to items in node.kwargs args = tuple( node.kwargs[arg] if isinstance(arg, str) else arg for arg in args ) graph.call_function(function, args=args) ================================================ FILE: python/sglang/srt/compilation/fx_utils.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/fx_utils.py import operator from collections.abc import Iterable, Iterator from typing import Optional from torch import fx from torch._higher_order_ops.auto_functionalize import auto_functionalized from torch._ops import OpOverload def is_func(node: fx.Node, target) -> bool: return node.op == "call_function" and node.target == target def is_auto_func(node: fx.Node, op: OpOverload) -> bool: return is_func(node, auto_functionalized) and node.args[0] == op # Returns the first specified node with the given op (if it exists) def find_specified_fn_maybe( nodes: Iterable[fx.Node], op: OpOverload ) -> Optional[fx.Node]: for node in nodes: if node.target == op: return node return None # Returns the first specified node with the given op def find_specified_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node: node = find_specified_fn_maybe(nodes, op) assert node is not None, f"Could not find {op} in nodes {nodes}" return node # Returns the first auto_functionalized node with the given op (if it exists) def find_auto_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> Optional[fx.Node]: for node in nodes: if is_func(node, auto_functionalized) and node.args[0] == op: # noqa return node return None # Returns the first auto_functionalized node with the given op def find_auto_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node: node = find_auto_fn_maybe(nodes, op) assert node is not None, f"Could not find {op} in nodes {nodes}" return node # Returns the getitem node that extracts the idx-th element from node # (if it exists) def find_getitem_maybe(node: fx.Node, idx: int) -> Optional[fx.Node]: for user in node.users: if is_func(user, operator.getitem) and user.args[1] == idx: return user return None # Returns the getitem node that extracts the idx-th element from node def find_getitem(node: fx.Node, idx: int) -> fx.Node: ret = find_getitem_maybe(node, idx) assert ret is not None, f"Could not find getitem {idx} in node {node}" return ret # An auto-functionalization-aware utility for finding nodes with a specific op def find_op_nodes(op: OpOverload, graph: fx.Graph) -> Iterator[fx.Node]: if not op._schema.is_mutable: yield from graph.find_nodes(op="call_function", target=op) for n in graph.find_nodes(op="call_function", target=auto_functionalized): if n.args[0] == op: yield n # Asserts that the node only has one user and returns it # Even if a node has only 1 user, it might share storage with another node, # which might need to be taken into account. def get_only_user(node: fx.Node) -> fx.Node: assert len(node.users) == 1 return next(iter(node.users)) ================================================ FILE: python/sglang/srt/compilation/inductor_pass.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/inductor_pass.py import hashlib import inspect import json import logging import time import types from contextlib import contextmanager from typing import Any, Callable, Optional, Union import torch from torch import fx from torch._dynamo.utils import lazy_format_graph_code from torch._inductor.custom_graph_pass import CustomGraphPass logger = logging.getLogger(__name__) _pass_context = None class PassContext: def __init__(self, runtime_shape: Optional[int]): self.runtime_shape = runtime_shape def get_pass_context() -> PassContext: """Get the current pass context.""" assert _pass_context is not None return _pass_context @contextmanager def pass_context(runtime_shape: Optional[int]): """A context manager that stores the current pass context, usually it is a list of sizes to specialize. """ global _pass_context prev_context = _pass_context _pass_context = PassContext(runtime_shape) try: yield finally: _pass_context = prev_context class InductorPass(CustomGraphPass): """ A custom graph pass that uses a hash of its source as the UUID. This is defined as a convenience and should work in most cases. """ def uuid(self) -> Any: """ Provide a unique identifier for the pass, used in Inductor code cache. This should depend on the pass implementation, so that changes to the pass result in recompilation. By default, the object source is hashed. """ return InductorPass.hash_source(self) @staticmethod def hash_source(*srcs: Union[str, Any]): """ Utility method to hash the sources of functions or objects. :param srcs: strings or objects to add to the hash. Objects and functions have their source inspected. :return: """ hasher = hashlib.sha256() for src in srcs: if isinstance(src, str): src_str = src elif isinstance(src, types.FunctionType): src_str = inspect.getsource(src) else: src_str = inspect.getsource(src.__class__) hasher.update(src_str.encode("utf-8")) return hasher.hexdigest() @staticmethod def hash_dict(dict_: dict[Any, Any]): """ Utility method to hash a dictionary, can alternatively be used for uuid. :return: A sha256 hash of the json rep of the dictionary. """ encoded = json.dumps(dict_, sort_keys=True).encode("utf-8") return hashlib.sha256(encoded).hexdigest() def is_applicable_for_shape(self, shape: Optional[int]): return True class CallableInductorPass(InductorPass): """ This class is a wrapper for a callable that automatically provides an implementation of the UUID. """ def __init__( self, callable: Callable[[fx.Graph], None], uuid: Optional[Any] = None ): self.callable = callable self._uuid = self.hash_source(callable) if uuid is None else uuid def __call__(self, graph: torch.fx.Graph): self.callable(graph) def uuid(self) -> Any: return self._uuid class SGLangInductorPass(InductorPass): def __init__( self, ): self.pass_name = self.__class__.__name__ def dump_graph(self, graph: torch.fx.Graph, stage: str): lazy_format_graph_code(stage, graph.owning_module) def begin(self): self._start_time = time.perf_counter_ns() def end_and_log(self): self._end_time = time.perf_counter_ns() duration_ms = float(self._end_time - self._start_time) / 1.0e6 logger.debug("%s completed in %.1f ms", self.pass_name, duration_ms) class PrinterInductorPass(SGLangInductorPass): def __init__(self, name: str): super().__init__() self.name = name def __call__(self, graph: torch.fx.Graph): self.dump_graph(graph, self.name) ================================================ FILE: python/sglang/srt/compilation/npu_piecewise_backend.py ================================================ from contextlib import ExitStack from typing import Any, Callable from unittest.mock import patch import torch import torch.fx as fx from sglang.srt.compilation.compilation_config import CompilationConfig from sglang.srt.compilation.compilation_counter import compilation_counter from sglang.srt.compilation.cuda_piecewise_backend import ( CUDAPiecewiseBackend, weak_ref_tensors, ) class NPUPiecewiseBackend(CUDAPiecewiseBackend): def __init__( self, graph: fx.GraphModule, compile_config: CompilationConfig, inductor_config: dict[str, Any], graph_pool: Any, piecewise_compile_index: int, total_piecewise_compiles: int, sym_shape_indices: list[int], compiled_graph_for_general_shape: Callable, sglang_backend, ): super().__init__( graph, compile_config, inductor_config, graph_pool, piecewise_compile_index, total_piecewise_compiles, sym_shape_indices, compiled_graph_for_general_shape, sglang_backend, ) def __call__(self, *args): runtime_shape = args[self.sym_shape_indices[0]] if runtime_shape not in self.concrete_size_entries: # we don't need to do anything for this shape return self.compiled_graph_for_general_shape(*args) entry = self.concrete_size_entries[runtime_shape] if entry.runnable is None: entry.runnable = self.compiled_graph_for_general_shape if entry.cudagraph is None: if entry.num_finished_warmup < 1: # noqa entry.num_finished_warmup += 1 return entry.runnable(*args) if self.compile_config.get_enable_debug_mode(): input_addresses = [ x.data_ptr() for x in args if isinstance(x, torch.Tensor) ] entry.input_addresses = input_addresses npugraph = torch.npu.NPUGraph() with ExitStack() as stack: if not self.is_first_graph: # during every model forward, we will capture # many pieces of cudagraphs (roughly one per layer). # running gc again and again across layers will # make the cudagraph capture very slow. # therefore, we only run gc for the first graph, # and disable gc for the rest of the graphs. stack.enter_context(patch("gc.collect", lambda: None)) stack.enter_context(patch("torch.npu.empty_cache", lambda: None)) # mind-exploding: carefully manage the reference and memory. with torch.npu.graph(npugraph, pool=self.graph_pool): # `output` is managed by pytorch's cudagraph pool output = entry.runnable(*args) if self.is_last_graph: # by converting it to weak ref, # the original `output` will immediately be released # to save memory. It is only safe to do this for # the last graph, because the output of the last graph # will not be used by any other cuda graph. output = weak_ref_tensors(output) # here we always use weak ref for the output # to save memory entry.output = weak_ref_tensors(output) entry.cudagraph = npugraph compilation_counter.num_cudagraph_captured += 1 # important: we need to return the output, rather than # the weak ref of the output, so that pytorch can correctly # manage the memory during cuda graph capture return output if self.compile_config.get_enable_debug_mode(): # check if the input addresses are the same new_input_addresses = [ x.data_ptr() for x in args if isinstance(x, torch.Tensor) ] assert new_input_addresses == entry.input_addresses, ( "Input addresses for cudagraphs are different during replay." f" Expected {entry.input_addresses}, got {new_input_addresses}" ) entry.cudagraph.replay() return entry.output ================================================ FILE: python/sglang/srt/compilation/pass_manager.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/pass_manager.py import logging from torch import fx as fx from sglang.srt.compilation.fix_functionalization import FixFunctionalizationPass from sglang.srt.compilation.inductor_pass import ( CustomGraphPass, InductorPass, SGLangInductorPass, get_pass_context, ) logger = logging.getLogger(__name__) class PostGradPassManager(CustomGraphPass): """ The pass manager for post-grad passes. It handles configuration, adding custom passes, and running passes. It supports uuid for the Inductor code cache. That includes torch<2.6 support using pickling (in .inductor_pass.CustomGraphPass). The order of the post-grad post-passes is: 1. passes (constructor parameter) 2. default passes (NoopEliminationPass, FusionPass) 3. config["post_grad_custom_post_pass"] (if it exists) 4. fix_functionalization This way, all passes operate on a functionalized graph. """ def __init__(self): self.passes: list[SGLangInductorPass] = [] def __call__(self, graph: fx.Graph): shape = get_pass_context().runtime_shape for pass_ in self.passes: if pass_.is_applicable_for_shape(shape): pass_(graph) # always run fix_functionalization last self.fix_functionalization(graph) def configure( self, ): self.pass_config = dict() self.fix_functionalization = FixFunctionalizationPass() def add(self, pass_: InductorPass): assert isinstance(pass_, InductorPass) self.passes.append(pass_) def uuid(self): """ The PostGradPassManager is set as a custom pass in the Inductor and affects compilation caching. Its uuid depends on the UUIDs of all dependent passes and the pass config. See InductorPass for more info. """ pass_manager_uuid = "fshdakhsa" state = {"pass_config": pass_manager_uuid, "passes": []} for pass_ in self.passes: state["passes"].append(pass_.uuid()) state["passes"].append(self.fix_functionalization.uuid()) return InductorPass.hash_dict(state) ================================================ FILE: python/sglang/srt/compilation/piecewise_context_manager.py ================================================ from __future__ import annotations import logging from contextlib import contextmanager from dataclasses import dataclass from typing import TYPE_CHECKING, Any, List, Optional import torch logger = logging.getLogger(__name__) if TYPE_CHECKING: from sglang.srt.model_executor.forward_batch_info import ForwardBatch _in_piecewise_cuda_graph = False _in_pcg_torch_compile = False _pcg_capture_stream = None def is_in_piecewise_cuda_graph(): return _in_piecewise_cuda_graph def is_in_pcg_torch_compile(): return _in_pcg_torch_compile def get_pcg_capture_stream(): return _pcg_capture_stream @contextmanager def enable_piecewise_cuda_graph_compile(): global _in_pcg_torch_compile _in_pcg_torch_compile = True yield _in_pcg_torch_compile = False @contextmanager def enable_piecewise_cuda_graph(): global _in_piecewise_cuda_graph _in_piecewise_cuda_graph = True try: yield except Exception as e: logger.error( "Piecewise CUDA Graph failed with error: %s\n%s", e, PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG, ) raise finally: _in_piecewise_cuda_graph = False @contextmanager def set_pcg_capture_stream(stream: torch.cuda.Stream): global _pcg_capture_stream _pcg_capture_stream = stream yield _pcg_capture_stream = None @dataclass class ForwardContext: def __init__(self): self.forward_batch = None self.attention_layers = None self.quant_config = None self.moe_layers = None self.moe_fusions = None def set_forward_batch(self, forward_batch: ForwardBatch): self.forward_batch = forward_batch def set_attention_layers(self, layers: List[Any]): self.attention_layers = layers def set_quant_config(self, quant_config: Any): self.quant_config = quant_config def set_moe_layers(self, layers: List[Any]): self.moe_layers = layers def set_moe_fusions(self, fusions: List[Any]): self.moe_fusions = fusions _forward_context: Optional[ForwardContext] = None def get_forward_context() -> Optional[ForwardContext]: if _forward_context is None: return None return _forward_context @contextmanager def set_forward_context( forward_batch: ForwardBatch, attention_layers: List[Any], quant_config: Any, moe_layers: List[Any], moe_fusions: List[Any], ): global _forward_context _forward_context = ForwardContext() _forward_context.set_forward_batch(forward_batch) _forward_context.set_attention_layers(attention_layers) _forward_context.set_quant_config(quant_config) _forward_context.set_moe_layers(moe_layers) _forward_context.set_moe_fusions(moe_fusions) try: yield finally: _forward_context = None PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG = ( "Piecewise CUDA Graph is enabled by default as an experimental feature.\n" "To work around this error, add --disable-piecewise-cuda-graph to your launch command.\n" "Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose" ) ================================================ FILE: python/sglang/srt/compilation/weak_ref_tensor.py ================================================ from typing import Any, Union import torch from sglang.srt.utils.common import is_cuda, is_hip, is_npu if is_cuda() or is_hip(): from sgl_kernel import weak_ref_tensor elif is_npu(): from torch_npu._C import _weak_ref_tensor as weak_ref_tensor else: raise NotImplementedError("weak_ref_tensor is implemented only for CUDA and NPU.") def weak_ref_tensors( tensors: Union[torch.Tensor, list[torch.Tensor], tuple[torch.Tensor]], ) -> Union[torch.Tensor, list[Any], tuple[Any], Any]: """ Convenience function to create weak references to tensors, for single tensor, list of tensors or tuple of tensors. """ if isinstance(tensors, torch.Tensor): return weak_ref_tensor(tensors) if isinstance(tensors, list): return [weak_ref_tensor(t) for t in tensors] if isinstance(tensors, tuple): return tuple(weak_ref_tensor(t) for t in tensors) raise ValueError("Invalid type for tensors") ================================================ FILE: python/sglang/srt/configs/__init__.py ================================================ from sglang.srt.configs.afmoe import AfmoeConfig from sglang.srt.configs.bailing_hybrid import BailingHybridConfig from sglang.srt.configs.chatglm import ChatGLMConfig from sglang.srt.configs.dbrx import DbrxConfig from sglang.srt.configs.deepseekvl2 import DeepseekVL2Config from sglang.srt.configs.dots_ocr import DotsOCRConfig from sglang.srt.configs.dots_vlm import DotsVLMConfig from sglang.srt.configs.exaone import ExaoneConfig from sglang.srt.configs.falcon_h1 import FalconH1Config from sglang.srt.configs.granitemoehybrid import GraniteMoeHybridConfig from sglang.srt.configs.janus_pro import MultiModalityConfig from sglang.srt.configs.jet_nemotron import JetNemotronConfig from sglang.srt.configs.jet_vlm import JetVLMConfig from sglang.srt.configs.kimi_k25 import KimiK25Config from sglang.srt.configs.kimi_linear import KimiLinearConfig from sglang.srt.configs.kimi_vl import KimiVLConfig from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig from sglang.srt.configs.lfm2 import Lfm2Config from sglang.srt.configs.lfm2_moe import Lfm2MoeConfig from sglang.srt.configs.longcat_flash import LongcatFlashConfig from sglang.srt.configs.nano_nemotron_vl import NemotronH_Nano_VL_V2_Config from sglang.srt.configs.nemotron_h import NemotronHConfig from sglang.srt.configs.olmo3 import Olmo3Config from sglang.srt.configs.qwen3_5 import Qwen3_5Config, Qwen3_5MoeConfig from sglang.srt.configs.qwen3_next import Qwen3NextConfig from sglang.srt.configs.step3_vl import ( Step3TextConfig, Step3VisionEncoderConfig, Step3VLConfig, ) from sglang.srt.configs.step3p5 import Step3p5Config __all__ = [ "AfmoeConfig", "BailingHybridConfig", "ExaoneConfig", "ChatGLMConfig", "DbrxConfig", "DeepseekVL2Config", "LongcatFlashConfig", "MultiModalityConfig", "KimiVLConfig", "MoonViTConfig", "Step3VLConfig", "Step3TextConfig", "Step3VisionEncoderConfig", "Olmo3Config", "KimiLinearConfig", "KimiK25Config", "Qwen3NextConfig", "Qwen3_5Config", "Qwen3_5MoeConfig", "DotsVLMConfig", "DotsOCRConfig", "FalconH1Config", "GraniteMoeHybridConfig", "Lfm2Config", "Lfm2MoeConfig", "NemotronHConfig", "NemotronH_Nano_VL_V2_Config", "JetNemotronConfig", "JetVLMConfig", "Step3p5Config", ] ================================================ FILE: python/sglang/srt/configs/afmoe.py ================================================ from typing import List, Optional from transformers import PretrainedConfig class AfmoeConfig(PretrainedConfig): model_type = "afmoe" def __init__( self, vocab_size: int = 32000, hidden_size: int = 4096, intermediate_size: int = 11008, moe_intermediate_size: int = 256, num_hidden_layers: int = 32, num_attention_heads: int = 32, num_key_value_heads: Optional[int] = None, head_dim: Optional[int] = None, hidden_act: str = "silu", max_position_embeddings: int = 131072, initializer_range: float = 0.02, rms_norm_eps: float = 1e-5, use_cache: bool = True, pad_token_id: Optional[int] = None, bos_token_id: int = 1, eos_token_id: int = 2, tie_word_embeddings: bool = False, rope_theta: float = 10000.0, rope_scaling: Optional[dict] = None, attention_bias: bool = False, attention_dropout: float = 0.0, # MoE parameters num_experts: Optional[int] = None, num_experts_per_tok: Optional[int] = None, num_shared_experts: int = 0, num_dense_layers: int = 0, # Routing parameters score_func: str = "sigmoid", route_norm: bool = True, route_scale: float = 1.0, n_group: int = 1, topk_group: int = 1, # Attention parameters sliding_window: Optional[int] = None, layer_types: Optional[List[str]] = None, global_attn_every_n_layers: int = 4, # muP scaling mup_enabled: bool = False, **kwargs, ): self.vocab_size = vocab_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.moe_intermediate_size = moe_intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads if num_key_value_heads is None: num_key_value_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.head_dim = ( head_dim if head_dim is not None else hidden_size // num_attention_heads ) self.hidden_act = hidden_act self.initializer_range = initializer_range self.rms_norm_eps = rms_norm_eps self.use_cache = use_cache self.rope_theta = rope_theta self.rope_scaling = rope_scaling self.attention_bias = attention_bias self.attention_dropout = attention_dropout # MoE parameters self.num_experts = num_experts self.num_experts_per_tok = num_experts_per_tok self.num_shared_experts = num_shared_experts self.num_dense_layers = num_dense_layers # Routing parameters self.score_func = score_func self.route_norm = route_norm self.route_scale = route_scale self.n_group = n_group self.topk_group = topk_group # Attention parameters self.sliding_window = sliding_window self.layer_types = layer_types self.global_attn_every_n_layers = global_attn_every_n_layers # muP scaling self.mup_enabled = mup_enabled super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs, ) ================================================ FILE: python/sglang/srt/configs/bailing_hybrid.py ================================================ # coding=utf-8 # Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """BailingHybrid model configuration""" import enum from transformers.configuration_utils import PretrainedConfig from transformers.utils import logging from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape logger = logging.get_logger(__name__) class HybridLayerType(enum.Enum): full_attention = "attention" linear_attention = "linear_attention" class BailingHybridConfig(PretrainedConfig): model_type = "bailing_hybrid" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size=157184, hidden_size=2048, intermediate_size=5120, num_hidden_layers=20, num_attention_heads=16, num_key_value_heads=4, hidden_act="silu", use_qkv_bias=False, # bailing only use_bias=False, # bailing only rms_norm_eps=1e-06, tie_word_embeddings=False, # PretrainedConfig key, here change default value. embedding_dropout=0.0, attention_dropout=0.0, output_dropout=0.0, initializer_range=0.02, max_position_embeddings=32768, rope_theta=600000.0, use_cache=True, max_window_layers=20, rope_scaling=None, pad_token_id=156892, eos_token_id=156892, num_experts=256, num_shared_experts=1, num_experts_per_tok=8, n_group=8, topk_group=4, moe_intermediate_size=512, first_k_dense_replace=1, head_dim=128, output_router_logits=False, use_qk_norm=True, num_nextn_predict_layers=0, mtp_loss_scaling_factor=0, moe_router_enable_expert_bias=True, routed_scaling_factor=1.0, layer_group_size=1, group_norm_size=1, linear_silu=False, kv_lora_rank=512, q_lora_rank=None, qk_rope_head_dim=64, v_head_dim=128, qk_nope_head_dim=128, rope_interleave=True, **kwargs, ): self.num_hidden_layers = num_hidden_layers self.vocab_size = vocab_size self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.hidden_act = hidden_act self.use_qkv_bias = use_qkv_bias self.use_bias = use_bias self.rms_norm_eps = rms_norm_eps self.embedding_dropout = embedding_dropout self.attention_dropout = attention_dropout self.output_dropout = output_dropout self.num_nextn_predict_layers = num_nextn_predict_layers self.mtp_loss_scaling_factor = mtp_loss_scaling_factor self.initializer_range = initializer_range self.max_position_embeddings = max_position_embeddings self.rope_theta = rope_theta self.use_cache = use_cache self.max_window_layers = max_window_layers self.head_dim = head_dim or self.hidden_size // self.num_attention_heads self.rope_scaling = rope_scaling self.use_qk_norm = use_qk_norm self.moe_router_enable_expert_bias = moe_router_enable_expert_bias self.routed_scaling_factor = routed_scaling_factor # MoE configs self.num_experts = num_experts self.num_shared_experts = num_shared_experts self.num_experts_per_tok = num_experts_per_tok self.n_group = n_group self.topk_group = topk_group self.moe_intermediate_size = moe_intermediate_size self.first_k_dense_replace = first_k_dense_replace self.output_router_logits = output_router_logits # Linear configs self.layer_group_size = layer_group_size self.group_norm_size = group_norm_size self.linear_silu = linear_silu self.num_linear_key_value_heads = num_attention_heads # mla self.kv_lora_rank = kv_lora_rank self.q_lora_rank = q_lora_rank self.qk_rope_head_dim = qk_rope_head_dim self.v_head_dim = v_head_dim self.qk_nope_head_dim = qk_nope_head_dim self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim self.rope_interleave = rope_interleave self.for_nextn_model = False super().__init__( pad_token_id=pad_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs, ) @property def layers_block_type(self): if self.for_nextn_model: return [HybridLayerType.full_attention.value] layer_type_list = [] for l in range(self.num_hidden_layers): if (l + 1) % self.layer_group_size == 0: layer_type_list.append(HybridLayerType.full_attention.value) else: layer_type_list.append(HybridLayerType.linear_attention.value) return layer_type_list @property def linear_layer_ids(self): return [ i for i, type_value in enumerate(self.layers_block_type) if type_value == HybridLayerType.linear_attention.value ] @property def full_attention_layer_ids(self): return [ i for i, type_value in enumerate(self.layers_block_type) if type_value == HybridLayerType.full_attention.value ] @property def mamba2_cache_params(self) -> Mamba2CacheParams: from sglang.srt.layers.dp_attention import get_attention_tp_size shape = Mamba2StateShape.create( tp_world_size=get_attention_tp_size(), intermediate_size=0, n_groups=0, num_heads=self.num_linear_key_value_heads, head_dim=self.head_dim, state_size=self.head_dim, conv_kernel=1, ) return Mamba2CacheParams(shape=shape, layers=self.linear_layer_ids) ================================================ FILE: python/sglang/srt/configs/chatglm.py ================================================ # Adapted from # https://github.com/THUDM/ChatGLM2-6B # https://github.com/vllm-project/vllm/blob/main/vllm/transformers_utils/configs/chatglm.py # ChatGLM2 and ChatGLM3 share the same config. # ChatGLM4 is officially supported by Huggingface # transformers >= 4.46.0 is required # https://huggingface.co/docs/transformers/en/model_doc/glm from transformers import PretrainedConfig class ChatGLMConfig(PretrainedConfig): model_type = "chatglm" attribute_map = { "num_hidden_layers": "num_layers", "n_head_kv": "multi_query_group_num", } def __init__( self, num_layers=28, padded_vocab_size=65024, hidden_size=4096, ffn_hidden_size=13696, kv_channels=128, num_attention_heads=32, seq_length=2048, hidden_dropout=0.0, attention_dropout=0.0, layernorm_epsilon=1e-5, rmsnorm=True, apply_residual_connection_post_layernorm=False, post_layer_norm=True, add_bias_linear=False, add_qkv_bias=False, interleaved_qkv=False, bias_dropout_fusion=True, multi_query_attention=False, multi_query_group_num=1, apply_query_key_layer_scaling=True, attention_softmax_in_fp32=True, fp32_residual_connection=False, quantization_bit=0, pre_seq_len=None, prefix_projection=False, **kwargs ): self.num_layers = num_layers self.vocab_size = padded_vocab_size self.padded_vocab_size = padded_vocab_size self.hidden_size = hidden_size self.ffn_hidden_size = ffn_hidden_size self.kv_channels = kv_channels self.num_attention_heads = num_attention_heads self.seq_length = seq_length # It is to be compatible with long lora. self.max_position_embeddings = seq_length self.hidden_dropout = hidden_dropout self.attention_dropout = attention_dropout self.layernorm_epsilon = layernorm_epsilon self.rmsnorm = rmsnorm self.apply_residual_connection_post_layernorm = ( apply_residual_connection_post_layernorm ) self.post_layer_norm = post_layer_norm self.add_bias_linear = add_bias_linear self.add_qkv_bias = add_qkv_bias self.bias_dropout_fusion = bias_dropout_fusion self.multi_query_attention = multi_query_attention self.multi_query_group_num = multi_query_group_num self.apply_query_key_layer_scaling = apply_query_key_layer_scaling self.attention_softmax_in_fp32 = attention_softmax_in_fp32 self.fp32_residual_connection = fp32_residual_connection self.quantization_bit = quantization_bit self.pre_seq_len = pre_seq_len self.prefix_projection = prefix_projection self.interleaved_qkv = interleaved_qkv super().__init__(**kwargs) ================================================ FILE: python/sglang/srt/configs/dbrx.py ================================================ # Adapted from # https://huggingface.co/databricks/dbrx-base/blob/main/configuration_dbrx.py # https://github.com/vllm-project/vllm/blob/main/vllm/transformers_utils/configs/dbrx.py """Dbrx configuration.""" from typing import Any, Optional from transformers.configuration_utils import PretrainedConfig from transformers.utils import logging logger = logging.get_logger(__name__) DBRX_PRETRAINED_CONFIG_ARCHIVE_MAP = {} # type: ignore class DbrxAttentionConfig(PretrainedConfig): """Configuration class for Dbrx Attention. [`DbrxAttention`] class. It is used to instantiate attention layers according to the specified arguments, defining the layers architecture. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. Args: attn_pdrop (`float`, *optional*, defaults to 0.0): The dropout probability for the attention layers. clip_qkv (`float`, *optional*, defaults to None): If not `None`, clip the queries, keys, and values in the attention layer to this value. kv_n_heads (Optional[int]): For grouped_query_attention only, allow user to specify number of kv heads. rope_theta (float): The base frequency for rope. """ def __init__( self, attn_pdrop: float = 0, clip_qkv: Optional[float] = None, kv_n_heads: int = 1, rope_theta: float = 10000.0, **kwargs: Any, ): super().__init__(**kwargs) self.attn_pdrop = attn_pdrop self.clip_qkv = clip_qkv self.kv_n_heads = kv_n_heads self.rope_theta = rope_theta for k in ["model_type"]: if k in kwargs: kwargs.pop(k) if len(kwargs) != 0: raise ValueError(f"Found unknown {kwargs=}") @classmethod def from_pretrained( cls, pretrained_model_name_or_path: str, **kwargs: Any ) -> "PretrainedConfig": cls._set_token_in_kwargs(kwargs) config_dict, kwargs = cls.get_config_dict( pretrained_model_name_or_path, **kwargs ) if config_dict.get("model_type") == "dbrx": config_dict = config_dict["attn_config"] if ( "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type ): logger.warning( "You are using a model of type %s to instantiate a model of " "type %s. This is not supported for all configurations of " "models and can yield errors.", config_dict["model_type"], cls.model_type, ) return cls.from_dict(config_dict, **kwargs) class DbrxFFNConfig(PretrainedConfig): """Configuration class for Dbrx FFN. [`DbrxFFN`] class. It is used to instantiate feedforward layers according to the specified arguments, defining the layers architecture. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. Args: ffn_act_fn (dict, optional): A dict specifying activation function for the FFN. The dict should have a key 'name' with the value being the name of the activation function along with any additional keyword arguments. ffn_hidden_size (int, optional): The hidden size of the feedforward network. moe_num_experts (int, optional): The number of experts in the mixture of experts layer. moe_top_k (int, optional): The number of experts to use in the mixture of experts layer. moe_jitter_eps (float, optional): The jitter epsilon for the mixture of experts layer. moe_loss_weight (float, optional): The loss weight for the mixture of experts layer. moe_normalize_expert_weights (float, optional): The normalization factor for the expert weights. uniform_expert_assignment (bool, optional): Whether to use uniform expert assignment. This should only be used for benchmarking purposes. """ def __init__( self, ffn_act_fn: Optional[dict] = None, ffn_hidden_size: int = 3584, moe_num_experts: int = 4, moe_top_k: int = 1, moe_jitter_eps: Optional[float] = None, moe_loss_weight: float = 0.01, moe_normalize_expert_weights: Optional[float] = 1, uniform_expert_assignment: bool = False, **kwargs: Any, ): super().__init__() if ffn_act_fn is None: ffn_act_fn = {"name": "silu"} self.ffn_act_fn = ffn_act_fn self.ffn_hidden_size = ffn_hidden_size self.moe_num_experts = moe_num_experts self.moe_top_k = moe_top_k self.moe_jitter_eps = moe_jitter_eps self.moe_loss_weight = moe_loss_weight self.moe_normalize_expert_weights = moe_normalize_expert_weights self.uniform_expert_assignment = uniform_expert_assignment for k in ["model_type"]: if k in kwargs: kwargs.pop(k) if len(kwargs) != 0: raise ValueError(f"Found unknown {kwargs=}") @classmethod def from_pretrained( cls, pretrained_model_name_or_path: str, **kwargs: Any ) -> "PretrainedConfig": cls._set_token_in_kwargs(kwargs) config_dict, kwargs = cls.get_config_dict( pretrained_model_name_or_path, **kwargs ) if config_dict.get("model_type") == "dbrx": config_dict = config_dict["ffn_config"] if ( "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type ): logger.warning( "You are using a model of type %s to instantiate a model of " "type %s. This is not supported for all " "configurations of models and can yield errors.", config_dict["model_type"], cls.model_type, ) return cls.from_dict(config_dict, **kwargs) class DbrxConfig(PretrainedConfig): """Configuration class for Dbrx. [`DbrxModel`]. It is used to instantiate a Dbrx model according to the specified arguments, defining the model architecture. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. Args: d_model (`int`, *optional*, defaults to 6144): Dimensionality of the embeddings and hidden states. n_heads (`int`, *optional*, defaults to 48): Number of attention heads for each attention layer in the Transformer encoder. n_layers (`int`, *optional*, defaults to 40): Number of hidden layers in the Transformer encoder. max_seq_len (`int`, *optional*, defaults to 32768): The maximum sequence length of the model. vocab_size (`int`, *optional*, defaults to 100352): Vocabulary size of the Dbrx model. Defines the maximum number of different tokens that can be represented by the `inputs_ids` passed when calling [`DbrxModel`]. resid_pdrop (`float`, *optional*, defaults to 0.0): The dropout probability applied to the attention output before combining with residual. emb_pdrop (`float`, *optional*, defaults to 0.0): The dropout probability for the embedding layer. attn_config (`dict`, *optional*): A dictionary used to configure the model's attention module. ffn_config (`dict`, *optional*): A dictionary used to configure the model's FFN module. use_cache (`bool`, *optional*, defaults to `False`): Whether or not the model should return the last key/values attentions (not used by all models). initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices. output_router_logits (`bool`, *optional*, defaults to `False`): Whether or not the router logits should be returned by the model. Enabling this will also allow the model to output the auxiliary loss. See [here]() for more details router_aux_loss_coef (`float`, *optional*, defaults to 0.001): The aux loss factor for the total loss. Example: ```python >>> from transformers import DbrxConfig, DbrxModel >>> # Initializing a Dbrx configuration >>> configuration = DbrxConfig() >>> # Initializing a model (with random weights) from the configuration >>> model = DbrxModel(configuration) >>> # Accessing the model configuration >>> configuration = model.config ``` """ model_type = "dbrx" attribute_map = { "num_attention_heads": "n_heads", "hidden_size": "d_model", "num_hidden_layers": "n_layers", "max_position_embeddings": "max_seq_len", } def __init__( self, d_model: int = 2048, n_heads: int = 16, n_layers: int = 24, max_seq_len: int = 2048, vocab_size: int = 32000, resid_pdrop: float = 0.0, emb_pdrop: float = 0.0, attn_config: Optional[DbrxAttentionConfig] = None, ffn_config: Optional[DbrxFFNConfig] = None, use_cache: bool = True, initializer_range: float = 0.02, output_router_logits: bool = False, router_aux_loss_coef: float = 0.05, **kwargs: Any, ): if attn_config is None: self.attn_config = DbrxAttentionConfig() elif isinstance(attn_config, dict): self.attn_config = DbrxAttentionConfig(**attn_config) else: self.attn_config = attn_config if ffn_config is None: self.ffn_config = DbrxFFNConfig() elif isinstance(ffn_config, dict): self.ffn_config = DbrxFFNConfig(**ffn_config) else: self.ffn_config = ffn_config self.d_model = d_model self.n_heads = n_heads self.n_layers = n_layers self.max_seq_len = max_seq_len self.vocab_size = vocab_size self.resid_pdrop = resid_pdrop self.emb_pdrop = emb_pdrop self.use_cache = use_cache self.initializer_range = initializer_range self.output_router_logits = output_router_logits self.router_aux_loss_coef = router_aux_loss_coef tie_word_embeddings = kwargs.pop("tie_word_embeddings", False) if tie_word_embeddings: raise ValueError("tie_word_embeddings is not supported for Dbrx models.") super().__init__( tie_word_embeddings=tie_word_embeddings, **kwargs, ) ================================================ FILE: python/sglang/srt/configs/deepseek_ocr.py ================================================ import math from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple import torch from PIL import Image, ImageOps from transformers import ( AutoProcessor, LlamaTokenizerFast, PretrainedConfig, ProcessorMixin, ) from sglang.srt.multimodal.customized_mm_processor_utils import ( register_customized_processor, ) from sglang.srt.sampling.custom_logit_processor import ( DeepseekOCRNoRepeatNGramLogitProcessor, ) BASE_SIZE = 1024 IMAGE_SIZE = 640 CROP_MODE = True MIN_CROPS = 2 MAX_CROPS = 6 # max:9; If your GPU memory is small, it is recommended to set it to 6. MAX_CONCURRENCY = 100 # If you have limited GPU memory, lower the concurrency count. NUM_WORKERS = 64 # image pre-process (resize/padding) workers PRINT_NUM_VIS_TOKENS = False SKIP_REPEAT = True MODEL_PATH = "deepseek-ai/DeepSeek-OCR" # change to your model path NGRAM_NO_REPEAT_SIZE = 30 NGRAM_NO_REPEAT_WINDOW = 90 # Whitelist `` and `` token ids to allow table structures. NGRAM_NO_REPEAT_WHITELIST = (128821, 128822) DEFAULT_CUSTOM_LOGIT_PROCESSOR = DeepseekOCRNoRepeatNGramLogitProcessor.to_str() def get_default_ngram_custom_params() -> Dict[str, Any]: """Return default custom params for the DeepSeek-OCR n-gram no repeat processor.""" return { "ngram_size": NGRAM_NO_REPEAT_SIZE, "window_size": NGRAM_NO_REPEAT_WINDOW, "whitelist_token_ids": list(NGRAM_NO_REPEAT_WHITELIST), } PROMPT = "\n<|grounding|>Convert the document to markdown." class DictOutput(object): def items(self): return self.__dict__.items() def keys(self): return self.__dict__.keys() def __getitem__(self, item): return self.__dict__[item] def __contains__(self, key): return key in self.__dict__ def __setitem__(self, key, value): self.__dict__[key] = value @dataclass class VLChatProcessorOutput(DictOutput): input_ids: torch.LongTensor target_ids: torch.LongTensor images_crop: torch.LongTensor pixel_values: ( torch.Tensor ) # rename from "images" to "pixel_values" for compatibility images_seq_mask: torch.BoolTensor images_spatial_crop: torch.LongTensor def __len__(self): return len(self.input_ids) class ImageTransform(object): def __init__( self, mean: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5), std: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5), normalize: bool = True, ): self.mean = mean self.std = std self.normalize = normalize # only load torchvision.transforms when needed try: import torchvision.transforms as T # FIXME: add version check for gguf except ImportError as err: raise ImportError( "Please install torchvision via `pip install torchvision` to use Deepseek-VL2." ) from err transform_pipelines = [T.ToTensor()] if normalize: transform_pipelines.append(T.Normalize(mean, std)) self.transform = T.Compose(transform_pipelines) def __call__(self, pil_img: Image.Image): x = self.transform(pil_img) return x def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size): best_ratio_diff = float("inf") best_ratio = (1, 1) area = width * height for ratio in target_ratios: target_aspect_ratio = ratio[0] / ratio[1] ratio_diff = abs(aspect_ratio - target_aspect_ratio) if ratio_diff < best_ratio_diff: best_ratio_diff = ratio_diff best_ratio = ratio elif ratio_diff == best_ratio_diff: if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]: best_ratio = ratio return best_ratio def dynamic_preprocess( image, min_num=MIN_CROPS, max_num=MAX_CROPS, image_size=640, use_thumbnail=False ): orig_width, orig_height = image.size aspect_ratio = orig_width / orig_height # calculate the existing image aspect ratio target_ratios = set( (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if i * j <= max_num and i * j >= min_num ) target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1]) # find the closest aspect ratio to the target target_aspect_ratio = find_closest_aspect_ratio( aspect_ratio, target_ratios, orig_width, orig_height, image_size ) # calculate the target width and height target_width = image_size * target_aspect_ratio[0] target_height = image_size * target_aspect_ratio[1] blocks = target_aspect_ratio[0] * target_aspect_ratio[1] # resize the image resized_img = image.resize((target_width, target_height)) processed_images = [] for i in range(blocks): box = ( (i % (target_width // image_size)) * image_size, (i // (target_width // image_size)) * image_size, ((i % (target_width // image_size)) + 1) * image_size, ((i // (target_width // image_size)) + 1) * image_size, ) # split the image split_img = resized_img.crop(box) processed_images.append(split_img) assert len(processed_images) == blocks if use_thumbnail and len(processed_images) != 1: thumbnail_img = image.resize((image_size, image_size)) processed_images.append(thumbnail_img) return processed_images, target_aspect_ratio class DeepseekOCRProcessor(ProcessorMixin): tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast") attributes = ["tokenizer"] def __init__( self, tokenizer: LlamaTokenizerFast, candidate_resolutions: Tuple[Tuple[int, int]], patch_size: int, downsample_ratio: int, image_mean: Tuple[float, float, float] = (0.5, 0.5, 0.5), image_std: Tuple[float, float, float] = (0.5, 0.5, 0.5), normalize: bool = True, image_token: str = "", pad_token: str = "<｜▁pad▁｜>", add_special_token: bool = False, sft_format: str = "deepseek", mask_prompt: bool = True, ignore_id: int = -100, ocr2_mode: bool = False, **kwargs, ): self.candidate_resolutions = candidate_resolutions self.image_size = candidate_resolutions[0][0] self.patch_size = patch_size self.image_mean = image_mean self.image_std = image_std self.normalize = normalize self.downsample_ratio = downsample_ratio self.base_size = BASE_SIZE self.image_transform = ImageTransform( mean=image_mean, std=image_std, normalize=normalize ) self.tokenizer = tokenizer # must set this，padding side with make a difference in batch inference self.tokenizer.padding_side = "left" # add the pad_token as special token to use 'tokenizer.pad_token' and 'tokenizer.pad_token_id' if tokenizer.pad_token is None: self.tokenizer.add_special_tokens({"pad_token": pad_token}) # add image token image_token_id = self.tokenizer.vocab.get(image_token) if image_token_id is None: special_tokens = [image_token] special_tokens_dict = {"additional_special_tokens": special_tokens} self.tokenizer.add_special_tokens(special_tokens_dict) self.image_token_id = self.tokenizer.vocab.get(image_token) # add five special tokens for grounding-related tasks # <|ref|>, <|/ref|>, <|det|>, <|/det|>, <|grounding|> special_tokens = ["<|ref|>", "<|/ref|>", "<|det|>", "<|/det|>", "<|grounding|>"] special_tokens_dict = {"additional_special_tokens": special_tokens} self.tokenizer.add_special_tokens(special_tokens_dict) # add special tokens for SFT data special_tokens = ["<|User|>", "<|Assistant|>"] special_tokens_dict = {"additional_special_tokens": special_tokens} self.tokenizer.add_special_tokens(special_tokens_dict) self.image_token = image_token self.pad_token = pad_token self.add_special_token = add_special_token self.sft_format = sft_format self.mask_prompt = mask_prompt self.ignore_id = ignore_id self.ocr2_mode = ocr2_mode super().__init__( tokenizer, **kwargs, ) def format_messages_v2(self, messages: str, pil_images, max_req_input_len=-1): """play the role of format_messages_v2 and get_images_info in the last version""" tokenized_data = [] masked_tokenized_data = [] # labels images_list = [] images_seq_mask = [] images_spatial_crop = [] image_index = 0 image_token_cnt = messages.count(self.image_token) ( input_ids, images, images_crop, seq_mask, spatial_crop, num_image_tokens, image_shapes, ) = self.tokenize_with_images( messages, pil_images[image_index : image_index + image_token_cnt], bos=True, eos=True, cropping=len(pil_images) <= 2, ) image_index = image_token_cnt images_list += images images_seq_mask += seq_mask images_spatial_crop = spatial_crop return ( input_ids, masked_tokenized_data, images_list, images_seq_mask, images_spatial_crop, images_crop, ) @property def bos_id(self): return self.tokenizer.bos_token_id @property def eos_id(self): return self.tokenizer.eos_token_id @property def pad_id(self): return self.tokenizer.pad_token_id def encode(self, text: str, bos: bool = True, eos: bool = False): t = self.tokenizer.encode(text, add_special_tokens=False) if bos: t = [self.bos_id] + t if eos: t = t + [self.eos_id] return t def decode(self, t: List[int], **kwargs) -> str: return self.tokenizer.decode(t, **kwargs) def process_one( self, prompt: str = None, conversations: List[Dict[str, str]] = None, images: List[Image.Image] = None, apply_sft_format: bool = False, inference_mode: bool = True, system_prompt: str = "", max_req_input_len: int = -1, cropping: bool = True, **kwargs, ): """ Args: prompt (str): the formatted prompt; conversations (List[Dict]): conversations with a list of messages; images (List[ImageType]): the list of images; apply_sft_format (bool): if prompt is not None, then apply the SFT format to prompt; if conversations is not None, then it will always apply the SFT format to conversations; inference_mode (bool): if True, then remove the last eos token; system_prompt (str): the system prompt; **kwargs: Returns: outputs (BaseProcessorOutput): the output of the processor, - input_ids (torch.LongTensor): [N + image tokens] - target_ids (torch.LongTensor): [N + image tokens] - images (torch.FloatTensor): [n_images, 3, H, W] - image_id (int): the id of the image token - num_image_tokens (List[int]): the number of image tokens """ prompt = conversations or prompt ( input_ids, masked_tokenized_str, images_list, images_seq_mask, images_spatial_crop, images_crop, ) = self.format_messages_v2(prompt, images, max_req_input_len) target_ids = torch.LongTensor(masked_tokenized_str) has_images = len(images_list) > 0 has_local_crops = False if len(images_spatial_crop) > 0: has_local_crops = any( crop[0] > 1 or crop[1] > 1 for crop in images_spatial_crop ) if len(images_list) == 0: images = torch.zeros((1, 3, self.image_size, self.image_size)) else: images = torch.stack(images_list, dim=0) images_spatial_crop = torch.stack( [images_spatial_crop], dim=0 ) # stack the tensor to make it a batch of 1 prepare = VLChatProcessorOutput( input_ids=input_ids, target_ids=target_ids, images_crop=images_crop, pixel_values=images, images_seq_mask=images_seq_mask, images_spatial_crop=images_spatial_crop, ) prepare.has_images = has_images prepare.has_local_crops = has_local_crops return prepare def __call__( self, *, prompt: str = None, conversations: List[Dict[str, str]] = None, images: List[Image.Image] = None, apply_sft_format: bool = False, inference_mode: bool = True, system_prompt: str = "", max_req_input_len: int = -1, text: list[str] = None, **kwargs, ): assert text is None or isinstance(text, list) if text is not None: text = text[0] prepare = self.process_one( prompt=prompt or text, conversations=conversations, images=images, apply_sft_format=apply_sft_format, inference_mode=inference_mode, system_prompt=system_prompt, max_req_input_len=max_req_input_len, ) return prepare def find_all_indices(self, messages, target_value): indices = [] for index, item in enumerate(messages): if item == target_value: indices.append(index) return indices def tokenize_with_images( self, conversation: str, images: List[Image.Image], bos: bool = True, eos: bool = True, cropping: bool = True, ): """Tokenize text with tags.""" conversation = conversation assert conversation.count(self.image_token) == len(images) text_splits = conversation.split(self.image_token) images_list, images_crop_list, images_seq_mask, images_spatial_crop = ( [], [], [], [], ) image_shapes = [] num_image_tokens = [] tokenized_str = [] for text_sep, image in zip(text_splits, images): """encode text_sep""" tokenized_sep = self.encode(text_sep, bos=False, eos=False) tokenized_str += tokenized_sep images_seq_mask += [False] * len(tokenized_sep) image_shapes.append(image.size) if image.size[0] <= 640 and image.size[1] <= 640: crop_ratio = [1, 1] else: if cropping: images_crop_raw, crop_ratio = dynamic_preprocess( image, image_size=IMAGE_SIZE ) else: crop_ratio = [1, 1] """process the global view""" if self.image_size <= 640 and not cropping: image = image.resize((self.image_size, self.image_size)) global_view = ImageOps.pad( image, (self.base_size, self.base_size), color=tuple(int(x * 255) for x in self.image_transform.mean), ) images_list.append(self.image_transform(global_view)) num_width_tiles, num_height_tiles = crop_ratio images_spatial_crop.append([num_width_tiles, num_height_tiles]) if num_width_tiles > 1 or num_height_tiles > 1: for i in range(len(images_crop_raw)): images_crop_list.append(self.image_transform(images_crop_raw[i])) """add image tokens""" num_queries = math.ceil( (self.image_size // self.patch_size) / self.downsample_ratio ) num_queries_base = math.ceil( (self.base_size // self.patch_size) / self.downsample_ratio ) if self.ocr2_mode: tokenized_image = [] if num_width_tiles > 1 or num_height_tiles > 1: tokenized_image += [self.image_token_id] * ( num_queries * num_width_tiles * num_queries * num_height_tiles ) tokenized_image += [self.image_token_id] * ( num_queries_base * num_queries_base ) # One extra token for the view separator. tokenized_image += [self.image_token_id] else: tokenized_image = ( [self.image_token_id] * num_queries_base + [self.image_token_id] ) * num_queries_base tokenized_image += [self.image_token_id] if num_width_tiles > 1 or num_height_tiles > 1: tokenized_image += ( [self.image_token_id] * (num_queries * num_width_tiles) + [self.image_token_id] ) * (num_queries * num_height_tiles) tokenized_str += tokenized_image images_seq_mask += [True] * len(tokenized_image) num_image_tokens.append(len(tokenized_image)) """process the last text split""" tokenized_sep = self.encode(text_splits[-1], bos=False, eos=False) tokenized_str += tokenized_sep images_seq_mask += [False] * len(tokenized_sep) """add the bos and eos tokens""" if bos: tokenized_str = [self.bos_id] + tokenized_str images_seq_mask = [False] + images_seq_mask if eos: tokenized_str = tokenized_str + [self.eos_id] images_seq_mask = images_seq_mask + [False] assert len(tokenized_str) == len( images_seq_mask ), f"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}" masked_tokenized_str = [] for token_index in tokenized_str: if token_index != self.image_token_id: masked_tokenized_str.append(token_index) else: masked_tokenized_str.append(self.ignore_id) assert ( len(tokenized_str) == len(images_seq_mask) == len(masked_tokenized_str) ), ( f"tokenized_str's length {len(tokenized_str)}, input_ids' length {len(masked_tokenized_str)}, " f"imags_seq_mask's length {len(images_seq_mask)}, are not equal" ) input_ids = torch.LongTensor(tokenized_str) target_ids = torch.LongTensor(masked_tokenized_str) images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool) # set input_ids < 0 | input_ids == self.image_token_id as ignore_id target_ids[(input_ids < 0) | (input_ids == self.image_token_id)] = ( self.ignore_id ) input_ids[input_ids < 0] = self.pad_id inference_mode = True if inference_mode: # Remove the ending eos token assert input_ids[-1] == self.eos_id input_ids = input_ids[:-1] target_ids = target_ids[:-1] images_seq_mask = images_seq_mask[:-1] if len(images_list) == 0: pixel_values = torch.zeros((1, 3, self.base_size, self.base_size)) images_spatial_crop = torch.zeros((1, 1), dtype=torch.long) images_crop = torch.zeros( (1, 3, self.image_size, self.image_size) ).unsqueeze(0) else: pixel_values = torch.stack(images_list, dim=0) images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long) if images_crop_list: images_crop = torch.stack(images_crop_list, dim=0).unsqueeze(0) else: images_crop = torch.zeros( (1, 3, self.image_size, self.image_size) ).unsqueeze(0) input_ids = input_ids.unsqueeze(0) return ( input_ids, pixel_values, images_crop, images_seq_mask, images_spatial_crop, num_image_tokens, image_shapes, ) class VisionEncoderConfig(PretrainedConfig): model_type: str = "vision" model_name: str = "vit_so400m_patch14_siglip_384.webli" image_size: int = 384 patch_size: int = 16 width: int = 1024 layers: int = 24 heads: int = 16 mlp_ratio: int = 4 global_pool: str = "map" ignore_head: bool = True class_token: bool = False num_classes: int = 0 use_checkpoint: bool = False weight_init: str = "skip" deterministic: bool = False num_recomputing_layers: int = 0 def __init__( self, model_name: str = "vit_so400m_patch14_siglip_384.webli", image_size: int = 384, patch_size: int = 16, width: int = 1024, layers: int = 24, heads: int = 16, mlp_ratio: int = 4, global_pool: str = "map", ignore_head: bool = True, class_token: bool = False, num_classes: int = 0, use_checkpoint: bool = False, **kwargs, ): self.model_name = model_name self.image_size = image_size self.patch_size = patch_size self.width = width self.layers = layers self.heads = heads self.mlp_ratio = mlp_ratio self.global_pool = global_pool self.ignore_head = ignore_head self.class_token = class_token self.num_classes = num_classes self.use_checkpoint = use_checkpoint super().__init__(**kwargs) class MlpProjectorConfig(PretrainedConfig): model_type = "mlp_projector" projector_type: str = "downsample_mlp_gelu" input_dim: int = 1152 n_embed: int = 2048 depth: int = 2 mlp_ratio: int = 1 downsample_ratio: int = 2 token_pooling: bool = False def __init__( self, projector_type: str = "downsample_mlp_gelu", input_dim: int = 1152, n_embed: int = 2048, depth: int = 2, mlp_ratio: int = 1, downsample_ratio: int = 2, **kwargs, ): self.projector_type = projector_type self.input_dim = input_dim self.n_embed = n_embed self.depth = depth self.mlp_ratio = mlp_ratio self.downsample_ratio = downsample_ratio super().__init__(**kwargs) class DeepseekV2Config(PretrainedConfig): model_type = "deepseek_v2" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size=102400, hidden_size=4096, intermediate_size=11008, moe_intermediate_size=1407, num_hidden_layers=30, num_attention_heads=32, num_key_value_heads=32, n_shared_experts=None, n_routed_experts=None, ep_size=1, routed_scaling_factor=1.0, kv_lora_rank=512, q_lora_rank=1536, qk_rope_head_dim=64, v_head_dim=128, qk_nope_head_dim=128, topk_method="gready", n_group=None, topk_group=None, num_experts_per_tok=None, moe_layer_freq=1, first_k_dense_replace=0, norm_topk_prob=False, scoring_func="softmax", aux_loss_alpha=0.001, seq_aux=True, hidden_act="silu", max_position_embeddings=2048, initializer_range=0.02, rms_norm_eps=1e-6, use_cache=True, pad_token_id=None, bos_token_id=100000, eos_token_id=100001, pretraining_tp=1, tie_word_embeddings=False, rope_theta=10000.0, rope_scaling=None, attention_bias=False, attention_dropout=0.0, use_mla=True, **kwargs, ): self.vocab_size = vocab_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.moe_intermediate_size = moe_intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.n_shared_experts = n_shared_experts self.n_routed_experts = n_routed_experts self.ep_size = ep_size self.routed_scaling_factor = routed_scaling_factor self.kv_lora_rank = kv_lora_rank self.q_lora_rank = q_lora_rank self.qk_rope_head_dim = qk_rope_head_dim self.v_head_dim = v_head_dim self.qk_nope_head_dim = qk_nope_head_dim self.topk_method = topk_method self.n_group = n_group self.topk_group = topk_group self.num_experts_per_tok = num_experts_per_tok self.moe_layer_freq = moe_layer_freq self.first_k_dense_replace = first_k_dense_replace self.norm_topk_prob = norm_topk_prob self.scoring_func = scoring_func self.aux_loss_alpha = aux_loss_alpha self.seq_aux = seq_aux # for backward compatibility if num_key_value_heads is None: num_key_value_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.hidden_act = hidden_act self.initializer_range = initializer_range self.rms_norm_eps = float(rms_norm_eps) self.pretraining_tp = pretraining_tp self.use_cache = use_cache self.rope_theta = rope_theta self.rope_scaling = rope_scaling self.attention_bias = attention_bias self.attention_dropout = attention_dropout self.use_mla = use_mla super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs, ) @register_customized_processor(processor_class=DeepseekOCRProcessor) class DeepseekVLV2Config(PretrainedConfig): # model_type = "deepseek_vl_v2" model_type = "deepseek-ocr" vision_config: VisionEncoderConfig = None projector_config: MlpProjectorConfig = None tile_tag: str = "2D" global_view_pos: str = "head" candidate_resolutions: tuple[tuple[int, int]] = ((384, 384),) customized_processor_type: type[Any] = DeepseekOCRProcessor def __init__( self, tile_tag: str = "tile_tag", global_view_pos: str = "head", candidate_resolutions: tuple[tuple[int, int]] = ((384, 384),), **kwargs, ): super().__init__(**kwargs) vision_config = kwargs.get("vision_config", {}) self.vision_config = VisionEncoderConfig(**vision_config) projector_config = kwargs.get("projector_config", {}) self.projector_config = MlpProjectorConfig(**projector_config) language_config = kwargs.get("language_config", {}) self.text_config = DeepseekV2Config(**language_config) self.tile_tag = tile_tag self.global_view_pos = global_view_pos self.candidate_resolutions = candidate_resolutions self.vocab_size = self.text_config.vocab_size self.hidden_size = self.text_config.hidden_size AutoProcessor.register(DeepseekVLV2Config, DeepseekOCRProcessor) ================================================ FILE: python/sglang/srt/configs/deepseekvl2.py ================================================ import math from dataclasses import dataclass from typing import Dict, List, Optional, Tuple import torch from PIL import Image, ImageOps from transformers import ( AutoProcessor, LlamaTokenizerFast, PretrainedConfig, ProcessorMixin, ) def select_best_resolution(image_size, candidate_resolutions): # used for cropping original_width, original_height = image_size best_fit = None max_effective_resolution = 0 min_wasted_resolution = float("inf") for width, height in candidate_resolutions: scale = min(width / original_width, height / original_height) downscaled_width, downscaled_height = int(original_width * scale), int( original_height * scale ) effective_resolution = min( downscaled_width * downscaled_height, original_width * original_height ) wasted_resolution = (width * height) - effective_resolution if effective_resolution > max_effective_resolution or ( effective_resolution == max_effective_resolution and wasted_resolution < min_wasted_resolution ): max_effective_resolution = effective_resolution min_wasted_resolution = wasted_resolution best_fit = (width, height) return best_fit class DictOutput(object): def items(self): return self.__dict__.items() def keys(self): return self.__dict__.keys() def __getitem__(self, item): return self.__dict__[item] def __contains__(self, key): return key in self.__dict__ def __setitem__(self, key, value): self.__dict__[key] = value @dataclass class VLChatProcessorOutput(DictOutput): input_ids: torch.LongTensor target_ids: torch.LongTensor pixel_values: ( torch.Tensor ) # rename from "images" to "pixel_values" for compatibility images_seq_mask: torch.BoolTensor images_spatial_crop: torch.LongTensor def __len__(self): return len(self.input_ids) class ImageTransform(object): def __init__( self, mean: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5), std: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5), normalize: bool = True, ): self.mean = mean self.std = std self.normalize = normalize # only load torchvision.transforms when needed try: import torchvision.transforms as T # FIXME: add version check for gguf except ImportError as err: raise ImportError( "Please install torchvision via `pip install torchvision` to use Deepseek-VL2." ) from err transform_pipelines = [T.ToTensor()] if normalize: transform_pipelines.append(T.Normalize(mean, std)) self.transform = T.Compose(transform_pipelines) def __call__(self, pil_img: Image.Image): x = self.transform(pil_img) return x class DeepseekVLV2Processor(ProcessorMixin): tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast") attributes = ["tokenizer"] def __init__( self, tokenizer: LlamaTokenizerFast, candidate_resolutions: Tuple[Tuple[int, int]], patch_size: int, downsample_ratio: int, image_mean: Tuple[float, float, float] = (0.5, 0.5, 0.5), image_std: Tuple[float, float, float] = (0.5, 0.5, 0.5), normalize: bool = True, image_token: str = "", pad_token: str = "<｜▁pad▁｜>", add_special_token: bool = False, sft_format: str = "deepseek", mask_prompt: bool = True, ignore_id: int = -100, **kwargs, ): self.candidate_resolutions = candidate_resolutions self.image_size = candidate_resolutions[0][0] self.patch_size = patch_size self.image_mean = image_mean self.image_std = image_std self.normalize = normalize self.downsample_ratio = downsample_ratio self.image_transform = ImageTransform( mean=image_mean, std=image_std, normalize=normalize ) self.tokenizer = tokenizer # must set this，padding side with make a difference in batch inference self.tokenizer.padding_side = "left" # add the pad_token as special token to use 'tokenizer.pad_token' and 'tokenizer.pad_token_id' if tokenizer.pad_token is None: self.tokenizer.add_special_tokens({"pad_token": pad_token}) # add image token image_token_id = self.tokenizer.vocab.get(image_token) if image_token_id is None: special_tokens = [image_token] special_tokens_dict = {"additional_special_tokens": special_tokens} self.tokenizer.add_special_tokens(special_tokens_dict) self.image_token_id = self.tokenizer.vocab.get(image_token) # add five special tokens for grounding-related tasks # <|ref|>, <|/ref|>, <|det|>, <|/det|>, <|grounding|> special_tokens = ["<|ref|>", "<|/ref|>", "<|det|>", "<|/det|>", "<|grounding|>"] special_tokens_dict = {"additional_special_tokens": special_tokens} self.tokenizer.add_special_tokens(special_tokens_dict) # add special tokens for SFT data special_tokens = ["<|User|>", "<|Assistant|>"] special_tokens_dict = {"additional_special_tokens": special_tokens} self.tokenizer.add_special_tokens(special_tokens_dict) self.image_token = image_token self.pad_token = pad_token self.add_special_token = add_special_token self.sft_format = sft_format self.mask_prompt = mask_prompt self.ignore_id = ignore_id super().__init__( tokenizer, **kwargs, ) def format_messages_v2(self, messages, pil_images, max_req_input_len=-1): """play the role of format_messages_v2 and get_images_info in the last version""" tokenized_data = [] masked_tokenized_data = [] # labels images_list = [] images_seq_mask = [] images_spatial_crop = [] image_index = 0 image_token_cnt = messages.count(self.image_token) tokenized_str, images, seq_mask, spatial_crop = self.tokenize_with_images( messages, pil_images[image_index : image_index + image_token_cnt], bos=True, eos=True, cropping=len(pil_images) <= 2, max_req_input_len=max_req_input_len, ) image_index = image_token_cnt tokenized_data += tokenized_str if self.mask_prompt: masked_tokenized_data += [self.ignore_id] * len(tokenized_str) else: masked_tokenized_data += tokenized_str images_list += images images_seq_mask += seq_mask images_spatial_crop += spatial_crop assert len(tokenized_data) == len( images_seq_mask ), f"format_messages_v2: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}" return ( tokenized_data, masked_tokenized_data, images_list, images_seq_mask, images_spatial_crop, ) @property def bos_id(self): return self.tokenizer.bos_token_id @property def eos_id(self): return self.tokenizer.eos_token_id @property def pad_id(self): return self.tokenizer.pad_token_id def encode(self, text: str, bos: bool = True, eos: bool = False): t = self.tokenizer.encode(text, add_special_tokens=False) if bos: t = [self.bos_id] + t if eos: t = t + [self.eos_id] return t def decode(self, t: List[int], **kwargs) -> str: return self.tokenizer.decode(t, **kwargs) def process_one( self, prompt: str = None, conversations: List[Dict[str, str]] = None, images: List[Image.Image] = None, apply_sft_format: bool = False, inference_mode: bool = True, system_prompt: str = "", max_req_input_len: int = -1, **kwargs, ): """ Args: prompt (str): the formatted prompt; conversations (List[Dict]): conversations with a list of messages; images (List[ImageType]): the list of images; apply_sft_format (bool): if prompt is not None, then apply the SFT format to prompt; if conversations is not None, then it will always apply the SFT format to conversations; inference_mode (bool): if True, then remove the last eos token; system_prompt (str): the system prompt; **kwargs: Returns: outputs (BaseProcessorOutput): the output of the processor, - input_ids (torch.LongTensor): [N + image tokens] - target_ids (torch.LongTensor): [N + image tokens] - images (torch.FloatTensor): [n_images, 3, H, W] - image_id (int): the id of the image token - num_image_tokens (List[int]): the number of image tokens """ assert ( prompt is None or conversations is None ), "prompt and conversations cannot be used at the same time." ( tokenized_str, masked_tokenized_str, images_list, images_seq_mask, images_spatial_crop, ) = self.format_messages_v2(conversations, images, max_req_input_len) assert ( len(tokenized_str) == len(images_seq_mask) == len(masked_tokenized_str) ), ( f"tokenized_str's length {len(tokenized_str)}, input_ids' length {len(masked_tokenized_str)}, " f"imags_seq_mask's length {len(images_seq_mask)}, are not equal" ) input_ids = torch.LongTensor(tokenized_str) target_ids = torch.LongTensor(masked_tokenized_str) images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool) # set input_ids < 0 | input_ids == self.image_token_id as ignore_id target_ids[(input_ids < 0) | (input_ids == self.image_token_id)] = ( self.ignore_id ) input_ids[input_ids < 0] = self.pad_id if inference_mode: assert input_ids[-1] == self.eos_id input_ids = input_ids[:-1] target_ids = target_ids[:-1] images_seq_mask = images_seq_mask[:-1] if len(images_list) == 0: images = torch.zeros((1, 3, self.image_size, self.image_size)) images_spatial_crop = torch.zeros((1, 2), dtype=torch.long) else: images = torch.stack(images_list, dim=0) images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long) images_spatial_crop = torch.stack( [images_spatial_crop], dim=0 ) # stack the tensor to make it a batch of 1 prepare = VLChatProcessorOutput( input_ids=input_ids, target_ids=target_ids, pixel_values=images, images_seq_mask=images_seq_mask, images_spatial_crop=images_spatial_crop, ) return prepare def __call__( self, *, prompt: str = None, conversations: List[Dict[str, str]] = None, images: List[Image.Image] = None, apply_sft_format: bool = False, inference_mode: bool = True, system_prompt: str = "", max_req_input_len: int = -1, **kwargs, ): prepare = self.process_one( prompt=prompt, conversations=conversations, images=images, apply_sft_format=apply_sft_format, inference_mode=inference_mode, system_prompt=system_prompt, max_req_input_len=max_req_input_len, ) return prepare def find_all_indices(self, messages, target_value): indices = [] for index, item in enumerate(messages): if item == target_value: indices.append(index) return indices def tokenize_with_images( self, conversation: str, images: List[Image.Image], bos: bool = True, eos: bool = True, cropping: bool = True, max_req_input_len: int = -1, ): """Tokenize text with tags.""" images_list, images_seq_mask, images_spatial_crop = [], [], [] text_splits = conversation.split(self.image_token) tokenized_str = [] for text_sep, image in zip(text_splits, images): """encode text_sep""" tokenized_sep = self.encode(text_sep, bos=False, eos=False) tokenized_str += tokenized_sep images_seq_mask += [False] * len(tokenized_sep) """select best resolution for anyres""" if cropping: best_width, best_height = select_best_resolution( image.size, self.candidate_resolutions ) else: best_width, best_height = self.image_size, self.image_size # print(image.size, (best_width, best_height)) # check the select_best_resolutions func """process the global view""" global_view = ImageOps.pad( image, (self.image_size, self.image_size), color=tuple(int(x * 255) for x in self.image_transform.mean), ) images_list.append(self.image_transform(global_view)) """process the local views""" local_view = ImageOps.pad( image, (best_width, best_height), color=tuple(int(x * 255) for x in self.image_transform.mean), ) for i in range(0, best_height, self.image_size): for j in range(0, best_width, self.image_size): images_list.append( self.image_transform( local_view.crop( (j, i, j + self.image_size, i + self.image_size) ) ) ) """record height / width crop num""" num_width_tiles, num_height_tiles = ( best_width // self.image_size, best_height // self.image_size, ) images_spatial_crop.append([num_width_tiles, num_height_tiles]) """add image tokens""" h = w = math.ceil( (self.image_size // self.patch_size) / self.downsample_ratio ) # global views tokens h * (w + 1), 1 is for line separator tokenized_image = [self.image_token_id] * h * (w + 1) # add a separator between global and local views tokenized_image += [self.image_token_id] # local views tokens, (num_height_tiles * h) * (num_width_tiles * w + 1) tokenized_image += ( [self.image_token_id] * (num_height_tiles * h) * (num_width_tiles * w + 1) ) tokenized_str += tokenized_image images_seq_mask += [True] * len(tokenized_image) # print(width_crop_num, height_crop_num, len(tokenized_image)) # test the correctness of the number of image-related tokens """process the last text split""" tokenized_sep = self.encode(text_splits[-1], bos=False, eos=False) # deal with video, limit with request len if max_req_input_len > -1: if max_req_input_len < len(tokenized_sep) + len(tokenized_str) - 1: rest = max_req_input_len - len(tokenized_sep) - 1 - 1024 tokenized_str = tokenized_str[:rest] images_seq_mask = images_seq_mask[:rest] tokenized_str += tokenized_sep images_seq_mask += [False] * len(tokenized_sep) """add the bos and eos tokens""" if bos: tokenized_str = [self.bos_id] + tokenized_str images_seq_mask = [False] + images_seq_mask if eos: tokenized_str = tokenized_str + [self.eos_id] images_seq_mask = images_seq_mask + [False] assert len(tokenized_str) == len( images_seq_mask ), f"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}" return tokenized_str, images_list, images_seq_mask, images_spatial_crop class DeepseekVL2VisionEncoderConfig(PretrainedConfig): model_type: str = "vision" model_name: str = "siglip_large_patch16_384" image_size: int = 384 patch_size: int = 16 width: int = 1024 layers: int = 24 heads: int = 16 mlp_ratio: int = 4 global_pool: str = "map" ignore_head: bool = True class_token: bool = False num_classes: int = 0 use_checkpoint: bool = False weight_init: str = "skip" deterministic: bool = False num_recomputing_layers: int = 0 def __init__( self, model_name: str = "siglip_large_patch16_384", image_size: int = 384, patch_size: int = 16, width: int = 1024, layers: int = 24, heads: int = 16, mlp_ratio: int = 4, global_pool: str = "map", ignore_head: bool = True, class_token: bool = False, num_classes: int = 0, use_checkpoint: bool = False, **kwargs, ): self.model_name = model_name self.image_size = image_size self.patch_size = patch_size self.width = width self.layers = layers self.heads = heads self.mlp_ratio = mlp_ratio self.global_pool = global_pool self.ignore_head = ignore_head self.class_token = class_token self.num_classes = num_classes self.use_checkpoint = use_checkpoint super().__init__(**kwargs) class DeepseekVL2MlpProjectorConfig(PretrainedConfig): model_type = "mlp_projector" projector_type: str = "downsample_mlp_gelu" input_dim: int = 1152 n_embed: int = 2048 depth: int = 2 mlp_ratio: int = 1 downsample_ratio: int = 2 token_pooling: bool = False def __init__( self, projector_type: str = "downsample_mlp_gelu", input_dim: int = 1152, n_embed: int = 2048, depth: int = 2, mlp_ratio: int = 1, downsample_ratio: int = 2, **kwargs, ): self.projector_type = projector_type self.input_dim = input_dim self.n_embed = n_embed self.depth = depth self.mlp_ratio = mlp_ratio self.downsample_ratio = downsample_ratio super().__init__(**kwargs) class DeepseekV2Config(PretrainedConfig): model_type = "deepseek_v2" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size=102400, hidden_size=4096, intermediate_size=11008, moe_intermediate_size=1407, num_hidden_layers=30, num_attention_heads=32, num_key_value_heads=32, n_shared_experts=None, n_routed_experts=None, ep_size=1, routed_scaling_factor=1.0, kv_lora_rank=512, q_lora_rank=1536, qk_rope_head_dim=64, v_head_dim=128, qk_nope_head_dim=128, topk_method="gready", n_group=None, topk_group=None, num_experts_per_tok=None, moe_layer_freq=1, first_k_dense_replace=0, norm_topk_prob=False, scoring_func="softmax", aux_loss_alpha=0.001, seq_aux=True, hidden_act="silu", max_position_embeddings=2048, initializer_range=0.02, rms_norm_eps=1e-6, use_cache=True, pad_token_id=None, bos_token_id=100000, eos_token_id=100001, pretraining_tp=1, tie_word_embeddings=False, rope_theta=10000.0, rope_scaling=None, attention_bias=False, attention_dropout=0.0, use_mla=True, **kwargs, ): self.vocab_size = vocab_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.moe_intermediate_size = moe_intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.n_shared_experts = n_shared_experts self.n_routed_experts = n_routed_experts self.ep_size = ep_size self.routed_scaling_factor = routed_scaling_factor self.kv_lora_rank = kv_lora_rank self.q_lora_rank = q_lora_rank self.qk_rope_head_dim = qk_rope_head_dim self.v_head_dim = v_head_dim self.qk_nope_head_dim = qk_nope_head_dim self.topk_method = topk_method self.n_group = n_group self.topk_group = topk_group self.num_experts_per_tok = num_experts_per_tok self.moe_layer_freq = moe_layer_freq self.first_k_dense_replace = first_k_dense_replace self.norm_topk_prob = norm_topk_prob self.scoring_func = scoring_func self.aux_loss_alpha = aux_loss_alpha self.seq_aux = seq_aux # for backward compatibility if num_key_value_heads is None: num_key_value_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.hidden_act = hidden_act self.initializer_range = initializer_range self.rms_norm_eps = float(rms_norm_eps) self.pretraining_tp = pretraining_tp self.use_cache = use_cache self.rope_theta = rope_theta self.rope_scaling = rope_scaling self.attention_bias = attention_bias self.attention_dropout = attention_dropout self.use_mla = use_mla super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs, ) class DeepseekVL2Config(PretrainedConfig): model_type = "deepseek_vl_v2" vision_config: DeepseekVL2VisionEncoderConfig = None projector_config: DeepseekVL2MlpProjectorConfig = None language_config: DeepseekV2Config = None tile_tag: str = "2D" global_view_pos: str = "head" candidate_resolutions: Tuple[Tuple[int, int]] = ((384, 384),) def __init__( self, tile_tag: str = "tile_tag", global_view_pos: str = "head", candidate_resolutions: Tuple[Tuple[int, int]] = ((384, 384),), **kwargs, ): super().__init__(**kwargs) vision_config = kwargs.get("vision_config", {}) self.vision_config = DeepseekVL2VisionEncoderConfig(**vision_config) projector_config = kwargs.get("projector_config", {}) self.projector_config = DeepseekVL2MlpProjectorConfig(**projector_config) language_config = kwargs.get("language_config", {}) if isinstance(language_config, DeepseekV2Config): self.language_config = language_config else: self.language_config = DeepseekV2Config(**language_config) self.tile_tag = tile_tag self.global_view_pos = global_view_pos self.candidate_resolutions = candidate_resolutions self.architectures = ["DeepseekVL2ForCausalLM"] AutoProcessor.register(DeepseekVL2Config, DeepseekVLV2Processor) ================================================ FILE: python/sglang/srt/configs/device_config.py ================================================ import logging from typing import Optional import torch logger = logging.getLogger(__name__) SUPPORTED_DEVICES = ["cuda", "xpu", "hpu", "cpu", "npu", "musa", "mps"] class DeviceConfig: device: Optional[torch.device] gpu_id: Optional[int] def __init__(self, device: str = "cuda", gpu_id: int = -1) -> None: if device in SUPPORTED_DEVICES: self.device_type = device else: raise RuntimeError(f"Not supported device type: {device}") self.device = torch.device(self.device_type) self.gpu_id = gpu_id ================================================ FILE: python/sglang/srt/configs/dots_ocr.py ================================================ from typing import Optional from transformers import AutoProcessor, Qwen2_5_VLProcessor from transformers.image_processing_utils import BaseImageProcessor from transformers.models.qwen2 import Qwen2Config from sglang.srt.configs.dots_vlm import DotsVisionConfig class DotsOCRConfig(Qwen2Config): model_type = "dots_ocr" def __init__( self, image_token_id=151665, video_token_id=151656, vision_config: Optional[dict] = None, *args, **kwargs ): super().__init__(*args, **kwargs) self.image_token_id = image_token_id self.video_token_id = video_token_id self.vision_config = DotsVisionConfig(**(vision_config or {})) def save_pretrained(self, save_directory, **kwargs): self._auto_class = None super().save_pretrained(save_directory, **kwargs) class DummyVideoProcessor(BaseImageProcessor): model_input_names = ["pixel_values"] def __call__(self, *args, **kwargs): return None class DotsVLProcessor(Qwen2_5_VLProcessor): def __init__( self, image_processor=None, tokenizer=None, video_processor=None, chat_template=None, **kwargs ): if video_processor is None: video_processor = DummyVideoProcessor() super().__init__( image_processor, tokenizer, video_processor, chat_template=chat_template ) self.image_token = ( "<|imgpad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token ) self.image_token_id = ( tokenizer.image_token_id if getattr(tokenizer, "image_token_id", None) is not None else tokenizer.convert_tokens_to_ids(self.image_token) ) AutoProcessor.register(DotsOCRConfig, DotsVLProcessor) ================================================ FILE: python/sglang/srt/configs/dots_vlm.py ================================================ from transformers import AutoProcessor, PretrainedConfig from transformers.processing_utils import ProcessingKwargs try: from transformers import Qwen2_5_VLProcessor except ImportError: raise ImportError( "Qwen2_5_VLProcessor can not be found. Please upgrade your transformers version." ) from sglang.srt.configs.deepseekvl2 import DeepseekV2Config class DotsVisionConfig(PretrainedConfig): model_type: str = "dots_vit" def __init__( self, embed_dim: int = 1536, # vision encoder embed size hidden_size: int = 1536, # after merger hidden size intermediate_size: int = 4224, num_hidden_layers: int = 42, num_attention_heads: int = 12, num_channels: int = 3, patch_size: int = 14, spatial_merge_size: int = 2, temporal_patch_size: int = 1, rms_norm_eps: float = 1e-5, use_bias: bool = False, attn_implementation="flash_attention_2", # "eager","sdpa","flash_attention_2" initializer_range=0.02, init_merger_std=0.02, is_causal=False, # ve causal forward post_norm=True, gradient_checkpointing=False, **kwargs, ): super().__init__(**kwargs) self.embed_dim = embed_dim self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_channels = num_channels self.patch_size = patch_size self.spatial_merge_size = spatial_merge_size self.temporal_patch_size = temporal_patch_size self.rms_norm_eps = rms_norm_eps self.use_bias = use_bias self.attn_implementation = attn_implementation self.initializer_range = initializer_range self.init_merger_std = init_merger_std self.is_causal = is_causal self.post_norm = post_norm self.gradient_checkpointing = gradient_checkpointing class DotsVLMConfig(PretrainedConfig): model_type = "dots_vlm" def __init__(self, **kwargs): super().__init__(**kwargs) vision_config = kwargs.get("vision_config", {}) self.im_span_id = kwargs.get("image_token_id", 128815) self.video_span_id = kwargs.get("video_token_id", 128836) self.vision_config = DotsVisionConfig(**vision_config) self.language_config = DeepseekV2Config(**kwargs) self.architectures = ["DotsVLMForCausalLM"] class DotsVLMProcessorKwargs(ProcessingKwargs, total=False): _defaults = { "text_kwargs": { "padding": False, }, } class DotsVLMProcessor(Qwen2_5_VLProcessor): r""" Constructs a DotsVLM processor which derives from Qwen2_5_VLProcessor, but overrides the image and video token ids. Besides, its tokenizer is a LlamaTokenizerFast instead of Qwen2TokenizerFast. [`DotsVLMProcessor`] offers all the functionalities of [`DotsVisionConfig`] and [`LlamaTokenizerFast`]. See the [`~DotsVLMProcessor.__call__`] and [`~DotsVLMProcessor.decode`] for more information. Args: image_processor ([`Qwen2VLImageProcessor`], *optional*): The image processor is a required input. tokenizer ([`LlamaTokenizerFast`], *optional*): The tokenizer is a required input. chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages in a chat into a tokenizable string. """ attributes = ["image_processor", "tokenizer"] valid_kwargs = ["chat_template"] tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast") def __init__( self, image_processor=None, tokenizer=None, chat_template=None, **kwargs ): super().__init__(image_processor, tokenizer, chat_template=chat_template) self.image_token = ( "<|imgpad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token ) self.video_token = ( "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token ) self.img_token = ( "<|img|>" if not hasattr(tokenizer, "img_token") else tokenizer.img_token ) self.endofimg_token = ( "<|endofimg|>" if not hasattr(tokenizer, "endofimg_token") else tokenizer.endofimg_token ) self.image_token_id = ( tokenizer.image_token_id if getattr(tokenizer, "image_token_id", None) else tokenizer.encode(self.image_token)[0] ) self.video_token_id = ( tokenizer.video_token_id if getattr(tokenizer, "video_token_id", None) else tokenizer.encode(self.video_token)[0] ) AutoProcessor.register(DotsVLMConfig, DotsVLMProcessor) ================================================ FILE: python/sglang/srt/configs/exaone.py ================================================ # coding=utf-8 # Copyright 2024 The LG AI Research EXAONE Lab. All rights reserved. # Copyright 2024 The LG CNS AI Engineering Team. # Copyright 2023-2024 SGLang Team. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """EXAONE model configuration""" from typing import Any, Dict from transformers.configuration_utils import PretrainedConfig from transformers.utils import logging logger = logging.get_logger(__name__) EXAONE_PRETRAINED_CONFIG_ARCHIVE_MAP: Dict[str, Any] = {} # ruff: noqa: E501 class ExaoneConfig(PretrainedConfig): r""" This is the configuration class to store the configuration of a :class:`~transformers.ExaoneModel`. It is used to instantiate a EXAONE model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the Exaone Configuration objects inherit from :class:`~transformers.PretrainedConfig` and can be used to control the model outputs. Read the documentation from :class:`~transformers.PretrainedConfig` for more information. Args: vocab_size (:obj:`int`, `optional`, defaults to 102400): Vocabulary size of the EXAONE model. Defines the number of different tokens that can be represented by the :obj:`inputs_ids` passed when calling :class:`~transformers.ExaoneModel`. Vocabulary size of the model. Defines the different tokens that can be represented by the `inputs_ids` passed to the forward method of :class:`~transformers.EXAONEModel`. max_position_embeddings (:obj:`int`, `optional`, defaults to 2048): The maximum sequence length that this model might ever be used with. Typically set this to something large just in case (e.g., 512 or 1024 or 2048). hidden_size (:obj:`int`, `optional`, defaults to 2048): Dimensionality of the encoder layers and the pooler layer. num_layers (:obj:`int`, `optional`, defaults to 32): Number of hidden layers in the Transformer encoder. num_attention_heads (:obj:`int`, `optional`, defaults to 32): Number of attention heads for each attention layer in the Transformer decoder. num_key_value_heads (:obj:`int`, `optional`): This is the number of key_value heads that should be used to implement Grouped Query Attention. If `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed by meanpooling all the original heads within that group. For more details checkout [this paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `num_attention_heads`. intermediate_size (:obj:`int`, `optional`, defaults to `hidden_size * 4`): Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder. activation_function (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"silu"`): The non-linear activation function (function or string) in the decoder. rope_theta (:obj:`float`, `optional`, defaults to 10000.0): The base period of the RoPE embeddings. rope_scaling (:obj:`Dict`, `optional`): Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value accordingly. Expected contents: `rope_type` (:obj:`str`): The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope', 'llama3'], with 'default' being the original RoPE implementation. `factor` (:obj:`float`, `optional`): Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In most scaling types, a `factor` of x will enable the model to handle sequences of length x * original maximum pre-trained length. `original_max_position_embeddings` (:obj:`int`, `optional`): Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during pretraining. `attention_factor` (:obj:`float`, `optional`): Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention computation. If unspecified, it defaults to value recommended by the implementation, using the `factor` field to infer the suggested value. `beta_fast` (:obj:`float`, `optional`): Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear ramp function. If unspecified, it defaults to 32. `beta_slow` (:obj:`float`, `optional`): Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear ramp function. If unspecified, it defaults to 1. `short_factor` (:obj:`List[float]`, `optional`): Only used with 'longrope'. The scaling factor to be applied to short contexts (< `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden size divided by the number of attention heads divided by 2 `long_factor` (:obj:`List[float]`, `optional`): Only used with 'longrope'. The scaling factor to be applied to long contexts (< `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden size divided by the number of attention heads divided by 2 `low_freq_factor` (:obj:`float`, `optional`): Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE `high_freq_factor` (:obj:`float`, `optional`): Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE embed_dropout (:obj:`float`, `optional`, defaults to 0.0): The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler. attention_dropout (:obj:`float`, `optional`, defaults to 0.0): The dropout ratio for the attention probabilities. layer_norm_epsilon (:obj:`float`, `optional`, defaults to 1e-5): The epsilon used by the layer normalization layers. initializer_range (:obj:`float`, `optional`, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices. use_cache (:obj:`bool`, `optional`, defaults to :obj:`True`): Whether or not the model should return the last key/values attentions (not used by all models). Only relevant if ``configs.is_decoder=True``. bos_token_id (:obj:`int`, `optional`, defaults to 0): Beginning of stream token id. eos_token_id (:obj:`int`, `optional`, defaults to 2): End of stream token id. tie_word_embeddings (:obj:`bool`, `optional`, defaults to :obj:`True`): Whether to tie weight embeddings gradient_checkpointing (:obj:`bool`, `optional`, defaults to :obj:`False`): If True, use gradient checkpointing to save memory at the expense of slower backward pass. Example:: >>> from transformers import EXAONEModel, ExaoneConfig >>> # Initializing a EXAONE configuration >>> configuration = ExaoneConfig() >>> # Initializing a model from configuration >>> model = EXAONEModel(configuration) >>> # Accessing the model configuration >>> configuration = model.configs """ model_type = "exaone" keys_to_ignore_at_inference = ["past_key_values"] attribute_map = {"num_hidden_layers": "num_layers"} def __init__( self, vocab_size=102400, max_position_embeddings=2048, hidden_size=2048, num_layers=32, num_attention_heads=32, num_key_value_heads=None, intermediate_size=None, activation_function="silu", rope_theta=10000.0, rope_scaling=None, embed_dropout=0.0, attention_dropout=0.0, layer_norm_epsilon=1e-5, initializer_range=0.02, use_cache=True, bos_token_id=0, eos_token_id=2, tie_word_embeddings=True, **kwargs ): self.vocab_size = vocab_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.num_layers = num_layers self.num_attention_heads = num_attention_heads self.num_hidden_layers = num_layers if num_key_value_heads is None: num_key_value_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads if intermediate_size: self.intermediate_size = intermediate_size else: self.intermediate_size = hidden_size * 4 self.activation_function = activation_function self.embed_dropout = embed_dropout self.attention_dropout = attention_dropout self.layer_norm_epsilon = layer_norm_epsilon self.initializer_range = initializer_range self.use_cache = use_cache self.rope_theta = rope_theta self.rope_scaling = rope_scaling self.bos_token_id = bos_token_id self.eos_token_id = eos_token_id super().__init__( bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs ) ================================================ FILE: python/sglang/srt/configs/falcon_h1.py ================================================ # coding=utf-8 # Copyright 2024 TII and the HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Falcon-H1 model configuration""" from transformers.configuration_utils import PretrainedConfig from transformers.utils import logging from sglang.srt.configs.mamba_utils import ( Mamba2CacheParams, Mamba2StateShape, mamba2_state_dtype, ) logger = logging.get_logger(__name__) class FalconH1Config(PretrainedConfig): r""" This is the configuration class to store the configuration of a [`FalconH1Model`]. It is used to instantiate a FalconH1Model model according to the specified arguments, defining the model architecture. Instantiating a configuration with defaults taken from [ibm-fms/FalconH1-9.8b-2.2T-hf](https://huggingface.co/ibm-fms/FalconH1-9.8b-2.2T-hf). The FalconH1Model is a hybrid [mamba2](https://github.com/state-spaces/mamba) architecture with SwiGLU. The checkpoints are jointly trained by IBM, Princeton, and UIUC. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. Args: vocab_size (`int`, *optional*, defaults to 128000): Vocabulary size of the FalconH1 model. Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling [`FalconH1Model`] tie_word_embeddings (`bool`, *optional*, defaults to `False`): Whether the model's input and output word embeddings should be tied. Note that this is only relevant if the model has a output word embedding layer. hidden_size (`int`, *optional*, defaults to 4096): Dimension of the hidden representations. intermediate_size (`int`, *optional*, defaults to 14336): Dimension of the MLP representations. num_hidden_layers (`int`, *optional*, defaults to 32): Number of hidden layers in the Transformer encoder. num_attention_heads (`int`, *optional*, defaults to 32): Number of attention heads for each attention layer in the Transformer encoder. num_key_value_heads (`int`, *optional*, defaults to 8): This is the number of key_value heads that should be used to implement Grouped Query Attention. If `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed by meanpooling all the original heads within that group. For more details, check out [this paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `8`. hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): The non-linear activation function (function or string) in the decoder. initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices. rms_norm_eps (`float`, *optional*, defaults to 1e-05): The epsilon used by the rms normalization layers. use_cache (`bool`, *optional*, defaults to `True`): Whether or not the model should return the last key/values attentions (not used by all models). Only relevant if `config.is_decoder=True`. num_logits_to_keep (`int` or `None`, *optional*, defaults to 1): Number of prompt logits to calculate during generation. If `None`, all logits will be calculated. If an integer value, only last `num_logits_to_keep` logits will be calculated. Default is 1 because only the logits of the last prompt token are needed for generation. For long sequences, the logits for the entire sequence may use a lot of memory so, setting `num_logits_to_keep=1` will reduce memory footprint significantly. pad_token_id (`int`, *optional*, defaults to 0): The id of the padding token. bos_token_id (`int`, *optional*, defaults to 1): The id of the "beginning-of-sequence" token. eos_token_id (`int`, *optional*, defaults to 2): The id of the "end-of-sequence" token. max_position_embeddings (`int`, *optional*, defaults to 8192): Max cached sequence length for the model attention_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the attention probabilities. mamba_d_ssm (`int`, *optional*, defaults to 1024): The dimension of the SSM state space latents. mamba_n_heads (`int`, *optional*, defaults to 128): The number of mamba heads used in the v2 implementation. mamba_d_head (`int`, *optional*, defaults to `"auto"`): Head embedding dimension size mamba_n_groups (`int`, *optional*, defaults to 1): The number of the mamba groups used in the v2 implementation. mamba_d_state (`int`, *optional*, defaults to 256): The dimension the mamba state space latents mamba_d_conv (`int`, *optional*, defaults to 4): The size of the mamba convolution kernel mamba_expand (`int`, *optional*, defaults to 2): Expanding factor (relative to hidden_size) used to determine the mamba intermediate size mamba_chunk_size (`int`, *optional*, defaults to 256): The chunks in which to break the sequence when doing prefill/training mamba_conv_bias (`bool`, *optional*, defaults to `True`): Flag indicating whether or not to use bias in the convolution layer of the mamba mixer block. mamba_proj_bias (`bool`, *optional*, defaults to `False`): Flag indicating whether or not to use bias in the input and output projections (["in_proj", "out_proj"]) of the mamba mixer block mamba_norm_before_gate (`bool`, *optional*, defaults to `True`): Whether to use RMSNorm before the gate in the Mamba block mamba_rms_norm (`bool`, *optional*, defaults to `False`): Whether to use RMSNorm instead of LayerNorm in the Mamba block projectors_bias (`bool`, *optional*, defaults to `False`): Flag indicating whether or not to use bias in the input and output projections (["in_proj", "out_proj"]) of the attention block rope_theta (`float`, *optional*, defaults to 100000.0): The theta value used for the RoPE embeddings. rope_scaling (`float`, *optional*): The scaling value used for the RoPE embeddings. If `None`, no scaling is applied. lm_head_multiplier (`float`, *optional*, defaults to 1.0): The multiplier for the LM head. This is used to scale the output of the LM head. embedding_multiplier (`float`, *optional*, defaults to 1.0): The multiplier for the embedding layer. This is used to scale the output of the embedding layer. mlp_multipliers (`list[float]`, *optional*): The multipliers for the MLP layers. This is used to scale the output of the MLP layers. The first value is the multiplier of gate layer, the second value is the multiplier of the down_proj layer. key_multiplier (`float`, *optional*): The multiplier for the key layer. This is used to scale the output of the key layer. attention_out_multiplier (`float`, *optional*): The multiplier for the attention output layer. This is used to scale the output of the attention output attention_in_multiplier (`float`, *optional*): The multiplier for the attention input layer. This is used to scale the output of the attention input layer. ssm_multipliers (`list[float]`, *optional*): The multipliers for the SSM layers. This is used to scale the output of the SSM layers. ssm_in_multiplier (`float`, *optional*): The multiplier for the SSM input layer. This is used to scale the output of the SSM input layer. ssm_out_multiplier (`float`, *optional*): The multiplier for the SSM output layer. This is used to scale the output of the SSM output layer. """ model_type = "falcon_h1" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size=128000, tie_word_embeddings=False, hidden_size=4096, intermediate_size=14336, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=8, hidden_act="silu", initializer_range=0.02, rms_norm_eps=1e-5, use_cache=True, num_logits_to_keep=1, pad_token_id=0, bos_token_id=1, eos_token_id=2, max_position_embeddings=8192, attention_dropout=0.0, mamba_d_ssm=1024, mamba_n_heads=128, mamba_d_head="auto", mamba_n_groups=1, mamba_d_state=256, mamba_d_conv=4, mamba_expand=2, mamba_chunk_size=256, mamba_conv_bias=True, mamba_proj_bias=False, mamba_norm_before_gate=True, mamba_rms_norm=False, projectors_bias=False, rope_theta=100000.0, rope_scaling=None, lm_head_multiplier=1.0, embedding_multiplier=1.0, mlp_multipliers=None, key_multiplier=None, attention_out_multiplier=None, attention_in_multiplier=None, ssm_multipliers=None, ssm_in_multiplier=None, ssm_out_multiplier=None, **kwargs, ): self.vocab_size = vocab_size self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.max_position_embeddings = max_position_embeddings self.attention_dropout = attention_dropout self.attention_bias = False self.mlp_bias = False # for backward compatibility if num_key_value_heads is None: num_key_value_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.hidden_act = hidden_act self.initializer_range = initializer_range self.rms_norm_eps = rms_norm_eps self.use_cache = use_cache self.num_logits_to_keep = num_logits_to_keep self.rope_theta = rope_theta self.rope_scaling = None self.rope_scaling = rope_scaling self.projectors_bias = projectors_bias self.mamba_intermediate = mamba_intermediate = ( mamba_expand * hidden_size if mamba_d_ssm is None else mamba_d_ssm ) if mamba_intermediate % mamba_n_heads != 0: raise ValueError("mamba_n_heads must divide mamba_expand * hidden_size") # for the mamba_v2, must satisfy the following if mamba_d_head == "auto": mamba_d_head = mamba_intermediate // mamba_n_heads if mamba_d_head * mamba_n_heads != mamba_intermediate: raise ValueError( "The dimensions for the Mamba head state do not match the model intermediate_size" ) self.mamba_d_ssm = mamba_d_ssm self.mamba_n_heads = mamba_n_heads self.mamba_d_head = mamba_d_head self.mamba_n_groups = mamba_n_groups self.mamba_d_state = mamba_d_state self.mamba_d_conv = mamba_d_conv self.mamba_expand = mamba_expand self.mamba_chunk_size = mamba_chunk_size self.mamba_conv_bias = mamba_conv_bias self.mamba_proj_bias = mamba_proj_bias self.mamba_norm_before_gate = mamba_norm_before_gate self.mamba_rms_norm = mamba_rms_norm self.lm_head_multiplier = lm_head_multiplier self.embedding_multiplier = embedding_multiplier if mlp_multipliers is not None: self.mlp_multipliers = mlp_multipliers else: self.mlp_multipliers = [1.0, 1.0] if attention_out_multiplier is not None: self.attention_out_multiplier = attention_out_multiplier else: self.attention_out_multiplier = 1.0 if attention_in_multiplier is not None: self.attention_in_multiplier = attention_in_multiplier else: self.attention_in_multiplier = 1.0 if key_multiplier is not None: self.key_multiplier = key_multiplier else: self.key_multiplier = 1.0 if ssm_multipliers is not None: self.ssm_multipliers = ssm_multipliers else: self.ssm_multipliers = [1.0, 1.0, 1.0, 1.0, 1.0] if ssm_in_multiplier is not None: self.ssm_in_multiplier = ssm_in_multiplier else: self.ssm_in_multiplier = 1.0 if ssm_out_multiplier is not None: self.ssm_out_multiplier = ssm_out_multiplier else: self.ssm_out_multiplier = 1.0 super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs, ) @property def layers_block_type(self): return ["falcon_h1" for i in range(self.num_hidden_layers)] @property def full_attention_layer_ids(self): # For Falcon-H1, we do have attention on all layers return range(self.num_hidden_layers) @property def linear_layer_ids(self): # For Falcon-H1, we do have mamba on all layers return range(self.num_hidden_layers) @property def mamba2_cache_params(self): from sglang.srt.layers.dp_attention import get_attention_tp_size shape = Mamba2StateShape.create( tp_world_size=get_attention_tp_size(), intermediate_size=self.mamba_intermediate, n_groups=self.mamba_n_groups, num_heads=self.mamba_n_heads, head_dim=self.mamba_d_head, state_size=self.mamba_d_state, conv_kernel=self.mamba_d_conv, ) return Mamba2CacheParams( shape=shape, layers=self.linear_layer_ids, dtype=mamba2_state_dtype(self) ) ================================================ FILE: python/sglang/srt/configs/granitemoehybrid.py ================================================ # coding=utf-8 # Copyright 2025 IBM and the HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """GraniteMoeHybrid model configuration""" from transformers.configuration_utils import PretrainedConfig from transformers.utils import logging from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape logger = logging.get_logger(__name__) MAMBA = "mamba" ATTENTION = "attention" class GraniteMoeHybridConfig(PretrainedConfig): r""" This is the configuration class to store the configuration of a [`GraniteMoeHybridModel`]. It is used to instantiate a GraniteMoeHybrid model according to the specified arguments, defining the model architecture. The GraniteMoeHybrid is a hybrid architecture combining Mamba2 layers with attention layers, developed by IBM. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. Args: vocab_size (`int`, *optional*, defaults to 100352): Vocabulary size of the GraniteMoeHybrid model. Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling [`GraniteMoeHybridModel`] tie_word_embeddings (`bool`, *optional*, defaults to `True`): Whether the model's input and output word embeddings should be tied. Note that this is only relevant if the model has a output word embedding layer. hidden_size (`int`, *optional*, defaults to 2048): Dimension of the hidden representations. intermediate_size (`int`, *optional*, defaults to 8192): Dimension of the MLP representations. num_hidden_layers (`int`, *optional*, defaults to 40): Number of hidden layers in the model. layer_types (`list[str]`, *optional*): List of layer types for each layer. Each element should be either "mamba" or "attention". If not provided, defaults to alternating pattern based on num_hidden_layers. num_attention_heads (`int`, *optional*, defaults to 32): Number of attention heads for each attention layer. num_key_value_heads (`int`, *optional*, defaults to 8): This is the number of key_value heads that should be used to implement Grouped Query Attention. If `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): The non-linear activation function (function or string) in the decoder. initializer_range (`float`, *optional*, defaults to 0.1): The standard deviation of the truncated_normal_initializer for initializing all weight matrices. rms_norm_eps (`float`, *optional*, defaults to 1e-05): The epsilon used by the rms normalization layers. normalization_function (`str`, *optional*, defaults to `"rmsnorm"`): The normalization function to use. Currently only "rmsnorm" is supported. use_cache (`bool`, *optional*, defaults to `True`): Whether or not the model should return the last key/values attentions (not used by all models). Only relevant if `config.is_decoder=True`. pad_token_id (`int`, *optional*, defaults to 100256): The id of the padding token. bos_token_id (`int`, *optional*, defaults to 100257): The id of the "beginning-of-sequence" token. eos_token_id (`int`, *optional*, defaults to 100257): The id of the "end-of-sequence" token. max_position_embeddings (`int`, *optional*, defaults to 131072): Max cached sequence length for the model attention_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the attention probabilities. attention_bias (`bool`, *optional*, defaults to `False`): Whether to use bias in attention layers. position_embedding_type (`str`, *optional*, defaults to `"nope"`): Type of position embedding. Can be "nope" (no position embedding) or "rope". rope_theta (`float`, *optional*, defaults to 10000.0): The theta value used for the RoPE embeddings. rope_scaling (`dict`, *optional*): The scaling configuration for the RoPE embeddings. If `None`, no scaling is applied. mamba_d_state (`int`, *optional*, defaults to 128): The dimension of the mamba state space latents mamba_d_conv (`int`, *optional*, defaults to 4): The size of the mamba convolution kernel mamba_expand (`int`, *optional*, defaults to 2): Expanding factor (relative to hidden_size) used to determine the mamba intermediate size mamba_d_head (`int`, *optional*, defaults to 64): Head embedding dimension size for Mamba mamba_n_heads (`int`, *optional*, defaults to 64): The number of mamba heads mamba_n_groups (`int`, *optional*, defaults to 1): The number of the mamba groups mamba_chunk_size (`int`, *optional*, defaults to 256): The chunks in which to break the sequence when doing prefill/training mamba_conv_bias (`bool`, *optional*, defaults to `True`): Flag indicating whether or not to use bias in the convolution layer of the mamba mixer block. mamba_proj_bias (`bool`, *optional*, defaults to `False`): Flag indicating whether or not to use bias in the input and output projections of the mamba mixer block embedding_multiplier (`float`, *optional*, defaults to 12.0): The multiplier for the embedding layer. This is used to scale the output of the embedding layer. logits_scaling (`float`, *optional*, defaults to 8.0): The scaling factor for the logits. attention_multiplier (`float`, *optional*, defaults to 0.015625): The multiplier for the attention layers. residual_multiplier (`float`, *optional*, defaults to 0.22): The multiplier for residual connections. num_local_experts (`int`, *optional*, defaults to 0): Number of local experts in MoE layers. num_experts_per_tok (`int`, *optional*, defaults to 0): Number of experts to use per token in MoE layers. shared_intermediate_size (`int`, *optional*, defaults to 8192): Intermediate size for shared experts. output_router_logits (`bool`, *optional*, defaults to `False`): Whether to output router logits. router_aux_loss_coef (`float`, *optional*, defaults to 0.01): Auxiliary loss coefficient for the router. """ model_type = "granitemoehybrid" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size=100352, tie_word_embeddings=True, hidden_size=2048, intermediate_size=8192, num_hidden_layers=40, layer_types=None, num_attention_heads=32, num_key_value_heads=8, hidden_act="silu", initializer_range=0.1, rms_norm_eps=1e-5, normalization_function="rmsnorm", use_cache=True, pad_token_id=100256, bos_token_id=100257, eos_token_id=100257, max_position_embeddings=131072, attention_dropout=0.0, attention_bias=False, position_embedding_type="nope", rope_theta=10000.0, rope_scaling=None, mamba_d_state=128, mamba_d_conv=4, mamba_expand=2, mamba_d_head=64, mamba_n_heads=64, mamba_n_groups=1, mamba_chunk_size=256, mamba_conv_bias=True, mamba_proj_bias=False, embedding_multiplier=12.0, logits_scaling=8.0, attention_multiplier=0.015625, residual_multiplier=0.22, num_local_experts=0, num_experts_per_tok=0, shared_intermediate_size=8192, output_router_logits=False, router_aux_loss_coef=0.01, **kwargs, ): self.vocab_size = vocab_size self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers # Set layer types - if not provided, create default pattern if layer_types is None: # Default pattern: mamba layers with attention every 6th layer (roughly) self.layer_types = [] for i in range(num_hidden_layers): if (i + 1) % 6 == 0: self.layer_types.append(ATTENTION) else: self.layer_types.append(MAMBA) else: self.layer_types = layer_types # Validate layer_types if len(self.layer_types) != self.num_hidden_layers: raise ValueError( f"layer_types must have length equal to num_hidden_layers ({num_hidden_layers}), " f"but got {len(self.layer_types)}" ) for layer_type in self.layer_types: if layer_type not in [MAMBA, ATTENTION]: raise ValueError( f"Each element in layer_types must be either '{MAMBA}' or '{ATTENTION}', " f"but got '{layer_type}'" ) self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.hidden_act = hidden_act self.initializer_range = initializer_range self.rms_norm_eps = rms_norm_eps self.normalization_function = normalization_function self.use_cache = use_cache self.max_position_embeddings = max_position_embeddings self.attention_dropout = attention_dropout self.attention_bias = attention_bias self.position_embedding_type = position_embedding_type self.rope_theta = rope_theta self.rope_scaling = rope_scaling # Mamba configuration self.mamba_d_state = mamba_d_state self.mamba_d_conv = mamba_d_conv self.mamba_expand = mamba_expand self.mamba_d_head = mamba_d_head self.mamba_n_heads = mamba_n_heads self.mamba_n_groups = mamba_n_groups self.mamba_chunk_size = mamba_chunk_size self.mamba_conv_bias = mamba_conv_bias self.mamba_proj_bias = mamba_proj_bias # Calculate mamba intermediate size self.mamba_intermediate_size = mamba_expand * hidden_size # Validate mamba configuration if self.mamba_intermediate_size % mamba_n_heads != 0: raise ValueError( f"mamba_intermediate_size ({self.mamba_intermediate_size}) must be divisible by " f"mamba_n_heads ({mamba_n_heads})" ) if mamba_d_head * mamba_n_heads != self.mamba_intermediate_size: raise ValueError( f"mamba_d_head ({mamba_d_head}) * mamba_n_heads ({mamba_n_heads}) must equal " f"mamba_intermediate_size ({self.mamba_intermediate_size})" ) # Scaling factors self.embedding_multiplier = embedding_multiplier self.logits_scaling = logits_scaling self.attention_multiplier = attention_multiplier self.residual_multiplier = residual_multiplier # MoE configuration self.num_local_experts = num_local_experts self.num_experts_per_tok = num_experts_per_tok self.shared_intermediate_size = shared_intermediate_size self.output_router_logits = output_router_logits self.router_aux_loss_coef = router_aux_loss_coef super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs, ) @property def mamba_layer_ids(self): """Returns the indices of layers that are Mamba layers.""" return [ i for i in range(self.num_hidden_layers) if self.layer_types[i] == MAMBA ] @property def attention_layer_ids(self): """Returns the indices of layers that are attention layers.""" return [ i for i in range(self.num_hidden_layers) if self.layer_types[i] == ATTENTION ] @property def full_attention_layer_ids(self): """Alias for attention_layer_ids for compatibility.""" return self.attention_layer_ids @property def mamba2_cache_params(self): """Returns the Mamba2 cache parameters for this configuration.""" from sglang.srt.layers.dp_attention import get_attention_tp_size shape = Mamba2StateShape.create( tp_world_size=get_attention_tp_size(), intermediate_size=self.mamba_intermediate_size, n_groups=self.mamba_n_groups, num_heads=self.mamba_n_heads, head_dim=self.mamba_d_head, state_size=self.mamba_d_state, conv_kernel=self.mamba_d_conv, ) return Mamba2CacheParams(shape=shape, layers=self.mamba_layer_ids) ================================================ FILE: python/sglang/srt/configs/internvl.py ================================================ import copy import os from shutil import copyfile from typing import Any, Dict, List, Optional, Tuple, Union import sentencepiece as spm from transformers import ( TOKENIZER_MAPPING, GptOssConfig, LlamaConfig, PretrainedConfig, PreTrainedTokenizer, Qwen2Config, Qwen3Config, Qwen3MoeConfig, ) from sglang.utils import logger # Copied from: https://github.com/OpenGVLab/InternVL/blob/34a81000402bf8f716bab8c9b57aff1f6b436bd0/internvl_chat/internvl/model/internvl_chat/configuration_internvl_chat.py#L21 VOCAB_FILES_NAMES = {"vocab_file": "./tokenizer.model"} PRETRAINED_VOCAB_FILES_MAP = {} # Modified from transformers.model.llama.configuration_llama.LlamaConfig class InternLM2Config(PretrainedConfig): r""" This is the configuration class to store the configuration of a [`InternLM2Model`]. It is used to instantiate an InternLM2 model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the InternLM2-7B. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. Args: vocab_size (`int`, *optional*, defaults to 32000): Vocabulary size of the InternLM2 model. Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling [`InternLM2Model`] hidden_size (`int`, *optional*, defaults to 4096): Dimension of the hidden representations. intermediate_size (`int`, *optional*, defaults to 11008): Dimension of the MLP representations. num_hidden_layers (`int`, *optional*, defaults to 32): Number of hidden layers in the Transformer encoder. num_attention_heads (`int`, *optional*, defaults to 32): Number of attention heads for each attention layer in the Transformer encoder. num_key_value_heads (`int`, *optional*): This is the number of key_value heads that should be used to implement Grouped Query Attention. If `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed by meanpooling all the original heads within that group. For more details checkout [this paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `num_attention_heads`. hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): The non-linear activation function (function or string) in the decoder. max_position_embeddings (`int`, *optional*, defaults to 2048): The maximum sequence length that this model might ever be used with. Typically set this to something large just in case (e.g., 512 or 1024 or 2048). initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices. rms_norm_eps (`float`, *optional*, defaults to 1e-12): The epsilon used by the rms normalization layers. use_cache (`bool`, *optional*, defaults to `True`): Whether or not the model should return the last key/values attentions (not used by all models). Only relevant if `config.is_decoder=True`. tie_word_embeddings(`bool`, *optional*, defaults to `False`): Whether to tie weight embeddings Example: """ model_type = "internlm2" _auto_class = "AutoConfig" def __init__( # pylint: disable=W0102 self, vocab_size=103168, hidden_size=4096, intermediate_size=11008, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=None, hidden_act="silu", max_position_embeddings=2048, initializer_range=0.02, rms_norm_eps=1e-6, use_cache=True, pad_token_id=0, bos_token_id=1, eos_token_id=2, tie_word_embeddings=False, bias=True, rope_theta=10000, rope_scaling=None, attn_implementation="eager", **kwargs, ): self.vocab_size = vocab_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.bias = bias if num_key_value_heads is None: num_key_value_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.hidden_act = hidden_act self.initializer_range = initializer_range self.rms_norm_eps = rms_norm_eps self.use_cache = use_cache self.rope_theta = rope_theta self.rope_scaling = rope_scaling self._rope_scaling_validation() self.attn_implementation = attn_implementation if self.attn_implementation is None: self.attn_implementation = "eager" super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs, ) def _rope_scaling_validation(self): """ Validate the `rope_scaling` configuration. """ if self.rope_scaling is None: return if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2: raise ValueError( "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, " f"got {self.rope_scaling}" ) rope_scaling_type = self.rope_scaling.get("type", None) rope_scaling_factor = self.rope_scaling.get("factor", None) if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]: raise ValueError( f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}" ) if ( rope_scaling_factor is None or not isinstance(rope_scaling_factor, (float, int)) or rope_scaling_factor < 1.0 ): raise ValueError( f"`rope_scaling`'s factor field must be a float|int >= 1, got {rope_scaling_factor=}, {type(rope_scaling_factor)=}" ) if isinstance(rope_scaling_factor, int): rope_scaling_factor = float(rope_scaling_factor) class InternVisionConfig(PretrainedConfig): r""" This is the configuration class to store the configuration of a [`InternVisionModel`]. It is used to instantiate a vision encoder according to the specified arguments, defining the model architecture. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. Args: num_channels (`int`, *optional*, defaults to 3): Number of color channels in the input images (e.g., 3 for RGB). patch_size (`int`, *optional*, defaults to 14): The size (resolution) of each patch. image_size (`int`, *optional*, defaults to 224): The size (resolution) of each image. qkv_bias (`bool`, *optional*, defaults to `False`): Whether to add a bias to the queries and values in the self-attention layers. hidden_size (`int`, *optional*, defaults to 3200): Dimensionality of the encoder layers and the pooler layer. num_attention_heads (`int`, *optional*, defaults to 25): Number of attention heads for each attention layer in the Transformer encoder. intermediate_size (`int`, *optional*, defaults to 12800): Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder. qk_normalization (`bool`, *optional*, defaults to `True`): Whether to normalize the queries and keys in the self-attention layers. num_hidden_layers (`int`, *optional*, defaults to 48): Number of hidden layers in the Transformer encoder. use_flash_attn (`bool`, *optional*, defaults to `True`): Whether to use flash attention mechanism. hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`): The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`, `"relu"`, `"selu"` and `"gelu_new"` ``"gelu"` are supported. layer_norm_eps (`float`, *optional*, defaults to 1e-6): The epsilon used by the layer normalization layers. dropout (`float`, *optional*, defaults to 0.0): The dropout probability for all fully connected layers in the embeddings, encoder, and pooler. drop_path_rate (`float`, *optional*, defaults to 0.0): Dropout rate for stochastic depth. attention_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the attention probabilities. initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices. initializer_factor (`float`, *optional*, defaults to 0.1): A factor for layer scale. """ model_type = "intern_vit_6b" def __init__( self, num_channels=3, patch_size=14, image_size=224, qkv_bias=False, hidden_size=3200, num_attention_heads=25, intermediate_size=12800, qk_normalization=True, num_hidden_layers=48, use_flash_attn=True, hidden_act="gelu", layer_norm_eps=1e-6, dropout=0.0, drop_path_rate=0.0, attention_dropout=0.0, initializer_range=0.02, initializer_factor=0.1, **kwargs, ): super().__init__(**kwargs) self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.dropout = dropout self.drop_path_rate = drop_path_rate self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_channels = num_channels self.patch_size = patch_size self.image_size = image_size self.initializer_range = initializer_range self.initializer_factor = initializer_factor self.attention_dropout = attention_dropout self.layer_norm_eps = layer_norm_eps self.hidden_act = hidden_act self.qkv_bias = qkv_bias self.qk_normalization = qk_normalization self.use_flash_attn = use_flash_attn @classmethod def from_pretrained( cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs ) -> "PretrainedConfig": config_dict, kwargs = cls.get_config_dict( pretrained_model_name_or_path, **kwargs ) if "vision_config" in config_dict: config_dict = config_dict["vision_config"] if ( "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type ): logger.warning( f"You are using a model of type {config_dict['model_type']} to instantiate a model of type " f"{cls.model_type}. This is not supported for all configurations of models and can yield errors." ) return cls.from_dict(config_dict, **kwargs) class InternVLChatConfig(PretrainedConfig): model_type = "internvl_chat" is_composition = True def __init__( self, vision_config=None, llm_config=None, use_backbone_lora=0, use_llm_lora=0, pad2square=False, select_layer=-1, force_image_size=None, downsample_ratio=0.5, template=None, dynamic_image_size=False, use_thumbnail=False, ps_version="v1", min_dynamic_patch=1, max_dynamic_patch=6, **kwargs, ): super().__init__(**kwargs) if vision_config is None: vision_config = {"architectures": ["InternVisionModel"]} logger.info( "vision_config is None. Initializing the InternVisionConfig with default values." ) if llm_config is None: llm_config = {"architectures": ["InternLM2ForCausalLM"]} logger.info( "llm_config is None. Initializing the LlamaConfig config with default values (`LlamaConfig`)." ) self.vision_config = InternVisionConfig(**vision_config) if llm_config.get("architectures")[0] == "LlamaForCausalLM": self.llm_config = LlamaConfig(**llm_config) elif llm_config.get("architectures")[0] == "InternLM2ForCausalLM": self.llm_config = InternLM2Config(**llm_config) elif llm_config.get("architectures")[0] == "Qwen2ForCausalLM": self.llm_config = Qwen2Config(**llm_config) elif llm_config.get("architectures")[0] == "Qwen3MoeForCausalLM": self.llm_config = Qwen3MoeConfig(**llm_config) elif llm_config.get("architectures")[0] == "Qwen3ForCausalLM": self.llm_config = Qwen3Config(**llm_config) elif llm_config.get("architectures")[0] == "GptOssForCausalLM": self.llm_config = GptOssConfig(**llm_config) else: raise ValueError( "Unsupported architecture: {}".format( llm_config.get("architectures")[0] ) ) self.use_backbone_lora = use_backbone_lora self.use_llm_lora = use_llm_lora self.pad2square = pad2square self.select_layer = select_layer self.force_image_size = force_image_size self.downsample_ratio = downsample_ratio self.template = template self.dynamic_image_size = dynamic_image_size self.use_thumbnail = use_thumbnail self.ps_version = ps_version # pixel shuffle version self.min_dynamic_patch = min_dynamic_patch self.max_dynamic_patch = max_dynamic_patch self.hidden_size = self.llm_config.hidden_size # By default, we use tie_word_embeddings=False for models of all sizes. self.tie_word_embeddings = False self.llm_config.tie_word_embeddings = self.tie_word_embeddings def to_dict(self): """ Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`]. Returns: `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance, """ output = copy.deepcopy(self.__dict__) output["vision_config"] = self.vision_config.to_dict() output["llm_config"] = self.llm_config.to_dict() output["model_type"] = self.__class__.model_type output["use_backbone_lora"] = self.use_backbone_lora output["use_llm_lora"] = self.use_llm_lora output["select_layer"] = self.select_layer output["force_image_size"] = self.force_image_size output["downsample_ratio"] = self.downsample_ratio output["template"] = self.template output["dynamic_image_size"] = self.dynamic_image_size output["use_thumbnail"] = self.use_thumbnail output["ps_version"] = self.ps_version output["min_dynamic_patch"] = self.min_dynamic_patch output["max_dynamic_patch"] = self.max_dynamic_patch return output # # Modified from transformers.model.llama.tokenization_llama_fast.LlamaTokenizerFast -> InternLM2TokenizerFast # class InternLM2TokenizerFast(PreTrainedTokenizerFast): # vocab_files_names = VOCAB_FILES_NAMES # slow_tokenizer_class = InternLM2Tokenizer # padding_side = 'left' # model_input_names = ['input_ids', 'attention_mask'] # _auto_class = 'AutoTokenizer' # # def __init__( # self, # vocab_file, # unk_token='', # bos_token='', # eos_token='', # pad_token='', # sp_model_kwargs: Optional[Dict[str, Any]] = None, # add_bos_token=True, # add_eos_token=False, # decode_with_prefix_space=False, # clean_up_tokenization_spaces=False, # **kwargs, # ): # super().__init__( # vocab_file=vocab_file, # unk_token=unk_token, # bos_token=bos_token, # eos_token=eos_token, # pad_token=pad_token, # sp_model_kwargs=sp_model_kwargs, # add_bos_token=add_bos_token, # add_eos_token=add_eos_token, # decode_with_prefix_space=decode_with_prefix_space, # clean_up_tokenization_spaces=clean_up_tokenization_spaces, # **kwargs, # ) # self._add_bos_token = add_bos_token # self._add_eos_token = add_eos_token # self.update_post_processor() # self.vocab_file = vocab_file # # @property # def can_save_slow_tokenizer(self) -> bool: # return os.path.isfile(self.vocab_file) if self.vocab_file else False # # def update_post_processor(self): # """ # Updates the underlying post processor with the current `bos_token` and `eos_token`. # """ # bos = self.bos_token # bos_token_id = self.bos_token_id # if bos is None and self.add_bos_token: # raise ValueError('add_bos_token = True but bos_token = None') # # eos = self.eos_token # eos_token_id = self.eos_token_id # if eos is None and self.add_eos_token: # raise ValueError('add_eos_token = True but eos_token = None') # # single = f"{(bos + ':0 ') if self.add_bos_token else ''}$A:0{(' ' + eos + ':0') if self.add_eos_token else ''}" # pair = f"{single}{(' ' + bos + ':1') if self.add_bos_token else ''} $B:1{(' ' + eos + ':1') if self.add_eos_token else ''}" # # special_tokens = [] # if self.add_bos_token: # special_tokens.append((bos, bos_token_id)) # if self.add_eos_token: # special_tokens.append((eos, eos_token_id)) # self._tokenizer.post_processor = processors.TemplateProcessing( # single=single, pair=pair, special_tokens=special_tokens # ) # # @property # def add_eos_token(self): # return self._add_eos_token # # @property # def add_bos_token(self): # return self._add_bos_token # # @add_eos_token.setter # def add_eos_token(self, value): # self._add_eos_token = value # self.update_post_processor() # # @add_bos_token.setter # def add_bos_token(self, value): # self._add_bos_token = value # self.update_post_processor() # # def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]: # if not self.can_save_slow_tokenizer: # raise ValueError( # 'Your fast tokenizer does not have the necessary information to save the vocabulary for a slow ' # 'tokenizer.' # ) # # if not os.path.isdir(save_directory): # logger.error(f'Vocabulary path ({save_directory}) should be a directory') # return # out_vocab_file = os.path.join( # save_directory, (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file'] # ) # # if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file): # copyfile(self.vocab_file, out_vocab_file) # # return (out_vocab_file,) # Modified from transformers.model.llama.tokenization_llama.LlamaTokenizer class InternLM2Tokenizer(PreTrainedTokenizer): """ Construct a InternLM2 tokenizer. Based on byte-level Byte-Pair-Encoding. Args: vocab_file (`str`): Path to the vocabulary file. """ vocab_files_names = VOCAB_FILES_NAMES pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP model_input_names = ["input_ids", "attention_mask"] _auto_class = "AutoTokenizer" def __init__( self, vocab_file, unk_token="", bos_token="", eos_token="", pad_token="", sp_model_kwargs: Optional[Dict[str, Any]] = None, add_bos_token=True, add_eos_token=False, decode_with_prefix_space=False, clean_up_tokenization_spaces=False, **kwargs, ): print("register succeed") self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs self.vocab_file = vocab_file self.add_bos_token = add_bos_token self.add_eos_token = add_eos_token self.decode_with_prefix_space = decode_with_prefix_space self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs) self.sp_model.Load(vocab_file) self._no_prefix_space_tokens = None super().__init__( bos_token=bos_token, eos_token=eos_token, unk_token=unk_token, pad_token=pad_token, clean_up_tokenization_spaces=clean_up_tokenization_spaces, **kwargs, ) @property def no_prefix_space_tokens(self): if self._no_prefix_space_tokens is None: vocab = self.convert_ids_to_tokens(list(range(self.vocab_size))) self._no_prefix_space_tokens = { i for i, tok in enumerate(vocab) if not tok.startswith("▁") } return self._no_prefix_space_tokens @property def vocab_size(self): """Returns vocab size""" return self.sp_model.get_piece_size() @property def bos_token_id(self) -> Optional[int]: return self.sp_model.bos_id() @property def eos_token_id(self) -> Optional[int]: return self.sp_model.eos_id() def get_vocab(self): """Returns vocab as a dict""" vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)} vocab.update(self.added_tokens_encoder) return vocab def _tokenize(self, text): """Returns a tokenized string.""" return self.sp_model.encode(text, out_type=str) def _convert_token_to_id(self, token): """Converts a token (str) in an id using the vocab.""" return self.sp_model.piece_to_id(token) def _convert_id_to_token(self, index): """Converts an index (integer) in a token (str) using the vocab.""" token = self.sp_model.IdToPiece(index) return token def _maybe_add_prefix_space(self, tokens, decoded): if tokens and tokens[0] not in self.no_prefix_space_tokens: return " " + decoded else: return decoded def convert_tokens_to_string(self, tokens): """Converts a sequence of tokens (string) in a single string.""" current_sub_tokens = [] out_string = "" prev_is_special = False for token in tokens: # make sure that special tokens are not decoded using sentencepiece model if token in self.all_special_tokens: if not prev_is_special: out_string += " " out_string += self.sp_model.decode(current_sub_tokens) + token prev_is_special = True current_sub_tokens = [] else: current_sub_tokens.append(token) prev_is_special = False out_string += self.sp_model.decode(current_sub_tokens) out_string = self._maybe_add_prefix_space(tokens=tokens, decoded=out_string) return out_string[1:] def save_vocabulary( self, save_directory, filename_prefix: Optional[str] = None ) -> Tuple[str]: """ Save the vocabulary and special tokens file to a directory. Args: save_directory (`str`): The directory in which to save the vocabulary. Returns: `Tuple(str)`: Paths to the files saved. """ if not os.path.isdir(save_directory): logger.error(f"Vocabulary path ({save_directory}) should be a directory") return out_vocab_file = os.path.join( save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"], ) if os.path.abspath(self.vocab_file) != os.path.abspath( out_vocab_file ) and os.path.isfile(self.vocab_file): copyfile(self.vocab_file, out_vocab_file) elif not os.path.isfile(self.vocab_file): with open(out_vocab_file, "wb") as fi: content_spiece_model = self.sp_model.serialized_model_proto() fi.write(content_spiece_model) return (out_vocab_file,) def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None): if self.add_bos_token: bos_token_ids = [self.bos_token_id] else: bos_token_ids = [] output = bos_token_ids + token_ids_0 if token_ids_1 is not None: output = output + token_ids_1 if self.add_eos_token: output = output + [self.eos_token_id] return output def get_special_tokens_mask( self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False, ) -> List[int]: """ Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding special tokens using the tokenizer `prepare_for_model` method. Args: token_ids_0 (`List[int]`): List of IDs. token_ids_1 (`List[int]`, *optional*): Optional second list of IDs for sequence pairs. already_has_special_tokens (`bool`, *optional*, defaults to `False`): Whether or not the token list is already formatted with special tokens for the model. Returns: `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. """ if already_has_special_tokens: return super().get_special_tokens_mask( token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True, ) if token_ids_1 is None: return [1] + ([0] * len(token_ids_0)) + [1] return [1] + ([0] * len(token_ids_0)) + [1, 1] + ([0] * len(token_ids_1)) + [1] def create_token_type_ids_from_sequences( self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None ) -> List[int]: """ Create a mask from the two sequences passed to be used in a sequence-pair classification task. T5 does not make use of token type ids, therefore a list of zeros is returned. Args: token_ids_0 (`List[int]`): List of IDs. token_ids_1 (`List[int]`, *optional*): Optional second list of IDs for sequence pairs. Returns: `List[int]`: List of zeros. """ eos = [self.eos_token_id] if token_ids_1 is None: return len(token_ids_0 + eos) * [0] return len(token_ids_0 + eos + token_ids_1 + eos) * [0] TOKENIZER_MAPPING.register( InternVLChatConfig, (InternLM2Tokenizer, None), exist_ok=True ) ================================================ FILE: python/sglang/srt/configs/janus_pro.py ================================================ # Adapted from: # https://github.com/deepseek-ai/Janus/tree/main/janus/models from dataclasses import dataclass from typing import Dict, List, Tuple, Union import numpy as np import PIL import torch from PIL.Image import Image from transformers import ( BaseImageProcessor, BatchFeature, LlamaConfig, LlamaTokenizerFast, PretrainedConfig, ProcessorMixin, ) from transformers.image_utils import to_numpy_array from sglang.srt.configs.utils import register_image_processor, register_processor from sglang.srt.multimodal.mm_utils import expand2square class DictToObject(dict): def __init__(self, dictionary): super(self).__init__(dictionary) for key, value in dictionary.items(): if isinstance(value, dict): value = DictToObject(value) setattr(self, key, value) class VisionConfig(PretrainedConfig): model_type = "vision" cls: str = "" params = {} def __init__(self, **kwargs): super().__init__(**kwargs) self.cls = kwargs.get("cls", "") if not isinstance(self.cls, str): self.cls = self.cls.__name__ self.params = kwargs.get("params", {}) class GenAlignerConfig(PretrainedConfig): model_type = "gen_aligner" cls: str = "" params = {} def __init__(self, **kwargs): super().__init__(**kwargs) self.cls = kwargs.get("cls", "") if not isinstance(self.cls, str): self.cls = self.cls.__name__ self.params = kwargs.get("params", {}) class GenHeadConfig(PretrainedConfig): model_type = "gen_head" cls: str = "" params = {} def __init__(self, **kwargs): super().__init__(**kwargs) self.cls = kwargs.get("cls", "") if not isinstance(self.cls, str): self.cls = self.cls.__name__ self.params = kwargs.get("params", {}) class AlignerConfig(PretrainedConfig): model_type = "aligner" cls: str = "" params = {} def __init__(self, **kwargs): super().__init__(**kwargs) self.cls = kwargs.get("cls", "") if not isinstance(self.cls, str): self.cls = self.cls.__name__ self.params = kwargs.get("params", {}) class GenVisionConfig(PretrainedConfig): model_type = "gen_vision" cls: str = "" params = {} def __init__(self, **kwargs): super().__init__(**kwargs) self.cls = kwargs.get("cls", "") if not isinstance(self.cls, str): self.cls = self.cls.__name__ self.params = kwargs.get("params", {}) @dataclass class SigLIPVisionCfg: width: int = 1152 layers: Union[Tuple[int, int, int, int], int] = 27 heads: int = 16 patch_size: int = 14 image_size: Union[Tuple[int, int], int] = 336 global_pool: str = "map" mlp_ratio: float = 3.7362 class_token: bool = False num_classes: int = 0 use_checkpoint: bool = False class MultiModalityConfig(PretrainedConfig): model_type = "multi_modality" vision_config: VisionConfig = None aligner_config: AlignerConfig = None gen_vision_config: GenVisionConfig = None gen_aligner_config: GenAlignerConfig = None gen_head_config: GenHeadConfig = None language_config: LlamaConfig = None def __init__(self, **kwargs): super().__init__(**kwargs) vision_config = kwargs.get("vision_config", {}) self.vision_config = VisionConfig(**vision_config) aligner_config = kwargs.get("aligner_config", {}) self.aligner_config = AlignerConfig(**aligner_config) gen_vision_config = kwargs.get("gen_vision_config", {}) self.gen_vision_config = GenVisionConfig(**gen_vision_config) gen_aligner_config = kwargs.get("gen_aligner_config", {}) self.gen_aligner_config = GenAlignerConfig(**gen_aligner_config) gen_head_config = kwargs.get("gen_head_config", {}) self.gen_head_config = GenHeadConfig(**gen_head_config) language_config = kwargs.get("language_config", {}) if isinstance(language_config, LlamaConfig): self.language_config = language_config else: self.language_config = LlamaConfig(**language_config) class VLMImageProcessor(BaseImageProcessor): model_input_names = ["pixel_values"] def __init__( self, image_size: int, min_size: int = 14, image_mean: Union[Tuple[float, float, float], List[float]] = ( 0.48145466, 0.4578275, 0.40821073, ), image_std: Union[Tuple[float, float, float], List[float]] = ( 0.26862954, 0.26130258, 0.27577711, ), rescale_factor: float = 1.0 / 255.0, do_normalize: bool = True, **kwargs, ): super().__init__(**kwargs) self.image_size = image_size self.rescale_factor = rescale_factor self.image_mean = image_mean self.image_std = image_std self.min_size = min_size self.do_normalize = do_normalize if image_mean is None: self.background_color = (127, 127, 127) else: self.background_color = tuple([int(x * 255) for x in image_mean]) def resize(self, pil_img: Image) -> np.ndarray: """ Args: pil_img (PIL.Image): [H, W, 3] in PIL.Image in RGB Returns: x (np.ndarray): [3, self.image_size, self.image_size] """ width, height = pil_img.size max_size = max(width, height) size = [ max(int(height / max_size * self.image_size), self.min_size), max(int(width / max_size * self.image_size), self.min_size), ] if width <= 0 or height <= 0 or size[0] <= 0 or size[1] <= 0: # print(f"orig size = {pil_img.size}, new size = {size}") raise ValueError("Invalid size!") def resize( pil_img, size, interpolation=PIL.Image.Resampling.BICUBIC, antialias=True ): if isinstance(size, int): w, h = pil_img.size if (w <= h and w == size) or (h <= w and h == size): return pil_img if w < h: ow = size oh = int(size * h / w) else: oh = size ow = int(size * w / h) size = (ow, oh) else: size = (size[1], size[0]) return pil_img.resize( size, resample=interpolation, reducing_gap=None if antialias else 3.0 ) pil_img = resize( pil_img, size, interpolation=PIL.Image.Resampling.BICUBIC, antialias=True ) pil_img = expand2square(pil_img, self.background_color) x = to_numpy_array(pil_img) # [H, W, 3] -> [3, H, W] x = np.transpose(x, (2, 0, 1)) return x def preprocess(self, images, return_tensors: str = "pt", **kwargs) -> BatchFeature: # resize and pad to [self.image_size, self.image_size] # then convert from [H, W, 3] to [3, H, W] if not isinstance(images, list): images = [images] images: List[np.ndarray] = [self.resize(image) for image in images] images = [image[:3, ...] for image in images] # rescale from [0, 255] -> [0, 1] images = [ self.rescale( image=image, scale=self.rescale_factor, input_data_format="channels_first", ) for image in images ] # normalize if self.do_normalize: images = [ self.normalize( image=image, mean=self.image_mean, std=self.image_std, input_data_format="channels_first", ) for image in images ] data = {"pixel_values": images} return BatchFeature(data=data, tensor_type=return_tensors) @property def default_shape(self): return [3, self.image_size, self.image_size] class DictOutput(object): def items(self): return self.__dict__.items() def keys(self): return self.__dict__.keys() def __getitem__(self, item): return self.__dict__[item] def __contains__(self, key): return key in self.__dict__ def __setitem__(self, key, value): self.__dict__[key] = value @dataclass class VLChatProcessorOutput(DictOutput): sft_format: str input_ids: torch.Tensor pixel_values: torch.Tensor num_image_tokens: torch.IntTensor def __len__(self): return len(self.input_ids) @dataclass class BatchedVLChatProcessorOutput(DictOutput): sft_format: List[str] input_ids: torch.Tensor pixel_values: torch.Tensor attention_mask: torch.Tensor images_seq_mask: torch.BoolTensor images_emb_mask: torch.BoolTensor # FIXME: had to place Official Processor here, since image_processor module would not be imported in all threads, # hence AutoProcessor registration would not be affective in some cases class VLChatProcessor(ProcessorMixin): image_processor_class = "AutoImageProcessor" tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast") attributes = ["image_processor", "tokenizer"] def __init__( self, image_processor: VLMImageProcessor, tokenizer: LlamaTokenizerFast, image_tag: str = "", image_start_tag: str = "", image_end_tag: str = "", pad_tag: str = "<｜▁pad▁｜>", num_image_tokens: int = 576, add_special_token: bool = False, sft_format: str = "deepseek", mask_prompt: bool = True, ignore_id: int = -100, **kwargs, ): self.image_processor = image_processor self.tokenizer = tokenizer image_id = self.tokenizer.vocab.get(image_tag) if image_id is None: special_tokens = [image_tag] special_tokens_dict = {"additional_special_tokens": special_tokens} self.tokenizer.add_special_tokens(special_tokens_dict) # print(f"Add image tag = {image_tag} to the tokenizer") self.image_tag = image_tag self.image_start_tag = image_start_tag self.image_end_tag = image_end_tag self.pad_tag = pad_tag self.num_image_tokens = num_image_tokens self.add_special_token = add_special_token self.sft_format = sft_format self.ignore_id = ignore_id super().__init__( image_processor, tokenizer, **kwargs, ) @property def image_token(self): return self.image_tag @property def image_id(self) -> int: image_id = self.tokenizer.vocab.get(self.image_tag) return image_id @property def image_start_id(self): image_start_id = self.tokenizer.vocab.get(self.image_start_tag) return image_start_id @property def image_end_id(self): image_end_id = self.tokenizer.vocab.get(self.image_end_tag) return image_end_id @property def image_start_token(self): return self.image_start_tag @property def image_end_token(self): return self.image_end_tag @property def pad_id(self): pad_id = self.tokenizer.vocab.get(self.pad_tag) return pad_id def add_image_token( self, image_indices: List[int], input_ids: torch.LongTensor, ): """ Args: image_indices (List[int]): [index_0, index_1, ..., index_j] input_ids (torch.LongTensor): [N] Returns: input_ids (torch.LongTensor): [N + image tokens] num_image_tokens (torch.IntTensor): [n_images] """ input_slices = [] start = 0 for index in image_indices: if self.add_special_token: end = index + 1 else: end = index # original text tokens input_slices.append(input_ids[start:end]) # add boi, image tokens, eoi and set the mask as False input_slices.append(self.image_start_id * torch.ones((1), dtype=torch.long)) input_slices.append( self.image_id * torch.ones((self.num_image_tokens,), dtype=torch.long) ) input_slices.append(self.image_end_id * torch.ones((1), dtype=torch.long)) start = index + 1 # the left part input_slices.append(input_ids[start:]) # concat all slices input_ids = torch.cat(input_slices, dim=0) num_image_tokens = torch.IntTensor([self.num_image_tokens] * len(image_indices)) return input_ids, num_image_tokens def process_one( self, prompt: str = None, images: List[Image] = None, **kwargs, ): """ Args: prompt (str): the formatted prompt; images (List[ImageType]): the list of images; **kwargs: Returns: outputs (BaseProcessorOutput): the output of the processor, - input_ids (torch.LongTensor): [N + image tokens] - target_ids (torch.LongTensor): [N + image tokens] - images (torch.FloatTensor): [n_images, 3, H, W] - image_id (int): the id of the image token - num_image_tokens (List[int]): the number of image tokens """ sft_format = prompt # tokenize input_ids = self.tokenizer.encode(sft_format) input_ids = torch.LongTensor(input_ids) # add image tokens to the input_ids image_token_mask: torch.Tensor = (input_ids == self.image_id).to(torch.bool) image_indices = image_token_mask.nonzero() input_ids, num_image_tokens = self.add_image_token( image_indices=image_indices, input_ids=input_ids, ) # load images images_outputs = self.image_processor(images, return_tensors="pt") prepare = VLChatProcessorOutput( sft_format=sft_format, input_ids=input_ids, pixel_values=images_outputs.pixel_values, num_image_tokens=num_image_tokens, ) return prepare def __call__( self, *, prompt: str = None, conversations: List[Dict[str, str]] = None, images: List[Image] = None, force_batchify: bool = True, **kwargs, ): """ Args: prompt (str): the formatted prompt; conversations (List[Dict]): conversations with a list of messages; images (List[ImageType]): the list of images; force_batchify (bool): force batchify the inputs; **kwargs: Returns: outputs (BaseProcessorOutput): the output of the processor, - input_ids (torch.LongTensor): [N + image tokens] - images (torch.FloatTensor): [n_images, 3, H, W] - image_id (int): the id of the image token - num_image_tokens (List[int]): the number of image tokens """ prepare = self.process_one( prompt=prompt, conversations=conversations, images=images ) if force_batchify: prepare = self.batchify([prepare]) return prepare def batchify( self, prepare_list: List[VLChatProcessorOutput] ) -> BatchedVLChatProcessorOutput: """ Preprocesses the inputs for multimodal inference. Args: prepare_list (List[VLChatProcessorOutput]): A list of VLChatProcessorOutput. Returns: BatchedVLChatProcessorOutput: A dictionary of the inputs to use for multimodal inference. """ batch_size = len(prepare_list) sft_format = [] n_images = [] seq_lens = [] for prepare in prepare_list: n_images.append(len(prepare.num_image_tokens)) seq_lens.append(len(prepare)) input_token_max_len = max(seq_lens) max_n_images = max(1, max(n_images)) batched_input_ids = torch.full( (batch_size, input_token_max_len), self.pad_id ).long() # FIXME batched_attention_mask = torch.zeros((batch_size, input_token_max_len)).long() batched_pixel_values = torch.zeros( (batch_size, max_n_images, *self.image_processor.default_shape) ).float() batched_images_seq_mask = torch.zeros((batch_size, input_token_max_len)).bool() batched_images_emb_mask = torch.zeros( (batch_size, max_n_images, self.num_image_tokens) ).bool() for i, prepare in enumerate(prepare_list): input_ids = prepare.input_ids seq_len = len(prepare) n_image = len(prepare.num_image_tokens) # left-padding batched_attention_mask[i, -seq_len:] = 1 batched_input_ids[i, -seq_len:] = torch.LongTensor(input_ids) batched_images_seq_mask[i, -seq_len:] = input_ids == self.image_id if n_image > 0: batched_pixel_values[i, :n_image] = prepare.pixel_values for j, n_image_tokens in enumerate(prepare.num_image_tokens): batched_images_emb_mask[i, j, :n_image_tokens] = True sft_format.append(prepare.sft_format) batched_prepares = BatchedVLChatProcessorOutput( input_ids=batched_input_ids, attention_mask=batched_attention_mask, pixel_values=batched_pixel_values, images_seq_mask=batched_images_seq_mask, images_emb_mask=batched_images_emb_mask, sft_format=sft_format, ) return batched_prepares class VLMImageProcessorConfig(PretrainedConfig): model_type = "deepseek_vlm" image_size: int = None min_size: int = None image_mean: Union[Tuple[float, float, float], List[float]] = None image_std: Union[Tuple[float, float, float], List[float]] = None rescale_factor: float = None do_normalize: bool = None def __init__( self, image_size: int, min_size: int = 14, image_mean: Union[Tuple[float, float, float], List[float]] = ( 0.48145466, 0.4578275, 0.40821073, ), image_std: Union[Tuple[float, float, float], List[float]] = ( 0.26862954, 0.26130258, 0.27577711, ), rescale_factor: float = 1.0 / 255.0, do_normalize: bool = True, **kwargs, ): self.image_size = image_size self.min_size = min_size self.image_mean = image_mean self.image_std = image_std self.rescale_factor = rescale_factor self.do_normalize = do_normalize super().__init__(**kwargs) register_processor(MultiModalityConfig, VLChatProcessor) register_image_processor(MultiModalityConfig, VLMImageProcessor) ================================================ FILE: python/sglang/srt/configs/jet_nemotron.py ================================================ from dataclasses import dataclass from typing import Any from transformers.configuration_utils import PretrainedConfig from sglang.srt.configs.mamba_utils import ( Mamba2CacheParams, Mamba2StateShape, mamba2_state_dtype, ) @dataclass class JetBlockConfig: mode: str expand_v: float num_heads: int head_dim: int norm_eps: str conv_size: int dconv_generator_reduction: int dconv_implementation: str class JetNemotronConfig(PretrainedConfig): model_type: str = "jet_nemotron" efficient_attention_config: dict[str, dict[str, Any]] = None hidden_act: str = None hidden_size: int = None initializer_range: float = None intermediate_size: int = None layer_types: list[str] = None max_position_embeddings: int = None num_attention_heads: int = None num_key_value_heads: int = None rms_norm_eps: float = None rope_scaling: None = None rope_theta: float = None @property def full_attention_layer_ids(self) -> list[int]: return [ idx for idx, layer_type in enumerate(self.layer_types) if layer_type in ("attn", "swa") ] @property def linear_layer_ids(self) -> list[int]: return [ idx for idx, layer_type in enumerate(self.layer_types) if layer_type == "jet" ] @property def mamba2_cache_params(self) -> Mamba2CacheParams: from sglang.srt.layers.dp_attention import get_attention_tp_size jet_block_config = JetBlockConfig(**self.efficient_attention_config["jet"]) num_heads = jet_block_config.num_heads head_k_dim = jet_block_config.head_dim head_v_dim = int(head_k_dim * jet_block_config.expand_v) total_v_dim = num_heads * head_v_dim shape = Mamba2StateShape.create( tp_world_size=get_attention_tp_size(), intermediate_size=total_v_dim, n_groups=num_heads, num_heads=num_heads, head_dim=head_v_dim, state_size=head_k_dim, conv_kernel=jet_block_config.conv_size, ) return Mamba2CacheParams( shape=shape, layers=self.linear_layer_ids, dtype=mamba2_state_dtype(self) ) ================================================ FILE: python/sglang/srt/configs/jet_vlm.py ================================================ from typing import Any from transformers.configuration_utils import PretrainedConfig from transformers.models.siglip import SiglipVisionConfig from sglang.srt.configs.jet_nemotron import JetNemotronConfig from sglang.srt.configs.mamba_utils import Mamba2CacheParams class JetVLMConfig(PretrainedConfig): model_type = "jet_vlm" sub_configs = { "text_config": JetNemotronConfig, "vision_config": SiglipVisionConfig, } _auto_class = "AutoConfig" def __init__( self, *, text_config: dict[str, Any] | None = None, vision_config: dict[str, Any] | None = None, image_token_id: int | None = None, video_token_id: int | None = None, **kwargs, ): self.text_config = ( JetNemotronConfig(**text_config) if text_config is not None else JetNemotronConfig() ) self.vision_config = ( SiglipVisionConfig(**vision_config) if vision_config is not None else SiglipVisionConfig() ) self.image_token_id = image_token_id if image_token_id is not None else -1 self.video_token_id = video_token_id if video_token_id is not None else -1 super().__init__(**kwargs) @property def full_attention_layer_ids(self) -> list[int]: return self.text_config.full_attention_layer_ids @property def linear_layer_ids(self) -> list[int]: return self.text_config.linear_layer_ids @property def mamba2_cache_params(self) -> Mamba2CacheParams: return self.text_config.mamba2_cache_params ================================================ FILE: python/sglang/srt/configs/kimi_k25.py ================================================ """ Kimi K25 Model Configuration. """ from transformers import DeepseekV3Config from transformers.configuration_utils import PretrainedConfig class KimiK25VisionConfig(PretrainedConfig): """Vision configuration for K2-VL (vision tower + mm projector). Args: Vision Tower Parameters: patch_size: Patch size for vision tower. init_pos_emb_height: Initial position embedding height. init_pos_emb_width: Initial position embedding width. init_pos_emb_time: Initial position embedding time dimension. pos_emb_type: Type of position embedding. num_attention_heads: Number of attention heads in vision tower. num_hidden_layers: Number of hidden layers in vision tower. hidden_size: Hidden size of vision tower. intermediate_size: Intermediate size in vision tower FFN. merge_kernel_size: Kernel size for spatial patch merging. video_attn_type: Type of video attention. merge_type: Type of merge operation. MM Projector Parameters: mm_projector_type: Type of multimodal projector. mm_hidden_size: Hidden size for projector (defaults to hidden_size). projector_hidden_act: Activation function for projector. projector_ln_eps: Layer norm epsilon for projector. """ model_type = "kimi_k25" def __init__( self, # Vision Tower patch_size: int = 14, init_pos_emb_height: int = 64, init_pos_emb_width: int = 64, init_pos_emb_time: int = 4, pos_emb_type: str = "divided_fixed", num_attention_heads: int = 16, num_hidden_layers: int = 27, hidden_size: int = 1152, intermediate_size: int = 4304, merge_kernel_size: tuple[int, int] = (2, 2), video_attn_type: str = "spatial_temporal", merge_type: str = "sd2_tpool", # MM Projector mm_projector_type: str = "patchmerger", mm_hidden_size: int | None = None, projector_hidden_act: str = "gelu", projector_ln_eps: float = 1e-5, text_hidden_size: int = 7168, **kwargs, ): super().__init__(**kwargs) # Vision Tower self.patch_size = patch_size self.init_pos_emb_height = init_pos_emb_height self.init_pos_emb_width = init_pos_emb_width self.init_pos_emb_time = init_pos_emb_time self.pos_emb_type = pos_emb_type self.num_attention_heads = num_attention_heads self.num_hidden_layers = num_hidden_layers self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.merge_kernel_size = merge_kernel_size self.video_attn_type = video_attn_type self.merge_type = merge_type # MM Projector self.mm_projector_type = mm_projector_type if mm_hidden_size is not None: self.mm_hidden_size = mm_hidden_size else: self.mm_hidden_size = hidden_size self.projector_hidden_act = projector_hidden_act self.projector_ln_eps = projector_ln_eps self.text_hidden_size = text_hidden_size class KimiK25Config(PretrainedConfig): """K2-VL model configuration. K2-VL extends Kimi-VL with video support using video-chunks. A video-chunk consists of multiple consecutive frames (default: 4) that are processed together with temporal pooling. Args: text_config: Configuration for the text model (DeepseekV3). Vision Tower Parameters: patch_size: Patch size for vision tower. init_pos_emb_height: Initial position embedding height. init_pos_emb_width: Initial position embedding width. init_pos_emb_time: Initial position embedding time dimension. pos_emb_type: Type of position embedding. vt_num_attention_heads: Number of attention heads in vision tower. vt_num_hidden_layers: Number of hidden layers in vision tower. vt_hidden_size: Hidden size of vision tower. vt_intermediate_size: Intermediate size in vision tower FFN. merge_kernel_size: Kernel size for spatial patch merging. video_attn_type: Type of video attention. merge_type: Type of merge operation. Video-Chunk Parameters: temporal_merge_kernel_size: Number of frames per video chunk. Default is 4, meaning 4 frames are merged into 1 chunk. sample_fps: Video sampling frame rate. timestamp_mode: Format for chunk timestamps. MM Projector Parameters: mm_projector_type: Type of multimodal projector. mm_hidden_size: Hidden size from vision tower. projector_hidden_act: Activation function for projector. projector_ln_eps: Layer norm epsilon for projector. Other Parameters: ignore_index: The ignore index for the loss function. media_placeholder_token_id: The token ID for media placeholders. pad_token_id: The token ID for padding. """ model_type = "kimi_k25" def __init__( self, text_config: dict | DeepseekV3Config | None = None, vision_config: dict | KimiK25VisionConfig | None = None, # Other parameters ignore_index: int = -100, media_placeholder_token_id: int = 163605, pad_token_id: int = 0, use_unified_vision_chunk: bool = False, video_placeholder: str = "<|kimi_k25_video_placeholder|>", **kwargs, ): if text_config is None: text_config = DeepseekV3Config() elif isinstance(text_config, dict): text_config = DeepseekV3Config(**text_config) if vision_config is None: vision_config = KimiK25VisionConfig() elif isinstance(vision_config, dict): vision_config = KimiK25VisionConfig(**vision_config) self.vision_config = vision_config self.text_config = text_config # Other config self.ignore_index = ignore_index self.media_placeholder_token_id = media_placeholder_token_id self.use_unified_vision_chunk = use_unified_vision_chunk self.video_placeholder = video_placeholder # Propagate quantization config from text model if getattr(self.text_config, "quantization_config", None) is not None: self.quantization_config = self.text_config.quantization_config super().__init__(pad_token_id=pad_token_id, **kwargs) @property def hidden_size(self) -> int: """Get hidden size from text config for compatibility.""" return self.text_config.hidden_size @property def vocab_size(self) -> int: """Get vocab size from text config for compatibility.""" return self.text_config.vocab_size ================================================ FILE: python/sglang/srt/configs/kimi_linear.py ================================================ # Adapted from: https://github.com/vllm-project/vllm/blob/0384aa7150c4c9778efca041ffd1beb3ad2bd694/vllm/transformers_utils/configs/kimi_linear.py from transformers.configuration_utils import PretrainedConfig from sglang.srt.configs.mamba_utils import KimiLinearCacheParams, KimiLinearStateShape class KimiLinearConfig(PretrainedConfig): model_type = "kimi_linear" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, model_type="kimi_linear", vocab_size=163840, hidden_size=4096, head_dim=None, intermediate_size=11008, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=None, hidden_act="silu", initializer_range=0.02, rms_norm_eps=1e-6, use_cache=True, pad_token_id=0, bos_token_id=1, eos_token_id=2, rope_theta=10000.0, rope_scaling=None, tie_word_embeddings=False, moe_intermediate_size: int | None = None, moe_renormalize: bool = True, moe_router_activation_func: str = "sigmoid", num_experts: int | None = None, num_experts_per_token: int | None = None, num_shared_experts: int = 0, routed_scaling_factor: float = 1.0, first_k_dense_replace: int = 0, moe_layer_freq: int = 1, use_grouped_topk: bool = True, num_expert_group: int = 1, topk_group: int = 1, q_lora_rank: int | None = None, kv_lora_rank: int | None = None, qk_nope_head_dim: int | None = None, qk_rope_head_dim: int | None = None, v_head_dim: int | None = None, mla_use_nope: bool | None = False, num_nextn_predict_layers: int = 0, linear_attn_config: dict | None = None, **kwargs, ): self.model_type = model_type self.vocab_size = vocab_size self.hidden_size = hidden_size self.head_dim = ( head_dim if head_dim is not None else hidden_size // num_attention_heads ) self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads # for backward compatibility if num_key_value_heads is None: num_key_value_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.hidden_act = hidden_act self.initializer_range = initializer_range self.rms_norm_eps = rms_norm_eps self.use_cache = use_cache self.rope_theta = rope_theta self.rope_scaling = rope_scaling self.q_lora_rank = q_lora_rank self.kv_lora_rank = kv_lora_rank self.qk_nope_head_dim = qk_nope_head_dim self.qk_rope_head_dim = qk_rope_head_dim self.v_head_dim = v_head_dim self.mla_use_nope = mla_use_nope # moe config self.n_routed_experts = self.num_experts = num_experts self.num_experts_per_token = num_experts_per_token self.moe_renormalize = moe_renormalize self.num_shared_experts = num_shared_experts self.routed_scaling_factor = routed_scaling_factor self.moe_router_activation_func = moe_router_activation_func assert self.moe_router_activation_func in ("softmax", "sigmoid") self.moe_intermediate_size = moe_intermediate_size self.first_k_dense_replace = first_k_dense_replace self.moe_layer_freq = moe_layer_freq self.use_grouped_topk = use_grouped_topk self.num_expert_group = num_expert_group self.topk_group = topk_group self.num_nextn_predict_layers = num_nextn_predict_layers if linear_attn_config is not None: assert linear_attn_config["kda_layers"] is not None assert linear_attn_config["full_attn_layers"] is not None self.linear_attn_config = linear_attn_config super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs, ) @property def is_mla(self): return ( self.q_lora_rank is not None or self.kv_lora_rank is not None or self.qk_nope_head_dim is not None or self.qk_rope_head_dim is not None or self.v_head_dim is not None or self.mla_use_nope is True ) @property def is_moe(self): return self.num_experts is not None @property def is_linear_attn(self) -> bool: return not ( self.linear_attn_config is None or ( isinstance(self.linear_attn_config, dict) and self.linear_attn_config["kda_layers"] is not None and len(self.linear_attn_config["kda_layers"]) == 0 ) ) def is_kda_layer(self, layer_idx: int): return ( self.linear_attn_config is not None and (layer_idx + 1) in self.linear_attn_config["kda_layers"] ) @property def linear_layer_ids(self): return [i for i in range(self.num_hidden_layers) if self.is_kda_layer(i)] @property def full_attention_layer_ids(self): return [i for i in range(self.num_hidden_layers) if not self.is_kda_layer(i)] @property def mamba2_cache_params(self) -> KimiLinearCacheParams: from sglang.srt.layers.dp_attention import get_attention_tp_size shape = KimiLinearStateShape.create( tp_world_size=get_attention_tp_size(), num_heads=self.linear_attn_config["num_heads"], head_dim=self.linear_attn_config["head_dim"], conv_kernel_size=self.linear_attn_config["short_conv_kernel_size"], ) return KimiLinearCacheParams(shape=shape, layers=self.linear_layer_ids) ================================================ FILE: python/sglang/srt/configs/kimi_vl.py ================================================ # SPDX-License-Identifier: Apache-2.0 # Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/configuration_kimi_vl.py from typing import Optional, Union from transformers.configuration_utils import PretrainedConfig from sglang.srt.configs.deepseekvl2 import DeepseekV2Config from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig class KimiVLConfig(PretrainedConfig): model_type = "kimi_vl" def __init__( self, vision_config: Optional[Union[dict, MoonViTConfig]] = None, text_config: Optional[Union[dict, DeepseekV2Config]] = None, ignore_index: int = -100, media_placeholder_token_id: int = 163605, pad_token_id: int = 0, **kwargs ): if vision_config is None: vision_config = MoonViTConfig() elif isinstance(vision_config, dict): vision_config = MoonViTConfig(**vision_config) self.vision_config = vision_config if text_config is None: text_config = DeepseekV2Config() elif isinstance(text_config, dict): text_config = DeepseekV2Config(**text_config) self.text_config = text_config self.ignore_index = ignore_index self.media_placeholder_token_id = media_placeholder_token_id super().__init__(pad_token_id=pad_token_id, **kwargs) ================================================ FILE: python/sglang/srt/configs/kimi_vl_moonvit.py ================================================ # SPDX-License-Identifier: Apache-2.0 # Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/configuration_kimi_vl.py from transformers.configuration_utils import PretrainedConfig class MoonViTConfig(PretrainedConfig): model_type = "moonvit" def __init__( self, patch_size: int = 14, init_pos_emb_height: int = 64, init_pos_emb_width: int = 64, num_attention_heads: int = 16, num_hidden_layers: int = 27, hidden_size: int = 1152, intermediate_size: int = 4304, merge_kernel_size: tuple[int, int] = (2, 2), **kwargs, ): super().__init__(**kwargs) self.patch_size = patch_size # Positional embedding config self.init_pos_emb_height = init_pos_emb_height self.init_pos_emb_width = init_pos_emb_width # Transformer config self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.hidden_size = hidden_size self.intermediate_size = intermediate_size # Patch merger config self.merge_kernel_size = merge_kernel_size ================================================ FILE: python/sglang/srt/configs/lfm2.py ================================================ # coding=utf-8 # Copyright 2024 Liquid AI and the HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """LFM2 (Liquid Foundation Model 2) configuration""" from typing import List, Optional from transformers import CONFIG_MAPPING from transformers import Lfm2Config as HFLfm2Config from transformers.utils import logging from sglang.srt.configs.mamba_utils import ( Mamba2CacheParams, Mamba2StateShape, mamba2_state_dtype, ) logger = logging.get_logger(__name__) class Lfm2Config(HFLfm2Config): """ SGLang configuration for LFM2 models. Extends HuggingFace's Lfm2Config with hybrid model properties needed by SGLang. LFM2 uses a hybrid architecture mixing full attention and ShortConv layers. """ @property def full_attention_layer_ids(self) -> List[int]: """Return indices of attention layers for KV cache.""" return [i for i, lt in enumerate(self.layer_types) if lt == "full_attention"] @property def linear_layer_ids(self) -> List[int]: """Return indices of conv layers for conv state cache.""" return [ i for i, lt in enumerate(self.layer_types) if lt in ("conv", "short_conv") ] @property def mamba_chunk_size(self) -> int: """Return chunk size for Mamba2 backend. LFM2 doesn't use chunking, return 1.""" return 1 @property def mamba2_cache_params(self) -> Optional[Mamba2CacheParams]: """ Get cache params for HybridReqToTokenPool initialization. LFM2 uses ShortConv layers with a small fixed-size cache (kernel_size - 1). Unlike full Mamba2 models, LFM2 only uses the conv state, not SSM temporal state. """ from sglang.srt.layers.dp_attention import get_attention_tp_size conv_layer_ids = self.linear_layer_ids if not conv_layer_ids: return None hidden_size = self.hidden_size conv_kernel = int(self.conv_L_cache) # get_attention_tp_size() requires initialization, default to 1 if not available try: tp_size = get_attention_tp_size() except (AssertionError, RuntimeError): tp_size = 1 # For ShortConv layers, we use a simplified Mamba2StateShape # LFM2 doesn't use SSM state (state_size=0), only conv state # We pass num_heads=tp_size so divide(tp_size, tp_size)=1 always works. # Since state_size=0, the temporal state shape has zero elements anyway. shape = Mamba2StateShape.create( tp_world_size=tp_size, intermediate_size=hidden_size, n_groups=1, # ShortConv doesn't use grouping num_heads=tp_size, # Ensures divide works; temporal state is empty anyway head_dim=hidden_size, # Conv operates on full hidden dim state_size=0, # No SSM temporal state for ShortConv conv_kernel=conv_kernel, ) return Mamba2CacheParams( shape=shape, layers=conv_layer_ids, dtype=mamba2_state_dtype(self), ) # Override HuggingFace's Lfm2Config with our extended version # Cannot use .register() because lfm2 is already registered by transformers # Directly modify the internal _extra_content dict instead CONFIG_MAPPING._extra_content["lfm2"] = Lfm2Config ================================================ FILE: python/sglang/srt/configs/lfm2_moe.py ================================================ # Copyright 2025 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """LFM2-MoE (Liquid Foundation Model 2 - Mixture of Experts) configuration Note: HF transformers has Lfm2MoeConfig in v5.0.0rc2 (unreleased). Once released, we could inherit from it like Lfm2Config does with HFLfm2Config. For now, we define a standalone config to support the model immediately. """ from typing import List, Optional from transformers import CONFIG_MAPPING from transformers.configuration_utils import PretrainedConfig from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape class Lfm2MoeConfig(PretrainedConfig): """ Configuration for LFM2-MoE models (e.g., LiquidAI/LFM2-8B-A1B). LFM2-MoE is a hybrid architecture with: - Attention layers and ShortConv layers (like dense LFM2) - MoE (Mixture of Experts) FFN layers with sigmoid routing Key MoE specifics: - First `num_dense_layers` use dense MLP, rest use MoE - Sigmoid routing (not softmax) with expert_bias for load balancing - expert_bias is fp32 for numerical stability """ model_type = "lfm2_moe" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size: int = 65536, hidden_size: int = 2048, intermediate_size: int = 7168, moe_intermediate_size: int = 1792, num_hidden_layers: int = 32, num_attention_heads: int = 32, num_key_value_heads: int = 8, max_position_embeddings: int = 128000, initializer_range: float = 0.02, norm_eps: float = 1e-5, use_cache: bool = True, pad_token_id: int = 0, bos_token_id: int = 1, eos_token_id: int = 2, tie_word_embeddings: bool = True, rope_parameters: Optional[dict] = None, conv_bias: bool = False, conv_L_cache: int = 3, # MoE-specific parameters num_dense_layers: int = 2, num_experts: int = 32, num_experts_per_tok: int = 4, use_expert_bias: bool = True, routed_scaling_factor: float = 1.0, norm_topk_prob: bool = True, # Layer types layer_types: Optional[List[str]] = None, **kwargs, ): self.vocab_size = vocab_size self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.moe_intermediate_size = moe_intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.max_position_embeddings = max_position_embeddings self.initializer_range = initializer_range self.norm_eps = norm_eps self.use_cache = use_cache # Conv parameters self.conv_bias = conv_bias self.conv_L_cache = conv_L_cache # MoE parameters self.num_dense_layers = num_dense_layers self.num_experts = num_experts self.num_experts_per_tok = num_experts_per_tok self.use_expert_bias = use_expert_bias self.routed_scaling_factor = routed_scaling_factor self.norm_topk_prob = norm_topk_prob # Layer types (attention vs conv) self.layer_types = layer_types # RoPE parameters self.rope_parameters = rope_parameters # Validate layer_types length matches num_hidden_layers if layer_types is not None and len(layer_types) != num_hidden_layers: raise ValueError( f"layer_types length ({len(layer_types)}) must match " f"num_hidden_layers ({num_hidden_layers})" ) # Handle tie_embedding alias from original config tie_word_embeddings = kwargs.pop("tie_embedding", tie_word_embeddings) super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs, ) @property def full_attention_layer_ids(self) -> List[int]: """Return indices of attention layers for KV cache.""" if self.layer_types is None: return [] return [i for i, lt in enumerate(self.layer_types) if lt == "full_attention"] @property def linear_layer_ids(self) -> List[int]: """Return indices of conv layers for conv state cache.""" if self.layer_types is None: return [] return [ i for i, lt in enumerate(self.layer_types) if lt in ("conv", "short_conv") ] @property def mamba_chunk_size(self) -> int: """Return chunk size for Mamba2 backend. LFM2 doesn't use chunking.""" return 1 @property def mamba2_cache_params(self) -> Optional[Mamba2CacheParams]: """ Get cache params for HybridReqToTokenPool initialization. LFM2-MoE uses ShortConv layers with a small fixed-size cache. """ from sglang.srt.layers.dp_attention import get_attention_tp_size conv_layer_ids = self.linear_layer_ids if not conv_layer_ids: return None hidden_size = self.hidden_size # conv_L_cache in config is kernel_size (e.g., 3) conv_kernel = int(self.conv_L_cache) # actual cache size is kernel_size - 1 (e.g., 2 for kernel=3) try: tp_size = get_attention_tp_size() except (AssertionError, RuntimeError): tp_size = 1 shape = Mamba2StateShape.create( tp_world_size=tp_size, intermediate_size=hidden_size, n_groups=1, num_heads=tp_size, # Ensures divide works; temporal state is empty anyway head_dim=hidden_size, state_size=0, conv_kernel=conv_kernel, ) # Uses default mamba2_state_dtype() which reads SGLANG_MAMBA_CONV_DTYPE env var # (defaults to bfloat16). Set SGLANG_MAMBA_CONV_DTYPE=float16 for fp16 inference. return Mamba2CacheParams( shape=shape, layers=conv_layer_ids, ) # Register with transformers CONFIG_MAPPING so AutoConfig.from_pretrained() # can instantiate our config class when loading models with model_type="lfm2_moe" try: CONFIG_MAPPING.register("lfm2_moe", Lfm2MoeConfig) except Exception: # Already registered or registration failed - use direct assignment CONFIG_MAPPING._extra_content["lfm2_moe"] = Lfm2MoeConfig ================================================ FILE: python/sglang/srt/configs/load_config.py ================================================ # Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py import enum import logging from dataclasses import dataclass, field from typing import Any, List, Optional, Union import orjson from sglang.srt.configs.modelopt_config import ModelOptConfig from sglang.srt.utils import is_hip logger = logging.getLogger(__name__) class LoadFormat(str, enum.Enum): AUTO = "auto" PT = "pt" SAFETENSORS = "safetensors" NPCACHE = "npcache" DUMMY = "dummy" SHARDED_STATE = "sharded_state" GGUF = "gguf" BITSANDBYTES = "bitsandbytes" MISTRAL = "mistral" LAYERED = "layered" FLASH_RL = "flash_rl" # For RL training with quantized models JAX = "jax" REMOTE = "remote" REMOTE_INSTANCE = "remote_instance" RDMA = "rdma" LOCAL_CACHED = "local_cached" FASTSAFETENSORS = "fastsafetensors" PRIVATE = "private" @dataclass class LoadConfig: """ download_dir: Directory to download and load the weights, default to the default cache directory of huggingface. load_format: The format of the model weights to load: "auto" will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available. "pt" will load the weights in the pytorch bin format. "safetensors" will load the weights in the safetensors format. "npcache" will load the weights in pytorch format and store a numpy cache to speed up the loading. "dummy" will initialize the weights with random values, which is mainly for profiling. "bitsandbytes" will load nf4 type weights. "flash_rl" will load weights with support for RL training with quantized models, enabling efficient weight reloading. ignore_patterns: The list of patterns to ignore when loading the model. Default to "original/**/*" to avoid repeated loading of llama's checkpoints. decryption_key_file: If set, decrypts the output files with a password read from this file (after PBKDF2). decrypt_max_concurrency: The maximum number of concurrent processes to decrypt the safetensor files. -1 means no limit. # ModelOpt-specific loading options modelopt_checkpoint_restore_path: Optional[str] = None modelopt_checkpoint_save_path: Optional[str] = None modelopt_export_path: Optional[str] = None """ load_format: Union[str, LoadFormat] = LoadFormat.AUTO download_dir: Optional[str] = None model_loader_extra_config: Optional[Union[str, dict]] = field(default_factory=dict) ignore_patterns: Optional[Union[List[str], str]] = None decryption_key_file: Optional[str] = None decrypt_max_concurrency: int = -1 tp_rank: Optional[int] = None remote_instance_weight_loader_seed_instance_ip: Optional[str] = None remote_instance_weight_loader_seed_instance_service_port: Optional[int] = None remote_instance_weight_loader_send_weights_group_ports: Optional[List[int]] = None remote_instance_weight_loader_backend: Optional[str] = None remote_instance_weight_loader_transfer_engine: Optional[Any] = None modelexpress_url: Optional[str] = None modelexpress_model_name: Optional[str] = None # ModelOpt-specific loading options modelopt_checkpoint_restore_path: Optional[str] = None modelopt_checkpoint_save_path: Optional[str] = None modelopt_export_path: Optional[str] = None # ModelOpt configuration object modelopt_config: Optional[ModelOptConfig] = None # QuantizedRL-specific options (for FlashRL-style quantization) rl_quant_profile: Optional[str] = ( None # Path to rollout quantization profile (e.g., /root/profile.7b.pt) ) # For multi-layer MTP draft_model_idx: Optional[int] = None def __post_init__(self): model_loader_extra_config = self.model_loader_extra_config or {} if isinstance(model_loader_extra_config, str): self.model_loader_extra_config = orjson.loads(model_loader_extra_config) self._verify_load_format() if self.ignore_patterns is not None and len(self.ignore_patterns) > 0: logger.info( "Ignoring the following patterns when downloading weights: %s", self.ignore_patterns, ) else: self.ignore_patterns = ["original/**/*"] # Create ModelOptConfig if not provided if self.modelopt_config is None: self.modelopt_config = ModelOptConfig( checkpoint_restore_path=self.modelopt_checkpoint_restore_path, checkpoint_save_path=self.modelopt_checkpoint_save_path, export_path=self.modelopt_export_path, ) def _verify_load_format(self) -> None: if not isinstance(self.load_format, str): return load_format = self.load_format.lower() self.load_format = LoadFormat(load_format) rocm_not_supported_load_format: List[str] = [] if is_hip() and load_format in rocm_not_supported_load_format: rocm_supported_load_format = [ f for f in LoadFormat.__members__ if (f not in rocm_not_supported_load_format) ] raise ValueError( f"load format '{load_format}' is not supported in ROCm. " f"Supported load formats are " f"{rocm_supported_load_format}" ) ================================================ FILE: python/sglang/srt/configs/longcat_flash.py ================================================ from transformers.configuration_utils import PretrainedConfig from transformers.utils import logging logger = logging.get_logger(__name__) FLASH_PRETRAINED_CONFIG_ARCHIVE_MAP = {} class LongcatFlashConfig(PretrainedConfig): model_type = "longcat_flash" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size=131072, hidden_size=6144, intermediate_size=None, ffn_hidden_size=12288, expert_ffn_hidden_size=2048, num_layers=28, num_hidden_layers=None, num_attention_heads=64, ep_size=1, kv_lora_rank=512, q_lora_rank=1536, qk_rope_head_dim=128, qk_nope_head_dim=128, v_head_dim=128, n_routed_experts=512, moe_topk=12, norm_topk_prob=False, max_position_embeddings=131072, rms_norm_eps=1e-05, use_cache=True, pad_token_id=None, bos_token_id=1, eos_token_id=2, pretraining_tp=1, tie_word_embeddings=False, rope_theta=10000000.0, rope_scaling=None, attention_bias=False, attention_dropout=0.0, mla_scale_q_lora=True, mla_scale_kv_lora=True, torch_dtype="bfloat16", params_dtype="bfloat16", rounter_params_dtype="float32", router_bias=False, topk_method=None, routed_scaling_factor=6.0, zero_expert_num=256, zero_expert_type="identity", nextn_use_scmoe=False, num_nextn_predict_layers=1, ngram_vocab_size_ratio=None, emb_neighbor_num=None, emb_split_num=None, **kwargs, ): super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, torch_dtype=torch_dtype, params_dtype=params_dtype, rounter_params_dtype=rounter_params_dtype, topk_method=topk_method, router_bias=router_bias, nextn_use_scmoe=nextn_use_scmoe, num_nextn_predict_layers=num_nextn_predict_layers, **kwargs, ) self.vocab_size = vocab_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.num_hidden_layers = ( num_hidden_layers if num_hidden_layers is not None else num_layers ) self.intermediate_size = ( intermediate_size if intermediate_size is not None else ffn_hidden_size ) self.moe_intermediate_size = expert_ffn_hidden_size self.num_attention_heads = num_attention_heads self.ep_size = ep_size self.kv_lora_rank = kv_lora_rank self.q_lora_rank = q_lora_rank self.qk_rope_head_dim = qk_rope_head_dim self.v_head_dim = v_head_dim self.qk_nope_head_dim = qk_nope_head_dim self.n_routed_experts = n_routed_experts self.moe_topk = moe_topk self.norm_topk_prob = norm_topk_prob self.rms_norm_eps = rms_norm_eps self.pretraining_tp = pretraining_tp self.use_cache = use_cache self.rope_theta = rope_theta self.rope_scaling = rope_scaling self.attention_bias = attention_bias self.attention_dropout = attention_dropout self.mla_scale_q_lora = mla_scale_q_lora self.mla_scale_kv_lora = mla_scale_kv_lora self.zero_expert_num = zero_expert_num self.zero_expert_type = zero_expert_type self.routed_scaling_factor = routed_scaling_factor self.hidden_act = "silu" self.use_ngram_embedding = ngram_vocab_size_ratio is not None if self.use_ngram_embedding: self.ngram_embedding_m = int(ngram_vocab_size_ratio * vocab_size) self.ngram_embedding_n = emb_neighbor_num self.ngram_embedding_k = emb_split_num ================================================ FILE: python/sglang/srt/configs/mamba_utils.py ================================================ # Copyright 2025 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Common config utils for mamba2 - NemotronH, FalconH1, Qwen3Next, LFM2, etc.""" import logging from abc import ABC from dataclasses import dataclass, field from typing import List, Optional import numpy as np import torch from sglang.srt.distributed.utils import divide from sglang.srt.environ import envs logger = logging.getLogger(__name__) def extra_groups_for_head_shards(ngroups: int, tp_size: int): """Compute the increase in group numbers to account for replication in order to accompany the head shards.""" # in the case ngoups % tp_size == 0, this will be zero if ngroups % tp_size == 0: return 0 # for n_groups == 1, this is exactly tp_size - n_groups return tp_size - ngroups @dataclass(kw_only=True, frozen=True) class Mamba2StateDType: conv: torch.dtype temporal: torch.dtype def mamba2_state_dtype(config=None) -> Mamba2StateDType: """ Get mamba2 state dtype from config or environment variable. Priority (from highest to lowest): 1. Environment variable SGLANG_MAMBA_SSM_DTYPE 2. Config file (config.mamba_ssm_dtype or config.text_config.mamba_ssm_dtype) 3. Default "float32" Args: config: Optional config object (PretrainedConfig). If provided, will read mamba_ssm_dtype from it. For VL models, reads from text_config. Returns: Mamba2StateDType with conv and temporal dtypes """ dtype_map = { "float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16, } conv_dtype = dtype_map.get(envs.SGLANG_MAMBA_CONV_DTYPE.get(), torch.bfloat16) # Get SSM dtype: default -> config -> env var ssm_dtype = torch.float32 # Step 1: Default value # Step 2: Try to read from config if config is not None: config_dtype = None if hasattr(config, "text_config") and hasattr( config.text_config, "mamba_ssm_dtype" ): # VL model: read from text_config config_dtype = config.text_config.mamba_ssm_dtype elif hasattr(config, "mamba_ssm_dtype"): # Text model: read from root config config_dtype = config.mamba_ssm_dtype if config_dtype is not None: if config_dtype not in dtype_map: logger.warning( f"Invalid mamba_ssm_dtype '{config_dtype}' in config. " f"Must be one of {list(dtype_map.keys())}. Using default 'float32'." ) else: ssm_dtype = dtype_map[config_dtype] # Step 3: Check environment variable, if not None, override env_ssm_dtype = envs.SGLANG_MAMBA_SSM_DTYPE.get() if env_ssm_dtype is not None: if env_ssm_dtype not in dtype_map: logger.warning( f"Invalid mamba_ssm_dtype '{env_ssm_dtype}' from environment variable. " f"Must be one of {list(dtype_map.keys())}. Using default 'float32'." ) else: ssm_dtype = dtype_map[env_ssm_dtype] logger.debug(f"Mamba2 state dtype: conv_dtype={conv_dtype}, ssm_dtype={ssm_dtype}") return Mamba2StateDType(conv=conv_dtype, temporal=ssm_dtype) @dataclass(kw_only=True, frozen=True) class BaseLinearStateParams(ABC): dtype: Mamba2StateDType = field(default_factory=lambda: mamba2_state_dtype(None)) layers: list[int] @property def mamba_cache_per_req(self) -> int: conv_numel = int( np.sum([np.prod(conv_shape) for conv_shape in self.shape.conv]) ) ssm_numel = int(np.prod(self.shape.temporal)) return ( conv_numel * self.dtype.conv.itemsize + ssm_numel * self.dtype.temporal.itemsize ) * len(self.layers) @dataclass(kw_only=True, frozen=True) class Mamba2StateShape: conv: list[tuple[int, int]] temporal: tuple[int, int, int] intermediate_size: int conv_dim: int ssm_state_size: int num_heads: int head_dim: int state_size: int conv_kernel: int @staticmethod def create( *, tp_world_size: int, intermediate_size: int, n_groups: int, num_heads: int, head_dim: int, state_size: int, conv_kernel: int, ) -> "Mamba2StateShape": # if n_groups is not divisible by world_size, need to extend the shards # to ensure all groups needed by a head is sharded along with it if n_groups % tp_world_size != 0: extra_groups = extra_groups_for_head_shards(n_groups, tp_world_size) n_groups += extra_groups # heads and n_groups are TP-ed conv_dim = intermediate_size + 2 * n_groups * state_size # contiguous along 'dim' axis conv_state_shape = divide(conv_dim, tp_world_size), conv_kernel - 1 # These are not TP-ed as they depend on A, dt_bias, D # - they are typically small # e.g., QWen3-Next: (32, 128, 128) temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, state_size) return Mamba2StateShape( conv=[conv_state_shape], temporal=temporal_state_shape, intermediate_size=intermediate_size, conv_dim=conv_dim, ssm_state_size=state_size, num_heads=num_heads, head_dim=head_dim, state_size=state_size, conv_kernel=conv_kernel, ) @dataclass(kw_only=True, frozen=True) class Mamba2CacheParams(BaseLinearStateParams): shape: Mamba2StateShape @dataclass(kw_only=True, frozen=True) class KimiLinearStateShape: conv: List[tuple[int, int]] temporal: tuple[int, int, int] num_heads: int head_dim: int num_k_heads: int head_k_dim: int conv_kernel: int num_spec: int @staticmethod def create( *, tp_world_size: int, num_heads: int, head_dim: int, num_k_heads: Optional[int] = None, head_k_dim: Optional[int] = None, conv_kernel_size: int = 4, num_spec: int = 0, ) -> "KimiLinearStateShape": if num_k_heads is None: num_k_heads = num_heads if head_k_dim is None: head_k_dim = head_dim proj_size = num_heads * head_dim proj_k_size = num_k_heads * head_k_dim conv_state_shape = (divide(proj_size, tp_world_size), conv_kernel_size - 1) conv_state_k_shape = (divide(proj_k_size, tp_world_size), conv_kernel_size - 1) temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim) conv_state_shape = ( conv_state_shape[1], conv_state_shape[0] + conv_state_k_shape[0] * 2, ) return KimiLinearStateShape( conv=[conv_state_shape], temporal=temporal_state_shape, num_heads=num_heads, head_dim=head_dim, num_k_heads=num_k_heads, head_k_dim=head_k_dim, conv_kernel=conv_kernel_size, num_spec=num_spec, ) @dataclass(kw_only=True, frozen=True) class KimiLinearCacheParams(BaseLinearStateParams): shape: KimiLinearStateShape ================================================ FILE: python/sglang/srt/configs/model_config.py ================================================ # Copyright 2023-2024 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== import json import logging import math import os from enum import Enum, IntEnum, auto from pathlib import Path from typing import Any, List, Optional, Set, Union import torch from transformers import PretrainedConfig from sglang.srt.environ import envs from sglang.srt.layers.quantization import QUANTIZATION_METHODS from sglang.srt.server_args import ServerArgs from sglang.srt.utils import is_hip, is_sm100_supported, retry from sglang.srt.utils.hf_transformers_utils import ( get_config, get_context_length, get_generation_config, get_hf_text_config, get_sparse_attention_config, ) from sglang.utils import is_in_ci logger = logging.getLogger(__name__) class AttentionArch(IntEnum): MLA = auto() MHA = auto() class ModelImpl(str, Enum): AUTO = "auto" SGLANG = "sglang" TRANSFORMERS = "transformers" MINDSPORE = "mindspore" def is_deepseek_nsa(config) -> bool: architectures = ( config.get("architectures") if isinstance(config, dict) else getattr(config, "architectures", None) ) index_topk = ( config.get("index_topk") if isinstance(config, dict) else getattr(config, "index_topk", None) ) return ( architectures is not None and architectures[0] in [ "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", "DeepseekV3ForCausalLMNextN", "MistralLarge3ForCausalLM", "PixtralForConditionalGeneration", "GlmMoeDsaForCausalLM", ] and index_topk is not None ) def get_nsa_index_head_dim(config: PretrainedConfig) -> int: assert is_deepseek_nsa(config) return config.index_head_dim def get_nsa_index_topk(config: PretrainedConfig) -> int: assert is_deepseek_nsa(config) return config.index_topk def get_nsa_index_n_heads(config: PretrainedConfig) -> int: assert is_deepseek_nsa(config) return config.index_n_heads class ModelConfig: def __init__( self, model_path: str, trust_remote_code: bool = True, revision: Optional[str] = None, context_length: Optional[int] = None, model_override_args: str = "{}", is_embedding: Optional[bool] = None, enable_multimodal: Optional[bool] = None, dtype: str = "auto", quantization: Optional[str] = None, override_config_file: Optional[str] = None, is_draft_model: bool = False, model_impl: Union[str, ModelImpl] = ModelImpl.AUTO, sampling_defaults: str = "openai", quantize_and_serve: bool = False, is_multi_layer_eagle: bool = False, encoder_only: bool = False, language_only: bool = False, disable_hybrid_swa_memory: bool = False, ) -> None: # Parse args self.model_path = model_path self.revision = revision self.quantization = quantization self.is_draft_model = is_draft_model self.model_impl = model_impl self.sampling_defaults = sampling_defaults self.quantize_and_serve = quantize_and_serve self.is_multi_layer_eagle = is_multi_layer_eagle self.disable_hybrid_swa_memory = disable_hybrid_swa_memory # Validate quantize_and_serve configuration self._validate_quantize_and_serve_config() # Get hf config self._maybe_pull_model_tokenizer_from_remote() self.model_override_args = json.loads(model_override_args) kwargs = {} if override_config_file and override_config_file.strip(): kwargs["_configuration_file"] = override_config_file.strip() self.hf_config = get_config( self.model_path, trust_remote_code=trust_remote_code, revision=revision, model_override_args=self.model_override_args, **kwargs, ) self.hf_text_config = get_hf_text_config(self.hf_config) self.hf_generation_config = get_generation_config( self.model_path, trust_remote_code=trust_remote_code, revision=revision, **kwargs, ) # Set enable_multimodal if enable_multimodal is None: mm_disabled_models = [ "Gemma3ForConditionalGeneration", "Llama4ForConditionalGeneration", "Step3VLForConditionalGeneration", ] if self.hf_config.architectures[0] in mm_disabled_models: enable_multimodal = False logger.info( f"Multimodal is disabled for {self.hf_config.model_type}. To enable it, set --enable-multimodal." ) else: enable_multimodal = True # Config draft model self._config_draft_model() # Check model type self.attention_chunk_size = getattr( self.hf_text_config, "attention_chunk_size", None ) self.sliding_window_size = self._get_sliding_window_size() self.is_generation = is_generation_model( self.hf_config.architectures, is_embedding ) self.is_multimodal = enable_multimodal and is_multimodal_model( self.hf_config.architectures ) self.is_multimodal_gen = enable_multimodal and is_multimodal_gen_model( self.hf_config.architectures ) self.is_image_gen = enable_multimodal and is_image_gen_model( self.hf_config.architectures ) self.is_audio_model = enable_multimodal and is_audio_model( self.hf_config.architectures ) # TODO: requires further polishing self.is_image_understandable_model = enable_multimodal and hasattr( self.hf_config, "vision_config" ) self.is_audio_understandable_model = enable_multimodal and hasattr( self.hf_config, "audio_config" ) self.is_multimodal_chunked_prefill_supported = ( enable_multimodal and is_multimodal_chunked_prefill_supported(self.hf_config.architectures) ) self.is_encoder_decoder = is_encoder_decoder_model(self.hf_config.architectures) self.is_local_attention_model = is_local_attention_model( self.hf_config.architectures ) self.use_ngram_embedding = getattr(self.hf_config, "use_ngram_embedding", False) self.is_piecewise_cuda_graph_disabled_model = ( is_piecewise_cuda_graph_disabled_model(self.hf_config.architectures) or is_deepseek_nsa(self.hf_text_config) ) self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype) # Derive context length and model shapes self._derive_context_length(context_length) self._derive_model_shapes() # Update hybrid model self._derive_hybrid_model() # Verify quantization self._verify_quantization() self._verify_transformers_version() # Verify dual-chunk attention config self._verify_dual_chunk_attention_config() # Cache attributes self.hf_eos_token_id = self._get_hf_eos_token_id() # multimodal self.image_token_id = getattr( self.hf_config, "image_token_id", None ) or getattr(self.hf_config, "image_token_index", None) self.hf_config.encoder_only = encoder_only self.hf_config.language_only = language_only # matryoshka embeddings self.matryoshka_dimensions = getattr( self.hf_config, "matryoshka_dimensions", None ) self.is_matryoshka = self.matryoshka_dimensions or getattr( self.hf_config, "is_matryoshka", False ) @staticmethod def from_server_args( server_args: ServerArgs, model_path: str = None, model_revision: str = None, is_draft_model: bool = False, **kwargs, ): quantization = ( server_args.speculative_draft_model_quantization if is_draft_model else server_args.quantization ) override_config_file = ( server_args.decrypted_draft_config_file if is_draft_model else server_args.decrypted_config_file ) return ModelConfig( model_path=model_path or server_args.model_path, trust_remote_code=server_args.trust_remote_code, revision=model_revision or server_args.revision, context_length=server_args.context_length, model_override_args=server_args.json_model_override_args, is_embedding=server_args.is_embedding, enable_multimodal=server_args.enable_multimodal, dtype=server_args.dtype, quantization=quantization, model_impl=server_args.model_impl, sampling_defaults=server_args.sampling_defaults, quantize_and_serve=server_args.quantize_and_serve, override_config_file=override_config_file, is_multi_layer_eagle=server_args.enable_multi_layer_eagle, language_only=server_args.language_only, encoder_only=server_args.encoder_only, is_draft_model=is_draft_model, disable_hybrid_swa_memory=server_args.disable_hybrid_swa_memory, **kwargs, ) def _config_draft_model(self): is_draft_model = self.is_draft_model if is_draft_model and self.hf_config.architectures[0] in [ "DeepseekV3ForCausalLM", "GlmMoeDsaForCausalLM", ]: self.hf_config.architectures[0] = "DeepseekV3ForCausalLMNextN" if is_draft_model and self.hf_config.architectures[0] in [ "Glm4MoeForCausalLM", "Glm4MoeLiteForCausalLM", ]: self.hf_config.architectures[0] = "Glm4MoeForCausalLMNextN" if is_draft_model and self.hf_config.architectures[0] in [ "GlmOcrForConditionalGeneration", ]: self.hf_config.architectures[0] = "GlmOcrForConditionalGenerationNextN" if ( is_draft_model and self.hf_config.architectures[0] == "LongcatFlashForCausalLM" ): self.hf_config.architectures[0] = "LongcatFlashForCausalLMNextN" self.hf_config.num_hidden_layers = self.hf_config.num_nextn_predict_layers if is_draft_model and self.hf_config.architectures[0] == "MiMoForCausalLM": self.hf_config.architectures[0] = "MiMoMTP" if ( is_draft_model and self.hf_config.architectures[0] == "MiMoV2FlashForCausalLM" ): self.hf_config.architectures[0] = "MiMoV2MTP" if is_draft_model and self.hf_config.architectures[0] == "Step3p5ForCausalLM": self.hf_config.architectures[0] = "Step3p5MTP" if is_draft_model and self.hf_config.architectures[0] in [ "BailingMoeV2ForCausalLM", "BailingMoeForCausalLM", "BailingMoeV2_5ForCausalLM", ]: self.hf_config.architectures[0] = "BailingMoeForCausalLMNextN" if ( is_draft_model and self.hf_config.architectures[0] == "Ernie4_5_MoeForCausalLM" ): self.hf_config.architectures[0] = "Ernie4_5_MoeForCausalLMMTP" if is_draft_model and self.hf_config.architectures[0] == "Qwen3NextForCausalLM": self.hf_config.architectures[0] = "Qwen3NextForCausalLMMTP" self.hf_config.num_nextn_predict_layers = 1 if is_draft_model and self.hf_config.architectures[0] in [ "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration", ]: self.hf_config.architectures[0] = "Qwen3_5ForCausalLMMTP" self.hf_config.num_nextn_predict_layers = 1 if is_draft_model and self.hf_config.architectures[0] == "ExaoneMoEForCausalLM": self.hf_config.architectures[0] = "ExaoneMoEForCausalLMMTP" self.hf_config.num_nextn_predict_layers = 1 if is_draft_model and self.hf_config.architectures[0] == "NemotronHForCausalLM": self.hf_config.architectures[0] = "NemotronHForCausalLMMTP" self.hf_config.num_nextn_predict_layers = 1 def _derive_hybrid_model(self): # Use self.context_len after it has been initialized to prevent using context_len which may be None. self.is_hybrid_swa = ( is_hybrid_swa_model(self.hf_config.architectures) and not self.disable_hybrid_swa_memory ) if self.is_hybrid_swa: self.swa_attention_layer_ids, self.full_attention_layer_ids = ( get_hybrid_layer_ids( self.hf_config.architectures, self.hf_text_config, ) ) self.is_hybrid_swa_compress = self.hf_config.architectures[0] in [ "MiMoV2FlashForCausalLM", "MiMoV2MTP", ] def _derive_context_length(self, context_length: int): is_draft_model = self.is_draft_model derived_context_len = get_context_length(self.hf_text_config) if context_length is not None: if context_length > derived_context_len: reason = "Target model's" if is_draft_model else "User-specified" msg = ( f"Warning: {reason} context_length ({context_length}) is greater than the derived context_length ({derived_context_len}). " f"This may lead to incorrect model outputs or CUDA errors. Note that the derived context_length may differ from max_position_embeddings in the model's config." ) if ( envs.SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN.get() or is_in_ci() # FIXME: fix this special case ): logger.warning(msg) self.context_len = context_length if is_draft_model: self.hf_text_config.max_position_embeddings = context_length logger.warning( f"Overriding the draft model's max_position_embeddings to {context_length}." ) else: raise ValueError( f"{msg} To allow overriding this maximum, set the env var SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1" ) else: self.context_len = context_length else: self.context_len = derived_context_len # Transfer context_len to HuggingFace config so models can access it self.hf_config.context_len = self.context_len def _derive_model_shapes(self): # Unify the config keys for hf_text_config self.head_dim = getattr( self.hf_text_config, "head_dim", self.hf_text_config.hidden_size // self.hf_text_config.num_attention_heads, ) self.v_head_dim = getattr( self.hf_text_config, "v_head_dim", self.head_dim, ) self.swa_head_dim = getattr( self.hf_text_config, "swa_head_dim", self.head_dim, ) self.swa_v_head_dim = getattr( self.hf_text_config, "swa_v_head_dim", self.v_head_dim, ) # FIXME: temporary special judge for MLA architecture if ( "DeepseekV2ForCausalLM" in self.hf_config.architectures or "DeepseekV32ForCausalLM" in self.hf_config.architectures or "DeepseekV3ForCausalLM" in self.hf_config.architectures or "DeepseekV3ForCausalLMNextN" in self.hf_config.architectures or "Glm4MoeLiteForCausalLM" in self.hf_config.architectures or "GlmMoeDsaForCausalLM" in self.hf_config.architectures or "LongcatFlashForCausalLM" in self.hf_config.architectures or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures or "DotsVLMForCausalLM" in self.hf_config.architectures or "MistralLarge3ForCausalLM" in self.hf_config.architectures or "PixtralForConditionalGeneration" in self.hf_config.architectures or "MistralLarge3ForCausalLMEagle" in self.hf_config.architectures or "KimiK25ForConditionalGeneration" in self.hf_config.architectures ): self.head_dim = 256 self.attention_arch = AttentionArch.MLA self.kv_lora_rank = self.hf_text_config.kv_lora_rank self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim self.v_head_dim = self.hf_text_config.v_head_dim self.index_head_dim = ( get_nsa_index_head_dim(self.hf_text_config) if is_deepseek_nsa(self.hf_text_config) else None ) # Handle rope scaling self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) # in transformers v5, rope_scaling is just rope_parameters for backward compatibility rope_scaling = self.hf_text_config.rope_scaling if rope_scaling: # v5 uses "rope_type", v4 uses "type" rope_type = ( rope_scaling.get("rope_type") or rope_scaling.get("type") or "default" ) if rope_type != "default": self.scaling = compute_mla_mscale_scaling( rope_scaling, self.scaling ) elif "MiniCPM3ForCausalLM" in self.hf_config.architectures: self.head_dim = 128 self.attention_arch = AttentionArch.MLA self.kv_lora_rank = self.hf_config.kv_lora_rank self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim elif "DeepseekVL2ForCausalLM" in self.hf_config.architectures and getattr( self.hf_text_config, "use_mla", True ): self.head_dim = 256 self.attention_arch = AttentionArch.MLA self.kv_lora_rank = self.hf_text_config.kv_lora_rank self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim elif "KimiVLForConditionalGeneration" in self.hf_config.architectures: self.head_dim = 256 self.attention_arch = AttentionArch.MLA self.kv_lora_rank = self.hf_text_config.kv_lora_rank self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim self.v_head_dim = self.hf_text_config.v_head_dim self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim elif "KimiLinearForCausalLM" in self.hf_config.architectures: self.head_dim = 72 self.attention_arch = AttentionArch.MLA self.kv_lora_rank = self.hf_config.kv_lora_rank self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim self.v_head_dim = self.hf_config.v_head_dim self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim elif ( "BailingMoeV2_5ForCausalLM" in self.hf_config.architectures or "BailingMoeForCausalLMNextN" in self.hf_config.architectures ): self.head_dim = self.hf_text_config.head_dim self.attention_arch = AttentionArch.MLA self.kv_lora_rank = self.hf_text_config.kv_lora_rank self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim self.v_head_dim = self.hf_config.v_head_dim # Handle rope scaling with yarn self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) if self.hf_config.rope_scaling: self.scaling = compute_mla_mscale_scaling( self.hf_config.rope_scaling, self.scaling ) elif "SarvamMLAForCausalLM" in self.hf_config.architectures: self.head_dim = ( self.hf_config.qk_nope_head_dim + self.hf_config.qk_rope_head_dim ) self.attention_arch = AttentionArch.MLA self.kv_lora_rank = self.hf_config.kv_lora_rank self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim self.v_head_dim = self.hf_config.v_head_dim self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) if self.hf_config.rope_scaling: self.scaling = compute_mla_mscale_scaling( self.hf_config.rope_scaling, self.scaling ) else: if ( "MistralModel" in self.hf_config.architectures or "MixtralForCausalLM" in self.hf_config.architectures or "MistralForCausalLM" in self.hf_config.architectures ): if getattr(self, "head_dim", None) is None: self.head_dim = ( self.hf_config.hidden_size // self.hf_config.num_attention_heads ) # In transformers==4.52.3, the head_dim is null in MistralConfig if ( not hasattr(self.hf_text_config, "head_dim") or self.hf_text_config.head_dim is None ): setattr(self.hf_text_config, "head_dim", self.head_dim) elif "BaichuanForCausalLM" in self.hf_config.architectures: self.use_alibi = self.hf_config.hidden_size != 4096 self.attention_arch = AttentionArch.MHA self.num_attention_heads = self.hf_text_config.num_attention_heads self.num_key_value_heads = getattr( self.hf_text_config, "num_key_value_heads", None ) # for Dbrx and MPT models if self.hf_config.model_type in ["dbrx", "mpt"]: self.num_key_value_heads = getattr( self.hf_config.attn_config, "kv_n_heads", None ) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads self.hidden_size = self.hf_text_config.hidden_size self.num_hidden_layers = self.hf_text_config.num_hidden_layers self.num_attention_layers = self.num_hidden_layers if "LongcatFlashForCausalLM" in self.hf_config.architectures: self.num_attention_layers = self.num_hidden_layers * 2 if "IQuestLoopCoderForCausalLM" in self.hf_config.architectures: loop_num = getattr(self.hf_text_config, "loop_num", 1) self.num_attention_layers = int(self.num_hidden_layers * int(loop_num)) if "WhisperForConditionalGeneration" in self.hf_config.architectures: # Whisper has unique layer ID scheme: # - Encoder self-attention: 0 to encoder_layers-1 (no KV cache) # - Decoder self-attention: encoder_layers to encoder_layers+decoder_layers-1 (uses KV cache) # - Decoder cross-attention: encoder_layers+decoder_layers to encoder_layers+2*decoder_layers-1 # Even though cross-attention doesn't save KV cache, attention backend needs buffer to exist encoder_layers = getattr(self.hf_text_config, "encoder_layers", 0) decoder_layers = getattr( self.hf_text_config, "decoder_layers", self.num_hidden_layers ) self.num_attention_layers = encoder_layers + 2 * decoder_layers self.num_nextn_predict_layers = getattr( self.hf_text_config, "num_nextn_predict_layers", None ) self.vocab_size = self.hf_text_config.vocab_size def get_total_num_attention_heads(self) -> int: return self.num_attention_heads def get_num_attention_heads(self, tensor_parallel_size) -> int: total_num_attention_heads = self.num_attention_heads return max(1, total_num_attention_heads // tensor_parallel_size) # adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py#L289 def get_total_num_kv_heads(self) -> int: """Returns the total number of KV heads.""" # For GPTBigCode & Falcon: # NOTE: for falcon, when new_decoder_architecture is True, the # multi_query flag is ignored and we use n_head_kv for the number of # KV heads. falcon_model_types = ["falcon", "RefinedWeb", "RefinedWebModel"] new_decoder_arch_falcon = ( self.hf_config.model_type in falcon_model_types and getattr(self.hf_config, "new_decoder_architecture", False) ) if not new_decoder_arch_falcon and getattr( self.hf_text_config, "multi_query", False ): # Multi-query attention, only one KV head. # Currently, tensor parallelism is not supported in this case. return 1 # For DBRX and MPT if self.hf_config.model_type in ["mpt"]: if "kv_n_heads" in self.hf_config.attn_config: return self.hf_config.attn_config["kv_n_heads"] return self.hf_config.num_attention_heads if self.hf_config.model_type in ["dbrx"]: return getattr( self.hf_config.attn_config, "kv_n_heads", self.hf_config.num_attention_heads, ) if self.hf_config.model_type in ["nemotron-nas"]: nkvh = { self.hf_config.num_attention_heads // block.attention.n_heads_in_group for block in self.hf_config.block_configs if not block.attention.no_op } if len(nkvh) == 0: raise RuntimeError("Couldn't determine number of kv heads") if len(nkvh) > 1: raise ValueError( "Variable GQA (VGQA) is not yet supported for nemotron-nas in sglang" ) return next(iter(nkvh)) attributes = [ # For Falcon: "n_head_kv", "num_kv_heads", # For LLaMA-2: "num_key_value_heads", # For ChatGLM: "multi_query_group_num", # For Step3 "num_attention_groups", ] for attr in attributes: num_kv_heads = getattr(self.hf_text_config, attr, None) if num_kv_heads is not None: return num_kv_heads # For non-grouped-query attention models, the number of KV heads is # equal to the number of attention heads. return self.hf_text_config.num_attention_heads def get_num_kv_heads(self, tensor_parallel_size) -> int: """Returns the number of KV heads per GPU.""" total_num_kv_heads = self.get_total_num_kv_heads() # If tensor parallelism is used, we divide the number of KV heads by # the tensor parallel size. We will replicate the KV heads in the # case where the number of KV heads is smaller than the tensor # parallel size so each GPU has at least one KV head. return max(1, total_num_kv_heads // tensor_parallel_size) def get_swa_num_kv_heads(self, tensor_parallel_size) -> int: """Similar to get_num_kv_heads(), but for SWA.""" if hasattr(self.hf_text_config, "swa_num_key_value_heads"): total_num_kv_heads = self.hf_text_config.swa_num_key_value_heads return max(1, total_num_kv_heads // tensor_parallel_size) elif hasattr(self.hf_text_config, "attention_other_setting"): # For step3p5 total_num_kv_heads = self.hf_text_config.attention_other_setting.get( "num_attention_groups" ) return max(1, total_num_kv_heads // tensor_parallel_size) else: return self.get_num_kv_heads(tensor_parallel_size) # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py def _parse_quant_hf_config(self): quant_cfg = getattr(self.hf_config, "quantization_config", None) if quant_cfg is not None and not isinstance(quant_cfg, dict): quant_cfg = quant_cfg.to_dict() if quant_cfg is not None: # Identify modelopt quantization if ( "quant_method" not in quant_cfg or quant_cfg["quant_method"] == "modelopt" ): parsed_cfg = self._parse_modelopt_quant_config( {"quantization": quant_cfg} ) if parsed_cfg: quant_cfg.update(parsed_cfg) if quant_cfg is None: # compressed-tensors uses a "compression_config" key quant_cfg = getattr(self.hf_config, "compression_config", None) if quant_cfg is None: # check if is modelopt or mixed-precision model -- Both of them don't have corresponding field # in hf `config.json` but has a standalone `hf_quant_config.json` in the root directory # example: https://huggingface.co/nvidia/Llama-3.1-8B-Instruct-FP8/tree/main # example: https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/tree/main is_local = os.path.exists(self.model_path) if not is_local: # Conditional import based on SGLANG_USE_MODELSCOPE environment variable if envs.SGLANG_USE_MODELSCOPE.get(): from modelscope import HubApi, model_file_download hf_api = HubApi() else: import huggingface_hub from huggingface_hub import HfApi, hf_hub_download hf_api = HfApi() try: # In offline mode, skip file_exists check to avoid OfflineModeIsEnabled error # Instead, directly try to download/read from cache with local_files_only file_exists = False # Initialize to avoid UnboundLocalError if not huggingface_hub.constants.HF_HUB_OFFLINE: # Online mode: check if file exists before attempting download (optimization) file_exists = retry( lambda: hf_api.file_exists( self.model_path, "hf_quant_config.json" ), max_retry=2, initial_delay=1.0, max_delay=5.0, ) if not file_exists: # File doesn't exist on hub, no need to try downloading return quant_cfg # None # Download (online mode) or read from cache (offline mode) if envs.SGLANG_USE_MODELSCOPE.get(): quant_config_file = model_file_download( model_id=self.model_path, file_path="hf_quant_config.json", revision=self.revision, ) else: quant_config_file = hf_hub_download( repo_id=self.model_path, filename="hf_quant_config.json", revision=self.revision, local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, ) with open(quant_config_file) as f: quant_config_dict = json.load(f) quant_cfg = self._parse_modelopt_quant_config(quant_config_dict) except huggingface_hub.errors.LocalEntryNotFoundError: # Offline mode and file not in cache - this is normal for non-quantized models logger.debug( f"hf_quant_config.json not found in cache for {self.model_path} " "(offline mode, normal for non-quantized models)" ) except huggingface_hub.errors.OfflineModeIsEnabled: # Should not reach here after our changes, but keep for safety logger.warning( "Offline mode is enabled, skipping hf_quant_config.json check" ) except Exception as e: logger.warning( "Failed to load hf_quant_config.json for model %s: %s", self.model_path, e, ) elif os.path.exists(os.path.join(self.model_path, "hf_quant_config.json")): quant_config_file = os.path.join( self.model_path, "hf_quant_config.json" ) with open(quant_config_file) as f: quant_config_dict = json.load(f) quant_cfg = self._parse_modelopt_quant_config(quant_config_dict) return quant_cfg def _find_quant_modelslim_config(self): quant_config_file = Path(self.model_path, "quant_model_description.json") quant_cfg = None if quant_config_file.is_file(): with open(quant_config_file) as f: quant_cfg = json.load(f) # This field is required for flagless model loading but is not present in # modelslim model description, so we're adding it here manually. quant_cfg["quant_method"] = "modelslim" return quant_cfg def _parse_modelopt_quant_config(self, quant_config_dict: dict) -> Optional[dict]: """Parse ModelOpt quantization config and return the appropriate quant_method.""" json_quant_configs = quant_config_dict["quantization"] quant_algo = json_quant_configs.get("quant_algo", None) if quant_algo == "MIXED_PRECISION": architectures = getattr(self.hf_config, "architectures", []) or [] if getattr(self.hf_config, "model_type", None) == "nemotron_h" or any( arch.startswith("NemotronH") for arch in architectures ): return {"quant_method": "modelopt_mixed", "quant_algo": quant_algo} return {"quant_method": "w4afp8", "quant_algo": quant_algo} elif quant_algo and ("FP4" in quant_algo or "NVFP4" in quant_algo): return {"quant_method": "modelopt_fp4", "quant_algo": quant_algo} elif quant_algo and "FP8" in quant_algo: return {"quant_method": "modelopt_fp8", "quant_algo": quant_algo} else: return None def get_quantization_config_log_str(self) -> Optional[str]: """ Get a concise string representation of the quantization config for logging. Returns something like "quant=fp8, fmt=e4m3" or "quant=gptq, bits=4". """ try: quant_cfg = self._parse_quant_hf_config() if not quant_cfg: return None quant_method = quant_cfg.get("quant_method", "quantized") log_str = f"quant={quant_method}" # Append interesting fields if they exist for field in ["bits", "quant_algo", "fmt"]: if field in quant_cfg: log_str += f", {field}={quant_cfg[field]}" return log_str except Exception: return None def _is_already_quantized(self) -> bool: """Check if the model is already quantized based on config files.""" # Check for quantization in hf_config (config.json) if getattr(self.hf_config, "quantization_config", None) or getattr( self.hf_config, "compression_config", None ): return True # Check for HuggingFace quantization config from sglang.srt.utils import has_hf_quant_config return has_hf_quant_config(self.model_path) def _get_modelopt_quant_type(self) -> str: """Extract ModelOpt quantization type from unified quantization flag.""" if self.quantization == "modelopt_fp8": return "fp8" elif self.quantization == "modelopt_fp4": return "nvfp4" elif self.quantization == "modelopt_mixed": raise ValueError( "modelopt_mixed is only supported for pre-quantized checkpoints." ) elif self.quantization == "modelopt": # Auto-detect from model config quant_cfg = self._parse_quant_hf_config() if quant_cfg: quant_method = quant_cfg.get("quant_method", "").lower() if "fp4" in quant_method: return "fp4" elif "fp8" in quant_method: return "fp8" # Default to fp8 if can't detect return "fp8" else: return "fp8" # Default fallback def _get_sliding_window_size(self) -> Optional[int]: sliding_window_size = getattr(self.hf_text_config, "sliding_window_size", None) if sliding_window_size is None: sliding_window_size = getattr(self.hf_text_config, "sliding_window", None) return sliding_window_size def _validate_quantize_and_serve_config(self): """Validate quantize_and_serve configuration.""" if not self.quantize_and_serve: return # Check if ModelOpt quantization is specified _MODELOPT_QUANTIZATION_METHODS = [ "modelopt", "modelopt_fp8", "modelopt_fp4", "modelopt_mixed", ] modelopt_quantization_specified = ( self.quantization in _MODELOPT_QUANTIZATION_METHODS ) if not modelopt_quantization_specified: raise ValueError( "quantize_and_serve requires ModelOpt quantization (set with --quantization " f"{{{', '.join(sorted(_MODELOPT_QUANTIZATION_METHODS))}}})" ) # quantize_and_serve is disabled due to compatibility issues raise NotImplementedError( "quantize_and_serve functionality is currently disabled due to compatibility issues. " "Please use the separate quantize-then-deploy workflow instead. " "Step 1: Quantize and export model. " "Step 2: Deploy the exported model." ) # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py def _verify_quantization(self) -> None: supported_quantization = [*QUANTIZATION_METHODS] rocm_supported_quantization = [ "awq", "gptq", "fp8", "compressed_tensors", "compressed-tensors", "fbgemm_fp8", "w8a8_fp8", "petit_nvfp4", "quark", "mxfp4", "auto-round", "quark_int4fp8_moe", ] optimized_quantization_methods = [ "fp8", "marlin", "modelopt_fp8", "modelopt_fp4", "modelopt_mixed", "gptq_marlin_24", "gptq_marlin", "awq_marlin", "fbgemm_fp8", "compressed_tensors", "compressed-tensors", "experts_int8", "w8a8_int8", "w8a8_fp8", "moe_wna16", "qoq", "w4afp8", "petit_nvfp4", "quark", "modelslim", ] compatible_quantization_methods = { "modelopt_fp8": ["modelopt"], "modelopt_fp4": ["modelopt"], "modelopt_mixed": ["modelopt"], "petit_nvfp4": ["modelopt"], "w8a8_int8": ["compressed-tensors", "compressed_tensors"], "w8a8_fp8": ["compressed-tensors", "compressed_tensors"], } if self.quantization is not None: self.quantization = self.quantization.lower() # Parse quantization method from the HF and ModelSlim model config, if available. # Only one function should return config, other should return None. cfg_list = [] hf_config = self._parse_quant_hf_config() modelslim_config = self._find_quant_modelslim_config() quant_config = modelslim_config or hf_config if quant_config is not None: cfg_list.append(quant_config) # Filter out None values cfg_list = [item for item in cfg_list if item is not None] if len(cfg_list) > 1: raise ValueError( "Config list contains configs from 2 methods, must be only 1" ) quant_cfg = cfg_list[0] if cfg_list else None if quant_cfg is not None: quant_method = quant_cfg.get( "quant_method", "" if not self.quantization else self.quantization ).lower() # Detect which checkpoint is it for _, method in QUANTIZATION_METHODS.items(): quantization_override = method.override_quantization_method( quant_cfg, self.quantization ) if quantization_override: quant_method = quantization_override self.quantization = quantization_override break # Verify quantization configurations. if self.quantization is None: self.quantization = quant_method elif self.quantization != quant_method: # Check if the CLI-specified quantization is compatible with HF config's quant_method is_compatible = ( self.quantization in compatible_quantization_methods and quant_method in compatible_quantization_methods[self.quantization] ) if is_compatible: # Keep the CLI-specified quantization (e.g., modelopt_fp4) even if # HF config says "modelopt" - they are compatible logger.info( f"Using CLI-specified quantization ({self.quantization}) which is " f"compatible with HF config quant_method ({quant_method})." ) elif self.is_draft_model: # Allow auto-detection of quantization from checkpoint for draft model # only if the CLI quantization is not compatible logger.info( f"Draft model quantization ({quant_method}) differs from " f"main model quantization ({self.quantization}). " f"Using draft model's detected quantization: {quant_method}" ) self.quantization = quant_method else: raise ValueError( "Quantization method specified in the model config " f"({quant_method}) does not match the quantization " f"method specified in the `quantization` argument " f"({self.quantization})." ) # Check if the scale_fmt is ue8m0, and warn user if deepgemm is enabled for non-ue8m0 models on blackwell self.use_scale_ue8m0 = quant_cfg.get("scale_fmt", None) == "ue8m0" from sglang.srt.layers import deep_gemm_wrapper if not self.use_scale_ue8m0 and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: logger.warning( "DeepGemm is enabled but the scale_fmt of checkpoint is not ue8m0. This might cause accuracy degradation on Blackwell." ) if self.quantization is not None: if self.quantization not in supported_quantization: raise ValueError( f"Unknown quantization method: {self.quantization}. Must " f"be one of {supported_quantization}." ) if is_hip() and self.quantization not in rocm_supported_quantization: raise ValueError( f"{self.quantization} quantization is currently not " f"supported in ROCm." ) if self.quantization not in optimized_quantization_methods: # Don't warn for MXFP4 on SM100 since it has optimized kernels if not (self.quantization == "mxfp4" and is_sm100_supported()): logger.warning( "%s quantization is not fully " "optimized yet. The speed can be slower than " "non-quantized models.", self.quantization, ) def _verify_dual_chunk_attention_config(self) -> None: if hasattr(self.hf_config, "dual_chunk_attention_config"): # Try loading the sparse attention config sparse_attn_config = get_sparse_attention_config(self.model_path) if not sparse_attn_config: return self.hf_config.dual_chunk_attention_config["sparse_attention_config"] = ( sparse_attn_config ) if ( "sparse_attention_enabled" not in self.hf_config.dual_chunk_attention_config ): self.hf_config.dual_chunk_attention_config[ "sparse_attention_enabled" ] = True def _verify_transformers_version(self): import transformers from packaging import version tf_version_str = getattr(transformers, "__version__", None) if tf_version_str is None: return vision_config = getattr(self.hf_config, "vision_config", None) is_glm_46vmoe = "glm-4.6v" in self.model_path.lower() or ( vision_config is not None and getattr(vision_config, "model_type", None) == "glm4v_moe_vision" # The vision config model type for GLM-4.5v is 'glm4v_moe', # while for GLM-4.6v, it is 'glm4v_moe_vision'. ) needs_tf_v5 = is_glm_46vmoe tf_version = version.parse(tf_version_str) required_version = version.parse("5.0.0dev0") if tf_version < required_version: if needs_tf_v5: raise ValueError( f"Transformers version {tf_version_str} is not supported for model {self.model_path} " f"or model type {self.hf_config.model_type}. " "Please upgrade transformers to >= 5.0.0." ) elif not needs_tf_v5: logger.warning( f"Transformers version {tf_version_str} is used for model type {self.hf_config.model_type}. " "If you experience issues related to RoPE parameters, " "they may be due to incompatibilities between Transformers >=5.0.0 and some models. " "You can try downgrading to transformers==4.57.1 as a workaround." ) def _get_hf_eos_token_id(self) -> Optional[Set[int]]: eos_ids = getattr(self.hf_config, "eos_token_id", None) if eos_ids is not None: # it can be either int or list of int eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids) if eos_ids is None: eos_ids = set() if self.hf_generation_config: generation_eos_ids = getattr( self.hf_generation_config, "eos_token_id", None ) if generation_eos_ids: generation_eos_ids = ( {generation_eos_ids} if isinstance(generation_eos_ids, int) else set(generation_eos_ids) ) eos_ids = eos_ids | generation_eos_ids return eos_ids def get_default_sampling_params(self) -> dict[str, Any]: """ Get default sampling parameters from the model's generation config. This method returns non-default sampling parameters from the model's generation_config.json when sampling_defaults is set to "model". Returns: A dictionary containing the non-default sampling parameters. """ if self.sampling_defaults != "model": return {} if self.hf_generation_config is None: return {} config = self.hf_generation_config.to_dict() available_params = [ "repetition_penalty", "temperature", "top_k", "top_p", "min_p", ] default_sampling_params = { p: config.get(p) for p in available_params if config.get(p) is not None } return default_sampling_params def _maybe_pull_model_tokenizer_from_remote(self) -> None: """ Pull the model config files to a temporary directory in case of remote. Args: model: The model name or path. """ from sglang.srt.connector import create_remote_connector from sglang.srt.utils import is_remote_url if is_remote_url(self.model_path): logger.info("Pulling model configs from remote...") # BaseConnector implements __del__() to clean up the local dir. # Since config files need to exist all the time, so we DO NOT use # with statement to avoid closing the client. client = create_remote_connector(self.model_path) if is_remote_url(self.model_path): client.pull_files(allow_pattern=["*config.json"]) self.model_weights = self.model_path self.model_path = client.get_local_dir() # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py _STR_DTYPE_TO_TORCH_DTYPE = { "half": torch.float16, "float16": torch.float16, "float": torch.float32, "float32": torch.float32, "bfloat16": torch.bfloat16, } # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py def _get_and_verify_dtype( config: PretrainedConfig, dtype: Union[str, torch.dtype], ) -> torch.dtype: # NOTE: getattr(config, "torch_dtype", torch.float32) is not correct # because config.torch_dtype can be None. if isinstance(config, dict): config_dtype = config.get("dtype", None) or config.get("torch_dtype", None) model_type = config.get("model_type", "") else: config_dtype = getattr(config, "dtype", None) model_type = getattr(config, "model_type", "") if isinstance(config_dtype, str): config_dtype = _STR_DTYPE_TO_TORCH_DTYPE.get(config_dtype, None) if config_dtype is None: config_dtype = torch.float32 if isinstance(dtype, str): dtype = dtype.lower() if dtype == "auto": if config_dtype == torch.float32: if model_type.startswith("gemma"): if model_type == "gemma": gemma_version = "" else: gemma_version = model_type[5] logger.info( f"For Gemma {gemma_version}, we downcast float32 to bfloat16 instead " "of float16 by default. Please specify `dtype` if you " "want to use float16." ) torch_dtype = torch.bfloat16 else: # Following the common practice, we use float16 for float32 # models. torch_dtype = torch.float16 else: torch_dtype = config_dtype else: if dtype not in _STR_DTYPE_TO_TORCH_DTYPE: raise ValueError(f"Unknown dtype: {dtype}") torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype] elif isinstance(dtype, torch.dtype): torch_dtype = dtype else: raise ValueError(f"Unknown dtype: {dtype}") # Verify the dtype. if torch_dtype != config_dtype: if torch_dtype == torch.float32: # Upcasting to float32 is allowed. logger.info("Upcasting %s to %s.", config_dtype, torch_dtype) pass elif config_dtype == torch.float32: # Downcasting from float32 to float16 or bfloat16 is allowed. logger.info("Downcasting %s to %s.", config_dtype, torch_dtype) pass else: # Casting between float16 and bfloat16 is allowed with a warning. logger.warning("Casting %s to %s.", config_dtype, torch_dtype) return torch_dtype def is_generation_model(model_architectures: List[str], is_embedding: bool = False): # We have two ways to determine whether a model is a generative model. # 1. Check the model architecture # 2. check the `is_embedding` server args if ( "LlamaEmbeddingModel" in model_architectures or "MistralModel" in model_architectures or "LlamaForSequenceClassification" in model_architectures or "LlamaForSequenceClassificationWithNormal_Weights" in model_architectures or "InternLM2ForRewardModel" in model_architectures or "Qwen2ForRewardModel" in model_architectures or "Qwen3ForRewardModel" in model_architectures or "Qwen2ForSequenceClassification" in model_architectures or "Qwen3ForSequenceClassification" in model_architectures or "CLIPModel" in model_architectures or "BertModel" in model_architectures or "Contriever" in model_architectures or "BertForSequenceClassification" in model_architectures or "XLMRobertaModel" in model_architectures or "XLMRobertaForSequenceClassification" in model_architectures or "Gemma2ForSequenceClassification" in model_architectures ): return False else: return not is_embedding multimodal_model_archs = [ "CLIPModel", "DeepseekVL2ForCausalLM", "Ernie4_5_VLMoeForConditionalGeneration", "Gemma3ForConditionalGeneration", "Gemma3nForConditionalGeneration", "Glm4vForConditionalGeneration", "Glm4vMoeForConditionalGeneration", "GlmOcrForConditionalGeneration", "GlmAsrForConditionalGeneration", "Grok1VForCausalLM", "Grok1AForCausalLM", "LlavaLlamaForCausalLM", "Llama4ForConditionalGeneration", "LlavaMistralForCausalLM", "LlavaQwenForCausalLM", "LlavaForConditionalGeneration", "LlavaVidForCausalLM", "LightOnOCRForConditionalGeneration", "MiniCPMO", "MiniCPMV", "Mistral3ForConditionalGeneration", "MultiModalityCausalLM", "MllamaForConditionalGeneration", "NemotronH_Nano_VL_V2", "PixtralForConditionalGeneration", "Qwen2AudioForConditionalGeneration", "Qwen2VLForConditionalGeneration", "Qwen2_5_VLForConditionalGeneration", "Qwen3VLForConditionalGeneration", "Qwen3VLMoeForConditionalGeneration", "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration", "Qwen3OmniMoeForConditionalGeneration", "KimiVLForConditionalGeneration", "InternVLChatModel", "InternS1ForConditionalGeneration", "InternS1ProForConditionalGeneration", "Phi4MMForCausalLM", "WhisperForConditionalGeneration", "Step3VLForConditionalGeneration", "POINTSV15ChatModel", "DotsVLMForCausalLM", "DotsOCRForCausalLM", "Sarashina2VisionForCausalLM", "NVILAForConditionalGeneration", "NVILALiteForConditionalGeneration", "DeepseekOCRForCausalLM", "JetVLMForConditionalGeneration", "PaddleOCRVLForConditionalGeneration", "MiDashengLMModel", "StepVLForConditionalGeneration", "KimiK25ForConditionalGeneration", ] piecewise_cuda_graph_disabled_model_archs = [ "DeepseekV32ForCausalLM", "Qwen3NextForCausalLM", "GlmMoeDsaForCausalLM", "BailingMoeV2_5ForCausalLM", "LLaDAModelLM", ] if external_mm_model_arch := envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.get(): multimodal_model_archs.append(external_mm_model_arch) def is_multimodal_model(model_architectures: List[str]): if any( multi_model_arch in model_architectures for multi_model_arch in multimodal_model_archs ): return True else: return False def is_multimodal_gen_model(model_architectures: List[str]): return False def is_image_gen_model(model_architectures: List[str]): return False def is_audio_model(model_architectures: List[str]): models = [ "WhisperForConditionalGeneration", ] return any(model in model_architectures for model in models) def is_encoder_decoder_model(model_architectures: List[str]): models = [ "WhisperForConditionalGeneration", "MllamaForConditionalGeneration", ] return any(model in model_architectures for model in models) def is_local_attention_model(model_architectures: List[str]): return "Llama4ForConditionalGeneration" in model_architectures def is_multimodal_chunked_prefill_supported(model_architectures: List[str]): """Check if chunked prefill is supported for a MultiModal model.""" unsupported = [ "Grok1VForCausalLM", "Grok1AForCausalLM", "LlavaLlamaForCausalLM", "MllamaForConditionalGeneration", "CLIPModel", ] if any(multi_model_arch in unsupported for multi_model_arch in model_architectures): return False else: return True def is_piecewise_cuda_graph_disabled_model(model_architectures: List[str]): return any( arch in piecewise_cuda_graph_disabled_model_archs for arch in model_architectures ) def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: if scale <= 1: return 1.0 return 0.1 * mscale * math.log(scale) + 1.0 def compute_mla_mscale_scaling(rope_scaling: dict, base_scaling: float) -> float: """Compute MLA attention scaling factor from rope_scaling with mscale. Used by DeepSeek, BailingMoe, SarvamMLA and similar MLA models. Warns if 'factor' is missing from rope_scaling (common in v5 configs). """ mscale_all_dim = rope_scaling.get("mscale_all_dim", False) if "factor" not in rope_scaling: logger.warning( "rope_scaling missing 'factor', defaulting to 1.0. " "Check model accuracy.", ) scaling_factor = rope_scaling.get("factor", 1.0) mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) return base_scaling * mscale * mscale def is_hybrid_swa_model(model_architectures: List[str]): hybrid_swa_archs = { "Llama4ForConditionalGeneration", "GptOssForCausalLM", "MiMoV2FlashForCausalLM", "MiMoV2MTP", "Step3p5ForCausalLM", "Step3p5MTP", } return any(arch in hybrid_swa_archs for arch in model_architectures) def get_hybrid_layer_ids( model_architectures: List[str], hf_text_config: PretrainedConfig, ): num_hidden_layers = hf_text_config.num_hidden_layers if "Llama4ForConditionalGeneration" in model_architectures: swa_attention_layer_ids = [ i for i in range(num_hidden_layers) if (i + 1) % 4 != 0 ] full_attention_layer_ids = [ i for i in range(num_hidden_layers) if (i + 1) % 4 == 0 ] elif "GptOssForCausalLM" in model_architectures: layer_types = getattr(hf_text_config, "layer_types", None) swa_attention_layer_ids = [ i for i, x in enumerate(layer_types) if x == "sliding_attention" ] full_attention_layer_ids = [ i for i, x in enumerate(layer_types) if x == "full_attention" ] elif "MiMoV2FlashForCausalLM" in model_architectures: hybrid_layer_pattern = getattr(hf_text_config, "hybrid_layer_pattern", None) swa_attention_layer_ids = [ i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 1 ] full_attention_layer_ids = [ i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 0 ] elif "MiMoV2MTP" in model_architectures: swa_attention_layer_ids = [0] full_attention_layer_ids = [] elif "Step3p5ForCausalLM" in model_architectures: layer_types = hf_text_config.layer_types swa_attention_layer_ids = [ i for i, x in enumerate(layer_types) if x == "sliding_attention" and i < num_hidden_layers ] full_attention_layer_ids = [ i for i, x in enumerate(layer_types) if x == "full_attention" and i < num_hidden_layers ] elif "Step3p5MTP" in model_architectures: swa_attention_layer_ids = [0] full_attention_layer_ids = [] else: swa_attention_layer_ids = None full_attention_layer_ids = None return swa_attention_layer_ids, full_attention_layer_ids ================================================ FILE: python/sglang/srt/configs/modelopt_config.py ================================================ # Configuration for NVIDIA ModelOpt quantization integration from dataclasses import dataclass from typing import Optional @dataclass class ModelOptConfig: """Configuration for NVIDIA ModelOpt quantization operations. This configuration class holds parameters for ModelOpt quantization, checkpoint management, and model export operations. Args: quant: Quantization method/type (e.g., "fp8", "fp4") checkpoint_restore_path: Path to restore ModelOpt checkpoint from checkpoint_save_path: Path to save ModelOpt checkpoint to export_path: Path to export quantized model in HuggingFace format quantize_and_serve: Whether to quantize and serve in one step """ quant: Optional[str] = None checkpoint_restore_path: Optional[str] = None checkpoint_save_path: Optional[str] = None export_path: Optional[str] = None quantize_and_serve: bool = False def __post_init__(self): """Validate configuration after initialization.""" # Add any validation logic if needed pass ================================================ FILE: python/sglang/srt/configs/nano_nemotron_vl.py ================================================ # Copyright 2025 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== # Adapted from https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/blob/cb5a65ff10232128389d882d805fa609427544f1/configuration.py from typing import Any from transformers.configuration_utils import PretrainedConfig from sglang.srt.configs.nemotron_h import NemotronHConfig from sglang.srt.configs.radio import RadioConfig from sglang.srt.multimodal.internvl_utils import IMAGENET_MEAN, IMAGENET_STD def float_triplet(seq: Any): a, b, c = tuple(seq) assert ( isinstance(a, float) and isinstance(b, float) and isinstance(c, float) ), "expected three floats" return a, b, c class NemotronH_Nano_VL_V2_Config(PretrainedConfig): model_type = "NemotronH_Nano_VL_V2" is_composition = True def __init__( self, vision_config=None, llm_config=None, force_image_size: int = 512, patch_size: int = 16, downsample_ratio=0.5, template=None, ps_version="v2", image_tag_type="internvl", projector_hidden_size=4096, vit_hidden_size=1280, video_pruning_rate: float = 0.0, video_context_token: str = "